/// LSU EE 4702-1 (Fall 2024), GPU Programming
//

 /// More Shaders

/// Purpose
//
//   Demonstrate use of Geometry Shaders
//   Demonstrates use of buffer storage objects.
//   See demo-10-shdr-simple.cc and demo-10-geo-code.cc for shader
//     program source code.

 ///  Note: Requires Vulkan 1.2

/// References
//
// :ogl45: OpenGL Specification Version 4.5
//         http://www.opengl.org/registry/doc/glspec45.compatibility.pdf
//
// :sl45:  The OpenGL Shading Language - Language Version 4.5
//         http://www.opengl.org/registry/doc/glspec45.compatibility.pdf

#if 0
/// Background

 /// New Material in Demo 10
 //

 /// VPipeline::record_draw_indexed
//
//
//   Like VPipeline::record_draw, except that an index array is used.




 /// Accessing Storage Buffers in Shader Programs
 //

//   Within shader code, buffer objects accessed using same syntax
//   as arrays in the language C.
//
     vertex_o.xyz = helix_coord[hidx].xyz + radius * gl_Normal;


//   The declaration of buffer objects in shader code is much different
//   than C.
//
     layout ( binding = 7 ) buffer Helix_Coord  { vec4  helix_coord[];  };



 /// Interpolation Qualifiers
 //
 //  :sl45: Section 4.5
 //
 //  Used on fragment shader inputs and vertex or geometry shader outputs.
 //
 //  They specify how variables written by vertex or geometry shader ..
 //  .. should be interpolated to obtain values at input to fragment shader.
 //
 //  Options:
 //
 //    smooth:
 //       Linearly interpolate based on object-space coordinates.
 //       The default for FP types, produces best looking results.
 //
 //    noperspective:
 //       Linearly interpolate based on window-space coordinates.
 //       Won't look as good, but computationally less demanding.
 //
 //    flat:
 //       Don't interpolate, instead use value of provoking (last)
 //       vertex. The default for integer types.
 //
 //  For an example, see demo-10-shdr-simple.cc.


 /// Geometry Shader Coding
//
//   Overview
//
//     The geometry shader reads in one primitive ...
//     ... and writes zero or more primitives.
//
//     One possible application:
//       Split input triangle into many smaller triangles so that
//       object appears smoother, *if necessary*.  "If necessary"
//       means that such a geometry shader will use window-space
//       coordinates to decide whether the input triangle needs to
//       be split.  For example, if the maximum distance between
//       the primitive's vertices is less than 10 pixels than don't
//       split.
//

//   Major Steps to use a Geometry Shader

//     -- Declare type of input primitive.

//   Geometry Shader Input
//
//     Geometry Shader Input Type.  (OGL 4.5 Section 4.4)
//
//       Declared in shader code using layout qualifier

layout ( TYPE_OF_PRIMITIVE ) in;

//       where TYPE_OF_PRIMITIVE can be:
//         points, lines, lines_adjacency, triangles, triangles_adjacency.
//
//       Adjacency indicates that nearby vertices are also included.
//
//       Common use:

layout ( triangles ) in;

//     Geometry Shader Input Size
//
//      The inputs to the geometry shader are all arrays ..
//      .. the array size is determined by the primitive type:
//
//         Layout               Array Size
//         points               1
//         lines                2
//         lines_adjacency      4
//         triangles            3
//         triangles_adjacency  6

//      Geometry Shader Input Declaration
//
//        The geometry shader inputs ..
//        .. MUST match the vertex shader outputs ..
//        .. plus an array dimension.

//        Example:


out Data  // Vertex shader output
{
  int hidx;
  vec3 normal_o;
  vec4 color;
};

in Data  // Geometry shader input.
{
  int hidx;
  vec3 normal_o;
  vec4 color;
} In[];


//   Geometry Shader Output
//     Usually triangles or triangle strips.
//     DOES NOT have to match input primitive to current vertex shader.
//
//     Output is a single vertex, declared of course as scalars.

layout ( triangle_strip, max_vertices = 4 ) out;

//     Other output types: points, line_strip.


out Data_GF
{
  vec3 var_normal_e;
  vec4 var_vertex_e;
  flat vec4 color;
};




#endif

