# based on https://github.com/tremou/RGB_MSUastro

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

def sky_median_sig_clip(input_arr, sig_fract, percent_fract, max_iter=100):
	"""Estimates median sky value for a given number of iterationsself.

    Parameters
    ----------
	input_arr: numpy array
	   Image data array
	sig_fract: float
	   Fraction of sigma clipping.
	percent_fract: float
        Convergence fraction.
	max_iter: int, default `100`
        Maximum number of iterations

    Returns
    -------
	(new_sky, iteration): tuple
	   Median sky value and number of iteration
	"""
	work_arr = np.ravel(input_arr)
	old_sky = np.median(work_arr)
	sig = work_arr.std()
	upper_limit = old_sky + sig_fract * sig
	lower_limit = old_sky - sig_fract * sig
	indices = np.where((work_arr < upper_limit) & (work_arr > lower_limit))
	work_arr = work_arr[indices]
	new_sky = np.median(work_arr)
	iteration = 0

	while ((math.fabs(old_sky - new_sky)/new_sky) > percent_fract) and (iteration < max_iter) :
		iteration += 1
		old_sky = new_sky
		sig = work_arr.std()
		upper_limit = old_sky + sig_fract * sig
		lower_limit = old_sky - sig_fract * sig
		indices = np.where((work_arr < upper_limit) & (work_arr > lower_limit))
		work_arr = work_arr[indices]
		new_sky = np.median(work_arr)

	return (new_sky, iteration)

def sky_mean_sig_clip(input_arr, sig_fract, percent_fract, max_iter=100):
	"""Estimates mean sky value for a given number of iterations.

	Parameters
    ----------
	input_arr: numpy array
	   Image data array
	sig_fract: float
	   Fraction of sigma clipping.
	percent_fract: float
        Convergence fraction.
	max_iter: int, default `100`
        Maximum number of iterations

    Returns
    -------
	(new_sky, iteration): tuple
	   Mean sky value and number of iteration
	"""
	work_arr = np.ravel(input_arr)
	old_sky = np.mean(work_arr)
	sig = work_arr.std()
	upper_limit = old_sky + sig_fract * sig
	lower_limit = old_sky - sig_fract * sig
	indices = np.where((work_arr < upper_limit) & (work_arr > lower_limit))
	work_arr = work_arr[indices]
	new_sky = np.mean(work_arr)
	iteration = 0

	while ((math.fabs(old_sky - new_sky)/new_sky) > percent_fract) and (iteration < max_iter) :
		iteration += 1
		old_sky = new_sky
		sig = work_arr.std()
		upper_limit = old_sky + sig_fract * sig
		lower_limit = old_sky - sig_fract * sig
		indices = np.where((work_arr < upper_limit) & (work_arr > lower_limit))
		work_arr = work_arr[indices]
		new_sky = np.mean(work_arr)

	return (new_sky, iteration)

def linear(inputArray, scale_min=None, scale_max=None):
	"""Performs linear scaling of the input numpy array.

    Parameters
    ----------
	inputArray: numpy array
	   Image data array.
	scale_min: float, default `None`
	   Minimum data value.
	scale_max: float, default `None`
	   Maximum data value.

    Returns
    -------
	imageData: numpy array
	   Scaled image data array.
	"""
	imageData=np.array(inputArray, copy=True)

	if scale_min == None:
		scale_min = np.nanmin(imageData)
	if scale_max == None:
		scale_max = np.nanmax(imageData)

	imageData = imageData.clip(min=scale_min, max=scale_max)
	imageData = (imageData -scale_min) / (scale_max - scale_min)
	indices = np.where(imageData < 0)
	imageData[indices] = 0.0
	indices = np.where(imageData > 1)
	imageData[indices] = 1.0

	return imageData

def sqrt(inputArray, scale_min=None, scale_max=None):
	"""Performs sqrt scaling of the input numpy array.

	Parameters
    ----------
	inputArray: numpy array
	   Image data array.
	scale_min: float, default `None`
	   Minimum data value.
	scale_max: float, default `None`
	   Maximum data value.

    Returns
    -------
	imageData: numpy array
	   Scaled image data array.
	"""
	imageData=np.array(inputArray, copy=True)

	if scale_min == None:
		scale_min = np.nanmin(imageData)
	if scale_max == None:
		scale_max = np.nanmax(imageData)

	imageData = imageData.clip(min=scale_min, max=scale_max)
	imageData = imageData - scale_min
	indices = np.where(imageData < 0)
	imageData[indices] = 0.0
	imageData = np.sqrt(imageData)
	imageData = imageData / math.sqrt(scale_max - scale_min)

	return imageData

