/// LSU EE 4702-1 (Fall 2025), GPU Programming
//
//
 /// Homework 1 -- SOLUTION
 //
 //  Assignment: https://www.ece.lsu.edu/gpup/2025/hw01.pdf
 //
 //  This code is based on the solution to 2023 Homework 1.
 //  For the solution to 2023 Homework 1 look 
 //  at file ../../2023/hw01/hw01-sol.cc
 //  or visit https://www.ece.lsu.edu/koppel/gpup/2023/hw01-sol.cc.html
 //

#include <gp/coord.h>
#include <numbers>
#include "frame_buffer.h"

// Definitions for iVect, iCoor
#include "hw01-misc.h"

template<class T>
class safe_vector : public vector<T> {
public:
  T& operator [] ( size_t idx )
  {
    return vector<T>::at(idx);
  }
};

struct HW01_Data {
  HW01_Data():inited(false){};
  bool inited;
  bool tryout0, tryout1;
  // Note: Safe vector defined in hw01-misc.h
  safe_vector<uint32_t> fb_area_dup;

  bool opt_circle_hw01;

  float gamma;
  int mag_r;
};


void
circle_draw_hw01
( pFrame_Buffer& fb, HW01_Data& hw01_data,
  int ctr_x, int ctr_y,
  int r_outer, int r_inner, uint32_t co)
{
  const int win_width [[maybe_unused]] = fb.width_get();
  const int win_height [[maybe_unused]] = fb.height_get();

  /// SOLUTION -- Problem 1

  // Pre-compute the squares of the inner and outer radii.
  //
  int r_o_sq = r_outer * r_outer;
  int r_i_sq = r_inner * r_inner;

  // The term "ring" refers to the area from r_inner to r_outer.
  //
  // Technically ring refers to the parts of the circle at
  // (ctr_x,ctr_y) of radius r_outer which are not also parts of the
  // circle at (ctr_x,ctr_y) of radius r_inner.

  // Compute frame buffer array index of center of circle.
  //
  int idx_ctr = ctr_y * win_width + ctr_x;

  // Idea 1 (and 3): Iterate over circle-local coordinates, ly and lx.
  // Center of circle is at lx=0, ly=0.
  //
  for ( int ly = 0;  ly <= r_outer;  ly++ )
    {
      // Idea 2: (with x and y swapped)
      // Compute range of x values over the ring at ly.
      //
      int lx0 = sqrtf( max( 0, r_i_sq - ly * ly ) );
      int lx1 = sqrtf( max( 0, r_o_sq - ly * ly ) );

      // Compute frame buffer array index distance from center of
      // circle to the pixel at (lx=0, ly ):
      //
      int ly_ww = ly * win_width;

      for ( int lx = lx0;  lx <= lx1;  lx++ )
        {
          // Idea 3: Draw a pixel in each of four quadrants of the circle.
          //
          fb[ idx_ctr + ly_ww + lx ] = co;
          fb[ idx_ctr + ly_ww - lx ] = co;
          fb[ idx_ctr - ly_ww + lx ] = co;
          fb[ idx_ctr - ly_ww - lx ] = co;
        }
    }
}

void
circle_draw_parametric
( pFrame_Buffer& fb, HW01_Data& hw01_data,
  int ctr_x, int ctr_y,
  int r_outer, int r_inner, uint32_t co)
{
  const int win_width = fb.width_get();

  /// This routine DOES NOT have to be modified.
  //
  //  It can be modified to experiment or to help in understanding
  //  the code, but no problems in 2025 Homework 1 require that this
  //  routine be modified.

  // Choose delta_theta so that a loop will write every pixel along
  // the outer edge.
  float delta_theta = asinf(1.0f/r_outer);
  for ( float theta = 0; theta < 2*numbers::pi; theta += delta_theta )
    {
      float cth = cosf(theta),       sth = sinf(theta);
      for ( int r = r_inner; r<=r_outer; r++ )
        {
          int x = ctr_x + r * cth,   y = ctr_y + r * sth;
          fb[ y * win_width + x ] = co;
        }
    }
}

