#include <cmath>
#include "data.hpp"
#include "ViolaJones_device.hpp"

constexpr static inline void add_empty_feature(const np::Array<uint8_t>& feats, size_t& n) noexcept {
	memset(&feats[n], 0, 4 * sizeof(uint8_t));
	n += 4;
}

constexpr static inline void add_right_feature(const np::Array<uint8_t>& feats, size_t& n, const uint16_t& i, const uint16_t& j, const uint16_t& w, const uint16_t& h) noexcept {
	feats[n++] = i + w;
	feats[n++] = j;
	feats[n++] = w;
	feats[n++] = h;
}

constexpr static inline void add_immediate_feature(const np::Array<uint8_t>& feats, size_t& n, const uint16_t& i, const uint16_t& j, const uint16_t& w, const uint16_t& h) noexcept {
	feats[n++] = i;
	feats[n++] = j;
	feats[n++] = w;
	feats[n++] = h;
}

constexpr static inline void add_bottom_feature(const np::Array<uint8_t>& feats, size_t& n, const uint16_t& i, const uint16_t& j, const uint16_t& w, const uint16_t& h) noexcept {
	feats[n++] = i;
	feats[n++] = j + h;
	feats[n++] = w;
	feats[n++] = h;
}

constexpr static inline void add_right2_feature(const np::Array<uint8_t>& feats, size_t& n, const uint16_t& i, const uint16_t& j, const uint16_t& w, const uint16_t& h) noexcept {
	feats[n++] = i + 2 * w;
	feats[n++] = j;
	feats[n++] = w;
	feats[n++] = h;
}

constexpr static inline void add_bottom2_feature(const np::Array<uint8_t>& feats, size_t& n, const uint16_t& i, const uint16_t& j, const uint16_t& w, const uint16_t& h) noexcept {
	feats[n++] = i;
	feats[n++] = j + 2 * h;
	feats[n++] = w;
	feats[n++] = h;
}

constexpr static inline void add_bottom_right_feature(const np::Array<uint8_t>& feats, size_t& n, const uint16_t& i, const uint16_t& j, const uint16_t& w, const uint16_t& h) noexcept {
	feats[n++] = i + w;
	feats[n++] = j + h;
	feats[n++] = w;
	feats[n++] = h;
}

/**
 * @brief Initialize the features based on the input shape.
 *
 * @param width Width of the image
 * @param height Height of the image
 * @return The initialized features
 */
np::Array<uint8_t> build_features(const uint16_t& width, const uint16_t& height) noexcept {
	size_t n = 0;
	uint16_t w, h, i, j;
	for (w = 1; w < width; ++w)
		for (h = 1; h < height; ++h)
			for (i = 0; i < width - w; ++i)
				for (j = 0; j < height - h; ++j) {
					if (i + 2 * w < width)	++n;
					if (j + 2 * h < height)	++n;
					if (i + 3 * w < width)	++n;
					if (j + 3 * h < height)	++n;
					if (i + 2 * w < width && j + 2 * h < height) ++n;
				}

	np::Array<uint8_t> feats = np::empty<uint8_t>({ n, 2, 2, 4 });

	n = 0;
	for (w = 1; w < width; ++w)
		for (h = 1; h < height; ++h)
			for (i = 0; i < width - w; ++i)
				for (j = 0; j < height - h; ++j) {
					if (i + 2 * w < width) {
						add_right_feature(feats, n, i, j, w, h);
						add_empty_feature(feats, n);
						add_immediate_feature(feats, n, i, j, w, h);
						add_empty_feature(feats, n);
					}

					if (j + 2 * h < height) {
						add_immediate_feature(feats, n, i, j, w, h);
						add_empty_feature(feats, n);
						add_bottom_feature(feats, n, i, j, w, h);
						add_empty_feature(feats, n);
					}

					if (i + 3 * w < width) {
						add_right_feature(feats, n, i, j, w, h);
						add_empty_feature(feats, n);
						add_right2_feature(feats, n, i, j, w, h);
						add_immediate_feature(feats, n, i, j, w, h);
					}

					if (j + 3 * h < height) {
						add_bottom_feature(feats, n, i, j, w, h);
						add_empty_feature(feats, n);
						add_bottom2_feature(feats, n, i, j, w, h);
						add_immediate_feature(feats, n, i, j, w, h);
					}

					if (i + 2 * w < width && j + 2 * h < height) {
						add_right_feature(feats, n, i, j, w, h);
						add_bottom_feature(feats, n, i, j, w, h);
						add_immediate_feature(feats, n, i, j, w, h);
						add_bottom_right_feature(feats, n, i, j, w, h);
					}
				}
	return feats;
}

