#include <iostream>
#include "../includeCU/sfx_compute_esm.cuh"
#include "../includeCU/sfx_surface.cuh"

#define gpuErrchk(ans) { gpuAssert((ans), __FILE__, __LINE__); }
inline void gpuAssert(cudaError_t code, const char *file, int line, bool abort=true)
{
   if (code != cudaSuccess) 
   {
      fprintf(stderr,"GPUassert: %s %s %d\n", cudaGetErrorString(code), file, line);
      if (abort) exit(code);
   }
}

template<typename T>
__device__ void get_charnock_roughness(T &z0_m, T &u_dyn0,
    const T h, const T U,
    const T kappa, 
    const T h_charnock, const T c1_charnock, const T c2_charnock, 
    const int maxiters)
{
    T Uc, a, b, c, c_min, f;

    Uc = U;
    a = 0.0;
    b = 25.0;
    c_min = log(h_charnock) / kappa;

    for (int i = 0; i < maxiters; i++)
    {
        f = c1_charnock - 2.0 * log(Uc);
        for (int j = 0; j < maxiters; j++)
        {
            c = (f + 2.0 * log(b)) / kappa;
            if (U <= 8.0e0) 
                a = log(1.0 + c2_charnock * ( pow(b / Uc, 3) ) ) / kappa;
            c = max(c - a, c_min);
            b = c;
        }
        z0_m = h_charnock * exp(-c * kappa);
        z0_m = max(z0_m, T(0.000015e0));
        Uc = U * log(h_charnock / z0_m) / log(h / z0_m);
    }
    
    u_dyn0 = Uc / c;
}

template __device__ void get_charnock_roughness(float &z0_m, float &u_dyn0,
    const float h, const float U,
    const float kappa, 
    const float h_charnock, const float c1_charnock, const float c2_charnock, 
    const int maxiters);
template __device__ void get_charnock_roughness(double &z0_m, double &u_dyn0, 
    const double h, const double U,
    const double kappa, 
    const double h_charnock, const double c1_charnock, const double c2_charnock,
    const int maxiters);
    
    
template<typename T>
__device__ 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 T Pr_t_inf_inv, const T Pr_t_0_inv,
    const T alpha_h, const T alpha_m, const T alpha_h_fix)
{
    T psi_m, psi_h, f_m, f_h, c;

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

    f_m_lim = pow(1.0 - alpha_m * zeta_lim, 0.25);
    f_h_lim = sqrt(1.0 - 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 - alpha_m * f_m, 0.25);
    f_h = sqrt(1.0 - 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))) / Pr_t_0_inv;

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

template __device__ void get_convection_lim(float &zeta_lim, float &Rib_lim, float &f_m_lim, float &f_h_lim,
    const float h0_m, const float h0_t, const float B,
    const float Pr_t_inf_inv, const float Pr_t_0_inv,
    const float alpha_h, const float alpha_m, const float alpha_h_fix);
template __device__ void get_convection_lim(double &zeta_lim, double &Rib_lim, double &f_m_lim, double &f_h_lim,
    const double h0_m, const double h0_t, const double B,
    const double Pr_t_inf_inv, const double Pr_t_0_inv,
    const double alpha_h, const double alpha_m, const double alpha_h_fix);

template<typename T>
void __device__ 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 T Pr_t_0_inv, const T beta_m)
{
    T Rib_coeff, psi0_m, psi0_h, phi, c;
    
    psi0_m = log(h0_m);
    psi0_h = B / psi0_m;

    Rib_coeff = beta_m * Rib;
    c = (psi0_h + 1.0) / 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 * beta_m * (1.0 - Rib_coeff));

    phi = beta_m * zeta;

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

template __device__ void get_psi_stable(float &psi_m, float &psi_h, float &zeta,
    const float Rib, const float h0_m, const float h0_t, const float B,
    const float Pr_t_0_inv, const float beta_m);
template __device__ void get_psi_stable(double &psi_m, double &psi_h, double &zeta,
    const double Rib, const double h0_m, const double h0_t, const double B,
    const double Pr_t_0_inv, const double beta_m);

template<typename T>
void __device__ 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 T Pr_t_0_inv,
    const T alpha_h, const T alpha_m, const T alpha_h_fix,
    const int maxiters)
{
    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 + 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 - alpha_m * zeta0_m, 0.25);
        f0_h = sqrt(1.0 - 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) / Pr_t_0_inv;

        if (i == maxiters + 1) 
            break;

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

