
#pragma once

#include "Solver.h"
#include "DummyTask.h"

#include <TNL/FileName.h>
#include <TNL/Timer.h>
#include <TNL/Meshes/Grid.h>
#include <TNL/Pointers/SharedPointer.h>
#include <TNL/Functions/MeshFunction.h>
#include <TNL/Images/PNGImage.h>
#include <TNL/Images/RegionOfInterest.h>

static std::vector< TNL::String > dimensionIds = { "x-dimension", "y-dimension", "z-dimension" };
static std::vector< TNL::String > kernelSizeIds = { "kernel-x-size", "kernel-y-size", "kernel-z-size" };
static std::vector< TNL::String > kernels = { "identity",        "gauss3x3",      "gauss5x5",
                                              "sobelHorizontal", "sobelVertical", "edgeDetection" };

class ImageSolver : public Solver< 2, TNL::Devices::Cuda >
{
public:
   constexpr static int Dimension = 2;
   using Device = TNL::Devices::Cuda;

   using Base = Solver< Dimension, Device >;
   using Vector = TNL::Containers::StaticVector< Dimension, int >;
   using DataStore = TNL::Containers::Vector< float, Device, int >;
   using HostDataStore = TNL::Containers::Vector< float, TNL::Devices::Host, int >;

   using GridType = TNL::Meshes::Grid< 2, float, Device, int >;
   using GridPointer = TNL::Pointers::SharedPointer< GridType >;
   using MeshFunctionType = TNL::Functions::MeshFunction< GridType >;

   virtual void
   start( const TNL::Config::ParameterContainer& parameters ) const override
   {
      GridPointer grid;
      MeshFunctionType meshFunction;
      TNL::Images::PNGImage< int > image;
      TNL::Images::RegionOfInterest< int > roi;

      meshFunction.setMesh( grid );

      auto output = parameters.getParameter< TNL::String >( "output" );

      if( ! this->readImage( parameters, grid, meshFunction, image, roi ) || ! this->convolve( parameters, meshFunction )
          || ! this->write( parameters, image, meshFunction ) )
         return;
   }

   virtual TNL::Config::ConfigDescription
   makeInputConfig() const override
   {
      TNL::Config::ConfigDescription config = Base::makeInputConfig();

      config.addDelimiter( "Image settings:" );

      config.addEntry< TNL::String >( "input", "PNG image" );
      config.addEntry< TNL::String >( "output", "PNG image" );
      config.addEntry< TNL::String >( "kernel", "A kernel to apply", kernels[ 0 ] );

      for( const auto& kernel : kernels )
         config.addEntryEnum( kernel );

      config.addDelimiter( "Roi settings:" );

      config.addEntry< int >( "roi-top", "Top (smaller number) line of the region of interest.", -1 );
      config.addEntry< int >( "roi-bottom", "Bottom (larger number) line of the region of interest.", -1 );
      config.addEntry< int >( "roi-left", "Left (smaller number) column of the region of interest.", -1 );
      config.addEntry< int >( "roi-right", "Right (larger number) column of the region of interest.", -1 );

      return config;
   }

   template< typename Image >
   bool
   readImage( const TNL::Config::ParameterContainer& parameters,
              GridPointer& grid,
              MeshFunctionType& meshFunction,
              Image& image,
              TNL::Images::RegionOfInterest< int >& roi ) const
   {
      auto input = parameters.getParameter< TNL::String >( "input" );

      if( image.openForRead( input ) ) {
         if( ! roi.setup( parameters, &image ) ) {
            std::cout << "Invalid image roi.";
            image.close();
            return false;
         }

         std::cout << image.getWidth() << " " << image.getHeight() << std::endl;

         auto meshPointer = meshFunction.getMeshPointer();

         meshPointer->setDimensions( image.getWidth(), image.getHeight() );

         meshFunction.setMesh( meshPointer );

         if( ! image.read( roi, meshFunction ) ) {
            std::cout << "Invalid image size" << std::endl;

            image.close();
            return false;
         }

         image.close();

         std::cout << "Image read was successful: " << meshFunction.getData().getSize() << " elements count" << std::endl;
         return true;
      }

      std::cout << "Image open for read failed. Please check file path" << std::endl;

      return false;
   }

