"""
Spectral Guard Implementation
============================

Monitors spectral properties of model weights to detect instabilities.
Provides spectral control mechanisms for maintaining numerical stability.
"""

import math
import time
from collections import defaultdict
from datetime import datetime
from typing import Any, TypedDict

try:
    from typing import NotRequired
except ImportError:  # Python 3.10 fallback
    from typing import NotRequired

import numpy as np
import torch

from invarlock.cli._evidence import maybe_dump_guard_evidence
from invarlock.core.api import Guard

from ._contracts import guard_assert


class SpectralPolicy(TypedDict, total=False):
    """Type definition for spectral guard policy configuration."""

    sigma_quantile: float
    contraction: NotRequired[float]  # Backward compatibility alias
    kappa: NotRequired[float]  # Legacy alias
    deadband: float
    scope: str
    correction_enabled: bool
    family_caps: dict[str, dict[str, float]]
    ignore_preview_inflation: bool
    max_caps: int
    multiple_testing: dict[str, Any]


def _default_family_caps() -> dict[str, dict[str, float]]:
    """Default per-family spectral z-score caps."""
    return {
        "ffn": {"kappa": 2.5},
        "attn": {"kappa": 2.8},
        "embed": {"kappa": 3.0},
        "other": {"kappa": 3.0},
    }


def _normalize_family_caps(
    caps: Any, *, default: bool = True
) -> dict[str, dict[str, float]]:
    """Normalize family cap configuration into canonical mapping."""

    if not isinstance(caps, dict) or not caps:
        return _default_family_caps() if default else {}

    normalized: dict[str, dict[str, float]] = {}
    for family, values in caps.items():
        entry: dict[str, float] = {}
        if isinstance(values, dict):
            for key, val in values.items():
                if isinstance(val, int | float) and math.isfinite(float(val)):
                    entry[str(key)] = float(val)
        elif isinstance(values, int | float) and math.isfinite(float(values)):
            entry["kappa"] = float(values)
        if entry:
            normalized[str(family)] = entry

    if normalized:
        return normalized

    return _default_family_caps() if default else {}


