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

###########################################################################################################

#############   Algorithm which from ion intensities calculate the coefitients to estimate electron t
#############   temperature and impurity density. 

###########################################################################################################

names =  ['HI','OI','OII', 'OIII', 'HeI', 'CII', 'CIII', 'NI','NII', 'NIII','Mystery1', 'CIV+NIV+OIV' ]
ion_names = ['H','O','He','C','N','M1','M2']
ions_index = [[0], [1,2,3],  [4], [5,6], [7,8,9], [10],[11]]

from numpy import *
from numpy.fft import fft,rfft
from matplotlib.pyplot import *
from numpy.linalg import norm,lstsq
from scipy.linalg import qr,inv,pinv,eig,cholesky,svd,cho_factor,solve_triangular
from scipy.stats.stats import pearsonr
from scipy.sparse import block_diag
from scipy.optimize import fmin_cg,leastsq,fmin_ncg,fmin_bfgs,fmin_powell,anneal
from scipy.stats.mstats import mquantiles
from scipy.stats import nanmedian
from numpy.matlib import repmat
from time import time
from scipy.signal import medfilt,order_filter


#Load data
name = 'vylepseny odhad y0'

n_ions = 12
n_elemets = 7
k = 2  #polynome order

def _sparseness(x):
    """Hoyer's measure of sparsity for a vector"""
    x = np.ravel(x)
    sqrt_n = np.sqrt(len(x))
    spars = (sqrt_n - norm(x, 1) / norm(x)) / (sqrt_n - 1)
    if np.isnan(spars):
	return 1
    else:
	return spars


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]
M_limited = (M[:,bound_inter_i])[bound_inter_j,:]

iM = linalg.pinv(M)



eigval,eigvec = eig(M_limited,left = True, right = False)
null_ind = abs(eigval) < 1e-6  
null_space = eigvec[:,null_ind]    #optimalizaion constraion 
    
#load the data from old shots
try:
    data = load('projection.npz')
    p = data['projections']
    pE = data['projectionsErr']
    shots = data['shotnum']

except:
	
    p = zeros((1,12))
    pE = zeros((1,12))
    shots = zeros(1)
    for shot in xrange(7136,10300):
	try:
	    c = load('./data/projection'+str(shot)+'.npy')
	    ce = load('./data/projectionError'+str(shot)+'.npy')
	    
	except:
	    continue
	n_specta = size(c,0)  
	shot_space = shot + linspace(0,1,n_specta, endpoint = False)
	p = vstack((p,c))
	pE = vstack((pE,ce))
	shots = hstack((shots,shot_space))

    
    s = sum(p[:,(1,2,3,5,6,7,8,9)],axis = 1) > 0.1 #remove dark and very weak spectra

    p = p[s,:]
    pE = pE[s,:]
    shots = shots[s]
    savez_compressed('projection', projections = p, projectionErr = pE, shotnum = shots)




pE[pE <= 0] = infty
pE[pE < 0.01] = 0.01  #změna intenzity během vývoje záření je mnohem větší 

#CALCULATE INITIAL GUESS OF THE TEMPERATURE
n_ions = size(p,1)
p_ = p[:,(1,2,3,5,6,7,8,9)]
pE_ = pE[:,(1,2,3,5,6,7,8,9)]


wp = (p_ == 0.0 ) | (isnan(p_))

p_[wp] = 1

divE = pE_/p_
#log ratios
R = vstack((p_[:,0]/p_[:,1],p_[:,1]/p_[:,2],p_[:,0]/p_[:,2],p_[:,3]/p_[:,4],p_[:,4]/p_[:,3], p_[:,5]/p_[:,6],p_[:,6]/p_[:,7],p_[:,5]/p_[:,7]))
R = log(R)
#ratio log error
Re = vstack((hypot(divE[:,0],divE[:,1]),hypot(divE[:,1],divE[:,2]),hypot(divE[:,2],divE[:,0]),hypot(divE[:,3],divE[:,4]),hypot(divE[:,3],divE[:,4]),
	    hypot(divE[:,5],divE[:,6]),hypot(divE[:,6],divE[:,7]),hypot(divE[:,7],divE[:,5])))
