"""For each mainshock, fetch ComCat events (M>=2.5) within the largest forecast radius used for it, from 1 min before
the mainshock to now. Writes raw/cat/<eid>.csv."""
import csv, os, time, urllib.request, collections
from concurrent.futures import ThreadPoolExecutor
os.chdir('/work/aftershock-forecasts'); os.makedirs('raw/cat', exist_ok=True)
from math import radians, sin, cos, asin, sqrt
def km(a, b, c, d):
    a, b, c, d = map(radians, (a, b, c, d))
    return 2 * 6371 * asin(sqrt(sin((c - a) / 2) ** 2 + cos(a) * cos(c) * sin((d - b) / 2) ** 2))
R = collections.defaultdict(float); info = {}
for v in csv.DictReader(open('versions.csv')):
    if v['rkm']:
        off = km(float(v['lat']), float(v['lon']), float(v['rlat']), float(v['rlon']))
        R[v['eid']] = max(R[v['eid']], float(v['rkm']) + off); info[v['eid']] = (v['lat'], v['lon'], int(v['time_ms']))
def go(eid):
    p = f'raw/cat/{eid}.csv'
    if os.path.exists(p): return
    lat, lon, t = info[eid]
    out = b''
    now = time.time()
    t0 = t / 1000 - 60
    while t0 < now:   # one-year chunks, each paged 20000 at a time; any failed request aborts this event
        t1 = min(t0 + 365.25 * 86400, now + 86400)
        st = time.strftime('%Y-%m-%dT%H:%M:%S', time.gmtime(t0)); en = time.strftime('%Y-%m-%dT%H:%M:%S', time.gmtime(t1))
        url = (f'https://earthquake.usgs.gov/fdsnws/event/1/query?format=csv&latitude={lat}&longitude={lon}'
               f'&maxradiuskm={R[eid] + 1:.1f}&starttime={st}&endtime={en}&minmagnitude=2.5&orderby=time-asc&limit=20000')
        off = 1
        while True:
            data = None
            for i in range(5):
                try:
                    data = urllib.request.urlopen(url + f'&offset={off}', timeout=300).read(); break
                except Exception:
                    time.sleep(5 * (i + 1))
            if data is None: raise RuntimeError(f'fetch failed: {eid} {st} offset {off}')
            lines = data.splitlines(True)
            out += data if not out else b''.join(lines[1:])
            if len(lines) - 1 < 20000: break
            off += 20000
        t0 = t1
    open(p, 'wb').write(out)
with ThreadPoolExecutor(3) as ex: list(ex.map(go, list(info)))
print(len(os.listdir('raw/cat')))