from __future__ import annotations
from copy import deepcopy
from typing import Tuple, Callable
import numpy as np
from IMLearn import BaseEstimator


def cross_validate(estimator: BaseEstimator, X: np.ndarray, y: np.ndarray,
                   scoring: Callable[[np.ndarray, np.ndarray, ...], float], cv: int = 5) -> Tuple[float, float]:
    """
    Evaluate metric by cross-validation for given estimator

    Parameters
    ----------
    estimator: BaseEstimator
        Initialized estimator to use for fitting the data

    X: ndarray of shape (n_samples, n_features)
       Input data to fit

    y: ndarray of shape (n_samples, )
       Responses of input data to fit to

    scoring: Callable[[np.ndarray, np.ndarray, ...], float]
        Callable to use for evaluating the performance of the cross-validated model.
        When called, the scoring function receives the true- and predicted values for each sample
        and potentially additional arguments. The function returns the score for given input.

    cv: int
        Specify the number of folds.

    Returns
    -------
    train_score: float
        Average train score over folds

    validation_score: float
        Average validation score over folds
    """
    rows, cols = X.shape
    group_length = int(rows / cv)
    train_error = []
    test_error = []
    for i in range(cv):
        fold_delimiter = (i + 1) * group_length
        S_divided_by_S_i_X = np.concatenate((X[:(fold_delimiter - group_length)], X[fold_delimiter:]))
        S_divided_by_S_i_y = np.concatenate((y[:(fold_delimiter - group_length)], y[fold_delimiter:]))
        estimator.fit(S_divided_by_S_i_X, S_divided_by_S_i_y)
        train_pred = estimator.predict(S_divided_by_S_i_X)
        test_pred = estimator.predict(X[(fold_delimiter - group_length):fold_delimiter])
        train_error.append(scoring(S_divided_by_S_i_y, train_pred))
        test_error.append(scoring(y[(fold_delimiter - group_length):fold_delimiter], test_pred))
    return np.average(train_error), np.average(test_error)
