"""Download the P-Alert public dataset used to train the intensity GMM.

Stages (all cached, safe to re-run):
  1. events   -> data/raw/events.json
  2. stations -> data/raw/stations.json
  3. per-event PGAList -> data/cache/pga/<event>.json
  4. merged table -> data/dataset.csv

Usage:
  python fetch_all.py                 # full 2024-01-01 .. today
  python fetch_all.py --from 2024-01-01 --to 2024-06-30
  python fetch_all.py --delay 0.4
"""
import argparse
import json
import os
import time

import pandas as pd

from palert.client import fetch_events, fetch_pga, fetch_stations
from palert.paths import CACHE_DIR, DATA_DIR, RAW_DIR

PGA_CACHE = os.path.join(CACHE_DIR, "pga")
os.makedirs(PGA_CACHE, exist_ok=True)

EVENTS_PATH = os.path.join(RAW_DIR, "events.json")
STATIONS_PATH = os.path.join(RAW_DIR, "stations.json")
DATASET_PATH = os.path.join(DATA_DIR, "dataset.csv")


def _load_json(path):
    if os.path.exists(path):
        with open(path, encoding="utf-8") as fh:
            return json.load(fh)
    return None


def _save_json(path, obj):
    tmp = path + ".tmp"
    with open(tmp, "w", encoding="utf-8") as fh:
        json.dump(obj, fh, ensure_ascii=False)
    os.replace(tmp, path)


def event_key(dt):
    return dt.replace(":", "").replace("-", "").replace("T", "_")


def get_events(date_from, date_to):
    cached = _load_json(EVENTS_PATH)
    if cached is not None:
        print(f"[events] cached: {len(cached)} events")
        return cached
    print(f"[events] querying {date_from} .. {date_to}")
    evs = fetch_events(date_from, date_to)
    evs = [e for e in evs if e.get("hasPGA")]
    evs.sort(key=lambda e: e["DateUTC"])
    _save_json(EVENTS_PATH, evs)
    print(f"[events] saved {len(evs)} events with PGA")
    return evs


def get_stations():
    cached = _load_json(STATIONS_PATH)
    if cached is not None:
        print(f"[stations] cached: {len(cached)} stations")
        return cached
    print("[stations] querying onlineAll")
    stas = fetch_stations("onlineAll")
    _save_json(STATIONS_PATH, stas)
    print(f"[stations] saved {len(stas)} stations")
    return stas


def get_pga(event_dt, delay):
    path = os.path.join(PGA_CACHE, event_key(event_dt) + ".json")
    if os.path.exists(path):
        with open(path, encoding="utf-8") as fh:
            return json.load(fh)
    rows = fetch_pga(event_dt)
    _save_json(path, rows)
    if delay:
        time.sleep(delay)
    return rows


def build_dataset(events, stations, delay):
    sta_meta = {s["station"]: s for s in stations}
    records = []
    total = len(events)
    for i, ev in enumerate(events, 1):
        dt = ev["DateUTC"]
        rows = get_pga(dt, delay)
        for r in rows:
            rec = {
                "event": dt,
                "staCode": r.get("staCode"),
                "M": ev["ML"],
                "depth": ev["Depth"],
                "ev_lat": ev["Latitude"],
                "ev_lon": ev["Longitude"],
                "dist_km": r.get("dist"),
                "az": r.get("az"),
                "pga": r.get("pga"),
                "pgv": r.get("pgv"),
            }
            sm = sta_meta.get(r.get("staCode"))
            if sm is not None:
                rec["t_lat"] = sm.get("lat")
                rec["t_lon"] = sm.get("lon")
                rec["area"] = sm.get("area")
            records.append(rec)
        if i % 25 == 0 or i == total:
            print(f"  [pga] {i}/{total} events, {len(records)} rows")
    df = pd.DataFrame.from_records(records)
    return df


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--from", dest="date_from", default="2024-01-01")
    ap.add_argument("--to", dest="date_to", default=time.strftime("%Y-%m-%d"))
    ap.add_argument("--delay", type=float, default=0.3, help="seconds between PGAList calls")
    args = ap.parse_args()

    events = get_events(args.date_from, args.date_to)
    stations = get_stations()
    df = build_dataset(events, stations, args.delay)

    df = df.dropna(subset=["pga", "pgv", "t_lat", "t_lon", "dist_km", "az"])
    # -12345 is the server's "no geometry" sentinel for dist/az
    df = df[(df["dist_km"] >= 0) & (df["az"] >= 0)]
    df = df[(df["pga"] > 0) & (df["pgv"] > 0)]
    df = df.drop_duplicates(subset=["event", "staCode"])
    df.to_csv(DATASET_PATH, index=False)
    print(f"[dataset] {len(df)} rows, {df['event'].nunique()} events -> {DATASET_PATH}")
    print(df.head())


if __name__ == "__main__":
    main()