"""One-command checks for all five visual claim cards."""
import json
import hashlib
import math
import csv
import random
import argparse
import statistics
from pathlib import Path
from economics import compute

ROOT=Path(__file__).resolve().parents[2]
HERE=Path(__file__).resolve().parent
E=ROOT/'evidence/my-model-my-data-my-cloud'

def load():
    return (json.loads((E/'results.json').read_text()),
        json.loads((E/'economics.json').read_text()),json.loads((HERE/'sources.json').read_text()))

def validate_dataset(r):
    path=ROOT/'assets/my-model-my-data-my-cloud/test.csv'
    raw=path.read_bytes();manifest=json.loads((HERE/'assets.json').read_text())
    digest=hashlib.sha256(raw).hexdigest()
    if digest!=r['dataset_sha256'] or digest!=manifest['files'][0]['sha256']:
        raise ValueError('DATASET_HASH_MISMATCH')
    rows=list(csv.DictReader(path.open()))
    groups={'CARD':[],'TRANSFER':[]}
    for idx,row in enumerate(rows):
        card='card' in row['category'].lower();transfer='transfer' in row['category'].lower()
        if card!=transfer:groups['CARD' if card else 'TRANSFER'].append(idx)
    if {k:len(v) for k,v in groups.items()}!=r['pool_counts']:raise ValueError('POOL_COUNT_MISMATCH')
    rng=random.Random(r['seed'])
    picked=[(label,idx) for label in groups for idx in rng.sample(groups[label],32)]
    rng.shuffle(picked)
    if picked != [(x['true'],x['row_index']) for x in r['records']]:raise ValueError('SAMPLE_MISMATCH')

def validate_records(r,full):
    for x in r['records']:
        expected=x['raw'].strip().upper()
        expected=expected if expected in ('CARD','TRANSFER') else 'INVALID'
        if x['pred']!=expected:raise ValueError('PARSER_MISMATCH')
    manifest=json.loads((HERE/'assets.json').read_text())
    if r['model_revision']!=manifest['model']['revision'] or r['model_asset_sha256']!=manifest['model']['sha256']:
        raise ValueError('MODEL_RECORD_MISMATCH')
    if full:
        model_path=ROOT/'assets/my-model-my-data-my-cloud/model'
        for name,digest in manifest['model']['sha256'].items():
            if hashlib.sha256((model_path/name).read_bytes()).hexdigest()!=digest:
                raise ValueError('MODEL_HASH_MISMATCH')

def quality(r,e,s):
    n=len(r['records']);good=sum(x['pred']==x['true'] for x in r['records'])
    z=1.959963984540054;p=good/n;d=1+z*z/n
    centre=(p+z*z/(2*n))/d
    half=z*math.sqrt(p*(1-p)/n+z*z/(4*n*n))/d
    interval=[round(centre-half,4),round(centre+half,4)]
    counts={label:sum(x['true']==label for x in r['records']) for label in ('CARD','TRANSFER')}
    ok=(n==64 and counts=={'CARD':32,'TRANSFER':32} and good==r['correct']
        and math.isclose(p,r['accuracy']) and interval==r['accuracy_wilson95'])
    return ok,f"{good}/{n} = {100*good/n:.1f}% (Wilson 95% {100*r['accuracy_wilson95'][0]:.1f}–{100*r['accuracy_wilson95'][1]:.1f}%)"

def speed(r,e,s):
    records=r['records'];seconds=sum(x['seconds'] for x in records)
    inputs=sum(x['input_tokens'] for x in records)
    outputs=sum(x['output_tokens'] for x in records)
    total=inputs+outputs
    ok=(inputs==r['input_tokens'] and outputs==r['output_tokens']
        and math.isclose(seconds,r['inference_seconds'],abs_tol=0.0001)
        and abs(total/seconds-r['total_tokens_per_second'])<0.001
        and abs(outputs/seconds-r['output_tokens_per_second'])<0.001
        and math.isclose(statistics.median(x['seconds'] for x in records),
            r['median_query_seconds'],abs_tol=0.000001))
    return ok,f"{total} tokens / {seconds:.1f}s = {r['total_tokens_per_second']:.2f} total tokens/s"

def memory(r,e,s):
    peak=max(x['rss_bytes'] for x in r['records']);baseline=r['baseline_rss_bytes']
    ok=peak==r['peak_rss_bytes'] and peak>baseline and peak<r['hardware']['total_memory_bytes']
    return ok,f"max sampled process RSS {peak/2**30:.2f} GiB; after-load {baseline/2**30:.2f} GiB"

def economics(r,e,s):
    expected=compute(r,s)
    ok=expected==e and e['projected_monthly_tokens_at_observed_speed']>0
    return ok,f"assumed 25% use: ${e['illustrative_local_usd_per_m']:.2f}/M combined tokens; fixed ${e['fixed_usd_per_month']:.2f}/month"

def citations(r,e,s):
    facts={x['id']:x for x in s['checked']}
    ok=('All system providers shall ensure' in facts['rbi']['snapshot_excerpt']
        and 'payment systems operated by them' in facts['rbi']['snapshot_excerpt']
        and 'only in India' in facts['rbi']['snapshot_excerpt']
        and 'by notification, restrict the transfer' in facts['dpdp']['snapshot_excerpt']
        and 'does not impose universal localization' in facts['dpdp']['fact']
        and len(s['checked'])==9)
    return ok,'saved RBI provider-scope and DPDP §16 excerpts present; current applicability unchecked'

CHECKS=[('Quality',quality),('Throughput',speed),('Memory',memory),('Cost',economics),('India context',citations)]

def main():
    parser=argparse.ArgumentParser()
    parser.add_argument('--arithmetic-only',action='store_true',
        help='read-only checks of committed results, economics and source excerpts; skips dataset hash/sample verification')
    args=parser.parse_args()
    r,e,s=load();rows=[]
    if not args.arithmetic_only:validate_dataset(r)
    validate_records(r,not args.arithmetic_only)
    print('| Claim | Mark | Actual output |');print('| --- | --- | --- |')
    for name,fn in CHECKS:
        passed,output=fn(r,e,s)
        rows.append({'claim':name,'passed':passed,'output':output})
        print(f'| {name} | {"PASS" if passed else "FAIL"} | {output} |')
    if not args.arithmetic_only:
        (E/'checks.json').write_text(json.dumps(rows,indent=2)+'\n')
        (E/'check-table.md').write_text('| Claim | Mark | Actual output |\n| --- | --- | --- |\n'+''.join(
            f'| {x["claim"]} | {"PASS" if x["passed"] else "FAIL"} | {x["output"]} |\n' for x in rows))
        paths=[E/'results.json',E/'economics.json',HERE/'sources.json',HERE/'checks.py',HERE/'economics.py',HERE/'measure.py']
        (E/'check-run.json').write_text(json.dumps({'run_id':r['run_id'],
            'sha256':{str(p.relative_to(ROOT)):hashlib.sha256(p.read_bytes()).hexdigest() for p in paths},
            'passes':sum(x['passed'] for x in rows),'failures':sum(not x['passed'] for x in rows)},indent=2)+'\n')
    if not all(x['passed'] for x in rows):raise SystemExit(1)

if __name__=='__main__':main()