import h5py
import numpy
from ewoks import execute_graph
from silx.io.url import DataUrl
import logging

logger = logging.getLogger(__name__)


def check_h5groups_equivalent(
    url_reference,
    url_test,
):
    dataurl_reference = DataUrl(url_reference)
    dataurl_test = DataUrl(url_test)

    with h5py.File(dataurl_reference.file_path(), "r") as file_reference:
        with h5py.File(dataurl_test.file_path(), "r") as file_test:
            reference_group = file_reference[dataurl_reference.data_path()]
            test_group = file_test[dataurl_test.data_path()]

            def check_equivalence(name_reference_item, h5item_reference):
                if "TitleExtension" in name_reference_item:
                    # Only the ewoks files before ~October25 do not contain TitleExtension
                    # After that, the header includes TitleExtensionTemplate #TODO add the option TitleExtensionTemplate
                    return

                if "interpreted" in h5item_reference.name:
                    if name_reference_item == "epoch":
                        # ewoksid02 does not generate epoch group, it's artificial
                        return

                    # When it's an offline processing, the counter names in interpreted group cannot be rebuilt fully because this information is not in the RAW_DATA file (only in blissdata)
                    if name_reference_item not in test_group:
                        logger.warning(
                            f"{name_reference_item} not in {dataurl_test.path()}, but could have a different name"
                        )
                        return

                assert (
                    name_reference_item in test_group
                ), f"{name_reference_item} not in {dataurl_test.path()}"
                logger.debug(f"{name_reference_item} found in {test_group}")

                if isinstance(h5item_reference, h5py.Dataset):
                    check_equivalence_datasets(
                        dataset_test=test_group[name_reference_item],
                        dataset_reference=h5item_reference,
                    )

            def check_equivalence_datasets(
                dataset_test: h5py.Dataset, dataset_reference: h5py.Dataset
            ):

                data_reference = dataset_reference[()]
                data_test = dataset_test[()]

                if isinstance(data_reference, numpy.ndarray):
                    if dataurl_reference.data_slice():
                        data_reference = data_reference[
                            dataurl_reference.data_slice()[
                                0
                            ] : dataurl_reference.data_slice()[1]
                        ]
                    if dataurl_test.data_slice():
                        data_test = data_test[
                            dataurl_test.data_slice()[0] : dataurl_test.data_slice()[1]
                        ]

                    if any(_ in dataset_test.name for _ in ("HS32C", "HS32V")):
                        # Because in ewoksid02, we save also aux1, aux2 counters
                        data_test = numpy.concatenate(
                            (data_test[:, 0:10], data_test[:, 12:]), axis=1
                        )
                        data_reference = numpy.concatenate(
                            (data_reference[:, 0:10], data_reference[:, 12:]), axis=1
                        )

                    if "HS32N" in dataset_test.name:
                        # HS32N is the only array (so far) made of string elements
                        data_test = numpy.array([_.decode() for _ in data_test])
                        data_reference = numpy.array(
                            [_.decode() for _ in data_reference]
                        )
                        for i, j in zip(data_test, data_reference):
                            assert i == j
                        return

                    assert numpy.allclose(
                        data_test, data_reference, equal_nan=True
                    ), f"{dataset_reference.name} is not equivalent"
                elif isinstance(data_reference, bytes):
                    assert (
                        data_reference.decode() == data_test.decode()
                    ), f"{dataset_reference.name} is not equivalent"
                else:
                    assert (
                        data_reference == data_test
                    ), f"{dataset_reference.name} is not equivalent"
                logger.debug(f"{dataset_reference.name} are equivalent")

            if isinstance(reference_group, h5py.Dataset) and isinstance(
                test_group, h5py.Dataset
            ):
                check_equivalence_datasets(
                    dataset_test=test_group,
                    dataset_reference=reference_group,
                )
            elif isinstance(reference_group, h5py.Group):
                reference_group.visititems(func=check_equivalence)


