// File:  SigCalc.cc
// Glen Cowan
// RHUL Physics

// Implementation of SigCalc class

#include <iostream>
#include <string>
#include <vector>
#include <cmath>
#include <TRandom3.h>
#include "SigCalc.h"
#include "fitPar.h"
 
SigCalc::SigCalc(double n, double s, vector<double> mVec,  
                 vector<double> tauVec, vector<double> S1Vec,
                 vector<double> S2Vec, int option){
  m_n = n;
  m_s = s;
  m_m = mVec;
  m_tau = tauVec;
  m_S1 = S1Vec;
  m_S2 = S2Vec;
  m_numBck = mVec.size();
  m_option = option;

}

//-------------------------------------------------------------------------
//
// For this constructor, input parameter values and generate
// a data set.

SigCalc::SigCalc(double s, vector<double> tauVec, double mu, 
                 vector<double> bVec, vector<double> aVec,
		 vector<double> xiVec, TRandom3* ran, int option){
  
  m_s = s;
  m_tau = tauVec;
  m_option = option;
  m_numBck = bVec.size();

  // Generate a data set.

  double btot = 0.;
  for (int i=0; i<m_numBck; i++){
    btot += bVec[i];
  }
  int n = ran->Poisson(mu*s + btot);
  cout << "Generating data set from input parameters: " << endl;
  cout << "mu, s, btot, n = " << mu << "  " << s << "  "
       << btot << "  " << n << endl;

  double omega = 1.;          // exact by construction
  vector<double> mVec;
  vector<double> S1Vec;
  vector<double> S2Vec;
  for (int i=0; i<m_numBck; i++){
    int m = ran->Poisson(tauVec[i]*bVec[i]/omega);
    mVec.push_back(m);
    double wsum = 0.;
    double w2sum = 0.;
    double lsum = 0.;
    double l2sum = 0.;
    double md = static_cast<double>(m);
    vector<double> wVec(m);
    vector<double> lVec(m);
    for (int j=0; j<m; j++){
      double r = ran->Rndm();
      double x = aVec[i]*r;
      double w = aVec[i] * (1./xiVec[i])*exp(-x/xiVec[i]) / 
                 (1 - exp(-aVec[i]/xiVec[i]));
      double lnw = log(w);
      wVec[j] = w;
      lVec[j] = lnw;
      wsum += w;
      w2sum += w*w;
      lsum += log(w);
      l2sum += lnw*lnw;
    }

    cout << "background " << i << ", m = " << m << endl;
    cout << "w, lnw: " << endl;
    for (int j=0; j<m; j++){
      cout << wVec[j] << "    " << lVec[j] << endl;
    }
    cout << endl;

    if ( option < 100 ) {
      S1Vec.push_back(lsum);
      S2Vec.push_back(l2sum);
    }
    else if ( option >= 100 && option < 300 ) {
      S1Vec.push_back(wsum);
      S2Vec.push_back(w2sum);
    }

    double z = 1. - exp(-aVec[i]/xiVec[i]);
    double sW2 = (0.5*aVec[i]/xiVec[i]) * 
                 (1. - exp(-2.*aVec[i]/xiVec[i])) / (z*z) - 1.; 
    double sigmaW = sqrt(sW2); 
    double lambda = log (aVec[i]/(xiVec[i]*z)) - 0.5*aVec[i]/xiVec[i];
    double sigmaL = (aVec[i]/xiVec[i]) / sqrt(12.);

    cout << "bkg#, a, xi = " << i << "  " << aVec[i] << "  " 
         << xiVec[i] << endl;
    cout << "sigmaW, lambda, sigmaL = " << sigmaW << "  "
         << lambda << "  " << sigmaL << endl;
    cout << "m, S1, S2, = " << m << "  " << S1Vec[i] << "  " 
         << S2Vec[i] << endl;

  }

  // assign to rest of member variables

  m_n = n;
  m_m = mVec;
  m_S1 = S1Vec;
  m_S2 = S2Vec;

}

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

