"""Tests for FHD_cal object."""
import nose.tools as nt
import os
from pyuvdata import UVCal
import pyuvdata.utils as uvutils
import pyuvdata.tests as uvtest
from pyuvdata.data import DATA_PATH
import numpy as np

# set up FHD file list
testdir = os.path.join(DATA_PATH, 'fhd_cal_data/')
testfile_prefix = '1061316296_'
obs_testfile = os.path.join(testdir, testfile_prefix + 'obs.sav')
cal_testfile = os.path.join(testdir, testfile_prefix + 'cal.sav')
settings_testfile = os.path.join(testdir, testfile_prefix + 'settings.txt')


def test_ReadFHDcalWriteReadcalfits():
    """
    FHD cal to calfits loopback test.

    Read in FHD cal files, write out as calfits, read back in and check for object
    equality.
    """
    fhd_cal = UVCal()
    calfits_cal = UVCal()
    if not uvtest.scipy_warnings:
        fhd_cal.read_fhd_cal(cal_testfile, obs_testfile, settings_file=settings_testfile)
    else:
        # numpy 1.14 introduced a new deprecation warning.
        # Should be fixed when the next scipy version comes out.
        # The number of replications of the warning varies some and must be
        # empirically discovered. It it defaults to the most common number.
        n_scipy_warnings, scipy_warn_list, scipy_category_list = \
            uvtest.get_scipy_warnings(n_scipy_warnings=605)
        uvtest.checkWarnings(fhd_cal.read_fhd_cal, [cal_testfile, obs_testfile],
                             {'settings_file': settings_testfile},
                             message=scipy_warn_list, category=scipy_category_list,
                             nwarnings=n_scipy_warnings)
    outfile = os.path.join(DATA_PATH, 'test/outtest_FHDcal_1061311664.calfits')
    fhd_cal.write_calfits(outfile, clobber=True)
    calfits_cal.read_calfits(outfile)
    nt.assert_equal(fhd_cal, calfits_cal)

    # do it again with fit gains (rather than raw)
    if not uvtest.scipy_warnings:
        fhd_cal.read_fhd_cal(cal_testfile, obs_testfile,
                             settings_file=settings_testfile, raw=False)
    else:
        # numpy 1.14 introduced a new deprecation warning.
        # Should be fixed when the next scipy version comes out.
        # The number of replications of the warning varies some and must be
        # empirically discovered. It it defaults to the most common number.
        n_scipy_warnings, scipy_warn_list, scipy_category_list = \
            uvtest.get_scipy_warnings(n_scipy_warnings=605)
        uvtest.checkWarnings(fhd_cal.read_fhd_cal, [cal_testfile, obs_testfile],
                             {'settings_file': settings_testfile, 'raw': False},
                             message=scipy_warn_list, category=scipy_category_list,
                             nwarnings=n_scipy_warnings)
    outfile = os.path.join(DATA_PATH, 'test/outtest_FHDcal_1061311664.calfits')
    fhd_cal.write_calfits(outfile, clobber=True)
    calfits_cal.read_calfits(outfile)
    nt.assert_equal(fhd_cal, calfits_cal)


