#include "cross_validation.h"

#include <map>
#include <set>
#include <algorithm>
#include "common.h"
#include "string"
#include <sstream>


const string StratifiedKFoldParameters::xml_root_name = "StratifiedKFoldParameters";
const string StratifiedKFoldParameters::xml_n_folds_name = "n_folds";
const string StratifiedKFoldParameters::xml_shuffle_name = "shuffle";
const string StratifiedKFoldParameters::xml_seed_name = "seed";


TCXMLnode* 
StratifiedKFoldParameters::to_node(void) const {
	TCXMLnode* rootnode = CxmlCreateNode(xml_root_name.c_str(), ntElement, nullptr);
	CxmlAddChild(rootnode, 
				 CxmlCreateNode(xml_n_folds_name.c_str(), ntAttribute, std::to_string(n_folds).c_str()));
	CxmlAddChild(rootnode,
				 CxmlCreateNode(xml_shuffle_name.c_str(), ntAttribute, common::bool_to_str(shuffle).c_str()));
	CxmlAddChild(rootnode,
				 CxmlCreateNode(xml_seed_name.c_str(), ntAttribute, std::to_string(seed).c_str()));

	return rootnode;
}

StratifiedKFoldParameters& 
StratifiedKFoldParameters::from_node(TCXMLnode* node) {
	n_folds = std::stoi(CxmlGetAttribute(node, &(string(xml_n_folds_name))[0]));
	shuffle = common::str_to_bool(CxmlGetAttribute(node, &(string(xml_shuffle_name))[0]));
	seed = std::stoi(CxmlGetAttribute(node, &(string(xml_seed_name))[0]));
	return *this;
}

vector<size_t>
cmp_fold_sizes(unsigned n_folds, size_t count) {
    vector<size_t> folds_sizes(n_folds);
    for(size_t i = 0; i < n_folds; i++)
        folds_sizes[i] = count / n_folds;
    for(size_t i = 0; i < count % n_folds; i++)
        folds_sizes[i]++;
    return folds_sizes;
}

//key - label, value - list of indices in y
map<label_t, vector<size_t> >
cmp_ind_per_label(const vector<label_t>& y) {
    map<label_t, vector<size_t> > ind_per_label;
    for(size_t i = 0; i < y.size(); i++) {
        map<label_t, vector<size_t> >::iterator item = ind_per_label.find(y[i]);
        if(item == ind_per_label.end()) {
            vector<size_t> vec; vec.push_back(i);
            ind_per_label.insert(std::make_pair(y[i], vec));
        }
        else
            item->second.push_back(i);
    }
    return ind_per_label;
}

void
StratifiedKFold::cmp_folds(const Matrix<data_t>& X, const vector<label_t>& y, const InteropNegotiator& negotiator) {
	_X = X;
	_y = y;
    if(_params.n_folds < 2 || _params.n_folds > y.size())
        throw std::invalid_argument("Number of folds should be in [2; y.Length]");

    size_t n_folds = _params.n_folds;
    bool shuffle = _params.shuffle;

    if(negotiator.message)
        negotiator.message(negotiator.control_frame_handle, "Computing folds...");

    map<label_t, vector<size_t> > ind_per_label = cmp_ind_per_label(y);

    //shuffle indices of each label
    if(shuffle)
        for(map<label_t, vector<size_t> >::iterator it = ind_per_label.begin(); it != ind_per_label.end(); it++)
            std::random_shuffle(it->second.begin(), it->second.end());

    vector<set<size_t> > test_folds_ind; //test folds indices
    for(size_t i = 0; i < n_folds; i++) {
        set<size_t> test_fold_ind;
        for(map<label_t, vector<size_t> >::iterator it = ind_per_label.begin(); it != ind_per_label.end(); it++) {
            size_t label_count = it->second.size();
            if(label_count < n_folds) {
                if(negotiator.message)
                    negotiator.message(negotiator.control_frame_handle, "StratifiedKFold Warning: there are too few number of samples for this number of folds");
            }
            vector<size_t> folds_sizes = cmp_fold_sizes(n_folds, label_count);
            vector<size_t>::iterator idx_start = it->second.begin() + i    * folds_sizes[i];
            vector<size_t>::iterator idx_stop = it->second.begin() + (i + 1) * folds_sizes[i];
            for(vector<size_t>::iterator idx_it = idx_start; idx_it != idx_stop; idx_it++)
                test_fold_ind.insert(*idx_it);

        }
        test_folds_ind.push_back(test_fold_ind);
		_test_folds_inds = test_folds_ind;
    }
}


vector<set<size_t> >  
StratifiedKFold::train_folds_ind() {
	vector<set<size_t>> res;
	for (size_t idx = 0; idx < this->n_folds(); idx++) {
		//cmp train fold indices using test fold indices (set difference)
		set<size_t> train_fold_inds;
		for (size_t i = 0; i < _y.size(); i++) {
			if (_test_folds_inds[idx].find(i) == _test_folds_inds[idx].end())
				_train_fold_inds.insert(i);
		}
		res.push_back(train_fold_inds);
	}
	return res;
}

StratifiedKFold::StratifiedKFold(const Matrix<data_t>& X, const vector<label_t>& y, const StratifiedKFoldParameters& params) 
	: _params(params),  _X(X), _y(y) {
    if(_params.n_folds < 2 || _params.n_folds > y.size())
        throw std::invalid_argument("Number of folds should be in [2; y.Length]");
    std::srand(_params.seed);
    InteropNegotiator negotiator = {nullptr, nullptr, nullptr, nullptr};
    cmp_folds(X, y, negotiator);
}

