import gradio as gr
import numpy as np
import h5py
import faiss
import json
from transformers import AutoTokenizer, AutoModel
from sentence_transformers import SentenceTransformer, models
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import re
from collections import Counter
import spacy
import joblib

# Load Spacy model for advanced NLP
try:
    nlp = spacy.load("en_core_web_sm")
except IOError:
    print("Downloading spacy model...")
    spacy.cli.download("en_core_web_sm")
    nlp = spacy.load("en_core_web_sm")

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()

# Load BERT model for encoding search queries
try:
    tokenizer = AutoTokenizer.from_pretrained('anferico/bert-for-patents')
    bert_model = AutoModel.from_pretrained('anferico/bert-for-patents')
    word_embedding_model = models.Transformer(model_name='anferico/bert-for-patents', tokenizer=tokenizer, model=bert_model)
    pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension())
    model = SentenceTransformer(modules=[word_embedding_model, pooling_model])
except Exception as e:
    print(f"Error loading anferico/bert-for-patents: {e}")
    print("Falling back to a general-purpose model.")
    model = SentenceTransformer('all-MiniLM-L6-v2')

# Check if the embedding dimensions match
if embeddings.shape[1] != model.get_sentence_embedding_dimension():
    print("Embedding dimensions do not match. Rebuilding FAISS index.")
    # Rebuild embeddings using the new model
    embeddings = np.array([model.encode(text) for text in texts])
    embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)

# 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)

# Create and save TF-IDF vectorizer
tfidf_vectorizer = TfidfVectorizer(stop_words='english')
tfidf_matrix = tfidf_vectorizer.fit_transform(texts)
joblib.dump(tfidf_vectorizer, 'tfidf_vectorizer.joblib')

def extract_key_features(text):
    # Use Spacy to extract noun phrases and key phrases
    doc = nlp(text)
    noun_phrases = [chunk.text.lower() for chunk in doc.noun_chunks]
    feature_phrases = [sent.text.lower() for sent in doc.sents if re.search(r'(comprising|including|consisting of)', sent.text, re.IGNORECASE)]
    
    all_features = noun_phrases + feature_phrases
    return list(set(all_features))

def compare_features(query_features, patent_features):
    common_features = set(query_features) & set(patent_features)
    similarity_score = len(common_features) / max(len(query_features), len(patent_features))
    return common_features, similarity_score

def hybrid_search(query, top_k=5):
    print(f"Original query: {query}")
    
    query_features = extract_key_features(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