import gradio as gr
import numpy as np
import h5py
import faiss
import json
from sentence_transformers import SentenceTransformer
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

def load_data():
    try:
        with h5py.File('patent_embeddings.h5', 'r') as f:
            embeddings = f['embeddings'][:]
            patent_numbers = f['patent_numbers'][:]
        
        metadata = {}
        texts = []
        with open('patent_metadata.jsonl', 'r') as f:
            for line in f:
                data = json.loads(line)
                metadata[data['patent_number']] = data
                texts.append(data['text'])
        
        print(f"Embedding shape: {embeddings.shape}")
        print(f"Number of patent numbers: {len(patent_numbers)}")
        print(f"Number of metadata entries: {len(metadata)}")
        
        return embeddings, patent_numbers, metadata, texts
    except FileNotFoundError as e:
        print(f"Error: Could not find file. {e}")
        raise
    except Exception as e:
        print(f"An unexpected error occurred while loading data: {e}")
        raise

embeddings, patent_numbers, metadata, texts = load_data()

# Normalize embeddings for cosine similarity
embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)

# Create FAISS index for cosine similarity
index = faiss.IndexFlatIP(embeddings.shape[1])
index.add(embeddings)

# Load BERT model for encoding search queries
model = SentenceTransformer('all-mpnet-base-v2')

# Create TF-IDF vectorizer
tfidf_vectorizer = TfidfVectorizer(stop_words='english')
tfidf_matrix = tfidf_vectorizer.fit_transform(texts)

def hybrid_search(query, top_k=5):
    print(f"Searching for: {query}")
    
    # Encode the query using the transformer model
    query_embedding = model.encode([query])[0]
    query_embedding = query_embedding / np.linalg.norm(query_embedding)
    
    # Perform semantic similarity search
    semantic_distances, semantic_indices = index.search(np.array([query_embedding]), top_k * 2)
    
    # Perform TF-IDF based search
    query_tfidf = tfidf_vectorizer.transform([query])
    tfidf_similarities = cosine_similarity(query_tfidf, tfidf_matrix).flatten()
    tfidf_indices = tfidf_similarities.argsort()[-top_k * 2:][::-1]
    
    # Combine and rank results
    combined_results = {}
    for i, idx in enumerate(semantic_indices[0]):
        patent_number = patent_numbers[idx].decode('utf-8')
        combined_results[patent_number] = semantic_distances[0][i]
    
    for idx in tfidf_indices:
        patent_number = patent_numbers[idx].decode('utf-8')
        if patent_number in combined_results:
            combined_results[patent_number] += tfidf_similarities[idx]
        else:
            combined_results[patent_number] = tfidf_similarities[idx]
    
    # Sort and get top results
    top_results = sorted(combined_results.items(), key=lambda x: x[1], reverse=True)[:top_k]
    
    results = []
    for patent_number, score in top_results:
        if patent_number not in metadata:
            print(f"Warning: Patent number {patent_number} not found in metadata")
            continue
        patent_data = metadata[patent_number]
        result = f"Patent Number: {patent_number}\n"
        text = patent_data.get('text', 'No text available')
        result += f"Text: {text[:200]}...\n"
        result += f"Combined Score: {score:.4f}\n\n"
        results.append(result)
    
    return "\n".join(results)

# Create Gradio interface
iface = gr.Interface(
    fn=hybrid_search,
    inputs=gr.Textbox(lines=2, placeholder="Enter your search query here..."),
    outputs=gr.Textbox(lines=10, label="Search Results"),
    title="Patent Similarity Search",
    description="Enter a query to find similar patents based on their content."
)

if __name__ == "__main__":
    iface.launch()