#!/usr/bin/env catnip
# Carte de contacts Cα et diagramme de Ramachandran de la crambine (PDB 1CRN)
#
# La crambine est une petite protéine hydrophobe de 46 résidus (chaîne A).
# Deux mesures géométriques classiques la décrivent :
#
#  - Carte de contacts : pour chaque paire de résidus (i, j), on prend la distance
#    entre leurs carbones Cα (le carbone central du squelette peptidique). Un
#    « contact » est une paire proche dans l'espace (< SEUIL) mais éloignée le long
#    de la séquence (|i - j| >= SÉPARATION). Exclure les voisins de séquence isole
#    les rapprochements dus au repliement, pas à la simple continuité de la chaîne.
#
#  - Angles de torsion φ (phi) et ψ (psi) : rotations autour des liaisons N–Cα et
#    Cα–C du squelette. Le diagramme de Ramachandran (Ramachandran et al. 1963)
#    porte φ en abscisse et ψ en ordonnée ; les régions peuplées correspondent aux
#    structures secondaires (hélice α, feuillet β).
#
# gemmi lit le mmCIF et calcule φ/ψ ; on extrait tôt name/seqid/x/y/z dans des
# structs Catnip. numpy/scipy calculent la matrice de distances ; Catnip organise
# la sélection des contacts, la classification et le rapport.
#
# Provenance des données :
#   PDB 1CRN — crambin, Crambe abyssinica, résolution 0.54 Å (chaîne A, 46 résidus)
#   Citation primaire : Teeter M.M. (1984) PNAS 81:6014-6018
#   Source : RCSB PDB / wwPDB — format PDBx/mmCIF — licence CC0 (domaine public)
#
# Portée de l'exemple : seuls les Cα de la chaîne A sont analysés. Sont exclus les
# eaux, hétéro-atomes et ligands (absents ici : la chaîne A est 100 % acides
# aminés). find_atom('CA', '*') retourne le conformère primaire ; ce dépôt ne
# contient aucun altloc secondaire (colonne altloc entièrement vide). Aucune
# symétrie cristallographique n'est appliquée : distances sur l'unité asymétrique.
#
# DEPS: gemmi numpy scipy matplotlib

gemmi = import('gemmi')
numpy = import('numpy')
math = import('math')
sp_dist = import('scipy.spatial.distance')
mpl = import('matplotlib')
mpl.use('Agg')  # backend headless : rendu fichier, aucune fenêtre
plt = import('matplotlib.pyplot')
import('pathlib', 'Path')

script_dir = Path(META.file).parent
output_dir = script_dir / 'output'
output_dir.mkdir(exist_ok=True)

# Paramètres géométriques (documentés, unités explicites)
SEQ_SEP = 4            # séparation minimale en séquence : |i - j| >= 4 (non-local)
THRESHOLD = 8.0        # seuil de contact Cα-Cα en Ångström (convention type CASP)
DEG = 180.0 / math.pi  # radians (gemmi) -> degrés (affichage/Ramachandran)

PDB_ID = '1CRN'
CITATION = "Teeter M.M. (1984) PNAS 81:6014-6018"
PROVENANCE = "RCSB PDB / wwPDB, PDBx/mmCIF, licence CC0"

# Un résidu réduit à son identité et à la position de son Cα. On ne broadcaste
# jamais sur les objets gemmi bruts : tout est copié ici dans des scalaires.
struct Residue {
    name: str; seqid: int; x: float; y: float; z: float;
}

# Une paire de résidus en contact spatial mais distants en séquence.
struct Contact {
    i: int; j: int; dist: float;
}

# Régions grossières du plan de Ramachandran (bornes en degrés). Ce sont des
# boîtes indicatives, PAS une assignation de structure secondaire (qui exige un
# algorithme dédié type DSSP). Les régions favorisées précises proviennent
# d'enquêtes statistiques (Lovell et al. 2003).
union RamaRegion {
    alpha_r; beta; alpha_l; other

    label(self): str => {
        match self {
            RamaRegion.alpha_r => { "hélice α droite" }
            RamaRegion.beta    => { "feuillet β" }
            RamaRegion.alpha_l => { "hélice α gauche" }
            RamaRegion.other   => { "autre / coude" }
        }
    }
}

# Un résidu interne avec ses deux angles de torsion et sa région heuristique.
struct Torsion {
    seqid: int; name: str; phi: float; psi: float; region;
}

in_box = (v: float, lo: float, hi: float): bool => { v >= lo and v <= hi }

classify_rama = (phi: float, psi: float): RamaRegion => {
    match True {
        _ if in_box(phi, -135.0, -25.0) and in_box(psi, -70.0, 20.0) => { RamaRegion.alpha_r }
        _ if in_box(phi, -180.0, -45.0) and in_box(psi, 90.0, 180.0) => { RamaRegion.beta }
        _ if in_box(phi, 25.0, 90.0) and in_box(psi, -20.0, 80.0)    => { RamaRegion.alpha_l }
        _                                                            => { RamaRegion.other }
    }
}

