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

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

from numpy import *
import os
import sys
from pygolem_lite.modules import list2array, multiplot, get_data
from pygolem_lite import Shot, saveconst, save_adv,load_adv


channels = [2,1]


  
def LoadData():
  
  
    Data = Shot()


    tvec,dataTe = Data['niturbo']
    dataTe = dataTe[:,channels]

    Bt_trigger = Data['Tb']
    plasma_start = Data['plasma_start']
    plasma_end = Data['plasma_end']
    plasma = Data['plasma']
  

    return tvec, dataTe, Bt_trigger,plasma_end,plasma_start, plasma


    
def Calc_Temperature(file_type):
   
    
    tvec, dataTe,Bt_trigger,plasma_end,plasma_start, plasma  = LoadData()

    if not plasma:
	print "============= no plasma ==============="
	return
    

    
    offset_interv = (tvec > Bt_trigger)*(tvec < plasma_start)
    offset = median(dataTe[offset_interv,:], axis = 0)
    dataTe[:,0]-=offset[0]
    dataTe[:,1]-=offset[1]
    dataTe *= 100    # I dont know why... 
    dU = dataTe[:,0]-dataTe[:,1]
    Te = dU/log(2);   #result of the double (triple) probe theory
    
    plasma_interv = (tvec > plasma_start)*(tvec < plasma_end)
    Te[~plasma_interv] = 0
    Te[Te < 0] = 0
    
    save_adv('Te_probe', tvec, Te)
    save_adv('dU_probe', tvec, dU)
    save_adv('dataTe', tvec, dataTe)


    class MyFormatter(plt.ScalarFormatter): 
	def __call__(self, x, pos=None): 
	    if pos==0: 
		return '' 
	    else: return plt.ScalarFormatter.__call__(self, x, pos)  
	
	
	
    fig = plt.figure(num=None, figsize=(6, 8), 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*1e3,    dU,   label = '$U_1-U_2$ ')
    plt.plot(tvec*1e3,   dataTe[:,0] ,   label = '$U_1$ ')
    plt.plot(tvec*1e3,   dataTe[:,1] ,   label = '$U_2$ ')
    leg = plt.legend(loc='upper left', fancybox=True)
    leg.get_frame().set_alpha(0.5)
    plt.ylim(0, None)
    plt.xlim(plasma_start*1e3, plasma_end*1e3)

    plt.ylabel('U [V]')
    
    
    ax = fig.add_subplot(2,1,2)
    plt.plot(tvec*1e3,    Te,   label = 'Probe $T_e$ ')

    leg = plt.legend(loc='upper left', fancybox=True)
    leg.get_frame().set_alpha(0.5)
    plt.ylim(0,None)
    plt.xlim(plasma_start*1e3, plasma_end*1e3)

    plt.ylabel('[eV]')
    plt.xlabel('t [ms]')
    plt.savefig('./Te_new.'+file_type,bbox_inches='tight')

    plt.close()
    
    
    

    title = "Temperature"
    out = [get_data('dU_probe', '$U_1-U_2$ ' , "U [V]"),    # , smoothing = 1e3), 
	    get_data('dataTe', ['$U_1$ ', '$U_2$ '] , "U [V]" ),
	    get_data('Te', '$Probe Te$ ' , "T [eV]" ) ]

    multiplot(out, title , 'temperature', (10,len(out)*2), 100, 'vertical', file_type)

def main():

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



if __name__ == "__main__":
    main()