import numpy as np
import data
import matplotlib.pyplot as plt


def fcann2_train(X,Y_,param_delta=0.1,param_niter=100000,param_lambda=0.001):
    C = max(Y_) + 1
    n=np.shape(X)[1]
    N = np.shape(X)[0]
    L=5
    b1=np.zeros((L,))
    W1=np.random.randn(n,L)
    b2=np.zeros((C,))
    W2=np.random.randn(L,C)
    # one-hot Y
    Yoh=data.class_to_onehot(Y_) # NxC


    for i in range(param_niter):

        scores1 = np.dot(X,W1) + b1 # N x L
        #hidden layer
        h1=np.maximum(0,scores1) # N X L

        scores2=np.dot(h1,W2) + b2 # N x C

        #softmax
        expscores = np.exp(scores2)

        sumexp = np.sum(expscores,axis=1).reshape(-1,1)

        #probability
        probs=expscores/sumexp # NxC
        logprobs = np.log(probs)  # N x C

        # gubitak
        l_2=param_lambda/2
        l1=l_2*np.sum(W1**2)
        l2=l_2*np.sum(W2**2)
        loss  = - np.mean(np.sum(Yoh*logprobs, axis=1)) + l1 + l2 # scalar


        if i % 10000 == 0:
          print("iteration {}: loss {}".format(i, loss))


        scores1[scores1 <= 0] = 0
        scores1[scores1 > 0] = 1

        dL_ds2 = probs - Yoh # NxC dscores
        dL_dh1 = np.dot(dL_ds2,np.transpose(W2)) # NxL
        dL_ds1 = dL_dh1 * scores1 #NxL



        #gradients
        grad_W2 = (1/N)*np.transpose(np.dot(np.transpose(dL_ds2),h1))# CxL
        grad_b2 = (1/N)*np.sum(np.transpose(dL_ds2), axis = 1) # Cx1
        grad_W1 = (1/N)*np.transpose(np.dot(np.transpose(dL_ds1),X)) #Lxn
        grad_b1 = (1/N)*np.sum(np.transpose(dL_ds1), axis=1) #Lx1

        #reg
        grad_W2 += param_lambda*W2
        grad_W1 += param_lambda*W1


        #poboljšani parametri
        W1 += -param_delta * grad_W1
        b1 += -param_delta * grad_b1
        W2 += -param_delta * grad_W2
        b2 += -param_delta * grad_b2

    return W1, b1, W2, b2

def fcann2_classify(X,W1,W2,b1,b2):

    scores1 = np.dot(X,W1) + b1
    h1=np.maximum(0,scores1)
    scores2=np.dot(h1,W2) + b2
    expscores = np.exp(scores2)
    sumexp = np.sum(expscores,axis=1).reshape(-1,1)
    return expscores/sumexp

def fcann2_decfun(W1, W2, b1, b2):
    def classify(X):
      return np.argmax(fcann2_classify(X, W1, W2, b1, b2),axis=1)
    return classify

if __name__=="__main__":
    np.random.seed(100)

    # get the training dataset
    X,Y_ = data.sample_gmm_2d(6,2,10)

    # train the model
    W1, b1, W2, b2 = fcann2_train(X, Y_)


    # evaluate the model on the training dataset
    probs = fcann2_classify(X, W1, W2, b1, b2)
    Y=np.argmax(probs, axis=1)


    # graph the decision surface
    decfun = fcann2_decfun(W1, W2, b1, b2)
    bbox=(np.min(X, axis=0), np.max(X, axis=0))
    data.graph_surface(decfun, bbox, offset=0.5)

  # graph the data points
    data.graph_data(X, Y_, Y, special=[])
  # show the plot
    plt.show()