//np::Array<int32_t> select_percentile(const np::Array<uint8_t> X_feat, const np::Array<uint8_t> y) noexcept {
//	std::vector<float64_t> class_0, class_1;
//
//	const int32_t im_size = X_feat.shape[0] / y.shape[0];
//	int32_t idy = 0, n_samples_per_class_0 = 0, n_samples_per_class_1 = 0;
//	for (size_t i = 0; i < X_feat.shape[0]; i += im_size) {
//		if (y[idy] == 0) {
//			++n_samples_per_class_0;
//			class_0.push_back(static_cast<float64_t>(X_feat[i]));
//		}
//		else {
//			++n_samples_per_class_1;
//			class_1.push_back(static_cast<float64_t>(X_feat[i]));
//		}
//		++idy;
//	}
//	const int32_t n_samples = n_samples_per_class_0 + n_samples_per_class_1;
//
//	float64_t ss_alldata_0 = 0;
//	for (int32_t i = 0;i < n_samples_per_class_0;++i)
//		ss_alldata_0 += (class_0[i] * class_0[i]);
//
//	float64_t ss_alldata_1 = 0;
//	for (int32_t i = 0;i < n_samples_per_class_1;++i)
//		ss_alldata_1 += (class_1[i] * class_1[i]);
//
//	const float64_t ss_alldata = ss_alldata_0 + ss_alldata_1;
//
//	float64_t sums_classes_0 = 0;
//	for (int32_t i = 0;i < n_samples_per_class_0;++i)
//		sums_classes_0 += class_0[i];
//
//	float64_t sums_classes_1 = 0;
//	for (int32_t i = 0;i < n_samples_per_class_1;++i)
//		sums_classes_1 += class_1[i];
//
//	float64_t sq_of_sums_alldata = sums_classes_0 + sums_classes_1;
//	sq_of_sums_alldata *= sq_of_sums_alldata;
//
//	const float64_t sq_of_sums_args_0 = sums_classes_0 * sums_classes_0;
//	const float64_t sq_of_sums_args_1 = sums_classes_1 * sums_classes_1;
//	const float64_t ss_tot = ss_alldata - sq_of_sums_alldata / n_samples;
//	const float64_t sqd_sum_bw_n = sq_of_sums_args_0 / n_samples_per_class_0 + sq_of_sums_args_1 / n_samples_per_class_1 - sq_of_sums_alldata / n_samples;
//	const float64_t ss_wn = ss_tot - sqd_sum_bw_n;
//	const int32_t df_wn = n_samples - 2;
//	const float64_t msw = ss_wn / df_wn;
//	const float64_t f_values = sqd_sum_bw_n / msw;
//
//	const np::Array<int32_t> res = np::empty<int32_t>({ static_cast<size_t>(std::ceil(static_cast<float64_t>(im_size) / 10.0)) });
//	// TODO Complete code
//	return res;
//}

/**
 * @brief Initialize the weights of the weak classifiers based on the training labels.
 *
 * @param y_train Training labels
 * @return The initialized weights
 */
np::Array<float64_t> init_weights(const np::Array<uint8_t>& y_train) noexcept {
	np::Array<float64_t> weights = np::empty<float64_t>(y_train.shape);
	const uint16_t t = np::sum(np::astype<uint16_t>(y_train));
	return map(weights, static_cast<std::function<float64_t(const size_t&, const float64_t&)>>(
			[&t, &y_train](const size_t& i, const float64_t&) -> float64_t {
		return 1.0 / (2 * (y_train[i] == 0 ? t : y_train.shape[0] - t));
	}));
}

/**
 * @brief Classify the integrated features based on polarity and threshold.
 *
 * @param X_feat_i Integrated features
 * @param j Index of the classifier
 * @param threshold Trained threshold
 * @param polarity Trained polarity
 * @return Classified features
 */