double SigCalc::lnL(double mu, vector<double> bVec, 
       vector<double> lambdaVec, vector<double> sigmaVec){

  // option < 100 is log-normal model
  // 100 <= option < 200 is Gaussian; lambdaVec actually contains omega
  // 200 <= option < 300 is Gaussin in bHat
  // 300 <= option < 400 is for all events with w = 1

  int option = this->option();
  double logL = -9999.;
  if ( option < 100 ) {
    logL = this->lnLLogNorm(mu, bVec, lambdaVec, sigmaVec);
  }
  else if ( option >= 100 && option < 200 ) { 
    logL = this->lnLGauss(mu, bVec, lambdaVec, sigmaVec);
  }
  else if ( option >= 200 && option < 300 ) {
    logL = this->lnLGaussBhat(mu, bVec);
  }
  else if ( option >= 300 && option < 400 ) {
    logL = this->lnLw1(mu, bVec);
  }

  return logL;

}

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

double SigCalc::lnLLogNorm(double mu, vector<double> bVec, 
       vector<double> lambdaVec, vector<double> sigmaVec){

// The log-likelihood function.
// mu = global strength parameter
// here bVec is the vector of adjustable background values (not
// to be confused with the m values in the SigCalc object,
// which are estimates based on MC or sidebands.

  double btot = 0;
  for (int i=0; i<bVec.size(); i++){
    btot += bVec[i];
  }  
  double nu = mu*this->s() + btot;

  double logL = 0.;
  if ( nu > 0. ) {
    logL = psi(this->n(), nu);
  }
  else {
    logL = -9999;
    //     cout << "warning, nu = " << nu << endl;
  }
   
  for (int i=0; i<this->numBck(); i++){
    double s2i = sigmaVec[i]*sigmaVec[i];

    // if m = 0, then treat lack of events as lack of unweighted (i.e. w=1)
    // events, therefore in this case likelihood should be the
    // same as for unweighted events, with m ~ Poisson(tau*b)
    double oi;
    if ( this->m(i) > 0 ){
      oi = exp(lambdaVec[i] + 0.5*s2i);
    }
    else {
      oi = 1.;
    }
    logL += psi( this->m(i), this->tau(i)*bVec[i]/oi );
  }
  
 // Only include following part of lnL if m >  0

  for (int i=0; i<this->numBck(); i++){
    double l = lambdaVec[i];
    double s = sigmaVec[i];
    if ( m(i) > 0 ) {
      double deltaLnL = -0.5*this->m(i)*l*l/(s*s) - this->m(i)*log(s) + 
                        (2*l*this->S1(i) - this->S2(i))/(2*s*s);
      logL += deltaLnL;
    }
    //    cout << "l, s, m, S1, S2, dL = " << l << "  " << s 
    //    << "  " << this->m(i) << "  " <<
    //  this->S1(i) << "  " << this->S2(i) << "  " << deltaLnL << endl;
  }

  // cout << "logL = " << logL << endl;
  return logL;

}

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

double SigCalc::lnLGauss(double mu, vector<double> bVec, 
       vector<double> omegaVec, vector<double> sigmaVec){

// The log-likelihood function based on w ~ Gaussian
// mu = global strength parameter
// here bVec is the vector of adjustable background values (not
// to be confused with the m values in the SigCalc object,
// which are estimates based on MC or sidebands.

// This model assumes S1 = sum of weights, S2 = sum of w^2

  double btot = 0;
  for (int i=0; i<bVec.size(); i++){
    btot += bVec[i];
  }  
  double nu = mu*this->s() + btot;

  double logL = 0.;
  if ( nu > 0. ) {
    logL = psi(this->n(), nu);
  }
  else {
    logL = -9999;
    //     cout << "warning, nu = " << nu << endl;
  }
   
  for (int i=0; i<this->numBck(); i++){
    double s2i = sigmaVec[i]*sigmaVec[i];

    // if m = 0, then treat lack of events as lack of unweighted (i.e. w=1)
    // events, therefore in this case likelihood should be the
    // same as for unweighted events, with m ~ Poisson(tau*b)
    double oi;
    if ( this->m(i) > 0 ){
      oi = omegaVec[i];
    }
    else {
      oi = 1.;
    }
    logL += psi( this->m(i), this->tau(i)*bVec[i]/oi );
  }
  
 // Only include following part of lnL if m >  0

  for (int i=0; i<this->numBck(); i++){
    double o = omegaVec[i];
    double s = sigmaVec[i];
    if ( m(i) > 0 ) {
      double deltaLnL = -0.5*this->m(i)*o*o/(s*s) - this->m(i)*log(s) + 
                        (2*o*this->S1(i) - this->S2(i))/(2*s*s);
      logL += deltaLnL;
    }
    //    cout << "o, s, m, S1, S2, dL = " << o << "  " << s 
    //    << "  " << this->m(i) << "  " <<
    //  this->S1(i) << "  " << this->S2(i) << "  " << deltaLnL << endl;
  }

  // cout << "logL = " << logL << endl;
  return logL;

}

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

