import astropy
import datetime
import numpy as np
from edibles.utils.edibles_spectrum import EdiblesSpectrum


def testEdiblesSpectrum(filename="tests/HD170740_w860_redl_20140915_O12.fits"):

    # Spectrum information
    sp = EdiblesSpectrum(filename=filename, fully_featured=True, noDATADIR=True)
    assert isinstance(sp.header, astropy.io.fits.header.Header)
    assert isinstance(sp.target, str)
    assert isinstance(sp.date, str)
    assert isinstance(sp.datetime, datetime.datetime)
    assert isinstance(sp.v_bary, float)
    assert isinstance(sp.wave_units, str)
    assert isinstance(sp.flux_units, str)

    # Raw
    assert isinstance(sp.raw_wave, np.ndarray)
    assert isinstance(sp.raw_bary_wave, np.ndarray)
    assert isinstance(sp.raw_flux, np.ndarray)
    assert len(sp.raw_wave) == len(sp.raw_bary_wave)
    assert len(sp.raw_wave) == len(sp.raw_flux)

    assert isinstance(sp.raw_grid, np.ndarray)
    assert len(sp.raw_grid) == 200443  # print(len(sp.raw_grid))

    assert isinstance(sp.raw_sky_wave, np.ndarray)
    assert isinstance(sp.raw_sky_flux, np.ndarray)
    assert len(sp.raw_sky_wave) == len(sp.raw_sky_flux)

    assert isinstance(sp.wave, np.ndarray)
    assert isinstance(sp.bary_wave, np.ndarray)
    assert isinstance(sp.flux, np.ndarray)

    # getSpectrum
    xmin = 7660
    xmax = 7680
    sp.getSpectrum(xmin=xmin, xmax=xmax)
    assert xmin == sp.xmin
    assert xmax == sp.xmax
    assert isinstance(sp.wave, np.ndarray)
    assert isinstance(sp.flux, np.ndarray)
    assert len(sp.wave) == len(sp.flux)
    assert np.min(sp.wave) > sp.xmin
    assert np.max(sp.wave) < sp.xmax

    assert isinstance(sp.bary_wave, np.ndarray)
    assert isinstance(sp.bary_flux, np.ndarray)
    assert len(sp.bary_wave) == len(sp.bary_flux)
    assert np.min(sp.bary_wave) > sp.xmin
    assert np.max(sp.bary_wave) < sp.xmax

    assert isinstance(sp.grid, np.ndarray)
    assert isinstance(sp.interp_flux, np.ndarray)
    assert isinstance(sp.interp_bary_flux, np.ndarray)
    assert len(sp.grid) == len(sp.interp_flux)
    assert len(sp.grid) == len(sp.interp_bary_flux)
    assert np.min(sp.grid) > sp.xmin
    assert np.max(sp.grid) < sp.xmax

    assert isinstance(sp.sky_wave, np.ndarray)
    assert isinstance(sp.sky_flux, np.ndarray)
    assert len(sp.sky_wave) == len(sp.sky_flux)
    assert np.min(sp.sky_wave) > sp.xmin
    assert np.max(sp.sky_wave) < sp.xmax

    # shift
    zoom_xmin = 7661
    zoom_xmax = 7679
    shift = 0.05
    sp.shift(shift=shift, zoom_xmin=zoom_xmin, zoom_xmax=zoom_xmax)

    assert isinstance(sp.wave, np.ndarray)
    assert isinstance(sp.flux, np.ndarray)
    assert len(sp.wave) == len(sp.flux)
    assert np.min(sp.wave) > sp.xmin
    assert np.max(sp.wave) < sp.xmax

    assert isinstance(sp.bary_wave, np.ndarray)
    assert isinstance(sp.bary_flux, np.ndarray)
    assert len(sp.bary_wave) == len(sp.bary_flux)
    assert np.min(sp.bary_wave) > sp.xmin
    assert np.max(sp.bary_wave) < sp.xmax

    assert isinstance(sp.grid, np.ndarray)
    assert isinstance(sp.interp_flux, np.ndarray)
    assert isinstance(sp.interp_bary_flux, np.ndarray)
    assert len(sp.grid) == len(sp.interp_flux)
    assert len(sp.grid) == len(sp.interp_bary_flux)
    assert np.min(sp.grid) > sp.xmin
    assert np.max(sp.grid) < sp.xmax

    assert isinstance(sp.sky_wave, np.ndarray)
    assert isinstance(sp.sky_flux, np.ndarray)
    assert len(sp.sky_wave) == len(sp.sky_flux)
    assert np.min(sp.sky_wave) > sp.xmin
    assert np.max(sp.sky_wave) < sp.xmax


if __name__ == "__main__":

    filename = "HD170740_w860_redl_20140915_O12.fits"
    testEdiblesSpectrum(filename=filename)

[(10, 'edibles.utils.edibles_spectrum.EdiblesSpectrum', 'EdiblesSpectrum', ([], {'filename': 'filename', 'fully_featured': '(True)', 'noDATADIR': '(True)'}), False, 'from edibles.utils.edibles_spectrum import EdiblesSpectrum\n'), (46, 'numpy.min', 'np.min', (['sp.wave'], {}), True, 'import numpy as np\n'), (47, 'numpy.max', 'np.max', (['sp.wave'], {}), True, 'import numpy as np\n'), (52, 'numpy.min', 'np.min', (['sp.bary_wave'], {}), True, 'import numpy as np\n'), (53, 'numpy.max', 'np.max', (['sp.bary_wave'], {}), True, 'import numpy as np\n'), (60, 'numpy.min', 'np.min', (['sp.grid'], {}), True, 'import numpy as np\n'), (61, 'numpy.max', 'np.max', (['sp.grid'], {}), True, 'import numpy as np\n'), (66, 'numpy.min', 'np.min', (['sp.sky_wave'], {}), True, 'import numpy as np\n'), (67, 'numpy.max', 'np.max', (['sp.sky_wave'], {}), True, 'import numpy as np\n'), (78, 'numpy.min', 'np.min', (['sp.wave'], {}), True, 'import numpy as np\n'), (79, 'numpy.max', 'np.max', (['sp.wave'], {}), True, 'import numpy as np\n'), (84, 'numpy.min', 'np.min', (['sp.bary_wave'], {}), True, 'import numpy as np\n'), (85, 'numpy.max', 'np.max', (['sp.bary_wave'], {}), True, 'import numpy as np\n'), (92, 'numpy.min', 'np.min', (['sp.grid'], {}), True, 'import numpy as np\n'), (93, 'numpy.max', 'np.max', (['sp.grid'], {}), True, 'import numpy as np\n'), (98, 'numpy.min', 'np.min', (['sp.sky_wave'], {}), True, 'import numpy as np\n'), (99, 'numpy.max', 'np.max', (['sp.sky_wave'], {}), True, 'import numpy as np\n')]