%**************************************************************************
%This function generates the field from an EPU quadrant in the coordinate
%system where z is in the beam direction, v is away from the magnets of the
%quadrant, and u is along the magnets perpendicular to v and z.  z=0 is at
%the first pole.
%
%Input:
%lambda_u, undulator period (m)
%num_per, number of periods in the undulator, including ends
%phi_0, peak magnetic scalar potential (Tm)
%u, at each z-position, u-position to give the field at (m)
%v, at each z-position, v-position to give the field at (m)
%z, array of z-positions to give the field at (m)
%
%Output:
%bu, Bu field at each z-position (T)
%bv, Bv field at each z-position (T)
%bz, Bz field at each z-position (T)
%
%Zachary Wolf
%9/20/14

function [bu, bv, bz] = gen_qn_field(lambda_u, num_per, phi_0, u, v, z)

%Define the undulator parameters
ku = 2 * pi / lambda_u;

%First calculate the fields for all z, afterward put in ends
%phi = phi_0 exp(-ku v) cos(ku z)
%b = grad(phi)
bu = zeros(1, length(z));
bv = -phi_0 * ku * exp(-ku * v) .* cos(ku * z);
bz = -phi_0 * ku * exp(-ku * v) .* sin(ku * z);

%Make the ends

%z before undulator
z_start_xy = -lambda_u / 4;
z_start_z = 0;
index_array_xy = find(z < z_start_xy);
index_array_z = find(z < z_start_z);
bu(index_array_xy) = 0;
bv(index_array_xy) = 0;
bz(index_array_z) = 0;

%z in first half period
index_array_xy = find((z >= z_start_xy) & (z < z_start_xy + lambda_u / 2));
index_array_z = find((z >= z_start_z) & (z < z_start_z + lambda_u / 2));
bu(index_array_xy) = 0.25 * bu(index_array_xy);
bv(index_array_xy) = 0.25 * bv(index_array_xy);
bz(index_array_z) = 0.25 * bz(index_array_z);

%z in second half period
index_array_xy = find((z >= z_start_xy + lambda_u / 2) & (z < z_start_xy + lambda_u));
index_array_z = find((z >= z_start_z + lambda_u / 2) & (z < z_start_z + lambda_u));
bu(index_array_xy) = 0.75 * bu(index_array_xy);
bv(index_array_xy) = 0.75 * bv(index_array_xy);
bz(index_array_z) = 0.75 * bz(index_array_z);

%z after undulator
z_end_xy = num_per * lambda_u - lambda_u / 4;
z_end_z = num_per * lambda_u;
index_array_xy = find(z > z_end_xy);
index_array_z = find(z > z_end_z);
bu(index_array_xy) = 0;
bv(index_array_xy) = 0;
bz(index_array_z) = 0;

%z in first half period from end
index_array_xy = find((z <= z_end_xy) & (z > z_end_xy - lambda_u / 2));
index_array_z = find((z <= z_end_z) & (z > z_end_z - lambda_u / 2));
bu(index_array_xy) = 0.25 * bu(index_array_xy);
bv(index_array_xy) = 0.25 * bv(index_array_xy);
bz(index_array_z) = 0.25 * bz(index_array_z);

%z in second half period from end
index_array_xy = find((z <= z_end_xy - lambda_u / 2) & (z > z_end_xy - lambda_u));
index_array_z = find((z <= z_end_z - lambda_u / 2) & (z > z_end_z - lambda_u));
bu(index_array_xy) = 0.75 * bu(index_array_xy);
bv(index_array_xy) = 0.75 * bv(index_array_xy);
bz(index_array_z) = 0.75 * bz(index_array_z);

end
