"""Train four scale-specific pairs of joint multi-output random forests."""

import json
import os
import random
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))
from model.rf_climparam import FORMAT_VERSION, MODEL_NAME, build_pair


def load_scale(path, scale):
    data = np.load(path)
    expected = {"tend_inputs": 145, "tend_targets": 144, "diff_inputs": 62, "diff_targets": 17}
    if str(data["format_version"]) != FORMAT_VERSION or str(data["scale"]) != scale:
        raise ValueError(f"invalid metadata in {path}")
    count = len(data["tend_inputs"])
    ny, nx = map(int, data["grid_shape"])
    if count != ny * nx or int(data["snapshot_count"]) != 1:
        raise ValueError(f"{scale} must contain one complete [{ny},{nx}] snapshot")
    linear = data["grid_row"].astype(np.int64) * nx + data["grid_column"].astype(np.int64)
    if not np.array_equal(linear, np.arange(count)):
        raise ValueError(f"{scale} grid cannot be reversibly flattened")
    for name, width in expected.items():
        value = data[name]
        if value.shape != (count, width) or value.dtype != np.float32 or not np.isfinite(value).all():
            raise ValueError(f"{scale}/{name} requires finite float32 [{count},{width}]")
    if np.any(data["diff_targets"][:, :15] < 0):
        raise ValueError("Dbar training targets must be nonnegative")
    return data


def sample_training_columns(data, columns_per_latitude, seed):
    ny, nx = map(int, data["grid_shape"])
    if not 1 <= columns_per_latitude <= nx:
        raise ValueError("train_columns_per_latitude must be between 1 and grid width")
    rng = np.random.default_rng(seed)
    selected = []
    for row in range(ny):
        selected.extend(row * nx + rng.choice(nx, columns_per_latitude, replace=False))
    return np.asarray(selected, dtype=np.int64)


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    seed = int(config["seed"])
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    if distributed:
        torch.distributed.init_process_group("gloo")
    rank = torch.distributed.get_rank() if distributed else 0
    world = torch.distributed.get_world_size() if distributed else 1
    scales = list(config["data"]["scales"])
    assigned = [scale for i, scale in enumerate(scales) if i % world == rank]
    local_states, local_metrics = {}, {}
    for scale_index, scale in enumerate(assigned):
        data = load_scale(ROOT / config["data"]["root"] / f"{scale}.npz", scale)
        selected = sample_training_columns(
            data, int(config["data"]["train_columns_per_latitude"]), seed + scales.index(scale))
        pair = build_pair(config["model"], seed + scales.index(scale) * 100)
        pair["rf_tend"].fit(data["tend_inputs"][selected], data["tend_targets"][selected])
        pair["rf_diff"].fit(data["diff_inputs"][selected], data["diff_targets"][selected])
        local_states[scale] = {name: model.state_dict() for name, model in pair.items()}
        local_metrics[scale] = {"complete_grid_points": len(data["tend_inputs"]),
                                "train_columns": len(selected),
                                "columns_per_latitude": int(config["data"]["train_columns_per_latitude"])}
        print(f"rank={rank} trained={scale} sampled_columns={len(selected)}")
    if distributed:
        gathered_states, gathered_metrics = [None] * world, [None] * world
        torch.distributed.all_gather_object(gathered_states, local_states)
        torch.distributed.all_gather_object(gathered_metrics, local_metrics)
        states = {key: value for item in gathered_states for key, value in item.items()}
        metrics = {key: value for item in gathered_metrics for key, value in item.items()}
    else:
        states, metrics = local_states, local_metrics
    if rank == 0:
        if set(states) != set(scales):
            raise RuntimeError("not all scales were trained")
        checkpoint = {
            "model": states,
            "model_name": MODEL_NAME,
            "model_config": {"dimensions": {"rf_tend_input": 145, "rf_tend_output": 144,
                                               "rf_diff_input": 62, "rf_diff_output": 17},
                             "engineering": config["model"]["engineering"],
                             "paper_model": config["paper_model"], "scales": scales},
            "format_version": FORMAT_VERSION, "seed": seed}
        path = ROOT / config["paths"]["checkpoint"]
        path.parent.mkdir(parents=True, exist_ok=True)
        torch.save(checkpoint, path)
        metrics_path = ROOT / config["paths"]["training_metrics"]
        metrics_path.parent.mkdir(parents=True, exist_ok=True)
        metrics_path.write_text(json.dumps({"format_version": FORMAT_VERSION, "scales": metrics}, indent=2) + "\n")
        print(f"checkpoint={path.relative_to(ROOT)} scales={len(states)}")
    if distributed:
        torch.distributed.barrier(); torch.distributed.destroy_process_group()


if __name__ == "__main__":
    main()