#!/usr/bin/env python3
"""
To Run Command

Generate production run scripts for adapter minimize output using OpenMM or AMBER.
"""

import sys
import shutil
from pathlib import Path

import click

GROMACS_INCLUDE_DIR = '/usr/local/gromacs/share/gromacs/top'


def get_system_name_from_yaml(yaml_path: str) -> str:
    """Read system_name from YAML config file."""
    import yaml
    path = Path(yaml_path)
    if not path.exists():
        raise FileNotFoundError(f"YAML file not found: {yaml_path}")

    with open(path, 'r') as f:
        config = yaml.safe_load(f)

    if 'system_name' not in config:
        raise KeyError(f"'system_name' field not found in YAML: {yaml_path}")

    return config['system_name']


def detect_gro_file(minimize_dir: Path) -> tuple[str, str]:
    """Detect which coordinate file to use."""
    gro_file = minimize_dir / "minimize_final_solvated.gro"
    pdb_file = minimize_dir / "minimize_final.pdb"

    if gro_file.exists():
        return ("minimize_final_solvated.gro", str(gro_file))
    elif pdb_file.exists():
        return ("minimize_final.pdb", str(pdb_file))
    else:
        raise FileNotFoundError(
            f"Neither minimize_final_solvated.gro nor minimize_final.pdb found in {minimize_dir}"
        )


def setup_production_folder(
    minimize_dir: Path,
    production_dir: Path,
    gro_name: str,
    gro_path: str,
) -> None:
    """Setup production folder with appropriate files."""
    production_dir.mkdir(parents=True, exist_ok=True)
    click.echo(f"Created production folder: {production_dir}")

    # Copy coordinate file
    shutil.copy2(gro_path, production_dir / gro_name)
    click.echo(f"Copied: {gro_name}")

    # Copy topology file
    top_file = minimize_dir / "topol.top"
    if top_file.exists():
        shutil.copy2(top_file, production_dir / "topol.top")
        click.echo(f"Copied: topol.top")
    else:
        click.echo(f"Warning: topol.top not found in {minimize_dir}")

    # Copy .ff folders (force field directories)
    for item in minimize_dir.iterdir():
        if item.is_dir() and item.name.endswith('.ff'):
            dest = production_dir / item.name
            shutil.copytree(item, dest, dirs_exist_ok=True)
            print(f"Copied: {item.name}/")


