text:
# Copyright (c) 2021, Eric Sabo
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

import Base: ==, show

#############################
         # General
#############################

mutable struct Vertex
    label::BigInt
    prev::Int64
    next::Int64
    value::Float64
    edgeloc::Int64
    polynomial::Vector{Vector{Int64}}
end

# could redo this for classical to remove sign
# for a trellis with 10 million edges, this would save 10 MB
# to do this, make EdgeC and EdgeQ and change all following functions to Union both
mutable struct Edge
    label::Union{fq_nmod, fq_nmod_mat}
    weight::Float64
    outvertex::Int64
    sign::Int8
end

mutable struct Trellis
    vertices::Vector{Vector{Vertex}}
    edges::Vector{Vector{Vector{Edge}}}
end

vertices(T::Trellis) = T.vertices
edges(T::Trellis) = T.edges

function ==(a::Vertex, b::Vertex)
    return a.label == b.label && a.prev == b.prev && a.next == b.next &&
        a.value == b.value && a.edgeloc == b.edgeloc && a.polynomial == b.polynomial
end

function ==(a::Vector{Vertex}, b::Vector{Vertex})
    if length(a) != length(b)
        return false
    end

    for i in 1:length(a)
        if a[i] != b[i]
            return false
        end
    end
    return true
end

function ==(a::Edge, b::Edge)
    return a.label == b.label && a.weight == b.weight && a.outvertex == b.outvertex
end

function ==(a::Vector{Edge}, b::Vector{Edge})
    if length(a) != length(b)
        return false
    end

    for i in 1:length(a)
        if a[i] != b[i]
            return false
        end
    end

    return true
end

function ==(a::Vector{Vector{Edge}}, b::Vector{Vector{Edge}})
    if length(a) != length(b)
        return false
    end

    for i in 1:length(a)
        if a[i] != b[i]
            return false
        end
    end

    return true
end

function ==(a::Vector{Vector{Vector{Edge}}}, b::Vector{Vector{Vector{Edge}}})
    if length(a) != length(b)
        return false
    end

    for i in 1:length(a)
        if a[i] != b[i]
            return false
        end
    end

    return true
end

function isisomorphic(V1::Vector{Vector{Vertex}}, V2::Vector{Vector{Vertex}})
    if length(V1) != length(V2)
        return false
    else
        for i in 1:length(V1)
            if length(V1[i]) != length(V2[i])
                return false
            end
        end
        return true
    end
end

function isisomorphic(E1::Vector{Vector{Vector{Edge}}}, E2::Vector{Vector{Vector{Edge}}})
    if length(E1) != length(E2)
        return false
    else
        for i in 1:length(E1)
            if length(E1[i]) != length(E2[i]) || length(E1[i][1]) != length(E2[i][1])
                return false
            end
        end
        return true
    end
end

function isisomorphic(V1::Vector{Vector{Vertex}}, E1::Vector{Vector{Vector{Edge}}},
    V2::Vector{Vector{Vertex}}, E2::Vector{Vector{Vector{Edge}}})

    isoV = isisomorphic(V1, V2)
    if !isoV
        return false
    end
    return isisomorphic(E1, E2)
end

function isisomorphic(V1::Vector{Vector{Vertex}}, V2::Vector{Vector{Vertex}},
    E1::Vector{Vector{Vector{Edge}}}, E2::Vector{Vector{Vector{Edge}}})

    return isisomorphic(V1, E1, V2, E2)
end

# provide "sorted" option
function isequal(V1::Vector{Vector{Vertex}}, V2::Vector{Vector{Vertex}}, verbose=false)
    isiso = isisomorphic(V1, V2)
    verbose && println("V1 and V2 are isomorphic.")

    if !isiso
        verbose && println("V1 and V2 are not isomorphic.")
        return false
    end

    for (i, V1i) in enumerate(V1)
        verbose && println("Checking V1[", i, "]")

        for (j, v1) in enumerate(V1i)
            foundv = false
            for (k, v2) in enumerate(V2[i])
                if v1.label == v2.label
                    foundv = true
                    verbose && println("V1[", i, "][", j, "] = V2[", i, "][", k, "]")
                    break
                end
            end

            if !foundv
                verbose && println("Vertex V1[", i, "][", j, "] not found in V2[", i, "]")
                return false
            end
        end
    end

    return true
end

function ==(V1::Vector{Vector{Vertex}}, V2::Vector{Vector{Vertex}})
    return isequal(V1, V2, false)
end

function isequal(V1::Vector{Vector{Vertex}}, E1::Vector{Vector{Vector{Edge}}},
    V2::Vector{Vector{Vertex}}, E2::Vector{Vector{Vector{Edge}}},
    verbose=false, rhssorted=false)

    isiso = isisomorphic(V1, E1, V2, E2)
    verbose && println("V1 and V2 and E1 and E2 are isomorphic, respectively.")
    if !isiso
        verbose && println("Not  isomorphic.")
        return false
    end

    for (i, V1i) in enumerate(V1)
        lenright = length(V1i)
        verbose && println("Checking V1[", i, "]")

        for (j, v1) in enumerate(V1i)
            foundv = false
            vloc = 0
            if rhssorted
                binsearchleft = 1
                binsearchright = lenright
                while binsearchleft <= binsearchright
                    mid = fld((binsearchleft + binsearchright), 2)
                    if V2[i][mid].label < v1.label
                        binsearchleft = mid + 1
                    elseif V2[i][mid].label > v1.label
                        binsearchright = mid - 1
                    else
                        foundv = true
                        vloc = mid
                        verbose && println("V1[", i, "][", j, "] = V2[", i, "][", mid, "]")
                        break
                    end
                end
            else
                for (k, v2) in enumerate(V2[i])
                    if v1.label == v2.label
                        foundv = true
                        vloc = k
                        verbose && println("V1[", i, "][", j, "] = V2[", i, "][", k, "]")
                        break
                    end
                end
            end

            if !foundv
                verbose && println("Vertex V1[", i, "][", j, "] not found in V2[", i, "]")
                return false
            else
                if i != 1
                    verbose && println("Checking for equality in corresponding edges")
                    for (k, e1) in enumerate(E1[i - 1][j])
                        founde = false
                        for (l, e2) in enumerate(E2[i - 1][vloc])
                            if e1.label == e2.label && e1.weight == e2.weight && V1[i - 1][e1.outvertex].label == V2[i - 1][e2.outvertex].label
                                founde = true
                                verbose && println("E1[", i - 1, "][", j, "][", k, "] = E2[", i - 1, "][", vloc, "][", l, "]")
                                break
                            end
                        end

                        if !founde
                            verbose && println("Edge E1[", i - 1, "][", j, "][", k, "] not found in E2[", i - 1, "][", vloc, "]")
                            return false
                        end
                    end
                end
            end
        end
    end
    return true
end

function isequal(V1::Vector{Vector{Vertex}}, V2::Vector{Vector{Vertex}},
    E1::Vector{Vector{Vector{Edge}}}, E2::Vector{Vector{Vector{Edge}}},
    verbose=false, rhssorted=false)

    return isequal(V1, E1, V2, E2, verbose, rhssorted)
end

function _sortbyleftindex(A::fq_nmod_mat)
    numrows = size(A, 1)
    numcols = size(A, 2)
    arr = []
    for r = 1:numrows
        for c in 1:numcols
            if !iszero(A[r, c])
                push!(arr, [c, A[r, :]])
                break
            end
        end
        push!(arr, [numcols, A[r, :]])
    end
    sort!(arr, by=x->x[1])
    # vcat is faster and cheaper than the following
    # return A[[arr2[i][2] for i in 1:numrows], :]
    return vcat([arr[i][2] for i in 1:numrows]...)
end

function _sortbyrightindex(A::fq_nmod_mat)
    numrows = size(A, 1)
    numcols = size(A, 2)
    arr = []
    for r = 1:numrows
        for c in numcols:-1:1
            if !iszero(A[r, c])
                push!(arr, [c, A[r, :]])
                break
            end
        end
        push!(arr, [1, A[r, :]])
    end
    sort!(arr, by=x->x[1]) #, rev=true
    return vcat([arr[i][2] for i in 1:numrows]...)
end

function _leftrightindices(A::fq_nmod_mat)
    # iseven(size(A, 2)) || error("Vectors in leftrightindices must have even length.")
    # n = div(size(A, 2), 2)

#take quadratic only here

    left = Vector{Int64}()
    right = Vector{Int64}()
    for i in 1:size(A, 1)
        for j in 1:size(A, 2)
            if !iszero(A[i, j])
                push!(left, j)
                break
            end
        end
    end
    for i in 1:size(A, 1)
        for j in size(A, 2):-1:1
            if !iszero(A[i, j])
                push!(right, j)
                break
            end
        end
    end
    return left, right
end

function _findactive(A::fq_nmod_mat, edges::Bool=false)
    # need to make sure quadratic extension
    n = size(A, 2)
    k = size(A, 1)
    leftright = [[0, 0] for i in 1:k]
    for i in 1:k
        for j in 1:n
            if !iszero(A[i, j])
                leftright[i][1] = j
                break
            end
        end
    end
    for i in 1:k
        for j in n:-1:1
            if !iszero(A[i, j])
                leftright[i][2] = j
                break
            end
        end
    end

    active = Vector{Vector{Int64}}()
    if edges
        for i in 1:n
            arri = Vector{Int64}()
            for j in 1:k
                if leftright[j][1] <= i <= leftright[j][2]
                    push!(arri, j)
                end
            end
            if !isempty(arri)
                push!(active, arri)
            end
        end
    else
        for i in 1:n
            arri = Vector{Int64}()
            for j in 1:k
                if leftright[j][1] <= i < leftright[j][2]
                    push!(arri, j)
                end
            end
            if !isempty(arri)
                push!(active, arri)
            end
        end
    end
    return active
end

# finds all elements of B which have zero symplectic inner product with all of A
function _symplectickernel(A::fq_nmod_mat, B::fq_nmod_mat, returnker::Bool=false)
    size(A, 2) == size(B, 2) || error("Length mismatch in computedimkernel.")
    iseven(size(A, 2)) || error("Vectors in symplectickernel must have even length.")

    if returnker
        ker = []
        for rb in 1:size(B, 1)
            for ra in 1:size(A, 1)
                @views if !iszero(symplecticinnerproduct(A[ra, :], B[rb, :]))
                    push!(ker, B[rb, :])
                    break
                end
            end
        end
        return length(ker), vcat(ker...)
    else
        AEuc = hcat(A[:, div(size(A, 2), 2) + 1:end], -A[:, 1:div(size(A, 2), 2)])
        prod = AEuc * transpose(B)
        iszero(prod) && return 0
        count = 0
        for i in 1:size(prod, 2)
            if !iszero(prod[:, i])
                count += 1
            end
        end
        return count
    end
end

function _pastfuture(A::fq_nmod_mat)
    past = zeros(Int64, size(A, 2) + 1)
    future = zeros(Int64, size(A, 2) + 1)
    left, right = _leftrightindices(A)
    past[1] = 0
    future[1] = size(A, 1)
    for i in 2:size(A, 2) + 1
        past[i] = length(right[right .<= i - 1])
        future[i] = length(left[left .> i - 1])
    end
    return past, future
end

function loadbalancedecode(profile::Vector{Int64})
    leftsum = 0
    leftloc = 1
    rightsum = 0
    rightloc = length(profile)
    while rightloc - leftloc > 2
        if leftsum <= rightsum
            leftsum += profile[leftloc]
            leftloc += 1
        else
            rightsum += profile[rightloc]
            rightloc -= 1
        end
    end
    # println(leftsum, ", ", rightsum)
    return leftloc, rightloc
end

function loadbalancedcode(profiles::Vector{Vector{Int64}})
    length(profiles) == 4 || error("Expected a length 4 profile vector. Pass in all or just the edges.")
    return loadbalancedcode(profiles[2])
end

#############################
        # Classical
#############################

