# -*- coding: future_fstrings -*-
# This module contains a set of functions and classes that are used in several different Python scripts in the Database.
from astropy.io import fits
from astropy.wcs import WCS

from pyFAT_astro.Support.fat_errors import FunctionCallError,BadHeaderError,\
                                        BadCubeError,BadMaskError
import pyFAT_astro.Support.support_functions as sf

#from pyFAT_astro.Support.read_functions import obtain_border_pix
import pyFAT_astro
from scipy import ndimage
import numpy as np
import copy
import warnings
import os

#Check that the mask only contains the selected sources
def check_mask(Configuration,id,Fits_Files):
    sf.print_log(f'''CHECK_MASK: Checking the mask to contain only the correct source.
''',Configuration, case =['debug_start','verbose'])
    mask = fits.open(f"{Configuration['FITTING_DIR']}/{Fits_Files['MASK']}",uint = False, do_not_scale_image_data=True,ignore_blank = True, output_verify= 'ignore')

    if float(id) not in mask[0].data:
        sf.print_log(f'''CHECK_MASK: We cannot find the selected source in the mask. This will lead to errors. Aborting the fit.
''',Configuration,case=['main','screen'])
        BadMaskError(f" We can not find the sofia source id in the mask.")
    else:
        data = copy.deepcopy(mask[0].data)
        data[data != float(id)] = 0.
        diff = data-mask[0].data
        neg_index = np.where(diff < 0.)[0]
        if len(neg_index) != 0:
            sf.print_log(f'''CHECK_MASK: The initial mask had more than a single source. redoing the mask.
''',Configuration, case = ['verbose'])
            mask[0].data = data
            fits.writeto(f"{Configuration['FITTING_DIR']}{Fits_Files['MASK']}",mask[0].data,mask[0].header, overwrite = True)
# to ensure compatible units and calculations with th models we make the maps ourselves
    del mask[0].header['C*3']
    mask[0].data[mask[0].data> 0.5] = 1.
    chan_map = np.array(np.nansum(mask[0].data,axis=0),dtype =float)
    mask[0].header['BITPIX'] = -32
    fits.writeto(f"{Configuration['FITTING_DIR']}/Sofia_Output/{Configuration['BASE_NAME']}_chan.fits",chan_map,header=mask[0].header,overwrite = True)
    mask.close()
    if Configuration['SOFIA_RAN']:
        make_moments(Configuration, Fits_Files,fit_type='Generic_Initialize',vel_unit = 'm/s')
check_mask.__doc__ =f'''
 NAME:
    check_mask

 PURPOSE:
    Check that the mask only contains the source we want and not parts of other nearby sources.
    Subsequently make the channel_map and the moment maps.

 CATEGORY:
    fits_functions

 INPUTS:
    Configuration = Standard FAT configuration
    id = the SoFiA ID of the source we are interested in
    Fits_Files = Standard FAT dictionary with filenames
    Tirific_Template = Standard FAT Tirific Template

 OPTIONAL INPUTS:
    fit_type = 'Undefined'

 OUTPUTS:
    moment maps and the channel map

 OPTIONAL OUTPUTS:
    a cleaned mask

 PROCEDURES CALLED:
    Unspecified

 NOTE:
'''




# Create a cube suitable for FAT
def create_fat_cube(Configuration, Fits_Files,sofia_catalogue=False,\
        id='No default',name='No default'):
    #First get the cubes
    if sofia_catalogue:
        Cube = fits.open(f"{Configuration['SOFIA_DIR']}{name}_{id}_cube.fits",\
            uint = False, do_not_scale_image_data=True,ignore_blank = True, output_verify= 'ignore')
    else:
        Cube = fits.open(Configuration['FITTING_DIR']+Fits_Files['ORIGINAL_CUBE'],\
            uint = False, do_not_scale_image_data=True,ignore_blank = True, output_verify= 'ignore')
    data = Cube[0].data
    hdr = Cube[0].header
    if hdr['NAXIS'] == 4:
        data = data[0,:,:,:]
        del hdr['NAXIS4']
        hdr['NAXIS'] = 3
    # clean the header
    hdr = sf.clean_header(Configuration,hdr)
    data = prep_cube(Configuration,hdr,data)
    # and write our new cube

    sf.print_log(f'''CREATE_FAT_CUBE: We are writing a FAT modified cube to be used for the fitting. This cube is called {Configuration['FITTING_DIR']+Fits_Files['FITTING_CUBE']}
''',Configuration)
    if sofia_catalogue:
        fits.writeto(f"{Configuration['MAIN_DIRECTORY']}{name}_FAT_cubelets/{name}_{id}/{name}_{id}_FAT.fits",data,hdr)
    else:
        fits.writeto(Configuration['FITTING_DIR']+Fits_Files['FITTING_CUBE'],data,hdr)
    # Release the arrays

    Cube.close()

    data = []
    hdr = []
create_fat_cube.__doc__ =f'''
 NAME:
    create_fat_cube

 PURPOSE:
    As we do not want to work from the original cube
    this function copies that cube into a new cube
    that we can then happily modify and work with.
    To keep disk space overhead this cube might later on also
    be cut down to size.

 CATEGORY:
    fits_functions

 INPUTS:
    Configuration = Standard FAT configuration
    Fits_Files = Standard FAT dictionary with filenames

 OPTIONAL INPUTS:

 OUTPUTS:
    The FAT input cube

 OPTIONAL OUTPUTS:

 PROCEDURES CALLED:
    Unspecified

 NOTE:
'''

