import numpy as np

try:
    import cupy as cp

except (ModuleNotFoundError, ImportError):
    pass

import warnings
from scipy.optimize import brentq
from few.utils.utility import get_separatrix
from scipy.constants import year

def cornish_lisa_psd(f, sky_averaged=False, use_gpu=False):
    """PSD from https://arxiv.org/pdf/1803.01944.pdf

    Power Spectral Density for the LISA detector assuming it has been active for a year.
    I found an analytic version in one of Niel Cornish's paper which he submitted to the arXiv in
    2018. I evaluate the PSD at the frequency bins found in the signal FFT.

    PSD obtained from: https://arxiv.org/pdf/1803.01944.pdf

    """

    if use_gpu:
        xp = cp
    else:
        xp = np

    if sky_averaged:
        sky_averaging_constant = 20 / 3

    else:
        sky_averaging_constant = 1.0  # set to one for one source

    L = 2.5 * 10 ** 9  # Length of LISA arm
    f0 = 19.09 * 10 ** (-3)  # transfer frequency

    # Optical Metrology Sensor
    Poms = ((1.5e-11) * (1.5e-11)) * (1 + xp.power((2e-3) / f, 4))

    # Acceleration Noise
    Pacc = (
        (3e-15)
        * (3e-15)
        * (1 + (4e-4 / f) * (4e-4 / f))
        * (1 + xp.power(f / (8e-3), 4))
    )

    # constants for Galactic background after 1 year of observation
    alpha = 0.171
    beta = 292
    k = 1020
    gamma = 1680
    f_k = 0.00215

    # Galactic background contribution
    Sc = (
        9e-45
        * xp.power(f, -7 / 3)
        * xp.exp(-xp.power(f, alpha) + beta * f * xp.sin(k * f))
        * (1 + xp.tanh(gamma * (f_k - f)))
    )

    # PSD
    PSD = (sky_averaging_constant) * (
        (10 / (3 * L * L))
        * (Poms + (4 * Pacc) / (xp.power(2 * np.pi * f, 4)))
        * (1 + 0.6 * (f / f0) * (f / f0))
        + Sc
    )

    return PSD


def get_sensitivity(f, sens_fn="lisasens", return_type="PSD", *args, **kwargs):
    """Generic sensitivity generator

    Same interface to many sensitivity curves.

    Args:
        f (1D double np.ndarray): Array containing frequency  values.
        sens_fn (str, optional): String that represents the name of the desired
            SNR function. Options are "cornish_lisa_psd" or any sensitivity
            function found in tdi.py from the MLDC gitlab. Default is the
            LISA sensitivity from the tdi.py.
        return_type (str, optional): Described the desired output. Choices are ASD,
            PSD, or char_strain (characteristic strain). Default is ASD.

        *args (list or tuple, optional): Any additional arguments for the sensitivity function.
        **kwargs (dict, optional): Keyword arguments to pass to sensitivity function.

    """

    try:
        sensitivity = globals()[sens_fn]
    except KeyError:
        raise ValueError("{} sensitivity is not available.".format(sens_fn))

    PSD = sensitivity(f, *args, **kwargs)

    if return_type == "PSD":
        return PSD

    elif return_type == "ASD":
        return PSD ** (1 / 2)

    elif return_type == "char_strain":
        return (f * PSD) ** (1 / 2)

    else:
        raise ValueError("return_type must be PSD, ASD, or char_strain.")


