# -*- coding: utf-8 -*-
"""
Spyder Editor
Created on Fri Aug  7 15:11:50 2020. Last modified May 2021
@author: piaastone
"""

#from scipy import stats
import matplotlib.pyplot as plot
import numpy as np
from scipy.fft import fft 

#Ora andiamo dal dominio tempo a quello della frequenza
#dt e'il tempo di campionamento. tmax il tempo complessivo, in secondi
#Attenzione: dt deve essere minore di 1/(2*f) condizione sufficente (non necessaria..)
dt=0.001
tmax=1024 #1024 , 128 ad esempio
time= np.arange(0, tmax, dt)  
f1=100 
omega=2*np.pi*f1
AMP=1
# Le ampiezze sotto potete variarle e vedere cosa succede. Anche le frequenze
AMP11=3
AMP12=1.5
f11=100
f12=150
omega1=2*np.pi*f11
omega2=2*np.pi*f12
amplitude   = AMP*np.sin(omega*time)
amplitudeS=AMP11*np.sin(omega1*time)+AMP12*np.sin(omega2*time)

#Define the parameters needed to go from time to frequency
#Apply FFT= fast Fourier Transform

dnu=1/(tmax)  
freq=np.arange(0, 1/dt,dnu)

spA = fft(amplitude,norm="ortho")
spS = fft(amplitudeS,norm="ortho")


#Figure sinusoide singola e somma di sinusoidi
plot.figure()
plot.plot(time, amplitude)
plot.title('One sinusoidal signal')
plot.ylabel('Amplitude')
plot.xlabel('Time[s]')
plot.grid(True, which='both')
plot.xlim(0, 0.3)  #Ad hoc ZOOM
plot.figure()
plot.plot(time, amplitudeS)
plot.title('Sum of two sinusoidal signals')
plot.ylabel('Amplitude')
plot.xlabel('Time[s]')
plot.grid(True, which='both')
plot.xlim(0, 0.3)  #Ad hoc ZOOM

# plot.figure()
# plot.plot(freq, spS.real*spS.real+spS.imag*spS.imag)
# plot.xlim(f11-100, f11+200)
# plot.grid('True')
# plot.xlabel('Freq [Hz]')
# plot.title('zoom FFT squared modulus')
# plot.show

##White noise
#amplitude and standard deviation
AMPN=10
STDN=10
mediaN=0 #0
standN= STDN
NN=amplitude.size #mettiamo tanti elementi quanti nella sinusoide
noiseN =AMPN* np.random.normal(mediaN,standN,NN)
plot.figure()
plot.plot(time, noiseN)
plot.title('White Noise')
plot.ylabel('Amplitude')
plot.xlabel('Time [s]')
plot.grid(True, which='both')
#FFFT
FSN=fft(noiseN,norm="ortho")
ValoreassolutoN=np.absolute(FSN)*np.absolute(FSN)
plot.figure()
plot.semilogy(freq, ValoreassolutoN)
plot.grid('True')
plot.xlabel('Freq [Hz]')
plot.title('Noise only. Squared modulus of the FFT')
plot.show

#ANow add the sinusoidal signal to noise
i=1
ncases=2
while i <=ncases:
    if i ==1:
        amplitudeU=amplitude  ####
        spU=spA
    if i ==2:
        amplitudeU=amplitudeS  ####
        spU=spS
    
    SandN=amplitudeU+noiseN
    plot.figure()
    plot.plot(time, SandN)
    plot.plot(time, amplitudeU,'r')
    if i == 1:
         plot.title('Sinusoidal signal in white noise')
    if i == 2:
         plot.title('Two Sinusoidal signals in white noise')
    plot.ylabel('Amplitude')
    plot.xlabel('Time [s]')
    plot.grid(True, which='both')
    #zoom
    plot.figure()
    plot.plot(time, noiseN)
    plot.plot(time, amplitudeU,'r')
    plot.title('ZOOM')
    plot.ylabel('Amplitude')
    plot.xlabel('Time [s]')
    plot.grid(True, which='both')
    plot.xlim(0,0.5)
    #zoom
    plot.figure()
    plot.plot(time, noiseN)
    plot.plot(time, amplitudeU,'r')
    plot.title('Another ZOOM')
    plot.ylabel('Amplitude')
    plot.xlabel('Time [s]')
    plot.grid(True, which='both')
    plot.xlim(0,0.5)
    plot.ylim(-20,20)
    #Ora la FFT
    SpSN=fft(SandN,norm="ortho")
    Valoreassoluto=SpSN.real*SpSN.real+SpSN.imag*SpSN.imag
    ValS=spU.real*spU.real+spU.imag*spU.imag
    #ZOOM
    bord=10
    fini=f11
    if i == 1:
        fend=f11
    if i == 2:
        fend=f12
    plot.figure()
    plot.semilogy(freq, Valoreassoluto)
    plot.xlim(fini-bord, fend+bord)
    plot.ylim(np.median(Valoreassoluto), np.max(Valoreassoluto*1.5))
    plot.grid('True')
    plot.xlabel('Freq [Hz]')
    if i ==1:
        plot.title('FFT one signal and noise')
    if i ==2:
        plot.title('FFT two signals and noise')
    plot.show
    i+=1  #end of the while