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

function getDataFromSetup3D(setup,scale)
    # -----------------------------------------------------------------------
    # A(nel,1)        = cross-sectional area of elements
    # E(nel,1)        = Young's modulus of elements
    # f(ndf,nnp)      = prescribed nodal forces
    # g(ndf,nnp)      = prescribed nodal displacements
    # idb(ndf,nnp)    = 1 if the degree of freedom is prescribed, 0 otherwise
    # ien(nen,nel)    = element connectivities
    # ndf             = number of degrees of freedom per node
    # nel             = number of elements
    # nen             = number of element equations
    # nnp             = number of nodes
    # nsd             = number of spacial dimensions
    # xn(nsd,nnp)     = nodal coordinates

    # =======================================================================
    nodes=setup["nodes"]
    edges=setup["edges"]


    # ---- MESH -------------------------------------------------------------
    nsd =  3;  # number of spacial dimensions
    ndf =  6;  # number of degrees of freedom per node 
    nen =  2;  # number of element nodes

    nel =  length(edges);  # number of elements 
    nnp =  length(nodes);  # number of nodal points


    # ---- NODAL COORDINATES ------------------------------------------------
    # xn(i,N) = coordinate i for node N
    # N = 1,...,nnp
    # i = 1,...,nsd
    # -----------------------------------------------------------------------

    xn         = zeros(nsd,nnp);
    for i in 1:nnp
        xn[1:nsd, i]  = [(nodes[i]["position"]["x"]/scale)  (nodes[i]["position"]["y"]/scale) (nodes[i]["position"]["z"]/scale)]';
    end
    
    # ---- NODAL COORDINATES ------------------------------------------------
    # ien(a,e) = N
    # N: global node number - N = 1,...,nnp
    # e: element number     - e = 1,...,nel
    # a: local node number  - a = 1,...,nen
    # -----------------------------------------------------------------------

    ien = zeros(nen,nel);
    for i in 1:nel
        ien[1:2,i]  = [(edges[i]["source"]+1)   (edges[i]["target"]+1)]' ;
    end


    len=zeros(nel);
    for i in 1:nel
        x1=(nodes[(edges[i]["source"]+1)]["position"]["x"]/scale)
        x2=(nodes[(edges[i]["target"]+1)]["position"]["x"]/scale)

        y1=(nodes[(edges[i]["source"]+1)]["position"]["y"]/scale)
        y2=(nodes[(edges[i]["target"]+1)]["position"]["y"]/scale)

        z1=(nodes[(edges[i]["source"]+1)]["position"]["z"]/scale)
        z2=(nodes[(edges[i]["target"]+1)]["position"]["z"]/scale)

        len[i]  = sqrt((x1-x2)^2+(y1-y2)^2+(z1-z2)^2);
    end
    

    # ---- MATERIAL PROPERTIES ----------------------------------------------
    # E(e) = E_mat
    # e: element number     - e = 1,...,nel
    # -----------------------------------------------------------------------
    E_mat = 29000*1000;         # Young's modulus # lbf/in^2
    E = E_mat*ones(nel,1);
    # todo change to make it parameter

    # ---- GEOMETRIC PROPERTIES ---------------------------------------------
    # A(e) = A_bar
    # e: element number     - e = 1,...,nel
    # -----------------------------------------------------------------------
    #A = ones(nel);
    #A[:] .= 400; # mm^2

    # ---- BOUNDARY CONDITIONS ----------------------------------------------
    # prescribed displacement flags (essential boundary condition)
    #
    # idb(i,N) = 1 if the degree of freedom i of the node N is prescribed
    #          = 0 otherwise
    #
    # 1) initialize idb to 0
    # 2) enter the flag for prescribed displacement boundary conditions
    # -----------------------------------------------------------------------
    idb = zeros(ndf,nnp);
    for i in 1:nnp
        if nodes[i]["restrained_degrees_of_freedom"][1]
            idb[1,i]=1; 
        end
        if nodes[i]["restrained_degrees_of_freedom"][2]
            idb[2,i]=1; 
        end
        if nodes[i]["restrained_degrees_of_freedom"][3]
            idb[3,i]=1; 
        end
        if nodes[i]["restrained_degrees_of_freedom"][4]
            idb[4,i]=1; 
        end
        if nodes[i]["restrained_degrees_of_freedom"][5]
            idb[5,i]=1; 
        end
        if nodes[i]["restrained_degrees_of_freedom"][6]
            idb[6,i]=1; 
        end
    end


    # ---- BOUNDARY CONDITIONS: PRESCRIBED NODAL DISPLACEMENTS --------------
    # g(i,N) = prescribed displacement magnitude
    # N = 1,...,nnp
    # i = 1,...,nsd
    #
    # 1) initialize g to 0
    # 2) enter the values
    # -----------------------------------------------------------------------
    g = zeros(ndf,nnp);


    # ---- BOUNDARY CONDITIONS: PRESCRIBED NODAL FORCES ---------------------
    # f(i,N) = prescribed force magnitude
    # N = 1,...,nnp
    # i = 1,...,nsd
    #
    # 1) initialize f to 0
    # 2) enter the values
    # -----------------------------------------------------------------------
    P = 20.0* 10000;     # lbf
    f = zeros(ndf,nnp);
    for i in 1:nnp
        f[1,i] = nodes[i]["force"]["x"]*P;
        f[2,i] = nodes[i]["force"]["y"]*P;
        f[3,i] = nodes[i]["force"]["z"]*P;
    end

    Ls=[]
    for i in 1:nnp
        if (nodes[i]["fixedDisplacement"]["x"]!=0)||(nodes[i]["fixedDisplacement"]["y"]!=0)||(nodes[i]["fixedDisplacement"]["z"]!=0)
            L=zeros(ndf,nnp);
            L[1,i] = nodes[i]["fixedDisplacement"]["x"];
            L[2,i] = nodes[i]["fixedDisplacement"]["y"];
            L[3,i] = nodes[i]["fixedDisplacement"]["z"];
            append!(Ls,[L]);
        end
    end
    
    A=ones(nel)
    
    # ---- NUMBER THE EQUATIONS ---------------------------------------------
    # line 380
    id,neq = number_eq(idb,ndf,nnp);
    
    # ---- FORM THE ELEMENT STIFFNESS MATRICES ------------------------------
    # line 324
    nee = ndf*nen;                   # number of element equations
    
    Ke  = zeros(nee,nee,nel);
    # Te  = zeros(nen*1,nen*nsd,nel);  # *1 is specific to frame
    Te  = zeros(nee,nee,nel);  # *1 is specific to frame

    for i = 1:nel
        Ke[:,:,i],Te[:,:,i] = Ke_frame(A[i],E[i],ien[:,i],nee,nsd,xn);
    end    

    return E,f,g,idb,ien,ndf,nel,nen,nnp,nsd,xn,len,Ke,Te,Ls