#wrong points
wp = vstack((wp[:,0]|wp[:,1],wp[:,1]|wp[:,2],wp[:,2]|wp[:,0],wp[:,3]|wp[:,4],wp[:,3]|wp[:,4],wp[:,5]|wp[:,6],wp[:,6]|wp[:,7],wp[:,7]|wp[:,5]))
rnames  = ('OI/OII','OII/OIII','OI/OIII','CII/CIII','CIII/CII', 'NI/NII','NII/NIII','NI/NIII')

Re[wp] = inf  

n = size(R,1)
m = size(R,0)
print ' m ',m,' n ',n, ' k ', k


#estimate polynom coeffitients of the temperature polynome
def CalcCoeff(R,Re,T,Te):
    
    coeff = zeros((m,k))
    coeff_error = zeros((m,k))
    N = vstack([T**i for i in range(k)]).T

    q,r = qr(N,mode =  'economic')
    F = matrix(dot(inv(r),q.T))
	
	
    total_res = 0
    for (f,e,x,xe) in zip(R,Re,coeff,coeff_error):

	x[:],res,_,_ =  lstsq(dot(diag(1/e/Te),N),f/e/Te)
	chi2 = float(res/len(e))
	q,r = qr(dot(diag(1/e),N),mode =  'economic')
	F = matrix(dot(inv(r),q.T))
	xe[:] = sqrt(diag( chi2*(F*F.T)))
	total_res+=res
    

    print total_res/(size(R)-size(T)-size(coeff))
    return coeff,coeff_error



#def CalcT(R,Re,coeff):
    #I = matrix(ones(n))

    #R = R-dot(matrix(A).T,I)
    #R = array(R)
    #B = array(B)
    #T = zeros(n)
    #total_res = 0
    #for i in xrange(n):  #jak to udělat vektorově? 
	#x = R[:,i]/Re[:,i]
	#y = B/Re[:,i]

	#y/=norm(y)

	#T[i] = dot(x,y).T
	#total_res += norm(x-T[i]*y)**2
	
    #print total_res/size(R)
    #return T
    
    
Te = ones(n)
Te[Re[2,:] == inf] = inf
I = (ones(n))

iterator = 0

def f_optim1(x):
    global iterator
    iterator+=1
    x = double(x)
    coeff = reshape(x[:k*m], (m,k))

    T = x[-n:]

    a = mean(coeff[:,0])-1
    b = mean(coeff[:,1])
    T = (1+a)+T*b
    coeff[:,0]-=mean(coeff[:,0])
    coeff[:,1]/=mean(coeff[:,1])


    N = vstack([T**i for i in range(k)]).T
    Resid  = (R-dot(coeff,N.T))/Re

    #robust least squares estimator
    Resid **= 2
    chi2 = sum(Resid/sqrt(1+Resid))/(size(Resid)-size(x))
    
    ##boundary condition 
    boundary = norm(dot(coeff.T,null_space))**2   
    if iterator%1000 == 0:	    
	print chi2,  boundary
    if iterator%200000 == 0:	    
	for i in  range(m):
	    normal = (Re[i,:] != inf)
	    plot(T[normal],R[i,normal],'.',markersize = 0.5)
	    plot(T[normal],dot(coeff, N.T)[i,normal].T,'.',markersize = 0.3,label = rnames[i])
	    xlabel('f(T)')
	    ylabel('ln('+rnames[i]+')')
	    xlim(-1.5,4)
	    ylim(-2,6)		
	    savefig('grafy_advance/_'+str(i)+name+'.png')
	    clf()

    return chi2 + boundary



try:
    #x
    x = load('X'+name+'.npy')
except:
    print 'temperature initial guess'
    T = R[2,:]  #počáteční odhad

    coeff,_= CalcCoeff(R,Re,T,Te)  

    #využívám volnosti v definici, abych to zafixoval  mena A = 0, mean B = -1, a ostatní na nulu
    a = mean(coeff[:,0])-1
    b = mean(coeff[:,1])
    T = (1+a)+T*b
    coeff[:,0]-=mean(coeff[:,0])
    coeff[:,1]/=mean(coeff[:,1])
    x0 = hstack((ravel(coeff),T))
    x = fmin_powell(f_optim1, x0)
    
    coeff = reshape(x[:k*m], (m,k))
   
    T = x[k*m:]
    N = vstack([T**i for i in range(k)]).T

    Resid = R-dot(coeff,N.T)
    Resid /=Re

    wrong_shots    = median(abs(Resid), axis = 0) > 4   #hodně nadhodnocené aby tam něco i zbylo 
    print 'part of the wrongly identified shots', sum(wrong_shots)/float(n)
    T[wrong_shots] = median(T[~wrong_shots])    
    T = medfilt(T,11)
    x[k*m:] = T
    
    
    x = fmin_powell(f_optim1, x)
    save('X'+name,x)
    
    