def cut_cubes(Configuration, Fits_Files, galaxy_box):
    sf.print_log(f'''CUT_CUBES: Starting to cut the cube_size.
''', Configuration, case= ['debug_start'])
    cube_edge= [6.,5.*round(Configuration['BEAM_IN_PIXELS'][0]),5.*round(Configuration['BEAM_IN_PIXELS'][0])]
    cube_size= []
    new_cube = []
    for i in [2,1,0]:
        cube_size.append(Configuration['NAXES'][i])
        new_cube.append([0,Configuration['NAXES'][i]])
    new_cube = np.array(new_cube,dtype=int)
    cut  = False
    for i,limit in enumerate(galaxy_box):
        if limit[0] > cube_edge[i]:
            cut = True
            new_cube[i,0] = limit[0]-int(cube_edge[i])
        if limit[1] < cube_size[i] - cube_edge[i]:
            cut = True
            new_cube[i,1] = limit[1]+int(cube_edge[i])
    if 'restart_fitting' in [x.lower() for x in Configuration['FITTING_STAGES']]:
        if 'create_fat_cube' in [x.lower() for x in Configuration['FITTING_STAGES']]:
            files_to_cut = [Fits_Files['FITTING_CUBE']]
        else:
            cut = False
    else:
        files_to_cut = [Fits_Files['FITTING_CUBE'],Fits_Files['MASK'],\
                        Fits_Files['MOMENT0'],\
                        Fits_Files['MOMENT1'],\
                        Fits_Files['MOMENT2'],\
                        Fits_Files['CHANNEL_MAP'],\
                        ]
    if cut:
        sf.print_log(f'''CUT_CUBES: Your input cube is significantly larger than the detected source.
{"":8s}CUT_CUBES: we will cut to x-axis = [{new_cube[2,0]},{new_cube[2,1]}] y-axis = [{new_cube[1,0]},{new_cube[1,1]}]
{"":8s}CUT_CUBES: z-axis = [{new_cube[2,0]},{new_cube[2,1]}].
{"":8s}CUT_CUBES: We will cut the following files:
{"":8s}CUT_CUBES: {', '.join(files_to_cut)}
''', Configuration)

        for file in files_to_cut:
            sf.print_log(f'''CUT_CUBES: We are cutting the file {file}
''', Configuration)
            if os.path.exists(f"{Configuration['FITTING_DIR']}{file}"):
                cutout_cube(Configuration,file,new_cube)
            else:
                if file == Fits_Files['CHANNEL_MAP']:
                    pass
                else:
                    raise FunctionCallError(f'We are trying to cut {file} but it does not exist')


    #We want to check if the cube has a decent number of pixels per beam.
    if not os.path.exists(f"{Configuration['FITTING_DIR']}/{Fits_Files['OPTIMIZED_CUBE']}"):
        optimized_cube(Configuration, Fits_Files)
    return new_cube
cut_cubes.__doc__ =f'''
 NAME:
    cut_cubes

 PURPOSE:
    Cut all FAT related output back in size to a system that fits snugly around the sofia detection

 CATEGORY:
    fits_functions

 INPUTS:
    Configuration = Standard FAT configuration
    Fits_Files = Standard FAT dictionary with filenames
    galaxy_box = array that contains the new size as
                 [[z_min,z_max],[y_min,y_max], [x_min,x_max]]
                 adhering to fits' idiotic way of reading fits files

 OPTIONAL INPUTS:

 OUTPUTS:

 OPTIONAL OUTPUTS:

 PROCEDURES CALLED:
    cutout_cube, optimized_cube

 NOTE: The cut cubes will overwrite the _FAT products
       down to size and overwrite the existing files,
       User supplied SoFIA products are hence copied as fat products.
       To keep them safe.
'''

def cutout_cube(Configuration,filename,sub_cube):
    Cube = fits.open(Configuration['FITTING_DIR']+filename,uint = False, do_not_scale_image_data=True,ignore_blank = True, output_verify= 'ignore')
    hdr = Cube[0].header

    if hdr['NAXIS'] == 3:
        data = Cube[0].data[sub_cube[0,0]:sub_cube[0,1],sub_cube[1,0]:sub_cube[1,1],sub_cube[2,0]:sub_cube[2,1]]
        hdr['NAXIS1'] = sub_cube[2,1]-sub_cube[2,0]
        hdr['NAXIS2'] = sub_cube[1,1]-sub_cube[1,0]
        hdr['NAXIS3'] = sub_cube[0,1]-sub_cube[0,0]
        hdr['CRPIX1'] = hdr['CRPIX1'] -sub_cube[2,0]
        hdr['CRPIX2'] = hdr['CRPIX2'] -sub_cube[1,0]
        hdr['CRPIX3'] = hdr['CRPIX3'] -sub_cube[0,0]
        #Only update when cutting the cube
        Configuration['NAXES'] =[ hdr['NAXIS1'],hdr['NAXIS2'],hdr['NAXIS3']]
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            coordinate_frame = WCS(hdr)
            xlow,ylow,zlow = coordinate_frame.wcs_pix2world(1,1,1., 1.)
            xhigh,yhigh,zhigh = coordinate_frame.wcs_pix2world(*Configuration['NAXES'], 1.)
            xlim = np.sort([xlow,xhigh])
            ylim = np.sort([ylow,yhigh])
            zlim =np.sort([zlow,zhigh])/1000.
            sf.set_boundaries(Configuration,'VSYS',*zlim,input=True)
            sf.set_boundaries(Configuration,'XPOS',*xlim,input=True)
            sf.set_boundaries(Configuration,'YPOS',*ylim,input=True)

    elif hdr['NAXIS'] == 2:
        data = Cube[0].data[sub_cube[1,0]:sub_cube[1,1],sub_cube[2,0]:sub_cube[2,1]]
        hdr['NAXIS1'] = sub_cube[2,1]-sub_cube[2,0]
        hdr['NAXIS2'] = sub_cube[1,1]-sub_cube[1,0]
        hdr['CRPIX1'] = hdr['CRPIX1'] -sub_cube[2,0]
        hdr['CRPIX2'] = hdr['CRPIX2'] -sub_cube[1,0]

    Cube.close()
    fits.writeto(Configuration['FITTING_DIR']+filename,data,hdr,overwrite = True)
