import os
import streamlit as st
from pinecone import Pinecone
from sentence_transformers import SentenceTransformer

# Title of the Streamlit App
st.title("Pinecone Query Search on 'pubmed-splade' Index")

# Initialize Pinecone globally
index = None

# Function to initialize Pinecone
def initialize_pinecone():
    api_key = os.getenv('PINECONE_API_KEY')  # Get Pinecone API key from environment variable
    if api_key:
        # Initialize Pinecone client using the new class instance method
        pc = Pinecone(api_key=api_key)
        return pc
    else:
        st.error("Pinecone API key not found! Please set the PINECONE_API_KEY environment variable.")
        return None

# Function to connect to the 'pubmed-splade' index
def connect_to_index(pc):
    index_name = 'pubmed-splade'  # Hardcoded index name
    # Connect to the 'pubmed-splade' index
    if index_name in pc.list_indexes().names():
        st.info(f"Successfully connected to index '{index_name}'")
        index = pc.Index(index_name)
        return index
    else:
        st.error(f"Index '{index_name}' not found!")
        return None

# Function to encode query using sentence transformers model
def encode_query(model, query_text):
    return model.encode(query_text).tolist()

# Initialize Pinecone
pc = initialize_pinecone()

# If Pinecone initialized successfully, proceed with index management
if pc:
    # Connect directly to 'pubmed-splade' index
    index = connect_to_index(pc)

    # Model for query encoding
    model = SentenceTransformer('msmarco-bert-base-dot-v5')

    # Query input
    query_text = st.text_input("Enter a Query to Search", "Can clinicians use the PHQ-9 to assess depression?")
    
    # Button to encode query and search the Pinecone index
    if st.button("Search Query"):
        if query_text and index:
            dense_vector = encode_query(model, query_text)
            st.write(f"Encoded Query Vector: {dense_vector}")
            
            # Search the index (sparse values can be added here as well)
            results = index.query(
                vector=dense_vector,
                top_k=5,
                include_metadata=True
            )
            
            st.write("Search Results:")
            for match in results.matches:
                st.write(f"ID: {match.id}, Score: {match.score}, Metadata: {match.metadata}")
        else:
            st.error("Please enter a query and ensure the index is initialized.")