class SpectralGuard(Guard):
    """
    Spectral guard for monitoring weight matrix spectral properties.

    Tracks singular values and spectral norms to detect numerical instabilities.
    Provides automatic spectral control when violations are detected.
    """

    name = "spectral"

    def __init__(self, **kwargs):
        """Initialize spectral guard."""
        self.config = dict(kwargs)
        self.prepared = False
        self.baseline_metrics = {}
        self.events = []
        self.current_metrics = {}
        self.violations = []

        # Default configuration
        sigma_quantile = self.config.get("sigma_quantile")
        if sigma_quantile is None:
            for alias in ("contraction", "kappa"):
                if self.config.get(alias) is not None:
                    sigma_quantile = self.config[alias]
                    break
        if sigma_quantile is None:
            sigma_quantile = 0.95
        self.sigma_quantile = float(sigma_quantile)
        self.config["sigma_quantile"] = self.sigma_quantile
        self.config.pop("contraction", None)
        self.config.pop("kappa", None)
        self.deadband = kwargs.get("deadband", 0.10)
        self.scope = kwargs.get("scope", "all")  # 'all', 'ffn', 'attn'
        self.max_spectral_norm = kwargs.get("max_spectral_norm", 10.0)
        if self.max_spectral_norm is not None:
            self.max_spectral_norm = float(self.max_spectral_norm)
        self.config["max_spectral_norm"] = self.max_spectral_norm
        self.min_condition_number = kwargs.get("min_condition_number", 1e-12)
        self.correction_enabled = kwargs.get("correction_enabled", True)
        self.family_caps = _normalize_family_caps(
            kwargs.get("family_caps"), default=True
        )
        self.ignore_preview_inflation = kwargs.get("ignore_preview_inflation", True)
        self.max_caps = kwargs.get("max_caps", 5)
        self.multiple_testing = kwargs.get(
            "multiple_testing", {"method": "bh", "alpha": 0.05, "m": 4}
        )

        # Baseline and tracking structures
        self.baseline_sigmas: dict[str, float] = {}
        self.baseline_family_stats: dict[str, dict[str, float]] = {}
        self.module_family_map: dict[str, str] = {}
        self.latest_z_scores: dict[str, float] = {}
        self.pre_edit_z_scores: dict[str, float] = {}

    def _log_event(
        self, operation: str, level: str = "INFO", message: str = "", **data
    ):
        """Log an event with timestamp."""
        event = {
            "timestamp": datetime.utcnow().isoformat(),
            "component": "spectral_guard",
            "operation": operation,
            "level": level,
            "message": message,
            "data": data,
        }
        self.events.append(event)

    def _serialize_policy(self) -> dict[str, Any]:
        """Snapshot current guard policy for report serialization."""
        return {
            "scope": self.scope,
            "sigma_quantile": float(self.sigma_quantile),
            "deadband": float(self.deadband),
            "max_caps": int(self.max_caps),
            "max_spectral_norm": (
                float(self.max_spectral_norm)
                if self.max_spectral_norm is not None
                else None
            ),
            "family_caps": self.family_caps,
            "multiple_testing": self.multiple_testing,
            "correction_enabled": bool(self.correction_enabled),
            "ignore_preview_inflation": bool(self.ignore_preview_inflation),
        }

    def prepare(
        self, model: Any, adapter: Any, calib: Any, policy: dict[str, Any]
    ) -> dict[str, Any]:
        """
        Prepare spectral guard by capturing baseline spectral properties.

        Args:
            model: Model to prepare for
            adapter: ModelAdapter instance
            calib: Calibration data (unused for spectral analysis)
            policy: Policy configuration

        Returns:
            Preparation results
        """
        start_time = time.time()

        # Update configuration from policy
        if policy:
            sigma_value = policy.get("sigma_quantile")
            if sigma_value is None:
                alias_value = policy.get("contraction", policy.get("kappa"))
                if alias_value is not None:
                    sigma_value = alias_value
            if sigma_value is not None:
                self.sigma_quantile = float(sigma_value)
                policy["sigma_quantile"] = self.sigma_quantile
            policy.pop("contraction", None)
            policy.pop("kappa", None)
            self.config["sigma_quantile"] = self.sigma_quantile

            for key in [
                "sigma_quantile",
                "deadband",
                "scope",
                "max_spectral_norm",
                "correction_enabled",
                "max_caps",
            ]:
                if key in policy:
                    setattr(self, key, policy[key])
                    self.config[key] = policy[key]

            if self.max_spectral_norm is not None:
                self.max_spectral_norm = float(self.max_spectral_norm)
            self.config["max_spectral_norm"] = self.max_spectral_norm

            if "family_caps" in policy:
                self.family_caps = _normalize_family_caps(
                    policy["family_caps"], default=True
                )
                self.config["family_caps"] = self.family_caps

            if "ignore_preview_inflation" in policy:
                self.ignore_preview_inflation = bool(policy["ignore_preview_inflation"])
                self.config["ignore_preview_inflation"] = self.ignore_preview_inflation

            # Optional hydration of baseline stats from policy (e.g., baseline certificate)
            if "baseline_family_stats" in policy and isinstance(
                policy["baseline_family_stats"], dict
            ):
                self.baseline_family_stats = {
                    family: stats.copy()
                    for family, stats in policy["baseline_family_stats"].items()
                    if isinstance(stats, dict)
                }
                self.config["baseline_family_stats"] = self.baseline_family_stats
            mt_policy = policy.get("multiple_testing")
            if mt_policy is None:
                mt_policy = policy.get("multipletesting")
            if isinstance(mt_policy, dict):
                self.multiple_testing = mt_policy.copy()
                policy["multiple_testing"] = self.multiple_testing
                self.config["multiple_testing"] = self.multiple_testing
            policy.pop("multipletesting", None)

        self._log_event(
            "prepare",
            message=(
                f"Preparing spectral guard with scope={self.scope}, "
                f"sigma_quantile={self.sigma_quantile}"
            ),
        )

        try:
            # Capture baseline spectral properties
            self.baseline_sigmas = capture_baseline_sigmas(model, scope=self.scope)
            self.module_family_map = classify_model_families(
                model, scope=self.scope, existing=self.module_family_map
            )
            if not self.baseline_family_stats:
                self.baseline_family_stats = compute_family_stats(
                    self.baseline_sigmas, self.module_family_map
                )

            # Compute additional baseline metrics
            baseline_stats = scan_model_gains(model, scope=self.scope)
            summarized = _summarize_sigmas(self.baseline_sigmas)
            baseline_stats.update(summarized)

            # Store target sigma value
            self.target_sigma = auto_sigma_target(model, percentile=self.sigma_quantile)
            baseline_stats["target_sigma"] = self.target_sigma

            baseline_stats["family_stats"] = {
                family: stats.copy()
                for family, stats in self.baseline_family_stats.items()
            }
            baseline_stats["family_caps"] = {
                family: caps.copy() for family, caps in self.family_caps.items()
            }
            baseline_stats["module_sigmas"] = self.baseline_sigmas.copy()

            self.baseline_metrics = baseline_stats

            self.prepared = True
            preparation_time = time.time() - start_time

            self._log_event(
                "prepare_success",
                message=f"Prepared spectral guard with {len(self.baseline_metrics)} baseline metrics",
                baseline_metrics_count=len(self.baseline_metrics),
                target_sigma=self.target_sigma,
                preparation_time=preparation_time,
            )

            return {
                "ready": True,
                "baseline_metrics": self.baseline_metrics,
                "target_sigma": self.target_sigma,
                "scope": self.scope,
                "preparation_time": preparation_time,
            }

        except Exception as e:
            self.prepared = False
            self._log_event(
                "prepare_failed",
                level="ERROR",
                message=f"Failed to prepare spectral guard: {str(e)}",
                error=str(e),
            )

            return {
                "ready": False,
                "error": str(e),
                "preparation_time": time.time() - start_time,
            }

    def before_edit(self, model: Any) -> None:
        """Execute before edit (capture pre-edit state)."""
        if not self.prepared:
            self._log_event(
                "before_edit_skipped",
                level="WARN",
                message="Spectral guard not prepared, skipping pre-edit capture",
            )
            return

        # Capture pre-edit spectral state for comparison
        self.pre_edit_metrics = capture_baseline_sigmas(model, scope=self.scope)
        self.pre_edit_z_scores = compute_z_scores(
            self.pre_edit_metrics,
            self.baseline_family_stats,
            self.module_family_map,
            self.baseline_sigmas,
            deadband=self.deadband,
        )
        self._log_event("before_edit", message="Captured pre-edit spectral state")

    def after_edit(self, model: Any) -> None:
        """Execute after edit (detect violations and apply control)."""
        if not self.prepared:
            self._log_event(
                "after_edit_skipped",
                level="WARN",
                message="Spectral guard not prepared, skipping post-edit analysis",
            )
            return

        try:
            # Capture current spectral state
            self.current_metrics = capture_baseline_sigmas(model, scope=self.scope)

            # Detect violations
            violations = self._detect_spectral_violations(
                model, self.current_metrics, phase="after_edit"
            )
            self.violations = violations

            # Apply spectral control if violations detected and correction enabled
            if violations and self.correction_enabled:
                control_result = apply_spectral_control(
                    model,
                    policy={
                        "sigma_quantile": self.sigma_quantile,
                        "scope": self.scope,
                        "baseline_sigmas": self.baseline_sigmas,
                        "target_sigma": self.target_sigma,
                    },
                )

                self._log_event(
                    "spectral_control_applied",
                    message=f"Applied spectral control, violations: {len(violations)}",
                    violations_count=len(violations),
                    control_result=control_result,
                )

            self._log_event(
                "after_edit",
                message=f"Post-edit analysis complete, {len(violations)} violations detected",
            )

        except Exception as e:
            self._log_event(
                "after_edit_failed",
                level="ERROR",
                message=f"Post-edit spectral analysis failed: {str(e)}",
                error=str(e),
            )

    def _detect_spectral_violations(
        self, model: Any, metrics: dict[str, float], phase: str = "finalize"
    ) -> list[dict[str, Any]]:
        """Detect spectral property violations using per-family z-score caps."""
        violations: list[dict[str, Any]] = []
        latest_z: dict[str, float] = {}

        for name, module in model.named_modules():
            if not self._should_check_module(name, module):
                continue

            try:
                if hasattr(module, "weight") and module.weight.ndim == 2:
                    sigma_max = metrics.get(name)
                    if sigma_max is None:
                        sigma_max = compute_sigma_max(module.weight)

                    baseline_sigma = self.baseline_sigmas.get(name, self.target_sigma)
                    family = self.module_family_map.get(name)
                    if family is None:
                        family = classify_module_family(name, module)
                        self.module_family_map[name] = family

                    family_stats = self.baseline_family_stats.get(family, {})
                    cap_config = self.family_caps.get(family, {})
                    kappa_cap = float(cap_config.get("kappa", self.sigma_quantile))

                    z_score = compute_z_score_for_value(
                        sigma_max,
                        family_stats,
                        fallback_value=baseline_sigma,
                        deadband=self.deadband,
                    )
                    latest_z[name] = z_score

                    # Skip preview inflation if configured and not in final phase
                    if self.ignore_preview_inflation and phase == "after_edit":
                        continue

                    if z_score > kappa_cap:
                        violations.append(
                            {
                                "type": "family_z_cap",
                                "module": name,
                                "family": family,
                                "z_score": float(z_score),
                                "kappa": kappa_cap,
                                "sigma": float(sigma_max),
                                "baseline_sigma": float(baseline_sigma),
                                "message": (
                                    f"Family '{family}' z-score {z_score:.2f}"
                                    f" exceeds cap {kappa_cap:.2f}"
                                ),
                            }
                        )

                    if (
                        self.max_spectral_norm is not None
                        and sigma_max > self.max_spectral_norm
                    ):
                        threshold = float(self.max_spectral_norm)
                        violations.append(
                            {
                                "type": "max_spectral_norm",
                                "module": name,
                                "family": family,
                                "current_sigma": float(sigma_max),
                                "threshold": threshold,
                                "message": f"Spectral norm {sigma_max:.3f} exceeds maximum {threshold}",
                            }
                        )

                    # Condition number monitoring (warn only)
                    try:
                        U, S, V = torch.svd(module.weight.float())
                        if len(S) > 0:
                            condition_number = S[0].item() / max(S[-1].item(), 1e-12)
                            if S[-1].item() < self.min_condition_number:
                                violations.append(
                                    {
                                        "type": "ill_conditioned",
                                        "module": name,
                                        "family": family,
                                        "condition_number": float(condition_number),
                                        "min_singular_value": float(S[-1].item()),
                                        "threshold": float(self.min_condition_number),
                                        "message": f"Matrix is ill-conditioned, min singular value: {S[-1].item():.2e}",
                                    }
                                )
                    except Exception:
                        pass  # SVD failure is not a violation

            except Exception as e:
                self._log_event(
                    "violation_check_error",
                    level="WARN",
                    message=f"Failed to check module {name}: {str(e)}",
                    module=name,
                    error=str(e),
                )

        self.latest_z_scores = latest_z
        return violations

    def _should_check_module(self, name: str, module: Any) -> bool:
        """Determine if a module should be checked based on scope."""
        if not hasattr(module, "weight") or module.weight.ndim != 2:
            return False

        if self.scope == "all":
            return True
        elif self.scope == "attn":
            return any(
                keyword in name.lower()
                for keyword in ["attn", "attention", "self_attn"]
            )
        elif self.scope == "ffn":
            return any(
                keyword in name.lower()
                for keyword in ["mlp", "ffn", "feed_forward", "fc"]
            )

        return True

    def _compute_family_observability(
        self,
    ) -> tuple[dict[str, dict[str, float]], dict[str, list[dict[str, Any]]]]:
        """Generate per-family quantiles and top-|z| listings from latest z-scores."""
        family_scores: dict[str, list[float]] = defaultdict(list)
        family_modules: dict[str, list[tuple[float, str]]] = defaultdict(list)

        for module_name, z_value in (self.latest_z_scores or {}).items():
            family = self.module_family_map.get(module_name)
            if family is None:
                continue
            try:
                z_abs = abs(float(z_value))
            except (TypeError, ValueError):
                continue
            family_scores.setdefault(family, []).append(z_abs)
            family_modules.setdefault(family, []).append((z_abs, module_name))

        def _quantile(sorted_values: list[float], quantile: float) -> float:
            if not sorted_values:
                return 0.0
            if len(sorted_values) == 1:
                return sorted_values[0]
            position = (len(sorted_values) - 1) * quantile
            lower = math.floor(position)
            upper = math.ceil(position)
            if lower == upper:
                return sorted_values[int(position)]
            fraction = position - lower
            return (
                sorted_values[lower]
                + (sorted_values[upper] - sorted_values[lower]) * fraction
            )

        family_quantiles: dict[str, dict[str, float]] = {}
        for family, scores in family_scores.items():
            sorted_scores = sorted(scores)
            family_quantiles[family] = {
                "q95": _quantile(sorted_scores, 0.95),
                "q99": _quantile(sorted_scores, 0.99),
                "max": sorted_scores[-1] if sorted_scores else 0.0,
                "count": len(sorted_scores),
            }

        top_z_scores: dict[str, list[dict[str, Any]]] = {}
        for family, module_entries in family_modules.items():
            module_entries.sort(key=lambda item: item[0], reverse=True)
            top_entries: list[dict[str, Any]] = []
            for z_abs, module_name in module_entries[:3]:
                top_entries.append(
                    {"module": module_name, "z": float(z_abs), "family": family}
                )
            top_z_scores[family] = top_entries

        return family_quantiles, top_z_scores

    def validate(
        self, model: Any, adapter: Any, context: dict[str, Any]
    ) -> dict[str, Any]:
        """
        Validate model spectral properties.

        Args:
            model: Model to validate
            adapter: ModelAdapter instance
            context: Validation context

        Returns:
            Dictionary with validation results
        """
        try:
            if not self.prepared:
                # Auto-prepare if needed
                self.prepare(model, adapter, None, {})

            # Capture current spectral state
            current_metrics = capture_baseline_sigmas(model, scope=self.scope)

            # Detect violations (final validation phase)
            violations = self._detect_spectral_violations(
                model, current_metrics, phase="validate"
            )

            # Determine if passed under budget/fatal rules
            fatal_violation_types = {"max_spectral_norm", "ill_conditioned"}
            budgeted_violations = [
                violation
                for violation in violations
                if violation.get("type") not in fatal_violation_types
            ]
            fatal_violations = [
                violation
                for violation in violations
                if violation.get("type") in fatal_violation_types
            ]

            caps_applied = len(budgeted_violations)
            caps_exceeded = caps_applied > int(self.max_caps)
            passed = not fatal_violations and not caps_exceeded
            if fatal_violations or caps_exceeded:
                action = "abort"
            elif caps_applied > 0:
                action = "warn"
            else:
                action = "continue"

            # Compute overall metrics
            family_summary = summarize_family_z_scores(
                self.latest_z_scores, self.module_family_map, self.family_caps
            )
            metrics = {
                "modules_checked": len(current_metrics),
                "violations_found": len(violations),
                "budgeted_violations": caps_applied,
                "fatal_violations": len(fatal_violations),
                "max_spectral_norm": max(current_metrics.values())
                if current_metrics
                else 0.0,
                "mean_spectral_norm": np.mean(list(current_metrics.values()))
                if current_metrics
                else 0.0,
                "stability_score": 1.0
                - min(len(violations) / max(len(current_metrics), 1), 1.0),
                "family_z_summary": family_summary,
                "family_caps": self.family_caps,
                "sigma_quantile": float(self.sigma_quantile),
                "deadband": float(self.deadband),
                "max_caps": int(self.max_caps),
                "caps_applied": caps_applied,
                "caps_exceeded": caps_exceeded,
                "multiple_testing": self.multiple_testing,
            }

            family_quantiles, top_z_scores = self._compute_family_observability()
            if family_quantiles:
                metrics["family_z_quantiles"] = family_quantiles
            if top_z_scores:
                metrics["top_z_scores"] = top_z_scores

            if passed:
                message = (
                    "Spectral validation passed with "
                    f"{len(violations)} violations "
                    f"(caps_applied={caps_applied}, max_caps={self.max_caps})"
                )
            else:
                reason = (
                    "fatal spectral violation detected"
                    if fatal_violations
                    else "cap budget exceeded"
                )
                message = (
                    f"Spectral validation failed: {reason} "
                    f"(caps_applied={caps_applied}, max_caps={self.max_caps})"
                )

            # Runtime contracts (lightweight)
            mt = self.multiple_testing or {}
            try:
                alpha = float(mt.get("alpha", 0.05)) if isinstance(mt, dict) else 0.05
            except Exception:
                alpha = 0.05
            guard_assert(self.deadband >= 0.0, "spectral.deadband must be >= 0")
            guard_assert(
                0.0 < alpha <= 1.0, "spectral.multiple_testing.alpha out of range"
            )
            guard_assert(self.max_caps >= 0, "spectral.max_caps must be >= 0")

            return {
                "passed": passed,
                "action": action,
                "metrics": metrics,
                "violations": violations,
                "message": message,
                "policy": self._serialize_policy(),
                "final_z_scores": self.latest_z_scores.copy(),
                "module_family_map": dict(self.module_family_map),
            }

        except Exception as e:
            return {
                "passed": False,
                "action": "warn",
                "error": str(e),
                "metrics": {},
                "message": f"Spectral validation failed: {e}",
            }

    def finalize(self, model: Any) -> dict[str, Any]:
        """
        Finalize spectral guard and return comprehensive results.

        Args:
            model: The final model state

        Returns:
            Dictionary with spectral guard results
        """
        if not self.prepared:
            return {
                "passed": False,
                "metrics": {},
                "warnings": ["Spectral guard not properly prepared"],
                "errors": ["Preparation failed or not called"],
                "events": self.events,
            }

        # Final spectral analysis
        final_metrics = capture_baseline_sigmas(model, scope=self.scope)
        final_violations = self._detect_spectral_violations(
            model, final_metrics, phase="finalize"
        )
        final_z_summary = summarize_family_z_scores(
            self.latest_z_scores, self.module_family_map, self.family_caps
        )
        final_family_stats = compute_family_stats(final_metrics, self.module_family_map)

        family_quantiles, top_z_scores = self._compute_family_observability()

        # Determine overall status based on budgeted vs fatal violations
        fatal_violation_types = {"max_spectral_norm", "ill_conditioned"}
        budgeted_violations = [
            violation
            for violation in final_violations
            if violation.get("type") not in fatal_violation_types
        ]
        fatal_violations = [
            violation
            for violation in final_violations
            if violation.get("type") in fatal_violation_types
        ]

        caps_applied = len(budgeted_violations)
        caps_exceeded = caps_applied > int(self.max_caps)
        passed = not fatal_violations and not caps_exceeded

        # Compute comprehensive metrics
        metrics = {
            "modules_analyzed": len(final_metrics),
            "violations_detected": len(final_violations),
            "budgeted_violations": caps_applied,
            "fatal_violations": len(fatal_violations),
            "baseline_modules": len(self.baseline_metrics),
            "scope": self.scope,
            "max_spectral_norm_final": max(final_metrics.values())
            if final_metrics
            else 0.0,
            "mean_spectral_norm_final": np.mean(list(final_metrics.values()))
            if final_metrics
            else 0.0,
            "spectral_stability_score": 1.0
            - min(len(final_violations) / max(len(final_metrics), 1), 1.0),
            "target_sigma": self.target_sigma,
            "correction_applied": len(final_violations) > 0 and self.correction_enabled,
            "family_caps": self.family_caps,
            "family_z_summary": final_z_summary,
            "family_stats": final_family_stats,
            "sigma_quantile": float(self.sigma_quantile),
            "deadband": float(self.deadband),
            "max_caps": int(self.max_caps),
            "caps_applied": caps_applied,
            "caps_exceeded": caps_exceeded,
            "multiple_testing": self.multiple_testing,
            "family_z_quantiles": family_quantiles,
            "top_z_scores": top_z_scores,
        }

        # Categorize violations
        warnings = []
        errors = []

        for violation in final_violations:
            if violation["type"] in ["max_spectral_norm", "ill_conditioned"]:
                errors.append(violation["message"])
            else:
                warnings.append(violation["message"])

        result = {
            "passed": passed,
            "metrics": metrics,
            "warnings": warnings,
            "errors": errors,
            "violations": final_violations,
            "events": self.events,
            "baseline_metrics": self.baseline_metrics,
            "final_metrics": final_metrics,
            "final_z_scores": self.latest_z_scores,
            "module_family_map": dict(self.module_family_map),
            "policy": self._serialize_policy(),
        }

        # Env-gated tiny evidence dump for auditors
        try:
            payload = {
                "spectral": {
                    "sigma_quantile": float(self.sigma_quantile),
                    "deadband": float(self.deadband),
                    "max_caps": int(self.max_caps),
                    "multiple_testing": self.multiple_testing.get("method")
                    if isinstance(self.multiple_testing, dict)
                    else None,
                    "evaluated": True,
                }
            }
            maybe_dump_guard_evidence(".", payload)
        except Exception:
            pass

        return result