cutout_cube.__doc__ =f'''
 NAME:
    cutout_cube

 PURPOSE:
    Cut filename back to the size of subcube, update the header and write back to disk.

 CATEGORY:
    fits_functions

 INPUTS:
    Configuration = Standard FAT configuration
    filename = name of the cube to be cut
    sub_cube = array that contains the new size as
                [[z_min,z_max],[y_min,y_max], [x_min,x_max]]
                adhering to fits' idiotic way of reading fits files.

 OPTIONAL INPUTS:

 OUTPUTS:
    the cut cube is written to disk.

 OPTIONAL OUTPUTS:

 PROCEDURES CALLED:
    Unspecified

 NOTE:
'''

# Extract a PV-Diagrams
def extract_pv(Configuration,cube_in,angle,center=[-1,-1,-1],finalsize=[-1,-1],convert=-1):
    sf.print_log(f'''EXTRACT_PV: We are the extraction of a PV-Diagram
{'':8s} PA = {angle}
{'':8s} center = {center}
{'':8s} finalsize = {finalsize}
{'':8s} convert = {convert}
''', Configuration, case = ['debug_start'])


    cube = copy.deepcopy(cube_in)
    hdr = copy.deepcopy(cube[0].header)
    TwoD_hdr= copy.deepcopy(cube[0].header)
    data = copy.deepcopy(cube[0].data)
    #Because astro py is even dumber than Python
    try:
        if hdr['CUNIT3'].lower() == 'km/s':
            hdr['CUNIT3'] = 'm/s'
            hdr['CDELT3'] = hdr['CDELT3']*1000.
            hdr['CRVAL3'] = hdr['CRVAL3']*1000.
        elif hdr['CUNIT3'].lower() == 'm/s':
            hdr['CUNIT3'] = 'm/s'
    except KeyError:
        hdr['CUNIT3'] = 'm/s'
    if center[0] == -1:
        center = [hdr['CRVAL1'],hdr['CRVAL2'],hdr['CRVAL3']]
        xcenter,ycenter,zcenter = hdr['CRPIX1'],hdr['CRPIX2'],hdr['CRPIX3']
    else:
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            coordinate_frame = WCS(hdr)
        xcenter,ycenter,zcenter = coordinate_frame.wcs_world2pix(center[0], center[1], center[2], 1.)

    nz, ny, nx = data.shape
    if finalsize[0] != -1:
        if finalsize[1] >nz:
            finalsize[1] = nz
        if finalsize[0] > nx:
            finalsize[0] = nx
    # if the center is not set assume the crval values

    sf.print_log(f'''EXTRACT_PV: The shape of the output
{'':8s} nz = {nz}
{'':8s} ny = {ny}
{'':8s} nx = {nx}
''', Configuration,case=['verbose'])
    x1,x2,y1,y2 = sf.obtain_border_pix(Configuration,angle,[xcenter,ycenter])
    linex,liney,linez = np.linspace(x1,x2,nx), np.linspace(y1,y2,nx), np.linspace(0,nz-1,nz)
    #This only works when ny == nx hence nx is used in liney
    new_coordinates = np.array([(z,y,x)
                        for z in linez
                        for y,x in zip(liney,linex)
                        ],dtype=float).transpose().reshape((-1,nz,nx))
    #spatial_resolution = abs((abs(x2-x1)/nx)*np.sin(np.radians(angle)))+abs(abs(y2-y1)/ny*np.cos(np.radians(angle)))
    PV = ndimage.map_coordinates(data, new_coordinates,order=1)
    if hdr['CDELT1'] < 0:
        PV = PV[:,::-1]

    if finalsize[0] == -1:
        # then lets update the header
        # As python is stupid making a simple copy will mean that these changes are still applied to hudulist
        TwoD_hdr['NAXIS2'] = nz
        TwoD_hdr['NAXIS1'] = nx

        TwoD_hdr['CRPIX2'] = hdr['CRPIX3']
        TwoD_hdr['CRPIX1'] = xcenter+1
    else:
        zstart = sf.set_limits(int(zcenter-finalsize[1]/2.),0,int(nz))
        zend = sf.set_limits(int(zcenter+finalsize[1]/2.),0,int(nz))
        xstart = sf.set_limits(int(xcenter-finalsize[0]/2.),0,int(nx))
        xend = sf.set_limits(int(xcenter+finalsize[0]/2.),0,int(nx))

        PV =  PV[zstart:zend, xstart:xend]
        TwoD_hdr['NAXIS2'] = int(finalsize[1])
        TwoD_hdr['NAXIS1'] = int(finalsize[0])
        TwoD_hdr['CRPIX2'] = hdr['CRPIX3']-int(nz/2.-finalsize[1]/2.)
        TwoD_hdr['CRPIX1'] = int(finalsize[0]/2.)+1

    if convert !=-1:
        TwoD_hdr['CRVAL2'] = hdr['CRVAL3']/convert
        TwoD_hdr['CDELT2'] = hdr['CDELT3']/convert
    else:
        TwoD_hdr['CRVAL2'] = hdr['CRVAL3']
        TwoD_hdr['CDELT2'] = hdr['CDELT3']
    TwoD_hdr['CTYPE2'] = hdr['CTYPE3']
    try:
        if hdr['CUNIT3'].lower() == 'm/s' and convert == 1000. :
            TwoD_hdr['CDELT2'] = hdr['CDELT3']/1000.
            TwoD_hdr['CRVAL2'] = hdr['CRVAL3']/1000.
            TwoD_hdr['CUNIT2'] = 'km/s'
            del (TwoD_hdr['CUNIT3'])
        elif  convert != -1:
            del (TwoD_hdr['CUNIT3'])
            try:
                del (TwoD_hdr['CUNIT2'])
            except KeyError:
                pass
        else:
            TwoD_hdr['CUNIT2'] = hdr['CUNIT3']
            del (TwoD_hdr['CUNIT3'])
    except KeyError:
        sf.print_log(f'''EXTRACT_PV: We could not find units in the header for the 3rd axis.
''', Configuration)
    del (TwoD_hdr['CRPIX3'])
    del (TwoD_hdr['CRVAL3'])
    del (TwoD_hdr['CDELT3'])
    del (TwoD_hdr['CTYPE3'])
    del (TwoD_hdr['NAXIS3'])
    TwoD_hdr['NAXIS'] = 2
    TwoD_hdr['CRVAL1'] = 0.
    #Because we used nx in the linspace for liney we also use it here
    TwoD_hdr['CDELT1'] = np.sqrt(((x2-x1)*abs(hdr['CDELT1'])/nx)**2+((y2-y1)*abs(hdr['CDELT2'])/nx)**2)*3600.

    TwoD_hdr['CTYPE1'] = 'OFFSET'
    TwoD_hdr['CUNIT1'] = 'ARCSEC'
    TwoD_hdr['HISTORY'] = f'EXTRACT_PV: PV diagram extracted with angle {angle} and center {center}'
    # Then we change the cube and rteturn the PV construct
    cube[0].header = TwoD_hdr
    cube[0].data = PV

    return cube
