from __future__ import annotations

import matplotlib
import numpy as np
import pandas as pd

matplotlib.use("Agg")
import matplotlib.pyplot as plt


def normalize_signal(x: np.ndarray, mode: str = "maxabs") -> np.ndarray:
    x = np.asarray(x, dtype=float)

    if mode == "maxabs":
        denom = np.max(np.abs(x))
        return x / denom if denom > 0 else x

    if mode == "minmax":
        lo, hi = np.min(x), np.max(x)
        return (x - lo) / (hi - lo) if hi > lo else np.zeros_like(x)

    raise ValueError(f"Unknown normalization mode: {mode}")


def gradient_to_pseudo_ms1(
    grads: np.ndarray,
    mz_start: float = 200.0,
    bin_size: float = 0.1,
    take_abs: bool = True,
    normalize: bool = False,
) -> pd.DataFrame:
    grads = np.asarray(grads, dtype=float)

    signal = np.abs(grads) if take_abs else grads.copy()
    if normalize:
        signal = normalize_signal(signal, mode="maxabs")

    mz = mz_start + np.arange(len(signal)) * bin_size

    return pd.DataFrame(
        {
            "mz": mz,
            "intensity": signal,
        }
    )


def save_pseudo_ms1_csv(
    pseudo_df: pd.DataFrame,
    out_csv: str,
) -> None:
    pseudo_df.to_csv(out_csv, index=False)


def plot_pseudo_ms1(
    pseudo_df: pd.DataFrame,
    title: str = "Pseudo-MS1 from ProXAI Saliency",
    out_png: str | None = None,
    figsize: tuple[int, int] = (14, 5),
    linewidth: float = 0.8,
) -> None:
    plt.figure(figsize=figsize)
    plt.plot(pseudo_df["mz"], pseudo_df["intensity"], linewidth=linewidth)
    plt.xlabel("m/z")
    plt.ylabel("Intensity")
    plt.title(title)
    plt.tight_layout()

    if out_png:
        plt.savefig(out_png, dpi=200, bbox_inches="tight")

    plt.close()