import numpy as np
import pytest
from pathlib import Path
from itertools import product

from astropy import units as u
from astropy.table import Table
from astropy.io import fits
from astropy.wcs import WCS
from astroquery.mast import Mast, MastMissions
from gwcs import WCS as GWCS
from specutils import SpectralRegion, Spectrum

from jdaviz.core.registries import loader_resolver_registry
from jdaviz.core.loaders.resolvers import find_matching_resolver
from jdaviz.utils import cached_uri


def test_loaders_registry(specviz_helper):
    # built-in loaders: file, url, etc
    assert len(specviz_helper.loaders) == len(loader_resolver_registry.members)


def test_open_close(specviz_helper):
    specviz_helper.app.state.dev_loaders = True

    assert specviz_helper.app.state.drawer_content == ''

    loader = specviz_helper.loaders['file']
    loader.open_in_tray()
    assert specviz_helper.app.state.drawer_content == 'loaders'
    loader.close_in_tray()
    assert specviz_helper.app.state.drawer_content == 'loaders'
    loader.close_in_tray(close_sidebar=True)
    assert specviz_helper.app.state.drawer_content == ''

    subset_plg = specviz_helper.plugins['Subset Tools']
    subset_plg.open_in_tray()
    assert specviz_helper.app.state.drawer_content == 'plugins'

    assert subset_plg._obj.loader_panel_ind is None  # loader panel not open
    subset_plg._obj.loaders['url'].open_in_tray()
    assert subset_plg._obj.loader_panel_ind == 0  # loader panel open
    subset_plg._obj.loaders['url'].close_in_tray()
    assert subset_plg._obj.loader_panel_ind is None  # loader panel not open
    subset_plg._obj.loaders['url'].close_in_tray(close_sidebar=True)
    assert specviz_helper.app.state.drawer_content == ''


def test_resolver_matching(specviz_helper):
    sp = Spectrum(spectral_axis=np.array([1, 2, 3])*u.nm,
                  flux=np.array([1, 2, 3])*u.Jy)

    res_sp = find_matching_resolver(specviz_helper.app, sp)
    assert res_sp._obj._registry_label == 'object'
    assert res_sp.format == '1D Spectrum'

    specviz_helper._load(sp)
    assert len(specviz_helper.app.data_collection) == 1


def test_dbg_access(deconfigged_helper):
    test_data = np.array([1, 2, 3])
    deconfigged_helper.loaders['object'].object = test_data

    # test that _get_loader is successful to grab the object
    # and expose the underlying error
    with pytest.raises(ValueError, match="Cannot load as image with ndim=1"):  # noqa
        deconfigged_helper._get_loader('object', 'object', 'Image')


def test_trace_importer(specviz2d_helper, spectrum2d):
    specviz2d_helper._load(spectrum2d, format='2D Spectrum')

    trace = specviz2d_helper.plugins['2D Spectral Extraction'].export_trace()

    res_sp = find_matching_resolver(specviz2d_helper.app, trace)
    assert res_sp._obj._registry_label == 'object'
    assert res_sp.format == 'Trace'

    # import through loader API
    ldr = specviz2d_helper.loaders['object']
    ldr.object = trace
    assert ldr.format == 'Trace'
    ldr.importer.data_label = 'Trace 1'
    ldr.load()
    assert specviz2d_helper.app.data_collection[-1].label == 'Trace 1'

    # import through load method
    specviz2d_helper._load(trace, data_label='Trace 2')
    assert specviz2d_helper.app.data_collection[-1].label == 'Trace 2'


def test_spectrum2d_viewer_options(deconfigged_helper, spectrum2d):
    ldr = deconfigged_helper.loaders['object']
    ldr.object = spectrum2d

    assert ldr.importer.viewer.create_new == '2D Spectrum'
    assert ldr.importer.viewer.new_label == '2D Spectrum'
    assert ldr.importer.ext_viewer.create_new == '1D Spectrum'
    assert ldr.importer.ext_viewer.new_label == '1D Spectrum'

    ldr.importer.ext_viewer.create_new = ''
    assert ldr.importer.ext_viewer == []

    assert ldr.importer.auto_extract
    ldr.load()

    # created 2D Spectrum viewer, did auto-extract,
    # but did not create 1D Spectrum viewer
    assert len(deconfigged_helper.viewers) == 1
    assert len(deconfigged_helper.app.data_collection) == 2


