import os
import re
import numpy as np
import matplotlib.pyplot as plt
from scipy import stats

# 物理常数（SI单位）
kB = 1.380649e-23  # 玻尔兹曼常数 (J/K)
q = 1.602176634e-19  # 元电荷 (C，Li+的电荷)
ps_to_s = 1e-12  # 皮秒转秒
angstrom2_to_m2 = 1e-20  # 埃²转米²
angstrom3_to_m3 = 1e-30  # 埃³转米³


def extract_temperature(folder_name):
    """从文件夹名（如320k）提取温度（K）"""
    match = re.search(r'(\d+)K', folder_name)
    if match:
        return int(match.group(1))
    return None


def read_msd(msd_path):
    """读取msd_li.dat，返回时间（ps）和总MSD（Ų）"""
    try:
        data = np.loadtxt(msd_path, comments='#')
    except Exception as e:
        raise ValueError(f"读取MSD文件失败：{e}")
    
    if data.size == 0:
        raise ValueError("MSD文件内容为空")
    
    # 时间步转实际时间（ps）：TimeStep × 0.001
    time = data[:, 0] * 0.001  # 第一列是TimeStep，乘以0.001转为ps
    msd_total = data[:, -1]    # 最后一列是总MSD
    
    # 打印时间范围（关键调试信息）
    print(f"  MSD时间范围：{time.min():.1f} ~ {time.max():.1f} ps")
    return time, msd_total


def calculate_diffusion(time, msd, t_start, t_end):
    """根据MSD计算扩散系数D（单位：m²/s）"""
    # 筛选时间范围内的数据
    mask = (time >= t_start) & (time <= t_end)
    t_filtered = time[mask] * ps_to_s  # 转为秒
    msd_filtered = msd[mask] * angstrom2_to_m2  # 转为m²
    
    # 检查筛选后的数据是否为空
    if len(t_filtered) < 2:
        raise ValueError(f"筛选后的数据点不足（仅{len(t_filtered)}个），请调整t_start/t_end")
    
    # 线性拟合：MSD vs 时间，斜率 = 6D
    slope, _, r_value, _, _ = stats.linregress(t_filtered, msd_filtered)
    D = slope / 6
    
    # 打印拟合信息（调试用）
    print(f"  MSD拟合：斜率 = {slope:.6e}，R² = {r_value**2:.4f}")
    return D


def parse_lammps_data(data_path):
    """解析LAMMPS data文件，提取盒子体积和Li+数量（Li的原子类型为1）"""
    with open(data_path, 'r') as f:
        lines = f.readlines()
    
    # 提取原子总数
    atom_count_line = next((line for line in lines if 'atoms' in line), None)
    if not atom_count_line:
        raise ValueError("无法在data文件中找到原子总数")
    total_atoms = int(atom_count_line.split()[0])
    
    # 提取盒子尺寸
    x_line = next((line for line in lines if 'xlo xhi' in line), None)
    y_line = next((line for line in lines if 'ylo yhi' in line), None)
    z_line = next((line for line in lines if 'zlo zhi' in line), None)
    
    if not (x_line and y_line and z_line):
        raise ValueError("无法在data文件中找到盒子尺寸")
    
    xlo, xhi = map(float, x_line.split()[:2])
    ylo, yhi = map(float, y_line.split()[:2])
    zlo, zhi = map(float, z_line.split()[:2])
    
    # 计算盒子体积（ų）
    volume_angstrom3 = (xhi - xlo) * (yhi - ylo) * (zhi - zlo)
    
    # 手动指定Li的原子类型为1
    li_type = 1
    
    # 统计Li原子数量
    atoms_section_start = next((i for i, line in enumerate(lines) if 'Atoms' in line), None)
    if atoms_section_start is None:
        raise ValueError("无法在data文件中找到Atoms部分")
    
    li_count = 0
    for i in range(atoms_section_start + 1, atoms_section_start + total_atoms + 1):
        if i >= len(lines):
            break
        parts = lines[i].split()
        if len(parts) >= 2 and int(parts[1]) == li_type:
            li_count += 1
    
    if li_count == 0:
        raise ValueError(f"未找到原子类型为{li_type}的Li原子")
    
    return li_count, volume_angstrom3


