#!/usr/bin/env python3
"""Run autoregressive FengWu-W2S inference and save forecast fields."""

from __future__ import annotations

import argparse
import json
import sys
from datetime import datetime, timedelta
from pathlib import Path

import numpy as np
import torch
import yaml

PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT))

from model.fengwu_w2s import FengWuW2S
from scripts.data_loader import ERA5WindowDataset, read_metadata, resolve_data_dir


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", default=str(PROJECT_ROOT / "conf/config.yaml"))
    parser.add_argument("--data-dir")
    parser.add_argument("--checkpoint")
    parser.add_argument("--output-dir")
    parser.add_argument("--split", choices=("train", "val", "test"), default="test")
    parser.add_argument("--steps", type=int)
    parser.add_argument("--limit", type=int, help="Limit windows for a quick smoke test")
    parser.add_argument("--stochastic", action="store_true")
    parser.add_argument("--seed", type=int, default=0)
    return parser.parse_args()


def _resolve(path: str | Path, root: Path = PROJECT_ROOT) -> Path:
    candidate = Path(path).expanduser()
    return candidate if candidate.is_absolute() else (root / candidate).resolve()


def _build_model(config: dict) -> FengWuW2S:
    model_cfg = config["model"]
    data_cfg = config["data"]
    return FengWuW2S(
        in_channels=len(data_cfg["channels"]),
        group_indices=data_cfg["groups"],
        hidden_channels=int(model_cfg.get("hidden_channels", 32)),
        latent_channels=int(model_cfg.get("latent_channels", 64)),
        patch_size=int(model_cfg.get("patch_size", 4)),
        num_blocks=int(model_cfg.get("num_blocks", 2)),
        perturbation_scale=float(model_cfg.get("perturbation_scale", 0.01)),
    )


def main() -> None:
    args = parse_args()
    with Path(args.config).open(encoding="utf-8") as source:
        config = yaml.safe_load(source)
    data_cfg = config["data"]
    model_cfg = config["model"]
    data_dir = resolve_data_dir(args.data_dir or data_cfg["data_dir"], PROJECT_ROOT)
    checkpoint = _resolve(args.checkpoint or model_cfg.get("default_checkpoint", "./data/checkpoints/model_bak.pth"))
    if not checkpoint.exists():
        raise FileNotFoundError(f"Checkpoint not found: {checkpoint}; run scripts/train.py first")
    split_years = {
        "train": data_cfg["train_years"],
        "val": data_cfg["val_years"],
        "test": data_cfg["test_years"],
    }[args.split]
    steps = int(args.steps or model_cfg.get("inference_steps", 1))
    dataset = ERA5WindowDataset(
        data_dir=data_dir,
        years=split_years,
        channels=data_cfg["channels"],
        input_steps=int(model_cfg.get("input_steps", 2)),
        rollout_steps=max(1, steps),
    )
    metadata = read_metadata(data_dir, data_cfg["channels"])
    means = torch.from_numpy(metadata["means"])
    stds = torch.from_numpy(metadata["stds"])
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = _build_model(config).to(device)
    state = torch.load(checkpoint, map_location=device, weights_only=False)
    model.load_state_dict(state["model_state_dict"])
    model.eval()
    output_dir = _resolve(args.output_dir or "./result/output")
    output_dir.mkdir(parents=True, exist_ok=True)
    index = []
    limit = len(dataset) if args.limit is None else min(len(dataset), max(0, int(args.limit)))
    for sample_index in range(limit):
        inputs, _, timestamp = dataset[sample_index]
        initial = inputs.unsqueeze(0).to(device)
        with torch.no_grad():
            prediction = model.rollout(initial, steps=steps, stochastic=args.stochastic, seed=args.seed + sample_index)
        prediction = prediction.squeeze(0).cpu() * stds + means
        year = timestamp[:4]
        year_dir = output_dir / year
        year_dir.mkdir(parents=True, exist_ok=True)
        start = datetime.strptime(timestamp, "%Y%m%d%H")
        for lead, field in enumerate(prediction):
            valid_time = start + timedelta(hours=lead * int(metadata["time_step"]))
            valid_timestamp = valid_time.strftime("%Y%m%d%H")
            path = year_dir / f"{timestamp}_lead{lead:03d}.npy"
            np.save(path, field.numpy().astype(np.float32))
            index.append({
                "source_timestamp": timestamp,
                "valid_timestamp": valid_timestamp,
                "lead": lead,
                # Keep the manifest relocatable when the project is copied to
                # a compute node or downloaded from ModelScope.
                "path": str(path.relative_to(output_dir)),
            })
    (output_dir / "index.json").write_text(json.dumps(index, indent=2), encoding="utf-8")
    print(f"Saved {len(index)} forecast fields under {output_dir}")


if __name__ == "__main__":
    main()