/*
 * objloader.cpp - Parametric surface plotting example scene.
 *
 * (C) Copyright 2013-2022 John E. Stone
 * SPDX-License-Identifier: BSD-3-Clause
 *
 * $Id: objloader.cpp,v 1.11 2022/04/19 02:54:21 johns Exp $
 *
 */

/**
 *  \file objloader.cpp
 *  \brief Loader for Wavefront OBJ scenes, using the "tiny_obj_loader"
 *         header-only loader library.
 *     
 *         The Wavefront OBJ file format is described in detail here:     
 *           http://paulbourke.net/dataformats/mtl/
 */

//
// OBJ loader example
// John E. Stone, Mar 2022
//

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

#define TINYOBJLOADER_IMPLEMENTATION // define this in only *one* .cc
// Optional requires C++11: TINYOBJLOADER_USE_MAPBOX_EARCUT robust trinagulation
//#define TINYOBJLOADER_USE_MAPBOX_EARCUT
#include "tiny_obj_loader.h"

#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.h"
#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"

#include "TachyonOptiX.h"

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


#if 0
static int replacefileextension(char * s,
                                const char * oldextension,
                                const char * newextension) {
  int sz, extsz;
  sz = strlen(s);
  extsz = strlen(oldextension);

  if (strlen(newextension) != strlen(oldextension))
   return -1;

  if (extsz > sz)
    return -1;

//  if (strupncmp(s + (sz - extsz), oldextension, extsz)) {
  if (strncmp(s + (sz - extsz), oldextension, extsz)) {
    return -1;
  }

  strcpy(s + (sz - extsz), newextension);

  return 0;
}


static char * getmaterialfilename(const char* objfilename) {
  char *mtlfilename = strdup(objfilename);
  if (replacefileextension(mtlfilename, ".obj", ".mtl")) {
    free(mtlfilename);
    return NULL;
  }
 
  return mtlfilename;
}
#endif


static char * getmaterialdir(const char* filename) {
#if defined(_MSC_VER)
  char slash = '\\';
#else
  char slash = '/';
#endif

  char *dirname = strdup(filename);
  char *sep = strrchr(dirname, slash); 
  if (sep != NULL) {
    *sep = '\0'; // terminate after last seperator
  }

  return dirname;
}


static unsigned char * loadimage_rgba4u(const char *searchpath,
                                        const char *filename,
                                        int &xres, int &yres, int &chinfile) {
  int newlen = strlen(searchpath) + strlen(filename) + 2;
  char *fullpath = (char *) calloc(1, newlen);

  strcpy(fullpath, searchpath);
  strcat(fullpath, "/");
  strcat(fullpath, filename);

  int fplen = strlen(fullpath);
  for (int i=0; i<fplen; i++) {
    if (fullpath[i] == '\\')
      fullpath[i] = '/';  
  }

//  printf("Load: %s\n", fullpath); 
  unsigned char *img = NULL;
  chinfile=0;
  img = stbi_load(fullpath, &xres, &yres, &chinfile, 4);

  return img;
}


void print_coords(float *coords, int numpts, int count) {
  printf("Coordinate dump:\n");
  if (count > numpts)
    count = numpts;

  int i;
  for (i=0; i<count; i++) {
    int ind = i * 3;
    printf("[%d]: %.3f  %.3f  %.3f\n", 
           i, coords[ind], coords[ind+1], coords[ind+2]);

  }
  printf("...\n");
  for (i=numpts-count; i<numpts; i++) {
    int ind = i * 3;
    printf("[%d]: %.3f  %.3f  %.3f\n", 
           i, coords[ind], coords[ind+1], coords[ind+2]);
  }
}


