# coding: utf-8

# # Golem Workshop Kiten 2018 - Data Processing
# 
# This is a notebook for data processing of the Golem operation domain scan.

# In[1]:

get_ipython().magic('matplotlib inline')

# At first, load some basic libraries.

import numpy as np
from urllib.request import urlopen
import shutil
import os 
from matplotlib import pyplot as plt
from scipy.interpolate import interp1d

from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm

# Plot settings, can be changed to globally influence the appearance

plt.rc('text', usetex=True)
plt.rc('font', family='sans-serif')
plt.rc('font', size=22)          # controls default text sizes
plt.rc('axes', titlesize=22)     # fontsize of the axes title
plt.rc('axes', labelsize=20)    # fontsize of the x and y labels
plt.rc('xtick', labelsize=20)    # fontsize of the tick labels
plt.rc('ytick', labelsize=20)    # fontsize of the tick labels
plt.rc('legend', fontsize=22)    # legend fontsize
plt.rc('figure', titlesize=22)  # fontsize of the figure title


# ## Golem Database Access
# 
# Thanks to Golem staff for the golem_data function.
# 
# This should not be changed :).

# In[2]:

def golem_data(shot, diagn):
    """
     Simple interface for GOLEM database
      Use:
	obj = golem_data(10011, 'loop_voltage')
	plot(obj.tvec, obj.data)

	d - object containing all availible informations
	d.tvec  - time vector Nx1
	d.data - data matrix  NxM

      Example:
      from golem_data import golem_data
      from matplotlib.pyplot import *
      obj1 = golem_data(10689 , 'electron_temperature')
      obj2 = golem_data(10689 , 'spectrometr:temperature')
      plot(obj1.tvec,obj1.data,  label=obj1.labels)  %
      errorbar(obj2.tvec, obj2.data, xerr=obj2.tvec_err, yerr=[obj2.data_err[:,0], obj2.data_err[:,1]],  label=obj2.labels )
      xlabel(obj2.ax_labels[0])
      ylabel(obj2.ax_labels[1])
      legend([obj1.name, obj2.name])
      axis([obj1.plasma_start, obj1.plasma_end, 0, None])
      title(obj1.info)
      show()
    """
      
    remote_file = "http://golem.fjfi.cvut.cz/utils/data/"+str(shot)+"/"+diagn+'.npz'
    
    gfile = np.DataSource().open(remote_file, mode='rb')  # remote file
    try:
        d =  dict(np.load(gfile, encoding='bytes', fix_imports=True))
    except IOError:
        raise IOError('Missing diagnostic ' + str(diagn) + ' in shot ' + str(shot))
    
    if not 'tvec' in d and 't_start' in d:
        d['tvec'] = np.linspace( d.pop('t_start'), d.pop('t_end'), len(d['data']) )
    
    try:
        if 'scale' in d:
            d['data'] = np.double(d['data']) * d.pop('scale')  # rescale data 
    except:
        pass
    return Data( **d )


class Data(object):
    """ Simple data handling object"""
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)
    def __repr__(self):
        return "Data object, keys:\n" + str(self.__dict__.keys())
    def __getitem__(self, key):
        return getattr(self, key)
	


# ## Some Useful Scripts
# 
# ### Validate
# Excludes bad shots from a shot list.

# In[3]:

def Validate(shot_list, bad_shots):
    valid = [];
    for shot_no in shot_list:
        try:
            if not(shot_no in bad_shots): 
                b_v = golem_data(shot_no, 'breakdown_voltage').data
                valid.append(shot_no)
        except Exception:
            pass
    return valid;


# ### GetValAtTime
# 
# Takes a diagnostic obtained by `golem_data` and interpolates its value for a specific time in the shot.
# 
# Calling with `plot_progress = True` will show the interpolation on the graph.

# In[4]:

