"""Write a lightweight lazy-data protocol and content-determining coefficients."""

import argparse
from pathlib import Path

import numpy as np
import yaml


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


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())
    data, generator = config["data"], config["generator"]
    output = ROOT / data["root"] / data["protocol_file"]
    output.parent.mkdir(parents=True, exist_ok=True)
    if output.exists() and not args.force:
        print(f"exists={output.relative_to(ROOT)}")
        return
    rng = np.random.default_rng(int(config["seed"]))
    experiments = len(data["experiments"])
    growth = float(generator["growth_rate_per_hour"]) * np.asarray(generator["experiment_growth_multipliers"])
    np.savez_compressed(
        output,
        format_version=np.asarray(data["format_version"]),
        generation_method=np.asarray("deterministic_structured_lazy_low_rank_v1"),
        experiments=np.asarray(data["experiments"]),
        times_hours=np.arange(int(data["time_steps"]), dtype=np.int16) * int(data["step_hours"]),
        pressure_levels_hpa=np.asarray(data["pressure_levels_hpa"], dtype=np.int16),
        latitudes_degrees=np.linspace(90.0, -90.0, int(data["latitude_points"])),
        longitudes_degrees=np.linspace(0.0, 360.0, int(data["longitude_points"]), endpoint=False),
        field_shape=np.asarray(data["field_shape"], dtype=np.int64),
        spectral_shape=np.asarray(data["spectral_shape"], dtype=np.int64),
        ensemble_size=np.asarray(data["ensemble_size"]),
        triangular_truncation=np.asarray(data["triangular_truncation"]),
        nsp=np.asarray(data["nsp"]),
        phase=rng.uniform(0, 2 * np.pi, experiments),
        amplitude=float(generator["perturbation_ms"]) * np.asarray(generator["compensation_scales"]),
        growth_rate=growth,
        spectral_phase=rng.uniform(0, 2 * np.pi, experiments),
        spectral_amplitude=rng.uniform(0.8, 1.2, experiments) * 1e-5,
        base_wind_ms=np.asarray(generator["base_wind_ms"]),
        spectral_slope=np.asarray(generator["spectral_slope"]),
        coefficient_seed=np.asarray(config["seed"]),
        raw_fields_materialized=np.asarray(False),
    )
    print(f"protocol={output.relative_to(ROOT)} bytes={output.stat().st_size} logical_field={data['field_shape']} logical_spectral={data['spectral_shape']}")


if __name__ == "__main__":
    main()