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


""" CREATED: 7/2012
    AUTHOR: MICHAL ODSTRČIL
"""


import matplotlib
matplotlib.rcParams['backend'] = 'Agg'
matplotlib.rc('font',  size='10')
matplotlib.rc('text', usetex=True)  # FIXME !! nicer but slower !!!

import pygolem_lite

from pygolem_lite import Shot


from numpy import *
from pygolem_lite.config import *
from pygolem_lite.modules import *
from matplotlib.pyplot import *
import time
from shutil import copy, move
import os
from scipy.interpolate import interp1d


def analysis():
    #pass
    try:
	data = Shot()
	

	I = data['plasma_current_mean'] * 1e-6
	Bt = data['toroidal_field_mean']
	P = data['input_power_mean']*1e-3
	n = data['electron_density_mean'] * 1e-19
	M  = 1 # hydrogen !!
	R = MajorRadius
	eps = Aspect
	kappa = 1
	const = 5.62e-2
	print "I", I
	print "B", Bt
	print "P", P
	print "n", n
	print "R", R
	print "Aspect", Aspect
	
	t_98 = const*I**0.93*Bt**0.15*P**-0.69*n**0.41*M**0.18*R**1.97*eps**0.58*kappa**0.78
	
	print "confinent time = ", t_98 * 1e6
	
	saveconst('t_98', t_98)
    except Exception, e:
	print " confinent time failed ", str(e)


