#include <optix.h>
#include <optix_function_table_definition.h>
#include <optix_stack_size.h>
#include <optix_stubs.h>

#include <cuda_runtime.h>
#include "cuda/whitted.h"

#include <sampleConfig.h>

#include <sutil/CUDAOutputBuffer.h>
#include <sutil/Exception.h>
#include <sutil/sutil.h>
#include <sutil/Camera.h>
#include <sutil/Trackball.h>
#include <sutil/Scene.h>

#include "optixTriangle.h"
#include "generateMesh.h"

#include <array>
#include <iomanip>
#include <iostream>
#include <string>

struct RaycastingState
{
    int                         width                       = 0;
    int                         height                      = 0;

    OptixDeviceContext          context                     = 0;
    sutil::Scene                scene                       = {};

    OptixPipelineCompileOptions pipeline_compile_options    = {};
    OptixModule                 ptx_module = 0;
    //OptixPipeline               pipeline_1 = 0;
    //OptixPipeline               pipeline_2 = 0;
    OptixPipeline               pipeline                    = 0;

    OptixProgramGroup           raygen_prog_group           = 0;
    OptixProgramGroup           miss_prog_group             = 0;
    OptixProgramGroup           hit_prog_group              = 0;

    Params                      params                      = {};
    //Params                      params_translated = {};
    OptixShaderBindingTable     sbt                         = {};

    sutil::Texture              mask                        = {};
};

template <typename T>
struct SbtRecord
{
    __align__( OPTIX_SBT_RECORD_ALIGNMENT ) char header[OPTIX_SBT_RECORD_HEADER_SIZE];
    T data;
};

typedef SbtRecord<RayGenData>               RayGenSbtRecord;
typedef SbtRecord<MissData>                 MissSbtRecord;
typedef SbtRecord<TriangleMeshSBTData>      HitGroupSbtRecord;

void configureCamera( sutil::Camera& cam, const uint32_t width, const uint32_t height )
{
    // --- Sets the ray origin (eye)
    cam.setEye( {0.0f, 2.0f, 2.0f} );
    cam.setLookat( {0.0f, 0.0f, 0.0f} );
    cam.setUp( {0.0f, 1.0f, 3.0f} );
    cam.setFovY( 45.0f );
    cam.setAspectRatio( (float)width / (float)height );
}


void printUsageAndExit( const char* argv0 )
{
    std::cerr << "Usage  : " << argv0 << " [options]\n";
    std::cerr << "Options: --file | -f <filename>      Specify file for image output\n";
    std::cerr << "         --help | -h                 Print this usage message\n";
    std::cerr << "         --dim=<width>x<height>      Set image dimensions; defaults to 512x384\n";
    exit( 1 );
}


static void context_log_cb( unsigned int level, const char* tag, const char* message, void* /*cbdata */)
{
    std::cerr << "[" << std::setw( 2 ) << level << "][" << std::setw( 12 ) << tag << "]: "
              << message << "\n";
}

OptixDeviceContext context = nullptr;

OptixTraversableHandle  gas_handle;
CUdeviceptr             d_gas_output_buffer;

OptixModule             module                          = nullptr;
OptixModule             sphere_module                   = nullptr;
OptixPipelineCompileOptions pipeline_compile_options    = {};

OptixProgramGroup       raygen_prog_group               = nullptr;
OptixProgramGroup       miss_prog_group                 = nullptr;
OptixProgramGroup       hitgroup_prog_group             = nullptr;

OptixPipeline pipeline = nullptr;

OptixShaderBindingTable sbt = {};

CUdeviceptr             d_vertices                      = 0;
CUdeviceptr             d_indices                       = 0;

//OptixPipelineCompileOptions pipeline_compile_options = {};


//
// Initialize CUDA and create OptiX context
//
//OptixDeviceContext context = nullptr;
/*********************************/
/* CUDA AND OPTIX INITIALIZATION */
/*********************************/
void init() {

    // Initialize CUDA
    CUDA_CHECK(cudaFree(0));

    // Initialize the OptiX API, loading all API entry points
    OPTIX_CHECK(optixInit());

    // Specify context options
    OptixDeviceContextOptions options = {};
    options.logCallbackFunction = &context_log_cb;
    options.logCallbackLevel = 4;

    // Associate a CUDA context (and therefore a specific GPU) with this
    // device context
    CUcontext cuCtx = 0;  // zero means take the current context
    OPTIX_CHECK(optixDeviceContextCreate(cuCtx, &options, &context));
}