int main(int argc, const char **argv) {
  PROFILE_INITIALIZE();

  // some sane defaults
  int imgSize[2] = {3840, 2160 }; // W x H
  const char *objfilename = NULL;
  int ambientocclusion = 1;
  int denoiser = 0;
  int dof = 0;
  int warmup = 0; 

  printf("Tachyon OBJ scene viewer demo\n");

  //
  // camera defaults
  //

#if 0
  // XXX some default camera settings for the "San Miguel" test scene.
  float campos[3] = {  22.0f,   14.0f,   10.0f};
  float camV[3]   = {   0.0f,    1.0f,    0.0f};
  float camat[3]  = {  10.0f,    0.0f,    0.0f};
  float camzoom   = 0.5f;
  float aomaxdist = 15;
#else
  // XXX some default camera settings for the "Sponza" test scene.
  float campos[3] = { 800.0f,  400.0f,    0.0f};
  float camV[3]   = {   0.0f,    1.0f,    0.0f};
  float camat[3]  = {   0.0f,  400.0f,    0.0f};
  float camzoom   = 0.5f;
  float aomaxdist = 700;
#endif

  // 
  // parse args
  //
  if (argc == 1) {
    printf("Usage: %s objfilename.obj [optional flags]\n", argv[0]);
    printf("  optional flags:  -ao: add renderer-specific AO lighting\n");
    printf("                   -res XXXX YYYY: override default image res\n");
    return -1;
  }

  // parse surface type
  if (argc > 1) {
    objfilename = argv[1];
  }

  // parse remaining optional parameter flags
  if (argc > 2) {
    for (int i=2; i<argc; i++) {
      if (!strcmp("-ao", argv[i])) {
        ambientocclusion = 1;
        printf("Enabling renderer-specific AO lighting.\n");
        continue;
      }

      if (!strcmp("-camlookat", argv[i])) {
        if ((argc - i) >= 3) {
          camat[0] = atof(argv[++i]);
          camat[1] = atof(argv[++i]);
          camat[2] = atof(argv[++i]);
        }
        continue;
      }

      if (!strcmp("-campos", argv[i])) {
        if ((argc - i) >= 3) {
          campos[0] = atof(argv[++i]);
          campos[1] = atof(argv[++i]);
          campos[2] = atof(argv[++i]);
        }
        continue;
      }

      if (!strcmp("-camup", argv[i])) {
        if ((argc - i) >= 3) {
          camV[0] = atof(argv[++i]);
          camV[1] = atof(argv[++i]);
          camV[2] = atof(argv[++i]);
        }
        continue;
      }

      if (!strcmp("-camzoom", argv[i])) {
        if ((argc - i) >= 1) {
          camzoom = atof(argv[++i]);
        }
        continue;
      }

      if (!strcmp("-aomaxdist", argv[i])) {
        if ((argc - i) >= 1) {
          aomaxdist = atof(argv[++i]);
        }
        printf("** aomaxdist: %.2f\n", aomaxdist);
        continue;
      }

      if (!strcmp("-noao", argv[i])) {
        ambientocclusion = 0;
        printf("Disabling renderer-specific AO lighting.\n");
        continue;
      }

      if (!strcmp("-res", argv[i])) {
        if ((argc - i) >= 2) {
          imgSize[0] = atoi(argv[++i]);
          imgSize[1] = atoi(argv[++i]);
          printf("Image resolution set to: %d x %d\n", imgSize[0], imgSize[1]);
        }
        continue;
      }

      if (!strcmp("-denoiser", argv[i])) {
        denoiser = 1;
        continue;
      }

      if (!strcmp("-dof", argv[i])) {
        dof = 1;
        continue;
      }

#if 0
      if (!strcmp("-pause", argv[i])) {
        sleep(10);
      }
#endif

      if (!strcmp("-warmup", argv[i])) {
        warmup = 1;
        printf("Enabling profiling warm-up pass and timing.\n");
        continue;
      }

      printf("Unrecognized flag: '%s'.\n", argv[i]);
    }
  } 

  printf("All command line flags processed.\n");

#if defined(TACHYON_USEPINNEDMEMORY)
  printf("USING PINNED HOST MEMORY ALLOCATIONS\n");
#endif

  PROFILE_PUSH_RANGE("Loading OBJ", 0);

  std::string inputfile = objfilename;
  tinyobj::ObjReaderConfig reader_config;
  tinyobj::ObjReader reader;

  // get path to material files
  char *materialdir = getmaterialdir(objfilename);
  if (materialdir) {
    printf("MTL directory: %s\n", materialdir);
    reader_config.mtl_search_path = materialdir;
  } else {
    reader_config.mtl_search_path = "./";
  }


  if (!reader.ParseFromFile(inputfile, reader_config)) {
    if (!reader.Error().empty()) {
      printf("TinyObjReader Error: %s\n", reader.Error().c_str());
    }
    return -1;
  }

  if (!reader.Warning().empty()) {
    printf("TinyObjReader Warning: %s\n", reader.Warning().c_str());
  }

  PROFILE_POP_RANGE();

  auto& attrib = reader.GetAttrib();
  auto& shapes = reader.GetShapes();
  auto& materials = reader.GetMaterials();

  printf("OBJ materials: %ld\n", materials.size());
#if 0
  for (size_t m = 0; m < materials.size(); m++) {
    printf("  mat[%ld]: '%s'\n", m, materials[m].name.c_str());
  } 
  printf("\n");
#endif

  printf("OBJ shapes: %ld\n", shapes.size());
#if 0
  // Loop over shapes
  for (size_t s = 0; s < shapes.size(); s++) {
    printf("  shape[%ld]: %ld faces\n", 
           s, shapes[s].mesh.num_face_vertices.size());

#if 0
    // Loop over faces(polygon)
    size_t index_offset = 0;
    for (size_t f = 0; f < shapes[s].mesh.num_face_vertices.size(); f++) {
      size_t fv = size_t(shapes[s].mesh.num_face_vertices[f]);

      index_offset += fv;
    }
#endif
  }
#endif


  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();
  printf("\n");

  PROFILE_POP_RANGE();


  PROFILE_PUSH_RANGE("Build Scene", 0);

  //
  // Build scene
  //

  // 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.4, 0.6, 0.7 };
  float rtbggradbotcolor[] = { 0.5, 0.1, 0.2 };

  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(512);
  rt->shadows_enable(1);

  if (ambientocclusion) {
    rt->set_ao_samples(1);
    rt->set_ao_ambient(1.2);
    rt->set_ao_direct(0.4);
    rt->set_ao_maxdist(aomaxdist);
  }

  float lightdir0[] = {  0.3f, 1.0f,  0.3f };
  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);
