"""Batched, answer-blind candidate selection extracted from measured code."""
import ast
from collections import Counter
import re
import time
import numpy as np
import torch
import diffusion_runtime as base

def forbidden_tokens(context, window, ngram, max_count, eos):
    """Constraints only inspect resolved left context; EOS is always permitted."""
    context = context[-window:] if window else []
    banned = {token for token, count in Counter(context).items() if max_count and count >= max_count}
    if ngram and len(context) >= ngram - 1:
        suffix = tuple(context[-(ngram - 1):]) if ngram > 1 else ()
        for i in range(len(context) - ngram + 1):
            if tuple(context[i:i + ngram - 1]) == suffix:
                banned.add(context[i + ngram - 1])
    return banned - set(eos)


@torch.inference_mode()
def sample_pool(sampler, prompt, cfg, length, seed):
    """Batched candidate generation, with CPU rejection resampling included in timing."""
    model = sampler.model
    eos_raw = model.config.eos_token_id
    eos = set(eos_raw if isinstance(eos_raw,list) else [eos_raw])
    eos_one = min(eos)
    pad = model.config.pad_token_id
    p = prompt.shape[1]
    count = cfg["best_of"]
    wave_size = min(count, max(1,4096//(p+length)))
    rng = np.random.default_rng(seed)
    prefix_list = prompt[0].tolist()
    outputs, forwards = [], 0
    for start in range(0,count,wave_size):
        base.check_time(30)
        batch = min(wave_size,count-start)
        prefix = prompt.expand(batch,-1)
        canvas = torch.full((batch,length),pad,device="cuda",dtype=torch.long)
        mask = torch.ones_like(canvas,dtype=torch.bool)
        cpu_ids = np.full((batch,length),-1,dtype=np.int64)
        cpu_lp = np.zeros((batch,length),dtype=np.float64)
        steps = min(cfg["steps"],length)
        for step in range(steps):
            base.check_time(30)
            ids = torch.cat((prefix,canvas),1)
            fullmask = torch.cat((torch.zeros_like(prefix,dtype=torch.bool),mask),1)
            h = sampler.reference(ids,fullmask,p)[:,p:][mask]
            logits = model.lm_head(h).float()
            logits[:,pad] = -torch.inf
            k = min(cfg["top_k"],logits.shape[-1]-1)
            values, tokens = logits.topk(k,dim=-1)
            # EOS is a genuine logit candidate even if outside top-k; never synthesize it as a hard-ban violation.
            eos_values = logits[:,eos_one:eos_one+1]
            eos_tokens = torch.full_like(tokens[:,:1],eos_one)
            duplicate = tokens.eq(eos_one).any(-1,keepdim=True)
            eos_values = eos_values.masked_fill(duplicate,-torch.inf)
            values = torch.cat((values,eos_values),-1)
            tokens = torch.cat((tokens,eos_tokens),-1)
            normalizer = torch.logsumexp(logits/.9,-1).cpu().numpy()
            raw_values = values.cpu().numpy()
            token_values = tokens.cpu().numpy()
            del logits, values, tokens, h
            forwards += 1
            remaining_after = length*(steps-step-1)//steps
            row_offset = 0
            for row in range(batch):
                positions = np.flatnonzero(cpu_ids[row]<0)
                commit = len(positions)-remaining_after
                # Confidence selects slots; ascending positional commits enforce same-step bans.
                maximum = raw_values[row_offset:row_offset+len(positions),0]/.9-normalizer[row_offset:row_offset+len(positions)]
                chosen = np.argsort(maximum)[-commit:]
                for local in sorted(chosen.tolist(),key=lambda x:int(positions[x])):
                    pos = int(positions[local])
                    idx = row_offset+local
                    left = pos-1
                    while left>=0 and cpu_ids[row,left]>=0:
                        left -= 1
                    context = cpu_ids[row,left+1:pos].tolist()
                    if left<0:
                        context = prefix_list+context
                    window_context = context[-cfg["window"]:] if cfg["window"] else []
                    banned = forbidden_tokens(context,cfg["window"],cfg["ngram"],cfg["max_repeat"],eos)
                    scores = raw_values[idx].astype(np.float64).copy()
                    if cfg["penalty"] != 1:
                        counts = Counter(window_context)
                        for j,token in enumerate(token_values[idx]):
                            if int(token) in counts and int(token) not in eos:
                                penalty = cfg["penalty"]**min(counts[int(token)],3)
                                scores[j] = scores[j]/penalty if scores[j]>=0 else scores[j]*penalty
                    scores /= .9
                    for j,token in enumerate(token_values[idx]):
                        if int(token) in banned:
                            scores[j] = -np.inf
                    assert np.isfinite(scores).any(), "EOS must remain selectable"
                    probability = np.exp(scores-np.max(scores))
                    probability /= probability.sum()
                    choice = int(rng.choice(len(scores),p=probability))
                    token = int(token_values[idx,choice])
                    assert token not in banned
                    cpu_ids[row,pos] = token
                    cpu_lp[row,pos] = float(raw_values[idx,choice]/.9-normalizer[idx])
                row_offset += len(positions)
            known = cpu_ids>=0
            canvas = torch.as_tensor(np.where(known,cpu_ids,pad),device="cuda",dtype=torch.long)
            mask = torch.as_tensor(~known,device="cuda")
        for row in range(batch):
            sequence = cpu_ids[row].tolist()
            stop = next((i for i,token in enumerate(sequence) if token in eos),len(sequence))
            assert all(t>=0 for t in sequence[:stop])
            outputs.append({"ids":sequence[:stop],"ended":stop<len(sequence),
                            "mean_logprob":float(cpu_lp[row,:max(1,stop)].mean())})
    return outputs,forwards,wave_size


def code_candidate(text, repair):
    blocks = re.findall(r"```(?:python)?\s*\n(.*?)```",text,flags=re.S)
    source = blocks[0] if blocks else text.strip()
    choices = [source]
    if repair:
        lines = source.splitlines()
        choices += ["\n".join(lines[:i]) for i in range(len(lines)-1,max(0,len(lines)-17),-1)]
    for candidate in choices:
        try:
            tree = ast.parse(candidate)
            definitions = [n for n in tree.body if isinstance(n,ast.FunctionDef)]
            forbidden = any(isinstance(n,(ast.Import,ast.ImportFrom,ast.Attribute,ast.ClassDef)) for n in ast.walk(tree))
            if definitions and not forbidden:
                return candidate, [n.name for n in definitions], candidate!=source
        except (SyntaxError,ValueError):
            pass
    return source,[],False


def rank_candidate(record, tokenizer, prompt_text, repair):
    """No answer or tests argument: selector cannot access evaluation truth."""
    text = tokenizer.decode(record["ids"],skip_special_tokens=True).strip()
    match = re.search(r"defining\s+([A-Za-z_]\w*)",prompt_text)
    score = record["mean_logprob"]
    changed = False
    if match:
        text,names,changed = code_candidate(text,repair)
        score += 100 if match.group(1) in names else -100
        score += 10 if names else 0
    elif "only" in prompt_text.lower():
        words = text.split()
        score += 20 if 1<=len(words)<=4 and len(text)<60 else -20
    health = base.health(record["ids"])
    score -= max(0,health["max_run"]-3)*.5
    score -= (1-health["distinct4"])*5
    if not text:
        score -= 1000
    return score,text,changed


def generate_selected(sampler, tokenizer, prompt_text, cfg, seed, long=False):
    initial = base.encode(tokenizer,prompt_text)
    torch.cuda.synchronize()
    start = time.perf_counter()
    generated, gross, calls, chunks, repaired, max_batch = [],0,0,0,False,0
    selected_text = ""
    target = max(1024,cfg["block"]) if long else min(256,cfg["block"])
    while len(generated)<target:
        n = min(cfg["block"],target-len(generated))
        history = initial.new_tensor([generated[-256:]])
        prompt = torch.cat((initial,history),1)
        pool,forwards,wave = sample_pool(sampler,prompt,cfg,n,seed+chunks*1009)
        ranked = [(rank_candidate(row,tokenizer,prompt_text,cfg["repair"] and not long),index) for index,row in enumerate(pool)]
        ranking,index = max(ranked,key=lambda item:item[0][0])
        selected = pool[index]
        selected_text,changed = ranking[1],ranking[2]
        repaired |= changed
        gross += sum(len(row["ids"]) for row in pool)
        calls += forwards
        max_batch = max(max_batch,wave)
        generated.extend(selected["ids"])
        chunks += 1
        if not long or selected["ended"]:
            break
    if long:
        selected_text = tokenizer.decode(generated,skip_special_tokens=True)
        delivered = len(generated)
    else:
        # Post-processed output tokens, not discarded syntax-repair suffixes.
        delivered = len(tokenizer(selected_text,add_special_tokens=False)["input_ids"])
    torch.cuda.synchronize()
    elapsed = time.perf_counter()-start
    return selected_text,{"selected_tokens":delivered,"gross_candidate_tokens":gross,"seconds":elapsed,
                          "selected_tokens_per_second":delivered/elapsed,"gross_tokens_per_second":gross/elapsed,
                          "best_of":cfg["best_of"],"max_parallel_candidates":max_batch,"forward_calls":calls,
                          "chunks":chunks,"syntax_repaired":repaired,**base.health(generated)}


def cfg(block,steps=1,best_of=1,window=0,ngram=0,max_repeat=0,penalty=1.,top_k=50,repair=False):
    return dict(block=block,steps=steps,best_of=best_of,window=window,ngram=ngram,max_repeat=max_repeat,
                penalty=penalty,top_k=top_k,repair=repair,temperature=.9)