"""Reproducible, nonparametric dry-spell probabilities. Run with NumPy and pandas."""
from pathlib import Path
from collections import Counter
import calendar, datetime, gzip, hashlib, json
import numpy as np
import pandas as pd

ROOT = Path(__file__).resolve().parents[1]
DATES = pd.date_range('1961-01-01','2025-12-31',freq='D')
YEARS = DATES.year.to_numpy()
MONTHS = DATES.month.to_numpy()
PERIODS = {'all':(1961,2025),'early':(1961,1990),'recent':(1991,2025)}
SEASONS = {'warm':list(range(4,10)),'all':list(range(1,13)),'cold':[10,11,12,1,2,3],'summer':[6,7,8]}
AGES = [0,3,7,14,21,30,45,60,90]
WINDOWS = [1,7,14,28]
THRESHOLDS = [1,5,10]
SEED = 20260904
BOOTSTRAPS = 1000
DEFAULT = {'station':'ROE00108892','period':'all','season':'summer','threshold':5,'d':30,'w':7}

NAMES = {'ARAD':'Arad','BAIA MARE':'Baia Mare','BUCURESTI FILARET':'București Filaret',
 'BUZAU':'Buzău','CALARASI':'Călărași','CLUJ NAPOCA':'Cluj-Napoca','DROBETA TURNU SEVERIN':'Drobeta-Turnu Severin',
 'TG JIU':'Târgu Jiu','TURNU MAGURELE':'Turnu Măgurele','BACAU':'Bacău','BOTOSANI':'Botoșani',
 'BUCURESTI-BANEASA':'București Băneasa','CARANSEBES':'Caransebeș','CEAHLAU TOACA':'Ceahlău Toaca',
 'CONSTANTA':'Constanța','CRAIOVA':'Craiova','DEVA':'Deva','GALATI':'Galați','IASI':'Iași',
 'MIERCUREA CIUC':'Miercurea Ciuc','OCNA SUGATAG':'Ocna Șugatag','RAMNICU VALCEA':'Râmnicu Vâlcea',
 'ROSIORI DE VEDE':'Roșiori de Vede','SIBIU':'Sibiu','TULCEA':'Tulcea','STEFAN CEL MARE':'Ștefan cel Mare',
 'BISTRITA':'Bistrița','TRAIAN VUIA':'Timișoara Traian Vuia','VARFU OMUL':'Vârful Omu','SULINA':'Sulina'}

def parse_daily(path):
    """Reject quality flags and presumed zero; never interpolate missing days."""
    x = np.full(len(DATES), np.nan)
    source = np.full(len(DATES), '', dtype='U1')
    excluded = Counter()
    seen = set()
    for line in path.read_text().splitlines():
        if line[17:21] != 'PRCP': continue
        year,month = int(line[11:15]),int(line[15:17])
        if year<1961 or year>2025: continue
        assert (year,month) not in seen, 'duplicate station month'
        seen.add((year,month))
        start = (pd.Timestamp(year,month,1)-DATES[0]).days
        for day in range(calendar.monthrange(year,month)[1]):
            raw=line[21+8*day:29+8*day]
            value=int(raw[:5]); measurement,quality,origin=raw[5:8]
            if value == -9999: excluded['missing_recorded']+=1
            elif value<0: excluded['negative']+=1
            elif quality != ' ': excluded['quality_'+quality]+=1
            elif measurement == 'P': excluded['presumed_zero']+=1
            else:
                x[start+day]=value/10
                source[start+day]=origin
    return x,source,dict(excluded)

def age_since_wet(x, threshold):
    """Age at end of current date. Unknown after a gap until a wet day occurs."""
    age=np.full(len(x),-1,dtype=int)
    current=-1
    for i,v in enumerate(x):
        if not np.isfinite(v): current=-1
        elif v>threshold: current=0
        elif current>=0: current+=1
        age[i]=current
    return age

