utils.py

"""
Utility functions for the Multi-Method RAG System.

Directory Layout:
/data/         # Original PDFs, HTML
/embeddings/   # FAISS, Chroma, DPR vector stores
/graph/        # Graph database files
/metadata/     # Image metadata (SQLite or MongoDB)
"""

import os
import json
import pickle
import sqlite3
import base64
from pathlib import Path
from typing import List, Dict, Tuple, Optional, Any, Union
from dataclasses import dataclass
import logging

import pymupdf4llm
import pymupdf
import numpy as np
import pandas as pd
from PIL import Image
import requests
from bs4 import BeautifulSoup

# Vector stores and search
import faiss
import chromadb
from rank_bm25 import BM25Okapi
import networkx as nx

# ML models
from openai import OpenAI
from sentence_transformers import SentenceTransformer, CrossEncoder
import torch
# import clip

# Text processing
from sklearn.feature_extraction.text import TfidfVectorizer
import tiktoken

from config import *

logger = logging.getLogger(__name__)

@dataclass
class DocumentChunk:
    """Data structure for document chunks."""
    text: str
    metadata: Dict[str, Any]
    chunk_id: str
    embedding: Optional[np.ndarray] = None
    
@dataclass
class ImageData:
    """Data structure for image metadata."""
    image_path: str
    image_id: str
    classification: Optional[str] = None
    embedding: Optional[np.ndarray] = None
    metadata: Optional[Dict[str, Any]] = None

class DocumentLoader:
    """Load and extract text from various document formats."""
    
    def __init__(self):
        self.client = OpenAI(api_key=OPENAI_API_KEY)
        validate_api_key()
    
    def load_pdf_documents(self, pdf_paths: List[Union[str, Path]]) -> List[Dict[str, Any]]:
        """Load text from PDF files using pymupdf4llm."""
        documents = []
        
        for pdf_path in pdf_paths:
            try:
                pdf_path = Path(pdf_path)
                logger.info(f"Loading PDF: {pdf_path}")
                
                # Extract text using pymupdf4llm
                text = pymupdf4llm.to_markdown(str(pdf_path))
                
                # Extract images if present
                images = self._extract_pdf_images(pdf_path)
                
                doc = {
                    'text': text,
                    'source': str(pdf_path.name),
                    'path': str(pdf_path),
                    'type': 'pdf',
                    'images': images,
                    'metadata': {
                        'file_size': pdf_path.stat().st_size,
                        'modified': pdf_path.stat().st_mtime
                    }
                }
                documents.append(doc)
                
            except Exception as e:
                logger.error(f"Error loading PDF {pdf_path}: {e}")
                continue
                
        return documents
    
    def _extract_pdf_images(self, pdf_path: Path) -> List[Dict[str, Any]]:
        """Extract images from PDF using pymupdf."""
        images = []
        
        try:
            doc = pymupdf.open(str(pdf_path))
            
            for page_num in range(len(doc)):
                page = doc[page_num]
                image_list = page.get_images(full=True)
                
                for img_index, img in enumerate(image_list):
                    try:
                        # Extract image
                        xref = img[0]
                        pix = pymupdf.Pixmap(doc, xref)
                        
                        # Skip if pixmap is invalid or has no colorspace
                        if not pix or pix.colorspace is None:
                            if pix:
                                pix = None
                            continue
                        
                        # Only process images with valid color channels
                        if pix.n - pix.alpha < 4:  # GRAY or RGB
                            image_id = f"{pdf_path.stem}_p{page_num}_img{img_index}"
                            image_path = IMAGES_DIR / f"{image_id}.png"
                            
                            # Convert to RGB if grayscale or other formats
                            if pix.n == 1:  # Grayscale
                                rgb_pix = pymupdf.Pixmap(pymupdf.csRGB, pix)
                                pix = None  # Clean up original
                                pix = rgb_pix
                            elif pix.n == 4 and pix.alpha == 0:  # CMYK
                                rgb_pix = pymupdf.Pixmap(pymupdf.csRGB, pix)
                                pix = None  # Clean up original
                                pix = rgb_pix
                            
                            # Save image
                            pix.save(str(image_path))
                            
                            images.append({
                                'image_id': image_id,
                                'image_path': str(image_path),
                                'page': page_num,
                                'source': str(pdf_path.name)
                            })
                        
                        pix = None
                        
                    except Exception as e:
                        logger.warning(f"Error extracting image {img_index} from page {page_num}: {e}")
                        if 'pix' in locals() and pix:
                            pix = None
                        continue
            
            doc.close()
            
        except Exception as e:
            logger.error(f"Error extracting images from {pdf_path}: {e}")
        
        return images
    
    def load_html_documents(self, html_sources: List[Dict[str, str]]) -> List[Dict[str, Any]]:
        """Load text from HTML sources."""
        documents = []
        
        for source in html_sources:
            try:
                logger.info(f"Loading HTML: {source.get('title', source['url'])}")
                
                # Fetch HTML content
                response = requests.get(source['url'], timeout=30)
                response.raise_for_status()
                
                # Parse with BeautifulSoup
                soup = BeautifulSoup(response.text, 'html.parser')
                
                # Extract text
                text = soup.get_text(separator=' ', strip=True)
                
                doc = {
                    'text': text,
                    'source': source.get('title', source['url']),
                    'path': source['url'],
                    'type': 'html',
                    'images': [],
                    'metadata': {
                        'url': source['url'],
                        'title': source.get('title', ''),
                        'year': source.get('year', ''),
                        'category': source.get('category', ''),
                        'format': source.get('format', 'HTML')
                    }
                }
                documents.append(doc)
                
            except Exception as e:
                logger.error(f"Error loading HTML {source['url']}: {e}")
                continue
                
        return documents
    
    def load_text_documents(self, data_dir: Path = DATA_DIR) -> List[Dict[str, Any]]:
        """Load all supported document types from data directory."""
        documents = []
        
        # Load PDFs
        pdf_files = list(data_dir.glob("*.pdf"))
        if pdf_files:
            documents.extend(self.load_pdf_documents(pdf_files))
        
        # Load HTML sources (from config)
        if DEFAULT_HTML_SOURCES:
            documents.extend(self.load_html_documents(DEFAULT_HTML_SOURCES))
        
        logger.info(f"Loaded {len(documents)} documents total")
        return documents

