import numpy as np
import chainer
from chainer import Chain, Variable
import chainer.functions as F
import chainer.links as L
import openpyxl

class NeuralNet(chainer.Chain):
    def __init__(self, n_units, n_out):
        super().__init__(
            l1=L.Linear(None, n_units),
            l2=L.Linear(n_units, n_out),
        )

    def __call__(self, x):
        h1 = F.relu(self.l1(x))
        return self.l2(h1)

def check_accuracy(model, xs, ts):
    ys = model(xs)
    loss = F.softmax_cross_entropy(ys, ts)
    ys = np.argmax(ys.data, axis=1)
    cors = (ys == ts)
    num_cors = sum(cors)
    accuracy = num_cors / ts.shape[0]
    return accuracy, loss

def main():
    wb = openpyxl.load_workbook('excel/omomi40.xlsx')
    wb.create_sheet(index=1, title='ac')
    wb.create_sheet(index=2, title='l1_W_1')
    wb.create_sheet(index=3, title='l1_b_1')
    wb.create_sheet(index=4, title='l2_W_1')
    wb.create_sheet(index=5, title='l2_b_1')

    wb.create_sheet(index=6, title='l1_W_5')
    wb.create_sheet(index=7, title='l1_b_5')
    wb.create_sheet(index=8, title='l2_W_5')
    wb.create_sheet(index=9, title='l2_b_5')

    ac = wb['ac']
    l1_W_1 = wb['l1_W_1']
    l1_b_1 = wb['l1_b_1']
    l2_W_1 = wb['l2_W_1']
    l2_b_1 = wb['l2_b_1']

    l1_W_5 = wb['l1_W_5']
    l1_b_5 = wb['l1_b_5']
    l2_W_5 = wb['l2_W_5']
    l2_b_5 = wb['l2_b_5']

    for aaa in range(784):
        l1_W_1.cell(row=1, column=aaa + 2).value = str(int(aaa / 784) + 1) + '-' + str((aaa % 784) + 1)

    for aaa in range(3920):
        l1_W_5.cell(row=1, column=aaa + 2).value = str(int(aaa / 784) + 1) + '-' + str((aaa % 784) + 1)

    for bbb in range(1):
        l1_b_1.cell(row=1, column=bbb + 2).value = str(bbb + 1)

    for bbb in range(5):
        l1_b_5.cell(row=1, column=bbb + 2).value = str(bbb + 1)

    for ccc in range(10):
        l2_W_1.cell(row=1, column=ccc + 2).value = str(int(ccc / 10) + 1) + '-' + str((ccc % 10) + 1)

    for ccc in range(50):
        l2_W_5.cell(row=1, column=ccc + 2).value = str(int(ccc / 10) + 1) + '-' + str((ccc % 10) + 1)

    for ddd in range(10):
        l2_b_1.cell(row=1, column=ddd + 2).value = str(ddd + 1)
        l2_b_5.cell(row=1, column=ddd + 2).value = str(ddd + 1)

    for count in range(2):
        if (count == 0):
            model = NeuralNet(1, 10)
        elif (count == 1):
            model = NeuralNet(5, 10)
        optimizer = chainer.optimizers.Adam()
        optimizer.setup(model)

        train, test = chainer.datasets.get_mnist()
        xs, ts = train._datasets
        txs, tts = test._datasets

        bm = 1000
        a = 2

        for i in range(400):
            l1_W_1.cell(row=i + 2, column=1).value = str(i + 1)
            l1_b_1.cell(row=i + 2, column=1).value = str(i + 1)
            l2_W_1.cell(row=i + 2, column=1).value = str(i + 1)
            l2_b_1.cell(row=i + 2, column=1).value = str(i + 1)

            l1_W_5.cell(row=i + 2, column=1).value = str(i + 1)
            l1_b_5.cell(row=i + 2, column=1).value = str(i + 1)
            l2_W_5.cell(row=i + 2, column=1).value = str(i + 1)
            l2_b_5.cell(row=i + 2, column=1).value = str(i + 1)

            for j in range(60):
                model.zerograds()
                x = xs[(j * bm):((j + 1) * bm)]
                t = ts[(j * bm):((j + 1) * bm)]
                t = Variable(np.array(t, "i"))
                y = model(x)
                loss = F.softmax_cross_entropy(y, t)
                loss.backward()
                optimizer.update()

            accuracy_train, loss_train = check_accuracy(model, xs, ts)
            accuracy_test, _ = check_accuracy(model, txs, tts)

            print("Epoch %d loss(train) = %f, accuracy(train) = %f, accuracy(test) = %f" % (
            i + 1, loss_train.data, accuracy_train, accuracy_test))
            ac.cell(row=2 + i, column=count + 1).value = accuracy_test

            if (count == 0):
                for k in range(1):
                    for t in range(784):
                        l1_W_1.cell(row=i + 2, column=a).value = model.l1.W.data[k][t]
                        a += 1
                a = 2

                for kk in range(1):
                    l1_b_1.cell(row=i + 2, column=kk + 2).value = model.l1.b.data[kk]

                for kkk in range(10):
                    for tt in range(1):
                        l2_W_1.cell(row=i + 2, column=a).value = model.l2.W.data[kkk][tt]
                        a+=1
                a=2

                for kkkk in range(10):
                    l2_b_1.cell(row=i + 2, column=kkkk + 2).value = model.l2.b.data[kkkk]

            elif (count == 1):
                for k in range(5):
                    for t in range(784):
                        l1_W_5.cell(row=i + 2, column=a).value = model.l1.W.data[k][t]
                        a += 1
                a = 2

                for kk in range(5):
                    l1_b_5.cell(row=i + 2, column=kk + 2).value = model.l1.b.data[kk]

                for kkk in range(10):
                    for tt in range(5):
                        l2_W_5.cell(row=i + 2, column=a).value = model.l2.W.data[kkk][tt]
                        a += 1
                a = 2


                for kkkk in range(10):
                    l2_b_5.cell(row=i + 2, column=kkkk + 2).value = model.l2.b.data[kkkk]

    wb.save('excel/omomi40.xlsx')


if __name__ == '__main__':
    main()