#include "BinnedAnalysis.hh"

/*! fill histogram with pseudo-data for the given sigstrength */

void BinnedAnalysis2D::fill_pseudo_dataset( vector<TH2D>& data, double sigstrength )
{
    TIMEFUNC

    for (unsigned i =0 ; i < data.size() ; i++ )
    { 	
        TH2D& h                  = data[i];
        const TH2D& h_signal     = signal_histos[i];
        const TH2D& h_background = background_histos[i];

        for (int ix = 1; ix<= h.GetNbinsX() ; ix ++ ) 
            for (int iy = 1; iy<= h.GetNbinsX() ; iy ++ ) 
                {
                    double s = h_signal.GetBinContent( ix,iy);
                    double b = h_background.GetBinContent( ix,iy );
                    double mu = b + sigstrength * s;
                    if ( mu == 0 ) { if (verb>2 ) print ("mu ==0", ix, iy); mu = mumin; }

                    int k = gRandom->Poisson( mu );
                    h.SetBinContent( ix,iy, k );
                }
        }
}


/*! Log of the poisson likelihood of the dataset, given the signal strength */

double BinnedAnalysis2D::log_likelihood( const vector<TH2D>& data , double sigstrength )
{
    double r = 0;
    for (unsigned i =0 ; i < data.size() ; i++ )
    { 	
        const TH2D& h            = data[i];
        const TH2D& h_signal     = signal_histos[i];
        const TH2D& h_background = background_histos[i];

        for (int ix = 1; ix<= h.GetNbinsX() ; ix ++ ) 
            for (int iy = 1; iy<= h.GetNbinsX() ; iy ++ ) 
                {
                    int    k = (int) h.GetBinContent(ix,iy);
                    double s = h_signal.GetBinContent( ix,iy);
                    double b = h_background.GetBinContent( ix,iy );

                    double mu = b + sigstrength * s;
                    if ( mu == 0 ) { if (verb>2) print ("mu ==0", ix, iy); mu = mumin; }

                    r += k * log( mu ) - mu;
                }
    }

    if (verb>1) print( "lik for", sigstrength, " = ", r );
    return r;
}

struct sparse_sum // for only keeping the k>0 terms.
{
    double stot, btot;
    vector<double> ss;
    vector<double> bs;

    double log_likelihood( double sigstrength )
    {
        TIMEFUNC
        double c = 0;
        for (unsigned i =0 ; i < ss.size() ; i++)
        {
           c+= log ( bs[i]+ sigstrength * ss[i] );
        }

    return c - sigstrength * stot - btot;
    }

    void set(const vector<TH2D>& data, const vector<TH2D>& signal_histos,  const vector<TH2D>& background_histos ) 
    {
        TIMEFUNC

        ss.clear(); bs.clear(); stot=btot=0;

        for (unsigned i =0 ; i < data.size() ; i++ )
        { 	
            const TH2D& h            = data[i];
            const TH2D& h_signal     = signal_histos[i];
            const TH2D& h_background = background_histos[i];

            for (int ix = 1; ix<= h.GetNbinsX() ; ix ++ ) 
                for (int iy = 1; iy<= h.GetNbinsX() ; iy ++ ) 
                    {
                        int    k = (int) h.GetBinContent(ix,iy);
                        double s = h_signal.GetBinContent( ix,iy);
                        double b = h_background.GetBinContent( ix,iy );

                        stot += s;
                        btot += b;
                        for (int j=0;j<k;j++) { bs.push_back(b); ss.push_back(s);}
                    }
        }
       // print ("initialised sparse struct with", bs.size(), ss.size() );
    }
};


/*! The log likelihood ratio of the data as function of the tested sig. strength */


double BinnedAnalysis2D::log_likelihood_ratio( const vector<TH2D>& data, double sigstrength ) 
{
    return log_likelihood( data, sigstrength) - log_likelihood( data , 0 );
}


/*! Fit the signal-strength, optimising the log-likelihood */

double BinnedAnalysis2D::optimize_sigstrength( const vector<TH2D>& data , 
                                               double startvalue , double xmin, double xmax, double stepsize)
{
    TIMEFUNC

    TMinuitMinimizer minimizer; 
    ROOT::Math::Functor score;
    sparse_sum spsum;

    if ( sparse )
    {
        spsum.set( data , signal_histos, background_histos ); 
        score =  ROOT::Math::Functor(  [&] ( const double* x ) { return -spsum.log_likelihood( x[0]); }, 1 );
    }
    else
    {
        score = ROOT::Math::Functor(  [&] ( const double* x ) { return -log_likelihood( data, x[0]); }, 1 );   
    }
    
    minimizer.SetFunction(score);
    minimizer.SetLimitedVariable(0, "sigstrength", startvalue, stepsize, 0.1 , 1.E7 ); 
    minimizer.Minimize();

    // return the result
    double ss = minimizer.X()[0];	
    if ( verb > 0 ) print ("best sigstrength-value=", ss );
    return ss;
}


/*! Return a histogram with the distribution of the test-statistic for the given signal strength */ 

TH1D BinnedAnalysis2D::test_statistic_distribution( double sig_strength , int n_pe  , 
                                                    int nbins, double xmin , double xmax  )
{
    TStopwatch s;

    TH1D r( Form("h_%3.2f", sig_strength),Form("h_%3.2f", sig_strength), nbins, xmin, xmax  );

    vector<TH2D> pe = signal_histos; // copy

    for (int i=0; i< n_pe; i++ )
    {
        fill_pseudo_dataset( pe , sig_strength );
        double ts =  test_statistic( pe );
    
        double tsprime = sqrt(fabs(ts)) *  (ts<0?-1:1);

        r.Fill( ts );
    }

    print ("generated", r.Integral(), "pe's in ", s.RealTime () ,"s.");
    return r;
}