#include "library.h"

const int win_width = 600, 
          win_height = 600;
double range_min_real = -2.25,
       range_max_real = 0.75,
       range_min_imag = -1.5, 
       range_max_imag = 1.5;

struct complex
{ double r, i; };

complex make_c(double r0, double i0)
{ complex x;
  x.r = r0;
  x.i = i0;
  return x; }

complex add(complex a, complex b)
{ complex x;
  x.r = a.r + b.r;
  x.i = a.i + b.i;
  return x; }

complex sub(complex a, complex b)
{ complex x;
  x.r = a.r - b.r;
  x.i = a.i - b.i;
  return x; }

complex mul(complex a, complex b)
{ complex x;
  x.r = a.r * b.r - a.i * b.i;
  x.i = a.r * b.i + a.i * b.r;
  return x; }

complex div(complex a, complex b)
{ double bottom = b.r * b.r + b.i * b.i;
  if (bottom == 0)
    cout << "Look out! Division by zero\n";
  complex x;
  x.r = (a.r * b.r + a.i * b.i) / bottom;
  x.i = (a.i * b.r - a.r * b.i) / bottom;
  return x; }

double magnitude(complex a)
{ return sqrt(a.r * a.r + a.i * a.i); }

ostream & operator<<(ostream & out, complex a)
{ out << a.r;
  if (a.i > 0)
    out << "+" << a.i << "i";
  else if (a.i < 0)
    out << a.i << "i";
  return out; }

bool operator==(complex a, complex b)
{ return a.r == b.r && a.i == b.i; }

bool operator!=(complex a, complex b)
{ return a.r != b.r || a.i != b.i; }

double r_to_x(double r)
{ double p = (r - range_min_real) / (range_max_real - range_min_real);
  return p * win_width; }

double x_to_r(double x)
{ double p = x / win_width;
  return range_min_real + p * (range_max_real - range_min_real); }

double i_to_y(double i)
{ double p = (i - range_min_imag) / (range_max_imag - range_min_imag);
  return win_height - p * win_height; }

double y_to_i(double y)
{ double p = (win_height - y) / win_height;
  return range_min_imag + p * (range_max_imag - range_min_imag); }

struct point
{ double x, y; };

point make_p(double x, double y)
{ point p;
  p.x = x;
  p.y = y;
  return p; }

ostream & operator<<(ostream & out, point p)
{ out << "(" << p.x << ", " << p.y << ")";
  return out; }

point complex_to_point(complex a)
{ point p;
  p.x = r_to_x(a.r);
  p.y = i_to_y(a.i);
  return p; }

complex point_to_complex(point p)
{ complex c;
  c.r = x_to_r(p.x);
  c.i = y_to_i(p.y);
  return c; }

void move_to(complex c)
{ point p = complex_to_point(c);
  move_to(p.x, p.y); }

void draw_to(complex c)
{ point p = complex_to_point(c);
  draw_to(p.x, p.y); }

void draw_point(complex c)
{ point q = complex_to_point(c);
  draw_point(q.x, q.y); }

void draw_axes()
{ set_pen_width(1);
  set_pen_color(color::grey);
  complex c;
  double axis_step = (range_max_real - range_min_real) / 10;
  for (c.r = range_min_real; c.r <= range_max_real; c.r += axis_step)
  { c.i = range_min_imag;
    move_to(c);
    c.i = range_max_imag;
    draw_to(c); }
  axis_step = (range_max_imag - range_min_imag) / 10;
  for (c.i = range_min_imag; c.i <= range_max_imag; c.i += axis_step)
  { c.r = range_min_real;
    move_to(c);
    c.r = range_max_real;
    draw_to(c); }
  set_pen_color(color::black);
  if (range_min_real <= 0 && 0 <= range_max_real)
  { move_to(make_c(0.0, range_min_imag));
    draw_to(make_c(0.0, range_max_imag)); }
  if (range_min_imag <= 0 && 0 <= range_max_imag)
  { move_to(make_c(range_min_real, 0.0));
    draw_to(make_c(range_max_real, 0.0)); } }

