import numpy as np
from standard_tool import groupby
import sys
import warnings

# This is the version that the queue is only generated by the first node
# These is no tanh in the output layer
def AddIntercept(inputs):
    return np.hstack([inputs, np.ones([inputs.shape[0], 1])])

def Forward_1(inputs, weights):
    return np.tanh(np.dot(inputs, weights))


def Forward_dict(inputs, weights_dict):
    input_update = inputs
    for weights in weights_dict.values():
        input_update = AddIntercept(Forward_1(input_update, weights))
    return input_update[:,:-1]

#end need for outside

#need for test, only used when w0 and w1 are both simple vectors
def Forward_all(inputs_raw,w0,w1):
    return np.outer(np.tanh(np.dot(inputs_raw,w0)),w1)

#end need for test


def Forward_2(inputs, weights):
    return np.dot(inputs, weights)

def Loss_Calc_Greedy(inputs, w0, w1,static1,AF,index_list):
    #This function will be run for each node
    y = Forward_all(inputs, w0, w1)
    y = y + AF*static1
    n_sample = len(index_list)
    return (0.5 * np.linalg.norm(y[index_list] - inputs[index_list,:-1]) ** 2 / n_sample, y)


def calcdist_row(a, b):
    return 0.5 * np.linalg.norm(a - b) ** 2

def Loss_index(inputs,weights0_temp,weights1_temp):
    #This function is only run for the first node
    #It returns the index of the sorted loss
    output = Forward_all(inputs, weights0_temp, weights1_temp)
    loss = map(calcdist_row, inputs[:,:-1], output)
    # from matplotlib import pylab as plt
    # plt.hist(loss)
    # plt.show()
    # sys.exit(1)
    return np.argsort(loss)

def FindBadIndex_list(loss_index,d_next):
    #Need an algorithm to balance the node, but for now maybe we can use the simplest version
    return np.array_split(loss_index,d_next-1)

def Train_SGD(inputs, response, weights0, IO_agent=None,alpha=0.001,epoch_limit=300,AF=1.0):
    #This response should be a vector, can be just a string

    index_list_class = groupby(response)
    n_class = len(index_list_class)

    n_sample = inputs.shape[0]

    #It is just a simple 3-layer neural network

    if (IO_agent is None):
        from IO_Wrapper import IO_Wrapper_stdout
        printlog = IO_Wrapper_stdout.printlog
        printlog_loss = IO_Wrapper_stdout.printlog
    else:
        printlog = IO_agent.printlog
        printlog_loss = IO_agent.printlog_loss


    printlog('Amnesia Factor: ', AF)
    # It is just a simple 3-layer neural network
    d_current = inputs.shape[1]  # with intercepts
    d_next = weights0.shape[1]
    #weights1 = (np.random.rand(d_current - 1, d_next) - 0.5)/256
    weights1 = (np.random.rand(d_next, d_current - 1) - 0.5) / (d_current)
    node_list = np.arange(d_next)

    #Need to know how much partition I need to break each class for each node
    n_each_node = d_next/n_class
    assert d_next%n_class == 0
    sorted_table = sorted(list(np.unique(response)))
    index_list_split = [np.array_split(np.array(index_list_class[x],dtype=int), n_each_node) for x in sorted_table]

    static1 = np.zeros([n_sample,d_current-1]) #the result from last fp

    # now we do SGD
    printlog('Epoch_Limit: ', epoch_limit)

    alphaX = alpha

    currentmeans = []
    net1 = np.zeros([n_sample,d_next])

    if(epoch_limit == 0):
        return (epoch_limit, weights1,np.zeros(d_current-1))


    for i in node_list:
        weights0_temp = weights0[:,i]
        weights1_temp = weights1[i]

        epoch = 0

        index_list = index_list_split[i%n_class][i/n_class]
        epoch_limit_s = epoch_limit

        if(len(index_list)*epoch_limit<4.0/alpha):
            
            epoch_limit_s = np.ceil(4.0/alpha/len(index_list))
            warnings.warn('The total updates for the currect node is less than 4.0/alpha, \
                the epoch_limit is reset to meet that requirement.')


        currentmean = np.mean(inputs[index_list],axis=0)
        data_to_use = inputs[index_list] - np.append(currentmean[:-1],0)

        #printlog("Size for node ",i,":",len(index_list))
        # print index_list
        while (True):
            if (epoch > 100):
                alpha = alphaX / (1.0 + epoch / 1000.0)
            for x,index in  zip(data_to_use,index_list):

                target = inputs[index]

                output0_local = np.tanh(np.dot(x,weights0_temp))
                if(i == 0):
                    output1 = output0_local*weights1_temp
                else:
                    output1_local = output0_local*weights1_temp
                    output1 = output1_local + static1[index]*AF


                delta1 = (output1 - target[:-1]) #delta1 shape [d_input]

                #IMPORTANT!!!
                delta0 = np.dot(weights1_temp, delta1) * (1 - output0_local ** 2)
                grad1_sub = delta1 * output0_local 
                grad0_sub = delta0 * x 
                alpha1 = alpha
                alpha0 = alpha
                step1 = -alpha1 * grad1_sub
                step0 = -alpha0 * grad0_sub

                weights0_temp += step0 
                weights1_temp += step1


                # loss,static1= Loss_Calc_Greedy(inputs,weights0_temp, weights1_temp,static1,(epoch+1)==epoch_limit)

            #loss_prev = loss
            if (epoch ==(epoch_limit_s-1)):
                recenteredinputs = inputs-np.append(currentmean[:-1],0)
                loss,static1= Loss_Calc_Greedy(recenteredinputs, weights0_temp,weights1_temp,static1,AF,index_list)
                net = np.dot(recenteredinputs,weights0_temp)
                net1[:,i] = net
                printlog(i, epoch, loss)
                break
            epoch = epoch + 1


        #end of the optimization for ith node
        # save the weights


    net2 = np.dot(inputs,weights0)
    biasbias = np.mean(net1-net2,axis=0)
    weights0[-1] += biasbias 

    # loss = Loss_Calc(inputs, weights0, weights1)
    # printlog('Final Loss: ', loss)
    return (epoch, weights1)


if __name__ == '__main__':
    inputs = np.array([[-1.0, -0.5, 0.5, 1.0, 1]
        , [-0.5, -1.0, 1.0, 0.5, 1]
        , [0.5, 1.0, -0.5, -1.0, 1]
        , [1.0, 0.5, -1.0, -0.5, 1]])
    response = ["1","0","0","1"]
    np.random.seed(33)
    weights0 = np.random.rand(5, 4) - 0.5
    print weights0
    # weights0 = np.zeros([4,5])
    crap, w1 = Train_SGD(inputs, response,weights0,alpha=0.0)
    print "after training"
    print weights0
    print Forward_2(Forward_1(inputs, weights0), w1)
