# -*- coding: utf-8 -*-
import numpy as np


def fit_error(y, y_predicted, n_parameters=2):
    """Calculate the fit error for a model.

    Also specific and direct access functions can be used, such as `fit_mse()`, `fit_rmse()` and `fit_r2()`.

    Parameters
    ----------
    y : Union[list, np.array, pd.Series]
        The response variable (the y axis).
    y_predicted : Union[list, np.array, pd.Series]
        The fitted data generated by a model.
    n_parameters : int
        Number of model parameters (for the degrees of freedom used in R2).


    Returns
    -------
    dict
        A dictionary containing different indices of fit error.

    See Also
    --------
    fit_mse, fit_rmse, fit_r2

    Examples
    --------
    >>> import neurokit2 as nk
    >>>
    >>> y = np.array([-1.0, -0.5, 0, 0.5, 1])
    >>> y_predicted = np.array([0.0, 0, 0, 0, 0])
    >>>
    >>> # Master function
    >>> x = nk.fit_error(y, y_predicted)
    >>> x #doctest: +SKIP
    >>>
    >>> # Direct access
    >>> nk.fit_mse(y, y_predicted) #doctest: +ELLIPSIS
    0.5
    >>>
    >>> nk.fit_rmse(y, y_predicted) #doctest: +ELLIPSIS
    0.7071067811865476
    >>>
    >>> nk.fit_r2(y, y_predicted, adjusted=False) #doctest: +ELLIPSIS
    0.7071067811865475
    >>>
    >>> nk.fit_r2(y, y_predicted, adjusted=True, n_parameters=2) #doctest: +ELLIPSIS
    0.057190958417936755

    """

    # Get information
    SSE, n, df = _fit_error_prepare(y, y_predicted, n_parameters)

    # Mean squared error
    MSE = SSE / n

    # Root mean squared error
    RMSE = np.sqrt(SSE / n)

    # Adjusted r-squared
    # For optimization use 1 - adjR2 since we want to minimize the function
    SST = np.std(y) * n

    # Get R2
    if SST == 0:
        R2 = 1
    else:
        R2 = SSE / SST

    # R2 adjusted
    R2_adjusted = 1 - (1 - (1 - R2)) * (n - 1) / df

    return {"SSE": SSE, "MSE": MSE, "RMSE": RMSE, "R2": R2, "R2_adjusted": R2_adjusted}


# =============================================================================
# Direct accessors
# =============================================================================
def fit_mse(y, y_predicted):
    """Compute Mean Square Error (MSE)."""
    return fit_error(y, y_predicted)["MSE"]


def fit_rmse(y, y_predicted):
    """Compute Root Mean Square Error (RMSE)."""
    return fit_error(y, y_predicted)["RMSE"]


def fit_r2(y, y_predicted, adjusted=True, n_parameters=2):
    """Compute R2."""
    if adjusted is True:
        return fit_error(y, y_predicted, n_parameters=n_parameters)["R2_adjusted"]
    return fit_error(y, y_predicted, n_parameters=n_parameters)["R2"]


# =============================================================================
# Internals
# =============================================================================
def _fit_error_prepare(y, y_predicted, n_parameters=2):

    # n, i.e., how many observations (signal length)
    n = len(y)

    # Sanitize
    if n != len(y_predicted):
        raise TypeError("NeuroKit error: fit_error(): 'y' and 'y_predicted' are not of the same length.")

    # Residual, i.e. the difference between data and model
    residual = y - y_predicted

    # Degrees of freedom, i.e., number of observations (length of signal) minus number of parameters
    df = n - n_parameters

    # Calculate sum of squared errors
    SSE = np.sum(residual ** 2)
    return SSE, n, df
