#include <stdio.h>
#define MAIN_INCLUDE
#include <vhelper.h>
#include <coord.h>
#include <colors.h>
#include <vutil-pipeline.h>
#include <vutil-texture.h>

#include "shapes.h"
#include "vutil-raytrace.h"

#include "v3-common.h"

class World
{
public:
  World(pVulkan_Helper& vh_)
    :vh(vh_),ff_state(vh.qs),raytrace(vh.qs),
     shapes(ff_state), sphere(ff_state), cone(ff_state),
     rt_geo_strip(), rt_geo_cube(), rt_geo_sphere(),
     transform(vh.qs)
  {
  };

  World() = delete;

  pVulkan_Helper& vh;
  VFixed_Function_State_Manager ff_state;
  VRaytrace raytrace;

  Shapes shapes;
  Sphere sphere;
  Cone cone;
  RT_Shader_Set ss_alt, ss_sphere;
  RT_Geo_Package rt_geo_strip, rt_geo_cube, rt_geo_sphere;
  VPipeline pipe, pipe2;

  bool opt_want_raytrace;
  bool opt_tryout1, opt_tryout2, opt_tryout3;
  float opt_tryoutf;
  bool opt_strip_motion;
  pMatrix box_xform;

  /// Textures
  VTexture tex_img, tex_img2, tex_img_sphere;
  vk::Sampler sampler;

  void setup_and_run();
  void record(vk::CommandBuffer& cb);
  void pre_device_destroy();

  pVariable_Control variable_control;

  void animate();
  void cb_keyboard();

  double time_start_s, time_world_s;
  double time_frame_prev_s;
  deque<double> ftimes;
  float opt_omega;
  int n_slices, n_slices_setup;
  int n_hcycles, n_hcycles_setup;

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

  pCoor eye_location;
  pVect eye_direction;

  pCoor box_location;

  vk::Extent2D extent_used;
  bool global_transform_stale;
  VTransform transform;
  VBufferV<Uni_Lighting> uni_light, *puni_light;
  VVertex_Buffer_Set bufset, bset_strip, bset_sphere;
};

