"""
Formato npz
Rasterizzazione spettrogramma per formato eps
"""

import numpy as np
import scipy as sp
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap

npulse = 84682

"""
# Lines used to produce the datafile

M = msecompare2.MSEcompare(npulse, cxuid='JETPPF', cxdda='CXGM')

#Bthr = M.W.lines.Bthr
tkc1m = M.W.kc1m.time
fkc1m = M.W.kc1m.F
Akc1m = M.W.kc1m.A0
Nkc1m = M.W.kc1m.Ntor

np.savez('data84682.npz', tkc1m=tkc1m, fkc1m=fkc1m, Akc1m=Akc1m, Nkc1m=Nkc1m)

"""

fff = np.load('data_5b.npz')


Bthr = 5.e-7

# plot parameters
xlim = [4.0, 7.5]
ylim = [0.0, 45.0]
title = '84682'

tkc1m = fff['tkc1m']
fkc1m = fff['fkc1m']
Akc1m = fff['Akc1m']
Nkc1m = fff['Nkc1m']

tlim = [4, 9]
idt = (tkc1m >= tlim[0] ) & (tkc1m <= tlim[1])

Ntor = Nkc1m[idt,:].T    #toroidal numbers
t = tkc1m[idt]
F = fkc1m * 1e-3 # in KHz

A0 = Akc1m[idt,:].T        # raw spectral amplitudes
dfreq = fkc1m[1]- fkc1m[0]
# A1 spectral amplitudes in Volts
A1 = A0.copy()
A1[1:-1,:] = np.sqrt(2*dfreq * (A0[2:,:] + A0[1:-1,:] +  A0[:-2,:])) 
A1[0,:] = A1[1,:]
A1[-1,:] = A1[-2,:]
# A2 spectral amplitudes in Tesla
FF = F.copy()
FF[0:17] = FF[17] # frequencies below 2 kHz are altered to reduce ELMs noise
A2 = A1.T/.06/2.e3/np.pi/FF
# A3 mask for the toroidal number spectrum
A3 = A2.T 
Nmasked = np.ma.masked_where(A3 < Bthr, Ntor) 

# plot figure (a) with rasterization
plt.figure(facecolor='w')
plt.imshow(np.log10(A3), interpolation='nearest', 
    aspect='auto', origin='lower', extent=(t.min(), 
    t.max(), F.min(), F.max()), zorder=-10)
ax = plt.gca()
cbar = plt.colorbar()
cbar.ax.set_ylabel('log amplitude')
plt.xlabel('Time (s)', fontsize=14)
plt.ylabel('Frequency (kHz)', fontsize=14)
plt.xlim(xlim)
plt.ylim(ylim)
plt.title(title, fontsize=14)
plt.text(4.2, 40.0, '(a)', fontsize=16, color='w')
plt.text(7.2, 40.0, '(a)', fontsize=16)
plt.xticks(fontsize=14)
plt.yticks(fontsize=14)

ax.set_rasterization_zorder(-4)
plt.savefig('Output5b.eps')
plt.show()