"""CodeSOTA declarative multiple-choice pilot v1. Data field mappings only; no generated code."""
import hashlib,importlib.metadata,json,platform,sys,time
from pathlib import Path

def field(row,path):
 value=row
 for key in path.split('.'):value=value[key]
 return value

def example(row,recipe):
 question=field(row,recipe['question_field']);choices=field(row,recipe['choices_field']);answer=field(row,recipe['answer_field'])
 if not isinstance(question,str) or not 1<=len(question)<=12000:raise ValueError('Invalid question')
 if not isinstance(choices,list) or not 2<=len(choices)<=10 or not all(isinstance(c,str) and 0<len(c)<=5000 for c in choices):raise ValueError('Invalid choices')
 kind=recipe['answer_type']
 if kind=='index':target=int(answer)
 elif kind=='label':target=[str(x) for x in field(row,recipe['choice_labels_field'])].index(str(answer))
 elif kind=='text':target=choices.index(str(answer))
 else:raise ValueError('Unknown answer mapping')
 if not 0<=target<len(choices):raise ValueError('Answer outside choices')
 return question,choices,target

def rows(task):
 from huggingface_hub import hf_hub_download
 import pyarrow.parquet as pq
 file=hf_hub_download(task['dataset'],task['recipe']['data_file'],repo_type='dataset',revision=task['dataset_revision'])
 for batch in pq.ParquetFile(file).iter_batches(batch_size=128):
  yield from batch.to_pylist()

def evaluate(job,output):
 import numpy as np
 import torch
 from transformers import AutoTokenizer,AutoModelForCausalLM
 task=job['task'];torch.manual_seed(42);np.random.seed(42);torch.set_num_threads(4)
 device='cuda:0' if torch.cuda.is_available() else 'cpu'
 selected=[]
 for i,row in enumerate(rows(task)):
  question,choices,target=example(row,task['recipe']);selected.append((i,question,choices,target))
  if len(selected)==task['limit']:break
 if len(selected)!=task['limit']:raise ValueError('Incomplete pinned dataset')
 start=time.monotonic()
 tokenizer=AutoTokenizer.from_pretrained(job['model'],revision=job['revision'],trust_remote_code=False)
 model=AutoModelForCausalLM.from_pretrained(job['model'],revision=job['revision'],trust_remote_code=False,torch_dtype=torch.float16 if device.startswith('cuda') else torch.float32).to(device).eval()
 samples=[]
 for index,question,choices,target in selected:
  prefix=tokenizer.encode('Question: '+question+'\nAnswer:',add_special_tokens=False)[-1800:]
  scores=[]
  for answer in choices:
   suffix=tokenizer.encode(' '+answer,add_special_tokens=False)
   if len(suffix)>512:raise ValueError('Choice exceeds token budget')
   ids=torch.tensor([prefix+suffix],device=device)
   with torch.inference_mode():logits=model(ids).logits[0,len(prefix)-1:-1].float()
   scores.append(torch.log_softmax(logits,dim=-1).gather(1,ids[0,len(prefix):,None]).mean().item())
  pred=int(np.argmax(scores));samples.append(dict(id=str(index),reference=target,prediction=pred,score=float(pred==target),choices=choices,log_likelihood_mean=scores))
 result=dict(score=float(np.mean([s['score'] for s in samples])),n=len(samples),model_revision=job['revision'],dataset_revision=task['dataset_revision'],runner_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),seconds=time.monotonic()-start,peak_gpu_bytes=torch.cuda.max_memory_allocated() if torch.cuda.is_available() else 0,selection='first 32 examples in pinned parquet order',seed=42,device=device,protocol=task,measured_cost_usd=None,notes='32-example CodeSOTA declarative multiple-choice pilot. Not an official full-benchmark score.')
 environment=dict(python=sys.version,platform=platform.platform(),torch=torch.__version__,cuda=torch.version.cuda,gpu=torch.cuda.get_device_name() if torch.cuda.is_available() else None,packages={d.metadata['Name']:d.version for d in importlib.metadata.distributions()})
 output.parent.mkdir(parents=True,exist_ok=True);output.write_text(json.dumps(dict(result=result,samples=samples,environment=environment),allow_nan=False))
if __name__=='__main__':evaluate(json.loads(Path(sys.argv[1]).read_text()),Path(sys.argv[2]))