import streamlit as st
import tensorflow as tf
import tensorflow_gnn as tfgnn
from tensorflow_gnn.models import mt_albis
import networkx as nx
import matplotlib.pyplot as plt

# Define the model function
def model_fn(graph_tensor_spec: tfgnn.GraphTensorSpec):
    graph = inputs = tf.keras.Input(type_spec=graph_tensor_spec)
    
    # Encode input features
    graph = tfgnn.keras.layers.MapFeatures(
        node_sets_fn=lambda node_set, node_set_name: tf.keras.layers.Dense(64)(node_set['features'])
    )(graph)
    
    # Message passing layers
    for _ in range(2):
        graph = mt_albis.MtAlbisGraphUpdate(
            units=128, message_dim=64,
            attention_type="none", simple_conv_reduce_type="mean",
            normalization_type="layer", next_state_type="residual",
            state_dropout_rate=0.2, l2_regularization=1e-5,
        )(graph)
    
    return tf.keras.Model(inputs, graph)

# Function to create a sample graph
def create_sample_graph():
    num_nodes = 10
    num_edges = 15
    
    graph = nx.gnm_random_graph(num_nodes, num_edges, directed=True)
    
    # Create a GraphTensor
    node_features = tf.random.normal((num_nodes, 64))
    edge_features = tf.random.normal((num_edges, 32))
    
    graph_tensor = tfgnn.GraphTensor.from_pieces(
        node_sets={
            "papers": tfgnn.NodeSet.from_fields(
                sizes=[num_nodes],
                features={"features": node_features}
            )
        },
        edge_sets={
            "cites": tfgnn.EdgeSet.from_fields(
                sizes=[num_edges],
                adjacency=tfgnn.Adjacency.from_indices(
                    source=("papers", tf.cast(list(e[0] for e in graph.edges()), tf.int32)),
                    target=("papers", tf.cast(list(e[1] for e in graph.edges()), tf.int32))
                ),
                features={"features": edge_features}
            )
        }
    )
    
    return graph, graph_tensor

# Streamlit app
def main():
    st.title("Graph Neural Network Architecture Visualization")
    
    # Create sample graph
    nx_graph, graph_tensor = create_sample_graph()
    
    # Create and compile the model
    model = model_fn(graph_tensor.spec)
    model.compile(optimizer='adam', loss='binary_crossentropy')
    
    # Display model summary
    st.subheader("Model Summary")
    model.summary(print_fn=lambda x: st.text(x))
    
    # Visualize the graph
    st.subheader("Sample Graph Visualization")
    fig, ax = plt.subplots(figsize=(10, 8))
    pos = nx.spring_layout(nx_graph)
    nx.draw(nx_graph, pos, with_labels=True, node_color='lightblue', 
            node_size=500, arrowsize=20, ax=ax)
    st.pyplot(fig)
    
    # Display graph tensor info
    st.subheader("Graph Tensor Information")
    st.text(f"Number of nodes: {graph_tensor.node_sets['papers'].total_size}")
    st.text(f"Number of edges: {graph_tensor.edge_sets['cites'].total_size}")
    st.text(f"Node feature shape: {graph_tensor.node_sets['papers']['features'].shape}")
    st.text(f"Edge feature shape: {graph_tensor.edge_sets['cites']['features'].shape}")

if __name__ == "__main__":
    main()