def forward_outcome(x, threshold, window):
    """Outcome after anchor, excluding anchor itself. Require full window."""
    valid=np.isfinite(x)
    wet=valid & (x>threshold)
    cv=np.r_[0,np.cumsum(valid)]
    cw=np.r_[0,np.cumsum(wet)]
    complete=np.zeros(len(x),dtype=bool)
    outcome=np.zeros(len(x),dtype=bool)
    i=np.arange(len(x)-window)
    complete[i]=valid[i] & (cv[i+window+1]-cv[i+1]==window)
    outcome[i]=(cw[i+window+1]-cw[i+1])>0
    return complete,outcome

def interval(values):
    clean=values[np.isfinite(values)]
    if len(clean)<BOOTSTRAPS*.95: return None
    return [round(float(v),6) for v in np.quantile(clean,[.025,.975])]

def write_json(path,obj):
    path.parent.mkdir(parents=True,exist_ok=True)
    path.write_text(json.dumps(obj,ensure_ascii=False,separators=(',',':'),allow_nan=False)+'\n')

def historical_cohort(x, station, source):
    """Dated observations for the fixed opening case, using its complete cohort."""
    threshold,duration=DEFAULT['threshold'],DEFAULT['d']
    first,last=PERIODS[DEFAULT['period']]
    xx=x.copy();xx[(YEARS<first)|(YEARS>last)]=np.nan
    age=age_since_wet(xx,threshold)
    complete,_=forward_outcome(xx,threshold,max(WINDOWS))
    anchors=np.flatnonzero((age==duration)&np.isin(MONTHS,SEASONS[DEFAULT['season']])&complete)
    rows=[]
    for i in anchors:
        future_wet=np.flatnonzero(xx[i+1:i+max(WINDOWS)+1]>threshold)
        wait=int(future_wet[0]+1) if len(future_wet) else None
        rows.append({'anchor':str(DATES[i].date()),'lastWet':str(DATES[i-duration].date()),
                     'totalBeforeMM':round(float(xx[i-duration+1:i+1].sum()),1),
                     'nextWet':str(DATES[i+wait].date()) if wait is not None else None,
                     'rainMM':float(xx[i+wait]) if wait is not None else None,'waitDays':wait})
    totals=[row['totalBeforeMM'] for row in rows]
    return {'station':{'id':station['id'],'name':station['name']},
            'settings':{'period':DEFAULT['period'],'years':[first,last],'season':DEFAULT['season'],
                        'anchorMonths':SEASONS[DEFAULT['season']],'threshold':threshold,'d':duration,'maxWindow':max(WINDOWS)},
            'source':source,'n':len(rows),'anchorYears':int(len(set(YEARS[anchors]))),
            'totalBeforeSummary':{'minimum':min(totals) if totals else None,
                                  'median':round(float(np.median(totals)),2) if totals else None,
                                  'maximum':max(totals) if totals else None,
                                  'zeroRainEpisodes':sum(total==0 for total in totals)},'rows':rows}

