from keras.datasets import mnist
import numpy as np
import operator


class NN:
    def __init__(self, input_size, hidden_size, output_size, rate=0.1):
        self.i_dim = input_size
        self.h_dim = hidden_size
        self.o_dim = output_size
        self.W_ih = np.random.rand(self.i_dim, self.h_dim)*0.2-0.1
        self.W_ho = np.random.rand(self.h_dim, self.o_dim)*0.2-0.1
        self.delta_ih = np.zeros(self.h_dim)
        self.delta_ho = np.zeros(self.o_dim)
        self.learning_rate = rate
        self.avg_cost = 0

    def forwardPass(self, sample):
        # update neuron activation state with one sample
        activation_hidden = self.tanh(np.dot(sample, self.W_ih))
        activation_output = self.sigmoid(np.dot(activation_hidden, self.W_ho))
        return activation_hidden, activation_output

    def _objective(self, x, y):
        _output = self.forwardPass(x)[1]
        return np.sum((np.array(_output) - np.array(y))**2)*0.5, np.array(_output)-np.array(y)

    def objective(self, x, y):
        # using batch
        val, vec = 0., np.zeros(self.o_dim)
        for i in range(len(x)):
            t = self._objective(x[i], y[i])
            val += t[0]
            vec += t[1]
        return val, vec

    def nes(self, x, y, x_test, y_test, npop=10, sigma=0.1, alpha=0.001, seed=0):
        np.random.seed(seed)
        for i in range(10000):
            #if i%10==0:
            print('Epoch=', i, 'Error=', self.objective(x, y)[0]/len(x))
            if i%20==0:
                self.test(x_test, y_test)
            N_ih = np.random.randn(npop, self.i_dim, self.h_dim)
            N_ho = np.random.randn(npop, self.h_dim, self.o_dim)
            R = np.zeros(npop)
            for j in range(npop):
                new_nn = NN(self.i_dim, self.h_dim, self.o_dim)
                new_nn.W_ih = self.W_ih + sigma*N_ih[j]
                new_nn.W_ho = self.W_ho + sigma*N_ho[j]
                R[j] = self.solution(new_nn, x, y)

            A = (R - np.mean(R))/ np.std(R)
            self.W_ih += alpha / (npop*sigma) * np.dot(N_ih.T, A).T
            self.W_ho += alpha / (npop*sigma) * np.dot(N_ho.T, A).T

    def _bp(self, x, y):
        cost, cost_vec = self._objective(x, y)
        activation_hidden, activation_output = self.forwardPass(x)
        delta_o = cost_vec*activation_output*(1-activation_output)
        delta_h = np.dot(self.W_ho, delta_o)*(1-activation_hidden**2)
        self.W_ho -= self.learning_rate*delta_o*activation_hidden[:,np.newaxis]
        self.W_ih -= self.learning_rate*delta_h*np.array(x)[:,np.newaxis]
        return self.W_ho, self.W_ih, cost

    def bp(self, x, y):
        self.avg_cost = 0
        t = np.arange(len(x))
        np.random.shuffle(t)
        for i in t[:100]:
            W_ho, W_ih, cost = self._bp(x[i], y[i])
            self.avg_cost += cost
        self.avg_cost /= len(x)
        return self.avg_cost

    def train(self, x, y, x_test, y_test, epoch=300):
        for i in range(epoch):
            self.test(x_test, y_test)
            self.bp(x, y)
            print('Epoch = {} , Average loss = {}'.format(i+1, self.avg_cost))

    def test(self, x, y):
        accuracy = 0
        for i in range(len(x)):
            output = self.forwardPass(x[i])[1]
            index, _ = max(enumerate(output), key=operator.itemgetter(1))
            if y[i][index] == 1:
                accuracy += 1
        print('Test accuracy ={}'.format(accuracy / len(x)))
        return float(accuracy)/len(x)

    @staticmethod
    def sigmoid(x):
        # used in hidden to output connections
        return 1/(1+np.exp(-x))

    @staticmethod
    def tanh(x):
        # used in input to hidden connections
        return (1-np.exp(-2*x))/(1+np.exp(-2*x))

    @staticmethod
    def solution(nn, x, y):
        return -nn.objective(x, y)[0]


