"""Summarize the AlphaFold Database PAE matrix of human calmodulin-1.

Reproduces the block statistics in content/guides/pae.mdx:

    python analyze.py

Inputs in this directory:
- AF-P0DP23-F1-predicted_aligned_error_v6.json, downloaded unchanged from
  https://alphafold.ebi.ac.uk/files/AF-P0DP23-F1-predicted_aligned_error_v6.json
- plddt.csv, the Cα B-factor column (pLDDT) of
  https://alphafold.ebi.ac.uk/files/AF-P0DP23-F1-model_v6.pdb

Lobes use the UniProt P0DP23 EF-hand domain boundaries: EF-hands 1 and 2 at
residues 8 to 79, EF-hands 3 and 4 at residues 81 to 149. Matrix rows are
aligned residues and columns are scored residues.
"""

import csv
import json
from pathlib import Path
from statistics import mean

HERE = Path(__file__).resolve().parent
LOBES = {"N-lobe": (8, 79), "C-lobe": (81, 149)}
CONFIDENT_PAE = 5


def residues(span):
    start, end = span
    return range(start, end + 1)


def main():
    record = json.loads(
        (HERE / "AF-P0DP23-F1-predicted_aligned_error_v6.json").read_text()
    )[0]
    pae = record["predicted_aligned_error"]
    with (HERE / "plddt.csv").open(newline="") as handle:
        plddt = {int(row["residue"]): float(row["plddt"]) for row in csv.DictReader(handle)}

    blocks = []
    for aligned_name, aligned_span in LOBES.items():
        for scored_name, scored_span in LOBES.items():
            values = [
                pae[i - 1][j - 1]
                for i in residues(aligned_span)
                for j in residues(scored_span)
            ]
            blocks.append(
                {
                    "aligned": aligned_name,
                    "scored": scored_name,
                    "mean_pae": round(mean(values), 1),
                    "percent_below_5": round(
                        100 * sum(value < CONFIDENT_PAE for value in values) / len(values), 1
                    ),
                }
            )

    size = len(pae)
    asymmetry = [
        abs(pae[i][j] - pae[j][i]) for i in range(size) for j in range(i + 1, size)
    ]
    summary = {
        "entry": "AF-P0DP23-F1",
        "residues": size,
        "max_predicted_aligned_error": record["max_predicted_aligned_error"],
        "mean_plddt": round(mean(plddt.values()), 1),
        "lobe_mean_plddt": {
            name: round(mean(plddt[r] for r in residues(span)), 1)
            for name, span in LOBES.items()
        },
        "linker_min_plddt": min(
            (plddt[r], r) for r in residues((LOBES["N-lobe"][1], LOBES["C-lobe"][0]))
        ),
        "blocks": blocks,
        "mean_absolute_asymmetry": round(mean(asymmetry), 1),
    }
    (HERE / "summary.json").write_text(json.dumps(summary, indent=2) + "\n")
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()