/// LSU EE 4702-1 (Fall 2022), GPU Programming
//
 /// Homework 4 -- SOLUTION
//
//   Most of the code is based on Homework 3 and Midterm Exam Problem 1b
//   
//   Handout: https://www.ece.lsu.edu/koppel/gpup/2022/hw04.pdf


// Include files provided for this course.
//
#define MAIN_INCLUDE
#include <vhelper.h>

#include <gp/coord.h>
#include <gp/pstring.h>
#include <gp/misc.h>
#include <gp/colors.h>

#include <vutil-texture.h>
#include <vutil-pipeline.h>
#include "shapes.h"

struct Uni_Light_Simple {
  vec4 position;
  vec4 color;
};

struct Uni_Misc {
  ivec4 opt_tryout;
};


constexpr int ncolors = 10;
struct MT_Colors {
  pColor front[ncolors], back[ncolors];
};


struct Fan_Info {
  pCoor p1, p2, p3, p4;
  int n_blades;

  pCoor pos_axis_top;
  pCoor pos_axis_bottom;
  pVect axis_to_cyl;
  float outer_radius;

  pNorm az;

  pColor blade_0_front_upper, blade_0_front_lower;
  pColor blade_0_back_upper, blade_0_back_lower;
  pColor blade_i_front_upper, blade_i_front_lower;
  pColor blade_i_back_upper, blade_i_back_lower;

  int serial;
};


class World {
public:
  World(pVulkan_Helper &vh)
    :vh(vh),ff_state(vh.qs),frame_timer(vh.frame_timer),shapes(ff_state),
     transform(vh.qs){}
  void setup_and_run();
  void render(vk::CommandBuffer& cb);
  void keyboard_handle();

  // Class providing utilities, such as showing text.
  //
  pVulkan_Helper& vh;
  VFixed_Function_State_Manager ff_state;

  // Class for showing frame timing.
  //
  pFrame_Timer& frame_timer;

  Shapes shapes;

  // Class for easy keyboard control of variables.
  //
  pVariable_Control variable_control;

  pCoor light_location;
  float opt_light_intensity;
  enum { MI_Eye, MI_Light, MI_Ball, MI_Ball_V, MI_COUNT } opt_move_item;

  pCoor eye_location, eye_initial_location;
  pVect eye_direction;
  pNorm eye_initial_direction;

  VTransform transform;
  VBufferV<Uni_Light_Simple> uni_light_simple;
  VBufferV<Uni_Misc> uni_misc;
  VPipeline pipe_plain;
  VVertex_Buffer_Set bset_plain;

  bool opt_tryout1, opt_tryout2;  // For ad-hoc experiments.

  /// Homework 4

  Fan_Info fan_info;
  vector<pMatrix> fans_xforms;
  void fan_setup(bool rand);
  void scene_setup_1();
  void scene_setup_2();
  void scene_setup_3();
  double world_time, time_last_update;
  bool opt_pause;

  VPipeline pipe_gear_base;
  VPipeline pipe_gear_hw04;
  VVertex_Buffer_Set bset_mta;
  VVertex_Buffer_Set bset_gear_list;
  VVertex_Buffer_Set bset_gear_strip;

  VBufferV<MT_Colors> uni_colors;
  int bset_p2_serial;

  bool opt_tri_strip;
  bool opt_tri_shader;

  int opt_n_blades;
  int colors_serial;

};




void
World::scene_setup_1()
{
  fan_setup(true);
}

void
World::scene_setup_2()
{
  fan_setup(false);
}

void
World::scene_setup_3()
{
  // Used for the assignment screenshots.
  vh.clearValues[0].color =
    vk::ClearColorValue( array<float, 4>( { { 0.4, .4, .4, 1.0 } } ) );
  eye_location = pCoor(1.0,0.1,2.1);
  eye_direction = pVect(.37,-.19,-.91);
  opt_light_intensity = 14.51;
  fan_setup(false);
}

float
prand(float min, float max)
{
  double r = random() * (1.0 / RAND_MAX );
  return min + r * ( max - min );
}

