## -*- texinfo -*-
## @deftypefn  {Function File} dsp_plot_pz (@var{ax}, @var{b}, @var{a})
## @deftypefnx {Function File} dsp_plot_pz (@var{ax}, @var{b}, @var{a}, @var{stab})
## @deftypefnx {Function File} dsp_plot_pz (@var{ax}, @var{b}, @var{a}, @var{stab}, @var{show_legend})
## Draw a @code{zplane}-style pole-zero diagram of the filter @var{b}/@var{a}
## on the axes @var{ax}.  Used by the detailed analysis window and by the
## HTML report so the two always agree.
##
## Two things distinguish this from plotting @code{roots (b)} and
## @code{roots (a)} directly:
##
## @itemize
## @item Both polynomials are padded with roots at the origin up to the
## filter order, so an FIR filter shows the @var{N} poles it has at
## @math{z = 0} (its denominator is the scalar 1, whose root set is empty)
## and an all-pole IIR shows its origin zeros.
## @item Coincident roots are drawn as a single marker annotated with their
## occurrence number, rather than as markers stacked on top of each other.
## @end itemize
##
## @var{stab} is a struct from @code{dsp_check_stability} (recomputed when
## omitted); it supplies the maximum-pole-radius circle.  @var{show_legend}
## defaults to true.
##
## Returns the legend handle (empty when @var{show_legend} is false) so a
## caller can reposition it -- the equal-aspect plot is often small, and an
## in-axes legend covers a quadrant of it.
##
## @seealso{dsp_check_stability, dsp_filter_order, dsp_draw_zplane}
## @end deftypefn

function leg = dsp_plot_pz(ax, b, a, stab, show_legend)
    if nargin < 4 || isempty(stab)
        stab = dsp_check_stability(b, a);
    end
    if nargin < 5 || isempty(show_legend)
        show_legend = true;
    end

    axes(ax);
    cla;

    % Pad both polynomials to the filter order so the roots at the origin
    % that the shorter one implies are represented explicitly.  This is what
    % puts an FIR filter's N poles at z = 0 on the plot.
    order = dsp_filter_order(b, a);
    [zeros_z, zero_mult] = group_roots(pad_poly(b, order));
    [poles_z, pole_mult] = group_roots(pad_poly(a, order));

    radii = abs(poles_z);
    theta = linspace(0, 2*pi, 256);

    if stab.stable
        stab_color = [0.13, 0.55, 0.13];
    else
        stab_color = [0.85, 0.15, 0.15];
    end

    plot(cos(theta), sin(theta), "k--", "LineWidth", 0.8); hold on;
    leg_items = {"Unit Circle"};

    if ~isempty(zeros_z)
        plot(real(zeros_z), imag(zeros_z), "bo", "MarkerSize", 9, "LineWidth", 2);
        leg_items{end+1} = sprintf("Zeros (%d)", sum(zero_mult));
    end

    safe = poles_z(radii < 0.95);
    warn = poles_z(radii >= 0.95 & radii < 1 - 1e-10);
    crit = poles_z(radii >= 1 - 1e-10);

    if ~isempty(safe)
        plot(real(safe), imag(safe), "x", "MarkerSize", 9, "LineWidth", 2, ...
             "Color", [0.13, 0.55, 0.13]);
        leg_items{end+1} = "Poles (safe)";
    end
    if ~isempty(warn)
        plot(real(warn), imag(warn), "x", "MarkerSize", 10, "LineWidth", 2.5, ...
             "Color", [0.93, 0.69, 0.13]);
        leg_items{end+1} = "Poles (near boundary)";
    end
    if ~isempty(crit)
        plot(real(crit), imag(crit), "x", "MarkerSize", 11, "LineWidth", 2.5, ...
             "Color", [0.85, 0.15, 0.15]);
        leg_items{end+1} = "Poles (unstable)";
    end

    if stab.max_radius > 0
        r = stab.max_radius;
        plot(r*cos(theta), r*sin(theta), ":", "LineWidth", 1.2, "Color", stab_color);
        leg_items{end+1} = sprintf("Max radius = %.3f", r);
    end

    % Occurrence numbers for coincident roots (zplane convention).
    annotate_mult(zeros_z, zero_mult, [0.00, 0.00, 0.80]);
    annotate_mult(poles_z, pole_mult, [0.60, 0.00, 0.00]);

    hold off;
    axis equal; grid on;

    pts = [zeros_z(:); poles_z(:)];
    lim = 1.2;
    if ~isempty(pts)
        lim = max(lim, 1.15 * max([abs(real(pts)); abs(imag(pts))]));
    end
    xlim([-lim, lim]);
    ylim([-lim, lim]);

    title("Pole-Zero Diagram");
    xlabel("Real Part");
    ylabel("Imaginary Part");

    leg = [];
    if show_legend
        leg = legend(leg_items{:}, "location", "northeast");
        set(leg, "fontsize", 8);
    end
end

% ------------------------------------------------------------------
% Helpers
% ------------------------------------------------------------------

function p = pad_poly(p, order)
    % Append zero coefficients so deg(p) == order; the extra roots sit at
    % the origin, which is exactly what the missing z^k factor represents.
    p = p(:).';
    p = [p, zeros(1, order - (numel(p) - 1))];
end

function [reps, mult] = group_roots(p)
    % Return the distinct roots of p and how many times each occurs.
    %
    % Root multiplicity cannot be read off the numerical root scatter: an
    % 8-fold zero of a Butterworth numerator comes back from roots() spread
    % over ~1.6e-2, while the 8 genuinely distinct poles of the matching
    % denominator sit ~2.3e-2 apart.  So candidate clusters are formed
    % generously and then *verified* algebraically -- p is deflated by
    % (z - c)^m and the group is only accepted as an m-fold root when the
    % remainder vanishes.  For a true multiple root the centroid c is
    % accurate to machine precision (it is a coefficient ratio, unlike the
    % individual roots), which makes the residual ~1e-15 against ~1e-3 for
    % a cluster of distinct roots.
    cluster_tol = 0.05;
    verify_tol  = 1e-6;

    p = p(:).';
    r = roots(p);
    r = r(:);
    reps = zeros(0, 1);
    mult = zeros(0, 1);
    if isempty(r)
        return;
    end

    used = false(numel(r), 1);
    for k = 1:numel(r)
        if used(k); continue; end

        % Transitively grow a cluster of near-coincident roots.
        grp = ~used & (abs(r - r(k)) <= cluster_tol);
        changed = true;
        while changed
            changed = false;
            for j = find(grp)'
                add = ~used & ~grp & (abs(r - r(j)) <= cluster_tol);
                if any(add)
                    grp = grp | add;
                    changed = true;
                end
            end
        end
        used = used | grp;

        m = sum(grp);
        c = mean(r(grp));
        if m > 1 && ~divides(p, c, m, verify_tol)
            % Not a genuine multiple root -- keep the members separate.
            reps = [reps; r(grp)];
            mult = [mult; ones(m, 1)];
        else
            reps = [reps; c];
            mult = [mult; m];
        end
    end
end

function tf = divides(p, c, m, verify_tol)
    [~, rem] = deconv(p, poly(repmat(c, m, 1)));
    tf = norm(rem) <= verify_tol * norm(p);
end

function annotate_mult(reps, mult, col)
    idx = find(mult > 1);
    for k = idx(:)'
        text(real(reps(k)) + 0.07, imag(reps(k)) + 0.07, ...
             sprintf("%d", mult(k)), ...
             "Color", col, "FontSize", 9, "FontWeight", "bold");
    end
end
