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

  PFourier.h

  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.             *
 ***************************************************************************/

#ifndef _PFOURIER_H_
#define _PFOURIER_H_

#include <vector>

#include "fftw3.h"

#include <TH1F.h>

#include "Minuit2/FCNBase.h"

#include "PMusr.h"

//-------------------------------------------------------------
/**
 * <p>Apodization (windowing) strength constants.
 *
 * <p>Apodization applies a window function to time-domain data before
 * Fourier transformation to reduce spectral leakage (Gibbs phenomenon).
 * Stronger apodization improves frequency resolution but reduces amplitude
 * accuracy.
 */
/// No apodization (rectangular window)
#define F_APODIZATION_NONE   1
/// Weak apodization (gentle roll-off at edges)
#define F_APODIZATION_WEAK   2
/// Medium apodization (moderate roll-off)
#define F_APODIZATION_MEDIUM 3
/// Strong apodization (heavy roll-off for best frequency resolution)
#define F_APODIZATION_STRONG 4

//-------------------------------------------------------------
/**
 * <p>Phase correction optimizer for Fourier transforms.
 *
 * <p>This class performs automatic phase correction on complex Fourier spectra
 * to maximize the real component and minimize the imaginary component. Phase
 * errors arise from:
 * - Uncertain time-zero determination
 * - Detector time offsets
 * - Signal dispersion
 *
 * <p><b>Algorithm:</b> Minimizes a combined entropy-penalty functional using
 * Minuit2, finding optimal phase parameters (constant + linear dispersion):
 * φ(ω) = c₀ + c₁·ω
 *
 * <p><b>Applications:</b>
 * - Improving signal clarity in real Fourier spectra
 * - Identifying field distributions in vortex lattices
 * - Resolving closely-spaced frequency components
 *
 * <p><b>Usage:</b> Specify frequency range for optimization to focus on
 * signal peaks while avoiding noise regions.
 */
class PFTPhaseCorrection : public ROOT::Minuit2::FCNBase
{
  public:
    /**
     * <p>Constructor for phase correction with default Fourier data.
     *
     * @param minBin Minimum frequency bin for optimization (-1 = use all)
     * @param maxBin Maximum frequency bin for optimization (-1 = use all)
     */
    PFTPhaseCorrection(const Int_t minBin=-1, const Int_t maxBin=-1);

    /**
     * <p>Constructor with explicit Fourier data.
     *
     * @param reFT Real part of Fourier transform
     * @param imFT Imaginary part of Fourier transform
     * @param minBin Minimum frequency bin for optimization
     * @param maxBin Maximum frequency bin for optimization
     */
    PFTPhaseCorrection(std::vector<Double_t> &reFT, std::vector<Double_t> &imFT, const Int_t minBin=-1, const Int_t maxBin=-1);

    virtual ~PFTPhaseCorrection() {}

    /// Returns true if phase correction initialized successfully
    /// @return Validity status
    virtual Bool_t IsValid() { return fValid; }

    /**
     * <p>Performs phase correction minimization.
     *
     * <p>Uses Minuit2 to find optimal phase parameters that maximize
     * the real spectrum while minimizing imaginary components.
     */
    virtual void Minimize();

    /// Sets the gamma balancing parameter between entropy and penalty
    /// @param gamma Balancing factor (typical range: 0.1 to 10)
    virtual void SetGamma(const Double_t gamma) { fGamma = gamma; }

    /// Sets phase correction parameters manually
    /// @param c0 Constant phase offset in degrees
    /// @param c1 Linear phase dispersion coefficient
    virtual void SetPh(const Double_t c0, const Double_t c1) { fPh_c0 = c0; fPh_c1 = c1; CalcPhasedFT(); CalcRealPhFTDerivative(); }

    /// Returns the gamma parameter
    /// @return Balancing factor between entropy and penalty
    virtual Double_t GetGamma() { return fGamma; }

    /**
     * <p>Gets phase correction parameter.
     *
     * @param idx Parameter index (0=c₀, 1=c₁)
     * @return Phase parameter value
     */
    virtual Double_t GetPhaseCorrectionParam(UInt_t idx);

    /// Returns the minimum value of the optimization functional
    /// @return Minimum value achieved
    virtual Double_t GetMinimum();