static np::Array<uint8_t> classify_weak_clf(const np::Array<int32_t>& X_feat_i, const size_t& j, const float64_t& threshold, const float64_t& polarity) noexcept {
	np::Array<uint8_t> res = np::empty<uint8_t>({ X_feat_i.shape[1] });
	for(size_t i = 0; i < res.shape[0]; ++i)
		res[i] = polarity * X_feat_i[j * X_feat_i.shape[1] + i] < polarity * threshold ? 1 : 0;
	return res;
}

/**
 * @brief Classify the trained classifiers on the given features.
 *
 * @param alphas Trained alphas
 * @param classifiers Trained classifiers
 * @param X_feat integrated features
 * @return Classification results
 */
np::Array<uint8_t> classify_viola_jones(const np::Array<float64_t>& alphas, const np::Array<float64_t>& classifiers, const np::Array<int32_t>& X_feat) noexcept {
	np::Array<float64_t> total = np::zeros<float64_t>({ X_feat.shape[1] });

	float64_t clf = 0.0, threshold = 0.0, polarity = 0.0;
	for(size_t i = 0; i < alphas.shape[0]; ++i){
		clf = classifiers[i * 3]; threshold = classifiers[i * 3 + 1]; polarity = classifiers[i * 3 + 2];
		//total += alphas * np::astype<float64_t>(classify_weak_clf(X_feat, clf, threshold, polarity));
		const np::Array<uint8_t> res = classify_weak_clf(X_feat, clf, threshold, polarity);
		for(size_t j = 0; j < res.shape[0]; ++j)
			total[j] += alphas[i] * res[j];
	}

	np::Array<uint8_t> y_pred = np::empty<uint8_t>({ X_feat.shape[1] });
	const float64_t alphas_sum = np::sum(alphas);
	for(size_t i = 0; i < X_feat.shape[1]; ++i)
		y_pred[i] = total[i] >= 0.5 * alphas_sum ? 1 : 0;

	return y_pred;
}

/**
 * @brief Select the best classifer given their predictions.
 *
 * @param classifiers The weak classifiers
 * @param weights Trained weights of each classifiers
 * @param X_feat Integrated features
 * @param y Features labels
 * @return Index of the best classifier, the best error and the best accuracy
 */
std::tuple<int32_t, float64_t, np::Array<float64_t>> select_best(const np::Array<float64_t>& classifiers, const np::Array<float64_t>& weights, const np::Array<int32_t>& X_feat, const np::Array<uint8_t>& y) noexcept {
	std::tuple<int32_t, float64_t, np::Array<float64_t>> res = { -1, np::inf, np::empty<float64_t>({ X_feat.shape[0] }) };

	for(size_t j = 0; j < classifiers.shape[0]; ++j){
		const np::Array<float64_t> accuracy = np::abs(np::astype<float64_t>(classify_weak_clf(X_feat, j, classifiers[j * 2], classifiers[j * 2 + 1])) - y);
		// const float64_t error = np::mean(weights * accuracy);
		float64_t error = 0.0;
		for(size_t i = 0; i < weights.shape[0]; ++i)
			error += weights[i] * accuracy[i];
		error /= weights.shape[0];

		if (error < std::get<1>(res))
			res = { j, error, accuracy };
	}
	return res;
}

/**
 * @brief Train the weak calssifiers.
 *
 * @param T Number of weak classifiers
 * @param X_feat Integrated features
 * @param X_feat_argsort Sorted indexes of the integrated features
 * @param y Features labels
 * @return List of trained alphas and the list of the final classifiers
 */
