#!/usr/bin/env python3
"""
CompressionRouter - 4-Bucket Codec Routing with Role Prior Bias (Plane B)

Phase 3.5: Routes tensors to optimal compression strategy based on:
  - Energy concentration (E@10%) - low-rank vs high-rank
  - Structural coherence (C_struct) - structured vs unstructured
  - Structural axis (struct_axis) - row, col, isotropic, none
  - Role prior bias (layer_band + projection type)

Bucket Strategy (with 2A/2B axis-aware split):
  | Bucket | Condition                           | Codec              | Axis-Aware? |
  |--------|-------------------------------------|--------------------|-------------|
  | 1      | E@10% > thresh, C_struct < 0.30    | svd_lowrank        | No          |
  | 2A     | E@10% > thresh, C >= 0.30, row     | svd_row_aware      | Yes (row)   |
  | 2B     | E@10% > thresh, C >= 0.30, col     | svd_col_aware      | Yes (col)   |
  | 3      | E@10% <= thresh, C_struct < 0.30   | kmeans_vq          | No          |
  | 4      | E@10% <= thresh, C_struct >= 0.30  | dct_quant          | Yes         |

Role Prior Bias (Phase 3.5):
  - attn_q in early layers gets +0.10 bonus (more likely to hit SVD)
  - FFN projections get -0.10 penalty (never low-rank in our data)
  - Based on real GGUF analysis: only 2/224 tensors hit Bucket 2, both attn_q

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

import hashlib
import json
import re

import numpy as np
from dataclasses import dataclass, asdict
from typing import Dict, Any, List, Optional, Union, TYPE_CHECKING

if TYPE_CHECKING:
    from .tensor_descriptor import TensorDescriptor

# Try to import TensorProfile from various locations
try:
    from experiments.cymatic_compression.core.tensor_classifier import TensorProfile
except ImportError:
    TensorProfile = None  # Will be passed as parameter


# =============================================================================
# Bucket Names (Phase 3.5: Split bucket 2 into 2A/2B)
# =============================================================================

BUCKET_NAMES = {
    1: "low_rank_unstructured",
    "2A": "low_rank_row_structured",
    "2B": "low_rank_col_structured",
    2: "low_rank_structured",  # Legacy/fallback for isotropic
    3: "high_rank_unstructured",
    4: "high_rank_structured",
}


# =============================================================================
# Role Prior Bias (Phase 3.5)
# =============================================================================

# Role prior bias: (module_family, projection, layer_band) -> E@10% threshold adjustment
# Positive = bonus (lowers effective threshold, more likely to hit SVD)
# Negative = penalty (raises effective threshold, less likely to hit SVD)
#
# Based on real GGUF analysis (2026-01-26):
#   - Only 2/224 tensors hit Bucket 2 (SVD), both attn_q
#   - attn_q avg E@10% = 18%, all others 3-6%
#   - Later layers go full-rank (E@10% drops after layer 2)

ROLE_PRIORS = {
    # attn_q in early layers gets bonus (more likely to hit SVD)
    ("attn", "q", "early"): +0.10,   # 0.70 -> 0.60 effective threshold
    ("attn", "q", "mid"): +0.05,     # 0.70 -> 0.65
    ("attn", "q", "late"): 0.0,      # no adjustment

    # Other attention projections get penalty (rarely low-rank)
    ("attn", "k", "any"): -0.05,     # 0.70 -> 0.75
    ("attn", "v", "any"): -0.05,
    ("attn", "o", "any"): -0.05,

    # FFN projections get penalty (never low-rank in our data)
    ("ffn", "gate", "any"): -0.10,   # 0.70 -> 0.80
    ("ffn", "up", "any"): -0.10,
    ("ffn", "down", "any"): -0.10,

    # Embeddings and LM head - conservative
    ("embed", "wte", "any"): -0.10,
    ("embed", "wpe", "any"): -0.10,
    ("lm_head", "output", "any"): -0.10,

    # Norms are 1D, shouldn't hit this path, but be conservative
    ("norm", "any", "any"): -0.15,
}


def get_layer_band(layer_idx: int, total_layers: int = 32) -> str:
    """
    Map layer index to band: early/mid/late.

    Based on observed E@10% decay pattern in Mistral-7B:
      - Layers 0-7 (early): Higher chance of low-rank Q
      - Layers 8-23 (mid): Transitioning
      - Layers 24-31 (late): Full-rank

    Args:
        layer_idx: Block layer index (0-based), -1 for non-block tensors
        total_layers: Total number of transformer blocks (default 32)

    Returns:
        Layer band: "early", "mid", "late", or "any" for non-block tensors
    """
    if layer_idx < 0:
        return "any"  # Non-block tensors (embeddings, etc.)

    if layer_idx < total_layers // 4:      # 0-7 for 32-layer model
        return "early"
    elif layer_idx < 3 * total_layers // 4:  # 8-23 for 32-layer model
        return "mid"
    else:                                    # 24-31 for 32-layer model
        return "late"


def get_role_prior(
    descriptor: Optional["TensorDescriptor"],
    total_layers: int = 32
) -> float:
    """
    Get E@10% threshold adjustment for this tensor's role.

    Returns positive values (bonus) for roles that are likely low-rank,
    negative values (penalty) for roles that are rarely low-rank.

    Args:
        descriptor: TensorDescriptor with role information
        total_layers: Total number of transformer blocks

    Returns:
        Threshold adjustment (positive = bonus, negative = penalty)
    """
    if descriptor is None:
        return 0.0

    band = get_layer_band(descriptor.layer_idx, total_layers)

    # Try specific band first
    key = (descriptor.module_family, descriptor.projection, band)
    if key in ROLE_PRIORS:
        return ROLE_PRIORS[key]

    # Try "any" band
    key_any = (descriptor.module_family, descriptor.projection, "any")
    if key_any in ROLE_PRIORS:
        return ROLE_PRIORS[key_any]

    # Try module family with "any" projection and band
    key_family = (descriptor.module_family, "any", "any")
    if key_family in ROLE_PRIORS:
        return ROLE_PRIORS[key_family]

    # Unknown role - no adjustment
    return 0.0


def detect_total_layers(tensor_names: List[str]) -> int:
    """
    Auto-detect total_layers from tensor names (max blk.N + 1).

    Phase 3.5.1: Remove hardcoded 32-layer assumption.

    Args:
        tensor_names: List of tensor names from the model

    Returns:
        Total number of transformer blocks, or 32 as default if no blocks found
    """
    max_layer = -1
    for name in tensor_names:
        match = re.match(r"blk\.(\d+)\.", name)
        if match:
            max_layer = max(max_layer, int(match.group(1)))
    return max_layer + 1 if max_layer >= 0 else 32  # Default 32 if no blocks found


def get_router_version_hash() -> str:
    """
    Compute SHA256 hash of role priors config for drift detection.

    Phase 3.5.1: Include this in receipts to detect when routing config changes.

    Returns:
        First 16 chars of SHA256 hash of ROLE_PRIORS JSON
    """
    # Convert tuple keys to string keys for JSON serialization
    serializable_priors = {
        f"{k[0]}_{k[1]}_{k[2]}": v for k, v in ROLE_PRIORS.items()
    }
    priors_json = json.dumps(serializable_priors, sort_keys=True)
    return hashlib.sha256(priors_json.encode()).hexdigest()[:16]


# Phase 3.5.1: Compute router version hash at module load time
ROUTER_VERSION_HASH = get_router_version_hash()


@dataclass
class CompressionDecision:
    """
    Compression decision for a single tensor.

    This captures the routing decision and all parameters needed to reproduce
    the compression exactly during regen (Plane C).

    Phase 3.5: bucket can be int (1, 3, 4) or str ("2A", "2B") for axis-aware SVD.
    """
    bucket: Union[int, str]        # 1, "2A", "2B", 3, 4 (see BUCKET_NAMES)
    bucket_name: str               # Human-readable bucket name
    codec_id: str                  # "svd_lowrank" | "svd_row_aware" | "svd_col_aware" | "kmeans_vq" | "dct_quant"
    params: Dict[str, Any]         # Codec-specific parameters
    struct_axis: str               # "row" | "col" | "isotropic" | "none"
    routing_reason: str            # Human-readable explanation
    routing_metrics: Dict[str, float]  # Metrics used for routing decision

    # Phase 3.5: Role prior tracking
    role_prior: float = 0.0        # Threshold adjustment applied
    effective_threshold: float = 0.70  # Actual threshold used (base - prior)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for manifest serialization."""
        d = asdict(self)
        # Phase 3.5.1: Include router version hash for drift detection
        d["router_version"] = ROUTER_VERSION_HASH
        return d

    @property
    def bucket_int(self) -> int:
        """Return bucket as integer (2A/2B -> 2) for backward compat."""
        if isinstance(self.bucket, str):
            return int(self.bucket[0])  # "2A" -> 2, "2B" -> 2
        return self.bucket


