import numpy as np
from matplotlib import pyplot as plt
from astropy.io import fits
from glob import glob
import os

#Author: Trevor Dorn-Wallenstein
#Date: 9.20.17
#Author's note: the data reduction portion of this does almost exactly what acronym 
#(Weisenburger et al. 2017, https://github.com/kweis/acronym) does, with the exception of 
#performing an overscan subtraction before trimming. I wrote this to be slightly more explicit 
#(read: slower) so that a PreMAP student can run any portion of the reduction and see what it does.
#If you want to reduce ARCTIC data and you are not my PreMAP student, PLEASE use Kolby's code!

#Hell, even if you are my PreMAP student, try Kolby's code too!

#If you want to run the pipeline from the command line without any customization,
#do python adapt.py datadir caldir reddir
#where datadir is where the data are, caldir is where master cals go, and reddir is
#where the reduced data go

def split_data_overscan(hdu):
    """
    Search through the image header and return python indices
    for where the overscan region is
    
    Parameters
    ---------
    hdu : `~astropy.io.fits.hdu.image.PrimaryHDU`
        HDU object containing the raw FITS image from which overscan and data arrays 
        are generated.
        
    Returns
    -------
    data : `~numpy.ndarray`
        Trimmed data
    overscan : `~numpy.ndarray`
        Overscan region
        
    """
    
    #Give me the overscan and the header
    image = hdu.data
    header = hdu.header
    
    #Some string manipulation to get the actual values.
    overscan_str = hdu.header['BSEC11'].lstrip('[').rstrip(']').split(',
')
    data_str = hdu.header['DSEC11'].lstrip('[').rstrip(']').split(',')
    overscan_x_min,overscan_x_max = overscan_str[1].split(':')
    overscan_y_min,overscan_y_max = overscan_str[0].split(':')
    data_x_min,data_x_max = data_str[1].split(':')
    data_y_min,data_y_max = data_str[0].split(':')
    
    #Remember, FITS is 1-indexed, and sections are inclusive,
    #whereas python is
 0-indexed, and exlusive to the end of a section
    data = image[int(data_x_min)-1:int(data_x_max),int(data_y_min)-1:int(data_y_max)]
    overscan = image[int(overscan_x_min)-1:int(overscan_x_max),int(overscan_y_min)-1:int(overscan_y_max)]
    
    return data, overscan

def trim_subtract_overscan(hdu,fit_degree = 8):
    """
    Search through the image header and return python indices
    for where the overscan region is
    
    Parameters
    ---------
    hdu : `~astropy.io.fits.hdu.image.PrimaryHDU`
        HDU object containing the raw FITS image from which overscan is trimmed, then
        fit and subtracted
    fit_degree: int, optional
        Order of the polynomial used to fit the overscan
        
    Returns
    -------
    data_subtracted : `~numpy.ndarray`
        Trimmed and overscan-subtracted data
    header : `~astropy.io.fits.header.Header`
        `~astropy.io.fits.header.Header` object of original fits file,
        modified to say the data have been trimmed and overscan subtracted
        
    """
    
    #Trim data
    data,overscan = split_data_overscan(hdu)
    header = hdu.header
    
    #Average along columns
    avg_overscan = np.mean(overscan,axis=1)
    
    #Index array, then fit!
    idx = np.arange(len(avg_overscan))
    p = np.polyfit(idx,avg_overscan,deg=fit_degree)
    #Calculate array from fit, then transpose into a column
    fit_overscan = np.poly1d(p)(idx)
    fit_overscan_col = fit_overscan[:,np.newaxis]
    #Subtract column!
    data_subtracted = data - fit_overscan_col
    
    #Edit the header
    header.set('COMMENT','Overscan Subtracted + Trimmed')
    header.set('COMMENT','Overscan Fit Order = {0}'.format(fit_degree))
    
    return data_subtracted,header

def master_bias(biaslist,overscan_fit_degree = 8, caldir = None, overwrite = False):
    """
    Construct a master bias using median combination
    
    Parameters
    ---------
    biaslist : list
        List of filenames, should be complete filenames. Use glob to construct. 
        If the list is empty, nothing will happen
    overscan_fit_degree : int, optional
        Order of polynomial to fit overscan with
    caldir : str, optional
        Directory to place master bias into.
    overwrite : bool, optional
        If True, and caldir/master_bias.fits exists, it will be overwritten
        
    Returns
    -------
    masterbias : `~numpy.ndarray`
        Bias array. Note: this will be saved as outdir/master_bias.fits
        
    """
    
    if len(biaslist) == 0:
        print('feed me biases!')
        return None
    
    master_biases =[]
    for bias_name in biaslist:
        hdu = fits.open(bias_name)[0]
        data,header = trim_subtract_overscan(hdu,fit_degree=overscan_fit_degree)
        master_biases.append(data)
        
    master_biases = np.array(master_biases)
    master_bias = np.median(master_biases,axis=0)
    header.set('COMMENT','Biases median-combined')
    header.set('COMMENT','Composed of raw bias frames:')
    for bias_name in biaslist:
        header.set('COMMENT',bias_name)
        
    bias_hdu = fits.PrimaryHDU(master_bias,header)
    if caldir == None:
        bias_hdu.writeto('master_bias.fits', overwrite=overwrite)
    else:
        bias_hdu.writeto(caldir+'master_bias.fits', overwrite=overwrite)
        
    print('Master bias constructed')
    
    return master_bias

def master_dark(darklist,exptime,overscan_fit_degree = 8, caldir = None, overwrite = False):
    """
    Construct a master dark frame using median combination
    
    Parameters
    ---------
    darklist : list
        List of filenames, should be complete filenames. Use glob to construct. 
        If the list is empty, nothing will happen
    exptime : float
        float of exposure time for the dark in seconds. Will be appended to the filename
    overscan_fit_degree : int, optional
        Order of polynomial to fit overscan with
    caldir : str, optional
        Directory to place master dark into.
    overwrite : bool, optional
        If True, and caldir/master_dark_exptime.fits exists, it will be overwritten
        
    Returns
    -------
    masterdark : `~numpy.ndarray`
        dark array. Note: this will be saved as outdir/master_dark_exptime.fits
        
    """
    
    if len(darklist) == 0:
        print('feed me darks!')
        return None
    
    if caldir == None:
        bias = fits.getdata('master_bias.fits')
    else:
        bias = fits.getdata(caldir+'master_bias.fits')
        
    master_darks = []
    for dark_name in darklist:
        hdu = fits.open(dark_name)[0]
        data,header = trim_subtract_overscan(hdu,fit_degree=overscan_fit_degree)
        data -= bias
        master_darks.append(data)
        
    master_darks = np.array(master_darks)
    master_dark = np.median(master_darks,axis=0)
    
    #Some bookkeeping
    header.set('COMMENT','Darks median-combined')
    header.set('COMMENT','Composed of raw dark frames:')    
    for dark_name in darklist:
        header.set('COMMENT',dark_name)
    if caldir == None:
        header.set('COMMENT', 'Bias subtraction done with master_bias.fits')
    else:
        header.set('COMMENT', 'Bias subtraction done with {0}master_bias.fits'.format(caldir))
        
    dark_hdu = fits.PrimaryHDU(master_dark,header)
    if caldir == None:
        dark_hdu.writeto('master_dark_{0}.fits'.format(exptime), overwrite=overwrite)
    else:
        dark_hdu.writeto(caldir+'master_dark_{0}.fits'.format(exptime), overwrite=overwrite)
        
    print('Master dark for {0}s constructed'.format(exptime))
    
    return master_dark

def get_dark(exptime,caldir = None):
    """
    Fetch the appropriate dark frame! If it doesn't exist, scale the longest dark
    
    Parameters
    ---------
    exptime : str
        float of exposure time for the bias in seconds. Will be appended to the filename
    caldir : str, optional
       Directory to search for master dark. 
        
    Returns
    -------
    dark : `~np.ndarray`
        master dark array.
    darkname : str
        name of the file for later reference
        
    """
    
    #Search for