class TextPreprocessor:
    """Preprocess text for different retrieval methods."""
    
    def __init__(self):
        self.encoding = tiktoken.get_encoding("cl100k_base")
    
    def chunk_text_by_tokens(self, text: str, chunk_size: int = CHUNK_SIZE, 
                           overlap: int = CHUNK_OVERLAP) -> List[str]:
        """Split text into chunks by token count."""
        tokens = self.encoding.encode(text)
        chunks = []
        
        start = 0
        while start < len(tokens):
            end = start + chunk_size
            chunk_tokens = tokens[start:end]
            chunk_text = self.encoding.decode(chunk_tokens)
            chunks.append(chunk_text)
            start = end - overlap
            
        return chunks
    
    def chunk_text_by_sections(self, text: str, method: str = "vanilla") -> List[str]:
        """Split text by sections based on method requirements."""
        if method in ["vanilla", "dpr"]:
            return self.chunk_text_by_tokens(text)
        elif method == "bm25":
            # BM25 works better with paragraph-level chunks
            paragraphs = [p.strip() for p in text.split('\n\n') if p.strip()]
            return paragraphs
        elif method == "graph":
            # Graph method uses larger sections
            return self.chunk_text_by_tokens(text, chunk_size=CHUNK_SIZE*2)
        elif method == "context_stuffing":
            # Context stuffing uses full documents
            return [text]
        else:
            return self.chunk_text_by_tokens(text)
    
    def preprocess_for_method(self, documents: List[Dict[str, Any]], 
                            method: str) -> List[DocumentChunk]:
        """Preprocess documents for specific retrieval method."""
        chunks = []
        
        for doc in documents:
            text_chunks = self.chunk_text_by_sections(doc['text'], method)
            
            for i, chunk_text in enumerate(text_chunks):
                chunk_id = f"{doc['source']}_{method}_chunk_{i}"
                
                chunk = DocumentChunk(
                    text=chunk_text,
                    metadata={
                        'source': doc['source'],
                        'path': doc['path'],
                        'type': doc['type'],
                        'chunk_index': i,
                        'method': method,
                        **doc.get('metadata', {})
                    },
                    chunk_id=chunk_id
                )
                chunks.append(chunk)
        
        logger.info(f"Created {len(chunks)} chunks for method '{method}'")
        return chunks

