# task1.py
import numpy as np

class GNN1:
    def __init__(self, D, T, W):
        ...
    ...
    @staticmethod
    def f(x):
        return np.frompyfunc(relu, 1, 1)(x)

    @staticmethod
    def aggregate1(G, x):
        a = np.dot(G, x)
        return a
    def aggregate2(self, a):
        x = self.f(np.dot(a, self.W))
        return x

    def readout(self, G):
        N = len(G)
        x = self.x0(N)
        for _ in range(self.T):
            a = self.aggregate1(G, x)
            x = self.aggregate2(a)
        return x.sum(axis=0)

# task2.py
from task1 import *

class GNN2(GNN1):
    def __init__(self, D, T, sigma=0.4,
                 epsilon=1.0e-3): # 数値微分の微小変分
        ...
    ...

    # 学習パラメータの初期化。Thetaと(W,A,b)でメモリを共有する。
    # bも1要素numpy配列にしておくこと。
    def init_Theta(self, sigma):
        self.Theta = np.random.normal(0,sigma,self.D*self.D+self.D +1)
        self.W, self.A, self.b = self.decode_Theta()
    def decode_Theta(self):
        return (self.Theta[:self.D*self.D].reshape((self.D,self.D)),
                self.Theta[self.D*self.D:self.D*(self.D+1)],
                self.Theta[self.D*(self.D+1):self.D*(self.D+1)+1])
        # numpy配列のスライスやreshapeはもとのオブジェクトとメモリを共有している。
        # （メモリのコピーを作成しない。）
        # 不意の副作用を気にすることが多いが、ここでは積極的に利用。

    # GNN1で定義したreadout()を用いて陽性確率／損失関数を計算する
    def s(self, G):
        h = self.readout(G)
        return np.dot(h,self.A) + self.b[0]
    def p(self, G):
        return sigmoid(self.s(G))
    def loss(self,G,y):
        return binary_cross_entropy(y, self.s(G))

    # 損失関数loss()の数値微分
    def grad(self, G, y):
        L0 = self.loss(G,y)
        K = len(self.Theta)
        dL = np.empty(K)
        for i in range(K):
            self.Theta[i] += self.epsilon
            dL[i] = (self.loss(G,y) - L0)/self.epsilon
            self.Theta[i] -= self.epsilon
        return dL

    # 勾配降下
    def shift_Theta(self, dTheta):
        self.Theta += dTheta
    def descendant_update(self, G, y, alpha):
        self.shift_Theta(-alpha*self.grad(G,y))

# task3.py
from task2 import *

class GNN3(GNN2):
    ...

    def fit(self,
            x, y,   # 学習データ。
            optimizer,      # Optimizer派生クラスのオブジェクトを与える。
            epochs=1, batchsize=16,
            validation=None     # 検証データ。タプル (vx,vy) で与える。
            ):
        n = len(y)
        ...

        for e in range(epochs):
            # epoch毎にデータをランダムに並べ替える。
            order = np.random.permutation(n)

            for b in range((n-1)//batchsize+1):
                start = batchsize * b
                end = min(start+batchsize, n)
                # 勾配ミニバッチ平均
                grad = np.mean([ self.grad(x[i],y[i])
                                 for i in order[start:end]],axis=0)
                # optimizerで変化量を求め、Thetaを動かす。
                self.shift_Theta(optimizer.update(grad))
                ...
        ...
        # loss, accの履歴を返す。
        return batch_losses, epoch_losses, epoch_accs, batch_vlosses, epoch_vlosses, epoch_vaccs

# task3.py つづき
### ほぼダミーの基底クラス。（抽象クラスにした方がすっきりするが、さぼった）
class Optimizer():
    @staticmethod
    def update(grad):
        # 勾配を受け取り、内部状態を更新し、変化量を返すメソッド。
        # （派生クラスではそのような関数に再定義する。）
        return -grad

### 確率的勾配降下法
# Optimizerを継承してupdate()を上書き再定義。
class SGD(Optimizer):
    def __init__(self, alpha=1.0e-4):
        self.alpha = alpha

    def update(self, grad):
        # SGDは更新される内部状態を持たない。
        return -self.alpha * grad

### 運動量付き確率的勾配降下法
class MomentumSGD(Optimizer):
    def __init__(self, alpha=1.0e-4, eta=0.9):
        self.alpha = alpha
        self.eta = eta
        self.w = None

    def update(self, grad):
        # 内部のwを更新しながら変化量を返す。
        if self.w is None:
            self.w = np.zeros(len(grad))
        self.w *= self.eta
        self.w -= self.alpha * grad
        return self.w

# task4a.py
from task3 import *

# 原論文 https://arxiv.org/abs/1412.6980 の通り、愚直に実装する。
class Adam(Optimizer):
    def __init__(self,
                 alpha=1.0e-3, beta1=0.9, beta2=0.999,
                 epsilon=1.0e-8):
        # 内部状態ベクトル m,v を初期化
        ...

    def update(self, grad):
        # m,v を更新して変化量を返す
        ...

# task4b.py
from task3 import *

class GNN4(GNN3):
    def __init__(self, D, T, sigma=0.4,
                 Nw=2): # Wのlayer数。 W.shape=(Nw,D,D)
        ...

    ### パラメータ初期化の関数2つを上書きする。
    def init_Theta(self, sigma):
        self.Theta = np.random.normal(0,sigma,Nw*self.D*self.D+self.D+1)
        self.W, self.A, self.b = self.decode_Theta()

    def decode_Theta(self):
        Wsize, Asize = self.Nw*self.D*self.D, self.D
        return (self.Theta[:Wsize].reshape((self.Nw,self.D,self.D)),
                self.Theta[Wsize:Wsize+Asize],
                self.Theta[Wsize+Asize:Wsize+Asize+1])

    ### 集約2を多層ニューラルネットに変更。
    def aggregate2(self, a):
        x = a
        for i in range(self.Nw):
            x = self.f(np.dot(x, self.W[i]))
        return x