/*
	This program has been created by Stefan Zwaard, based on the work of Paul Baker. Current version: 0.5, 16-06-2020.
	The (Distributed) Model Training and Aggregation (DMTLA) program allows for training and aggregation of SVM (HOG based) and ERT models.
	The program is part of the EyeBlink Project, and works together with the Image processing and EAR calculation program to proces images using the newly created models.
	The program was also used as a proof of concept and for results gathering for the proposed Mean Weight Matrix Aggregation (MWMA) distributed training for L-SVM models, and Weighted Bin Aggregation (WBA) distributed training of ERT models,
	in the paper: "Privacy-Preserving Algorithms for Object Detection & Localization Using Distributed Machine Learning".
	Running the resulting program requires opencv_world400.dll to be present in the same folder. In order to compile this code, ensure both the OpenCV and DLib libraries are linked.
	Always use the AVX extended instructions on Release x64 for best preformance of program
*/

#include <dlib/svm_threaded.h>
#include <dlib/string.h>
#include <dlib/gui_widgets.h>
#include <dlib/image_processing.h>
#include <dlib/data_io.h>
#include <dlib/cmd_line_parser.h>
#include <dlib/image_io.h>

#include <iostream>
#include <fstream>
#include <stdlib.h>
#include <stdio.h>
#include <cstdio>
#include <math.h>
#include <string>
#include <filesystem>

#include "Window.h"
#include "CERT.h"

#include <dlib/image_processing/frontal_face_detector.h>
#include <dlib/image_processing/render_face_detections.h>
#include <dlib/image_keypoint/draw_surf_points.h>
#include <opencv2/core/core.hpp>
#include <opencv2/features2d/features2d.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/calib3d/calib3d.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/opencv.hpp>
#include <dlib/opencv.h>

using namespace std;
using namespace cv;
using namespace dlib;

// Define scnaner type for SVM training
typedef scan_fhog_pyramid<pyramid_down<6> > image_scanner_type;

// Here follows a few function definitions from DLib for SVM and ERT training
// ----------------------------------------------------------------------------------------

std::vector<std::vector<double> > get_interocular_distances(
	const std::vector<std::vector<full_object_detection> >& objects
);

void pick_best_window_size(
	const std::vector<std::vector<dlib::rectangle> >& boxes,
	unsigned long& width,
	unsigned long& height,
	const unsigned long target_size
)
/*!
	ensures
		- Finds the average aspect ratio of the elements of boxes and outputs a width
		  and height such that the aspect ratio is equal to the average and also the
		  area is equal to target_size.  That is, the following will be approximately true:
			- #width*#height == target_size
			- #width/#height == the average aspect ratio of the elements of boxes.
!*/
{
	// find the average width and height
	running_stats<double> avg_width, avg_height;
	for (unsigned long i = 0; i < boxes.size(); ++i)
	{
		for (unsigned long j = 0; j < boxes[i].size(); ++j)
		{
			avg_width.add(boxes[i][j].width());
			avg_height.add(boxes[i][j].height());
		}
	}

	// Adjusting the box size sp it is about target_pizels pixels in size
	double size = avg_width.mean()*avg_height.mean();
	double scale = std::sqrt(target_size / size);

	width = (unsigned long)(avg_width.mean()*scale + 0.5);
	height = (unsigned long)(avg_height.mean()*scale + 0.5);
	// make sure the width and height never round to zero.
	if (width == 0)
		width = 1;
	if (height == 0)
		height = 1;
}

bool contains_any_boxes(
	const std::vector<std::vector<dlib::rectangle> >& boxes
)
{
	for (unsigned long i = 0; i < boxes.size(); ++i)
	{
		if (boxes[i].size() != 0)
			return true;
	}
	return false;
}

void throw_invalid_box_error_message(
	const std::string& dataset_filename,
	const std::vector<std::vector<dlib::rectangle> >& removed,
	const unsigned long target_size
)
{
	image_dataset_metadata::dataset data;
	load_image_dataset_metadata(data, dataset_filename);

	std::ostringstream sout;
	sout << "Error!  An impossible set of object boxes was given for training. ";
	sout << "All the boxes need to have a similar aspect ratio and also not be ";
	sout << "smaller than about " << target_size << " pixels in area. ";
	sout << "The following images contain invalid boxes:\n";
	std::ostringstream sout2;
	for (unsigned long i = 0; i < removed.size(); ++i)
	{
		if (removed[i].size() != 0)
		{
			const std::string imgname = data.images[i].filename;
			sout2 << "  " << imgname << "\n";
		}
	}
	throw dlib::error("\n" + wrap_string(sout.str()) + "\n" + sout2.str());
}

// ----------------------------------------------------------------------------------------

// General system variables, mostly used for the main GUI window
bool Start = 0;
int TrainSVMConfig = 0;
int AggregateSVMConfig = 0;
std::string SVMModelsInPath = "E:\\EyeBlink\\Models\\ToAggregate.svm";
std::string SVMModelOutPath = "E:\\EyeBlink\\Models\\NewModel.svm";
int TrainERTConfig = 0;
int AggregateERTConfig = 0;
std::string ERTModelsInPath = "E:\\EyeBlink\\Models\\ToAggregate.dat";
std::string ERTModelOutPath = "E:\\EyeBlink\\Models\\NewModel.dat";
std::string SVMTrainingDatasetPath = "E:\\EyeBlink\\Data\\SVMTrainingdata.xml";
std::string SVMTestDatasetPath = "E:\\EyeBlink\\Data\\SVMTestingdata.xml";
std::string ERTTrainingDatasetPath = "E:\\EyeBlink\\Data\\ERTTrainingdata.xml";
std::string ERTTestDatasetPath = "E:\\EyeBlink\\Data\\ERTTestingdata.xml";

// Variable decleration for SVM training settings, changed and set with the SVM training GUI
bool SVMDone = 0;
int SVMThreads = 6;
double C = 100.0;
double eps = 0.01;
unsigned int num_folds = 3;
unsigned long target_size = 80 * 80;
unsigned long upsample_amount = 2;

// Variable decleration for ERT training settings, changed and set with the ERT training GUI
bool ERTDone = 0;
unsigned int ERTThreads = 6;
unsigned long oversampling_amount = 20; // Random deformations of input data to help overcome generalization
double nu = 0.1; // 0,1 for default, the lower the more generalized training, 1 for complete fitting
unsigned long tree_depth = 5; // Default 2 for normal, 8 for high accuracy, verry slow on higher values, also severly increases model size be extreemly carefull
unsigned int feature_pool_size = 400; // Number of pixels used for ERT feature processing in each cascade. 400 default
unsigned int num_test_splits = 20; // Default 50. Changes accuracy of training, but slows training speed
unsigned int cascade_depth = 10; // Number of iterations
unsigned long num_trees_per_cascade_level = 500; // Default is 500
double lambda = 0.1;