def GetValAtTime(diag, time, plot_progress = False):
    f = interp1d(diag.tvec, diag.data)
    # change to true if you want to see the plots
    if plot_progress:
        fig = plt.figure()
        plt.plot(diag.tvec, diag.data, linewidth=2)
        t = np.linspace(min(diag.tvec), max(diag.tvec), 50)
        plt.plot(t, f(t), linewidth=1)
        plt.plot([time, time], [min(diag.data), max(diag.data)], '--')
        plt.plot(time, f(time), 'o')
        plt.xlabel('$t$ [s]')
        plt.ylabel('$B_\mathrm{t}$ [T]')
        plt.grid(True, axis='both')

        plt.show()
        plt.close()
    return f(time)


# ### Load
# 
# Loads all specified diagnostics for a specified shot number.
# 
# Any newly requested diagnostics should be added here.
# 
# **To use the diagnostic in the scan, one point should be selected (e.g. by interpolation for a specific time).** Toroidal field is now treated like this.

# In[5]:

# Diagnostics should be added here (according to the sample fro)
def Load(shot_no):
    data_dict = {
        'shot_no': shot_no,
        'pressure_request': golem_data(shot_no, 'pressure_request').data,
        'pressure': golem_data(shot_no, 'pressure').data,
        'ucd': golem_data(shot_no, 'ucd').data,
        'electron_temperature': golem_data(shot_no, 'electron_temperature_max').data,
        'breakdown_voltage': golem_data(shot_no, 'breakdown_voltage').data,
        'toroidal_field': GetValAtTime(golem_data(shot_no, 'toroidal_field'), 
                                       golem_data(shot_no, 'plasma_start').data)
        }
    return data_dict


# ### GetScan1, GetScan2
# 
# Sorts the data for further processing. Poorly written, should be somewhat more refined (to allow for arbitrary N-dimensional scans).

# In[6]:

def GetScan1(data_list, scan_var, diag_var):
    d_l = sorted(data_list, key=lambda d: d[scan_var])
    return {
            'x': [d[scan_var] for d in d_l],
            'val': [d[diag_var] for d in d_l],
            }

def GetScan2(data_list, scan_var, diag_var):
    d_l = sorted(data_list, key=lambda d: (d[scan_var[0]], d[scan_var[1]]))
    return {
            'x': [d[scan_var[0]] for d in d_l],
            'y': [d[scan_var[1]] for d in d_l],
            'val': [d[diag_var] for d in d_l],
            }


# ### ScanPlot2D
# 
# Plots a 2D scan (scaled colormap).

# In[7]:

def ScanPlot2D(golem_scan, scan_id, scan_description):
    fig = plt.figure(figsize=(8, 6), dpi=80)

    plt.tricontourf(golem_scan[scan_id]['x'], golem_scan[scan_id]['y'], 
                    golem_scan[scan_id]['val'])
    plt.plot(golem_scan[scan_id]['x'], golem_scan[scan_id]['y'], 'ko', ms=3)
    plt.xlabel('$p$ [mPa]')
    plt.ylabel('$U_\mathrm{CD}$ [V]')
    c = plt.colorbar()
    c.set_label(scan_description)
    return fig


# ### ScanPlotLines
# 
# Plots a 2D scan as a series of lines grouped by one of scan variables (can be selected by setting `group_by` to either `x` or `y`.

# In[8]:

def ScanPlotLines(golem_scan, group_by, scan_id, scan_description):
    fig = plt.figure(figsize=(8, 6), dpi=80)

    x_axis = 'x' if group_by == 'y' else 'y';
    
    x_label = '$p$ [mPa]' if x_axis == 'x' else '$U_\mathrm{CD}$ [V]'
    lines_label = 'mPa' if x_axis == 'y' else 'V'
    
    unique = set(golem_scan[scan_id][group_by])
    unique = sorted(list(unique))
    for u_v in unique:
        data_x = []
        data_y = []
        for i in range(len(golem_scan[scan_id][x_axis])):
            if golem_scan[scan_id][group_by][i] == u_v:
                data_x.append(golem_scan[scan_id][x_axis][i])
                data_y.append(golem_scan[scan_id]['val'][i])
        
        plt.plot(data_x, data_y, 'o--', label = '{0} {1}'.format(u_v, lines_label))
    
    plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)
    plt.xlabel(x_label)
    plt.ylabel(scan_description)
    plt.grid()
    
    return fig