# currently untested for characteristic > 2
function trellisorientedformC(A::fq_nmod_mat)
    E = base_ring(A)
    isone(degree(E)) || error("The trellis oriented form has not yet been implemented for higher degree extensions.")
    A = _sortbyleftindex(A)
    for c in 1:size(A, 2)
        left, right = _leftrightindices(A)
        rows = findall(x->x==c, left)
        if length(rows) == 1
            @views A[rows[1], :] *= inv(A[rows[1], c])
        elseif length(rows) > 1
            edges = []
            for row in rows
                push!(edges, (row, A[row, c]))
            end

            pivot = false
            if !isempty(edges)
                pivot = true
                row, X = edges[1]
                @views A[row, :] *= inv(X)
                for i in 2:length(edges)
                    @views A[edges[i][1], :] -= edges[i][2] * A[row, :]
                end
            end
        end
    end
    A = _sortbyleftindex(A)
    left, right = _leftrightindices(A)

    for c in size(A, 2):-1:1
        rows = findall(x->x==c, right)
        if length(rows) == 1
            @views A[rows[1], :] *= inv(A[rows[1], c])
        elseif length(rows) > 1
            edges = []
            for row in rows
                push!(edges, (row, A[row, c]))
            end

            pivot = false
            if !isempty(edges)
                pivot = true
                row, X = edges[end]
                @views A[row, :] *= inv(X)
                # no problems here if this is only length 1
                for i in length(edges) - 1:-1:1
                    @views A[edges[i][1], :] -= edges[i][2] * A[row, :]
                end
            end
        end
        left, right = _leftrightindices(A)
    end
    return A
end

function _trellisprofilesC(wrtV::fq_nmod_mat, wrtE::fq_nmod_mat, boundaries::Union{Vector{Int64}, Missing})
    if ismissing(boundaries)
        bds = [1:size(wrtV, 2) + 1...]
    else
        bds = deepcopy(boundaries)
        if 0 ∈ bds
            bds .+= 1
        end
    end

    n = length(bds) - 1
    stateprofile = zeros(Int64, n + 1)
    branchprofile = zeros(Int64, n)
    indegrees = zeros(Int64, n)
    outdegrees = zeros(Int64, n)
    past, future = _pastfuture(wrtV)
    past = past[bds]
    future = future[bds]

    p = Int64(characteristic(base_ring(wrtV)))
    for i in 1:n + 1
    #     # same formula works wrtE but still needs a separate loop for the n + 1
    #     # if only doing wrtE need to think about role of dimker
        stateprofile[i] = p^(size(wrtV, 1) - past[i] - future[i]) # - dimker)
    end

    left, right = _leftrightindices(wrtE)
    past, future = _pastfuture(wrtE)
    past = past[bds]
    future = future[bds]
    for i in 1:n
        dimparallel = 0
        branchprofile[i] = p^(size(wrtE, 1) - past[i] - future[i + 1] - dimparallel) # - dimker
        indegrees[i] = div(branchprofile[i], stateprofile[i + 1])
        outdegrees[i] = div(branchprofile[i], stateprofile[i])
    end
    # stateprofile[n + 1] = p^(size(wrtE, 1) - past[n + 1] - future[n + 1] - dimker)
    return [stateprofile, branchprofile, indegrees, outdegrees]
end

# TODO: remove dictionaries, iterate once to find left, once for right
# keep first bipartite structure and shift it as a coset to find the next ones - fix for sectionalization
function _syndrometrellisC(profiles::Vector{Vector{T}}, boundaries::Union{Vector{Int64}, Missing},
    wrtV::fq_nmod_mat, wrtE::fq_nmod_mat, verbose=false) where T <: Integer

    K = base_ring(wrtE)
    base_ring(wrtV) == K || error("Vertices and edges must have the same base ring.")
    if ismissing(boundaries)
        bds = [0:size(wrtV, 2)...]
    else
        bds = deepcopy(boundaries)
    end

    p = Int64(characteristic(K))
    n = length(profiles[1]) - 1
    V = Vector{Vertex}[Vertex[] for i = 1:n + 1]
    Threads.@threads for i = 1:n + 1
        V[i] = [Vertex(999, 0, 0, 0.0, 0, [[0, 0, 0]]) for j = 1:profiles[1][i]]
    end
    V[1] = [Vertex(0, 0, 0, 0.0, 0, [[1, 0, 0]])]
    V[end] = [Vertex(0, 0, 0, 0.0, 0, [[0, 0, 0]])]
    verbose && println("Vertex preallocation completed.")

    E = Vector{Vector{Edge}}[[Edge[]] for i = 1:n]
    Threads.@threads for i in 1:n
        # the j-th element of Ei is going to be all of the edges going into Vi[j]
        E[i] = [[Edge(K(0), 0.0, 0, 1) for j = 1:profiles[3][i]] for k = 1:profiles[1][i + 1]]
    end
    verbose && println("Edge preallocation completed.")

    biz = BigInt(0)
    bio = BigInt(1)
    synlen = size(wrtV, 1)
    active = _findactive(wrtV)
    active = active[bds[2:end - 1]]
    Threads.@threads for i = 2:n
        Visize = profiles[1][i]
        for num in 0:Visize - 1
            bin = reverse(digits(num, base=2, pad=length(active[i - 1])))
            templabel = zeros(Int64, synlen)
            loc = 1
            for j in active[i - 1]
                templabel[j] = bin[loc]
                loc += 1
            end

            curlabel = biz
            for (shift, val) in enumerate(reverse(templabel, dims=1))
                if val == 1
                    curlabel += bio << (shift - 1)
                end
            end
            V[i][num + 1].label = curlabel
        end
    end
    verbose && println("Vertex construction completed.")

    # there has to be a one-liner for the below
    left, right = _leftrightindices(wrtE)
    active = _findactive(wrtE, true)
    if ismissing(boundaries)
        activetemp = active
    else
        activetemp = Vector{Vector{Int64}}()
        for i in 1:length(bds) - 1
            temp = Vector{Int64}()
            for j in bds[i] + 1:bds[i + 1]
                if !(bds[i] <= left[j] && right[j] <= bds[i + 1])
                    append!(temp, active[j])
                end
            end
            push!(activetemp, sort!(unique!(temp)))
        end
    end

    H = FpmattoJulia(wrtV)
    G = FpmattoJulia(wrtE)
    Threads.@threads for i = n:-1:1
        verbose && println("Starting E[$i]")
        seclen = bds[i + 1] - bds[i]
        validedges = Vector{fq_nmod_mat}()
        edgecontrib = Dict{fq_nmod_mat, Vector{Int64}}()
        contribedge = Dict{Vector{Int64}, fq_nmod_mat}()

        for a in activetemp[i]
            temp = wrtE[a, bds[i] + 1:bds[i + 1]]
            if !iszero(temp)
                push!(validedges, temp)
            end
        end
        unique!(validedges)

        for iter in Nemo.AbstractAlgebra.ProductIterator(collect(0:Int64(characteristic(K)) - 1), length(validedges))
            e = K(iter[1]) * validedges[1]
            for r in 2:length(validedges)
                if !iszero(iter[r])
                    e += K(iter[r]) * validedges[r]
                end
            end

            P = zeros(Int64, n)
            for (j, k) in enumerate(e)
                P[bds[i] + j] = coeff(k, 0)
            end
            syn = H * P .% p
            edgecontrib[e] = syn
            contribedge[syn] = e
        end
        verbose && println("Edges dictionaries completed for E[$i].")

        Vleft = V[i]
        Vright = V[i + 1]
        lenleft = length(Vleft)
        lenright = length(Vright)
        Vrightlocs = trues(lenright)
        startingrightindex = 1
        # keep below here instead of preallocating above or else the answer comes out wrong
        blank = K(0)

        while startingrightindex <= lenright
            startingrightv = Vright[startingrightindex].label
            leftvertices = Vector{Tuple{Int64, BigInt, Vector{Int64}}}()
            rightvertices = Vector{Tuple{Int64, BigInt, Vector{Int64}}}()
            sizehint!(leftvertices, profiles[4][i])
            sizehint!(rightvertices, profiles[3][i])

            startingrightvsyn = reverse(digits(startingrightv, base=2, pad=synlen))
            push!(rightvertices, (startingrightindex, startingrightv, startingrightvsyn))
            connectingstarts = blank
            startingleftv = biz

            # start with a fixed right vertex and find all left vertices
            for lab in keys(edgecontrib)
                temp = (startingrightvsyn .- edgecontrib[lab])
                for t in 1:length(temp)
                    if temp[t] < 0
                        temp[t] = p + temp[t]
                    end
                end
                temp = temp .% p

                leftlabel = biz
                for (shift, val) in enumerate(reverse(temp, dims=1))
                    if val == 1
                        leftlabel += bio << (shift - 1)
                    end
                end

                binsearchleft = 1
                binsearchright = lenleft
                while binsearchleft <= binsearchright
                    mid = fld((binsearchleft + binsearchright), 2)
                    if Vleft[mid].label < leftlabel
                        binsearchleft = mid + 1
                    elseif Vleft[mid].label > leftlabel
                        binsearchright = mid - 1
                    else
                        push!(leftvertices, (mid, leftlabel, temp))
                        if connectingstarts == blank
                            connectingstarts = lab
                            startingleftv = temp
                        end
                        break
                    end
                end

                if length(leftvertices) == profiles[3][i]
                    break
                end
            end

            # start with first left vertex and find all right vertices
            if length(rightvertices) != profiles[4][i]
                for lab in keys(edgecontrib)
                    if lab != connectingstarts
                        temp = (startingleftv .+ edgecontrib[lab]) .% p
                        rightlabel = biz
                        for (shift, val) in enumerate(reverse(temp, dims=1))
                            if val == 1
                                rightlabel += bio << (shift - 1)
                            end
                        end

                        binsearchleft = 1
                        binsearchright = lenright
                        while binsearchleft <= binsearchright
                            mid = fld((binsearchleft + binsearchright), 2)
                            if Vright[mid].label < rightlabel
                                binsearchleft = mid + 1
                            elseif Vright[mid].label > rightlabel
                                binsearchright = mid - 1
                            else
                                push!(rightvertices, (mid, rightlabel, temp))
                                break
                            end
                        end
                    end

                    if length(rightvertices) == profiles[4][i]
                        break
                    end
                end
            end

            # can probably skip this recalculation of temp by immediately storing
            # instead of building right and left vertex lists
            # should now have all vertices
            for (rightindex, rightlabel, rightsyn) in rightvertices
                count = 1
                for (leftindex, leftlabel, leftsyn) in leftvertices
                    temp = rightsyn .- leftsyn
                    for t in 1:length(temp)
                        if temp[t] < 0
                            temp[t] = p + temp[t]
                        end
                    end
                    temp = temp .% p

                    E[i][rightindex][count].label = contribedge[temp]
                    E[i][rightindex][count].outvertex = leftindex
                    count += 1
                end
                Vrightlocs[rightindex] = false
            end

            # should have ==
            while startingrightindex <= lenright
                startingrightindex += 1
                if startingrightindex <= lenright && !Vrightlocs[startingrightindex]
                    startingrightindex += 1
                else
                    break
                end
            end

        end
        verbose && println("E[$i] complete")
    end
    return Trellis(V, E)
end

# should probably return missing to unify if statements in the function below
function trellisprofiles(C::AbstractLinearCode, sect::Bool=false)
    GTOF = trellisorientedformC(generatormatrix(C))
    HTOF = trellisorientedformC(paritycheckmatrix(C))
    if sect
        opt, _ = optimalsectionalizationC(HTOF, GTOF)
        return _trellisprofilesC(HTOF, GTOF, opt), opt
    end
    return _trellisprofilesC(HTOF, GTOF, missing)
end

# recopy above so I don't call it and redo the TOF multiple times
function syndrometrellis(C::AbstractLinearCode, sect::Bool=false)
    GTOF = trellisorientedformC(generatormatrix(C))
    HTOF = trellisorientedformC(paritycheckmatrix(C))
    if sect
        opt, _ = optimalsectionalizationC(HTOF, GTOF)
        profiles = _trellisprofilesC(HTOF, GTOF, opt)
        return _syndrometrellisC(profiles, opt, HTOF, GTOF, false)
    else
        profiles = _trellisprofilesC(HTOF, GTOF, missing)
        return _syndrometrellisC(profiles, missing, HTOF, GTOF, false)
        # return _syndrometrellisC(profiles, missing, HTOF, GTOF, [1 for _ in 1:length(C)], ' ', false)
    end
end

#############################
         # Quantum
#############################