int main(int argc, char** argv)
{
	try
	{
		// Showing main GUI
		win my_window(&Start, &TrainSVMConfig, &AggregateSVMConfig, &SVMModelsInPath, &SVMModelOutPath, &TrainERTConfig, &AggregateERTConfig, &ERTModelsInPath, &ERTModelOutPath, &SVMTrainingDatasetPath, &SVMTestDatasetPath, &ERTTrainingDatasetPath, &ERTTestDatasetPath);
		while (true)
		{
			while (Start == 0)
			{
				sleep(100); // Sleep while system waits for user to start program, Idle running can be implemented here if neeeded
			}
			cout << "Program has started" << endl;

			// If *.svm or *.dat is used to select more then one model of each type, for either multi model aggregation or testing, each model is seperated into its own file.
			std::vector<String> VectorSVMInputModelPaths;
			std::vector<String> VectorERTInputModelPaths; // Keep in mind, this path can also be used to load CERT models for testing only (option 3 in main GUI)
			std::vector<String> VectorCERTInputModelPaths; // Not used at the moment, could be used later if adding or removing Subdivisions from CERT model using other ERT models is needed. For now the ERT input is used.
			std::vector < object_detector<image_scanner_type>> SVMModels; // Vector for all loaded SVM models
			std::vector < shape_predictor> ERTModels; // Vector for all loaded ERT models
			std::vector < CERT > CERTModels; // Vector for all loaded CERT models

			if (AggregateSVMConfig == 1 || AggregateSVMConfig == 2) // Load SVM models from input path if needed
			{
				glob(SVMModelsInPath, VectorSVMInputModelPaths); // If a incursive function is wanted, where all subfolders of the target are also searched, add '1' to the arguments
				cout << "Number of .svm models detected on SVM model(s) input path:" << VectorSVMInputModelPaths.size() << endl;
				for (int i = 0; i < VectorSVMInputModelPaths.size(); i++)
				{
					ifstream fin(VectorSVMInputModelPaths[i], ios::binary);
					if (fin)
					{
						object_detector<image_scanner_type> TempDetector;
						cout << "Loaded .svm file for either multi model aggregation or test only testing, file:" << VectorSVMInputModelPaths[i] << endl;
						deserialize(TempDetector, fin);
						SVMModels.push_back(TempDetector);
					}
					else
						cout << "Error reading .svm file on SVM model(s) input path, file was:" << VectorSVMInputModelPaths[i] << endl;
				}
			}

			if (AggregateERTConfig == 1 || AggregateERTConfig == 2) // Load ERT models from input path if needed
			{
				glob(ERTModelsInPath, VectorERTInputModelPaths); // If a incursive function is wanted, where all subfolders of the target are also searched, add '1' to the arguments
				cout << "Number of .dat models detected on ERT model(s) input path:" << VectorERTInputModelPaths.size() << endl;
				for (int i = 0; i < VectorERTInputModelPaths.size(); i++)
				{
					ifstream fin(VectorERTInputModelPaths[i], ios::binary);
					if (fin)
					{
						shape_predictor TempDetector;
						cout << "Loaded .dat file for either Weighted Bin Aggregation or test only testing, file:" << VectorERTInputModelPaths[i] << endl;
						deserialize(TempDetector, fin);
						ERTModels.push_back(TempDetector);
					}
					else
						cout << "Error reading .dat file on ERT model(s) input path, file was:" << VectorERTInputModelPaths[i] << endl;
				}
			}

			if (AggregateERTConfig == 3) // Load CERT models from input path if needed (Uses ERT input path at the moment)
			{
				glob(ERTModelsInPath, VectorERTInputModelPaths); // If a incursive function is wanted, where all subfolders of the target are also searched, add '1' to the arguments
				cout << "Number of .CERT models detected on ERT model(s) input path:" << VectorERTInputModelPaths.size() << endl;
				for (int i = 0; i < VectorERTInputModelPaths.size(); i++)
				{
					ifstream fin(VectorERTInputModelPaths[i], ios::binary);
					if (fin)
					{
						CERT TempDetector;
						cout << "Loaded .CERT file for either WBA CERT model only testing, file:" << VectorERTInputModelPaths[i] << endl;
						TempDetector.deserialize(fin);
						CERTModels.push_back(TempDetector);
					}
					else
						cout << "Error reading .CERT file on CERT model(s) input path, file was:" << VectorERTInputModelPaths[i] << endl;
				}
			}

			if (TrainSVMConfig == 1) // Getting SVM training settings from user if needed
			{
				cout << "Showing SVM model training window now" << endl;
				SVMsettingsWin SVMWindow(&SVMDone, &SVMThreads, &C, &eps, &num_folds, &target_size, &upsample_amount);
				while (SVMDone == 0)
				{
					sleep(100); // Wait while user finishes SVM training input
				}
				SVMDone = 0;
				// Callout settings back to user
				cout << "SVM model training configured, using the following settings:" << endl;
				cout << "threads:" << SVMThreads << endl;
				cout << "C:" << C << endl;
				cout << "eps:" << eps << endl;
				cout << "num_folds:" << num_folds << endl;
				cout << "target_size:" << target_size << endl;
				cout << "upsample_amount" << upsample_amount << endl;
			}

			if (TrainERTConfig == 1) // Getting ERT training settings from user if needed
			{
				cout << "Showing ERT model training window now" << endl;
				ERTsettingsWin ERTWindow(&ERTDone, &ERTThreads, &oversampling_amount, &nu, &tree_depth, &feature_pool_size, &num_test_splits, &cascade_depth, &num_trees_per_cascade_level, &lambda);
				while (ERTDone == 0)
				{
					sleep(100); // Wait while user finishes SVM training input
				}
				ERTDone = 0;
				// Callout settings back to user
				cout << "ERT model training configured, using the following settings:" << endl;
				cout << "threads:" << ERTThreads << endl;
				cout << "oversampling_amount:" << oversampling_amount << endl;
				cout << "nu:" << nu << endl;
				cout << "tree_depth:" << tree_depth << endl;
				cout << "feature_pool_size:" << feature_pool_size << endl;
				cout << "num_test_splits:" << num_test_splits << endl;
				cout << "cascade_depth:" << cascade_depth << endl;
				cout << "num_trees_per_cascade_level:" << num_trees_per_cascade_level << endl;
				cout << "lambda:" << lambda << endl;
			}

			// SVM Model training if needed (based on default DLib Training)
			object_detector<image_scanner_type> NewSVMDetector; // Placeholder for the new created SVM model
			if (TrainSVMConfig == 1) // Training new .svm model based on user settings
			{
				cout << "Starting preperations for SVM model training." << endl;
				dlib::array<array2d<unsigned char> > images;
				std::vector<std::vector<dlib::rectangle> > object_locations, ignore;

				cout << "Loading training image dataset from dataset .xml file " << SVMTrainingDatasetPath << endl;
				ignore = load_image_dataset(images, object_locations, SVMTrainingDatasetPath);
				cout << "Number of images loaded: " << images.size() << endl;

				// Check if there are more folds than there are images.
				if (num_folds > images.size())
					num_folds = images.size();

				// Upsampling if needed
				for (unsigned long i = 0; i < upsample_amount; ++i)
					upsample_image_dataset<pyramid_down<2> >(images, object_locations, ignore);

				image_scanner_type scanner;

				unsigned long width, height;
				pick_best_window_size(object_locations, width, height, target_size);
				scanner.set_detection_window_size(width, height);
				structural_object_detection_trainer<image_scanner_type> trainer(scanner);

				trainer.set_num_threads(SVMThreads);
				trainer.be_verbose(); // This is always set, this gives a progress overview in the CLI during training, remove this is unwanted, or make it depending on user settings.
				trainer.set_c(C);
				trainer.set_epsilon(eps);

				// Making sure input trainingsdata has boxes usable by detector
				std::vector<std::vector<dlib::rectangle> > removed;
				removed = remove_unobtainable_rectangles(trainer, images, object_locations);
				// Error handeling for if a box does not match (this happens a lot if there is no upsampling)
				if (contains_any_boxes(removed))
				{
					unsigned long scale = upsample_amount + 1;
					scale = scale * scale;
					throw_invalid_box_error_message(SVMTrainingDatasetPath, removed, target_size / scale);
				}
				cout << "Preperations completed. Starting training on new SVM model" << endl;
				NewSVMDetector = trainer.train(images, object_locations, ignore);
				cout << "Training of new SVM  model has completed." << endl;
			}

			// ERT model training if needed (based on default DLib Training)
			shape_predictor NewERTPredictor; // Placeholder for the new ERT model created
			if (TrainERTConfig == 1) // training new .dat model based on user settings
			{
				cout << "Starting preperations for ERT model training." << endl;
				cout << "Loading training image dataset from dataset .xml file " << ERTTrainingDatasetPath << endl;

				dlib::array<array2d<unsigned char> > images_train;
				std::vector<std::vector<full_object_detection> > faces_train;
				load_image_dataset(images_train, faces_train, ERTTrainingDatasetPath);
				cout << "Number of images loaded: " << images_train.size() << endl;

				shape_predictor_trainer trainer;

				// Applying ERT training settings
				trainer.set_nu(nu);
				trainer.set_tree_depth(tree_depth);
				trainer.set_num_threads(ERTThreads);
				if (oversampling_amount > 0)
					trainer.set_oversampling_amount(oversampling_amount);
				trainer.set_cascade_depth(cascade_depth);
				trainer.set_feature_pool_size(feature_pool_size);
				trainer.set_num_test_splits(num_test_splits);
				trainer.set_num_trees_per_cascade_level(num_trees_per_cascade_level);
				trainer.set_lambda(lambda);

				trainer.be_verbose(); // This is always set, this gives a progress overview in the CLI during training, remove this is unwanted, or make it depending on user settings.

				cout << "Preperations completed. Starting training on new ERT model" << endl;

				NewERTPredictor = trainer.train(images_train, faces_train);
				cout << "Training of new ERT  model has completed." << endl;
				cout << endl << "Number of parts of ERT Model: " << NewERTPredictor.num_parts();
				cout << endl << "Number of feutures of of ERT Model" << NewERTPredictor.num_features() << endl;
			}

			// SVM Model aggregation if needed, using the Mean Weight Matrix Aggregation (MWMA) algorithm, see paper for detailed explenation
			if (AggregateSVMConfig == 1) // Aggregating new .svm model based on user settings, aggregates any trained .svm model as well
			{
				cout << "Starting SVM multi model aggregation process" << endl;
				cout << "Loaded " << SVMModels.size() << " SVM models from input path for aggregation proces" << endl;
				if (TrainSVMConfig == 1)
					cout << "Trained SVM model shall be included in aggregation proces" << endl;

				cout << "Starting SVM multi aggregation" << endl;

				matrix<double, 0, 1> TempW; // MWMA step 1: New temp matrix for aggregation
				TempW.set_size((SVMModels[0].get_w()).size());
				TempW = SVMModels[0].get_w(); // MWMA step 1: init of w using first model
				for (int i = 1; i < SVMModels.size(); i++)
				{
					TempW += SVMModels[i].get_w(); // MWMA step 2a: Adding all other matrix values
				}
				if (TrainSVMConfig == 1)
				{
					TempW += NewSVMDetector.get_w(); // MWMA step 2a: Inclusion of trained SVM model in aggregation
					TempW /= (SVMModels.size() + 1); // MWMA step 2b: averaging all matrix values including trained model
				}
				else
					TempW /= SVMModels.size(); // MWMA step 2b: averageing all matrix values

				NewSVMDetector = object_detector<image_scanner_type>(SVMModels[0].get_scanner(), SVMModels[0].get_overlap_tester(), TempW); // MWMA step 3: Feuture extraction transfer and saving
				cout << "SVM Multi aggregation completed" << endl;
			}

			// ERT Model aggregation if needed, using tthe Weighted Bin Aggregation (MWMA) algorithm, this combines several ERT models into a new Comnbined-ERT model (CERT), see paper for detailed explenation
			CERT NewCERT;
			if (AggregateERTConfig == 1) // Aggregating new ERT model based on user settings, aggregates trained .dat model as well
			{
				cout << "Starting ERT Weighted Bin Aggregation process" << endl;
				cout << "Loaded " << ERTModels.size() << " ERT models from input path for aggregation proces" << endl;
				if (TrainERTConfig == 1)
				{
					cout << "Trained ERT model shall be included in aggregation proces" << endl;
					ERTModels.push_back(NewERTPredictor);
				}

				// WBA Step 1: Forrest combination
				cout << "ERT WBA: Forrest combination" << endl;
				std::vector<dlib::shape_predictor> NewSubDivisions;
				for (int i = 0; i < ERTModels.size(); i++)
				{
					NewSubDivisions.push_back(ERTModels[i]);
				}

				// WBA Step 2: Devider calculation
				cout << "ERT WBA: Devider calculation" << endl;
				std::vector<double> NewDevider; // The devider represents how much the subdevision partakes in the calculation of the end result and should total 1
				std::vector<double> DeviderOffSet; // This vector contains all offset values in order to bias end outcome towards a certain model. This could be filled using user config later, but for now its hardcoded for euqal participation (by filling it with zero's only.
				double TotalDeviationValue = 0; // This number is needed to determin how the end result index of 1 is reached. If all moddels have equal weight, this is the number of models (to ensure mean average). If a model has more or less weight, then this model should be changed.
				for (int i = 0; i < ERTModels.size(); i++) // Offset calculation
				{
					DeviderOffSet.push_back(1); // Adding ones first always, a 1 means 100% base of this model is used for final weight calculation, so place in vector is filled even if no offset is given, 1's only means all models partake equaly.
					DeviderOffSet[i] += 0; // Add specific model offset here (later trough GUI settings, but now hardcoded). For 20% more bias towards a model, add 0.2 (this ofcourse also means that other models lose participation value towards this model. Therefore If all models get the same increase in bias, it has no effect; 2's for all modes results in same for 1's in all.
					TotalDeviationValue += DeviderOffSet[i]; // If a model has increased or decreased weight, the toal deviation should always be changed, to ensure the later total result index always equals 1.
				}
				for (int i = 0; i < ERTModels.size(); i++) // Final devider calculation, by pushing back the final deviation values now instead of in the first loop, changes can be made depending on the TotalDeviationValue if needed.
				{
					double TempBinDeviation;
					TempBinDeviation = DeviderOffSet[i];
					NewDevider.push_back(TempBinDeviation);
				}

				// WBA Step 3: Weighted Bins Creation
				cout << "ERT WBA: Weighted bins creation" << endl;
				std::vector<std::pair<dlib::full_object_detection, double>> NewWeightedBins;
				for (int i = 0; i < ERTModels.size(); i++) // Creation of all bins, one bin for each subdevision, this is used to store sub results and used in combination with the corresponding devider for end result calculation
				{
					dlib::full_object_detection NewBin;
					std::pair<dlib::full_object_detection, double> NewWeightedBin(NewBin, NewDevider[i]);
					NewWeightedBins.push_back(NewWeightedBin);
				}

				// WBA step 4: Averaging of init shape model
				cout << "ERT WBA: Aggregating shapes"<< endl;
				matrix<float, 0, 1> NewShape; // New init Shape for averaging
				NewShape.set_size((ERTModels[0].get_shape()).size());
				NewShape = ERTModels[0].get_shape();
				for (int i = 1; i < ERTModels.size(); i++)
				{
					NewShape += ERTModels[i].get_shape();

				}
				NewShape /= ERTModels.size(); // Averageing all matrix values

				NewCERT = CERT(NewSubDivisions, TotalDeviationValue, NewWeightedBins, NewShape); // Creation of new combi ERT model

				cout << "ERT WBA completed" << endl;
			}

			// Model testing and saving
			if (AggregateSVMConfig == 1  || TrainSVMConfig == 1) // Saving and testing new .svm model if one has been created
			{
				cout << "Saving newly created SVM model to disk." << endl;
				serialize(SVMModelOutPath) << NewSVMDetector; // saving new SVM file to disk

				// Testing newly created SVM model
				dlib::array<array2d<unsigned char> > images;
				std::vector<std::vector<dlib::rectangle> > object_locations, ignore;
				cout << "Loading Test image dataset from dataset.xml file " << SVMTestDatasetPath << endl;
				ignore = load_image_dataset(images, object_locations, SVMTestDatasetPath);
				cout << "Number of images loaded: " << images.size() << endl;
				cout << "Testing new SVM model on specified test dataset." << endl;
				cout << "Accuracy of SVM model (precision,recall,AP): " << test_object_detection_function(NewSVMDetector, images, object_locations, ignore) << endl;
			}

			if (AggregateSVMConfig == 2) // Testing selected input .svm models if test only mode is used
			{
				dlib::array<array2d<unsigned char> > images;
				std::vector<std::vector<dlib::rectangle> > object_locations, ignore;
				cout << "Loading SVM Test image dataset from dataset.xml file " << SVMTestDatasetPath << endl;
				ignore = load_image_dataset(images, object_locations, SVMTestDatasetPath);
				cout << "Number of images loaded: " << images.size() << endl;

				cout << "Testing:" << SVMModels.size() << " SVM model(s) on specified test dataset." << endl;
				for (int i = 0; i < SVMModels.size(); i++)
				{
					cout << "Testing SVM model:" << VectorSVMInputModelPaths[i] << endl;
					cout << "Accuracy of SVM model  (precision,recall,AP): " << test_object_detection_function(SVMModels[i], images, object_locations, ignore) << endl;
				}
				cout << "Testing of SVM models completed." << endl;
			}

			if (TrainERTConfig == 1) // Saving and testing new .dat model if one has been trained
			{
				cout << "Saving newly trained ERT model to disk." << endl;
				serialize(ERTModelOutPath) << NewERTPredictor; // saving new .dat file to disk

				cout << "Loading ERT Test image dataset from dataset.xml file " << ERTTestDatasetPath << endl;
				dlib::array<array2d<unsigned char> > images_test;
				std::vector<std::vector<full_object_detection> > faces_test;
				load_image_dataset(images_test, faces_test, ERTTestDatasetPath);
				cout << "Number of images loaded: " << images_test.size() << endl;

				cout << "Mean testing error of new ERT model: " <<
					test_shape_predictor(NewERTPredictor, images_test, faces_test, get_interocular_distances(faces_test)) << endl;  // Test ERT model using internal test software
			}

			if (AggregateERTConfig == 1) // Saving and testing new .CERT model if one has been created
			{
				cout << "Saving newly aggregated CERT model to disk." << endl;
				ofstream fout(ERTModelOutPath, ios::binary);
				if (fout)
				{
					NewCERT.serialize(fout);
				}
				else
					cout << "Error saving CERT model. Model output path was:" << ERTModelOutPath << endl;

				cout << "Loading ERT Test image dataset from dataset.xml file " << ERTTestDatasetPath << endl;
				dlib::array<array2d<unsigned char> > images_test;
				std::vector<std::vector<full_object_detection> > faces_test;
				load_image_dataset(images_test, faces_test, ERTTestDatasetPath);
				cout << "Number of images loaded: " << images_test.size() << endl;

				const std::vector<std::vector<double> > scales = get_interocular_distances(faces_test);
				running_stats<double> rs;
				for (unsigned long i = 0; i < faces_test.size(); ++i)
				{
					for (unsigned long j = 0; j < faces_test[i].size(); ++j)
					{
						const double scale = scales.size() == 0 ? 1 : scales[i][j]; // Scale needed for testing
					    full_object_detection Shape = NewCERT.PredictFinalShape(images_test[i], faces_test[i][j].get_rect()); // This provides a full_object_detection from a CERT model, use this function if a CERT model is used where normaly a ERT model is used.
						for (unsigned long k = 0; k < Shape.num_parts(); k++) // The new CERT model has no internal test function like the default ERT code from DLib, therefore the Mean Error Rate is calculated here below.
						{
							double score = length(Shape.part(k) - faces_test[i][j].part(k)) / scale;
							rs.add(score);
						}
					}
				}
				cout << "Mean testing error of new ERT model: " << rs.mean() << endl;

			}

			if (AggregateERTConfig == 2) // Testing selected input .dat models if test only mode is used. This tests ERT models, not the Combined-ERT models, can be used to see most suitable models for aggregation
			{
				cout << "Loading ERT Test image dataset from dataset.xml file " << ERTTestDatasetPath << endl;
				dlib::array<array2d<unsigned char> > images_test;
				std::vector<std::vector<full_object_detection> > faces_test;
				load_image_dataset(images_test, faces_test, ERTTestDatasetPath);
				cout << "Number of images loaded: " << images_test.size() << endl;

				cout << "Testing:" << ERTModels.size() << " ERT model(s) on specified test dataset." << endl;

				for (int i = 0; i < ERTModels.size(); i++)
				{
					cout << "Testing ERT model:" << VectorERTInputModelPaths[i] << endl;
					cout << "Mean testing error of ERT model: " << test_shape_predictor(ERTModels[i], images_test, faces_test, get_interocular_distances(faces_test)) << endl; // Test ERT model using internal test software
				}

			}

			if (AggregateERTConfig == 3) // Testing selected input CERT .dat models if test only mode is used
			{
				cout << "Loading ERT Test image dataset from dataset.xml file " << ERTTestDatasetPath << endl;
				dlib::array<array2d<unsigned char> > images_test;
				std::vector<std::vector<full_object_detection> > faces_test;
				load_image_dataset(images_test, faces_test, ERTTestDatasetPath);
				cout << "Number of images loaded: " << images_test.size() << endl;

				cout << "Testing:" << CERTModels.size() << " CERT model(s) on specified test dataset." << endl;

				for (int m = 0; m < CERTModels.size(); m++)
				{
					cout << "Testing CERT model:" << VectorERTInputModelPaths[m] << endl;
					const std::vector<std::vector<double> > scales = get_interocular_distances(faces_test);
					running_stats<double> rs;
					for (unsigned long i = 0; i < faces_test.size(); ++i)
					{
						for (unsigned long j = 0; j < faces_test[i].size(); ++j)
						{
							const double scale = scales.size() == 0 ? 1 : scales[i][j]; // Scale needed for testing
							full_object_detection Shape = CERTModels[m].PredictFinalShape(images_test[i], faces_test[i][j].get_rect()); // This provides a full_object_detection from a CERT model, use this function if a CERT model is used where normaly a ERT model is used.
							for (unsigned long k = 0; k < Shape.num_parts(); k++) // The new CERT model has no internal test function like the default ERT code from DLib, therefore the Mean Error Rate is calculated here below.
							{
								double score = length(Shape.part(k) - faces_test[i][j].part(k)) / scale;
								rs.add(score);
							}
						}
					}
					cout << "Mean testing error of CERT model: " << rs.mean() << endl;
				}
			}

			Start = 0; // System is done and shall await next usage
			cout << "Program has completed task, a new task can now be started using the start button, please change settings where needed." << endl;
		}
	}
	catch (exception& e)
	{
		cout << "\nexception thrown!" << endl;
		cout << e.what() << endl;
		cout << "Restart of program needed, close window or terminate program from CLI." << endl;
	}
}