void
World::fan_setup(bool rand)
{
  world_time = 0;
  time_last_update = time_wall_fp();
  fan_info.serial++;

  auto prand = [=](float min, float max)
  { return rand ? ::prand(min,max) : min + (max-min)/2; };

  int n_fans = rand ? prand(7,14) : 4;

  float arr_r = 3.5;
  pCoor pos_ref = eye_initial_location
    + ( 20 + 10 ) * eye_initial_direction + pVect(0,0,0);
  pNorm ax = cross( pVect(0,1,0), eye_initial_direction );
  pNorm ay(0,1,prand(0,.4));
  pNorm az = cross(ax,ay);

  float cyl_ht = rand ? min( 2.0, M_PI * arr_r / n_fans ) : 4;
  float delta_theta = 2 * M_PI / n_fans;

  fan_info.p4 = fan_info.pos_axis_bottom = pCoor(0,0,0);
  fan_info.az = pNorm(0,1,0);
  fan_info.p1 = fan_info.pos_axis_top =
    fan_info.pos_axis_bottom + fan_info.az * cyl_ht;
  float inner_r = arr_r;
  fan_info.axis_to_cyl = pVect( inner_r,0,0);
  fan_info.p2 = fan_info.p1 + fan_info.axis_to_cyl;
  fan_info.outer_radius = inner_r * 1.3;
  fan_info.p3 = fan_info.p1 + 1.3 * fan_info.axis_to_cyl;
  fan_info.n_blades = opt_n_blades;

  if ( rand && false )
    {
    }
  else
    {
      fan_info.blade_0_front_upper = color_lsu_spirit_purple;
      fan_info.blade_0_front_lower = color_lsu_business_purple;
      fan_info.blade_0_back_upper = color_lsu_spirit_gold;
      fan_info.blade_0_back_lower = color_gold;
      fan_info.blade_i_front_upper = color_aquamarine;
      fan_info.blade_i_front_lower = color_aquamarine;
      fan_info.blade_i_back_upper = color_salmon;
      fan_info.blade_i_back_lower = color_salmon;
    }

  fans_xforms.clear();

  for ( int i=0; i<n_fans; i++ )
    {
      float theta = M_PI + i * delta_theta;
      pCoor pos_axis_top =
        pos_ref  + az * 10 * cosf(theta) + ax * 10 * sin(theta);
      pCoor pos_axis_bottom = pos_axis_top - ay * cyl_ht;
      pMatrix_Translate trans_to_global(pos_axis_bottom);
      pMatrix_Cols rot_to_global(ax,ay,az);
      fans_xforms << trans_to_global * rot_to_global;
    }
}


void
World::setup_and_run()
{
  // Setup Vulkan context, etc.
  vh.init();
  vh.display_cb_set([&](){});
  vh.cbs_cmd_record.push_back( [&](vk::CommandBuffer& cb){ render(cb); });

  vh.clearValues[0].color =
    vk::ClearColorValue( array<float, 4>( { { 0.4, 0.4, 0.4, 1.0 } } ) );

  opt_tryout1 = opt_tryout2 = false;
  eye_initial_location = pCoor( 3, .5, 8.6) ;
  eye_location = eye_initial_location;
  eye_initial_direction = pVect(0,0,-1);
  eye_direction = eye_initial_direction;

  opt_n_blades = 16;
  variable_control.insert(opt_n_blades,"Num Blades",1,3);

  opt_light_intensity = 13.1;
  light_location = pCoor( -5, 3.2, -2.7 );
  variable_control.insert(opt_light_intensity,"Light Intensity");

  opt_move_item = MI_Eye;

  uni_misc.init(vh.qs,vk::BufferUsageFlagBits::eUniformBuffer);
  uni_misc->opt_tryout = ivec4(opt_tryout1,opt_tryout2,0,0);

  uni_light_simple.init(vh.qs,vk::BufferUsageFlagBits::eUniformBuffer);

  fan_info.serial = 0;

  uni_colors.init(vh.qs,vk::BufferUsageFlagBits::eUniformBuffer);
  bset_p2_serial = 0;

  opt_pause = false;

  opt_tri_strip = false;
  opt_tri_shader = false;
  colors_serial = -1;

  scene_setup_2();

  // Start the graphics. The function below does not return until the
  // user exits by closing the window.
  //
  vh.message_loop_spin();
  //
  // At this point the user exited and so it's time to
  // clean up.

  uni_light_simple.destroy();
  uni_misc.destroy();
  pipe_plain.destroy();
  pipe_gear_base.destroy();
  pipe_gear_hw04.destroy();
  bset_plain.destroy();
  bset_mta.destroy();
  bset_gear_list.destroy();
  bset_gear_strip.destroy();
  shapes.destroy();
  transform.destroy();

  uni_colors.destroy();

  vh.finish();
}