///  Keyboard Commands
 //
 /// Object (Eye, Light, Ball) Location or Push
 //   Arrows, Page Up, Page Down
 //   Will move object or push ball, depending on mode:
 //   'e': Move eye.
 //   'l': Move light.
 //
 /// Eye Direction
 //   Home, End, Delete, Insert
 //   Turn the eye direction.
 //   Home should rotate eye direction up, End should rotate eye
 //   down, Delete should rotate eye left, Insert should rotate eye
 //   right.  The eye direction vector is displayed in the upper left.

 /// Simulation Options
 //  (Also see variables below.)
 //
 //  's'    Switch between different shaders in forward direction.
 //  'S'    Switch between different shaders in reverse direction.
 //  'F11'  Change size of text.
 //  'F12'  Write screenshot to file.

 /// Variables
 //   Selected program variables can be modified using the keyboard.
 //   Use "Tab" to cycle through the variable to be modified, the
 //   name of the variable is displayed next to "VAR" on the bottom
 //   line of green text.

 //  'Tab' Cycle to next variable.
 //  '`'   Cycle to previous variable.
 //  '+'   Increase variable value.
 //  '-'   Decrease variable value.
 //
 //  VAR Light Intensity - The light intensity.
 //      Helix Segs Per Rev - The number of segments in one revolution of helix.
 //         A smaller number means fewer primitives.
 //      Wire Segs Per Rev - The number of segments in 1 revolution around wire.
 //         A smaller number means fewer primitives.


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

#include <vstroke.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"

enum Shader_Program
  { SP_Plain, SP_Vtx_Animation, SP_Geo_Shade,
    SP_Geo_Animation, SP_ENUM_SIZE };
const char* const shader_program[] =
  { "SP_Plain", "SP_Vtx_Animation", "SP_Geo_Shade",
    "SP_Geo_Animation", "SP_ENUM_SIZE" };


struct Buldge_Info
{
  float bulge_pos;
  float bulge_dist_thresh;
  float wire_radius;
  bool opt_texture;
};


class World {
public:
  World(pVulkan_Helper &fb)
    :vh(fb), ff_state(fb.qs), frame_timer(vh.frame_timer), shapes(ff_state),
     transform(fb.qs)
  {init();}
  World() = delete;
  void init();
  void run();
  void render(vk::CommandBuffer& cb);
  void cb_keyboard();
  void modelview_update();

  // Class providing utilities, such as showing text.
  //
  pVulkan_Helper& vh;
  VFixed_Function_State_Manager ff_state;
  pFrame_Timer& frame_timer;    // Class for showing frame timing.
  Shapes shapes;

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

  pCoor light_location;
  float opt_light_intensity;
  VBufferV<Uni_Lighting> uni_light;

  pCoor helix_location;
  float helix_radius;   // Radius of helix.
  float wire_radius;    // Radius of wire forming helix.
  int seg_per_helix_revolution;
  int seg_per_wire_revolution;

  float opt_bulge_dist_thresh;

  bool coords_stale;

  VVertex_Buffer_Set bset_helix, bset_lonely;
  VBufferVV<int32_t> buf_wire_surface_indices;
  VBufferVV<pCoor> buf_helix_coords; // Storage Buffer

  VBufferV<Buldge_Info> buf_bulge_info;

  int opt_shader;

  VPipeline pipe_plain;          // Phong shading.
  VPipeline pipe_vtx_animation;  // Compute animation using vertex shader.
  VPipeline pipe_geo_shade;      // Geometry shader to color adjacent triangles.
  VPipeline pipe_geo_animation;  // Geometry shader to help with animation.

  enum { MI_Eye, MI_Light, MI_Ball, MI_Ball_V, MI_COUNT } opt_move_item;

  bool opt_sim_paused;
  double last_evolve_time;
  double sim_time;

  pCoor eye_location;
  pVect eye_direction;
  pMatrix modelview;

  VTransform transform;

  vk::Sampler sampler;
  VTexture texture_id_syllabus;
  bool opt_texture;
};

