#!/usr/bin/env python3
"""
Debug Transformer Recommendation Engine

Debug version of TransformerRecommendationEngine with detailed print statements
at each layer to help identify issues with transformer recommendation inference.
"""

import tensorflow as tf
import numpy as np
import pandas as pd
import pickle
import faiss
from typing import Dict, List, Tuple, Optional
import os

from debug_transformer_user_tower import DebugTransformerUserTower
from debug_transformer_item_tower import DebugTransformerItemTower


class DebugTransformerRecommendationEngine:
    """Debug transformer-based recommendation engine with detailed layer output logging."""

    def __init__(self, artifacts_path: str = "src/artifacts/transformers/"):
        self.artifacts_path = artifacts_path
        self.transformer_user_tower = None
        self.transformer_item_tower = None
        self.rating_model = None
        self.two_tower_model = None
        self.faiss_index = None
        self.items_df = None
        self.price_normalizer = None
        self.vocabularies = None
        self.faiss_metadata = None

        print(f"\n🚀 INITIALIZING DEBUG TRANSFORMER RECOMMENDATION ENGINE")
        print(f"   Artifacts path: {artifacts_path}")

        # Load components
        self._load_all_components()

    def _load_all_components(self):
        """Load all required components for transformer inference with debug logging."""

        print(f"\n🔄 LOADING TRANSFORMER COMPONENTS...")

        # Load vocabularies
        self._load_vocabularies()

        # Load price normalizer
        self._load_price_normalizer()

        # Load items dataframe for metadata
        self.items_df = pd.read_csv("datasets/items.csv")
        print(f"✅ Loaded items dataset: {len(self.items_df)} items")

        # Load trained models
        self._load_transformer_user_tower()
        self._load_transformer_item_tower()
        self._load_rating_model()

        # Load FAISS index
        self._load_faiss_index()

        print(f"\n✅ ALL DEBUG TRANSFORMER COMPONENTS LOADED SUCCESSFULLY!")

    def _load_vocabularies(self):
        """Load vocabularies from transformer artifacts with debug info."""
        vocab_path = f"{self.artifacts_path}/transformer_vocabularies.pkl"
        try:
            with open(vocab_path, 'rb') as f:
                self.vocabularies = pickle.load(f)
            print(f"✅ Loaded transformer vocabularies:")
            for vocab_name, vocab_dict in self.vocabularies.items():
                print(f"     {vocab_name}: {len(vocab_dict)} entries")
        except Exception as e:
            print(f"❌ Warning: Could not load transformer vocabularies: {e}")
            # Fallback to default vocabularies
            self.vocabularies = {
                'item_vocab': {},
                'category_vocab': {},
                'category_code_vocab': {},
                'brand_vocab': {},
                'user_vocab': {}
            }

    def _load_price_normalizer(self):
        """Load price normalizer from transformer artifacts with debug info."""
        normalizer_path = f"{self.artifacts_path}/transformer_price_normalizer.pkl"
        try:
            with open(normalizer_path, 'rb') as f:
                self.price_normalizer = pickle.load(f)
            print(f"✅ Loaded transformer price normalizer")
            # Extract price statistics for item tower
            self.price_mean = float(self.price_normalizer.mean_[0]) if hasattr(self.price_normalizer, 'mean_') else 0.0
            self.price_std = float(self.price_normalizer.scale_[0]) if hasattr(self.price_normalizer, 'scale_') else 1.0
            print(f"   Price statistics: mean={self.price_mean:.4f}, std={self.price_std:.4f}")
        except Exception as e:
            print(f"❌ Warning: Could not load transformer price normalizer: {e}")
            # Create dummy normalizer with default stats
            from sklearn.preprocessing import StandardScaler
            self.price_normalizer = StandardScaler()
            self.price_mean = 0.0
            self.price_std = 1.0

    def _load_transformer_user_tower(self):
        """Load trained transformer user tower using debug version."""

        print(f"\n🧑 LOADING DEBUG TRANSFORMER USER TOWER...")

        # Create debug transformer user tower with same config as training
        self.transformer_user_tower = DebugTransformerUserTower(
            embedding_dim=128,
            transformer_layers=2,
            transformer_heads=4,
            transformer_ff_dim=256,
            hidden_dims=[256, 128],
            dropout_rate=0.2,
            max_sequence_length=100,
            compatibility_mode=False
        )

        # Build user tower with dummy input
        dummy_input = {
            'age': tf.constant([2]),
            'gender': tf.constant([1]),
            'income': tf.constant([2]),
            'profession': tf.constant([0]),
            'location': tf.constant([0]),
            'education_level': tf.constant([2]),
            'marital_status': tf.constant([1]),
            'item_history_embeddings': tf.constant([[[0.0] * 128] * 100]),
            'attention_masks': tf.constant([[1.0] * 100]),
            'sequence_lengths': tf.constant([10])
        }
        _ = self.transformer_user_tower(dummy_input)

        # Try to load weights (best first, then fallback)
        weight_paths = [
            f"{self.artifacts_path}/transformer_user_tower_weights_best",
            f"{self.artifacts_path}/transformer_user_tower_weights"
        ]

        for weight_path in weight_paths:
            try:
                self.transformer_user_tower.load_weights(weight_path)
                print(f"✅ Loaded debug user tower weights from {weight_path}")
                break
            except Exception as e:
                print(f"⚠️ Could not load from {weight_path}: {e}")
        else:
            print(f"❌ Warning: Could not load any debug user tower weights")

    def _load_transformer_item_tower(self):
        """Load trained transformer item tower using debug version."""

        print(f"\n🏪 LOADING DEBUG TRANSFORMER ITEM TOWER...")

        if not self.vocabularies:
            print(f"❌ Cannot load debug item tower without vocabularies")
            return

        if not hasattr(self, 'price_mean') or not hasattr(self, 'price_std'):
            print(f"❌ Warning: Price normalization parameters not loaded, using defaults")
            self.price_mean = 0.0
            self.price_std = 1.0

        # Create debug transformer item tower with transformer config and price normalization
        self.transformer_item_tower = DebugTransformerItemTower(
            item_vocab_size=len(self.vocabularies['item_vocab']),
            category_vocab_size=len(self.vocabularies['category_vocab']),
            category_code_vocab_size=len(self.vocabularies['category_code_vocab']),
            brand_vocab_size=len(self.vocabularies['brand_vocab']),
            embedding_dim=128,
            hidden_dims=[256, 128],
            dropout_rate=0.2,
            price_mean=self.price_mean,
            price_std=self.price_std
        )

        # Build transformer item tower with dummy input
        dummy_input = {
            'product_id': tf.constant([0]),
            'category_id': tf.constant([0]),
            'category_code_id': tf.constant([0]),
            'brand_id': tf.constant([0]),
            'price': tf.constant([0.0])
        }
        _ = self.transformer_item_tower(dummy_input)

        # Try to load transformer item tower weights (best first, then fallback)
        weight_paths = [
            f"{self.artifacts_path}/transformer_item_tower_weights_finetuned_best",
            f"{self.artifacts_path}/transformer_item_tower_weights",
        ]

        for weight_path in weight_paths:
            try:
                self.transformer_item_tower.load_weights(weight_path)
                print(f"✅ Loaded debug item tower weights from {weight_path}")
                break
            except Exception as e:
                print(f"⚠️ Could not load from {weight_path}: {e}")
        else:
            print(f"❌ Warning: Could not load any debug item tower weights")

    def _load_rating_model(self):
        """Load trained rating prediction model with debug info."""

        print(f"\n⭐ LOADING RATING MODEL...")

        # Create rating model with same architecture as training
        self.rating_model = tf.keras.Sequential([
            tf.keras.layers.Dense(64, activation="relu"),
            tf.keras.layers.Dropout(0.2),
            tf.keras.layers.Dense(1, activation="sigmoid")
        ])

        # Build model with dummy input
        dummy_input = tf.constant([[0.0] * 256])
        _ = self.rating_model(dummy_input)

        # Try to load transformer rating model weights
        weight_paths = [
            f"{self.artifacts_path}/transformer_rating_model_weights_best",
            f"{self.artifacts_path}/transformer_rating_model_weights"
        ]

        for weight_path in weight_paths:
            try:
                self.rating_model.load_weights(weight_path)
                print(f"✅ Loaded rating model weights from {weight_path}")
                break
            except Exception as e:
                print(f"⚠️ Could not load from {weight_path}: {e}")
        else:
            print(f"❌ Warning: Could not load any rating model weights")

    def _load_faiss_index(self):
        """Load FAISS index for item similarity search with debug info."""

        print(f"\n🔍 LOADING FAISS INDEX...")

        try:
            # Load FAISS index
            index_path = f"{self.artifacts_path}/transformer_faiss_index.bin"
            if os.path.exists(index_path):
                self.faiss_index = faiss.read_index(index_path)
                print(f"✅ Loaded transformer FAISS index: {self.faiss_index.ntotal} items")
            else:
                print(f"⚠️ FAISS index not found at {index_path}")

            # Load metadata
            metadata_path = f"{self.artifacts_path}/transformer_faiss_metadata.pkl"
            if os.path.exists(metadata_path):
                with open(metadata_path, 'rb') as f:
                    self.faiss_metadata = pickle.load(f)
                print(f"✅ Loaded transformer FAISS metadata: {len(self.faiss_metadata.get('item_id_to_idx', {}))} item mappings")
            else:
                print(f"⚠️ FAISS metadata not found at {metadata_path}")

        except Exception as e:
            print(f"❌ Warning: Could not load transformer FAISS components: {e}")

    # Categorization methods (same as original)
    def categorize_age(self, age: float) -> int:
        """Categorize age into 6 demographic groups."""
        if age < 18:
            return 0  # Teen
        elif age < 26:
            return 1  # Young Adult
        elif age < 36:
            return 2  # Adult
        elif age < 51:
            return 3  # Middle Age
        elif age < 66:
            return 4  # Mature
        else:
            return 5  # Senior

    def categorize_income(self, income: float) -> int:
        """Categorize income based on quintiles."""
        if income < 30000:
            return 0
        elif income < 50000:
            return 1
        elif income < 75000:
            return 2
        elif income < 100000:
            return 3
        else:
            return 4

    def categorize_profession(self, profession: str) -> int:
        """Categorize profession into numeric categories."""
        profession_map = {
            "Technology": 0,
            "Healthcare": 1,
            "Education": 2,
            "Finance": 3,
            "Retail": 4,
            "Manufacturing": 5,
            "Services": 6,
            "Other": 7
        }
        return profession_map.get(profession, 7)

    def categorize_location(self, location: str) -> int:
        """Categorize location into numeric categories."""
        location_map = {
            "Urban": 0,
            "Suburban": 1,
            "Rural": 2
        }
        return location_map.get(location, 0)

    def categorize_education_level(self, education: str) -> int:
        """Categorize education level into numeric categories."""
        education_map = {
            "High School": 0,
            "Some College": 1,
            "Bachelor's": 2,
            "Master's": 3,
            "PhD+": 4
        }
        return education_map.get(education, 0)

    def categorize_marital_status(self, marital_status: str) -> int:
        """Categorize marital status into numeric categories."""
        marital_map = {
            "Single": 0,
            "Married": 1,
            "Divorced": 2,
            "Widowed": 3
        }
        return marital_map.get(marital_status, 0)

    def prepare_user_features(self,
                            age: int,
                            gender: str,
                            income: float,
                            profession: str = "Other",
                            location: str = "Urban",
                            education_level: str = "High School",
                            marital_status: str = "Single",
                            interaction_history: List[int] = None) -> Dict[str, tf.Tensor]:
        """Prepare user features for transformer inference with debug logging."""

        print(f"\n👤 PREPARING USER FEATURES:")
        print(f"   Age: {age}")
        print(f"   Gender: {gender}")
        print(f"   Income: {income}")
        print(f"   Profession: {profession}")
        print(f"   Location: {location}")
        print(f"   Education: {education_level}")
        print(f"   Marital status: {marital_status}")
        print(f"   Interaction history length: {len(interaction_history) if interaction_history else 0}")

        if interaction_history is None:
            interaction_history = []

        # Convert gender
        gender_numeric = 1 if gender.lower() == 'male' else 0

        # Categorize all demographics
        age_category = self.categorize_age(age)
        income_category = self.categorize_income(income)
        profession_category = self.categorize_profession(profession)
        location_category = self.categorize_location(location)
        education_category = self.categorize_education_level(education_level)
        marital_category = self.categorize_marital_status(marital_status)

        print(f"\n📊 CATEGORICAL MAPPINGS:")
        print(f"   Age {age} → {age_category}")
        print(f"   Gender {gender} → {gender_numeric}")
        print(f"   Income {income} → {income_category}")
        print(f"   Profession {profession} → {profession_category}")
        print(f"   Location {location} → {location_category}")
        print(f"   Education {education_level} → {education_category}")
        print(f"   Marital {marital_status} → {marital_category}")

        # Get item embeddings for history
        print(f"\n📚 PROCESSING INTERACTION HISTORY:")
        history_embeddings = []
        valid_sequence_length = 0

        for i, item_id in enumerate(interaction_history):
            embedding = self._get_item_embedding_from_faiss(item_id)
            if embedding is not None:
                history_embeddings.append(embedding)
                valid_sequence_length += 1
                if i < 3:  # Only log first few items
                    print(f"   Item {item_id}: found embedding (norm={np.linalg.norm(embedding):.4f})")
            else:
                # Use zero embedding for unknown items
                history_embeddings.append(np.zeros(128))
                if i < 3:
                    print(f"   Item {item_id}: using zero embedding (unknown item)")

        print(f"   Valid interactions: {valid_sequence_length}/{len(interaction_history)}")

        # Pad or truncate to max_history_length
        max_history_length = 100
        attention_mask = []

        if len(history_embeddings) < max_history_length:
            # Create attention mask: 1 for valid tokens, 0 for padding
            attention_mask = [1.0] * len(history_embeddings) + [0.0] * (max_history_length - len(history_embeddings))
            # Add padding at the END
            padding = [np.zeros(128)] * (max_history_length - len(history_embeddings))
            history_embeddings = history_embeddings + padding
        else:
            # Keep most recent interactions
            history_embeddings = history_embeddings[-max_history_length:]
            attention_mask = [1.0] * max_history_length
            valid_sequence_length = max_history_length

        print(f"   Final sequence length: {valid_sequence_length}")
        print(f"   Attention mask valid tokens: {sum(attention_mask)}")

        history_embeddings = np.array(history_embeddings, dtype=np.float32)
        attention_mask = np.array(attention_mask, dtype=np.float32)

        # Prepare features for transformer user tower
        user_features = {
            'age': tf.constant([age_category]),
            'gender': tf.constant([gender_numeric]),
            'income': tf.constant([income_category]),
            'profession': tf.constant([profession_category]),
            'location': tf.constant([location_category]),
            'education_level': tf.constant([education_category]),
            'marital_status': tf.constant([marital_category]),
            'item_history_embeddings': tf.constant([history_embeddings]),
            'attention_masks': tf.constant([attention_mask]),
            'sequence_lengths': tf.constant([valid_sequence_length])
        }

        return user_features

    def prepare_item_features(self, item_ids: List[int]) -> Dict[str, tf.Tensor]:
        """Prepare item features for inference with debug logging."""

        print(f"\n🏪 PREPARING ITEM FEATURES:")
        print(f"   Item IDs: {item_ids}")

        features = {
            'product_id': [],
            'category_id': [],
            'category_code_id': [],
            'brand_id': [],
            'price': []
        }

        for item_id in item_ids:
            # Find item in dataframe
            item_row = self.items_df[self.items_df['product_id'] == item_id]

            if len(item_row) > 0:
                item_row = item_row.iloc[0]

                vocab_product_id = self.vocabularies['item_vocab'].get(item_id, 0)
                vocab_category_id = self.vocabularies['category_vocab'].get(item_row['category_id'], 0)
                vocab_category_code_id = self.vocabularies['category_code_vocab'].get(item_row['category_code'], 0)
                vocab_brand_id = self.vocabularies['brand_vocab'].get(item_row['brand'], 0)
                item_price = float(item_row['price'])

                print(f"   Item {item_id}:")
                print(f"     Product vocab ID: {vocab_product_id}")
                print(f"     Category ID: {item_row['category_id']} → vocab {vocab_category_id}")
                print(f"     Category code: {item_row['category_code']} → vocab {vocab_category_code_id}")
                print(f"     Brand: {item_row['brand']} → vocab {vocab_brand_id}")
                print(f"     Price: {item_price}")

                features['product_id'].append(vocab_product_id)
                features['category_id'].append(vocab_category_id)
                features['category_code_id'].append(vocab_category_code_id)
                features['brand_id'].append(vocab_brand_id)
                features['price'].append(item_price)
            else:
                print(f"   Item {item_id}: UNKNOWN ITEM - using default values")
                # Unknown item
                features['product_id'].append(0)
                features['category_id'].append(0)
                features['category_code_id'].append(0)
                features['brand_id'].append(0)
                features['price'].append(0.0)

        # Convert to tensors
        return {k: tf.constant(v) for k, v in features.items()}

    def get_user_embedding(self,
                          age: int,
                          gender: str,
                          income: float,
                          profession: str = "Other",
                          location: str = "Urban",
                          education_level: str = "High School",
                          marital_status: str = "Single",
                          interaction_history: List[int] = None) -> np.ndarray:
        """Get user embedding from transformer user tower with debug logging."""

        if self.transformer_user_tower is None:
            raise RuntimeError("Debug transformer user tower not loaded")

        print(f"\n🧑 GENERATING USER EMBEDDING:")

        user_features = self.prepare_user_features(
            age, gender, income, profession, location, education_level, marital_status, interaction_history
        )
        user_embedding = self.transformer_user_tower(user_features, training=False)

        print(f"\n✨ FINAL USER EMBEDDING:")
        print(f"   Shape: {user_embedding.shape}")
        print(f"   L2 norm: {tf.norm(user_embedding).numpy():.6f}")
        print(f"   First 10 values: {user_embedding.numpy()[0][:10]}")

        return user_embedding.numpy()[0]

    def _get_item_embedding_from_faiss(self, item_id: int) -> Optional[np.ndarray]:
        """Get item embedding from FAISS index with debug logging."""

        if self.faiss_metadata is None or self.faiss_index is None:
            return None

        item_id_to_idx = self.faiss_metadata.get('item_id_to_idx', {})

        if item_id in item_id_to_idx:
            idx = item_id_to_idx[item_id]
            try:
                embedding = self.faiss_index.reconstruct(idx)
                return embedding
            except:
                return None

        return None

    def get_item_embedding(self, item_id: int) -> Optional[np.ndarray]:
        """Get item embedding from FAISS index or item tower with debug logging."""

        print(f"\n🏪 GENERATING ITEM EMBEDDING for item {item_id}:")

        # First try FAISS index (faster)
        embedding = self._get_item_embedding_from_faiss(item_id)
        if embedding is not None:
            print(f"   Found in FAISS index")
            print(f"   Shape: {embedding.shape}")
            print(f"   L2 norm: {np.linalg.norm(embedding):.6f}")
            print(f"   First 5 values: {embedding[:5]}")
            return embedding

        # Fall back to transformer item tower for new items
        if self.transformer_item_tower is not None:
            print(f"   Not in FAISS, using item tower")
            item_features = self.prepare_item_features([item_id])
            item_embedding = self.transformer_item_tower(item_features, training=False)

            print(f"\n✨ FINAL ITEM EMBEDDING:")
            print(f"   Shape: {item_embedding.shape}")
            print(f"   L2 norm: {tf.norm(item_embedding).numpy():.6f}")
            print(f"   First 5 values: {item_embedding.numpy()[0][:5]}")

            return item_embedding.numpy()[0]

        return None

    def search_similar_items(self, user_embedding: np.ndarray, k: int = 10) -> List[Tuple[int, float]]:
        """Search for similar items using FAISS index with debug logging."""

        print(f"\n🔍 SEARCHING SIMILAR ITEMS:")
        print(f"   User embedding shape: {user_embedding.shape}")
        print(f"   User embedding L2 norm: {np.linalg.norm(user_embedding):.6f}")
        print(f"   Searching for top {k} similar items")

        if self.faiss_index is None or self.faiss_metadata is None:
            print(f"   ❌ FAISS index or metadata not available")
            return []

        # Ensure embedding is the right shape
        if len(user_embedding.shape) == 1:
            user_embedding = user_embedding.reshape(1, -1)

        print(f"   FAISS index size: {self.faiss_index.ntotal} items")

        # Search FAISS index
        similarities, indices = self.faiss_index.search(user_embedding.astype('float32'), k)

        print(f"   Raw similarities: {similarities[0]}")
        print(f"   Raw indices: {indices[0]}")

        # Convert to item IDs
        idx_to_item_id = self.faiss_metadata.get('idx_to_item_id', {})
        results = []

        for i, (similarity, idx) in enumerate(zip(similarities[0], indices[0])):
            if idx in idx_to_item_id:
                item_id = idx_to_item_id[idx]
                results.append((item_id, float(similarity)))
                print(f"   #{i+1}: Item {item_id}, similarity = {similarity:.6f}")

        print(f"\n✅ Found {len(results)} similar items")
        return results

    def recommend_items(self,
                       age: int,
                       gender: str,
                       income: float,
                       profession: str = "Other",
                       location: str = "Urban",
                       education_level: str = "High School",
                       marital_status: str = "Single",
                       interaction_history: List[int] = None,
                       k: int = 10,
                       exclude_history: bool = True) -> List[Tuple[int, float, Dict]]:
        """Generate recommendations using transformer-based approach with debug logging."""

        print(f"\n🎯 GENERATING TRANSFORMER RECOMMENDATIONS")
        print(f"   Target recommendations: {k}")
        print(f"   Exclude history: {exclude_history}")

        # Get user embedding from transformer user tower
        user_embedding = self.get_user_embedding(
            age, gender, income, profession, location, education_level, marital_status, interaction_history
        )

        # Search for similar items
        similar_items = self.search_similar_items(user_embedding, k * 2)

        # Filter out interaction history
        history_set = set(interaction_history) if (exclude_history and interaction_history) else set()
        if exclude_history and interaction_history:
            print(f"\n🚫 FILTERING HISTORY: Excluding {len(history_set)} items")

        recommendations = []
        for item_id, score in similar_items:
            if item_id in history_set:
                print(f"   Skipping item {item_id} (in history)")
                continue

            item_info = self._get_item_info(item_id)
            recommendations.append((item_id, score, item_info))
            print(f"   ✅ Added item {item_id}: {item_info['brand']} - ${item_info['price']:.2f} (score: {score:.4f})")

            if len(recommendations) >= k:
                break

        print(f"\n🎉 FINAL RECOMMENDATIONS: {len(recommendations)} items")
        return recommendations

    def _get_item_info(self, item_id: int) -> Dict:
        """Get item metadata."""

        item_row = self.items_df[self.items_df['product_id'] == item_id]

        if len(item_row) > 0:
            item_row = item_row.iloc[0]
            return {
                'product_id': int(item_id),
                'category_id': int(item_row['category_id']),
                'category_code': str(item_row['category_code']),
                'brand': str(item_row['brand']) if pd.notna(item_row['brand']) else 'Unknown',
                'price': float(item_row['price'])
            }
        else:
            return {
                'product_id': int(item_id),
                'category_id': 0,
                'category_code': 'unknown',
                'brand': 'Unknown',
                'price': 0.0
            }


