import os
import torch
import numpy as np
from PIL import Image
from transformers import AutoImageProcessor, AutoModel
from torchvision import transforms, datasets
import streamlit as st
import faiss
import time
import gc
import psutil
from pathlib import Path

class ImageSimilaritySearch:
    def __init__(self):
        self.MODEL_CHECKPOINT = "google/vit-large-patch16-224-in21k"
        self.DEFAULT_DATASET_PATH = r"G:\thedeathofyouth"  # Windows path format
        self.MAX_IMAGE_WIDTH = 300
        self.device = None
        self.processor = None
        self.model = None
        self.transformation_chain = None
        self.class_names = []
        self.index = None
        self.gpu_resources = None
        self._init_session_state()
        
    def _init_session_state(self):
        keys = [
            'run', 'clear_embeddings', 'query_image', 'last_results',
            'candidate_paths', 'candidate_labels', 'index_built', 'total_images'
        ]
        for key in keys:
            if key not in st.session_state:
                st.session_state[key] = None if key in ['last_results', 'total_images'] else False
        
    def setup_sidebar(self):
        st.sidebar.title("⚙️ Configuration")
        self.DATASET_PATH = st.sidebar.text_input("📂 Dataset Path", self.DEFAULT_DATASET_PATH)
        self.SEED = st.sidebar.number_input("🎲 Random Seed", value=42, min_value=0)
        self.BATCH_SIZE = st.sidebar.slider("📦 Batch Size", 64, 512, 256, step=64)
        self.TOP_K = st.sidebar.slider("🔍 Similar Images to Show", 1, 50, 12)
        self.COLUMNS = st.sidebar.slider("Number of Columns", 2, 8, 4)
        
        st.sidebar.markdown("### 🔧 Hardware Status")
        if torch.cuda.is_available():
            gpu_name = torch.cuda.get_device_name(0)
            gpu_mem = f"{torch.cuda.get_device_properties(0).total_memory/1024**3:.1f} GB"
            st.sidebar.markdown(f"- **GPU**: {gpu_name} ({gpu_mem} VRAM)")
        st.sidebar.markdown(f"- **CPU**: Ryzen 9 5900X (12-core)")
        st.sidebar.markdown(f"- **RAM**: 64GB")
        st.sidebar.markdown(f"- **Storage**: NVMe SSD")
        
        st.sidebar.markdown("### 🚀 Performance Settings")
        self.INDEX_TYPE = st.sidebar.selectbox(
            "🔧 Search Algorithm",
            ["GPU-Powered Exact Search", "HNSW (Fast Approximate)"],
            index=0
        )
        self.USE_HALF_PRECISION = st.sidebar.checkbox("Use FP16 Precision", value=True)
        
        if st.sidebar.button("🚀 Build New Index", use_container_width=True):
            st.session_state.clear_embeddings = True
            st.session_state.query_image = None
            st.session_state.index_built = False
            st.rerun()
            
    @st.cache_resource
    def load_model(_self):
        device = "cuda" if torch.cuda.is_available() else "cpu"
        st.toast(f"✅ Using device: {device}", icon="🚀")
        
        processor = AutoImageProcessor.from_pretrained(_self.MODEL_CHECKPOINT)
        model = AutoModel.from_pretrained(_self.MODEL_CHECKPOINT)
        
        # Use half precision if supported
        if _self.USE_HALF_PRECISION and device == "cuda":
            model = model.half()
            st.toast("Using FP16 precision for faster inference", icon="⚡")
        
        model = model.to(device)
        model.eval()
        
        return device, processor, model

    def setup_model(self):
        self.device, self.processor, self.model = self.load_model()
        self.transformation_chain = self._get_transforms()
        
        # Initialize GPU resources for FAISS
        if self.device == "cuda" and faiss.get_num_gpus() > 0:
            self.gpu_resources = faiss.StandardGpuResources()
            st.toast(f"FAISS using {faiss.get_num_gpus()} GPUs", icon="🎮")
    
    def _get_transforms(self):
        normalize = transforms.Normalize(
            mean=self.processor.image_mean, 
            std=self.processor.image_std
        )
        size = self.processor.size.get("shortest_edge", 224)
        
        return transforms.Compose([
            transforms.Lambda(lambda img: img.convert("RGB")),
            transforms.Resize(size, interpolation=transforms.InterpolationMode.BICUBIC),
            transforms.CenterCrop(size),
            transforms.ToTensor(),
            normalize,
        ])
    
    def load_dataset(self, dataset_path: str):
        """Load dataset with Windows path compatibility"""
        try:
            # Convert to Path object for Windows compatibility
            dataset_path = Path(dataset_path)
            if not dataset_path.exists():
                raise FileNotFoundError(f"Path not found: {dataset_path}")
                
            full_dataset = datasets.ImageFolder(root=str(dataset_path))
            class_names = full_dataset.classes
            st.session_state.total_images = len(full_dataset)
            return full_dataset, class_names
        except Exception as e:
            st.error(f"❌ Dataset loading failed: {str(e)}")
            st.stop()
    
    def batch_extract_embeddings(self, images: list):
        """Batch process images to embeddings with memory optimization"""
        try:
            # Convert images to tensor batch
            images_transformed = torch.stack(
                [self.transformation_chain(img) for img in images]
            ).to(self.device)
            
            # Use half precision if enabled
            if self.USE_HALF_PRECISION and self.device == "cuda":
                images_transformed = images_transformed.half()
            
            with torch.no_grad():
                with torch.cuda.amp.autocast(enabled=self.USE_HALF_PRECISION):
                    embeddings = self.model(images_transformed).last_hidden_state[:, 0]
            
            # Move to CPU and convert to float32 for FAISS
            return embeddings.cpu().float().numpy()
        except RuntimeError as e:
            if "out of memory" in str(e).lower():
                st.warning("⚠️ GPU memory overflow! Reducing batch size...")
                return None
            raise
    
    def build_faiss_index(self, embeddings: np.ndarray):
        """Create optimized FAISS index for Windows/CUDA"""
        dim = embeddings.shape[1]
        embeddings = embeddings.astype('float32')
        
        # Normalize embeddings for cosine similarity
        faiss.normalize_L2(embeddings)
        
        if self.INDEX_TYPE == "GPU-Powered Exact Search" and self.device == "cuda":
            # Use GPU Flat index for exact search
            index_flat = faiss.IndexFlatIP(dim)
            index = faiss.index_cpu_to_gpu(self.gpu_resources, 0, index_flat)
            index.add(embeddings)
            return index
        else:
            # Use HNSW for approximate search
            index = faiss.IndexHNSWFlat(dim, 32)
            index.hnsw.efConstruction = 100
            index.hnsw.efSearch = 128
            index.add(embeddings)
            return index
    
    def generate_embeddings(self):
        """Process dataset with Windows path handling"""
        full_dataset, self.class_names = self.load_dataset(self.DATASET_PATH)
        n_images = st.session_state.total_images
        
        if n_images == 0:
            st.error("❌ No images found in dataset")
            return
        
        # Memory estimation
        embedding_size = 1024 * 4  # float32 = 4 bytes
        estimated_mem = (n_images * embedding_size) / (1024**3)
        st.info(f"💾 Estimated memory: {estimated_mem:.2f} GB for embeddings")
        
        # Batch processing
        all_embeddings = []
        st.session_state.candidate_paths = []
        st.session_state.candidate_labels = []
        
        progress_bar = st.progress(0)
        status_text = st.empty()
        ram_usage = st.empty()
        gpu_stats = st.empty()
        
        # Get all samples with Windows-compatible paths
        all_samples = [(str(Path(path)), label) for path, label in full_dataset.samples]
        
        # Process in batches with adaptive batch sizing
        current_batch_size = self.BATCH_SIZE
        processed = 0
        
        while processed < n_images:
            batch_end = min(processed + current_batch_size, n_images)
            batch_samples = all_samples[processed:batch_end]
            
            # Load images
            images = []
            for path, _ in batch_samples:
                try:
                    img = Image.open(path)
                    if img.mode != "RGB":
                        img = img.convert("RGB")
                    images.append(img)
                except Exception as e:
                    st.warning(f"Skipping corrupted image: {path}")
                    continue
            
            # Extract embeddings
            embeddings = self.batch_extract_embeddings(images)
            
            if embeddings is None:  # OOM occurred
                current_batch_size = max(32, current_batch_size // 2)
                st.toast(f"Reduced batch size to {current_batch_size}", icon="⚠️")
                torch.cuda.empty_cache()
                gc.collect()
                continue
            
            all_embeddings.append(embeddings)
            st.session_state.candidate_labels.extend([label for _, label in batch_samples])
            st.session_state.candidate_paths.extend([path for path, _ in batch_samples])
            
            processed = batch_end
            progress = processed / n_images
            
            # Update progress
            progress_bar.progress(progress)
            status_text.text(f"🔄 Processed {processed}/{n_images} images")
            
            # Memory stats
            if self.device == "cuda":
                mem_used = torch.cuda.memory_allocated() / (1024**3)
                mem_total = torch.cuda.get_device_properties(0).total_memory / (1024**3)
                gpu_stats.text(f"🎮 GPU Memory: {mem_used:.2f}/{mem_total:.2f} GB")
            
            ram_usage.text(f"💻 RAM Usage: {psutil.Process().memory_info().rss / (1024**3):.2f} GB")
        
        # Build FAISS index
        if all_embeddings:
            status_text.text("⚡ Building search index...")
            all_embeddings = np.vstack(all_embeddings)
            
            start_time = time.time()
            self.index = self.build_faiss_index(all_embeddings)
            build_time = time.time() - start_time
            
            st.session_state.index_built = True
            status_text.text(f"✅ Index built in {build_time:.1f}s | {all_embeddings.shape[0]} vectors")
            st.toast(f"Index ready! Search time: ~1-5ms per query", icon="🚀")
        else:
            st.warning("⚠️ No valid images found in dataset to compute embeddings.")
        
        progress_bar.empty()
        ram_usage.empty()
        gpu_stats.empty()
    
    def search_similar(self, query_image: Image.Image, top_k: int):
        """Find similar images using FAISS with GPU acceleration"""
        # Extract query embedding
        query_tensor = self.transformation_chain(query_image).unsqueeze(0)
        
        if self.USE_HALF_PRECISION and self.device == "cuda":
            query_tensor = query_tensor.half()
        
        with torch.no_grad():
            with torch.cuda.amp.autocast(enabled=self.USE_HALF_PRECISION):
                query_embed = self.model(query_tensor.to(self.device)).last_hidden_state[:, 0].cpu().numpy()
        
        # Normalize and search
        faiss.normalize_L2(query_embed)
        
        # Time the search
        start_time = time.perf_counter()
        distances, indices = self.index.search(query_embed.astype('float32'), top_k)
        search_time = (time.perf_counter() - start_time) * 1000  # in ms
        
        # Convert to similarity scores (0-1 range)
        scores = (1 + distances[0]) / 2  # Convert from [-1,1] to [0,1]
        
        # Retrieve results
        paths = [st.session_state.candidate_paths[i] for i in indices[0]]
        labels = [st.session_state.candidate_labels[i] for i in indices[0]]
        
        return paths, labels, scores.tolist(), search_time
    
    def display_results(self, paths: list, labels: list, scores: list, search_time: float):
        """Display results in an optimized grid layout"""
        if not paths:
            st.warning("⚠️ No similar images found")
            return
            
        st.subheader(f"🔍 Found {len(paths)} similar images in {search_time:.2f}ms")
        cols = st.columns(self.COLUMNS)
        
        # Pre-calculate positions to avoid modulus in loop
        positions = [(i % self.COLUMNS, i // self.COLUMNS) for i in range(len(paths))]
        
        # Use a single container for the grid
        grid_container = st.container()
        
        with grid_container:
            for idx, (path, label, score) in enumerate(zip(paths, labels, scores)):
                col_idx, row_idx = positions[idx]
                with cols[col_idx]:
                    try:
                        img = Image.open(path)
                        width = min(self.MAX_IMAGE_WIDTH, img.width)
                        st.image(img, 
                                caption=f"{self.class_names[label]} ({score:.3f})", 
                                width=width,
                                use_column_width=True)
                        
                        if st.button(f"🔍 Use as Query #{idx+1}", key=f"query_{idx}", use_container_width=True):
                            st.session_state.query_image = img
                            st.session_state.last_results = None
                            st.rerun()
                    except Exception as e:
                        st.error(f"⚠️ Error loading image: {e}")
    
    def run(self):
        """Main application flow"""
        st.set_page_config(layout="wide", page_title="Image Similarity Search")
        st.title("🚀 GPU-Powered Image Similarity Search (Windows Optimized)")
        
        with st.expander("💡 How to use", expanded=True):
            st.markdown("""
            1. **Build Index**: First run will automatically index your images
            2. **Upload**: Select a query image to find similar ones
            3. **Refine**: Click any result to use it as the next query
            """)
        
        self.setup_sidebar()
        self.setup_model()
        
        # Build index if needed
        if st.session_state.get('clear_embeddings') or not st.session_state.get('index_built'):
            self.generate_embeddings()
            st.session_state.clear_embeddings = False
        
        # Query interface
        if st.session_state.query_image is None:
            col1, col2 = st.columns([1, 3])
            with col1:
                st.subheader("Upload Query Image")
                uploaded_file = st.file_uploader(
                    " ",  # Empty label for cleaner UI
                    type=["jpg", "jpeg", "png", "bmp", "webp"],
                    label_visibility="collapsed"
                )
                if uploaded_file:
                    st.session_state.query_image = Image.open(uploaded_file)
                    st.rerun()
            with col2:
                if st.session_state.index_built:
                    st.subheader("System Status")
                    st.info(f"✅ Index ready with {st.session_state.total_images} images")
                    st.code("Click 'Build New Index' in sidebar to refresh", language="bash")
                else:
                    st.warning("Index not built yet - upload an image to start")
        else:
            # Display query image
            st.subheader("Current Query")
            col1, col2 = st.columns([3, 1])
            with col1:
                st.image(
                    st.session_state.query_image, 
                    caption="🔍 Query Image", 
                    width=min(self.MAX_IMAGE_WIDTH, st.session_state.query_image.width)
                )
            with col2:
                if st.button("❌ Clear Query", use_container_width=True):
                    st.session_state.query_image = None
                    st.session_state.last_results = None
                    st.rerun()
                if st.button("🔄 New Search with Same", use_container_width=True):
                    st.session_state.last_results = None
                    st.rerun()
            
            # Perform search
            if st.session_state.last_results is None:
                with st.spinner("🔍 Searching similar images..."):
                    try:
                        paths, labels, scores, search_time = self.search_similar(
                            st.session_state.query_image,
                            self.TOP_K
                        )
                        st.session_state.last_results = (paths, labels, scores, search_time)
                    except Exception as e:
                        st.error(f"Search failed: {str(e)}")
            
            # Display results
            if st.session_state.last_results:
                self.display_results(*st.session_state.last_results)

if __name__ == "__main__":
    app = ImageSimilaritySearch()
    app.run()