%% Calculates transport in the phloem and xylem

%axial area and radius
A_ax_p  =   A_ax_p0 * exp((pp-0)/e_phloem);           %axial area phloem
r_p     =   (A_ax_p/pi).^0.5;                         %phloem radius

%Conductivity and viscosity
psi     = sp*V_./(sp*V_+wp/rho);                      %volume fraction sucrose in phloem 
mu_p    = mu_x*exp((4.68*0.956*psi)./(1-0.956*psi));  %dynamic visc,MPa s
k_p     = (ap_p*A_ax_p)./mu_p;                        %phloem axial hydraulic conductivity, m4 MPa-1 s-1 
k_x     = (r_x^2)/(8*mu_x)* A_ax_x;                   %xylem axial hydraulic conductivity,  m4 MPa-1 s-1 

%Osmotic pressure and water potential
m       =   cp ./ (rho *(1- cp*V_));
pip     = - rho * R * T * (0.9988*m + 0.089*m.*m); %non-linear osmotic pressure, MPa
phi_st  =   pp + pip;                              %phloem water potential, MPa
phi_xyl =   px;                                    %xylem water potential, MPa

%Boundary conditions xylem           
JW_load = -0.018/rho;                               %Evaporation Loading zon                     
JW_1    = k_x*(px(Loc_2_un(end))-0)/x_2_un;         %Uptake sink 1
JW_2    = 0.002/rho;                                %Evaporation sink 2


%Loading zone
i=2:Loc_load(end); 
JWp_in  = [0; k_p(i-1)/ x_load .* (pp(i-1)-pp(i))];   
JWx_in  = [JW_load; k_x/ x_load .* (px(i-1)-px(i))];        
A_rad_load =  2*pi*r_p(1:N_load)*x_load; JWlat_p =  Lr*(phi_xyl(1:N_load)-phi_st(1:N_load)).*A_rad_load;

%Transport equations sink 1
i=Loc_1; 
JWp_in  = [JWp_in; k_p(i-1)/ x_1 .* (pp(i-1)-pp(i))]; 
JWx_in  = [JWx_in; k_x     / x_1 .* (px(i-1)-px(i))]; 
    % Branching point
    JWp_in(Loc_1(1))  = [k_p(Loc_load(end))/ x_1 .* (pp(Loc_load(end))-pp((Loc_1(1)))) ];      
    JWx_in(Loc_1(1))  = [k_x               / x_1 .* (px(Loc_load(end))-px((Loc_1(1)))) ];  
A_rad_1 =   2*pi*r_p(i)*x_1  ;  JWlat_p = [JWlat_p ;Lr*(phi_xyl(i)-phi_st(i)).* A_rad_1];  

i=Loc_1_un; 
JWp_in  = [JWp_in; k_p(i-1)/ x_1_un .* (pp(i-1)-pp(i))]; 
JWx_in  = [JWx_in; k_x     / x_1_un .* (px(i-1)-px(i))]; 
A_rad_1_un =   2*pi*r_p(i)*x_1_un  ;  JWlat_p = [JWlat_p ;Lr*(phi_xyl(i)-phi_st(i)).* A_rad_1_un];  

%Transport equations sink 2
i=Loc_2; 
JWp_in  = [JWp_in; k_p(i-1)/ x_2 .* (pp(i-1)-pp(i))]; 
JWx_in  = [JWx_in; k_x     / x_2 .* (px(i-1)-px(i))]; 
    % Branching point
    JWp_in(Loc_2(1))  = [k_p(Loc_load(end))/ x_2 .* (pp(Loc_load(end))-pp((Loc_2(1)))) ];      
    JWx_in(Loc_2(1))  = [k_x               / x_2 .* (px(Loc_load(end))-px((Loc_2(1)))) ];  
A_rad_2 =   2*pi*r_p(i)*x_2  ;  JWlat_p = [JWlat_p ;Lr*(phi_xyl(i)-phi_st(i)).* A_rad_2];  

i=Loc_2_un; 
JWp_in  = [JWp_in; k_p(i-1)/ x_2_un .* (pp(i-1)-pp(i))]; 
JWx_in  = [JWx_in; k_x     / x_2_un .* (px(i-1)-px(i))]; 
A_rad_2_un =   2*pi*r_p(i)*x_2_un  ;  JWlat_p = [JWlat_p ;Lr*(phi_xyl(i)-phi_st(i)).* A_rad_2_un];  

%% Connecting outflows & Solute flows
JWp_out = [JWp_in(2:Loc_1_un(end)); 0]; 
JWx_out = [JWx_in(2:Loc_1_un(end)); JW_2]; 

JWp_out = [JWp_out; JWp_in(Loc_1_un(end)+2:end); 0]; 
JWx_out = [JWx_out; JWx_in(Loc_1_un(end)+2:end); JW_1]; 

JWp_out(Loc_load(end)) = JWp_in(Loc_1(1)) + JWp_in(Loc_2(1)); 
JWx_out(Loc_load(end)) = JWx_in(Loc_1(1)) + JWx_in(Loc_2(1));   
  
%Lateral water flow xylem
JWlat_x = - JWlat_p; 
    
%Solute transport   
JSp_in      = JWp_in.*[0; cp((2:Loc_2_un(end))-1)];
JSp_out     = JWp_out.* cp ; 
  
    JSp_in(Loc_1(1)) =  JWp_in(Loc_1(1)) * cp(Loc_load(end));
    JSp_in(Loc_2(1)) =  JWp_in(Loc_2(1)) * cp(Loc_load(end));
    JSp_out(Loc_load(end)) = JSp_in(Loc_1(1)) + JSp_in(Loc_2(1));     
  
 JSp = JSp_in - JSp_out;