def main():
    """Demo the debug transformer recommendation engine."""

    # Initialize debug transformer recommendation engine
    print(f"\n🚀 INITIALIZING DEBUG TRANSFORMER RECOMMENDATION ENGINE...")
    engine = DebugTransformerRecommendationEngine()

    # Demo user profile
    demo_user = {
        'age': 32,
        'gender': 'male',
        'income': 75000,
        'profession': 'Technology',
        'location': 'Urban',
        'education_level': "Bachelor's",
        'marital_status': 'Married',
        'interaction_history': [1000978, 1001588, 1001618]  # Sample item IDs
    }

    print(f"\n👤 DEMO USER PROFILE:")
    for k, v in demo_user.items():
        print(f"   {k}: {v}")

    # Generate transformer recommendations
    print(f"\n🎯 GENERATING DEBUG TRANSFORMER RECOMMENDATIONS...")
    transformer_recs = engine.recommend_items(**demo_user, k=5)

    print(f"\n📊 TRANSFORMER RECOMMENDATIONS SUMMARY:")
    for i, (item_id, score, info) in enumerate(transformer_recs, 1):
        print(f"   {i}. Item {item_id}: {info['brand']} - ${info['price']:.2f} (Score: {score:.4f})")

    print(f"\n✅ DEBUG TRANSFORMER RECOMMENDATION ENGINE DEMO COMPLETED!")


if __name__ == "__main__":
    main()