import numpy as np
import matplotlib.pyplot as plt

import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import linregress


def plot_and_get_final_msd(file_path, ion_name='Ion',timestep_ps = 0.2):
    """
    读取LAMMPS输出的MSD数据文件，绘制MSD-时间图，并返回最后一个MSD值。

    参数:
    file_path (str): msd_*.dat文件的路径。
    ion_name (str, optional): 离子的名称，用于图表标题和标签。默认为'Ion'。

    返回:
    float: 文件中记录的最后一个MSD值。
    """
    # LAMMPS输出的MSD文件通常会跳过前几行注释，使用np.loadtxt可以轻松处理
    try:
        # 尝试直接读取，通常第一列是Timestep，第四列是总MSD
        # 文件格式: Timestep, Time, MSD_x, MSD_y, MSD_z, MSD_total
        # 但fix ave/time的输出是: Timestep, c_msd[1], c_msd[2], c_msd[3], c_msd[4]
        # c_msd[4] 就是总MSD
        data = np.loadtxt(file_path, comments='#')
        timesteps = data[:, 0]
        msd_values = data[:, -1]  # 最后一列通常是总MSD
    except Exception as e:
        print(f"读取文件时出错: {e}")
        print("请检查文件格式是否正确。")
        return None

    # 从in.lmp文件我们知道时间步长是0.001 ps

    time_ps = timesteps * timestep_ps

    # --- 绘图 ---
    plt.figure(figsize=(10, 6))
    plt.plot(time_ps/1000, msd_values, label=f'{ion_name} MSD')

    plt.xlabel('Time (ps)', fontsize=14)
    plt.ylabel('MSD (Ų)', fontsize=14)
    plt.title(f'Mean Squared Displacement (MSD) of {ion_name}', fontsize=16)
    plt.grid(True)
    plt.legend()
    plt.show()

    # --- 输出并返回最终值 ---
    final_msd_value = msd_values[-1]
    final_time = time_ps[-1]

    print(f"文件路径: {file_path}")
    print(f"总模拟时间: {final_time/1000:.2f} ps")
    print(f"最终MSD值: {final_msd_value:.4f} Ų")

    return final_msd_value


def calculate_conductivity_from_msd(
        msd_file_path,
        log_file_path,
        ion_name,
        charge,
        num_ions,
        fit_start_fraction=0.5,
        fit_end_fraction=0.5
):
    """
    从MSD数据计算电导率 (v2)。
    自动从 log.lammps 文件提取温度和体积。

    参数:
    msd_file_path (str): msd_*.dat 文件的路径。
    log_file_path (str): log.lammps 文件的路径。
    ion_name (str): 离子名称 (例如 'Li⁺')。
    charge (int): 离子的电荷数 (例如 Li⁺ 为 1)。
    num_ions (int): 模拟盒子中该离子的总数量。
    fit_fraction (float): 用于线性拟合的数据比例 (使用最后部分)。

    返回:
    dict: 包含计算结果的字典。
    """
    # --- 1. 从 log 文件自动提取参数 ---
    sim_params = extract_params_from_log(log_file_path)
    if not sim_params or not sim_params["avg_volume"] or not sim_params["avg_temp"]:
        print("无法从 log 文件获取必要参数，计算中止。")
        return None

    temp_K = sim_params["avg_temp"]
    volume_A3 = sim_params["avg_volume"]

    # --- 2. 读取和处理 MSD 数据 ---
    try:
        data = np.loadtxt(msd_file_path, comments='#')
        timesteps = data[:, 0]
        msd_values = data[:, -1]  # 总MSD (Ų)
    except Exception as e:
        print(f"读取文件 '{msd_file_path}' 时出错: {e}")
        return None

    timestep_ps = 0.001  # 时间步长 (ps)
    time_ps = timesteps * timestep_ps

    # --- 3. 线性拟合计算扩散系数 ---
    fit_start_index = int(len(time_ps) * fit_start_fraction)
    fit_end_index = int(len(time_ps) * fit_end_fraction)
    fit_time_ps = time_ps[fit_start_index:fit_end_index]
    fit_msd_values = msd_values[fit_start_index:fit_end_index]

    if len(fit_time_ps) < 2:
        print("错误: 用于拟合的数据点不足 (少于2个)，无法进行线性回归。")
        return None

    slope, intercept, r_value, _, _ = linregress(fit_time_ps, fit_msd_values)

    diff_coeff_A2_ps = slope / 6
    diff_coeff_cm2_s = diff_coeff_A2_ps * 1e-4
    diff_coeff_m2_s = diff_coeff_A2_ps * 1e-8

    # --- 4. 计算电导率 (Nernst-Einstein) ---
    e_charge = 1.60217663e-19  # (C)
    kB = 1.380649e-23  # (J/K)

    q = charge * e_charge
    n = num_ions / (volume_A3 * 1e-30)  # (m⁻³)

    conductivity_S_m = (n * q ** 2 * diff_coeff_m2_s) / (kB * temp_K)

    # --- 5. 可视化 ---
    plt.figure(figsize=(12, 7))
    plt.plot(time_ps, msd_values, 'b.', markersize=4, label='MSD Data')
    fit_line = slope * fit_time_ps + intercept
    plt.plot(fit_time_ps, fit_line, 'r-', linewidth=3, label=f'Linear Fit (R² = {r_value ** 2:.4f})')

    plt.xlabel('Time (ps)', fontsize=14)
    plt.ylabel('MSD (Ų)', fontsize=14)
    plt.title(f'MSD Analysis for {ion_name}', fontsize=16)
    plt.legend(fontsize=12)
    plt.grid(True)

    info_text = (
        f"Diffusion Coefficient (D): {diff_coeff_cm2_s:.2e} cm²/s\n"
        f"Conductivity (σ): {conductivity_S_m:.4f} S/m\n"
        f"Avg Temp: {temp_K:.2f} K\n"
        f"Avg Volume: {volume_A3:.2f} ų"
    )
    plt.text(0.05, 0.95, info_text, transform=plt.gca().transAxes, fontsize=12,
             verticalalignment='top', bbox=dict(boxstyle='round,pad=0.5', fc='wheat', alpha=0.5))

    plt.show()

    # --- 6. 返回结果 ---
    results = {
        "ion": ion_name,
        "diffusion_coefficient_cm2/s": diff_coeff_cm2_s,
        "conductivity_S/m": conductivity_S_m,
        "fit_R2": r_value ** 2,
        "avg_temp_K": temp_K,
        "avg_volume_A3": volume_A3
    }

    print("--- 最终计算结果 ---")
    for key, value in results.items():
        print(f"{key}: {value}")

    return results



