#!/usr/bin/env python3

import sys
import numpy as np
import matplotlib.pyplot as plt

# Parameters
baseband_freq    =  10000  # Center freq. of detector.
fs               = 250000
detector_bw      =  20000
integration_time = 1

# Internal stuff
n = int(fs * integration_time)  # 
target_bin_low  = int(n * (baseband_freq-detector_bw/2)/fs)
target_bin_high = int(n * (baseband_freq+detector_bw/2)/fs)
if target_bin_low < 0:
    # We don't allow this because we don't want to rearrange the frequency bins and
    # doing this right for detectors that overlap with baseband DC will be annoyingly tricky.
    print('Detector must be fully contained in the positive half of the baseband spectrum. Aborting.')
    sys.exit(1)

FFT_WINDOW = np.ones(n) # Use this for rectangular FFT window
#FFT_WINDOW  = np.hamming(n)
FFT_WINDOW /= np.sum(FFT_WINDOW) # Normalize

def get_carrier_iq(t, freq):
    return  np.exp(2j*np.pi*freq*t)

def get_fm_iq(t, carrier_freq, modulation_freq):
    # First we need a carrier at the target frequency
    carrier = get_carrier_iq(t, baseband_freq)
    # Then we generate some AF, in this case just a tone
    modulation = np.sin(2*np.pi*modulation_freq*t)
    # We would have to integrate this to a cumulative phase angle but since
    # sinusoids integrate in a circular fashion, we can skip this here. In
    # particular we don't really care about the precise modulation index
    # Instead, we modulate around 0 Hz
    (baseband_i, baseband_q) = (np.cos(modulation), np.sin(modulation))
    baseband_iq = baseband_i + 1j*baseband_q
    # Finally we shift this to the target carrier frequency and return
    return baseband_iq*carrier

t = np.linspace(0, integration_time, num=n, endpoint=False)
for amplitude in np.logspace(0, 5, 6):
    signal = amplitude**0.5 * get_carrier_iq(t, baseband_freq)
    #signal = amplitude**0.5 * get_fm_iq(t, baseband_freq, 1e3)

    # There is nicer ways to compute this but this form is the closest to the canonical RMS computation
    power_time_domain = 1./n * np.sum( np.abs(signal)**2 )

    spectrum = np.fft.fft(signal*FFT_WINDOW)                           # First we compute the FFT (complex)
    spectrum_windowed = spectrum[target_bin_low:target_bin_high]       # Then we extract the bins that are within the detector
    #power_bins = spectrum_windowed.real**2 + spectrum_windowed.imag**2 # We compute per-bin power (P = U**2 / R)
    voltage_bins = np.abs(spectrum_windowed)
    power_bins   = voltage_bins**2
    power_db   = 10*np.log10(np.sum(power_bins))
    voltage_db = 20*np.log10(np.sum(voltage_bins))
    print(10*np.log10(power_time_domain), power_db, voltage_db)