class Evolution:
    def __init__(self,
                 init_size,
                 learning_rate,
                 ratio,
                 power,
                 innov,
                 input_dim,
                 hidden_dim,
                 output_dim):
        self.size = init_size
        self.ratio = ratio
        self.power = power
        self.innov = innov
        self.rate = learning_rate
        self.current_herd = {}
        self.label = init_size+1
        self.i_dim, self.h_dim, self.o_dim = input_dim, hidden_dim, output_dim

    def initializa(self):
        for i in range(self.size):
            self.current_herd[i] = [NN(self.i_dim, self.h_dim, self.o_dim), float('inf'), i]
        return self.current_herd

    def _evaluate(self, nn, x, y):
        return nn._objective(x,y)[0]

    @staticmethod
    def _test_evaluate(nn, x, y):
        accuracy = 0
        for i in range(len(x)):
            output = nn.forwardPass(x[i])[1]
            index, _ = max(enumerate(output), key=operator.itemgetter(1))
            if y[i][index] == 1:
                accuracy += 1
        #print('Test accuracy ={}'.format(accuracy / len(x)))
        return float(accuracy) / len(x)

    def evaluate(self, x, y):
        avg_fitness = 0
        best = 0
        for i in self.current_herd:
            fitness = self._test_evaluate(self.current_herd[i][0],x, y)
            self.current_herd[i][1]=fitness
            avg_fitness+=fitness
            best = max(best, fitness)
        print("average fitness=", avg_fitness/self.size, "best fitness=", best)
        return self.current_herd

    def get_parents(self):
        l = sorted(self.current_herd.values(), key=lambda x: -x[1])[:int(self.ratio*self.size)]
        new_herd = {}
        for i in l:
            new_herd[i[2]] = i
        self.current_herd = new_herd
        return self.current_herd

    def generation(self):
        n = np.random.choice(len(self.current_herd), size=int((1-self.innov)*self.size)-len(self.current_herd))
        ind = list(self.current_herd.keys())
        for i in n:
            nn = self.current_herd[ind[i]][0]
            new_nn = NN(self.i_dim, self.h_dim, self.o_dim)
            delta_ih = np.random.normal(0,1,nn.W_ih.shape)
            delta_ih[abs(delta_ih)>self.power] = 0
            delta_ho = np.random.normal(0,1,nn.W_ho.shape)
            delta_ho[abs(delta_ho)>self.power] = 0
            new_nn.W_ih = nn.W_ih+delta_ih*self.rate
            new_nn.W_ho = nn.W_ho+delta_ho*self.rate
            self.current_herd[self.label+1] = [new_nn, float('inf'), self.label]
            self.label += 1
        for i in range(int(self.innov*self.size)):
            self.current_herd[self.label+1] = [NN(self.i_dim, self.h_dim, self.o_dim), float('inf'), self.label+1]
            self.label += 1
        return self.current_herd

    def evolution(self, n, x, y):
        self.initializa()
        for i in range(n):
            self.evaluate(x, y)
            self.get_parents()
            self.generation()
            #self.rate -= 0.0001*n

    def test(self, x, y):
        accuracy = 0
        for i in self.current_herd:
            accuracy += self.current_herd[i][0].test(x, y)
        print('Test accuracy ={}'.format(accuracy / self.size))

if __name__ == '__main__':
    (x_train, y_train), (x_test, y_test) = mnist.load_data()

    x_train = x_train.reshape(60000, 784)
    x_test = x_test.reshape(10000, 784)

    num_classes = 10

    x_train = x_train.astype('float32')
    x_test = x_test.astype('float32')
    x_train /= 255
    x_test /= 255

    t_train = np.zeros((y_train.shape[0], num_classes))
    t_train[np.arange(y_train.shape[0]), y_train] = 1
    y_train = t_train

    t_test = np.zeros((y_test.shape[0], num_classes))
    t_test[np.arange(y_test.shape[0]), y_test] = 1
    y_test = t_test

    nn = NN(784, 100, num_classes)
    nn.nes(x_train[:10000], y_train[:10000], x_test[:1000], y_test[:1000])