from functools import partial
import importlib
import pytest

from pathlib import Path
import zarr
import dask.array as dsa
import dask
import dask.core
import xarray
import numpy

from rechunker import api


_DIMENSION_KEY = "_ARRAY_DIMENSIONS"


def requires_import(module, *args):
    try:
        importlib.import_module(module)
    except ImportError:
        skip = True
    else:
        skip = False
    mark = pytest.mark.skipif(skip, reason=f"requires {module}")
    return pytest.param(*args, marks=mark)


requires_beam = partial(requires_import, "apache_beam")
requires_prefect = partial(requires_import, "prefect")
requires_pywren = partial(requires_import, "pywren_ibm_cloud")


@pytest.fixture(params=[(8000, 200), {"y": 8000, "x": 200}])
def target_chunks(request):
    return request.param


def test_invalid_executor():
    with pytest.raises(ValueError, match="unrecognized executor"):
        api._get_executor("unknown")


@pytest.mark.parametrize("shape", [(100, 50)])
@pytest.mark.parametrize("source_chunks", [(10, 50)])
@pytest.mark.parametrize("target_chunks", [(20, 10)])
@pytest.mark.parametrize("max_mem", ["10MB"])
@pytest.mark.parametrize("executor", ["dask"])
def test_rechunk_dataset(
    tmp_path, shape, source_chunks, target_chunks, max_mem, executor
):
    target_store = str(tmp_path / "target.zarr")
    temp_store = str(tmp_path / "temp.zarr")

    a = numpy.arange(numpy.prod(shape)).reshape(shape).astype("f4")
    a[-1] = numpy.nan
    ds = xarray.Dataset(
        dict(
            a=xarray.DataArray(
                a, dims=["x", "y"], attrs={"a1": 1, "a2": [1, 2, 3], "a3": "x"}
            ),
            b=xarray.DataArray(numpy.ones(shape[0]), dims=["x"]),
            c=xarray.DataArray(numpy.ones(shape[1]), dims=["y"]),
        ),
        coords=dict(
            cx=xarray.DataArray(numpy.ones(shape[0]), dims=["x"]),
            cy=xarray.DataArray(numpy.ones(shape[1]), dims=["y"]),
        ),
        attrs={"a1": 1, "a2": [1, 2, 3], "a3": "x"},
    )
    ds = ds.chunk(chunks=dict(zip(["x", "y"], source_chunks)))
    options = dict(
        a=dict(
            compressor=zarr.Blosc(cname="zstd"),
            dtype="int32",
            scale_factor=0.1,
            _FillValue=-9999,
        )
    )
    rechunked = api.rechunk(
        ds,
        target_chunks=dict(a=target_chunks, b=target_chunks[:1]),
        max_mem=max_mem,
        target_store=target_store,
        target_options=options,
        temp_store=temp_store,
        executor=executor,
    )
    assert isinstance(rechunked, api.Rechunked)
    rechunked.execute()

    # Validate encoded variables
    dst = xarray.open_zarr(target_store, decode_cf=False)
    assert dst.a.dtype == options["a"]["dtype"]
    assert all(dst.a.values[-1] == options["a"]["_FillValue"])
    assert dst.a.encoding["compressor"] is not None

    # Validate decoded variables
    dst = xarray.open_zarr(target_store, decode_cf=True)
    assert dst.a.data.chunksize == target_chunks
    assert dst.b.data.chunksize == target_chunks[:1]
    assert dst.c.data.chunksize == source_chunks[1:]
    xarray.testing.assert_equal(ds.compute(), dst.compute())
    assert ds.attrs == dst.attrs


