"""
PyTorch and Hugging Face implementation of Q-TensorFormer:
Closed-Loop Information-Adaptive Tensor Network LLM.

Features:
- TensorTrainLinear: Low-rank Tensor-Train contracted projection layer with zero-SVD dynamic rank slicing
- AdaptiveKVCache: Attention-sink preservation with sliding-window eviction & 4/8-bit dynamic quantization
- PIDResourceController: Closed-loop CMDP/PID token-by-token compute budgeter with hysteresis anti-chattering
- QTensorFormerAttention: Grouped-Query Attention (GQA) with RoPE and budget-adaptive projection
- QTensorFormerForCausalLM: Turnkey Hugging Face PreTrainedModel compatible with model.generate()
"""

import math
from typing import Any, Dict, List, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import CrossEntropyLoss

from transformers.modeling_utils import PreTrainedModel
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
from transformers.cache_utils import Cache
from transformers import AutoModelForCausalLM

from configuration_qtensorformer import QTensorFormerConfig, factorize_dim


# ==============================================================================
# 1. TENSOR-TRAIN LINEAR LAYER WITH DYNAMIC ZERO-SVD RANK SLICING
# ==============================================================================

class TensorTrainLinear(nn.Module):
    r"""
    Low-rank Tensor-Train contracted linear projection layer.
    Replaces dense weight matrix $W \in \mathbb{R}^{M \times N}$ with contracted TT-cores:
        $W \approx \mathcal{G}^{(1)} \times \mathcal{G}^{(2)} \times \dots \times \mathcal{G}^{(d)}$
    where each core $\mathcal{G}^{(k)} \in \mathbb{R}^{R_{k-1} \times m_k \times n_k \times R_k}$ with $R_0 = R_d = 1$.

    Provides:
    1. Direct contracted forward pass bypassing dense matrix materialization.
    2. Dynamic zero-SVD rank slicing: instantaneously truncates core ranks $R \rightarrow r$ at runtime.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        tt_ranks: Optional[List[int]] = None,
        in_shape: Optional[Tuple[int, ...]] = None,
        out_shape: Optional[Tuple[int, ...]] = None,
        bias: bool = False,
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.tt_ranks = list(tt_ranks) if tt_ranks is not None else [1, 16, 16, 1]
        self.num_cores = len(self.tt_ranks) - 1

        self.in_shape = in_shape if in_shape is not None else factorize_dim(in_features, num_cores=self.num_cores)
        self.out_shape = out_shape if out_shape is not None else factorize_dim(out_features, num_cores=self.num_cores)

        if math.prod(self.in_shape) != in_features:
            raise ValueError(f"Product of in_shape {self.in_shape} must equal in_features ({in_features})")
        if math.prod(self.out_shape) != out_features:
            raise ValueError(f"Product of out_shape {self.out_shape} must equal out_features ({out_features})")
        if len(self.in_shape) != self.num_cores or len(self.out_shape) != self.num_cores:
            raise ValueError("Core counts in in_shape, out_shape, and tt_ranks must match.")

        # Initialize TT-cores
        self.cores = nn.ParameterList()
        for k in range(self.num_cores):
            r_in = self.tt_ranks[k]
            m_k = self.in_shape[k]
            n_k = self.out_shape[k]
            r_out = self.tt_ranks[k + 1]
            std = (2.0 / (r_in * m_k + n_k * r_out)) ** 0.5
            core = nn.Parameter(torch.randn(r_in, m_k, n_k, r_out) * std)
            self.cores.append(core)

        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features))
        else:
            self.register_parameter("bias", None)

        self.active_rank: Optional[int] = None

    def set_active_rank(self, rank: Optional[int]):
        """Dynamically slice TT contraction rank for instant FLOP/memory reduction without SVD."""
        self.active_rank = rank

    def get_sliced_cores(self, active_rank: Optional[int] = None) -> List[torch.Tensor]:
        """Return TT cores sliced down to active_rank."""
        r = active_rank if active_rank is not None else self.active_rank
        if r is None:
            return list(self.cores)

        sliced = []
        for k, core in enumerate(self.cores):
            r_in = core.shape[0] if k == 0 else min(core.shape[0], r)
            r_out = core.shape[3] if k == self.num_cores - 1 else min(core.shape[3], r)
            sliced.append(core[:r_in, :, :, :r_out])
        return sliced

    def forward(self, x: torch.Tensor, active_rank: Optional[int] = None) -> torch.Tensor:
        """
        Forward pass contracting input tensor x against TT-cores.
        Shape: [..., in_features] -> [..., out_features]
        """
        orig_shape = x.shape
        batch_seq = orig_shape[:-1]
        x_flat = x.reshape(-1, *self.in_shape)  # [B*S, m_1, m_2, ..., m_d]
        batch_size = x_flat.shape[0]

        cores = self.get_sliced_cores(active_rank)

        # Step-by-step tensor-train contraction
        # Start: state is [batch_size, 1] representing boundary rank r_0 = 1
        res = x_flat.reshape(batch_size, self.in_shape[0], -1)  # [B*S, m_1, rem_m]
        # Core 1: G_1 is [1, m_1, n_1, r_1] -> contract m_1:
        # [B*S, m_1, rem_m] * [1, m_1, n_1, r_1] -> [B*S, n_1, r_1, rem_m]
        g0 = cores[0].squeeze(0)  # [m_1, n_1, r_1]
        curr = torch.einsum("b m r, m n k -> b n k r", res, g0)  # [B*S, n_1, r_1, rem_m]

        for i in range(1, self.num_cores):
            gi = cores[i]  # [r_{i-1}, m_i, n_i, r_i]
            m_i = self.in_shape[i]
            rem_after = -1
            curr = curr.reshape(batch_size, -1, gi.shape[0], m_i, max(1, math.prod(self.in_shape[i + 1:])))
            # curr: [B*S, prev_out, r_{i-1}, m_i, rem_after]
            # gi: [r_{i-1}, m_i, n_i, r_i]
            # Contract r_{i-1} and m_i:
            curr = torch.einsum("b o r m p, r m n k -> b o n k p", curr, gi)
            # Flatten outputs:
            curr = curr.reshape(batch_size, -1, gi.shape[3], max(1, math.prod(self.in_shape[i + 1:])))

        # Squeeze out boundary rank r_d = 1 and remainder
        out = curr.reshape(batch_size, self.out_features)
        if self.bias is not None:
            out = out + self.bias
        return out.reshape(*batch_seq, self.out_features)

    def reconstruct_weight(self) -> torch.Tensor:
        """Utility for test verification: reconstruct full dense W matrix from cores."""
        curr = self.cores[0]  # [1, m_1, n_1, r_1]
        for k in range(1, self.num_cores):
            curr = torch.einsum("... r, r m n s -> ... m n s", curr, self.cores[k])
        # Reshape to (prod(m), prod(n))
        W = curr.squeeze(0).squeeze(-1)
        # Permute modes: [m_1, n_1, m_2, n_2, ...] -> [m_1, m_2, ..., n_1, n_2, ...]
        perm_in = [2 * i for i in range(self.num_cores)]
        perm_out = [2 * i + 1 for i in range(self.num_cores)]
        W = W.permute(*perm_in, *perm_out).reshape(self.in_features, self.out_features)
        return W


# ==============================================================================
# 2. ADAPTIVE KV CACHE (ATTENTION SINK + SLIDING WINDOW + DYNAMIC QUANTIZATION)
# ==============================================================================

class AdaptiveKVCache(Cache):
    """
    Hardware-efficient KV cache supporting:
    1. Attention Sink Preservation: initial tokens (e.g., 4) are kept permanently.
    2. Sliding Window Eviction: keeps most recent local context tokens.
    3. Cold-Line Dynamic Quantization: compresses evicted/older tokens to INT4/INT8.
    """

    def __init__(
        self,
        sink_size: int = 4,
        window_size: int = 1024,
        quantize_cold: bool = True,
        quant_bits: int = 4,
    ):
        super().__init__()
        self.sink_size = sink_size
        self.window_size = window_size
        self.quantize_cold = quantize_cold
        self.quant_bits = quant_bits

        self.key_cache: List[torch.Tensor] = []
        self.value_cache: List[torch.Tensor] = []
        self._seen_tokens = 0

    def get_seq_length(self, layer_idx: Optional[int] = 0) -> int:
        if len(self.key_cache) <= layer_idx:
            return 0
        return self.key_cache[layer_idx].shape[-2]

    def get_max_length(self) -> Optional[int]:
        return self.sink_size + self.window_size

    def update(
        self,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        layer_idx: int,
        cache_kwargs: Optional[Dict[str, Any]] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        if len(self.key_cache) <= layer_idx:
            # First insertion
            self.key_cache.append(key_states)
            self.value_cache.append(value_states)
        else:
            prev_k = self.key_cache[layer_idx]
            prev_v = self.value_cache[layer_idx]
            self.key_cache[layer_idx] = torch.cat([prev_k, key_states], dim=-2)
            self.value_cache[layer_idx] = torch.cat([prev_v, value_states], dim=-2)

        # Apply attention-sink + sliding window retention if limit exceeded
        total_len = self.key_cache[layer_idx].shape[-2]
        max_allowed = self.sink_size + self.window_size

        if total_len > max_allowed:
            # Preserve [0 : sink_size] and slice recent [total_len - window_size : total_len]
            sink_k = self.key_cache[layer_idx][:, :, : self.sink_size, :]
            sink_v = self.value_cache[layer_idx][:, :, : self.sink_size, :]

            recent_k = self.key_cache[layer_idx][:, :, total_len - self.window_size :, :]
            recent_v = self.value_cache[layer_idx][:, :, total_len - self.window_size :, :]

            # Optional cold quantization on evicted segment
            if self.quantize_cold:
                middle_k = self.key_cache[layer_idx][:, :, self.sink_size : total_len - self.window_size, :]
                middle_v = self.value_cache[layer_idx][:, :, self.sink_size : total_len - self.window_size, :]
                # Simulate dynamic low-precision quantization (quantize & dequantize for numerical stability)
                q_k = self._quantize_dequantize(middle_k, self.quant_bits)
                q_v = self._quantize_dequantize(middle_v, self.quant_bits)
                # Keep compact slice of representative cold tokens
                sub_k = q_k[:, :, -min(8, q_k.shape[-2]) :, :]
                sub_v = q_v[:, :, -min(8, q_v.shape[-2]) :, :]
                self.key_cache[layer_idx] = torch.cat([sink_k, sub_k, recent_k], dim=-2)
                self.value_cache[layer_idx] = torch.cat([sink_v, sub_v, recent_v], dim=-2)
            else:
                self.key_cache[layer_idx] = torch.cat([sink_k, recent_k], dim=-2)
                self.value_cache[layer_idx] = torch.cat([sink_v, recent_v], dim=-2)

        if layer_idx == 0:
            self._seen_tokens += key_states.shape[-2]

        return self.key_cache[layer_idx], self.value_cache[layer_idx]

    @staticmethod
    def _quantize_dequantize(tensor: torch.Tensor, bits: int = 4) -> torch.Tensor:
        """Fast dynamic symmetric uniform quantization & dequantization."""
        if tensor.numel() == 0:
            return tensor
        qmax = (2 ** (bits - 1)) - 1
        abs_max = torch.amax(tensor.abs(), dim=-1, keepdim=True).clamp(min=1e-8)
        scale = abs_max / qmax
        quant = torch.clamp(torch.round(tensor / scale), -qmax, qmax)
        return quant * scale


# ==============================================================================
# 3. CLOSED-LOOP PID / CMDP DYNAMIC RESOURCE CONTROLLER
# ==============================================================================

class PIDResourceController(nn.Module):
    r"""
    Hardware-closed-loop Proportional-Integral-Derivative (PID) dual controller.
    Tracks 8D token information state $\mathbf{z}_t$, minimizes latency/energy budget drift:
        $e(t) = \text{Target\_Energy\_Budget} - \text{Estimated\_Cost}(t)$
        $u(t) = K_p e(t) + K_i \int e(\tau)d\tau + K_d \frac{de(t)}{dt}$
    Applies hysteresis deadband $\tau$ to prevent high-frequency routing chattering.
    """

    def __init__(
        self,
        kp: float = 0.1,
        ki: float = 0.01,
        kd: float = 0.05,
        hysteresis_threshold: float = 0.15,
        target_budget: float = 1.0,
        anti_windup_max: float = 5.0,
    ):
        super().__init__()
        self.kp = kp
        self.ki = ki
        self.kd = kd
        self.hysteresis_threshold = hysteresis_threshold
        self.target_budget = target_budget
        self.anti_windup_max = anti_windup_max

        self.register_buffer("integral_err", torch.tensor(0.0))
        self.register_buffer("prev_err", torch.tensor(0.0))
        self.register_buffer("prev_u", torch.tensor(0.0))
        self.register_buffer("active_rank_idx", torch.tensor(2))  # Default full rank

        # Predefined rank hierarchy (e.g. 4, 8, 16)
        self.rank_levels = [4, 8, 16]

    def reset_state(self):
        self.integral_err.zero_()
        self.prev_err.zero_()
        self.prev_u.zero_()
        self.active_rank_idx.fill_(2)

    def compute_8d_token_state(
        self,
        logits: torch.Tensor,
        hidden_states: torch.Tensor,
        position_idx: int,
        max_seq_len: int,
    ) -> torch.Tensor:
        """
        Extracts 8-dimensional token information vector z_t:
        [entropy, top1_prob, top2_margin, pos_ratio, hidden_norm, latency_drift, inst_cost, ema_cost]
        """
        probs = F.softmax(logits.detach().float()[:, -1, :], dim=-1)
        top_probs, _ = torch.topk(probs, k=min(2, probs.shape[-1]), dim=-1)
        top1 = top_probs[:, 0].mean()
        top2 = top_probs[:, 1].mean() if top_probs.shape[-1] > 1 else torch.tensor(0.0, device=logits.device)
        margin = top1 - top2

        entropy = -(probs * torch.log(probs + 1e-9)).sum(dim=-1).mean()
        pos_ratio = torch.tensor(position_idx / max(1, max_seq_len), device=logits.device)
        h_norm = hidden_states.norm(dim=-1).mean() / (hidden_states.shape[-1] ** 0.5)

        drift = torch.tensor(0.0, device=logits.device)
        inst_cost = 1.0 - (top1 * 0.5)
        ema_cost = inst_cost

        return torch.stack([entropy, top1, margin, pos_ratio, h_norm, drift, inst_cost, ema_cost])

    def forward(
        self,
        logits: Optional[torch.Tensor] = None,
        hidden_states: Optional[torch.Tensor] = None,
        position_idx: int = 0,
        max_seq_len: int = 4096,
        measured_cost: Optional[float] = None,
    ) -> Dict[str, Any]:
        """Compute control signal u(t) and select active contraction rank."""
        if measured_cost is None and logits is not None and hidden_states is not None:
            z_t = self.compute_8d_token_state(logits, hidden_states, position_idx, max_seq_len)
            cost = z_t[6].item()
        else:
            cost = measured_cost if measured_cost is not None else 0.5

        # Error signal
        err = self.target_budget - cost

        # PID terms
        p_term = self.kp * err
        self.integral_err = torch.clamp(
            self.integral_err + self.ki * err,
            -self.anti_windup_max,
            self.anti_windup_max,
        )
        d_term = self.kd * (err - self.prev_err.item())
        self.prev_err = torch.tensor(err, device=self.prev_err.device)

        u = p_term + self.integral_err.item() + d_term

        # Hysteresis band check: only switch mode if delta > threshold
        delta_u = abs(u - self.prev_u.item())
        if delta_u >= self.hysteresis_threshold:
            self.prev_u = torch.tensor(u, device=self.prev_u.device)
            if u > 0.3:
                # Plentiful budget -> increase rank
                self.active_rank_idx = torch.tensor(min(2, self.active_rank_idx.item() + 1), device=self.active_rank_idx.device)
            elif u < -0.3:
                # Tight budget -> decrease rank
                self.active_rank_idx = torch.tensor(max(0, self.active_rank_idx.item() - 1), device=self.active_rank_idx.device)

        active_rank = self.rank_levels[int(self.active_rank_idx.item())]
        return {
            "u_t": float(u),
            "error": float(err),
            "active_rank": active_rank,
            "chattering_suppressed": delta_u < self.hysteresis_threshold,
        }


# ==============================================================================
# 4. ROTARY POSITION EMBEDDINGS (RoPE)
# ==============================================================================

class QTensorFormerRotaryEmbedding(nn.Module):
    def __init__(self, dim: int, max_position_embeddings: int = 4096, base: float = 10000.0):
        super().__init__()
        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float() / self.dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def forward(self, x: torch.Tensor, seq_len: int) -> Tuple[torch.Tensor, torch.Tensor]:
        t = torch.arange(seq_len, device=x.device, dtype=self.inv_freq.dtype)
        freqs = torch.outer(t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        return emb.cos(), emb.sin()


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    cos = cos.unsqueeze(0).unsqueeze(0)  # [1, 1, seq_len, dim]
    sin = sin.unsqueeze(0).unsqueeze(0)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed


# ==============================================================================
# 5. RMS NORM
# ==============================================================================

class QTensorFormerRMSNorm(nn.Module):
    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * hidden_states.to(input_dtype)


# ==============================================================================
# 6. GROUPED-QUERY ATTENTION (GQA) WITH BUDGET-ADAPTIVE PROJECTIONS
# ==============================================================================

class QTensorFormerAttention(nn.Module):
    """
    Grouped-Query Attention (GQA) layer with RoPE and budget-adaptive projection.
    """

    def __init__(self, config: QTensorFormerConfig, layer_idx: Optional[int] = None):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = config.head_dim
        self.num_key_value_heads = config.num_key_value_heads
        self.num_key_value_groups = config.num_key_value_groups

        # Standard dense or TT-factorized projections
        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)

        self.rotary_emb = QTensorFormerRotaryEmbedding(
            self.head_dim,
            max_position_embeddings=config.max_position_embeddings,
            base=config.rope_theta,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_value: Optional[Cache] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
        **kwargs,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
        bsz, q_len, _ = hidden_states.size()

        query_states = self.q_proj(hidden_states)
        key_states = self.k_proj(hidden_states)
        value_states = self.v_proj(hidden_states)

        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)

        kv_seq_len = key_states.shape[-2]
        if past_key_value is not None:
            kv_seq_len += past_key_value.get_seq_length(self.layer_idx)

        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
        if past_key_value is not None:
            # Offset cos/sin for decoding tokens
            offset = past_key_value.get_seq_length(self.layer_idx)
            cos = cos[offset : offset + q_len]
            sin = sin[offset : offset + q_len]

        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

        if past_key_value is not None:
            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx)

        # Zero-copy GQA view broadcasting: keep key_states and value_states in (bsz, num_key_value_heads, 1, kv_seq_len, head_dim)
        grouped_5d = self.num_key_value_groups > 1
        if grouped_5d:
            q_for_attn = query_states.reshape(bsz, self.num_key_value_heads, self.num_key_value_groups, q_len, self.head_dim)
            k_for_attn = key_states.unsqueeze(2)
            v_for_attn = value_states.unsqueeze(2)
        else:
            q_for_attn = query_states
            k_for_attn = key_states
            v_for_attn = value_states

        self.last_kv_storage_elements = key_states.numel()
        self.last_kv_shared_storage = (
            k_for_attn.untyped_storage().data_ptr() == key_states.untyped_storage().data_ptr()
            and v_for_attn.untyped_storage().data_ptr() == value_states.untyped_storage().data_ptr()
        )

        # Prepare attention mask for F.scaled_dot_product_attention
        attn_mask = None
        is_causal = False
        if attention_mask is not None:
            if attention_mask.dim() == 2:
                # Check if all tokens are attended (common generation without padding)
                if (attention_mask == 1).all():
                    if q_len > 1:
                        is_causal = True
                        attn_mask = None
                    else:
                        is_causal = False
                        attn_mask = None
                else:
                    bool_mask = attention_mask[:, None, None, -key_states.shape[-2]:].bool()
                    if q_len > 1:
                        offset = max(0, key_states.shape[-2] - q_len)
                        causal_tril = torch.tril(
                            torch.ones((q_len, key_states.shape[-2]), dtype=torch.bool, device=hidden_states.device),
                            diagonal=offset,
                        )
                        attn_mask = bool_mask & causal_tril
                        is_causal = False
                    else:
                        attn_mask = bool_mask
                        is_causal = False
                    # Guard against completely masked rows (e.g. left-padded query tokens) to prevent SDPA NaN
                    row_has_any = attn_mask.any(dim=-1, keepdim=True)
                    if not row_has_any.all():
                        fallback_col = torch.zeros_like(attn_mask)
                        fallback_col[..., 0] = True
                        attn_mask = torch.where(row_has_any, attn_mask, fallback_col)
            else:
                if attention_mask.dtype in (torch.int64, torch.int32):
                    attn_mask = attention_mask.bool()
                elif attention_mask.dtype == torch.bool:
                    attn_mask = attention_mask
                else:
                    attn_mask = attention_mask.to(query_states.dtype)
            if attn_mask is not None and grouped_5d and attn_mask.dim() == 4:
                attn_mask = attn_mask.unsqueeze(2)
        elif q_len > 1:
            is_causal = True

        attn_output = F.scaled_dot_product_attention(
            q_for_attn,
            k_for_attn,
            v_for_attn,
            attn_mask=attn_mask,
            dropout_p=0.0,
            is_causal=is_causal,
        )
        if grouped_5d:
            attn_output = attn_output.reshape(bsz, self.num_heads, q_len, self.head_dim)

        attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, q_len, self.hidden_size)
        attn_output = self.o_proj(attn_output)

        return attn_output, None, past_key_value


# ==============================================================================
# 7. TENSOR-TRAIN MLP (SwiGLU WITH DYNAMIC RANK CONSTRACTION)
# ==============================================================================

class QTensorFormerMLP(nn.Module):
    """SwiGLU feed-forward network powered by low-rank TensorTrainLinear layers."""

    def __init__(self, config: QTensorFormerConfig):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size

        self.gate_proj = TensorTrainLinear(
            self.hidden_size,
            self.intermediate_size,
            tt_ranks=config.tt_ranks,
            bias=False,
        )
        self.up_proj = TensorTrainLinear(
            self.hidden_size,
            self.intermediate_size,
            tt_ranks=config.tt_ranks,
            bias=False,
        )
        self.down_proj = TensorTrainLinear(
            self.intermediate_size,
            self.hidden_size,
            tt_ranks=config.tt_ranks,
            bias=False,
        )

    def forward(self, x: torch.Tensor, active_rank: Optional[int] = None) -> torch.Tensor:
        gate = F.silu(self.gate_proj(x, active_rank=active_rank))
        up = self.up_proj(x, active_rank=active_rank)
        return self.down_proj(gate * up, active_rank=active_rank)


# ==============================================================================
# 8. TRANSFORMER DECODER BLOCK
# ==============================================================================

class QTensorFormerBlock(nn.Module):
    def __init__(self, config: QTensorFormerConfig, layer_idx: Optional[int] = None):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.input_layernorm = QTensorFormerRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.self_attn = QTensorFormerAttention(config, layer_idx=layer_idx)
        self.post_attention_layernorm = QTensorFormerRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.mlp = QTensorFormerMLP(config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_value: Optional[Cache] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
        active_rank: Optional[int] = None,
        **kwargs,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]:
        # Self-Attention
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states, attn_weights, present_key_value = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_value=past_key_value,
            output_attentions=output_attentions,
            use_cache=use_cache,
            **kwargs,
        )
        hidden_states = residual + hidden_states

        # MLP with dynamic Tensor-Train contraction
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states, active_rank=active_rank)
        hidden_states = residual + hidden_states

        outputs = (hidden_states,)
        if output_attentions:
            outputs += (attn_weights,)
        if use_cache:
            outputs += (present_key_value,)
        return outputs


# ==============================================================================
# 9. BASE MODEL & CAUSAL LANGUAGE MODEL
# ==============================================================================

class QTensorFormerPreTrainedModel(PreTrainedModel):
    config_class = QTensorFormerConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["QTensorFormerBlock"]

    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()


class QTensorFormerModel(QTensorFormerPreTrainedModel):
    def __init__(self, config: QTensorFormerConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        self.layers = nn.ModuleList(
            [QTensorFormerBlock(config, layer_idx=i) for i in range(config.num_hidden_layers)]
        )
        self.norm = QTensorFormerRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        # Dynamic closed-loop PID controller
        self.pid_controller = PIDResourceController(
            kp=config.pid_kp,
            ki=config.pid_ki,
            kd=config.pid_kd,
            hysteresis_threshold=config.hysteresis_threshold,
            target_budget=config.target_energy_budget,
        )

        self.post_init()

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, value):
        self.embed_tokens = value

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutputWithPast]:
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            inputs_embeds = self.embed_tokens(input_ids)

        # Setup custom AdaptiveKVCache if use_cache and not provided
        if use_cache and past_key_values is None:
            past_key_values = AdaptiveKVCache(
                sink_size=self.config.attention_sink_size,
                window_size=self.config.kv_cache_window_size,
                quantize_cold=self.config.quantize_kv_cold,
                quant_bits=self.config.kv_quant_bits,
            )

        hidden_states = inputs_embeds
        seq_len = hidden_states.shape[1]

        # Determine active TT-rank via PID controller if dynamic routing is active
        active_rank = None
        if self.config.enable_dynamic_routing and seq_len == 1:
            # During decode token generation, step the PID controller
            pid_info = self.pid_controller(
                hidden_states=hidden_states,
                position_idx=past_key_values.get_seq_length(0) if past_key_values else 0,
                max_seq_len=self.config.max_position_embeddings,
            )
            active_rank = pid_info["active_rank"]

        for decoder_layer in self.layers:
            layer_outputs = decoder_layer(
                hidden_states,
                attention_mask=attention_mask,
                position_ids=position_ids,
                past_key_value=past_key_values,
                output_attentions=output_attentions,
                use_cache=use_cache,
                active_rank=active_rank,
            )
            hidden_states = layer_outputs[0]

        hidden_states = self.norm(hidden_states)

        if not return_dict:
            return tuple(v for v in [hidden_states, past_key_values] if v is not None)

        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=past_key_values if use_cache else None,
        )


class QTensorFormerForCausalLM(QTensorFormerPreTrainedModel):
    _tied_weights_keys = ["lm_head.weight"]

    def __init__(self, config: QTensorFormerConfig):
        super().__init__(config)
        self.model = QTensorFormerModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        self.post_init()

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, value):
        self.model.embed_tokens = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        hidden_states = outputs[0]
        logits = self.lm_head(hidden_states)

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss_fct = CrossEntropyLoss()
            shift_logits = shift_logits.view(-1, self.config.vocab_size)
            shift_labels = shift_labels.view(-1)
            shift_labels = shift_labels.to(shift_logits.device)
            loss = loss_fct(shift_logits, shift_labels)

        if not return_dict:
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    def prepare_inputs_for_generation(
        self,
        input_ids,
        past_key_values=None,
        attention_mask=None,
        inputs_embeds=None,
        **kwargs,
    ):
        if past_key_values is not None:
            # During decoding with KV cache, only process the last token
            past_length = past_key_values.get_seq_length()
            if input_ids.shape[1] > past_length:
                input_ids = input_ids[:, past_length:]
            else:
                input_ids = input_ids[:, -1:]

        model_inputs = {
            "input_ids": input_ids,
            "past_key_values": past_key_values,
            "use_cache": kwargs.get("use_cache", True),
            "attention_mask": attention_mask,
        }
        return model_inputs

    def _reorder_cache(self, past_key_values, beam_idx):
        if past_key_values is None:
            return None
        reordered_past = []
        for layer_past in past_key_values:
            reordered_past.append(
                tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
            )
        return reordered_past


# Register with Hugging Face AutoModelForCausalLM
AutoModelForCausalLM.register(QTensorFormerConfig, QTensorFormerForCausalLM)