def test_markers_specviz2d_unit_conversion(specviz2d_helper, spectrum2d):
    data = np.zeros((5, 10))
    data[3] = np.arange(10)
    spectrum2d = Spectrum(flux=data*u.MJy, spectral_axis=data[3]*u.AA)
    specviz2d_helper.load_data(spectrum2d)


@pytest.mark.remote_data
@pytest.mark.filterwarnings(r"ignore::astropy.wcs.wcs.FITSFixedWarning")
def test_fits_spectrum2d(deconfigged_helper):
    uri = cached_uri('mast:jwst/product/jw02123-o001_v000000353_nirspec_f170lp-g235h_s2d.fits')
    if 'mast' in uri:
        ldr = deconfigged_helper.loaders['url']
        ldr.cache = True
        ldr.url = uri
    else:
        ldr = deconfigged_helper.loaders['file']
        ldr.filepath = uri

    # Default format may not be 2D Spectrum
    assert 'Image' in ldr.format.choices
    assert '2D Spectrum' in ldr.format.choices
    ldr.format = '2D Spectrum'
    assert ldr.importer._obj._parser.__class__.__name__ == 'FITSParser'

    ldr.load()

    # ensure get_data works, retrieves a Spectrum1D object, and has spectral WCS attached correctly
    sp2d = deconfigged_helper.get_data('jw02123-o001_v000000353_nirspec_f170lp-g235h_s2d')  # noqa
    assert isinstance(sp2d, Spectrum)
    assert str(sp2d.spectral_axis.unit) == 'um'

    sp1d = deconfigged_helper.get_data('jw02123-o001_v000000353_nirspec_f170lp-g235h_s2d (auto-ext)')  # noqa
    assert isinstance(sp1d, Spectrum)
    assert str(sp1d.spectral_axis.unit) == 'um'


@pytest.mark.remote_data
@pytest.mark.filterwarnings(r"ignore::astropy.wcs.wcs.FITSFixedWarning")
@pytest.mark.filterwarnings(r"ignore::asdf.exceptions.AsdfPackageVersionWarning")
def test_jwst_wfss_bsub(deconfigged_helper):
    uri = 'https://stsci.box.com/shared/static/k5l446jid348jk343m8r4vvbkzk4syom.fits'
    ldr = deconfigged_helper.loaders['url']
    ldr.cache = True
    ldr.url = uri

    assert ldr.format == '2D Spectrum'

    ldr.load()

    sp1d = deconfigged_helper.get_data('k5l446jid348jk343m8r4vvbkzk4syom (auto-ext)')  # noqa
    assert isinstance(sp1d, Spectrum)
    assert str(sp1d.spectral_axis.unit) == 'pix'


# TODO: Skip until file is uploaded to box
@pytest.mark.skip
def test_fits_spectrum_list_L3_wfss(deconfigged_helper):
    ldr = deconfigged_helper.loaders['url']
    # placeholder for WFSS L3 URI
    ldr.url = None

    # ldr = deconfigged_helper.loaders['file']
    # ldr.filepath = './jdaviz/notebooks/WFSS_fits/jw01076-o103_t0000_nircam_f356w-grismr_x1d.fits'  # noqa
    ldr.format = '1D Spectrum List'

    # 1_117 is completely masked
    sources_obj = ldr.importer.sources
    number_combos = product((1, 2), (9, 17, 23))
    sources_obj.selected = [f'Exposure {e_num}, Source ID: {s_id}'
                            for e_num, s_id in number_combos]
    ldr.load()

    assert len(deconfigged_helper.data_labels) == len(sources_obj.selected)
    dc = deconfigged_helper.app.data_collection
    assert len(dc) == len(sources_obj.selected)
    assert len(deconfigged_helper.viewers) == 1

    filestem = Path(ldr.filepath).stem
    for e_num, s_id in number_combos:
        spec = deconfigged_helper.get_data(f'{filestem}_EXP-{e_num}_ID-{s_id}')
        assert isinstance(spec, Spectrum)
        assert str(spec.spectral_axis.unit) == 'um'


