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

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

  ///
 /// 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;

    /// SOLUTION -- Problem 2
    topology_strip = false;
    /// SOLUTION -- Problem 3
    vtx_normals = nullptr;
  };

  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;
  }

  /// SOLUTION -- Problem 2
  //
  // 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;
  }

  /// SOLUTION -- Problem 3
  //
  //  Specify that no vertex normals will be provided for the next
  //  draw_rasterization. Instead, lighting code will compute a
  //  triangle normal using the provided triangle vertices.
  //
  Our_3D& vtx_normals_set()
  {
    vtx_normals = nullptr;
    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;
  vector<pColor> *vtx_colors;
  bool opt_tryout1, opt_tryout2;

  /// SOLUTION -- Problem 2
  //
  // Provide a variable to hold vertex grouping setting.
  //
  bool topology_strip;

  /// SOLUTION -- Problem 3
  //
  //  Provide a variable to hold pointer to normals, if any.
  //
  pVect4 *vtx_normals;

};

 ///
 ///  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;
  pCoor *adj_location;
  pVect eye_direction;
  bool opt_strip, opt_normals;
  bool opt_tryout1,opt_tryout2;

};

void
World::init_scene()
{
  eye_location = pCoor(-1.2,.5,10.2);
  light_location = pCoor(-2.2,1.5,10.2);
  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;
  opt_normals = opt_strip = false;
}

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

  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': opt_normals = !opt_normals; break;
  case 's': case 'S': opt_strip = !opt_strip; break;
  }

  gc.frame_buffer.fbprintf
    ("Strip %s ('s')  Normals %s ('n')  Tryout1 %s ('y')  Tryout2 %s ('Y') \n",
     opt_strip ? "ON " : "OFF", opt_normals ? "ON " : "OFF",
     opt_tryout1 ? "ON " : "OFF", opt_tryout2 ? "ON " : "OFF");

  gc.light_location_set( light_location ).light_color_set( color_white * 100 );

  ///
  /// 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 in object space and list of colors.
  //
  vector<pCoor> coors_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_white << color_white << color_white;

  // Insert Second Triangle
  //
  // Add coordinates and color of another triangle
  //
  coors_os << pCoor(7,2,1) << pCoor(-3,3,-3.5) << pCoor(9,0,0);
  colors << color_lsu_spirit_purple << color_lsu_spirit_purple
         << color_lsu_spirit_purple;

  // 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_lsu_spirit_gold << color_lsu_spirit_gold
         << color_lsu_spirit_gold;

  coors_os << pCoor(-4,2,-3) << pCoor(-2,2,-3) << pCoor(-2,0,-3);
  colors << color_blue << color_green << color_red;

 // Add a strip.
  int n_segs = 20;
  pCoor center(1,-1,-3);
  pVect sz( 0, 0.75, 0.05 );
  pNorm ax( 0, -sz.z, sz.y ), ay = cross(sz,ax);
  vector<pCoor> strip_coor(n_segs+1);
  float r = 4, d_th = 2 * M_PI / n_segs;

  /// SOLUTION -- Problem 1
  //
  // Draw the red triangle. (Okay, it's orange_red. Close enough.)
  //
  coors_os << pCoor(-2,2,-3) << pCoor(-2,0,-3) << center + sz;
  colors << color_orange_red << color_orange_red << color_orange_red;

  /// SOLUTION -- Problem 2
  //
  // Tell Our_3D that in the next draws vertices will not be grouped
  // as triangle strips.
  //
  gc.topology_strip_set(false);

  /// SOLUTION -- Problem 3
  //
  // Tell Our_3D that in the next draws normals will not be provided.
  //
  gc.vtx_normals_set();

  gc.vtx_coors_set(coors_os).vtx_colors_set(colors);
  gc.draw_rasterization();
  coors_os.clear(); colors.clear();

  /// SOLUTION -- Problem 3
  //
  vector<pVect4> strip_norms(n_segs+1);

  for ( int i=0; i<=n_segs; i++ )
    {
      /// SOLUTION -- Problem 3
      //
      // Compute normal for segment i.
      //
      strip_norms[i] = ax * cosf(i*d_th) + ay * sinf(i*d_th);

      strip_coor[i] = center + r*strip_norms[i];
    }
  vector<pVect4> norms_os;

  if ( opt_strip )
    {
      for ( int i=0; i<=n_segs; i++ )
        {
          /// SOLUTION -- Problem 2
          //
          //  Insert vertex coordinates and colors in triangle-strip
          //  order. Just two!
          //
          coors_os << strip_coor[i] + sz << strip_coor[i];
          colors << color_cyan << color_cyan;

          /// SOLUTION -- Problem 3
          //
          //  Insert appropriate normal for each vertex.
          //
          norms_os << strip_norms[i] << strip_norms[i];
        }
    }
  else
    {
      for ( int i=0; i<n_segs; i++ )
        {
          coors_os << strip_coor[i] << strip_coor[i+1] << strip_coor[i] + sz;
          coors_os << strip_coor[i]+sz << strip_coor[i+1]+sz << strip_coor[i+1];

          for ( int i=0; i<6; i++ ) colors << color_cyan;

          /// SOLUTION -- Problem 3
          //
          //  Insert appropriate normal for each vertex.
          //
          norms_os << strip_norms[i] << strip_norms[i+1] << strip_norms[i];
          norms_os << strip_norms[i] << strip_norms[i+1] << strip_norms[i+1];
        }
    }

  /// SOLUTION -- Problem 3
  //
  //  Provide normals to Our_3D or tell Our_3D not to expect them.
  //
  if ( opt_normals ) gc.vtx_normals_set(norms_os); else gc.vtx_normals_set();

  /// SOLUTION -- Problem 2
  //
  //  Tell Our_3D whether the vertices are grouped for triangle strips.
  //
  gc.topology_strip_set(opt_strip);

  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;

  // Homework 3: DON'T MODIFY THIS LINE.
  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;

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

  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];
      pCoor_Homogenized
        w0( ws_from_os * o0 ), w1( ws_from_os * o1 ), w2( ws_from_os * o2 );

      pColor c0 = colors[i+0], c1 = colors[i+1], c2 = colors[i+2];

      // Compute triangle normal (for lighting).
      //
      pNorm tn = cross(o0,o1,o2);

      /// SOLUTION -- Problem 3
      //
      //  Extract normals from list.
      //
      pVect4 n0, n1, n2;
      if ( vtx_normals )
        {
          n0 = vtx_normals[i+0];
          n1 = vtx_normals[i+1];
          n2 = vtx_normals[i+2];
        }

      pVect v20(w2,w0), v21(w2,w1);
      const float db0 = 1/max(fabs(v20.x),fabs(v20.y));
      const float db1 = 1/max(fabs(v21.x),fabs(v21.y));

      // Used for computing perspective-correct barycentric coords.
      const float w0ow2z = w0.z / w2.z, w1ow2z = w1.z / w2.z;

      // Rasterize the Triangle.
      // Iterate over triangle using barycentric coordinates.
      //
      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;

            // Each iteration operates on one fragment of triangle.

            pCoor wf = b0*w0 + b1*w1 + b2*w2;

            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;

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

            zbuffer[ idx ] = zinv;

            // Compute perspective-correct barycentric coordinates.
            //
            float bc0 = b0 ? 1/( w0ow2z * (1-b0)/b0 + 1 ) : 0;
            float bc1 = b1 ? 1/( w1ow2z * (1-b1)/b1 + 1 ) : 0;
            float bc2 = 1 - bc0 - bc1;

            /// SOLUTION -- Problem 3
            //
            //  Either blend provided vertex normals or use tn (the
            //  triangle normal computed above).
            //
            pVect4 n = vtx_normals ? bc0*n0 + bc1*n1 + bc2*n2 : tn;

            // Compute Lighted Color
            //

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

            // Find approximate object-space coordinate of fragment.
            //
            pCoor of =     bc0*o0 + bc1*o1 + bc2*o2;

            pNorm f_to_l(of,light_location);
            float phase = fabs(dot(n,f_to_l));
            pColor lighted_color = phase / f_to_l.mag_sq * light_color * color;

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