using Graphs.Experimental.Traversals
using BipartiteGraphs: set_neighbors!

function extreme_var(var_to_diff, v, level = nothing, ::Val{descend} = Val(true);
    callback = _ -> nothing) where {descend}
g = descend ? invview(var_to_diff) : var_to_diff
callback(v)
while (v′ = g[v]) !== nothing
    v::Int = v′
    callback(v)
    if level !== nothing
        descend ? (level -= 1) : (level += 1)
    end
end
level === nothing ? v : (v => level)
end

function structural_singularity_removal!(
        state::TransformationState, ::Val{ReturnPivots} = Val{false}();
        variable_underconstrained! = force_var_to_zero!, kwargs...
    ) where {ReturnPivots}
    mm = linear_subsys_adjmat!(state; kwargs...)
    if size(mm, 1) == 0
        if ReturnPivots
            return mm, PivotInfo(0, 0, Int[])
        else
            return mm # No linear subsystems
        end
    end

    (; graph, var_to_diff, solvable_graph) = state.structure
    mm, pivotinfo = structural_singularity_removal!(state, mm, Val{true}(); variable_underconstrained!)
    s = state.structure
    for (ei, e) in enumerate(mm.nzrows)
        set_neighbors!(s.graph, e, mm.row_cols[ei])
    end
    if s.solvable_graph isa BipartiteGraph{Int, Nothing}
        for (ei, e) in enumerate(mm.nzrows)
            set_neighbors!(s.solvable_graph, e, mm.row_cols[ei])
        end
    end

    if ReturnPivots
        return mm, pivotinfo
    else
        return mm
    end
end

# For debug purposes
function aag_bareiss(sys)
    state = TearingState(sys)
    complete!(state.structure)
    mm = linear_subsys_adjmat!(state)
    return aag_bareiss!(state.structure.graph, state.structure.var_to_diff, mm)
end


"""
$(SIGNATURES)

Find a variable (column) and its coefficient (value) in `M` among equations (rows) in `range`,
filtering out equations for which `mask` is `false` (`mask` can be `nothing` to avoid masking).
In case of a tie, `var_priorities` is used to choose a variable with lower priority. In case
priorities are tied, it will prefer the row with fewer elements.
"""
@inline function find_first_linear_variable(
        M::SparseMatrixCLIL,
        range,
        mask,
        var_priorities = nothing
    )
    eadj = M.row_cols
    candidate_i = 0
    candidate_v = 0
    candidate_val = 0
    candidate_nnz = 0
    @inbounds for i in range
        vertices = eadj[i]
        nnz = length(vertices)
        if !iszero(candidate_v) && var_priorities === nothing && nnz >= candidate_nnz
            continue
        end
        for (j, v) in enumerate(vertices)
            mask === nothing || mask[v] || continue
            # Prefer, in order
            # 1. Lower priority pivots
            # 2. Rows with fewer elements
            if iszero(candidate_v) || var_priorities === nothing && nnz < candidate_nnz ||
                    var_priorities !== nothing && (
                    var_priorities[v] < var_priorities[candidate_v] ||
                        var_priorities[v] == var_priorities[candidate_v] && nnz < candidate_nnz
                )
                candidate_i = i
                candidate_v = v
                candidate_val = M.row_vals[i][j]
                candidate_nnz = nnz
            end
        end
    end
    iszero(candidate_v) || return CartesianIndex(candidate_i, candidate_v), candidate_val
    return nothing
end

@inline function find_first_linear_variable(M::AbstractMatrix,
        range,
        mask,
        var_priorities = nothing)
    candidate_i = 0
    candidate_v = 0
    candidate_val = 0
    candidate_nnz = 0
    @inbounds for i in range
        row = @view M[i, :]
        nnz = count(!izero, row)
        if !izero(candidate_v) && var_priorities === nothing && nnz >= candidate_nnz
            continue
        end
        for (v, val) in enumerate(row)
            mask === nothing || mask[v] || continue
            if iszero(candidate_v) || (var_priorities === nothing && nnz < candidate_nnz || var_priorities !== nothing && (var_priorities[v] < var_priorities[candidate_v] || var_priorities[v] == var_priorities[candidate_v] && nnz < candidate_nnz))
                candidate_i = i
                candidate_v = v
                candidate_val = val
                candidate_nnz = nnz
            end
        end
    end
    return nothing
end

function find_masked_pivot(variables, M, k, var_priorities)
    return find_first_linear_variable(M, k:size(M, 1), variables, var_priorities)
end

count_nonzeros(a::AbstractArray) = count(!iszero, a)

