# Authors: Hubert Banville <hubert.jbanville@gmail.com>
#          Bruno Aristimunha <b.aristimunha@gmail.com>
# License: BSD-3

import os
from unittest import mock

import mne
import numpy as np
import h5py
import pytest
import torch
import tempfile

from sklearn.utils import check_random_state
from numpy.testing import assert_array_equal, assert_allclose


from braindecode.util import _cov_and_var_to_corr, _cov_to_corr, \
    corr, create_mne_dummy_raw, \
    read_all_file_names, set_random_seeds, th_to_np, cov, np_to_th, \
    get_balanced_batches


def test_create_mne_dummy_raw(tmp_path):
    n_channels, n_times, sfreq = 2, 10000, 100
    raw, fnames = create_mne_dummy_raw(
        n_channels, n_times, sfreq, savedir=tmp_path, save_format=["fif", "hdf5"]
    )

    assert isinstance(raw, mne.io.RawArray)
    assert len(raw.ch_names) == n_channels
    assert raw.n_times == n_times
    assert raw.info["sfreq"] == sfreq
    assert isinstance(fnames, dict)
    assert os.path.isfile(fnames["fif"])
    assert os.path.isfile(fnames["hdf5"])

    _ = mne.io.read_raw_fif(fnames["fif"], preload=False, verbose=None)
    with h5py.File(fnames["hdf5"], "r") as hf:
        _ = np.array(hf["fake_raw"])


def test_set_random_seeds_raise_value_error():
    with pytest.raises(
        ValueError, match="cudnn_benchmark expected to be bool or None, got 'abc'"
    ):
        set_random_seeds(100, True, "abc")


def test_set_random_seeds_warning():
    torch.backends.cudnn.benchmark = True
    with pytest.warns(
        UserWarning,
        match="torch.backends.cudnn.benchmark was set to True which may results in "
        "lack of reproducibility. In some cases to ensure reproducibility you "
        "may need to set torch.backends.cudnn.benchmark to False.",
    ):
        set_random_seeds(100, True)


def test_set_random_seeds_with_valid_cudnn_benchmark():
    with mock.patch("torch.backends.cudnn") as mock_cudnn:
        # Test with cudnn_benchmark = True
        set_random_seeds(42, cuda=True, cudnn_benchmark=True)
        assert mock_cudnn.benchmark is True

        # Test with cudnn_benchmark = False
        set_random_seeds(42, cuda=True,
                         cudnn_benchmark=False)
        assert mock_cudnn.benchmark is False


def test_set_random_seeds_with_invalid_cudnn_benchmark():
    with pytest.raises(ValueError):
        set_random_seeds(42, cuda=True,
                         cudnn_benchmark='invalid_type')


def test_th_to_np_data_preservation():
    # Test with different data types
    for dtype in [torch.float32, torch.int32]:
        tensor = torch.tensor([1, 2, 3], dtype=dtype)
        np_array = th_to_np(tensor)
        assert np_array.dtype == tensor.numpy().dtype
        # Corrected attribute access
        assert_array_equal(np_array, tensor.numpy())


def test_th_to_np_on_cpu():
    # Create a tensor on CPU
    cpu_tensor = torch.tensor([1, 2, 3], dtype=torch.float32)
    np_array = th_to_np(cpu_tensor)
    # Check the type and data of the numpy array
    assert isinstance(np_array, np.ndarray)
    assert np_array.dtype == cpu_tensor.numpy().dtype
    # Correct way to check dtype
    assert_array_equal(np_array, cpu_tensor.numpy())


def test_cov_basic():
    # Create two simple identical arrays
    a = np.array([[1, 2, 3], [4, 5, 6]])
    b = np.array([[1, 2, 3], [4, 5, 6]])
    expected_cov = np.array([[1, 1], [1, 1]])  # Calculated expected covariance
    computed_cov = cov(a, b)
    assert_allclose(computed_cov, expected_cov, rtol=1e-5)


def test_cov_dimension_mismatch():
    # Arrays with mismatched sample size should raise an error
    a = np.array([[1, 2], [3, 4]])
    b = np.array([[1, 2, 3]])
    with pytest.raises(ValueError):
        cov(a, b)


def test_np_to_th_basic_conversion():
    # Convert a simple list to tensor
    data = [1, 2, 3]
    tensor = np_to_th(data)
    assert torch.is_tensor(tensor)
    assert_array_equal(tensor.numpy(), np.array(data))


