/* ======================================================================== *\
!
! *
! * This file is part of MARS, the MAGIC Analysis and Reconstruction
! * Software. It is distributed to you in the hope that it can be a useful
! * and timesaving tool in analysing Data of imaging Cerenkov telescopes.
! * It is distributed WITHOUT ANY WARRANTY.
! *
! * Permission to use, copy, modify and distribute this software and its
! * documentation for any purpose is hereby granted without fee,
! * provided that the above copyright notice appear in all copies and
! * that both that copyright notice and this permission notice appear
! * in supporting documentation. It is provided "as is" without express
! * or implied warranty.
! *
!
!
!   Author(s): Robert Wagner   <magicsoft@rwagner.de> 10/2002
!              Wolfgang Wittek <wittek@mppmu.mpg.de>  02/2003
!
!   Copyright: MAGIC Software Development, 2000-2002
!
!
\* ======================================================================== */

/////////////////////////////////////////////////////////////////////////////
//                                                                         //
//  MPadSchweizer                                                          //
//                                                                         //
//  This task applies padding such that for a given pixel and for a given  //
//  Theta bin the resulting distribution of the pedestal sigma is identical//
//  to the distributions given by fHSigmaPixTheta and fHDiffPixTheta.      //
//                                                                         //
//  The number of photons, its error and the pedestal sigmas are altered.  //
//  On average, the number of photons added is zero.                       //
//                                                                         //
//  The formulas used can be found in Thomas Schweizer's Thesis,           //
//                                    Section 2.2.1                        //
//                                                                         //
//  There are 2 options for the padding :                                  //
//                                                                         //
//  1) fPadFlag = 1 :                                                      //
//     Generate first a Sigmabar using the 2D-histogram Sigmabar vs. Theta //
//     (fHSigmaTheta). Then generate a pedestal sigma for each pixel using //
//     the 3D-histogram Theta, pixel no., Sigma^2-Sigmabar^2               //
//     (fHDiffPixTheta).                                                   //
//                                                                         //
//     This is the preferred option as it takes into account the           //
//     correlations between the Sigma of a pixel and Sigmabar.             //
//                                                                         //
//  2) fPadFlag = 2 :                                                      //
//     Generate a pedestal sigma for each pixel using the 3D-histogram     //
//     Theta, pixel no., Sigma (fHSigmaPixTheta).                          //
//                                                                         //
//                                                                         //
//  The padding has to be done before the image cleaning because the       //
//  image cleaning depends on the pedestal sigmas.                         //
//                                                                         //
//                                                                         //
//  This implementation has been tested for CT1 data. For MAGIC some       //
//  modifications are necessary.                                           //
//                                                                         //
/////////////////////////////////////////////////////////////////////////////
#include "MPadSchweizer.h"

#include <stdio.h>

#include "TH1.h"
#include "TH2.h"
#include "TH3.h"
#include "TProfile.h"
#include "TRandom.h"
#include "TCanvas.h"

#include "MBinning.h"
#include "MSigmabar.h"
#include "MMcEvt.hxx"
#include "MLog.h"
#include "MLogManip.h"
#include "MParList.h"
#include "MGeomCam.h"
#include "MCerPhotPix.h"
#include "MCerPhotEvt.h"
#include "MPedestalPix.h"

ClassImp(MPadSchweizer);

// --------------------------------------------------------------------------
//
// Default constructor. 
//
MPadSchweizer::MPadSchweizer(const char *name, const char *title,
                             TH2D *fHist2, TH3D *fHist3, TH3D *fHist3Diff) : 
               fRunType(0), fGroup(0)
{
  fName  = name  ? name  : "MPadSchweizer";
  fTitle = title ? title : "Task for the padding (Schweizer)";

  fHSigmaTheta    = fHist2;
  fHSigmaPixTheta = fHist3;
  fHDiffPixTheta  = fHist3Diff;

  Print();
}

// --------------------------------------------------------------------------
//
// Destructor. 
//
MPadSchweizer::~MPadSchweizer()
{
  //nothing yet
}