def check_h5groups_common(
    filename_reference: str,
    filename_test: str,
    detector_name: str = "eiger2",
):

    check_h5groups_equivalent(
        url_reference=f"silx://{filename_reference}?path=/entry_0000/PyFAI/TFG/&slice=0,2",
        url_test=f"silx://{filename_test}?path=/entry_0000/PyFAI/TFG/",
    )
    check_h5groups_equivalent(
        url_reference=f"silx://{filename_reference}?path=/entry_0000/PyFAI/parameters/&slice=0,2",
        url_test=f"silx://{filename_test}?path=/entry_0000/PyFAI/parameters/",
    )
    check_h5groups_equivalent(
        url_reference=f"silx://{filename_reference}?path=/entry_0000/PyFAI/{detector_name}",
        url_test=f"silx://{filename_test}?path=/entry_0000/PyFAI/{detector_name}",
    )
    check_h5groups_equivalent(
        url_reference=f"silx://{filename_reference}?path=/entry_0000/PyFAI/MCS/ExposureTime/&slice=0,2",
        url_test=f"silx://{filename_test}?path=/entry_0000/PyFAI/MCS/ExposureTime",
    )
    check_h5groups_equivalent(
        url_reference=f"silx://{filename_reference}?path=/entry_0000/PyFAI/MCS/HS32C/&slice=0,2",
        url_test=f"silx://{filename_test}?path=/entry_0000/PyFAI/MCS/HS32C",
    )
    check_h5groups_equivalent(
        url_reference=f"silx://{filename_reference}?path=/entry_0000/PyFAI/MCS/HS32F",
        url_test=f"silx://{filename_test}?path=/entry_0000/PyFAI/MCS/HS32F",
    )
    check_h5groups_equivalent(
        url_reference=f"silx://{filename_reference}?path=/entry_0000/PyFAI/MCS/HS32N",
        url_test=f"silx://{filename_test}?path=/entry_0000/PyFAI/MCS/HS32N",
    )
    check_h5groups_equivalent(
        url_reference=f"silx://{filename_reference}?path=/entry_0000/PyFAI/MCS/HS32V/&slice=0,2",
        url_test=f"silx://{filename_test}?path=/entry_0000/PyFAI/MCS/HS32V",
    )
    check_h5groups_equivalent(
        url_reference=f"silx://{filename_reference}?path=/entry_0000/PyFAI/MCS/HS32Z",
        url_test=f"silx://{filename_test}?path=/entry_0000/PyFAI/MCS/HS32Z",
    )
    check_h5groups_equivalent(
        url_reference=f"silx://{filename_reference}?path=/entry_0000/PyFAI/MCS/Intensity0ShutCor/&slice=0,2",
        url_test=f"silx://{filename_test}?path=/entry_0000/PyFAI/MCS/Intensity0ShutCor",
    )
    check_h5groups_equivalent(
        url_reference=f"silx://{filename_reference}?path=/entry_0000/PyFAI/MCS/Intensity0UnCor/&slice=0,2",
        url_test=f"silx://{filename_test}?path=/entry_0000/PyFAI/MCS/Intensity0UnCor",
    )
    check_h5groups_equivalent(
        url_reference=f"silx://{filename_reference}?path=/entry_0000/PyFAI/MCS/Intensity1ShutCor/&slice=0,2",
        url_test=f"silx://{filename_test}?path=/entry_0000/PyFAI/MCS/Intensity1ShutCor",
    )
    check_h5groups_equivalent(
        url_reference=f"silx://{filename_reference}?path=/entry_0000/PyFAI/MCS/Intensity1UnCor/&slice=0,2",
        url_test=f"silx://{filename_test}?path=/entry_0000/PyFAI/MCS/Intensity1UnCor",
    )
    check_h5groups_equivalent(
        url_reference=f"silx://{filename_reference}?path=/entry_0000/PyFAI/MCS/raw/&slice=0,2",
        url_test=f"silx://{filename_test}?path=/entry_0000/PyFAI/MCS/raw/&slice=0,2",
    )
    check_h5groups_equivalent(
        url_reference=f"silx://{filename_reference}?path=/entry_0000/PyFAI/MCS/interpreted/&slice=0,2",
        url_test=f"silx://{filename_test}?path=/entry_0000/PyFAI/MCS/interpreted",
    )
    for grp in (
        "HSI0",
        "HSI1",
        "HSI0Factor",
        "HSI1Factor",
        "HSTime",
        "ShutterClosingTime",
        "ShutterOpeningTime",
    ):
        check_h5groups_equivalent(
            url_reference=f"silx://{filename_reference}?path=/entry_0000/PyFAI/MCS/{grp}",
            url_test=f"silx://{filename_test}?path=/entry_0000/PyFAI/MCS/{grp}",
        )


