# -*- coding: utf-8 -*-

import unittest

import pytest
import yaml
import pandas as pd
import numpy as np

from gordo_components.model.transformers.imputer import InfImputer
from gordo_components import serializer

from sklearn.preprocessing import FunctionTransformer
from sklearn.pipeline import Pipeline
from sklearn.decomposition import PCA


class GordoFunctionTransformerFuncsTestCase(unittest.TestCase):
    """
    Test all functions within gordo_components meants for use in a Scikit-Learn
    FunctionTransformer work as expected
    """

    def _validate_transformer(self, transformer):
        """
        Inserts a transformer into the middle of a pipeline and runs it
        """
        pipe = Pipeline([("pca1", PCA()), ("custom", transformer), ("pca2", PCA())])
        X = np.random.random(size=100).reshape(10, 10)
        pipe.fit_transform(X)

    def test_multiply_by_function_transformer(self):
        from gordo_components.model.transformer_funcs.general import multiply_by

        # Provide a require argument
        tf = FunctionTransformer(func=multiply_by, kw_args={"factor": 2})
        self._validate_transformer(tf)

        # Ignore the required argument
        tf = FunctionTransformer(func=multiply_by)
        with self.assertRaises(TypeError):
            self._validate_transformer(tf)


@pytest.mark.parametrize("strategy", ["extremes", "minmax"])
def test_infimputer_basic(strategy):
    """
    Functionality of the InfImputer
    """
    base_x = np.random.random((100, 10)).astype(np.float32)

    flat_view = base_x.ravel()

    pos_inf_idxs = np.random.randint(0, len(flat_view), size=100)
    neg_inf_idxs = np.random.randint(0, len(flat_view), size=100)

    flat_view[pos_inf_idxs] = np.inf
    flat_view[neg_inf_idxs] = -np.inf

    # Our base x should now be littered with pos/neg inf values
    assert np.isposinf(base_x).sum() > 0, "Expected some positive infinity values here"
    assert np.isneginf(base_x).sum() > 0, "Expected some negative infinity values here"

    imputer = InfImputer(strategy=strategy, delta=2.0)

    # Test imputer on numpy array
    X = base_x.copy()
    X = imputer.fit_transform(X)
    assert np.isposinf(X).sum() == 0, "Expected no positive infinity values here"
    assert np.isneginf(X).sum() == 0, "Expected no negative infinity values here"

    if strategy == "extremes":
        # All pos infs in base_x should be filled in X now as the max dtype value
        assert np.all(X[np.where(np.isposinf(base_x))] == np.finfo(X.dtype).max)

        # All neg infs in base_x should be filled in X now as the min dtype value
        assert np.all(X[np.where(np.isneginf(base_x))] == np.finfo(X.dtype).min)

    # min max is a bit more difficult to assert however.
    elif strategy == "minmax":

        # Identify the features in the base array which have pos infs
        features_with_pos_infs = np.where(np.isposinf(base_x))[1]

        # Get the maxes of those features in the imputed X
        filled_maxes = X[:, features_with_pos_infs].max(axis=0)

        # Get the previous maxes
        previous_maxes = np.ma.masked_invalid(base_x[:, features_with_pos_infs]).max(
            axis=0
        )

        # Compare that each new max is the previous max + 2.
        # These were previously 'inf' values, replaced by previous max + 2
        for filled_max, previous_max in zip(filled_maxes, previous_maxes):
            assert np.isclose(filled_max, previous_max + 2.0)

        ### Repeat the process for negative infs ###

        # Identify the features in the base array which have pos infs
        features_with_neg_infs = np.where(np.isneginf(base_x))[1]

        # Get the maxes of those features in the imputed X
        filled_mins = X[:, features_with_neg_infs].min(axis=0)

        # Get the previous maxes
        previous_mins = np.ma.masked_invalid(base_x[:, features_with_neg_infs]).min(
            axis=0
        )

        # Now compare, that each new min is the min - 2. wq
        for filled_min, previous_min in zip(filled_mins, previous_mins):
            assert np.isclose(filled_min, previous_min - 2.0)

    # Test imputer on pandas dataframe
    X = pd.DataFrame(base_x.copy())
    X = imputer.fit_transform(X)
    assert np.isposinf(X).sum() == 0, "Expected no positive infinity values here"
    assert np.isneginf(X).sum() == 0, "Expected no negative infinity values here"


def test_infimputer_fill_values():
    """
    InfImputer when fill values are provided
    """
    base_x = np.random.random((100, 10)).astype(np.float32)

    flat_view = base_x.ravel()

    pos_inf_idxs = [1, 2, 3, 4, 5]
    neg_inf_idxs = [6, 7, 8, 9, 10]

    flat_view[pos_inf_idxs] = np.inf
    flat_view[neg_inf_idxs] = -np.inf

    # Our base x should now be littered with pos/neg inf values
    assert np.isposinf(base_x).sum() > 0, "Expected some positive infinity values here"
    assert np.isneginf(base_x).sum() > 0, "Expected some negative infinity values here"

    imputer = InfImputer(inf_fill_value=9999.0, neg_inf_fill_value=-9999.0)
    X = imputer.fit_transform(base_x)
    np.equal(
        X.ravel()[[pos_inf_idxs]], np.array([9999.0, 9999.0, 9999.0, 9999.0, 9999.0])
    )
    np.equal(
        X.ravel()[[neg_inf_idxs]],
        np.array([-9999.0, -9999.0, -9999.0, -9999.0, -9999.0]),
    )


@pytest.mark.parametrize(
    "config_str",
    [
        """
    sklearn.pipeline.Pipeline:
      steps:
        - gordo_components.model.transformers.imputer.InfImputer
    """,
        """
    sklearn.pipeline.Pipeline:
      steps:
        - gordo_components.model.transformers.imputer.InfImputer:
            inf_fill_value: 10
    """,
        """gordo_components.model.transformers.imputer.InfImputer""",
    ],
)
def test_imputer_from_definition(config_str: str):
    """
    Ensure it plays well with the gordo serializer
    """
    config = yaml.safe_load(config_str)
    model = serializer.pipeline_from_definition(config)

    if isinstance(model, Pipeline):
        assert isinstance(model.steps[-1][1], InfImputer)
    else:
        assert isinstance(model, InfImputer)

    serializer.pipeline_from_definition(serializer.pipeline_into_definition(model))
