import plotly.graph_objects as go
import plotly.io as pio 
import numpy as np
from scipy.interpolate import RegularGridInterpolator
import matplotlib.pyplot as plt 

# read CHGCAR
with open('CHGCAR.copy', 'r') as f:
    system = f.readline()
    scale = float(f.readline())
    tmp = []
    for i in range(0,3):
        tmp.append( [float(x) for x in f.readline().split()] )
    lattice_matrix = np.array( tmp )
    atoms = [ x for x in f.readline().split()]
    atoms_count = [ int(x) for x in f.readline().split()]
    if "Direct" in f.readline():
        fractional=True
    else:
        fractional=False
    position = []
    for i in range(0,len(atoms)):
        for j in range(0,atoms_count[i]):
            position.append( [float(x) for x in f.readline().split()] )
    f.readline()
    gridc = [int(x) for x in f.readline().split()]
    tmp = []
    n_lines = int(np.ceil(gridc[0]*gridc[1]*gridc[2]/5))
    for _ in range(0,1):
        for __ in range(0,n_lines):
            line = f.readline()
            for x in line.split():
                tmp.append(float(x))
        f.readline()
        #f.readline()
    locpot = np.array(tmp)

# charge density
print(np.size(locpot), gridc[2],gridc[1],gridc[0])
b = np.reshape(locpot,(1,gridc[2],gridc[1],gridc[0]))
v0 = np.array(b[0])
v0 = np.transpose(v0, (2, 1, 0))

x_vals = np.linspace(0,1,gridc[0])
y_vals = np.linspace(0,1,gridc[1])
z_vals = np.linspace(0,1,gridc[2])

print("x=", np.shape(x_vals))
print("y=", np.shape(y_vals))
print("z=", np.shape(z_vals))

# uncomment to plot unit cell in fractional coordinates
#lattice_matrix = np.array([[1,0,0],[0,1,0],[0,0,1]])
a, b, c = scale*lattice_matrix
x, y, z = [], [], []
n=0
xx, yy, zz=[], [], []
for k in range(0,gridc[0]):
    xx.append([])
    yy.append([])
    zz.append([])
    for j in range(0,gridc[1]):
        xx[k].append([])
        yy[k].append([])
        zz[k].append([])
        for i in range(0,gridc[2]):
            pos = (i/(gridc[0]-1)*c+j/(gridc[1]-1)*b+k/(gridc[2]-1)*a)
            x.append(pos[0])
            y.append(pos[1])
            z.append(pos[2])
            xx[k][j].append(pos[0])
            yy[k][j].append(pos[1])
            zz[k][j].append(pos[2])
            n = n + 1 
xx, yy, zz = np.array(xx), np.array(yy), np.array(zz)

"""
xx, yy, zz = np.meshgrid(x_vals, y_vals, z_vals, indexing='ij')
print("Fractional")
print(np.shape(xx))
#print(zz)
"""

x_slice = int(0.0*gridc[2])

z_slice = zz[x_slice, :, :]
y_slice = yy[x_slice, :, :]
v0_slice =  v0[x_slice, :, :]
print(np.size(y_slice),np.size(z_slice), np.size(v0_slice))
plt.figure(figsize=(24,20))
plt.contourf(y_slice, z_slice, v0_slice, alpha=0.5)
plt.colorbar()
plt.xlabel("y")

plt.ylabel("z")
plt.title("Charge density (rho) in the yz-plane at x = 0.0")

#plt.xlim([-0.5, 4.5])
#plt.ylim([-0.5, 6.6])

plt.savefig("fig_chgcar-slice.png")

