from __future__ import annotations

import numpy as np
import tensorflow as tf

EPS = 1e-8


def _group_logodds_grad_for_model(
    model: tf.keras.Model,
    x: np.ndarray,
    pos_label: int = 1,
    neg_label: int = 0,
) -> np.ndarray:
    """
    Compute the gradient of the positive-vs-negative log-odds for one model.
    Returns the mean gradient across samples.
    """

    x_tensor = tf.convert_to_tensor(x)

    with tf.GradientTape() as tape:
        tape.watch(x_tensor)
        preds = model(x_tensor, training=False)

        pos = preds[:, pos_label]
        neg = preds[:, neg_label]
        log_odds = tf.math.log(pos + EPS) - tf.math.log(neg + EPS)

    grads = tape.gradient(log_odds, x_tensor)
    return grads.numpy().mean(axis=0)


def compute_avg_group_logodds_gradient(
    models: list,
    x: np.ndarray,
    pos_label: int = 1,
    neg_label: int = 0,
) -> np.ndarray:
    """
    Average the log-odds gradients across an ensemble of models.
    """

    all_grads = []
    for model in models:
        g = _group_logodds_grad_for_model(
            model,
            x,
            pos_label=pos_label,
            neg_label=neg_label,
        )
        all_grads.append(g)

    return np.mean(all_grads, axis=0)


def separate_gradient_polarity(grads: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """
    Split one gradient vector into positive and negative contributions.

    Returns
    -------
    positive_grads : np.ndarray
        Positive contribution only, negative values set to 0.
    negative_grads : np.ndarray
        Magnitude of negative contribution only, positive values set to 0.
        Returned as non-negative values for direct plotting / pseudo-MS1 export.
    """

    grads = np.asarray(grads, dtype=float)
    positive_grads = np.clip(grads, a_min=0.0, a_max=None)
    negative_grads = np.clip(-grads, a_min=0.0, a_max=None)
    return positive_grads, negative_grads



def compute_avg_group_logodds_gradient_split(
    models: list,
    x: np.ndarray,
    pos_label: int = 1,
    neg_label: int = 0,
) -> dict[str, np.ndarray]:
    """
    Compute ensemble-average gradients and return signed, positive-only,
    and negative-only versions.
    """

    grads = compute_avg_group_logodds_gradient(
        models=models,
        x=x,
        pos_label=pos_label,
        neg_label=neg_label,
    )
    positive_grads, negative_grads = separate_gradient_polarity(grads)

    return {
        "signed": grads,
        "positive": positive_grads,
        "negative": negative_grads,
    }
