#!/usr/bin/env python
# ./dump2cube.py -i 050Li3PS4-050LiI.lammpstrj -m 160 160 160 -cut 8
import numpy as np
import argparse
import re
import os
import time

description = """This is a test program"""
par = argparse.ArgumentParser(description=description)

par.add_argument('-i', '--trjfile', default="", required=True,
                 help='input trjfile')
par.add_argument('-m', '--mesh', default=(30, 30, 30), required=False,
                 nargs=3, type=int, help='mesh grid')
par.add_argument('-cut', '--cutoff', default=7, required=False,
                 type=float, help='mesh grid')

args = par.parse_args()
ang_borr = 0.5291772109217
dirname = os.path.dirname(__file__)


class LammpsTrj():
    def __init__(self):
        """
        trjファイルを読み込んで、原子数、lattice、ステップ数を格納
        """
        with open(args.trjfile) as o:
            data = o.read().split()
        self.atoms = int(data[7])  # 原子数
        data = np.array(data).reshape(-1, self.atoms*7+32)
        self.lattices = data[:, 17:23].astype(float)  # (step, 6)
        self.data = data[:, 32:]  # 座標データ
        self.steps = data.shape[0]
        self.mesh = args.mesh
        self.cutoff = args.cutoff

    def setLattice(self, lat):
        """
        latticeの形によって操作を分岐
        self.M, self.M_を作成
        """
        if lat.shape[0] == 6:
            M = np.array([[lat[1]-lat[0], 0, 0],
                          [0, lat[3]-lat[2], 0],
                          [0, 0, lat[5]-lat[4]]]).astype(float)
        if lat.shape[0] == 9:
            xlo_bound, xhi_bound, xy = lat[0, 0], lat[0, 1], lat[0, 2]
            ylo_bound, yhi_bound, xz = lat[1, 0], lat[1, 1], lat[1, 2]
            zlo_bound, zhi_bound, yz = lat[2, 0], lat[2, 1], lat[2, 2]
            xlo = xlo_bound - np.min([0.0, xy, xz, xy+xz])
            xhi = xhi_bound - np.max([0.0, xy, xz, xy+xz])
            ylo = ylo_bound - np.min([0.0, yz])
            yhi = yhi_bound - np.max([0.0, yz])
            zlo = zlo_bound
            zhi = zhi_bound
            lx = xhi - xlo
            ly = yhi - ylo
            lz = zhi - zlo
            a = lx
            b = np.sqrt(ly**2 + xy**2)
            c = np.sqrt(lz**2 + xz**2 + yz**2)
            alpha = np.arccos((xy*xz + ly*yz)/b/c)
            beta = np.arccos(xz/c)
            gamma = np.arccos(xy/b)
            v1 = [a, 0, 0]
            v2 = [b*np.cos(gamma), b*np.sin(gamma), 0]
            v3 = [c*np.cos(beta),
                  c*(np.cos(alpha)-np.cos(beta)*np.cos(gamma))/np.sin(gamma),
                  c*np.sqrt(1+2*np.cos(alpha)*np.cos(beta)*np.cos(gamma)
                            - np.cos(alpha)**2-np.cos(beta)**2
                            - np.cos(gamma)**2 / np.sin(gamma))]
            M = np.array([v1, v2, v3])
        return M

    def getOnestep(self, step):
        """
        引数のstepにおける座標を、原子ごとにself.elems(辞書)に格納する
        self.elemsの座標は分率座標
        """
        step_data = self.data[step, :].reshape(self.atoms, -1)  # (atoms, 6)
        self.M = self.setLattice(self.lattices[step])
        self.M_ = np.linalg.inv(self.M)
        self.elems = {}
        for e in np.unique(step_data[:, 2]):
            d = step_data[step_data[:, 2] == e, 3:].astype(float)
            dxyz = d[:, 0:3]
            dxyz = dxyz @ self.M_
            dxyz = dxyz - np.floor(dxyz)
            d[:, 0:3] = dxyz
            self.elems[e] = d  # 元素ごとに結晶座標とmol番号を格納
        return self.elems

    def makeCube(self, trjfile):
        """
        すべてのP原子について処理を行う。
        P原子ごとに最も近い4つのI原子を選出し、回転・アライメントを行う。
        """
        # Li用の座標リスト
        self.li_coords_list = []

        # I用のリスト (全てまとめて格納)
        self.i_coords_list = []

        # ヒストグラム結果格納用
        self.li_hist = None
        self.i_hist = None

        count_p = 0  # 処理したP原子の総数カウント用

        for step in range(self.steps):
            print(f"processing {step} step")
            self.elems = self.getOnestep(step)
            elems_ = {k: v.copy() for k, v in self.elems.items()}  # copy

            # 全てのP原子に対してループ処理を行う
            for p_data in elems_["P"]:

                # 全てのI原子の中から、このP原子に最も近い4つのI原子を特定する
                i_diffs = elems_["I"][:, 0:3] - p_data[0:3]
                i_diffs = i_diffs - np.around(i_diffs)  # 周期境界条件の考慮(分率)
                i_diffs_abs = i_diffs @ self.M         # 絶対座標へ変換
                distances = np.linalg.norm(i_diffs_abs, axis=1)

                sorted_indices = np.argsort(distances)
                closest_Is = elems_["I"][sorted_indices[:4]]  # 近い順に4つ抽出

                nearest_I_coords = closest_Is[0]  # 回転定義に使用する「最も近いI」

                # S原子の処理 (同じmol番号のSを取得)
                s_data = elems_["S"][elems_["S"][:, 3] == p_data[3]]

                # 1つのPS4に注目して、Pを原点に配置するための距離計算
                si_diff = s_data[:, 0:3] - nearest_I_coords[0:3]
                si_diff = si_diff - np.around(si_diff)
                si_diff = np.linalg.norm(si_diff, axis=1)

                # 最も遠いSをz軸上に配置するため、そのインデックスを取得
                n_S_idx = np.argmax(si_diff)
                # その他のインデックスを取得
                s_indices = [0, 1, 2, 3]
                s_indices.remove(n_S_idx)
                n2_S_idx = s_indices[0]

                s_xyz = s_data[:, 0:3] - p_data[0:3]
                s_xyz = s_xyz - np.round(s_xyz)
                s_xyz = s_xyz @ self.M  # 絶対座標へ

                # 1つのSをyz平面上に配置(z軸周り回転)
                theta = np.arctan2(s_xyz[n_S_idx][0], s_xyz[n_S_idx][1])
                self.Mat_z = np.array([[np.cos(-theta), np.sin(-theta), 0],
                                       [-np.sin(-theta), np.cos(-theta), 0],
                                       [0, 0, 1]])
                s_xyz = (self.Mat_z @ s_xyz.T).T

                # 1つのSをz軸上に配置(x軸周り回転)
                theta2 = np.arctan2(s_xyz[n_S_idx][1], s_xyz[n_S_idx][2])
                self.Mat_x = np.array([[1, 0, 0],
                                       [0, np.cos(-theta2), np.sin(-theta2)],
                                       [0, -np.sin(-theta2), np.cos(-theta2)]])
                s_xyz = (self.Mat_x @ s_xyz.T).T

                # もう1つのSをyz平面上に配置(z軸周り回転)
                theta3 = np.arctan2(s_xyz[n2_S_idx][0], s_xyz[n2_S_idx][1])
                self.Mat_z2 = np.array([[np.cos(-theta3), np.sin(-theta3), 0],
                                        [-np.sin(-theta3), np.cos(-theta3), 0],
                                        [0, 0, 1]])
                s_xyz = (self.Mat_z2 @ s_xyz.T).T
                self.s_xyz = s_xyz

                # PS4座標保存 (Output用)
                p_xyz = np.array([0, 0, 0])
                self.ps4_coord = np.vstack((p_xyz, s_xyz))
                self.ps4_coord = self.ps4_coord / ang_borr

                # 座標回転関数の定義
                def rotate_coords(coords_fractional):
                    if len(coords_fractional) == 0:
                        return np.array([])
                    rot_xyz = coords_fractional[:, 0:3] - p_data[0:3]
                    rot_xyz = rot_xyz - np.round(rot_xyz)
                    rot_xyz = rot_xyz @ self.M
                    rot_xyz = (self.Mat_z @ rot_xyz.T).T
                    rot_xyz = (self.Mat_x @ rot_xyz.T).T
                    rot_xyz = (self.Mat_z2 @ rot_xyz.T).T
                    return rot_xyz

                # ヒストグラム範囲
                bounds = [[-(self.cutoff)/ang_borr,
                           (self.cutoff)/ang_borr]] * 3

                # --- Liの処理: 座標をリストに追加 ---
                segment_xyz_Li = []
                for xyz in self.elems["Li"]:
                    diff = xyz[0:3] - p_data[0:3]
                    diff = diff - np.around(diff)
                    diff = diff @ self.M
                    diff = np.linalg.norm(diff)
                    for xyzdata in xyz[diff < self.cutoff, :]:
                        segment_xyz_Li.append(xyzdata)
                segment_xyz_Li = np.array(segment_xyz_Li, dtype=float)

                if len(segment_xyz_Li) > 0:
                    rotationed_Li = rotate_coords(segment_xyz_Li)
                    self.li_coords_list.extend(rotationed_Li.tolist())

                # --- Iの処理: 抽出した4つのIを回転させてまとめてリストに追加 ---
                rotationed_Is = rotate_coords(closest_Is)
                self.i_coords_list.extend(rotationed_Is.tolist())

                count_p += 1

        # ループ終了後にまとめてヒストグラム計算
        print("Calculating Histograms...")
        volume = (self.cutoff*2)**3
        bounds = [[-(self.cutoff)/ang_borr, (self.cutoff)/ang_borr]] * 3

        # Li
        if len(self.li_coords_list) > 0:
            print("num_Li: ", len(self.li_coords_list))
            li_arr = np.array(self.li_coords_list) / ang_borr
            li_hist, _ = np.histogramdd(li_arr, bins=self.mesh, range=bounds)
            self.li_hist = li_hist.ravel() / len(self.li_coords_list) / volume
            print("sum_Li: ", np.sum(self.li_hist) * volume)

        # I (All 4 atoms combined)
        if len(self.i_coords_list) > 0:
            print("num_I: ", len(self.i_coords_list))
            i_arr = np.array(self.i_coords_list) / ang_borr
            i_hist, _ = np.histogramdd(i_arr, bins=self.mesh, range=bounds)
            self.i_hist = i_hist.ravel() / len(self.i_coords_list) / volume
            print("sum_I: ", np.sum(self.i_hist) * volume)

    def outputCube(self):
        output_start_time = time.time()
        base = re.match(r"(\d{3}.*?LiI).*?",
                        args.trjfile.split("/")[-1]).group(1)
        self.atomsDic = {'I': '53', 'Li': '3', 'P': '15', 'S': '16'}

        if self.li_hist is None and self.i_hist is None:
            print("No data found.")
            return

        def get_header():
            body = f"created from {__file__}, {args}\n"
            body += "Contains the selected quantity on a FFT grid\n"
            origin = [-(self.cutoff)/ang_borr] * 3
            body += "{:>5d}{:>12.7f}{:>12.7f}{:>12.7f}\n".format(
                5, *origin)
            body += "{:>5d}{:>12.7f}{:>12.7f}{:>12.7f}\n".format(
                self.mesh[0], (self.cutoff*2)/self.mesh[0]/ang_borr, 0, 0)
            body += "{:>5d}{:>12.7f}{:>12.7f}{:>12.7f}\n".format(
                self.mesh[1], 0, (self.cutoff*2)/self.mesh[1]/ang_borr, 0)
            body += "{:>5d}{:>12.7f}{:>12.7f}{:>12.7f}\n".format(
                self.mesh[2], 0, 0, (self.cutoff*2)/self.mesh[2]/ang_borr)
            # Pの座標
            body += "{:>5d}{:>12.7f}{:>12.7f}{:>12.7f}{:>12.7f}\n".format(
                int(self.atomsDic["P"]), float(self.atomsDic["P"]), *self.ps4_coord[0])
            # Sの座標
            for s_coord in self.ps4_coord[1:5]:
                body += "{:>5d}{:>12.7f}{:>12.7f}{:>12.7f}{:>12.7f}\n".format(
                    int(self.atomsDic["S"]), float(self.atomsDic["S"]), *s_coord)
            return body

        # Liの出力
        if self.li_hist is not None:
            body = get_header()
            for idx, r in enumerate(self.li_hist):
                if idx % 6 == 5:
                    body += "{:>13.5E}\n".format(r)
                else:
                    body += "{:>13.5E}".format(r)

            outfile = f"{dirname}/{base}_PS4_Li.cube"
            with open(outfile, "w") as o:
                o.write(body)
            print(f"{outfile} was created.")

        # Iの出力 (まとめて1ファイルに出力)
        if self.i_hist is not None:
            body = get_header()
            for idx, r in enumerate(self.i_hist):
                if idx % 6 == 5:
                    body += "{:>13.5E}\n".format(r)
                else:
                    body += "{:>13.5E}".format(r)

            outfile = f"{dirname}/{base}_PS4_I.cube"
            with open(outfile, "w") as o:
                o.write(body)
            print(f"{outfile} was created.")

        output_end_time = time.time()
        print(f"output_time : {output_end_time - output_start_time} s")


trj = LammpsTrj()
trj.makeCube(args.trjfile)
trj.outputCube()
