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


""" CREATED: 4/2013
    AUTHOR: TOMÁŠ ODSTRČIL
"""


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

from scipy.signal import *

from numpy import *
from matplotlib.pyplot import *
from stft import stft
from matplotlib.ticker import MultipleLocator, FormatStrFormatter,LogLocator
import os
import sys

from pygolem_lite import load_adv, saveconst, Shot
from  multiprocessing import Process, Pool, cpu_count
from scipy.stats.mstats import mquantiles


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 Spectrogram((frequenciCutOff,tvec, signal,name,start, end, plasma)):
    
    
    
    
    
    
    print name

    signal = double(signal)

    startAdv =  max(start*0.8, tvec[0]*1e3)
    endAdv =  min((end-start)*1.2+start, tvec[-1]*1e3)   

    # cut-off signal 
    ind = (tvec > startAdv*1e-3) & (tvec < endAdv*1e-3)
    signal = signal[ind]
    tvec = tvec[ind]
    
    
    dt = mean(diff(tvec))

    N = len(tvec)

    for i in range(3):
	filtered = medfilt(signal, 2*floor(N/500)+1)

	res = abs(signal - filtered)
        ind = argsort(res)
        n_out = floor(N/400)
        ind_2 = in1d(range(N), ind[-n_out:]) & (res > 0.2*amax(filtered))
	signal[ind_2] =  filtered[ind_2]    #remove "n_out" most distant points from moving medianZ
                                               
    print "computing spectrum"
    #spec, scale = cwt(signal, dt, dj=0.05, wf='morlet', p=omega0, extmethod='none', extlength='powerof2')
    
    fmin = 3000
    fmax = frequenciCutOff
    
    spec, freq = stft(signal, dt, dTime=1e-5,  p=1000, extmethod='periodic', extlength='powerof2',res = 200,fmin=fmin,fmax=fmax)
    
    print "done"



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

    contrast = 10
    #spec = spec[:,(tvec > startAdv )*( tvec <endAdv)] 

    pole = log(1+contrast*abs(spec))



    #pole = pole[freq > frequenciCutOff, : ] 
    #freq = freq[freq > frequenciCutOff] 

    fig = figure()



    ax = fig.add_axes([0.1, 0.1, 0.8, 0.85])
    freq/=1000

    img = ax.imshow(pole, extent=[startAdv,endAdv ,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([startAdv,endAdv,amin(freq), amax(freq)])

    ax.set_xlabel('time [ms]')


    ax.set_ylabel('Frequency [kHz]')
    savefig('data/spectrogram_'+name+'.png')
    clf()

    ind = (tvec > startAdv*1e-3 ) & ( tvec<endAdv*1e-3)
    signal= signal[ ind ]
    plot(1e3*tvec[ ind ],signal,'k',linewidth=0.1)
    #show()

    
    #axvline(linewidth=4)
    #minimum = -max(abs(amin(signal)), abs(amax(signal)))
    #maximum = -minimum
    minimum = mquantiles(signal,0.01)
    maximum = mquantiles(signal, 0.99)

    minimum -= (maximum - minimum)*0.2
    maximum += (maximum - minimum)*0.2

    
    if maximum == minimum:
	axis([startAdv,endAdv,None, None])
    else:
	axis([startAdv,endAdv,minimum, maximum])

    #axis([startAdv,endAdv,None, None])
    

    if plasma:
	plot([start, start], [minimum, maximum], 'r--')
	plot([end, end],     [minimum, maximum], 'r--')
	

    xlabel('time [ms]')
    ylabel('Intensity [a.u.]')
    #show() 
    savefig('data/signal_'+name+'.png')
    clf()

Data = Shot()
plasma = Data['plasma']

if plasma:
    start = Data['plasma_start']*1e3
    end = Data['plasma_end']*1e3
else:
    start = 0
    end = 40
    

tvec, data = Shot()['nistandard6132']
tvec = float_(tvec)

frequenciCutOff = 8e4

try:
    os.mkdir('data')
except:
    pass

Ndim = size(data,1)

try:
    p = Pool(cpu_count())
    p.map( Spectrogram, [ (frequenciCutOff,tvec, data[:,i],str(i),start,end, plasma) for i in range(Ndim)])
    p.close()
    
    saveconst('status', 0)

except:
    saveconst('status', 1)

    raise

    
os.system('convert -resize 150 data/spectrogram_1.png icon.png')