#!/usr/bin/env python
""" Plot the spectrum of a QSO in the IR overlayed with atmospheric
absorption and sky background emission.
"""
from barak.sed import get_bands
from barak.spec import qso_template
from barak.utilities import between, get_data_path
from barak.io import readtxt

import pylab as pl
import numpy as np
import sys

if 1:
    z = float(sys.argv[1])
    wa = np.logspace(np.log10(10000), np.log10(30000), 1e4)
    fl = qso_template(wa, z)
    datadir = get_data_path()
    J,H,K = get_bands('NIRI', ['niri_J.txt','niri_H.txt','niri_K.txt'])

    sky = readtxt(datadir + '/sky/ir_sky.dat', names='wa,em')
    sky.wa /= 1000.
    
    atm = readtxt(datadir + '/atmos_extinction/atm_trans_am1.0.dat',
                  names='wa,tr')

    fig = pl.figure(figsize=(14, 6))
    fig.subplots_adjust(left=0.05, right=0.95)
    pl.plot(atm.wa, atm.tr, 'c', lw=0.5, label='sky transmission')
    c0 = between(sky.wa, 1, 3)
    pl.plot(sky.wa[c0], 5*sky.em[c0]/sky.em[c0].max(),'y', lw=0.5,
            label='sky emission')
    pl.plot(J.wa / 1e4, J.tr, 'r', lw=2)
    pl.plot(H.wa / 1e4, H.tr, 'r', lw=2)
    pl.plot(K.wa / 1e4, K.tr, 'r', lw=2)
    pl.text(J.effective_wa / 1e4, 0.7, 'J')
    pl.text(H.effective_wa / 1e4, 0.7, 'H')
    pl.text(K.effective_wa / 1e4, 0.7, 'K')
    pl.plot(wa/ 1e4, fl/fl.max(),'b', label='QSO z=%.3f' % z)
    pl.xlim(1, 2.7)
    pl.ylim(-0.2, 1.4)
    pl.minorticks_on()
    pl.legend(frameon=0)
    pl.xlabel('$\lambda$ (microns)')
    pl.show()
