from typing import List

import h5py
import numpy

from ..io.blissconcat import concatenate_bliss_scans


def test_concatenate_basic(tmp_path):
    infile = tmp_path / "dataset.h5"
    outfile = tmp_path / "concat.h5"

    values1 = numpy.arange(5)
    values2 = numpy.arange(5, 12)
    _create_test_scan(infile, "1.1", values1, scalars=["a", "b", "c"])
    _create_test_scan(infile, "2.1", values2, scalars=["a", "b"])

    result_uri = concatenate_bliss_scans(
        [f"{infile}::/1.1", f"{infile}::/2.1"],
        f"{outfile}::/1.1",
        virtual_axes={"sy": "<samy> + <sampy>"},
        axes_units={"sampy": "um", "sy": "cm"},
    )
    assert result_uri == f"{outfile}::/1.1"

    with h5py.File(outfile, "r") as f:
        samy = f["/1.1/measurement/samy"][()]
        sampy = f["/1.1/measurement/sampy"][()]
        sy = f["/1.1/measurement/sy"][()]
        positioners = set(f["/1.1/instrument/positioners"])

    expected = numpy.arange(12)  # mm
    assert numpy.allclose(samy, expected)

    expected = numpy.arange(12)  # um
    assert numpy.allclose(sampy, expected)

    expected = numpy.arange(12) / 10 + numpy.arange(12) / 10000  # cm
    assert numpy.allclose(sy, expected)

    assert positioners == {"samy", "sampy", "sy", "a", "b"}


def _create_test_scan(
    infile: str, scan: str, values: numpy.ndarray, scalars: List[str]
) -> None:
    with h5py.File(infile, "a") as h5file:
        grp = h5file.create_group(scan)
        grp.attrs["NX_class"] = "NXentry"

        instrument = grp.create_group("instrument")
        instrument.attrs["NX_class"] = "NXinstrument"

        pos = instrument.create_group("positioners")
        pos.attrs["NX_class"] = "NXcollection"

        meas = grp.create_group("measurement")
        meas.attrs["NX_class"] = "NXcollection"

        dset = pos.create_dataset("samy", data=values)
        dset.attrs["units"] = "mm"
        meas["samy"] = h5py.SoftLink(dset.name)

        dset = pos.create_dataset("sampy", data=values)
        # Missing unit
        # dset.attrs["units"] = "um"
        meas["sampy"] = h5py.SoftLink(dset.name)

        for scalar in scalars:
            pos[scalar] = 0