# Default routing thresholds
DEFAULT_THRESHOLDS = {
    "low_rank_energy": 0.70,      # E@10% > 0.70 = low-rank (before role prior)
    "struct_coherence": 0.30,     # C_struct > 0.30 = structured
    "total_layers": 32,           # For layer band calculation
}

# Default codec parameters per bucket (Phase 3.5: 2A/2B split)
DEFAULT_CODEC_PARAMS = {
    1: {"codec_id": "svd_lowrank", "rank": 64},
    "2A": {"codec_id": "svd_row_aware", "rank": 64, "axis": "row", "residual_mode": "row_blocks"},
    "2B": {"codec_id": "svd_col_aware", "rank": 64, "axis": "col", "residual_mode": "col_blocks"},
    2: {"codec_id": "svd_axis_aware", "rank": 64, "axis": None},  # Legacy fallback for isotropic
    3: {"codec_id": "kmeans_vq", "n_bins": 256, "percentile": 99.9},
    4: {"codec_id": "dct_quant", "fraction": 0.15, "axis": None},  # axis filled from struct_axis
}


def route_compression(
    tensor: np.ndarray,
    descriptor: Optional["TensorDescriptor"],
    profile: Any,  # TensorProfile
    thresholds: Optional[Dict[str, float]] = None,
    codec_overrides: Optional[Dict[Union[int, str], Dict[str, Any]]] = None,
) -> CompressionDecision:
    """
    Route tensor to optimal compression strategy (Plane B).

    Phase 3.5: Uses role prior bias to adjust E@10% threshold per tensor role.
    Uses (energy_at_10pct, C_struct, struct_axis, role_prior) to select codec.

    Args:
        tensor: The weight matrix to compress
        descriptor: TensorDescriptor with role information (layer, projection, etc.)
                   If None, routes to conservative bucket 3 (k-means VQ)
        profile: TensorProfile with cheap stats from tensor_classifier
        thresholds: Override default routing thresholds
        codec_overrides: Override default codec parameters per bucket

    Returns:
        CompressionDecision with bucket, codec, and parameters

    Invariant C-ROUTE-001:
        No descriptor -> conservative codec (bucket 3: k-means VQ)
    """
    thresholds = thresholds or DEFAULT_THRESHOLDS
    codec_overrides = codec_overrides or {}

    LOW_RANK_THRESHOLD_BASE = thresholds.get("low_rank_energy", 0.70)
    STRUCT_THRESHOLD = thresholds.get("struct_coherence", 0.30)
    TOTAL_LAYERS = thresholds.get("total_layers", 32)

    # Handle missing descriptor (Invariant C-ROUTE-001)
    if descriptor is None:
        return CompressionDecision(
            bucket=3,
            bucket_name=BUCKET_NAMES[3],
            codec_id="kmeans_vq",
            params={"n_bins": 256, "percentile": 99.9, "fallback": True},
            struct_axis="none",
            routing_reason="No TensorDescriptor - fallback to conservative k-means VQ",
            routing_metrics={"descriptor": "missing"},
            role_prior=0.0,
            effective_threshold=LOW_RANK_THRESHOLD_BASE,
        )

    # Phase 3.5: Calculate role prior bias
    role_prior = get_role_prior(descriptor, TOTAL_LAYERS)
    effective_threshold = LOW_RANK_THRESHOLD_BASE - role_prior  # Bonus lowers threshold
    # Phase 3.5.1: Clamp to valid range to prevent degenerate routing
    effective_threshold = min(0.99, max(0.01, effective_threshold))

    # Phase 3.5.1: Warn if layer_idx >= total_layers (likely misconfiguration)
    layer_band = get_layer_band(descriptor.layer_idx, TOTAL_LAYERS)
    if descriptor.layer_idx >= 0 and descriptor.layer_idx >= TOTAL_LAYERS:
        import warnings
        warnings.warn(
            f"layer_idx {descriptor.layer_idx} >= total_layers {TOTAL_LAYERS}, "
            f"forcing band='late' for {descriptor.tensor_name}"
        )
        layer_band = "late"

    # Extract routing metrics from profile
    E10 = profile.energy_at_10pct
    row_coh = profile.row_coherence
    col_coh = profile.col_coherence
    C_struct = max(row_coh, col_coh)
    struct_axis = getattr(profile, "struct_axis", "none")

    routing_metrics = {
        "energy_at_10pct": float(E10),
        "row_coherence": float(row_coh),
        "col_coherence": float(col_coh),
        "C_struct": float(C_struct),
        "struct_axis": struct_axis,
        "condition_number": float(getattr(profile, "condition_number", 0.0)),
        "sparsity_ratio": float(getattr(profile, "sparsity_ratio", 0.0)),
        # Phase 3.5: Add role prior info to metrics
        "role_prior": float(role_prior),
        "effective_threshold": float(effective_threshold),
        "layer_band": layer_band,  # Phase 3.5.1: Use pre-computed band (may be forced to "late")
    }

    # Route to bucket (using effective threshold with role prior)
    is_low_rank = E10 > effective_threshold
    is_structured = C_struct >= STRUCT_THRESHOLD

    if is_low_rank:
        if is_structured:
            # Phase 3.5: Split bucket 2 into 2A (row) and 2B (col)
            if struct_axis == "row":
                bucket = "2A"
                codec_id = "svd_row_aware"
                base_params = codec_overrides.get("2A", DEFAULT_CODEC_PARAMS["2A"]).copy()
                reason = (
                    f"Low-rank (E@10%={E10:.0%}>{effective_threshold:.0%}) + "
                    f"row-structured (C={C_struct:.2f}>={STRUCT_THRESHOLD})"
                )
            elif struct_axis == "col":
                bucket = "2B"
                codec_id = "svd_col_aware"
                base_params = codec_overrides.get("2B", DEFAULT_CODEC_PARAMS["2B"]).copy()
                reason = (
                    f"Low-rank (E@10%={E10:.0%}>{effective_threshold:.0%}) + "
                    f"col-structured (C={C_struct:.2f}>={STRUCT_THRESHOLD})"
                )
            else:  # isotropic or none - use legacy bucket 2
                bucket = 2
                codec_id = "svd_axis_aware"
                base_params = codec_overrides.get(2, DEFAULT_CODEC_PARAMS[2]).copy()
                base_params["axis"] = struct_axis
                reason = (
                    f"Low-rank (E@10%={E10:.0%}>{effective_threshold:.0%}) + "
                    f"structured (C={C_struct:.2f}>={STRUCT_THRESHOLD}) on {struct_axis} axis"
                )

            # Add role prior explanation if applied
            if role_prior != 0:
                prior_str = f"+{role_prior:.2f}" if role_prior > 0 else f"{role_prior:.2f}"
                reason += f" [prior={prior_str}]"
        else:
            bucket = 1
            codec_id = "svd_lowrank"
            base_params = codec_overrides.get(1, DEFAULT_CODEC_PARAMS[1]).copy()
            reason = (
                f"Low-rank (E@10%={E10:.0%}>{effective_threshold:.0%}) + "
                f"unstructured (C={C_struct:.2f}<{STRUCT_THRESHOLD})"
            )
            if role_prior != 0:
                prior_str = f"+{role_prior:.2f}" if role_prior > 0 else f"{role_prior:.2f}"
                reason += f" [prior={prior_str}]"
    else:
        if is_structured:
            bucket = 4
            codec_id = "dct_quant"
            base_params = codec_overrides.get(4, DEFAULT_CODEC_PARAMS[4]).copy()
            base_params["axis"] = struct_axis
            reason = (
                f"High-rank (E@10%={E10:.0%}<={effective_threshold:.0%}) + "
                f"structured (C={C_struct:.2f}>={STRUCT_THRESHOLD}) on {struct_axis} axis"
            )
        else:
            bucket = 3
            codec_id = "kmeans_vq"
            base_params = codec_overrides.get(3, DEFAULT_CODEC_PARAMS[3]).copy()
            reason = (
                f"High-rank (E@10%={E10:.0%}<={effective_threshold:.0%}) + "
                f"unstructured (C={C_struct:.2f}<{STRUCT_THRESHOLD}) -> VQ"
            )

    return CompressionDecision(
        bucket=bucket,
        bucket_name=BUCKET_NAMES.get(bucket, f"bucket_{bucket}"),
        codec_id=codec_id,
        params=base_params,
        struct_axis=struct_axis,
        routing_reason=reason,
        routing_metrics=routing_metrics,
        role_prior=role_prior,
        effective_threshold=effective_threshold,
    )


