'''Create multi-band AWGN.

Mike Markowski, mike.ab3ap@gmail.com
Mar 2021
'''

import numpy as np

def polar2iq(r, theta):
    return r * np.exp(1j * theta)

def iq2polar(z):
    return (np.abs(z), np.angle(z))

def bands(fc_Hz, bw_Hz, fs_Hz, n):
    '''Make bands of noise whose span is centered at 0 Hz.

    XXX Maybe better to use normalized parameters??

    Inputs:
      fc_Hz (float) : center frequencies of AWGN bands, in Hz.
      bw_Hz (float) : bandwidths corresponding to fc_Hz of AWGN, in Hz.
      fs_Hz (float) : sample rate, in Hz.
      n (int) : number of samples to generate.

    Output:
      (complex[]) : AWGN array of I+jQ of length n.
    '''

    # Normalize parameter values.
    fc_norm = np.array(fc_Hz) / fs_Hz
    bw_norm = np.array(bw_Hz) / fs_Hz

    # Slide bands to start at zero.
    leftmost_norm = min(fc_norm - bw_norm / 2) # Leftmost of all AWGN bands.
    fc_norm -= leftmost_norm # Slide leftmost edge down to 0.
    span_norm = max(fc_norm + bw_norm / 2) # Rightmost edge of all AWGN bands.
    left_norm = fc_norm - bw_norm / 2 # Left edges of bands.
    right_norm = left_norm + bw_norm # Right edges of bands.

    # Convert normalized units to array indexes.
    spanInd = int(np.round(span_norm * n)) # Scalar.
    s = list(map(int, np.round(left_norm * n))) # Start indexes of bands.
    e = list(map(int, np.round(right_norm * n))) # End indexes of bands.

    # Create bands of AWGN as polar mag /_ phase.
    mag = np.zeros(n)
    phase = np.zeros(n)
    for band in range(len(fc_norm)): # Create AWGN band by band.
        si, ei = s[band], e[band] # Indexes for this band.
        mag[si:ei] = n # Constant magnitude and random phase.
        phase[si:ei] = np.pi * (2 * np.random.rand(ei-si) - 1)

    # Convert to i/q.
    awgn_iq = polar2iq(mag, phase) # Polar to Cartesian.
    awgn_iq = np.roll(awgn_iq, -spanInd//2) # Center bands at 0 Hz.
    return np.fft.ifft(awgn_iq) # Time domain.