"""Deadline-constrained grid-average carbon accounting; Python standard library only."""
import calendar
import csv
import datetime as dt
import gzip
import hashlib
import io
import json
import math
from collections import defaultdict
from fractions import Fraction
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
DAY = dt.timedelta(days=1)
DURATIONS = (1, 6, 12)
WAITS = (0, 2, 6, 12, 24)
POLICIES = ('immediate', 'annual', 'seasonal', 'recent', 'hindsight')

def season(date):
    return ('DJF' if date.month in (12, 1, 2) else
            'MAM' if date.month in (3, 4, 5) else
            'JJA' if date.month in (6, 7, 8) else 'SON')

def json_write(path, value):
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, indent=2, sort_keys=True, allow_nan=False) + '\n')

def merge_records(records):
    unique = {}
    duplicates = 0
    for record in records:
        start = dt.datetime.fromisoformat(record['from'].replace('Z', '+00:00'))
        end = dt.datetime.fromisoformat(record['to'].replace('Z', '+00:00'))
        if start.utcoffset() != dt.timedelta(0) or end.utcoffset() != dt.timedelta(0):
            raise ValueError('Expected UTC timestamps')
        if start.minute not in (0,30) or start.second or start.microsecond or end-start != dt.timedelta(minutes=30):
            raise ValueError('Misaligned interval')
        if start in unique:
            if record != unique[start]:
                raise ValueError('Contradictory duplicate')
            duplicates += 1
        else:
            unique[start] = record
    days = defaultdict(dict)
    invalid = []
    for start, record in sorted(unique.items()):
        if not 2023 <= start.year <= 2025:
            continue
        value = record['intensity'].get('actual')
        slot = start.hour * 2 + start.minute // 30
        if isinstance(value, bool) or not isinstance(value, (int,float)) or not math.isfinite(value) or value < 0:
            invalid.append({'from':record['from'], 'reason':'actual_missing_or_invalid'})
            continue
        exact = Fraction(str(value))
        days[start.date()][slot] = exact.numerator if exact.denominator == 1 else exact
    complete = {date:[slots[i] for i in range(48)] for date,slots in days.items() if len(slots)==48}
    exclusions = []
    date = dt.date(2023,1,1)
    while date < dt.date(2026,1,1):
        if date not in complete:
            exclusions.append({'date':str(date), 'missing_slots':sorted(set(range(48))-days[date].keys())})
        date += DAY
    return complete, {'duplicate_intervals':duplicates,'invalid_intervals':invalid,'incomplete_days':exclusions}

def load_source(root=ROOT):
    records = []
    sources = json.loads((root/'data/SOURCES.json').read_text())
    for source in sources:
        compressed = (root/source['path']).read_bytes()
        assert len(compressed)==source['compressed_bytes']
        assert 'sha256:'+hashlib.sha256(compressed).hexdigest()==source['compressed_sha256']
        raw = gzip.decompress(compressed)
        assert len(raw)==source['bytes']
        assert 'sha256:'+hashlib.sha256(raw).hexdigest()==source['sha256']
        data = json.loads(raw)['data']
        assert len(data)==source['records']
        records.extend(data)
    complete, audit = merge_records(records)
    audit['source_files'] = len(sources)
    return complete,audit

def mean_profile(days):
    if not days:
        raise ValueError('No training days')
    return [sum((Fraction(d[i]) for d in days),Fraction())/len(days) for i in range(48)]

def prior_dates(date, complete):
    return [date-dt.timedelta(days=lag) for lag in range(29,1,-1) if date-dt.timedelta(days=lag) in complete]

def window_sums(values, duration):
    prefix = [0]
    for value in values:
        prefix.append(prefix[-1]+value)
    return [prefix[i+duration]-prefix[i] for i in range(len(values)-duration+1)]

def choose(scores, release, wait):
    # min is stable in candidate order, implementing earliest exact tie breaking.
    return min(range(release,release+wait+1),key=scores.__getitem__)

def new_stats():
    return {'n':0,'total':0,'baseline':0,'oracle':0,'delay':0,'max_delay':0,'lower':0,'equal':0,'higher':0}