class EmbeddingGenerator:
    """Generate embeddings using various models."""
    
    def __init__(self):
        self.openai_client = OpenAI(api_key=OPENAI_API_KEY)
        self.sentence_transformer = None
        # self.clip_model = None
        # self.clip_preprocess = None
        
    def _get_sentence_transformer(self):
        """Lazy loading of sentence transformer."""
        if self.sentence_transformer is None:
            self.sentence_transformer = SentenceTransformer(SENTENCE_TRANSFORMER_MODEL)
            if DEVICE == "cuda":
                self.sentence_transformer = self.sentence_transformer.to(DEVICE)
        return self.sentence_transformer
    
    # def _get_clip_model(self):
    #     """Lazy loading of CLIP model."""
    #     if self.clip_model is None:
    #         self.clip_model, self.clip_preprocess = clip.load(CLIP_MODEL, device=DEVICE)
    #     return self.clip_model, self.clip_preprocess
    
    def embed_text_openai(self, texts: List[str]) -> np.ndarray:
        """Generate embeddings using OpenAI API."""
        embeddings = []
        
        # Process in batches
        for i in range(0, len(texts), EMBEDDING_BATCH_SIZE):
            batch = texts[i:i + EMBEDDING_BATCH_SIZE]
            
            try:
                response = self.openai_client.embeddings.create(
                    model=OPENAI_EMBEDDING_MODEL,
                    input=batch
                )
                
                batch_embeddings = [data.embedding for data in response.data]
                embeddings.extend(batch_embeddings)
                
            except Exception as e:
                logger.error(f"Error generating OpenAI embeddings: {e}")
                raise
        
        return np.array(embeddings)
    
    def embed_text_sentence_transformer(self, texts: List[str]) -> np.ndarray:
        """Generate embeddings using sentence transformers."""
        model = self._get_sentence_transformer()
        
        try:
            embeddings = model.encode(texts, convert_to_numpy=True, 
                                    show_progress_bar=True, batch_size=32)
            return embeddings
            
        except Exception as e:
            logger.error(f"Error generating sentence transformer embeddings: {e}")
            raise
    
    def embed_image_clip(self, image_paths: List[str]) -> np.ndarray:
        """Generate image embeddings using CLIP."""
        # model, preprocess = self._get_clip_model()
        # embeddings = []
        
        # for image_path in image_paths:
        #     try:
        #         image = preprocess(Image.open(image_path)).unsqueeze(0).to(DEVICE)
        #         
        #         with torch.no_grad():
        #             image_features = model.encode_image(image)
        #             image_features /= image_features.norm(dim=-1, keepdim=True)
        #             
        #         embeddings.append(image_features.cpu().numpy().flatten())
        #         
        #     except Exception as e:
        #         logger.error(f"Error embedding image {image_path}: {e}")
        #         continue
        
        # return np.array(embeddings) if embeddings else np.array([])
        
        # Placeholder for CLIP embeddings
        logger.warning("CLIP embeddings not implemented - returning dummy embeddings")
        return np.random.rand(len(image_paths), 512)