void
World::setup_and_run()
{
  vh.display_cb_set( [&](){ animate(); } );
  vh.want_raytrace().init(); // inst, phys, logical dev, win, surf, render_pass
  vh.opt_record_every_time = true;

  global_transform_stale = true;
  opt_want_raytrace = false;
  opt_tryout1 = opt_tryout2 = opt_tryout3 = false;
  opt_strip_motion = true;

  vh.cbs_cmd_record.push_back([&](vk::CommandBuffer& cb){ record(cb); });

#if 0
  vk::ApplicationInfo ai("My App", 0, "My Engine", 0, VK_MAKE_VERSION(1,2,0));
  vh.instance_create(ai);

  // Skip physical device.  Could choose which GPU.
  // Skip logical device. Could choose: # of queues. Phys device features.

  // Window and surface.
  vh.fb_create( 640, 480, "my window" );

  // Render pass. Could choose formats. Stencil?

  VBufferSet<pos,color,normal,tcord> rpx(qs);
  pipe1 = pipe_make(rpx,triangles);

  pipe3 = pipe_make<pos,color,normal,triangles>(qs);
  bs3 = pipe3.bs_make;

  VBuffer vbuf = vh.vtx_buffer_create( umm_vtx_color );

  pipe1 = vh.pipe_new( triangles, vtx, norm, color, tex );

#endif

  /// Sampler
  //
  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,
        false, // anisotropyEnable,
        16.0f,
        false,
        vk::CompareOp::eNever,
        0.0f, VK_LOD_CLAMP_NONE,  // min and max LOD
        vk::BorderColor::eFloatOpaqueBlack } );

  string syl_path("../gpup/gpup.png");

  P_Image_Read image_syl(syl_path,255);
  tex_img2.init(vh.qs,image_syl);

  P_Image_Read image_sp("../vulkan/mult.png");
  tex_img_sphere.init(vh.qs,image_sp);
  sphere.texture_set(sampler,tex_img_sphere);

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

  time_frame_prev_s = time_start_s = time_wall_fp();
  time_world_s = 0;

  opt_pause = false;

  n_slices = 300;
  n_hcycles = 10;
  n_slices_setup = 0;
  n_hcycles_setup = 0;

  variable_control.insert(n_slices,"N Slices",50,5);
  variable_control.insert(n_hcycles,"N H Cycles");

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

  opt_omega = .2;
  variable_control.insert(opt_omega,"Rotation Rate");

  opt_tryoutf = 1.0;
  variable_control.insert(opt_tryoutf,"Tryout F");

  light_location = pCoor(-1.5,0,2.4);
  eye_location = pCoor(0,3,7.5);
  eye_direction = pNorm(0,-.3,-1);
  opt_move_item = MI_Eye;
  box_location = pCoor(0,0,0);

  pColor black(0,0,0,0);
  uni_light->cgl_LightModel.ambient = pColor(.1,.1,.1,1);
  uni_light->cgl_LightSource[0].position = light_location;
  uni_light->cgl_LightSource[0].diffuse = pColor(1,1,1,1);
  uni_light->cgl_LightSource[0].ambient = black;
  uni_light->cgl_LightSource[0].specular = black;
  uni_light->cgl_LightSource[0].constantAttenuation = .3;
  uni_light->cgl_LightSource[0].linearAttenuation = 1;
  uni_light->cgl_LightSource[0].quadraticAttenuation = 0;

  bset_strip.usage_set
    ( vk::BufferUsageFlagBits::eVertexBuffer
      | vk::BufferUsageFlagBits::eStorageBuffer );
  bufset.usage_set
    ( vk::BufferUsageFlagBits::eVertexBuffer
      | vk::BufferUsageFlagBits::eStorageBuffer );
  bset_sphere.usage_set
    ( vk::BufferUsageFlagBits::eVertexBuffer
      | vk::BufferUsageFlagBits::eStorageBuffer ).init(vh.qs);

  sphere.transform = &transform;
  puni_light = &uni_light;
  sphere.ppuni_light = &puni_light;
  sphere.init(40);

  // Ray Tracing

  if ( vh.have_raytrace )
    {
      raytrace.init();

      raytrace.buf_light_set(uni_light);
      raytrace.buf_uni_common->opt_tryout =
        ivec4( opt_tryout1, opt_tryout2, 0, 0);
      raytrace.buf_uni_common->opt_tryoutf = vec4( opt_tryoutf, 0, 0, 0 );
      raytrace.buf_uni_common.to_dev();

      raytrace.inited = true;
      raytrace.stale_as = raytrace.stale_ds = true;

      ss_alt.init().path("v3-shdr-rt.cc").closest_hit("main_alt");
      rt_geo_strip
        .init(raytrace).shaders(ss_alt)
        .color_set( color_white, color_red );

      ss_sphere.init().path("rt-shdr-sphere.cc").closest_hit().intersection();
      rt_geo_sphere
        .init(raytrace)
        .shaders(ss_sphere)
        .color_set( color_lsu_spirit_gold, color_red );
      bset_sphere.setup_pos().setup_rot();

      rt_geo_cube.init(raytrace);
    }

  vh.message_loop_spin();
  pre_device_destroy();
  vh.finish();
}

void
World::animate()
{
  cb_keyboard();
  if ( vh.have_raytrace ) vh.opt_want_raytrace_now = opt_want_raytrace;
}