def route_compression_batch(
    tensors: Dict[str, np.ndarray],
    descriptors: Dict[str, "TensorDescriptor"],
    profiles: Dict[str, Any],
    thresholds: Optional[Dict[str, float]] = None,
) -> Dict[str, CompressionDecision]:
    """
    Route a batch of tensors to compression strategies.

    Phase 3.5.1: Auto-detects total_layers from tensor names if not provided.

    Args:
        tensors: Dict mapping name -> tensor
        descriptors: Dict mapping name -> TensorDescriptor
        profiles: Dict mapping name -> TensorProfile
        thresholds: Override default routing thresholds

    Returns:
        Dict mapping name -> CompressionDecision
    """
    # Phase 3.5.1: Auto-detect total_layers if not provided
    if thresholds is None:
        thresholds = DEFAULT_THRESHOLDS.copy()
    else:
        thresholds = thresholds.copy()

    if "total_layers" not in thresholds or thresholds["total_layers"] == 32:
        detected = detect_total_layers(list(tensors.keys()))
        thresholds["total_layers"] = detected

    decisions = {}
    for name, tensor in tensors.items():
        descriptor = descriptors.get(name)
        profile = profiles.get(name)

        if profile is None:
            # Can't route without profile - use conservative default
            decisions[name] = CompressionDecision(
                bucket=3,
                bucket_name=BUCKET_NAMES[3],
                codec_id="kmeans_vq",
                params={"n_bins": 256, "percentile": 99.9, "fallback": True},
                struct_axis="none",
                routing_reason="No TensorProfile - fallback to k-means VQ",
                routing_metrics={"profile": "missing"},
                role_prior=0.0,
                effective_threshold=0.70,
            )
        else:
            decisions[name] = route_compression(
                tensor, descriptor, profile, thresholds
            )

    return decisions