# ## Data Load
# 
# Let's load the data from selected shots.

# In[9]:

# Kiten scan
shot_list = range(27454, 27491, 1)
bad_shots = [27481, 27482, 27483]

# Second scan (Prague 28/06)
shot_list = [27617, 27618, 27619, 27620, 27621, 27622, 27623, 27624, 27625, 27626, 27627, 27628, 27629, 27630, 
             27633, 27634, 27635, 27636, 27637, 27638, 27639, 27640, 27641, 27642, 27643, 27644, 27645, 27646, 
             27647, 27648, 27649, 27650, 27651, 27652, 27653, 27654, 27655, 27656, 27657, 27658, 27659, 27660, 
             27661, 27662, 27663, 27664, 27665, 27666, 27667, 27668, 27669, 27670, 27671, 27672, 27673, 27674]
bad_shots = []

# Do the loading
valid_shots = Validate(shot_list, bad_shots)
data_list = [Load(s_no) for s_no in valid_shots]


# ## Data Processing
# 
# Three scans variables (`['breakdown_voltage', 'electron_temperature', 'toroidal_field']`) are selected and a scans in `pressure` and `ucd` (current drive voltage) are generated.
# 
# A second scan set is generated for requested pressure instead of real pressure.

# In[10]:

scan_vars = ['breakdown_voltage', 'electron_temperature', 'toroidal_field']

# the main scan
golem_scan = {};
for s_id in scan_vars:
    golem_scan[s_id] = GetScan2(data_list, ['pressure', 'ucd'], s_id)
    
# the auxiliary scan
golem_scan_request = {};
for s_id in scan_vars:
    golem_scan_request[s_id] = GetScan2(data_list, ['pressure_request', 'ucd'], s_id)


# ### Breakdown voltage
# We should see some minima on the Paschen curves.

# In[11]:

fig = ScanPlotLines(golem_scan, 'y', 'breakdown_voltage', '$U_\mathrm{b}$ [V]');
fig.savefig("grouped-p-U_bd.pdf", bbox_inches="tight");
fig.savefig("grouped-p-U_bd.png", bbox_inches="tight");


# ### 2D plots
# 
# * Breakdown voltage
# * Electron temperature
# * Toroidal field at the beginning of the discharge (plasma)

# In[12]:

fig = ScanPlot2D(golem_scan, 'breakdown_voltage', '$U_\mathrm{b}$ [V]');
fig.savefig("scan2d-p-U_cd-U_bd.pdf", bbox_inches="tight");
fig.savefig("scan2d-p-U_cd-U_bd.png", bbox_inches="tight");


# In[13]:

fig = ScanPlot2D(golem_scan, 'electron_temperature', '$T_\mathrm{e}$ [eV]');
fig.savefig("scan2d-p-U_cd-T_e.pdf", bbox_inches="tight");
fig.savefig("scan2d-p-U_cd-T_e.png", bbox_inches="tight");


# In[14]:

fig = ScanPlot2D(golem_scan, 'toroidal_field', '$B_\mathrm{t}$ [T]');
fig.savefig("scan2d-p-U_cd-B_t.pdf", bbox_inches="tight");
fig.savefig("scan2d-p-U_cd-B_t.png", bbox_inches="tight");


# ### Toroidal field 
# 
# Grouped by the requested pressure.

# In[15]:

fig = ScanPlotLines(golem_scan_request, 'x', 'toroidal_field', '$B_\mathrm{t}$ [T]');
fig.savefig("grouped-p-B_t.pdf", bbox_inches="tight");
fig.savefig("grouped-p-B_t.png", bbox_inches="tight");


# In[ ]: