function varargout = saw_full(varargin)
%%% calculation of profiles after sawtooth crash - full reconnection model
%% % saw_full module

% #codegen
coder.extrinsic('legend');
coder.extrinsic('warning');

if nargin == 0,
  module_params = struct(...
    'check', false...
    );
  varargout{1} = module_params;
  return %empty call, probably to get default structures
elseif nargin==7;
  x = varargin{1};
  g = varargin{2};
  v = varargin{3};
  %u = varargin{4};
  %it = varargin{5};
  model  = varargin{6};
  saw_params = varargin{7}; % distribute inputs
else
  error('must call with 0 or 7 inputs');
end

xnew = x;

% options
reconnect_ne = saw_params.reconnect_ne && strcmp(model.ne.method,'state');
reconnect_te = saw_params.reconnect_te && strcmp(model.te.method,'state');
reconnect_ni = saw_params.reconnect_ni && strcmp(model.ni.method,'state');
reconnect_ti = saw_params.reconnect_ti && strcmp(model.ti.method,'state');

% general geometry info
geop = geometry_profiles(g,0*g,model);

% metric
rhogauss = model.rgrid.rhogauss;
rho = model.rgrid.rho;
rho_both_unsorted=[rhogauss;rho];
[rho_both,irhogaus_rho] = sort(rho_both_unsorted);

% finer grid
% Add rho and rhogauss points so that linear interpolation of any profile outside mixing radius will yield
% same input points if they are in the grid, should ensure to get same values on rhogauss
% for outside points at the end

