import torch


def torch_hilbert(x_real, n_fft=None):
    """
    Obtain imaginary counterpart to a real signal such that there are no negative frequency
    components when represented as a complex signal. This is done by using the Hilbert transform.
    We end up with an analytic signal and return only the imaginary part. Most importantly,
    this procedure is fully differentiable. Adapted from the SciPy signal.hilbert function.

    Parameters
    ----------
    x_real : Tensor (F x T)
      Real counterpart of an analytic signal,
      F - number of independent signals
      T - number of time steps (samples)
    n_fft : int
      Number of Fourier components

    Returns
    ----------
    x_imag : Tensor (F x T)
      Imaginary counterpart of an analytic signal,
      F - number of independent signals
      T - number of time steps (samples)
    """

    # Default to the length of the input signal
    if n_fft is None:
        n_fft = x_real.size(-1)

    # Create the transfer function for an analytic signal
    h = torch.zeros(n_fft).to(x_real.device)
    if n_fft % 2 == 0:
        h[0] = h[n_fft // 2] = 1
        h[1 : n_fft // 2] = 2
    else:
        h[0] = 1
        h[1 : (n_fft + 1) // 2] = 2

    # Take the Fourier transform of the real part
    Xf = torch.fft.fft(x_real, n=n_fft, dim=-1)
    # Apply the transfer function to the Fourier transform
    Xfh = Xf * h.unsqueeze(-2)
    # Take the inverse Fourier Transform to obtain the analytic signal
    x_alyt = torch.fft.ifft(Xfh, dim=-1)
    # Take the imaginary part of the analytic signal to obtain the Hilbert transform
    x_imag = x_alyt.imag

    return x_imag