# -*- coding: utf-8 -*-
"""
Created on Mon Aug 16 11:51:36 2021

@author: Joschka
"""

import numpy as np
from scipy import special as sp
from src import physical_constants as cs


#HH = +1 --> Helmholtz configuration, HH = -1 --> Anti Helmholtz configuration
def B_field_ideal(N_windings,I_current,R_radius,d_distance,HH,z_arg): 
    """ Calculate B-field for ideal point like wires with R_radius and d_distance
    HH = +1 --> Helmholtz configuration, HH = -1 --> Anti Helmholtz configuration
    """
    B = cs.mu_0*N_windings*I_current/2 * R_radius**2 * (1/(R_radius**2+(z_arg-d_distance/2)**2)**(3/2) +HH* 1/(R_radius**2+(z_arg+d_distance/2)**2)**(3/2))
    B *= 1e4 #conversion Gauss
    return  B

#Argument for elliptic integrals
def k_sq(R_loop, z_loc, r, z):
    """ Argument for elliptic integrals"""
    return (4*R_loop*r)/((R_loop+r)**2 + (z-z_loc)**2)



def B_z_loop(I_current, R_loop, z_loc, r, z):
    """calculate z-component of B-field at position r and z for each individual loop"""
    B_z = 2e-7*I_current * 1/( (R_loop+r)**2 + (z-z_loc)**2 )**(1/2) *(sp.ellipk(k_sq(R_loop,z_loc,r,z))+ sp.ellipe(k_sq(R_loop,z_loc,r,z))*(R_loop**2 - r**2 - (z-z_loc)**2)/((R_loop-r)**2 + (z-z_loc)**2))
    B_z *= 1e4 #conversion to gauss
    return B_z


def B_r_loop(I_current, R_loop, z_loc, r, z):
    """calculate r-component of B-field at position r and z for each individual loop""" 
    B_r = 2e-7*I_current/r * (z-z_loc)/( (R_loop+r)**2 + (z-z_loc)**2 )**(1/2) *(-sp.ellipk(k_sq(R_loop,z_loc,r,z))+ sp.ellipe(k_sq(R_loop,z_loc,r,z))*(R_loop**2 + r**2 + (z-z_loc)**2)/((R_loop-r)**2 + (z-z_loc)**2))

    B_r *= 1e4 #conversion to gauss
    return B_r

def B_multiple(I_current, HH, R_inner, distance_coils, layers, windings, wire_width, wire_height, x, z):
    """Returns Bz along z-axis and B_r along r-axis
    HH = +1 --> Helmholtz configuration, HH = -1 --> Anti Helmholtz configuration"""
    z_start = (distance_coils/2 - windings * wire_height/2 + wire_height/2)*1e-3
    R_start = (R_inner + wire_width/2 )*1e-3

    
    if x[0] <=0:
        x_neg = np.abs([el for el in x if el<0])
        x_pos = np.array([el for el in x if el>0])
        x_zero = np.array([el for el in x if el == 0])
        
        
    B_z = np.zeros(len(z))
    B_x_pos = np.zeros(len(x_pos))
    B_x_neg = np.zeros(len(x_neg))
    B_x_zero = x_zero #If there is a zero in the x-array it is assured that the x component of the field at this point is zero
    
    
    
    for xx in range(0,layers):
        for zz in range(0,windings):
            z_pos = z_start + zz * wire_height*1e-3
            R_pos = R_start + xx * wire_width*1e-3

            B_z += B_z_loop(I_current,R_pos,z_pos,0,z) + B_z_loop(HH*I_current,R_pos,-z_pos,0,z)            
            B_x_pos += B_r_loop(I_current,R_pos,z_pos,x_pos,0) + B_r_loop(HH*I_current,R_pos,-z_pos,x_pos,0)
            B_x_neg += B_r_loop(I_current,R_pos,z_pos,x_pos,0) + B_r_loop(HH*I_current,R_pos,-z_pos,x_pos,0)
    
    B_x_neg *=-1 #B_x_neg is pointing in opposite x_direction
    B_x = np.concatenate((B_x_neg,B_x_zero,B_x_pos))
    return B_z,B_x