# Code entier stable pour le comptage (évite de comparer des libellés affichés).
region_code = (r: RamaRegion): int => {
    match r {
        RamaRegion.alpha_r => { 0 }
        RamaRegion.beta    => { 1 }
        RamaRegion.alpha_l => { 2 }
        RamaRegion.other   => { 3 }
    }
}

rama_color = (r: RamaRegion): str => {
    match r {
        RamaRegion.alpha_r => { '#d62728' }
        RamaRegion.beta    => { '#1f77b4' }
        RamaRegion.alpha_l => { '#2ca02c' }
        RamaRegion.other   => { '#7f7f7f' }
    }
}

# Rapport des grandeurs CALCULÉES (distinctes des faits extraits et des
# classifications heuristiques ci-dessous).
struct AnalysisReport {
    n_residues: int; n_contacts: int; seq_sep: int; threshold: float;
    longest: str; max_ctrl_err: float; n_torsions: int; rg: float;

    display(self) => {
        print("⇒ Mesures calculées (géométrie, pas d'interprétation biologique)")
        print(f"  Contacts non-locaux : {self.n_contacts} paires")
        print(f"    définition : |i - j| >= {self.seq_sep} et distance Cα-Cα < {self.threshold} Å")
        print(f"  Contact le plus longue-portée : {self.longest}")
        print(f"  Rayon de giration (Cα) : {round(self.rg, 3)} Å")
        print(f"  Angles de torsion φ/ψ complets : {self.n_torsions} résidus internes")
        print(f"  Contrôle numérique (cdist numpy vs gemmi dist) : erreur max = {self.max_ctrl_err} Å")
    }
}

# --- Chargement et extraction (FAITS lus dans le fichier) ---

cif_path = script_dir / 'data/1crn.cif'
st = gemmi.read_structure(str(cif_path))
chain = st[0][0]  # modèle 0, chaîne A
n = len(chain)
idx = numpy.arange(n).tolist()

# Broadcast sur des index entiers (pas sur les résidus gemmi) : chaque itération
# lit son Cα et copie name/seqid/x/y/z dans un struct Catnip.
residues = idx.[(i) => {
    r = chain[int(i)]
    ca = r.find_atom('CA', '*')  # conformère primaire (aucun altloc ici)
    Residue(r.name, r.seqid.num, float(ca.pos.x), float(ca.pos.y), float(ca.pos.z))
}]

sequence = residues.[(r) => { r.name }]

# --- Matrice de distances et contacts (MESURES) ---

coords = numpy.array(residues.[(r) => { list(r.x, r.y, r.z) }])
dmat = sp_dist.cdist(coords, coords)  # (n, n) distances Cα-Cα en Å

# triu_indices(n, SEQ_SEP) : couples (i, j) avec j - i >= SEQ_SEP, donc i < j et
# séparation de séquence suffisante. Indexation à plat car Catnip ne parse pas
# l'indexation numpy à deux axes dmat[ii, jj].
tri = numpy.triu_indices(n, SEQ_SEP)
ii = tri[0]
jj = tri[1]
dv = dmat.ravel()[ii * n + jj]
sel = numpy.less(dv, THRESHOLD)
sel_i = ii[sel].tolist()
sel_j = jj[sel].tolist()
sel_d = dv[sel].tolist()

contacts = numpy.arange(len(sel_i)).tolist().[(k) => {
    Contact(int(sel_i[int(k)]), int(sel_j[int(k)]), float(sel_d[int(k)]))
}]

# Contrôle numérique : la distance de scipy (cdist) confrontée à la distance
# calculée indépendamment par gemmi (pos.dist, C++). Deux implémentations de la
# norme euclidienne sur les mêmes coordonnées : l'écart max doit être ~0.
ctrl_err = contacts.[(c) => {
    a = chain[c.i].find_atom('CA', '*')
    b = chain[c.j].find_atom('CA', '*')
    abs(float(a.pos.dist(b.pos)) - c.dist)
}]
max_ctrl_err = float(numpy.max(numpy.array(ctrl_err)))

# Contact le plus longue-portée : plus grande séparation de séquence.
seps = numpy.array(contacts.[(c) => { c.j - c.i }])
lc = contacts[int(numpy.argmax(seps))]
ri = residues[lc.i]
rj = residues[lc.j]
longest = f"{ri.name}{ri.seqid}–{rj.name}{rj.seqid} à {round(lc.dist, 2)} Å (|i-j| = {lc.j - lc.i})"

# Rayon de giration : dispersion RMS des Cα autour du centroïde (compacité).
centroid = numpy.mean(coords, axis=0)
rg = float(numpy.sqrt(numpy.mean(numpy.sum(numpy.square(coords - centroid), axis=1))))

