#from builtins import range
try:
   from astropy.io import fits as pyfits
except:
   import pyfits
import numpy as np
import glob
import os
import copy
import re
import string
import sys

from .log import Log, ERROR, WARNING, NOTICE
log = Log().log

#
# Load files and define usefull instance (ex: tag)
#

###
def get_tag (h):
    '''Get the SOF tag for this header'''
    # Has a PRO.CATG
    if 'HIERARCH ESO PRO CATG' in h:
        procatg = h['HIERARCH ESO PRO CATG']
        if procatg=='JSDC_CAT':
            return 'DIAMETER_CAT'
        elif re.match('TF_*.',procatg):
            return procatg[3:]+"_TF"
        elif re.match('VIS_*.RAW',procatg):
            return procatg[4:-4]+"_VIS"
        else:
            return procatg
    # Has no PRO.CATG, return TYPE_RAW
    elif 'HIERARCH ESO DPR TYPE' in h:
        tag = h['HIERARCH ESO DPR TYPE']
        if 'SKY' in tag:
            tag = tag[4:]+'_SKY'
        elif 'OBJECT' in tag:
            tag = tag[7:]+"_"+h['HIERARCH ESO DPR CATG'][:3]
        elif 'STD' in tag:
            tag = tag[4:]+"_"+h['HIERARCH ESO DPR CATG'][:3]
        elif 'WAVE,SC' in tag:
            tag = 'WAVESC'
        elif 'WAVE,LAMP' in tag:
            tag = 'WAVELAMP'
        return tag+'_RAW'
    else:
        return 'UNKNOWN'

###
def get_shutters (h):
    '''Get the shutter string for this header'''
    try:
       out = ''
       for i in range(1,5):
          out += 'T' if h['HIERARCH ESO INS SHUT1%i ST'%i] else 'F'
       return out
    except:
       return 'TTTT'

###
def add_gravity (headers,files):
    '''Add gravity headers to the list'''
    
    files.sort()
    for f in files:
        if len(headers)>0 and f in [h.name for h in headers]:
            continue

        # Get entire FITS header
        try:
           hdulist = pyfits.open(os.path.realpath(f))
        except IOError as e:
           raise IOError("IOError in file " +os.path.realpath(f)+": "+str(e))
        h = hdulist[0].header.copy()
        hdulist.close ()
        
        # Set additional information in header
        h.name = f
        h.tag = get_tag(h)
        h.shutters = get_shutters(h)
        h.is_raw = False if 'HIERARCH ESO PRO CATG' in h else True
        h.mode = 'DUAL' if 'DUAL' in h.tag else 'SINGLE'
        
        # Keep or reject
        if 'TEST' in h.tag or 'UNKNOWN' in h.tag:
           log ('Skip %s as %s'%(h.name,h.tag), 1, NOTICE)
        else:
           log ('Load %s as %s'%(h.name,h.tag), 1, NOTICE)
           headers.append(h)
        
    return headers

#
# Search for headers with a given tag and
# matching criteria
#

opti_ft_setup=['HIERARCH ESO FT POLA MODE']
opti_sc_setup = ['HIERARCH ESO INS SPEC RES', 'HIERARCH ESO INS POLA MODE']
opti_setup = opti_sc_setup + opti_ft_setup

det_setup_ftgain = ['HIERARCH ESO INS DET3 GAIN']
det_setup_ftdit = ['HIERARCH ESO DET3 SEQ1 DIT']
det_setup_scdit = ['HIERARCH ESO DET2 SEQ1 DIT']
det_setup_win = ['HIERARCH ESO INS FDDL WINDOW']
det_setup=det_setup_ftdit+det_setup_ftgain+det_setup_scdit+det_setup_win

calib_setup = opti_setup + ['ESO PRO NIGHT OBS']
tpl_setup = ['HIERARCH ESO TPL START']

met_setup = ['HIERARCH ESO INS MET MODE']

