import streamlit as st
import graphviz as gv
import json
import os

FILE_NAME = 'saved_data.json'

# Function to create the Graphviz graph
def create_graph():
    g = gv.Digraph('G', engine='dot', format='png')
    
    # Create the first box
    box1_label = 'Source Sequence: Question{s1,s2,s3}\nContext: {sx,sy,sz}\nAnswer: {s1,s2,s3}'
    g.node('box1', label=box1_label, shape='box', style='rounded')
    
    # Create the second box
    box2_label = 'Target Sequence: The answer to the question given the context is yes.'
    g.node('box2', label=box2_label, shape='box', style='rounded')
    
    # Add the line connecting the two boxes
    g.edge('box1', 'box2')

    return g

def save_data(data):
    with open(FILE_NAME, 'w') as f:
        json.dump(data, f)

def load_data():
    if not os.path.exists(FILE_NAME):
        return {}
    with open(FILE_NAME, 'r') as f:
        return json.load(f)

# Create the graph
graph = create_graph()

# Streamlit app
st.title("In Context Learning - Prompt Targeting QA Pattern")
st.subheader("The Question / Answer pattern below can be used in concert with a LLM to do real time in context learning using general intelligence.")
st.graphviz_chart(graph)

data = load_data()

# Input fields
st.header("Enter your data")
question = st.text_input("Question:")
context = st.text_input("Context:")
answer = st.text_input("Answer:")
target_sequence = st.text_input("Target Sequence:")

if st.button("Save"):
    if question and context and answer and target_sequence:
        data["question"] = question
        data["context"] = context
        data["answer"] = answer
        data["target_sequence"] = target_sequence
        save_data(data)
        st.success("Data saved successfully.")
    else:
        st.error("Please fill in all fields.")

st.header("Saved data")
st.write(data)