import os
import random
import itertools
from typing import Set

import h5py
import numpy
import pytest
from fabio.edfimage import EdfImage

from blisswriter.io import nexus
from blisswriter.io import h5_config
from blisswriter.utils.array_order import Order
from blisswriter.writer.datasets import Hdf5DatasetWriter

from .helpers.nxw_test_math import asproduct


@pytest.mark.parametrize("scan_shape", (tuple(), (0,), (8,), (4, 8), (2, 4, 6)))
@pytest.mark.parametrize("detector_shape", (tuple(), (3,), (3, 5)))
@pytest.mark.parametrize("source_file_type", ("hdf5", "multi-hdf5", "edf", None))
def test_dataset_proxy_shapes_and_formats(
    scan_shape, detector_shape, source_file_type, tmpdir
):
    nframes_per_file = 3
    fulluri = source_file_type and "multi-" in source_file_type
    if fulluri:
        source_file_type = source_file_type.replace("multi-", "")

    mainfile = str(tmpdir.join("main.h5"))
    with nexus.nxRoot(mainfile, mode="w") as nxroot:
        nexus.nxEntry(nxroot, "entry0000")
    dsetname_generator = name_generator("dataset{:04d}")
    sourcedir_generator = directory_generator(tmpdir)

    scan_save_ndims = 0, 1, 2
    saveorders = (
        Order(order="C"),
        Order(order="F"),
        Order(order="C", caxes="all"),
        Order(order="F", caxes="all"),
    )
    publishorders = tuple(saveorders)
    nextras = 0, -2, 2

    options = [saveorders, publishorders, nextras, scan_save_ndims]
    for params in itertools.product(*options):
        saveorder, publishorder, nextra, scan_save_ndim = params
        if source_file_type == "edf" and saveorder.forder:
            continue

        # Number of points published: npoints + nextra
        if not detector_shape and source_file_type:
            continue
        if not scan_shape:
            if nextra > 0:
                continue
            elif nextra < 0:
                nextra = -1

        # Scan part of the saved dataset shape
        npoints = numpy.prod(scan_shape, dtype=int)
        if npoints:
            # Fixed length scan
            lst = list(asproduct(npoints, scan_save_ndim, includeone=False))
            if not lst:
                continue
            escan_save_shape = scan_save_shape = random.choice(lst)
        else:
            # Variable length scan
            scan_save_shape = scan_shape
            escan_save_shape = tuple(n if n else i for i, n in enumerate(scan_shape, 4))
            npoints = numpy.prod(escan_save_shape, dtype=int)

        # Expected shape after publish + reshape
        edetector_shape = detector_shape
        if saveorder.corder:
            escan_shape = escan_save_shape + edetector_shape
        else:
            escan_shape = edetector_shape + escan_save_shape
        if fulluri:
            external_images_per_file = nframes_per_file
        else:
            external_images_per_file = None

        err_msg = str(
            {
                "scan_shape": scan_shape,
                "scan_save_shape": scan_save_shape,
                "detector_shape": detector_shape,
                "saveorder": str(saveorder),
                "publishorder": str(publishorder),
                "nextra": nextra,
                "npoints": npoints,
            }
        )

        name = next(dsetname_generator)
        dproxy = Hdf5DatasetWriter(
            filename=mainfile,
            parent="/entry0000",
            device={"data_name": name, "data_alias": None, "data_info": {"units": "s"}},
            scan_shape=scan_shape,
            scan_save_shape=scan_save_shape,
            detector_shape=detector_shape,
            external_images_per_file=external_images_per_file,
            external_uri_from_file=not fulluri,
            dtype=float,
            saveorder=saveorder,
            publishorder=publishorder,
        )

        # Add data
        source_dir = next(sourcedir_generator)
        source_dir.mkdir()
        if source_file_type == "hdf5":
            data_generator = hdf5_data_generator(
                source_dir,
                "file",
                detector_shape,
                nframes=nframes_per_file,
                fulluri=fulluri,
            )
        elif source_file_type == "edf":
            data_generator = edf_data_generator(
                source_dir, "file", detector_shape, nframes=nframes_per_file
            )
        else:
            data_generator = mem_data_generator(detector_shape)
        add_data(dproxy, data_generator, npoints + nextra)

        if source_file_type == "hdf5" and fulluri:
            # Creating a VDS: no need to open the files
            os.chmod(str(source_dir), 0)

        # Make sure it exists with the expected shape
        if dproxy.is_external:
            dproxy.reshape(escan_save_shape)
            dproxy.flush()
        else:
            dproxy.flush()
            dproxy.reshape(escan_save_shape)
        assert dproxy.current_scan_save_shape == escan_save_shape, err_msg
        assert dproxy.current_shape == escan_shape, err_msg
        assert dproxy.current_detector_shape == edetector_shape, err_msg

        if source_file_type == "hdf5" and fulluri:
            # Allow access again to check the content
            os.chmod(str(source_dir), 0o755)

        _validate_data(dproxy, npoints, nextra, escan_save_shape, err_msg)