def compute_sigma_max(weight_matrix: Any) -> float:
    """
    Compute maximum singular value of a weight matrix.

    Args:
        weight_matrix: Weight matrix to analyze

    Returns:
        Maximum singular value
    """
    try:
        if isinstance(weight_matrix, torch.Tensor):
            # Handle different tensor types
            if weight_matrix.dtype in [torch.int8]:
                # Skip quantized weights
                return 1.0

            # Ensure float type for SVD
            W = weight_matrix.float()

            # Handle edge cases
            if W.numel() == 0 or W.shape[0] == 0 or W.shape[1] == 0:
                return 0.0

            # Compute singular values using deterministic backend when available
            try:
                singular_values = torch.linalg.svdvals(W)
            except RuntimeError:
                # Fallback for older backends without svdvals
                singular_values = torch.linalg.svd(W, full_matrices=False).S

            return singular_values[0].item() if singular_values.numel() > 0 else 0.0
        else:
            return 1.0  # Fallback for non-tensor inputs

    except Exception:
        return 1.0  # Fallback on any error


def auto_sigma_target(model: Any, percentile: float = 0.95, **kwargs: Any) -> float:
    """
    Automatically determine sigma target for a model.

    Args:
        model: Model to analyze
        percentile: Scale factor (target percentile of spectral norms)

    Returns:
        Target sigma value
    """
    if "kappa" in kwargs and percentile == 0.95:
        try:
            percentile = float(kwargs["kappa"])
        except (TypeError, ValueError):
            pass
    try:
        # Collect all spectral norms
        spectral_norms = []

        for _name, module in model.named_modules():
            if hasattr(module, "weight") and module.weight.ndim == 2:
                sigma = compute_sigma_max(module.weight)
                if sigma > 0:
                    spectral_norms.append(sigma)

        if spectral_norms:
            # Use kappa-percentile as target
            target = np.percentile(spectral_norms, percentile * 100)
            return float(target)
        else:
            return percentile  # Fallback to requested sigma quantile

    except Exception:
        return percentile  # Default fallback


