# uv (preferred) — all three publish binary wheels on PyPI, no conda needed
uv add openmm "mdanalysis>=2.10" pdbfixer
# If a pdbfixer wheel is unavailable for your platform, install from source:
uv pip install "pdbfixer @ git+https://github.com/openmm/pdbfixer.git"
 
# conda-forge alternative (pulls CUDA builds on Linux):
# conda install -c conda-forge openmm mdanalysis pdbfixer nglview

from openmm.app import *
from openmm import *
from openmm.unit import *
import sys
def prepare_system_from_pdb(pdb_file, forcefield_name="amber14-all.xml",
                              water_model="amber14/tip3pfb.xml"):
    """
    Prepare an OpenMM system from a PDB file.
    Args:
        pdb_file: Path to cleaned PDB file (use PDBFixer for raw PDB files)
        forcefield_name: Force field XML file
        water_model: Water model XML file
    Returns:
        pdb, forcefield, system, topology
    """
    # Load PDB
    pdb = PDBFile(pdb_file)
    # Load force field
    forcefield = ForceField(forcefield_name, water_model)
    # Add hydrogens and solvate
    modeller = Modeller(pdb.topology, pdb.positions)
    modeller.addHydrogens(forcefield)
    # Add solvent box (10 Å padding, 150 mM NaCl)
    modeller.addSolvent(
        forcefield,
        model='tip3p',
        padding=10*angstroms,
        ionicStrength=0.15*molar
    )
    print(f"System: {modeller.topology.getNumAtoms()} atoms, "
          f"{modeller.topology.getNumResidues()} residues")
    # Create system
    system = forcefield.createSystem(
        modeller.topology,
        nonbondedMethod=PME,         # Particle Mesh Ewald for long-range electrostatics
        nonbondedCutoff=1.0*nanometer,
        constraints=HBonds,           # Constrain hydrogen bonds (allows 2 fs timestep)
        rigidWater=True,
        ewaldErrorTolerance=0.0005
    )
    return modeller, system

from openmm.app import *
from openmm import *
from openmm.unit import *
def minimize_energy(modeller, system, output_pdb="minimized.pdb",
                     max_iterations=1000, tolerance=10.0):
    """
    Energy minimize the system to remove steric clashes.
    Args:
        modeller: Modeller object with topology and positions
        system: OpenMM System
        output_pdb: Path to save minimized structure
        max_iterations: Maximum minimization steps
        tolerance: Convergence criterion in kJ/mol/nm
    Returns:
        simulation object with minimized positions
    """
    # This integrator is reused by the NVT/NPT stages below, so use 2 fs: with only
    # HBonds constraints (no hydrogen mass repartitioning) a 4 fs step is unstable.
    integrator = LangevinMiddleIntegrator(300*kelvin, 1/picosecond, 0.002*picoseconds)
    # Create simulation
    # Use GPU if available (CUDA or OpenCL), fall back to CPU
    try:
        platform = Platform.getPlatformByName('CUDA')
        properties = {'DeviceIndex': '0', 'Precision': 'mixed'}
    except Exception:
        try:
            platform = Platform.getPlatformByName('OpenCL')
            properties = {}
        except Exception:
            platform = Platform.getPlatformByName('CPU')
            properties = {}
    simulation = Simulation(
        modeller.topology, system, integrator,
        platform, properties
    )
    simulation.context.setPositions(modeller.positions)
    # Check initial energy
    state = simulation.context.getState(getEnergy=True)
    print(f"Initial energy: {state.getPotentialEnergy()}")
    # Minimize
    simulation.minimizeEnergy(
        tolerance=tolerance*kilojoules_per_mole/nanometer,
        maxIterations=max_iterations
    )
    state = simulation.context.getState(getEnergy=True, getPositions=True)
    print(f"Minimized energy: {state.getPotentialEnergy()}")
    # Save minimized structure
    with open(output_pdb, 'w') as f:
        PDBFile.writeFile(simulation.topology, state.getPositions(), f)
    return simulation