@pytest.mark.remote_data
def test_resolver_url(deconfigged_helper, fake_classes_in_registries):

    def ignore_custom_loaders(existing_choices):
        return [choice for choice in existing_choices if choice not in fake_classes_in_registries]

    loader = deconfigged_helper.loaders['url']

    # no url, no valid formats
    assert len(loader.format.choices) == 0

    # non-valid input
    loader.url = 'not-valid-url'
    assert len(loader.format.choices) == 0

    # s3 input
    loader.url = "s3://stpubdata/jwst/public/jw02727/L3/t/o002/jw02727-o002_t062_nircam_clear-f277w_i2d.fits"  # noqa: E501
    assert loader._obj.url_scheme == 's3'
    assert len(loader.format.choices) > 0

    # https valid input
    loader.url = 'https://stsci.box.com/shared/static/exnkul627fcuhy5akf2gswytud5tazmw.fits'  # noqa

    # may change with future importers
    assert len(ignore_custom_loaders(loader.format.choices)) == 4
    assert loader.format.selected == 'Image'  # default may change with future importers

    # test target filtering
    assert len(loader.target.choices) > 1
    assert loader.target.selected == 'Any'
    loader.target = '1D Spectrum'

    # may change with future importers
    assert len(ignore_custom_loaders(loader.format.choices)) == 2
    assert loader.format == '1D Spectrum List'  # default may change with future importers
    assert loader.importer.data_label == 'exnkul627fcuhy5akf2gswytud5tazmw'  # noqa

    loader.target = 'Any'
    assert len(ignore_custom_loaders(loader.format.choices)) == 4
    loader.format = '2D Spectrum'
    assert loader.importer.data_label == 'exnkul627fcuhy5akf2gswytud5tazmw'  # noqa

    assert len(deconfigged_helper.app.data_collection) == 0
    assert len(deconfigged_helper.viewers) == 0

    loader.load()

    # 2D spectrum and auto-extracted 1D spectrum
    assert len(deconfigged_helper.app.data_collection) == 2
    assert len(deconfigged_helper.viewers) == 2

    with pytest.raises(ValueError, match="Failed query for URI"):
        # NOTE: this test will attempt to reach out to MAST via astroquery
        # even if cache is available.
        deconfigged_helper.load('mast:invalid')


def test_resolver_table_as_query(deconfigged_helper):
    ldr = deconfigged_helper.loaders['object']

    obs_table = Table({'id': [1, 2, 3], 'fileSetName': ['a', 'b', 'c']})
    file_table = Table({'id': [1, 2, 3], 'url': ['a', 'b', 'c']})
    invalid_table = Table({'id': [1, 2, 3], 'name': ['a', 'b', 'c']})

    assert ldr._obj.parsed_input_is_query is False

    ldr.object = obs_table
    assert ldr._obj.parsed_input_is_query is True
    assert ldr._obj.observation_table_populated is True
    assert ldr._obj.file_table_populated is False

    ldr.object = file_table
    assert ldr._obj.parsed_input_is_query is True
    assert ldr._obj.observation_table_populated is False
    assert ldr._obj.file_table_populated is True

    ldr.object = invalid_table
    assert ldr._obj.parsed_input_is_query is False
    assert ldr._obj.observation_table_populated is False
    assert ldr._obj.file_table_populated is False


@pytest.mark.filterwarnings("ignore::UserWarning")
@pytest.mark.remote_data
def test_resolver_table_as_query_astroquery(deconfigged_helper, tmp_path):
    ldr = deconfigged_helper.loaders['object']

    mast = Mast()
    coords = mast.resolve_object("M101")

    jwst = MastMissions(mission='JWST')

    datasets = jwst.query_region(coords,
                                 radius=3,
                                 select_cols=["sci_stop_time", "sci_targname",
                                              "sci_start_time", "sci_status"],
                                 sort_by=['sci_targname'])

    ldr.object = datasets
    assert ldr._obj.parsed_input_is_query is True
    assert ldr._obj.observation_table_populated is True
    assert ldr._obj.file_table_populated is False
    assert len(ldr.format.choices) == 0

    ldr.observation_table.select_rows(0)

    assert ldr._obj.file_table_populated is True
    assert ldr._obj.get_selected_url() is None

    ldr.file_table.select_rows(0)
    assert ldr._obj.get_selected_url() is not None
    # we can't guarantee any specific format choices since
    # this depends on what data is returned from MAST
    # but let's at least make sure the download was successful
    # and points to a local temporary cache file
    assert len(ldr._obj.output) > 0