def apply_weight_rescale(
    model: Any, scale_factor: float = 1.0, scope: str = "all"
) -> dict[str, Any]:
    """
    Apply weight rescaling to model parameters.

    Args:
        model: Model to rescale
        scale_factor: Scaling factor to apply
        scope: Which modules to rescale ('all', 'attn', 'ffn')

    Returns:
        Rescaling results
    """
    try:
        rescaled_modules = []
        failed_modules = []

        for name, module in model.named_modules():
            if not _should_process_module(name, module, scope):
                continue

            try:
                if hasattr(module, "weight") and module.weight.ndim == 2:
                    # Skip quantized weights
                    if hasattr(module.weight, "dtype") and module.weight.dtype in [
                        torch.int8,
                    ]:
                        continue

                    # Apply rescaling
                    with torch.no_grad():
                        module.weight.mul_(scale_factor)
                        if hasattr(module, "bias") and module.bias is not None:
                            module.bias.mul_(scale_factor)

                    rescaled_modules.append(name)

            except Exception as e:
                failed_modules.append((name, str(e)))

        return {
            "applied": len(rescaled_modules) > 0,
            "scale_factor": scale_factor,
            "rescaled_modules": rescaled_modules,
            "failed_modules": failed_modules,
            "message": f"Rescaled {len(rescaled_modules)} modules with factor {scale_factor}",
        }

    except Exception as e:
        return {
            "applied": False,
            "error": str(e),
            "message": f"Weight rescaling failed: {e}",
        }