# seeks to eliminate edges of the form a + bω in a quadratic extension
# can't check that it's quadratic directly so note this will fail if a higher degree
# should work for degree 1 extensions given the way AbstractAlgebra stores coefficints?
# currently untested for characteristic > 2
function trellisorientedformQ(A::fq_nmod_mat)
    # check for quadratic field extension
    E = base_ring(A)
    A = _sortbyleftindex(A)
    for c in 1:size(A, 2)
        left, right = _leftrightindices(A)
        rows = findall(x->x==c, left)
        if length(rows) == 1
            if !iszero(coeff(A[rows[1], c], 0)) && !iszero(coeff(A[rows[1], c], 1))
                @views A[rows[1], :] *= inv(E(coeff(A[rows[1], c], 0)))
            elseif !iszero(coeff(A[rows[1], c], 0))
                @views A[rows[1], :] *= inv(E(coeff(A[rows[1], c], 0)))
            else
                @views A[rows[1], :] *= inv(E(coeff(A[rows[1], c], 1)))
            end
        elseif length(rows) > 1
            Xedges = []
            Zedges = []
            mixededges = []
            for row in rows
                if !iszero(coeff(A[row, c], 0)) && !iszero(coeff(A[row, c], 1))
                    push!(mixededges, (row, A[row, c]))
                elseif !iszero(coeff(A[row, c], 0)) && iszero(coeff(A[row, c], 1))
                    push!(Xedges, (row, A[row, c]))
                elseif iszero(coeff(A[row, c], 0)) && !iszero(coeff(A[row, c], 1))
                    push!(Zedges, (row, A[row, c]))
                end
            end

            if length(Xedges) <= 1 && length(Zedges) <= 1 && length(mixededges) <= 1 && (length(Xedges) + length(Zedges) + length(mixededges)) <= 2
                continue
            else
                Xpivot = false
                Zpivot = false
                if !isempty(Xedges)
                    Xpivot = true
                    row, X = Xedges[1]
                    @views A[row, :] *= inv(E(coeff(X, 0)))
                    # no problems here if this is only length 1
                    for i in 2:length(Xedges)
                        @views A[Xedges[i][1], :] -= E(coeff(Xedges[i][2], 0)) * A[row, :]
                    end
                end
                if !isempty(Zedges)
                    Zpivot = true
                    row, Z = Zedges[1]
                    @views A[row, :] *= inv(E(coeff(Z, 1)))
                    # no problems here if this is only length 1
                    for i in 2:length(Zedges)
                        @views A[Zedges[i][1], :] -= E(coeff(Zedges[i][2], 1)) * A[row, :]
                    end
                end
                if !isempty(mixededges)
                    if Xpivot && Zpivot
                        for i in 1:length(mixededges)
                            @views A[mixededges[i][1], :] -= E(coeff(mixededges[i][2], 0)) * A[Xedges[1][1], :] + E(coeff(mixededges[i][2], 1)) * A[Zedges[1][1], :]
                        end
                    elseif Xpivot
                        @views A[mixededges[1][1], :] -= E(coeff(mixededges[1][2], 0)) * A[Xedges[1][1], :]
                        # no problems here if this is only length 1
                        for i in 2:length(mixededges)
                            @views A[mixededges[i][1], :] -= E(coeff(mixededges[i][2], 0)) * A[Xedges[1][1], :] + E(coeff(mixededges[i][2], 1)) * A[mixededges[1][1], :]
                        end
                    elseif Zpivot
                        @views A[mixededges[1][1], :] -= E(coeff(mixededges[1][2], 1)) * A[Zedges[1][1], :]
                        # no problems here if this is only length 1
                        for i in 2:length(mixededges)
                            @views A[mixededges[i][1], :] -= E(coeff(mixededges[i][2], 0)) * A[mixededges[1][1], :] + E(coeff(mixededges[i][2], 1)) * A[Zedges[1][1], :]
                        end
                    else
                        @views A[mixededges[1][1], :] *= inv(E(coeff(mixededges[1][2], 0)))
                        if length(mixededges) > 1
                            @views A[mixededges[2][1], :] -= E(coeff(mixededges[2][2], 0)) * A[mixededges[1][1], :]
                            @views A[mixededges[2][1], :] *= inv(E(coeff(A[mixededges[2][1], c], 1)))
                            if length(mixededges) > 2
                                @views A[mixededges[3][1], :] -= E(coeff(mixededges[3][2], 1)) * A[mixededges[2][1], :]
                                @views A[mixededges[3][1], :] *= inv(E(coeff(A[mixededges[3][1], c], 0)))
                                # no problems here if this is only length 3
                                for i in 3:length(mixededges)
                                    @views A[mixededges[i][1], :] -= E(coeff(mixededges[i][2], 0)) * A[mixededges[3][1], :] + E(coeff(mixededges[i][2], 1)) * A[mixededges[2][1], :]
                                end
                            end
                        end
                    end
                end
            end
        end
    end
    A = _sortbyleftindex(A)
    left, right = _leftrightindices(A)

    for c in size(A, 2):-1:1
        rows = findall(x->x==c, right)
        if length(rows) == 1
            if !iszero(coeff(A[rows[1], c], 0)) && !iszero(coeff(A[rows[1], c], 1))
                @views A[rows[1], :] *= inv(E(coeff(A[rows[1], c], 0)))
            elseif !iszero(coeff(A[rows[1], c], 0))
                @views A[rows[1], :] *= inv(E(coeff(A[rows[1], c], 0)))
            else
                @views A[rows[1], :] *= inv(E(coeff(A[rows[1], c], 1)))
            end
        elseif length(rows) > 1
            Xedges = []
            Zedges = []
            mixededges = []
            for row in rows
                if !iszero(coeff(A[row, c], 0)) && !iszero(coeff(A[row, c], 1))
                    push!(mixededges, (row, A[row, c]))
                elseif !iszero(coeff(A[row, c], 0)) && iszero(coeff(A[row, c], 1))
                    push!(Xedges, (row, A[row, c]))
                elseif iszero(coeff(A[row, c], 0)) && !iszero(coeff(A[row, c], 1))
                    push!(Zedges, (row, A[row, c]))
                end
            end

            if length(Xedges) <= 1 && length(Zedges) <= 1 && length(mixededges) <= 1 && (length(Xedges) + length(Zedges) + length(mixededges)) <= 2
                continue
            else
                Xpivot = false
                Zpivot = false
                if !isempty(Xedges)
                    Xpivot = true
                    row, X = Xedges[end]
                    @views A[row, :] *= inv(E(coeff(X, 0)))
                    # no problems here if this is only length 1
                    for i in length(Xedges) - 1:-1:1
                        @views A[Xedges[i][1], :] -= E(coeff(Xedges[i][2], 0)) * A[row, :]
                    end
                end
                if !isempty(Zedges)
                    Zpivot = true
                    row, Z = Zedges[end]
                    @views A[row, :] *= inv(E(coeff(Z, 1)))
                    # no problems here if this is only length 1
                    for i in length(Zedges) - 1:-1:1
                        @views A[Zedges[i][1], :] -= E(coeff(Zedges[i][2], 1)) * A[row, :]
                    end
                end
                if !isempty(mixededges)
                    if Xpivot && Zpivot
                        for i in 1:length(mixededges)
                            @views A[mixededges[i][1], :] -= E(coeff(mixededges[i][2], 0)) * A[Xedges[end][1], :] + E(coeff(mixededges[i][2], 1)) * A[Zedges[end][1], :]
                        end
                    elseif Xpivot
                        @views A[mixededges[end][1], :] -= E(coeff(mixededges[1][2], 0)) * A[Xedges[end][1], :]
                        # no problems here if this is only length 1
                        for i in length(mixededges) - 1:-1:1
                            @views A[mixededges[i][1], :] -= E(coeff(mixededges[i][2], 0)) * A[Xedges[end][1], :] + E(coeff(mixededges[i][2], 1)) * A[mixededges[end][1], :]
                        end
                    elseif Zpivot
                        @views A[mixededges[end][1], :] -= E(coeff(mixededges[1][2], 1)) * A[Zedges[end][1], :]
                        # no problems here if this is only length 1
                        for i in length(mixededges) - 1:-1:1
                            @views A[mixededges[i][1], :] -= E(coeff(mixededges[i][2], 0)) * A[mixededges[end][1], :] + E(coeff(mixededges[i][2], 1)) * A[Zedges[end][1], :]
                        end
                    else
                        @views A[mixededges[end][1], :] *= inv(E(coeff(mixededges[end][2], 0)))
                        if length(mixededges) > 1
                            @views A[mixededges[end - 1][1], :] -= E(coeff(mixededges[end - 1][2], 0)) * A[mixededges[end][1], :]
                            @views A[mixededges[end - 1][1], :] *= inv(E(coeff(A[mixededges[end - 1][1], c], 1)))
                            if length(mixededges) > 2
                                @views A[mixededges[end - 2][1], :] -= E(coeff(mixededges[end - 2][2], 1)) * A[mixededges[end - 1][1], :]
                                @views A[mixededges[end - 2][1], :] *= inv(E(coeff(A[mixededges[end - 2][1], c], 0)))
                                # no problems here if this is only length 3
                                for i in length(mixededges) - 1:-1:1
                                    @views A[mixededges[i][1], :] -= E(coeff(mixededges[i][2], 0)) * A[mixededges[end - 2][1], :] + E(coeff(mixededges[i][2], 1)) * A[mixededges[end - 1][1], :]
                                end
                            end
                        end
                    end
                end
            end
        end
        left, right = _leftrightindices(A)
    end
    return A
end

# only valid for quantum codes
function optimalsectionalizationQ(wrtV::fq_nmod_mat, wrtE::fq_nmod_mat)
    K = base_ring(wrtE)
    base_ring(wrtV) == K || error("Vertices and edges must have the same base ring.")

    n = size(wrtV, 2)
    p = Int64(characteristic(K))
    V = [Vertex(i, 0, 0, 0.0, 0, [[0, 0, 0, 0]]) for i in 1:n + 1]
    E = [[Edge(K(0), 0.0, 0, 1) for i in 1:j] for j in n:-1:1]

    symV = quadratictosymplectic(wrtV)
    symE = quadratictosymplectic(wrtE)
    dimker = _symplectickernel(symV, symE)
    past, future = _pastfuture(wrtE)
    for i in 1:n
        for j in i:n
            # arbitrary size cutoff
            if size(wrtE, 1) - past[i] - future[j + 1] - dimker > 50
                E[i][j - i + 1].weight = Inf
            else
                E[i][j - i + 1].weight = p^(size(wrtE, 1) - past[i] - future[j + 1] - dimker)
            end
        end
    end

    for i in 1:n
        leftb = 1
        rightb = i
        arr = [E[leftb][rightb].weight + V[leftb].value]

        while leftb < i
            leftb += 1
            rightb -= 1
            append!(arr, E[leftb][rightb].weight + V[leftb].value)
        end

        m, arg = findmin(arr)
        V[i + 1].prev = arg
        V[i + 1].value = m
    end

    sectboundaries = [n]
    next = V[end].prev
    # val = V[end].value
    while next > 0
        append!(sectboundaries, next - 1)
        # val += V[next].value
        next = V[next].prev
    end
    return reverse(sectboundaries), Int(V[end].value) #Int(val)
end

function _trellisprofilesQ(wrtV::fq_nmod_mat, wrtE::fq_nmod_mat, boundaries::Union{Vector{Int64}, Missing})
    if ismissing(boundaries)
        bds = [1:size(wrtV, 2) + 1...]
    else
        bds = deepcopy(boundaries)
        if 0 ∈ bds
            bds .+= 1
        end
    end

    symV = quadratictosymplectic(wrtV)
    symE = quadratictosymplectic(wrtE)
    dimker = _symplectickernel(symV, symE)

    n = length(bds) - 1
    stateprofile = zeros(Int64, n + 1)
    branchprofile = zeros(Int64, n)
    indegrees = zeros(Int64, n)
    outdegrees = zeros(Int64, n)
    past, future = _pastfuture(wrtV)
    past = past[bds]
    future = future[bds]

    p = Int64(characteristic(base_ring(wrtV)))
    for i in 1:n + 1
    #     # same formula works wrtE but still needs a separate loop for the n + 1
    #     # if only doing wrtE need to think about role of dimker
        stateprofile[i] = p^(size(wrtV, 1) - past[i] - future[i] - dimker)
    end

    left, right = _leftrightindices(wrtE)
    past, future = _pastfuture(wrtE)
    past = past[bds]
    future = future[bds]
    for i in 1:n
        dimparallel = 0
        # for k in 1:length(left)
        #     # maybw subtract 1 from all of these?
        #     if bds[i] <= left[k] && right[k] <= bds[i + 1]
        #         # println("Going from V[$(boundaries[i])] to V[$(boundaries[i + 1])] and claiming stabilizer $k has left index $(left[k]) and right index $(right[k]).")
        #         dimparallel += 1
        #     end
        # end
        # stateprofile[i] = p^(size(wrtE, 1) - past[i] - future[i] - dimker)
        branchprofile[i] = p^(size(wrtE, 1) - past[i] - future[i + 1] - dimker - dimparallel)
        indegrees[i] = div(branchprofile[i], stateprofile[i + 1])
        outdegrees[i] = div(branchprofile[i], stateprofile[i])
    end
    # stateprofile[n + 1] = p^(size(wrtE, 1) - past[n + 1] - future[n + 1] - dimker)
    return [stateprofile, branchprofile, indegrees, outdegrees]
