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 
import sys
from ase.io.vasp import read_vasp
import matplotlib.colors
import matplotlib.cm as cm


num = int(sys.argv[1]) # Take every point nth point (e.g. n=1 would be every point, n=2 would be every other point)
height = float(sys.argv[2]) # How far through to cut the plane in direct coordinates, e.g. 0.5 would be half way through the cell
contour_spacing = float(sys.argv[3]) # Spacing in ppm between the contours in the final plot. I tend to use 5.0 ppm

# read NICSCAR
with open("NICS", 'r') as f:
    #PROCESS STRUCTURAL INFORMATION
    system = f.readline()
    scale = float(f.readline())

    lattice = []
    for i in range(0,3):
        lattice.append( [float(x) for x in f.readline().split()] )
    lattice_matrix = np.array(lattice)

    atoms = [ x for x in f.readline().split()]
    atoms_count = [ int(x) for x in f.readline().split()]

    fractional = True if "Direct" in f.readline() else 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()

    #READ NICS INFO
    gridc = [int(x) for x in f.readline().split()]
    tmp = []
    n_lines = int(np.ceil(gridc[0]*gridc[1]*gridc[2]/5))
    #print(n_lines)
    for _ in range(0,9):
        for __ in range(0,n_lines):
            line = f.readline()
            for x in line.split():
                tmp.append(float(x))
        f.readline()
        #f.readline()
    nics_den = np.array(tmp)


# components of nics
print(np.size(nics_den), gridc[2],gridc[1],gridc[0])
nics_tensors = np.reshape(nics_den,(9,gridc[2],gridc[1],gridc[0]))
print(nics_tensors.shape)
nics_xx, nics_yy, nics_zz = np.array(nics_tensors[0]), np.array(nics_tensors[4]), np.array(nics_tensors[8])

clight = 137.037
clightv = 1973.3036311408
momtomom = 0.529177249/137.037**2/2
magmomtoenergy = 1/137.037**2*0.529177249**3*13.605826
nics_iso = (nics_xx + nics_yy +nics_zz)/3

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


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

# 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])*c+j/(gridc[1])*b+k/(gridc[2])*a)
            z.append(pos[0])
            y.append(pos[1])
            x.append(pos[2])
            zz[k][j].append(pos[0])
            yy[k][j].append(pos[1])
            xx[k][j].append(pos[2])
            #print(pos[0], pos[1], pos[2])
            n = n + 1 
xx, yy, zz = np.array(xx), np.array(yy), np.array(zz)

x_slice = int(height*gridc[2])
#z_slice = zz[x_slice, ::num, ::num]
z_slice = zz[::num, ::num, x_slice]
#y_slice = yy[x_slice, ::num, ::num]
y_slice = yy[::num, ::num, x_slice]
#nics_iso_slice = nics_iso[x_slice, ::num, ::num]
nics_iso_slice = nics_iso[::num, ::num, x_slice]

# Prepare plot
plt.rcParams['figure.figsize'] = [12, 10]
fig, ax = plt.subplots()

# Spacing for dashed contours - set to every 5.0 ppm
spacings = np.arange(nics_iso_slice.min(), nics_iso_slice.max(), contour_spacing)

cvals  = [nics_iso_slice.min(), 0.0, nics_iso_slice.max()] # NICS value limits for contour plot (z-axis)
colors = ["#a82c35", "w", "#3E70EA"] # Define color settings for contour plot gradient (z-axis)

norm=plt.Normalize(min(cvals),max(cvals)) # Normalise z-axis limits
tuples = list(zip(map(norm,cvals), colors)) # Associates colors and z-axis limits
cmap = matplotlib.colors.LinearSegmentedColormap.from_list("", tuples) # Prepare color map

CS = ax.contourf(y_slice, z_slice, nics_iso_slice, 1000, cmap=cmap, norm=norm, corner_mask=True) # Plot contor map - colors
CS2 = ax.contour(y_slice, z_slice, nics_iso_slice, levels=spacings, colors='black') # Plot contor lines
ax.clabel(CS2, inline=True, fontsize=10)
cbar = fig.colorbar(CS)
ax.set_title('Shielding contour plot in ppm')
ax.set_xlabel('y in Å')
ax.set_ylabel('z in Å')

atoms = read_vasp("./POSCAR")

# Draw the structure of C4H4
x, y, z = [], [], []
for a in atoms.positions:
    x.append(a[0])
    y.append(a[1])
    z.append(a[2])
#print(x, y)

# Draw the C-C bonds of C4H4
a, b = [], []
for i in range(0, 4):
    a.append(y[i])
    b.append(z[i])
a.append(y[0])
b.append(z[0])
plt.plot(a, b, 'w-') # 'ko-' for dots as atoms

# Draw the C-H bonds of C4H4
c, d = [], []
for i in range(0, 4):
    plt.plot([y[i],y[i-4]], [z[i],z[i-4]], 'w-') # 'ko-' for dots as atoms'''

plt.scatter(y, z, color='black', zorder=2)

plt.savefig('./nics_c4h4.png')

