# Identify high-error regions for targeted follow-up
def select_targeted_queries(surrogate, target_cache, candidates, top_k=1000):
    errors = []
    for prompt, target_response in target_cache.items():
        surrogate_output = surrogate.generate(prompt, max_tokens=len(target_response))
        divergence = compute_token_divergence(surrogate_output, target_response)
        errors.append((prompt, divergence))
    errors.sort(key=lambda x: x[1], reverse=True)
    # Generate variants of high-error prompts
    return expand_prompt_variants(errors[:top_k], candidates)

# Active learning query selection with MMR diversity
def select_queries_mmr(candidates, surrogate, budget, lambda_div=0.3):
    entropies = [surrogate.compute_output_entropy(c) for c in candidates]
    embeddings = surrogate.encode(candidates)
    selected = []
    remaining = list(range(len(candidates)))
    for _ in range(budget):
        scores = []
        for idx in remaining:
            relevance = entropies[idx]
            if selected:
                max_sim = max(cosine_sim(embeddings[idx], embeddings[s]) for s in selected)
            else:
                max_sim = 0
            score = (1 - lambda_div) * relevance - lambda_div * max_sim
            scores.append((idx, score))
        best = max(scores, key=lambda x: x[1])[0]
        selected.append(best)
        remaining.remove(best)
    return [candidates[i] for i in selected]

def sparse_distillation_loss(student_logits, teacher_top_k, temperature=2.0):
    student_log_probs = log_softmax(student_logits / temperature, dim=-1)
    # Build sparse teacher distribution from top-k
    teacher_dist = torch.zeros_like(student_log_probs)
    for idx, logprob in teacher_top_k:
        teacher_dist[idx] = exp(logprob)
    teacher_dist = teacher_dist / teacher_dist.sum()  # renormalize
    mask = teacher_dist > 0
    loss = -(teacher_dist[mask] * student_log_probs[mask]).sum()
    return loss * (temperature ** 2)

def reconstruct_teacher_dist(teacher_top_k, student_logits, vocab_size):
    full_dist = torch.zeros(vocab_size)
    top_k_mass = 0
    for idx, logprob in teacher_top_k:
        prob = exp(logprob)
        full_dist[idx] = prob
        top_k_mass += prob
    # Distribute residual mass using student prior
    residual = 1.0 - top_k_mass
    student_probs = softmax(student_logits, dim=-1)
    for idx, logprob in teacher_top_k:
        student_probs[idx] = 0
    student_probs = student_probs / student_probs.sum()
    full_dist += residual * student_probs
    return full_dist