/*
 * knot.cpp - Knot plotting example scene.
 *
 * (C) Copyright 2013-2022 John E. Stone
 * SPDX-License-Identifier: BSD-3-Clause
 *
 * $Id: knot.cpp,v 1.8 2022/04/04 05:01:24 johns Exp $
 *
 */

/**
 *  \file knot.cpp
 *  \brief Knot plotting example scene.
 *
 */

//
// Knot plotting example
// John E. Stone, Dec 2021
//

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include "ProfileHooks.h"

#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.h"

#include "TachyonOptiX.h"

#if defined(TACHYON_USEPINNEDMEMORY)
#include <cuda_runtime.h>
#endif


// http://katlas.org/wiki/36_Torus_Knots
// https://www.maa.org/sites/default/files/images/upload_library/23/stemkoski/knots/page5.html
float * pq_torus_knot(int p, int q, int points) {
  float * coords = (float *) malloc(points * 3 * sizeof(float));

  int i;
  for (i=0; i<points; i++) {
    float theta = 2.0f * M_PI * float(i) / float(points);
    float r = 2.0f + cos(p * theta);
    int ind = 3*i;
    coords[ind    ] = r * cos(q * theta); // X
    coords[ind + 1] = r * sin(q * theta); // Y
    coords[ind + 2] = -sin(p * theta);    // Z
  }
 
  return coords; 
}


void HSItoRGB(float h, float s, float i, float &r, float &g, float &b) {
  float t=2.0f * M_PI * h;
  float scale = i / 2.0f;
  r=(1.0f + s*sin(t - 2.0f*M_PI/3.0f)) * scale;
  g=(1.0f + s*sin(t)) * scale;
  b=(1.0f + s*sin(t + 2.0f*M_PI/3.0f)) * scale;
}


float *knot_colors(int numpts) {
  float *colors = (float *) malloc(numpts * 3 * sizeof(float));

  for (int i=0; i<numpts; i++) {
    int ind = i * 3;
    float h = float(i) / float(numpts);
    float s = 1.0f;
    float l = (fmodf(h*40.0f, 1.0f) > 0.5f) ? 0.85f : 0.35f;
    HSItoRGB(h, s, l, colors[ind], colors[ind + 1], colors[ind + 2]);
  }

  return colors;
}


//
// Draw the knot geometry using an array of spheres
//
void gen_knot(TachyonOptiX *rt, float *coords, float *colors, 
              int numpts, float radius, int mat) {

  SphereArray spheres;
  spheres.center.resize(numpts);
  spheres.radius.resize(numpts);
  spheres.primcolors3f.resize(numpts);

  float3 *verts = spheres.center.data();
  float *radii = spheres.radius.data();
  float3 *cols = spheres.primcolors3f.data();

  memcpy(verts, coords, numpts * 3 * sizeof(float));
  memcpy(cols, colors, numpts * 3 * sizeof(float));
  for (int i=0; i<numpts; i++)
    radii[i]=radius;

  rt->add_spherearray(spheres, mat);
}


//
// Draw a quad or triangle mesh for the floor
//
void gen_floor(TachyonOptiX *rt, float width, float height, float length, int mat) {
  float vertex[] = {
      -100.0f,  4.0f, -100.0f,
      -100.0f,  4.0f,  100.0f,
       100.0f,  4.0f, -100.0f,
       100.0f,  4.0f,  100.0f
  };
  float color[] = {
      1.0f, 1.0f, 1.0f, 1.0f,
      1.0f, 1.0f, 1.0f, 1.0f,
      1.0f, 1.0f, 1.0f, 1.0f,
      1.0f, 1.0f, 1.0f, 1.0f
  };
  int index[] = {
      0, 1, 2,                                  // triangle-1
      1, 2, 3                                   // triangle-2
  };

  vertex[ 0] = -width / 2.0f;
  vertex[ 3] = -width / 2.0f;
  vertex[ 6] =  width / 2.0f;
  vertex[ 9] =  width / 2.0f;

  vertex[ 1] = height;
  vertex[ 4] = height;
  vertex[ 7] = height;
  vertex[10] = height;

  vertex[ 2] = -length / 2.0f;
  vertex[ 5] =  length / 2.0f;
  vertex[ 8] = -length / 2.0f;
  vertex[11] =  length / 2.0f;

  TriangleMesh mesh;
  mesh.vertices.resize(4);
  mesh.vertcolors3f.resize(4);
  mesh.indices.resize(2*3);

  float3 *verts = mesh.vertices.data();
  float3 *cols = mesh.vertcolors3f.data();
  int3 *indices = mesh.indices.data();

  memcpy(verts, vertex, 4 * 3 * sizeof(float));
  memcpy(cols, color, 4 * 3 * sizeof(float));
  memcpy(indices, index, 2 * 3 * sizeof(int));

  rt->add_trimesh(mesh, mat);
}


