import MDAnalysis as mda
from MDAnalysis.analysis import rdf
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

# ==================== 1. 全局参数配置 ====================
DATA_FILE = "annealing.data"  # 拓扑文件
TRAJ_FILE = "run2.lammpstrj"   # 轨迹文件
NBINS = 300                    # 柱状图数量
R_RANGE = (0.0, 12.0)          # 统计距离范围 (埃)

# 🚀 【核心配置区】现在完全基于 LAMMPS 的 type 来定义原子
ION_TASKS = {
    "Mg": {
        "ion_type": "1",         # 🌟 Mg 离子的 LAMMPS type (请根据你的 data 文件修改)
        "water_o_type": "4",      # 对应周围原子的 type
        "color": "#1f77b4"        # 绘图颜色：蓝色
    },
    "Li": {
        "ion_type": "2",         # 🌟 Li 离子的 LAMMPS type (请根据你的 data 文件修改)
        "water_o_type": "4",      # 对应周围原子的 type
        "color": "#ff7f0e"        # 绘图颜色：橙色
    },
    "Cl": {
        "ion_type": "3",         # 🌟 Cl 离子的 LAMMPS type (请根据你的 data 文件修改)
        "water_o_type": "4",      # 🌟 若算 Cl-水氢，记得改成氢的 type
        "color": "#2ca02c"        # 绘图颜色：绿色
    }
}

# ==================== 2. 计算与绘图核心逻辑 ====================
def run_all_rdf_analysis():
    print("正在加载 Universe (全局只需加载一次)...")
    u = mda.Universe(DATA_FILE, TRAJ_FILE, format="LAMMPSDUMP")
    
    # 建立一个字典，用来存下所有离子的计算结果，方便最后画合并图
    all_results = {}

    # ---------------- 迭代计算每种离子 ----------------
    for ion_name, config in ION_TASKS.items():
        print(f"\n⚡ 正在计算 {ion_name} 的 RDF...")
        
        # 🌟 核心修改：全部基于 type 进行原子选择
        ions = u.select_atoms(f"type {config['ion_type']}")
        water_o = u.select_atoms(f"type {config['water_o_type']}")

        if len(ions) == 0 or len(water_o) == 0:
            print(f"❌ 错误 ({ion_name}): 未找到原子！离子数: {len(ions)}, 目标原子数: {len(water_o)}")
            continue

        # 计算 RDF
        rdf_analyzer = rdf.InterRDF(ions, water_o, nbins=NBINS, range=R_RANGE)
        rdf_analyzer.run()

        bins = rdf_analyzer.results.bins
        r_values = rdf_analyzer.results.rdf
        
        # 缓存结果供后续合并图使用
        all_results[ion_name] = (bins, r_values)

        # 导出该离子的 CSV 数据
        output_csv = f"{ion_name}_rdf_output.csv"
        pd.DataFrame({'r_Angstrom': bins, 'g_r': r_values}).to_csv(output_csv, index=False)
        print(f"💾 数据已导出至: {output_csv}")

        # 输出第一峰特征信息
        max_idx = np.argmax(r_values)
        print(f"✨ {ion_name} 第一峰位置: {bins[max_idx]:.2f} Å, 高度: {r_values[max_idx]:.2f}")

        # 📊 绘制【分开的图】
        plt.figure(figsize=(8, 5), dpi=100)
        plt.plot(bins, r_values, label=f'{ion_name} System', color=config['color'], linewidth=2)
        plt.axhline(y=1, color='black', linestyle='--', alpha=0.3, label='Baseline (g=1)')
        plt.title(f'Radial Distribution Function - {ion_name}', fontsize=12)
        plt.xlabel(r'Distance $r$ ($\AA$)')
        plt.ylabel(r'$g(r)$')
        plt.xlim(0, R_RANGE[1])
        plt.ylim(0, max(r_values) * 1.1)
        plt.grid(True, linestyle=':', alpha=0.6)
        plt.legend()
        
        # 保存并显示分开的图
        output_single_img = f"{ion_name}_rdf.png"
        plt.savefig(output_single_img)
        print(f"🖼️ 单独图已保存为 '{output_single_img}'")
        plt.show()  
        plt.close() # 必须关闭当前画布，防止下一张图污染

    # ---------------- 绘制【合在一起的图】 ----------------
    if all_results:
        print("\n📊 正在生成多离子总对比图...")
        plt.figure(figsize=(10, 6), dpi=120)
        
        # 将所有离子的曲线画在同一个 figure 里
        for ion_name, (bins, r_values) in all_results.items():
            config = ION_TASKS[ion_name]
            plt.plot(bins, r_values, label=f'{ion_name} - Water', color=config['color'], linewidth=2)
            
        plt.axhline(y=1, color='black', linestyle='--', alpha=0.3, label='Baseline (g=1)')
        plt.title('Comparison of Radial Distribution Functions $g(r)$', fontsize=14)
        plt.xlabel(r'Distance $r$ ($\AA$)', fontsize=12)
        plt.ylabel(r'$g(r)$', fontsize=12)
        plt.xlim(0, R_RANGE[1])
        
        # 动态调整合并图的纵坐标范围，防止曲线封顶
        all_max_y = max([max(res[1]) for res in all_results.values()])
        plt.ylim(0, all_max_y * 1.1)
        
        plt.legend(fontsize=11)
        plt.grid(True, linestyle=':', alpha=0.6)
        
        # 保存并显示合在一起的图
        output_combined_img = "all_ions_rdf_comparison.png"
        plt.savefig(output_combined_img)
        print(f"🖼️ 总对比图已保存为 '{output_combined_img}'")
        plt.show()
        plt.close()

# ==================== 3. 执行 ====================
if __name__ == "__main__":
    run_all_rdf_analysis()