#!/usr/bin/env python
import numpy as np
import io
import re
import pickle
import glob
import pandas as pd


class LammpsTrj():
    def __init__(self, trjfile):
        self.fp = open(trjfile)
        lines = [self.fp.readline() for i in range(4)]
        self.atoms = int(lines[-1])
        self.fp.seek(0)
        self.getOnestep()
        # print(self.df)
        self.fp.seek(0)
        self.molecules = np.unique(self.df["mol"]).shape[0]

    def __del__(self):
        self.fp.close()

    def getOnestep(self):
        block = self.atoms + 9
        lines = [self.fp.readline() for i in range(block)]
        if lines[0] == "":
            return 0
        self.step = int(lines[1])
        self.lx = float(lines[5].split()[1])
        self.ly = float(lines[6].split()[1])
        self.lz = float(lines[7].split()[1])
        header = lines[8].split()[2:]
        data = io.StringIO(" ".join(lines[9:]))
        self.df = pd.read_csv(data, delimiter=r"\s+", header=None)
        self.df.columns = header
        return 1

    def getOrientation(self, benzene_id):
        xyz = self.df[["xu", "yu", "zu"]].to_numpy()
        xyz = xyz.reshape((self.molecules, -1, 3))
        benzene_id = np.array(benzene_id, dtype=int) - 1
        xyz = xyz[:, benzene_id, :]
        g = xyz.mean(axis=1)
        p = xyz[:, [1, 2], :].mean(axis=1)
        vec = p - g
        angle = np.arctan2(vec[:, 1], vec[:, 0])
        return angle
        # for i, xyz_i in enumerate(xyz):
        #     body = "{}\n\n".format(benzene_id.shape[0])
        #     for r in xyz_i:
        #         body += "C {} {} {}\n".format(*r)
        #     outfile = f"{i:04d}.xyz"
        #     with open(outfile, "w") as o:
        #         o.write(body)

    def getZposition(self, benzene_id):
        xyz = self.df[["xu", "yu", "zu"]].to_numpy()
        xyz = xyz.reshape((self.molecules, -1, 3))
        benzene_id = np.array(benzene_id, dtype=int) - 1
        xyz = xyz[:, benzene_id, :]
        zpos = xyz.mean(axis=1)[:, 2]
        return zpos


def run(trjfile):
    benzene_id = [24, 25, 26, 27, 28, 29]
    lmp = LammpsTrj(trjfile)
    orientation = np.empty((0, lmp.molecules))
    zpos = np.empty((0, lmp.molecules))
    while lmp.getOnestep():
        # print(lmp.step)
        orientation_ = lmp.getOrientation(benzene_id)
        orientation = np.vstack((orientation, orientation_))
        zpos_ = lmp.getZposition(benzene_id)
        zpos = np.vstack((zpos, zpos_))
    hist, x = np.histogram(orientation, bins=np.linspace(-np.pi, np.pi, 200))
    df = pd.DataFrame()
    df["angle"] = x[0:-1] + (x[1]-x[0])*0.5
    df["angle_y"] = hist
    hist, x = np.histogram(zpos, bins=np.linspace(24, 25, 200))
    df["zpos"] = x[0:-1] + (x[1]-x[0])*0.5
    df["zpos_y"] = hist
    temp = int(re.search(r"_(\d+)K\.", trjfile).group(1))
    dset[temp] = {}
    dset[temp]["orientation"] = df


for k in ["D3d-o", "D3d-p"]:
    dset = {}
    files = sorted(glob.glob(f"{k}_*K.lammpstrj"))
    for f in files:
        print(f)
        run(f)
    outfile = f"{k}.pkl"
    with open(outfile, "wb") as f:
        pickle.dump(dset, f)
    print(f"{outfile} was created.")