def add(stats, value, baseline, oracle, delay):
    stats['n'] += 1
    stats['total'] += value
    stats['baseline'] += baseline
    stats['oracle'] += oracle
    stats['delay'] += delay
    stats['max_delay'] = max(stats['max_delay'],delay)
    stats['lower' if value<baseline else 'higher' if value>baseline else 'equal'] += 1

def rounded(value):
    return round(float(value),6)

def finalize(stats, duration):
    n = stats['n']
    difference = stats['baseline']-stats['total']
    gap = stats['baseline']-stats['oracle']
    return {'jobs':n,
            'mean_intensity_gco2_per_kwh':rounded(stats['total']/n/duration),
            'absolute_reduction_gco2_per_kwh':rounded(difference/n/duration),
            'percent_reduction':rounded(100*difference/stats['baseline']) if stats['baseline'] else None,
            'mean_delay_hours':rounded(stats['delay']/n/2),
            'max_delay_hours':stats['max_delay']/2,
            'lower_jobs':stats['lower'],'equal_jobs':stats['equal'],'higher_jobs':stats['higher'],
            'lower_percent':rounded(100*stats['lower']/n),
            'equal_percent':rounded(100*stats['equal']/n),
            'higher_percent':rounded(100*stats['higher']/n),
            'hindsight_gap_fraction':rounded(difference/gap) if gap>0 else None}

