"""
Weight modification optimizer using BaseOptimizer.

Optimizes weight modification parameters (strength, max_weight, min_weight, position)
to achieve target metrics like compliance rate or task accuracy.

Supports checkpointing for long-running optimizations:
- Save checkpoint after each trial (or every N trials)
- Resume from checkpoint if interrupted
- Save best model periodically
"""

from __future__ import annotations

import json
import os
from dataclasses import dataclass
from typing import Any, Callable, Literal

import optuna
import torch

from wisent.core.utils.services.optimization.core.atoms import BaseOptimizer, Direction, HPOConfig, HPORun
from wisent.core.utils.config_tools.constants import (
    WEIGHT_MIN_DISTANCE_FRACTION,
    OPTI_WEIGHTS_STRENGTH_RANGE,
    OPTI_WEIGHTS_MAX_WEIGHT_RANGE,
    OPTI_WEIGHTS_MIN_WEIGHT_RANGE,
    OPTI_WEIGHTS_POSITION_RANGE,
)

from wisent.core.utils.services.optimization.methods._opti_weights_checkpointing import WeightsCheckpointingMixin

__all__ = [
    "WeightsOptimizer",
    "WeightsOptimizerConfig",
]


@dataclass(slots=True)
class WeightsOptimizerConfig:
    """
    Configuration for weight modification optimization.

    attributes:
        strength_range:
            Range for strength parameter (min, max).
        max_weight_range:
            Range for max_weight parameter (min, max).
        min_weight_range:
            Range for min_weight parameter (min, max).
        position_range:
            Range for max_weight_position as ratio of layers (min, max).
        method:
            Weight modification method: "directional" or "bake".
        components:
            Number of components for weight modification.
        norm_preserve:
            Whether to preserve norms during directional projection.
        optimize_direction_index:
            Whether to optimize direction index as well.
        target_metric:
            Metric to optimize (e.g., "compliance_rate", "accuracy").
        target_value:
            Target value to achieve (for early stopping).
    """
    strength_range: tuple[float, float] = OPTI_WEIGHTS_STRENGTH_RANGE
    max_weight_range: tuple[float, float] = OPTI_WEIGHTS_MAX_WEIGHT_RANGE
    min_weight_range: tuple[float, float] = OPTI_WEIGHTS_MIN_WEIGHT_RANGE
    position_range: tuple[float, float] = OPTI_WEIGHTS_POSITION_RANGE
    method: Literal["directional", "bake"] = "directional"
    components: int = 1
    norm_preserve: bool = True
    optimize_direction_index: bool = False
    target_metric: str = "compliance_rate"
    target_value: float | None = None