/*****************/
/* GENERATE MESH */
/*****************/


    //// Triangle build input: simple list of three vertices
    ///*const std::array<float3, 6> vertices[1];
    //const std::array<float3, 3> vertex =
    //{ {
    //      { -0.5f, -0.5f, 0.0f },
    //      {  0.5f, -0.5f, 0.0f },
    //      {  0.0f,  0.5f, 0.0f }
    //} };

    //vertices[0] = vertex;*/

    ///*const std::array<float3, 6> vertices =
    //{ {
    //      { -0.25f,  0.5f, 0.0f },
    //      { -0.25f, -0.5f, 0.0f },
    //      { -0.75f, -0.5f, 0.0f },
    //      {  0.25f,  0.5f, 0.0f },
    //      {  0.25f, -0.5f, 0.0f },
    //      {  0.75f, -0.5f, 0.0f },
    //} };*/

    ////float vertices_temp[] = { 0.5f,  0.5f, 0.0f, -0.5f,  0.5f, 0.0f, -0.5f, -0.5f, 0.0f, 0.5f, -0.5f, 0.0f, 0.5f,  0.5f, 1.0f, -0.5f,  0.5f, 1.0f, -0.5f, -0.5f, 1.0f, 0.5f, -0.5f, 1.0f };
    //std::vector<float3> vertices;
    //vertices.push_back({  0.5f,  0.5f, 0.0f });
    //vertices.push_back({ -0.5f,  0.5f, 0.0f });
    //vertices.push_back({ -0.5f, -0.5f, 0.0f });
    //vertices.push_back({  0.5f, -0.5f, 0.0f });
    //vertices.push_back({  0.5f,  0.5f, 1.0f });
    //vertices.push_back({ -0.5f,  0.5f, 1.0f });
    //vertices.push_back({ -0.5f, -0.5f, 1.0f });
    //vertices.push_back({  0.5f, -0.5f, 1.0f });
    ////for (int i = 0; i < 8; i++) {
    ////    float3 vertt;
    ////    vertt.x = vertices_temp[3 * i + 0];
    ////    vertt.y = vertices_temp[3 * i + 1];
    ////    vertt.z = vertices_temp[3 * i + 2];
    ////    vertices.push_back(vertt);
    ////}

    //int indices_temp[] = { 1,2,3, 0,1,3, 5,6,7, 4,5,7, 4,7,0, 7,3,0, 6,2,3, 6,7,3, 5,1,2, 5,6,2, 1,0,5, 5,4,0 };
    //std::vector<int3> indices;
    //for (int i = 0; i < 12; i++) {
    //    int3 indt;
    //    indt.x = indices_temp[3 * i + 0];
    //    indt.y = indices_temp[3 * i + 1];
    //    indt.z = indices_temp[3 * i + 2];
    //    indices.push_back(indt);
    //}
