__author__ = "sibirrer"

import numpy as np
import math
import scipy.special

import lenstronomy.Util.param_util as param_util

from lenstronomy.Util.package_util import exporter

export, __all__ = exporter()


@export
class ShapeletsPolar(object):
    """2D polar Shapelets, see Massey & Refregier 2005."""

    param_names = ["amp", "beta", "n", "m", "center_x", "center_y"]
    param_names_latex = {r"$I_0$", r"$\beta$", r"$n$", r"$m$", r"$x_0$", r"$y_0$"}
    lower_limit_default = {
        "amp": 0,
        "beta": 0,
        "n": 0,
        "m": 0,
        "center_x": -100,
        "center_y": -100,
    }
    upper_limit_default = {
        "amp": 100,
        "beta": 100,
        "n": 150,
        "m": 150,
        "center_x": 100,
        "center_y": 100,
    }

    def __init__(self):
        """Load interpolation of the Hermite polynomials in a range [-30,30] in order
        n<= 150 :return:"""
        pass

    def function(self, x, y, amp, beta, n, m, complex_bool, center_x, center_y):
        """

        :param x: x-coordinate, numpy array
        :param y: y-ccordinate, numpy array
        :param amp: amplitude normalization
        :param beta: shaplet scale
        :param n: order of polynomial
        :param m: rotational invariance
        :param complex_bool: boolean; if True uses complex value of function _chi_n_m()
        :param center_x: center of shapelet
        :param center_y: center of shapelet
        :return: amplitude of shapelet at possition (x, y)
        """
        r, phi = param_util.cart2polar(x, y, center_x, center_y)
        if complex_bool is True:
            return amp * self._chi_n_m(r, beta, n, m) * np.exp(-1j * m * phi).imag
        else:
            return amp * self._chi_n_m(r, beta, n, m) * np.exp(-1j * m * phi).real

    @staticmethod
    def _chi_n_m(r, beta, n, m):
        """

        :param r: radius / beta
        :param beta: shapelet scale
        :param n: non-negative integer
        :param m: integer, running from -n to n in steps of two
        :return: value of function (8) in Massey & Refregier, complex numbers
        """
        m_abs = int(abs(m))
        p = int((n - m_abs) / 2)
        p2 = int((n + m_abs) / 2)
        if p % 2 == 0:  # if p is even
            prefac = 1
        else:
            prefac = -1
        prefactor = (
            prefac
            / beta ** (m_abs + 1)
            * np.sqrt(math.factorial(p) / (np.pi * math.factorial(p2)))
        )
        poly = scipy.special.genlaguerre(
            n=p, alpha=m_abs
        )  # lower part, upper part of definition in Massey & Refregier
        r_ = (r / beta) ** 2
        l_n_alpha = poly(r_)
        return prefactor * r**m_abs * l_n_alpha * np.exp(-((r / beta) ** 2) / 2)

    @staticmethod
    def _index2n(index):
        """

        :param index: index of convention
        :return: n order of Laguerre
        """
        n_float = (-3 + np.sqrt(9 + 8 * index)) / 2
        n_int = int(n_float)
        if n_int == n_float:
            n = n_int
        else:
            n = n_int + 1
        return n

    def index2poly(self, index):
        """Manages the convention from an iterative index to the specific polynomial n,
        m, (real/imaginary part)

        :param index: int, index of list
        :return: n, m bool
        """

        n = self._index2n(index)

        num_prev = n * (n + 1) / 2
        num = index + 1
        delta = int(num - num_prev - 1)
        if n % 2 == 0:
            if delta == 0:
                m = delta
                complex_bool = False
            elif delta % 2 == 0:
                complex_bool = True
                m = delta
            else:
                complex_bool = False
                m = delta + 1
        else:
            if delta % 2 == 0:
                complex_bool = False
                m = delta + 1
            else:
                complex_bool = True
                m = delta
        return n, m, complex_bool

    @staticmethod
    def poly2index(n, m, complex_bool):
        """

        :param n: non-negative integer
        :param m: integer, running from -n to n in steps of two
        :param complex_bool: bool, if True, assigns complex part
        :return:
        """
        index = n * (n + 1) / 2
        if complex_bool is True:
            if m == 0:
                raise ValueError("m=0 can not have imaginary part!")
        if n % 2 == 0:
            if m % 2 == 0:
                if m == 0:
                    index += 1
                else:
                    index += m
                if complex_bool is True:
                    index += 1
            else:
                raise ValueError("m needs to be even for even n!")
        else:
            if complex_bool is True:
                index += m + 1
            else:
                index += m
        return int(index - 1)

    @staticmethod
    def num_param(n_max):
        """

        :param n_max: maximal polynomial order
        :return: number of basis components
        """
        return int((n_max + 1) * (n_max + 2) / 2)