void
World::record(vk::CommandBuffer& cb)
{
  const double tnow = time_wall_fp();
  const double delta_t = tnow - time_frame_prev_s;
  if ( !opt_pause ) time_world_s += delta_t;
  time_frame_prev_s = tnow;
  const double time_elapsed_s = tnow - time_start_s;
  ftimes.push_back(tnow);
  const int n_samples = 40;

  const double time_now = time_wall_fp();
  const bool blink_visible = int64_t(time_now*3) & 1;
# define BLINK(txt,pad) ( blink_visible ? txt : pad )

  double delta_th = opt_omega * time_world_s * 3.1415926535;
  vh.fbprintf("%s\n",vh.frame_timer.frame_rate_text_get());
  vh.fbprintf("Every Time: %s\n",
                       vh.opt_record_every_time ? "YES" : "NO ");
  vh.fbprintf
    ("Want Raytrace: %s ('r')  RT Ready: %s  "
     "Tryout 1: %s  ('y')  Tryout 2: %s  ('Y') T3: %s ('z')\n",
     opt_want_raytrace ? "YES" : "NO ",
     !vh.have_raytrace ? "CAN'T" : raytrace.rt_ready ? "YES" : "NO ",
     opt_tryout1 ? BLINK("ON ","   ") : "OFF",
     opt_tryout2 ? BLINK("ON ","   ") : "OFF",
     opt_tryout3 ? BLINK("ON ","   ") : "OFF");

  vh.fbprintf
    ("Eye: %.1f,%.1f,%.1f %.1f,%.1f,%.1f  "
     "Light: %.1f,%.1f,%.1f  Box: %.1f,%.1f,%.1f\n",
     eye_location.x, eye_location.y, eye_location.z,
     eye_direction.x, eye_direction.y, eye_direction.z,
     light_location.x, light_location.y, light_location.z,
     box_location.x, box_location.y, box_location.z);
  if ( ftimes.size() > n_samples )
    {
      const double tf = ftimes.front(); ftimes.pop_front();
      const double rate = n_samples / ( tnow - tf );
      vh.gstroke.fb_printf(" %6.2f FPS\n", rate);
    }

  vh.fbprintf
    ("t = %5.3f s  world %5.3f s\n", time_elapsed_s, time_world_s);
  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());

  /// Global Transformation Matrices
  //
  if ( extent_used != vh.s_extent ) global_transform_stale = true;
  if ( global_transform_stale )
    {
      pMatrix_Rotation_Shortest m_eye_turn(eye_direction,pVect(0,0,-1));
      const float aspect = float(vh.s_extent.width) / vh.s_extent.height;
      pMatrix_Translate m_eye_loc(-eye_location);
      pMatrix_Frustum frust(-1,1,-1/aspect,1/aspect,1,5000);

      pMatrix mvglo = m_eye_turn * m_eye_loc;

      extent_used = vh.s_extent;
      global_transform_stale = false;
      transform.eye_from_global_set( mvglo );
      transform.clip_from_eye_set( frust );
    }


  /// Lighting
  //
  uni_light->cgl_LightSource[0].diffuse =
    pColor(1,1,1,1) * opt_light_intensity;
  uni_light->cgl_LightSource[0].position =
    transform.eye_from_global * light_location;
  uni_light.to_dev();


  /// The Cube
  //
  if ( !pipe )
    {
      string name("../vulkan/mult.png");
      P_Image_Read image(name);
      tex_img.init(vh.qs,image);

      pipe.init(ff_state)
        .lighting_on()
        .use_uni_light(uni_light)
        .color_uniform_back_set( color_red )
        .topology_set( vk::PrimitiveTopology::eTriangleList )
        .use_texture(sampler,tex_img)
        .create();

      pVect vx(1,0,0), vy(0,1,0), vz(0,0,1);
      vector<pVect> dims = {vx,vy,vz,vx,vy};
      vector<pColor> colors =
      //  {color_white, color_white, color_white};
      {color_lsu_spirit_gold, color_lsu_business_purple, color_white };

      bufset.reset(pipe);

      for ( int d=0; d<3; d++ )
        {
          for ( int s: {1,-1} )
            {
              pVect lz = dims[d] * s;
              pVect lx = dims[d+1] * s;
              pVect ly = dims[d+2];
              pVect n = lz;

              pCoor p00 = - ( lz + lx + ly );
              pCoor p10 = p00 + 2 * lx;
              pCoor p01 = p00 + 2 * ly;
              pCoor p11 = p10 + 2 * ly;
              pTCoor t00{1,1}, t01{1,0}, t10{0,1}, t11{0,0};
              pColor color = colors[d];
              bufset << p11 << t11 << color << n;
              bufset << p10 << t10 << color << n;
              bufset << p00 << t00 << color << n;
              bufset << p01 << t01 << color << n;
              bufset << p11 << t11 << color << n;
              bufset << p00 << t00 << color << n;
            }
        }

      bufset.to_dev();
    }

  pMatrixRows rot
    ( { cosf(delta_th), 0, sinf(delta_th)  },
      { 0,              1, 0               },
      { -sinf(delta_th), 0, cosf(delta_th) } );
  pMatrix_Translate box_loc(box_location);
  box_xform = box_loc * rot;

  transform.global_from_local_mult_for( box_xform, pipe );
  if ( cb )
    pipe.record_draw(cb,bufset);


  //
  /// The Spheres
  //

  const int n_spheres = 15;
  float sp_rad = 0.2;
  const float sp_ring_radius = 2;

  const float delta_theta = 2 * M_PI / n_spheres;
  sphere.color = color_lsu_spirit_gold;

  // Check whether a new set of sphere AABB's should be prepared.
  //
  const bool fill_sphere_aabbs =
    opt_want_raytrace && raytrace.inited &&
    ( opt_tryout3 || !rt_geo_sphere.buf_aabb.size() );

  if ( fill_sphere_aabbs )
    {
      rt_geo_sphere.buf_aabb.clear();
      bset_sphere.reset();
    }

  for ( int i=0; i<n_spheres; i++ )
    {
      const float theta = i * delta_theta;

      // Sphere location in "box" local space.
      pCoor sp_loc_l =
        pCoor( sp_ring_radius * sinf(theta), 0, sp_ring_radius * cosf(theta) );

      // Sphere location in global space.
      pCoor sp_loc_g = box_xform * sp_loc_l;

      // Sphere orientation.
      pMatrix_Rotation rot( pVect(0,1,0), theta );

      sphere.render(cb, sp_rad, sp_loc_g, rot);

      if ( !fill_sphere_aabbs ) continue;

      pCoor sp_loc = opt_tryout3 ? sp_loc_g : sp_loc_l;

      // Append a bounding box for the sphere.
      //
      // This will be read by the Vulkan code that constructs the acceleration
      // structure.
      //
      rt_geo_sphere << pAABB( sp_loc, sp_rad );

      // Append the sphere's location and radius.
      //
      // This will be read by the intersection shader and hit shader.
      //
      bset_sphere << pCoor(sp_loc,sp_rad) << rot;
      //
      // Note: The w component is set to the sphere radius for convenience.
      // It is not actually the w component of the sphere's center.
    }

  if ( fill_sphere_aabbs )
    {
      rt_geo_sphere.buf_aabb.to_dev();
      bset_sphere.to_dev();
    }


  //
  /// The Strip
  //

  // Initialize the Strip's Rasterization Pipeline, if necessary.
  //
  if ( !pipe2 )
    {
      pipe2.init(ff_state)
        .topology_set( vk::PrimitiveTopology::eTriangleStrip )
        .use_texture(sampler,tex_img2)
        .color_uniform_back_set( color_dark_red )
        .color_uniform_front_set( color_dim_gray )
        .use_uni_light(uni_light)
        .create();
    }

  // Initialize or Update the Strip's bset (Pipeline Input Buffer Set)
  //
  if ( n_slices != n_slices_setup || n_hcycles != n_hcycles_setup )
    {
      n_slices_setup = n_slices;
      n_hcycles_setup = n_hcycles;
      bset_strip.reset(pipe2);
      raytrace.stale_as = raytrace.stale_ds = true;

      const float delta_theta = 2 * M_PI / n_slices;
      const float delta_eta = delta_theta * n_hcycles;
      pCoor ctr(0,0,0);
      const float r0 = 1.5, r1 = 1.5;
      pVect ax = pVect(1,0,0), ay = pVect(0,1,0), az = pVect(0,0,1);
      pVect vy = 0.2 * ay;
      for ( int s=0; s<=n_slices; s++ )
        {
          const float theta = -s * delta_theta;
          const float eta = s * delta_eta;
          const float r = r0 + r1 + r1 * sinf(eta);
          pVect n0 = ax * cosf(theta) + az * sinf(theta);
          pNorm dpdt =
            ( -n_hcycles * r1 * cosf(eta) ) * n0 +
            r * ( -ax * sinf(theta) + az * cosf(theta) );
          pNorm n = cross(dpdt,ay);
          pCoor p1 = ctr + r * n0;

          bset_strip << n << pTCoor(s,0) << p1 + vy;
          bset_strip << n << pTCoor(s,1) << p1;
        }
      bset_strip.to_dev();

      vh.cmd_buffers_stale = true; // Ask after class.
    }

  if ( opt_strip_motion )
    {
      // Change the position of the strip just along the y axis.
      // This is done to provide an easy case for acceleration structure
      // update.

      float omega = 5 * 2 * M_PI;
      float dy = 2 * .01 * sinf(omega * time_elapsed_s );
      for ( pCoor& c: bset_strip.pos.vals ) c.y += dy;
      bset_strip.pos.to_dev();
    }

  transform.use_for( pipe2 );
  if ( cb )
    pipe2.record_draw(cb, bset_strip);

  // Render Marker for Light Source
  //
  shapes.record_tetrahedron(cb,transform,light_location,0.05);
  shapes.record_tetrahedron(cb,transform,light_location+pVect(0.2,0,0),0.1);
  //
  // The shapes object will render using rasterization or ray tracing.

  // Update Ray Tracing Information
  //
  if ( opt_want_raytrace && raytrace.inited )
    {
      if ( raytrace.stale_as )
        {
          rt_geo_strip
            .bind( bset_strip, RT_Geometry_Triangle_Strip )
            .bind( VR_BUF_IDX_sampler, pipe2.imageInfo );

          rt_geo_cube
            .bind( bufset, RT_Geometry_Triangles )
            .bind( VR_BUF_IDX_sampler, pipe.imageInfo );

          rt_geo_sphere
            .bind( bset_sphere, RT_Geometry_AABBs )
            .bind( VR_BUF_IDX_sampler, sampler, tex_img_sphere );
        }

      if ( raytrace.stale_as || opt_strip_motion )
        {
          rt_geo_cube.global_from_local = box_xform;
          rt_geo_sphere.global_from_local = opt_tryout3 ? pMatrix() : box_xform;
        }

      raytrace.tform_update(transform);
      raytrace.buf_uni_common.to_dev();
    }

}


