/* ======================================================================== *\
!
! *
! * 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): Oscar Blanch 12/2001 <mailto:blanch@ifae.es>
!   Author(s): Thomas Bretz 08/2002 <mailto:tbretz@astro.uni-wuerzburg.de>
!
!   Copyright: MAGIC Software Development, 2000-2004
!
!
\* ======================================================================== */

/////////////////////////////////////////////////////////////////////////////
//
//  MBadPixelsTreat
//
//  You can use MBadPixelsTreat::SetUseInterpolation to replaced the
//  bad pixels by the average of its neighbors instead of unmapping
//  them. If you want to include the central pixel use
//  MBadPixelsTreat::SetUseCentralPixel. The bad pixels are taken from
//  an existing MBadPixelsCam.
//
//  It check if there are enough neighbors to calculate the mean
//  If not, unmap the pixel. The minimum number of good neighbors
//  should be set using SetNumMinNeighbors
//
//  If you want to interpolate unreliable pixels and  unsuitable
//  (broken) pixels use SetHardTreatment().
//
//
//  Input Containers:
//   MCerPhotEvt
//   MPedPhotCam
//   MBadPixelsCam
//   [MGeomCam]
//
//  Output Containers:
//   MCerPhotEvt
//
/////////////////////////////////////////////////////////////////////////////
#include "MBadPixelsTreat.h"

#include <fstream>

#include <TEnv.h>
#include <TArrayD.h>

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

#include "MParList.h"

#include "MGeomPix.h"
#include "MGeomCam.h"

#include "MCerPhotPix.h"
#include "MCerPhotEvt.h"

#include "MPedPhotPix.h"
#include "MPedPhotCam.h"

#include "MBadPixelsPix.h"
#include "MBadPixelsCam.h"

ClassImp(MBadPixelsTreat);

using namespace std;

static const TString gsDefName  = "MBadPixelsTreat";
static const TString gsDefTitle = "Task to treat bad pixels (interpolation, unmapping)";

// --------------------------------------------------------------------------
//
// Default constructor.
//
MBadPixelsTreat::MBadPixelsTreat(const char *name, const char *title)
  : fFlags(0), fNumMinNeighbors(3), fNamePedPhotContainer("MPedPhotCam")
{
    fName  = name  ? name  : gsDefName.Data();
    fTitle = title ? title : gsDefTitle.Data();
}

// --------------------------------------------------------------------------
//
// Returns the status of a pixel. If kHardTreatment is set a pixel must
// be unsuitable or uriliable to be treated. If not it is treated only if 
// it is unsuitable
// (IsBad() checks for any flag)
//
Bool_t MBadPixelsTreat::IsPixelBad(Int_t idx) const
{
    return TESTBIT(fFlags, kHardTreatment) ? (*fBadPixels)[idx].IsBad():(*fBadPixels)[idx].IsUnsuitable() ;
}

// --------------------------------------------------------------------------
//
//  - Try to find or create MBlindPixels in parameter list.
//  - get the MCerPhotEvt from the parlist (abort if missing)
//  - if no pixels are given by the user try to determin the starfield
//    from the monte carlo run header.
//
Int_t MBadPixelsTreat::PreProcess (MParList *pList)
{
    fBadPixels = (MBadPixelsCam*)pList->FindObject(AddSerialNumber("MBadPixelsCam"));
    if (!fBadPixels)
    {
        *fLog << err << "MBadPixelsCam not found... aborting." << endl;
        return kFALSE;
    }

    fPedPhot = (MPedPhotCam*)pList->FindObject(AddSerialNumber(fNamePedPhotContainer), "MPedPhotCam");
    if (!fPedPhot)
    {
        *fLog << err << "MPedPhotCam not found... aborting." << endl;
        return kFALSE;
    }
    
    fEvt = (MCerPhotEvt*)pList->FindObject(AddSerialNumber("MCerPhotEvt"));
    if (!fEvt)
    {
        *fLog << err << "MCerPhotEvt not found... aborting." << endl;
        return kFALSE;
    }

    fGeomCam = (MGeomCam*)pList->FindObject(AddSerialNumber("MGeomCam"));
    if (!fGeomCam && TESTBIT(fFlags, kUseInterpolation))
    {
        *fLog << err << "MGeomCam not found... can't use interpolation." << endl;
        return kFALSE;
    }

    return kTRUE;
}

