import numpy as np
from sklearn.neighbors import NearestNeighbors
from sklearn.metrics.pairwise import cosine_similarity
from tag_identification import TagIdentifier
import pandas as pd
import pickle
import os


class ImageSimilarity:
    def __init__(self, expert_data_path, community_data_path):
        self.identifier = TagIdentifier()
        
        # Load both datasets
        self.expert_data = pd.read_csv(expert_data_path)
        self.community_data = pd.read_csv(community_data_path)
        
        # Combine and prepare data
        self._prepare_image_data()
        
        # Try to load precomputed embeddings, otherwise compute
        self.embeddings_file = 'data/image_embeddings_200.pkl'
        if os.path.exists(self.embeddings_file):
            self._load_embeddings()
        else:
            self._compute_and_save_embeddings()
        
        # Build nearest neighbors index
        self._build_index()
        
    def _prepare_image_data(self):
        """Combine expert and community data and limit to 200 images"""
        # Prepare expert data
        expert_cols = ['brand_name', 'appraisal_value', 'front_tag', 'status']
        if 'year' in self.expert_data.columns:
            expert_cols.append('year')
        
        expert_images = self.expert_data[expert_cols].copy()
        expert_images = expert_images.rename(columns={'front_tag': 'image_url'})
        expert_images['source'] = 'expert'
        
        # Prepare community data
        community_cols = ['brand_name', 'appraisal_value', 'front_tag', 'status']
        if 'year' in self.community_data.columns:
            community_cols.append('year')
            
        community_images = self.community_data[community_cols].copy()
        community_images = community_images.rename(columns={'front_tag': 'image_url'})
        community_images['source'] = 'community'
        
        # Combine datasets
        self.image_data = pd.concat([expert_images, community_images], ignore_index=True)
        self.image_data = self.image_data.dropna(subset=['image_url'])
        
        # LIMIT TO 200 IMAGES FOR TESTING
        print(f"Total images available: {len(self.image_data)}")
        self.image_data = self.image_data.head(200)
        print(f"Limited to first 200 images for testing")
        
        # Reset index for consistent indexing
        self.image_data = self.image_data.reset_index(drop=True)
        
    def _compute_and_save_embeddings(self):
        """Compute CLIP embeddings for 200 images and save to file"""
        print(f"Computing embeddings for {len(self.image_data)} images...")
        embeddings = []
        valid_indices = []
        
        for idx, row in self.image_data.iterrows():
            try:
                image = self.identifier.load_image(row['image_url'])
                features = self.identifier.get_image_features(image)
                embeddings.append(features)
                valid_indices.append(idx)
                
                # Progress update every 25 images
                if len(embeddings) % 25 == 0:
                    print(f"Processed {len(embeddings)}/200 images...")
                    
            except Exception as e:
                print(f"Failed to process image {idx}: {e}")
                continue
        
        # Store embeddings and update data to only include valid images
        self.embeddings = np.array(embeddings)
        self.image_data = self.image_data.iloc[valid_indices].reset_index(drop=True)
        
        # Save to file
        with open(self.embeddings_file, 'wb') as f:
            pickle.dump({
                'embeddings': self.embeddings,
                'image_data': self.image_data
            }, f)
        
        print(f"✅ Computed and saved embeddings for {len(embeddings)} images")
    
    def _load_embeddings(self):
        """Load precomputed embeddings from file"""
        print("Loading precomputed embeddings...")
        with open(self.embeddings_file, 'rb') as f:
            data = pickle.load(f)
            self.embeddings = data['embeddings']
            self.image_data = data['image_data']
        print(f"✅ Loaded embeddings for {len(self.embeddings)} images")
    
    def _build_index(self):
        """Build nearest neighbors index for fast similarity search"""
        if len(self.embeddings) > 0:
            self.nn_model = NearestNeighbors(
                n_neighbors=min(10, len(self.embeddings)),
                metric='cosine'
            )
            self.nn_model.fit(self.embeddings)
        else:
            self.nn_model = None
    
    def find_similar(self, image_url, top_n=5):
        """Find visually similar images using CLIP embeddings"""
        if self.nn_model is None or len(self.embeddings) == 0:
            return []
        
        try:
            # Get features for query image
            query_image = self.identifier.load_image(image_url)
            query_features = self.identifier.get_image_features(query_image)
            
            # Find similar images
            distances, indices = self.nn_model.kneighbors(
                [query_features], 
                n_neighbors=min(top_n, len(self.embeddings))
            )
            
            results = []
            for i, idx in enumerate(indices[0]):
                row = self.image_data.iloc[idx]
                
                result = {
                    'image_url': row['image_url'],
                    'brand_name': row.get('brand_name', 'unknown'),
                    'appraisal_value': float(row.get('appraisal_value', 0)),
                    'status': int(row.get('status', 1)),
                    'source': row.get('source', 'unknown'),
                    'similarity_score': float(1 - distances[0][i]),  # Convert distance to similarity
                    'year': self._extract_year_from_metadata(row)
                }
                results.append(result)
                
            return results
            
        except Exception as e:
            print(f"Error in similarity search: {e}")
            return []

    def _extract_year_from_metadata(self, item):
        """Extract year information from various possible fields"""
        # Try different year fields
        for year_field in ['year', 'from_decade', 'to_decade', 'manufacturing_year']:
            if year_field in item and pd.notna(item[year_field]):
                year_val = item[year_field]
                if isinstance(year_val, (int, float)) and year_val > 1900:
                    return str(int(year_val))
                elif isinstance(year_val, str) and year_val.isdigit():
                    return year_val
        
        # Try to extract from brand_name or other text fields
        if 'brand_name' in item and pd.notna(item['brand_name']):
            text = str(item['brand_name'])
            # Look for 4-digit years in text
            import re
            years = re.findall(r'\b(19|20)\d{2}\b', text)
            if years:
                return years[0]
        
        return "unknown"
