"""Generate small structured wildfire sequences with physical correlations."""

import argparse
from pathlib import Path

import numpy as np
import yaml


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


def make_split(path, count, config, seed, day_offset):
    rng = np.random.default_rng(seed)
    data_config = config["data"]
    time = int(data_config["sequence_days"])
    height = int(data_config["patch_height"])
    width = int(data_config["patch_width"])
    yy, xx = np.mgrid[:height, :width].astype(np.float32)
    inputs = np.empty((count, time, 25, height, width), dtype=np.float32)
    danger_scores = np.empty(count, dtype=np.float32)
    for sample in range(count):
        center_y = height / 2 + rng.uniform(-3, 3)
        center_x = width / 2 + rng.uniform(-3, 3)
        hotspot = np.exp(-((yy - center_y) ** 2 + (xx - center_x) ** 2) / (2 * rng.uniform(4, 7) ** 2))
        elevation = np.clip(0.25 + 0.018 * yy + 0.012 * xx + rng.normal(0, 0.02, (height, width)), 0, 1)
        slope = np.clip(np.hypot(*np.gradient(elevation)) * 15, 0, 1)
        road = np.clip(np.abs(xx - rng.uniform(5, 20)) / 20, 0, 1)
        water = np.clip(np.abs(yy - (height / 2 + 2 * np.sin(xx / 4))) / 18, 0, 1)
        population = np.exp(-((xx - rng.uniform(5, 20)) ** 2 + (yy - rng.uniform(5, 20)) ** 2) / 60)
        cover_logits = rng.normal(0, 0.8, (10, height, width))
        cover_logits += np.stack([np.sin((xx + index) / (3 + index / 3)) for index in range(10)])
        cover = np.exp(cover_logits - cover_logits.max(axis=0, keepdims=True))
        cover /= cover.sum(axis=0, keepdims=True)
        weather = rng.normal(0, 0.45)
        for day in range(time):
            weather = 0.82 * weather + rng.normal(0, 0.25)
            drying = day / max(time - 1, 1)
            spatial_noise = rng.normal(0, 0.025, (height, width))
            temperature = 0.50 + 0.16 * weather + 0.20 * drying + 0.16 * hotspot + spatial_noise
            wind = 0.32 + 0.12 * weather + 0.10 * hotspot + rng.normal(0, 0.035, (height, width))
            humidity = 0.62 - 0.19 * weather - 0.20 * drying - 0.14 * hotspot + spatial_noise
            precipitation = np.clip(0.30 - 0.13 * weather - 0.18 * drying - 0.10 * hotspot + spatial_noise, 0, 1)
            dewpoint = 0.55 * temperature + 0.40 * humidity
            pressure = 0.55 - 0.06 * weather + 0.02 * hotspot + spatial_noise
            ndvi = np.clip(0.62 - 0.14 * drying - 0.08 * hotspot + 0.08 * cover[2], 0, 1)
            day_lst = np.clip(temperature + 0.10 * hotspot, 0, 1)
            night_lst = np.clip(temperature - 0.16 + 0.04 * hotspot, 0, 1)
            soil_moisture = np.clip(0.58 * humidity + 0.42 * precipitation - 0.10 * drying, 0, 1)
            dynamic = [temperature, wind, humidity, precipitation, dewpoint, pressure,
                       ndvi, day_lst, night_lst, soil_moisture]
            static = [road, water, population, elevation, slope, *cover]
            inputs[sample, day] = np.stack(dynamic + static).astype(np.float32)
        cy, cx = height // 2, width // 2
        latest = inputs[sample, -1, :, cy, cx]
        danger_scores[sample] = (1.7 * latest[0] + 1.1 * latest[1] - 1.5 * latest[2]
                                 - 1.2 * latest[9] - 0.35 * latest[10]
                                 + 0.25 * latest[12] + rng.normal(0, 0.12))
    labels = (danger_scores >= np.median(danger_scores)).astype(np.float32)[:, None]
    timestamps = (np.datetime64("2018-06-01") + (np.arange(count) + day_offset).astype("timedelta64[D]"))
    timestamps = timestamps.astype("datetime64[s]").astype(np.int64)
    latitude = rng.uniform(34.0, 43.0, count).astype(np.float32)
    longitude = rng.uniform(19.0, 30.0, count).astype(np.float32)
    np.savez_compressed(
        path, inputs=inputs, labels=labels, timestamps_unix_s=timestamps,
        coords=np.column_stack((latitude, longitude)).astype(np.float32),
        format_version=np.asarray(data_config["format_version"]),
        data_source=np.asarray("structured_synthetic"), input_layout=np.asarray("BTCHW"),
    )


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


if __name__ == "__main__":
    main()