@export
class ShapeletsPolarExp(object):
    """2D exponential shapelets, Berge et al.

    2019
    """

    param_names = ["amp", "beta", "n", "m", "center_x", "center_y"]
    lower_limit_default = {
        "amp": 0,
        "beta": 0,
        "n": 0,
        "m": 0,
        "center_x": -100,
        "center_y": -100,
    }
    upper_limit_default = {
        "amp": 100,
        "beta": 100,
        "n": 150,
        "m": 150,
        "center_x": 100,
        "center_y": 100,
    }

    def __init__(self):
        """Load interpolation of the Hermite polynomials in a range [-30,30] in order
        n<= 150 :return:"""
        pass

    def function(self, x, y, amp, beta, n, m, complex_bool, center_x, center_y):
        """

        :param x: x-coordinate, numpy array
        :param y: y-ccordinate, numpy array
        :param amp: amplitude normalization
        :param beta: shaplet scale
        :param n: order of polynomial
        :param m: rotational invariance
        :param complex_bool: boolean; if True uses complex value of function _chi_n_m()
        :param center_x: center of shapelet
        :param center_y: center of shapelet
        :return: amplitude of shapelet at possition (x, y)
        """
        r, phi = param_util.cart2polar(x, y, center_x, center_y)
        if complex_bool is True:
            return amp * self._chi_n_m(r, beta, n, m) * np.exp(-1j * m * phi).imag
        else:
            return amp * self._chi_n_m(r, beta, n, m) * np.exp(-1j * m * phi).real

    @staticmethod
    def _chi_n_m(r, beta, n, m):
        """

        :param r: radius
        :param beta: shapelet scale
        :param n: non-negative integer
        :param m: integer, running from -n to n in steps of two
        :return: value of function (20) in Berge et al. 2019 without complex coefficient
        """
        m_abs = int(abs(m))
        p = int(n - m_abs)
        p2 = int(n + m_abs)
        if p % 2 == 0:  # if p is even
            prefac = 1
        else:
            prefac = -1
        prefactor = prefac * np.sqrt(
            2.0
            / (beta * np.pi * (2 * n + 1) ** 3)
            * math.factorial(p)
            / math.factorial(p2)
        )
        poly = scipy.special.genlaguerre(
            n=p, alpha=2 * m_abs
        )  # lower part, upper part of definition in Massey & Refregier
        x_ = 2.0 * r / (beta * (2 * n + 1))
        l_n_alpha = poly(x_)
        return prefactor * x_**m_abs * l_n_alpha * np.exp(-x_ / 2) / np.sqrt(beta)

    @staticmethod
    def num_param(n_max):
        """

        :param n_max: maximal polynomial order
        :return: number of basis components
        """
        return int((n_max + 1) ** 2)

    @staticmethod
    def _index2n(index):
        """

        :param index: index convention
        :return: n
        """
        n_float = np.sqrt(index + 1) - 1
        n_int = int(n_float)
        if n_int == n_float:
            n = n_int
        else:
            n = n_int + 1
        return n

    def index2poly(self, index):
        """

        :param index:
        :return:
        """
        n = self._index2n(index)
        num_prev = n**2
        delta = index + 1 - num_prev
        if delta % 2 == 1:
            m = int((delta - 1) / 2)
            complex_bool = False
        else:
            m = int(delta / 2)
            complex_bool = True
        return n, m, complex_bool

    @staticmethod
    def poly2index(n, m, complex_bool):
        """

        :param n:
        :param m:
        :param complex_bool:
        :return: index convention, integer
        """
        index = n**2
        index += 2 * m
        if complex_bool is True:
            index -= 1
        return int(index)


