#include "HistMan.h"

#include <TKey.h>
#include <TROOT.h>
#include <TFile.h>
#include <TH1.h>
#include <TH2.h>

#include <string>
#include <vector>
#include <iostream>
#include <cassert>
using namespace std;

  

vector<string> parse_path(const char* path)
{
    vector<string> vs;
    string p(path);
    string::iterator it1=p.begin(), it2, done=p.end();
    for (it2=find(p.begin(),done,'/'); it2 != done; it2=find(it1,done,'/')) {
        if (it1==it2) {
            ++it1;
            continue;
        }
        string dir(it1,it2);
        if (dir != ".")
            vs.push_back(dir);
        it1 = it2+1;
    }
    if (it1 != it2) {
        string dir(it1,it2);
        if (dir != ".")
            vs.push_back(string(it1,it2));
    }
    return vs;
}

TFolder& mkdir_p(TFolder& top, const char* path)
{
    if (!path || path[0] == '\0') return top;

    vector<string> vp = parse_path(path);
    TFolder* folder = &top;

    for (unsigned int ind=0; ind < vp.size(); ++ind) {
        const char* name = vp[ind].c_str();
        TFolder* tmp = dynamic_cast<TFolder*>(folder->FindObject(name));
        if (tmp) folder = tmp;
        else folder = folder->AddFolder(name,name);
    }
    return *folder;
}

static int histcount = 0;
HistMan::HistMan(TFolder* folder, bool own)
    : fFolder(folder)
    , fOwn(own)
{
    if (!fFolder) {
        fFolder = new TFolder(Form("HistFolder%d",histcount),
                              Form("Histogram Folder #%d",histcount));
    }
    fFolder->SetOwner(true);
}

HistMan::HistMan(const char* base_directory)
{
    fFolder = &this->BaseFolder();
    fFolder = &mkdir_p(*fFolder,base_directory);
    fOwn = false;
    fFolder->SetOwner(true);
}

static TFolder* directory_to_folder(TDirectory& directory)
{
    TFolder* folder = new TFolder(directory.GetName(),directory.GetName());

    TList *l = directory.GetListOfKeys();
    if (!l) return folder;

    TIter it(l->MakeIterator());
    TObject *obj=0;
    while ( (obj=it()) ) {
	TKey* key = dynamic_cast<TKey*>(obj);
	assert(key);
	obj = directory.Get(key->GetName());
	TDirectory* dir = dynamic_cast<TDirectory*>(obj);
	if (dir) {
	    TFolder* fol = directory_to_folder(*dir);
	    //cerr << "Adding folder " << fol->GetName()
	    //<< " to folder " << folder->GetName()
	    //<< " from directory " << dir->GetName() << endl;
	    folder->Add(fol);
	}
	else {
	    //cerr << "Adding object " << obj->GetName()
	    //<< " to folder " << folder->GetName() << endl;
	    folder->Add(obj);
	}
    }
    return folder;
}



HistMan::HistMan(TFile& file)
{
    TDirectory* dir = dynamic_cast<TDirectory*>(file.Get("HistMan"));
    if (!dir) {
	cerr << "Failed to find HistMan folder in " << file.GetName() << endl;
	fFolder = &this->BaseFolder();
	fFolder = &mkdir_p(*fFolder,"");
	fOwn = false;
	fFolder->SetOwner(true);
	return;
    }

    fFolder = directory_to_folder(*dir);
    gROOT->Add(fFolder);
}

HistMan::~HistMan()
{
    if (fOwn) delete fFolder; fFolder = 0;
}

TFolder& HistMan::BaseFolder()
{
    TFolder* folder = dynamic_cast<TFolder*>(gROOT->FindObjectAny("HistMan"));
    if (folder) return *folder;

    folder = new TFolder("HistMan","Base Histogram Manager Folder");
    gROOT->Add(folder);
    return *folder;
}
void HistMan::RegisterWithRoot()
{
    if (!fOwn) return;
    fOwn = false;

    TFolder& base = this->BaseFolder();
    base.Add(fFolder);
}

static void folder_to_directory(TFolder* folder, TDirectory& directory)
{
    TDirectory*& prevdir = gDirectory, *dir = 0;

    TObject* obj = directory.Get(folder->GetName());
    if (obj) {
	dir = dynamic_cast<TDirectory*>(obj);
	if (!dir) {
	    cerr << "Non-directory object \"" << folder->GetName()
		 << "\" already in directory\n";
	    return;
	}
    }
    else
	dir = directory.mkdir(folder->GetName());

    dir->cd();

    TCollection* sf = folder->GetListOfFolders();
    if (!sf) {
	prevdir->cd();
	return;
    }

    TIter it(sf->MakeIterator());
    obj=0;
    while ( (obj=it()) ) {
	TFolder *fol = dynamic_cast<TFolder*>(obj);
	if (fol) folder_to_directory(fol,*dir);
	else {
	    dir->cd();
	    obj->Write();
	}
    }
    prevdir->cd();
}

void HistMan::WriteOut(TFile& file)
{
    if (!fFolder) return;
    folder_to_directory(fFolder,file);
}
void HistMan::WriteOut(const char* filename)
{
    TFile f(filename,"recreate");
    this->WriteOut(f);
    f.Close();
}


TObject* HistMan::Adopt(const char* dirpath, TObject* obj)
{
    TH1* hist = dynamic_cast<TH1*>(obj);
    if (hist)
	hist->SetDirectory(0);
    TFolder* folder = &mkdir_p(*fFolder,dirpath);
    if (folder->FindObject(obj->GetName())) { // FindObject(Object*) not implemented!
        cerr << "Object: " << dirpath << "/" 
             << obj->GetName() << " already exists\n";
        delete obj;
        return 0;
    }
    folder->Add(obj);
    return obj;
}
bool HistMan::Fill1d(const char* pathname, Axis_t x, Stat_t w)
{
    TObject* o = fFolder->FindObject(pathname);
    if (!o) o = fFolder->FindObjectAny(pathname);
    TH1* h = dynamic_cast<TH1*>(o);
    if (!h) {
        cerr << "Fill1d(\""<<pathname<<"\") failed lookup\n";
        return false;
    }
    h->Fill(x,w);
    return true;
}
bool HistMan::Fill2d(const char* pathname, Axis_t x, Axis_t y, Stat_t w)
{
    TObject* o = fFolder->FindObject(pathname);
    if (!o) o = fFolder->FindObjectAny(pathname);
    TH2* h = dynamic_cast<TH2*>(o);
    if (!h) {
        cerr << "Fill2d(\""<<pathname<<"\") failed lookup\n";
        return false;
    }
    h->Fill(x,y,w);
    return true;
}


TObject* HistMan::GetObject(const char* pathname)
{
    //cerr << "Getting object at " << pathname << endl;
    if (!fFolder) return 0;
    vector<string> path = parse_path(pathname);

    TFolder* folder = fFolder;

    for (size_t ind=0; ind < path.size(); ++ind) {
	//cerr << "Checking " << path[ind] << endl;
	TObject* o = folder->FindObject(path[ind].c_str());
	TFolder* f = dynamic_cast<TFolder*>(o);
	if (!f) return o;
	folder = f;
    }
    return 0;
}