#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Jan  4 09:40:26 2018

@author: jaroslav
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
from scipy import interpolate
import os

# Define line position
old_line = [706.52,728.13,667.82]
new_line = [504.8,587.6,501.6]

# Define constants
t_ind = [5,8,10]# time index
n = 3           # number of lines
Lr = np.array([[1030,1043],[1115,1130],[1164,1179]]) # pixel ranges for lines

# Define paths
path_tables = ['Tables/M668.0_728.33.npy','Tables/M728.33_706.76.npy',\
               'Tables/ne.npy','Tables/Te.npy']
path_fig = './figures/neTe%s.png'
diag = ['ne','Te','Te2'] # interferometry, spitzer's formula, spectrometer
path_diag = ['./basic/electron_density_','./basic/electron_temperature_',\
             './basic/electron_temperature2_']
url_base = 'http://golem.fjfi.cvut.cz/utils/data/%s/'
url_diag =['electron_density','electron_temperature','spectrometr:temperature']

# Define model function to be used to fit to the data:
def gauss(x, *p):
    A,B, mu, sigma = p
    return A*np.exp(-(x-mu)**2/(2.*sigma**2))+B

# Load spectrometer file from folder or GDB
def load_spect(shot):
    filename = './spectra/spectra' + str(shot) + '.txt'
    try:
        f =  open(filename)
    except:
        url = 'http://golem.fjfi.cvut.cz/shots/' + str(shot) + \
        '//diagnostics/Radiation/1111Spectrometer.ON/data/spectra.txt_Data_.txt'
        os.system('wget -O ./spectra/spectra' + str(shot) + '.txt '+ url\
                  + ' spectra.txt')
        f =  open(filename)
    for i in range(0,9):
        line=f.readline()
        if i==8:
            t = np.fromstring(line[19:-2], dtype=float, sep=' ')
    f.close()
    spect = np.loadtxt(filename,skiprows=11)
    return spect,t

# Load diagnostics data from folder or GDB
def load_basic(shot,name):
    filename = path_diag[diag.index(name)] + str(shot)
    try:
        f =  open(filename)
        f.close()
    except:
        url = url_base % str(shot) + url_diag[diag.index(name)]
        os.system('wget -O ' + path_diag[diag.index(name)] + str(shot) +\
                                         ' ' + url)
    return np.loadtxt(filename)

def show_spectrum(shot):
    spect,t = load_spect(shot)
    plt.figure()
    for i in t_ind:
        plt.plot(spect[:,0],spect[:,i],label = str(t[i])+' ms')
    plt.xlabel('Wavelength [nm]')
    for l in old_line:
        plt.axvline(l, color = 'r', ls = '--', label = str(l)+' nm')
    for l in new_line:
        plt.axvline(l, color = 'b', ls = '--', label = str(l)+' nm')
    plt.legend()
        
# Calculates amplitude of lines, and ratio of density and temperature lines
def line_ratio(shot):    
    spect,t = load_spect(shot)
    L = np.zeros((len(t),n))    # line amplitude
    L_err = np.zeros((len(t),n))# line amplitude error

    # Fitting
    for i in range(len(t)):
        for j in range(n):
            # Initial values
            p0 = [spect[int((Lr[j,0]+Lr[j,1])/2),i],10, \
                        spect[int((Lr[j,0]+Lr[j,1])/2),0], 1.]
            # Fit the data using a Gaussian with offset
            try:
                coeff, var_matrix = curve_fit(gauss, spect[Lr[j,0]:Lr[j,1],0],\
                        spect[Lr[j,0]:Lr[j,1],i], p0=p0)
                perr = np.sqrt(np.diag(var_matrix))
            except:
                coeff = [None,None,None,None]
            L[i,j] = coeff[0]
            L_err[i,j] = np.sqrt(perr[0]**2+perr[1]**2) #TODO nekorelovane?
            
    ne_r = L[:,0]/L[:,2]
    ne_r_err = np.sqrt(((1/L[:,2])**2)*L_err[:,0]**2+\
                       ((-L[:,0]/(L[:,2]**2))**2)*L_err[:,2]**2)
    Te_r = L[:,2]/L[:,1]  
    Te_r_err = np.sqrt(((1/L[:,1])**2)*L_err[:,2]**2+\
                       ((-L[:,2]/(L[:,1]**2))**2)*L_err[:,1]**2)
    return ne_r, Te_r, ne_r_err, Te_r_err, L

