//////////////////////////////////////////////////////////////////////////
//                                                                      //
// GFMultiDmFitter                                                      //
//
//   Author:      Gero Flucke
//   Date:        June 10th, 2003
//   last update: $Date: 2005/11/15 17:27:52 $
//   by:          $Author: flucke $
//                                                                      //
// fitting several delta m histogramms using DstarDmFitter              //
//                                                                      //
//////////////////////////////////////////////////////////////////////////

#include "GFMultiDmFitter.h"
#include <iostream>
using namespace std;
#include "TH1.h"
#include "TH2.h"
#include "TF1.h"
#include "GFUtils/GFHistArray.h"
#include "GFUtils/GFHistManip.h"
#include "GFUtils/GFDstarRunPeriods.h"
#include "GFAnalysis/GFDstarCount.h"

#include "/afs/desy.de/user/f/flucke/h1/root/DstarDmFitter/DstarDmFitter.h"



ClassImp(GFMultiDmFitter)

 GFMultiDmFitter::GFMultiDmFitter()
{
  this->Init();
}

 GFMultiDmFitter::~GFMultiDmFitter()
{
  delete fFitter;
  delete fHists1D;   // really?
  delete fHists1DWc; // dito?
}

//__________________________________________________________________
 void GFMultiDmFitter::SetSmartMode(Bool_t s)
{
  fFitter->SetSmartMode(s);
}

//__________________________________________________________________
 void GFMultiDmFitter::Init()
{
  fFitter = new DstarDmFitter;
  fFitter->SetSmartMode(kFALSE); // fixme
  fStartUsqr = -1.; // < 0 indicates: ignore and use default from DstarDmFitter
  fFitModeAll = 0;
  fFitModeBins = 1011;

  fFitOrNotFlag = 0;
  fMergeNBins = 1;

  fOldParams[0] = fOldParams[1] = fOldParams[2] = fOldParams[3] = 0.;
  fOldParams[4] = 1.;

  fStartDmSignal = GFDstarRunPeriods::fgStartDmSignal;
  fEndDmSignal   = GFDstarRunPeriods::fgEndDmSignal;
  fStartDmBackgr = GFDstarRunPeriods::fgStartDmBackgr;
  fEndDmBackgr   = GFDstarRunPeriods::fgEndDmBackgr;

  fHists1D = NULL;
  fHists1DWc = NULL;
  fHistSignalBack = NULL;
}


