"""Flatten raw/fc/*.json into forecasts.csv (one row per version x window x magnitude) and versions.csv (one row per version)."""
import json, glob, os, csv
os.chdir('/work/aftershock-forecasts')
ev = {f['id']: f for f in json.load(open('raw/eventlist.json'))['features']}
V, F = [], []
for path in sorted(glob.glob('raw/fc/*.json')):
    eid, upd, _ = os.path.basename(path).rsplit('_', 2)
    try: d = json.load(open(path))
    except Exception: continue
    p = d['model'].get('parameters', {})
    e = ev[eid]; lon, lat, dep = e['geometry']['coordinates']
    vid = f'{eid}_{upd}'
    obs = {o['magnitude']: o['count'] for o in d.get('observations', [])}
    V.append(dict(vid=vid, eid=eid, mag=e['properties']['mag'], magType=e['properties']['magType'], time_ms=e['properties']['time'],
                  lat=lat, lon=lon, depth=dep, net=e['properties']['net'], update_ms=int(upd), creation_ms=d['creationTime'],
                  model=d['model']['name'], advisory=d.get('advisoryTimeFrame'), regionType=p.get('regionType'),
                  rlat=p.get('regionCenterLat'), rlon=p.get('regionCenterLon'), rkm=p.get('regionRadius'), Mcat=p.get('Mcat', p.get('Mc')),
                  magMain=p.get('magMain'), aSigma=p.get('aSigma'), a=p.get('a'), b=p.get('b'), p=p.get('p'), c=p.get('c'),
                  obs3=obs.get(3), obs4=obs.get(4), obs5=obs.get(5), obs6=obs.get(6), obs7=obs.get(7)))
    for w in d['forecast']:
        for b in w['bins']:
            F.append(dict(vid=vid, eid=eid, window=w['label'], t0_ms=w['timeStart'], t1_ms=w['timeEnd'], M=b['magnitude'],
                          prob=b['probability'], median=b.get('median'), lo95=b.get('p95minimum'), hi95=b.get('p95maximum'),
                          fractiles=' '.join(str(x) for x in b.get('fractileValues', []))))
for name, rows in (('versions.csv', V), ('forecasts.csv', F)):
    with open(name, 'w', newline='') as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0])); w.writeheader(); w.writerows(rows)
    print(name, len(rows))