# SPDX-FileCopyrightText: Copyright © 2025 Idiap Research Institute <contact@idiap.ch>
#
# SPDX-License-Identifier: GPL-3.0-or-later

"""Solver CLI for demographically fair machine learning systems.

This module defines the ``solve`` command, which loads one or more system score output
files in JSON format, validates them against a predefined data model, and computes
fairness-related indicators such as accuracy and equalized odds. It outputs JSON files
containing solutions (or operating-modes) following the chosen metrics space, and
thresholds.

Input files must conform to a structured data model including scores, ground-truth
labels, and protected attributes. The CLI provides options to customize metrics, and
select thresholds.
"""

import logging
import pathlib
import typing

logger = logging.getLogger(__name__.split(".", 1)[0])

import click

from ..metrics import supported_metrics
from ..scores import Scores
from ..solutions import Solutions
from . import utils


@click.command(
    epilog="""Examples:

\b
  1. Compute and store solutions from two systems on the current directory:

     .. code:: sh

        fairical solve scores/system1.json scores/system2.json

  2. Compute and store a-priori solutions for system1, using the thresholds
     from the non-dominated solutions of system2:

     .. code:: sh

        fairical solve scores/system1.json@solutions/system2_nds.json

""",
)
@click.argument(
    "system",
    type=str,
    required=True,
    nargs=-1,
    callback=utils.validate_scores_json,
)
@click.option(
    "-d/-D",
    "--deduplicate/--no-deduplicate",
    default=True,
    show_default=True,
    help="If we should prune (deduplicate) similar solutions.",
)
@click.option(
    "--ds/--nds",
    "dominated",
    default=None,
    flag_value=True,
    help="This flag only concerns *a priori* analysis. If --ds is set, "
    "then only applies only dominated solutions to input score set, if "
    "--nds, only applies non-dominated solutions. If neither of these "
    "flags is set, then apply all solutions from the input solution set "
    "during analysis.",
)
@click.option(
    "-o",
    "--output-path",
    type=click.Path(
        dir_okay=True, file_okay=False, writable=True, path_type=pathlib.Path
    ),
    required=True,
    default=pathlib.Path(),
    show_default=True,
    help="Directory where to store system solutions.",
)
@click.option(
    "-m",
    "--metric",
    type=str,
    multiple=True,
    help=f"Metric to consider for solution evaluation. Values are considered in the "
    f"order they are provided forming the axes of a potentially multi-dimensional "
    f"surface where dominated and non-dominated solutions are evaluated. Valid values "
    f"are: {', '.join(supported_metrics())}, where ``<attr>`` refers to attributes in "
    f"the provided sensitive attribute dictionary, and ``<util>`` refers to any "
    f"supported utility metrics.  This option is ignored for *a priori* analysis.",
    default=["acc", "eod+gender"],
    callback=utils.validate_metrics,
    show_default=True,
)
@click.option(
    "-T",
    "--thresholds",
    type=click.IntRange(2),
    help="If set, then run for a fixed number of equally-spaced thresholds in the "
    "interval [0.0, 1.0].  Here, you only specify the number of thresholds to use, "
    "with the minimum being 2. If you want explicitly set the thresholds, you should "
    "use the API directly. If unset, then let scikit-learn calculate the number of "
    "sensible thresholds based the scores and changes on the ROC or PR-curve "
    "depending on utility metrics chosen.",
    default=None,
    show_default=True,
    callback=utils.validate_thresholds,
)
@utils.verbosity_option(logger)
def solve(
    system: list[
        tuple[pathlib.Path, Scores]
        | tuple[tuple[pathlib.Path, Scores], tuple[pathlib.Path, Solutions]]
    ],
    deduplicate: bool,
    dominated: bool | None,
    output_path: pathlib.Path,
    metric: list[str],
    thresholds: list[float] | None,
    verbose: int,
) -> None:  # numpydoc ignore=PR01
    """Find operating-modes (a.k.a. solutions) for ML systems.

    This module defines the ``solve`` command, which loads one or more system score
    output files in JSON format, validates them against a predefined data model, and
    computes fairness-related indicators such as accuracy and equalized odds. It outputs
    JSON files containing solutions (or operating-modes) following the chosen metrics
    space, and thresholds.

    Input JSON files must be formatted as described in the documentation.
    """
    # 0. validate that the selected protected attributes do exist inside JSON system
    # TODO: validate metrics against user loaded systems

    # 2. for each system provided as input, get solutions
    for k in system:
        if isinstance(k[0], pathlib.Path):
            # solutions chosen a posteriori
            score_fn, scores = typing.cast(tuple[pathlib.Path, Scores], k)

            message = "Calculating solutions *a posteriori* for system "
            message += f"`{str(score_fn)}`, using metrics "
            message += f"`{', '.join(metric)}`..."
            logger.info(message)

            solutions = scores.solutions_a_posteriori(metric, thresholds=thresholds)

            output_filename = (output_path / k[0].stem).with_suffix(".json")

        else:
            # solutions chosen a priori
            (score_fn, scores), (solution_fn, solutions_a_priori) = typing.cast(
                tuple[tuple[pathlib.Path, Scores], tuple[pathlib.Path, Solutions]], k
            )

            metric = list(solutions_a_priori.points.keys())

            message = f"Calculating solutions *a priori* for system `{str(score_fn)}`"
            message += f" based on solutions from `{str(solution_fn)}`, using metrics "
            message += f"`{', '.join(metric)}`..."
            logger.info(message)

            solutions = scores.solutions_a_priori(metric, solutions_a_priori, dominated)
            solutions.metadata["prior-solution-from"] = str(solution_fn)

            output_filename = (
                output_path / f"{score_fn.stem}@{solution_fn.stem}"
            ).with_suffix(".json")

        if deduplicate:
            logger.info(f"Deduplicating solutions from {output_filename}")
            solutions = solutions.deduplicate()

        utils.save_json_with_backup(output_filename, solutions.model_dump())
        logger.info(f"Saved `{output_filename}`")
