from __future__ import annotations

from numbers import Number
from typing import TypeVar

import numpy as np
from math import sqrt, log, exp, ceil

from scipy.special import gamma

T = TypeVar('T')

__all__ = ['mlf']

# calculate Mittag-Leffler based on MATLAB code
# https://de.mathworks.com/matlabcentral/fileexchange/48154-the-mittag-leffler-function
# R. Garrappa, Numerical evaluation of two and three parameter Mittag-Leffler functions,
# SIAM Journal of Numerical Analysis, 2015, 53(3), 1350-1369


def mlf(z: T, a: float, b: float = 1, g: float = 1) -> T:
    if b == 1 and gamma == 1:
        if a == 0:
            return 1 / (1 - z)
        if a == 1:
            return np.exp(z)
        if a == 2:
            return np.cosh(z)

    if g > 1 and a > 1:
        raise ValueError('Three parameter Mittag-Leffler needs 0<alpha<1')

    if isinstance(z, Number):
        return _mlf_single(z, a, b, g)

    else:
        ret_val = np.zeros(z.size)
        for i, zz in enumerate(z.ravel()):
            ret_val[i] = _mlf_single(zz, a, b, g)

        return ret_val.reshape(z.shape)


def _mlf_single(z: int | float | complex, a: float, b: float, g: float):
    if abs(z) < 1.0e-15:
        ret_val = 1 / gamma(b)
    else:
        ret_val = laplace_inversion(z, a, b, g)

    return ret_val


def laplace_inversion(lamb, a, b, g):
    pi = np.pi
    theta = np.angle(lamb)
    j_bar = np.arange(np.ceil(-a/2 - theta/2/pi), np.floor(a/2 - theta/2/pi)+1)  # eq (2.4)
    s_star = np.abs(lamb)**(1/a) * np.exp(1j * (theta + 2*j_bar*pi/a))  # eq (2.3)

    # sort according to phi(s*) = (Re(s) + abs(s)) / 2
    phi_s = s_star.real + np.abs(s_star) / 2
    idx_phi_s = np.argsort(phi_s)
    s_star = s_star[idx_phi_s]
    phi_s = phi_s[idx_phi_s]

    # remove phi(s) = 0 to avoid
    idx_zero = phi_s > 1e-15
    s_star = s_star[idx_zero]
    phi_s = phi_s[idx_zero]

    # add origin and fictitious phi(s_j+1)=infty
    phi_s = np.r_[0, phi_s, np.infty]
    s_star = np.r_[0, s_star]

    # strength of the singularities
    p = g * np.ones(phi_s.shape, dtype=np.float)
    p[0] = max(0, -2*(a*g-b + 1))
    q = g * np.ones(phi_s.shape, dtype=np.float)
    q[-1] = np.inf

    log_epsilon = log(1e-15)  # limit for errors
    # Looking for the admissible regions with respect to round-off errors admissible_regions = find( ...
    admissable_regions = np.nonzero((phi_s[:-1] < (log_epsilon-log(np.finfo(np.float).eps))) &
                                    (phi_s[:-1] < phi_s[1:]))[0]

    # initialize vectors for optimal parameters
    last_region = admissable_regions[-1]
    mu_list = np.ones(last_region+1) * np.inf
    num_intsteps_list = np.ones(last_region+1) * np.inf
    h_list = np.ones(last_region+1) * np.inf
    
    # Evaluation of parameters for inversion of LT in each admissible region
    len_s_star = len(s_star)-1
    while True:
        for j in admissable_regions:
            if j < len_s_star:
                mu_j, h_j, n_j = optimum_bounded_region(phi_s[j], phi_s[j+1], p[j], q[j], log_epsilon)
            else:
                mu_j, h_j, n_j = optimum_unbounded_region(phi_s[j], p[j], log_epsilon)

            mu_list[j] = mu_j
            h_list[j] = h_j
            num_intsteps_list[j] = n_j

        if min(num_intsteps_list) > 200:
            log_epsilon += log(10)
        else:
            break

    # selection of admissible region for integration which involves the minimum number of nodes
    num_intsteps_min_idx = np.argmin(num_intsteps_list)
    num_intsteps = num_intsteps_list[num_intsteps_min_idx]
    mu = mu_list[num_intsteps_min_idx]
    h = h_list[num_intsteps_min_idx]

    # Evaluation of the inverse Laplace transform (eq. (2.5)
    u = h * np.arange(-num_intsteps, num_intsteps+1)
    z = mu * (1j*u + 1)**2  # eq. (3.1)
    z_deriv = -2*mu*u + 2j * mu
    g_u = np.exp(z) * z**(a*g - b) * z_deriv / (z**a - lamb)**g  # eq (2.7)
    integration_part = h * np.sum(g_u) / 2 / pi / 1j

    # evaluation of residues
    ss_star = s_star[num_intsteps_min_idx+1:]
    residues = np.sum(1. / a * ss_star**(1 - b) * np.exp(ss_star))

    # Evaluation of the ML function
    ret_val = residues + integration_part
    if np.imag(lamb) == 0.:
        ret_val = np.real(ret_val)

    return ret_val


