#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Analysis-related functions
"""

# Third-party modules
import numpy as np
from astropy.modeling import fitting, models, Fittable2DModel
from astropy.modeling.parameters import Parameter
from astropy.time import Time

# Custom modules
from utils.logger import log

# Subclass Moffat2D to include separate gammas
class Moffat2DXY(Fittable2DModel):
    amplitude = Parameter("amplitude", default=2000, min=0, max=65535)
    x_0 = Parameter("x_0", default=0)
    y_0 = Parameter("y_0", default=0)
    gamma_x = Parameter("gamma_x", default=9)
    gamma_y = Parameter("gamma_y", default=9)
    alpha = Parameter("alpha", default=3)

    @staticmethod
    def evaluate(x, y, amplitude, x_0, y_0, gamma_x, gamma_y, alpha):
        dx = (x - x_0) / gamma_x
        dy = (y - y_0) / gamma_y
        return amplitude * (1 + dx**2 + dy**2 ) ** (-alpha)
    
    @staticmethod
    def fit_deriv(x, y, amplitude, x_0, y_0, gamma_x, gamma_y, alpha):
        dx = (x - x_0) / gamma_x
        dy = (y - y_0) / gamma_y
        norm = (1 + dx**2 + dy**2) ** (-alpha - 1)
        
        d_amplitude = norm
        d_x_0 = amplitude * (alpha + 1) * dx * norm / gamma_x
        d_y_0 = amplitude * (alpha + 1) * dy * norm / gamma_y
        d_gamma_x = amplitude * (alpha + 1) * dx**2 * norm / gamma_x
        d_gamma_y = amplitude * (alpha + 1) * dy**2 * norm / gamma_y
        d_alpha = amplitude * np.log(1 + dx**2 + dy**2) * norm
        return [d_amplitude, d_x_0, d_y_0, d_gamma_x, d_gamma_y, d_alpha]
    

class EllipticalMoffat(Fittable2DModel):
    """
    From: http://www.aspylib.com/doc/aspylib_fitting.html
    """
    
    amplitude = Parameter("amplitude", default=2000, min=0, max=65535)
    x_0 = Parameter("x_0", default=0)
    y_0 = Parameter("y_0", default=0)
    fwhm_x = Parameter("fwhm_x", default=9)
    fwhm_y = Parameter("fwhm_y", default=9)
    alpha = Parameter("alpha", default=5, max=10)
    theta = Parameter("theta", default=0, min=-np.pi/2, max=np.pi/2)
    
    @staticmethod
    def evaluate(x, y, amplitude, x_0, y_0, fwhm_x, fwhm_y, alpha, theta):

        alpha_x = fwhm_x / (2* np.sqrt(2**(1/alpha) - 1 ))
        alpha_y = fwhm_y / (2* np.sqrt(2**(1/alpha) - 1 ))
        
        A = (np.cos(theta) / alpha_x)**2  + (np.sin(theta) / alpha_y)**2
        B = (np.sin(theta) / alpha_x)**2  + (np.cos(theta) / alpha_y)**2 
        C = 2 * np.sin(theta) * np.cos(theta) * (1/alpha_x**2 - 1/alpha_y**2)
        
        return amplitude / ( 1 + A*(x-x_0)**2 + B*(y-y_0)**2 + C*(x-x_0)*(y-y_0) )**alpha
    
    
def fit_star(data, fwhm=10, background_estimation=900, model="moffat", display=False):

    # Coordinates of the max of the image
    y_star, x_star = np.unravel_index(data.argmax(), data.shape)

    # Value of the max of the image
    data_max = data[y_star, x_star]

    # Grid of indices
    y, x = np.indices(data.shape)

    # Create Background model
    surface = models.Const2D(amplitude=background_estimation)

    # Choose custom PSF model
    if model == "moffat":
        psf = models.Moffat2D(amplitude=data_max,
                              x_0=x_star,
                              y_0=y_star,
                              gamma=fwhm,
                              alpha=3.0,
                              bounds={'amplitude': (0, 65565)})

    elif model == "xy":
        psf = Moffat2DXY(amplitude=data_max,
                         x_0=x_star,
                         y_0=y_star,
                         gamma_x=fwhm,
                         gamma_y=fwhm,
                         alpha=3.0)

    elif model == "xytheta":                             
        psf = EllipticalMoffat(amplitude=data_max,
                               x_0=x_star,
                               y_0=y_star,
                               fwhm_x=fwhm,
                               fwhm_y=fwhm,
                               alpha=3.0,
                               theta=np.deg2rad(-43))

    elif model == "gaussian":
        psf = models.Gaussian2D(amplitude=data_max,
                                x_mean=x_star,
                                y_mean=y_star,
                                x_stddev=fwhm,
                                y_stddev=fwhm,
                                theta=0,
                                bounds={'amplitude': (0, 65565),
                                        'theta': (-np.pi/2, np.pi/2)})

    else:
        log.error(f"Wrong model name: {model}")
        return

    # Sum background model and custom psf model
    compound_model = surface + psf

    # Fit the data to the model
    fitter = fitting.LevMarLSQFitter()
    fitted_model = fitter(compound_model, x, y, data, maxiter=100)

    #log.debug(f"Guesses: data_max:{data_max}, x_star:{x_star}, y_star:{y_star}, fwhm:{fwhm}, bkg:{background_estimation}")
    #log.debug(fitted_model)

    # Calculate the residuals
    residuals = data - fitted_model(x, y)

    if display:

        import gnuplotlib as gp

        gp.plotimage(

            (data, dict(
                xlabel="data")),

            (fitted_model(x, y), dict(
                xlabel="model")),

            (residuals, dict(
                xlabel="Residuals")),
            
            unset="grid",
            set="size square",
            terminal='x11 1',
            multiplot="layout 1,3",
            )

    if model != "gaussian":
        if model == "moffat":
            fwhm = 2*np.abs(fitted_model.gamma_1)*np.sqrt(2**(1/fitted_model.alpha_1)-1)
        elif model == "ellip":
            fwhm = fitted_model.fwhm_x_1
        else:
            fwhm = 2*np.abs(fitted_model.gamma_x_1)*np.sqrt(2**(1/fitted_model.alpha_1)-1)

        xc = fitted_model.x_0_1
        yc = fitted_model.y_0_1
    else:
        fwhm = 2*np.sqrt(2*np.log(2))*(fitted_model.x_stddev_1)
        xc = fitted_model.x_mean_1
        yc = fitted_model.y_mean_1

    fitted_model.xc = xc.value
    fitted_model.yc = yc.value
    fitted_model.fwhm = fwhm
    fitted_model.background = fitted_model.amplitude_0.value
    fitted_model.peak = fitted_model.amplitude_1.value
    
    #log.debug(f"Fitted : data_max:{fitted_model.peak}, x_star:{xc.value}, y_star:{yc.value}, fwhm:{fwhm}, bkg:{fitted_model.background}")

    return fitted_model


def skyflat_exptime(average_counts_per_exptime, expected_counts=50000,
                    binning=1, repeat=20, min_exp=0.1, max_exp=60):
    """Calculate the exptime of sky flat images to obtain a given number
    of expected counts. Based on Tyson&Gal
    """

    now = Time.now()

    # k: relative time rate of change of sky surface brightness (sec^-1)
    # hour < 12: 'twilight', hour > 12: 'dawn'
    hour = now.datetime.hour
    k = 0.091/60

    if hour > 12:
        k = -k

    # tau: duration of astronomical twilight normalized to spring equinox at lat = +-30
    season = now.datetime.month

    if season in [12, 1, 2]:  # 'winter'
        tau = 1.33
    elif season in [6, 7, 8]:  # 'summer'
        tau = 1.91
    else:  # 'sping' or 'autumn'
        tau = 1.23

    readout_time = 14.0/(binning**2)  # In sec

    # Let's say we took an image with a given average counts / exptime
    # calculating first Texp / first exposure
    t_1 = expected_counts / average_counts_per_exptime

    a = 10**(k/tau)
    S = (a**(t_1)-1)/np.log(a)

    # Requested flats
    exptimes = []
    for _ in range(repeat):

        t_2 = np.log(a**(t_1+readout_time)+S*np.log(a))/np.log(a)
        # exptime so that (counts / pixel) is constant (in sec)
        exptime = t_2 - (t_1 + readout_time)

        if exptime < min_exp or exptime > max_exp:
            break

        exptimes.append(round(exptime, 2))
        t_1 = t_2  # Iterate

    return exptimes