// --------------------------------------------------------------------------
//
// Set the option for the padding 
//
void MPadSchweizer::SetPadFlag(Int_t padflag)
{
  fPadFlag = padflag;
  *fLog << "MPadSchweizer::SetPadFlag(); choose option " << fPadFlag << endl; 
}

// --------------------------------------------------------------------------
//
//  Set the pointers and prepare the histograms
//
Bool_t MPadSchweizer::PreProcess(MParList *pList)
{
  fRnd = new TRandom3(0);

  fMcEvt = (MMcEvt*)pList->FindObject("MMcEvt");
  if (!fMcEvt)
    {
       *fLog << err << dbginf << "MPadSchweizer::PreProcess : MMcEvt not found... aborting." << endl;
       return kFALSE;
     }
  
   fPed = (MPedestalCam*)pList->FindObject("MPedestalCam");
   if (!fPed)
     {
       *fLog << dbginf << "MPadSchweizer::PreProcess : MPedestalCam not found... aborting." << endl;
       return kFALSE;
     }
  
   fCam = (MGeomCam*)pList->FindObject("MGeomCam");
   if (!fCam)
     {
       *fLog << dbginf << "MPadSchweizer::PreProcess : MGeomCam not found (no geometry information available)... aborting." << endl;
       return kFALSE;
     }
  
   fEvt = (MCerPhotEvt*)pList->FindObject("MCerPhotEvt");
   if (!fEvt)
     {
       *fLog << dbginf << "MPadSchweizer::PreProcess : MCerPhotEvt not found... aborting." << endl;
       return kFALSE;
     }

   fSigmabar = (MSigmabar*)pList->FindCreateObj("MSigmabar");
   if (!fSigmabar)
     {
       *fLog << dbginf << "MPadSchweizer::PreProcess : MSigmabar not found... aborting." << endl;
       return kFALSE;
     }
   
   // Get Theta Binning  
   MBinning* binstheta  = (MBinning*)pList->FindObject("BinningTheta");
   if (!binstheta)
     {
       *fLog << err << dbginf << "MPadSchweizer::PreProcess : BinningTheta not found... aborting." << endl;
       return kFALSE;      
     }

   // Get Sigma Binning  
   MBinning* binssigma  = (MBinning*)pList->FindObject("BinningSigmabar");
   if (!binssigma)
     {
       *fLog << err << dbginf << "MPadSchweizer::PreProcess : BinningSigmabar not found... aborting." << endl;
       return kFALSE;      
     }

   // Get binning for (sigma^2-sigmabar^2)  
   MBinning* binsdiff  = (MBinning*)pList->FindObject("BinningDiffsigma2");
   if (!binsdiff)
     {
       *fLog << err << dbginf << "MHSigmaTheta::SetupFill : BinningDiffsigma2 not found... aborting." << endl;
       return kFALSE;      
     }

   // Get binning for pixel number
   UInt_t npix = fPed->GetSize();
   MBinning binspix("BinningPixel");
   MBinning* binspixel = &binspix;
   binspixel->SetEdges(npix, -0.5, ((float)npix)-0.5 ); 


   //--------------------------------------------------------------------
   // histograms for checking the padding
   //
   // plot pedestal sigmas
   fHSigmaPedestal = new TH2D("SigPed","Sigma: after vs. before padding", 
                     100, 0.0, 5.0, 100, 0.0, 5.0);
   fHSigmaPedestal->SetXTitle("orig. Pedestal sigma");
   fHSigmaPedestal->SetYTitle("padded Pedestal sigma");

   // plot no.of photons (before vs. after padding) 
   fHPhotons = new TH2D("Photons","Photons: after vs.before padding", 
                        100, -10.0, 90.0, 100, -10, 90);
   fHPhotons->SetXTitle("no.of photons before padding");
   fHPhotons->SetYTitle("no.of photons after padding");

   // plot of added NSB
   fHNSB = new TH1D("NSB","Distribution of added NSB", 100, -10.0, 10.0);
   fHNSB->SetXTitle("no.of added NSB photons");
   fHNSB->SetYTitle("no.of pixels");

   fHSigbarTheta = new TH2D();
   fHSigbarTheta->SetNameTitle("fHSigmabar",
                            "2D : Sigmabar, \\Theta (after padding)");
   MH::SetBinning(fHSigbarTheta, binstheta, binssigma);
   fHSigbarTheta->SetXTitle("Theta");
   fHSigbarTheta->SetYTitle("Sigmabar");

   fHSigPixTh = new TH3D();
   fHSigPixTh->SetNameTitle("fHSigPixTh","3D: \\Theta, pixel no., Sigma (after padding)");
   MH::SetBinning(fHSigPixTh, binstheta, binspixel, binssigma);
   fHSigPixTh->SetXTitle("\\Theta [\\circ]");
   fHSigPixTh->SetYTitle("pixel number");
   fHSigPixTh->SetZTitle("Sigma");

   fHDiffPixTh = new TH3D();
   fHDiffPixTh->SetNameTitle("fHDiffPixTh","3D: \\Theta, pixel no., Sigma^2-Sigmabar^2 (after padding)");
   MH::SetBinning(fHDiffPixTh, binstheta, binspixel, binsdiff);
   fHDiffPixTh->SetXTitle("\\Theta [\\circ]");
   fHDiffPixTh->SetYTitle("pixel number");
   fHDiffPixTh->SetZTitle("Sigma^2-Sigmabar^2");

   return kTRUE;
}

