"""
    Autor: Borek Leitl FJFI CVUT
    last update 5.11.2017
"""

#from pygolem_lite import Shot
import numpy as np
import os, sys
import urllib

import matplotlib
matplotlib.rcParams['backend'] = 'Agg'
matplotlib.rc('font',  size='10')
import matplotlib.image as mpimg

import matplotlib.pyplot as plt
from scipy.ndimage import median_filter
from matplotlib.pyplot import axvline
import time
from scipy.optimize import curve_fit

def download(url, fileData = False):
    """
    downloads data to file if given = want to save data and work with it later or to numpy table if the filename missing
    substitutes all previous download functions
    """
    
    if fileData:
        try:
            urllib.request.urlretrieve(url, fileData)
        except:
            print("can not get data")
            return -1
    
    #download value
    else:
        try:
            with urllib.request.urlopen(url) as response:
                val = response.read().splitlines()
                val = float(val[0].decode(encoding="utf-8", errors="strict"))
                return val
        except:
            print("can not get data")
            return -1

#============================================================== FIT AXUV DATA IN SPECIFIC TIME - RESULT IS ESTIMATED PLASMA POSITION
def Gauss_function(x, a, x0, sigma,c):
    return a * np.exp(-(x - x0)**2 / (2 * sigma**2)) +c

def Gauss_fit(x,y,baseMean):
    """
    fit given data with gauss function - standart least square method is implicitly used by curve_fit from scipy.optimize
    for test if the fit is appropriate, the use of student test or so is necessary 
    """
    #test if x,y have same dimension
    if len(x) != len(y):
        return print("Gauss_fit problem - wrong dimmensions of fitting data")
    
    mean = sum(x * y) / sum(y)
    sigma = np.sqrt(sum(y * (x - mean)**2) / sum(y))
    popt, pcov = curve_fit(Gauss_function, x, y, p0=[max(y), mean, sigma,baseMean])
    
    return popt, pcov


def get_verticalPosition(data,plasmaStart):
    """
    """
    yCalib = -np.loadtxt("y_calibration.txt")/10 #midplane distance from detector axis to crossection of detector chord and midplane
    x = np.linspace(-10,10,len(data[0,:]))
    stampnum = 100
    dataCh = data[::stampnum,:] #reduces data to 100 timepoints
    time = np.linspace(0,len(data[0,:]),len(dataCh[:,0]))+plasmaStart*1e3
    z = np.zeros(len(dataCh[:,0]))

    for i in range(len(dataCh[:,0])):
        try:
            popt, pcov = Gauss_fit(yCalib,dataCh[i,:],baseMean = 0)
        except:
            popt = [0,0,0,0]

        xmax = yCalib[np.argmax(dataCh[i,:])]
        if abs(popt[1]) > 10:
            popt[1] =  10*abs(popt[1])/popt[1]
        z[i] = popt[1]
        
    
    return [time, z]

#def fitAllData(data)
    #"""
    #fit all given data and return matrix of fit parameters
    #"""
    ##TODO - nejaky jiny zpusob animace - bylo by dost sileny ukladat tolik obrazku na disk

    #fitParMatrix = 

    #return fitParMatrix
    #TODO rozhodovac� kriterium spravnost fitovani

#============================================================== FIT 
        
def make_img(data, fname, interpolation, start, end,shotno, label,title):
    plt.figure(figsize=[20, 8])          #size in inches
    plt.imshow(data,
        aspect='auto',
        extent = [start * 1e3, end * 1e3,   #time extents in ms
                   20, 1,],                   #channel idx
        interpolation=interpolation,
        cmap=plt.cm.hot,
        )
    
    hx = plt.colorbar()
    hx.set_label(label)
    plt.xlabel('time [ms]')
    plt.ylabel('AXUV1 CHANNEL NUMBER')
    plt.text(1, 2, 'TOP', backgroundcolor = "w") #TODO white background of the text - same for bottom
    plt.text(1, 18, 'BOTTOM', backgroundcolor = "w")
    #cx = axvline(x = 16,linewidth=5, color='black')
    #cx.set_label('16 ms')
    plt.title(title)
    plt.yticks(np.arange(1, 19+1))
    plt.savefig(fname)
    plt.close()                           #close this figure
    
