#!/usr/bin/env python3
"""
TensorDescriptor - Role-Aware Tensor Identity for Routing Decisions

Phase 3 Part 1: Structured parsing of Mistral/Llama tensor naming conventions.

Tensor naming convention:
    blk.{layer_idx}.{module}_{projection}.weight

Examples:
    blk.0.attn_q.weight    -> layer=0, family=attn, projection=q
    blk.0.ffn_down.weight  -> layer=0, family=ffn, projection=down
    token_embd.weight      -> layer=-1, family=embed, projection=wte
    output.weight          -> layer=-1, family=lm_head, projection=output
    output_norm.weight     -> layer=-1, family=norm, projection=output_norm

Invariant C-ROUTE-001:
    Compression routing MUST include TensorDescriptor. No descriptor -> route
    to conservative codec (bucket 3: k-means VQ).
"""

import re
from dataclasses import dataclass, asdict
from typing import Tuple, Optional, Dict, Any


@dataclass
class TensorDescriptor:
    """
    Role-aware tensor identity for routing decisions.

    This provides structured metadata about a tensor's role in the model,
    enabling role-aware codec selection.
    """
    tensor_name: str              # Full name: "blk.2.attn_q.weight"
    layer_idx: int                # Block number: 2 (or -1 for non-block tensors)
    module_family: str            # "attn" | "ffn" | "embed" | "norm" | "lm_head"
    projection: str               # "q" | "k" | "v" | "o" | "gate" | "up" | "down" | "wte" | "wpe" | "output" | etc
    shape: Tuple[int, ...]        # Tensor shape
    dtype: str                    # "float16" | "float32" | "bfloat16"

    # Derived flags
    is_attention: bool = False    # True for Q/K/V/O projections
    is_ffn: bool = False          # True for gate/up/down projections
    is_embedding: bool = False    # True for token/position embeddings
    is_norm: bool = False         # True for layer norms
    is_lm_head: bool = False      # True for output projection

    @classmethod
    def from_tensor_name(
        cls,
        name: str,
        shape: Tuple[int, ...],
        dtype: str
    ) -> "TensorDescriptor":
        """
        Parse Mistral/Llama tensor naming convention.

        Patterns:
            blk.{N}.attn_{q,k,v,output}.weight  -> attention projections
            blk.{N}.ffn_{gate,up,down}.weight   -> FFN projections
            blk.{N}.attn_norm.weight            -> attention norm
            blk.{N}.ffn_norm.weight             -> FFN norm
            token_embd.weight                   -> token embedding
            output_norm.weight                  -> final layer norm
            output.weight                       -> LM head

        Args:
            name: Full tensor name
            shape: Tensor shape tuple
            dtype: Data type as string

        Returns:
            TensorDescriptor with parsed fields
        """
        layer_idx = -1
        module_family = "unknown"
        projection = "unknown"

        is_attention = False
        is_ffn = False
        is_embedding = False
        is_norm = False
        is_lm_head = False

        # Pattern 1: Block tensors - blk.{N}.{module}_{type}.weight
        blk_pattern = re.match(r"blk\.(\d+)\.(\w+?)(?:_(\w+))?\.weight", name)
        if blk_pattern:
            layer_idx = int(blk_pattern.group(1))
            module_part = blk_pattern.group(2)
            sub_part = blk_pattern.group(3) or ""

            # Attention projections
            if module_part == "attn":
                if sub_part in ("q", "k", "v", "output"):
                    module_family = "attn"
                    projection = sub_part if sub_part != "output" else "o"
                    is_attention = True
                elif sub_part == "norm":
                    module_family = "norm"
                    projection = "attn_norm"
                    is_norm = True
                else:
                    module_family = "attn"
                    projection = sub_part or "unknown"
                    is_attention = True

            # FFN projections
            elif module_part == "ffn":
                if sub_part in ("gate", "up", "down"):
                    module_family = "ffn"
                    projection = sub_part
                    is_ffn = True
                elif sub_part == "norm":
                    module_family = "norm"
                    projection = "ffn_norm"
                    is_norm = True
                else:
                    module_family = "ffn"
                    projection = sub_part or "unknown"
                    is_ffn = True

            else:
                # Other block components
                module_family = module_part
                projection = sub_part or module_part

        # Pattern 2: Token embedding
        elif name in ("token_embd.weight", "tok_embeddings.weight", "wte.weight"):
            layer_idx = -1
            module_family = "embed"
            projection = "wte"
            is_embedding = True

        # Pattern 3: Position embedding (if present)
        elif name in ("position_embd.weight", "pos_embeddings.weight", "wpe.weight"):
            layer_idx = -1
            module_family = "embed"
            projection = "wpe"
            is_embedding = True

        # Pattern 4: Output norm
        elif name in ("output_norm.weight", "norm.weight", "ln_f.weight"):
            layer_idx = -1
            module_family = "norm"
            projection = "output_norm"
            is_norm = True

        # Pattern 5: LM head
        elif name in ("output.weight", "lm_head.weight"):
            layer_idx = -1
            module_family = "lm_head"
            projection = "output"
            is_lm_head = True

        # Pattern 6: Fallback - try to extract what we can
        else:
            # Handle alternative naming conventions
            alt_match = re.match(r"layers\.(\d+)\.(\w+)\.(\w+)\.weight", name)
            if alt_match:
                layer_idx = int(alt_match.group(1))
                module_part = alt_match.group(2)
                proj_part = alt_match.group(3)

                if "attn" in module_part or "attention" in module_part:
                    module_family = "attn"
                    is_attention = True
                elif "mlp" in module_part or "ffn" in module_part:
                    module_family = "ffn"
                    is_ffn = True
                else:
                    module_family = module_part

                projection = proj_part
            else:
                # Give up - mark as unknown
                module_family = "unknown"
                projection = name.replace(".weight", "")

        return cls(
            tensor_name=name,
            layer_idx=layer_idx,
            module_family=module_family,
            projection=projection,
            shape=shape,
            dtype=dtype,
            is_attention=is_attention,
            is_ffn=is_ffn,
            is_embedding=is_embedding,
            is_norm=is_norm,
            is_lm_head=is_lm_head,
        )

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for JSON serialization."""
        return asdict(self)

    @property
    def is_weight_matrix(self) -> bool:
        """True if this is a 2D weight matrix (not norm/bias)."""
        return len(self.shape) == 2 and not self.is_norm

    @property
    def role_tag(self) -> str:
        """Short human-readable role tag for logging."""
        if self.layer_idx >= 0:
            return f"L{self.layer_idx}_{self.module_family}_{self.projection}"
        else:
            return f"{self.module_family}_{self.projection}"

    def __str__(self) -> str:
        return f"TensorDescriptor({self.tensor_name}, layer={self.layer_idx}, {self.module_family}/{self.projection})"


def parse_tensor_batch(
    tensor_names: list,
    shapes: dict,
    dtypes: dict
) -> Dict[str, TensorDescriptor]:
    """
    Parse a batch of tensor names into descriptors.

    Args:
        tensor_names: List of tensor names
        shapes: Dict mapping name -> shape
        dtypes: Dict mapping name -> dtype string

    Returns:
        Dict mapping name -> TensorDescriptor
    """
    return {
        name: TensorDescriptor.from_tensor_name(name, shapes[name], dtypes[name])
        for name in tensor_names
    }


# Test cases for validation
_TEST_CASES = [
    ("blk.0.attn_q.weight", (4096, 4096), "float16",
     {"layer_idx": 0, "module_family": "attn", "projection": "q", "is_attention": True}),
    ("blk.2.attn_k.weight", (1024, 4096), "float16",
     {"layer_idx": 2, "module_family": "attn", "projection": "k", "is_attention": True}),
    ("blk.5.attn_v.weight", (1024, 4096), "float16",
     {"layer_idx": 5, "module_family": "attn", "projection": "v", "is_attention": True}),
    ("blk.0.attn_output.weight", (4096, 4096), "float16",
     {"layer_idx": 0, "module_family": "attn", "projection": "o", "is_attention": True}),
    ("blk.0.ffn_gate.weight", (14336, 4096), "float16",
     {"layer_idx": 0, "module_family": "ffn", "projection": "gate", "is_ffn": True}),
    ("blk.0.ffn_up.weight", (14336, 4096), "float16",
     {"layer_idx": 0, "module_family": "ffn", "projection": "up", "is_ffn": True}),
    ("blk.0.ffn_down.weight", (4096, 14336), "float16",
     {"layer_idx": 0, "module_family": "ffn", "projection": "down", "is_ffn": True}),
    ("blk.0.attn_norm.weight", (4096,), "float16",
     {"layer_idx": 0, "module_family": "norm", "projection": "attn_norm", "is_norm": True}),
    ("blk.0.ffn_norm.weight", (4096,), "float16",
     {"layer_idx": 0, "module_family": "norm", "projection": "ffn_norm", "is_norm": True}),
    ("token_embd.weight", (32000, 4096), "float16",
     {"layer_idx": -1, "module_family": "embed", "projection": "wte", "is_embedding": True}),
    ("output_norm.weight", (4096,), "float16",
     {"layer_idx": -1, "module_family": "norm", "projection": "output_norm", "is_norm": True}),
    ("output.weight", (32000, 4096), "float16",
     {"layer_idx": -1, "module_family": "lm_head", "projection": "output", "is_lm_head": True}),
]


def run_self_test() -> bool:
    """Run internal test cases. Returns True if all pass."""
    passed = 0
    failed = 0

    for name, shape, dtype, expected in _TEST_CASES:
        td = TensorDescriptor.from_tensor_name(name, shape, dtype)

        all_ok = True
        for key, val in expected.items():
            actual = getattr(td, key)
            if actual != val:
                print(f"FAIL: {name}.{key} = {actual!r}, expected {val!r}")
                all_ok = False

        if all_ok:
            passed += 1
        else:
            failed += 1

    print(f"\nTensorDescriptor self-test: {passed} passed, {failed} failed")
    return failed == 0


if __name__ == "__main__":
    import sys

    # Run self-test
    print("Running TensorDescriptor self-test...")
    success = run_self_test()

    # Demo parsing
    print("\n" + "=" * 60)
    print("Demo: Parsing Mistral tensor names")
    print("=" * 60)

    demo_names = [
        "blk.0.attn_q.weight",
        "blk.2.attn_k.weight",
        "blk.10.ffn_down.weight",
        "token_embd.weight",
        "output.weight",
    ]

    for name in demo_names:
        td = TensorDescriptor.from_tensor_name(name, (4096, 4096), "float16")
        print(f"  {name}")
        print(f"    -> {td.role_tag}")
        print(f"    -> layer={td.layer_idx}, family={td.module_family}, proj={td.projection}")
        print()

    sys.exit(0 if success else 1)