void
World::init()
{
  vh.init();
  vh.opt_record_every_time = true;
  vh.display_cb_set([&](){});
  vh.cbs_cmd_record.push_back( [&](vk::CommandBuffer& cb){ render(cb); });

  opt_sim_paused = false;
  sim_time = 0;
  last_evolve_time = time_wall_fp();
  frame_timer.work_unit_set("Steps / s");

  coords_stale = true;

  seg_per_helix_revolution = 40;
  seg_per_wire_revolution = 20;
  opt_bulge_dist_thresh = 2;
  variable_control.insert(seg_per_helix_revolution,"Helix Seg Per Rev");
  variable_control.insert(seg_per_wire_revolution,"Wire Seg Per Rev");
  variable_control.insert(opt_bulge_dist_thresh,"Buldge Distance");

  eye_location = pCoor(2.6,5.7,15);
  eye_direction = pVect(0,0,-1);

  opt_light_intensity = 1.5;
  light_location = pCoor(7,4.0,-0.3);
  uni_light.init(vh.qs,vk::BufferUsageFlagBits::eUniformBuffer);


  helix_location = pCoor(0,0,-5);
  helix_radius = 5;
  wire_radius = 0.5;

  variable_control.insert(opt_light_intensity,"Light Intensity");

  opt_move_item = MI_Eye;

  texture_id_syllabus.init(vh.qs, P_Image_Read("gpup.png",255) );

  opt_texture = true;
  sampler = vh.qs.dev.createSampler
    ( { {},
        vk::Filter::eLinear, vk::Filter::eLinear,
        vk::SamplerMipmapMode::eLinear,
        // Also: eRepeat, eMirroredRepeat, eClampToEdge, etc.
        vk::SamplerAddressMode::eRepeat,
        vk::SamplerAddressMode::eRepeat,
        vk::SamplerAddressMode::eRepeat,
        0.0f,
        true, // anisotropyEnable,
        16.0f,
        false,
        vk::CompareOp::eNever,
        0.0f, VK_LOD_CLAMP_NONE,  // min and max LOD
        vk::BorderColor::eFloatOpaqueBlack } );

  buf_bulge_info.init
    (vh.dev_phys, vh.dev, vk::BufferUsageFlagBits::eUniformBuffer );

  buf_helix_coords.init
    (vh.dev_phys, vh.dev, vk::BufferUsageFlagBits::eStorageBuffer );

  buf_wire_surface_indices.init
    (vh.dev_phys, vh.dev, vk::BufferUsageFlagBits::eIndexBuffer );

  pipe_plain
    .init( vh.qs ).ds_use_texture().ds_set_material().ds_use( uni_light )
    .ds_uniform_use( "BIND_BULGE_INFO", buf_bulge_info )
    .shader_inputs_info_set<pCoor,pNorm,pTCoor>()
    .shader_code_set
    ("demo-10-shdr-ngeo.cc",  // File holding shader program.
     "vs_main_basic(); ",     // Name of vertex shader main routine.
     "fs_main_phong();"       // Name of fragment shader main routine.
     )
    .topology_set( vk::PrimitiveTopology::eTriangleStrip )
    .create();

  pipe_vtx_animation
    .init( vh.qs ).ds_use_texture().ds_set_material().ds_use( uni_light )
    .ds_uniform_use( "BIND_BULGE_INFO", buf_bulge_info )
    .ds_storage_follow( "BIND_HELIX_COORDS", buf_helix_coords )
    .shader_inputs_info_set<int,pNorm,pTCoor>()
    .shader_code_set
    ("demo-10-shdr-ngeo.cc",  // File holding shader program.
     "vs_main_helix();",      // Name of vertex shader main routine.
     "fs_main_phong();"       // Name of fragment shader main routine.
     )
    .topology_set( vk::PrimitiveTopology::eTriangleStrip )
    .create();

  pipe_geo_shade
    .init( vh.qs ).ds_use_texture().ds_set_material().ds_use( uni_light )
    .ds_uniform_use( "BIND_BULGE_INFO", buf_bulge_info )
    .ds_storage_follow( "BIND_HELIX_COORDS", buf_helix_coords )
    .shader_inputs_info_set<int,pNorm,pTCoor>()
    .shader_code_set
    ("demo-10-shdr-simple.cc", // File holding shader program.
     "vs_main_helix();",       // Name of vertex shader main routine.
     "gs_main_helix();",       // Name of geometry shader main routine.
     "fs_main_phong();"        // Name of fragment shader main routine.
     )
    .topology_set( vk::PrimitiveTopology::eTriangleStrip )
    .create();

  pipe_geo_animation
    .init( vh.qs ).ds_use_texture().ds_set_material().ds_use( uni_light )
    .ds_uniform_use( "BIND_BULGE_INFO", buf_bulge_info )
    .ds_storage_follow( "BIND_HELIX_COORDS", buf_helix_coords )
    .shader_inputs_info_set<int,pNorm,pTCoor>()
    .shader_code_set
    ("demo-10-shdr-geo.cc",   // File holding shader program.
     "vs_main_helix();",      // Name of vertex shader main routine.
     "gs_main_helix();",      // Name of geometry shader main routine.
     "fs_main_phong();"       // Name of fragment shader main routine.
     )
    .topology_set( vk::PrimitiveTopology::eTriangleStrip )
    .create();

  opt_shader = SP_Geo_Animation;

  modelview_update();
}

