
#ifndef COMPARISON_BENCHMARKS_BASE_HEAT_H
#define COMPARISON_BENCHMARKS_BASE_HEAT_H

#include <vector>
#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>
class heat_diffusion {
  // Center portion is SIZExSIZE, borders are fixed temperature values.
  // Should be a matrix of (SIZE + 2) x (SIZE + 2)
  using matrix = std::vector<T>;
 public:
  size_t size_;

  // 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_)};

  matrix data_1_;
  matrix data_2_;

  matrix *current_data_;
  matrix *next_data_;

  explicit heat_diffusion(size_t size) : size_{size},
                                         data_1_((size + 2) * (size + 2)),
                                         data_2_((size + 2) * (size + 2)) {
    current_data_ = &data_1_;
    next_data_ = &data_2_;

    reset_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();
    }
  }

  T &current_data_at(int row, int column) {
    return (*current_data_)[row * (size_ + 2) + column];
  }

  T &next_data_at(int row, int column) {
    return (*next_data_)[row * (size_ + 2) + column];
  }

  void update_element(int row, int column) {
    next_data_at(row, column) = current_data_at(row, column) + ((c_ * d_t_) / (d_s_ * d_s_)) *
        current_data_at(row + 1, column) + current_data_at(row - 1, column)
        - 4 * current_data_at(row, column)
        + current_data_at(row, column + 1) + current_data_at(row, column - 1);
  }

  void swap_data_arrays() {
    std::swap(current_data_, next_data_);
  }

  void reset_data() {
    for (int row = 0; row < size_ + 2; row++) {
      for (int column = 0; column < size_ + 2; column++) {
        current_data_at(row, column) = 0.0;
        next_data_at(row, column) = 0.0;

        // Edges are a fixed, hot temperature
        if (row == 0 || row == size_ + 1) {
          current_data_at(row, column) = 1.0;
          next_data_at(row, column) = 1.0;
        }
      }
    }
  }
};

template<typename T>
std::ostream &operator<<(std::ostream &strm, const heat_diffusion<T> &simulation) {
  for (int i = 0; i < simulation.size_ + 2; i++) {
    for (int j = 0; j < simulation.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
