#!/usr/bin/env python3
"""Train month-by-target ML-MODIS forests, optionally task-parallel under torchrun."""

from __future__ import annotations

import argparse
import json
import os
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, feature_names, regression_metrics, validate_multimodal_keys


def args_parser() -> argparse.Namespace:
    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("--paper-model", action="store_true")
    parser.add_argument("--trees", type=int, default=None)
    return parser.parse_args()


def distributed_context() -> tuple[int, int]:
    world = int(os.environ.get("WORLD_SIZE", "1"))
    rank = int(os.environ.get("RANK", "0"))
    if world > 1:
        torch.distributed.init_process_group(backend="gloo")
    return rank, world


def main() -> None:
    args = args_parser()
    config = yaml.safe_load(Path(args.config).read_text())
    settings = dict(config["model"])
    if args.paper_model:
        settings.update(config["paper_model"])
    if args.trees is not None:
        settings["trees"] = args.trees
    data_path = ROOT / (args.data or config["data"]["path"])
    with np.load(data_path) as archive:
        data = {key: archive[key] for key in archive.files}
    validate_multimodal_keys(data)
    rank, world = distributed_context()
    months = config["data"]["months"]
    targets = config["data"]["variables"]["targets"]["names"]
    tasks = [(int(month), target_index, target) for month in months
              for target_index, target in enumerate(targets)]
    local_models = {}
    for task_index, (month, target_index, target) in enumerate(tasks):
        if task_index % world != rank:
            continue
        mask = (data["month"] == month) & (data["year"] != config["train"]["excluded_year"])
        x, y = data["X"][mask], data["Y"][mask, target_index]
        model = BootstrapRandomForestRegressor(
            n_trees=int(settings["trees"]), min_leaf=int(settings["min_leaf"]),
            max_features=int(settings["max_features"]), bootstrap_fraction=float(settings["bootstrap_fraction"]),
            max_depth=settings["max_depth"], split_candidates=int(settings["split_candidates"]),
            seed=int(config["runtime"]["seed"] + task_index * 1009),
        ).fit(x, y)
        oob, counts = model.oob_predict(x)
        local_models[f"{month}:{target}"] = {
            "state": model.state_dict(), "oob_metrics": regression_metrics(y[counts > 0], oob[counts > 0]),
            "train_samples": int(mask.sum()), "excluded_year": int(config["train"]["excluded_year"]),
        }
        print(f"rank={rank} trained month={month} target={target} samples={mask.sum()}", flush=True)
    if world > 1:
        gathered = [None] * world if rank == 0 else None
        torch.distributed.gather_object(local_models, gathered, dst=0)
        if rank == 0:
            local_models = {key: value for shard in gathered for key, value in shard.items()}
    if rank == 0:
        checkpoint = ROOT / (args.checkpoint or config["paths"]["checkpoint"])
        checkpoint.parent.mkdir(parents=True, exist_ok=True)
        model_config = {
            "architecture": "BootstrapRandomForestRegressor", "settings": settings,
            "targets": targets, "months": months, "input_features": 114,
            "feature_names": feature_names(), "excluded_year": int(config["train"]["excluded_year"]),
        }
        torch.save({"model": local_models, "model_config": model_config,
                    "format_version": config["format_version"],
                    "training": {"paper_model": args.paper_model, "distributed_world_size": world}}, checkpoint)
        summary = {key: value["oob_metrics"] for key, value in sorted(local_models.items())}
        metrics_path = ROOT / config["paths"]["training_metrics"]
        metrics_path.parent.mkdir(parents=True, exist_ok=True)
        metrics_path.write_text(json.dumps({"format_version": config["format_version"],
                                            "models": summary}, indent=2, allow_nan=False) + "\n")
        print(json.dumps({"checkpoint": str(checkpoint), "models": len(local_models), "oob": summary}, indent=2))
    if world > 1:
        torch.distributed.destroy_process_group()


if __name__ == "__main__":
    main()