end

# TODO: remove dictionaries, iterate once to find left, once for right
# keep first bipartite structure and shift it as a coset to find the next ones - fix for sectionalization
function _syndrometrellisQ(profiles::Vector{Vector{T}}, boundaries::Union{Vector{Int64}, Missing},
    wrtV::fq_nmod_mat, wrtE::fq_nmod_mat, charvec::Vector{Int64}, Pauli::Char=' ',
    verbose=false) where T <: Integer

    Pauli ∈ [' ', 'X', 'Z'] || error("Pauli parameter needs to be ' ', 'X', or 'Z'; received $Pauli.")
    K = base_ring(wrtE)
    base_ring(wrtV) == K || error("Vertices and edges must have the same base ring.")
    if ismissing(boundaries)
        bds = [0:size(wrtV, 2)...]
    else
        bds = deepcopy(boundaries)
    end

    ω = gen(K)
    p = Int64(characteristic(K))
    n = length(profiles[1]) - 1
    symsize = size(wrtV, 2)
    V = Vector{Vertex}[Vertex[] for i = 1:n + 1]
    Threads.@threads for i = 1:n + 1
        V[i] = [Vertex(999, 0, 0, 0.0, 0, [[0, 0, 0, 0]]) for j = 1:profiles[1][i]]
    end
    V[1] = [Vertex(0, 0, 0, 0.0, 0, [[1, 0, 0, 0]])]
    V[end] = [Vertex(0, 0, 0, 0.0, 0, [[0, 0, 0, 0]])]
    verbose && println("Vertex preallocation completed.")

    E = Vector{Vector{Edge}}[[Edge[]] for i = 1:n]
    Threads.@threads for i in 1:n
        # the j-th element of Ei is going to be all of the edges going into Vi[j]
        E[i] = [[Edge(K(0), 0.0, 0, 1) for j = 1:profiles[3][i]] for k = 1:profiles[1][i + 1]]
    end
    verbose && println("Edge preallocation completed.")

    biz = BigInt(0)
    bio = BigInt(1)
    synlen = size(wrtV, 1)
    active = _findactive(wrtV)
    active = active[bds[2:end - 1]]
    Threads.@threads for i = 2:n
        Visize = profiles[1][i]
        for num in 0:Visize - 1
            bin = reverse(digits(num, base=2, pad=length(active[i - 1])))
            templabel = zeros(Int64, synlen)
            loc = 1
            for j in active[i - 1]
                templabel[j] = bin[loc]
                loc += 1
            end

            curlabel = biz
            for (shift, val) in enumerate(reverse(templabel, dims=1))
                if val == 1
                    curlabel += bio << (shift - 1)
                end
            end
            V[i][num + 1].label = curlabel
        end
    end
    verbose && println("Vertex construction completed.")

    # there has to be a one-liner for the below
    left, right = _leftrightindices(wrtE)
    active = _findactive(wrtE, true)
    if ismissing(boundaries)
        activetemp = active
    else
        activetemp = Vector{Vector{Int64}}()
        for i in 1:length(bds) - 1
            temp = Vector{Int64}()
            for j in bds[i] + 1:bds[i + 1]
                if !(bds[i] <= left[j] && right[j] <= bds[i + 1])
                    append!(temp, active[j])
                end
            end
            push!(activetemp, sort!(unique!(temp)))
        end
    end

    symwrtV = quadratictosymplectic(wrtV)
    G = FpmattoJulia(hcat(symwrtV[:, symsize + 1:end], -symwrtV[:, 1:symsize]))
    Threads.@threads for i = n:-1:1
        verbose && println("Starting E[$i]")
        seclen = bds[i + 1] - bds[i]
        validedges = Vector{fq_nmod_mat}()
        edgecontrib = Dict{fq_nmod_mat, Vector{Int64}}()
        contribedge = Dict{Vector{Int64}, fq_nmod_mat}()

        for a in activetemp[i]
            temp = wrtE[a, bds[i] + 1:bds[i + 1]]
            if !iszero(temp)
                push!(validedges, temp)
            end
        end
        unique!(validedges)

        for iter in Nemo.AbstractAlgebra.ProductIterator(collect(0:p - 1), length(validedges))
            e = K(iter[1]) * validedges[1]
            for r in 2:length(validedges)
                if !iszero(iter[r])
                    e += K(iter[r]) * validedges[r]
                end
            end

            P = zeros(Int64, 2 * symsize)
            for (j, k) in enumerate(e)
                P[bds[i] + j] = coeff(k, 0)
                P[bds[i] + j + symsize] = coeff(k, 1)
            end
            syn = G * P .% p
            edgecontrib[e] = syn
            contribedge[syn] = e
        end
        verbose && println("Edges dictionaries completed for E[$i].")

        Vleft = V[i]
        Vright = V[i + 1]
        lenleft = length(Vleft)
        lenright = length(Vright)
        Vrightlocs = trues(lenright)
        startingrightindex = 1
        # keep below here instead of preallocating above or else the answer comes out wrong
        blank = K(0)

        while startingrightindex <= lenright
            startingrightv = Vright[startingrightindex].label
            leftvertices = Vector{Tuple{Int64, BigInt, Vector{Int64}}}()
            rightvertices = Vector{Tuple{Int64, BigInt, Vector{Int64}}}()
            sizehint!(leftvertices, profiles[4][i])
            sizehint!(rightvertices, profiles[3][i])

            startingrightvsyn = reverse(digits(startingrightv, base=2, pad=synlen))
            push!(rightvertices, (startingrightindex, startingrightv, startingrightvsyn))
            connectingstarts = blank
            startingleftv = biz

            # start with a fixed right vertex and find all left vertices
            for lab in keys(edgecontrib)
                temp = (startingrightvsyn .- edgecontrib[lab])
                for t in 1:length(temp)
                    if temp[t] < 0
                        temp[t] = p + temp[t]
                    end
                end
                temp = temp .% p

                leftlabel = biz
                for (shift, val) in enumerate(reverse(temp, dims=1))
                    if val == 1
                        leftlabel += bio << (shift - 1)
                    end
                end

                binsearchleft = 1
                binsearchright = lenleft
                while binsearchleft <= binsearchright
                    mid = fld((binsearchleft + binsearchright), 2)
                    if Vleft[mid].label < leftlabel
                        binsearchleft = mid + 1
                    elseif Vleft[mid].label > leftlabel
                        binsearchright = mid - 1
                    else
                        push!(leftvertices, (mid, leftlabel, temp))
                        if connectingstarts == blank
                            connectingstarts = lab
                            startingleftv = temp
                        end
                        break
                    end
                end

                if length(leftvertices) == profiles[3][i]
                    break
                end
            end

            # start with first left vertex and find all right vertices
            if length(rightvertices) != profiles[4][i]
                for lab in keys(edgecontrib)
                    if lab != connectingstarts
                        temp = (startingleftv .+ edgecontrib[lab]) .% p
                        rightlabel = biz
                        for (shift, val) in enumerate(reverse(temp, dims=1))
                            if val == 1
                                rightlabel += bio << (shift - 1)
                            end
                        end

                        binsearchleft = 1
                        binsearchright = lenright
                        while binsearchleft <= binsearchright
                            mid = fld((binsearchleft + binsearchright), 2)
                            if Vright[mid].label < rightlabel
                                binsearchleft = mid + 1
                            elseif Vright[mid].label > rightlabel
                                binsearchright = mid - 1
                            else
                                push!(rightvertices, (mid, rightlabel, temp))
                                break
                            end
                        end
                    end

                    if length(rightvertices) == profiles[4][i]
                        break
                    end
                end
            end

            # can probably skip this recalculation of temp by immediately storing
            # instead of building right and left vertex lists
            # should now have all vertices
            for (rightindex, rightlabel, rightsyn) in rightvertices
                count = 1
                for (leftindex, leftlabel, leftsyn) in leftvertices
                    temp = rightsyn .- leftsyn
                    for t in 1:length(temp)
                        if temp[t] < 0
                            temp[t] = p + temp[t]
                        end
                    end
                    temp = temp .% p
                    lab = contribedge[temp]
                    sign = 1 # should be K(1) when implementing as roots of unity or in \C?
                    for (j, k) in enumerate(lab)
                        if !iszero(coeff(k, 0))
                            sign *= charvec[bds[i] + j]
                        end
                        if !iszero(coeff(k, 1))
                            sign *= charvec[bds[i] + j + symsize]
                        end
                    end

                    E[i][rightindex][count].label = lab
                    E[i][rightindex][count].outvertex = leftindex
                    E[i][rightindex][count].sign = sign
                    count += 1
                end
                Vrightlocs[rightindex] = false
            end

            # should have ==
            while startingrightindex <= lenright
                startingrightindex += 1
                if startingrightindex <= lenright && !Vrightlocs[startingrightindex]
                    startingrightindex += 1
                else
                    break
                end
            end

        end
        verbose && println("E[$i] complete")
    end
    return Trellis(V, E)
end

# error models need to take CSS combinations into account
# Pauli == 'X'
# I -> I + X
# Z -> Z + Y
# Pauli == 'Z'
# I -> I + Z
# X -> X + Y
function weightQ!(T::Trellis, Ps::fq_nmod_mat, err_models::Vector{Dict{String, Float64}},
    weighttype::String="additive")

    weighttype ∈ ["additive", "multiplicative"] || error("Weight type needs to be 'additive' or 'multiplicative'.")

    V = vertices(T)
    E = edges(T)
    for i in 1:length(E)
        model = err_models[i]
        for (j, v) in enumerate(V[i + 1])
            Threads.@threads for e in E[i][j]
                if weighttype == "additive"
                    weight = 0.0
                    for (j, k) in enumerate(e.label)
                        weight += model[j]
                    end
                else
                    weight = 1.0
                    for (j, k) in enumerate(e.label)
                        weight += model[j]
                    end
                end
                e.weight = weight
            end
        end
    end
end

# error models need to take CSS combinations into account
# Pauli == 'X'
# I -> I + X
# Z -> Z + Y
# Pauli == 'Z'
# I -> I + Z
# X -> X + Y
# remove wrtV here
function shiftandweightQ!(T::Trellis, Ps::fq_nmod_mat, boundaries::Union{Vector{Int64}, Missing},
    err_models::Vector{Dict{String, Float64}}, charvec::Vector{Int64}, Pauli::Char=' ',
    weighttype::String="additive")

    Pauli ∈ [' ', 'X', 'Z'] || error("Pauli parameter needs to be ' ', 'X', or 'Z'; received $Pauli.")
    weighttype ∈ ["additive", "multiplicative"] || error("Weight type needs to be 'additive' or 'multiplicative'.")
    length(charvec) == 2 * length(err_models) || error("Lengths of character vector and error models are not consistent.")

    K = base_ring(Ps)
    V = vertices(T)
    E = edges(T)
    coden = length(err_models)
    if ismissing(boundaries)
        bds = [0:coden...]
    else
        bds = deepcopy(boundaries)
    end

    for i in 1:length(E)
        model = err_models[i]
        for (j, v) in enumerate(V[i + 1])
            Threads.@threads for e in E[i][j]
                e.label += Ps[bds[i] + 1:bds[i + 1]]
                if Pauli == 'X'
                    for k in e.label
                        k -= K(coeff(k, 0))
                    end
                elseif Pauli == 'Z'
                    for k in e.label
                        k -= K(coeff(k, 1))
                    end
                end

                sign = 1 # should be K(1) when implementing as roots of unity or in \C?
                if weighttype == "additive"
                    weight = 0.0
                else
                    weight = 1.0
                end
                for (j, k) in enumerate(e.label)
                    if !iszero(coeff(k, 0))
                        sign *= charvec[bds[i] + j]
                    end
                    if !iszero(coeff(k, 1))
                        sign *= charvec[bds[i] + j + coden]
                    end
                    if weighttype == "additive"
                        weight += model[j]
                    else
                        weight *= model[j]
                    end
                end
                e.sign = sign
                e.weight = weight
            end
        end
    end