   bool
   convolve( const TNL::Config::ParameterContainer& parameters, MeshFunctionType& meshFunction ) const
   {
      auto imageData = meshFunction.getData().getConstView();

      Vector kernelSize;
      DataStore kernel;

      kernel = getKernel( parameters, kernelSize );

      DataStore result;

      result.setLike( imageData );
      result = 0;

      TNL::Timer timer;

      timer.start();

      std::cout << imageData.getSize() << " " << result.getSize() << std::endl;

      launchConvolution(
         imageData, kernel.getConstView(), result.getView(), meshFunction.getMeshPointer()->getDimensions(), kernelSize );

      timer.stop();

      result.forAllElements(
         [] __cuda_callable__( int i, float& value )
         {
            value = TNL::max( TNL::min( value, 1.0 ), 0.0 );
         } );

      meshFunction.getData() = result;

      std::cout << "Image convolution was successful. Time: " << timer.getRealTime() << " sec" << std::endl;

      return true;
   }

   template< typename Image >
   bool
   write( const TNL::Config::ParameterContainer& parameters, Image& image, MeshFunctionType& meshFunction ) const
   {
      auto output = parameters.getParameter< TNL::String >( "output" );
      GridType grid = meshFunction.getMesh();

      if( image.openForWrite( output, grid ) ) {
         if( ! image.write( meshFunction ) ) {
            std::cout << "Image write failed" << std::endl;

            image.close();
            return false;
         }

         image.close();

         return true;
      }

      std::cout << "Image open for write failed. Please check file path" << std::endl;

      return false;
   }

   HostDataStore
   getKernel( const TNL::Config::ParameterContainer& parameters, Vector& kernelDimension ) const
   {
      auto kernel = parameters.getParameter< TNL::String >( "kernel" );

      if( kernel == "identity" ) {
         kernelDimension = { 3, 3 };

         return { 0, 0, 0, 0, 1, 0, 0, 0, 0 };
      }

      if( kernel == "gauss3x3" ) {
         kernelDimension = { 3, 3 };

         HostDataStore kernel = { 1, 2, 1, 2, 4, 2, 1, 2, 1 };

         kernel /= 16;

         return kernel;
      }

      if( kernel == "gauss5x5" ) {
         kernelDimension = { 5, 5 };

         HostDataStore kernel = { 1, 4, 7, 4, 1, 4, 16, 26, 16, 4, 7, 26, 41, 26, 7, 4, 16, 26, 16, 4, 1, 4, 7, 4, 1 };

         kernel /= 273;

         return kernel;
      }

      if( kernel == "sobelHorizontal" ) {
         kernelDimension = { 3, 3 };

         return { 1, 2, 1, 0, 0, 0, -1, -2, -1 };
      }

      if( kernel == "sobelVertical" ) {
         kernelDimension = { 3, 3 };

         return { 1, 0, -1, 2, 0, -2, 1, 0, -1 };
      }

      if( kernel == "edgeDetection" ) {
         kernelDimension = { 3, 3 };

         return { -1, -1, -1, -1, 8, -1, -1, -1, -1 };
      }

      std::cout << "Unknown kernel " << kernel << ". Exit" << std::endl;
      exit( 1 );
   }

   void
   launchConvolution( DataStore::ConstViewType image,
                      DataStore::ConstViewType kernel,
                      DataStore::ViewType result,
                      const GridType::CoordinatesType& imageDimension,
                      const GridType::CoordinatesType& kernelDimension ) const
   {
      DummyTask< int, float, Dimension, Device >::exec( imageDimension, kernelDimension, image, result, kernel );
   }
};