coeff = reshape(x[:k*m], (m,k))
T = x[k*m:]



a = mean(coeff[:,0])-1
b = mean(coeff[:,1])
T = (1+a)+T*b
coeff[:,0]-=mean(coeff[:,0])
coeff[:,1]/=mean(coeff[:,1])

N = vstack([T**i for i in range(k)]).T
Resid  = R-dot(coeff,N.T)
Resid /= Re

#robust least squares estimator
Resid[abs(Resid) < 0.01] = 0.01
Resid **= 2
chi2 = sum(1/sqrt(1/Resid+1/Resid**2))/(size(Resid)-size(x))

print 'vypočtené chi2 ', chi2




Resid[Resid == 0] = nan   #body s nekonečnou chybou



_,coeff_err =  CalcCoeff(R,Re,T,Te)


###plot resultes
#chi2_tot = 0
#ind_sort = argsort(T)
#for i in  range(m):
    #normal = (Re[i,:] != inf) #& (Re[2,:] != inf)

    #ind_sort_i = ind_sort[normal]
    #errorbar(-T[ind_sort_i],R[i,ind_sort_i],Re[i,ind_sort_i], fmt='.',capsize = 0,linewidth = 0.1,markersize = 0.5)
    #plot(-T[ind_sort_i],dot(coeff, N.T)[i,ind_sort_i].T,'-' ,markersize = 0.3,label = rnames[i])

    #xlabel('f(T)')
    #ylabel('ln('+rnames[i]+')')
    #xlim(-4,1.5)
    #ylim(-2,6)
    
    #chi2 = norm(((R-dot(coeff, N.T))/Re)[i,:])**2/size(R[i,:])
    #title('chi2 = %1.3f A = %1.3f$\pm$ %1.3f B = %1.3f$\pm$%1.3f '%(chi2,coeff[i,0],coeff_err[i,0],coeff[i,1],coeff_err[i,1]))
    #legend(loc = 'best')
    #chi2_tot+=chi2
    #savefig(str(i)+name+'.pdf')
    #clf()
#print 'chi2_tot',chi2_tot/m






#######Druhý stupeň optimalizace, odhad hustoty##################
#TODO normovat dlouhodobý harmonický průměr hustoty na jedničku (log na nulu)?? nebo to udělat až v postporocessingu?? 

name = 'density estimate '
print 'name: ', name
single_ion = [0,2,5,6]
multi_ion = [1,3,4]

N = vstack([T**i for i in range(k)]).T
bound_inter = zeros(15, dtype = 'bool')
bound_inter[bound_inter_j] = True


#prepare initila guess for the optimalization

density0 = (p[:,[0,2,4,5,8,10,11]]+1e-3)
density0 = order_filter(density0, ones((101,1)), 90)
density0 = log(density0)

T0 = array([1,0.5,0.5,-0.25,-0.4,-0.5,-0.5])
a0 = -ones(n_elemets)