def calculate_conductivity(D, T, n):
    """根据Nernst-Einstein方程计算离子电导σ（单位：S/m）"""
    sigma = (n * q**2 * D) / (kB * T)
    return sigma

def arrhenius_plot(temperatures, D_list, sigma_list):
    """绘制阿伦尼乌斯图，拟合400K前后两条直线，基于320-400K数据外推300K电导"""
    plt.rcParams["axes.unicode_minus"] = False  # 确保负号正常显示
    temps = np.array(temperatures, dtype=float)
    sigmas = np.array(sigma_list, dtype=float)
    inv_T = 1000 / temps  # 1000/T (1/K)

    # 数据分组：320-400K（用于拟合和外推）和>400K（仅拟合）
    mask_below = (temps >= 320) & (temps <= 400)  # 320-400K区间
    mask_above = temps > 400                     # 400K以上区间

    # 检查低温度区间数据量
    if np.sum(mask_below) < 2:
        raise ValueError(f"320-400K区间数据点不足（仅{np.sum(mask_below)}个），无法拟合直线")

    # 1. 拟合320-400K区间（用于外推300K）
    inv_T_below = inv_T[mask_below]
    ln_sigma_below = np.log(sigmas[mask_below])
    slope_below, intercept_below, r_below, _, _ = stats.linregress(inv_T_below, ln_sigma_below)
    Ea_below = -slope_below * kB * 1000  # 活化能（J）
    Ea_below_eV = Ea_below / q           # 转换为eV

    # 2. 拟合400K以上区间（若有足够数据）
    slope_above = None
    intercept_above = None
    Ea_above_eV = None
    r_above = None
    if np.sum(mask_above) >= 2:
        inv_T_above = inv_T[mask_above]
        ln_sigma_above = np.log(sigmas[mask_above])
        slope_above, intercept_above, r_above, _, _ = stats.linregress(inv_T_above, ln_sigma_above)
        Ea_above = -slope_above * kB * 1000
        Ea_above_eV = Ea_above / q

    # 3. 外推300K的离子电导（基于320-400K拟合线）
    T_target = 300  # 目标温度300K
    inv_T_target = 1000 / T_target  # 1000/300 (1/K)
    ln_sigma_target = slope_below * inv_T_target + intercept_below  # 用低温度区间拟合线外推
    sigma_target = np.exp(ln_sigma_target)  # 转换为电导值

    # 4. 绘图
    fig, ax1 = plt.subplots(figsize=(10, 6))

    # 绘制数据点
    ax1.scatter(inv_T_below, ln_sigma_below, label='320-400K Data', color='blue', s=60, zorder=3)
    if np.sum(mask_above) >= 1:
        ax1.scatter(inv_T[mask_above], np.log(sigmas[mask_above]), label='>400K Data', color='red', s=60, zorder=3)

    # 绘制拟合线（低温度区间）
    x_fit_below = np.linspace(min(inv_T_below), max(inv_T_below), 100)
    y_fit_below = slope_below * x_fit_below + intercept_below
    ax1.plot(x_fit_below, y_fit_below, '--', color='blue', 
             label=f'320-400K Fit (Eₐ={Ea_below_eV:.3f} eV, R²={r_below**2:.3f})', zorder=2)

    # 绘制拟合线（高温度区间，若有）
    if slope_above is not None:
        x_fit_above = np.linspace(min(inv_T[mask_above]), max(inv_T[mask_above]), 100)
        y_fit_above = slope_above * x_fit_above + intercept_above
        ax1.plot(x_fit_above, y_fit_above, '--', color='red', 
                 label=f'>400K Fit (Eₐ={Ea_above_eV:.3f} eV, R²={r_above**2:.3f})', zorder=2)

    # 标记300K外推点
    ax1.scatter([inv_T_target], [ln_sigma_target], color='green', marker='*', s=150, 
                label=f'300K Extrapolated', zorder=4)

    # 坐标轴设置
    ax1.set_xlabel('1000/T (1/K)')
    ax1.set_ylabel('ln(σ) (S/m)')
    ax1.set_title('Arrhenius Plot with Dual Fitting')
    ax1.legend()
    ax1.grid(alpha=0.3)

    # 顶部温度坐标轴
    ax2 = ax1.twiny()
    x1_min, x1_max = ax1.get_xlim()
    ax2.set_xlim(x1_min, x1_max)
    ax2.set_xticks(1000 / np.unique(temps))  # 显示现有温度刻度
    ax2.set_xticklabels([f"{int(t)}" for t in np.unique(temps)])
    ax2.set_xlabel('Temperature (K)')

    plt.tight_layout()
    plt.savefig('arrhenius_dual_fit.png', dpi=300)
    plt.show()

    # 输出结果
    print("\n===== 拟合结果 =====")
    print(f"320-400K区间：")
    print(f"  活化能 Eₐ = {Ea_below_eV:.3f} eV (R² = {r_below**2:.3f})")
    if Ea_above_eV is not None:
        print(f">400K区间：")
        print(f"  活化能 Eₐ = {Ea_above_eV:.3f} eV (R² = {r_above**2:.3f})")
    
    print(f"\n===== 300K外推结果（基于320-400K拟合） =====")
    print(f"  离子电导率 σ = {sigma_target:.6e} S/m")
    print(f"  转换单位：σ = {sigma_target * 10:.6e} mS/cm")  # 1 S/m = 0.1 mS/cm，乘以10转换


