import os
import glob
import shutil
import yaml

from IPython import embed

import pytest

import numpy as np

from pypeit.par.util import parse_pypeit_file
from pypeit.pypeitsetup import PypeItSetup
from pypeit.tests.tstutils import dev_suite_required, data_path
from pypeit.metadata import PypeItMetaData
from pypeit.spectrographs.util import load_spectrograph
from pypeit.scripts.setup import Setup

def test_read_combid():

    # ------------------------------------------------------------------
    # In case of failed tests
    setup_dir = data_path('setup_files')
    if os.path.isdir(setup_dir):
        shutil.rmtree(setup_dir)
    config_dir = data_path('shane_kast_blue_A')
    if os.path.isdir(config_dir):
        shutil.rmtree(config_dir)
    # ------------------------------------------------------------------

    # Generate the pypeit file with the comb_id
    droot = data_path('b')
    pargs = Setup.parse_args(['-r', droot, '-s', 'shane_kast_blue', '-c=all', '-b',
                             '--extension=fits.gz', '--output_path={:s}'.format(data_path(''))])
    Setup.main(pargs)
    shutil.rmtree(setup_dir)

    pypeit_file = os.path.join(config_dir, 'shane_kast_blue_A.pypeit')
    cfg_lines, data_files, frametype, usrdata, setups, _ = parse_pypeit_file(pypeit_file)

    # Get the spectrograph
    spectrograph = None
    for l in cfg_lines:
        if 'spectrograph' in l:
            spectrograph = load_spectrograph(l.split(' ')[-1])
            break
    assert spectrograph is not None, 'Did not appropriately read spectrograph'

    # Set the metadata
    pmd = PypeItMetaData(spectrograph, spectrograph.default_pypeit_par(), files=data_files,
                         usrdata=usrdata, strict=False)

    indx = pmd['filename'] == 'b27.fits.gz'
    assert pmd['comb_id'][indx] == [1], 'Incorrect combination group ID'
    assert pmd['comb_id'][np.where(~indx)[0]][0] == -1, 'Incorrect combination group ID'

    shutil.rmtree(config_dir)

@dev_suite_required
def test_lris_red_multi_400():
    file_list = glob.glob(os.path.join(os.environ['PYPEIT_DEV'], 'RAW_DATA', 'keck_lris_red',
                          'multi_400_8500_d560', '*.fits.gz'))
    cfg_lines = ['[rdx]',
                 'spectrograph = keck_lris_red']
    ps = PypeItSetup(file_list, cfg_lines=cfg_lines)
    ps.build_fitstbl()
    ps.get_frame_types(flag_unknown=True)
    cfgs = ps.fitstbl.unique_configurations()
    ps.fitstbl.set_configurations(cfgs)
    ps.fitstbl.set_calibration_groups() #global_frames=['bias', 'dark'])
    # Test
    assert np.all(ps.fitstbl['setup'] == 'A')


@dev_suite_required
def test_lris_red_multi():
    file_list = glob.glob(os.path.join(os.environ['PYPEIT_DEV'], 'RAW_DATA', 'keck_lris_red',
                          'multi*', '*.fits*'))
    cfg_lines = ['[rdx]',
                 'spectrograph = keck_lris_red']
    ps = PypeItSetup(file_list, cfg_lines=cfg_lines)
    ps.build_fitstbl()
    ps.get_frame_types(flag_unknown=True)
    cfgs = ps.fitstbl.unique_configurations()
    ps.fitstbl.set_configurations(cfgs)
    ps.fitstbl.set_calibration_groups() #global_frames=['bias', 'dark'])


