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

# Topology Optimization Implementation in Julia from various sources
# based on https://paulino.ce.gatech.edu/conferences/papers/11pereira_anefficientandcompact.pdf and http://www.topopt.mek.dtu.dk/apps-and-software and https://github.com/blademwang11/Topopt/blob/master/top.jl


function topologyOptimization(nelx,nely,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
    # dofs:
    ndof = 2*(nelx+1)*(nely+1)
    # Allocate design variables (as array), initialize and allocate sens.
    x=volfrac * ones(Float64,nely,nelx)
    xold=copy(x)
    xPhys=copy(x)
    g=0 # must be initialized to use the NGuyen/Paulino OC approach
    dc=zeros(Float64,(nely,nelx))
    
    # FE: Build the index vectors for the for coo matrix format.
    KE=lk()
    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 = repeat(edofVec,1,8).+repeat([0 1 2*nely.+[2 3 0 1] -2 -1],nelx*nely,1)
    iK = convert(Array{Int},reshape(kron(edofMat,ones(8,1))',64*nelx*nely,1))
    jK = convert(Array{Int},reshape(kron(edofMat,ones(1,8))',64*nelx*nely,1))
    # DEFINE LOADS AND SUPPORTS (HALF MBB-BEAM)
    F = sparse([2],[1],[-1.0],2*(nely+1)*(nelx+1),1)
    U = zeros(2*(nely+1)*(nelx+1),1)
    fixeddofs = union(1:2:2*(nely+1),2*(nelx+1)*(nely+1))
    alldofs = 1:(2*(nely+1)*(nelx+1))
    freedofs = setdiff(alldofs,fixeddofs)
    # Prepare filter
    iH = ones(convert(Int,nelx*nely*(2*(ceil(rmin)-1)+1)^2),1)
    jH = ones(Int,size(iH))
    sH = zeros(size(iH))
    k = 0;
    for i1 = 1:nelx
        for j1 = 1:nely
            e1 = (i1-1)*nely+j1
            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 = (i2-1)*nely+j2
                    k = k+1
                    iH[k] = e1
                    jH[k] = e2
                    sH[k] = max(0,rmin-sqrt((i1-i2)^2+(j1-j2)^2))
                end
            end
        end
    end
    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)),64*nelx*nely,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)
            c = sum(sum((Emin.+xPhys.^penal*(Emax-Emin)).*ce))
            
            dc = -penal*(Emax-Emin)*xPhys.^(penal-1).*ce
            dv = ones(nely,nelx)
            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 loop<20||mod(loop,10)==0
                xxx=1 .- clamp.(xPhys,0,1)
                display(heatmap(xxx,xaxis=nothing,yaxis=nothing,legend=nothing,fc=:grays,clims=(0.0, 1.0),aspect_ratio=:equal))
                heatmap(xxx,xaxis=nothing,yaxis=nothing,legend=nothing,fc=:grays,clims=(0.0, 1.0),aspect_ratio=:equal)
                frame(anim)
            end

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

function topologyOptimizationMMA(nelx,nely,volfrac,rmin,penal,maxEval)
    
    display("Minimum compliance problem with MMA")
    display("ndes: $nelx x $nely")
    display("volfrac: $volfrac rmin: $rmin penal: $penal")
    # Max and min stiffness
    Emin=1e-9
    Emax=1.0
    # dofs:
    ndof = 2*(nelx+1)*(nely+1)
    # Allocate design variables (as array), initialize and allocate sens.
    x=volfrac * ones(Float64,nely,nelx)
    xold=copy(x)
    xPhys=copy(x)
    g=0 # must be initialized to use the NGuyen/Paulino OC approach
    dc=zeros(Float64,(nely,nelx))

    # FE: Build the index vectors for the for coo matrix format.
    KE=lk()
    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 = repeat(edofVec,1,8).+repeat([0 1 2*nely.+[2 3 0 1] -2 -1],nelx*nely,1)
    iK = convert(Array{Int},reshape(kron(edofMat,ones(8,1))',64*nelx*nely,1))
    jK = convert(Array{Int},reshape(kron(edofMat,ones(1,8))',64*nelx*nely,1))
    # DEFINE LOADS AND SUPPORTS (HALF MBB-BEAM)
    F = sparse([2],[1],[-1.0],2*(nely+1)*(nelx+1),1)
    U = zeros(2*(nely+1)*(nelx+1),1)
    fixeddofs = union(1:2:2*(nely+1),2*(nelx+1)*(nely+1))
    alldofs = 1:(2*(nely+1)*(nelx+1))
    freedofs = setdiff(alldofs,fixeddofs)
    # Prepare filter
    iH = ones(convert(Int,nelx*nely*(2*(ceil(rmin)-1)+1)^2),1)
    jH = ones(Int,size(iH))
    sH = zeros(size(iH))
    k = 0;
    for i1 = 1:nelx
        for j1 = 1:nely
            e1 = (i1-1)*nely+j1
            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 = (i2-1)*nely+j2
                    k = k+1
                    iH[k] = e1
                    jH[k] = e2
                    sH[k] = max(0,rmin-sqrt((i1-i2)^2+(j1-j2)^2))
                end
            end
        end
    end
    H = sparse(vec(iH),vec(jH),vec(sH))
    Hs = sum(H,dims=2);  
    
    nel=nely*nelx
    

    function FA(x::Vector, grad::Vector)

        xPhys=reshape(x,nely,nelx)

        sK = reshape(KE[:]*(Emin.+xPhys[:]'.^penal*(Emax-Emin)),64*nelx*nely,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)
        c = sum(sum((Emin.+xPhys.^penal*(Emax-Emin)).*ce))

        dc = -penal*(Emax-Emin)*xPhys.^(penal-1).*ce
        dc[:] = H*(dc[:]./Hs)

        grad[:] .= dc[:]

        return c
    end

    function G(x::Vector, grad::Vector)
        dv = ones(nely,nelx)
        dv[:] = H*(dv[:]./Hs)

        grad[:] .=  dv[:]

        return (sum(x) - volfrac*nel)
    end

    FA(ones(nel)*volfrac, fill(volfrac,nel))
    G(ones(nel)*volfrac, fill(volfrac,nel))

    opt = Opt(:LD_MMA, nel)
    opt.lower_bounds = fill(1e-9,nel)
    opt.upper_bounds = fill(1,nel)
    opt.xtol_rel = 1e-4
    opt.maxeval = maxEval

    opt.min_objective = FA
    inequality_constraint!(opt, (x,gg) -> G(x,gg), 1e-4)

    display(@time (minf,minx,ret) = optimize(opt, ones(nel).*volfrac))
    numevals = opt.numevals # the number of function evaluations
    display("got $minf after $numevals iterations (returned $ret)")
    
    xPhys=reshape(minx,nely,nelx)
    display(heatmap(xPhys, aspect_ratio=:equal, legend=false, axis=nothing, foreground_color_subplot=colorant"white",fc=:grays))


    return xPhys
end
#########################################################################################################

function CompliantTopologyOptimization(nelx,nely,volfrac,rmin,penal,maxIter,Load,Support,Spring,DOut)
    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
    change = 1
    # dofs:
    ndof = 2*(nelx+1)*(nely+1)
    # Allocate design variables (as array), initialize and allocate sens.
    x=volfrac * ones(Float64,nely,nelx)
    xold=copy(x)
    xPhys=copy(x)
    g=0 # must be initialized to use the NGuyen/Paulino OC approach
    dc=zeros(Float64,(nely,nelx))
    
    # FE: Build the index vectors for the for coo matrix format.
    KE=lk()
    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 = repeat(edofVec,1,8).+repeat([0 1 2*nely.+[2 3 0 1] -2 -1],nelx*nely,1)
    iK = convert(Array{Int},reshape(kron(edofMat,ones(8,1))',64*nelx*nely,1))
    jK = convert(Array{Int},reshape(kron(edofMat,ones(1,8))',64*nelx*nely,1))
    
    
    # DEFINE LOADS AND SUPPORTS 
    F = sparse(2 .*Load[:,1] .-2 .+ Load[:,2],ones(size(Load,1)),Load[:,3],2*(nely+1)*(nelx+1),2)
    DofDOut = 2 * DOut[1] - 2 +DOut[2]; #only one
    F=Array(F)
    F[DofDOut,2]=-1
    fixeddofs = 2 .*Support[:,1] .-2 .+ Support[:,2]
    NSpring = size(Spring,1);
    s = sparse(2 .*Spring[:,1] .-2 .+ Spring[:,2],ones(size(Spring,1)),Spring[:,3],2*(nely+1)*(nelx+1),2)
    m=Array(s)[:,1]
    S= sparse(diagm(m))
    
    U = zeros(2*(nely+1)*(nelx+1),2)

    alldofs = 1:(2*(nely+1)*(nelx+1))
    freedofs = setdiff(alldofs,fixeddofs)
    
    # Prepare filter
    iH = ones(convert(Int,nelx*nely*(2*(ceil(rmin)-1)+1)^2),1)
    jH = ones(Int,size(iH))
    sH = zeros(size(iH))
    k = 0;
    for i1 = 1:nelx
        for j1 = 1:nely
            e1 = (i1-1)*nely+j1
            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 = (i2-1)*nely+j2
                    k = k+1
                    iH[k] = e1
                    jH[k] = e2
                    sH[k] = max(0,rmin-sqrt((i1-i2)^2+(j1-j2)^2))
                end
            end
        end
    end
    H = sparse(vec(iH),vec(jH),vec(sH))
    Hs = sum(H,dims=2)
    
    change= 1
    loop = 0
    changeStable=0
    # Start iteration
    for penal in 1.0:0.5:4
        display(" Penalty: $penal")
        loop = 0
        for i =1:maxIter
            if (change < 0.01)
                changeStable+=1
                display(" Change Stable for $changeStable iterations")
            else
                changeStable=0
            end
            if (changeStable<10)
                # Start iteration
                loop += 1
                # FE-ANALYSIS
                sK = reshape(KE[:]*(Emin.+xPhys[:]'.^penal*(Emax-Emin)),64*nelx*nely,1)
                K = sparse(vec(iK),vec(jK),vec(sK))+S;K = (K+K')/2
                @timed U[freedofs,:] = K[freedofs,freedofs] \ Array(F[freedofs,:])
                U1=U[:,1]
                U2=U[:,2]
                # Objective function and sensitivity analysis
                #ce = reshape(sum((U1[edofMat]*KE).*U2[edofMat],dims=2),nely,nelx)

                ce=reshape(-sum((U1[edofMat]*KE).*U2[edofMat],dims=2),nely,nelx)

                #c = sum((Emin.+xPhys[:].^penal*(Emax-Emin)).*ce)

                c=U[DofDOut,1]

                dc = -penal.*(Emax-Emin).*xPhys.^(penal-1).*ce
                dv = ones(nely,nelx)
                dc[:] = H*(dc[:]./Hs)
                dv[:] = H*(dv[:]./Hs)

                # OPTIMALITY CRITERIA UPDATE OF DESIGN VARIABLES AND PHYSICAL DENSITIES
                l1 = 0; l2 = 1e9; move = 0.05; xnew = 0 #move=0.2
                while (l2-l1)/(l2+l1) > 1e-4 && l2>1e-40
                #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.(max.(1e-10,-dc./dv./lmid))))))
                    #xnew = max.(0,max.(x.-move,min.(1,min.(x.+move,x.*sqrt.(-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

                # print result
                m=mean(xPhys[:])
                display(" It:$loop Obj:$c Vol:$m ch:$change ")
                if loop<20||mod(loop,5)==0
                    xxx=vcat(xPhys[end:-1:1,:],xPhys)
                    xxx=1 .- xxx
                    display(heatmap(xxx,xaxis=nothing,yaxis=nothing,legend=nothing,fc=:grays,clims=(0.0, 1.0)))
                    heatmap(xxx,xaxis=nothing,yaxis=nothing,legend=nothing,fc=:grays,clims=(0.0, 1.0))
                    frame(anim)
                end

                xPhys = copy(x)
            end
        end
    end
    # heatmap(xPhys, legend=false, axis=nothing, foreground_color_subplot=colorant"white",fc=:grays,clims=(0.0, 1.0))
    # heatmap(xPhys, legend=false,fc=:grays,clims=(0.0, 1.0))
    return xPhys,anim
end