import numpy as np
import random
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D


def logistic_z(z):
    return 1.0/(1.0+np.exp(-z))


def logistic_wx(w,x):
    return logistic_z(np.inner(w,x))


def classify(w,x):
    x = np.hstack(([1],x))
    return 0 if (logistic_wx(w,x)<0.5) else 1


# x_train = [number_of_samples,number_of_features] = number_of_samples x \in R^number_of_features
def stochast_train_w(x_train,y_train,learn_rate=0.1,niter=1000):
    x_train = np.hstack((np.array([1]*x_train.shape[0]).reshape(x_train.shape[0],1),x_train))
    dim = x_train.shape[1]
    num_n = x_train.shape[0]
    w = np.random.rand(dim)
    index_lst = []
    for it in xrange(niter):
        if len(index_lst) == 0:
            index_lst = random.sample(xrange(num_n), k=num_n)
        xy_index = index_lst.pop()
        x = x_train[xy_index,:]
        y = y_train[xy_index]
        for i in xrange(dim):
            update_grad = 1 ### something needs to be done here
            w[i] = w[i] + learn_rate ### something needs to be done here
    return w


def batch_train_w(x_train,y_train,learn_rate=0.1,niter=1000):
    x_train = np.hstack((np.array([1]*x_train.shape[0]).reshape(x_train.shape[0],1),x_train))
    dim = x_train.shape[1]
    num_n = x_train.shape[0]
    w = np.random.rand(dim)
    index_lst = []
    for it in xrange(niter):
        for i in xrange(dim):
            update_grad=0.0
            for n in xrange(num_n):
                update_grad+=(-logistic_wx(w,x_train[n])+y_train[n])# something needs to be done here
            w[i] = w[i] + learn_rate * update_grad/num_n
    return w


def train_and_plot(xtrain, ytrain, xtest, ytest, training_method, learn_rate=0.1, niter=10):
    plt.figure()
    #train data
    data = pd.DataFrame(np.hstack((xtrain,ytrain.reshape(xtrain.shape[0],1))),columns=['x','y','lab'])
    ax=data.plot(kind='scatter',x='x',y='y',c='lab',cmap=cm.copper,edgecolors='black')

    #train weights
    w=training_method(xtrain,ytrain,learn_rate,niter)
    error=[]
    y_est=[]
    for i in xrange(len(ytest)):
        error.append(np.abs(classify(w,xtest[i])-ytest[i]))
        y_est.append(classify(w,xtest[i]))
    y_est=np.array(y_est)
    data_test = pd.DataFrame(np.hstack((xtest,y_est.reshape(xtest.shape[0],1))),columns=['x','y','lab'])
    data_test.plot(kind='scatter',x='x',y='y',c='lab',ax=ax,cmap=cm.coolwarm,edgecolors='black')
    print "error=",np.mean(error)
    return w


def l_simple(weights):
    return (logistic_wx(weights, [1,0]) - 1) ** 2 + (logistic_wx(weights, [0,1])) ** 2 + (logistic_wx(weights, [1, 1]) - 1) ** 2


def i_simple_loss():
    r = 13
    w_grid = np.zeros((r, r, 3))
    for row in range(r):
        for col in range(r):
            x = row - r/2
            y = col - r/2
            w_grid[row][col] = [x, y, l_simple([x, y])]

    lowest_point = [0, 0 , float("inf")]
    number_of_lowest_points = 1
    for row in w_grid:
        for col in row:
            if col[2] < lowest_point[2]:
                lowest_point = col
            elif col[2] == lowest_point[2]:
                number_of_lowest_points += 1

    print "Lowest point: ", lowest_point, "\nNumber of points equally low: ", number_of_lowest_points

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    Axes3D.set_xlabel(ax, "x")
    Axes3D.set_ylabel(ax, "y")
    Axes3D.set_zlabel(ax, "loss")
    Axes3D.plot_trisurf(ax, w_grid[:,:,0].flatten(), w_grid[:,:,1].flatten(), w_grid[:,:,2].flatten())
    plt.show()


def der_l_simple(w):
    def expo(w, x):
        return np.exp(-np.inner(w, x))
    w1_p1 = (logistic_wx(w, [1,0]) - 1) * (logistic_wx(w, [1,0])**2) * expo(w, [1,0])
    w2_p1 = (logistic_wx(w, [0,1])) * (logistic_wx(w, [0,1])**2) * expo(w, [0,1])
    shared_p2 = (logistic_wx(w, [1,1]) - 1) * (logistic_wx(w, [1,1])**2) * expo(w, [1,1])
    return [w1_p1 + shared_p2, w2_p1 + shared_p2]

def i_gradient_descent(eta, iterations, plot=False):
    weights = [-10, 10]
    t = np.arange(0, iterations+1, 1)
    l = np.zeros(iterations+1)
    l[0] = l_simple(weights)
    for i in range(iterations):
        weights = np.array(weights) - eta * np.array(der_l_simple(weights))
        l[i+1] = l_simple(weights)
    if plot:
        print weights, l_simple(weights)
        plt.plot(t, l)
        plt.xlabel("iteration")
        plt.ylabel("loss (logscale)")
        plt.yscale("log")
        plt.show()
    return l_simple(weights)


def i_plot_gd():
    iterations = 10000
    doublings = 30
    etas = np.zeros(doublings)
    etas[0] = 0.000001
    for i in range(len(etas) - 1):
        etas[i+1] = etas[i]*2
    l_mins = np.zeros(doublings)
    for i in range(len(l_mins)):
        l_mins[i] = i_gradient_descent(etas[i], iterations)
    plt.plot(etas, l_mins)
    plt.xscale("log")
    plt.xlabel('eta (logscale)')
    plt.ylabel('loss minimum')
    plt.grid(True)
    plt.show()



i_plot_gd()
#i_gradient_descent(0.1, 10000, plot=True)