@pytest.mark.parametrize("shape", [(8000, 8000)])
@pytest.mark.parametrize("source_chunks", [(200, 8000)])
@pytest.mark.parametrize("dtype", ["f4"])
@pytest.mark.parametrize("max_mem", [25600000, "25.6MB"])
@pytest.mark.parametrize(
    "executor",
    [
        "dask",
        "python",
        requires_beam("beam"),
        requires_prefect("prefect"),
        requires_pywren("pywren"),
    ],
)
@pytest.mark.parametrize(
    "dims,target_chunks",
    [
        (None, (8000, 200)),
        # would be nice to support this syntax eventually
        pytest.param(None, (-1, 200), marks=pytest.mark.xfail),
        (["y", "x"], (8000, 200)),
        (["y", "x"], {"y": 8000, "x": 200}),
        # can't infer missing dimension chunk specification
        pytest.param(["y", "x"], {"x": 200}, marks=pytest.mark.xfail),
        # can't use dict syntax without array dims
        pytest.param(None, {"y": 8000, "x": 200}, marks=pytest.mark.xfail),
    ],
)
def test_rechunk_array(
    tmp_path, shape, source_chunks, dtype, dims, target_chunks, max_mem, executor
):

    ### Create source array ###
    store_source = str(tmp_path / "source.zarr")
    source_array = zarr.ones(
        shape, chunks=source_chunks, dtype=dtype, store=store_source
    )
    # add some attributes
    source_array.attrs["foo"] = "bar"
    if dims:
        source_array.attrs[_DIMENSION_KEY] = dims

    ### Create targets ###
    target_store = str(tmp_path / "target.zarr")
    temp_store = str(tmp_path / "temp.zarr")

    rechunked = api.rechunk(
        source_array,
        target_chunks,
        max_mem,
        target_store,
        temp_store=temp_store,
        executor=executor,
    )
    assert isinstance(rechunked, api.Rechunked)

    target_array = zarr.open(target_store)

    if isinstance(target_chunks, dict):
        target_chunks_list = [target_chunks[d] for d in dims]
    else:
        target_chunks_list = target_chunks
    assert target_array.chunks == tuple(target_chunks_list)
    assert dict(source_array.attrs) == dict(target_array.attrs)

    result = rechunked.execute()
    assert isinstance(result, zarr.Array)
    a_tar = dsa.from_zarr(target_array)
    assert dsa.equal(a_tar, 1).all().compute()


@pytest.mark.parametrize("shape", [(8000, 8000)])
@pytest.mark.parametrize("source_chunks", [(200, 8000), (800, 8000)])
@pytest.mark.parametrize("dtype", ["f4"])
@pytest.mark.parametrize("max_mem", [25600000])
@pytest.mark.parametrize(
    "target_chunks", [(200, 8000), (800, 8000), (8000, 200), (400, 8000),],
)
def test_rechunk_dask_array(
    tmp_path, shape, source_chunks, dtype, target_chunks, max_mem
):

    ### Create source array ###
    source_array = dsa.ones(shape, chunks=source_chunks, dtype=dtype)

    ### Create targets ###
    target_store = str(tmp_path / "target.zarr")
    temp_store = str(tmp_path / "temp.zarr")

    rechunked = api.rechunk(
        source_array, target_chunks, max_mem, target_store, temp_store=temp_store
    )
    assert isinstance(rechunked, api.Rechunked)

    target_array = zarr.open(target_store)

    assert target_array.chunks == tuple(target_chunks)

    result = rechunked.execute()
    assert isinstance(result, zarr.Array)
    a_tar = dsa.from_zarr(target_array)
    assert dsa.equal(a_tar, 1).all().compute()


