function [A, Gamma, B, ok, logS] = smith_invariant_f2d_twopass(Gin, opts)
% SMITH_INVARIANT_F2D_twopass  Smith invariant-factor form over F2[D] 
%    (low->high coeff vectors).
% [A, Gamma, B, ok, logS] = smith_invariant_f2d_twopass(Gin, opts);
% Produces Gamma with a k x k diagonal and all trailing (n-k) columns zero.
% Returns undo operators A,B so that A*Gamma*B == original G.
% J.Cioffi and ChatGPT - September 24, 2025
%
% INPUT
%   Gin  : k x n cell array; entries are 0, 1, or row-vectors [c0 c1 ...] over GF(2), low->high in D.
%   opts : struct (all optional)
%          .verbose_steps  (default 12)   -- print first N steps
%          .print_poly     (default true) -- pretty polynomial printing
%          .max_iters      (default 40000)
%
% OUTPUT
%   A, B : unimodular over F2[D] s.t. A*Gamma*B == Gin
%   Gamma: diagonal (Smith invariant factors) with last (n-k) columns exactly zero
%   ok   : boolean reconstruction check
%   logS : struct log with fields type,i,j,q,E_or_F,G_after,side,stage,step
%
% Always runs TWO passes of the F2[D] Smith reduction and composes A,B,
% then does optional polish passes if the leading block is not strictly diagonal.
% ========================================================================

    if nargin < 2
        opts = struct;
    end
    if ~isfield(opts,'verbose_steps'); opts.verbose_steps = 12; end
    if ~isfield(opts,'print_poly');    opts.print_poly    = true; end
    if ~isfield(opts,'max_iters');     opts.max_iters     = 40000; end
    if ~isfield(opts,'max_polish');    opts.max_polish    = 2; end
    if ~isfield(opts,'step0');         opts.step0         = 0; end

    % -------- PASS 1 --------
    if opts.verbose_steps > 0
        fprintf('=== PASS 1 ===\n');
    end
    [A1, G1, B1, log1, step_end] = smith_core(Gin, opts, opts.step0);

    % -------- PASS 2 (forced) --------
    if opts.verbose_steps > 0
        fprintf('=== PASS 2 (start step %d) ===\n', step_end+1);
    end
    [A2, G2, B2, log2, step_end] = smith_core(G1, opts, step_end);

    % Compose totals
    Aacc = pmat_mul(A1, A2);
    Bacc = pmat_mul(B2, B1);
    Gacc = G2;
    logS = [log1, log2];

    % -------- Optional polish passes --------
    passes = 0;
    [bad,i0,j0,val] = find_first_offdiag_or_tail(Gacc);
    while bad && passes < opts.max_polish
        passes = passes + 1;
        if opts.verbose_steps > 0
            fprintf('--- POLISH pass %d (start step %d); offending G(%d,%d)=%s ---\n', ...
                passes, step_end+1, i0, j0, poly_to_str(val));
        end
        [Aa, Ga, Ba, loga, step_end] = smith_core(Gacc, opts, step_end);
        Aacc = pmat_mul(Aacc, Aa);
        Bacc = pmat_mul(Ba,  Bacc);
        Gacc = Ga;
        logS = [logS, loga];
        [bad,i0,j0,val] = find_first_offdiag_or_tail(Gacc);
    end

    Gamma = pmat_normalize(Gacc);
    A = Aacc;
    B = Bacc;
    ok = pmat_equal( pmat_mul(pmat_mul(A, Gamma), B), pmat_normalize(Gin) );

    if opts.verbose_steps > 0
        fprintf('Final check ok = %d\n', ok);
        if opts.print_poly
            disp('Gamma (k-diagonal, tail columns zero):');
            pmat_disp(Gamma, true);
        end
    end
end

