import matplotlib.pyplot as plt
import math
import numpy as np
import pandas as pd
from pylab import *
from scipy.optimize import curve_fit

name = str(input("What is the name of the file?"))
file = open(name, "r")
content = file.read()
data = content.split()

#to be used later
def func(x, a, c):
    return a*np.exp(c*x)

if(name[len(name)-2]=="u"): #slurm file .out
	i=0
	a=0
	while(i<len(data)): #pull beginning of data
		if(data[i]=="TotEng"):
			if(data[i-1]=="KinEng"):
				break
		i+=1
	i+=1
	while(a<len(data)):
		if(data[a]=="srun:" or data[a]=="slurmstepd:"): #pull end
			break
		a+=1
		
	#print(i)
	#print(a)
	#print(len(data))
	#print(data[i:a])
	x = np.array([0.0]*int(len(data[i:a])/5))
	t = np.array([0.0]*len(x))
	pe = np.array([0.0]*len(x))
	ke = np.array([0.0]*len(x))
	te = np.array([0.0]*len(x))
	
	j=0
	while(j<len(x)):
		x[j] = int(data[5*j+i])/100000 #convert from timestep = 0.01 fs to ps
		t[j] = float(data[5*j+i+1]) #temp
		pe[j] = float(data[5*j+i+2]) #PotEng
		ke[j] = float(data[5*j+i+3]) #KinEng
		te[j] = float(data[5*j+i+4]) #TotEng
		j+=1
	
	plt.scatter(x, t)
	plt.title("Temperature")
	plt.xlabel("Time (ps)")
	plt.ylabel("Temperature (K)")
	plt.show()
	plt.scatter(x, pe)
	plt.title("Potential Energy")
	plt.xlabel("Time (ps)")
	plt.ylabel("Energy (kcal/mol)")
	plt.show()
	plt.scatter(x, ke)
	plt.title("Kinetic Energy")
	plt.xlabel("Time (ps)")
	plt.ylabel("Energy (kcal/mol)")
	plt.show()
	plt.scatter(x, te)
	plt.title("Total Energy")
	plt.xlabel("Time (ps)")
	plt.ylabel("Energy (kcal/mol)")
	plt.show()
	
	
else: #it is a len file



	x = np.array([0.0]*int(len(data)/31))
	y = np.array([0.0]*int(len(data)/31))

	#calculate distance from x,y,z inputs
	i=0
	while(i<len(x)):
		x[i] = int(data[2+31*i])/100000 #convert from timestep = 0.01 fs to ps
		y[i] = math.sqrt((float(data[25+31*i])-float(data[28+31*i]))**2+(float(data[26+31*i])-float(data[29+31*i]))**2+(float(data[27+31*i])-float(data[30+31*i]))**2) #distance
		i+=1

	plt.scatter(x, y)
	plt.title("Length")
	plt.xlabel("Time (ps)")
	plt.ylabel("Distance (Å)")
	plt.show()
	
	autocorr = True#bool(input("Would you like to autocorrelate? (y/n) ")=="y")
	
	if(autocorr):
		start = float(input("What time would you like to treat as t0? (ps) ")) #to accurately choose a subsection of data to minimize noise
		end = float(input("At what time would you like to cut off? (ps) "))
		st = int((start/x[len(x)-1])*len(x)) 
		en = int((end/x[len(x)-1])*len(x)) 
		
		y2 = y[st:en]
		results = np.array([0.0]*(int(len(y2)/10)))
		x2=np.array([0.0]*len(results))
		
		#autocorrelate!
		series = pd.Series(y2,copy=False)
		i=0
		while(i<len(results)):
			results[i]=series.autocorr(i)
			x2[i]=x[i]-x[0]
			i+=1
			
		#show plot
		plt.scatter(x2, results)
		plt.title("Autocorrelation")
		plt.xlabel("Lag (ps)")
		plt.ylabel("Autocorrelation")
		plt.show()
		
		#often oscillatory modes can cause noise after a cutoff so it's useful to fit only starting data
		cutoff = float(input("At what lag would you like to cut off for exponential fit? "))
		ind = (cutoff/x2[len(x2)-1])*len(x2) 
		
		#essentially what proportion in the data the mult by index length of data
		results2 = np.array([0.0]*int(ind))
		x3=np.array([0.0]*len(results2))
		
		#print(ind)
		#print(len(results2))
		
		i=0
		while(i<len(results2)):
			results2[i] = results[int(i)]
			x3[i] = x2[int(i)]
			i+=1
		
		#fit exponential
		popt, pcov = curve_fit(func, x3, results2)	
		print(popt)
		print(pcov)
		
		#plot fitted exponential
		xM=np.linspace(-0.02,x2[len(x2)-1],10000)
		yM = func(xM, *popt)
		yP = func(x3, *popt)
		plt.plot(xM,yM, color='black')
		
		#calculate R^2
		yA = np.average(results2)
		SSr=0
		SSt=0
		
		i=0
		while(i<len(yP)):
			SSr += (results2[i]-yP[i])**2
			SSt += (results2[i]-yA)**2
			i+=1
			
		rsq = 1-SSr/SSt
		print(rsq)
		
		#plot
		plt.scatter(x2, results)
		plt.title("Autocorrelation")
		plt.xlabel("Lag (ps)")
		plt.ylabel("Autocorrelation")
		plt.show()