import numpy as np
from functools import reduce
import matplotlib.pyplot as plt
import time


def readcenxyz(fname):
    file1 = open(fname)
    frames = []
    while 1:
        newline = file1.readline()
        if newline:
            pframe = []
            atomnum = int(newline.split()[1])
            for ix in range(atomnum):
                pline = file1.readline().split()[1:]
                pframe.append([float(jx) for jx in pline])
            frames.append(pframe)
        else:
            return np.array(frames)


def xshift(watt1, watt2):
    return (watt1 - watt2)*(watt1 - watt2)


def msdxy(data1, max_time=500):
    fin_dip_up = []
    fin_dip_down = []
    for time in range(0, min([int(data1.shape[0] / 2) + 1, max_time])):
        fin_dip_1 = []
        fin_dip_2 = []
        for numb in range(data1.shape[0] - time):
            diplist_1 = []
            diplist_2 = []
            for wat_ind in range(data1.shape[1]):
                diplist_1.append(xshift(watt1=data1[numb][wat_ind], watt2=data1[numb + time][wat_ind]))
                diplist_2.append(1)
            fin_dip_1.append(reduce(lambda x, y: x + y, diplist_1))
            fin_dip_2.append(reduce(lambda x, y: x + y, diplist_2))
        fin_dip_up.append(reduce(lambda x, y: x + y, fin_dip_1))
        fin_dip_down.append(reduce(lambda x, y: x + y, fin_dip_2))
    return fin_dip_up, fin_dip_down


def getmsd(data1, tstep):
    xup, xdn = msdxy(data1, 2500) #迭代最大值
    finy = np.array([xup[ix]/xdn[ix] for ix in range(len(xup))])
    finx = np.array([ix*tstep for ix in range(len(xup))])
    resmsd = np.polyfit(finx, finy, 1)
    return resmsd[0], finy


t1 = time.time()
datas = readcenxyz("newcen")
timestep = 1 #1=1ps
molist = [1, 20000, 225, 225, 32, 100, 20, 56, 46, 36, 140, 80, 64, 20, 20, 10, 14, 14, 4, 24, 18]
molsum = [sum(molist[:ix+1]) for ix in range(len(molist))]
molsum.append(0)
allmsd = []
sqr_dis = []
mass = [0, 18, 23, 35.5, 170, 226, 282, 226, 282, 196, 196, 176, 212, 202, 202, 167, 223, 217, 186, 170, 184]
for ix in range(len(molsum)-1):  #for ix in range(5, len(molsum)-1):  从第五个分子开始算
    msddata = getmsd(datas[:, molsum[ix-1]:molsum[ix], 0], timestep) #0=x,1=y,2=z
    allmsd.append(msddata[0])
    sqr_dis.append(msddata[1])
t2 = time.time()
print("time:", t2-t1)
fileout1 = open("allmsd", "w")
fileout2 = open("sqdis", "w")
for ix in range(len(mass)):                         #for ix in range(len(allmsd)):
    print(mass[ix], allmsd[ix], file=fileout1)      #    print(mass[5+ix], allmsd[ix], file=fileout1) 
pltx = [ix*timestep for ix in range(len(sqr_dis[0]))]
for ix in range(len(sqr_dis[0])):
    print(pltx[ix], end=' ', file=fileout2)
    for jx in range(len(sqr_dis)):
        print("%f" % sqr_dis[jx][ix], end=' ', file=fileout2)
    print("\n", end='', file=fileout2)
fileout1.close()
fileout2.close()

# plt.figure(1)
# plt.scatter(mass, allmsd)
# plt.figure(2)
# for ix in range(len(sqr_dis)):
#     plt.plot(pltx, sqr_dis[ix])
# plt.show()
# 算出的单方向msd系数未除2
