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

"""
Memory estimation and reporting for AXIOM-1 inference.

Helps determine whether a given inference scenario fits within
a target VRAM budget (2GB / 4GB / 8GB / 16GB).

AXIOM-1 is memory-efficient by design: one shared block means
model weights are small. The main variable is KV cache size,
which scales with num_loops × seq_len × n_kv_heads × head_dim.
"""

import torch
import torch.nn as nn

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


DTYPE_BYTES = {
    'bfloat16': 2,
    'float16': 2,   # Listed for estimation only — never use float16 for AXIOM-1
    'float32': 4,
    'int8': 1,
}


def estimate_inference_vram(
    cfg: AXIOM1Config,
    batch_size: int = 1,
    seq_len: int = 512,
    dtype: str = 'bfloat16'
) -> dict[str, float]:
    """
    Estimate VRAM usage for a given inference scenario.

    Provides a breakdown of memory consumption and checks against
    common VRAM budgets.

    Args:
        cfg: AXIOM1Config instance
        batch_size: Number of sequences in batch
        seq_len: Sequence length for inference
        dtype: Data type string ('bfloat16', 'float32', 'int8')

    Returns:
        Dict with 'embeddings_mb', 'block_mb', 'kv_cache_mb',
        'activations_mb', 'total_mb', and 'fits_in' budget results.
    """
    bpe = DTYPE_BYTES.get(dtype, 2)

    # Embedding table: vocab_size × d_model
    # Weight-tied, so lm_head shares this
    embed_params = cfg.vocab_size * cfg.d_model
    embed_bytes = embed_params * bpe

    # Loop step embeddings: max_loops × d_model
    lse_bytes = cfg.max_loops * cfg.d_model * bpe

    # Shared block parameters (one instance):
    #   GQA: q_proj + k_proj + v_proj + o_proj
    #   SwiGLU: gate + up + down
    #   2 × RMSNorm
    d_inner = int(cfg.d_ff * 2 / 3)
    attn_params = (
        cfg.d_model * cfg.d_model +                    # q_proj
        cfg.d_model * cfg.n_kv_heads * cfg.head_dim +  # k_proj
        cfg.d_model * cfg.n_kv_heads * cfg.head_dim +  # v_proj
        cfg.d_model * cfg.d_model                       # o_proj
    )
    ffn_params = (
        cfg.d_model * d_inner +  # gate
        cfg.d_model * d_inner +  # up
        d_inner * cfg.d_model    # down
    )
    norm_params = cfg.d_model * 2  # ln1 + ln2
    block_bytes = (attn_params + ffn_params + norm_params) * bpe

    # KV cache: per-step, per-head, per-token
    # num_loops × 2(k+v) × batch × n_kv_heads × seq_len × head_dim
    kv_cache_bytes = (
        cfg.num_loops * 2 * batch_size *
        cfg.n_kv_heads * seq_len * cfg.head_dim * bpe
    )

    # Activations estimate: batch × seq_len × d_model × num_loops
    # (rough upper bound — actual depends on implementation)
    act_bytes = batch_size * seq_len * cfg.d_model * bpe * 2  # ~2 buffers

    total_bytes = embed_bytes + lse_bytes + block_bytes + kv_cache_bytes + act_bytes

    # Convert to MB
    to_mb = lambda b: b / (1024 ** 2)

    result = {
        'embeddings_mb': to_mb(embed_bytes + lse_bytes),
        'block_mb': to_mb(block_bytes),
        'kv_cache_mb': to_mb(kv_cache_bytes),
        'activations_mb': to_mb(act_bytes),
        'total_mb': to_mb(total_bytes),
        'fits_2gb': total_bytes < 2 * 1024**3,
        'fits_4gb': total_bytes < 4 * 1024**3,
        'fits_8gb': total_bytes < 8 * 1024**3,
        'fits_16gb': total_bytes < 16 * 1024**3,
    }

    return result


def memory_breakdown(model: AXIOM1) -> dict[str, float]:
    """
    Report actual memory per component of a loaded model.

    Args:
        model: AXIOM1 model instance

    Returns:
        Dict with component sizes in MB: 'token_embed', 'loop_step_embed',
        'block', 'final_norm', 'lm_head_effective', 'total_unique'.
    """
    def param_bytes(module: nn.Module) -> int:
        return sum(p.nelement() * p.element_size() for p in module.parameters())

    to_mb = lambda b: b / (1024 ** 2)

    embed_bytes = param_bytes(model.token_embed)
    lse_bytes = param_bytes(model.loop_step_embed)
    block_bytes = param_bytes(model.block)
    norm_bytes = (
        param_bytes(model.pre_loop_norm) +
        param_bytes(model.final_norm)
    )

    # lm_head is weight-tied to token_embed, so effective cost = 0
    total_unique = embed_bytes + lse_bytes + block_bytes + norm_bytes

    result = {
        'token_embed_mb': to_mb(embed_bytes),
        'loop_step_embed_mb': to_mb(lse_bytes),
        'block_mb': to_mb(block_bytes),
        'norms_mb': to_mb(norm_bytes),
        'lm_head_effective_mb': 0.0,  # Weight-tied
        'total_unique_mb': to_mb(total_unique),
    }

    return result


def print_memory_report(
    cfg: AXIOM1Config,
    model: AXIOM1 | None = None,
    batch_size: int = 1,
    seq_len: int = 512
) -> None:
    """
    Print a human-readable memory report for AXIOM-1 inference.

    Args:
        cfg: AXIOM1Config instance
        model: Optional loaded model for actual breakdown
        batch_size: Batch size for estimation
        seq_len: Sequence length for estimation
    """
    print("=" * 55)
    print("  AXIOM-1 Memory Report | EGen Labs")
    print("=" * 55)

    est = estimate_inference_vram(cfg, batch_size, seq_len)
    print(f"\n  Estimated VRAM (batch={batch_size}, seq={seq_len}, bfloat16):")
    print(f"    Embeddings:   {est['embeddings_mb']:8.1f} MB")
    print(f"    Block:        {est['block_mb']:8.1f} MB")
    print(f"    KV Cache:     {est['kv_cache_mb']:8.1f} MB")
    print(f"    Activations:  {est['activations_mb']:8.1f} MB")
    print(f"    ─────────────────────────────")
    print(f"    Total:        {est['total_mb']:8.1f} MB")

    print(f"\n  VRAM Budget Fit:")
    print(f"    2 GB:  {'✓' if est['fits_2gb'] else '✗'}")
    print(f"    4 GB:  {'✓' if est['fits_4gb'] else '✗'}")
    print(f"    8 GB:  {'✓' if est['fits_8gb'] else '✗'}")
    print(f"    16 GB: {'✓' if est['fits_16gb'] else '✗'}")

    if model is not None:
        bd = memory_breakdown(model)
        print(f"\n  Actual Model Breakdown:")
        print(f"    token_embed:     {bd['token_embed_mb']:8.2f} MB")
        print(f"    loop_step_embed: {bd['loop_step_embed_mb']:8.2f} MB")
        print(f"    block:           {bd['block_mb']:8.2f} MB")
        print(f"    norms:           {bd['norms_mb']:8.2f} MB")
        print(f"    lm_head:         {bd['lm_head_effective_mb']:8.2f} MB (weight-tied)")
        print(f"    ─────────────────────────────")
        print(f"    Total unique:    {bd['total_unique_mb']:8.2f} MB")

    print("=" * 55)