def log(inputArray, scale_min=None, scale_max=None):
	"""Performs log10 scaling of the input numpy array.

	Parameters
    ----------
	inputArray: numpy array
	   Image data array.
	scale_min: float, default `None`
	   Minimum data value.
	scale_max: float, default `None`
	   Maximum data value.

    Returns
    -------
	imageData: numpy array
	   Scaled image data array.
	"""
	imageData=np.array(inputArray, copy=True)

	if scale_min == None:
		scale_min = np.nanmin(imageData)
	if scale_max == None:
		scale_max = np.nanmax(imageData)

	factor = math.log10(scale_max - scale_min)
	indices0 = np.where(imageData < scale_min)
	indices1 = np.where((imageData >= scale_min) & (imageData <= scale_max))
	indices2 = np.where(imageData > scale_max)
	imageData[indices0] = 0.0
	imageData[indices2] = 1.0
	try :
		imageData[indices1] = np.log10(imageData[indices1])/factor
	except :
		print("Error on math.log10 for ", (imageData[i][j] - scale_min))

	return imageData

def asinh(inputArray, scale_min=None, scale_max=None, non_linear=2.0):
	"""Performs asinh scaling of the input numpy array.

    Parameters
    ----------
	inputArray: numpy array
	   Image data array.
	scale_min: float, default `None`
	   Minimum data value.
	scale_max: float, default `None`
	   Maximum data value.
    non_linear: float, default `2.0`
  	   Non-linearity factor.

    Returns
    -------
	imageData: numpy array
	   Scaled image data array.
	"""
	imageData=np.array(inputArray, copy=True)

	if scale_min == None:
		scale_min = np.nanmin(imageData)
	if scale_max == None:
		scale_max = np.nanmax(imageData)

	factor = np.arcsinh((scale_max - scale_min)/non_linear)
	indices0 = np.where(imageData < scale_min)
	indices1 = np.where((imageData >= scale_min) & (imageData <= scale_max))
	indices2 = np.where(imageData > scale_max)
	imageData[indices0] = 0.0
	imageData[indices2] = 1.0
	imageData[indices1] = np.arcsinh((imageData[indices1] - \
	scale_min)/non_linear)/factor

	return imageData

def create_RGB_image(outfile=None, survey='PS1', filters='zir', images_dir='',
                                scaling='linear', scaling_params=None):
    """Creates an RGB image.

    Parameters
    ----------
    outfile: str, default `None`
        Output file name. If `None`, no output is saved.
    survey: str, default `PS1`
        Survey used for the filters.
    filters: str, defaul `zir`
        Filters used to create the RGB image (in that order). Must be three filters.
    images_dir: str, default ``
        Directory where to find the images.
    scaling: str
        Type of scaling. Choose between `linear`, `sqrt`, `log` and `asinh`.
    scaling_params: dict
        Dictionary with the parameters for the `scaling` function. If `None`,
        use the default values.

    Example of scaling_params:

        scaling_params = {'R':{'min':r_med/100, 'max':r_med*30},
                          'G':{'min':g_med/1000, 'max':g_med*150},
                          'B':{'min':b_med/30, 'max':b_med*100}
        }
    """

    assert len(filters)==3, 'Three filters should be given.'
    scalings = ['linear', 'sqrt', 'log', 'asinh']
    assert scaling in scalings, f'Not a valid scaling (choose between {scalings})'

    rgb_files = []
    for filt in filters:
        filt_image = os.path.join(images_dir, f'{survey}_{filt}.fits')
        rgb_files.append(filt_image)

    r_img = fits.getdata(rgb_files[0], ext=0)
    r_med = np.nanmedian(r_img)
    g_img = fits.getdata(rgb_files[1], ext=0)
    g_med = np.nanmedian(g_img)
    b_img = fits.getdata(rgb_files[2], ext=0)
    b_med = np.nanmedian(b_img)
    imgs = [r_img, g_img, b_img]

    if scaling_params is None:
        scaling_params = {'R':{'min':r_med/100, 'max':r_med*30},
                          'G':{'min':g_med/1000, 'max':g_med*150},
                          'B':{'min':b_med/30, 'max':b_med*100}
        }
    scaling_func = eval(scaling)  # scaling function

    img = np.zeros((r_img.shape[0], r_img.shape[1], 3), dtype=float)
    for i, filt_img in enumerate(imgs):
        img_scalings = list(scaling_params.values())[i]
        scale_min, scale_max = img_scalings.values()
        img[:,:,i] = scaling_func(filt_img, scale_min=scale_min, scale_max=scale_max)

    fig = plt.figure(figsize=(12, 12))
    plt.clf()
    plt.imshow(img, origin='lower')
    plt.title(f'Blue = {filters[0]}, Green = {filters[1]}, Red = {filters[2]}')
    if outfile is not None:
        plt.savefig(outfile)
