"""Same held-out cases through several autoencoders, compared where the SSIM scores are: the original scan.

Each model runs in the geometry it was trained with; its reconstruction is de-normalised and mapped back
onto the scan's own 1 mm grid with MONAI's inverse transforms (as compare_autoencoders.py scores it). One
row per case; for each plane the original, then each model, cut through the tumour centre and shown in
one intensity window per row. Each model's foreground SSIM at 1 mm is printed on its axial panel.

python render_same_cases.py <autoencoder app_files> <images dir> <out.png> CASE:SEQUENCE[,CASE:SEQUENCE...] \
    NAME:GEOMETRY:CHECKPOINT:CONFIG_YAML [...]
"""

import os
import sys
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import monai.transforms as mt  # noqa: E402
import numpy as np  # noqa: E402
import torch  # noqa: E402
import yaml  # noqa: E402
from monai.data import MetaTensor  # noqa: E402
from monai.networks.nets import AutoencoderKL  # noqa: E402

app, images, out = Path(sys.argv[1]), Path(sys.argv[2]), Path(sys.argv[3])
cases = [tuple(item.split(":")) for item in sys.argv[4].split(",")]
specs = [spec.split(":") for spec in sys.argv[5:]]
sys.path.insert(0, str(app))
from trainer import foreground_ssim  # noqa: E402

device = torch.device(os.environ.get("DEVICE", "cpu"))
PLANES = (("axial", 2, False), ("coronal", 1, False), ("sagittal", 0, True))
to_ras = mt.Orientation(axcodes="RAS")


def geometry_chain(geometry: str, document: dict) -> mt.Compose:
    shape = document["training"]["spatial_shape"]
    steps = [mt.LoadImaged(keys=["image"], image_only=True), mt.EnsureChannelFirstd(keys=["image"], channel_dim="no_channel"),
             mt.Orientationd(keys=["image"], axcodes="RAS")]
    if geometry == "squash":
        steps.append(mt.Resized(keys=["image"], spatial_size=shape))
    else:
        steps += [mt.Spacingd(keys=["image"], pixdim=document["training"]["voxel_spacing"], mode="bilinear"),
                  mt.ResizeWithPadOrCropd(keys=["image"], spatial_size=shape)]
    return mt.Compose(steps)


def load_model(checkpoint: str, document: dict) -> AutoencoderKL:
    model = AutoencoderKL(**document["autoencoder"]).to(device).eval()
    state = torch.load(checkpoint, map_location="cpu", weights_only=True)
    model.load_state_dict({k.removeprefix("autoencoder."): v for k, v in state.items()})
    return model


def back_to_original(path: Path, chain: mt.Compose, model: AutoencoderKL) -> torch.Tensor:
    data = chain({"image": str(path)})
    grid = data["image"]
    mean, std = grid.mean(), grid.std()
    with torch.no_grad():
        reconstruction = model(((grid - mean) / std).unsqueeze(0).to(device))[0].float().cpu()
    data["image"] = grid.clone().copy_(reconstruction[0] * std + mean)
    return torch.as_tensor(chain.inverse(data)["image"]).float()


def cut(volume: np.ndarray, axis: int, index: int, flip_lr: bool) -> np.ndarray:
    plane = np.rot90(np.take(volume, index, axis=axis))
    return np.fliplr(plane) if flip_lr else plane


models = []
for name, geometry, checkpoint, config in specs:
    document = yaml.safe_load(Path(config).read_text())
    models.append((name, geometry_chain(geometry, document), load_model(checkpoint, document)))

columns = ["original"] + [name for name, _, _ in models]
figure, axes = plt.subplots(len(cases), len(PLANES) * len(columns),
                            figsize=(1.9 * len(PLANES) * len(columns), 2.1 * len(cases)), squeeze=False)
figure.patch.set_facecolor("black")
for row, (case, sequence) in enumerate(cases):
    folder = next(p.parent for p in images.glob(f"*/scans/label_{case}.nii.gz"))
    path = folder / f"input_{sequence}_{case}.nii.gz"
    truth = mt.LoadImage(image_only=True)(str(path))[None]
    label = mt.LoadImage(image_only=True)(str(folder / f"label_{case}.nii.gz"))[None]
    volumes, scores = [to_ras(truth)[0].numpy()], [None]
    for name, chain, model in models:
        back = back_to_original(path, chain, model)
        scores.append(foreground_ssim((back * (truth > 0))[None].float(), torch.as_tensor(truth)[None].float()))
        volumes.append(to_ras(MetaTensor(back, affine=truth.affine))[0].numpy())
        print(f"{case} {sequence} {name}: SSIM at 1 mm {scores[-1]:.3f}", flush=True)
    tumour = to_ras(label)[0].numpy() > 0
    centre = [int(round(c.mean())) for c in np.nonzero(tumour)]
    brain = volumes[0] > 0
    vmax = float(np.percentile(volumes[0][brain], 99.5))
    for p, (view, axis, flip_lr) in enumerate(PLANES):
        mask = cut(brain, axis, centre[axis], flip_lr)
        rows_, cols_ = np.nonzero(mask)
        box = (slice(max(rows_.min() - 4, 0), rows_.max() + 5), slice(max(cols_.min() - 4, 0), cols_.max() + 5))
        for c, (title, volume, score) in enumerate(zip(columns, volumes, scores)):
            handle = axes[row][p * len(columns) + c]
            handle.imshow(cut(volume, axis, centre[axis], flip_lr)[box], cmap="gray", vmin=0, vmax=vmax)
            handle.set_xticks([])
            handle.set_yticks([])
            for spine in handle.spines.values():
                spine.set_visible(False)
            if row == 0:
                handle.set_title(f"{title}\n{view}", color="0.85", fontsize=10)
            if p == 0 and score is not None:
                handle.text(0.03, 0.03, f"SSIM {score:.3f}", transform=handle.transAxes, color="yellow", fontsize=9,
                            bbox={"facecolor": "black", "edgecolor": "none", "alpha": 0.75, "pad": 2})
            if p == 0 and c == 0:
                handle.set_ylabel(f"{case}\n{sequence}", color="0.85", fontsize=10)
figure.tight_layout()
figure.savefig(out, facecolor=figure.get_facecolor(), dpi=110)
print("wrote", out)