#pragma once
#ifndef CROSS_VALIDATION_H
#define CROSS_VALIDATION_H


#include <vector>
#include <set>
#include <map>
#include <algorithm>
#include <random>
#include <string>
#include "cxml/cxml.h"
//#include <memory>
#include "matrix.h"
#include "abc.h"
#include "metrics.h"

using std::vector;
using std::set;
using std::map;
using std::size_t;
using std::pair;
using std::string;
using namespace metrics::classification;

//key - label, value - list of indices in y
map<label_t, vector<size_t> > cmp_ind_per_label(const vector<label_t>& y);

struct Fold {
    Matrix<data_t> train_data;
    vector<label_t> train_targets;
    Matrix<data_t> test_data;
    vector<label_t> test_targets;
};

class CrossValidationABC  {
public:
	virtual Fold get_fold(size_t idx) = 0;
	virtual size_t n_folds() = 0;
};

struct StratifiedKFoldParameters :  public XmlSerializableABC {
	const static string xml_root_name;
	const static string xml_n_folds_name;
	const static string xml_shuffle_name;
	const static string xml_seed_name;

	size_t n_folds = 5;
	bool shuffle = true;
	unsigned seed = 1999;

	TCXMLnode* to_node(void) const;
	StratifiedKFoldParameters& from_node(TCXMLnode* node);
};


class StratifiedKFold : public CrossValidationABC{
private:
    StratifiedKFoldParameters _params;
	Matrix<data_t> _X;
	vector<label_t> _y;
    vector<set<size_t> > _test_folds_inds;
	set<size_t> _train_fold_inds;

    Fold _ind_to_fold(const Matrix<data_t>& X, const vector<label_t>& y, size_t i);
    StratifiedKFold();
public:
    StratifiedKFold(const StratifiedKFoldParameters& params);
	StratifiedKFold(const Matrix<data_t>& X, const vector<label_t>& y, const StratifiedKFoldParameters& params);
	void cmp_folds(const Matrix<data_t>& X, const vector<label_t>& y, const InteropNegotiator& negotiator);

    vector<set<size_t> >  test_folds_ind() { return _test_folds_inds; }
	vector<set<size_t> >  train_folds_ind();
	Fold get_fold(size_t idx);
	size_t n_folds();
    const StratifiedKFoldParameters& parameters() const { return _params;}

};

struct CrossValDetails {
	vector<vector<label_t>> train_y_pred; //fold x fold size
	vector<vector<label_t>> train_y_true;
	vector<vector<label_t>> test_y_true;
	vector<vector<label_t>> test_y_pred;
};

class CrossValCancellable {
public:
    bool cancellation_pending = false;
    CrossValDetails cross_val(ClassifierABC* clf, CrossValidationABC* cv, const InteropNegotiator& negotiator, bool predict_train=false);
    //returns concatenated true and predicted label
    std::pair<vector<label_t>, vector<label_t> > cross_val_predict(ClassifierABC* clf, CrossValidationABC* cv,  const InteropNegotiator& negotiator, bool report_progress=true);
    score_t cross_val_score(ClassifierABC* clf, CrossValidationABC* cv, multiclass_metric_fptr scoring_func, const InteropNegotiator& negotiator, bool report_progress=true);
};

CrossValDetails cross_val(ClassifierABC* clf, CrossValidationABC* cv, bool predict_train=false);
std::pair<vector<label_t>, vector<label_t> > cross_val_predict(ClassifierABC* clf, CrossValidationABC* cv);
score_t cross_val_score(ClassifierABC* clf, CrossValidationABC* cv, multiclass_metric_fptr scoring_func);


#endif