#!/usr/bin/env python 
# -*- coding: utf-8 -*-
import time
t = time.time()
import os,sys
import urllib2
from CalcIonProjections import GapsFilling,CalcProjections
from IDA import IntegLightAnalyz
from CalcDensTemp import CalcDensTemp,plotDensTemp
from SpectrometerControl import Tokamak, OceanOptics_Spec
from numpy import *
import matplotlib
from scipy.linalg import  norm
from scipy.signal import fftconvolve
matplotlib.rcParams['backend'] = 'Agg'
#matplotlib.rc('font',  size='10')
#matplotlib.rc('text', usetex=True)  # FIXME !! nicer but slower !!!
#from matplotlib.pyplot import *
import matplotlib.pyplot as plt
from pygolem_lite.modules import load_adv,get_page_paths,saveconst,save_adv
from pygolem_lite import Shot


print 'including time',  time.time()-t

def LoadData_web(shot_num):
    print 'downloading '+str(shot_num)


    url = 'http://golem.fjfi.cvut.cz/operation/shots/'+str(shot_num)+'/diagnostics/Radiation/1111Spectrometer.ON/data/spectra.txt_Data_.txt'
    try:
	#print './data/PlasmaStart_'+str(shot_num)
	plasma_start = loadtxt('./data/PlasmaStart_'+str(shot_num))
	plasma_end= loadtxt('./data/PlasmaEnd_'+str(shot_num))
    except:
	print 'missing plasma start or end = no plasma?'
	plasma_start = 0
	plasma_end = inf

    #print os.path.isfile('data_'+str(shot_num)+'.txt')
    for i in range(10):
	try:
	    if not os.path.isfile('data_'+str(shot_num)+'.txt'):		
		u = urllib2.urlopen(url)
		localFile = open('data_'+str(shot_num)+'.txt', 'w')
		localFile.write(u.read())
		localFile.close()
	    break
	except:
	    print 'spectra not found, waiting ...'
	    time.sleep(1)
    print 'data waiting time %ds'%i
	

    GOLEM = Tokamak('GOLEM')
    HR2000 = OceanOptics_Spec('HR+C1886',GOLEM)
    shot = HR2000.loadData('','Python_Data_file','data_'+str(shot_num))

    return HR2000, shot, plasma_start,plasma_end

def LoadData():
    
    Data = Shot()
    
    plasma_start = Data['plasma_start']
    plasma_end = Data['plasma_end']
    plasma = Data['plasma']
    GOLEM = Tokamak('GOLEM')

    HR2000 = OceanOptics_Spec('HR+C1886',GOLEM)

    file_path = Shot().get_data('spectrometr:data', return_path=True)
    nmax = 100
    for i in range(nmax):
	try:
	    shot = HR2000.loadData('','Python_Data_file',file_path[:-4])  
	    break
	except:
	    print 'no spectroscopic data '+str(i)
	    time.sleep(10)

    if i == nmax-1:
	raise Exception('no data found!')

    return HR2000, GOLEM, shot, plasma_start,plasma_end



def LoadData_offline(shot_num):

    print 'Loading shot ', shot_num
    #try:
	##print './data/PlasmaStart_'+str(shot_num)
	#plasma_start = loadtxt('./data/PlasmaStart_'+str(shot_num))
	#plasma_end= loadtxt('./data/PlasmaEnd_'+str(shot_num))
    #except:
	#print 'missing plasma start or end = no plasma?'
    plasma_start = 0
    plasma_end = inf



    GOLEM = Tokamak('GOLEM')
    HR2000 = OceanOptics_Spec('HR+C1886',GOLEM)
    shot = HR2000.loadData('','Python_Data_file','./net_data/data_'+str(shot_num))

    return HR2000, shot, plasma_start,plasma_end