void circle_draw
( pFrame_Buffer& fb, HW01_Data& hw01_data,
  int ctr_x, int ctr_y,
  int r_outer, int r_inner, uint32_t co)
{
  if ( hw01_data.opt_circle_hw01 )
    circle_draw_hw01(fb,hw01_data,ctr_x,ctr_y,r_outer, r_inner,co);
  else
    circle_draw_parametric(fb,hw01_data,ctr_x,ctr_y,r_outer, r_inner,co);
}


void
line_draw
( pFrame_Buffer& fb, iCoor p0, iCoor p1, uint32_t co)
{
  int x0 = p0.x, y0 = p0.y;
  int x1 = p1.x, y1 = p1.y;

  const iVect delta = p1 - p0;
  const int width = fb.width_get();

  if ( abs( delta.x ) < abs( delta.y ) )
    {
      // Iterate over y
      if ( y0 > y1 ) { swap(x0,x1); swap(y0,y1); }

      float dxdy = float(delta.x) / delta.y;
      float x = x0;

      for ( int yi = y0;  yi < y1;  yi++ )
        {
          fb[ yi * width + x ] = co;
          x += dxdy;
        }
    }
  else
    {
      // Iterate over x
      if ( x0 > x1 ) { swap(x0,x1); swap(y0,y1); }

      float dydx = float(delta.y) / delta.x;
      float y = y0;

      for ( int xi = x0;  xi < x1;  xi++ )
        {
          fb[ int(y) * width + xi ] = co;
          y += dydx;
        }
    }
}


void
parallel_lines_draw
( pFrame_Buffer& fb, iCoor p0, iCoor pb, iVect v, uint32_t co)
{
  line_draw( fb, p0, p0 + v, co );
  line_draw( fb, pb, pb + v, co );
}

void
parallelogram_draw
( pFrame_Buffer& fb, iCoor p00, iVect vx,  iVect vy, uint32_t color )
{
  parallel_lines_draw(fb, p00, p00+vy, vx, color);
  parallel_lines_draw(fb, p00, p00+vx, vy, color);
}

void
aa_rectangle_draw
( pFrame_Buffer& fb, iCoor p00, int wd, int ht, uint32_t color )
{
  // Window Size, in Pixels
  //
  const int win_width [[maybe_unused]] = fb.width_get();
  const int win_height [[maybe_unused]] = fb.height_get();

  iVect vx(wd,0);
  iVect vy(0,ht);
  iCoor p10 = p00 + vx;
  iCoor p01 = p00 + vy;
  iCoor p11 = p10 + vy;

  // Draw vertical lines.
  line_draw( fb, p00, p01, color );
  line_draw( fb, p10, p11, color );

  // Draw horizontal lines.
  line_draw( fb, p00, p10, color );
  line_draw( fb, p01, p11, color );

}

void kbd_input_handle(pFrame_Buffer& fb, HW01_Data& hw01_data);


