"""Generate structured annual forcing windows and climate responses."""

import argparse
from pathlib import Path

import numpy as np
import yaml


ROOT = Path(__file__).resolve().parents[1]


def make_split(path: Path, count: int, config: dict, seed: int, test: bool) -> None:
    rng = np.random.default_rng(seed)
    data = config["data"]
    time, height, width = int(data["time_steps"]), int(data["height"]), int(data["width"])
    lat = np.linspace(-90, 90, height, dtype=np.float32)
    lon = np.linspace(0, 360, width, endpoint=False, dtype=np.float32)
    lat2d, lon2d = np.meshgrid(lat, lon, indexing="ij")
    windows, responses, years = [], [], []
    for sample in range(count):
        end_year = 2100 if test else 2060 + 10 * sample
        annual_years = np.arange(end_year - time + 1, end_year + 1, dtype=np.int32)
        progress = (annual_years - 1850).astype(np.float32)
        co2 = 0.90 * progress + 0.0018 * progress ** 2
        ch4 = 0.42 * progress + 12.0 * np.sin(progress / 38.0)
        industrial = np.exp(-((lat2d - 35) / 22) ** 2) * (0.65 + 0.35 * np.cos(np.deg2rad(lon2d - 90)))
        tropical = np.exp(-(lat2d / 20) ** 2) * (0.65 + 0.35 * np.sin(np.deg2rad(2 * lon2d)))
        forcing = np.empty((time, 4, height, width), dtype=np.float32)
        for index, year_progress in enumerate(progress):
            forcing[index, 0] = co2[index]
            forcing[index, 1] = ch4[index]
            forcing[index, 2] = industrial * (32.0 * np.exp(-((year_progress - 125) / 75) ** 2))
            forcing[index, 3] = tropical * (5.0 + 2.0 * np.sin(year_progress / 27.0))
        forcing += rng.normal(0, 0.01, forcing.shape).astype(np.float32)
        co2_level = forcing[-1, 0].mean() / 300.0
        methane = forcing[-3:, 1].mean() / 120.0
        aerosol = forcing[-3:, 2:].mean(axis=(0, 1)) / 25.0
        arctic = 1.0 + 1.3 * (np.abs(lat2d) / 90.0) ** 2
        land_pattern = np.cos(np.deg2rad(2 * lon2d)) * np.cos(np.deg2rad(lat2d))
        itcz = np.sin(np.deg2rad(lon2d)) * np.exp(-(lat2d / 17) ** 2)
        tas = 1.35 * co2_level * arctic + 0.22 * methane - 0.30 * aerosol + 0.05 * land_pattern
        dtr = 0.10 * co2_level + 0.23 * aerosol * land_pattern - 0.04 * methane
        pr = 0.16 * co2_level * np.cos(np.deg2rad(lat2d)) + 0.20 * itcz - 0.07 * aerosol
        pr90 = 1.35 * pr + 0.08 * co2_level * np.exp(-(lat2d / 28) ** 2)
        response = np.stack([tas, dtr, pr, pr90]).astype(np.float32)
        response += rng.normal(0, 0.003, response.shape).astype(np.float32)
        windows.append(forcing)
        responses.append(response)
        years.append(annual_years)
    evaluation = config["evaluation"]
    np.savez_compressed(
        path, inputs=np.stack(windows).astype(np.float32), targets=np.stack(responses).astype(np.float32),
        years=np.stack(years), latitude=lat, longitude=lon,
        channel_names=np.asarray(data["channels"]), target_names=np.asarray(data["targets"]),
        format_version=np.asarray(data["format_version"]), storage_layout=np.asarray("NTCHW"),
        temporal_resolution=np.asarray("annual"), data_source=np.asarray("structured_synthetic"),
        scenario=np.asarray(evaluation["scenario"] if test else "synthetic_training_scenarios"),
        target_aggregation=np.asarray("2080-2100 climatological mean" if test else "annual response"),
        evaluation_start_year=np.asarray(evaluation["start_year"], dtype=np.int32),
        evaluation_end_year=np.asarray(evaluation["end_year"], dtype=np.int32),
    )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--force", action="store_true")
    args = parser.parse_args()
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    output = ROOT / config["data"]["root"]
    output.mkdir(parents=True, exist_ok=True)
    for offset, (name, count, test) in enumerate((("train.npz", config["data"]["train_samples"], False),
                                                  ("test.npz", config["data"]["test_samples"], True))):
        path = output / name
        if args.force or not path.exists():
            make_split(path, int(count), config, int(config["seed"]) + offset, test)
        print(f"generated={path.relative_to(ROOT)} samples={count} layout=NTCHW annual=true")


if __name__ == "__main__":
    main()