#!/usr/bin/env python
# ./voronoi_count_diff.py -i *trj
import numpy as np
import matplotlib.pyplot as plt
import argparse
import os

"""
複数の組成のtrjファイルを読み込み、
横軸：組成（系のP/P+I比率）
縦軸：
  青：系のP/P+Iの比率
  緑：Pに近いLi/Liの総数の比率（step平均）
をプロットするプログラム
"""


class LoadData():
    def __init__(self, trjfile):
        self.trjfile = trjfile
        self.loadtrj()

    def loadtrj(self):
        print(f"Loading {os.path.basename(self.trjfile)} ...")
        with open(self.trjfile) as o:
            d = o.read()
        lines = d.splitlines()    
        self.atoms = int(lines[3])

        x_bnd = [float(x) for x in lines[5].split()]
        y_bnd = [float(x) for x in lines[6].split()]
        z_bnd = [float(x) for x in lines[7].split()]
        self.box_size = np.array(
            [x_bnd[1]-x_bnd[0], y_bnd[1]-y_bnd[0], z_bnd[1]-z_bnd[0]])

        data = np.array(" ".join(lines).split())
        self.data = data.reshape(-1, 28+self.atoms*6)
        self.step = self.data.shape[0]
        data = self.data[:, 28:].reshape(self.step, self.atoms, -1)

        self.li_data = data[data[:, :, 2] == "Li"].reshape(self.step, -1, 6)
        self.li_xyz = self.li_data[:, :, 3:].astype(float)

        self.p_data = data[data[:, :, 2] == "P"].reshape(self.step, -1, 6)
        self.p_xyz = self.p_data[:, :, 3:].astype(float)
        
        self.i_data = data[data[:, :, 2] == "I"].reshape(self.step, -1, 6)
        self.i_xyz = self.i_data[:, :, 3:].astype(float)

        self.s_data = data[data[:, :, 2] == "S"].reshape(self.step, -1, 6)
        self.s_xyz = self.s_data[:, :, 3:].astype(float)
        
        return data

    def get_pbc_distance(self, xyz1, xyz2):
        d = xyz1[:, np.newaxis, :] - xyz2[np.newaxis, :, :]
        d -= self.box_size * np.round(d / self.box_size)
        return np.linalg.norm(d, axis=-1)

    def calc_ratios(self):
        p_ratios = []

        print(f"Calculating for {os.path.basename(self.trjfile)} ...")
        
        # 系のP/P+Iの比率はステップ間で変わらないため、最初のステップで計算
        p_num = self.p_xyz[0].shape[0]
        i_num = self.i_xyz[0].shape[0]
        total_num = p_num + i_num
        system_p_ratio = p_num / total_num if total_num > 0 else 0.0

        # 各ステップのPに近いLiの比率を計算
        for step in range(self.step):
            li_pos = self.li_xyz[step]
            p_pos = self.p_xyz[step]
            i_pos = self.i_xyz[step]

            # 周期境界条件を考慮し、距離を計算
            diff_li_p = self.get_pbc_distance(li_pos, p_pos)
            diff_li_i = self.get_pbc_distance(li_pos, i_pos)

            # 各Liについて、最も近いPとIまでの距離を取得
            min_diff_p = np.min(diff_li_p, axis=1)
            min_diff_i = np.min(diff_li_i, axis=1)

            # Pの方が近いLiの数、Iの方が近いLiの数をカウント
            count_p = np.sum(min_diff_p < min_diff_i)
            count_i = np.sum(min_diff_p > min_diff_i)

            # Pが一番近いLiの比率を計算
            total = count_p + count_i
            ratio = count_p / total if total > 0 else 0.0
            p_ratios.append(ratio)

        # ステップ平均を計算
        avg_near_p_ratio = np.mean(p_ratios)
        
        return system_p_ratio, avg_near_p_ratio

    
if __name__ == "__main__":
    description = """This is a test program"""
    par = argparse.ArgumentParser(description=description)

    par.add_argument('-i', '--trjfiles', default="", required=True, nargs="+",
                     help='input file')
    args = par.parse_args()

    results = []

    for trjfile in args.trjfiles:
        trj = LoadData(trjfile)
        sys_ratio, avg_near_ratio = trj.calc_ratios()
        results.append((sys_ratio, avg_near_ratio))

    comps = [int(r[0]*100) for r in results][::-1]
    y_sys_ratios = [r[0] for r in results][::-1] 
    y_near_ratios = [r[1] for r in results][::-1]
    
    # --------------------------------------------------    
    # プロット
    # --------------------------------------------------
    print("Plotting results...")
    fig, ax = plt.subplots(figsize=(5, 3.5))

    color1, color2 = "tab:blue", "tab:green"
    
    ax.set_xlabel("Composition")
    ax.set_ylabel("Ratio of P")

    xtick_labels = [f"LPSI{comp:02d}" for comp in comps]
    ax.set_xticks(comps)
    ax.set_xticklabels(xtick_labels, fontsize=11)
    
    # マーカーをつけてプロット（複数の組成点が分かりやすいように）
    ax.plot(comps, y_sys_ratios, color=color1, marker='o', linestyle='-',
            label="Ratio of P (system)")
    ax.plot(comps, y_near_ratios, color=color2, marker='s', linestyle='-',
            label="Ratio of P (near Li, step avg)")
    
    ax.legend(loc='best')
    ax.grid(True, linestyle='--', alpha=0.7)
    
    fig.tight_layout()
    fig.savefig("voronoi_count_diff.pdf", dpi=300)
    plt.show()