#!/usr/bin/env python
"""
Explicit-solvent MD of a free PROTAC (no POI, no E3 protein), following the
protocol of "Linker-Dependent Folding Rationalizes PROTAC Cell Permeability"
(J. Med. Chem. 2022, DOI 10.1021/acs.jmedchem.2c00877): triplicate ~100 ns runs in
chloroform (membrane-interior mimic) and/or water, analyzed for Rg, 3D PSA,
intramolecular H-bonds and conformational clusters.

Pipeline:
  1. RDKit: SMILES -> add Hs -> ETKDGv3 multi-conformer -> MMFF -> lowest-E conf.
  2. OpenFF Sage (openff-2.2.0) + AM1-BCC charges via AmberTools for the PROTAC
     (and for chloroform, when used) through openmmforcefields.
  3. Solvate:
       water      -> OpenMM addSolvent (TIP3P, 0.15 M NaCl, neutralized)
       chloroform -> packmol via openff-interchange pack_box
  4. Minimize once, then run N independent replicas (different velocity seeds),
     each: NVT heat 50->T, NPT equilibration, NPT production.

Outputs per replica land in <out>/<solvent>/rep<k>/; analyze them with
analyze_conformations.py.

Environment: conda env "protacmd" (environment_md.yml). Needs setuptools<81:
openff.amber_ff_ports imports pkg_resources, which setuptools 81 dropped, and
the run dies inside SystemGenerator without it.

Usage:
    python protac_md.py --solvent chloroform --replicas 3 --prod-ns 100 --gpu 0
    python protac_md.py --solvent water      --replicas 3 --prod-ns 100 --gpu 0
    python protac_md.py --prep-only                       # just the 3D conformer
"""

import argparse
import os
import sys
import time

# Top-ranked generated PROTACs for EGFR-VHL, BRD9-VHL and CDK6-CRBN. Any other
# SMILES can be simulated with --smiles.
CANDIDATES = {
    "egfrvhl_rank4": "CCCC(C(=O)N[C@H](C(=O)N1C[C@H](O)C[C@H]1C(=O)NCc1ccc(-c2scnc2C)cc1)"
                     "C(C)(C)C)C(=O)N(c1ccc(F)c(Cl)c1)c1ncnc2cc(OC)c(OCCCCN3CCNCC3)cc12",
    "egfrvhl_rank5": "CCNC(=O)C(OC(=O)C(=O)N[C@H](C(=O)N1C[C@H](O)C[C@H]1C(=O)NCc1ccc(-c2scnc2C)cc1)"
                     "C(C)(C)C)C(=O)N(c1ccc(F)c(Cl)c1)c1ncnc2cc(OC)c(OCCCCN3CCNCC3)cc12",
    "brd9vhl_rank1": "CCC(C)C(C(=O)N[C@H](C(=O)N1C[C@H](O)C[C@H]1C(=O)NCc1ccc(-c2scnc2C)cc1)C(C)(C)C)"
                     "C(=O)N1CCN(Cc2c(OC)cc(-c3cn(C)c(=O)c4cnccc34)cc2OC)CC1",
    "cdk6crbn_rank2": "CC(=O)c1c(C)c2cnc(N(C(=O)CC(C)CC(=O)Nc3cccc4c3C(=O)N(C3CCC(=O)NC3=O)C4=O)"
                      "c3ccc(N4CCNCC4)nc3)nc2n(C2CCCC2)c1=O",
}
DEFAULT_CANDIDATE = "egfrvhl_rank4"
CHLOROFORM_SMILES = "C(Cl)(Cl)Cl"
CHLOROFORM_DENSITY_PER_NM3 = 7.51   # molecules/nm^3 at 1.49 g/mL, 300 K