end
# =======================================================================

function getSensitivites(Ke0,dcomp,ien,nel,len,λ)

    dfdA=zeros(nel)
    dgdA=zeros(nel)


    for i in 1:nel
        de=[dcomp[1,Int(ien[1,i])] ,dcomp[2,Int(ien[1,i])], dcomp[3,Int(ien[1,i])],dcomp[4,Int(ien[1,i])] ,dcomp[5,Int(ien[1,i])], dcomp[6,Int(ien[1,i])],
            dcomp[1,Int(ien[2,i])] ,dcomp[2,Int(ien[2,i])] ,dcomp[3,Int(ien[2,i])],dcomp[4,Int(ien[2,i])] ,dcomp[5,Int(ien[2,i])] ,dcomp[6,Int(ien[2,i])]
            ];
        # de=[dcomp[1,Int(ien[1,i])],dcomp[2,Int(ien[1,i])],dcomp[3,Int(ien[1,i])] ,dcomp[1,Int(ien[2,i])] ,dcomp[2,Int(ien[2,i])] ,dcomp[3,Int(ien[2,i])]]
        λe=[λ[1,Int(ien[1,i])], λ[2,Int(ien[1,i])],λ[3,Int(ien[1,i])],λ[4,Int(ien[1,i])], λ[5,Int(ien[1,i])],λ[6,Int(ien[1,i])],
            λ[1,Int(ien[2,i])] ,λ[2,Int(ien[2,i])],λ[3,Int(ien[2,i])],λ[4,Int(ien[2,i])] ,λ[5,Int(ien[2,i])],λ[6,Int(ien[2,i])]
            ]
        # println(de)
        # println(Ke0[:,:,i])
        dKedA=Ke0[:,:,i]
        dgdA[i]=(-1.0 .*λe)'*dKedA*de
        dfdA[i]=len[i]
    end
    return dfdA,dgdA
end
# =======================================================================

function getAdjoint(K,dcomp,L,free)
    L1=L[free]
    λ1=K\L1[:]
    λ=zeros(size(dcomp))
    λ[free].=reshape(λ1,size(L1))
    return λ,L1,λ1