% ========================================================================
%                               CORE PASS
% ========================================================================
function [Aret, G, Bret, logS, step] = smith_core(Gin, opts, step0)
    G = pmat_normalize(Gin);
    sz = size(G);
    k = sz(1); n = sz(2);
    Aret = pmat_eye(k);
    Bret = pmat_eye(n);
    step = step0;

    logS = struct('type',{},'i',{},'j',{},'q',{},'E_or_F',{},'G_after',{},'side',{},'stage',{},'step',{});

    smax  = min(k,n);
    iters = 0;

    for piv = 1:smax
        stage_k = piv;

        [pi, pj] = find_best_pivot(G, piv);
        if isempty(pi)
            break;
        end

        if pi ~= piv
            [E, G] = row_swap(G, pi, piv);
            Aret = pmat_mul(Aret, E);
            step = step + 1;
            logS(end+1) = packlog('row_swap',pi,piv,[],E,G,'L',stage_k,step);
            vprintf(opts,step,'[step %d | stage k=%d] LEFT: row_swap (%d <-> %d)\n', step, stage_k, pi, piv);
            vdispEFG(opts, step, 'E:', E, 'G:', G);
            vdispAB(opts, step, Aret, Bret);
        end

        if pj ~= piv
            [F, G] = col_swap(G, pj, piv);
            Bret = pmat_mul(F, Bret);
            step = step + 1;
            logS(end+1) = packlog('col_swap',pj,piv,[],F,G,'R',stage_k,step);
            vprintf(opts,step,'[step %d | stage k=%d] RIGHT: col_swap (%d <-> %d)\n', step, stage_k, pj, piv);
            vdispEFG(opts, step, 'F:', F, 'G:', G);
            vdispAB(opts, step, Aret, Bret);

            [G, Aret, logS, step] = clear_leading_block(G, Aret, Bret, 1, piv, opts, logS, step);
        end

        % Clear below pivot column
        changed = true;
        while changed
            changed = false;
            iters = iters + 1;
            if iters > opts.max_iters
                error('Exceeded max_iters');
            end
            for r = piv+1:k
                if ~p_is_zero(G{r,piv})
                    while ~p_is_zero(G{r,piv})
                        [q, ~, rmd] = p_divmod(G{r,piv}, G{piv,piv});
                        if ~p_is_zero(q)
                            [E, G] = row_add_multiple(G, r, piv, q);
                            Aret = pmat_mul(Aret, E);
                            step = step + 1;
                            logS(end+1) = packlog('row_add',r,piv,q,E,G,'L',stage_k,step);
                            vprintf(opts,step,'[step %d | stage k=%d] LEFT: row_add r=%d <- r + q*%d  (q=%s)\n', step, stage_k, r, piv, poly_to_str(q));
                            vdispEFG(opts, step, 'E:', E, 'G:', G);
                            vdispAB(opts, step, Aret, Bret);
                            changed = true;
                        end
                        if ~p_is_zero(rmd)
                            [E2, G] = row_swap(G, r, piv);
                            Aret = pmat_mul(Aret, E2);
                            step = step + 1;
                            logS(end+1) = packlog('row_swap',r,piv,[],E2,G,'L',stage_k,step);
                            vprintf(opts,step,'[step %d | stage k=%d] LEFT: row_swap promote remainder (r=%d <-> piv=%d)\n', step, stage_k, r, piv);
                            vdispEFG(opts, step, 'E:', E2, 'G:', G);
                            vdispAB(opts, step, Aret, Bret);
                            changed = true;
                        else
                            break;
                        end
                    end
                end
            end
        end

        % Clear to the right in pivot row
        for c = piv+1:n
            while ~p_is_zero(G{piv,c})
                iters = iters + 1;
                if iters > opts.max_iters
                    error('Exceeded max_iters');
                end
                [q, ~, rmd] = p_divmod(G{piv,c}, G{piv,piv});
                if ~p_is_zero(q)
                    [F, G] = col_add_multiple(G, c, piv, q);
                    Bret = pmat_mul(F, Bret);
                    step = step + 1;
                    logS(end+1) = packlog('col_add',c,piv,q,F,G,'R',stage_k,step);
                    vprintf(opts,step,'[step %d | stage k=%d] RIGHT: col_add c=%d <- c + q*%d  (q=%s)\n', step, stage_k, c, piv, poly_to_str(q));
                    vdispEFG(opts, step, 'F:', F, 'G:', G);
                    vdispAB(opts, step, Aret, Bret);
                end
                if ~p_is_zero(rmd)
                    [F2, G] = col_swap(G, c, piv);
                    Bret = pmat_mul(F2, Bret);
                    step = step + 1;
                    logS(end+1) = packlog('col_swap',c,piv,[],F2,G,'R',stage_k,step);
                    vprintf(opts,step,'[step %d | stage k=%d] RIGHT: col_swap promote remainder (c=%d <-> piv=%d)\n', step, stage_k, c, piv);
                    vdispEFG(opts, step, 'F:', F2, 'G:', G);
                    vdispAB(opts, step, Aret, Bret);

                    [G, Aret, logS, step] = clear_leading_block(G, Aret, Bret, 1, piv, opts, logS, step);
                else
                    break;
                end
            end
        end

        % Ensure pivot column is clean before moving on
        [G, Aret, logS, step] = clear_pivot_col(G, Aret, piv, stage_k, opts, logS, step);
    end
