/* ======================================================================== *\
!
! *
! * This file is part of MARS, the MAGIC Analysis and Reconstruction
! * Software. It is distributed to you in the hope that it can be a useful
! * and timesaving tool in analysing Data of imaging Cerenkov telescopes.
! * It is distributed WITHOUT ANY WARRANTY.
! *
! * Permission to use, copy, modify and distribute this software and its
! * documentation for any purpose is hereby granted without fee,
! * provided that the above copyright notice appear in all copies and
! * that both that copyright notice and this permission notice appear
! * in supporting documentation. It is provided "as is" without express
! * or implied warranty.
! *
!
!
!   Author(s): Thomas Bretz  9/2002 <mailto:tbretz@astro.uni-wuerzburg.de>
!
!   Copyright: MAGIC Software Development, 2000-2002
!
!
\* ======================================================================== */

/////////////////////////////////////////////////////////////////////////////
//                                                                         //
// MEnergyEstParam                                                         //
//                                                                         //
// Task to estimate the energy using a parametrization.                    //
// Make sure, that all source dependant containers are based on the same   //
// set of basic hillas parameters                                          //
//                                                                         //
/////////////////////////////////////////////////////////////////////////////
#include "MEnergyEstParam.h"

#include "MParList.h"

#include "MMcEvt.hxx"
#include "MHMatrix.h"
#include "MHillasSrc.h"
#include "MEnergyEst.h"

#include "MLog.h"
#include "MLogManip.h"

ClassImp(MEnergyEstParam);

// --------------------------------------------------------------------------
//
// Initialize the coefficients with (nonsense) values
//
void MEnergyEstParam::InitCoefficients()
{
    fA[0] =   39.667402; // [cm]
    fA[1] =  143.081912; // [cm]
    fA[2] = -395.501677; // [cm]
    fA[3] =    0.006193;

    fB[0] =   45.037867; // [GeV]
    fB[1] =    0.087222; // [GeV]
    fB[2] =   -0.208065; // [GeV/cm]
    fB[3] =   78.164942; // [GeV]
    fB[4] = -159.361283; // [GeV]
    fB[5] =   -0.130455; // [GeV/cm]
    fB[6] =    0.051139;
}

// --------------------------------------------------------------------------
//
// Default constructor. Give the name of the parameter container (MHillas)
// storing wisth, length and size (Default="MHillas").
// For the Zenith Angle MMcEvt.fTheta is used.
//
MEnergyEstParam::MEnergyEstParam(const char *hillas, const char *name, const char *title)
   : fMatrix(NULL), fA(4), fB(7)
{
    fName  = name  ? name  : "MEnergyEstParam";
    fTitle = title ? title : "Task to estimate the energy";

    fHillasName = hillas;

    fPairs     = new TList;
    fEnergy    = new TList;
    fHillasSrc = new TList;

    fPairs->SetOwner();

    InitCoefficients();

    AddToBranchList("MMcEvt.fTheta");
    AddToBranchList(fHillasName+"fSize");
    AddToBranchList(fHillasName+"fWidth");
    AddToBranchList(fHillasName+"fLength");
}


// --------------------------------------------------------------------------
//
// Destructor.
//
MEnergyEstParam::~MEnergyEstParam()
{
    delete fPairs;
    delete fEnergy;
    delete fHillasSrc;
}

// --------------------------------------------------------------------------
//
// Check for all necessary parameter containers.
//
Bool_t MEnergyEstParam::PreProcess(MParList *plist)
{
    if (!fMatrix)
    {
        fHillas = (MHillas*)plist->FindObject(fHillasName, "MHillas");
        if (!fHillas)
        {
            *fLog << err << dbginf << fHillasName << " [MHillas] not found... aborting." << endl;
            return kFALSE;
        }

        fMc = (MMcEvt*)plist->FindObject("MMcEvt");
        if (!fMc)
        {
            *fLog << err << dbginf << "MMcEvt not found... aborting." << endl;
            return kFALSE;
        }
    }

    TObject *obj = NULL;
    TIter Next(fPairs);

    while ((obj=Next()))
    {
        TObject *o = plist->FindCreateObj(obj->GetTitle(), "MEnergyEst");
        if (!o)
            return kFALSE;

        if (!fEnergy->FindObject(obj->GetTitle()))
            fEnergy->Add(o);

        if (fMatrix)
            continue;

        o = plist->FindObject(obj->GetName(), "MHillasSrc");
        if (!o)
        {
            *fLog << err << dbginf << obj->GetName() << " not found... aborting." << endl;
            return kFALSE;
        }
        fHillasSrc->Add(o);
    }

   return kTRUE;    
}

