#
# This file is part of pyS5p
#
# https://github.com/rmvanhees/pys5p.git
#
# Copyright (c) 2017-2025 SRON
#   All Rights Reserved
#
# License:  BSD-3-Clause
"""Routines to divide or add (partial) uncertainties."""

from __future__ import annotations

__all__ = ["unc_div", "unc_sum"]

import numpy as np


def unc_div(
    value_a: np.ndarray, sigma_a: np.ndarray, value_b: np.ndarray, sigma_b: np.ndarray
) -> np.ndarray:
    r"""Absolute error for parameter `a` divided by `b`.

    .. math:: (a / b) * \sqrt{(\sigma_a / a)^2 + (\sigma_b / b)^2}
    """
    if not value_a.shape == value_b.shape == sigma_a.shape == sigma_b.shape:
        raise TypeError("dimensions of input arrays are not the same")

    mask = (
        np.isfinite(value_a)
        & np.isfinite(sigma_a)
        & np.isfinite(value_b)
        & np.isfinite(sigma_b)
    )

    if np.sum(mask) == sigma_a.size:
        return (value_a / value_b) * np.sqrt(
            (sigma_a / value_a) ** 2 + (sigma_b / value_b) ** 2
        )

    res = np.full(sigma_a.shape, np.nan)
    res[mask] = (value_a[mask] / value_b[mask]) * np.sqrt(
        (sigma_a[mask] / value_a[mask]) ** 2 + (sigma_b[mask] / value_b[mask]) ** 2
    )
    return res


def unc_sum(sigma_a: np.ndarray, sigma_b: np.ndarray) -> np.ndarray:
    r"""Absolute error for the sum of the parameters `a` and `b`.

    .. math:: \sqrt{\sigma_a^2 + \sigma_b^2}
    """
    if sigma_a.shape != sigma_b.shape:
        raise TypeError("dimensions of sigma_a and sigma are not the same")

    mask = np.isfinite(sigma_a) & np.isfinite(sigma_b)

    if np.sum(mask) == sigma_a.size:
        return np.sqrt(sigma_a**2 + sigma_b**2)

    res = np.full(sigma_a.shape, np.nan)
    res[mask] = np.sqrt(sigma_a[mask] ** 2 + sigma_b[mask] ** 2)

    return res
