import cmath
import numpy as np

# When playing noise against a captured signal, it is often convenient
# to create the noise by specifying number of samples needed and sample
# rate of target signal.  The function awgn() provides this.
#
# Mike Markowski, mike.ab3ap@gmail.com
# Mar 21, 2015

# cart2pol
#
# Convert arrays of cartesian coordinates to arrays of polar coordinates.

def cart2pol(x, y):
    rho = np.hypot(x, y)
    phi = np.arctan2(y, x)
    return phi, rho

# pol2cart
#
# Convert arrays of polar coordinates to arrays of Cartesian coordinates.

def pol2cart(rho, phi):
    x = rho * np.cos(phi)
    y = rho * np.sin(phi)
    return x, y

# Generate multiband Gaussian noise waveform based on number of samples
# needed.
#
# Create white Gaussian noise.  The caller specifies a set of center
# frequencies along with corresponding noise bandwidths.  The subroutine
# creates a multiband noise curve following the specifications.
#
# Based on single band noise Matlab code by Vikas Singh.
#
# Inputs:
#  centers_Hz: scalar, list, or array of center frequencies in units of Hz.
#  widths_Hz: noise bandwidths corresponding to each center frequency.
#  sampleRate_Hz: sample rate in Hz of the noise signal.
#  numSamples: number of samples in the generated noise curve.
#
# Output:
#  Time domain multiband noise curve as specified by inputs.

def awgn(centers_Hz, widths_Hz, sampleRate_Hz, numSamples):

    if type(centers_Hz) != np.ndarray:
        centers_Hz = np.array(centers_Hz)
    if type(widths_Hz) != np.ndarray:
        widths_Hz = np.array(widths_Hz)

    # fft size == numSamples
    toneGap_Hz = float(sampleRate_Hz) / numSamples # FFT frequency bin step.

    halfWidth_index = np.ceil((widths_Hz / toneGap_Hz) / 2.)
    centers_index = centers_Hz / toneGap_Hz
    ranges_index = np.column_stack([
        np.floor(centers_index - halfWidth_index),
        np.ceil(centers_index + halfWidth_index)])

    noiseTones = np.zeros(numSamples)
    for range_index in ranges_index:
        lo = range_index[0]
        hi = range_index[1]
        if lo < 0:
            # Wrap noise tones, negative to right side.
            # 1 ---+          +---
            #      |          | wrapped
            # 0    +----------+
            noiseTones[0 : hi+1] = 1
            noiseTones[numSamples+lo : numSamples] = 1
        elif hi >= numSamples:
            # Wrap noise tones, too high to left side.
            #    1 ---+          +---
            # wrapped |          |
            #    0    +----------+
            noiseTones[0 : numSamples-hi] = 1
            noiseTones[lo : numSamples] = 1
        else:
            # Make band.
            # 1     +------+
            #       |      |
            # 0 ----+      +----
            noiseTones[lo : hi+1] = 1

    # Random phase noise, -pi to pi radians, for each noise tone.
    randomPhase = np.pi * (2 * np.random.rand(numSamples) - 1)

    # Convert polar mag/phase to Cartesian.
    x, y = pol2cart(noiseTones, randomPhase)
    c = np.vectorize(complex)(x, y)
    # Move to time domain.
    timeDomSig = np.fft.ifft(c)

#   If IQ needed to feed to ARB, use following line
#   IQdata = real(timeDomSig), imag(timeDomSig)
#   return IQdata

    return timeDomSig