#!/usr/bin/env python 
# -*- coding: utf-8 -*-

import matplotlib
#matplotlib.rcParams['backend'] = 'Qt4Agg'
matplotlib.rcParams['backend'] = 'Agg'


from scipy.signal import *

from numpy import *
from matplotlib.pyplot import *
import mlpy
#import cwt

from matplotlib.ticker import MultipleLocator, FormatStrFormatter,LogLocator


matplotlib.rcParams['xtick.direction'] = 'out'
matplotlib.rcParams['ytick.direction'] = 'out'

matplotlib.rcParams['xtick.major.size'] = 10
matplotlib.rcParams['xtick.minor.size'] = 7

matplotlib.rcParams['ytick.major.size'] = 10
matplotlib.rcParams['ytick.minor.size'] = 7

def PlotSpectrogram(spec,freq,name,start = None, end = None):
    contrast = 10		
    pole = log(1+contrast*abs(spec))
    if (start== None):
	start= 0
	end = 1
	
    fig = figure()
    ax = fig.add_axes([0.1, 0.1, 0.8, 0.85])
    img = ax.imshow(pole, extent=[start*1000,end*1000 ,freq[-1], freq[0]], aspect='auto') 
    ax.set_yscale('log', nonposy='clip')
    minorLocator   = MultipleLocator(1)
    ax.xaxis.set_minor_locator(minorLocator)
    #ax.plot([start, start], [amin(freq), amax(freq)], 'w--')
    #ax.plot([end  ,   end], [amin(freq), amax(freq)], 'w--')
    ax.axis([start*1000,end*1000,amin(freq), amax(freq)])
    ax.set_xlabel('time [ms]')
    ax.set_ylabel('Frequency [Hz]')
    fig.colorbar(img)
    savefig(name+'_spectrogram'+'.png')
    clf()

def Plotdata(tvec,signal,name):
    

    plot(tvec[0:100]*1000,signal[0:100]+1e-9,'k',linewidth=0.1)
    xlabel('time [ms]')
    ylabel('U [V]')
    savefig(name+'_signal'+'.png')
    clf()
    plot(tvec*1000,signal,'k',linewidth=0.1)
    xlabel('time [ms]')
    ylabel('U [V]')
    savefig(name+'_signal_full'+'.png')
    clf()
    

def Spectrogram(omega0,frequenciCutOff,signal,dt):

    #spec, scale = cwt.cwt(signal, dt, dj=0.05, wf='morlet', p=omega0, extmethod='periodic', extlength='powerof2',res = 2000)
    spec, scale = mlpy.cwt(signal, dt, dj=0.05, wf='morlet', p=omega0, extmethod='periodic', extlength='powerof2')


    # approximate scales through frequencies
    freq = (omega0 + sqrt(2.0 + omega0**2))/(4*pi * scale[1:])

    spec = spec[freq > frequenciCutOff, : ] 
    freq = freq[freq > frequenciCutOff] 



    return spec,freq

def FFTanalysis(signal,start, end, name,dt,frequenciCutOff):

    Fsignal = fft.rfft(signal)

    fvec = 2/dt/len(Fsignal)*arange(size(Fsignal))
    Fsignal = Fsignal[fvec > frequenciCutOff]
    fvec = fvec[fvec > frequenciCutOff]
    plot(fvec,abs(Fsignal),linewidth=0.1)
    xlabel('Frequency [Hz]')
    ylabel('Intensity [a.u.]')
    yscale('log', nonposy='clip')
    xscale('log', nonposy='clip')
    limits = sort(abs(Fsignal))[[len(Fsignal)*1e-3,len(Fsignal)*(1-1e-3 )]]
    axis([fvec[0],fvec[-1],limits[0], amax(abs (Fsignal))])
  
    savefig(name+'_fft'+'.png')
    

    
    
    clf()
    

    
def main():
   
    dir_list = ('.')


    name_list = ('NIbasic_01','NIturbo_08', 'PapouchSt_05', 'PapouchSt_08', 'PapouchSt_09', 'NI6251ext_01')
    for directory in dir_list:
	for file_name in name_list:
	    name = file_name
	    print 'processing ', name
	    signal = loadtxt(name)
	    ##print signal
	    
	    #print shape(signal)
	   
	    if all(signal == 0):
		continue
	    omega0 = 50
	    frequenciCutOff =1e3
	    dt = signal[1,0]-signal[0,0]
	    start = signal[0,0]
	    end   = signal[-1,0]
	    Plotdata(signal[:,0],signal[:,1],name)

	    try:
		signal[:,1] -=mean(signal[:,1])
		( spec,freq) = Spectrogram(omega0,frequenciCutOff,signal[:,1],dt)
		PlotSpectrogram(spec,freq,name,start, end)
		FFTanalysis(signal[:,1],start, end, name,dt,frequenciCutOff)
	    except ValueError:
		print 'chyba'
    
def main2():


    start = 0
    end = 0.04

    omega0 = 50	#čím větší tím lepší frekvenční rozlišení a opačně
    frequenciCutOff =1e4


    signal = loadtxt('textronix25MHz1')
    signal -=mean(signal)
    print size(signal)
    length = size(signal)
    end = size(signal)*4e-8
    t = linspace(start,end,length)
    signal = vstack((t,signal)).T

    name = 'textronix'
    
    #length 
    specLen = 2000.0
    step = int(length/specLen)
    signalPartLen = 10000.0
    nParts = int(length/signalPartLen)+1
    lenPart = int(2**14*0.8)
    dt = (end-start)/length
    overlap = 0.2
    spec = empty
    for i in range(nParts):
	s =   int(  i*lenPart)
	e = int((i+1)*lenPart)
	#if i <=0:
	    #s = 0
	#elif i>= nParts-1:
	    #e = length-1
	
	print s," ",e
	signaltmp = signal[s:e,:]	    
	print shape(signaltmp)
	( specTmp,freqTmp) = Spectrogram(omega0,frequenciCutOff,signaltmp,dt)
    
	specTmp = specTmp[:,s:e:step]
	
	    
	
	
    
    
    #PlotSpectrogram(spec,freq,name,start, end)


    s=   int(size(tvec)/pix)
    pole = pole[:,::s]






    
    
if __name__ == "__main__":
    main()