@dev_suite_required
def test_lris_red_multi_calib():
    file_list = glob.glob(os.path.join(os.environ['PYPEIT_DEV'], 'RAW_DATA', 'keck_lris_red',
                          'multi_400_8500_d560', '*.fits.gz'))
    cfg_lines = ['[rdx]',
                 'spectrograph = keck_lris_red']
    ps = PypeItSetup(file_list, cfg_lines=cfg_lines)
    ps.build_fitstbl()
    ps.get_frame_types(flag_unknown=True)
    cfgs = ps.fitstbl.unique_configurations()
    ps.fitstbl.set_configurations(cfgs)
    ps.fitstbl.set_calibration_groups() #global_frames=['bias', 'dark'])

    cfile = data_path('test.calib') 
    ps.fitstbl.write_calib(cfile)
    with open(cfile, 'r') as f:
        calib = yaml.load(f, Loader=yaml.FullLoader)

    assert np.array_equal(list(calib['A'].keys()), ['--', 1]), \
            'Calibrations dictionary read incorrectly.'

    os.remove(cfile)


@dev_suite_required
def test_lris_red_multi_run():
    # Perform the setup
    file_list = glob.glob(os.path.join(os.environ['PYPEIT_DEV'], 'RAW_DATA', 'keck_lris_red',
                          'multi*', '*.fits*'))
    cfg_lines = ['[rdx]',
                 'spectrograph = keck_lris_red']
    ps = PypeItSetup(file_list, cfg_lines=cfg_lines)
    ps.run(setup_only=True)

    # Test
    #assert len(ps.setup_dict) == 2, 'Should find two setups'
    assert len(ps.fitstbl) >= 40, 'Should find 40+ files'
    arcs = ps.fitstbl['filename'][ps.fitstbl.find_frames('arc')]
    assert len(arcs) >= 2, 'Should find two or more arcs'
    assert 'r170320_2017.fits.gz' in arcs, \
            'Should have identified r170320_2017.fits.gz as an arc'
    assert 'r170816_0057.fits' in ps.fitstbl['filename'][ps.fitstbl.find_frames('science')], \
            'Should have identified r170816_0057.fits as a science frame'

    # Clean-up
    #os.remove('keck_lris_red.lst')
    #os.remove('keck_lris_red.setups')
    os.remove('keck_lris_red.sorted')


@dev_suite_required
def test_lris_blue_pypeit_overwrite():
    f = os.path.join(os.environ['PYPEIT_DEV'],
                     'pypeit_files/keck_lris_blue_long_400_3400_d560.pypeit')
    assert os.path.isfile(f), 'Could not find pypeit file.'
        
    cfg_lines, data_files, frametype, usrdata, setups, _ = parse_pypeit_file(f, file_check=False)

    # Change the dev path
    for i in range(len(data_files)):
        path_list = data_files[i].split('/')
        for j,p in enumerate(path_list):
            if p == 'RAW_DATA':
                break
        data_files[i] = os.path.join(os.environ['PYPEIT_DEV'], '/'.join(path_list[j:]))

    # Read the fits table with and without the user data
    spectrograph = load_spectrograph('keck_lris_blue')
    par = spectrograph.default_pypeit_par()
    fitstbl = PypeItMetaData(spectrograph, par, files=data_files)
    fitstbl_usr = PypeItMetaData(spectrograph, par, files=data_files, usrdata=usrdata)

    assert fitstbl['target'][0] == 'unknown', 'Grating name changed in file header'
    assert fitstbl_usr['target'][0] == 'test', 'Grating name changed in pypeit file'
    assert fitstbl['target'][0] != fitstbl_usr['target'][0], \
            'Fits header value and input pypeit file value expected to be different.'


