#ifndef DUNE_DORIE_PARAM_TRANSPORT_HH
#define DUNE_DORIE_PARAM_TRANSPORT_HH

#include <cmath>
#include <vector>
#include <unordered_map>
#include <map>
#include <utility>
#include <functional>
#include <memory>

#include <yaml-cpp/yaml.h>

#include <dune/common/exceptions.hh>
#include <dune/grid/common/mcmgmapper.hh>

#include <dune/dorie/common/parameterization_factory.hh>
#include <dune/dorie/model/transport/parameterization/interface.hh>
#include <dune/dorie/model/transport/parameterization/hydrodynamic_dispersion/interface.hh>
#include <dune/dorie/model/transport/parameterization/hydrodynamic_dispersion/const.hh>
#include <dune/dorie/model/transport/parameterization/hydrodynamic_dispersion/power_law.hh>
#include <dune/dorie/model/transport/parameterization/hydrodynamic_dispersion/eff_hydromechanic_dispersion/interface.hh>
#include <dune/dorie/model/transport/parameterization/hydrodynamic_dispersion/eff_hydromechanic_dispersion/const.hh>
#include <dune/dorie/model/transport/parameterization/hydrodynamic_dispersion/eff_hydromechanic_dispersion/isotropic.hh>
#include <dune/dorie/model/transport/parameterization/hydrodynamic_dispersion/eff_hydromechanic_dispersion/longitudinal/interface.hh>
#include <dune/dorie/model/transport/parameterization/hydrodynamic_dispersion/eff_hydromechanic_dispersion/longitudinal/const.hh>
#include <dune/dorie/model/transport/parameterization/hydrodynamic_dispersion/eff_hydromechanic_dispersion/transverse/interface.hh>
#include <dune/dorie/model/transport/parameterization/hydrodynamic_dispersion/eff_hydromechanic_dispersion/transverse/const.hh>
#include <dune/dorie/model/transport/parameterization/hydrodynamic_dispersion/eff_diffusion/interface.hh>
#include <dune/dorie/model/transport/parameterization/hydrodynamic_dispersion/eff_diffusion/const.hh>
#include <dune/dorie/model/transport/parameterization/hydrodynamic_dispersion/eff_diffusion/millington_quirk_1.hh>
#include <dune/dorie/model/transport/parameterization/hydrodynamic_dispersion/eff_diffusion/millington_quirk_2.hh>
#include <dune/dorie/model/transport/parameterization/hydrodynamic_dispersion/superposition.hh>



