import numpy as np
import scipy.sparse as sp

from HPOlibConfigSpace.configuration_space import ConfigurationSpace
from HPOlibConfigSpace.conditions import EqualsCondition, InCondition
from HPOlibConfigSpace.hyperparameters import UniformFloatHyperparameter, \
    UniformIntegerHyperparameter, CategoricalHyperparameter, Constant

from autosklearn.pipeline.components.base import AutoSklearnRegressionAlgorithm
from autosklearn.pipeline.constants import *


class LinReg(AutoSklearnRegressionAlgorithm):

    def __init__(self, number_updates, batch_size, dropout_output,
                 learning_rate, solver, lambda2,
                 momentum=0.99, beta1=0.9, beta2=0.9, rho=0.95,
                 lr_policy='fixed', gamma=0.01, power=1.0, epoch_step=2,
                 random_state=None):
        self.number_updates = number_updates
        self.batch_size = batch_size
        self.dropout_output = dropout_output
        self.learning_rate = learning_rate
        self.lr_policy = lr_policy
        self.lambda2 = lambda2
        self.momentum = momentum
        self.beta1 = 1-beta1 if beta1 is not None else 0.9
        self.beta2 = 1-beta2 if beta2 is not None else 0.99
        self.rho = rho
        self.solver = solver
        self.gamma = gamma
        self.power = power
        self.epoch_step = epoch_step

        # Empty features and shape
        self.n_features = None
        self.input_shape = None
        self.m_issparse = False
        self.m_isregression = True
        self.m_isbinary = False
        self.m_ismultilabel = False

        self.estimator = None

    def _prefit(self, X, y):
        self.batch_size = int(self.batch_size)
        self.n_features = X.shape[1]
        self.input_shape = (self.batch_size, self.n_features)

        self.num_output_units = 1  # Regression
        # Normalize the output
        self.mean_y = np.mean(y)
        self.std_y = np.std(y)
        y = (y - self.mean_y) / self.std_y
        if len(y.shape) == 1:
            y = y[:, np.newaxis]

        self.m_issparse = sp.issparse(X)

        return X, y

    def fit(self, X, y):

        Xf, yf = self._prefit(X, y)

        epoch = (self.number_updates * self.batch_size)//X.shape[0]
        number_epochs = min(max(2, epoch), 110)  # Cap the max number of possible epochs

        from implementation import LogisticRegression
        self.estimator = LogisticRegression.LogisticRegression(batch_size=self.batch_size,
                                                               input_shape=self.input_shape,
                                                               num_output_units=self.num_output_units,
                                                               dropout_output=self.dropout_output,
                                                               learning_rate=self.learning_rate,
                                                               lr_policy=self.lr_policy,
                                                               lambda2=self.lambda2,
                                                               momentum=self.momentum,
                                                               beta1=self.beta1,
                                                               beta2=self.beta2,
                                                               rho=self.rho,
                                                               solver=self.solver,
                                                               num_epochs=number_epochs,
                                                               gamma=self.gamma,
                                                               power=self.power,
                                                               epoch_step=self.epoch_step,
                                                               is_sparse=self.m_issparse,
                                                               is_binary=self.m_isbinary,
                                                               is_multilabel=self.m_ismultilabel,
                                                               is_regression=self.m_isregression)
        self.estimator.fit(Xf, yf)
        return self

    def predict(self, X):
        if self.estimator is None:
            raise NotImplementedError
        preds = self.estimator.predict(X, self.m_issparse)
        return preds * self.std_y + self.mean_y

    def predict_proba(self, X):
        if self.estimator is None:
            raise NotImplementedError()
        return self.estimator.predict_proba(X, self.m_issparse)

    @staticmethod
    def get_properties(dataset_properties=None):
        return {'shortname': 'lin_reg',
                'name': 'Linear Regression',
                'handles_regression': True,
                'handles_classification': False,
                'handles_multiclass': False,
                'handles_multilabel': False,
                'is_deterministic': True,
                'input': (DENSE, SPARSE, UNSIGNED_DATA),
                'output': (PREDICTIONS,)}

    @staticmethod
    def get_hyperparameter_search_space(dataset_properties=None):

        policy_choices = ['fixed', 'inv', 'exp', 'step']

        batch_size = UniformIntegerHyperparameter("batch_size",
                                                  100, 3000,
                                                  log=True,
                                                  default=150)

        number_updates = UniformIntegerHyperparameter("number_updates",
                                                      500, 10500,
                                                      log=True,
                                                      default=500)

        dropout_output = UniformFloatHyperparameter("dropout_output", 0.0, 0.99,
                                                    default=0.5)

        lr = UniformFloatHyperparameter("learning_rate", 1e-6, 0.1,
                                        log=True,
                                        default=0.01)

        l2 = UniformFloatHyperparameter("lambda2", 1e-6, 1e-2,
                                        log=True,
                                        default=1e-3)

        solver = CategoricalHyperparameter(name="solver",
                                           choices=["sgd", "adam"],
                                           default="sgd")

        beta1 = UniformFloatHyperparameter("beta1", 1e-4, 0.1,
                                           log=True,
                                           default=0.1)

        beta2 = UniformFloatHyperparameter("beta2", 1e-4, 0.1,
                                           log=True,
                                           default=0.01)

        lr_policy = CategoricalHyperparameter(name="lr_policy",
                                              choices=policy_choices,
                                              default='fixed')

        gamma = UniformFloatHyperparameter(name="gamma",
                                           lower=1e-3, upper=1e-1,
                                           default=1e-2)

        power = UniformFloatHyperparameter("power",
                                           0.0, 1.0,
                                           default=0.5)

        epoch_step = UniformIntegerHyperparameter("epoch_step",
                                                  2, 20,
                                                  default=5)

        cs = ConfigurationSpace()
        cs.add_hyperparameter(number_updates)
        cs.add_hyperparameter(batch_size)
        cs.add_hyperparameter(dropout_output)
        cs.add_hyperparameter(lr)
        cs.add_hyperparameter(l2)
        cs.add_hyperparameter(solver)
        cs.add_hyperparameter(beta1)
        cs.add_hyperparameter(beta2)
        cs.add_hyperparameter(lr_policy)
        cs.add_hyperparameter(gamma)
        cs.add_hyperparameter(power)
        cs.add_hyperparameter(epoch_step)

        beta1_depends_on_solver = EqualsCondition(beta1, solver, "adam")
        beta2_depends_on_solver = EqualsCondition(beta2, solver, "adam")
        gamma_depends_on_policy = InCondition(child=gamma, parent=lr_policy,
                                              values=['inv', 'exp', 'step'])
        power_depends_on_policy = EqualsCondition(power, lr_policy, 'inv')
        epoch_step_depends_on_policy = EqualsCondition(epoch_step,
                                                       lr_policy, 'step')
        cs.add_condition(beta1_depends_on_solver)
        cs.add_condition(beta2_depends_on_solver)
        cs.add_condition(gamma_depends_on_policy)
        cs.add_condition(power_depends_on_policy)
        cs.add_condition(epoch_step_depends_on_policy)

        return cs

