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

  PFitter.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 _PFITTER_H_
#define _PFITTER_H_

#include <memory>

#include "Minuit2/MnUserParameters.h"
#include "Minuit2/FunctionMinimum.h"

#include "PMusr.h"
#include "PMsrHandler.h"
#include "PRunListCollection.h"
#include "PFitterFcn.h"

//-------------------------------------------------------------
/**
 * <p>Minuit2 command identifiers for COMMANDS block.
 *
 * <p>These constants map MSR file COMMANDS block keywords to internal
 * command identifiers. Commands control the fitting process, including
 * minimization algorithms, parameter fixing/releasing, error analysis,
 * and output options.
 *
 * <p><b>Categories:</b>
 * - <b>Minimizers:</b> MIGRAD (gradient descent), SIMPLEX (non-gradient), MINIMIZE (automatic)
 * - <b>Error analysis:</b> HESSE (covariance matrix), MINOS (asymmetric errors)
 * - <b>Parameter control:</b> FIX, RELEASE, RESTORE, SAVE
 * - <b>Exploration:</b> SCAN (1D/2D parameter space), CONTOURS (error ellipses)
 * - <b>Settings:</b> STRATEGY, MACHINE_PRECISION, PRINT
 * - <b>Analysis:</b> FIT_RANGE (time-window scan), SECTOR (chisq vs. time)
 */

/// Interactive mode (not implemented in musrfit)
#define PMN_INTERACTIVE       0
/// Contour plot command (2D error ellipses)
#define PMN_CONTOURS          1
/// Eigenvalue analysis (not implemented)
#define PMN_EIGEN             2
/// Fit range scan to find optimal fitting window
#define PMN_FIT_RANGE         3
/// Fix parameters at current values
#define PMN_FIX               4
/// Calculate Hessian matrix for error estimation
#define PMN_HESSE             5
/// Set machine precision for numerical derivatives
#define PMN_MACHINE_PRECISION 6
/// MIGRAD minimizer (gradient-based, recommended)
#define PMN_MIGRAD            7
/// MINIMIZE (automatic algorithm selection)
#define PMN_MINIMIZE          8
/// MINOS error analysis (asymmetric confidence intervals)
#define PMN_MINOS             9
/// Plot command (for scan/contour results)
#define PMN_PLOT              10
/// Release previously fixed parameters
#define PMN_RELEASE           11
/// Restore parameters from previous SAVE
#define PMN_RESTORE           12
/// Save current parameter state
#define PMN_SAVE              13
/// Parameter scan (1D or 2D)
#define PMN_SCAN              14
/// SIMPLEX minimizer (robust but slow)
#define PMN_SIMPLEX           15
/// Set minimization strategy (0=fast, 1=default, 2=careful)
#define PMN_STRATEGY          16
/// Set user covariance matrix (not implemented)
#define PMN_USER_COVARIANCE   17
/// Set user parameter state (not implemented)
#define PMN_USER_PARAM_STATE  18
/// Set print level for fit output (0=quiet, 1=normal, 2=verbose)
#define PMN_PRINT             19
/// Calculate χ² vs. time in sectors (time-window analysis)
#define PMN_SECTOR            20

//-----------------------------------------------------------------------------
/**
 * <p>Container for sector χ² (or max likelihood) analysis results.
 *
 * <p>The SECTOR command analyzes how χ² changes as the fit range is extended
 * from the first good bin (fgb) to progressively later times. This helps
 * identify:
 * - Optimal fitting windows
 * - Time regions with systematic deviations
 * - Data quality issues
 * - Relaxation component contributions
 *
 * <p>A "sector" is a time window [fgb, fLast] where fLast < lgb. The class
 * stores χ²/maxLH and NDF for each sector, both for the total fit and for
 * individual runs.
 *
 * <p><b>Use case:</b> If early-time data shows poor χ², the sector analysis
 * reveals whether to adjust t0, fgb, or background estimation.
 */
class PSectorChisq
{
  public:
    /**
     * <p>Constructor.
     *
     * @param noOfRuns Number of runs in the fit (allocates storage for per-run statistics)
     */
    PSectorChisq(UInt_t noOfRuns);