# --- 3D conformer -----------------------------------------------------------
def build_conformer(smiles, name, out_dir, n_confs=20, seed=0xF00D):
    from rdkit import Chem
    from rdkit.Chem import AllChem, Descriptors, rdMolDescriptors

    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        sys.exit(f"[fatal] RDKit could not parse SMILES: {smiles}")
    print(f"[prep] {name}: MW={Descriptors.MolWt(mol):.1f}  "
          f"charge={Chem.GetFormalCharge(mol):+d}  "
          f"rot.bonds={rdMolDescriptors.CalcNumRotatableBonds(mol)}")

    molh = Chem.AddHs(mol)
    params = AllChem.ETKDGv3()
    params.randomSeed = seed
    params.useRandomCoords = True
    cids = AllChem.EmbedMultipleConfs(molh, numConfs=n_confs, params=params)
    if len(cids) == 0:
        if AllChem.EmbedMolecule(molh, params) != 0:
            sys.exit("[fatal] 3D embedding failed.")
        cids = [0]
    if not AllChem.MMFFHasAllMoleculeParams(molh):
        print("[warn] MMFF missing some params; conformer quality may be lower.")
    results = AllChem.MMFFOptimizeMoleculeConfs(molh, maxIters=2000)
    best_cid, best_e = min(zip(cids, [e for _, e in results]), key=lambda x: x[1])
    print(f"[prep] lowest MMFF conformer {best_cid} (E={best_e:.1f} kcal/mol)")

    sdf = os.path.join(out_dir, f"{name}.sdf")
    w = Chem.SDWriter(sdf); molh.SetProp("_Name", name); w.write(molh, confId=best_cid); w.close()
    print(f"[prep] wrote {sdf}")
    keep = Chem.Mol(molh); keep.RemoveAllConformers()
    keep.AddConformer(molh.GetConformer(best_cid), assignId=True)
    return keep


# --- solvation --------------------------------------------------------------
def solvate_water(molecule, system_generator, cfg):
    from openmm import app, unit
    off_top = molecule.to_topology()
    modeller = app.Modeller(off_top.to_openmm(), molecule.conformers[0].to_openmm())
    modeller.addSolvent(system_generator.forcefield, model="tip3p",
                        padding=cfg.padding * unit.nanometer,
                        ionicStrength=cfg.ionic * unit.molar, neutralize=True)
    return modeller.topology, modeller.positions


def solvate_chloroform(molecule, chcl3, cfg):
    import numpy as np
    from openff.interchange.components._packmol import pack_box
    from openff.units import unit as ou
    from openff.units.openmm import to_openmm

    conf = molecule.conformers[0].m_as(ou.nanometer)
    extent = np.linalg.norm(conf - conf.mean(0), axis=1).max() * 2.0   # ~diameter, nm
    box = extent + 2.0 * cfg.padding
    n = int(round(0.90 * CHLOROFORM_DENSITY_PER_NM3 * box ** 3))       # under-fill; NPT compresses
    print(f"[md] chloroform box {box:.2f} nm, packing {n} CHCl3 (NPT will set final density)")
    box_vecs = np.eye(3) * box * ou.nanometer
    packed = pack_box(molecules=[chcl3], number_of_copies=[n],
                      solute=molecule.to_topology(), box_vectors=box_vecs,
                      center_solute=True, tolerance=2.0 * ou.angstrom)
    omm_top = packed.to_openmm()
    omm_top.setPeriodicBoxVectors(to_openmm(packed.box_vectors))
    positions = to_openmm(packed.get_positions())
    return omm_top, positions


# --- solvated system, minimized once ----------------------------------------
def build_system(rdkit_mol, name, out_dir, cfg):
    from openmm import app, unit, LangevinMiddleIntegrator, Platform
    from openff.toolkit import Molecule
    from openmmforcefields.generators import SystemGenerator

    molecule = Molecule.from_rdkit(rdkit_mol, allow_undefined_stereo=True)
    molecule.name = name

    extra = []
    chcl3 = None
    if cfg.solvent == "chloroform":
        chcl3 = Molecule.from_smiles(CHLOROFORM_SMILES)
        chcl3.generate_conformers(n_conformers=1)
        chcl3.name = "CL3"
        extra = [chcl3]

    forcefield_kwargs = {"constraints": app.HBonds, "rigidWater": True,
                         "removeCMMotion": True, "hydrogenMass": cfg.hmass * unit.amu}
    periodic_kwargs = {"nonbondedMethod": app.PME, "nonbondedCutoff": 1.0 * unit.nanometer}
    system_generator = SystemGenerator(
        forcefields=["amber/ff14SB.xml", "amber/tip3p_standard.xml",
                     "amber/tip3p_HFE_multivalent.xml"],
        small_molecule_forcefield=cfg.small_mol_ff,
        molecules=[molecule] + extra,
        forcefield_kwargs=forcefield_kwargs,
        periodic_forcefield_kwargs=periodic_kwargs,
        cache=os.path.join(out_dir, "ff_cache.json"))
    print(f"[md] solvent={cfg.solvent}  small-molecule FF={cfg.small_mol_ff} (AM1-BCC charges)")

    if cfg.solvent == "water":
        topology, positions = solvate_water(molecule, system_generator, cfg)
    else:
        topology, positions = solvate_chloroform(molecule, chcl3, cfg)

    n_atoms = topology.getNumAtoms()
    print(f"[md] solvated system: {n_atoms} atoms")
    app.PDBFile.writeFile(topology, positions,
                          open(os.path.join(out_dir, f"{name}_{cfg.solvent}_solvated.pdb"), "w"))

    base_system = system_generator.create_system(topology)

    # minimize once; replicas branch from here with fresh velocities
    integ = LangevinMiddleIntegrator(cfg.temperature * unit.kelvin, 1.0 / unit.picosecond,
                                     cfg.timestep * unit.femtosecond)
    plat, prop = _pick_platform(Platform, cfg)
    sim = app.Simulation(topology, base_system, integ, plat, prop)
    sim.context.setPositions(positions)
    print("[md] minimizing solvated system...")
    sim.minimizeEnergy(maxIterations=cfg.min_iters)
    min_pos = sim.context.getState(getPositions=True).getPositions()
    del sim, integ
    return topology, base_system, min_pos, plat, prop


