# -*- coding: utf-8 -*-
"""
Created on Sat Apr 24 22:32:26 2021

@author: Giacomo Ciani
"""

# This script, ispired by an earlier version by Michal Was, computes and
# displays the beat notes between different modulation sidebands up to the
# specified order, ignoring beatnotes between sidebands belonging to different
# carriers (ITF, green, SC, etc...)
#
# It then highlights:
# - the distributin of beatnotes in frequecy
# - any beatnote falling into the ITF band (i.e. <= 2 kHz)
# - any beatnote spaced less than a specified amount from a sideband frequency
#   or harmoinc up to a specified order

from itertools import product
import numpy as np
import matplotlib.pyplot as plt
#%% modulatin frequency and analysis parameters
# Modulation frequencies spcificed as a list of lists
# Info can be found in
# /virgoData/VirgoOnline/DAQ_requested_modulation_frequencies.cfg

modfreqs = [
    np.array([6270783, 8361045, 56436993, 4006184.1]), # sb of ITF carrier
    np.array([8245000.0, 11120000.0]), # sb of SC
    np.array([5529427.0])  # sb of green beam
    ]

sbgroups = [ # Labels for the sb groups above
    'ITF carrier',
    'SQZ Subcarrier',
    'Green'
    ]

colors = ['r','b','g'] # color to use to represent groups

allsbs = [freq for flist in modfreqs for freq in flist] # flattened list

maxOrder = 3 # max sideband harmonic to consider for beatnote generation
maxTotOrder = 8 # max total beatnote order to consider

hoftmax = 2e3; # minimum sideband frequency to be safe for hoft

maxOrderDem = 2 # max sideband harmonic to consider for checking nearby beatnotes
minSpacing = 50e3 #minimum spacing between beatnotes and sidebands considered safe

#%% Compute all possible beatnotes up to the specified order
orders = [] # each element in this list is an array containing all possible sb order combinations of on group that don't exceed the total order
beatnotes = []; # each element in this list is an array containing the beatnote frequency resulting from corresponding combination of sb in combs
for sbs in modfreqs:
    combs = np.array(list(product(list(range(-maxOrder,maxOrder+1)), repeat = len(sbs))))
    combs = combs[sum(abs(combs.transpose()))<=maxTotOrder]
    freqs = combs.dot(sbs)
    orders.append(combs)
    beatnotes.append(freqs)

#%% Plot beatnotes frequencies and order
maxf = maxOrderDem*max(allsbs)+2*minSpacing #only plots up to the highest considered demodulation frequency plus some margin
plt.figure()

for x, c in zip(modfreqs, colors): # plot demodulation frequencies colored according to group
    for f in x:
        for n in np.arange(maxOrderDem):
            plt.plot([(n+1)*f,(n+1)*f], [0, maxTotOrder+1], '--'+c, linewidth=3) 

for x, o, c in zip(beatnotes, orders, colors): # Plot beatnotes colored according to group
    plt.stem(x, 1+maxTotOrder-np.sum(abs(o),1), 'o'+c) 

plt.xlim(0, maxf)    
plt.ylim(0, maxTotOrder+1)
plt.title('Beatnotes up to total order ' + str(maxTotOrder) + '')
plt.xlabel('Frequcny [Hz]')
plt.ylabel(str(maxTotOrder) + ' - Beatnote order + 1')
plt.grid('on')

#%% List dangerous frequencies for h of t
hoft_danger_mask = []
print('List of beatnotes in ITF band (f < ' + str(hoftmax/1e3) +' kHz)')
for grp, freqs, combs in zip(sbgroups, beatnotes, orders):
    dm = (freqs < hoftmax) & (freqs > 0)
    print(grp, ': ', freqs[dm])
    hoft_danger_mask.append(dm)
print('')
    
#%% List dangerous frequencies for control loops
demod_danger_mask = []
for demOrd in np.arange(1, maxOrderDem+1):
    for demf in allsbs:
        for freqs, combs, grp in zip(beatnotes, orders, sbgroups):
            dm = (freqs!=demOrd*demf) & (abs(freqs-demOrd*demf) < minSpacing)
            if sum(dm) > 0:
                    print('List of ' + grp + '\'s beatnotes too close to ' + str(demOrd) + '*' + str(demf/1e6) + 'MHz)')
                    for f, cmb in zip(freqs[dm], combs[dm]):
                        print(f, cmb, "(difference = " + str(f-demOrd*demf) +" Hz)")
                    print('')
            demod_danger_mask.append(dm)