def check_allclose(
    url_reference,
    url_test,
):
    dataurl_reference = DataUrl(url_reference)
    dataurl_test = DataUrl(url_test)
    with h5py.File(dataurl_reference.file_path(), "r") as file_reference:
        with h5py.File(dataurl_test.file_path(), "r") as file_tested:
            data_reference = file_reference[dataurl_reference.data_path()][()]
            data_test = file_tested[dataurl_test.data_path()][()]
            if dataurl_reference.data_slice():
                data_reference = data_reference[
                    dataurl_reference.data_slice()[0] : dataurl_reference.data_slice()[
                        1
                    ]
                ]
            if dataurl_test.data_slice():
                data_test = data_test[
                    dataurl_test.data_slice()[0] : dataurl_test.data_slice()[1]
                ]

            assert numpy.allclose(data_reference, data_test, equal_nan=True)


# def allclose_data_sigma(
#     filename_processed_reference,
#     filename_processed_test,
#     range_index_read,
#     processing_type,
#     # atol=1e-6,
# ):
#     path_to_data_signal_processed = f"entry_0000/PyFAI/result_{processing_type}/data"
#     path_to_data_sigma_processed = (
#         f"entry_0000/PyFAI/result_{processing_type}/data_errors"
#     )
#     with h5py.File(filename_processed_reference, "r") as file_reference:
#         with h5py.File(filename_processed_test, "r") as file_tested:
#             data_reference = file_reference[path_to_data_signal_processed][
#                 range_index_read[0] : range_index_read[1]
#             ]
#             data_errors_reference = file_reference[path_to_data_sigma_processed][
#                 range_index_read[0] : range_index_read[1]
#             ]

#             data_processed = file_tested[path_to_data_signal_processed][:]
#             data_errors_processed = file_tested[path_to_data_sigma_processed][:]

#             assert numpy.allclose(data_reference, data_processed, equal_nan=True)
#             assert numpy.allclose(
#                 data_errors_reference, data_errors_processed, equal_nan=True
#             )


# def h5py_equivalence(
#     filename_processed_reference,
#     filename_processed_test,
#     path_to_parameters="/entry_0000/PyFAI/parameters",
#     path_to_mcs="/entry_0000/PyFAI/MCS",
#     path_to_tfg="/entry_0000/PyFAI/TFG",
#     nb_frames=None,
# ):
#     """
#     Compare two HDF5 files for equivalence.
#     """
#     # if task is not None:
#     #     file_reference = task.filename_metadata
#     #     file_output = task.processing_filename
#     #     path_to_parameters = task.path_to_parameters
#     #     path_to_mcs = task.path_to_mcs
#     #     path_to_tfg = task.path_to_tfg
#     #     nb_frames = task.range_index_read[1] - task.range_index_read[0]

#     def check_existence(name):
#         SKIP_GROUPS = [
#             "entry_0000/PyFAI/result_",
#             "entry_0000/PyFAI/MCS/raw",
#             "entry_0000/PyFAI/MCS/interpreted",
#             # "entry_0000/PyFAI/MCS/HS32C",
#             "TitleExtension",
#             "entry_0000/ewoks",
#             "entry_0000/pyFAI/benchmark",
#         ]
#         for gr in SKIP_GROUPS:
#             if gr in name:
#                 return
#         if "ave_limits" in name:
#             return
#         if name not in fproc:
#             assert False, f"Object {name} not found in processed file"

