"""Guided modification execution."""
from __future__ import annotations
import torch
import torch.nn.functional as F
from typing import Dict, List, Optional, TYPE_CHECKING
from wisent.core.utils.cli.cli_logger import setup_logger, bind
if TYPE_CHECKING:
    from torch import Tensor
    from torch.nn import Module
    from wisent.core.primitives.models.wisent_model import WisentModel
    from wisent.core.primitives.contrastive_pairs.core.pair import ContrastivePair
from wisent.core.utils.config_tools.constants import SEPARATOR_WIDTH_WIDE
from wisent.core.weight_modification.methods.guided import (
    GuidedModificationConfig, GuidedModificationResult, CollateralDamageReport)
from wisent.core.weight_modification.methods._guided_diagnostics import compute_layer_diagnostics
from wisent.core.weight_modification.methods._guided_scoring import compute_fisher_weights
_LOG = setup_logger(__name__)

def select_surgical_layers(
    diagnostics: Dict[int, LayerDiagnostics],
    config: GuidedModificationConfig,
) -> List[int]:
    """
    Select layers for surgical ablation.
    
    Instead of ablating all layers, we select only the top-k layers
    with the strongest signal. This minimizes collateral damage.
    """
    if not diagnostics:
        return []
    
    # Filter by minimum linear score
    valid_layers = {
        l: d for l, d in diagnostics.items()
        if d.linear_score >= config.min_linear_score
    }
    
    if not valid_layers:
        # Fallback: use best layer even if below threshold
        best_layer = max(diagnostics.keys(), key=lambda l: diagnostics[l].linear_score)
        return [best_layer]
    
    # Sort by linear score (descending)
    sorted_layers = sorted(
        valid_layers.keys(),
        key=lambda l: valid_layers[l].linear_score,
        reverse=True
    )
    
    # Take top-k
    return sorted_layers[:config.surgical_top_k]


def run_guided_modification(
    model: "Module",
    pairs: List["ContrastivePair"],
    wisent_model: "WisentModel",
    config: Optional[GuidedModificationConfig] = None,
    components: Optional[List[str]] = None,
    *,
    guided_variance_threshold: float,
    guided_strong_signal: float,
    guided_moderate_signal: float,
    blend_default: float,
    architecture_module_limit: int,
) -> GuidedModificationResult:
    """Run guided weight modification using linearity diagnostics."""
    from wisent.core.weight_modification.directional import project_weights
    
    cfg = config or GuidedModificationConfig()
    log = bind(_LOG)
    
    if cfg.verbose:
        print("\n" + "=" * SEPARATOR_WIDTH_WIDE)
        print("GUIDED WEIGHT MODIFICATION (Linearity-Driven)")
        print("=" * SEPARATOR_WIDTH_WIDE)
        print(f"Mode: {cfg.mode.value}")
        print(f"Use Fisher weights: {cfg.use_fisher_weights}")
        print(f"Validate collateral: {cfg.validate_collateral}")
        print("=" * SEPARATOR_WIDTH_WIDE + "\n")
    
    # Step 1: Compute layer diagnostics
    diagnostics = compute_layer_diagnostics(
        pairs=pairs,
        model=wisent_model,
        layers=None,  # All layers
        extraction_strategy=cfg.extraction_strategy,
        verbose=cfg.verbose,
        blend_default=blend_default,
        architecture_module_limit=architecture_module_limit,
    )
    
    if not diagnostics:
        raise ValueError("No valid layer diagnostics computed. Check your pairs.")
    
    # Step 2: Select layers based on mode
    if cfg.mode == AblationMode.SURGICAL:
        selected_layers = select_surgical_layers(diagnostics, cfg)
        mode_used = AblationMode.SURGICAL
    elif cfg.mode == AblationMode.FULL:
        selected_layers = [
            l for l, d in diagnostics.items()
            if d.linear_score >= cfg.min_linear_score
        ]
        mode_used = AblationMode.FULL
    else:  # ADAPTIVE — conservative: always use SURGICAL
        selected_layers = select_surgical_layers(diagnostics, cfg)
        mode_used = AblationMode.SURGICAL
    
    if not selected_layers:
        # Fallback: use best layer
        best_layer = max(diagnostics.keys(), key=lambda l: diagnostics[l].linear_score)
        selected_layers = [best_layer]
        if cfg.verbose:
            print(f"Warning: No layers above threshold, using best layer {best_layer}")
    
    if cfg.verbose:
        print(f"\nSelected layers ({mode_used.value}): {selected_layers}")
    
    # Step 3: Compute layer weights
    if cfg.use_fisher_weights:
        all_weights = compute_fisher_weights(diagnostics, cfg, blend_default=blend_default)
        layer_weights = {l: all_weights[l] for l in selected_layers if l in all_weights}
    else:
        # Use recommended weights from diagnostics
        layer_weights = {
            l: diagnostics[l].recommended_weight * cfg.base_strength
            for l in selected_layers
        }
    
    if cfg.verbose:
        print("\nLayer weights:")
        for l in sorted(layer_weights.keys()):
            print(f"  Layer {l:3d}: weight={layer_weights[l]:.3f}")
    
    # Step 4: Compute steering vectors for selected layers
    steering_vectors = _compute_steering_vectors(
        pairs=pairs,
        model=wisent_model,
        layers=selected_layers,
        extraction_strategy=cfg.extraction_strategy,
        normalize=cfg.normalize_vectors,
        architecture_module_limit=architecture_module_limit,
    )
    
    if cfg.verbose:
        print(f"\nComputed steering vectors for {len(steering_vectors)} layers")
    
    # Step 5: Apply weight modification
    if cfg.verbose:
        print("\nApplying weight modification...")
    
    stats = project_weights(
        model=model,
        steering_vectors=steering_vectors,
        harmless_vectors=None,
        components=components,
        layer_weights=layer_weights,
        strength=1.0,  # Weights already incorporate strength
        normalize_vectors=cfg.normalize_vectors,
        norm_preserve=True,
        use_biprojection=False,
        verbose=cfg.verbose,
    )
    
    # Step 6: Validate collateral damage (optional)
    collateral_report = None
    if cfg.validate_collateral:
        if cfg.verbose:
            print("\nValidating collateral damage...")
        # Note: Full implementation would require access to original model
        # and running diagnostics on unrelated benchmarks
        # For now, we just report that validation was requested
        collateral_report = CollateralDamageReport(
            benchmarks_tested=[],
            before_scores={},
            after_scores={},
            deltas={},
            max_degradation=0.0,
            mean_degradation=0.0,
            passed=True,
            details={"note": "Full validation requires benchmark evaluation"}
        )
    
    # Generate recommendation
    best_layer = max(selected_layers, key=lambda l: diagnostics[l].linear_score)
    best_score = diagnostics[best_layer].linear_score
    
    if best_score >= guided_strong_signal:
        recommendation = (
            f"Strong linear signal detected. Modified {len(selected_layers)} layers "
            f"with {mode_used.value} mode. Best layer: {best_layer} (score={best_score:.3f})"
        )
    elif best_score >= guided_moderate_signal:
        recommendation = (
            f"Moderate linear signal. Modified {len(selected_layers)} layers. "
            f"Consider verifying results with benchmark evaluation."
        )
    else:
        recommendation = (
            f"Weak linear signal (best={best_score:.3f}). Results may be suboptimal. "
            f"Consider using multi-directional methods (TECZA) or GROM."
        )
    
    return GuidedModificationResult(
        layers_modified=stats["layers_modified"],
        total_parameters_modified=stats["total_parameters_modified"],
        layer_diagnostics=diagnostics,
        layer_weights=layer_weights,
        mode_used=mode_used,
        steering_vectors=steering_vectors,
        collateral_report=collateral_report,
        recommendation=recommendation,
    )


