#!/usr/bin/env python3
"""Setun-Cortex — 采样工具"""

import torch
import torch.nn.functional as F


def top_k_top_p_filtering(logits, top_k=0, top_p=0.0, filter_value=-float('Inf')):
    """Filter logits using top-k and/or top-p (nucleus) sampling."""
    if top_k > 0:
        top_k = min(top_k, logits.size(-1))
        indices_to_remove = logits < torch.topk(logits, top_k)[0][:, -1:]
        logits[indices_to_remove] = filter_value
    if top_p > 0.0:
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
        sorted_indices_to_remove = cumulative_probs > top_p
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = 0
        indices_to_remove = sorted_indices[sorted_indices_to_remove]
        logits[..., indices_to_remove] = filter_value
    return logits


def sample_next_token(logits, temperature=0.8, top_k=40, top_p=0.9, repetition_penalty=1.0, prev_tokens=None):
    """从 logits 采样下一个 token"""
    logits = logits[:, -1, :] / max(temperature, 1e-5)

    # Repetition penalty
    if repetition_penalty != 1.0 and prev_tokens is not None:
        for token in prev_tokens[0].tolist():
            if logits[0, token] < 0:
                logits[0, token] *= repetition_penalty
            else:
                logits[0, token] /= repetition_penalty

    # Top-K / Top-P filtering
    logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
    probs = F.softmax(logits, dim=-1)
    return torch.multinomial(probs, 1)


@torch.no_grad()
def generate(model, prompt_ids, max_new_tokens=300, temperature=0.8,
             top_k=40, top_p=0.9, repetition_penalty=1.0, seq_len=64, alpha_scale=0.5,
             beam_width=1):
    """自回归生成文本; beam_width=1 → 采样模式, >1 → beam search (2026-08-05)"""
    model.eval()
    if beam_width <= 1:
        gen = prompt_ids.clone()
        for _ in range(max_new_tokens):
            ctx = gen[:, -seq_len:]
            logits = model(ctx, alpha_scale=alpha_scale)
            next_id = sample_next_token(logits, temperature, top_k, top_p,
                                        repetition_penalty, gen[:, -seq_len:] if repetition_penalty > 1 else None)
            gen = torch.cat([gen, next_id], dim=1)
        return gen
    # ===== beam search (v17c 2026-08-05: 消除采样重复死循环) =====
    import math
    beams = [(prompt_ids.clone(), 0.0)]
    for _ in range(max_new_tokens):
        candidates = []
        for seq, score in beams:
            ctx = seq[:, -seq_len:]
            logits = model(ctx, alpha_scale=alpha_scale)[:, -1, :] / max(temperature, 1e-5)
            if repetition_penalty != 1.0:
                for t in seq[0].tolist():
                    if logits[0, t] < 0:
                        logits[0, t] *= repetition_penalty
                    else:
                        logits[0, t] /= repetition_penalty
            logp = F.log_softmax(logits, dim=-1)[0]
            topk = torch.topk(logp, beam_width * 2)
            for v, lp in zip(topk.indices.tolist(), topk.values.tolist()):
                new_seq = torch.cat([seq, torch.tensor([[v]], dtype=torch.long, device=seq.device)], dim=1)
                candidates.append((new_seq, score + lp))
        candidates.sort(key=lambda x: x[1], reverse=True)
        beams = candidates[:beam_width]
    best = max(beams, key=lambda x: x[1] / ((len(x[0][0]) + 5) ** 0.8))
    return best[0]