def test_extra_history():
    """
    test that setting the extra_history keyword works
    """
    fhd_cal = UVCal()
    calfits_cal = UVCal()
    extra_history = 'Some extra history for testing\n'
    if not uvtest.scipy_warnings:
        fhd_cal.read_fhd_cal(cal_testfile, obs_testfile,
                             settings_file=settings_testfile,
                             extra_history=extra_history)
    else:
        # numpy 1.14 introduced a new deprecation warning.
        # Should be fixed when the next scipy version comes out.
        # The number of replications of the warning varies some and must be
        # empirically discovered. It it defaults to the most common number.
        n_scipy_warnings, scipy_warn_list, scipy_category_list = \
            uvtest.get_scipy_warnings(n_scipy_warnings=605)
        uvtest.checkWarnings(fhd_cal.read_fhd_cal, [cal_testfile, obs_testfile],
                             {'settings_file': settings_testfile,
                              'extra_history': extra_history},
                             message=scipy_warn_list, category=scipy_category_list,
                             nwarnings=n_scipy_warnings)
    outfile = os.path.join(DATA_PATH, 'test/outtest_FHDcal_1061311664.calfits')
    fhd_cal.write_calfits(outfile, clobber=True)
    calfits_cal.read_calfits(outfile)
    nt.assert_equal(fhd_cal, calfits_cal)
    nt.assert_true(extra_history in fhd_cal.history)

    # try again with a list of history strings
    extra_history = ['Some extra history for testing',
                     'And some more history as well']
    if not uvtest.scipy_warnings:
        fhd_cal.read_fhd_cal(cal_testfile, obs_testfile,
                             settings_file=settings_testfile,
                             extra_history=extra_history)
    else:
        # numpy 1.14 introduced a new deprecation warning.
        # Should be fixed when the next scipy version comes out.
        # The number of replications of the warning varies some and must be
        # empirically discovered. It it defaults to the most common number.
        n_scipy_warnings, scipy_warn_list, scipy_category_list = \
            uvtest.get_scipy_warnings(n_scipy_warnings=605)
        uvtest.checkWarnings(fhd_cal.read_fhd_cal, [cal_testfile, obs_testfile],
                             {'settings_file': settings_testfile,
                              'extra_history': extra_history},
                             message=scipy_warn_list, category=scipy_category_list,
                             nwarnings=n_scipy_warnings)
    outfile = os.path.join(DATA_PATH, 'test/outtest_FHDcal_1061311664.calfits')
    fhd_cal.write_calfits(outfile, clobber=True)
    calfits_cal.read_calfits(outfile)
    nt.assert_equal(fhd_cal, calfits_cal)
    for line in extra_history:
        nt.assert_true(line in fhd_cal.history)


def test_flags_galaxy():
    """
    test that files with time, freq and tile flags and galaxy models behave as expected
    """
    testdir = os.path.join(DATA_PATH, 'fhd_cal_data/flag_set')
    obs_testfile_flag = os.path.join(testdir, testfile_prefix + 'obs.sav')
    cal_testfile_flag = os.path.join(testdir, testfile_prefix + 'cal.sav')
    settings_testfile_flag = os.path.join(testdir, testfile_prefix + 'settings.txt')

    fhd_cal = UVCal()
    calfits_cal = UVCal()
    if not uvtest.scipy_warnings:
        fhd_cal.read_fhd_cal(cal_testfile_flag, obs_testfile_flag,
                             settings_file=settings_testfile_flag)
    else:
        # numpy 1.14 introduced a new deprecation warning.
        # Should be fixed when the next scipy version comes out.
        # The number of replications of the warning varies some and must be
        # empirically discovered. It it defaults to the most common number.
        n_scipy_warnings, scipy_warn_list, scipy_category_list = \
            uvtest.get_scipy_warnings(n_scipy_warnings=602)
        uvtest.checkWarnings(fhd_cal.read_fhd_cal, [cal_testfile_flag, obs_testfile_flag],
                             {'settings_file': settings_testfile_flag},
                             message=scipy_warn_list, category=scipy_category_list,
                             nwarnings=n_scipy_warnings)
    outfile = os.path.join(DATA_PATH, 'test/outtest_FHDcal_1061311664.calfits')
    fhd_cal.write_calfits(outfile, clobber=True)
    calfits_cal.read_calfits(outfile)
    nt.assert_equal(fhd_cal, calfits_cal)


def test_breakReadFHDcal():
    """Try various cases of missing files."""
    fhd_cal = UVCal()
    nt.assert_raises(StandardError, fhd_cal.read_fhd_cal, cal_testfile)  # Missing obs

    if not uvtest.scipy_warnings:
        uvtest.checkWarnings(fhd_cal.read_fhd_cal, [cal_testfile, obs_testfile],
                             message=['No settings file'])
    else:
        # numpy 1.14 introduced a new deprecation warning.
        # Should be fixed when the next scipy version comes out.
        # The number of replications of the warning varies some and must be
        # empirically discovered. It it defaults to the most common number.
        n_scipy_warnings, scipy_warn_list, scipy_category_list = \
            uvtest.get_scipy_warnings(n_scipy_warnings=605)
        warn_list = scipy_warn_list + ['No settings file']
        category_list = scipy_category_list + [UserWarning]
        uvtest.checkWarnings(fhd_cal.read_fhd_cal, [cal_testfile, obs_testfile],
                             message=warn_list, category=category_list,
                             nwarnings=n_scipy_warnings + 1)
    # Check only pyuvdata version history with no settings file
    nt.assert_equal(fhd_cal.history, '\n' + fhd_cal.pyuvdata_version_str)


