text: # 1/usr/bin/env python3
  
  """ Base module for performing fits with Minuit.
  
  .. code-author: Raymond Ehlers <raymond.ehlers@cern.ch>, Yale University
  """
  
  import functools
  import itertools
  import logging
  import time
  from dataclasses import dataclass
  from typing import (
      TYPE_CHECKING,
      Any,
      Callable,
      Dict,
      Iterable,
      List,
      Optional,
      Sequence,
      Tuple,
      Type,
      TypeVar,
      Union,
      cast,
  )
  
  import iminuit
  import numdifftools as nd
  import numpy as np
  import numpy.typing as npt
  
  from pachyderm import generic_class
  
  
  if TYPE_CHECKING:
      from pachyderm.fit import cost_function
  
  logger = logging.getLogger(__name__)
  
  # Typing
  T_FuncCode = TypeVar("T_FuncCode", bound="FuncCode")
  T_ArgumentPositions = List[List[int]]
  _T_FitResult = TypeVar("_T_FitResult", bound="FitResult")
  
  
  class FitFailed(Exception):
      """Raised if the fit failed. The message will include further details."""
  
      pass
  
  
  @dataclass
  class BaseFitResult:
      """Base fit result.
  
      This represents the most basic fit result.
  
      Attributes:
          parameters: Names of the parameters used in the fit.
          free_parameters: Names of the free parameters used in the fit.
          fixed_parameters: Names of the fixed parameters used in the fit.
          values_at_minimum: Contains the values of the full RP fit function at the minimum. Keys are the
              names of parameters, while values are the numerical values at convergence.
          errors_on_parameters: Contains the values of the errors associated with the parameters
              determined via the fit.
          covariance_matrix: Contains the values of the covariance matrix. Keys are tuples
              with (param_name_a, param_name_b), and the values are covariance between the specified parameters.
              Note that fixed parameters are _not_ included in this matrix.
          errors: Store the errors associated with the component fit function.
      """
  
      parameters: List[str]
      free_parameters: List[str]
      fixed_parameters: List[str]
      values_at_minimum: Dict[str, float]
      errors_on_parameters: Dict[str, float]
      covariance_matrix: Dict[Tuple[str, str], float]
      errors: npt.NDArray[Any]
  
      @property
      def correlation_matrix(self) -> Dict[Tuple[str, str], float]:
          """The correlation matrix of the free parameters.
  
          These values are derived from the covariance matrix values stored in the fit.
  
          Note:
              This property caches the correlation matrix value so we don't have to calculate it every time.
  
          Args:
              None
          Returns:
              The correlation matrix of the fit result.
          """
          try:
              # We attempt to cache the covariance matrix, so first try to return that.
              return self._correlation_matrix
          except AttributeError:
  
              def corr(i_name: str, j_name: str) -> float:
                  """Calculate the correlation matrix (definition from iminuit) from the covariance matrix."""
                  # The + 1e-100 is just to ensure that we don't divide by 0.
                  value = self.covariance_matrix[(i_name, j_name)] / (
                      np.sqrt(self.covariance_matrix[(i_name, i_name)] * self.covariance_matrix[(j_name, j_name)])
                      + 1e-100
                  )
                  # Need to explicitly cast to float. Otherwise, it will return a np.float64, which will cause problems
                  # for YAML...
                  return float(value)
  
              matrix: Dict[Tuple[str, str], float] = {}
              for i_name in self.free_parameters:
                  for j_name in self.free_parameters:
                      matrix[(i_name, j_name)] = corr(i_name, j_name)
  
              self._correlation_matrix = matrix
  
          return self._correlation_matrix
  
  
  @dataclass
  class FitResult(BaseFitResult):
      """Main fit result class.
  
      Note:
          free_parameters + fixed_parameters == parameters
  
      Attributes:
          parameters: Names of the parameters used in the fit.
          free_parameters: Names of the free parameters used in the fit.
          fixed_parameters: Names of the fixed parameters used in the fit.
          values_at_minimum: Contains the values of the full RP fit function at the minimum. Keys are the
              names of parameters, while values are the numerical values at convergence.
          errors_on_parameters: Contains the values of the errors associated with the parameters
              determined via the fit.
          covariance_matrix: Contains the values of the covariance matrix. Keys are tuples
              with (param_name_a, param_name_b), and the values are covariance between the specified parameters.
              Note that fixed parameters are _not_ included in this matrix.
          errors: Store the errors associated with the component fit function.
          x: x values where the fit result should be evaluated.
          n_fit_data_points: Number of data points used in the fit.
          minimum_val: Minimum value of the fit when it coverages. This is the chi squared value for a
              chi squared minimization fit.
      """
  
      x: npt.NDArray[Any]
      n_fit_data_points: int
      minimum_val: float
  
      @property
      def nDOF(self) -> int:
          """Number of degrees of freedom."""
          return self.n_fit_data_points - len(self.free_parameters)
  
      def effective_chi_squared(self, cost_func: "cost_function.DataComparisonCostFunction") -> float:
          """Calculate the effective chi squared value.
  
          If the fit was performed using a chi squared cost function, it's just equal to
          the ``minimal_val``. If it's log likelihood, one must calculate the effective
          chi squared.
  
          Note:
              We attempt to cache this value so we don't have to calculate it every time.
  
          Args:
              cost_function: Cost function used to create the fit function.
              data: Data to be used to calculate the chi squared.
          Returns:
              The effective chi squared value.
          """
          try:
              # We attempt to cache the chi squared, so first try to return that.
              return self._chi_squared
          except AttributeError:
              # Setup
              from pachyderm import histogram
              from pachyderm.fit import cost_function
  
              # Calculate the chi_squared
              self._chi_squared: float
              if isinstance(cost_func, (cost_function.ChiSquared, cost_function.BinnedChiSquared)):
                  self._chi_squared = self.minimum_val
              elif isinstance(cost_func, cost_function.BinnedLogLikelihood):
                  data = cost_func.data
                  # Help out mypy...
                  assert isinstance(data, histogram.Histogram1D)
                  # Calculate using the binned chi squared
                  self._chi_squared = cost_function.binned_chi_squared_safe_for_zeros(
                      data.x, data.y, data.errors, data.bin_edges, cost_func.f, *self.values_at_minimum.values()
                  )
              elif isinstance(cost_func, cost_function.SimultaneousFit):
                  chi_squared = 0
                  for cf, function_args in zip(
                      cost_function.unravel_simultaneous_fits(cost_func.cost_functions),
                      _function_arguments_from_argument_positions(
                          cost_func.argument_positions, *self.values_at_minimum.values()
                      ),
                  ):
                      data = cf.data
                      # Help out mypy...
                      assert isinstance(data, histogram.Histogram1D)
                      # Calculate using the binned chi squared. This assumes that we have binned data (which
                      # is probably fairly reasonable for our purpose).
                      chi_squared += cost_function.binned_chi_squared_safe_for_zeros(
                          data.x, data.y, data.errors, data.bin_edges, cf.f, *function_args
                      )
                  self._chi_squared = chi_squared
              else:
                  raise NotImplementedError("Needs to be implement for unbinned data.")
                  # data = cost_func.data
                  ## Help out mypy...
                  # assert isinstance(data, histogram.Histogram1D)
                  # self._chi_squared = cost_function._chi_squared(
                  #    data.x, data.y, data.errors, data.bin_edges, cost_func.f, *self.values_at_minimum.values()
                  # )
  
          return self._chi_squared
  
      @classmethod
      def from_minuit(
          cls: Type[_T_FitResult], minuit: iminuit.Minuit, cost_func: Callable[..., float], x: npt.NDArray[Any]
      ) -> _T_FitResult:
          """Create a fit result form the Minuit fit object.
  
          Args:
              minuit: Minuit fit object after performing the fit.
              cost_func: Cost function used to perform the fit.
          """
          # Validation
          if not minuit.valid:
              raise RuntimeError("The fit is invalid - unable to extract result!")
  
          # Determine the relevant fit parameters.
          fixed_parameters: List[str] = [k for k in minuit.parameters if minuit.fixed[k] is True]
          # We use the cost function because we want intentionally want to skip "x"
          parameters: List[str] = iminuit.util.describe(cost_func)
          # Can't just use set(parameters) - set(fixed_parameters) because set() is unordered!
          free_parameters = [p for p in parameters if p not in set(fixed_parameters)]
          # Store the result
          return cls(
              parameters=parameters,
              free_parameters=free_parameters,
              fixed_parameters=fixed_parameters,
              values_at_minimum={k: minuit.values[k] for k in minuit.parameters},
              errors_on_parameters={k: minuit.errors[k] for k in minuit.parameters},
              covariance_matrix=minuit.covariance,
              x=x,
              n_fit_data_points=len(x),
              minimum_val=minuit.fval,
              errors=np.array([]),
          )
  
      def __eq__(self, other: Any) -> bool:
          """Check for equality of members.
  
          This is pretty hacky, but the code is already brittle, and it passes the tests, so it's
          good enough for these purposes.
          """
          # Check identity to avoid needing to perform the (potentially costly) dict comparison.
          if self is other:
              return True
          # Compare via the member values.
          if type(other) is type(self):
              keys_agree = list(self.__dict__.keys()) == list(other.__dict__)
              try:
                  values_agree = [
                      v == other.__dict__[k] if not isinstance(v, np.ndarray) else np.allclose(v, other.__dict__[k])
                      for k, v in self.__dict__.items()
                  ]
              except KeyError:
                  values_agree = [False]
              return keys_agree and all(values_agree)
          return NotImplemented
  
  
  def extract_function_values(func: Callable[..., float], fit_result: BaseFitResult) -> Tuple[Dict[str, Any], List[str]]:
      """Extract the parameters relevant to the given function from a fit result.
  
      Note:
          The fit result may have more arguments at minimum and free parameters than the fit function that we've
          passed (for example, if we've calculating the background parameters for the inclusive signal fit), so
          we need to determine the free parameters here.
  
      Args:
          func: Function of interest.
          fit_result: Fit result from which we will extract the information.
      Returns:
          ((kw)args at minimum, free parameters)
      """
      # Setup the parameters needed to execute the function.
      # Determine relevant parameters for the given function
      func_parameters = iminuit.util.describe(func)
      # Determine the arguments for the fit function
      # NOTE: The fit result may have more arguments at minimum and free parameters than the fit function that we've
      #       passed (for example, if we've calculating the background parameters for the inclusive signal fit), so
      #       we need to determine the free parameters here.
      args_at_minimum = {k: v for k, v in fit_result.values_at_minimum.items() if k in func_parameters}
      # Retrieve the parameters to use in calculating the fit errors.
      free_parameters = [p for p in fit_result.free_parameters if p in func_parameters]
  
      return args_at_minimum, free_parameters
  
  
  def calculate_function_errors(
      func: Callable[..., float], fit_result: BaseFitResult, x: npt.NDArray[Any]
  ) -> npt.NDArray[Any]:
      """Calculate the errors of the given function based on values from the fit.
  
      Note:
          We don't take the x values for the fit_result as it may be desirable to calculate the errors for
          only a subset of x values. Plus, the component fit result doesn't store the x values, so it would
          complicate the validation. It's much easier to just require the user to pass the x values (and it takes
          little effort to do so).
  
      Args:
          func: Function to use in calculating the errors.
          fit_result: Fit result for which the errors will be calculated.
          x: x values where the errors will be evaluated.
      Returns:
          The calculated error values.
      """
      # Setup the parameters needed to execute the function.
      # Determine the arguments for the fit function and the parameters to use in calculating the fit errors.
      args_at_minimum, free_parameters = extract_function_values(func, fit_result)
      # To calculate the error, we need to match up the parameter names to their index in the arguments list
      args_at_minimum_keys = list(args_at_minimum)
      name_to_index = {name: args_at_minimum_keys.index(name) for name in free_parameters}
      logger.debug(f"args_at_minimum: {args_at_minimum}")
      logger.debug(f"free_parameters: {free_parameters}")
      logger.debug(f"name_to_index: {name_to_index}")
  
      # Evaluate the gradient
      partial_derivative_result = evaluate_gradient(func=func, fit_result=fit_result, x=x)
  
      # Finally, calculate the error by multiplying the matrix of gradient values by the covariance matrix values.
      error_vals = np.zeros(len(x))
      for i_name in free_parameters:
          for j_name in free_parameters:
              # Determine the error value
              # logger.debug(f"Calculating error for i_name: {i_name}, j_name: {j_name}")
              # Add error to overall error value
              # NOTE: This is a vector operation for the partial_derivative_result values.
              error_vals += (
                  partial_derivative_result[:, name_to_index[i_name]]
                  * partial_derivative_result[:, name_to_index[j_name]]
                  * fit_result.covariance_matrix[(i_name, j_name)]
              )
      # logger.debug("error_val: shape: {error_val.shape}, error_val: {error_val}")
  
      # We want the error itself, so we take the square root.
      res: npt.NDArray[Any] = np.sqrt(error_vals)
      return res
  
  
  def evaluate_gradient(func: Callable[..., float], fit_result: BaseFitResult, x: npt.NDArray[Any]) -> npt.NDArray[Any]:
      """Evaluate the gradient of the given function based on the fit values.
  
      For a function of 5 free parameters (7 total) and 10 x values, the returned result would be of the shape (10, 5).
  
      Args:
          func: Function to use in calculating the errors.
          fit_result: Fit result for which the errors will be calculated.
          x: x values where the errors will be evaluated.
      Returns:
          For each x value, the gradient is evaluated for each free parameter. It will be of the
          shape (len(x_values), len(free_parameters)).
      """
      # Setup the parameters needed to execute the function.
      # Determine the arguments for the fit function and the parameters to use in calculating the fit errors.
      args_at_minimum, free_parameters = extract_function_values(func, fit_result)
  
      # To take the gradient, ``numdifftools`` requires a particular function signature. The first argument
      # must contain a list of values that it will vary when taking the gradient. The rest of the args are
      # passed on to the function via *args and **kwargs, but they won't be varied.
      # To ensure the proper signature, we wrap the function and route the arguments.
      def func_wrap(args_to_vary: Sequence[float], x: npt.NDArray[Any]) -> float:
          """Wrap the given function to ensure that the arguments are routed properly for ``numdifftools``.
  
          To take the gradient, ``numdifftools`` requires a particular function signature. The first argument
          must contain a list of values that it will vary when taking the gradient. The rest of the args are
          passed on to the function via *args and **kwargs, but they won't be varied (we don't event use those
          generic arguments here). To ensure the proper signature given any function, we wrap the function and
          route the arguments.
  
          Args:
              args_to_vary: List of arguments to vary when taking the gradient. This should correspond
                  to the value of the free parameters.
              x: x value(s) where the function will be evaluated.
          Returns:
              Function evaluated at the given values.
          """
          # Need to expand the arguments
          return func(x, *args_to_vary)
  
      # Setup to compute the derivative
      partial_derivative_func = nd.Gradient(func_wrap)
  
      logger.debug("Calculating the gradient")
      # We time it to keep track of how long it takes to evaluate. Sometimes it can be a bit slow.
      start = time.time()
      # Actually evaluate the gradient.
      # It returns an array with dimensions (len(x), len(args_to_vary)) containing the derivatives with
      # respect to each parameter at each point.
      #
      # NOTE: In principle, we've doing some unnecessary work because we also calculate the gradient with
      #       respect to fixed parameters. But due to the argument requirements of ``numdifftools``, it would be
      #       quite difficult to tell it to only take the gradient with respect to a non-continuous selection of
      #       parameters. So we just accept the inefficiency.
      partial_derivative_result: npt.NDArray[Any] = partial_derivative_func(list(args_at_minimum.values()), x)
      end = time.time()
      logger.debug(f"Finished calculating the gradient in {end-start} seconds.")
  
      # If we are only in 1D, we need to promote to a 2D (shape of 1D, 1) to get the indexing correct below.
      # This only occurs if we only have one parameter that is varied.
      if partial_derivative_result.ndim == 1:
          logger.debug(f"shape before adding axis: {partial_derivative_result.shape}")
          partial_derivative_result = partial_derivative_result[:, np.newaxis]
          logger.debug(f"shape after adding axis : {partial_derivative_result.shape}")
  
      return partial_derivative_result
  
  
  class FuncCode(generic_class.EqualityMixin):
      """Minimal class to describe function arguments.
  
      Same approach as is taken in ``iminuit``. Note that the precise name of the parameters is
      extremely important.
  
      Args:
          args: List of function arguments.
      Attributes:
          co_varnames: Name of the function arguments.
          co_argcount: Number of function arguments.
      """
  
      __slots__ = ("co_varnames", "co_argcount")
  
      def __init__(self, args: List[str]):
          self.co_varnames = args
          self.co_argcount = len(args)
  
      def __repr__(self) -> str:
          return f"FuncCode({self.co_varnames})"
  
      @classmethod
      def from_function(
          cls: Type[T_FuncCode], func: Callable[..., float], leading_parameters_to_remove: int = 1
      ) -> T_FuncCode:
          """Create a func_code from a function.
  
          Args:
              func: Function for which we want a func_code.
              leading_parameters_to_remove: Number of leading parameters to remove in the func_code. Default: 1,
                  which corresponds to ``x`` as the first argument.
          """
          return cls(iminuit.util.describe(func)[leading_parameters_to_remove:])
  
  
  def merge_func_codes(
      functions: Iterable[Callable[..., float]],
      prefixes: Optional[Sequence[str]] = None,
      skip_prefixes: Optional[Sequence[str]] = None,
  ) -> Tuple[List[str], List[List[int]]]:
      """Merge the arguments of the given functions into one func_code.
  
      Note:
          This has very similar functionality and is heavily inspired by ``Probfit.merge_func_code...)``.
  
      Args:
          functions: Functions whose arguments are to be merged.
          prefixes: Prefix for arguments of each function. Default: None. If specified, there must
              be one prefix for each function.
          skip_prefixes: Prefixes to skip when assigning prefixes. As noted in probfit, this can be
              useful to mix prefixed and non-prefixed arguments. Default: None.
      Returns:
          Merged list of arguments, map from merged arguments to arguments for each individual function.
      """
      # Validation
      # Ensure that we don't exhaust the iterator during validation.
      funcs = list(functions)
      # Ensure that we have the proper number of prefixes.
      if prefixes:
          if len(funcs) != len(prefixes):
              raise ValueError("Number of prefixes ({len(prefixes)} doesn't match the number of functions: {len(funcs)}")
      else:
          # Create an empty prefix array so we can zip with it.
          prefixes = ["" for _ in funcs]
      skip_prefix = set(skip_prefixes if skip_prefixes else [])
  
      # Retrieve all of the args.
      args = []
      for f, pre in zip(funcs, prefixes):
          temp = []
          for arg in iminuit.util.describe(f):
              value = f"{pre}_{arg}" if pre and arg not in skip_prefix else arg
              temp.append(value)
          args.append(temp)
  
      # Determine the unique arguments.
      # We want to ensure that we maintain the oder, so we use dict.fromkeys to do so.
      merged_args = list(dict.fromkeys(itertools.chain.from_iterable(args)))
  
      # Determine the map from merged arguments to arguments for each individual function.
      argument_positions = []
      for func_args in args:
          positions = []
          for arg in func_args:
              positions.append(merged_args.index(arg))
          argument_positions.append(positions)
  
      logger.debug(f"funcs: {funcs}, args: {args}")
      logger.debug(f"merged args: {merged_args}")
      logger.debug(f"argument_positions: {argument_positions}")
  
      return merged_args, argument_positions
  
  
  def _function_arguments_from_argument_positions(
      argument_positions: T_ArgumentPositions, *args: Union[float, npt.NDArray[Any]]
  ) -> Iterable[List[Union[float, npt.NDArray[Any]]]]:
      """Extract the function arguments from a larger set of arguments given the argument positions.
  
      It determines the arguments for calling functions stored in ``SimultaneousFit`` or ``AddPDF`` classes. Use it
      something like:
  
      >>> pdf = AddPDF(f_1, f_2)
      >>> for func, function_args in zip(pdf.functions,
      ...                                _function_arguments_from_argument_positions(
      ...                                    pdf.argument_positions, *full_set_of_args
      ...                                )):
      ...     # Call the function
      ...     func(*function_args)
  
      Args:
          argument_positions: Map from merged arguments to arguments for each function.
          args: Arguments for the functions. Must include the x argument!
      Returns:
          List of arguments for the function that is described by the maps in argument_positions.
      """
      for arg_positions in argument_positions:
          # logger.debug(f"arg_positions: {arg_positions}, full args: {args}")
          # Determine the arguments for the given function using the arg positions map.
          function_args = []
          for v in arg_positions:
              # We don't skip over the x argument because it's expected to be supplied in the args.
              function_args.append(args[v])
          # logger.debug(f"arg_positions: {arg_positions}, function_args: {function_args}")
          yield function_args
  
  
  def call_list_of_callables_with_operation(
      operation: Callable[[float, float], float],
      functions: Iterable[Callable[..., float]],
      argument_positions: T_ArgumentPositions,
      *args: Union[float, npt.NDArray[Any]],
  ) -> float:
      """Call and add a list of callables with the given args.
  
      Args:
          functions: Functions to be evaluated.
          argument_positions: Map from merged arguments to arguments for each function.
          args: Arguments for the functions. Must include the x argument!
      Returns:
          Sum of the values of the functions.
      """
      values = []
      # We write it out explicitly instead of using a generator because it's long enough that it would be difficult
      # to understand.
      for func, function_args in zip(functions, _function_arguments_from_argument_positions(argument_positions, *args)):
          values.append(func(*function_args))
      return functools.reduce(operation, values)
  
  
  def chi_squared_probability(chi_2: float, ndf: float) -> float:
      """Calculate the probability that the
  
      This is just a thin wrapped around ``scipy.stats``, but it's convenient.
  
      Args:
          chi_2: Chi squared value.
          ndf: Number of degrees of freedom.
      Returns:
          Probability that the fit is consistent with the data.
      """
      import scipy.stats
  
      return cast(float, 1 - scipy.stats.chi2.cdf(chi_2, ndf))
  