#include <cmath>


#include "MessageService/MsgService.h"
#include "TFile.h"
#include "TH2D.h"
#include "TDirectory.h"
#include "TChain.h"
#include "GenericFactory.h"
#include "BeamChi2.h"
#include "Parsers.h"
#include "TEventList.h"

REGISTER_OBJECT(Chi2Func,BeamChi2)


  using namespace std;

CVSID("");
BeamChi2::BeamChi2() 
  :fDataInput(""),
   fDataHistName(""),
   fMCInput(""),
   fMCTreeName(""),
   fDataPOT(0.),
   fMCPOT(0.),
   fUseNeugen(false),
   fUseSmoothing(false),
   fSmoothWidth(0.),
   fSmoothYWidth(0.01),
   fSaveFitData(false),
   fFitDataFile(""),
   fFitDataTree("tree"),
   fLoadFitData(false),
   fLoadFitDataFile(""),
   fInitChi2(0.),
   fBestChi2(0.),
   fNDF(0),
   fName("")
{

  fSkzpCalc=&SkzpCalculator::Instance();

}

//**********************************************************************

bool BeamChi2::SetProperty(std::string prop) 
{
  // Set a particular property should be of the
  // form of Property=Value
  
  string property=prop.substr(0,prop.find("="));
  string val=prop.substr(prop.find("=")+1);
  if(property=="UseNeugen")         Parsers::parse_val(val,fUseNeugen);
  else if(property=="SmoothWidth")  Parsers::parse_val(val,fSmoothWidth);
  else if(property=="SmoothYWidth") Parsers::parse_val(val,fSmoothYWidth);
  else if(property=="UseSmoothing") Parsers::parse_val(val,fUseSmoothing);
  else if(property=="DataPOT")      Parsers::parse_val(val,fDataPOT);
  else if(property=="DataHistName")  Parsers::parse_val(val,fDataHistName);
  else if(property=="MCPOT")        Parsers::parse_val(val,fMCPOT);
  else if(property=="MCInput")      Parsers::parse_val(val,fMCInput);
  else if(property=="MCTreeName")   Parsers::parse_val(val,fMCTreeName);
  else if(property=="MCCut")        Parsers::parse_val(val,fMCCut);
  else if(property=="DataInput")      Parsers::parse_val(val,fDataInput);
  else if(property=="Vars")         Parsers::parse_val(val,fVars);
  else if(property=="SaveFitData")  Parsers::parse_val(val,fSaveFitData);
  else if(property=="FitDataFile")  Parsers::parse_val(val,fFitDataFile);
  else if(property=="LoadFitData")  Parsers::parse_val(val,fLoadFitData);
  else if(property=="LoadFitDataFile")     Parsers::parse_val(val,fLoadFitDataFile);
  else if(property=="Beam") {
    string dummy;
    Parsers::parse_val(val,dummy);
    fBeam=FitBeam::StringToEnum(dummy);
  }
  else {
    MSG("BeamChi2",Msg::kError)<<"Cannot set property "<<property
			       <<" for BeamChi2 function "<<fName<<endl;
    return false;
  } 

  return true;
  
}

//**********************************************************************

bool BeamChi2::GetData()
{
  TDirectory *initial=gDirectory;

  TFile *file= TFile::Open(fDataInput.c_str(),"READ");

  if(!file) {
    MSG("BeamChi2",Msg::kError)<<"Cannot open file "<<fDataInput
			       <<" for BeamChi2 function "<<fName<<endl;
    return false;
  }

  std::vector<std::string>::iterator iter=fVars.begin();
  std::vector<std::string>::iterator  vars_end=fVars.end();

  for( ; iter!=vars_end ; ++iter) {

    std::string var=*iter;
    std::string histname=var+"_"+fDataHistName+"_data";

    TH1 *h= dynamic_cast<TH1*>(file->Get(histname.c_str()));
    if(!h) {
      MSG("BeamChi2",Msg::kError)<<"Histogram  "<<histname << " does not exist in "
				 <<fDataInput
				 <<" for BeamChi2 function "<<fName<<endl;
      return false;
    }

    // Create the histograms based on what was found in the data
    TH1 *data_clone = dynamic_cast<TH1*>(h->Clone());
    data_clone->SetDirectory(0);
    data_clone->SetName(Form("reco_%s_%s",var.c_str(),
                             fName.c_str()));
    
    TH1 *mc_clone = dynamic_cast<TH1*>(h->Clone());
    mc_clone->Reset();
    mc_clone->Sumw2();
    mc_clone->SetDirectory(0);
    mc_clone->SetName(Form("mc_%s_%s",var.c_str(),
                           fName.c_str()));
    TH1 *rw_clone = dynamic_cast<TH1*>(h->Clone());
    rw_clone->Reset();
    rw_clone->Sumw2();
    rw_clone->SetDirectory(0);
    rw_clone->SetName(Form("rw_%s_%s",var.c_str(),
                           fName.c_str()));
    
    fRWHists.push_back(rw_clone);
    fMCHists.push_back(mc_clone);
    fDataHists.push_back(data_clone);

  }
  
  file->Close();
  initial->cd();
  return true;
}

