//Renata Kopecna

#include <event.hh>

#include <vector>
#include <iostream>
#include <assert.h>

#include "TFile.h"
#include "TTree.h"
#include "TMath.h"

#include <helpers.hh>
#include <paths.hh>
#include <bu2kstarmumu_loader.hh>
#include <options.hh>

#include <spdlog.h>

void fcnc::print_event(const event& meas) 
{
  spdlog::info("Event has" );
  spdlog::info( " m:                  {0:f}", meas.m);
  spdlog::info( " sigma_m:            {0:f}", meas.sigma_m);
  spdlog::info( " weight:             {0:f}", meas.weight);
  spdlog::info( " delta_weight:       {0:f}", meas.delta_weight);
  spdlog::info( " max_mu_pt:          {0:f}", meas.max_mu_pt);
  spdlog::info( " kaon_pt:            {0:f}", meas.kaon_pt);
  spdlog::info( " event_type:         {0:d}", meas.event_type);
  spdlog::info( " Magnet polarity:    {0:d}", meas.magnet );
  spdlog::info( " cos(thetal):        {0:f}", meas.costhetal );
  spdlog::info( " cos(thetak):        {0:f}", meas.costhetak );
  spdlog::info( " phi:                {0:f}", meas.phi );

  spdlog::info( " q2:                 {0:f}", meas.q2 );
  spdlog::info( " p2:                 {0:f}", meas.p2 );
  spdlog::info( " cp conjugate:       " +(meas.cp_conjugate) );
  
  //spdlog::trace( " tagging_decision: {0:f}", meas.tagging_decision );
  //spdlog::trace( " tagging_dilution:         {0:f}", meas.tagging_dilution );
  //spdlog::trace( " efficiency:       {0:f}", meas.efficiency );
}

void fcnc::save_events(std::string filename, const std::vector<fcnc::event>& events){
  TFile* output = new TFile(filename.c_str(), "RECREATE");
  output->cd();
  TTree* tree = new TTree("Events", "Events");
  //init and link parameters
  bool cp_conjugate;
  int year, magnet, event_type;
  double m, mkpi, sigma_m, costhetal, costhetak, phi, q2, p2, weight, delta_weight, kaon_pt, max_mu_pt, pseudo_q2;
  tree->Branch("year",&year,"year/I"); 
  tree->Branch("m",&m,"m/D");
  tree->Branch("mkpi",&mkpi,"mkpi/D");
  tree->Branch("sigma_m",&sigma_m,"sigma_m/D");
  tree->Branch("costhetal",&costhetal,"costhetal/D");
  tree->Branch("costhetak",&costhetak,"costhetak/D");
  tree->Branch("phi",&phi,"phi/D");
  tree->Branch("magnet", &magnet, "magnet/I");
  tree->Branch("q2",&q2,"q2/D");
  tree->Branch("p2",&p2,"p2/D");
  tree->Branch("cp_conjugate",&cp_conjugate,"cp_conjugate/O");
  tree->Branch("weight", &weight, "weight/D");
  tree->Branch("delta_weight", &delta_weight, "delta_weight/D");
  tree->Branch("kaon_pt", &kaon_pt, "kaon_pt/D");
  tree->Branch("max_mu_pt", &max_mu_pt, "max_mu_pt/D");
  tree->Branch("pseudo_q2", &pseudo_q2, "pseudo_q2/D");
  tree->Branch("event_type", &event_type, "event_type/I"); //TODO
  spdlog::info( "START to fill events into TTree and save to file!" );

  for (unsigned int i=0; i< events.size(); i++){
      //get event from vector
      fcnc::event meas = events.at(i);
      //copy all values to the linked objects to save in TBranches
      year = meas.year;
      m = meas.m;
      mkpi = meas.mkpi;
      sigma_m = meas.sigma_m;
      costhetal = meas.costhetal;
      costhetak = meas.costhetak;
      phi = meas.phi;
      magnet = meas.magnet;
      q2 = meas.q2;
      p2 = meas.p2;
      cp_conjugate = meas.cp_conjugate;
      weight = meas.weight;
      delta_weight = meas.delta_weight;
      kaon_pt = meas.kaon_pt;
      max_mu_pt = meas.max_mu_pt;
      pseudo_q2 = meas.pseudo_q2;
      event_type = meas.event_type;
      //write to TTree
      tree->Fill();
  }
  Int_t N = tree->GetEntries();
  tree->Write();
  output->Write();
  output->Close();
  spdlog::info("DONE: Successfully saved {0:d} (out of {1:d}) events to file {2:s}", N, events.size(), filename);
  return;

}

