%**************************************************************************
%This function takes Hall probe magnetic center results, and it takes probe
%position results, and then it calculates the position of the magnetic center
%of the undulator in a coordinate system whose axis is the probe position
%at the two ends of the undulator.  A linear fit to the magnetic center
%gives the beam axis.  Fields are calculated on the beam axis.
%
%Input:
%param, structure containing program parameters
%dat_file, file to write data to
%plt_file, file to send plots to
%b_meas, structure containing the measured fields
%field, structure containing field analysis results
%mag_ctr, structure containing magnetic center information
%probe_pos, structure containing the probe position information
%
%Output:
%beam_axis, structure containing information about the beam axis
%
%Zachary Wolf
%9/18/13

function [beam_axis] = hall2_anal_beam_axis(param, dat_file, plt_file, b_meas, field, mag_ctr, probe_pos)

%Get the magnetic center information
z = b_meas.z;
bx1 = b_meas.bx1;
by1 = b_meas.by1;
bz1 = b_meas.bz1;
z_peak_ctr = mag_ctr.z_peak_ctr;
x1_minus_xc = mag_ctr.x1_minus_xc;
y1_minus_yc = mag_ctr.y1_minus_yc;
kx = mag_ctr.kx;
ky = mag_ctr.ky;
ku = mag_ctr.ku;
x_probe_laser = probe_pos.x;
y_probe_laser = probe_pos.y;
z_probe_laser = probe_pos.z;

%Plot x1_minus_xc, y1_minus_yc
% name = ['Und Pol Mode: ' , strrep(param.und_pol_mode, '_', '-'), ', Probe Offset Dir: ', strrep(param.probe2_offset_dir, '_', '-')];
% hall2_beam_axis_plot_x1_minus_xc(plt_file, z_peak_ctr, x1_minus_xc, name);
% hall2_beam_axis_plot_y1_minus_yc(plt_file, z_peak_ctr, y1_minus_yc, name);

%Plot x_probe_laser, y_probe_laser
name = 'Probe Position From The Laser System';
hall2_beam_axis_plot_x_probe_laser(plt_file, z_probe_laser, x_probe_laser, name);
hall2_beam_axis_plot_y_probe_laser(plt_file, z_probe_laser, y_probe_laser, name);

%Find the probe position relative to a line through the positions at the two ends of the undulator
[mm, ii] = min(abs(z_probe_laser - param.align_z_axis_enter));
x_entr = x_probe_laser(ii);
y_entr = y_probe_laser(ii);
z_entr = z_probe_laser(ii);
%
[mm, ii] = min(abs(z_probe_laser - param.align_z_axis_exit));
x_exit = x_probe_laser(ii);
y_exit = y_probe_laser(ii);
z_exit = z_probe_laser(ii);
%
z_probe = z_probe_laser;
%
x_line = x_entr + ((x_exit - x_entr) / (z_exit - z_entr)) * (z_probe - z_entr);
y_line = y_entr + ((y_exit - y_entr) / (z_exit - z_entr)) * (z_probe - z_entr);
%
x_probe = x_probe_laser - x_line;
y_probe = y_probe_laser - y_line;

%Plot x_probe, y_probe
% name = 'Probe Position Relative To Reference Axis';
% hall2_beam_axis_plot_x_probe(plt_file, z_probe, x_probe, name);
% hall2_beam_axis_plot_y_probe(plt_file, z_probe, y_probe, name);

%Find x_probe, y_probe at the z-locations of the magnetic center measurements
x_probe_at_zc = spline(z_probe, x_probe, z_peak_ctr);
y_probe_at_zc = spline(z_probe, y_probe, z_peak_ctr);

%Plot x_probe, y_probe at the magnetic center locations
%name = 'Probe Positions At The Magnetic Center Locations';
%hall2_beam_axis_plot_x_probe(plt_file, z_peak_ctr, x_probe_at_zc, name);
%hall2_beam_axis_plot_y_probe(plt_file, z_peak_ctr, y_probe_at_zc, name);

