from pathlib import Path
import sys

import numpy as np

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from model.pbl_emulator import INPUT_NAMES, OUTPUT_NAMES, load_yaml


def main() -> None:
    config = load_yaml(ROOT / "conf/config.yaml")
    count, seed = int(config["data"]["samples"]), int(config["data"]["seed"])
    interval_hours = int(config["data"]["interval_hours"])
    rng = np.random.default_rng(seed)
    hours = (np.arange(count) * interval_hours) % 24
    days = (np.arange(count) * interval_hours / 24) % 365
    daylight = np.maximum(0.0, np.sin(np.pi * (hours - 6) / 12))
    season = np.sin(2 * np.pi * (days - 172) / 365)
    swdown = (760 + 110 * season) * daylight
    t2 = 284 + 10 * season + 5 * np.sin(2 * np.pi * (hours - 14) / 24) + rng.normal(0, 0.8, count)
    q2 = np.clip(0.0065 + 0.004 * (season + 1) / 2 + 0.0012 * daylight + rng.normal(0, 0.0003, count), 0.001, None)
    ug = 5.5 + 1.8 * np.sin(2 * np.pi * days / 9) + rng.normal(0, 0.5, count)
    vg = 2.0 + 1.5 * np.cos(2 * np.pi * days / 13) + rng.normal(0, 0.5, count)
    u10 = 0.62 * ug + rng.normal(0, 0.45, count); v10 = 0.62 * vg + rng.normal(0, 0.45, count)
    hfx = 12 + 150 * daylight * (0.75 + 0.25 * season) + rng.normal(0, 5, count)
    lh = 25 + 110 * daylight * (0.8 + 0.2 * season) + rng.normal(0, 5, count)
    pblh = np.clip(170 + 1200 * daylight * (0.75 + 0.25 * season) + rng.normal(0, 45, count), 80, 1800)
    ust = np.clip(0.12 + 0.0014 * hfx + 0.022 * np.hypot(u10, v10), 0.05, 1.2)
    tsk = t2 + 1.5 + 4.0 * daylight; tslb = 283 + 7.0 * season
    smois = np.clip(0.25 + 0.04 * np.sin(2 * np.pi * days / 40) - 0.025 * season + rng.normal(0, 0.008, count), 0.08, 0.45)
    grdflx = 0.12 * hfx - 8 * (1 - daylight); glw = 285 + 2.2 * (t2 - 273.15) + 450 * q2
    x = np.column_stack((q2, t2, u10, v10, grdflx, swdown, glw, lh, hfx, pblh, ust, tsk, tslb, smois, ug, vg)).astype(np.float32)
    heights = np.array([10, 35, 70, 110, 160, 220, 290, 370, 460, 560, 680, 820, 980, 1160, 1370, 1600, 1880], dtype=np.float32)
    z = heights[None, :]; mix = np.exp(-z / pblh[:, None])
    u = ug[:, None] + (u10 - ug)[:, None] * mix + 0.18 * np.sin(z / 260)
    v = vg[:, None] + (v10 - vg)[:, None] * mix + 0.14 * np.cos(z / 300)
    w = 0.20 * daylight[:, None] * np.sin(np.pi * np.minimum(z / pblh[:, None], 1)) - 0.015 * (1 - daylight[:, None]) + rng.normal(0, 0.008, (count, 17))
    tk = t2[:, None] - 0.0062 * z + 0.9 * mix * daylight[:, None]
    qv = np.clip(q2[:, None] * np.exp(-z / 2100) * (0.88 + 0.12 * mix) + 1e-4 * smois[:, None], 1e-5, None)
    y = np.stack((u, v, w, tk, qv), axis=-1).astype(np.float32)
    y += rng.normal(size=y.shape).astype(np.float32) * np.array([0.035, 0.035, 0.003, 0.04, 0.000015], dtype=np.float32)
    n_train, n_val = int(0.70 * count), int(0.15 * count)
    slices = {"train": slice(0, n_train), "val": slice(n_train, n_train + n_val), "test": slice(n_train + n_val, count)}
    output = ROOT / config["paths"]["data"]; output.parent.mkdir(parents=True, exist_ok=True)
    payload = {"heights_m": heights, "input_names": np.asarray(INPUT_NAMES), "output_names": np.asarray(OUTPUT_NAMES)}
    for name, section in slices.items():
        payload.update({f"x_{name}": x[section], f"y_{name}": y[section], f"time_{name}": (np.arange(count)[section] * interval_hours).astype(np.int64), f"pblh_{name}": pblh[section].astype(np.float32)})
    np.savez_compressed(output, **payload)
    print(f"saved {output}: x={x.shape}, y={y.shape}, split={[len(range(*s.indices(count))) for s in slices.values()]}")


if __name__ == "__main__":
    main()