# ─────────────────────────────────────────────────────────
# EGen Labs | ErebusTN
# AXIOM-1 — EGen-Core/AXIOM-1
# File: axiom1/inference/generate.py
# ─────────────────────────────────────────────────────────

"""
Autoregressive text generation for AXIOM-1.

Always uses bfloat16 autocast. Never uses float16.
Supports top-k, top-p, and temperature sampling.
"""

import torch
import torch.nn.functional as F

from axiom1.model.config import AXIOM1Config
from axiom1.model.model import AXIOM1


@torch.no_grad()
def generate(
    model: AXIOM1,
    prompt_ids: torch.Tensor,
    cfg: AXIOM1Config,
    max_new_tokens: int = 512,
    temperature: float = 0.8,
    top_p: float = 0.9,
    top_k: int = 50,
) -> torch.Tensor:
    """
    Autoregressive text generation for AXIOM-1.

    Always uses bfloat16 autocast. Never uses float16.

    Args:
        model: AXIOM1 model instance
        prompt_ids: Prompt token IDs (batch, prompt_len)
        cfg: AXIOM1Config instance
        max_new_tokens: Maximum number of tokens to generate
        temperature: Sampling temperature (0.8 default)
        top_p: Nucleus sampling threshold (0.9 default)
        top_k: Top-k filtering (50 default)

    Returns:
        Full sequence including prompt and generated tokens (batch, total_len)
    """
    model.eval()
    device = prompt_ids.device
    ids = prompt_ids.clone()

    for _ in range(max_new_tokens):
        # For long sequences, only process last max_seq_len tokens
        input_ids = ids[:, -cfg.max_seq_len:]

        device_type = 'cuda' if input_ids.is_cuda else 'cpu'
        with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
            logits = model(input_ids)

        # Sample from the last position
        next_logits = logits[:, -1, :].float() / max(temperature, 1e-5)

        # Top-k filter
        if top_k > 0:
            vals, _ = torch.topk(next_logits, top_k)
            next_logits[next_logits < vals[:, -1:]] = float('-inf')

        # Top-p filter
        probs = torch.softmax(next_logits, dim=-1)
        sorted_probs, sorted_idx = torch.sort(probs, descending=True)
        cumprobs = sorted_probs.cumsum(dim=-1)
        remove = cumprobs - sorted_probs > top_p
        sorted_probs[remove] = 0.0
        sorted_probs /= sorted_probs.sum(dim=-1, keepdim=True)
        probs.scatter_(-1, sorted_idx, sorted_probs)

        next_token = torch.multinomial(probs, num_samples=1)
        ids = torch.cat([ids, next_token], dim=1)

        if next_token.item() == cfg.eos_token_id:
            break

    return ids
