/// LSU EE 4702-1 (Fall 2024), GPU Programming
//
 /// Homework 3 -- SOLUTION
 //
 //  Assignment: https://www.ece.lsu.edu/gpup/2024/hw03.pdf
 //
 /// Based on cpu-only/demo-06-rend-pipe.cc

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

enum HW3_Lines {
  Lines_None, Lines_Pixel, Lines_World, Lines_ENUM_SIZE };
const char* const hw3_lines_str[] =
  { "NONE", "PIXEL", "WORLD" };

  ///
 /// Our_3D: Low-Level 3D Graphics Library.
 ///
class Our_3D {
public:
  Our_3D(pFrame_Buffer& fbp)
    :frame_buffer(fbp)
  {
    vtx_colors = nullptr;
    vtx_coors = nullptr;
    vtx_normals = nullptr;

    topology_strip = false;
    opt_light_frag = true;
    opt_lines = Lines_None;
  };

  Our_3D& transform_eye_from_object_set(const pMatrix& m)
  {
    eye_from_object = m;
    return *this;
  }
  Our_3D& transform_clip_from_eye_set(const pMatrix& m)
  {
    clip_from_eye = m;
    return *this;
  }
  Our_3D& light_location_set(const pCoor c)
  {
    light_location = c;
    return *this;
  }
  Our_3D& light_color_set(const pColor c)
  {
    light_color = c;
    return *this;
  }

  Our_3D& vtx_coors_set(vector<pCoor>& c)
  {
    vtx_coors = &c;
    return *this;
  }

  Our_3D& vtx_colors_set(vector<pColor>& c)
  {
    vtx_colors = &c;
    return *this;
  }

  // Specify whether vertices are grouped for individual triangles or
  // a triangle strip in the next call to draw_rasterization.
  //
  Our_3D& topology_strip_set( bool strip = true )
  {
    topology_strip = strip;
    return *this;
  }

  //  Specify vertex normals to by used by the next draw_rasterization.
  //
  Our_3D& vtx_normals_set(vector<pVect4>& c)
  {
    vtx_normals = c.data();
    return *this;
  }

  Our_3D& draw_rasterization();

  pFrame_Buffer& frame_buffer;
  pMatrix eye_from_object, clip_from_eye;
  pCoor light_location;
  pColor light_color;
  vector<pCoor> *vtx_coors;
  pVect4 *vtx_normals;
  vector<pColor> *vtx_colors;
  bool topology_strip;
  bool opt_tryout1, opt_tryout2;
  bool opt_light_frag;
  int opt_lines;
};

 ///
 ///  World: An application.
 ///
class World : public World_Base {
public:
  World(pFrame_Buffer& fb):World_Base(fb),gc(fb){}
  Our_3D gc;

  // Called each time a frame needs to be rendered.
  virtual void render(bool first_time)
  {
    if ( first_time ) init_scene();
    render_scene();
  }

  // Called Once
  void init_scene();

  // Called for each frame.
  void render_scene();

  pCoor eye_location, light_location;
  pColor light_color;
  pCoor *adj_location;
  pVect eye_direction;
  bool opt_tryout1,opt_tryout2;
};

void
World::init_scene()
{
  eye_location = pCoor(-1.2,.5,10.2);
  light_location = pCoor(0.7, 1.9, -0.6);
  eye_direction = pVect(.1,0,-1);
  adj_location = &eye_location;
  opt_tryout1 = opt_tryout2 = false;
  gc.opt_tryout1 = opt_tryout1;
  gc.opt_tryout2 = opt_tryout2;
  light_color = color_white * 6.3;
}