end
# =======================================================================

function getSensitivitesSIMP(Ke0,dcomp,ien,nel,len,λ,η,X)

    dfdA=zeros(nel)
    dgdA=zeros(nel)


    for i in 1:nel
        de=[dcomp[1,Int(ien[1,i])] ,dcomp[2,Int(ien[1,i])], dcomp[3,Int(ien[1,i])],dcomp[4,Int(ien[1,i])] ,dcomp[5,Int(ien[1,i])], dcomp[6,Int(ien[1,i])],
            dcomp[1,Int(ien[2,i])] ,dcomp[2,Int(ien[2,i])] ,dcomp[3,Int(ien[2,i])],dcomp[4,Int(ien[2,i])] ,dcomp[5,Int(ien[2,i])] ,dcomp[6,Int(ien[2,i])]
            ];
        # de=[dcomp[1,Int(ien[1,i])],dcomp[2,Int(ien[1,i])],dcomp[3,Int(ien[1,i])] ,dcomp[1,Int(ien[2,i])] ,dcomp[2,Int(ien[2,i])] ,dcomp[3,Int(ien[2,i])]]
        λe=[λ[1,Int(ien[1,i])], λ[2,Int(ien[1,i])],λ[3,Int(ien[1,i])],λ[4,Int(ien[1,i])], λ[5,Int(ien[1,i])],λ[6,Int(ien[1,i])],
            λ[1,Int(ien[2,i])] ,λ[2,Int(ien[2,i])],λ[3,Int(ien[2,i])],λ[4,Int(ien[2,i])] ,λ[5,Int(ien[2,i])],λ[6,Int(ien[2,i])]
            ]
        # println(de)
        # println(Ke0[:,:,i])
        dKedA=Ke0[:,:,i].*(η).*X[i]^(η-1)
        dgdA[i]=(-1.0 .*λe)'*dKedA*de
        dfdA[i]=len[i]
    end
    return dfdA,dgdA
end
# =======================================================================


function optimizeCompliantMechanism1(problem,Ls,free,dmax,totalVolFactor,maxeval=500)
    
    E,f,g,idb,ien,ndf,nel,nen,nnp,nsd,xn,len,Ke,Te=problem
    
    nel=length(len)
    
    function FA(x::Vector, grad::Vector)

        K,F,d,stress,dcomp,g=FEM_frame(problem,x);
        λ,L1,λ1=getAdjoint(copy(K),copy(dcomp),copy(Ls[1]),copy(free))

        grad[:] .=getSensitivites(Ke,dcomp,ien,nel,len,copy(λ))[2]

        # g=(L1[:]'*d)[1] .- dmax - (λ1'*(K*d -F))[1]

        goal=L1[:]'*d -λ1'*K*d +λ1'*F .-dmax
        # display(goal )
        # display(sum(x) )
        
        return goal[1]
    end

    function GA(x::Vector, grad::Vector)
        grad[:] .=len[:]
        return (sum(x .* len )) - totalVolFactor
    end
    FA(ones(nel), fill(0.25,nel))
    GA(ones(nel), fill(0.25,nel))
    

    opt = Opt(:LD_MMA, nel)
    opt.lower_bounds = fill(1e-3,nel)
    opt.xtol_rel = 1e-6
    opt.maxeval = maxeval

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

    display(@time (minf,minx,ret) = optimize(opt, ones(nel)))
    numevals = opt.numevals # the number of function evaluations
    display("got $minf after $numevals iterations (returned $ret)")
    
    return minx
end
# =======================================================================

