"""Generate compact hourly data retaining all 12 stations and six pollutants."""

import sys
from pathlib import Path

import numpy as np
import yaml

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model.meteonorm_rf import BASE_FEATURES, FORMAT_VERSION, POLLUTANTS, STATIONS


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    rng = np.random.default_rng(int(config["seed"]))
    days = int(config["data"]["synthetic_days"])
    hours = np.tile(np.arange(days * 24), len(STATIONS))
    station = np.repeat(np.arange(len(STATIONS)), days * 24)
    day = hours // 24
    hour = hours % 24
    doy = day + int(config["data"]["start_day_of_year"])
    weekend = ((day + 1) % 7 >= 5).astype(np.float32)
    seasonal = np.sin(2 * np.pi * (doy - 172) / 365.25)
    diurnal = np.sin(2 * np.pi * (hour - 14) / 24)
    synoptic = np.sin(2 * np.pi * day / 6.5 + station * 0.17)
    wind_speed = np.maximum(0.2, 2.7 + 1.1 * synoptic + rng.normal(0, .55, len(hours)))
    wind_direction = np.mod(185 + 75 * synoptic + station * 9 + rng.normal(0, 28, len(hours)), 360)
    temperature = 13 + 12 * seasonal + 5.5 * diurnal + rng.normal(0, 1.4, len(hours))
    rh = np.clip(52 - 18 * seasonal - 10 * diurnal - 3 * wind_speed + rng.normal(0, 7, len(hours)), 15, 98)
    pressure = 1013 - .32 * temperature + 4 * np.cos(2 * np.pi * day / 8) + rng.normal(0, 1.8, len(hours))
    ttrend = 2013.0 + hours / (24 * 365.25)
    emission = 1.0 - .22 * day / max(days - 1, 1)
    stagnation = np.exp(-wind_speed / 3.2) * (0.65 + rh / 150)
    rush = ((hour >= 7) & (hour <= 10)) | ((hour >= 17) & (hour <= 21))
    station_factor = 0.82 + station * .025
    pm25 = (22 + 78 * stagnation + 16 * rush) * emission * station_factor
    pm10 = 28 + 1.22 * pm25 + 7 * wind_speed
    no2 = (15 + 37 * stagnation + 19 * rush) * emission * station_factor
    so2 = (4 + 25 * stagnation) * emission ** 1.8 * (1.1 - .15 * seasonal)
    o3 = np.maximum(3, 30 + 29 * np.maximum(diurnal, -.4) + 1.2 * temperature - .28 * no2)
    co = (0.25 + .012 * pm25 + .008 * no2) * emission
    pollution = np.column_stack([pm25, pm10, no2, so2, o3, co])
    pollution *= rng.lognormal(0, .09, pollution.shape)
    output = ROOT / config["data"]["path"]
    output.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(output, format_version=np.array(FORMAT_VERSION),
                        station_names=np.asarray(STATIONS), pollutant_names=np.asarray(POLLUTANTS),
                        station_id=station.astype(np.int16), timestamp_hour=hours.astype(np.int32),
                        ttrend=ttrend.astype(np.float32), day_of_year=doy.astype(np.float32),
                        weekend=weekend, hour=hour.astype(np.float32),
                        wind_speed=wind_speed.astype(np.float32), wind_direction=wind_direction.astype(np.float32),
                        pressure=pressure.astype(np.float32), temperature=temperature.astype(np.float32),
                        relative_humidity=rh.astype(np.float32), pollution=pollution.astype(np.float32),
                        feature_names=np.asarray(BASE_FEATURES))
    print(f"data={output.relative_to(ROOT)} rows={len(hours)} stations=12 pollutants=6 hourly_days={days}")


if __name__ == "__main__":
    main()