'''Metodi per studiare Argon con LJ'''


import numpy as np
import conversions
import copy

class argon:
    def __init__(self,n,density):
        self.n=n#numero di atomi
        self.density=density#in kg/m^3
        massa=np.float64(self.n*39.948*1.6605400e-27)
        volume=massa/density
        self.massa_ar=np.float64(39.948*1.6605400e-27)
        self.lato=volume**(1.0/3.)#lato in m
        self.sigma=np.float64(3.4e-10)
        self.eps=np.float64(120*conversions.KB_SI)
        

    def initialize(self,temp):
        self.pos0=np.zeros((3,self.n),dtype=np.float64)
        self.count_pbc=np.zeros((3,self.n),dtype=np.int16)
        self.temp=temp
        nn=int(self.n**(1.0/3.0))
        if(nn**3< self.n):
            nn+=1
        h=self.lato/nn
        print('LATO (Ang):\t'+str(self.lato))
        ipos=0
        for i in range(nn):
            x=i*h
            for j in range(nn):
                y=j*h
                for k in range(nn):
                    z=k*h
                    if(ipos<self.n):
                        self.pos0[0,ipos]=x
                        self.pos0[1,ipos]=y
                        self.pos0[2,ipos]=z
                    ipos+=1

        self.vel0=np.zeros((3,self.n),dtype=np.float64)

        pmax, vpmax=conversions.maxwell_boltzmann_max(self.massa_ar,temp)


        dir=np.zeros((3),dtype=np.float64)
        for i in range(self.n):
            for j in range(3):
                dir[j]=np.random.random()-0.5
            dir=dir/np.sqrt(np.sum(dir**2))
            accepted=False
            while(not accepted):
                v=3*vpmax*np.random.random()
                p=np.random.random()
                if(p<conversions.maxwell_boltzmann(self.massa_ar,temp,v)/pmax):
                    accepted=True
                    self.vel0[0:3,i]=dir[0:3]*v
        
        vtot=np.zeros((3),dtype=np.float64)
        for k in range(3):
            vtot[k]=np.sum(self.vel0[k,:])
            self.vel0[k]-=vtot[k]/self.n

    def energy(self):
        kin=0
        for i in range(self.n):
            kin+=0.5*self.massa_ar*np.sum(self.vel0[0:3,i]**2)
        jpos=np.zeros((3),dtype=np.float64)
        pot=0.0
       
        for i in range(2,self.n):
            for j in range(i-1):
                distmin=1e10   
                for jx in range (-1,2):
                    for jy in range(-1,2):
                        for jz in range(-1,2):
                            jpos[0:3]=self.pos0[0:3,j]+np.array([jx*self.lato,jy*self.lato,jz*self.lato],dtype=np.float64)
                            dist=np.sum((jpos[0:3]-self.pos0[0:3,i])**2)
                            if(dist< distmin):
                                distmin=dist
                distmin=np.sqrt(distmin)
            
                pot+=4*self.eps*((self.sigma/distmin)**12-(self.sigma/distmin)**6)

        temp=2*kin/(3*self.n-3)/conversions.KB_SI

        energy=kin+pot
        return energy, kin, pot, temp
    


    def energy2(self):
        kin=0
        for i in range(self.n):
            kin+=0.5*self.massa_ar*np.sum(self.vel0[0:3,i]**2)
        jpos=np.zeros((3),dtype=np.float64)
        pot=0.0
        sigmaq=self.sigma**2

        for i in range(1,self.n):
            for j in range(i):
                jpos[0:3]=self.pos0[0:3,i]-self.pos0[0:3,j]
                for k in range(3):
                    if(jpos[k]>self.lato/2):
                        jpos[k]-=self.lato
                    elif(jpos[k]<-self.lato/2):
                        jpos[k]+=self.lato
                distmin=np.sum(jpos[0:3]**2)
                #ene=4*self.eps*((sigmaq/distmin)**6-(sigmaq/distmin)**3)
                #if(ene>0):
                #    print(i,j,ene)
            
                pot+=4*self.eps*((sigmaq/distmin)**6-(sigmaq/distmin)**3)

        temp=2*kin/(3*self.n-3)/conversions.KB_SI

        energy=kin+pot
        return energy, kin, pot, temp
    


    def force(self):
        self.for0=np.zeros((3,self.n),dtype=np.float64)
        #print(self.for0)
        jpos=np.zeros((3),dtype=np.float64)
        jposmin=np.zeros((3),dtype=np.float64)
        for i in range(self.n):           
            for j in range(self.n):
                distmin=1e10   
                for jx in range (-1,2):
                    for jy in range(-1,2):
                        for jz in range(-1,2):
                            jpos[0:3]=self.pos0[0:3,j]+np.array([jx*self.lato,jy*self.lato,jz*self.lato],dtype=np.float64)
                            dist=np.sum(np.power((jpos[0:3]-self.pos0[0:3,i]),2,dtype=np.float64),dtype=np.float64)
                            if(dist< distmin):
                                distmin=dist
                                jposmin=copy.deepcopy(jpos)
                distmin=np.sqrt(distmin,dtype=np.float64)
                #print((i,j,distmin))
                if(i != j):
                    self.for0[0:3,i]+=(24*self.eps/distmin**2)*(2*(self.sigma/distmin)**12-(self.sigma/distmin)**6)*(self.pos0[0:3,i]-jposmin[0:3])
        #print(self.for0)

    def force2(self):
        self.for0=np.zeros((3,self.n),dtype=np.float64)
        #print(self.for0)
        jpos=np.zeros((3),dtype=np.float64)
        sigmaq=self.sigma**2
        for i in range(self.n):           
            for j in range(self.n):
                jpos[0:3]=self.pos0[0:3,i]-self.pos0[0:3,j]
                for k in range(3):
                    if(jpos[k]>self.lato/2):
                        jpos[k]-=self.lato
                    elif(jpos[k]<-self.lato/2):
                        jpos[k]+=self.lato
                distmin=np.sum(np.power(jpos[0:3],2.0,dtype=np.float64),dtype=np.float64)
                if(i != j):
                    self.for0[0:3,i]+=(24*self.eps/distmin)*(2*(sigmaq/distmin)**6-(sigmaq/distmin)**3)*jpos[0:3]

        for k in range(3):
            ftot=np.sum(self.for0[k,:],dtype=np.float64)
            ftot/=self.n
            self.for0[k,:]-=ftot
        #print(np.max(abs(self.pos0)))



        
    def step_velocity_verlet(self,dt):
        ar_next=copy.deepcopy(self)
        for i in range(self.n):
            ar_next.pos0[0:3,i]=self.pos0[0:3,i]+self.vel0[0:3,i]*dt+(0.5/self.massa_ar)*self.for0[0:3,i]*dt*dt
        ar_next.pbc()
        ar_next.force2()
        for i in range(self.n):
            ar_next.vel0[0:3,i]=self.vel0[0:3,i]+(0.5/self.massa_ar)*(ar_next.for0[0:3,i]+self.for0[0:3,i])*dt
        for k in range(3):
            ftot=np.sum(ar_next.vel0[k,:],dtype=np.float64)
            ftot/=self.n
            ar_next.vel0[k,:]-=ftot
        """print(np.max(abs(self.vel0)))
        print(np.max(abs(ar_next.vel0)))
        velmod=np.zeros((self.n),dtype=np.float64)
        for k in range(self.n):
            velmod[k]=np.sum(np.power(ar_next.vel0[:,k],2))
        print(np.argmax(velmod))
        print(self.count_pbc[:,np.argmax(velmod)])
        print(self.vel0[:,np.argmax(velmod)])
        print(ar_next.vel0[:,np.argmax(velmod)])
        print(self.for0[:,np.argmax(velmod)])
        print(ar_next.for0[:,np.argmax(velmod)])"""


        return ar_next

        

    def pbc(self):
        for i in range(self.n):
            for j in range(3):
                if(self.pos0[j,i]>self.lato):
                    self.pos0[j,i]-=self.lato
                    self.count_pbc[j,i]+=1
                elif(self.pos0[j,i]<0):
                    self.pos0[j,i]+=self.lato     
                    self.count_pbc[j,i]-=1