import numpy as np
from scipy.special import gamma as gamma_func

# from scipy.special import zeta as zeta_func
from scipy.integrate import quad
from collections import namedtuple
from functools import partial

# import mpmath
import warnings
import traceback
import mpmath

# import mypath
# import convenientMath as cm
# import persistentdata as pd
# import binfootprint as bf
# import fcSpline

import logging

log = logging.getLogger(__name__)

#############################################################
#
# parameter types
#
#############################################################


def zeta_func(x, q):
    return np.complex128(mpmath.zeta(x, q))


obcf_param_type = namedtuple(
    typename="obcf_param_type", field_names=["s", "eta", "gamma"]
)

obcf_param_KAST_ANKERHOLD_type = namedtuple(
    typename="obcf_param_KAST_ANKERHOLD_type", field_names=["s", "alpha", "omega_c"]
)

#############################################################
#
# class definitiona
#
#############################################################


class OSD(object):
    r"""ohmic spectral density

    J(w) = eta w^s exp(-w / gamma)
    """

    def __init__(self, s, eta, gamma, normed=False, mp=False):
        """
        represents an super / sub Ohmic spectral density function

        normed: if set True, scaled eta, such that int_0^infty J(w) = 1

        mp: if set True, use mpmath for evaluation

        this class implements set_state and get_state and is suitable for binfootprint

        make sure, that s, eta, and gamma are regular float objects
        """

        # the "args" attributes are use for get_state and set_state
        self.s_arg = s
        self.gamma_arg = gamma
        self.normed_arg = normed
        self.mp_arg = mp
        if normed:
            self.eta_arg = None
        else:
            self.eta_arg = eta

        # non non math version
        if not mp:
            self._exp = np.exp
            self._gamma_func = gamma_func
            self.s = s
            self.gamma = gamma
            self.eta = eta

        # mpmath version
        else:
            self._exp = mpmath.exp
            self._gamma_func = mpmath.gamma
            self.s = mpmath.mpf(s)
            self.gamma = mpmath.mpf(gamma)
            self.eta = mpmath.mpf(eta)

        if normed:
            self.eta = self.eta / self.integral()
        self.c = self.eta * self._gamma_func(self.s)

    def __call__(self, w):
        if isinstance(w, np.ndarray):
            res = np.empty_like(w)
            idx_l0 = np.where(w < 0)
            res[idx_l0] = 0
            idx_ge0 = np.where(w >= 0)
            res[idx_ge0] = (
                self.eta * w[idx_ge0] ** self.s * self._exp(-w[idx_ge0] / self.gamma)
            )
            return res
        else:
            if w < 0:
                return 0
            else:
                return self.eta * w ** self.s * self._exp(-w / self.gamma)

    def maximum_at(self):
        return self.s * self.gamma

    def maximum_val(self):
        return self(self.maximum_at())

    def integral(self):
        return self.eta * self.gamma ** (self.s + 1) * self._gamma_func(self.s + 1)

    def reorganization_energy(self):
        """
        int_0^infity J(w) / w
        """
        return self.eta * self.gamma ** self.s * self._gamma_func(self.s)

    def shifted_bath_influence(self, t):
        """
        int_0^infity J(w) / w exp(-i w t) = eta [ gamma/(1+i gamma t) ]**s * Gamma(s)
        """
        return self.c * (self.gamma / (1 + 1j * self.gamma * t)) ** self.s

    def __str__(self):
        return (
            "J(w) = eta w^s exp(-w / w_c)\n"
            + "eta  = {}\n".format(self.eta)
            + "s    = {}\n".format(self.s)
            + "w_c  = {}\n".format(self.gamma)
        )

    def __getstate__(self):
        return self.s_arg, self.eta_arg, self.gamma_arg, self.normed_arg, self.mp_arg

    def __setstate__(self, state):
        self.__init__(*state)


class MakeItReal(object):
    def __init__(self, func):
        self.func = func

    def __call__(self, x):
        return np.real(self.func(x))

    def __bfkey__(self):
        return self.func


