from __future__ import division

import numpy as np
from pylab import *
import csv
from collections import defaultdict
from scipy.signal import lfilter,welch,hann
from scipy import fft
import sys

import os
cppsimsharedhome = os.getenv("CPPSIMSHAREDHOME")
if cppsimsharedhome != None:
   CPPSIMSHARED_PATH = '%s' % cppsimsharedhome
else:
   HOME = os.getenv("HOME")
   CPPSIMSHARED_PATH = '%s/CppSim/CppSimShared' % HOME

sys.path.append(CPPSIMSHARED_PATH + '/Python')
from cppsimdata import *

def calc_win_fft(sd_out,fs):
   N = len(sd_out)
   hwin = hann(N)

   # assume sd_out is Delta-Sigma sequence alternating between 0 and 1
   sd_out_win = sd_out*hwin
   # remove low frequency spectral bleeding
   sd_out_minus_mean = sd_out - mean(sd_out_win)/mean(hwin)
   sd_out_win = sd_out_minus_mean*hwin

   s_out = (4/sum(hwin)*absolute(fft(sd_out_win,N)))**2
   freq = arange(N)*1/N*fs

   return freq,s_out

def calc_sig_noise_over_freq_range(s_out,freq,fstart,fend): 
   i_start = argwhere(freq >= fstart)[0]
   i_end = argwhere(freq >= fend)[0]
   f_ind_bw = arange(i_start,i_end+1)

   freq_bw = freq[f_ind_bw]
   s_out_bw = s_out[f_ind_bw]

   s_ind_max = argmax(s_out_bw)
   num_bw_samples = len(s_out_bw)
   fs = freq[-1]
#   print("    BW for SNR calc: %3.1fHz to %3.1fkHz for fs = %5.3fMHz (%d samples)" % (fstart,fend*1e-3,fs*1e-6,num_bw_samples)) 

   # calculate indices corresponding to signal (assume 3 bins)
   signal_indices = arange(-1,1+1) + s_ind_max

   # computation of noise power is slightly off since the values in the 3 signal bins
   # are assumed to be 0
   # -> need 'num_bw_samples' above to be reasonable large for this to be a small error   
   signal_power_raw = sum(s_out_bw[signal_indices])
   noise_power_raw = sum(s_out_bw) - signal_power_raw

   # scale signal and noise power by 2/3 due to Hann window
   signal_power = 2/3*signal_power_raw
   noise_power = 2/3*noise_power_raw

   return signal_power,noise_power,s_out_bw,freq_bw,signal_indices

close('all')

# look at frequencies from 20Hz to 20kHz
fstart = 20
fend = 20e3

fig_snr = figure(1,figsize=(14,8))
fig_snr.clf()
# set line thickness in plots to 2
rcParams['lines.linewidth'] = 2

dir_name = os.getcwd()
file_basename = 'test_globals'
file_suffix_val = 'tr0'
file_name= '%s/%s.%s' % (dir_name,file_basename,file_suffix_val)
if os.path.isfile(file_name):
   print("... Opening global param file '%s'..." % file_name)
else:
   print("... Stop:  global param file '%s' does not exist..." % file_name)
   sys.exit()
data = CppSimData(file_name)
amp_in = data.evalsig('amp_in_gl')

file_basename = 'test'
sig_name = 'out_ord2'
sig_name2 = 'out_ord3'
amp_in_db = zeros(len(amp_in))
sndr_sig = zeros(len(amp_in))
sndr_sig2 = zeros(len(amp_in))

fig_snr = figure(1,figsize=(10,6))
fig_snr.clf()
# set line thickness in plots to 2
rcParams['lines.linewidth'] = 2

file_count = 0
for amp_in_val in amp_in:
   file_suffix_val = 'tr%d' % file_count
   file_name= '%s/%s.%s' % (dir_name,file_basename,file_suffix_val)
   if os.path.isfile(file_name):
      print("... Opening alter file '%s'..." % file_name)
   else:
      print("... Stop:  alter file '%s' does not exist..." % file_name)
      sys.exit()

   amp_in_db[file_count] = 20*log10(amp_in_val)

   data = CppSimData(file_name)
   sd_out = data.evalsig(sig_name) 
   t = data.evalsig('TIME')
   t_delta = lfilter([1,-1],[1],t)
   Ts = mean(t_delta[1:-1])
   fs = 1/Ts

   # assume sd_out corresponds to 1-bit Delta-Sigma sequence alternating between 0 and 1
   sd_out_mean = mean(sd_out)
   sd_out[sd_out < sd_out_mean] = 0
   sd_out[sd_out > sd_out_mean] = 1

   freq,s_out = calc_win_fft(sd_out,fs)
   
   signal_power,noise_power,s_out_bw,freq_bw,signal_indices = calc_sig_noise_over_freq_range(s_out,freq,fstart,fend)
   sndr_sig[file_count] = 10*log10(signal_power/noise_power)  

   ############################################

   sd_out = data.evalsig(sig_name2) 

   # assume sd_out corresponds to 1-bit Delta-Sigma sequence alternating between 0 and 1
   sd_out_mean = mean(sd_out)
   sd_out[sd_out < sd_out_mean] = 0
   sd_out[sd_out > sd_out_mean] = 1

   freq,s_out = calc_win_fft(sd_out,fs)
   
   signal_power,noise_power,s_out_bw,freq_bw,signal_indices = calc_sig_noise_over_freq_range(s_out,freq,fstart,fend)
   sndr_sig2[file_count] = 10*log10(signal_power/noise_power) 

   file_count = file_count + 1

plot(amp_in_db,sndr_sig,'m-o',label='Ord2')
plot(amp_in_db,sndr_sig2,'b-o',label='Ord3')

#axis([-63, 3, 20, 130])
xlabel('Input Amplitude (dB FS)')
ylabel('SNDR (dB)')
fs = freq[-1]
title('SNDR versus Input Amplitude')
grid(True,which='both')
legend(loc='upper left',prop={'size':14})
fig_snr.show()

