#!/usr/bin/python

import sys
import os
import time
from optparse import OptionParser
try:
    import pyfits
    from pyfits import getval
    from pyfits import getheader
except:
    from astropy.io import fits as fits
    from fits import getval
    from fits import getheader
import glob
import logging
import subprocess
from subprocess import Popen, list2cmdline
import shutil


#######################################################################################
###     Files Selection function
#######################################################################################
def keyword_based_file_selection( filelist, keyword ):
    "Choose Which Template needs to be reduced"

    # Get the different keyword values encountered
    keyword_values = list(set(map(lambda x: x[keyword], filelist)))

    # One set - nothing to choose
    if len(keyword_values) == 1:
        return keyword_values[0]

    logging.info("User selection amongst {0} choices".format(len(keyword_values)))
    for index in range(len(keyword_values)):
        # Extract this file list
        extracted_files = [x for x in filelist if x[keyword] == keyword_values[index]]
        if keyword == 'tpl_start':
            print_filelist_tpl_candidate(index, extracted_files)
        elif keyword == 'obs_start': 
            print_filelist_ob_candidate(index, extracted_files)
        else:
            raise NameError("Unsupported keyword {0}".format(keyword))

    index = input('Choose Data Set Index: ')
    if index >= 0 and index < len(keyword_values):
        return keyword_values[index]
    else :
        raise NameError("Entered index out of range")

#######################################################################################
###     Printing / Logging Functions
#######################################################################################
def print_filelist_tpl_candidate( index, filelist ):
    "Prints relevant informations about a Template for user choice"

    if len(filelist) < 1:
        message = "\tTemplate Index: {0} - Empty Files list".format(index)
    else:
        firstfile = filelist[0]
        do_catgs = list(set(map(lambda x: do_catg_from_file(x), filelist)))
        message = "\tTemplate Index: {0} - {1} {2}/{3} ".format(index,firstfile['tpl_start'], len(filelist),
                firstfile['tpl_nexp'])+','.join(do_catgs) 
    multiple_log(message)
    return

def print_filelist_ob_candidate( index, filelist ):
    "Prints relevant informations about an OB for user choice"

    if len(filelist) < 1:
        message = "\tObs Block Index: {0} - Empty Files list".format(index)
    else:
        # Extract dark files - Count different TPL STARTs
        filelist_darks = [x for x in filelist if x['tpl_id'] == "KMOS_spec_cal_dark"]
        n_dark_tpls = len(list(set(map(lambda x: x['tpl_start'], filelist_darks))))

        # Extract Flat files - Count different TPL STARTs - Get different bands
        filelist_flats = [x for x in filelist if x['tpl_id'] == "KMOS_spec_cal_calunitflat"]
        n_flat_tpls = len(list(set(map(lambda x: x['tpl_start'], filelist_flats))))
        flat_bands = list(set(map(lambda x: x['band'], filelist_flats)))
        # Extract Wave files - Count different TPL STARTs - Get different bands
        filelist_waves = [x for x in filelist if x['tpl_id'] == "KMOS_spec_cal_wave"]
        n_wave_tpls = len(list(set(map(lambda x: x['tpl_start'], filelist_waves))))
        wave_bands = list(set(map(lambda x: x['band'], filelist_waves)))
        
        message = "\tObs Block Index: {0} ({1}) - {2} Dark TPLs, {3} Flat TPLs ({4}), {5} Wavecal TPLs ({6})".format(index,
            filelist[0]['obs_start'], n_dark_tpls, n_flat_tpls, ','.join(flat_bands), n_wave_tpls, ','.join(wave_bands)) 
    multiple_log(message)
    return

def multiple_log_ob_selected( filelist ):
    "Prints relevant informations about a selected OB"

    if len(filelist) < 1:
        message = "\t-> Selected : Empty Files list"
    else:
        firstfile = filelist[0]
        do_catgs = list(set(map(lambda x: do_catg_from_file(x), filelist)))
        message = "\t-> Selected : {0}".format(firstfile['obs_start'])
    multiple_log(message)
    return

