import numpy as np
import copy as cp


class CubeTranslation:
    NumPoints=None
    Direction=None

class CubeAtom:
    Index=None
    Nuclear_charge=None
    Coordinates=None

class Cube:
    Title1=None
    Title2=None
    Origin=None
    Translations=None
    Atoms=None
    Data=None

    def __init__(self,filename):
        with open(filename,'r') as f:
            splitline=f.readline().split()
            splitline=f.readline().split()
            splitline=f.readline().split()
            natoms=int(splitline[0])
            self.Origin=np.array(splitline[1:],"float")
            self.Translations=[CubeTranslation() for i in range(3)]
            for translation in self.Translations:
                splitline=f.readline().split()
                translation.NumPoints=int(splitline[0])
                translation.Direction=np.array(splitline[1:],"float")
            self.Atoms=[CubeAtom() for i in range(natoms)]
            for atom in self.Atoms:
                splitline=f.readline().split()
                atom.Index=int(splitline[0])
                atom.Nuclear_charge=float(splitline[1])
                atom.Coordinates=np.array(splitline[2:],"float")
            npoints=self.Translations[0].NumPoints*self.Translations[1].NumPoints*self.Translations[2].NumPoints
            data=np.zeros(npoints)
            pointer=0
            while True:
                line=f.readline()
                if not line:
                    break
                splitline=line.split()
                nelements=len(splitline)
                data[pointer:pointer+nelements]=np.array(splitline,"float")
                pointer+=nelements
            self.Data=np.reshape(data,[self.Translations[0].NumPoints,self.Translations[1].NumPoints,self.Translations[2].NumPoints])

    def Differentiate(self):
        grad_dot_a=np.zeros_like(self.Data)
        grad_dot_a[0,:,:]=self.Data[1,:,:]-self.Data[0,:,:]
        grad_dot_a[-1,:,:]=self.Data[-1,:,:]-self.Data[-2,:,:]
        grad_dot_a[1:-1,:,:]=0.5*(self.Data[2:,:,:]-self.Data[0:-2,:,:])
        grad_dot_b=np.zeros_like(self.Data)
        grad_dot_b[:,0,:]=self.Data[:,1,:]-self.Data[:,0,:]
        grad_dot_b[:,-1,:]=self.Data[:,-1,:]-self.Data[:,-2,:]
        grad_dot_b[:,1:-1,:]=0.5*(self.Data[:,2:,:]-self.Data[:,0:-2,:])
        grad_dot_c=np.zeros_like(self.Data)
        grad_dot_c[:,:,0]=self.Data[:,:,1]-self.Data[:,:,0]
        grad_dot_c[:,:,-1]=self.Data[:,:,-1]-self.Data[:,:,-2]
        grad_dot_c[:,:,1:-1]=0.5*(self.Data[:,:,2:]-self.Data[:,:,0:-2])
        grad_norm=cp.deepcopy(self)
        grad_x=cp.deepcopy(self)
        grad_y=cp.deepcopy(self)
        grad_z=cp.deepcopy(self)
        A=np.zeros([3,3])
        for i in range(3):
            A[i,:]=self.Translations[i].Direction
        for i in range(self.Translations[0].NumPoints):
            for j in range(self.Translations[1].NumPoints):
                for k in range(self.Translations[2].NumPoints):
                    B=[grad_dot_a[i,j,k],grad_dot_b[i,j,k],grad_dot_c[i,j,k]]
                    grad_x.Data[i,j,k],grad_y.Data[i,j,k],grad_z.Data[i,j,k]=np.linalg.solve(A,B)
        grad_norm.Data=np.sqrt(grad_x.Data**2+grad_y.Data**2+grad_z.Data**2)
        return grad_norm,grad_x,grad_y,grad_z

    def Function(self,function):
        self.Data=np.zeros([self.Translations[i].NumPoints for i in range(3)])
        for i in range(self.Translations[0].NumPoints):
            for j in range(self.Translations[1].NumPoints):
                for k in range(self.Translations[2].NumPoints):
                    x,y,z=self.Origin+i*self.Translations[0].Direction+j*self.Translations[1].Direction+k*self.Translations[2].Direction
                    self.Data[i,j,k]=function(x,y,z)

    def Export(self,filename):
        with open(filename,'w') as f:
            f.write("Generated by Orbaplaw\n" if self.Title1 is None else (self.Title1+'\n'))
            f.write(("Total %d grids\n" % int(self.Data.size)) if self.Title2 is None else (self.Title2+'\n'))
            f.write("%d %f %f %f\n" % (len(self.Atoms),*self.Origin))
            for translation in self.Translations:
                f.write("%d %f %f %f\n" % (translation.NumPoints,*translation.Direction))
            for atom in self.Atoms:
                f.write("%d %f %f %f %f\n" % (atom.Index,atom.Nuclear_charge,*atom.Coordinates))
            for i in range(self.Translations[0].NumPoints):
                for j in range(self.Translations[1].NumPoints):
                    for k in range(self.Translations[2].NumPoints):
                        f.write("%13.5E" % self.Data[i,j,k])
                        if k%6==5 or k+1==self.Translations[2].NumPoints:
                            f.write('\n')

