#include "decisiontree.h"

DecisionTree::DecisionTree()
{
}

/**
 * @brief The countInstances function counting number of instances in a parametr
 * @param vec - data for counting
 * @return integer - number of instances
 */
int DecisionTree::countInstances(std::vector < std::vector<std::string> > vec) {
    return vec.size();
}

/**
 * @brief The divideData function dividing data to vector, where are only class and given atribute
 * @param vec - data for dividing
 * @param Atr - which atribute is divided based on actual name
 * @return vector - diveded data
 */
std::vector < std::vector<std::string> > DecisionTree::divideData(std::vector < std::vector<std::string> > vec, std::string Atr){

    std::vector<std::string> pomToDivide;
    std::vector < std::vector<std::string> > dividedData;

     for (size_t i = 0; i < vec.size(); i++) {
       if (Atr == vec[i][0]){ //finds out if which atribut to divide
         for (size_t j = 0; j < vec[i].size(); j++) {
           pomToDivide.push_back(vec[i][j]); //adding atribute to 1d help vector
          }
         dividedData.push_back(pomToDivide); //adding atributes to 2d vector
         pomToDivide.clear();
        }
    }
    return dividedData;
}

/**
 * @brief The returnDividedDatafunction dividing data to vector, where are only class and given atribute
 * @param vec - data for dividing
 * @param column - which atribute is divided based on index
 * @return vector - diveded data
 */
std::vector < std::vector<std::string> > DecisionTree::returnDividedData(std::vector < std::vector<std::string> > vec, int column){

    std::vector<std::string> pomToDivide;
    std::vector < std::vector<std::string> > DataToDivide;

     for (size_t i = 0; i < vec.size(); i++) {
          pomToDivide.push_back(vec[i][column]); //adding atribute to vector based on index of atribute
          pomToDivide.push_back(vec[i][vec[i].size()-1]); //adding class to vector
         DataToDivide.push_back(pomToDivide); // adding both atributes to one vector
         pomToDivide.clear();
    }

    std::sort(DataToDivide.begin(), DataToDivide.end()); //sorting divided data

    return DataToDivide;
}

/**
 * @brief The countOccurances function counting occurences in data
 * @param DataToDivide - data for counting
 * @return map - counted data
 */
std::map <std::string, double> DecisionTree::countOccurances(std::vector < std::vector<std::string> > DataToDivide){

    std::map <std::string, double> firstColumn;

    for (size_t i = 0; i < DataToDivide.size(); i++) {
         firstColumn.insert(std::pair<std::string, double>(DataToDivide[i][0], firstColumn[DataToDivide[i][0]]++));
    }
    return firstColumn;
}


/**
 * @brief The countInfoGain function counting information gain in given column
 * @param vec - data for counting
 * @param column - index of column to count
 * @param cls - map which holds class counts
 * @return double - information gain
 */
double DecisionTree::countInfoGain(std::vector < std::vector<std::string> > vec, int column, std::map <std::string, double> cls) {

    std::map <std::string, double> firstColumn;

    std::vector < std::vector<std::string> > dividedData;
    std::vector < std::vector<std::string> > DataToDivide;

    double entropy, infoGain, totalInfoGain;
    double totalEntropy = countClassEntropy(cls);

    DataToDivide = returnDividedData(vec, column);

    firstColumn = countOccurances(DataToDivide);

    for(std::map<std::string, double>::iterator it = firstColumn.begin(); it != firstColumn.end(); ++it) {
     for (size_t i = 0; i < DataToDivide.size(); i++) {
         if (DataToDivide[i][0] == it->first) {
            dividedData = divideData(DataToDivide, it->first);
         }
       }
    entropy = countEntropy(countProbabilities(dividedData, cls)); // counting probabilities and entropy for diveded data
    infoGain = infoGain + (it->second/countInstances(vec)*entropy); //counting information gain for diveded data

    }

    totalInfoGain = totalEntropy - infoGain; // counting total information gain

    return totalInfoGain;
}

/**
 * @brief The countEntropy function counting entropy in given vector
 * @param vec - data for counting
 * @return double - entropy
 */
double DecisionTree::countEntropy(std::vector <probabilityPair> vec) {

    double entropy = 0;

    for (size_t i = 0; i < vec.size(); ++i) {
      entropy = entropy - (vec[i].second * log2f(vec[i].second)); //counting entropy based on logaritmic function
    }
   return entropy;
}

/**
 * @brief The countProbabilities function counting probabilities in given vector
 * @param vec - data for counting
 * @param cls - map which holds class counts
 * @return vector - holds atribute and probability of that atribute
 */
std::vector <probabilityPair> DecisionTree::countProbabilities(std::vector < std::vector<std::string> > vec, std::map <std::string, double> cls) {

    std::vector <probabilityPair> DataProbabilities;

    std::map <std::string, double> countMap;

    int totalCounts = 0;
    int count = 0;

    for(std::map<std::string, double>::iterator it = cls.begin(); it != cls.end(); ++it) {
        for (size_t i = 0; i < vec.size(); i++) {
           if (vec[i][1] == it->first) {
              count = countMap[vec[i][1]]++; // counting ocurrences in vector
              countMap.insert(std::pair<std::string, double> (vec[i][1], count)); //creating map with name of atribute and number od occurences in vector
              totalCounts = vec.size();
           }
        }
    }

    for(std::map<std::string, double>::iterator it = countMap.begin(); it != countMap.end(); ++it) {
        DataProbabilities.push_back(std::make_pair(it->first, it->second/totalCounts)); //counting probabilites for vector
    }
    return DataProbabilities;
}

/**
 * @brief The countClassEntropy function counting entropy for whole data set
 * @param cls - map which holds class counts
 * @return double - class entropy
 */
double DecisionTree::countClassEntropy(std::map <std::string, double> cls) {

    double entropy = 0;

    for(std::map<std::string, double>::iterator it = cls.begin(); it != cls.end(); ++it) {
       entropy = entropy - (it->second * log2f(it->second));
    }
   return entropy;
}

/**
 * @brief The divideWhole function dividing whole data to vector, where are only class and given atribute
 * @param vec - data for dividing
 * @param Atr - which atribute is divided based on actual name
 * @return vector - diveded data
 */
std::vector < std::vector<std::string> > DecisionTree::divideWhole(std::vector < std::vector<std::string> > Data, std::string atribute, int column){

  std::vector < std::vector<std::string> > cutData;

  for (size_t j = 0; j < Data.size(); j++) {
   // for (size_t i = 0; i < Data[j].size()-1; i++) {
             if (Data[j][column] == atribute) { // if instance is same as atribute then divide
                 cutData.push_back(Data[j]);
             }
   //     }
   }
    return cutData;
}
