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

#script for the creation of the plasma position animation. Intensity of the poloidal magnetics field around plasma is plotted. 
#Autor: Tomas Odstrcil

from numpy import *
from matplotlib.pyplot import *
import time
from numpy.linalg import norm
from scipy.signal import  fftconvolve
from MagFieldCalc import *
import matplotlib.animation as manimation

matplotlib.rcParams['xtick.direction'] = 'out'
matplotlib.rcParams['ytick.direction'] = 'out'

global mu_0,a,R_0

a = 0.085		#[m]
R_0 = 0.4		#[m]
mu_0 = 4*pi*1e-7


 

    
    
def PlotProfile(tvec,I_p,pos_p,name, I_coils = (), pos_coils = ()):
    # time vector, plasma current, plasma position, list of the coils currents, list of their positions 
    n = 100

    r = linspace(-a*2+R_0,a*2+R_0,n)
    z = linspace(-a*2,2*a,n)
    R, Z = np.meshgrid(r,z)
    
    
    

    phi = linspace(0,2*pi, 100)
    c = cos(phi)
    s = sin(phi)
    
    FFMpegWriter = manimation.writers['ffmpeg']
    metadata = dict(title='Movie Test', artist='Matplotlib',
	    comment='Movie support!')
    writer = FFMpegWriter(fps=10, metadata=metadata,codec='mpeg4',bitrate=2000 )

    
    
    fig = figure(figsize=(5,5))
    ax = fig.add_subplot(111)

    #external coils
    B_coil = ()
    for coordin_coil in pos_coils:
	B_coil.append(Bloop_analytic(R,Z,coordin_coil[0],coordin_coil[1]))
    
    plot(a*c+R_0, a*s, 'r--')
    plasma_edge_plot, = plot([], [], 'g--')
    mag_img = imshow(zeros((n/2,n/2)), extent=[-a+R_0,a+R_0,-a,a])
    CS = contour(R,R,R,colors='k')
    cl = clabel(CS, inline=1, fontsize=10)

  
    
    #B = Bloop_analytic(R,Z,R_0,0)
    #norm = sqrt(B[...,1]**2+B[...,0]**2)
    
    with writer.saving(fig, "video/animation.wmv", len(tvec)):
	for i in range(len(tvec)):  
	    t = time.time()
	    ax.set_title('Pol. Mag. field [mT]; time %1.2f ms'%(1e3*tvec[i]))

	    #calculate profile of the current loop
	    B = Bloop_analytic(R,Z,pos_p[i,0],pos_p[i,1])
	    
	    #external coils
	    for B_ext, I_ext in zip(B_coil,I_coils) :
		B+= B_ext*I_ext
		
	    
	    
	    norm = sqrt(B[...,1]**2+B[...,0]**2)
	    j,ao = CurrentProfile(R,Z,pos_p[i,0],pos_p[i,1],I_p[i], 0.2,centered = True)
	    #pcolor(j)  #kreslení použitého proudového profilu
	    #show()
	    #calculate profile of the current loop
	    #B = Bloop_analytic(R,Z,R_0,0)
	    
	    #BUG bude to dělat chybu s externími cívkami 
	    normconv = fftconvolve(norm,j,mode='same')*(r[1]-r[0])*(z[1]-z[0]) # neplatí to zcela přesně!! ale je to rychlé
	    
	    plasma_edge_plot.set_data(c*ao+pos_p[i,0],s*ao+pos_p[i,1])
	    
	    normconv*=1e3  # convert testa to mili tesla
	    mag_img.set_data(normconv[-n/4:n/4:-1,n/4:-n/4 ])

	    ax.autoscale_view(tight=True)
	    mag_img.autoscale()
	
	    #remove CS + clabels
	    for coll in CS.collections:
		ax.collections.remove(coll)
		
	    for text in CS.labelTexts:
		text.remove()
	    
		
	    CS = contour(R[n/4:-n/4,n/4:-n/4],Z[n/4:-n/4,n/4:-n/4],normconv[n/4:-n/4,n/4:-n/4], colors='k') # kvůli okrajovým efektům


	    cl = CS.clabel( inline=1, fontsize=10)
	    #quiver(R,Z,B[:,:,0]/norm,B[:,:,1]/norm )  #vykresí to šipečky

	    ax.relim()   #to update the axes limits if desired.

	    writer.grab_frame()	
	    
	    if i == len(tvec)/2:
		fig.savefig('icon.png',bbox_inches='tight',figsize=(4,3),dpi=40)
		fig.savefig('large_icon.png',bbox_inches='tight',figsize=(5,5),dpi=40)
	    #print  time.time()-t