function [r,pol,res,zer,err,S] = aaaz(f,deg,nl,tol,mero,plt)
%AAAZ   Rational approximation on unit circle or disk.
%   [r,pol,res,zer,err,S] = aaaz(f,deg,nlawson,tol,mero,plt)
%Examples:
%   aaaz(@exp); aaaz(@exp,6); aaaz(@exp,6,20);
%   [r,pol] = aaaz(@(z) tan(pi*z),150,0,1e-13,1); pol
%   aaaz(@(z) abs(real(z)),150,0,1e-8,1);

                                               % INITIALIZATIONS
if nargin<2, deg = 150; end                    % arg 2: (max) degree
if nargin<3, nl = 0; end                       % arg 3: Lawson steps, eg 20
if nargin<4, tol = 1e-13; end                  % arg 4: relative tolerance
if nargin<5, mero = 0; end                     % arg 5: allow poles in disk
if nargin<6, plt = 1; end                      % arg 6: plt=0 stops plotting
if plt, clf, end
blue = [0 0 .8]; red = [.8 0 0];
grn = [.5 .8 .5];  MS = 'markersize';
CO = 'color'; LW = 'linewidth'; RT = 'right';
HA = 'horizontalalignment'; FS = 'fontsize';
S = [-1 1];                                    % row vector of support pts
f0 = f([S.'; ZS(S,10)]); err = std(f0);
if (err==0) | (err/abs(mean(f0))<=tol) ...     % constant function or deg 0:
            | (deg==0)                         % return (no plot)
   r = @(z) mean(f0) + 0*z; return
end 
m0 = 2; S0 = S; w0 = [1; -1]; err0 = err;

while true                                     % MAIN AAA LOOP
   m = length(S);
   Z = ZS(S,max(3,16-m));                      % column vector of sample pts
   C = 1./(Z-S);                               % Cauchy matrix
   A = (f(Z)-f(S)).*C;                         % Loewner matrix
   [~,~,V] = svd(A,0);                         % SVD
   w = V(:,end);                               % barycentric wt vector
   R = (C*(w.*f(S.')))./(C*w);                 % rat approx at sample pts
   err = norm(f(Z)-R,inf);                     % abs error of this approx
   pol = prz(S.',f(S.'),w);                    % poles
   bad = any(abs(pol)<=1);                     % check for bad poles
   if mero~=0, bad = false; end
   co = blue;  
   fmax = norm(f([S.';Z]),inf);
   if ~bad & err < err0
      m0 = m; S0 = S; w0 = w; err0 = err;      % save latest success
   end
   if plt
      subplot(231), if bad, co = red; end
      semilogy(m-1,err,'.',MS,6,CO,co), hold on
      title convergence, ylim([1e-16 1e4]), grid on
      set(gca,'ytick',10.^(-16:8:0),FS,7), xlabel('degree (=m-1)')
   end
   if ~bad & (err/fmax<=tol) | ...             % stop if converged
      (m == deg+1) | ...                       % stop if max degree reached
      ((m-m0>=10) & (err0/fmax<1e-2))          % stop if stagnated
      break
   end 
   [~,j] = max(abs(f(Z)-R)); S = [S Z(j)];     % next support point
end
m = m0; S = S0; w = w0; err = err0; 
[pol,res,zer] = prz(S.',f(S.'),w);
r = @(z) reval(z,S.',f(S.'),w);                % function handle for r

if nl > 0                                      % LAWSON ITERATION
   Z = ZS(S,20); n = length(Z);                % switch to finer grid
   wt = ones(m+n,1);                           % Lawson wt vector
   C = 1./(Z-S);                               % Cauchy matrix
   R = (C*(w.*f(S.')))./(C*w);                 % rat approx at sample pts
   F = [f(Z); f(S.')]; R = [R; f(S.')];        % f and r values
   L = f(Z).*C;
   A = [C; eye(m)]; B = [L; diag(f(S))];       % numer and f*denom rows
   for stepno = 1:nl                           % take nl Lawson steps
      [~,~,V] = svd(sqrt(wt).*[A B/fmax],0);
      c = V(:,2*m);
      a = c(1:m); w = -c(m+1:2*m)/fmax;        % numer wts, denom wts
      R = (A*a)./(A*w);
      wt = wt.*abs(F-R);                       % iterative reweighting
      wt = wt/norm(wt,inf);
   end
   r = @(z) reval(z,S.',a./w,w);               % function handle for r
end

zz = exp(1i*sort(angle([ZS(S,50); S.'])));     % COMPUTE ERROR AND PLOT
ee = f(zz)-r(zz); err = norm(ee,inf);
if plt
   semilogy(m-1,max(err,1e-16),'o',LW,1,CO,grn)
   axis square, hold off, subplot(232), co = [0 .45 .74];
   if nl > 0, co = [.7 0 .9]; end, plot(ee,LW,.6,CO,co)
   s = sprintf('           err%8.1e',err); scl = 1.5*(err+realmin);
   title(s), grid on, axis(scl*[-1 1 -1 1]), axis square
   tt = unwrap(angle(ee)); if nl>0, wn = round((tt(end)-tt(1))/(2*pi));
   s = sprintf('winding no %d',wn);
   text(0.95*scl,-.83*scl,s,HA,RT,FS,7), end, set(gca,FS,7)
   subplot(233), scl = 2.5; if length(pol)>0, scl = 1.5+min(abs(pol)); end
   plot(zz([1:end 1]),'k',LW,.8), axis(scl*[-1 1 -1 1]), axis square, hold on
   grid on, plot(pol+1e-12i,'.',MS,5,CO,red)
   title poles, hold off, set(gca,FS,7)
end
end

function ZS = ZS(S,p)                          % SAMPLE POINTS
T = sort(angle(S)); d = (1:p)/(p+1);
m = length(T);
dT = diff([T T(1)+2*pi]);
ZT = T + d'.*dT; ZS = exp(1i*ZT(:));
end

function [pol,res,zer] = prz(zj,fj,wj)         % POLES, RESIDUES, ZEROS
m = length(wj);                                %    (from Chebfun)
if any(wj==0)
   ii = find(wj~=0); m = length(ii);
   zj = zj(ii); fj = fj(ii); wj = wj(ii);
end   
B = eye(m+1); B(1,1) = 0;
E = [0 wj.'; ones(m,1) diag(zj)];
pol = eig(E,B); pol = pol(~isinf(pol));
if nargout >= 2
   N = @(t)(1./(t-zj.'))*(fj.*wj);
   Ddiff = @(t) -((1./(t-zj.')).^2)*wj;
   res = N(pol)./Ddiff(pol);
   E = [0 (wj.*fj).'; ones(m,1) diag(zj)];
   zer = eig(E,B); zer = zer(~isinf(zer));
end
end

function r = reval(zz,zj,fj,wj)                % FUNCTION HANDLE FOR R
z = zz(:); C = 1./(z-zj.');                    %    (from Chebfun)    
r = (C*(wj.*fj))./(C*wj);
r(isinf(z)) = sum(wj.*fj)./sum(wj);
ii = find(isnan(r));
for j = 1:length(ii)
   r(ii(j)) = fj(z(ii(j))==zj);
end
r = reshape(r,size(zz));
end