/*************************/
/* ACCELERATION HANDLING */
/*************************/
void accelerationHandling(std::vector<float3>& vertices, std::vector<int3>& indices) {

    OptixAccelBuildOptions accel_options    = {};
    // --- Use default options for simplicity.  In a real use case we would want to enable compaction, etc
    //accel_options.buildFlags                = OPTIX_BUILD_FLAG_NONE;
    accel_options.buildFlags                = OPTIX_BUILD_FLAG_ALLOW_COMPACTION;
    accel_options.operation                 = OPTIX_BUILD_OPERATION_BUILD;

    // --- AABB build input
    //OptixAabb   aabb = { -1.5f, -1.5f, -1.5f, 1.5f, 1.5f, 1.5f };
    //CUdeviceptr d_aabb_buffer;
    //CUDA_CHECK(cudaMalloc(reinterpret_cast<void**>(&d_aabb_buffer), sizeof(OptixAabb)));
    //CUDA_CHECK(cudaMemcpy(reinterpret_cast<void*>(d_aabb_buffer), &aabb, sizeof(OptixAabb), cudaMemcpyHostToDevice));
    
    //OptixBuildInput aabb_input = {};
    //aabb_input.type = OPTIX_BUILD_INPUT_TYPE_CUSTOM_PRIMITIVES;
    //aabb_input.customPrimitiveArray.aabbBuffers = &d_aabb_buffer;
    //aabb_input.customPrimitiveArray.numPrimitives = 2 * indices.size();

    //uint32_t aabb_input_flags[1] = { OPTIX_GEOMETRY_FLAG_NONE };
    //aabb_input.customPrimitiveArray.flags = aabb_input_flags;
    //aabb_input.customPrimitiveArray.numSbtRecords = 1;

    const size_t    vertices_size               = sizeof(float3) * vertices.size();
    //CUdeviceptr     d_vertices                  = 0;
    CUDA_CHECK(cudaMalloc(reinterpret_cast<void**>(&d_vertices), vertices_size));
    CUDA_CHECK(cudaMemcpy(reinterpret_cast<void*>(d_vertices), vertices.data(), vertices_size, cudaMemcpyHostToDevice));

    const size_t    indices_size                = sizeof(int3) * indices.size();
    //CUdeviceptr     d_indices                   = 0;
    CUDA_CHECK(cudaMalloc(reinterpret_cast<void**>(&d_indices), indices_size));
    CUDA_CHECK(cudaMemcpy(reinterpret_cast<void*>(d_indices), indices.data(), indices_size, cudaMemcpyHostToDevice));

    // Our build input is a simple list of non-indexed triangle vertices
    OptixBuildInput triangle_input                  = {};
    triangle_input.triangleArray.vertexFormat       = OPTIX_VERTEX_FORMAT_FLOAT3;
    triangle_input.type                             = OPTIX_BUILD_INPUT_TYPE_TRIANGLES;
    triangle_input.triangleArray.numVertices        = static_cast<uint32_t>(vertices.size());
    triangle_input.triangleArray.vertexBuffers      = &d_vertices;

    const uint32_t triangle_input_flags[1]          = { OPTIX_GEOMETRY_FLAG_NONE };
    triangle_input.triangleArray.flags              = triangle_input_flags;
    triangle_input.triangleArray.numSbtRecords      = 1;
    triangle_input.triangleArray.vertexStrideInBytes = sizeof(float3);
    //triangle_input.triangleArray.sbtIndexOffsetBuffer = 0;
    //triangle_input.triangleArray.sbtIndexOffsetSizeInBytes = 0;
    //triangle_input.triangleArray.sbtIndexOffsetStrideInBytes = 0;

    triangle_input.triangleArray.indexFormat        = OPTIX_INDICES_FORMAT_UNSIGNED_INT3;
    triangle_input.triangleArray.indexStrideInBytes = sizeof(int3);
    triangle_input.triangleArray.numIndexTriplets   = static_cast<uint32_t>(indices.size());
    triangle_input.triangleArray.indexBuffer        = d_indices;

    //// in this example we have one SBT entry, and no per-primitive
    //// materials:

    OptixAccelBufferSizes gas_buffer_sizes;
    OPTIX_CHECK(optixAccelComputeMemoryUsage(context, &accel_options, &triangle_input, 1, // Number of build inputs
        &gas_buffer_sizes));
    //OPTIX_CHECK(optixAccelComputeMemoryUsage(context, &accel_options, &aabb_input, 1, &gas_buffer_sizes));
    CUdeviceptr d_temp_buffer_gas;
    CUDA_CHECK(cudaMalloc(reinterpret_cast<void**>(&d_temp_buffer_gas), gas_buffer_sizes.tempSizeInBytes));
    CUDA_CHECK(cudaMalloc(
        reinterpret_cast<void**>(&d_gas_output_buffer),
        gas_buffer_sizes.outputSizeInBytes
    ));

    // non-compacted output
    //CUdeviceptr d_buffer_temp_output_gas_and_compacted_size;
    //size_t      compactedSizeOffset = roundUp<size_t>(gas_buffer_sizes.outputSizeInBytes, 8ull);
    //CUDA_CHECK(cudaMalloc(reinterpret_cast<void**>(&d_buffer_temp_output_gas_and_compacted_size), compactedSizeOffset + 8));

    //OptixAccelEmitDesc emitProperty = {};
    //emitProperty.type = OPTIX_PROPERTY_TYPE_COMPACTED_SIZE;
    //emitProperty.result = (CUdeviceptr)((char*)d_buffer_temp_output_gas_and_compacted_size + compactedSizeOffset);

    // --- E' questa routine che dà problemi
    //OPTIX_CHECK(optixAccelBuild(context,
    //    0,                  // CUDA stream
    //    &accel_options,
    //    &aabb_input,
    //    1,                  // num build inputs
    //    d_temp_buffer_gas,
    //    gas_buffer_sizes.tempSizeInBytes,
    //    d_buffer_temp_output_gas_and_compacted_size,
    //    gas_buffer_sizes.outputSizeInBytes,
    //    &gas_handle,
    //    &emitProperty,      // emitted property list
    //    1                   // num emitted properties
    //));

    OPTIX_CHECK(optixAccelBuild(
        context,
        0,                  // CUDA stream
        &accel_options,
        &triangle_input,
        1,                  // num build inputs
        d_temp_buffer_gas,
        gas_buffer_sizes.tempSizeInBytes,
        d_gas_output_buffer,
        gas_buffer_sizes.outputSizeInBytes,
        &gas_handle,
        nullptr,            // emitted property list
        0                   // num emitted properties
    ));


    // We can now free the scratch space buffer used during build and the vertex
    // inputs, since they are not needed by our trivial shading method
    CUDA_CHECK(cudaFree(reinterpret_cast<void*>(d_temp_buffer_gas)));
    //CUDA_CHECK(cudaFree((void*)d_aabb_buffer));
    // CUDA_CHECK(cudaFree(reinterpret_cast<void*>(d_vertices)));

    //size_t compacted_gas_size;
    //CUDA_CHECK(cudaMemcpy(&compacted_gas_size, (void*)emitProperty.result, sizeof(size_t), cudaMemcpyDeviceToHost));

    //if (compacted_gas_size < gas_buffer_sizes.outputSizeInBytes)
    //{
    //    CUDA_CHECK(cudaMalloc(reinterpret_cast<void**>(&d_gas_output_buffer), compacted_gas_size));

    //    // use handle as input and output
    //    OPTIX_CHECK(optixAccelCompact(context, 0, gas_handle, d_gas_output_buffer, compacted_gas_size, &gas_handle));

    //    CUDA_CHECK(cudaFree((void*)d_buffer_temp_output_gas_and_compacted_size));
    //}
    //else
    //{
    //    d_gas_output_buffer = d_buffer_temp_output_gas_and_compacted_size;
    //}

}