//  rt->set_camera_type(TachyonOptiX::RT_EQUIRECTANGULAR);
//  rt->set_camera_type(TachyonOptiX::RT_OCTOHEDRAL);

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

  if (denoiser)
    rt->denoiser_enable(1);

  if (dof)
    rt->camera_dof_enable(1);

  rt->set_camera_dof_fnumber(80.0f);
  rt->set_camera_dof_focal_dist(500.0f);

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


  printf("Loading materials and texture images\n");
  for (size_t m = 0; m < materials.size(); m++) {
#if 1
    float ambient = 0.0f;
//      (materials[m].ambient[0] + 
//       materials[m].ambient[1] +
//       materials[m].ambient[2]) / 3.0f;

    float diffuse = 
      (materials[m].diffuse[0] + 
       materials[m].diffuse[1] +
       materials[m].diffuse[2]) / 3.0f;

    float specular =
      (materials[m].specular[0] + 
       materials[m].specular[1] +
       materials[m].specular[2]) / 3.0f;

    float shininess = 
      materials[m].shininess;

    float reflectivity = 0.0f;

    float opacity = 
      materials[m].dissolve;

//      (materials[m].dissolve[0] +
//       materials[m].dissolve[1] +
//       materials[m].dissolve[2]) / 3.0f;

    float outline = 0.0f;
    float outlinewidth = 0.0f;
    int transmode = 0;

    // skipped material properties:
    //   materials[m].transmittance
    //   materials[m].emission
    //   materials[m].ior
    //   materials[m].illum   
    //   materials[m].ambient_texname.c_str()
    //   materials[m].diffuse_texname.c_str()
    //   materials[m].specular_texname.c_str()
    //   materials[m].specular_highlight_texname.c_str()
    //
    //   materials[m].bump_texname.c_str()
    //   materials[m].bump_texopt.bump_multiplier
    //   materials[m].alpha_texname.c_str()
    //   materials[m].displacement_texname.c_str()
    //
    // PBR extensions: 
    //   materials[m].roughness
    //   materials[m].metallic  
    //   materials[m].sheen    
    //   materials[m].clearcoat_thickness
    //   materials[m].clearcoat?
    //   materials[m].anisotropy
    //   materials[m].anisotropy_rotation
    //   materials[m].emissive_texname.c_str()
    //   materials[m].roughness_texname.c_str()
    //   materials[m].metallic_texname.c_str()
    //   materials[m].sheen_texname.c-str()
    //   materials[m].normal_texname.c_str()
    //
    //
    //  // handle unrecognized/non-standard property names
    //  std::map<std::string, std::string>::const_iterator it(
    //      materials[i].unknown_parameter.begin());
    //  std::map<std::string, std::string>::const_iterator itEnd(
    //      materials[i].unknown_parameter.end());
    //
    //  for (; it != itEnd; it++) {
    //    printf("  material.%s = %s\n", it->first.c_str(),it->second.c_str());
    //  }
    //
#else
    float ambient = 0.1f;
    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;
#endif

//    printf("mat[%ld]: a:%.1f d:%.2f s:%.2f o:%.2f sh:%.2f\n",
//           m, ambient, diffuse, specular, opacity, shininess);

    unsigned char *img = NULL;
    int xres=0, yres=0, chinfile=0;
    const char *texfilename = materials[m].diffuse_texname.c_str();
    if ((texfilename != NULL) && (texfilename[0] != '\0')) {
      printf("mat[%lu] '%s'\n", m, texfilename);
      img = loadimage_rgba4u(materialdir, texfilename, xres, yres, chinfile);
      if (img != NULL) {
        printf("  %4d x %4d, %d channels\n", xres, yres, chinfile);
      } else {
        printf("  Failed to load image!\n");
      }
    }

    if (img != NULL) {
      int texflags = (chinfile == 4) ? RT_TEX_ALPHA : RT_TEX_NONE;

      rt->add_tex2d_rgba4u(img, xres, yres, texflags | RT_COLORSPACE_sRGB, m);

      rt->add_material_textured(ambient, diffuse, specular, 
                                shininess, reflectivity, opacity, 
                                outline, outlinewidth, transmode, m, m);
    } else {
      rt->add_material(ambient, diffuse, specular, shininess, reflectivity, 
                       opacity, outline, outlinewidth, transmode, m);
    }
  } 
  printf("\n");

  // XXX Start Nsight Compute Profiles here...
  PROFILE_START();

  PROFILE_POP_RANGE();

  if (warmup) {
    rt->set_verbose_mode(TachyonOptiX::RT_VERB_MIN);
    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();
  }

  rt->set_verbose_mode(TachyonOptiX::RT_VERB_TIMING);