// Implementaiton of remaining DLib functions for testing and training of ERT models, needed to calculate the scale
// ----------------------------------------------------------------------------------------

double interocular_distance(
	const full_object_detection& det
)
{
	dlib::vector<double, 2> l, r;
	double cnt = 0;
	// Find the center of the left eye by averaging the points around
	// the eye.
	for (unsigned long i = 36; i <= 41; ++i)
	{
		l += det.part(i);
		++cnt;
	}
	l /= cnt;

	// Find the center of the right eye by averaging the points around
	// the eye.
	cnt = 0;
	for (unsigned long i = 42; i <= 47; ++i)
	{
		r += det.part(i);
		++cnt;
	}
	r /= cnt;

	// Now return the distance between the centers of the eyes
	return length(l - r);
}

std::vector<std::vector<double> > get_interocular_distances(
	const std::vector<std::vector<full_object_detection> >& objects
)
{
	std::vector<std::vector<double> > temp(objects.size());
	for (unsigned long i = 0; i < objects.size(); ++i)
	{
		for (unsigned long j = 0; j < objects[i].size(); ++j)
		{
			temp[i].push_back(interocular_distance(objects[i][j]));
		}
	}
	return temp;
}


//          Copyright Gavin Band 2008 - 2012.
// Distributed under the Boost Software License, Version 1.0.
//    (See accompanying file LICENSE_1_0.txt or copy at
//          http://www.boost.org/LICENSE_1_0.txt)

#include <string>
#include <vector>
#include <boost/function.hpp>
#include <boost/bind.hpp>
#include <boost/math/distributions/chi_squared.hpp>
#include "genfile/Error.hpp"
#include "genfile/ToGP.hpp"
#include "genfile/VariantDataReader.hpp"
#include "metro/likelihood/Multinomial.hpp"
#include "components/SNPSummaryComponent/SNPHWE.hpp"
#include "components/SNPSummaryComponent/SNPSummaryComputation.hpp"
#include "components/SNPSummaryComponent/HWEComputation.hpp"
#include "components/SNPSummaryComponent/IntensitySummaryComputation.hpp"
#include "components/SNPSummaryComponent/ClusterFitComputation.hpp"

// #define DEBUG_SNP_SUMMARY_COMPUTATION 1

namespace snp_summary_component {

	namespace {
		struct AlleleCountClient: public genfile::VariantDataReader::PerSampleSetter {
			AlleleCountClient( std::vector< double >* counts ):
				m_counts( counts ),
				m_number_of_alleles( 0 )
			{
				assert( counts != 0 ) ;
			}

			~AlleleCountClient() throw() {}

			void set_counts( std::vector< double >* counts ) {
				assert( counts != 0 ) ;
				m_counts = counts ;
			}

			void initialise( std::size_t number_of_samples, std::size_t number_of_alleles ) {
				m_number_of_alleles = number_of_alleles ;
				assert( m_counts->size() == number_of_alleles ) ;
			}

			bool set_sample( std::size_t i ) {
				m_ploidy = 0 ;
				m_table = 0 ;
				return true ;
			}

			void set_number_of_entries( uint32_t ploidy, std::size_t n, OrderType const order_type, ValueType const value_type ) {
				assert( order_type == genfile::ePerUnorderedGenotype ) ;
				assert( value_type == genfile::eProbability ) ;

				std::map< uint32_t, genfile::impl::Enumeration >::iterator where = m_tables.find( ploidy ) ;
				if( where == m_tables.end() ) {
					std::pair< std::map< uint32_t, genfile::impl::Enumeration >::iterator, bool >
						result = m_tables.insert( std::make_pair( ploidy, genfile::impl::enumerate_unphased_genotypes( ploidy ) )) ;
					assert( result.second ) ;
					where = result.first ;
				}
				m_ploidy = ploidy ;
				m_table = &(where->second) ;
			}

			void set_value( std::size_t value_i, genfile::MissingValue const value ) {
				set_value( value_i, 0.0 ) ;
			}

			void set_value( std::size_t value_i, double const value ) {
				assert( value_i <= std::size_t( std::numeric_limits< uint16_t >::max() )) ;
				genfile::impl::Enumeration const& enumeration = *m_table ;
				std::size_t const& maxAlleles = enumeration.first.second ;
				assert( m_number_of_alleles <= maxAlleles ) ;
				uint16_t const encodedGenotype = enumeration.second.second[ uint16_t( value_i ) ] ;
				uint32_t const bitsPerAllele = enumeration.first.first ;
				uint16_t mask = uint16_t( 0xFFFF ) >> ( 16 - bitsPerAllele ) ;
				// Dosage of the 1st allele in the genotype is not encoded directly.
				// We compute it as the ploidy minus the dosage of other alleles.
				uint16_t dosage_of_nonref_alleles = 0 ;
				for( std::size_t allele = 1; allele < m_number_of_alleles; ++allele ) {
					uint16_t const dosage = ( encodedGenotype >> ( (allele-1) * bitsPerAllele )) & mask ;
					(*m_counts)[allele] += value * dosage ;
					dosage_of_nonref_alleles += dosage ;
				}
				(*m_counts)[0] += value * (m_ploidy - dosage_of_nonref_alleles) ;
			}

			void finalise() {}
		private:
			std::vector< double >* m_counts ;
			std::size_t m_number_of_alleles ;
			uint32_t m_ploidy ;
			std::map< uint32_t, genfile::impl::Enumeration > m_tables ;
			genfile::impl::Enumeration* m_table ;
		} ;
	}

	struct AlleleCountComputation: public SNPSummaryComputation {
	public:

		AlleleCountComputation():
			m_counter( &m_counts )
		{}

		void list_variables( NameCallback callback ) const {
			using genfile::string_utils::to_string ;
			// By default we list 10 alleles
			callback( "number_of_alleles" ) ;
			for( std::size_t i = 0; i < 10; ++i ) {
				callback( "allele" + to_string(i+1) + "_count" ) ;
			}
		}

		void operator()(
			VariantIdentifyingData const& snp,
			Genotypes const& genotypes,
			Ploidy const& ploidy,
			genfile::VariantDataReader& data_reader,
			ResultCallback callback
		) {
			using genfile::string_utils::to_string ;
			m_counts = std::vector< double >( snp.number_of_alleles(), 0.0 ) ;
			m_counter.set_counts( &m_counts ) ;
			callback( "number_of_alleles", int64_t( snp.number_of_alleles() )) ;
			data_reader.get( ":genotypes:", genfile::to_GP_unphased( m_counter ) ) ;
			std::size_t i = 0 ;
			for( ; i < snp.number_of_alleles(); ++i ) {
				callback( "allele" + to_string(i+1) + "_count", m_counts[i] ) ;
			}
			for( ; i < 10; ++i ) {
				callback( "allele" + to_string(i+1) + "_count", genfile::MissingValue() ) ;
			}
		}

		std::string get_summary( std::string const& prefix, std::size_t column_width ) const {
			return prefix + "AlleleCountComputation" ;
		}
	private:
		std::vector< double > m_counts ;
		AlleleCountClient m_counter ;
	} ;

	struct AlleleFrequencyComputation: public SNPSummaryComputation
	{
		AlleleFrequencyComputation( std::string const& what ):
			m_compute_counts( what == "everything" || what == "counts" ),
			m_compute_frequencies( what == "everything" )
		{
			assert( what == "counts" || what == "everything" ) ;
		}
		void operator()(
			VariantIdentifyingData const& snp,
			Genotypes const& genotypes,
			Ploidy const& ploidy,
			genfile::VariantDataReader&,
			ResultCallback callback
		) {
			genfile::Chromosome const& chromosome = snp.get_position().chromosome() ;
			if( snp.number_of_alleles() == 2 ) {
				bool const allDiploid = ( ploidy.array() == 2 ).cast< int >().sum() == ploidy.size() ;
				if( allDiploid ) {
					compute_autosomal_frequency( snp, genotypes, callback ) ;
				} else {
					compute_sex_chromosome_frequency( snp, genotypes, ploidy, callback ) ;
				}
			}
		}

