/***************************************************************************

  PPrepFourier.cpp

  Author: Andreas Suter
  e-mail: andreas.suter@psi.ch

***************************************************************************/

/***************************************************************************
 *   Copyright (C) 2007-2025 by Andreas Suter                              *
 *   andreas.suter@psi.ch                                                  *
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.,                                       *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
 ***************************************************************************/

#include <cmath>

#include "PPrepFourier.h"

//--------------------------------------------------------------------------
// Constructor
//--------------------------------------------------------------------------
/**
 * \brief Default constructor that initializes to default values.
 *
 * Sets background range to unused (-1, -1) and packing to 1 (no rebinning).
 * Background values and data must be set separately using setter methods.
 */
PPrepFourier::PPrepFourier()
{
  fBkgRange[0] = -1;
  fBkgRange[1] = -1;
  fPacking = 1;
}

//--------------------------------------------------------------------------
// Constructor
//--------------------------------------------------------------------------
/**
 * \brief Full constructor that initializes all configuration parameters.
 *
 * Creates a PPrepFourier instance with complete configuration. Background
 * can be specified either via range (for automatic calculation) or via
 * explicit values (one per data set to be added later).
 *
 * \param packing Rebinning factor (1=no rebinning, N=combine N bins)
 * \param bkgRange Background range [start, end] in bins (-1 if not used)
 * \param bkg Vector of explicit background values (one per data set)
 */
PPrepFourier::PPrepFourier(const Int_t packing, const Int_t *bkgRange, PDoubleVector bkg) :
  fPacking(packing)
{
  SetBkgRange(bkgRange);
  SetBkg(bkg);
}

//--------------------------------------------------------------------------
// Destructor
//--------------------------------------------------------------------------
/**
 * \brief Destructor that cleans up internal data structures.
 *
 * Clears both raw data and processed data vectors. Note that TH1F objects
 * returned by GetData() are owned by the caller and not cleaned up here.
 */
PPrepFourier::~PPrepFourier()
{
  fRawData.clear();
  fData.clear();
}

//--------------------------------------------------------------------------
// SetBkgRange
//--------------------------------------------------------------------------
/**
 * \brief Sets the background range for automatic background calculation.
 *
 * Specifies the bin range [start, end] for calculating background by averaging.
 * Values of -1 indicate that background range is not used (explicit background
 * values should be provided instead). Validates that both values are ≥-1 and
 * prints warnings if invalid values are encountered.
 *
 * \param bkgRange Array [start, end] with background bin range (both must be ≥-1)
 */
void PPrepFourier::SetBkgRange(const Int_t *bkgRange)
{
  Int_t err=0;
  if (bkgRange[0] >= -1) {
    fBkgRange[0] = bkgRange[0];
  } else {
    err = 1;
  }
  if (bkgRange[1] >= -1) {
    fBkgRange[1] = bkgRange[1];
  } else {
    if (err == 0)
      err = 2;
    else
      err = 3;
  }

  TString errMsg("");
  switch (err) {
    case 1:
    errMsg = TString::Format("start bkg range < 0 (given: %d), will ignore it.", bkgRange[0]);
      break;
    case 2:
      errMsg = TString::Format("end bkg range < 0 (given: %d), will ignore it.", bkgRange[1]);
      break;
    case 3:
      errMsg = TString::Format("start/end bkg range < 0 (given: %d/%d), will ignore it.", bkgRange[0], bkgRange[1]);
      break;
    default:
      errMsg = TString("??");
      break;
  }

  if (err != 0) {
    std::cerr << std::endl << ">> PPrepFourier::SetBkgRange: **WARNING** " << errMsg << std::endl;
  }
}

//--------------------------------------------------------------------------
// SetBkg
//--------------------------------------------------------------------------
/**
 * \brief Sets explicit background values for all data sets.
 *
 * Provides pre-calculated background values (one per data set) instead of
 * calculating them from a background range. This allows different background
 * values for each histogram.
 *
 * \param bkg Vector of background values (should match number of data sets)
 */
void PPrepFourier::SetBkg(PDoubleVector bkg)
{
  for (UInt_t i=0; i<bkg.size(); i++)
    fBkg.push_back(bkg[i]);
}

//--------------------------------------------------------------------------
// SetPacking
//--------------------------------------------------------------------------
/**
 * \brief Sets the rebinning/packing factor for data reduction.
 *
 * Specifies how many adjacent bins to combine during DoPacking(). A value
 * of 1 means no rebinning, 2 combines pairs of bins, etc. Invalid values
 * (≤0) are rejected with a warning.
 *
 * \param packing Number of bins to combine (must be > 0)
 */