def apply_relative_spectral_cap(
    model: Any,
    cap_ratio: float = 2.0,
    scope: str = "all",
    baseline_sigmas: dict[str, float] | None = None,
) -> dict[str, Any]:
    """
    Apply relative spectral capping to model weights.

    Args:
        model: Model to cap
        cap_ratio: Maximum allowed ratio relative to baseline
        scope: Which modules to cap ('all', 'attn', 'ffn')
        baseline_sigmas: Mapping of module name to pre-edit sigma values

    Returns:
        Capping results
    """
    try:
        if baseline_sigmas is None:
            baseline_sigmas = capture_baseline_sigmas(model, scope=scope)

        capped_modules = []
        failed_modules = []

        for name, module in model.named_modules():
            if not _should_process_module(name, module, scope):
                continue

            try:
                if hasattr(module, "weight") and module.weight.ndim == 2:
                    # Skip quantized weights
                    if hasattr(module.weight, "dtype") and module.weight.dtype in [
                        torch.int8,
                    ]:
                        continue

                    current_sigma = compute_sigma_max(module.weight)
                    baseline_sigma = baseline_sigmas.get(name, current_sigma)
                    max_allowed = baseline_sigma * cap_ratio

                    if current_sigma > max_allowed:
                        # Apply spectral capping using SVD
                        scale_factor = max_allowed / current_sigma

                        with torch.no_grad():
                            module.weight.mul_(scale_factor)

                        capped_modules.append(
                            {
                                "module": name,
                                "original_sigma": current_sigma,
                                "capped_sigma": max_allowed,
                                "scale_factor": scale_factor,
                            }
                        )

            except Exception as e:
                failed_modules.append((name, str(e)))

        return {
            "applied": len(capped_modules) > 0,
            "cap_ratio": cap_ratio,
            "capped_modules": capped_modules,
            "failed_modules": failed_modules,
            "message": f"Applied spectral capping to {len(capped_modules)} modules",
        }

    except Exception as e:
        return {
            "applied": False,
            "error": str(e),
            "message": f"Spectral capping failed: {e}",
        }


