import argparse
from pathlib import Path
import sys

import numpy as np
import torch
import yaml

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from model.spectralgpt import SpectralGPT


def main():
    parser = argparse.ArgumentParser(description="Run SpectralGPT reconstruction")
    parser.add_argument("--config", default="conf/config.yaml")
    parser.add_argument("--checkpoint")
    parser.add_argument("--batch-size", type=int)
    args = parser.parse_args()
    with open(args.config, encoding="utf-8") as handle:
        config = yaml.safe_load(handle)
    requested = config["runtime"]["device"]
    device = torch.device("cuda" if torch.cuda.is_available() and requested != "cpu" else "cpu")
    torch.manual_seed(config["runtime"]["seed"])
    stage = config["stages"][-1]
    model = SpectralGPT(image_size=stage["image_size"], **config["model"]).to(device)
    checkpoint_path = args.checkpoint or config["training"]["checkpoint"]
    if not Path(checkpoint_path).exists():
        raise FileNotFoundError(
            f"Missing checkpoint: {checkpoint_path}. Run `python scripts/train.py` first."
        )
    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
    if checkpoint.get("stage") != stage["name"] or checkpoint.get("image_size") != stage["image_size"]:
        raise ValueError("Checkpoint is not the configured final stage2 checkpoint")
    model.load_state_dict(checkpoint["model"])
    model.eval()
    data_path = Path(config["data"]["test_path"])
    if not data_path.exists():
        raise FileNotFoundError(
            f"Missing inference data: {data_path}. Run `python scripts/fake_data.py` first."
        )
    with np.load(data_path) as data:
        images = data["images"].copy()
        data_source = str(data["data_source"]) if "data_source" in data.files else "unknown"
        protocol = str(data["protocol"]) if "protocol" in data.files else "unknown"
        normalization = str(data["normalization"]) if "normalization" in data.files else "unknown"
        scale_factors = data["scale_factors"].copy() if "scale_factors" in data.files else np.ones(len(images), np.float32)
        stored_stage = str(data["stage"]) if "stage" in data.files else "unknown"
    expected = (config["model"]["in_channels"], stage["image_size"], stage["image_size"])
    if images.dtype != np.float32 or images.ndim != 4 or tuple(images.shape[1:]) != expected:
        raise ValueError(f"Expected float32 stage2 test [N,{','.join(map(str, expected))}], got {images.dtype} {images.shape}")
    if stored_stage != stage["name"]:
        raise ValueError(f"Expected test stage {stage['name']}, got {stored_stage}")
    batch_size = args.batch_size or config["training"]["batch_size"]
    collected = {name: [] for name in ("reconstruction", "prediction_image", "mask", "mask_image")}
    with torch.inference_mode():
        for start in range(0, len(images), batch_size):
            batch = torch.from_numpy(images[start:start + batch_size]).to(device)
            output = model(batch)
            for name in collected:
                collected[name].append(output[name].cpu().numpy())
    output_dir = Path(config["runtime"]["output_dir"])
    output_dir.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(output_dir / "reconstruction.npz",
                         inputs=images,
                         reconstructions=np.concatenate(collected["reconstruction"]),
                         predictions=np.concatenate(collected["prediction_image"]),
                         masks=np.concatenate(collected["mask"]),
                         mask_images=np.concatenate(collected["mask_image"]),
                         data_source=np.asarray(data_source),
                         protocol=np.asarray(protocol), normalization=np.asarray(normalization),
                         scale_factors=scale_factors, stage=np.asarray(stored_stage))
    print(
        f"saved: {output_dir / 'reconstruction.npz'} "
        f"data_source={data_source} protocol={protocol}"
    )


if __name__ == "__main__":
    main()