import pandas as pd
import numpy as np
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union

import ipytest
import pytest
from typeguard import typechecked
from skmultilearn.model_selection.iterative_stratification import IterativeStratification

# enables ipytest notebook magic
ipytest.autoconfig()

logger = logging.getLogger(__name__)

@typechecked
def stratify_shuffle_split_subsets(
    full_dataset: Dict[str, Callable], output_partition_name: str, train_fraction: float
) -> Tuple[pd.DataFrame, Dict[str, pd.DataFrame], Dict[str, pd.DataFrame]]:
    """
    Stratify-shuffle-split the a multi-class multi-label dataset into
    train/dev/validation sets.

    Args:
        full_dataset: the full supervised dataset. One column is the img uris, and the rest are binary labels.
        output_partition_name: the name of the output partition
        train_fraction: the fraction of data to reserve for the training dataset. The remaining data will be evenly
            split into the dev and validation subsets.

    Returns:
        the supervised dataset, split into train/test/dev subsets. The train subset is not saved as a partitionedDataset
        yet because we are doing label smoothing in the next processing step.
    """
    if not 0 < train_fraction < 1:
        raise ValueError(f"Training fraction must be value between 0 and 1, got {train_fraction}")

    loader = full_dataset[output_partition_name]
    full_dataset = loader()
    original_cols = full_dataset.columns
    # pandas documentation says to use .to_numpy() instead of .values for consistency
    img_urls = full_dataset["img_url"].to_numpy()

    # sanity check: no duplicate labels
    if not len(img_urls) == len(set(img_urls)):
        raise ValueError("Duplicate image keys detected.")

    labels = full_dataset.drop(columns=["img_url"]).to_numpy().astype(int)
    # NOTE generators are replicated across workers. do stratified shuffle split beforehand
    logger.info("stratifying dataset iteratively. this may take a while.")
    # NOTE: splits >2 broken; https://github.com/scikit-multilearn/scikit-multilearn/issues/209
    # so, do 2 rounds of iterative splitting
    stratifier = IterativeStratification(
        n_splits=2, order=2, sample_distribution_per_fold=[1.0 - train_fraction, train_fraction],
    )
    # this class is a generator that produces k-folds. we just want to iterate it once to make a single static split
    # NOTE: needs to be computed on hard labels.
    train_indexes, everything_else_indexes = next(stratifier.split(X=img_urls, y=labels))

    num_overlapping_samples = len(set(train_indexes).intersection(set(everything_else_indexes)))
    if num_overlapping_samples != 0:
        raise ValueError(f"First splitting failed, {num_overlapping_samples} overlapping samples detected")

    # s3url array shape (N_samp,)
    x_train, x_else = img_urls[train_indexes], img_urls[everything_else_indexes]
    # labels array shape (N_samp, n_classes)
    Y_train, Y_else = labels[train_indexes, :], labels[everything_else_indexes, :]

    # now, split the "else" data evenly into dev/val splits
    stratifier = IterativeStratification(n_splits=2, order=2)  # splits evenly by default
    dev_indexes, validation_indexes = next(stratifier.split(x_else, Y_else))

    num_overlapping_samples = len(set(dev_indexes).intersection(set(validation_indexes)))
    if num_overlapping_samples != 0:
        raise ValueError(f"Second splitting failed, {num_overlapping_samples} overlapping samples detected")

    x_dev, x_val = (x_else[dev_indexes], x_else[validation_indexes])
    Y_dev, Y_val = (Y_else[dev_indexes, :], Y_else[validation_indexes, :])

    for subset_name, frac, encodings_collection in [
        ("train", train_fraction, Y_train),
        ("dev", (1.0 - train_fraction) / 2, Y_dev),
        ("val", (1.0 - train_fraction) / 2, Y_val),
    ]:
        # column-wise sum. sum(counts) > n_samples due to imgs with >1 class
        count_values = np.sum(encodings_collection, axis=0)
        # skip first col, which is the image key, not a class ID
        counts = {class_id: count_val for class_id, count_val in zip(full_dataset.columns[1:], count_values)}
        logger.info(f" {subset_name} subset ({frac * 100:.1f}%) encodings counts after stratification: {counts}")

    # combine (x,y) data into dataframes
    train_subset = pd.DataFrame(Y_train)
    train_subset.insert(0, "img_url", pd.Series(x_train))
    train_subset.columns = original_cols

    dev_subset = pd.DataFrame(Y_dev)
    dev_subset.insert(0, "img_url", pd.Series(x_dev))
    dev_subset.columns = original_cols

    val_subset = pd.DataFrame(Y_val)
    val_subset.insert(0, "img_url", pd.Series(x_val))
    val_subset.columns = original_cols

    # TODO create Great Expectations suite for this node instead
    assert "img_url" in dev_subset.columns
    assert "img_url" in val_subset.columns
    assert "img_url" in train_subset.columns

    return (
        train_subset,
        {output_partition_name: dev_subset},
        {output_partition_name: val_subset},
    )