def test_invoke_from_plugin(specviz_helper, spectrum1d, tmp_path):
    s = SpectralRegion(5*u.um, 6*u.um)
    local_path = str(tmp_path / 'spectral_region.ecsv')
    s.write(local_path)

    specviz_helper.load_data(spectrum1d)

    loader = specviz_helper.plugins['Subset Tools'].loaders['file']

    assert len(loader.format.choices) == 0
    loader.filepath = local_path
    assert len(loader.format.choices) > 0

    loader.load()


@pytest.mark.parametrize('order', ([0, 1, 2], [0, 2, 1], [2, 0, 1], [2, 1, 0], [1, 2, 0]))
def test_mult_data_types(deconfigged_helper, image_nddata_wcs, spectrum2d, spectrum1d, order):
    datas = [image_nddata_wcs, spectrum2d, spectrum1d]
    load_kwargs = [{'format': 'Image'},
                   {'format': '2D Spectrum', 'auto_extract': True},
                   {'format': '1D Spectrum'}]

    for i in order:
        deconfigged_helper.load(datas[i], **load_kwargs[i])

    # NOTE: 2D Spectrum will also result in auto-extracted 1D Spectrum
    assert len(deconfigged_helper.app.data_collection) == 4
    assert len(deconfigged_helper.viewers) == 3


def test_freq_wavelength_linking(deconfigged_helper, spectrum1d):
    deconfigged_helper.load(spectrum1d, format='1D Spectrum', data_label='sp_wavelength')
    sp1d_freq = Spectrum(spectral_axis=spectrum1d.spectral_axis.to(u.Hz, equivalencies=u.spectral()),  # noqa
                         flux=spectrum1d.flux,
                         uncertainty=spectrum1d.uncertainty,
                         mask=spectrum1d.mask,
                         meta=spectrum1d.meta)
    deconfigged_helper.load(sp1d_freq, format='1D Spectrum', data_label='sp_frequency')

    # flux <> flux, uncertainty <> uncertainty, wavelength <> freq, Pixel 0[x] <> Pixel 1[x]
    assert len(deconfigged_helper.app.data_collection.external_links) == 4


def test_load_image_mult_sci_extension(imviz_helper):
    # test loading an image with multiple SCI extensions and
    # ensure that automatic parenting logic is handled correctly
    arr = np.zeros((2, 2), dtype=np.float32)
    hdul = fits.HDUList([fits.PrimaryHDU(),
                        fits.ImageHDU(arr, name='SCI', ver=1),
                        fits.ImageHDU(arr, name='ERR', ver=1),
                        fits.ImageHDU(arr, name='SCI', ver=2),
                        fits.ImageHDU(arr, name='ERR', ver=2)
                         ])

    # imviz_helper._load(hdul, extension=('SCI,1', 'SCI,2', 'ERR,2'))
    imviz_helper.load_data(hdul, ext=('SCI,1', 'SCI,2', 'ERR,2'))

    assert len(imviz_helper.app.data_collection) == 3
    assert [d.label for d in imviz_helper.app.data_collection] == ['Image[SCI,1]', 'Image[SCI,2]', 'Image[ERR,2]']  # noqa

    assert imviz_helper.app._get_assoc_data_children('Image[SCI,1]') == []
    assert imviz_helper.app._get_assoc_data_children('Image[SCI,2]') == ['Image[ERR,2]']
    assert imviz_helper.app._get_assoc_data_parent('Image[ERR,2]') == 'Image[SCI,2]'