@pytest.mark.parametrize(
    "executor",
    [
        "dask",
        "python",
        requires_beam("beam"),
        requires_prefect("prefect"),
        requires_pywren("pywren"),
    ],
)
def test_rechunk_group(tmp_path, executor):
    store_source = str(tmp_path / "source.zarr")
    group = zarr.group(store_source)
    group.attrs["foo"] = "bar"
    # 800 byte chunks
    a = group.ones("a", shape=(5, 10, 20), chunks=(1, 10, 20), dtype="f4")
    a.attrs["foo"] = "bar"
    b = group.ones("b", shape=(20,), chunks=(10,), dtype="f4")
    b.attrs["foo"] = "bar"

    target_store = str(tmp_path / "target.zarr")
    temp_store = str(tmp_path / "temp.zarr")

    max_mem = 1600  # should force a two-step plan for a
    target_chunks = {"a": (5, 10, 4), "b": (20,)}

    rechunked = api.rechunk(
        group,
        target_chunks,
        max_mem,
        target_store,
        temp_store=temp_store,
        executor=executor,
    )
    assert isinstance(rechunked, api.Rechunked)

    target_group = zarr.open(target_store)
    assert "a" in target_group
    assert "b" in target_group
    assert dict(group.attrs) == dict(target_group.attrs)

    rechunked.execute()
    for aname in target_chunks:
        assert target_group[aname].chunks == target_chunks[aname]
        a_tar = dsa.from_zarr(target_group[aname])
        assert dsa.equal(a_tar, 1).all().compute()


def sample_xarray_dataset():
    return xarray.Dataset(
        dict(
            a=xarray.DataArray(
                dsa.ones(shape=(10, 20, 40), chunks=(5, 10, 4), dtype="f4"),
                dims=("x", "y", "z"),
                attrs={"foo": "bar"},
            ),
            b=xarray.DataArray(
                dsa.ones(shape=(8000,), chunks=(200,), dtype="f4"),
                dims="w",
                attrs={"foo": "bar"},
            ),
        ),
        attrs={"foo": "bar"},
    )


def sample_zarr_group(tmp_path):
    path = str(tmp_path / "source.zarr")
    group = zarr.group(path)
    group.attrs["foo"] = "bar"
    # 800 byte chunks
    a = group.ones("a", shape=(10, 20, 40), chunks=(5, 10, 4), dtype="f4")
    a.attrs["foo"] = "bar"
    b = group.ones("b", shape=(8000,), chunks=(200,), dtype="f4")
    b.attrs["foo"] = "bar"
    return group


def sample_zarr_array(tmp_path):
    shape = (8000, 8000)
    source_chunks = (200, 8000)
    dtype = "f4"
    dims = None

    path = str(tmp_path / "source.zarr")
    array = zarr.ones(shape, chunks=source_chunks, dtype=dtype, store=path)
    # add some attributes
    array.attrs["foo"] = "bar"
    if dims:
        array.attrs[_DIMENSION_KEY] = dims
    return array


@pytest.fixture(params=["Array", "Group", "Dataset"])
def rechunk_args(tmp_path, request):
    if request.param == "Dataset":
        ds = sample_xarray_dataset()

        target_store = str(tmp_path / "target.zarr")
        temp_store = str(tmp_path / "temp.zarr")
        max_mem = 1600  # should force a two-step plan for a and b
        target_chunks = {"a": (10, 5, 4), "b": (100,)}

        args = dict(
            source=ds,
            target_chunks=target_chunks,
            max_mem=max_mem,
            target_store=target_store,
            temp_store=temp_store,
        )
    elif request.param == "Group":
        group = sample_zarr_group(tmp_path)

        target_store = str(tmp_path / "target.zarr")
        temp_store = str(tmp_path / "temp.zarr")
        max_mem = 1600  # should force a two-step plan for a and b
        target_chunks = {"a": (10, 5, 4), "b": (100,)}

        args = dict(
            source=group,
            target_chunks=target_chunks,
            max_mem=max_mem,
            target_store=target_store,
            temp_store=temp_store,
        )
    else:
        array = sample_zarr_array(tmp_path)

        target_store = str(tmp_path / "target.zarr")
        temp_store = str(tmp_path / "temp.zarr")
        max_mem = 25600000
        target_chunks = (8000, 200)

        args = dict(
            source=array,
            target_chunks=target_chunks,
            max_mem=max_mem,
            target_store=target_store,
            temp_store=temp_store,
        )
    return args


