"""Push docking to credibility: Meeko ligand prep + in-place spyrmsd, on the clean HDAC2 target.

Isolates the two cheap suspects behind the 7.9 A redocking failure:
  1. ligand prep — replace bare obabel --gen3d with Meeko (correct rotatable bonds / AD typing)
  2. RMSD metric — replace obrms (superimposes) with spyrmsd in-place, symmetry-corrected
Receptor reused (obabel-protonated 4LXZ with catalytic Zn kept). If redock RMSD drops <2 A, the
docking pipeline is validated and the earlier failure was prep+metric, not docking itself.
"""

from __future__ import annotations

import subprocess
from pathlib import Path

import numpy as np
from rdkit import Chem, RDLogger
from rdkit.Chem import AllChem
from meeko import MoleculePreparation, PDBQTWriterLegacy
from spyrmsd import io as spyio, rmsd as spyrmsd

import sys
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from scripts.dock_positive_controls import STRUCT, VINA, WORK, pubchem_smiles  # noqa: E402
from scripts.dock_validate import dock_pose, extract_crystal_ligand  # noqa: E402

RDLogger.DisableLog("rdApp.*")
PDB, RES, DRUG = "4LXZ", "SHH", "vorinostat"


def prep_receptor_keep_zn() -> tuple[Path, np.ndarray, np.ndarray]:
    atoms, lig, chosen = [], [], None
    for ln in (STRUCT / f"{PDB}.pdb").read_text().splitlines():
        if ln.startswith("ATOM") or (ln.startswith("HETATM") and ln[17:20].strip() == "ZN"):
            atoms.append(ln)
        elif ln.startswith("HETATM") and ln[17:20].strip() == RES:
            key = (ln[21], ln[22:26]); chosen = chosen or key
            if key == chosen:
                lig.append(ln)
    recpdb = WORK / f"{PDB}_rec.pdb"; recpdb.write_text("\n".join(atoms) + "\nEND\n")
    recpdbqt = WORK / f"{PDB}_rec.pdbqt"
    subprocess.run(["obabel", str(recpdb), "-O", str(recpdbqt), "-xr", "-p", "7.4"], capture_output=True)
    xyz = np.array([[float(l[30:38]), float(l[38:46]), float(l[46:54])] for l in lig])
    return recpdbqt, xyz.mean(0), np.clip((xyz.max(0) - xyz.min(0)) + 8, 16, 26)


def meeko_ligand(smiles: str) -> Path:
    mol = Chem.AddHs(Chem.MolFromSmiles(smiles))
    AllChem.EmbedMolecule(mol, randomSeed=0xC0FFEE)
    AllChem.MMFFOptimizeMolecule(mol)
    setups = MoleculePreparation().prepare(mol)
    pdbqt, ok, err = PDBQTWriterLegacy.write_string(setups[0])
    if not ok:
        raise RuntimeError(f"meeko ligand prep failed: {err}")
    out = WORK / f"lig_{DRUG}_meeko.pdbqt"; out.write_text(pdbqt)
    return out


def inplace_rmsd(crystal_pdb: Path, docked_pdbqt: Path) -> float:
    cref = WORK / "xtal.sdf"; cdoc = WORK / "docked.sdf"
    subprocess.run(["obabel", str(crystal_pdb), "-O", str(cref)], capture_output=True)
    subprocess.run(["obabel", str(docked_pdbqt), "-O", str(cdoc), "-f", "1", "-l", "1"], capture_output=True)
    ref, mol = spyio.loadmol(str(cref)), spyio.loadmol(str(cdoc))
    ref.strip(); mol.strip()
    r = spyrmsd.rmsdwrapper(ref, mol, symmetry=True, minimize=False)
    return float(r[0] if isinstance(r, (list, tuple, np.ndarray)) else r)


def main() -> None:
    rec, center, size = prep_receptor_keep_zn()
    smi = pubchem_smiles(DRUG)
    lig = meeko_ligand(smi)
    print(f"meeko ligand pdbqt: {lig.name}; box center {center.round(1)} size {size.round(1)}")
    pose = dock_pose(rec, lig, center, size)  # writes WORK/redock_out.pdbqt
    raw = WORK / "redock_out.pdbqt"
    xtal = extract_crystal_ligand(PDB, RES)
    rmsd = inplace_rmsd(xtal, raw)
    verdict = "PASS (<2A)" if rmsd < 2 else ("MARGINAL (<3A)" if rmsd < 3 else "FAIL")
    print(f"\nvorinostat -> HDAC2 redock | in-place spyrmsd RMSD = {rmsd:.2f} A  {verdict}")
    print("(prior obabel-ligand + obrms gave 7.9 A)")


if __name__ == "__main__":
    main()