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

"""
 Estimate electron temperature and ion concentration based on their the spectral lines intensity. 
 
 Ions intensity is descrimed by a very simple model with minimal number of parameters:
 Expectations:
  1) that line intesity is proportional to impurity density and fnction of the temperature
  2) ratio of lines of the different ionization states of the impuritie as 
    function of the temperature is polynome of order k (k = 1 now)
  3) shape of f(T_e) of the first ionozation stage is gaussian with width a and center T0
    (for k = 1  => all other are also gaussians)
  4) this model can not exactly estimate the temperature, only something which is 
    monotonous function of the T_E => caibration from the spitzer resistivity was necassary. 
  5) it is expected that during the shot is impurity density as constant as possible. 
  
  Advantages:
  simple model => robust
  most of the spectra are explained very well even by this simple model 

  
  
  Disadvantages:
  
   OI and NI can be observed in low temperatures a  deviation from model, f(T) if this
  ions don't have a gaussian shape, it radiates strongyl alte low temperatures and the the radiation
  is alomost constant over large range of temperatures. 

 In helium shots is probably a different physic or a week he lines mess estimate of the NI density. 
 NI radiation is significantly underestimated. 
 
 Autor :Tomas Odstrcil
 Date 12.12.2012


"""

import matplotlib 
matplotlib.rcParams['backend'] = 'Agg'
from numpy import *
import matplotlib.pyplot as plt

from numpy.linalg import norm
from scipy.linalg import inv,pinv,cholesky,solve_triangular
from scipy.stats.stats import pearsonr,hmean
from scipy.sparse import block_diag
from numpy.matlib import repmat
from time import time
from pygolem_lite.modules import save_adv,load_adv
from pygolem_lite import saveconst
import numexpr as ne
from scipy.optimize import fmin_bfgs


ion_names =  ['HI','OI','OII', 'OIII', 'HeI', 'CII', 'CIII', 'NI','NII', 'NIII','M1', 'MIV' ]
element_names = ['H','O','He','C','N','M1','MIV']

n_ions = len(ion_names)
n_elemets = len(element_names)
k = 2  #polynome order

ions_index = [[0], [1,2,3],  [4], [5,6], [7,8,9], [10],[11]]
single_ion = [0,2,5,6]
multi_ion  = [1,3,4]



M1 = matrix([1,])                                       #H   0
M2 = matrix([[1,-1,0],[0,1,-1],[1,0,-1],[1,0,0]])       #O   1
M3 = matrix([1,])                                       #He  2
M4 = matrix([[1,-1],[-1,1],[1,0]])                      #C   3
M5 = matrix([[1,-1,0],[0,1,-1],[1,0,-1],[0,1,0]])       #N   4
M6 = matrix([1,])                                       #M1  5
M7 = matrix([1,])                                       #MIV 6
M  = block_diag((M1,M2,M3,M4,M5,M6,M7))
M  = M.todense()

bound_inter_i =  [1,2,3,5,6,7,8,9]
bound_inter_j =  [1,2,3,6,7,9,10,11]

bound_inter = zeros(15, dtype = 'bool')
bound_inter[bound_inter_j] = True


M_limited = (M[:,bound_inter_i])[bound_inter_j,:]
iM = array(linalg.pinv(M))


def CalibrateTemperature(T,Terr):  #basend on coeffitients from CompareSpitzerResistivity.py
    nl_coeff = loadtxt('model_constants/nonlin_tranform_coeff.txt')

    a = nl_coeff[0]
    b = nl_coeff[1]
    c = nl_coeff[2]
    d = nl_coeff[3]
    
    def T2T_eV(T):
        T = copy(T)
        ind = where(T <=c*0.95)
        T[ind] = c*0.95+c*0.05*(1-exp((T[ind]-c*0.05)))

        ind = where(T>2.9)
        T_ = (tan((-T-c)/d)-b)/a
        T_[ind] = exp(-T[ind]+2.9)*(tan((-2.9-c)/d)-b)/a
        return T_

    
    T_ = T2T_eV(T)

    Terr_u = T2T_eV(T-Terr)-T_
    Terr_d = -T2T_eV(T+Terr)+T_

   
    return T_,Terr_u,Terr_d