def LoadDataIDA(plasma_start,plasma_end,spectrometer ):		#TODO není dokončené načítání 




    #LOAD photodiod ========================================

    photoLabel = list()
    photoData = list()
    photoError = list()
    photoTvec = list()
    photoFilter = list()
    photoSensitivity = list()
    photoTvecError = list()
    correl_list = list()  # correlation of the errors in time vectors
    
    
    Data = Shot()
    
    photo_list = ('photodiode','photodiode_alpha','photodiode alpha') 
    for photo_name in photo_list:
	try:
	    photodiod = array(Data[photo_name], copy = False).T
	except:
	    print photo_name, 'was not found'
	    continue
	#correct abnormalilies
	photo_tvec = photodiod[:,0]
	photodiod = photodiod[:,1]
	#photodiod[photodiod < 0] = 0
	ind = (photo_tvec < plasma_end) & (photo_tvec > plasma_start)
	photodiod[~ind] = 0  
	med = median(photodiod[ind])
	photodiod[photodiod>3*med] = 3*med  #opraví se tím "EM disrupce" na konci
	SNR = 100
	background = 0.005
	photoTvec.append(photo_tvec)
	photoData.append(photodiod)
	photoError.append(hypot(photodiod/SNR,background))
	photoLabel.append(photo_name)
	photoSensitivity.append(loadtxt('bpy47.txt'))  #Photodiod_BPY47 Leybold!
	
	foto_filter = array([[200,1],[1200,1]])			#BUG  temporal solution, photodiode properties are not loaded from database
	if photo_name in ('photodiode_alpha', 'photodiode alpha'):
	    foto_filter = loadtxt('./filters/656.txt')
	    foto_filter[:,1] /= 100
	photoFilter.append(foto_filter)
	photoTvecError.append(1e-4) #[s]
	correl_list.append(1)

    for j,camera in enumerate(('camera\color_emiss_1','camera\color_emiss_2')):   #BUG !!!
	try:
	    tvec,data = Data[camera]
	    if size(data,-1) != 3:
		continue
	except:
	    continue

	tvec/= 1000
	tvec += plasma_start

	camera_sensitivity = loadtxt('filters/camera_filters.txt')                          #TODO loadovat to ze serveru


	for i,c in enumerate(('R','G','B')):
	    camera_gaps = isnan(data[:,i])
	    camera_filter = array(((350,1),(750,1)))  #no filter
	    photoTvec.append(copy(tvec))
	    photoData.append(data[:,i])
	    SNR = 10
	    background = median(data[:,i])/10
	    photoError.append(hypot(data[:,i]/SNR,background))

	    photoLabel.append('camera %s%d'%(c,j))
	    photoSensitivity.append(camera_sensitivity[:,2*i:2*(i+1)])
	    photoFilter.append(camera_filter)
	    photoTvecError.append(1e-2) #[s]

	correl_list.append(ones((3,3)))
	
    
 
    PhotodiodesData = (photoData,photoError,correl_list,photoTvec,photoTvecError,photoFilter,photoSensitivity,photoLabel)

    #LOAD spectrometer======================================
    tvec,projection = load_adv('./data/projection')
    projectionError = projection.data_err
    tvecError = projection.tvec_err
  
    SpectrSensitiv = spectrometer.relative_calibration
    SpectrometerData = (tvec,tvecError,SpectrSensitiv,projection,projectionError)


    return SpectrometerData,PhotodiodesData

