import gradio as gr
import numpy as np
import torch
import torch.nn as nn
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.patches import Circle

class SechGraphConv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.linear = nn.Linear(in_ch, out_ch)
        self.inv_e = 1.0 / np.e
    
    def forward(self, x, edge_index, edge_weight=None):
        N = x.size(0)
        if edge_weight is None:
            edge_weight = torch.ones(edge_index.size(1))
        sech_w = 1.0 / torch.cosh(edge_weight)
        row, col = edge_index
        out = torch.zeros_like(x)
        for i in range(N):
            mask = row == i
            if mask.sum() > 0:
                out[i] = (x[col[mask]] * sech_w[mask].view(-1,1)).sum(0)
        deg = torch.zeros(N)
        for i in range(N):
            deg[i] = (row == i).sum().float()
        deg_inv_sqrt = deg.pow(-0.5)
        deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
        out = deg_inv_sqrt.view(-1,1) * out
        out = self.linear(out)
        return out * (1.0 / torch.cosh(self.inv_e * out))

class SentinelGNN(nn.Module):
    def __init__(self, in_ch, hidden, out_ch, layers=2):
        super().__init__()
        self.convs = nn.ModuleList()
        self.convs.append(SechGraphConv(in_ch, hidden))
        for _ in range(layers-2):
            self.convs.append(SechGraphConv(hidden, hidden))
        self.convs.append(SechGraphConv(hidden, out_ch))
    
    def forward(self, x, edge_index, ew=None):
        for i, conv in enumerate(self.convs):
            x = conv(x, edge_index, ew)
        return x

def generate_graph(n_nodes, p_edge, n_features, n_classes):
    """Generate random graph."""
    x = torch.randn(n_nodes, n_features)
    edge_list = []
    for i in range(n_nodes):
        for j in range(i+1, n_nodes):
            if np.random.rand() < p_edge:
                edge_list.append([i, j])
                edge_list.append([j, i])
    edge_index = torch.tensor(edge_list, dtype=torch.long).t() if edge_list else torch.zeros(2,0, dtype=torch.long)
    y = torch.randint(0, n_classes, (n_nodes,))
    return x, edge_index, y

def run_gnn(n_nodes, p_edge, n_features, n_classes, n_layers, hidden):
    """Run Sentinel GNN."""
    x, edge_index, y = generate_graph(n_nodes, p_edge, n_features, n_classes)
    
    model = SentinelGNN(n_features, hidden, n_classes, n_layers)
    out = model(x, edge_index)
    pred = out.argmax(dim=1)
    
    # Visualize
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    
    # Graph layout (simple spring)
    pos = np.random.randn(n_nodes, 2) * 2
    for _ in range(50):
        forces = np.zeros_like(pos)
        for i in range(n_nodes):
            for j in range(n_nodes):
                if i != j:
                    diff = pos[i] - pos[j]
                    dist = np.linalg.norm(diff) + 1e-8
                    if dist < 3:
                        forces[i] += diff / dist * (3 - dist) * 0.1
        pos += forces
    
    # Draw graph
    ax = axes[0]
    for i in range(edge_index.size(1)):
        u, v = edge_index[0, i].item(), edge_index[1, i].item()
        if u < v:
            ax.plot([pos[u,0], pos[v,0]], [pos[u,1], pos[v,1]], 'gray', alpha=0.3, linewidth=0.5)
    
    colors = plt.cm.tab10(np.linspace(0, 1, n_classes))
    for c in range(n_classes):
        mask = y.numpy() == c
        ax.scatter(pos[mask,0], pos[mask,1], c=[colors[c]], s=100, label=f'Class {c}', edgecolors='black')
    ax.set_title(f'Graph (N={n_nodes}, E={edge_index.size(1)//2})')
    ax.legend()
    ax.set_aspect('equal')
    ax.axis('off')
    
    # Embedding visualization
    ax = axes[1]
    emb = out.detach().numpy()
    for c in range(n_classes):
        mask = y.numpy() == c
        if mask.sum() > 0:
            ax.scatter(emb[mask, 0], emb[mask, 1] if emb.shape[1] > 1 else emb[mask, 0] * 0,
                      c=[colors[c]], s=100, label=f'Class {c}', edgecolors='black', alpha=0.7)
    ax.set_title('Sentinel GNN Embeddings')
    ax.legend()
    ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('/tmp/gnn_viz.png', dpi=150)
    plt.close()
    
    stats = f"""
## Sentinel GNN Results

| Property | Value |
|----------|-------|
| Nodes | {n_nodes} |
| Edges | {edge_index.size(1)//2} |
| Features | {n_features} |
| Classes | {n_classes} |
| Layers | {n_layers} |
| Hidden | {hidden} |
| Parameters | {sum(p.numel() for p in model.parameters()):,} |

### Key Innovation
**Hyperbolic message passing**: sech(‖x−y‖) is the natural distance kernel
for hyperbolic geometry — matching brain connectome structure.
"""
    return '/tmp/gnn_viz.png', stats

with gr.Blocks(title="Sentinel Graph Neural Network") as demo:
    gr.Markdown("""
    # 🌐 Sentinel Graph Neural Network
    
    **Hyperbolic message passing with sech kernel.**
    
    Brain connectomes and social networks have hyperbolic structure.
    The sech kernel is the natural distance function in hyperbolic space.
    """)
    
    with gr.Row():
        with gr.Column():
            n_nodes = gr.Slider(10, 500, value=100, step=10, label="Nodes")
            p_edge = gr.Slider(0.01, 0.5, value=0.05, label="Edge Probability")
            n_features = gr.Slider(2, 64, value=8, step=2, label="Node Features")
            n_classes = gr.Slider(2, 10, value=3, step=1, label="Classes")
            n_layers = gr.Slider(1, 5, value=2, step=1, label="GNN Layers")
            hidden = gr.Slider(8, 128, value=32, step=8, label="Hidden Dim")
        
        with gr.Column():
            btn = gr.Button("Generate & Run", variant="primary")
            output_img = gr.Image()
            output_stats = gr.Markdown()
    
    btn.click(run_gnn, [n_nodes, p_edge, n_features, n_classes, n_layers, hidden], [output_img, output_stats])
    
    gr.Markdown("""
    ## About Sentinel GNN
    
    - **Message kernel**: sech(‖x−y‖) (hyperbolic geometry)
    - **Activation**: σ(x) = x·sech(x/e) (theorem-backed gradient)
    - **Applications**: Brain connectomes, social networks, molecules, knowledge graphs
    
    [Model Repo](https://huggingface.co/5dimension/sentinel-graph-neural-network)
    """)

if __name__ == "__main__":
    demo.launch()