import os
import os.path
import matplotlib.pyplot as plt
import time as ti
import numpy as np

def main(): 
    data_folder = os.path.join("OHNWCHEM")
    file_input = "T3Fmolecule copy.xyz"
    data_file = os.path.join(data_folder, file_input)
    
    ts = []
    Hs = []
    Os = []
    
    #Opens a file, goes through every line, and pulls the Ns and Os
    with open(data_file, "r") as f:
        f_lines = f.readlines()
        count = 0 
        Hcount = 0
        for line in f_lines:
            l=line.split()
            if l[0] == "HX":
                Hs.append(np.array((float(l[1]),float(l[2]),float(l[3]))))
            elif l[0] == "OX":
                Os.append(np.array((float(l[1]),float(l[2]),float(l[3]))))
            elif l[0].isnumeric() == True and len(l)>1:
                ts.append(float(l[0]))
    
                
    #Turns a vector position into a vector matrix (one big list)
    f.close()
    Hs = np.vstack(Hs)
    Os = np.vstack(Os)
   # print(ts)    #include when you want to see the number of timesteps
    
    #Going down the list for x,y,z for N and O and then plotting it
    r = distCalc(Hs[:,0],Hs[:,1],Hs[:,2],Os[:,0],Os[:,1],Os[:,2])
    float_ts = list(np.float_(ts)*0.242)
    plt.plot(float_ts,r)
    plt.xlabel("Time (fs)")
    plt.ylabel("O-H Distance (nm)")
    plt.grid()    
    plt.show()
    
    
#Performs the Distance Calculation for O-N
def distCalc (a1,b1,c1,a2,b2,c2):
    return pow(sum((pow((a2-a1)*0.1,2),pow((b2-b1)*0.1,2),pow((c2-c1)*0.1,2))),0.5)

main()