#!/usr/bin/env python3
import argparse
import json
from pathlib import Path

import numpy as np

parser = argparse.ArgumentParser(description="Generate deterministic hydrometeorological smoke-test data")
parser.add_argument("--config", default="conf/config.yaml")
args = parser.parse_args()
with open(args.config, encoding="utf-8") as handle:
    config = json.load(handle)

rng = np.random.default_rng(config["seed"])
data = config["data"]
gauges, train_n, val_n = len(data["gauges"]), data["train_samples"], data["validation_samples"]
cases, leads, history, variables = data["forecast_cases"], data["forecast_steps"], 28, 23
means = np.array([0.37, 8.24, 20.5, 3.28, 4.47, 49.0, 26.0, 0.14, 1.19, 8.55], dtype=np.float32)


def forcing(gauge, count, phase_offset=0.0):
    phase = rng.uniform(0, 2 * np.pi, (count, 1)) + phase_offset
    t = np.arange(history, dtype=np.float32)[None] / 4.0
    x = rng.normal(0, 0.35, (count, history, variables)).astype(np.float32)
    temperature = 9 + 12 * np.sin(2 * np.pi * (t / 365 + phase / (2 * np.pi))) - gauge * 0.35
    precipitation = rng.gamma(1.2, 0.8, (count, history)) * (0.5 + gauge / 18)
    x[:, :, 2] = temperature
    x[:, :, 5] = precipitation
    x[:, :, 6] = np.maximum(0, precipitation * 0.55 + rng.normal(0, 0.1, precipitation.shape))
    x[:, :, 10:14] = np.clip(0.45 + np.cumsum(precipitation[:, :, None] * 0.004, axis=1), 0, 1)
    response = means[gauge] + means[gauge] * (0.08 * precipitation[:, -8:].mean(1) + 0.025 * x[:, -1, 6])
    response += rng.normal(0, max(float(means[gauge]) * 0.03, 0.02), count)
    raw = np.maximum(0, response * (1.12 - gauge * 0.012) + rng.normal(0, max(float(means[gauge]) * 0.05, 0.03), count))
    x[:, :, 20], x[:, :, 21], x[:, :, 22] = raw[:, None], (0.92 * raw)[:, None], means[gauge]
    return x, np.maximum(0, response).astype(np.float32)


train_x = np.empty((gauges, train_n, history, variables), np.float32)
train_y = np.empty((gauges, train_n), np.float32)
val_x = np.empty((gauges, val_n, history, variables), np.float32)
val_y = np.empty((gauges, val_n), np.float32)
forecast_x = np.empty((gauges, cases, leads, history, variables), np.float32)
forecast_y = np.empty((gauges, cases, leads), np.float32)
persistence = np.empty_like(forecast_y)
glofas = np.empty_like(forecast_y)
for g in range(gauges):
    train_x[g], train_y[g] = forcing(g, train_n)
    val_x[g], val_y[g] = forcing(g, val_n, 0.3)
    for case in range(cases):
        base_x, base_y = forcing(g, 1, case / 7)
        initial = float(base_y[0])
        persistence[g, case] = initial
        state = initial
        for lead in range(leads):
            window, target = forcing(g, 1, case / 7 + lead / 80)
            # Forecast meteorology evolves, while observed flow is frozen at issue time.
            window[:, :, 20] = initial
            window[:, :, 21] *= 1.0 + 0.003 * lead
            state = 0.75 * state + 0.25 * float(target[0])
            forecast_x[g, case, lead] = window[0]
            forecast_y[g, case, lead] = state
            glofas[g, case, lead] = max(0, state * (1.10 - 0.002 * lead) + rng.normal(0, max(float(means[g]) * 0.06, 0.03)))

path = Path(data["path"])
path.parent.mkdir(parents=True, exist_ok=True)
np.savez_compressed(path, train_x=train_x, train_y=train_y, val_x=val_x, val_y=val_y,
                    forecast_x=forecast_x, forecast_y=forecast_y, persistence=persistence,
                    glofas=glofas, gauges=np.array(data["gauges"]), lead_hours=np.arange(1, leads + 1) * 6)
print(f"wrote {path}: train={train_x.shape}, forecast={forecast_x.shape}")