from __future__ import annotations

import h5py
import numpy
import pytest
from silx.io.url import DataUrl

from tomoscan.framereducer.reducedframesinfos import ReducedFramesInfos
from tomoscan.framereducer.ReduceFrameSaver import ReduceFrameSaver
from tomoscan.framereducer.utils import (
    read_reduced_frames,
    read_reduced_frames_metadata,
)
from tomoscan.framereducer.version import version as frame_reducer_current_version


@pytest.fixture(scope="function")
def reduced_flats_urls(tmp_path) -> tuple[DataUrl, DataUrl]:
    """Two urls. One pointing to a HDF5 dataset with one reduced flats, second on associated metadata"""
    output_file = tmp_path / "my_reduced_flats.h5"
    frames_metadata = ReducedFramesInfos()
    frames_metadata.machine_current = (12.3, 12.4)
    frames_metadata.count_time = (1.0, 1.1)
    url_frames = DataUrl(
        file_path=output_file,
        data_path="entry0000/flats/{index}",
        scheme="silx",
    )
    url_frames_metadata = DataUrl(
        file_path=output_file,
        data_path="entry0000/flats/",
        scheme="silx",
    )
    saver = ReduceFrameSaver(
        frames={
            0: numpy.ones((10, 10)),
            10: numpy.ones((10, 10)) * 2,
        },
        output_urls=(url_frames,),
        metadata_output_urls=(url_frames_metadata,),
        frames_metadata=frames_metadata,
    )
    saver.save()
    return url_frames, url_frames_metadata


@pytest.fixture(scope="function")
def reduced_darks_urls(tmp_path) -> tuple[DataUrl, DataUrl]:
    """Two urls. One pointing to a HDF5 dataset with one reduced dark, second on associated metadata"""
    output_file = tmp_path / "my_reduced_darks.h5"

    frames_metadata = ReducedFramesInfos()
    frames_metadata.machine_current = (12.2,)
    frames_metadata.count_time = (2.0,)
    url_frames = DataUrl(
        file_path=output_file,
        data_path="entry0000/darks/{index}",
        scheme="silx",
    )
    url_frames_metadata = DataUrl(
        file_path=output_file,
        data_path="entry0000/darks/",
        scheme="silx",
    )
    saver = ReduceFrameSaver(
        frames={
            1: numpy.zeros((10, 10)),
        },
        output_urls=(url_frames,),
        metadata_output_urls=(url_frames_metadata,),
        frames_metadata=frames_metadata,
    )
    saver.save()
    return url_frames, url_frames_metadata


def test_read_reduced(reduced_flats_urls, reduced_darks_urls):
    """test to make sure we can read reduced frames and metadata directly from 'read_reduced_frame' and 'read_reduced_frame_metadata'"""
    reduced_flats_urls, reduced_flats_metadata_url = reduced_flats_urls
    reduced_darks_urls, reduced_darks_metadata_url = reduced_darks_urls
    reduced_darks = {}
    reduced_darks_infos = ReducedFramesInfos()
    reduced_flats = {}
    reduced_flats_infos = ReducedFramesInfos()
    # test reading data
    read_reduced_frames(
        url=reduced_darks_urls,
        res_data=reduced_darks,
        res_metadata=reduced_darks_infos,
        read_as_url=False,
    )
    assert len(reduced_darks) == 1
    numpy.testing.assert_almost_equal(reduced_darks[1], numpy.zeros((10, 10)))
    read_reduced_frames(
        url=reduced_flats_urls,
        res_data=reduced_flats,
        res_metadata=reduced_flats_infos,
        read_as_url=False,
    )
    assert len(reduced_flats) == 2
    numpy.testing.assert_almost_equal(reduced_flats[0], numpy.ones((10, 10)))
    numpy.testing.assert_almost_equal(reduced_flats[10], numpy.ones((10, 10)) * 2)

    # test reading metadata
    read_reduced_frames_metadata(
        url=reduced_darks_metadata_url,
        res_metadata=reduced_darks_infos,
    )
    numpy.testing.assert_almost_equal(reduced_darks_infos.machine_current, (12.2,))
    read_reduced_frames_metadata(
        url=reduced_flats_metadata_url,
        res_metadata=reduced_flats_infos,
    )
    numpy.testing.assert_almost_equal(reduced_flats_infos.machine_current, (12.3, 12.4))
    with h5py.File(reduced_flats_urls.file_path(), "r") as h5s:
        assert h5s["entry0000/flats"].attrs["version"] == frame_reducer_current_version

    with h5py.File(reduced_darks_urls.file_path(), "r") as h5s:
        assert h5s["entry0000/darks"].attrs["version"] == frame_reducer_current_version