double SigCalc::lnLGaussBhat(double mu, vector<double> bVec){

  double btot = 0;
  for (int i=0; i<bVec.size(); i++){
    btot += bVec[i];
  }  
  double nu = mu*this->s() + btot;

  double logL = 0.;
  if ( nu > 0. ) {
    logL = psi(this->n(), nu);
  }
  else {
    logL = -9999;
    //     cout << "warning, nu = " << nu << endl;
  }

  for (int i=0; i<this->numBck(); i++){
    double bHat = this->S1(i)/tau(i);
    double sigma2bHat = this->S2(i)/(tau(i)*tau(i));
    double z = (bVec[i] - bHat);
    logL += -0.5 * z * z / sigma2bHat;
  }

  return logL;

}

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

double SigCalc::lnLw1(double mu, vector<double> bVec){

  // This version is for when all events have w=1.

  double btot = 0;
  for (int i=0; i<bVec.size(); i++){
    btot += bVec[i];
  }  
  double nu = mu*this->s() + btot;
  double logL = psi(this->n(), nu);

  for (int i=0; i<this->numBck(); i++){
    logL += psi( this->m(i), this->tau(i)*bVec[i] );
  }

  return logL;

}

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

double SigCalc::qmu(double mu){

// Returns qmu = - 2 ln lambda(mu), where lambda = profile likelihood ratio.
// For this version the only argument is mu.  If you
// need access to estimates of parameters, 
// use the version of qmu below.

  double muHat;
  vector<double> bHat;
  vector<double> bHatHat;
  vector<double> lambdaHat;
  vector<double> lambdaHatHat;
  vector<double> sigmaHat;
  vector<double> sigmaHatHat;
  double q = this->qmu(mu, muHat, bHat, bHatHat, lambdaHat, 
                       lambdaHatHat, sigmaHat, sigmaHatHat);
  return q;

}

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