class VectorStoreManager:
    """Manage vector stores for different methods."""
    
    def __init__(self):
        self.embedding_generator = EmbeddingGenerator()
    
    def build_faiss_index(self, chunks: List[DocumentChunk], method: str = "vanilla") -> Tuple[Any, List[Dict]]:
        """Build FAISS index for vanilla or DPR method."""
        
        # Generate embeddings
        texts = [chunk.text for chunk in chunks]
        
        if method == "vanilla":
            embeddings = self.embedding_generator.embed_text_openai(texts)
        elif method == "dpr":
            embeddings = self.embedding_generator.embed_text_sentence_transformer(texts)
        else:
            raise ValueError(f"Unsupported method for FAISS: {method}")
        
        # Build FAISS index
        dimension = embeddings.shape[1]
        index = faiss.IndexFlatIP(dimension)  # Inner product for cosine similarity
        
        # Ensure embeddings are float32 and normalize for cosine similarity
        embeddings = embeddings.astype(np.float32)
        faiss.normalize_L2(embeddings)
        index.add(embeddings)
        
        # Store chunk metadata
        metadata = []
        for i, chunk in enumerate(chunks):
            metadata.append({
                'chunk_id': chunk.chunk_id,
                'text': chunk.text,
                'metadata': chunk.metadata,
                'embedding': embeddings[i].tolist()
            })
        
        logger.info(f"Built FAISS index with {index.ntotal} vectors for method '{method}'")
        return index, metadata
    
    def build_chroma_index(self, chunks: List[DocumentChunk], method: str = "vanilla") -> Any:
        """Build Chroma vector database."""
        
        # Initialize Chroma client
        chroma_client = chromadb.PersistentClient(path=str(CHROMA_PATH / method))
        collection = chroma_client.get_or_create_collection(
            name=f"{method}_collection",
            metadata={"method": method}
        )
        
        # Prepare data for Chroma
        texts = [chunk.text for chunk in chunks]
        ids = [chunk.chunk_id for chunk in chunks]
        metadatas = [chunk.metadata for chunk in chunks]
        
        # Add to collection (Chroma handles embeddings internally)
        collection.add(
            documents=texts,
            ids=ids,
            metadatas=metadatas
        )
        
        logger.info(f"Built Chroma collection with {collection.count()} documents for method '{method}'")
        return collection
    
    def build_bm25_index(self, chunks: List[DocumentChunk]) -> BM25Okapi:
        """Build BM25 index for keyword search."""
        
        # Tokenize texts
        tokenized_corpus = []
        for chunk in chunks:
            tokens = chunk.text.lower().split()
            tokenized_corpus.append(tokens)
        
        # Build BM25 index
        bm25 = BM25Okapi(tokenized_corpus, k1=BM25_K1, b=BM25_B)
        
        logger.info(f"Built BM25 index with {len(tokenized_corpus)} documents")
        return bm25
    
    def build_graph_index(self, chunks: List[DocumentChunk]) -> nx.Graph:
        """Build NetworkX graph for graph-based retrieval."""
        
        # Create graph
        G = nx.Graph()
        
        # Generate embeddings for similarity calculation
        texts = [chunk.text for chunk in chunks]
        embeddings = self.embedding_generator.embed_text_openai(texts)
        
        # Add nodes (convert embeddings to lists for GML serialization)
        for i, chunk in enumerate(chunks):
            G.add_node(chunk.chunk_id, 
                      text=chunk.text,
                      metadata=chunk.metadata,
                      embedding=embeddings[i].tolist())  # Convert to list for serialization
        
        # Add edges based on similarity
        threshold = 0.7  # Similarity threshold
        for i in range(len(chunks)):
            for j in range(i + 1, len(chunks)):
                # Calculate cosine similarity
                sim = np.dot(embeddings[i], embeddings[j]) / (
                    np.linalg.norm(embeddings[i]) * np.linalg.norm(embeddings[j])
                )
                
                if sim > threshold:
                    G.add_edge(chunks[i].chunk_id, chunks[j].chunk_id, 
                              weight=float(sim))
        
        logger.info(f"Built graph with {G.number_of_nodes()} nodes and {G.number_of_edges()} edges")
        return G
    
    def save_index(self, index: Any, metadata: Any, method: str):
        """Save index and metadata to disk."""
        
        if method == "vanilla":
            faiss.write_index(index, str(VANILLA_FAISS_INDEX))
            with open(VANILLA_METADATA, 'wb') as f:
                pickle.dump(metadata, f)
                
        elif method == "dpr":
            faiss.write_index(index, str(DPR_FAISS_INDEX))
            with open(DPR_METADATA, 'wb') as f:
                pickle.dump(metadata, f)
                
        elif method == "bm25":
            with open(BM25_INDEX, 'wb') as f:
                pickle.dump({'index': index, 'texts': metadata}, f)
                
        elif method == "context_stuffing":
            with open(CONTEXT_DOCS, 'wb') as f:
                pickle.dump(metadata, f)
                
        elif method == "graph":
            nx.write_gml(index, str(GRAPH_FILE))
            
        logger.info(f"Saved {method} index to disk")