    /// Sets the first good bin time for a specific run
    /// @param first First good bin time in microseconds
    /// @param idx Run index
    void SetRunFirstTime(Double_t first, UInt_t idx);

    /// Sets the sector end time (last bin included in this sector)
    /// @param last Sector end time in microseconds
    void SetSectorTime(Double_t last) { fLast = last; }

    /// Sets the total χ² (or max likelihood) for this sector
    /// @param chisq Chi-squared or maximum likelihood value
    void SetChisq(Double_t chisq) { fChisq = chisq; }

    /// Sets the χ² for a specific run in this sector
    /// @param chisq Chi-squared value for run
    /// @param idx Run index
    void SetChisq(Double_t chisq, UInt_t idx);

    /// Sets the expected total χ² for this sector
    /// @param expChisq Expected chi-squared value
    void SetExpectedChisq(Double_t expChisq) { fExpectedChisq = expChisq; }

    /// Sets the expected χ² for a specific run
    /// @param chisq Expected chi-squared value
    /// @param idx Run index
    void SetExpectedChisq(Double_t chisq, UInt_t idx);

    /// Sets the total number of degrees of freedom for this sector
    /// @param ndf Degrees of freedom
    void SetNDF(UInt_t ndf) { fNDF = ndf; }

    /// Sets the NDF for a specific run
    /// @param ndf Degrees of freedom for run
    /// @param idx Run index
    void SetNDF(UInt_t ndf, UInt_t idx);

    /// Gets the first good bin time for a specific run
    /// @param idx Run index
    /// @return First good bin time in microseconds
    Double_t GetTimeRangeFirst(UInt_t idx);

    /// Returns the sector end time
    /// @return Last bin time in microseconds
    Double_t GetTimeRangeLast() { return fLast; }

    /// Returns the total χ² for this sector
    /// @return Chi-squared or max likelihood value
    Double_t GetChisq() { return fChisq; }

    /// Returns the χ² for a specific run
    /// @param idx Run index
    /// @return Chi-squared value
    Double_t GetChisq(UInt_t idx);

    /// Returns the expected total χ²
    /// @return Expected chi-squared value
    Double_t GetExpectedChisq() { return fExpectedChisq; }

    /// Returns the expected χ² for a specific run
    /// @param idx Run index
    /// @return Expected chi-squared value
    Double_t GetExpectedChisq(UInt_t idx);

    /// Returns the total NDF for this sector
    /// @return Degrees of freedom
    UInt_t   GetNDF() { return fNDF; }

    /// Returns the NDF for a specific run
    /// @param idx Run index
    /// @return Degrees of freedom
    UInt_t   GetNDF(UInt_t idx);

    /// Returns the number of runs
    /// @return Number of runs in the analysis
    UInt_t   GetNoRuns() { return fNoOfRuns; }

  private:
    UInt_t   fNoOfRuns; ///< number of runs presesent
    Double_t fLast;  ///< requested time stamp
    Double_t fChisq; ///< chisq or maxLH for the sector
    Double_t fExpectedChisq; ///< keep the expected chisq or maxLH for the sector
    UInt_t   fNDF;   ///< NDF for the sector
    PDoubleVector fFirst; ///< time stamp for fgb for a given run
    PDoubleVector fChisqRun; ///< chisq or maxLH for the sector and run
    PDoubleVector fExpectedChisqRun; ///< expected chisq or maxLH for the sector and run
    PUIntVector fNDFRun; ///< NDF for the sector and run
};

