"""Generate 21 structured, physically interpretable full-grid response pairs."""

import sys
from pathlib import Path

import numpy as np
import yaml


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model.climemu_s2l import FORMAT_VERSION, GRID_SHAPE


SCENARIOS = [
    ("PDRMIP_2xCO2", 2.5, "global", 0, 0),
    ("PDRMIP_3xCH4", 1.5, "global", 0, 0),
    ("PDRMIP_10xCFC12", 1.1, "global", 0, 0),
    ("PDRMIP_solar", 1.0, "global", 0, 0),
    ("PDRMIP_5xSO4", -1.6, "global", 0, 0),
    ("PDRMIP_10xBC", 1.2, "global", 0, 0),
    ("PDRMIP_10xSO4_Europe", -1.0, "regional", 52, 15),
    ("PDRMIP_10xSO4_Asia", -1.1, "regional", 35, 105),
    ("PDRMIP_preindustrial_SO4", 0.8, "global", 0, 0),
    ("ECLIPSE_CH4_minus20", -0.7, "global", 0, 0),
    ("ECLIPSE_2xCO2", 2.3, "global", 0, 0),
    ("ECLIPSE_BC_minus100", -0.6, "global", 0, 0),
    ("ECLIPSE_SO2_minus100", 1.0, "global", 0, 0),
    ("ECLIPSE_CO_minus100", -0.45, "global", 0, 0),
    ("KASOAR_SO2_NHML_minus100", 0.75, "regional", 42, 30),
    ("KASOAR_BC_NHML_minus100", -0.55, "regional", 42, 30),
    ("KASOAR_SO2_China_minus100", 0.70, "regional", 34, 105),
    ("KASOAR_SO2_EastAsia_minus100", 0.65, "regional", 40, 125),
    ("KASOAR_SO2_Europe_minus100", 0.62, "regional", 52, 15),
    ("KASOAR_SO2_US_minus100", 0.58, "regional", 40, 260),
    ("PDRMIP_SO4_Asia_alt", -0.72, "regional", 25, 80),
]


def wrapped_distance(longitude, centre):
    return (longitude - centre + 180.0) % 360.0 - 180.0


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    ny, nx = map(int, config["data"]["grid"])
    if (ny, nx) != GRID_SHAPE or len(SCENARIOS) != 21:
        raise ValueError("protocol requires exactly 21 scenarios on a 145x192 grid")
    latitude = np.linspace(-90.0, 90.0, ny, dtype=np.float64)
    longitude = np.linspace(0.0, 360.0, nx, endpoint=False, dtype=np.float64)
    lat, lon = np.meshgrid(latitude, longitude, indexing="ij")
    polar = np.sin(np.deg2rad(lat)) ** 4
    land_wave = np.cos(np.deg2rad(2.0 * lon - 0.6 * lat)) * np.cos(np.deg2rad(lat)) ** 2
    short_fields, long_fields = [], []
    for index, (scenario_id, amplitude, forcing_type, centre_lat, centre_lon) in enumerate(SCENARIOS):
        local = np.exp(-0.5 * ((lat - centre_lat) / 16.0) ** 2
                       -0.5 * (wrapped_distance(lon, centre_lon) / 28.0) ** 2)
        if forcing_type == "global":
            local = 0.25 * np.cos(np.deg2rad(lat)) ** 2
        phase = 2.0 * np.pi * index / len(SCENARIOS)
        forcing_region = amplitude * local
        global_warming = amplitude * (0.52 + 0.34 * polar)
        circulation = 0.13 * amplitude * np.cos(np.deg2rad(lat * 2.0) + phase) * np.sin(np.deg2rad(lon) - phase)
        short = global_warming + 0.72 * forcing_region + 0.10 * amplitude * land_wave + circulation
        zonal = short.mean(axis=1, keepdims=True)
        remote_east = np.roll(short, nx // 5, axis=1)
        cross_equatorial = np.flip(zonal, axis=0)
        planetary_wave = np.cos(np.deg2rad(lon * 2.0 + centre_lon)) * np.cos(np.deg2rad(lat))
        long = (1.42 * short + 0.28 * remote_east + 0.23 * cross_equatorial
                + short.mean() * (0.48 + 0.62 * polar) + 0.09 * amplitude * planetary_wave)
        short_fields.append(short.astype(np.float32))
        long_fields.append(long.astype(np.float32))
    short_response = np.stack(short_fields)
    long_response = np.stack(long_fields)
    if short_response.shape != (21, 145, 192) or not np.isfinite(long_response).all():
        raise ValueError("invalid synthetic response fields")
    output = ROOT / config["data"]["path"]
    output.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(output, format_version=np.array(FORMAT_VERSION),
                        scenario_ids=np.asarray([item[0] for item in SCENARIOS]),
                        forcing_amplitude=np.asarray([item[1] for item in SCENARIOS], dtype=np.float32),
                        latitude_deg=latitude.astype(np.float32), longitude_deg=longitude.astype(np.float32),
                        short_response=short_response, long_response=long_response)
    print(f"data={output.relative_to(ROOT)} shape={short_response.shape} features={ny * nx}")


if __name__ == "__main__":
    main()