# N.B.: Ordinarily sparse vectors allow zero stored elements.
# Here we have a guarantee that they won't, so we can make this identification
count_nonzeros(a::CLILVector) = nnz(a)

"""
    $(TYPEDSIGNATURES)

Mark variable `v` as non-linear and propagate that marking transitively to all
variables that co-appear with `v` in a linear equation.  `stack` is a
pre-allocated work buffer (cleared and reused by this function).  `var_to_lineq`
maps each variable to the set of linear equations it appears in.
"""
function mark_not_linear!(linear_variables::BitVector, stack::Vector{Int},
        var_to_lineq::Dict{Int, BitSet}, graph, v::Int)
    linear_variables[v] = false
    push!(stack, v)
    while !isempty(stack)
        v = pop!(stack)
        eqs = get(var_to_lineq, v, nothing)
        eqs === nothing && continue
        for eq in eqs, v′ in 𝑠neighbors(graph, eq)
            if linear_variables[v′]
                linear_variables[v′] = false
                push!(stack, v′)
            end
        end
    end
end

# Linear variables are highest order differentiated variables that only appear
# in linear equations with only linear variables. Also, if a variable's any
# derivatives is nonlinear, then all of them are not linear variables.
function find_linear_variables(graph, linear_equations, var_to_diff, irreducibles)
    stack = Int[]
    linear_variables = falses(length(var_to_diff))
    var_to_lineq = Dict{Int, BitSet}()
    for eq in linear_equations, v in 𝑠neighbors(graph, eq)
        linear_variables[v] = true
        vlineqs = get!(() -> BitSet(), var_to_lineq, v)
        push!(vlineqs, eq)
    end
    for v in irreducibles
        lv = extreme_var(var_to_diff, v)
        while true
            mark_not_linear!(linear_variables, stack, var_to_lineq, graph, lv)
            lv = var_to_diff[lv]
            lv === nothing && break
        end
    end

    linear_equations_set = BitSet(linear_equations)
    for (v, islinear) in enumerate(linear_variables)
        islinear || continue
        lv = extreme_var(var_to_diff, v)
        oldlv = lv
        remove = invview(var_to_diff)[v] !== nothing
        while !remove
            for eq in 𝑑neighbors(graph, lv)
                if !(eq in linear_equations_set)
                    remove = true
                end
            end
            lv = var_to_diff[lv]
            lv === nothing && break
        end
        lv = oldlv
        remove && while true
            mark_not_linear!(linear_variables, stack, var_to_lineq, graph, lv)
            lv = var_to_diff[lv]
            lv === nothing && break
        end
    end

    return linear_variables
end

"""
    $TYPEDEF

Comparator intended for use when sorting the rows of a `SparseMatrixCLIL` before
running Bareiss.
"""
struct MMSortKey{M <: SparseMatrixCLIL}
    mm::M
end

function (k::MMSortKey)(i::Integer)
    (; mm) = k
    # Sort by sparsity, earlier variable coefficient, smaller coefficient and finally by
    # equation index.
    return (length(mm.row_cols[i]), mm.row_cols[i], mm.row_vals[i], mm.nzrows[i])
end

"""
    $TYPEDSIGNATURES

Sort the rows of `mm` to obtain reproducible results from Bareiss. Uses `MMSortKey` as
the metric to sort on. The heuristic depends upon the variable ordering.
"""
function sort_mm_rows(mm::SparseMatrixCLIL{T, Ti}) where {T, Ti}
    rperm = sortperm(eachindex(mm.nzrows); by = MMSortKey(mm))
    return SparseMatrixCLIL{T, Ti}(
        mm.nparentrows, mm.ncols, mm.nzrows[rperm], mm.row_cols[rperm], mm.row_vals[rperm]
    )
end

"""
    $(TYPEDSIGNATURES)

Return `true` if variable `v` is purely algebraic, i.e. it has no derivative
and is not itself a derivative of another variable.
"""
is_algebraic(var_to_diff, v::Int)::Bool =
    var_to_diff[v] === nothing === invview(var_to_diff)[v]