# Show specific lines
def show_lines(shot):
    spect,t = load_spect(shot)
    
    fig, ax = plt.subplots(1,3,sharey = 'row', figsize=(9,4),squeeze='False')
    for i in t_ind:
        for j in range(n):
            ax[j].plot(spect[Lr[j,0]:Lr[j,1],0],spect[Lr[j,0]:Lr[j,1],i],'.',\
              label='{:.1f}'.format(t[i]) + ' ms')
            p0 = [spect[int((Lr[j,0]+Lr[j,1])/2),i],10,\
              spect[int((Lr[j,0]+Lr[j,1])/2),0], 1.]
            coeff, var_matrix = curve_fit(gauss, spect[Lr[j,0]:Lr[j,1],0],\
              spect[Lr[j,0]:Lr[j,1],i], p0=p0)
            x = np.arange(spect[Lr[j,0],0],spect[Lr[j,1],0],0.001)
            data_fit = gauss(x, *coeff)
            ax[j].plot(x,data_fit,label = 'fit')
    
    for j in range(n):
        ax[j].set_xlabel('Wavelength [nm]')   
        ax[j].grid()
        ax[j].legend()

    ax[0].set_ylabel('Level [a.u.]')
    ax[1].set_title('#' + str(shot))
    plt.tight_layout()
    fig.subplots_adjust(wspace=0)

def neTe(shot): # TODO predelat
    spect,t = load_spect(shot)
    ne_r, Te_r, ne_r_err, Te_r_err, L = line_ratio(shot)
    netable = np.load(path_tables[0])
    Tetable = np.load(path_tables[1])
    neaxis = np.load(path_tables[2])
    Teaxis = np.load(path_tables[3])
    
    ne = np.zeros(len(t))
    ne_err = np.zeros(len(t))
    Te = np.zeros(len(t))
    Te_err = np.zeros(len(t))

    f_ne = interpolate.interp2d(Teaxis, neaxis, netable, kind='cubic')
    f_Te = interpolate.interp2d(Teaxis, neaxis, Tetable, kind='cubic')
    Teaxis = np.arange(0,100,0.05)
    neaxis = np.arange(0,4e19,1e16)
    netable = f_ne(Teaxis,neaxis)
    Tetable = f_Te(Teaxis,neaxis)
    
    for i in range(len(t)):
        try:
            netecut = f_ne(Teaxis,neaxis)
            netecut[(netable > ne_r[i]+ne_r_err[i]) | \
                    (netable < ne_r[i]-ne_r_err[i]) | \
                    (Tetable > Te_r[i]+Te_r_err[i]) | \
                    (Tetable < Te_r[i]-Te_r_err[i])] = None
            ne_ind = ~np.isnan(np.nanmean(netecut,axis=1))
            ne_max = np.max(neaxis[ne_ind])
            ne_min = np.min(neaxis[ne_ind])
            ne[i] = (ne_max+ne_min)/2
            ne_err[i] = (ne_max-ne_min)/2
            Te_ind = ~np.isnan(np.nanmean(netecut,axis=0))
            Te_max = np.max(Teaxis[Te_ind])
            Te_min = np.min(Teaxis[Te_ind])
            Te[i] = (Te_max+Te_min)/2
            Te_err[i] = (Te_max-Te_min)/2
        except:
            ne[i] = 0
            ne_err[i] = 0
            Te[i] = 0
            Te_err[i] = 0
    
    return ne,Te,ne_err,Te_err
    
def show_neTe(shot):
    spect,t = load_spect(shot)
    ne,Te,ne_err,Te_err = neTe(shot)
    neshot = load_basic(shot,'ne')
    Teshot = load_basic(shot,'Te')
    Te2shot = load_basic(shot,'Te2')
    fig, ax = plt.subplots(2,1,sharex = 'col', figsize=(9,4),squeeze='False')
    ax[0].fill_between(t, ne-ne_err, ne+ne_err,color='gray')
    ax[0].plot(t,ne,'xk-',label='#'+str(shot) + ', Helium')
    ax[0].plot(neshot[:,0]*1e3,neshot[:,1],label='interferometer')
    ax[1].fill_between(t, Te-Te_err, Te+Te_err,color='gray')
    ax[1].plot(t,Te,'xk-',label='#'+str(shot)+', Helium')
    ax[1].plot(Teshot[:,0]*1e3,Teshot[:,1],label='Spitzer\'s formula')
    ax[1].plot(Te2shot[:,0]*1e3,Te2shot[:,1],'+',label='spectrometer')
    ax[0].set(xlim = [5,20], ylim = [0,2.5e19],ylabel = '$n_e$ [m$^{-3}$]')
    ax[1].set(ylim = [0,70],ylabel = '$T_e$ [eV]', xlabel = 'Time [ms]')
    ax[0].legend()
    ax[1].legend()
    ax[0].grid()
    ax[1].grid()
    plt.tight_layout()
    fig.subplots_adjust(hspace=0)
    plt.savefig(path_fig % str(shot))
    
def show_ratio(shot):
    spect,t = load_spect(shot)
    ne_r, Te_r, ne_r_err, Te_r_err, L = line_ratio(shot)
    fig, ax = plt.subplots(2,1,sharex = 'col', figsize=(9,4),squeeze='False')
    ax[0].plot(t,ne_r,'x-')
    ax[1].plot(t,Te_r,'x-')
    ax[0].set(xlim=[5,20],ylim=[0,10],ylabel='$n_e$~668/728 [-]',grid='on')
    ax[1].set(ylim=[0,0.28],ylabel='$T_e$~728/708 [-]',xlabel='Time [ms]')
    ax[0].grid()
    ax[1].grid()
    plt.tight_layout()
    fig.subplots_adjust(hspace=0)