//**********************************************************************

bool BeamChi2::FillMC()  
{
  TDirectory *initial=gDirectory;
  int i=0;
  if(!fLoadFitData) {
    
    
    TChain *mctree= new TChain(fMCTreeName.c_str(),"");
    mctree->Add(fMCInput.c_str());
    
    
    TFile *ofile=0;
    TTree *outtree=0;
    FitData *data=new FitData;
    if(fSaveFitData) {
      ofile=new TFile(fFitDataFile.c_str(),"recreate");
      outtree=new TTree(fFitDataTree.c_str(),""); 
      outtree->Branch("fit",&data);
    }
    
    // Define PAN variables
    int pass;
    int ntrack;
    int is_fid;
    float reco_enu;
    float true_enu;
    float reco_emu;
    float reco_eshw;
  float true_eshw;
  float tpx;
  float tpy;
  float tpz;
  
  float ppdxdz;
  float ppdydz;   
  float pppz;
  int tptype;
  int ptype;
   
  float reco_y;
  float nu_px;
  float nu_py;
  float nu_pz;
  float tar_e;
  float tar_px;
  float tar_py;
  float tar_pz;
  float true_y;
  float true_x;
  float true_q2;
  float true_w2;
  int inu;
  int process;
  int initial_state;
  int nucleus;
  int resnum;
  int had_fs;
  int cc_nc;
  int is_cev;

    
  // set branch addresses to mc tree
  mctree->SetBranchAddress("pass",&pass);
  mctree->SetBranchAddress("ntrack",&ntrack);
  mctree->SetBranchAddress("is_fid",&is_fid);

  mctree->SetBranchAddress("reco_enu",&reco_enu);
  mctree->SetBranchAddress("true_enu",&true_enu);
  mctree->SetBranchAddress("reco_emu",&reco_emu);
  mctree->SetBranchAddress("true_eshw",&true_eshw);
  mctree->SetBranchAddress("reco_y",&reco_y);
  mctree->SetBranchAddress("true_x",&true_x);
  mctree->SetBranchAddress("true_y",&true_y);
  mctree->SetBranchAddress("true_q2",&true_q2);
  mctree->SetBranchAddress("true_w2",&true_w2);
  mctree->SetBranchAddress("process",&process);
  mctree->SetBranchAddress("initial_state",&initial_state);
  mctree->SetBranchAddress("nucleus",&nucleus);
  mctree->SetBranchAddress("cc_nc",&cc_nc);
  mctree->SetBranchAddress("is_cev",&is_cev);
  mctree->SetBranchAddress("ntrack",&ntrack);
   
  mctree->SetBranchAddress("reco_eshw",&reco_eshw);
  mctree->SetBranchAddress("tpx",&tpx);
  mctree->SetBranchAddress("tpy",&tpy);
  mctree->SetBranchAddress("tpz",&tpz);
  mctree->SetBranchAddress("tptype",&tptype);
  mctree->SetBranchAddress("ptype",&ptype);
  mctree->SetBranchAddress("pppz",&pppz);
  mctree->SetBranchAddress("ppdydz",&ppdydz);
  mctree->SetBranchAddress("ppdxdz",&ppdxdz);
   
  mctree->SetBranchAddress("nu_px",&nu_px);
  mctree->SetBranchAddress("nu_py",&nu_py);
  mctree->SetBranchAddress("nu_pz",&nu_pz);
  mctree->SetBranchAddress("tar_px",&tar_px);
  mctree->SetBranchAddress("tar_py",&tar_py);
  mctree->SetBranchAddress("tar_pz",&tar_pz);
  mctree->SetBranchAddress("tar_e",&tar_e);
  mctree->SetBranchAddress("inu",&inu);
  
  mctree->SetBranchAddress("resnum",&resnum);
  mctree->SetBranchAddress("resnum",&had_fs);
     
  i=0;

    if(!fMCCut.empty()) {
      mctree->Draw(">>elist",fMCCut.c_str());
    }
    else mctree->Draw(">>elist","");

    TEventList *elist = (TEventList*)gDirectory->Get("elist");
    if(!elist) {
      MSG("BeamChi2",Msg::kError)<<"Cannot get event list "<<endl;
      return false;
    }
    mctree->SetEventList(elist);

    
    while(mctree->GetEntry((elist->GetEntry(i))) > 0) {
      
      ++i;
      
      FitData fit_data;
      fit_data.RecoEnergy(reco_enu);
      fit_data.TrueEnergy(true_enu);
      fit_data.RecoTrkEnergy(reco_emu);
      fit_data.RecoShwEnergy(reco_eshw);
      fit_data.TargetPt(sqrt(tpx*tpx+tpy*tpy));
      fit_data.TargetPz(tpz);
      fit_data.NuType(inu);
      fit_data.ParentType(tptype);
      fit_data.InteractionType(cc_nc);
      fit_data.InteractionProcess(process);
      fit_data.IsContained(is_cev);
      fit_data.BeamType(fBeam);
      
      // Calculate shower scale factor for 1-sigma
      double x0=10;
      double x = TMath::Min((double)true_eshw,x0);
      double err_a = 1.04/1.0 - 1.0;
      double err_b = (1.1+2.*(1-1.1)*(x/x0)+(1.1-1)*(x/x0)*(x/x0))/(1.0) - 1.0;
      double err_c = 0.057;
      double total_error=sqrt(pow(err_a,2)+pow(err_b,2)+pow(err_c,2));
      fit_data.ShwError(total_error);
      


    double ge=1.;
    MCEventInfo ei;
    if(fUseNeugen) {
	
      ei.nuE=true_enu;
      ei.nuPx=nu_px;
      ei.nuPy=nu_py;
      ei.nuPz=nu_pz;
      ei.tarE=tar_e;
      ei.tarPx=tar_px;
      ei.tarPy=tar_py;
      ei.tarPz=tar_pz;
      ei.y=true_y;
      ei.x=true_x;
      ei.q2=true_q2;
      ei.w2=true_w2;
      ei.iaction=cc_nc;
      ei.inu=inu;
      ei.iresonance=process;
      ei.initial_state=initial_state;
      ei.nucleus=nucleus;
      ei.had_fs=had_fs;
      if(ei.iresonance!=1005){
	ge = fSkzpCalc->GetNeugenWeight(&ei);
      }
    }
    fit_data.GeneratorError(ge-1.0);

    fMCList.push_back(fit_data);
    
    
      if(fSaveFitData) {
	data=&fit_data;
	outtree->Fill();
      }



    if(!fUseSmoothing) {
	
      for(int j=0; j<(int)fVars.size(); j++) {
	
	if( fVars[j]=="enu") {
	  fMCHists[j]->Fill(reco_enu,fDataPOT/fMCPOT);
	}
	else if( fVars[j]=="emu") {
	  fMCHists[j]->Fill(reco_emu,fDataPOT/fMCPOT);
	}
	else if( fVars[j]=="eshw") {
	  fMCHists[j]->Fill(reco_eshw,fDataPOT/fMCPOT);
	}
	else if( fVars[j]=="enu_y") {
	  TH2D *h2=(TH2D*) fMCHists[j];
	  h2->Fill(reco_enu,reco_eshw/reco_enu,fDataPOT/fMCPOT);
	}
      }
    }
    else {
      for(int j=0; j<(int)fVars.size(); j++) {
	if( fVars[j]=="enu") {
	  fSkzpCalc->FillExactSmooth(fMCHists[j],reco_enu,fSmoothWidth*reco_enu,
				     fDataPOT/fMCPOT);
	}
	else if( fVars[j]=="emu") {
	  fSkzpCalc->FillExactSmooth(fMCHists[j],reco_emu,fSmoothWidth*reco_emu,
				     fDataPOT/fMCPOT);
	}
	else if( fVars[j]=="eshw") {
	  fSkzpCalc->FillExactSmooth(fMCHists[j],reco_eshw,fSmoothWidth*reco_eshw,
				     fDataPOT/fMCPOT);
	}
	else if( fVars[j]=="enu_y") {
	  TH2D *h2=(TH2D*) fMCHists[j];
	  fSkzpCalc->FillExactSmooth2(h2,reco_enu,reco_eshw/reco_enu,reco_enu*fSmoothWidth,
				      fSmoothYWidth,fDataPOT/fMCPOT);
	}	
      }
    } // end else for smooth filling
      
  } // end event loop
    
    if(fSaveFitData) {
      ofile->cd();
      outtree->Write();
      ofile->Close();
    }
    
  } 
   else {
     TFile *file= TFile::Open(fLoadFitDataFile.c_str(),"READ");
    
     if(!file) {
       MSG("BeamChi2",Msg::kError)<<"Cannot open file "<<fLoadFitDataFile
				  <<" for BeamChi2 function "<<fName<<endl;
       return false;
       
     }
     TTree *mctree= dynamic_cast<TTree*> (file->Get(fFitDataTree.c_str()));
     if(!mctree) {
       MSG("BeamChi2",Msg::kError)<<"Cannot find tree "<< fFitDataTree
				  <<" in file "<<fMCInput
				  <<" for BeamChi2 function "<<fName<<endl;
       return false;
     }
     
     FitData *fit_data=new FitData;
     mctree->SetBranchAddress("fit",&fit_data);
     
     i=0;
     while(mctree->GetEntry(i) > 0) {
       ++i;
       
       fMCList.push_back(*fit_data);
       
       if(!fUseSmoothing) {
	 
	 for(int j=0; j<(int)fVars.size(); j++) {
	   
	   if( fVars[j]=="enu") {
	     fMCHists[j]->Fill(fit_data->RecoEnergy(),fDataPOT/fMCPOT);
	   }
	   else if( fVars[j]=="emu") {
	     fMCHists[j]->Fill(fit_data->RecoTrkEnergy(),fDataPOT/fMCPOT);
	   }
	   else if( fVars[j]=="eshw") {
	     fMCHists[j]->Fill(fit_data->RecoShwEnergy(),fDataPOT/fMCPOT);
	   }
	   else if( fVars[j]=="enu_y") {
	     TH2D *h2=(TH2D*) fMCHists[j];
	     h2->Fill(fit_data->RecoEnergy(),fit_data->RecoShwEnergy()/fit_data->RecoEnergy(),
		      fDataPOT/fMCPOT);
	   }
	 }
       }
       else {
	 for(int j=0; j<(int)fVars.size(); j++) {
	   if( fVars[j]=="enu") {
	     fSkzpCalc->FillExactSmooth(fMCHists[j],fit_data->RecoEnergy(),
					fSmoothWidth*fit_data->RecoEnergy(),
					fDataPOT/fMCPOT);
	   }
	   else if( fVars[j]=="emu") {
	     fSkzpCalc->FillExactSmooth(fMCHists[j],fit_data->RecoTrkEnergy(),
					fSmoothWidth*fit_data->RecoTrkEnergy(),
					fDataPOT/fMCPOT);
	   }
	   else if( fVars[j]=="eshw") {
	     fSkzpCalc->FillExactSmooth(fMCHists[j],fit_data->RecoShwEnergy(),
					fSmoothWidth*fit_data->RecoShwEnergy(),
					fDataPOT/fMCPOT);
	   }
	   else if( fVars[j]=="enu_y") {
	     
	     TH2D *h2=(TH2D*) fMCHists[j];
	     fSkzpCalc->FillExactSmooth2(h2,fit_data->RecoEnergy(),
					 fit_data->RecoShwEnergy()/fit_data->RecoEnergy(),
					 fit_data->RecoEnergy()*fSmoothWidth,
					 fSmoothYWidth,
					 fDataPOT/fMCPOT);
	   }	
	 }
       } // end else for smooth filling
     }
   }

    
  cout<<"Added "<<i<<" MC events"<<endl;  
  initial->cd();
  return true;
  
} 