end

# do I actually care about updating the signs here?
function shiftanddecodeQ!(T::Trellis, Ps::fq_nmod_mat, boundaries::Union{Vector{Int64}, Missing},
    err_models::Vector{Dict{fq_nmod, Float64}}, charvec::Vector{Int64}, Pauli::Char=' ',
    weighttype::String="additive")

    # Pauli ∈ [' ', 'X', 'Z'] || error("Pauli parameter needs to be ' ', 'X', or 'Z'; received $Pauli.")
    # weighttype ∈ ["additive", "multiplicative"] || error("Weight type needs to be 'additive' or 'multiplicative'.")
    # length(charvec) == 2 * length(err_models) || error("Lengths of character vector and error models are not consistent.")

    K = base_ring(Ps)
    V = vertices(T)
    E = edges(T)
    coden = length(err_models)
    if ismissing(boundaries)
        bds = [0:coden...]
    else
        bds = deepcopy(boundaries)
    end

    for i in 1:length(E)
        model = err_models[i]
        for (j, v) in enumerate(V[i + 1])
            Threads.@threads for e in E[i][j]
                e.label += Ps[1, bds[i] + 1:bds[i + 1]]
                if Pauli == 'X'
                    for k in e.label
                        k -= K(coeff(k, 0))
                    end
                elseif Pauli == 'Z'
                    for k in e.label
                        k -= K(coeff(k, 1))
                    end
                end

                # sign = 1 # should be K(1) when implementing as roots of unity or in \C?
                if weighttype == "additive"
                    weight = 0.0
                else
                    weight = 1.0
                end
                for (j, k) in enumerate(e.label)
                    # if !iszero(coeff(k, 0))
                    #     sign *= charvec[bds[i] + j]
                    # end
                    # if !iszero(coeff(k, 1))
                    #     sign *= charvec[bds[i] + j + coden]
                    # end
                    if weighttype == "additive"
                        weight += model[k]
                    else
                        weight *= model[k]
                    end
                end
                # e.sign = sign
                e.weight = weight
            end

            # ignoring random tie breaker, breaks ties here by order of entered into trellis
            w, loc = findmin([e.weight + V[i][e.outvertex].value for e in E[i][j]])
            v.value = E[i][j][loc].weight + V[i][E[i][j][loc].outvertex].value
            v.prev = E[i][j][loc].outvertex
            v.edgeloc = loc
        end
    end

    # redo here for the two-sided approach
    path = zero_matrix(base_ring(Ps), 1, coden)
    curr = coden
    prev = 1
    for i in length(E) + 1:-1:2
        elabel = E[i - 1][prev][V[i][prev].edgeloc].label
        path[1, (curr - length(elabel) + 1):curr] = elabel
        prev = V[i][prev].prev
        curr -= length(elabel)
    end
    return path
end

function shift!(T::Trellis, Ps::fq_nmod_mat, boundaries::Union{Vector{Int64}, Missing},
    err_models::Vector{Dict{fq_nmod, Float64}}, charvec::Vector{Int64}, Pauli::Char=' ',
    weighttype::String="additive")

    # Pauli ∈ [' ', 'X', 'Z'] || error("Pauli parameter needs to be ' ', 'X', or 'Z'; received $Pauli.")
    # weighttype ∈ ["additive", "multiplicative"] || error("Weight type needs to be 'additive' or 'multiplicative'.")
    # length(charvec) == 2 * length(err_models) || error("Lengths of character vector and error models are not consistent.")

    K = base_ring(Ps)
    V = vertices(T)
    E = edges(T)
    coden = length(err_models)
    if ismissing(boundaries)
        bds = [0:coden...]
    else
        bds = deepcopy(boundaries)
    end

    for i in 1:length(E)
        model = err_models[i]
        for j in 1:length(V[i + 1])
            Threads.@threads for e in E[i][j]
                e.label += Ps[1, bds[i] + 1:bds[i + 1]]
                if Pauli == 'X'
                    for k in e.label
                        k -= K(coeff(k, 0))
                    end
                elseif Pauli == 'Z'
                    for k in e.label
                        k -= K(coeff(k, 1))
                    end
                end

                # sign = 1 # should be K(1) when implementing as roots of unity or in \C?
                if weighttype == "additive"
                    weight = 0.0
                else
                    weight = 1.0
                end
                for (j, k) in enumerate(e.label)
                    # if !iszero(coeff(k, 0))
                    #     sign *= charvec[bds[i] + j]
                    # end
                    # if !iszero(coeff(k, 1))
                    #     sign *= charvec[bds[i] + j + coden]
                    # end
                    if weighttype == "additive"
                        weight += model[k]
                    else
                        weight *= model[k]
                    end
                end
                # e.sign = sign
                e.weight = weight
            end
        end
    end
end

# think of more scenarios
# could allow general trellises given partial stabilizers for use in trellis product
function trellisprofiles(Q::AbstractQuantumCode, type::String="weight", Pauli::Char=' ',
    sect::Bool=false)

    type ∈ ["weight", "decoding"] || error("Unknown type parameter in trellisprofiles.")
    (Pauli != ' ' && typeof(Q) <: CSSCode) && error("Pauli parameter is non-empty but the code is not CSS.")
    Pauli ∈ [' ', 'X', 'Z'] || error("Unknown Pauli parameter $Pauli; must be ' ', 'X', or 'Z'.")

    if type == "weight"
        if Pauli == ' '
            STOF = trellisorientedformQ(stabilizers(Q))
            nTOF = trellisorientedformQ(normalizermatrix(Q))
            if sect
                opt, _ = optimalsectionalizationQ(nTOF, STOF)
                return _trellisprofilesQ(nTOF, STOF, opt), opt
            end
            return _trellisprofilesQ(nTOF, STOF, missing)
        elseif Pauli == 'X'
            _, _, Zperp, _, _, _ = splitsymplecticstabilizers(quadratictosymplectic(normalizermatrix(Q)), ones(Int64, size(normalizermatrix(Q), 1)))
            X = hcat(Xstabilizers(Q), zero_matrix(field(Q), size(Xstabilizers(Q), 1), size(Xstabilizers(Q), 2)))
            ZperpTOF = trellisorientedformQ(symplectictoquadratic(Zperp))
            XTOF = trellisorientedformQ(symplectictoquadratic(X))
            if sect
                opt, _ = optimalsectionalizationQ(ZperpTOF, XTOF)
                return _trellisprofilesQ(ZperpTOF, XTOF, opt), opt
            end
            return _trellisprofilesQ(ZperpTOF, XTOF, missing)
        else
            Xperp, _, _, _, _, _ = splitsymplecticstabilizers(quadratictosymplectic(normalizermatrix(Q)), ones(Int64, size(normalizermatrix(Q), 1)))
            Z = hcat(zero_matrix(field(Q), size(Zstabilizers(Q), 1), size(Zstabilizers(Q), 2)), Zstabilizers(Q))
            XperpTOF = trellisorientedformQ(symplectictoquadratic(Xperp))
            ZTOF = trellisorientedformQ(symplectictoquadratic(Z))
            if sect
                opt, _ = optimalsectionalizationQ(XperpTOF, ZTOF)
                return _trellisprofilesQ(XperpTOF, ZTOF, opt), opt
            end
            return _trellisprofilesQ(XperpTOF, ZTOF, missing)
        end
    else
        if Pauli == ' '
            STOF = trellisorientedformQ(stabilizers(Q))
            nTOF = trellisorientedformQ(normalizermatrix(Q))
            if sect
                opt, _ = optimalsectionalizationQ(STOF, nTOF)
                return _trellisprofilesQ(STOF, nTOF, opt), opt
            end
            return _trellisprofilesQ(STOF, nTOF, missing)
        elseif Pauli == 'X'
            _, _, Zperp, _, _, _ = splitsymplecticstabilizers(quadratictosymplectic(normalizermatrix(Q)), ones(Int64, size(normalizermatrix(Q), 1)))
            X = hcat(Xstabilizers(Q), zero_matrix(field(Q), size(Xstabilizers(Q), 1), size(Xstabilizers(Q), 2)))
            ZperpTOF = trellisorientedformQ(symplectictoquadratic(Zperp))
            XTOF = trellisorientedformQ(symplectictoquadratic(X))
            if sect
                opt, _ = optimalsectionalizationQ(XTOF, ZperpTOF)
                return _trellisprofilesQ(XTOF, ZperpTOF, opt), opt
            end
            return _trellisprofilesQ(XTOF, ZperpTOF, missing)
        else
            Xperp, _, _, _, _, _ = splitsymplecticstabilizers(quadratictosymplectic(normalizermatrix(Q)), ones(Int64, size(normalizermatrix(Q), 1)))
            Z = hcat(zero_matrix(field(Q), size(Zstabilizers(Q), 1), size(Zstabilizers(Q), 2)), Zstabilizers(Q))
            XperpTOF = trellisorientedformQ(symplectictoquadratic(Xperp))
            ZTOF = trellisorientedformQ(symplectictoquadratic(Z))
            if sect
                opt, _ = optimalsectionalizationQ(ZTOF, XperpTOF)
                return _trellisprofilesQ(ZTOF, XperpTOF, opt), opt
            end
            return _trellisprofilesQ(ZTOF, XperpTOF, missing)
        end
    end
end

function syndrometrellis(Q::AbstractQuantumCode, type::String="weight", Pauli::Char=' ',
    sect::Bool=false)

    type ∈ ["weight", "decoding"] || error("Unknown type parameter in syndrometrellis.")
    (Pauli != ' ' && typeof(Q) <: CSSCode) && error("Pauli parameter is non-empty but the code is not CSS.")
    Pauli ∈ [' ', 'X', 'Z'] || error("Unknown Pauli parameter $Pauli; must be ' ', 'X', or 'Z'.")

    if type == "weight"
        if Pauli == ' '
            STOF = trellisorientedformQ(stabilizers(Q))
            nTOF = trellisorientedformQ(normalizermatrix(Q))
            if sect
                profiles, opt = trellisprofiles(Q, type, Pauli, sect)
                return _syndrometrellisQ(profiles, opt, nTOF, STOF, charactervector(Q), Pauli, false)
            else
                profiles = trellisprofiles(Q, type, Pauli)
                return _syndrometrellisQ(profiles, missing, nTOF, STOF, charactervector(Q), Pauli, false)
            end
        elseif Pauli == 'X'
            _, _, Zperp, _, _, _ = splitsymplecticstabilizers(quadratictosymplectic(normalizermatrix(Q)), ones(Int64, size(normalizermatrix(Q), 1)))
            X = hcat(Xstabilizers(Q), zero_matrix(field(Q), size(Xstabilizers(Q), 1), size(Xstabilizers(Q), 2)))
            ZperpTOF = trellisorientedformQ(symplectictoquadratic(Zperp))
            XTOF = trellisorientedformQ(symplectictoquadratic(X))
            if sect
                profiles, opt = trellisprofiles(Q, type, Pauli, sect)
                return _syndrometrellisQ(profiles, opt, ZperpTOF, XTOF, charactervector(Q), 'Z', false)
            else
                profiles = trellisprofiles(Q, type, Pauli)
                return _syndrometrellisQ(profiles, missing, ZperpTOF, XTOF, charactervector(Q), 'Z', false)
            end
        else
            Xperp, _, _, _, _, _ = splitsymplecticstabilizers(quadratictosymplectic(normalizermatrix(Q)), ones(Int64, size(normalizermatrix(Q), 1)))
            Z = hcat(zero_matrix(field(Q), size(Zstabilizers(Q), 1), size(Zstabilizers(Q), 2)), Zstabilizers(Q))
            XperpTOF = trellisorientedformQ(symplectictoquadratic(Xperp))
            ZTOF = trellisorientedformQ(symplectictoquadratic(Z))
            if sect
                profiles, opt = trellisprofiles(Q, type, Pauli, sect)
                return _syndrometrellisQ(profiles, opt, XperpTOF, ZTOF, charactervector(Q), 'X', false)
            else
                profiles = trellisprofiles(Q, type, Pauli)
                return _syndrometrellisQ(profiles, missing, XperpTOF, ZTOF, charactervector(Q), 'X', false)
            end
        end
    else
        if Pauli == ' '
            STOF = trellisorientedformQ(stabilizers(Q))
            nTOF = trellisorientedformQ(normalizermatrix(Q))
            if sect
                profiles, opt = trellisprofiles(Q, type, Pauli, sect)
                return _syndrometrellisQ(profiles, opt, STOF, nTOF, charactervector(Q), Pauli, false)
            else
                profiles = trellisprofiles(Q, type, Pauli)
                return _syndrometrellisQ(profiles, missing, STOF, nTOF, charactervector(Q), Pauli, false)
            end
        elseif Pauli == 'X'
            _, _, Zperp, _, _, _ = splitsymplecticstabilizers(quadratictosymplectic(normalizermatrix(Q)), ones(Int64, size(normalizermatrix(Q), 1)))
            X = hcat(Xstabilizers(Q), zero_matrix(field(Q), size(Xstabilizers(Q), 1), size(Xstabilizers(Q), 2)))
            ZperpTOF = trellisorientedformQ(symplectictoquadratic(Zperp))
            XTOF = trellisorientedformQ(symplectictoquadratic(X))
            if sect
                profiles, opt = trellisprofiles(Q, type, Pauli, sect)
                return _syndrometrellisQ(profiles, opt, XTOF, ZperpTOF, charactervector(Q), Pauli, false)
            else
                profiles = trellisprofiles(Q, type, Pauli)
                return _syndrometrellisQ(profiles, missing, XTOF, ZperpTOF, charactervector(Q), Pauli, false)
            end
        else
            Xperp, _, _, _, _, _ = splitsymplecticstabilizers(quadratictosymplectic(normalizermatrix(Q)), ones(Int64, size(normalizermatrix(Q), 1)))
            Z = hcat(zero_matrix(field(Q), size(Zstabilizers(Q), 1), size(Zstabilizers(Q), 2)), Zstabilizers(Q))
            XperpTOF = trellisorientedformQ(symplectictoquadratic(Xperp))
            ZTOF = trellisorientedformQ(symplectictoquadratic(Z))
            if sect
                profiles, opt = trellisprofiles(Q, type, Pauli, sect)
                return _syndrometrellisQ(profiles, opt, ZTOF, XperpTOF, charactervector(Q), Pauli, false)
            else
                profiles = trellisprofiles(Q, type, Pauli)
                return _syndrometrellisQ(profiles, missing, ZTOF, XperpTOF, charactervector(Q), Pauli, false)
            end
        end
    end
