% returns the lattice beta function evolution from the start of el_i to the
%   end of el_j
% el count starts from 0  (following placet octave element numbering)
% beta count: 1: before first element, 2: after first element, 3:
%   after second element etc. etc..
%
% phase-advances are calculated analytically
%
% dispersion: only implemented for quadrupoles, dipoles and sbend
% and not yet fully tested
%
function [s, beta_x, beta_y, alpha_x, alpha_y, mu_x, mu_y, Dx, Dy] = placet_evolve_beta_function(beta0_x, alpha0_x, beta0_y, alpha0_y, el_i, el_j, D0x, D0y, D0xp, D0yp)
%function [s, beta_x, beta_y, Dx, Dy, mu_x, mu_y] = placet_evolve_beta_function(beta0_x, alpha0_x, beta0_y, alpha0_y, el_i, el_j, D0x, D0y, D0xp, D0yp)
% initial dispersion parameters are option, 0 by default
if( nargin < 10 ), D0yp=0; end; 
if( nargin < 9 ), D0xp=0; end; 
if( nargin < 8 ), D0y=0; end; 
if( nargin < 7 ), D0x=0; end; 
B_s(1) = placet_element_get_attribute("$mysimname", el_i, "s") - placet_element_get_attribute("$mysimname", el_i, "length");
B_arr_x(:,:,1) = [beta0_x -alpha0_x 0; -alpha0_x (1+alpha0_x^2)/beta0_x 0; 0 0 1];
B_arr_y(:,:,1) = [beta0_y -alpha0_y 0; -alpha0_y (1+alpha0_y^2)/beta0_y 0; 0 0 1];
D_arr_x =[D0x D0xp 1]';
D_arr_y =[D0y D0yp 1]';
M_tot = eye(6,6);
mu_x(1) = 0;
mu_y(1) = 0;


for n=el_i:el_j
  M = placet_get_transfer_matrix("$mysimname", n,n);
  M3_x= [M(1:2,1:2) [0; 0]; 0 0 1];
  M3_y= [M(3:4,3:4) [0; 0]; 0 0 1];
  
  % evolve length
  B_s(n-el_i+2) = B_s(n-el_i+1) + placet_element_get_attribute("$mysimname", n, "length");
  %s = placet_element_get_attribute("$mysimname", n, "s");
  %B_s(n-el_i+2) = s; 
  
  % evolve beta function
  B_arr_x(:,:,n-el_i+2) = M3_x * B_arr_x(:,:,n-el_i+1) * M3_x';
  B_arr_y(:,:,n-el_i+2) = M3_y * B_arr_y(:,:,n-el_i+1) * M3_y';
  %M = placet_get_transfer_matrix("$mysimname", n,n) % debug
  %My= M(3:4, 3:4) % debug
  %B_arr_y(:,:,n-el_i+2) % debug


  % evolve phase-advance
  temp_x = M(1,1)*B_arr_x(1,1,n-el_i+1) - M(1,2)*(-B_arr_x(1,2,n-el_i+1));
  delta_mu_x = atan2(M(1,2), temp_x );
  mu_x(n-el_i+2) =  mu_x(n-el_i+1) + delta_mu_x;
  temp_y = M(3,3)*B_arr_y(1,1,n-el_i+1) - M(3,4)*(-B_arr_y(1,2,n-el_i+1));
  delta_mu_y = atan2(M(3,4), temp_y );
  mu_y(n-el_i+2) =  mu_y(n-el_i+1) + delta_mu_y;

  
  % evolve dispersion
  %   % dispersion driving for the moment limited to only Dipoles,
  %   Quadrupole offsets
  %   and Sbends
  D3_x = M3_x;
  D3_y = M3_y;
  if( strcmp(placet_element_get_attribute("$mysimname", n, "type_name"), "quadrupole") )
    strength = placet_element_get_attribute("$mysimname", n, "strength");
    off_x = placet_element_get_attribute("$mysimname", n, "x");
    off_y = placet_element_get_attribute("$mysimname", n, "y");
    s = placet_element_get_attribute("$mysimname", n, "length");
    ref_energy = placet_element_get_attribute("$mysimname", n, "e0");
    if(ref_energy == 0) 
      disp('WARNING: ref_energy not set, ignoring element of type Quadrupole, dispersion might be wrong');
      ref_energy = 1e99;
    end
    kick_x = strength*off_x / ref_energy;
    kick_y = -strength*off_y / ref_energy;
    D3_x(1,3) = s*kick_x/2;
    D3_x(2,3) = kick_x;
    D3_y(1,3) = s*kick_y/2;
    D3_y(2,3) = kick_y;
  end% if
  if( strcmp(placet_element_get_attribute("$mysimname", n, "type_name"), "dipole") );
    strength_x = placet_element_get_attribute("$mysimname", n, "strength_x");
    strength_y = placet_element_get_attribute("$mysimname", n, "strength_y");
    s = placet_element_get_attribute("$mysimname", n, "length");
    ref_energy = placet_element_get_attribute("$mysimname", n, "e0");
    if(ref_energy == 0) 
      disp('WARNING: ref_energy not set, ignoring element of type Dipole, dispersion might be wrong');
      ref_energy = 1e99;
    end
    kick_x = strength_x / ref_energy;
    kick_y = strength_y / ref_energy;
    D3_x(1,3) = s*kick_x/2;
    D3_x(2,3) = kick_x;
    D3_y(1,3) = s*kick_y/2;
    D3_y(2,3) = kick_y;
  end% if
  if( strcmp(placet_element_get_attribute("$mysimname", n, "type_name"), "sbend" ) )
    s = placet_element_get_attribute("$mysimname", n, "length");
    angle = placet_element_get_attribute("$mysimname", n, "angle");
    e0 = placet_element_get_attribute("$mysimname", n, "e0");
    if(angle != 0) 
      D3_x(1,3) = s/angle*(1-cos(angle));
      D3_x(2,3) = sin(angle);
    end
  end% if
  D_arr_x(:,n-el_i+2) = D3_x*D_arr_x(:,n-el_i+1);
  D_arr_y(:,n-el_i+2) = D3_y*D_arr_y(:,n-el_i+1);
end% for

s = B_s;
beta_x = squeeze(B_arr_x(1,1,:));
beta_y = squeeze(B_arr_y(1,1,:));
alpha_x = -squeeze(B_arr_x(1,2,:));
alpha_y = -squeeze(B_arr_y(1,2,:));
Dx = D_arr_x(1,:);
Dy = D_arr_y(1,:);

return;