StratifiedKFold::StratifiedKFold(const StratifiedKFoldParameters& params) : _params(params){ 
    std::srand(_params.seed);
}

Fold
StratifiedKFold::_ind_to_fold(const Matrix<data_t>& X, const vector<label_t>& y, size_t i) {
	_X = X;
	_y = y;
    Fold fold;
    vector<label_t> test_targets(_test_folds_inds[i].size());
    vector<label_t> train_targets(_train_fold_inds.size());
    vector<data_t> test_data(_test_folds_inds[i].size() * X.n_cols());
    vector<data_t> train_data(_train_fold_inds.size() * X.n_cols());

    size_t k = 0;
    for(set<size_t>::iterator it = _test_folds_inds[i].begin(); it != _test_folds_inds[i].end(); it++) {
        test_targets[k] = y[*it];
		for (size_t j = 0; j < X.n_cols(); j++) {
			test_data[k * X.n_cols() + j] = X(*it, j);
		}
        k++;
    }
    fold.test_targets = test_targets;
    fold.test_data = Matrix<data_t>(test_data, fold.test_targets.size(), X.n_cols());

    k = 0;
    for(set<size_t>::iterator it = _train_fold_inds.begin(); it != _train_fold_inds.end(); it++) {
        train_targets[k] = y[*it];
		for (size_t j = 0; j < X.n_cols(); j++) {
			train_data[k * X.n_cols() + j] = X(*it, j);
		}
        k++;
    }
    fold.train_targets = train_targets;
    fold.train_data = Matrix<data_t>(train_data, fold.train_targets.size(), X.n_cols());
    return fold;
}

Fold 
StratifiedKFold::get_fold(size_t idx) {
	//cmp train fold indices using test fold indices (set difference)
	_train_fold_inds.clear();
	for (size_t i = 0; i < _y.size(); i++){
		if (_test_folds_inds[idx].find(i) == _test_folds_inds[idx].end())
			_train_fold_inds.insert(i);
	}
	return _ind_to_fold(_X, _y, idx);
}

size_t 
StratifiedKFold::n_folds() {
	return _params.n_folds;
}

CrossValDetails 
CrossValCancellable::cross_val(ClassifierABC* clf, CrossValidationABC* cv, const InteropNegotiator& negotiator,  bool predict_train) {
	CrossValDetails cv_details;
	// todo: reserve place in result
	for (size_t i = 0; i < cv->n_folds(); i++) {
        if(cancellation_pending) {
			return CrossValDetails();
        }
		if (negotiator.report_progress)
			negotiator.report_progress(negotiator.control_frame_handle,(i + 1) / static_cast<float>(cv->n_folds()));

		Fold f = cv->get_fold(i);
		cv_details.train_y_true.push_back(f.train_targets);
		cv_details.test_y_true.push_back(f.test_targets);
	
		clf->fit(f.train_data, f.train_targets);

		if(predict_train) 
			cv_details.train_y_pred.push_back(clf->predict(f.train_data));
		cv_details.test_y_pred.push_back(clf->predict(f.test_data));
	}
	return cv_details;
}

std::pair<vector<label_t>, vector<label_t> >
CrossValCancellable::cross_val_predict(ClassifierABC* clf, CrossValidationABC* cv,  const InteropNegotiator& negotiator, bool report_progress){
	vector<label_t> y_pred;
	vector<label_t> y_true;
	// todo: reserve place in result
	for (size_t i = 0; i < cv->n_folds(); i++) {
        if(cancellation_pending) {
			return std::pair<vector<label_t>, vector<label_t>>();
        }
		if (report_progress)
			if (negotiator.report_progress)
				negotiator.report_progress(negotiator.control_frame_handle,(i + 1) / static_cast<float>(cv->n_folds()));

		Fold f = cv->get_fold(i);
		y_true.insert(y_true.end(), f.test_targets.begin(), f.test_targets.end());
		clf->fit(f.train_data, f.train_targets);
		vector<label_t> fold_y_pred = clf->predict(f.test_data);
		y_pred.insert(y_pred.end(), fold_y_pred.begin(), fold_y_pred.end());
	}
	return std::pair<vector<label_t>, vector<label_t>>(y_true, y_pred);
}

score_t
CrossValCancellable::cross_val_score(ClassifierABC* clf, CrossValidationABC* cv, multiclass_metric_fptr scoring_func, const InteropNegotiator& negotiator, bool report_progress) {
    auto res = cross_val_predict(clf, cv, negotiator, report_progress);
    return scoring_func(res.first, res.second);
}

CrossValDetails cross_val(ClassifierABC* clf, CrossValidationABC* cv, bool predict_train) {
    InteropNegotiator negotiator = {nullptr, nullptr, nullptr, nullptr};
    return CrossValCancellable().cross_val(clf, cv, negotiator, predict_train);
}

std::pair<vector<label_t>, vector<label_t> > cross_val_predict(ClassifierABC* clf, CrossValidationABC* cv) {
    InteropNegotiator negotiator = {nullptr, nullptr, nullptr, nullptr};
    return CrossValCancellable().cross_val_predict(clf, cv, negotiator, false);
}

score_t cross_val_score(ClassifierABC* clf, CrossValidationABC* cv, multiclass_metric_fptr scoring_func) {
    InteropNegotiator negotiator = {nullptr, nullptr, nullptr, nullptr};
    return CrossValCancellable().cross_val_score(clf, cv, scoring_func, negotiator, false);
}