def print_routing_summary(decisions: Dict[str, CompressionDecision]) -> None:
    """Print a summary of routing decisions. Handles Phase 3.5 buckets 2A/2B."""
    # Phase 3.5: Track both numeric and 2A/2B buckets
    bucket_counts: Dict[Union[int, str], int] = {1: 0, "2A": 0, "2B": 0, 2: 0, 3: 0, 4: 0}
    for decision in decisions.values():
        bucket = decision.bucket
        if bucket in bucket_counts:
            bucket_counts[bucket] += 1
        else:
            # Unknown bucket - map to closest
            bucket_counts[3] += 1

    print("\n" + "=" * 60)
    print("COMPRESSION ROUTING SUMMARY (Phase 3.5)")
    print("=" * 60)

    # Print in order: 1, 2A, 2B, 2, 3, 4
    bucket_order = [1, "2A", "2B", 2, 3, 4]
    total_bucket_2 = bucket_counts["2A"] + bucket_counts["2B"] + bucket_counts[2]

    for bucket in bucket_order:
        count = bucket_counts[bucket]
        if count == 0 and bucket in (2, "2A", "2B"):
            # Skip zero-count bucket 2 variants to reduce clutter
            continue

        name = BUCKET_NAMES.get(bucket, f"bucket_{bucket}")
        pct = 100 * count / len(decisions) if decisions else 0

        # Get codec for this bucket
        params = DEFAULT_CODEC_PARAMS.get(bucket, {"codec_id": "unknown"})
        codec = params.get("codec_id", "unknown")

        bar = "#" * int(pct / 2)
        bucket_str = str(bucket).rjust(2)
        print(f"  Bucket {bucket_str} ({name:25s}): {count:3d} ({pct:5.1f}%) |{bar}")

    # Show combined bucket 2 total if there are any
    if total_bucket_2 > 0:
        pct = 100 * total_bucket_2 / len(decisions) if decisions else 0
        print(f"  -------- Bucket 2 total (SVD): {total_bucket_2:3d} ({pct:5.1f}%)")

    print()
    print(f"  Total tensors: {len(decisions)}")