void
World::cb_keyboard()
{
  auto& ogl_helper = vh;
  const int key = ogl_helper.keyboard_key_get();
  if ( !key ) return;
  pVect adjustment(0,0,0);
  pVect user_rot_axis(0,0,0);
  const bool kb_mod_s = ogl_helper.keyboard_shift;
  const bool kb_mod_c = ogl_helper.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 'b': opt_move_item = MI_Ball; break;
  case 'B': opt_move_item = MI_Ball_V; break;
  case 'e': case 'E': opt_move_item = MI_Eye; break;
  case 'l': case 'L': opt_move_item = MI_Light; break;
  case 'm': opt_strip_motion = !opt_strip_motion; break;
  case 'p': case 'P': opt_pause = !opt_pause; break;
  case 'r': case 'R': opt_want_raytrace = !opt_want_raytrace; break;

  case 'x': raytrace.stale_as = true; break;

  case 'y':
    opt_tryout1 = !opt_tryout1;
    raytrace.buf_uni_common->opt_tryout.x = opt_tryout1;
    break;
  case 'Y':
    opt_tryout2 = !opt_tryout2;
    raytrace.buf_uni_common->opt_tryout.y = opt_tryout2;
    break;
  case 'z':
    opt_tryout3 = !opt_tryout3;
    raytrace.buf_uni_common->opt_tryout.z = opt_tryout3;
    break;
  case FB_KEY_TAB:
    if ( !kb_mod_s ) { variable_control.switch_var_right(); break; }
  case 96: variable_control.switch_var_left(); break; // `, until S-TAB works.
  case '-':case '_': variable_control.adjust_lower(); break;
  case '+':case '=': variable_control.adjust_higher(); break;

  default: printf("Unknown key, %d\n",key); break;
  }

  if ( opt_tryoutf != raytrace.buf_uni_common.get_val().opt_tryoutf.x )
    raytrace.buf_uni_common->opt_tryoutf.x = opt_tryoutf;

  // Update eye_direction based on keyboard command.
  //
  if ( user_rot_axis.x || user_rot_axis.y )
    {
      global_transform_stale = true;
      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 )
    {
      if ( opt_move_item != MI_Ball && opt_move_item != MI_Ball_V )
        global_transform_stale = true;
      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_Ball: case MI_Ball_V:
        box_location += adjustment; break;
      case MI_Light: light_location += adjustment; break;
      case MI_Eye: eye_location += adjustment; break;
      default: break;
      }
    }
}

void
World::pre_device_destroy()
{
  sphere.destroy();
  shapes.destroy();
  cone.destroy();
  pipe.destroy();
  pipe2.destroy();

  uni_light.destroy();

  bufset.destroy();
  bset_strip.destroy();
  bset_sphere.destroy();

  rt_geo_strip.destroy();
  rt_geo_cube.destroy();
  rt_geo_sphere.destroy();

  raytrace.destroy();

  vh.dev.destroySampler(sampler);
  tex_img.destroy();
  tex_img2.destroy();
  tex_img_sphere.destroy();
}


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

  world.setup_and_run();

  return 0;
}