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


#include <unistd.h>
#include <ios>
#include <iostream>
#include <fstream>
#include <string>

//////////////////////////////////////////////////////////////////////////////
//
// process_mem_usage(double &, double &) - takes two doubles by reference,
// attempts to read the system-dependent data for a process' virtual memory
// size and resident set size, and return the results in KB.
//
// On failure, returns 0.0, 0.0

void process_mem_usage(double& vm_usage, double& resident_set)
{
   using std::ios_base;
   using std::ifstream;
   using std::string;

   vm_usage     = 0.0;
   resident_set = 0.0;

   // 'file' stat seems to give the most reliable results
   //
   ifstream stat_stream("/proc/self/stat",ios_base::in);

   // dummy vars for leading entries in stat that we don't care about
   //
   string pid, comm, state, ppid, pgrp, session, tty_nr;
   string tpgid, flags, minflt, cminflt, majflt, cmajflt;
   string utime, stime, cutime, cstime, priority, nice;
   string O, itrealvalue, starttime;

   // the two fields we want
   //
   unsigned long vsize;
   long rss;

   stat_stream >> pid >> comm >> state >> ppid >> pgrp >> session >> tty_nr
               >> tpgid >> flags >> minflt >> cminflt >> majflt >> cmajflt
               >> utime >> stime >> cutime >> cstime >> priority >> nice
               >> O >> itrealvalue >> starttime >> vsize >> rss; // don't care about the rest

   stat_stream.close();

   long page_size_kb = sysconf(_SC_PAGE_SIZE) / 1024; // in case x86-64 is configured to use 2MB pages
   vm_usage     = vsize / 1024.0;
   resident_set = rss * page_size_kb;
}



// 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 filename2 = "anEventTree2.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);
    chain->AddFile(filename2, 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;
        // bufStuf->evMap[0] = new AEvent1();
        // bufStuf->evMap[1] = new AEvent2();
        // bufStuf->evMap[2] = new AEvent3();
        // bufStuf->evMap[3] = new AEvent4();
        // bufStuf->evMap[4] = new AEvent5();
        // bufStuf->evMap[5] = new AEvent6();
        AEvent1 *ev1 = nullptr;
        AEvent2 *ev2 = nullptr;
        AEvent3 *ev3 = nullptr;
        AEvent4 *ev4 = nullptr;
        AEvent5 *ev5 = nullptr;
        AEvent6 *ev6 = nullptr;
        bufStuf->evMap[0] = ev1;
        bufStuf->evMap[1] = ev2;
        bufStuf->evMap[2] = ev3;
        bufStuf->evMap[3] = ev4;
        bufStuf->evMap[4] = ev5;
        bufStuf->evMap[5] = ev6;
        

        std::cout << "Chain file number " << chain->GetFileNumber() << std::endl;
        
        for (auto &el : bufStuf->evMap)
        {
            // chain->SetBranchAddress(branchNamesMap[el.first].Data(), &(el.second));
            chain->SetBranchAddress(branchNamesMap[el.first].Data(), &( bufStuf->evMap[el.first] ));
        }

        if (i % (nentries / 10) == 0)
        {
            Double_t vm_usage, resident_set;
            process_mem_usage(vm_usage, resident_set);
            std::cout << "Writting " << i << "/" << nentries << " VM: "<< vm_usage/1024. << " RM: "<< resident_set/1024. << 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.);
        // for (std::size_t is = 0; is < (static_cast<AEvent1 *>(bufStuf->evMap[0]))->value.size(); ++is){
        //     std::cout << (static_cast<AEvent1 *>(bufStuf->evMap[0]))->value.at(is) << std::endl;
        // }

        (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();
        chain->ResetBranchAddresses();
        
        delete bufStuf;
        bufStuf = nullptr;

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

    newfile.Write();
}