  private:
    Bool_t fValid;

    std::vector<Double_t> fReal; /// original real Fourier data set
    std::vector<Double_t> fImag; /// original imag Fourier data set
    mutable std::vector<Double_t> fRealPh;  /// phased real Fourier data set
    mutable std::vector<Double_t> fImagPh;  /// phased imag Fourier data set
    mutable std::vector<Double_t> fRealPhD; /// 1st derivative of fRealPh

    Int_t fMinBin; /// minimum bin from Fourier range to be used for the phase correction estimate
    Int_t fMaxBin; /// maximum bin from Fourier range to be used for the phase correction estimate
    mutable Double_t fPh_c0; /// constant part of the phase dispersion used for the phase correction
    mutable Double_t fPh_c1; /// linear part of the phase dispersion used for the phase correction
    Double_t fGamma; /// gamma parameter to balance between entropy and penalty
    Double_t fMin; /// keeps the minimum of the entropy/penalty minimization

    virtual void Init();
    virtual void CalcPhasedFT() const;
    virtual void CalcRealPhFTDerivative() const;
    virtual Double_t Penalty() const;
    virtual Double_t Entropy() const;

    virtual Double_t Up() const { return 1.0; }
    virtual Double_t operator()(const std::vector<Double_t>&) const;
};

//-------------------------------------------------------------
/**
 * <p>Fourier transform engine for μSR time-domain data.
 *
 * <p>PFourier converts time-domain μSR signals to frequency domain,
 * revealing:
 * - Muon precession frequencies (field measurements)
 * - Internal field distributions (superconductors, magnets)
 * - Multiple muon stopping sites
 * - Dynamic frequency fluctuations
 *
 * <p><b>Key features:</b>
 * - Uses FFTW3 library for efficient FFT computation
 * - DC offset removal (for baseline correction)
 * - Zero-padding (improves frequency interpolation)
 * - Apodization/windowing (reduces spectral leakage)
 * - Multiple output formats (real, imaginary, power, phase)
 * - Unit conversion (field ↔ frequency)
 *
 * <p><b>Workflow:</b>
 * 1. Create PFourier with time histogram and settings
 * 2. Call Transform() with desired apodization
 * 3. Retrieve results: GetRealFourier(), GetPowerFourier(), etc.
 *
 * <p><b>Unit conversions:</b>
 * - Gauss: ω(MHz) = γ_μ/(2π) × B(G) = 0.01355 × B(G)
 * - Tesla: ω(MHz) = γ_μ/(2π) × B(T) = 135.54 × B(T)
 *
 * <p><b>Example:</b> TF-μSR measurement at 100 G produces a peak at
 * ~1.36 MHz in the Fourier spectrum.
 */
class PFourier
{
  public:
    /**
     * <p>Constructor for Fourier transformation.
     *
     * @param data Time histogram to transform
     * @param unitTag Output units (1=Gauss, 2=Tesla, 3=MHz, 4=Mc/s)
     * @param startTime Start time for transform in microseconds (0=from t0)
     * @param endTime End time for transform in microseconds (0=to end)
     * @param dcCorrected If true, remove DC offset before FFT
     * @param zeroPaddingPower Zero-pad to 2^N points (0=no padding)
     */
    PFourier(TH1F *data, Int_t unitTag,
             Double_t startTime = 0.0, Double_t endTime = 0.0,
             Bool_t dcCorrected = false, UInt_t zeroPaddingPower = 0);

    virtual ~PFourier();

    /**
     * <p>Performs the Fourier transformation.
     *
     * <p>Applies optional apodization, computes FFT using FFTW3,
     * and prepares output histograms in requested units.
     *
     * @param apodizationTag Apodization strength (0/1=none, 2=weak, 3=medium, 4=strong)
     */
    virtual void Transform(UInt_t apodizationTag = 0);

    /// Returns the original data histogram title
    /// @return Title string
    virtual const char* GetDataTitle() { return fData->GetTitle(); }

    /// Returns the output unit tag (1=G, 2=T, 3=MHz, 4=Mc/s)
    /// @return Unit identifier
    virtual const Int_t GetUnitTag() { return fUnitTag; }