def generate_openmm_script(
    production_dir: Path,
    gro_file: str,
) -> None:
    """Generate OpenMM run script in production folder using default parameters."""
    # Default parameters
    TOTAL_TIME_NS = 100.0
    TIMESTEP_FS = 2.0
    TEMPERATURE_K = 300.0
    RUN_NPT = True
    PRESSURE_BAR = 1.0
    CUTOFF_NM = 1.2
    TRAJ_INTERVAL_PS = 100.0  # 0.1 ns
    LOG_INTERVAL_PS = 100.0
    DO_MINIMIZE = True
    MIN_TOLERANCE = 1.0
    MIN_MAX_STEPS = 5000

    # Calculate total steps: time(ns) * 1e6 / timestep(fs)
    n_steps = int(TOTAL_TIME_NS * 1_000_000 / TIMESTEP_FS)

    # Convert intervals from ns to steps
    traj_interval_steps = int(TRAJ_INTERVAL_PS * 1000 / TIMESTEP_FS)
    log_interval_steps = int(LOG_INTERVAL_PS * 1000 / TIMESTEP_FS)

    script_content = f'''#!/usr/bin/env python3
"""
OpenMM production run script for {production_dir.name}

Auto-generated by adapter to_run
Default parameters: 100ns, 2fs, 300K, NPT, 1.2nm cutoff
"""

# =============================================================================
# DEFAULT CONFIGURATION
# =============================================================================

# Input files
GRO_FILE = '{gro_file}'
TOP_FILE = 'topol.top'
GROMACS_INCLUDE_DIR = '{GROMACS_INCLUDE_DIR}'

# Simulation parameters
TIMESTEP_FS = {TIMESTEP_FS}          # Timestep in femtoseconds
TEMPERATURE_K = {TEMPERATURE_K}      # Temperature in Kelvin
N_STEPS = {n_steps}          # Total number of steps ({TOTAL_TIME_NS} ns)

# Ensemble type: True for NPT, False for NVT
RUN_NPT = {str(RUN_NPT)}
PRESSURE_BAR = {PRESSURE_BAR}         # Pressure in bar (only used if RUN_NPT=True)
BAROSTAT_INTERVAL = 25    # Barostat frequency

# Nonbonded parameters
CUTOFF_NM = {CUTOFF_NM}           # Nonbonded cutoff in nm

# Output settings
TRAJ_FILE = 'trajectory.xtc'
LOG_FILE = 'simulation.log'
REPORT_INTERVAL = {traj_interval_steps}    # Steps between trajectory writes
LOG_INTERVAL = {log_interval_steps}       # Steps between log output

# Energy minimization
DO_MINIMIZE = {str(DO_MINIMIZE)}
MIN_TOLERANCE = {MIN_TOLERANCE}       # Minimization tolerance
MIN_MAX_STEPS = {MIN_MAX_STEPS}      # Maximum minimization steps

# =============================================================================

import sys

from openmm import app, unit, LangevinMiddleIntegrator, MonteCarloBarostat
from openmm.app import Simulation, GromacsGroFile, GromacsTopFile, PME, HBonds

# Try to import XTC Reporter (available in OpenMM 7.9+)
try:
    from openmm.app import XTCReporter
    HAS_XTC = True
except ImportError:
    from openmm.app import DCDReporter as XTCReporter
    HAS_XTC = False

from openmm.app import StateDataReporter


def create_system(top):
    system = top.createSystem(
        nonbondedMethod=PME,
        nonbondedCutoff=CUTOFF_NM * unit.nanometer,
        constraints=HBonds,
        rigidWater=True,
        removeCMMotion=True
    )
    return system


def create_integrator():
    temperature = TEMPERATURE_K * unit.kelvin
    timestep = TIMESTEP_FS * unit.femtoseconds

    friction = 1.0 / unit.picosecond
    integrator = LangevinMiddleIntegrator(temperature, friction, timestep)
    return integrator


def add_barostat(system):
    if RUN_NPT:
        temperature = TEMPERATURE_K * unit.kelvin
        pressure = PRESSURE_BAR * unit.bar
        barostat = MonteCarloBarostat(pressure, temperature, BAROSTAT_INTERVAL)
        system.addForce(barostat)
        return True
    return False


def main():
    print(f"Loading input files: {{GRO_FILE}}, {{TOP_FILE}}")
    gro = GromacsGroFile(GRO_FILE)
    top = GromacsTopFile(
        TOP_FILE,
        periodicBoxVectors=gro.getPeriodicBoxVectors(),
        includeDir=GROMACS_INCLUDE_DIR
    )

    print("Creating system...")
    system = create_system(top)

    is_npt = add_barostat(system)
    ensemble = "NPT" if is_npt else "NVT"
    print(f"Running {{ensemble}} ensemble at {{TEMPERATURE_K}} K")

    print("Creating integrator...")
    integrator = create_integrator()

    simulation = Simulation(top.topology, system, integrator)
    simulation.context.setPositions(gro.positions)

    if gro.getPeriodicBoxVectors() is not None:
        simulation.context.setPeriodicBoxVectors(*gro.getPeriodicBoxVectors())

    if DO_MINIMIZE:
        print(f"Minimizing energy (tolerance={{MIN_TOLERANCE}}, max_steps={{MIN_MAX_STEPS}})...")
        simulation.minimizeEnergy(tolerance=MIN_TOLERANCE, maxIterations=MIN_MAX_STEPS)

    print(f"Using trajectory format: {{'XTC' if HAS_XTC else 'DCD'}}")
    simulation.reporters.append(
        XTCReporter(TRAJ_FILE, REPORT_INTERVAL)
    )

    with open(LOG_FILE, 'w') as f:
        simulation.reporters.append(
            StateDataReporter(
                f,
                LOG_INTERVAL,
                step=True,
                time=True,
                potentialEnergy=True,
                kineticEnergy=True,
                totalEnergy=True,
                temperature=True,
                volume=True if is_npt else False,
                density=True if is_npt else False
            )
        )

    simulation.reporters.append(
        StateDataReporter(
            sys.stdout,
            LOG_INTERVAL,
            step=True,
            potentialEnergy=True,
            temperature=True,
            volume=True if is_npt else False
        )
    )

    print(f"Starting simulation: {{N_STEPS}} steps (dt={{TIMESTEP_FS}} fs)")
    simulation.step(N_STEPS)
    print("Simulation complete!")


if __name__ == '__main__':
    main()
'''

    script_path = production_dir / "openmm_run.py"
    with open(script_path, 'w') as f:
        f.write(script_content)

    click.echo(f"Generated: openmm_run.py")
    click.echo(f"  - Total time: {TOTAL_TIME_NS} ns")
    click.echo(f"  - Trajectory interval: {TRAJ_INTERVAL_PS} ps ({traj_interval_steps} steps)")
    click.echo(f"  - Log interval: {LOG_INTERVAL_PS} ps ({log_interval_steps} steps)")


