import numpy as np
import pytest
from functools import lru_cache
from importlib.util import module_from_spec, spec_from_file_location
from pathlib import Path


@lru_cache(maxsize=1)
def _wdm_instance():
    from wdm_wavelet.wdm import WDM
    return WDM(32, 64, 6, 10)


def _timeseries_values(ts):
    if hasattr(ts, "value"):
        return np.asarray(ts.value)
    return np.asarray(ts)


def _timeseries_scalar_attr(ts, name):
    value = getattr(ts, name)
    return value.value if hasattr(value, "value") else value


def _available_backends():
    from wdm_wavelet.core.t2w import HAS_JAX as HAS_T2W_JAX, HAS_NUMBA as HAS_T2W_NUMBA
    from wdm_wavelet.core.w2t import HAS_JAX as HAS_W2T_JAX, HAS_NUMBA as HAS_W2T_NUMBA

    backends = []
    if HAS_T2W_JAX and HAS_W2T_JAX:
        backends.append("jax")
    if HAS_T2W_NUMBA and HAS_W2T_NUMBA:
        backends.append("numba")
    return backends


@lru_cache(maxsize=1)
def _benchmark_module():
    script_path = Path(__file__).resolve().parents[1] / "scripts" / "benchmark_t2w_w2t.py"
    spec = spec_from_file_location("benchmark_t2w_w2t", script_path)
    module = module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def test_wdm():
    rng = np.random.default_rng(12345)
    fake_strain = rng.standard_normal(4096)

    wdm = _wdm_instance()
    tf_map = wdm.t2w(fake_strain, sample_rate=4096.)

    assert len(tf_map.data) == tf_map.n_freq
    assert tf_map.df == 4096./32/2.


def test_wdm_end_to_end_reconstruction_regression():
    """
    End-to-end regression:
    x(t) -> WDM map -> reconstructed x_hat(t)

    The transform is approximately orthonormal due to finite filter truncation,
    so this test guards a practical relative-error envelope instead of exact equality.
    """
    rng = np.random.default_rng(20260222)
    signal = rng.standard_normal(4096)
    sample_rate = 4096.0
    t0 = 100.25

    wdm = _wdm_instance()
    tf_map = wdm.t2w(signal, sample_rate=sample_rate, t0=t0)
    reconstructed_ts = wdm.w2t(tf_map)

    reconstructed = _timeseries_values(reconstructed_ts)
    relative_error = np.linalg.norm(reconstructed - signal) / np.linalg.norm(signal)

    assert reconstructed.shape == signal.shape
    assert relative_error < 0.25
    assert np.isclose(_timeseries_scalar_attr(reconstructed_ts, "sample_rate"), sample_rate)
    assert np.isclose(_timeseries_scalar_attr(reconstructed_ts, "t0"), t0)


def test_w2t_returns_original_signal_length_when_input_is_unaligned():
    """
    C++ `WDM::w2t` copies back only `nSTS` samples (original input length).
    Guard that Python `WDM.w2t` follows the same behavior when t2w used
    internal alignment/padding.
    """
    rng = np.random.default_rng(20260223)
    signal = rng.standard_normal(4099)
    sample_rate = 4096.0

    wdm = _wdm_instance()
    tf_map = wdm.t2w(signal, sample_rate=sample_rate)
    reconstructed_ts = wdm.w2t(tf_map)
    reconstructed = _timeseries_values(reconstructed_ts)

    assert reconstructed.shape == signal.shape


def test_backend_switches_match_numpy_reference():
    """Explicit backend='jax' instance produces identical output to default instance."""
    from wdm_wavelet.wdm import WDM

    rng = np.random.default_rng(20260224)
    signal = rng.standard_normal(4096)
    sample_rate = 4096.0

    wdm_default = _wdm_instance()
    tf_ref = wdm_default.t2w(signal, sample_rate=sample_rate)
    rec_ref = _timeseries_values(wdm_default.w2t(tf_ref))

    for backend in _available_backends():
        wdm_explicit = WDM(wdm_default.M, wdm_default.K, wdm_default.beta_order, wdm_default.precision, backend=backend)
        tf = wdm_explicit.t2w(signal, sample_rate=sample_rate)
        np.testing.assert_allclose(tf.data, tf_ref.data, rtol=1e-12, atol=1e-12)

        rec = _timeseries_values(wdm_explicit.w2t(tf))
        np.testing.assert_allclose(rec, rec_ref, rtol=1e-12, atol=1e-12)


def test_w2tq_backend_switches_match_numpy_reference():
    """Explicit backend='jax' instance produces identical w2tQ output to default instance."""
    from wdm_wavelet.wdm import WDM

    rng = np.random.default_rng(20260226)
    signal = rng.standard_normal(4096)
    sample_rate = 4096.0

    wdm_default = _wdm_instance()
    tf_ref = wdm_default.t2w(signal, sample_rate=sample_rate)
    recq_ref = _timeseries_values(wdm_default.w2tQ(tf_ref))

    for backend in _available_backends():
        wdm_explicit = WDM(wdm_default.M, wdm_default.K, wdm_default.beta_order, wdm_default.precision, backend=backend)
        tf = wdm_explicit.t2w(signal, sample_rate=sample_rate)
        recq = _timeseries_values(wdm_explicit.w2tQ(tf))
        np.testing.assert_allclose(recq, recq_ref, rtol=1e-12, atol=1e-12)


def test_w2tq_rejects_power_only_map():
    rng = np.random.default_rng(20260227)
    signal = rng.standard_normal(1024)
    wdm = _wdm_instance()
    tf_power = wdm.t2w(signal, sample_rate=4096.0, MM=0)

    with pytest.raises(ValueError):
        _ = wdm.w2tQ(tf_power)


def test_invalid_backend_raises():
    """Creating a WDM instance with an unsupported backend raises ValueError at construction."""
    from wdm_wavelet.wdm import WDM

    with pytest.raises(ValueError, match="Unsupported backend"):
        _ = WDM(32, 64, 6, 10, backend="invalid-backend")


def test_benchmark_end_to_end_smoke():
    rng = np.random.default_rng(20260225)
    signal = rng.standard_normal(2048)
    sample_rate = 4096.0

    wdm = _wdm_instance()
    backends = tuple(_available_backends())
    benchmark = _benchmark_module()
    results = benchmark.benchmark_python_backends(
        signal=signal,
        sample_rate=sample_rate,
        M=wdm.M,
        K=wdm.K,
        beta_order=wdm.beta_order,
        precision=wdm.precision,
        MM=-1,
        backends=list(backends),
        n_runs=2,
        warmup=0,
    )

    assert set(results.keys()) == set(backends)
    for backend in backends:
        summary = results[backend]
        assert summary["t2w_backend"] == backend
        assert summary["w2t_backend"] == backend
        assert summary["n_runs"] == 2
        assert summary["t2w_mean_s"] >= 0.0
        assert summary["w2t_mean_s"] >= 0.0
        assert summary["total_mean_s"] >= 0.0
