# -*- coding: utf-8 -*-
import numpy as np
import sys
from dataset.mnist import load_mnist
from lib.TwoLayerNet import TwoLayerNet

## @file        453_train_neuralnet.py
#  @brief       Chapter 4: Train MNIST dataset usin two-layered neuralnet
#  @author      tsutaj
#  @date        November 8, 2018
#  @note        This script uses slower method for calculating gradient (\p network.numerical_gradient).
#  @note        The implementation using fast method will be also available. (TODO)

def main():
    output_filename = '453_result.out'
    (x_train, t_train), (x_test, t_test) = load_mnist(normalize=True, one_hot_label=True)

    iters_num     = 10000
    train_size    = x_train.shape[0]
    batch_size    = 100
    learning_rate = 0.1

    train_loss_list = []
    train_acc_list  = []
    test_acc_list   = []
    iter_per_epoch  = max(train_size / batch_size, 1)

    network = TwoLayerNet(input_size=784, hidden_size=50, output_size=10)

    with open(output_filename, 'w') as fp:
        for i in range(iters_num):
            print('iter {} ...'.format(i), file=sys.stderr)
            batch_mask = np.random.choice(train_size, batch_size)
            x_batch = x_train[batch_mask]
            t_batch = t_train[batch_mask]
            
            grad = network.numerical_gradient(x_batch, t_batch)
            
            for key in ('W1', 'b1', 'W2', 'b2'):
                network.params[key] -= learning_rate * grad[key]

                loss = network.loss(x_batch, t_batch)
                train_loss_list.append(loss)

            if i % iter_per_epoch == 0:
                train_acc = network.accuracy(x_train, t_train)
                test_acc  = network.accuracy(x_test , t_test )
                train_acc_list.append(train_acc)
                test_acc_list.append(test_acc)
                fp.write('train acc, test acc | {0}, {1}\n'.format(train_acc, test_acc))
                fp.flush() # write immediately
                
if __name__ == '__main__':
    main()
