r"""
Implements turbulence closure model stability functions.

.. math::
    S_m &= S_m(\alpha_M, \alpha_N) \\
    S_\rho &= S_\rho(\alpha_M, \alpha_N)

where :math:`\alpha_M, \alpha_N` are the normalized shear and buoyancy frequency

    .. math::
        \alpha_M &= \frac{k^2}{\varepsilon^2} M^2 \\
        \alpha_N &= \frac{k^2}{\varepsilon^2} N^2

The following stability functions have been implemented

- Canuto A
- Canuto B
- Kantha-Clayson
- Cheng

References:

Umlauf, L. and Burchard, H. (2005). Second-order turbulence closure models
for geophysical boundary layers. A review of recent work. Continental Shelf
Research, 25(7-8):795--827. http://dx.doi.org/10.1016/j.csr.2004.08.004

Burchard, H. and Bolding, K. (2001). Comparative Analysis of Four
Second-Moment Turbulence Closure Models for the Oceanic Mixed Layer. Journal of
Physical Oceanography, 31(8):1943--1968.
http://dx.doi.org/10.1175/1520-0485(2001)031

Umlauf, L. and Burchard, H. (2003). A generic length-scale equation for
geophysical turbulence models. Journal of Marine Research, 61:235--265(31).
http://dx.doi.org/10.1357/002224003322005087

"""
from __future__ import absolute_import
import numpy as np
from abc import ABC, abstractmethod
from scipy.optimize import minimize
from .log import print_output


__all__ = [
    'StabilityFunctionBase',
    'StabilityFunctionCanutoA',
    'StabilityFunctionCanutoB',
    'StabilityFunctionCheng',
    'GOTMStabilityFunctionCanutoA',
    'GOTMStabilityFunctionCanutoB',
    'GOTMStabilityFunctionCheng',
    'GOTMStabilityFunctionKanthaClayson',
    'compute_normalized_frequencies'
]


def compute_normalized_frequencies(shear2, buoy2, k, eps, verbose=False):
    r"""
    Computes normalized buoyancy and shear frequency squared.

    .. math::
        \alpha_M &= \frac{k^2}{\varepsilon^2} M^2 \\
        \alpha_N &= \frac{k^2}{\varepsilon^2} N^2

    From Burchard and Bolding (2001).

    :arg shear2: :math:`M^2`
    :arg buoy2: :math:`N^2`
    :arg k: turbulent kinetic energy
    :arg eps: TKE dissipation rate
    """
    alpha_buoy = k**2/eps**2*buoy2
    alpha_shear = k**2/eps**2*shear2
    if verbose:
        print_output('{:8s} {:8.3e} {:8.3e}'.format('M2', shear2.min(), shear2.max()))
        print_output('{:8s} {:8.3e} {:8.3e}'.format('N2', buoy2.min(), buoy2.max()))
        print_output('{:8s} {:10.3e} {:10.3e}'.format('a_buoy', alpha_buoy.min(), alpha_buoy.max()))
        print_output('{:8s} {:10.3e} {:10.3e}'.format('a_shear', alpha_shear.min(), alpha_shear.max()))
    return alpha_buoy, alpha_shear


