#pragma once
#include "sfx-template-parameters.h"
#include "model-base.h"
#include "sfx-data.h"

template<typename T, MemType memIn, MemType memOut, MemType RunMem >
class FluxShebaBase : public ModelBase<T, memIn, memOut, RunMem>
{
public:
    using ModelBase<T, memIn, memOut, RunMem>::res_sfx;
    using ModelBase<T, memIn, memOut, RunMem>::sfx;
    using ModelBase<T, memIn, memOut, RunMem>::meteo;
    using ModelBase<T, memIn, memOut, RunMem>::grid_size;
    using ModelBase<T, memIn, memOut, RunMem>::ifAllocated;
    using ModelBase<T, memIn, memOut, RunMem>::allocated_size;

    sfx_surface_param surface;
    sfx_phys_constants phys;
    sfx_sheba_param_C model;
    sfx_sheba_numericsType_C numerics;

    FluxShebaBase(sfxDataVecTypeC* sfx,
                meteoDataVecTypeC* meteo,
                const sfx_sheba_param_C model, 
                const sfx_surface_param surface,
                const sfx_sheba_numericsType_C numerics,
                const sfx_phys_constants phys,
                const int grid_size);
    ~FluxShebaBase();
};

template<typename T, MemType memIn, MemType memOut, MemType RunMem >
class FluxSheba : public FluxShebaBase<T, memIn, memOut, RunMem>
{};

template<typename T, MemType memIn, MemType memOut >
class FluxSheba<T, memIn, memOut, MemType::CPU> : public FluxShebaBase<T, memIn, memOut, MemType::CPU>
{
    using FluxShebaBase<T, memIn, memOut, MemType::CPU>::res_sfx;
    using FluxShebaBase<T, memIn, memOut, MemType::CPU>::sfx;
    using FluxShebaBase<T, memIn, memOut, MemType::CPU>::meteo;
    using FluxShebaBase<T, memIn, memOut, MemType::CPU>::surface;
    using FluxShebaBase<T, memIn, memOut, MemType::CPU>::phys;
    using FluxShebaBase<T, memIn, memOut, MemType::CPU>::grid_size;
    using FluxShebaBase<T, memIn, memOut, MemType::CPU>::ifAllocated;
    using FluxShebaBase<T, memIn, memOut, MemType::CPU>::allocated_size;
    using FluxShebaBase<T, memIn, memOut, MemType::CPU>::model;
    using FluxShebaBase<T, memIn, memOut, MemType::CPU>::numerics;
public:
    FluxSheba(sfxDataVecTypeC* sfx,
                meteoDataVecTypeC* meteo,
                const sfx_sheba_param_C model, 
                const sfx_surface_param surface,
                const sfx_sheba_numericsType_C numerics,
                const sfx_phys_constants phys,
                const int grid_size) : FluxShebaBase<T, memIn, memOut, MemType::CPU>(sfx, meteo, model, 
                                       surface, numerics, phys, grid_size) {}
    ~FluxSheba() = default;
    void compute_flux();
};

#ifdef INCLUDE_CUDA
template<typename T, MemType memIn, MemType memOut >
class FluxSheba<T, memIn, memOut, MemType::GPU> : public FluxShebaBase<T, memIn, memOut, MemType::GPU>
{
    using FluxShebaBase<T, memIn, memOut, MemType::GPU>::res_sfx;
    using FluxShebaBase<T, memIn, memOut, MemType::GPU>::sfx;
    using FluxShebaBase<T, memIn, memOut, MemType::GPU>::meteo;
    using FluxShebaBase<T, memIn, memOut, MemType::GPU>::surface;
    using FluxShebaBase<T, memIn, memOut, MemType::GPU>::phys;
    using FluxShebaBase<T, memIn, memOut, MemType::GPU>::grid_size;
    using FluxShebaBase<T, memIn, memOut, MemType::GPU>::ifAllocated;
    using FluxShebaBase<T, memIn, memOut, MemType::GPU>::allocated_size;
    using FluxShebaBase<T, memIn, memOut, MemType::GPU>::model;
    using FluxShebaBase<T, memIn, memOut, MemType::GPU>::numerics;
public:
    FluxSheba(sfxDataVecTypeC* sfx,
                meteoDataVecTypeC* meteo,
                const sfx_sheba_param_C model, 
                const sfx_surface_param surface,
                const sfx_sheba_numericsType_C numerics,
                const sfx_phys_constants phys,
                const int grid_size) : FluxShebaBase<T, memIn, memOut, MemType::GPU>(sfx, meteo, model, 
                                       surface, numerics, phys, grid_size) {}
    ~FluxSheba() = default;
    void compute_flux();
};
#endif