# -*- coding: utf-8 -*-
"""Weight 1 harmonic sums."""

import numba as nb
import numpy as np

from .polygamma import cern_polygamma


@nb.njit(cache=True)
def S1(N):
    r"""Computes the harmonic sum :math:`S_1(N)`.

    .. math::
      S_1(N) = \sum\limits_{j=1}^N \frac 1 j = \psi_0(N+1)+\gamma_E

    with :math:`\psi_0(N)` the digamma function and :math:`\gamma_E` the
    Euler-Mascheroni constant.

    Parameters
    ----------
        N : complex
            Mellin moment

    Returns
    -------
        S_1 : complex
            (simple) Harmonic sum :math:`S_1(N)`

    See Also
    --------
        eko.harmonics.polygamma.cern_polygamma : :math:`\psi_k(N)`
    """
    return cern_polygamma(N + 1.0, 0) + np.euler_gamma


@nb.njit(cache=True)
def Sm1(N, hS1, is_singlet=None):
    r"""Analytic continuation of harmonic sum :math:`S_{-1}(N)`.

    .. math::
      S_{-1}(N) = \sum\limits_{j=1}^N \frac {(-1)^j} j

    Parameters
    ----------
        N : complex
            Mellin moment
        hS1:  complex
            Harmonic sum :math:`S_{1}(N)`
        is_singlet: bool, None
            symmetry factor: True for singlet like quantities (:math:`\eta=(-1)^N = 1`),
            False for non-singlet like quantities (:math:`\eta=(-1)^N=-1`)
    Returns
    -------
        Sm1 : complex
            Harmonic sum :math:`S_{-1}(N)`

    See Also
    --------
        eko.anomalous_dimension.w1.S1 : :math:`S_1(N)`
    """
    if is_singlet is None:
        return (
            (1 - (-1) ** N) / 2 * S1((N - 1) / 2)
            + ((-1) ** N + 1) / 2 * S1(N / 2)
            - hS1
        )
    if is_singlet:
        return S1(N / 2) - hS1
    return S1((N - 1) / 2) - hS1