from openmm.app import *
from openmm import *
from openmm.unit import *
def run_nvt_equilibration(simulation, n_steps=50000, temperature=300,
                            report_interval=1000, output_prefix="nvt"):
    """
    NVT equilibration: constant N, V, T.
    Equilibrate velocities to target temperature.
    Args:
        simulation: OpenMM Simulation (after minimization)
        n_steps: Number of MD steps (50000 × 2fs = 100 ps)
        temperature: Temperature in Kelvin
        report_interval: Steps between data reports
        output_prefix: File prefix for trajectory and log
    """
    # Add position restraints for backbone during NVT
    # (Optional: restraint heavy atoms)
    # Set temperature
    simulation.context.setVelocitiesToTemperature(temperature*kelvin)
    # Add reporters
    simulation.reporters = []
    # Log file
    simulation.reporters.append(
        StateDataReporter(
            f"{output_prefix}_log.txt",
            report_interval,
            step=True,
            potentialEnergy=True,
            kineticEnergy=True,
            temperature=True,
            volume=True,
            speed=True
        )
    )
    # DCD trajectory (compact binary format)
    simulation.reporters.append(
        DCDReporter(f"{output_prefix}_traj.dcd", report_interval)
    )
    print(f"Running NVT equilibration: {n_steps} steps ({n_steps*2/1000:.1f} ps)")
    simulation.step(n_steps)
    print("NVT equilibration complete")
    return simulation

def run_npt_production(simulation, n_steps=500000, temperature=300, pressure=1.0,
                        report_interval=5000, output_prefix="npt"):
    """
    NPT production run: constant N, P, T.
    Args:
        n_steps: Production steps (500000 × 2fs = 1 ns)
        temperature: Temperature in Kelvin
        pressure: Pressure in bar
        report_interval: Steps between reports
    """
    # Add Monte Carlo barostat for pressure control
    system = simulation.context.getSystem()
    system.addForce(MonteCarloBarostat(pressure*bar, temperature*kelvin, 25))
    simulation.context.reinitialize(preserveState=True)
    # Update reporters
    simulation.reporters = []
    simulation.reporters.append(
        StateDataReporter(
            f"{output_prefix}_log.txt",
            report_interval,
            step=True,
            potentialEnergy=True,
            temperature=True,
            density=True,
            speed=True
        )
    )
    simulation.reporters.append(
        DCDReporter(f"{output_prefix}_traj.dcd", report_interval)
    )
    # Save checkpoints
    simulation.reporters.append(
        CheckpointReporter(f"{output_prefix}_checkpoint.chk", 50000)
    )
    print(f"Running NPT production: {n_steps} steps ({n_steps*2/1000000:.2f} ns)")
    simulation.step(n_steps)
    print("Production MD complete")
    return simulation

import MDAnalysis as mda
from MDAnalysis.analysis import rms, align, contacts
import numpy as np
import matplotlib.pyplot as plt
def load_trajectory(topology_file, trajectory_file):
    """
    Load an MD trajectory with MDAnalysis.
    Args:
        topology_file: PDB, PSF, or other topology file
        trajectory_file: DCD, XTC, TRR, or other trajectory
    """
    u = mda.Universe(topology_file, trajectory_file)
    print(f"Universe: {u.atoms.n_atoms} atoms, {u.trajectory.n_frames} frames")
    print(f"Time range: 0 to {u.trajectory.totaltime:.0f} ps")
    return u

def compute_rmsd(u, selection="backbone", reference_frame=0):
    """
    Compute RMSD of selected atoms relative to reference frame.
    Args:
        u: MDAnalysis Universe
        selection: Atom selection string (MDAnalysis syntax)
        reference_frame: Frame index for reference structure
    Returns:
        numpy array of (time, rmsd) values
    """
    # Pre-align the trajectory (rewrites coordinates in memory) so that any
    # downstream per-atom analysis sees superimposed frames.
    aligner = align.AlignTraj(u, u, select=selection, in_memory=True)
    aligner.run()
    # rms.RMSD does its own optimal superposition against ref_frame, so the
    # RMSD numbers below do not depend on the AlignTraj step above.
    R = rms.RMSD(u, select=selection, ref_frame=reference_frame)
    R.run()
    rmsd_data = R.results.rmsd  # columns: frame, time, RMSD
    return rmsd_data
def plot_rmsd(rmsd_data, title="RMSD over time", output_file="rmsd.png"):
    """Plot RMSD over simulation time."""
    fig, ax = plt.subplots(figsize=(10, 4))
    ax.plot(rmsd_data[:, 1] / 1000, rmsd_data[:, 2], 'b-', linewidth=0.5)
    ax.set_xlabel("Time (ns)")
    ax.set_ylabel("RMSD (Å)")
    ax.set_title(title)
    ax.axhline(rmsd_data[:, 2].mean(), color='r', linestyle='--',
               label=f'Mean: {rmsd_data[:, 2].mean():.2f} Å')
    ax.legend()
    plt.tight_layout()
    plt.savefig(output_file, dpi=150)
    return fig