int main(int argc, const char **argv) {
  // some sane defaults
  int imgSize[2] = {4096, 4096 }; // W x H
  int numpts = 1000;
  int p = 5;
  int q = 3; 

  //
  // parse args
  //
  if (argc == 1) {
    printf("Usage: %s P Q\n", argv[0]);
    printf("  Knot params P and Q control the knot type generated.\n");
    printf("Using defaults, P=5 and Q=3\n");
    // return -1;
  }

  if (argc > 2) {
    p = atoi(argv[1]);
    q = atoi(argv[2]);
  }
  if (argc > 3) {
    numpts = atoi(argv[3]);
  }


  //
  // camera defaults
  //
#if 0
  float cam_pos[] =  {-2.0f,   0.0f,   -8.0f};   // look at origin from -Z
  float cam_up[] =   {0.0f,   1.0f,    0.0f};	// Y-up
  float cam_view[3];

  // look at origin
  float invlen = 1.0f / sqrtf(cam_pos[0]*cam_pos[0] +
                              cam_pos[1]*cam_pos[1] +
                              cam_pos[2]*cam_pos[2]);
  cam_view[0] = -cam_pos[0] * invlen;
  cam_view[1] = -cam_pos[1] * invlen;
  cam_view[2] = -cam_pos[2] * invlen;
#endif



  //
  // Build scene
  //
  printf("Calculating knot, p: %d, q: %d, points: %d\n", p, q, numpts);
  float *coords = pq_torus_knot(p, q, numpts);
  float *colors = knot_colors(numpts);


  PROFILE_PUSH_RANGE("Initialize Tachyon", 0);
  printf("Initializing TachyonOptiX...");

  /// static methods for querying OptiX-supprted GPU hardware independent
  /// of whether we actually have an active context.
  unsigned int devcount = TachyonOptiX::device_count();
  unsigned int optixversion = TachyonOptiX::optix_version();

  printf("Found %u OptiX devices\n", devcount);
  printf("OptiX version used for build: %d.%d.%d (%u)\n",
         optixversion/10000,
         (optixversion%10000)/100,
         (optixversion%100),
         optixversion);

  TachyonOptiX *rt = new TachyonOptiX();
  PROFILE_POP_RANGE();

  PROFILE_PUSH_RANGE("Build Scene", 0);

  // create and setup camera
  rt->framebuffer_colorspace(RT_COLORSPACE_sRGB);
  rt->framebuffer_resize(imgSize[0], imgSize[1]);
//  rt->set_verbose_mode(TachyonOptiX::RT_VERB_MIN);
//  rt->set_verbose_mode(TachyonOptiX::RT_VERB_DEBUG);
  float rtbgcolor[] = { 1.0, 1.0, 1.0 };
  float rtbggradtopcolor[] = { 0.6, 0.0, 0.0 };
  float rtbggradbotcolor[] = { 0.0, 0.0, 0.6 };

  rt->set_bg_color(rtbgcolor);
  rt->set_bg_color_grad_top(rtbggradtopcolor);
  rt->set_bg_color_grad_bot(rtbggradbotcolor);

  float bggradient[] = { 0.0f, 1.0f, 0.0f };
  rt->set_bg_gradient(bggradient);
  rt->set_bg_gradient_topval(1.0f);
  rt->set_bg_gradient_botval(-1.0f);

//  rt->set_bg_mode(TachyonOptiX::RT_BACKGROUND_TEXTURE_SOLID);
rt->set_bg_mode(TachyonOptiX::RT_BACKGROUND_TEXTURE_SKY_SPHERE);
//  rt->set_bg_mode(TachyonOptiX::RT_BACKGROUND_TEXTURE_SKY_SPHERE);

  rt->set_aa_samples(16);
  rt->shadows_enable(1);

  //
  // Optionally add AO
  //
#if 1
  rt->set_ao_samples(16);
  rt->set_ao_ambient(0.9);
  rt->set_ao_direct(0.2);
  rt->set_ao_maxdist(100.2);
#endif

  rt->camera_dof_enable(0);


  float lightdir0[] = { -0.5f, 0.5f, -1.0f };
  float lightcolor0[] = { 1.0f, 1.0f, 1.0f };
  rt->add_directional_light(lightdir0, lightcolor0);

  // set camera params
  rt->set_camera_type(TachyonOptiX::RT_PERSPECTIVE);
//  rt->set_camera_type(TachyonOptiX::RT_ORTHOGRAPHIC);

  float campos[3] = {0.0f, 10.0f, 12.0f};
//  float camU[3]   = {1.0f,  0.0f,  0.0f};
  float camV[3]   = {0.0f,  1.0f,  0.0f};
//  float camW[3]   = {0.0f,  0.0f,  1.0f};
  float camat[3]  = {0.0f,  0.0f,  0.0f};

  rt->set_camera_pos(campos);
//  rt->set_camera_ONB(camU, camV, camW);
  rt->set_camera_lookat(camat, camV);

  rt->set_camera_zoom(0.5f);
  rt->set_camera_dof_fnumber(64.0f);
  rt->set_camera_dof_focal_dist(0.7f);
  // set stereoscopic display parameters
  rt->set_camera_stereo_eyesep(0.6f);
  rt->set_camera_stereo_convergence_dist(10.0f);

  // set depth cueing parameters
  float start = 1.0f;
  float end = 30.0f;
  float density = 0.33f;
//    rt->set_cue_mode(TachyonOptiX::RT_FOG_LINEAR, start, end, density);
//    rt->set_cue_mode(TachyonOptiX::RT_FOG_EXP, start, end, density);
//    rt->set_cue_mode(TachyonOptiX::RT_FOG_EXP2, start, end, density);
  rt->set_cue_mode(TachyonOptiX::RT_FOG_NONE, start, end, density);


  //
  // Set the material rendering parameters
  //
  int mat = 0;
  int mat2 = 1;
  float ambient = 0.0f;
  float diffuse = 0.7f;
  float specular = 0.0f;
  float shininess = 0.0f;
  float reflectivity = 0.0f;
  float opacity = 1.0f;
  float outline = 0.0f;
  float outlinewidth = 0.0f;
  int transmode = 0;

  rt->add_material(ambient, diffuse, specular, shininess, reflectivity, 
                   opacity, outline, outlinewidth, transmode, mat);

  rt->add_material(ambient, diffuse, specular, shininess, reflectivity, 
                   opacity, outline, outlinewidth, transmode, mat2);
  PROFILE_POP_RANGE();

  rt->set_verbose_mode(TachyonOptiX::RT_VERB_MIN);
#if 1
  PROFILE_PUSH_RANGE("Renderer Warmup Passes", 0);
  // force warmup passes on an empty scene so our timings of subsequent
  // scene data are way more realistic
  for (int w=0; w<100; w++) {
    rt->render();
    rt->framebuffer_clear();
  }
  PROFILE_POP_RANGE();
#endif
  rt->set_verbose_mode(TachyonOptiX::RT_VERB_TIMING);

  PROFILE_PUSH_RANGE("Generate Scene", 0);
  gen_knot(rt, coords, colors, numpts, 0.33f, mat);
  gen_floor(rt, 200.0f, -3.5f, 200.0f, mat2);
  PROFILE_POP_RANGE();


  char filename[1024];
  sprintf(filename, "knot-%d-%d.png", p, q);
  printf("rendering 100 accumulated frames to '%s'...\n", filename);

  PROFILE_PUSH_RANGE("Render Scene", 0);

  // render 100 accumulated frames
//  const char *statestr = "|/-\\.";
//  int state=0;
//  for (int frames = 0; frames < 100; frames++,state=(state+1) & 3) {

    rt->render();

//    printf("%c\r", statestr[state]); 
//    fflush(stdout);
//  }
//  printf("\n");

  PROFILE_POP_RANGE();

  rt->print_raystats_info();

  PROFILE_PUSH_RANGE("Write Output Image", 0);

  // access frame and write its content as PNG file
  if (filename != NULL) {
    rt->framebuffer_get_size(imgSize[0], imgSize[1]);
    size_t bufsz = imgSize[0] * imgSize[1] * sizeof(int);
    unsigned char *rgb4u = (unsigned char *) calloc(1, bufsz);
    rt->framebuffer_download_rgb4u(rgb4u);

#if 0
    if (writealpha) {
printf("Writing rgba4u alpha channel output image 2\n");
      if (write_image_file_rgba4u(filename, rgb4u, imgSize[0], imgSize[1]))
        printf("Failed to write image '%s'!!\n", filename);
    } else {
      if (write_image_file_rgb4u(filename, rgb4u, imgSize[0], imgSize[1]))
        printf("Failed to write image '%s'!!\n", filename);
    }
#else
    stbi_flip_vertically_on_write(1);

    stbi_write_png(filename, imgSize[0], imgSize[1], 4, rgb4u, imgSize[0] * sizeof(int));
#endif

    free(rgb4u);
  }

  free(coords);
  free(colors);

  delete rt;

  PROFILE_POP_RANGE();

  return 0;
}