/*****************/
/* CREATE MODULE */
/*****************/
void createModule() {
    
    char log[2048]; // For error reporting from OptiX creation functions
    size_t sizeof_log = sizeof(log);

    OptixModuleCompileOptions module_compile_options    = {};
    module_compile_options.maxRegisterCount             = OPTIX_COMPILE_DEFAULT_MAX_REGISTER_COUNT;
    module_compile_options.optLevel                     = OPTIX_COMPILE_OPTIMIZATION_DEFAULT;
    module_compile_options.debugLevel                   = OPTIX_COMPILE_DEBUG_LEVEL_LINEINFO;

    pipeline_compile_options.usesMotionBlur             = false;
    //pipeline_compile_options.traversableGraphFlags      = OPTIX_TRAVERSABLE_GRAPH_FLAG_ALLOW_SINGLE_LEVEL_INSTANCING;
    pipeline_compile_options.traversableGraphFlags      = OPTIX_TRAVERSABLE_GRAPH_FLAG_ALLOW_SINGLE_GAS;
    //pipeline_compile_options.numPayloadValues           = 3;
    pipeline_compile_options.numPayloadValues           = 4;
    //pipeline_compile_options.numAttributeValues         = 3;
    pipeline_compile_options.numAttributeValues         = 2;
#ifdef DEBUG // Enables debug exceptions during optix launches. This may incur significant performance cost and should only be done during development.
    pipeline_compile_options.exceptionFlags = OPTIX_EXCEPTION_FLAG_DEBUG | OPTIX_EXCEPTION_FLAG_TRACE_DEPTH | OPTIX_EXCEPTION_FLAG_STACK_OVERFLOW;
#else
    //pipeline_compile_options.exceptionFlags             = OPTIX_EXCEPTION_FLAG_NONE;
    pipeline_compile_options.exceptionFlags             = OPTIX_EXCEPTION_FLAG_STACK_OVERFLOW;
#endif
    pipeline_compile_options.pipelineLaunchParamsVariableName = "params";
    pipeline_compile_options.usesPrimitiveTypeFlags     = OPTIX_PRIMITIVE_TYPE_FLAGS_TRIANGLE;

    const std::string ptx = sutil::getPtxString(OPTIX_SAMPLE_NAME, OPTIX_SAMPLE_DIR, "optixTriangle.cu");
    //const std::string ptx = sutil::getPtxString(OPTIX_SAMPLE_NAME, OPTIX_SAMPLE_DIR, "optixPathTracer.cu");

    OPTIX_CHECK_LOG(optixModuleCreateFromPTX(context, &module_compile_options, &pipeline_compile_options, ptx.c_str(), ptx.size(), log, &sizeof_log, &module));
}