class StabilityFunctionBase(ABC):
    """
    Base class for all stability functions
    """
    @property
    @abstractmethod
    def name(self):
        pass

    def __init__(self, lim_alpha_shear=True, lim_alpha_buoy=True,
                 smooth_alpha_buoy_lim=True, alpha_buoy_crit=-1.2):
        r"""
        :kwarg bool lim_alpha_shear: limit maximum :math:`\alpha_M` values
            (see Umlauf and Burchard (2005) eq. 44)
        :kwarg bool lim_alpha_buoy: limit minimum (negative) :math:`\alpha_N` values
            (see Umlauf and Burchard (2005))
        :kwarg bool smooth_alpha_buoy_lim: if :math:`\alpha_N` is limited, apply a
            smooth limiter (see Burchard and Bolding (2001) eq. 19). Otherwise
            :math:`\alpha_N` is clipped at minimum value.
        :kwarg float alpha_buoy_crit: parameter for :math:`\alpha_N` smooth limiter
        """
        self.lim_alpha_shear = lim_alpha_shear
        self.lim_alpha_buoy = lim_alpha_buoy
        self.smooth_alpha_buoy_lim = smooth_alpha_buoy_lim
        self.alpha_buoy_crit = alpha_buoy_crit
        # for plotting and such
        self.description = []
        if self.lim_alpha_shear:
            self.description += ['lim', 'shear']
        if self.lim_alpha_buoy:
            self.description += ['lim', 'alpha']
            if self.smooth_alpha_buoy_lim:
                self.description += ['smooth']
                self.description += ['ac'+str(self.alpha_buoy_crit)]
        self.description = ' '.join(self.description)

        # denominator
        self.d0 = 1.0
        self.d1 = 1.0
        self.d2 = 1.0
        self.d3 = 1.0
        self.d4 = 1.0
        self.d5 = 1.0

        # c_mu
        self.n0 = 0.0
        self.n1 = 0.0
        self.n2 = 0.0

        # c_mu_p
        self.nb0 = 0.0
        self.nb1 = 0.0
        self.nb2 = 0.0

    def compute_alpha_shear_steady(self, ri_st, analytical=True):
        r"""
        Compute the steady-state :math:`\alpha_M`.

        Under steady-state conditions, the stability functions satisfy:

        .. math::
            S_m \alpha_M - S_\rho \alpha_M Ri_{st} = 1.0

        (Umlauf and Buchard, 2005, eq A.15) from which :math:`\alpha_M` can be
        solved for given gradient Richardson number :math:`Ri_{st}`.

        :arg float ri_st: Gradient Richardson number
        :kwarg bool analytical: If True (default), solve analytically using the
            coefficients of the stability function. Otherwise, solve
            :math:`\alpha_M` numerically from the equilibrium condition.
        """
        if not analytical:
            # A) solve numerically
            # use equilibrium equation (Umlauf and Buchard, 2005, eq A.15)
            # s_m*a_shear - s_h*a_shear*ri_st = 1.0
            # to solve a_shear at equilibrium
            # NOTE may fail/return incorrect solution for ri_st < -4

            def cost(a_shear):
                a_buoy = ri_st*a_shear
                s_m, s_h = self.eval_funcs(a_buoy, a_shear)
                res = s_m*a_shear - s_h*a_buoy - 1.0
                return res**2
            p = minimize(cost, 1.0)
            assert p.success, 'solving alpha_shear failed, Ri_st={:}'.format(ri_st)
            a_shear = p.x[0]
        else:
            # B) solve analytically
            # compute alpha_shear for equilibrium condition (Umlauf and Buchard, 2005, eq A.19)
            # aM^2 (-d5 + n2 - (d3 - n1 + nb2 )Ri - (d4 + nb1)Ri^2) + aM (-d2 + n0 - (d1 + nb0)Ri) - d0 = 0
            # NOTE this is more robust method
            a = -self.d5 + self.n2 - (self.d3 - self.n1 + self.nb2)*ri_st - (self.d4 + self.nb1)*ri_st**2
            b = -self.d2 + self.n0 - (self.d1 + self.nb0)*ri_st
            c = -self.d0
            a_shear = (-b + np.sqrt(b**2 - 4*a*c))/2/a

        return a_shear

    def compute_c3_minus(self, c1, c2, ri_st):
        r"""
        Compute c3_minus parameter from c1 and c2 parameters.

        c3_minus is solved from equation

        .. math::
            Ri_{st} = \frac{s_m}{s_h} \frac{c2 - c1}{c2 - c3_{-}}

        where :math:`Ri_{st}` is the steady state gradient Richardson number.
        (see Burchard and Bolding, 2001, eq 32)
        """
        a_shear = self.compute_alpha_shear_steady(ri_st, analytical=False)

        # compute aN from Ri_st and aM, Ri_st = aN/aM
        a_buoy = ri_st*a_shear

        # evaluate stability functions for equilibrium conditions
        s_m, s_h = self.eval_funcs(a_buoy, a_shear)

        # compute c3 minus from Umlauf and Burchard (2005)
        c3_minus = c2 - (c2 - c1)*s_m/s_h/ri_st

        # check error in ri_st
        err = s_m/s_h*(c2 - c1)/(c2 - c3_minus) - ri_st
        assert np.abs(err) < 1e-5, 'steady state gradient Richardson number does not match'

        return c3_minus

    def compute_cmu0(self, analytical=True):
        r"""
        Compute parameter :math:`c_\mu^0`

        See: Umlauf and Buchard (2005) eq A.22

        :kwarg bool analytical: If True (default), solve analytically using the
            coefficients of the stability function. Otherwise, solve
            :math:`\alpha_M` numerically from the equilibrium condition.
        """
        a_buoy = 0.0
        if analytical:
            a = self.d5 - self.n2
            b = self.d2 - self.n0
            c = self.d0
            a_shear = (-b - np.sqrt(b**2 - 4*a*c))/2/a
            s_m, s_h = self.eval_funcs(a_buoy, a_shear)
            cm0 = s_m**0.25
        else:

            def cost(a_shear):
                s_m, s_h = self.eval_funcs(a_buoy, a_shear)
                res = s_m*a_shear - 1.0
                return res**2
            p = minimize(cost, 1.0)
            assert p.success, 'solving alpha_shear failed'
            a_shear = p.x[0]
            s_m, s_h = self.eval_funcs(a_buoy, a_shear)
            cm0 = s_m**0.25
        return cm0

    def compute_kappa(self, sigma_psi, cmu0, n, c1, c2):
        r"""
        Computes von Karman constant from the Psi Schmidt number.

        See: Umlauf and Burchard (2003) eq (14)

        :arg sigma_psi: Psi Schmidt number
        :arg cmu0, n, c1, c2: GLS model parameters
        """
        return cmu0 / np.abs(n) * np.sqrt(sigma_psi * (c2 - c1))

    def compute_sigma_psi(self, kappa, cmu0, n, c1, c2):
        r"""
        Computes the Psi Schmidt number from von Karman constant.

        See: Umlauf and Burchard (2003) eq (14)

        :arg kappa: von Karman constant
        :arg cmu0, n, c1, c2: GLS model parameters
        """
        return (n * kappa)**2 / (cmu0**2 * (c2 - c1))

    def compute_length_clim(self, cmu0, ri_st):
        r"""
        Computes the Galpering length scale limit.

        :arg cmu0: parameter :math:`c_\mu^0`
        :arg ri_st: gradient Richardson number
        """
        a_shear = self.compute_alpha_shear_steady(ri_st)

        # compute aN from Ri_st and aM, Ri_st = aN/aM
        a_buoy = ri_st*a_shear

        clim = cmu0**3.0 * np.sqrt(a_buoy/2)
        return clim

    def get_alpha_buoy_min(self):
        r"""
        Compute minimum normalized buoyancy frequency :math:`\alpha_N`

        See: Umlauf and Buchard (2005), Table 3
        """
        # G = epsilon case, this is used in GOTM
        an_min = 0.5*(np.sqrt((self.d1 + self.nb0)**2. - 4.*self.d0*(self.d4 + self.nb1)) - (self.d1 + self.nb0))/(self.d4 + self.nb1)
        # eq. (47)
        # an_min = self.alpha_buoy_min
        return an_min

    def get_alpha_shear_max(self, alpha_buoy, alpha_shear):
        r"""
        Compute maximum normalized shear frequency :math:`\alpha_M`

        from Umlauf and Buchard (2005) eq (44)

        :arg alpha_buoy: normalized buoyancy frequency :math:`\alpha_N`
        :arg alpha_shear: normalized shear frequency :math:`\alpha_M`
        """
        as_max_n = (self.d0*self.n0
                    + (self.d0*self.n1 + self.d1*self.n0)*alpha_buoy
                    + (self.d1*self.n1 + self.d4*self.n0)*alpha_buoy**2
                    + self.d4*self.n1*alpha_buoy**3)
        as_max_d = (self.d2*self.n0
                    + (self.d2*self.n1 + self.d3*self.n0)*alpha_buoy
                    + self.d3*self.n1*alpha_buoy**2)
        return as_max_n/as_max_d

    def get_alpha_buoy_smooth_min(self, alpha_buoy):
        r"""
        Compute smoothed minimum for normalized buoyancy frequency

        See: Burchard and Petersen (1999), eq (19)

        :arg alpha_buoy: normalized buoyancy frequency :math:`\alpha_N`
        """
        return alpha_buoy - (alpha_buoy - self.alpha_buoy_crit)**2/(alpha_buoy + self.get_alpha_buoy_min() - 2*self.alpha_buoy_crit)

    def eval_funcs(self, alpha_buoy, alpha_shear):
        r"""
        Evaluate (unlimited) stability functions

        See: Burchard and Petersen (1999) eqns (30) and (31)

        :arg alpha_buoy: normalized buoyancy frequency :math:`\alpha_N`
        :arg alpha_shear: normalized shear frequency :math:`\alpha_M`
        :returns: :math:`S_m`, :math:`S_\rho`
        """
        den = self.d0 + self.d1*alpha_buoy + self.d2*alpha_shear + self.d3*alpha_buoy*alpha_shear + self.d4*alpha_buoy**2 + self.d5*alpha_shear**2