%Find the magnetic center relative to the measurement axis
xc = x_probe_at_zc - x1_minus_xc;
yc = y_probe_at_zc - y1_minus_yc;

%Fit the magnetic center with a line to determine the beam axis
pxb = polyfit(z_peak_ctr, xc, 1);
xb = polyval(pxb, z_peak_ctr);
%
pyb = polyfit(z_peak_ctr, yc, 1);
yb = polyval(pyb, z_peak_ctr);

%Plot the magnetic center and the beam axis
name = 'Magnetic Center Rel Reference Axis';
hall2_beam_axis_plot_x_ctr_x_beam(plt_file, z_peak_ctr, xc, z_peak_ctr, xb, name);
hall2_beam_axis_plot_y_ctr_y_beam(plt_file, z_peak_ctr, yc, z_peak_ctr, yb, name);

%Perform fits, use the z-positions of the measurements
fit_order = 2;

%Find the x-position of the beam axis at the measurement z-positions
z_enter = z(1) : .001 : z_peak_ctr(1) - .001;
xb_enter = xb(1) * ones(1, length(z_enter));
z_exit = z_peak_ctr(end) + .001 : .001 : z(end);
xb_exit = xb(end) * ones(1, length(z_exit));
%xb_fit = spline([z_enter, z_peak_ctr, z_exit], [xb_enter, xb, xb_exit], z);
%pxb = polyfit([z_enter, z_peak_ctr, z_exit], [xb_enter, xb, xb_exit], 1);
pxb = polyfit(z_peak_ctr, xb, 1);
xb_fit = polyval(pxb, z);

%Find the y-position of the beam axis at the measurement z-positions
z_enter = z(1) : .001 : z_peak_ctr(1) - .001;
yb_enter = yb(1) * ones(1, length(z_enter));
z_exit = z_peak_ctr(end) + .001 : .001 : z(end);
yb_exit = yb(end) * ones(1, length(z_exit));
%yb_fit = spline([z_enter, z_peak_ctr, z_exit], [yb_enter, yb, yb_exit], z);
%pyb = polyfit([z_enter, z_peak_ctr, z_exit], [yb_enter, yb, yb_exit], 1);
pyb = polyfit(z_peak_ctr, yb, 1);
yb_fit = polyval(pyb, z);

%Fit the x-position of the magnetic center
z_enter = z(1) : .001 : z_peak_ctr(1) - .001;
xc_enter = xc(1) * ones(1, length(z_enter));
z_exit = z_peak_ctr(end) + .001 : .001 : z(end);
xc_exit = xc(end) * ones(1, length(z_exit));
%xc_fit = spline([z_enter, z_peak_ctr, z_exit], [xc_enter, xc, xc_exit], z);
%pxc = polyfit([z_enter, z_peak_ctr, z_exit], [xc_enter, xc, xc_exit], 3);
pxc = polyfit(z_peak_ctr, xc, fit_order);
xc_fit = polyval(pxc, z);

%Fit the y-position of the magnetic center
z_enter = z(1) : .001 : z_peak_ctr(1) - .001;
yc_enter = yc(1) * ones(1, length(z_enter));
z_exit = z_peak_ctr(end) + .001 : .001 : z(end);
yc_exit = yc(end) * ones(1, length(z_exit));
%yc_fit = spline([z_enter, z_peak_ctr, z_exit], [yc_enter, yc, yc_exit], z);
%pyc = polyfit([z_enter, z_peak_ctr, z_exit], [yc_enter, yc, yc_exit], 3);
pyc = polyfit(z_peak_ctr, yc, fit_order);
yc_fit = polyval(pyc, z);

