"""Generate 1x3 QA screenshots (axial/coronal/sagittal) per volume.

Each proc volume is first mapped to a canonical body array (S, P, L) using the
direction matrix in procmeta (columns = image-axis directions, LPS basis), so
the panels are correct for any source orientation (LPS axial, LIP coronal,
PIR sagittal, ...). Radiological display conventions:
  axial:    anterior on top, patient's left on right
  coronal:  superior on top, patient's left on right
  sagittal: superior on top, anterior on left
Layout: single row, axial | coronal | sagittal; each panel is scaled so its
maximal physical dimension (mm) is the same across the row (panels are
centered vertically, widths follow the physical aspect of each cut).
Usage:
  python scripts/make_screenshots.py --keys k1 k2 ...
  python scripts/make_screenshots.py --sample 20 [--keys ...]
Output: data/screenshots/<patient>/<rest>.png (per-patient, mirrors data/lee_nii)
"""
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import argparse
import json
import random
import glob

import numpy as np
import SimpleITK as sitk
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from src.common import ROOT, d, read_nii_img

B_L = np.array([1.0, 0, 0])
B_P = np.array([0.0, 1, 0])
B_S = np.array([0.0, 0, 1])


def vol_to_body(a, M):
    """a: array (z,y,x) over image axes; M: 3x3 LPS, columns = axis dirs.
    Returns (body, ok, kof): body = (S, P, L)-indexed array, kof maps each
    body axis to its image axis index (0=x, 1=y, 2=z); kof is None if not
    body-aligned."""
    M = np.asarray(M).reshape(3, 3)
    v = a.transpose(2, 1, 0)  # (i0, i1, i2) image axes
    kof, sgn = {}, {}
    for name, b in (("L", B_L), ("P", B_P), ("S", B_S)):
        dots = M.T @ b
        k = int(np.argmax(np.abs(dots)))
        # 0.92 ~= 23 deg: tolerates tilted clinical protocols (e.g. oblique
        # sagittal MP-RAGE ~21 deg from P-S)
        if abs(dots[k]) < 0.92:
            return None, False, None
        kof[name], sgn[name] = k, (1 if dots[k] > 0 else -1)
    out = np.transpose(v, [kof["S"], kof["P"], kof["L"]])
    for ax, name in enumerate(("S", "P", "L")):
        if sgn[name] < 0:
            out = np.flip(out, ax)
    return np.ascontiguousarray(out), True, kof


def panels_from_body(b):
    s, p, l = b.shape
    sm, pm, lm = s // 2, p // 2, l // 2
    axial = b[sm]                      # (P, L): top=anterior, right=left
    coronal = b[:, pm, :][::-1]        # (S-rev, L): top=superior, right=left
    sagittal = b[:, :, lm][::-1]       # (S-rev, P): top=superior, left=anterior
    return axial, coronal, sagittal, (sm, pm, lm)


def shot_path(out_dir, key):
    """Per-patient screenshot path: <out_dir>/<patient>/<rest>.png, where the
    patient id is the 2nd underscore token of the key (lee_/m6_/ntuh_ prefixes)."""
    parts = key.split("_")
    if len(parts) < 3:
        return os.path.join(out_dir, key + ".png")
    return os.path.join(out_dir, parts[1], "_".join(parts[2:]) + ".png")