def compute_rmsf(u, selection="backbone", start_frame=0):
    """
    Compute per-residue RMSF (flexibility).
    Returns:
        resids, rmsf_values arrays
    """
    atoms = u.select_atoms(selection)
    # RMSF does no superposition of its own — align the trajectory to the average
    # structure first, otherwise the fluctuations include rigid-body motion.
    average = align.AverageStructure(u, u, select=selection, ref_frame=0).run()
    align.AlignTraj(u, average.results.universe, select=selection, in_memory=True).run()
    R = rms.RMSF(atoms).run(start=start_frame)
    # R.results.rmsf is indexed by the atom's position WITHIN `atoms` (0..n-1),
    # not by global atom index — use the group's per-atom residue indices to average.
    resindices = atoms.resindices
    unique_res = np.unique(resindices)
    rmsf_per_res = np.array([R.results.rmsf[resindices == ri].mean() for ri in unique_res])
    return u.residues[unique_res].resids, rmsf_per_res

from MDAnalysis.lib import distances as mda_distances
def analyze_contacts(u, protein_sel="protein", ligand_sel="resname LIG",
                      radius=4.5, start_frame=0):
    """
    Track which protein residues are in contact with the ligand per frame.
    Args:
        radius: Contact distance cutoff in Angstroms
    Returns:
        list of sets, one per analyzed frame, of contacting protein resids
    """
    protein = u.select_atoms(protein_sel)
    ligand = u.select_atoms(ligand_sel)
    contact_frames = []
    for ts in u.trajectory[start_frame:]:
        # Pairwise protein-ligand distance matrix (PBC-aware via box=u.dimensions),
        # then threshold into a boolean contact matrix.
        # NOTE: contacts.contact_matrix(d, radius) takes a *distance matrix* d,
        # not raw coordinates — passing positions directly is a common bug.
        dmat = mda_distances.distance_array(
            protein.positions, ligand.positions, box=u.dimensions
        )
        contact_mask = contacts.contact_matrix(dmat, radius)  # shape (n_prot, n_lig)
        # A protein atom is in contact if it is close to any ligand atom.
        contacting_atoms = protein.atoms[contact_mask.any(axis=1)]
        contact_frames.append(set(contacting_atoms.resids))
    return contact_frames

from pdbfixer import PDBFixer
from openmm.app import PDBFile
def fix_pdb(input_pdb, output_pdb, ph=7.0):
    """Fix common PDB issues: missing residues, atoms, add H, standardize."""
    fixer = PDBFixer(filename=input_pdb)
    fixer.findMissingResidues()
    fixer.findNonstandardResidues()
    fixer.replaceNonstandardResidues()
    fixer.removeHeterogens(True)    # Remove water/ligands
    fixer.findMissingAtoms()
    fixer.addMissingAtoms()
    fixer.addMissingHydrogens(ph)
    with open(output_pdb, 'w') as f:
        PDBFile.writeFile(fixer.topology, fixer.positions, f)
    return output_pdb

# conda/mamba install -c conda-forge openff-toolkit openff-interchange
# (the .offxml files live in openff-forcefields, which is conda-only — not on PyPI,
#  so a pip/uv-only install of openff-toolkit cannot load the Sage force field.)
# openff-2.2.1.offxml is OpenFF "Sage" — a SMIRNOFF force field, NOT GAFF2 (the newest
# Sage line is openff-2.3.0.offxml). For actual GAFF2, use AmberTools antechamber/ACPYPE.
from openff.toolkit import Molecule, ForceField as OFFForceField
def parameterize_ligand(smiles, ff_name="openff-2.2.1.offxml"):
    """Generate OpenFF (Sage) parameters for a small molecule as an Interchange."""
    mol = Molecule.from_smiles(smiles)
    mol.generate_conformers(n_conformers=1)
    mol.assign_partial_charges("am1bcc")  # AM1-BCC charges, OpenFF default
    off_ff = OFFForceField(ff_name)
    interchange = off_ff.create_interchange(mol.to_topology())
    return interchange  # -> interchange.to_openmm() for an OpenMM System