@pytest.fixture
def small_compression_limit():
    limit = h5_config.DEFAULT_COMPRESSION_LIMIT_NBYTES
    h5_config.DEFAULT_COMPRESSION_LIMIT_NBYTES = 0
    yield
    h5_config.DEFAULT_COMPRESSION_LIMIT_NBYTES = limit


@pytest.mark.parametrize(
    "compression_scheme",
    (
        None,
        "none",
        "gzip",
        "bitshuffle",
        "gzip-byteshuffle",
        "lz4-bitshuffle",
        "wrong_value",
    ),
)
def test_dataset_proxy_compression(compression_scheme, small_compression_limit, tmpdir):
    mainfile = str(tmpdir.join("main.h5"))
    with nexus.nxRoot(mainfile, mode="w") as nxroot:
        nexus.nxEntry(nxroot, "entry0000")

    npoints = 1
    detsize = 1024

    data = numpy.arange(npoints * detsize).reshape((npoints, detsize))
    chunk_options = {"compression_scheme": compression_scheme}

    if compression_scheme == "wrong_value":
        with pytest.raises(ValueError):
            dproxy = Hdf5DatasetWriter(
                filename=mainfile,
                parent="/entry0000",
                device={"data_name": "data", "data_alias": None},
                scan_shape=(0,),
                detector_shape=(detsize,),
                dtype=data.dtype,
                chunk_options=chunk_options,
            )
        return

    dproxy = Hdf5DatasetWriter(
        filename=mainfile,
        parent="/entry0000",
        device={"data_name": "data", "data_alias": None},
        scan_shape=(0,),
        detector_shape=(detsize,),
        dtype=data.dtype,
        chunk_options=chunk_options,
    )
    dproxy.add(data)
    dproxy.flush()  # flush chunk buffer

    with nexus.nxRoot(mainfile, mode="r") as nxroot:
        dset = nxroot["/entry0000/data"]
        compression = hdf5_compression_name(dset)
        if compression_scheme in (None, "gzip-byteshuffle"):
            assert compression == {"deflate", "shuffle"}
        elif compression_scheme == "none":
            assert not compression
        elif compression_scheme == "gzip":
            assert compression == {"deflate"}
        elif compression_scheme == "bitshuffle":
            assert compression == {
                "bitshuffle; see https://github.com/kiyo-masui/bitshuffle"
            }
        elif compression_scheme == "lz4-bitshuffle":
            assert compression == {
                "bitshuffle; see https://github.com/kiyo-masui/bitshuffle"
            }
        else:
            raise ValueError(compression_scheme)
        numpy.testing.assert_array_equal(data, dset[()])


def directory_generator(tmpdir):
    i = 1
    while True:
        yield tmpdir.join(f"data{i}")
        i += 1


def _validate_data(dproxy, npoints, nextra, escan_save_shape, err_msg):
    if dproxy.csaveorder:
        detector_save_axes = tuple(
            range(dproxy.scan_save_ndim, dproxy.scan_save_ndim + dproxy.detector_ndim)
        )
    else:
        detector_save_axes = tuple(range(dproxy.detector_ndim))

    with dproxy.open() as dset:
        data = dset[()].max(axis=detector_save_axes)
        expected = numpy.arange(1, npoints + 1, dtype=dproxy.dtype)
        if nextra < 0:
            expected[nextra:] = dproxy.fillvalue
        if escan_save_shape:
            expected = dproxy.saveorder.reshape(expected, escan_save_shape)
        numpy.testing.assert_array_equal(data, expected, err_msg=err_msg)