/*************************/
/* CREATE PROGRAM GROUPS */
/*************************/
void createProgramGroups() {
    char log[2048]; // For error reporting from OptiX creation functions
    size_t sizeof_log = sizeof(log);

    OptixProgramGroupOptions program_group_options      = {}; // Initialize to zeros

    // --- RENDERING
    OptixProgramGroupDesc raygen_prog_group_desc        = {}; 
    raygen_prog_group_desc.kind                         = OPTIX_PROGRAM_GROUP_KIND_RAYGEN;
    raygen_prog_group_desc.raygen.module                = module;
    raygen_prog_group_desc.raygen.entryFunctionName     = "__raygen__rg";
    // --- 1 is the number of program groups
    OPTIX_CHECK_LOG(optixProgramGroupCreate(context, &raygen_prog_group_desc, 1, &program_group_options, log, &sizeof_log, &raygen_prog_group));

    // --- RAY MISS
    OptixProgramGroupDesc miss_prog_group_desc          = {};
    miss_prog_group_desc.kind                           = OPTIX_PROGRAM_GROUP_KIND_MISS;
    miss_prog_group_desc.miss.module                    = module;
    miss_prog_group_desc.miss.entryFunctionName         = "__miss__ms";
    // --- 1 is the number of program groups
    OPTIX_CHECK_LOG(optixProgramGroupCreate(context, &miss_prog_group_desc, 1, &program_group_options, log, &sizeof_log, &miss_prog_group));

    // --- RAY HIT
    OptixProgramGroupDesc hitgroup_prog_group_desc      = {};
    hitgroup_prog_group_desc.kind                       = OPTIX_PROGRAM_GROUP_KIND_HITGROUP;
    hitgroup_prog_group_desc.hitgroup.moduleCH          = module;
    hitgroup_prog_group_desc.hitgroup.entryFunctionNameCH = "__closesthit__ch";
    // --- 1 is the number of program groups
    OPTIX_CHECK_LOG(optixProgramGroupCreate(context, &hitgroup_prog_group_desc, 1, &program_group_options, log, &sizeof_log, &hitgroup_prog_group));
}

/*******************/
/* CREATE PIPELINE */
/*******************/
void createPipeline() {
    char log[2048]; // For error reporting from OptiX creation functions
    size_t sizeof_log = sizeof(log);

    const uint32_t    max_trace_depth = 1;
    OptixProgramGroup program_groups[] = { raygen_prog_group, miss_prog_group, hitgroup_prog_group };

    OptixPipelineLinkOptions pipeline_link_options              = {};
    pipeline_link_options.maxTraceDepth                         = max_trace_depth;
    pipeline_link_options.debugLevel                            = OPTIX_COMPILE_DEBUG_LEVEL_FULL;
    OPTIX_CHECK_LOG(optixPipelineCreate(context, &pipeline_compile_options, &pipeline_link_options, program_groups, sizeof(program_groups) / sizeof(program_groups[0]), log, &sizeof_log, &pipeline));

    OptixStackSizes stack_sizes                                 = {};
    for (auto& prog_group : program_groups) { OPTIX_CHECK(optixUtilAccumulateStackSizes(prog_group, &stack_sizes)); }

    uint32_t direct_callable_stack_size_from_traversal;
    uint32_t direct_callable_stack_size_from_state;
    uint32_t continuation_stack_size;
    // --- 0 is maxCCDepth
    // --- 0 is maxDCDEpth
    OPTIX_CHECK(optixUtilComputeStackSizes(&stack_sizes, max_trace_depth, 0, 0, &direct_callable_stack_size_from_traversal, &direct_callable_stack_size_from_state, &continuation_stack_size));
    // --- 1 is maxTraversableDepth
    OPTIX_CHECK(optixPipelineSetStackSize(pipeline, direct_callable_stack_size_from_traversal, direct_callable_stack_size_from_state, continuation_stack_size, 1));
}