void
World::render_scene()
{
  pFrame_Buffer& frame_buffer = gc.frame_buffer;

  //
  /// Process User Key Presses
  //
  switch ( frame_buffer.keyboard_key ) {
  case FB_KEY_LEFT: adj_location->x -= 0.1; break;
  case FB_KEY_RIGHT: adj_location->x += 0.1; break;
  case FB_KEY_PAGE_UP: adj_location->y += 0.1; break;
  case FB_KEY_PAGE_DOWN: adj_location->y -= 0.1; break;
  case FB_KEY_UP: adj_location->z -= 0.1; break;
  case FB_KEY_DOWN: adj_location->z += 0.1; break;
  case 'e': case 'E': adj_location = &eye_location; break;
  case 'l': case 'L': adj_location = &light_location; break;
  case 'y': gc.opt_tryout1 = opt_tryout1 = !opt_tryout1; break;
  case 'Y': gc.opt_tryout2 = opt_tryout2 = !opt_tryout2; break;
  case 'n': case 'N':
    gc.opt_lines++;
    if ( gc.opt_lines >= Lines_ENUM_SIZE ) gc.opt_lines = Lines_None;
    break;
  case 'p': case 'P': gc.opt_light_frag = !gc.opt_light_frag; break;
  case '+': case '=':
    if ( !frame_buffer.keyboard_control ) light_color *= 1.1;
    break;
  case '-':
    if ( !frame_buffer.keyboard_control ) light_color *= 1/1.1;
    break;
  }

  ///
  /// Print "Green Text" Information
  ///
  gc.frame_buffer.fbprintf
    ("Lines %s ('n')  Lighting per %s ('p')  "
     "Tryout1 %s ('y')  Tryout2 %s ('Y') \n",
     hw3_lines_str[gc.opt_lines],
     gc.opt_light_frag ? "FRAGMENT" : "VERTEX  ",
     opt_tryout1 ? "ON " : "OFF", opt_tryout2 ? "ON " : "OFF");
  gc.frame_buffer.fbprintf
    ( "Light Location [ %.2f, %.2f, %.2f ],  Light Intensity %.2f ('+-')  "
      "Eye Location [ %.2f, %.2f, %.2f ]\n",
      light_location.x, light_location.y, light_location.z,
      light_color.r,
      eye_location.x, eye_location.y, eye_location.z );

  gc.light_location_set( light_location ).light_color_set( light_color );


  /// Specifications for Cylinder (Ring)
  //
  pCoor cyl_center(-1,-1,-3);
  int n_segs = 50;
  pVect cyl_axis( 0, 0.75, 0.05 );
  float cyl_radius = 4;
  //
  // The cylinder is constructed much further below.

  ///
  /// Prepare Transformations
  ///

  // Specifications of User's Monitor
  //
  const uint win_width = frame_buffer.width_get();
  const uint win_height = frame_buffer.height_get();
  //
  // win_width and win_height are pixel-space coordinates.
  //
  const float aspect = 1.0 * win_width / win_height;
  const float width_m = 0.8;  // Width of user's monitor.
  const float height_m = width_m / aspect;
  //
  // width_m and height_m are in eye-space coordinates.

  // Specify Transformation from Object to Eye Space: eye_from_object.
  //
  pMatrix_Translate center_eye(-eye_location);
  pMatrix_Rotation rotate_eye(eye_direction,pVect(0,0,-1));
  pMatrix eye_from_object = rotate_eye * center_eye;

  gc.transform_eye_from_object_set(eye_from_object);

  // Transform from Eye to Clip
  //
  pMatrix_Frustum clip_from_eye
    ( -width_m/2, width_m/2,  -height_m/2, height_m/2,  1, 5000 );

  gc.transform_clip_from_eye_set(clip_from_eye);


  ///
  /// Prepare Triangles
  ///

  // List of coordinates and normals in object space and list of colors.
  //
  vector<pCoor> coors_os;
  vector<pVect4> norms_os;
  vector<pColor> colors;

  // Insert First Triangle
  //
  // Insert coordinates of a triangle into coordinate list,
  // and insert its color into a color list.
  //
  coors_os << pCoor( 0, 0, 0 ) << pCoor( 9, 6, -5 ) << pCoor( 0, 7, -3 );
  colors << color_dark_olive_green << color_dark_olive_green
         << color_dark_olive_green;

  // Insert Second Triangle
  //
  // Add coordinates and color of another triangle
  //
  coors_os << pCoor(0.3,1.5, -.1) << pCoor(4,2,-5) << pCoor(5,0,0);
  pColor gray = 0.2 * color_white;
  colors << gray << gray << gray;

  // Insert Two More Triangles
  //
  // Add a square consisting of a gold and multicolored triangle.
  //
  coors_os << pCoor(-4,0,-3) << pCoor(-4,2,-3) << pCoor(-2,0,-3);
  colors << color_gold << color_gold << color_gold;
  //
  coors_os << pCoor(-4,2,-3) << pCoor(-2,2,-3) << pCoor(-2,0,-3);
  colors << color_red << color_green << color_blue;

  /// SOLUTION -- Problem 1 a
  //
  //  Insert purple triangle.
  //
  coors_os << pCoor( -2, 2, -3) << pCoor( -2, 1, -3 ) << pCoor( 0, 2, -3 );
  colors << color_lsu_spirit_purple << color_lsu_spirit_purple
         << color_lsu_spirit_purple;

  // Insert a triangle at light location, facing user. Light is top point.
  pVect light_to_eye( light_location, eye_location );
  float l_sz = 0.1;
  pVect lax = cross( pVect( 0, 1, 0 ), light_to_eye ).normal() * l_sz;
  pVect lay = cross( light_to_eye, lax ).normal() * l_sz;
  coors_os << light_location
           << light_location - lax - lay << light_location + lax - lay;
  colors << color_white << color_white << color_white;

  /// SOLUTION -- Problem 1 b
  //
  //  Insert a triangle with vertices at the top and bottom of the
  //  cylinder axis and the third vertex towards the light.
  //
  pCoor cyl_center_top = cyl_center + cyl_axis;
  pNorm cct_to_l( cyl_center_top, light_location);
  coors_os << cyl_center << cyl_center_top << cyl_center_top + 2 * cct_to_l;
  colors << color_salmon << color_salmon << color_salmon;

  // Now compute the triangle normals.
  for ( size_t i = 0;  i < coors_os.size();  i += 3 )
    {
      pCoor p1 = coors_os[i], p2 = coors_os[i+1], p3 = coors_os[i+2];
      pVect4 n = cross( p1, p2, p3 ).normal();
      norms_os << n << n << n;
    }

  ///
  /// Send Triangles to Our3D for Rendering
  ///
  // Tell Our_3D that in the next draws vertices will not be grouped
  // as triangle strips.
  //
  gc.topology_strip_set(false);

  gc.vtx_coors_set(coors_os);
  gc.vtx_colors_set(colors);
  gc.vtx_normals_set(norms_os);
  gc.draw_rasterization();


  ///
  /// Prepare Ring (Cylinder)
  ///

  // Empty the containers so they can be used for the next draw.
  //
  coors_os.clear(); colors.clear(); norms_os.clear();

  // Enter primitives for a cylinder (a short one, looks like a ring).
  //
  pNorm ax(0,-cyl_axis.z,cyl_axis.y);  // Find an axis orthogonal to cyl_axis.
  pNorm ay = cross( cyl_axis, ax );
  float delta_theta = 2 * M_PI / n_segs;

  for ( int i=0; i<=n_segs; i++ )
    {
      float theta = i * delta_theta;
      pVect4 n = ax * cosf(theta) + ay * sinf(theta);
      pCoor c = cyl_center + cyl_radius * n;

      coors_os << c + cyl_axis << c;
      norms_os << n << n;
      colors << color_cyan << color_cyan;
    }

  gc.topology_strip_set(true);
  gc.vtx_normals_set(norms_os);
  gc.vtx_coors_set(coors_os).vtx_colors_set(colors);
  gc.draw_rasterization();
}



 ///
 /// Code Below (in Our3D) is Part of (Fictional) 3D Graphics Library
 ///