class OBCF(object):
    r"""ohmic bath correlation functions (BCF)

    consider ohmic spectral density
    J(w) = eta w^s exp(-w / gamma)

    general BCF
    alpha(tau) = 1/pi int_0^infty d w J(w) exp(-i w tau)

    -> ohmic BCF
    alpha(tau) = eta/pi gamma^(s+1) int_0^infty d x x^s exp(-x) exp(-i x gamma t)
               = eta/pi gamma^(s+1) (1 + i gamma t)^(-s-1) gamma_func(s+1)
    """

    def __init__(self, s, eta, gamma, normed=False):
        self.s = float(s)
        self.gamma = float(gamma)

        if normed:
            self.eta = None
            self._c1 = 1 / np.pi
        else:
            self.eta = eta
            self._c1 = self.eta * gamma_func(s + 1) * gamma ** (s + 1) / np.pi

    def __call__(self, tau):
        # with warnings.catch_warnings():
        #     warnings.filterwarnings('error')
        #     try:
        #         a1 = 1 + 1j*self.gamma * tau
        #         a2 = a1**(-(self.s+1))
        #         return self._c1 * a2
        #     except:
        #         traceback.print_stack()
        #         print("gamma", self.gamma)
        #         print("tau", tau)
        #         print('a1', a1)
        #         print('a2', a2)
        #         raise
        return self._c1 * (1 + 1j * self.gamma * tau) ** (-(self.s + 1))

    def div1(self, tau):
        return (
            -self._c1
            * (self.s + 1)
            * (1 + 1j * self.gamma * tau) ** (-(self.s + 2))
            * 1j
            * self.gamma
        )

    def __getstate__(self):
        return self.s, self.eta, self.gamma

    def __setstate__(self, state):
        eta = state[1]
        if eta is None:
            self.__init__(*state, normed=True)
        else:
            self.__init__(*state, normed=False)

    def __eq__(self, other):
        return (
            (self.s == other.s)
            and (self.eta == other.eta)
            and (self.gamma == other.gamma)
        )

    def __repr__(self):
        s = "\nbcf with J(w) = eta w^s exp(-w / gamma)"
        s += "\ns    :{}".format(self.s)
        s += "\neta  :{}".format(self.eta)
        s += "\ngamma:{}".format(self.gamma)

        return s


class OBCF_FT(object):
    def __init__(self, s, eta, gamma, beta):
        self.s_p1 = s + 1
        self.eta = eta
        self.gamma = gamma
        self.beta = beta
        self._c = self.eta / self.beta ** (self.s_p1) * gamma_func(self.s_p1) / np.pi
        self._beta_gamma = self.beta * self.gamma

    def __call__(self, tau):
        if isinstance(tau, np.ndarray):
            res = np.empty(shape=tau.shape, dtype=np.complex128)
            res_flat = res.flat
            tau_flat = tau.flat
            for i, ti in enumerate(tau_flat):
                zf = zeta_func(
                    self.s_p1,
                    (1 + self._beta_gamma + 1j * self.gamma * ti) / self._beta_gamma,
                )
                res_flat[i] = self._c * (
                    (self._beta_gamma / (1 + 1j * self.gamma * ti)) ** self.s_p1
                    + zf
                    + np.conj(zf)
                )
            return res
        else:
            zf = zeta_func(
                self.s_p1,
                (1 + self._beta_gamma + 1j * self.gamma * tau) / self._beta_gamma,
            )
            return self._c * (
                (self._beta_gamma / (1 + 1j * self.gamma * tau)) ** self.s_p1
                + zf
                + np.conj(zf)
            )

    def __bfkey__(self):
        return self.s_p1, self.eta, self.gamma, self.beta


class OBCF_FT_mp(object):
    def __init__(self, s, eta, gamma, beta):
        self.s_p1 = mpmath.mpf(s) + 1
        self.eta = mpmath.mpf(eta)
        self.gamma = mpmath.mpf(gamma)
        self.beta = mpmath.mpf(beta)
        self._c = (
            self.eta / self.beta ** (self.s_p1) * mpmath.gamma(self.s_p1) / mpmath.pi
        )
        self._beta_gamma = self.beta * self.gamma

    def __call__(self, tau):
        zf = mpmath.zeta(
            self.s_p1, (1 + self._beta_gamma + 1j * self.gamma * tau) / self._beta_gamma
        )
        return self._c * (
            (self._beta_gamma / (1 + 1j * self.gamma * tau)) ** self.s_p1
            + zf
            + mpmath.conj(zf)
        )


class OFTCorr(object):
    def __init__(self, s, eta, gamma, beta):
        self.s_p1 = s + 1
        self.eta = eta
        self.gamma = gamma
        self.beta = beta
        self._c = self.eta / self.beta ** (self.s_p1) * gamma_func(self.s_p1) / np.pi
        self._beta_gamma = self.beta * self.gamma

    def __call__(self, tau):
        if isinstance(tau, np.ndarray):
            res = np.empty(shape=tau.shape, dtype=np.complex128)
            res_flat = res.flat
            tau_flat = tau.flat
            for i, ti in enumerate(tau_flat):
                res_flat[i] = self._c * zeta_func(
                    self.s_p1,
                    (1 + self._beta_gamma + 1j * self.gamma * ti) / self._beta_gamma,
                )
            return res
        else:
            return self._c * zeta_func(
                self.s_p1,
                (1 + self._beta_gamma + 1j * self.gamma * tau) / self._beta_gamma,
            )

    def __bfkey__(self):
        return self.s_p1, self.eta, self.gamma, self.beta