end

% ========================================================================
%                       Utility: clear leading block
% ========================================================================
function [G, Aret, logS, step] = clear_leading_block(G, Aret, Bret, p_from, p_to, opts, logS, step)
    for pchk = p_from:p_to
        [G, Aret, logS, step] = clear_pivot_col(G, Aret, pchk, pchk, opts, logS, step);
    end
end

% ========================================================================
%             Utility: clear pivot column (below and above)
% ========================================================================
function [G, Aret, logS, step] = clear_pivot_col(G, Aret, p, stage_k, opts, logS, step)
    sz = size(G);
    k = sz(1);

    % below
    for rr = p+1:k
        while ~p_is_zero(G{rr,p})
            [q2, ~, rmd2] = p_divmod(G{rr,p}, G{p,p});
            if ~p_is_zero(q2)
                [E3, G] = row_add_multiple(G, rr, p, q2);
                Aret = pmat_mul(Aret, E3);
                step = step + 1;
                logS(end+1) = packlog('row_add',rr,p,q2,E3,G,'L',stage_k,step);
                vprintf(opts,step,'[step %d | stage k=%d] LEFT: row_add (clear col %d) r=%d <- r + q*%d  (q=%s)\n', step, stage_k, p, rr, p, poly_to_str(q2));
                vdispEFG(opts, step, 'E:', E3, 'G:', G);
            end
            if ~p_is_zero(rmd2)
                [E4, G] = row_swap(G, rr, p);
                Aret = pmat_mul(Aret, E4);
                step = step + 1;
                logS(end+1) = packlog('row_swap',rr,p,[],E4,G,'L',stage_k,step);
                vprintf(opts,step,'[step %d | stage k=%d] LEFT: row_swap (promote remainder) r=%d <-> piv=%d\n', step, stage_k, rr, p);
                vdispEFG(opts, step, 'E:', E4, 'G:', G);
            else
                break;
            end
        end
    end

    % above
    for rr = 1:p-1
        while ~p_is_zero(G{rr,p})
            [q2, ~, rmd2] = p_divmod(G{rr,p}, G{p,p});
            if ~p_is_zero(q2)
                [E3, G] = row_add_multiple(G, rr, p, q2);
                Aret = pmat_mul(Aret, E3);
                step = step + 1;
                logS(end+1) = packlog('row_add',rr,p,q2,E3,G,'L',stage_k,step);
                vprintf(opts,step,'[step %d | stage k=%d] LEFT: row_add (clear above col %d) r=%d <- r + q*%d  (q=%s)\n', step, stage_k, p, rr, p, poly_to_str(q2));
                vdispEFG(opts, step, 'E:', E3, 'G:', G);
            end
            if ~p_is_zero(rmd2)
                [E4, G] = row_swap(G, rr, p);
                Aret = pmat_mul(Aret, E4);
                step = step + 1;
                logS(end+1) = packlog('row_swap',rr,p,[],E4,G,'L',stage_k,step);
                vprintf(opts,step,'[step %d | stage k=%d] LEFT: row_swap (promote above) r=%d <-> piv=%d\n', step, stage_k, rr, p);
                vdispEFG(opts, step, 'E:', E4, 'G:', G);
            else
                break;
            end
        end
    end
end

% ========================================================================
%                         Pivot selection
% ========================================================================
function [ri, cj] = find_best_pivot(G, piv)
    sz = size(G);
    k = sz(1); n = sz(2);
    best_deg = inf;
    ri = [];
    cj = [];
    for r = piv:k
        for c = piv:n
            if ~p_is_zero(G{r,c})
                d = p_deg(G{r,c});
                if d < best_deg
                    best_deg = d;
                    ri = r;
                    cj = c;
                    if best_deg == 0
                        return;
                    end
                end
            end
        end
    end
end

% ========================================================================
%                              Logging record
% ========================================================================
function rec = packlog(type,i,j,q,EF,Gnow,side,stage,step)
    rec = struct();
    rec.type   = type;
    rec.i      = i;
    rec.j      = j;
    rec.q      = p_copy(q);
    rec.E_or_F = pmat_copy_and_trim(EF);
    rec.G_after= pmat_copy_and_trim(Gnow);
    rec.side   = side;
    rec.stage  = stage;
    rec.step   = step;