//**********************************************************************

bool BeamChi2::ReweightMC()  
{
  
  std::vector<TH1*>::iterator hist_it=fRWHists.begin();
  std::vector<TH1*>::iterator hist_end=fRWHists.end();
  for( ; hist_it!=hist_end ; ++hist_it) {
    if(*hist_it) (*hist_it)-> Reset();
  }
  
  list<FitData>::const_iterator mc_iter=fMCList.begin();
  list<FitData>::const_iterator mc_end=fMCList.end();
  int i=0;
  for( ; mc_iter != mc_end; ++mc_iter) {
    
    double weight=1.;
    FitData data=*mc_iter;
    
    weight=fSkzpCalc->GetSkzpWeight(data);

    MSG("BeamChi2",Msg::kVerbose)<<"Event "<<i<<" Total Weight="
				 <<weight<<endl;
    i++;
    
    if(!fUseSmoothing) {
      
      for(int i=0; i<(int)fVars.size(); i++) {
	
	if( fVars[i]=="enu") {
	  fRWHists[i]->Fill(data.RecoEnergy(),fDataPOT/fMCPOT*weight);
	}
	else if( fVars[i]=="emu") {
	  fRWHists[i]->Fill(data.RecoTrkEnergy(),fDataPOT/fMCPOT*weight);
	}
	else if( fVars[i]=="eshw") {
	  fRWHists[i]->Fill(data.RecoShwEnergy(),fDataPOT/fMCPOT*weight);
	}
	else if( fVars[i]=="enu_y") {
	  TH2D *h2=(TH2D*) fRWHists[i];
	  h2->Fill(data.RecoEnergy(),data.RecoShwEnergy()/data.RecoEnergy(),
		   fDataPOT/fMCPOT*weight);
	}
	  }
    }
    else {
      
      for(int i=0; i<(int)fVars.size(); i++) {
	if( fVars[i]=="enu") {
	  fSkzpCalc->FillExactSmooth(fRWHists[i],data.RecoEnergy(),
				     data.RecoEnergy()*fSmoothWidth,fDataPOT/fMCPOT*weight);
	}
	else if( fVars[i]=="emu") {
	  fSkzpCalc->FillExactSmooth(fRWHists[i],data.RecoTrkEnergy(),
				     data.RecoTrkEnergy()*fSmoothWidth,fDataPOT/fMCPOT*weight);
	}
	else if( fVars[i]=="eshw") {
	  fSkzpCalc->FillExactSmooth(fRWHists[i],data.RecoShwEnergy(),
				     data.RecoShwEnergy()*fSmoothWidth,fDataPOT/fMCPOT*weight);
	}
	else if( fVars[i]=="enu_y") {
	  TH2D *h2=(TH2D*) fRWHists[i];
	  fSkzpCalc->FillExactSmooth2(h2,data.RecoEnergy(),
				      data.RecoShwEnergy()/data.RecoEnergy(),
				      data.RecoEnergy()*fSmoothWidth,
				      fSmoothYWidth,
				      fDataPOT/fMCPOT*weight);
	  
	}	
      }
    } // end else for smooth filling
    
  }// end loop over mc

  return true;
}