def multiple_log_tpl_selected( filelist ):
    "Prints relevant informations about a selected Template"

    if len(filelist) < 1:
        message = "\t-> Selected : Empty Files list"
    else:
        firstfile = filelist[0]
        do_catgs = list(set(map(lambda x: do_catg_from_file(x), filelist)))
        message = "\t-> Selected : {0} {1}/{2} ".format(firstfile['tpl_start'], len(filelist),
                firstfile['tpl_nexp'])+','.join(do_catgs) 
    multiple_log(message)
    return

#######################################################################################
###     Various Utilities
#######################################################################################
def tpl_id_to_recipe_name( tpl_id ) :
    "Get the recipe name from the tpl.id"
    if tpl_id == "KMOS_spec_cal_dark" :
        return "kmos_dark"
    if tpl_id == "KMOS_spec_cal_calunitflat" :
        return "kmos_flat"
    if tpl_id == "KMOS_spec_cal_wave" :
        return "kmos_wave_cal"
    return "Unknown_TPL_ID"
    
def do_catg_from_file(file_entry):
    "Get the DO_CATG for a given file"

    if file_entry['tpl_id'] == "KMOS_spec_cal_dark" :
        return "DARK"
    if file_entry['tpl_id'] == "KMOS_spec_cal_calunitflat" :
        if file_entry['dpr_type'] == "FLAT,OFF" : 
            return "FLAT_OFF"
        if file_entry['dpr_type'] == "FLAT,LAMP" : 
            return "FLAT_ON"
    if file_entry['tpl_id'] == "KMOS_spec_cal_wave" :
        if file_entry['dpr_type'] == "WAVE,OFF" : 
            return "ARC_OFF"
        if file_entry['dpr_type'] == "WAVE,LAMP" : 
            return "ARC_ON"
    raise NameError("Unknown DO_CATG")
    return

#######################################################################################
###     SOF Creation
#######################################################################################
def create_sof( filelist, mydir ):
    "Creates a SOF file in mydir directory"
    
    ###############
    # Verifications
    if len(filelist) < 1 :
        raise NameError("Empty List")
        return
    # Same band ?
    if len(set(map(lambda x: x['band'], filelist))) != 1 :
        raise NameError("Different bands in this sof")
        return
    # Same tpl_id ?
    if len(set(map(lambda x: x['tpl_id'], filelist))) != 1 :
        raise NameError("Different TPL_IDs in this sof")
        return
    # Same tpl_start ?
    if len(set(map(lambda x: x['tpl_start'], filelist))) != 1 :
        raise NameError("Different TPL_STARTs in this sof")
        return
    ###############

    # SOF file name
    sof_basename = tpl_id_to_recipe_name(filelist[0]['tpl_id'])+"_"+filelist[0]['tpl_start']+".sof"
    sof_name = mydir+"/"+sof_basename

    # Open the file
    sof = open(sof_name, "wb")

    # Write RAW Files
    for file_entry in filelist:
        sof.write(file_entry['name']+" "+do_catg_from_file(file_entry)+"\n")

    band_lc = filelist[0]['band'].lower()
    band_3uc = filelist[0]['band']*3

    # Add Calibrations
    if filelist[0]['tpl_id'] == "KMOS_spec_cal_calunitflat" :
        sof.write(os.getcwd()+"/BADPIXEL_DARK.fits BADPIXEL_DARK\n")

    if filelist[0]['tpl_id'] == "KMOS_spec_cal_wave" :
        sof.write(os.environ['KMOS_CALIB']+"/kmos_wave_ref_table.fits  REF_LINES\n")
        sof.write(os.environ['KMOS_CALIB']+"/kmos_wave_band.fits       WAVE_BAND\n")
        sof.write(os.environ['KMOS_CALIB']+"/kmos_ar_ne_list_"+band_lc+".fits    ARC_LIST\n")
        sof.write(os.getcwd()+"/"+filelist[0]['band']+"/FLAT_EDGE_"+band_3uc+".fits   FLAT_EDGE\n")
        sof.write(os.getcwd()+"/"+filelist[0]['band']+"/XCAL_"+band_3uc+".fits        XCAL\n")
        sof.write(os.getcwd()+"/"+filelist[0]['band']+"/YCAL_"+band_3uc+".fits        YCAL\n")

    # Close sof file
    sof.close()
    return sof_basename
   