def make_multiplot(data,dataUloop,plasmaCurrent,plasmaPosition,plasmaPositionAXUV1,DataLight,  fname, interpolation, start, end,shotno, label,title1, cam):
    """
    plot bolo and fast camera image
    """
    #AXUV1, H-alpha, Camera image , plasma position Mirnov vs plasma position bolo
    plot_num = 5
    if cam == -1:
        plot_num =plot_num-1
    
    plot_numAct = 1
    fig = plt.figure(figsize=[20, 4*plot_num])
    
    t = np.arange(len(dataUloop))/1e3
    index_start = np.where(start*1e3 <= t)[0][0]
    index_end = np.where(end*1e3 <= t)[0][0]
    
    #plot AXUV1
    ax=fig.add_subplot(plot_num,1,1)
    #plt.figure(figsize=[100, 20])             #TODO size in inches try to change text size,...
    imgplot = plt.imshow(data,
        aspect='auto',
        extent = [start * 1e3, end * 1e3,   #time extents in ms
                20, 1,],                    #channel idx
        interpolation=interpolation,
        cmap=plt.cm.hot,
        )
    ax.set_title('Before')
    hx = plt.colorbar()
    hx.set_label(label)
    plt.xlabel('time [ms]')
    plt.ylabel('CHANNEL NUMBER')
    plt.text(1, 2, 'TOP', backgroundcolor = "w") #white background of the text - same for bottom
    plt.text(1, 18, 'BOTTOM', backgroundcolor = "w")
    #cx = axvline(t = 16,linewidth=5, color='black')
    #cx.set_label('16 ms')
    plt.title(title1)
    plt.yticks(np.arange(1, 19+1))
    fig.subplots_adjust(hspace=0.30)
    plot_numAct +=1
    
    #plot cam
    if cam != -1:
        ax=fig.add_subplot(plot_num,1,plot_numAct)
        box = ax.get_position()
        ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])
        img2 = mpimg.imread('PLOTS/0211FastCamera.png')
        imgplot = plt.imshow(img2,aspect='auto',extent = [start * 1e3, end * 1e3, 20, 1,])   #need to be extended same way as bolo pic
        ax.set_title('Fast Camera - corrected image')
        plt.axis("off")
        plot_numAct +=1
    
    #plot Uloop, Plasma current
    #index_min = np.where(0.1 <= dataUloop)[0][0]
    #index_max = np.where(dataUloop >= 0.1)[0][-1]
    #t = index_min/1e3 + np.arange(len(dataUloop))/1000
    
    ax=fig.add_subplot(plot_num,1,plot_numAct)
    ax.set_title('Uloop - SHOT: %s' %shotno)
    box = ax.get_position()
    
    plt.ylim((0,np.max(dataUloop)))

    ax.plot(t[index_start:index_end],dataUloop[index_start:index_end], "r-", label = "Uloop")
    ax.plot(np.nan, "b-", label = "Plasma current") #agent for ax2 label for legend all in one
    plt.legend(loc = 2)
    #plt.axvline(start*1e3, color = "black", linestyle = "dashed")
    #plt.axvline(end*1e3, color = "black", linestyle = "dashed")
    plt.xlabel("time [ms]")
    ax.set_ylabel("U [V]")
    ax2 = ax.twinx()
    ax2.plot(t[index_start:index_end],plasmaCurrent[index_start:index_end]/1e3, "b-", label = "Plasma current") #/1e3 to kA
    ax2.set_ylabel("I [kA]")
    ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])
    ax2.set_position([box.x0, box.y0, box.width * 0.8, box.height])
    plt.xlim(start*1e3,end*1e3)
    plot_numAct +=1
    
    #add position - comparison of fitted position-bolo with estimation from Mirnov coils data
    ax=fig.add_subplot(plot_num,1,plot_numAct)
    box = ax.get_position()
    ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])

    plt.xlabel("time [ms]")
    plt.ylabel("z [cm]")
    
    z = plasmaPositionAXUV1[1]
    time = plasmaPositionAXUV1[0]
    plt.plot(time,-z,"rx", label = "Z - AXUV1")
    ax.set_title('Plasma vertical position AXUV1 - SHOT: %s' %shotno)
    
    if type(plasmaPosition) != int: #added only if data from Mirnov coils exists
        plt.plot(plasmaPosition[:,0]*1e3,plasmaPosition[:,1]*1e2, label = "Z - MirnovJK")
        ax.set_title('Plasma vertical position from Mirnov coils JK script vs AXUV1 - SHOT: %s' %shotno)
    plt.legend(loc = 1)
    plt.xlim(start*1e3,end*1e3)
    plt.ylim(-10,10)
    plot_numAct +=1
        
    #plot H-alpha and Visible light
    ax = fig.add_subplot(plot_num,1,plot_numAct)
    box = ax.get_position()
    ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])
    
    plt.plot(DataLight[:,0],-DataLight[:,1],"g-",label = "HAlpha")
    plt.plot(DataLight[:,0],-DataLight[:,2],"b-",label = "Visible light")
    plt.legend(loc = 1)
    ax.set_label("time [ms]")
    ax.set_ylabel("intensity [a.u.]")
    ax.set_title('H-alpha and Visible light - SHOT: %s' %shotno)
    #plt.axvline(start*1e3, color = "black", linestyle = "dashed")
    #plt.axvline(end*1e3, color = "black", linestyle = "dashed")
    plt.xlim(start*1e3,end*1e3)
    #TODO - pridat ohraniceni shotu pro kazde zobrazeni
    
    plt.savefig(fname)
    plt.close() #close this figure
    
    #Optional:
            #vertical line - if whole graph wtht data restrictions is used:
            #plt.axvline(start*1e3, color = "black", linestyle = "dashed")
            #plt.axvline(end*1e3, color = "black", linestyle = "dashed")

