// File: corpus_test.cpp
#include "crn_core.h"
#include "corpus_test.h"
#include "crn_find_files.h"
#include "crn_console.h"
#include "crn_image_utils.h"
#include "crn_hash.h"
#include "crn_hash_map.h"
#include "crn_radix_sort.h"
#include "crn_mipmapped_texture.h"

namespace crnlib
{
   corpus_tester::corpus_tester()
   {
      m_bad_block_img.resize(256, 256);
      m_next_bad_block_index = 0;
      m_total_bad_block_files = 0;
   }

   void corpus_tester::print_comparative_metric_stats(const command_line_params& cmd_line_params, const crnlib::vector<image_utils::error_metrics>& stats1, const crnlib::vector<image_utils::error_metrics>& stats2, uint num_blocks_x, uint num_blocks_y)
   {
      num_blocks_y;

      crnlib::vector<uint> better_blocks;
      crnlib::vector<uint> equal_blocks;
      crnlib::vector<uint> worse_blocks;
      crnlib::vector<float> delta_psnr;

      for (uint i = 0; i < stats1.size(); i++)
      {
         //uint bx = i % num_blocks_x;
         //uint by = i / num_blocks_x;

         const image_utils::error_metrics& em1 = stats1[i];
         const image_utils::error_metrics& em2 = stats2[i];

         if (em1.mPeakSNR < em2.mPeakSNR)
         {
            worse_blocks.push_back(i);
            delta_psnr.push_back((float)(em2.mPeakSNR - em1.mPeakSNR));
         }
         else if (fabs(em1.mPeakSNR - em2.mPeakSNR) < .001f)
            equal_blocks.push_back(i);
         else
            better_blocks.push_back(i);
      }  

      console::printf("Num worse blocks: %u, %3.3f%%", worse_blocks.size(), worse_blocks.size() * 100.0f / stats1.size());          
      console::printf("Num equal blocks: %u, %3.3f%%", equal_blocks.size(), equal_blocks.size() * 100.0f / stats1.size());
      console::printf("Num better blocks: %u, %3.3f%%", better_blocks.size(), better_blocks.size() * 100.0f / stats1.size());
      console::printf("Num equal+better blocks: %u, %3.3f%%", equal_blocks.size()+better_blocks.size(), (equal_blocks.size()+better_blocks.size()) * 100.0f / stats1.size());

      if (!cmd_line_params.has_key("nobadblocks"))
      {
         crnlib::vector<uint> indices[2];
         indices[0].resize(worse_blocks.size());
         indices[1].resize(worse_blocks.size());

         uint* pSorted_indices = NULL;
         if (worse_blocks.size())
         {
            pSorted_indices = indirect_radix_sort(worse_blocks.size(), &indices[0][0], &indices[1][0], &delta_psnr[0], 0, sizeof(float), true);

            console::printf("List of worse blocks sorted by delta PSNR:");
            for (uint i = 0; i < worse_blocks.size(); i++)
            {
               uint block_index = worse_blocks[pSorted_indices[i]];
               uint bx = block_index % num_blocks_x;
               uint by = block_index / num_blocks_x;

               console::printf("%u. [%u,%u] %3.3f %3.3f %3.3f", 
                  i,
                  bx, by,
                  stats1[block_index].mPeakSNR,
                  stats2[block_index].mPeakSNR,
                  stats2[block_index].mPeakSNR - stats1[block_index].mPeakSNR);
            }
         }         
      }         
   }

   void corpus_tester::print_metric_stats(const crnlib::vector<image_utils::error_metrics>& stats, uint num_blocks_x, uint num_blocks_y)
   {
      num_blocks_y;

      image_utils::error_metrics best_metrics;
      image_utils::error_metrics worst_metrics;
      worst_metrics.mPeakSNR = 1e+6f;

      vec2I best_loc;
      vec2I worst_loc;
      utils::zero_object(best_loc);
      utils::zero_object(worst_loc);

      double psnr_total = 0.0f;
      double psnr2_total = 0.0f;
      uint num_non_inf = 0;
      uint num_inf = 0;

      for (uint i = 0; i < stats.size(); i++)
      {
         uint bx = i % num_blocks_x;
         uint by = i / num_blocks_x;

         const image_utils::error_metrics& em = stats[i];

         if ((em.mPeakSNR < 200.0f) && (em > best_metrics)) { best_metrics = em; best_loc.set(bx, by); }
         if (em < worst_metrics) { worst_metrics = em; worst_loc.set(bx, by); }

         if (em.mPeakSNR < 200.0f)
         {
            psnr_total += em.mPeakSNR;
            psnr2_total += em.mPeakSNR*em.mPeakSNR;
            num_non_inf++;
         }
         else
         {
            num_inf++;
         }
      }         

      console::printf("Number of infinite PSNR blocks: %u", num_inf);
      console::printf("Number of non-infinite PSNR blocks: %u", num_non_inf);
      if (num_non_inf)
      {
         psnr_total /= num_non_inf;
         psnr2_total /= num_non_inf;

         double psnr_std_dev = sqrt(psnr2_total - psnr_total * psnr_total);

         console::printf("Average Non-Inf PSNR: %3.3f, Std dev: %3.3f", psnr_total, psnr_std_dev);
         console::printf("Worst PSNR: %3.3f, Block Location: %i,%i", worst_metrics.mPeakSNR, worst_loc[0], worst_loc[1]);
         console::printf("Best Non-Inf PSNR: %3.3f, Block Location: %i,%i", best_metrics.mPeakSNR, best_loc[0], best_loc[1]);
      }         
   }

