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

"""    
    Algorithm for integrated data analyze of the visible light signal from GOLEM tokamak. 
    Many different diagnotics with different wavelength resolution, time resolution, time shift error, 
    and signal error is availible on the tokamak.This algirithm puts them together and finds simplest 
    solution which corespondes to the measured signal. Usually most of the wavelength information is provided 
    by spectrometer and the higth freqency part is provided by fast cameras and photodiodes. 
    The problem is solved as least squares fit of the data by nonparametrical model sim maximal simplisity. 
    
	Autor: Tomas Odstrcil
	date:12.12.2012
	
	

"""    

#TODO jak jsou ošetřená přepálená data z foťáku?? dát tam inf???
from numpy import *
from matplotlib.pyplot import *
from  SpectrometerControl  import *
from scikits.sparse.cholmod import cholesky, analyze, cholesky_AAt,CholmodError
from scipy.signal import fftconvolve,gaussian
from numpy import linalg
import time
from scipy import sparse
from numpy.matlib import repmat
from scipy.sparse import *
from  scipy.linalg import  norm, inv,qr,pinv

from pygolem_lite.modules import save_adv,load_adv
from CalcIonProjections import Resample,upsample_smooth,GapsFilling

names =  ('HI','OI','OII', 'OIII', 'HeI', 'CII', 'CIII', 'NI','NII', 'NIII','Mystery1', 'CIV+NIV+OIV' )
styles = ('b-','r-','r--', 'r-.' , 'g-' ,  'k--', 'k-.' ,  'y-','y--', 'y-.' , 'c-'     , 'm:'      )


def WeightedMean(values,error,correlMatrix = None): #implicit expectation of the gaussin error distribution
    n = len(values)
    if correlMatrix == None:
	correlMatrix = np.identity(n)
    covarMatrix = (diag(error)*matrix(correlMatrix**2)*diag(error))
    icovarMatrix = linalg.pinv(covarMatrix)

    s = 1/sum(icovarMatrix) 
    m = s*sum(values*icovarMatrix)    
    chi2 = double((values-m)*((values-m)*icovarMatrix).T)

    s*= chi2/(n-1)
    return m,sqrt(s)

def WeightedMeanVec(vectors,vec_errors,correlMatrixes): #implicit expectation of the gaussin error distribution
    
    n = len(vectors[0])
    #http://en.wikipedia.org/wiki/Weighted_mean#Vector-valued_estimates

    s = zeros((n,n))
    x = zeros((n,1))

    for (cor,err,vec) in zip(correlMatrixes,vec_errors,vectors):
	cov = diag(err)*matrix(cor**2)*diag(err)
	icov = linalg.pinv(cov)
	s += icov	
	x += dot(icov,vec).T
    s = linalg.pinv(s)
    x = dot(s,x).T
    
    chi2 = 0
    
    for (cor,err,vec) in zip(correlMatrixes,vec_errors,vectors):
	cov = diag(err)*matrix(cor**2)*diag(err)	
	icov = linalg.pinv(cov)
	chi2 += double((vec-x)*((vec-x)*icov).T)
    

    s*= chi2/len(vectors)
    
    return squeeze(x),sqrt(diag(s))


    
    
    