@export
class ShapeletSetPolar(object):
    """Class to operate on entire shapelet set."""

    param_names = ["amp", "n_max", "beta", "center_x", "center_y"]
    lower_limit_default = {"beta": 0, "center_x": -100, "center_y": -100}
    upper_limit_default = {"beta": 100, "center_x": 100, "center_y": 100}

    def __init__(self, exponential=False):
        if exponential is True:
            self.shapelets = ShapeletsPolarExp()
        else:
            self.shapelets = ShapeletsPolar()
        self._exponential = exponential

    def function(self, x, y, amp, n_max, beta, center_x=0, center_y=0):
        """

        :param x:
        :param y:
        :param amp:
        :param n_max:
        :param beta:
        :param center_x:
        :param center_y:
        :return:
        """
        num_param = self.shapelets.num_param(n_max)
        f_ = np.zeros(len(np.atleast_1d(x)))
        l_list = self._pre_calc(x, y, beta, n_max, center_x, center_y)
        for i in range(num_param):
            out = self._pre_calc_function(l_list, i) * amp[i]
            f_ += out
        try:
            len(x)
        except:
            f_ = f_[0]
        return np.nan_to_num(f_)

    def function_split(self, x, y, amp, n_max, beta, center_x=0, center_y=0):
        num_param = self.shapelets.num_param(n_max)
        A = []
        L_list = self._pre_calc(x, y, beta, n_max, center_x, center_y)
        for i in range(num_param):
            A.append(self._pre_calc_function(L_list, i) * amp[i])
        return A

    def _pre_calc(self, x, y, beta, n_max, center_x, center_y):
        """

        :param x:
        :param y:
        :param beta:
        :param n_max:
        :param center_x:
        :param center_y:
        :return:
        """
        L_list = []
        # polar coordinates
        r, phi = param_util.cart2polar(x, y, center_x, center_y)
        num_param = self.shapelets.num_param(n_max)

        # compute real and imaginary part of the complex angles in the range n_max
        theta_m_real_list, theta_m_imag_list = [], []
        for i in range(n_max + 1):
            m = i
            exp_complex = np.exp(-1j * m * phi)
            theta_m_real_list.append(exp_complex.real)
            theta_m_imag_list.append(exp_complex.imag)

        # compute the Laguerre polynomials in n, m
        chi_n_m_list = [[0 for _ in range(n_max + 1)] for _ in range(n_max + 1)]
        for n in range(n_max + 1):
            for m in range(n + 1):
                if (n - m) % 2 == 0 or self._exponential is True:
                    chi_n_m_list[n][m] = self.shapelets._chi_n_m(r, beta, n, m)

        # combine together the pre-computed components
        for index in range(num_param):
            n, m, complex_bool = self.shapelets.index2poly(index)
            if complex_bool is True:
                L_i = chi_n_m_list[n][m] * theta_m_imag_list[m]
            else:
                L_i = chi_n_m_list[n][m] * theta_m_real_list[m]
            L_list.append(L_i)
        return L_list

    @staticmethod
    def _pre_calc_function(L_list, i):
        """Evaluates the shapelet function based on the pre-calculated components in
        _pre_calc()

        :param L_list: pre-calculated components
        :param i: index conventions of the sequence of basis components
        :return: shaplet basis at the pre-calculated positions
        """
        return L_list[i]

    def decomposition(self, image, x, y, n_max, beta, deltaPix, center_x=0, center_y=0):
        """Decomposes an image into the shapelet coefficients in same order as for the
        function call :param image:

        :param x:
        :param y:
        :param n_max:
        :param beta:
        :param center_x:
        :param center_y:
        :return:
        """
        num_param = self.shapelets.num_param(n_max)
        param_list = np.zeros(num_param)
        amp_norm = 1.0 * deltaPix**2
        L_list = self._pre_calc(x, y, beta, n_max, center_x, center_y)
        for i in range(num_param):
            base = self._pre_calc_function(L_list, i) * amp_norm
            param = np.sum(image * base)
            n, m, complex_bool = self.shapelets.index2poly(i)
            if m != 0:
                param *= 2
            param_list[i] = param
        return param_list

    def index2poly(self, index):
        """

        :param index: index of coefficient in the convention here
        :return: n, m, complex_bool
        """
        return self.shapelets.index2poly(index)
