import csv, os, numpy as np
from datetime import datetime, timezone
os.chdir('/work/aftershock-forecasts')
def hav(lat1, lon1, lat2, lon2):
    lat1, lon1, lat2, lon2 = map(np.radians, (lat1, lon1, lat2, lon2))
    return 2 * 6371.0 * np.arcsin(np.sqrt(np.sin((lat2 - lat1) / 2) ** 2 + np.cos(lat1) * np.cos(lat2) * np.sin((lon2 - lon1) / 2) ** 2))
_cat = {}
def catalog(eid):
    """ComCat events (M>=2.5) near the mainshock, mainshock removed: arrays t_ms, lat, lon, mag."""
    if eid not in _cat:
        p = f'raw/cat/{eid}.csv'
        if not os.path.exists(p): return None
        R = [r for r in csv.DictReader(open(p)) if r['id'] != eid and r['mag']]
        t = np.array([datetime.fromisoformat(r['time'].replace('Z', '+00:00')).timestamp() * 1000 for r in R])
        _cat[eid] = dict(t=t, lat=np.array([float(r['latitude']) for r in R]), lon=np.array([float(r['longitude']) for r in R]),
                         mag=np.array([float(r['mag']) for r in R]), id=[r['id'] for r in R], type=[r['type'] for r in R])
        _cat[eid]['eq'] = np.array([t == 'earthquake' for t in _cat[eid]['type']], bool)
    return _cat[eid]
def count(eid, rlat, rlon, rkm, t0, t1, M, eps=0.05, eq_only=True):
    """ComCat events (type 'earthquake' unless eq_only=False) with mag >= M - eps in the circle and t0 <= t < t1."""
    c = catalog(eid)
    sel = (c['t'] >= t0) & (c['t'] < t1) & (c['mag'] >= M - eps) & (hav(rlat, rlon, c['lat'], c['lon']) <= rkm)
    if eq_only: sel &= c['eq']
    return int(sel.sum())