end

max_stars_repo_name: esabo/CodingTheory

max_issues_repo_name: esabo/CodingTheory

max_forks_repo_name: esabo/CodingTheory


text:
# Question to answer, what if you orthogonalize the residual vectors against all
# historical residual vectors???

# ------------------------------------------------------------------------------
# A original flavor of the Conjugate Gradient Algorithm.
# * It performs the Conjugate Gradient and record all the residuals vectors for
#   computing the lanczos iterations.

using LinearAlgebra
using Logging
using Plots

mutable struct CG
    r
    rnew
    d
    A::Function
    x
    b
    itr

    residuals
    reorthogonalize::Bool
    function CG(A::Function, b, x0=nothing)
        this = new()
        this.A = A
        this.x = x0 === nothing ? b .+ 0.1  : x0  # just to handle matrix A that has eigenvalue of exactly 1.
        this.r = b - A(this.x)
        this.rnew = similar(this.r)
        this.d = this.r
        this.itr = 0
        this.residuals = Vector{typeof(this.r)}()
        push!(this.residuals, this.r)
        this.reorthogonalize = true
        return this
    end

    function CG(A::AbstractArray, b::AbstractArray)
        return CG((x)->A*x, b)
    end

end

function (this::CG)()
    r = this.r
    if r == 0
        return 0 # The problem is solved already.
    end
    A = this.A
    d = this.d
    Ad = A(d)

    a = dot(r, r)/dot(d, Ad)
    if a < 0
        error("CG got a non-definite matrix")
    end

    this.x += a*d
    this.rnew = r - a*Ad                # update rnew
    if this.reorthogonalize
        for residual in this.residuals
            this.rnew -= (dot(residual, this.rnew)/dot(residual, residual))*residual
        end
        push!(this.residuals, this.rnew)
    end
    b = dot(this.rnew, this.rnew)/dot(r, r)
    # @assert abs(dot(rnew + β*d, Ad)) < 1e-8 "Not conjugate"
    this.d = this.rnew + b*d
    this.r = this.rnew                      # Override
    this.itr += 1
    return convert(Float64, norm(this.rnew))
end

function GetCurrentResidualNorm(this::CG)
    return norm(this.r)
end

"""
    Turn on the reorthogonalizations using the residual vectors,
    and add current residual to the list of residuals.
"""
function TurnOnReorthgonalize(this::CG)
    this.reorthogonalize = true
    push!(this.residuals, this.r)
    return
end

"""
    Turn off the reorthogonalization on the residual vectors. And then
    clear all the stored residual vectors.
"""
function TurnOffReorthgonalize(this::CG)
    this.reorthogonalize = false
    empty!(this.residuals)
    return
end

### -----------------------Testing Utility -------------------------------------

"""
    Accept a linear system, and a list of guesses from the Lanczos Algorithm,
    it will return a list energy norm of the
    error vector for the linear system.

    Parameters:
        A: The matrix
        b: The vector on the RHS
        Xs: All the guesses vector from the CG Algorithm, including the initial guess!
"""
function EnergyErrorNorm(A, b, Xs)
    XStar = A\b
    InitialError = dot(Xs[1] - XStar, A, Xs[1] - XStar)
    err = (x) -> dot(x - XStar, A, x - XStar)/InitialError

    return err.(Xs)
end

"""
    Bad eigen value distribution, they are distributed from 0.001 to 1 like a
    geometric series.

    * A lot of small eigenvalues are clustered close to zero, few larger ones are
    far away from the other Eigenvalues.

    Parameters:
        rho: Number:
            An number between 0 and 1 that paramaterize the distribution of the
            eigenvalues for the PSD matrix.
        N=20:
            The size of the PSD matrix.

    returns:
        a diagonal matrix.
"""
function GetNastyPSDMatrix(rho::Number, N=20, inverted=false)
    @assert rho <= 1 && rho >= 0
    EigenValues = zeros(N)
    EigenMin, EigenMax = 0.001, 1    # Min Max Eigenvalues.
    EigenValues[1] = EigenMin
    for IdexI in 2:N
        EigenValues[IdexI] = EigenMin +
            ((IdexI - 1)/(N - 1))*(EigenMax - EigenMin)*rho^(N - IdexI)  # formulas
    end
    if inverted
        return I - diagm(EigenValues)
    end
    return diagm(EigenValues)
end

N = 100
A = SymTridiagonal(randn(N), randn(N - 1))
A = A'*A
b = rand(N)
cg1 = CG(A, b)
cg2 = CG(A, b)
cg1Soln = Vector(); push!(cg1Soln, cg1.x)
cg2Soln = Vector(); push!(cg2Soln, cg2.x)

TurnOffReorthgonalize(cg2)
for itr in 1: 1.5*N
    Res2Norm1 = cg1()
    push!(cg1Soln, cg1.x)
    Res2Norm2 = cg2()
    push!(cg2Soln, cg2.x)
    println("With Reorthogonalization: $Res2Norm1, Without: $Res2Norm2")
end

relError1 = EnergyErrorNorm(A, b, cg1Soln)
relError2 = EnergyErrorNorm(A, b, cg2Soln)

fig1 = plot(1:length(relError1), log10.(relError1), label="With Re-Orthogonalization", title="CG Relative Error")
plot!(fig1, 1:length(relError2), log10.(relError2), label="Without Re-Orthogonalization")

ext: jl


text:
function ManifoldTests.test_forwarddiff(M::Manifold, pts, tv)
    return for (p, X) in zip(pts, tv)
        exp_f(t) = distance(M, p, exp(M, p, t[1] * X))
        d12 = norm(M, p, X)
        for t in 0.1:0.1:0.9
            Test.@test d12 ≈ ForwardDiff.derivative(exp_f, t)
        end

        retract_f(t) = distance(M, p, retract(M, p, t[1] * X))
        for t in 0.1:0.1:0.9
            Test.@test ForwardDiff.derivative(retract_f, t) ≥ 0
        end
    end
end


text:
# Conditional statements in Julia
# Navendu Pottekkat | 2020 | Apache-2.0 License
# https://www.apache.org/licenses/LICENSE-2.0

if true
    println("It's true!")
else
    println("It's false!")
end


