#include "TFile.h"
#include "TRandom.h"
#include "TFile.h"
#include "TTree.h"
#include "TObjString.h"
#include "TObject.h"

#include <iostream>
#include <vector>
#include <functional>
#include <cmath>
#include <thread>
#include <mutex>

class TreeStructure {

  public:

    void clear() {
        etaVec.clear();
        phiVec.clear();
        ptVec.clear();
    }

    std::vector<float> etaVec;
    std::vector<float> phiVec;
    std::vector<float> ptVec;
};

using StringAttributes = std::vector<std::pair<const char *, std::string>>;

class Cache {

  public:

    Cache(const std::string filename, const StringAttributes& attributes)
        : file_(TFile::Open((filename + ".root").c_str(), "RECREATE"))
        , tree_(new TTree("tree", "tree"))
    {
        tree_->Branch("eta", &ts_.etaVec);
        tree_->Branch("phi", &ts_.phiVec);
        tree_->Branch("pt" , &ts_.ptVec);

        for(auto const& attr : attributes) {
            TObjString* s = new TObjString(attr.second.c_str());
            s->Write(attr.first);
        }

        // Try to write an integer to the Root file
        TObjString* n = new TObjString(std::to_string(7).c_str());
        n->Write("N_EVENTS");
    }

    ~Cache() {

        tree_->Print();
        tree_->Write();

        delete file_;
    }

    void fill(TreeStructure && ts, std::vector<TreeStructure>& buffer) {
        buffer.push_back(std::move(ts));
        if (buffer.size() > maxBufferSize_) write(buffer);
    }

    void write(std::vector<TreeStructure>& buffer) {

        std::lock_guard<std::mutex> lock(mutex_);

        size_t n = buffer.size();
        for(auto& entry : buffer) {
            ts_ = std::move(entry);
            tree_->Fill();
        }
        if(n > 0) buffer.clear();
    }

  private:

    TFile * const file_;
    TTree * const tree_;

    mutable std::mutex mutex_;

    TreeStructure ts_;

    const size_t maxBufferSize_ = 512;
};


void fillEvents(int nEvents, Cache* cache)
{
    TRandom *eventGenerator = new TRandom();
    std::vector<TreeStructure> buffer;

    for(int iEvent = 0; iEvent < nEvents; ++iEvent)
    {
        TreeStructure ts;

        int nLeptons = 4 + eventGenerator->Integer(4);
        for(int iLepton = 0; iLepton < nLeptons; ++iLepton)
        {
            ts.phiVec.push_back(eventGenerator->Uniform() * 2*M_PI - M_PI);
            ts.etaVec.push_back(eventGenerator->Uniform() * 5.0 - 2.5    );
            ts.ptVec .push_back(eventGenerator->Exp(10) + 20.            );
        }
        cache->fill(std::move(ts), buffer);
    }

    cache->write(buffer);
}

int main()
{
    int nEvents = 100000;
    int nThreads = 8;

    auto cache = std::make_shared<Cache>("data", StringAttributes{
            {"jonas", "idiot"}
        });

    auto fillEventsThreaded = std::bind(fillEvents, nEvents/nThreads, cache.get());

    std::vector<std::thread> threads;

    for(int i = 0; i < nThreads; ++i) threads.emplace_back(fillEventsThreaded);
    for(auto& thread : threads) thread.join();
}