#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os, sys
from scipy.interpolate import InterpolatedUnivariateSpline
import numpy as np
import fitsio

#
#
#
def load_throughput(filename):
    """
    Load throughputs copied and pasted from summary throughput document
    DESI-0347v2.  First row is wavelengths in nanometers; subsequent rows
    are throughput components which should be multiplied together and then
    interpolated with a cubic spline.

    Returns InterpolatedUnivariateSpline instance.
    """
    tmp = np.loadtxt(filename).T
    wavelength = tmp[0]*10  #- nm -> Angstroms
    
    #- Throughput, removing spectrograph*CCD contribution
    throughput = tmp[1] / tmp[2]

    return InterpolatedUnivariateSpline(wavelength, throughput, k=3)
#
#
#
def load_fiberinput(filename):
    """
    Load fiberinput as calculated by fiberloss.py

    Returns InterpolatedUnivariateSpline instance.
    """
    tmp = np.loadtxt(filename).T
    wavelength = tmp[0]  #- nm -> Angstroms
    throughput = tmp[1]
    
    return InterpolatedUnivariateSpline(wavelength, throughput, k=3)
#
#
#
def load_spec_throughput(filename, column=-1):
    """
    Spectrograph throughputs from DESI-0334 have wavelength [nm] in the
    first column and total throughput in the last column.  However, the NIR
    throughput has the total in the second-to-the-last column, thus
    the support for the column= keyword.

    Returns InterpolatedUnivariateSpline instance.
    """
    tmp = np.loadtxt(filename)
    wavelength = tmp[:, 0] * 10  #- nm -> Angstroms
    throughput = tmp[:, column]
    return InterpolatedUnivariateSpline(wavelength, throughput, k=3)

def load_blue_spec_throughput():
    """
    HACK FOR THIS BRANCH:
    
    Load blue spectrograph throughputs, but remove e2v CCD QE and replace
    with STA type A QE.
    """
    filename = os.getenv('DESIMODEL_DIR')+'/data/inputs/throughput/DESI-0334-spectro/blue-thru.txt'
    tmp = np.loadtxt(filename)
    wavelength = tmp[:, 0] * 10  #- nm -> Angstroms
    throughput = tmp[:, -1]
    e2vqe = tmp[:, -2]
    
    #- Load STA throughput
    stafilename = os.getenv('DESIMODEL_DIR')+'/data/inputs/DESI-0336-ccd/staqe.txt'
    tmp = np.loadtxt(filename).T
    stawave = tmp[0] * 10  #- nm -> Angstroms
    sp = InterpolatedUnivariateSpline(stawave, tmp[1], k=3)
    staqe = sp(wavelength)
    
    newthru = (throughput / e2vqe) * staqe
        
    return InterpolatedUnivariateSpline(wavelength, newthru, k=3)
#
#
#   
def get_waveminmax(psffile):
    """
    return wmin, wmax as taken from the header of a PSF file
    """
    hdr = fitsio.read_header(psffile)
    return hdr['WAVEMIN'], hdr['WAVEMAX']
#
#-------------------------------------------------------------------------
#
def main():
    """
    Combine the various throughput parameters and generate a Specter
    compatible throughput fits file.
    """
    import optparse
    import yaml
    import os
    
    parser = optparse.OptionParser(usage = "%prog [options]")
    parser.add_option("-m", "--modeldir", type="string",  help="model directory")
    parser.add_option("-o", "--outdir", type="string",  help="output directory")
    # parser.add_option("-x", "--xxx",   help="some flag", action="store_true")
    opts, args = parser.parse_args()

    if opts.modeldir is None:
        if 'DESIMODEL_DIR' in os.environ:
            datadir = os.environ['DESIMODEL_DIR'] + '/data/'
        elif 'DESIMODEL' in os.environ:
            datadir = os.environ['DESIMODEL'] + '/data/'
        else:
            print >> sys.stderr, 'ERROR: you must give --modeldir or set DESIMODEL_DIR'
            exit(1)
    else:
        datadir = opts.modeldir + '/data/'

    if opts.outdir is None:
        opts.outdir = '.'

    #- Load telescope parameters
    fx = open(datadir + "/desi.yaml")
    params = yaml.load(fx)
    fx.close()

    #- Telescope geometric area m^2 -> cm^2
    params['area']['geometric_area'] *= 100**2

    #- Load atmospheric extinction
    d = fitsio.read(datadir+'/inputs/throughput/ZenithExtinction-KPNO.fits', 'EXTINCTION')
    extinction = InterpolatedUnivariateSpline(d['WAVELENGTH'], d['EXTINCTION'])

    #- Load pre-spectrograph throughputs
    thru = load_throughput(datadir + '/inputs/throughput/DESI-0347-throughput.txt')
    fiberinput = dict()
    for objtype in ['elg', 'lrg', 'sky', 'star']:
        fiberinput[objtype] = load_fiberinput(datadir + '/throughput/fiberloss-{}.dat'.format(objtype))

    #- Spectrograph throughputs
    spec = dict()
    spec['b'] = load_blue_spec_throughput()  #- HACK for altccd branch
    spec['r'] = load_spec_throughput(datadir+'/inputs/throughput/DESI-0334-spectro/red-thru.txt')
    
    #- HACK for altccd branch: use last column 250um CCD rather than 500um CCD
    spec['z'] = load_spec_throughput(datadir+'/inputs/throughput/DESI-0334-spectro/nir-thru.txt', -1)

    #- Load wavemin / wavemax from PSF files
    wmin = dict()
    wmax = dict()
    wmin['b'], wmax['b'] = get_waveminmax(datadir+'/specpsf/psf-b.fits')
    wmin['r'], wmax['r'] = get_waveminmax(datadir+'/specpsf/psf-r.fits')
    wmin['z'], wmax['z'] = get_waveminmax(datadir+'/specpsf/psf-z.fits')

    for channel in ('b', 'r', 'z'):
        dw = 0.1
        ww = np.arange(wmin[channel], wmax[channel]+dw/2, dw)
        tt = thru(ww) * spec[channel](ww)

        data = dict(wavelength=ww, throughput=tt,
                    extinction=extinction(ww),
                    fiberinput=fiberinput['elg'](ww) )
        hdr = list()
        hdr.append(dict(name='EXPTIME', value=params['exptime'], comment='default exposure time [sec]'))
        hdr.append(dict(name='GEOMAREA', value=params['area']['geometric_area'], comment='geometric area of mirror - obscurations'))
        hdr.append(dict(name='FIBERDIA', value=params['fibers']['diameter_arcsec'], comment='average fiber diameter [arcsec]'))
        hdr.append(dict(name='WAVEMIN', value=wmin[channel], comment='Minimum wavelength [Angstroms]'))
        hdr.append(dict(name='WAVEMAX', value=wmax[channel], comment='Maximum wavelength [Angstroms]'))

        outfile = opts.outdir + '/thru-{0}.fits'.format(channel)
        fitsio.write(outfile, data, header=hdr, clobber=True, extname='THROUGHPUT')
        
        #- Write another header with fiberinput for multiple object types
        data = np.rec.fromarrays([ww, fiberinput['elg'](ww), fiberinput['lrg'](ww),
                                      fiberinput['star'](ww), fiberinput['sky'](ww)],
                                names='wavelength,elg,lrg,star,sky')
        fitsio.write(outfile, data, extname='FIBERINPUT')
        
    return 0
#
#
#
sys.exit(main())