#!/usr/bin/python

import sys
import os
import time
from optparse import OptionParser
try:
    import pyfits
    from pyfits import getheader
except:
    from astropy.io import fits as fits
    from fits import getheader
import glob
import logging

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

def log_filelist ( filelist ):
    "Log filelist informations"

    if len(filelist) < 1:
        logging.info("Empty Files list")
    else:
        for my_file in filelist:
            logging.info("Name: {0} / n ext: {1} / size: {2} / PRO.CATG: {3} / band: {4}".format(
                    my_file['name'], my_file['n_ext'], my_file['size'], my_file['pro_catg'], my_file['band']))
            logging.info("\tExtnames : "+','.join(my_file['extnames']))
            logging.info("\tAngles : "+','.join(map(str, my_file['angles'])))
    return

#######################################################################################
###     Main Programm
#######################################################################################
# Setup logging
logging.basicConfig(filename='kmos_verify.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
usage = "usage: %prog [options] result_dir"
parser = OptionParser(usage=usage, description="KMOS Calibration Data Verification Script")
parser.add_option("-d", "--description", action="store_true", dest="description", default=False, help="Detailed Description")
(options, args) = parser.parse_args()

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

\033[1mPurpose:\033[0m
The script runs basic verifications on the calibration files generated by mos_calib.py
The results/ directory created by kmos_calib.py si supposed to contain the following files:
- 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

The scripts checks the consistency between the xcal, ycal and lcal files.

\033[1mVerifications:\033[0m
- Check that there are 15 files
- Check that all bands [H,K,HK,IZ,YJ] have exactly 1 xcal, 1ycal, 1lcal
- Check that all found files have the same size(+/-5%), nb of extensions, angles list, 
    extnames list
- Check that the angles list contains 6 different values (3 times each)
''' 
if options.description:
    print description
    sys.exit(0)

logging.info("Input Directory: {0}".format(args[0]))
logging.info("Description required: {0}".format(options.description))

# Loop on all files in the input directory
files_list = []
for fname in glob.glob(args[0]+"/*cal*.fits"):
    # Open the file
    hdulist = pyfits.open(fname)
    # Initialize the angles list
    angles_list = []
    extnames_list = []
    for ix in range(len(hdulist)):
        if (ix == 0) :
            continue
        angles_list.append(hdulist[ix].header['HIERARCH ESO PRO ROT NAANGLE'])
        extnames_list.append(hdulist[ix].header['EXTNAME'])
    # Store the file informations
    files_list.append({  
            'name': fname, 
            'n_ext': len(hdulist),
            'size': os.path.getsize(fname),
            'pro_catg': hdulist[0].header['HIERARCH ESO PRO CATG'],
            'band': hdulist[0].header['HIERARCH ESO INS GRAT1 ID'],
            'extnames' : extnames_list,
            'angles' : angles_list})
    # Close the file
    hdulist.close()

# Display the stored information
log_filelist(files_list)

# Define expected values
detected_problems = 0 
expected_bands = ['H','K','HK','IZ','YJ']
expected_procatg = ['XCAL', 'YCAL', 'LCAL']
nb_expected_files = 3*len(expected_bands)
nb_expected_angles = 6 ;

# Check that there are 15 files
if len(files_list) == nb_expected_files:
    logging.info("Number of files ... ok")
else:
    multiple_log("Number of files ... Nok - {0} files instead of {1}".format(len(files_list), nb_expected_files))
    detected_problems+=1

# Check that all bands [H,K,HK,IZ,YJ] have exactly 1 xcal, 1ycal, 1lcal
for band in expected_bands:
    for procatg in expected_procatg:
        # Extract current procatg/band files 
        cur_list = [x for x in files_list if x['pro_catg'] == procatg and x['band'] == band]

        # Chaeck that 1 file is found
        if len(cur_list) == 1:
            logging.info("Check if 1 file in band {0} with PRO.CATG {1} ... ok".format(band, procatg))
        else:
            multiple_log("Check if 1 file in band {0} with PRO.CATG {1} ... Nok - {2} files found".format(
                        band, procatg, len(cur_list)))
            detected_problems+=1

# Check that all found files have the same size(+/-5%), nb of extensions, angles list, extnames list
ref_size = files_list[0]['size']
ref_n_ext = files_list[0]['n_ext']
ref_angles = files_list[0]['angles']
ref_extnames = files_list[0]['extnames']
tol = 0.05
for ix in range(len(files_list)):
    if (ix == 0) :
        continue
    # Check size
    if files_list[ix]['size'] > ref_size*0.95 and files_list[ix]['size'] < ref_size*1.05 :
        logging.info("Check if {0} size is the same (+/-5%) as the first file ({1})... ok".format(
                    files_list[ix]['name'], ref_size))
    else:
        multiple_log("Check if {0} size is the same (+/-5%) as the first file ({1})... Nok - {2}".format(
                    files_list[ix]['name'], ref_size, files_list[ix]['size']))
        detected_problems+=1

    # Check nb of extensions
    if files_list[ix]['n_ext'] == ref_n_ext :
        logging.info("Check if {0} nb of extensions is the same as for the first file ({1})... ok".format(
                    files_list[ix]['name'], ref_n_ext))
    else:
        multiple_log("Check if {0} nb of extensions is the same as for the first file ({1})... Nok - {2}".format(
                    files_list[ix]['name'], ref_n_ext, files_list[ix]['n_ext']))
        detected_problems+=1

    # Check Angles list
    if cmp(files_list[ix]['angles'], ref_angles) == 0 :
        logging.info("Check if {0} angles are the same as for the first file ... ok".format(files_list[ix]['name']))
    else:
        multiple_log("Check if {0} angles are the same as for the first file ... Nok - Check the log file".format(
                    files_list[ix]['name']))
        logging.info("\tRef Angles : "+','.join(map(str, ref_angles)))
        logging.info("\tCur Angles : "+','.join(map(str, files_list[ix]['angles'])))
        detected_problems+=1

    # Check extnames list
    if cmp(files_list[ix]['extnames'], ref_extnames) == 0 :
        logging.info("Check if {0} ext names are the same as for the first file ... ok".format(files_list[ix]['name']))
    else:
        multiple_log("Check if {0} ext names are the same as for the first file ... Nok - Check the log file".format(
                    files_list[ix]['name']))
        logging.info("\tRef ExtNames : "+','.join(map(str, ref_extnames)))
        logging.info("\tCur ExtNames : "+','.join(map(str, files_list[ix]['extnames'])))
        detected_problems+=1

# Check that the angles list contains 6 different values (3 times each)
angles = list(set(ref_angles))
if len(angles) == nb_expected_angles:
    logging.info("Number of angles ... ok")
else:
    multiple_log("Number of angles ... Nok - {0} angles instead of {1}".format(len(angles), nb_expected_angles))
    detected_problems+=1

# Check that there are 3 extension per angle
for angle in angles:
    cur_angles = [x for x in ref_angles if x == angle]
    if len(cur_angles) == 3:
        logging.info("Number of extensions for angle {0} ... ok".format(angle))
    else:
        multiple_log("Number of extensions for angle {0} ...  Nok - {1} angles instead of 3".format(angle, len(cur_angles)))
        detected_problems+=1

# Summary
if detected_problems == 0:
    multiple_log("... No problem detected")
else:
    multiple_log("WARNING - {0} problem(s) detected".format(detected_problems))
    
