# Copyright 2019 Kristopher L. Kuhlman <klkuhlm _at_ sandia _dot_ gov>

# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

class InverseLaplaceTransform(object):

    def __init__(self):
        self.talbot_cache = {}
        self.stehfest_cache = {}

    def clear(self):
        self.talbot_cache = {}
        self.stehfest_cache = {}

    def calc_laplace_parameter(self,t,**kwargs):
        raise NotImplementedError

    def calc_time_domain_solution(self,fp):
        raise NotImplementedError

class FixedTalbot(InverseLaplaceTransform):
    def calc_laplace_parameter(self,t,**kwargs):

        import numpy as np
        
        # required
        # ------------------------------
        # time of desired approximation
        self.t = t

        # optional
        # ------------------------------
        # maximum time desired (used for scaling) default is requested
        # time.
        self.tmax = kwargs.get('tmax',self.t)

        self.degree = int(kwargs.get('degree',14))
        M = self.degree

        # Abate & Valko rule of thumb for r parameter
        if 'r' in kwargs:
            self.r = kwargs['r']
            r_default = False
        else:
            self.r = 2.0/5.0*M
            r_default = True

        if r_default and self.degree in self.talbot_cache:
            self.theta,self.cot_theta,self.delta = self.talbot_cache[self.degree]
        else:
            self.theta = np.linspace(0.0, np.pi, M+1, dtype=np.float64)
            self.cot_theta = np.empty((M,), dtype=np.float64)

            self.cot_theta[0] = 0.0
            self.cot_theta[1:] = 1.0/np.tan(self.theta[1:-1])

            # all but time-dependent part of p
            self.delta = np.empty((M,), dtype=np.complex64)
            self.delta[0] = self.r
            self.delta[1:] = self.r*self.theta[1:-1]*(self.cot_theta[1:] + 1j)

            self.talbot_cache[self.degree] = self.theta,self.cot_theta,self.delta
            
        self.p = self.delta/self.tmax

        # NB: p is complex

    def calc_time_domain_solution(self,fp,t):

        import numpy as np
        
        # required
        # ------------------------------
        self.t = t

        # assume fp was computed from p matrix returned from
        # calc_laplace_parameter()

        # these were computed in previous call to
        # calc_laplace_parameter()
        theta = self.theta
        delta = self.delta
        M = self.degree
        p = self.p
        r = self.r

        ans = np.empty((M,), dtype=np.complex64)
        ans[0] = np.exp(delta[0])*fp[0]/2.0

        ans[1:] = np.exp(delta[1:])*fp[1:]
        ans[1:] *= (1.0 + 1j*theta[1:-1]*(1.0 + self.cot_theta[1:]**2) -
                    1j*self.cot_theta[1:])

        result = 2.0/5.0*np.sum(ans)/self.t

        # ignore any small imaginary part
        return result.real

# ****************************************

class Stehfest(InverseLaplaceTransform):

    def calc_laplace_parameter(self,t,**kwargs):

        import numpy as np
        
        # required
        # ------------------------------
        # time of desired approximation
        self.t = t

        # optional
        # ------------------------------

        self.degree = int(kwargs.get('degree',16))

        self.ln2 = np.log(2.0)
        
        # _coeff routine requires even degree
        if self.degree%2 > 0:
            self.degree += 1

        M = self.degree

        if self.degree in self.stehfest_cache:
            self.V = self.stehfest_cache[self.degree]
        else:
            self.V = kwargs.get('V',self._coeff())
            self.stehfest_cache[self.degree] = self.V
            
        self.p = np.arange(1,M+1)*self.ln2/self.t
        
        # NB: p is real

    def _coeff(self):
        r"""Salzer summation weights (aka, "Stehfest coefficients")
        only depend on the approximation order (M) and the precision"""

        import numpy as np
        from scipy.misc import factorial
        
        M = self.degree
        M2 = int(M/2.0) # checked earlier that M is even

        V = np.empty((M,), dtype=np.float64)

        fac = lambda x: float(factorial(x,exact=True))
        
        # Salzer summation weights
        # get very large in magnitude and oscillate in sign,
        # if the precision is not high enough, there will be
        # catastrophic cancellation
        for k in range(1,M+1):
            z = np.zeros((min(k,M2)+1,), dtype=np.float64)
            for j in range(int((k+1)/2.0),min(k,M2)+1):
                z[j] = (j**M2*fac(2*j)/
                        (fac(M2-j)*fac(j)*fac(j-1)*fac(k-j)*fac(2*j-k)))
            V[k-1] = (-1)**(k+M2)*np.sum(z)

        return V

    def calc_time_domain_solution(self,fp,t):

        import numpy as np
        
        # required
        self.t = t

        # assume fp was computed from p matrix returned from
        # calc_laplace_parameter()

        result = np.dot(self.V,fp)*self.ln2/self.t

        # ignore any small imaginary part
        return result.real

# ****************************************