@pytest.fixture()
def rechunked(rechunk_args):
    return api.rechunk(**rechunk_args)


def test_repr(rechunked):
    assert isinstance(rechunked, api.Rechunked)
    repr_str = repr(rechunked)

    assert repr_str.startswith("<Rechunked>")
    assert all(thing in repr_str for thing in ["Source", "Intermediate", "Target"])


def test_repr_html(rechunked):
    rechunked._repr_html_()  # no exceptions


def _is_collection(source):
    assert isinstance(
        source,
        (dask.array.Array, zarr.core.Array, zarr.hierarchy.Group, xarray.Dataset),
    )
    return isinstance(source, (zarr.hierarchy.Group, xarray.Dataset))


def _wrap_options(source, options):
    if _is_collection(source):
        options = {v: options for v in source}
    return options


def test_rechunk_option_overwrite(rechunk_args):
    api.rechunk(**rechunk_args).execute()
    # TODO: make this match more reliable based on outcome of
    # https://github.com/zarr-developers/zarr-python/issues/605
    with pytest.raises(ValueError, match=r"path .* contains an array"):
        api.rechunk(**rechunk_args).execute()
    options = _wrap_options(rechunk_args["source"], dict(overwrite=True))
    api.rechunk(**rechunk_args, target_options=options).execute()


def test_rechunk_passthrough(rechunk_args):
    # Verify that no errors are raised when the target chunks == source chunks
    if _is_collection(rechunk_args["source"]):
        rechunk_args["target_chunks"] = {v: None for v in rechunk_args["source"]}
    else:
        rechunk_args["target_chunks"] = None
    api.rechunk(**rechunk_args).execute()


def test_rechunk_no_temp_dir_provided_error(rechunk_args):
    # Verify that the correct error is raised when no temp_store is given
    # and the chunks to write differ from the chunks to read
    args = {k: v for k, v in rechunk_args.items() if k != "temp_store"}
    with pytest.raises(ValueError, match="A temporary store location must be provided"):
        api.rechunk(**args).execute()


def test_rechunk_option_compression(rechunk_args):
    def rechunk(compressor):
        options = _wrap_options(
            rechunk_args["source"], dict(overwrite=True, compressor=compressor)
        )
        rechunked = api.rechunk(**rechunk_args, target_options=options)
        rechunked.execute()
        return sum(
            file.stat().st_size
            for file in Path(rechunked._target.store.path).rglob("*")
        )

    size_uncompressed = rechunk(None)
    size_compressed = rechunk(
        zarr.Blosc(cname="zstd", clevel=9, shuffle=zarr.Blosc.SHUFFLE)
    )
    assert size_compressed < size_uncompressed


def test_rechunk_invalid_option(rechunk_args):
    if isinstance(rechunk_args["source"], xarray.Dataset):
        # Options are essentially unbounded for Xarray (for CF encoding params),
        # so check only options with special error cases
        options = _wrap_options(rechunk_args["source"], {"chunks": 10})
        with pytest.raises(
            ValueError,
            match="Chunks must be provided in ``target_chunks`` rather than options",
        ):
            api.rechunk(**rechunk_args, target_options=options)
    else:
        for o in ["shape", "chunks", "dtype", "store", "name", "unknown"]:
            options = _wrap_options(rechunk_args["source"], {o: True})
            with pytest.raises(ValueError, match=f"Zarr options must not include {o}"):
                api.rechunk(**rechunk_args, temp_options=options)
            with pytest.raises(ValueError, match=f"Zarr options must not include {o}"):
                api.rechunk(**rechunk_args, target_options=options)


def test_rechunk_bad_target_chunks(rechunk_args):
    if not _is_collection(rechunk_args["source"]):
        return
    rechunk_args = dict(rechunk_args)
    rechunk_args["target_chunks"] = (10, 10)
    with pytest.raises(
        ValueError, match="You must specify ``target-chunks`` as a dict"
    ):
        api.rechunk(**rechunk_args)