def inner_product(
        sig1,
        sig2,
        dt=None,
        df=None,
        f_arr=None,
        PSD="lisasens",
        PSD_args=(),
        PSD_kwargs={},
        normalize=False,
        use_gpu=False,
):
    if use_gpu:
        xp = cp

    else:
        xp = np

    if df is None and dt is None and f_arr is None:
        raise ValueError("Must provide either df, dt or f_arr keyword arguments.")

    if isinstance(sig1, list) is False:
        sig1 = [sig1]

    if isinstance(sig2, list) is False:
        sig2 = [sig2]

    if len(sig1) != len(sig2):
        raise ValueError(
            "Signal 1 has {} channels. Signal 2 has {} channels. Must be equal.".format(
                len(sig1), len(sig2)
            )
        )

    if dt is not None:

        if len(sig1[0]) != len(sig2[0]):
            warnings.warn(
                "The two signals are two different lengths in the time domain. Zero padding smaller array."
            )

            length = len(sig1[0]) if len(sig1[0]) > len(sig2[0]) else len(sig2[0])

            sig1 = [xp.pad(sig, (0, length - len(sig1[0]))) for sig in sig1]
            sig2 = [xp.pad(sig, (0, length - len(sig2[0]))) for sig in sig2]

        length = len(sig1[0])

        freqs = xp.fft.rfftfreq(length, dt)[1:]

        ft_sig1 = [
            xp.fft.rfft(sig)[1:] * dt for sig in sig1
        ]  # remove DC / dt factor helps to cast to proper dimensionality
        ft_sig2 = [xp.fft.rfft(sig)[1:] * dt for sig in sig2]  # remove DC

    else:
        ft_sig1 = sig1
        ft_sig2 = sig2

        if df is not None:
            freqs = (xp.arange(len(sig1[0])) + 1) * df  # ignores DC component using + 1

        else:
            freqs = f_arr

    # get PSD weighting
    if isinstance(PSD, str):
        PSD_arr = get_sensitivity(freqs, sens_fn=PSD, *PSD_args, **PSD_kwargs)

    elif isinstance(PSD, xp.ndarray):
        PSD_arr = PSD

    elif PSD is None:
        PSD_arr = 1.0

    else:
        raise ValueError(
            "PSD must be a string giving the sens_fn or a predetermimed array or None if noise weighting is included in a signal."
        )

    out = 0.0
    # assumes right summation rule
    x_vals = xp.zeros(len(PSD_arr))

    x_vals[1:] = xp.diff(freqs)
    x_vals[0] = x_vals[1]

    # account for hp and hx if included in time domain signal
    for temp1, temp2 in zip(ft_sig1, ft_sig2):
        y = xp.real(temp1.conj() * temp2) / PSD_arr  # assumes right summation rule

        out += 4 * xp.sum(x_vals * y)

    normalization_value = 1.0
    if normalize is True:
        norm1 = inner_product(
            sig1,
            sig1,
            dt=dt,
            df=df,
            f_arr=f_arr,
            PSD=PSD,
            PSD_args=PSD_args,
            PSD_kwargs=PSD_kwargs,
            use_gpu=use_gpu,
            normalize=False,
        )
        norm2 = inner_product(
            sig2,
            sig2,
            dt=dt,
            df=df,
            f_arr=f_arr,
            PSD=PSD,
            PSD_args=PSD_args,
            PSD_kwargs=PSD_kwargs,
            use_gpu=use_gpu,
            normalize=False,
        )

        normalization_value = np.sqrt(norm1 * norm2)

    elif isinstance(normalize, str):
        if normalize == "sig1":

            sig_to_normalize = sig1

        elif normalize == "sig2":
            sig_to_normalize = sig2

        else:
            raise ValueError(
                "If normalizing with respect to sig1 or sig2, normalize kwarg must either be 'sig1' or 'sig2'."
            )

        normalization_value = inner_product(
            sig_to_normalize,
            sig_to_normalize,
            dt=dt,
            df=df,
            f_arr=f_arr,
            PSD=PSD,
            PSD_args=PSD_args,
            PSD_kwargs=PSD_kwargs,
            normalize=False,
        )

    elif normalize is not False:
        raise ValueError("Normalize must be True, False, 'sig1', or 'sig2'.")

    out /= normalization_value
    return out