// --------------------------------------------------------------------------
//
// Set the four coefficients for the estimation of the impact parameter.
// Argument must ba a TArrayD of size 4.
//
void MEnergyEstParam::SetCoeffA(TArrayD arr)
{
    if (arr.GetSize() != fA.GetSize())
    {
        *fLog << err << dbginf << "Error - Wrong number of coefficients!" << endl;
        return;
    }

    fA = arr;
}

// --------------------------------------------------------------------------
//
// Set the four coefficients for the estimation of the energy.
// Argument must ba a TArrayD of size 7.
//
void MEnergyEstParam::SetCoeffB(TArrayD arr)
{
    if (arr.GetSize() != fB.GetSize())
    {
        *fLog << err << dbginf << "Error - Wrong number of coefficients!" << endl;
        return;
    }

    fB = arr;
}

// --------------------------------------------------------------------------
//
// Returns the mapped value from the Matrix
//
Double_t MEnergyEstParam::GetVal(Int_t i) const
{
    return (*fMatrix)[fMap[i]];
}

// --------------------------------------------------------------------------
//
// You can use this function if you want to use a MHMatrix instead of the
// given containers. This function adds all necessary columns to the
// given matrix. Afterward you should fill the matrix with the corresponding
// data (eg from a file by using MHMatrix::Fill). If you now loop
// through the matrix (eg using MMatrixLoop) MEnergyEstParam::Process
// will take the values from the matrix instead of the containers.
//
void MEnergyEstParam::InitMapping(MHMatrix *mat)
{
    if (fMatrix)
        return;

    fMatrix = mat;

    fMap[0] = fMatrix->AddColumn("MMcEvt.fTheta");
    fMap[1] = fMatrix->AddColumn(fHillasName+".fWidth");
    fMap[2] = fMatrix->AddColumn(fHillasName+".fLength");
    fMap[3] = fMatrix->AddColumn(fHillasName+".fSize");

    Int_t col = 4;
    TIter Next(fPairs);
    TObject *o = NULL;
    while ((o=Next()))
        fMap[col++] = fMatrix->AddColumn(TString(o->GetName())+".fDist");
}

// --------------------------------------------------------------------------
//
// Add a pair of input/output containers.
//   eg. Add("MHillasSrc", "MEnergyEst");
// Usefull if you want to estimate the stuff for the source and antisource
//
void MEnergyEstParam::Add(const TString hillas, const TString energy)
{
    fPairs->Add(new TNamed(hillas, energy));

    AddToBranchList(hillas+".fDist");
}

// --------------------------------------------------------------------------
//
//  Estimates the impact parameter and energy using a parameterization
//  (see code)
//
Bool_t MEnergyEstParam::Process()
{
    const Double_t theta  = fMatrix ? GetVal(0) : fMc->GetTheta();
    const Double_t width  = fMatrix ? GetVal(1) : fHillas->GetWidth();
    const Double_t length = fMatrix ? GetVal(2) : fHillas->GetLength();
    const Double_t size   = fMatrix ? GetVal(3) : fHillas->GetSize();

    const Double_t k   = 1/cos(theta);
    const Double_t k2  = k*k;

    const Double_t i0 = k * (1+fA[3]*k)/(1+fA[3]);
    const Double_t i1 = fA[0] + fA[2]*width;

    const Double_t e0 = k2 * (1+fB[6]*k2)/(1+fB[6]);
    const Double_t e1 = fB[0] + fB[1]*size + fB[3]*length + fB[4]*size*length;
    const Double_t e2 = fB[2] + fB[5]*length;

    TIter NextH(fHillasSrc);
    TIter NextE(fEnergy);

    Int_t col = 4;

    while (1)
    {
        MEnergyEst *est = (MEnergyEst*)NextE();
        if (!est)
            break;

        const Double_t dist = fMatrix ? GetVal(col++) : ((MHillasSrc*)NextH())->GetDist();

        const Double_t ir =  i0 * (i1 + fA[1]*dist); // [cm]
        const Double_t er = -e0 * (e1 + e2*ir);      // [GeV]

        est->SetEnergy(er);
        est->SetImpact(ir);
    }

    return kTRUE;
}