'''
Transform Gaussian models
-------------------------
TrHermite
TrOchi
'''
#-------------------------------------------------------------------------------
# Name:        transform.models
# Purpose:
#
# Author:      pab
#
# Created:     24.11.2008
# Copyright:   (c) pab 2008
# Licence:     <your licence>
#-------------------------------------------------------------------------------
#!/usr/bin/env python
from __future__ import division
from scipy.optimize import brentq
from numpy import (sqrt, atleast_1d, abs, imag, sign, where, cos, arccos, ceil, #@UnresolvedImport
    expm1, log1p, pi) #@UnresolvedImport
import numpy as np
import warnings
from core import TrCommon
__all__ = ['TrHermite', 'TrLinear', 'TrOchi']

_example = '''
    >>> std = 7./4
    >>> g = <generic>(sigma=std, ysigma=std)
    
    Simulate a Transformed Gaussian process:
    >>> import numpy as np
    >>> import wafo.spectrum.models as sm
    >>> Sj = sm.Jonswap(Hm0=4*std, Tp=11)
    >>> w = np.linspace(0,4,256)
    >>> S = Sj.tospecdata(w) #Make spectrum object from numerical values
    >>> ys = S.sim(ns=15000) # Simulated in the Gaussian world
    
    >>> me, va, sk, ku = S.stats_nl(moments='mvsk')
    >>> g2 = <generic>(mean=me, var=va, skew=sk, kurt=ku, ysigma=std)
    >>> xs = g2.gauss2dat(ys[:,1:]) # Transformed to the real world 
    '''