		void compute_sex_chromosome_frequency(
			VariantIdentifyingData const& snp,
			Genotypes const& genotypes,
			Ploidy const& ploidy,
			ResultCallback callback
		) {
			Genotypes haploid_genotypes = genotypes ;
			Genotypes diploid_genotypes = genotypes ;
			assert( std::size_t( genotypes.rows() ) == ploidy.size() ) ;

			for( int i = 0; i < genotypes.rows(); ++i ) {
				if( ploidy(i) != 1 ) {
					haploid_genotypes.row(i).setZero() ;
				}
				if( ploidy(i) != 2 ) {
					diploid_genotypes.row(i).setZero() ;
				}
			}

			double const a_allele_count = haploid_genotypes.col(0).sum()
				+ ( ( 2.0 * diploid_genotypes.col(0).sum() ) + diploid_genotypes.col(1).sum() ) ;
			double const b_allele_count = haploid_genotypes.col(1).sum()
				+ ( ( 2.0 * diploid_genotypes.col(2).sum() ) + diploid_genotypes.col(1).sum() ) ;

			if( m_compute_counts ) {
				callback( "alleleA_count", a_allele_count ) ;
				callback( "alleleB_count", b_allele_count ) ;
			}

			if( m_compute_frequencies ) {
				double const total_allele_count = ( haploid_genotypes.sum() + 2.0 * diploid_genotypes.sum() ) ;
				double const a_allele_freq = a_allele_count / total_allele_count ;
				double const b_allele_freq = b_allele_count / total_allele_count ;

				callback( "alleleA_frequency", a_allele_freq ) ;
				callback( "alleleB_frequency", b_allele_freq ) ;

				if( a_allele_freq < b_allele_freq ) {
					callback( "minor_allele_frequency", a_allele_freq ) ;
					callback( "minor_allele", snp.get_allele(0) ) ;
					callback( "major_allele", snp.get_allele(1) ) ;
				}
				else if( a_allele_freq > b_allele_freq ) {
					callback( "minor_allele_frequency", b_allele_freq ) ;
					callback( "minor_allele", snp.get_allele(1) ) ;
					callback( "major_allele", snp.get_allele(0) ) ;
				} else {
					callback( "minor_allele_frequency", a_allele_freq ) ;
				}
			}
		}

		void compute_autosomal_frequency( VariantIdentifyingData const& snp, Genotypes const& genotypes, ResultCallback callback ) {
			double const a_allele_count = ( 2.0 * genotypes.col(0).sum() ) + genotypes.col(1).sum() ;
			double const b_allele_count = ( 2.0 * genotypes.col(2).sum() ) + genotypes.col(1).sum() ;

			if( m_compute_counts ) {
				callback( "alleleA_count", a_allele_count ) ;
				callback( "alleleB_count", b_allele_count ) ;
			}

			if( m_compute_frequencies ) {
				double const total_allele_count = ( 2.0 * genotypes.sum() ) ;
				double const a_allele_freq = a_allele_count / total_allele_count ;
				double const b_allele_freq = b_allele_count / total_allele_count ;

				callback( "alleleA_frequency", a_allele_freq ) ;
				callback( "alleleB_frequency", b_allele_freq ) ;

				if( a_allele_freq < b_allele_freq ) {
					callback( "minor_allele_frequency", a_allele_freq ) ;
					callback( "minor_allele", snp.get_allele(0) ) ;
					callback( "major_allele", snp.get_allele(1) ) ;
				}
				else if( a_allele_freq > b_allele_freq ) {
					callback( "minor_allele_frequency", b_allele_freq ) ;
					callback( "minor_allele", snp.get_allele(1) ) ;
					callback( "major_allele", snp.get_allele(0) ) ;
				} else {
					callback( "minor_allele_frequency", a_allele_freq ) ;
				}
			}
		}

		std::string get_summary( std::string const& prefix, std::size_t column_width ) const {
			return prefix + "AlleleFrequencyComputation" ;
		}
	private:
		bool const m_compute_counts ;
		bool const m_compute_frequencies ;
	} ;

	// What proportion of the mass on a genotype is due to high-confidence calls?
	struct CallMassComputation: public SNPSummaryComputation
	{

		CallMassComputation( double const threshhold = 0.9 ):
			m_threshhold(threshhold)
		{}

		void operator()( VariantIdentifyingData const& snp, Genotypes const& genotypes, Ploidy const& ploidy, genfile::VariantDataReader&, ResultCallback callback ) {
			genfile::Chromosome const& chromosome = snp.get_position().chromosome() ;
			bool const allDiploid = ( ploidy.array() == 2 ).cast< int >().sum() == ploidy.size() ;
			if( allDiploid ) {
				compute_autosomal_call_mass( snp, genotypes, callback ) ;
			} else {
				return ;
			}
		}

		void compute_autosomal_call_mass( VariantIdentifyingData const& snp, Genotypes const& genotypes, ResultCallback callback ) {
			Eigen::VectorXd const masses = genotypes.colwise().sum() ;
			m_hcGenotypes = ( genotypes.array() > m_threshhold ).cast< double >()  * genotypes.array() ;
			Eigen::VectorXd const hcMasses = m_hcGenotypes.colwise().sum() ;

			callback( "AA_mass_propn", hcMasses(0)/masses(0) ) ;
			callback( "AB_mass_propn", hcMasses(1)/masses(1) ) ;
			callback( "BB_mass_propn", hcMasses(2)/masses(2) ) ;

			callback( "non-AA-mass_propn", (hcMasses(1)+hcMasses(2))/(masses(1)+masses(2))) ;
			callback( "non-BB-mass_propn", (hcMasses(1)+hcMasses(0))/(masses(1)+masses(0))) ;
		}

		std::string get_summary( std::string const& prefix, std::size_t column_width ) const {
			return prefix + "CallMassComputation" ;
		}
	private:
		double const m_threshhold ;
		Genotypes m_hcGenotypes ;
	} ;

	struct MissingnessComputation: public SNPSummaryComputation {
		MissingnessComputation( double call_threshhold = 0.9 ): m_call_threshhold( call_threshhold ) {}
		void operator()(
			VariantIdentifyingData const& snp,
			Genotypes const& genotypes,
			Ploidy const& ploidy,
			genfile::VariantDataReader&,
			ResultCallback callback
		) {
			assert( std::size_t( genotypes.rows() ) == ploidy.size() ) ;

			double missingness = double( genotypes.rows() ) - genotypes.array().sum() ;
			callback( "missing_proportion", missingness / double( genotypes.rows() ) ) ;

			bool const allDiploid = ( ploidy.array() == 2 ).cast< int >().sum() == ploidy.size() ;
			if( allDiploid ) {
				callback( "A", 0 ) ;
				callback( "B", 0 ) ;
				if( snp.number_of_alleles() == 2 ) {
					callback( "AA", genotypes.col(0).sum() ) ;
					callback( "AB", genotypes.col(1).sum() ) ;
					callback( "BB", genotypes.col(2).sum() ) ;
				}
				callback( "NULL", genotypes.rows() - genotypes.sum() ) ;
			} else {
				compute_haploid_diploid_counts( snp, genotypes, ploidy, callback ) ;
			}
			callback( "total", genfile::VariantEntry::Integer( genotypes.rows() )) ;
		}

		void compute_haploid_diploid_counts(
			VariantIdentifyingData const& snp,
			Genotypes const& genotypes,
			Ploidy const& ploidy,
			ResultCallback callback
		) {
			std::map< int, Eigen::VectorXd > counts ;
			std::map< int, double > null_counts ;
			counts[ -1 ] = Eigen::VectorXd::Zero( 3 ) ;
			counts[ 0 ] = Eigen::VectorXd::Zero( 3 ) ;
			counts[ 1 ] = Eigen::VectorXd::Zero( 3 ) ;
			counts[ 2 ] = Eigen::VectorXd::Zero( 3 ) ;
			char countKey[5] = { '.', '.', 'm', 'f', '.' } ;
			std::map< int, std::size_t > sample_counts ;

			for( std::size_t i = 0; i < ploidy.size(); ++i ) {
				counts[ ploidy(i) ] += genotypes.row( i ) ;
				null_counts[ ploidy(i) ] += ( 1 - genotypes.row(i).sum() ) ;
				++sample_counts[ ploidy(i) ] ;
#if DEBUG_SNP_SUMMARY_COMPUTATION
				if( ploidy(i) == 1 && genotypes(i,2) != 0 ) {
					std::cerr << "! ( MissingnessComputation::compute_sex_chromosome_counts() ): individual " << (i+1) << "is male but has genotype " << genotypes.row(i) << "!!\n" ;
				}
#endif
			}

			callback( "A", counts[ 1 ]( 0 ) ) ;
			callback( "B", counts[ 1 ]( 1 ) ) ;
			callback( "AA", counts[ 2 ]( 0 ) ) ;
			callback( "AB", counts[ 2 ]( 1 ) ) ;
			callback( "BB", counts[ 2 ]( 2 ) ) ;
			callback( "NULL", null_counts[ 'm' ] + null_counts[ 'f' ] ) ;
			callback( "unknown_ploidy", counts[ -1 ].sum() + null_counts[ -1 ] ) ;
			assert( counts[ 1 ]( 2 ) == 0 ) ;
		}

		std::string get_summary( std::string const& prefix = "", std::size_t column_width = 20 ) const { return prefix + "MissingnessComputation" ; }
	private:
		double const m_call_threshhold ;
	} ;

	struct InfoComputation: public SNPSummaryComputation {
		void operator()(
			VariantIdentifyingData const& snp,
			Genotypes const& genotypes,
			Ploidy const& ploidy,
			genfile::VariantDataReader&,
			ResultCallback callback
		) {
			if( snp.number_of_alleles() == 2 ) {
				bool const allDiploid = ( ploidy.array() == 2 ).cast< int >().sum() == ploidy.size() ;
				if( allDiploid ) {
					compute_autosomal_info( snp, genotypes, callback ) ;
				} else {
					compute_sex_chromosome_info( snp, genotypes, ploidy, callback ) ;
				}
			} else {
				// we don't compute for multiallelics currently
			}
		}

		void compute_autosomal_info( VariantIdentifyingData const& snp, Genotypes const& genotypes, ResultCallback callback ) {
			double const theta_mle = ( genotypes.col( 1 ).sum() + 2.0 * genotypes.col( 2 ).sum() ) / ( 2.0 * genotypes.sum() ) ;
			double const theta_est = theta_mle ;

			Eigen::VectorXd const impute_fallback_distribution = Eigen::VectorXd::Zero( 3 ) ;
			Eigen::VectorXd fallback_distribution = Eigen::VectorXd::Zero( 3 ) ;
			fallback_distribution( 0 ) = ( 1 - theta_est ) * ( 1 - theta_est ) ;
			fallback_distribution( 1 ) = 2.0 * theta_est * ( 1 - theta_est ) ;
			fallback_distribution( 2 ) = theta_est * theta_est ;

			//std::cerr << "theta = " << theta_est << ", fallback_distribution = " << fallback_distribution.transpose() << ".\n" ;

			Eigen::VectorXd const levels = Eigen::VectorXd::LinSpaced( 3, 0, 2 ) ;

			double const info = 1.0 - (
				compute_sum_of_variances( levels, genotypes, fallback_distribution )
				/ ( genotypes.rows() * 2.0 * theta_est * ( 1 - theta_est ) )
			) ;

			double const impute_info = 1.0 - (
				compute_sum_of_variances( levels, genotypes, impute_fallback_distribution )
				/ ( genotypes.sum() * 2.0 * theta_mle * ( 1 - theta_mle ) )
			) ;

			callback( "info", info ) ;
			callback( "impute_info", impute_info ) ;
		}

