function varargout = elastic_mk(model, opt)
%ELASTIC_MK Compute mass and stiffness matrices of the elastic FE model
%   [M K DOF] = ELASTIC_MK(MODEL) Computes the mass matrix M, stiffness
%   matrix K of a FE/IE model MODEL. The output vector DOF relates the 
%   elements of the output matrices to the degrees of
%   freedom of the model.
%   [M K DOF] = MODEL_MK(MODEL)
%   [I J M K DOF] = MODE_MK(MODEL, 'ind')

% Peter Fiala and Peter Rucz
% 2009 - 2018

% Argument check
narginchk(1,2);
nargoutchk(1,6);

% Process arguments
if nargin < 2
    % Choose matrix assembly as default
    opt = 'mat';
end
    

% Preprocessing
% Drop non-elastic elements
MatTypeElastic = 2;
ElMat = model.Materials(model.Materials(:,2) == MatTypeElastic,1);
model.Elements = model.Elements(ismember(model.Elements(:,3), ElMat),:);

% drop User IDs
Elements = drop_IDs(model);
% drop PML and Infinite elements
Elements = Elements(Elements(:,2) < 100,:);

% Preallocating space for the matrices
% The elements of the sparse matrices are collected in arrays. The matrix
% values of each element are stored in non overlapping parts of the arrays.
% do not calculate infinite elements
nNode = sum(Elements(Elements(:,2),5:end)~=0, 2);   % number of nodes per element
nDim = floor(Elements(:,2) / 10);       % number of dimensions of elements
nDof = size(model.Nodes,1) * 3;
if (any(nDim ~= 3))
    error('elastic_mk: Only 3D elements are supported');
end

nVal = (nDim .* nNode).^2;               % number of matrix entries per elem

Start = [0; cumsum(nVal)];              % starting indices of elements

N = Start(end); % total number of matrix entries (non overlapped)
I  = zeros(N,1);  % preallocating row indices
J  = zeros(N,1);  % preallocating column indices
M  = zeros(N,1);  % preallocating mass values
K  = zeros(N,1);  % preallocating stiffness values

% Matrix assembly 
nElem = size(Elements,1);               % number of elements
iElem = 0;                              % 

