"""Utilitary functions for testing fast waveform accuracy."""

import numpy as np
from jax import Array

from jaxgb.params import GBObject


def overlap(A1: Array, A2: Array) -> float:
    r"""Return overlap of two signals.

    .. math::

        O(A_1|A_2) = \frac{1}{\sqrt{ N_1 N_2}} \Re \int_{f_{\min}}^{f_{\max}}
        \Sum_i \tilde{A_1^i}(f) \tilde{A_2^i}^+(f) \mathrm{d}f,

    where

    .. :math::

        N^i = \int_{f_{\min}}^{f_{\max}} \tilde{A^i}^2(f) \mathrm{d}f}

    Parameters
    ----------
    A1
        TDI channels of 1st signal (nsamples, nchannels)
    A2
        TDI channels of 2nd signal (nsamples, nchannels)

    Returns
    -------
    float
        overlap of the 2 signals (between 0 and 1)
    """
    n1 = np.sum(np.sum(np.abs(A1) ** 2, axis=0), axis=0)
    n2 = np.sum(np.sum(np.abs(A2) ** 2, axis=0), axis=0)
    olap = np.sum(np.sum(np.real(A1 * A2.conj()), axis=0), axis=0)
    return np.abs(olap) / np.sqrt(n1 * n2)


def mismatch(A1: Array, A2: Array) -> float:
    r"""Return a well behaved mismatch (from the numerical point of view)

    .. math::

        m = \frac{1}{2} \left\Vert\frac{A_2}{\Vert A_2\Vert} -
        \frac{A_1}{\Vert A_1\Vert}\right\Vert^2

    Parameters
    ----------
    A1
        TDI channels of 1st signal (nsamples, nchannels)
    A2
        TDI channels of 2nd signal (nsamples, nchannels)

    Returns
    -------
    float
        mismatch of the 2 signals
    """
    n2 = A2 / np.sum(A2 * A2) ** 0.5
    n1 = A1 / np.sum(A1 * A1) ** 0.5
    return 0.5 * np.sum((n2 - n1) ** 2)


def get_source(
    name: str,
    as_dict: bool = False,
    as_gwresponse_dict: bool = False,
    as_gbgpu: bool = False,
    as_fastgb: bool = False,
) -> Array | dict:
    """Return source parameters from name.

    Parameters
    ----------
    name
        name of a known source (like `highfreq` for a high SNR, high frequency source).
    as_dict
        return parameters as a dictionary (LDC names)
    as_gwresponse_dict
        return parameters as a dictionary (lisagwresponse names)
    as_gbgpu
        return parameters as an array in the order expected by GBGPU
    as_fastgb
        return parameters as an array in the order expected by fastgb

    Returns
    -------
    array or dict of parameters of the chosen source.
    """
    gbobject = None
    if name == "highfreq":
        gbobject = GBObject(
            f0=0.018038,  # f0 Hz
            fdot=1.211329e-13,  # fdot "Hz^2
            A=2.765173e-22,  # ampl strain
            ra=4.22546,  # equatorial longitude (ra) radian
            dec=-0.823607,  # equatorial latitude (dec) radian
            psi=2.426183,  # polarization radian
            iota=1.052777,  # inclination radian
            phi0=3.020977,  # initial phase radian
            t_init=0.0,  # epoch in seconds
        )  # snr 130
    if as_dict:
        return gbobject.to_ldc_dict()
    if as_gwresponse_dict:
        return gbobject.to_gwresponse_dict()
    if as_gbgpu:
        return gbobject.to_gbgpu_array()
    if as_fastgb:
        return gbobject.to_fastgb_array()
    return gbobject.to_jaxgb_array(t0=0.0)