[((52, 22, 52, 32), 'numpy.mean', 'np.mean', ({(52, 30, 52, 31): 'y'}, {}), '(y)', True, 'import numpy as np\n'), ((53, 21, 53, 30), 'numpy.std', 'np.std', ({(53, 28, 53, 29): 'y'}, {}), '(y)', True, 'import numpy as np\n'), ((58, 26, 58, 40), 'scipy.sparse.issparse', 'sp.issparse', ({(58, 38, 58, 39): 'X'}, {}), '(X)', True, 'import scipy.sparse as sp\n'), ((70, 25, 89, 97), 'implementation.LogisticRegression.LogisticRegression', 'LogisticRegression.LogisticRegression', (), '', False, 'from implementation import LogisticRegression\n'), ((121, 21, 124, 62), 'HPOlibConfigSpace.hyperparameters.UniformIntegerHyperparameter', 'UniformIntegerHyperparameter', (), '', False, 'from HPOlibConfigSpace.hyperparameters import UniformFloatHyperparameter, UniformIntegerHyperparameter, CategoricalHyperparameter, Constant\n'), ((126, 25, 129, 66), 'HPOlibConfigSpace.hyperparameters.UniformIntegerHyperparameter', 'UniformIntegerHyperparameter', (), '', False, 'from HPOlibConfigSpace.hyperparameters import UniformFloatHyperparameter, UniformIntegerHyperparameter, CategoricalHyperparameter, Constant\n'), ((131, 25, 132, 64), 'HPOlibConfigSpace.hyperparameters.UniformFloatHyperparameter', 'UniformFloatHyperparameter', (), '', False, 'from HPOlibConfigSpace.hyperparameters import UniformFloatHyperparameter, UniformIntegerHyperparameter, CategoricalHyperparameter, Constant\n'), ((134, 13, 136, 53), 'HPOlibConfigSpace.hyperparameters.UniformFloatHyperparameter', 'UniformFloatHyperparameter', (), '', False, 'from HPOlibConfigSpace.hyperparameters import UniformFloatHyperparameter, UniformIntegerHyperparameter, CategoricalHyperparameter, Constant\n'), ((138, 13, 140, 53), 'HPOlibConfigSpace.hyperparameters.UniformFloatHyperparameter', 'UniformFloatHyperparameter', (), '', False, 'from HPOlibConfigSpace.hyperparameters import UniformFloatHyperparameter, UniformIntegerHyperparameter, CategoricalHyperparameter, Constant\n'), ((142, 17, 144, 57), 'HPOlibConfigSpace.hyperparameters.CategoricalHyperparameter', 'CategoricalHyperparameter', (), '', False, 'from HPOlibConfigSpace.hyperparameters import UniformFloatHyperparameter, UniformIntegerHyperparameter, CategoricalHyperparameter, Constant\n'), ((146, 16, 148, 55), 'HPOlibConfigSpace.hyperparameters.UniformFloatHyperparameter', 'UniformFloatHyperparameter', (), '', False, 'from HPOlibConfigSpace.hyperparameters import UniformFloatHyperparameter, UniformIntegerHyperparameter, CategoricalHyperparameter, Constant\n'), ((150, 16, 152, 56), 'HPOlibConfigSpace.hyperparameters.UniformFloatHyperparameter', 'UniformFloatHyperparameter', (), '', False, 'from HPOlibConfigSpace.hyperparameters import UniformFloatHyperparameter, UniformIntegerHyperparameter, CategoricalHyperparameter, Constant\n'), ((154, 20, 156, 62), 'HPOlibConfigSpace.hyperparameters.CategoricalHyperparameter', 'CategoricalHyperparameter', (), '', False, 'from HPOlibConfigSpace.hyperparameters import UniformFloatHyperparameter, UniformIntegerHyperparameter, CategoricalHyperparameter, Constant\n'), ((158, 16, 160, 56), 'HPOlibConfigSpace.hyperparameters.UniformFloatHyperparameter', 'UniformFloatHyperparameter', (), '', False, 'from HPOlibConfigSpace.hyperparameters import UniformFloatHyperparameter, UniformIntegerHyperparameter, CategoricalHyperparameter, Constant\n'), ((162, 16, 164, 55), 'HPOlibConfigSpace.hyperparameters.UniformFloatHyperparameter', 'UniformFloatHyperparameter', (), '', False, 'from HPOlibConfigSpace.hyperparameters import UniformFloatHyperparameter, UniformIntegerHyperparameter, CategoricalHyperparameter, Constant\n'), ((166, 21, 168, 60), 'HPOlibConfigSpace.hyperparameters.UniformIntegerHyperparameter', 'UniformIntegerHyperparameter', (), '', False, 'from HPOlibConfigSpace.hyperparameters import UniformFloatHyperparameter, UniformIntegerHyperparameter, CategoricalHyperparameter, Constant\n'), ((170, 13, 170, 33), 'HPOlibConfigSpace.configuration_space.ConfigurationSpace', 'ConfigurationSpace', ({}, {}), '()', False, 'from HPOlibConfigSpace.configuration_space import ConfigurationSpace\n'), ((184, 34, 184, 72), 'HPOlibConfigSpace.conditions.EqualsCondition', 'EqualsCondition', ({(184, 50, 184, 55): 'beta1', (184, 57, 184, 63): 'solver', (184, 65, 184, 71): '"""adam"""'}, {}), "(beta1, solver, 'adam')", False, 'from HPOlibConfigSpace.conditions import EqualsCondition, InCondition\n'), ((185, 34, 185, 72), 'HPOlibConfigSpace.conditions.EqualsCondition', 'EqualsCondition', ({(185, 50, 185, 55): 'beta2', (185, 57, 185, 63): 'solver', (185, 65, 185, 71): '"""adam"""'}, {}), "(beta2, solver, 'adam')", False, 'from HPOlibConfigSpace.conditions import EqualsCondition, InCondition\n'), ((186, 34, 187, 76), 'HPOlibConfigSpace.conditions.InCondition', 'InCondition', (), '', False, 'from HPOlibConfigSpace.conditions import EqualsCondition, InCondition\n'), ((188, 34, 188, 74), 'HPOlibConfigSpace.conditions.EqualsCondition', 'EqualsCondition', ({(188, 50, 188, 55): 'power', (188, 57, 188, 66): 'lr_policy', (188, 68, 188, 73): '"""inv"""'}, {}), "(power, lr_policy, 'inv')", False, 'from HPOlibConfigSpace.conditions import EqualsCondition, InCondition\n'), ((189, 39, 190, 73), 'HPOlibConfigSpace.conditions.EqualsCondition', 'EqualsCondition', ({(189, 55, 189, 65): 'epoch_step', (190, 55, 190, 64): 'lr_policy', (190, 66, 190, 72): '"""step"""'}, {}), "(epoch_step, lr_policy, 'step')", False, 'from HPOlibConfigSpace.conditions import EqualsCondition, InCondition\n')]