import numpy as _np
import pandas as _pd
import tensorflow as _tf
import tensorflow.keras as _ks
import so_ml_tools as _soml


def is_multiclass_classification(y_prob: any) -> bool:
    """Return True if the y value is a multiclass classification where the values represent the probability of the result
    to be that specific class, for example, when we have three possible outcomes:
    [
        [0.33, 0.78, 0.23],
        [0.64, 0.32, 0.11]
    ]

    Args:
        y_prob: the label to check if it is a multiclass classification prediction

    Returns:
        True if it is a multiclass classification, False if not.

    Raises:
        TypeError: If `y_prob` is neither a 'list', 'tf.Tensor' or 'np.array'.
    """
    assert y_prob is not None, "y_prob is null"

    if isinstance(y_prob, _pd.DataFrame):
        y_prob = y_prob.to_numpy()
    elif isinstance(y_prob, _pd.Series):
        y_prob = _np.expand_dims(y_prob.to_numpy(), axis=-1)

    if isinstance(y_prob, list):
        return len(y_prob) > 0 and (isinstance(y_prob[0], list) or isinstance(y_prob[0], _np.ndarray)) and len(y_prob[0]) > 1
    elif _tf.is_tensor(y_prob):
        return y_prob.get_shape().ndims == 2 and y_prob.shape[1] > 1
    # elif isinstance(y_prob, _pd.DataFrame):
    #     return y_prob.ndims == 2 and y_prob.shape[1] > 1
    elif isinstance(y_prob, _np.ndarray):
        return y_prob.ndim == 2 and y_prob.shape[1] > 1

    raise TypeError('y should be of type tf.Tensor, np.array or pd.DataFrame.')


def is_binary_classification(y_prob: any) -> bool:
    """Return True if the y value is a binary classification prediction where the value represents a probability between 0 and 1,
    for example, when we have three different predictions:
    [
        [0.56],
        [0.32],
        [0.98]
    ]

    Args:
        y_prob: the label to check if it is a binary classification prediction

    Returns:
        True if it is a binary classification, False if not.

    Raises:
        TypeError: If `y_prob` is neither a  'list', 'tf.Tensor' or 'np.array'.
    """
    assert y_prob is not None, "y_prob is null"

    if isinstance(y_prob, _pd.DataFrame):
        y_prob = y_prob.to_numpy()
    elif isinstance(y_prob, _pd.Series):
        y_prob = _np.expand_dims(y_prob.to_numpy(), axis=-1)

    if isinstance(y_prob, list):
        if len(y_prob) > 0 and (isinstance(y_prob[0], list) or isinstance(y_prob[0], _np.ndarray)):
            if len(y_prob[0]) == 1: # [[1], [0], [1]]
                if 0 <= _np.max(y_prob) <= 2:
                    return True
        elif len(y_prob) > 0: # [0, 1, 1, 0, 1]
            if 0 <= max(y_prob) <= 2:
                return True

        return False
    if _tf.is_tensor(y_prob):
        if y_prob.get_shape().ndims == 2 and y_prob.shape[1] == 1 and 0 <= _tf.math.reduce_max(y_prob) <= 2:  # [[1], [0], [1]]
            return True
        elif y_prob.get_shape().ndims == 1 and 0 <= _tf.math.reduce_max(y_prob) <= 2:  # [0, 1, 1, 0, 1]
            return True

        return False
    # elif isinstance(y_prob, _pd.DataFrame):
    #     if y_prob.ndims == 2 and y_prob.shape[1] == 1 and 0 <= y_prob.max() <= 2:  # [[1], [0], [1]]
    #         return True
    #     elif y_prob.ndims == 1 and 0 <= y_prob.max() <= 2:  # [0, 1, 1, 0, 1]
    #         return True
    #
    #     return False
    elif isinstance(y_prob, _np.ndarray):
        if y_prob.ndim == 2 and y_prob.shape[1] == 1 and 0 <= _np.max(y_prob) <= 2:  # [[1], [0], [1]]
            return True
        elif y_prob.ndim == 1 and 0 <= _np.max(y_prob) <= 2:  # [0, 1, 1, 0, 1]
            return True

        return False

    raise TypeError('y_prob should be of type tf.Tensor, np.array, list or pd.DataFrame, pd.Series.')


def probability_to_class(y_prob: any) -> any:
    """Converts multiclass dense label predictions to sparse labels:

    Probability:
    [
        [0.33, 0.78, 0.23],
        [0.64, 0.32, 0.11],
        [0.22, 0.36, 0.76]
    ]

    Class:
    [1, 0, 2]

    Args:
        y_prob: the probabilities matrix either a 'list', 'tf.Tensor' or a 'np.array'

    Returns:
        The sparse encoded classes

    Raises:
        TypeError: If `y_prob` is neither a 'list', 'tf.Tensor' or 'np.array'.
    """
    assert y_prob is not None, "y_prob is null"

    if isinstance(y_prob, list):
        return list(_np.argmax(y_prob, axis=1))
    elif _tf.is_tensor(y_prob):
        return _tf.argmax(y_prob, axis=1)
    elif isinstance(y_prob, _pd.DataFrame):
        return y_prob.arg.idxmax(axis=1)
    elif isinstance(y_prob, _np.ndarray):
        return _np.argmax(y_prob, axis=1)

    raise TypeError('y_prob should be of type tf.Tensor, np.array, list or pd.DataFrame.')