#######################################################################################
###     Commands Execution functions
#######################################################################################
def exec_commands_seq(cmds):
    "Exec commands sequentially"
    for cmd in cmds:
        multiple_log("\tRun in {0} : {1} ...".format(cmd['dir'], list2cmdline(cmd['cmd'])))
        cwd = os.getcwd()
        os.chdir(cmd['dir'])
        subprocess.check_call(cmd['cmd'], stdout=open(os.devnull, 'wb'), stderr=open(os.devnull, 'wb'))
        os.chdir(cwd)

def exec_commands_par(cmds):
    "Exec commands in parallel in multiple process"
    def done(p):
        return p.poll() is not None
    def success(p):
        return p.returncode == 0
    def fail():
        sys.exit(1)

    max_task = os.sysconf('SC_NPROCESSORS_ONLN')
    processes = []
    while True:
        while cmds and len(processes) < max_task:
            cmd = cmds.pop()
            multiple_log("\tRun in {0} : {1} ...".format(cmd['dir'], list2cmdline(cmd['cmd'])))
            processes.append(Popen(cmd['cmd'],cwd=cmd['dir'],stdout=open(os.devnull, 'wb'),stderr=open(os.devnull, 'wb')))

        for p in processes:
            if done(p):
                if success(p):
                    processes.remove(p)
                else:
                    fail()

        if not processes and not cmds:
            break
        else:
            time.sleep(0.05)

def exec_commands(cmds, par):
    "Exec commands"
    # Empty list
    if not cmds: 
        return 
    if par:
        exec_commands_par(cmds)
    else:
        exec_commands_seq(cmds)

def multiple_log(message):
    print message
    logging.info(message)

#######################################################################################
###     Main Programm
#######################################################################################
# Setup logging
logging.basicConfig(filename='kmos_calib.log', level=logging.DEBUG, datefmt='%m/%d/%Y %I:%M:%S %p', 
        format='%(asctime)s - %(levelname)s - %(message)s')

# Define the command line Options and Parameters

# Whenever argparse is available on the pipeline ws, use argparse instead of deprecated optparse
# import argparse
# parser = argparse.ArgumentParser(description='KMOS Calibration')
# parser.add_argument('dir', metavar='dir', type=str, help='The reference directory')
# parser.add_argument('--parallel', action='store_true', default=False, help='Parallel execution of esorex')
# args = parser.parse_args()
usage = "usage: %prog [options] data_dir"
parser = OptionParser(usage=usage, description="KMOS Calibration Data Generation Script")
parser.add_option("-p", "--parallel", action="store_true", dest="parallel", default=False, help="Parallel execution of esorex")
parser.add_option("-d", "--description", action="store_true", dest="description", default=False, help="Detailed Description")
parser.add_option("-b", "--band", default="All", help="Band that needs to be reduced (H, K, HK, YJ, IZ) [default: %default]")
(options, args) = parser.parse_args()