# Finding optimal parameters in a right-bounded region
def optimum_bounded_region(phi_s_j, phi_s_j1, pj, qj, log_epsilon):
    log_eps = -36.043653389117154
    fac = 1.01
    conservative_error_analysis = False

    # Maximum value of fbar as the ration between tolerance and round-off unit
    f_max = exp(log_epsilon - log_eps)
    f_bar = np.nan
    
    # Evaluation of the starting values for sq_phi_star_j and sq_phi_star_j1
    sq_phi_j = sqrt(phi_s_j)
    threshold = 2 * sqrt(log_epsilon - log_eps)
    sq_phi_j1 = min(sqrt(phi_s_j1), threshold - sq_phi_j)

    if pj < 1.0e-14:
        if qj < 1.0e-14:
            # Zero or negative values of pj and qj
            sq_phibar_j = sq_phi_j
            sq_phibar_j1 = sq_phi_j1

        else:
            # Zero or negative values of just pj
            sq_phibar_j = sq_phi_j
            if sq_phi_j > 0:
                f_min = fac * (sq_phi_j / (sq_phi_j1-sq_phi_j))**qj
            else:
                f_min = fac

            if f_min < f_max:
                f_bar = f_min + f_min/f_max*(f_max-f_min)
                fq = f_bar**(-1/qj)
                sq_phibar_j1 = (2*sq_phi_j1 - fq*sq_phi_j)/(2+fq)
            else:
                return 0, 0, np.inf

    else:
        if qj < 1e-14:
            # zero or negative values of just q
            sq_phibar_j1 = sq_phi_j1
            f_min = fac*(sq_phi_j1/(sq_phi_j1-sq_phi_j))**pj
            if f_min < f_max:
                f_bar = f_min + f_min/f_max*(f_max-f_min)
                fp = f_bar**(-1/pj)
                sq_phibar_j = (2*sq_phi_j+fp*sq_phi_j1)/(2-fp)
            else:
                return 0, 0, np.inf

        else:
            # positive values of both pj and qj
            f_min = fac*(sq_phi_j+sq_phi_j1) / (sq_phi_j1-sq_phi_j)**max(pj, qj)
            if f_min < f_max:
                f_min = max(f_min, 1.5)
                f_bar = f_min + f_min/f_max*(f_max-f_min)
                fp = f_bar**(-1/pj)
                fq = f_bar**(-1/qj)
                if conservative_error_analysis:
                    w = -2*phi_s_j1 / (log_epsilon - phi_s_j1)
                else:
                    w = -phi_s_j1/log_epsilon
                den = 2+w - (1+w)*fp + fq
                sq_phibar_j = ((2+w+fq)*sq_phi_j + fp*sq_phi_j1)/den
                sq_phibar_j1 = (-(1+w)*fq*sq_phi_j + (2+w-(1+w)*fp)*sq_phi_j1)/den
            else:
                return 0, 0, np.inf

    log_epsilon = log_epsilon - log(f_bar)
    if conservative_error_analysis:
        w = -2*sq_phibar_j1**2 / (log_epsilon-sq_phibar_j1**2)
    else:
        w = -sq_phibar_j1**2 / log_epsilon

    mu_j = (((1+w)*sq_phibar_j + sq_phibar_j1)/(2+w))**2
    h_j = -2*np.pi / log_epsilon*(sq_phibar_j1-sq_phibar_j) / ((1+w)*sq_phibar_j + sq_phibar_j1)
    n_j = ceil(sqrt(1-log_epsilon/mu_j)/h_j)

    return mu_j, h_j, n_j


# Finding optimal parameters in a right-unbounded region
def optimum_unbounded_region(phi_s_j, pj, log_epsilon):
    # evaluation of the starting values for sq_phi_star_j
    sq_phi_s_j = sqrt(phi_s_j)
    phibar_j = phi_s_j*1.01 if phi_s_j > 0 else 0.01

    sq_phibar_j = sqrt(phibar_j)

    # Definition of some constants
    f_min = 1
    f_max = 10
    f_tar = 5

    # Iterative process to look for fbar in [f_min,f_max]
    while True:
        phi_t = phibar_j
        log_eps_phi_t = log_epsilon/phi_t
        n_j = ceil(phi_t/np.pi * (1 - 3*log_eps_phi_t/2 + sqrt(1-2*log_eps_phi_t)))
        a = np.pi*n_j/phi_t
        sq_mu_j = sq_phibar_j * abs(4-a) / abs(7-sqrt(1+12*a))
        fbar = ((sq_phibar_j-sq_phi_s_j)/sq_mu_j)**(-pj)

        if (pj < 1.0e-14) or (f_min < fbar < f_max):
            break

        sq_phibar_j = f_tar**(-1/pj)*sq_mu_j + sq_phi_s_j
        phibar_j = sq_phibar_j**2

    mu_j = sq_mu_j**2
    h_j = (-3*a - 2 + 2*sqrt(1+12*a))/(4-a)/n_j

    # Adjusting integration parameters to keep round-off errors under control
    log_eps = log(np.finfo(float).eps)
    threshold = (log_epsilon - log_eps)
    if mu_j > threshold:
        q = 0 if abs(pj) < 1.0e-14 else f_tar**(-1/pj)*sqrt(mu_j)
        phibar_j = (q + sqrt(phi_s_j))**2
        if phibar_j < threshold:
            w = sqrt(log_eps/(log_eps-log_epsilon))
            u = sqrt(-phibar_j/log_eps)
            mu_j = threshold
            n_j = ceil(w*log_epsilon/2/np.pi/(u*w-1))
            h_j = sqrt(log_eps/(log_eps - log_epsilon))/n_j
        else:
            n_j = np.inf
            h_j = 0

    return mu_j, h_j, n_j