text:
@testset "Zoom" begin
    @test (Zoom <: Augmentor.AffineOperation) == false
    @test typeof(@inferred(Zoom(1))) <: Zoom <: Augmentor.ImageOperation
    @testset "constructor" begin
        @test_throws MethodError Zoom()
        @test_throws MethodError Zoom(())
        # T not defined:
        @test_throws UndefVarError Zoom{0}(())
        @test_throws MethodError Zoom(:a)
        @test_throws MethodError Zoom([:a])
        @test_throws ArgumentError Zoom([])
        @test_throws ArgumentError Zoom{1}(([:a],))
        @test_throws ArgumentError Zoom(1:3,1:2)
        @test_throws ArgumentError Zoom(1:2,1:2,1)
        @test @inferred(Zoom(1,2)) == @inferred(Zoom{2}((1,2)))
        @test @inferred(Zoom(1)) == @inferred(Zoom{2}((1:1,1:1)))
        @test @inferred(Zoom((1,))) == Zoom{1}((1:1,))
        @test @inferred(Zoom(1,2)) == Zoom{2}((1:1,2:2))
        @test @inferred(Zoom(1.)) == Zoom{2}((1.0:1:1,1.0:1:1))
        @test @inferred(Zoom(1,2,3.)) == Zoom{3}((1.0:1,2.0:2,3.0:3))
        @test @inferred(Zoom(1,2,3)) == Zoom{3}((1:1,2:2,3:3))
        @test @inferred(Zoom(1:1,2,3)) == Zoom{3}((1:1,2:2,3:3))
        @test @inferred(Zoom((1,2,3))) == Zoom{3}((1:1,2:2,3:3))
        @test @inferred(Zoom(1:2,2:3)) == Zoom{2}((1:2,2:3))
        @test @inferred(Zoom((1:2,2:3))) == Zoom{2}((1:2,2:3))
        op = @inferred(Zoom([1,2],2:3))
        @test op.factors == ([1.,2.],[2.,3.])
        @test str_show(op) == "Augmentor.Zoom{2}(([1.0, 2.0], [2.0, 3.0]))"
        @test str_showconst(op) == "Zoom([1.0, 2.0], [2.0, 3.0])"
        @test str_showcompact(op) == "Zoom by I ∈ {1.0×2.0, 2.0×3.0}"
        op = @inferred(Zoom([1,2],[2.,3]))
        @test op.factors == ([1.,2.],[2.,3.])
        op = @inferred(Zoom([1,2]))
        @test op.factors == ([1.,2.],[1.,2.])
        op = @inferred(Zoom(1))
        @test str_show(op) == "Augmentor.Zoom{2}((1, 1))"
        @test str_showconst(op) == "Zoom(1, 1)"
        @test str_showcompact(op) == "Zoom by 1×1"
        op = @inferred(Zoom(0.8,0.9))
        @test str_show(op) == "Augmentor.Zoom{2}((0.8, 0.9))"
        @test str_showconst(op) == "Zoom(0.8, 0.9)"
        @test str_showcompact(op) == "Zoom by 0.8×0.9"
        op = @inferred(Zoom(1:2,3:4))
        @test str_show(op) == "Augmentor.Zoom{2}((1:2, 3:4))"
        @test str_showconst(op) == "Zoom(1:2, 3:4)"
        @test str_showcompact(op) == "Zoom by I ∈ {1×3, 2×4}"
        op = @inferred(Zoom(1:3,3:5))
        @test str_show(op) == "Augmentor.Zoom{2}((1:3, 3:5))"
        @test str_showconst(op) == "Zoom(1:3, 3:5)"
        @test str_showcompact(op) == "Zoom by I ∈ {1×3, 2×4, 3×5}"
        op = @inferred(Zoom(1,2,3))
        @test str_show(op) == "Augmentor.Zoom{3}((1, 2, 3))"
        @test str_showconst(op) == "Zoom(1, 2, 3)"
        @test str_showcompact(op) == "Zoom by 1×2×3"
    end
    @testset "randparam" begin
        @test @inferred(Augmentor.randparam(Zoom(45), rect)) === (45.0, 45.0)
        @test @inferred(Augmentor.randparam(Zoom(20.0, 15.0), rect)) === (20.0, 15.0)
        @test @inferred(Augmentor.randparam(Zoom(20.0, 15.0), (rect, square2))) === (20.0, 15.0)
        @test @inferred(Augmentor.randparam(Zoom(1:2,2:3), rect)) in [(1,2), (2,3)]
        @test @inferred(Augmentor.toaffinemap(Zoom(2,3), rect, (2.,3.))) ≈ AffineMap([2. 0.; 0. 3.], [-1.5,-4.0])
        @test @inferred(Augmentor.toaffinemap(Zoom([0.9,0.4],[0.8,0.3]), rect, (0.9,0.8))) ≈ AffineMap([.9 0.; 0. .8], [0.15,0.4])
    end
    @testset "eager" begin
        @test_throws MethodError Augmentor.applyeager(Zoom(10), nothing)
        @test Augmentor.supports_eager(Zoom) === false
        # TODO: actual content tests (maybe test_reference)
        img_out1 = @inferred Augmentor.applyeager(Zoom(1.5), square2)
        img_out2 = @inferred Augmentor.applyeager(Zoom(0.2), square2)
        @test axes(img_out1) == (1:4, 1:4)
        @test axes(img_out2) == (1:4, 1:4)
        imgs = [
            (square2),
            (view(square2, :, :)),
            (Augmentor.prepareaffine(square2)),
            (OffsetArray(square2, -1, -1)),
            (view(square2, IdentityRange(1:4), IdentityRange(1:4))),
        ]
        @testset "fixed parameter" begin
            for img_in in imgs
                res = @inferred(Augmentor.applyeager(Zoom(1.5), img_in))
                @test parent(res) == parent(img_out1)
                @test typeof(res) == typeof(img_out1)
                res = @inferred(Augmentor.applyeager(Zoom(0.2), img_in))
                @test parent(res) == parent(img_out2)
                @test typeof(res) == typeof(img_out2)
                # test same with tuple of images
                res1, res2 = @inferred(Augmentor.applyeager(Zoom(1.5), (img_in, N0f8.(img_in))))
                @test parent(res1) == parent(img_out1)
                @test parent(res2) == parent(img_out1)
                @test typeof(res1) == typeof(img_out1)
                @test typeof(res2) <: OffsetArray{N0f8}
                res1, res2 = @inferred(Augmentor.applyeager(Zoom(0.2), (img_in, N0f8.(img_in))))
                @test parent(res1) == parent(img_out2)
                @test parent(res2) == parent(img_out2)
                @test typeof(res1) == typeof(img_out2)
                @test typeof(res2) <: OffsetArray{N0f8}
            end
            # check that the affine map is computed for each image
            res1, res2 = @inferred(Augmentor.applyeager(Zoom(1.5), (square, OffsetArray(square,-5,-5))))
            @test collect(res1) == collect(res2)
            @test axes(res1) != axes(res2)
            res1, res2 = @inferred(Augmentor.applyeager(Zoom(1.5), (square, square2)))
            @test res1 == Augmentor.applyeager(Zoom(1.5), square)
            @test res2 == Augmentor.applyeager(Zoom(1.5), square2)
        end
        @testset "random parameter" begin
            for img_in in imgs
                res1, res2 = @inferred(Augmentor.applyeager(Zoom(.5:.1:2., .6:.1:2.1), (img_in, N0f8.(img_in))))
                # make sure same scales are used
                @test res1 == res2
                @test typeof(res1) == typeof(img_out1)
                @test typeof(res2) <: OffsetArray{N0f8}
            end
        end
    end
    @testset "affine" begin
        @test Augmentor.supports_affine(Zoom) === false
    end
    ref = @inferred Augmentor.applyeager(Zoom(2,3), square)
    @testset "affineview" begin
        @test Augmentor.supports_affineview(Zoom) === true
        @test_throws MethodError Augmentor.applyaffineview(Zoom(90), nothing)
        wv = @inferred Augmentor.applyaffineview(Zoom(2,3), Augmentor.prepareaffine(square))
        # TODO: better tests
        @test wv == ref
        @test eltype(wv) == eltype(square)
        @test typeof(wv) <: SubArray
        @test typeof(wv.indices) <: Tuple{Vararg{IdentityRange}}
        @test typeof(parent(wv)) <: InvWarpedView
        @test parent(parent(wv)).itp.coefs === square
    end
    @testset "lazy" begin
        @test Augmentor.supports_lazy(Zoom) === true
        wv = @inferred Augmentor.applylazy(Zoom(2,3), square)
        # TODO: better tests
        @test wv == ref
        @test eltype(wv) == eltype(square)
        @test typeof(wv) <: SubArray
        @test typeof(wv.indices) <: Tuple{Vararg{IdentityRange}}
        @test typeof(parent(wv)) <: InvWarpedView
        @test parent(parent(wv)).itp.coefs === square
    end
    @testset "view" begin
        @test Augmentor.supports_view(Zoom) === false
    end
    @testset "stepview" begin
        @test Augmentor.supports_stepview(Zoom) === false
    end
    @testset "permute" begin
        @test Augmentor.supports_permute(Zoom) === false
    end
end


text:
module A1R1
export GameEnv, GameSpec
include("game.jl")
end


text:
import MultivariatePolynomials
const MP = MultivariatePolynomials

import MultivariateBases
const MB = MultivariateBases

@testset "Monomial selection for certificate" begin
    @polyvar x y z
    @ncpolyvar a b
    @testset "Multipartite error not commutative" for parts in [([a],), ([a], [b])]
        err = ArgumentError("Multipartite Newton polytope not supported with noncommutative variables.")
        @test_throws err SumOfSquares.Certificate.monomials_half_newton_polytope([a*b, b^2], parts)
    end
    @testset "Multipartite error not disjoint: $parts" for parts in [([x], [x]),
                                                                     ([x], [x, y]),
                                                                     ([x], [y], [x, y])]
        err = ArgumentError("Parts are not disjoint in multipartite Newton polytope estimation: $parts.")
        @test_throws err SumOfSquares.Certificate.monomials_half_newton_polytope([x*y, y^2], parts)
    end
    @testset "Unipartite" begin
        @test SumOfSquares.Certificate.monomials_half_newton_polytope([x*y, y^2], tuple()) == [y]
        @test isempty(SumOfSquares.Certificate.monomials_half_newton_polytope([x, y], tuple()))
        @test SumOfSquares.Certificate.monomials_half_newton_polytope([x^2, y^2], tuple()) == [x, y]
        @test SumOfSquares.Certificate.monomials_half_newton_polytope([x^2, y^2], ([x, y],)) == [x, y]
        @test SumOfSquares.Certificate.monomials_half_newton_polytope([x^2, y^2], ([y, x],)) == [x, y]
        @test SumOfSquares.Certificate.monomials_half_newton_polytope([x^2, x^3*y^2, x^4*y^4], tuple()) == [x^2*y^2, x]
    end
    @testset "Non-commutative" begin
        @test SumOfSquares.Certificate.monomials_half_newton_polytope([a^4, a^3*b, a*b*a^2, a*b*a*b], tuple()) == [a^2, a*b]
        @test SumOfSquares.Certificate.monomials_half_newton_polytope([a^2, a^10*b^20*a^11, a^11*b^20*a^10, a^10*b^20*a^20*b^20*a^10], tuple()) == [a^10*b^20*a^10, a]
    end
    @testset "Multipartite" begin
        # In the part [y, z], the degree is between 0 and 2
        X = [x^4, x^2*y^2, x^2*z^2, x^2*y*z, y*z]
        @test SumOfSquares.Certificate.monomials_half_newton_polytope(X, tuple(), apply_post_filter = false) == [x^2, x*y, x*z, y*z, x, y, z]
        function full_test(X, Y, part1, part2)
            @test SumOfSquares.Certificate.monomials_half_newton_polytope(X, (part1,), apply_post_filter = false) == Y
            @test SumOfSquares.Certificate.monomials_half_newton_polytope(X, (part2,), apply_post_filter = false) == Y
            a = SumOfSquares.Certificate.monomials_half_newton_polytope(X, (part2,), apply_post_filter = false)
            @test SumOfSquares.Certificate.monomials_half_newton_polytope(X, (part1, part2), apply_post_filter = false) == Y
            @test SumOfSquares.Certificate.monomials_half_newton_polytope(X, (part2, part1), apply_post_filter = false) == Y
        end
        full_test(X, monovec([x^2, x*y, x*z,      x, y, z]), [x], [y, z])
        full_test(X, monovec([x^2, x*y, x*z, y*z, x,    z]), [y], [x, z])
        full_test(X, monovec([x^2, x*y, x*z, y*z, x, y   ]), [z], [x, y])
        # FIXME: With recursive merging, it should give [x^2, x*y, x*z, x]
        @test SumOfSquares.Certificate.monomials_half_newton_polytope([x^4, x^2*y^2, x^2*z^2, x^2*y*z, y*z], ([x], [y], [z]), apply_post_filter = false) == [x^2, x*y, x*z, y*z, x, y, z]
    end
end

@testset "Random SOS should be SOS" begin
    @polyvar x y
    x = [1, x, y, x^2, y^2, x*y]
    @test_throws ArgumentError randsos(x, monotype=:Unknown)
    for i in 1:10
        for monotype in [:Classic, :Gram]
            p = randsos(x, monotype=monotype)
            @test p isa GramMatrix
            @test isposdef(Matrix(p.Q))
        end
    end
end

function _certificate_api(certificate::Certificate.AbstractCertificate)
    @test Certificate.get(certificate, Certificate.Cone()) isa SumOfSquares.SOSLikeCone
    @test SumOfSquares.matrix_cone_type(typeof(certificate)) <: MOI.AbstractVectorSet
end
function _basis_check_each(basis::MB.AbstractPolynomialBasis, basis_type)
    @test basis isa basis_type
    if basis isa MB.AbstractMonomialBasis
        # This fails if `basis` is `Vector{Monomial{true}}` instead of `MonomialVector{true}`
        # for DynamicPolynomials. This is important as
        # `polynomial(::AbstractMatrix, ::MonomialVector, ::Type)` is implemented but
        # `polynomial(Q::AbstractMatrix, X::AbstractVector, ::Type)` falls back to
        # `dot(X, Q * X)` which does not work (`promote_operation` calls `zero(eltype(X))`
        # which gives `Polynomial{true, Int}` which then tries to multiply a
        # `ScalarAffineFunction{Float64}` with an `Int`).
        monos = basis.monomials
        @test typeof(monos) == typeof(monovec(monos))
        @test issorted(monos, rev=true)
    end
end
function _basis_check(basis, basis_type)
    @test basis isa MB.AbstractPolynomialBasis || basis isa Vector{<:MB.AbstractPolynomialBasis}
    if basis isa Vector
        # FIXME `basis_type` is `Vector{MB.MonomialBasis}` instead of `Vector{MB.MonomialBasis{...}}`
        # Once this is fixed, we should check
        # @test basis isa basis_type
        for b in basis
            _basis_check_each(b, eltype(basis_type))
        end
    else
        _basis_check_each(basis, basis_type)
    end
end

function certificate_api(certificate::Certificate.AbstractIdealCertificate)
    _certificate_api(certificate)
    @polyvar x
    poly = x + 1
    domain = @set x == 1
    @test Certificate.get(certificate, Certificate.ReducedPolynomial(), poly, domain) isa MP.AbstractPolynomial
    _basis_check(Certificate.get(certificate, Certificate.GramBasis(), poly),
                 Certificate.get(typeof(certificate), Certificate.GramBasisType()))
    zbasis = Certificate.zero_basis(certificate)
    @test zbasis <: MB.AbstractPolynomialBasis
    @test zbasis == Certificate.zero_basis_type(typeof(certificate))
end

function certificate_api(certificate::Certificate.AbstractPreorderCertificate)
    _certificate_api(certificate)
    @polyvar x
    poly = x + 1
    domain = @set x >= 1
    processed = Certificate.get(certificate, Certificate.PreprocessedDomain(), domain, poly)
    for idx in Certificate.get(certificate, Certificate.PreorderIndices(), processed)
        _basis_check(Certificate.get(certificate, Certificate.MultiplierBasis(), idx, processed),
                     Certificate.get(typeof(certificate), Certificate.MultiplierBasisType()))
        @test Certificate.get(certificate, Certificate.Generator(), idx, processed) isa MP.AbstractPolynomial
    end
    icert = Certificate.get(certificate, Certificate.IdealCertificate())
    @test icert isa Certificate.AbstractIdealCertificate
    @test typeof(icert) == Certificate.get(typeof(certificate), Certificate.IdealCertificate())