# --- one replica: heat, NPT equilibrate, NPT production ---------------------
def run_replica(topology, base_system, min_pos, plat, prop, name, rep_dir, cfg, seed):
    import openmm
    from openmm import app, unit, XmlSerializer

    os.makedirs(rep_dir, exist_ok=True)
    temperature = cfg.temperature * unit.kelvin
    pressure = cfg.pressure * unit.bar
    dt = cfg.timestep * unit.femtosecond

    # independent copy per replica: each adds its own barostat below, so sharing
    # base_system would stack one barostat per replica
    system = XmlSerializer.deserialize(XmlSerializer.serialize(base_system))
    integrator = openmm.LangevinMiddleIntegrator(temperature, 1.0 / unit.picosecond, dt)
    integrator.setRandomNumberSeed(seed)
    sim = app.Simulation(topology, system, integrator, plat, prop)
    sim.context.setPositions(min_pos)
    sim.context.setVelocitiesToTemperature(50 * unit.kelvin, seed)

    # NVT heating 50 -> T
    heat_steps = int(cfg.heat_ps * 1000 / cfg.timestep)
    n_stages = 25
    for i in range(n_stages):
        t = 50 + (cfg.temperature - 50) * (i + 1) / n_stages
        integrator.setTemperature(t * unit.kelvin)
        sim.step(max(1, heat_steps // n_stages))

    # add barostat -> NPT equilibration
    baro = openmm.MonteCarloBarostat(pressure, temperature, 25)
    baro.setRandomNumberSeed(seed)
    system.addForce(baro)
    sim.context.reinitialize(preserveState=True)
    sim.step(int(cfg.equil_ps * 1000 / cfg.timestep))
    _write_pdb(app, sim, os.path.join(rep_dir, f"{name}_equilibrated.pdb"))

    # production
    steps_per_ns = int(round(1e6 / cfg.timestep))
    prod_steps = int(cfg.prod_ns * steps_per_ns)
    report_every = max(1, int(cfg.report_ps * 1000 / cfg.timestep))
    sim.currentStep = 0   # so the progress reporter excludes heat + equilibration
    dcd = os.path.join(rep_dir, f"{name}_prod.dcd")
    sim.reporters.append(app.DCDReporter(dcd, report_every))
    sim.reporters.append(app.StateDataReporter(
        os.path.join(rep_dir, f"{name}_prod.csv"), report_every, step=True, time=True,
        potentialEnergy=True, temperature=True, density=True, volume=True, speed=True))
    sim.reporters.append(app.StateDataReporter(
        sys.stdout, report_every * 20, step=True, temperature=True, density=True,
        speed=True, progress=True, remainingTime=True, totalSteps=prod_steps))
    sim.reporters.append(app.CheckpointReporter(os.path.join(rep_dir, f"{name}_prod.chk"),
                                                report_every * 20))
    print(f"[md] {os.path.basename(rep_dir)}: production {cfg.prod_ns} ns -> {dcd}")
    t0 = time.time()
    sim.step(prod_steps)
    print(f"[md] {os.path.basename(rep_dir)} done in {(time.time()-t0)/3600:.2f} h")
    _write_pdb(app, sim, os.path.join(rep_dir, f"{name}_final.pdb"))


def run_md(rdkit_mol, name, out_dir, cfg):
    solvent_dir = os.path.join(out_dir, cfg.solvent)
    os.makedirs(solvent_dir, exist_ok=True)
    topology, base_system, min_pos, plat, prop = build_system(rdkit_mol, name, solvent_dir, cfg)
    for k in range(cfg.replicas):
        rep_dir = os.path.join(solvent_dir, f"rep{k+1}")
        run_replica(topology, base_system, min_pos, plat, prop, name, rep_dir, cfg,
                    seed=cfg.seed + 1000 * (k + 1))
    print(f"[md] all {cfg.replicas} replicas ({cfg.solvent}) complete -> {solvent_dir}")
    print(f"[md] analyze with:  python analyze_conformations.py --dir {solvent_dir} --name {name}")


def _pick_platform(Platform, cfg):
    avail = {Platform.getPlatform(i).getName() for i in range(Platform.getNumPlatforms())}
    order = [cfg.platform] if cfg.platform != "auto" else ["CUDA", "OpenCL", "CPU"]
    for pn in order:
        if pn in avail:
            prop = {"Precision": "mixed", "DeviceIndex": str(cfg.gpu)} if pn in ("CUDA", "OpenCL") else {}
            print(f"[md] platform = {pn} {prop}")
            return Platform.getPlatformByName(pn), prop
    sys.exit(f"[fatal] platform '{cfg.platform}' unavailable; have {sorted(avail)}")


def _write_pdb(app, sim, path):
    st = sim.context.getState(getPositions=True, enforcePeriodicBox=True)
    with open(path, "w") as fh:
        app.PDBFile.writeFile(sim.topology, st.getPositions(), fh)


def parse_args():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    m = p.add_argument_group("molecule")
    m.add_argument("--candidate", choices=sorted(CANDIDATES), default=DEFAULT_CANDIDATE)
    m.add_argument("--smiles", default=None, help="custom PROTAC SMILES (e.g. a reference PROTAC)")
    m.add_argument("--name", default=None)
    m.add_argument("--out-dir", default=None)
    m.add_argument("--n-confs", type=int, default=20)

    r = p.add_argument_group("simulation")
    r.add_argument("--prep-only", action="store_true", help="only build 3D conformer (rdkit only)")
    r.add_argument("--solvent", choices=["chloroform", "water"], default="chloroform",
                   help="chloroform = membrane-interior mimic (permeability); water = polar")
    r.add_argument("--replicas", type=int, default=3, help="independent runs (different seeds)")
    r.add_argument("--seed", type=int, default=1234)
    r.add_argument("--small-mol-ff", dest="small_mol_ff", default="openff-2.2.0")
    r.add_argument("--padding", type=float, default=1.2, help="solvent padding, nm")
    r.add_argument("--ionic", type=float, default=0.15, help="NaCl conc (water only), M")
    r.add_argument("--temperature", type=float, default=300.0)
    r.add_argument("--pressure", type=float, default=1.0)
    r.add_argument("--timestep", type=float, default=2.0)
    r.add_argument("--hmass", type=float, default=1.5)
    r.add_argument("--min-iters", type=int, default=5000)
    r.add_argument("--heat-ps", type=float, default=100.0)
    r.add_argument("--equil-ps", type=float, default=500.0)
    r.add_argument("--prod-ns", type=float, default=100.0)
    r.add_argument("--report-ps", type=float, default=10.0)

    h = p.add_argument_group("hardware")
    h.add_argument("--platform", default="auto", choices=["auto", "CUDA", "OpenCL", "CPU"])
    h.add_argument("--gpu", type=int, default=0)
    return p.parse_args()


def main():
    cfg = parse_args()
    smiles = cfg.smiles or CANDIDATES[cfg.candidate]
    name = cfg.name or (cfg.candidate if not cfg.smiles else "custom_protac")
    out_dir = cfg.out_dir or os.path.join(os.path.dirname(os.path.abspath(__file__)), name)
    os.makedirs(out_dir, exist_ok=True)
    print(f"[run] name={name}  solvent={cfg.solvent}  replicas={cfg.replicas}  out={out_dir}")
    print(f"[run] SMILES={smiles}")

    rdkit_mol = build_conformer(smiles, name, out_dir, n_confs=cfg.n_confs)
    if cfg.prep_only:
        print("[run] --prep-only: done."); return
    run_md(rdkit_mol, name, out_dir, cfg)
    print("[run] all done.")


if __name__ == "__main__":
    main()