import os
from huggingface_hub import login
from transformers import pipeline
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
import streamlit as st

# Authenticate with Hugging Face using the API key from environment variables
# If running on Hugging Face Spaces, make sure the API key is stored as a secret.
login(os.environ["HF_API_KEY"])

# Initialize a question-answering model
question_answerer = pipeline("question-answering", model="distilbert-base-cased-distilled-squad")

# Load or create data on Pakistan's economic and population growth trends
documents = [
    {"id": 1, "text": "Pakistan's population growth rate is approximately 2%, making it one of the fastest-growing populations in South Asia."},
    {"id": 2, "text": "The youth population in Pakistan is significant, with over 60% of the population under the age of 30."},
    {"id": 3, "text": "Pakistan's economy relies heavily on agriculture, with about 20% of GDP coming from this sector."},
    {"id": 4, "text": "In recent years, Pakistan has been investing in infrastructure projects, such as the China-Pakistan Economic Corridor (CPEC), to boost economic growth."},
    {"id": 5, "text": "Urbanization is rapidly increasing in Pakistan, with cities like Karachi and Lahore seeing substantial population inflows."},
    {"id": 6, "text": "Remittances from overseas Pakistanis play a critical role in supporting the country's economy."},
    {"id": 7, "text": "Pakistan's literacy rate has improved over the years but remains lower than the regional average."},
    {"id": 8, "text": "The government is focusing on initiatives for digital economy growth, particularly in the technology and freelancing sectors."},
    {"id": 9, "text": "Pakistan’s unemployment rate is a concern, especially among young people entering the job market."},
    {"id": 10, "text": "The fertility rate in Pakistan has been declining but remains above the replacement rate."},
]

# Embed documents for retrieval using SentenceTransformer
embedder = SentenceTransformer('all-MiniLM-L6-v2')
document_embeddings = [embedder.encode(doc['text']) for doc in documents]

# Convert embeddings to a FAISS index for similarity search
index = faiss.IndexFlatL2(384)  # Dimension of embeddings
index.add(np.array(document_embeddings))

# Define the RAG retrieval function
def retrieve_documents(query, top_k=3):
    query_embedding = embedder.encode(query).reshape(1, -1)
    distances, indices = index.search(query_embedding, top_k)
    return [documents[i]['text'] for i in indices[0]]

# Implement the question-answering function with retrieval
def ask_question(question):
    # Retrieve relevant documents
    retrieved_docs = retrieve_documents(question)
    # Combine retrieved documents into a single context
    context = " ".join(retrieved_docs)
    
    # Generate an answer based on retrieved context
    answer = question_answerer(question=question, context=context)
    return answer['answer']

# Streamlit Interface
def streamlit_interface():
    # Set title and description
    st.title("Pakistan Economic and Population Growth Advisor")
    st.write("Ask questions related to Pakistan's economic and population growth. This app uses retrieval-augmented generation to provide answers based on relevant documents about Pakistan.")

    # Input: User enters a question
    question = st.text_input("Ask a question:")

    if question:
        # Get the answer using the RAG system
        answer = ask_question(question)
        # Output the answer
        st.write("Answer:", answer)

if __name__ == "__main__":
    # Run the Streamlit app
    streamlit_interface()