from collections import Counter, defaultdict
from itertools import chain
from typing import Generator, Tuple, Union
import numpy as np

from luma.interface.util import Matrix, Vector

FoldType = Union["KFold", "StratifiedKFold", "GroupKFold"]


__all__ = ("KFold", "StratifiedKFold", "GroupKFold")


class KFold:
    """
    K-Fold cross-validation is a model evaluation method that divides
    the dataset into `k` equal or nearly equal sized folds. In each
    of `k` iterations, one fold is used as the test set, and the
    remaining `k-1` folds are combined to form the training set. This
    process ensures that every data point gets to be in the test set
    exactly once and in the training set `k-1` times. It's widely used
    to assess the performance of a model with limited data, providing
    a robust estimate of its generalization capability.

    Parameters
    ----------
    `X` : Input data
    `y` : Target data
    `n_folds` : Number of folds
    `shuffle` : Whether to shuffle the dataset
    `random_state` : Seed for random shuffling

    Properties
    ----------
    ```py
    @property
    def split(self)
    ```
    This returns a generator with the type:
    ```py
    Generator[Tuple[Vector, Vector], None, None]
    ```
    yielding train indices and test indices.

    Examples
    --------
    Usage of the generator returned by the property `split`:

    ```py
    kfold = KFold(X, y, n_folds=5, shuffle=True)

    for train_indices, test_indices in kfold.split:
        X_train, y_train = X[train_indices], y[train_indices]
        ...
    ```
    """

    def __init__(
        self,
        X: Matrix,
        y: Vector,
        n_folds: int = 5,
        shuffle: bool = True,
        random_state: int = None,
    ) -> None:
        self.X = X
        self.y = y
        self.shuffle = shuffle
        self.random_state = random_state
        self.n_folds = n_folds

        self.set_param_ranges({"n_folds": ("0,+inf", int)})
        self.check_param_ranges()

    @property
    def split(self) -> Generator[Tuple[Vector, Vector], None, None]:
        np.random.seed(self.random_state)
        m, _ = self.X.shape
        fold_size = m // self.n_folds

        indices = np.arange(m)
        if self.shuffle:
            np.random.shuffle(indices)

        for i in range(self.n_folds):
            start = i * fold_size
            end = start + fold_size if i < self.n_folds - 1 else m

            test_indices = indices[start:end]
            train_indices = np.concatenate((indices[:start], indices[end:]))

            yield train_indices, test_indices


class StratifiedKFold:
    """
    Stratified K-Fold cross-validation is a model evaluation method similar
    to KFold but it divides the dataset in a way that preserves the
    percentage of samples for each class. This is especially useful for
    handling imbalances in the dataset. It ensures each fold is a good
    representative of the whole by maintaining the same class distribution
    in each fold as in the complete dataset.

    Parameters
    ----------
    `X` : Input data
    `y` : Target data
    `n_folds` : Number of folds
    `shuffle` : Whether to shuffle the dataset
    `random_state` : Seed for random shuffling

    Properties
    ----------
    ```py
    @property
    def split(self)
    ```
    This returns a generator with the type:
    ```py
    Generator[Tuple[Vector, Vector], None, None]
    ```
    yielding train indices and test indices.

    Examples
    --------
    Usage of the generator returned by the property `split`:

    ```py
    kfold = StratifiedKFold(X, y, n_folds=5, shuffle=True)

    for train_indices, test_indices in kfold.split:
        X_train, y_train = X[train_indices], y[train_indices]
        ...
    ```
    """

    def __init__(
        self,
        X: Matrix,
        y: Vector,
        n_folds: int = 5,
        shuffle: bool = True,
        random_state: int = None,
    ) -> None:
        self.X = X
        self.y = y
        self.n_folds = n_folds
        self.shuffle = shuffle
        self.random_state = random_state

        self.set_param_ranges({"n_folds": ("0,+inf", int)})
        self.check_param_ranges()

    @property
    def split(self) -> Generator[Tuple[Vector, Vector], None, None]:
        np.random.seed(self.random_state)

        indices_per_class = defaultdict(list)
        for idx, class_ in enumerate(self.y):
            indices_per_class[class_].append(idx)

        if self.shuffle:
            for indices in indices_per_class.values():
                np.random.shuffle(indices)

        folds = defaultdict(list)
        for class_, indices in indices_per_class.items():
            fold_size = len(indices) // self.n_folds
            for i in range(self.n_folds):
                start = i * fold_size
                end = start + fold_size if i < self.n_folds - 1 else len(indices)
                folds[i].extend(indices[start:end])

        if self.shuffle:
            for fold_indices in folds.values():
                np.random.shuffle(fold_indices)

        for i in range(self.n_folds):
            test_indices = folds[i]
            train_indices = list(
                chain.from_iterable(folds[j] for j in range(self.n_folds) if j != i)
            )

            yield Vector(train_indices), Vector(test_indices)


class GroupKFold:
    """
    Group K-Fold cross-validation provides train/test indices to split data into
    train/test sets. Each set contains group samples that are entirely in the set
    of training or test. This cross-validation object is a variation of KFold
    that ensures samples from the same group are not split between training and
    testing sets.

    Parameters
    ----------
    `X` : Input data
    `y` : Target data
    `groups` : Array of group labels for the samples
    `n_folds` : Number of folds. Must be at least 2.
    `shuffle` : Whether to shuffle the dataset
    `random_state` : Seed for random shuffling

    Properties
    ----------
    ```py
    @property
    def split(self)
    ```
    This returns a generator with the type:
    ```py
    Generator[Tuple[Vector, Vector], None, None]
    ```
    yielding train indices and test indices.

    Examples
    --------
    Usage of the generator returned by the property `split`:

    ```py
    kfold = GroupKFold(X, y, n_folds=5, shuffle=True)

    for train_indices, test_indices in kfold.split:
        X_train, y_train = X[train_indices], y[train_indices]
        ...
    ```
    """

    def __init__(
        self,
        X: Matrix,
        y: Vector,
        groups: Vector,
        n_folds: int = 5,
        shuffle: bool = True,
        random_state: int = None,
    ) -> None:
        self.X = X
        self.y = y
        self.groups = groups
        self.n_folds = n_folds
        self.shuffle = shuffle
        self.random_state = random_state

        self.set_param_ranges({"n_folds": ("0,+inf", int)})
        self.check_param_ranges()

    @property
    def split(self) -> Generator[Tuple[Vector, Vector], None, None]:
        np.random.seed(self.random_state)
        _, groups_indices = np.unique(self.groups, return_inverse=True)

        group_counts = Counter(groups_indices)
        group_indices = defaultdict(list)
        for idx, group_idx in enumerate(groups_indices):
            group_indices[group_idx].append(idx)

        sorted_groups = sorted(group_counts, key=group_counts.get, reverse=True)

        if self.shuffle:
            np.random.shuffle(sorted_groups)

        folds = [[] for _ in range(self.n_folds)]
        for group in sorted_groups:
            smallest = min(range(self.n_folds), key=lambda x: len(folds[x]))
            folds[smallest].extend(group_indices[group])

        for i in range(self.n_folds):
            test_indices = folds[i]
            train_indices = list(set(range(len(self.X))) - set(test_indices))

            yield Vector(train_indices), Vector(test_indices)