def B_multiple_raster(I_current, HH, R_inner, distance_coils, layers, windings, wire_width, wire_height, x, z):
    """Returns Bz along z-axis and B_r along r-axis
    HH = +1 --> Helmholtz configuration, HH = -1 --> Anti Helmholtz configuration"""
    z_start = (distance_coils/2 - windings * wire_height/2 + wire_height/2)*1e-3
    R_start = (R_inner + wire_width/2 )*1e-3
    
    if x[0] <=0:
        x_neg = np.abs([el for el in x if el<0])
        x_pos = np.array([el for el in x if el>0])
        x_zero = np.array([el for el in x if el == 0])
        
        
    B_z = np.zeros(len(z))
    B_x_pos = np.zeros(len(x_pos))
    B_x_neg = np.zeros(len(x_neg))
    B_x_zero = x_zero #If there is a zero in the x-array it is assured that the x component of the field at this point is zero
    
    
    
    #dividing each wire into 10 x 10 raster
    raster = 10
    I_current /=raster**2
    for xx in range(0,layers):
        for zz in range(0,windings):
            for xx_in in range(0,raster):
                for zz_in in range(0,raster):
                    z_pos = z_start + (zz * wire_height - wire_height/2 + zz_in * wire_height/(raster-1) )*1e-3
                    R_pos = R_start + (xx * wire_width - wire_width/2 + xx_in * wire_width/(raster-1) )*1e-3

                    B_z += B_z_loop(I_current,R_pos,z_pos,0,z) + B_z_loop(HH*I_current,R_pos,-z_pos,0,z)            
                    B_x_pos += B_r_loop(I_current,R_pos,z_pos,x_pos,0) + B_r_loop(HH*I_current,R_pos,-z_pos,x_pos,0)
                    B_x_neg += B_r_loop(I_current,R_pos,z_pos,x_pos,0) + B_r_loop(HH*I_current,R_pos,-z_pos,x_pos,0)
    
    B_x_neg *=-1 #B_x_neg is pointing in opposite x_direction
    B_x = np.concatenate((B_x_neg,B_x_zero,B_x_pos))
    return B_z,B_x

def B_multiple_raster_test(I_current, HH, R_inner, distance_coils, layers, windings, wire_width, wire_height, x, z):
    """Returns Bz along z-axis and B_r along r-axis
    HH = +1 --> Helmholtz configuration, HH = -1 --> Anti Helmholtz configuration"""
    z_start = (distance_coils/2 - windings * wire_height/2 + wire_height/2)*1e-3
    R_start = (R_inner + wire_width/2 )*1e-3
    
    if x[0] <=0:
        x_neg = np.abs([el for el in x if el<0])
        x_pos = np.array([el for el in x if el>0])
        x_zero = np.array([el for el in x if el == 0])
        
        
    B_z = np.zeros(len(z))
    B_x_pos = np.zeros(len(x_pos))
    B_x_neg = np.zeros(len(x_neg))
    B_x_zero = x_zero #If there is a zero in the x-array it is assured that the x component of the field at this point is zero
    
    
    
    #dividing each wire into 10 x 10 raster
    raster = 10
    I_current /=raster
    for xx in range(0,layers):
        for zz in range(0,windings):
            for xx_in in range(0,1):
                for zz_in in range(0,raster):
                    z_pos = z_start + (zz * wire_height - wire_height/2 + zz_in * wire_height/(raster-1) )*1e-3
                    R_pos = R_start + (xx * wire_width)*1e-3

                    B_z += B_z_loop(I_current,R_pos,z_pos,0,z) + B_z_loop(HH*I_current,R_pos,-z_pos,0,z)            
                    B_x_pos += B_r_loop(I_current,R_pos,z_pos,x_pos,0) + B_r_loop(HH*I_current,R_pos,-z_pos,x_pos,0)
                    B_x_neg += B_r_loop(I_current,R_pos,z_pos,x_pos,0) + B_r_loop(HH*I_current,R_pos,-z_pos,x_pos,0)
    
    B_x_neg *=-1 #B_x_neg is pointing in opposite x_direction
    B_x = np.concatenate((B_x_neg,B_x_zero,B_x_pos))
    return B_z,B_x