/******************************************************************************
   * Copyright (c) 2011-2018, NVIDIA CORPORATION.  All rights reserved.
   *
   * Redistribution and use in source and binary forms, with or without
   * modification, are permitted provided that the following conditions are met:
   *     * Redistributions of source code must retain the above copyright
   *       notice, this list of conditions and the following disclaimer.
   *     * Redistributions in binary form must reproduce the above copyright
   *       notice, this list of conditions and the following disclaimer in the
   *       documentation and/or other materials provided with the distribution.
   *     * Neither the name of the NVIDIA CORPORATION nor the
   *       names of its contributors may be used to endorse or promote products
   *       derived from this software without specific prior written permission.
   *
   * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
   * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
   * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
   * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
   * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
   * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
   * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
   * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
   * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
   * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
   *
   ******************************************************************************/
  
  #include <stdio.h>
  #include <map>
  #include <vector>
  #include <algorithm>
  #include <cstdio>
  #include <fstream>
  
  #include "histogram/histogram_gmem_atomics.h"
  #include "histogram/histogram_smem_atomics.h"
  #include "histogram/histogram_cub.h"
  
  #include <cub/util_allocator.cuh>
  #include <test/test_util.h>
  
  using namespace cub;
  
  //---------------------------------------------------------------------
  // Globals, constants, and type declarations
  //---------------------------------------------------------------------
  
  // Ensure printing of CUDA runtime errors to console
  #define CUB_STDERR
  
  bool                    g_verbose = false;  // Whether to display input/output to console
  bool                    g_report = false;   // Whether to display a full report in CSV format
  CachingDeviceAllocator  g_allocator(true);  // Caching allocator for device memory
  
  struct less_than_value
  {
      inline bool operator()(
          const std::pair<std::string, double> &a,
          const std::pair<std::string, double> &b)
      {
          return a.second < b.second;
      }
  };
  
  
  //---------------------------------------------------------------------
  // Targa (.tga) image file parsing
  //---------------------------------------------------------------------
  
  /**
   * TGA image header info
   */
  struct TgaHeader
  {
      char idlength;
      char colormaptype;
      char datatypecode;
      short colormaporigin;
      short colormaplength;
      char colormapdepth;
      short x_origin;
      short y_origin;
      short width;
      short height;
      char bitsperpixel;
      char imagedescriptor;
  
      void Parse (FILE *fptr)
      {
          idlength = fgetc(fptr);
          colormaptype = fgetc(fptr);
          datatypecode = fgetc(fptr);
          fread(&colormaporigin, 2, 1, fptr);
          fread(&colormaplength, 2, 1, fptr);
          colormapdepth = fgetc(fptr);
          fread(&x_origin, 2, 1, fptr);
          fread(&y_origin, 2, 1, fptr);
          fread(&width, 2, 1, fptr);
          fread(&height, 2, 1, fptr);
          bitsperpixel = fgetc(fptr);
          imagedescriptor = fgetc(fptr);
      }
  
      void Display (FILE *fptr)
      {
          fprintf(fptr, "ID length:           %d
  ", idlength);
          fprintf(fptr, "Color map type:      %d
  ", colormaptype);
          fprintf(fptr, "Image type:          %d
  ", datatypecode);
          fprintf(fptr, "Color map offset:    %d
  ", colormaporigin);
          fprintf(fptr, "Color map length:    %d
  ", colormaplength);
          fprintf(fptr, "Color map depth:     %d
  ", colormapdepth);
          fprintf(fptr, "X origin:            %d
  ", x_origin);
          fprintf(fptr, "Y origin:            %d
  ", y_origin);
          fprintf(fptr, "Width:               %d
  ", width);
          fprintf(fptr, "Height:              %d
  ", height);
          fprintf(fptr, "Bits per pixel:      %d
  ", bitsperpixel);
          fprintf(fptr, "Descriptor:          %d
  ", imagedescriptor);
      }
  };
  
  
  /**
   * Decode image byte data into pixel
   */
  void ParseTgaPixel(uchar4 &pixel, unsigned char *tga_pixel, int bytes)
  {
      if (bytes == 4)
      {
          pixel.x = tga_pixel[2];
          pixel.y = tga_pixel[1];
          pixel.z = tga_pixel[0];
          pixel.w = tga_pixel[3];
      }
      else if (bytes == 3)
      {
          pixel.x = tga_pixel[2];
          pixel.y = tga_pixel[1];
          pixel.z = tga_pixel[0];
          pixel.w = 0;
      }
      else if (bytes == 2)
      {
          pixel.x = (tga_pixel[1] & 0x7c) << 1;
          pixel.y = ((tga_pixel[1] & 0x03) << 6) | ((tga_pixel[0] & 0xe0) >> 2);
          pixel.z = (tga_pixel[0] & 0x1f) << 3;
          pixel.w = (tga_pixel[1] & 0x80);
      }
  }
  
  
  /**
   * Reads a .tga image file
   */
  void ReadTga(uchar4* &pixels, int &width, int &height, const char *filename)
  {
      // Open the file
      FILE *fptr;
      if ((fptr = fopen(filename, "rb")) == NULL)
      {
          fprintf(stderr, "File open failed
  ");
          exit(-1);
      }
  
      // Parse header
      TgaHeader header;
      header.Parse(fptr);
  //    header.Display(stdout);
      width = header.width;
      height = header.height;
  
      // Verify compatibility
      if (header.datatypecode != 2 && header.datatypecode != 10)
      {
          fprintf(stderr, "Can only handle image type 2 and 10
  ");
          exit(-1);
      }
      if (header.bitsperpixel != 16 && header.bitsperpixel != 24 && header.bitsperpixel != 32)
      {
          fprintf(stderr, "Can only handle pixel depths of 16, 24, and 32
  ");
          exit(-1);
      }
      if (header.colormaptype != 0 && header.colormaptype != 1)
      {
          fprintf(stderr, "Can only handle color map types of 0 and 1
  ");
          exit(-1);
      }
  
      // Skip unnecessary header info
      int skip_bytes = header.idlength + (header.colormaptype * header.colormaplength);
      fseek(fptr, skip_bytes, SEEK_CUR);
  
      // Read the image
      int pixel_bytes = header.bitsperpixel / 8;
  
      // Allocate and initialize pixel data
      size_t image_bytes = width * height * sizeof(uchar4);
      if ((pixels == NULL) && ((pixels = (uchar4*) malloc(image_bytes)) == NULL))
      {
          fprintf(stderr, "malloc of image failed
  ");
          exit(-1);
      }
      memset(pixels, 0, image_bytes);
  
      // Parse pixels
      unsigned char   tga_pixel[5];
      int             current_pixel = 0;
      while (current_pixel < header.width * header.height)
      {
          if (header.datatypecode == 2)
          {
              // Uncompressed
              if (fread(tga_pixel, 1, pixel_bytes, fptr) != pixel_bytes)
              {
                  fprintf(stderr, "Unexpected end of file at pixel %d  (uncompressed)
  ", current_pixel);
                  exit(-1);
              }
              ParseTgaPixel(pixels[current_pixel], tga_pixel, pixel_bytes);
              current_pixel++;
          }
          else if (header.datatypecode == 10)
          {
              // Compressed
              if (fread(tga_pixel, 1, pixel_bytes + 1, fptr) != pixel_bytes + 1)
              {
                  fprintf(stderr, "Unexpected end of file at pixel %d (compressed)
  ", current_pixel);
                  exit(-1);
              }
              int run_length = tga_pixel[0] & 0x7f;
              ParseTgaPixel(pixels[current_pixel], &(tga_pixel[1]), pixel_bytes);
              current_pixel++;
  
              if (tga_pixel[0] & 0x80)
              {
                  // RLE chunk
                  for (int i = 0; i < run_length; i++)
                  {
                      ParseTgaPixel(pixels[current_pixel], &(tga_pixel[1]), pixel_bytes);
                      current_pixel++;
                  }
              }
              else
              {
                  // Normal chunk
                  for (int i = 0; i < run_length; i++)
                  {
                      if (fread(tga_pixel, 1, pixel_bytes, fptr) != pixel_bytes)
                      {
                          fprintf(stderr, "Unexpected end of file at pixel %d (normal)
  ", current_pixel);
                          exit(-1);
                      }
                      ParseTgaPixel(pixels[current_pixel], tga_pixel, pixel_bytes);
                      current_pixel++;
                  }
              }
          }
      }
  
      // Close file
      fclose(fptr);
  }
  
  
  
  //---------------------------------------------------------------------
  // Random image generation
  //---------------------------------------------------------------------
  
  /**
   * Generate a random image with specified entropy
   */
  void GenerateRandomImage(uchar4* &pixels, int width, int height, int entropy_reduction)
  {
      int num_pixels = width * height;
      size_t image_bytes = num_pixels * sizeof(uchar4);
      if ((pixels == NULL) && ((pixels = (uchar4*) malloc(image_bytes)) == NULL))
      {
          fprintf(stderr, "malloc of image failed
  ");
          exit(-1);
      }
  
      for (int i = 0; i < num_pixels; ++i)
      {
          RandomBits(pixels[i].x, entropy_reduction);
          RandomBits(pixels[i].y, entropy_reduction);
          RandomBits(pixels[i].z, entropy_reduction);
          RandomBits(pixels[i].w, entropy_reduction);
      }
  }
  
  
  
  //---------------------------------------------------------------------
  // Histogram verification
  //---------------------------------------------------------------------
  
  // Decode float4 pixel into bins
  template <int NUM_BINS, int ACTIVE_CHANNELS>
  void DecodePixelGold(float4 pixel, unsigned int (&bins)[ACTIVE_CHANNELS])
  {
      float* samples = reinterpret_cast<float*>(&pixel);
  
      for (int CHANNEL = 0; CHANNEL < ACTIVE_CHANNELS; ++CHANNEL)
          bins[CHANNEL] = (unsigned int) (samples[CHANNEL] * float(NUM_BINS));
  }
  
  // Decode uchar4 pixel into bins
  template <int NUM_BINS, int ACTIVE_CHANNELS>
  void DecodePixelGold(uchar4 pixel, unsigned int (&bins)[ACTIVE_CHANNELS])
  {
      unsigned char* samples = reinterpret_cast<unsigned char*>(&pixel);
  
      for (int CHANNEL = 0; CHANNEL < ACTIVE_CHANNELS; ++CHANNEL)
          bins[CHANNEL] = (unsigned int) (samples[CHANNEL]);
  }
  
  // Decode uchar1 pixel into bins
  template <int NUM_BINS, int ACTIVE_CHANNELS>
  void DecodePixelGold(uchar1 pixel, unsigned int (&bins)[ACTIVE_CHANNELS])
  {
      bins[0] = (unsigned int) pixel.x;
  }
  
  
  // Compute reference histogram.  Specialized for uchar4
  template <
      int         ACTIVE_CHANNELS,
      int         NUM_BINS,
      typename    PixelType>
  void HistogramGold(PixelType *image, int width, int height, unsigned int* hist)
  {
      memset(hist, 0, ACTIVE_CHANNELS * NUM_BINS * sizeof(unsigned int));
  
      for (int i = 0; i < width; i++)
      {
          for (int j = 0; j < height; j++)
          {
              PixelType pixel = image[i + j * width];
  
              unsigned int bins[ACTIVE_CHANNELS];
              DecodePixelGold<NUM_BINS>(pixel, bins);
  
              for (int CHANNEL = 0; CHANNEL < ACTIVE_CHANNELS; ++CHANNEL)
              {
                  hist[(NUM_BINS * CHANNEL) + bins[CHANNEL]]++;
              }
          }
      }
  }
  
  
  //---------------------------------------------------------------------
  // Test execution
  //---------------------------------------------------------------------
  
  /**
   * Run a specific histogram implementation
   */
  template <
      int         ACTIVE_CHANNELS,
      int         NUM_BINS,
      typename    PixelType>
  void RunTest(
      std::vector<std::pair<std::string, double> >&   timings,
      PixelType*                                      d_pixels,
      const int                                       width,
      const int                                       height,
      unsigned int *                                  d_hist,
      unsigned int *                                  h_hist,
      int                                             timing_iterations,
      const char *                                    long_name,
      const char *                                    short_name,
      double (*f)(PixelType*, int, int, unsigned int*, bool))
  {
      if (!g_report) printf("%s ", long_name); fflush(stdout);
  
      // Run single test to verify (and code cache)
      (*f)(d_pixels, width, height, d_hist, !g_report);
  
      int compare = CompareDeviceResults(h_hist, d_hist, ACTIVE_CHANNELS * NUM_BINS, true, g_verbose);
      if (!g_report) printf("\t%s
  ", compare ? "FAIL" : "PASS"); fflush(stdout);
  
      double elapsed_ms = 0;
      for (int i = 0; i < timing_iterations; i++)
      {
          elapsed_ms += (*f)(d_pixels, width, height, d_hist, false);
      }
      double avg_us = (elapsed_ms / timing_iterations) * 1000;    // average in us
      timings.push_back(std::pair<std::string, double>(short_name, avg_us));
  
      if (!g_report)
      {
          printf("Avg time %.3f us (%d iterations)
  ", avg_us, timing_iterations); fflush(stdout);
      }
      else
      {
          printf("%.3f, ", avg_us); fflush(stdout);
      }
  
      AssertEquals(0, compare);
  }
  
  
  /**
   * Evaluate corpus of histogram implementations
   */
  template <
      int         NUM_CHANNELS,
      int         ACTIVE_CHANNELS,
      int         NUM_BINS,
      typename    PixelType>
  void TestMethods(
      PixelType*  h_pixels,
      int         height,
      int         width,
      int         timing_iterations,
      double      bandwidth_GBs)
  {
      // Copy data to gpu
      PixelType* d_pixels;
      size_t pixel_bytes = width * height * sizeof(PixelType);
      CubDebugExit(g_allocator.DeviceAllocate((void**) &d_pixels, pixel_bytes));
      CubDebugExit(cudaMemcpy(d_pixels, h_pixels, pixel_bytes, cudaMemcpyHostToDevice));
  
      if (g_report) printf("%.3f, ", double(pixel_bytes) / bandwidth_GBs / 1000);
  
      // Allocate results arrays on cpu/gpu
      unsigned int *h_hist;
      unsigned int *d_hist;
      size_t histogram_bytes = NUM_BINS * ACTIVE_CHANNELS * sizeof(unsigned int);
      h_hist = (unsigned int *) malloc(histogram_bytes);
      g_allocator.DeviceAllocate((void **) &d_hist, histogram_bytes);
  
      // Compute reference cpu histogram
      HistogramGold<ACTIVE_CHANNELS, NUM_BINS>(h_pixels, width, height, h_hist);
  
      // Store timings
      std::vector<std::pair<std::string, double> > timings;
  
      // Run experiments
      RunTest<ACTIVE_CHANNELS, NUM_BINS>(timings, d_pixels, width, height, d_hist, h_hist, timing_iterations,
          "CUB", "CUB", run_cub_histogram<NUM_CHANNELS, ACTIVE_CHANNELS, NUM_BINS, PixelType>);
      RunTest<ACTIVE_CHANNELS, NUM_BINS>(timings, d_pixels, width, height, d_hist, h_hist, timing_iterations,
          "Shared memory atomics", "smem atomics", run_smem_atomics<ACTIVE_CHANNELS, NUM_BINS, PixelType>);
      RunTest<ACTIVE_CHANNELS, NUM_BINS>(timings, d_pixels, width, height, d_hist, h_hist, timing_iterations,
          "Global memory atomics", "gmem atomics", run_gmem_atomics<ACTIVE_CHANNELS, NUM_BINS, PixelType>);
  
      // Report timings
      if (!g_report)
      {
          std::sort(timings.begin(), timings.end(), less_than_value());
          printf("Timings (us):
  ");
          for (int i = 0; i < timings.size(); i++)
          {
              double bandwidth = height * width * sizeof(PixelType) / timings[i].second / 1000;
              printf("\t %.3f %s (%.3f GB/s, %.3f%% peak)
  ", timings[i].second, timings[i].first.c_str(), bandwidth, bandwidth / bandwidth_GBs * 100);
          }
          printf("
  ");
      }
  
      // Free data
      CubDebugExit(g_allocator.DeviceFree(d_pixels));
      CubDebugExit(g_allocator.DeviceFree(d_hist));
      free(h_hist);
  }
  
  
  /**
   * Test different problem genres
   */
  void TestGenres(
      uchar4*     uchar4_pixels,
      int         height,
      int         width,
      int         timing_iterations,
      double      bandwidth_GBs)
  {
      int num_pixels = width * height;
  
      {
          if (!g_report) printf("1 channel uchar1 tests (256-bin):
  
  "); fflush(stdout);
  
          size_t      image_bytes     = num_pixels * sizeof(uchar1);
          uchar1*     uchar1_pixels   = (uchar1*) malloc(image_bytes);
  
          // Convert to 1-channel (averaging first 3 channels)
          for (int i = 0; i < num_pixels; ++i)
          {
              uchar1_pixels[i].x = (unsigned char)
                  (((unsigned int) uchar4_pixels[i].x +
                    (unsigned int) uchar4_pixels[i].y +
                    (unsigned int) uchar4_pixels[i].z) / 3);
          }
  
          TestMethods<1, 1, 256>(uchar1_pixels, width, height, timing_iterations, bandwidth_GBs);
          free(uchar1_pixels);
          if (g_report) printf(", ");
      }
  
      {
          if (!g_report) printf("3/4 channel uchar4 tests (256-bin):
  
  "); fflush(stdout);
          TestMethods<4, 3, 256>(uchar4_pixels, width, height, timing_iterations, bandwidth_GBs);
          if (g_report) printf(", ");
      }
  
      {
          if (!g_report) printf("3/4 channel float4 tests (256-bin):
  
  "); fflush(stdout);
          size_t      image_bytes     = num_pixels * sizeof(float4);
          float4*     float4_pixels   = (float4*) malloc(image_bytes);
  
          // Convert to float4 with range [0.0, 1.0)
          for (int i = 0; i < num_pixels; ++i)
          {
              float4_pixels[i].x = float(uchar4_pixels[i].x) / 256;
              float4_pixels[i].y = float(uchar4_pixels[i].y) / 256;
              float4_pixels[i].z = float(uchar4_pixels[i].z) / 256;
              float4_pixels[i].w = float(uchar4_pixels[i].w) / 256;
          }
          TestMethods<4, 3, 256>(float4_pixels, width, height, timing_iterations, bandwidth_GBs);
          free(float4_pixels);
          if (g_report) printf("
  ");
      }
  }
  
  
  /**
   * Main
   */
  int main(int argc, char **argv)
  {
      // Initialize command line
      CommandLineArgs args(argc, argv);
      if (args.CheckCmdLineFlag("help"))
      {
          printf(
              "%s "
              "[--device=<device-id>] "
              "[--v] "
              "[--i=<timing iterations>] "
              "
  \t"
                  "--file=<.tga filename> "
              "
  \t"
                  "--entropy=<-1 (0%), 0 (100%), 1 (81%), 2 (54%), 3 (34%), 4 (20%), ..."
                  "[--height=<default: 1080>] "
                  "[--width=<default: 1920>] "
              "
  ", argv[0]);
          exit(0);
      }
  
      std::string         filename;
      int                 timing_iterations   = 100;
      int                 entropy_reduction   = 0;
      int                 height              = 1080;
      int                 width               = 1920;
  
      g_verbose = args.CheckCmdLineFlag("v");
      g_report = args.CheckCmdLineFlag("report");
      args.GetCmdLineArgument("i", timing_iterations);
      args.GetCmdLineArgument("file", filename);
      args.GetCmdLineArgument("height", height);
      args.GetCmdLineArgument("width", width);
      args.GetCmdLineArgument("entropy", entropy_reduction);
  
      // Initialize device
      CubDebugExit(args.DeviceInit());
  
      // Get GPU device bandwidth (GB/s)
      int device_ordinal, bus_width, mem_clock_khz;
      CubDebugExit(cudaGetDevice(&device_ordinal));
      CubDebugExit(cudaDeviceGetAttribute(&bus_width, cudaDevAttrGlobalMemoryBusWidth, device_ordinal));
      CubDebugExit(cudaDeviceGetAttribute(&mem_clock_khz, cudaDevAttrMemoryClockRate, device_ordinal));
      double bandwidth_GBs = double(bus_width) * mem_clock_khz * 2 / 8 / 1000 / 1000;
  
      // Run test(s)
      uchar4* uchar4_pixels = NULL;
      if (!g_report)
      {
          if (!filename.empty())
          {
              // Parse targa file
              ReadTga(uchar4_pixels, width, height, filename.c_str());
              printf("File %s: width(%d) height(%d)
  
  ", filename.c_str(), width, height); fflush(stdout);
          }
          else
          {
              // Generate image
              GenerateRandomImage(uchar4_pixels, width, height, entropy_reduction);
              printf("Random image: entropy-reduction(%d) width(%d) height(%d)
  
  ", entropy_reduction, width, height); fflush(stdout);
          }
  
          TestGenres(uchar4_pixels, height, width, timing_iterations, bandwidth_GBs);
      }
      else
      {
          // Run test suite
          printf("Test, MIN, RLE CUB, SMEM, GMEM, , MIN, RLE_CUB, SMEM, GMEM, , MIN, RLE_CUB, SMEM, GMEM
  ");
  
          // Entropy reduction tests
          for (entropy_reduction = 0; entropy_reduction < 5; ++entropy_reduction)
          {
              printf("entropy reduction %d, ", entropy_reduction);
              GenerateRandomImage(uchar4_pixels, width, height, entropy_reduction);
              TestGenres(uchar4_pixels, height, width, timing_iterations, bandwidth_GBs);
          }
          printf("entropy reduction -1, ");
          GenerateRandomImage(uchar4_pixels, width, height, -1);
          TestGenres(uchar4_pixels, height, width, timing_iterations, bandwidth_GBs);
          printf("
  ");
  
          // File image tests
          std::vector<std::string> file_tests;
          file_tests.push_back("animals");
          file_tests.push_back("apples");
          file_tests.push_back("sunset");
          file_tests.push_back("cheetah");
          file_tests.push_back("nature");
          file_tests.push_back("operahouse");
          file_tests.push_back("austin");
          file_tests.push_back("cityscape");
  
          for (int i = 0; i < file_tests.size(); ++i)
          {
              printf("%s, ", file_tests[i].c_str());
              std::string filename = std::string("histogram/benchmark/") + file_tests[i] + ".tga";
              ReadTga(uchar4_pixels, width, height, filename.c_str());
              TestGenres(uchar4_pixels, height, width, timing_iterations, bandwidth_GBs);
          }
      }
  
      free(uchar4_pixels);
  
      CubDebugExit(cudaDeviceSynchronize());
      printf("
  
  ");
  
      return 0;
  }