"""4x4 grid of chest X-rays from the final global model of the X-ray diffusion tutorial.

The first n samples drawn from one seed, not picked, with the tutorial's own scheduler and sampler (the full
1000-step DDPM), shown as the tutorial's debug grids are: one min-max over the whole grid. Runs on the CPU so it
leaves the GPU to the queue. Also saves the samples as a plain tensor.

python sample_xray.py <app custom dir> <FL_global_model.pt> <out.png> [n=16] [seed=0] [threads=8]
"""

import json
import sys
import time
from pathlib import Path

import torch
from torchvision.utils import make_grid, save_image

app, ckpt, out = Path(sys.argv[1]), Path(sys.argv[2]), Path(sys.argv[3])
n, seed, threads = (int(sys.argv[i]) if len(sys.argv) > i else d for i, d in ((4, 16), (5, 0), (6, 8)))
sys.path.insert(0, str(app))
from models import get_model  # noqa: E402
from monai.inferers import DiffusionInferer  # noqa: E402
from monai.networks.schedulers import DDPMScheduler  # noqa: E402

torch.set_num_threads(threads)
config = json.loads((app / "config.json").read_text())
model = get_model()
blob = torch.load(ckpt, map_location="cpu", weights_only=False)  # our own run's output
model.load_state_dict(blob["model"] if isinstance(blob.get("model"), dict) else blob, strict=True)
model.eval()

scheduler = DDPMScheduler(num_train_timesteps=1000, schedule="scaled_linear_beta", clip_sample=False,
                          prediction_type="epsilon", beta_start=0.0015, beta_end=0.0195)  # as trainer.py
inferer = DiffusionInferer(scheduler=scheduler)
shape = [n, config["net_config"]["diffusion_model"]["in_channels"], *config["spatial_shape"]]
generator = torch.Generator().manual_seed(seed)
start = time.time()
with torch.no_grad():
    samples = inferer.sample(input_noise=torch.randn(shape, generator=generator), diffusion_model=model.diffusion_model,
                             scheduler=scheduler, save_intermediates=False, verbose=False)
print(f"sampled {n} in {time.time() - start:.0f} s; range {samples.min():.3f}..{samples.max():.3f}")
torch.save(samples.as_subclass(torch.Tensor).clone(), out.with_suffix(".pt"))
save_image(make_grid(samples, nrow=int(n**0.5), normalize=True, scale_each=False, padding=2), str(out))
print("wrote", out)