def apply_spectral_control(model: Any, policy: dict[str, Any]) -> dict[str, Any]:
    """
    Apply spectral control based on policy.

    Args:
        model: Model to control
        policy: Spectral control policy

    Returns:
        Control results
    """
    try:
        results: dict[str, Any] = {
            "rescaling_applied": False,
            "capping_applied": False,
            "modules_processed": 0,
            "corrections": [],
        }

        scope = policy.get("scope", "all")
        baseline_sigmas = policy.get("baseline_sigmas")

        # Apply relative spectral capping if needed
        cap_ratio = policy.get("cap_ratio", 2.0)
        cap_result = apply_relative_spectral_cap(
            model,
            cap_ratio=cap_ratio,
            scope=scope,
            baseline_sigmas=baseline_sigmas,
        )

        if cap_result["applied"]:
            results["capping_applied"] = True
            results["corrections"].extend(cap_result["capped_modules"])

        # Apply rescaling if target sigma is specified
        if "rescale_factor" in policy:
            rescale_result = apply_weight_rescale(
                model, scale_factor=policy["rescale_factor"], scope=scope
            )

            if rescale_result["applied"]:
                results["rescaling_applied"] = True
                results["modules_processed"] += len(rescale_result["rescaled_modules"])

        results["applied"] = results["rescaling_applied"] or results["capping_applied"]
        results["policy"] = policy
        results["message"] = (
            f"Spectral control applied: capping={results['capping_applied']}, rescaling={results['rescaling_applied']}"
        )

        return results

    except Exception as e:
        return {
            "applied": False,
            "error": str(e),
            "policy": policy,
            "message": f"Spectral control failed: {e}",
        }


