import streamlit as st
import openai

def chunk_text(text, chunk_size=2000):
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunk = text[start:end]
        chunks.append(chunk)
        start = end
    return chunks

# Streamlit Session State
if 'learning_objectives' not in st.session_state:
    st.session_state.learning_objectives = ""

# Streamlit User Input Form
st.title("Patent Claims Extraction")

# API Key Input
api_key = st.text_input("Enter your OpenAI API Key:", type="password")

# Model Selection Dropdown
model_choice = st.selectbox(
    "Select the model you want to use:",
    ["gpt-3.5-turbo-0301", "gpt-3.5-turbo-0613", "gpt-3.5-turbo", "gpt-4-0314", "gpt-4-0613", "gpt-4"]
)

# Context, Subject, and Level
context = "You are a patent claims identifier and extractor. You will freeform text, identify any claims contained therein that may be patentable. You identify , extract, print such claims, briefly explain why each claim is patentable."
userinput = st.text_input("Input Text:", "Freeform text here!")

# Initialize OpenAI API
if api_key:
    openai.api_key = api_key

# Learning Objectives
st.write("### Patentable Claims:")
# Initialize autogenerated objectives
claims_extraction = ""
# Initialize status placeholder
learning_status_placeholder = st.empty()
disable_button_bool = False
if userinput and api_key and st.button("Extract Claims", key="claims_extraction", disabled=disable_button_bool):
    # Split the user input into chunks
    input_chunks = chunk_text(userinput)
    
    # Initialize a variable to store the extracted claims
    all_extracted_claims = ""

    for chunk in input_chunks:
        # Display status message for the current chunk
        learning_status_placeholder.text(f"Extracting Patentable Claims for chunk {input_chunks.index(chunk) + 1}...")

        # API call to generate objectives for the current chunk
        claims_extraction_response = openai.ChatCompletion.create(
            model=model_choice,
            messages=[
                {"role": "user", "content": f"Extract any patentable claims from the following: \n {chunk}. \n extract each claim. Briefly explain why you extracted this word phrase. Exclude any additional commentary."}
            ]
        )

        # Extract the generated objectives from the API response
        claims_extraction = claims_extraction_response['choices'][0]['message']['content']

        # Append the extracted claims from the current chunk to the overall results
        all_extracted_claims += claims_extraction.strip()

    # Save the generated objectives to session state
    st.session_state.claims_extraction = all_extracted_claims

    # Display generated objectives for all chunks
    learning_status_placeholder.text(f"Patentable Claims Extracted!\n{all_extracted_claims.strip()}")

# Generate Lesson Plan Button
if st.button("Extract Claims") and api_key:
    
    # Construct the prompt as a dictionary
    prompt_dict = {
        "context": context,
        "userinput": userinput,
        "claims_extraction": claims_extraction.strip(),  # Use the claims_extraction variable
        "tasks": [
            {"task": "Extract Claims", "objective": "extract any wordphrases in the text provided that could be considered a patentable claim"},
            {"task": "Extract Every Claim", "objective": "Ensure each and every wordphrase with a claim is evaluated whether or not it is patentable"},
            {"task": "Explain Your Choice", "objective": "Briefly explain why you have retained each claim"},
            {"task": "Check Your Work", "objective": "Check your work to assure you have not forgotten any claims"}
        ],
        "output_format": """Present in a structured format.
        \nClaim: 
        \nExplanation:
        \nNotes:
        
        """
    }

    
    # Convert the dictionary to a string
    prompt_str = str(prompt_dict)

    # API call to generate the lesson plan
    lesson_plan_response = openai.ChatCompletion.create(
        model=model_choice,
        messages=[
            {"role": "user", "content": f"Create a lesson plan based on the following parameters: {prompt_str}"}
        ]
    )

    # Display status message
    lesson_plan=st.text("Extracting Patentable Claims...")

    # Extract and display
    assistant_reply = claims_extraction_response['choices'][0]['message']['content']
    claims_extraction=st.text(assistant_reply.strip())

# Citation 
st.markdown("<sub>This app was created by [Taylor Ennen](https://github.com/taylor-ennen/GPT-Streamlit-MVP) & [Tonic](https://huggingface.co./tonic)</sub>", unsafe_allow_html=True)