extract_pv.__doc__ = '''
 NAME:
    extract_pv

 PURPOSE:
    extract a PV diagram from a cube object. Angle is the PA and center the central location. The profile is extracted over the full length of the cube and afterwards cut back to the finalsize.

 CATEGORY:
     fits_functions

 INPUTS:
    Configuration = Standard FAT Configuration
    cube_in = is a fits cube object
    angle = Pa of the slice

 OPTIONAL INPUTS:
    center = [-1,-1,-1]
    the central location of the slice in WCS map_coordinates [RA,DEC,VSYS], default is the CRVAL values in the header

    finalsize = [-1,-1,-1]
    final size of the PV-diagram in pixels, default is no cutting

    convert=-1
    conversion factor for velocity axis, default no conversion

 KEYWORD PARAMETERS:

 OUTPUTS:

 OPTIONAL OUTPUTS:

 PROCEDURES CALLED:
    Unspecified

 NOTE:
'''


def prep_cube(Configuration,hdr,data):
    sf.print_log( f'''PREPROCESSING: starting the preprocessing of the cube
''',Configuration, case = ['debug_start'])

    if hdr['CDELT3'] < -1:
        sf.print_log( f'''PREPROCESSING: Your velocity axis is declining with increasing channels
{"":8s}PREPROCESSING: We reversed the velocity axis.
''', Configuration)
        hdr['CDELT3'] = abs(hdr['CDELT3'])
        hdr['CRPIX3'] = hdr['NAXIS3']-hdr['CRPIX3']+1
        data = data[::-1,:,:]
    #Check for zeros
    if np.where(data == 0.):
        sf.print_log(f'''PREPROCESSING: Your cube contains values exactly 0. If this is padding these should be blanks.
{"":8s}PREPROCESSING: We have changed them to blanks.
''', Configuration)
        data[np.where(data == 0.)] = float('NaN')

    # check for blank channels and noise statistics
    cube_ok = False
    prev_first_comparison = 0.
    prev_last_comparison = 0.
    times_cut_first_channel = 0
    times_cut_last_channel = 0
    corner_box_size = int(np.floor(2*hdr['BMAJ']/abs(hdr['CDELT1'])))
    if corner_box_size > np.mean([hdr['NAXIS1'],hdr['NAXIS2']])/10.:
        corner_box_size=int(np.floor(np.mean([hdr['NAXIS1'],hdr['NAXIS2']])/10.))
    while not cube_ok:
        while np.isnan(data[0,:,:]).all():
            data=data[1:,:,:]
            hdr['NAXIS3'] = hdr['NAXIS3']-1
            hdr['CRPIX3'] = hdr['CRPIX3'] - 1
            if hdr['NAXIS3'] < 5:
                sf.print_log(f'''PREPROCESSING: This cube has too many blanked channels.
''', Configuration,case = ['main', 'screen'])
                raise BadCubeError("The cube has too many blanked channels")
            log_statement = f'''PREPROCESSING: We are cutting the cube as the first channel is completely blank.
'''
            sf.print_log(f'''PREPROCESSING: We are cutting the cube as the first channel is completely blank.
''', Configuration)
        while np.isnan(data[-1,:,:]).all():
            data=data[:-1,:,:]
            hdr['NAXIS3'] = hdr['NAXIS3']-1
            if hdr['NAXIS3'] < 5:
                sf.print_log(f'''PREPROCESSING: This cube has too many blanked channels.
''', Configuration,case = ['main', 'screen'])
                raise BadCubeError("The cube has too many blanked channels")
            sf.print_log(f'''PREPROCESSING: We are cutting the cube as the last channel is completely blank.
''', Configuration)
        #Then check the noise statistics
        noise_first_channel = np.nanstd(data[0,:,:])
        noise_last_channel = np.nanstd(data[-1,:,:])
        noise_bottom_right =  np.nanstd(data[:,:corner_box_size,-corner_box_size:])
        noise_top_right =  np.nanstd(data[:,-corner_box_size:,-corner_box_size:])
        noise_bottom_left =  np.nanstd(data[:,:corner_box_size,:corner_box_size])
        noise_top_left =  np.nanstd(data[:,-corner_box_size:,:corner_box_size])
        noise_corner = np.nanmean([noise_top_right,noise_bottom_right,noise_bottom_left,noise_top_left])
        channel_noise = np.nanmean([noise_first_channel,noise_last_channel])
        if ~np.isfinite(noise_corner):
            noise_corner = copy.deepcopy(channel_noise)
        difference = abs((noise_first_channel-noise_last_channel)/noise_first_channel)
        if noise_corner/channel_noise >1.5 and difference < 0.2:
            noise_corner = copy.deepcopy(channel_noise)
        difference2 = abs(channel_noise-noise_corner)/channel_noise
        if difference < 0.2 and np.isfinite(difference) and difference2 < 0.25:
            cube_ok = True
        else:
            sf.print_log(f'''PREPROCESSING: We are cutting the cube as clearly the noise statistics are off.
{"":8s}PREPROCESSING: Noise in the first channel is {noise_first_channel}
{"":8s}PREPROCESSING: Noise in the last channel is {noise_last_channel}
{"":8s}PREPROCESSING: Noise in the corners is {noise_corner}
''',Configuration,case= ['verbose'])
            first_comparison = abs(noise_first_channel-noise_corner)/noise_corner
            last_comparison = abs(noise_last_channel-noise_corner)/noise_corner
            if prev_first_comparison == 0:
                prev_first_comparison = first_comparison
            if prev_last_comparison == 0.:
                prev_last_comparison = last_comparison
            if (first_comparison > last_comparison and \
               abs((first_comparison-prev_first_comparison)/first_comparison) >= abs((last_comparison-prev_last_comparison)/last_comparison) and \
               times_cut_first_channel < 0.1*hdr['NAXIS3']) or \
               ~np.isfinite(noise_first_channel):
                prev_first_comparison = first_comparison
                times_cut_first_channel += 1
                data=data[1:,:,:]
                hdr['CRPIX3'] = hdr['CRPIX3'] - 1
                hdr['NAXIS3'] = hdr['NAXIS3'] - 1
            else:
                if times_cut_last_channel < 0.1*hdr['CRPIX3'] or \
                    ~np.isfinite(noise_last_channel):
                    prev_last_comparison = last_comparison
                    times_cut_last_channel += 1
                    data=data[:-1,:,:]
                    hdr['NAXIS3'] = hdr['NAXIS3'] - 1
                else:
                    cube_ok = True

            if times_cut_first_channel >= 8 and times_cut_last_channel >= 8:
                cube_ok = True
                sf.print_log( f'''PREPROCESSING: This cube has non-uniform noise statistics.'
''',Configuration)
            if hdr['NAXIS3'] < 5:
                sf.print_log(f'''PREPROCESSING: This cube has noise statistics that cannot be dealt with.
''',Configuration,case = ['main','screen'])
                raise BadCubeError('The Cube has noise statistics that cannot be dealt with')
    if ~np.isfinite(noise_corner):
        sf.print_log(f'''PREPROCESSING: This cube has noise statistics that cannot be dealt with.
''',Configuration,case = ['main','screen'])
        raise BadCubeError('The Cube has noise statistics that cannot be dealt with')
    #hdr['FATNOISE'] = noise_corner
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        low_noise_indices = np.where(data < -10*noise_corner)
    if len(low_noise_indices) > 0.0001*hdr['NAXIS1']*hdr['NAXIS2']*hdr['NAXIS3']:
        data[low_noise_indices] = float('NaN')
        sf.print_log(f'''PREPROCESSING: Your cube had a significant amount of values below -10*sigma. If you do not have a central absorption source there is something seriously wrong with the cube.
{"":8s}PREPROCESSING: We blanked these values.
''',Configuration)

    #Let's make sure that a central absorption source is treated the same regardless od the size of the cube
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        buffer=5.*hdr['BMAJ']/np.mean([abs(hdr['CDELT1']),abs(hdr['CDELT2'])])
        central = data[:,int(hdr['NAXIS2']/2.-buffer):int(hdr['NAXIS2']/2.+buffer),\
                        int(hdr['NAXIS1']/2.-buffer):int(hdr['NAXIS1']/2.+buffer)]

        central_noise_indices = np.where(central < -10*noise_corner)
    if len(central_noise_indices) > 0.0001*4*buffer*hdr['NAXIS3']:
        central[central_noise_indices] = float('NaN')
        data[:,int(hdr['NAXIS2']/2.-buffer):int(hdr['NAXIS2']/2.+buffer),\
                        int(hdr['NAXIS1']/2.-buffer):int(hdr['NAXIS1']/2.+buffer)] =\
                        central
        sf.print_log(f'''PREPROCESSING: Your cube had a significant amount of values below -10*sigma. If you do not have a central absorption source there is something seriously wrong with the cube.
{"":8s}PREPROCESSING: We blanked these values.
''',Configuration)


    # Check whether any channels got fully blanked
    blanked_channels = []
    for z in range(hdr['NAXIS3']):
        if np.isnan(data[z,:,:]).all():
            blanked_channels.append(str(z))
    if len(blanked_channels) == 0.:
        blanked_channels = ['-1']
    hdr['BL_CHAN'] = ','.join(blanked_channels)
    #Previously SoFiA was not dealing with blanks properly and the statement went here
    #finally we want to get the noise level from the negative in the final cube

    new_noise = np.nanstd(np.hstack([data[data<0],-1.*data[data<0]]))
    #If we have noiseles cubes this will be far to low
    if abs(new_noise/np.mean([channel_noise,noise_corner])) > 4. \
        or abs(new_noise/np.mean([channel_noise,noise_corner])) < 0.25:
            sf.print_log(f'''PREPROCESSING: There is something odd in the noise statistics of your cube. We are not using the nehgative values
''',Configuration)
            #new_noise = np.mean([channel_noise,noise_corner])
    hdr['FATNOISE'] = new_noise

    return data
