import numpy as np
import spikeinterface as si
import spikeinterface.extractors as se
from mountainsort5.core.compute_pca_features import compute_pca_features
from mountainsort5.core.compute_templates import compute_templates
from mountainsort5.core.detect_spikes import detect_spikes
from mountainsort5.core.extract_snippets import extract_snippets, extract_snippets_in_channel_neighborhood
from mountainsort5.core.get_sampled_recording_for_training import get_sampled_recording_for_training
from mountainsort5.core.get_block_recording_for_scheme3 import get_block_recording_for_scheme3
from mountainsort5.core.cluster_snippets import cluster_snippets
from mountainsort5.core.pairwise_merge_step import pairwise_merge_step
from mountainsort5.core.get_times_labels_from_sorting import get_times_labels_from_sorting


def test_compute_pca_features():
    x = np.random.normal(size=(100, 20))
    features = compute_pca_features(x, npca=10)
    assert features.shape == (100, 10)

def test_compute_templates():
    L = 1000
    T = 20
    M = 10
    snippets = np.random.normal(size=(L, T, M))
    labels = np.random.randint(1, 11, size=(L,))
    templates = compute_templates(snippets, labels)
    assert templates.shape == (10, T, M)

def test_detect_spikes():
    N = 100
    M = 3
    traces= np.zeros((N, M))
    channel_locations = np.zeros((M, 2))
    traces[8, 0] = -10 # in margin
    traces[9, 1] = -4 # in margin
    traces[30, 1] = -4
    traces[31, 1] = -10
    traces[40, 2] = -10
    traces[41, 2] = -4
    traces[98, 0] = -10 # in margin
    times, channel_indices = detect_spikes(
        traces,
        channel_locations=channel_locations,
        time_radius=10,
        channel_radius=50,
        detect_threshold=3,
        detect_sign=-1,
        margin_left=10,
        margin_right=10,
        verbose=False
    )
    assert len(times) == len(channel_indices) == 2
    assert times[0] == 31
    assert channel_indices[0] == 1
    assert times[1] == 40
    assert channel_indices[1] == 2

def test_extract_snippets():
    N = 1000
    M = 4
    L = 100
    T1 = 20
    T2 = 20
    T = T1 + T2
    traces = np.random.normal(size=(N, M))
    times = np.random.randint(T1, N - T2, size=(L,))
    snippets = extract_snippets(traces, times=times, channel_locations=None, mask_radius=None, channel_indices=None, T1=T1, T2=T2)
    assert snippets.shape == (L, T, M)

def test_extract_snippets_in_channel_neighborhood():
    N = 1000
    M = 4
    L = 100
    T1 = 20
    T2 = 20
    T = T1 + T2
    traces = np.random.normal(size=(N, M))
    times = np.random.randint(T1, N - T2, size=(L,))
    neighborhood = [0, 2]
    snippets = extract_snippets_in_channel_neighborhood(
        traces=traces,
        times=times,
        neighborhood=neighborhood,
        T1=T1,
        T2=T2
    )
    assert snippets.shape == (L, T, len(neighborhood))

