#!/usr/bin/env python3
"""Run all serialized trees and retain ensemble and per-tree predictions."""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

import numpy as np
import torch
import yaml

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "model"))
from ml_modis import BootstrapRandomForestRegressor, validate_multimodal_keys


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default=str(ROOT / "conf/config.yaml"))
    parser.add_argument("--data", default=None)
    parser.add_argument("--checkpoint", default=None)
    parser.add_argument("--output", default=None)
    args = parser.parse_args()
    config = yaml.safe_load(Path(args.config).read_text())
    with np.load(ROOT / (args.data or config["data"]["path"])) as archive:
        data = {key: archive[key] for key in archive.files}
    validate_multimodal_keys(data)
    checkpoint = torch.load(ROOT / (args.checkpoint or config["paths"]["checkpoint"]), map_location="cpu", weights_only=False)
    if checkpoint.get("format_version") != config["format_version"]:
        raise ValueError("Checkpoint format_version does not match configuration")
    targets = list(checkpoint["model_config"]["targets"])
    tree_count = len(next(iter(checkpoint["model"].values()))["state"]["trees"])
    tree_predictions = np.full((data["X"].shape[0], len(targets), tree_count), np.nan, dtype=np.float32)
    for month in checkpoint["model_config"]["months"]:
        mask = data["month"] == month
        for target_index, target in enumerate(targets):
            model = BootstrapRandomForestRegressor.from_state_dict(checkpoint["model"][f"{month}:{target}"]["state"])
            tree_predictions[mask, target_index, :] = model.predict_trees(data["X"][mask])
    prediction = tree_predictions.mean(axis=2)
    safe_prediction = np.where(np.abs(prediction) > 1e-8, prediction, np.nan)
    ratio = data["Y"] / safe_prediction
    if not np.isfinite(prediction).all() or not np.isfinite(ratio).all():
        raise FloatingPointError("Inference produced non-finite values")
    output = ROOT / (args.output or config["paths"]["predictions"])
    output.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(output, pred=prediction, pred_trees=tree_predictions, obs=data["Y"],
                        obs_over_pred=ratio, relative_response=ratio - 1.0,
                        year=data["year"], month=data["month"], platform=data["platform"],
                        latitude=data["latitude"], longitude=data["longitude"],
                        target_names=np.asarray(targets))
    print(f"output={output.relative_to(ROOT)} samples={prediction.shape[0]} "
          f"targets={targets} trees_per_prediction={tree_count}")


if __name__ == "__main__":
    main()