# Amira Abdel-Rahman
# (c) Massachusetts Institute of Technology 2020


#########################################3d################################################################
function topologyOptimization3d(nelx,nely,nelz,volfrac,rmin,penal)
    anim=Animation()
    
    display("Minimum compliance problem with OC")
    display("ndes: $nelx x $nely")
    display("volfrac: $volfrac rmin: $rmin penal: $penal")
    # Max and min stiffness
    Emin=1e-9
    Emax=1.0
    nu=0.3
    # dofs:
    ndof = 3*(nelx+1)*(nely+1)*(nelz+1)
    # Allocate design variables (as array), initialize and allocate sens.
    x=volfrac * ones(Float64,nely,nelx,nelz)
    xold=copy(x)
    xPhys=copy(x)
    g=0 # must be initialized to use the NGuyen/Paulino OC approach
    dc=zeros(Float64,(nely,nelx,nelz))
    
    # FE: Build the index vectors for the for coo matrix format.
    KE=lk_H8(nu)
    nele = nelx*nely*nelz
    
    nodenrs = reshape(1:(1+nelx)*(1+nely)*(1+nelz),1+nely,1+nelx,1+nelz)
    edofVec = reshape(3*nodenrs[1:end-1,1:end-1,1:end-1].+1,nelx*nely*nelz,1)
    edofMat = repeat(edofVec,1,24).+repeat([0 1 2 3*nely.+[3 4 5 0 1 2] -3 -2 -1 3*(nely+1)*(nelx+1).+[0 1 2 3*nely.+[3 4 5 0 1 2] -3 -2 -1]],nelx*nely*nelz,1)
        
    iK = convert(Array{Int},reshape(kron(edofMat,ones(24,1))',24*24*nele,1))
    jK = convert(Array{Int},reshape(kron(edofMat,ones(1,24))',24*24*nele,1))
    
    
    
    
    # USER-DEFINED LOAD DOFs
    m= Matlab.meshgrid(nelx:nelx, 0:0, 0:nelz)
    il=m[1]
    jl=m[2]
    kl=m[3]
    loadnid = kl.*(nelx+1).*(nely+1).+il.*(nely+1).+(nely+1 .-jl)
    loaddof = 3 .*loadnid[:] .- 1; 
    # USER-DEFINED SUPPORT FIXED DOFs
    m= Matlab.meshgrid(0:0,0:nely,0:nelz)# Coordinates
    iif=m[1]
    jf=m[2]
    kf=m[3]
    fixednid = kf.*(nelx+1).*(nely+1) .+iif .*(nely .+1) .+(nely .+1 .-jf) # Node IDs
    fixeddof = [3 .*fixednid[:]; 3 .*fixednid[:].-1; 3 .*fixednid[:].-2] # DOFs
    
    
    # DEFINE LOADS AND SUPPORTS (HALF MBB-BEAM)
    F= sparse(loaddof,fill(1,length(loaddof)),fill(-1,length(loaddof)),ndof,1);
    U = zeros(ndof,1)
    alldofs = 1:ndof
    freedofs = setdiff(1:ndof,fixeddof)
    
    
    
    # Prepare filter
    #iH = ones(convert(Int,nele*(2*(ceil(rmin)-1)+1)^2),1)
    #jH = ones(Int,size(iH))
    #sH = zeros(size(iH))
    iH=[]
    jH=[]
    sH=[]
    k = 0

    for k1 = 1:nelz
        for i1 = 1:nelx
            for j1 = 1:nely
                e1 = (k1-1)*nelx*nely + (i1-1)*nely+j1
                for k2 = max(k1-(ceil(rmin)-1),1):min(k1+(ceil(rmin)-1),nelz)
                    for i2 = max(i1-(ceil(rmin)-1),1):min(i1+(ceil(rmin)-1),nelx)
                        for j2 = max(j1-(ceil(rmin)-1),1):min(j1+(ceil(rmin)-1),nely)
                            e2 = (k2-1)*nelx*nely + (i2-1)*nely+j2;
                            k = k+1;
                            append!(iH, e1  )
                            append!(jH, e2  )
                            append!(sH, max(0.0,rmin-sqrt((i1-i2)^2+(j1-j2)^2+(k1-k2)^2) ))
                            
                        end
                    end
                end
            end
        end
    end
    iH=reshape(iH,length(iH),1)
    jH=reshape(jH,length(jH),1)
    sH=reshape(convert(Array{Float64}, sH),length(sH),1)
    H = sparse(vec(iH),vec(jH),vec(sH))
    Hs = sum(H,dims=2)
    ###################################################
    loop = 0
    change = 1
    maxIter=1000
    # Start iteration
    for i =1:maxIter
        if (change > 0.01)
            # Start iteration
            loop += 1
            # FE-ANALYSIS
            sK = reshape(KE[:]*(Emin.+xPhys[:]'.^penal*(Emax-Emin)),24*24*nelx*nely*nelz,1)
            K = sparse(vec(iK),vec(jK),vec(sK)); K = (K+K')/2
            @timed U[freedofs] = K[freedofs,freedofs] \ Array(F[freedofs])
            # Objective function and sensitivity analysis
            ce = reshape(sum((U[edofMat]*KE).*U[edofMat],dims=2),nely,nelx,nelz)
            c = sum(sum(sum((Emin.+xPhys.^penal*(Emax-Emin)).*ce)))
            
            dc = -penal*(Emax-Emin)*xPhys.^(penal-1).*ce
            dv = ones(nely,nelx,nelz)
            dc[:] = H*(dc[:]./Hs)
            dv[:] = H*(dv[:]./Hs)
            # OPTIMALITY CRITERIA UPDATE OF DESIGN VARIABLES AND PHYSICAL DENSITIES
            l1 = 0; l2 = 1e9; move = 0.2; xnew = 0
            while (l2-l1)/(l1+l2) > 1e-3
                lmid = 0.5*(l2+l1)
                xnew = max.(0,max.(x.-move,min.(1,min.(x.+move,x.*sqrt.((0.0.-dc)./dv./lmid)))))
                xPhys[:] = (H*xnew[:])./Hs
                if sum(xPhys[:]) > volfrac*nelx*nely
                    l1 = lmid
                else
                    l2 = lmid
                end
            end
            change = maximum(abs.(xnew[:].-x[:]))
            x = xnew
            m=mean(xPhys[:])
            display(" It:$loop Obj:$c Vol:$m ch:$change ")
            if mod(loop,2)==0
                display(topplot3d(xPhys))
                topplot3d(xPhys)
                frame(anim)
            end

            xPhys = copy(x)
        end
    end
    return xPhys,anim
end

#########################################3d KE################################################################
function topologyOptimization3dKE(nelx,nely,nelz,volfrac,rmin,penal,CE)
    anim=Animation()
    
    display("Minimum compliance problem with OC")
    display("ndes: $nelx x $nely")
    display("volfrac: $volfrac rmin: $rmin penal: $penal")
    # Max and min stiffness
    Emin=1e-9
    Emax=1.0
    nu=0.3
    # dofs:
    ndof = 3*(nelx+1)*(nely+1)*(nelz+1)
    # Allocate design variables (as array), initialize and allocate sens.
    x=volfrac * ones(Float64,nely,nelx,nelz)
    xold=copy(x)
    xPhys=copy(x)
    g=0 # must be initialized to use the NGuyen/Paulino OC approach
    dc=zeros(Float64,(nely,nelx,nelz))
    
    # FE: Build the index vectors for the for coo matrix format.
    # KE=lk_H8(nu)

    KE = elementMatVec3D(0.5, 0.5, 0.5, CE);

    nele = nelx*nely*nelz
    
    nodenrs = reshape(1:(1+nelx)*(1+nely)*(1+nelz),1+nely,1+nelx,1+nelz)
    edofVec = reshape(3*nodenrs[1:end-1,1:end-1,1:end-1].+1,nelx*nely*nelz,1)
    edofMat = repeat(edofVec,1,24).+repeat([0 1 2 3*nely.+[3 4 5 0 1 2] -3 -2 -1 3*(nely+1)*(nelx+1).+[0 1 2 3*nely.+[3 4 5 0 1 2] -3 -2 -1]],nelx*nely*nelz,1)
        
    iK = convert(Array{Int},reshape(kron(edofMat,ones(24,1))',24*24*nele,1))
    jK = convert(Array{Int},reshape(kron(edofMat,ones(1,24))',24*24*nele,1))
    
    
    # USER-DEFINED LOAD DOFs
    m= Matlab.meshgrid(nelx:nelx, 0:0, 0:nelz)
    il=m[1]
    jl=m[2]
    kl=m[3]
    loadnid = kl.*(nelx+1).*(nely+1).+il.*(nely+1).+(nely+1 .-jl)
    loaddof = 3 .*loadnid[:] .- 1; 
    # USER-DEFINED SUPPORT FIXED DOFs
    m= Matlab.meshgrid(0:0,0:nely,0:nelz)# Coordinates
    iif=m[1]
    jf=m[2]
    kf=m[3]
    fixednid = kf.*(nelx+1).*(nely+1) .+iif .*(nely .+1) .+(nely .+1 .-jf) # Node IDs
    fixeddof = [3 .*fixednid[:]; 3 .*fixednid[:].-1; 3 .*fixednid[:].-2] # DOFs
    
    
    # DEFINE LOADS AND SUPPORTS (HALF MBB-BEAM)
    F= sparse(loaddof,fill(1,length(loaddof)),fill(-1,length(loaddof)),ndof,1);
    U = zeros(ndof,1)
    alldofs = 1:ndof
    freedofs = setdiff(1:ndof,fixeddof)
    
    
    
    # Prepare filter
    #iH = ones(convert(Int,nele*(2*(ceil(rmin)-1)+1)^2),1)
    #jH = ones(Int,size(iH))
    #sH = zeros(size(iH))
    iH=[]
    jH=[]
    sH=[]
    k = 0

    for k1 = 1:nelz
        for i1 = 1:nelx
            for j1 = 1:nely
                e1 = (k1-1)*nelx*nely + (i1-1)*nely+j1
                for k2 = max(k1-(ceil(rmin)-1),1):min(k1+(ceil(rmin)-1),nelz)
                    for i2 = max(i1-(ceil(rmin)-1),1):min(i1+(ceil(rmin)-1),nelx)
                        for j2 = max(j1-(ceil(rmin)-1),1):min(j1+(ceil(rmin)-1),nely)
                            e2 = (k2-1)*nelx*nely + (i2-1)*nely+j2;
                            k = k+1;
                            append!(iH, e1  )
                            append!(jH, e2  )
                            append!(sH, max(0.0,rmin-sqrt((i1-i2)^2+(j1-j2)^2+(k1-k2)^2) ))
                            
                        end
                    end
                end
            end
        end
    end
    iH=reshape(iH,length(iH),1)
    jH=reshape(jH,length(jH),1)
    sH=reshape(convert(Array{Float64}, sH),length(sH),1)
    H = sparse(vec(iH),vec(jH),vec(sH))
    Hs = sum(H,dims=2)
    ###################################################
    loop = 0
    change = 1
    maxIter=1000
    # Start iteration
    for i =1:maxIter
        if (change > 0.01)
            # Start iteration
            loop += 1
            # FE-ANALYSIS
            sK = reshape(KE[:]*(Emin.+xPhys[:]'.^penal*(Emax-Emin)),24*24*nelx*nely*nelz,1)
            K = sparse(vec(iK),vec(jK),vec(sK)); K = (K+K')/2
            @timed U[freedofs] = K[freedofs,freedofs] \ Array(F[freedofs])
            # Objective function and sensitivity analysis
            ce = reshape(sum((U[edofMat]*KE).*U[edofMat],dims=2),nely,nelx,nelz)
            c = sum(sum(sum((Emin.+xPhys.^penal*(Emax-Emin)).*ce)))
            
            dc = -penal*(Emax-Emin)*xPhys.^(penal-1).*ce
            dv = ones(nely,nelx,nelz)
            dc[:] = H*(dc[:]./Hs)
            dv[:] = H*(dv[:]./Hs)
            # OPTIMALITY CRITERIA UPDATE OF DESIGN VARIABLES AND PHYSICAL DENSITIES
            l1 = 0; l2 = 1e9; move = 0.2; xnew = 0
            while (l2-l1)/(l1+l2) > 1e-3
                lmid = 0.5*(l2+l1)
                xnew = max.(0,max.(x.-move,min.(1,min.(x.+move,x.*sqrt.((0.0.-dc)./dv./lmid)))))
                xPhys[:] = (H*xnew[:])./Hs
                if sum(xPhys[:]) > volfrac*nelx*nely
                    l1 = lmid
                else
                    l2 = lmid
                end
            end
            change = maximum(abs.(xnew[:].-x[:]))
            x = xnew
            m=mean(xPhys[:])
            display(" It:$loop Obj:$c Vol:$m ch:$change ")
            if mod(loop,5)==0
                display(topplot3d(xPhys))
                topplot3d(xPhys)
                frame(anim)
            end

            xPhys = copy(x)
        end
    end
    return xPhys,anim
end

#########################################################################################################

function compliantOptimization3d(nelx,nely,nelz,volfrac,rmin,penal,maxIter=500,getProblem=inverter)
    anim=Animation()

    display("Compliance sythesyz problem with OC")
    display("ndes: $nelx x $nely x $nelz")
    display("volfrac: $volfrac rmin: $rmin penal: $penal")
    # Max and min stiffness
    Emin=1e-9
    Emax=1.0
    nu=0.3
    # dofs:
    ndof = 3*(nelx+1)*(nely+1)*(nelz+1)
    # Allocate design variables (as array), initialize and allocate sens.
    x=volfrac * ones(Float64,nely,nelx,nelz)
    xold=copy(x)
    xPhys=copy(x)
    g=0 # must be initialized to use the NGuyen/Paulino OC approach
    dc=zeros(Float64,(nely,nelx,nelz))
    
    # FE: Build the index vectors for the for coo matrix format.
    KE=lk_H8(nu)
    nele = nelx*nely*nelz
    
    nodenrs = reshape(1:(1+nelx)*(1+nely)*(1+nelz),1+nely,1+nelx,1+nelz)
    edofVec = reshape(3*nodenrs[1:end-1,1:end-1,1:end-1].+1,nelx*nely*nelz,1)
    edofMat = repeat(edofVec,1,24).+repeat([0 1 2 3*nely.+[3 4 5 0 1 2] -3 -2 -1 3*(nely+1)*(nelx+1).+[0 1 2 3*nely.+[3 4 5 0 1 2] -3 -2 -1]],nelx*nely*nelz,1)
        
    iK = convert(Array{Int},reshape(kron(edofMat,ones(24,1))',24*24*nele,1))
    jK = convert(Array{Int},reshape(kron(edofMat,ones(1,24))',24*24*nele,1))
    
    
    
    
    # # USER-DEFINED LOAD DOFs
    # m= Matlab.meshgrid(nelx:nelx, 0:0, 0:nelz)
    # il=m[1]
    # jl=m[2]
    # kl=m[3]
    # loadnid = kl.*(nelx+1).*(nely+1).+il.*(nely+1).+(nely+1 .-jl)
    # loaddof = 3 .*loadnid[:] .- 1; 
    # # USER-DEFINED SUPPORT FIXED DOFs
    # m= Matlab.meshgrid(0:0,0:nely,0:nelz)# Coordinates
    # iif=m[1]
    # jf=m[2]
    # kf=m[3]
    # fixednid = kf.*(nelx+1).*(nely+1) .+iif .*(nely .+1) .+(nely .+1 .-jf) # Node IDs
    # fixeddof = [3 .*fixednid[:]; 3 .*fixednid[:].-1; 3 .*fixednid[:].-2] # DOFs
    
    
    # # DEFINE LOADS AND SUPPORTS (HALF MBB-BEAM)
    # F= sparse(loaddof,fill(1,length(loaddof)),fill(-1,length(loaddof)),ndof,1);
    # U = zeros(ndof,1)
    # alldofs = 1:ndof
    # freedofs = setdiff(1:ndof,fixeddof)
    # USER-DEFINED LOAD DOFs
    # il = [0 nelx]; jl = [nely nely]; kl = [0 0];
    # loadnid = kl .*(nelx+1) .*(nely+1) .+il .*(nely .+1) .+(nely+1 .-jl);
    # loaddof = 3 .*loadnid[:] .- 2;
    # din = loaddof[1]; dout = loaddof[2];
    # F = sparse(loaddof,[1;2],[1;-1],ndof,2);
    
    # # USER-DEFINED SUPPORT FIXED DOFs
    # # Top face
    # m = Matlab.meshgrid(0:nelx,0:nelz)
    # iif=m[1]
    # kf=m[2]
    # fixednid = kf .*(nelx+1) .*(nely+1) .+iif .*(nely+1) .+1;
    # fixeddof_t = 3 .*fixednid[:] .-1;
    # # Side face
    # m = Matlab.meshgrid(0:nelx,0:nely)
    # iif=m[1]
    # jf=m[2]
    # fixednid = iif .*(nely+1) .+(nely+1 .-jf);
    # fixeddof_s = 3*fixednid[:];
    # # Two pins
    # iif = [0;0]; jf = [0;1]; kf = [nelz;nelz];
    # fixednid = kf .*(nelx+1) .*(nely .+1) .+iif .*(nely .+1) .+(nely+1 .-jf)
    # fixeddof_p = [3 .*fixednid[:]; 3 .*fixednid[:].-1; 3 .*fixednid[:].-2] # DOFs
    # # Fixed DOFs
    # fixeddof = union(fixeddof_t,fixeddof_s);
    # fixeddof = union(fixeddof, fixeddof_p);
    # fixeddof = sort(fixeddof);
    
    # # Displacement vector
    # U = zeros(ndof,2);
    # freedofs = setdiff(1:ndof,fixeddof)

    U,F,freedofs,din,dout=getProblem(nelx,nely,nelz)

    
    
    # Prepare filter
    #iH = ones(convert(Int,nele*(2*(ceil(rmin)-1)+1)^2),1)
    #jH = ones(Int,size(iH))
    #sH = zeros(size(iH))
    iH=[]
    jH=[]
    sH=[]
    k = 0

    for k1 = 1:nelz
        for i1 = 1:nelx
            for j1 = 1:nely
                e1 = (k1-1)*nelx*nely + (i1-1)*nely+j1
                for k2 = max(k1-(ceil(rmin)-1),1):min(k1+(ceil(rmin)-1),nelz)
                    for i2 = max(i1-(ceil(rmin)-1),1):min(i1+(ceil(rmin)-1),nelx)
                        for j2 = max(j1-(ceil(rmin)-1),1):min(j1+(ceil(rmin)-1),nely)
                            e2 = (k2-1)*nelx*nely + (i2-1)*nely+j2;
                            k = k+1;
                            append!(iH, e1  )
                            append!(jH, e2  )
                            append!(sH, max(0.0,rmin-sqrt((i1-i2)^2+(j1-j2)^2+(k1-k2)^2) ))
                            
                        end
                    end
                end
            end
        end
    end

    iH=reshape(iH,length(iH),1)
    jH=reshape(jH,length(jH),1)
    sH=reshape(convert(Array{Float64}, sH),length(sH),1)
    H = sparse(vec(iH),vec(jH),vec(sH))
    Hs = sum(H,dims=2)
    ###################################################
    loop = 0
    change = 1
    
    # Start iteration
    for i =1:maxIter
        if (change > 0.01)
            # Start iteration
            loop += 1
            # FE-ANALYSIS
            sK = reshape(KE[:]*(Emin.+xPhys[:]'.^penal*(Emax-Emin)),24*24*nelx*nely*nelz,1)
            K = sparse(vec(iK),vec(jK),vec(sK)) 
            K[din,din] = K[din,din] + 0.1
            K[dout,dout] = K[dout,dout] + 0.1
            K = (K+K')/2
            
            
            # @timed U[freedofs] = K[freedofs,freedofs] \ Array(F[freedofs])
            @timed U[freedofs,:] = K[freedofs,freedofs] \ Array(F[freedofs,:])
            U1 = U[:,1]
            U2 = U[:,2]
            
            # Objective function and sensitivity analysis
            
            #ce = reshape(sum((U[edofMat]*KE).*U[edofMat],dims=2),nely,nelx,nelz)
            ce = reshape(sum((U1[edofMat]*KE).*U2[edofMat],dims=2),nely,nelx,nelz);
            
            #c = sum(sum(sum((Emin.+xPhys.^penal*(Emax-Emin)).*ce)))
            c  = U[dout,1]
            
            dc = penal*(Emax-Emin)*xPhys.^(penal-1).*ce
            dv = ones(nely,nelx,nelz)
            dc[:] = H*(dc[:]./Hs)
            dv[:] = H*(dv[:]./Hs)
            # OPTIMALITY CRITERIA UPDATE OF DESIGN VARIABLES AND PHYSICAL DENSITIES
            #l1 = 0; l2 = 1e9; move = 0.2; xnew = 0
            l1 = 0; l2 = 1e9; move = 0.1;xnew = 0
            #while (l2-l1)/(l1+l2) > 1e-3
            while (l2-l1)/(l1+l2) > 1e-4 && l2 > 1e-40
                lmid = 0.5*(l2+l1)
                #xnew = max.(0,max.(x.-move,min.(1,min.(x.+move,x.*sqrt.((0.0.-dc)./dv./lmid)))))
                xnew = max.(0,max.(x.-move,min.(1,min.(x.+move,x.*(max.(1e-10,-dc./dv./lmid)).^0.3))));

                xPhys[:] = (H*xnew[:])./Hs
                if sum(xPhys[:]) > volfrac*nelx*nely
                    l1 = lmid
                else
                    l2 = lmid
                end
            end
            change = maximum(abs.(xnew[:].-x[:]))
            x = xnew
            m=mean(xPhys[:])
            display(" It:$loop Obj:$c Vol:$m ch:$change ")
            if mod(loop,5)==0
                display(topplot3d(xPhys))
                topplot3d(xPhys)
                frame(anim)
            end

            xPhys = copy(x)
        end
    end
    return xPhys,anim
end