"""SinhalaMMLU E4B V2 standalone inference + automatic invalid-answer retry.

CANDIDATE reconstructed from archived research code and original RETRY32.
NOT certified to regenerate byte-identical historical 2,238 predictions.
The exact first-pass 2,206-answer production script is not provided.

Organizers change ONLY DATA_REPO. All model resources come from Hugging Face.
Never access reference predictions or test answer keys during generation.
"""
# ===================== ORGANIZERS CHANGE ONLY THIS LINE =====================
DATA_REPO = "SinhalaMMLU-SharedTask/Sinhala-MMLU-SharedTask-Test-Dataset"
# ===========================================================================

MODEL_REPO = "google/gemma-4-E4B-it"
REVISION = "ee0ef6023621cff504d758262d4e04895a5af4a2"
WEIGHT_SHA256 = "cfbd3d2f1cd71bd471c37fe2bf8546d5028d41e5736f64e1ca6c6b8893125503"
EXPECTED_PARAMS = 7_941_100_832
PROMPT_SHA256 = "f53159bafe22a5d0063a065b6d3aef8167e7a6d6b998ce65e224feeeca68287d"
BATCH_SIZE = 4
FIRST_PASS_TOKENS = 2048
RETRY_TOKENS = 4096
SEED = 2026
OUTPUT_FILE = "predictions.jsonl"
OUTPUT_DIR = "inference_records"

import hashlib
import json
import logging
import os
import random
import re
import time
from pathlib import Path

import numpy as np
import torch
from huggingface_hub import HfFileSystem, list_repo_files, snapshot_download

V2_PROMPT = (
    'Answer this Sinhala multiple-choice question. '
    'Use your reasoning to select the correct option. '
    'In your final response, provide a brief explanation '
    'in English, approximately 90 words. '
    'Then put ANSWER: X on the last line, '
    'where X is A, B, C, or D. '
    'Do not change the selected answer.'
)
assert hashlib.sha256(V2_PROMPT.encode('utf-8')).hexdigest() == PROMPT_SHA256
LETTER_TO_NUMBER = {'A':'1','B':'2','C':'3','D':'4'}
FIELDS = ('q_no','subject','category','level','answer')


def sha_file(path):
    h = hashlib.sha256()
    with open(path, 'rb') as f:
        for block in iter(lambda: f.read(8*1024*1024), b''):
            h.update(block)
    return h.hexdigest()


def save_atomic(path, data):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    temp = path.with_suffix(path.suffix+'.tmp')
    temp.write_text(json.dumps(data,ensure_ascii=False,indent=2)+'\n',encoding='utf-8')
    os.replace(temp,path)


def init_seed():
    random.seed(SEED)
    np.random.seed(SEED)
    torch.manual_seed(SEED)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(SEED)


def load_questions(data_repo):
    """Load every .json file per organizer template; preserve source position."""
    fs = HfFileSystem()
    files = sorted(f for f in list_repo_files(data_repo, repo_type='dataset') if f.endswith('.json'))
    rows = []
    for f in files:
        with fs.open(f'datasets/{data_repo}/{f}','r',encoding='utf-8') as fh:
            raw = json.load(fh)
        if isinstance(raw,dict):
            raw = next(v for v in raw.values() if isinstance(v,list))
        assert isinstance(raw,list),f
        for position,r in enumerate(raw):
            r = dict(r)
            r['_file'] = f
            r['_position'] = position
            rows.append(r)
    if not rows:
        raise ValueError('No questions found in DATA_REPO')
    return rows


def get_level(r):
    return (r.get('metadata') or {}).get('difficulty') or r['_file'].split('/')[0]


def canonical_order(rows):
    """Historical final file: Easy, Medium, Hard; source basename, row order.

    Do not hard-code question count, index or test answers.
    This is deliberately separate from the organizer's loader.
    """
    priority = {'easy':0,'medium':1,'hard':2}
    for row in rows:
        if get_level(row) not in priority:
            raise ValueError(f'Unknown difficulty {get_level(row)}')
        if len(row['choices']) != 4 or not all(isinstance(c,str) for c in row['choices']):
            raise ValueError('Expected exactly four string choices')
    return sorted(rows,key=lambda r:(priority[get_level(r)],Path(r['_file']).name.casefold(),r['_position']))


