import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Degree-normalized neighbor aggregation.

    Computes a GCN-style sparse aggregation where each incoming message is
    scaled by 1 / sqrt(deg(dst) * deg(src)). This captures a common spectral
    graph normalization pattern over CSR adjacency.
    """

    def __init__(self):
        super().__init__()

    def forward(
        self,
        row_ptr: torch.Tensor,
        col_idx: torch.Tensor,
        node_feat: torch.Tensor,
        degrees: torch.Tensor,
    ) -> torch.Tensor:
        num_nodes = row_ptr.numel() - 1
        feat_dim = node_feat.shape[1]
        out = torch.zeros(num_nodes, feat_dim, dtype=node_feat.dtype, device=node_feat.device)

        for dst in range(num_nodes):
            start = int(row_ptr[dst].item())
            end = int(row_ptr[dst + 1].item())
            if end <= start:
                continue
            src = col_idx[start:end].long()
            norm = 1.0 / torch.sqrt(degrees[dst] * degrees[src]).unsqueeze(-1)
            out[dst] = (node_feat[src] * norm).sum(dim=0)

        return out


num_nodes = 640
avg_degree = 16
feat_dim = 80




def get_inputs():
    return list(popcorn_pri.make_csr_degree_normalized_graph(num_nodes, avg_degree, feat_dim))



def get_init_inputs():
    return []