@pytest.fixture
def unsplit_supervised_learning_dataset_partition_name():
    """PartitionName to simulate Kedro PartitionDataset"""
    return "pytest_unsplit_supervised_learning_dataset_partition_name"


@pytest.fixture
def unsplit_supervised_learning_dataset(unsplit_supervised_learning_dataset_partition_name: str) -> Dict[str, Callable]:
    """
    Generates a multi-class multilabel dataset.

    Args:
        unsplit_supervised_learning_dataset_partition_name:

    Returns:
        a multi-class multi-label df, packed in a lambda to mimic Kedro PartitionDataset
    """
    df_rows = []

    n_total = 60
    n_created = 0
    n_class1 = n_total // 3
    n_class2 = n_total // 3
    n_class2_plus_3 = n_total // 3

    for _ in range(n_class1):
        df_rows.append({"img_url": f"s3://test{n_created}", "class0": 1, "class1": 0, "class2": 0})
        n_created += 1

    for _ in range(n_class2):
        df_rows.append({"img_url": f"s3://test{n_created}", "class0": 0, "class1": 1, "class2": 0})
        n_created += 1

    for _ in range(n_class2_plus_3):
        df_rows.append({"img_url": f"s3://test{n_created}", "class0": 0, "class1": 1, "class2": 1})
        n_created += 1

    df = pd.DataFrame(df_rows)
    data_loader = lambda: df

    return {unsplit_supervised_learning_dataset_partition_name: data_loader}


@pytest.mark.parametrize("train_fraction", [-1.0, 0, 1, 5])
def test_train_fraction_range(
    unsplit_supervised_learning_dataset, train_fraction, unsplit_supervised_learning_dataset_partition_name
):
    """
    Tests that train_fraction is validated between 0, 1.
    """
    with pytest.raises(ValueError):
        stratify_shuffle_split_subsets(
            full_dataset=unsplit_supervised_learning_dataset,
            output_partition_name=unsplit_supervised_learning_dataset_partition_name,
            train_fraction=train_fraction,
        )


def test_duplicate_uris_detected(unsplit_supervised_learning_dataset_partition_name):
    """Tests whether duplicate image keys are properly detected"""
    # setup

    df = pd.DataFrame(
        [
            {"img_url": "s3://duplicate1", "class0": 1, "class1": 0, "class2": 0},
            {"img_url": "s3://duplicate1", "class0": 1, "class1": 0, "class2": 0},
            {"img_url": "s3://duplicate1", "class0": 1, "class1": 0, "class2": 0},
            {"img_url": "s3://duplicate2", "class0": 0, "class1": 1, "class2": 0},
            {"img_url": "s3://duplicate2", "class0": 0, "class1": 1, "class2": 0},
        ]
    )
    data_loader = lambda: df
    partition_dset = {unsplit_supervised_learning_dataset_partition_name: data_loader}

    # test
    with pytest.raises(ValueError):
        stratify_shuffle_split_subsets(
            full_dataset=partition_dset,
            output_partition_name=unsplit_supervised_learning_dataset_partition_name,
            train_fraction=0.6,
        )

        