def plot_cut(data, fname,shotno):
    plt.plot(data[:,5000],'r+',markersize = 10)
    plt.savefig(fname) 
    plt.xlabel('CHANNEL NUMBER')
    plt.ylabel('P[W/sr/m$^2$]')
    plt.title('SHOT NUMBER ' '%d'%shotno)

def make_plots():
    #shot = Shot()
    
    #t, data_ko = shot['papouch_ko']
    #t, data_ja = shot['papouch_ja']
    
    shotno = "0"
    #shotno = "16131"
    url_shotno = "http://golem.fjfi.cvut.cz/utils/data/%s/shotno" %(shotno)
    shotNum = int(download(url_shotno))
    shotno = str(shotNum)
    
    url_shotno = "http://golem.fjfi.cvut.cz/utils/data/%s/shotno" %(shotno)
    url_start =  "http://golem.fjfi.cvut.cz/utils/data/%s/plasma_start" %(shotno)
    url_end = "http://golem.fjfi.cvut.cz/utils/data/%s/plasma_end" %(shotno)
    url_Cam = "http://golem.fjfi.cvut.cz/shots/%s/diagnostics/Radiation/0211FastCamera.ON/1/CorrectedRGB.png" %shotno
    #url_Ko = "http://golem.fjfi.cvut.cz/shots/%s//DAS/%s.ON/Papouch_All.npz" %(shotno,"1113Papouch_Za")
    #url_Ja = "http://golem.fjfi.cvut.cz/shots/%s//DAS/%s.ON/Papouch_All.npz" %(shotno,"0314Papouch_Ko")
        
    url_Ko = "http://golem.fjfi.cvut.cz/shots/%s//DAS/%s.ON/Papouch_All.npz" %(shotno,"1113Papouch_Ko")
    url_Ja = "http://golem.fjfi.cvut.cz/shots/%s//DAS/%s.ON/Papouch_All.npz" %(shotno,"0314Papouch_Ja")
    url_Uloop = "http://golem.fjfi.cvut.cz/utils/data/%s/loop_voltage" %(shotno)
    url_PlasmaCurrent = "http://golem.fjfi.cvut.cz/utils/data/%s/plasma_current" %(shotno)
    url_PlasmaPosJK = "http://golem.fjfi.cvut.cz/shots/%s/diagnostics/Magnetic/0316PlasmaPosition_JK.ON/plasma_position.txt" %(shotno) #vertical plasma position
    url_Halpha = "http://golem.fjfi.cvut.cz/utils/data/%s/haphoto" %(shotno)
    url_Visible = "http://golem.fjfi.cvut.cz/utils/data/%s/photo" %(shotno)
    #url_HXR = http://golem.fjfi.cvut.cz/shots/24869/diagnostics/Magnetic/0316PlasmaPosition_JK.ON/plasma_position.txt
    
    start = download(url_start)
    end = download(url_end)
    #shotNum = int(download(url_shotno))
    
    fileData = "DATA/"
    filePlots = "PLOTS/"
    
    try:
        os.mkdir("DATA/")
    except: pass  
    try:
        os.mkdir("PLOTS/")
    except: pass

    download(url_Ko, fileData + "Papouch_Ko_All.npz")
    download(url_Ja, fileData + "Papouch_Ja_All.npz")
    download(url_Uloop, fileData + "Uloop_%s" %shotno)
    download(url_PlasmaCurrent, fileData + "PlasmaCurrent_%s" %shotno)
    download(url_PlasmaPosJK, fileData + "PlasmaPosJK_%s" %shotno)
    download(url_Halpha, fileData + "Halpha_%s" %shotno)
    download(url_Visible, fileData + "Visible_%s" %shotno)
    cam = download(url_Cam,filePlots + "0211FastCamera.png")
    
    data_ko = np.load(fileData + "Papouch_Ko_All.npz")["data"]
    data_ja = np.load(fileData + "Papouch_Ja_All.npz")["data"]
    dataUloop = np.loadtxt(fileData + "Uloop_%s" %shotno)[:,1]
    plasmaCurrent = np.loadtxt(fileData + "PlasmaCurrent_%s" %shotno)[:,1]
    
    halpha = np.loadtxt(fileData + "Halpha_%s" %shotno)
    halpha[:,0] = halpha[:,0]*1e3
    visible = np.loadtxt(fileData + "Visible_%s" %shotno)
    DataLight = np.hstack((halpha,np.array([visible[:,1]]).T)) #time, halpha, visible
    
    try:
        plasmaPosition = np.loadtxt(fileData + "PlasmaPosJK_%s" %shotno)
    except:
        plasmaPosition = -1
    
    #t = np.load(fileData + "Papouch_Ja_All.npz")["data"][:,0]
    
    #start = np.load("Papouch_Ja_All.npz")["t_start"]
    #end = np.load("Papouch_Ja_All.npz")["t_end"]
    t = np.linspace(0,len(data_ko[:,0])/1e6, len(data_ko[:,0]))
    #vector_coef = np.genfromtxt('coeficients_AXUV2.txt') #TODO set after calibration is done
    start_idx = np.where(start <= t)[0][0]
    end_idx = np.where(t <= end)[0][-1]
        
    data = np.hstack([data_ko,data_ja]) # NOTE: make sure this is the right order
    data = data[start_idx:end_idx,0:20]
    data = np.delete(data,18,axis = 1) #if Jakoubek is used - 7th channel malfunction
    
    #from calibration AMP channels meets photodiodes in this way: 1 = 1, 2 <> 3, 4<>5,...
    a = np.arange(0,17,2)
    for i in a:
        pom = data[:,i+1]*1
        data[:,i+1] = data[:,i+2]*1
        data[:,i+2] = pom
    
    plasmaPositionAXUV1 = get_verticalPosition(data,start)
    
    data = data.T#/100*np.max(data)                    # transpose to have time as x axis; 
    data_calib = data
    
    make_multiplot(median_filter(data, size=(3, 333)),dataUloop,plasmaCurrent,plasmaPosition,plasmaPositionAXUV1,DataLight, filePlots+ 'AXUV1_compare_all.png', 'bicubic', start, end,shotno, 'Amplified signal filtered [V]','AXUV1 data - hot, median_filter (3,333), SHOT NUMBER: %d' %shotNum, cam)
        
    make_img(data_calib,filePlots+ 'AXUV1.png', 'none', start, end,shotno, 'U[V]', "AXUV1 - raw data, SHOT NUMBER: %d" %shotNum)
    make_img(median_filter(data, size=(3, 333)),filePlots+ 'AXUV1_medfit-bicubic.png', 'bicubic', start, end,shotno, 'U[V]', "AXUV1_medfit")
    #make_multiplot(median_filter(data_calib, size=(3, 333)),filePlots+ 'AXUV1_calib_medfit.png', 'bicubic', start, end,shotno, 'P[W/sr/m$^2$]')
    
    for i in range(60):
        cam = download(url_Cam,filePlots + "0211FastCamera.png")
        time.sleep(1)
        if cam != -1:
            break
    
    make_multiplot(median_filter(data, size=(3, 333)),dataUloop,plasmaCurrent,plasmaPosition,plasmaPositionAXUV1,DataLight, filePlots+ 'AXUV1_compare_all.png', 'bicubic', start, end,shotno, 'Amplified signal filtered [V]','AXUV1 data - hot, median_filter (3,333), SHOT NUMBER: %d' %shotNum, cam)

def print_shotno(shotno):
    print(shotno)

if __name__ == '__main__':
    make_plots()