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)

[((228, 299), 'edibles.utils.edibles_spectrum.EdiblesSpectrum', 'EdiblesSpectrum', ([], {'filename': 'filename', 'fully_featured': '(True)', 'noDATADIR': '(True)'}), '(filename=filename, fully_featured=True, noDATADIR=True)\n', (243, 299), False, 'from edibles.utils.edibles_spectrum import EdiblesSpectrum\n'), ((1564, 1579), 'numpy.min', 'np.min', (['sp.wave'], {}), '(sp.wave)\n', (1570, 1579), True, 'import numpy as np\n'), ((1601, 1616), 'numpy.max', 'np.max', (['sp.wave'], {}), '(sp.wave)\n', (1607, 1616), True, 'import numpy as np\n'), ((1785, 1805), 'numpy.min', 'np.min', (['sp.bary_wave'], {}), '(sp.bary_wave)\n', (1791, 1805), True, 'import numpy as np\n'), ((1827, 1847), 'numpy.max', 'np.max', (['sp.bary_wave'], {}), '(sp.bary_wave)\n', (1833, 1847), True, 'import numpy as np\n'), ((2117, 2132), 'numpy.min', 'np.min', (['sp.grid'], {}), '(sp.grid)\n', (2123, 2132), True, 'import numpy as np\n'), ((2154, 2169), 'numpy.max', 'np.max', (['sp.grid'], {}), '(sp.grid)\n', (2160, 2169), True, 'import numpy as np\n'), ((2334, 2353), 'numpy.min', 'np.min', (['sp.sky_wave'], {}), '(sp.sky_wave)\n', (2340, 2353), True, 'import numpy as np\n'), ((2375, 2394), 'numpy.max', 'np.max', (['sp.sky_wave'], {}), '(sp.sky_wave)\n', (2381, 2394), True, 'import numpy as np\n'), ((2683, 2698), 'numpy.min', 'np.min', (['sp.wave'], {}), '(sp.wave)\n', (2689, 2698), True, 'import numpy as np\n'), ((2720, 2735), 'numpy.max', 'np.max', (['sp.wave'], {}), '(sp.wave)\n', (2726, 2735), True, 'import numpy as np\n'), ((2904, 2924), 'numpy.min', 'np.min', (['sp.bary_wave'], {}), '(sp.bary_wave)\n', (2910, 2924), True, 'import numpy as np\n'), ((2946, 2966), 'numpy.max', 'np.max', (['sp.bary_wave'], {}), '(sp.bary_wave)\n', (2952, 2966), True, 'import numpy as np\n'), ((3236, 3251), 'numpy.min', 'np.min', (['sp.grid'], {}), '(sp.grid)\n', (3242, 3251), True, 'import numpy as np\n'), ((3273, 3288), 'numpy.max', 'np.max', (['sp.grid'], {}), '(sp.grid)\n', (3279, 3288), True, 'import numpy as np\n'), ((3453, 3472), 'numpy.min', 'np.min', (['sp.sky_wave'], {}), '(sp.sky_wave)\n', (3459, 3472), True, 'import numpy as np\n'), ((3494, 3513), 'numpy.max', 'np.max', (['sp.sky_wave'], {}), '(sp.sky_wave)\n', (3500, 3513), True, 'import numpy as np\n')]