def test_col_names_same(unsplit_supervised_learning_dataset, unsplit_supervised_learning_dataset_partition_name):
    """
    Tests that train_fraction is validated between 0, 1.
    """
    # setup
    train_split: pd.DataFrame
    dev_split: Dict[str, pd.DataFrame]
    val_split: Dict[str, pd.DataFrame]
    # set train fraction to 0.6 to make math easier later
    train_fraction = 0.6
    data_loader = unsplit_supervised_learning_dataset[unsplit_supervised_learning_dataset_partition_name]
    original_df = data_loader()
    original_df_cols = original_df.columns

    # run
    train_df, dev_split, val_split = stratify_shuffle_split_subsets(
        full_dataset=unsplit_supervised_learning_dataset,
        output_partition_name=unsplit_supervised_learning_dataset_partition_name,
        train_fraction=train_fraction,
    )

    # test setup
    dev_df, val_df = (
        dev_split[unsplit_supervised_learning_dataset_partition_name],
        val_split[unsplit_supervised_learning_dataset_partition_name],
    )
    train_df_cols, dev_df_cols, val_df_cols = train_df.columns, dev_df.columns, val_df.columns

    # tests
    assert all(train_df_cols == original_df_cols)
    assert all(dev_df_cols == original_df_cols)
    assert all(val_df_cols == original_df_cols)

    
def test_stratify_shuffle_split_subsets(
    unsplit_supervised_learning_dataset, unsplit_supervised_learning_dataset_partition_name
):
    """Tests whether a multi-class multi-label dataset gets properly stratify-shuffle-split."""
    # setup
    train_split: pd.DataFrame
    dev_split: Dict[str, pd.DataFrame]
    val_split: Dict[str, pd.DataFrame]
    # set train fraction to 0.6 to make math easier later
    train_fraction = 0.6
    ratio_train_to_val = train_fraction / ((1.0 - train_fraction) / 2)
    data_loader = unsplit_supervised_learning_dataset[unsplit_supervised_learning_dataset_partition_name]
    original_dset_len = len(data_loader())

    # run
    train_df, dev_split, val_split = stratify_shuffle_split_subsets(
        full_dataset=unsplit_supervised_learning_dataset,
        output_partition_name=unsplit_supervised_learning_dataset_partition_name,
        train_fraction=train_fraction,
    )

    # test setup
    dev_df, val_df = (
        dev_split[unsplit_supervised_learning_dataset_partition_name],
        val_split[unsplit_supervised_learning_dataset_partition_name],
    )
    train_keys, dev_keys, val_keys = train_df["img_url"], dev_df["img_url"], val_df["img_url"]

    uniq_train_keys, uniq_dev_keys, uniq_val_keys = set(train_keys), set(dev_keys), set(val_keys)

    # Tests

    ## test keys

    # test no duplicate keys within subsets
    assert len(train_keys) == len(uniq_train_keys)
    assert len(val_keys) == len(uniq_val_keys)
    assert len(dev_keys) == len(uniq_dev_keys)

    # test that all subsets have mutually exclusive S3 URIs
    assert len(uniq_train_keys.intersection(uniq_dev_keys)) == 0
    assert len(uniq_train_keys.intersection(uniq_val_keys)) == 0
    assert len(uniq_val_keys.intersection(uniq_dev_keys)) == 0

    ## test ratios
    # test that dev and val subsets get same number of samples
    assert len(dev_df) == len(val_df)
    # test that test set is 2x size of dev
    assert len(train_df) / (len(train_df) + len(dev_df) + len(val_df)) == pytest.approx(train_fraction)
    # test that all samples are used
    assert len(train_df) + len(dev_df) + len(val_df) == original_dset_len

    ## test distributions

    ### class 0
    # test whether dev/val set got same amount of class 0
    assert len(dev_df[dev_df.class0 == 1]) == len(val_df[val_df.class0 == 1])
    # test whether train set got 2x class0 as val set
    assert len(train_df[train_df.class0 == 1]) / len(val_df[val_df.class0 == 1]) == pytest.approx(ratio_train_to_val)

    ### class 1
    # test whether dev/val set got same amount of class 1
    assert len(dev_df[dev_df.class1 == 1]) == len(val_df[val_df.class1 == 1])
    # test whether train set got 2x class1 as val set
    assert len(train_df[train_df.class1 == 1]) / len(val_df[val_df.class1 == 1]) == pytest.approx(ratio_train_to_val)

    ### class 2
    # test whether dev/val set got same amount of class 2
    assert len(dev_df[dev_df.class2 == 1]) == len(val_df[val_df.class2 == 1])
    # test whether train set got 2x class1 as val set
    assert len(train_df[train_df.class2 == 1]) / len(val_df[val_df.class2 == 1]) == pytest.approx(ratio_train_to_val)

ipytest.run()