%Fit x1_minus_xc
z_enter = z(1) : .001 : z_peak_ctr(1) - .001;
x1_minus_xc_enter = x1_minus_xc(1) * ones(1, length(z_enter));
z_exit = z_peak_ctr(end) + .001 : .001 : z(end);
x1_minus_xc_exit = x1_minus_xc(end) * ones(1, length(z_exit));
%x1_minus_xc_fit = spline([z_enter, z_peak_ctr, z_exit], [x1_minus_xc_enter, x1_minus_xc, x1_minus_xc_exit], z);
%px1_minus_xc = polyfit([z_enter, z_peak_ctr, z_exit], [x1_minus_xc_enter, x1_minus_xc, x1_minus_xc_exit], 3);
px1_minus_xc = polyfit(z_peak_ctr, x1_minus_xc, fit_order);
x1_minus_xc_fit = polyval(px1_minus_xc, z);

%Fit y1_minus_yc
z_enter = z(1) : .001 : z_peak_ctr(1) - .001;
y1_minus_yc_enter = y1_minus_yc(1) * ones(1, length(z_enter));
z_exit = z_peak_ctr(end) + .001 : .001 : z(end);
y1_minus_yc_exit = y1_minus_yc(end) * ones(1, length(z_exit));
%y1_minus_yc_fit = spline([z_enter, z_peak_ctr, z_exit], [y1_minus_yc_enter, y1_minus_yc, y1_minus_yc_exit], z);
%py1_minus_yc = polyfit([z_enter, z_peak_ctr, z_exit], [y1_minus_yc_enter, y1_minus_yc, y1_minus_yc_exit], 3);
py1_minus_yc = polyfit(z_peak_ctr, y1_minus_yc, fit_order);
y1_minus_yc_fit = polyval(py1_minus_yc, z);

%Plot the fits
% hall2_beam_axis_plot_fit(plt_file, z_peak_ctr, xb, z, xb_fit, 'Xb vs Z');
% hall2_beam_axis_plot_fit(plt_file, z_peak_ctr, yb, z, yb_fit, 'Yb vs Z');
% hall2_beam_axis_plot_fit(plt_file, z_peak_ctr, xc, z, xc_fit, 'Xc vs Z');
% hall2_beam_axis_plot_fit(plt_file, z_peak_ctr, yc, z, yc_fit, 'Yc vs Z');
% hall2_beam_axis_plot_fit(plt_file, z_peak_ctr, x1_minus_xc, z, x1_minus_xc_fit, '(X1 - Xc) vs Z');
% hall2_beam_axis_plot_fit(plt_file, z_peak_ctr, y1_minus_yc, z, y1_minus_yc_fit, '(Y1 - Yc) vs Z');

%Use average values for kx and ky
kx_sq_ave = mean(kx.^2);
ky_sq_ave = mean(ky.^2);

%Calculate factors to find the fields on the beam axis from the measured fields
if strcmp(param.und_pol_mode, 'LIN_POL_VERT_FIELD')
    %Calculate the correction factors
    bx1_corr = by1 .* kx_sq_ave .* ((y1_minus_yc_fit .* ((xb_fit - xc_fit) - x1_minus_xc_fit)) + (x1_minus_xc_fit .* ((yb_fit - yc_fit) - y1_minus_yc_fit)) + (((xb_fit - xc_fit) - x1_minus_xc_fit) .* ((yb_fit - yc_fit) - y1_minus_yc_fit)));
    by1_corr = by1 .* (1/2) .* (kx_sq_ave .* (xb_fit - xc_fit).^2 + ky_sq_ave .* (yb_fit - yc_fit).^2 - kx_sq_ave .* x1_minus_xc_fit.^2 - ky_sq_ave .* y1_minus_yc_fit.^2);
    by1_deriv = (by1(3 : end) - by1(1 : end-2)) ./ (z(3 : end) - z(1 : end-2));
    by1_deriv = [by1_deriv(1), by1_deriv, by1_deriv(end)];
    bz1_corr = by1_deriv .* ((yb_fit - yc_fit) - y1_minus_yc_fit);