   void corpus_tester::flush_bad_blocks()
   {
      if (!m_next_bad_block_index)
         return;

      dynamic_string filename(cVarArg, "badblocks_%u.tga", m_total_bad_block_files);
      console::printf("Writing bad block image: %s", filename.get_ptr());
      image_utils::write_to_file(filename.get_ptr(), m_bad_block_img, image_utils::cWriteFlagIgnoreAlpha);

      m_bad_block_img.set_all(color_quad_u8::make_black());

      m_total_bad_block_files++;

      m_next_bad_block_index = 0;
   }

   void corpus_tester::add_bad_block(image_u8& block)
   {
      uint num_blocks_x = m_bad_block_img.get_block_width(4);
      uint num_blocks_y = m_bad_block_img.get_block_height(4);
      uint total_blocks = num_blocks_x * num_blocks_y;

      m_bad_block_img.blit((m_next_bad_block_index % num_blocks_x) * 4, (m_next_bad_block_index / num_blocks_x) * 4, block);
      m_next_bad_block_index++;

      if (m_next_bad_block_index == total_blocks)
         flush_bad_blocks();
   }
   
   static bool progress_callback(uint percentage_complete, void* pUser_data_ptr)
   {
      static int s_prev_percentage_complete = -1;
      pUser_data_ptr;
      if (s_prev_percentage_complete != static_cast<int>(percentage_complete))
      {
         console::progress("%u%%", percentage_complete);
         s_prev_percentage_complete = percentage_complete;
      }
      return true;
   }

