# Source code for noctua.utils.analysis

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

    """
    Analysis-related functions
    """

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

    # Other templates
    from .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:

            # Third-party modules
            import gnuplotlib as gp

            gp.plotimage(

                (data, {"xlabel": "data"}),

                (fitted_model(x, y), {"xlabel": "model"}),

                (residuals, {"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:  # 'spring' 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