###
def assoc (h, allh, tag, keys, prefered_key=[], which='closest', required=0):
    '''Search for headers with tag and matching criteria'''

    # Keep only the requested tag matching the criteria
    atag = [a for a in allh if a.tag==tag]
    out = []
    for a in atag:
        tmp = True
        for k in keys:
            tmp *= (h.get(k,0) == a.get(k,0))
        if tmp:
            out.append(a)

    # Check required
    if len (out) < required:
        log ('Cannot find %i %s for %s'%(required,tag,h.name), 1, WARNING)
    
    # if more than required, check for frames with addition of the prefered_key
    if len (out) > required:
        out2 = []
        for a in atag:
            tmp = True
            for k in keys+prefered_key:
                tmp *= (h.get(k,0) == a.get(k,0))
            if tmp:
                out2.append(a)

    # Check if out2 is equal or above required, if so, use outE
        if len (out2) >= max([required,1]):
            out=out2
    
    # Check closest
    if len (out) > required and which=='closest':
        # Case need closest and more than 1 not supported yet
        if required < 2:
            time_diffs = np.array([o['MJD-OBS'] - h['MJD-OBS'] for o in out])
            out = [out[np.abs(time_diffs).argmin()]]

    # Check closest
    if len (out) > required and which=='before':
        # Case need closest and more than 1 not supported yet
        if required < 2:
            time_diffs = np.array([o['MJD-OBS'] - h['MJD-OBS'] for o in out])
            time_diffs_before=np.argwhere(time_diffs < 0 );
            if len (time_diffs_before) > 0:
               out = [out[time_diffs_before[(np.abs(time_diffs[time_diffs_before])).argmin()][0]]];
            else:
               log ('Cannot find %s before, thus use after'%tag, 1, WARNING);
               out = [out[(np.abs(time_diffs)).argmin()]];
        else:
            log ('Case with "before" and more than 1 is not supported yet (report)', 1, WARNING);
        
    return out


## other functions
def new_file_in_dir(directory, prevhash_point):
    """ check if there is a new file in directory
        prevhash_point must be a list with one ellement (act like a pointer) and
        is modified by the function
    """
    prevhash  = prevhash_point[0]
    newhash = tuple( name for name in os.listdir(directory) if os.path.isfile(name) )
    prevhash_point[0] = newhash
    return prevhash != newhash
    
#
# Execute the SOF
#

###
def run_sof (sof, name=None, recipe='unknown', options=None, overwrite=False, outputdir='reduced'):

    # Define the name of SOF and of executable files
    master = os.path.splitext(os.path.basename(name))[0]
    sof_filename = outputdir+'/'+master+'_esorex.sof'
    cmd_filename = outputdir+'/'+master+'_esorex.sh'

    # Check outputdir exist
    if not os.path.exists(outputdir):
        os.mkdir(outputdir)

    # Check if product exist (defined as outputdir/master_*fits)
    if (not overwrite) and len(glob.glob("%s/%s_*.fits"%(outputdir,master))):
        log ("Skiping %s : product already exist"%master, 1, NOTICE)
        return

    # Write the sof
    log ("Write %s "%sof_filename, 1, NOTICE)
    sof_file = open(sof_filename, "w")
    for s in sof:
        sof_file.write('%s    %s\n'%(s.name,s.tag))
    sof_file.close()

    # Esorex options, returned as dictionary {'opt-name':'value'}
    e_opt = get_recipe_options (options, 'esorex')
    e_opt['suppress-prefix']='TRUE'
    e_opt['log-dir']=outputdir
    e_opt['log-file']=master+'_esorex.log'
    e_opt['output-dir']=outputdir
    e_opt['output-prefix']=master
    e_opt['check-sof-exist']='TRUE'

    # Build the esorex option string
    e_options  = ' '.join(['--'+k+'='+e_opt[k] for k in list(e_opt.keys()) if e_opt[k] is not None])

    # Recipe options, returned as dictionary {'opt-name':'value'}
    r_opt = get_recipe_options (options, recipe)

    # Build the recipe option string
    r_options = ' '.join(['--'+k+'='+r_opt[k] for k in list(r_opt.keys()) if r_opt[k] is not None])

    # Build the esorex command
    cmd = 'esorex %s %s %s %s'%(e_options,recipe,r_options,sof_filename)
    
    # Write the esorex command and make it executable
    log ("Write %s "%cmd_filename, 1, NOTICE)
    cmd_file = open(cmd_filename, "w")
    cmd_file.write(cmd)
    cmd_file.close()
    os.chmod(cmd_filename, 0o775)
    
    # Execute esorex command
    # log ("Execute %s "%cmd_filename, 1, NOTICE)
    sys.stdout.flush()
    erc = os.system(cmd_filename)
    os.utime(sof_filename, None)
    os.utime(cmd_filename, None)
    if erc:
        log ("shell esorex command returned error code %d"%erc, 1, WARNING)
    return erc

