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

template<typename T>
FUCNTION_DECLARATION_SPECIFIER void get_thermal_roughness(T &z0_t, T &B,
    const T z0_m, const T Re, 
    const struct sfx_surface_param& param,
    const int surface_type)
{
    // --- define B = log(z0_m / z0_t)
    B = (Re <= param.Re_rough_min) ? param.B1_rough * log(param.B3_rough * Re) + param.B2_rough : 
                                                 param.B4_rough * (pow(Re, param.B2_rough));

    //   --- apply max restriction based on surface type
    if (surface_type == param.surface_ocean)  B = sfx_math::min(B, param.B_max_ocean);
    if (surface_type == param.surface_lake)   B = sfx_math::min(B, param.B_max_lake);
    if (surface_type == param.surface_land)   B = sfx_math::min(B, param.B_max_land);

    // --- define roughness [thermal]
    z0_t = z0_m / exp(B);
}

template<typename T>
FUCNTION_DECLARATION_SPECIFIER void get_charnock_roughness(T &z0_m, T &u_dyn0,
    const T U, const T h,
    const sfx_surface_param& surface_param,
    const int maxiters)
{
    T Uc, a, b, c, c_min, f;

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

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