src/recommendation_engine.py

import os
import pandas as pd
from typing import Dict, List, Optional, Tuple

class RecommendationEngine:
    def __init__(self):
        """Initialize the recommendation engine with data directory path."""
        self.data_dir = os.path.join(os.path.dirname(__file__), "..", "data")

        # Organization to logo filename mapping
        self.org_to_logo = {
            'OpenAI': 'openai.png',
            'Google': 'google.png',
            'Anthropic': 'anthropic.png',
            'Meta': 'meta.png',
            'Microsoft': 'microsoft.png',
            'Mistral': 'mistral.png',
            'Cohere': 'cohere.png',
            'xAI': 'xai.png',
            'DeepSeek': 'deepseek.png',
            'Alibaba': 'alibaba.png',
            'Z.ai': 'zai.png',
            'Moonshot': 'moonshot.png',
            'Tencent': 'tencent.png',
            'LG AI Research': 'lg.png',
            'Perplexity': 'perplexity.png',
            'AI21 Labs': 'ai21.png',
            'Reka': 'reka.png',
            'Inflection': 'inflection.png',
            'Stability AI': 'stability.png',
        }

        # Mapping from questionnaire choices to CSV files (with pricing data)
        self.use_case_mapping = {
            'conversational_knowledge': ['lmarena_text_with_pricing_or.csv'],
            'productivity_information': ['lmarena_text_with_pricing_or.csv'],
            'creative_content': ['lmarena_text_with_pricing_or.csv'],
            'technical_developer': ['lmarena_webdev_with_pricing_or.csv'],
            'advanced_automation': ['lmarena_text_with_pricing_or.csv'],
            'visual_ai': ['lmarena_vision_with_pricing_or.csv', 'lmarena_image_with_pricing.csv', 'lmarena_image-edit_with_pricing.csv']
        }

        # Visual AI sub-category mapping
        self.visual_ai_mapping = {
            'image_understanding': ['lmarena_vision_with_pricing_or.csv'],
            'image_generation': ['lmarena_image_with_pricing.csv'],
            'image_editing': ['lmarena_image-edit_with_pricing.csv']
        }

    def get_relevant_csv_files(self, use_case: str, visual_ai_type: Optional[str] = None) -> List[str]:
        """
        Map user's use case choice to relevant CSV files.

        Args:
            use_case: Primary use case selected by user
            visual_ai_type: Visual AI type if use_case is 'visual_ai'

        Returns:
            List of CSV filenames to use for recommendations
        """
        if use_case == 'visual_ai' and visual_ai_type:
            # If Visual AI is selected, use the specific visual AI mapping
            return self.visual_ai_mapping.get(visual_ai_type, [])
        else:
            # For all other use cases, use the primary mapping
            return self.use_case_mapping.get(use_case, [])

    def load_leaderboard_data(self, csv_files: List[str]) -> Dict[str, pd.DataFrame]:
        """
        Load leaderboard data from specified CSV files.

        Args:
            csv_files: List of CSV filenames to load

        Returns:
            Dictionary mapping filename to DataFrame
        """
        data = {}

        for csv_file in csv_files:
            file_path = os.path.join(self.data_dir, csv_file)

            if os.path.exists(file_path):
                try:
                    df = pd.read_csv(file_path)
                    data[csv_file] = df
                    print(f"✅ Loaded {csv_file}: {len(df)} models")
                except Exception as e:
                    print(f"❌ Error loading {csv_file}: {e}")
            else:
                print(f"⚠️  File not found: {csv_file}")

        return data

    def get_top_recommendations(self, use_case: str, visual_ai_type: Optional[str] = None, model_type: Optional[str] = None, priority: Optional[str] = None, top_n: int = 3) -> List[Dict]:
        """
        Get top N model recommendations for a given use case.

        Args:
            use_case: Primary use case selected by user
            visual_ai_type: Visual AI type if use_case is 'visual_ai'
            model_type: Model type preference ('open_only', 'proprietary_only_enterprise', 'no_preference')
            priority: User priority ('lower_cost' or 'better_performance')
            top_n: Number of top recommendations to return (default: 3)

        Returns:
            List of dictionaries containing model information
        """
        # Get relevant CSV files
        csv_files = self.get_relevant_csv_files(use_case, visual_ai_type)

        if not csv_files:
            print(f"❌ No CSV files mapped for use case: {use_case}")
            return []

        # Load the data
        loaded_data = self.load_leaderboard_data(csv_files)

        if not loaded_data:
            print(f"❌ No data loaded for use case: {use_case}")
            return []

        # Get the first CSV file's data (primary ranking source)
        primary_csv = csv_files[0]
        df = loaded_data.get(primary_csv)

        if df is None or df.empty:
            print(f"❌ No data in primary CSV: {primary_csv}")
            return []

        # Apply model type filter based on user preference
        if model_type == 'open_only':
            # Filter to exclude proprietary models
            df = df[df['license'] != 'Proprietary']
            print(f"🔓 Filtered to open weights models: {len(df)} models")
        elif model_type == 'proprietary_only_enterprise':
            # Filter to include only proprietary models
            df = df[df['license'] == 'Proprietary']
            print(f"🔒 Filtered to proprietary models: {len(df)} models")
        # If 'no_preference', no filtering is applied

        if df.empty:
            print(f"❌ No models found after applying model_type filter: {model_type}")
            return []

        # Apply sorting based on user's priority
        if priority == 'lower_cost':
            # Sort by pricing (ascending) - lower cost first
            # Convert pricing to float, handling non-numeric values
            df['pricing_numeric'] = pd.to_numeric(df['pricing'], errors='coerce')
            df = df.sort_values('pricing_numeric', ascending=True)
            print(f"💰 Sorted by cost (lowest first)")
        else:
            # Default: Sort by rank (better performance first)
            # CSV is already sorted by rank, but ensure it
            df = df.sort_values('rank', ascending=True)
            print(f"🏆 Sorted by performance (rank)")

        # Get top N models
        top_models = df.head(top_n)

        # Convert to list of dictionaries with relevant fields
        recommendations = []
        for idx, row in top_models.iterrows():
            organization = row.get('organization', 'Unknown')
            # Get logo filename from mapping
            logo_filename = self.org_to_logo.get(organization, '')
            logo_url = f'/static/images/logos/{logo_filename}' if logo_filename else ''

            model_info = {
                'rank': int(row.get('rank', idx + 1)),
                'model': row.get('displayed_name', row.get('model', 'Unknown')),  # Use displayed_name if available
                'arena_score': row.get('arena_score', 'N/A'),
                'votes': row.get('votes', 'N/A'),
                'organization': organization,
                'license': row.get('license', 'Unknown'),
                'url': row.get('url', ''),
                'ci': row.get('95_pct_ci', 'N/A'),
                'logo_url': logo_url,
                'knowledge_cutoff': row.get('knowledge_cutoff', 'N/A'),
                'pricing': row.get('pricing', 'N/A')  # Add pricing information
            }
            recommendations.append(model_info)

        print(f"✅ Generated {len(recommendations)} recommendations")
        return recommendations

    def process_user_input(self, questionnaire_data: Dict) -> Tuple[List[str], Dict[str, pd.DataFrame]]:
        """
        Process user questionnaire input and return relevant data.

        Args:
            questionnaire_data: Dictionary containing user's questionnaire responses

        Returns:
            Tuple of (csv_files_used, loaded_dataframes)
        """
        # Extract primary use case
        use_case = questionnaire_data.get('use_case')
        visual_ai_type = questionnaire_data.get('visual_ai_type')

        print(f"🎯 User selected use case: {use_case}")
        if visual_ai_type:
            print(f"🖼️  Visual AI type: {visual_ai_type}")

        # Get relevant CSV files
        csv_files = self.get_relevant_csv_files(use_case, visual_ai_type)

        if not csv_files:
            print(f"❌ No CSV files mapped for use case: {use_case}")
            return [], {}

        print(f"📊 Mapped to CSV files: {csv_files}")

        # Load the data
        loaded_data = self.load_leaderboard_data(csv_files)

        return csv_files, loaded_data

    def test_mapping(self):
        """Test the mapping logic with sample inputs."""
        print("🧪 Testing Use Case Mapping:")
        print("=" * 50)

        test_cases = [
            {'use_case': 'conversational_knowledge'},
            {'use_case': 'productivity_information'},
            {'use_case': 'creative_content'},
            {'use_case': 'technical_developer'},
            {'use_case': 'advanced_automation'},
            {'use_case': 'visual_ai', 'visual_ai_type': 'image_understanding'},
            {'use_case': 'visual_ai', 'visual_ai_type': 'image_generation'},
            {'use_case': 'visual_ai', 'visual_ai_type': 'image_editing'},
        ]

        for test_case in test_cases:
            print(f"\nTest case: {test_case}")
            csv_files, data = self.process_user_input(test_case)
            print(f"Result: {len(data)} datasets loaded")
            print("-" * 30)


if __name__ == "__main__":
    # Initialize and test the recommendation engine
    engine = RecommendationEngine()
    engine.test_mapping()