#ifndef ASIMOVHYPOTHESIS_INCLUDED_HH
#define ASIMOVHYPOTHESIS_INCLUDED_HH

#include <vector>
#include <string>
#include <TMath.h>
#include <map>
#include <iostream>
#include <TH1D.h>
#include <TFile.h>
#include <TRandom.h>
#include "Math/Minimizer.h"
#include "Math/Factory.h"
#include "Math/Functor.h"
#include "Math/ProbFunc.h" // gaussian cdf 
#include "Math/QuantFuncMathCore.h" // gaussian quantile
#include "TMinuitMinimizer.h"

using namespace std;


/*! obtain p value from Z in std dev
 * slide 3 https://idpasc.lip.pt/uploads/talk/file/192/Stats-Tutorial2-2015.pdf
 * */
inline double pvalue_from_Z( double Z, bool onesided = true )
{	
	if (onesided) return ROOT::Math::gaussian_cdf(-Z,1,0); // sigma 1, mean 0
	else return ROOT::Math::gaussian_cdf(-Z,1,0)*2;
}

/*! obtain Z in std dev from p value */
inline double Z_from_pvalue( double p_value, bool onesided = true )
{	
	if (onesided) return ROOT::Math::gaussian_quantile(1-p_value,1); // sigma 1, mean 0
	else return ROOT::Math::gaussian_quantile(1-p_value/2,1);
}


/*! 
 * Perform asimov estimates of sensitivity and discovery potential
 * p_value function can cross check your limit/discovery by throwing some pseudo experiments
 * based on https://arxiv.org/pdf/1007.1727.pdf
 * */
struct AsimovHypothesis
{

	vector<pair<double,double>> asimov_data; /*! channels with expected signal (first) background (second) */
	vector<double> random_dataset; /*! vector can be used to store random dataset */
  int random_dataset_nobs; /*! number of observations in random data set */

  // histograms for storing pseudo experiments for cross checking results
	TH1D h_h0; /*! test statistic h0 hypothesis */
	TH1D h_h1; /*! test statistic h1 hypothesis */

	TH1D h_mu_hat_h0; /*! maximum likelihood estimator signal strength h0 */
	TH1D h_mu_hat_h1; /*! maximum likelihood estimator signal strength h1 */

  TH1D h_nobs_h0; /*! number of observations h0 */
  TH1D h_nobs_h1; /*! number of observations h1 */

	AsimovHypothesis( vector<double> _s, vector<double> _b, TString name = "x", double seed = 10 )
	{
		gRandom->SetSeed(seed);

		h_h0 = TH1D("h_h0_" + name,";test statistic;pseudo experiments", 20000, -40, 180);
		h_h1 = TH1D("h_h1_" + name,";test statistic;pseudo experiments", 20000, -40, 180);
		h_mu_hat_h0 = TH1D("h_mu_hat_h0_" + name,";mu hat h0;pseudo experiments", 10000, -3, 6);
		h_mu_hat_h1 = TH1D("h_mu_hat_h1_" + name,";mu hat h1;pseudo experiments", 10000, -3, 6);
		h_nobs_h0 = TH1D("h_nobs_h0_" + name,";observations;pseudo experiments", 10000, -0.5, 200.5);
		h_nobs_h1 = TH1D("h_nobs_h1_" + name,";observations;pseudo experiments", 10000, -0.5, 200.5);

		string channel;
		pair<double,double> s_b;
		pair<string,pair<double,double>> data_tmp;

		for ( unsigned int i = 0; i < _s.size(); i++ )
		{
			s_b = make_pair(_s[i], _b[i]);
			asimov_data.push_back(s_b);
		}
	}

	/*! combine two asimov hypotheses */
	void add( AsimovHypothesis p )
	{
		asimov_data.insert( asimov_data.end(), p.asimov_data.begin(), p.asimov_data.end() );

	}

	/*! poisson likelihood */
	double likelihood( double k, double l)
	{
		return TMath::Poisson(k,l);
	}

