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

# BASED ON https://github.com/jonhiller/Voxelyze


function updateNodes!(dt,currentTimeStep,N_position, N_restrained,N_displacement,N_angle,N_currPosition,N_linMom,N_angMom,
        N_intForce,N_intMoment,N_force,N_fixedDisplacement,N_moment,N_orient,N_edgeID,N_edgeFirst,N_material,N_poissonStrain,
        E_intForce1,E_intMoment1,E_intForce2,E_intMoment2,E_axis, E_strain,E_material,E_maxStrain)

    index = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    stride = blockDim().x * gridDim().x
    ## @cuprintln("thread $index, block $stride")
    N,M=size(N_edgeID)
    for i = index:stride:N
        
        isFloorEnabled=floorEnabled();#todo make these arguments ??
        
        # @inbounds if lengthVector3(Vector3(convert(Float64,N_displacement[i].x), 
        #             convert(Float64,N_displacement[i].y),convert(Float64,N_displacement[i].z))) >1000.0 #todo ?? change this according to ratio
        #     # @cuprintln("DIVERGEDNODE!!!!!!!!!!")
        #     return 
        # end
        
        
        
        @inbounds if (N_restrained[i].x && N_restrained[i].y && N_restrained[i].z)
            @inbounds N_linMom[i]=Vector3(0,0,0)
            @inbounds N_angMom[i]=Vector3(0,0,0)
            @inbounds if(N_material[i].poisson)
                @inbounds N_poissonStrain[i]=strain(i,N_material[i].poisson,i, N_edgeID, N_edgeFirst, E_axis, E_strain, E_material,E_maxStrain, N_restrained[i], N_material[i])
            end
            return
        elseif N_fixedDisplacement[i].x != 0.0 || N_fixedDisplacement[i].y != 0.0  ||  N_fixedDisplacement[i].z != 0.0
            @inbounds N_linMom[i]=Vector3(0,0,0)
            @inbounds N_angMom[i]=Vector3(0,0,0)
            @inbounds if(N_material[i].poisson)
                @inbounds N_poissonStrain[i]=strain(i,N_material[i].poisson,i, N_edgeID, N_edgeFirst, E_axis, E_strain, E_material,E_maxStrain, N_restrained[i], N_material[i])
            end
            @inbounds translate=externalDisplacement(currentTimeStep,N_position,N_fixedDisplacement[i])
            @inbounds N_currPosition[i]=N_currPosition[i]+translate
            @inbounds N_displacement[i]=N_displacement[i]+translate

        else
            
            for j in 1:M
                @inbounds if (N_edgeID[i,j]!=-1 && !isFailed(E_maxStrain[N_edgeID[i,j]],E_material[N_edgeID[i,j]]))
                    #@cuprintln("i $i, j $j, N_edgeID[i,j] $temp")
                    @inbounds N_intForce[i]=ifelse(N_edgeFirst[i,j], N_intForce[i]+E_intForce1[N_edgeID[i,j]], N_intForce[i]+E_intForce2[N_edgeID[i,j]] )
                    @inbounds N_intMoment[i]=ifelse(N_edgeFirst[i,j], N_intMoment[i]+E_intMoment1[N_edgeID[i,j]], N_intMoment[i]+E_intMoment2[N_edgeID[i,j]] )
                end
            
            end
            
            FloorStaticFriction=false; 
            
            prec=1e8
            #get properties
            @inbounds E=roundd(N_material[i].E,prec)
            @inbounds nomSize=roundd(N_material[i].nomSize,prec)
            @inbounds mass=roundd(N_material[i].mass,prec)
            @inbounds massInverse=roundd(N_material[i].massInverse,prec)
            
            @inbounds curForce = force(i,N_intForce[i],N_orient[i],N_force[i],N_position[i],currentTimeStep,N_material[i],N_linMom[i])

            fricForce = Vector3(curForce.x,curForce.y,curForce.z);
            if (isFloorEnabled)
                @inbounds curForce,FloorStaticFriction=floorForce!(dt,curForce,N_currPosition[i],
                    Vector3(N_linMom[i].x,N_linMom[i].y ,N_linMom[i].z),FloorStaticFriction,N_material[i])
                
            end
            
            fricForce = curForce-fricForce;

            # @cuprintln("dt $dt")
            # @cuprintln("massInverse $massInverse")
            
            
            #########################################
            
            # @inbounds N_intForce[i]=Vector3(0,0,0) #??
        
            
            @inbounds N_linMom[i]=N_linMom[i]+curForce*Vector3(dt,dt,dt) #todo make sure right
            @inbounds translate=N_linMom[i]*Vector3((dt*massInverse),(dt*massInverse),(dt*massInverse)) # ??massInverse

            
            
            ############################
            
            
            #we need to check for friction conditions here (after calculating the translation) and stop things accordingly
            @inbounds if (isFloorEnabled && floorPenetration( convert(Float64,N_currPosition[i].x),convert(Float64,N_currPosition[i].y),N_material[i].nomSize)>= 0.0 )#we must catch a slowing voxel here since it all boils down to needing access to the dt of this timestep.
                work =convert(Float64,fricForce.x*translate.x + fricForce.z*translate.z); #F dot disp
                @inbounds hKe = 0.5*massInverse*convert(Float64,(N_linMom[i].x*N_linMom[i].x + N_linMom[i].z*N_linMom[i].z)); #horizontal kinetic energy

                if((hKe + work) <= 0.0) 
                    FloorStaticFriction=true; #this checks for a change of direction according to the work-energy principle
                end

                if(FloorStaticFriction)
                    #if we're in a state of static friction, zero out all horizontal motion
                    N_linMom[i]=Vector3(0.0 ,convert(Float64,N_linMom[i].y) ,0.0)
                    translate=Vector3(0.0 ,convert(Float64,translate.y) ,0.0)
                end
            else
                FloorStaticFriction=false
            end

            if N_restrained[i].x
                translate=Vector3(0.0 ,convert(Float64,translate.y) ,convert(Float64,translate.z))
            end
            if N_restrained[i].y
                translate=Vector3(convert(Float64,translate.x) ,0.0 ,convert(Float64,translate.z))
            end
            if N_restrained[i].z
                translate=Vector3(convert(Float64,translate.x) ,convert(Float64,translate.y) ,0.0)
            end
            
            @inbounds N_currPosition[i]=N_currPosition[i]+translate
            @inbounds N_displacement[i]=N_displacement[i]+translate

            
            
            
            # Rotation
            @inbounds curMoment = moment(N_intMoment[i],N_orient[i],N_moment[i],N_material[i],N_angMom[i]) 
            
            @inbounds N_intMoment[i]=Vector3(0,0,0) # do i really need it?
            @inbounds N_intForce[i]=Vector3(0,0,0) # do i really need it?
            
            @inbounds N_angMom[i]=N_angMom[i]+curMoment*Vector3(dt,dt,dt)
            
            @inbounds momentInertiaInverse=convert(Float64,N_material[i].momentInertiaInverse)
            # if i==10
            #     @cuprintln("momentInertiaInverse $momentInertiaInverse")
            # end
            # momentInertiaInverse=1.92e-6

            @inbounds temp=FromRotationVector(N_angMom[i]*Vector3((dt*momentInertiaInverse),(dt*momentInertiaInverse),(dt*momentInertiaInverse)))

            
            @inbounds N_orient[i]=multiplyQuaternions(temp,N_orient[i])

            @inbounds N_angle[i]=ToRotationVector(N_orient[i])
        
            ####poisson
            @inbounds if(N_material[i].poisson)
                @inbounds N_poissonStrain[i]=strain(i,N_material[i].poisson,i, N_edgeID, N_edgeFirst, E_axis, E_strain, E_material,E_maxStrain, N_restrained[i], N_material[i])
            end
        end
    end
    return
