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

// Uses TMinuit to fit the MLEs muHat, bHat and conditional MLEs bHatHat
// used in profile likelihood ratio.
// Inputs:     
// mu    signal strength parameter (0 = background only, 1 = nominal).
// n     number of events seen in data, e.g. generate from Poisson(mu*s+b)
// s     expected number of signal events
// m     vector of numbers of events seen in subsidiary measurements.
// tau   defined by m_i ~ Poisson(tau_i*b_i), e.g., ratio of MC lumonisity
//       to that of data sample.

#include <iostream>
#include <fstream>
#include <sstream>
#include <vector>
#include <cstdlib>
#include <cmath>

#include "SigCalc.h"
#include "fitPar.h"
#include <TMinuit.h>

SigCalc* scGlobal;   // needs to be global to communicate with fcn (below)
bool debug;

int fitPar (SigCalc* sc, vector<bool> freePar, vector<double>& parVec){
  
// set up start values, step sizes, etc. for fit

  scGlobal = sc;                  // communicate to fcn via global
  int npar = parVec.size();

  const int maxpar = 100;
  double par[maxpar];
  string parName[maxpar];
  double stepSize[maxpar];
  double minVal[maxpar];
  double maxVal[maxpar];

  for (int i=0; i<npar; i++){
    par[i] = parVec[i];
  }

  for (int i=0; i<npar; i++){    // some of these overridden below
    stepSize[i] = 0.1;
    minVal[i] = 0.;
    maxVal[i] = 100000.;
  }
  minVal[0] = -100000.;             // allow muHat to be negative
  for (int i=0; i<(npar-1)/3; i++){
    minVal[3*i+2] = -100000.;       // allow lambdaHat to be negative
    minVal[3*i+3] = 1.e-9;          // minimum sigma
  }


  parName[0] = "mu";
  int numBkg = (npar-1)/3;
  for (int i=0; i<numBkg; i++){
    stringstream ss;
    string bckNumber;
    ss << i;
    ss >> bckNumber;
    parName[3*i + 1] = "b" +  bckNumber;
    parName[3*i + 2] = "lambda" + bckNumber;
    parName[3*i + 3] = "sigma" + bckNumber;
  }

  for (int i=0; i<npar; i++){
    if ( par[i] != 0 ) {
      stepSize[i] = abs( par[i] ) * 0.1;
    }
    else {
      stepSize[i] = 0.1;
    }
  }

// Create the TMinuit object and define parameters

  TMinuit* minuit = new TMinuit(npar);

  minuit->SetPrintLevel(-1);
  // minuit->SetPrintLevel(0);
  int ierr = 0;
  minuit->mnexcm("SET NOWarnings",0,0,ierr);

  minuit->SetFCN(fcn);
  for (int i=0; i<npar; i++){
    minuit->DefineParameter(i, parName[i].c_str(), 
    par[i], stepSize[i], minVal[i], maxVal[i]);
  }

  for (int i=0; i<npar; i++){
    // cout << "i, freepar = " << i << "  " << freePar[i] << endl;
    if ( !freePar[i] ){
      minuit->FixParameter(i);
    }
  }

// do the fit and extract results
   

// KLUDGE try first fixing b, then release

  minuit->FixParameter(1);
  int status = minuit->Command("MIGRAD");
  minuit->Release(1);
  status = minuit->Command("MIGRAD");

  if ( status != 0 ) {
    debug = true;
    status = minuit->Command("SIMPLEX");
    status = minuit->Command("MIGRAD");
    debug = false;

    /*
    cout << "status = " << status << endl;
    cout << "n = " << sc->n() << endl;
    cout << "m = " << sc->m(0) << endl;
    cout << "S1 = " << sc->S1(0) << endl;
    cout << "S2 = " << sc->S2(0) << endl;
    */

  }

  parVec.clear();
  int nFreePar = minuit->GetNumFreePars();
  double fitpar[nFreePar], err[nFreePar];
  for (int i=0; i<npar; i++){
    minuit->GetParameter(i, fitpar[i], err[i]);
    parVec.push_back(fitpar[i]);
  }

  delete minuit;
  return status;     // 0 = normal completion, 4 = MIGRAD not converged

}

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

// fcn must be non-member function, uses global SigCalc object scGlobal

void fcn(int& npar, double* deriv, double& f, double par[], int flag){

  int numBck = scGlobal->numBck();
  double mu = par[0];
  vector<double> bVec;
  vector<double> lambdaVec;
  vector<double> sigmaVec;
  for (int i=0; i<numBck; i++){
    bVec.push_back(par[3*i+1]);
    lambdaVec.push_back(par[3*i+2]);
    sigmaVec.push_back(par[3*i+3]);
  }
  f = -2.*scGlobal->lnL(mu, bVec, lambdaVec, sigmaVec);

  /*
  double n = scGlobal->n();
  double m = scGlobal->m(0);
  double S1 = scGlobal->S1(0);
  double S2 = scGlobal->S2(0);
  cout << "data: " << n << "  " << m << "  " << S1 << "  " << S2 << endl;
  cout << " pars, f: " << par[0] << "  " << par[1] << "  " << par[2] << "  " 
        << par[3] << "  " << f << endl;
  */

  // if ( debug ) {
  // for (int i=0; i<npar; i++){
  //   cout << "i, par = " << i << "  " << par[i] << endl;
  // }
  // cout << "f = " << f << endl;
  // cout << endl;
  // }

}