def test_np_to_th_dtype_conversion():
    # Convert and specify dtype
    data = [1.0, 2.0, 3.0]
    tensor = np_to_th(data, dtype=np.float32)
    assert tensor.dtype == torch.float32


def test_np_to_th_requires_grad():
    # Check requires_grad attribute
    data = np.array([1.0, 2.0, 3.0])
    tensor = np_to_th(data, requires_grad=True)
    assert tensor.requires_grad is True


@pytest.mark.skipif(not torch.cuda.is_available(),
                    reason="Requires CUDA support")
def test_np_to_th_pin_memory():
    # Create a numpy array
    data = np.array([1, 2, 3])

    # Convert the numpy array to a tensor with pin_memory=True
    tensor = np_to_th(data, pin_memory=True)
    # Check if the tensor is pinned in memory
    assert tensor.is_pinned() is True
    # Convert the numpy array to a tensor with pin_memory=False
    tensor = np_to_th(data, pin_memory=False)
    # Check if the tensor is not pinned in memory
    assert tensor.is_pinned() is False

def test_np_to_th_requires_grad_unsupported_dtype():
    # Attempt to set requires_grad on an unsupported dtype (integers)
    data = np.array([1, 2, 3])
    with pytest.raises(RuntimeError,
                       match="Only Tensors of floating point and complex dtype can require gradients"):
        np_to_th(data, requires_grad=True)


def test_np_to_th_tensor_options():
    # Additional tensor options like device
    data = [1, 2, 3]
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    tensor = np_to_th(data, device=device)
    assert tensor.device.type == device


def test_np_to_th_single_number_conversion():
    # Single number conversion to tensor
    data = 42
    tensor = np_to_th(data)
    assert torch.is_tensor(tensor)
    assert tensor.item() == 42  # Check if the value is correct


def test_cov_and_var_to_corr_zero_variance():
    # Scenario with zero variance, expecting result to handle divide by zero
    this_cov = np.array([[1, 0],
                         [0, 1]])
    var_a = np.array([0, 1])  # Zero variance in the first variable
    var_b = np.array([1, 0])  # Zero variance in the second variable
    calculated_corr = _cov_and_var_to_corr(this_cov, var_a, var_b)

    # Expected correlation matrix
    expected_corr = np.array([[np.inf, np.nan],
                               [0, np.inf]])
    assert_array_equal(calculated_corr, expected_corr)


def test_cov_and_var_to_corr_single_element():
    # Testing with single-element arrays
    this_cov = np.array([[1]])
    var_a = np.array([1])
    var_b = np.array([1])
    expected_corr = np.array([[1]])
    calculated_corr = _cov_and_var_to_corr(this_cov, var_a, var_b)
    assert_array_equal(calculated_corr, expected_corr)



def test_cov_to_corr_unbiased():
    # Create datasets a and b with known covariance and variance characteristics
    a = np.array([[1, 2, 3, 4],
                  [2, 3, 4, 5]])
    b = np.array([[1, 3, 5, 7],
                  [5, 6, 7, 8]])
    # Covariance between the features of a and b
    # Calculating covariance manually for known values
    demeaned_a = a - np.mean(a, axis=1, keepdims=True)
    demeaned_b = b - np.mean(b, axis=1, keepdims=True)
    this_cov = np.dot(demeaned_a, demeaned_b.T) / (b.shape[1] - 1)

    # Compute expected correlation using standard formulas for correlation
    var_a = np.var(a, axis=1, ddof=1)
    var_b = np.var(b, axis=1, ddof=1)
    expected_divisor = np.outer(np.sqrt(var_a), np.sqrt(var_b))
    expected_corr = this_cov / expected_divisor

    # Compute correlation using the function
    calculated_corr = _cov_to_corr(this_cov, a, b)

    # Assert that the calculated correlation matches the expected correlation
    assert_allclose(calculated_corr, expected_corr, rtol=1e-5)


