# Tests related to writing dask arrays to FITS files in an efficient way

import pytest

import numpy as np
from astropy.io import fits
from astropy.io.fits import ImageHDU, PrimaryHDU

da = pytest.importorskip("dask.array")


@pytest.fixture
def dask_array_in_mem():
    return da.random.uniform(-1000, 1000, (1322, 755)).rechunk((59, 55))


def test_construct_image_hdu(dask_array_in_mem):
    hdu = ImageHDU(data=dask_array_in_mem)
    assert isinstance(hdu.data, da.Array)


def test_construct_hdulist(dask_array_in_mem):
    hdu = ImageHDU(data=dask_array_in_mem)
    hdulist = fits.HDUList([hdu])
    assert isinstance(hdulist[0].data, da.Array)


def test_save_primary_hdu(dask_array_in_mem, tmp_path):

    # Saving a Primary HDU directly

    filename = tmp_path / 'test.fits'

    hdu = PrimaryHDU(data=dask_array_in_mem)
    hdu.writeto(filename)

    with fits.open(filename) as hdulist_new:
        assert isinstance(hdulist_new[0].data, np.ndarray)
        np.testing.assert_allclose(hdulist_new[0].data, dask_array_in_mem.compute())


def test_save_image_hdu(dask_array_in_mem, tmp_path):

    # Saving an image HDU directly

    filename = tmp_path / 'test.fits'

    hdu = ImageHDU(data=dask_array_in_mem)
    hdu.writeto(filename)

    with fits.open(filename) as hdulist_new:
        assert isinstance(hdulist_new[1].data, np.ndarray)
        np.testing.assert_allclose(hdulist_new[1].data, dask_array_in_mem.compute())


def test_save_hdulist(dask_array_in_mem, tmp_path):

    # Saving an HDUList

    filename = tmp_path / 'test.fits'

    hdu1 = PrimaryHDU(data=dask_array_in_mem)
    hdu2 = ImageHDU(data=np.random.random((128, 128)))
    hdu3 = ImageHDU(data=dask_array_in_mem * 2)
    hdulist = fits.HDUList([hdu1, hdu2, hdu3])
    assert isinstance(hdulist[0].data, da.Array)
    hdulist.writeto(filename)

    with fits.open(filename) as hdulist_new:
        assert isinstance(hdulist_new[0].data, np.ndarray)
        np.testing.assert_allclose(hdulist_new[0].data, dask_array_in_mem.compute())
        assert isinstance(hdulist_new[1].data, np.ndarray)
        np.testing.assert_allclose(hdulist_new[1].data, hdu2.data)
        assert isinstance(hdulist_new[2].data, np.ndarray)
        np.testing.assert_allclose(hdulist_new[2].data, dask_array_in_mem.compute() * 2)


def test_long_header(dask_array_in_mem, tmp_path):

    # Make sure things work correctly if there is a long header in the HDU.

    filename = tmp_path / 'test.fits'

    # NOTE: we deliberately set up a long header here rather than add the
    # keys one by one to hdu.header as adding the header in one go used to
    # cause issues, so this acts as a regression test.
    header = fits.Header()
    for index in range(2048):
        header[f'KEY{index:x}'] = 0.

    hdu = PrimaryHDU(data=dask_array_in_mem, header=header)
    hdu.writeto(filename)

    with fits.open(filename) as hdulist_new:
        assert len(hdulist_new[0].header) == 2053
        assert isinstance(hdulist_new[0].data, np.ndarray)
        np.testing.assert_allclose(hdulist_new[0].data, dask_array_in_mem.compute())


VALID_DTYPES = ('>i2', '<i2', '>i4', '<i4', '>i8', '<i8', '>f4', '<f4', '>f8', '<f8')


@pytest.mark.parametrize('dtype', VALID_DTYPES)
def test_dtypes(dask_array_in_mem, tmp_path, dtype):

    filename = tmp_path / 'test.fits'

    array = dask_array_in_mem.astype(dtype)

    hdu = PrimaryHDU(data=array)
    hdu.writeto(filename)

    with fits.open(filename) as hdulist_new:
        assert isinstance(hdulist_new[0].data, np.ndarray)
        np.testing.assert_allclose(hdulist_new[0].data, array.compute())


def test_scaled(dask_array_in_mem, tmp_path):

    filename = tmp_path / 'test.fits'

    hdu = PrimaryHDU(data=dask_array_in_mem)
    hdu.scale('int32', bzero=-1000, bscale=1e-6)
    hdu.writeto(filename)

    with fits.open(filename) as hdulist_new:
        assert isinstance(hdulist_new[0].data, np.ndarray)
        np.testing.assert_allclose(hdulist_new[0].data, dask_array_in_mem.compute(), atol=1e-5)


def test_scaled_minmax(dask_array_in_mem, tmp_path):

    filename = tmp_path / 'test.fits'

    hdu = PrimaryHDU(data=dask_array_in_mem)
    hdu.scale('int32',option='minmax')
    hdu.writeto(filename)

    with fits.open(filename) as hdulist_new:
        assert isinstance(hdulist_new[0].data, np.ndarray)
        np.testing.assert_allclose(hdulist_new[0].data, dask_array_in_mem.compute(), atol=1e-5)


def test_append(dask_array_in_mem, tmp_path):

    # Test append mode

    filename = tmp_path / 'test.fits'

    fits.append(filename, dask_array_in_mem)
    fits.append(filename, np.arange(10))

    with fits.open(filename) as hdulist_new:
        assert isinstance(hdulist_new[0].data, np.ndarray)
        np.testing.assert_allclose(hdulist_new[0].data, dask_array_in_mem.compute())
        assert isinstance(hdulist_new[1].data, np.ndarray)
        np.testing.assert_allclose(hdulist_new[1].data, np.arange(10))


# @pytest.mark.parametrize('mode', ['rb+', 'ab', 'ab+', 'wb', 'wb+'])
@pytest.mark.parametrize('mode', ['wb', 'wb+'])
def test_file_handle(mode, dask_array_in_mem, tmp_path):

    filename = tmp_path / 'test.fits'
    hdu1 = PrimaryHDU(data=dask_array_in_mem)
    hdu2 = ImageHDU(data=np.arange(10))
    hdulist = fits.HDUList([hdu1, hdu2])

    with filename.open(mode=mode) as fp:
        hdulist.writeto(fp)

    with fits.open(filename) as hdulist_new:
        assert isinstance(hdulist_new[0].data, np.ndarray)
        np.testing.assert_allclose(hdulist_new[0].data, dask_array_in_mem.compute())
        assert isinstance(hdulist_new[1].data, np.ndarray)
        np.testing.assert_allclose(hdulist_new[1].data, np.arange(10))