class TrHermite(TrCommon):
    __doc__ = TrCommon.__doc__.replace('<generic>', 'Hermite') + """
    pardef : scalar, integer
        1  Winterstein et. al. (1994) parametrization [1]_ (default)
        2  Winterstein (1988) parametrization [2]_

    Description
    -----------
    The hermite transformation model is monotonic cubic polynomial, calibrated
    such that the first 4 moments of the transformed model G(y)=g^-1(y) match
    the moments of the true process. The model is given as:

        g(x) =  xn - c3(xn**2-1) - c4*(xn**3-3*xn)

    for kurt<3 (hardening model) where
        xn = (x-mean)/sigma
        c3 = skew/6
        c4 = (kurt-3)/24.

    or
        G(y) = mean + K*sigma*[ y + c3(y**2-1) + c4*(y**3-3*y) ]

    for kurt>=3 (softening model) where
        y  = g(x) = G**-1(x)
        K  = 1/sqrt(1+2*c3^2+6*c4^2)
        If pardef = 1 :
            c3  = skew/6*(1-0.015*abs(skew)+0.3*skew^2)/(1+0.2*(kurt-3))
            c4  = 0.1*((1+1.25*(kurt-3))^(1/3)-1)*c41
            c41 = (1-1.43*skew^2/(kurt-3))^(1-0.1*(kurt)^0.8)
        If pardef = 2 :
            c3 = skew/(6*(1+6*c4))
            c4 = [sqrt(1+1.5*(kurt-3))-1]/18


    Example:
    --------
    """ + _example.replace('<generic>', 'TrHermite') + """
    >>> g.dist2gauss()
     3.9858776379926808
    
    See also
    --------  
    spec2skew, ochitr, lc2tr, dat2tr

    References
    ----------
    .. [1] Winterstein, S.R, Ude, T.C. and Kleiven, G. (1994)
           "Springing and slow drift responses:
           predicted extremes and fatigue vs. simulation"
           In Proc. 7th International behaviour of Offshore structures, (BOSS)
           Vol. 3, pp.1-15
    .. [2] Winterstein, S.R. (1988)
           'Nonlinear vibration models for extremes and fatigue.'
           J. Engng. Mech., ASCE, Vol 114, No 10, pp 1772-1790
    """
    def __init__(self, *args, **kwds):
        super(TrHermite, self).__init__(*args, **kwds)
        self.pardef = kwds.get('pardef', 1)
        self._c3 = None
        self._c4 = None
        self._forward = None
        self._backward = None
        self._x_limit = None
        self.set_poly()

    def _poly_par_from_stats(self):
        skew = self.skew
        ga2 = self.kurt - 3.0
        if ga2 <= 0:
            self._c4 = ga2 / 24.
            self._c3 = skew / 6.
        elif self.pardef == 2:
            # Winterstein 1988 parametrization
            if skew ** 2 > 8 * (ga2 + 3.) / 9.:
                warnings.warn('Kurtosis too low compared to the skewness')

            self._c4 = (sqrt(1. + 1.5 * ga2) - 1.) / 18.
            self._c3 = skew / (6. * (1 + 6. * self._c4))
        else:
            # Winterstein et. al. 1994 parametrization intended to
            # apply for the range:  0 <= ga2 < 12 and 0<= skew^2 < 2*ga2/3
            if skew ** 2 > 2 * (ga2) / 3:
                warnings.warn('Kurtosis too low compared to the skewness')

            if (ga2 < 0) or (12 < ga2):
                warnings.warn('Kurtosis must be between 0 and 12')

            self._c3 = skew / 6 * (1 - 0.015 * abs(skew) + 0.3 * skew ** 2) / (1 + 0.2 * ga2)
            if ga2 == 0.:
                self._c4 = 0.0
            else:
                c41 = (1. - 1.43 * skew ** 2. / ga2) ** (1. - 0.1 * (ga2 + 3.) ** 0.8)
                self._c4 = 0.1 * ((1. + 1.25 * ga2) ** (1. / 3.) - 1.) * c41

        if not np.isfinite(self._c3) or not np.isfinite(self._c4):
            raise ValueError('Unable to calculate the polynomial')
        
    def set_poly(self):
        '''
        Set poly function from stats (i.e., mean, sigma, skew and kurt) 
        '''

        if self._c3 is None:
            self._poly_par_from_stats()
        eps = np.finfo(float).eps
        c3 = self._c3
        c4 = self._c4
        ma = self.mean
        sa = self.sigma
        if abs(c4) < sqrt(eps):
            c4 = 0.0

        #gdef = self.kurt-3.0
        if self.kurt < 3.0:
            p = np.poly1d([-c4, -c3, 1. + 3. * c4, c3]) # forward, g
            self._forward = p
            self._backward = None
        else:
            Km1 = np.sqrt(1. + 2. * c3 ** 2 + 6 * c4 ** 2)
            p = np.poly1d(np.r_[c4, c3, 1. - 3. * c4, -c3] / Km1)   # backward G
            self._forward = None
            self._backward = p

        #% Check if it is a strictly increasing function.
        dp = p.deriv(m=1) #% Derivative
        r = dp.r         #% Find roots of the derivative
        r = r[where(abs(imag(r)) < eps)] # Keep only real roots

        if r.size > 0:
            # Compute where it is possible to invert the polynomial
            if self.kurt < 3.:
                self._x_limit = r
            else:
                self._x_limit = sa * p(r) + ma
            txt1 = ''' 
                The polynomial is not a strictly increasing function.
                The derivative of g(x) is infinite at x = %g''' % self._x_limit
            warnings.warn(txt1)
        return
    def check_forward(self, x):
        if not (self._x_limit is None):
            x00 = self._x_limit
            txt2 = 'for the given interval x = [%g, %g]' % (x[0], x[-1])

            if any(np.logical_and(x[0] <= x00, x00 <= x[-1])):
                cdef = 1
            else:
                cdef = sum(np.logical_xor(x00 <= x[0] , x00 <= x[-1]))

            if np.mod(cdef, 2):
                errtxt = 'Unable to invert the polynomial \n %s' % txt2
                raise ValueError(errtxt)
            np.disp('However, successfully inverted the polynomial\n %s' % txt2)


    def _dat2gauss(self, x, *xi):
        if len(xi) > 0:
            raise ValueError('Transforming derivatives is not implemented!')
        xn = atleast_1d(x)
        self.check_forward(xn)

        xn = (xn - self.mean) / self.sigma

        if self._forward is None:
            #Inverting the polynomial
            yn = self._poly_inv(self._backward, xn)
        else:
            yn = self._forward(xn)
        return yn * self.ysigma + self.ymean

    def _gauss2dat(self, y, *yi):
        if len(yi) > 0:
            raise ValueError('Transforming derivatives is not implemented!')
        yn = (atleast_1d(y) - self.ymean) / self.ysigma
        #self.check_forward(y)

        if self._backward is None:
            #% Inverting the polynomial
            #%~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
            xn = self._poly_inv(self._forward, yn)
        else:
            xn = self._backward(yn)
        return self.sigma * xn + self.mean
    
    def _poly_inv(self, p, xn):
        '''
        Invert polynomial
        '''
        if p.order < 2:
            return xn
        elif p.order == 2:
            # Quadratic: Solve a*u**2+b*u+c = xn
            coefs = p.coeffs
            a = coefs[0]
            b = coefs[1]
            c = coefs[2] - xn
            t = 0.5 * (b + sign(b) * sqrt(b ** 2 - 4 * a * c))
            #so1 = t/a # largest solution
            so2 = -c / t # smallest solution
            return so2
        elif p.order == 3:
            # Solve
            # K*(c4*u^3+c3*u^2+(1-3*c4)*u-c3) = xn = (x-ma)/sa
            # -c4*xn^3-c3*xn^2+(1+3*c4)*xn+c3 = u
            coefs = p.coeffs[1::] / p.coeffs[0]
            a = coefs[0]
            b = coefs[1]
            c = coefs[2] - xn / p.coeffs[0]

            x0 = a / 3.
            #% substitue xn = z-x0  and divide by c4 => z^3 + 3*p1*z+2*q0  = 0
            p1 = b / 3 - x0 ** 2
            #p1 = (b-a**2/3)/3


            #q0 = (c + x0*(2.*x0/3.-b))/2.
            #q0 = x0**3 -a*b/6 +c/2
            q0 = x0 * (x0 ** 2 - b / 2) + c / 2
