// +-------------------------------- DISCLAIMER ---------------------------------+
// |                                                                             |
// | This program is provided to you free of charge "as such".                   |
// | You should not assume that this program is error-free or suitable for any   |
// | purpose whatsoever.                                                         |
// |                                                                             |
// | EURESYS does not give any representation, warranty or undertaking that this |
// | program is free of any defect or error or suitable for any purpose. EURESYS |
// | shall not be liable, in contract, in torts or otherwise, for any damages,   |
// | loss, costs, expenses or other claims for compensation, including those     |
// | asserted by third parties, arising out of or in connection with the use of  |
// | this program.                                                               |
// +-----------------------------------------------------------------------------+

#include "CameraHandler.h"
#include <QPainter>
#include <QApplication>
#include <QScreen>
#include <sstream>
#include <iomanip>

#include "QtDrawAdapter.h"
#include "ProcessingHandler.h"

#include "EasyDeepLearning.h"

using namespace Euresys::Open_eVision;
using namespace Euresys::Open_eVision::EasyDeepLearning;


ProcessingHandler::ProcessingHandler(const std::string& deepLearningModel, CameraHandler* cameraHandler) : abort_(false), zoom_(1.0f), model_(deepLearningModel), cameraHandler_(cameraHandler)
{
}

ProcessingHandler::~ProcessingHandler()
{
  if (processingThread_.joinable())
  {
    abort_ = true;
    // waits for ProcessingHandler::run to end

    // notifies canProcessImage_ once in case a problem occured with the camera thread
    cameraHandler_->processImageMutex_.lock();
    cameraHandler_->canProcessImage_.notify_one();
    cameraHandler_->processImageMutex_.unlock();

    processingThread_.join();
  }
  std::cout << "Destroyed ProcessingHandler" << std::endl;
}

void ProcessingHandler::start()
{
  // run main loop in a separate thread
  processingThread_ = std::thread(&ProcessingHandler::run, this);
}

void ProcessingHandler::run()
{
  try
  {
    std::shared_ptr<EDeepLearningTool> tool;

    // Look model using the factory method of EDeepLearningTool
    tool = std::move(EDeepLearningTool::Create(model_));

    Easy::SetMaxNumberOfProcessingThreads(std::max(1, Easy::GetNumberOfAvailableProcessorCores() - 1));

    EDeepLearningExecutionSettings settings;
    for (const auto& device: settings.GetAvailableDevices())
    {
      if (device.GetDeviceType() == EasyDeepLearning::EDeepLearningDeviceType_GPU)
      {
        settings.SetDevice(device);
        break;
      }
    }

    tool->SetExecutionSettings(settings);

    QSize screenSize = qApp->primaryScreen()->availableSize();

    EImageC24 src;

    while(!abort_)
    {
      {
        std::unique_lock<std::mutex> lock(cameraHandler_->processImageMutex_);
        cameraHandler_->canProcessImage_.wait(lock);
        src.InitializeFromUnalignedBuffer(cameraHandler_->width_, cameraHandler_->height_, cameraHandler_->data_, cameraHandler_->pitch_);
      }

      if (cameraHandler_->width_ > screenSize.width() * zoom_ || cameraHandler_->height_ > screenSize.height() * zoom_)
      {
        // Compute a zoom factor such that the image fits in the screen.
        zoom_ = std::min((float)screenSize.width() / cameraHandler_->width_, (float)screenSize.height() / cameraHandler_->height_);
      }

      // draw the image with overlay
      QImage image(zoom_ * cameraHandler_->width_, zoom_ * cameraHandler_->height_, QImage::Format_ARGB32_Premultiplied);
      QPainter painter(&image);

      QtDrawAdapter drawAdapter(&painter);

      ERGBColor color = ERGBColor(0, 255, 255);
      EPen mypen(color, 5, EPenStyle_Solid);

      drawAdapter.SetPen(mypen);

      src.Draw(&drawAdapter, zoom_, zoom_, 0, 0);

      switch (tool->GetToolType())
      {
      case EDeepLearningToolType_EasyClassify:
        ProcessImage(src, std::dynamic_pointer_cast<EClassifier>(tool), &drawAdapter);
        break;
      case EDeepLearningToolType_EasySegmentSupervised:
        ProcessImage(src, std::dynamic_pointer_cast<ESupervisedSegmenter>(tool), &drawAdapter);
        break;
      case EDeepLearningToolType_EasySegmentUnsupervised:
        ProcessImage(src, std::dynamic_pointer_cast<EUnsupervisedSegmenter>(tool), &drawAdapter);
        break;
      case EDeepLearningToolType_EasyLocate:
      case EDeepLearningToolType_EasyLocateInterestPoint:
        ProcessImage(src, std::dynamic_pointer_cast<ELocatorBase>(tool), &drawAdapter);
        break;
      }

      Q_EMIT renderedImage(image, "");

    }
  }
  catch (EException e)
  {
    std::cerr << "OeV exception caught: " << e.what() << std::endl;
    abort_ = true;
  }
}