void fcnc::lhcbtotheory(const bool b_plus, const double ctl_lhcb, const double ctk_lhcb, const double phi_lhcb, double& ctl_th, double& ctk_th, double& phi_th)
{
  ctl_th = b_plus ? ctl_lhcb : -ctl_lhcb;
  ctk_th = -ctk_lhcb;
  if (b_plus)
    phi_th = phi_lhcb > 0 ? TMath::Pi() - phi_lhcb : -TMath::Pi() - phi_lhcb;
  else
    phi_th = phi_lhcb; 
}

void fcnc::theorytolhcb(const bool b_plus, const double ctl_th, const double ctk_th, const double phi_th, double& ctl_lhcb, double& ctk_lhcb, double& phi_lhcb)
{
  ctl_lhcb = b_plus ? ctl_th : -ctl_th;
  ctk_lhcb = -ctk_th;
  if (b_plus)
    phi_lhcb = phi_th > 0 ? TMath::Pi() - phi_th : -TMath::Pi() - phi_th;	      
  else
    phi_lhcb = phi_th;
}

void fcnc::save_eos_events(std::string filename, const std::vector<fcnc::event>& events)
{
  TFile* output = new TFile(filename.c_str(), "RECREATE");
  output->cd();
  TTree* tree = new TTree("Events", "Events");
  fcnc::event meas;  
  tree->Branch("m",&meas.m,"m/D");
  tree->Branch("mkpi",&meas.mkpi,"mkpi/D");
  //double costhetal_lhcb, costhetak_lhcb, phi_lhcb;
  double costhetal_th, costhetak_th, phi_th;

  // tree->Branch("costhetal",&costhetal_lhcb,"costhetal/D");
  // tree->Branch("costhetak",&costhetak_lhcb,"costhetak/D");
  // tree->Branch("phi",&phi_lhcb,"phi/D");

  // tree->Branch("costhetal_theory",&meas.costhetal,"costhetal_theory/D");
  // tree->Branch("costhetak_theory",&meas.costhetak,"costhetak_theory/D");
  // tree->Branch("phi_theory",&meas.phi,"phi_theory/D");

  tree->Branch("costhetal",&meas.costhetal,"costhetal/D");
  tree->Branch("costhetak",&meas.costhetak,"costhetak/D");
  tree->Branch("phi",&meas.phi,"phi/D");

  tree->Branch("costhetal_theory",&costhetal_th,"costhetal_theory/D");
  tree->Branch("costhetak_theory",&costhetak_th,"costhetak_theory/D");
  tree->Branch("phi_theory",&phi_th,"phi_theory/D");

  tree->Branch("q2",&meas.q2,"q2/D");
  tree->Branch("p2",&meas.p2,"p2/D");
  tree->Branch("cp_conjugate",&meas.cp_conjugate,"cp_conjugate/O");
  int bid;
  tree->Branch("bid",&bid,"bid/I");
  tree->Branch("bkgcat", &meas.event_type, "bkgcat/I");

  for (unsigned int i=0; i< events.size(); i++)
    {
      meas = events.at(i);
      bid = meas.cp_conjugate ? +511 : -511;

      //convert from theory convention to lhcb convention
      lhcbtotheory(bid > 0, meas.costhetal, meas.costhetak, meas.phi, costhetal_th, costhetak_th, phi_th);
      tree->Fill();
    }
  tree->Write();
  output->Write();
  output->Close();
  delete output;
}

//Removes the resonances from the data sample
bool isResonance(double q2){
    for (auto bin: EXCLUDED_Q2){
        if (bin[0]<q2 && q2<bin[1]) return true;
    }
    return false;
}

std::vector<fcnc::event> fcnc::filterResonances(std::vector<fcnc::event> events){
    std::vector<fcnc::event> filtered;
    for (auto evt: events){
        if (isResonance(evt.q2)) continue;
        filtered.push_back(evt);
    }
    return filtered;
}

