#  **************************************************************************  #
"""
#   BornAgain: simulate and fit reflection and scattering
#
#   @file      Wrap/Python/src/bornagain/ba_fit.py
#   @brief     Pure-Python FitObjective and IterationInfo, replacing Sim/Fitting C++.
#
#   @homepage  http://apps.jcns.fz-juelich.de/BornAgain
#   @license   GNU General Public License v3 or higher (see COPYING)
#   @copyright Forschungszentrum Juelich GmbH 2026
#   @authors   Scientific Computing Group at MLZ (see CITATION, AUTHORS)
"""
#  **************************************************************************  #

import numpy as np


class IterationInfo:
    """Stores fit iteration info to track fit flow from observers."""

    def __init__(self):
        self._count = 0
        self._chi2 = 0.0
        self._params = {}

    def update(self, params, chi2):
        self._count += 1
        self._chi2 = chi2
        self._params = dict(params)

    def iterationCount(self):
        """Returns current number of minimizer iterations."""
        return self._count

    def chi2(self):
        return self._chi2

    def parameterMap(self):
        """Returns map of fit parameter names and current values."""
        return self._params


class FitObjective:
    """
    Orchestrates fitting: runs simulations, compares to data, computes chi².

    Supports simultaneous fitting of several data sets, as demonstrated
    in example fit/scatter2d/multiple_datasets.py.
    """

    def __init__(self):
        self._pairs = []
        self._observers = []
        self._info = IterationInfo()
        self._metric = 'poisson'
        self._norm = 'l2'
        self._interrupted = False

    def addFitPair(self, sim_fn, exp_data, weight=1.0):
        """
        Adds a simulation builder and experimental data to the fit objective.

        Arguments:

        sim_fn:   a user-defined function that takes a parameter dict as input
                  and returns a BornAgain Simulation object.

        exp_data: a Datafield containing experimental data.

        weight:   user-defined weight of the dataset (default 1.0).
        """
        self._pairs.append({
            'sim_fn': sim_fn,
            'raw_data': exp_data,
            'weight': weight,
            'sim_result': None,
            'exp_result': None,
        })

    def setObjectiveMetric(self, metric, norm='l2'):
        """
        Sets the objective metric.

        metric: 'poisson' (default), 'chi2', 'log', or 'reldiff'
        norm:   'l2' (default) or 'l1'
        """
        self._metric = metric.lower()
        self._norm = norm.lower()

    def evaluate(self, params):
        """Runs simulations and returns the scalar metric value."""
        if self._interrupted:
            raise RuntimeError("Fitting was interrupted by the user.")
        if not self._pairs:
            raise RuntimeError(
                "Cannot start fit as no simulation/data pairs are defined."
            )

        params_dict = FitObjective._convert_params(params)
        self._exec_simulations(params_dict)
        chi2 = self._compute_metric()
        self._info.update(params_dict, chi2)

        count = self._info.iterationCount()
        for every_nth, callback in self._observers:
            if count % every_nth == 0:
                callback(self)

        return chi2

    def evaluate_residuals(self, params):
        """
        Runs simulations and returns the array of residuals exp - sim.

        This is the function to pass to lmfit.minimize for
        Levenberg-Marquardt fitting.
        """
        self.evaluate(params)
        exp = np.asarray(list(self.flatExpData()))
        sim = np.asarray(list(self.flatSimData()))
        return exp - sim

    def initPrint(self, every_nth):
        """Initializes printing to standard output on every_nth fit iteration."""
        from bornagain import ba_fitmonitor
        self.initPlot(every_nth, ba_fitmonitor.Printer())

    def initPlot(self, every_nth, callback):
        """Registers a callback to be called every every_nth fit iteration."""
        self._observers.append((every_nth, callback))

    def simulationResult(self, i=0):
        """Returns simulation result for the i-th dataset as a Datafield."""
        return self._pairs[i]['sim_result']

    def experimentalData(self, i=0):
        """Returns experimental data for the i-th dataset as a Datafield."""
        return self._pairs[i]['exp_result']

    def relativeDifference(self, i=0):
        """Returns relative difference between sim and exp as a Datafield."""
        sim = self._pairs[i]['sim_result']
        exp = self._pairs[i]['exp_result']
        sim_v = np.asarray(list(sim.flatVector()))
        exp_v = np.asarray(list(exp.flatVector()))
        denom = np.abs(sim_v) + np.abs(exp_v)
        rel = np.where(denom > 0, 2.0*np.abs(sim_v - exp_v)/denom, 0.0)
        import bornagain as ba
        return ba.Datafield(sim.frame(), list(rel))

    def absoluteDifference(self, i=0):
        """Returns absolute difference between sim and exp as a Datafield."""
        sim = self._pairs[i]['sim_result']
        exp = self._pairs[i]['exp_result']
        sim_v = np.asarray(list(sim.flatVector()))
        exp_v = np.asarray(list(exp.flatVector()))
        diff = np.abs(sim_v - exp_v)
        import bornagain as ba
        return ba.Datafield(sim.frame(), list(diff))

    def flatSimData(self):
        """Returns merged simulated intensities from all datasets."""
        result = []
        for pair in self._pairs:
            result.extend(pair['sim_result'].flatVector())
        return result

    def flatExpData(self):
        """Returns merged experimental data from all datasets."""
        result = []
        for pair in self._pairs:
            result.extend(pair['exp_result'].flatVector())
        return result

    def nPairs(self):
        """Returns number of simulation/data pairs."""
        return len(self._pairs)

    def iterationInfo(self):
        """Returns current IterationInfo."""
        return self._info

    def interruptFitting(self):
        """Requests that fitting be stopped after the current iteration."""
        self._interrupted = True

    def isInterrupted(self):
        """Returns True if fitting was interrupted."""
        return self._interrupted

    @staticmethod
    def _convert_params(params):
        """Converts lmfit.Parameters or plain dict to a plain Python dict."""
        if str(getattr(params, '__module__', '')) == "lmfit.parameter":
            return {p: params[p].value for p in params}
        if hasattr(params, 'items'):
            return {k: float(v) for k, v in params.items()}
        return dict(params)

    def _exec_simulations(self, params_dict):
        for pair in self._pairs:
            try:
                sim = pair['sim_fn'](params_dict)
            except Exception as exc:
                raise RuntimeError(
                    f"FitObjective: simulation builder raised "
                    f"{type(exc).__name__}: {exc}") from exc
            pair['sim_result'] = sim.simulate()
            if pair['exp_result'] is None:
                pair['exp_result'] = _reposition_exp(sim, pair['raw_data'])

    def _compute_metric(self):
        metric = self._metric
        norm = self._norm
        eps = np.finfo(float).tiny
        total = 0.0
        for pair in self._pairs:
            sim_v = np.asarray(list(pair['sim_result'].flatVector()))
            exp_v = np.asarray(list(pair['exp_result'].flatVector()))
            errs = np.asarray(list(pair['exp_result'].errorSigmas()))
            has_unc = len(errs) == len(sim_v) and len(errs) > 0

            mask = exp_v >= 0.0

            if metric == 'chi2':
                if has_unc:
                    unc_mask = mask & (errs > 0.0)
                    r = np.where(unc_mask, (exp_v - sim_v)/errs, 0.0)
                else:
                    r = np.where(mask, exp_v - sim_v, 0.0)
            elif metric in ('poisson', 'poisson-like'):
                if has_unc:
                    unc_mask = mask & (errs > 0.0)
                    r = np.where(unc_mask, (exp_v - sim_v)/errs, 0.0)
                else:
                    variance = np.maximum(1.0, sim_v)
                    r = np.where(mask, (sim_v - exp_v)/np.sqrt(variance),
                                 0.0)
            elif metric == 'log':
                if has_unc:
                    ln10 = np.log(10.0)
                    unc_mask = mask & (errs > 0.0)
                    sim_s = np.maximum(eps, sim_v)
                    exp_s = np.maximum(eps, exp_v)
                    safe_errs = np.where(unc_mask, errs, 1.0)
                    val = (np.log10(sim_s) -
                           np.log10(exp_s))*exp_s*ln10/safe_errs
                    r = np.where(unc_mask, val, 0.0)
                else:
                    sim_s = np.maximum(eps, sim_v)
                    exp_s = np.maximum(eps, exp_v)
                    r = np.where(mask,
                                 np.log10(sim_s) - np.log10(exp_s), 0.0)
            elif metric in ('reldiff', 'relative'):
                if has_unc:
                    unc_mask = mask & (errs > 0.0)
                    r = np.where(unc_mask, (exp_v - sim_v)/errs, 0.0)
                else:
                    sim_s = np.maximum(eps, sim_v)
                    exp_s = np.maximum(eps, exp_v)
                    r = np.where(mask, (exp_s - sim_s)/(exp_s + sim_s),
                                 0.0)
            else:
                raise ValueError(
                    f"Unknown metric: {metric!r}. "
                    "Valid options: 'chi2', 'poisson', 'log', 'reldiff'.")

            if norm == 'l2':
                contrib = float(np.sum(r*r))
            elif norm == 'l1':
                contrib = float(np.sum(np.abs(r)))
            else:
                raise ValueError(
                    f"Unknown norm: {norm!r}. Valid options: 'l2', 'l1'.")

            total += pair['weight']*contrib

        return total