double SigCalc::qmu(double mu, double& muHat, 
                 vector<double>& bHat, vector<double>& bHatHat, 
                 vector<double>& lambdaHat, vector<double>& lambdaHatHat, 
                 vector<double>& sigmaHat, vector<double>& sigmaHatHat){

// This version of qmu passes back estimates of the parameters

  int numBck = this->numBck();
  int npar = 1 + 3*numBck;   // KLUDGE -- for now true for all options
  vector<double> parVec(npar);
  vector<bool> freePar(npar);

// Fix mu to input value and fit b (gives bHatHat)

  parVec[0] = mu;
  freePar[0] = false;
  this->setStartValues(parVec, freePar);
  int status = fitPar(this, freePar, parVec);

  if ( status != 0 ) {
    cout << "HatHat: status, muHat = " << status << "  " << parVec[0] << endl;
    for (int i=1; i<numBck; i++){
      cout << parVec[1] << "  " << parVec[2] << "  " << parVec[3] << endl;
      cout << "n, m, S1, S2 = " << this->n() << "  " << this->m(i) << 
	"  " << this->S1(i) << "  " << this->S2(i) << endl;
    }
    cout << endl;
  }

  bHatHat.clear();
  lambdaHatHat.clear();
  sigmaHatHat.clear();
  for (int i=0; i<numBck; i++){
    bHatHat.push_back(parVec[3*i+1]);
    lambdaHatHat.push_back(parVec[3*i+2]);
    sigmaHatHat.push_back(parVec[3*i+3]);
  }
  double lnLmubHatHat = this->lnL(mu, bHatHat, 
			lambdaHatHat, sigmaHatHat);

// Fit all parameters

  freePar[0] = true;
  this->setStartValues(parVec, freePar);
  status = fitPar(this, freePar, parVec);

  if ( status != 0 ) {
    cout << "Hat: status, muHat = " << status << "  " << parVec[0] << endl;
    for (int i=1; i<numBck; i++){
      cout << parVec[1] << "  " << parVec[2] << "  " << parVec[3] << endl;
      cout << "n, m, S1, S2 = " << this->n() << "  " << this->m(i) << 
	"  " << this->S1(i) << "  " << this->S2(i) << endl;
    }
    cout << endl;
  }

  muHat = parVec[0];
  bHat.clear();
  lambdaHat.clear();
  sigmaHat.clear();
  for (int p=0; p<numBck; p++){
    bHat.push_back(parVec[3*p + 1]);
    lambdaHat.push_back(parVec[3*p + 2]);
    sigmaHat.push_back(parVec[3*p + 3]);
  }
  double lnLmuHatbHat = this->lnL(muHat, bHat, lambdaHat, sigmaHat);

  double qVal = 0;
  if ( muHat < mu ) {
    qVal = 2.*(lnLmuHatbHat - lnLmubHatHat);
    qVal = max(0., qVal);           // KLUDGE protect against numerical error.
  }
  return qVal;
  

}

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

double SigCalc::q0(){

// Returns q0 = - 2 ln lambda(0), where lambda = profile likelihood ratio.
// For this version there are no arguments.  If you
// need access to muHat, bHat, bHatHat, use the version of q0 below.

  double muHat;
  vector<double> bHat;
  vector<double> bHatHat;
  vector<double> lambdaHat;
  vector<double> lambdaHatHat;
  vector<double> sigmaHat;
  vector<double> sigmaHatHat;
  double qVal = this->q0(muHat, bHat, bHatHat, lambdaHat, 
			 lambdaHatHat, sigmaHat, sigmaHatHat);
  return qVal;

}

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