[((547, 571), 'pypeit.tests.tstutils.data_path', 'data_path', (['"""setup_files"""'], {}), "('setup_files')\n", (556, 571), False, 'from pypeit.tests.tstutils import dev_suite_required, data_path\n'), ((579, 603), 'os.path.isdir', 'os.path.isdir', (['setup_dir'], {}), '(setup_dir)\n', (592, 603), False, 'import os\n'), ((655, 685), 'pypeit.tests.tstutils.data_path', 'data_path', (['"""shane_kast_blue_A"""'], {}), "('shane_kast_blue_A')\n", (664, 685), False, 'from pypeit.tests.tstutils import dev_suite_required, data_path\n'), ((693, 718), 'os.path.isdir', 'os.path.isdir', (['config_dir'], {}), '(config_dir)\n', (706, 718), False, 'import os\n'), ((888, 902), 'pypeit.tests.tstutils.data_path', 'data_path', (['"""b"""'], {}), "('b')\n", (897, 902), False, 'from pypeit.tests.tstutils import dev_suite_required, data_path\n'), ((1088, 1105), 'pypeit.scripts.setup.Setup.main', 'Setup.main', (['pargs'], {}), '(pargs)\n', (1098, 1105), False, 'from pypeit.scripts.setup import Setup\n'), ((1110, 1134), 'shutil.rmtree', 'shutil.rmtree', (['setup_dir'], {}), '(setup_dir)\n', (1123, 1134), False, 'import shutil\n'), ((1154, 1206), 'os.path.join', 'os.path.join', (['config_dir', '"""shane_kast_blue_A.pypeit"""'], {}), "(config_dir, 'shane_kast_blue_A.pypeit')\n", (1166, 1206), False, 'import os\n'), ((1266, 1296), 'pypeit.par.util.parse_pypeit_file', 'parse_pypeit_file', (['pypeit_file'], {}), '(pypeit_file)\n', (1283, 1296), False, 'from pypeit.par.util import parse_pypeit_file\n'), ((1949, 1974), 'shutil.rmtree', 'shutil.rmtree', (['config_dir'], {}), '(config_dir)\n', (1962, 1974), False, 'import shutil\n'), ((2268, 2311), 'pypeit.pypeitsetup.PypeItSetup', 'PypeItSetup', (['file_list'], {'cfg_lines': 'cfg_lines'}), '(file_list, cfg_lines=cfg_lines)\n', (2279, 2311), False, 'from pypeit.pypeitsetup import PypeItSetup\n'), ((2558, 2592), 'numpy.all', 'np.all', (["(ps.fitstbl['setup'] == 'A')"], {}), "(ps.fitstbl['setup'] == 'A')\n", (2564, 2592), True, 'import numpy as np\n'), ((2868, 2911), 'pypeit.pypeitsetup.PypeItSetup', 'PypeItSetup', (['file_list'], {'cfg_lines': 'cfg_lines'}), '(file_list, cfg_lines=cfg_lines)\n', (2879, 2911), False, 'from pypeit.pypeitsetup import PypeItSetup\n'), ((3432, 3475), 'pypeit.pypeitsetup.PypeItSetup', 'PypeItSetup', (['file_list'], {'cfg_lines': 'cfg_lines'}), '(file_list, cfg_lines=cfg_lines)\n', (3443, 3475), False, 'from pypeit.pypeitsetup import PypeItSetup\n'), ((3713, 3736), 'pypeit.tests.tstutils.data_path', 'data_path', (['"""test.calib"""'], {}), "('test.calib')\n", (3722, 3736), False, 'from pypeit.tests.tstutils import dev_suite_required, data_path\n'), ((3984, 4000), 'os.remove', 'os.remove', (['cfile'], {}), '(cfile)\n', (3993, 4000), False, 'import os\n'), ((4304, 4347), 'pypeit.pypeitsetup.PypeItSetup', 'PypeItSetup', (['file_list'], {'cfg_lines': 'cfg_lines'}), '(file_list, cfg_lines=cfg_lines)\n', (4315, 4347), False, 'from pypeit.pypeitsetup import PypeItSetup\n'), ((5008, 5041), 'os.remove', 'os.remove', (['"""keck_lris_red.sorted"""'], {}), "('keck_lris_red.sorted')\n", (5017, 5041), False, 'import os\n'), ((5111, 5210), 'os.path.join', 'os.path.join', (["os.environ['PYPEIT_DEV']", '"""pypeit_files/keck_lris_blue_long_400_3400_d560.pypeit"""'], {}), "(os.environ['PYPEIT_DEV'],\n    'pypeit_files/keck_lris_blue_long_400_3400_d560.pypeit')\n", (5123, 5210), False, 'import os\n'), ((5239, 5256), 'os.path.isfile', 'os.path.isfile', (['f'], {}), '(f)\n', (5253, 5256), False, 'import os\n'), ((5356, 5394), 'pypeit.par.util.parse_pypeit_file', 'parse_pypeit_file', (['f'], {'file_check': '(False)'}), '(f, file_check=False)\n', (5373, 5394), False, 'from pypeit.par.util import parse_pypeit_file\n'), ((5764, 5799), 'pypeit.spectrographs.util.load_spectrograph', 'load_spectrograph', (['"""keck_lris_blue"""'], {}), "('keck_lris_blue')\n", (5781, 5799), False, 'from pypeit.spectrographs.util import load_spectrograph\n'), ((5858, 5909), 'pypeit.metadata.PypeItMetaData', 'PypeItMetaData', (['spectrograph', 'par'], {'files': 'data_files'}), '(spectrograph, par, files=data_files)\n', (5872, 5909), False, 'from pypeit.metadata import PypeItMetaData\n'), ((5928, 5996), 'pypeit.metadata.PypeItMetaData', 'PypeItMetaData', (['spectrograph', 'par'], {'files': 'data_files', 'usrdata': 'usrdata'}), '(spectrograph, par, files=data_files, usrdata=usrdata)\n', (5942, 5996), False, 'from pypeit.metadata import PypeItMetaData\n'), ((613, 637), 'shutil.rmtree', 'shutil.rmtree', (['setup_dir'], {}), '(setup_dir)\n', (626, 637), False, 'import shutil\n'), ((728, 753), 'shutil.rmtree', 'shutil.rmtree', (['config_dir'], {}), '(config_dir)\n', (741, 753), False, 'import shutil\n'), ((2053, 2160), 'os.path.join', 'os.path.join', (["os.environ['PYPEIT_DEV']", '"""RAW_DATA"""', '"""keck_lris_red"""', '"""multi_400_8500_d560"""', '"""*.fits.gz"""'], {}), "(os.environ['PYPEIT_DEV'], 'RAW_DATA', 'keck_lris_red',\n    'multi_400_8500_d560', '*.fits.gz')\n", (2065, 2160), False, 'import os\n'), ((2668, 2760), 'os.path.join', 'os.path.join', (["os.environ['PYPEIT_DEV']", '"""RAW_DATA"""', '"""keck_lris_red"""', '"""multi*"""', '"""*.fits*"""'], {}), "(os.environ['PYPEIT_DEV'], 'RAW_DATA', 'keck_lris_red',\n    'multi*', '*.fits*')\n", (2680, 2760), False, 'import os\n'), ((3217, 3324), 'os.path.join', 'os.path.join', (["os.environ['PYPEIT_DEV']", '"""RAW_DATA"""', '"""keck_lris_red"""', '"""multi_400_8500_d560"""', '"""*.fits.gz"""'], {}), "(os.environ['PYPEIT_DEV'], 'RAW_DATA', 'keck_lris_red',\n    'multi_400_8500_d560', '*.fits.gz')\n", (3229, 3324), False, 'import os\n'), ((3820, 3856), 'yaml.load', 'yaml.load', (['f'], {'Loader': 'yaml.FullLoader'}), '(f, Loader=yaml.FullLoader)\n', (3829, 3856), False, 'import yaml\n'), ((4104, 4196), 'os.path.join', 'os.path.join', (["os.environ['PYPEIT_DEV']", '"""RAW_DATA"""', '"""keck_lris_red"""', '"""multi*"""', '"""*.fits*"""'], {}), "(os.environ['PYPEIT_DEV'], 'RAW_DATA', 'keck_lris_red',\n    'multi*', '*.fits*')\n", (4116, 4196), False, 'import os\n'), ((1067, 1080), 'pypeit.tests.tstutils.data_path', 'data_path', (['""""""'], {}), "('')\n", (1076, 1080), False, 'from pypeit.tests.tstutils import dev_suite_required, data_path\n'), ((1881, 1896), 'numpy.where', 'np.where', (['(~indx)'], {}), '(~indx)\n', (1889, 1896), True, 'import numpy as np\n')]