from matplotlib import pyplot as plt
import streamlit as st
import hypernetx as hnx
from utils.hypergraph_drawer import draw_hypergraph, pyplot_fig_to_buffer

def hypergraph_visualization_component(hypergraph:hnx.Hypergraph, 
                         draw_width:int, 
                         draw_height:int):
    """
    Visualize a hypergraph using Streamlit.
    
    Parameters:
    hypergraph (dict): A dictionary representing the hypergraph, where keys are nodes and values are lists of connected nodes.
    """
    st.header("超图可视化")

    draw_dual = st.toggle("对偶翻转", value=False, key="draw_dual")
    euler_first = st.toggle("先看Euler Diagram", value=True, key="euler_first")
    if draw_dual:
        hypergraph = hypergraph.dual()
    def euler_visualization(hypergraph):
        st.subheader("Euler Diagram 可视化")
        col1, col2, col3, col4 = st.columns(4)
        with col1:
            fill_edges  = st.toggle("填充超边", value=True, key="fill_edges")
        with col2:
            toplexes    = st.toggle("只保留未被包含的边（toplexes）", value=False, key="toplexes")
        with col3:
            with_edge_labels = st.toggle("不看边名", value=False, key="with_edge_labels")
            

        col1, col2 = st.columns(2)
        captions = ["超图", "对偶超图"]

        if 'poses' not in st.session_state:
            st.session_state.poses = [None, None]
        if draw_dual:
            captions = ["对偶超图", "超图"]
            # st.session_state.poses = [st.session_state.poses[1], st.session_state.poses[0]]

        with col1:
            hypergraph_image, st.session_state.poses[0] = draw_hypergraph(hypergraph, draw_dual=False, 
                                            fill_edges=fill_edges, toplexes=toplexes,
                                            with_edge_labels=with_edge_labels
                                            # pos=st.session_state.poses[0]
                                            ,draw_width=draw_width, draw_height=draw_height
                                            )
            st.image(hypergraph_image, caption = captions[0], use_container_width=True)

        with col2:
            hypergraph_image, st.session_state.poses[1] = draw_hypergraph(hypergraph, draw_dual=True, 
                                            fill_edges=fill_edges, toplexes=toplexes,
                                            with_edge_labels=with_edge_labels
                                            # pos=st.session_state.poses[1]
                                            ,draw_width=draw_width, draw_height=draw_height
                                            )
            st.image(hypergraph_image, caption =captions[1], use_container_width=True)

    def new_visualization(hypergraph):
        st.subheader("新型可视化")
        col1, col2 = st.columns(2)
        with col1:
            fig, ax = plt.subplots(figsize=(draw_width, draw_height))
            hnx.draw_bipartite_using_euler(hypergraph)
            st.image(pyplot_fig_to_buffer(plt.gcf()), caption="双列二分图可视化", use_container_width=True)
        with col2:
            fig, ax = plt.subplots(figsize=(draw_width, draw_height))
            hnx.draw_incidence_upset(hypergraph)
            st.image(pyplot_fig_to_buffer(plt.gcf()), caption="Incidence/UpSet 可视化", use_container_width=True)

    if euler_first:
        euler_visualization(hypergraph)
        new_visualization(hypergraph)
    else:
        new_visualization(hypergraph)
        euler_visualization(hypergraph)