def _test_routing():
    """Internal test of routing logic including Phase 3.5 features."""
    from dataclasses import dataclass

    @dataclass
    class MockProfile:
        energy_at_10pct: float
        row_coherence: float
        col_coherence: float
        struct_axis: str
        condition_number: float = 1.0
        sparsity_ratio: float = 0.0

    @dataclass
    class MockDescriptor:
        tensor_name: str
        layer_idx: int
        module_family: str
        projection: str
        shape: tuple
        dtype: str

    # Phase 3.5: Test cases with 2A/2B bucket split
    # (E@10%, row_coh, col_coh, struct_axis, expected_bucket)
    test_cases = [
        (0.85, 0.10, 0.15, "none", 1),        # Low-rank + unstructured -> bucket 1
        (0.82, 0.45, 0.20, "row", "2A"),      # Low-rank + row-structured -> bucket 2A
        (0.75, 0.15, 0.50, "col", "2B"),      # Low-rank + col-structured -> bucket 2B
        (0.78, 0.40, 0.40, "isotropic", 2),   # Low-rank + isotropic -> bucket 2 (legacy)
        (0.50, 0.10, 0.10, "none", 3),        # High-rank + unstructured -> bucket 3
        (0.40, 0.45, 0.15, "row", 4),         # High-rank + row-structured -> bucket 4
        (0.30, 0.10, 0.55, "col", 4),         # High-rank + col-structured -> bucket 4
    ]

    print("Testing compression routing (Phase 3.5)...")
    print("\n--- Basic 2A/2B Bucket Split Tests ---")
    passed = 0
    failed = 0

    for E10, row_coh, col_coh, struct_axis, expected_bucket in test_cases:
        profile = MockProfile(
            energy_at_10pct=E10,
            row_coherence=row_coh,
            col_coherence=col_coh,
            struct_axis=struct_axis,
        )
        descriptor = MockDescriptor(
            tensor_name="test",
            layer_idx=0,
            module_family="test",
            projection="test",
            shape=(100, 100),
            dtype="float16",
        )
        tensor = np.zeros((100, 100))

        decision = route_compression(tensor, descriptor, profile)

        if decision.bucket == expected_bucket:
            passed += 1
            print(f"  PASS: E@10%={E10:.0%}, C_struct={max(row_coh,col_coh):.2f}, axis={struct_axis:9s} -> bucket {decision.bucket}")
        else:
            failed += 1
            print(f"  FAIL: E@10%={E10:.0%}, C_struct={max(row_coh,col_coh):.2f}, axis={struct_axis}")
            print(f"        Expected bucket {expected_bucket}, got bucket {decision.bucket}")
            print(f"        Reason: {decision.routing_reason}")

    # Test missing descriptor (Invariant C-ROUTE-001)
    print("\n--- Invariant C-ROUTE-001 (No Descriptor) ---")
    profile = MockProfile(energy_at_10pct=0.9, row_coherence=0.5, col_coherence=0.5, struct_axis="col")
    decision = route_compression(np.zeros((100, 100)), None, profile)
    if decision.bucket == 3 and decision.params.get("fallback"):
        passed += 1
        print(f"  PASS: No descriptor -> bucket 3 (fallback)")
    else:
        failed += 1
        print(f"  FAIL: No descriptor should route to bucket 3 with fallback=True")

    # Phase 3.5: Test role prior bias
    print("\n--- Role Prior Bias Tests ---")

    # attn_q at layer 0 (early) should get +0.10 bonus
    # With E@10% = 0.65, base threshold = 0.70, effective = 0.60
    # 0.65 > 0.60 -> should hit SVD (bucket 1 or 2)
    profile = MockProfile(energy_at_10pct=0.65, row_coherence=0.10, col_coherence=0.10, struct_axis="none")
    descriptor = MockDescriptor(
        tensor_name="blk.0.attn_q.weight",
        layer_idx=0,
        module_family="attn",
        projection="q",
        shape=(4096, 4096),
        dtype="float16",
    )
    decision = route_compression(np.zeros((100, 100)), descriptor, profile)
    expected_prior = 0.10
    if decision.bucket == 1 and abs(decision.role_prior - expected_prior) < 0.001:
        passed += 1
        print(f"  PASS: attn_q layer 0 (E@10%=65%) -> bucket {decision.bucket} [prior={decision.role_prior:+.2f}]")
    else:
        failed += 1
        print(f"  FAIL: attn_q layer 0 (E@10%=65%) should get prior=+0.10 and hit bucket 1")
        print(f"        Got bucket={decision.bucket}, prior={decision.role_prior}")

    # ffn_gate at layer 0 should get -0.10 penalty
    # With E@10% = 0.72, base threshold = 0.70, effective = 0.80
    # 0.72 <= 0.80 -> should NOT hit SVD (bucket 3)
    profile = MockProfile(energy_at_10pct=0.72, row_coherence=0.10, col_coherence=0.10, struct_axis="none")
    descriptor = MockDescriptor(
        tensor_name="blk.0.ffn_gate.weight",
        layer_idx=0,
        module_family="ffn",
        projection="gate",
        shape=(14336, 4096),
        dtype="float16",
    )
    decision = route_compression(np.zeros((100, 100)), descriptor, profile)
    expected_prior = -0.10
    if decision.bucket == 3 and abs(decision.role_prior - expected_prior) < 0.001:
        passed += 1
        print(f"  PASS: ffn_gate layer 0 (E@10%=72%) -> bucket {decision.bucket} [prior={decision.role_prior:+.2f}]")
    else:
        failed += 1
        print(f"  FAIL: ffn_gate layer 0 (E@10%=72%) should get prior=-0.10 and hit bucket 3")
        print(f"        Got bucket={decision.bucket}, prior={decision.role_prior}")

    # Test layer band: attn_q at layer 20 (mid) should get +0.05 bonus
    profile = MockProfile(energy_at_10pct=0.66, row_coherence=0.10, col_coherence=0.10, struct_axis="none")
    descriptor = MockDescriptor(
        tensor_name="blk.20.attn_q.weight",
        layer_idx=20,
        module_family="attn",
        projection="q",
        shape=(4096, 4096),
        dtype="float16",
    )
    decision = route_compression(np.zeros((100, 100)), descriptor, profile)
    expected_prior = 0.05
    # effective_threshold = 0.70 - 0.05 = 0.65, so 0.66 > 0.65 -> bucket 1
    if decision.bucket == 1 and abs(decision.role_prior - expected_prior) < 0.001:
        passed += 1
        print(f"  PASS: attn_q layer 20 (E@10%=66%) -> bucket {decision.bucket} [prior={decision.role_prior:+.2f}]")
    else:
        failed += 1
        print(f"  FAIL: attn_q layer 20 (E@10%=66%) should get prior=+0.05 and hit bucket 1")
        print(f"        Got bucket={decision.bucket}, prior={decision.role_prior}")

    # Test layer band: attn_q at layer 28 (late) should get no bonus
    profile = MockProfile(energy_at_10pct=0.69, row_coherence=0.10, col_coherence=0.10, struct_axis="none")
    descriptor = MockDescriptor(
        tensor_name="blk.28.attn_q.weight",
        layer_idx=28,
        module_family="attn",
        projection="q",
        shape=(4096, 4096),
        dtype="float16",
    )
    decision = route_compression(np.zeros((100, 100)), descriptor, profile)
    expected_prior = 0.0
    # effective_threshold = 0.70 - 0.0 = 0.70, so 0.69 <= 0.70 -> bucket 3
    if decision.bucket == 3 and abs(decision.role_prior - expected_prior) < 0.001:
        passed += 1
        print(f"  PASS: attn_q layer 28 (E@10%=69%) -> bucket {decision.bucket} [prior={decision.role_prior:+.2f}]")
    else:
        failed += 1
        print(f"  FAIL: attn_q layer 28 (E@10%=69%) should get prior=0.00 and hit bucket 3")
        print(f"        Got bucket={decision.bucket}, prior={decision.role_prior}")

    print(f"\nRouting tests: {passed} passed, {failed} failed")
    return failed == 0


