"""Generate structured station-day PCA predictors with realistic surge drivers."""

import argparse
from pathlib import Path

import numpy as np
import yaml


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


def correlated_features(rng, latent, count, phase):
    weights = rng.normal(0, 0.35, (latent.shape[1], count))
    weights[0, :12] += np.linspace(1.5, 0.4, 12)
    weights[1, 12:24] += np.linspace(1.2, 0.3, 12)
    features = latent @ weights + rng.normal(0, 0.12, (len(latent), count))
    features += 0.08 * np.sin(np.arange(count)[None, :] * 0.13 + phase)
    return features.astype(np.float32)


def make_split(path, count, config, seed, day_offset):
    rng = np.random.default_rng(seed)
    day = np.arange(count, dtype=np.float32) + day_offset
    latitude = rng.uniform(-62, 62, count).astype(np.float32)
    longitude = rng.uniform(-180, 180, count).astype(np.float32)
    annual = np.sin(2 * np.pi * day / 365.25)
    synoptic = np.zeros(count, np.float32)
    for index in range(1, count):
        synoptic[index] = 0.72 * synoptic[index - 1] + rng.normal(0, 0.7)
    cyclone = np.maximum(synoptic - 0.55, 0) ** 2
    tropical = (np.abs(latitude) < 30).astype(np.float32)
    latent = np.column_stack([synoptic, cyclone, annual, np.cos(np.deg2rad(latitude)), tropical]).astype(np.float32)
    data = config["data"]
    rs_daily = correlated_features(rng, latent, int(data["rs_daily_features"]), 0.1)
    rs_lagged = correlated_features(rng, latent, int(data["rs_lagged_features"]), 0.4)
    ar_daily = correlated_features(rng, latent, int(data["ar_daily_features"]), 0.7)
    ar_lagged = correlated_features(rng, latent, int(data["ar_lagged_features"]), 1.0)
    surge = (0.035 * synoptic + 0.055 * cyclone + 0.012 * annual +
             0.015 * rs_lagged[:, 0] - 0.010 * rs_lagged[:, 13])
    surge += rng.normal(0, 0.012 + 0.010 * tropical, count)
    gtsr = 0.72 * surge + rng.normal(0, 0.025, count)
    timestamps = (np.datetime64("2000-01-01") + day.astype("timedelta64[D]")).astype("datetime64[s]").astype(np.int64)
    np.savez_compressed(
        path, format_version=np.asarray(data["format_version"]), data_source=np.asarray("structured_synthetic"),
        rs_daily=rs_daily, rs_lagged=rs_lagged, ar_daily=ar_daily, ar_lagged=ar_lagged,
        targets_m=surge.astype(np.float32)[:, None], gtsr_m=gtsr.astype(np.float32)[:, None],
        latitude_degrees=latitude, longitude_degrees=longitude, timestamps_unix_s=timestamps,
        feature_normalization=np.asarray("PCA scores centered and scaled per training feature"),
    )


def main():
    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, day_offset) in enumerate((("train.npz", config["data"]["train_samples"], 0),
                                                         ("test.npz", config["data"]["test_samples"], 1000))):
        path = output / name
        if args.force or not path.exists():
            make_split(path, int(count), config, int(config["seed"]) + offset, day_offset)
        print(f"generated={path.relative_to(ROOT)} samples={count} daily=50 lagged=300 target=1")


if __name__ == "__main__":
    main()