# app.py

import streamlit as st
st.set_page_config(page_title = 'Climate Policy Vulnerabilities RAG App', 
                   initial_sidebar_state='expanded', layout="wide") 

import os
import re
import numpy as np
import pandas as pd
from langchain.docstore.document import Document
from openai import OpenAI
import utils  # Import the utils module

openai_api_key = os.environ["OPENAI_API_KEY"]

# Initialize OpenAI client
client = utils.get_openai_client()

# Get Pinecone API key from environment variables
pinecone_api_key = os.getenv("PINECONE_API_KEY")

# Initialize vector store and embeddings
vector_store, embeddings = utils.initialize_vector_store(
    pinecone_api_key, index_name="cpv-full-southern-africa-test"
)

# Prompt template
prompt_template = (
    """
    Answer the given question in an academic tone using the following documents. 
    Provide example quotes and citations using extracted text from the documents. 
    Use facts and numbers from the documents in your answer. 
    ALWAYS include citations for information used from documents at the end of each applicable sentence \
    using the format: '[ref. #]', where '[ref. #]' is included in the text provided at the start of each document \
    (demarcated by the pattern '- &&& [ref. #] document_name &&&:')'.
    Do not include page numbers in the citations.
    Do not include a title, overall introduction or conclusion in your response.
    If no relevant information to answer the question is present in the documents, just say you don't have enough information to answer.
    """
)

# Dropdown options
country_options = [
    'All Countries', 'Angola', 'Botswana', 'Lesotho', 'Kenya', 'Malawi',
    'Mozambique', 'Namibia', 'Rwanda', 'South Africa', 'Zambia', 'Zimbabwe'
]

vulnerability_options = [
    'All Categories', 'Agricultural communities', 'Children', 'Coastal communities',
    'Ethnic, racial or other minorities', 'Fishery communities', 'Informal sector workers',
    'Members of indigenous and local communities', 'Migrants and displaced persons',
    'Older persons', 'Persons living in poverty', 'Persons with disabilities',
    'Persons with pre-existing health conditions', 'Residents of drought-prone regions',
    'Rural populations', 'Sexual minorities (LGBTQI+)', 'Urban populations',
    'Women and other genders', 'Other'
]

examples = [
    "-",
    "What specific initiatives are presented in the context to address the needs of groups such as women and children to the effects climate change?",
    "In addition to gender, children, and youth, is there any mention of other groups facing disproportional impacts from climate change due to their geographic location, socio-economic status, age, gender, health, and occupation?"
]

# Sidebar Filters
with st.sidebar:
    country = st.multiselect('Filter by country:', country_options)
    vulnerabilities_cat = st.multiselect('Filter by vulnerabilities category:', vulnerability_options)
    with st.expander("ℹ️ - About filters", expanded=False):
        st.markdown(
            """
            * *These selections will filter the data matched against your query.*
            * *For a comparative analysis of multiple countries or vulnerability categories, select the items you require or select **'All Countries'** or **'All Categories'***.
            * *Be careful in using the vulnerabilities category filter, as many of the categories are not well represented in the documents. Therefore, this will severely limit the data available for analysis.*
            """
        )

# Main Window Title
st.markdown("<h2 style='text-align: center;'> Climate Policy Documents: Vulnerabilities Analysis Q&A </h2>", unsafe_allow_html=True)
st.write(' ')

# Main Window Instructions
with st.expander("ℹ️ - About this app", expanded=False):
    st.write(
        """
        This tool seeks to provide an interface for querying national climate policy documents (NDCs, LTS etc.). The current version is powered by chatGPT (3.5). The document store is limited to 10 Southern African countries (Angola, Botswana, Eswatini, Lesotho, Malawi, Mozambique, Namibia, South Africa, Zambia, Zimbabwe), as well as Kenya and Rwanda. The intended use case is to allow users to interact with the documents and obtain valuable insights on various vulnerable groups affected by climate change.
        
        **DISCLAIMER:** *This prototype tool based on LLMs (Language Models) is provided "as is" for experimental and exploratory purposes only, and should not be used for critical or production applications. Users are advised that the tool may contain errors, bugs, or limitations and should be used with caution and awareness of potential risks, and the developers make no warranties or guarantees regarding its performance, reliability, or suitability for any specific purpose.*
        """
    )
    selected_example = st.radio("Example questions", examples)
    st.write(
        """
        You can request comparative analyses between countries by filtering by country, and using more advanced prompts. For example:
        
        *Provide a comparative analysis between Angola and Kenya with regard to specific initiatives presented in the context to address the needs of groups such women and children to the effects climate change.*
    
        Make sure your filters match the countries you have specified for the analysis!
        """
    )

# Model Selection (Currently fixed to chatGPT)
model_sel = "chatGPT"

# Prompt Logic
if model_sel == "chatGPT":
    model_name = "gpt-3.5-turbo"

# Input Text Area
if selected_example == "-":
    text = st.text_area('Enter your question in the text box below using natural language or select an example from above:')
else:
    text = st.text_area('Enter your question in the text box below using natural language or select an example from above:', value=selected_example)

# Submit Button
if st.button('Submit'):
    st.markdown("----")
    st.markdown('**RESPONSE:**')
    res_box = st.empty()
    
    # Fetch documents based on user input and filters
    docs = utils.get_docs(vector_store, embeddings, text, country=country, vulnerability_cat=vulnerabilities_cat)
    
    # Construct the prompt
    prompt = utils.get_prompt(prompt_template, docs, text)
    
    # Run the query and get references
    references = utils.run_query(client, prompt, docs, res_box)
    
    # Display references
    st.markdown("----")
    st.markdown('**REFERENCES:**')
    with st.expander("ℹ️ - About References", expanded=False):
        st.write(
            """
            *References are based on text automatically extracted from climate policy documents.
            These extracts may contain non-legible characters or disjointed text as an artifact of the extraction procedure.*
            """,
            unsafe_allow_html=True
        )
    
    st.markdown(references, unsafe_allow_html=True)