#include "method1.hh"
#include "common.hh"
#include "drawing.hh"
#include <algorithm>
#include <array>
#include <thread>
#include <vector>

const auto thread_nbr = std::thread::hardware_concurrency();
std::mutex numbers_mutex;

using randint = std::uniform_int_distribution<>;
using Color = std::array<uchar, 3>;
using ColorSet = std::vector<Color>;

void newSquare1(cv::Mat &t_process_img, std::mt19937 &t_gen,
                randint &t_rand_pos) {
  const int square_size = t_rand_pos(t_gen);
  auto square_top_left = cv::Point{t_rand_pos(t_gen), t_rand_pos(t_gen)};
  draw_shape(t_process_img, square_top_left, square_size, random_color(t_gen),
             Shapes::Square);
}

void method1(cv::Mat &t_reference, cv::Mat &t_output, int t_iterations,
             std::mt19937 &t_gen) {
  auto diff = euclidian_distance(t_reference, t_output);
  auto const max_size =
      std::max(t_reference.size().width, t_reference.size().height);
  randint dist(0, max_size);
  spdlog::info("Beginning method1, initial difference: {}", diff);
  while (t_iterations > 0) {
    auto temp_image = t_output.clone();
    newSquare1(temp_image, t_gen, dist);
    if (auto new_diff = euclidian_distance(t_reference, temp_image);
        new_diff < diff) {
      diff = new_diff;
      temp_image.copyTo(t_output);
      --t_iterations;
      spdlog::info("Iteration {}: diff {}", t_iterations, diff);
    }
  }
}

void threadedGetColor(cv::Mat &t_reference, ColorSet &t_colors, int t_h) {
  if (t_h > t_reference.size().height)
    return;
  for (int w = 0; w < t_reference.size().width; w += 3) {
    Color temp = {t_reference.at<uchar>(t_h, w),
                  t_reference.at<uchar>(t_h, w + 1),
                  t_reference.at<uchar>(t_h, w + 2)};
    auto pos = std::find(std::begin(t_colors), std::end(t_colors), temp);
    if (pos == std::end(t_colors)) {
      numbers_mutex.lock();
      t_colors.push_back(temp);
      numbers_mutex.unlock();
    }
  }
}

ColorSet getColorSet(cv::Mat &t_reference) {
  ColorSet res{};
  for (int h = 0; h < t_reference.size().height; h += thread_nbr) {
    std::vector<std::thread> thread_list{};
    for (int i = 0; i < thread_nbr; ++i) {
      thread_list.push_back(std::thread(threadedGetColor, std::ref(t_reference),
                                        std::ref(res), h + i));
    }
    for (auto &th : thread_list)
      th.join();
  }
  res.shrink_to_fit();
  return res;
}

void newSquare2(cv::Mat &t_process_img, std::mt19937 &t_gen,
                ColorSet const &t_colors, randint &t_rand_pos,
                randint &t_rand_color) {
  int const square_size = t_rand_pos(t_gen);
  auto square_top_left = cv::Point{t_rand_pos(t_gen), t_rand_pos(t_gen)};
  const auto color = t_colors[t_rand_color(t_gen)];
  draw_shape(t_process_img, square_top_left, square_size,
             cv::Scalar{static_cast<double>(color[0]),
                        static_cast<double>(color[1]),
                        static_cast<double>(color[2])},
             Shapes::Square);
}

void method2(cv::Mat &t_reference, cv::Mat &t_output, int t_iterations,
             std::mt19937 &t_gen) {
  auto diff = euclidian_distance(t_reference, t_output);
  auto const max_size =
      std::max(t_reference.size().width, t_reference.size().height);
  randint dist(0, max_size);
  spdlog::info("Beginning method2, initial difference: {}", diff);
  auto const colors = getColorSet(t_reference);
  spdlog::info("Running {} threads.", thread_nbr);
  spdlog::info("{} colors detected.", colors.size());
  randint rand_color(0, colors.size());
  while (t_iterations > 0) {
    auto temp_image = t_output.clone();
    newSquare2(temp_image, t_gen, colors, dist, rand_color);
    if (auto new_diff = euclidian_distance(t_reference, temp_image);
        new_diff < diff) {
      diff = new_diff;
      temp_image.copyTo(t_output);
      --t_iterations;
      spdlog::info("Iteration {}: diff {}", t_iterations, diff);
    }
  }
}