template __device__ void get_psi_convection(float &psi_m, float &psi_h, float &zeta, 
    const float Rib, const float h0_m, const float h0_t, const float B, 
    const float zeta_conv_lim, const float f_m_conv_lim, const float f_h_conv_lim,
    const float Pr_t_0_inv,
    const float alpha_h, const float alpha_m, const float alpha_h_fix,
    const int maxiters);
template __device__ void get_psi_convection(double &psi_m, double &psi_h, double &zeta, 
    const double Rib, const double h0_m, const double h0_t, const double B, 
    const double zeta_conv_lim, const double f_m_conv_lim, const double f_h_conv_lim,
    const double Pr_t_0_inv,
    const double alpha_h, const double alpha_m, const double alpha_h_fix,
    const int maxiters);

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

template __device__ void get_psi_neutral(float &psi_m, float &psi_h, float &zeta,
    const float h0_m, const float h0_t, const float B,
    const float Pr_t_0_inv);
template __device__ void get_psi_neutral(double &psi_m, double &psi_h, double &zeta,
    const double h0_m, const double h0_t, const double B,
    const double Pr_t_0_inv);

template<typename T>
void __device__ 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 T Pr_t_0_inv,
    const T alpha_m, const T alpha_h_fix,
    const int maxiters)
{
    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 * Pr_t_0_inv * psi_m * psi_m / psi_h;

    for (int i = 1; i <= maxiters + 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 - alpha_m * zeta, 0.25e0);
        f_h = sqrt(1.0 - alpha_h_fix * zeta);

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

        f0_m = max(f0_m, T(1.000001e0));
        f0_h = 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))) / Pr_t_0_inv;

        if (i == maxiters + 1) 
            break;

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

template __device__ void get_psi_semi_convection(float &psi_m, float &psi_h, float &zeta,
    const float Rib, const float h0_m, const float h0_t, const float B, 
    const float Pr_t_0_inv,
    const float alpha_m, const float alpha_h_fix,
    const int maxiters);
template __device__ void get_psi_semi_convection(double &psi_m, double &psi_h, double &zeta,
    const double Rib, const double h0_m, const double h0_t, const double B, 
    const double Pr_t_0_inv,
    const double alpha_m, const double alpha_h_fix,
    const int maxiters);

