import streamlit as st
from streamlit_flow import streamlit_flow
from streamlit_flow.elements import StreamlitFlowNode, StreamlitFlowEdge
from streamlit_flow.state import StreamlitFlowState
from streamlit_flow.layouts import LayeredLayout
import base64

class FlowChart:
    def __init__(self):
        # Initialize the flowchart state in session state if it doesn't exist
        if "flowchart_state" not in st.session_state:
            st.session_state.flowchart_state = self.create_default_state()

    def create_default_state(self):
        """Initialize default nodes and edges."""
        
        nodes = [
            # Source
            StreamlitFlowNode("bundestag",(0, 0), {"content": "Bundestag"}, "input", "right"),
            StreamlitFlowNode("faz", (0, 0), {'content': "FAZ"}, 'input', 'right'),
            StreamlitFlowNode("taz", (0, 0), {"content": "taz"}, "input", "right"),
            StreamlitFlowNode("spiegel", (0, 0), {"content": "Spiegel"}, "input", "right"),
            StreamlitFlowNode("sz", (0, 0), {"content": "Süddeutsche Zeitung"}, "input", "right"),
            StreamlitFlowNode("bild", (0, 0), {"content": "Bild"}, "input", "right"),
            StreamlitFlowNode("zeit", (0, 0), {"content": "zeit"}, "input", "right"),
            StreamlitFlowNode("nzz", (0, 0), {"content": "Neu Zürcher Zeitung"}, "input", "right"),

            # corpus       
            StreamlitFlowNode("corpus",(0, 0), {'content': "corpus"}, 'default', 'right', 'left'),
            # BERTopic
            StreamlitFlowNode("bertopic",(0, 0), {'content': """### BERToptic"""}, 'default', 'right', 'left'),   
            # spacy
            StreamlitFlowNode("spacy",(0, 0), {'content': """### spacy"""}, 'default', 'right', 'left'), 
            # inception
            StreamlitFlowNode("inception",(0, 0), {'content': """### INCEpTION"""}, 'default', 'right', 'left'), 
            # inception
            StreamlitFlowNode("causalbert",(0, 0), {'content': """### CausalBERT"""}, 'default', 'right', 'left'),     

        ]

        edges = [
            StreamlitFlowEdge("Bundestag-corpus", "bundestag", "corpus", animated=True),
            StreamlitFlowEdge("FAZ-corpus", "faz", "corpus", animated=True),
            StreamlitFlowEdge("taz-corpus", "taz", "corpus", animated=True),
            StreamlitFlowEdge("spiegel-corpus", "spiegel", "corpus", animated=True),
            StreamlitFlowEdge("sz-corpus", "sz", "corpus", animated=True),
            StreamlitFlowEdge("bild-corpus", "bild", "corpus", animated=True),
            StreamlitFlowEdge("zeit-corpus", "zeit", "corpus", animated=True),
            StreamlitFlowEdge("nzz-corpus", "nzz", "corpus", animated=True),
            StreamlitFlowEdge("corpus-BERTopic", "corpus", "bertopic", animated=True),
            StreamlitFlowEdge("corpus-spacy", "corpus", "spacy", animated=True),
            StreamlitFlowEdge("spacy-inception", "spacy", "inception", animated=True),
            StreamlitFlowEdge("inception-causalbert", "inception", "causalbert", animated=True)
        ]
        return StreamlitFlowState(nodes, edges)

    def render(self):
        """Render the flowchart with TreeLayout and restricted interactions."""
        st.session_state.flowchart_state = streamlit_flow(
            "markdown_node_flow",
            st.session_state.flowchart_state,
            layout=LayeredLayout(direction='right'),
            fit_view=True,
            show_minimap=False,
            show_controls=True,
            hide_watermark=True,
            allow_new_edges=False,
            enable_node_menu=False,
            enable_edge_menu=False,
            enable_pane_menu=False
        )