#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Simple S/N calculator for DESI spectra.
Use quicksim.py --help for instructions on running this program.
Based on the IDL code pro/desi_quicksim.pro by D. Schlegel (LBL)

Created 23-Jun-2014 by David Kirkby (dkirkby@uci.edu)
"""

from __future__ import absolute_import, division, print_function, unicode_literals
# The line above will help with 2to3 support.

import argparse
import os
import os.path
import numpy as np

import desimodel.simulate as sim

def main():
    # parse command-line arguments
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('-v','--verbose', action = 'store_true',
        help = 'provide verbose output on progress')
    parser.add_argument('--model', choices=['lrg','elg','star','perfect'],
        help = 'throughput model for light entering the fiber')
    parser.add_argument('--infile', type = str, default = None,
        help = 'name of file containing wavelength (Ang) and flux (1e-17 erg/cm^2/s/Ang) columns')
    parser.add_argument('--infile-skiprows', type = int, default = 2,
        help = 'number of rows to skip at the top of INFILE')
    parser.add_argument('--airmass', type = float, default = 1.0,
        help = 'airmass for atmospheric extinction')
    parser.add_argument('--exptime', type = float, default = None,
        help = 'overrides exposure time specified in the parameter file (secs)')
    parser.add_argument('--outfile', type = str, default = None,
        help = 'optional output file name')
    parser.add_argument('--downsampling', type = int, default = 5,
        help = 'wavelength downsampling factor to use for SNR calculations')
    args = parser.parse_args()

    # We require that the DESIMODEL environment variable is set
    if 'DESIMODEL' not in os.environ:
        raise RuntimeError('The environment variable DESIMODEL must be set.')
    root = os.environ['DESIMODEL']

    # Check for required arguments.
    if args.infile is None:
        print('A source spectrum must be specified with the --infile option.')
        return -1
    if args.model is None:
        print('A source model type must be specified with the --model option.')
        return -1

    # Load the source spectrum to use.
    if not os.path.isfile(args.infile):
        print('No such infile: %s' % args.infile)
        return -1
    with open(args.infile) as stream:
        srcSpectrum = np.transpose(np.loadtxt(stream,skiprows = args.infile_skiprows))

    # Initialize the atmosphere model.
    if args.verbose:
        print('Initializing atmosphere model...')
    atmos = sim.Atmosphere(
        os.path.join(root,'data','spectra','spec-sky.dat'),
        os.path.join(root,'data','spectra','ZenithExtinction-KPNO.dat'))

    # Initialize the instrument model.
    if args.verbose:
        print('Initializing instrument model...')
    instr = sim.Instrument(
        os.path.join(root,'data','desi.yaml'),
        os.path.join(root,'data','throughput'),
        os.path.join(root,'data','specpsf','psf-quicksim.fits'))

    # Perform a quick simulation of the observed spectrum.
    if args.verbose:
        print('Running quick simulation...')
    qsim = sim.Quick(atmos,instr)
    results = qsim.simulate(sourceType=args.model,sourceSpectrum=srcSpectrum,
        airmass=args.airmass,expTime=args.exptime,downsampling=args.downsampling)

    # Calculate and print the median SNR and total SNR^2.
    snr = results[:,3]
    medianSNR = np.median(snr)
    totalSNR2 = np.sum(snr**2)
    print('Median S/N = %.3f, Total (S/N)^2 = %.1f' % (medianSNR,totalSNR2))

    # Save the results if requested (using the same format as the original IDL code).
    if args.outfile:
        if args.verbose:
            print('Saving results to %s' % args.outfile)
        with open(args.outfile,'w') as out:
            print('# INFILE=',args.infile,file=out)
            print('# AIRMASS=',qsim.airmass,file=out)
            print('# MODEL=',qsim.sourceType,file=out)
            print('# EXPTIME=',qsim.expTime,file=out)
            print('#',file=out)
            print('# Median (S/N)=',medianSNR,file=out)
            print('# Total (S/N)^2=',totalSNR2,file=out)
            print('#',file=out)
            print('# Wave    Flux        Invvar      S/N         Counts_obj  Counts_sky  Counts_read  FWHM',
                file=out)
            print('# [Ang] [e-17 erg/s/cm^2/Ang] [1/(e-17 erg/s/cm^2/Ang)^2] []',
                '[electrons] [electrons] [electrons] [Ang]',file=out)
            for (wave,flux,ivar,snr,nobj,nsky,rdnoise,psf) in results:
                print('%9.2f %11.4e %11.4e %11.4e %11.4e %11.4e %11.4e %7.3f'
                    % (wave,flux,ivar,snr,nobj,nsky,rdnoise,psf),file=out)

if __name__ == '__main__':
    main()