# SPDX-License-Identifier: MIT
"""YS-EN-0003 · the encoding gallery: one training curve per class, and the
images the networks actually see.

Single-channel fields are shown through a perceptually uniform colour map
(the models get the raw values, not the colours); the stack is shown as the
RGB image it is (R = GASF, G = GADF, B = MTF).

  python build_gallery.py   # -> assets/img/ys-en-0003-enc-<class>-<encoding>.webp
"""
from __future__ import annotations

from pathlib import Path

import numpy as np
import torch
from matplotlib import colormaps
from PIL import Image

import lcdata

HERE = Path(__file__).resolve().parent
IMG = HERE.parents[2] / "assets" / "img"
SIZE = 224
ENCS = ["gasf", "gadf", "mtf", "rp", "line", "stack"]
SLUG = {0: "cepheid", 1: "eclipsing", 2: "rrlyrae"}


def pick_examples(X, y):
    """The most typical training curve of each class: closest to the class median."""
    out = {}
    for c in range(3):
        Xc = X[y == c]
        med = np.median(Xc, 0)
        out[c] = Xc[np.argmin(((Xc - med) ** 2).sum(1))]
    return out


def to_rgb(img: np.ndarray, enc: str) -> np.ndarray:
    """(3, m, m) encoding in [-1, 1] -> (m, m, 3) uint8 for display."""
    if enc == "stack":
        return ((img.transpose(1, 2, 0) + 1) / 2 * 255).round().clip(0, 255).astype(np.uint8)
    v = (img[0] + 1) / 2
    cmap = colormaps["gray_r" if enc == "line" else "magma"]
    return (cmap(v)[:, :, :3] * 255).round().astype(np.uint8)


def main():
    X, y, _, _ = lcdata.load(0)
    ex = pick_examples(X, y)
    IMG.mkdir(parents=True, exist_ok=True)
    for c, x in ex.items():
        xt = torch.tensor(x[None])
        for enc in ENCS:
            img = lcdata.encode(xt, enc, SIZE)[0].numpy()
            Image.fromarray(to_rgb(img, enc)).save(IMG / f"ys-en-0003-enc-{SLUG[c]}-{enc}.webp", quality=92, method=6)
    np.savez(HERE / "results" / "gallery_examples.npz", **{SLUG[c]: x for c, x in ex.items()})
    print("wrote", len(ex) * len(ENCS), "images to", IMG)


if __name__ == "__main__":
    main()