#include "TFile.h"
#include "TTree.h"
#include "TObject.h"
#include "TChain.h"
#include <iostream>
#include "TRandom3.h"

// Step 1: Define a class to be stored in the tree
class BaseEvent : public TObject
{
public:
    BaseEvent() {}
    virtual ~BaseEvent() {}
};

class AEvent1 : public BaseEvent
{
public:
    std::vector<Double_t> value;

    AEvent1() : value({}) {}
    virtual ~AEvent1() { value.clear(); }

    ClassDef(AEvent1, 1); // ROOT macro for I/O
};

class AEvent2 : public BaseEvent
{
public:
    std::vector<Double_t> value;

    AEvent2() : value({}) {}
    virtual ~AEvent2() { value.clear(); }

    ClassDef(AEvent2, 1); // ROOT macro for I/O
};

class AEvent3 : public BaseEvent
{
public:
    std::vector<Double_t> value;

    AEvent3() : value({}) {}
    virtual ~AEvent3() { value.clear(); }

    ClassDef(AEvent3, 1); // ROOT macro for I/O
};


class AEvent4 : public BaseEvent
{
public:
    std::vector<Double_t> value;

    AEvent4() : value({}) {}
    virtual ~AEvent4() { value.clear(); }

    ClassDef(AEvent4, 1); // ROOT macro for I/O
};

class AEvent5 : public BaseEvent
{
public:
    std::vector<Double_t> value;

    AEvent5() : value({}) {}
    virtual ~AEvent5() { value.clear(); }

    ClassDef(AEvent5, 1); // ROOT macro for I/O
};

class AEvent6 : public BaseEvent
{
public:
    std::vector<Double_t> value;

    AEvent6() : value({}) {}
    virtual ~AEvent6() { value.clear(); }

    ClassDef(AEvent6, 1); // ROOT macro for I/O
};



class BufferStuff
{
public:
    BufferStuff() {};
    ~BufferStuff()
    {
        for (auto el : evMap)
        {
            delete el.second;
            el.second = nullptr;
        }
        evMap.clear();
    }

public:
    std::map<Int_t, BaseEvent *> evMap;
};

// Step 2: Linkdef file is needed if compiled as shared lib; skip here

void write_tree_with_AnEvent()
{
    // Step 3: Create the output file
    TFile *outfile = new TFile("anEventTree.root", "RECREATE");

    // Step 4: Create the tree
    TTree *tree = new TTree("tree", "Tree storing AnEvent");

    TRandom3 rnmd(42);
    // Step 5: Create an AnEvent object and connect it to the tree
    AEvent1 *aevent1 = new AEvent1();
    AEvent2 *aevent2 = new AEvent2();
    AEvent3 *aevent3 = new AEvent3();
    AEvent4 *aevent4 = new AEvent4();
    AEvent5 *aevent5 = new AEvent5();
    AEvent6 *aevent6 = new AEvent6();
    tree->Branch("AEvent1", &aevent1);
    tree->Branch("AEvent2", &aevent2);
    tree->Branch("AEvent3", &aevent3);
    tree->Branch("AEvent4", &aevent4);
    tree->Branch("AEvent5", &aevent5);
    tree->Branch("AEvent6", &aevent6);
    
    // Step 6: Fill the tree
    Long64_t nEntries = 100;
    for (Long64_t i = 0; i < nEntries; ++i)
    {
        if (i % (nEntries / 10) == 0)
        {
            std::cout << "Writting " << i << "/" << nEntries << std::endl;
        }
        for (int j = 0; j < 100; ++j)
        {
            aevent1->value.push_back(rnmd.Uniform()); // assign a different value
            aevent2->value.push_back(rnmd.Uniform()); // assign a different value
            aevent3->value.push_back(rnmd.Uniform()); // assign a different value
            aevent4->value.push_back(rnmd.Uniform()); // assign a different value
            aevent5->value.push_back(rnmd.Uniform()); // assign a different value
            aevent6->value.push_back(rnmd.Uniform()); // assign a different value
        }

        tree->Fill();
        aevent1->value.clear();
        aevent2->value.clear();
        aevent3->value.clear();
        aevent4->value.clear();
        aevent5->value.clear();
        aevent6->value.clear();
    }

    // Step 7: Write and clean up
    tree->Write();
    outfile->Close();
    delete outfile;
    delete aevent1;
    delete aevent2;
    delete aevent3;
    delete aevent4;
    delete aevent5;
    delete aevent6;
}

