import os
import joblib
import pandas as pd
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import load_model
from functools import lru_cache

# Custom FMLayer definition for Factorization Machine component in DeepFM
class FMLayer(tf.keras.layers.Layer):
    def call(self, inputs):
        square_of_sum = tf.square(tf.reduce_sum(inputs, axis=1))
        sum_of_square = tf.reduce_sum(tf.square(inputs), axis=1)
        return 0.5 * (square_of_sum - sum_of_square)

class RecommenderManager:
    def __init__(self, models_dir="Models", data_path="Processed/train_full.parquet"):
        self.models_dir = models_dir
        self.data_path = data_path
        
        # Placeholders
        self.model = None
        self.user_encoder = None
        self.movie_encoder = None
        self.scaler = None
        self.watch_history_df = None
        self.movie_embeddings = None
        self.loaded = False
        
    def load_resources(self):
        """
        Loads all encoders, scaler, watch history data, and TensorFlow model.
        """
        if self.loaded:
            return
            
        print("Loading recommender encoders and scaler...")
        try:
            self.user_encoder = joblib.load(os.path.join(self.models_dir, "user_encoder_full.pkl"))
            self.movie_encoder = joblib.load(os.path.join(self.models_dir, "movie_encoder_full.pkl"))
            self.scaler = joblib.load(os.path.join(self.models_dir, "scaler_full.pkl"))
        except Exception as e:
            print(f"Error loading encoders/scaler: {e}")
            
        print("Loading training watch history...")
        try:
            if os.path.exists(self.data_path):
                self.watch_history_df = pd.read_parquet(self.data_path, columns=["user_id", "movie_id"])
            else:
                print(f"Watch history not found at {self.data_path}, initializing empty")
                self.watch_history_df = pd.DataFrame(columns=["user_id", "movie_id"])
        except Exception as e:
            print(f"Error loading watch history: {e}")
            self.watch_history_df = pd.DataFrame(columns=["user_id", "movie_id"])
            
        print("Loading Keras DeepFM model (deepfm_full.keras)...")
        try:
            model_path = os.path.join(self.models_dir, "deepfm_full.keras")
            if not os.path.exists(model_path):
                model_path = os.path.join(self.models_dir, "deepfm_model.keras")
                
            self.model = load_model(model_path, custom_objects={"FMLayer": FMLayer})
            print("Model loaded successfully")
            
            # Extract movie embeddings
            try:
                self.movie_embeddings = self.model.get_layer("embedding_1").get_weights()[0]
                print(f"Extracted movie embeddings with shape {self.movie_embeddings.shape}")
            except Exception as emb_err:
                print(f"Could not extract embeddings from model layer: {emb_err}")
        except Exception as e:
            print(f"Error loading model: {e}")
            
        self.loaded = True

    def _mmr_rerank(self, candidates_df, top_n, lambda_param=0.6, seed=None, relevance_col="pred_rating"):
        """
        Maximal Marginal Relevance re-ranking to balance relevance and diversity.
        Uses movie embeddings for similarity when available, falls back to year-based.
        """
        if len(candidates_df) <= top_n:
            return candidates_df.copy()

        pool = candidates_df.copy().reset_index(drop=True)
        n = len(pool)

        if self.movie_embeddings is not None and "movie_idx" in pool.columns:
            idx_min, idx_max = int(pool["movie_idx"].min()), int(pool["movie_idx"].max())
            if idx_min >= 0 and idx_max < len(self.movie_embeddings):
                emb = self.movie_embeddings[pool["movie_idx"].values.astype(int)]
                norms = np.linalg.norm(emb, axis=1, keepdims=True)
                norms[norms == 0] = 1e-9
                sim_matrix = np.dot(emb, emb.T) / (norms * norms.T + 1e-9)
                sim_matrix = np.clip(sim_matrix, 0.0, 1.0)
            else:
                sim_matrix = self._fallback_sim_matrix(pool)
        else:
            sim_matrix = self._fallback_sim_matrix(pool)

        relevance = pool[relevance_col].values
        if seed is not None:
            rng = np.random.RandomState(seed)
            noise = rng.normal(0, 0.25 * (relevance.max() - relevance.min() + 1e-6), size=n)
            relevance = relevance + noise

        selected_mask = np.zeros(n, dtype=bool)
        remaining = np.ones(n, dtype=bool)

        for _ in range(min(top_n, n)):
            mmr_scores = np.full(n, -np.inf)
            for i in np.where(remaining)[0]:
                rel = relevance[i]
                if selected_mask.any():
                    div = float(sim_matrix[i, selected_mask].max())
                else:
                    div = 0.0
                mmr_scores[i] = lambda_param * rel - (1.0 - lambda_param) * div

            best = int(np.argmax(mmr_scores))
            selected_mask[best] = True
            remaining[best] = False

        return pool[selected_mask].copy().reset_index(drop=True)

    def _fallback_sim_matrix(self, pool):
        """Build similarity matrix from year when embeddings unavailable."""
        years = pool["year"].values.astype(float)
        year_diff = np.abs(years[:, None] - years[None, :])
        sim = 1.0 - (year_diff / 50.0)
        return np.clip(sim, 0.0, 1.0)

    @lru_cache(maxsize=128)
    def recommend_for_user(self, user_id, top_n=10, movie_metadata_file="movie_metadata.parquet"):
        """
        Generates predictions and recommendations for a given User ID.
        Uses caching to avoid redundant model runs.
        """
        if not self.loaded:
            self.load_resources()
            
        movie_df = pd.read_parquet(movie_metadata_file)
        
        is_cold_start = True
        if self.user_encoder is not None and user_id in self.user_encoder.classes_:
            is_cold_start = False
            
        if is_cold_start:
            movie_df["popularity_score"] = movie_df["avg_rating"] * np.log1p(movie_df["rating_count"])
            pool_size = min(300, len(movie_df))
            pool = movie_df.sort_values(by="popularity_score", ascending=False).head(pool_size).copy()
            pool["pred_rating"] = pool["popularity_score"] / (pool["popularity_score"].max() + 1e-9)
            if self.movie_encoder is not None:
                valid = pool["movie_id"].isin(self.movie_encoder.classes_)
                pool = pool[valid].copy()
                if not pool.empty:
                    pool["movie_idx"] = self.movie_encoder.transform(pool["movie_id"])
            diverse = self._mmr_rerank(pool, top_n, lambda_param=0.25, seed=int(user_id))
            diverse["rank"] = range(1, len(diverse) + 1)
            diverse["match_score"] = 0.85
            diverse["pred_rating"] = diverse["avg_rating"]
            diverse["reason"] = diverse.apply(
                lambda x: "Popular choice recommended for new profiles" if x["rating_count"] > 5000 
                else "Highly rated classic choice", axis=1
            )
            return diverse, True
            
        user_idx = self.user_encoder.transform([user_id])[0]
        
        watched_movies = set()
        if self.watch_history_df is not None and not self.watch_history_df.empty:
            watched_movies = set(self.watch_history_df.loc[self.watch_history_df["user_id"] == user_id, "movie_id"])
            
        candidates = movie_df[~movie_df["movie_id"].isin(watched_movies)].copy()
        if candidates.empty:
            candidates = movie_df.copy()
            
        if self.movie_encoder is not None:
            valid_movies_mask = candidates["movie_id"].isin(self.movie_encoder.classes_)
            candidates = candidates[valid_movies_mask].copy()
            if candidates.empty:
                candidates = movie_df.copy()
                
            candidates["movie_idx"] = self.movie_encoder.transform(candidates["movie_id"])
            
        candidates["user_idx"] = user_idx

        # Drop rows with NaN features (e.g., missing year)
        candidates = candidates.dropna(subset=["year", "avg_rating", "rating_count"]).copy()
        if candidates.empty:
            candidates = movie_df.copy()
            candidates["movie_idx"] = self.movie_encoder.transform(candidates["movie_id"])
            candidates["user_idx"] = user_idx
        
        scale_df = pd.DataFrame({
            "year_x": candidates["year"],
            "movie_avg_rating": candidates["avg_rating"],
            "movie_rating_count": candidates["rating_count"]
        })
        
        scaled_features = self.scaler.transform(scale_df)
        
        if self.model is not None:
            user_input = np.array(candidates["user_idx"]).reshape(-1, 1)
            movie_input = np.array(candidates["movie_idx"]).reshape(-1, 1)
            
            predictions = self.model.predict(
                [user_input, movie_input, scaled_features],
                batch_size=4096,
                verbose=0
            ).flatten()
            
            candidates["pred_rating"] = predictions
        else:
            candidates["pred_rating"] = candidates["avg_rating"]
            
        # Drop any NaN predictions
        candidates = candidates.dropna(subset=["pred_rating"]).copy()
        if candidates.empty:
            candidates = movie_df.copy()
            candidates["pred_rating"] = candidates["avg_rating"] / 5.0

        # Compute user-specific surprise: pred_rating vs expected (avg_rating/5)
        # This is user-specific because pred_rating varies per user
        candidates["global_expect"] = candidates["avg_rating"] / 5.0
        candidates["user_surprise"] = candidates["pred_rating"] - candidates["global_expect"]

        # Quality gate: consider movies above ~25th percentile by avg_rating
        min_quality = candidates["avg_rating"].quantile(0.25)
        quality_pool = candidates[candidates["avg_rating"] >= min_quality].copy()
        if len(quality_pool) < top_n * 20:
            quality_pool = candidates.copy()

        # Use raw user_surprise as primary ranking signal (breaks 96% prediction correlation)
        pool = quality_pool.nlargest(600, "user_surprise").copy()
        diverse = self._mmr_rerank(pool, top_n, lambda_param=0.15, seed=int(user_id), relevance_col="pred_rating")
        diverse["rank"] = range(1, len(diverse) + 1)
        
        if self.model is not None:
            raw_probs = diverse["pred_rating"].copy()
            diverse["match_score"] = np.clip(raw_probs, 0.0, 1.0)
            diverse["pred_rating"] = 1.0 + 4.0 * raw_probs
        else:
            diverse["match_score"] = np.clip((diverse["pred_rating"] - 1.0) / 4.0, 0.0, 1.0)
        
        def get_explanation(row):
            score = row["match_score"]
            avg_r = row["avg_rating"]
            count = row["rating_count"]
            if score >= 0.85:
                return "Highly matched with your implicit interest fingerprint."
            elif avg_r >= 4.3:
                return "Highly rated globally matching your cinematic quality standard."
            elif count >= 3000:
                return "Popular blockbuster choice aligning with your content depth."
            return "Recommended based on similar user preferences."
            
        diverse["reason"] = diverse.apply(get_explanation, axis=1)
        
        return diverse, False

    def get_recommendation_analytics(self, recommendations_df):
        """
        Calculates confidence, diversity, and personalization metrics.
        """
        if recommendations_df.empty:
            return {
                "avg_confidence": 0.0,
                "diversity_score": 0.0,
                "personalization_score": 0.0
            }
            
        avg_confidence = float(recommendations_df["match_score"].mean() * 100)
        
        years = recommendations_df["year"].dropna()
        if len(years) > 1:
            year_std = float(years.std())
            diversity = min(100.0, (year_std / 20.0) * 100.0)
        else:
            diversity = 50.0
            
        if "pred_rating" in recommendations_df.columns:
            rating_diffs = np.abs(recommendations_df["pred_rating"] - recommendations_df["avg_rating"])
            personalization = min(100.0, float(rating_diffs.mean() * 60.0 + 40.0))
        else:
            personalization = 50.0
            
        return {
            "avg_confidence": avg_confidence,
            "diversity_score": diversity,
            "personalization_score": personalization
        }

    @lru_cache(maxsize=128)
    def get_similar_movies(self, movie_id, top_n=10, movie_metadata_file="movie_metadata.parquet"):
        """
        Retrieves top N similar movies using cosine similarity on DeepFM embeddings.
        """
        if not self.loaded:
            self.load_resources()
            
        movie_df = pd.read_parquet(movie_metadata_file)
        
        is_known = False
        if (self.movie_encoder is not None 
                and self.movie_embeddings is not None 
                and movie_id in self.movie_encoder.classes_):
            is_known = True
            
        if not is_known:
            print(f"Fallback Similarity lookup for movie ID: {movie_id}")
            match_row = movie_df[movie_df["movie_id"] == movie_id]
            if not match_row.empty:
                ref_year = match_row.iloc[0]["year"]
                decade_start = (ref_year // 10) * 10
                candidates = movie_df[
                    (movie_df["year"] >= decade_start) & 
                    (movie_df["year"] < decade_start + 10) & 
                    (movie_df["movie_id"] != movie_id)
                ].copy()
                candidates["similarity_score"] = 0.70 + 0.15 * (candidates["avg_rating"] / 5.0)
                recs = candidates.sort_values(by="avg_rating", ascending=False).head(top_n)
                return recs
            else:
                candidates = movie_df.copy()
                candidates["similarity_score"] = 0.60
                return candidates.sort_values(by="rating_count", ascending=False).head(top_n)
                
        movie_idx = self.movie_encoder.transform([movie_id])[0]
        movie_vector = self.movie_embeddings[movie_idx]
        
        dot_products = np.dot(self.movie_embeddings, movie_vector)
        norms = np.linalg.norm(self.movie_embeddings, axis=1)
        ref_norm = np.linalg.norm(movie_vector)
        
        if ref_norm == 0:
            similarities = np.zeros(len(self.movie_embeddings))
        else:
            similarities = dot_products / (norms * ref_norm + 1e-9)
            
        sim_df = pd.DataFrame({
            "movie_idx": range(len(self.movie_embeddings)),
            "similarity_score": similarities
        })
        
        sim_df["movie_id"] = self.movie_encoder.inverse_transform(sim_df["movie_idx"])
        sim_df = sim_df[sim_df["movie_id"] != movie_id]
        
        merged = pd.merge(sim_df, movie_df, on="movie_id")
        top_recs = merged.sort_values(by="similarity_score", ascending=False).head(top_n).copy()
        top_recs["similarity_score"] = np.clip(top_recs["similarity_score"], 0.0, 1.0)
        
        return top_recs

    @lru_cache(maxsize=64)
    def get_movie_detailed_profile(self, movie_id, ratings_file="Processed/ratings_filtered.parquet"):
        """
        Loads all raw ratings for a movie using parquet filters,
        calculates percentiles, rating distributions, timelines, and generates semantic summaries.
        """
        # Load catalog for percentiles
        movie_df = pd.read_parquet("movie_metadata.parquet")
        match_row = movie_df[movie_df["movie_id"] == movie_id]
        
        if match_row.empty:
            return None
            
        ref_movie = match_row.iloc[0]
        title = str(ref_movie["title"])
        year = int(ref_movie["year"])
        avg_rating = float(ref_movie["avg_rating"])
        rating_count = int(ref_movie["rating_count"])
        
        # Calculate global rating percentile
        percentile = float((movie_df["avg_rating"] < avg_rating).mean() * 100)
        
        # Query raw ratings for detailed metrics
        ratings_df = pd.DataFrame(columns=["rating", "date"])
        rating_std = 1.0
        
        try:
            if os.path.exists(ratings_file):
                # Using parquet filter for efficient query
                ratings_df = pd.read_parquet(ratings_file, columns=["rating", "date"], filters=[("movie_id", "==", movie_id)])
                if not ratings_df.empty:
                    rating_std = float(ratings_df["rating"].std())
        except Exception as e:
            print(f"Error querying ratings timeline for movie {movie_id}: {e}")
            
        # Generate semantic AI-style movie summary
        decade_str = f"released in {year}" if year > 0 else "from our catalog"
        if 1990 <= year < 2000:
            decade_str = "released during the iconic 1990s era of filmmaking"
        elif 1980 <= year < 1990:
            decade_str = "released during the nostalgic 1980s cinematic wave"
        elif 2000 <= year < 2010:
            decade_str = "released in the early 2000s transition period of cinema"
            
        if avg_rating >= 4.5:
            quality = "widely regarded as a cinematic masterpiece of exceptional caliber"
        elif avg_rating >= 4.0:
            quality = "recognized as a highly polished, critically acclaimed production"
        elif avg_rating >= 3.5:
            quality = "offering a solid, engaging narrative with general audience appeal"
        else:
            quality = "a lighter, casual watch with mixed viewer receptions"
            
        if rating_std > 1.2:
            consensus = "a polarizing work that continues to spark passionate debates among film critics and enthusiasts alike"
        elif rating_std < 0.8:
            consensus = "a universally praised entry with a highly cohesive consensus and solid backing across all user demographics"
        else:
            consensus = "a stable favorite that maintains consistent appeal and strong, standard ratings"
            
        if rating_count >= 10000:
            popularity = "As a massive blockbuster sensation, it has amassed an enormous viewer base and high cultural footprint."
        elif rating_count >= 2000:
            popularity = "With a healthy, popular following, it remains a common recommendation in modern streaming circles."
        else:
            popularity = "Operating as a hidden gem, it appeals greatly to niche enthusiasts looking for unique narrative directions."
            
        summary_text = f"<strong>\"{title}\"</strong> is {quality}, {decade_str}. It is valued as {consensus}. {popularity}"
        
        # Process popularity timeline: group/sort by date to calculate cumulative votes
        timeline_data = pd.DataFrame(columns=["date", "cumulative_votes"])
        if not ratings_df.empty:
            ratings_df["date"] = pd.to_datetime(ratings_df["date"])
            sorted_ratings = ratings_df.sort_values(by="date").copy()
            sorted_ratings["cumulative_votes"] = range(1, len(sorted_ratings) + 1)
            
            # Sub-sample to keep Plotly line charts light
            if len(sorted_ratings) > 500:
                indices = np.linspace(0, len(sorted_ratings) - 1, 500, dtype=int)
                timeline_data = sorted_ratings.iloc[indices][["date", "cumulative_votes"]].copy()
            else:
                timeline_data = sorted_ratings[["date", "cumulative_votes"]].copy()
                
        # Process rating counts for histogram
        dist_counts = {str(i): 0 for i in range(1, 6)}
        if not ratings_df.empty:
            freqs = ratings_df["rating"].value_counts().to_dict()
            for star, val in freqs.items():
                dist_counts[str(int(star))] = int(val)
                
        return {
            "title": title,
            "year": year,
            "avg_rating": avg_rating,
            "rating_count": rating_count,
            "percentile": percentile,
            "summary": summary_text,
            "timeline": timeline_data,
            "distribution": dist_counts,
            "std": rating_std
        }