"""Render side-by-side C-alpha traces of an example input and the generated design (matplotlib, no extra viewer).

```bash
python render.py --inputs-dir inputs --designs-dir outputs --out images
```
"""

import argparse
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from biotite.structure import filter_amino_acids
from biotite.structure.io import pdbx


# Atoms that are not part of a C-alpha trace (ligands, motif side chains) are drawn as red dots.
GREY, COLORS = "#9aa0a6", ["#1a73e8", "#e8710a", "#188038", "#d93025", "#9334e6", "#b06000"]


def load(path: Path):
    """Return `(ca_traces, hetero, other, n_atoms)`: per-chain C-alpha coordinates, non-polymer atoms, other resolved
    non-C-alpha atoms and the number of resolved atoms."""
    atoms = pdbx.get_structure(pdbx.CIFFile.read(path), model=1)
    atoms = atoms[np.isfinite(atoms.coord).all(axis=1)]
    is_ca = filter_amino_acids(atoms) & (atoms.atom_name == "CA")
    traces = {chain: atoms.coord[is_ca & (atoms.chain_id == chain)] for chain in dict.fromkeys(atoms.chain_id[is_ca])}
    hetero = atoms.hetero & ~is_ca
    return traces, atoms.coord[hetero], atoms.coord[~is_ca & ~hetero], len(atoms)


def draw(ax, traces, dots, title, grey, center):
    for index, (chain, coords) in enumerate(traces.items()):
        color = GREY if grey else COLORS[index % len(COLORS)]
        if len(coords) > 1:
            ax.plot(*(coords - center).T, color=color, lw=1.6)
        else:
            ax.scatter(*(coords - center).T, color=color, s=18)
    if len(dots):
        ax.scatter(*(dots - center).T, color="#d93025", s=22, alpha=0.95, depthshade=False)
    ax.set_title(title, fontsize=10)
    ax.set_axis_off()


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--inputs-dir", type=Path, required=True)
    parser.add_argument("--designs-dir", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--names", nargs="*")
    args = parser.parse_args()
    args.out.mkdir(parents=True, exist_ok=True)
    for design in sorted(args.designs_dir.glob("*.cif")):
        if args.names and design.stem not in args.names:
            continue
        source = args.inputs_dir / ("inverse_folding.cif" if design.stem == "partial_diffusion" else f"{design.stem}.cif")
        (given, given_hetero, given_other, given_atoms), (generated, generated_hetero, _, _) = load(source), load(design)
        # Sparse conditioning (a motif or a ligand) is drawn atom by atom and zoomed; dense conditioning as a trace.
        sparse = given_atoms <= 400
        panels = [
            (given, np.concatenate([given_hetero, given_other]) if sparse else given_hetero, True, sparse),
            (generated, generated_hetero, False, False),
        ]
        titles = ["input (conditioning, zoomed)" if sparse else "input (conditioning)", "generated design"]
        figure = plt.figure(figsize=(8, 4), dpi=150)
        centers, extents = [], []
        for traces, dots, _, _ in panels:
            points = np.concatenate([*traces.values(), dots]) if len(traces) or len(dots) else np.zeros((1, 3))
            centers.append(points.mean(axis=0))
            extents.append(max(1.0, np.abs(points - centers[-1]).max() * 1.1))
        shared = max(extent for extent, (_, _, _, zoom) in zip(extents, panels) if not zoom)
        for position, ((traces, dots, grey, zoom), center) in enumerate(zip(panels, centers)):
            ax = figure.add_subplot(1, 2, position + 1, projection="3d")
            draw(ax, traces, dots, titles[position], grey, center)
            limit = extents[position] if zoom else shared
            ax.set_xlim(-limit, limit), ax.set_ylim(-limit, limit), ax.set_zlim(-limit, limit)
        figure.suptitle(design.stem.replace("_", " "), fontsize=11)
        figure.tight_layout()
        figure.savefig(args.out / f"{design.stem}.png")
        plt.close(figure)
        print("rendered", design.stem)


if __name__ == "__main__":
    main()