import numpy as np
import typing
from collections import defaultdict


def kfold_split(
    num_objects: int, num_folds: int
) -> list[tuple[np.ndarray, np.ndarray]]:
    all_indices = np.arange(num_objects)
    fold_size = num_objects // num_folds

    splits = []
    for fold_idx in range(num_folds):
        fold_start = fold_idx * fold_size
        fold_end = (
            num_objects if fold_idx == num_folds - 1 else (fold_idx + 1) * fold_size
        )

        validation = all_indices[fold_start:fold_end]
        training = np.concatenate([all_indices[:fold_start], all_indices[fold_end:]])

        splits.append((training, validation))

    return splits


def knn_cv_score(
    X: np.ndarray,
    y: np.ndarray,
    parameters: dict[str, list],
    score_function: callable,
    folds: list[tuple[np.ndarray, np.ndarray]],
    knn_class: object,
) -> dict[str, float]:
    cv_results = {}

    normalizer_configs = parameters.get("normalizers", [(None, None)])
    neighbor_counts = parameters.get("n_neighbors", [5])
    distance_metrics = parameters.get("metrics", ["euclidean"])
    weight_schemes = parameters.get("weights", ["uniform"])

    for norm_obj, norm_label in normalizer_configs:
        for num_neighbors in neighbor_counts:
            for distance_metric in distance_metrics:
                for weight_scheme in weight_schemes:
                    scores_per_fold = []

                    for train_indices, val_indices in folds:
                        X_tr, X_va = X[train_indices], X[val_indices]
                        y_tr, y_va = y[train_indices], y[val_indices]

                        if norm_obj is not None:
                            fitted_normalizer = norm_obj.fit(X_tr)
                            X_tr = fitted_normalizer.transform(X_tr)
                            X_va = fitted_normalizer.transform(X_va)

                        classifier = knn_class(
                            n_neighbors=num_neighbors,
                            metric=distance_metric,
                            weights=weight_scheme,
                        )
                        classifier.fit(X_tr, y_tr)
                        predictions = classifier.predict(X_va)
                        fold_score = score_function(y_va, predictions)
                        scores_per_fold.append(fold_score)

                    avg_score = np.mean(scores_per_fold)
                    param_key = (
                        norm_label,
                        num_neighbors,
                        distance_metric,
                        weight_scheme,
                    )
                    cv_results[param_key] = avg_score

    return cv_results