void PPrepFourier::SetPacking(const Int_t packing)
{
  if (packing > 0) {
    fPacking = packing;
  } else {
    std::cerr << std::endl << ">> PPrepFourier::SetPacking: **WARNING** found packing=" << packing << " < 0, will ignore it." << std::endl;
  }
}

//--------------------------------------------------------------------------
// AddData
//--------------------------------------------------------------------------
/**
 * \brief Adds a time-domain data set to the internal collection.
 *
 * Stores the raw data and metadata for later processing. Multiple data sets
 * can be added and will be processed together by DoBkgCorrection(), DoPacking(),
 * and DoLifeTimeCorrection().
 *
 * \param data musrFT_data structure containing raw histogram data and metadata
 */
void PPrepFourier::AddData(musrFT_data &data)
{
  fRawData.push_back(data);
}

//--------------------------------------------------------------------------
// DoBkgCorrection
//--------------------------------------------------------------------------
/**
 * \brief Applies background correction to all data sets.
 *
 * Subtracts background from each data point. The background can be specified
 * in two ways:
 * 1. Background range: Calculates average over the specified bin range
 * 2. Explicit values: Uses the background values set via SetBkg()
 *
 * If fData is not yet initialized, calls InitData() first. If neither
 * background range nor explicit values are provided, no correction is applied.
 * Validates that background range is within data bounds and that the number
 * of explicit background values matches the number of data sets.
 */
void PPrepFourier::DoBkgCorrection()
{
  // make sure fData are already present, and if not create the necessary data sets
  if (fData.size() != fRawData.size()) {
    InitData();
  }

  // if no bkg-range is given, nothing needs to be done
  if ((fBkgRange[0] == -1) && (fBkgRange[1] == -1) && (fBkg.size() == 0)) {
    return;
  }

  if ((fBkgRange[0] != -1) && (fBkgRange[1] != -1)) { // background range is given
    // make sure that the bkg range is ok
    for (UInt_t i=0; i<fRawData.size(); i++) {
      if ((fBkgRange[0] >= static_cast<Int_t>(fRawData[i].rawData.size())) || (fBkgRange[1] >= static_cast<Int_t>(fRawData[i].rawData.size()))) {
        std::cerr << std::endl << "PPrepFourier::DoBkgCorrection() **ERROR** bkg-range out of data-range!";
        return;
      }
    }

    Double_t bkg=0.0;
    for (UInt_t i=0; i<fRawData.size(); i++) {
      // calculate the bkg for the given range
      for (Int_t j=fBkgRange[0]; j<=fBkgRange[1]; j++) {
        bkg += fRawData[i].rawData[j];
      }
      bkg /= (fBkgRange[1]-fBkgRange[0]+1);
      std::cout << "info> background " << i << ": " << bkg << std::endl;

      // correct data
      for (UInt_t j=0; j<fData[i].size(); j++)
        fData[i][j] -= bkg;
    }
  } else { // there might be an explicit background list
    // check if there is a background list
    if (fBkg.size() == 0)
      return;

    // check if there are as many background values than data values
    if (fBkg.size() != fData.size()) {
      std::cerr << std::endl << "PPrepFourier::DoBkgCorrection() **ERROR** #bkg values != #histos. Will do nothing here." << std::endl;
      return;
    }

    for (UInt_t i=0; i<fData.size(); i++)
      for (UInt_t j=0; j<fData[i].size(); j++)
        fData[i][j] -= fBkg[i];
  }
}

//--------------------------------------------------------------------------
// DoPacking
//--------------------------------------------------------------------------
/**
 * \brief Applies rebinning/packing to reduce the number of data points.
 *
 * Combines adjacent bins according to the packing factor to improve statistics
 * and reduce data size. Bins are summed in groups of fPacking size. If
 * fPacking=1, no rebinning is performed.
 *
 * If fData is not yet initialized, calls InitData() first. The time resolution
 * should be adjusted accordingly (multiplied by fPacking) when creating histograms.
 */
void PPrepFourier::DoPacking()
{
  // make sure fData are already present, and if not create the necessary data sets
  if (fData.size() != fRawData.size()) {
    InitData();
  }

  if (fPacking == 1) // nothing to be done
    return;

  PDoubleVector tmpData;
  Double_t dval = 0.0;
  for (UInt_t i=0; i<fData.size(); i++) {
    tmpData.clear();
    dval = 0.0;
    for (UInt_t j=0; j<fData[i].size(); j++) {
      if ((j % fPacking == 0) && (j != 0)) {
        tmpData.push_back(dval);
        dval = 0.0;
      } else {
        dval += fData[i][j];
      }
    }
    // change the original data set with the packed one
    fData[i].clear();
    fData[i] = tmpData;
  }
}

