import gradio as gr
import torch
from torch_geometric.data import Data
from torch_geometric.utils import from_networkx
import networkx as nx

# Usamos el modelo GCN previamente entrenado (podrías cambiar por GAT si lo prefieres)
model.eval()

def predict_mutagenicity():
    # Creamos un grafo de prueba simple (3 nodos conectados)
    G = nx.Graph()
    G.add_edges_from([(0, 1), (1, 2)])
    nx.set_node_attributes(G, {i: [1, 0, 0, 1, 0, 1, 0] for i in G.nodes}, "x")  # vector ficticio

    # Convertimos a objeto PyG
    pyg_data = from_networkx(G)
    pyg_data.x = torch.tensor(list(nx.get_node_attributes(G, 'x').values()), dtype=torch.float)
    pyg_data.edge_index = pyg_data.edge_index
    pyg_data.batch = torch.tensor([0] * pyg_data.num_nodes)

    pyg_data = pyg_data.to(device)
    with torch.no_grad():
        out = model(pyg_data.x, pyg_data.edge_index, pyg_data.batch)
        pred = out.argmax(dim=1).item()

    return "Mutagénico" if pred == 1 else "No mutagénico"

gr.Interface(fn=predict_mutagenicity, inputs=[], outputs="text",
             title="Clasificador de Moléculas con GNN",
             description="Demo simple de GCN sobre grafos moleculares (MUTAG)").launch()