###
def get_recipe_options(argoptions, recipe):
    options = {None:None}
    if argoptions is None:
        return options
    
    dic = vars(argoptions)
    for o in dic:
        if o.split('.')[0] == recipe and dic[o] is not None:
            options[o.split('.')[1]] = dic[o]
    return options

###
def implement_recipe_options(parser,options):
    for o in options:
        recipe = o[0].split('.')[0].replace('esorex','')
        parser.add_argument('--'+o[0], dest=o[0], default=o[1],
                            help='(see esorex -h '+recipe+')',
                            metavar='')

###
def verbose_sof(i,sofs):
    log ("", 1, NOTICE)
    log ("", 1, NOTICE)
    log ("***** Now reducing %s %i over %i *****"%(sofs[i].tag, i+1,len(sofs)), 1, NOTICE)
    log ("", 1, NOTICE)

    
#
# High-level routine to reduce the DARK_RAW, P2VM_RAW,
# DISP_RAW, OBJECT_RAW...
#
# All build following the same sctructure:
# - search for the list of individual reduction, then for each
# - associate RAW and CALIB files to have a SOF
# - run esorex
# 

###
def reduce_piezotf_raw (raw, outputdir, options=None):
    piezotfs = [ h  for h in raw if h.tag=="PIEZOTF_RAW" and h['HIERARCH ESO TPL ID']=='GRAVITY_gen_cal_piezoTF' ]
    for i,piezotf in enumerate(piezotfs):
        verbose_sof (i,piezotfs)
        sof = [piezotf]
        run_sof (sof, name=piezotf.name, recipe='gravity_piezo', options=options, outputdir=outputdir)


###
def reduce_dark_raw (raw, outputdir, options=None):
    darks = [ h  for h in raw if h.tag=="DARK_RAW" and h['HIERARCH ESO TPL ID']=='GRAVITY_gen_cal_dark' ]
    for i,dark in enumerate(darks):
        verbose_sof (i,darks)
        sof = [dark]
        run_sof (sof, name=dark.name, recipe='gravity_dark', options=options, outputdir=outputdir)

###
# reduce WAVE (but why recipe p2vm??)
###
def reduce_wave_raw (raw, calibs, outputdir, options=None):
    waves = [ h  for h in raw if h.tag=="WAVE_RAW" and h['HIERARCH ESO TPL ID']=='GRAVITY_gen_cal_wave']
    for i,wave in enumerate(waves):
        verbose_sof (i,waves)
        sof = [wave]
        # FIXME: shall we check for the same DIT or same TPL.START in the DARK ??
        sof += assoc (wave, calibs, 'DARK', opti_setup, required=1)
        sof += assoc (wave, calibs, "BAD", opti_setup, required=1)
        sof += assoc (wave, calibs, "FLAT", opti_setup, required=1)
        sof += assoc (wave, raw, 'WAVESC_RAW', opti_setup+tpl_setup, required=1)
        sof += assoc (wave, raw+calibs, "WAVE_PARAM", [])
        run_sof (sof, name=wave.name, recipe='gravity_p2vm', options=options, outputdir=outputdir)
                        
        
