"""Stochastic smoothing for score matching validity."""

from __future__ import annotations

import numpy as np

from babappa.utils.random import deterministic_seed


def smooth_matrix(matrix: np.ndarray, bandwidth: float, seed: int) -> np.ndarray:
    if bandwidth < 0:
        raise ValueError("Smoothing bandwidth must be non-negative")
    if bandwidth == 0:
        return matrix.copy()
    rng = np.random.default_rng(seed)
    noise = rng.normal(loc=0.0, scale=bandwidth, size=matrix.shape)
    return matrix + noise


def smooth_matrix_with_site_seed_schedule(
    matrix: np.ndarray,
    bandwidth: float,
    base_seed: int,
    *parts: object,
) -> np.ndarray:
    """Apply Gaussian smoothing with a deterministic per-site seed schedule."""
    if bandwidth < 0:
        raise ValueError("Smoothing bandwidth must be non-negative")
    if matrix.ndim != 2:
        raise ValueError("Matrix must be 2D [n_sites, dim]")
    out = matrix.copy()
    if bandwidth == 0:
        return out

    for site_idx in range(out.shape[0]):
        site_seed = deterministic_seed(base_seed, *parts, site_idx + 1)
        rng = np.random.default_rng(site_seed)
        out[site_idx, :] = out[site_idx, :] + rng.normal(
            loc=0.0,
            scale=bandwidth,
            size=out.shape[1],
        )
    return out