def _reposition_exp(simulation, data):
    """
    Repositions experimental data to match the simulation's detector ROI frame.

    For ScatteringSimulation, crops or maps the data to the active ROI pixels.
    For specular and other simulations, returns the data unchanged.
    """
    if not hasattr(simulation, 'detector'):
        return data

    det = simulation.detector()
    clipped_frame = det.clippedFrame()
    active_indices = list(det.activeIndices())

    n = clipped_frame.size()
    flat_data = list(data.flatVector())
    flat_errs = list(data.errorSigmas())
    has_errors = len(flat_errs) == len(flat_data) and len(flat_errs) > 0

    values = [0.0]*n
    errors = [0.0]*n if has_errors else []

    if data.frame().hasSameSizes(clipped_frame):
        # Data is already ROI-sized
        for i in active_indices:
            values[i] = flat_data[i]
            if has_errors:
                errors[i] = flat_errs[i]
    else:
        # Expect data to have full detector shape
        data_matches_det = (data.rank() == 2
                            and data.axis(0).size() == det.axis(0).size()
                            and data.axis(1).size() == det.axis(1).size())
        if not data_matches_det:
            raise RuntimeError(
                "FitObjective: Detector and experimental data have different shape"
            )
        for i in active_indices:
            full_i = det.roiToFullIndex(i)
            values[i] = flat_data[full_i]
            if has_errors:
                errors[i] = flat_errs[full_i]

    import bornagain as ba
    if has_errors:
        return ba.Datafield(clipped_frame, values, errors)
    return ba.Datafield(clipped_frame, values)