std::array<np::Array<float64_t>, 2> train_viola_jones(const size_t& T, const np::Array<int32_t>& X_feat, const np::Array<uint16_t>& X_feat_argsort, const np::Array<uint8_t>& y) noexcept {
	np::Array<float64_t> weights = init_weights(y);
	np::Array<float64_t> alphas = np::empty<float64_t>({ T });
	np::Array<float64_t> final_classifier = np::empty<float64_t>({ T, 3 });

	for(size_t t = 0; t < T; ++t ){
		weights /= np::sum(weights);
		const np::Array<float64_t> classifiers = train_weak_clf(X_feat, X_feat_argsort, y, weights);
		const auto [ clf, error, accuracy ] = select_best(classifiers, weights, X_feat, y);
		float64_t beta = error / (1.0 - error);
		weights *= np::pow(beta, (1.0 - accuracy));
		alphas[t] = std::log(1.0 / beta);
		final_classifier[t * 3] = clf;
		final_classifier[t * 3 + 1] = classifiers[clf * 2];
		final_classifier[t * 3 + 2] = classifiers[clf * 2 + 1];
	}
	return { alphas, final_classifier };
}

/**
 * @brief Compute the accuracy score i.e. how a given set of measurements are close to their true value.
 *
 * @param y Ground truth labels
 * @param y_pred Predicted labels
 * @return computed accuracy score
 */
float64_t accuracy_score(const np::Array<uint8_t>& y, const np::Array<uint8_t>& y_pred) noexcept {
	float64_t res = 0.0;
	for(size_t i = 0; i < y.shape[0]; ++i)
		if(y[i] == y_pred[i])
			++res;
	return res / y.shape[0];
}

/**
 * @brief Compute the precision score i.e. how a given set of measurements are close to each other.
 *
 * @param y Ground truth labels
 * @param y_pred Predicted labels
 * @return computed precision score
 */
float64_t precision_score(const np::Array<uint8_t>& y, const np::Array<uint8_t>& y_pred) noexcept {
	uint16_t true_positive = 0, false_positive = 0;
	for(size_t i = 0; i < y.shape[0]; ++i)
		if(y[i] == 1){
			if(y[i] == y_pred[i])
				++true_positive;
			else
				++false_positive;
		}
	return static_cast<float64_t>(true_positive) / (true_positive + false_positive);
}

/**
 * @brief Compute the recall score i.e. the ratio (TP / (TP + FN)) where TP is the number of true positives and FN the number of false negatives.
 *
 * @param y Ground truth labels
 * @param y_pred Predicted labels
 * @return computed recall score
 */
float64_t recall_score(const np::Array<uint8_t>& y, const np::Array<uint8_t>& y_pred) noexcept {
	uint16_t true_positive = 0, false_negative = 0;
	for(size_t i = 0; i < y.shape[0]; ++i)
		if(y[i] == 0) {
			if(y[i] != y_pred[i])
				++false_negative;
		} else {
			if(y[i] == y_pred[i])
				++true_positive;
		}
	return static_cast<float64_t>(true_positive) / (true_positive + false_negative);
}

/**
 * @brief Compute the F1 score aka balanced F-score or F-measure.
 *
 * F1 = (2 * TP) / (2 * TP + FP + FN)
 * where TP is the true positives,
 * FP is the false positives,
 * and FN is the false negatives
 *
 * @param y Ground truth labels
 * @param y_pred Predicted labels
 * @return computed F1 score
 */
float64_t f1_score(const np::Array<uint8_t>& y, const np::Array<uint8_t>& y_pred) noexcept {
	const float64_t precision = precision_score(y, y_pred);
	const float64_t recall = recall_score(y, y_pred);
	return 2 * (precision * recall) / (precision + recall);
}

/**
 * @brief Compute the confusion matrix to evaluate a given classification.
 *
 * A confusion matrix of a binary classification consists of a 2x2 matrix containing
 * | True negatives  | False positives |
 * | False negatives | True positives  |
 *
 * @param y Ground truth labels
 * @param y_pred Predicted labels
 * @return computed confusion matrix
 */
std::tuple<uint16_t, uint16_t, uint16_t, uint16_t> confusion_matrix(const np::Array<uint8_t>& y, const np::Array<uint8_t>& y_pred) noexcept {
	uint16_t true_positive = 0, false_positive = 0, true_negative = 0, false_negative = 0;
	for(size_t i = 0; i < y.shape[0]; ++i)
		if(y[i] == 0)
			if(y[i] == y_pred[i])
				++true_negative;
			else
				++false_negative;
		else
			if(y[i] == y_pred[i])
				++true_positive;
			else
				++false_positive;
	return std::make_tuple(true_negative, false_positive, false_negative, true_positive);
}