"""
Layer Graph Analysis — Determine adjacent layer pairs for paired transforms.

Obfuscation transforms (permutation, rotation, scaling) require knowing which
weight matrices are adjacent: the output of one layer feeds into the next.
Applying a permutation to layer L without the inverse on layer L+1 would break
the model. This module figures out those pairings.
"""

from __future__ import annotations

import re
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple

import numpy as np

from model_osmosis.format.osmo_format import ArchitectureSpec, LayerType


@dataclass
class TransformGroup:
    """A group of parameters that must be transformed together.

    Represents a hidden dimension shared between two adjacent weight matrices.
    Any permutation/rotation/scaling applied to this dimension must be
    consistently applied to all tensors in the group.
    """

    layer_name: str
    """Human-readable identifier for this group."""

    weight_key: str
    """Parameter key for the current layer's weight matrix."""

    bias_key: Optional[str]
    """Parameter key for the current layer's bias (if any)."""

    next_weight_key: str
    """Parameter key for the next layer's weight matrix."""

    parallel_weight_keys: List[str] = field(default_factory=list)
    """Additional weight matrices that share the same output dimension
    and must receive the same transform (e.g. up_proj in SwiGLU)."""

    norm_keys: List[str] = field(default_factory=list)
    """Parameter keys for normalization layers (gamma/beta) in this dimension."""

    dim_size: int = 0
    """Size of the shared hidden dimension."""

    axis_current: int = 0
    """Axis of weight_key along which to transform (0=rows for PyTorch Linear)."""

    axis_next: int = 1
    """Axis of next_weight_key along which to apply inverse transform."""