// --------------------------------------------------------------------------
//
// Do the Padding
// idealy the events to be padded should have been generated without NSB
// 
Bool_t MPadSchweizer::Process()
{
  //*fLog << "Entry MPadSchweizer::Process();" << endl;

  const UInt_t npix = fEvt->GetNumPixels();

  //$$$$$$$$$$$$$$$$$$$$$$$$$$
  // to simulate the situation that before the padding the NSB and 
  // electronic noise are zero : set Sigma = 0 for all pixels
  for (UInt_t i=0; i<npix; i++) 
  {   
    MCerPhotPix &pix = fEvt->operator[](i);      
    Int_t j = pix.GetPixId();

    MPedestalPix &ppix = fPed->operator[](j);
    ppix.SetMeanRms(0.0);
  }
  //$$$$$$$$$$$$$$$$$$$$$$$$$$

  //-------------------------------------------
  // Calculate average sigma of the event
  //
  Double_t SigbarOld = fSigmabar->Calc(*fCam, *fPed, *fEvt);
  //fSigmabar->Print("");

  //if (SigbarOld > 0.0)
  //{
    //*fLog << "MPadSchweizer::Process(); Sigmabar of event to be padded is > 0 : "
    //      << SigbarOld << ". Stop event loop " << endl;
    // input data should have Sigmabar = 0; stop event loop
    // return kFALSE; 
  //}

  Double_t Theta  = kRad2Deg*fMcEvt->GetTelescopeTheta();
  // *fLog << "Theta = " << Theta << endl;


  //-------------------------------------------
  // for the current Theta, 
  // generate a Sigmabar according to the histogram fHSigmaTheta
  //
  Double_t Sigmabar;
  Int_t binNumber = fHSigmaTheta->GetXaxis()->FindBin(Theta);

  if ( binNumber < 1  ||  binNumber > fHSigmaTheta->GetNbinsX() )
  {
    *fLog << "MPadSchweizer::Process(); binNumber out of range : Theta, binNumber = "
          << Theta << ",  " << binNumber << ";  Skip event " << endl;
    // event cannot be padded; skip event
    return kCONTINUE;
  }
  else
  {
    TH1D* fHSigma = 
          fHSigmaTheta->ProjectionY("", binNumber, binNumber, "");
    if ( fHSigma->GetEntries() == 0.0 )
    {
      *fLog << "MPadSchweizer::Process(); projection for Theta bin " 
            << binNumber << " has no entries; Skip event " << endl;
      // event cannot be padded; skip event
      delete fHSigma;
      return kCONTINUE;
    }
    else
    {
      Sigmabar = fHSigma->GetRandom(); 

      //*fLog << "Theta, bin number = " << Theta << ",  " << binNumber 
      //      << ",  Sigmabar = " << Sigmabar << endl;
    }
    delete fHSigma;
  }


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

  //*fLog << "MPadSchweizer::Process(); SigbarOld, Sigmabar = " 
  //      << SigbarOld << ",  "<< Sigmabar << endl;

  // Skip event if target Sigmabar is <= SigbarOld
  if (Sigmabar <= SigbarOld)
  {
    *fLog << "MPadSchweizer::Process(); target Sigmabar is less than SigbarOld : "
          << Sigmabar << ",  " << SigbarOld << ",   Skip event" << endl;
    return kCONTINUE;     
  }


  //-------------------------------------------
  //
  // Calculate average number of NSB photons to be added (lambdabar)
  // from the value of Sigmabar, 
  //  - making assumptions about the average electronic noise (elNoise2) and
  //  - using a fixed value (F2excess)  for the excess noise factor
  
  Double_t elNoise2;         // [photons]
  Double_t F2excess  = 1.3;
  Double_t lambdabar;        // [photons]


  Int_t binTheta = fHDiffPixTheta->GetXaxis()->FindBin(Theta);
  if (binTheta != binNumber)
  {
    cout << "The binnings of the 2 histograms are not identical; aborting"
         << endl;
    return kERROR;
  }

  // Get RMS of (Sigma^2-Sigmabar^2) in this Theta bin.
  // The average electronic noise (to be added) has to be well above this RMS,
  // otherwise the electronic noise of an individual pixel (elNoise2Pix)
  // may become negative
  TH1D* fHNoise = fHDiffPixTheta->ProjectionZ("", binTheta, binTheta,
                                                          0, 9999, "");
  Double_t RMS = fHNoise->GetRMS(1);  
  delete fHNoise;
  elNoise2 = TMath::Min(RMS,  Sigmabar*Sigmabar - SigbarOld*SigbarOld);
  //*fLog << "elNoise2 = " << elNoise2 << endl; 

  lambdabar = (Sigmabar*Sigmabar - SigbarOld*SigbarOld - elNoise2) 
              / F2excess;    
  // This value of lambdabar is the same for all pixels;
  // note that lambdabar is normalized to the area of pixel 0


  //----------   start loop over pixels   ---------------------------------
  // do the padding for each pixel
  //
  // pad only pixels   - which are used (before image cleaning)
  //                   - and for which the no.of photons is != 0.0
  //
  Double_t Sig         = 0.0;
  Double_t Sigma2      = 0.0;
  Double_t Diff        = 0.0;
  Double_t addSig2     = 0.0;
  Double_t elNoise2Pix = 0.0;


  for (UInt_t i=0; i<npix; i++) 
  {   
    MCerPhotPix &pix = fEvt->operator[](i);      
    if ( !pix.IsPixelUsed() )
      continue;

    if ( pix.GetNumPhotons() == 0.0)
    {
      *fLog << "MPadSchweizer::Process(); no.of photons is 0 for used pixel" 
            << endl;
      continue;
    }

    Int_t j = pix.GetPixId();
    Double_t Area = fCam->GetPixRatio(j);

    MPedestalPix &ppix = fPed->operator[](j);
    Double_t oldsigma = ppix.GetMeanRms();

    //---------------------------------
    // throw the Sigma for this pixel 
    //
    Int_t binPixel = fHDiffPixTheta->GetYaxis()->FindBin( (Double_t)j );

    Int_t count;
    Int_t OK;
    TH1D* fHDiff;
    TH1D* fHSig;

    switch (fPadFlag)
    {
    case 1 :
      // throw the Sigma for this pixel from the distribution fHDiffPixTheta
      fHDiff = 
            fHDiffPixTheta->ProjectionZ("", binTheta, binTheta,
                                             binPixel, binPixel, "");
      if ( fHDiff->GetEntries() == 0.0 )
      {
        *fLog << "MPadSchweizer::Process(); projection for Theta bin " 
              << binTheta << "  and pixel bin " << binPixel  
              << " has no entries;  aborting " << endl;
        return kERROR;
      }

      count = 0;
      OK = 0;
      for (Int_t m=0; m<20; m++)
      {
        count += 1;
        Diff = fHDiff->GetRandom();
        // the following condition ensures that elNoise2Pix > 0.0 
        if ( (Diff+Sigmabar*Sigmabar-oldsigma*oldsigma/Area-
              lambdabar*F2excess) > 0.0 )
	{
          OK = 1;
          break;
	}
      }
      if (OK == 0) 
      {
        *fLog << "Theta, j, count, Sigmabar, Diff = " << Theta << ",  " 
              << j << ",  " << count << ",  " << Sigmabar << ",  " 
              << Diff << endl;
        Diff = lambdabar*F2excess + oldsigma*oldsigma/Area 
                                  - Sigmabar*Sigmabar;
      }
      delete fHDiff;
      Sigma2 = Diff + Sigmabar*Sigmabar;
      break;

    case 2 :
      // throw the Sigma for this pixel from the distribution fHSigmaPixTheta
      fHSig = 
            fHSigmaPixTheta->ProjectionZ("", binTheta, binTheta,
                                             binPixel, binPixel, "");
      if ( fHSig->GetEntries() == 0.0 )
      {
        *fLog << "MPadSchweizer::Process(); projection for Theta bin " 
              << binTheta << "  and pixel bin " << binPixel  
              << " has no entries;  aborting " << endl;
        return kERROR;
      }

      count = 0;
      OK = 0;
      for (Int_t m=0; m<20; m++)
      {
        count += 1;
        Sig = fHSig->GetRandom();
        Sigma2 = Sig*Sig/Area;
        // the following condition ensures that elNoise2Pix > 0.0 
        if ( (Sigma2-oldsigma*oldsigma/Area-lambdabar*F2excess) > 0.0 )
	{
          OK = 1;
          break;
	}
      }
      if (OK == 0) 
      {
        *fLog << "Theta, j, count, Sigmabar, Sig = " << Theta << ",  " 
              << j << ",  " << count << ",  " << Sigmabar << ",  " 
              << Sig << endl;
        Sigma2 = lambdabar*F2excess + oldsigma*oldsigma/Area; 
      }
      delete fHSig;
      break;
    }

    //---------------------------------
    // get the additional sigma^2 for this pixel (due to the padding)
    addSig2 = Sigma2*Area - oldsigma*oldsigma;


    //---------------------------------
    // get the additional electronic noise for this pixel
    elNoise2Pix = addSig2 - lambdabar*F2excess*Area;


    //---------------------------------
    // throw actual number of additional NSB photons (NSB)
    //       and its RMS (sigmaNSB) 
    Double_t NSB0 = fRnd->Poisson(lambdabar*Area);
    Double_t arg  = NSB0*(F2excess-1.0) + elNoise2Pix;
    Double_t sigmaNSB0;

    if (arg >= 0.0)
    {
      sigmaNSB0 = sqrt( arg );
    }
    else
    {
      *fLog << "MPadSchweizer::Process(); argument of sqrt < 0.0 : " 
            << arg << endl;
      sigmaNSB0 = 0.0000001;      
    }


    //---------------------------------
    // smear NSB0 according to sigmaNSB0
    // and subtract lambdabar because of AC coupling
    Double_t NSB = fRnd->Gaus(NSB0, sigmaNSB0) - lambdabar*Area;

    //---------------------------------
 
    // add additional NSB to the number of photons
    Double_t oldphotons = pix.GetNumPhotons();
    Double_t newphotons = oldphotons + NSB;
    pix.SetNumPhotons(	newphotons );

    fHNSB->Fill( NSB/sqrt(Area) );
    fHPhotons->Fill( newphotons/sqrt(Area), oldphotons/sqrt(Area) );


    // error: add sigma of padded noise quadratically
    Double_t olderror = pix.GetErrorPhot();
    Double_t newerror = sqrt( olderror*olderror + addSig2 );
    pix.SetErrorPhot( newerror );


    Double_t newsigma = sqrt( oldsigma*oldsigma + addSig2 ); 
    ppix.SetMeanRms( newsigma );

    fHSigmaPedestal->Fill( oldsigma, newsigma );
    fHSigPixTh-> Fill( Theta, (Double_t) j, newsigma );
  } 
  //----------   end of loop over pixels   ---------------------------------

  // Calculate Sigmabar again and crosscheck
  Double_t SigbarNew = fSigmabar->Calc(*fCam, *fPed, *fEvt);
  //fSigmabar->Print("");

  fHSigbarTheta->Fill( Theta, SigbarNew );

  // this loop is only for filling the histogram fHDiffPixTh
  for (UInt_t i=0; i<npix; i++) 
  {   
    MCerPhotPix &pix = fEvt->operator[](i);      
    if ( !pix.IsPixelUsed() )
      continue;

    if ( pix.GetNumPhotons() == 0.0)
    {
      *fLog << "MPadSchweizer::Process(); no.of photons is 0 for used pixel" 
            << endl;
      continue;
    }

    Int_t j = pix.GetPixId();
    Double_t Area = fCam->GetPixRatio(j);

    MPedestalPix &ppix = fPed->operator[](j);
    Double_t newsigma = ppix.GetMeanRms();

    fHDiffPixTh->Fill( Theta, (Double_t) j, 
                              newsigma*newsigma/Area-SigbarNew*SigbarNew);
  }


  //*fLog << "Exit MPadSchweizer::Process();" << endl;

  return kTRUE;

}