	/*! calculate log likelihood of internal dataset given expected mu_hat */
	double log_likelihood_dataset_internal( const double* x )
	{

		double mu_hat = x[0];

		double lik = 0;
		double n_expected;

		for ( int i = 0; i < random_dataset.size(); i++)
		{
			n_expected = mu_hat*asimov_data[i].first + asimov_data[i].second;
			if ( n_expected < 0 ) lik+= log(1e-9); // watch out for negative expected number of events!
			else lik+= log( likelihood( random_dataset[i], n_expected ) );
		}

		return -lik;
	}

	/*! cross check your found discovery/limit by doing pseudo experiments\
	 * 	for limit: p_value( found_limit, found_limit, 0.0, nexps) 
	 *  for disc : p_value( 0.0, 0.0, found_discovery, nexps)
	 *  */
	double p_value( double _mu_hypothesis, double _mu_h0, double _mu_h1, double nexps = 5 )
	{

		h_h0.Reset();
		h_h1.Reset();
		h_mu_hat_h0.Reset();
		h_mu_hat_h1.Reset();

		double mu_hypothesis[1] = { _mu_hypothesis }; // mu of hypothesis that is tested
		double mu_h0[1] = { _mu_h0 }; // mu for generating h0 datasets
		double mu_h1[1] = { _mu_h1 }; // mu for generating h1 datasets

		auto min = new TMinuitMinimizer();

		// create funciton wrapper for minmizer
		ROOT::Math::Functor f( this , &AsimovHypothesis::log_likelihood_dataset_internal, 1 );

		double step[1] = {0.001};
		double variable[1] = { 0.01 }; 		// starting point
		double mu_hat[1]; // ML estimator of mu
		double q_obs[1]; // observed test statistic

		// generate pseudo experiments for h0 and h1 dataset
		for (int i = 0; i < nexps; i++)
		{

			// h0 (bg only)
			generate_dataset_internal( mu_h0[0] );
			variable[0] = mu_h0[0] + 0.005;
			min->SetFunction(f);
			min->SetLimitedVariable(0, "x", variable[0] , step[0], -3, -3 + 1.E7);  // make sure the minimiser doesnt get stuck at minus large values
			min->Minimize(); // do the minimization
			const double *xs = min->X();
			mu_hat[0] = xs[0];
			if ( mu_hypothesis[0] > 0.0 && mu_hat[0] > mu_hypothesis[0] ) q_obs[0] = 0; // limit setting
			else if ( mu_hypothesis[0] == 0.0 && mu_hat[0] < 0.0 ) q_obs[0] = 0; // discovery
			else q_obs[0]  = -2*( -log_likelihood_dataset_internal( mu_hypothesis ) + log_likelihood_dataset_internal( mu_hat ) ); // extra minus sign from the minus sign in score function

			h_mu_hat_h0.Fill( mu_hat[0] );
      h_h0.Fill(q_obs[0]);
      h_nobs_h0.Fill( random_dataset_nobs );

			// h1 (sig+bg)
			generate_dataset_internal( mu_h1[0] );
			variable[0] = mu_h1[0] + 0.005;
			min->SetFunction(f);
			min->SetLimitedVariable(0, "x", variable[0] , step[0], -3, -3 + 1.E7);  // use 1.E7 which will make TMinuit happy
			min->Minimize(); // do the minimization
			xs = min->X();
			mu_hat[0] = xs[0];
			if ( mu_hypothesis[0] > 0.0 && mu_hat[0] > mu_hypothesis[0] ) q_obs[0] = 0; // limit setting
			else if ( mu_hypothesis[0] == 0.0 && mu_hat[0] < 0.0 ) q_obs[0] = 0; // discovery
			else q_obs[0]  = -2*( -log_likelihood_dataset_internal( mu_hypothesis ) + log_likelihood_dataset_internal( mu_hat ) );

			h_mu_hat_h1.Fill( mu_hat[0] );
			h_h1.Fill(q_obs[0]);
      h_nobs_h1.Fill( random_dataset_nobs );

		}

		// find the p value
		double prob[1] = {0.5};
		double q_median[1] = {0.0};
		h_h1.GetQuantiles(1, q_median, prob);
		int bin_obs = h_h0.FindBin( q_median[0] );
		double p_value = h_h0.Integral( bin_obs, h_h0.GetNbinsX() + 1 )/h_h0.Integral();

		delete min; // clean

		return p_value;

	}