def fingerprint(r):
    keys = ('q_no','subject','category','question','choices')
    content = {k:r[k] for k in keys}
    content['level'] = get_level(r)
    content['file'] = r['_file']
    content['source_position'] = r['_position']
    return hashlib.sha256(json.dumps(content,ensure_ascii=False,sort_keys=True,separators=(',',':')).encode('utf-8')).hexdigest()


def make_prompt(r):
    opts = '\n'.join(f'{letter}. {choice}' for letter,choice in zip('ABCD',r['choices']))
    return V2_PROMPT + '\n\n' + r['question'] + '\n\n' + opts


def split_tokens(tokenizer, sequence, width, pad_id, stop_ids):
    actual=[]
    ended=False
    for token in sequence[width:].tolist():
        if token == pad_id:
            break
        actual.append(token)
        if token in stop_ids:
            ended=True
            break
    return actual,ended


def parse_retry_strict(raw, ended, token_count, limit):
    """Verbatim logic from RETRY32 for final answer and validity."""
    parts = re.split(r'<\|channel>final|<channel\|>',raw)
    final = re.sub(r'<[^>\n]+>','',parts[-1]).strip()
    lines = [x.strip() for x in final.splitlines() if x.strip()]
    last = lines[-1] if lines else ''
    m = re.fullmatch(r'\*{0,2}\s*ANSWER\s*:\s*([ABCD])\s*[.!]?\s*\*{0,2}',last,flags=re.I)
    pred = m.group(1).upper() if m else None
    hit_limit = not ended and token_count>=limit
    valid = pred in LETTER_TO_NUMBER and ended and token_count>0
    status = ('completed' if valid else 'token_limit' if hit_limit else
              'padding_only' if not token_count else 'unparseable_or_unterminated')
    return pred if valid else None,status,last


def parse_first_pass(raw):
    """Candidate first-pass parser; exact historical 2206 parser not supplied.

    Do not parse thought tokens if final channel missing.
    """
    if ('<|channel>thought' in raw or '<channel|>thought' in raw) and not (
        '<|channel>final' in raw or '<channel|>final' in raw
    ):
        return None
    parts=re.split(r'<\|channel>final|<channel\|>',raw)
    final=re.sub(r'<[^>\n]+>','',parts[-1]).strip()
    lines=[x.strip() for x in final.splitlines() if x.strip()]
    if not lines:
        return None
    m=re.fullmatch(r'\*{0,2}\s*ANSWER\s*:\s*([ABCD])\s*[.!]?\s*\*{0,2}',lines[-1],flags=re.I)
    return m.group(1).upper() if m else None


def model_device_map(model_dir):
    from transformers import AutoConfig, AutoModelForMultimodalLM
    from accelerate import init_empty_weights
    cfg=AutoConfig.from_pretrained(str(model_dir),local_files_only=True)
    with init_empty_weights():
        skeleton=AutoModelForMultimodalLM.from_config(cfg,torch_dtype=torch.float16)
    root='model.language_model.layers'
    assert len(dict(skeleton.named_modules())[root])==42
    decoders={f'{root}.{i}' for i in range(42)}
    dm={}
    def includes_decoder(name):
        prefix=name+'.' if name else ''
        return any(d==name or d.startswith(prefix) for d in decoders)
    def assign(mod,name=''):
        if name in decoders:
            dm[name]='cuda:1'
            return
        if name and not includes_decoder(name):
            dm[name]='cuda:0'
            return
        assert not list(mod.named_parameters(recurse=False))
        assert not list(mod.named_buffers(recurse=False))
        for childname,child in mod.named_children():
            assign(child,f'{name}.{childname}' if name else childname)
    assign(skeleton)
    assert len(dm)==53 and '' not in dm
    for name,_ in list(skeleton.named_parameters())+list(skeleton.named_buffers()):
        assert sum(name==k or name.startswith(k+'.') for k in dm)==1,name
    del skeleton
    return dm


