clc;
% This script provides a guide on how to extract designs from the dataset,
% set up an equilibrium system based on its problem formulation and calculate
% its compliance. 

file = fullfile('data.h5');

%% Specify sample index 
idx = 1;

%% Extract sample parameters
nx = double(h5read(file,['/',num2str(idx),'/nx']));         % number of elements in x-direction
ny = double(h5read(file,['/',num2str(idx),'/ny']));         % number of elements in y-direction
nElem = h5read(file,['/',num2str(idx),'/nElem']);           % total number of elements
nNodes = h5read(file,['/',num2str(idx),'/nNodes']);         % total number of nodes 
scale = h5read(file,['/',num2str(idx),'/scale']);           % Scaling parameters for the L-beam design domain

%% Extract designs
x = double(h5read(file,['/',num2str(idx),'/design/x']));    % non-fail-safe design
y = double(h5read(file,['/',num2str(idx),'/design/y']));    % fail-safe design

%% Extract supports and determine free degrees of freedom
supports = h5read(file,['/',num2str(idx),'/supports']);     % support conditions [node index, direction, displacement]
freedofs = freeDegreesOfFreedom(nNodes,supports);           % free degrees of freedom 

%% Extract loads and set up load vector 
loads = h5read(file,['/',num2str(idx),'/loads']);           % load cases [node index, direction, magnitude]
F = loadVector(nNodes,loads);                               % load vector 

%% Set up stiffness matrix
% Material Parameters
nu = 0.3;                                                   % Poisson's ratio
Emin = 1e-9;                                                % Void Young's modulus
E0 = 1;                                                     % Material Young's modulus
% SIMP
penal = 3;                                                  % Interpolation parameter
% element stiffness matrix 
KE = elementStiffnessMatrix;                                % element stiffness matrix 
% element degrees of freedom 
edofMat = elementDegreesOfFreedom(nElem,nx,ny,scale);       % element degrees of freedom [element index, node index]
%% global stiffness matrix          
K = assembleStiffnessMatrix(y,E0,Emin,penal,KE,edofMat);    % stiffness matrix 

%% Solve equilibrium system  
U = zeros(size(K,1),size(F,2));
U(freedofs,:) = K(freedofs,freedofs)\F(freedofs,:);         % Nodal displacements

%% Calculate structural compliance and its sensitivity
[c,dc] = compliance(y,U,edofMat,KE,E0,Emin,penal);          % Compliance and gradient 

%% Functions
function freedofs = freeDegreesOfFreedom(nNodes,supports)
    supports = sortrows(supports);
    dpnsum = cumsum(2*ones(nNodes,1))-1;
    fixeddofs = dpnsum(supports(:,1))+(supports(:,2))-1;
    alldofs = 1:2*nNodes;
    freedofs = setdiff(alldofs,fixeddofs);
end

function F = loadVector(nNodes,loads)
    numLoads = size(loads,1);
    dpnsum = cumsum(2*ones(nNodes,1))-1;
    indexL = double(dpnsum(loads(:,1)))+loads(:,2)-1; 
    indexLC = repelem(1:numLoads,1)';
    F = sparse(indexL,indexLC,loads(:,3),2*nNodes,numLoads);
end

function K = assembleStiffnessMatrix(v,E0,Emin,penal,KE,edofMat)
    iK = reshape(kron(edofMat,ones(8,1))',1,64*size(edofMat,1));
    jK = reshape(kron(edofMat,ones(1,8))',1,64*size(edofMat,1),1);
    sK = reshape(KE(:)*(Emin+v'.^penal*(E0-Emin)),64*size(edofMat,1),1);
    K = sparse(iK,jK,sK);
    K = (K+K')/2;
end

function KE = elementStiffnessMatrix
    nu = 0.3;
    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]);
end

function [c,dc] = compliance(v,U,edofMat,KE,E0,Emin,penal)
    c=0;
    dc = 0;
    for i = 1:size(U,2)
        Ui = U(:,i);
        ce = reshape(sum((Ui(edofMat)*KE).*Ui(edofMat),2),[],1);
        c = c + sum(sum((Emin+v.^penal*(E0-Emin)).*ce));
        dc = dc - penal*(E0-Emin)*v.^(penal-1).*ce;
    end    
end

function edofMat = elementDegreesOfFreedom(nElem,nx,ny,scale)
    if nElem ~= ny*nx
        a = scale(1)*nx;
        b = scale(2)*ny;
        nodenrs = zeros(nx+1,ny+1);
        nodenrs(1:nx*scale(1)+1,1:ny+1) = reshape(1:(nx*scale(1)+1)*(ny+1),ny+1,nx*scale(1)+1)';
        nodenrs(nx*scale(1)+2:end,1:ny*scale(2)+1) = reshape((nx*scale(1)+1)*(ny+1)+1:(ny+1)*(nx+1)-(nx-a)*(ny-b),ny*scale(2)+1,nx-nx*scale(1))';        
        edofVec1 = 2*nodenrs(1:nx*scale(1),1:ny)+1;
        edofVec2 = 2*nodenrs(nx*scale(1)+1:end-1,1:ny*scale(2))+1;
        edofVec1 = edofVec1(:);
        edofVec2 = edofVec2(:);
        edofMat1 = repmat(edofVec1,1,8)+repmat([-2 -1 2*ny+[0 1 2 3] 0 1], size(edofVec1,1),1);
        edofMat2 = repmat(edofVec2,1,8)+repmat([-2 -1 2*ny*scale(2)+[0 1 2 3] 0 1], size(edofVec2,1),1);
        edofMat2(1:nx-nx*scale(1):end,3:6) = edofMat2(1:nx-nx*scale(1):end,3:6)+2*(ny-ny*scale(2));  
        edofMat = [edofMat1;edofMat2];
    else        
        nodenrs = reshape(1:(1+nx)*(1+ny),1+nx,1+ny);
        edofVec = reshape(2*nodenrs(1:end-1,1:end-1)+1,nx*ny,1);
        edofMat = repmat(edofVec,1,8)+repmat([-2 -1 0 1 2*nx+[2 3 0 1]], nx*ny,1);
    end
end