#pragma once
#include "sfx_template_parameters.h"
#include <cstddef>

template<typename T, MemType memIn, MemType memOut, MemType RunMem >
class Flux
{
private:
    T *U, *dT, *Tsemi, *dQ, *h, *in_z0_m;
    T *zeta, *Rib, *Re, *B, *z0_m, *z0_t, *Rib_conv_lim, *Cm, *Ct, *Km, *Pr_t_inv;

    T 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;

    int grid_size;
    bool ifAllocated;
    size_t allocated_size;
public:
    Flux();
    void set_params(const int grid_size, 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);
    ~Flux();

    void compute_flux(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_,
    T *U_, T *dT_, T *Tsemi_, T *dQ_, T *h_, T *in_z0_m_, 
    const int maxiters_charnock, const int maxiters_convection);

private:
    void set_data( 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_,
    T *U_, T *dT_, T *Tsemi_, T *dQ_, T *h_, T *in_z0_m_);
};