// --------------------------------------------------------------------------
//
//
Bool_t MPadSchweizer::PostProcess()
{
    TCanvas &c = *(MH::MakeDefCanvas("PadSchweizer", "", 900, 1200)); 
    c.Divide(3, 4);
    gROOT->SetSelectedPad(NULL);


    c.cd(1);
    fHSigmaTheta->SetTitle("2D : Sigmabar, \\Theta (reference sample)");
    fHSigmaTheta->DrawClone();     
    fHSigmaTheta->SetBit(kCanDelete);     

      //c.cd(1);
      //fHSigmaPixTheta->DrawClone();     
      //fHSigmaPixTheta->SetBit(kCanDelete);     

    c.cd(4);
    fHSigbarTheta->DrawClone();     
    fHSigbarTheta->SetBit(kCanDelete);     


    c.cd(7);
    fHNSB->DrawClone();
    fHNSB->SetBit(kCanDelete);    


    //--------------------------------------------------------------------
    // draw the 3D histogram : Theta, pixel, Sigma^2-Sigmabar^2

    TH2D *l;

    c.cd(2);
    l = (TH2D*) ((TH3*)fHDiffPixTh)->Project3D("zx");

    l->SetTitle("Sigma^2-Sigmabar^2 vs. \\Theta (all pixels)");
    l->SetXTitle("\\Theta [\\circ]");
    l->SetYTitle("Sigma^2-Sigmabar^2");

    *fLog << "before box" << endl;

    l->Draw("box");
    l->SetBit(kCanDelete);;

    c.cd(5);
    l = (TH2D*) ((TH3*)fHDiffPixTh)->Project3D("zy");
    l->SetTitle("Sigma^2-Sigmabar^2 vs. pixel number (all \\Theta)");
    l->SetXTitle("pixel");
    l->SetYTitle("Sigma^2-Sigmabar^2");

    l->Draw("box");
    l->SetBit(kCanDelete);;

    c.cd(8);
    ((TH2*)fHDiffPixTh)->DrawCopy();


    //--------------------------------------------------------------------
    // draw the 3D histogram : Theta, pixel, Sigma

    TH2D *k;

    c.cd(3);
    k = (TH2D*) ((TH3*)fHSigPixTh)->Project3D("zx");
    k->SetTitle("Sigma vs. \\Theta (all pixels)");
    k->SetXTitle("\\Theta [\\circ]");
    k->SetYTitle("Sigma");

    k->Draw("box");
    k->SetBit(kCanDelete);

    c.cd(6);
    k = (TH2D*) ((TH3*)fHSigPixTh)->Project3D("zy");
    k->SetTitle("Sigma vs. pixel number (all \\Theta)");
    k->SetXTitle("pixel");
    k->SetYTitle("Sigma");

    k->Draw("box");
    k->SetBit(kCanDelete);;

    c.cd(9);
    ((TH2*)fHSigPixTh)->DrawCopy();

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

    c.cd(10);
    fHSigmaPedestal->DrawClone();
    fHSigmaPedestal->SetBit(kCanDelete);    

    c.cd(11);
    fHPhotons->DrawClone();
    fHPhotons->SetBit(kCanDelete);    

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

  return kTRUE;
}

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