//-----------------------------------------------------------------------------
/**
 * <p>Main fitting engine interfacing with ROOT Minuit2.
 *
 * <p>PFitter orchestrates the entire fitting process for musrfit:
 * - Initializes Minuit2 with parameters from MSR file
 * - Executes COMMANDS block directives (MIGRAD, HESSE, MINOS, etc.)
 * - Manages parameter fixing, releasing, and boundaries
 * - Performs error analysis (Hessian, MINOS)
 * - Conducts parameter scans and contour plots
 * - Calculates sector χ² for time-window analysis
 * - Generates fit output and statistics
 *
 * <p><b>Fitting workflow:</b>
 * 1. Initialize parameters and set boundaries
 * 2. Execute COMMANDS block sequentially:
 *    - MIGRAD: Find minimum using gradient descent
 *    - HESSE: Calculate symmetric errors from covariance matrix
 *    - MINOS: Calculate asymmetric errors (optional, slower)
 *    - SAVE: Save parameter state
 * 3. Update MSR file with fitted parameters and statistics
 *
 * <p><b>Minimization modes:</b>
 * - <b>χ² minimization:</b> Standard least-squares fitting
 * - <b>Maximum likelihood:</b> Poisson statistics (better for low counts)
 *
 * <p><b>Example COMMANDS block:</b>
 * @code
 * COMMANDS
 * SET PRINT 1
 * MIGRAD
 * MINOS
 * SAVE
 * @endcode
 */
class PFitter
{
  public:
    /**
     * <p>Constructor for fitting engine.
     *
     * @param runInfo Pointer to MSR handler containing fit configuration
     * @param runListCollection Pointer to collection of data runs to fit
     * @param chisq_only If true, only calculate χ² without fitting
     * @param yaml_out If true, generate YAML output file with fit results
     */
    PFitter(PMsrHandler *runInfo, PRunListCollection *runListCollection, Bool_t chisq_only = false, Bool_t yaml_out = false);

    virtual ~PFitter();

    /// Returns true if fitter initialized successfully
    /// @return Validity status
    Bool_t IsValid() { return fIsValid; }

    /// Returns true if only parameter scan requested (no minimization)
    /// @return Scan-only flag
    Bool_t IsScanOnly() { return fIsScanOnly; }

    /// Returns true if fit converged successfully
    /// @return Convergence status
    Bool_t HasConverged() { return fConverged; }

    /**
     * <p>Executes the complete fitting procedure.
     *
     * <p>Processes all commands from the COMMANDS block sequentially,
     * performs the fit, calculates errors, and prepares output statistics.
     *
     * @return true if fit completed successfully, false on error
     */
    Bool_t DoFit();

  private:
    // State flags
    Bool_t fIsValid;     ///< Overall validity flag: true if fitter initialized successfully
    Bool_t fIsScanOnly;  ///< Scan mode flag: true if only parameter scans requested (no minimization)
    Bool_t fConverged;   ///< Convergence flag: true if fit converged to a valid minimum
    Bool_t fChisqOnly;   ///< Evaluation-only flag: true to calculate χ² without fitting
    Bool_t fYamlOut;     ///< Output flag: true to generate YAML output file (MINUIT2.OUTPUT → yaml)
    Bool_t fUseChi2;     ///< Fit mode: true = χ² minimization, false = log-max-likelihood
    UInt_t fPrintLevel;  ///< Verbosity level: 0=quiet, 1=normal, 2=verbose (Minuit output)

    UInt_t fStrategy; ///< Minuit2 strategy: 0=fast/low-accuracy, 1=default, 2=careful/high-accuracy

    // Core data structures
    PMsrHandler *fRunInfo; ///< Pointer to MSR file handler (parameters, theory, commands)
    PRunListCollection *fRunListCollection; ///< Pointer to preprocessed run data collection

    PMsrParamList fParams; ///< Copy of parameter list from MSR file

    PMsrLines      fCmdLines; ///< Raw command lines from MSR COMMANDS block
    PIntPairVector fCmdList;  ///< Parsed commands: first=command ID, second=line number

    std::unique_ptr<PFitterFcn> fFitterFcn; ///< Objective function for Minuit2 minimization

    ROOT::Minuit2::MnUserParameters fMnUserParams; ///< Minuit2 parameter state (values, errors, limits)
    std::unique_ptr<ROOT::Minuit2::FunctionMinimum> fFcnMin; ///< Minuit2 function minimum result

    // Scan and contour analysis
    Bool_t fScanAll;           ///< Multi-parameter scan flag: false=1D scan, true=2D scan (not fully implemented)
    UInt_t fScanParameter[2];  ///< Parameter indices: [0]=primary scan/contour, [1]=secondary (contours only)
    UInt_t fScanNoPoints;      ///< Number of scan/contour evaluation points (default=41)
    Double_t fScanLow;         ///< Scan lower bound: 0.0 = auto (2σ below current value)
    Double_t fScanHigh;        ///< Scan upper bound: 0.0 = auto (2σ above current value)
    PDoublePairVector fScanData; ///< Scan results: (parameter_value, χ²) pairs

