"""ACE variable order and channel contract."""

from __future__ import annotations

import torch

PROGNOSTIC_3D = tuple(f"{name}_{level}" for name in ("T", "qT", "u", "v") for level in range(8))
PROGNOSTIC_SURFACE = ("Ts_land_or_seaice", "ps")
PROGNOSTIC_CHANNELS = PROGNOSTIC_3D + PROGNOSTIC_SURFACE
FORCING_CHANNELS = ("DSW_RF_toa", "Ts_open_ocean", "zs", "fl", "fo", "fsi")
DIAGNOSTIC_CHANNELS = (
    "USW_RF_toa", "ULW_RF_toa", "USW_RF_sfc", "ULW_RF_sfc",
    "DSW_RF_sfc", "DLW_RF_sfc", "P", "dTWP_adv_dt", "LHF", "SHF",
)
PROGNOSTIC_UNITS = ("K", "kg/kg", "m/s", "m/s") * 8 + ("K", "Pa")
FORCING_UNITS = ("W/m2", "K", "m", "1", "1", "1")
DIAGNOSTIC_UNITS = ("W/m2",) * 6 + ("kg/m2/s", "kg/m2/s", "W/m2", "W/m2")
INPUT_UNITS = PROGNOSTIC_UNITS + FORCING_UNITS
OUTPUT_UNITS = PROGNOSTIC_UNITS + DIAGNOSTIC_UNITS
INPUT_CHANNELS = PROGNOSTIC_CHANNELS + FORCING_CHANNELS
OUTPUT_CHANNELS = PROGNOSTIC_CHANNELS + DIAGNOSTIC_CHANNELS


def ledger() -> dict:
    return {
        "input_channels": list(INPUT_CHANNELS),
        "output_channels": list(OUTPUT_CHANNELS),
        "input_units": list(INPUT_UNITS),
        "output_units": list(OUTPUT_UNITS),
        "prognostic_channels": len(PROGNOSTIC_CHANNELS),
        "forcing_channels": len(FORCING_CHANNELS),
        "diagnostic_channels": len(DIAGNOSTIC_CHANNELS),
        "input_total": len(INPUT_CHANNELS),
        "output_total": len(OUTPUT_CHANNELS),
    }


def validate_channels(x: torch.Tensor, expected: int, name: str = "tensor") -> None:
    if x.ndim != 4:
        raise ValueError(f"{name} must have shape [B,C,H,W], got {tuple(x.shape)}")
    if x.shape[1] != expected:
        raise ValueError(f"{name} must have {expected} channels, got {x.shape[1]}")


def split_input(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    validate_channels(x, len(INPUT_CHANNELS), "input")
    n = len(PROGNOSTIC_CHANNELS)
    return x[:, :n], x[:, n:]


def split_output(y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    validate_channels(y, len(OUTPUT_CHANNELS), "output")
    n = len(PROGNOSTIC_CHANNELS)
    return y[:, :n], y[:, n:]