def analyze_configuration(x, threshold, period, season, weights):
    first,last=PERIODS[period]
    psel=(YEARS>=first)&(YEARS<=last)
    xx=x.copy(); xx[~psel]=np.nan
    age=age_since_wet(xx,threshold)
    seasonal=np.isin(MONTHS,SEASONS[season])
    anchor_year=(YEARS-first).clip(0,last-first)
    n_years=last-first+1
    rows=[]; draws={}
    common_complete,_=forward_outcome(xx,threshold,max(WINDOWS))
    for window in WINDOWS:
        _,outcome=forward_outcome(xx,threshold,window)
        complete=common_complete
        base=complete & seasonal
        # Whole-calendar-year resampling preserves overlapping forward windows.
        base_n=np.zeros((n_years,12));base_k=np.zeros_like(base_n)
        np.add.at(base_n,(anchor_year[base],MONTHS[base]-1),1)
        np.add.at(base_k,(anchor_year[base],MONTHS[base]-1),outcome[base])
        bn=weights@base_n;bk=weights@base_k
        rates=np.divide(base_k.sum(axis=0),base_n.sum(axis=0),out=np.zeros(12),where=base_n.sum(axis=0)>0)
        bootrates=np.divide(bk,bn,out=np.zeros_like(bk),where=bn>0)
        year_month_rates=np.divide(base_k,base_n,out=np.zeros_like(base_k),where=base_n>0)
        for duration in AGES:
            candidates=(age==duration)&seasonal&psel
            take=candidates&complete
            n=int(take.sum());k=int(outcome[take].sum())
            yn=np.zeros(n_years);yk=np.zeros(n_years);ym=np.zeros((n_years,12))
            np.add.at(yn,anchor_year[take],1)
            np.add.at(yk,anchor_year[take],outcome[take])
            np.add.at(ym,(anchor_year[take],MONTHS[take]-1),1)
            yrs=int((yn>0).sum())
            supported=n>=30 and yrs>=10
            row={'d':duration,'w':window,'n':n,'k':k,'years':yrs,
                 'excludedWindows':int(candidates.sum()-n),'supported':supported}
            if supported:
                boot_n=weights@yn
                boot_p=np.divide(weights@yk,boot_n,out=np.full(BOOTSTRAPS,np.nan),where=boot_n>0)
                boot_month=np.divide(((weights@ym)*bootrates).sum(axis=1),boot_n,out=np.full(BOOTSTRAPS,np.nan),where=boot_n>0)
                expected_by_year=(ym*year_month_rates).sum(axis=1)
                boot_base=np.divide(weights@expected_by_year,boot_n,out=np.full(BOOTSTRAPS,np.nan),where=boot_n>0)
                p=k/n; b=float(expected_by_year.sum()/n);bm=float((ym.sum(axis=0)*rates).sum()/n)
                boundary=k in (0,n)
                row.update(p=round(p,6),ci=None if boundary else interval(boot_p),baseline=round(b,6),
                    delta=round(p-b,6),deltaCI=None if boundary else interval(boot_p-boot_base),
                    baselineMonth=round(bm,6),deltaMonth=round(p-bm,6),
                    deltaMonthCI=None if boundary else interval(boot_p-boot_month),boundary=boundary)
                if not boundary:draws[(duration,window)]=(boot_p,boot_p-boot_base)
            else: row.update(p=None,ci=None,baseline=None,delta=None,deltaCI=None,baselineMonth=None,deltaMonth=None,deltaMonthCI=None,boundary=False)
            rows.append(row)
    # Predefined contrast; this is not the source paper's spline derivative at p99.
    contrast={}
    if (7,1) in draws and (30,1) in draws:
        left=next(r for r in rows if r['d']==7 and r['w']==1)
        right=next(r for r in rows if r['d']==30 and r['w']==1)
        contrast={'from':7,'to':30,'raw':round(right['p']-left['p'],6),
            'rawCI':interval(draws[(30,1)][0]-draws[(7,1)][0]),
            'adjusted':round(right['delta']-left['delta'],6),
            'adjustedCI':interval(draws[(30,1)][1]-draws[(7,1)][1])}
    # Completed interval lengths, including consecutive wet dates (length one).
    wet=np.flatnonzero(np.isfinite(xx)&(xx>threshold))
    durations=[]
    for a,b in zip(wet[:-1],wet[1:]):
        if age[b-1]>=0 and age[b-1]==b-a-1 and seasonal[b]:durations.append(int(b-a))
    tail=None
    if durations:
        p95,p99=np.quantile(durations,[.95,.99],method='higher').astype(int)
        tail={'intervalP95':int(p95),'intervalP99':int(p99),'completeIntervals':len(durations),
            'atOrAboveP99':int(np.sum(np.array(durations)>=p99))}
    return {'rows':rows,'contrast':contrast,'tail':tail}