    /// Returns the frequency resolution (bin width in output units)
    /// @return Frequency resolution
    virtual Double_t GetResolution() { return fResolution; }

    /**
     * <p>Returns the maximum frequency (Nyquist frequency).
     *
     * @return Maximum frequency in output units
     */
    virtual Double_t GetMaxFreq();

    /**
     * <p>Gets real part of Fourier transform as histogram.
     *
     * @param scale Scaling factor for amplitudes (default=1.0)
     * @return Pointer to TH1F histogram (caller must delete)
     */
    virtual TH1F* GetRealFourier(const Double_t scale = 1.0);

    /**
     * <p>Gets imaginary part of Fourier transform as histogram.
     *
     * @param scale Scaling factor for amplitudes (default=1.0)
     * @return Pointer to TH1F histogram (caller must delete)
     */
    virtual TH1F* GetImaginaryFourier(const Double_t scale = 1.0);

    /**
     * <p>Gets power spectrum |F(ω)|² as histogram.
     *
     * <p>Power spectrum is always positive and shows signal strength
     * at each frequency, useful for identifying dominant frequencies.
     *
     * @param scale Scaling factor for power (default=1.0)
     * @return Pointer to TH1F histogram (caller must delete)
     */
    virtual TH1F* GetPowerFourier(const Double_t scale = 1.0);

    /**
     * <p>Gets phase spectrum arg(F(ω)) as histogram.
     *
     * @param scale Scaling factor (default=1.0)
     * @return Pointer to TH1F histogram (caller must delete)
     */
    virtual TH1F* GetPhaseFourier(const Double_t scale = 1.0);

    /**
     * <p>Static method for phase-optimized real Fourier spectrum.
     *
     * <p>Applies phase correction to maximize real component using
     * provided phase parameters.
     *
     * @param re Real part of Fourier transform
     * @param im Imaginary part of Fourier transform
     * @param phase Phase correction parameters [c₀, c₁]
     * @param scale Scaling factor (default=1.0)
     * @param min Minimum frequency for correction (-1=all)
     * @param max Maximum frequency for correction (-1=all)
     * @return Pointer to phase-corrected TH1F histogram
     */
    static TH1F* GetPhaseOptRealFourier(const TH1F *re, const TH1F *im, std::vector<Double_t> &phase,
                                        const Double_t scale = 1.0, const Double_t min = -1.0, const Double_t max = -1.0);

    /// Returns true if Fourier transform is ready
    /// @return Validity status
    virtual Bool_t IsValid() { return fValid; }

  private:
    TH1F *fData; ///< data histogram to be Fourier transformed.

    Bool_t fValid; ///< true = all boundary conditions fullfilled and hence a Fourier transform can be performed.
    Int_t  fUnitTag; ///< 1=Field Units (G), 2=Field Units (T), 3=Frequency Units (MHz), 4=Angular Frequency Units (Mc/s)

    Int_t fApodization; ///< 0=none, 1=weak, 2=medium, 3=strong

    Double_t fTimeResolution; ///< time resolution of the data histogram in (us)
    Double_t fStartTime; ///< start time of the data histogram
    Double_t fEndTime; ///< end time of the data histogram
    Bool_t fDCCorrected; ///< if true, removed DC offset from signal before Fourier transformation, otherwise not
    UInt_t fZeroPaddingPower; ///< power for zero padding, if set < 0 no zero padding will be done
    Double_t fResolution; ///< Fourier resolution (field, frequency, or angular frequency)

    UInt_t fNoOfData; ///< number of bins in the time interval between fStartTime and fStopTime
    UInt_t fNoOfBins; ///< number of bins to be Fourier transformed. Might be different to fNoOfData due to zero padding
    fftw_plan fFFTwPlan; ///< fftw plan (see FFTW3 User Manual)
    fftw_complex *fIn; ///< real part of the Fourier transform
    fftw_complex *fOut; ///< imaginary part of the Fourier transform

//as    PFTPhaseCorrection *fPhCorrectedReFT;

    virtual void PrepareFFTwInputData(UInt_t apodizationTag);
    virtual void ApodizeData(Int_t apodizationTag);
};

#endif // _PFOURIER_H_