class LayerGraph:
    """Analyzes parameter names and shapes to find adjacent layer pairs.

    Supports:
    - Sequential models (numbered layers: 0.weight, 2.weight, ...)
    - Transformer patterns (q_proj, k_proj, fc1→fc2, etc.)
    - Expert-prefixed MoE models (expert_0/layer.weight)
    - Generic shape-based adjacency detection
    """

    def __init__(
        self,
        parameters: Dict[str, np.ndarray],
        architecture: Optional[ArchitectureSpec] = None,
    ):
        self.parameters = parameters
        self.architecture = architecture
        self._weight_keys: List[str] = []
        self._bias_keys: Dict[str, str] = {}
        self._norm_keys: Dict[str, List[str]] = {}
        self._classify_parameters()

    def _classify_parameters(self) -> None:
        """Separate weights, biases, and normalization parameters."""
        for key in self.parameters:
            arr = self.parameters[key]
            if arr.ndim < 2:
                # 1D tensor — could be bias or norm parameter
                if self._is_bias_key(key):
                    # Map bias to its weight key
                    weight_key = self._bias_to_weight_key(key)
                    if weight_key:
                        self._bias_keys[weight_key] = key
                elif self._is_norm_key(key):
                    # Group norm params by their layer prefix
                    prefix = self._get_layer_prefix(key)
                    self._norm_keys.setdefault(prefix, []).append(key)
            else:
                # 2D+ tensor — weight matrix
                self._weight_keys.append(key)

    def _is_bias_key(self, key: str) -> bool:
        return key.endswith(".bias") or key.endswith("_bias")

    def _is_norm_key(self, key: str) -> bool:
        norm_patterns = [
            "layernorm", "layer_norm", "ln_", "rmsnorm",
            "norm.weight", "norm.bias", "norm_weight", "norm_bias",
            "gamma", "beta",
        ]
        lower = key.lower()
        return any(p in lower for p in norm_patterns)

    def _bias_to_weight_key(self, bias_key: str) -> Optional[str]:
        """Convert a bias key to its corresponding weight key."""
        if bias_key.endswith(".bias"):
            candidate = bias_key[:-5] + ".weight"
            if candidate in self.parameters:
                return candidate
        if bias_key.endswith("_bias"):
            candidate = bias_key[:-5] + "_weight"
            if candidate in self.parameters:
                return candidate
        return None

    def _get_layer_prefix(self, key: str) -> str:
        """Extract the layer prefix from a parameter name."""
        # "layers.0.attention.q_proj.weight" → "layers.0.attention.q_proj"
        parts = key.rsplit(".", 1)
        return parts[0] if len(parts) > 1 else key

    def build_transform_groups(self) -> List[TransformGroup]:
        """Build the list of transform groups from parameter analysis.

        Returns:
            List of TransformGroup objects, each representing an adjacent
            weight pair that can be jointly transformed.
        """
        groups: List[TransformGroup] = []

        # Try expert-based grouping first (for stacked/MoE models)
        expert_groups = self._find_expert_groups()
        if expert_groups:
            for prefix, keys in expert_groups.items():
                sub_groups = self._find_sequential_pairs(keys, prefix)
                groups.extend(sub_groups)
            # Also handle non-expert parameters
            non_expert = [k for k in self._weight_keys
                          if not any(k.startswith(p) for p in expert_groups)]
            if non_expert:
                groups.extend(self._find_sequential_pairs(non_expert, "shared"))
            return self._deduplicate_groups(groups)

        # Try transformer-style pairing
        transformer_groups = self._find_transformer_pairs()
        if transformer_groups:
            return self._deduplicate_groups(transformer_groups)

        # Try embedding ↔ lm_head pairing
        embedding_groups = self._find_embedding_pairs()
        if embedding_groups:
            groups.extend(embedding_groups)

        # Fallback to shape-based pairing if nothing found
        if not groups:
            groups = self._find_shape_based_pairs()

        return self._deduplicate_groups(groups)

    def _find_expert_groups(self) -> Dict[str, List[str]]:
        """Detect expert-prefixed parameters (expert_0/, expert_1/, etc.)."""
        expert_pattern = re.compile(r"^(expert[_s]?\d+)[./]")
        groups: Dict[str, List[str]] = {}
        for key in self._weight_keys:
            match = expert_pattern.match(key)
            if match:
                prefix = match.group(1)
                groups.setdefault(prefix, []).append(key)
        return groups

    def _find_transformer_pairs(self) -> List[TransformGroup]:
        """Find pairs in transformer-style architectures.

        Detects patterns like:
        - fc1/fc2 (FFN intermediate dimension)
        - gate_proj/up_proj → down_proj (GLU-style FFN)
        - q_proj, k_proj, v_proj sharing input dimension
        """
        groups: List[TransformGroup] = []

        # Group parameters by transformer layer
        layer_params: Dict[str, Dict[str, str]] = {}
        layer_pattern = re.compile(
            r"((?:layers?|blocks?|transformer\.h|model\.layers)"
            r"[._]\d+)"
        )

        for key in self._weight_keys:
            match = layer_pattern.search(key)
            if match:
                layer_id = match.group(1)
                # Extract the sub-key (e.g., "mlp.fc1.weight")
                sub_key = key[match.end():].lstrip("./")
                layer_params.setdefault(layer_id, {})[sub_key] = key

        for layer_id, sub_keys in layer_params.items():
            # --- Standard FFN: fc1 → fc2 ---
            # GPT-2: mlp.c_fc → mlp.c_proj
            # BERT:  intermediate.dense → output.dense
            # Falcon: mlp.dense_h_to_4h → mlp.dense_4h_to_h
            # T5:    DenseReluDense.wi → DenseReluDense.wo
            fc1_key = self._find_subkey(sub_keys, [
                "mlp.fc1", "mlp.c_fc", "mlp.dense_h_to_4h",
                "ffn.fc1", "feed_forward.w1", "intermediate.dense",
                "DenseReluDense.wi", "fc_in",
            ])
            fc2_key = self._find_subkey(sub_keys, [
                "mlp.fc2", "mlp.c_proj", "mlp.dense_4h_to_h",
                "ffn.fc2", "feed_forward.w2", "output.dense",
                "DenseReluDense.wo", "fc_out",
            ])
            if fc1_key and fc2_key:
                fc1_full = sub_keys[fc1_key]
                fc2_full = sub_keys[fc2_key]
                fc1_arr = self.parameters[fc1_full]
                fc2_arr = self.parameters[fc2_full]

                # Auto-detect axis convention:
                # Standard Linear: fc1=[out, in], fc2=[out, in] → shared dim is fc1.shape[0] = fc2.shape[1]
                # GPT-2 Conv1D:    fc1=[in, out], fc2=[in, out] → shared dim is fc1.shape[1] = fc2.shape[0]
                ax_cur, ax_next, dim = self._infer_shared_dim(fc1_arr, fc2_arr)
                if dim > 0:
                    groups.append(TransformGroup(
                        layer_name=f"{layer_id}.ffn",
                        weight_key=fc1_full,
                        bias_key=self._bias_keys.get(fc1_full),
                        next_weight_key=fc2_full,
                        dim_size=dim,
                        axis_current=ax_cur,
                        axis_next=ax_next,
                    ))

            # --- GLU-style FFN: gate + up → down ---
            # LLaMA/Mistral/Qwen: mlp.gate_proj + mlp.up_proj → mlp.down_proj
            # T5 GLU: DenseReluDense.wi_0 + DenseReluDense.wi_1 → DenseReluDense.wo
            # Phi:    mlp.fc1 + mlp.fc2 → mlp.fc_out (handled above as fc1/fc2)
            gate_key = self._find_subkey(sub_keys, [
                "mlp.gate_proj", "mlp.w1", "feed_forward.gate",
                "DenseReluDense.wi_0", "mlp.gate",
            ])
            up_key = self._find_subkey(sub_keys, [
                "mlp.up_proj", "mlp.w3", "feed_forward.up",
                "DenseReluDense.wi_1", "mlp.up",
            ])
            down_key = self._find_subkey(sub_keys, [
                "mlp.down_proj", "mlp.w2", "feed_forward.down",
                "DenseReluDense.wo", "mlp.down",
            ])
            if gate_key and down_key and not fc1_key:
                gate_full = sub_keys[gate_key]
                down_full = sub_keys[down_key]
                gate_arr = self.parameters[gate_full]
                down_arr = self.parameters[down_full]

                ax_cur, ax_next, dim = self._infer_shared_dim(gate_arr, down_arr)
                if dim > 0:
                    parallel = []
                    if up_key:
                        parallel.append(sub_keys[up_key])
                    groups.append(TransformGroup(
                        layer_name=f"{layer_id}.glu_ffn",
                        weight_key=gate_full,
                        bias_key=self._bias_keys.get(gate_full),
                        next_weight_key=down_full,
                        parallel_weight_keys=parallel,
                        dim_size=dim,
                        axis_current=ax_cur,
                        axis_next=ax_next,
                    ))

            # --- Attention: q/k/v_proj → o_proj ---
            # The output dimension of q/k/v feeds into o_proj's input.
            # q_proj and k_proj/v_proj may have different output dims (GQA),
            # so we only pair q_proj → o_proj (same hidden dim) and treat
            # k_proj/v_proj as parallel only if they share the same output dim.
            q_key = self._find_subkey(sub_keys, [
                "self_attn.q_proj", "attention.q_proj", "attn.q_proj",
                "self_attention.query", "attention.query",
            ])
            o_key = self._find_subkey(sub_keys, [
                "self_attn.o_proj", "attention.o_proj", "attn.o_proj",
                "self_attention.dense", "attention.dense",
                "self_attn.out_proj", "attention.out_proj", "attn.out_proj",
            ])
            if q_key and o_key:
                q_full = sub_keys[q_key]
                o_full = sub_keys[o_key]
                q_arr = self.parameters[q_full]
                o_arr = self.parameters[o_full]

                ax_cur, ax_next, dim = self._infer_shared_dim(q_arr, o_arr)
                if dim > 0:
                    # Find k/v projections that share the same output dim
                    parallel = []
                    for kv_patterns in [
                        ["self_attn.k_proj", "attention.k_proj", "attn.k_proj",
                         "self_attention.key", "attention.key"],
                        ["self_attn.v_proj", "attention.v_proj", "attn.v_proj",
                         "self_attention.value", "attention.value"],
                    ]:
                        kv_key = self._find_subkey(sub_keys, kv_patterns)
                        if kv_key:
                            kv_full = sub_keys[kv_key]
                            kv_arr = self.parameters[kv_full]
                            # Check the same axis has the same dim
                            if kv_arr.shape[ax_cur] == dim:
                                parallel.append(kv_full)

                    groups.append(TransformGroup(
                        layer_name=f"{layer_id}.attn",
                        weight_key=q_full,
                        bias_key=self._bias_keys.get(q_full),
                        next_weight_key=o_full,
                        parallel_weight_keys=parallel,
                        dim_size=dim,
                        axis_current=ax_cur,
                        axis_next=ax_next,
                    ))

        return groups

    @staticmethod
    def _infer_shared_dim(
        w_first: np.ndarray,
        w_second: np.ndarray,
    ) -> Tuple[int, int, int]:
        """Infer the shared intermediate dimension between two weight matrices.

        For adjacent layers (e.g. fc1 → fc2), the shared dimension is the
        *intermediate* (expanded) dimension — typically the larger one.

        Handles both conventions:
        - Standard Linear: W1=[out, in], W2=[out, in] → shared = W1.rows, W2.cols
        - GPT-2 Conv1D:    W1=[in, out], W2=[in, out] → shared = W1.cols, W2.rows

        When both conventions match (square-ish matrices), prefer the larger
        shared dimension since FFN intermediate dims are typically expanded.

        Returns (axis_first, axis_second, dim_size). Returns (0, 0, 0) if
        no shared dimension is found.
        """
        if w_first.ndim < 2 or w_second.ndim < 2:
            return (0, 0, 0)

        candidates = []

        # Standard: out_dim of first (axis 0) matches in_dim of second (axis 1)
        if w_first.shape[0] == w_second.shape[1]:
            candidates.append((0, 1, w_first.shape[0]))

        # Conv1D: out_dim of first (axis 1) matches in_dim of second (axis 0)
        if w_first.shape[1] == w_second.shape[0]:
            candidates.append((1, 0, w_first.shape[1]))

        if not candidates:
            return (0, 0, 0)

        # Prefer the larger shared dimension (intermediate/expanded dim)
        return max(candidates, key=lambda c: c[2])

    def _find_subkey(
        self,
        sub_keys: Dict[str, str],
        patterns: List[str],
    ) -> Optional[str]:
        """Find a sub-key matching any of the given patterns."""
        for pattern in patterns:
            for sk in sub_keys:
                # Match with .weight suffix stripped or present
                clean = sk.replace(".weight", "").replace("_weight", "")
                if clean == pattern or clean.endswith("." + pattern):
                    return sk
        return None

    def _find_sequential_pairs(
        self,
        weight_keys: List[str],
        group_prefix: str,
    ) -> List[TransformGroup]:
        """Find sequential weight matrix pairs based on naming or order.

        Works for:
        - Sequential models: "0.weight", "2.weight", "4.weight"
        - Simple numbered: "layer_0.weight", "layer_1.weight"
        """
        groups: List[TransformGroup] = []
        if len(weight_keys) < 2:
            return groups

        # Try to extract numeric indices for ordering
        indexed: List[Tuple[float, str]] = []
        for key in weight_keys:
            arr = self.parameters[key]
            if arr.ndim < 2:
                continue
            # Extract all numbers from the key
            nums = re.findall(r"\d+", key)
            if nums:
                # Use the last number as the ordering index
                idx = float(nums[-1])
                # Add fractional based on position within same number
                idx += float(nums[0]) * 0.0001 if len(nums) > 1 else 0
                indexed.append((idx, key))
            else:
                indexed.append((len(indexed), key))

        # Sort by index
        indexed.sort(key=lambda x: x[0])

        # Create pairs from consecutive weight matrices where dimensions match
        for i in range(len(indexed) - 1):
            _, key_a = indexed[i]
            _, key_b = indexed[i + 1]
            arr_a = self.parameters[key_a]
            arr_b = self.parameters[key_b]

            if arr_a.ndim < 2 or arr_b.ndim < 2:
                continue

            # PyTorch convention: Linear weight is [out_features, in_features]
            out_dim_a = arr_a.shape[0]
            in_dim_b = arr_b.shape[1]

            if out_dim_a == in_dim_b:
                # Find norm keys for this dimension
                norm_keys = self._find_norm_keys_for_dim(key_a, out_dim_a)
                groups.append(TransformGroup(
                    layer_name=f"{group_prefix}.pair_{i}",
                    weight_key=key_a,
                    bias_key=self._bias_keys.get(key_a),
                    next_weight_key=key_b,
                    norm_keys=norm_keys,
                    dim_size=out_dim_a,
                    axis_current=0,
                    axis_next=1,
                ))

        return groups

    def _find_embedding_pairs(self) -> List[TransformGroup]:
        """Find embedding ↔ lm_head pairs that share a vocabulary dimension.

        In many language models, the embedding and output projection (lm_head)
        share the same weight matrix or have matching vocab dimensions.
        We can permute the embedding dimension if they are separate tensors.
        """
        groups: List[TransformGroup] = []

        embed_patterns = ["embed_tokens", "wte", "word_embeddings", "embedding"]
        head_patterns = ["lm_head", "output.weight", "cls.predictions"]

        embed_key = None
        head_key = None

        for key in self._weight_keys:
            lower = key.lower()
            if any(p in lower for p in embed_patterns):
                embed_key = key
            if any(p in lower for p in head_patterns):
                head_key = key

        if embed_key and head_key and embed_key != head_key:
            embed_arr = self.parameters[embed_key]
            head_arr = self.parameters[head_key]

            # Embedding: [vocab_size, hidden_dim]
            # lm_head:   [vocab_size, hidden_dim] (or transposed)
            if embed_arr.ndim == 2 and head_arr.ndim == 2:
                # Check if they share the hidden dimension
                # Embedding rows = vocab, cols = hidden
                # lm_head rows = vocab (out), cols = hidden (in) in PyTorch
                if embed_arr.shape[1] == head_arr.shape[1]:
                    # Shared hidden dimension — can permute it
                    groups.append(TransformGroup(
                        layer_name="embedding_head",
                        weight_key=embed_key,
                        bias_key=self._bias_keys.get(embed_key),
                        next_weight_key=head_key,
                        dim_size=embed_arr.shape[1],
                        axis_current=1,  # cols of embedding
                        axis_next=1,     # cols of lm_head (input dim)
                    ))

        return groups

    def _find_shape_based_pairs(self) -> List[TransformGroup]:
        """Fallback: pair weight matrices purely by matching dimensions."""
        groups: List[TransformGroup] = []
        used_as_next: set = set()

        for key_a in self._weight_keys:
            arr_a = self.parameters[key_a]
            if arr_a.ndim < 2:
                continue
            out_dim_a = arr_a.shape[0]

            for key_b in self._weight_keys:
                if key_b == key_a or key_b in used_as_next:
                    continue
                arr_b = self.parameters[key_b]
                if arr_b.ndim < 2:
                    continue
                in_dim_b = arr_b.shape[1]

                if out_dim_a == in_dim_b and out_dim_a > 1:
                    groups.append(TransformGroup(
                        layer_name=f"shape_pair.{key_a}->{key_b}",
                        weight_key=key_a,
                        bias_key=self._bias_keys.get(key_a),
                        next_weight_key=key_b,
                        dim_size=out_dim_a,
                        axis_current=0,
                        axis_next=1,
                    ))
                    used_as_next.add(key_b)
                    break  # Only pair each key_a once

        return groups

    def _find_norm_keys_for_dim(self, weight_key: str, dim: int) -> List[str]:
        """Find normalization parameters that share the given dimension."""
        prefix = self._get_layer_prefix(weight_key)
        result = []

        # Check norm keys with similar prefix
        for norm_prefix, keys in self._norm_keys.items():
            if self._prefixes_overlap(prefix, norm_prefix):
                for nk in keys:
                    if nk in self.parameters and self.parameters[nk].shape == (dim,):
                        result.append(nk)

        return result

    def _prefixes_overlap(self, prefix_a: str, prefix_b: str) -> bool:
        """Check if two parameter prefixes belong to the same layer block."""
        # Extract the common transformer layer prefix
        # e.g., "layers.0.mlp.fc1" and "layers.0.post_attention_layernorm"
        # both share "layers.0"
        parts_a = prefix_a.split(".")
        parts_b = prefix_b.split(".")
        # Check if they share at least 2 levels
        common = 0
        for pa, pb in zip(parts_a, parts_b):
            if pa == pb:
                common += 1
            else:
                break
        return common >= 2

    def _deduplicate_groups(self, groups: List[TransformGroup]) -> List[TransformGroup]:
        """Remove duplicate groups — each weight key appears at most once."""
        seen_weight: set = set()
        seen_next: set = set()
        unique: List[TransformGroup] = []
        for g in groups:
            if g.weight_key not in seen_weight and g.next_weight_key not in seen_next:
                seen_weight.add(g.weight_key)
                seen_next.add(g.next_weight_key)
                unique.append(g)
        return unique
