#include <iostream>
#include "TFile.h"
#include "TPad.h"
#include "TH1F.h"
#include "TCanvas.h"
#include "TLegend.h"
#include "TChain.h"
#include "TStyle.h"
#include "TGaxis.h"
#include "TLatex.h"
#include "TText.h"
#include "Registry/Registry.h"
#include "MCReweight/MCReweight.h"
#include "MCReweight/MCEventInfo.h"
#include "MCReweight/NeugenWeightCalculator.h"
#include "MCReweight/Zbeam.h"
#include "MCReweight/Zfluk.h"

using namespace std;

void MedReweight(char* panFile)
{

  gStyle->SetOptStat(0);

  const Int_t NBEAMS=6;

  MCReweight *mcr = &MCReweight::Instance();
  NeugenWeightCalculator *n=new NeugenWeightCalculator();
  mcr->AddWeightCalculator(n);

  Registry *rwtconfig = new Registry();
  rwtconfig->UnLockValues();
  rwtconfig->UnLockKeys();
  rwtconfig->Clear();
  rwtconfig->Set("neugen:config_name","MODBYRS");
  rwtconfig->Set("neugen:config_no",3);
  rwtconfig->LockValues();
  rwtconfig->LockKeys();

  Zbeam *zbeam= new Zbeam();
  Double_t beam_sys[5]={0.756841,0.458754,-0.000287821,3.30865e-07,-2.60897};  // 12 parameter fit
  // Double_t beam_sys[5]={-0.139562,0.435836,0.00257711,0.861925,0.36293};    // 15 parameter fit
  Zfluk *zfluk = new Zfluk(); 
  Double_t had_prod[7]={1.59562,-8.13304,1.12333,0.177879,0.938968,2.11178,0.883997};   // 12 parameter fit
  // Double_t had_prod[7]={2.22749,-4.28995,1.06742,1.23044,0.996444,2.74841,0.844169}; // 15 parameter fit
  
  Int_t det=1;
  Int_t Ibeam[NBEAMS]={6,2,7,3,4,8};
  // {le10 low horn, le10, le10 high horn, pME, pHE, horn off}

  TFile* f = new TFile(panFile,"update");

  TTree* pan = (TTree*)f->Get("pan");

  Float_t reco_emu,reco_eshw;
  Float_t tpx,tpy,tpz;
  Int_t tptype;
  Float_t nu_px;
  Float_t nu_py;
  Float_t nu_pz;
  Float_t tar_e;
  Float_t tar_px;
  Float_t tar_py;
  Float_t tar_pz;
  Float_t true_enu;
  Float_t true_y;
  Float_t true_x;
  Float_t true_q2;
  Float_t true_w2;
  Int_t inu;
  Int_t process;
  Int_t initial_state;
  Int_t nucleus;
  Int_t resnum;
  Int_t cc_nc;
  
  pan->SetBranchAddress("reco_emu",&reco_emu);
  pan->SetBranchAddress("reco_eshw",&reco_eshw);
  pan->SetBranchAddress("tpx",&tpx);
  pan->SetBranchAddress("tpy",&tpy);
  pan->SetBranchAddress("tpz",&tpz);
  pan->SetBranchAddress("tptype",&tptype);
  pan->SetBranchAddress("true_enu",&true_enu);
  pan->SetBranchAddress("true_x",&true_x);
  pan->SetBranchAddress("true_y",&true_y);
  pan->SetBranchAddress("true_q2",&true_q2);
  pan->SetBranchAddress("true_w2",&true_w2);
  pan->SetBranchAddress("process",&process);
  pan->SetBranchAddress("initial_state",&initial_state);
  pan->SetBranchAddress("nucleus",&nucleus);
  pan->SetBranchAddress("cc_nc",&cc_nc);
  pan->SetBranchAddress("inu",&inu);
  pan->SetBranchAddress("resnum",&resnum);
  pan->SetBranchAddress("nu_px",&nu_px);
  pan->SetBranchAddress("nu_py",&nu_py);
  pan->SetBranchAddress("nu_pz",&nu_pz);
  pan->SetBranchAddress("tar_px",&tar_px);
  pan->SetBranchAddress("tar_py",&tar_py);
  pan->SetBranchAddress("tar_pz",&tar_pz);
  pan->SetBranchAddress("tar_e",&tar_e);

  Float_t neuweight;
  Float_t beamweight;
  Float_t hadweight;
  
  TBranch* newBranch1 = pan->Branch("neuweight",&neuweight,"neuweight/F",32000);
  TBranch* newBranch2 = pan->Branch("beamweight",&beamweight,"beamweight/F",32000);
  TBranch* newBranch3 = pan->Branch("hadweight",&hadweight,"hadweight/F",32000);

  TH1F* enu = new TH1F("enu","enu",120,0,60);
  TH1F* rwenu = new TH1F("rwenu","rwenu",120,0,60);

  Int_t z=0;
  Long64_t nents = pan->GetEntriesFast();
  while(pan->GetEntry(z)){
    if(z%1000==0){
      cout << float(z)*100./float(nents) << "% done" << endl;
    }

    MCEventInfo ei;
    ei.UseStoredXSec(true);
    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=resnum;
    NuParent *np=0;
    Double_t gw=1.;
    if(ei.iresonance!=1005){
      gw = mcr->ComputeWeight(&ei,np,rwtconfig); // MODBYRS3 reweight
    }
    
    //Ibeam array element 0: le10 low horn I
    //                    1: le10
    //                    2: le10 high horn I
    //                    3: pME
    //                    4: pHE
    //                    5: horn off
    //Change according to MC file and only do spot size reweight for pME and pHE

    Double_t beamsysw=1.;
    //reweight for spot size in me and he
    beamsysw*=zbeam->GetWeight(det,Ibeam[4],6,0,true_enu);
    //horn 1 offset
    beamsysw*=zbeam->GetWeight(det,Ibeam[4],1,beam_sys[0],true_enu);
    //baffle scraping
    beamsysw*=zbeam->GetWeight(det,Ibeam[4],2,beam_sys[1],true_enu);
    //pot
    beamsysw*=zbeam->GetWeight(det,Ibeam[4],3,beam_sys[2],true_enu);
    //horn current miscalibration
    beamsysw*=zbeam->GetWeight(det,Ibeam[4],4,beam_sys[3],true_enu);
    //horn current distribution
    beamsysw*=zbeam->GetWeight(det,Ibeam[4],5,beam_sys[4],true_enu);
    
    //now change pt and pz
    Double_t pt = sqrt(tpx*tpx+tpy*tpy);
    Double_t pz = 1.0*tpz;      
    Double_t func = zfluk->GetWeight(tptype,pt,pz,had_prod);
    
    Double_t weight = gw*beamsysw*func;
    
    neuweight = gw;
    newBranch1->Fill();
    beamweight = beamsysw;
    newBranch2->Fill();
    hadweight = func;
    newBranch3->Fill();

    enu->Fill(reco_emu+reco_eshw);
    rwenu->Fill(reco_emu+reco_eshw,weight);

    z++;
    
  }
  
  cout << "Finished." << endl;

  pan->Write("",TObject::kOverwrite);

  TCanvas* c = new TCanvas("c","c");
  enu->Draw();
  rwenu->SetLineColor(2);
  rwenu->Draw("sames");
  c->Update();

}