end

@testset "API" begin
    @polyvar x
    cone = SumOfSquares.SOSCone()
    BT = MB.MonomialBasis
    maxdegree = 2
    function _test(certificate::Certificate.AbstractIdealCertificate)
        certificate_api(certificate)
        preorder = Certificate.Putinar(certificate, cone, BT, maxdegree)
        certificate_api(preorder)
        sparsities = Sparsity[VariableSparsity()]
        if certificate isa Certificate.MaxDegree
            push!(sparsities, MonomialSparsity(ChordalCompletion(), 1))
        end
        @testset "$(typeof(sparsity))" for sparsity in sparsities
            certificate_api(Certificate.SparsePreorder(sparsity, preorder))
        end
    end
    basis = BT(monovec([x^2, x]))
    @testset "$(typeof(certificate))" for certificate in [
        Certificate.MaxDegree(cone, BT, maxdegree),
        Certificate.FixedBasis(cone, basis),
        Certificate.Newton(cone, BT, tuple())
    ]
        _test(certificate)
        _test(Certificate.Remainder(certificate))
        if certificate isa Certificate.MaxDegree
            _test(Certificate.SparseIdeal(VariableSparsity(), certificate))
        end
        @testset "$(typeof(sparsity))" for sparsity in [SignSymmetry(), MonomialSparsity(ChordalCompletion(), 1)]
            _test(Certificate.SparseIdeal(sparsity, certificate))
            _test(Certificate.SparseIdeal(sparsity, Certificate.Remainder(certificate)))
        end
    end
end

include("ceg_test.jl")
include("csp_test.jl")
include("sparsity.jl")


text:
# This file was generated, do not modify it. # hide
dtc = DecisionTreeClassifier()
r   = range(dtc, :max_depth, lower=1, upper=5)

max_stars_repo_name: tlienart/MLJTutorials

max_issues_repo_name: tlienart/MLJTutorials

max_forks_repo_name: tlienart/MLJTutorials


text:
using Turing
using Test

obs = [0,1,0,1,1,1,1,1,1,1]

@model constrained_test(obs) = begin
  p ~ Beta(2,2)
  for i = 1:length(obs)
    obs[i] ~ Bernoulli(p)
  end
  p
end

chain = sample(constrained_test(obs), HMC(1000, 1.5, 3)) # using a large step size (1.5)

check_numerical(chain, [:p], [10/14], eps=0.1)


text:
#=
Compute the fixed point of a given operator T, starting from
specified initial condition v.

@author : Spencer Lyon <spencer.lyon@nyu.edu>

@date: 2014-07-05

References
----------

https://lectures.quantecon.org/jl/optgrowth.html
=#

doc"""
Repeatedly apply a function to search for a fixed point

Approximates ``T^∞ v``, where ``T`` is an operator (function) and ``v`` is an initial
guess for the fixed point. Will terminate either when `T^{k+1}(v) - T^k v <
err_tol` or `max_iter` iterations has been exceeded.

Provided that ``T`` is a contraction mapping or similar,  the return value will
be an approximation to the fixed point of ``T``.

##### Arguments

* `T`: A function representing the operator ``T``
* `v::TV`: The initial condition. An object of type ``TV``
* `;err_tol(1e-3)`: Stopping tolerance for iterations
* `;max_iter(50)`: Maximum number of iterations
* `;verbose(2)`: Level of feedback (0 for no output, 1 for warnings only, 2
        for warning and convergence messages during iteration)
* `;print_skip(10)` : if `verbose == 2`, how many iterations to apply between
        print messages

##### Returns
---

* '::TV': The fixed point of the operator ``T``. Has type ``TV``

##### Example

```julia
using QuantEcon
T(x, μ) = 4.0 * μ * x * (1.0 - x)
x_star = compute_fixed_point(x->T(x, 0.3), 0.4)  # (4μ - 1)/(4μ)
```

"""
function compute_fixed_point(T::Function,
                            v::TV;
                            err_tol=1e-4,
                            max_iter=100,
                            verbose=2,
                            print_skip=10) where TV

    if !(verbose in (0, 1, 2))
        throw(ArgumentError("verbose should be 0, 1 or 2"))
    end

    iterate = 0
    err = err_tol + 1
    while iterate < max_iter && err > err_tol
        new_v = T(v)::TV
        iterate += 1
        err = Base.maximum(abs, new_v - v)
        if verbose == 2
            if iterate % print_skip == 0
                println("Compute iterate $iterate with error $err")
            end
        end
        v = new_v
    end

    if verbose >= 1
        if err > err_tol
            warn("max_iter attained in compute_fixed_point")
        elseif verbose == 2
            println("Converged in $iterate steps")
        end
    end

    return v
end


text:
const julia_version = v"1.5.3"

include("../common.jl")

max_stars_repo_name: jameskermode/Yggdrasil

max_issues_repo_name: jameskermode/Yggdrasil

max_forks_repo_name: jameskermode/Yggdrasil


text:
using Dates
using StructTypes

struct Bar
    instrument_token::Int
    interval::String
    timestamp::DateTime
    open::Number
    high::Number
    low::Number
    close::Number
    volume::Int
end

StructTypes.StructType(::Type{Bar}) = StructTypes.Struct()

"""
    historical_data(instrument_token::Integer, interval::String,
                    from_dt::DateTime, to_dt::DateTime, access_token::String)

Get historical data
"""
function historical_data(instrument_token::Integer, interval::String,
                         from_dt::DateTime, to_dt::DateTime, access_token::String)
    url_fragment = "instruments/historical/$instrument_token/$interval"
    fd = Dates.format(from_dt, "yyyy-mm-dd+HH:MM:SS")
    td = Dates.format(to_dt, "yyyy-mm-dd+HH:MM:SS")
    res = http_get(url_fragment * "?from=$fd&to=$td", access_token)

    # return a list of Bar objects
    bars = []
    for candle in res["data"]["candles"]
        bar = Bar(
            instrument_token, interval,
            DateTime(candle[1], "yyyy-mm-ddTH:M:S+0530"),
            candle[2], candle[3], candle[4], candle[5],
            candle[6])

        push!(bars, bar)
    end

    # Remove the last bar. Why?
    # Say we want data from 10:00 to 11:00 i.e. for a duration of 60 minutes, at an interval
    # of 15 minutes.
    # This will give 5 bars
    # 1. 10:00 ohlc
    # 2. 10:15 ohlc
    # 3. 10:30 ohlc
    # 4. 10:45 ohlc
    # 5. 11:00 ohlc
    #
    # The last bar contains data from 11:00 onwards. We don't what this, so we can discard this
    # bar
    if length(bars) > 0
        pop!(bars)
    end

    return bars
end

function accumulate_historical_data(bars, period)
    period_2 = period

    ret = []
    for idx in 1:period:length(bars)-period
        inside_bars = bars[idx:period_2]
        bar1 = inside_bars[1]

        high = max(map(b -> b.high, inside_bars)...)
        low = min(map(b -> b.low, inside_bars)...)
        vol = sum(map(b -> b.volume, inside_bars))

        bar = Bar(
            bar1.instrument_token,
            bar1.interval,
            bar1.timestamp,
            bar1.open,
            high,
            low,
            inside_bars[end].close,
            vol
        )

        push!(ret, bar)
        period_2 += period
    end

    return ret
end


text:
"""
Usage example:

```
using Distributions, LinearAlgebra, TopOpt

f1 = RandomMagnitude([0, -1], Uniform(0.5, 1.5))
f2 = RandomMagnitude(normalize([1, -1]), Uniform(0.5, 1.5))
f3 = RandomMagnitude(normalize([-1, -1]), Uniform(0.5, 1.5))

base_problem = PointLoadCantilever(Val{:Linear}, (160, 40), (1.0, 1.0), 1.0, 0.3, 1.0)
problem = MultiLoad(base_problem, [(160, 20) => f1, (80, 40) => f2, (120, 0) => f3], 10000)
"""
struct MultiLoad{dim, T, TP <: StiffnessTopOptProblem{dim, T}, TF} <: StiffnessTopOptProblem{dim, T}
    problem::TP
    F::TF
end
@forward_property MultiLoad problem
for F in (:getE, :getν, :nnodespercell, :getcloaddict, :getdim, :getpressuredict, :getfacesets)
    @eval $F(p::MultiLoad) = $F(p.problem)
end
function MultiLoad(problem::StiffnessTopOptProblem, N::Int, load_rules::Vector{<:Pair})
    I = Int[]
    J = Int[]
    V = Float64[]
    for (pos, f) in load_rules
        dofs = find_nearest_dofs(problem, pos)
        for i in 1:N
            load = f()
            append!(I, dofs)
            push!(J, fill(i, length(dofs))...)
            append!(V, load)
        end
    end
    F = sparse(I, J, V, ndofs(problem.ch.dh), N)
    return MultiLoad(problem, F)
end
function MultiLoad(
    problem::StiffnessTopOptProblem,
    N::Int,
    dist::Distributions.Distribution = Uniform(-2, 2),
)
    F = generate_random_loads(problem, N, dist, random_direction)
    return MultiLoad(problem, F)
end

function find_nearest_dofs(problem, p)
    grid = problem.ch.dh.grid
    shortest = Inf
    closest = 0
    for (i, n) in enumerate(grid.nodes)
        dist = norm(n.x .- p)
        if dist < shortest
            shortest = dist
            closest = i
        end
    end
    @assert closest != 0
    return problem.metadata.node_dofs[:, closest]
end

struct RandomMagnitude{Tf, Tdist} <: Function
    f::Tf
    dist::Tdist
end
(rm::RandomMagnitude)() = rm.f .* rand(rm.dist)

function random_direction()
    theta = rand() * 2 * π
    return [cos(theta), sin(theta)]
end

function get_surface_dofs(problem::StiffnessTopOptProblem)
	dh = problem.ch.dh
	boundary_matrix = dh.grid.boundary_matrix
	interpolation = dh.field_interpolations[1]
	celliterator = Ferrite.CellIterator(dh)
	node_dofs = problem.metadata.node_dofs

	faces, cells, _ = findnz(boundary_matrix)
    surface_node_inds = Int[]
    for i in 1:length(cells)
	    cellind = cells[i]
	    faceind = faces[i]
	    face = [Ferrite.faces(interpolation)[faceind]...]
	    Ferrite.reinit!(celliterator, cellind)
        nodeinds = celliterator.nodes[face]
        append!(surface_node_inds, nodeinds)
    end
    unique!(surface_node_inds)
    return setdiff(node_dofs[:, surface_node_inds], problem.ch.prescribed_dofs)
end

function generate_random_loads(
    problem::StiffnessTopOptProblem,
    N::Int,
    scalar::Distributions.Distribution = Distributions.Uniform(-2, 2),
    direction::Function = random_direction,
)
    loadrule = () -> direction() .* rand(scalar)
    surface_dofs = get_surface_dofs(problem)

    FI = Int[]
	FJ = Int[]
	FV = Float64[]
    nodeinds = rand(1:size(surface_dofs, 2), N)
	for i in 1:N
        load = loadrule()
        dofs = surface_dofs[:, nodeinds[i]]
	    append!(FI, dofs)
	    push!(FJ, i, i)
	    append!(FV, load)
	end
	return sparse(FI, FJ, FV)
end


text:
struct TestStruct
    data1::Vector
    data2::Matrix
end

function Base.getindex(data::TestStruct, idx::Int)::Tuple
    # 例外処理
    1 <= idx <= length(data.data1) || throw(BoundsError(data.data1, idx))
    1 <= idx <= length(data.data2) || throw(BoundsError(data.data2, idx))

    return (data.data1[idx], data.data2[idx, idx])
end

function Base.setindex!(data::TestStruct, val::Real, idx::Int)
    # 例外処理
    1 <= idx <= length(data.data1) || throw(BoundsError(data.data1, idx))
    1 <= idx <= length(data.data2) || throw(BoundsError(data.data2, idx))

    data.data1[idx] = val
    data.data2[idx, idx] = val

    return
end

hoge = TestStruct(rand(100), rand(100, 100))


text:
using Polynomials
using Jacobi

include("testjacobi.jl")
include("testgauss.jl")