class ImageProcessor:
    """Process and classify images."""
    
    def __init__(self):
        self.embedding_generator = EmbeddingGenerator()
        self.openai_client = OpenAI(api_key=OPENAI_API_KEY)
        self._init_database()
    
    def _init_database(self):
        """Initialize SQLite database for image metadata."""
        conn = sqlite3.connect(IMAGES_DB)
        cursor = conn.cursor()
        
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS images (
                image_id TEXT PRIMARY KEY,
                image_path TEXT NOT NULL,
                classification TEXT,
                metadata TEXT,
                embedding BLOB,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        conn.commit()
        conn.close()
    
    def classify_image(self, image_path: str) -> str:
        """Classify image using GPT-5 Vision."""
        try:
            # Convert image to base64
            with open(image_path, "rb") as image_file:
                image_b64 = base64.b64encode(image_file.read()).decode()
            
            messages = [{
                "role": "user",
                "content": [
                    {"type": "text", "text": "Classify this image in 1-2 words (e.g., 'machine guard', 'press brake', 'conveyor belt', 'safety sign')."},
                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_b64}", "detail": "low"}}
                ]
            }]
            
            # For GPT-5 vision, temperature must be default (1.0)
            response = self.openai_client.chat.completions.create(
                model=OPENAI_CHAT_MODEL,
                messages=messages,
                max_completion_tokens=50
            )
            
            return response.choices[0].message.content.strip()
            
        except Exception as e:
            logger.error(f"Error classifying image {image_path}: {e}")
            return "unknown"
    
    def should_filter_image(self, image_path: str) -> tuple[bool, str]:
        """
        Check if image should be filtered out based on height and black image criteria.
        
        Args:
            image_path: Path to the image file
            
        Returns:
            Tuple of (should_filter: bool, reason: str)
        """
        try:
            from PIL import Image
            import numpy as np
            
            # Open and analyze the image
            with Image.open(image_path) as img:
                # Convert to RGB if needed
                if img.mode != 'RGB':
                    img = img.convert('RGB')
                
                width, height = img.size
                
                # Filter 1: Height less than 40 pixels
                if height < 40:
                    return True, f"height too small ({height}px)"
                
                # Filter 2: Check if image is mostly black
                img_array = np.array(img)
                mean_brightness = np.mean(img_array)
                
                # If mean brightness is very low (mostly black)
                if mean_brightness < 10:  # Adjust threshold as needed
                    return True, "mostly black image"
                
        except Exception as e:
            logger.warning(f"Error analyzing image {image_path}: {e}")
            # If we can't analyze it, don't filter it out
            return False, "analysis failed"
        
        return False, "passed all filters"
    
    def store_image_metadata(self, image_data: ImageData):
        """Store image metadata in database."""
        conn = sqlite3.connect(IMAGES_DB)
        cursor = conn.cursor()
        
        # Serialize metadata and embedding
        metadata_json = json.dumps(image_data.metadata) if image_data.metadata else None
        embedding_blob = image_data.embedding.tobytes() if image_data.embedding is not None else None
        
        cursor.execute('''
            INSERT OR REPLACE INTO images 
            (image_id, image_path, classification, metadata, embedding)
            VALUES (?, ?, ?, ?, ?)
        ''', (image_data.image_id, image_data.image_path, 
              image_data.classification, metadata_json, embedding_blob))
        
        conn.commit()
        conn.close()
    
    def get_image_metadata(self, image_id: str) -> Optional[ImageData]:
        """Retrieve image metadata from database."""
        conn = sqlite3.connect(IMAGES_DB)
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT image_id, image_path, classification, metadata, embedding
            FROM images WHERE image_id = ?
        ''', (image_id,))
        
        row = cursor.fetchone()
        conn.close()
        
        if row:
            image_id, image_path, classification, metadata_json, embedding_blob = row
            
            metadata = json.loads(metadata_json) if metadata_json else None
            embedding = np.frombuffer(embedding_blob, dtype=np.float32) if embedding_blob else None
            
            return ImageData(
                image_path=image_path,
                image_id=image_id,
                classification=classification,
                embedding=embedding,
                metadata=metadata
            )
        
        return None

def load_text_documents() -> List[Dict[str, Any]]:
    """Convenience function to load all text documents."""
    loader = DocumentLoader()
    return loader.load_text_documents()

def embed_image_clip(image_paths: List[str]) -> np.ndarray:
    """Convenience function to embed images with CLIP."""
    generator = EmbeddingGenerator()
    return generator.embed_image_clip(image_paths)

def store_image_metadata(image_data: ImageData):
    """Convenience function to store image metadata."""
    processor = ImageProcessor()
    processor.store_image_metadata(image_data)

def classify_image(image_path: str) -> str:
    """Convenience function to classify an image."""
    processor = ImageProcessor()
    return processor.classify_image(image_path)