"""Score extracted autoencoders on held-out scans, in their own grid and back in the original 1 mm space.

Each model is run through the geometry it was trained with: ``squash`` (resize the whole scan to the
grid, the chain before resampling) or ``spacing`` (resample to millimetres, then centre crop or pad).
The reconstruction is de-normalised and mapped back onto the scan's own 240x240x155 grid with MONAI's
inverse transforms, then scored against the untouched scan — so both models answer one question: how
much of the original brain survives the round trip, resampling included.

python compare_autoencoders.py <autoencoder app_files> <images dir> <held-out cases file> \
    NAME:GEOMETRY:CHECKPOINT:CONFIG_YAML [...]

CHECKPOINT ``none`` skips the model: the geometry's round trip alone, the ceiling any autoencoder on
that grid is held under.
"""

import json
import os
import statistics
import sys
from pathlib import Path

import monai.transforms as mt
import numpy as np
import torch
import yaml
from monai.networks.nets import AutoencoderKL

app, images, cases_file = Path(sys.argv[1]), Path(sys.argv[2]), Path(sys.argv[3])
sys.path.insert(0, str(app))
from trainer import foreground_ssim  # noqa: E402

cases = Path(cases_file).read_text().split()
# DEVICE=cpu scores without the GPU (in float32), e.g. while a training run holds it.
device = torch.device(os.environ.get("DEVICE", "cuda"))


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(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


results = {}
per_volume = {}
for spec in sys.argv[4:]:
    name, geometry, checkpoint, config = spec.split(":")
    document = yaml.safe_load(Path(config).read_text())
    chain = geometry_chain(geometry, document)
    model = None if checkpoint == "none" else load(checkpoint, document)
    own, original, paths = [], [], []
    for case in cases:
        for path in sorted(images.glob(f"*/scans/input_*_{case}.nii.gz")):
            data = chain({"image": str(path)})
            grid = data["image"]
            mean, std = grid.mean(), grid.std()
            with torch.no_grad(), torch.autocast(device_type=device.type, enabled=device.type == "cuda"):
                x = ((grid - mean) / std).unsqueeze(0).to(device)
                reconstruction = x if model is None else model(x)[0]
            reconstruction = reconstruction.float()
            own.append(foreground_ssim(reconstruction, x.float()))
            data["image"] = grid.clone().copy_(reconstruction[0].cpu() * std + mean)
            # The inverse undoes the orientation too, so the reference is the file exactly as stored.
            back = torch.as_tensor(chain.inverse(data)["image"]).unsqueeze(0).to(device).float()
            truth = torch.as_tensor(mt.LoadImage(image_only=True)(str(path))).float()[None, None].to(device)
            assert back.shape == truth.shape, (back.shape, truth.shape)
            back = back * (truth > 0)  # the skull-stripped background is exactly zero; score the brain
            original.append(foreground_ssim(back, truth))
            paths.append(path.name)
    results[name] = (own, original)
    per_volume[name] = [{"volume": p, "own": o, "original": r} for p, o, r in zip(paths, own, original)]
    print(f"{name:<6} {len(own)} volumes  own-grid SSIM {statistics.mean(own):.4f} ± {statistics.stdev(own):.4f}   "
          f"original 1 mm SSIM {statistics.mean(original):.4f} ± {statistics.stdev(original):.4f}", flush=True)
models = [name for name, spec in zip(results, sys.argv[4:]) if not spec.split(":")[2] == "none"]
if len(models) == 2:
    first, second = (np.array(results[name][1]) for name in models)
    diff = second - first
    print(f"{models[1]} - {models[0]} at 1 mm: mean {diff.mean():+.4f}, better on {int((diff > 0).sum())} of {len(diff)} volumes")
if os.environ.get("OUT"):
    Path(os.environ["OUT"]).write_text(json.dumps(per_volume, indent=1))