function optimizeCompliantMechanism(problem,Ls,free,dmax,amax,maxeval=500,SIMP=false)
    
    E,f,g,idb,ien,ndf,nel,nen,nnp,nsd,xn,len,Ke,Te=problem
    
    nel=length(len)

    η=3.0
    
    function GA(x::Vector, grad::Vector,num::Int)

        K,F,d,stress,dcomp,g=FEM_frame(problem,x);
        λ,L1,λ1=getAdjoint(copy(K),copy(dcomp),copy(Ls[num]),copy(free))

        grad[:] .=getSensitivites(Ke,dcomp,ien,nel,len,copy(λ))[2]

        # g=(L1[:]'*d)[1] .- dmax - (λ1'*(K*d -F))[1]

        goal=L1[:]'*d -λ1'*K*d +λ1'*F .-dmax
        # display(goal )
        # display(sum(x) )
        
        return goal[1]
    end

    function GASIMP(x::Vector, grad::Vector,num::Int)

        K,F,d,stress,dcomp,g=FEM_frame(problem,x.^η);
        λ,L1,λ1=getAdjoint(copy(K),copy(dcomp),copy(Ls[num]),copy(free))

        grad[:] .=getSensitivitesSIMP(Ke,dcomp,ien,nel,len,copy(λ),η,x)[2]

        # g=(L1[:]'*d)[1] .- dmax - (λ1'*(K*d -F))[1]

        goal=L1[:]'*d -λ1'*K*d +λ1'*F .-dmax
        # display(goal )
        # display(sum(x) )
        
        return goal[1]
    end

    function FA(x::Vector, grad::Vector)
        grad[:] .=len[:]
        return (sum(x .* len )) 
    end

    function FASIMP(x::Vector, grad::Vector)
        grad[:] .=len[:]
        return (sum(x.^η .* len )) 
    end
    

    opt = Opt(:LD_MMA, nel)
    opt.lower_bounds = fill(1e-6,nel)
    opt.upper_bounds = fill(amax,nel)
    opt.xtol_rel = 1e-8
    opt.maxeval = maxeval

    if (SIMP)
        opt.min_objective = FASIMP
        println(size(Ls)[1])
        for i =1:size(Ls)[1]
            inequality_constraint!(opt, (x,gg) -> GASIMP(x,gg,i), 1e-8)
        end
    else
        opt.min_objective = FA
        println(size(Ls)[1])
        for i =1:size(Ls)[1]
            inequality_constraint!(opt, (x,gg) -> GA(x,gg,i), 1e-8)
        end
    end

    display(@time (minf,minx,ret) = optimize(opt, ones(nel).*amax*0.5))
    numevals = opt.numevals # the number of function evaluations
    display("got $minf after $numevals iterations (returned $ret)")
    
    return minx
end
# =======================================================================

function simulateAndExport(setup,X,dcomp,fileName,threshold)
    nodes=setup["nodes"]
    edges=setup["edges"]

    nodesNew=[]
    edgesNew=[]

    nel =  length(edges);  # number of elements 
    nnp =  length(nodes);  # number of nodal points

    for i in 1:nnp
        nodes[i]["in"]=false
        nodes[i]["displacement"]["x"]=dcomp[1,i];
        nodes[i]["displacement"]["y"]=dcomp[2,i];
        nodes[i]["displacement"]["z"]=dcomp[3,i]; 
    end

    idCount=0
    for i in 1:nel
        if X[i]>threshold
            if (!nodes[(edges[i]["source"]+1)]["in"])
                nodes[(edges[i]["source"]+1)]["in"]=true
                nodes[(edges[i]["source"]+1)]["id"]="n$(idCount)"
                nodes[(edges[i]["source"]+1)]["idd"]=idCount
                append!(nodesNew,[nodes[(edges[i]["source"]+1)]])
                id1=idCount
                idCount=idCount+1
            else
                id1=nodes[(edges[i]["source"]+1)]["idd"]
            end

            if (!nodes[(edges[i]["target"]+1)]["in"])
                nodes[(edges[i]["target"]+1)]["in"]=true
                nodes[(edges[i]["target"]+1)]["id"]="n$(idCount)"
                nodes[(edges[i]["target"]+1)]["idd"]=idCount
                append!(nodesNew,[nodes[(edges[i]["target"]+1)]])
                id2=idCount
                idCount=idCount+1
            else
                id2=nodes[(edges[i]["target"]+1)]["idd"]
            end
            edges[i]["source"]=id1
            edges[i]["target"]=id2

            append!(edgesNew,[edges[i]])

        end
    end


    setup["nodes"]=nodesNew;
    setup["edges"]=edgesNew;

    stringdata = JSON.json(setup)# pass data as a json string (how it shall be displayed in a file)

    # write the file with the stringdata variable information
    open(fileName, "w") do f
            write(f, stringdata)
    end


    # E,f,g,idb,ien,ndf,nel,nen,nnp,nsd,xn,len,Ke,Te,Ls=getDataFromSetup3D(setup,scale);

    # problem1=E,f,g,idb,ien,ndf,nel,nen,nnp,nsd,xn,len,Ke,Te;

    # nel=length(edgesNew)

    # display(a)

    # X=ones(nel)
    # K,F,d,stress,dcomp,g=FEM_frame(problem1,X);
    # display(dcomp)
    

    # display(plotTrussDeformed3D(problem1,copy(X),scale,0.1,1))


    
end