"""Tests for spectrogram visualization across RadioDataset sample layouts.

Regression: split-I/Q samples with shape ``(2, T)`` previously reported a length
of 2 (the I/Q axis) instead of ``T``, so every such dataset was rejected with
"doesn't have sufficient signal data for spectrogram visualization".
"""

import numpy as np
import pandas as pd
import pytest

from ria_toolkit_oss.viz.radio_dataset import _to_complex_1d, sample_spectrogram_plot


class _FakeDataset:
    """Minimal RadioDataset stand-in: a one-row metadata frame + a fixed sample."""

    def __init__(self, sample):
        self._sample = np.asarray(sample)
        self.metadata = pd.DataFrame({"modulation": ["bpsk"]})

    def __getitem__(self, idx):
        return self._sample


def _has_spectrogram(fig):
    """True when fig is a real spectrogram (a Heatmap trace), not an error card."""
    return any(getattr(tr, "type", None) == "heatmap" for tr in fig.data)


# --- _to_complex_1d ---------------------------------------------------------


@pytest.mark.parametrize(
    "sample, expected_len",
    [
        (np.exp(1j * np.linspace(0, 1, 1024)), 1024),  # complex 1-D
        (np.ones((2, 1024)), 1024),  # split I/Q rows (2, T)
        (np.ones((1024, 2)), 1024),  # split I/Q cols (T, 2)
        (np.ones(1024), 1024),  # real 1-D
        (np.ones((2, 4, 256)), 1024),  # multi-channel I/Q rows
    ],
)
def test_to_complex_1d_normalizes(sample, expected_len):
    sig = _to_complex_1d(sample)
    assert sig is not None
    assert sig.ndim == 1
    assert sig.size == expected_len
    assert np.iscomplexobj(sig)


def test_to_complex_1d_combines_iq_rows():
    arr = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])  # I row, Q row
    assert np.allclose(_to_complex_1d(arr), np.array([1 + 4j, 2 + 5j, 3 + 6j]))


def test_to_complex_1d_combines_iq_cols():
    arr = np.array([[1.0, 4.0], [2.0, 5.0], [3.0, 6.0]])  # (T, 2)
    assert np.allclose(_to_complex_1d(arr), np.array([1 + 4j, 2 + 5j, 3 + 6j]))


@pytest.mark.parametrize("sample", [None, np.array([])])
def test_to_complex_1d_returns_none_for_empty(sample):
    assert _to_complex_1d(sample) is None


# --- sample_spectrogram_plot ------------------------------------------------


@pytest.mark.parametrize(
    "sample",
    [
        np.exp(1j * np.linspace(0, 10, 1024)),  # complex 1-D
        np.random.randn(2, 1024),  # split I/Q rows  <-- the reported bug
        np.random.randn(1024, 2),  # split I/Q cols
        np.random.randn(1024),  # real 1-D
        np.random.randn(2, 4, 256),  # multi-channel I/Q
    ],
)
def test_sample_spectrogram_renders(sample):
    fig = sample_spectrogram_plot(_FakeDataset(sample), sample_idx=0)
    assert _has_spectrogram(fig), "expected a real spectrogram, got an error/unavailable figure"


@pytest.mark.parametrize("sample", [np.random.randn(2), np.array([])])
def test_sample_spectrogram_too_short_returns_error(sample):
    fig = sample_spectrogram_plot(_FakeDataset(sample), sample_idx=0)
    assert not _has_spectrogram(fig), "expected the 'Not Available' figure for too-short/empty samples"