"""Run multi-scale reconstruction inference."""
import importlib.util
from pathlib import Path
import numpy as np, torch, yaml
import argparse
import random

ROOT = Path(__file__).resolve().parents[1]


def validate_config(cfg):
    data, model = cfg["data"], cfg["model"]
    if data["image_size"] != model["image_size"] or data["channels"] != model["in_channels"]:
        raise ValueError("data and model image shape settings must match")
    if data["scales"] != model["scales"]:
        raise ValueError("data.scales and model.scales must match")

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", type=Path, default=ROOT / "conf/config.yaml")
    parser.add_argument("--seed", type=int, default=None)
    parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default=None)
    parser.add_argument("--data", type=Path, default=None)
    parser.add_argument("--checkpoint", type=Path, default=None)
    parser.add_argument("--output-dir", type=Path, default=None)
    args = parser.parse_args()
    cfg = yaml.safe_load(args.config.read_text())
    validate_config(cfg)
    seed = cfg["seed"] if args.seed is None else args.seed
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    spec = importlib.util.spec_from_file_location("satmaepp", ROOT / "model/satmaepp.py"); module = importlib.util.module_from_spec(spec); spec.loader.exec_module(module)
    model_cfg = {k: v for k, v in cfg["model"].items() if k not in {"architecture", "runtime_profile"}}
    requested = args.device or cfg["runtime"]["device"]
    if requested == "cuda" and not torch.cuda.is_available(): raise RuntimeError("CUDA was requested but is unavailable")
    device = torch.device("cuda" if requested == "cuda" or (requested == "auto" and torch.cuda.is_available()) else "cpu")
    model = module.SatMAEPP(**model_cfg).to(device)
    checkpoint = args.checkpoint or ROOT / cfg["paths"]["checkpoint"]
    if not checkpoint.exists(): raise FileNotFoundError("Run training before inference")
    model.load_state_dict(torch.load(checkpoint, map_location="cpu", weights_only=False)["model"]); model.eval()
    archive = np.load(args.data or ROOT / cfg["data"]["root"] / "test.npz")
    images = torch.from_numpy(archive["images"])
    expected = (cfg["data"]["channels"], cfg["data"]["image_size"], cfg["data"]["image_size"])
    if images.ndim != 4 or tuple(images.shape[1:]) != expected:
        raise ValueError(f"test images must have shape [N, {expected[0]}, {expected[1]}, {expected[2]}]")
    targets = {}
    for scale in cfg["model"]["scales"]:
        if scale != 1:
            field = f"images_{scale}x"
            if field not in archive:
                raise ValueError(f"test dataset is missing native target {field}")
            targets[str(scale)] = torch.from_numpy(archive[field])
    images, targets = images.to(device), {key: value.to(device) for key, value in targets.items()}
    with torch.no_grad(): output = model(images, high_resolution_targets=targets)
    out = args.output_dir or ROOT / cfg["paths"]["inference_dir"]; out.mkdir(parents=True, exist_ok=True)
    payload = {"target": images.cpu().numpy(), "prediction": output["reconstruction"].cpu().numpy(), "mask": output["mask"].cpu().numpy(), "labels": archive["labels"]}
    for scale, value in output["predictions"].items():
        payload[f"prediction_{scale}"] = value.cpu().numpy()
        payload[f"target_{scale}"] = output["targets"][scale].cpu().numpy()
    np.savez_compressed(out / "reconstruction.npz", **payload)
    print("inference=", out / "reconstruction.npz")

if __name__ == "__main__": main()