"""WICHER runtime hooks for subspace-projected Newton steering."""

from __future__ import annotations

import torch
from typing import Dict, Optional, TYPE_CHECKING
from wisent.core.cli.cli_logger import setup_logger, bind

if TYPE_CHECKING:
    from torch.nn import Module

_LOG = setup_logger(__name__)


class WicherRuntimeHooks:
    """
    Runtime hook system for WICHER Newton-based steering.

    Installs forward hooks on transformer layers that run Newton
    iterations in the SVD concept subspace using the lm_head weight.
    """

    def __init__(
        self,
        model: Module,
        concept_directions: Dict[int, torch.Tensor],
        concept_bases: Dict[int, torch.Tensor],
        component_variances: Dict[int, torch.Tensor],
        layer_variance: Dict[int, float],
        lm_head_weight: torch.Tensor,
        num_steps: int = 3,
        alpha: float = 5e-3,
        eta: float = 0.5,
        top_k: int = 256,
        beta: float = 0.0,
        alpha_decay: float = 1.0,
        armijo: bool = False,
        prob_mass_threshold: float = 0.0,
        base_strength: float = 1.0,
    ):
        self.model = model
        self.concept_directions = concept_directions
        self.concept_bases = concept_bases
        self.component_variances = component_variances
        self.layer_variance = layer_variance
        self.lm_head_weight = lm_head_weight.detach().float()
        self.num_steps = num_steps
        self.alpha = alpha
        self.eta = eta
        self.top_k = top_k
        self.beta = beta
        self.alpha_decay = alpha_decay
        self.armijo = armijo
        self.prob_mass_threshold = prob_mass_threshold
        self.base_strength = base_strength
        self._hooks = []

        self._variance_weights: Dict[int, float] = {}
        if self.layer_variance:
            total = sum(self.layer_variance.values())
            if total > 0:
                n_layers = len(self.layer_variance)
                for layer, var in self.layer_variance.items():
                    self._variance_weights[layer] = (var / total) * n_layers

        if hasattr(model, "model"):
            self._layers = model.model.layers
        elif hasattr(model, "transformer"):
            self._layers = model.transformer.h
        else:
            self._layers = model.layers

    def install(self) -> None:
        """Install forward hooks on steered layers."""
        self.remove()
        for layer_idx in self.concept_directions:
            if layer_idx < len(self._layers):
                hook = self._layers[layer_idx].register_forward_hook(
                    lambda mod, inp, out, li=layer_idx: self._wicher_hook(
                        mod, inp, out, li,
                    ),
                )
                self._hooks.append(hook)

    def remove(self) -> None:
        """Remove all installed hooks."""
        for hook in self._hooks:
            hook.remove()
        self._hooks = []

    def _wicher_hook(self, module, input, output, layer_idx: int):
        """Apply Newton-based steering to layer output."""
        hidden_states = output[0] if isinstance(output, tuple) else output
        rest = output[1:] if isinstance(output, tuple) else None

        h_new = self._apply_newton(hidden_states, layer_idx)

        if rest is not None:
            return (h_new,) + rest
        return h_new

    def _apply_newton(
        self, hidden_states: torch.Tensor, layer_idx: int,
    ) -> torch.Tensor:
        """Run Newton iterations in SVD concept subspace."""
        from wisent.core.steering_methods.methods.wicher.newton import (
            wicher_newton_step,
        )

        concept_dir = self.concept_directions[layer_idx].float()
        basis = self.concept_bases[layer_idx].float()
        comp_var = self.component_variances[layer_idx].float()

        original_shape = hidden_states.shape
        original_dtype = hidden_states.dtype

        effective_strength = self.base_strength
        if layer_idx in self._variance_weights:
            effective_strength *= self._variance_weights[layer_idx]

        if hidden_states.dim() == 1:
            h = hidden_states.unsqueeze(0).float()
        elif hidden_states.dim() == 2:
            h = hidden_states.float()
        else:
            b, s, hd = hidden_states.shape
            h = hidden_states.reshape(b * s, hd).float()

        lm_w = self.lm_head_weight.to(h.device)
        cd = concept_dir.to(h.device)
        basis_dev = basis.to(h.device)
        comp_var_dev = comp_var.to(h.device)

        h_new = wicher_newton_step(
            h, lm_w, cd * effective_strength,
            concept_basis=basis_dev,
            component_variances=comp_var_dev,
            num_steps=self.num_steps,
            alpha=self.alpha,
            eta=self.eta,
            top_k=self.top_k,
            beta=self.beta,
            alpha_decay=self.alpha_decay,
            armijo=self.armijo,
            prob_mass_threshold=self.prob_mass_threshold,
        )

        h_new = h_new.to(original_dtype)
        return h_new.reshape(original_shape)

    @property
    def steered_layers(self) -> list[int]:
        """Return list of layer indices being steered."""
        return sorted(self.concept_directions.keys())


def project_weights_wicher(
    model: Module,
    steering_obj,
    base_strength: float = 1.0,
    components: list[str] | None = None,
    verbose: bool = True,
) -> dict[str, int]:
    """Static projection: bake concept directions into model weights."""
    from wisent.core.weight_modification.methods.additive import (
        bake_steering_into_weights,
    )
    from wisent.core.activations.core.atoms import LayerActivations
    import torch.nn.functional as F

    if components is None:
        components = ["self_attn.o_proj", "mlp.down_proj"]

    direction_vectors = {}
    for layer_idx, cd in steering_obj.concept_directions.items():
        direction_vectors[layer_idx] = F.normalize(cd.float(), p=2, dim=-1)

    if verbose:
        print(f"\n{'='*60}")
        print("WICHER WEIGHT MODIFICATION (ADDITIVE — concept direction)")
        print(f"{'='*60}")
        print(f"Layers: {len(direction_vectors)}, Strength: {base_strength}")

    steering_vectors = LayerActivations(direction_vectors)
    stats = bake_steering_into_weights(
        model=model, steering_vectors=steering_vectors,
        alpha=base_strength, components=components, verbose=verbose,
    )
    stats["wicher_layers"] = len(steering_obj.concept_directions)
    stats["method"] = "additive"
    return stats
