"""Explicitly load the v2 research preview and super-resolve a saved RGBN reflectance tensor."""
import argparse
from pathlib import Path

import torch
from omegaconf import OmegaConf
from opensr_model import SRLatentDiffusion


def load_model(device=None):
    root = Path(__file__).resolve().parent
    config = OmegaConf.load(root / "inference.yaml")
    device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
    model = SRLatentDiffusion(config, device=device)
    model.load_pretrained(str(root.parent / "opensr-ldsrs2_v2_0_0_preview.ckpt"))
    model.eval()
    return model, config


@torch.inference_mode()
def super_resolve(model, config, rgbn, *, seed=42, sampler="dpm_solver++", steps=None):
    if not isinstance(rgbn, torch.Tensor):
        raise TypeError("Input must be a torch.Tensor")
    if rgbn.ndim not in (3, 4) or rgbn.shape[-3] != 4:
        raise ValueError("Expected [4,H,W] or [B,4,H,W] in R,G,B,NIR order")
    if not rgbn.is_floating_point() or not torch.isfinite(rgbn).all():
        raise ValueError("Input must contain finite floating-point reflectance")
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    steps = steps if steps is not None else (3 if sampler == "dpm_solver++" else 100)
    return model(
        rgbn.to(device=model.device, dtype=torch.float32),
        sampler=sampler,
        sampling_steps=steps,
        sampling_eta=0.0 if sampler == "dpm_solver++" else 0.95,
        spectral_correction="back_projection",
        back_projection=config.back_projection,
    )


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    source = parser.add_mutually_exclusive_group(required=True)
    source.add_argument("--input", type=Path, help="A saved float RGBN tensor (.pt)")
    source.add_argument("--smoke", action="store_true", help="Use synthetic input")
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--device", default=None)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--sampler", choices=["dpm_solver++", "ddim"], default="dpm_solver++")
    parser.add_argument("--steps", type=int)
    args = parser.parse_args()
    torch.manual_seed(args.seed)
    rgbn = torch.rand(1, 4, 128, 128) if args.smoke else torch.load(
        args.input, map_location="cpu", weights_only=True
    )
    model, config = load_model(args.device)
    sr = super_resolve(model, config, rgbn, seed=args.seed, sampler=args.sampler, steps=args.steps)
    if not torch.isfinite(sr).all():
        raise RuntimeError("Non-finite model output")
    args.output.parent.mkdir(parents=True, exist_ok=True)
    torch.save(sr.cpu(), args.output)
    print(f"Saved {tuple(sr.shape)} to {args.output}")


if __name__ == "__main__":
    main()