double SigCalc::q0(double& muHat, 
       vector<double>& bHat, vector<double>& bHatHat, 
       vector<double>& lambdaHat, vector<double>& lambdaHatHat, 
       vector<double>& sigmaHat, vector<double>& sigmaHatHat){

// This version of q0 passes back estimates of the parameters

  double mu = 0.;
  int numBck = this->numBck();
  int option = this->option();

  int npar = 1 + 3*numBck;   // KLUDGE -- for now true for all options
  vector<double> parVec(npar);
  vector<bool> freePar(npar);

  // Set up start values.  Minor kludge -- for w ~ Guassian,
  // the parameter "lambda" really holds sigma, and S1 and S2
  // are sums of w and w^2, rather than l and l^2.  The prescriptions
  // for the starting values are designed for the log-normal model,
  // but with minor mod for m=0, try to see if they work for the
  // Gaussian version as well.

  parVec[0] = mu;  
  freePar[0] = false; 
  this->setStartValues(parVec, freePar);

  //   for (int i=0; i<parVec.size(); i++){
  //     cout << "about to call fitPar  i, par, freePar... " 
  //     << i << "  " << parVec[i] << "  " << freePar[i] << endl;
  //   }

  vector<double> startPar(npar);
  for (int i=0; i<parVec.size(); i++){
    startPar[i] = parVec[i];
  }

  int status = fitPar(this, freePar, parVec);

  //  if ( status != 0 ) {
  //  cout << "HatHat(0) status= " << status << endl;
  //  cout << "par 1,2,3 = " << 
  //  parVec[1] << "  " << parVec[2] << "  " << parVec[3] << endl;
  //  cout << "start par = " << startPar[0] << "  " << startPar[1] <<
  //    "  " << startPar[2] << "  " << startPar[3] << endl;
  //  for (int i=0; i<numBck; i++){
  //    cout << "n, m, S1, S2 = " << this->n() << "  " << this->m(i) << 
  //	"  " << this->S1(i) << "  " << this->S2(i) << endl;
  //  }
  //  cout << endl;
  //  }


  bHatHat.clear();
  lambdaHatHat.clear();
  sigmaHatHat.clear();
  for (int i=0; i<numBck; i++){
    bHatHat.push_back(parVec[3*i + 1]);
    lambdaHatHat.push_back(parVec[3*i + 2]);
    sigmaHatHat.push_back(parVec[3*i + 3]);
    // cout << "HatHat:  " << i << "  " << bHatHat[i] << "  " <<
    //       lambdaHatHat[i] << "  " << sigmaHatHat[i] << endl;
  }
  double lnLmubHatHat = this->lnL(mu, bHatHat, 
         lambdaHatHat, sigmaHatHat);
  // cout << "lnLmubHatHat = " << lnLmubHatHat << endl;

// Fit all parameters

  freePar[0] = true; 
  this->setStartValues(parVec, freePar);

  for (int i=0; i<parVec.size(); i++){
    startPar[i] = parVec[i];
  }

  status = fitPar(this, freePar, parVec);
  
  //  if ( status != 0 ) {
  //  cout << "Hat: status, muHat = " << status << "  " << parVec[0] << endl;
  //  cout << "par 1,2,3 = " <<
  //           parVec[1] << "  " << parVec[2] << "  " << parVec[3] << endl;
  //  cout << "start par = " << startPar[0] << "  " << startPar[1] <<
  //    "  " << startPar[2] << "  " << startPar[3] << endl;
  //  for (int i=0; i<numBck; i++){
  //    cout << "n, m, S1, S2 = " << this->n() << "  " << this->m(i) << 
  //     "  " << this->S1(i) << "  " << this->S2(i) << endl;
  //  }
  //  cout << endl;
  //  }

  muHat = parVec[0];
  // cout << "muHat = " << muHat << endl;
  bHat.clear();
  lambdaHat.clear();
  sigmaHat.clear();
  for (int i=0; i<numBck; i++){
    bHat.push_back(parVec[3*i + 1]);
    lambdaHat.push_back(parVec[3*i + 2]);
    sigmaHat.push_back(parVec[3*i + 3]);
    // cout << "Hat:  " << i << "  " << bHat[i] << "  " <<
    // lambdaHat[i] << "  " << sigmaHat[i] << endl;
  }
  double lnLmuHatbHat = this->lnL(muHat, bHat, lambdaHat, sigmaHat);
  // cout << "lnLmuHatbHat = " << lnLmuHatbHat << endl; 

  double qVal = 0;
  if ( muHat > 0 ) {
    qVal = 2.*(lnLmuHatbHat - lnLmubHatHat);
    qVal = max(0., qVal);           // KLUDGE protect against numerical error.
  }
  //  cout << "q0 = " << qVal << endl;

  return qVal;

}

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

double SigCalc::psi(double n, double nu){

// The function psi returns log of Poisson probability nu^n exp(-nu)
// (the n! term is dropped).  For nu -> 0 and n = 0 the probability is 1,
// i.e., psi returns zero.  The function avoids evaluating log(nu) for 
// nu <= epsilon.

  const double epsilon = 1.e-100;
  const double logOfSmallValue = -300.;
  double val;
  if ( n <= epsilon && nu <= epsilon ) {
    val = 0.;
    // cout << "psi = 0 " << n << "  " << nu << endl;
  }
  else if ( n > epsilon && nu <= epsilon ) {
    val = - n * logOfSmallValue; 
    // cout << "psi = small " << n << "  " << nu << endl;
  }
  else {
    val = n*log(nu) - nu;
  }
  return val;
}

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