def test_loaders_extension_select(imviz_helper):
    # tests internal logic of SelectFileExtensionComponent
    arr = np.zeros((2, 2), dtype=np.float32)
    hdul = fits.HDUList([fits.PrimaryHDU(),
                        fits.ImageHDU(arr, name='SCI', ver=1),
                        fits.ImageHDU(arr, name='ERR', ver=1),
                        fits.ImageHDU(arr, name='SCI', ver=2),
                        fits.ImageHDU(arr, name='ERR', ver=2)
                         ])

    ldr = imviz_helper.loaders['object']
    ldr.object = hdul

    # test setting extension by index: [name, ver], [name,ver], and name,ver
    ldr.importer.extension = ['1: [SCI,1]', '[SCI,2]', 'ERR,2']
    assert ldr.importer.extension.selected == ['1: [SCI,1]',
                                               '3: [SCI,2]',
                                               '4: [ERR,2]']
    assert ldr.importer.extension.selected_index == [1, 3, 4]

    # if not providing ver, first match will be used
    ldr.importer.extension = 'SCI'
    assert ldr.importer.extension.selected == ['1: [SCI,1]']

    # case of providing index
    ldr.importer.extension = [1, 3]
    assert ldr.importer.extension.selected == ['1: [SCI,1]', '3: [SCI,2]']


def test_load_image_align_by(deconfigged_helper, image_nddata_wcs):
    ldr = deconfigged_helper.loaders['object']
    ldr.object = image_nddata_wcs
    assert 'Image' in ldr.format.choices
    ldr.format = 'Image'

    assert ldr.importer.align_by == 'Pixels'

    ldr.load()

    assert deconfigged_helper.plugins['Orientation'].align_by.selected == 'Pixels'

    ldr.importer.align_by = 'WCS'
    ldr.load()

    assert deconfigged_helper.plugins['Orientation'].align_by.selected == 'WCS'


@pytest.mark.remote_data
@pytest.mark.parametrize(
    ('gwcs_to_fits_sip', 'expected_cls'),
    ((True, WCS), (False, GWCS),), ids=('True-WCS', 'False-GWCS'))
@pytest.mark.filterwarnings("ignore:Some non-standard WCS keywords were excluded")
def test_gwcs_to_fits_sip(gwcs_to_fits_sip, expected_cls, deconfigged_helper):
    """Test gwcs_to_fits_sip through the importer API."""
    ldr = deconfigged_helper.loaders['url']
    ldr.url = 'https://data.science.stsci.edu/redirect/JWST/jwst-data_analysis_tools/imviz_test_data/jw00042001001_01101_00001_nrcb5_cal.fits'  # noqa
    ldr.importer.gwcs_to_fits_sip = gwcs_to_fits_sip

    ldr.load()

    data = deconfigged_helper.app.data_collection[0]
    assert isinstance(data.coords, expected_cls)


@pytest.mark.remote_data
@pytest.mark.parametrize('helper', ['deconfigged_helper', 'specviz_helper'])
def test_roman_1d_spectrum(helper, request):
    helper = request.getfixturevalue(helper)
    ldr = helper.loaders['url']
    # wfi_spec_combined_1d_r0000201001001001001_0002_WFI01.asdf
    ldr.url = 'https://stsci.box.com/shared/static/rgasl942so9hno2rq9f1xgdeolav59o5.asdf'
    assert len(ldr.importer.extension.choices) > 1
    ldr.format = '1D Spectrum'

    ldr.load()
    assert len(helper.app.data_collection) == 1


@pytest.mark.remote_data
@pytest.mark.parametrize('helper', ['deconfigged_helper', 'specviz2d_helper'])
def test_roman_2d_spectrum(helper, request):
    helper = request.getfixturevalue(helper)
    ldr = helper.loaders['url']
    # wfi_spec_decontaminated_2d_r0000201001001001001_0002_WFI01.asdf
    ldr.url = 'https://stsci.box.com/shared/static/g8yb9hxguy3aedveef9su67lesgd8c6w.asdf'
    assert len(ldr.importer.extension.choices) > 1
    ldr.format = '2D Spectrum'

    ldr.load()
    # 2D spectrum and auto-extracted 1D spectrum
    assert len(helper.app.data_collection) == 2
