import reflex 
import sys 
import os
import itertools
import json
import warnings
import copy_reg
from types import *
import numpy as np
import shutil
import time
import argparse
import logging

try:
  from astropy.io import fits as pyfits
  from scipy.ndimage.filters import minimum_filter
except ImportError:
    import_ok = 0


logging.basicConfig(format='[%(levelname)s] %(message)s', level=logging.INFO,
                    stream=sys.stdout)
logger = logging.getLogger(__name__)



def crea_mask(image,min,max):

  try:
        where_are_NaNs = np.isnan(image)

        vector = image[~where_are_NaNs]
        num = len(vector)
        vector.sort()
        sel = vector[num*min : num*max]
        min_sel = sel.min()
        max_sel = sel.max()

        mask =  np.where((image >=min_sel) & (image<= max_sel),0,1)
  except:
        logger.warning('could not apply limits')
        mask= image * 0 + 1

  mask[where_are_NaNs] = 1

  return mask

def add_neibors(image,n_grow):
  mask= minimum_filter(image, size=n_grow)
  return mask

def add_skyboxes(image,skyboxes):
  ydim= len(image)
  xdim= len(image[0])
  for skybox in skyboxes:
          try:
                xmin, xmax, ymin, ymax = map(int,skybox.replace(",", " ").split())
                if xmin >= 0 and xmax <= xdim and ymin >= 0 and ymax <= ydim:
                        image[ymin:ymax,xmin:xmax]= 0
          except:
                logger.warning("can't use skybox " + skybox)

if __name__ == '__main__':
  warnings.filterwarnings("ignore")
 #getting the inputs and defining the outputs
  parser = reflex.ReflexIOParser()
  parser.add_option("-i", "--in_sof", dest="in_sof")
  parser.add_output("-o", "--out_sof", dest="out_sof")
  parser.add_input("-p", "--in_sop", dest="in_sop", type='string')
  parser.add_output("-m", "--messages", dest="messages")
  
  inputs  = parser.get_inputs()
  outputs = parser.get_outputs()
  
  in_sof = inputs.in_sof
  pattern = '--products-dir'
  for arg in sys.argv:
     if arg.split("=")[0] == pattern:
         output_dir = arg.split("=")[1]

  

  files=in_sof.files
  params  = inputs.in_sop
  values= {}
  for param in params:
          values[param.name]= param.value

  min_frac = values['muse_crea_mask.min_frac']
  max_frac = values['muse_crea_mask.max_frac']
  n_grow   = values['muse_crea_mask.n_grow']
  use_input_mask = values['muse_crea_mask.use_input_mask']
  skyboxes= []

  if min_frac < 0:
    min_frac = 0.0
    outputs.messages= 'Warning min_frac < 0. Setting it to 0.'
  if max_frac > 1:
    max_frac = 1.0
    outputs.messages= 'Warning max_frac < 1. Setting it to 1.'

  if min_frac < 0 and max_frac > 1:
    outputs.messages= ('Warning min_frac < 0 . Setting it to 0.\n'
                       'Warning max_frac > 1 . Setting it to 1.')

  if min_frac < 0 and max_frac < 0:
    outputs.messages= 'Warning min_frac and max_frac < 0. Setting it to 0'
  
  if min_frac > 1 and max_frac >1:
    outputs.messages= 'Warning min_frac and max_frac > 1. Setting it to 0'
  
  if min_frac > max_frac:
    min_frac = 0.0
    max_frac = 0.0
    message1 = 'Warning, min_frac > max_frac. Setting them to 0.'

  for key in values.keys():
          if key[0:21] == 'muse_crea_mask.skybox':
                  skyboxes.append(values[key])

  
                  
  bad_parameters = 0
  # read input images
  input_mask_provided = False
  for file in files:
      if file.category == 'IMAGE_FOV':
        hdu_fov  = pyfits.open(file.name)
        image_fov = hdu_fov[1].data 
      if file.category == 'MASK_CUBE':
        input_mask_name = file.name
        hdu  = pyfits.open(input_mask_name)
        input_mask = hdu[0].data
        input_mask_provided = True


  # create mask with different options:
  if input_mask_provided and  use_input_mask[0] == 'u': 
        output_name= input_mask_name
  else: 
        if input_mask_provided and  use_input_mask[0] == 'a':
                first_mask= input_mask
        else:
                first_mask = crea_mask(image_fov,min_frac,max_frac)
        mask = add_neibors(first_mask, n_grow)
        add_skyboxes(mask, skyboxes)

        #check for bad parameters and create a fake mask + error message
        if np.min(mask) != 0 or np.max(mask) !=1:
           mask=mask*0.0+1.
           mask[0]= 0
           outputs.messages=('Warning, invalid mask created. Check the recipe parameters')

        # create output mask:
        output_name = output_dir + '/MASK_FINAL_CUBE.fits'
        hdu2 = pyfits.PrimaryHDU(mask)
        hdulist = pyfits.HDUList([hdu2])
        hdulist[0].header = hdu_fov[0].header
        hdulist[0].header['HIERARCH ESO PRO CATG'] = 'MASK_FINAL_CUBE'
           
           
        i=0
        for param in params:
          i=i+1
          hkey1 = 'HIERARCH ESO MASK REC'+str(i)+' NAME'
          hkey2 = 'HIERARCH ESO MASK REC'+str(i)+' VALUE'
          hdulist[0].header[hkey1] = param.name
          hdulist[0].header[hkey2] = param.value
        hdulist.writeto(output_name,clobber=True)
         
  

  # create output SOF:
  files = list()
  output_purpose=file.purposes 
  output_datasetname=in_sof.datasetName
  output_category = 'MASK_FINAL_CUBE' #Category will be changed (even if I do not edit the user-provided mask)
  files.append(reflex.FitsFile(output_name,output_category,None,output_purpose))
  newsof = reflex.SetOfFiles(output_datasetname,files)
  outputs.out_sof = newsof

# emitt output:
  parser.write_outputs()
  sys.exit()
 


