import numpy as np
import h5py
from sfm.shape_metrics import offset_from_samples, offset_from_h5


def test_flags_constant_dc_floor():
    # A geophone channel sitting at a constant +0.05 in/s across the whole record
    # is a DC offset: baseline off zero AND flat across pre/mid/end thirds.
    n = 300
    chans = {"Tran": np.full(n, 0.05), "Vert": np.zeros(n), "Long": np.zeros(n)}
    r = offset_from_samples(chans, pretrig_n=50)
    assert r["offset"] is True
    assert r["axis"] == "Tran"
    assert abs(r["pre"] - 0.05) < 1e-6
    assert r["spread"] < 0.02


def test_transient_rejected_by_spread():
    # Off-zero pre-trigger but the baseline SETTLES back over the record — a
    # transient, not a constant offset. The spread test must reject it.
    x = np.concatenate([np.full(100, 0.05), np.full(100, 0.025), np.zeros(100)])
    chans = {"Tran": x, "Vert": np.zeros(300), "Long": np.zeros(300)}
    r = offset_from_samples(chans, pretrig_n=100)
    assert r["offset"] is False


def test_clean_oscillation_not_offset():
    t = np.arange(300)
    x = 0.4 * np.sin(2 * np.pi * t / 20)   # oscillates around zero — baseline IS zero
    chans = {"Tran": x, "Vert": np.zeros(300), "Long": np.zeros(300)}
    r = offset_from_samples(chans, pretrig_n=50)
    assert r["offset"] is False


def test_below_floor_not_offset_but_reports_pre():
    # A flat baseline below the floor is not an offset; still report the axis/pre
    # for tuning transparency.
    n = 300
    chans = {"Tran": np.full(n, 0.01), "Vert": np.zeros(n), "Long": np.zeros(n)}
    r = offset_from_samples(chans, pretrig_n=50)
    assert r["offset"] is False
    assert r["axis"] == "Tran"
    assert abs(r["pre"] - 0.01) < 1e-6


def test_none_when_no_geo_channels():
    assert offset_from_samples({"MicL": np.full(300, 0.05)}, pretrig_n=50) is None


def test_pretrig_fallback_when_invalid():
    # pretrig_n of 0 (missing/unusable) falls back to the first third.
    n = 300
    chans = {"Tran": np.full(n, 0.05), "Vert": np.zeros(n), "Long": np.zeros(n)}
    r = offset_from_samples(chans, pretrig_n=0)
    assert r["offset"] is True


def test_flags_offset_on_any_axis():
    # Offset on Vert alone still flags the event, and Vert is reported.
    n = 300
    chans = {"Tran": np.zeros(n), "Vert": np.full(n, -0.06), "Long": np.zeros(n)}
    r = offset_from_samples(chans, pretrig_n=50)
    assert r["offset"] is True
    assert r["axis"] == "Vert"


def _write_h5(path, chans, pretrig_n):
    with h5py.File(path, "w") as f:
        g = f.create_group("samples")
        for k, v in chans.items():
            g.create_dataset(k, data=np.asarray(v, dtype="float32"))
        if pretrig_n is not None:
            f.attrs["pretrig_samples"] = pretrig_n


def test_offset_from_h5_reads_pretrig_attr(tmp_path):
    p = tmp_path / "ev.h5"
    n = 300
    _write_h5(p, {"Tran": np.full(n, 0.05), "Vert": np.zeros(n), "Long": np.zeros(n)},
              pretrig_n=50)
    r = offset_from_h5(str(p))
    assert r["offset"] is True and r["axis"] == "Tran"


def test_offset_from_h5_missing_pretrig_attr_falls_back(tmp_path):
    p = tmp_path / "noattr.h5"
    n = 300
    _write_h5(p, {"Tran": np.full(n, 0.05), "Vert": np.zeros(n), "Long": np.zeros(n)},
              pretrig_n=None)
    assert offset_from_h5(str(p))["offset"] is True   # falls back to first-third


def test_offset_from_h5_missing_file_is_none(tmp_path):
    assert offset_from_h5(str(tmp_path / "nope.h5")) is None