from __future__ import print_function, division
import unittest


class TestSanityOfPyaslExt1(unittest.TestCase):

    def setUp(self):
        pass

    def tearDown(self):
        pass

    def testsanity_ExampleTransmissionCurves(self):
        """
        Checking example for 'TransmissionCurves'
        """
        # from __future__ import print_function
        from PyAstronomy import pyasl
        import numpy as np
        import matplotlib.pylab as plt
        # Import six for Python 2/3 compatibility
        import six

        # Get transmission curve object
        tcs = pyasl.TransmissionCurves()

        # Add passbands from Spitzer IRAC
        tcs.addSpitzerIRACPassbands()

        print("Available bands: ", tcs.availableBands())

        # Wavelength axis
        wvl = np.linspace(3000, 10000, 10000)

        # Plot transmission curves for Bessel b, v, and r bands
        for (b, c) in six.iteritems({"b": "b", "v": "k", "r": "r"}):
            tc = tcs.getTransCurve("Bessel " + b)
            trans = tc(wvl)
            plt.plot(wvl, trans, c+'-', label="Bessel " + b)

        # Plot transmission curves for Johnson U, B, and V bands
        for (b, c) in six.iteritems({"U": "m", "B": "b", "V": "k"}):
            tc = tcs.getTransCurve("Johnson " + b)
            trans = tc(wvl)
            plt.plot(wvl, trans, c+'--', label="Johnson " + b)

        plt.legend()
        plt.xlabel(r"Wavelength [$\AA$]")
        plt.ylabel("Transmission")
        # plt.show()

        # Create Planck spectrum ...
        wvl = np.arange(3000., 10000., 1.0)
        spec = pyasl.planck(T=5777., lam=wvl*1e-10)
        # ... and convolve with Johnson V band
        vbs = tcs.convolveWith(wvl, spec, "Johnson V")

        plt.plot(wvl, spec, 'b-', label='Input spectrum')
        plt.plot(wvl, vbs, 'r--', label='Convolution with Johnson V band')
        plt.legend()
        # plt.show()
        
    def testsanity_DownloadTransmissionCurves(self):
        """
        Checking downloads of 'TransmissionCurves'
        """
        from PyAstronomy import pyasl
        tcs = pyasl.TransmissionCurves()
        tcs.addSpitzerIRACPassbands(forceDownload=True)
        tcs.addKeplerPassband(forceDownload=True)
        tcs.addTESSPassband(forceDownload=True)