/*****************/
/* SET UP SHADER */
/*****************/
void createSBT(float3 color) {
    
    // --- RENDERING
    CUdeviceptr  d_raygen_record            = 0;
    const size_t raygen_record_size         = sizeof(RayGenSbtRecord);
    CUDA_CHECK(cudaMalloc(reinterpret_cast<void**>(&d_raygen_record), raygen_record_size));
    RayGenSbtRecord rg_sbt;
    OPTIX_CHECK(optixSbtRecordPackHeader(raygen_prog_group, &rg_sbt));
    CUDA_CHECK(cudaMemcpy(reinterpret_cast<void*>(d_raygen_record), &rg_sbt, raygen_record_size, cudaMemcpyHostToDevice));

    // --- RAY MISS
    CUdeviceptr d_miss_record               = 0;
    size_t      miss_record_size            = sizeof(MissSbtRecord);
    CUDA_CHECK(cudaMalloc(reinterpret_cast<void**>(&d_miss_record), miss_record_size));
    MissSbtRecord ms_sbt;
    ms_sbt.data = { 0.3f, 0.1f, 0.2f };     // --- Miss ray color
    OPTIX_CHECK(optixSbtRecordPackHeader(miss_prog_group, &ms_sbt));
    CUDA_CHECK(cudaMemcpy(reinterpret_cast<void*>(d_miss_record), &ms_sbt, miss_record_size, cudaMemcpyHostToDevice));

    // --- RAY HIT
    HitGroupSbtRecord hg_sbt;
    hg_sbt.data.vertex  = (float3*)d_vertices;
    hg_sbt.data.index   = (int3*)  d_indices;
    hg_sbt.data.color   = color;

    CUdeviceptr d_hitgroup_record           = 0;
    size_t      hitgroup_record_size        = sizeof(HitGroupSbtRecord);
    CUDA_CHECK(cudaMalloc(reinterpret_cast<void**>(&d_hitgroup_record), hitgroup_record_size));
    OPTIX_CHECK(optixSbtRecordPackHeader(hitgroup_prog_group, &hg_sbt));
    CUDA_CHECK(cudaMemcpy(reinterpret_cast<void*>(d_hitgroup_record), &hg_sbt, hitgroup_record_size, cudaMemcpyHostToDevice));


    //// --- Packing pointers on the host side
    //std::vector<HitGroupSbtRecord> hitgroupRecords;
    //HitGroupSbtRecord hg_sbt;
    //hg_sbt.data.vertex  = (float3*)d_vertices;
    //hg_sbt.data.index   = (float3*)d_indices;
    ////hg_sbt.data.color = model.color;
    //hitgroupRecords.push_back(hg_sbt);

    //// --- Moving the packed pointers from host to device
    //CUdeviceptr d_hitgroup_record;
    //size_t      hitgroup_record_size = sizeof(HitGroupSbtRecord);
    //CUDA_CHECK(cudaMalloc(reinterpret_cast<void**>(&d_hitgroup_record), hitgroup_record_size));
    //CUDA_CHECK(cudaMemcpy(reinterpret_cast<void*>(d_hitgroup_record), &hg_sbt, hitgroup_record_size, cudaMemcpyHostToDevice));
    //
    //OPTIX_CHECK(optixSbtRecordPackHeader(hitgroup_prog_group, &hg_sbt));

    //sbt.hitgroupRecordStrideInBytes = sizeof(HitGroupSbtRecord);
    //sbt.hitgroupRecordCount         = (int)hitgroupRecords.size();
    //sbt.hitgroupRecordBase          = d_hitgroup_record;

    sbt.raygenRecord                        = d_raygen_record;
    sbt.missRecordBase                      = d_miss_record;
    sbt.missRecordStrideInBytes             = sizeof(MissSbtRecord);
    sbt.missRecordCount                     = 1;
    sbt.hitgroupRecordBase                  = d_hitgroup_record;
    sbt.hitgroupRecordStrideInBytes         = sizeof(HitGroupSbtRecord);
    sbt.hitgroupRecordCount                 = 1;
}

