#include <iostream>
#include <fstream>
#include <cmath>
//~ #include <set>
#include <numeric>

using std::cout;
using std::endl;

#include "TMath.h"
#include "TH2.h"
#include "TStyle.h"
#include "TLatex.h"
#include "TTree.h"
#include "TPaveStats.h"
#include "TFile.h"
#include "TCanvas.h"
#include "TROOT.h"

const UInt_t kWidth=1200;//Canvas size for png output resolution
const UInt_t kHeight=kWidth*0.75;

void ResizeCanvas(TCanvas* const c, const UInt_t width, const UInt_t height)
{
    if(!c) return;
    if(gROOT->IsBatch())
    {
        c->SetCanvasSize(width,height);
    }
    else
    {
        c->SetWindowSize(width,height);
    }
}

TCanvas* CreateCanvas(const TString name)
{
    TCanvas* const c = new TCanvas(name,name,100 ,100, kWidth, kHeight);
    ResizeCanvas(c,kWidth,kHeight);

    c->SetFillColor(0);
    c->SetFillStyle(0);
    c->SetFrameFillColor(0);
    c->SetFrameFillStyle(0);
    c->Draw();

    return c;
}

void reformat_stats(TH1* h)
{
    gPad->Update();
    TPaveStats *st = ((TPaveStats*)h->FindObject("stats"));
    st->SetFillStyle(0);
    st->SetX1NDC(0.69);
    st->SetX2NDC(0.89);
    st->SetY1NDC(0.85);
    st->SetY2NDC(0.99);
}

void draw(const TString outputDir="/tmp/")
{
    const ULong64_t nChunks = 200;
    const ULong64_t chunkSize = 512ULL*1024*1024; // 512 MB was the readout chunk size in the whole ELBE campaign
    const ULong64_t chunkSamples = chunkSize/2;
    const ULong64_t windowSamples = 250; // samples
    const ULong64_t preTrigger = 10; // samples
    const double time_step_ns = 0.4; // ns
    const double bin_width_time = 0.001; //s

    TH2D* h2_trig_raw = new TH2D("h2_trig_raw", "",  windowSamples, -time_step_ns*preTrigger, (windowSamples-preTrigger)*time_step_ns, 16384, 0, 16384);
    h2_trig_raw->SetContour(255);
    h2_trig_raw->GetXaxis()->SetTitle("Time relative to trigger / ns");
    h2_trig_raw->GetYaxis()->SetTitle("ADC");
    h2_trig_raw->SetEntries(10);
    const UInt_t nbins_time = TMath::Nint( nChunks * chunkSamples * time_step_ns / 1e9 / bin_width_time);
    const double length_x_time = bin_width_time * nbins_time; // s
    TH1D* h_abs_time = new TH1D("h_abs_time", "", nbins_time, 0, length_x_time);
    h_abs_time->GetXaxis()->SetTitle("Time / s");
    h_abs_time->GetYaxis()->SetTitle("Count rate / cps");
    h_abs_time->Fill(1e12);
    TH2D* h2_abs_raw = new TH2D("h2_abs_raw", "", nbins_time, 0, length_x_time, 16384, 0, 16384);
    h2_abs_raw->SetContour(255);
    h2_abs_raw->GetXaxis()->SetTitle("Time / s");
    h2_abs_raw->GetYaxis()->SetTitle("ADC");

    TLatex* text = new TLatex(0.2, 0.04, "/a/b/cdee/fff/");
    text->SetNDC();  // Set NDC coordinate system
    text->SetTextAlign(21);  // Left alignment (10) + bottom (3) = 13 (https://root.cern/doc/v608/TAttText_8h_source.html line 58)
    text->SetTextSize(0.02);
    TFile* out = new TFile(outputDir+"/plots.root", "RECREATE");
    auto c2 = CreateCanvas("c_trig_raw");
    c2->SetSupportGL(true);
    c2->SetGridx();
    c2->SetGridy();
    c2->SetLogz();
    h2_trig_raw->Draw("COLZ");
    reformat_stats(h2_trig_raw);
    text->Draw();
    c2->SaveAs(outputDir + "/" + c2->GetName() + ".png");
    c2->Write();
    c2 = CreateCanvas("c_abs_time");
    c2->SetSupportGL(true);
    c2->SetGridx();
    c2->SetGridy();
    c2->SetLogy();
    h_abs_time->Draw();
    reformat_stats(h_abs_time);
    text->Draw();
    c2->SaveAs(outputDir + "/" + c2->GetName() + ".png");
    c2->Write();
    c2 = CreateCanvas("c_abs_raw");
    c2->SetSupportGL(true);
    c2->SetGridx();
    c2->SetGridy();
    c2->SetLogz();
    h2_abs_raw->Draw("COLZ");
    reformat_stats(h2_abs_raw);
    text->Draw();
    c2->SaveAs(outputDir + "/" + c2->GetName() + ".png");
    c2->Write();
    out->Close();
}