// --------------------------------------------------------------------------
//
//  Replaces each pixel (signal, signal error, pedestal, pedestal rms)
//  by the average of its surrounding pixels.
//  If TESTBIT(fFlags, kUseCentralPixel) is set the central pixel is also
//  included.
//
void MBadPixelsTreat::InterpolateSignal() const
{
    const UShort_t entries = fGeomCam->GetNumPixels();

    //
    // Create arrays (FIXME: Check if its possible to create it only once)
    //
    TArrayD nphot(entries);
    TArrayD perr(entries);
 
    //
    // Loop over all pixels
    //
    for (UShort_t i=0; i<entries; i++)
    {
        MCerPhotPix *pix = fEvt->GetPixById(i);

        //
        // Check whether pixel with idx i is blind
        //
        if (pix && !IsPixelBad(i))
            continue;

        //
        // Get a pointer to this pixel. If it is not yet existing
        // create a new entry for this pixel in MCerPhotEvt
        //
        if (!pix)
        {
            pix = fEvt->AddPixel(i, 0, 0);
            (*fBadPixels)[i].SetUnsuitable(MBadPixelsPix::kUnsuitableEvt);
        }

        //
        // Get the corresponding geometry and pedestal
        //
        const MGeomPix &gpix = (*fGeomCam)[i];

        // Do Not-Use-Central-Pixel
        const Bool_t nucp = !TESTBIT(fFlags, kUseCentralPixel);

        Int_t num = nucp ? 0 : 1;

        nphot[i]  = nucp ? 0 : pix->GetNumPhotons();
        perr[i]   = nucp ? 0 : Pow2(pix->GetErrorPhot());

        //
	// The values are rescaled to the small pixels area for the right comparison
        //
        const Double_t ratio = fGeomCam->GetPixRatio(i);

        nphot[i] *= ratio;
        perr[i]  *= ratio;

        //
        // Loop over all its neighbors
        //
        const Int_t n = gpix.GetNumNeighbors();
        for (int j=0; j<n; j++)
        {
            const UShort_t nidx = gpix.GetNeighbor(j);

            //
            // Do not use blind neighbors
            //
            if (IsPixelBad(nidx))
                continue;

            //
            // Check whether the neighbor has a signal stored
            //
            const MCerPhotPix *evtpix = fEvt->GetPixById(nidx);
            if (!evtpix)
                continue;

            //
            // Get the geometry for the neighbor
            //
            const Double_t nratio = fGeomCam->GetPixRatio(nidx);

            //
	    //The error is calculated as the quadratic sum of the errors
	    //
            nphot[i] += nratio*evtpix->GetNumPhotons();
            perr[i]  += nratio* Pow2(evtpix->GetErrorPhot());

            num++;
        }

	// Check if there are enough neighbors to calculate the mean
        // If not, unmap the pixel. The maximum number of blind neighbors
        // should be 2
        if (num<fNumMinNeighbors)
        {
            pix->SetPixelUnmapped();
            continue;
        }

        //
        // Now the mean is calculated and the values rescaled back
        // to the pixel area
        //
	nphot[i] /= (num*ratio);
        perr[i]   = TMath::Sqrt(perr[i]/(num*ratio));
 
        pix->Set(nphot[i], perr[i]);
    }
}

// --------------------------------------------------------------------------
//
void MBadPixelsTreat::InterpolatePedestals() const
{
    const Int_t entries = fPedPhot->GetSize();

    // Create arrays (FIXME: Check if its possible to create it only once)
    TArrayD ped(entries);
    TArrayD rms(entries);

    //
    // Loop over all pixels
    //
    for (UShort_t i=0; i<entries; i++)
    {
        //
        // Check whether pixel with idx i is blind
        //
        if (!IsPixelBad(i))
            continue;

        //
        // Get the corresponding geometry and pedestal
        //
        const MGeomPix    &gpix = (*fGeomCam)[i];
        const MPedPhotPix &ppix = (*fPedPhot)[i];

        // Do Not-Use-Central-Pixel
        const Bool_t nucp = !TESTBIT(fFlags, kUseCentralPixel);

        Int_t num = nucp ? 0 : 1;

        ped[i] = nucp ? 0 : ppix.GetMean();
        rms[i] = nucp ? 0 : Pow2(ppix.GetRms());

        //
        // The values are rescaled to the small pixels area for the right comparison
        //
        const Double_t ratio = fGeomCam->GetPixRatio(i);

        ped[i] *= ratio;
        rms[i] *= ratio;

        //
        // Loop over all its neighbors
        //
        const Int_t n = gpix.GetNumNeighbors();
        for (int j=0; j<n; j++)
        {
            const UShort_t nidx = gpix.GetNeighbor(j);

            //
            // Do not use blind neighbors
            //
            if (IsPixelBad(nidx))
                continue;

            //
            // Get the geometry for the neighbor
            //
            const Double_t    nratio = fGeomCam->GetPixRatio(nidx);
            const MPedPhotPix &nppix = (*fPedPhot)[nidx];

            //
            //The error is calculated as the quadratic sum of the errors
            //
            ped[i] += nratio*nppix.GetMean();
            rms[i] += nratio*Pow2(nppix.GetRms());

            num++;
        }

        // Check if there are enough neighbors to calculate the mean
        // If not, unmap the pixel. The minimum number of good neighbors
        // should be fNumMinNeighbors
        if (num < fNumMinNeighbors)
        {
            MCerPhotPix *pix =fEvt->GetPixById(i);
            if (!pix)
                pix = fEvt->AddPixel(i, 0, 0);
            pix->SetPixelUnmapped();
            continue;
        }

        //
        // Now the mean is calculated and the values rescaled back
        // to the pixel area
        //
        ped[i] /=  (num*ratio);
        rms[i]  = TMath::Sqrt(rms[i]/(num*ratio));

        (*fPedPhot)[i].Set(ped[i], rms[i]);
    }
}

