# program to calculate the B-field for a magnetic coil

import numpy as np
import matplotlib.pyplot as plt


def biot_savart(wire1, wire2, point, current = 1):
    '''
    current is given in units of inverse permeability 
    '''
    dl = wire2-wire1
    r = point - 0.5 * (wire2 + wire1)

    return current * np.cross(dl, r) / (4 * np.pi * np.linalg.norm(r)**3)

def coil(center, radius, theta, elements):
    '''
    theta is angle between the plane of the normal and the x-y plane (chosen for numerical reasons)
    '''
    points = np.zeros((elements, 3))
    rad = np.linspace(0, 2 * np.pi, elements, endpoint=True)
    for i in range(elements):
        points[i,:] = radius*np.array([np.cos(rad[i])*np.cos(theta),np.sin(rad[i])*np.cos(theta),np.sin(theta)])
    points = [point + center for point in points]

    return points
        

def single_coil():
    '''
    elem    : Number of elements to use in the coil
    num     : number of test points along z axis
    R       : Radius of the coil
    theta   : Angle between the normal of the coil and x-y plane
    '''
    elem = 100
    num = 10

    R = 1
    theta = 0
    wire = coil(np.zeros(3), R, theta, elem + 1)

    points = np.zeros((num,3))
    z = np.linspace(0, 1, num)
    points[:,2] = z

    B = np.zeros_like(points)

    for i in range(num):
        for j in range(elem-1):
            B[i,:] += biot_savart(wire[j], wire[j+1], points[i,:])
    
    B_true = np.zeros_like(B)
    B_true[:,2] = 0.5 / (1+z**2)**1.5

    print(B_true)
    print(B_true-B)

    pass

def helmholtz_coil():
    '''
    elem    : Number of elements to use in the coil
    num     : number of test points along each grid direction
    R       : Radius of the coil
    theta   : Angle between the normal of the coil and x-y plane
    '''
    elem = 100
    num = 10
    R = 1
    theta = 0
    wire_l = coil(np.array([0,0,R/2]), R, theta, elem)
    wire_r = coil(np.array([0,0,-R/2]), R, theta, elem)

    # x = np.linspace(-0.5, 0.5, num)
    size = 0.05
    x = np.zeros(1)
    y = np.linspace(-size, size, num)
    z = np.linspace(-size, size, num)
    xx, yy, zz = np.meshgrid(x,y,z)
    points = np.vstack([xx.reshape(-1), yy.reshape(-1), zz.reshape(-1)]).T

    B = np.zeros_like(points)

    for i in range(points.shape[0]):
        for k in range(elem-1):
            B[i,:] += biot_savart(wire_l[k], wire_l[k+1], points[i,:])
            B[i,:] += biot_savart(wire_r[k], wire_r[k+1], points[i,:])


    fig, ax = plt.subplots(1,1)
    ax.quiver(points[:,1], points[:,2], B[:,1], B[:,2])
    plt.show()

    B0 = 0.8**1.5
    delta_B = np.sqrt((B[:,2]-B0)**2 + B[:,1]**2 + B[:,0]**2)/B0
    print('The maximum B field deviation is {:.2}%'.format(max(delta_B)*100))

    pass


if __name__ == "__main__":
    # single_coil()
    helmholtz_coil()