#include <stdio.h>
#include <iostream>

#include <vector>
#include <string>
#include <string.h>

#include "TFile.h"
#include "TH1D.h"
#include "TCanvas.h"
#include "TPad.h"
#include "TLegend.h"
#include "TLatex.h"



double pdf_model ( double* var, double* par )
{
	double variable = var[0];
	if ( !par[1] && !par[3] )   return 0;
	if ( !par[0] && !par[2] )   return 0;
	double bkg_slope  = fabs( par[0] );
	double sig_slope  = fabs( par[2] );
	double bkg_norm   = bkg_slope / ( 1-exp(-2*bkg_slope) );
	double sig_norm   = sig_slope / ( 1-exp(-2*sig_slope) );
	double func_scale = par[1] + par[3];
	double bkg_scale  = par[1] * bkg_norm / func_scale;
	double sig_scale  = par[3] * sig_norm / func_scale;
	double f_bkg = bkg_scale * exp( -bkg_slope * ( variable+1.0 ) );
	double f_sig = sig_scale * exp( sig_slope * ( variable-1.0 ) );
	if ( bkg_slope*sig_slope != 0 )
		return ( fabs(variable)<=1 ) * ( f_bkg + f_sig );
	if ( bkg_slope == 0 )
		return ( fabs(variable)<=1 ) * ( f_sig );
	return ( fabs(variable)<=1 ) * ( f_bkg );
}

TF1* custom_function (
		const char* func_name,
		double var_min, double var_max,
		std::vector<double> parameter_list )
{
	int n_var = parameter_list.size( );
	TF1 *func = new TF1( func_name, pdf_model, var_min, var_max, n_var );
	for ( int i=0; i<n_var; i++ )
		func -> SetParameter( i, parameter_list[i] );
	return func;
}

TH1D* event_generator(
		TF1* model, const char* hist_name,
		long n_trials, TRandom3* rgen = nullptr )
{
	if ( !rgen )  rgen = new TRandom3( 0 );
	TH1D *hist = new TH1D( hist_name, "", 10, -1.0, 1.0 );
	double fmax_value = 2 * model->GetMaximum( );
	while ( n_trials )
	{
		double rand_x = rgen -> Uniform( -1.0, 1.0 );
		double rand_y = rgen -> Uniform( 0.0, fmax_value );
		double f_value = model -> Eval( rand_x );
		if ( rand_y > f_value )  continue;
		double eff_x = rand_x;
		if ( fabs(rand_x) >= 1.0 )
			eff_x = 0.999 * rand_x / fabs(rand_x);
		hist -> Fill( eff_x );
		n_trials--;
	}
	return hist;
}


int test_tfractionfitter ( )
{
	TH1::AddDirectory( kFALSE );

	double lower_limit = -1.2;
	double upper_limit = 1.2;
	// The following definition declare 20% of signal in the mix sample
	TF1 *mix_model = custom_function (
			"mix_model", lower_limit, upper_limit,
			std::vector<double>{ -2.0, 0.8, 1.6, 0.2 } );
	TF1 *signal_model = custom_function (
			"signal_model", lower_limit, upper_limit,
			std::vector<double>{ 0.0, 0.0, 1.6, 1.0 } );
	TF1 *background_model = custom_function (
			"background_model", lower_limit, upper_limit,
			std::vector<double>{ -2.0, 1.0, 0.0, 0.0 } );

	double n_entries = 100000;
	TH1D *hist_mix = event_generator ( mix_model, "mixed", n_entries );
	TH1D *hist_sig = event_generator ( signal_model, "signal", n_entries );
	TH1D *hist_bkg = event_generator ( background_model, "background", n_entries );
	hist_sig -> Scale( 1.0 / hist_sig->Integral( ) );
	hist_bkg -> Scale( 1.0 / hist_bkg->Integral( ) );

	double n_total = hist_mix -> Integral( );
	TObjArray* components = new TObjArray( );
	components -> Add( hist_sig );
	components -> Add( hist_bkg );
	TFractionFitter *fitter = new TFractionFitter( hist_mix, components );
	//fitter -> Constrain( 0, 0.05, 0.50 );
	//fitter -> Constrain( 1, 0.05, 0.90 );
	fitter -> SetRangeX( 1, 10 );

	int fit_status = fitter -> Fit( );
	double par_value, par_error;
	fitter -> GetResult( 0, par_value, par_error );
	printf( "fit result for signal\n" );
	printf( "|-- signal yield = %.0f\n", n_total * par_value );
	printf( "`-- yield error =  %.0f\n", n_total * par_error );

	return 0;
}