end

% ========================================================================
%                         Pretty printing helpers
% ========================================================================
function vprintf(opts, step, varargin)
    if step <= opts.verbose_steps
        fprintf(varargin{:});
    end
end

function vdispEFG(opts, step, lblEF, EF, lblG, G)
    if step <= opts.verbose_steps
        disp(lblEF);
        pmat_disp(EF, opts.print_poly);
        disp(lblG);
        pmat_disp(G,  opts.print_poly);
    end
end

function vdispAB(opts, step, A, B)
    if step <= opts.verbose_steps
        disp('Cumulative A (returned):');
        pmat_disp(A, true);
        disp('Cumulative B (returned):');
        pmat_disp(B, true);
    end
end

% ========================================================================
%               Detector: off-diagonals / tail nonzeros
% ========================================================================
function [bad,i0,j0,val] = find_first_offdiag_or_tail(G)
    sz = size(G);
    k = sz(1); n = sz(2);
    kk = min(k,n);
    bad = false; i0 = []; j0 = []; val = [];

    % off-diagonals in leading kk x kk
    for i = 1:kk
        for j = 1:kk
            if i ~= j && ~p_is_zero(G{i,j})
                bad = true; i0 = i; j0 = j; val = G{i,j};
                return;
            end
        end
    end
    % tail columns
    for j = kk+1:n
        for i = 1:k
            if ~p_is_zero(G{i,j})
                bad = true; i0 = i; j0 = j; val = G{i,j};
                return;
            end
        end
    end
end

% ========================================================================
%                      Elementary poly-matrix ops
% ========================================================================
function I = pmat_eye(m)
    I = cell(m,m);
    ii = 1;
    while ii <= m
        jj = 1;
        while jj <= m
            if ii == jj
                I{ii,jj} = poly_one();
            else
                I{ii,jj} = poly_zero();
            end
            jj = jj + 1;
        end
        ii = ii + 1;
    end
end

function Z = pmat_zero(m,n)
    Z = cell(m,n);
    ii = 1;
    while ii <= m
        jj = 1;
        while jj <= n
            Z{ii,jj} = poly_zero();
            jj = jj + 1;
        end
        ii = ii + 1;
    end
end

function G2 = pmat_mul(A,B)
    szA = size(A);
    szB = size(B);
    m = szA(1); p = szA(2);
    p2 = szB(1); n = szB(2);
    if p ~= p2
        error('Dimension mismatch in pmat_mul');
    end
    G2 = pmat_zero(m,n);
    i = 1;
    while i <= m
        j = 1;
        while j <= n
            acc = poly_zero();
            t = 1;
            while t <= p
                if ~p_is_zero(A{i,t}) && ~p_is_zero(B{t,j})
                    acc = p_add(acc, p_mul(A{i,t}, B{t,j}));
                end
                t = t + 1;
            end
            G2{i,j} = p_norm(acc);
            j = j + 1;
        end
        i = i + 1;
    end
end

function G2 = pmat_normalize(G)
    sz = size(G);
    m = sz(1); n = sz(2);
    G2 = cell(m,n);
    i = 1;
    while i <= m
        j = 1;
        while j <= n
            G2{i,j} = p_norm(G{i,j});
            j = j + 1;
        end
        i = i + 1;
    end
end

function tf = pmat_equal(A,B)
    szA = size(A);
    m = szA(1); n = szA(2);
    if ~isequal(size(B),[m,n])
        tf = false; return;
    end
    tf = true;
    i = 1;
    while i <= m
        j = 1;
        while j <= n
            if ~p_equal(A{i,j}, B{i,j})
                tf = false; return;
            end
            j = j + 1;
        end
        i = i + 1;
    end
end

function Xout = pmat_copy_and_trim(Xin)
    Xout = Xin;
end

function [E, G2] = row_swap(G, i, j)
    sz = size(G);
    k = sz(1);
    E = pmat_eye(k);
    E{i,i} = poly_zero();
    E{j,j} = poly_zero();
    E{i,j} = poly_one();
    E{j,i} = poly_one();
    G2 = pmat_mul(E, G);
end

