#!/usr/bin/env python2
# -*- coding: utf8 -*-
"""
    SAMI XJoin

    This script simply joins the four existing extensions inside a FITS file
    created during observations with SAMI (SAM Imager). During the process,
    it also fits a 2nd degree polynomium to the OVERSCAN region that is
    subtracted from the corresponding image.

    The user also may want to add flags in order to process the images
    according to the following options (in order):

    - BIAS subtraction;
    - DARK subtraction;
    - Remove hot pixels and cosmic rays;
    - Remove overglow using a long exposure DARK image;
    - Divide by the FLAT;
    - Divide by the exposure time;

    The documentation for each process is shown in the corresponding function.

    Bruno Quint (bquint at ctio.noao.edu)
    May 2016

    Thanks to Andrei Tokovinin and Claudia M. de Oliveira for the ideas that
    were implemented here.
"""

from __future__ import division, print_function

import astropy.io.fits as pyfits
import argparse
import logging as log
import numpy as np
import os

from scipy import ndimage
from scipy import signal

try:
    # noinspection PyUnresolvedReferences,PyUnboundLocalVariable
    xrange
except NameError:
    # noinspection PyShadowingBuiltins
    xrange = range

# Piece of code from cosmics.py
# We define the laplacian kernel to be used
laplkernel = np.array([[0.0, -1.0, 0.0], [-1.0, 4.0, -1.0], [0.0, -1.0, 0.0]])

# Other kernels :
growkernel = np.ones((3, 3))

# dilation structure for some morphological operations
dilstruct = np.ones((5, 5))
dilstruct[0, 0] = 0
dilstruct[0, 4] = 0
dilstruct[4, 0] = 0
dilstruct[4, 4] = 0


# noinspection PyPep8Naming
class SAMI_XJoin:
    def __init__(self, list_of_files, bias_file=None, clean=False,
                 cosmic_rays=False, dark_file=None, debug=False,
                 flat_file=None, glow_file=None, time=False, verbose=False):

        self.set_verbose(verbose)
        self.set_debug(debug)
        self.main(list_of_files, bias_file=bias_file, clean=clean,
                  cosmic_rays=cosmic_rays, dark_file=dark_file,
                  flat_file=flat_file, glow_file=glow_file, time=time)

        return

    @staticmethod
    def clean_column(_data, x0, y0, yf, n=5):
        t1 = _data[y0:yf, x0 - n:x0]
        t2 = _data[y0:yf, x0 + 1:x0 + n]
        t = np.hstack((t1, t2))
        _data[y0:yf, x0] = np.median(t, axis=1)
        return _data

    def clean_columns(self, _data):
        bad_columns = [
            [167, 0, 513],
            [476, 0, 513],
            [602, 0, 513],
            [671, 0, 513],
            [673, 475, 513],
            [810, 0, 513],
            [213, 513, 1024]
        ]

        for column in bad_columns:
            x0 = column[0]
            y0 = column[1]
            yf = column[2]
            _data = self.clean_column(_data, x0, y0, yf)
        return _data

    @staticmethod
    def clean_line(_data, x0, xf, y, n=5):
        t1 = _data[y - n:y, x0:xf]
        t2 = _data[y + 1:y + n, x0:xf]
        t = np.vstack((t1, t2))
        _data[y, x0:xf] = np.median(t, axis=0)
        return _data

    def clean_hot_columns_and_lines(self, data, header, prefix, clean):

        if clean is True:
            data = self.clean_columns(data)
            data = self.clean_lines(data)
            header.add_history('Cleaned bad columns and lines.')
            prefix = 'c' + prefix

        return data, header, prefix

    def clean_lines(self, _data):
        bad_lines = [
            [214, 239, 688],
            [477, 516, 490],
            [387, 429, 455],
            [574, 603, 494],
            [574, 603, 493],
            [640, 672, 388],
            [604, 671, 388]
        ]

        for line in bad_lines:
            x0 = line[0]
            xf = line[1]
            y = line[2]
            _data = self.clean_line(_data, x0, xf, y)
        return _data

    @staticmethod
    def dark_subtraction(data, header, prefix, dark_file):

        if dark_file is not None:
            dark = pyfits.getdata(dark_file)
            data -= dark
            header['DARKFILE'] = dark_file
            prefix = 'd' + prefix
            header.add_history('Dark subtracted')

        return data, header, prefix

    @staticmethod
    def divide_by_flat(data, header, prefix, flat_file):

        if flat_file is not None:
            flat = pyfits.getdata(flat_file)
            data /= flat
            header['FLATFILE'] = flat_file
            header.add_history('Flat normalized')
            prefix = 'f' + prefix

        return data, header, prefix

    @staticmethod
    def divide_by_exposuretime(data, header, prefix, time):

        if time is True:
            try:
                exptime = float(header['EXPTIME'])
                data /= exptime
                header['UNITS'] = 'COUNTS/s'
                header.add_history('Divided by exposure time.')
                prefix = 't' + prefix
            except KeyError:
                pass

        return data, header, prefix

    @staticmethod
    def get_header(filename):

        fits_file = pyfits.open(filename)
        h0 = fits_file[0].header
        # noinspection PyUnusedLocal
        h1 = fits_file[1].header

        # TODO - Multiple header inheritance.
        # If there's any card that should be passed from the extentions
        # header to the main header, uncomment bellow.
        # for key in h1:
        #     if key not in h0:
        #         h0.set(key, value=h1[key], comment=h1.comments[key])

        h0.append('UNITS')
        h0.set('UNITS', value='COUNTS', comment='Pixel intensity units.')

        return h0

    @staticmethod
    def get_joined_data(filename):

        fits_file = pyfits.open(filename)
        w, h = str2pixels(fits_file[1].header['DETSIZE'])

        log.info(' > %s' % filename)

        # Correct for binning
        bin_size = np.array(fits_file[1].header['CCDSUM'].split(' '),