"""
Tensor-Train decomposed linear layers with Nested-Core Dynamic Execution.

Architecture:
  Replaces dense linear matrix W with d TT-cores:
    W_{i1...ik, o1...ok} = G^(1) · G^(2) ··· G^(d)

Key Innovations:
  1. Nested-Core Dynamic Execution:
     Allocates cores up to max_rank (default 8). At runtime, dynamically
     slices cores to active rank r in {1, 2, 4, 8} with ZERO runtime SVD overhead!
  2. Memory Traffic Instrumentation:
     Tracks actual bytes read (weights + inputs) and bytes written (activations)
     per forward pass for true hardware bandwidth optimization.
  3. Batched Vectorized Contraction:
     Hardware-efficient einsum/bmm contraction path.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Tuple, Optional, Dict, List


def factorize_dim(dim: int, max_factors: int = 4) -> Tuple[int, ...]:
    """
    Factorize a dimension into factors >= 2 to avoid trivial padding cores.
    """
    if dim <= 1:
        return (1,)
    factors = []
    remaining = dim
    for p in [2, 2, 3, 2, 5, 2, 3, 7]:
        while remaining % p == 0 and len(factors) < max_factors - 1:
            factors.append(p)
            remaining //= p
        if remaining == 1:
            break
    if remaining > 1 and len(factors) < max_factors:
        factors.append(remaining)
    while len(factors) < 2:
        val = factors[0] if factors else dim
        root = int(math.isqrt(val))
        for d in range(root, 1, -1):
            if val % d == 0:
                factors = [d, val // d]
                break
        else:
            factors = [1, val]
    return tuple(factors[:max_factors])


class TTLinear(nn.Module):
    """
    Tensor-Train decomposed linear layer with nested rank slicing.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        max_rank: int = 8,
        bias: bool = True,
        rank: Optional[int] = None,
        **kwargs,
    ):
        super().__init__()
        if rank is not None:
            effective_max_rank = max(max_rank, rank, 8)
            init_active = rank
        else:
            effective_max_rank = max(max_rank, 8)
            init_active = max_rank

        self.in_features = in_features
        self.out_features = out_features
        self.max_rank = effective_max_rank
        self.active_rank = init_active

        # Factorize input and output dimensions
        in_factors = list(factorize_dim(in_features))
        out_factors = list(factorize_dim(out_features))
        self.ndim = max(len(in_factors), len(out_factors))

        while len(in_factors) < self.ndim:
            in_factors.append(1)
        while len(out_factors) < self.ndim:
            out_factors.append(1)

        self.in_shape = tuple(in_factors)
        self.out_shape = tuple(out_factors)

        # Allocate nested TT-cores: G^(k) of shape (r_k, out_k, in_k, r_{k+1})
        self.cores = nn.ParameterList()
        for k in range(self.ndim):
            r_left = 1 if k == 0 else max_rank
            r_right = 1 if k == self.ndim - 1 else max_rank
            core = torch.empty(r_left, out_factors[k], in_factors[k], r_right)
            fan = max(1, r_left * in_factors[k] + r_right * out_factors[k])
            bound = math.sqrt(6.0 / fan)
            nn.init.uniform_(core, -bound, bound)
            self.cores.append(core)

        self.bias = nn.Parameter(torch.zeros(out_features)) if bias else None

        # Statistics
        self.dense_params = in_features * out_features
        self._last_bytes_read = 0
        self._last_bytes_written = 0

    @property
    def rank(self) -> int:
        return self.active_rank

    @rank.setter
    def rank(self, value: int):
        self.set_rank(value)

    def set_rank(self, rank: int):
        """Set active rank for nested core slicing via zero-copy tensor stride extraction."""
        self.active_rank = max(1, min(self.max_rank, int(rank)))

    def _slice_cores_for_rank(self, r: int) -> List[torch.Tensor]:
        """Slice nested cores according to specified rank."""
        sliced_cores = []
        for k, core in enumerate(self.cores):
            r_left = 1 if k == 0 else min(r, core.shape[0])
            r_right = 1 if k == self.ndim - 1 else min(r, core.shape[3])
            sliced_cores.append(core[:r_left, :, :, :r_right])
        return sliced_cores

    def get_active_cores(self) -> List[torch.Tensor]:
        """Slice nested cores according to current active rank."""
        return self._slice_cores_for_rank(self.active_rank)

    def measure_slicing_overhead_us(self, n_runs: int = 1000) -> float:
        """Empirically measure microsecond overhead of core slice extraction."""
        import time
        t0 = time.perf_counter()
        for _ in range(n_runs):
            _ = self._slice_cores_for_rank(4)
        return ((time.perf_counter() - t0) / n_runs) * 1e6

    @property
    def active_param_count(self) -> int:
        """Parameters actively used under current rank."""
        cores = self.get_active_cores()
        total = sum(c.numel() for c in cores)
        if self.bias is not None:
            total += self.bias.numel()
        return total

    @property
    def compression_ratio(self) -> float:
        """Effective compression ratio vs dense linear layer."""
        return self.dense_params / max(1, self.active_param_count)

    def _contract_with_cores(self, x_flat: torch.Tensor, cores: List[torch.Tensor]) -> torch.Tensor:
        """Contract 2D input (N, in_features) through given sliced cores."""
        N = x_flat.shape[0]
        state = x_flat.reshape(N, *self.in_shape)
        for k in range(self.ndim):
            core = cores[k]
            r_k, o_k, i_k, r_kp1 = core.shape
            if k == 0:
                rest = math.prod(self.in_shape[1:]) if self.ndim > 1 else 1
                s = state.reshape(N, i_k, rest)
                cm = core.squeeze(0).permute(1, 0, 2).reshape(i_k, o_k * r_kp1)
                s = torch.bmm(s.transpose(1, 2), cm.unsqueeze(0).expand(N, -1, -1))
                s = s.reshape(N, rest, o_k, r_kp1).permute(0, 3, 2, 1)
                state = s.reshape(N, r_kp1, -1)
            elif k == self.ndim - 1:
                prev_os = math.prod(self.out_shape[:k]) if k > 0 else 1
                s = state.reshape(N, r_k, prev_os, i_k)
                cm = core.squeeze(-1)
                s = torch.einsum("brpi,roi->bpo", s, cm)
                state = s.reshape(N, prev_os * o_k)
            else:
                prev_os = math.prod(self.out_shape[:k]) if k > 0 else 1
                rest_in = math.prod(self.in_shape[k + 1:])
                s = state.reshape(N, r_k, prev_os, i_k, rest_in)
                s = torch.einsum("brpix,roiq->bpoqx", s, core)
                s = s.permute(0, 3, 1, 2, 4)
                state = s.reshape(N, r_kp1, prev_os * o_k * rest_in)
        out = state.reshape(N, self.out_features)
        if self.bias is not None:
            out = out + self.bias
        return out

    def to_matrix(self, rank: Optional[int] = None) -> torch.Tensor:
        """Reconstruct the effective dense weight matrix (out_features, in_features) for rank r."""
        r = self.active_rank if rank is None else rank
        cores = self._slice_cores_for_rank(r)
        device = self.cores[0].device
        dtype = self.cores[0].dtype
        eye = torch.eye(self.in_features, device=device, dtype=dtype)
        bias_backup = self.bias
        self.bias = None
        try:
            wt = self._contract_with_cores(eye, cores)
        finally:
            self.bias = bias_backup
        return wt.t()


    def forward(self, x: torch.Tensor, token_ranks: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        Forward pass with nested core contraction. Supports token-grouped dynamic rank execution.
        """
        batch_shape = x.shape[:-1]
        B = math.prod(batch_shape) if batch_shape else 1
        x_flat = x.reshape(B, self.in_features)

        if token_ranks is not None and token_ranks.numel() == B:
            flat_ranks = token_ranks.reshape(-1)
            out = torch.empty(B, self.out_features, device=x.device, dtype=x.dtype)
            total_read = 0
            for r in flat_ranks.unique():
                r_int = int(r.item())
                idx = (flat_ranks == r).nonzero(as_tuple=True)[0]
                cores_r = self._slice_cores_for_rank(r_int)
                out[idx] = self._contract_with_cores(x_flat[idx], cores_r)
                total_read += sum(c.numel() * c.element_size() for c in cores_r) + x_flat[idx].numel() * x_flat.element_size()
            self._last_bytes_read = total_read
            self._last_bytes_written = out.numel() * out.element_size()
            return out.reshape(*batch_shape, self.out_features)
        else:
            cores = self.get_active_cores()
            self._last_bytes_read = sum(c.numel() * c.element_size() for c in cores) + x_flat.numel() * x_flat.element_size()
            out = self._contract_with_cores(x_flat, cores)
            self._last_bytes_written = out.numel() * out.element_size()
            return out.reshape(*batch_shape, self.out_features)

    def flops(self, batch_size: int = 1, seq_len: int = 1) -> int:
        """Compute active FLOPs for current rank."""
        r = self.active_rank
        avg_dim = (sum(self.in_shape) + sum(self.out_shape)) / (2 * self.ndim)
        return int(2 * r**2 * self.ndim * avg_dim * batch_size * seq_len)

    def get_memory_traffic(self) -> Dict[str, int]:
        """Return bytes read and written in last forward pass."""
        return {
            "bytes_read": self._last_bytes_read,
            "bytes_written": self._last_bytes_written,
            "total_bytes": self._last_bytes_read + self._last_bytes_written,
        }


class TTFeedForward(nn.Module):
    """
    Tensor-Train Feed-Forward Network with nested dynamic ranks.
    """

    def __init__(
        self,
        hidden_dim: int,
        ff_multiplier: int = 4,
        rank: int = 8,
        activation=F.gelu,
    ):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.expanded_dim = hidden_dim * ff_multiplier
        self.max_rank = rank

        self.up_proj = TTLinear(hidden_dim, self.expanded_dim, max_rank=rank, bias=True)
        self.down_proj = TTLinear(self.expanded_dim, hidden_dim, max_rank=rank, bias=True)
        self.activation = activation

    def forward(self, x: torch.Tensor, token_ranks: Optional[torch.Tensor] = None) -> torch.Tensor:
        return self.down_proj(self.activation(self.up_proj(x, token_ranks=token_ranks)), token_ranks=token_ranks)

    def set_rank(self, rank: int):
        self.up_proj.set_rank(rank)
        self.down_proj.set_rank(rank)

    @property
    def rank(self) -> int:
        return self.up_proj.active_rank

    @property
    def total_params(self) -> int:
        return sum(p.numel() for p in self.parameters())

    @property
    def active_params(self) -> int:
        return self.up_proj.active_param_count + self.down_proj.active_param_count

    def flops(self, batch_size: int = 1, seq_len: int = 1) -> int:
        return self.up_proj.flops(batch_size, seq_len) + self.down_proj.flops(batch_size, seq_len)

    def get_memory_traffic(self) -> Dict[str, int]:
        up = self.up_proj.get_memory_traffic()
        down = self.down_proj.get_memory_traffic()
        return {
            "bytes_read": up["bytes_read"] + down["bytes_read"],
            "bytes_written": up["bytes_written"] + down["bytes_written"],
            "total_bytes": up["total_bytes"] + down["total_bytes"],
        }