##            # z^3+3*p1*z+2*q0=0

##            c3 = self._c3
##            c4 = self._c4
##            b1 = 1./(3.*c4)
##            #x0 = c3*b1 
##            #% substitue u = z-x0  and divide by c4 => z^3 + 3*c*z+2*q0  = 0
##            #p1  = b1-1.-x0**2.
##            Km1 = np.sqrt(1.+2.*c3**2+6*c4**2)
##            q0 = x0**3-1.5*b1*(x0+xn*Km1)
            #q0 = x0**3-1.5*b1*(x0+xn)
            if not (self._x_limit is None):  # % Three real roots
                d = sqrt(-p1)
                theta1 = arccos(-q0 / d ** 3) / 3
                th2 = np.r_[0, -2 * pi / 3, 2 * pi / 3]
                x1 = abs(2 * d * cos(theta1[ceil(len(xn) / 2)] + th2) - x0)
                ix = x1.argmin() #   % choose the smallest solution
                return 2. * d * cos(theta1 + th2[ix]) - x0
            else:                # %Only one real root exist
                q1 = sqrt((q0) ** 2 + p1 ** 3)
                #% Find the real root of the monic polynomial
                A0 = (q1 - q0) ** (1. / 3.)
                B0 = -(q1 + q0) ** (1. / 3.)
                return A0 + B0 - x0  #% real root
                #%% The other complex roots are given by
                #%x= -(A0+B0)/2+(A0-B0)*sqrt(3)/2-x0
                #%x=-(A0+B0)/2+(A0-B0)*sqrt(-3)/2-x0



class TrLinear(TrCommon):
    __doc__ = TrCommon.__doc__.replace('<generic>', 'Linear') + """
    Description
    -----------
    The linear transformation model is monotonic linear polynomial, calibrated
    such that the first 2 moments of the transformed model G(y)=g^-1(y) match
    the moments of the true process. 

    Example:
    --------
    """ + _example.replace('<generic>', 'TrLinear') + """
    >>> g.dist2gauss()
    0.0
    
    See also
    --------  
    spec2skew, ochitr, lc2tr, dat2tr

    """
    def _dat2gauss(self, x, *xi):
        sratio = atleast_1d(self.ysigma / self.sigma)
        y = (atleast_1d(x) - self.mean) * sratio + self.ymean
        if len(xi) > 0:
            y = [y, ] + [ ix * sratio for ix in xi]
        return y

    def _gauss2dat(self, y, *yi):
        sratio = atleast_1d(self.sigma / self.ysigma)
        x = (atleast_1d(y) - self.ymean) * sratio + self.mean
        if len(yi) > 0:
            x = [x, ] + [iy * sratio for iy in yi]
        return x
    
   