def IntegLightAnalyz(wavelength,plasma_start,plasma_end,SpectrometerData,PhotodiodesData, upsample = 10):
    print 'IDA start'

    wavelength,components = load_adv('./data/components')  #the most critical part, weeks of my work

    (spec_tvec,spec_tvecError,SpectrSensitiv,projection,projectionError) = SpectrometerData
    (photoData,photoError,correlList,photoTvec,photoTvecPrecis,photoFilter,photoSensitivity,photoLabel)  = PhotodiodesData
    #photoFilter - musí být všechny stejně normované aby se nepokazily vzájemé citlivosti
    SpectrSensitiv = interp(wavelength, SpectrSensitiv[:,0], SpectrSensitiv[:,1])


    
    for i,(pfilter,rsPhotodiod) in enumerate(zip(photoFilter,photoSensitivity )):	
	photoFilter[i] = interp(wavelength, pfilter[:,0], pfilter[:,1], left = 0, right = 0)
	photoSensitivity[i]=interp(wavelength, rsPhotodiod[:,0], rsPhotodiod[:,1], left = 0, right = 0)

         
    n_features    = size(components,1)
    n_components  = size(components,0)
    n_measurement = size(projection,0)
    n_photodiodes = len(photoData)
    dt = mean(diff(spec_tvec))/upsample
    


    #this matrix make a projection of the spectrometer data to the photodiodes data
    FilterProjMatrix = empty((n_components, n_photodiodes)) 
    
    for i,(pfilter,rsPhotodiod) in enumerate(zip(photoFilter,photoSensitivity)):
	FilterProjMatrix[:,i] = dot(components,1/SpectrSensitiv*pfilter*rsPhotodiod)
    
    
    photoData_down = list()
    photoError_down = list()
    photoTvec_down = list()
    #downsample data from photodiods
    for i,(pdata,perror,ptvec) in enumerate(zip(photoData,photoError,photoTvec)): 
	#nanind = array(~isfinite(pdata))

	if not all(isfinite(pdata)):
	    perror,_ = GapsFilling(perror)
	    pdata,_  = GapsFilling(pdata)
	    
	t_min = amin(spec_tvec)
	t_max = amax(spec_tvec+dt*(upsample-1))
	minerr= amin(perror)
	#plot(perror)
	#savefig('err%d.png'%i)
	#clf()
	tvec,data  = Resample(ptvec, pdata,n_measurement*upsample,t_min=t_min,t_max=t_max)
	tvec,err   = Resample(ptvec,perror,n_measurement*upsample,
				    t_min=t_min,t_max=t_max,left=minerr,right=minerr) 
	#plot(ptvec,nanind)
	#plot(ptvec,pdata)
	#plot(tvec,data)
	#savefig('data'+str(i)+'.png')
	#clf()
	photoTvec_down.append(tvec)
	photoData_down.append(data)
	photoError_down.append(err)
	##_,nanind = Resample(ptvec, nanind,n_measurement*upsample,t_min=t_min,t_max=t_max)
	#plot(photoTvec[i],nanind)
	#nanind = nanind > 0.3
	#plot(photoTvec[i],nanind)
	#xlim(plasma_start,plasma_end)
	#savefig('nanind.pdf')

	#(photoError[i])[nanind] = 1e6  #big number insted of infinity

    #upsample spectrometer data
    upsampled_proj = empty((n_components, upsample*n_measurement))
    for i, proj in enumerate(projection.T ):	
	upsampled_proj[i,:],_ = upsample_smooth(proj,15,1,sampleInteg = True,upsample = upsample)

    
    #calculate expected photodiods signal from spectrometer
    PhotodFromSpec = dot(FilterProjMatrix.T,upsampled_proj )
    PhotodFromSpec2 = dot(FilterProjMatrix.T,projection.T  )


    #calculate the mutual renormalization of the intensity and time shift of the time axis if the diagnostics
    shift = zeros(n_photodiodes)  
    amplitude = zeros(n_photodiodes) #TODO měly by být taky všechny stejné, až to bude dobře zkalibrované
    for i, (photod,spect) in enumerate(zip(photoData_down,PhotodFromSpec)): 
	convolution = fftconvolve(photod,spect[::-1], mode = 'same')
	shift[i] = len(convolution)/2-argmax(convolution)
	amplitude[i] = sum(spect)/sum(photod)
    
    tshift = shift*dt
    #estimate correction of the spectrometer timeaxis
    correlList.insert(0,1)  # spectrometer
    M = sparse.block_diag(correlList)
    correlMatrix = M.todense()

    spec_tshift,spec_tshift_err =  WeightedMean(np.hstack((0,tshift)), 
				    np.hstack((spec_tvecError,photoTvecPrecis)), 
				    correlMatrix=correlMatrix)
    

    spec_tvec += -spec_tshift+dt*upsample/2  #korekce subixelové chyby    
    tshift -= spec_tshift
    #correstion of the timeaaxis of the photodiods
    tshift,tshift_err =  WeightedMeanVec((zeros(n_photodiodes),tshift),
		    (photoTvecPrecis,spec_tshift_err*ones(n_photodiodes)),
		    (correlMatrix[1:,1:],correlMatrix[1:,1:]))
		    
    tshift = array(tshift,ndmin=1)
    tshift_err = array(tshift_err,ndmin=1)



    #make a coresponding shift of the photodiod signal
    for i, (shift,shift_err,photod,photodE,photoT,spect) \
	    in enumerate(zip(tshift,tshift_err,photoData_down,photoError_down,photoTvec_down,PhotodFromSpec)): 
	

	err_shift = hypot(shift_err,spec_tshift)
	photoTvec[i] += shift
	photoData[i] *= amplitude[i]
	
	shift += spec_tshift
	photod[:]  = interp(photoT,photoT+shift,photod )*amplitude[i] 
	photodE[:] = interp(photoT,photoT+shift,photodE)*amplitude[i]
	photodE[1:]+= err_shift*abs(diff(photodE))/dt  #error in x variable
	photoT -= spec_tshift

	
	#plot(photoData[i])
	#plot(spect)
	#savefig('_'+str(i)+'.png')
	#clf()	
	#errorbar(photoT,photod,photodE)
	#title(photoLabel[i])
	#plot(spec_tvec,PhotodFromSpec2[i,:])
	#axvline(x = plasma_start,c = 'r')
	#axvline(x = plasma_end,c = 'r')
	#savefig(str(i)+'.png')
	#clf()
    
    # x is solution    
    n_b = n_components*n_measurement
    n_x = n_b*upsample
    x = zeros(n_x)    

  
	
    # b are measured data
    b1 = reshape(projection, (-1,1),order='F')  #data from spectrometer
    b2 = array(photoData_down,copy=False)		 #data from photodiode
    b2 = reshape(b2, (-1,1),order='C')
    b = squeeze(np.vstack((b1,b2)))
   
    #e is expected error in data_file
    e1 = reshape( projectionError, (-1,1), order='F') #error from spectrometer    
    e2 = array(photoError_down,copy=False)		 #error from photodiode
    
    e2 = reshape( e2, (-1,1),order='C')    

    e = squeeze(np.vstack((e1,e2)))

    e+= 1e-4   #nonzeroes errors
    
    #The vogel matrix     
    ReductMatrix =  kron(eye(n_b,n_b),ones((1,upsample))/upsample,format='csr')

    FiltMatrix =  kron(FilterProjMatrix.T,
		    identity(n_measurement*upsample,format='dia'),format='csr')
 
    T = vstack((ReductMatrix, FiltMatrix))
    iE = diags(1/e, 0)
    T = iE*T
    b = iE*b
    tvec = photoTvec_down[0]

    #Smoothing matrix
    interv_noplasma = (tvec< plasma_start)+(tvec> plasma_end)
    interv_noplasma = squeeze(repmat(interv_noplasma,1,n_components))

 
    diag_data = ones((2,n_x))
    diag_data[1,:] *= -1
    diag_data[1,interv_noplasma] = 0

    std = sum(~interv_noplasma)/50
    gauss_win = gaussian(std*3,std)
    gauss_win/= sum(gauss_win)
    diag_data[1,:] = fftconvolve(diag_data[1,:],gauss_win,mode='same')


 
    diag_data[1,::(n_measurement*upsample)] = 0 
    D = spdiags(diag_data, (0,1), n_x, n_x,format='csr')
    D = D.T*D
    DD = D.T*D

    
    TT = T.T*T
    
    W = ones(n_x) #weight matrix
    factor = analyze(TT+DD)

    for j in range(5):
	W = D.T*sparse.spdiags(W,0, n_x,n_x)*D    
	lam = 1e4  #initial guess, just estimated "by eye"

	for i in range(5):
	    factor.cholesky_inplace(TT +lam*W)
	    g = squeeze(factor(T.T*b))
	    chi2 = norm(T*g-b)**2/len(b)
	    print lam, chi2
	    if chi2 < 1.3: #very simple root finder
		lam*= 2
	    else:
		break
	    #elif chi2 < 1:
		#lam*= 3
	    #else:
		#break
	g_tmp = copy(g)
	g_tmp[g_tmp < 0.001] = 0.001  # remove negative points (result must be positive)
	W = mean(g_tmp)/(g_tmp)   # new weight matrix

    f = T*g
    resid = f-b
    chi2 = norm(resid)**2/len(b)
    print 'chi2',chi2

    E = diags(e, 0)
    g = reshape(g, (-1,n_components),order='F')
   
    save_adv('./data/IDA_ions_projections',tvec,g )
   
    #plot results
    fig = figure(num=None, figsize=(8, 6), dpi=80, facecolor='w', edgecolor='k')
    ax = fig.add_subplot(111)
    for i in range(n_components):
	plot(tvec*1e3,g[:,i], styles[i], label=names[i]) 
    ax.text(0.9,0.9,'$\\chi^2/doF = %2.1f$'%chi2,horizontalalignment='right',
	verticalalignment='bottom',transform=ax.transAxes,backgroundcolor='w')
    leg = legend(loc='upper left', fancybox=True)
    leg.get_frame().set_alpha(0.7)
    axvline(x = plasma_start*1e3,ls = '-.',c = 'g',)
    axvline(x = plasma_end*1e3  ,ls = '-.',c = 'g',)
    ax.set_xlabel('t [ms]')
    ax.set_ylabel('relative intensity [-]')
    ax.axis('tight')
    ax.set_xlim(plasma_start*1e3-1,plasma_end*1e3+1)
    fig.savefig('./graphs/IDA_ions_projections.png',bbox_inches='tight')
    fig.clf()
    
    
    #plot retrofit 
    retroSpec = (E*f)[:n_measurement*n_components]
    retroSpec = reshape(retroSpec, (n_components,-1),order='C')
    chi2 = norm(resid[:n_b])**2/n_b
    space = 0.1
    ax = fig.add_subplot(111)

    for i in range(n_components):
	plot(spec_tvec*1e3,retroSpec[i,:]+i*space, styles[i], label=names[i])
	#errorbar(spec_tvec*1e3,projection[:,i]+i*space,
	    #yerr=projectionError[:,i],fmt=styles[i][0]+'.')#intensity with errorbars
	
    ax.text(0.9,0.9,'$\\chi^2/doF = %2.1f$'%chi2,horizontalalignment='right',
	verticalalignment='bottom',transform=ax.transAxes,backgroundcolor='w')

    ax.set_xlabel('t [ms]')
    ax.set_ylabel('shift+intensity [a.u.]')
    ##TODO dokoncit BUG  
    #show only labels of lines, not errobars
    #labels, handles = ax.get_legend_handles_labels()
    #N = len(handles)/2
    #leg = legend(labels[:N], handles[N:],loc='upper left', fancybox=True)
    
    leg = legend(loc='upper left', fancybox=True)
    leg.get_frame().set_alpha(0.7)
    for i in range(n_components):
	#plot(spec_tvec*1e3,retroSpec[i,:]+i*space, styles[i], label=names[i])
	errorbar(spec_tvec*1e3,projection[:,i]+i*space,
	    yerr=projectionError[:,i],fmt=styles[i][0]+'.')#intensity with errorbars
    
    ax.axis('tight')
    ax.set_xlim(plasma_start*1e3-1,plasma_end*1e3+1)
    fig.savefig('./graphs/spectrometer_retrofit.png',bbox_inches='tight')
    fig.clf()


  
    retroDiods = (E*f)
    cm = get_cmap('gist_rainbow')
    color = (cm(i/(1.+n_photodiodes)) for i in range(n_photodiodes+1))
    ax = fig.add_subplot(111)
    chi2 = norm(resid[n_b:])**2/(len(b)-n_b)
    ymax = 0

    index = n_measurement*n_components
    for (ptvec,pdata,perr,ptvec_full,pdata_full,lab,pc) in zip(photoTvec_down,photoData_down,photoError_down,photoTvec,photoData,photoLabel,color):
	plot(ptvec_full*1e3,pdata_full, color=pc, linestyle='-')
	fill(np.concatenate([ptvec*1e3, ptvec[::-1]*1e3]), \
	    np.concatenate([pdata - perr,(pdata + perr)[::-1]]), \
	    alpha=.4, fc=pc, ec='None')
	
	ymax = amax(pdata) if amax(pdata) > ymax else ymax
	plot(ptvec*1e3,retroDiods[index:index+len(pdata)]+index/100.,'--', color=pc,label=lab)
	index+= len(pdata)
	#fig.savefig('./graphs/photodiodes_retrofit'+str(index)+'.png',bbox_inches='tight')
    ax.text(0.9,0.9,'$\\chi^2/doF = %2.1f$'%chi2,horizontalalignment='right',
	verticalalignment='bottom',transform=ax.transAxes,backgroundcolor='w')
    leg = legend(loc='upper left', fancybox=True)
    leg.get_frame().set_alpha(0.7)
    ax.set_xlim(plasma_start*1e3-1,plasma_end*1e3+1)
    ax.set_ylim(-ymax/20,ymax*1.1)

    ax.set_xlabel('t [ms]')
    ax.set_ylabel('intensity [a.u.]')

    fig.savefig('./graphs/photodiodes_retrofit.png',bbox_inches='tight')
    fig.clf()