#include "TFile.h"
#include "TTree.h"
#include "TH1F.h"
#include "TStyle.h"
#include "TH2F.h"
#include "TCanvas.h"
#include "TTreeReader.h"

#include "AutoFit.C"
#include "SaveTH1IntoText.C"

#include "TROOT.h"
#include "ROOT/TThreadedObject.hxx"
#include "ROOT/TTreeProcessorMT.hxx" 
#include <memory>
#include <vector>

// mod = 0, 1, 2
// row = 1, 2, 3
// pstrip = 0 - 127

#define MAXMOD 3
#define MAXROW 3
#define MAXPStrip 128

const double ExRange[2] = {-0.5, 3.5};
const double ExRel = 0.05; //MeV
const int nBin = (ExRange[1] - ExRange[0])/ExRel;

#define SAVETREE true

//^====================================
TH1F * haha[MAXMOD][MAXROW][MAXPStrip];

TH2F * dudu[MAXMOD][MAXROW];

TH1F * hEx;

//^====================================


void parallel(){


  /// Load peaks
  std::vector<double> peaks[MAXMOD][MAXROW][MAXPStrip];

  FILE * inList = fopen("fit_result.txt", "r");

  char buffer[256];

  while( fgets(buffer, sizeof(buffer), inList) ){

    std::vector<std::string> haha = SplitStrAF(buffer, " ");

    // printf("%zu| %s", haha.size(), buffer);
    int m = atoi(haha[0].c_str());
    int r = atoi(haha[1].c_str()) - 1;
    int p = atoi(haha[2].c_str());

    for( int i = 3; i < (int) haha.size(); i++ ) peaks[m][r][p].push_back(atof(haha[i].c_str()));
  }

  fclose(inList);

  //^=========================== histogram
  for( int i = 0; i < MAXMOD; i++ ){
    for( int j = 0; j < MAXROW; j++ ){

      dudu[i][j] = new TH2F(Form("dudu_%d%d", i, j+1), Form("%d-%d; pStrip; Ex", i, j+1), 128, 0, 128, nBin, ExRange[0], ExRange[1]);

      for( int p = 0; p < MAXPStrip; p ++ ){
        haha[i][j][p] = new TH1F(Form("h_%d%d_%d", i, j+1, p), Form("%d-%d-%d", i, j+1, p), nBin, ExRange[0], ExRange[1]);
      }
    }
  }

  hEx = new TH1F("hEx", "Ex", 10*nBin, ExRange[0], ExRange[1]);

  ROOT::TThreadedObject<TH1F> phEx("hEx", "Ex", 10*nBin, ExRange[0], ExRange[1]);

  //^=========================== tree
  int nThreads = 4;
  ROOT::EnableImplicitMT(nThreads);


  TString inFileName = "R9-172_summed_tree.root";

  ROOT::TTreeProcessorMT tp(inFileName, "rxtree", nThreads);

  const double y1 = 0.000;
  const double ya = 0.854;
  const double y2 = 2.005;

  TFile * newFile = nullptr;

  ROOT::TThreadedObject<TTree> pTree;
  std::shared_ptr<TTree> pTree_a[nThreads];

  Double_t Ex_org;
  Double_t Ex_C;
  Int_t mod_C;
  Int_t row_C;
  Int_t pStrip_C;
  Double_t z;
  Double_t thetaCM;

  if( SAVETREE ){  
    newFile = new TFile("test.root", "recreate");

    // newTree = new TTree("tree", "Calibred Ex");
    // newTree->Branch("Ex_org", &Ex_org, "Ex_org/D");
    // newTree->Branch("mod",    &mod_C, "mod/I");
    // newTree->Branch("row",    &row_C, "row/I");
    // newTree->Branch("pstrip", &pStrip_C, "pstrip/I");
    // newTree->Branch("Ex", &Ex_C, "Ex/D");
    // newTree->Branch("z", &z, "z/D");
    // newTree->Branch("thetaCM", &thetaCM, "thetaCM/D");

    for( int i = 0; i < nThreads; i++ ){

      pTree_a[i] = std::make_shared<TTree>("tree", "Calibrated Ex");
    
      pTree_a[i]->Branch("Ex_org", &Ex_org, "Ex_org/D");
      pTree_a[i]->Branch("mod",    &mod_C, "mod/I");
      pTree_a[i]->Branch("row",    &row_C, "row/I");
      pTree_a[i]->Branch("pstrip", &pStrip_C, "pstrip/I");
      pTree_a[i]->Branch("Ex", &Ex_C, "Ex/D");
      pTree_a[i]->Branch("z", &z, "z/D");
      pTree_a[i]->Branch("thetaCM", &thetaCM, "thetaCM/D");

      pTree.SetAtSlot(i, pTree_a[i]);

    }
  }

  std::mutex treeMuTex;

  //^========================= define the process in each thread
  auto myFunction = [&](TTreeReader &reader) {

    TTreeReaderValue<Double_t>    Ex = {reader, "Ex"};
    TTreeReaderValue<Int_t>     mod = {reader, "mod"};
    TTreeReaderValue<Int_t>     row = {reader, "row"};
    TTreeReaderValue<Int_t>  pStrip = {reader, "pstrip"};
    TTreeReaderValue<Int_t>    ebis = {reader, "ebis_on"};
    TTreeReaderValue<Double_t>    ThetaCM = {reader, "ThetaCM"};
    TTreeReaderValue<Double_t>    Zmeasured = {reader, "Zmeasured"};

    while( reader.Next() ){

      if( *ebis == 0 ) continue;
      if( *row == 0 ) continue;

      int m = *mod;
      int r = *row -1;
      int p = *pStrip;

      if(  reader.GetCurrentEntry()  < 100) printf("%llu | %d %d %d | %f\n", reader.GetCurrentEntry(),  m, r, p, *Ex/1000.);

      if( peaks[m][r][p].size() == 0 ) continue;

      double slope = 1;
      double offset = 0;

      if( peaks[m][r][p].size() == 1 ) {
        if( peaks[m][r][p][0] < 1 ){
          offset = -peaks[m][r][p][0];
        }else{
          offset = y2 - peaks[m][r][p][0];
        }
      }
      if( peaks[m][r][p].size() >= 2 ){
        double x1 = peaks[m][r][p][0];
        double x2 = peaks[m][r][p].back();

        if( x1 < 0.2 ) {
          if( x2 < 1.8 ) {
            slope = (ya-y1)/(x2-x1);
          }else{
            slope = (y2-y1)/(x2-x1);
          }
          offset = y1 - slope * x1;
        }
        if( x1 >= 0.2 ) {
          // printf("%d %d %d| x1 : %f, ya : %f \n", m, r, p, x1, ya);
          slope = (y2-ya)/(x2-x1);
          offset = ya - slope * x1;
        }
      }

      Ex_org = *Ex/1000.;
      mod_C = m;
      row_C = r;
      pStrip_C = p;
      z = *Zmeasured;
      thetaCM = *ThetaCM;

      Ex_C = slope * (*Ex/1000.) + offset;


      treeMuTex.lock();
      
      haha[m][r][p]->Fill(Ex_C);

      dudu[m][r]->Fill(p, Ex_C);

      hEx->Fill( *Ex/1000.);

      // if( SAVETREE ) {
      //   newFile->cd();
      //   treeMuTex.lock();
      //   newTree->Fill();
      // }

      treeMuTex.unlock();

      phEx->Fill(Ex_C);

      if( SAVETREE ) {
        pTree->Fill();
      }

    }

  };

  //^========================= process data
  tp.Process(myFunction);

  if( SAVETREE ) {
    newFile->cd();

    TList list;
    for (int i = 0; i < pTree.GetNSlots(); ++i) {
      if( pTree.GetAtSlot(i)->GetEntries() == 0 ) continue;
      list.Add(pTree.GetAtSlot(i).get());
    }

    TTree* mergedTree = TTree::MergeTrees(&list);
    mergedTree->Write("", TObject::kWriteDelete);
    printf("============== %llu\n", mergedTree->GetEntries());
    
    newFile->Close();
  }


  auto kEx = phEx.Merge();


  // kEx->SetLineColor(2);
  // kEx->DrawCopy();
  // hEx->Draw("same");

  // gStyle->SetOptStat("");

  // TCanvas * canvas = new TCanvas("canvas", "canvas", 2100, 1500);
  // canvas->Divide(3,3);
  
  // canvas->cd(1); dudu[0][0]->Draw("colz");
  // canvas->cd(2); dudu[0][1]->Draw("colz");
  // canvas->cd(3); dudu[0][2]->Draw("colz");

  // canvas->cd(4); dudu[1][0]->Draw("colz");
  // canvas->cd(5); dudu[1][1]->Draw("colz");
  // canvas->cd(6); dudu[1][2]->Draw("colz");
  
  // canvas->cd(7); dudu[2][0]->Draw("colz");
  // canvas->cd(8); dudu[2][1]->Draw("colz");
  // canvas->cd(9); dudu[2][2]->Draw("colz");

  // TCanvas * canvasEx = new TCanvas("canvasEx", "canvasEx", 800, 400);
  // canvasEx->cd();
  // hEx->Draw();
  
  // TCanvas * chaha = new TCanvas("chaha", "c haah", 500, 500);
  // haha[0][1][100]->Draw();

}