def CalcIntensity(T,density,N,T0,a,coeff,index):
 
    density[density>10] = 10
    coeff[:,0]-=mean(coeff[:,0])
    coeff[:,1]/=mean(coeff[:,1])
    n = size(density,1)
    F = empty((15,n), dtype= 'double')
    F[bound_inter,:] = dot(coeff,N.T)

    T =T[:,newaxis]- T0[newaxis,:]

    T **= 2
    T *= a
    T = T.T
    T += density

    F[~bound_inter,:] = T
    C = exp(iM[index,:]*F-1).T
    C[C > 1e5] = 0
    return array(C, copy = False)
  


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


    #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)

	plot(-T_model, c2[:,k], colour[k],label = names[(ions_index[i])[k]])
    xlabel('f(T) [-]')
    ylabel('intensity [a.u.]')
    leg = legend(loc='upper left', fancybox=True)
    leg.get_frame().set_alpha(0.7)
    title(ion_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')

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

    for k in range(len(ions_index[i])):
	ax = subplot(1,1,1)

    	r,_ = pearsonr(p[:,(ions_index[i])[k]] , C[:,k])
    	x = C[:,k]
	y = p[:,(ions_index[i])[k]] 
	y_err = pE[:,(ions_index[i])[k]] 
	noisy = y_err > y*2
	chi2 = norm((x-y)/y_err)**2/(size(Resid))

	plot(x[~noisy],y[~noisy],'.'+colour[k],markersize = 0.5)
	plot(arange(-10,10),arange(-10,10), colour[k])
	ax.text(0.2,0.8,'$\\chi^2$ = %2.1f \n $r$ =  %1.2f'% (chi2,r),horizontalalignment='left',verticalalignment='bottom')
	title(names[(ions_index[i])[k]])
	ylim(0,1)
	xlim(0,1)
	#show()
	savefig('grafy_advance/correl_'+names[(ions_index[i])[k]]+' '+name+'.png')

        clf()





density_all = zeros((n_elemets,n))
for i,index in enumerate(ions_index):
    #if i != 0:
	#continue
    density0_ = density0[:,i]
    print 'iont '+ion_names[i]

    a0_ = a0[i]
    T0_ = T0[i]
    p_  = array(p[:,index], ndmin = 2)
    pE_  = array(pE[:,index], ndmin = 2)
    #pE_[pE_ > p_] = inf   #exteramly bad points
    name_ = [names[j] for j in index]

    x0 = hstack((T0_,a0_,ravel(density0_)))
    x0 = single(x0)
    #n_elemets = 1
    n_ions = len(index)
	

    def f_optim2(x):
	x = double(x)
	global iterator,t0

	T0 = x[:1] #basic temperature shift
	a  = x[1:2]#width of the gaussian
	density = reshape(x[2:],(n,1)).T#logarithm density
	#print shape(density), shape(p)
	C = CalcIntensity(T,density,N,T0,a,coeff,index)
	#print where(~isfinite(C))
	#exit()
	Resid = (p_-C)/pE_
	max_Resid = double(amax(abs(Resid)))
	#max_Resid  je aby se zabránilo přetečení
	Resid/= max_Resid
	#print amax(C)
	#exit()
	Resid **= 2
	chi2 = max_Resid*sum(Resid/sqrt((1/max_Resid)**2 + Resid))/(size(Resid))

	boundary  = norm(a[a > -0.5]+0.5)
	boundary += norm(a[a < -3]+3)**2
	boundary += norm(T0[T0 > 2]-2)**2
	
	d = diff(exp(density),1)

	new_shots = diff(floor(shots)) > 0

	d[:,new_shots] = 0
	
	d = norm(d)
	if iterator%1000 == 0:
	    t2 = time()
	    print chi2,d,boundary
	    t0 = t2
	if iterator%10000 == 0:
	    plot(exp(density.T), linewidth = 0.1)
	    title(ion_names[i])
	    ylim([0,5])
	    savefig('grafy_advance/density'+ion_names[i]+'.pdf')
	    clf()
	    print 'a: %2.2f, T0: %2.2f' %(a,T0)

	    print 'profile plots', ion_names[i]
	    density_all[i,:] = density
	    a_all  = zeros(n_elemets)
	    T0_all = zeros(n_elemets)
	    a_all[i] = a
	    T0_all[i] = T0
	    plotTprofiles(T0_all, a_all, density_all, T, coeff,p,pE,i)
	    save('y'+name+str(i)+'.npy',x)


	iterator+=1
	return chi2+d*lam+    boundary

    try:
	
	y = load('y'+name+str(i)+'.npy')
    except: 
	print 'density initial guess'

	lam = 0.01
	iterator = 0

	L = len(x0)

	t0 = time()
	y = fmin_powell(f_optim2, x0)
	try:
	    y = fmin_powell(f_optim2, x0)
	except:
	    try:
		y = load('y'+name+str(i)+'.npy')
	    except:
		pass


	lam = 0.2
	#try:
	y = fmin_powell(f_optim2, y)
	#except:
	    #y = load('y'+name+str(i)+'.npy')
	
	save('y'+name+str(i)+'.npy',y)










#######Třetí stupeň optimalizace, nový odhad teploty ##################


n_elemets = len(ion_names)

for i in range(n_elemets):
    y = load('y'+name+str(i)+'.npy')
    T0[i] = y[:1] #basic temperature shift
    a0[i]  = y[1:2]#width of the gaussian
    #print shape(y[2:]),n
    density0[:,i] = reshape(y[2:],(n,1)).T#logarithm density
density0 = density0.T
density0[density0<-3] = -3





name = '3. cyklus'
print 'name: ', name


x0 = hstack((ravel(coeff),T0,a0,T))
x0 = single(x0)
    

def f_optim3(x):
    x = double(x)
    global iterator,t0
    coeff = reshape(x[:k*m], (m,k))

    T0 = x[k*m:k*m+n_elemets] #basic temperature shift
    a  = x[k*m+n_elemets:k*m+2*n_elemets]#width of the gaussian
    T = x[-n:]
    N = vstack([T**j for j in range(k)]).T

    A = mean(coeff[:,0])-1
    B = mean(coeff[:,1])
    T = (1+A)+T*B
    coeff[:,0]-=mean(coeff[:,0])
    coeff[:,1]/=mean(coeff[:,1])
    #exit()
    C = CalcIntensity(T,density0,N,T0,a,coeff,range(12))

    
    Resid = (p-C)/pE
    max_Resid = double(amax(abs(Resid)))
    #max_Resid  je aby se zabránilo přetečení
    Resid/= max_Resid
    #exit()

    Resid **= 2
    chi2 = max_Resid*sum(Resid/sqrt((1/max_Resid)**2 + Resid))/(size(Resid)-size(x)-size(density0))

    #chi2 = norm(Resid)**2/(size(Resid)-size(x))
    boundary  = norm(a[a > -0.5]+0.5)
    boundary += norm(a[a < -3]+3)**2
    boundary += norm(T0[T0 > 2]-2)**2
    boundary += norm(dot(coeff.T,null_space))**2   


    
    if iterator%1000 == 0:
	t2 = time()
	print chi2,boundary
	t0 = t2
    if iterator%50000 == 0:
	plot(-T, linewidth = 0.1)
	ylim(-2, None)
	savefig('grafy_advance/T '+name+'.pdf')
	clf()
	for i in range(n_elemets):
	    plotTprofiles(T0, a, density0, T, coeff,p,pE,i)
	save('y'+name+'.npy',x)

    iterator+=1
    return chi2 + boundary


try:
    y = load('y'+name+'.npy')
except: 
    print 'temperatute second guess'

    lam = 0.01
    iterator = 0

    L = len(x0)

    A = eye(L, dtype = 'single')
    A = A[::-1,:]
    t0 = time()
    try:
	y = fmin_powell(f_optim3, x0, direc = A)
    except:
	y = load('y'+name+'.npy')

    lam = 0.01
    try:
	y = fmin_powell(f_optim3, y, direc = A)
    except:
	y = load('y'+name+'.npy')
    
    save('y'+name+'.npy',y)















#######Poslední stupeň optimalizace, globání optimalizace ##################


#final stage -- full optimalization
#prepare initial guess

#name = '_4. cyklus'

x = load('y'+name+'.npy')

#stary
#coeff = reshape(x[:k*m], (m,k))
#T0 = x[k*m:k*m+n_elemets] #basic temperature shift
#a  = x[k*m+n_elemets:k*m+2*n_elemets]#width of the gaussian
#T = x[-n:]
#print shape(x), shape(density0)
x0 = r_[x,ravel(density0.T)]
#print shape(density0), (n,n_elemets)
#novy
#coeff = reshape(x[:k*m], (m,k))
#T0 = x[k*m:k*m+n_elemets] #basic temperature shift
#a  = x[k*m+n_elemets:k*m+2*n_elemets]#width of the gaussian
#T  = x[-(1+n_elemets)*n:-n_elemets*n]# temperature
#print shape(x[-n*n_elemets:]), n_elemets*n
#density = reshape(x[-n*n_elemets:],(n,n_elemets)).T#logarithm density
#exit()
name = '4. cyklus 2 norm'
print 'name: ', name


x0 = single(x0)
    

def f_optim4(x):
    x = double(x)
    global iterator,t0
    coeff = reshape(x[:k*m], (m,k))

    T0 = x[k*m:k*m+n_elemets] #basic temperature shift
    a  = x[k*m+n_elemets:k*m+2*n_elemets]#width of the gaussian
    T  = x[-(1+n_elemets)*n:-n_elemets*n]# temperature
    density = reshape(x[-n*n_elemets:],(n,n_elemets)).T#logarithm density

    N = vstack([T**j for j in range(k)]).T

    C = CalcIntensity(T,density,N,T0,a,coeff,range(12))

    Resid = (p-C)/pE
    max_Resid = double(amax(abs(Resid)))
    #max_Resid  je aby se zabránilo přetečení
    #Resid/= max_Resid

    #Resid **= 2
    #chi2 = max_Resid*sum(Resid/sqrt((1/max_Resid)**2 + Resid))/(size(Resid)-size(x))
    chi2 = norm(Resid)**2/(size(p)-size(x))
    #boundary  = norm(a[a > -0.5]+0.5)
    boundary  = norm(a[a < -3]+3)**2
    boundary += norm(T0[T0 > 2]-2)**2
    boundary += norm(dot(coeff.T,null_space))**2   
    boundary += abs(mean(coeff[:,0]))*10 + abs(mean(coeff[:,1])-1)*10

    d = diff((density),1, axis = 1)
    new_shots = diff(floor(shots)) > 0
    d[:,new_shots] = 0
    d = norm(d)
    
    spars = norm(ravel(exp(density)/10))/100
    density[density<-5] = -5
    
    
    if iterator%1000 == 0:
	t2 = time()
	print chi2,d,spars,boundary,mean(coeff[:,0]),mean(coeff[:,1])-1
	t0 = t2
    if iterator%50000 == 0:
	plot(-T, linewidth = 0.1)
	ylim(-2, None)
	savefig('grafy_advance/T '+name+'.pdf')
	clf()
	#print 'profile plots'
	for i in range(n_elemets):
	    plot(exp(density[i,:]), linewidth = 0.1)
	    ylim(0, 10)
	    title(ion_names[i])
	    savefig('grafy_advance/dens'+str(i)+name+'.pdf')
	    clf()
	    
	    plotTprofiles(T0, a, density, T, coeff,p,pE,i)
	save('y'+name+'.npy',x)


    iterator+=1
    return chi2 + boundary+lam*d+spars

try:
    ss
    y = load('y'+name+'.npy')
except: 
    print 'final step temperature+density'

    lam = 0.06
    iterator = 0
    L = len(x0)

    print L
    A = eye(L, dtype = 'single')
    A = A[::-1,:]
    t0 = time()
    y = fmin_powell(f_optim4, x0, direc = A)

    try:
	x0 = load('y'+name+'.npy')

	y = fmin_powell(f_optim4, x0, direc = A)
    except:
	y = load('y'+name+'.npy')

    lam = 0.07
    try:
	y = fmin_powell(f_optim4, y, direc = A)
    except:
	y = load('y'+name+'.npy')
    
    save('y'+name+'.npy',y)
    
    
x = load('y'+name+'.npy')





#estimate statistical errors: 


def RetrofitFunction(x):
    coeff = reshape(x[:k*m], (m,k))
    T0 = x[k*m:k*m+n_elemets] #basic temperature shift
    a  = x[k*m+n_elemets:k*m+2*n_elemets]#width of the gaussian
    T  = x[-(1+n_elemets)*n:-n_elemets*n]# temperature
    density = reshape(x[-n*n_elemets:],(n,n_elemets)).T#logarithm density
    C = CalcIntensity(T,density,N,T0,a,coeff,range(12))
    return ravel(C)




def CalcJacobi(x0,fun):  #Bude se tu muset udělat něco jako pseudoimverze?? 
    n2 = size(p)
    N = len(x0)
    dx = 1e-6
    J = zeros((N,n2), dtype = 'float32')
    for i in range(N):
	xl = copy(x0)
	xl[i]-= dx
	xr = copy(x0)
	xr[i]+= dx
	J[i,:] = (fun(xl)-fun(xr))/(2*dx)
	
	
	print i, N
    return J


#first step - calc jacobi matrix
try:
    J = load('J.npy', mmap_mode='r')
except:
    J = CalcJacobi(x0,RetrofitFunction)
    save('J.npy',J)


try:
    JJ = load('JJ.npy',mmap_mode = 'r')
except:
    N = size(J,0)
    JJ = zeros((N,N), dtype = 'float32')
    #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
    intervals = array_split(range(N), 100)
    for i,inter in enumerate(intervals): #teoreticky by se tu dalo využít symetrie matice JJ 
	print i, 100
	A = dot(J,J[inter,:].T)
	JJ[:,inter] = A

    del J
    save('JJ.npy',JJ)
    
    


#calc cholesky decomposition of the L*L.H = J.T*J
try: 
    #hh
    L = load('ires_L.npy', mmap_mode='r')
except:
    N = size(J,0)

    JJ.flat[::N+1]+= 1e-4  #regularization of the singular matrix

    JJ = JJ[-(1+n_elemets)*n:, -(1+n_elemets)*n:]

    L = cholesky(JJ, lower=False,overwrite_a=True) #vytvoří to asi dollní trojúhelníkovu matici rezkladu
    save('ires_L.npy',L)
   
#invert L
try:
    iL = load('iL.npy', mmap_mode='r')
except:
    
    N = L.shape[0]
    I = identity(N, dtype = 'float32')
    intervals = array_split(range(N), 100)
    #I= solve_triangular(L,I , trans=0, lower=True, overwrite_b=True)

    for i,inter in enumerate(intervals):
	print i, 100
	I[:,inter] = solve_triangular(L,I[:,inter] , trans=0, lower=True, overwrite_b=True)
    save('iL', I)
    
try:
    err = load('err.npy')
except:
    iL = absolute(iL, out = iL)
    iL **= 2
    err = sqrt(sum(iL, 0))
    save('err.npy',err)


coeff = reshape(x[:k*m], (m,k))
T0 = x[k*m:k*m+n_elemets] #basic temperature shift
a  = x[k*m+n_elemets:k*m+2*n_elemets]#width of the gaussian
T  = x[-(1+n_elemets)*n:-n_elemets*n]# temperature
density = reshape(x[-n*n_elemets:],(n,n_elemets)).T#logarithm density





N = vstack([T**j for j in range(k)]).T

C = CalcIntensity(T,density,N,T0,a,coeff,range(12))



#save final resultes
savetxt('model_constants/B.txt',mean(density, axis = 1),fmt = '%.4f')
savetxt('model_constants/C.txt',coeff,fmt = '%.4f')
savetxt('model_constants/A.txt',a,fmt = '%.4f')
savetxt('model_constants/T0.txt',T0,fmt = '%.4f')



#BUG!!!!
err/= 3
#BUG!!!!

errT = err[-(1+n_elemets)*n:-n_elemets*n]
errDensity = reshape(err[-n*n_elemets:],(n,n_elemets)).T




for i in range(n_elemets):
    y = exp( density[i,:])
    y_err_up   =  (exp( errDensity[i,:])-1)*y
    y_err_down = -(exp(-errDensity[i,:])-1)*y


    errorbar(arange(n), y,yerr=[y_err_down,y_err_up],markersize = 0.5)
    title(ion_names[i])
    ylim(0,10)
    show()
    
errorbar(arange(n), -T, errT,markersize = 0.5)
ylim(-3,1.5)
show()




coeffE = reshape(err[:k*m], (m,k))
T0E = err[k*m:k*m+n_elemets] #basic temperature shift
aE  = err[k*m+n_elemets:k*m+2*n_elemets]#width of the gaussian
TE  = err[-(1+n_elemets)*n:-n_elemets*n]# temperature
densityE = reshape(err[-n*n_elemets:],(n,n_elemets)).T#logarithm density
#print shape(shots), shape
x = arange(len(T))
save('shots', shots)

save('density', density)
save('temperature', T)