def screenshot_from_volume(arr, direction, key, out_dir, dpi=110, spacing=None,
                           series_desc=None, protocol=None):
    """Render the 1x3 QA screenshot (axial | coronal | sagittal) from an
    image-axis array (z,y,x), a 3x3 LPS direction matrix (columns =
    image-axis directions) and the image-axis spacing in mm, (x, y, z)
    (None = isotropic 1mm). series_desc / protocol (optional) are shown
    on a second title line. Panel cells follow the physical (mm) aspect
    of each cut. Output: <out_dir>/<patient>/<rest>.png (mirrors the
    data/lee_nii layout). Returns the output path, or None if the volume
    is not body-aligned."""
    if spacing is None:
        spacing = (1.0, 1.0, 1.0)
    body, ok, kof = vol_to_body(arr, direction)
    if not ok:
        return None
    sp = {n: float(spacing[kof[n]]) for n in ("S", "P", "L")}
    lo, hi = np.percentile(body[body > 0], [1, 99.5])
    axial, coronal, sagittal, _ = panels_from_body(body)
    s, p, l = body.shape
    # (image, title, aspect = mm per row / mm per col, n_rows, n_cols)
    panels = [
        (axial, "axial", sp["P"] / sp["L"], p, l),
        (coronal, "coronal", sp["S"] / sp["L"], s, l),
        (sagittal, "sagittal", sp["S"] / sp["P"], s, p),
    ]
    max_in, gap_in, m_in, top_in, bot_in = 4.4, 0.25, 0.3, 1.5, 0.2
    # each panel is scaled so its maximal physical dimension spans max_in
    w_in, h_in = [], []
    for _, _, a, nrow, ncol in panels:
        pw, ph = ncol, nrow * a  # physical extent along display x / y
        m = max(pw, ph)
        w_in.append(max_in * pw / m)
        h_in.append(max_in * ph / m)
    plot_h = max_in
    fig_w = sum(w_in) + 2 * gap_in + 2 * m_in
    fig_h = plot_h + top_in + bot_in
    fig = plt.figure(figsize=(fig_w, fig_h))
    left = m_in / fig_w
    for (im, title, a, _, _), w, h in zip(panels, w_in, h_in):
        axi = fig.add_axes([left, (bot_in + (plot_h - h) / 2) / fig_h,
                            w / fig_w, h / fig_h])
        axi.imshow(im, cmap="gray", vmin=lo, vmax=hi, origin="upper", aspect=a)
        axi.set_title(title, color="w", fontsize=11)
        axi.set_xticks([]); axi.set_yticks([])
        for spine in axi.spines.values():
            spine.set_edgecolor("0.35")
        left += (w + gap_in) / fig_w
    mx = arr.shape[::-1]  # (x, y, z)
    tkey = key.split("_", 1)[1] if "_" in key else key  # drop dataset prefix (lee_/m6_/ntuh_)
    title = (f"{tkey}   {mx[0]}×{mx[1]}×{mx[2]} @ "
             f"{spacing[0]:.3f}×{spacing[1]:.3f}×{spacing[2]:.3f}mm")
    def _clean(s):
        # drop chars the Agg font cannot render (e.g. CJK in old study descs)
        return " ".join(s.split()).encode("ascii", "ignore").decode()

    sd = _clean(series_desc) if series_desc else ""
    pt = _clean(protocol) if protocol else ""
    info = []
    if sd:
        info.append("series: " + sd[:60])
    if pt and pt[:10] != sd[:10]:
        info.append("protocol: " + pt[:60])
    if info:
        title += "\n" + "    ".join(info)
    fig.suptitle(title, color="w", fontsize=12)
    fig.patch.set_facecolor("k")
    out = shot_path(out_dir, key)
    os.makedirs(os.path.dirname(out), exist_ok=True)
    fig.savefig(out, facecolor="k", dpi=dpi)
    plt.close(fig)
    return out


def screenshot(key, out_dir, dpi=110):
    p = os.path.join(d("data/proc"), key + ".nii.gz")
    if not os.path.exists(p):
        return None
    arr = np.asarray(sitk.GetArrayFromImage(read_nii_img(p)), dtype=np.float32)
    meta = json.load(open(os.path.join(d("data/procmeta"), key + ".json")))
    return screenshot_from_volume(arr, meta["direction"], key, out_dir, dpi)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--keys", nargs="*", default=[])
    ap.add_argument("--sample", type=int, default=0)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--out", default=d("data/screenshots"))
    args = ap.parse_args()
    os.makedirs(args.out, exist_ok=True)
    allkeys = [os.path.basename(f)[:-7] for f in glob.glob(os.path.join(d("data/proc"), "*.nii.gz"))
               if not f.endswith("_label.nii.gz")]
    keys = list(args.keys)
    if args.sample:
        rnd = random.Random(args.seed)
        extra = [k for k in allkeys if k not in keys]
        keys += rnd.sample(extra, min(args.sample - len(keys), len(extra)))
    n = 0
    for k in keys:
        r = screenshot(k, args.out)
        if r:
            n += 1
            print("wrote", r, flush=True)
        else:
            print("skip", k, flush=True)
    print(f"done: {n}/{len(keys)} screenshots in {args.out}")


if __name__ == "__main__":
    main()