namespace Dune {
namespace Dorie {


template<class Traits>
struct TransportParameterizationFactory 
    : public ParameterizationFactory<Dorie::Parameterization::Transport<Traits>>
{
  /// Return a default-initialized parameterization object
  /** \param type The type indicating the parameterization implementation
   *  \param name The name of the parameterization object (layer type)
   */
  std::shared_ptr<Dorie::Parameterization::Transport<Traits>> 
  selector(
      const YAML::Node& type_node,
      const std::string name) const override
  {
    std::shared_ptr<Parameterization::Transport<Traits>> hd_dips;

    const auto type = type_node["type"].as<std::string>();

    if (type == Parameterization::ConstHydrodynamicDispersion<Traits>::type) {
        hd_dips = std::make_shared<Parameterization::ConstHydrodynamicDispersion<Traits>>(name);
    } else if (type == Parameterization::PowerLawDispersion<Traits>::type) {
        hd_dips = std::make_shared<Parameterization::PowerLawDispersion<Traits>>(name);
    } else if (type == Parameterization::HydrodynamicDispersionSuperposition<Traits>::type) {

      // Select type for effective diffusion
      const YAML::Node& eff_diff_node = type_node["eff_diff"];
      auto eff_diff_type = eff_diff_node["type"].as<std::string>();
      std::shared_ptr<Parameterization::EffectiveDiffusion<Traits>> eff_diff;
  
      // Build parametrization for effective diffusion
      if (eff_diff_type == Parameterization::ConstEffectiveDiffusion<Traits>::type) {
          eff_diff = std::make_shared<Parameterization::ConstEffectiveDiffusion<Traits>>(name);
      } else if (eff_diff_type == Parameterization::MillingtonQuirk1<Traits>::type) {
          eff_diff = std::make_shared<Parameterization::MillingtonQuirk1<Traits>>(name);
      } else if (eff_diff_type == Parameterization::MillingtonQuirk2<Traits>::type) {
          eff_diff = std::make_shared<Parameterization::MillingtonQuirk2<Traits>>(name);
      } else {
          DUNE_THROW(NotImplemented,"Invalid eff_diff");
      }

      // Select type for effective hydromechanic dispersion
      const YAML::Node& eff_hm_disp_node = type_node["eff_hydromechanic_disp"];
      auto eff_hm_disp_type = eff_hm_disp_node["type"].as<std::string>();
      std::shared_ptr<Parameterization::EffectiveHydromechanicDispersion<Traits>> eff_hm_disp;

      // Build parametrization for effective hydromechanic dispersion
      if (eff_hm_disp_type == Parameterization::ConstEffectiveHydromechanicDispersion<Traits>::type) {
          eff_hm_disp = std::make_shared<Parameterization::ConstEffectiveHydromechanicDispersion<Traits>>(name);
      } else if (eff_hm_disp_type == Parameterization::IsotropicEffectiveHydromechanicDispersion<Traits>::type) {
        
        // Select type for longitudinal and transverse dispersivity
        const YAML::Node& long_disp_node = eff_hm_disp_node["long_disp"];
        const YAML::Node& trans_disp_node = eff_hm_disp_node["trans_disp"];
        auto long_disp_type = long_disp_node["type"].as<std::string>();
        auto trans_disp_type = trans_disp_node["type"].as<std::string>();
        std::shared_ptr<Parameterization::ConstLongitudinalDispersivity<Traits>> lambda_l;
        std::shared_ptr<Parameterization::ConstTransverseDispersivity<Traits>> lambda_t;

        // Build parametrization for longitudinal dispersivity
        if (long_disp_type == Parameterization::ConstLongitudinalDispersivity<Traits>::type) {
            lambda_l = std::make_shared<Parameterization::ConstLongitudinalDispersivity<Traits>>(name);
        } else {
            DUNE_THROW(NotImplemented,"Invalid long_disp");
        }

        // Build parametrization for transverse dispersivity
        if (trans_disp_type == Parameterization::ConstTransverseDispersivity<Traits>::type) {
            lambda_t = std::make_shared<Parameterization::ConstTransverseDispersivity<Traits>>(name);
        } else {
            DUNE_THROW(NotImplemented,"Invalid trans_disp");
        }

        eff_hm_disp = std::make_shared<Parameterization::IsotropicEffectiveHydromechanicDispersion<Traits>>(name,lambda_l,lambda_t);
      } else {
          DUNE_THROW(NotImplemented,"Invalid eff_hydromech_disp");
      }

      hd_dips = std::make_shared<Parameterization::HydrodynamicDispersionSuperposition<Traits>>(name,eff_diff,eff_hm_disp);
    } else {
            DUNE_THROW(NotImplemented,"Invalid hydrodyn_disp" );
    }
    return hd_dips;
  }
};

/// Top-level parameterization interface for the local operator.
/** This class loads and stores the parameters, maps them to grid entities,
 *  and provides functions to access the parameterization. For doing so, the
 *  FlowParameters::bind function has to be called first to cache the
 *  parameters of the given grid entity. Parameters are only stored for entities
 *  of codim 0.
 *
 *  The parameterization consists of three structures:
 *      - This top-level class serves as interface and storage. It also casts to
 *          data types used by the DUNE operators
 *      - The Dorie::Parameterization::Transport interface that provides strong
 *          data types for the base parameters and purely virtual functions that
 *          need to be defined by derived parameterizations
 *      - The actual parameterization defining all parameters and functions
 *
 *  \author Lukas Riedel
 *  \author Santiago Ospina
 *  \date 2018-2019
 *  \ingroup RichardsParam
 */
template <class Traits>
class SoluteParameters
{
private:
    using Grid = typename Traits::Grid;
    using LevelGridView = typename Grid::LevelGridView;
    using Mapper = typename Dune::MultipleCodimMultipleGeomTypeMapper<
        LevelGridView
    >;
    using Scalar = typename Traits::Scalar;
    using Vector = typename Traits::Vector;
    using Tensor = typename Traits::Tensor;

    // Define the storage and cache for parameterizations and scaling
    /// Index type used by the element mapper
    using Index = typename Mapper::Index;
    /// Base class of the parameterization
    using ParameterizationType = typename Dorie::Parameterization::Transport<Traits>;
    /// Parameterization factory
    using ParameterizationFactory = Dorie::TransportParameterizationFactory<Traits>;

    /// Storage for parameters
    using ParameterStorage = std::unordered_map<Index,std::shared_ptr<ParameterizationType>>;

    /// Need this gruesome typedef because 'map::value_type' has const key
    using Cache = std::shared_ptr<ParameterizationType>;

    using ConstCache = std::shared_ptr<const ParameterizationType>;
    /// Configuration file tree
    const Dune::ParameterTree& _config;
    /// Grid view of the coarsest grid configuration (level 0)
    const LevelGridView _gv;
    /// Logger for this instance
    const std::shared_ptr<spdlog::logger> _log;
    /// Mapper for mapping grid elements to indices
    Mapper _mapper;
    /// Map for storing parameterization information for each grid entity
    ParameterStorage _param;
    /// Currently cached parameterization (bound to element)
    mutable Cache _cache;

private:

