#include <thread>
#include <vector>
#include <iostream>
#include <chrono>
#include <complex>
#include <atomic>
#include <unistd.h>
#include <sstream>
#include "libpng_wrapper.hpp"
#include "mthread.hpp"
#include "test.hpp"

using namespace std;

//Argument defaults. Really not used to using constants over preprocessor defines; this feels cursed
const uint32_t DEFAULT_WIDTH = 1920;
const uint32_t DEFAULT_HEIGHT = 1080;
const int DEFAULT_JOBS = 1;
const string DEFAULT_IMG_PATH = "out.png";

const complex<double> DEFAULT_MIN_CORD (-0.74364386269 - 0.00000003000, 0.13182590271 - 0.00000003000);
const complex<double> DEFAULT_MAX_CORD (-0.74364386269 + 0.00000003000, 0.13182590271 + 0.00000003000);

const unsigned int DEFAULT_ITERS = 50000;
const unsigned int DEFAULT_BAILOUT = 256;

void print_help(char *arg, bool error) {
  stringstream help_text; 
  help_text.precision(numeric_limits<double>::max_digits10);
  help_text << "Usage: " << arg << " [options]\n"
    "Options:\n"
    "\t-h\tthis cruft\n"
    "\t-w\timage width \t\t\t\tdefault: " << DEFAULT_WIDTH << "\n"
    "\t-H\timage height\t\t\t\tdefault: " << DEFAULT_HEIGHT << "\n"
    "\t-o\timage output path\t\t\tdefault: " << DEFAULT_IMG_PATH << "\n"
    "\t-j\tjobs -- set this to your corecount\tdefault: " << DEFAULT_JOBS << "\n"
    "\t-c\tcomplex bottom border\t\t\tdefault: " << DEFAULT_MIN_CORD << "\n"
    "\t-C\tcomplex top border\t\t\tdefualt: " << DEFAULT_MAX_CORD << " \n"
    "\t-i\tfractal iterations\t\t\tdefault: " << DEFAULT_ITERS << "\n"
    "\t-I\tbailout value\t\t\t\tdefault: " << DEFAULT_BAILOUT << "\n"
    "\nFOR COMPLEX NUMBERS: if you want to input, say, 2-3i, your option argument will be \"(2,-3)\".\n";
  cout << help_text.str() << endl;
  exit(error);
}

template <class t>
bool getopt_int(t& number, char *optarg, char opt, char *arg) {
  try { number = stoi(optarg); }
  catch(invalid_argument const&) {
    cout << "You must supply an integer for option -" << opt << "." << endl;
    print_help(arg, true);
    return true;
  }
  catch(out_of_range const&) {
    cout << "You must supply an integer under " << numeric_limits<t>::max() << " for option -" << opt << "." << endl;
    print_help(arg, true);
    return true;
  }
  return false;
}

/** colors using linear histrogram method with generic hue shift **/
//TODO figure out how to pass constant arrays without allowing pointed material to be modified
void calc_colors(const double *value_map, unsigned int *histogram, png &image, unsigned int max_iters) {
  double below, above, hue;
  int c;
  uint32_t x, y;
  int rgb[3];
  unsigned int histogram_sum = 0;
  double current_hue = 0;
  double *freq_hue = new double[max_iters]();
  
  //find the sum of all histogram values, we could ajust this to increase or decrease contrast
  for(unsigned int p = 0; p < max_iters; p++) histogram_sum += histogram[p];

  for(unsigned int i = 0; i < max_iters; i++) {
    current_hue += histogram[i] / (double)histogram_sum;
    freq_hue[i] = current_hue;
  }

  for(y = 0; y < image.height(); y++) {
    for(x = 0; x < image.width(); x++) {
      below = freq_hue[(int)value_map[(y * image.width()) + x]];
      above = freq_hue[(int)ceil(value_map[((y * image.width()) + x) + 1])];
      hue = (((above - below) * fmod(value_map[(y * image.width()) + x], 1.0)) + below);
      rgb[0] = 255 * cos((M_PI * hue) - M_PI);
      rgb[1] = 255 * cos((M_PI * hue) - ((M_PI) / 2.0));
      rgb[2] = 255 * cos(M_PI * hue);
      for(c = 0; c < 3; c++) if(rgb[c] < 0) rgb[c] = 0;
      image.set_pixel(x, y, (png_byte)rgb[0], (png_byte)rgb[1], (png_byte)rgb[2]);
    }
  }
}