template<typename T>
__global__ void kernel_compute_flux_esm_gpu(T *zeta_, T *Rib_, T *Re_, T *B_, T *z0_m_, T *z0_t_, T *Rib_conv_lim_, T *Cm_, T *Ct_, T *Km_, T *Pr_t_inv_,
    const T *U_, const T *dT_, const T *Tsemi_, const T *dQ_, const T *h_, const T *in_z0_m_,
    const T kappa, const T Pr_t_0_inv, const T Pr_t_inf_inv, 
    const T alpha_m, const T alpha_h, const T alpha_h_fix, 
    const T beta_m, const T beta_h, const T Rib_max, const T Re_rough_min, 
    const T B1_rough, const T B2_rough,
    const T B_max_land, const T B_max_ocean, const T B_max_lake,
    const T gamma_c, const T Re_visc_min,
    const T Pr_m, const T nu_air, const T g, 
    const int maxiters_charnock, const int maxiters_convection, 
    const int grid_size)
{
    const int index = blockIdx.x * blockDim.x + threadIdx.x;

    T h, U, dT, Tsemi, dQ, z0_m;
    T Re, z0_t, B, h0_m, h0_t, u_dyn0, zeta, Rib, zeta_conv_lim, Rib_conv_lim, f_m_conv_lim, f_h_conv_lim, psi_m, psi_h, phi_m, phi_h, Km, Pr_t_inv, Cm, Ct;
    int surface_type;
    T fval;

    const T B3_rough = kappa * Pr_m, B4_rough =( 0.14 * ( pow(30.0, B2_rough) ) ) * (pow(Pr_m, 0.8));
    const T h_charnock = 10.0, c1_charnock = log(h_charnock * (g / gamma_c)), c2_charnock = Re_visc_min * nu_air * c1_charnock;

    if(index < grid_size)
    {
        U = U_[index];
        Tsemi = Tsemi_[index];
        dT = dT_[index];
        dQ = dQ_[index];
        h = h_[index];
        z0_m = in_z0_m_[index];

        if (z0_m < 0.0) surface_type = 0;
        else            surface_type = 1;

        if (surface_type == 0)
        {
            get_charnock_roughness(z0_m, u_dyn0, h, U, kappa, h_charnock, c1_charnock, c2_charnock, maxiters_charnock);
            h0_m = h / z0_m;
        }
        if (surface_type == 1) 
        {
            h0_m = h / z0_m;
            u_dyn0 = U * kappa / log(h0_m);
        }

        Re = u_dyn0 * z0_m / nu_air;

        if(Re <= Re_rough_min) B = B1_rough * log(B3_rough * Re) + B2_rough;
        else                   B = B4_rough * (pow(Re, B2_rough));

        if (surface_type == 0)  B = min(B, B_max_ocean);
        if (surface_type == 1)   B = min(B, B_max_land);
        if (surface_type == 2)   B = min(B, B_max_lake);

        z0_t = z0_m / exp(B);
        h0_t = h / z0_t;
        Rib = (g / Tsemi) * h * (dT + 0.61e0 * Tsemi * dQ) / (U*U);

        get_convection_lim(zeta_conv_lim, Rib_conv_lim, f_m_conv_lim, f_h_conv_lim, h0_m, h0_t, B, Pr_t_inf_inv, Pr_t_0_inv, alpha_h, alpha_m, alpha_h_fix);


        if (Rib > 0.0) 
        {
            Rib = min(Rib, Rib_max);
            get_psi_stable(psi_m, psi_h, zeta, Rib, h0_m, h0_t, B, Pr_t_0_inv, beta_m);

            fval = beta_m * zeta;
            phi_m = 1.0 + fval;
            phi_h = 1.0/Pr_t_0_inv + fval;
        }

        else if (Rib < Rib_conv_lim) 
        {
            get_psi_convection(psi_m, psi_h, zeta, Rib, h0_m, h0_t, B, zeta_conv_lim, f_m_conv_lim, f_h_conv_lim, Pr_t_0_inv, alpha_h, alpha_m, alpha_h_fix, maxiters_convection);

            fval = pow(zeta_conv_lim / zeta, 1.0/3.0);
            phi_m = fval / f_m_conv_lim;
            phi_h = fval / (Pr_t_0_inv * f_h_conv_lim);
        }
        else if (Rib > -0.001) 
        {
            get_psi_neutral(psi_m, psi_h, zeta, h0_m, h0_t, B, Pr_t_0_inv);
        
            phi_m = 1.0;
            phi_h = 1.0 / Pr_t_0_inv;
        }
        else
        {
            get_psi_semi_convection(psi_m, psi_h, zeta, Rib, h0_m, h0_t, B, Pr_t_0_inv, alpha_m, alpha_h_fix, maxiters_convection);
            
            phi_m = pow(1.0 - alpha_m * zeta, -0.25);
            phi_h = 1.0 / (Pr_t_0_inv * sqrt(1.0 - alpha_h_fix * zeta));
        }

        Cm = kappa / psi_m;
        Ct = kappa / psi_h;

        Km = kappa * Cm * U * h / phi_m;
        Pr_t_inv = phi_m / phi_h;

        zeta_[index]         = zeta;
        Rib_[index]          = Rib;
        Re_[index]           = Re;
        B_[index]            = B;
        z0_m_[index]         = z0_m;
        z0_t_[index]         = z0_t;
        Rib_conv_lim_[index] = Rib_conv_lim;
        Cm_[index]           = Cm;
        Ct_[index]           = Ct;
        Km_[index]           = Km;
        Pr_t_inv_[index]     = Pr_t_inv;
    }
}

template __global__ void kernel_compute_flux_esm_gpu(float *zeta_, float *Rib_, float *Re_, float *B_, float *z0_m_, float *z0_t_, float *Rib_conv_lim_, float *Cm_, float *Ct_, float *Km_, float *Pr_t_inv_,
    const float *U, const float *dt, const float *T_semi, const float *dq, const float *H, const float *in_z0_m,
    const float kappa, const float Pr_t_0_inv, const float Pr_t_inf_inv, 
    const float alpha_m, const float alpha_h, const float alpha_h_fix, 
    const float beta_m, const float beta_h, const float Rib_max, const float Re_rough_min, 
    const float B1_rough, const float B2_rough,
    const float B_max_land, const float B_max_ocean, const float B_max_lake,
    const float gamma_c, const float Re_visc_min,
    const float Pr_m, const float nu_air, const float g, 
    const int maxiters_charnock, const int maxiters_convection, 
    const int grid_size);
