#include "dataset.h"
#include "common.h"

#include <sstream> 


const string Sample::xml_root_name = "Sample";
const string Sample::xml_label_name = "label";
const string Sample::xml_path_name = "path";
const string Sample::xml_data_name = "data";

const string ClassLabel::xml_root_name = "ClassLabel";
const string ClassLabel::xml_name_name = "name";
const string ClassLabel::xml_code_name = "code";
const string ClassLabel::xml_enabled_name = "enabled";

const string MetaData::xml_root_name = "MetaData";
const string MetaData::xml_features_name = "features";
const string MetaData::xml_classes_name = "classes";

const string Dataset::xml_root_name = "Dataset";
const string Dataset::xml_samples_name = "samples";


TCXMLnode* 
Sample::to_node(void) const{
	TCXMLnode * rootnode = CxmlCreateNode(xml_root_name.c_str(), ntElement, NULL);
	CxmlAddChild(rootnode, CxmlCreateNode(xml_label_name.c_str(), ntAttribute, label.c_str()));
	CxmlAddChild(rootnode, CxmlCreateNode(xml_path_name.c_str(), ntAttribute, path.c_str()));
	std::string data_str = common::vec_to_str(data, ';');
	TCXMLnode * n;
	CxmlAddChild(rootnode, n = CxmlCreateNode(xml_data_name.c_str(), ntElement, NULL));
	CxmlAddChild(n, CxmlCreateNode(nullptr, ntText, data_str.c_str()));
	return rootnode;
}

Sample&
Sample::from_node(TCXMLnode* node){
	data.clear();
	label = CxmlGetAttribute(node, &string(xml_label_name)[0]);
	path = CxmlGetAttribute(node, &string(xml_path_name)[0]);
	TCXMLnode *data_node = CxmlGetSubElemnetN(node, &string(xml_data_name)[0]);
	vector<string> data_str = common::split(CxmlGetSubTextIText(data_node, 0), ';');
	data = common::str_to_vec_double(CxmlGetSubTextIText(data_node, 0), ';');
	return *this;
}


TCXMLnode* 
ClassLabel::to_node(void) const{
	TCXMLnode * rootnode = CxmlCreateNode(xml_root_name.c_str(), ntElement, NULL);
	CxmlAddChild(rootnode, CxmlCreateNode(xml_code_name.c_str(), ntAttribute, std::to_string(code).c_str()));
	CxmlAddChild(rootnode, 
		CxmlCreateNode(xml_enabled_name.c_str(),
					   ntAttribute,
					   common::bool_to_str(enabled).c_str()));
	CxmlAddChild(rootnode, CxmlCreateNode("count", ntAttribute, std::to_string(count).c_str()));
	CxmlAddChild(rootnode, CxmlCreateNode(xml_name_name.c_str(), ntAttribute, name.c_str()));
	return rootnode;
}

ClassLabel& 
ClassLabel::from_node(TCXMLnode* node){
	name = CxmlGetAttribute(node, &string(xml_name_name)[0]);
	code = std::stoi(CxmlGetAttribute(node, &string(xml_code_name)[0]));
	enabled = common::str_to_bool(CxmlGetAttribute(node, &string(xml_enabled_name)[0]));
	return *this;
}


const set<label_t> 
MetaData::get_enabled_class_codes() const {
	set<label_t> res;
	for (auto v : _classes) {
		if (v.second.enabled)
			res.insert(v.second.code);
	}
	return res;
}

const set<string> 
MetaData::get_classes(label_t class_code) const {
	set<string> res;
	for (auto v : _classes) {
		if (v.second.code == class_code)
			res.insert(v.second.name);
	}
	return res;
}

bool 
MetaData::add_class(const ClassLabel& cl){
	if (class_exists(cl.name))
		return false;
	else{
        _classes.insert(std::make_pair(cl.name, cl));
		return true;
	}
}

bool MetaData::remove_class(const string& name) {
    if(!class_exists(name))
        return false;
    _classes.erase(_classes.find(name));
    return true;
}

bool inline
MetaData::class_exists(const string& name) const {
    return _classes.find(name) != _classes.end();
}


TCXMLnode* 
MetaData::to_node(void) const{
	TCXMLnode * rootnode = CxmlCreateNode(xml_root_name.c_str(), ntElement, NULL);

	CxmlAddChild(rootnode, CxmlCreateNode(
		xml_features_name.c_str(), 
		ntAttribute, 
		common::vec_to_str(_features).c_str()));
	TCXMLnode *n;
	CxmlAddChild(rootnode, n=CxmlCreateNode(xml_classes_name.c_str(), ntElement, NULL));
	CxmlAddChild(n, CxmlCreateNode("count", ntAttribute, std::to_string(_classes.size()).c_str()));
	for (auto v: _classes)
		CxmlAddChild(n,  v.second.to_node());
	return rootnode;
}