def load_system():
    from transformers import AutoProcessor, AutoModelForMultimodalLM
    assert torch.cuda.device_count()==2, 'Original system requires T4 x2'
    model_dir = Path(OUTPUT_DIR)/'model_download'
    model_dir.mkdir(parents=True,exist_ok=True)
    # Downloads original checkpoint and processor from HF (never from Kaggle Input/Drive).
    snapshot_download(repo_id=MODEL_REPO,revision=REVISION,local_dir=str(model_dir),token=os.getenv('HF_TOKEN') or None)
    weights = model_dir/'model.safetensors'
    assert weights.is_file() and weights.stat().st_size==15992595884
    assert sha_file(weights)==WEIGHT_SHA256,'Pinned model weights SHA256 mismatch'
    processor=AutoProcessor.from_pretrained(str(model_dir),local_files_only=True)
    processor.tokenizer.padding_side='left'
    dm=model_device_map(model_dir)
    model=AutoModelForMultimodalLM.from_pretrained(
        str(model_dir),local_files_only=True,dtype=torch.float16,
        attn_implementation='sdpa',device_map=dm,low_cpu_mem_usage=True).eval()
    assert type(model).__name__=='Gemma4ForConditionalGeneration'
    assert sum(p.numel() for p in model.parameters())==EXPECTED_PARAMS
    assert not model.training
    tok=processor.tokenizer
    pad=tok.pad_token_id
    turn=tok.convert_tokens_to_ids('<turn|>')
    assert pad is not None and tok.decode([turn],skip_special_tokens=False)=='<turn|>'
    eos=model.generation_config.eos_token_id
    stop=({eos} if isinstance(eos,int) else set(eos or []))|{turn}
    return processor,model,pad,stop


def run_batch(system, rows, limit):
    processor,model,pad,stop=system
    tok=processor.tokenizer
    msgs=[[{'role':'user','content':[{'type':'text','text':make_prompt(row)}]}] for row in rows]
    encoded=processor.apply_chat_template(msgs,tokenize=True,add_generation_prompt=True,
        enable_thinking=True,padding=True,return_dict=True,return_tensors='pt').to('cuda:0')
    width=encoded['input_ids'].shape[1]
    for i in range(2):torch.cuda.synchronize(i)
    started=time.perf_counter()
    with torch.inference_mode():
        seqs=model.generate(**encoded,max_new_tokens=limit,do_sample=False,
                            use_cache=True,pad_token_id=pad)
    for i in range(2):torch.cuda.synchronize(i)
    elapsed=time.perf_counter()-started
    results=[]
    for j,row in enumerate(rows):
        ids,ended=split_tokens(tok,seqs[j],width,pad,stop)
        raw=tok.decode(ids,skip_special_tokens=False)
        if limit==RETRY_TOKENS:
            pred,status,last=parse_retry_strict(raw,ended,len(ids),limit)
        else:
            pred=parse_first_pass(raw)
            status='completed' if pred else 'invalid_first_pass'
            last=''
        results.append({'fingerprint':fingerprint(row),'prediction':pred,'status':status,
            'raw_output':raw,'generated_tokens':len(ids),'terminated':bool(ended),
            'max_new_tokens':limit,'batch_size':len(rows),'elapsed_batch_seconds':round(elapsed,3),
            'last_line':last})
    return results