function [F, G2] = col_swap(G, i, j)
    sz = size(G);
    n = sz(2);
    F = pmat_eye(n);
    F{i,i} = poly_zero();
    F{j,j} = poly_zero();
    F{i,j} = poly_one();
    F{j,i} = poly_one();
    G2 = pmat_mul(G, F);
end

function [E, G2] = row_add_multiple(G, i, j, q)
    sz = size(G);
    k = sz(1);
    E = pmat_eye(k);
    E{i,j} = p_copy(q);
    G2 = pmat_mul(E, G);
end

function [F, G2] = col_add_multiple(G, i, j, q)
    sz = size(G);
    n = sz(2);
    F = pmat_eye(n);
    F{j,i} = p_copy(q);
    G2 = pmat_mul(G, F);
end

% ========================================================================
%                           GF(2)[D] poly ops
% ========================================================================
function p = poly_zero()
    p = 0;
end

function p = poly_one()
    p = 1;
end

function q = p_copy(p)
    if nargin == 0
        q = [];
    else
        q = p;
    end
end

function p = p_norm(p)
    if isempty(p)
        p = 0;
        return;
    end
    p = double(p(:).');
    p = mod(p,2);
    idx = find(p~=0,1,'last');
    if isempty(idx)
        p = 0;
    else
        p = p(1:idx);
    end
end

function tf = p_is_zero(p)
    if isempty(p)
        tf = true;
        return;
    end
    if isscalar(p) && p == 0
        tf = true;
        return;
    end
    tf = isrow(p) && all(p == 0);
end

function tf = p_equal(a,b)
    aa = p_norm(a);
    bb = p_norm(b);
    tf = isequal(aa, bb);
end

function d = p_deg(p)
    pp = p_norm(p);
    if p_is_zero(pp)
        d = -inf;
    else
        d = numel(pp) - 1;
    end
end

function c = p_add(a,b)
    aa = p_norm(a);
    bb = p_norm(b);
    L = max(numel(aa), numel(bb));
    A = zeros(1,L);
    B = zeros(1,L);
    A(1:numel(aa)) = aa;
    B(1:numel(bb)) = bb;
    c = p_norm(mod(A + B, 2));
end

function c = p_mul(a,b)
    aa = p_norm(a);
    bb = p_norm(b);
    if p_is_zero(aa) || p_is_zero(bb)
        c = 0;
        return;
    end
    c = p_norm(mod(conv(aa, bb), 2));
end

function [q, d, r] = p_divmod(a,b)
    aa = p_norm(a);
    bb = p_norm(b);
    if p_is_zero(bb)
        error('Division by zero polynomial');
    end
    if p_is_zero(aa)
        q = 0; r = 0; d = [];
        return;
    end
    ra = aa;
    rb = bb;
    qa = 0;
    while ~p_is_zero(ra) && p_deg(ra) >= p_deg(rb)
        shift = p_deg(ra) - p_deg(rb);
        term  = p_Dpow(shift);
        qa = p_add(qa, term);
        ra = p_add(ra, p_mul(rb, term));
    end
    q = p_norm(qa);
    r = p_norm(ra);
    d = [];
end

function t = p_Dpow(s)
    if ~isscalar(s) || s < 0 || ~isfinite(s)
        error('p_Dpow: nonneg int required');
    end
    t = zeros(1, s+1);
    t(end) = 1;
end

% ========================================================================
%                               Display
% ========================================================================
function s = poly_to_str(p)
    pp = p_norm(p);
    if p_is_zero(pp)
        s = '0';
        return;
    end
    parts = {};
    i = 1;
    while i <= numel(pp)
        if pp(i) ~= 0
            if i == 1
                parts{end+1} = '1'; %#ok<AGROW>
            elseif i == 2
                parts{end+1} = 'D'; %#ok<AGROW>
            else
                parts{end+1} = sprintf('D^%d', i-1); %#ok<AGROW>
            end
        end
        i = i + 1;
    end
    parts = fliplr(parts);
    s = strjoin(parts, ' + ');
end

function pmat_disp(M, pretty)
    if nargin < 2
        pretty = true;
    end
    sz = size(M);
    m = sz(1); n = sz(2);
    i = 1;
    while i <= m
        line = '';
        j = 1;
        while j <= n
            if pretty
                token = poly_to_str(M{i,j});
            else
                token = mat2str(p_norm(M{i,j}));
            end
            line = [line, sprintf('%-14s', token)]; %#ok<AGROW>
            j = j + 1;
        end
        disp(['  ', line]);
        i = i + 1;
    end
end