		void compute_sex_chromosome_info(
			VariantIdentifyingData const& snp,
			Genotypes const& genotypes,
			Ploidy const& ploidy,
			ResultCallback callback
		) {
			Eigen::MatrixXd hap_or_diploid( genotypes.rows(), 2 ) ;
			for( std::size_t i = 0; i < ploidy.size(); ++i ) {
				if( ploidy(i) == 1 ) {
					hap_or_diploid(i,0) = 1 ;
					hap_or_diploid(i,1) = 0 ;
				}
				else if( ploidy[ i ] == 2 ) {
					hap_or_diploid(i,0) = 0 ;
					hap_or_diploid(i,1) = 1 ;
				}
				else {
					// Don't understand, so treat ploidy as missing.
					// individuals with missing sex do not contribute to the computation.
					// I think this is the least surprising thing to do.
					hap_or_diploid(i,0) = 0 ;
					hap_or_diploid(i,1) = 0 ;
				}
			}

			double const a_allele_count_diploid = (
					( genotypes.col( 1 ) + 2.0 * genotypes.col( 0 ) ).array() * ( hap_or_diploid.col(1).array() )
				).sum() ;
			double const b_allele_count_diploid = (
					( genotypes.col( 1 ) + 2.0 * genotypes.col( 2 ) ).array() * ( hap_or_diploid.col(1).array() )
				).sum() ;
			double const a_allele_count_haploid = 	(
					genotypes.col( 0 ).array() * hap_or_diploid.col(0).array()
				).sum() ;
			double const b_allele_count_haploid = 	(
					genotypes.col( 1 ).array() * hap_or_diploid.col(0).array()
				).sum() ;

			// MLE estimate of allele frequency
			double const theta_mle = ( b_allele_count_diploid + b_allele_count_haploid )
				/ ( a_allele_count_diploid + b_allele_count_diploid + a_allele_count_haploid + b_allele_count_haploid ) ;

			// For the new info measure, could regularise by data augmentation, adding one allele of each type.
			// This makes info better behaved at low frequencies.
			// Note adding one of each allele constitutes minimal prior information that the variant is polymorphic.
			// double const theta_est = ( b_allele_count_diploid + b_allele_count_haploid + 1 )
			//	/ ( a_allele_count_diploid + b_allele_count_diploid + a_allele_count_haploid + b_allele_count_haploid + 2 ) ;

			// ...but don't actually do that; it is less conservative at rare SNPs.
			double const theta_est = theta_mle ;

			Eigen::VectorXd diploid_fallback_distribution = Eigen::VectorXd::Zero( 3 ) ;
			diploid_fallback_distribution( 0 ) = ( 1 - theta_est ) * ( 1 - theta_est ) ;
			diploid_fallback_distribution( 1 ) = 2.0 * theta_est * ( 1 - theta_est ) ;
			diploid_fallback_distribution( 2 ) = theta_est * theta_est ;

			Eigen::VectorXd haploid_fallback_distribution = Eigen::VectorXd::Zero( 3 ) ;
			haploid_fallback_distribution( 0 ) = 1 - theta_est ;
			haploid_fallback_distribution( 1 ) = theta_est ;

			//std::cerr << "theta = " << theta_mle << ", fallback_distribution = " << fallback_distribution.transpose() << ".\n" ;

			Eigen::VectorXd const levels = Eigen::VectorXd::LinSpaced( 3, 0, 2 ) ;

			double info = 0.0 ;
			Eigen::VectorXd const haploids = ( hap_or_diploid.col( 0 ).array() == 1 ).cast< double >() ;
			Eigen::VectorXd const diploids = ( hap_or_diploid.col( 1 ).array() == 1 ).cast< double >() ;
			{
				if( haploids.sum() > 0 ) {
					info -= compute_sum_of_variances( levels, genotypes, haploid_fallback_distribution, haploids )
						/ ( theta_est * ( 1 - theta_est ) ) ;
				}

				if( diploids.sum() > 0 ) {
					info -= compute_sum_of_variances( levels, genotypes, diploid_fallback_distribution, diploids )
						/ ( 2.0 * theta_est * ( 1 - theta_est ) ) ;
				}

				double const totalProb = (
					( genotypes.block( 0, 0, genotypes.rows(), 2 ).rowwise().sum().array() * haploids.array() ).sum()
					+ ( genotypes.rowwise().sum().array() * diploids.array() ).sum()
				) ;

				info = 1.0 + (info / (haploids.sum() + diploids.sum() )) ;
			}

			{
				// IMPUTE-style MLE, treat all samples as diploid ignoring gender.
				// Haploids are encoded in the 0/1 columns here so we have to reverse-engineer the
				// original info computation to allow for this.
				double const autosomal_theta_mle = (
					( 2.0 * genotypes.col( 1 ).array() * haploids.array() )
					+ (( genotypes.col( 1 ) + 2.0 * genotypes.col( 2 ) ).array() * ( 1.0 - haploids.array() ))
				).sum() / ( 2.0 * genotypes.sum() ) ;

#if DEBUG
				std::cerr << "autosomal_theta_mle  = " << std::setprecision(10) << autosomal_theta_mle << ".\n" ;
				std::cerr << genotypes << ".\n" ;
#endif
				Eigen::VectorXd haploidLevels = Eigen::VectorXd::Zero(3) ;
				haploidLevels(1) = 2 ;
				double const impute_info = 1.0 - (
					(
						compute_sum_of_variances( haploidLevels, genotypes, Eigen::VectorXd::Zero( 3 ), haploids )
						+ compute_sum_of_variances( levels, genotypes, Eigen::VectorXd::Zero( 3 ), 1.0 - haploids.array() )
					) / ( genotypes.sum() * 2.0 * autosomal_theta_mle * ( 1 - autosomal_theta_mle ) )
				) ;

				callback( "info", info ) ;
				callback( "impute_info", impute_info ) ;
			}
		}

		std::string get_summary( std::string const& prefix = "", std::size_t column_width = 20 ) const { return prefix + "InfoComputation" ; }

	private:
		double compute_sum_of_variances( Eigen::VectorXd const& levels, Eigen::MatrixXd const& probabilities, Eigen::VectorXd const& fallback ) const {
			return compute_sum_of_variances( levels, probabilities, fallback, Eigen::VectorXd::Constant( probabilities.rows(), 1 ) ) ;
		}

		// treat the rows of the probabilities matrix as probabilities.
		// distribution for individual i is taken as a mixture of the distribution given by row i of probabilities,
		// and the fallback distribution if the sum of row i < 1, appropriately weighted.
		// Only individuals with inclusion = 1 are used.
		double compute_sum_of_variances( Eigen::VectorXd const& levels, Eigen::MatrixXd const& probabilities, Eigen::VectorXd const& fallback, Eigen::VectorXd const& inclusion ) const {
			assert( levels.size() == fallback.size() ) ;
			assert( levels.size() == probabilities.cols() ) ;
			assert( inclusion.size() == probabilities.rows() ) ;
			Eigen::VectorXd levels_squared = ( levels.array() * levels.array() ) ;

			double result = 0.0 ;
			for( int i = 0; i < probabilities.rows(); ++i ) {
				if( inclusion( i ) == 1 ) {
					double const c = probabilities.row( i ).sum() ;
					result += compute_variance( levels, levels_squared, probabilities.row( i ).transpose() + ( 1 - c ) * fallback ) ;
				}
			}
			return result ;
		}

		double compute_variance( Eigen::VectorXd const& levels, Eigen::VectorXd const& levels_squared, Eigen::VectorXd const& probs ) const {
			double const mean = ( probs.transpose() * levels )(0) ;
			double const variance = ( probs.transpose() * levels_squared )(0) - ( mean * mean ) ;
			// std::cerr << "compute_variance: levels = " << levels.transpose() << ", levels_squared = " << levels_squared.transpose() << ", probs = " << probs.transpose() << ", variance = " << variance << ".\n" ;
			return variance ;
		}
	} ;
}

SNPSummaryComputation::UniquePtr SNPSummaryComputation::create(
	std::string const& name
) {
	UniquePtr result ;
	if( name == "allele-frequencies" ) { result.reset( new snp_summary_component::AlleleFrequencyComputation( "everything" )) ; }
	else if( name == "allele-counts" ) { result.reset( new snp_summary_component::AlleleFrequencyComputation( "counts" )) ; }
	else if( name == "HWE" ) { result.reset( new snp_summary_component::HWEComputation()) ; }
	else if( name == "missingness" ) { result.reset( new snp_summary_component::MissingnessComputation()) ; }
	else if( name == "info" ) { result.reset( new snp_summary_component::InfoComputation()) ; }
	else if( name == "call-mass-proportion" ) { result.reset( new snp_summary_component::CallMassComputation()) ; }
	else if( name == "intensity-stats" ) { result.reset( new snp_summary_component::IntensitySummaryComputation() ) ; }
	else if( name == "multi-allele-counts" ) { result.reset( new snp_summary_component::AlleleCountComputation() ) ; }
	else {
		throw genfile::BadArgumentError( "SNPSummaryComputation::create()", "name=\"" + name + "\"" ) ;
	}
	return result ;
}


/*****************************************************************************
*
*     Program: EPTlib
*     Author: Alessandro Arduino <a.arduino@inrim.it>
*
*  MIT License
*
*  Copyright (c) 2020-2021  Alessandro Arduino
*  Istituto Nazionale di Ricerca Metrologica (INRiM)
*  Strada delle cacce 91, 10135 Torino
*  ITALY
*
*  Permission is hereby granted, free of charge, to any person obtaining a copy
*  of this software and associated documentation files (the "Software"), to deal
*  in the Software without restriction, including without limitation the rights
*  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
*  copies of the Software, and to permit persons to whom the Software is
*  furnished to do so, subject to the following conditions:
*
*  The above copyright notice and this permission notice shall be included in all
*  copies or substantial portions of the Software.
*
*  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
*  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
*  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
*  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
*  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
*  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
*  SOFTWARE.
*
*****************************************************************************/

#include "eptlib/ept_convreact.h"

#include <complex>
#include <limits>

#include <iostream>
#include <fstream>

#include <Eigen/Sparse>

using namespace eptlib;

namespace {
	static double nand = std::numeric_limits<double>::quiet_NaN();
	static std::complex<double> nancd = std::complex<double>(nand,nand);
}

// EPTConvReact constructor
EPTConvReact::
EPTConvReact(const double freq, const std::array<int,NDIM> &nn,
    const std::array<double,NDIM> &dd, const Shape &shape) :
    EPTInterface(freq,nn,dd, 1,1), dir_epsr_(1.0), dir_sigma_(0.0),
    plane_idx_(nn[2]/2), is_volume_(false), diff_coeff_(0.0),
    thereis_diff_(false), fd_filter_(shape),
    solver_iterations_(0), solver_residual_(0.0) {
    return;
}

// EPTConvReact destructor
EPTConvReact::
~EPTConvReact() {
    return;
}

// EPTConvReact run
EPTlibError EPTConvReact::
Run() {
    if (thereis_tx_sens_.all()) {
        thereis_epsr_ = true;
        if (is_volume_) {
            epsr_ = Image<double>(nn_[0],nn_[1],nn_[2]);
        } else {
            epsr_ = Image<double>(nn_[0],nn_[1]);
        }
    }
    if (thereis_trx_phase_.all()) {
        thereis_sigma_ = true;
        if (is_volume_) {
            sigma_ = Image<double>(nn_[0],nn_[1],nn_[2]);
        } else {
            sigma_ = Image<double>(nn_[0],nn_[1]);
        }
    } else {
        return EPTlibError::MissingData;
    }
    if (thereis_epsr_) {
        // ...complete convection-reaction EPT
        return CompleteEPTConvReact();
    } else {
        // ...phase-based convection-reaction EPT
        return PhaseEPTConvReact();
    }
    return EPTlibError::Success;
}

// EPTConvReact set volume tomography
bool EPTConvReact::
Toggle3D() {
    is_volume_ = !is_volume_;
    return is_volume_;
}

// EPTConvReact set the selected plane index
EPTlibError EPTConvReact::
SelectSlice(const int slice_idx) {
    int r2 = fd_filter_.GetShape().GetSize()[2]/2;
    if (slice_idx<r2||slice_idx>nn_[2]-1-r2) {
        return EPTlibError::OutOfRange;
    }
    plane_idx_ = slice_idx;
    return EPTlibError::Success;
}

// EPTConvReact set the dirichlet boundary condition
EPTlibError EPTConvReact::
SetDirichlet(const double dir_epsr, const double dir_sigma) {
    if (dir_epsr<1.0||dir_sigma<0.0) {
        return EPTlibError::WrongDataFormat;
    }
    dir_epsr_ = dir_epsr;
    dir_sigma_ = dir_sigma;
    return EPTlibError::Success;
}

// EPTConvReact set artificial diffusion
void EPTConvReact::
SetArtificialDiffusion(const double diff_coeff) {
    diff_coeff_ = diff_coeff;
    thereis_diff_ = true;
    return;
}
// EPTConvReact unset artificial diffusion
void EPTConvReact::
UnsetArtificialDiffusion() {
    thereis_diff_ = false;
    return;
}

// EPTConvReact get the number of iterations
int EPTConvReact::
GetSolverIterations() {
    return solver_iterations_;
}
// EPTConvReact get estimated error
double EPTConvReact::
GetSolverResidual() {
    return solver_residual_;
}

namespace { // details
    template <typename T>
    void FillDoF(std::vector<int> *dof, int *idx_dof, int *n_dof, int *n_dop,
        const int n_dim, const int i2, const std::array<int,NDIM> &step,
        const std::vector<T> &beta) {
        int idx = step[2]*i2;
        for (int idx_out = 0; idx_out<step[2]; ++idx_out) {
            if (beta[idx]==beta[idx]) {
                // if beta is not a NaN, it could be a DoF...
                (*dof)[(*idx_dof)++] = ++(*n_dof);
                for (int d = 0; d<n_dim; ++d) {
                    if (beta[idx+step[d]]!=beta[idx+step[d]] || beta[idx-step[d]]!=beta[idx-step[d]]) {
                        // ...unless it is near to a NaN
                        (*dof)[*idx_dof-1] = --(*n_dop);
                        --(*n_dof);
                        break;
                    }
                }
            } else {
                (*dof)[(*idx_dof)++] = --(*n_dop);
            }
            ++idx;
        }
        return;
    }
}  // namespace

