"""Helpers for persisted FFT background profiles."""

from __future__ import annotations

from pathlib import Path

import numpy as np


def validate_fft_background(background: np.ndarray) -> np.ndarray:
    """Validate a saved FFT background payload."""
    values = np.asarray(background)
    if values.ndim != 1:
        raise ValueError("FFT background must be a 1D array")
    if not np.issubdtype(values.dtype, np.number):
        raise ValueError("FFT background must be numeric")
    values = np.asarray(values, dtype=np.float32).reshape(-1)
    if values.size == 0:
        raise ValueError("FFT background is empty")
    return values


def _normalize_background_path(path: str | Path) -> Path:
    out = Path(path).expanduser()
    if out.suffix.lower() != ".npy":
        out = out.with_suffix(".npy")
    return out


def save_fft_background(path: str | Path, background: np.ndarray) -> str:
    """Persist an FFT background profile as a .npy file."""
    normalized_path = _normalize_background_path(path)
    values = validate_fft_background(background)
    np.save(normalized_path, values.astype(np.float32, copy=False))
    return str(normalized_path)


def load_fft_background(path: str | Path) -> np.ndarray:
    """Load and validate an FFT background profile from a .npy file."""
    normalized_path = _normalize_background_path(path)
    loaded = np.load(normalized_path, allow_pickle=False)
    return validate_fft_background(loaded)


def subtract_fft_background(signal_mag: np.ndarray, background_mag: np.ndarray) -> np.ndarray:
    """Subtract a background profile from FFT magnitudes in linear amplitude."""
    signal = np.asarray(signal_mag, dtype=np.float32)
    background = validate_fft_background(background_mag)
    if signal.ndim == 1:
        if signal.size != background.size:
            raise ValueError("FFT background size does not match signal size")
        valid = np.isfinite(signal) & np.isfinite(background)
        out = np.full_like(signal, np.nan, dtype=np.float32)
        if np.any(valid):
            out[valid] = np.maximum(signal[valid] - background[valid], 0.0)
        return out

    if signal.ndim == 2:
        if signal.shape[0] != background.size:
            raise ValueError("FFT background size does not match signal rows")
        background_2d = background[:, None]
        valid = np.isfinite(signal) & np.isfinite(background_2d)
        diff = signal - background_2d
        return np.where(valid, np.maximum(diff, 0.0), np.nan).astype(np.float32, copy=False)

    raise ValueError("FFT background subtraction supports only 1D or 2D signals")