#include <iostream>

#include <TString.h>
#include <TStyle.h>
#include <TMath.h>
#include <TFile.h>
#include <TList.h>
#include <TH3.h>
#include <TF1.h>

namespace
{

//_______________________________________________________________________
TH3 * clone_sizes(const TH3 * src)
{
    TString name(src->GetName() + TString("_clone"));
    //
    const double xMin = src->GetXaxis()->GetXmin();
    const double xMax = src->GetXaxis()->GetXmax();
    const int xBins = src->GetNbinsX();
    //
    const double yMin = src->GetYaxis()->GetXmin();
    const double yMax = src->GetYaxis()->GetXmax();
    const int yBins = src->GetNbinsY();
    //
    const double zMin = src->GetZaxis()->GetXmin();
    const double zMax = src->GetZaxis()->GetXmax();
    const int zBins = src->GetNbinsZ();

    return new TH3D(name.Data(), name.Data(), xBins, xMin, xMax, yBins, yMin, yMax, zBins, zMin, zMax);
}

//_______________________________________________________________________
TH3 * transform_hist(const TH3 * src, double (*map)(double))
{
    TH3 * dst = clone_sizes(src);

    for(int i = 1; i <= src->GetNbinsX(); ++i)
    {
        for(int j = 1; j <= src->GetNbinsY(); ++j)
        {
            for(int k = 1; k <= src->GetNbinsZ(); ++k)
            {
                const double val = src->GetBinContent(i, j, k);
                dst->SetBinContent(i, j, k, map(val));
            }
        }
    }

    return dst;
}

//_______________________________________________________________________
double log_scale(double val)
{
    if(val <= 0.)
        val = 10e-50;
    return TMath::Log(val);
}

//_______________________________________________________________________
double transfer_function(const double * px, const Double_t *)
{
    const double x = *px;
    if(x > -25.)
        return 0.5;
    if(x < -25. && x > -35.)
        return 0.2;
    if(x < -35. && x > -50.)
        return 0.1;
    return 0.04;
}

}

//_______________________________________________________________________
void th3()
{
    TFile input("rel_energy.root", "READ");
    TH3 * energyHist = dynamic_cast<TH3*>(input.Get("energy_hist"));
    if(!energyHist)
    {
        std::cout<<"Could not find energy_hist or type is not TH3 derived\n";
        return;
    }

    gStyle->SetCanvasPreferGL(1);
    TH3 * energyHistModified = transform_hist(energyHist, log_scale);

    //Transfer function.
    TList * lf = energyHistModified->GetListOfFunctions();
    if(lf)
    {
        TF1 * tf = new TF1("TransferFunction", transfer_function);
        lf->Add(tf);
    }

    gStyle->SetPalette(1);
    energyHistModified->SetDirectory(0);
    energyHistModified->Draw("glcolz");
}