from pathlib import Path
import math
import yaml
import numpy as np


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


def main():
    cfg = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    rng = np.random.default_rng(cfg["seed"])
    ncol = cfg["data"]["columns"]
    nwin = cfg["data"]["windows_per_column"]
    steps = cfg["data"]["window_steps"]
    levels = cfg["data"]["levels"]
    dt = cfg["data"]["step_hours"] * 3600.0
    total_steps = nwin * steps + 64
    sigma = np.linspace(0.02, 0.995, levels, dtype=np.float32)
    pressure = 100000.0 * sigma
    interfaces = np.linspace(0.0, 100000.0, levels + 1, dtype=np.float32)
    layer_mass = np.diff(interfaces) / 9.80665
    time = np.arange(total_steps + 1, dtype=np.float32)
    states = np.empty((ncol, total_steps + 1, 68), dtype=np.float32)
    surface = np.empty((ncol, total_steps, 3), dtype=np.float32)
    advection = np.empty((ncol, total_steps + 1, 68), dtype=np.float32)
    for c in range(ncol):
        phase = 2.0 * math.pi * c / ncol
        sl0 = 2.82e5 + 1.25e5 * (1.0 - sigma) + 1400.0 * np.sin(math.pi * sigma + phase)
        qt0 = 0.0175 * sigma ** 2.6 + 0.00012 + 0.0004 * np.sin(2 * math.pi * sigma + phase)
        states[c, 0] = np.concatenate((sl0, np.maximum(qt0, 2e-5)))
        solar = np.maximum(0.0, np.sin(2.0 * math.pi * (time[:-1] % 8.0) / 8.0 + phase))
        surface[c, :, 0] = 30.0 + 18.0 * np.sin(2 * math.pi * time[:-1] / 8.0 + phase)
        surface[c, :, 1] = 92.0 + 34.0 * np.sin(2 * math.pi * time[:-1] / 8.0 + phase - 0.5)
        surface[c, :, 2] = 760.0 * solar
        vertical = np.sin(math.pi * sigma)
        for t in range(total_steps + 1):
            wave = math.sin(2 * math.pi * t / 16.0 + phase)
            advection[c, t, :levels] = (0.010 * wave * vertical + rng.normal(0, 0.0003, levels))
            advection[c, t, levels:] = (1.4e-9 * wave * vertical + rng.normal(0, 4e-11, levels))
        for t in range(total_steps):
            adv = 0.5 * (advection[c, t] + advection[c, t + 1])
            sl, qt = np.split(states[c, t], 2)
            shf, lhf, solin = surface[c, t]
            heat = 2.0e-7 * (300000.0 - sl) + 1.0e-5 * shf * sigma + 1.0e-7 * solin * vertical
            moist = 1.4e-7 * (0.012 * sigma ** 2.4 - qt) + 2.5e-11 * lhf * sigma
            next_state = states[c, t] + dt * (adv + np.concatenate((heat, moist)))
            next_state[levels:] = np.maximum(next_state[levels:], 1e-6)
            states[c, t + 1] = next_state
    initial, targets, surf_windows, adv_windows, source = [], [], [], [], []
    for c in range(ncol):
        for w in range(nwin):
            start = w * steps
            initial.append(states[c, start])
            targets.append(states[c, start:start + steps + 1])
            surf_windows.append(surface[c, start:start + steps])
            adv_windows.append(advection[c, start:start + steps + 1])
            source.append(c)
    out = ROOT / cfg["data"]["file"]
    out.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(out, initial=np.array(initial), target=np.array(targets),
                        surface=np.array(surf_windows), advection=np.array(adv_windows),
                        layer_mass=layer_mass, pressure=pressure, source=np.array(source),
                        long_state=states, long_surface=surface, long_advection=advection,
                        dt_seconds=np.float32(dt), format_version=cfg["data"]["format_version"])
    print(f"saved {out}: windows={len(initial)}, T={steps}, levels={levels}")


if __name__ == "__main__":
    main()