%**************************************************************************
%This function generates fake EPU field data.
%It takes the quadrant row positions as input and calculates Bx and By on
%the undulator axis.  See LCLS-TN-14-2.
%It writes the field data to an output file.
%
%Zachary Wolf
%9/17/14

function [] = gen_epu()

%Clear the workspace
clear

%Get the program parameters
[param] = gen_epu_param();

%Extract the program parameters
output_file_dir = param.output_file_dir;
lambda_u = param.lambda_u;
num_per = param.num_per;
phi_0 = param.phi_0;
z_01 = param.z_01;
z_02 = param.z_02;
z_03 = param.z_03;
z_04 = param.z_04;
z = param.z;
u = zeros(1, length(z));
v = zeros(1, length(z));

%Make the output file name
dat_file = strcat(output_file_dir, 'gen_epu.dat');
plt_file = strcat(output_file_dir, 'gen_epu.ps');
mat_file = strcat(output_file_dir, 'gen_epu.mat');

%Initialize the data file
util_dat_init(dat_file);

%Write a message to the data file
msg{1} = ' ';
msg{2} = 'Generate Fake EPU Field Measurements';
msg{3} = '************************************';
util_dat_msg(dat_file, msg);

%Write the parameter values to the data file
gen_epu_dat_param(dat_file, param);

%Initialize the plot file
util_plot_init(plt_file, msg);

%Calculate the fields from quadrants 1 and 2
[bu, bv, bz] = gen_qn_field(lambda_u, num_per, phi_0, u, v, z);

%Rotate the quadrant measurements to find the fields in the undulator
[q1_bx, q1_by] = sup_calc_rot_fields(bu, bv, 135 * pi / 180);
q1_bz = bz;
[q2_bx, q2_by] = sup_calc_rot_fields(bu, bv, -135 * pi / 180);
q2_bz = bz;

%Plot the fields and rotated fields from the quadrants
% sup_quad_plot_b_meas_b_und(plt_file, 'Q1', z, bu, bv, q1_bx, q1_by);
% sup_quad_plot_b_meas_b_und(plt_file, 'Q2', z, bu, bv, q2_bx, q2_by);

%Calculate the fields from quadrants 3 and 4
[bu, bv, bz] = gen_qn_field(lambda_u, num_per, -phi_0, u, v, z);

%Rotate the quadrant measurements to find the fields in the undulator
[q3_bx, q3_by] = sup_calc_rot_fields(bu, bv, -45 * pi / 180);
q3_bz = bz;
[q4_bx, q4_by] = sup_calc_rot_fields(bu, bv, 45 * pi / 180);
q4_bz = bz;

%Plot the fields and rotated fields from the quadrants
% sup_quad_plot_b_meas_b_und(plt_file, 'Q3', z, bu, bv, q3_bx, q3_by);
% sup_quad_plot_b_meas_b_und(plt_file, 'Q4', z, bu, bv, q4_bx, q4_by);

%Shift the quadrants according to the row positions
q1u_bx = spline(z + z_01, q1_bx, z);
q1u_by = spline(z + z_01, q1_by, z);
q1u_bz = spline(z + z_01, q1_bz, z);
%
q2u_bx = spline(z + z_02, q2_bx, z);
q2u_by = spline(z + z_02, q2_by, z);
q2u_bz = spline(z + z_02, q2_bz, z);
%
q3u_bx = spline(z + z_03, q3_bx, z);
q3u_by = spline(z + z_03, q3_by, z);
q3u_bz = spline(z + z_03, q3_bz, z);
%
q4u_bx = spline(z + z_04, q4_bx, z);
q4u_by = spline(z + z_04, q4_by, z);
q4u_bz = spline(z + z_04, q4_bz, z);

%Plot the shifted rotated fields from the quadrants
sup_quad_plot_b_meas_b_und(plt_file, 'Q1', z, bu, bv, q1u_bx, q1u_by);
sup_quad_plot_b_meas_b_und(plt_file, 'Q2', z, bu, bv, q2u_bx, q2u_by);
sup_quad_plot_b_meas_b_und(plt_file, 'Q3', z, bu, bv, q3u_bx, q3u_by);
sup_quad_plot_b_meas_b_und(plt_file, 'Q4', z, bu, bv, q4u_bx, q4u_by);

%Make the undulator fields by superposing the quadrant fields
bx = q1u_bx + q2u_bx + q3u_bx + q4u_bx;
by = q1u_by + q2u_by + q3u_by + q4u_by;
bz = q1u_bz + q2u_bz + q3u_bz + q4u_bz;

%Add a corrector field
if param.add_corr == true
    bx_corr_array = zeros(1, length(z));
    by_corr_array = zeros(1, length(z));
    index_array = find((z > 0) & (z < param.num_per * param.lambda_u));
    bx_corr_array(index_array) = param.bx_corr;
    by_corr_array(index_array) = param.by_corr;
    bx = bx + bx_corr_array;
    by = by + by_corr_array;
end

%Remove end effects
index_array = find((z > z(1) + lambda_u) & (z < z(end) - lambda_u));
z = z(index_array);
bx = bx(index_array);
by = by(index_array);
bz = bz(index_array);

%Plot the fields in the undulator
sup_quad_plot_b_und(plt_file, z, bx, by);

%Write the results to the data file
gen_dat_zbxyz(dat_file, z, bx, by, bz);

%Output the results to the mat file
save(mat_file, 'param', 'z', 'bx', 'by', 'bz');

%Done