def predict_all(system,rows):
    root=Path(OUTPUT_DIR)
    root.mkdir(parents=True,exist_ok=True)
    config={'model':MODEL_REPO,'revision':REVISION,'prompt_sha256':PROMPT_SHA256,
        'parameters':EXPECTED_PARAMS,'native_thinking':True,'batch_size':BATCH_SIZE,
        'first_pass_tokens':FIRST_PASS_TOKENS,'retry_tokens':RETRY_TOKENS,
        'decoding':'greedy','seed':SEED,'rows':len(rows),
        'input_hashes':[fingerprint(r) for r in rows]}
    config_path=root/'configuration.json'
    if config_path.exists():
        assert json.loads(config_path.read_text('utf-8'))==config,'Resume configuration differs'
    else:save_atomic(config_path,config)
    answers={}
    for p in (root/'first_pass').glob('q_*.json'):
        obj=json.loads(p.read_text('utf-8'))
        i=obj['index']
        assert 0<=i<len(rows) and obj['fingerprint']==fingerprint(rows[i])
        assert i not in answers
        answers[i]=obj['prediction']
    for start in range(0,len(rows),BATCH_SIZE):
        # Recreate the ORIGINAL full batch when resuming a partially saved batch.
        # Dropping completed members would change padding and potentially outputs.
        idx=list(range(start,min(start+BATCH_SIZE,len(rows))))
        if all(i in answers for i in idx):continue
        results=run_batch(system,[rows[i] for i in idx],FIRST_PASS_TOKENS)
        for i,obj in zip(idx,results):
            obj['index']=i
            if i in answers:
                assert answers[i]==obj['prediction'],(
                    f'Replayed batch disagreed at index {i}; preserve old checkpoint')
            else:
                save_atomic(root/'first_pass'/f'q_{i:05}.json',obj)
                answers[i]=obj['prediction']
        logging.info('First pass %d/%d saved; valid=%d',len(answers),len(rows),sum(v in LETTER_TO_NUMBER for v in answers.values()))
    assert len(answers)==len(rows)
    failed=[i for i in range(len(rows)) if answers[i] not in LETTER_TO_NUMBER]
    logging.info('First pass valid=%d failed=%d',len(rows)-len(failed),len(failed))
    recovered={}
    for i in failed:
        p=root/'retry'/f'q_{i:05}.json'
        if p.exists():
            obj=json.loads(p.read_text('utf-8'))
            assert obj['fingerprint']==fingerprint(rows[i])
        else:
            obj=run_batch(system,[rows[i]],RETRY_TOKENS)[0]
            obj['index']=i
            save_atomic(p,obj)
        recovered[i]=obj['prediction']
        logging.info('Retry idx=%d status=%s valid=%s',i,obj['status'],obj['prediction'] in LETTER_TO_NUMBER)
    combined=[recovered[i] if i in recovered else answers[i] for i in range(len(rows))]
    unresolved=[i for i,pred in enumerate(combined) if pred not in LETTER_TO_NUMBER]
    save_atomic(root/'completion_status.json',{'questions':len(rows),'first_pass_valid':len(rows)-len(failed),
        'retry_targets':len(failed),'recovered_valid':len(failed)-len(unresolved),'still_invalid':unresolved})
    if unresolved:
        raise RuntimeError(f'{len(unresolved)} unresolved final answers; refusing to guess: {unresolved[:30]}')
    return [LETTER_TO_NUMBER[pred] for pred in combined]


def write_submission(rows,answers):
    assert len(rows)==len(answers)
    path=Path(OUTPUT_FILE)
    tmp=path.with_suffix('.jsonl.tmp')
    with tmp.open('w',encoding='utf-8') as f:
        for r,answer in zip(rows,answers):
            assert answer in ('1','2','3','4')
            item={'q_no':r['q_no'],'subject':r['subject'],'category':r['category'],
                  'level':get_level(r),'answer':answer}
            f.write(json.dumps(item,ensure_ascii=False,separators=(',',':'))+'\n')
    os.replace(tmp,path)
    logging.info('Saved %s sha256=%s',path,sha_file(path))


if __name__=='__main__':
    logging.basicConfig(filename='inference.log',level=logging.INFO,
                        format='%(asctime)s UTC %(levelname)s %(message)s')
    init_seed()
    rows=canonical_order(load_questions(DATA_REPO))
    logging.info('Questions loaded: %d',len(rows))
    system=load_system()
    answers=predict_all(system,rows)
    write_submission(rows,answers)
    print('Complete:',len(answers),'predictions; sha256=',sha_file(OUTPUT_FILE),flush=True)