//
//  GaussianNB.cpp
//
//  Created by Marc Melikyan on 1/17/21.
//

#include "gaussian_nb.h"
#include "../lin_alg/lin_alg.h"
#include "../stat/stat.h"
#include "../utilities/utilities.h"

#include <algorithm>
#include <iostream>
#include <random>


GaussianNB::GaussianNB(std::vector<std::vector<double>> inputSet, std::vector<double> outputSet, int class_num) :
		inputSet(inputSet), outputSet(outputSet), class_num(class_num) {
	y_hat.resize(outputSet.size());
	Evaluate();
	LinAlg alg;
}

std::vector<double> GaussianNB::modelSetTest(std::vector<std::vector<double>> X) {
	std::vector<double> y_hat;
	for (int i = 0; i < X.size(); i++) {
		y_hat.push_back(modelTest(X[i]));
	}
	return y_hat;
}

double GaussianNB::modelTest(std::vector<double> x) {
	Stat stat;
	LinAlg alg;

	double score[class_num];
	double y_hat_i = 1;
	for (int i = class_num - 1; i >= 0; i--) {
		y_hat_i += std::log(priors[i] * (1 / sqrt(2 * M_PI * sigma[i] * sigma[i])) * exp(-(x[i] * mu[i]) * (x[i] * mu[i]) / (2 * sigma[i] * sigma[i])));
		score[i] = exp(y_hat_i);
	}
	return std::distance(score, std::max_element(score, score + sizeof(score) / sizeof(double)));
}

double GaussianNB::score() {
	Utilities util;
	return util.performance(y_hat, outputSet);
}

void GaussianNB::Evaluate() {
	Stat stat;
	LinAlg alg;

	// Computing mu_k_y and sigma_k_y
	mu.resize(class_num);
	sigma.resize(class_num);
	for (int i = class_num - 1; i >= 0; i--) {
		std::vector<double> set;
		for (int j = 0; j < inputSet.size(); j++) {
			for (int k = 0; k < inputSet[j].size(); k++) {
				if (outputSet[j] == i) {
					set.push_back(inputSet[j][k]);
				}
			}
		}
		mu[i] = stat.mean(set);
		sigma[i] = stat.standardDeviation(set);
	}

	// Priors
	priors.resize(class_num);
	for (int i = 0; i < outputSet.size(); i++) {
		priors[int(outputSet[i])]++;
	}
	priors = alg.scalarMultiply(double(1) / double(outputSet.size()), priors);

	for (int i = 0; i < outputSet.size(); i++) {
		double score[class_num];
		double y_hat_i = 1;
		for (int j = class_num - 1; j >= 0; j--) {
			for (int k = 0; k < inputSet[i].size(); k++) {
				y_hat_i += std::log(priors[j] * (1 / sqrt(2 * M_PI * sigma[j] * sigma[j])) * exp(-(inputSet[i][k] * mu[j]) * (inputSet[i][k] * mu[j]) / (2 * sigma[j] * sigma[j])));
			}
			score[j] = exp(y_hat_i);
			std::cout << score[j] << std::endl;
		}
		y_hat[i] = std::distance(score, std::max_element(score, score + sizeof(score) / sizeof(double)));
		std::cout << std::distance(score, std::max_element(score, score + sizeof(score) / sizeof(double))) << std::endl;
	}
}