#!/usr/bin/env python3
"""Render two figures comparing human TYMS (1HVY chain A) with Plasmodium
falciparum bifunctional DHFR-TS TS domain (1J3I chain C).

  human_vs_pf_overlay.png
      structural superposition of the two TS domains, human in
      design-blue (#2563EB), Pf in design-magenta (#B8327E). Active-site
      Cys195 (human) and Cys490 (Pf) shown as sticks.

  human_vs_pf_cavity18.png
      zoom on cavity 18: human in slate, Pf in magenta. Cavity-18
      residues that differ between species (21 mutations from
      cavity18_mutations_per_taxon.json) shown as labelled sticks.
"""
import json, subprocess, shutil
from pathlib import Path

REPO = Path(__file__).resolve().parents[2]
OUT = REPO / "14_inhibitor_design" / "presentation" / "figures"
OUT.mkdir(parents=True, exist_ok=True)
PYMOL = shutil.which("pymol") or "/opt/homebrew/bin/pymol"

HUMAN_PDB = REPO / "03_structure" / "1hvy.pdb"
PF_PDB    = "/tmp/pf_1J3I.pdb"

# Cavity-18 mutations vs human (human residue numbering)
MUT_JSON = REPO / "14_inhibitor_design" / "04_allosteric" / "cavity18_evidence" / "downloads" / "cavity18_mutations_per_taxon.json"
muts = json.loads(MUT_JSON.read_text()).get("Plasmodium_falciparum", [])
# Parse "P26K" → (26, 'P', 'K')
parsed = []
for m in muts:
    wt = m[0]; new = m[-1]
    pos = int(m[1:-1])
    parsed.append((pos, wt, new))

# Plasmodium TS domain starts at residue 281, so human residue N maps to
# Pf residue N+~248 (approximate; PyMOL super will align structurally).
# We mark the human-residue positions on the human chain only — the
# overlay shows where the Pf side chain differs *spatially*.
human_pos_list = "+".join(str(p) for p, _, _ in parsed)

# --- Overlay figure -------------------------------------------------------
overlay_pml = f"""
reinitialize
bg_color white
load {HUMAN_PDB}, human
load {PF_PDB}, pf
# Pf TS domain = chain C residues 281-608
remove pf and not (chain C and resi 281-608)
# Human apo monomer = chain A
remove human and not chain A
# super-align
super pf, human
hide everything
show cartoon, human
show cartoon, pf
color marine, human
color hotpink, pf
# Active-site catalytic Cys (human 195, Pf 490 ≈ 195+295)
show sticks, human and resi 195 and not name H*
show sticks, pf and resi 490 and not name H*
util.cnc human and resi 195
util.cnc pf and resi 490
set stick_radius, 0.22
set cartoon_transparency, 0.0
set ray_opaque_background, 1
set ray_shadow, 0
set light_count, 4
orient human
zoom human, 8
ray 1700, 1100
png {OUT / 'human_vs_pf_overlay.png'}, dpi=160
"""

# --- Cavity-18 close-up figure --------------------------------------------
cavity_pml = f"""
reinitialize
bg_color white
load {HUMAN_PDB}, human
load {PF_PDB}, pf
remove pf and not (chain C and resi 281-608)
remove human and not chain A
super pf, human
hide everything
show cartoon, human
show cartoon, pf
color grey80, human
color grey60, pf
set cartoon_transparency, 0.5, human
set cartoon_transparency, 0.5, pf

# Cavity-18 mutation sites on the HUMAN chain — sticks coloured by physico-chemistry change
select cav_human, human and resi {human_pos_list}
show sticks, cav_human and not name H*
color marine, cav_human
util.cnc cav_human
# Label every position
"""
for p, wt, new in parsed:
    cavity_pml += f"label human and resi {p} and name CA, '{wt}{p}/{new}'\n"

cavity_pml += f"""
set label_size, 18
set label_color, black
set label_font_id, 7
set label_outline_color, white
set label_shadow_mode, 2
set label_position, (3, 3, 3)
set stick_radius, 0.22
set ray_opaque_background, 1
set ray_shadow, 0
set light_count, 4
orient cav_human
zoom cav_human, 6
ray 1700, 1100
png {OUT / 'human_vs_pf_cavity18.png'}, dpi=160
"""

for name, script in [("overlay", overlay_pml), ("cavity", cavity_pml)]:
    pml_path = OUT / f"_pf_{name}.pml"
    pml_path.write_text(script)
    print(f"  rendering {name}…")
    r = subprocess.run([PYMOL, "-cq", str(pml_path)],
                       capture_output=True, timeout=240)
    png = OUT / f"human_vs_pf_{name}.png"
    if name == "cavity": png = OUT / "human_vs_pf_cavity18.png"
    if png.exists():
        print(f"  ✓ {png}  ({png.stat().st_size/1024:.0f} KB)")
    else:
        print(f"  ✗ failed: {r.stderr.decode()[:300]}")