#pragma once 
#include "sfx_data.h"
#include "sfx_math.cuh"

template<typename T, class sfx_param>
FUCNTION_DECLARATION_SPECIFIER void get_convection_lim(T &zeta_lim, T &Rib_lim, T &f_m_lim, T &f_h_lim,
    const T h0_m, const T h0_t, const T B,
    const sfx_param& param)
{
    T psi_m, psi_h, f_m, f_h, c;

    c = pow(param.Pr_t_inf_inv / param.Pr_t_0_inv, 4);
    zeta_lim = (2.0 * param.alpha_h - c * param.alpha_m - 
                sqrt( (c * param.alpha_m)*(c * param.alpha_m) + 4.0 * c * param.alpha_h * (param.alpha_h - param.alpha_m))) / (2.0 * param.alpha_h*param.alpha_h);

    f_m_lim = pow(1.0 - param.alpha_m * zeta_lim, 0.25);
    f_h_lim = sqrt(1.0 - param.alpha_h * zeta_lim);

    f_m = zeta_lim / h0_m;
    f_h = zeta_lim / h0_t;
    if (fabs(B) < 1.0e-10) f_h = f_m;

    f_m = pow(1.0 - param.alpha_m * f_m, 0.25);
    f_h = sqrt(1.0 - param.alpha_h_fix * f_h);

    psi_m = 2.0 * (atan(f_m_lim) - atan(f_m)) + log((f_m_lim - 1.0) * (f_m + 1.0)/((f_m_lim + 1.0) * (f_m - 1.0)));
    psi_h = log((f_h_lim - 1.0) * (f_h + 1.0)/((f_h_lim + 1.0) * (f_h - 1.0))) / param.Pr_t_0_inv;

    Rib_lim = zeta_lim * psi_h / (psi_m * psi_m);
}

template<typename T, class sfx_param>
FUCNTION_DECLARATION_SPECIFIER void get_psi_stable(T &psi_m, T &psi_h, T &zeta,
    const T Rib, const T h0_m, const T h0_t, const T B,
    const sfx_param& param)
{
    T Rib_coeff, psi0_m, psi0_h, phi, c;
    
    psi0_m = log(h0_m);
    psi0_h = B / psi0_m;

    Rib_coeff = param.beta_m * Rib;
    c = (psi0_h + 1.0) / param.Pr_t_0_inv - 2.0 * Rib_coeff;
    zeta = psi0_m * (sqrt(c*c + 4.0 * Rib_coeff * (1.0 - Rib_coeff)) - c) / (2.0 * param.beta_m * (1.0 - Rib_coeff));

    phi = param.beta_m * zeta;

    psi_m = psi0_m + phi;
    psi_h = (psi0_m + B) / param.Pr_t_0_inv + phi;
}

template<typename T, class sfx_param>
FUCNTION_DECLARATION_SPECIFIER void get_psi_convection(T &psi_m, T &psi_h, T &zeta, 
    const T Rib, const T h0_m, const T h0_t, const T B, 
    const T zeta_conv_lim, const T f_m_conv_lim, const T f_h_conv_lim,
    const sfx_param& param,
    const int maxiters_convection)
{
    T zeta0_m, zeta0_h, f0_m, f0_h, p_m, p_h, a_m, a_h, c_lim, f;

    p_m = 2.0 * atan(f_m_conv_lim) + log((f_m_conv_lim - 1.0) / (f_m_conv_lim + 1.0));
    p_h = log((f_h_conv_lim - 1.0) / (f_h_conv_lim + 1.0));

    zeta = zeta_conv_lim;

    for (int i = 1; i <= maxiters_convection + 1; i++)
    {
        zeta0_m = zeta / h0_m;
        zeta0_h = zeta / h0_t;
        if (fabs(B) < 1.0e-10) 
            zeta0_h = zeta0_m;

        f0_m = pow(1.0 - param.alpha_m * zeta0_m, 0.25);
        f0_h = sqrt(1.0 - param.alpha_h_fix * zeta0_h);

        a_m = -2.0*atan(f0_m) + log((f0_m + 1.0)/(f0_m - 1.0));
        a_h = log((f0_h + 1.0)/(f0_h - 1.0));

        c_lim = pow(zeta_conv_lim / zeta, 1.0 / 3.0);
        f = 3.0 * (1.0 - c_lim);

        psi_m = f / f_m_conv_lim + p_m + a_m;
        psi_h = (f / f_h_conv_lim + p_h + a_h) / param.Pr_t_0_inv;

        if (i == maxiters_convection + 1) 
            break;

        zeta = Rib * psi_m * psi_m / psi_h;
    }
}

