seed = 123

np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.backends.cudnn.deterministic = True
torch.use_deterministic_algorithms = True

NUM_EPOCH = 100
BATCH_SIZE = 64

model = ST_GCN(
    num_classes=10,
    in_channels=3,
    t_kernel_size=9,
    hop_size=2,
)

# モデルを定義
class ST_GCN(nn.Module):
    def __init__(self, num_classes, in_channels, t_kernel_size, hop_size):
        super().__init__()

        # 誰?
        graph = Graph(hop_size)
        ...

    def forward(self, x):
        ...

# 隣接行列を作成
class Graph:
    def __init__(self, hop_size, bvhp):
        skeleton = bvhp.get_skeleton()
        node_num = len(skeleton)
        self.edge = self.__get_edge(bvhp)

        # hop数分離れた関節を取得
        hop_dis = self.__get_hop_distance(self.node_num, self.edge, hop_size)

        # 隣接行列を作成
        self.A = self.__get_adjacency_mat(hop_dis, hop_size)
        
    ...

class STGC_block(nn.Module):
    def __init__(
        self, in_channels, out_channels, stride, t_kernel_size, A_size, dropout=0.5
    ):
        super().__init__()
        # 空間グラフの畳み込み
        self.sgc = SpatialGraphConvolution(
            in_channels=in_channels, out_channels=out_channels, s_kernel_size=A_size[0]
        )

        # 重要なエッジを学習してくれるパラメータ
        self.M = nn.Parameter(torch.ones(A_size))

        # 時間畳み込み
        self.tgc = nn.Sequential(
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Conv2d(
                out_channels,
                out_channels,
                (t_kernel_size, 1),
                (stride, 1),
                ((t_kernel_size - 1) // 2, 0),
            ),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
        )

    def forward(self, x, A):
        # 実際に畳み込む
        x = self.tgc(self.sgc(x, A * self.M))
        return x

# モデルを定義
class ST_GCN(nn.Module):
    def __init__(self, num_classes, in_channels, t_kernel_size, hop_size):
        ...

        # 空間・時間で畳み込むやつ. 何度もやる理由はわからない
        self.stgc1 = STGC_block(in_channels, 32, 1, t_kernel_size, A_size)
        self.stgc2 = STGC_block(32, 32, 1, t_kernel_size, A_size)
        self.stgc3 = STGC_block(32, 32, 1, t_kernel_size, A_size)
        self.stgc4 = STGC_block(32, 64, 2, t_kernel_size, A_size)
        self.stgc5 = STGC_block(64, 64, 1, t_kernel_size, A_size)
        self.stgc6 = STGC_block(64, 64, 1, t_kernel_size, A_size)
        ...

    def forward(self, x):
        # データの次元をリシェイプしつつ, バッチ正規化を行い, 元の次元に戻す
        N, C, T, V = x.size()  # batch, channel, frame, node
        x = x.permute(0, 3, 1, 2).contiguous().view(N, V * C, T)
        x = self.bn(x)
        x = x.view(N, V, C, T).permute(0, 2, 3, 1).contiguous()

        # 特徴を深くまで学習している
        x = self.stgc1(x, self.A)
        x = self.stgc2(x, self.A)
        x = self.stgc3(x, self.A)
        x = self.stgc4(x, self.A)
        x = self.stgc5(x, self.A)
        x = self.stgc6(x, self.A)

        # 予測
        ## 平均プーリング
        x = F.avg_pool2d(x, x.size()[2:])
        x = x.view(N, -1, 1, 1)

        ## 全結合層の適用
        x = self.fc(x)
        x = x.view(x.size(0), -1)

        return x

model = ST_GCN(
    num_classes=10,
    in_channels=3,
    t_kernel_size=9,
    hop_size=2,
    bvhp=bvhp, # Graph をBVHファイルから生成するように変更した
)

# 確率的勾配降下法を使う
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

criterion = torch.nn.CrossEntropyLoss()

data_loader = dict()
data_loader["train"] = torch.utils.data.DataLoader(
    # 誰?
    dataset=Feeder(data_path="data/train_data.npy", label_path="data/train_label.npy"),
    batch_size=BATCH_SIZE,
    shuffle=True,
)
data_loader["test"] = torch.utils.data.DataLoader(
    dataset=Feeder(data_path="data/test_data.npy", label_path="data/test_label.npy"),
    batch_size=BATCH_SIZE,
    shuffle=False,
)

class Feeder(torch.utils.data.Dataset):
    def __init__(self, data_path, label_path):
        super().__init__()

        # 良い感じにデータをロードする
        self.motion_df = self.__load_bvh(bvh_path)
        self.label_df = self.__load_label(label_path)

    # データの数を返す
    def __len__(self):
        return len(self.label)

    # index番目のデータを返す
    def __getitem__(self, index):
        data = np.array(self.data[index])
        label = self.label[index]

        return data, label

model.train()

model.eval()

# エポックをループ
for epoch in range(1, NUM_EPOCH + 1):
    correct = 0  # 正しく分類されたサンプル数をカウント
    sum_loss = 0 # 累積損失を計算

    # ミニバッチ(データとラベルの組み合わせ) でループ
    for batch_idx, (data, label) in enumerate(data_loader["train"]):
        
        # GPU に転送する ※MacBookは無理
        data = data.cuda()
        label = label.cuda()

        # 予測を行う
        output = model(data)

        # 損失を計算
        loss = criterion(output, label)
        
        # オプティマイザ の勾配を初期化
        optimizer.zero_grad()
        
        # 逆伝播で損失を計算
        loss.backward()
        
        # モデルのパラメータを更新
        optimizer.step()

        # 累積損失に加算
        sum_loss += loss.item()
        
        # 予測結果から正解数を出す
        _, predict = torch.max(output.data, 1)
        correct += (predict == label).sum().item()

    print(
        "# Epoch: {} | Loss: {:.4f} | Accuracy: {:.4f}".format(
            epoch,
            sum_loss / len(data_loader["train"].dataset),
            (100.0 * correct / len(data_loader["train"].dataset)),
        )
    )