
#ifndef COMPARISON_BENCHMARKS_BASE_HEAT_H
#define COMPARISON_BENCHMARKS_BASE_HEAT_H

#include <array>
#include <iostream>
#include <memory>

namespace comparison_benchmarks {
namespace base {
namespace heat {

const int DIFFUSION_SIZE = 256;
const int DIFFUSION_STEPS = 256;

const int NUM_ITERATIONS = 100;
const int WARMUP_ITERATIONS = 20;

template<typename T, int SIZE>
class heat_diffusion {
  // Center portion is SIZExSIZE, borders are fixed temperature values
  using matrix = std::array<std::array<T, SIZE + 2>, SIZE + 2>;

 protected:
  // Sane default values for the simulation (form paper).
  // This is not about perfect simulation results but the speedup of the workload.
  double c = 0.1;
  double d_s = 1.0 / (SIZE + 1);
  double d_t = (d_s * d_s) / (4 * c);

 public:
  matrix *current_data;
  matrix *next_data;

  explicit heat_diffusion() {
    current_data = new matrix;
    next_data = new matrix;
    reset_data();
  }

  ~heat_diffusion() {
    delete current_data;
    delete next_data;
  }

  virtual void run_simulation(int n) {
    for (int i = 0; i < n; i++) {
      for (int row = 1; row <= SIZE; row++) {
        for (int column = 1; column <= SIZE; column++) {
          update_element(row, column);
        }
      }

      // Synchronization point needed to coordinate the calculation!
      swap_data_arrays();
    }
  }

 protected:
  void update_element(int row, int column) {
    (*next_data)[row][column] = (*current_data)[row][column] + ((c * d_t) / (d_s * d_s)) *
        ((*current_data)[row + 1][column] + (*current_data)[row - 1][column]
            - 4 * (*current_data)[row][column]
            + (*current_data)[row][column + 1] + (*current_data)[row][column - 1]);
  }

  void swap_data_arrays() {
    matrix *tmp = current_data;
    current_data = next_data;
    next_data = tmp;
  }

  void reset_data() {
    for (int row = 0; row < SIZE + 2; row++) {
      for (int column = 0; column < SIZE + 2; column++) {
        (*current_data)[row][column] = 0.0;
        (*next_data)[row][column] = 0.0;

        // Edges are a fixed, hot temperature
        if (row == 0 || row == SIZE + 1) {
          (*current_data)[row][column] = 1.0;
          (*next_data)[row][column] = 1.0;
        }
      }
    }
  }
};

template<typename T, int SIZE>
std::ostream &operator<<(std::ostream &strm, const heat_diffusion<T, SIZE> &simulation) {
  for (int i = 0; i < SIZE + 2; i++) {
    for (int j = 0; j < SIZE + 2; j++) {
      // 'color' our output according to temperature
      char out;
      if (simulation.current_data[i][j] < 0.1) {
        out = ' ';
      } else if (simulation.current_data[i][j] < 0.2) {
        out = '-';
      } else if (simulation.current_data[i][j] < 0.5) {
        out = '=';
      } else {
        out = '#';
      }

      strm << out << "\t";
    }
    strm << std::endl;
  }

  return strm;
}

}
}
}

#endif //COMPARISON_BENCHMARKS_BASE_HEAT_H