def generate_amber_script(
    production_dir: Path,
    gro_file: str,
) -> None:
    """Generate AMBER input files using ParmEd."""
    try:
        import parmed
    except ImportError:
        raise ImportError(
            "ParmEd is required for AMBER conversion. "
            "Please install it: mamba install -c conda-forge parmed"
        )

    # Load Gromacs files using ParmEd
    top_file = production_dir / "topol.top"
    gro_path = production_dir / gro_file

    if not top_file.exists():
        raise FileNotFoundError(f"topol.top not found in {production_dir}")

    click.echo(f"Loading Gromacs files: {gro_file}, topol.top")

    # Load structure using ParmEd
    structure = parmed.load_file(str(top_file), xyz=str(gro_path))

    # Save AMBER files
    prmtop_path = production_dir / "system.prmtop"
    inpcrd_path = production_dir / "system.inpcrd"

    structure.save(str(prmtop_path), overwrite=True)
    structure.save(str(inpcrd_path), overwrite=True)

    click.echo(f"Generated: system.prmtop")
    click.echo(f"Generated: system.inpcrd")

    click.echo(f"\n  Note: No .mdin file is generated.")
    click.echo(f"  Please provide your own AMBER input file to run the simulation.")
    click.echo(f"  Example: pmemd -O -i run.npt.mdin -c system.inpcrd -p system.prmtop -r restart.ncrst -x trajectory.nc")


@click.command('to_run')
@click.option(
    '--input-file', '-f',
    type=click.Path(exists=True),
    required=True,
    help='Configuration YAML (same as CG stage)',
)
@click.option(
    '--engine', '-e',
    type=click.Choice(['openmm', 'amber']),
    default='openmm',
    help='Simulation engine: openmm or amber',
)
@click.option(
    '--input', '-i',
    type=click.Path(),
    required=False,
    default=None,
    help='Input: minimize output directory. Default: {system_name}_minimize',
)
def to_run_command(
    input_file,
    engine,
    input,
):
    """\b
    Generate production run scripts for adapter minimize output.

    \b
    This command reads the minimize output directory and generates
    appropriate run scripts for molecular dynamics simulation.

    \b
    Supported engines:
        - openmm: Generate OpenMM Python script
        - amber:  Generate AMBER prmtop/inpcrd files (user must provide .mdin)

    \b
    Examples:
        adapter to_run -f config.yaml                    # Generate OpenMM script
        adapter to_run -f config.yaml -e amber           # Generate AMBER files
        adapter to_run -f config.yaml -i custom_minimize # Custom input dir
    """
    from ...minimize.config_loader import load_config_from_yaml

    # Get system_name from YAML
    system_name, _ = load_config_from_yaml(input_file)

    # Determine minimize directory
    if input is None:
        input = f"{system_name}_minimize"

    minimize_dir = Path(input)
    if not minimize_dir.exists():
        raise click.BadParameter(
            f'Minimize directory "{input}" does not exist. '
            f'Run adapter minimize first or provide correct --input path.'
        )

    # Determine output directory
    production_dir = Path(f"{system_name}_production")

    click.echo(f"\n{'=' * 60}")
    click.echo(f"Setup Production Run (engine: {engine})")
    click.echo(f"{'=' * 60}")

    click.echo(f"\n  System: {system_name}")
    click.echo(f"  Input: {minimize_dir}")
    click.echo(f"  Output: {production_dir}")
    click.echo(f"  Engine: {engine}")

    # Detect coordinate file
    try:
        gro_name, gro_path = detect_gro_file(minimize_dir)
    except FileNotFoundError as e:
        raise click.BadParameter(str(e))

    click.echo(f"  Coordinate file: {gro_name}")

    # Setup production folder
    setup_production_folder(minimize_dir, production_dir, gro_name, gro_path)

    # Generate run script based on engine
    if engine == 'openmm':
        generate_openmm_script(
            production_dir=production_dir,
            gro_file=gro_name,
        )
    elif engine == 'amber':
        generate_amber_script(
            production_dir=production_dir,
            gro_file=gro_name,
        )

    click.echo(f"\n  Success! Production folder ready at:")
    click.echo(f"    {production_dir}")
    click.echo(f"\n  To run the simulation:")
    click.echo(f"    cd {production_dir}")
    if engine == 'openmm':
        click.echo(f"    python openmm_run.py")
    elif engine == 'amber':
        click.echo(f"    # Please provide your own run.npt.mdin file")
        click.echo(f"    pmemd -O -i run.npt.mdin -c system.inpcrd -p system.prmtop -r restart.ncrst -x trajectory.nc")
    click.echo()
