"""No privileged outcomes in model construction; charged, selected-bit replay."""
from __future__ import annotations

import time
import torch

from unionmind.decision_sufficiency.core import CONTEXTS, count_samples, encode, estimate, plan, push


class Timer:
    def __init__(self, device): self.device = device
    def __enter__(self):
        torch.cuda.synchronize(self.device)
        self.start = time.perf_counter()
        return self
    def __exit__(self, *args):
        torch.cuda.synchronize(self.device)
        self.seconds = time.perf_counter()-self.start


def construct(samples, model_name, method):
    """Only labeled fitting samples enter the model/representation/planner."""
    counts = count_samples(samples[None], 6)
    model = estimate(counts, model_name)
    code, k = encode(model, method)
    compact = model if method == 'identity' else push(model, code, k)
    return model,code.to(torch.uint8),compact,plan(compact)


def select_samples(archive, regime, seed):
    generator = torch.Generator(device=archive['source'].device).manual_seed(seed)
    source = archive['source'][torch.randperm(len(archive['source']),device=archive['source'].device,
                                             generator=generator)[:1024]]
    adaptation = archive['adaptation'][torch.randperm(len(archive['adaptation']),
        device=archive['source'].device,generator=generator)]
    if regime == 'source_only': return source
    if regime == 'target_64': return adaptation[:64]
    if regime == 'target_256':
        return None if len(adaptation) < 256 else adaptation[:256]
    assert regime == 'source_plus_target_64'
    return torch.cat((source,adaptation[:64]))


def replay(policy, code, current, future, budget, cost_index):
    """Return final action, paid query count, and query-mask for each record.

    `future` is a retained offline payload. Only the bit selected by a paid
    query is indexed. It is never passed to a model or decision table lookup.
    """
    assert policy['shape'][0] == 1
    r = policy['shape'][1]
    table = policy['choices'].reshape(3,1,r,len(policy['costs']),9)[:,0]
    z = code[0,current].long()
    context = torch.zeros_like(current)
    action = torch.zeros_like(current)
    count = torch.zeros_like(current)
    seen = torch.zeros_like(current)
    live = torch.ones_like(current,dtype=torch.bool)
    child = torch.zeros((9,2,2),device=current.device,dtype=torch.long)
    for t,c in enumerate(CONTEXTS):
        for q in (0,1):
            if c[q] == -1:
                for o in (0,1):
                    child[t,q,o] = CONTEXTS.index(tuple(o if i == q else v for i,v in enumerate(c)))
    for left in range(budget,-1,-1):
        choice = table[left,z,cost_index,context]
        querying = live & (choice >= 3)
        q = (choice-3).clamp(0,1)
        selected = querying.nonzero().flatten()
        revealed = torch.zeros_like(current)
        revealed[selected] = (future[selected] >> (1-q[selected])) & 1
        context = torch.where(querying,child[context,q,revealed],context)
        action = torch.where(live & ~querying,choice,action)
        count += querying.long()
        seen |= torch.where(querying,1 << q,0)
        live = querying
    assert not bool(live.any())
    return action,count,seen