/*
* This code implements the algorithm described in our paper "Semantic Background Subtraction", ICIP 2017. See file "README.txt" for usage.
* Copyright - Marc Braham - August 2017
*/

#include <iostream>
#include <stdexcept>
#include <string>
#include <sstream>
#include <ctime>
#include <boost/program_options.hpp>
#include <opencv2/core/core.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/video/video.hpp>
#include <opencv2/highgui/highgui.hpp>

#define SIZE_PSEUDO_RANDOM_SEQUENCE 500000

using namespace std;
using namespace boost::program_options;
using namespace cv;

///Allocate a 2D array of booleans
bool** get2DArrayBool(int SIZE_1D, int SIZE_2D)
{
  bool** array = new bool* [SIZE_1D];
  for (int i = 0; i < SIZE_1D; i++)
  {
    array[i] = new bool[SIZE_2D];
    for (int j = 0; j < SIZE_2D; j++)
    {
      array[i][j] = false;
    }
  }
  return array;
}
///Deallocate a 2D array of booleans
void delete2DArrayBool(bool** array, int SIZE_1D)
{
  for(int i = 0; i < SIZE_1D; i++)
  {
    delete array[i];
  }
  delete array;
}

int main(int argc, char const *argv[])
{
    ///Parsing command-line arguments 
	options_description optDesc(
    string("Copyright - Marc Braham - 2017\n\n")
    );
    
    optDesc.add_options()
      ("help,h", "print this help message")
      ("bgs,b", value<string>(), "path to the folder of BGS masks")
      ("semantics,s", value<string>(), "path to the folder of semantic masks")
      ("input,i", value<string>(), "path to the folder of input images (optional)")
      ("groundtruth,g", value<string>(), "path to the folder of groundtruth images (optional)")
      ("results,r", value<string>(), "path to the folder of results (optional)")
      ("visualization,v", "visualize everything in windows (optional)")
      ("tau_BG,d", value<int>()->default_value(240), "tau BG threshold for 16-bit images (value should belong to [-1,65535]). Note that the value -1 deactivates rule 1.")
      ("tau_FG,u", value<int>()->default_value(226), "tau FG threshold for 8-bit images (value should belong to [0,256]). Note that the value 256 deactivates rule 2.")
      ("modulo_update,m", value<int>()->default_value(4096), "set the update rate of the semantic BG model")
    ;

    variables_map varsMap;
    store(parse_command_line(argc, argv, optDesc), varsMap);
    notify(varsMap);

    /* Help message. */
    if (varsMap.count("help")) 
    {
      cout << optDesc << endl;
      return EXIT_SUCCESS;
    }

    /* Sanity check. */
    if (!varsMap.count("bgs"))
      throw runtime_error("path to directory containing the original background subtraction maps should be set");

    if (!varsMap.count("semantics"))
      throw runtime_error("path to the directory containing the semantic maps should be set");

    ///Setting variables
    String folder_BGS, folder_Semantic, folder_Input, folder_Groundtruth;
    vector<String> filenames_BGS, filenames_Semantic, filenames_Input, filenames_Groundtruth;

    folder_BGS = varsMap["bgs"].as<string>() + "/*.png";
    folder_Semantic = varsMap["semantics"].as<string>() + "/*.png";
    if (varsMap.count("input"))
      folder_Input = varsMap["input"].as<string>() + "/*.jpg";
    if (varsMap.count("groundtruth"))
      folder_Groundtruth = varsMap["groundtruth"].as<string>() + "/*.png";

    glob(folder_BGS, filenames_BGS);
    glob(folder_Semantic, filenames_Semantic);
    if (varsMap.count("input"))
      glob(folder_Input, filenames_Input);
    if (varsMap.count("groundtruth"))
      glob(folder_Groundtruth, filenames_Groundtruth);
    
    stringstream ss_Output;
    vector<int> compression_params;
    compression_params.push_back(CV_IMWRITE_PNG_COMPRESSION);
    compression_params.push_back(0);

    Mat BGS[1];
    Mat Semantic[1];
    Mat Input[1];
    Mat Groundtruth[1];
    Mat Output[1];

    int height,width,al;

    double TP_BGS(0),TN_BGS(0),FP_BGS(0),FN_BGS(0);
    double Pr_BGS,TPR_BGS,TNR_BGS;
    double Accuracy_BGS,F1_BGS;
    double TP_Semantic_BGS(0),TN_Semantic_BGS(0),FP_Semantic_BGS(0),FN_Semantic_BGS(0);
    double Pr_Semantic_BGS,TPR_Semantic_BGS,TNR_Semantic_BGS;
    double Accuracy_Semantic_BGS,F1_Semantic_BGS;


    ///Setting look-up tables for faster operations
    unsigned short* convert_8bit = new unsigned short[65536];
    for (int s = 0 ; s < 65536 ; s++)
    {
      convert_8bit[s] = s/256;
    }

    bool* rule_BG = new bool[65536];
    for (int s = 0 ; s < 65536 ; s++)
    {
       if (s <= varsMap["tau_BG"].as<int>())
         rule_BG[s] = true;
       else
         rule_BG[s] = false;
    }

    bool** rule_FG = get2DArrayBool(256,256);
    for (int s1 = 0 ; s1 < 256 ; s1++)
    {
       for (int s2 = 0 ; s2 < 256 ; s2++)
       {
          if ((s2 - s1) >= varsMap["tau_FG"].as<int>())
            rule_FG[s1][s2] = true;
          else
            rule_FG[s1][s2] = false;
       }
    }

    srand(0);
    bool* update_sequence = new bool[SIZE_PSEUDO_RANDOM_SEQUENCE];
    int*  next_index = new int[SIZE_PSEUDO_RANDOM_SEQUENCE];
    for (int index = 0 ; index < SIZE_PSEUDO_RANDOM_SEQUENCE ; index++)
    {
       al =  rand()%varsMap["modulo_update"].as<int>();
       if (al == 0)
         update_sequence[index] = true;
       else
         update_sequence[index] = false;
       next_index[index] = index+1;
    }
    next_index[SIZE_PSEUDO_RANDOM_SEQUENCE - 1] = 0;
    int index(0);

    ///Initialization of the semantic BG model
    Semantic[0] = imread(filenames_Semantic[0],CV_LOAD_IMAGE_ANYDEPTH);
    if(Semantic[0].empty())
    {
       cout << "Path to the first semantic mask is incorrect" << endl;
       exit(EXIT_FAILURE);
    }
    height = Semantic[0].rows;
    width = Semantic[0].cols;

    Mat_<unsigned short> BG_Semantic(height,width);
    for (int i = 0 ; i < height ; i++)
    {
       for (int j = 0 ; j < width ; j++)
       {
           BG_Semantic.at<unsigned short>(i,j) = Semantic[0].at<unsigned short>(i,j);
       }
    }

    ///Setting windows positions
    namedWindow("Semantic map");
    namedWindow("Original background subtraction map");
    namedWindow("Final result (after using our method)");
    if (varsMap.count("input"))
    {
       namedWindow("Input frame");
    }
    if (varsMap.count("groundtruth"))
    {
       namedWindow("Groundtruth mask");
    }

    moveWindow("Semantic map", 20,20);
    moveWindow("Original background subtraction map",width+40,20);
    moveWindow("Final result (after using our method)",2*(width+20)+20,20);
    if (varsMap.count("input"))
    {
       moveWindow("Input frame",20,height+120);
    }
    if (varsMap.count("groundtruth"))
    {
       moveWindow("Groundtruth mask",width+40,height+120);
    }


    ///Semantic BGS
    for (int k = 0 ; k < filenames_BGS.size() ; k++)
    {
       BGS[0] = imread(filenames_BGS[k],CV_LOAD_IMAGE_GRAYSCALE);
       Semantic[0] = imread(filenames_Semantic[k], CV_LOAD_IMAGE_ANYDEPTH);
       if(BGS[0].empty())
       {
          cout << "Error in loading BGS mask number " << k + 1 << endl;
          exit(EXIT_FAILURE);
       }
       if(Semantic[0].empty())
       {
          cout << "Error in loading semantic mask number " << k + 1 << endl;
          exit(EXIT_FAILURE);
       }

       if (varsMap.count("input"))
       {
          Input[0] = imread(filenames_Input[k],CV_LOAD_IMAGE_COLOR);
          if(Input[0].empty())
          {
             cout << "Error in loading input mask number " << k + 1 << endl;
             exit(EXIT_FAILURE);
          }
       }
       if (varsMap.count("groundtruth"))
       {
          Groundtruth[0] = imread(filenames_Groundtruth[k],CV_LOAD_IMAGE_GRAYSCALE);
          if(Groundtruth[0].empty())
          {
             cout << "Error in loading groundtruth mask number " << k + 1 << endl;
             exit(EXIT_FAILURE);
          }
       }

       Output[0] = BGS[0].clone();
       for (int i = 0 ; i < height ; i++)
       {
          for (int j = 0 ; j < width ; j++)
          {
              if (rule_BG[Semantic[0].at<unsigned short>(i,j)])
                Output[0].at<uchar>(i,j) = 0;
              else if (rule_FG[convert_8bit[BG_Semantic.at<unsigned short>(i,j)]][convert_8bit[Semantic[0].at<unsigned short>(i,j)]])
                Output[0].at<uchar>(i,j) = 255;

              if (Output[0].at<uchar>(i,j) == 0)
              {
                  index = next_index[index];
                  if (update_sequence[index])
                    BG_Semantic.at<unsigned short>(i,j) = Semantic[0].at<unsigned short>(i,j);
              }
          }
       }

       if (varsMap.count("groundtruth"))
       {
          for (int i = 0 ; i < height ; i++)
          {
             for (int j = 0 ; j < width ; j++)
             {
                if ((Groundtruth[0].at<uchar>(i,j) == 0) or (Groundtruth[0].at<uchar>(i,j) == 50))
                {
                   if (Output[0].at<uchar>(i,j) == 0)
                     TN_Semantic_BGS += 1;
                   else
                     FP_Semantic_BGS += 1;
                    if (BGS[0].at<uchar>(i,j) == 0)
                     TN_BGS += 1;
                   else
                     FP_BGS += 1;
                }
                else if (Groundtruth[0].at<uchar>(i,j) == 255)
                {
                   if (Output[0].at<uchar>(i,j) == 0)
                     FN_Semantic_BGS += 1;
                   else
                     TP_Semantic_BGS += 1;
                   if (BGS[0].at<uchar>(i,j) == 0)
                     FN_BGS += 1;
                   else
                     TP_BGS += 1;
                }
             }
          }
       }

       if (varsMap.count("results"))
       {
            ss_Output.clear();
            ss_Output.str("");
            if (k+1 < 10)
              ss_Output << varsMap["results"].as<string>() << "/out00000" << k+1 << ".png";
            else if (k+1 < 100)
              ss_Output << varsMap["results"].as<string>() << "/out0000" << k+1 << ".png";
            else if (k+1 < 1000)
              ss_Output << varsMap["results"].as<string>() << "/out000" << k+1 << ".png";
            else if (k+1 < 10000)
              ss_Output << varsMap["results"].as<string>() << "/out00" << k+1 << ".png";
            else if (k+1 < 100000)
              ss_Output << varsMap["results"].as<string>() << "/out0" << k+1 << ".png";
            else if (k+1 < 1000000)
              ss_Output << varsMap["results"].as<string>() << "/out" << k+1 << ".png";
            imwrite(ss_Output.str(),Output[0], compression_params);
       }

       if (varsMap.count("visualization"))
       {
           imshow("Original background subtraction map", BGS[0]);
           imshow("Semantic map", Semantic[0]);
           if (varsMap.count("input"))
             imshow("Input frame", Input[0]);
           if (varsMap.count("groundtruth"))
             imshow("Groundtruth mask", Groundtruth[0]);
           imshow("Final result (after using our method)", Output[0]);
           waitKey(1);
       }  
    }

    if (varsMap.count("groundtruth"))
    {
       Pr_BGS = TP_BGS*1.0/(TP_BGS+FP_BGS);
       TPR_BGS = TP_BGS*1.0/(TP_BGS+FN_BGS);
       TNR_BGS = TN_BGS*1.0/(TN_BGS+FP_BGS);
       Accuracy_BGS = (TP_BGS+TN_BGS)*1.0/(TP_BGS+TN_BGS+FP_BGS+FN_BGS);
       F1_BGS = 2.0*Pr_BGS*TPR_BGS/(Pr_BGS+TPR_BGS);

       Pr_Semantic_BGS = TP_Semantic_BGS*1.0/(TP_Semantic_BGS+FP_Semantic_BGS);
       TPR_Semantic_BGS = TP_Semantic_BGS*1.0/(TP_Semantic_BGS+FN_Semantic_BGS);
       TNR_Semantic_BGS = TN_Semantic_BGS*1.0/(TN_Semantic_BGS+FP_Semantic_BGS);
       Accuracy_Semantic_BGS = (TP_Semantic_BGS+TN_Semantic_BGS)*1.0/(TP_Semantic_BGS+TN_Semantic_BGS+FP_Semantic_BGS+FN_Semantic_BGS);
       F1_Semantic_BGS = 2.0*Pr_Semantic_BGS*TPR_Semantic_BGS/(Pr_Semantic_BGS+TPR_Semantic_BGS);

       cout << endl << "RESULTS" << endl << endl;
       
       cout << "True Positive Rate (TPR) of the original background subtraction algorithm  = " << 100.0*TPR_BGS << "%" << endl;   
       cout << "True Positive Rate (TPR) after using our method based on semantic          = " << 100.0*TPR_Semantic_BGS << "%" << endl << endl;
       
       cout << "True Negative Rate (TNR) of the original background subtraction algorithm  = " << 100.0*TNR_BGS << "%" << endl;
       cout << "True Negative Rate (TNR) after using our method based on semantic          = " << 100.0*TNR_Semantic_BGS << "%" << endl << endl;

       cout << "Accuracy of the original background subtraction algorithm = " << 100.0*Accuracy_BGS << "%" << endl;
       cout << "Accuracy after using our method based on semantic         = " << 100.0*Accuracy_Semantic_BGS << "%" << endl << endl;

       cout << "F1 score of the original background subtraction algorithm = " << F1_BGS << endl;
       cout << "F1 score after using our method based on semantic         = " << F1_Semantic_BGS << endl << endl;
    }

    BG_Semantic.release();
    BGS[0].release();
    Semantic[0].release();
    Output[0].release();
    Input[0].release();
    Groundtruth[0].release();
    delete[] rule_BG;
    delete[] convert_8bit;
    delete2DArrayBool(rule_FG,256);
    delete[] update_sequence;
    delete[] next_index;

    return 0;
}