end


function run_updateNodes!(dt,currentTimeStep,N_position, N_restrained,N_displacement, N_angle,N_currPosition,
        N_linMom,N_angMom,N_intForce,N_intMoment,N_force,N_fixedDisplacement,N_moment,N_orient,N_edgeID,N_edgeFirst,N_material,N_poissonStrain,
        E_intForce1,E_intMoment1,E_intForce2,E_intMoment2,E_axis, E_strain,E_material,E_maxStrain)
    N=length(N_intForce)
    numblocks = ceil(Int, N/256)
    CUDA.@sync begin
        @cuda threads=256 blocks=numblocks updateNodes!(dt,currentTimeStep,N_position, N_restrained,N_displacement, 
            N_angle,N_currPosition,N_linMom,N_angMom,N_intForce,N_intMoment,N_force,N_fixedDisplacement,N_moment,N_orient,N_edgeID,N_edgeFirst,N_material,N_poissonStrain,
            E_intForce1,E_intMoment1,E_intForce2,E_intMoment2,E_axis, E_strain,E_material,E_maxStrain)
    end
end


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

function strain(i,poissonStrain,id,N_edgeID,N_edgeFirst, E_axis, E_strain,E_material,E_maxStrain,N_restrained,mat)
	
	###if no connections in the positive and negative directions of a particular axis, strain is zero
	###if one connection in positive or negative direction of a particular axis, strain is that strain - ?? and force or constraint?
	###if connections in both the positive and negative directions of a particular axis, strain is the average. 
	
    # intStrRet= [0.0, 0.0, 0.0] ###intermediate strain return value. axes according to linkAxis enum
	# numBondAxis = [0,0,0]; ###number of bonds in this axis (0,1,2). axes according to linkAxis enum
    # tension = [false,false,false];

    intStrRet1=0.0
    intStrRet2=0.0
    intStrRet3=0.0

    numBondAxis1=0.0
    numBondAxis2=0.0
    numBondAxis3=0.0

    tension1=false
    tension2=false
    tension3=false

    fixed=N_restrained.x

    N,M=size(N_edgeID)
    for j in 1:M
        @inbounds linkID=N_edgeID[id,j]
        # @cuprintln("id $id, j $j, linkID $linkID")
        if (linkID!=-1 && !isFailed(E_maxStrain[linkID],E_material[linkID]))
            # if (E_strain[j]!=0.0)
                @inbounds axis= normalizeVector3(E_axis[linkID])
                @inbounds strain1=axialStrain( N_edgeFirst[id,j],E_strain[linkID],E_material[linkID])
                # if (E_strain[j]!=0.0)
                #     @cuprint("strain1 $strain1")
                # end
                x= abs(convert(Float64,axis.x))*strain1 #todo change for general axis
                y= abs(convert(Float64,axis.y))*strain1 #todo change for general axis
                z= abs(convert(Float64,axis.z))*strain1 #todo change for general axis
                if (axis.x!=0.0)
                    intStrRet1+=x
                    numBondAxis1+=1
                    # if(i==32)
                    #     @cuprintln("strain1 x $(strain1)")
                    # end
                end
                if (axis.y!=0.0)
                    intStrRet2+=y
                    numBondAxis2+=1
                end
                if (axis.z!=0.0)
                    intStrRet3+=z
                    numBondAxis3+=1
                end
            # end
        end
    end
    # @cuprintln("intStrRet 11 x $(intStrRet1*1e6) *1e-6, y $(intStrRet2*1e6) *1e-6, z $(intStrRet3*1e6) *1e-6")
    # for i =1:3
        if (numBondAxis1>1.0) 
            intStrRet1=intStrRet1 /numBondAxis1 ###average
        end
        if (poissonStrain)
            tension1 = ((numBondAxis1==2.0) || (fixed&&numBondAxis1==1.0))
                #(ext && (numBondAxis[i]==1 && (ext.isFixed((dofComponent)(1<<i)) || ext.force()[i] != 0)))); #if both sides pulling, or just one side and a fixed or forced voxel...
        end

        if (numBondAxis2>1.0) 
            intStrRet2=intStrRet2 /numBondAxis2 ###average
        end
        if (poissonStrain)
            tension2 = ((numBondAxis2==2.0) || (fixed&&numBondAxis2==1.0))
                #(ext && (numBondAxis[i]==1 && (ext.isFixed((dofComponent)(1<<i)) || ext.force()[i] != 0)))); #if both sides pulling, or just one side and a fixed or forced voxel...
        end

        if (numBondAxis3>1.0) 
            intStrRet3=intStrRet3 /numBondAxis3 ###average
        end
        if (poissonStrain)
            tension3 = ((numBondAxis3==2.0) || (fixed&&numBondAxis3==1.0))
                #(ext && (numBondAxis[i]==1 && (ext.isFixed((dofComponent)(1<<i)) || ext.force()[i] != 0)))); #if both sides pulling, or just one side and a fixed or forced voxel...
        end
    # end

    # if(i==32)
    #     @cuprintln("intStrRet 1 x $(intStrRet1) , y $(intStrRet2) , z $(intStrRet3)")
    #     @cuprintln("numBondAxis x $(numBondAxis1) , y $(numBondAxis2) , z $(numBondAxis3) ")
    # end

	if (poissonStrain)
		if (!(tension1 && tension2 && tension3 )) ###if at least one isn't in tension
			add = 0.0;
            # for i =1:3
                if (tension1) 
                    add=add+intStrRet1
                end
                if (tension2) 
                    add=add+intStrRet2
                end
                if (tension3) 
                    add=add+intStrRet3
                end
            # end
            # @cuprintln("add x $(add*1e6)")
            # @cuprintln("mat.nu x $(mat.nu)")
            # value =  pow((1.0 + add),(-mat.nu))-1.0 #((1.0 + add)^(-mat.nu)) - 1.0;
            value = CUDA.pow( 1.0 + add, -mat.nu) - 1.0;

            # if(i==32)
            #     @cuprintln("value x $(value)")
            #     @cuprintln("mat.nu x $(mat.nu)")
            #     @cuprintln("add x $(add)")
            # end


            # value =((1.0 + add)^(-roundd(mat.nu,10e5))) - 1.0;
            # @cuprintln("value x $(value*1e6)")
            # value=0.0
            # for i =1:3
                if (!tension1) 
                    intStrRet1=value;
                end
                if (!tension2) 
                    intStrRet2=value;
                end
                if (!tension3) 
                    intStrRet3=value;
                end
            # end   
        end
    end

    # if(i==32)
    #     @cuprintln("intStrRet 2 x $(intStrRet1*1e6) *1e-6, y $(intStrRet2*1e6) *1e-6, z $(intStrRet3*1e6) *1e-6")
    # end

    intStrRet= Vector3(intStrRet1,intStrRet2,intStrRet3); ###intermediate strain return value. axes according to linkAxis enum
    # @cuprintln("intStrRet x $(intStrRet1*1e6) *1e-6, y $(intStrRet2*1e6) *1e-6, z $(intStrRet3*1e6) *1e-6")
    # @cuprintln("numBondAxis x $(numBondAxis1) , y $(numBondAxis2) , z $(numBondAxis3) ")
    # if tension1
    #     tension1=1.0
    # else
    #     tension1=0.0
    # end
    # if tension2
    #     tension2=1.0
    # else
    #     tension2=0.0
    # end
    # if tension3
    #     tension3=1.0
    # else
    #     tension3=0.0
    # end
    # @cuprintln("tension x $(tension1) , y $(tension2) , z $(tension3) ")
	return intStrRet
end

function axialStrain( positiveEnd,strain,mat)
    #strainRatio = pVPos->material()->E/pVNeg->material()->E;
    strainRatio=convert(Float64,mat.strainRatio)
    # positiveEnd=true
    return positiveEnd ? 2.0 *strain*strainRatio/(1.0+strainRatio) : 2.0*strain/(1.0+strainRatio)
end