function aag_bareiss!(structure, mm_orig::SparseMatrixCLIL{T, Ti}) where {T, Ti}
    (; graph, var_to_diff) = structure
    mm = copy(mm_orig)
    linear_equations_set = BitSet(mm_orig.nzrows)

    # All unassigned (not a pivot) algebraic variables that only appears in
    # linear algebraic equations can be set to 0.
    #
    # For all the other variables, we can update the original system with
    # Bareiss'ed coefficients as Gaussian elimination is nullspace preserving
    # and we are only working on linear homogeneous subsystem.

    is_linear_variables = Base.Fix1(is_algebraic, var_to_diff).(1:length(var_to_diff))
    is_highest_diff = computed_highest_diff_variables(structure)
    for i in 𝑠vertices(graph)
        # only consider linear algebraic equations
        (i in linear_equations_set &&
            all(Base.Fix1(is_algebraic, var_to_diff), 𝑠neighbors(graph, i))) &&
            continue
        for j in 𝑠neighbors(graph, i)
            is_linear_variables[j] = false
        end
    end
    solvable_variables = findall(is_linear_variables)
    sp = has_state_priorities(structure) ? get_state_priorities(structure) : nothing
    cr = has_canonical_ranks(structure) ? get_canonical_ranks(structure) : nothing
    var_priorities = if cr === nothing
        sp
    elseif sp === nothing
        cr
    else
        big = maximum(cr; init = 0) + 1
        Int[sp[i] * big + cr[i] for i in eachindex(sp)]
    end

    # `do_bareiss!` performs a *global* pivot search over the whole matrix,
    # which will happily swap rows across block boundaries when the matrix is
    # block-diagonal (e.g. several independent submodels that share no
    # variables). Such cross-block swaps don't change correctness, but they
    # perturb the within-block elimination order relative to what that block
    # would get if it were factored in isolation, which can synthesize
    # fill-in-heavy (and, downstream during tearing, near-singular) reduced
    # rows for blocks that would otherwise reduce cleanly on their own. Since
    # rows that share no column are mathematically independent linear
    # subsystems, we instead partition rows into connected components (rows
    # connected iff they share a column) and Bareiss-eliminate each component
    # separately; this is exact and immune to cross-component perturbation.
    nrows = size(mm, 1)
    parent = collect(1:nrows)
    col_to_row = Dict{Ti, Int}()
    for i in 1:nrows
        for c in mm.row_cols[i]
            j = get(col_to_row, c, 0)
            if j == 0
                col_to_row[c] = i
            else
                _uf_union!(parent, i, j)
            end
        end
    end
    components = Dict{Int, Vector{Int}}()
    for i in 1:nrows
        r = _uf_find!(parent, i)
        push!(get!(() -> Int[], components, r), i)
    end

    if length(components) <= 1
        mm = sort_mm_rows(mm)
        bar = do_bareiss!(mm, mm_orig, is_linear_variables, is_highest_diff, var_priorities)
        return mm, solvable_variables, bar
    end

    # Deterministic component order: since rows were visited in increasing
    # order above, each component's row list is already sorted, so `first`
    # gives its smallest row index.
    comp_rows = sort(collect(values(components)); by = first)

    all_nzrows = Ti[]
    all_row_cols = Vector{Vector{Ti}}()
    all_row_vals = Vector{Vector{T}}()
    tier1_pivots = Int[]
    tier2_pivots = Int[]
    tier3_pivots = Int[]
    rank1 = 0
    rank2_extra = 0
    rank3_extra = 0
    for rows in comp_rows
        sort!(rows; by = MMSortKey(mm))
        sub = SparseMatrixCLIL{T, Ti}(
            mm.nparentrows, mm.ncols, mm.nzrows[rows], mm.row_cols[rows], mm.row_vals[rows])
        (rank1_c, rank2_c, rank3_c, pivots_c) = do_bareiss!(
            sub, nothing, is_linear_variables, is_highest_diff, var_priorities)

        append!(all_nzrows, sub.nzrows)
        append!(all_row_cols, sub.row_cols)
        append!(all_row_vals, sub.row_vals)

        append!(tier1_pivots, @view pivots_c[1:rank1_c])
        append!(tier2_pivots, @view pivots_c[(rank1_c + 1):rank2_c])
        append!(tier3_pivots, @view pivots_c[(rank2_c + 1):rank3_c])
        rank1 += rank1_c
        rank2_extra += rank2_c - rank1_c
        rank3_extra += rank3_c - rank2_c
    end
    pivots_merged = vcat(tier1_pivots, tier2_pivots, tier3_pivots)
    rank2 = rank1 + rank2_extra
    rank3 = rank2 + rank3_extra

    mm = SparseMatrixCLIL{T, Ti}(mm.nparentrows, mm.ncols, all_nzrows, all_row_cols, all_row_vals)

    return mm, solvable_variables, (rank1, rank2, rank3, pivots_merged)
end

"Column-swap no-op passed to `bareiss!` when using the virtual column-swap strategy."
noop_colswap(M, i, j) = nothing

"""
    $(TYPEDEF)

A callable that swaps rows `i` and `j` in the working matrix `M` and, if
`Mold` is not `nothing`, simultaneously applies the same swap to `Mold`.
Used to keep the original matrix in sync with the Bareiss working copy during
row operations.
"""
struct SyncedSwapRows{M}
    Mold::M