def probability_to_binary(y_prob: any) -> any:
    """
    Converts scaled classification prediction labels to binarized classification prediction labels:

    Scaled:
    [
        [0.56],
        [0.32],
        [0.98]
    ]

    Ordinal:
    [
        [1],
        [0],
        [1]
    ]

    Args:
        y_prob: the probabilities matrix either a 'list', 'tf.Tensor', 'pd.DataFrame' or a 'np.array'

    Returns:
        Binarized prediction labels

    Raises:
        TypeError: If `y_prob` is neither a 'list', 'tf.Tensor', 'pd.DataFrame' or an 'np.array'
    """
    assert y_prob is not None, "y_prob is null"

    if isinstance(y_prob, list):
        return list(_np.round(y_prob))
    if _tf.is_tensor(y_prob):
        y1 = _tf.round(y_prob)
        return _tf.cast(y1, dtype=_tf.int8)
    elif isinstance(y_prob, _pd.DataFrame):
        y1 = y_prob.round(decimals=0)
        return y1.astype(_np.int64)
    elif isinstance(y_prob, _np.ndarray):
        y1 = _np.round(a=y_prob, decimals=0)
        return y1.astype(_np.int64)

    raise TypeError('y_prob should be of type tf.Tensor, np.array, list or pd.DataFrame.')


def to_prediction_custom_threshold(y_prob: any, custom_threshold: float = 0.5) -> any:
    """
    Does the same as the to_prediction method except that it allows for a custom threshold, by default
    the to_prediction method will use the standard rounding function, meaning a threshold of 0.5, using this
    method you can define your own threshold.

    Args:
        y_prob: the probabilities matrix either a 'list' or an 'np.array'
        custom_threshold: threshold value for when the result is either a 0 or a 1 ( >= threshold)

    Raises:
        TypeError: If `y_prob` is neither a 'list' nor 'np.array'
    """
    if is_binary_classification(y_prob=y_prob):
        if isinstance(y_prob, list):
            return [[1.0] if y[0] >= custom_threshold else [0.0] for y in y_prob]
        elif isinstance(y_prob, _np.ndarray):
            return _np.where(y_prob > custom_threshold, 1, 0)

    raise TypeError('y_prob should be of type np.array or list.')


def to_prediction(y_prob: any) -> any:
    """
    Determines if the probability is a multiclass or binary classification and then will
    return the prediction in either class or binary form.

    Args:
        y_prob: the probabilities matrix either a 'list', 'tf.Tensor' or a 'np.array'

    Return:
        A tensor or numpy array with the prediction.
    """
    assert y_prob is not None, "y_prob is null"

    if is_multiclass_classification(y_prob=y_prob):
        y_prob = probability_to_class(y_prob=y_prob)
    elif is_binary_classification(y_prob=y_prob):
        y_prob = probability_to_binary(y_prob=y_prob)

    return y_prob


def is_one_hot(x):
    """
    Check if the given x is a one-hot encoded value.

    Args:
        x: the value to check
    """
    return (x.sum(axis=1)-_np.ones(x.shape[0])).sum() == 0


def to_categorical(y, num_classes=None, dtype=None):
    """
    Converts a class vector (integers) to binary class matrix.

    E.g. for use with `categorical_crossentropy`.

    Args:
        y: Array-like with class values to be converted into a matrix
            (integers from 0 to `num_classes - 1`).
        num_classes: Total number of classes. If `None`, this would be inferred
          as `max(y) + 1`.
        dtype: The data type expected by the input. Default: `'float32'`.

    Returns:
        A binary matrix representation of the input as a NumPy array. The class
        axis is placed last.
    """
    return _ks.utils.to_categorical(y, num_classes, dtype=dtype)


def to_ordinal(y, dtype=None):
    """
        Converts a matrix of categorical values to ordinal values:

        For example:

        [
            [0., 1., 0., 0., 0.],
            [0., 0., 1., 0., 0.],
            [0., 0., 0., 1., 0.],
            [0., 0., 0., 0., 1.]
        ]

        To:

        [1, 2, 3, 4]
    """
    if dtype is None:
        return _np.argmax(y, axis=1).astype(dtype)

    return _np.argmax(y, axis=1)


def diff_indexes(y_true, y_pred=None, y_prob=None):
    if isinstance(y_true, _tf.data.Dataset):
        y_true = _soml.tf.dataset.get_labels(dataset=y_true)

    y_true = to_prediction(y_prob=y_true)
    if y_pred is None and y_prob is not None:
        y_pred = to_prediction(y_prob=y_prob)
    elif y_pred is None and y_prob is None:
        raise "Must specify 'y_pred' or 'y_prob'"

    compare_to = y_pred

    # Check if the shapes match, if not, try to fix it.
    if y_pred.shape != y_true.shape:
        compare_to = _np.reshape(y_pred, newshape=y_true.shape)

    compared_result = y_true == compare_to
    return _np.where(compared_result == False)[0]