class WeightsOptimizer(BaseOptimizer, WeightsCheckpointingMixin):
    """
    Optuna optimizer for weight modification parameters.

    Finds optimal parameters (strength, max_weight, min_weight, position) for
    applying steering vectors as permanent weight modifications.

    arguments:
        model:
            The model to optimize (will be modified in-place during trials).
        base_state_dict:
            Original model state dict for restoration between trials.
        steering_vectors:
            Dict mapping layer indices to steering vector tensors.
        evaluator:
            Callable that takes (model, tokenizer) and returns dict with metrics.
        tokenizer:
            Tokenizer for the model.
        config:
            WeightsOptimizerConfig with search space and method settings.
        num_layers:
            Number of layers in the model.

    example:
        >>> from wisent.optimization.methods.opti_weights import WeightsOptimizer, WeightsOptimizerConfig
        >>> from wisent.optimization.core.atoms import HPOConfig
        >>>
        >>> config = WeightsOptimizerConfig(
        ...     strength_range=(0.5, 2.0),
        ...     max_weight_range=(0.1, 1.0),
        ...     target_metric="compliance_rate",
        ... )
        >>>
        >>> optimizer = WeightsOptimizer(
        ...     model=model,
        ...     base_state_dict=base_state_dict,
        ...     steering_vectors=steering_vectors,
        ...     evaluator=my_evaluator,
        ...     tokenizer=tokenizer,
        ...     config=config,
        ...     num_layers=32,
        ... )
        >>>
        >>> result = optimizer.optimize(HPOConfig(n_trials=50, direction="maximize"))
        >>> print("Best params:", result.best_params)
        >>> print("Best score:", result.best_value)
    """

    name = "weights-optimizer"
    direction: Direction = "maximize"

    def __init__(
        self,
        model: Any,
        base_state_dict: dict[str, torch.Tensor],
        steering_vectors: dict[int, torch.Tensor],
        evaluator: Callable[[Any, Any], dict[str, float]],
        tokenizer: Any,
        config: WeightsOptimizerConfig,
        num_layers: int,
    ) -> None:
        self.model = model
        self.base_state_dict = base_state_dict
        self.steering_vectors = steering_vectors
        self.evaluator = evaluator
        self.tokenizer = tokenizer
        self.config = config
        self.num_layers = num_layers

        # Set direction based on target metric
        if config.target_metric in ["refusal_rate", "kl_divergence"]:
            self.direction = "minimize"
        else:
            self.direction = "maximize"

    def _objective(self, trial: optuna.Trial) -> float:
        """
        One trial: apply weight modification with suggested params and evaluate.

        arguments:
            trial: Optuna trial object.

        returns:
            float, value of the target metric.
        """
        # Suggest parameters
        params = {
            "strength": trial.suggest_float(
                "strength",
                self.config.strength_range[0],
                self.config.strength_range[1],
            ),
            "max_weight": trial.suggest_float(
                "max_weight",
                self.config.max_weight_range[0],
                self.config.max_weight_range[1],
            ),
            "min_weight": trial.suggest_float(
                "min_weight",
                self.config.min_weight_range[0],
                self.config.min_weight_range[1],
            ),
            "max_weight_position": trial.suggest_float(
                "max_weight_position",
                self.config.position_range[0],
                self.config.position_range[1],
            ),
        }

        # Optional direction index optimization
        if self.config.optimize_direction_index:
            params["direction_index"] = trial.suggest_float(
                "direction_index",
                0.0,
                self.num_layers - 1,
            )

        # Restore base model weights
        # Use strict=False because bake_steering may add bias parameters
        # that didn't exist in the original model
        self._restore_base_weights()

        # Apply weight modification
        self._apply_weight_modification(params)

        # Evaluate
        eval_result = self.evaluator(self.model, self.tokenizer)

        # Get target metric score
        score = eval_result.get(
            self.config.target_metric,
            eval_result.get("score", 0.0),
        )

        # Check for early stopping if target is set
        if self.config.target_value is not None:
            if self.direction == "maximize" and score >= self.config.target_value:
                trial.study.stop()
            elif self.direction == "minimize" and score <= self.config.target_value:
                trial.study.stop()

        return float(score)

    def _apply_weight_modification(self, params: dict[str, float]) -> None:
        """
        Apply weight modification with given parameters.

        arguments:
            params: Dict with strength, max_weight, min_weight, max_weight_position.
        """
        from wisent.core.weight_modification import (
            project_with_kernel,
            bake_steering_with_kernel,
        )

        # Convert position ratio to layer index
        max_weight_position = params["max_weight_position"] * (self.num_layers - 1)

        # Compute min_weight_distance from position
        min_weight_distance = WEIGHT_MIN_DISTANCE_FRACTION * (self.num_layers - 1)

        if self.config.method == "directional":
            project_with_kernel(
                self.model,
                self.steering_vectors,
                max_weight=params["max_weight"] * params["strength"],
                max_weight_position=max_weight_position,
                min_weight=params["min_weight"],
                min_weight_distance=min_weight_distance,
                components=self.config.components,
                norm_preserve=self.config.norm_preserve,
                verbose=False,
            )
        elif self.config.method in ("additive", "grom", "tecza", "tetno"):
            # Direct additive: add steering vector directly to weight matrices
            # This modifies weights directly, not biases, so it persists when saved
            # Used for additive and multi-direction methods (grom/tecza/tetno)
            self._apply_direct_additive(params)
        else:
            # Default fallback - use bake_steering_with_kernel
            # Note: This adds biases which may not load correctly for some architectures
            bake_steering_with_kernel(
                self.model,
                self.steering_vectors,
                max_alpha=params["max_weight"] * params["strength"],
                max_alpha_position=max_weight_position,
                min_alpha=params["min_weight"],
                components=self.config.components,
                verbose=False,
            )