class TrOchi(TrCommon):
    __doc__ = TrCommon.__doc__.replace('<generic>', 'Ochi') + """

    Description
    -----------
    The Ochi transformation model is a monotonic exponential function,
    calibrated such that the first 3 moments of the transformed model
    G(y)=g^-1(y) match the moments of the true  process. However, the
    skewness is limited by ABS(SKEW)<2.82. According to Ochi it is
    appropriate for a process with very strong non-linear characteristics.
    The model is given as:
        g(x) = ((1-exp(-gamma*(x-mean)/sigma))/gamma-mean2)/sigma2
    where
        gamma  = 1.28*a  for x>=mean
                 3*a     otherwise
        mean,
        sigma  = standard deviation and mean, respectively, of the process.
        mean2,
        sigma2 = normalizing parameters in the transformed world, i.e., to
                make the gaussian process in the transformed world is N(0,1).

    The unknown parameters a, mean2 and sigma2 are found by solving the
    following non-linear equations:

        a*(sigma2^2+mean2^2)+mean2 = 0
           sigma2^2-2*a^2*sigma2^4 = 1
    2*a*sigma2^4*(3-8*a^2*sigma2^2) = skew

    Note
    ----
    Transformation, g, does not have continous derivatives of 2'nd order or higher.

    Example
    -------
    """ + _example.replace('<generic>', 'TrOchi') + """
    >>> g.dist2gauss()
    5.9322684525265501
    
    See also
    --------
    spec2skew, hermitetr, lc2tr, dat2tr

    References
    ----------
    Ochi, M.K. and Ahn, K. (1994)
    'Non-Gaussian probability distribution of coastal waves.'
    In Proc. 24th Conf. Coastal Engng, Vol. 1, pp 482-496

    Michel K. Ochi (1998),
    "OCEAN WAVES, The stochastic approach",
    OCEAN TECHNOLOGY series 6, Cambridge, pp 255-275.
    """

    def __init__(self, *args, **kwds):
        super(TrOchi, self).__init__(*args, **kwds)
        self.kurt = None
        self._phat = None
        self._par_from_stats()
        
    def _par_from_stats(self):
        skew = self.skew
        if abs(skew) > 2.82842712474619:
            raise ValueError('Skewness must be less than 2.82842')
        
        mean1 = self.mean
        sigma1 = self.sigma

        if skew == 0:
            self._phat = [sigma1, mean1, 0, 0, 1, 0]
            return

        # Solve the equations to obtain the gamma parameters:
        #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
        #          a*(sig2^2+ma2^2)+ma2 = 0
        #           sig2^2-2*a^2*sig2^4 = E(y^2) % =1
        #   2*a*sig2^4*(3-8*a^2*sig2^2) = E(y^3) % = skew

        # Let x = [a sig2^2 ]
        # Set up the 2D non-linear equations for a and sig2^2:
        # g1='[x(2)-2.*x(1).^2.*x(2).^2-P1, 2.*x(1).*x(2).^2.*(3-8.*x(1).^2.*x(2))-P2  ]'
        # Or solve the following 1D non-linear equation for sig2^2:
        g2 = lambda x:-sqrt(abs(x - 1) * 2) * (3. * x - 4 * abs(x - 1)) + abs(skew)


        a1 = 1. # % Start interval where sig2^2 is located.
        a2 = 2.

        sig22 = brentq(g2, a1, a2) #% smallest solution for sig22
        a = sign(skew) * sqrt(abs(sig22 - 1) / 2) / sig22
        gam_a = 1.28 * a
        gam_b = 3 * a
        sigma2 = sqrt(sig22)

        #% Solve the following 2nd order equation to obtain ma2
        #%        a*(sig2^2+ma2^2)+ma2 = 0
        my2 = (-1. - sqrt(1. - 4. * a ** 2 * sig22)) / a  #% Largest mean
        mean2 = a * sig22 / my2                   #% choose the smallest mean

        self._phat = [sigma1, mean1, gam_a, gam_b, sigma2, mean2]
        return
    
    def _get_par(self):
        '''
        Returns ga, gb, sigma2, mean2
        '''
        if (self._phat is None or self.sigma != self._phat[0] 
            or self.mean != self._phat[1]): 
            self._par_from_stats()
        #sigma1 = self._phat[0]
        #mean1 = self._phat[1]
        ga = self._phat[2]
        gb = self._phat[3]
        sigma2 = self._phat[4]
        mean2 = self._phat[5]
        return ga, gb, sigma2, mean2

    def _dat2gauss(self, x, *xi):
        if len(xi) > 0:
            raise ValueError('Transforming derivatives is not implemented!')
        ga, gb, sigma2, mean2 = self._get_par()
        mean = self.mean
        sigma = self.sigma
        xn = atleast_1d(x)
        xn = (xn - mean) / sigma
        igp, = where(0 <= xn)
        igm, = where(xn < 0)
        
        g = xn.copy()

        if ga != 0:
            np.put(g, igp, (-expm1(-ga * xn[igp])) / ga)

        if gb != 0:
            np.put(g, igm, (-expm1(-gb * xn[igm])) / gb)

        return (g - mean2) * self.ysigma / sigma2 + self.ymean

    def _gauss2dat(self, y, *yi):
        if len(yi) > 0:
            raise ValueError('Transforming derivatives is not implemented!')
        
        ga, gb, sigma2, mean2 = self._get_par()
        mean = self.mean
        sigma = self.sigma
        
        yn = (atleast_1d(y) - self.ymean) / self.ysigma
        xn = sigma2 * yn + mean2

        igp, = where(0 <= xn)
        igm, = where(xn < 0)

        if ga != 0:
            np.put(xn, igp, -log1p(-ga * xn[igp]) / ga)

        if gb != 0:
            np.put(xn, igm, -log1p(-gb * xn[igm]) / gb)

        return  sigma * xn + mean

def main():
    import pylab
    g = TrHermite(skew=0.1, kurt=3.01)
    g.dist2gauss()
    #g = TrOchi(skew=0.56)
    x = np.linspace(-5, 5)
    y = g(x)
    pylab.plot(np.abs(x - g.gauss2dat(y)))
    #pylab.plot(x,y,x,x,':',g.gauss2dat(y),y,'r')

    pylab.show()
    np.disp('finito')
    
if __name__ == '__main__':
    if True : # False: #
        import doctest
        doctest.testmod()
    else:
        main()