def extract_params_from_log(log_file_path):
    """
    从 log.lammps 文件中提取 NVT 生产阶段的平均参数。

    参数:
    log_file_path (str): log.lammps 文件的路径。

    返回:
    dict: 包含提取参数的字典 (avg_temp, avg_volume, total_atoms)，如果失败则返回 None。
    """
    try:
        with open(log_file_path, 'r',encoding='utf-8',errors='ignore') as f:
            lines = f.readlines()
    except FileNotFoundError:
        print(f"错误: log 文件 '{log_file_path}' 未找到。")
        return None

    # 初始化变量
    params = {
        "avg_temp": None,
        "avg_volume": None,
        "total_atoms": None
    }

    # --- 1. 提取总原子数 ---
    # 通常在文件开头附近
    for line in lines:
        if "atoms" in line and "atom types" in line:
            parts = line.split()
            params["total_atoms"] = int(parts[0])
            break  # 找到后就停止

    # --- 2. 定位并提取 NVT 生产阶段的热力学数据 ---
    # 我们通过寻找 "run 2000000" 来定位生产阶段的开始
    # 这是根据您提供的 in.lmp 文件
    nvt_data_lines = []
    in_nvt_section = False

    for i, line in enumerate(lines):
        # 寻找 NVT 生产阶段的开始标志
        if "fix           nvt_prod all nvt" in line:
            # 找到 fix 命令后，热力学数据头通常在几行之后
            # 我们从下一行的 "run" 命令开始找
            if "run           2000000" in lines[i + 1]:
                # 再往下找 "Step Time Temp..." 这样的表头
                for j in range(i + 2, len(lines)):
                    if lines[j].strip().startswith("Step"):
                        header = lines[j].split()
                        in_nvt_section = True
                        break
                if in_nvt_section:
                    continue  # 跳转到下一个循环，开始读取数据

        # 如果在 NVT 区域，并且行看起来像数据行，则读取
        if in_nvt_section:
            # 数据行以数字开头，并且列数与表头匹配
            parts = line.strip().split()
            if parts and parts[0].isdigit() and len(parts) == len(header):
                nvt_data_lines.append(parts)
            # "Loop time" 标志着一个 run 命令的结束
            elif line.strip().startswith("Loop time"):
                in_nvt_section = False
                break  # 生产阶段数据读取完毕

    if not nvt_data_lines:
        print("警告: 未能在 log 文件中找到 NVT 生产阶段的热力学数据。")
        print("请检查 in.lmp 中的 run 命令是否为 'run 2000000'。")
        return params  # 即使没找到，也可能返回原子数

    # --- 3. 计算平均值 ---
    # 将数据转换为 numpy 数组以便于计算
    try:
        nvt_data = np.array(nvt_data_lines).astype(float)
    except ValueError as e:
        print("错误：在将日志数据转换为数字时失败。")
        print("这通常意味着数据行中混入了非数字字符。")
        print(f"原始错误信息: {e}")
        # 可以打印出有问题的数据行来帮助调试
        # for i, row in enumerate(nvt_data_lines):
        #     try:
        #         [float(x) for x in row]
        #     except ValueError:
        #         print(f"问题行号 {i}: {row}")
        #         break
        return params
    # 根据表头找到 Temp 和 Volume 所在的列索引
    try:
        temp_col = header.index("Temp")
        vol_col = header.index("Volume")
    except ValueError as e:
        print(f"错误: 在日志表头中找不到列: {e}")
        return params
    print(nvt_data[1:5,temp_col])
    print(nvt_data[1:5,vol_col])
    params["avg_temp"] = np.mean(nvt_data[1:20, temp_col])
    params["avg_volume"] = np.mean(nvt_data[1:20, vol_col])

    print("--- 从 log.lammps 提取的参数 ---")
    print(f"总原子数: {params['total_atoms']}")
    print(f"NVT阶段平均温度: {params['avg_temp']:.2f} K")
    print(f"NVT阶段平均体积: {params['avg_volume']:.4f} ų")
    print("------------------------------------")

    return params