def _summarize_sigmas(sigmas: dict[str, float]) -> dict[str, float]:
    """Compute summary statistics for a sigma dictionary."""
    if not sigmas:
        return {
            "max_spectral_norm": 0.0,
            "mean_spectral_norm": 0.0,
            "min_spectral_norm": 0.0,
        }

    values = np.array(list(sigmas.values()), dtype=float)
    return {
        "max_spectral_norm": float(values.max()),
        "mean_spectral_norm": float(values.mean()),
        "min_spectral_norm": float(values.min()),
    }


def compute_z_score_for_value(
    sigma: float,
    family_stats: dict[str, float],
    fallback_value: float,
    deadband: float,
) -> float:
    """Compute per-family z-score for a spectral norm with sensible fallbacks."""
    mean = float(family_stats.get("mean", 0.0) or 0.0)
    std = float(family_stats.get("std", 0.0) or 0.0)

    if std > 0:
        return float((sigma - mean) / std)

    # Fallback: scale relative change by deadband width
    denom = fallback_value if fallback_value > 0 else 1.0
    rel_change = (sigma / denom) - 1.0

    if abs(rel_change) <= deadband:
        return 0.0

    scale = deadband if deadband > 0 else 1.0
    return float(rel_change / scale)


def compute_z_scores(
    metrics: dict[str, float],
    baseline_family_stats: dict[str, dict[str, float]],
    module_family_map: dict[str, str],
    baseline_sigmas: dict[str, float],
    deadband: float,
) -> dict[str, float]:
    """Compute z-scores for all modules given baseline family stats."""
    z_scores: dict[str, float] = {}
    for name, sigma in metrics.items():
        family = module_family_map.get(name, "other")
        family_stats = baseline_family_stats.get(family, {})
        fallback_value = baseline_sigmas.get(name, family_stats.get("mean", sigma))
        z_scores[name] = compute_z_score_for_value(
            float(sigma),
            family_stats,
            float(fallback_value),
            deadband=deadband,
        )
    return z_scores


def summarize_family_z_scores(
    z_scores: dict[str, float],
    module_family_map: dict[str, str],
    family_caps: dict[str, dict[str, float]],
) -> dict[str, dict[str, float]]:
    """Summarize z-scores per family, including violation counts."""
    family_values: dict[str, list[float]] = defaultdict(list)
    for name, z in z_scores.items():
        family = module_family_map.get(name, "other")
        family_values[family].append(float(z))

    summary: dict[str, dict[str, float]] = {}
    for family, values in family_values.items():
        if not values:
            continue
        arr = np.array(values, dtype=float)
        cap = family_caps.get(family, {}).get("kappa")
        violations = 0
        if cap is not None:
            violations = int(np.sum(arr > float(cap)))
        summary[family] = {
            "max": float(arr.max()),
            "mean": float(arr.mean()),
            "count": len(values),
            "violations": violations,
        }
        if cap is not None:
            summary[family]["kappa"] = float(cap)
    return summary


def compute_family_stats(
    sigmas: dict[str, float], family_map: dict[str, str]
) -> dict[str, dict[str, float]]:
    """Compute per-family statistics (mean/std/min/max/count)."""
    buckets: dict[str, list[float]] = defaultdict(list)
    for name, sigma in sigmas.items():
        family = family_map.get(name, "other")
        buckets[family].append(float(sigma))

    stats: dict[str, dict[str, float]] = {}
    for family, values in buckets.items():
        if not values:
            continue
        arr = np.array(values, dtype=float)
        stats[family] = {
            "count": len(values),
            "mean": float(arr.mean()),
            "std": float(arr.std(ddof=0)),
            "min": float(arr.min()),
            "max": float(arr.max()),
        }
    return stats


def classify_model_families(
    model: Any, scope: str = "all", existing: dict[str, str] | None = None
) -> dict[str, str]:
    """Build or update a module→family map for the provided model."""
    family_map = dict(existing) if existing else {}
    for name, module in model.named_modules():
        if _should_process_module(name, module, scope):
            family_map[name] = classify_module_family(name, module)
    return family_map