void SigCalc::setStartValues(vector<double>& parVec,
                             vector<bool>& freePar){

  int option = this->option();
  int numBck = this->numBck();

  if ( option < 100 ) {         // w ~ log normal

    for (int i=0; i<numBck; i++){
      double mi = this->m(i);
      double s1i = this->S1(i);
      double s2i = this->S2(i);
      double lambdaStart_i, sigmaStart_i, bStart_i;
      if ( mi == 0 ) {
        lambdaStart_i = 0.;
        sigmaStart_i = 0.;
        bStart_i = 0.;            // avoid boundary?
        freePar[3*i+1] = true;    //  b
        freePar[3*i+2] = false;   // lambda
        freePar[3*i+3] = false;   // sigma
      }
      else if ( mi == 1 ) {
        lambdaStart_i = s1i;         // since only one event
        sigmaStart_i = abs(s1i);     // since E[l] = 0
        double w_i = exp(s1i);     // since only one event
        bStart_i = w_i / this->tau(i);
        freePar[3*i+1] = true;     // b
        freePar[3*i+2] = true;     // lambda
	freePar[3*i+3] = false;    // sigma
        // freePar[3*i+3] = true;   // KLUDGE -- try this
      }
      else {
        lambdaStart_i = s1i/mi;
        sigmaStart_i = sqrt(s2i/(mi-1) - s1i*s1i/(mi*(mi-1)));
        double omega_i = exp(lambdaStart_i + 0.5*sigmaStart_i*sigmaStart_i);
        bStart_i = this->m(i)*omega_i/this->tau(i);
        freePar[3*i+1] = true;   // b
        freePar[3*i+2] = true;   // lambda
        freePar[3*i+3] = true;   // sigma
      }
      parVec[3*i+1] = bStart_i;
      parVec[3*i+2] = lambdaStart_i;
      parVec[3*i+3] = sigmaStart_i;
    }        // loop over background components

  }

  else if ( option >= 100 && option < 200 ) {

    for (int i=0; i<numBck; i++){
      double mi = this->m(i);
      double s1i = this->S1(i);
      double s2i = this->S2(i);
      double omegaStart_i, sigmaStart_i, bStart_i;
      if ( mi == 0 ) {
        omegaStart_i = 1.;
        sigmaStart_i = 0.;
        bStart_i = 0.;            // avoid boundary?
        freePar[3*i+1] = true;    //  b
        freePar[3*i+2] = false;   // omega
        freePar[3*i+3] = false;   // sigma
      }
      else if ( mi == 1 ) {
        double w_i = s1i;     // since only one event
        omegaStart_i = w_i;
        sigmaStart_i = abs(w_i-1.);     // since E[w] = 1
        bStart_i = w_i / this->tau(i);
        freePar[3*i+1] = true;     // b
        freePar[3*i+2] = true;     // omega
        freePar[3*i+3] = false;    // sigma
      }
      else {
        omegaStart_i = s1i/mi;
        sigmaStart_i = sqrt(s2i/(mi-1) - s1i*s1i/(mi*(mi-1)));
        bStart_i = s1i / this->tau(i);
        freePar[3*i+1] = true;   // b
        freePar[3*i+2] = true;   // omega
        freePar[3*i+3] = true;   // sigma
      }
      parVec[3*i+1] = bStart_i;
      parVec[3*i+2] = omegaStart_i;
      parVec[3*i+3] = sigmaStart_i;
    }        // loop over background components

  }
  else if ( option >= 200 && option < 400 ) {

    for (int i=0; i<numBck; i++){
      double mi = this->m(i);
      double s1i = this->S1(i);
      double s2i = this->S2(i);
      double bStart_i;
      if ( mi == 0 ) {
        bStart_i = 0.;      // avoid boundary?
      }
      else {
        bStart_i = s1i / this->tau(i);
      }
      freePar[3*i+1] = true;
      freePar[3*i+2] = false;
      freePar[3*i+3] = false;
      parVec[3*i+1] = bStart_i;
      parVec[3*i+2] = 0.;        // dummy value
      parVec[3*i+3] = 0.;        // dummy value
    }          // loop over background components

  }         // selection of option

  // Initial mu is same for all options

  if ( freePar[0] ) {           // also fit mu
    double btot = 0.;
    for (int i=0; i<numBck; i++){
      btot += parVec[3*i + 1];
    }
    double muHat = (this->n() - btot)/this->s();
    parVec[0] = max(muHat, 0.);
  }


}