prep_cube.__doc__ =f'''
 NAME:
    prep_cube

 PURPOSE:
    Clean up the cube  and make sure it has all the right
    characteristics and blanks

 CATEGORY:
    fits_functions

 INPUTS:
    Configuration = Standard FAT configuration
    hdr = header to be cleaned
    data = data array to be cleaned

 OPTIONAL INPUTS:

 OUTPUTS:
    the updated header

 OPTIONAL OUTPUTS:

 PROCEDURES CALLED:
    Unspecified

 NOTE:
'''

#Create an optimized cube if required
def optimized_cube(Configuration,Fits_Files):
    pix_per_beam = Configuration['BEAM_IN_PIXELS'][1]
    if pix_per_beam > Configuration['OPT_PIXEL_BEAM']:
        cube = fits.open(Configuration['FITTING_DIR']+Fits_Files['FITTING_CUBE'])
        data = cube[0].data
        hdr = cube[0].header
        if f"{abs(hdr['CDELT1']):.16f}" != f"{abs(hdr['CDELT2']):.16f}":
            sf.print_log(f'''OPTIMIZED_CUBE: Your input cube does not have square pixels.
{"":8s}OPTIMIZED_CUBE: FAT cannot optimize your cube.
''', Configuration)
        cube.close()
        required_cdelt = hdr['BMIN']/int(Configuration['OPT_PIXEL_BEAM'])
        ratio = required_cdelt/abs(hdr['CDELT2'])
        opt_data,opt_hdr = regrid_cube(data, hdr, ratio)

        fits.writeto(Configuration['FITTING_DIR']+Fits_Files['OPTIMIZED_CUBE'], opt_data,opt_hdr)

        Configuration['OPTIMIZED'] = True
        sf.print_log(f'''OPTIMIZED_CUBE: Your input cube has {pix_per_beam} pixels along the minor FWHM.
{"":8s}OPTIMIZED_CUBE: You requested { Configuration['OPT_PIXEL_BEAM']} therefore we regridded the cube into a new cube.
{"":8s}OPTIMIZED_CUBE: We are using the pixel size of {required_cdelt}.
''', Configuration)

    else:
        sf.print_log(f'''OPTIMIZED_CUBE: Your input cube has {pix_per_beam} pixels along the minor FWHM.
{"":8s}OPTIMIZED_CUBE: You requested {Configuration['OPT_PIXEL_BEAM']} but we cannot improve the resolution.
{"":8s}OPTIMIZED_CUBE: We are using the pixel size of the original cube.
''', Configuration)
optimized_cube.__doc__ =f'''
 NAME:
    optimized_cube

 PURPOSE:
    Check the amount of pixels that are in the minor axis FWHM,
    if more than OPT_PIXELBEAM regrid the cube to less pixels
    and write FAT_opt cube and set Configuration['OPTIMIZED'] = True

 CATEGORY:
    fits_functions

 INPUTS:
    Configuration = Standard FAT configuration
    Fits_Files = Standard FAT dictionary with filenames

 OPTIONAL INPUTS:

 OUTPUTS:
    An optimized cube is created and written to disk.

 OPTIONAL OUTPUTS:

 PROCEDURES CALLED:
    Unspecified

 NOTE:
'''