if __name__ == "__main__":
    import sys

    print("CompressionRouter - Phase 3.5 (Role Prior Bias + 2A/2B Split)")
    print("=" * 60)

    success = _test_routing()

    # Demo routing
    print("\n" + "=" * 60)
    print("ROUTING DEMO (Phase 3.5)")
    print("=" * 60)

    from dataclasses import dataclass

    @dataclass
    class MockProfile:
        energy_at_10pct: float
        row_coherence: float
        col_coherence: float
        struct_axis: str
        condition_number: float = 1.0
        sparsity_ratio: float = 0.0

    @dataclass
    class MockDescriptor:
        tensor_name: str
        layer_idx: int
        module_family: str
        projection: str
        shape: tuple
        dtype: str

    # Simulate different tensor types - showcasing role prior and axis-aware routing
    demo_tensors = {
        # attn_q at layer 0 with col structure -> 2B + bonus
        "blk.0.attn_q.weight": {
            "profile": MockProfile(0.82, 0.12, 0.48, "col"),
            "descriptor": MockDescriptor("blk.0.attn_q.weight", 0, "attn", "q", (4096, 4096), "float16"),
        },
        # attn_q at layer 0 with row structure -> 2A + bonus
        "blk.0.attn_q.weight_row": {
            "profile": MockProfile(0.75, 0.45, 0.12, "row"),
            "descriptor": MockDescriptor("blk.0.attn_q.weight", 0, "attn", "q", (4096, 4096), "float16"),
        },
        # attn_k gets penalty -> borderline goes to bucket 3
        "blk.0.attn_k.weight": {
            "profile": MockProfile(0.73, 0.15, 0.12, "none"),
            "descriptor": MockDescriptor("blk.0.attn_k.weight", 0, "attn", "k", (1024, 4096), "float16"),
        },
        # FFN gets penalty -> definitely bucket 3
        "blk.0.ffn_down.weight": {
            "profile": MockProfile(0.72, 0.10, 0.12, "none"),
            "descriptor": MockDescriptor("blk.0.ffn_down.weight", 0, "ffn", "down", (4096, 14336), "float16"),
        },
        # Late layer attn_q -> no bonus, goes to bucket 3
        "blk.28.attn_q.weight": {
            "profile": MockProfile(0.55, 0.18, 0.22, "isotropic"),
            "descriptor": MockDescriptor("blk.28.attn_q.weight", 28, "attn", "q", (4096, 4096), "float16"),
        },
    }

    decisions = {}
    for name, data in demo_tensors.items():
        tensor = np.zeros(data["descriptor"].shape)
        decision = route_compression(tensor, data["descriptor"], data["profile"])
        decisions[name] = decision
        print(f"\n{name}:")
        print(f"  -> Bucket {decision.bucket} ({decision.bucket_name})")
        print(f"  -> Codec: {decision.codec_id}")
        print(f"  -> Role prior: {decision.role_prior:+.2f}, effective threshold: {decision.effective_threshold:.0%}")
        print(f"  -> Reason: {decision.routing_reason}")

    print_routing_summary(decisions)

    sys.exit(0 if success else 1)
