import time
import pandas as pd
from pathlib import Path
from qspoc.descriptor import QSPOCDesc


def run_benchmark(label, qspoc_obj, geo_opt_params=None, sp_params=None, desc_params=None):
    """
    运行完整的 QSPOC 流程并记录每一步的时间
    """
    print(f"\n>>> 开始测试配置: {label}")
    results = {"Config": label}

    # 1. 初始化结构计时
    start = time.perf_counter()
    qspoc_obj.get_init_sturct()
    results["init_struct_time(s)"] = round(time.perf_counter() - start, 2)
    print(f"[{label}] 初始结构生成完成: {results['init_struct_time(s)']}s")

    # 2. 几何优化计时
    start = time.perf_counter()
    geo_opt_params = geo_opt_params or {}
    qspoc_obj.geometric_opt(**geo_opt_params)
    results["geometric_opt_time(s)"] = round(time.perf_counter() - start, 2)
    print(f"[{label}] 几何优化完成: {results['geometric_opt_time(s)']}s")

    # 3. 单点能计算计时
    start = time.perf_counter()
    sp_params = sp_params or {}
    qspoc_obj.singlepoint_energy(**sp_params)
    results["singlepoint_energy_time(s)"] = round(time.perf_counter() - start, 2)
    print(f"[{label}] 单点能计算完成: {results['singlepoint_energy_time(s)']}s")

    # 4. 描述符计算计时
    start = time.perf_counter()
    desc_params = desc_params or {}
    qspoc_obj.descriptor_calc(**desc_params)
    results["descriptor_calc_time(s)"] = round(time.perf_counter() - start, 2)
    print(f"[{label}] 描述符计算完成: {results['descriptor_calc_time(s)']}s")

    # 总耗时
    results["total_time(s)"] = round(
        results["init_struct_time(s)"]
        + results["geometric_opt_time(s)"]
        + results["singlepoint_energy_time(s)"]
        + results["descriptor_calc_time(s)"],
        2,
    )
    return results


# --- 路径设置 ---
BASE_DIR = Path(__file__).parent
Multiwfn_path = BASE_DIR.parent / "bin/Multiwfn/Multiwfn_3.8_bin_Linux_noGUI/Multiwfn_noGUI"
smiles_file = BASE_DIR / "test_file/speed_test_smiles.csv"
output_dir = BASE_DIR / "test_file/speed_test"
output_dir.mkdir(parents=True, exist_ok=True)

all_stats = []

# --- 测试 1: Normal (xtb only) ---
qspoc1 = QSPOCDesc(multiwfn_path=Multiwfn_path, save_dir=str(output_dir / "1"), exe_path_dict={"xtb": "xtb"}, acc="normal")
qspoc1.load_data_from_file(file_path=smiles_file, smiles_tag="SMILES")
stats1 = run_benchmark("Normal_Acc_xtb", qspoc1)
all_stats.append(stats1)

# --- 测试 2: High (xtb opt + g16 sp) ---
qspoc2 = QSPOCDesc(multiwfn_path=Multiwfn_path, save_dir=str(output_dir / "2"), exe_path_dict={"xtb": "xtb", "g16": "g16"}, acc="high")
qspoc2.load_data_from_file(file_path=smiles_file, smiles_tag="SMILES")
stats2 = run_benchmark(
    "High_Acc_mixed",
    qspoc2,
    geo_opt_params={"method": "xtb"},
    sp_params={"method": "g16", "calc_command": "b3lyp/def2svp em=gd3bj"},
    desc_params={"cdft_command": "b3lyp/def2tzvp"},
)
all_stats.append(stats2)

# --- 测试 3: Ultra (g16 only) ---
qspoc3 = QSPOCDesc(multiwfn_path=Multiwfn_path, save_dir=str(output_dir / "3"), exe_path_dict={"xtb": "xtb", "g16": "g16"}, acc="ultra")
qspoc3.load_data_from_file(file_path=smiles_file, smiles_tag="SMILES")
stats3 = run_benchmark(
    "Ultra_Acc_g16",
    qspoc3,
    geo_opt_params={"method": "g16", "calc_command": "b3lyp/def2svp"},
    sp_params={"method": "g16", "calc_command": "wb97xd/def2tzvp"},
    desc_params={"cdft_command": "wb97xd/def2tzvp"},
)
all_stats.append(stats3)

# --- 保存统计结果 ---
df_stats = pd.DataFrame(all_stats)
summary_path = output_dir / "speed_test_summary.csv"
df_stats.to_csv(summary_path, index=False)

print(f"\n{'='*30}")
print(f"所有测试已完成！")
print(f"计时统计结果已保存至: {summary_path}")
print(df_stats)
