# データセットを読み込む
dataset = Planetoid(root="./Cora", name="Cora")
# グラフ構造の数
print(len(dataset)) # -> 1
# クラス数
print(dataset.num_classes) # -> 7
# 徴量の次元数(1433種類の特定のワードが論文中に含まれているか. 0:ない,1:ある)
print(dataset.num_node_features) # -> 1433
# ノード数
print(dataset[0].num_nodes) # -> 2708

# ノードの数
num_nodes = dataset[0].num_nodes
# 特徴量の次元数
in_channels = dataset.num_node_features
#
out_channels = dataset.num_classes
# 特徴量行列 (迷走中)
# X = rearrange(dataset[0].x, "V C -> C 1 V")
# X = dataset[0].x
X = rearrange(dataset[0].x, "V C -> V 1 C")
# ラベル
y = dataset[0].y
# 隣接行列
A = create_adjacency_matrix(dataset[0].edge_index, num_nodes)
# 次数行列
D = create_degree_matrix(A)

# DAD行列
DAD = D @ A @ D

def create_adjacency_matrix(edge_index, num_nodes):
    adjacency_matrix = torch.zeros((num_nodes, num_nodes), dtype=torch.float32)

    for i in range(edge_index.size(1)):
        source = edge_index[0, i]
        target = edge_index[1, i]
        adjacency_matrix[source, target] = 1

    return adjacency_matrix


def create_degree_matrix(adjacency_matrix):
    # 各行の要素の合計を計算し、それを逆数にする
    degree_vector = torch.sum(adjacency_matrix, dim=1)
    # 逆数を計算し、ゼロ除算を防ぐためにepsを加算
    inv_degree_vector = 1.0 / (degree_vector + torch.finfo(torch.float32).eps)
    # 対角行列として設定
    degree_matrix = torch.diag(inv_degree_vector)

    return degree_matrix

# モデルの定義
model = GraphConv(in_channels, out_channels)

class GraphConv(nn.Module):
    def __init__(self, in_features, out_features):
        super(GraphConv, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.conv = nn.Conv2d(in_features, out_features, kernel_size=1)

    def forward(self, input, adj):
        """
        Args:
            input (Tensor): graph feature
                            input.size() = (N, V, C)
            adj (Tensor): normalized adjacency matrix.
                          e.g. DAD or DA
                          input.size() = (V, V)
        Returns:
            Tensor: out.size() = (N, V, C_out)
        """
        input = rearrange(input, "N V C -> N C 1 V")
        XW = self.conv(input)
        DADXW = torch.einsum("NCTV,VW->NCTW", XW, adj)
        DADXW = rearrange(DADXW, "N C 1 V -> N V C")
        return DADXW

# 結果の図示
fig, ax = plt.subplots(1, 2, width_ratios=[4, 8])
ax[0].pcolor(X[0], cmap=plt.cm.Blues)
ax[0].set_aspect('equal', 'box')
ax[0].set_title('X', fontsize=10)
ax[0].invert_yaxis()

ax[1].pcolor(new_X[0], cmap=plt.cm.Blues)
ax[1].set_aspect('equal', 'box')
ax[1].set_title('new_X', fontsize=10)
ax[1].invert_yaxis()

dataset[0].x.shape # -> torch.Size([2708, 1433])

X.shape # -> (10, 5, 4)