% should construct r in initialization of sawtooth model otherwise nb points not fixed
nr = 401;
%r = linspace(rhogauss(1), rhogauss(end), nr);
r = unique(sort([linspace(0., 1., nr)'; rhogauss])); % need rho=0 for psistar(axis) which defines rmix

%% get profiles on dense grid
% iota
[iotarhogauss] = eval_iota(x,g,v,model,true);
[iotarho] = eval_iota(x,g,v,model,false);
iota_both = [iotarhogauss;iotarho];
iota_both = iota_both(irhogaus_rho);

% temperature
[te_rhogauss,dte_dx_eff] = eval_te(x,g,v,model,true);
te_rho = eval_te(x,g,v,model,false);
te_both=[te_rhogauss;te_rho];
te_both=te_both(irhogaus_rho);

% ion temperature
ti_rhogauss = eval_ti(x,g,v,model,true);
ti_rho = eval_ti(x,g,v,model,false);
ti_both=[ti_rhogauss;ti_rho];
ti_both=ti_both(irhogaus_rho);

% density
ne_rhogauss = eval_ne(x,g,v,model,true);
ne_rho = eval_ne(x,g,v,model,false);
ne_both = [ne_rhogauss;ne_rho];
ne_both = ne_both(irhogaus_rho);

% ion density (total number of ions)
ni_rhogauss = eval_nitot(x,g,v,model,true);
ni_rho = eval_nitot(x,g,v,model,false);
ni_both = [ni_rhogauss;ni_rho];
ni_both = ni_both(irhogaus_rho);

% phi_rho
psi = eval_psi(x,[],[],model,false);
psi0 = psi(1);
Phib = eval_Phib([],g,[],model,[]); % edge flux from g
phi_rho_both = Phib * rho_both.^2;

%% find q==1
% must use for loop: find gives variable size result that codegen does not like
%iq1 = find(iota >= 1., 1, 'last'); % outermost point with iota >=1
for iq1 = numel(iota_both):-1:1
  if iota_both(iq1)>=1; break; end
end

% workaround for reversed shear q (similar to Kim's option 1 in ASTRA)
% check for region of reverse-shear inside q surface, substitute
iota_bothfix = iota_both;
for irev = iq1:-1:1
  if iota_bothfix(irev)<iota_bothfix(irev+1)
    iota_bothfix(irev) = iota_bothfix(irev+1);
  end
end

if ~all(iota_bothfix == iota_both) % if some change occurred
  % reassign Psihat with new (corrected) iota_both
  iota_both = iota_bothfix; % reassign;
end

%%

%interpolate on a denser r grid
phi_r = interp1(rho_both, phi_rho_both, r, 'linear');
iota_r = interp1(rho_both, iota_both, r, 'linear');
% $$$ phi_r = interpos(rho_both, phi_rho_both, r, 0.);
% $$$ iota_r = interpos(rho_both, iota_both, r,0.);

[imax,imix,r_mix,psis] = mixing_region(r,phi_r,iota_r);

% get interval within rho grid modified by mixing, to check volume integrals later
indgaus_inrmix = false(model.rgrid.nrhogauss,1);
for ij = 1:numel(rho)
  if rho(ij)>r_mix;
    indgaus_inrmix = (model.rgrid.rhogauss<=model.rgrid.rho(ij));
    break
  end
end

%% reconnect psi
[rk,r1,r2,psisnew] = reconnect_psis(r,imax,imix,psis,saw_params.check);

% psi_star_new in original grid
psis_f = interp1(rk,psisnew,r);
% New psi on rhogauss
phi = r.^2.*Phib; % toroidal flux
psi_f = psis_f + phi + psi0;
psi_f_rhogauss = interp1(r, psi_f, rhogauss, 'linear');
% $$$ psi_f_rhogauss = interpos(r', psi_f, rhogauss);
xnew(model.psi.xind) = model.psi.Lamgauss \ psi_f_rhogauss;

%% Profiles on 'in' grid
[r_in,vol_in,volk] = geom_in(rho_both,r,r_mix,rk,imix,irhogaus_rho,g,model);

[ne_in,pe_in,ne_in_int,pe_in_int] = profiles_in(rho_both,ne_both,te_both,r,r_mix,imix,vol_in);
[ni_in,pi_in,ni_in_int,pi_in_int] = profiles_in(rho_both,ni_both,ti_both,r,r_mix,imix,vol_in);


%% Reconnect density profiles
if reconnect_ne
  ne_f_all = reconnect_profile(r_in,ne_in,ne_in_int,volk,r1,r2,imix,imax);
  intne_in_all = int_Vtot(ne_rhogauss,geop,model);
  ne_f_rho = conserve_particles(ne_f_all,intne_in_all,indgaus_inrmix,rk,model,geop);
  xnew(model.ne.xind) = model.ne.Lamgauss \ ne_f_rho;
end
% now has ne_f_rho as relevant ne on rhogauss

if reconnect_ni
  ni_f_all = reconnect_profile(r_in,ni_in,ni_in_int,volk,r1,r2,imix,imax);
  intni_in_all = int_Vtot(ni_rhogauss,geop,model);
  ni_f_rho = conserve_particles(ni_f_all,intni_in_all,indgaus_inrmix,rk,model,geop);
  xnew(model.ni.xind) = model.ni.Lamgauss \ ni_f_rho;
end
% now has ne_f_rho as relevant ne on rhogauss

% re-evaluate ne,ni with potentially new xnew
ne_f_rho = eval_ne(xnew,g,v,model,true);

% need all ion species to conserve energy
ni_f_rho = eval_ni(xnew,g,v,model,true) + ...
  eval_n1(xnew,g,v,model,true) + ...
  eval_n2(xnew,g,v,model,true) + ...
  eval_n3(xnew,g,v,model,true);


%% reconnect pressure profiles
if reconnect_te
  pe_new = reconnect_profile(r_in,pe_in,pe_in_int,volk,r1,r2,imix,imax);
  % total pe integral
  pe_rhogauss = ne_rhogauss.*te_rhogauss;
  intpe_in_all  = int_Vtot(pe_rhogauss,geop,model);
  Te_f_rho = conserve_energy(pe_rhogauss,pe_new,intpe_in_all,ne_f_rho,indgaus_inrmix,rk,model,geop);
  if isequal(dte_dx_eff(:,model.te.xind),model.te.Lamgauss)
    xnew(model.te.xind) = model.te.Lamgauss \ Te_f_rho;
  else
    % Cannot use te_coeffs_tmp = dte_dx_eff(:,model.te.xind) \ Te_f_rho it does not return the full lamgauss_hmode that is needed here
    % te_coeffs_tmp = dte_dx_eff(:,model.te.xind) \ Te_f_rho;
    te_coeffs_tmp = model.te.Lamgauss_hmode \ Te_f_rho;
    xnew(model.te.xind) = te_coeffs_tmp(1:numel(model.te.xind));
  end
end

if reconnect_ti
  pi_new = reconnect_profile(r_in,pi_in,pi_in_int,volk,r1,r2,imix,imax);
  pi_rhogauss = ni_rhogauss.*ti_rhogauss;
  intpi_in_all  = int_Vtot(pi_rhogauss,geop,model);
  Ti_f_rho = conserve_energy(pi_rhogauss,pi_new,intpi_in_all,ni_f_rho,indgaus_inrmix,rk,model,geop);
  xnew(model.ti.xind) = model.ti.Lamgauss \ Ti_f_rho;
end

if any(isnan(xnew))
  error('NaNs in xnew')
end

varargout{1} = xnew;

return


function [r_in,vol_in,volk] = geom_in(rho_both,r,r_mix,rk,imix,irhogaus_rho,g,model)

r_in = zeros(numel(r)+1,1); % init
r_in(1:imix) = r(1:imix); r_in(imix+1) = r_mix; r_in(imix+2:end) = r(imix+1:end);

% use both rhogrid and rhotor points to integrate Vp with more accuracy
Vpgau = eval_Vp([],g,[],model,true);
Vprho = eval_Vp([],g,[],model,false);
vvp = [Vpgau;Vprho];
vol_rho2 = cumtrapz(rho_both,vvp(irhogaus_rho));
% $$$ [~,~,~,vol_rho2] = interpos(rho_both,vvp(irhogaus_rho));
vol_in = interp1(rho_both,vol_rho2,r_in,'linear');
% $$$ vol_in = interpos(rho_both,vol_rho2,r_in);

%%% new volume with new radial profile
volk = interp1(r_in, vol_in, rk);
% $$$ volk = interpos(r_in', vol_in, rk);
volk(1) = 0; % only because r_in(1) is 0
return

function [ne_in,pe_in,ne_in_int,pe_in_int] = profiles_in(rho_both,ne_both,te_both,r,r_mix,imix,vol_in)
ne_in = calc_profile_in(rho_both,ne_both,r,r_mix,imix);
te_in = calc_profile_in(rho_both,te_both,r,r_mix,imix);
pe_in = ne_in .* te_in; % on r_in (one extra point than r)

% integration of n dV, p dV
% Vp = model.geom.LamGgauss*g(model.geom.ind_Vp);
% vol_rho = int_V(ones(model.rgrid.nrhogauss,1),g,model);
% vol1 = interp1(rho,vol_rho,r,'linear');

ne_in_int = cumtrapz(vol_in,ne_in);
pe_in_int = cumtrapz(vol_in,pe_in);
% $$$ tension=0.;
% $$$ [~,~,~,ne_in_int] = interpos(vol_in, ne_in, tension);
% $$$ [~,~,~,pe_in_int] = interpos(vol_in, pe_in, tension);

return

function ne_f_rho = conserve_particles(ne_f1_all,intne_int,indgaus_inrmix,rk,model,geop)

rhogauss = model.rgrid.rhogauss;
ne_f_rho = interp1(rk,ne_f1_all,rhogauss, 'linear', 'extrap');
% fix ne new on rhogauss to ensure particle conservation
ne_new_rhogauss = model.ne.Lamgauss*(model.ne.Lamgauss \ ne_f_rho);

% From this good approximation of xnew(ne), rescale ne part in rmix only to match total int(ne vp drho)
intne_new = int_Vtot(ne_new_rhogauss,geop,model);

nee = ne_new_rhogauss; nee(~indgaus_inrmix)=0;
intne_new_rmix = int_Vtot(nee,geop,model);

corr_ne=1+(intne_int-intne_new)./intne_new_rmix;
ne_new_rhogauss(indgaus_inrmix) = corr_ne.*ne_new_rhogauss(indgaus_inrmix);
ne_f_rho = ne_new_rhogauss;

return


function Te_f_rho = conserve_energy(pe_rhogauss,pe_new,intpe_in_all,ne_f_rho,indgaus_inrmix,rk,model,geop)
rhogauss = model.rgrid.rhogauss;
% integral up to rmix
pepr = pe_rhogauss; pepr(~indgaus_inrmix) = 0;
intpe_in_rmix = int_Vtot(pepr,geop,model);

% integral up to rmix
% interpolate to rhogauss
pe_new_rhogauss = interp1(rk,pe_new,rhogauss, 'linear', 'extrap');
pepr_new = pe_new_rhogauss; pepr_new(~indgaus_inrmix) = 0;
intpe_new_rmix = int_Vtot(pepr_new,geop,model);

% scale inner part to match correct integral;
pe_new_rhogauss(indgaus_inrmix) = pe_new_rhogauss(indgaus_inrmix).*intpe_in_rmix./intpe_new_rmix;

% project on xnew
Te_f_rho = pe_new_rhogauss./ne_f_rho;
te_new_rhogauss = model.te.Lamgauss*(model.te.Lamgauss \ Te_f_rho);
% Note te_new_rhogauss is a spline approx of a sharp drop in Te_f_rho, thus not the same as Te_f_rho

% From this good approximation of xnew, rescale Te part in rmix only to match total int(pe)
% total integral of pe_new
pe_new_rhogauss = ne_f_rho.*te_new_rhogauss;
intpe_new_all = int_Vtot(pe_new_rhogauss,geop,model);

% integral only up to rmix
pep = pe_new_rhogauss; pep(~indgaus_inrmix)=0;
intpe_new_rmix = int_Vtot(pep,geop,model);

corr_te = 1+(intpe_in_all-intpe_new_all)./intpe_new_rmix;
te_new_rhogauss(indgaus_inrmix) = corr_te.*te_new_rhogauss(indgaus_inrmix);

Te_f_rho = te_new_rhogauss;
return

function [imax,imix,r_mix,psis] = mixing_region(r,phi_r,iota_r)
% helical flux
% psis_rho = (psi_rho - psi_rho(1)) - (phi_rho - phi_rho(1)); %convergence problems

psis = cumtrapz(phi_r, iota_r - 1);
% $$$ [~,~,~,psis] = interpos(phi_r', iota_r - 1);

%% define regions:
% 1: 0 to psistar max
% 2: psistar max to psistar=0 (mixing radius)
%disp([imix,imix2]); % should be the same

% define r1
[~,imax] = max(psis);

%imix2 = find((psis(1:end-1)>0)&(diff(psis)<0),1,'last')+1;
% must use for loop to avoid variable size from find in codegen.
imix = numel(psis); % init
for iimix = (numel(psis)-1):-1:imax % search from out to in
  if psis(iimix)>0 && diff(psis(iimix:iimix+1))<0;
    imix = iimix;
    break
  end
  if imix == imax; plot(r, psis); error('no good imix found'); end
end
%disp([imix,imix2]); % should be the same

r_mix = r(imix) + (0-psis(imix))*(r(imix+1)-r(imix))/(psis(imix+1)-psis(imix));
return


function [rk,r1,r2,psisnew] = reconnect_psis(r,imax,imix,psis,check)
%%
nr = numel(r);
% define r2
ii2 = imax:imix+1; % use imix+1 point to avoid extrapolation
ppp = psis(imix+1+imax-(ii2)); % use imix+1 point to avoid extrapolation
rrr = r(imix+1+imax-(ii2));

assert(all(diff(ppp)>0),'ppp not monotonic')
r1 = r(1:imax); % radius from 0 to max
r2 = interp1(ppp,rrr,psis(1:imax), 'linear', 'extrap'); % find r corresponding to each psis(r1)
% $$$ r2 = interpos(ppp,rrr,psis(1:imax)); r2=reshape(r2,1,imax);

%% define rk and corresponding psi_star
rk = zeros(nr-imix+imax,1);  % init
rk(1:imax) = sqrt(abs(r2(imax:-1:1).^2-r1(imax:-1:1).^2)); % use abs to avoid complex values
rk(imax+1:nr-imix+imax) = r(imix+1:nr);

assert(all(diff(rk)>0),'rk not monotonic')

psisnew = zeros(1,nr-imix+imax); % init
psisnew(1:imax) = psis(imax:-1:1);
psisnew(imax+1:nr-imix+imax) = psis(imix+1:nr);

%   hs=subplot(222);
%   plot(rrr,ppp,'ro',r,psis,'r.',...
%     [r1,r2]',[1;1]*psis(1:imax)','bx-',...
%     rk,psisnew,'g-');

return

function prof_new = reconnect_profile(r_in,prof_in,prof_in_int,volk,r1,r2,imix,imax)
nr = numel(r_in)-1;
ii2 = imax:imix+1; % use imix+1 point to avoid extrapolation

% interpolation of pe on r1
peint1 = interp1(r_in,prof_in_int,r1);
% $$$ peint1 = interpos(r_in',pe_in_int,r1,0.);

% interpolation of pe on r2
peint2 = interp1(r_in(ii2),prof_in_int(ii2),r2, 'linear', 'extrap');
% $$$ peint2 = interpos(r_in(ii2)',pe_in_int(ii2),r2,0.);
cal = peint2(end:-1:1)-peint1(end:-1:1);
% $$$ pe_f1(2:imax) = diff(cal)./diff(volk(1:imax));pe_f1(1)=pe_f1(2);
% this 5 points central derivative gives some smoothing but interpos below yields much better
% however later fit to splines leads to smoothing in any case probably
pe_f1=zeros(size(cal));
pe_f1(3:end-2) = ((2.*(cal(4:end-1)-cal(2:end-3))./((volk(4:imax-1)-volk(2:imax-3))/2))+ ...
  (cal(5:end)-cal(1:end-4))./((volk(5:imax)-volk(1:imax-4))./4))/8;
pe_f1(1:2) = pe_f1(3);
pe_f1(end-1:end) = pe_f1(end-2);
% $$$ [pe_fit2,pe_f1] = interpos(volk(1:imax),cal,-1e3,[2 2],[cal(1) cal(imax)]);

prof_new = zeros(nr-imix+imax,1); % init
prof_new(1:numel(pe_f1)) = pe_f1;
prof_new(imax+1:nr-imix+imax)  = prof_in(imix+1+1:nr+1); % shited by 1 since r_in has r_mix added

return

function prof_in = calc_profile_in(rho_both,prof_both,r,r_mix,p)
pro = interp1(rho_both, prof_both, r, 'linear');
% $$$ ne = interpos(rho_both, ne_both, r);
prof_mix = interp1(r, pro, r_mix);
% $$$ ne_mix = interpos(r, ne, r_mix,0.);
prof_in = insert_1point(pro,prof_mix,p);
return


function q_out = insert_1point(q,q1,point1)
q_out = zeros(numel(q)+1,1); % init
q_out(1:numel(q)+1) = [q(1:point1);q1;q(point1+1:end)];

return

% function q = insert_2points(q,q1,point1,q2,point2)
% q = [q(1:point1),q1,q(point1+1:point2),q2,q(point2+1:end)];
% return
