import numpy as np

__doc__ = r"""
Calculates the Fourier transform of real-valued function on a logarithmic scale:

.. math::
    F(\omega) = \int_0^\infty f(t) \exp(i\omega t) \,dt
    
For a function f(t) sampled at :math:`N` discrete timepoints, the integrals are approximated by 

.. math::
    \integral_0^\infty f(t) \cos(\omega t) \approx \sum_{i=0}^{N-1} 
    \frac{f(t_{i+i}) - f(t_i)}{f(t_{i+i}) - f(t_i)}{t_{i+1} - t_i} 
    \frac{\cos(\omega t_{i+1}) - \cos(\omega t_i)}{\omega^2} 
    
and 

.. math::
    \integral_0^\infty f(t) \sin(\omega t) \approx \frac{f(t_0)}{\omega}
    \sum_{i=0}^{N-1} \frac{f(t_{i+i}) - f(t_i)}{f(t_{i+i}) - f(t_i)}{t_{i+1} - t_i} 
    \frac{\sin(\omega t_{i+1}) - \sin(\omega t_i)}{\omega^2}
    
More details can be found in [Blo]

[Blo] Blochowicz, Th.: Ph.D. Thesis, Universität Bayreuth (2003)

"""


def logft_cos(x, y, new_x=None):
    n = x.size
    if new_x is None:
        new_x = 2*np.pi * np.geomspace(1/np.max(x), 1/np.min(x), num=n)

    dydx = np.diff(y) / np.diff(x)

    # broadcasting faster for lower array sizes, calculation of cosine makes looping as fast
    # and memory needs are more a concern
    if n < 5000:
        wt = new_x[:, None] * x[None, :]
        np.cos(wt, out=wt)
        ret_val = np.sum(dydx * np.diff(wt, axis=1), axis=1)
    else:
        ret_val = np.zeros(n)
        for i, o in enumerate(new_x):
            cos_wt = np.cos(o * x)
            ret_val[i] = np.sum(dydx * np.diff(cos_wt))

    ret_val /= new_x**2

    return new_x, ret_val


def logft_sin(x, y, new_x=None):
    n = x.size
    if new_x is None:
        new_x = 2*np.pi * np.geomspace(1/np.max(x), 1/np.min(x), num=n)

    dydx = np.diff(y) / np.diff(x)

    # broadcasting faster for lower array sizes, calculation of cosine makes looping as fast
    # and memory needs are more a concern
    if n < 5000:
        wt = new_x[:, None] * x[None, :]
        np.sin(wt, out=wt)
        ret_val = np.sum(dydx * np.diff(wt, axis=1), axis=1)
    else:
        ret_val = np.zeros(n)
        for i, o in enumerate(new_x):
            sin_wt = np.sin(o*x)
            ret_val[i] = np.sum(dydx * np.diff(sin_wt))

    ret_val /= new_x**2
    ret_val += y[0] / new_x

    return new_x, ret_val


def logft_cmplx(x, y, new_x=None, backward: bool = False):
    n = x.size
    if new_x is None:
        new_x = 2 * np.pi * np.geomspace(1/np.max(x), 1/np.min(x), num=n)

    dydx = np.diff(y) / np.diff(x)

    sign = 2*backward - 1
    if n < 5000:
        wt = new_x[:, None] * x[None, :]
        exp_wt = np.exp(sign*1j*wt)

        ret_val = np.sum(dydx * np.diff(exp_wt, axis=1), axis=1)
    else:
        ret_val = np.zeros(n, dtype=complex)
        for i, o in enumerate(new_x):
            exp_wt = np.exp(sign*1j*o*x)
            ret_val[i] = np.sum(dydx * np.diff(exp_wt))

    ret_val /= new_x**2
    ret_val += sign * 1j*y[0]/new_x

    return new_x, ret_val


def logft(x, y, new_x=None, mode='cos'):
    if mode not in ['cos', 'sin', 'complex']:
        raise ValueError(f'Unknown mode {mode}, use "cos", "sin", "complex".')

    ft_func = {'cos': logft_cos, 'sin': logft_sin, 'complex': logft_cmplx}[mode]

    return ft_func(x, y, new_x=new_x)


def logift(x, y, new_x=None, mode='cos'):
    if mode not in ['cos', 'sin', 'complex']:
        raise ValueError(f'Unknown mode {mode}, use "cos", "sin", "complex".')

    if mode == 'complex':
        return logft_cmplx(x, y, new_x=new_x, backward=True)
    elif mode == 'cos':
        return logft_cos(x, y, new_x=new_x)
    else:
        return logft_sin(x, y, new_x=new_x)