def make_graphs():
    

    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:'      )
    tvec,plasma =   load_adv('./data/plasma') 
    if not any(plasma):
	print 'no radiation observed'
	return 
    tvec,intensities = load_adv('./data/intensities')
    tvec,proj_dict = load_adv('./data/projection')
    projection = proj_dict
    projectionError = proj_dict.data_err
    tvec,totalPower = load_adv('./data/TotalPower')
    totalPower_err = totalPower.data_err
    totalPower = totalPower

    wavelength,components = load_adv('./data/components')
    energy_constants = load('./data/energy_constants.npy') 
    chi2ion = loadtxt('./results/IonChi2.txt')
    n_features    = size(components,1)
    n_components  = size(components,0)
    n_measurement = size(intensities,1)
    chi2 = load('./results/TotalChi2.npy')
    tvec,plasma = load_adv('./data/plasma')  #BUG opravit to i jinde!!!
    plasma = bool_(plasma)

    
    class MyFormatter(plt.ScalarFormatter): 
	def __call__(self, x, pos=None): 
	    if pos==0: 
		return '' 
	    else: return plt.ScalarFormatter.__call__(self, x, pos)      

    
    majorLocator   = plt.MultipleLocator(50)
    majorFormatter = plt.FormatStrFormatter('%d')
    minorLocator   = plt.MultipleLocator(10)
    
 
    
    fit = mean(dot(projection[plasma,:],components).T, axis=1)
    mean_data = mean(intensities[:,plasma], axis=1)
    resid = mean_data-fit
    save_adv('./data/residual_spectra', wavelength,resid)
    
    fig = plt.figure(figsize=(15, 3), dpi=80, facecolor='w', edgecolor='k')
    ax = fig.add_subplot(111)
    ax.plot(wavelength,resid-10,'c',label = 'residuum',lw = 0.3)
    ax.plot(wavelength,mean_data,'r',label = 'data',lw = 0.3)
    ax.plot(wavelength,fit,'b', label = 'retrofit',lw = 0.3)
    ax.text(0.05,0.8,'$\\chi^2$ = %2.1f'%chi2,horizontalalignment='left',
	    verticalalignment='bottom',transform=ax.transAxes,backgroundcolor = 'w')
    ax.xaxis.set_major_locator(majorLocator)
    ax.xaxis.set_major_formatter(majorFormatter)

    #for the minor ticks, use no labels; default NullFormatter
    ax.xaxis.set_minor_locator(minorLocator)
    ax.set_xlabel('$\lambda$ [nm]')
    ax.set_ylabel('counts/read noise [-]')
    leg = ax.legend(loc='upper right', fancybox=True)
    leg.get_frame().set_alpha(0.7)

    ax.axis('tight')
    ax.set_xlim(220,880)
    ax.set_ylim(-20, 20)  #BUG dát tam třeba průměr spektra? 
    ax.set_title('Total retrofit')
    fig.savefig('graphs/spectrum_retrofit.png',bbox_inches='tight')
    fig.clf()
    
    #plasma = bool_(plasma)
    
    if not any(plasma):
	plasma[:] = True

    #plot line intensites normalized by their mean value in the database
    fig = plt.figure( figsize=(10, 6), dpi=80, facecolor='w', edgecolor='k')
    ax = fig.add_subplot(111)
    for i in range(n_components):   
	ax.errorbar(tvec[plasma]*1e3,projection[plasma,i],fmt = styles[i],\
		    yerr=projectionError[plasma,i], label=names[i])   #intensity with errorbars


    ax.set_title('ion radiation intensity')
    leg = ax.legend(loc='upper left', fancybox=True)
    leg.get_frame().set_alpha(0.7)
    ax.set_ylabel('intenzity')
    ax.set_xlabel('t [ms]')
    ax.axis('tight')
    ax.set_ylim(0,None)
    fig.savefig('graphs/RelativeIntensity.png',bbox_inches='tight')
    
    fig.clf()
    ax = fig.add_subplot(111)

	
    projection*= energy_constants
    projectionError *= energy_constants

    for i in range(n_components):   
	ax.errorbar(tvec[plasma]*1e3,projection[plasma,i]/1e3, 
		    yerr=projectionError[plasma,i]/1e3,fmt=styles[i], label=names[i])   #intensity with errorbars

    ax.set_title('radiation power in range 200-1200 nm')
    leg = ax.legend(loc='upper left', fancybox=True)
    leg.get_frame().set_alpha(0.7)
    ax.set_ylabel('P [kW]')
    ax.set_xlabel('t [ms]')
    ax.axis('tight')
    ax.set_ylim(0,None)
    fig.savefig('graphs/RadiatedEnergy.png',bbox_inches='tight')
    fig.clf()
    #show()


    tvec_temp,temp_dict = load_adv('./data/temperature')
    tmin = amin(tvec_temp)*1e3
    tmax = amax(tvec_temp)*1e3

    fig = plt.figure(figsize=(10, 3), dpi=80, facecolor='w', edgecolor='k')
    ax = fig.add_subplot(111)

    ax.errorbar(tvec*1e3,totalPower/1e3, yerr=totalPower_err/1e3)  
    ax.set_title('Total radiated power in range 200-1200 nm')
    ax.set_ylabel('P [kW]')
    ax.set_xlabel('t [ms]')
    #plt.axis('tight')
    ax.set_xlim(tmin-1, tmax+1)
    ax.set_ylim(0,None)
    fig.savefig('graphs/TotalPower.png',bbox_inches='tight')
    fig.clf()


    ##MAIN PLOT
    #data_max  = amax(projection)
    #colours = ('b','r','g','c')
    #fig = figure(num=None, figsize=(10, 10), dpi=80, facecolor='w', edgecolor='k')
    #subplots_adjust(hspace=0, wspace = 0)


    ##Hydrogen
    #ax = fig.add_subplot(6,1,1)
    #title('radiation energy in range 200-1200 nm')

    #ax.xaxis.set_major_formatter( NullFormatter() )
    #ax.yaxis.set_major_formatter( MyFormatter() )
    #i = 0
    #chi2 = chi2ion[i]/n_measurement
    #ax.text(.05,.1,'$\\chi^2$ = %2.2f'% chi2,horizontalalignment='left',verticalalignment='bottom',transform=ax.transAxes)

    #errorbar(tvec*1e3,projection[:,i], yerr=projectionError[ :,i], label = names[i],fmt=styles[i])   #intensity with errorbars
    #ylim(0, data_max)
    #ylabel('I normalized')

    #leg = legend(loc='upper right', fancybox=True)
    #leg.get_frame().set_alpha(0.5)


    ##Helium
    #ax = fig.add_subplot(6,1,2)
    #i = 4
    #ax.xaxis.set_major_formatter( NullFormatter() )
    #ax.yaxis.set_major_formatter( MyFormatter() )
    #chi2 = chi2ion[i]/n_measurement
    #ax.text(.05,.1,'$\\chi^2$ = %2.2f'% chi2,horizontalalignment='left',verticalalignment='bottom',transform=ax.transAxes)


    #errorbar(tvec*1e3,projection[:,i], yerr=projectionError[ :,i], label = names[i],fmt=styles[i])   #intensity with errorbars
    #ylim(0, data_max)
    #ylabel('I normalized')

    #leg = legend(loc='upper right', fancybox=True)
    #leg.get_frame().set_alpha(0.5)



    ##carbon 
    #ax = fig.add_subplot(6,1,3)

    #index = [5,6]
    #ax.xaxis.set_major_formatter( NullFormatter() )
    #ax.yaxis.set_major_formatter( MyFormatter() )
    #chi2 = chi2ion[i]/n_measurement
    #ax.text(.05,.1,'$\\chi^2$ = %2.2f'% chi2,horizontalalignment='left',verticalalignment='bottom',transform=ax.transAxes)

    #for i in index:  
	#errorbar(tvec*1e3,projection[:,i], yerr=projectionError[ :,i], label = names[i],fmt=styles[i])   #intensity with errorbars

    #ylim(0, data_max)
    #leg = legend(loc='upper right', fancybox=True)
    #leg.get_frame().set_alpha(0.5)
    #ylabel('I normalized')



    ##oxygen 
    #ax = fig.add_subplot(6,1,4)

    #index = [1,2,3]
    #ax.xaxis.set_major_formatter( NullFormatter() )
    #ax.yaxis.set_major_formatter( MyFormatter() )
    #chi2 = chi2ion[i]/n_measurement
    #ax.text(.05,.1,'$\\chi^2$ = %2.2f'% chi2,horizontalalignment='left',verticalalignment='bottom',transform=ax.transAxes)

    #for i in index:  
	#errorbar(tvec*1e3,projection[:,i], yerr=projectionError[ :,i], label = names[i],fmt=styles[i])   #intensity with errorbars

    #ylim(0, data_max)
    #leg = legend(loc='upper right', fancybox=True)
    #leg.get_frame().set_alpha(0.5)
    #ylabel('I normalized')



    ##Nitrogen 
    #ax = fig.add_subplot(6,1,5)
    #ax.xaxis.set_major_formatter( NullFormatter() )
    #ax.yaxis.set_major_formatter( MyFormatter() )
    #index = [7,8,9]
    ##index = [7,8]

    #chi2 = chi2ion[i]/n_measurement
    #ax.text(.05,.1,'$\\chi^2$ = %2.2f'% chi2,horizontalalignment='left',verticalalignment='bottom',transform=ax.transAxes)

    #for i in index:  
	#errorbar(tvec*1e3,projection[:,i], yerr=projectionError[ :,i], label = names[i],fmt=styles[i])   #intensity with errorbars

    #ylim(0, data_max)
    #leg = legend(loc='upper right', fancybox=True)
    #leg.get_frame().set_alpha(0.5)
    #ylabel('I normalized')





    ##Mystery 
    #ax = fig.add_subplot(6,1,6)

    #index = [10,11]
    ##index = [9,10]

    #chi2 = chi2ion[i]/n_measurement
    #ax.text(.05,.1,'$\\chi^2$ = %2.2f'% chi2,horizontalalignment='left',verticalalignment='bottom',transform=ax.transAxes)

    #for i in index:  
	#errorbar(tvec*1e3,projection[:,i], yerr=projectionError[ :,i], label = names[i],fmt=styles[i])   #intensity with errorbars

    #ylim(0, data_max)
    #leg = legend(loc='upper right', fancybox=True)
    #leg.get_frame().set_alpha(0.5)
    #ylabel('I normalized')
    #xlabel('t [ms]')

    #savefig('./graphs/MultiPlot.png', bbox_inches='tight')
    ##show()
    #clf()