//__________________________________________________________________
 TH1* GFMultiDmFitter::Fit(TH2* hist2D, TH1* histRef, TH2* hist2DWc, TH1* histWcRef)
{
  // hist2D has x-axis bins in delta m!
  // If histAll given it will be taken as reference hist for fixing parameters,
  // otherwise the projection of hist2D is taken.
  // This method allocates memory:
  //   - the returned histogram
  //   - the 1-D hists that can be accessed via   GFHistArray* GetHists1D()
  //     (but the array is owned by this class!)
  //   - dito for WC 1-D hists if hist2DWc is given (accessed via   GFHistArray* GetHists1DWc()
  //   - the histogram of the signal/background ratios, to be accessed via GetHistSignalBack()


  if(!hist2D) return NULL;
  if(hist2DWc && hist2D->GetNbinsY() != hist2DWc->GetNbinsY()){
    this->Error("Fit", "WC hist has %d instead of %d bins in Y", 
		hist2DWc->GetNbinsY(), hist2D->GetNbinsY());
    return NULL;
  }
  Bool_t setSmart = kTRUE;
  if(hist2DWc) setSmart = fFitter->SetSmartMode(kFALSE);


  TH1* refHistInt = NULL;
  if(histRef) refHistInt = histRef;
  else refHistInt = hist2D->ProjectionX("tmpRefHist", -1, -1, "e");
  TH1* refHistIntWc = NULL;
  if(histWcRef || !hist2DWc) refHistIntWc = histWcRef;
  else refHistIntWc = hist2DWc->ProjectionX("tmpRefHistWc", -1, -1, "e");

  if(fFitOrNotFlag == 0 || fFitOrNotFlag == 5) {
    this->FitAll(refHistInt, refHistIntWc);
  }
  if(fHists1D) delete fHists1D;

  fHists1D = GFHistManip::CreateHistArray1D(hist2D, hist2D->GetYaxis()->GetTitle(), NULL, 
					    Form("m%dou", fMergeNBins));

  if(fHists1DWc) delete fHists1DWc;
  fHists1DWc = (hist2DWc ? 
		GFHistManip::CreateHistArray1D(hist2DWc, hist2DWc->GetYaxis()->GetTitle(),
					       NULL, Form("m%dou", fMergeNBins)) : NULL);

  TH1* histNumDstar = GFHistManip::CreateHistYBins(hist2D, "", fMergeNBins);
  fHistSignalBack = GFHistManip::CreateHistYBins(hist2D, "", fMergeNBins);
  fHistSignalBack->SetTitle(Form("S/B;%s;S/B", fHistSignalBack->GetXaxis()->GetTitle()));

// if over-/underflow are empty, remove them:
  if (fHists1D->First()->Integral() == 0. && fHists1D->Last()->Integral() == 0.) {
    delete fHists1D->RemoveAt(0);
    delete fHists1D->Remove(fHists1D->Last());
    if (fHists1DWc) {
      delete fHists1DWc->RemoveAt(0);
      delete fHists1DWc->Remove(fHists1DWc->Last());
    }
  } else { // else merge them in first hist
    GFHistManip::Add2ndTo1st(fHists1D->First(), fHists1D->Last(), 1.);
    fHists1D->First()->SetTitle(Form("%s || %s", fHists1D->First()->GetTitle(),
				     fHists1D->Last()->GetTitle()));
    delete fHists1D->Remove(fHists1D->Last());
    if (fHists1DWc) {
      GFHistManip::Add2ndTo1st(fHists1DWc->First(), fHists1DWc->Last(), 1.);
      fHists1DWc->First()->SetTitle(Form("%s || %s", fHists1DWc->First()->GetTitle(),
					 fHists1DWc->Last()->GetTitle()));
      delete fHists1DWc->Remove(fHists1DWc->Last());
    }
  }

  Double_t sumw2 = 0.;
  for (Int_t i = 0; i < fHists1D->GetEntriesFast(); ++i) {
    if (!fHists1D->At(i)) continue; // empty slots due to removal of under/overflow
    TH1* hWc = (fHists1DWc ? fHists1DWc->At(i) : NULL);
    TArrayD nErrSB = this->FitSingle(fHists1D->At(i), hWc);
    histNumDstar->SetBinContent(i, nErrSB[0]);
    histNumDstar->SetBinError(i, nErrSB[1]);
    fHistSignalBack->SetBinContent(i, nErrSB[2]);
    sumw2 += nErrSB[1]*nErrSB[1];
  }
  fHists1D->Compress(); // remove empty slots
  if (fHists1DWc) fHists1DWc->Compress();

  histNumDstar->SetEntries(histNumDstar->Integral());
  Stat_t s[11];
  histNumDstar->GetStats(s);
  s[0] = histNumDstar->GetEntries(); 
  s[1] = sumw2;
  histNumDstar->PutStats(s);

  if(histRef != refHistInt) delete refHistInt;
  if(histWcRef != refHistIntWc) delete refHistIntWc;
//   fOldParams[0] = fOldParams[1] = fOldParams[2] = fOldParams[3] = 0.;
//   fOldParams[4] = 1.;

  if(hist2DWc && setSmart) fFitter->SetSmartMode(kTRUE);

  return histNumDstar;
}

