"""Generated #0-#2 per sequence from E5, E7, E10 and E11 beside three held-out patients, mid-axial.

python four_runs.py <ldm app_files> <images dir> <heldout_cases.txt> <out.png> E5=<pt> E7=<pt> E10=<pt> E11=<pt>
"""

import sys
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch

app_files, images, cases_file, out = map(Path, sys.argv[1:5])
runs = dict(a.split("=", 1) for a in sys.argv[5:])
sys.path.insert(0, str(app_files))
from transforms import get_brain_mri_transforms  # noqa: E402

M = ["FLAIR", "T1w", "T1Gd", "T2w"]
N = 3


def every(t: torch.Tensor) -> np.ndarray:
    t = t.as_subclass(torch.Tensor)
    count = t.untyped_storage().nbytes() // (t[0].numel() * t.element_size())
    return torch.empty(0, dtype=t.dtype).set_(t.untyped_storage(), 0, (count, *t.shape[1:])).numpy()[:, 0]


gen = {}
for name, pt in runs.items():
    saved = torch.load(pt, map_location="cpu", weights_only=False)  # our own eval output
    gen[name] = {m: every(saved[m])[:N] for m in M}
chain = get_brain_mri_transforms(is_validation=True)
folders = {p.name.removeprefix("label_").removesuffix(".nii.gz"): p.parent for p in images.glob("*/scans/label_*.nii.gz")}
cases = cases_file.read_text().split()[:N]
gen["held-out real"] = {
    m: [chain({"image": str(folders[c] / f"input_{m}_{c}.nii.gz")})["image"].as_tensor().numpy()[0] for c in cases]
    for m in M}

groups = list(gen)
fig, axes = plt.subplots(len(M), len(groups) * N, figsize=(len(groups) * N * 1.05, len(M) * 1.25), facecolor="black")
for r, m in enumerate(M):
    for g, name in enumerate(groups):
        for i in range(N):
            ax, v = axes[r, g * N + i], gen[name][m][i]
            ax.imshow(np.rot90(v[:, :, v.shape[2] // 2]), cmap="gray", vmin=np.percentile(v, 1),
                      vmax=np.percentile(v, 99.5))
            ax.set_xticks([]), ax.set_yticks([]), ax.set_frame_on(False)
            if r == 0 and i == 1:
                ax.set_title(name, color="white", fontsize=11)
        axes[r, 0].set_ylabel(m, color="white", fontsize=11)
plt.subplots_adjust(left=0.03, right=0.995, top=0.93, bottom=0.01, wspace=0.03, hspace=0.04)
plt.savefig(out, dpi=90, facecolor="black")
print("saved", out)