def capture_baseline_sigmas(model: Any, scope: str = "all") -> dict[str, float]:
    """
    Capture baseline singular values for model layers.

    Args:
        model: Model to analyze
        scope: Which modules to analyze ('all', 'attn', 'ffn')

    Returns:
        Dictionary of layer name to max singular value
    """
    try:
        baseline_sigmas = {}

        for name, module in model.named_modules():
            if _should_process_module(name, module, scope):
                if hasattr(module, "weight") and module.weight.ndim == 2:
                    sigma = compute_sigma_max(module.weight)
                    baseline_sigmas[name] = sigma

        return baseline_sigmas

    except Exception:
        return {}


def scan_model_gains(model: Any, scope: str = "all") -> dict[str, Any]:
    """
    Scan model for gain values and spectral statistics.

    Args:
        model: Model to scan
        scope: Which modules to scan ('all', 'attn', 'ffn')

    Returns:
        Gain analysis results
    """
    try:
        results: dict[str, Any] = {
            "total_layers": 0,
            "scanned_modules": 0,
            "spectral_norms": [],
            "condition_numbers": [],
            "weight_statistics": {},
        }

        for name, module in model.named_modules():
            results["total_layers"] += 1

            if _should_process_module(name, module, scope):
                if hasattr(module, "weight") and module.weight.ndim == 2:
                    results["scanned_modules"] += 1

                    # Compute spectral norm
                    sigma_max = compute_sigma_max(module.weight)
                    results["spectral_norms"].append(sigma_max)

                    # Compute condition number if possible
                    try:
                        U, S, V = torch.svd(module.weight.float())
                        if len(S) > 1:
                            condition_num = (S[0] / S[-1]).item()
                            results["condition_numbers"].append(condition_num)
                    except Exception:
                        pass

                    # Basic weight statistics
                    try:
                        weight_stats = {
                            "mean": module.weight.mean().item(),
                            "std": module.weight.std().item(),
                            "min": module.weight.min().item(),
                            "max": module.weight.max().item(),
                        }
                        results["weight_statistics"][name] = weight_stats
                    except Exception:
                        pass

        # Compute summary statistics
        if results["spectral_norms"]:
            results["mean_spectral_norm"] = np.mean(results["spectral_norms"])
            results["max_spectral_norm"] = np.max(results["spectral_norms"])
            results["min_spectral_norm"] = np.min(results["spectral_norms"])

        if results["condition_numbers"]:
            results["mean_condition_number"] = np.mean(results["condition_numbers"])
            results["max_condition_number"] = np.max(results["condition_numbers"])

        results["message"] = (
            f"Scanned {results['scanned_modules']} modules out of {results['total_layers']} total layers"
        )

        return results

    except Exception as e:
        return {
            "total_layers": sum(1 for _ in model.named_modules()),
            "scanned_modules": 0,
            "error": str(e),
            "message": f"Model scanning failed: {e}",
        }


def _should_process_module(name: str, module: Any, scope: str) -> bool:
    """Helper function to determine if a module should be processed based on scope."""
    if not hasattr(module, "weight") or module.weight.ndim != 2:
        return False

    if scope == "all":
        return True
    elif scope == "attn":
        return any(
            keyword in name.lower()
            for keyword in ["attn", "attention", "self_attn", "c_attn", "c_proj"]
        )
    elif scope == "ffn":
        return any(
            keyword in name.lower()
            for keyword in ["mlp", "ffn", "feed_forward", "fc", "c_fc"]
        )
    elif scope == "ffn+proj":
        lname = name.lower()
        return any(
            keyword in lname
            for keyword in [
                "mlp",
                "ffn",
                "feed_forward",
                "fc",
                "c_fc",
                "c_proj",
                "projection",
            ]
        )

    return True


def classify_module_family(name: str, module: Any) -> str:
    """Classify module into a spectral family for policy purposes."""
    lname = name.lower()

    # MoE router/gating
    if any(
        tok in lname
        for tok in ("router", "routing", "gate", "gating", "dispatch", "switch")
    ):
        return "router"
    # MoE expert FFN
    if any(tok in lname for tok in ("experts", "expert", "moe", "mixture_of_experts")):
        return "expert_ffn"

    if "mlp" in lname or "ffn" in lname or "feed_forward" in lname:
        return "ffn"

    if (
        "attn" in lname
        or "attention" in lname
        or any(
            token in lname
            for token in ["q_proj", "k_proj", "v_proj", "o_proj", "c_attn"]
        )
    ):
        return "attn"

    if "embed" in lname or "wte" in lname or "embedding" in lname:
        return "embed"

    module_type = module.__class__.__name__.lower()
    if "embedding" in module_type:
        return "embed"
    if "conv1d" in module_type or "linear" in module_type:
        if "attn" in lname:
            return "attn"
        if "mlp" in lname or "ffn" in lname:
            return "ffn"

    return "other"


# Export the main components
__all__ = [
    "SpectralGuard",
    "SpectralPolicy",
    "compute_sigma_max",
    "auto_sigma_target",
    "apply_weight_rescale",
    "apply_relative_spectral_cap",
    "apply_spectral_control",
    "capture_baseline_sigmas",
    "scan_model_gains",
    "compute_family_stats",
    "summarize_family_z_scores",
    "classify_module_family",
]