end
(s::SyncedSwapRows{Nothing})(M, i::Int, j::Int) = Base.swaprows!(M, i, j)
(s::SyncedSwapRows)(M, i::Int, j::Int) = (Base.swaprows!(s.Mold, i, j); Base.swaprows!(M, i, j))

"""
    $TYPEDEF

Lazy `&&` of two boolean masks. Only implements whatever is required for `find_masked_pivot`.
"""
struct LazyMaskAnd{V1 <: AbstractVector{Bool}, V2 <: AbstractVector{Bool}}
    mask1::V1
    mask2::V2
end

Base.getindex(lma::LazyMaskAnd, i::Integer) = lma.mask1[i] && lma.mask2[i]

"""
    $(TYPEDEF)

Mutable state threaded through the Bareiss factorization callbacks.

- `rank1`/`rank2`: record where each tier of pivot selection was exhausted,
  used to compute the three rank values returned by `do_bareiss!`.
- `pivots`: accumulates the column index of every pivot chosen during elimination.
- `is_linear_variables`/`is_highest_diff`: masks used for the tiered pivot search.
"""
mutable struct BareissContext{V1 <: AbstractVector{Bool}, V2 <: AbstractVector{Bool}, P <: Union{Nothing, AbstractVector{Int}}}
    rank1::Union{Nothing, Int}
    rank2::Union{Nothing, Int}
    pivots::Vector{Int}
    is_linear_variables::V1
    is_highest_diff::V2
    valid_pivot_mask::BitVector
    var_priorities::P
end

function BareissContext(is_linear_variables, is_highest_diff, var_priorities = nothing)
    return BareissContext(
        nothing, nothing, Int[], is_linear_variables, is_highest_diff,
        trues(length(is_linear_variables)), var_priorities
    )
end

"""
    $(TYPEDSIGNATURES)

Tiered pivot selection for Bareiss factorization.  Tries three strategies in
order, recording the rank at which each tier is exhausted:
1. A pivot among `is_linear_variables` columns (rank1 boundary).
2. A pivot among `is_highest_diff` columns (rank2 boundary).
3. Any pivot in the remaining matrix.
The column index of every selected pivot is appended to `ctx.pivots`.
"""
function (ctx::BareissContext)(M, k::Int)
    if ctx.rank1 === nothing
        mask = LazyMaskAnd(ctx.is_linear_variables, ctx.valid_pivot_mask)
        r = find_masked_pivot(ctx.is_linear_variables, M, k, ctx.var_priorities)
        if r !== nothing
            push!(ctx.pivots, r[1][2])
            return r
        end
        ctx.rank1 = k - 1
    end
    if ctx.rank2 === nothing
        mask = LazyMaskAnd(ctx.is_highest_diff, ctx.valid_pivot_mask)
        r = find_masked_pivot(ctx.is_highest_diff, M, k, ctx.var_priorities)
        if r !== nothing
            push!(ctx.pivots, r[1][2])
            return r
        end
        ctx.rank2 = k - 1
    end
    # TODO: It would be better to sort the variables by
    # derivative order here to enable more elimination
    # opportunities.
    r = find_masked_pivot(nothing, M, k, ctx.var_priorities)
    r !== nothing && push!(ctx.pivots, r[1][2])
    return r
end

struct BareissContextUpdate{C <: BareissContext, F}
    context::C
    inner_update::F
end

function (bcu::BareissContextUpdate)(zero!, M, k, swapto, pivot, last_pivot; kw...)
    ctx = bcu.context
    col = swapto[2]
    ctx.valid_pivot_mask[col] = false
    return bcu.inner_update(zero!, M, k, swapto, pivot, last_pivot; kw...)
end

function do_bareiss!(M, Mold, is_linear_variables::AbstractVector{Bool},
        is_highest_diff::AbstractVector{Bool}, var_priorities = nothing)
    ctx = BareissContext(is_linear_variables, is_highest_diff, var_priorities)
    update! = BareissContextUpdate(ctx, bareiss_update_virtual_colswap_mtk!)
    bareiss_ops = (noop_colswap, SyncedSwapRows(Mold), update!, bareiss_zero!)
    rank3, = bareiss!(M, bareiss_ops; find_pivot = ctx)
    rank2 = something(ctx.rank2, rank3)
    rank1 = something(ctx.rank1, rank2)
    (rank1, rank2, rank3, ctx.pivots)
end

