from keras import optimizers
from keras.constraints import maxnorm
from keras.layers import Dense, Dropout
from keras.models import Sequential
from keras.regularizers import l2

import numpy as np
import pandas as pd

from sklearn.model_selection import train_test_split
from sklearn.multioutput import MultiOutputRegressor
from sklearn.svm import SVR

import evaluation


class Ones():
    """
    Naive model that predicts '1' no matter the input.
    """
    def predict(self, data):
        l, w = data.shape
        return np.ones((l, w))


def get_xcols(dataset):
    """Return training columns.

    :param dataset: dataset
    :return: training columns
    """
    return [col for col in dataset.columns if col.startswith("signal")]


def get_ycols(dataset):
    """Return target columns.

    :param dataset: dataset
    :return: target columns
    """
    return [col for col in dataset.columns if col.startswith("distance")]


def train_svr(data, k):
    """Perform random search of SVR hyperparamaters.

    Split data into training and validation set. Do a random search of SVR
    hyperparameters. Return the SVR model with best performance on validation
    set.

    :param data: training dataset
    :param k: input size
    :return: SVR model
    """
    iterations = 10
    train, valid = train_test_split(data, test_size=0.2)
    xcols, ycols = get_xcols(data), get_ycols(data)
    best_loss = 10000
    best_model = ()

    for i in range(iterations):
        kern = np.random.choice(["rbf", "linear"], size=1)[0]
        c = np.random.uniform(0, 32)
        g = np.random.uniform(0, 4)
        e = np.random.uniform(0, 32)

        svr = MultiOutputRegressor(SVR(kernel=kern, C=c, gamma=g, epsilon=e))
        svr.fit(train[xcols], train[ycols])

        loss = evaluation.evaluate(svr, valid, k)

        if loss < best_loss:
            best_loss = loss
            best_model = svr

    best_model.fit(valid[xcols], valid[ycols])
    return best_model


def build_model(input_size, neurons, layers, activation):
    """Build neural network model.

    :param input_size: size of input and output
    :param neurons: number of neurons
    :param layers: number of hidden layers
    :param activation: activation function
    :return: compiled keras model
    """
    model = Sequential()
    model.add(Dense(units=layers*neurons,
                    activation=activation,
                    input_dim=input_size,
                    kernel_constraint=maxnorm(3),
                    activity_regularizer=l2(0.01)))

    model.add(Dropout(0.2))

    for i in range(layers-1):
        model.add(Dense(units=(layers-i-1)*neurons,
                        activation=activation,
                        kernel_constraint=maxnorm(3),
                        activity_regularizer=l2(0.01)))
        model.add(Dropout(0.2))

    model.add(Dense(units=input_size,
                    activation="linear",
                    kernel_constraint=maxnorm(3),
                    activity_regularizer=l2(0.01)))

    opt = optimizers.Adam()
    model.compile(loss="huber_loss", optimizer=opt)

    return model


def train_nn(data, k):
    """Perform random search of neural network hyperparamaters.

    Split data into training and validation set. Do a random search of
    hyperparameters. Return the model with best performance on validation set.

    :param data: training dataset
    :param k: input size
    :return: neural network model
    """
    iterations = 10
    train, valid = train_test_split(data, test_size=0.2)
    xcols, ycols = get_xcols(data), get_ycols(data)
    best_loss = 10000
    best_model = ()

    for i in range(iterations):
        neurons = np.random.randint(low=k, high=64, size=1)[0]
        layers = np.random.randint(low=1, high=4, size=1)[0]
        activation = np.random.choice(["tanh", "relu", "linear", "sigmoid"], size=1)[0]

        if k == 1:
            neurons = 1
        if activation == "linear":
            layers = 1

        model = build_model(k, neurons, layers, activation)
        model.fit(train[xcols], train[ycols], batch_size=32, epochs=20, verbose=0)
        loss = evaluation.evaluate(model, valid, k)

        if loss < best_loss:
            best_loss = loss
            best_model = model

    best_model.fit(valid[xcols], valid[ycols], batch_size=32, epochs=20, verbose=0)
    return best_model
