"""Real CPU-only, seeded binary customer-query routing experiment."""
import csv
import hashlib
import json
import os
import platform
import random
import statistics
import sys
import time
from pathlib import Path

import psutil
import torch
import transformers
import tokenizers
from transformers import AutoModelForCausalLM, AutoTokenizer

ROOT=Path(__file__).resolve().parents[2]
HERE=Path(__file__).resolve().parent
ASSETS=ROOT/'assets/my-model-my-data-my-cloud'
RESULT=ROOT/'evidence/my-model-my-data-my-cloud/results.json'
SEED=30030

def sample():
    rows=list(csv.DictReader((ASSETS/'test.csv').open()))
    groups={'CARD':[],'TRANSFER':[]}
    for idx,row in enumerate(rows):
        category=row['category'].lower()
        card='card' in category
        transfer='transfer' in category
        if card != transfer: groups['CARD' if card else 'TRANSFER'].append((idx,row))
    rng=random.Random(SEED)
    picked=[(label,idx,row) for label in groups for idx,row in rng.sample(groups[label],32)]
    rng.shuffle(picked)
    return picked,{k:len(v) for k,v in groups.items()}

def wilson(correct,n):
    z=1.959963984540054
    p=correct/n; d=1+z*z/n
    centre=(p+z*z/(2*n))/d
    half=z*((p*(1-p)/n+z*z/(4*n*n))**0.5)/d
    return [round(centre-half,4),round(centre+half,4)]

def main():
    torch.set_num_threads(4)
    torch.manual_seed(SEED)
    manifest=json.loads((HERE/'assets.json').read_text())
    dataset_digest=hashlib.sha256((ASSETS/'test.csv').read_bytes()).hexdigest()
    if dataset_digest!=manifest['files'][0]['sha256']:
        raise ValueError('DATASET_HASH_MISMATCH')
    picked,pool=sample()
    model_path=Path((ASSETS/'model-path.txt').read_text().strip())
    model_hashes={name:hashlib.sha256((model_path/name).read_bytes()).hexdigest()
        for name in manifest['model']['sha256']}
    if model_hashes!=manifest['model']['sha256']:
        raise ValueError('MODEL_HASH_MISMATCH')
    load_started=time.perf_counter()
    tokenizer=AutoTokenizer.from_pretrained(model_path,local_files_only=True)
    model=AutoModelForCausalLM.from_pretrained(model_path,local_files_only=True,torch_dtype=torch.float32)
    model.eval()
    load_seconds=time.perf_counter()-load_started
    process=psutil.Process()
    baseline=process.memory_info().rss
    records=[]
    with torch.inference_mode():
        for label,idx,row in picked:
            messages=[{'role':'system','content':'Route a banking support query. Answer with exactly CARD or TRANSFER.'},
                {'role':'user','content':row['text']}]
            prompt=tokenizer.apply_chat_template(messages,tokenize=False,add_generation_prompt=True)
            tokens=tokenizer(prompt,return_tensors='pt')
            start=time.perf_counter()
            generated=model.generate(**tokens,max_new_tokens=4,do_sample=False,pad_token_id=tokenizer.eos_token_id)
            elapsed=time.perf_counter()-start
            out=generated[0,tokens['input_ids'].shape[1]:]
            raw=tokenizer.decode(out,skip_special_tokens=True).strip()
            pred=raw.upper() if raw.upper() in ('CARD','TRANSFER') else 'INVALID'
            records.append({'row_index':idx,'true':label,'pred':pred,'raw':raw,
                'input_tokens':int(tokens['input_ids'].numel()),'output_tokens':int(out.numel()),
                'seconds':round(elapsed,6),'rss_bytes':process.memory_info().rss})
            print(f"{len(records)}/{len(picked)} true={label} pred={pred} seconds={elapsed:.2f}",flush=True)
    correct=sum(x['true']==x['pred'] for x in records)
    total_seconds=sum(x['seconds'] for x in records)
    input_tokens=sum(x['input_tokens'] for x in records)
    output_tokens=sum(x['output_tokens'] for x in records)
    result={'run_id':'30-mmmdmc-20261010-01','seed':SEED,'sample_size':len(records),
        'pool_counts':pool,'accuracy':correct/len(records),'correct':correct,
        'accuracy_wilson95':wilson(correct,len(records)),
        'input_tokens':input_tokens,'output_tokens':output_tokens,
        'inference_seconds':round(total_seconds,6),'model_load_seconds':round(load_seconds,6),
        'median_query_seconds':round(statistics.median(x['seconds'] for x in records),6),
        'total_tokens_per_second':round((input_tokens+output_tokens)/total_seconds,3),
        'output_tokens_per_second':round(output_tokens/total_seconds,3),
        'baseline_rss_bytes':baseline,'peak_rss_bytes':max(x['rss_bytes'] for x in records),
        'hardware':{'machine':platform.machine(),'processor':platform.processor(),
            'cpu_count':os.cpu_count(),'total_memory_bytes':psutil.virtual_memory().total,
            'torch_threads':torch.get_num_threads(),'gpu_used':False},
        'dataset_sha256':dataset_digest,
        'model_revision':manifest['model']['revision'],
        'model_asset_sha256':model_hashes,
        'runtime_versions':{'python':sys.version.split()[0], 'torch':torch.__version__,
            'transformers':transformers.__version__, 'tokenizers':tokenizers.__version__,
            'psutil':psutil.__version__},
        'records':records}
    RESULT.parent.mkdir(parents=True,exist_ok=True)
    RESULT.write_text(json.dumps(result,indent=2)+'\n')
    print(json.dumps({k:v for k,v in result.items() if k!='records'},indent=2))

if __name__=='__main__':main()