/***************************************************************************
                          OperatorFunction.h  -  description
                             -------------------
    begin                : Dec 31, 2015
    copyright            : (C) 2015 by oberhuber
    email                : tomas.oberhuber@fjfi.cvut.cz
 ***************************************************************************/

/* See Copyright Notice in tnl/Copyright */

#pragma once

#include <type_traits>
#include <TNL/Devices/Cuda.h>
#include <TNL/Functions/MeshFunction.h>
#include <TNL/Solvers/PDE/BoundaryConditionsSetter.h>

namespace TNL {
namespace Functions {   

/***
 * This class evaluates given operator on given preimageFunction. If the flag
 * EvaluateOnFly is set on true, the values on particular mesh entities
 * are computed just  when operator() is called. If the EvaluateOnFly flag
 * is 'false', values on all mesh entities are evaluated by calling a method
 * refresh() they are stores in internal mesh preimageFunction and the operator()
 * just returns precomputed values. If BoundaryConditions are void then the
 * values on the boundary mesh entities are undefined. In this case, the mesh
 * preimageFunction evaluator evaluates this preimageFunction only on the INTERIOR mesh entities.
 */
   
template< typename Operator,
          typename Function = typename Operator::FunctionType,
          typename BoundaryConditions = void,
          bool EvaluateOnFly = false,
          bool IsAnalytic = ( Function::getDomainType() == SpaceDomain || Function::getDomainType() == NonspaceDomain ) >
class OperatorFunction{};

/****
 * Specialization for 'On the fly' evaluation with the boundary conditions does not make sense.
 */
template< typename Operator,
          typename MeshFunctionT,
          typename BoundaryConditions,
          bool IsAnalytic >
class OperatorFunction< Operator, MeshFunctionT, BoundaryConditions, true, IsAnalytic >
 : public Domain< Operator::getDimension(), MeshDomain >
{
};

/****
 * Specialization for 'On the fly' evaluation and no boundary conditions.
 */
template< typename Operator,
          typename MeshFunctionT,
          bool IsAnalytic >
class OperatorFunction< Operator, MeshFunctionT, void, true, IsAnalytic >
 : public Domain< Operator::getDomainDimension(), Operator::getDomainType() >
{
   public:
 
      static_assert( MeshFunctionT::getDomainType() == MeshDomain ||
                     MeshFunctionT::getDomainType() == MeshInteriorDomain ||
                     MeshFunctionT::getDomainType() == MeshBoundaryDomain,
         "Only mesh preimageFnctions may be used in the operator preimageFunction. Use ExactOperatorFunction instead of OperatorFunction." );
      static_assert( std::is_same< typename Operator::MeshType, typename MeshFunctionT::MeshType >::value,
          "Both, operator and mesh preimageFunction must be defined on the same mesh." );
 
      typedef Operator OperatorType;
      typedef MeshFunctionT FunctionType;
      typedef typename OperatorType::MeshType MeshType;
      typedef typename OperatorType::RealType RealType;
      typedef typename OperatorType::DeviceType DeviceType;
      typedef typename OperatorType::IndexType IndexType;
      typedef typename OperatorType::ExactOperatorType ExactOperatorType;
      typedef MeshFunction< MeshType, OperatorType::getPreimageEntitiesDimension() > PreimageFunctionType;
      typedef SharedPointer< MeshType, DeviceType > MeshPointer;
      
      static constexpr int getEntitiesDimension() { return OperatorType::getImageEntitiesDimension(); };     
      
      OperatorFunction( const OperatorType& operator_ )
      :  operator_( operator_ ), preimageFunction( 0 ){};
 
      OperatorFunction( const OperatorType& operator_,
                           const FunctionType& preimageFunction )
      :  operator_( operator_ ), preimageFunction( &preimageFunction ){};
 
      const MeshType& getMesh() const
      {
         TNL_ASSERT_TRUE( this->preimageFunction, "The preimage function was not set." );
         return this->preimageFunction->getMesh();
      };
      
      const MeshPointer& getMeshPointer() const
      { 
         TNL_ASSERT_TRUE( this->preimageFunction, "The preimage function was not set." );
         return this->preimageFunction->getMeshPointer(); 
      };

      
      void setPreimageFunction( const FunctionType& preimageFunction ) { this->preimageFunction = &preimageFunction; }
 
      Operator& getOperator() { return this->operator_; }
 
      const Operator& getOperator() const { return this->operator_; }
 
      bool refresh( const RealType& time = 0.0 ) { return true; };
 
      bool deepRefresh( const RealType& time = 0.0 ) { return true; };
 
      template< typename MeshEntity >
      __cuda_callable__
      RealType operator()(
         const MeshEntity& meshEntity,
         const RealType& time = 0.0 ) const
      {
         TNL_ASSERT_TRUE( this->preimageFunction, "The preimage function was not set." );
         return operator_( *preimageFunction, meshEntity, time );
      }
 
   protected:
 
      const Operator& operator_;
 
      const FunctionType* preimageFunction;
 
      template< typename, typename > friend class MeshFunctionEvaluator;
};

/****
 * Specialization for precomputed evaluation and no boundary conditions.
 */
template< typename Operator,
          typename PreimageFunction,
          bool IsAnalytic >
class OperatorFunction< Operator, PreimageFunction, void, false, IsAnalytic >
 : public Domain< Operator::getDomainDimension(), Operator::getDomainType() >
{
   public:
 
      static_assert( PreimageFunction::getDomainType() == MeshDomain ||
                     PreimageFunction::getDomainType() == MeshInteriorDomain ||
                     PreimageFunction::getDomainType() == MeshBoundaryDomain,
         "Only mesh preimageFunctions may be used in the operator preimageFunction. Use ExactOperatorFunction instead of OperatorFunction." );
      static_assert( std::is_same< typename Operator::MeshType, typename PreimageFunction::MeshType >::value,
          "Both, operator and mesh preimageFunction must be defined on the same mesh." );
 
      typedef Operator OperatorType;
      typedef typename OperatorType::MeshType MeshType;
      typedef typename OperatorType::RealType RealType;
      typedef typename OperatorType::DeviceType DeviceType;
      typedef typename OperatorType::IndexType IndexType;
      typedef PreimageFunction PreimageFunctionType;
      typedef Functions::MeshFunction< MeshType, Operator::getImageEntitiesDimension() > ImageFunctionType;
      typedef OperatorFunction< Operator, PreimageFunction, void, true > OperatorFunctionType;
      typedef typename OperatorType::ExactOperatorType ExactOperatorType;
      typedef SharedPointer< MeshType, DeviceType > MeshPointer;
      
      static constexpr int getEntitiesDimension() { return OperatorType::getImageEntitiesDimension(); };     
      
      OperatorFunction( OperatorType& operator_,
                           const MeshPointer& mesh )
      :  operator_( operator_ ), imageFunction( mesh )
      {};
 
      OperatorFunction( OperatorType& operator_,
                           PreimageFunctionType& preimageFunction )
      :  operator_( operator_ ), imageFunction( preimageFunction.getMeshPointer() ), preimageFunction( &preimageFunction )
      {};
 
      const MeshType& getMesh() const { return this->imageFunction.getMesh(); };
      
      const MeshPointer& getMeshPointer() const { return this->imageFunction.getMeshPointer(); };
      
      ImageFunctionType& getImageFunction() { return this->imageFunction; };
 
      const ImageFunctionType& getImageFunction() const { return this->imageFunction; };
 
      void setPreimageFunction( PreimageFunctionType& preimageFunction )
      {
         this->preimageFunction = &preimageFunction;
         this->imageFunction.setMesh( preimageFunction.getMeshPointer() );
      };
 
      const PreimageFunctionType& getPreimageFunction() const { return *this->preimageFunction; };
 
      Operator& getOperator() { return this->operator_; }
 
      const Operator& getOperator() const { return this->operator_; }

      bool refresh( const RealType& time = 0.0 )
      {
         OperatorFunction operatorFunction( this->operator_, *preimageFunction );
         this->operator_.setPreimageFunction( *this->preimageFunction );
         if( ! this->operator_.refresh( time ) ||
             ! operatorFunction.refresh( time )  )
             return false;
         this->imageFunction = operatorFunction;
         return true;
      };
 
      bool deepRefresh( const RealType& time = 0.0 )
      {
         if( ! this->preimageFunction->deepRefresh( time ) )
            return false;
         return this->refresh( time );
      };
 
      template< typename MeshEntity >
      __cuda_callable__
      RealType operator()(
         const MeshEntity& meshEntity,
         const RealType& time = 0 ) const
      {
         return imageFunction[ meshEntity.getIndex() ];
      }
 
      __cuda_callable__
      RealType operator[]( const IndexType& index ) const
      {
         return imageFunction[ index ];
      }
 
   protected:
 
      Operator& operator_;
 
      PreimageFunctionType* preimageFunction;
 
      ImageFunctionType imageFunction;
 
      template< typename, typename > friend class MeshFunctionEvaluator;
};

/****
 * Specialization for precomputed evaluation and with boundary conditions.
 */
template< typename Operator,
          typename Function,
          typename BoundaryConditions,
          bool IsAnalytic >
class OperatorFunction< Operator, Function, BoundaryConditions, false, IsAnalytic >
  : public Domain< Operator::getDimension(), MeshDomain >
{
   public:
 
      static_assert( Function::getDomainType() == MeshDomain ||
                     Function::getDomainType() == MeshInteriorDomain ||
                     Function::getDomainType() == MeshBoundaryDomain,
         "Only mesh preimageFunctions may be used in the operator preimageFunction. Use ExactOperatorFunction instead of OperatorFunction." );
      static_assert( std::is_same< typename Operator::MeshType, typename Function::MeshType >::value,
          "Both, operator and mesh preimageFunction must be defined on the same mesh." );
      static_assert( std::is_same< typename BoundaryConditions::MeshType, typename Operator::MeshType >::value,
         "The operator and the boundary conditions are defined on different mesh types." );
 
      typedef Operator OperatorType;
      typedef typename OperatorType::MeshType MeshType;
      typedef SharedPointer< MeshType > MeshPointer;
      typedef typename OperatorType::RealType RealType;
      typedef typename OperatorType::DeviceType DeviceType;
      typedef typename OperatorType::IndexType IndexType;
      typedef Function PreimageFunctionType;
      typedef Functions::MeshFunction< MeshType, Operator::getImageEntitiesDimension() > ImageFunctionType;
      typedef BoundaryConditions BoundaryConditionsType;
      typedef OperatorFunction< Operator, Function, void, true > OperatorFunctionType;
      typedef typename OperatorType::ExactOperatorType ExactOperatorType;
 
      static constexpr int getEntitiesDimension() { return OperatorType::getImageEntitiesDimension(); };
 
      OperatorFunction( OperatorType& operator_,
                           const BoundaryConditionsType& boundaryConditions,
                           const MeshPointer& meshPointer )
      :  operator_( operator_ ),
         boundaryConditions( boundaryConditions ),
         imageFunction( meshPointer )//,
         //preimageFunction( 0 )
      {
         this->preimageFunction = NULL;
      };
      
      OperatorFunction( OperatorType& operator_,
                           const BoundaryConditionsType& boundaryConditions,
                           const PreimageFunctionType& preimageFunction )
      :  operator_( operator_ ),
         boundaryConditions( boundaryConditions ),
         imageFunction( preimageFunction.getMeshPointer() ),
         preimageFunction( &preimageFunction )
      {};
 
      const MeshType& getMesh() const { return imageFunction.getMesh(); };
      
      const MeshPointer& getMeshPointer() const { return imageFunction.getMeshPointer(); };
      
      void setPreimageFunction( const PreimageFunctionType& preimageFunction )
      {
         this->preimageFunction = &preimageFunction;
      }
 
      const PreimageFunctionType& getPreimageFunction() const
      {
         TNL_ASSERT_TRUE( this->preimageFunction, "The preimage function was not set." );
         return *this->preimageFunction;
      };
 
      PreimageFunctionType& getPreimageFunction()
      {
         TNL_ASSERT_TRUE( this->preimageFunction, "The preimage function was not set." );
         return *this->preimageFunction;
      };
 
      const ImageFunctionType& getImageFunction() const { return this->imageFunction; };
 
      ImageFunctionType& getImageFunction() { return this->imageFunction; };
 
      Operator& getOperator() { return this->operator_; }
 
      const Operator& getOperator() const { return this->operator_; }

      bool refresh( const RealType& time = 0.0 )
      {
         OperatorFunctionType operatorFunction( this->operator_, *this->preimageFunction );
         this->operator_.setPreimageFunction( *this->preimageFunction );
         if( ! this->operator_.refresh( time ) ||
             ! operatorFunction.refresh( time )  )
             return false;
         this->imageFunction = operatorFunction;
         Solvers::PDE::BoundaryConditionsSetter< ImageFunctionType, BoundaryConditionsType >::apply( this->boundaryConditions, time, this->imageFunction );
         return true;
      };
 
      bool deepRefresh( const RealType& time = 0.0 )
      {
         return preimageFunction->deepRefresh( time ) &&
                this->refresh( time );
      };
 
      template< typename MeshEntity >
      __cuda_callable__
      const RealType& operator()(
         const MeshEntity& meshEntity,
         const RealType& time = 0 ) const
      {
         return imageFunction[ meshEntity.getIndex() ];
      }
 
      __cuda_callable__
      const RealType& operator[]( const IndexType& index ) const
      {
         return imageFunction[ index ];
      }
 
   protected:
 
      Operator& operator_;
 
      const PreimageFunctionType* preimageFunction;
 
      ImageFunctionType imageFunction;
 
      const BoundaryConditionsType& boundaryConditions;
 
      template< typename, typename > friend class MeshFunctionEvaluator;
};

/****
 * Specialization for precomputed evaluation and with boundary conditions.
 */
template< typename Operator,
          typename Function >
class OperatorFunction< Operator, Function, void, false, true >
  : public Domain< Function::getDomainDimension(), Function::getDomainType() >
{
   public:
      
      typedef Function FunctionType;
      typedef typename FunctionType::RealType RealType;
      typedef typename FunctionType::PointType PointType;
      typedef Operator OperatorType;
      
      bool setup( const Config::ParameterContainer& parameters,
                  const String& prefix = "" )
      {
         return( this->function.setup( parameters, prefix) &&
                 this->operator_.setup( parameters, prefix ) );
      }
      
      __cuda_callable__
      FunctionType& getFunction()
      {
         return this->function;
      }
      
      __cuda_callable__
      const FunctionType& getFunction() const
      {
         return this->function;
      }
      
      __cuda_callable__
      OperatorType& getOperator()
      {
         return this->operator_;
      }
      
      __cuda_callable__
      const OperatorType& getOperator() const
      {
         return this->operator_;
      }
      
      __cuda_callable__
      RealType operator()( const PointType& v,
                           const RealType& time = 0.0 ) const
      {
         return this->operator_( this->function, v, time );
      }
      
   protected:
      
      Function function;
      
      Operator operator_;
 
};

} // namespace Functions
} // namespace TNL