template<typename T, class sfx_param>
FUCNTION_DECLARATION_SPECIFIER void get_psi_neutral(T &psi_m, T &psi_h, T &zeta,   
    const T h0_m, const T h0_t, const T B,
    const sfx_param& param)
{
    zeta = 0.0;
    psi_m = log(h0_m);
    psi_h = log(h0_t) / param.Pr_t_0_inv;
    if (fabs(B) < 1.0e-10) 
        psi_h = psi_m / param.Pr_t_0_inv;
}

template<typename T, class sfx_param>
FUCNTION_DECLARATION_SPECIFIER void get_psi_semi_convection(T &psi_m, T &psi_h, T &zeta,
    const T Rib, const T h0_m, const T h0_t, const T B, 
    const sfx_param& param,
    const int maxiters_convection)
{
    T zeta0_m, zeta0_h, f0_m, f0_h, f_m, f_h;

    psi_m = log(h0_m);
    psi_h = log(h0_t);

    if (fabs(B) < 1.0e-10) 
        psi_h = psi_m;

    zeta = Rib * param.Pr_t_0_inv * psi_m * psi_m / psi_h;

    for (int i = 1; i <= maxiters_convection + 1; i++)
    {
        zeta0_m = zeta / h0_m;
        zeta0_h = zeta / h0_t;
        if (fabs(B) < 1.0e-10) 
            zeta0_h = zeta0_m;

        f_m = pow(1.0 - param.alpha_m * zeta, 0.25e0);
        f_h = sqrt(1.0 - param.alpha_h_fix * zeta);

        f0_m = pow(1.0 - param.alpha_m * zeta0_m, 0.25e0);
        f0_h = sqrt(1.0 - param.alpha_h_fix * zeta0_h);

        f0_m = sfx_math::max(f0_m, T(1.000001e0));
        f0_h = sfx_math::max(f0_h, T(1.000001e0));

        psi_m = log((f_m - 1.0e0)*(f0_m + 1.0e0)/((f_m + 1.0e0)*(f0_m - 1.0e0))) + 2.0e0*(atan(f_m) - atan(f0_m));
        psi_h = log((f_h - 1.0e0)*(f0_h + 1.0e0)/((f_h + 1.0e0)*(f0_h - 1.0e0))) / param.Pr_t_0_inv;

        if (i == maxiters_convection + 1) 
            break;

        zeta = Rib * psi_m * psi_m / psi_h;
    }
}

template<typename T, class sfx_param>
FUCNTION_DECLARATION_SPECIFIER void get_psi(T &psi_m, T &psi_h,
    const T zeta,
    const sfx_param& param)
{
    T x_m, x_h;
    T q_m, q_h;

    if (zeta >= 0.0) 
    {
        q_m = pow((1.0 - param.b_m) / param.b_m, 1.0 / 3.0);
        q_h = sqrt(param.c_h * param.c_h - 4.0);

        x_m = pow(1.0 + zeta, 1.0 / 3.0);
        x_h = zeta;

        psi_m = -3.0 * (param.a_m / param.b_m) * (x_m - 1.0) + 0.5 * (param.a_m / param.b_m) * q_m * (2.0 * log((x_m + q_m) / (1.0 + q_m)) - log((x_m * x_m - x_m * q_m + q_m * q_m) / (1.0 - q_m + q_m * q_m)) + 2.0 * sqrt(3.0) * (atan((2.0 * x_m - q_m) / (sqrt(3.0) * q_m)) - atan((2.0 - q_m) / (sqrt(3.0) * q_m))));

        psi_h = -0.5 * param.b_h * log(1.0 + param.c_h * x_h + x_h * x_h) + ((-param.a_h / q_h) + ((param.b_h * param.c_h) / (2.0 * q_h))) * (log((2.0 * x_h + param.c_h - q_h) / (2.0 * x_h + param.c_h + q_h)) - log((param.c_h - q_h) / (param.c_h + q_h)));
    }
    else
    {
        x_m = pow(1.0 - param.alpha_m * zeta, 0.25);
        x_h = pow(1.0 - param.alpha_h * zeta, 0.25); 

        psi_m = (4.0 * atan(1.0) / 2.0) + 2.0 * log(0.5 * (1.0 + x_m)) + log(0.5 * (1.0 + x_m * x_m)) - 2.0 * atan(x_m);
        psi_h = 2.0 * log(0.5 * (1.0 + x_h * x_h));
    }
}