    /// Check if an entity has been cached
    void verify_cache () const
    {
        if (not _cache) {
            _log->error("Parameterization cache is empty. Call 'bind' before "
                        "accessing the cache or the parameterization");
            DUNE_THROW(Dune::InvalidStateException,
                        "No parameterization cached");
        }
    }

public:

    /// Copy constructor for solute parameters.
    /** Create a level grid view of level 0, create a mapper on it, and store
     *  the config tree.
     */
    SoluteParameters (
        const SoluteParameters& solute_param
    ):
        _config(solute_param._config),
        _gv(solute_param._gv),
        _log(solute_param._log),
        _mapper(solute_param._mapper)
    {
        // copy parameterization map
        this->_param.clear();
        for(const auto& [index, p] : solute_param._param) {
            // make a hard copy of parameterization
            std::shared_ptr<Parameterization::Transport<Traits>> _p = p->clone();
            this->_param.emplace(index,_p);
        }
    }

    /// Create this object and load the data.
    /** Create a level grid view of level 0, create a mapper on it, and store
     *  the config tree.
     *  \param config Configuration file tree
     *  \param grid Shared pointer to the grid
     *  \param element_index_map The mapping from grid entity index to
     *      parameterization index.
     */
    SoluteParameters (
        const Dune::ParameterTree& config,
        const std::shared_ptr<Grid> grid,
        const std::vector<int>& element_index_map
    ):
        _config(config),
        _gv(grid->levelGridView(0)),
        _log(Dorie::get_logger(log_transport)),
        _mapper(_gv, Dune::mcmgElementLayout())
    {
        build_parameterization(element_index_map);
    }

    /// Return a copy of the current cache with constant parameterization
    ConstCache cache() const
    {
        verify_cache();
        return _cache;
    }

    /// Return the current cache
    const Cache& cache()
    {
        verify_cache();
        return _cache;
    }

    /// Bind to a grid entity. Required to call parameterization functions!
    /** \param entity Grid entity (codim 0) to bind to
     */
    template<typename Entity>
    void bind (Entity entity) const
    {
        // retrieve index of top father element of chosen entity
        while(entity.hasFather()) {
            entity = entity.father();
        }
        const auto index = _mapper.index(entity);
        _cache = _param.find(index)->second;
    }

    /// Return the hydrodynamic dispersion function
    /** Uses the function of the underlying parameterization. Cast to
     *  operator-internal Tensor.
     *  \return Function: {WaterFlux,WaterContent} -> HydrdyDisp
     */
    std::function<Tensor(Vector,Scalar)> hydrodynamic_dispersion_f () const
    {
        verify_cache();
        const auto hd_dips = _cache->hydrodynamic_dispersion_f();

        using WaterFluxType = typename ParameterizationType::WaterFluxType;
        using WaterContentType = typename ParameterizationType::WaterContentType;
        return [hd_dips](const Vector water_flux, const Scalar water_content) {
            return hd_dips(WaterFluxType{water_flux},WaterContentType{water_content}).value;
        };
    }

    /// Return the microscopic peclet fuction 
    /** Uses the function of the underlying parameterization. Cast to
     *  operator-internal RF.
     *  \return Function: {WaterFlux,WaterContent} -> MicroPeclet
     */
    auto peclet_f () const
    {
        verify_cache();
        const auto peclet = _cache->peclet_f();

        using WaterFluxType = typename ParameterizationType::WaterFluxType;
        using WaterContentType = typename ParameterizationType::WaterContentType;
        return [peclet](const Vector water_flux, const Scalar water_content) {
            return peclet(WaterFluxType{water_flux},WaterContentType{water_content}).value;
        };
    }

private:

    /// Build the parameterization from an element mapping and the input file.
    void build_parameterization (const std::vector<int>& element_index_map)
    {
        // Open the YAML file
        const auto param_file_name = _config.get<std::string>("parameters.file");
        _log->info("Loading parameter data file: {}", param_file_name);
        YAML::Node param_file = YAML::LoadFile(param_file_name);

        // create the parameterization data
        ParameterizationFactory factory;
        const auto parameterization_map = factory.reader(param_file,"transport",_log);

        // insert parameterization for each element
        for (auto&& elem : elements(_gv)) {
            const auto index = _mapper.index(elem);
            auto p = parameterization_map.at(element_index_map.at(index));
            _param.emplace(index,p);
        }
        // check that mapper can map to parameterization
        assert(_param.size() == _mapper.size());
    }
};

} // namespace Dune
} // namespace Dorie

#endif // DUNE_DORIE_PARAM_TRANSPORT_HH