int main(int argc, char **argv) {



  //argument options
  uint32_t width = DEFAULT_WIDTH;
  uint32_t height = DEFAULT_HEIGHT;
  unsigned int jobs = DEFAULT_JOBS;
  string img_path = DEFAULT_IMG_PATH;
  complex<double> min_cord = DEFAULT_MIN_CORD;
  complex<double> max_cord = DEFAULT_MAX_CORD;
  unsigned int m_iters = DEFAULT_ITERS;
  unsigned int bailout = DEFAULT_BAILOUT;
  bool jobs_set = false;
  int arg; //for getopt

  stringstream complex_str_buffer;
  while((arg = getopt(argc, argv, "hw:H:o:j:c:C:i:I:")) != -1) {
    switch(arg) {
      case 'h':
        print_help(argv[0], false);
        break;
      case 'w':
        getopt_int(width, optarg, 'w', argv[0]);
        break;
      case 'H':
        getopt_int(height, optarg, 'H', argv[0]);
        break;
      case 'o':
        img_path = optarg;
        break;
      case 'j':
        getopt_int(jobs, optarg, 'j', argv[0]);
        jobs_set = true;
        break;
      case 'c':
        complex_str_buffer << optarg;
        complex_str_buffer >> min_cord;
        break;
      case 'C':
        complex_str_buffer << optarg;
        complex_str_buffer >> max_cord;
        break;
      case 'i':
        getopt_int(m_iters, optarg, 'i', argv[0]);
        break;
      case 'I':
        getopt_int(bailout, optarg, 'I', argv[0]);
        break;
      default:
        cout << "Invalid option." << endl;
        print_help(argv[0], true);
        exit(1);
        break;
    }
  }

  if(!jobs_set) {
    cout << "\nPERFORMANCE TIP: for best preformance, set jobs to the number of cores in your CPU.\n"
      "See " << argv[0] << " -h for help.\n" << endl;
  }

  //dependent on getopt settings
  double *bail_map = new double[width * height];
  unsigned int *histogram = new unsigned int[m_iters]();
  unsigned int width_per_job = width / jobs;
  atomic<uint32_t> progress(0);
  png image(img_path, width, height);
  thread threads[jobs];
  

  //allocate worker threads, spawn workers
  mthread** worker_objects = (mthread **)malloc(sizeof(mthread) * jobs);
  

  for(unsigned int j = 0; j < jobs - 1; j++) {
    worker_objects[j] = new mthread(j * width_per_job, (j + 1) * width_per_job, 
        min_cord, max_cord, bailout, m_iters, 
        image, bail_map, histogram, worker_objects, j, jobs, progress);
  }

  //last worker thread needs the width to go all the way to the edge of the screen,
  //regardless of rounding issues
  worker_objects[jobs - 1] = new mthread((jobs - 1) * width_per_job, width - 1, 
      min_cord, max_cord, bailout, m_iters, 
      image, bail_map, histogram, worker_objects, jobs - 1, jobs, progress);
  
  
  for(unsigned int j = 0; j < jobs; j++) worker_objects[j]->dispatch();

  //the progress variables is simply how many pixels we have calculated
  while(progress < (height * jobs)) {
    cout << "\033[2K\033[0GCalculating pixel values... " << ((float)progress / (height * jobs)) * 100 << "\% complete" << flush;
    this_thread::sleep_for(chrono::milliseconds(100));
  }
  cout << endl;

  for(unsigned int j = 0; j < jobs; j++) worker_objects[j]->join();
  
  //now to color the image
  cout << "Coloring image... (this shouldn't take more then a few seconds)" << endl;

  //actually writes image as well
  calc_colors(bail_map, histogram, image, m_iters);

  for(unsigned int j = 0; j < jobs; j++) delete worker_objects[j];

  //image will be deleted (and written) upon funciton exit
  cout << "Image exported to " << img_path << "." << endl;

}