template __global__ void kernel_compute_flux_esm_gpu(double *zeta_, double *Rib_, double *Re_, double *B_, double *z0_m_, double *z0_t_, double *Rib_conv_lim_, double *Cm_, double *Ct_, double *Km_, double *Pr_t_inv_,
    const double *U, const double *dt, const double *T_semi, const double *dq, const double *H, const double *in_z0_m, 
    const double kappa, const double Pr_t_0_inv, const double Pr_t_inf_inv, 
    const double alpha_m, const double alpha_h, const double alpha_h_fix, 
    const double beta_m, const double beta_h, const double Rib_max, const double Re_rough_min, 
    const double B1_rough, const double B2_rough,
    const double B_max_land, const double B_max_ocean, const double B_max_lake,
    const double gamma_c, const double Re_visc_min,
    const double Pr_m, const double nu_air, const double g, 
    const int maxiters_charnock, const int maxiters_convection, 
    const int grid_size);

template<typename T>
void compute_flux_esm_gpu(T *zeta_, T *Rib_, T *Re_, T *B_, T *z0_m_, T *z0_t_, T *Rib_conv_lim_, T *Cm_, T *Ct_, T *Km_, T *Pr_t_inv_,
    const T *U_, const T *dT_, const T *Tsemi_, const T *dQ_, const T *h_, const T *in_z0_m_,
    const T kappa, const T Pr_t_0_inv, const T Pr_t_inf_inv, 
    const T alpha_m, const T alpha_h, const T alpha_h_fix, 
    const T beta_m, const T beta_h, const T Rib_max, const T Re_rough_min, 
    const T B1_rough, const T B2_rough,
    const T B_max_land, const T B_max_ocean, const T B_max_lake,
    const T gamma_c, const T Re_visc_min,
    const T Pr_m, const T nu_air, const T g, 
    const int maxiters_charnock, const int maxiters_convection, 
    const int grid_size)
{
    const int BlockCount = int(ceil(float(grid_size) / 1024.0));
    dim3 cuBlock = dim3(1024, 1, 1);
	dim3 cuGrid = dim3(BlockCount, 1, 1);

    kernel_compute_flux_esm_gpu<<<cuGrid, cuBlock>>>(zeta_, Rib_, Re_, B_, z0_m_, z0_t_, Rib_conv_lim_, Cm_, Ct_, Km_, Pr_t_inv_,
    U_, dT_, Tsemi_, dQ_, h_, in_z0_m_,
    kappa, Pr_t_0_inv, Pr_t_inf_inv, 
    alpha_m, alpha_h, alpha_h_fix, 
    beta_m, beta_h, Rib_max, Re_rough_min, 
    B1_rough, B2_rough,
    B_max_land, B_max_ocean, B_max_lake,
    gamma_c, Re_visc_min,
    Pr_m, nu_air, g, 
    maxiters_charnock, maxiters_convection, 
    grid_size);

    gpuErrchk( cudaPeekAtLastError() );
}

template void compute_flux_esm_gpu(float *zeta_, float *Rib_, float *Re_, float *B_, float *z0_m_, float *z0_t_, float *Rib_conv_lim_, float *Cm_, float *Ct_, float *Km_, float *Pr_t_inv_,
    const float *U, const float *dt, const float *T_semi, const float *dq, const float *H, const float *in_z0_m,
    const float kappa, const float Pr_t_0_inv, const float Pr_t_inf_inv, 
    const float alpha_m, const float alpha_h, const float alpha_h_fix, 
    const float beta_m, const float beta_h, const float Rib_max, const float Re_rough_min, 
    const float B1_rough, const float B2_rough,
    const float B_max_land, const float B_max_ocean, const float B_max_lake,
    const float gamma_c, const float Re_visc_min,
    const float Pr_m, const float nu_air, const float g, 
    const int maxiters_charnock, const int maxiters_convection, 
    const int grid_size);
template void compute_flux_esm_gpu(double *zeta_, double *Rib_, double *Re_, double *B_, double *z0_m_, double *z0_t_, double *Rib_conv_lim_, double *Cm_, double *Ct_, double *Km_, double *Pr_t_inv_,
    const double *U, const double *dt, const double *T_semi, const double *dq, const double *H, const double *in_z0_m, 
    const double kappa, const double Pr_t_0_inv, const double Pr_t_inf_inv, 
    const double alpha_m, const double alpha_h, const double alpha_h_fix, 
    const double beta_m, const double beta_h, const double Rib_max, const double Re_rough_min, 
    const double B1_rough, const double B2_rough,
    const double B_max_land, const double B_max_ocean, const double B_max_lake,
    const double gamma_c, const double Re_visc_min,
    const double Pr_m, const double nu_air, const double g, 
    const int maxiters_charnock, const int maxiters_convection, 
    const int grid_size);