def test_balanced_batches_basic():
    n_trials = 100
    seed = 42
    rng = check_random_state(seed)
    n_batches = 10
    batches = get_balanced_batches(n_trials, rng, shuffle=False,
                                   n_batches=n_batches)

    # Check correct number of batches
    assert len(batches) == n_batches

    # Check balanced batch sizes
    all_batch_sizes = [len(batch) for batch in batches]
    max_size = max(all_batch_sizes)
    min_size = min(all_batch_sizes)
    assert max_size - min_size <= 1

    # Check if all indices are unique and accounted for
    all_indices = np.concatenate(batches)
    assert np.array_equal(np.sort(all_indices), np.arange(n_trials))


def test_balanced_batches_with_batch_size():
    n_trials = 105
    seed = 42
    rng = check_random_state(seed)
    batch_size = 20
    batches = get_balanced_batches(n_trials, rng, shuffle=False,
                                   batch_size=batch_size)

    # Check the modified batch size condition
    expected_n_batches = int(np.round(n_trials / float(batch_size)))
    assert len(batches) == expected_n_batches

    # Checking the total number of indices
    all_indices = np.concatenate(batches)
    assert len(all_indices) == n_trials


def test_balanced_batches_shuffle():
    n_trials = 50
    seed = 42
    rng = check_random_state(seed)
    batches_no_shuffle = get_balanced_batches(n_trials, rng, shuffle=False,
                                              batch_size=10)
    rng = check_random_state(seed)
    batches_with_shuffle = get_balanced_batches(n_trials, rng, shuffle=True,
                                                batch_size=10)

    # Check that shuffling changes the order of indices
    assert not np.array_equal(np.concatenate(batches_no_shuffle),
                              np.concatenate(batches_with_shuffle))


def test_corr_correlation_computation():
    # Create two 2D arrays with known correlation
    a = np.array([[1, 2, 3], [4, 5, 6]])
    b = np.array([[1, 2, 3], [4, 5, 6]])

    # Call the corr function
    corr_result = corr(a, b)

    # Compute the known correlation
    known_corr = np.array([np.corrcoef(a[i], b[i]) for i in range(a.shape[0])])

    # Extract the correlation computation from the corr function
    computed_corr = _cov_to_corr(cov(a, b), a, b)

    # Assert that the computed correlation matches the known correlation
    assert np.allclose(computed_corr, known_corr)
    assert np.allclose(corr_result, computed_corr)


def test_get_balanced_batches_zero_batches():
    # Create a scenario where n_batches is 0
    n_trials = 10
    rng = check_random_state(0)
    shuffle = False
    n_batches = 0
    batch_size = None

    # Call the get_balanced_batches function
    batches = get_balanced_batches(n_trials, rng, shuffle, n_batches, batch_size)

    # Check if the function returns a single batch with all trials
    assert len(batches) == 1
    assert len(batches[0]) == n_trials


def test_get_balanced_batches_i_batch_less_than_n_batches_with_extra_trial():
    # Create a scenario where i_batch < n_batches_with_extra_trial
    n_trials = 10
    rng = check_random_state(0)
    shuffle = False
    n_batches = 6
    batch_size = None

    # Call the get_balanced_batches function
    batches = get_balanced_batches(n_trials, rng, shuffle, n_batches, batch_size)

    # Check if the first batch has one more trial than the last batch
    assert len(batches[0]) > len(batches[-1])


def test_read_all_file_names():
    # Create a temporary directory
    with tempfile.TemporaryDirectory() as tmpdir:
        temp_files = []
        try:
            # Create some temporary files with .txt extension
            for i in range(5):
                temp_file = os.path.join(tmpdir, f'temp{i}.txt')
                with open(temp_file, 'w') as f:
                    f.write('This is a temporary file.')
                temp_files.append(temp_file)

            # Call the read_all_file_names function
            file_paths = read_all_file_names(tmpdir, '.txt')

            # Check if the function found all the temporary files
            assert len(file_paths) == 5

            # Check if the paths returned by the function are correct
            for i in range(5):
                assert os.path.join(tmpdir, f'temp{i}.txt') in file_paths
        finally:
            # Delete the temporary files
            for temp_file in temp_files:
                os.remove(temp_file)


def test_read_all_file_names_error():
    with pytest.raises(AssertionError):
        # Call the read_all_file_names function with a non-existent directory
        read_all_file_names('non_existent_dir', '.txt')


def test_read_all_files_not_extension():
    with pytest.raises(AssertionError):
        # Call the read_all_file_names function with a non-existent directory
        read_all_file_names('non_existent_dir', 'txt')