	/*! return median likelihood estimator of h0 and h1 histograms */
	vector<double> median_mu_hat()
	{
		double prob[1] = {0.5};
		double q_median_h0[1] = {0.0};
		double q_median_h1[1] = {0.0};
		h_mu_hat_h0.GetQuantiles(1, q_median_h0, prob);
		h_mu_hat_h1.GetQuantiles(1, q_median_h1, prob);

		return {q_median_h0[0], q_median_h1[0]};
	}

	/*! return median nobs of h0 and h1 histograms */
  vector<double> median_nobs()
  {
    double prob[1] = {0.5};
    double q_median_h0[1] = {0.0};
    double q_median_h1[1] = {0.0};
    h_nobs_h0.GetQuantiles(1, q_median_h0, prob);
    h_nobs_h1.GetQuantiles(1, q_median_h1, prob);

    return {q_median_h0[0], q_median_h1[0]};
  }

	/*! generate random dataset using signal strength mu_hat */
	void generate_dataset_internal( double mu_hat )
	{
		random_dataset.clear();
    random_dataset_nobs = 0;

		double nobs;

		for ( auto channel : asimov_data)
		{
			nobs = gRandom->Poisson( mu_hat*channel.first + channel.second );
			random_dataset.push_back(nobs);
      random_dataset_nobs+=nobs;
		}
	}

	/*! general test statistic for asimov sensitivity estimation */
	double test_statistic_asimov( double mu, double mu_hat )
	{
		double log_L_mu     = 0;
		double log_L_mu_hat = 0;

		for (auto channel : asimov_data)
		{
			log_L_mu     += log( likelihood( mu_hat*channel.first + channel.second, mu    *channel.first + channel.second) + 1e-15 ); // 1e-15: no log of 0
			log_L_mu_hat += log( likelihood( mu_hat*channel.first + channel.second, mu_hat*channel.first + channel.second) + 1e-15 );
		}

		return - 2 * ( log_L_mu - log_L_mu_hat );
	}

	/*! 
  Test statistic for limit setting (equation 14, https://arxiv.org/pdf/1007.1727.pdf)
  Using likelihood ratio (equation 6)
  Likelihood is not only a function of mu, but also of the data (mu_hat)
	*/
	double q_mu_asimov( double mu, double mu_hat )
	{
		if ( mu_hat > mu ) return 0;

		return test_statistic_asimov( mu, mu_hat);

	}

	/*! 
  Test statistic for discovery (equation 12), https://arxiv.org/pdf/1007.1727.pdf)
	*/
  double q_0_asimov( double mu_hat )
	{
		if ( mu_hat < 0 ) return 0;

		double mu = 0;

		return test_statistic_asimov( mu, mu_hat);

	}

	/*! 
  Find mu for discovery or limit setting that corresponds to a given Z
  confidence levels for limit setting need to be transformed to a p value: 1 - cl, and then to Z in std dev
	*/
	double find_mu( double Z, bool discovery = true )
	{

		double mu = 0.0, Z_obs = 0.0, q_obs = 0.0;
		
		// boundaries to look for mu
    double mu_low = 0; 
    double mu_high = 10;

    int iterations = 0;
		while ( abs(Z_obs - Z) > 0.01 )
		{
      iterations+=1;
			mu = (mu_high - mu_low)/2 + mu_low;
			
      if ( discovery ) q_obs = q_0_asimov( mu );
      else q_obs = q_mu_asimov( mu, 0.0 );

			Z_obs = sqrt(q_obs);

			if ( Z_obs > Z ) mu_high = mu;
      else mu_low = mu;

		}

		return mu;
	}

	// print output help functions
	TH1D bg_histogram(  )
	{

		TH1D h_bg = TH1D("h_bg",";channel;expected events", asimov_data.size(), 0.5, asimov_data.size() + 0.5 );
		for ( int i = 0; i < asimov_data.size(); i++)
		{
			h_bg.SetBinContent(i + 1, asimov_data[i].second );
		}
		return h_bg;
	}