std::vector<fcnc::event> fcnc::load_events(std::string filename, std::string treename, int neventsmax)
{
    spdlog::info("Reading " + filename + " into fcnc event vector.");
    TFile* input = new TFile(filename.c_str(), "READ");    
    TTree* tree = dynamic_cast<TTree*>(input->Get(treename.c_str()));
    spdlog::debug("Treename: " + std::string(tree->GetName()));
    fcnc::event meas;

    tree->SetBranchAddress("year", &meas.year);
    tree->SetBranchAddress("m",&meas.m);
    tree->SetBranchAddress("mkpi",&meas.mkpi);
    tree->SetBranchAddress("sigma_m",&meas.sigma_m);

    tree->SetBranchAddress("costhetal",&meas.costhetal);
    tree->SetBranchAddress("costhetak",&meas.costhetak);
    tree->SetBranchAddress("phi",&meas.phi);

    tree->SetBranchAddress("q2",&meas.q2);
    tree->SetBranchAddress("p2",&meas.p2);
    tree->SetBranchAddress("cp_conjugate",&meas.cp_conjugate);
    //tree->SetBranchAddress("bkgcat",&meas.event_type); //ONLY USED IN TOYS
    tree->SetBranchAddress("magnet", &meas.magnet);
    tree->SetBranchAddress("weight", &meas.weight);
    tree->SetBranchAddress("delta_weight", &meas.delta_weight);
    tree->SetBranchAddress("kaon_pt", &meas.kaon_pt);
    tree->SetBranchAddress("max_mu_pt", &meas.max_mu_pt);

    std::vector<fcnc::event> result;
    unsigned int nmax = tree->GetEntries();
    if (neventsmax > 0){
        unsigned int n = TMath::Abs(neventsmax);
        if(n < nmax)nmax = n;
    }
    for (unsigned int i=0; i< nmax; i++){
        tree->GetEntry(i);
        if (!isEvtInAngleRange(&meas)) continue;
        meas.costhetal_fa = meas.costhetal;
        meas.costhetak_fa = meas.costhetak;
        meas.phi_fa = meas.phi;
        result.push_back(meas);
    }
    spdlog::debug("Read events: {0:d}",result.size());
    input->Close();
    //delete tree;
    delete input;
    return result;
};

std::vector<fcnc::event> merge_two_evtVecs(std::vector<fcnc::event> one, std::vector<fcnc::event> two){
    std::vector<fcnc::event> merged_vec;
    merged_vec.insert(merged_vec.end(), one.begin(), one.end());
    merged_vec.insert(merged_vec.end(), two.begin(), two.end());
    return merged_vec;
}

//run over the complete tree and create a tree with only the branches needed for FCNC in the correct nomenclature of the FCNCfitter
//the data and MC samples are divided into 6 sub-samples: 2 per Run, 3 per sub-decay (K+pi0, KS0pi+(DD) and KS0pi-(LL))
int convert_tuples(int job_id, std::string theFCNCpath, fcnc::options opts,std::vector<int> years){

    spdlog::info( "Converting "+get_sample_from_jobID(job_id)+" tuples of Run {0:d} to FCNC format",opts.run);

    reset_spdlog(); // back to default format

    //load the options
    fcnc::bu2kstarmumu_loader loader(&opts);

    //initialize a vector of this special class of events
    std::vector<fcnc::event> theEvents;

    if(job_id==4 || job_id ==5) years={2017}; //Only one year available for GenLvl
    //read the full tuple for each year
    std::string treePath;
    for(auto yr: years){        
        //Set the tree name
        treePath = getSelectedTuplePath(job_id,opts.run,yr);
        spdlog::debug("Processing " +treePath);
        //Load the event vector
        std::vector<fcnc::event> Events = loader.read_full_tuple(yr, treePath, get_inputTree_name(job_id),
                                                                 job_id>0, job_id>0 && opts.useMC, job_id>=4, -1);
        spdlog::debug("Inserting the vector into theEvents...");
        //Add
        theEvents.insert(theEvents.end(), Events.begin(), Events.end());
    }

    //Some log messages
    spdlog::debug("Event vector lenght: {0:d}",theEvents.size());
    if (theEvents.size()==0){
        spdlog::error("Empty event vector! Abort.");
        assert(0);
    }
    spdlog::info("Completed loading sample: Run {0:d} with years {1:s}", opts.run, convert_vector_to_string(years));

    //save events to standardised output filepath
    fcnc::save_events(theFCNCpath, theEvents);
    spdlog::info( "Completed saving sample: Run {0:d}.", opts.run );

    return 0;
}