//**********************************************************************

double BeamChi2::GetHistChi2()
{
  
  fNDF=0; 
  double chi2=0;
  
  //double data,md,edata,emc;
  
  for(int k=0; k<(int) fVars.size(); ++k) {

    if(fRWHists[k]->InheritsFrom("TH1D") ||
       fRWHists[k]->InheritsFrom("TH1F") ) {
	  
      for(int i=1; i<=fRWHists[k]->GetNbinsX() ; ++i) {
	double data=fDataHists[k]->GetBinContent(i);
	double mc=fRWHists[k]->GetBinContent(i);
	double edata=fDataHists[k]->GetBinError(i);
	double emc=fRWHists[k]->GetBinError(i);
	    
	if(edata*edata+emc*emc==0){
	  continue;
	}
	chi2+=(data-mc)*(data-mc)/(edata*edata+emc*emc);
	if( data!=0 && mc!=0) fNDF++;
      }
    }
    else if(fRWHists[k]->InheritsFrom("TH2")) {
	  
      for(int i=1; i<=fRWHists[k]->GetNbinsX() ; ++i)
	for(int j=1; j<=fRWHists[k]->GetNbinsY() ; ++j) {
	  double data=fDataHists[k]->GetBinContent(i,j);
	  double mc=fRWHists[k]->GetBinContent(i,j);
	  double edata=fDataHists[k]->GetBinError(i,j);
	  double emc=fRWHists[k]->GetBinError(i,j);
	      
	  if(edata*edata+emc*emc==0){
	    continue;
	  }
	  chi2+=(data-mc)*(data-mc)/(edata*edata+emc*emc);
	  if( data!=0 && mc!=0) fNDF++;
	}
    }
  }
  return chi2;
}

//**********************************************************************

bool BeamChi2::Init()
{
  cout<<"Adding "<<fName<<endl;
  
  
  bool s=GetData();
  if(!s) return false;
  s=FillMC();
  
  if(s) return true;
  else return false;
}

//**********************************************************************

double BeamChi2::GetChi2()  
{
  ReweightMC();
  return GetHistChi2();
}

//**********************************************************************

void BeamChi2::Write() 
{
  for(int i=0; i<(int)fVars.size(); i++) {

    fDataHists[i]->Write();
    fRWHists[i]->Write();
    fMCHists[i]->Write();
  }

  
}

  
    

////////////////////////////////////////////////////////////////////////