# --- Angles de torsion φ/ψ (MESURES) ---
# Résidus internes seulement : le premier n'a pas de φ (pas de prev), le dernier
# pas de ψ (pas de next). Ces extrémités sont donc exclues du diagramme, qui
# n'accepte que les couples (φ, ψ) complets.
inner = numpy.arange(1, n - 1).tolist()
torsions = inner.[(i) => {
    r = chain[int(i)]
    pp = gemmi.calculate_phi_psi(chain[int(i) - 1], r, chain[int(i) + 1])  # radians
    phi = float(pp[0]) * DEG
    psi = float(pp[1]) * DEG
    Torsion(r.seqid.num, r.name, phi, psi, classify_rama(phi, psi))
}]

# Comptage par région (CLASSIFICATION heuristique).
codes = numpy.array(torsions.[(t) => { region_code(t.region) }])
region_counts = numpy.arange(4).tolist().[(c) => { int(numpy.sum(numpy.equal(codes, int(c)))) }]

# --- Rapport ---

print(f"⇒ Faits extraits du fichier ({PDB_ID})")
print(f"  Provenance : {PROVENANCE}")
print(f"  Citation primaire : {CITATION}")
print(f"  Structure : {st.name}, chaîne {chain.name}, {n} résidus (Cα)")
print(f"  Séquence : {reduce(sequence.[(nm) => { nm }], (a, b) => { a + '-' + b })}")
print()

report = AnalysisReport(n, len(contacts), SEQ_SEP, THRESHOLD, longest, max_ctrl_err, len(torsions), rg)
report.display()
print()

print("⇒ Classification heuristique des angles φ/ψ (indicative, non DSSP)")
region_labels = list(
    RamaRegion.alpha_r.label(),
    RamaRegion.beta.label(),
    RamaRegion.alpha_l.label(),
    RamaRegion.other.label(),
)
numpy.arange(4).tolist().[(c) => {
    print(f"  {region_labels[int(c)]} : {region_counts[int(c)]} résidus")
}]
print()

# --- Figure : carte de contacts + diagramme de Ramachandran ---

fig, axes = plt.subplots(1, 2, figsize=tuple(13, 6))

# Panneau gauche : matrice de distances (heatmap) + contacts non-locaux (points).
im = axes[0].imshow(dmat, cmap='viridis', origin='lower')
fig.colorbar(im, ax=axes[0], label="distance Cα-Cα (Å)", fraction=0.046, pad=0.04)
ci = numpy.array(contacts.[(c) => { c.i }])
cj = numpy.array(contacts.[(c) => { c.j }])
# Les deux triangles : le contact (i, j) est symétrique.
axes[0].scatter(cj, ci, s=14, c='white', edgecolors='black', linewidths=0.4)
axes[0].scatter(ci, cj, s=14, c='white', edgecolors='black', linewidths=0.4)
axes[0].set_title(f"Carte de contacts Cα ({PDB_ID}) — seuil {THRESHOLD} Å, |i-j| >= {SEQ_SEP}")
axes[0].set_xlabel("index résidu")
axes[0].set_ylabel("index résidu")

# Panneau droit : diagramme de Ramachandran, points colorés par région.
phi_arr = numpy.array(torsions.[(t) => { t.phi }])
psi_arr = numpy.array(torsions.[(t) => { t.psi }])
pt_colors = torsions.[(t) => { rama_color(t.region) }]
axes[1].axhline(0.0, color='#cccccc', linewidth=0.8)
axes[1].axvline(0.0, color='#cccccc', linewidth=0.8)
axes[1].scatter(phi_arr, psi_arr, c=pt_colors, s=36, edgecolors='black', linewidths=0.4)
axes[1].set_xlim(-180.0, 180.0)
axes[1].set_ylim(-180.0, 180.0)
axes[1].set_xticks(list(-180, -90, 0, 90, 180))
axes[1].set_yticks(list(-180, -90, 0, 90, 180))
axes[1].set_title("Diagramme de Ramachandran (φ vs ψ, résidus internes)")
axes[1].set_xlabel("φ (degrés)")
axes[1].set_ylabel("ψ (degrés)")

fig.tight_layout()
output_path = output_dir / 'crambin_contact_map.png'
fig.savefig(str(output_path), dpi=120, bbox_inches='tight')

print(f"⇒ Figure → {output_path}")

# Aperçu navigateur : sert la figure une fois puis rend la main.
# --no-browser garde un chemin headless (le PNG reste écrit ci-dessus).
if '--no-browser' not in import('sys').argv {
    http = import('http')
    b64 = import('base64').b64encode(output_path.read_bytes()).decode('ascii')
    http.serve(f'<!doctype html><meta charset="utf-8"><title>Crambin contact map</title><body style="margin:0;background:#0d1117"><img style="max-width:100%;display:block;margin:0 auto" src="data:image/png;base64,{b64}">',
        0, 'text/html; charset=utf-8', True)
}