def make_moments(Configuration,Fits_Files,fit_type = 'Undefined',
                 moments = [0,1,2],overwrite = False, level=None,
                  vel_unit= None):
    sf.print_log(f'''MAKE_MOMENTS: We are starting to create the moment maps.
''',Configuration, case = ['debug_start'])
    if fit_type == 'Generic_Initialize':
        filename = f"{Configuration['FITTING_DIR']}{Fits_Files['FITTING_CUBE']}"
        basename = f"{Configuration['BASE_NAME']}"
        directory = f"{Configuration['FITTING_DIR']}/Sofia_Output/"
        mask_cube = f"{Configuration['FITTING_DIR']}/{Fits_Files['MASK']}"
    elif fit_type == 'Generic_Final':
        filename = f"{Configuration['FITTING_DIR']}/Finalmodel/Finalmodel.fits"
        basename = 'Finalmodel'
        directory = f"{Configuration['FITTING_DIR']}/Finalmodel/"
        mask_cube = f"{Configuration['FITTING_DIR']}/{Fits_Files['MASK']}"
    else:
        filename = f"{Configuration['FITTING_DIR']}{fit_type}/{fit_type}.fits"
        basename = fit_type
        directory = f"{Configuration['FITTING_DIR']}{fit_type}"
        if not level:
            mask_cube = f"{Configuration['FITTING_DIR']}/{Fits_Files['MASK']}"

    cube = fits.open(filename)
    if vel_unit:
        cube[0].header['CUNIT3'] = vel_unit
    if mask_cube:
        mask = fits.open(mask_cube)
        with np.errstate(invalid='ignore', divide='ignore'):
            cube[0].data[mask[0].data < 0.5] = float('NaN')
        mask.close()
    else:
        if not level:
            level = 3.*np.mean([np.nanstd(cube[0].data[0:2,:,:]),np.nanstd(cube[0].data[-3:-1,:,:])])
        with np.errstate(invalid='ignore', divide='ignore'):
            cube[0].data[cube[0].data < level] = float('NaN')
    try:
        if cube[0].header['CUNIT3'].lower().strip() == 'm/s':
            sf.print_log(f"MAKE_MOMENTS: We convert your m/s to km/s. \n", Configuration, case=['verbose'])
            cube[0].header['CUNIT3'] = 'km/s'
            cube[0].header['CDELT3'] = cube[0].header['CDELT3']/1000.
            cube[0].header['CRVAL3'] = cube[0].header['CRVAL3']/1000.
        elif cube[0].header['CUNIT3'].lower().strip() == 'km/s':
            pass
        else:
            sf.print_log(f"MAKE_MOMENTS: Your Velocity unit {cube[0].header['CUNIT3']} is weird. Your units could be off", Configuration)
    except KeyError:
        sf.print_log(f"MAKE_MOMENTS: Your CUNIT3 is missing, that is bad practice. We'll add a blank one but we're not guessing the value", Configuration)
        cube[0].header['CUNIT3'] = 'Unknown'
    #Make a 2D header to use
    hdr2D = copy.deepcopy(cube[0].header)
    hdr2D.remove('NAXIS3')
    hdr2D['NAXIS'] = 2
    # removing the third axis means we cannot correct for varying platescale, Sofia does so this is and issue so let's not do this
    hdr2D.remove('CDELT3')
    hdr2D.remove('CTYPE3')
    hdr2D.remove('CUNIT3')
    hdr2D.remove('CRPIX3')
    hdr2D.remove('CRVAL3')

    # we need a moment 0 for the moment 2 as well
    if 0 in moments:
        hdr2D['BUNIT'] = f"{cube[0].header['BUNIT']}*{cube[0].header['CUNIT3']}"
        moment0 = np.nansum(cube[0].data, axis=0) * cube[0].header['CDELT3']
        moment0[np.invert(np.isfinite(moment0))] = float('NaN')
        try:
            hdr2D['DATAMAX'] = np.nanmax(moment0)
            hdr2D['DATAMIN'] = np.nanmin(moment0)
            fits.writeto(f"{directory}/{basename}_mom0.fits",moment0,hdr2D,overwrite = overwrite)
        except ValueError:
            sf.print_log(f"MAKE_MOMENTS: Your Moment 0 has bad data and we could not write the moment 0 fits file", Configuration)


        
    if 1 in moments or 2 in moments:
        zaxis = cube[0].header['CRVAL3'] + (np.arange(cube[0].header['NAXIS3'])+1 \
              - cube[0].header['CRPIX3']) * cube[0].header['CDELT3']
        c=np.transpose(np.resize(zaxis,[cube[0].header['NAXIS1'],cube[0].header['NAXIS2'],len(zaxis)]),(2,1,0))
        hdr2D['BUNIT'] = f"{cube[0].header['CUNIT3']}"
        # Remember Python is stupid so z,y,x
        with np.errstate(invalid='ignore', divide='ignore'):
            moment1 = np.nansum(cube[0].data*c, axis=0)/ np.nansum(cube[0].data, axis=0)
        moment1[np.invert(np.isfinite(moment1))] = float('NaN')
        try:
            hdr2D['DATAMAX'] = np.nanmax(moment1)
            hdr2D['DATAMIN'] = np.nanmin(moment1)
            if 1 in moments:
                fits.writeto(f"{directory}/{basename}_mom1.fits",moment1,hdr2D,\
                    overwrite = overwrite)

        except ValueError:
            sf.print_log(f"MAKE_MOMENTS: Your Moment 1 has bad data and we could not write the moment 0 fits file", Configuration)

        if 2 in moments:
            d = c - np.resize(moment1,[len(zaxis),cube[0].header['NAXIS2'],cube[0].header['NAXIS1']])
            with np.errstate(invalid='ignore', divide='ignore'):
                moment2 = np.sqrt(np.nansum(cube[0].data*d**2, axis=0)/ np.nansum(cube[0].data, axis=0))
            moment2[np.invert(np.isfinite(moment1))] = float('NaN')
            try:
                hdr2D['DATAMAX'] = np.nanmax(moment2)
                hdr2D['DATAMIN'] = np.nanmin(moment2)
                fits.writeto(f"{directory}/{basename}_mom2.fits",\
                    moment2,hdr2D,overwrite = overwrite)

            except ValueError:
                sf.print_log(f"MAKE_MOMENTS: Your Moment 1 has bad data and we could not write the moment 0 fits file", Configuration)

    cube.close()

