import numpy as np
from scipy.sparse import csr_matrix
from typing import List, Optional, Dict, Any, Tuple
from ..base_recommender import BaseRecommender

class BPRMFBase(BaseRecommender):
    """
    Bayesian Personalized Ranking with Matrix Factorization (BPRMF) for implicit feedback.
    
    Based on the paper:
    "BPR: Bayesian Personalized Ranking from Implicit Feedback" by Rendle et al.
    
    This model optimizes for correct ranking of items using a Bayesian approach
    with matrix factorization as the underlying prediction model.
    
    Parameters:
    -----------
    factors : int
        Number of latent factors
    learning_rate : float
        Learning rate for SGD
    regularization : float
        Regularization strength
    iterations : int
        Number of training iterations
    batch_size : int
        Size of mini-batches
    num_neg_samples : int
        Number of negative samples per positive sample
    seed : Optional[int]
        Random seed for reproducibility
    """
    def __init__(
        self,
        factors: int = 100,
        learning_rate: float = 0.05,
        regularization: float = 0.01,
        iterations: int = 100,
        batch_size: int = 1000,
        num_neg_samples: int = 1,
        seed: Optional[int] = None
    ):
        self.factors = factors
        self.learning_rate = learning_rate
        self.regularization = regularization
        self.iterations = iterations
        self.batch_size = batch_size
        self.num_neg_samples = num_neg_samples
        self.seed = seed
        
        self.user_map = {}
        self.item_map = {}
        self.reverse_user_map = {}
        self.reverse_item_map = {}
        self.user_factors = None
        self.item_factors = None
        self.item_bias = None
    
    def _create_mappings(self, user_ids: List[int], item_ids: List[int]) -> None:
        """Create mappings between original IDs and matrix indices"""
        self.user_map = {user_id: idx for idx, user_id in enumerate(user_ids)}
        self.item_map = {item_id: idx for idx, item_id in enumerate(item_ids)}
        self.reverse_user_map = {idx: user_id for user_id, idx in self.user_map.items()}
        self.reverse_item_map = {idx: item_id for item_id, idx in self.item_map.items()}
    
    def _init_params(self, n_users: int, n_items: int) -> None:
        """Initialize model parameters"""
        if self.seed is not None:
            np.random.seed(self.seed)
        
        # Initialize factors with small random values
        self.user_factors = np.random.normal(0, 0.1, (n_users, self.factors))
        self.item_factors = np.random.normal(0, 0.1, (n_items, self.factors))
        
        # Initialize item bias
        self.item_bias = np.zeros(n_items)
    
    def _predict(self, user_idx: int, item_idx: int) -> float:
        """Make prediction for a user-item pair"""
        # Compute prediction using dot product and item bias
        return np.dot(self.user_factors[user_idx], self.item_factors[item_idx]) + self.item_bias[item_idx]
    
    def _update_factors(self, user_idx: int, pos_item_idx: int, neg_item_idx: int) -> None:
        """Update factors using BPR optimization"""
        # Compute predictions
        pos_pred = self._predict(user_idx, pos_item_idx)
        neg_pred = self._predict(user_idx, neg_item_idx)
        
        # Compute BPR loss gradient (sigmoid of difference)
        diff = pos_pred - neg_pred
        sigmoid = 1.0 / (1.0 + np.exp(diff))
        
        # Compute gradients
        user_grad = sigmoid * (self.item_factors[pos_item_idx] - self.item_factors[neg_item_idx])
        pos_item_grad = sigmoid * self.user_factors[user_idx]
        neg_item_grad = -sigmoid * self.user_factors[user_idx]
        
        pos_bias_grad = sigmoid
        neg_bias_grad = -sigmoid
        
        # Update factors with regularization
        self.user_factors[user_idx] += self.learning_rate * (
            user_grad - self.regularization * self.user_factors[user_idx]
        )
        
        self.item_factors[pos_item_idx] += self.learning_rate * (
            pos_item_grad - self.regularization * self.item_factors[pos_item_idx]
        )
        
        self.item_factors[neg_item_idx] += self.learning_rate * (
            neg_item_grad - self.regularization * self.item_factors[neg_item_idx]
        )
        
        # Update biases
        self.item_bias[pos_item_idx] += self.learning_rate * (
            pos_bias_grad - self.regularization * self.item_bias[pos_item_idx]
        )
        
        self.item_bias[neg_item_idx] += self.learning_rate * (
            neg_bias_grad - self.regularization * self.item_bias[neg_item_idx]
        )
    
    def fit(self, user_ids: List[int], item_ids: List[int], ratings: List[float]) -> None:
        """
        Train the model on the given data.
        
        Parameters:
        -----------
        user_ids : List[int]
            List of user IDs
        item_ids : List[int]
            List of item IDs
        ratings : List[float]
            List of ratings (only used to determine positive interactions)
        """
        # Create mappings
        unique_user_ids = list(set(user_ids))
        unique_item_ids = list(set(item_ids))
        self._create_mappings(unique_user_ids, unique_item_ids)
        
        # Map IDs to indices
        user_indices = [self.user_map[user_id] for user_id in user_ids]
        item_indices = [self.item_map[item_id] for item_id in item_ids]
        
        # Create user-item matrix (binary interactions)
        n_users = len(self.user_map)
        n_items = len(self.item_map)
        self.user_item_matrix = csr_matrix(
            (np.ones(len(ratings)), (user_indices, item_indices)),
            shape=(n_users, n_items)
        )
        
        # Initialize parameters
        self._init_params(n_users, n_items)
        
        # Create positive interactions
        pos_interactions = list(zip(user_indices, item_indices))
        
        # Training loop
        for iteration in range(self.iterations):
            # Shuffle positive interactions
            if self.seed is not None:
                np.random.seed(self.seed + iteration)
            np.random.shuffle(pos_interactions)
            
            # Mini-batch training
            for i in range(0, len(pos_interactions), self.batch_size):
                batch = pos_interactions[i:i+self.batch_size]
                
                for user_idx, pos_item_idx in batch:
                    # Get items the user has interacted with
                    user_items = set(self.user_item_matrix[user_idx].indices)
                    
                    # Sample negative items
                    for _ in range(self.num_neg_samples):
                        # Sample a negative item (not interacted with by the user)
                        while True:
                            neg_item_idx = np.random.randint(0, n_items)
                            if neg_item_idx not in user_items:
                                break
                        
                        # Update factors
                        self._update_factors(user_idx, pos_item_idx, neg_item_idx)
    
    def recommend(self, user_id: int, top_n: int = 10, exclude_seen: bool = True) -> List[int]:
        """
        Generate top-N recommendations for a specific user.
        
        Parameters:
        -----------
        user_id : int
            ID of the user to generate recommendations for
        top_n : int
            Number of recommendations to generate
        exclude_seen : bool
            Whether to exclude items the user has already interacted with
            
        Returns:
        --------
        List[int] : List of recommended item IDs
        """
        if self.user_factors is None or self.item_factors is None:
            raise ValueError("Model has not been trained. Call fit() first.")
        
        # Map user_id to internal index
        if user_id not in self.user_map:
            raise ValueError(f"User ID {user_id} not found in training data")
            
        user_idx = self.user_map[user_id]
        
        # Calculate scores for all items
        user_vector = self.user_factors[user_idx]
        scores = np.dot(self.item_factors, user_vector) + self.item_bias
        
        # If requested, exclude items the user has already interacted with
        if exclude_seen:
            seen_items = self.user_item_matrix[user_idx].indices
            scores[seen_items] = float('-inf')
        
        # Get top-n item indices
        top_item_indices = np.argsort(-scores)[:top_n]
        
        # Map indices back to original item IDs
        top_items = [self.reverse_item_map[idx] for idx in top_item_indices]
        
        return top_items
    
    def save_model(self, filepath: str) -> None:
        """Save the model to a file"""
        if self.user_factors is None or self.item_factors is None:
            raise ValueError("Model has not been trained. Call fit() first.")
            
        # Save model data
        model_data = {
            'user_factors': self.user_factors,
            'item_factors': self.item_factors,
            'item_bias': self.item_bias,
            'user_map': self.user_map,
            'item_map': self.item_map,
            'reverse_user_map': self.reverse_user_map,
            'reverse_item_map': self.reverse_item_map,
            'params': {
                'factors': self.factors,
                'learning_rate': self.learning_rate,
                'regularization': self.regularization,
                'iterations': self.iterations,
                'batch_size': self.batch_size,
                'num_neg_samples': self.num_neg_samples,
                'seed': self.seed
            }
        }
        np.save(filepath, model_data, allow_pickle=True)
    
    @classmethod
    def load_model(cls, filepath: str) -> 'BPRMFBase':
        """Load a model from a file"""
        model_data = np.load(filepath, allow_pickle=True).item()
        
        # Create an instance with the saved parameters
        instance = cls(
            factors=model_data['params']['factors'],
            learning_rate=model_data['params']['learning_rate'],
            regularization=model_data['params']['regularization'],
            iterations=model_data['params']['iterations'],
            batch_size=model_data['params']['batch_size'],
            num_neg_samples=model_data['params']['num_neg_samples'],
            seed=model_data['params']['seed']
        )
        
        # Restore instance variables
        instance.user_factors = model_data['user_factors']
        instance.item_factors = model_data['item_factors']
        instance.item_bias = model_data['item_bias']
        instance.user_map = model_data['user_map']
        instance.item_map = model_data['item_map']
        instance.reverse_user_map = model_data['reverse_user_map']
        instance.reverse_item_map = model_data['reverse_item_map']
        
        return instance 