"""
Soft prompt engine for training prompt embeddings.

This module implements the core soft prompt training functionality
using SPSA (Simultaneous Perturbation Stochastic Approximation) and
other black-box optimization techniques.
"""

import numpy as np
from typing import Optional, Dict, Any, List, TYPE_CHECKING
from dataclasses import dataclass

from ..core.config import TrainingConfig
from ..utils.logger import get_logger
from ..utils.exceptions import TrainingError

if TYPE_CHECKING:
    from ..core.config import PromptConfig


@dataclass
class SoftPromptState:
    """State of a soft prompt during training."""
    prompt: np.ndarray
    learning_rate: float
    momentum: Optional[np.ndarray]
    iteration: int
    loss_history: List[float]


class SoftPromptEngine:
    """
    Engine for training soft prompts using black-box optimization.
    
    Uses SPSA (Simultaneous Perturbation Stochastic Approximation) to
    learn prompt embeddings that improve LLM performance on specific tasks.
    """
    
    def __init__(
        self,
        dimension: int,
        length: int,
        config: TrainingConfig,
        prompt_config: Optional['PromptConfig'] = None,
        existing_prompt: Optional[np.ndarray] = None
    ):
        """
        Initialize the soft prompt engine.
        
        Args:
            dimension: Embedding dimension
            length: Soft prompt length
            config: Training configuration
            prompt_config: Prompt configuration
            existing_prompt: Pre-existing soft prompt to continue training from
        """
        self.dimension = dimension
        self.length = length
        self.config = config
        self.prompt_config = prompt_config
        self.logger = get_logger(f"{__name__}.SoftPromptEngine")
        
        # Initialize the prompt - use existing if provided
        if existing_prompt is not None and existing_prompt.shape == (length, dimension):
            self.logger.info(f"Using existing soft prompt for continued training (shape: {existing_prompt.shape})")
            self.current_prompt = existing_prompt.copy()
        else:
            self.current_prompt = self._initialize_prompt()
        
        # Training state
        self.best_prompt = self.current_prompt.copy()
        self.best_loss = float('inf')
        self.momentum = None
        self.iteration = 0
        self.loss_history = []
        
        self.logger.info(f"SoftPromptEngine initialized: dim={dimension}, length={length}")
    
    def _initialize_prompt(self) -> np.ndarray:
        """Initialize soft prompt embeddings."""
        # CRITICAL FIX: Check if we already have a trained prompt (retraining case)
        if hasattr(self, 'current_prompt') and self.current_prompt is not None:
            self.logger.info(f"Using existing soft prompt for continued training (shape: {self.current_prompt.shape})")
            return self.current_prompt
        
        # Get initialization method from prompt config or default to random
        init_method = "random"
        if self.prompt_config and hasattr(self.prompt_config, 'soft_prompt_init'):
            init_method = self.prompt_config.soft_prompt_init
        
        if init_method == "random":
            # Random initialization with small values
            self.logger.info(f"Initializing new soft prompt with random values (dim={self.dimension}, length={self.length})")
            return np.random.normal(0, 0.01, (self.length, self.dimension))
        elif init_method == "zeros":
            self.logger.info(f"Initializing new soft prompt with zeros (dim={self.dimension}, length={self.length})")
            return np.zeros((self.length, self.dimension))
        elif init_method == "ones":
            self.logger.info(f"Initializing new soft prompt with small ones (dim={self.dimension}, length={self.length})")
            return np.ones((self.length, self.dimension)) * 0.01
        else:
            # Default to random
            self.logger.info(f"Initializing new soft prompt with random values (dim={self.dimension}, length={self.length})")
            return np.random.normal(0, 0.01, (self.length, self.dimension))
    
    async def update_from_feedback(
        self,
        input_text: str,
        expected_output: str,
        actual_output: Optional[str] = None,
        loss_fn: Optional[callable] = None,
        analysis_score: Optional[float] = None,
        suggestions: Optional[List[str]] = None
    ):
        """
        Update soft prompt based on feedback.
        
        Args:
            input_text: Original input
            expected_output: Desired output
            actual_output: What was actually produced
            loss_fn: Custom loss function
            analysis_score: Quality score from self-analysis (0-10)
            suggestions: List of improvement suggestions
        """
        if loss_fn is None:
            # Enhanced loss computation using analysis score if available
            if analysis_score is not None:
                # Convert 0-10 score to loss (higher score = lower loss)
                score_loss = max(0.0, (10.0 - analysis_score) / 10.0)
                text_loss = self._compute_text_similarity_loss(expected_output, actual_output or "")
                # Weighted combination of score-based and text-based loss
                loss = 0.7 * score_loss + 0.3 * text_loss
            else:
                loss = self._compute_text_similarity_loss(expected_output, actual_output or "")
        else:
            loss = loss_fn(expected_output, actual_output)
        
        # Enhanced SPSA update with adaptive learning
        await self._enhanced_spsa_update(loss, analysis_score, suggestions)
        
        self.logger.debug(f"Updated soft prompt from feedback, loss: {loss:.4f}, score: {analysis_score}")
    
    async def _enhanced_spsa_update(self, loss: float, analysis_score: Optional[float] = None, suggestions: Optional[List[str]] = None):
        """
        Enhanced SPSA update with improved parameter adaptation sensitivity.
        
        Args:
            loss: Current loss value
            analysis_score: Quality score from self-analysis
            suggestions: List of improvement suggestions
        """
        # Generate random perturbation with enhanced magnitude
        perturbation = np.random.choice([-1, 1], size=self.current_prompt.shape)
        perturbation = perturbation.astype(np.float32)
        
        # Enhanced perturbation scaling based on feedback quality
        base_epsilon = self.config.perturbation_scale
        
        # Increase perturbation for poor performance (low scores)
        if analysis_score is not None:
            # Scale perturbation inversely with score (worse performance = bigger changes)
            score_factor = max(0.5, (10.0 - analysis_score) / 5.0)  # Range: 0.5 to 2.0
            epsilon = base_epsilon * score_factor
        else:
            epsilon = base_epsilon
        
        # Additional scaling based on loss magnitude
        loss_factor = min(3.0, max(0.5, loss * 2.0))  # Range: 0.5 to 3.0
        epsilon *= loss_factor
        
        # Estimate gradient using finite differences with enhanced sensitivity
        gradient_estimate = (loss / epsilon) * perturbation
        
        # Enhanced gradient scaling for better adaptation
        gradient_magnitude = np.linalg.norm(gradient_estimate)
        
        # Apply gradient clipping with adaptive threshold
        adaptive_clip_threshold = self.config.gradient_clipping * (1.0 + loss)
        if gradient_magnitude > adaptive_clip_threshold:
            gradient_estimate = gradient_estimate * (adaptive_clip_threshold / gradient_magnitude)
        
        # Enhanced momentum with adaptive decay
        if self.momentum is not None:
            # Adaptive momentum based on performance
            momentum_factor = self.config.momentum
            if analysis_score is not None and analysis_score < 6.0:
                # Reduce momentum for poor performance to allow faster adaptation
                momentum_factor *= 0.7
            
            self.momentum = momentum_factor * self.momentum + (1 - momentum_factor) * gradient_estimate
            update = self.momentum
        else:
            update = gradient_estimate
        
        # Enhanced learning rate with adaptive scaling
        base_lr = self._get_adaptive_learning_rate(loss, analysis_score)
        
        # Apply the update with enhanced magnitude
        update_magnitude = np.linalg.norm(update)
        if update_magnitude > 0:
            # Ensure minimum update magnitude for meaningful adaptation
            min_update_magnitude = 1e-6
            if update_magnitude < min_update_magnitude:
                update = update * (min_update_magnitude / update_magnitude)
        
        # Apply the update
        self.current_prompt -= base_lr * update
        
        # Track loss history
        self.loss_history.append(loss)
        
        # Update best prompt if improved
        if loss < self.best_loss:
            self.best_loss = loss
            self.best_prompt = self.current_prompt.copy()
            self.logger.debug(f"New best prompt found at iteration {self.iteration}, loss: {loss:.4f}")
        
        # Log adaptation metrics
        adaptation_magnitude = np.linalg.norm(base_lr * update)
        self.logger.debug(f"Parameter adaptation magnitude: {adaptation_magnitude:.8f}")
        
        self.iteration += 1
    
    def _get_adaptive_learning_rate(self, loss: float, analysis_score: Optional[float] = None) -> float:
        """Get adaptive learning rate based on performance and feedback."""
        base_lr = self.config.learning_rate
        
        # Apply standard decay if enabled
        if self.config.adaptive_learning_rate:
            decay_factor = 1.0 / (1.0 + 0.01 * self.iteration)
            base_lr *= decay_factor
        
        # Adaptive scaling based on analysis score
        if analysis_score is not None:
            if analysis_score < 5.0:
                # Poor performance: increase learning rate for faster adaptation
                lr_multiplier = 2.0
            elif analysis_score < 7.0:
                # Moderate performance: standard learning rate
                lr_multiplier = 1.0
            else:
                # Good performance: reduce learning rate for fine-tuning
                lr_multiplier = 0.5
            
            base_lr *= lr_multiplier
        
        # Additional scaling based on loss magnitude
        loss_scaling = min(2.0, max(0.5, loss + 0.5))
        base_lr *= loss_scaling
        
        # Ensure minimum learning rate for meaningful updates
        min_lr = 1e-5
        return max(min_lr, base_lr)
    
    def _compute_text_similarity_loss(self, expected: str, actual: str) -> float:
        """
        Compute simple text similarity loss.
        
        Args:
            expected: Expected text
            actual: Actual text
            
        Returns:
            Loss value (0 = perfect match, higher = worse)
        """
        if not expected or not actual:
            return 1.0
        
        # Simple word-level Jaccard similarity
        expected_words = set(expected.lower().split())
        actual_words = set(actual.lower().split())
        
        if not expected_words and not actual_words:
            return 0.0
        
        intersection = len(expected_words.intersection(actual_words))
        union = len(expected_words.union(actual_words))
        
        jaccard_similarity = intersection / union if union > 0 else 0.0
        
        # Convert similarity to loss (1 - similarity)
        return 1.0 - jaccard_similarity
    
    def get_current_prompt(self) -> Optional[np.ndarray]:
        """Get current soft prompt."""
        return self.current_prompt
    
    def get_best_prompt(self) -> Optional[np.ndarray]:
        """Get best soft prompt found so far."""
        return self.best_prompt
    
    def set_prompt(self, prompt: np.ndarray):
        """Set the current prompt."""
        if prompt.shape != (self.length, self.dimension):
            raise TrainingError(f"Prompt shape {prompt.shape} doesn't match expected {(self.length, self.dimension)}")
        
        self.current_prompt = prompt.copy()
        self.logger.debug("Soft prompt updated manually")
    
    def get_training_stats(self) -> Dict[str, Any]:
        """Get training statistics."""
        # Use the last loss from history, or 0.5 as default
        current_loss = self.loss_history[-1] if self.loss_history else 0.5
        
        return {
            "iteration": self.iteration,
            "current_loss": current_loss,
            "best_loss": self.best_loss,
            "loss_history": self.loss_history[-10:],  # Last 10 losses
            "prompt_shape": self.current_prompt.shape,
            "learning_rate": self._get_adaptive_learning_rate(current_loss)
        }
    
    def reset(self):
        """Reset the soft prompt engine."""
        self.current_prompt = self._initialize_prompt()
        self.best_prompt = self.current_prompt.copy()
        self.best_loss = float('inf')
        self.iteration = 0
        self.loss_history = []
        
        if self.momentum is not None:
            self.momentum = np.zeros_like(self.current_prompt)
        
        self.logger.info("Soft prompt engine reset")
    
    def save_state(self) -> Dict[str, Any]:
        """Save current state for checkpointing."""
        state = {
            "current_prompt": self.current_prompt.tolist(),
            "best_prompt": self.best_prompt.tolist(),
            "best_loss": self.best_loss,
            "iteration": self.iteration,
            "loss_history": self.loss_history,
            "dimension": self.dimension,
            "length": self.length
        }
        
        if self.momentum is not None:
            state["momentum"] = self.momentum.tolist()
        
        return state
    
    def load_state(self, state: Dict[str, Any]):
        """Load state from checkpoint."""
        self.current_prompt = np.array(state["current_prompt"])
        self.best_prompt = np.array(state["best_prompt"])
        self.best_loss = state["best_loss"]
        self.iteration = state["iteration"]
        self.loss_history = state["loss_history"]
        
        if "momentum" in state and self.momentum is not None:
            self.momentum = np.array(state["momentum"])
        
        self.logger.info(f"Loaded soft prompt state at iteration {self.iteration}")
    
    def get_prompt_as_text(self, tokenizer=None) -> str:
        """
        Convert current prompt to text representation.
        
        Args:
            tokenizer: Optional tokenizer for decoding
            
        Returns:
            Text representation of the prompt
        """
        if tokenizer is None:
            # Return a basic representation
            return f"<soft_prompt_dim_{self.dimension}_len_{self.length}>"
        
        # If tokenizer provided, attempt to decode
        # This is model-specific and would need proper implementation
        try:
            # Placeholder for actual token decoding
            return "<decoded_soft_prompt>"
        except Exception as e:
            self.logger.warning(f"Failed to decode prompt to text: {str(e)}")
            return f"<soft_prompt_dim_{self.dimension}_len_{self.length}>" 