from pathlib import Path
import numpy as np

from breesy import Event, recording
from breesy.load import load_recording
from breesy.processing import get_frequency_spectrum_channel, get_frequency_spectrum, downsample

REPO_ROOT = Path(__file__).parent.parent


def test_get_frequency_spectrum_channel():
    # Arrange
    datapath = REPO_ROOT / "tests" / "data" / "alpha" / "subject_01.mat"
    rec = load_recording(str(datapath))
    first_channel_name = rec.channel_names[0]

    # Act
    f, spectrum = get_frequency_spectrum_channel(rec, first_channel_name)

    # Assert
    assert f.ndim == 1
    assert spectrum.ndim == 1
    assert len(f) == len(spectrum)
    assert np.all(f >= 0)
    assert np.all(spectrum > 0)
    assert np.all(np.isfinite(spectrum))


def test_get_frequency_spectrum_channel_matches_full():
    # Arrange
    datapath = REPO_ROOT / "tests" / "data" / "alpha" / "subject_01.mat"
    rec = load_recording(str(datapath))
    first_channel_name = rec.channel_names[0]

    # Act
    f_full, spectrum_full = get_frequency_spectrum(rec)
    f_single, spectrum_single = get_frequency_spectrum_channel(rec, first_channel_name)

    # Assert
    assert np.array_equal(f_full, f_single)
    assert np.allclose(spectrum_full[0], spectrum_single)

def test_downsample_integer_factor():
    # Arrange
    original_rate = 512
    target_rate = 128  # factor of 4 (0.25)

    datapath = REPO_ROOT / "tests" / "data" / "alpha" / "subject_01.mat"
    rec = load_recording(str(datapath), sample_rate=original_rate)
    
    test_events = [
        Event(name="start", index=512),  # Expect at 128
        Event(name="middle", index=1024), # Expect at 256
        Event(name="end", index=2048) # Expect at 512
    ]
    rec = recording.add_events(rec, test_events)
    
    # Act
    downsampled = downsample(rec, new_sample_rate=target_rate)
    
    # Assert
    assert downsampled.sample_rate == target_rate
    assert downsampled.channel_names == rec.channel_names
    assert downsampled.number_of_channels == rec.number_of_channels
    
    assert downsampled.number_of_samples == rec.number_of_samples / 4
    
    assert not np.array_equal(downsampled.data[:, :50], rec.data[:, :50])
    
    # Events
    assert len(downsampled.events) == len(test_events)

    assert downsampled.events[0].index == 128
    assert downsampled.events[1].index == 256
    assert downsampled.events[2].index == 512


def test_downsample_float_factor():
    # Arrange
    original_rate = 512
    target_rate = 384  # factor of 1.33(3) (0.75)
    
    datapath = REPO_ROOT / "tests" / "data" / "alpha" / "subject_01.mat"
    rec = load_recording(str(datapath), sample_rate=original_rate)
    
    test_events = [
        Event(name="start", index=512),  # Expect at 384
        Event(name="middle", index=1024),  # Expect at 768
        Event(name="end", index=2048)  # Expect at 1536
    ]
    rec = recording.add_events(rec, test_events)
    
    # Act
    downsampled = downsample(rec, new_sample_rate=target_rate)
    
    # Assert
    assert downsampled.sample_rate == target_rate
    assert downsampled.channel_names == rec.channel_names
    assert downsampled.number_of_channels == rec.number_of_channels
    
    assert downsampled.number_of_samples == rec.number_of_samples * 3/4
    
    assert not np.array_equal(downsampled.data[:, :50], rec.data[:, :50])
    
    # Events
    assert len(downsampled.events) == len(test_events)
    
    assert downsampled.events[0].index == 384
    assert downsampled.events[1].index == 768
    assert downsampled.events[2].index == 1536