def PlotDetailResiduum(GOLEM):
    print 'PlotDetailResiduum'

   
    style_dict = {'H':'b-','OI':'r-','OII':'r--', 'OIII':'r-.','OIV':'r:', 'HeI':'g-','HeII':'g--',\
    'CI':'k-','CII':'k--', 'CIII':'k-.','CIV':'k:', 'NI':'y-','NII':'y--', 'NIII':'y-.','NIV':'y:','unknow':'c-'}
    
    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:'      )
    
    wavelength,components = load_adv('./data/components')
    tvec,proj_dict = load_adv('./data/projection')
    tvec,plasma = load_adv('./data/plasma')  #BUG opravit to i jinde!!!
    plasma = bool_(plasma)
   

    projection = (proj_dict)[plasma,:]
    projectionError = proj_dict.data_err[plasma,:]
    projection = mean(projection, axis=0)
    tvec,intensities = load_adv('./data/intensities')
    intensities = mean(intensities[:,plasma],axis=1)
    norm_components = projection[:,None]*components


    fit = mean(dot(projection,components), axis = 0)
    fit = sum(norm_components, axis = 0)
    resid = intensities- fit
    
    
    majorLocator   = plt.MultipleLocator(10)
    majorFormatter = plt.FormatStrFormatter('%d')
    minorLocator   = plt.MultipleLocator(1)

    #fig = plt.figure(figsize=(10,5))
    fig = plt.figure(figsize=(200,15))

    ax = fig.add_axes([0.1, 0.1, 0.8, 0.8])

    shift = 0.7 - (wavelength-206)/183*0.1  #by eye

    ax.plot(wavelength-shift, resid, 'b:', linewidth=1,label='fit residuum')
    ax.plot(wavelength-shift, intensities, 'k:', linewidth=1, label='data')


    for shot_spectrum,n,s in zip(norm_components, names, styles):
        ind  = shot_spectrum!= 0
	ax.plot(wavelength[ind]-shift[ind],shot_spectrum[ind],s,linewidth=1, label=n)
	



    #plot lines from database
    for i, element in enumerate(GOLEM.elements):
	s = style_dict[element[1]]
	ax.axvline(x = 0 , label = element[1], c = s[0], ls = s[1:])
	for j,line in enumerate(element[0]):	
	
	    if line< amin(wavelength) or line> amax(wavelength):
		continue
	    
	    ax.axvline(x=line , c=s[0], ls=s[1:])
	    ax.text(line , i+j, element[1],color=s[0])
	    
	    
    ax.set_xlim(amin(wavelength),amax(wavelength))
    ax.set_ylim(-1,None)
    ax.set_yscale('symlog')
    leg = ax.legend(loc='center right', fancybox=True)
    leg.get_frame().set_alpha(0.7)
    ax.xaxis.set_major_locator(majorLocator)
    
    ax.xaxis.set_major_formatter(majorFormatter)

    #for the minor ticks, use no labels; default NullFormatter
    ax.xaxis.set_minor_locator(minorLocator)
    
    print 'saving '
    
    fig.savefig('./graphs/spectra.svgz',bbox_inches='tight')
    plt.close()
    
    
    