def main():
    manifest=json.loads((ROOT/'data/manifest.json').read_text())
    out=ROOT/'public/data'
    index={'version':'1.1.0','sourceVersion':(ROOT/'data/raw/ghcnd-version.txt').read_text().splitlines()[0],
        'accessed':manifest['accessed'],'periods':PERIODS,'seasons':SEASONS,'ages':AGES,
        'windows':WINDOWS,'thresholds':THRESHOLDS,'seed':SEED,'bootstrapReplicates':BOOTSTRAPS,
        'default':DEFAULT,
        'stations':[]}
    weights={}
    for p,(a,b) in PERIODS.items():
        n=b-a+1
        weights[p]=np.random.default_rng(SEED+n).multinomial(n,np.full(n,1/n),size=BOOTSTRAPS).astype(float)
    all_daily=[];audit=[];flat=[]
    for station in manifest['stations']:
        x,source,excluded=parse_daily(ROOT/'data/raw'/(station['id']+'.dly'))
        coverage={}
        for p,(a,b) in PERIODS.items():
            valid=np.isfinite(x)&(YEARS>=a)&(YEARS<=b)
            years_complete=int(sum(np.isfinite(x[YEARS==y]).mean()>=.90 for y in range(a,b+1)))
            fraction=float(valid.sum()/((YEARS>=a)&(YEARS<=b)).sum())
            coverage[p]={'valid':int(valid.sum()),'total':int(((YEARS>=a)&(YEARS<=b)).sum()),
                'fraction':round(fraction,6),'completeYears':years_complete,
                'eligible':bool(fraction>=.90 and years_complete>=25)}
        metadata=dict(station,name=NAMES.get(station['name'],station['name'].title()),coverage=coverage,
            sources=dict(Counter(source[source!=''])),excluded=excluded)
        metadata['annualCoverage']=[{'year':int(y),'valid':int(np.isfinite(x[YEARS==y]).sum()),
            'total':int((YEARS==y).sum()),'sourceFlags':dict(Counter(source[(YEARS==y)&(source!='')]))} for y in range(1961,2026)]
        analysis={}
        for p in PERIODS:
            if not coverage[p]['eligible']:continue
            for season in SEASONS:
                for threshold in THRESHOLDS:
                    key=f'{p}/{season}/{threshold}'
                    config=analyze_configuration(x,threshold,p,season,weights[p])
                    analysis[key]=config
                    for r in config['rows']:
                        flat.append(dict(station=station['id'],period=p,season=season,threshold=threshold,**r))
        defaults=analysis.get(f"{DEFAULT['period']}/{DEFAULT['season']}/{DEFAULT['threshold']}")
        if defaults:
            metadata['defaultResult']=next(r for r in defaults['rows'] if r['d']==DEFAULT['d'] and r['w']==DEFAULT['w'])
            metadata['contrast']=defaults['contrast'];metadata['tail']=defaults['tail']
        if station['id']==DEFAULT['station']:
            history=historical_cohort(x,metadata,{'name':'NOAA GHCN-Daily','version':index['sourceVersion'],
                'accessed':index['accessed'],'rawSha256':hashlib.sha256((ROOT/'data/raw'/(station['id']+'.dly')).read_bytes()).hexdigest()})
            assert history['n']==metadata['defaultResult']['n']
            for row in defaults['rows']:
                if row['d']==DEFAULT['d']:
                    assert row['k']==sum(r['waitDays'] is not None and r['waitDays']<=row['w'] for r in history['rows'])
            write_json(out/'history.json',history)
        index['stations'].append(metadata)
        write_json(out/'stations'/(station['id']+'.json'),{'station':metadata,'configurations':analysis})
        audit.append(metadata)
        for date,v,origin in zip(DATES,x,source):
            if np.isfinite(v):all_daily.append([station['id'],date.strftime('%Y-%m-%d'),float(v),origin])
        print(station['id'],metadata['name'],len(analysis),'configs',flush=True)
    index['eligibleDefault']=sum(s['coverage']['all']['eligible'] for s in index['stations'])
    index['validDefaultDays']=sum(s['coverage']['all']['valid'] for s in index['stations'] if s['coverage']['all']['eligible'])
    index['eligibleDefaultEpisodes']=sum(s['defaultResult']['n'] for s in index['stations'] if 'defaultResult' in s)
    write_json(out/'index.json',index)
    write_json(ROOT/'research/data-quality.json',{'rules':{'coverage':.9,'completeYears':25,'minEpisodes':30,'minYears':10},'stations':audit})
    pd.DataFrame(all_daily,columns=['station','date','precipitation_mm','source_flag']).to_csv(out/'daily-clean.csv.gz',index=False,compression={'method':'gzip','mtime':0})
    # Lists are JSON text inside CSV, not an alternative estimate.
    pd.DataFrame(flat).to_csv(out/'estimates.csv',index=False)
    print(json.dumps({k:index[k] for k in ['eligibleDefault','validDefaultDays','eligibleDefaultEpisodes']}),flush=True)

if __name__=='__main__':main()