//__________________________________________________________________
 TArrayD GFMultiDmFitter::FitAll(TH1* hist, TH1* histWc)
{
  // fit hist with fFitModeAll (histWc as possible WC-background hist),
  // return array {N(D*), err(ND*), S/B-ratio, scale to be applied to histWc} 
  // to get same normalisation of the background part 
  // if fit not successful, result is an empty array
  // (fOldParams is filled with result of this fit - if successful!)

  fFitter->ResetFitParameters();
  if(fStartUsqr >= 0.) fFitter->SetBackgrSqr(fStartUsqr);
  // reset old params, should be set in the following for flag 0 and 5:
  fOldParams[0] = fOldParams[1] = fOldParams[2] = fOldParams[3] = 0.;
  fOldParams[4] = 1.;

  TArrayD result(4);
  switch(fFitOrNotFlag){
  case 0:
    if(fFitter->Fit(fFitModeAll, hist, histWc)){
	result[0] = fFitter->GetNDstar();
	result[1] = fFitter->GetNDstarErr();
	result[2] = fFitter->GetSignalBack();
	if(histWc) result[3] = fFitter->GetWcScale();
	fOldParams[0] = fFitter->GetMean(); // store parameters of reference fit
	fOldParams[1] = fFitter->GetSigma();
	fOldParams[2] = fFitter->GetBackgrExp();
	fOldParams[3] = fFitter->GetBackgrSqr();
	fOldParams[4] = result[3];//fFitter->GetWcScale();
	break;
      }
      this->Warning("FitAll", "not fitted: count signal region and subtract background est.");
  case 1: // count signal region, but subtract background estimate (latter may be from hWc)
    result = this->CountSignal(hist, histWc, kTRUE);
    break;
  case 2: // count signal region
    result = this->CountSignal(hist, NULL, kFALSE);
    break;
  case 3: // count whole delta m
    for(Int_t j = 1; j <= hist->GetNbinsX(); ++j){
      result[0] += hist->GetBinContent(j);
      result[1] += (hist->GetBinError(j) * hist->GetBinError(j));
    }
    result[1] = TMath::Sqrt(result[1]);
    break;
  case 4: // subtract the fit to scaled background in signal region
    result = GFDstarCount::NDstarBackgrFuncSubtracted(hist, histWc, fStartDmSignal,fEndDmSignal,
						      fStartDmBackgr, fEndDmBackgr);
    break;
  case 5: // fit, subtract background function integral from counted signal (in 3-sigma)
    if(fFitter->Fit(fFitModeAll, hist, histWc)){
      const Int_t bin3SigmaL = hist->FindBin(fFitter->GetMean() - 3.* fFitter->GetSigma());
      const Int_t bin3SigmaU = hist->FindBin(fFitter->GetMean() + 3.* fFitter->GetSigma());
      result[0] = hist->Integral(bin3SigmaL, bin3SigmaU);
      result[1] = (result[0] ? TMath::Sqrt(result[0])/result[0] : 0.); // relative error so far
      TF1 backGr(*(fFitter->GetFittedBackgrFunc()));
      Double_t subtract = backGr.Integral(hist->GetXaxis()->GetBinLowEdge(bin3SigmaL), 
					  hist->GetXaxis()->GetBinUpEdge(bin3SigmaU));
      subtract /= fFitter->GetHistBinSize();
      result[0] -= subtract;
      result[1] *= result[0]; // relative error of SR count
      this->Warning("FitAll", "flag 5: relative error from count of SR, only 1 try to fit");
      result[2] = (subtract ? result[0]/subtract : 0.); // S/B
      if(histWc) result[3] = fFitter->GetWcScale();
      fOldParams[0] = fFitter->GetMean(); // store parameters of reference fit
      fOldParams[1] = fFitter->GetSigma();
      fOldParams[2] = fFitter->GetBackgrExp();
      fOldParams[3] = fFitter->GetBackgrSqr();
      fOldParams[4] = result[3];//fFitter->GetWcScale();
    } else {
      this->Warning("FitAll", "flag 5: fit did not succeed");
    }
    break;
  default:
    this->Warning("FitAll", "don't know fitOrNotFlag = %d", fFitOrNotFlag); 
  }
  
//   if(fFitter->Fit(fFitModeAll, hist, histWc)){
//     fOldParams[0] = fFitter->GetMean(); // store parameters of reference fit
//     fOldParams[1] = fFitter->GetSigma();
//     fOldParams[2] = fFitter->GetBackgrExp();
//     fOldParams[3] = fFitter->GetBackgrSqr();
//     fOldParams[4] = fFitter->GetWcScale();
    
//     result[0] = fFitter->GetNDstar();
//     result[1] = fFitter->GetNDstarErr();
//     result[2] = fFitter->GetSignalBack();
//     if(histWc) result[3] = fOldParams[4];
//   }
  return result;
}