def main():
    #print 'start 1',sys.argv

    
    for path in ['graphs', 'results','data','HistoricalAnalysis']:
	if not os.path.exists(path):
	    os.mkdir(path)
	  
	    

    #print 'start 2'

    if sys.argv[1] ==  "analysis":
	#print 'load'
	spectrometr,tok, actualShot,plasma_start,plasma_end = LoadData()
	print "loaded data"
	noplasma = CalcProjections(spectrometr,actualShot,plasma_start,plasma_end)
	print "CalcProjections done"
	if noplasma:
	    print "no plasma"
	    return 
	CalcDensTemp()
    #print 'start 3'

    if sys.argv[1] ==  "plots":
	make_graphs()
	plotDensTemp()
	#print 'plotDensTemp'
	saveconst('status', 0)
	os.system('convert -resize 150x120\! ./graphs/density.png icon.png')

    #print 'start 4'

    if sys.argv[1] ==  "postanalysis":
	spectrometr,tok, actualShot,plasma_start,plasma_end = LoadData()          #BUG loadování je šíleně pomalé? neuložit si to někd předtím??
	
	wavelength, _ = actualShot.getData()
	SpectrometerData,PhotodiodesData = LoadDataIDA(plasma_start,plasma_end,spectrometr)
	
	

	IntegLightAnalyz(wavelength,plasma_start,plasma_end, SpectrometerData,PhotodiodesData)
	#PlotDetailResiduum(tok)

	
    ##print 'start 5'



if __name__ == "__main__":
    main()
    
    
#TODO 7307 - co je to tam za divné čáry kolem 600nm? 
#7960
#TODO 9298 . 750nm . není to CI??
#TODO 9156 .350-400nm co to je
#TODO 9664,9665,9670,9839,9906,10131,10136 úplně noé čáry
#TODO 10097 590 to fitne mizerně
#TODO 9504 - molekulární spektrum
#problémy - heliové výboje 
#neznámý prvek, podivné čáry 

#NOTE nejošlklivější výboje všech dob 11527,11528,11529