def EstimateT(p,pE,coeff):
    
    pE = pE#+0.1   #because than can be too low values, nonstatistical errors
    n = size(p,0)
    m = size(p,1)

    I = ones(n)
    cutM = array(M[bound_inter_j,:])
    R = dot(cutM,log(p+1e-4).T)
    m = size(R,1)

    Re = sqrt(dot(cutM**2,((pE/(p+1e-5))**2).T))
    
    R-= outer(coeff[:,0], I)
    B = coeff[:,1]
    B = tile(B,(m,1)).T
    T = average(R/B, axis = 0, weights=(B/Re)**2)
    T = sum(B*R/Re**2, axis = 0)/sum((B/Re)**2,axis = 0)

  
    B *= T[newaxis,:]
    chi2 = norm((R-B)/Re)**2/(size(R)-size(T))
    print 'chi2 temp:',chi2
    
    return T



def plotTprofiles(T0, a, density,N, T,B, coeff,p,pE,i):

    density_model = ones((100,n_elemets))
    T_model = linspace(-3,3,100)
    N_model = vstack((ones(100),T_model)).T
    c2 = CalcIntensity(T_model,density_model,N_model,T0,a,B,coeff,ions_index[i])/e
    I  = CalcIntensity(T,density,N,T0,a,B,coeff,ions_index[i])
    Resid = (p[:,ions_index[i]]-I)/pE[:,ions_index[i]]
    ax = plt.subplot(1,1,1)
    r_tot,_ = pearsonr(ravel(p[:,ions_index[i]]), ravel(I))


    #max_Resid  je aby se zabránilo přetečení
    chi2 = norm(Resid)**2/(size(Resid))
    
    max_Resid = double(amax(abs(Resid)))
    Resid/= max_Resid
    Resid **= 2
    chi2_robust = max_Resid*sum(Resid/sqrt((1/max_Resid)**2 + Resid))/(size(Resid))
    
    colour = ['b', 'r','k', 'y']
    for k in range(len(ions_index[i])):
        x = -T
        y = squeeze(p[:,(ions_index[i])[k]]/exp(density[:,i]))
        y_err = squeeze(pE[:,(ions_index[i])[k]]/exp(density[:,i]))
        noisy = y_err > y*2
        errorbar(x[~noisy],y[~noisy],y_err[~noisy], fmt=colour[k]+'.',capsize = 0,linewidth = 0.1,markersize = 0.5)

        plt.plot(-T_model, c2[:,k], colour[k],label = ion_names[(ions_index[i])[k]])
    plt.xlabel('f(T) [-]')
    plt.ylabel('intensity [a.u.]')
    leg = plt.legend(loc='upper left', fancybox=True)
    leg.get_frame().set_alpha(0.7)
    plt.title(element_names[i])
    ax.text(-1.9,1,'$\\chi^2$ = %2.1f \n$\\chi^2$ robust = %2.1f \n$r$ =  %1.2f'
        % (chi2,chi2_robust,r_tot),horizontalalignment='left',verticalalignment='bottom')

    plt.ylim(0,1.5)
    plt.xlim(-2,1)
    #plt.show()
    plt.savefig('grafy_advance/_'+element_names[i]+' '+name+'.png')
    fig.clf()
    
        
	

from scipy.spatial.distance import cdist
iM = copy(iM.T)
def CalcIntensity(T,density,N,T0,a,B,coeff,index):
    
    density[density>100] = 100
    n = size(density,0)
    F = empty((n,15), dtype= 'double')
    
    F[:,bound_inter] = dot(N,coeff.T)
    T = repmat(T, n_elemets,1).T
  
    T -= T0
    T *= T
    T *= a
    T -= B
    T += density
    
 

    F[:,~bound_inter] = T

    I = exp(dot(F,iM[:,index])-1)
    
    return I


def EstimateDensity(p,pE,T,T0,a,B,coeff):
    
    n = size(p,0)
    #pE = pE+0.1   #because than can be too low values, nonstatistical errors
    N = vstack([T**j for j in range(k)]).T
    density = zeros((n,n_elemets))

    I  = CalcIntensity(T,density,N,T0,a,B,coeff,slice(0,12))
    I[I<1e-3] = 1e-3
    r  = log(p/I+1e-3)
    re = pE/(p+1e-3)

    
    for i,ind in enumerate(ions_index):
        density[:,i] = average(r[:,ind], axis = 1, weights=1/re[:,ind])

    I = CalcIntensity(T,density,N,T0,a,B,coeff,slice(0,12))
    chi2 = norm((p-I)/pE)**2/(size(p)-size(density))

    #for i in range(n_ions):
        #plotTprofiles(T0, a, density,N, T, B,coeff,p,pE,i)
        
    #import IPython
    #IPython.embed()
    print 'chi2 dens:',chi2
    return density
    
    