// EPTConvReact complete EPT
EPTlibError EPTConvReact::
CompleteEPTConvReact() {
    const std::array<int,NDIM> step{1,nn_[0],nn_[0]*nn_[1]};
    int n_dim = is_volume_?NDIM:NDIM-1;
    int n_out = is_volume_?n_vox_:step[2];
    std::complex<double> dir_x = 1.0/std::complex<double>(dir_epsr_*EPS0,-dir_sigma_/omega_);
    // compute the gradient
    std::vector<std::complex<double> > tx_sens_c(n_vox_);
    std::array<std::vector<std::complex<double> >,NDIM> beta;
    for (int idx = 0; idx<n_vox_; ++idx) {
        tx_sens_c[idx] = (*tx_sens_[0])[idx]*std::exp(std::complex<double>(0.0,0.5*(*trx_phase_[0])[idx]));
    }
    for (int d = 0; d<n_dim; ++d) {
        beta[d].resize(n_vox_,std::numeric_limits<double>::quiet_NaN());
        DifferentialOperator diff_op = static_cast<DifferentialOperator>(d);
        EPTlibError error = fd_filter_.Apply(diff_op,beta[d].data(),tx_sens_c.data(),nn_,dd_);
        if (error!=EPTlibError::Success) {
            return error;
        }
    }
    for (int idx = 0; idx<n_vox_; ++idx) {
        beta[0][idx] = beta[0][idx] - std::complex<double>(0.0,1.0)*beta[1][idx];
        beta[1][idx] = std::complex<double>(0.0,1.0)*beta[0][idx];
    }
    if (!is_volume_) {
        beta[2].resize(n_vox_,std::numeric_limits<double>::quiet_NaN());
        DifferentialOperator diff_op = DifferentialOperator::GradientZZ;
        EPTlibError error = fd_filter_.Apply(diff_op,beta[2].data(),tx_sens_c.data(),nn_,dd_);
        if (error!=EPTlibError::Success) {
            return error;
        }
    }
    // select the degrees of freedom
    std::vector<int> dof(n_out);
    int idx_dof = 0;
    int n_dof = 0;
    int n_dop = 0;
    if (is_volume_) {
        for (int i2 = 0; i2<nn_[2]; ++i2) {
            ::FillDoF(&dof,&idx_dof,&n_dof,&n_dop, n_dim,i2,step,beta[0]);
        }
    } else {
        ::FillDoF(&dof,&idx_dof,&n_dof,&n_dop, n_dim,plane_idx_,step,beta[0]);
    }
    // build coefficient matrix and forcing term
    Eigen::SparseMatrix<std::complex<double> > A(n_dof,n_dof);
    std::vector<Eigen::Triplet<std::complex<double> > > A_trip(0);
    Eigen::VectorXcd b(n_dof);
    for (int idx_out = 0; idx_out<n_out; ++idx_out) {
        int idx = is_volume_?idx_out:idx_out+step[2]*plane_idx_;
        int idof = dof[idx_out];
        if (idof > 0) {
            b[idof-1] = 0.0;
            // coefficient matrix
            for (int d = 0; d<n_dim; ++d) {
                if (thereis_diff_) {
                    A_trip.push_back(Eigen::Triplet<std::complex<double> >(idof-1,idof-1,2.0*diff_coeff_/dd_[d]/dd_[d]));
                }
                int jdx_out = idx_out+step[d];
                int jdx = idx+step[d];
                int jdof = dof[jdx_out];
                if (jdof > 0) {
                    A_trip.push_back(Eigen::Triplet<std::complex<double> >(idof-1,jdof-1,beta[d][jdx]/2.0/dd_[d]));
                    if (thereis_diff_) {
                        A_trip.push_back(Eigen::Triplet<std::complex<double> >(idof-1,jdof-1,-diff_coeff_/dd_[d]/dd_[d]));
                    }
                } else {
                    b[idof-1] += -dir_x*beta[d][jdx]/2.0/dd_[d];
                    if (thereis_diff_) {
                        b[idof-1] += dir_x*diff_coeff_/dd_[d]/dd_[d];
                    }
                }
                jdx_out = idx_out-step[d];
                jdx = idx-step[d];
                jdof = dof[jdx_out];
                if (jdof > 0) {
                    A_trip.push_back(Eigen::Triplet<std::complex<double> >(idof-1,jdof-1,-beta[d][jdx]/2.0/dd_[d]));
                    if (thereis_diff_) {
                        A_trip.push_back(Eigen::Triplet<std::complex<double> >(idof-1,jdof-1,-diff_coeff_/dd_[d]/dd_[d]));
                    }
                } else {
                    b[idof-1] += dir_x*beta[d][jdx]/2.0/dd_[d];
                    if (thereis_diff_) {
                        b[idof-1] += dir_x*diff_coeff_/dd_[d]/dd_[d];
                    }
                }
            }
            if (!is_volume_) {
                A_trip.push_back(Eigen::Triplet<std::complex<double> >(idof-1,idof-1,beta[2][idx]));
            }
            // forcing term
            b[idof-1] += -omega_*omega_*MU0*tx_sens_c[idx];
        }
    }
    A.setFromTriplets(A_trip.begin(),A_trip.end());
    A.makeCompressed();
    // Solve the linear system
    Eigen::BiCGSTAB<Eigen::SparseMatrix<std::complex<double> > > solver;
    solver.compute(A);
    Eigen::VectorXcd x = solver.solve(b);
    solver_iterations_ = solver.iterations();
    solver_residual_ = solver.error();
    // Extract the electric properties from the result
    for (int idx = 0; idx<n_out; ++idx) {
        int idof = dof[idx];
        if (idof > 0) {
            std::complex<double> epsc = 1.0/x[idof-1];
            epsr_[idx] = epsc.real()/EPS0;
            sigma_[idx] = -epsc.imag()*omega_;
        } else {
            epsr_[idx] = ::nand;
            sigma_[idx] = ::nand;
        }
    }
    return EPTlibError::Success;
}

// EPTConvReact phase-based EPT
EPTlibError EPTConvReact::
PhaseEPTConvReact() {
    const std::array<int,NDIM> step{1,nn_[0],nn_[0]*nn_[1]};
    int n_dim = is_volume_?NDIM:NDIM-1;
    int n_out = is_volume_?n_vox_:step[2];
    double dir_x = 1.0/dir_sigma_;
    // compute the gradient
    std::array<std::vector<double>,NDIM> beta;
    for (int d = 0; d<n_dim; ++d) {
        beta[d].resize(n_vox_,std::numeric_limits<double>::quiet_NaN());
        DifferentialOperator diff_op = static_cast<DifferentialOperator>(d);
        EPTlibError error;
        if (PhaseIsWrapped()) {
            error = WrappedPhaseDerivative(diff_op,beta[d].data(),trx_phase_[0]->GetData().data(),nn_,dd_,fd_filter_);
        } else {
            error = fd_filter_.Apply(diff_op,beta[d].data(),trx_phase_[0]->GetData().data(),nn_,dd_);
        }
        if (error!=EPTlibError::Success) {
            return error;
        }
    }
    if (!is_volume_) {
        beta[2].resize(n_vox_,std::numeric_limits<double>::quiet_NaN());
        DifferentialOperator diff_op = DifferentialOperator::GradientZZ;
        EPTlibError error;
        if (PhaseIsWrapped()) {
            error = WrappedPhaseDerivative(diff_op,beta[2].data(),trx_phase_[0]->GetData().data(),nn_,dd_,fd_filter_);
        } else {
            error = fd_filter_.Apply(diff_op,beta[2].data(),trx_phase_[0]->GetData().data(),nn_,dd_);
        }
        if (error!=EPTlibError::Success) {
            return error;
        }
    }
    // select the degrees of freedom
    std::vector<int> dof(n_out);
    int idx_dof = 0;
    int n_dof = 0;
    int n_dop = 0;
    if (is_volume_) {
        for (int i2 = 0; i2<nn_[2]; ++i2) {
            ::FillDoF(&dof,&idx_dof,&n_dof,&n_dop, n_dim,i2,step,beta[0]);
        }
    } else {
        ::FillDoF(&dof,&idx_dof,&n_dof,&n_dop, n_dim,plane_idx_,step,beta[0]);
    }
    // build coefficient matrix and forcing term
    Eigen::SparseMatrix<double> A(n_dof,n_dof);
    std::vector<Eigen::Triplet<double> > A_trip(0);
    Eigen::VectorXd b(n_dof);
    for (int idx_out = 0; idx_out<n_out; ++idx_out) {
        int idx = is_volume_?idx_out:idx_out+step[2]*plane_idx_;
        int idof = dof[idx_out];
        if (idof > 0) {
            b[idof-1] = 0.0;
            // coefficient matrix
            for (int d = 0; d<n_dim; ++d) {
                // diffusion (central term)
                if (thereis_diff_) {
                    A_trip.push_back(Eigen::Triplet<double>(idof-1,idof-1,2.0*diff_coeff_/dd_[d]/dd_[d]));
                }
                std::array<int,2> sides{+1,-1};
                for (int s = 0; s<2; ++s) {
                    int jdx_out = idx_out+sides[s]*step[d];
                    int jdx = idx+sides[s]*step[d];
                    int jdof = dof[jdx_out];
                    // convection
                    if (beta[d][idx]*beta[d][jdx]>0) {
                        if (sides[s]*beta[d][idx]>0) {
                            double A_tmp = sides[s]*beta[d][idx]/dd_[d];
                            A_trip.push_back(Eigen::Triplet<double>(idof-1,idof-1,A_tmp));
                        } else {
                            double A_tmp = sides[s]*beta[d][jdx]/dd_[d];
                            if (jdof > 0) {
                                A_trip.push_back(Eigen::Triplet<double>(idof-1,jdof-1,A_tmp));
                            } else {
                                b[idof-1] += -dir_x*A_tmp;
                            }
                        }
                    } else {
                        double A_tmp = sides[s]*beta[d][idx]/dd_[d]/2.0;
                        A_trip.push_back(Eigen::Triplet<double>(idof-1,idof-1,A_tmp));
                        A_tmp = sides[s]*beta[d][jdx]/dd_[d]/2.0;
                        if (jdof > 0) {
                            A_trip.push_back(Eigen::Triplet<double>(idof-1,jdof-1,A_tmp));
                        } else {
                            b[idof-1] += -dir_x*A_tmp;
                        }
                    }
                    // diffusion
                    if (thereis_diff_) {
                        if (jdof > 0) {
                            A_trip.push_back(Eigen::Triplet<double>(idof-1,jdof-1,-diff_coeff_/dd_[d]/dd_[d]));
                        } else {
                            b[idof-1] += dir_x*diff_coeff_/dd_[d]/dd_[d];
                        }
                    }
                }
            }
            if (!is_volume_) {
                A_trip.push_back(Eigen::Triplet<double>(idof-1,idof-1,beta[2][idx]));
            }
            // forcing term
            b[idof-1] += 2*omega_*MU0;
        }
    }
    A.setFromTriplets(A_trip.begin(),A_trip.end());
    A.makeCompressed();
    // Solve the linear system
    Eigen::BiCGSTAB<Eigen::SparseMatrix<double> > solver;
    solver.compute(A);
    Eigen::VectorXd x = solver.solve(b);
    solver_iterations_ = solver.iterations();
    solver_residual_ = solver.error();
    // Extract the electric properties from the result
    for (int idx = 0; idx<n_out; ++idx) {
        int idof = dof[idx];
        if (idof > 0) {
            sigma_[idx] = 1.0/x[idof-1];
        } else {
            sigma_[idx] = ::nand;
        }
    }
    return EPTlibError::Success;
}


/*
 Copyright (C) 2016 Quaternion Risk Management Ltd
 All rights reserved.

 This file is part of ORE, a free-software/open-source library
 for transparent pricing and risk analysis - http://opensourcerisk.org

 ORE is free software: you can redistribute it and/or modify it
 under the terms of the Modified BSD License.  You should have received a
 copy of the license along with this program.
 The license is also available online at <http://opensourcerisk.org>

 This program is distributed on the basis that it will form a useful
 contribution to risk analytics and model standardisation, but WITHOUT
 ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 FITNESS FOR A PARTICULAR PURPOSE. See the license for more details.
*/

#include <qle/termstructures/optionletstripper2.hpp>
#include <qle/termstructures/spreadedoptionletvolatility.hpp>

#include <ql/instruments/makecapfloor.hpp>
#include <ql/math/solvers1d/brent.hpp>
#include <ql/pricingengines/capfloor/bacheliercapfloorengine.hpp>
#include <ql/pricingengines/capfloor/blackcapfloorengine.hpp>
#include <ql/termstructures/volatility/optionlet/strippedoptionletadapter.hpp>

#include <boost/make_shared.hpp>