point click_to_point()
{ wait_for_mouse_click();
  point p;
  p.x = get_click_x();
  p.y = get_click_y();
  return p; }

complex click_to_complex()
{ return point_to_complex(click_to_point()); }

int steps_to_explosion(complex z, int max_steps)
{ int steps = 0;
  complex c = z;
  while (true)
  { if (magnitude(z) > 5)
      break;
    if (steps >= max_steps)
      break;
    z = add(mul(z, z), c);
    steps += 1; }
  return steps; }

void produce_shape(int max_steps, int & least_steps, int & most_steps)
{ cout << "range " << make_c(range_min_real, range_min_imag) << " to " << make_c(range_max_real, range_max_imag)
       << " (size " << range_max_real - range_min_real << "), max steps " << max_steps << "\n";
  set_pen_color(color::black);
  move_to(0, win_height - 1);
  draw_to(win_width, win_height - 1);
  int l_steps = INT_MAX, m_steps = 1;
  for (int x = 0; x < win_width; x += 1)
    for (int y = win_height; y >= 0; y -= 1)
    { complex z = point_to_complex(make_p(x, y));
      int s = steps_to_explosion(z, max_steps);
      if (s > m_steps)
        m_steps = s;
      if (s < l_steps)
        l_steps = s;
      int c = make_color_hls(0.666666 + 0.333333 * (s - least_steps) / (most_steps - least_steps), 0.5, 1);
      set_pixel_color(x, y, c);
      set_pixel_color(x, win_height - 2, color::black);
      set_pixel_color(x, win_height - 1, color::white); }
  least_steps = l_steps;
  most_steps = m_steps;
  cout << "least steps = " << least_steps << ", most steps = " << most_steps << "\n"; }

const int normal = 1, set_max_steps = 2, reset = 3, enhance_contrast = 4, try_again = 5;

int deal_with_key_press(char kt)
{ if (kt == 'z')
  { cout << "click on opposite corners\n";
    complex c1 = click_to_complex();
    complex c2 = click_to_complex();
    double minr = min(c1.r, c2.r), maxr = max(c1.r, c2.r), mini = min(c1.i, c2.i), maxi = max(c1.i, c2.i);
    double range = max(maxr - minr, maxi - mini);
    range_min_real = minr;
    range_max_real = minr + range;
    range_min_imag = mini;
    range_max_imag = mini + range;
    return normal; }
  else if (kt == 'r')
  { range_min_real = -2.25;
    range_max_real = 0.75;
    range_min_imag = -1.5;
    range_max_imag = 1.5;
    return reset; }
  else if (kt == 'm')
  { double size;
    cout << "select this window then enter min-real min-imag size: ";
    cin >> range_min_real >> range_min_imag >> size;
    range_max_real = range_min_real + size;
    range_max_imag = range_min_imag + size;
    return normal; }
  else if (kt == 's')
    return set_max_steps;
  else if (kt == 'a')
  { draw_axes();
    return try_again; }
  else if (kt == 'e')
    return enhance_contrast;
  else
  { cout << "key " << kt << " meaningless\n";
    return try_again; } }
 
void main()
{ make_window(win_width, win_height);
  cout << "type z to zoom in, r to reset, m to manually enter details, s to change max steps, a to draw axes, e to enhance contrast\n";
  set_pen_width(1);
  int action = normal, min_steps = 2, max_steps = 100, new_min_steps = 1, new_max_steps = 100;
  while (true)
  { if (action == enhance_contrast)
      produce_shape(new_max_steps, new_min_steps, new_max_steps);
    else if (action != try_again)
    { int old_max_steps = max_steps;
      produce_shape(max_steps, new_min_steps, new_max_steps);
      max_steps = old_max_steps; }
    cout << "ready\n";
    char c = wait_for_key_typed();
    action = deal_with_key_press(c);
    if (action == reset)
      max_steps = 100;
    else if (action == set_max_steps)
    { cout << "select this window then enter maximum steps: ";
      cin >> max_steps;
      action = normal; } } }