import numpy as np
import matplotlib.pyplot as plt
import networkx as nx

# --- データの読み込み ---
G = nx.karate_club_graph()

print(f"ノード数: {G.number_of_nodes()}")
print(f"エッジ数: {G.number_of_edges()}")
print(f"平均次数: {2 * G.number_of_edges() / G.number_of_nodes():.2f}")

# クラブのラベル（Mr. Hi派 or Officer派）
club_labels = [G.nodes[i]['club'] for i in G.nodes()]
color_map = ['red' if label == 'Mr. Hi' else 'blue' for label in club_labels]

# --- グラフの可視化 ---
fig, axes = plt.subplots(1, 2, figsize=(16, 6))

# Spring Layout
ax1 = axes[0]
pos_spring = nx.spring_layout(G, seed=42)
nx.draw(G, pos_spring, node_color=color_map, with_labels=True,
        node_size=300, font_size=8, font_color='white',
        edge_color='gray', alpha=0.8, ax=ax1)
ax1.set_title('Spring Layout')

# Kamada-Kawai Layout
ax2 = axes[1]
pos_kk = nx.kamada_kawai_layout(G)
nx.draw(G, pos_kk, node_color=color_map, with_labels=True,
        node_size=300, font_size=8, font_color='white',
        edge_color='gray', alpha=0.8, ax=ax2)
ax2.set_title('Kamada-Kawai Layout')

plt.suptitle('Zachary Karate Club (Red=Mr.Hi, Blue=Officer)', fontsize=14)
plt.tight_layout()
plt.show()

# --- 基本的な統計量 ---
degrees = dict(G.degree())
print(f"\n各ノードの次数:")
for node, deg in sorted(degrees.items(), key=lambda x: x[1], reverse=True)[:5]:
    print(f"  ノード {node}: 次数 {deg}")

# 次数分布
fig, axes = plt.subplots(1, 2, figsize=(12, 4))

ax1 = axes[0]
degree_values = list(degrees.values())
ax1.hist(degree_values, bins=range(min(degree_values), max(degree_values) + 2),
         alpha=0.7, color='steelblue', edgecolor='black')
ax1.set_xlabel('Degree')
ax1.set_ylabel('Count')
ax1.set_title('Degree Distribution')
ax1.grid(True, alpha=0.3)

# 隣接行列のヒートマップ
ax2 = axes[1]
adj_matrix = nx.adjacency_matrix(G).toarray()
ax2.imshow(adj_matrix, cmap='Blues', interpolation='nearest')
ax2.set_xlabel('Node')
ax2.set_ylabel('Node')
ax2.set_title('Adjacency Matrix')

plt.tight_layout()
plt.show()

# クラスタリング係数
clustering = nx.clustering(G)
avg_clustering = nx.average_clustering(G)
print(f"\n平均クラスタリング係数: {avg_clustering:.4f}")
print(f"グラフの直径: {nx.diameter(G)}")
print(f"平均最短経路長: {nx.average_shortest_path_length(G):.4f}")

import torch
from torch_geometric.datasets import KarateClub
from torch_geometric.utils import to_networkx

# --- データの読み込み ---
dataset = KarateClub()
data = dataset[0]

print(f"ノード数: {data.num_nodes}")
print(f"エッジ数: {data.num_edges}")
print(f"特徴量の次元: {data.num_node_features}")
print(f"ノード特徴量 x: {data.x.shape}")
print(f"エッジインデックス: {data.edge_index.shape}")
print(f"ラベル y: {data.y}")

# エッジインデックスの構造
print(f"\nエッジインデックスの例（最初の10辺）:")
print(f"  送信元: {data.edge_index[0, :10].tolist()}")
print(f"  受信先: {data.edge_index[1, :10].tolist()}")

# PyGのDataオブジェクトをNetworkXグラフに変換して可視化
G_pyg = to_networkx(data, to_undirected=True)

fig, ax = plt.subplots(figsize=(8, 6))
pos = nx.spring_layout(G_pyg, seed=42)
node_colors = data.y.numpy()
nx.draw(G_pyg, pos, node_color=node_colors, cmap=plt.cm.Set3,
        with_labels=True, node_size=300, font_size=8,
        edge_color='gray', alpha=0.8, ax=ax)
ax.set_title('Zachary Karate Club (PyG)')
plt.tight_layout()
plt.show()