# If the description is wished, print it and return
description = '''
\033[1mKMOS Calibration Data Generation Script\033[0m

\033[1mPurpose:\033[0m
Execute kmos_dark/_flat/_wave_cal on the data from the 3 relevant templates in the 5 bands in order to create the calibration files set needed for the image reconstruction on the instrument workstation.
The wished 37 files are:
- BADPIXEL_DARK.fits
- MASTER_DARK.fits
- BADPIXEL_FLAT_[HHH|KKK|HKHKHK|IZIZIZ|YJYJYJ].fits
- DET_IMG_WAVE_[HHH|KKK|HKHKHK|IZIZIZ|YJYJYJ].fits
- FLAT_EDGE_[HHH|KKK|HKHKHK|IZIZIZ|YJYJYJ].fits
- MASTER_FLAT_[HHH|KKK|HKHKHK|IZIZIZ|YJYJYJ].fits
- LCAL_[HHH|KKK|HKHKHK|IZIZIZ|YJYJYJ].fits
- XCAL_[HHH|KKK|HKHKHK|IZIZIZ|YJYJYJ].fits
- YCAL_[HHH|KKK|HKHKHK|IZIZIZ|YJYJYJ].fits

\033[1mUsage:\033[0m
- Create a new directory and cd to it (mkdir tmp ; cd tmp)
- Run the command, using the full path of the raw files directory:
  kmos_calib.py -p /.../raw_files/ (full path)
- Among the KMOS*.fits files, if there are several ObsBlocks, the user will be given the choice to select one
- Among the selected ObsBlock, if there are several templates executed, the user will be given the choice to select one
- Wait till the execution is finished (~9' in parallel, ~35' in sequential on wu1pl)
- The calibration files are in ./results/*.fits

\033[1mPrerequisites:\033[0m
- The script must be run in an empty directory
- KMOS_CALIB = /cal/kmos/cal
- ESOREX_PLUGIN_DIR = /usr/lib64/esopipes-plugins
- esorex needs to be in the PATH
- The Parallel option only works on Linux (use of specific command to get the number of processors)
- The argument is a reference directory containing KMOS FITS file (full path).

\033[1mAlgorithm\033[0m
1. Create the set of files filelist, composed by all files:
  - that are in the specified directory
  - that match KMOS*.fits 
  - are from one of the 3 relevant templates (use TPL.ID)
2. Extract from filelist (into filelist_ob) the files that are from the 
  ObsBlock chosen by the user (use OBS.START)
3. Extract from filelist_ob the files from the dark template.
   If several template executions (TPL.START), ask the user to choose one.
4. For each different band (INS.GRAT1.ID) identified in ObsSet:
  a. Extract from filelist_ob the files of this band and from the flat template
     If several template executions (TPL.START), ask the user to choose one.
  b. Extract from filelist_ob the files of this band and from the wave_cal template.
     If several template executions (TPL.START), ask the user to choose one.
5. Produce the sof files and the execution scripts
6. Run the scripts in Parallel or not
7. Verifications on the products
''' 
if options.description:
    print description
    sys.exit(0)

logging.info("Input Directory: {0}".format(args[0]))
logging.info("Parallel Mode: {0}".format(options.parallel))
logging.info("Description required: {0}".format(options.description))
logging.info("Wished Band: {0}".format(options.band))

# Environment Variables
my_var_name = "KMOS_CALIB"
my_var_value = "/cal/kmos/cal"
if not my_var_name in os.environ :
    multiple_log(my_var_name+" is not set - set it to "+my_var_value)
    os.environ[my_var_name] = my_var_value
else:
    multiple_log("Use "+my_var_name+"="+os.environ[my_var_name])

my_var_name = "ESOREX_PLUGIN_DIR"
my_var_value = "/usr/lib64/esopipes-plugins"
if not my_var_name in os.environ :
    multiple_log(my_var_name+" is not set - set it to "+my_var_value)
    os.environ[my_var_name] = my_var_value
else:
    multiple_log("Use "+my_var_name+"="+os.environ[my_var_name])