void
World::render(vk::CommandBuffer& cb)
{
  // This routine called whenever window needs to be updated.

  // Get any waiting keyboard commands.
  //
  keyboard_handle();

  /// Frame Buffer Informational Messages
  //
  //  Print messages using utility functions provided for this course.
  //
  vh.fbprintf("%s\n",frame_timer.frame_rate_text_get());
  const double time_now = time_wall_fp();
  const bool blink_visible = int64_t(time_now*2) & 1;
# define BLINK(txt,pad) ( blink_visible ? txt : pad )
  vh.fbprintf
    ("Compiled: %s\n",
#ifdef __OPTIMIZE__
     "WITH OPTIMIZATION"
#else
     BLINK("WITHOUT OPTIMIZATION","")
#endif
     );

  vh.fbprintf
    ("Eye location: [%5.1f, %5.1f, %5.1f]  "
     "Eye direction: [%+.2f, %+.2f, %+.2f]\n",
     eye_location.x, eye_location.y, eye_location.z,
     eye_direction.x, eye_direction.y, eye_direction.z);

  vh.fbprintf
    ("Light location: [%5.1f, %5.1f, %5.1f]\n",
     light_location.x, light_location.y, light_location.z);

  vh.fbprintf
    ("Topology: %s ('v')  Shader: %s ('V')  "
     "Tryout 1: %s  ('y')  Tryout 2: %s  ('Y')  \n",
     opt_tri_strip ? "TRI STRIP" : "TRI LIST ",
     opt_tri_shader ? "HW04 " : "BASE",
     opt_tryout1 ? BLINK("ON ","   ") : "OFF",
     opt_tryout2 ? BLINK("ON ","   ") : "OFF");


  pVariable_Control_Elt* const cvar = variable_control.current;
  vh.fbprintf("VAR %s = %.5f  (TAB or '`' to change, +/- to adjust)\n",
                      cvar->name,cvar->get_val());


  // -------------------------------------------------------------------------
  ///
  /// Specification of Transformation Matrices
  ///

  pMatrix eye_from_global =
    pMatrix_Rotation(eye_direction,pVect(0,0,-1))
    * pMatrix_Translate(-eye_location);

  transform.eye_from_global_set( eye_from_global );

  /// Setup Projection Transformation:  Eye Space -> Clip Space
  //
  const int win_width = vh.s_extent.width;
  const int win_height = vh.s_extent.height;
  const float aspect = float(win_width) / win_height;
  const float n_dist = 0.01;
  const float xr = .8 * n_dist;

  // Frustum: left, right, bottom, top, near, far
  transform.clip_from_eye_set
    ( pMatrix_Frustum( -xr, xr,                // left, right
                       -xr/aspect, xr/aspect,  // bottom, top
                       n_dist, 5000            // near, far 
                       ) );

  /// Lighting
  //
  uni_light_simple->color = color_white * opt_light_intensity;
  uni_light_simple->position = transform.eye_from_global * light_location;
  uni_light_simple.to_dev();

  uni_misc.to_dev();

  //
  // -------------------------------------------------------------------------

  if ( !pipe_plain )
    pipe_plain
      .init( vh.qs )
      .ds_uniform_use( "BIND_LIGHT_SIMPLE", uni_light_simple )
      .ds_uniform_use( "BIND_MISC", uni_misc )
      .shader_inputs_info_set<pCoor,pNorm,pColor>()
      .shader_code_set
      ("demo-03-shdr.cc", "vs_main();", nullptr, "fs_main();")
      .topology_set( vk::PrimitiveTopology::eTriangleList )
      .create();

  const auto prim_want = opt_tri_strip
    ? vk::PrimitiveTopology::eTriangleStrip
    : vk::PrimitiveTopology::eTriangleList;

  if ( pipe_gear_base && pipe_gear_base.p_in_asm_ci.topology != prim_want )
    pipe_gear_base.destroy();

  if ( !pipe_gear_base )
    pipe_gear_base
      .init( vh.qs )
      .ds_uniform_use( "BIND_LIGHT_SIMPLE", uni_light_simple )
      .ds_uniform_use( "BIND_MISC", uni_misc )
      .ds_uniform_use( "BIND_HW04", uni_colors )
      .shader_inputs_info_set<pCoor,pNorm,int>()
      .shader_code_set
      ("gear-base-shdr.cc",
       "vs_main_list();", "gs_main_list();", "fs_main_list();")
      .topology_set( prim_want )
      .create();

  if ( pipe_gear_hw04 && pipe_gear_hw04.p_in_asm_ci.topology != prim_want )
    pipe_gear_hw04.destroy();


  /// Homework 4 -- Problem 2
  //
  // Modify the pipeline initialization below so that unneeded data
  // is not used for vertex shader inputs.

  if ( !pipe_gear_hw04 )
    pipe_gear_hw04
      .init( vh.qs )
      .ds_uniform_use( "BIND_LIGHT_SIMPLE", uni_light_simple )
      .ds_uniform_use( "BIND_MISC", uni_misc )
      .ds_uniform_use( "BIND_HW04", uni_colors )
      /// SOLUTION
      // Remove pNorm and int from shader_inputs.
      .shader_inputs_info_set<pCoor>()
      .shader_code_set
      ("hw04-shdr-sol.cc",
       "vs_main_hw04();", "gs_main_hw04();", "fs_main_hw04();")
      .topology_set( prim_want )
      .create();


  // Reset buffers in this buffer set.
  //
  bset_plain.reset(pipe_plain);

  // Prepare a triangle.
  //
  pCoor p0 = { 9, 0,  -10 };
  pCoor p1 = { 18, 6, -4 };
  pCoor p2 = { 9, 5, -9 };
  pNorm tri_norm = cross( p0, p1, p2 );
  bset_plain << p0 << p1 << p2;
  bset_plain << color_green << color_red << color_blue;
  bset_plain << tri_norm << tri_norm << tri_norm;

  // Add a square consisting of a red and green triangle.
  //
  bset_plain << pCoor(-2,0,-4) << pCoor(-4,2,-4) << pCoor(-4,0,-4);
  bset_plain << color_red << color_red << color_red;
  bset_plain << pCoor(-2,0,-4) << pCoor(-2,2,-4) << pCoor(-4,2,-4);
  bset_plain << color_green << color_green << color_green;
  pNorm snorm = cross( pCoor(-4,0,-4), pCoor(-4,2,-4), pCoor(-2,0,-4) );
  bset_plain << snorm << snorm << snorm << snorm << snorm << snorm;

  bset_plain.to_dev();
  pipe_plain.ds_set( transform );
  pipe_plain.record_draw(cb, bset_plain);

  // Update theta_0, the angle by which the gear should be at.
  //
  if ( !opt_pause ) world_time += time_now - time_last_update;
  time_last_update = time_now;
  const float omega = 1;
  const float theta_0 = world_time * omega;

  if ( opt_n_blades != fan_info.n_blades )
    {
      fan_info.n_blades = opt_n_blades;
      fan_info.serial++;
    }

  const auto& f = fan_info;

  if ( colors_serial != f.serial )
    {
      colors_serial = f.serial;
      uni_colors->front[0] = f.blade_0_front_upper;
      uni_colors->back[0] = f.blade_0_back_upper;
      uni_colors->front[1] = f.blade_0_front_lower;
      uni_colors->back[1] = f.blade_0_back_lower;
      uni_colors->front[2] = f.blade_i_front_upper;
      uni_colors->back[2] = f.blade_i_back_upper;
      uni_colors->front[3] = f.blade_i_front_lower;
      uni_colors->back[3] = f.blade_i_back_lower;
      uni_colors.to_dev();
    }

  VPipeline& pipe = opt_tri_shader ? pipe_gear_hw04 : pipe_gear_base;

  // Would be nice to set stale to true only when bset needs an update.
  bool gear_list_stale = true;
  bool gear_strip_stale = true;

  if ( !opt_tri_strip && gear_list_stale )
    {
      gear_list_stale = false;
      bset_gear_list.reset(pipe);

      pCoor p1(f.p1),  p2(f.p2),  p3(f.p3),  p4(f.p4);
      float n_pieces = f.n_blades;

      pVect vx_inner(p1,p2), vx_outer(p1,p3), vz(p4,p1);
      pNorm az(vz);
      pVect vy_inner = cross( az, vx_inner );
      pVect vy_outer = cross( az, vx_outer );

      float delta_theta = 2 * M_PI / n_pieces;
      float delta_theta_h = delta_theta / 2;

      /// Homework 4 -- Do not modify this code.

      pVect v_prev = vx_inner;

      for ( int i=1; i<=2*n_pieces; i++ )
        {
          float theta = i * delta_theta_h;
          pVect v_curr =
            i & 1
            ? vx_outer * cosf(theta) + vy_outer * sinf(theta)
            : vx_inner * cosf(theta) + vy_inner * sinf(theta);

          pCoor pp_bot = p4 + v_prev;
          pCoor pc_bot = p4 + v_curr;
          pCoor pp_top = p1 + v_prev;
          pCoor pc_top = p1 + v_curr;

          // Emit Coordinates
          bset_gear_list << pc_top << pp_top << pp_bot; // Upper Triangle
          bset_gear_list << pc_bot << pc_top << pp_bot; // Lower Triangle

          // Emit Normals
          pNorm np = cross( vz, v_curr - v_prev );
          bset_gear_list << np << np << np << np << np << np;

          // Emit Colors
          if ( i < 3 ) bset_gear_list << 0 << 0 << 0 << 1 << 1 << 1;
          else         bset_gear_list << 2 << 2 << 2 << 3 << 3 << 3;

          v_prev = v_curr;
        }

      bset_gear_list.to_dev();
    }

  if ( opt_tri_strip && gear_strip_stale )
    {
      gear_strip_stale = false;
      bset_gear_strip.reset(pipe);

      pCoor p1(f.p1),  p2(f.p2),  p3(f.p3),  p4(f.p4);
      float n_pieces = f.n_blades;

      pVect vx_inner(p1,p2), vx_outer(p1,p3), vz(p4,p1);
      pNorm az(vz);
      pVect vy_inner = cross( az, vx_inner );
      pVect vy_outer = cross( az, vx_outer );

      float delta_theta = 2 * M_PI / n_pieces;
      float delta_theta_h = delta_theta / 2;

      /// Homework 4 -- Problem 1
      //
      //  Modify the code below so that the vertices in
      //  bset_gear_strip describe the gear for a triangle
      //  strip topology.

      pVect v_prev = vx_inner;

      for ( int i=0; i<=2*n_pieces; i++ )
        {
          float theta = i * delta_theta_h;
          pVect v_curr =
            i & 1
            ? vx_outer * cosf(theta) + vy_outer * sinf(theta)
            : vx_inner * cosf(theta) + vy_inner * sinf(theta);

          pCoor pc_top = p1 + v_curr;
          pCoor pc_bot = p4 + v_curr;

          /// SOLUTION

          // Emit Colors
          if ( bset_gear_strip.int1.selected )
            {
              if ( i < 2 ) bset_gear_strip << 0 << 1;
              else         bset_gear_strip << 2 << 3;
            }
          else
            {
              /// SOLUTION - Pack color index into w component.
              if ( i < 2 ) { pc_top.w = 0; pc_bot.w = 1; }
              else         { pc_top.w = 2; pc_bot.w = 3; }
            }

          // Emit Coordinates
          bset_gear_strip << pc_top << pc_bot;

          if ( bset_gear_strip.normal.selected )
            {
              // Emit Normals, If Needed
              pNorm np = cross( vz, v_curr - v_prev );
              bset_gear_strip << np << np;
            }
          v_prev = v_curr;
        }

      bset_gear_strip.to_dev();
    }

  pMatrix_Rotation rot(f.az,theta_0);

  // Apply the rotation before the global-from-local transform (fans_xforms[0])
  //
  pipe.ds_set( transform * fans_xforms[0] * rot );
  pipe.record_draw(cb, opt_tri_strip ? bset_gear_strip : bset_gear_list );

  // Render Marker for Light Source
  //
  shapes.record_tetrahedron(cb,transform,light_location,0.2);
}



