"""
Scripts to test cleanlab usage with various ML frameworks:
pytorch, skorch
"""

import pytest
import warnings

# pytest.mark.filterwarnings is unable to catch filterbuffers library DeprecationWarning
warnings.filterwarnings(action="ignore", category=DeprecationWarning)

import os
from copy import deepcopy
import random
import numpy as np

import torch
import skorch
import sklearn

from cleanlab.classification import CleanLearning

# Version check for sklearn compatibility
uses_sklearn_ge_1_6_0 = tuple(map(int, sklearn.__version__.split(".")[:2])) >= (1, 6)


def dataset_w_errors():
    num_classes = 2
    num_features = 3
    n = 50
    margin = 5
    X = np.vstack(
        [np.random.random((n, num_features)), np.random.random((n, num_features)) + margin]
    )
    X = (X - X.mean(axis=0)) / X.std(axis=0)  # normalize columns
    y = np.array([0] * n + [1] * n)
    y_og = np.array(y)
    # Introduce label errors
    error_indices = [n - 3, n - 2, n - 1, n, n + 1, n + 2]
    for idx in error_indices:
        y[idx] = 1 - y[idx]  # Flip label

    if os.name == "nt":  # running on Windows
        # numpy converts to int32 instead of int64 on Windows, incompatible with neural nets
        # https://github.com/numpy/numpy/issues/17640
        X = np.float64(X)
        y = np.int64(y)
        y_og = np.int64(y_og)

    return {
        "X": X,
        "y": y,
        "y_og": y_og,
        "error_indices": error_indices,
        "num_classes": num_classes,
        "num_features": num_features,
    }


def make_rare_label(data):
    """Makes one label really rare in the dataset."""
    data = deepcopy(data)
    y = data["y"]
    class0_inds = np.where(y == 0)[0]
    if len(class0_inds) < 1:
        raise ValueError("Class 0 too rare already")
    class0_inds_remove = class0_inds[1:]
    if len(class0_inds_remove) > 0:
        y[class0_inds_remove] = 1
    data["y"] = y
    return data


SEED = 1
np.random.seed(SEED)
random.seed(SEED)
torch.manual_seed(SEED)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
torch.cuda.manual_seed_all(SEED)

DATA = dataset_w_errors()
DATA_RARE_LABEL = make_rare_label(DATA)


def test_torch(data=DATA, hidden_units=128):
    dataset = torch.utils.data.TensorDataset(
        torch.from_numpy(data["X"]).float(), torch.from_numpy(data["y"])
    )

    class TorchNetwork(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.layers = torch.nn.Sequential(
                torch.nn.Linear(data["num_features"], hidden_units),
                torch.nn.ReLU(),
                torch.nn.Linear(hidden_units, data["num_classes"]),
            )

        def forward(self, X):
            return self.layers(X)

    # Test base model works:
    skorch_config = {"criterion": torch.nn.CrossEntropyLoss, "optimizer": torch.optim.Adam}
    net = skorch.NeuralNet(TorchNetwork, **skorch_config)
    net.fit(dataset, data["y"], epochs=2)
    preds_base = net.predict(dataset)

    # Test Cleanlearning performs well:
    net = skorch.NeuralNet(TorchNetwork, **skorch_config)
    cl = CleanLearning(net)
    cl.fit(dataset, data["y"], clf_kwargs={"epochs": 30}, clf_final_kwargs={"epochs": 60})

    preds = cl.predict(dataset).argmax(axis=1)
    err = np.sum(preds != data["y_og"]) / len(data["y_og"])
    issue_indices = list(cl.label_issues_df[cl.label_issues_df["is_label_issue"]].index.values)
    assert issue_indices == data["error_indices"]
    assert err < 1e-3


@pytest.mark.filterwarnings("ignore")
def test_torch_rarelabel(data=DATA_RARE_LABEL, hidden_units=8):
    dataset = torch.utils.data.TensorDataset(
        torch.from_numpy(data["X"]).float(), torch.from_numpy(data["y"])
    )

    class TorchNetwork(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.layers = torch.nn.Sequential(
                torch.nn.Linear(data["num_features"], hidden_units),
                torch.nn.ReLU(),
                torch.nn.Linear(hidden_units, data["num_classes"]),
            )

        def forward(self, X):
            return self.layers(X)

    # Test Cleanlearning works:
    net = skorch.NeuralNet(TorchNetwork, criterion=torch.nn.CrossEntropyLoss)
    cl = CleanLearning(net)
    cl.fit(dataset, data["y"], clf_kwargs={"epochs": 2})
    pred_probs = cl.predict(dataset)