#     with h5py.File(filename_processed_reference, "r") as fref:
#         with h5py.File(filename_processed_test, "r") as fproc:
#             # Traverse all objects (groups and datasets) in the reference file
#             fref.visit(check_existence)

#     parameters_group_equivalence(
#         file_reference=filename_processed_reference,
#         file_output=filename_processed_test,
#         path_to_parameters=path_to_parameters,
#     )

#     MCS_group_equivalence(
#         file_reference=filename_processed_reference,
#         file_output=filename_processed_test,
#         path_to_mcs=path_to_mcs,
#         nb_frames=nb_frames,
#     )

#     interpreted_group_equivalence(
#         file_reference=filename_processed_reference,
#         file_output=filename_processed_test,
#         path_to_interpreted=f"{path_to_mcs}/interpreted",
#     )

#     raw_group_equivalence(
#         file_reference=filename_processed_reference,
#         file_output=filename_processed_test,
#         path_to_raw=f"{path_to_mcs}/raw",
#     )

#     subscan1_group_equivalence(
#         file_reference=filename_processed_reference,
#         file_output=filename_processed_test,
#         path_to_subscan1=f"{path_to_mcs}/raw/subscan_1",
#     )

#     tfg_group_equivalence(
#         file_reference=filename_processed_reference,
#         file_output=filename_processed_test,
#         path_to_tfg=path_to_tfg,
#     )


# def parameters_group_equivalence(
#     file_reference,
#     file_output,
#     path_to_parameters,
# ):
#     """
#     Compare two HDF5 groups for equivalence.
#     """
#     with h5py.File(file_reference, "r") as fref:
#         with h5py.File(file_output, "r") as fproc:
#             for key in fref[path_to_parameters]:
#                 if key == "TitleExtension":
#                     continue
#                 assert key in fproc[path_to_parameters].keys()
#                 value_reference = fref[path_to_parameters][key][()]
#                 value_processed = fproc[path_to_parameters][key][()]
#                 assert value_reference == value_processed


# def MCS_group_equivalence(
#     file_reference,
#     file_output,
#     path_to_mcs,
#     nb_frames,
# ):
#     """
#     Compare all datasets in the MCS group of two HDF5 files for equivalence.
#     """
#     with h5py.File(file_reference, "r") as fref:
#         with h5py.File(file_output, "r") as fproc:
#             for key in fref[path_to_mcs]:
#                 if isinstance(fref[path_to_mcs][key], h5py.Group):
#                     continue
#                 assert key in fproc[path_to_mcs].keys()

#                 value_processed = fproc[path_to_mcs][key][()]

#                 if isinstance(value_processed, bytes):
#                     assert value_processed == fref[path_to_mcs][key][()]
#                     continue

#                 if value_processed.ndim > 0 and len(value_processed) == nb_frames:
#                     if fref[path_to_mcs][key].ndim == 2:
#                         value_reference = fref[path_to_mcs][key][
#                             0 : len(value_processed), :
#                         ]
#                     elif fref[path_to_mcs][key].ndim == 1:
#                         value_reference = fref[path_to_mcs][key][
#                             0 : len(value_processed)
#                         ]
#                 else:
#                     value_reference = numpy.array(fref[path_to_mcs][key][()])

#                 if value_processed.dtype == "O":
#                     assert numpy.array_equal(
#                         value_processed, value_reference
#                     ), f"Dataset {key} is not equal in the two files"
#                 else:
#                     assert numpy.array_equal(
#                         value_processed, value_reference, equal_nan=True
#                     ), f"Dataset {key} is not equal in the two files"


