"""
Obfuscation Verification — Confirm functional equivalence after transforms.

Compares the forward-pass outputs of original and obfuscated parameters
to ensure the model still behaves identically.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, List, Optional

import numpy as np

from model_osmosis.format.osmo_format import ArchitectureSpec


@dataclass
class VerificationResult:
    """Result of functional equivalence verification."""

    passed: bool
    """Whether outputs are within tolerance."""

    max_absolute_error: float
    """Largest element-wise difference across all samples."""

    mean_absolute_error: float
    """Average element-wise difference across all samples."""

    max_relative_error: float
    """Largest relative difference (|a-b| / max(|a|,|b|)) across all samples."""

    num_samples: int
    """Number of samples used for verification."""

    details: str
    """Human-readable summary."""


class ObfuscationVerifier:
    """Verify that obfuscated parameters produce identical outputs."""

    def verify(
        self,
        original_params: Dict[str, np.ndarray],
        obfuscated_params: Dict[str, np.ndarray],
        architecture: Optional[ArchitectureSpec],
        sample_data: np.ndarray,
        tolerance: float = 1e-4,
    ) -> VerificationResult:
        """Compare forward pass outputs between original and obfuscated params.

        For simple sequential/MLP models, reconstructs the forward pass from
        numpy parameters. For complex models, falls back to PyTorch if available.

        Args:
            original_params: Original parameter dictionary.
            obfuscated_params: Obfuscated parameter dictionary.
            architecture: Architecture specification (optional, helps ordering).
            sample_data: Input samples to run through both models.
            tolerance: Maximum allowed absolute error.

        Returns:
            VerificationResult with pass/fail and error metrics.
        """
        # Try numpy forward pass first
        try:
            original_out = self._numpy_forward(
                original_params, architecture, sample_data
            )
            obfuscated_out = self._numpy_forward(
                obfuscated_params, architecture, sample_data
            )
        except Exception:
            # Fallback to PyTorch if available
            try:
                original_out, obfuscated_out = self._torch_forward(
                    original_params, obfuscated_params, architecture, sample_data
                )
            except Exception as e:
                return VerificationResult(
                    passed=False,
                    max_absolute_error=float("inf"),
                    mean_absolute_error=float("inf"),
                    max_relative_error=float("inf"),
                    num_samples=len(sample_data),
                    details=f"Could not verify: {e}",
                )

        # Compare outputs
        diff = np.abs(original_out - obfuscated_out)
        max_err = float(np.max(diff))
        mean_err = float(np.mean(diff))

        # Relative error: |a-b| / max(|a|, |b|), avoiding division by zero
        abs_max = np.maximum(np.abs(original_out), np.abs(obfuscated_out))
        with np.errstate(divide="ignore", invalid="ignore"):
            rel = np.where(abs_max > 0, diff / abs_max, 0.0)
        max_rel_err = float(np.max(rel))

        passed = max_err <= tolerance

        details = (
            f"Verification {'PASSED' if passed else 'FAILED'}: "
            f"max_error={max_err:.2e}, mean_error={mean_err:.2e}, "
            f"max_relative_error={max_rel_err:.2e}, "
            f"tolerance={tolerance:.2e}, samples={len(sample_data)}"
        )

        return VerificationResult(
            passed=passed,
            max_absolute_error=max_err,
            mean_absolute_error=mean_err,
            max_relative_error=max_rel_err,
            num_samples=len(sample_data),
            details=details,
        )

    def _numpy_forward(
        self,
        params: Dict[str, np.ndarray],
        architecture: Optional[ArchitectureSpec],
        inputs: np.ndarray,
    ) -> np.ndarray:
        """Reconstruct a simple sequential forward pass from numpy params.

        Assumes PyTorch Linear convention: y = x @ W^T + b
        Applies ReLU activation between layers (except the last).
        """
        # Order weight matrices
        weight_keys = self._get_ordered_weights(params, architecture)

        if not weight_keys:
            raise ValueError("No weight matrices found for forward pass")

        x = inputs.astype(np.float64)

        for i, wk in enumerate(weight_keys):
            w = params[wk].astype(np.float64)

            # Find corresponding bias
            bias_key = None
            if wk.endswith(".weight"):
                candidate = wk[:-7] + ".bias"
                if candidate in params:
                    bias_key = candidate
            elif wk.endswith("_weight"):
                candidate = wk[:-7] + "_bias"
                if candidate in params:
                    bias_key = candidate

            # Linear transform: y = x @ W^T + b
            if w.ndim == 2:
                x = x @ w.T
            else:
                raise ValueError(f"Unsupported weight shape: {w.shape}")

            if bias_key:
                x = x + params[bias_key].astype(np.float64)

            # Apply ReLU between layers (not on the last layer)
            if i < len(weight_keys) - 1:
                x = np.maximum(x, 0)

        return x

    def _get_ordered_weights(
        self,
        params: Dict[str, np.ndarray],
        architecture: Optional[ArchitectureSpec],
    ) -> List[str]:
        """Get weight matrix keys in forward-pass order."""
        weight_keys = [
            k for k, v in params.items()
            if v.ndim >= 2 and not self._is_norm_weight(k)
        ]

        if not weight_keys:
            return []

        # Try to sort by numeric index in the key
        def sort_key(k: str):
            import re
            nums = re.findall(r"\d+", k)
            return tuple(int(n) for n in nums) if nums else (999,)

        weight_keys.sort(key=sort_key)
        return weight_keys

    def _is_norm_weight(self, key: str) -> bool:
        """Check if a key is a normalization weight (not a linear weight).

        Uses component-level matching to avoid false positives like
        'normalize', 'normal', or 'normalization' in parameter names.
        """
        # Split on '.' and '_' to get individual components
        lower = key.lower()
        components = lower.replace("_", ".").split(".")
        norm_components = {
            "layernorm", "layer_norm", "rmsnorm", "groupnorm",
            "batchnorm", "instancenorm", "ln", "norm",
        }
        return bool(norm_components.intersection(components))

    def _torch_forward(
        self,
        original_params: Dict[str, np.ndarray],
        obfuscated_params: Dict[str, np.ndarray],
        architecture: Optional[ArchitectureSpec],
        sample_data: np.ndarray,
    ) -> tuple:
        """Fallback: use PyTorch to build and run models."""
        import torch
        import torch.nn as nn

        weight_keys = self._get_ordered_weights(original_params, architecture)
        if not weight_keys:
            raise ValueError("No weight matrices found")

        def build_model(params: Dict[str, np.ndarray]) -> nn.Sequential:
            layers = []
            for i, wk in enumerate(weight_keys):
                w = params[wk]
                out_dim, in_dim = w.shape[0], w.shape[1]
                linear = nn.Linear(in_dim, out_dim, bias=False)
                linear.weight.data = torch.from_numpy(w.copy()).float()

                bias_key = None
                if wk.endswith(".weight"):
                    candidate = wk[:-7] + ".bias"
                    if candidate in params:
                        bias_key = candidate
                if bias_key:
                    linear.bias = nn.Parameter(
                        torch.from_numpy(params[bias_key].copy()).float()
                    )

                layers.append(linear)
                if i < len(weight_keys) - 1:
                    layers.append(nn.ReLU())

            return nn.Sequential(*layers)

        model_orig = build_model(original_params)
        model_obf = build_model(obfuscated_params)

        x = torch.from_numpy(sample_data).float()
        with torch.no_grad():
            out_orig = model_orig(x).numpy()
            out_obf = model_obf(x).numpy()

        return out_orig, out_obf
