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

"""
AirLLM-style CPU offloading for AXIOM-1 inference.

For low-VRAM inference only. Not used during training.
Uses bfloat16 autocast (never float16 — overflows under recurrence).

AXIOM-1 is ideal for offloading: only ONE shared block exists.
Unlike stacked models with N layers, AXIOM-1's block is moved to GPU once,
all 12 loop iterations execute, then it moves back to CPU.
"""

import torch

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


class AXIOM1_Offloaded:
    """
    AirLLM-style CPU offloading for AXIOM-1.

    Moves model components to GPU only when needed, then back to CPU.
    The shared LoopBlock is moved once and runs all N iterations
    before returning to CPU — this is the key advantage of LoopLLM
    for memory-efficient inference.

    Usage:
        model = AXIOM1(cfg)
        model.load_state_dict(checkpoint['model_state_dict'])
        offloaded = AXIOM1_Offloaded(model)
        ids = offloaded.generate(prompt_ids, cfg)
    """

    def __init__(self, model: AXIOM1, device: str = 'cuda') -> None:
        """
        Args:
            model: AXIOM1 model instance (will be moved to CPU)
            device: GPU device for computation
        """
        self.model = model.cpu()
        self.device = device

    @torch.no_grad()
    def generate(
        self,
        prompt_ids: torch.Tensor,
        cfg: AXIOM1Config,
        max_new_tokens: int = 256,
        temperature: float = 0.8,
        top_p: float = 0.9,
        top_k: int = 50
    ) -> torch.Tensor:
        """
        Generate text with CPU offloading.

        Args:
            prompt_ids: Prompt token IDs (batch, prompt_len)
            cfg: AXIOM1Config instance
            max_new_tokens: Maximum tokens to generate
            temperature: Sampling temperature
            top_p: Nucleus sampling threshold
            top_k: Top-k filtering (0 to disable)

        Returns:
            Full sequence including prompt and generated tokens
        """
        ids = prompt_ids.clone()

        for _ in range(max_new_tokens):
            logits = self._forward(ids[:, -cfg.max_seq_len:])
            next_tok = self._sample(
                logits[:, -1, :].float(), temperature, top_p, top_k
            )
            ids = torch.cat([ids, next_tok], dim=1)
            torch.cuda.empty_cache()
            if next_tok.item() == cfg.eos_token_id:
                break

        return ids

    def _forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        """
        Forward pass with component-level CPU↔GPU movement.

        The shared block is moved to GPU once, runs all loop iterations,
        then moves back — this is why LoopLLM is efficient for offloading.
        """
        m = self.model
        d = self.device

        # Embed
        m.token_embed.to(d)
        m.pre_loop_norm.to(d)
        h = m.token_embed(input_ids.to(d)) * (m.cfg.d_model ** 0.5)
        h = m.pre_loop_norm(h)
        m.token_embed.cpu()
        m.pre_loop_norm.cpu()
        torch.cuda.empty_cache()

        # Loop block (move once, run N times)
        m.block.to(d)
        m.loop_step_embed.to(d)
        with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
            for step in range(m.cfg.num_loops):
                lse = m.loop_step_embed(torch.tensor(step, device=d))
                h = h + lse.unsqueeze(0).unsqueeze(0)
                h = m.block(h)
        m.block.cpu()
        m.loop_step_embed.cpu()
        torch.cuda.empty_cache()

        # Output
        m.final_norm.to(d)
        m.lm_head.to(d)
        logits = m.lm_head(m.final_norm(h))
        m.final_norm.cpu()
        m.lm_head.cpu()
        torch.cuda.empty_cache()

        return logits.float()

    def _sample(
        self,
        logits: torch.Tensor,
        temperature: float,
        top_p: float,
        top_k: int = 50
    ) -> torch.Tensor:
        """Sample next token with temperature, top-k, and top-p filtering."""
        logits = logits / max(temperature, 1e-5)

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

        probs = torch.softmax(logits, dim=-1)
        sorted_probs, sorted_idx = torch.sort(probs, descending=True)
        cum = sorted_probs.cumsum(-1)
        sorted_probs[cum - sorted_probs > top_p] = 0.0
        sorted_probs /= sorted_probs.sum(-1, keepdim=True)
        probs.scatter_(-1, sorted_idx, sorted_probs)
        return torch.multinomial(probs, 1)