###
# Reduce P2VM
###
def reduce_p2vm_raw (raw, calibs, outputdir, options=None):
    # Hack for the SV new p2vms
    p2vms  = [h  for h in raw if h.tag=="P2VM_RAW" and h['HIERARCH ESO TPL ID']=='GRAVITY_gen_cal_p2vm' and h.shutters=='TTFF' and h["MJD-OBS"]<=57550]
    p2vms += [h  for h in raw if h.tag=="P2VM_RAW" and h['HIERARCH ESO TPL ID']=='GRAVITY_gen_cal_p2vmWAVESC' and h.shutters=='TTFF' and h["MJD-OBS"]>57550]
    p2vms += [h  for h in raw if h.tag=="P2VM_RAW" and h['HIERARCH ESO TPL ID']=='GRAVITY_gen_cal_p2vm' and h.shutters=='TTFF' and h["MJD-OBS"]>57610]
    for i,p2vm in enumerate(p2vms):
        verbose_sof (i,p2vms)
        sof = []
        sof += assoc (p2vm, raw, 'DARK_RAW',   tpl_setup+opti_setup, required=1)
        sof += assoc (p2vm, raw, 'FLAT_RAW',   tpl_setup+opti_setup, required=4)
        sof += assoc (p2vm, raw, 'P2VM_RAW',   tpl_setup+opti_setup, required=6)
        sof += assoc (p2vm, raw, 'WAVE_RAW',   tpl_setup+opti_setup, required=1)
        sof += assoc (p2vm, raw, 'WAVESC_RAW', tpl_setup+opti_setup, required=1)
        sof += assoc (p2vm, calibs, "WAVE_PARAM", [], which='before')        
        run_sof (sof, name=p2vm.name, recipe='gravity_p2vm', options=options, outputdir=outputdir)
        
###
# Reduce Wavelamp
###
def reduce_wavelamp_raw (raw, calibs, outputdir, options=None):
    wavelamps = [ h  for h in raw if h.tag=="WAVELAMP_RAW"]
    for i,wavelamp in enumerate(wavelamps):
        verbose_sof (i,wavelamps)
        sof = [wavelamp]
        sof += assoc (wavelamp, raw, 'DARK_RAW', opti_setup+det_setup+met_setup, required=1)
        sof += assoc (wavelamp, calibs, "BAD",  opti_setup, required=1)
        sof += assoc (wavelamp, calibs, "FLAT", opti_setup, required=1)
        sof += assoc (wavelamp, calibs, "WAVE", opti_setup, required=1)
        sof += assoc (wavelamp, calibs, "P2VM", opti_setup, required=1)
        sof += assoc (wavelamp, calibs, "WAVE_PARAM", [])  
        run_sof (sof, name=wavelamp.name, recipe='gravity_wavelamp', options=options, outputdir=outputdir)
                        
###
# Reduce dispersion
###
def reduce_disp_raw (raw, calibs, outputdir, options=None):
    disps = [ h  for h in raw if h.tag=='DISP_RAW' and h['HIERARCH ESO TPL EXPNO']==2 ]
    for i,disp in enumerate(disps):
        verbose_sof (i,disps)
        sof = []
        sof += assoc (disp, raw, "DISP_RAW", opti_setup+tpl_setup, which='all')
        # Search for a DARK with full-matching, if none, accept wrong FDDL.WINDOW
        sof += assoc (disp, raw, "DARK_RAW", opti_setup+det_setup+met_setup,required=1)
        sof += assoc (disp, calibs, "WAVELAMP", opti_setup, required=1)
        sof += assoc (disp, calibs, "BAD",  opti_setup, required=1)
        sof += assoc (disp, calibs, "FLAT", opti_setup, required=1)
        sof += assoc (disp, calibs, "WAVE", opti_setup, required=1)
        sof += assoc (disp, calibs, "P2VM", opti_setup, required=1)
        sof += assoc (disp, calibs, "WAVE_PARAM", [], required=1)
        run_sof (sof, name=disp.name, recipe='gravity_disp', options=options, outputdir=outputdir)
        