template<typename T, class sfx_param>
FUCNTION_DECLARATION_SPECIFIER void get_psi_mh(T &psi_m, T &psi_h,
    const T zeta_m, const T zeta_h,
    const sfx_param& param)
{
    T x_m, x_h;
    T q_m, q_h;

    if (zeta_m >= 0.0) 
    {
        q_m = pow((1.0 - param.b_m) / param.b_m, 1.0 / 3.0);
        x_m = pow(1.0 + zeta_m, 1.0 / 3.0);

        psi_m = -3.0 * (param.a_m / param.b_m) * (x_m - 1.0) + 0.5 * (param.a_m / param.b_m) * q_m * (2.0 * log((x_m + q_m) / (1.0 + q_m)) - log((x_m * x_m - x_m * q_m + q_m * q_m) / (1.0 - q_m + q_m * q_m)) + 2.0 * sqrt(3.0) * (atan((2.0 * x_m - q_m) / (sqrt(3.0) * q_m)) - atan((2.0 - q_m) / (sqrt(3.0) * q_m))));
    }                                
    else
    {    
        x_m = pow(1.0 - param.alpha_m * zeta_m, 0.25);
        psi_m = (4.0 * atan(1.0) / 2.0) + 2.0 * log(0.5 * (1.0 + x_m)) + log(0.5 * (1.0 + x_m * x_m)) - 2.0 * atan(x_m);
    }

    if (zeta_h >= 0.0)
    {    
        q_h = sqrt(param.c_h * param.c_h - 4.0);
        x_h = zeta_h;

        psi_h = -0.5 * param.b_h * log(1.0 + param.c_h * x_h + x_h * x_h) + ((-param.a_h / q_h) + ((param.b_h * param.c_h) / (2.0 * q_h))) * (log((2.0 * x_h + param.c_h - q_h) / (2.0 * x_h + param.c_h + q_h)) - log((param.c_h - q_h) / (param.c_h + q_h)));
    }
    else
    {
        x_h = pow(1.0 - param.alpha_h * zeta_h, 0.25);
        psi_h = 2.0 * log(0.5 * (1.0 + x_h * x_h));
    }
}


template<typename T, class sfx_param>
FUCNTION_DECLARATION_SPECIFIER void get_dynamic_scales(T &Udyn, T &Tdyn, T &Qdyn, T& zeta,
    const T U, const T Tsemi, const T dT, const T dQ, 
    const T z, const T z0_m, const T z0_t,
    const T beta,
    const sfx_param& param,
    const int maxiters)
{
    const T gamma = T(0.61);
    T psi_m, psi_h, psi0_m, psi0_h, Linv;

    Udyn = param.kappa * U / log(z / z0_m);
    Tdyn = param.kappa * dT * param.Pr_t_0_inv / log(z / z0_t);
    Qdyn = param.kappa * dQ * param.Pr_t_0_inv / log(z / z0_t);
    zeta = 0.0;

    // --- no wind
    if (Udyn < 1e-5) 
        return;

    Linv = param.kappa * beta * (Tdyn + gamma * Qdyn * Tsemi) / (Udyn * Udyn);
    zeta = z * Linv;

    // --- near neutral case
    if (Linv < 1e-5) 
        return;

    for (int i = 0; i < maxiters; i++)
    {
        get_psi(psi_m, psi_h, zeta, param);
        get_psi_mh(psi0_m, psi0_h, z0_m * Linv, z0_t * Linv, 
            param);

        Udyn = param.kappa * U / (log(z / z0_m) - (psi_m - psi0_m));
        Tdyn = param.kappa * dT * param.Pr_t_0_inv / (log(z / z0_t) - (psi_h - psi0_h));
        Qdyn = param.kappa * dQ * param.Pr_t_0_inv / (log(z / z0_t) - (psi_h - psi0_h));

        if (Udyn < 1e-5) 
            break;

        Linv = param.kappa * beta * (Tdyn + gamma * Qdyn * Tsemi) / (Udyn * Udyn);
        zeta = z * Linv;
    }
}

template<typename T, class sfx_param>
FUCNTION_DECLARATION_SPECIFIER void get_phi(T &phi_m, T &phi_h,
    const T zeta, 
    const sfx_param& param)
{
    if (zeta >= 0.0) 
    {
        phi_m = 1.0 + (param.a_m * zeta * pow(1.0 + zeta, 1.0 / 3.0) ) / (1.0 + param.b_m * zeta);
        phi_h = 1.0 + (param.a_h * zeta + param.b_h * zeta * zeta) / (1.0 + param.c_h * zeta + zeta * zeta);
    }
    else
    {
        phi_m = pow(1.0 - param.alpha_m * zeta, -0.25);
        phi_h = pow(1.0 - param.alpha_h * zeta, -0.5);
    }
}