namespace QuantExt {

OptionletStripper2::OptionletStripper2(const boost::shared_ptr<QuantExt::OptionletStripper>& optionletStripper,
                                       const Handle<QuantLib::CapFloorTermVolCurve>& atmCapFloorTermVolCurve,
                                       const Handle<YieldTermStructure>& discount, const VolatilityType type,
                                       const Real displacement)
    : OptionletStripper(optionletStripper->termVolSurface(), optionletStripper->iborIndex(), discount,
                        optionletStripper->volatilityType(), optionletStripper->displacement()),
      stripper_(optionletStripper), atmCapFloorTermVolCurve_(atmCapFloorTermVolCurve),
      dc_(stripper_->termVolSurface()->dayCounter()), nOptionExpiries_(atmCapFloorTermVolCurve->optionTenors().size()),
      atmCapFloorStrikes_(nOptionExpiries_), atmCapFloorPrices_(nOptionExpiries_), spreadsVolImplied_(nOptionExpiries_),
      caps_(nOptionExpiries_), maxEvaluations_(10000), accuracy_(1.e-6), inputVolatilityType_(type),
      inputDisplacement_(displacement) {

    registerWith(stripper_);
    registerWith(atmCapFloorTermVolCurve_);

    QL_REQUIRE(dc_ == atmCapFloorTermVolCurve->dayCounter(), "different day counters provided");
}

void OptionletStripper2::performCalculations() const {

    // optionletStripper data
    optionletDates_ = stripper_->optionletFixingDates();
    optionletPaymentDates_ = stripper_->optionletPaymentDates();
    optionletAccrualPeriods_ = stripper_->optionletAccrualPeriods();
    optionletTimes_ = stripper_->optionletFixingTimes();
    atmOptionletRate_ = stripper_->atmOptionletRates();
    for (Size i = 0; i < optionletTimes_.size(); ++i) {
        optionletStrikes_[i] = stripper_->optionletStrikes(i);
        optionletVolatilities_[i] = stripper_->optionletVolatilities(i);
    }

    // atmCapFloorTermVolCurve data
    const vector<Period>& optionExpiriesTenors = atmCapFloorTermVolCurve_->optionTenors();
    const vector<Time>& optionExpiriesTimes = atmCapFloorTermVolCurve_->optionTimes();

    // discount curve
    const Handle<YieldTermStructure>& discountCurve =
        discount_.empty() ? iborIndex_->forwardingTermStructure() : discount_;

    for (Size j = 0; j < nOptionExpiries_; ++j) {
        // Dummy strike, doesn't get used for ATM curve
        Volatility atmOptionVol = atmCapFloorTermVolCurve_->volatility(optionExpiriesTimes[j], 33.3333);

        // Create a cap for each pillar point on ATM curve and attach relevant pricing engine i.e. Black if
        // quotes are shifted lognormal and Bachelier if quotes are normal
        boost::shared_ptr<PricingEngine> engine;
        if (inputVolatilityType_ == ShiftedLognormal) {
            engine = boost::make_shared<BlackCapFloorEngine>(discountCurve, atmOptionVol, dc_, inputDisplacement_);
        } else if (inputVolatilityType_ == Normal) {
            engine = boost::make_shared<BachelierCapFloorEngine>(discountCurve, atmOptionVol, dc_);
        } else {
            QL_FAIL("unknown volatility type: " << volatilityType_);
        }

        // Using Null<Rate>() as strike => strike will be set to ATM rate. However, to calculate ATM rate, QL requires
        // a BlackCapFloorEngine to be set (not a BachelierCapFloorEngine)! So, need a temp BlackCapFloorEngine with a
        // dummy vol to calculate ATM rate. Needs to be fixed in QL.
        boost::shared_ptr<PricingEngine> tempEngine = boost::make_shared<BlackCapFloorEngine>(discountCurve, 0.01);
        caps_[j] = MakeCapFloor(CapFloor::Cap, optionExpiriesTenors[j], iborIndex_, Null<Rate>(), 0 * Days)
                       .withPricingEngine(tempEngine);

        // Now set correct engine and get the ATM rate and the price
        caps_[j]->setPricingEngine(engine);
        atmCapFloorStrikes_[j] = caps_[j]->atmRate(**discountCurve);
        atmCapFloorPrices_[j] = caps_[j]->NPV();
    }

    spreadsVolImplied_ = spreadsVolImplied(discountCurve);

    StrippedOptionletAdapter adapter(stripper_);
    adapter.enableExtrapolation();

    Volatility unadjustedVol, adjustedVol;
    for (Size j = 0; j < nOptionExpiries_; ++j) {
        for (Size i = 0; i < optionletVolatilities_.size(); ++i) {
            if (i <= caps_[j]->floatingLeg().size()) {
                unadjustedVol = adapter.volatility(optionletTimes_[i], atmCapFloorStrikes_[j]);
                adjustedVol = unadjustedVol + spreadsVolImplied_[j];

                // insert adjusted volatility
                vector<Rate>::const_iterator previous =
                    lower_bound(optionletStrikes_[i].begin(), optionletStrikes_[i].end(), atmCapFloorStrikes_[j]);
                Size insertIndex = previous - optionletStrikes_[i].begin();

                optionletStrikes_[i].insert(optionletStrikes_[i].begin() + insertIndex, atmCapFloorStrikes_[j]);
                optionletVolatilities_[i].insert(optionletVolatilities_[i].begin() + insertIndex, adjustedVol);
            }
        }
    }
}

vector<Volatility> OptionletStripper2::spreadsVolImplied(const Handle<YieldTermStructure>& discount) const {

    Brent solver;
    vector<Volatility> result(nOptionExpiries_);
    Volatility guess = 0.0001, minSpread = -0.1, maxSpread = 0.1;
    for (Size j = 0; j < nOptionExpiries_; ++j) {
        ObjectiveFunction f(stripper_, caps_[j], atmCapFloorPrices_[j], discount);
        solver.setMaxEvaluations(maxEvaluations_);
        Volatility root = solver.solve(f, accuracy_, guess, minSpread, maxSpread);
        result[j] = root;
    }
    return result;
}

vector<Volatility> OptionletStripper2::spreadsVol() const {
    calculate();
    return spreadsVolImplied_;
}

vector<Rate> OptionletStripper2::atmCapFloorStrikes() const {
    calculate();
    return atmCapFloorStrikes_;
}

vector<Real> OptionletStripper2::atmCapFloorPrices() const {
    calculate();
    return atmCapFloorPrices_;
}

// OptionletStripper2::ObjectiveFunction
OptionletStripper2::ObjectiveFunction::ObjectiveFunction(
    const boost::shared_ptr<QuantExt::OptionletStripper>& optionletStripper, const boost::shared_ptr<CapFloor>& cap,
    Real targetValue, const Handle<YieldTermStructure>& discount)
    : cap_(cap), targetValue_(targetValue), discount_(discount) {
    boost::shared_ptr<OptionletVolatilityStructure> adapter(new StrippedOptionletAdapter(optionletStripper));
    adapter->enableExtrapolation();

    // set an implausible value, so that calculation is forced
    // at first operator()(Volatility x) call
    spreadQuote_ = boost::shared_ptr<SimpleQuote>(new SimpleQuote(-1.0));

    boost::shared_ptr<OptionletVolatilityStructure> spreadedAdapter(
        new SpreadedOptionletVolatility(Handle<OptionletVolatilityStructure>(adapter), Handle<Quote>(spreadQuote_)));

    // Use the same volatility type as optionletStripper
    // Anything else would not make sense
    boost::shared_ptr<PricingEngine> engine;
    if (optionletStripper->volatilityType() == ShiftedLognormal) {
        engine = boost::make_shared<BlackCapFloorEngine>(
            discount_, Handle<OptionletVolatilityStructure>(spreadedAdapter), optionletStripper->displacement());
    } else if (optionletStripper->volatilityType() == Normal) {
        engine = boost::make_shared<BachelierCapFloorEngine>(discount_,
                                                             Handle<OptionletVolatilityStructure>(spreadedAdapter));
    } else {
        QL_FAIL("Unknown volatility type: " << optionletStripper->volatilityType());
    }

    cap_->setPricingEngine(engine);
}

Real OptionletStripper2::ObjectiveFunction::operator()(Volatility s) const {
    if (s != spreadQuote_->value())
        spreadQuote_->setValue(s);
    return cap_->NPV() - targetValue_;
}
} // namespace QuantExt


//| This file is a part of the sferes2 framework.
//| Copyright 2009, ISIR / Universite Pierre et Marie Curie (UPMC)
//| Main contributor(s): Jean-Baptiste Mouret, mouret@isir.fr
//|
//| This software is a computer program whose purpose is to facilitate
//| experiments in evolutionary computation and evolutionary robotics.
//|
//| This software is governed by the CeCILL license under French law
//| and abiding by the rules of distribution of free software.  You
//| can use, modify and/ or redistribute the software under the terms
//| of the CeCILL license as circulated by CEA, CNRS and INRIA at the
//| following URL "http://www.cecill.info".
//|
//| As a counterpart to the access to the source code and rights to
//| copy, modify and redistribute granted by the license, users are
//| provided only with a limited warranty and the software's author,
//| the holder of the economic rights, and the successive licensors
//| have only limited liability.
//|
//| In this respect, the user's attention is drawn to the risks
//| associated with loading, using, modifying and/or developing or
//| reproducing the software by the user in light of its specific
//| status of free software, that may mean that it is complicated to
//| manipulate, and that also therefore means that it is reserved for
//| developers and experienced professionals having in-depth computer
//| knowledge. Users are therefore encouraged to load and test the
//| software's suitability as regards their requirements in conditions
//| enabling the security of their systems and/or data to be ensured
//| and, more generally, to use and operate it in the same conditions
//| as regards security.
//|
//| The fact that you are presently reading this means that you have
//| had knowledge of the CeCILL license and that you accept its terms.

#ifndef MODIFIER_NOVELTY_HPP
#define MODIFIER_NOVELTY_HPP

#include <Eigen/Core>
#include "sferes/parallel.hpp"
#include <boost/uuid/uuid.hpp>            // uuid class

//#include <boost/accumulators/accumulators.hpp>
//#include <boost/accumulators/statistics/stats.hpp>
//
//// Headers specifics to the computations we need
//#include <boost/accumulators/statistics/max.hpp>
//#include <boost/accumulators/statistics/min.hpp>

namespace sferes {
  namespace modif {
    namespace novelty {

      // compute the matrix of distances
      template<typename Phen>
      struct _distance_f {
        typedef std::vector<boost::shared_ptr<Phen> > pop_t;
        const pop_t& _pop;
        const pop_t& _archive;
        Eigen::MatrixXf& distances;

        ~_distance_f() { }
        _distance_f(const pop_t& pop, const pop_t& archive, Eigen::MatrixXf& d) :
          _pop(pop), _archive(archive), distances(d) {}
        _distance_f(const _distance_f& ev) :
          _pop(ev._pop), _archive(ev._archive), distances(ev.distances) {}

        void operator() (const parallel::range_t& r) const {
          for (size_t i = r.begin(); i != r.end(); ++i) {
            for (size_t j = 0; j < _archive.size(); ++j)
            {
            	float d = 0;

            	// Check if they are not the same image
            	if ((_pop[i] != _archive[j]) && (_pop[i]->id() != _archive[j]->id()))
            	{
#ifdef EUCLIDEAN_DISTANCE
            		// Euclidean distance
            		d = _pop[i]->dist(*_archive[j]);

            		#else
            		// Regular distance
            		d = _pop[i]->fit().dist(*_archive[j]);
#endif
            	}

              distances(i, j) = d;
            }
          }
        }
      };
    }

    // The novelty score will be stored in the last objective 'slot'
    // If there is only one objective (this->objs().size() == 1), then
    // this modifiers is a standard Novelty search algorithm [1]
    // otherwise, it is a "novelty-based multi-objectivization" [2]
    // See [2] for more explanations of the differences between
    // novelty search, multi-objective novelty search, and behavioral diversity
    //
    // * References
    // [1] Lehman, Joel, and Kenneth O. Stanley. "Abandoning objectives:
    // Evolution through the search for novelty alone."
    // Evolutionary computation 19.2 (2011): 189-223.
    //
    // [2] Mouret, Jean-Baptiste. "Novelty-based multiobjectivization."
    // New Horizons in Evolutionary Robotics. Springer Berlin Heidelberg,
    // 2011. 139-154.
    template<typename Phen, typename Params, typename Exact = stc::Itself>
    class Novelty {
     public:
      typedef boost::shared_ptr<Phen> phen_t;
      typedef std::vector<phen_t> pop_t;
      Novelty() : _rho_min(Params::novelty::rho_min_init), _not_added(0) {}

      template<typename Ea>
      void apply(Ea& ea)
      {
//      	boost::accumulators::accumulator_set<double, boost::accumulators::stats<boost::accumulators::tag::max> > acc_max;
//      	boost::accumulators::accumulator_set<double, boost::accumulators::stats<boost::accumulators::tag::min> > acc_min;

        SFERES_CONST size_t k = Params::novelty::k;
        // merge the population and the archive in a single archive
        pop_t archive = _archive;
        archive.insert(archive.end(), ea.pop().begin(), ea.pop().end());

        // we compute all the distances from pop(i) to archive(j) and store them
        Eigen::MatrixXf distances(ea.pop().size(), archive.size());
        novelty::_distance_f<Phen> f(ea.pop(), archive, distances);
        parallel::init();
        parallel::p_for(parallel::range_t(0, ea.pop().size()), f);

        // compute the sparseness of each individual of the population
        // and potentially add to the archive
        int added = 0;

        for (size_t i = 0; i < ea.pop().size(); ++i) {
          size_t nb_objs = ea.pop()[i]->fit().objs().size();
          Eigen::VectorXf vd = distances.row(i);

          // Get the k-nearest neighbors
          double n = 0.0;
          std::partial_sort(vd.data(),
                            vd.data() + k,
                            vd.data() + vd.size());

          // Sum up the total of distances from k-nearest neighbors
          // This is the novelty score
          n = vd.head<k>().sum() / k;

          // Set the novelty score to the last objective in the array of objectives
          ea.pop()[i]->fit().set_obj(nb_objs - 1, n);

//          std::cout << "rho = " << _rho_min << "; n = " << n << "\n";
//          acc_max(n);
//          acc_min(n);

          // Check if this individual is already in the archive
          if (!in_archive_already(ea.pop()[i]->id()) &&
          		(n > _rho_min			// Check if the novelty score passes the threshold
              || misc::rand<float>() < Params::novelty::add_to_archive_prob))
          {
          	// Add this individual to the permanent archive
          	_archive.push_back(ea.pop()[i]);

          	// Print this indiv just added
//            std::cout << "----> n = " << n << "; label = " << ea.pop()[i]->fit().label() << "; score = " << ea.pop()[i]->fit().score() << "; id = " << ea.pop()[i]->id()  <<"\n";
//          	std::cout << "----> n = " << n << "; label = " << "; id = " << ea.pop()[i]->id()  <<"\n";

            _not_added = 0;
            ++added;
          } else {
          	// Do not add this individual to the archive
            ++_not_added;
          }
        } // end for all individuals

        // update rho_min
        if (_not_added > Params::novelty::stalled_tresh)//2500
        {
          _rho_min *= 0.95;	// Lower rho_min
          _not_added = 0;		// After lowering the stalled threshold, we need to reset the _not_added count
        }
        if (_archive.size() > Params::novelty::k && added > Params::novelty::adding_tresh)//4
        {
          _rho_min *= 1.05f;	// Increase rho_min
        }

        std::cout<<"a size:"<<_archive.size()<<std::endl;

        // DEBUG
//        std::cout << "a size:"<<_archive.size()
//        		<< "; max: " << boost::accumulators::max(acc_max)
//        		<< "; min: " << boost::accumulators::min(acc_min)
//        		<< std::endl;

      }