function force_var_to_zero!(structure::SystemStructure, ils::SparseMatrixCLIL, v::Int)
    (; graph, solvable_graph, eq_to_diff) = structure
    @set! ils.nparentrows += 1
    push!(ils.nzrows, ils.nparentrows)
    push!(ils.row_cols, [v])
    push!(ils.row_vals, [convert(eltype(ils), 1)])
    add_vertex!(graph, SRC)
    add_vertex!(solvable_graph, SRC)
    add_edge!(graph, ils.nparentrows, v)
    add_edge!(solvable_graph, ils.nparentrows, v)
    add_vertex!(eq_to_diff)
    return ils
end

"""
    $TYPEDSIGNATURES

Information about the pivots chosen by Bareiss during `structural_singularity_removal!`.
This can be returned from `structural_singularity_removal!` by passing `Val(true)` as the last
positional argument.

$TYPEDFIELDS
"""
struct PivotInfo
    """
    The length of the prefix of `pivots` that is variables which _only_ occur in linear
    equations of the sort considered by this pass. These variables must be solved for
    using the integer coefficient equations considered by this pass.
    """
    n_linear_vars::Int
    """
    Number of elements in `pivots` after `n_linear_vars` corresponding to highest order
    derivative variables.
    """
    n_highest_diff_vars::Int
    """
    The list of pivots chosen by the Bareiss algorithm.
    """
    pivots::Vector{Int}
end

"""
    $(TYPEDEF)

Tiered pivot search used for exact Bareiss factorization of pure-integer SCC
subsets during tearing. Unlike [`BareissContext`](@ref), there is no
unrestricted final tier: once `tier2` is exhausted the factorization stops.
This guarantees every chosen pivot column is an eligible (solvable-for) variable.
"""
mutable struct RestrictedBareissContext{V1 <: AbstractVector{Bool}, V2 <: AbstractVector{Bool}, P <: Union{Nothing, AbstractVector{Int}}}
    pivots::Vector{Int}
    tier1::V1
    tier2::V2
    valid_pivot_mask::BitVector
    var_priorities::P
    tier1_done::Bool
end

function RestrictedBareissContext(tier1, tier2, var_priorities = nothing)
    return RestrictedBareissContext(
        Int[], tier1, tier2, trues(length(tier1)), var_priorities, false)
end

function (ctx::RestrictedBareissContext)(M, k::Int)
    if !ctx.tier1_done
        r = find_masked_pivot(LazyMaskAnd(ctx.tier1, ctx.valid_pivot_mask), M, k, ctx.var_priorities)
        if r !== nothing
            push!(ctx.pivots, r[1][2])
            return r
        end
        ctx.tier1_done = true
    end
    r = find_masked_pivot(LazyMaskAnd(ctx.tier2, ctx.valid_pivot_mask), M, k, ctx.var_priorities)
    r === nothing && return nothing
    push!(ctx.pivots, r[1][2])
    return r
end

struct RestrictedContextUpdate{C <: RestrictedBareissContext, F}
    context::C
    inner_update::F
end

function (bcu::RestrictedContextUpdate)(zero!, M, k, swapto, pivot, last_pivot; kw...)
    bcu.context.valid_pivot_mask[swapto[2]] = false
    return bcu.inner_update(zero!, M, k, swapto, pivot, last_pivot; kw...)
end

function _uf_find!(parent::Vector{Int}, x::Int)
    while parent[x] != x
        parent[x] = parent[parent[x]]
        x = parent[x]
    end
    return x
end

function _uf_union!(parent::Vector{Int}, a::Int, b::Int)
    ra = _uf_find!(parent, a)
    rb = _uf_find!(parent, b)
    ra == rb || (parent[max(ra, rb)] = min(ra, rb))
    return nothing
end

function structural_singularity_removal!(
        state::TransformationState, ils::SparseMatrixCLIL, ::Val{ReturnPivots} = Val{false}();
        variable_underconstrained! = force_var_to_zero!
    ) where {ReturnPivots}
    (; structure) = state
    (; graph, solvable_graph, var_to_diff, eq_to_diff) = state.structure
    # Step 1: Perform Bareiss factorization on the adjacency matrix of the linear
    #         subsystem of the system we're interested in.
    #
    ils, solvable_variables, (rank1, rank2, rank3, pivots) = aag_bareiss!(structure, ils)

    ## Step 2: Simplify the system using the Bareiss factorization
    rk1vars = BitSet(@view pivots[1:rank1])
    for v in solvable_variables
        v in rk1vars && continue
        ils = variable_underconstrained!(structure, ils, v)
    end

    if ReturnPivots
        return ils, PivotInfo(rank1, rank2, pivots)
    else
        return ils
    end
end