def _compute_steering_vectors(
    pairs: List["ContrastivePair"],
    model: "WisentModel",
    layers: List[int],
    extraction_strategy: str,
    normalize: bool,
    *,
    architecture_module_limit: int,
) -> Dict[int, Tensor]:
    """Compute steering vectors (difference of means) for specified layers."""
    from wisent.core.primitives.model_interface.core.activations.activations_collector import ActivationCollector
    from wisent.core.primitives.model_interface.core.activations import ExtractionStrategy
    
    collector = ActivationCollector(model, architecture_module_limit=architecture_module_limit)
    strategy = ExtractionStrategy(extraction_strategy)
    
    pos_activations: Dict[int, List[Tensor]] = {l: [] for l in layers}
    neg_activations: Dict[int, List[Tensor]] = {l: [] for l in layers}
    
    for pair in pairs:
        try:
            pair_with_acts = collector.collect(
                pair,
                strategy=strategy,
                layers=[str(l) for l in layers],
                normalize=False,
            )
            
            pos_la = pair_with_acts.positive_response.layers_activations
            neg_la = pair_with_acts.negative_response.layers_activations
            
            if pos_la and neg_la:
                for layer in layers:
                    pos_t = pos_la.get(str(layer))
                    neg_t = neg_la.get(str(layer))
                    if pos_t is not None and neg_t is not None:
                        pos_activations[layer].append(pos_t.flatten())
                        neg_activations[layer].append(neg_t.flatten())
        except Exception:
            continue
    
    steering_vectors = {}
    for layer in layers:
        pos_list = pos_activations[layer]
        neg_list = neg_activations[layer]
        
        if len(pos_list) < 2 or len(neg_list) < 2:
            continue
        
        pos_mean = torch.stack(pos_list).mean(dim=0)
        neg_mean = torch.stack(neg_list).mean(dim=0)
        
        direction = pos_mean - neg_mean
        
        if normalize:
            direction = F.normalize(direction, p=2, dim=0)
        
        steering_vectors[layer] = direction
    
    return steering_vectors
