// File:  runSigCalc.cc
// Glen Cowan
// RHUL Physics
// 14 April 2008

// Simple driver program for the function getSignificance, which 
// returns the significance with which one can reject a hypothesized
// signal strength parameter mu.  (For discovery, one rejects mu = 0.)

// Input:  runSigCalc reads in a text file with the input data.  
// Lines beginning with # are comments.
// The first (non-comment) line contains the target luminosity.
// The second (non-comment) line contains the accepted number of
// signal events from MC and the luminosity of the signal MC sample.
//
// On subsequent lines the user lists for each background component the
// accepted number of background events m and the effective luminosity 
// L_MC used to generate the background sample.

#include <iostream>
#include <fstream>
#include <cstdlib>
#include <string>
#include <sstream>
#include <vector>
#include <TMath.h>
#include <TFile.h>
#include <TH1D.h>
#include <TRandom3.h>
#include "SigCalc.h"
#include "getDiscoverySignificance.h"
#include "getExclusionSignificance.h"

using namespace std;

int main(int argc, char **argv) {

  TFile* histFile = new TFile("SigCalc.root", "recreate");

  TH1D* h_q0 = new TH1D("h_q0", "q0", 200, 0., 100.);
  TH1D* h_muHat = new TH1D("h_muHat", "muHat", 200, -10., 10.);
  TH1D* h_lambdaHat = new TH1D("h_lambdaHat", "lambdaHat", 200, -5, 5);
  TH1D* h_sigmaHat = new TH1D("h_sigmaHat", "sigmaHat", 200, 0., 10.);
  TH1D* h_omegaHat = new TH1D("h_omegaHat", "omegaHat", 200, 0., 2.);

// versions of histograms with exponentially distributed weights ("ew")

  TH1D* h_q0_ew = new TH1D("h_q0_ew", "q0", 200, 0., 100.);
  TH1D* h_muHat_ew = new TH1D("h_muHat_ew", "muHat", 200, -10., 10.);
  TH1D* h_lambdaHat_ew = new TH1D("h_lambdaHat_ew", "lambdaHat", 200, -5, 5);
  TH1D* h_sigmaHat_ew = new TH1D("h_sigmaHat_ew", "sigmaHat", 200, 0., 10.);
  TH1D* h_omegaHat_ew = new TH1D("h_omegaHat_ew", "omegaHat", 200, 0., 2.);

  int seed = 12348;
  TRandom3* ran = new TRandom3(seed);

// Read input file

  string inputFileName;
  if ( argc == 1 ) {
    cout << "Enter name of input file: ";
    cin >> inputFileName;
    cin.ignore();         // clear buffer
  }
  else {
    inputFileName = argv[1];
  }

  ifstream f;
  f.open( inputFileName.c_str() );
  if ( f.fail() ){
    cout << "Sorry, couldn't open input file" << endl;
    exit(1);
  }

  int option;
  double nObs, s, mu;
  vector<double> mObs;
  vector<double> tau;
  vector<double> S1Obs;
  vector<double> S2Obs;
  vector<double> bVec;
  vector<double> aVec;
  vector<double> xiVec;
  vector<double> lambdaVec;
  vector<double> sigmaVec;

  istringstream instream;       // Declare an input string stream
  string line;
  int lineNum = 0;
  while ( getline(f, line) ) {
    bool useLine = !( line.substr(0,1) == "#" || line.substr(0,1) == "!");
    if ( useLine ) {
      instream.clear();          // Reset from possible previous errors
      instream.str(line);        // Use line as source of input
      lineNum++;
      if ( lineNum == 1 ) {
        instream >> option;
      }

      if ( option%100 == 0 ) {  // rest of info is a data set

        if ( lineNum == 2 ) {
          instream >> nObs >> s;
        }
        else if ( lineNum > 2 ) {
          double mObs_i, tau_i, S1Obs_i, S2Obs_i, wsum, w2sum;
          instream >> mObs_i >> tau_i >> S1Obs_i >> S2Obs_i >> wsum >> w2sum;
          mObs.push_back(mObs_i);
          tau.push_back(tau_i);
          S1Obs.push_back(S1Obs_i);
          S2Obs.push_back(S2Obs_i);
        }

      }     // option == 0
      else if ( option%100 == 1 ){    // rest of info are parameters

        if ( lineNum == 2 ) {
          instream >> mu >> s;
        }
        else if ( lineNum > 2 ) {
          double tau_i, b_i, a_i, xi_i;
          instream >> b_i >> tau_i >> a_i >> xi_i;
          tau.push_back(tau_i);
          bVec.push_back(b_i);
          aVec.push_back(a_i);
          xiVec.push_back(xi_i);
        }

      }      // option == 1
    }       // useline
  }         // while getline

  if ( option%100 == 1 ){         // generate data set
    SigCalc* sc = new SigCalc (s, tau, mu, bVec, aVec, xiVec, ran, option);
    nObs = sc->n();
    for (int i=0; i<bVec.size(); i++){
      mObs.push_back(sc->m(i));
      S1Obs.push_back(sc->S1(i));
      S2Obs.push_back(sc->S2(i));
      cout << tau[i] << "  " << bVec[i] << "  " << 
	aVec[i] << "  " << xiVec[i] << endl;
      cout << "nObs = " << nObs << endl;
      cout << "m, S1, S2 = " << 
              mObs[i] << "  " << S1Obs[i] << "  " << S2Obs[i] << endl;
      if ( option == 1 ) {       // w ~ log normal
        double z = 1. - exp(-aVec[i]/xiVec[i]);
        double lambda = log (aVec[i]/(xiVec[i]*z)) - 0.5*aVec[i]/xiVec[i];
        double sigma_l = (aVec[i]/xiVec[i]) / sqrt(12.);
        cout << "corresponding lambda, sigma_l = " << lambda 
             << "  " << sigma_l << endl;
        lambdaVec.push_back(lambda);
        sigmaVec.push_back(sigma_l);
      }
      else if ( option == 101 || option == 201 ) {    // w ~ Gauss
        double z = 1. - exp(-aVec[i]/xiVec[i]);
        double sigma_w = (0.5*aVec[i]/xiVec[i]) *
	  (1. - exp(-2*aVec[i]/xiVec[i])) / (z*z)  -  1.;
        double omega = 1.;

	// KLUDGE, use lambdaVec to hold the omega values.
	// Here sigma means sigma_w

       lambdaVec.push_back(omega);
       sigmaVec.push_back(sigma_w);

      }
    }
    delete sc;
  }

  int numBkg = mObs.size();
  cout << "numBkg = " << numBkg << endl;

  cout << "n = " << nObs << endl;
  cout << "s = " << s << endl;

  // Get prelim. estimate of btot.  Different prescription depending
  // on whether w ~ log-normal or Gaussian

  double btotHat = 0.;
  for (int i=0; i<numBkg; i++){ 

    double bHat_i = 0;
    if ( option < 100 ) {       // w ~ log-normal

      double lambdaHat_i;
      double sigma2Hat_i;
      if ( mObs[i] == 0 ) {
        sigma2Hat_i = 1.;
        lambdaHat_i = 1.;
      }
      else if ( mObs[i] == 1 ) {
        lambdaHat_i = S1Obs[i]/mObs[i]; 
        double lambda0 = 0;              // nominal ln w
        sigma2Hat_i = (lambdaHat_i - lambda0)*(lambdaHat_i - lambda0);
      }
      else {
        lambdaHat_i = S1Obs[i]/mObs[i]; 
        sigma2Hat_i = 
            (1./(mObs[i]-1)) * (S2Obs[i] - S1Obs[i]*S1Obs[i]/mObs[i]);
      }
      double omegaHat_i = exp(lambdaHat_i + 0.5*sigma2Hat_i);
      bHat_i = mObs[i]*omegaHat_i/tau[i];

    }
    else if ( option >= 100 && option < 200 ) {     // w ~ Gaussian
        bHat_i = S1Obs[i]/tau[i];
    }
    else if ( option >= 200 && option < 300 ) {     // bHat ~ Gaussian
        bHat_i = S1Obs[i]/tau[i];
    }
    btotHat += bHat_i;
    cout << "prelim. i,  bHat_i = " << i << "  " << bHat_i << endl;
  }
  cout << "prelim. estimated total background = " << btotHat << endl;
  cout << endl;

  double ZDisc = getDiscoverySignificance(nObs, s, mObs, tau, 
					  S1Obs, S2Obs, option);
  cout << "Asymptotic discovery significance Z   = " << ZDisc << endl;
  cout << "Corresponding p-value  = " << 1. - TMath::Freq(ZDisc) << endl;
  cout << endl;

  double muTest;
  // muTest = 1.0;      // for defining lambda(mu)
  // double ZExcl = getExclusionSignificance(muTest, nObs, s, mObs, tau, 
  //                S1Obs, S2Obs);
  // cout << "Asymptotic exclusion significance for mu = " << muTest 
  //      << " is Z = " << ZExcl << endl;
  // cout << "Corresponding p-value = " << 1. - TMath::Freq(ZExcl) << endl;

// Now redo discovery analysis without using asymptotic formulae.
// First step is to get HatHat(0) values for nuisance parameters

  muTest = 0.;

  SigCalc* sc = new SigCalc(nObs, s, mObs, tau, S1Obs, S2Obs, option);
  double muHatObs = (nObs - btotHat) / s;       // start value in fit
  vector<double> bHatObs(numBkg);
  vector<double> bHatHatObs(numBkg);
  vector<double> lambdaHatObs(numBkg);
  vector<double> lambdaHatHatObs(numBkg);
  vector<double> sigmaHatObs(numBkg);
  vector<double> sigmaHatHatObs(numBkg);
  double q0Obs = sc->q0(muHatObs, bHatObs, bHatHatObs,
                        lambdaHatObs, lambdaHatHatObs, 
                        sigmaHatObs, sigmaHatHatObs);
  cout << "muHatObs = " << muHatObs << endl;
  for (int i=0; i<numBkg; i++){
    cout << "i, bHatObs, lambdaHatObs, sigmaHatObs = " << i << 
      "  " <<  bHatObs[i] << "  " << lambdaHatObs[i] << "  " << 
      sigmaHatObs[i] << endl; 
    cout << "i, bHatHatObs, lambdaHatHatObs, sigmaHatHatObs = " << 
      i << "  " << bHatHatObs[i] << "  " << lambdaHatHatObs[i] << "  " << 
      sigmaHatHatObs[i] << endl; 
  }
  cout << "q0Obs    = " << q0Obs << endl;
  delete sc;

  double bHatHatObsTot = 0.;
  for (int j=0; j<numBkg; j++){
    bHatHatObsTot += bHatHatObs[j];
  }
  cout << "bHatHatObsTot = " << bHatHatObsTot << endl;

// Generate data according to the w ~ (log-)normal model.
// In the asymptotic limit Wilks' theorem should apply.

// KLUDGE -- if this was an MC generated data set, take
// the point in parameter space to correspond to the 
// input values of a and xi, not the HatHat values.  
// Implement by reassigning the HatHat values...

  if ( option%100 == 1 ){
    bHatHatObsTot = 0.;
    for (int j=0; j<numBkg; j++){
      bHatHatObs[j] = bVec[j];
      bHatHatObsTot += bHatHatObs[j];
      if ( option < 300 ) {
        lambdaHatHatObs[j] = lambdaVec[j];
        sigmaHatHatObs[j] = sigmaVec[j];    
        // cout << "New HatHat: b, lam, sig obs = " << bHatHatObs[j] << "  " 
        //    << lambdaHatHatObs[j] << "  " << sigmaHatHatObs[j] << endl;
      }
    }
    //     cout << "New bHatHatObsTot = " << bHatHatObsTot << endl;
  }             // option%100 == 1  (data generated)
    //     cout << "New bHatHatObsTot = " << bHatHatObsTot << endl;

  int numExp = 50000;
  int numRej = 0;
  for (int i=0; i<numExp; i++) {

    int nGen = ran->Poisson(muTest*s + bHatHatObsTot);
    double nd = static_cast<double>(nGen);
    vector<double> md;
    md.clear();
    vector<double> S1Gen;
    S1Gen.clear();
    vector<double> S2Gen;
    S2Gen.clear();
    for (int j=0; j<numBkg; j++){
      double lhho = lambdaHatHatObs[j];
      double shho = sigmaHatHatObs[j];

      // KLUDGE -- if w ~ Gauss, lambdaHatHatObs really stores omega

      double omega_j;
      if ( option < 100 ) {
        omega_j = exp(lhho + 0.5*shho*shho);
      }
      else if ( option >= 100 && option < 300 ){
        omega_j = lambdaHatHatObs[j];
      }
      else if ( option >= 300 && option < 400 ) {
        omega_j = 1.;
      }
      int mGen = ran->Poisson(tau[j]*bHatHatObs[j]/omega_j);
      md.push_back(static_cast<double>(mGen));

      // Replace this later with S1 ~ Gauss(n*lambda, sqrt(n)*sigma)
      // and S2 ~ chi2(n-1) * sigma2/(m-1)

      double S1 = 0.;
      double S2 = 0.;
      for (int k=0; k<mGen; k++){
        double l = ran->Gaus(lambdaHatHatObs[j], sigmaHatHatObs[j]);
        S1 += l;
        S2 += l*l;
      }
      S1Gen.push_back(S1);
      S2Gen.push_back(S2);
      // cout << "m, S1, S2... " << mGen << "  " << S1 << "  " << S2 << endl;
    }

    SigCalc* sc = new SigCalc(nd, s, md, tau, S1Gen, S2Gen, option);
    double muHat = (nd - bHatHatObsTot) / s;     // start value in fit
    vector<double> bHat(numBkg);
    vector<double> bHatHat(numBkg);
    vector<double> lambdaHat(numBkg);
    vector<double> lambdaHatHat(numBkg);
    vector<double> sigmaHat(numBkg);
    vector<double> sigmaHatHat(numBkg);
    double q0 = sc->q0(muHat, bHat, bHatHat, lambdaHat, lambdaHatHat, 
                       sigmaHat, sigmaHatHat);
    delete sc;
    h_q0->Fill(q0);
    h_muHat->Fill(muHat);
    h_lambdaHat->Fill(lambdaHat[0]);
    h_sigmaHat->Fill(sigmaHat[0]);

    // KLUDGE -- for w ~ Gauss, lambdaHat, lambdaHatHat actually
    // holds omegaHat, omegaHatHat.
    double omegaHat;
    if ( option < 100 ) {
      omegaHat = exp(lambdaHat[0] + 0.5*sigmaHat[0]*sigmaHat[0]);
    }
    else if ( option >= 100 && option < 200 ) {
      omegaHat = lambdaHat[0];
    }
    h_omegaHat->Fill(omegaHat);
    // cout << "runSigCalc: q0, q0Obs = " << q0 << "  " << q0Obs << endl;
    if (q0 >= q0Obs) {
      numRej++;
    }

    if ( i%(numExp/100) == 0 ) {
      cout << "generated experiment " << i << endl;
    }

  }

  cout << numRej << " events rejected out of " << numExp << 
                    " generated." << endl;
  double p0 = static_cast<double>(numRej)/static_cast<double>(numExp);
  double Z0 = TMath::NormQuantile(1. - p0);

  cout << "MC Z0 (data ~ log-normal weights) = " << Z0 << endl;

// Now do this again but generate the data according to the
// "real" model with a and xi (not quite w ~ log-normal).
// Only do this if a and xi available, i.e., data was generated

  if ( option%100 == 1 ) {    

    numRej = 0;
    for (int i=0; i<numExp; i++) {

      int nGen = ran->Poisson(muTest*s + bHatHatObsTot);
      double nd = static_cast<double>(nGen);
      vector<double> md;
      md.clear();
      vector<double> S1Gen;
      S1Gen.clear();
      vector<double> S2Gen;
      S2Gen.clear();
      for (int j=0; j<numBkg; j++){
        double lhho = lambdaHatHatObs[j];
        double shho = sigmaHatHatObs[j];

	// KLUDGE -- for w ~ Gauss, lambdaHatHatObs actually holds omega
        double omega_j;
        if ( option < 100 ){
          omega_j = exp(lhho + 0.5*shho*shho);
        }
        else if ( option >= 100 && option < 300) {
          omega_j = lambdaHatHatObs[j];
        }
        int mGen = ran->Poisson(tau[j]*bHatHatObs[j]/omega_j);
        md.push_back(static_cast<double>(mGen));

        double S1 = 0.;
        double S2 = 0.;
        for (int k=0; k<mGen; k++){
          double r = ran->Rndm();             // uniform in ]0,1]
          double x = aVec[j]*r;
          double w = aVec[j] * (1./xiVec[j])*exp(-x/xiVec[j]) 
                     / (1 - exp(-aVec[j]/xiVec[j]));
          double lnw = log(w); 
          if ( option < 100 ) {
            S1 += log(w);
            S2 += lnw*lnw;
          }
          else if ( option >= 100 && option < 300 ){
            S1 += w;
            S2 += w*w;
          }
          //      cout << w << "  " << lnw << endl;
        }

        S1Gen.push_back(S1);
        S2Gen.push_back(S2);
        // cout << "m, S1, S2... " << mGen << "  " << S1 << "  " << S2 << endl;
      }

      SigCalc* sc = new SigCalc(nd, s, md, tau, S1Gen, S2Gen, option);
      double muHat = (nd - bHatHatObsTot) / s;     // start value in fit
      vector<double> bHat(numBkg);
      vector<double> bHatHat(numBkg);
      vector<double> lambdaHat(numBkg);
      vector<double> lambdaHatHat(numBkg);
      vector<double> sigmaHat(numBkg);
      vector<double> sigmaHatHat(numBkg);
      double q0 = sc->q0(muHat, bHat, bHatHat, lambdaHat, lambdaHatHat, 
                         sigmaHat, sigmaHatHat);
      delete sc;
      h_q0_ew->Fill(q0);
      h_muHat_ew->Fill(muHat);
      h_lambdaHat_ew->Fill(lambdaHat[0]);
      h_sigmaHat_ew->Fill(sigmaHat[0]);

   // KLUDGE -- for w ~ Gauss, lambdaHat, lambdaHatHat actually
    // holds omegaHat, omegaHatHat.
      double omegaHat_ew;
      if ( option < 100 ) {
        omegaHat_ew = exp(lambdaHat[0] + 0.5*sigmaHat[0]*sigmaHat[0]);
      }
      else if ( option >= 100 && option < 200 ) {
        omegaHat_ew = lambdaHat[0];
      }
      h_omegaHat_ew->Fill(omegaHat_ew);

      // cout << "runSigCalc: q0, q0Obs = " << q0 << "  " << q0Obs << endl;
      if (q0 >= q0Obs) {
        numRej++;
      }

      if ( i%(numExp/10) == 0 ) {
        cout << "generated experiment " << i << endl;
      }

    }

    cout << "with exponentially distributed weights: " << endl;
    cout << numRej << " events rejected out of " << numExp << 
                      " generated." << endl;
    p0 = static_cast<double>(numRej)/static_cast<double>(numExp);
    Z0 = TMath::NormQuantile(1. - p0);
    cout << "MC Z0 (data ~ exponential weights) = " << Z0 << endl;

  }     // option == 1 (w ~ exp)

  cout << tau[0] << "     " << ZDisc << "     " << Z0 << endl;

  histFile->Write();
  histFile->Close();
  return 0;

}