#ifdef __ROOTCLING__
#pragma link C++ class AnEvent + ;
#endif

// Optional main for standalone compilation
int do_stuff()
{
    write_tree_with_AnEvent();
    return 0;
}

void copytree3()
{
    // Get old file, old tree and set top branch address
    //    TString dir = "$ROOTSYS/test/Event.root";
    //    gSystem->ExpandPathName(dir);
    const auto filename = "anEventTree.root";
    const auto treename = "tree";

    // TFile oldfile(filename);
    // TTree *oldtree;
    // oldfile.GetObject("tree", oldtree);
    TChain *chain = new TChain("Chain");
    chain->AddFile(filename, TTree::kMaxEntries, treename);
    const auto nentries = chain->GetEntries();

    // Create a new file + a clone of old tree in new file
    TFile newfile("small.root", "recreate");
    auto newtree = chain->CloneTree(0);
    newtree->SetDirectory(newfile.GetDirectory("/"));

    std::map<Int_t, TString> branchNamesMap({{0, "AEvent1"}, 
                                             {1, "AEvent2"}, 
                                             {2, "AEvent3"}, 
                                             {3, "AEvent4"}, 
                                             {4, "AEvent5"}, 
                                             {5, "AEvent6"}});

    for (auto i : ROOT::TSeqI(nentries))
    {
        std::map<Int_t, BaseEvent *> mapToCopy;
        mapToCopy[0] = nullptr; // aevent
        mapToCopy[1] = nullptr; // bevent
        mapToCopy[2] = nullptr; // aevent
        mapToCopy[3] = nullptr; // bevent
        mapToCopy[4] = nullptr; // aevent
        mapToCopy[5] = nullptr; // bevent

        BufferStuff *bufStuf = new BufferStuff();
        bufStuf->evMap = mapToCopy;

        for (auto &el : bufStuf->evMap)
        {
            chain->SetBranchAddress(branchNamesMap[el.first].Data(), &(el.second));
            // TBranch *br = oldtree->GetBranch(branchNamesMap[el.first].Data());
            // br->SetAddress(&(el.second));
        }

        if (i % (nentries / 10) == 0)
        {
            std::cout << "Writting " << i << "/" << nentries << std::endl;
        }

        chain->GetEntry(i);
        for (int i = 0; i < 6; ++i){
            std::cout << bufStuf->evMap[i]->ClassName() << " ";
        }
        std::cout << std::endl;
        (static_cast<AEvent1 *>(bufStuf->evMap[0]))->value.clear();
        (static_cast<AEvent1 *>(bufStuf->evMap[0]))->value.push_back(1.);
        (static_cast<AEvent2 *>(bufStuf->evMap[1]))->value.clear();
        (static_cast<AEvent2 *>(bufStuf->evMap[1]))->value.push_back(2.);
        (static_cast<AEvent3 *>(bufStuf->evMap[2]))->value.clear();
        (static_cast<AEvent3 *>(bufStuf->evMap[2]))->value.push_back(3.);
        (static_cast<AEvent4 *>(bufStuf->evMap[3]))->value.clear();
        (static_cast<AEvent4 *>(bufStuf->evMap[3]))->value.push_back(4.);
        (static_cast<AEvent5 *>(bufStuf->evMap[4]))->value.clear();
        (static_cast<AEvent5 *>(bufStuf->evMap[4]))->value.push_back(5.);
        (static_cast<AEvent6 *>(bufStuf->evMap[5]))->value.clear();
        (static_cast<AEvent6 *>(bufStuf->evMap[5]))->value.push_back(6.);

        newtree->Fill();
        delete bufStuf;
        bufStuf = nullptr;

        // event = nullptr;
        //   event = nullptr;
        //   oldtree->SetBranchAddress("event", &event);
    }

    newfile.Write();
}