    PDoublePairVector fOriginalFitRange; ///< Original fit ranges per run (saved for FIT_RANGE command)

    PStringVector fElapsedTime; ///< Timing information for each fit command

    // Sector χ² analysis
    Bool_t fSectorFlag; ///< SECTOR command present flag
    std::vector<PSectorChisq> fSector; ///< Sector analysis results (χ² vs. time windows)

    std::vector<bool> fPhase; ///< Phase parameter flags: true if parameter is a phase angle

    //----------------------------------------------------------------------
    // Phase parameter identification (private helpers)
    //----------------------------------------------------------------------

    /**
     * @brief Identifies which parameters represent phase angles.
     *
     * Scans the THEORY block to detect parameters used as phases in
     * standard functions (TFieldCos, bessel, etc.). Phase parameters
     * are constrained to [-360°, +360°] during fitting.
     */
    void GetPhaseParams();

    /**
     * @brief Extracts parameter numbers from a FUNCTIONS block entry.
     *
     * Parses "funX" references in theory lines to find all parameters
     * used in the function definition.
     *
     * @param funStr Function identifier string (e.g., "fun1", "fun23")
     * @return Vector of parameter numbers (1-indexed) used in the function
     */
    PIntVector GetParFromFun(const TString funStr);

    /**
     * @brief Extracts parameter numbers from a map reference.
     *
     * Parses "mapX" references to find mapped parameters across all runs.
     * Maps allow different runs to use different parameters for the same
     * theoretical component.
     *
     * @param mapStr Map identifier string (e.g., "map1", "map5")
     * @return Vector of parameter numbers (1-indexed) referenced by the map
     */
    PIntVector GetParFromMap(const TString mapStr);

    //----------------------------------------------------------------------
    // Command validation and execution (private methods)
    //----------------------------------------------------------------------

    /**
     * @brief Validates COMMANDS block syntax and builds execution queue.
     *
     * Parses all command lines, checks for syntax errors, extracts parameters,
     * and populates fCmdList for sequential execution.
     *
     * @return true if all commands are valid, false on syntax errors
     */
    Bool_t CheckCommands();

    /**
     * @brief Transfers MSR parameters to Minuit2 parameter state.
     *
     * Initializes fMnUserParams with values, errors, and bounds from the
     * MSR file's PARAMETERS block.
     *
     * @return true if parameters set successfully
     */
    Bool_t SetParameters();

    /**
     * @brief Executes CONTOURS command (2D error contours).
     *
     * Calculates confidence regions in 2D parameter space by evaluating
     * χ² on a grid around the minimum.
     *
     * @return true if contour calculation succeeded
     */
    Bool_t ExecuteContours();

    /**
     * @brief Executes FIT_RANGE command (optimal time-window search).
     *
     * Scans fit quality vs. fit start time to find the optimal first-good-bin.
     * Useful for determining when background subtraction is adequate.
     *
     * @param lineNo Command line number in MSR file
     * @return true if range scan succeeded
     */
    Bool_t ExecuteFitRange(UInt_t lineNo);

    /**
     * @brief Executes FIX command (freeze parameters).
     *
     * Prevents specified parameters from varying during subsequent minimization.
     *
     * @param lineNo Command line number in MSR file
     * @return true if parameters fixed successfully
     */
    Bool_t ExecuteFix(UInt_t lineNo);

    /**
     * @brief Executes HESSE command (calculate error matrix).
     *
     * Computes the covariance matrix by evaluating second derivatives at
     * the current minimum. Provides symmetric (parabolic) parameter errors.
     *
     * @return true if Hessian calculation succeeded
     */
    Bool_t ExecuteHesse();

    /**
     * @brief Executes MIGRAD command (gradient descent minimization).
     *
     * Runs Minuit2's MIGRAD algorithm, the recommended robust minimizer
     * using first derivatives and approximate Hessian updates.
     *
     * @return true if MIGRAD converged to a valid minimum
     */
    Bool_t ExecuteMigrad();

