%--------------------------
 % @Author: Jingqiao Hu
 % @Date: 2021-12-25 11:36:43
 % @LastEditTime: 2022-01-05 12:25:18

 % pre-compute the principle stress field
%--------------------------
function [stress, pmax] = stress_field_init(nelx, nely, scalar, optDesign, vF, e, nu)

    nelx = nelx * scalar;
    nely = nely * scalar;
    
    dx = 1 / scalar; 

    [freedofs, fixeddofs, F, loadnid] = designDomain(optDesign, nelx, nely, vF);

    A11 = [12  3 -6 -3;  3 12  3  0; -6  3 12 -3; -3  0 -3 12];
    A12 = [-6 -3  0  3; -3 -6 -3 -6;  0 -3 -6  3;  3 -6  3 -6];
    B11 = [-4  3 -2  9;  3 -4 -9  4; -2 -9 -4 -3;  9  4 -3 -4];
    B12 = [ 2 -3  4 -9; -3  2  9 -2;  4  9  2  3; -9 -2  3  2];
    KE = 1/(1-nu^2)/24*([A11 A12;A12' A11]+nu*[B11 B12;B12' B11]);
    nodenrs = reshape(1:(1+nelx)*(1+nely),1+nely,1+nelx);
    edofVec = reshape(2*nodenrs(1:end-1,1:end-1)+1,nelx*nely,1);
    edofMat = repmat(edofVec,1,8)+repmat([0 1 2*nely+[2 3 0 1] -2 -1],nelx*nely,1);
    iK = reshape(kron(edofMat,ones(8,1))',64*nelx*nely,1);
    jK = reshape(kron(edofMat,ones(1,8))',64*nelx*nely,1);

    xPhys = ones(nelx * nely, 1);

    U = zeros(2*(nely+1)*(nelx+1),1);
    sK = reshape(KE(:)*(xPhys(:)'*e),64*nelx*nely,1);
    K = sparse(iK,jK,sK); K = (K+K')/2;
    U(freedofs) = K(freedofs,freedofs)\F(freedofs);

    U1 = U(edofMat)'; % [8,N]

    [~, ~, D] = material_paras(e, nu);

    optKE = 2;
    S = intS(dx/2, dx/2, D, optKE);

    gp = 1;
    s0 = S{gp} * U1; % [3,n]
    s11 = s0(1, :);
    s22 = s0(2, :);
    s12 = s0(3, :);

    % plot
%     figure;
%     [gx, gy] = meshgrid(1:nelx, 1:nely);
%     gy = flipud(gy);

    stress = zeros(size(s0,2), 2);
    pmax   = zeros(size(s0,2), 1);

    for i = 1:size(s0,2)
        se = [s11(i), s12(i);
            s12(i), s22(i)];
        [v, dv] = eig(se);

        dv = abs(diag(dv) * 100);
        [a, j] = max(dv);

        vi = v(:, j);
        stress(i, :) = vi;
        pmax(i) = a;
                
%         ellipse(gx(i), gy(i), a, dv(3-j), atan2(vi(2), vi(1))); hold on;
    end
    stress = stress ./ vecnorm(stress, 2, 2);
    pmax = pmax ./ max(pmax);

%     figure; quiver(gx(:), gy(:), s1(:,1), s1(:,2), 1, 'LineWidth',1); axis equal; hold on;
end

function S = intS(a, b, DH, optKE)
    GaussNodes = [-1/sqrt(3); 1/sqrt(3)]; GaussWeigh = [1 1]; 
    if optKE==1
        L = [1 0 0 0; 0 0 0 1; 0 1/2 1/2 0]; 
    else
        L = [1 0 0 0; 0 0 0 1; 0 1 1 0]; 
    end
    S = cell(4, 1); 

    % (-1,-1), (-1,1), (1,-1), (1,1)
    for i = 1:length(GaussNodes) 
        for j = 1:length(GaussNodes)      
            GN_x = GaussNodes(i); 
            GN_y = GaussNodes(j);            
            
            dN_y = 1/4*[-(1-GN_x) -(1+GN_x) (1+GN_x)   (1-GN_x)];            
            dN_x = 1/4*[-(1-GN_y)  (1-GN_y) (1+GN_y)  -(1+GN_y)];

            J = [dN_x; dN_y]*[ -a  a  a  -a;  
                                -b  -b  b  b]';            
            G = [inv(J) zeros(size(J)); 
                zeros(size(J)) inv(J)];      

            dN(1,1:2:8) = dN_x; dN(2,1:2:8) = dN_y;            
            dN(3,2:2:8) = dN_x; dN(4,2:2:8) = dN_y;   

            Be = L*G*dN;        
            S{j + (i-1)*2, 1} = GaussWeigh(i)*GaussWeigh(j)*det(J)*DH*Be; 
        end  
    end      
end