#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@file               :runner.py
@Description:       :
@Date               :2023/12/04 14:34:36
@Author             :lyzeng
@Email              :pylyzeng@gmail.com
@version            :1.0
'''

import time
import logging
from dataclasses import dataclass, field
from pathlib import Path
import subprocess
import shutil
import os

# 设置日志记录
logging.basicConfig(level=logging.INFO, filename='simulation_log.log', filemode='a',
                    format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger()

@dataclass
class SimulationRunner:
    pdb_file: Path
    nsteps: int
    dt: float
    base_folder: Path
    bash_script: Path = None  # Bash脚本路径作为可选参数
    gmxrc_path: Path = None  # GMXRC文件路径作为可选参数
    runner_folder: Path = field(init=False)

    def __post_init__(self):
        self.runner_folder = self.base_folder / f"runner_{self.pdb_file.stem}"
        self.runner_folder.mkdir(exist_ok=True)
        self.bash_script = self.bash_script.absolute() or Path(__file__).resolve().parent / "md_gromacs.sh"
        # 设置 GMXRC_PATH 环境变量
        self.gmxrc_path = self.gmxrc_path or Path("/home/lingyuzeng/software/gmx2023.2/bin/GMXRC")

    def copy_pdb(self):
        shutil.copy(self.pdb_file, self.runner_folder / self.pdb_file.name)

    def run_simulation(self):
        start_time = time.time()
        result = None
        try:
            env_vars = {
                "NAME": self.pdb_file.stem,  # 首先将 NAME 设置为文件的 stem
                "NSTEPS": str(self.nsteps),
                "DT": str(self.dt),
                "GMXRC_PATH": str(self.gmxrc_path)
            }
            env_vars["HOME"] = os.environ["HOME"]
            logger.info(f"pdb_file: {self.pdb_file.name}")
            logger.info(f"Executing script at: {self.bash_script}")
            result = subprocess.run(["bash", str(self.bash_script)], env=env_vars, cwd=self.runner_folder,
                                    capture_output=True, text=True, check=True)
        except subprocess.CalledProcessError as e:
            logger.error(f"Error in simulation for {self.pdb_file.name}: {e}")
            if e.stdout:
                logger.error(f"Standard Output:\n{e.stdout}")
            if e.stderr:
                logger.error(f"Standard Error:\n{e.stderr}")

        end_time = time.time()
        duration = end_time - start_time

        if result:
            logger.info(f"Simulation for {self.pdb_file.name} completed successfully in {duration:.2f} seconds.")
            if result.stdout:
                logger.info(f"Shell Script Output:\n{result.stdout}")
            if result.stderr:
                logger.error(f"Shell Script Error Output:\n{result.stderr}")
        else:
            logger.error(f"Simulation for {self.pdb_file.name} failed in {duration:.2f} seconds.")

def main(simulation_steps, time_step, pdb_folder_path, bash_script_path, gmxrc_path):
    pdb_folder = Path(pdb_folder_path).resolve()
    for pdb_file in pdb_folder.glob("*.pdb"):
        runner = SimulationRunner(pdb_file, simulation_steps, time_step, pdb_folder, bash_script_path, gmxrc_path)
        runner.copy_pdb()
        logger.info(f"Running simulation for {pdb_file.name} in {runner.runner_folder}...")
        runner.run_simulation()
        logger.info(f"Finished simulation for {pdb_file.name}.")

if __name__ == "__main__":
    NSTEPS = 50000000  # Example: 50000000 steps
    DT = 0.002  # Example: 2 fs time step
    PDB_FOLDER_PATH = Path("./pdb_files")  # Assuming the PDB files are in a folder named 'pdb_files' in the current directory
    # 传入自定义的bash脚本路径
    CUSTOM_BASH_SCRIPT_PATH = Path('md_gromacs.sh')
    # 传入 GMXRC 文件的路径
    GMXRC_PATH = Path('/home/zenglingyu/software/gmx2023.2/bin/GMXRC')
    main(NSTEPS, DT, PDB_FOLDER_PATH, CUSTOM_BASH_SCRIPT_PATH, GMXRC_PATH)