def CalcDensTemp():

    C = loadtxt('model_constants/C.txt') #temperature polynome coeffitients
    A = loadtxt('model_constants/A.txt') #temperature profile gaussian widths
    T0= loadtxt('model_constants/T0.txt')#temperature profile gaussian center
    B = loadtxt('model_constants/B.txt') #temperature profile gaussian height
	
    #try:
    tvec,proj_dict = load_adv('./data/projection')
    

    p  = proj_dict
    pE = proj_dict.data_err

    #p  = load('./data/projection.npy')
    #pE = load('./data/projectionError.npy')
    #tvec= load('./data/tvec.npy')
    #except:
	#print 'Projection was not found'
	#return
    #print p.shape, type(p), p.ndim

    #exit()

    s = sum(p[:,bound_inter_i],axis = 1) > 0.01 #remove dark and very weak spectra

    shots_part = linspace(0,1,size(p,0), endpoint = False)
    p = p[s,:]
    pE = pE[s,:]
    tvec  = tvec[s] 
    shots_part = shots_part[s]
    n = size(p,0)
    
    if n == 0:
	return
    
    dp = zeros((n+1,n_ions))
    dp[1:-1,:] = abs(diff(p, axis = 0))
    dp_ = zeros_like(pE)
    for i in range(n_ions):
	x = arange(n)+0.5
	xp = arange(n+1)
	dp_[:,i] = interp(x,xp, dp[:,i])
	
	
    pE += 0.1*dp_  #error due to fast changes during integration
    pE[pE <= 1e-3] = 1e-3  #due to non-physically small errors

  
   


    n_infty = sum(pE == infty)


    def f_optim(x):
        #t = time()
	T  = x[:n]# temperature
	density = reshape(x[-n*n_elemets:],(n,n_elemets))	#logarithm dems
	
	N = vstack([T**j for j in range(k)]).T
        #t = time()
	I = CalcIntensity(T,density,N,T0,A,B,C,slice(0,12))
        #print time()-t
        size(pE)
        doF = (size(p)-size(x)-n_infty)

        Resid = ne.evaluate('((p-I)/(pE+1e-5))**2')
        chi2 = ne.evaluate('sum(Resid/sqrt(1 + Resid))')/doF
      

	b  = norm(T[T > 4])
	b += norm(T[T < -3])

	d = diff(density,1, axis = 0)  #BUG je to OK? 
	d[0,:]/=2 #allow a bigger changes at the start and end of the discharge where is the model not valid
	d[-1,:]/=2
     
	d = norm(d)/len(d)#+norm(exp(density[:,0]))+norm(exp(density[:,-1]))
	Td = diff(T,2)
	Td[Td > 0] = 0
	Td = norm(Td)
	spars = norm(ravel(exp(density)))  #suppress extremly high densities

        #print time()-t
        #print chi2 +spars*.1+lam*d*.6+Td/10+b
	return chi2 +spars*.1+lam*d*.6+Td/10+b


    
    #make a LSQ based estimates
    temp0 = EstimateT(p,pE,C)
    dens0 = EstimateDensity(p,pE,temp0,T0,A,B,C)
    #print dens0[:,element_names=='O'], temp0
    #dens0[:,array(element_names)=='O']
    


    x0 = hstack((temp0, ravel(dens0)))
    
    
    #import IPython
    #IPython.embed()
 
    lam = 10
    print 'fmin_bfgs'
    y= fmin_bfgs(f_optim, x0, gtol=1e-1,maxiter=1e6,disp=False)

    #print y
    #exit()
    #print 'fmin_bfgs finish'

    T  = y[:n]# temperature
    #print T
    N = vstack([T**j for j in range(k)]).T
    density = reshape(y[-n*n_elemets:],(n,n_elemets))	#logarithm density
    I = CalcIntensity(T,density,N,T0,A,B,C,range(12))
 
  
    
    #estimate statistical errors: 
    def RetrofitFunction(x):
	Temp  = x[:n]# temperature
	Dens = reshape(x[n:],(n,n_elemets))#logarithm density
	I = CalcIntensity(Temp,Dens,N,T0,A,B,C,range(12))
	return ravel(I)

    def CalcJacobi(x0,fun): 
	n2 = size(p)
	dx = 1e-6
	J = zeros((len(x0),n2), dtype = 'float32')
	for i in range(len(x0)):
	    xl,xr = copy(x0),copy(x0)
	    xl[i]-= dx
	    xr[i]+= dx
	    J[i,:] = (fun(xl)-fun(xr))/(2*dx)
	    
	    
	return J



    J = CalcJacobi(y,RetrofitFunction)
 
    chi2 = array([1,3,3,3,1,3,3,13,13,13,10,1])  #BUG must by set by hand
    correctedPE = copy(pE)
    correctedPE *= sqrt(chi2)
 
    
    
    err =  ravel(correctedPE)
    J/= err
    JJ= dot(J,J.T)  

    JJ.flat[::size(J,0)+1]+= 1e-3  #regularization of the sometimes singular matrix
        
    L = cholesky(JJ, lower=False,overwrite_a=True) #vytvoří to asi dolní trojúhelníkovu matici rozkladu
   
    #invert L    
    I = identity(L.shape[0])
    iL= solve_triangular(L,I , trans=0, lower=True, overwrite_b=True)

    iL = absolute(iL, out = iL)
    iL **= 2
    err = sqrt(sum(iL, 0))
   
    energy_constants=load('./data/energy_constants.npy')
    p_calib = p*energy_constants    
    P_total = sum(p_calib,1)
    
    temp_err = err[:n]
    dens_err = reshape(err[n:],(n,n_elemets))
    T_calib,T_calib_err_u,T_calib_err_d = CalibrateTemperature(T,temp_err)

    ind = argsort(P_total)[-2:]
    max_dens = average(density[ind,:],axis=0, weights=dens_err[ind,:])
    max_dens = exp(max_dens)

    for d,name in zip(max_dens, element_names):
	print 'ion %2s concentration = %1.2f'%(name,d)
	saveconst('./HistoricalAnalysis/'+name,d)
	
    flt_temp = average(T_calib[ind],weights=1/T_calib_err_d[ind]**2)

    print 'flattop temperatute %1.1f'%flt_temp
    saveconst('./HistoricalAnalysis/temperature',flt_temp)

    

    
    chi2 = zeros(n_elemets)
    for i in range(n_elemets):
	I = CalcIntensity(T,density,N,T0,A,B,C,ions_index[i])

	Resid = (p[:,ions_index[i]]-I)/pE[:,ions_index[i]]
	max_Resid = double(amax(abs(Resid)))
	Resid/= max_Resid
	Resid **= 2
	chi2[i] = max_Resid*sum(Resid/sqrt((1/max_Resid)**2 + Resid))/(size(Resid)-n+1)
    
    
    save_adv('./data/density', tvec,density,data_err=dens_err, tvec_err=1e-3)
    save_adv('./data/temperature', tvec,T_calib, data_err=[T_calib_err_d,T_calib_err_u],  tvec_err=1e-3)
    save_adv('./data/TempDensChi2', arange(n_elemets),chi2)





    #Plot basic plot
    
    try:
	shots = load('shots.npy')
	shot_index = where(abs(floor(shots)- shot_number) < 0.5)[0]
    except:
	pass
   
    fig = plt.figure(num=None, figsize=(8, 7), dpi=80, facecolor='w', edgecolor='k')
    ax = fig.add_subplot(311)

    I = CalcIntensity(T,density,N,T0,A,B,C,range(12))
    plt.errorbar(arange(size(p)), ravel(p.T), yerr=ravel(pE.T))
    plt.plot(ravel(I.T)) 

    
    chi2doF = norm((p-I)/pE)**2/size(p-size(y))
    ax.set_title('Spectral data retrofit')
    

    T_est = EstimateT(p,pE,C)
    d_est = EstimateDensity(p,pE,T_est,T0,A,B,C)
    I = CalcIntensity(T_est ,d_est,N,T0,A,B,C,range(12))
    plt.plot(ravel(I.T), 'r:')
    ax.axis('tight')
    ax.set_ylabel('intensity')
    ax.set_ylim(0,amax(p)*1.2)
    ax.text(2/2,amax(p)*0.5,'$\chi^2$/doF = %.2f'%chi2doF,horizontalalignment='left',verticalalignment='bottom')
    for i,name in enumerate(ion_names):
	ax.text(i*n+1,amax(p)*1.05,name,horizontalalignment='left',verticalalignment='bottom')
    
    
    ax = fig.add_subplot(312)  
    plt.errorbar(tvec*1e3, -T, temp_err,markersize = 0.5)
    plt.plot(tvec*1e3, -(T_est[:n]),':')
    try:	
	temperature_full  = load('temperature.npy')	
	plt.plot(shots[shot_index]-shot_number, -temperature_full[shot_index], '--')
    except:
	pass
    ax.set_ylim(-4,3)
    ax.set_ylabel('pseudo temperature')
   
    ax = fig.add_subplot(313)  
    plt.semilogy( tvec*1e3,exp(density))
    try:
	density_full = load('density.npy')
	density_full+= B[:,newaxis]
	plt.semilogy(shots[shot_index]-shot_number,exp(density_full[:,shot_index].T), '--')
    except:
	pass
    #ax.axis('tight')

    ax.set_ylim(1e-3,10)
    ax.set_ylabel('Density')
    ax.set_xlabel('t [ms]')

    fig.savefig('./graphs/proj_retrofit.png')
    fig.clf()
    return 

    
    