//__________________________________________________________________
 TArrayD GFMultiDmFitter::FitSingle(TH1* hist, TH1* hWc) const
{
  // returning TArrayD of length 4: 
  //   N(D*), error, S/B-ratio, scale between hist and hWc fit normalisation

  const Int_t minEntries = 28; // for fit
  const Int_t minEntFull = 200; // for fit with other option than 1111

  TArrayD result(4); // N(D*), error, S/B-ratio, scale for WC hist if existing

  Double_t histEntriesD = hist->Integral();
  const Int_t histEntries = static_cast<Int_t>(histEntriesD);
  if(TMath::Abs(histEntriesD - histEntries) > 1.e-10) {
    // weighted entries/scaled hists: use effective number of filling actions
    Stat_t s[11];
    hist->GetStats(s);//now s[1] is sum of squares of weights, s[0] is sum of weights
    histEntriesD  = s[0]*s[0]/s[1]; // effective number of entries, cf. LOOK manual
  }

  switch(fFitOrNotFlag){
  case 0:
    if(histEntriesD >= minEntries){ // real fit
      const Int_t fitMode = (histEntriesD >= minEntFull) ? fFitModeBins : 1111;
      if(fitMode != fFitModeBins) {
 	fFitter->SetMean(fOldParams[0]); // need original parameters for fit!
 	fFitter->SetSigma(fOldParams[1]);
 	fFitter->SetBackgrExp(fOldParams[2]);
 	fFitter->SetBackgrSqr(fOldParams[3]);
 	fFitter->SetWcScale(fOldParams[4]);
	this->Warning("FitSingle", "%s with %.1f (effective) entries: fit with mode 1111n", 
		      hist->GetName(), histEntriesD);
      }
      Bool_t fitOK = fFitter->Fit(fitMode, hist, hWc);
      if((!fitOK || fFitter->GetBackgrExp() < 0.) && fitMode != 1111){
	if(!fitOK) this->Warning("FitSingle", "fit not OK, try mode 1111");
	else {
	  this->Warning("FitSingle","U_exp = %.3f, fit with mode 1111",fFitter->GetBackgrExp());
	}
 	fFitter->SetMean(fOldParams[0]); // need original parameters for next fit!
 	fFitter->SetSigma(fOldParams[1]);
 	fFitter->SetBackgrExp(fOldParams[2]);
 	fFitter->SetBackgrSqr(fOldParams[3]);
 	fFitter->SetWcScale(fOldParams[4]);
	fitOK = fFitter->Fit(1111, hist, hWc);
      }
      if(fitOK){
	result[0] = fFitter->GetNDstar();
	result[1] = fFitter->GetNDstarErr();
	result[2] = fFitter->GetSignalBack();
	if(hWc) result[3] = fFitter->GetWcScale();
	break;
      }
    }
    this->Warning("FitSingle", "%s with %.1f (effective) entries not fitted:ncount signal "
		  "region and subtract background estimation!", hist->GetName(), histEntriesD);
  case 1: // count signal region, but subtract background estimate (latter may be from hWc)
    result = this->CountSignal(hist, hWc, kTRUE);
    break;
  case 2: // count signal region
    result = this->CountSignal(hist, NULL, kFALSE);
    break;
  case 3: // count whole delta m
    for(Int_t j = 1; j <= hist->GetNbinsX(); ++j){
      result[0] += hist->GetBinContent(j);
      result[1] += (hist->GetBinError(j) * hist->GetBinError(j));
    }
    result[1] = TMath::Sqrt(result[1]);
//     result[0] = hist->Integral();
//     Stat_t stats[11];
//     hist->GetStats(stats);
//     result[1] = TMath::Sqrt(stats[1]);
    break;
  case 4: // subtract the fit to scaled background in signal region
    result = GFDstarCount::NDstarBackgrFuncSubtracted(hist, hWc, fStartDmSignal, fEndDmSignal, 
						      fStartDmBackgr, fEndDmBackgr);
    break;
  case 5: // fit, subtract background function integral from counted signal (in 3-sigma)
    if(fFitter->Fit(fFitModeBins, hist, hWc)){
      const Int_t bin3SigmaL = hist->FindBin(fFitter->GetMean() - 3.* fFitter->GetSigma());
      const Int_t bin3SigmaU = hist->FindBin(fFitter->GetMean() + 3.* fFitter->GetSigma());
      result[0] = hist->Integral(bin3SigmaL, bin3SigmaU);
      result[1] = (result[0] ? TMath::Sqrt(result[0])/result[0] : 0.); // relative error so far
      TF1 backGr(*(fFitter->GetFittedBackgrFunc()));
      Double_t subtract = backGr.Integral(hist->GetXaxis()->GetBinLowEdge(bin3SigmaL), 
					  hist->GetXaxis()->GetBinUpEdge(bin3SigmaU));
      subtract /= fFitter->GetHistBinSize();
      result[0] -= subtract;
      result[1] *= result[0]; // relative error of SR count
      this->Warning("FitSingle", "flag 5: relative error from count of SR, only 1 try to fit");
      result[2] = (subtract ? result[0]/subtract : 0.); // S/B
      if(hWc) result[3] = fFitter->GetWcScale();
    } else {
      this->Warning("FitSingle", "flag 5: fit did not succeed");
    }
    break;
  default:
    this->Warning("FitSingle", "don't know fitOrNotFlag = %d", fFitOrNotFlag); 
  }
  
  return result;
}