void ProcessingHandler::ProcessImage(const EImageC24& img, std::shared_ptr<EClassifier> classifier, EDrawAdapter* drawAdapter)
{
  EClassificationResult r = classifier->Classify(img);

  std::stringstream resultDisplay;
  resultDisplay << "Label: " << r.GetBestLabel() << " (" << std::setprecision(4) << r.GetBestProbability() * 100.f << "%)";

  EFont font("Arial", 16);
  drawAdapter->SetFont(font);

  EBrush blackBrush(ERGBColor(0, 0, 0));
  EBrush whiteBrush(ERGBColor(255, 255, 255));

  drawAdapter->BackedText(resultDisplay.str(), 10, 10, 0, blackBrush, whiteBrush);
}

void ProcessingHandler::ProcessImage(const EImageC24& img, std::shared_ptr<EUnsupervisedSegmenter> segmenter, EDrawAdapter* drawAdapter)
{
  EUnsupervisedSegmenterResult r = segmenter->Apply(img);

  r.Draw(drawAdapter, zoom_, zoom_);

  std::string resultDisplay;
  EBrush fgBrush;
  if (r.IsGood())
  {
    resultDisplay = "Good";
    fgBrush.Color = ERGBColor(0, 255, 0); // Green for good
  }
  else
  {
    resultDisplay = "NG";
    fgBrush.Color = ERGBColor(0, 255, 0); // Red for NG
  }

  EFont font("Arial", 16);
  drawAdapter->SetFont(font);

  EBrush bgBrush(ERGBColor(255, 255, 255)); // White background

  drawAdapter->BackedText(resultDisplay, 10, 10, 0, fgBrush, bgBrush);
}

void ProcessingHandler::ProcessImage(const EImageC24& img, std::shared_ptr<ESupervisedSegmenter> segmenter, EDrawAdapter* drawAdapter)
{
  ESupervisedSegmenterResult r = segmenter->Apply(img);

  r.Draw(drawAdapter, zoom_, zoom_);

  std::string resultDisplay = "Detected " + std::to_string(r.GetNumBlobs()) + " blobs";

  EFont font("Arial", 16);
  drawAdapter->SetFont(font);

  EBrush fgBrush(ERGBColor(0, 0, 0)); // Black text
  EBrush bgBrush(ERGBColor(255, 255, 255)); // White background

  drawAdapter->BackedText(resultDisplay, 10, 10, 0, fgBrush, bgBrush);
}

void ProcessingHandler::ProcessImage(const EImageC24& img, std::shared_ptr<ELocatorBase> locator, EDrawAdapter* drawAdapter)
{
  ELocatorResult r = locator->Apply(img);

  std::string result;
  if (r.GetNumDetectedObjects() == 0)
    result = "No object detected";
  else
    result = "Detected " + std::to_string(r.GetNumDetectedObjects()) + " objects";

  r.Draw(drawAdapter, zoom_, zoom_);

  EFont font("Arial", 16);
  drawAdapter->SetFont(font);

  EBrush blackBrush(ERGBColor(0, 0, 0));
  EBrush whiteBrush(ERGBColor(255, 255, 255));

  drawAdapter->BackedText(result, 10, 10, 0, blackBrush, whiteBrush);
}
