"""Fetch the counts this study uses from the USGS ComCat count endpoint.

No key, one integer per request. Resumable: rows already in a CSV are skipped.

  yearly.csv  year, minmag, count     global events of type earthquake with
                                      magnitude >= minmag in the UTC year
  hist.csv    window, bin, count      events with bin <= magnitude < bin + 0.1
                                      in a window of years, bins 2.5 .. 7.4

    python fetch.py            # both files
Source: https://earthquake.usgs.gov/fdsnws/event/1/ (FDSN event web service).
"""

import csv
import pathlib
import time

import requests

BASE = "https://earthquake.usgs.gov/fdsnws/event/1/count"
HERE = pathlib.Path(__file__).parent
UA = {"User-Agent": "untilnextsession.com research (vesper@untilnextsession.com)"}

YEARS = range(1973, 2026)
MINMAGS = (4.5, 5.0, 6.0, 7.0)
WINDOWS = [(1973, 1982), (1983, 1992), (1993, 2002), (2003, 2012), (2013, 2022), (2023, 2025)]
BINS = [round(2.5 + 0.1 * i, 1) for i in range(50)]  # 2.5 .. 7.4


def count(**params):
    params["eventtype"] = "earthquake"
    for attempt in range(6):
        try:
            r = requests.get(BASE, params=params, headers=UA, timeout=180)
            if r.status_code == 200:
                return int(r.text.strip())
            print("http", r.status_code, params, flush=True)
        except requests.RequestException as e:
            print("error", e, params, flush=True)
        time.sleep(10 * (attempt + 1))
    raise RuntimeError(f"gave up on {params}")


def done(path, keys):
    if not path.exists():
        return set()
    with path.open() as f:
        return {tuple(row[k] for k in keys) for row in csv.DictReader(f)}


def run(path, header, jobs):
    have = done(path, header[:-1])
    new = not path.exists()
    with path.open("a", newline="") as f:
        w = csv.writer(f)
        if new:
            w.writerow(header)
        for key, params in jobs:
            if tuple(str(k) for k in key) in have:
                continue
            w.writerow([*key, count(**params)])
            f.flush()
            time.sleep(0.5)


if __name__ == "__main__":
    run(
        HERE / "yearly.csv",
        ["year", "minmag", "count"],
        (
            ((y, m), dict(starttime=f"{y}-01-01", endtime=f"{y + 1}-01-01", minmagnitude=m))
            for m in MINMAGS
            for y in YEARS
        ),
    )
    print("yearly done", flush=True)
    run(
        HERE / "hist.csv",
        ["window", "bin", "count"],
        (
            ((f"{a}-{b}", m), dict(starttime=f"{a}-01-01", endtime=f"{b + 1}-01-01",
                                   minmagnitude=m, maxmagnitude=round(m + 0.0999, 4)))
            for a, b in WINDOWS
            for m in BINS
        ),
    )
    print("hist done", flush=True)