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

from numpy import *
#from matplotlib.pyplot import *
from numpy.linalg import norm
from scipy.special import ellipk, ellipe
#import time


#  functions for the calculation of the magnetic field around the current loop, by function Bloop_integ 
#is values calculated by integration, in the function Bloop_analytic is used analytical formula


# obecný prodový profil může být nahrazen skupinou bodých proudových profilů  v těžištích příslušných mezikruží nesoucích proud 

global mu_0,a

a = 0.085		#[m]
R_0 = 0.4		#[m]
mu_0 = 4*pi*1e-7
p = 2		#píkovací faktor


    
def CurrentProfile(rho,zeta,Rc,zc,I, Shaf_shift, centered = 'False'):		
#parabolický profil, rho,zeta jsou souřadnice, Rc,zeta - střed plazmatu (ale ne proudového profilu!),
#Shaf_shift veličina udávají shafranův posuv  většinou je 0-0.2
    j= zeros(shape(rho)+shape(Rc)) #předpokládám že jeden z nich je číslo
    DRc = Rc-R_0
    Drho = rho-R_0
    ao = a-sqrt(DRc**2+zc**2)
    n = size(rho,0)

    if ao/a <=5./n:  #plasma escaped out of the chamber
	DRc/= sqrt(DRc**2+zc**2)*(1-5./n)*a
	zc/= sqrt(DRc**2+zc**2)*(1-5./n)*a
	ao = 5*a/n

    if centered:
	r2 = (zeta)**2+(Drho)**2   #centered profile
    else:
	r2 = (zeta-zc)**2+(Drho-DRc)**2   #shifted profile
    plazma = (r2 < ao**2)

    j[plazma] = (p+1)*I/(pi*ao**2)*(1-r2[plazma]/ao**2)**p
    j[plazma]*= (1+(j[plazma])/amax(j[plazma])*Shaf_shift*sin(pi*(rho[plazma]-Rc)/ao))

    return j,ao
    
    
    
def Bloop_integ(R,z,rho,zeta):	# toroidální vodič - vypočtenu přímou integrací
    #R,Z poloha detektoru, rho, zeta poloha vlákna

    m = 1000
    dphi = pi/m  #protože jsem integroval jen přes půlkruh
    phi_i = (arange(0,m))*dphi
    B = zeros(3)
    R_Q = zeros(3)
    R_Q[2] = zeta
    dR_Q = zeros(3)
    dR = zeros(3)
    R_D = zeros(3)
    R_D[0] = R
    R_D[2] = z
    S = sin(phi_i)
    C = cos(phi_i)
    N = arange(m)
    Bvec= empty((m,3))
    for i,c,s in zip(N,C,S):

	
	R_Q[0] = rho*c
	R_Q[1] = rho*s
	dR_Q[0] = -s
	dR_Q[1] = c
	dR = R_D-R_Q
	dB = cross(dR_Q, dR)/norm(dR)**3
	Bvec[i,:] = dB
	if i == 0 or i == m:   #simpson integration formula
	    B+= dB
	elif i%2 == 0:
	    B+= 2*dB
	else:
	    B+= 4*dB
    
     
    B*=mu_0/(4*pi)*rho*(dphi/3)
    B*= 2  #protože jsme integroval jen přes půlkruh
    B[1] = 0#protože jsme integroval jen přes půlkruh
    Br = B[0]
    Bz = B[2]
    return Bz, Br      
    
def Bloop_analytic(R,z,rho,zeta, out = None):	
    #R,Z poloha detektoru, rho, zeta poloha vlákna

    
    #http://www.netdenizen.com/emagnettest/offaxis/?offaxisloop

    r = abs(R)
    x = z-zeta
    
    B0 = mu_0/(2*pi*rho)
    alpha = r/rho
    beta_sq = (x/rho)**2
    gamma = x/r
    Q = (1+alpha)**2+beta_sq
    k_sq = 4*alpha/Q
    
    

    E_val =  ellipe(k_sq)
    K_val =  ellipk(k_sq)
    
    
    if out == None:
	B = empty(shape(rho)+shape(R)+ (2,)) # předpokládám že jedno z nich  je pole a druhé konstanta
    else:
	B = out
	
    B[...,1] =       B0/sqrt(Q)*(E_val*(1-alpha**2-beta_sq)/(Q-4*alpha)+K_val)# B_z
    B[...,0] = B0*gamma/sqrt(Q)*(E_val*(Q-2*alpha)/(Q-4*alpha)-K_val)# B_R
    
    return B