###
#
# Reduce GRAVITY DATA (gravity_vis recipe)
###
def reduce_object_raw (raw, calibs, outputdir, options=None, overwrite=False):
    objects = [ h  for h in raw if '_SCI_RAW' in h.tag or '_CAL_RAW' in h.tag]
    for i,obj in enumerate(objects):
        verbose_sof (i,objects)
        sof = [obj]
        # Search for a SKY with full-matching, if none, accept wrong SC_DIT
        sof += assoc (obj, raw, re.sub ('_.+_', '_SKY_',obj.tag),
        opti_ft_setup + det_setup_ftdit + det_setup_ftgain + met_setup,
        opti_sc_setup + det_setup_scdit, required=1)
        # Search for a DARK with full-matching, if none, accept wrong FT_CONFIG
        sof += assoc (obj, calibs, "DARK", opti_setup+det_setup_scdit+met_setup, det_setup, required=1)
        sof += assoc (obj, calibs, "BAD", opti_setup, det_setup_ftgain, required=1)
        sof += assoc (obj, calibs, "FLAT", opti_setup, det_setup_ftgain, required=1)
        sof += assoc (obj, calibs, "WAVE", opti_setup, det_setup_ftgain, required=1)
        sof += assoc (obj, calibs, "P2VM", opti_setup, det_setup_ftgain, required=1)
        sof += assoc (obj, raw+calibs, "DISP_MODEL", [], required=1)
        sof += assoc (obj, raw+calibs, "DIAMETER_CAT", [], required=1)
        sof += assoc (obj, raw+calibs, "EOP_PARAM", [], required=1)
        sof += assoc (obj, raw+calibs, "DIODE_POSITION", [], required=1)
        sof += assoc (obj, raw+calibs, "STATIC_PARAM", [], required=1, which='before')
        run_sof (sof, name=obj.name, recipe='gravity_vis', options=options, outputdir=outputdir, overwrite=overwrite)

###
def average_p2vmred (raw, outputdir, options=None, overwrite=False):
    p2vmred = [h for h in raw if re.match('.*P2VMRED.*',h.tag)]
    for i,p in enumerate(p2vmred):
        verbose_sof (i,p2vmred)
        sof = [p]
        run_sof (sof, name=p.name, recipe='gravity_vis_from_p2vmred', options=options, outputdir=outputdir, overwrite=overwrite)
    
###
def compute_tf (raw, calibs, outputdir, options=None, overwrite=False):
    viss = [h for h in raw if re.match('.*_CAL_VIS',h.tag)]
    for i,vis in enumerate(viss):
        verbose_sof (i,viss)
        sof = [vis]
        sof += assoc (vis, calibs, "DIAMETER_CAT", [])
        run_sof (sof, name=vis.name, recipe='gravity_viscal', options=options, outputdir=outputdir, overwrite=overwrite)

###
def calibrate_vis (raw, calibs, outputdir, options=None, overwrite=False):
    viss = [h for h in raw if re.match('.*_SCI_VIS',h.tag)]
    for i,vis in enumerate(viss):
        verbose_sof (i,viss)
        sof = [vis]
        sof += assoc (vis, calibs, 'SINGLE_CAL_TF', calib_setup, required=0, which='all')
        sof += assoc (vis, calibs, 'DUAL_CAL_TF', calib_setup, required=0, which='all')
        run_sof (sof, name=vis.name, recipe='gravity_viscal', options=options, outputdir=outputdir, overwrite=overwrite)