//__________________________________________________________________
 TArrayD GFMultiDmFitter::CountSignal(TH1* hist, TH1* hWc, Bool_t subtractBack) const
{
  // returning array of length 4: first 2 are number of D* in signal region and error
  // - if hWc != NULL: 
  //   subtract the entries in signal region (i.e. fStartDmSignal -> fEndDmSignal) 
  //   from the signal region of 'hist'. Scale the number from WC hist like the ratio
  //   of entries in background region (fStartDmBackgr -> fEndDmBackgr). 
  //   Assuming same binning for both hists!
  // - if subtractBack == kTRUE && hWc == NULL:
  //   subtracts 
  //   (mean bin content from fStartDmBackgr to fEndDmBackgr) times (number of signal bins)
  //   from number of entries in signal region (fStartDmSignal to fEndDmSignal)
  //   makes only sense for low statistics distributions with few entries
  //   in the background region (otherwise we get sometimes negative results)
  // - if subtractBack == kFALSE && hWc == NULL:
  //   just integrate full signal region (i.e. fStartDmSignal -> fEndDmSignal 
  // 
  // third position of array is ratio signal/(subtracted background) 
  // (empty if subtractBack == kFALSE && hWc == NULL)
  // fourth position is scale applied to hWc (empty if hWc == NULL)

  TArrayD result(4);
  if(!hist) return result;

  const Int_t firstSignalBin = hist->FindBin(fStartDmSignal);
  const Int_t lastSignalBin = hist->FindBin(fEndDmSignal);
  const Double_t numSignal = hist->Integral(firstSignalBin, lastSignalBin);

  const Int_t firstBackgrBin = hist->FindBin(fStartDmBackgr);
  const Int_t lastBackgrBin = hist->FindBin(fEndDmBackgr);

  if(hWc){
    const Double_t numBackGr   = hist->Integral(firstBackgrBin, lastBackgrBin);
    const Double_t numBackGrWc = hWc->Integral(firstBackgrBin, lastBackgrBin);
    const Double_t wcScale = (numBackGrWc ? numBackGr/numBackGrWc : 1.);
    const Double_t numSignalWc = hWc->Integral(firstSignalBin, lastSignalBin);
    result[0] = numSignal - numSignalWc * wcScale;
    const Double_t variance = numSignal + 3. * numSignalWc * wcScale; // cf. logbook 18.4.2004
    result[1] = (variance ? TMath::Sqrt(variance) : 0.);
    result[2] = ((numSignalWc * wcScale) ? result[0]/(numSignalWc * wcScale) : 0.);
    result[3] = wcScale;
  } else if(subtractBack){ 
    const Double_t numBackgr = hist->Integral(firstBackgrBin, lastBackgrBin);
    const Double_t backGrWidth = (lastBackgrBin - firstBackgrBin + 1);//+1: num. of bins
    const Double_t sigByBackWidth = (lastSignalBin - firstSignalBin + 1)/backGrWidth;
    const Double_t backGr = numBackgr * sigByBackWidth;
    result[0] = numSignal - backGr;
//     result[1] = (result[0] > 0. ? TMath::Sqrt(result[0]) : 0.);
// sigma^2 = sigma_sig^2 + (sigByBackWidth * sigma_back)^2
    result[1] = TMath::Sqrt(numSignal + numBackgr * sigByBackWidth*sigByBackWidth);
    result[2] = (backGr ? result[0]/backGr : 0.);
  } else {
    result[0] = numSignal;
    for(Int_t i = firstSignalBin; i <= lastSignalBin; ++i){
      const Stat_t error = hist->GetBinError(i);
      result[1] += (error * error);
    }
    result[1] = TMath::Sqrt(result[1]);
  }

  return result;  
}