def snr(sig1, *args, data=None, use_gpu=False, **kwargs):
    if use_gpu:
        xp = cp

    else:
        xp = np

    if data is None:
        sig2 = sig1

    else:
        sig2 = data

    return xp.sqrt(inner_product(sig1, sig2, *args, use_gpu=use_gpu, **kwargs))


def brentq_p_at_t(
        traj_module,
        t_out,
        traj_args,
        index_of_p=3,
        index_of_a=2,
        index_of_e=3,
        index_of_Y=4,
        traj_kwargs={},
        kerr_separatrix=True,
        xtol=2e-12,
        rtol=8.881784197001252e-16,
        bounds=None,
):
    """Find the value of p that will give a specific length inspiral using Brent's method.

    If you want to generate an inspiral that is a specific length, you
    can adjust p accordingly. This function tells you what that value of p
    is based on the trajectory module and other input parameters at a
    desired time of observation.

    This function uses scipy's brentq routine to find the (presumed only)
    value of p that gives a trajectory of duration t_out.

    arguments:
        traj_module (obj): Instantiated trajectory module. It must output
            the time array of the trajectory sparse trajectory as the first
            output value in the tuple.
        t_out (double): The desired length of time for the waveform.
        traj_args (list): List of arguments for the trajectory function.
            p is removed. **Note**: It must be a list, not a tuple because the
            new p values are inserted into the argument list.
        index_of_p (int, optional): Index where to insert the new p values in
            the :code:`traj_args` list. Default is 3.
        index_of_a (int, optional): Index of a in provided :code: `traj_args`. Default is 2.
        index_of_e (int, optional): Index of e0 in provided :code: `traj_args`. Default is 3.
        index_of_Y (int, optional): Index of Y0 in provided :code: `traj_args`. Default is 4.
        traj_kwargs (dict, optional): Keyword arguments for :code:`traj_module`.
            Default is an empty dict.
        kerr_separatrix (bool, optional): If True, the default lower bound of the root-finding will be
            (approximately) the Kerr separatrix. If False, uses the Schwarzchild separatrix 6 + 2e.
            Defaults to True.
        xtol (float, optional): Absolute tolerance of the brentq root-finding - see :code: `np.allclose()` for details.
            Defaults to 2e-12 (scipy default).
        rtol (float, optional): Relative tolerance of the brentq root-finding - see :code: `np.allclose()` for details.
            Defaults to ~8.8e-16 (scipy default).
        bounds (list, optional): Minimum and maximum values of p over which brentq will search for a root.
            If not given, will be set to [separatrix + 0.101, 50]. To supply only one of these two limits, set the
            other limit to None.
    returns:
        double: Value of p that creates the proper length trajectory.

    """

    if bounds is None:
        if kerr_separatrix:
            p_sep = get_separatrix(traj_args[index_of_a],traj_args[index_of_e],traj_args[index_of_Y]) # should be fairly close.
        else:
            p_sep = 6 + 2*traj_args[index_of_e]
        bounds = [p_sep + 0.101, 50]
    elif bounds[0] is None:
        if kerr_separatrix:
            p_sep = get_separatrix(traj_args[index_of_a],traj_args[index_of_e],traj_args[index_of_Y]) # should be fairly close.
        else:
            p_sep = 6 + 2*traj_args[index_of_e]
        bounds[0] = p_sep + 0.101
    elif bounds[1] is None:
        bounds[1] = 50

    def p_time_root(pval, traj, inj_args, t_outp, p_ind):
        """
        Function with one p root at T = t_outp, for brentq input.
        """
        inputs = inj_args.copy()
        inputs.insert(p_ind, pval)
        out = traj(*inputs, T=t_outp * 2, **traj_kwargs)
        return out[0][-1] - t_outp * year
    root = brentq(p_time_root, bounds[0], bounds[1], xtol=xtol, rtol=rtol, args=(traj_module, traj_args, t_out, index_of_p))
    return root
