#!/usr/bin/env python3
"""Reproduce the Colorado county report-date comparison; Python 3 standard library only.
Run: python3 reproduce.py [--verify-upstream /path/to/original.csv.gz]
The default calculation is offline. Results are compared with the frozen public CSVs.
"""
from pathlib import Path
import argparse, csv, datetime as dt, decimal, gzip, hashlib, io, json

B = Path(__file__).resolve().parent
D = decimal.Decimal
def read(path):
    with path.open(newline='') as f: return list(csv.DictReader(f))
def subset(rows):
    return [r for r in rows if r['STATE']=='COLORADO' and r['CZ_TYPE']=='C' and int(r['CZ_FIPS']) in (41,123) and r['EVENT_TYPE'] in ('Hail','Thunderstorm Wind')]
def date(r):
    value = dt.datetime.strptime(r['BEGIN_YEARMONTH']+r['BEGIN_DAY'].zfill(2),'%Y%m%d').date()
    assert value == dt.datetime.strptime(r['BEGIN_DATE_TIME'],'%d-%b-%y %H:%M:%S').date()
    assert value.year == 2025 and int(r['STATE_FIPS']) == 8 and r['CZ_TIMEZONE']=='MST-7'
    return value
def qualifies(r, variant):
    magnitude = D(r['MAGNITUDE']) if r['MAGNITUDE'] else D('-Infinity')
    if r['MAGNITUDE']: assert magnitude.is_finite()
    hail = r['EVENT_TYPE']=='Hail' and magnitude >= (D('2') if variant=='LargeHailPlusMeasuredWind' else D('1'))
    wind = r['EVENT_TYPE']=='Thunderstorm Wind' and magnitude >= D('50')
    measured = r['MAGNITUDE_TYPE'] in ('MG','MS')
    return hail or (variant!='HailOnly' and wind and (variant=='Base' or measured))
def encode(rows):
    f=io.StringIO(newline=''); w=csv.DictWriter(f,list(rows[0]),lineterminator='\n'); w.writeheader(); w.writerows(rows); return f.getvalue()
def calculate(rows):
    dates_out=[]; gaps_out=[]; summary=[]; sensitivity=[]
    assert len({r['EVENT_ID'] for r in rows})==len(rows)
    for r in rows: date(r)
    for county,code in [('Weld',123),('El Paso',41)]:
        countyrows=[r for r in rows if int(r['CZ_FIPS'])==code]
        for variant in ['MeasuredWind','Base','HailOnly','LargeHailPlusMeasuredWind']:
            rs=[r for r in countyrows if qualifies(r,variant)]
            ds=sorted({date(r) for r in rs}); gaps=[(b-a).days for a,b in zip(ds,ds[1:])]
            metric={'county':county,'variant':variant,'dates':len(ds),'rows':len(rs),'gaps':len(gaps),**{f'gaps_le{n}':sum(g<=n for g in gaps) for n in (3,7,14,30)}}
            sensitivity.append(metric)
            if variant=='MeasuredWind':
                assert all(r['BEGIN_DATE_TIME'].split()[0]==r['END_DATE_TIME'].split()[0] for r in rs)
                summary.append({'county':county,'qualifying_reports':len(rs),'hail_reports':sum(r['EVENT_TYPE']=='Hail' for r in rs),'measured_wind_reports':sum(r['EVENT_TYPE']=='Thunderstorm Wind' for r in rs),'dates':len(ds),'gaps':len(gaps),'gaps_le7':sum(g<=7 for g in gaps)})
                for i,d in enumerate(ds):
                    dates_out.append({'county':county,'date':d.isoformat(),'previous_gap_days':gaps[i-1] if i else ''})
                gaps_out += [{'county':county,'previous_date':a.isoformat(),'next_date':b.isoformat(),'gap_days':g,'at_most_seven_days':str(g<=7).lower()} for a,b,g in zip(ds,ds[1:],gaps)]
    assert [(r['dates'],r['gaps'],r['gaps_le7']) for r in summary]==[(20,19,14),(17,16,12)]
    return {'summary.csv':summary,'preferred-dates.csv':dates_out,'between-date-gaps.csv':gaps_out,'sensitivity.csv':sensitivity}
def main():
    parser=argparse.ArgumentParser(); parser.add_argument('--verify-upstream',type=Path); args=parser.parse_args()
    manifest=json.loads((B/'source-manifest.json').read_text())
    assert hashlib.sha256((B/'raw-county-reports.csv').read_bytes()).hexdigest()==manifest['county_subset_sha256']
    rows=read(B/'raw-county-reports.csv')
    if args.verify_upstream:
        assert hashlib.sha256(args.verify_upstream.read_bytes()).hexdigest()==manifest['upstream_sha256']
        with gzip.open(args.verify_upstream,'rt',newline='') as f: assert subset(list(csv.DictReader(f)))==rows
    outputs=calculate(rows)
    for name, data in outputs.items():
        assert (B/name).read_text()==encode(data), f'Mismatch: {name}'
    print(json.dumps({'verified':list(outputs),'summary':outputs['summary.csv'],'upstream_verified':bool(args.verify_upstream)},indent=2))
if __name__=='__main__': main()