def BOSE_distribution_single(x):
    """calc (exp(x)-1)^-1"""
    if x < 0:
        raise ValueError("x < 0")
    elif x > 40:
        return np.exp(-x)
    else:
        return 1 / np.expm1(x)


def BOSE_distribution(x):
    try:
        return np.asarray([BOSE_distribution_single(xi) for xi in x])
    except:
        return BOSE_distribution_single(x)


class OFTDens(object):
    def __init__(self, s, eta, gamma, beta):
        self.osd = OSD(s, eta, gamma)
        self.beta = beta

    def __call__(self, w):
        return BOSE_distribution(self.beta * w) * self.osd(w)

    def __bfkey__(self):
        return self.osd, self.beta


class BCF_aprx(object):
    r"""approximation of the bath correlation function using exponentials

    alpha(tau) = sum_i=1^n g_i exp(-omega_i tau)
    """

    def __init__(self, g, omega):
        g = np.asarray(g)
        omega = np.asarray(omega)
        self._n = g.shape
        assert self._n == omega.shape
        self.g = g
        self.omega = omega

    def __call__(self, tau):
        r"""return alpha(tau) = sum_i=1^n g_i exp(-w_i tau) for arbitrary shape of tau"""
        try:
            s_tau = tau.shape
            dims_tau = len(s_tau)

            tau_ = tau.reshape((1,) + s_tau)
            g_ = self.g.reshape(self._n + dims_tau * (1,))
            omega_ = self.omega.reshape(self._n + dims_tau * (1,))

            res = np.empty(shape=s_tau, dtype=np.complex128)
            idx_pos = np.where(tau >= 0)
            res[idx_pos] = np.sum(g_ * np.exp(-omega_ * tau_[(0,) + idx_pos]), axis=0)
            idx_neg = np.where(tau < 0)
            res[idx_neg] = np.conj(
                np.sum(g_ * np.exp(-omega_ * np.abs(tau_[(0,) + idx_neg])), axis=0)
            )
        except Exception as e:
            if tau >= 0:
                res = np.sum(self.g * np.exp(-self.omega * tau))
            else:
                res = np.sum(self.g * np.exp(self.omega * tau)).conj()

        return res

    def get_SD(self, t_max, n):
        """
        t0, t1, t2 ... ,tn=t_max

        al(t0), al(t1), al(t2), ... al(t n-1)    +    al(tn), al(t n+1) = al(t n-1)^ast, ... al(t 2n-1) = al(-1)^ast

        idx: 0 ... n-1   n .. -1

        -> dt = t_max / n-1

        """
        t, dt = np.linspace(0, t_max, n, retstep=True)
        alpha_t = self.__call__(t)
        alpha_t = np.hstack((alpha_t[:-1], np.conj(alpha_t[:0:-1])))

        N = 2 * (n - 1)
        t_, dt_ = np.linspace(0, 2 * t_max, N, endpoint=False, retstep=True)
        assert dt == dt_
        assert len(t_) == len(alpha_t)

        fft_alpha = np.fft.ifft(alpha_t) * dt * N

        dw = 2 * np.pi / N / dt

        freq = np.hstack((np.arange(N // 2), -np.arange(1, N // 2 + 1)[::-1])) * dw

        return 0.5 * np.fft.ifftshift(np.real(fft_alpha)), np.fft.ifftshift(freq)

    def get_SD_analyt(self, w):
        sd = np.asarray(
            [
                np.sum(
                    (
                        self.g.real * self.omega.real
                        - self.g.imag * (wi - self.omega.imag)
                    )
                    / (self.omega.real ** 2 + (wi - self.omega.imag) ** 2)
                )
                for wi in w
            ]
        )
        return sd

    def ft(self, om):
        if isinstance(om, np.ndarray):
            res = np.empty(shape=om.shape, dtype=np.complex128)
            res_flat_view = res.flat
            om_flat = om.flatten().reshape(-1, 1)
            res_flat_view[:] = np.sum(
                self.g
                * 2
                * self.omega
                / ((om_flat - 1j * self.omega) * (om_flat + 1j * self.omega)),
                axis=1,
            )
            return res

        else:
            return (
                self.g
                * 2
                * self.omega
                / ((om - 1j * self.omega) * (om + 1j * self.omega))
            )

    def n(self):
        return self._n[0]