// --------------------------------------------------------------------------
//
//  Replaces each pixel (signal, signal error, pedestal, pedestal rms)
//  by the average of its surrounding pixels.
//  If TESTBIT(fFlags, kUseCentralPixel) is set the central pixel is also
//  included.
//
//  NT: Informations about the interpolated pedestals are added. 
//      When the option Interpolated is called, the blind pixel with the new
//      values of signal and fluttuation is included in the calculation of
//      the Image Parameters.
//
/*
void MBadPixelsTreat::Interpolate() const
{
    const UShort_t entries = fGeomCam->GetNumPixels();

    //
    // Create arrays
    //
    TArrayD nphot(entries);
    TArrayD perr(entries);
    TArrayD ped(entries);
    TArrayD pedrms(entries);
 
    //
    // Loop over all pixels
    //
    for (UShort_t i=0; i<entries; i++)
    {
        MCerPhotPix *pix = fEvt->GetPixById(i);

        //
        // Check whether pixel with idx i is blind
        //
        if (pix && (*fBadPixels)[i].IsOK())
            continue;

        //
        // Get a pointer to this pixel. If it is not yet existing
        // create a new entry for this pixel in MCerPhotEvt
        //
        if (!pix)
        {
            pix = fEvt->AddPixel(i, 0, 0);
            (*fBadPixels)[i].SetUnsuitable(MBadPixelsPix::kUnsuitableEvt);
        }

        //
        // Get the corresponding geometry and pedestal
        //
        const MGeomPix &gpix    = (*fGeomCam)[i];
        const MPedPhotPix &ppix = (*fPedPhot)[i];

        // Do Not-Use-Central-Pixel
        const Bool_t nucp = !TESTBIT(fFlags, kUseCentralPixel);

        Int_t num = nucp ? 0 : 1;

        nphot[i]  = nucp ? 0 : pix->GetNumPhotons();
        ped[i]    = nucp ? 0 : ppix.GetMean();
        perr[i]   = nucp ? 0 : Pow2(pix->GetErrorPhot());
        pedrms[i] = nucp ? 0 : Pow2(ppix.GetRms());

        //
	// The values are rescaled to the small pixels area for the right comparison
        //
        const Double_t ratio = fGeomCam->GetPixRatio(i);

        if (nucp)
        {
            nphot[i]  *= ratio;
            perr[i]   *= ratio;
            ped[i]    *= ratio;
            pedrms[i] *= ratio;
        }

        //
        // Loop over all its neighbors
        //
        const Int_t n = gpix.GetNumNeighbors();
        for (int j=0; j<n; j++)
        {
            const UShort_t nidx = gpix.GetNeighbor(j);

            //
            // Do not use blind neighbors
            //
            if ((*fBadPixels)[nidx].IsBad())
                continue;

            //
            // Check whether the neighbor has a signal stored
            //
            const MCerPhotPix *evtpix = fEvt->GetPixById(nidx);
            if (!evtpix)
                continue;

            //
            // Get the geometry for the neighbor
            //
            const Double_t nratio = fGeomCam->GetPixRatio(nidx);
            MPedPhotPix &nppix    = (*fPedPhot)[nidx];

            //
	    // The error is calculated as the quadratic sum of the errors
	    //
            nphot[i]  += nratio*evtpix->GetNumPhotons();
            ped[i]    += nratio*nppix.GetMean();
            perr[i]   += nratio*Pow2(evtpix->GetErrorPhot());
	    pedrms[i] += nratio*Pow2(nppix.GetRms());

            num++;
        }

        if (num<fNumMinNeighbors)
        {
            pix->SetPixelUnmapped();
            nphot[i]  = 0;
            ped[i]    = 0;
            perr[i]   = 0;
            pedrms[i] = 0;
            continue;
        }

        //
	// Now the mean is calculated and the values rescaled back to the pixel area
        //
	nphot[i] /= num*ratio;
        ped[i]   /= num*ratio;
        perr[i]   = TMath::Sqrt(perr[i]/(num*ratio));
        pedrms[i] = TMath::Sqrt(pedrms[i]/(num*ratio));

    }

    //
    // Now the new pixel values are calculated and can be replaced in
    // the corresponding containers
    //
    for (UShort_t i=0; i<entries; i++)
    {
        //
        // Do not use blind neighbors
        //
        if ((*fBadPixels)[i].IsOK())
            continue;

        //
        // It must exist, we have created it in the loop before.
        //
        fEvt->GetPixById(i)->Set(nphot[i], perr[i]);
        (*fPedPhot)[i].Set(ped[i], pedrms[i]);
    }
}
*/