class deHoog(InverseLaplaceTransform):

    def calc_laplace_parameter(self,t,**kwargs):

        import numpy as np

        # NB: too many parameters, and simple p
        # nothing is cached here.
        
        self.t = t
        self.tmax = kwargs.get('tmax',self.t)

        self.degree = int(kwargs.get('degree',17))

        # 2*M+1 terms in approximation
        M = self.degree

        self.alpha = kwargs.get('alpha',1.0E-16)

        # desired tolerance (here simply related to alpha)
        self.tol = kwargs.get('tol',self.alpha*10.0)
        self.nump = 2*self.degree+1 # number of terms in approximation

        # scaling factor (likely tune-able, but 2 is typical)
        self.scale = kwargs.get('scale',2.0)
        self.T = kwargs.get('T',self.scale*self.tmax)

        self.gamma = self.alpha - np.log(self.tol)/(self.scale*self.T)
        self.p = self.gamma + 1j*np.pi*np.arange(self.nump)/self.T
        
        # NB: p is complex (mpc)

    def calc_time_domain_solution(self,fp,t):

        import numpy as np
        
        M = self.degree
        NP = self.nump
        T = self.T

        self.t = t

        # would it be useful to try re-using
        # space between e&q and A&B?
        e = np.empty((NP,M+1), dtype=np.complex64)
        q = np.empty((NP,M), dtype=np.complex64)
        d = np.empty((NP,), dtype=np.complex64)
        A = np.empty((NP+2,), dtype=np.complex64)
        B = np.empty((NP+2,), dtype=np.complex64)

        # initialize Q-D table
        e[0:2*M,0] = 0.0
        q[0,0] = fp[1]/(fp[0]/2.0)
        for i in range(1,2*M):
            q[i,0] = fp[i+1]/fp[i]

        # rhombus rule for filling triangular Q-D table (e & q)
        for r in range(1,M+1):
            # start with e, column 1, 0:2*M-2
            mr = 2*(M-r)
            e[0:mr,r] = q[1:mr+1,r-1] - q[0:mr,r-1] + e[1:mr+1,r-1]
            if not r == M:
                rq = r+1
                mr = 2*(M-rq)+1
                for i in range(mr):
                    q[i,rq-1] = q[i+1,rq-2]*e[i+1,rq-1]/e[i,rq-1]

        # build up continued fraction coefficients (d)
        d[0] = fp[0]/2.0
        for r in range(1,M+1):
            d[2*r-1] = -q[0,r-1] # even terms
            d[2*r]   = -e[0,r]   # odd terms

        # seed A and B for recurrence
        A[0] = 0.0 
        A[1] = d[0]
        B[0:2] = 1.0 

        # base of the power series
        z = np.exp(1j*np.pi*self.t/T) 

        # coefficients of Pade approximation (A & B)
        # using recurrence for all but last term
        for i in range(1,2*M):
            A[i+1] = A[i] + d[i]*A[i-1]*z
            B[i+1] = B[i] + d[i]*B[i-1]*z

        # "improved remainder" to continued fraction
        brem  = (1.0 + (d[2*M-1] - d[2*M])*z)/2.0
        rem = -brem*(1.0 - np.sqrt(1.0 + d[2*M]*z/brem**2))

        # last term of recurrence using new remainder
        A[NP] = A[2*M] + rem*A[2*M-1]
        B[NP] = B[2*M] + rem*B[2*M-1]

        # diagonal Pade approximation
        # F=A/B represents accelerated trapezoid rule
        result = np.exp(self.gamma*self.t)/T*(A[NP]/B[NP]).real

        return result

# ****************************************

# initialize classes
_fixed_talbot = FixedTalbot()
_stehfest = Stehfest()
_de_hoog = deHoog()

def invertlaplace(f, t, **kwargs):

    rule = kwargs.get('method','dehoog')
    if type(rule) is str:
        lrule = rule.lower()
        if lrule == 'talbot':
            rule = _fixed_talbot
        elif lrule == 'stehfest':
            rule = _stehfest
        elif lrule == 'dehoog':
            rule = _de_hoog
        else:
            raise ValueError("unknown invlap algorithm: %s" % rule)
    else:
        rule = rule()

    # determine the vector of Laplace-space parameter
    # needed for the requested method and desired time
    rule.calc_laplace_parameter(t,**kwargs)

    # compute the Laplace-space function evalutations
    # at the required abscissa.
    fp = [f(p) for p in rule.p]

    # compute the time-domain solution from the
    # Laplace-space function evaluations
    return rule.calc_time_domain_solution(fp,t)

# shortcuts for the above function for specific methods
def invlaptalbot(*args, **kwargs):
    kwargs['method'] = 'talbot'
    return invertlaplace(*args, **kwargs)

def invlapstehfest(*args, **kwargs):
    kwargs['method'] = 'stehfest'
    return invertlaplace(*args, **kwargs)

def invlapdehoog(*args, **kwargs):
    kwargs['method'] = 'dehoog'
    return invertlaplace(*args, **kwargs)