MetaData&
MetaData::from_node(TCXMLnode* node){
	_classes.clear();
	_features.clear();

	_features = common::split(CxmlGetAttribute(node, &string(xml_features_name)[0]), ';');
	TCXMLnode* classes_node = CxmlGetSubElemnetN(node, &string(xml_classes_name)[0]);
	size_t classes_count = std::stoi(CxmlGetAttribute(classes_node, "count"));
	
	for (size_t i = 0; i < classes_count; i++){
		ClassLabel cl;
		cl.from_node(CxmlGetSubElemnetI(classes_node, i));
		add_class(cl);
	}

	return *this;
}


Dataset::Dataset(const std::string& filename) {
    from_xml_file(filename);
}


const ClassLabel& 
Dataset::get_class_label(string label) const{
	return _meta._classes.at(label);
}

bool Dataset::add_sample(const Sample& s) {
    //check if exists
    if(_samples.find(s.path) != _samples.end())
        return false;
    //check if has new class
	if (!_meta.class_exists(s.label)) {
        ClassLabel cl;
        cl.code = 0;
        cl.enabled = false;
        cl.name = s.label;
		_meta.add_class(cl);
    }
    _samples.insert(std::make_pair(s.path, s));
    _meta._classes[s.label].count++;
    return true;
}

bool Dataset::remove_sample(const string& path) {
    auto s_it = _samples.find(path);
    if(s_it == _samples.end())
        return false;
    _meta._classes[s_it->second.label].count--;
    if(_meta._classes[s_it->second.label].count == 0)
        _meta.remove_class(s_it->second.label);
    _samples.erase(s_it);
    return true;
}

//enabled samples to X,y
const std::pair<Matrix<data_t>, vector<label_t> > 
Dataset::samples_to_Xy(const vector<Sample>& samples) const{
	Matrix<data_t> data(samples.size(), samples[0].data.size());
	vector<label_t> targets(samples.size());
	for (size_t i = 0; i < samples.size(); i++){
		data.set_row(samples[i].data, i);
		targets[i] = get_class_label(samples[i].label).code;
	}
	return std::make_pair(data, targets);
}

const vector<Sample>
Dataset::get_enabled_samples () const {
    vector<Sample> res;
    for(auto it = _samples.begin(); it != _samples.end(); it++) {
        if(_meta._classes.at(it->second.label).enabled) {
			res.push_back(it->second);
        }
    }
	return res;
}

const set<label_t> 
Dataset::get_enabled_class_codes() const {
    return _meta.get_enabled_class_codes();
}

const 
set<string> 
Dataset::get_classes(label_t class_code) const {
    return _meta.get_classes(class_code);
}

TCXMLnode*
Dataset::to_node(void) const{
	TCXMLnode * rootnode = CxmlCreateNode(xml_root_name.c_str(), ntElement, NULL);
	CxmlAddChild(rootnode, _meta.to_node());
	TCXMLnode * n;
	CxmlAddChild(rootnode, n=CxmlCreateNode(xml_samples_name.c_str(), ntElement, NULL));
	CxmlAddChild(n, CxmlCreateNode("count", ntAttribute, std::to_string(_samples.size()).c_str()));
	for (auto s : _samples)
		CxmlAddChild(n, s.second.to_node());
	return rootnode;
}

Dataset& 
Dataset::from_node(TCXMLnode* node){
	_samples.clear();
	TCXMLnode *meta_node = CxmlGetSubElemnetN(node, &string(MetaData::xml_root_name)[0]);
	_meta.from_node(meta_node);
	
	TCXMLnode *samples_node = CxmlGetSubElemnetN(node, &string(xml_samples_name)[0]);
	size_t samples_count = std::stoi(CxmlGetAttribute(samples_node, "count"));
	for (size_t i = 0; i < samples_count; i++){
		//if (i % 100 == 0 && InteropNegotiator::report_progress)
		//	InteropNegotiator::report_progress(InteropNegotiator::control_frame_handle,i / static_cast<float>(samples_count));
		Sample s;
		s.from_node(CxmlGetSubElemnetI(samples_node, i));
		add_sample(s);
	}

	return *this;
}

string
Dataset::info() const {
	std::stringstream ss;
	set<label_t> class_codes = get_enabled_class_codes();
	for (label_t label : class_codes) {
		ss << "label " << label << ": { ";
		size_t i = 0;
		for (const string& class_str : get_classes(label)) {
			if (++i % 4 == 0) ss << std::endl;
			ss << class_str << " ("
				<< "count: " << get_class_label(class_str).count << "), ";
		}
		ss.seekp(-2, ss.cur);
		ss << "}";
		ss << std::endl;
	}
	return ss.str();
}