# Loop on all files in the input directory
filelist_dir = []
for fname in glob.glob(args[0]+"/KMOS*.fits"):
    # Read the Primary header
    hdu = getheader(fname, 0)
    tpl_id = hdu['HIERARCH ESO TPL ID']
    # Only keep the proper TPL.ID
    if tpl_id in ["KMOS_spec_cal_dark","KMOS_spec_cal_calunitflat","KMOS_spec_cal_wave"] :
        filelist_dir.append({   'name': fname, 
                                'tpl_id': tpl_id, 
                                'tpl_start': hdu['HIERARCH ESO TPL START'],
                                'tpl_nexp': hdu['HIERARCH ESO TPL NEXP'],
                                'tpl_expno': hdu['HIERARCH ESO TPL EXPNO'],
                                'dpr_type': hdu['HIERARCH ESO DPR TYPE'],
                                'obs_start': hdu['HIERARCH ESO OBS START'],
                                'band': hdu['HIERARCH ESO INS GRAT1 ID']})

# Check that files have been selected
if len(filelist_dir) < 1:
    raise NameError("Could not find any relevant file")

# Select the Wished ObsBlock
obs_start = keyword_based_file_selection(filelist_dir, 'obs_start')

# Extract the Chosen ObsBlock files
filelist_ob = [x for x in filelist_dir if x['obs_start'] == obs_start]
multiple_log_ob_selected(filelist_ob) 

### Dark Template Selection ###
multiple_log("Dark files Selection")

# Extract the dark files   
dark_files = [x for x in filelist_ob if x['tpl_id'] == "KMOS_spec_cal_dark"]
# Check that some darks are found
if len(dark_files) == 0 :
    raise NameError("No Dark Files Found - Choose another OB")
# Which TPL START is whished ?
tpl_start = keyword_based_file_selection(dark_files, 'tpl_start')
# Extract the Chosen Template files 
dark_files_tpl = [x for x in dark_files if x['tpl_start'] == tpl_start]
# Sort
dark_files_tpl = sorted(dark_files_tpl, key=lambda x: x['name'])
multiple_log_tpl_selected(dark_files_tpl) 

# Get List of bands
bands = list(set(map(lambda x: x['band'], filelist_ob)))
multiple_log("Identified Bands: "+','.join(bands))

# Initialize selected_sets
selected_sets = []

# Loop on the different bands encountered 
for band_curr in bands:
    # Skip the bands that are not wished
    if options.band != 'All' and options.band != band_curr:
        continue
    
    ### Flatfield Template Selection ###
    multiple_log("Flat files Selection in Band: {0}".format(band_curr))

    # Extract the flat files in this band  
    flat_files_band = [x for x in filelist_ob if x['band'] == band_curr and x['tpl_id'] == "KMOS_spec_cal_calunitflat"]
    # Check that some flats are found
    if len(flat_files_band) == 0 :
        multiple_log("No Flat Field in this band - Skip it")
        continue
    # Which TPL START is whished ?
    tpl_start = keyword_based_file_selection(flat_files_band, 'tpl_start')
    # Extract the Chosen Template files 
    flat_files_tpl = [x for x in flat_files_band if x['tpl_start'] == tpl_start]
    flat_files_tpl = sorted(flat_files_tpl, key=lambda x: x['name'])
    # Log selected file lists
    multiple_log_tpl_selected(flat_files_tpl) 

    ### Wavecal Template Selection ###
    multiple_log("Wavecal files Selection in Band: {0}".format(band_curr))
    # Extract the wavecal files in this band  
    wavecal_files_band = [x for x in filelist_ob if x['band'] == band_curr and x['tpl_id'] == "KMOS_spec_cal_wave"]
    # Check that some wavecal are found
    if len(wavecal_files_band) == 0 :
        continue
    # Which TPL START is whished ?
    tpl_start = keyword_based_file_selection(wavecal_files_band, 'tpl_start')
    # Extract the Chosen Template files 
    wavecal_files_tpl = [x for x in wavecal_files_band if x['tpl_start'] == tpl_start]
    wavecal_files_tpl = sorted(wavecal_files_tpl, key=lambda x: x['name'])
    # Log selected file lists
    multiple_log_tpl_selected(wavecal_files_tpl) 

    # Store the flat_files_tpl and wavecal_files_tpl
    selected_sets.append({  'band': band_curr, 
                            'flat_files': flat_files_tpl, 
                            'wavecal_files': wavecal_files_tpl})
   