    /**
     * @brief Executes MINIMIZE command (automatic algorithm selection).
     *
     * Lets Minuit2 choose the best minimization strategy. Usually equivalent
     * to MIGRAD for well-behaved problems.
     *
     * @return true if minimization converged
     */
    Bool_t ExecuteMinimize();

    /**
     * @brief Executes MINOS command (asymmetric error analysis).
     *
     * Computes accurate asymmetric confidence intervals by scanning χ²
     * along each parameter axis. Slower but more accurate than HESSE.
     *
     * @return true if MINOS analysis completed
     */
    Bool_t ExecuteMinos();

    /**
     * @brief Executes PLOT command (visualize scan/contour results).
     *
     * Displays scan or contour data from previous SCAN/CONTOURS commands.
     *
     * @return true if plot generated successfully
     */
    Bool_t ExecutePlot();

    /**
     * @brief Executes PRINT command (set verbosity level).
     *
     * Controls Minuit2 output detail: 0=minimal, 1=normal, 2=debug.
     *
     * @param lineNo Command line number in MSR file
     * @return true if print level set successfully
     */
    Bool_t ExecutePrintLevel(UInt_t lineNo);

    /**
     * @brief Executes RELEASE command (unfreeze parameters).
     *
     * Allows previously fixed parameters to vary in subsequent fits.
     *
     * @param lineNo Command line number in MSR file
     * @return true if parameters released successfully
     */
    Bool_t ExecuteRelease(UInt_t lineNo);

    /**
     * @brief Executes RESTORE command (reload saved parameters).
     *
     * Restores parameter values from the last SAVE command.
     *
     * @return true if parameters restored successfully
     */
    Bool_t ExecuteRestore();

    /**
     * @brief Executes SCAN command (1D parameter space scan).
     *
     * Evaluates χ² along one or two parameter axes to visualize the
     * objective function landscape near the minimum.
     *
     * @return true if scan completed
     */
    Bool_t ExecuteScan();

    /**
     * @brief Executes SAVE command (store current parameters).
     *
     * Saves current parameter state for later RESTORE. Updates MSR file
     * statistics on first save (after final fit).
     *
     * @param first True if this is the first SAVE command in the session
     * @return true if parameters saved successfully
     */
    Bool_t ExecuteSave(Bool_t first);

    /**
     * @brief Executes SIMPLEX command (non-gradient minimization).
     *
     * Runs the Nelder-Mead simplex algorithm. Robust for rough objective
     * functions but slow to converge. Often used before MIGRAD for difficult fits.
     *
     * @return true if SIMPLEX found a minimum
     */
    Bool_t ExecuteSimplex();

    /**
     * @brief Prepares sector χ² analysis data structures.
     *
     * Initializes sector time windows and allocates storage for sector results.
     *
     * @param param Current parameter values
     * @param error Current parameter errors
     */
    void   PrepareSector(PDoubleVector &param, PDoubleVector &error);

    /**
     * @brief Executes SECTOR command (time-dependent χ² analysis).
     *
     * Calculates χ² for progressively wider time windows to identify
     * optimal fit ranges and systematic time-dependent effects.
     *
     * @param fout Output stream for sector analysis results
     * @return true if sector analysis completed
     */
    Bool_t ExecuteSector(std::ofstream &fout);

    //----------------------------------------------------------------------
    // Utility functions (private)
    //----------------------------------------------------------------------

    /**
     * @brief Returns current time in milliseconds.
     *
     * Used for timing fit commands and generating performance statistics.
     *
     * @return Timestamp in milliseconds since epoch
     */
    Double_t MilliTime();

    /**
     * @brief Rounds parameters for output with appropriate precision.
     *
     * Determines significant figures based on errors and formats parameters
     * for display in MSR file output.
     *
     * @param par Parameter values
     * @param err Parameter errors
     * @param ok Output flag: false if rounding failed
     * @return Rounded parameter values
     */
    PDoubleVector ParamRound(const PDoubleVector &par, const PDoubleVector &err, Bool_t &ok);
};

#endif // _PFITTER_H_