if __name__ == "__main__":
    # ------------ 必须根据实际数据调整的参数 ------------
    base_dir = "../data"  # 主文件夹路径
    t_start = 100  # MSD拟合起始时间（ps）
    t_end = 2000    # MSD拟合结束时间（ps）
    # ------------------------------------------------

    temperatures = []
    D_list = []
    sigma_list = []
    temp_folder_pattern = re.compile(r'^\d+K$')

    for folder in os.listdir(base_dir):
        folder_path = os.path.join(base_dir, folder)
        if not os.path.isdir(folder_path) or not temp_folder_pattern.match(folder):
            continue  # 跳过非温度文件夹

        T = extract_temperature(folder)
        print(f"\n处理文件夹: {folder}（温度 {T} K）")
        if T is None:
            print("  跳过：无法提取温度")
            continue

        # 检查文件是否存在
        msd_path = os.path.join(folder_path, "msd_li.dat")
        data_path = os.path.join(folder_path, "LYC.data")
        if not os.path.exists(msd_path):
            print(f"  跳过：缺少{msd_path}")
            continue
        if not os.path.exists(data_path):
            print(f"  跳过：缺少{data_path}")
            continue

        # 解析Li数量和体积
        try:
            li_count, volume_angstrom3 = parse_lammps_data(data_path)
            volume_m3 = volume_angstrom3 * angstrom3_to_m3
            n = li_count / volume_m3
            print(f"  Li+数量: {li_count}，盒子体积: {volume_angstrom3:.2f} ų")
            print(f"  载流子浓度: {n:.6e} m⁻³")
        except Exception as e:
            print(f"  跳过：解析data文件失败: {e}")
            continue

        # 读取MSD数据
        try:
            time, msd = read_msd(msd_path)
        except Exception as e:
            print(f"  跳过：读取MSD失败: {e}")
            continue

        # 计算扩散系数
        try:
            D = calculate_diffusion(time, msd, t_start, t_end)
            print(f"  扩散系数 D = {D:.6e} m²/s")
        except Exception as e:
            print(f"  跳过：计算扩散系数失败: {e}")
            continue

        # 计算离子电导
        try:
            sigma = calculate_conductivity(D, T, n)
            print(f"  离子电导 σ = {sigma:.6e} S/m")
        except Exception as e:
            print(f"  跳过：计算电导失败: {e}")
            continue

        # 保存结果
        temperatures.append(T)
        D_list.append(D)
        sigma_list.append(sigma)

    # 绘制阿伦尼乌斯图（双拟合线）
    if len(temperatures) >= 2:
        print("\n===== 计算完成，绘制双拟合阿伦尼乌斯图 =====")
        arrhenius_plot(temperatures, D_list, sigma_list)
    else:
        print(f"\n===== 错误：有效数据点不足（仅{len(temperatures)}个），无法绘图 =====")