void
render_hw01(pFrame_Buffer& fb, HW01_Data& hw01_data)
{
  // Window Size, in Pixels
  //
  const int win_width = fb.width_get();
  const int win_height = fb.height_get();

  // Adjust size and position of mag window based on keyboard input.
  kbd_input_handle(fb,hw01_data);

  /// Coordinates Used for Placement of Other Objects
  //
  const int sq_slen = min(win_width,win_height) * .45;
  const int sq_x0 = sq_slen * 0.1,   sq_y0 = win_height - 150;
  const int sq_x1 = sq_x0 + sq_slen, sq_y1 = sq_y0 - sq_slen;
  //
  // One vertex of the square is at (sq_x0,sq_y0) and the opposite
  // vertex is at (sq_x1,sq_y1).

  //                           RRGGBB
  const uint32_t color_red = 0xff0000;
  const uint32_t color_blue =    0xff;
  const uint32_t color_green = 0xff00;

  /// Demo Code: Draw Sine Waves  -- No need to modify this.
  //
  int num_waves = 4;
  float plot_ht = sq_slen/float(num_waves);
  for ( float
          yb = sq_y0 - plot_ht/2,  freq = 4 * M_PI / ( win_width - 40 - sq_x1 );
        yb > sq_y1;
        yb -= plot_ht,  freq *= 2 )
    for ( int x = 20;  x < win_width - 20;  x++ )
      {
        int y = yb + sin( x * freq ) * plot_ht * .45;
        fb[ y * win_width + x ] = 0xff00;
      }

  /// Demo Code: Draw Lines Radiating From a Point -- No need to modify.
  //
  //  This code calls the line_draw routine.
  //
  int n_rects = 24;
  int n_pts = n_rects * 3;
  float c_r1 = sq_slen/6;  // Inner radius
  float c_r2 = sq_slen/2;  // Outer radius.
  pCoor c_ctr( c_r2 + 20, c_r2 + 20 ); // Center of "circle"
  float delta_theta = 2 * M_PI / n_pts;
  vector<pVect> vecs;

  circle_draw(fb, hw01_data, c_ctr.x, c_ctr.y, c_r2+9, c_r2+1, color_green);
  circle_draw(fb, hw01_data, c_ctr.x, c_ctr.y, c_r1-1, c_r1-2, color_red);

  for ( int i=0; i<n_pts; i++ )
    vecs.push_back( pVect( cosf(i*delta_theta), sinf(i*delta_theta), 0 ) );
  for ( int i=0; i<n_pts; i+=3 )
    {
      pVect v1 = vecs[i+1] * ( c_r2 - c_r1 );
      pVect v2 = c_r1 * ( vecs[i+2] - vecs[i] );
      pCoor p1 = c_ctr + c_r1 * vecs[i];
      parallelogram_draw(fb, p1, v1, v2, 0xffffff );
    }

  fb.fb_stats_off(); // Turn off "redundancy" statistics.


  /// Homework 1 Problem 2

  // Coordinates of mouse pointer and circle center.
  //
  const iCoor ctr(fb.mouse_x,fb.mouse_y);
  //
  // See 2024 Homework 1 for a description of iCoor.

  // Radius of distorted area.
  //
  const int mag_r = hw01_data.mag_r;         // Radius of distorted area.

  // Distortion factor.
  //
  const float gamma [[maybe_unused]] = hw01_data.gamma;
  //
  // Gamma can be modified in the UI by pressing z and u.

  // Storage for a copy of the frame buffer in and around the blue circle.
  //
  auto& fb_dup [[maybe_unused]] = hw01_data.fb_area_dup;

  // Draw a blue circle.
  //
  circle_draw(fb, hw01_data,  ctr.x, ctr.y,  mag_r+2, mag_r+1, color_blue);

  /// Put Problem 2 Solution Below (and possibly above, but mostly below).
  //

  // [✔] Copy area inside the box to fb_dup.
  // [✔] Compute d' (see handout) 
  // [✔] Use d' to compute coordinates of part of frame buffer (actually fb_dup)
  //     to read from.
  // [y] Read from those coordinates and write to (x,y).

  /// SOLUTION -- Problem 2

  int idx_ctr = ctr.y * win_width + ctr.x;
  int mag_r_sq = mag_r * mag_r;
  int mag_2r = 2 * mag_r;

  // Copy area under blue circle to fb_dup.
  //
  for ( int x = -mag_r; x < mag_r; x++ )
    for ( int y = -mag_r; y < mag_r; y++ )
      fb_dup[ ( y + mag_r ) * mag_2r + x + mag_r ]
        = fb[ ( ctr.y + y ) * win_width + ctr.x + x ];
  //
  // This also copies pixels outside the circle. We don't need the
  // pixels outside the circle, but we copy them anyway because it
  // might take more time to determine whether a pixel is in the
  // circle than to copy it.

  const float mag_r_inv = 1.0 / mag_r;

  // [✔] Change x and y bounds so they iterate over a blue circle bounding box.
  //     Actually, iterate over one quadrant of circle in local coordinates.
  //
  for ( int lx = 0; lx < mag_r; lx++ )
    for ( int ly = 0; ly < mag_r; ly++ )
      {
        // Skip this iteration if we are not inside the circle.
        //
        int len_sq = lx*lx + ly*ly;
        if ( len_sq > mag_r_sq ) break;

        // Compute distance from pixel to circle center ..
        //
        float d = powf( len_sq, 0.5 );
        //
        // .. and the distorted distance.
        //
        float d_prime = powf( d * mag_r_inv, gamma ) * mag_r;

        // Scale lx and ly based on the distorted distance.
        //
        int lx_src = d ? lx * d_prime / d : 0;
        int ly_src = d ? ly * d_prime / d : 0;
        //
        // Note: because d_prime is a float lx is converted (promoted)
        // to a float before performing the multiplication lx * d_prime.

        // For each quadrant, read the pixel at lx_src, ly_src from fb_dup,
        // and write it to fp.
        //
        fb[ idx_ctr + ly * win_width + lx ] =
          fb_dup[ (mag_r+ly_src) * mag_2r + lx_src + mag_r];
        fb[ idx_ctr + ly * win_width - lx ] =
          fb_dup[ (mag_r+ly_src) * mag_2r - lx_src + mag_r];
        fb[ idx_ctr - ly * win_width - lx ] =
          fb_dup[ (mag_r-ly_src) * mag_2r - lx_src + mag_r];
        fb[ idx_ctr - ly * win_width + lx ] =
          fb_dup[ (mag_r-ly_src) * mag_2r + lx_src + mag_r];
        //
        // It would not be correct to both read and write from fb, since
        // the pixel read might be an already distorted pixel.
      }

}