//--------------------------------------------------------------------------
// DoLifeTimeCorrection
//--------------------------------------------------------------------------
/**
 * \brief Applies muon lifetime correction for theory-free Fourier analysis.
 *
 * Performs a theory-free lifetime correction by:
 * 1. Multiplying data by exp(t/τ_μ) to remove muon decay
 * 2. Estimating N0 as the average of the corrected data
 * 3. Subtracting N0 and normalizing by N0 to get asymmetry-like data
 *
 * This approach works well for high fields (>few kGauss) where depolarization
 * is minimal, but may be less accurate for low fields where significant
 * relaxation occurs during the muon lifetime.
 *
 * If fData is not yet initialized, calls InitData() first. Background correction
 * should typically be applied before lifetime correction.
 *
 * \param fudge Rescaling factor for estimated N0 (typically ~1.0, allows fine-tuning)
 */
void PPrepFourier::DoLifeTimeCorrection(Double_t fudge)
{
  // make sure fData are already present, and if not create the necessary data sets
  if (fData.size() != fRawData.size()) {
    InitData();
  }

  // calc exp(+t/tau)*N(t), where N(t) is already background corrected
  Double_t scale;
  for (UInt_t i=0; i<fData.size(); i++) {
    scale = fRawData[i].timeResolution / PMUON_LIFETIME;
    for (UInt_t j=0; j<fData[i].size(); j++) {
      fData[i][j] *= exp(j*scale);
    }
  }

  // calc N0
  Double_t dval;
  Double_t N0;
  for (UInt_t i=0; i<fData.size(); i++) {
    dval = 0.0;
    for (UInt_t j=0; j<fData[i].size(); j++) {
      dval += fData[i][j];
    }
    N0 = dval/fData[i].size();
    N0 *= fudge;
    for (UInt_t j=0; j<fData[i].size(); j++) {
      fData[i][j] -= N0;
      fData[i][j] /= N0;
    }
  }
}

//--------------------------------------------------------------------------
// GetInfo
//--------------------------------------------------------------------------
/**
 * \brief Returns the metadata string for a specific data set.
 *
 * \param idx Data set index
 * \return Metadata string from musrFT_data.info, or empty string if idx out of range
 */
TString PPrepFourier::GetInfo(const UInt_t idx)
{
  TString info("");

  if (idx < fRawData.size())
    info = fRawData[idx].info;

  return info;
}

//--------------------------------------------------------------------------
// GetDataSetTag
//--------------------------------------------------------------------------
/**
 * \brief Returns the data set tag identifier.
 *
 * The data set tag is used to label and group related data sets, particularly
 * for average-per-data-set operations in Fourier analysis.
 *
 * \param idx Data set index
 * \return Data set tag value, or -1 if idx out of range
 */
Int_t PPrepFourier::GetDataSetTag(const UInt_t idx)
{
  Int_t result = -1;

  if (idx < fRawData.size())
    result = fRawData[idx].dataSetTag;

  return result;
}

//--------------------------------------------------------------------------
// GetData
//--------------------------------------------------------------------------
/**
 * \brief Creates ROOT histograms for all processed data sets.
 *
 * Converts the processed data into TH1F histograms with proper time binning.
 * The time range, bin width, and histogram titles are taken from the
 * musrFT_data metadata. The packing factor is applied to the time resolution.
 *
 * \return Vector of TH1F pointers (caller owns and must delete)
 *
 * \note The caller is responsible for deleting the returned histograms.
 * \note Returns empty vector if no data is present.
 */