make_moments.__doc__ =f'''
 NAME:
    make_moments

 PURPOSE:
    Make the moment maps

 CATEGORY:
    fits_functions

 INPUTS:
    Configuration = Standard FAT configuration
    Fits_Files = Standard FAT dictionary with filenames


 OPTIONAL INPUTS:
    fit_type = 'Undefined'
    type of ftting

    moments = [0,1,2]
    moment maps to create

    overwrite = False
    overwrite existing maps

    level=None
    cutoff level to use, if set the mask will not be used

    vel_unit= none
    velocity unit of the input cube

 OUTPUTS:

 OPTIONAL OUTPUTS:

 PROCEDURES CALLED:
    Unspecified

 NOTE:
'''

def regrid_cube(data,hdr,ratio):
    regrid_hdr = copy.deepcopy(hdr)
    # First get the shape of the data
    shape = np.array(data.shape, dtype=float)
    #The new shape has to be integers
    new_shape = [int(x/ratio) for x in shape]
    # Which means our real ratio is
    real_ratio = shape/new_shape
    #* np.ceil(shape / ratio).astype(int)
    new_shape[0] = shape[0]
    # Create the zero-padded array and assign it with the old density
    regrid_data = regridder(data,new_shape)
    regrid_hdr['CRPIX1'] =   (regrid_hdr['CRPIX1']-1)/real_ratio[2]+1
    regrid_hdr['CRPIX2'] =   (regrid_hdr['CRPIX2']-1)/real_ratio[1]+1
    wcs_found = False
    try:
        regrid_hdr['CDELT1'] =   regrid_hdr['CDELT1']*real_ratio[2]
        regrid_hdr['CDELT2'] =   regrid_hdr['CDELT2']*real_ratio[1]
        wcs_found = True
    except KeyError:
        print("No CDELT found")
    try:
        regrid_hdr['CD1_1'] =   regrid_hdr['CD1_1']*real_ratio[2]
        regrid_hdr['CD2_2'] =   regrid_hdr['CD2_2']*real_ratio[1]
        wcs_found = True
    except KeyError:
        if not wcs_found:
            print("No CD corr matrix found")
    try:
        regrid_hdr['CD1_2'] =   regrid_hdr['CD1_2']*real_ratio[2]
        regrid_hdr['CD2_1'] =   regrid_hdr['CD2_1']*real_ratio[1]
        wcs_found = True
    except KeyError:
        if not wcs_found:
            print("No CD cross-corr matrix found")

    regrid_hdr['NAXIS1'] =   new_shape[2]
    regrid_hdr['NAXIS2'] =   new_shape[1]
    return regrid_data,regrid_hdr