void
kbd_input_handle(pFrame_Buffer& fb, HW01_Data& hw01_data)
{
  if ( !hw01_data.inited )
    {
      hw01_data.inited = true;
      hw01_data.gamma = 1;
      hw01_data.mag_r = 100;
      hw01_data.opt_circle_hw01 = true;
      int mag_2r = (hw01_data.mag_r+1) * 2;

      hw01_data.tryout0 = false;
      hw01_data.tryout1 = false;

      hw01_data.fb_area_dup.resize( mag_2r * mag_2r );
    }

  switch ( fb.keyboard_key ) {
#if 0
  case FB_KEY_RIGHT: case FB_KEY_LEFT:
    {
      const int inc = fb.keyboard_key == FB_KEY_LEFT ? -kbd_inc : kbd_inc;
      break;
    }
  case FB_KEY_UP: case FB_KEY_DOWN:
    {
      const int inc = fb.keyboard_key == FB_KEY_DOWN ? -kbd_inc : kbd_inc;
      break;
    }
#endif

  case 'c': case 'C':
    hw01_data.opt_circle_hw01 = !hw01_data.opt_circle_hw01; break;

  case 'u': hw01_data.gamma /= 1.1; break;
  case 'U': hw01_data.gamma /= 1.5; break;
  case 'z': hw01_data.gamma *= 1.1; break;
  case 'Z': hw01_data.gamma *= 1.5; break;

  case 'y':
    hw01_data.tryout0 = !hw01_data.tryout0; break;
  case 'Y':
    hw01_data.tryout1 = !hw01_data.tryout1; break;
  default:
    break;
  };

  fb.fbprintf("Gamma: %.2f ('z','u')  Circle Routine %s ('c')  "
              "Tryout 0,1:  %s,%s ('y','Y')\n",
              hw01_data.gamma,
              hw01_data.opt_circle_hw01
              ? "circle_draw_hw01" : "circle_draw_parametric",
              hw01_data.tryout0 ? "TRUE " : "FALSE",
              hw01_data.tryout1 ? "TRUE " : "FALSE");
}

int
main(int argc, char **argv)
{
  pFrame_Buffer demo_frame_buffer(argc,argv);

  HW01_Data hw01_data;
  demo_frame_buffer.show
    ( [&](pFrame_Buffer& fb){ render_hw01(fb,hw01_data); } );
  return 0;
}