def test_read_multi():
    """Test reading in multiple files."""
    testdir2 = os.path.join(DATA_PATH, 'fhd_cal_data/set2')
    obs_testfile_list = [obs_testfile, os.path.join(testdir2, testfile_prefix + 'obs.sav')]
    cal_testfile_list = [cal_testfile, os.path.join(testdir2, testfile_prefix + 'cal.sav')]
    settings_testfile_list = [settings_testfile, os.path.join(testdir2, testfile_prefix + 'settings.txt')]

    fhd_cal = UVCal()
    calfits_cal = UVCal()
    if not uvtest.scipy_warnings:
        uvtest.checkWarnings(fhd_cal.read_fhd_cal, [cal_testfile_list, obs_testfile_list],
                             {'settings_file': settings_testfile_list},
                             message='UVParameter diffuse_model does not match')
    else:
        # numpy 1.14 introduced a new deprecation warning.
        # Should be fixed when the next scipy version comes out.
        # The number of replications of the warning varies some and must be
        # empirically discovered. It it defaults to the most common number.
        n_scipy_warnings, scipy_warn_list, scipy_category_list = \
            uvtest.get_scipy_warnings(n_scipy_warnings=1210)
        warn_list = scipy_warn_list + ['UVParameter diffuse_model does not match']
        category_list = scipy_category_list + [UserWarning]
        uvtest.checkWarnings(fhd_cal.read_fhd_cal, [cal_testfile_list, obs_testfile_list],
                             {'settings_file': settings_testfile_list},
                             message=warn_list, category=category_list,
                             nwarnings=n_scipy_warnings + 1)
    outfile = os.path.join(DATA_PATH, 'test/outtest_FHDcal_1061311664.calfits')
    fhd_cal.write_calfits(outfile, clobber=True)
    calfits_cal.read_calfits(outfile)
    nt.assert_equal(fhd_cal, calfits_cal)


def test_break_read_multi():
    """Test errors for different numbers of files."""

    testdir2 = os.path.join(DATA_PATH, 'fhd_cal_data/set2')
    obs_testfile_list = [obs_testfile, os.path.join(testdir2, testfile_prefix + 'obs.sav')]
    cal_testfile_list = [cal_testfile, os.path.join(testdir2, testfile_prefix + 'cal.sav')]
    settings_testfile_list = [settings_testfile, os.path.join(testdir2, testfile_prefix + 'settings.txt')]

    fhd_cal = UVCal()
    nt.assert_raises(ValueError, fhd_cal.read_fhd_cal, cal_testfile_list,
                     obs_testfile_list[0], settings_file=settings_testfile_list)
    nt.assert_raises(ValueError, fhd_cal.read_fhd_cal, cal_testfile_list,
                     obs_testfile_list, settings_file=settings_testfile_list[0])
    nt.assert_raises(ValueError, fhd_cal.read_fhd_cal, cal_testfile_list,
                     obs_testfile_list + obs_testfile_list, settings_file=settings_testfile_list)
    nt.assert_raises(ValueError, fhd_cal.read_fhd_cal, cal_testfile_list,
                     obs_testfile_list, settings_file=settings_testfile_list + settings_testfile_list)
    nt.assert_raises(ValueError, fhd_cal.read_fhd_cal, cal_testfile_list[0],
                     obs_testfile_list, settings_file=settings_testfile_list[0])
    nt.assert_raises(ValueError, fhd_cal.read_fhd_cal, cal_testfile_list[0],
                     obs_testfile_list[0], settings_file=settings_testfile_list)