/**********/
/* LAUNCH */
/**********/
void launch(sutil::CUDAOutputBuffer<uchar4> &output_buffer, const int width, const int height) {
    CUstream stream;
    CUDA_CHECK(cudaStreamCreate(&stream));

    sutil::Camera cam;
    configureCamera(cam, width, height);

    Params params;
    params.image = output_buffer.map();
    params.image_width = width;
    params.image_height = height;
    params.handle = gas_handle;
    params.cam_eye = cam.eye();
    cam.UVWFrame(params.cam_u, params.cam_v, params.cam_w);

    CUdeviceptr d_param;
    CUDA_CHECK(cudaMalloc(reinterpret_cast<void**>(&d_param), sizeof(Params)));
    CUDA_CHECK(cudaMemcpy(
        reinterpret_cast<void*>(d_param),
        &params, sizeof(params),
        cudaMemcpyHostToDevice
    ));

    OPTIX_CHECK(optixLaunch(pipeline, stream, d_param, sizeof(Params), &sbt, width, height, /*depth=*/1));
    CUDA_SYNC_CHECK();

    output_buffer.unmap();
}

/*******************/
/* DISPLAY RESULTS */
/*******************/
void displayResults(sutil::CUDAOutputBuffer<uchar4>& output_buffer, std::string &outfile, char* argv[], const int width, const int height) {
    sutil::ImageBuffer buffer;
    buffer.data = output_buffer.getHostPointer();
    buffer.width = width;
    buffer.height = height;
    buffer.pixel_format = sutil::BufferImageFormat::UNSIGNED_BYTE4;
    if (outfile.empty())
        sutil::displayBufferWindow(argv[0], buffer);
    else
        sutil::saveImage(outfile.c_str(), buffer, false);
}

/***********/
/* CLEANUP */
/***********/
void cleanup() {
    CUDA_CHECK(cudaFree(reinterpret_cast<void*>(sbt.raygenRecord)));
    CUDA_CHECK(cudaFree(reinterpret_cast<void*>(sbt.missRecordBase)));
    CUDA_CHECK(cudaFree(reinterpret_cast<void*>(sbt.hitgroupRecordBase)));
    CUDA_CHECK(cudaFree(reinterpret_cast<void*>(d_gas_output_buffer)));

    OPTIX_CHECK(optixPipelineDestroy(pipeline));
    OPTIX_CHECK(optixProgramGroupDestroy(hitgroup_prog_group));
    OPTIX_CHECK(optixProgramGroupDestroy(miss_prog_group));
    OPTIX_CHECK(optixProgramGroupDestroy(raygen_prog_group));
    OPTIX_CHECK(optixModuleDestroy(module));

    OPTIX_CHECK(optixDeviceContextDestroy(context));
}


/********/
/* MAIN */
/********/
int main( int argc, char* argv[] )
{
    std::string outfile;
    int         width  = 1024;
    int         height =  768;

    for( int i = 1; i < argc; ++i )
    {
        const std::string arg( argv[i] );
        if( arg == "--help" || arg == "-h" )
        {
            printUsageAndExit( argv[0] );
        }
        else if( arg == "--file" || arg == "-f" )
        {
            if( i < argc - 1 )
            {
                outfile = argv[++i];
            }
            else
            {
                printUsageAndExit( argv[0] );
            }
        }
        else if( arg.substr( 0, 6 ) == "--dim=" )
        {
            const std::string dims_arg = arg.substr( 6 );
            sutil::parseDimensions( dims_arg.c_str(), width, height );
        }
        else
        {
            std::cerr << "Unknown option '" << arg << "'\n";
            printUsageAndExit( argv[0] );
        }
    }

    try
    {
        char log[2048]; // For error reporting from OptiX creation functions

        // --- CUDA and OptiX initialization
        init();

        std::vector<float3> vertices;
        std::vector<int3> indices;
        // --- Generate mesh
        generateMesh(vertices, indices);

        // --- Acceleration handling
        accelerationHandling(vertices, indices);

        // --- Create module
        createModule();

        // --- Create program groups
        createProgramGroups();

        // --- Link pipeline
        createPipeline();

        float3 color; color.x = .2f; color.y = .8f; color.z = .2f;
        
        // --- Set up shader binding table
        createSBT(color);

        sutil::CUDAOutputBuffer<uchar4> output_buffer( sutil::CUDAOutputBufferType::CUDA_DEVICE, width, height );

        // --- Launch
        launch(output_buffer, width, height);

        // --- Display results
        displayResults(output_buffer, outfile, argv, width, height);

        // --- Cleanup
        cleanup();
    }
    catch( std::exception& e )
    {
        std::cerr << "Caught exception: " << e.what() << "\n";
        return 1;
    }
    return 0;
}