elseif strcmp(param.und_pol_mode, 'LIN_POL_HORIZ_FIELD')
    %Calculate the correction factors
    bx1_corr = bx1 .* (1/2) .* (kx_sq_ave .* (xb_fit - xc_fit).^2 + ky_sq_ave .* (yb_fit - yc_fit).^2 - kx_sq_ave .* x1_minus_xc_fit.^2 - ky_sq_ave .* y1_minus_yc_fit.^2);
    by1_corr = bx1 .* ky_sq_ave .* ((y1_minus_yc_fit .* ((xb_fit - xc_fit) - x1_minus_xc_fit)) + (x1_minus_xc_fit .* ((yb_fit - yc_fit) - y1_minus_yc_fit)) + (((xb_fit - xc_fit) - x1_minus_xc_fit) .* ((yb_fit - yc_fit) - y1_minus_yc_fit)));
    bx1_deriv = (bx1(3 : end) - bx1(1 : end-2)) ./ (z(3 : end) - z(1 : end-2));
    bx1_deriv = [bx1_deriv(1), bx1_deriv, bx1_deriv(end)];
    bz1_corr = bx1_deriv .* ((xb_fit - xc_fit) - x1_minus_xc_fit);
elseif strcmp(param.und_pol_mode, 'CIRC_POL_RIGHT_HAND')
    bx1_corr = zeros(1, length(z));
    by1_corr = zeros(1, length(z));
    bz1_corr = -by1 * ku .* ((xb_fit - xc_fit) - x1_minus_xc_fit) + bx1 * ku .* ((yb_fit - yc_fit) - y1_minus_yc_fit);
elseif strcmp(param.und_pol_mode, 'CIRC_POL_LEFT_HAND')
    bx1_corr = zeros(1, length(z));
    by1_corr = zeros(1, length(z));
    bz1_corr = by1 * ku .* ((xb_fit - xc_fit) - x1_minus_xc_fit) - bx1 * ku .* ((yb_fit - yc_fit) - y1_minus_yc_fit);
else
    error('Unknown polarization mode');
end

%Plot the correction factors
hall2_beam_axis_plot_corr_fact(plt_file, z, bx1_corr, 'Bx1 Additive Corr');
hall2_beam_axis_plot_corr_fact(plt_file, z, by1_corr, 'By1 Additive Corr');
hall2_beam_axis_plot_corr_fact(plt_file, z, bz1_corr, 'Bz1 Additive Corr');

%Calculate the fields on the beam axis
bx_b = bx1 + bx1_corr;
by_b = by1 + by1_corr;
bz_b = bz1 + bz1_corr;

%Plot the fields on the beam axis
halln_meas_plot_bxyz_vs_z(plt_file, 'Beam Axis', z, bx_b, by_b, bz_b);

%Return the beam axis information
beam_axis.z_probe = z_probe;
beam_axis.x_probe = x_probe;
beam_axis.y_probe = y_probe;
beam_axis.z_peak_ctr = z_peak_ctr;
beam_axis.xc = xc;
beam_axis.yc = yc;
beam_axis.xb = xb;
beam_axis.yb = yb;
beam_axis.z = z;
beam_axis.xb_fit = xb_fit;
beam_axis.yb_fit = yb_fit;
beam_axis.xc_fit = xc_fit;
beam_axis.yc_fit = yc_fit;
beam_axis.x1_minus_xc_fit = x1_minus_xc_fit;
beam_axis.y1_minus_yc_fit = y1_minus_yc_fit;
beam_axis.bx1_corr = bx1_corr;
beam_axis.by1_corr = by1_corr;
beam_axis.bz1_corr = bz1_corr;
beam_axis.bx_b = bx_b;
beam_axis.by_b = by_b;
beam_axis.bz_b = bz_b;

%Done
