"""Isolated-process CPU policy benchmark and optional dynamic INT8 artifact."""
import argparse
import hashlib
import json
from pathlib import Path
import platform
import statistics
from time import perf_counter, process_time
import warnings
import psutil
import torch
from .features import encode
from .policy import LearnedPolicy
from .synthetic import load
from .train import logits, metrics


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--checkpoint',required=True)
    parser.add_argument('--data',default='datasets/synthetic-v1')
    parser.add_argument('--quantized',action='store_true')
    parser.add_argument('--output',required=True)
    args = parser.parse_args()
    torch.set_num_threads(2)
    torch.set_num_interop_threads(1)
    started = perf_counter()
    with warnings.catch_warnings():
        warnings.simplefilter('ignore',DeprecationWarning)
        policy = LearnedPolicy(args.checkpoint,quantized=args.quantized)
    load_ms = (perf_counter()-started)*1000
    root = Path(args.checkpoint)
    artifact = root/'model.safetensors'
    if args.quantized:
        artifact = root/'model-linear-int8.pt'
        torch.save(policy.model.state_dict(),artifact)
        # Verify serialization with restricted loading; no arbitrary pickle globals.
        state = torch.load(artifact,map_location='cpu',weights_only=True)
        policy.model.load_state_dict(state)
    evaluation = {}
    rows = None
    for split in ['validation','test','novel_wording']:
        rows = load(Path(args.data)/f'{split}.jsonl')
        inputs,a,t,_ = encode(rows,policy.vocab)
        la,lt = logits(policy.model,inputs)
        evaluation[split] = metrics(la,lt,a,t,policy.temperatures)
    wall, model_ms, cpu_ms, encode_ms = [], [], [], []
    process = psutil.Process()
    observed_rss = process.memory_info().rss
    with torch.inference_mode():
        for index,row in enumerate(rows[:240]):
            start,cpu = perf_counter(),process_time()
            inputs,*_ = encode([row],policy.vocab)
            encoded = perf_counter()
            policy.model(*inputs)
            finished = perf_counter()
            if index>=20:
                wall.append((finished-start)*1000)
                model_ms.append((finished-encoded)*1000)
                encode_ms.append((encoded-start)*1000)
                cpu_ms.append((process_time()-cpu)*1000)
            observed_rss = max(observed_rss,process.memory_info().rss)
    def stats(values):
        return dict(median=statistics.median(values),p95=sorted(values)[int(.95*len(values))])
    report = dict(checkpoint=args.checkpoint,quantization='dynamic INT8 Linear only; FP32 embeddings/encoder' if args.quantized else 'FP32',
        platform=platform.platform(),torch_version=torch.__version__,threads=2,
        load_ms=load_ms,disk_bytes=artifact.stat().st_size,
        artifact_sha256=hashlib.sha256(artifact.read_bytes()).hexdigest(),
        observed_process_rss_bytes=observed_rss,
        memory_scope='Observed Python RSS including training-library imports and evaluation tensors, not browser or exact peak',
        end_to_end_policy_ms=stats(wall),neural_forward_ms=stats(model_ms),
        feature_encoding_ms=stats(encode_ms),python_cpu_ms=stats(cpu_ms),
        evaluation=evaluation,target_vps_validated=False,
        scope='Synthetic single-step benchmark. Two torch threads, not two-vCPU CPU affinity or target EPYC.')
    Path(args.output).parent.mkdir(parents=True,exist_ok=True)
    Path(args.output).write_text(json.dumps(report,indent=2),encoding='utf-8')
    print(json.dumps(report,indent=2))


if __name__ == '__main__':
    main()