      /*
			 * Check if this individual is already in the archive.
			 */
			const bool in_archive_already ( const boost::uuids::uuid& id )
			{
				/* Comment out to save computation time
				for (int i = 0; i < _archive.size(); ++i)
				{
					// Check if this id is found in the archive.
					if (_archive[i]->id() == id)
					{
						return true;
					}

				}
				*/

				return false;
			}

      /*
       * Getter of archive.
       *
       * This function is created for statistics class to save the archive down to gen file.
       * 2015-01-21
       */
      const pop_t& archive() const
      {
      	return _archive;
      }

      /*
       * Setter of archive.
       *
       * This function is created for statistics class to save the archive down to gen file.
       * 2015-01-21
       */
      void set_archive(const pop_t& archive)
      {
      	_archive.clear();

      	// Assign items to archive
      	_archive.insert(_archive.end(), archive.begin(), archive.end());
      }

     protected:
      pop_t _archive;
      float _rho_min;
      size_t _not_added;
    };
  } // modif
} // sferes

#endif


// This file is part of MLDB. Copyright 2015 mldb.ai inc. All rights reserved.

/* separation_stats.cc
   Jeremy Barnes, 14 June 2011
   Copyright (c) 2011 mldb.ai inc.  All rights reserved.

*/

#include "separation_stats.h"
#include "mldb/arch/exception.h"
#include "mldb/base/exc_assert.h"
#include <boost/utility.hpp>
#include "mldb/vfs/filter_streams.h"

using namespace std;

namespace MLDB {

/*****************************************************************************/
/* BINARY STATS                                                              */
/*****************************************************************************/

double
BinaryStats::
rocAreaSince(const BinaryStats & other) const
{
    double tp1 = other.truePositiveRate(), fp1 = other.falsePositiveRate();
    double tp2 = truePositiveRate(), fp2 = falsePositiveRate();

    double result = (fp2 - fp1) * (tp1 + tp2) * 0.5;

    //cerr << "tp1 = " << tp1 << " tp2 = " << tp2 << " fp1 = " << fp1
    //     << " fp2 = " << fp2 << " area = " << result << endl;

    return result;
}

void
BinaryStats::
add(const BinaryStats & other, double weight)
{
    auto addCounts = [&] (      Array2D & local,
                          const Array2D & other,
                                       float w)

    {
        local[0][0] += w * other[0][0];
        local[0][1] += w * other[0][1];
        local[1][0] += w * other[1][0];
        local[1][1] += w * other[1][1];
    };

    addCounts(counts, other.counts, weight);
    addCounts(unweighted_counts, other.unweighted_counts, 1);

    threshold += weight * other.threshold;
}

Json::Value
BinaryStats::
toJson() const
{
    Json::Value result;
    result["population"]["included"] = includedPopulation();
    result["population"]["excluded"] = excludedPopulation();
    result["pr"]["accuracy"] = accuracy();
    result["pr"]["precision"] = precision();
    result["pr"]["recall"] = recall();
    result["pr"]["f1Score"] = f();
    result["mcc"] = mcc();
    result["gain"] = gain();
    result["counts"]["truePositives"] = truePositives();
    result["counts"]["falsePositives"] = falsePositives();
    result["counts"]["trueNegatives"] = trueNegatives();
    result["counts"]["falseNegatives"] = falseNegatives();
    result["threshold"] = threshold;
    return result;
}

/*****************************************************************************/
/* SCORED STATS                                                              */
/*****************************************************************************/

ScoredStats::
ScoredStats()
    : auc(1.0), isSorted(true)
{
}

BinaryStats
ScoredStats::
atThreshold(float threshold) const
{
    struct FindThreshold {
        bool operator () (const BinaryStats & e1, const BinaryStats & e2) const
        {
            return e1.threshold > e2.threshold;
        }

        bool operator () (const BinaryStats & e1, float e2) const
        {
            return e1.threshold > e2;
        }

        bool operator () (float e1, const BinaryStats & e2) const
        {
            return e1 > e2.threshold;
        }
    };

    if (!std::is_sorted(stats.begin(), stats.end(), FindThreshold()))
        throw Exception("stats not sorted on input");

    if (stats.empty())
        throw Exception("stats is empty");

    // Lower bound means strictly above
    auto lower
        = std::lower_bound(stats.begin(), stats.end(), threshold,
                           FindThreshold());

    if (lower == stats.end())
        return stats.back();

    return *lower;
}

BinaryStats
ScoredStats::
atPercentile(float percentile) const
{
    struct FindPercentile {
        bool operator () (const BinaryStats & e1, const BinaryStats & e2) const
        {
            return e1.proportionOfPopulation() < e2.proportionOfPopulation();
        }

        bool operator () (const BinaryStats & e1, float e2) const
        {
            return e1.proportionOfPopulation() < e2;
        }

        bool operator () (float e1, const BinaryStats & e2) const
        {
            return e1 < e2.proportionOfPopulation();
        }
    };

    for (unsigned i = 1;  i < stats.size();  ++i) {
        if (stats[i -1].proportionOfPopulation() > stats[i].proportionOfPopulation()) {
            cerr << "i = " << i << endl;
            cerr << "prev = " << stats[i - 1].toJson() << endl;
            cerr << "ours = " << stats[i].toJson() << endl;
            throw MLDB::Exception("really not sorted on input");
        }
    }

    if (!std::is_sorted(stats.begin(), stats.end(), FindPercentile()))
        throw Exception("stats not sorted on input");

    if (stats.empty())
        throw Exception("stats is empty");

    // Lower bound means strictly above
    auto it
        = std::lower_bound(stats.begin(), stats.end(), percentile,
                           FindPercentile());

    if (it == stats.begin())
        return *it;

    if (it == stats.end())
        return stats.back();

    BinaryStats upper = *it, lower = *std::prev(it);

    // Do an interpolation
    double atUpper = upper.proportionOfPopulation(),
           atLower = lower.proportionOfPopulation(),
           range = atUpper - atLower;

    ExcAssertLessEqual(percentile, atUpper);
    ExcAssertGreaterEqual(percentile, atLower);
    ExcAssertGreater(range, 0);

    double weight1 = (atUpper - percentile) / range;
    double weight2 = (percentile - atLower) / range;

    //cerr << "atUpper = " << atUpper << " atLower = " << atLower
    //     << " range = " << range << " weight1 = " << weight1
    //     << " weight2 = " << weight2 << endl;

    BinaryStats result;
    result.add(upper, weight2);
    result.add(lower, weight1);

    //cerr << "result prop = " << result.proportionOfPopulation() << endl;

    return result;
}

void
ScoredStats::
sort()
{
    // Go from highest to lowest score
    std::sort(entries.begin(), entries.end());
    isSorted = true;
}

void
ScoredStats::
calculate()
{
    BinaryStats current;

    for (unsigned i = 0;  i < entries.size();  ++i)
        current.counts[entries[i].label][false] += entries[i].weight;

    if (!isSorted)
        sort();

    bestF = current;
    bestMcc = current;

    double totalAuc = 0.0;

    stats.clear();

    // take the all point
    stats.push_back(BinaryStats(current, INFINITY));

    for (unsigned i = 0;  i < entries.size();  ++i) {
        const ScoredStats::ScoredEntry & entry = entries[i];

        if (i > 0 && entries[i - 1].score != entry.score) {
            totalAuc += current.rocAreaSince(stats.back());
            stats.push_back
                (BinaryStats(current, entries[i - 1].score, entry.key));

            if (current.f() > bestF.f())
                bestF = stats.back();
            if (current.mcc() > bestMcc.mcc())
                bestMcc = stats.back();
            if (current.specificity() > bestSpecificity.specificity())
                bestSpecificity = stats.back();

#if 0
            cerr << "entry " << i << ": score " << entries[i - 1].score
                 << " p " << current.precision() << " r " << current.recall()
                 << " f " << current.f() << " mcc " << current.mcc() << endl;
#endif
        }

        bool label = entry.label;

        // We transfer from a false positive to a true negative, or a
        // true positive to a false negative

        double weight = entry.weight;

        current.counts[label][false] -= weight;
        current.counts[label][true] += weight;

        current.unweighted_counts[label][false] -= 1;
        current.unweighted_counts[label][true] += 1;

    }

    totalAuc += current.rocAreaSince(stats.back());

    if (!entries.empty())
        stats.push_back(BinaryStats(current, entries.back().score));

    auc = totalAuc;
}

void
ScoredStats::
add(const ScoredStats & other)
{
    if (!isSorted)
        throw MLDB::Exception("attempt to add to non-sorted separation stats");
    if (!other.isSorted)
        throw MLDB::Exception("attempt to add non-sorted separation stats");

    size_t split = entries.size();
    entries.insert(entries.end(), other.entries.begin(), other.entries.end());
    std::inplace_merge(entries.begin(), entries.begin() + split,
                       entries.end());

    // If we had already calculated, we recalculate
    if (!stats.empty())
        calculate();
}

void
ScoredStats::
dumpRocCurveJs(std::ostream & stream) const
{
    if (stats.empty())
        throw MLDB::Exception("can't dump ROC curve without calling calculate()");

    stream << "this.data = {" << endl;
    stream << MLDB::format("  \"aroc\": %8.05f , ", auc) << endl;
    stream << "  \"model\":{";
    for(unsigned i = 0;  i < stats.size();  ++i) {
        const BinaryStats & x = stats[i];
        stream << MLDB::format("\n  \"%8.05f\": { ", x.threshold);
        stream << MLDB::format("\n  tpr : %8.05f,", x.recall());
        stream << MLDB::format("\n  accuracy : %8.05f,", x.accuracy());
        stream << MLDB::format("\n  precision : %8.05f,", x.precision());
        stream << MLDB::format("\n  fscore : %8.05f,", x.f());
        stream << MLDB::format("\n  fpr : %8.05f,", x.falsePositiveRate());
        stream << MLDB::format("\n  tp : %8.05f,", x.truePositives());
        stream << MLDB::format("\n  fp : %8.05f,", x.falsePositives());
        stream << MLDB::format("\n  fn : %8.05f,", x.falseNegatives());
        stream << MLDB::format("\n  tn : %8.05f,", x.trueNegatives());
        stream << MLDB::format("}");
        stream << ",";
    }
    stream << "\n}";
    stream << "};";
}

void
ScoredStats::
saveRocCurveJs(const std::string & filename) const
{
    filter_ostream stream(filename);
    dumpRocCurveJs(stream);
}

Json::Value
ScoredStats::
getRocCurveJson() const
{
    if (stats.empty())
        throw MLDB::Exception("can't dump ROC curve without calling calculate()");

    Json::Value modelJs;
    for(unsigned i = 0;  i < stats.size();  ++i) {
        const BinaryStats & x = stats[i];
        modelJs[MLDB::format("%8.05f", x.threshold)] = x.toJson();
    }

    Json::Value js;
    js["aroc"] = auc;
    js["model"] = modelJs;
    js["bestF1Score"] = bestF.toJson();
    js["bestMcc"] = bestMcc.toJson();

    return js;
}

void
ScoredStats::
saveRocCurveJson(const std::string & filename) const
{
    filter_ostream stream(filename);
    stream << getRocCurveJson().toStyledString() << endl;
}

Json::Value
ScoredStats::
toJson() const
{
    Json::Value result;
    result["auc"] = auc;
    result["bestF1Score"] = bestF.toJson();
    result["bestMcc"] = bestMcc.toJson();
    return result;
}

} // namespace MLDB