void
World::keyboard_handle()
{
  const int key = vh.keyboard_key_get();
  if ( !key ) return;
  pVect adjustment(0,0,0);
  pVect user_rot_axis(0,0,0);
  const bool kb_mod_s = vh.keyboard_shift;
  const bool kb_mod_c = vh.keyboard_control;
  const float move_amt = kb_mod_s ? 2.0 : kb_mod_c ? 0.08 : 0.4;

  switch ( key ) {
  case FB_KEY_LEFT: adjustment.x = -move_amt; break;
  case FB_KEY_RIGHT: adjustment.x = move_amt; break;
  case FB_KEY_PAGE_UP: adjustment.y = move_amt; break;
  case FB_KEY_PAGE_DOWN: adjustment.y = -move_amt; break;
  case FB_KEY_DOWN: adjustment.z = move_amt; break;
  case FB_KEY_UP: adjustment.z = -move_amt; break;
  case FB_KEY_DELETE: user_rot_axis.y = 1; break;
  case FB_KEY_INSERT: user_rot_axis.y =  -1; break;
  case FB_KEY_HOME: user_rot_axis.x = 1; break;
  case FB_KEY_END: user_rot_axis.x = -1; break;
  case '1': scene_setup_1(); break;
  case '2': scene_setup_2(); break;
  case '3': scene_setup_3(); break;
  case 'b': case 'B': opt_move_item = MI_Ball; break;
  case 'e': case 'E': opt_move_item = MI_Eye; break;
  case 'l': case 'L': opt_move_item = MI_Light; break;
  case 'p': case 'P': opt_pause = !opt_pause; break;
  case 'v': opt_tri_strip = !opt_tri_strip; break;
  case 'V': opt_tri_shader = !opt_tri_shader; break;
  case 'y': opt_tryout1 = !opt_tryout1;
    uni_misc->opt_tryout.x = int(opt_tryout1);
    break;
  case 'Y': opt_tryout2 = !opt_tryout2;
    uni_misc->opt_tryout.y = int(opt_tryout2);
    break;
  case FB_KEY_TAB:
    if ( !kb_mod_s ) { variable_control.switch_var_right(); break; }
  case 96: variable_control.switch_var_left(); break;
  case '-':case '_': variable_control.adjust_lower(); break;
  case '+':case '=': variable_control.adjust_higher(); break;
  default: printf("Unknown key, %d\n",key); break;
  }

  // Update eye_direction based on keyboard command.
  //
  if ( user_rot_axis.x || user_rot_axis.y )
    {
      pMatrix_Rotation rotall(eye_direction,pVect(0,0,-1));
      user_rot_axis *= invert(rotall);
      eye_direction *= pMatrix_Rotation(user_rot_axis, M_PI * 0.03);
    }

  // Update eye_location based on keyboard command.
  //
  if ( adjustment.x || adjustment.y || adjustment.z )
    {
      const double angle =
        fabs(eye_direction.y) > 0.99
        ? 0 : atan2(eye_direction.x,-eye_direction.z);
      pMatrix_Rotation rotall(pVect(0,1,0),-angle);
      adjustment *= rotall;

      switch ( opt_move_item ){
      case MI_Light: light_location += adjustment; break;
      case MI_Eye: eye_location += adjustment; break;
      //  case MI_Ball: object_location += adjustment; break;
      default: break;
      }
    }
}


int
main(int argv, char **argc)
{
  pVulkan_Helper pvulkan_helper(argv,argc);
  World world(pvulkan_helper);

  world.setup_and_run();

  return 0;
}