// --------------------------------------------------------------------------
//
//  Removes all blind pixels from the analysis by setting their state
//  to unused.
//
void MBadPixelsTreat::Unmap() const
{
    const UShort_t entries = fEvt->GetNumPixels();

    //
    // remove the pixels in fPixelsIdx if they are set to be used,
    // (set them to 'unused' state)
    //
    for (UShort_t i=0; i<entries; i++)
    {
        MCerPhotPix &pix = (*fEvt)[i];

        if (IsPixelBad(pix.GetPixId()))
            pix.SetPixelUnmapped();
    }
}

// --------------------------------------------------------------------------
//
// Interpolate Pedestals if kProcessRMS not set
//
Bool_t MBadPixelsTreat::ReInit(MParList *pList)
{
    if (!TESTBIT(fFlags, kProcessRMS))
        InterpolatePedestals();
    return kTRUE;
}

// --------------------------------------------------------------------------
//
// Treat the blind pixels
//
Int_t MBadPixelsTreat::Process()
{
    if (TESTBIT(fFlags, kUseInterpolation) && fGeomCam)
    {
        InterpolateSignal();
        if (TESTBIT(fFlags, kProcessRMS))
            InterpolatePedestals();
    }
    else
        Unmap();

    return kTRUE;
}

void MBadPixelsTreat::StreamPrimitive(ofstream &out) const
{
    out << "   MBadPixelsTreat " << GetUniqueName();
    if (fName!=gsDefName || fTitle!=gsDefTitle)
    {
        out << "(\"" << fName << "\"";
        if (fTitle!=gsDefTitle)
            out << ", \"" << fTitle << "\"";
        out <<")";
    }
    out << ";" << endl;

    if (TESTBIT(fFlags, kUseInterpolation))
        out << "   " << GetUniqueName() << ".SetUseInterpolation();" << endl;
    if (TESTBIT(fFlags, kUseCentralPixel))
        out << "   " << GetUniqueName() << ".SetUseCentralPixel();" << endl;
    if (TESTBIT(fFlags, kProcessRMS))
        out << "   " << GetUniqueName() << ".SetProcessRMS();" << endl;
    if (TESTBIT(fFlags, kHardTreatment))
        out << "   " << GetUniqueName() << ".SetHardTreatment();" << endl;
    if (fNumMinNeighbors!=3)
        out << "   " << GetUniqueName() << ".SetNumMinNeighbors(" << (int)fNumMinNeighbors << ");" << endl;
}

// --------------------------------------------------------------------------
//
// Read the setup from a TEnv, eg:
//   MBadPixelsTreat.UseInterpolation: no
//   MBadPixelsTreat.UseCentralPixel:  no
//   MBadPixelsTreat.HardTreatment:    no
//   MBadPixelsTreat.ProcessRMS:       no
//   MBadPixelsTreat.NumMinNeighbors:  3
//
Int_t MBadPixelsTreat::ReadEnv(const TEnv &env, TString prefix, Bool_t print)
{
    Bool_t rc = kFALSE;
    if (IsEnvDefined(env, prefix, "UseInterpolation", print))
    {
        rc = kTRUE;
        SetUseInterpolation(GetEnvValue(env, prefix, "UseInterpolation", IsUseInterpolation()));
    }
    if (IsEnvDefined(env, prefix, "UseCentralPixel", print))
    {
        rc = kTRUE;
        SetUseCentralPixel(GetEnvValue(env, prefix, "UseCentralPixel", IsUseCentralPixel()));
    }
    if (IsEnvDefined(env, prefix, "HardTreatment", print))
    {
        rc = kTRUE;
        SetHardTreatment(GetEnvValue(env, prefix, "HardTreatment", IsHardTreatment()));
    }
    if (IsEnvDefined(env, prefix, "ProcessRMS", print))
    {
        rc = kTRUE;
        SetProcessRMS(GetEnvValue(env, prefix, "ProcessRMS", IsProcessRMS()));
    }
    if (IsEnvDefined(env, prefix, "NumMinNeighbors", print))
    {
        rc = kTRUE;
        SetNumMinNeighbors(GetEnvValue(env, prefix, "NumMinNeighbors", fNumMinNeighbors));
    }
    return rc;
}