#!/usr/bin/python2
# -*- coding: utf-8 -*-
#THE SIMPLEST VERSION !!!!

####     Microwaves 2.0
#This algorithm calculate transformation of the sin signal from microwaves density measurement to the phase/amplitude space. 
# First step of the calculation is estimate of the base frequency and calculation of the complex exponential
#with the same frequency.In the second step is signal multiplied by this exponential
#and resulting low frequency signal is smoothed over Gaussian window. Finally complex phase and amplitude are calculated.  

# Authors: Tomas Odstrcil, Ondrej Grover

from time import time
t = time()

from numpy import *
from numpy.fft import fft, ifft
from scipy.signal import fftconvolve   
from scipy.constants import c,m_e,epsilon_0,e
from pygolem_lite.modules import save_adv,load_adv,saveconst
from pygolem_lite import Shot
import os
import sys
print 'include time ',time()-t 



def Demodulation(data,win):
    y = data-mean(data)
    t = time()
    n = len(y)
    fourier = fft(y) #calulcate the fourier transfrom for the sine data
    
    max_frequency_index =  argmax(abs(fourier))
    max_frequency = abs(fourier[max_frequency_index])
    fourier[:] = 0 #cancel out all other frequencies
    fourier[max_frequency_index] = max_frequency

    
    cmpl_exp = ifft(fourier) 

    gauss = exp(-arange(-3*win,3*win)**2/win**2)  
    
    signal = fftconvolve(y*cmpl_exp ,gauss,mode='same' )

    amplitude = abs(signal)     
  
    phase = angle(signal)

    #remove jumps in the signal
    #for i in range(3):
    #diff_phase = diff(phase)
    #phase[1:] -= cumsum(where(abs(abs(diff_phase) - 2*pi) < 1, diff_phase, 0), out = diff_phase)
    phase = unwrap(phase)

	
    print 'demodulated  ',time()-t
        
    return amplitude,phase

def LoadData():
    Data = Shot()

    gd = Shot().get_data
    tvec, density1  = gd('any', 'density')
    tvec, density2  = gd('any', 'density_2')
       
    return tvec,density1,density2


    
    
def graphs():
    
    import matplotlib 
    matplotlib.rcParams['backend'] = 'Agg'
    matplotlib.rc('font',  size='10')
    matplotlib.rc('text', usetex=True)  # FIXME !! nicer but slower !!!
    import matplotlib.pyplot as plt
    
    class MyFormatter(plt.ScalarFormatter): 
	def __call__(self, x, pos=None): 
	    if pos==0: 
		return '' 
	    else: return plt.ScalarFormatter.__call__(self, x, pos) 
	    
	    
    tvec, phase_pila = load_adv('results/phase_saw')
    tvec, phase_sinus = load_adv('results/phase_sinus')
    tvec, phase = load_adv('results/phase_substracted')
    tvec, amplitude = load_adv('results/amplitude_sinus')   
    tvec, n_e = load_adv('results/electron_density')

    
    
    fig = plt.figure(num=None, figsize=(10, 6), dpi=80, facecolor='w', edgecolor='k')
    plt.subplots_adjust(hspace=0, wspace = 0)
    
    ax = fig.add_subplot(2,1,1)
    ax.xaxis.set_major_formatter( plt.NullFormatter() )
    ax.yaxis.set_major_formatter( MyFormatter() )
    plt.plot(tvec*1000,-phase_pila+mean(phase_pila),'--', label = 'saw phase' )
    plt.plot(tvec*1000,-phase_sinus+mean(phase_sinus),'--', label = 'signal phase')
    plt.plot(tvec*1000,phase, 'k',label = 'substracted phase')
    plt.axis('tight')

    plt.xlim(0,None)
    plt.xlabel('time [ms]')
    plt.ylabel('phase [rad]')
    leg = plt.legend(loc='best', fancybox=True)
    leg.get_frame().set_alpha(0.5)
   
    ax = fig.add_subplot(2,1,2)
    plt.plot(tvec*1000,amplitude,label = 'amplitude')
    plt.xlim(0,None)
    plt.ylim(0,None)
    leg = plt.legend(loc='best', fancybox=True)
    leg.get_frame().set_alpha(0.5)
    
    plt.ylabel('amplitude [a.u.]')
    plt.savefig('graphs/demodulation.png',bbox_inches='tight')
    plt.close()
    
    
    Data = Shot()
    plasma_start = Data['plasma_start']
    plasma_end = Data['plasma_end']
    fig = plt.figure(num=None, figsize=(10, 3), dpi=80, facecolor='w', edgecolor='k')
    plt.plot(tvec*1000,n_e/1e19,label = '$n_e$')
    plt.ylabel('$<n_e>$ [$10^{19}\,m^{-3}$]')
    plt.xlabel('time [ms]')
    plt.xlim(0,20)
    plt.ylim(0,None)
    plt.axvline(x = 1000*plasma_start,linestyle = '--')
    plt.axvline(x = 1000*plasma_end, linestyle = '--')

    plt.savefig('graphs/electron_density.png',bbox_inches='tight')
    plt.close()
    
    
    

def main():



    
    for path in ['graphs', 'results' ]:
	if not os.path.exists(path):
	    os.mkdir(path)
	    
    if sys.argv[1] ==  "analysis":
	
	win = 30e-6 #[s]
	t = time()
	tvec,density1,density2 = LoadData()
	dt = (tvec[-1]-tvec[0])/len(tvec)
	print 'load time ', time()-t

	(amplitude2,phase_pila) = Demodulation(density2,win/dt)  
	(amplitude,phase_sinus) = Demodulation(density1,win/dt)

	downsample = int(win/dt/2)
	amplitude = amplitude[::downsample]
	phase_pila = phase_pila[::downsample]
	phase_sinus = phase_sinus[::downsample]
	tvec = tvec[::downsample]
	
	phase = phase_pila-phase_sinus
	phase -= median(phase)
	
	save_adv('results/phase_saw', tvec, phase_pila)
	save_adv('results/phase_sinus', tvec, phase_sinus)
	save_adv('results/phase_substracted', tvec, phase)
	save_adv('results/amplitude_sinus', tvec, amplitude)    

	a = 0.01   #[m]
	f_0 = 75e9 #[Hz]
	lambda_0 = c/f_0
	n_e = 4*pi*m_e*epsilon_0*c**2/(e**2*lambda_0*2*a)*phase
	save_adv('results/electron_density', tvec, n_e)

    if sys.argv[1] ==  "plots":
	graphs()
	saveconst('status', 0)



if __name__ == "__main__":
    main()