   bool corpus_tester::test(const char* pCmd_line)
   {
      console::printf("Command line:\n\"%s\"", pCmd_line);

      static const command_line_params::param_desc param_desc_array[] = 
      {
         { "corpus_test", 0, false },
         { "in", 1, true },
         { "deep", 0, false },
         { "alpha", 0, false },
         { "nomips", 0, false },
         { "perceptual", 0, false },
         { "endpointcaching", 0, false },
         { "multithreaded", 0, false },
         { "writehybrid", 0, false },
         { "nobadblocks", 0, false },
      };

      command_line_params cmd_line_params;
      if (!cmd_line_params.parse(pCmd_line, CRNLIB_ARRAY_SIZE(param_desc_array), param_desc_array, true))
         return false;

      double total_time1 = 0, total_time2 = 0;

      command_line_params::param_map_const_iterator it = cmd_line_params.begin();
      for ( ; it != cmd_line_params.end(); ++it)
      {
         if (it->first != "in")
            continue;
         if (it->second.m_values.empty())
         {
            console::error("Must follow /in parameter with a filename!\n");
            return false;
         }

         for (uint in_value_index = 0; in_value_index < it->second.m_values.size(); in_value_index++)
         {
            const dynamic_string& filespec = it->second.m_values[in_value_index];

            find_files file_finder;
            if (!file_finder.find(filespec.get_ptr(), find_files::cFlagAllowFiles | (cmd_line_params.has_key("deep") ? find_files::cFlagRecursive : 0)))
            {
               console::warning("Failed finding files: %s", filespec.get_ptr());
               continue;
            }

            if (file_finder.get_files().empty())
            {
               console::warning("No files found: %s", filespec.get_ptr());
               return false;
            }

            const find_files::file_desc_vec& files = file_finder.get_files();

            image_u8 o(4, 4), a(4, 4), b(4, 4);

            uint first_channel = 0;
            uint num_channels = 3;
            bool perceptual = cmd_line_params.get_value_as_bool("perceptual", false);
            if (perceptual)
            {
               first_channel = 0;
               num_channels = 0;
            }
            console::printf("Perceptual mode: %u", perceptual);

            for (uint file_index = 0; file_index < files.size(); file_index++)
            {
               const find_files::file_desc& file_desc = files[file_index];

               console::printf("-------- Loading image: %s", file_desc.m_fullname.get_ptr());

               image_u8 img;
               if (!image_utils::read_from_file(img, file_desc.m_fullname.get_ptr(), 0))
               {
                  console::warning("Failed loading image file: %s", file_desc.m_fullname.get_ptr());
                  continue;
               }

               if ((!cmd_line_params.has_key("alpha")) && img.is_component_valid(3))
               {
                  for (uint y = 0; y < img.get_height(); y++)
                     for (uint x = 0; x < img.get_width(); x++)
                        img(x, y).a = 255;

                  img.set_component_valid(3, false);
               }

               mipmapped_texture orig_tex;
               orig_tex.assign(crnlib_new<image_u8>(img));

               if (!cmd_line_params.has_key("nomips"))
               {
                  mipmapped_texture::generate_mipmap_params genmip_params;
                  genmip_params.m_srgb = true;

                  console::printf("Generating mipmaps");

                  if (!orig_tex.generate_mipmaps(genmip_params, false))
                  {
                     console::error("Mipmap generation failed!");
                     return false;
                  }
               }               

               console::printf("Compress 1");

               mipmapped_texture tex1(orig_tex);
               dxt_image::pack_params convert_params;
               convert_params.m_endpoint_caching = cmd_line_params.get_value_as_bool("endpointcaching", 0, false);
               convert_params.m_compressor = cCRNDXTCompressorCRN;
               convert_params.m_quality = cCRNDXTQualityNormal;
               convert_params.m_perceptual = perceptual;
               convert_params.m_num_helper_threads = cmd_line_params.get_value_as_bool("multithreaded", 0, true) ? (g_number_of_processors - 1) : 0;
               convert_params.m_pProgress_callback = progress_callback;
               timer t;
               t.start();
               if (!tex1.convert(PIXEL_FMT_ETC1, false, convert_params))
               {
                  console::error("Texture conversion failed!");
                  return false;
               }
               double time1 = t.get_elapsed_secs();
               total_time1 += time1;
               console::printf("Elapsed time: %3.3f", time1);

               console::printf("Compress 2");

               mipmapped_texture tex2(orig_tex);
               convert_params.m_endpoint_caching = false;
               convert_params.m_compressor = cCRNDXTCompressorCRN;
               convert_params.m_quality = cCRNDXTQualitySuperFast;
               t.start();
               if (!tex2.convert(PIXEL_FMT_ETC1, false, convert_params))
               {
                  console::error("Texture conversion failed!");
                  return false;
               }
               double time2 = t.get_elapsed_secs();
               total_time2 += time2;
               console::printf("Elapsed time: %3.3f", time2);

               image_u8 hybrid_img(img.get_width(), img.get_height());

               for (uint l = 0; l < orig_tex.get_num_levels(); l++)
               {
                  image_u8 orig_img, img1, img2;

                  image_u8* pOrig = orig_tex.get_level(0, l)->get_unpacked_image(orig_img, cUnpackFlagUncook | cUnpackFlagUnflip);
                  image_u8* pImg1 = tex1.get_level(0, l)->get_unpacked_image(img1, cUnpackFlagUncook | cUnpackFlagUnflip);
                  image_u8* pImg2 = tex2.get_level(0, l)->get_unpacked_image(img2, cUnpackFlagUncook | cUnpackFlagUnflip);

                  const uint num_blocks_x = pOrig->get_block_width(4);
                  const uint num_blocks_y = pOrig->get_block_height(4);

                  crnlib::vector<image_utils::error_metrics> metrics[2];

                  for (uint by = 0; by < num_blocks_y; by++)
                  {
                     for (uint bx = 0; bx < num_blocks_x; bx++)
                     {
                        pOrig->extract_block(o.get_ptr(), bx * 4, by * 4, 4, 4);
                        pImg1->extract_block(a.get_ptr(), bx * 4, by * 4, 4, 4);
                        pImg2->extract_block(b.get_ptr(), bx * 4, by * 4, 4, 4);

                        image_utils::error_metrics em1;
                        em1.compute(o, a, first_channel, num_channels);

                        image_utils::error_metrics em2;
                        em2.compute(o, b, first_channel, num_channels);

                        metrics[0].push_back(em1);
                        metrics[1].push_back(em2);

                        if (em1.mPeakSNR < em2.mPeakSNR)
                        {
                           add_bad_block(o);

                           hybrid_img.blit(bx * 4, by * 4, b);
                        }
                        else
                        {
                           hybrid_img.blit(bx * 4, by * 4, a);
                        }
                     }
                  }

                  if (cmd_line_params.has_key("writehybrid"))
                     image_utils::write_to_file("hybrid.tga", hybrid_img, image_utils::cWriteFlagIgnoreAlpha);

                  console::printf("---- Mip level: %u, Total blocks: %ux%u, %u", l, num_blocks_x, num_blocks_y, num_blocks_x * num_blocks_y);

                  console::printf("Compressor 1:");
                  print_metric_stats(metrics[0], num_blocks_x, num_blocks_y);

                  console::printf("Compressor 2:");
                  print_metric_stats(metrics[1], num_blocks_x, num_blocks_y);

                  console::printf("Compressor 1 vs. 2:");
                  print_comparative_metric_stats(cmd_line_params, metrics[0], metrics[1], num_blocks_x, num_blocks_y);

                  image_utils::error_metrics em;

                  em.compute(*pOrig, *pImg1, 0, perceptual ? 0 : 3);
                  em.print("Compressor 1: ");

                  em.compute(*pOrig, *pImg2, 0, perceptual ? 0 : 3);
                  em.print("Compressor 2: ");

                  em.compute(*pOrig, hybrid_img, 0, perceptual ? 0 : 3);
                  em.print("Best of Both: ");
               }
            }

         }  // file_index
      } 

      flush_bad_blocks();

      console::printf("Total times: %4.3f vs. %4.3f", total_time1, total_time2);

      return true;
   }

} // namespace crnlib