regrid_cube.__doc__ =f'''
 NAME:
    regrid_cube

 PURPOSE:
    Regrid a cube to an arbitrary smaller cube in the spatial plane with larger pixels

 CATEGORY:
    fits_functions

 INPUTS:
    data = the array containing the input cube.
    hdr = the header data
    ratio = the requested ratio of old/new

 OPTIONAL INPUTS:

 OUTPUTS:
    regriddata = regridded data array
    regrid_hdr = header corresponding to the regridded data

 OPTIONAL OUTPUTS:

 PROCEDURES CALLED:
    copy.deepcopy, regridder, np.array

 NOTE: This might not work well with extended projections.
'''

def regridder(oldarray, newshape):
    oldshape = np.array(oldarray.shape)
    newshape = np.array(newshape, dtype=float)
    ratios = oldshape/newshape
        # calculate new dims
    nslices = [ slice(0,j) for j in list(newshape) ]
    #make a list with new coord
    new_coordinates = np.mgrid[nslices]
    #scale the new coordinates
    for i in range(len(ratios)):
        new_coordinates[i] *= ratios[i]
    #create our regridded array
    newarray = ndimage.map_coordinates(oldarray, new_coordinates,order=1)

    return newarray
regridder.__doc__ =f'''
 NAME:
regridder
 PURPOSE:
Regrid an array into a new shape through the ndimage module
 CATEGORY:
    fits_functions

 INPUTS:
    oldarray = the larger array
    newshape = the new shape that is requested

 OPTIONAL INPUTS:

 OUTPUTS:
    newarray = regridded array

 OPTIONAL OUTPUTS:

 PROCEDURES CALLED:
    scipy.ndimage.map_coordinates, np.array, np.mgrid

 NOTE:
'''
