"""Generate correlated 113-feature WoFS ensemble storm-track examples."""

import argparse
from pathlib import Path

import numpy as np
import yaml


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


def correlated_block(rng, latent, count, noise, offset):
    projection_rng = np.random.default_rng(1000 + int(offset * 100))
    weights = projection_rng.normal(0, 0.35, (latent.shape[1], count))
    for index in range(min(latent.shape[1], count)):
        weights[index, index::latent.shape[1]] += 0.8
    values = latent @ weights + rng.normal(0, noise, (len(latent), count))
    values += 0.08 * np.sin(np.arange(count)[None, :] * 0.21 + offset)
    return values.astype(np.float32)


def make_split(path, samples, config, seed):
    rng = np.random.default_rng(seed)
    lead_start = rng.choice(np.arange(0, 125, 5), samples)
    lead_group = (lead_start > 60).astype(np.int64)
    organization = rng.normal(size=samples)
    rotation = 0.65 * organization + rng.normal(0, 0.75, samples)
    instability = rng.normal(size=samples)
    shear = 0.35 * instability + rng.normal(0, 0.9, samples)
    cold_pool = 0.45 * organization + rng.normal(0, 0.85, samples)
    spread = np.maximum(0.15, rng.lognormal(-0.2 + 0.2 * lead_group, 0.35, samples))
    latent = np.column_stack((organization, rotation, instability, shear, cold_pool, spread))
    amplitude = correlated_block(rng, latent, 30, 0.32, 0.0)
    spatial = correlated_block(rng, latent, 76, 0.48 + 0.08 * lead_group[:, None], 0.7)
    area = np.exp(0.45 * organization + rng.normal(5.8, 0.35, samples))
    eccentricity = 1 / (1 + np.exp(-(0.5 * shear + rng.normal(0, 0.6, samples))))
    major = np.sqrt(area) * (1.2 + eccentricity)
    minor = area / np.maximum(major, 1)
    orientation = np.arctan2(spatial[:, 4], spatial[:, 3]) / np.pi
    extent = np.clip(0.72 - 0.15 * spread + rng.normal(0, 0.08, samples), 0.15, 1)
    initialization_time = rng.uniform(0, 1, samples)
    object_properties = np.column_stack((np.log1p(area), eccentricity, orientation, major / 100,
                                         minor / 100, extent, initialization_time)).astype(np.float32)
    features = np.concatenate((amplitude, spatial, object_properties), axis=1).astype(np.float32)
    risk = np.column_stack((1.25 * rotation + 0.70 * shear + 0.35 * organization,
                            1.15 * instability + 0.55 * organization + 0.45 * amplitude[:, 5],
                            1.00 * cold_pool + 0.65 * shear + 0.35 * spatial[:, 12]))
    risk -= lead_group[:, None] * np.array((0.35, 0.22, 0.18))
    # Elevated synthetic rates make the tiny engineering dataset trainable; paper rates were about 1.2/2.5/4%.
    intercept = np.array((-2.35, -1.75, -1.45))
    probabilities = 1 / (1 + np.exp(-(risk + intercept)))
    targets = (rng.random((samples, 3)) < probabilities).astype(np.float32)
    np.savez_compressed(path, features=features, targets=targets, lead_group=lead_group,
                        lead_start_minutes=lead_start.astype(np.int16),
                        lead_end_minutes=(lead_start + 30).astype(np.int16),
                        format_version=np.asarray(config["data"]["format_version"]),
                        feature_group_sizes=np.asarray((30, 76, 7), dtype=np.int16),
                        hazards=np.asarray(config["data"]["hazards"]),
                        lead_group_names=np.asarray(config["data"]["lead_groups"]),
                        ensemble_members=np.asarray(18), grid_spacing_km=np.asarray(3),
                        forecast_window_minutes=np.asarray(30), forecast_interval_minutes=np.asarray(5),
                        data_source=np.asarray("structured_synthetic_elevated_event_rates"),
                        paper_event_rates=np.asarray((0.012, 0.025, 0.040), dtype=np.float32))
    return targets.mean(0)


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) in enumerate((("train.npz", config["data"]["train_samples"]),
                                             ("test.npz", config["data"]["test_samples"]))):
        path = output / name
        if args.force or not path.exists():
            rates = make_split(path, int(count), config, int(config["seed"]) + offset)
        else:
            rates = np.load(path)["targets"].mean(0)
        print(f"generated={path.relative_to(ROOT)} shape=({count},113) event_rates={rates.round(3).tolist()}")


if __name__ == "__main__":
    main()