def plotDensTemp():
    
    tvec,density = load_adv('./data/density')
    dens_err = density.data_err
  
    
    tvec,T = load_adv('./data/temperature')
    temp_err = T.data_err
    temp_err = [temp_err[:,0] ,temp_err[:,1]]
    _,chi2doF = load_adv('./data/TempDensChi2')




    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=(10, 10), dpi=80, facecolor='w', edgecolor='k')
    plt.subplots_adjust(hspace=0, wspace = 0)
    
    for i in range(0,n_elemets):
	ax = fig.add_subplot(n_elemets,1,i+1)
	if i == 0: 
	    ax.set_title('Relative ions concentration')
	
	ax.xaxis.set_major_formatter( plt.NullFormatter() )
	ax.yaxis.set_major_formatter( MyFormatter() )
	

	y = exp(density[:,i])
	y_err_up   =  expm1( dens_err[:,i])*y
	y_err_down = -expm1(-dens_err[:,i])*y
	y_med = average(y, weights=abs(y_err_up))
	y_err_med = (hmean(abs(y-y_med)+1e-3)+hmean(abs(y_err_up)))/2
	plt.errorbar(tvec*1e3,y,yerr=[y_err_down,y_err_up],markersize = 0.5, label = element_names[i])
	ax.yaxis.grid(True)

	ax.text(.05,.2,'$\\chi^2/doF = %2.1f$ \n $n = %1.2f\pm%1.1g$'% (chi2doF[i],y_med,y_err_med),horizontalalignment='left',verticalalignment='bottom',transform=ax.transAxes)
	ax.set_ylim(0,max(0.1,amax(y))*1.5)

	leg = plt.legend(loc='upper left', fancybox=True)
	leg.get_frame().set_alpha(0.5)


    ax.xaxis.set_major_formatter( plt.ScalarFormatter() )

  
    ax.set_xlabel('t [ms]')
    fig.savefig('./graphs/density.png',bbox_inches='tight')
    fig.clf()



    fig = plt.figure(num=None, figsize=(10, 3), dpi=80, facecolor='w', edgecolor='k')
    ax = fig.add_subplot(111)
    ax.set_title('Estimated temperature')
    plt.errorbar(tvec*1e3, T, yerr=temp_err,markersize = 0.5)
    ax.set_ylabel('$T_e$ [eV]')
    ax.set_xlabel('t [ms]')
    ax.set_ylim(0,80)
    ax.yaxis.grid(True)
    fig.savefig('./graphs/temperature.png',bbox_inches='tight')
    plt.close()
    
    
    try:
	fig = plt.figure(num=None, figsize=(10, 3), dpi=80, facecolor='w', edgecolor='k')
	ax = fig.add_subplot(111)
	data = load('./electronTemperatures/'+str(shot_number)+'/electron_temperature.npz')
	temp = data['data']*data['scale']
	
	tvec_spitz= linspace(data['t_start'], data['t_end'], len(temp))
	try:
	    plasma_start = loadtxt('./electronTemperatures/'+str(shot_number)+'/PlasmaStart')
	    plasma_end =   loadtxt('./electronTemperatures/'+str(shot_number)+'/PlasmaEnd')
	except:		
	    pass
	    #plasma_start =  amin(tvec[isfinite(temp)])
	    #plasma_end   =  amax(tvec[isfinite(temp)])
	if shot_number > 9300:
	    temp*= 2
	    
	else:
	    tvec_spitz/= 1e3
	    plasma_start-= 0.5
	    plasma_start/= 1000
	    plasma_end/=1000
	
	ax.set_title('Spitzer temperature')
	plt.plot(tvec_spitz*1e3, temp)

	ax.set_ylabel('$T_e$ [eV]')
	ax.set_xlabel('t [ms]')
	ax.set_ylim(0,80)
	ax.set_xlim(plasma_start*1000,plasma_end*1000)

	fig.savefig('./graphs/temperature'+str(shot_number)+'_spitz.png',bbox_inches='tight')
	plt.close()
	print 'Te plotted'
    except:
	print 'no electron temperature'
	pass