def add_data(dproxy, data_generator, npoints):
    newdata = []
    for i in range(npoints):
        newdata.append(next(data_generator))
        if random.choice([0, 0, 1]):
            if isinstance(newdata[0], tuple):
                filename, index = newdata[0]
                filename = nexus.splitUri(filename)[0]
                file_format = os.path.splitext(filename)[-1][1:]
                dproxy.add_external(newdata, file_format=file_format)
            else:
                newdata = numpy.array(newdata)
                dproxy.add_internal(newdata)
            newdata = []
    if newdata:
        if isinstance(newdata[0], tuple):
            filename, index = newdata[0]
            filename = nexus.splitUri(filename)[0]
            file_format = os.path.splitext(filename)[-1][1:]
            dproxy.add_external(newdata, file_format=file_format)
        else:
            newdata = numpy.array(newdata)
            dproxy.add_internal(newdata)


def mem_data_generator(detector_shape):
    ndim = len(detector_shape)
    if ndim:
        x = numpy.mgrid[tuple(slice(None, n) for n in detector_shape)]
        n = numpy.prod(detector_shape)
        x = x.reshape((ndim, n))
        mu = numpy.array(detector_shape) // 2
        sigma = numpy.array(detector_shape) / 6
        cov = sigma[:, numpy.newaxis].dot(sigma[numpy.newaxis, :]) / 10
        cov[range(ndim), range(ndim)] *= 10
        data = ndgaussian(x, mu, cov).reshape(detector_shape)
        m = 1
        while True:
            yield m * data
            m += 1
    else:
        m = 1
        while True:
            yield m
            m += 1


def hdf5_data_generator(tmpdir, basename, detector_shape, nframes=3, fulluri=False):
    fmt = str(tmpdir.join(basename + "{:04d}.hdf5"))
    filename_generator = name_generator(fmt)
    data_generator = mem_data_generator(detector_shape)
    iframe = nframes
    shape = (nframes,) + detector_shape
    while True:
        data = next(data_generator)
        if iframe == nframes:
            datafile = next(filename_generator)
            with nexus.nxRoot(datafile, mode="w") as nxroot:
                nxentry = nexus.nxEntry(nxroot, "entry0000")
                nxinstrument = nexus.nxInstrument(nxentry)
                nxdetector = nexus.nxDetector(nxinstrument, "detector")
                dtype = numpy.array(data).dtype
                dset = nxdetector.create_dataset(
                    "data", shape=shape, dtype=dtype, chunks=True, fillvalue=numpy.nan
                )
                dsetname = dset.name
                nexus.markDefault(dset)
                iframe = 0
                nxroot[dsetname][iframe] = data
        else:
            with nexus.nxRoot(datafile, mode="a") as nxroot:
                nxroot[dsetname][iframe] = data
        if fulluri:
            yield datafile + "::" + dsetname, iframe
        else:
            yield datafile, iframe
        iframe += 1


def edf_data_generator(tmpdir, basename, detector_shape, nframes=3):
    fmt = str(tmpdir.join(basename + "{:04d}.edf"))
    filename_generator = name_generator(fmt)
    if not detector_shape:
        detector_shape = 1, 1
    elif len(detector_shape) == 1:
        detector_shape = 1, detector_shape[0]
    data_generator = mem_data_generator(detector_shape)
    iframe = nframes
    while True:
        data = next(data_generator)
        if iframe == nframes:
            datafile = next(filename_generator)
            edffile = EdfImage(data=data)
            iframe = 0
        else:
            edffile.append_frame(data=data)
        edffile.write(datafile)
        yield datafile, iframe
        iframe += 1


def name_generator(fmt):
    i = 0
    while True:
        yield fmt.format(i)
        i += 1


def ndgaussian(x, mu, cov, normalize=False):
    """
    :param array x: ndim x n
    :param array mu: ndim
    :param array cov: ndim x ndim
    :param bool normalize:
    :returns array: n
    """
    ndim = x.shape[0]
    if ndim != len(mu) or (ndim, ndim) != cov.shape:
        raise ValueError("The dimensions of the input don't match")
    icov = numpy.linalg.inv(cov)
    result = x - mu[:, numpy.newaxis]
    result = result * icov.dot(result)
    result = numpy.exp(-result.sum(axis=0) / 2)
    if normalize:
        result /= (2 * numpy.pi) ** (ndim / 2) * numpy.linalg.det(cov) ** 0.5
    return result


def hdf5_compression_name(dset: h5py.Dataset) -> Set[str]:
    dcpl = dset.id.get_create_plist()
    filters = set()
    for index in range(dcpl.get_nfilters()):
        info = dcpl.get_filter(index)
        filters.add(info[-1].decode())
    return filters
