"""Stabilization Score (SS) and Stabilization Uplift (SU) as a Hugging Face `evaluate` module.

The metric logic lives in the `stabilization-uplift` package on PyPI; this module only wraps it.
"""

import datasets
import evaluate
from stabilization_uplift import stabilization_score, stabilization_uplift

_CITATION = """\
@inproceedings{varshavskiy2025mitigating,
  title     = {Mitigating Model Drift in Developing Economies Using Synthetic Data and Outliers},
  author    = {Varshavskiy, Ilyas and Boboeva, Bonu and Khalilbekov, Shuhrat and Azimi, Azizjon and Shulgin, Sergey and Nizamitdinov, Akhlitdin and S{\\'a}ez de Oc{\\'a}riz Borde, Haitz},
  booktitle = {NeurIPS 2025 Workshop on Generative AI in Finance},
  year      = {2025},
  url       = {https://openreview.net/forum?id=zfTaFD0B5Z}
}
"""

_DESCRIPTION = """\
Stabilization Score (SS) and Stabilization Uplift (SU) evaluate how stable a classifier stays under a sudden
distribution shift, such as a macroeconomic shock. SS measures how much one model's ROC AUC changed between the
base (pre-shock) and shock periods, normalized by the severity of the shift. SU measures whether model B
(e.g. trained with synthetic outliers) is more stable and not worse than model A (a baseline) under the shock.
"""

_KWARGS_DESCRIPTION = """
Each row is one comparison of model B with model A.

Args:
    auc_base_A (`list` of `float`): ROC AUC of model A on the base (pre-shock) test set, in (0, 1].
    auc_shock_A (`list` of `float`): ROC AUC of model A on the shock test set, in (0, 1].
    auc_base_B (`list` of `float`): ROC AUC of model B on the base test set, in (0, 1].
    auc_shock_B (`list` of `float`): ROC AUC of model B on the shock test set, in (0, 1].
    dist_shift (`list` of `float`): distribution shift between the base and shock data, normally in [0, 1].
        Compute it with `stabilization_uplift.distribution_shift` (pip install "stabilization-uplift[shift]").

Returns:
    stabilization_score_A (`list` of `float`): SS of model A for each row, in [0.5, 1]. 1 means no change in AUC.
    stabilization_score_B (`list` of `float`): SS of model B for each row, in [0.5, 1].
    stabilization_uplift (`list` of `float`): SU of B over A for each row, in [0, 1). 0 means no uplift.

Examples:
    >>> metric = evaluate.load("zyplai/stabilization-uplift")
    >>> results = metric.compute(auc_base_A=[0.80], auc_shock_A=[0.70],
    ...                          auc_base_B=[0.80], auc_shock_B=[0.81], dist_shift=[0.2])
    >>> {k: [round(v, 3) for v in vals] for k, vals in results.items()}
    {'stabilization_score_A': [0.915], 'stabilization_score_B': [0.992], 'stabilization_uplift': [0.725]}
"""

_FEATURE_NAMES = ["auc_base_A", "auc_shock_A", "auc_base_B", "auc_shock_B", "dist_shift"]


@evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
class StabilizationUplift(evaluate.Metric):
    def _info(self):
        return evaluate.MetricInfo(
            description=_DESCRIPTION,
            citation=_CITATION,
            inputs_description=_KWARGS_DESCRIPTION,
            features=datasets.Features({name: datasets.Value("float64") for name in _FEATURE_NAMES}),
            homepage="https://pypi.org/project/stabilization-uplift/",
            codebase_urls=["https://github.com/zypl-ai/stabilization_uplift"],
            reference_urls=["https://openreview.net/forum?id=zfTaFD0B5Z", "https://arxiv.org/abs/2510.09294"],
        )

    def _compute(self, auc_base_A, auc_shock_A, auc_base_B, auc_shock_B, dist_shift):
        rows = list(zip(auc_base_A, auc_shock_A, auc_base_B, auc_shock_B, dist_shift))
        return {
            "stabilization_score_A": [stabilization_score(a_base, a_shock, shift) for a_base, a_shock, _, _, shift in rows],
            "stabilization_score_B": [stabilization_score(b_base, b_shock, shift) for _, _, b_base, b_shock, shift in rows],
            "stabilization_uplift": [stabilization_uplift(*row) for row in rows],
        }