def test_get_sampled_recording_for_training():
    recording, _ = se.toy_example(duration=60, num_channels=4, num_units=10, sampling_frequency=30000, seed=0, num_segments=1)
    recording: si.BaseRecording = recording

    # test case where the training duration is longer than the recording duration
    recording1 = get_sampled_recording_for_training(
        recording=recording,
        training_duration_sec=100,
        mode='initial'
    )
    assert np.array_equal(
        recording1.get_traces(),
        recording.get_traces()
    )

    # test case where the training recording has only one chunk
    recording1b = get_sampled_recording_for_training(
        recording=recording,
        training_duration_sec=7,
        mode='uniform'
    )
    assert np.array_equal(
        recording1b.get_traces(),
        recording.get_traces(start_frame=0, end_frame=7 * recording.sampling_frequency)
    )

    # Test the mode='initial' case which is easier
    recording2 = get_sampled_recording_for_training(
        recording=recording,
        training_duration_sec=25,
        mode='initial'
    )
    assert recording2.get_num_frames() == 25 * recording.sampling_frequency
    assert np.array_equal(
        recording2.get_traces(),
        recording.get_traces(start_frame=0, end_frame=25 * recording.sampling_frequency)
    )

    # Important to test mode='uniform' which can be tricky
    recording3 = get_sampled_recording_for_training(
        recording=recording,
        training_duration_sec=25,
        mode='uniform'
    )
    assert recording3.get_num_frames() == 25 * recording.sampling_frequency
    assert np.array_equal(
        recording3.get_traces(start_frame=0, end_frame=10 * recording.sampling_frequency),
        recording.get_traces(start_frame=0, end_frame=10 * recording.sampling_frequency)
    )
    # 60 - 25 = 35
    # spacing = 35 / 2 = 17.5
    assert np.array_equal(
        recording3.get_traces(start_frame=10 * recording.sampling_frequency, end_frame=20 * recording.sampling_frequency),
        recording.get_traces(start_frame=int(27.5 * recording.sampling_frequency), end_frame=int(37.5 * recording.sampling_frequency))
    )
    # 37.5 + 17.5 = 55
    assert np.array_equal(
        recording3.get_traces(start_frame=20 * recording.sampling_frequency, end_frame=25 * recording.sampling_frequency),
        recording.get_traces(start_frame=int(55 * recording.sampling_frequency), end_frame=int(60 * recording.sampling_frequency))
    )

def test_get_block_recording_for_scheme3():
    recording, _ = se.toy_example(duration=60, num_channels=4, num_units=10, sampling_frequency=30000, seed=0, num_segments=1)
    recording: si.BaseRecording = recording

    recording2 = get_block_recording_for_scheme3(
        recording=recording,
        start_frame=10 * recording.sampling_frequency,
        end_frame=20 * recording.sampling_frequency
    )
    assert np.array_equal(
        recording.get_traces(start_frame=10 * recording.sampling_frequency, end_frame=20 * recording.sampling_frequency),
        recording2.get_traces()
    )

def test_cluster_snippets():
    N = 1000
    L = 100
    M = 4
    T1 = 20
    T2 = 20
    traces = np.random.normal(size=(N, M))
    traces[500:] += 10 # offset this so we get more than one cluster (important for coverage of cluster_snippets)
    times = np.random.randint(T1, N - T2, size=(L,))
    snippets = extract_snippets(traces, times=times, channel_locations=None, mask_radius=None, channel_indices=None, T1=T1, T2=T2)
    labels = cluster_snippets(
        snippets,
        npca_per_branch=10
    )
    assert np.min(labels) == 1
    assert len(labels) == L

def test_pairwise_merge_step():
    N = 1000
    L = 100
    M = 4
    T1 = 20
    T2 = 20
    traces = np.random.normal(size=(N, M))
    times = np.random.randint(T1, N - T2, size=(L,))
    snippets = extract_snippets(traces, times=times, channel_locations=None, mask_radius=None, channel_indices=None, T1=T1, T2=T2)
    labels = cluster_snippets(
        snippets,
        npca_per_branch=10
    )
    templates = compute_templates(snippets, labels)
    new_times, new_labels = pairwise_merge_step(
        snippets=snippets,
        labels=labels,
        templates=templates,
        times=times,
        detect_sign=-1,
        unit_ids=np.unique(labels),
        detect_time_radius=10
    )
    assert len(new_times) == len(new_labels)
    assert len(new_times) <= len(times)

def test_get_times_labels_from_sorting():
    recording, sorting = se.toy_example(duration=6, num_channels=4, num_units=10, sampling_frequency=30000, seed=0, num_segments=1)
    recording: si.BaseRecording = recording
    sorting: si.BaseSorting = sorting

    times, labels = get_times_labels_from_sorting(sorting)
    assert len(times) == len(labels)
    assert len(np.unique(labels)) == sorting.get_num_units()

    # test the case of an empty sorting
    sorting2 = se.NumpySorting(sampling_frequency=30000)
    times2, labels2 = get_times_labels_from_sorting(sorting2)
    assert len(times2) == 0
    assert len(labels2) == 0