% Note: 3D elements only
for type = [34 36 38]
    typeind = Elements(:,2) == type;
    if ~any(typeind)
        continue;
    end
    dim = floor(type/10);               % dimension number
    nnodes = mod(type,10);              % number of nodes of an element
    ndofs = nnodes*dim;
    elements = Elements(typeind,:);     % the elements of this type
    start = Start(typeind);             % start indices of elements
    [w, N, dN] = elem_quad(type);       % Gaussian quadarture for this type
    
    % Go through materials
    materials = unique(model.Materials(elements(:,3),1));
    for imat = 1 : length(materials)
        mat = materials(imat);
        rho = model.Materials(mat, 4);
        D = create_material(model, mat);
        
        % Elements of this material
        eleminds = find(elements(:,3) == mat);

        for ielemind = 1 : length(eleminds)
            ielem = eleminds(ielemind);
            elem = elements(ielem,:);
            nodes = elem(4+(1:nnodes));             % elem nodes
            vert = model.Nodes(nodes,1+(1:dim)).';  % corner vertices
            dofs = reshape(bsxfun(@plus, dim*(nodes(:)-1), 1 : dim).', [], 1);


            % Assemble element matrix
            [m, k] = elem_mk(type, vert, w, N, dN, rho, D);
            i = repmat(dofs, 1, ndofs);         % row indices
            j = i.';                                % column indices
            ind = start(ielem)+(1:ndofs^2);        % index positions
            I(ind) = i(:);
            J(ind) = j(:);
            M(ind) = m(:);
            K(ind) = k(:);
            iElem = iElem+1;
            progbar(1, nElem, iElem);
        end
    end
end

% Generate the sparse matrices from the arrays
allDof = 1 : nDof;
sel = ismember(allDof,I);            % Logical array of selected nodes
vNodes = sum(sel);                          % Number of value nodes
%index = zeros(vNodes,1);                    % Indices
%index(sel) = 1:vNodes;

% Renumber indices according to valid nodes
%I = index(I);
%J = index(J);

% Switch between output options
switch opt
    case {'mat','matrix'}
        M = sparse(I, J, M, vNodes, vNodes);
        K = sparse(I, J, K, vNodes, vNodes);
        varargout{1} = M;
        varargout{2} = K;
        if nargout > 2
            varargout{3} = DOF;
        end
    case {'ind','index'}
        varargout{1} = I;
        varargout{2} = J;
        varargout{3} = M;
        varargout{4} = K;
        varargout{5} = DOF;
        
end

end % of function model_mk


function [w, N, dNxi] = elem_quad(type)

% Compute Gaussian quadrature weights and shape function samples
nnod = mod(type, 10);

switch type
    case 12
        [xi, w] = gaussquad1(2);
    case {23, 24}
        [xi, w] = gaussquad2(4, nnod); %2
    case 34
        [xi, w] = gaussquad3(4, nnod); %2
    case {36, 38}
        [xi, w] = gaussquad3(5, nnod); %3
end

[N, dNxi] = shapefun(xi, type);
dNxi = reshape(shiftdim(dNxi,2),numel(dNxi)/nnod,nnod);

end

function [M, K] = elem_mk(type, xnode, w, N, dNxi, rho, D)

dim = floor(type/10);
ngauss = length(w);

% Coordinate transform
J = dNxi * xnode.';                 % Jacobian of the coordiate transform
dNx = zeros(size(dNxi));            % Gradient in xyz space
s = (dim * (dim+1)) / 2;
SN = zeros(s*ngauss, size(dNx,2)*dim);
DSN = zeros(s*ngauss, size(dNx,2)*dim);
j = zeros(ngauss,1);

switch dim
    case 2
        P = [1 0  0 0;
             0 0  0 1;
             0 1  1 0];
    case 3
        P = [1 0 0  0 0 0  0 0 0;
             0 0 0  0 1 0  0 0 0;
             0 0 0  0 0 0  0 0 1; 
             0 1 0  1 0 0  0 0 0;
             0 0 0  0 0 1  0 1 0;
             0 0 1  0 0 0  1 0 0];
end

for n = 1 : ngauss
    ind = dim*(n-1)+(1:dim);
    Jac = J(ind,:);
    dNx(ind,:) = (Jac) \ dNxi(ind,:);
    
    dNx_loc = zeros(size(dNx(ind,:))*dim);

    % Pre-expand shape function derivatives
    for i = 1 : dim
       row_inds = (i-1)*dim+(1 : dim);
       dNx_loc(row_inds, i : dim : end) = dNx(ind, :);
    end
    
    % Compute Strain-operator
    row_inds = s * (n-1) + (1 : s).';
    SN(row_inds, :) = P * dNx_loc;
    DSN(row_inds, :) = D * SN(row_inds,:);
    
    j(n) = abs(det(Jac));
end

% Numerical integration with Gaussian quadrature
% Use the same shape functions in all dimensions

% Expand shape functions for mass matrix
N = kron(N, eye(dim)); 

% Mass matrix assembply
q = kron(w.*j, ones(dim, 1));
M = N.'*diag(q)*N * rho;
M = .5 * (M + M.');

% Stiffness matrix assembly
q = kron(w.*j, ones(s, 1));
K = SN.' * diag(q) * DSN;
K = .5 * (K + K.');

end

function D = create_material(model, idx)
    material = model.Materials(idx, :);
    switch material(3)
        % Isotropic solid
        case 0
            D = isotropic_material(material(5), material(6));
        % Orthotropic solid
        case 1
            data = model.MaterialData{material(3)};
            [T, T_inv] = tensor_rotation(data{1});
            D = orthotropic_material(data{2}, data{3}, data{4});
            % Conversion from / to engineering notation
            En = diag([1 1 1 2 2 2]);
            En_inv = diag([1 1 1 .5 .5 .5]);
            % The material matrix 
            D = T * D * En * T_inv * En_inv;
    end
end

function [T, T_inv] = tensor_rotation(angles)
% TENSOR_ROTATION Rotation matrix for rotating symmetric tensors
%   [T, T_inv] = TENSOR_ROTATION(ANGLES) Computes the rotation matrix T and
%       its inverse T_inv that rotate the symmetric tensor by ANGLES around
%       the axes x, y, z in this order.

    Tx = @(x)[
        1           0            0       0          0        0
        0    cos(x)^2     sin(x)^2       0  -sin(2*x)        0
        0    sin(x)^2     cos(x)^2       0   sin(2*x)        0
        0           0            0  cos(x)          0  -sin(x)       
        0  sin(2*x)/2  -sin(2*x)/2       0   cos(2*x)        0
        0           0            0  sin(x)          0   cos(x)
    ];

    Ty = @(y)[
          cos(y)^2  0     sin(y)^2       0        0  -sin(2*y)
                 0  1            0       0        0          0
          sin(y)^2  0     cos(y)^2       0        0   sin(2*y)
                 0  0            0  cos(y)  -sin(y)          0
                 0  0            0  sin(y)   cos(y)          0
        sin(2*y)/2  0  -sin(2*y)/2       0        0   cos(2*y)
    ];

    Tz = @(z)[
          cos(z)^2     sin(z)^2  0  -sin(2*z)        0       0
          sin(z)^2     cos(z)^2  0   sin(2*z)        0       0
                 0            0  1          0        0       0
        sin(2*z)/2  -sin(2*z)/2  0   cos(2*z)        0       0
                 0            0  0          0   cos(z)  sin(z)
                 0            0  0          0  -sin(z)  cos(z)
    ];

    x = angles(1);
    y = angles(2);
    z = angles(3);

    T = Tz(z) * Ty(y) * Tx(x);
    T_inv = Tx(-x) * Ty(-y) * Tz(-z);
end

function D = isotropic_material(E, nu)
    mu = E / (2*(1 + nu));
    D_inv = [  1  -nu  -nu         0         0         0;
             -nu    1  -nu         0         0         0;
             -nu  -nu    1         0         0         0;
               0    0    0  2*(1+nu)         0         0;
               0    0    0         0  2*(1+nu)         0; 
               0    0    0         0         0  2*(1+nu)] / (2*mu*(1+nu));
    D = inv(D_inv);
end

function D = orthotropic_material(E, G, nu)
    % E = [E1 E2 E3]
    % G = [G23 G31 G12]
    % nu = [nu23 nu31 nu12]
    D_inv = [      1/E(1)  -nu(3)/E(1)  -nu(2)/E(3)       0       0       0
              -nu(3)/E(1)       1/E(2)  -nu(1)/E(2)       0       0       0
              -nu(2)/E(3)  -nu(1)/E(2)       1/E(3)       0       0       0
                        0            0            0  1/G(3)       0       0
                        0            0            0       0  1/G(1)       0
                        0            0            0       0       0  1/G(2)];
    D = inv(D_inv);
end