def test_rechunk_invalid_source(tmp_path):
    with pytest.raises(
        ValueError,
        match="Source must be a Zarr Array, Zarr Group, Dask Array or Xarray Dataset",
    ):
        api.rechunk(
            [[1, 2], [3, 4]], target_chunks=(10, 10), max_mem=100, target_store=tmp_path
        )


@pytest.mark.parametrize(
    "source,target_chunks",
    [
        (sample_xarray_dataset(), {"a": (10, 5, 4), "b": (100,)}),
        (dsa.ones((20, 10), chunks=(5, 5)), (10, 10)),
    ],
)
@pytest.mark.parametrize(
    "executor",
    [
        "python",
        requires_beam("beam"),
        requires_prefect("prefect"),
        requires_pywren("pywren"),
    ],
)
def test_unsupported_executor(tmp_path, source, target_chunks, executor):
    with pytest.raises(
        NotImplementedError, match="Executor type .* not supported for source",
    ):
        api.rechunk(
            source,
            target_chunks=target_chunks,
            max_mem=1600,
            target_store=str(tmp_path / "target.zarr"),
            temp_store=str(tmp_path / "temp.zarr"),
            executor=executor,
        )


def test_rechunk_no_target_chunks(rechunk_args):
    rechunk_args = dict(rechunk_args)
    if _is_collection(rechunk_args["source"]):
        rechunk_args["target_chunks"] = {v: None for v in rechunk_args["source"]}
    else:
        rechunk_args["target_chunks"] = None
    api.rechunk(**rechunk_args)


def test_no_intermediate():
    a = zarr.ones((4, 4), chunks=(2, 2))
    b = zarr.ones((4, 4), chunks=(4, 1))
    rechunked = api.Rechunked(None, None, source=a, intermediate=None, target=b)
    assert "Intermediate" not in repr(rechunked)
    rechunked._repr_html_()


def test_no_intermediate_fused(tmp_path):
    shape = (8000, 8000)
    source_chunks = (200, 8000)
    dtype = "f4"
    max_mem = 25600000
    target_chunks = (400, 8000)

    store_source = str(tmp_path / "source.zarr")
    source_array = zarr.ones(
        shape, chunks=source_chunks, dtype=dtype, store=store_source
    )

    target_store = str(tmp_path / "target.zarr")

    rechunked = api.rechunk(source_array, target_chunks, max_mem, target_store)

    num_tasks = len([v for v in rechunked.plan.dask.values() if dask.core.istask(v)])
    assert num_tasks < 20  # less than if no fuse


def test_pywren_function_executor(tmp_path):
    pytest.importorskip("pywren_ibm_cloud")
    from rechunker.executors.pywren import (
        pywren_local_function_executor,
        PywrenExecutor,
    )

    # Create a Pywren function exectutor that we manage ourselves
    # and pass in to rechunker's PywrenExecutor
    with pywren_local_function_executor() as function_executor:

        executor = PywrenExecutor(function_executor)

        shape = (8000, 8000)
        source_chunks = (200, 8000)
        dtype = "f4"
        max_mem = 25600000
        target_chunks = (400, 8000)

        ### Create source array ###
        store_source = str(tmp_path / "source.zarr")
        source_array = zarr.ones(
            shape, chunks=source_chunks, dtype=dtype, store=store_source
        )

        ### Create targets ###
        target_store = str(tmp_path / "target.zarr")
        temp_store = str(tmp_path / "temp.zarr")

        rechunked = api.rechunk(
            source_array,
            target_chunks,
            max_mem,
            target_store,
            temp_store=temp_store,
            executor=executor,
        )
        assert isinstance(rechunked, api.Rechunked)

        target_array = zarr.open(target_store)

        assert target_array.chunks == tuple(target_chunks)

        result = rechunked.execute()
        assert isinstance(result, zarr.Array)
        a_tar = dsa.from_zarr(target_array)
        assert dsa.equal(a_tar, 1).all().compute()