std::vector<TH1F*> PPrepFourier::GetData()
{
  std::vector<TH1F*> data;
  data.resize(fData.size());

  // if not data are present, just return an empty vector
  if (fData.size() == 0)
    return data;

  TString name("");
  Double_t dt=0.0;
  Double_t start=0.0;
  Double_t end=0.0;
  UInt_t size;
  UInt_t startIdx;
  UInt_t endIdx;

  for (UInt_t i=0; i<fData.size(); i++) {
    name = TString::Format("histo%2d", i);
    dt = fRawData[i].timeResolution*fPacking;
    start = fRawData[i].timeRange[0];
    end = fRawData[i].timeRange[1];

    // init size and start/end indices
    size = fData[i].size();
    startIdx = 1;
    endIdx = size;

    // time range given, hence calculate the proper size
    if (start != -1.0) {
      size = static_cast<UInt_t>((end-start)/dt);
      if (start >= 0.0) {
        startIdx = static_cast<UInt_t>(start/dt)+1;
        endIdx = static_cast<UInt_t>(end/dt)+1;
      } else {
        std::cerr << std::endl << ">> PPrepFourier::GetData **WARNING** found start time < 0.0, will set it to 0.0" << std::endl;
        endIdx = static_cast<UInt_t>(end/dt)+1;
      }
    }

    if (start == -1.0) { // no time range given, hence start == 0.0 - dt/2
      start = -dt/2.0;
    } else { // time range given
      start -= dt/2.0;
    }

    if (end == -1.0) { // no time range given, hence end == (fData[idx].size()-1)*dt + dt/2
      end = (fData[i].size()-1)*dt+dt/2.0;
    } else { // time range given
      end += dt/2.0;
    }

    data[i] = new TH1F(name.Data(), fRawData[i].info.Data(), size, start, end);
    for (UInt_t j=startIdx; j<endIdx; j++)
      data[i]->SetBinContent(j-startIdx+1, fData[i][j]);
  }

  return data;
}

//--------------------------------------------------------------------------
// GetData
//--------------------------------------------------------------------------
/**
 * \brief Creates a ROOT histogram for a specific processed data set.
 *
 * Converts a single processed data set into a TH1F histogram with proper
 * time binning. The time range, bin width, and histogram title are taken
 * from the musrFT_data metadata. The packing factor is applied to the
 * time resolution.
 *
 * \param idx Data set index
 * \return TH1F pointer (caller owns and must delete), or nullptr if idx out of range
 *
 * \note The caller is responsible for deleting the returned histogram.
 * \note Returns nullptr if no data is present or idx is out of range.
 */
TH1F *PPrepFourier::GetData(const UInt_t idx)
{
  if (fData.size() == 0) // no data present
    return nullptr;

  if (idx > fData.size()) // requested index out of range
    return nullptr;

  TString name = TString::Format("histo%2d", idx);
  Double_t dt = fRawData[idx].timeResolution*fPacking;
  Double_t start = fRawData[idx].timeRange[0];
  Double_t end = fRawData[idx].timeRange[1];
  UInt_t size = fData[idx].size();
  UInt_t startIdx = 1;
  UInt_t endIdx = size;

  // time range given, hence calculate the proper size
  if (start != -1.0) {
    size = static_cast<UInt_t>((end-start)/dt);
    if (start >= 0.0) {
      startIdx = static_cast<UInt_t>(start/dt)+1;
      endIdx = static_cast<UInt_t>(end/dt)+1;
    } else {
      std::cerr << std::endl << ">> PPrepFourier::GetData **WARNING** found start time < 0.0, will set it to 0.0" << std::endl;
      endIdx = static_cast<UInt_t>(end/dt)+1;
    }
  }

  if (start == -1.0) { // no time range given, hence start == 0.0 - dt/2
    start = -dt/2.0;
  } else { // time range given
    start -= dt/2.0;
  }

  if (end == -1.0) { // no time range given, hence end == (fData[idx].size()-1)*dt + dt/2
    end = (fData[idx].size()-1)*dt+dt/2.0;
  } else { // time range given
    end += dt/2.0;
  }

  TH1F *data = new TH1F(name.Data(), fRawData[idx].info.Data(), size, start, end);
  for (UInt_t i=startIdx; i<endIdx; i++)
    data->SetBinContent(i-startIdx+1, fData[idx][i]);

  return data;
}

//--------------------------------------------------------------------------
// InitData (private)
//--------------------------------------------------------------------------
/**
 * \brief Initializes processed data by copying from raw data starting at t0.
 *
 * Creates the fData vectors from fRawData, starting from the t0 bin for each
 * data set. This effectively removes pre-t0 bins and creates a working copy
 * for subsequent processing (background correction, packing, lifetime correction).
 *
 * If t0 is negative or not set, data is copied from bin 0. This method is
 * called automatically by processing methods if fData is not yet initialized.
 */
void PPrepFourier::InitData()
{
  fData.resize(fRawData.size());
  UInt_t t0;
  for (UInt_t i=0; i<fRawData.size(); i++) {
    if (fRawData[i].t0 >= 0)
      t0 = fRawData[i].t0;
    else
      t0 = 0;
    for (UInt_t j=t0; j<fRawData[i].rawData.size(); j++) {
      fData[i].push_back(fRawData[i].rawData[j]);
    }
  }
}