# Create Sub-Directories
for set_curr in selected_sets:
    # Skip the bands that are not wished
    if options.band != 'All' and options.band != set_curr['band']:
        continue
    
    # Create the subdirectory
    my_dir = set_curr['band']
    os.mkdir(my_dir) 

# Reduce Dark
my_dir = '.'
sof_name = create_sof(dark_files_tpl, my_dir)
recipe_name = tpl_id_to_recipe_name(dark_files_tpl[0]['tpl_id'])
log_file = "esorex_"+recipe_name+".log"
dark_cmd = ['esorex', '--suppress-prefix=TRUE', '--log-file='+log_file, '--log-dir=.', recipe_name, sof_name]
cmds = [ {'cmd': dark_cmd, 'dir': my_dir} ]
print("Dark Data Reduction")
exec_commands(cmds, options.parallel)

# Reduce Flats
cmds = []
for set_curr in selected_sets:
    # Skip the bands that are not wished
    if options.band != 'All' and options.band != set_curr['band']:
        continue

    # Generate exection command
    my_dir = set_curr['band']
    sof_name = create_sof(set_curr['flat_files'], my_dir)
    recipe_name = tpl_id_to_recipe_name(set_curr['flat_files'][0]['tpl_id'])
    log_file = "esorex_"+recipe_name+".log"
    flat_cmd = ['esorex', '--suppress-prefix=TRUE', '--log-file='+log_file, '--log-dir=.', recipe_name, sof_name]
    cmds.append({'cmd': flat_cmd, 'dir': my_dir})
print("Flatfield Data Reduction")
exec_commands(cmds, options.parallel)

# Reduce Wavecal
cmds = []
for set_curr in selected_sets:
    # Skip the bands that are not wished
    if options.band != 'All' and options.band != set_curr['band']:
        continue

    # Generate exection command
    my_dir = set_curr['band']
    sof_name = create_sof(set_curr['wavecal_files'], my_dir)
    recipe_name = tpl_id_to_recipe_name(set_curr['wavecal_files'][0]['tpl_id'])
    log_file = "esorex_"+recipe_name+".log"
    wavecal_cmd = ['esorex', '--suppress-prefix=TRUE', '--log-file='+log_file, '--log-dir=.', recipe_name, sof_name]
    cmds.append({'cmd': wavecal_cmd, 'dir': my_dir})
print("Wavelength Calibration Data Reduction")
exec_commands(cmds, options.parallel)

# Create results directory
if os.path.exists('results'):
    answer = input('Directory results already exists - Proceed (1) or Abort (0) ?')
    if answer != 1:
        sys.exit(0)
else:
    os.mkdir('results') 

# Collect results
os.rename(os.getcwd()+"/MASTER_DARK.fits", os.getcwd()+"/results/MASTER_DARK.fits")
os.rename(os.getcwd()+"/BADPIXEL_DARK.fits", os.getcwd()+"/results/BADPIXEL_DARK.fits")
for set_curr in selected_sets:
    # Skip the bands that are not wished
    if options.band != 'All' and options.band != set_curr['band']:
        continue

    # Collect the band results
    my_dir = set_curr['band']
    band_3uc = set_curr['band']*3
    filenames = [   "XCAL_{0}.fits".format(band_3uc), 
                    "YCAL_{0}.fits".format(band_3uc), 
                    "LCAL_{0}.fits".format(band_3uc),
                    "MASTER_FLAT_{0}.fits".format(band_3uc),
                    "BADPIXEL_FLAT_{0}.fits".format(band_3uc),
                    "DET_IMG_WAVE_{0}.fits".format(band_3uc),
                    "FLAT_EDGE_{0}.fits".format(band_3uc)]
    for filename in filenames :
        src_file = os.getcwd()+"/"+my_dir+"/"+filename
        dest_file = os.getcwd()+"/results/"+filename
        if os.path.isfile(src_file) :
            os.rename(src_file, dest_file)
