#include "naivebayes.h"

/**
 * @brief The NaiveBayes constructor constructs class and set up parametrs
 * @param classes - map which holds class counts
 */
NaiveBayes::NaiveBayes(std::map <std::string, double> classes)
{
    cls = classes;
}

/**
 * @brief The Train function trains the data making counts and probabilities of all instances
 * @param Data - Data for training
 * @return vector - trained data
 */
std::vector <probabilityPair> NaiveBayes::Train(std::vector < std::vector<std::string> > Data) {

   mainMap m;

  // std::vector <DataPair> DataCounts;
   std::vector <probabilityPair> DataProbabilities;

   int totalCounts = 0;
   int count = 0;

    //counting counts based on instances
    for (size_t j = 0; j < Data[0].size()-1; ++j) {
    for(std::map<std::string, double>::iterator it = cls.begin(); it != cls.end(); ++it) {
        for (size_t  i = 0; i < Data.size(); i++) {
            if (Data[i][Data[0].size()-1] == it->first) { // if instance is in class
                count = m[it->first][Data[i][j]]++; // counting counts of instances
                m.insert(std::pair<std::string, innerMap> (it->first, innerMap()));
                m[it->first].insert(std::pair<std::string, int> (Data[i][j], count));
            }  else m[it->first].insert(std::pair<std::string, int> (Data[i][j], 0));
        }
     }

   // DataCounts.push_back(make_pair(pompomTest[j], -1)); // adding first row for printing
   // DataProbabilities.push_back(make_pair(pompomTest[j], -1));
     for(std::map<std::string, innerMap>::iterator it = m.begin(); it != m.end(); ++it) {
       //  DataCounts.push_back(std::make_pair(/*"  "+*/it->first, 0)); // adding classes for printing
         DataProbabilities.push_back(std::make_pair(/*"  "+*/it->first, 0));
        for(std::map<std::string, int>::iterator inner_it = it->second.begin(); inner_it != it->second.end(); inner_it++) {
           totalCounts = totalCounts + (inner_it->second+1); // counting total counts in class and atribute
        }
       for(std::map<std::string, int>::iterator inner_it = it->second.begin(); inner_it != it->second.end(); inner_it++) {
        // DataCounts.push_back(std::make_pair(/*"    "+*/inner_it->first, inner_it->second+1)); // adding counts of instances
         DataProbabilities.push_back(std::make_pair(/*"    "+*/inner_it->first, (double)(inner_it->second+1)/(double)totalCounts)); //counting probabilities of instances
         }
     //   DataCounts.push_back(std::make_pair(/*"     [Total]"*/"[Total]", totalCounts)); // adding total counts for printing
      //  DataProbabilities.push_back(make_pair(/*"     [Total]"*/"[Total]", (float)totalCounts/(float)Data.size()));
        totalCounts = 0;
     }
     m.clear();
    }

  /* std::cout << "Counts" << std::endl;
   for (size_t j = 0; j < DataProbabilities.size(); ++j)
        {
            std::cout << DataProbabilities[j].first << " -> " << DataProbabilities[j].second << std::endl;
        }
*/
  // cout << endl;
  // cout << "Probabilities" << endl;

   return DataProbabilities;
}

/**
 * @brief The Classify function classifies and tests data
 * @param test - data for testing
 * @param dataProbabilities - trained data
 * @return vector - predicted classes
 */
std::vector <std::string> NaiveBayes::Classify(std::vector < std::vector<std::string> > test, std::vector <probabilityPair> DataProbabilities) {

    std::vector <double> pomCounts;
    std::vector <probabilityPair> rowProbabilities;
    std::vector <probabilityPair> totalProbabilities;
    std::vector <std::string> pomClass;
    std::vector <std::string> pomTest;

    int pom = 0;
    double percentage = 0;
    double totalPercentage = 1;

    // erasing class from testing data
    for (size_t j = 0; j < test.size(); j++) {
        test[j].erase(test[j].begin() + test[0].size());
    }

    // counting probabilities for testing instances
    for(std::map<std::string, double>::iterator it = cls.begin(); it != cls.end(); ++it) {
     for (size_t k = 0; k < test.size(); k++) {
       for (size_t j = 0; j < DataProbabilities.size(); j++) {
          for (size_t i = 0; i < test[0].size(); i++) {
              if (DataProbabilities[j].first == test[k][i]) {
                  percentage = DataProbabilities[j].second;
                  pomCounts.push_back(percentage);
              }
         }

         for(std::map<std::string, double>::iterator it = cls.begin(); it != cls.end(); ++it) {
             if (DataProbabilities[j].first == it->first) {
                pomClass.push_back(it->first);
             }
         }

        }
      }
    }

    for (size_t k = 0; k < pomClass.size(); k++) {
       totalProbabilities.push_back(std::make_pair(pomClass[k], pomCounts[k]));
    }
    // counting total probabilities for instances - multiply single probabilities in row
    for (size_t i = 1; i < test.size()+1; i++) {
     for (size_t k = 1; k < cls.size()+1; k++) {
       for (size_t j = 0; j < test[0].size(); j++) {
         totalPercentage = totalPercentage * totalProbabilities[pom].second;
         pom = pom + cls.size();
      }

      pom = 0;
      pom = pom + k;

      rowProbabilities.push_back(std::make_pair(totalProbabilities[k-1].first, totalPercentage));
      totalPercentage = 1;
     }

     pom = 0;
     totalProbabilities.erase(totalProbabilities.begin(), totalProbabilities.begin() + (cls.size()*test[0].size()));
    }

    // multiply total probabilities with probabilty of class
    for (size_t i = 0; i < rowProbabilities.size(); ++i) {
          for(std::map<std::string, double>::iterator it = cls.begin(); it != cls.end(); ++it) {
              if (rowProbabilities[i].first == it->first) {
                  rowProbabilities[i].second = rowProbabilities[i].second * it->second;
              }
          }
    }

   totalProbabilities.clear();

    // normalize counted probabilities
    percentage = 0;
    totalPercentage = 0;
    for (size_t i = 0; i < test.size(); i++) {
      for (size_t j = 0; j < cls.size(); j++) {
        percentage = percentage + rowProbabilities[j].second;
     }

     for (size_t j = 0; j < cls.size(); j++) {
       totalPercentage = rowProbabilities[j].second / percentage;
       totalProbabilities.push_back(std::make_pair(rowProbabilities[j].first, totalPercentage));
     }

     rowProbabilities.erase(rowProbabilities.begin(), rowProbabilities.begin() + cls.size());
     percentage = 0;
     totalPercentage = 0;
    }

    // adding predicted class to vector
    rowProbabilities.clear();
    for (size_t i = 0; i < test.size(); ++i) {
     for (size_t j = 0; j < cls.size(); j++) {
        rowProbabilities.push_back(std::make_pair(totalProbabilities[j].first, totalProbabilities[j].second));
     }
     pomTest.push_back(max(rowProbabilities));
     rowProbabilities.clear();
     totalProbabilities.erase(totalProbabilities.begin(), totalProbabilities.begin() + cls.size());
    }

    return pomTest;
}
