"""Does the diffusion model keep a tumour it is given? Real FLAIR latents with tumours are noised to step
T, denoised by the model, decoded; the bright-voxel fraction (tumour_check.py) before and after, plus a figure."""

import json
import os
import sys
from pathlib import Path

import matplotlib
import torch

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402

LDM = Path.home() / "scratch/ldm1305"
APP = Path.home() / "sandbox/londonaicentre/flip-wt-1305-ldm/fl-tutorials/nvflare/image_synthesis/latent_diffusion_model/app_files"
sys.path.insert(0, str(APP))
os.chdir(APP)
import models  # noqa: E402
from modality import one_hot_condition  # noqa: E402
from monai.networks.schedulers import DDPMScheduler  # noqa: E402

torch.manual_seed(0)
model = models.get_model().eval()
blob = torch.load(LDM / "ws/flip_fedavg/server/simulate_job/FL_global_model.pt", map_location="cpu", weights_only=False)
model.load_state_dict(blob["model"] if isinstance(blob.get("model"), dict) else blob)
scale = json.loads((LDM / "config.used.json").read_text())["LATENT_SCALE_FACTOR"]
real = torch.load(LDM / "tumour/e5_FLAIR.pt", map_location="cpu", weights_only=False)["real"][[0, 3, 4, 6]]  # labelled tumours


def bright(v, tau=1.5):
    v = v.float().flatten()
    lo, hi = torch.quantile(v[::7], torch.tensor([0.01, 0.995]))
    u = ((v - lo) / (hi - lo + 1e-8)).clamp(0, 1)
    brain = u > 0.1
    return float((brain & (u > tau * u[brain].median())).sum() / brain.sum())


cond = one_hot_condition(torch.zeros(len(real), dtype=torch.long), 4, "cpu")  # FLAIR
rows, out = {"real": real}, {"real": [bright(v) for v in real]}
with torch.no_grad():
    z0 = model.autoencoder.encode_stage_2_inputs(real) * scale
    rows["recon"] = model.autoencoder.decode_stage_2_outputs(z0 / scale)
    for T in (250, 500, 750):
        sched = DDPMScheduler(num_train_timesteps=1000, schedule="scaled_linear_beta", clip_sample=False,
                              prediction_type="epsilon", beta_start=0.0015, beta_end=0.0195)
        sched.set_timesteps(1000)
        z = sched.add_noise(z0, torch.randn_like(z0), torch.full((len(real),), T))
        for t in range(T, -1, -1):
            eps = model.diffusion_model(z, timesteps=torch.full((len(real),), t), context=cond)
            z, _ = sched.step(eps, t, z)
        rows[f"T={T}"] = model.autoencoder.decode_stage_2_outputs(z / scale)
        print(T, flush=True)
for k, vols in rows.items():
    out[k] = [bright(v) for v in vols]
print(json.dumps({k: [round(x, 3) for x in v] for k, v in out.items()}))
(LDM / "tumour/sdedit.json").write_text(json.dumps(out))

f, ax = plt.subplots(len(real), len(rows), figsize=(2.2 * len(rows), 2.3 * len(real)), facecolor="k")
for j, (k, vols) in enumerate(rows.items()):
    for i, v in enumerate(vols):
        v = v[0].float()
        lo, hi = torch.quantile(v.flatten()[::7], torch.tensor([0.01, 0.995]))
        ax[i, j].imshow(v[:, :, v.shape[2] // 2 + 4].T.flip(0), cmap="gray", vmin=lo, vmax=hi)
        ax[i, j].axis("off")
    ax[0, j].set_title(k, color="w", fontsize=10)
plt.tight_layout()
plt.savefig(LDM / "tumour/sdedit.png", dpi=80, facecolor="k")