from __future__ import annotations

from typing import Any

from .explain import compute_avg_group_logodds_gradient_split
from .plotting import gradient_to_pseudo_ms1, plot_pseudo_ms1, save_pseudo_ms1_csv
from .training import run_training


def run_proxai(
    csv_path: str,
    label_column: str,
    bin_column: str,
    bin_values: list[int | float],
    control_labels: list[Any],
    experiment_labels: list[Any],
    mz_start: float = 200.0,
    bin_size: float = 0.1,
    out_prefix: str = "proxai_output",
):
    """
    Full ProXAI pipeline for one or more bins.

    For each bin, trains a binary control-vs-experiment classifier using the
    provided label lists, then exports only positive and negative pseudo-MS1 outputs.
    """

    results_by_bin: dict[Any, dict[str, Any]] = {}

    for bin_value in bin_values:
        result = run_training(
            csv_path=csv_path,
            label_column=label_column,
            bin_column=bin_column,
            bin_value=bin_value,
            control_labels=control_labels,
            experiment_labels=experiment_labels,
        )

        grad_parts = compute_avg_group_logodds_gradient_split(
            result["models_a"],
            result["x"],
        )

        pseudo_positive = gradient_to_pseudo_ms1(
            grad_parts["positive"],
            mz_start=mz_start,
            bin_size=bin_size,
            take_abs=False,
            normalize=True,
        )
        pseudo_negative = gradient_to_pseudo_ms1(
            grad_parts["negative"],
            mz_start=mz_start,
            bin_size=bin_size,
            take_abs=False,
            normalize=True,
        )

        bin_prefix = f"{out_prefix}_bin_{bin_value}"
        positive_csv = f"{bin_prefix}_positive.csv"
        negative_csv = f"{bin_prefix}_negative.csv"
        positive_png = f"{bin_prefix}_positive.png"
        negative_png = f"{bin_prefix}_negative.png"

        save_pseudo_ms1_csv(pseudo_positive, positive_csv)
        save_pseudo_ms1_csv(pseudo_negative, negative_csv)

        plot_pseudo_ms1(
            pseudo_positive,
            title=f"ProXAI positive-gradient pseudo-MS1 (bin {bin_value})",
            out_png=positive_png,
        )
        plot_pseudo_ms1(
            pseudo_negative,
            title=f"ProXAI negative-gradient pseudo-MS1 (bin {bin_value})",
            out_png=negative_png,
        )

        results_by_bin[bin_value] = {
            "gradients": {
                "positive": grad_parts["positive"],
                "negative": grad_parts["negative"],
            },
            "pseudo_ms1": {
                "positive": pseudo_positive,
                "negative": pseudo_negative,
            },
            "csv": {
                "positive": positive_csv,
                "negative": negative_csv,
            },
            "png": {
                "positive": positive_png,
                "negative": negative_png,
            },
        }

    return {
        "bins": results_by_bin,
        "control_labels": control_labels,
        "experiment_labels": experiment_labels,
    }