void
World::modelview_update()
{
  pMatrix_Translate center_eye(-eye_location);
  pMatrix_Rotation rotate_eye(eye_direction,pVect(0,0,-1));
  modelview = rotate_eye * center_eye;
}

void
World::run()
{
  vh.message_loop_spin();

  bset_helix.destroy();
  bset_lonely.destroy();
  buf_wire_surface_indices.destroy();
  buf_helix_coords.destroy();
  buf_bulge_info.destroy();

  pipe_plain.destroy();
  pipe_vtx_animation.destroy();
  pipe_geo_shade.destroy();
  pipe_geo_animation.destroy();

  uni_light.destroy();

  texture_id_syllabus.destroy();

  vh.dev.destroySampler(sampler);
  shapes.destroy();
  transform.destroy();

  vh.finish();
}

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

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

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

  vh.fbprintf("%s\n",frame_timer.frame_rate_text_get());

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

  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());


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

  vh.fbprintf("Active Shader Program: %s  (s TO CHANGE)\n",
                      shader_program[opt_shader]);

  const int win_width = vh.get_width();
  const int win_height = vh.get_height();
  const float aspect = float(win_width) / win_height;
  const float n_dist = 0.01;
  const float xr = .8 * n_dist;

  transform.eye_from_global_set( modelview );

  // 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
  //
  auto& light0 = uni_light->cgl_LightSource[0];
  uni_light->cgl_LightModel.ambient = color_white * 0.5;
  light0.diffuse = opt_light_intensity * color_white;
  light0.position = transform.eye_from_global * light_location;
  light0.constantAttenuation = 0;
  light0.linearAttenuation = 1;
  light0.quadraticAttenuation = 0.05;
  light0.ambient = color_black;
  uni_light.to_dev();

  ///
  /// Paint Single Triangle.
  ///

  pColor color_tri(0x7815b6); // Red, Green, Blue

  if ( !bset_lonely )
    {
      bset_lonely.reset( pipe_plain );

      // Specify vertices for a triangle.
      //
      pCoor p1( 9,    6, -7 );
      pCoor p2( 0,    5, -3 );
      pCoor p3( 9.5, -5, -1.2 );
      pNorm triangle_normal = cross(p3,p2,p1);

      bset_lonely << triangle_normal << pTCoor(0.90,1.0) << p1
                  << triangle_normal << pTCoor(0.00,0.9) << p2
                  << triangle_normal << pTCoor(0.95,0.0) << p3;
      bset_lonely.to_dev();
    }

  pipe_plain.ds_set( color_tri, color_red );
  pipe_plain.ds_use(sampler, texture_id_syllabus );

  pipe_plain.ds_set( transform ).record_draw( cb, bset_lonely );

  ///
  /// Construct a Helix
  ///

  if ( coords_stale || !bset_helix )
    {
      // Recompute helix coordinates, etc.

      coords_stale = false;

      bset_helix.reset(vh.qs)
        .setup_pos().setup_normal().setup_tcoor().setup_int1();
      buf_wire_surface_indices.clear();
      buf_helix_coords.clear();

      // Number of times helix wraps around.
      const int revolutions_per_helix = 6;

      const int segments_per_helix =
        revolutions_per_helix * seg_per_helix_revolution;

      const double delta_eta = 2 * M_PI / seg_per_helix_revolution;
      const double delta_y = 4 * wire_radius / seg_per_helix_revolution;
      const double delta_theta = 2 * M_PI / seg_per_wire_revolution;

      for ( int i = 0; i < segments_per_helix; i++ )
        {
          const bool last_i_iteration = i + 1 == segments_per_helix;
          const double eta = i * delta_eta;
          const float cos_eta = cosf(eta);
          const float sin_eta = sinf(eta);

          pCoor p0( helix_radius * cos_eta,
                    i * delta_y,
                    helix_radius * sin_eta);

          buf_helix_coords << p0;

          // Compute axes for drawing surface.
          //
          pVect ax( -wire_radius * cos_eta, 0, -wire_radius * sin_eta);

          pNorm ay( delta_eta * helix_radius * -sin_eta,
                    delta_y,
                    delta_eta * helix_radius * cos_eta );

          pVect az = cross(ax,ay);

          for ( int j = 0; j < seg_per_wire_revolution; j++ )
            {
              const int idx = bset_helix.size();
              const float theta = j * delta_theta;

              pVect norm0 = cosf(theta) * ax + sinf(theta) * az;

              // Compute surface coordinate. This computation can also
              // be done by the vertex shader, when that feature is
              // turned on the two lines below are not needed.
              //
              pCoor s0 = p0 + norm0;
              bset_helix << s0;

              // Insert additional attributes into the buffer set.
              //
              bset_helix << norm0.normal() << i << pTCoor(theta/(2*M_PI),2*eta);

              if ( last_i_iteration ) continue;

              // Insert indices for triangle with one vertex on eta.
              buf_wire_surface_indices << idx << idx + seg_per_wire_revolution;
            }
        }

      bset_helix.to_dev();
      buf_wire_surface_indices.to_dev();
      buf_helix_coords.to_dev();
    }

  /// Switch to selected set of shader programs.
  //
  VPipeline& pipe =
    opt_shader == SP_Plain         ? pipe_plain         :
    opt_shader == SP_Vtx_Animation ? pipe_vtx_animation :
    opt_shader == SP_Geo_Shade     ? pipe_geo_shade     :
    opt_shader == SP_Geo_Animation ? pipe_geo_animation : pipe_geo_animation;

  ///
  /// Paint a Helix
  ///

  if ( !opt_sim_paused )
    {
      const double now = time_wall_fp();
      sim_time += now - last_evolve_time;
      last_evolve_time = now;
    }

  const float bulge_pos = fmodf( sim_time * 5.0f, buf_helix_coords.size());

  buf_bulge_info->bulge_pos = bulge_pos;
  buf_bulge_info->bulge_dist_thresh = opt_bulge_dist_thresh;
  buf_bulge_info->wire_radius = wire_radius;
  buf_bulge_info->opt_texture = opt_texture;
  buf_bulge_info.to_dev();

  pMatrix tform =
    pMatrix_Translate(helix_location) * pMatrix_Rotation(pVect(0,1,0),M_PI/3);

  pipe.ds_use(sampler, texture_id_syllabus );
  pipe.ds_set( transform * tform );
  pipe.ds_set( color_lsu_spirit_gold, color_red );

  pipe.record_draw_indexed( cb, bset_helix, buf_wire_surface_indices );

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


void
World::cb_keyboard()
{
  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 's':
    opt_shader++; if ( opt_shader == SP_ENUM_SIZE ) opt_shader = 0;
    break;
  case 'S':
    if ( opt_shader == 0 ) opt_shader = SP_ENUM_SIZE;
    opt_shader--;
    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':
    if ( opt_sim_paused ) last_evolve_time = time_wall_fp();
    opt_sim_paused = !opt_sim_paused;
    break;

  case 't': case 'T': opt_texture = !opt_texture;  coords_stale=true; break;
  case ' ': opt_sim_paused = true; sim_time += 0.1; 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();  coords_stale=true; break;
  case '+':case '=': variable_control.adjust_higher(); coords_stale=true; 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);
      modelview_update();
    }

  // 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: helix_location += adjustment; break;
      default: break;
      }
      modelview_update();
    }
}


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

  world.run();

  return 0;
}