def plot_data(ftype):


    data = Shot()
    plasma = data['plasma']
    start = data['plasma_start']
    end = data['plasma_end']
    tvec,_ = data['plasma_current']

    ind_plasma = (tvec > start) & ( tvec < end)

    nGW_limit = 3  # n/nGW < 3 
    calb_photo = 10 if data.shot_num < 10025 else 30 # guess of absolute calbration of photodiode !!

    if not plasma:
	t_cd, I_cd = data['current_cd_coils']   # current drive coils current
	#plot(t_cd, I_cd)
	#savefig('I_cd.png')
	#clf()
	_, I_cd, _ = DiffFilter(I_cd, mean(diff(t_cd)), 1000, 1e2)  # smooth the data 
	t_PHI, PHI = data['total_magnetic_flux']
	
	
	I_cd = interp(t_PHI, t_cd, squeeze(I_cd), left=0, right=0)

	data = [[get_data('total_magnetic_flux', 'Total magnetic flux  $\Phi$', '$\Psi$ [Vs]',  xlim = [0, 40], reduction = True),
		get_data([array([0, 40])*1e-3, ones(2)*MaxTransformatorSaturation], 'Max saturation', '$\Psi$ [Vs]' , ylim = [0,0.18], xlim = [0, 40], reduction = True)],
		get_data( [t_PHI, I_cd] , 'Current drive I', 'I [A]', xlabel = 'Time [ms]',  xlim = [0, 40], tvec_rescale=1e3),
		]


	paralel_multiplot(data, 'Total magnetic flux', 'magnetic_flux', (9,6),  file_type = ftype)

	plot( I_cd,PHI)
	#plot( PHI)

	savefig('I.png')
	clf()
	data = get_data( [I_cd, PHI ] , 'Hysteresis curve', '$\Phi$ [Vs]', xlabel = 'I [A]',xlim=[amin(I_cd), amax(I_cd)],  plot_limits = False, tvec_rescale=1)
	paralel_multiplot(data, 'Hysteresis curve', 'hysteresis', (9,3),  file_type = ftype)

	
	paralel_multiplot(data, "" , 'icon', (4,3), 40)

    
	return

    

    t0 = time.time()

    tvec, safety_factor = Shot()['safety_factor']

    try:
	ne_corr = r0/a*(pos_part(1-R/r0))**1.5
    except:
	tvec_n_e = tvec

    allow_corr = False  #  allow correction of plasma possition 
    try:
	tvec_r, plasma_radius =  Shot()['plasma_radius']
	reliability =  Shot()['plasma_position:reliability']
	assert reliability < 100, "Unreliable plasma position"
	plasma_radius_f = interp1d(tvec_r, plasma_radius, fill_value = 0, bounds_error  =  False)
	
	corr = (plasma_radius_f(tvec) / MeanPlasmaRadius ) ** 2  # corrected for plasma readius
	plot(tvec, corr)
	savefig('corr.png')
	close()
	
	corr[ corr < 0.3 ] = nan
	
	print "======= mean(isnan(corr)) ", mean(isnan(corr[ind_plasma]))
	if mean(isnan(corr[ind_plasma])) > 0.3:
	  del corr
	else:
	    allow_corr = True
    except Exception, e:
	print "correction on plasma position failed", str(e)
	allow_corr = False
      
    _, nGW =  Shot()['greenwald_density']
    nGW_med = max(nGW[ind_plasma])

      
    try:
	tvec_n_e, n_e =  Shot()['electron_density']
    except:
	#raise
	print "electron_density failed"
      
    try:
      	tvec_r, R_position =  Shot()['plasma_position_r']
	plasma_R_pos_f = interp1d(tvec_r, R_position, fill_value = 0, bounds_error  =  False)
	

	a  = 0.1 #[m]
	r0 = plasma_radius_f(tvec_n_e)
	R = plasma_R_pos_f(tvec_n_e)-0.4
	pos_part = lambda x: (x+abs(x))/2
	
	ne_corr_nu1   = r0/a*(pos_part(1-R/r0))**1.5  #peaking factor 1; profile = (1-(r/a)^2)
	ne_corr_nu0_5 = r0/a*(pos_part(1-R/r0))  #peaking factor 0.5; profile = sqrt(1-(r/a)^2)
	

	ne_corr_nu1[ne_corr_nu1>1.5] = nan
	ne_corr_nu0_5[ne_corr_nu0_5>1.5] = nan
	
	plot(tvec_n_e, ne_corr_nu1)
	plot(tvec_n_e, ne_corr_nu0_5)

	
	savefig("ne_corr_nu.png")
	close()
	
	
	if all(isnan(ne_corr_nu0_5)):
	    del ne_corr_nu0_5
	    del ne_corr_nu1
      
    except Exception, e:
      print "electron_density correction failed", str(e)
	
    safety_factor[safety_factor >  2* nanmedian(safety_factor[ind_plasma])] = nan

    #try:
	#temp_spec =  Shot()['spectrometr:temperature']
	#power_spec =  Shot()['spectrometr:total_radiated_power']
    #except:
	#print "missing data spectrometer"

    xlim = array([start, end])*1e3
    ylim = [0, None]
    figsize = (10,4)
    params = dict(xlim = xlim, ylim = ylim, reduction = True)
    data_1 = [
	[
	get_data([tvec,safety_factor], 'Safety factor edge (Q$_{edge}$)', 'Q [-]' , **params),
	get_data([tvec,safety_factor/3], '(Q$_{center}$) (peaking factor = 2)', 'Q [-]' , **params),
	get_data([tvec,corr*safety_factor], 'Q edge with radius', 'Q [-]' , fmt='b--', **params) if allow_corr else [],
	get_data([tvec,corr*safety_factor/3], 'Q center with radius', 'Q [-]' , fmt='r--', **params) if allow_corr else []
	],
	[
	get_data([tvec,1/safety_factor], 'Inverse Safety factor $\iota_{edge}$', '$\iota$ [-]' , **params),
	get_data([tvec,3/safety_factor], ' $\iota_{center}$ (peaking factor = 2)', '$\iota$ [-]' , **params),
	get_data([tvec,1/(safety_factor*corr)], ' $\iota_{edge}$ with radius', '$\iota$ [-]' ,   fmt='b--', **params) if allow_corr else [],
	get_data([tvec,3/(safety_factor*corr)], ' $\iota_{center}$ with radius', '$\iota$ [-]' , fmt='r--', **params) if allow_corr else [],
	get_data([array([start,end]), array([0.5, 0.5])], 'Q=2', '$\iota$ [-]' ,  fmt = "k:", **params)
	]
	]
    try:
	multiplot(data_1, 'Safety factor', 'plot_1', figsize,  file_type = ftype)
    except Exception, e:
	print 'Safety factor failed', str(e)
	
    data_2 = [[
	get_data('electron_temperature', 'Electron temperature', 'T [eV]' , **params),
	get_data('electron_temperature', '$T_e$ corrected by plasma radius', 'T [eV]' , data_rescale = 1/corr, fmt='r--', **params) if allow_corr else [], 
	get_data('spectrometr:temperature', '$T_e$ spectrometer', 'T [eV]' ,ylim = ylim,xlim = xlim, fmt='r--') , 
	]]

    try:
	multiplot(data_2, 'Electron temperature', 'plot_2', figsize,  file_type = ftype)
    except Exception, e:
	print 'Electron temperature failed', str(e)
    
    data_3 = [[
	get_data('input_power_total', 'Ohmic heating - total', 'P [kW]' ,xlim = xlim, data_rescale=1e-3, reduction = True),
	get_data('input_power_magnetic', 'Ohmic heating - magnetic', 'P [kW]' ,xlim = xlim, data_rescale=-1e-3, reduction = True),
	get_data('input_power_plasma', 'Ohmic heating - plasma', 'P [kW]' , xlim = xlim, data_rescale=-1e-3, reduction = True),
	get_data('input_power_chamber', 'Ohmic heating - chamber', 'P [kW]' ,xlim = xlim, data_rescale=-1e-3, reduction = True),
	get_data('photodiode', 'Output power (diod)', 'P [kW]' ,xlim = xlim, data_rescale = calb_photo , reduction = True, fmt = "k--"),
	get_data('spectrometr:total_radiated_power', 'Output power (spectrometer)', 'P [kW]' ,xlim = xlim, data_rescale=1e-3 ,  fmt = "r--"),
	get_data([array([start,end]), array([0, 0])], '', 'P [kW]' , xlim = xlim, reduction = True, fmt = "k:")
	]]
	
    try:
      multiplot(data_3, 'Input / Output power', 'plot_3', figsize,  file_type = ftype)
    except Exception, e:
      print 'Input / Output power', str(e)
      
    if 'n_e' in locals():
      data_4 = [[
	  get_data('electron_density', 'Electron density ($n_e$)', 'n$_{e}$ [10$^{19}\cdot$m$^{-3}$]' , data_rescale= 1e-19, **params) ,
	  get_data('electron_density', '$n_e$ - position corrected, $\\nu = 1$', 'n$_{e}$ [10$^{19}\cdot$m$^{-3}$]' ,data_rescale= 1e-19/ne_corr_nu1, fmt='k--', **params) if 'ne_corr_nu1' in locals() else [],
	  get_data('electron_density', '$n_e$ - position corrected, $\\nu = 0.5$', 'n$_{e}$ [10$^{19}\cdot$m$^{-3}$]', data_rescale= 1e-19/ne_corr_nu0_5, fmt='k:', **params) if 'ne_corr_nu0_5' in locals() else [] ,
	  get_data('greenwald_density', 'Greenwald density ($n_{GW}$)', 'n$_{GW}$ [10$^{19}\cdot$m$^{-3}$]' , data_rescale= 1e-19, **params),
	  get_data('greenwald_density', 'Greenwald density with radius ($n_{GW}$)', 'n$_{GW}$ [10$^{19}\cdot$m$^{-3}$]' ,data_rescale=1e-19/corr, fmt='r--', **params) if allow_corr else [], 
	  get_data([array([start,end]), array([nGW_med, nGW_med])*nGW_limit], '~GW limit', 'n$_{e}$ [10$^{19}\cdot$m$^{-3}$]' , data_rescale= 1e-19 , fmt = "k:", **params) if 'n_e' in locals() else [], 
	  ]]
      
      try:
	  multiplot(data_4, 'Density', 'plot_4', figsize,  file_type = ftype)
      except Exception, e:
	  print 'Density failed ', str(e)


    #os.system('convert -resize 150x120\! total_plot.png icon.png')

    print "===== plotting done ===== "
    data = data_1 + data_2 +data_3 + (data_4 if 'n_e' in locals() else [])
    
    #print data 
    
    try:
	multiplot(data, "" , 'icon', (4,3), 40)
    except Exception, e:
	print 'Icon failed', str(e)
	#raise



def main():
    if sys.argv[1] ==  "analysis":
	analysis()
	
    if sys.argv[1] ==  "plots":
	plot_data('png')
	saveconst('status', 0)

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


if __name__ == "__main__":
    main()