#!/usr/bin/env python
# ./calc_share_v2.py -i "*.lammpstrj"
import numpy as np
import argparse
import os
import glob
from collections import Counter
import pandas as pd


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

    def loadtrj(self):
        print(f" -> [処理中] {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.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 classify_I(self):
        I_S_CUTOFF = 4.7
        total_counts = Counter()
        total_I_count = 0 
        
        for step in range(self.step):
            p_xyz = self.p_xyz[step]
            s_xyz = self.s_xyz[step]
            i_xyz = self.i_xyz[step]
            
            total_I_count += len(i_xyz)
            
            dist_ps = self.get_pbc_distance(s_xyz, p_xyz)
            s_to_p = np.argmin(dist_ps, axis=1)
            dist_is = self.get_pbc_distance(i_xyz, s_xyz)
            
            for i in range(len(i_xyz)):
                close_s_indices = np.where(dist_is[i] <= I_S_CUTOFF)[0]
                associated_ps4 = s_to_p[close_s_indices]
                ps4_counts = Counter(associated_ps4)
                
                edge = 0
                corner = 0
                for p_idx, s_count in ps4_counts.items():
                    if s_count == 2:
                        edge += 1
                    elif s_count == 1:
                        corner += 1
                
                state = (edge, corner)
                total_counts[state] += 1
                
        avg_counts = {}
        for state, count in total_counts.items():
            column_name = f"edge_{state[0]}_corner_{state[1]}"
            if total_I_count > 0:
                avg_counts[column_name] = count / total_I_count
            else:
                avg_counts[column_name] = 0.0
            
        return avg_counts


if __name__ == "__main__":
    par = argparse.ArgumentParser()
    par.add_argument('-i', '--pattern', default="*.lammpstrj",
                     help='input file pattern (e.g., "*.lammpstrj")')
    par.add_argument('-o', '--output', default="summary_table.tex",
                     help='output tex file name')
    args = par.parse_args()

    filepaths = sorted(glob.glob(args.pattern))
    
    if not filepaths:
        print(f"エラー: パターン '{args.pattern}' に一致するファイルが見つかりませんでした。")
        exit()

    print(f"計 {len(filepaths)} 個のファイルを順番に処理します...")

    all_results = []

    for index, filepath in enumerate(filepaths, start=1):
        filename = os.path.splitext(os.path.basename(filepath))[0]
        print(f"[{index}/{len(filepaths)}] ファイル処理中...")
        
        try:
            calc = LoadData(filepath)
            result = calc.classify_I()
            
            result['Composition'] = filename
            all_results.append(result)
            print(f" -> [完了] {filename} の解析が成功しました。")
            
        except Exception as e:
            print(f" -> [エラー] {filepath} の処理中に問題が発生しました: {e}")

    # --- TeX形式の出力処理 ---
    if all_results:
        print("\nTeXコードの生成中...")
        df = pd.DataFrame(all_results)
        
        # 欠損値を0で埋め、Compositionをインデックスに設定
        df = df.fillna(0)
        df.set_index('Composition', inplace=True)
        
        # 列名（状態）をアルファベット・数字順にソート
        state_cols = sorted(df.columns)
        df = df[state_cols]
        
        # 行と列を入れ替える (Transpose)
        df_t = df.T
        df_t.index.name = 'Label'
        df_t = df_t.reset_index() # Labelをデータ列に戻す
        
        # 列数からフォーマット指定子を生成
        num_cols = len(df_t.columns)
        tabular_format = "wl{2cm}" + "wc{2cm}" * (num_cols - 1)
        
        # TeXテキストの構築
        tex_lines = []
        tex_lines.append("\\begin{table}[H]")
        tex_lines.append(f" \\begin{{tabular}}\n {{{tabular_format}}}")
        tex_lines.append(" \\hline")
        
        # ヘッダー行
        headers = [str(h).replace('_', r'\_') for h in df_t.columns]
        tex_lines.append("  " + " & ".join(headers) + " \\\\")
        tex_lines.append(" \\hline")
        
        # データ行
        for _, row in df_t.iterrows():
            row_vals = []
            for i, val in enumerate(row.values):
                if i == 0:
                    # Label列（状態名）のアンダースコアをエスケープ
                    row_vals.append(str(val).replace('_', r'\_'))
                else:
                    # 数値列の処理（百分率化、四捨五入、1%以下の処理）
                    try:
                        num = float(val)
                        pct = num * 100
                        if pct == 0.0:
                            row_vals.append("0.0\\%")
                        elif pct <= 1.0:
                            # 1%以下の場合は数値ではなく「trace（微量）」という英語にする
                            row_vals.append("trace")
                        else:
                            # 小数点1桁で四捨五入
                            row_vals.append(f"{pct:.1f}\\%")
                    except ValueError:
                        row_vals.append(str(val).replace('_', r'\_'))
            tex_lines.append("  " + " & ".join(row_vals) + " \\\\")
            
        tex_lines.append(" \\end{tabular}")
        tex_lines.append("\\end{table}")
        
        tex_output = "\n".join(tex_lines)
        
        # ファイルへの保存
        with open(args.output, "w", encoding="utf-8") as f:
            f.write(tex_output)
            
        print("\n=== TeX Output ===")
        print(tex_output)
        print(f"\n結果は '{args.output}' に保存されました。")