	TH1D sig_histogram( double mu )
	{

		TH1D h_sig = TH1D("h_sig",";channel;expected events", asimov_data.size(), 0.5, asimov_data.size() + 0.5 );
		for ( int i = 0; i < asimov_data.size(); i++)
		{
			h_sig.SetBinContent(i + 1, asimov_data[i].first*mu );
		}
		return h_sig;
	}

	TH1D sig_bg_histogram( double mu )
	{

		TH1D h_sig_bg = TH1D("h_sig_bg",";channel;expected events", asimov_data.size(), 0.5, asimov_data.size() + 0.5 );
		for ( int i = 0; i < asimov_data.size(); i++)
		{
			h_sig_bg.SetBinContent(i + 1, asimov_data[i].first*mu + asimov_data[i].second );
		}
		return h_sig_bg;
	}

	TH1D return_h_h0(  )
	{
		return h_h0;
	}

	TH1D return_h_h1(  )
	{
		return h_h1;
	}

	TH1D return_h_mu_hat_h0(  )
	{
		return h_mu_hat_h0;
	}

	TH1D return_h_mu_hat_h1(  )
	{
		return h_mu_hat_h1;
	}

  TH1D return_h_nobs_h0(  )
  {
    return h_nobs_h0;
  }

  TH1D return_h_nobs_h1(  )
  {
    return h_nobs_h1;
  }

	int asimov_data_size()
	{
		return asimov_data.size();
	}

	int random_data_size()
	{
		return random_dataset.size();
	}

	void print_asimov_data( double mu = 1.0 )
	{
		cout << "print asimov mu " << mu << endl;
		for ( int i = 0; i < asimov_data.size(); i++)
		{
			cout << "s " << asimov_data[i].first*mu << " b " << asimov_data[i].second << endl;
		}
	}

	void print_random_data( )
	{
		cout << "print dataset" << endl;
		for ( int i = 0; i < random_dataset.size(); i++)
		{
			cout << "n " << random_dataset[i] << endl;
		}
	}

	// // probably depricated when moving to c++ fully
	// // we dont have to return the dataset anymore, we keep it in the struct
	// vector<double> generate_dataset( double mu_hat )
	// {

	// 	vector<double> dataset;
	// 	double nobs;

	// 	for ( auto channel : asimov_data)
	// 	{
	// 		nobs = gRandom->Poisson( mu_hat*channel.first + channel.second );
	// 		// cout << nobs << endl;
	// 		dataset.push_back(nobs);
	// 	}

	// 	return dataset;
	// }

	// double log_likelihood_dataset( vector<double> dataset, double mu_hat )
	// {

	// 	double lik = 0;
	// 	double n_expected;

	// 	for ( int i = 0; i < dataset.size(); i++)
	// 	{
	// 		n_expected = mu_hat*asimov_data[i].first + asimov_data[i].second;
	// 		if ( n_expected < 0 ) lik+= log(1e-9); 
	// 		else lik+= log( likelihood( dataset[i], n_expected ) );
	// 	}
	// 	return -lik;
	// }

 //  double find_discovery_mu_hat( double Z )
 //  {

 //    double mu_hat_discovery = 0.4, Z_discovery = 0.0, q_obs = 0.0;

 //    int iterations = 0;
 //    while ( Z_discovery < Z )
 //    {
 //      iterations+=1;
 //      mu_hat_discovery += 0.01;
 //      q_obs = q_0_asimov( mu_hat_discovery );
 //      Z_discovery = sqrt(q_obs);
 //    }

 //    cout << "iterations " << iterations << endl;
 //    cout << "Z_discovery " << Z_discovery << endl;

 //    return mu_hat_discovery;
 //  }

 //  double find_limit_mu( double Z )
 //  {

 //    // Z = Z_from_pvalue( 1 - cl );

 //    double mu_limit = 0.0, Z_limit = 0.0, q_obs = 0.0;

 //    while ( Z_limit < Z )
 //    {
 //      mu_limit += 0.01;
 //      q_obs = q_mu_asimov( mu_limit, 0.0 );
 //      Z_limit = sqrt(q_obs);
 //    }

 //    return mu_limit;
 //  }



};


#endif