def run(root=ROOT):
    complete,audit = load_source(root)
    training = {d:v for d,v in complete.items() if d.year==2023}
    fixed = {'annual':mean_profile(list(training.values()))}
    fixed.update({s:mean_profile([v for d,v in training.items() if season(d)==s]) for s in ('DJF','MAM','JJA','SON')})
    exclusions = []
    dates = []
    date = dt.date(2024,1,1)
    while date < dt.date(2026,1,1):
        reasons = []
        if date not in complete:reasons.append('release_day_incomplete')
        if date+DAY not in complete:reasons.append('next_day_incomplete_or_outside_source')
        if len(prior_dates(date,complete))<21:reasons.append('fewer_than_21_recent_training_days')
        if reasons:exclusions.append({'date':str(date),'reasons':reasons})
        else:dates.append(date)
        date+=DAY
    coverage = {str(y):{'eligible_release_dates':sum(d.year==y for d in dates),
                        'calendar_dates':366 if calendar.isleap(y) else 365} for y in (2024,2025)}
    for value in coverage.values():
        value['coverage_percent']=rounded(100*value['eligible_release_dates']/value['calendar_dates'])
        assert value['coverage_percent']>=95
    yearly,monthly,slots,blocks = (defaultdict(new_stats) for _ in range(4))
    daily_buffer = io.StringIO()
    fields=['release_date','duration_hours','max_wait_hours','policy',*finalize({'n':1,'total':1,'baseline':1,'oracle':1,'delay':0,'max_delay':0,'lower':0,'equal':1,'higher':0},1).keys()]
    writer=csv.DictWriter(daily_buffer,fieldnames=fields,lineterminator='\n');writer.writeheader()
    samples=[]
    for date in dates:
        observations = complete[date]+complete[date+DAY]
        recent=mean_profile([complete[d] for d in prior_dates(date,complete)])
        profiles={'annual':fixed['annual']*2,
                  'seasonal':fixed[season(date)]+fixed[season(date+DAY)],
                  'recent':recent*2}
        for duration in DURATIONS:
            actual=window_sums(observations,duration)
            predictions={p:window_sums(values,duration) for p,values in profiles.items()}
            for wait in WAITS:
                daily={p:new_stats() for p in POLICIES}
                for release in range(48):
                    starts={'immediate':release,'hindsight':choose(actual,release,wait)}
                    starts.update({p:choose(pred,release,wait) for p,pred in predictions.items()})
                    baseline,oracle=actual[release],actual[starts['hindsight']]
                    assert oracle<=baseline
                    for policy,start in starts.items():
                        assert release<=start<=release+wait
                        args=(actual[start],baseline,oracle,start-release)
                        add(yearly[(str(date.year),duration,wait,policy)],*args)
                        add(monthly[(str(date.year),date.month,duration,wait,policy)],*args)
                        add(slots[(str(date.year),release,duration,wait,policy)],*args)
                        if duration==6 and wait==12:add(blocks[(str(date.year),release//12,policy)],*args)
                        add(daily[policy],*args)
                    # Deterministic verification sample spans months, runtimes, delays and release times.
                    if date.day==15 and release in (0,11,24,47):
                        samples.append({'date':str(date),'duration':duration,'wait':wait,'release':release,
                                        'starts':starts,'cost_sums':{p:float(actual[s]) for p,s in starts.items()}})
                for policy,stats in daily.items():
                    writer.writerow({'release_date':str(date),'duration_hours':duration/2,'max_wait_hours':wait/2,
                                     'policy':policy,**finalize(stats,duration)})
    years={}
    for (year,duration,wait,policy),stats in sorted(yearly.items()):
        years.setdefault(year,{}).setdefault('d'+str(duration),{}).setdefault('w'+str(wait),{})[policy]=finalize(stats,duration)
    for year in years:
        for duration in DURATIONS:
            oracle=[years[year]['d'+str(duration)]['w'+str(w)]['hindsight']['mean_intensity_gco2_per_kwh'] for w in WAITS]
            assert all(a>=b for a,b in zip(oracle,oracle[1:]))
    primary={}
    for year in years:
        x=years[year]['d6']['w12']
        primary[year]={'recent_minus_seasonal_gco2_per_kwh':rounded((yearly[(year,6,12,'recent')]['total']-yearly[(year,6,12,'seasonal')]['total'])/yearly[(year,6,12,'recent')]['n']/6),
                       'policies':x,'release_blocks':{str(b):{p:finalize(blocks[(year,b,p)],6) for p in POLICIES} for b in range(4)}}
    result={'config':{'durations_hours':[d/2 for d in DURATIONS],'max_wait_hours':[w/2 for w in WAITS],
                      'primary_duration_hours':3,'primary_max_wait_hours':6,'arrivals_per_date':48},
            'coverage':coverage,'training_days':{s:sum(season(d)==s for d in training) for s in ('DJF','MAM','JJA','SON')},
            'years':years,'primary':primary,'daily_aggregate_rows':len(dates)*len(DURATIONS)*len(WAITS)*len(POLICIES)}
    json_write(root/'results/R1.json',result)
    display = {'primary':{},'grid':[],'coverage':coverage,'config':result['config']}
    for year,values in primary.items():
        display['primary'][year]={'recent_minus_seasonal_gco2_per_kwh':round(values['recent_minus_seasonal_gco2_per_kwh'],2),
                                 'policies':{p:{k:round(v,2) if isinstance(v,float) else v for k,v in metrics.items()} for p,metrics in values['policies'].items()}}
    for duration in DURATIONS:
        for wait in WAITS:
            display['grid'].append({'duration_hours':duration/2,'max_wait_hours':wait/2,
                                    **{year:{p:round(years[year]['d'+str(duration)]['w'+str(wait)][p]['percent_reduction'],2) for p in POLICIES} for year in years}})
    json_write(root/'results/R2.json',display)
    json_write(root/'results/exclusions.json',{'raw':audit,'release_dates':exclusions})
    json_write(root/'results/monthly.json',{'|'.join(map(str,k)):finalize(v,k[2]) for k,v in sorted(monthly.items())})
    json_write(root/'results/release-slots.json',{'|'.join(map(str,k)):finalize(v,k[2]) for k,v in sorted(slots.items())})
    json_write(root/'results/verification-sample.json',samples)
    json_write(root/'results/profiles.json',{'fixed':{k:[str(x) for x in v] for k,v in fixed.items()}})
    (root/'results/daily.csv.gz').write_bytes(gzip.compress(daily_buffer.getvalue().encode(),mtime=0))
    print(json.dumps({'coverage':coverage,'primary':primary},indent=2))

if __name__=='__main__':run()