import tempfile, os
import numpy as np
import h5py
import pytest

def test_write_auv_mcap_group():
    from extract.extract_mcap import write_auv_mcap_group
    t = np.array([1000, 2000, 3000], dtype=np.int64)
    lat = np.array([43.1, 43.2, 43.3])
    lon = np.array([5.6, 5.61, 5.62])
    depth = np.array([5.0, 5.5, 6.0])

    with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as tmp:
        path = tmp.name
    try:
        write_auv_mcap_group(path, t, lat, lon, depth)
        with h5py.File(path, "r") as f:
            assert "auv_mcap" in f
            assert np.allclose(f["auv_mcap/lat"][:], lat)
            assert np.allclose(f["auv_mcap/lon"][:], lon)
            assert np.allclose(f["auv_mcap/depth_m"][:], depth)
            assert list(f["auv_mcap/t_ns"][:]) == list(t)
    finally:
        os.unlink(path)


def test_write_auv_mcap_group_with_altitude():
    from extract.extract_mcap import write_auv_mcap_group
    t     = np.array([1000, 2000, 3000], dtype=np.int64)
    lat   = np.array([43.1, 43.2, 43.3])
    lon   = np.array([5.6, 5.61, 5.62])
    depth = np.array([5.0, 5.5, 6.0])
    alt   = np.array([1.2, 1.3, 1.1])

    with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as tmp:
        path = tmp.name
    try:
        write_auv_mcap_group(path, t, lat, lon, depth, alt)
        with h5py.File(path, "r") as f:
            assert "altitude_m" in f["auv_mcap"]
            assert "seafloor_depth_m" in f["auv_mcap"]
            assert np.allclose(f["auv_mcap/altitude_m"][:], alt)
            assert np.allclose(f["auv_mcap/seafloor_depth_m"][:], depth + alt)
    finally:
        os.unlink(path)


def test_write_auv_mcap_group_without_altitude():
    from extract.extract_mcap import write_auv_mcap_group
    t     = np.array([1000, 2000], dtype=np.int64)
    lat   = np.zeros(2)
    lon   = np.zeros(2)
    depth = np.array([3.0, 4.0])

    with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as tmp:
        path = tmp.name
    try:
        write_auv_mcap_group(path, t, lat, lon, depth, altitude_m=None)
        with h5py.File(path, "r") as f:
            assert "altitude_m" not in f["auv_mcap"]
    finally:
        os.unlink(path)