#if 1
  int coordinit = 0;
  float3 vmin={ 0 }, vmax={ 0 };

  std::vector<TriangleMesh> meshes;

  // Loop over shapes
  int meshcnt = 0;
  for (size_t s = 0; s < shapes.size(); s++) {
    TriangleMesh tm;
    long numfaces = shapes[s].mesh.num_face_vertices.size();
    long numverts = numfaces * 3;
    tm.vertices.resize(numverts);
    tm.normals.resize(numverts);
    tm.tex2d.resize(numverts);
    int texcoords = 0;

    float3 *verts = tm.vertices.data();
    float3 *norms = tm.normals.data();
    float2 *tex2d = tm.tex2d.data();
    tm.uniform_color = make_float3(0.7); // fall-back

    // Loop over faces(polygon)
    size_t index_offset = 0;
    for (size_t f = 0; f < shapes[s].mesh.num_face_vertices.size(); f++) {
      size_t fv = size_t(shapes[s].mesh.num_face_vertices[f]);

      // Loop over vertices in the face.
      for (size_t v = 0; v < fv; v++) {
        // index into attribute buffer
        tinyobj::index_t idx = shapes[s].mesh.indices[index_offset + v];

        // access to vertex
        tinyobj::real_t vx = attrib.vertices[3*size_t(idx.vertex_index)+0];
        tinyobj::real_t vy = attrib.vertices[3*size_t(idx.vertex_index)+1];
        tinyobj::real_t vz = attrib.vertices[3*size_t(idx.vertex_index)+2];

        // Check if `normal_index` is zero or positive. 
        // negative = no normal data
        tinyobj::real_t nx=0, ny=0, nz=0;
        if (idx.normal_index >= 0) {
          nx = attrib.normals[3*size_t(idx.normal_index)+0];
          ny = attrib.normals[3*size_t(idx.normal_index)+1];
          nz = attrib.normals[3*size_t(idx.normal_index)+2];
        }

        // Check if `texcoord_index` is zero or positive. 
        // negative = no texcoord data
        tinyobj::real_t tx=0, ty=0;
        if (idx.texcoord_index >= 0) {
          texcoords=1;
          tx = attrib.texcoords[2*size_t(idx.texcoord_index)+0];
          ty = attrib.texcoords[2*size_t(idx.texcoord_index)+1];
        }

        // Optional: vertex colors
        // tinyobj::real_t red   = attrib.colors[3*size_t(idx.vertex_index)+0];
        // tinyobj::real_t green = attrib.colors[3*size_t(idx.vertex_index)+1];
        // tinyobj::real_t blue  = attrib.colors[3*size_t(idx.vertex_index)+2];

        float3 newv = make_float3(vx, vy, vz);
        float3 newn = make_float3(nx, ny, nz);

        if (!coordinit) {
          coordinit=1;
          vmin = vmax = newv;
        } else {
          vmin.x = fminf(vmin.x, newv.x);         
          vmin.y = fminf(vmin.y, newv.y);         
          vmin.z = fminf(vmin.z, newv.z);         
          vmax.x = fmaxf(vmax.x, newv.x);         
          vmax.y = fmaxf(vmax.y, newv.y);         
          vmax.z = fmaxf(vmax.z, newv.z);         
        }

        long vidx = index_offset+v;
        verts[vidx] = newv;
        norms[vidx] = newn;

        //
        // XXX flip vertical texcoord.ty axis to accomodate the 
        // memory layout of images loaded via stbi_load()...
        //
        tex2d[vidx] = make_float2(tx, 1.0f - ty);  
      }
      index_offset += fv;

      // per-face material
      shapes[s].mesh.material_ids[f];
    }

    if (!texcoords)
      tm.tex2d.clear();

    meshes.push_back(tm);
    rt->add_trimesh(tm, shapes[s].mesh.material_ids[0]);
    meshcnt++;
  }
  printf("Added %d trimeshes to scene...\n", meshcnt);
  printf("BBox: min: %.2f %.2f %.2f  max: %.2f %.2f %.2f\n",
         vmin.x, vmin.y, vmin.z, vmax.x, vmax.y, vmax.z);
  printf("\n");
#endif

  PROFILE_PUSH_RANGE("Render Scene", 0);
  printf("Rendering frames w/ accumulation buffer...\n");
  // render 100 accumulated frames
//  for (int frames = 0; frames < 3; frames++) {
    rt->render();
//  }
  PROFILE_POP_RANGE();

  rt->print_raystats_info();

  PROFILE_PUSH_RANGE("Write Output Image", 0);

  char filename[1024];
  sprintf(filename, "out.png");
  printf("Writing accumulated frames to '%s'...\n", filename);
  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);
  }

  delete rt;

  PROFILE_POP_RANGE();
  return 0;
}