# def interpreted_group_equivalence(
#     file_reference,
#     file_output,
#     path_to_interpreted,
# ):
#     """
#     Compare all datasets in the interpreted group of two HDF5 files for equivalence.
#     """
#     with h5py.File(file_reference, "r") as fref:
#         with h5py.File(file_output, "r") as fproc:
#             for key in fproc[path_to_interpreted]:
#                 value_processed = fproc[path_to_interpreted][key][()]
#                 value_reference = fref[path_to_interpreted][key][
#                     0 : len(value_processed)
#                 ]
#                 assert numpy.array_equal(
#                     value_processed, value_reference, equal_nan=True
#                 ), f"Dataset {key} is not equal in the two files"


# def raw_group_equivalence(
#     file_reference,
#     file_output,
#     path_to_raw,
# ):
#     """
#     Compare all datasets in the raw group of two HDF5 files for equivalence.
#     """
#     with h5py.File(file_reference, "r") as fref:
#         with h5py.File(file_output, "r") as fproc:
#             for key in fref[path_to_raw]:
#                 if isinstance(fref[path_to_raw][key], h5py.Group):
#                     continue

#                 value_processed = fproc[path_to_raw][key][()]
#                 value_reference = fref[path_to_raw][key][0 : len(value_processed)]
#                 assert numpy.array_equal(
#                     value_processed, value_reference, equal_nan=True
#                 ), f"Dataset {key} is not equal in the two files"


# def subscan1_group_equivalence(
#     file_reference,
#     file_output,
#     path_to_subscan1,
# ):
#     """
#     Compare all datasets in the subscan1 group of two HDF5 files for equivalence.
#     """
#     with h5py.File(file_reference, "r") as fref:
#         with h5py.File(file_output, "r") as fproc:
#             for key in fref[path_to_subscan1]:
#                 value_processed = fproc[path_to_subscan1][key][()]
#                 value_reference = fref[path_to_subscan1][key][0 : len(value_processed)]
#                 assert numpy.array_equal(
#                     value_processed, value_reference, equal_nan=True
#                 ), f"Dataset {key} is not equal in the two files"


# def tfg_group_equivalence(
#     file_reference,
#     file_output,
#     path_to_tfg,
# ):
#     """
#     Compare all datasets in the tfg group of two HDF5 files for equivalence.
#     """
#     with h5py.File(file_reference, "r") as fref:
#         with h5py.File(file_output, "r") as fproc:
#             for key in fref[path_to_tfg]:
#                 value_processed = fproc[path_to_tfg][key][()]
#                 if isinstance(value_processed, bytes):
#                     assert value_processed == fref[path_to_tfg][key][()]
#                     continue

#                 value_reference = fref[path_to_tfg][key][0 : len(value_processed)]
#                 assert numpy.array_equal(
#                     value_processed, value_reference, equal_nan=True
#                 ), f"Dataset {key} is not equal in the two files"


# def check_result(
#     reference_filename,
#     output_filename,
#     range_index_read,
#     processing_type,
#     **kwargs,
# ):
#     path_to_data_signal = f"entry_0000/PyFAI/result_{processing_type}/data"
#     with h5py.File(reference_filename, "r") as file_reference:
#         with h5py.File(output_filename, "r") as file_processed:
#             dset_signal_reference = file_reference[path_to_data_signal]

#             if dset_signal_reference.ndim == 3:
#                 data_reference = dset_signal_reference[
#                     range_index_read[0] : range_index_read[1], :, :
#                 ]

#             elif dset_signal_reference.ndim == 2:
#                 data_reference = dset_signal_reference[
#                     range_index_read[0] : range_index_read[1], :
#                 ]

#             data_processed = file_processed[path_to_data_signal][:]
#             assert numpy.allclose(
#                 data_reference, data_processed, equal_nan=True
#             ), f"Data signal is not equal: {data_reference.mean()} != {data_processed.mean()}"


def execute_ewoks(
    graph,
    inputs,
    engine="ppf",
    pool_type="thread",
):
    """
    Execute a graph with the given inputs and return the result.
    """
    result = execute_graph(
        graph=graph,
        inputs=inputs,
        engine=engine,
        pool_type=pool_type,
    )
    return result
