"""
Topological and Hypergraph structures for TopoHyper.

Implements:
- SimplicialComplex: boundary operators, Hodge Laplacians
- Hypergraph: incidence matrix, spectral operators
- build_topohyper_structure(): constructs both from features via k-NN
"""

import torch
import numpy as np
from scipy.spatial import Delaunay
from sklearn.neighbors import NearestNeighbors
from itertools import combinations


class SimplicialComplex:
    """
    Simplicial complex up to dimension 2 (nodes, edges, triangles).
    
    Computes boundary operators B1, B2 and Hodge Laplacians L0, L1.
    """
    
    def __init__(self, nodes, edges, triangles=None):
        self.nodes = list(nodes)
        self.edges = [tuple(sorted(e)) for e in edges]
        self.edges = list(set(self.edges))
        self.triangles = []
        if triangles:
            self.triangles = [tuple(sorted(t)) for t in triangles]
            self.triangles = list(set(self.triangles))
        
        self.node_idx = {n: i for i, n in enumerate(self.nodes)}
        self.edge_idx = {e: i for i, e in enumerate(self.edges)}
        
        self._build_boundary_operators()
    
    def _build_boundary_operators(self):
        """Build B1 (edges->nodes) and B2 (triangles->edges)."""
        n_nodes = len(self.nodes)
        n_edges = len(self.edges)
        n_triangles = len(self.triangles)
        
        # B1: n_nodes x n_edges
        self.B1 = torch.zeros(n_nodes, n_edges)
        for idx, (i, j) in enumerate(self.edges):
            ni, nj = self.node_idx[i], self.node_idx[j]
            self.B1[ni, idx] = -1.0
            self.B1[nj, idx] = 1.0
        
        # B2: n_edges x n_triangles
        self.B2 = torch.zeros(n_edges, n_triangles)
        for t_idx, (a, b, c) in enumerate(self.triangles):
            faces = [(a, b), (a, c), (b, c)]
            signs = [1.0, -1.0, 1.0]
            for face, sign in zip(faces, signs):
                if face in self.edge_idx:
                    self.B2[self.edge_idx[face], t_idx] = sign
    
    def hodge_laplacian(self, k=0):
        """
        Compute k-th Hodge Laplacian.
        L_k = B_k^T B_k + B_{k+1} B_{k+1}^T
        
        L0 = B1 B1^T (graph Laplacian)
        L1 = B1^T B1 + B2 B2^T
        """
        if k == 0:
            return self.B1 @ self.B1.t()
        elif k == 1:
            L_down = self.B1.t() @ self.B1
            L_up = self.B2 @ self.B2.t() if self.B2.shape[1] > 0 else torch.zeros(len(self.edges), len(self.edges))
            return L_down + L_up
        else:
            raise ValueError(f"Only k=0,1 supported, got k={k}")
    
    def adjacency_matrix(self):
        """Node adjacency from edges."""
        n = len(self.nodes)
        A = torch.zeros(n, n)
        for (i, j) in self.edges:
            ni, nj = self.node_idx[i], self.node_idx[j]
            A[ni, nj] = 1.0
            A[nj, ni] = 1.0
        return A


class Hypergraph:
    """
    Hypergraph H = (V, E, W).
    
    Implements spectral convolution operators:
    D_v^{-1/2} H W D_e^{-1} H^T D_v^{-1/2}
    """
    
    def __init__(self, num_nodes, hyperedges, weights=None):
        self.num_nodes = num_nodes
        self.hyperedges = hyperedges
        self.num_edges = len(hyperedges)
        
        if weights is None:
            self.weights = torch.ones(self.num_edges)
        else:
            self.weights = torch.tensor(weights, dtype=torch.float)
        
        self._build_incidence()
    
    def _build_incidence(self):
        """Build incidence matrix H (num_nodes x num_edges)."""
        self.H = torch.zeros(self.num_nodes, self.num_edges)
        for e_idx, he in enumerate(self.hyperedges):
            for v in he:
                self.H[v, e_idx] = 1.0
        
        # Degree matrices
        self.Dv = torch.diag(self.H @ self.weights)
        self.De = torch.diag(self.H.t().sum(dim=0))
    
    def propagation_matrix(self):
        """
        Compute D_v^{-1/2} H W D_e^{-1} H^T D_v^{-1/2}
        """
        Dv_inv_sqrt = torch.diag(1.0 / (torch.diagonal(self.Dv).sqrt() + 1e-7))
        De_inv = torch.diag(1.0 / (torch.diagonal(self.De) + 1e-7))
        W = torch.diag(self.weights)
        
        return Dv_inv_sqrt @ self.H @ W @ De_inv @ self.H.t() @ Dv_inv_sqrt
    
    def adjacency_matrix(self):
        """Node adjacency from hyperedges: A = H H^T - D_v."""
        A = self.H @ self.H.t()
        A = A - torch.diag(torch.diagonal(A))
        A = (A > 0).float()
        return A


def build_topohyper_structure(features, k=6):
    """
    Build both SimplicialComplex and Hypergraph from node features.
    
    1. Build k-NN graph
    2. Find triangles (3-cliques) for simplicial complex
    3. Create hyperedges from k-NN neighborhoods
    
    Args:
        features: (N, D) node feature tensor
        k: number of nearest neighbors
    
    Returns:
        sc: SimplicialComplex
        hg: Hypergraph
        edge_index: (2, E) edge index tensor
    """
    N = features.shape[0]
    feat_np = features.detach().numpy() if isinstance(features, torch.Tensor) else features
    
    # k-NN graph
    nn = NearestNeighbors(n_neighbors=min(k + 1, N), metric='euclidean')
    nn.fit(feat_np)
    distances, indices = nn.kneighbors(feat_np)
    
    # Build edges (undirected)
    edges = set()
    for i in range(N):
        for j in indices[i, 1:]:  # skip self
            edge = tuple(sorted([i, j]))
            edges.add(edge)
    edges = list(edges)
    
    # Find triangles (3-cliques)
    adj = {i: set() for i in range(N)}
    for (i, j) in edges:
        adj[i].add(j)
        adj[j].add(i)
    
    triangles = set()
    for (i, j) in edges:
        common = adj[i] & adj[j]
        for c in common:
            tri = tuple(sorted([i, j, c]))
            triangles.add(tri)
    triangles = list(triangles)
    
    # Build simplicial complex
    sc = SimplicialComplex(list(range(N)), edges, triangles)
    
    # Build hypergraph: each node's k-NN neighborhood is a hyperedge
    hyperedges = []
    for i in range(N):
        he = list(indices[i])  # includes self
        hyperedges.append(he)
    
    # Also add triangle-based hyperedges
    for tri in triangles:
        hyperedges.append(list(tri))
    
    hg = Hypergraph(N, hyperedges)
    
    # Edge index for GNN baselines
    src = [e[0] for e in edges] + [e[1] for e in edges]
    dst = [e[1] for e in edges] + [e[0] for e in edges]
    edge_index = torch.tensor([src, dst], dtype=torch.long)
    
    return sc, hg, edge_index