Our_3D&
Our_3D::draw_rasterization()
{
  ///
  /// Prepare Transformations
  ///

  // Specifications of User's Monitor
  //
  const uint win_width = frame_buffer.width_get();
  const uint win_height = frame_buffer.height_get();
  //
  // win_width and win_height are pixel-space coordinates.

  // Transform from Clip to Pixel
  //
  pMatrix_Translate recenter(pVect(1,1,0));
  pMatrix_Scale scale( win_width/2, win_height/2, 1 );
  pMatrix window_from_clip = scale * recenter;
  pMatrix ws_from_os = window_from_clip * clip_from_eye * eye_from_object;
  pCoor eye_location = -eye_from_object.col_get(3);
  pColor light_ambient = 0.4 * color_white;

  frame_buffer.n_vtx_frame += vtx_coors->size();

  ///
  /// Rasterize Triangle and Write Frame Buffer
  ///

  // Get pointer to the depth (z) buffer.
  //
  float* const zbuffer = frame_buffer.buffer_depth_get();
  //
  // The depth buffer uses the same index as the frame buffer.

  auto& coors_os = *vtx_coors;
  auto& colors = *vtx_colors;

  //  Set loop stride based on indicated grouping.
  //
  const size_t inc = topology_strip ? 1 : 3;

  /// SOLUTION -- Problem 2
  //
  //  Compute lighted colors *before* the main loop if opt_light_frag is false.
  //
  vector<pColor> lighted_colors;
  if ( !opt_light_frag )
    {
      // For performance reasons allocate space in advance, since we know
      // how much space we need.
      //
      lighted_colors.resize(coors_os.size());

      for ( size_t i=0; i < coors_os.size(); i++ )
        {
          pVect4 n = vtx_normals[i];
          pColor color = colors[i];
          pCoor of = coors_os[i];

          // Compute lighted color of vertex i.
          //
          pNorm frag_to_light(of,light_location);
          pVect frag_to_eye(of,eye_location);
          float cos_angle_eye = dot( n, frag_to_eye );
          float cos_angle_light = dot( n, frag_to_light );
          bool illumination_visible =
            ( cos_angle_light > 0 )  ==  ( cos_angle_eye > 0 );
          float clamped_phase =
            illumination_visible ? fabs(cos_angle_light) : 0;
          pColor lighted_color =
            ( light_ambient
              + clamped_phase / frag_to_light.magnitude * light_color )
            * color;

          // Write lighted color into new lighted_colors container.
          //
          lighted_colors[i] = lighted_color;
        }
    }

  ///
  /// Main Rasterization Loop
  ///
  //
  ///  Outer Loop: Iterate Over Triangles
  //
  for ( size_t i=0; i+2 < coors_os.size(); i += inc )
    {
      frame_buffer.n_tri_frame++; // Count of number of triangles.

      // Get next triangle's object space coordinates ..
      //
      pCoor o0 = coors_os[i+0],  o1 = coors_os[i+1],  o2 = coors_os[i+2];
      //
      // .. convert them into unhomogenized window-space coordinates ..
      //
      pCoor u0( ws_from_os * o0 ), u1( ws_from_os * o1 ), u2( ws_from_os * o2 );
      //
      // .. and homogenize them.
      //
      pCoor_Homogenized w0( u0 ), w1( u1 ), w2( u2 );

      //  Extract colors from list.
      //
      pColor c0 = colors[i], c1 = colors[i+1], c2 = colors[i+2];

      //  Extract normals from list.
      //
      pVect4 n0 = vtx_normals[i], n1 = vtx_normals[i+1], n2 = vtx_normals[i+2];

      // Compute increment that will advance by approximately one pixel along
      //  b0 and b1 dimensions.
      //
      pVect v20(w2,w0), v12(w1,w2);
      float db0 = 0.8/max(fabs(v20.x),fabs(v20.y));
      float db1 = 0.8/max(fabs(v12.x),fabs(v12.y));

      if ( db0 < 0.00025 || db1 < 0.00025 ) continue;

      // Used for computing perspective-correct barycentric coordinates.
      const float u0ou2 = u0.w/u2.w, u0ou1 = u0.w/u1.w;
      const float u1ou2 = u1.w/u2.w, u1ou0 = u1.w/u0.w;

      /// SOLUTION -- Problem 3
      //
      //  Prepare to draw 0.1-object-space-units-wide triangle edges.
      //
      //  Compute value of each object-space barycentric coordinate
      //  (bc0, bc1, pc2) that corresponds to a distance of 0.1 from
      //  the opposite edges.
      //
      float line_wid_o = 0.1;
      //
      // Compute vectors along each edge.
      //
      pVect o01(o0,o1), o12(o1,o2), o20(o2,o0);
      //
      // Compute triangle normal.
      //
      pVect nt = cross(o20,o01);
      //
      // Compute an orthogonal vector for each triangle edge.
      //
      pNorm Oo01(cross(nt,o01)), Oo12(cross(nt,o12)), Oo20(cross(nt,o20));
      //
      // Compute object-space distance from each vertex to opposite
      // edge. For example, dist_0_12 is distance from vertex 0 to
      // edge 12.
      //
      float dist_0_12 = dot( o20, Oo12 );
      float dist_1_20 = dot( o01, Oo20 );
      float dist_2_01 = dot( o12, Oo01 );
      //
      // Find object-space barycentric coordinate value at which
      // interpolated fragment is line_wid_o units from opposite edge.
      // For example, b0o is the value of object-space barycentric
      // coordinate bc0 at which the interpolated fragment is at
      // distance exactly line_wid_o object-space units from edge 12.
      //
      float b0o = min(1.0f, line_wid_o / dist_0_12 );
      float b1o = min(1.0f, line_wid_o / dist_1_20 );
      float b2o = min(1.0f, line_wid_o / dist_2_01 );
      //
      //
      //  Prepare to draw 2-pixel-wide triangle edges.
      //
      //  The method for finding the barycentric coordinate limits,
      //  b0w, b1w, b2w, is similar to the procedure for finding
      //  limits on object-space barycentric coordinates, except that
      //  these limits are computed using the window-space coordinates
      //  of the vertices.
      //
      float line_wid_p = 2;  // It doesn't have to be two. Adjust to taste.
      pVect v01(w0,w1);
      pVect wn = cross(v20,v01);
      pNorm Ow01(cross(wn,v01)), Ow12(cross(wn,v12)), Ow20(cross(wn,v20));
      float b0w = min(1.0f, line_wid_p / dot( v20, Ow12 ) );
      float b1w = min(1.0f, line_wid_p / dot( v01, Ow20 ) );
      float b2w = min(1.0f, line_wid_p / dot( v12, Ow01 ) );

      //
      /// Inner Loop: Iterate over Triangle Fragments
      //
      for ( float b0=0; b0<=1; b0 += db0 )
        for ( float b1=0; b1<=1-b0; b1 += db1 )
          {
            frame_buffer.n_frag_frame++; // Count of number of fragments.
            const float b2 = 1 - b0 - b1;

            // Compute window-space coordinate of fragment.
            //
            pCoor wf = b0*w0 + b1*w1 + b2*w2;

            // Check whether fragment is within window.
            //
            if ( uint(wf.x) >= win_width || uint(wf.y) >= win_height ) continue;

            // Compute the index into the frame and depth buffers.
            //
            const size_t idx = wf.x + int(wf.y) * win_width;

            // Compute perspective-correct barycentric coordinates.
            //
            float bc0 = b0 / ( b2 * u0ou2 + b1 * u0ou1 + b0 );
            float bc1 = b1 / ( b2 * u1ou2 + b0 * u1ou0 + b1 );
            float bc2 = 1 - bc0 - bc1;
            //
            // These should be used for interpolating everything except
            // pixel-space coordinates.

            switch ( opt_lines ){
            case Lines_None:
              break;

            case Lines_Pixel:

              /// SOLUTION -- Problem 3
              //
              //  If all of the window-space barycentric coordinates
              //  are too far away from an edge skip this fragment.
              //
              if (  b0 > b0w  &&   b1 > b1w  &&   b2 > b2w  ) continue;
              break;

            case Lines_World:

              /// SOLUTION -- Problem 3
              //
              //  If all of the object-space barycentric coordinates
              //  are too far away from an edge skip this fragment.
              //
              if ( bc0 > b0o  &&  bc1 > b1o  &&  bc2 > b2o  ) continue;
              break;

            default: assert( false ); break;
            }

            // Depth (Z) Test
            //
            float zinv = wf.z;
            if ( zinv < 0 || zinv > 1 ) continue;
            if ( zbuffer[ idx ] < zinv ) continue;

            zbuffer[ idx ] = zinv;


            /// SOLUTION -- Problem 2
            //
            //  Compute the lighted color if opt_light_frag is true,
            //  otherwise interpolate the values from lighted_colors.
            //
            pColor lighted_color;
            //
            if ( opt_light_frag )
              {
                // Find object-space coordinate of fragment.
                //
                pCoor of =     bc0*o0 + bc1*o1 + bc2*o2;

                // "Blend" normal of each vertex.
                //
                pNorm n =      bc0*n0 + bc1*n1 + bc2*n2;

                // Blend color of each vertex together.
                //
                pColor color = bc0*c0 + bc1*c1 + bc2*c2;

                // Compute Lighted Color
                //
                pNorm frag_to_light(of,light_location);
                pVect frag_to_eye(of,eye_location);
                float cos_angle_eye = dot( n, frag_to_eye );
                float cos_angle_light = dot( n, frag_to_light );
                bool illumination_visible =
                  ( cos_angle_light > 0 )  ==  ( cos_angle_eye > 0 );
                float clamped_phase =
                  illumination_visible ? fabs(cos_angle_light) : 0;
                lighted_color =
                  ( light_ambient
                    + clamped_phase / frag_to_light.magnitude * light_color )
                  * color;
              }
            else
              {
                /// SOLUTION -- Problem 2
                //
                //  Instead of computing lighted color for this fragment
                //  interpolate the vertices' lighted colors.
                //
                lighted_color =
                  bc0 * lighted_colors[i]
                  + bc1 * lighted_colors[i+1]
                  + bc2 * lighted_colors[i+2];
              }

            //
            // Write the frame (color) buffer with the lighted color
            //
            frame_buffer[ idx ] = lighted_color.int_rgb();
            frame_buffer.n_px_frame++; // Count of number of written pixels.
          }
    }
  return *this;
}


int
main(int argc, char **argv)
{
  pFrame_Buffer demo_frame_buffer(argc,argv);
  WRender render( demo_frame_buffer );
  World our_world( demo_frame_buffer );
  WRENDER_INSERT(render, our_world);
  render.run();
  return 0;
}