text:
export TAConstraint, TADiagonalConstraint, get_diagonal_constraints, set_clock_list!

struct TADiagonalConstraint{T}
    clock1::Union{Nothing, Symbol}
    clock2::Union{Nothing, Symbol}
    bound::TABound{T}
end
Base.hash(c::TADiagonalConstraint) = hash((hash(c.clock1), hash(c.clock2), hash(c.bound)))
Base.:(==)(c1::TADiagonalConstraint{T},c2::TADiagonalConstraint{T}) where T = c1.clock1==c2.clock1 && c1.clock2==c2.clock2 && c1.bound==c2.bound
Base.isequal(c1::TADiagonalConstraint,c2::TADiagonalConstraint) = c1==c2
function Base.:!(c::TADiagonalConstraint{T}) where T
    bound = TABound(!c.bound.strict,-c.bound.value)
    TADiagonalConstraint{T}(c.clock2, c.clock1, bound)
end

Base.iterate(x::TADiagonalConstraint) = (x, nothing)
Base.iterate(::TADiagonalConstraint, ::Any) = nothing
Base.length(x::TADiagonalConstraint) = 1
Base.Broadcast.broadcastable(c::TADiagonalConstraint) = Ref(c)
Base.adjoint(c::TADiagonalConstraint) = c

function Base.show(io::IO, M::MIME"text/plain", c::TADiagonalConstraint)
    res = if isnothing(c.clock1) && isnothing(c.clock2) && c.bound.value ≥ 0
        "⊤"
    elseif isnothing(c.clock1) && isnothing(c.clock2) && c.bound.value < 0
        "⊥"
    elseif isnothing(c.clock1)
        v = if isinf(c.bound) && c.bound.value > 0
            "-∞"
        elseif isinf(c.bound) && c.bound.value < 0
            "∞"
        else
            "$(-c.bound.value)"
        end
        "$(c.clock2) $(c.bound.strict ? :> : :≥) $(v)"
    elseif isnothing(c.clock2)
        v = if isinf(c.bound) && c.bound.value > 0
            "∞"
        elseif isinf(c.bound) && c.bound.value < 0
            "-∞"
        else
            "$(c.bound.value)"
        end
        "$(c.clock1) $(c.bound.strict ? :< : :≤) $(v)"
    else
        v = if isinf(c.bound) && c.bound.value > 0
            "∞"
        elseif isinf(c.bound) && c.bound.value < 0
            "-∞"
        else
            "$(c.bound.value)"
        end
        "$(c.clock1)-$(c.clock2) $(c.bound.strict ? :< : :≤) $(v)"
    end
    print(io,res)
end

mutable struct TAConstraint{T}
    clocks::Vector{Symbol}
    D::Matrix{TABound{T}}
end

function TAConstraint(terms::Vector{TADiagonalConstraint{T}}, clocks=Symbol[]) where T
    if isempty(clocks)
        for term ∈ terms
            if !isnothing(term.clock1) && !(term.clock1 in clocks)
                push!(clocks, term.clock1)
            end
            if !isnothing(term.clock2) && !(term.clock2 in clocks)
                push!(clocks, term.clock2)
            end
        end
    end
    D = make_dbm(terms, clocks)
    TAConstraint(clocks, D)
end

Base.copy(ta::TAConstraint) = TAConstraint(ta.clocks,ta.D)
Base.deepcopy(ta::TAConstraint) = TAConstraint(deepcopy(ta.clocks), deepcopy(ta.D))

function make_dbm(terms::Vector{TADiagonalConstraint{T}}, clocks) where T
    num_clocks = length(clocks)+1
    D = fill(typemax(TABound{T}), num_clocks, num_clocks)
    strict = zeros(Bool, num_clocks, num_clocks)
    lookup = Dict(zip(clocks,keys(clocks)))

    # Diagonal entries: all clocks are identical to themselves
    D[diagind(D)] .= zero(TABound{T})

    # first row: all clocks are positive
    D[1,:] .= zero(TABound{T})

    # Add specific constraints
    for term in terms
        key1 = isnothing(term.clock1) ? 1 : lookup[term.clock1]+1
        key2 = isnothing(term.clock2) ? 1 : lookup[term.clock2]+1

        # Refine constraints
        D[key1,key2] = min(term.bound,D[key1,key2])
    end
    return D
end

function Base.intersect!(c::TAConstraint{T}, g::TADiagonalConstraint) where T
    idx1 = isnothing(g.clock1) ? 1 : findfirst(==(g.clock1), c.clocks)+1
    idx2 = isnothing(g.clock2) ? 1 : findfirst(==(g.clock2), c.clocks)+1

    if c.D[idx2,idx1] < -g.bound || (isinf(g.bound) && g.bound.value < 0)
        c.D[1,1]=TABound(true,-one(T))
    elseif g.bound < c.D[idx1,idx2]
        c.D[idx1,idx2] = g.bound
        c.D .= min.(c.D, c.D[:,idx1] .+ c.D[idx1,:]', c.D[:,idx2] .+ c.D[idx2,:]')
    end
    c
end
Base.intersect(c::TAConstraint, g::TADiagonalConstraint) = intersect!(deepcopy(c), g)

function Base.intersect!(c1::TAConstraint, others::TAConstraint...)
    c1.D .= min.(c1.D, getfield.(others, :D)...)
    close!(c1)
    c1
end
Base.intersect(c1::TAConstraint, others::TAConstraint...) = intersect!(deepcopy(c1),others...)

Base.iterate(x::TAConstraint) = (x, nothing)
Base.iterate(::TAConstraint, ::Any) = nothing
Base.length(x::TAConstraint) = 1
Base.Broadcast.broadcastable(c::TAConstraint) = Ref(c)
Base.adjoint(c::TAConstraint) = c

function Base.issubset(c1::TAConstraint, c2::TAConstraint)
    @assert size(c1.D) == size(c2.D) "Must be of same size, got $(size(c1.D)) and $(size(c2.D))"
    for idx ∈ eachindex(c1.D)
        if c1.D[idx] > c2.D[idx]
            return false
        end
    end
    return true
end

Base.hash(c::TAConstraint) = hash((hash(c.clocks), hash(c.D)))
Base.:(==)(c1::TAConstraint, c2::TAConstraint) = (c1.clocks == c2.clocks) && (c1.D==c2.D)

function Base.isempty(c::TAConstraint)
    D  = c.D
    D′ = D'
    for idx ∈ eachindex(c.D)
        c1 = D[idx]
        c2 = D′[idx]
        # Situation: -m2 </≤ x-y </≤ m1
        # if either of the inequalities is strict, we have a problem if m1 ≤ -m2
        # if neither of the inequalities is strict, we only have a problem if m1 < -m2
        if (
                (c1.strict || c2.strict) && (
                    c1.value ≤ -c2.value ||
                    (isinf(c1.value) && c1.value < 0) ||
                    (isinf(c2.value) && c2.value < 0)
                )
            ) || (c1.value < -c2.value)
            return true
        end
    end
    return false
end

function set_clock_list!(ta::TAConstraint{T}, clocks::Vector{Symbol}; mapping=identity) where T
    clocks = copy(clocks)
    D = make_dbm(TADiagonalConstraint{T}[], clocks)

    @assert mapping.(ta.clocks) ⊆ clocks "Not all clocks of the constraint set ($(mapping.(ta.clocks))) are in the given list of clocks ($(clocks))."

    lookup = Dict(zip(clocks,keys(clocks)))
    for (i,clock1) ∈ enumerate([nothing; ta.clocks])
        idx1 = isnothing(clock1) ? 1 : lookup[mapping(clock1)]+1
        for (j,clock2) ∈ enumerate([nothing; ta.clocks])
            idx2 = isnothing(clock2) ? 1 : lookup[mapping(clock2)]+1
            D[idx1,idx2] = ta.D[i,j]
        end
    end

    ta.clocks = clocks
    ta.D = D
    return ta
end

function get_diagonal_constraints(b::TAConstraint{T}) where T
    if length(b.clocks) > 0
        TADiagonalConstraint{T}.([nothing; b.clocks], reshape([nothing; b.clocks],(1,:)), b.D)[2:end]
    else
        TADiagonalConstraint{T}[]
    end
end

function Base.show(io::IO, ::MIME"text/plain", b::TAConstraint)
    if isempty(b)
        print(io, "∅")
        return
    end
    constraints = String[]
    for j ∈ 1:length(b.clocks)+1
        for i ∈ j+1:length(b.clocks)+1
            c1 = b.D[i,j]
            c2 = b.D[j,i]
            clock = j==1 ? "$(b.clocks[i-1])" : "$(b.clocks[i-1])-$(b.clocks[j-1])"
            v1 = isinf(c1) ? (c1.value>0 ? "∞" : "-∞") : "$(c1.value)"
            v2 = isinf(c2) ? (c2.value>0 ? "-∞" : "∞") : "$(-c2.value)"
            if -c2 < c1
                if isinf(c1) && isinf(c2)
                    ""
                else
                    lhs = isinf(c2) || (j==1 && iszero(c2.value)) ? "" : "$(v2) $(c2.strict ? :< : :≤) "
                    rhs = isinf(c1) ? "" : " $(c1.strict ?  :< : :≤) $(v1)"
                    if !isempty(lhs) || !isempty(rhs)
                        push!(constraints, lhs*"$(clock)"*rhs)
                    end
                end
            elseif c1 < -c2
                push!(constraints, "⊥")
            else
                push!(constraints, "$(clock) == $(v1)")
            end
        end
    end
    print(io, join(constraints, "\n"))
end


text:
# Miscellaneous useful functions

export printover

"Copy a type, modifying specified fields"
function cpmod{T}(pp::T, di)
    di = !isa(di, Associative) ? Dict(di) : di
    ns = fieldnames(pp)
    args = Array{Any}(length(ns))
    for (i,n) in enumerate(ns)
        args[i] = get(di, n, getfield(pp, n))
    end
    T(args...)
end
cpmod{T}(pp::T; kws...) = cpmod(pp, kws)
# TODO: Remove uses of cpmod once default constructors for immutables are
# changed

"Map a function across rows of a 2D array"
maprows(f, m::Matrix) = mapslices(f, m, 2)

"Is this value not NaN?"
notnan(x) = !isnan(x)

"Does this contain any NaN values?"
hasnan(x) = any(isnan.(x))

"Wait for ENTER to be pressed before continuing"
function wait_for_enter()
    println("Press enter↵ to continue...")
    readline(STDIN)
end

"Print over a previous line"
function printover(line)
    if isdefined(Main, :IJulia)
        print("\r" * line)
    else
        print("\u1b[1G")   # go to first column
        print(line)
        print("\u1b[K")    # clear the rest of the line
    end
end

ext: jl


text:
@testset "ResidueRingMultGrp" begin

  function verify_order(g::Hecke.NfOrdQuoRingElem,o)
    g == 0 && return false
    g^o == 1 || return false
    for l in keys(factor(o).fac)
      g^div(o,l) == 1 && return false
    end
    return true
  end

  function verify_order(g,i,o)
    iszero(g) && return false
    powermod(g,o,i.gen_one) - 1 in i || return false
    for l in keys(factor(o).fac)
      powermod(g,div(o,l),i.gen_one) - 1  in i && return false
    end
    return true
  end

  @testset "multiplicative_group" begin
    Qx, x = PolynomialRing(FlintQQ, "x");

    @testset "K = Q" begin
      K, a = NumberField(x,"a");
      O = maximal_order(K)

      @testset "m0 = <$n>" for n in 1:50
        m0 = ideal(O,O(n))
        Q = NfOrdQuoRing(O,m0)
        G, M = multiplicative_group(Q)
        @test issnf(G)
        SNF = G.snf
        gens = M.generators
        @test length(SNF) == length(gens)
        for i in 1:length(gens)
          @test verify_order(gens[i],SNF[i])
        end
        @test M(G(G.snf)) == Q(O(1))
        for g in gens
          for exp in -5:10
            el = g^exp
            @test M(M\el) == el
          end
        end
      end

      @testset "m0 = <1361>^3000" begin
        m0 = ideal(O,O(1361))^100
        Q = NfOrdQuoRing(O,m0)
        G, M = multiplicative_group(Q)
        @test issnf(G)
        SNF = G.snf
        gens = M.generators
        @test length(SNF) == length(gens)
        for i in 1:length(gens)
          @test verify_order(gens[i],SNF[i])
        end
        @test M(G(G.snf)) == Q(O(1))
        for g in gens
          for exp in -5:10
            el = g^exp
            @test M(M\el) == el
          end
        end
      end

    end

    @testset "f = x^2-x+15" begin
      f = x^2-x+15
      K, a = NumberField(f,"a");
      O = maximal_order(K)

      @testset "m0 = <2>" begin
        m0 = ideal(O,O(2))
        Q = NfOrdQuoRing(O,m0)
        G, M = multiplicative_group(Q)
        @test issnf(G)
        SNF = G.snf
        gens = M.generators
        @test length(SNF) == length(gens)
        for i in 1:length(gens)
          @test verify_order(gens[i],SNF[i])
        end
        @test order(G) == 3
        structt = [3]
        H = DiagonalGroup(structt)
        @test Hecke.isisomorphic(G,H)
        @test M(G(G.snf)) == Q(O(1))
        for g in gens
          for exp in -5:10
            el = g^exp
            @assert M(M\el) == el
          end
        end
      end

      @testset "m0 = <4>" begin
        m0 = ideal(O,O(4))
        Q = NfOrdQuoRing(O,m0)
        G, M = multiplicative_group(Q)
        @test issnf(G)
        SNF = G.snf
        gens = M.generators
        @test length(SNF) == length(gens)
        for i in 1:length(gens)
          @test verify_order(gens[i],SNF[i])
        end
        @test order(G) == 12
        structt = [3,2,2]
        H = DiagonalGroup(structt)
        @test Hecke.isisomorphic(G,H)
        @test M(G(G.snf)) == Q(O(1))
        for g in gens
          for exp in -5:10
            el = g^exp
            @assert M(M\el) == el
          end
        end
      end

    end

    @testset "x^12+..." begin
      #= f = x^12+7*x^11-97*x^10-859*x^9+2558*x^8+38839*x^7+30012*x^6-710649*x^5-2189082*x^4+2534629*x^3+25314673*x^2+43623088*x+28168151 =#
      #= K, a = NumberField(f,"a"); =#
      #= O = maximal_order(K) =#

      #= @testset "m0 = <3>" begin =#
      #=   println("m0 = <3>") =#
      #=   m0 = ideal(O,O(3)) =#
      #=   Q = NfOrdQuoRing(O,m0) =#
      #=   G = domain(Hecke.multiplicative_group(Q)) =#
      #=   println(snf(G)[1].snf) =#
      #=   @test order(G) == 512000 =#
      #=   structt = [80,80,80] =#
      #=   H = DiagonalGroup(structt) =#
      #=   @test Hecke.is_isomorphic(G,H) =#
      #=   @test M(G(G.snf)) == Q(O(1)) =#
      #=   for g in gens =#
      #=     for exp in -5:10 =#
      #=       el = g^exp =#
      #=       @assert M(M\el) == el =#
      #=     end =#
      #=   end =#
      #= end =#

      #= @testset "m0 = <4>" begin =#
      #=   println("m0 = <4>") =#
      #=   m0 = ideal(O,O(4)) =#
      #=   Q = NfOrdQuoRing(O,m0) =#
      #=   G = domain(Hecke.multiplicative_group(Q)) =#
      #=   println(snf(G)[1].snf) =#
      #=   @test order(G) == 14745600 =#
      #=   structt = [15,2,2,2,2,15,4,4,4,4,2,2,2,2] =#
      #=   H = DiagonalGroup(structt) =#
      #=   @test Hecke.is_isomorphic(G,H) =#
      #=   @test M(G(G.snf)) == Q(O(1)) =#
      #=   for g in gens =#
      #=     for exp in -5:10 =#
      #=       el = g^exp =#
      #=       @assert M(M\el) == el =#
      #=     end =#
      #=   end =#
      #= end =#

      #= @testset "m0 = <5>" begin =#
      #=   println("m0 = <5>") =#
      #=   m0 = ideal(O,O(5)) =#
      #=   Q = NfOrdQuoRing(O,m0) =#
      #=   G = domain(Hecke.multiplicative_group(Q)) =#
      #=   println(snf(G)[1].snf) =#
      #=   @test order(G) == 187500000 =#
      #=   structt = [24,5,5,5,5,5,5,4,5,5,5] =#
      #=   H = DiagonalGroup(structt) =#
      #=   @test Hecke.is_isomorphic(G,H) =#
      #=   @test M(G(G.snf)) == Q(O(1)) =#
      #=   for g in gens =#
      #=     for exp in -5:10 =#
      #=       el = g^exp =#
      #=       @assert M(M\el) == el =#
      #=     end =#
      #=   end =#
      #= end =#

      #= @testset "m0 = <60>" begin =#
      #=   println("m0 = <60>") =#
      #=   m0 = ideal(O,O(60)) =#
      #=   Q = NfOrdQuoRing(O,m0) =#
      #=   G = domain(Hecke.multiplicative_group(Q)) =#
      #=   println(snf(G)[1].snf) =#
      #=   @test order(G) == 14155776*fmpz(10)^14 =#
      #=   structt = [15,2,2,2,2,15,4,4,4,4,2,2,2,2,80,80,80,24,5,5,5,5,5,5,4,5,5,5] =#
      #=   H = DiagonalGroup(structt) =#
      #=   @test Hecke.is_isomorphic(G,H) =#
      #=   @test M(G(G.snf)) == Q(O(1)) =#
      #=   for g in gens =#
      #=     for exp in -5:10 =#
      #=       el = g^exp =#
      #=       @assert M(M\el) == el =#
      #=     end =#
      #=   end =#
      #= end =#
    end
  end

  @testset "_multgrp" begin
    Qx, x = PolynomialRing(FlintQQ, "x");

    @testset "K = Q" begin
      K, a = NumberField(x,"a");
      O = maximal_order(K)

      @testset "m0 = <$n>" for n in 1:150
        m0 = ideal(O,O(n))
        Q = NfOrdQuoRing(O,m0)
        gens, structure, disc_log = Hecke._multgrp(Q)
        @test length(gens) == length(structure)
        @test typeof(gens) == Vector{NfOrdQuoRingElem}
        for i in 1:length(gens)
          @test verify_order(gens[i],structure[i])
        end
        G = DiagonalGroup(structure)
        H = Hecke.multgrp_of_cyclic_grp(n)
        @test Hecke.isisomorphic(G,H)
        #= # Test discrete logarithm =#
        for bas in gens
          for exp in [-1,0,1,6]
            el = bas^exp
            dlog = disc_log(el)
            prod = 1
            for j in 1:length(gens)
              prod *= gens[j]^dlog[j]
            end
            @test el == prod
          end
        end
      end
    end

    @testset "K = Q[√2]" begin
      K, a = NumberField(x^2 - 2,"a")
      O = maximal_order(K)

      @testset "m0 = <$n>" for n in 1:100
        m0 = ideal(O,O(n))
        Q = NfOrdQuoRing(O,m0)
        gens, structure, disc_log = Hecke._multgrp(Q)
        @test length(gens) == length(structure)
        for i in 1:length(gens)
          @test verify_order(gens[i],structure[i])
        end
        #= # Test discrete logarithm =#
        for bas in gens
          for exp in [-1,0,1,8]
            el = bas^exp
            dlog = disc_log(el)
            prod = 1
            for j in 1:length(gens)
              prod *= gens[j]^dlog[j]
            end
            @test el == prod
          end
        end
      end
    end

    @testset "f = x^2-x+15" begin
      f = x^2-x+15
      K, a = NumberField(f,"a");
      O = maximal_order(K)

      @testset "m0 = <2>" begin
        m0 = ideal(O,O(2))
        Q = NfOrdQuoRing(O,m0)
        gens, structure, disc_log = Hecke._multgrp(Q)
        @test typeof(structure)==Array{fmpz,1}
        @test typeof(gens)== Array{NfOrdQuoRingElem,1}
        @test length(gens) == length(structure)
        for i in 1:length(gens)
          @test verify_order(gens[i],structure[i])
        end
        G = DiagonalGroup(structure)
        @test order(G) == 3
        H = DiagonalGroup([3])
        @test Hecke.isisomorphic(G,H)
        #= # Test discrete logarithm =#
        for bas in gens
          for exp in [-11,4,12]
            el = bas^exp
            dlog = disc_log(el)
            prod = 1
            for j in 1:length(gens)
              prod *= gens[j]^dlog[j]
            end
            @test el == prod
          end
        end
      end

      @testset "m0 = <4>" begin
        m0 = ideal(O,O(4))
        Q = NfOrdQuoRing(O,m0)
        gens, structure, disc_log = Hecke._multgrp(Q)
        @test length(gens) == length(structure)
        for i in 1:length(gens)
          @test verify_order(gens[i],structure[i])
        end
        G = DiagonalGroup(structure)
        @test order(G) == 12
        H = DiagonalGroup([3,2,2])
        @test Hecke.isisomorphic(G,H)
        #= # Test discrete logarithm =#
        for bas in gens
          for exp in [-8,0,2]
            el = bas^exp
            dlog = disc_log(el)
            prod = 1
            for j in 1:length(gens)
              prod *= gens[j]^dlog[j]
            end
            @test el == prod
          end
        end
      end

      @testset "m0 = <5>" begin
        m0 = ideal(O,O(5))
        Q = NfOrdQuoRing(O,m0)
        gens, structure, disc_log = Hecke._multgrp(Q)
        @test length(gens) == length(structure)
        for i in 1:length(gens)
          @test verify_order(gens[i],structure[i])
        end
        G = DiagonalGroup(structure)
        @test order(G) == 16
        H = DiagonalGroup([4,4])
        @test Hecke.isisomorphic(G,H)
        #= # Test discrete logarithm =#
        for bas in gens
          for exp in [-3,2,4]
            el = bas^exp
            dlog = disc_log(el)
            prod = 1
            for j in 1:length(gens)
              prod *= gens[j]^dlog[j]
            end
            @test el == prod
          end
        end
      end

      @testset "m0 = <20>" begin
        m0 = ideal(O,O(20))
        Q = NfOrdQuoRing(O,m0)
        gens, structure, disc_log = Hecke._multgrp(Q)
        @test length(gens) == length(structure)
        for i in 1:length(gens)
          @test verify_order(gens[i],structure[i])
        end
        G = DiagonalGroup(structure)
        @test order(G) == 192
        H = DiagonalGroup([3,2,2,4,4])
        @test Hecke.isisomorphic(G,H)
        #= # Test discrete logarithm =#
        for bas in gens
          for exp in [-4,1,7]
            el = bas^exp
            dlog = disc_log(el)
            prod = 1
            for j in 1:length(gens)
              prod *= gens[j]^dlog[j]
            end
            @test el == prod
          end
        end
      end

    end

  end

  @testset "_multgrp_mod_pv" begin
    Qx, x = PolynomialRing(FlintQQ, "x");

    @testset "K = Q" begin
      K, a = NumberField(x,"a");
      O = maximal_order(K)

      @testset "i = <$pnum>^$v" for (pnum,v) in [(pnum,v) for pnum in [ x for x in 1:50 if isprime(fmpz(x))], v in [1,2,4,17]]
        #p = ideal(O,O(pnum))
        p = prime_decomposition(O, pnum)[1][1]
        #p = collect(keys(factor(p)))[1]
        g, d, disc_log = Hecke._multgrp_mod_pv(p,v)
        @test length(g) == length(d)
        for i in 1:length(g)
          @test verify_order(g[i],p^v,d[i])
        end
        G = DiagonalGroup(d)
        H = Hecke.multgrp_of_cyclic_grp(fmpz(pnum)^v)
        @test Hecke.isisomorphic(G,H)
        # Test discrete logarithm
        Q = NfOrdQuoRing(O,p^v)
        for bas in g
          for exp in [-1,1,6]
            el = Q(bas)^exp
            dlog = disc_log(el.elem)
            prod = 1
            for j in 1:length(g)
              prod *= Q(g[j])^dlog[j]
            end
            @test el == prod
          end
        end
      end
    end

    @testset "f = x^2-x+15" begin
      f = x^2-x+15
      K, a = NumberField(f,"a");
      O = maximal_order(K)
      i = ideal(O,O(20))

      @testset "p = <$(p.gen_one),$(p.gen_two)>, v = $v" for (p,v) in factor(i)
        gens , structure , disc_log = Hecke._multgrp_mod_pv(p,v)
        @test length(gens) == length(structure)
        for i in 1:length(gens)
          @test verify_order(gens[i],p^v,structure[i])
        end
        # Test discrete logarithm
        Q = NfOrdQuoRing(O,p^v)
        for bas in gens
          for exp in [-1,0,1,2]
            el = Q(bas)^exp
            dlog = disc_log(el.elem)
            prod = 1
            for j in 1:length(gens)
              prod *= Q(gens[j])^dlog[j]
            end
            @test el == prod
          end
        end
      end
    end

    @testset "f = x^3+8*x^2+6*x-17" begin
      f = x^3+8*x^2+6*x-17
      K, a = NumberField(f,"a");
      O = maximal_order(K)
      p = ideal(O,O(3))
      #p = collect(keys(factor(p)))[1]
      structures = Vector{fmpz}[[26],[26,3,3,3],[26,9,9,9],[26,27,27,27]]
      @testset "v = $v" for v in 1:length(structures)
        g, d, disc_log = Hecke._multgrp_mod_pv(p, v)
        @test length(g) == length(d)
        for i in 1:length(g)
          @test verify_order(g[i],p^v,d[i])
        end
        G = DiagonalGroup(d)
        @test order(G) == prod(structures[v])
        H = DiagonalGroup(structures[v])
        @test Hecke.isisomorphic(G,H)
        # Test discrete logarithm
        Q = NfOrdQuoRing(O,p^v)
        for bas in g
          for exp in [-2,1,3]
            el = Q(bas)^exp
            dlog = disc_log(el.elem)
            prod = 1
            for j in 1:length(g)
              prod *= Q(g[j])^dlog[j]
            end
            @test el == prod
          end
        end
      end
    end

    @testset "f = x^4+11*x^3-19*x^2-8*x+7" begin
      f = x^4+11*x^3-19*x^2-8*x+7
      K, a = NumberField(f,"a");
      O = maximal_order(K)
      p = ideal(O,fmpz(3),O(2+a+a^2))
      p = collect(keys(factor(p)))[1]
      structures = Vector{fmpz}[[8],[8,3,3],[8,3,3,3,3],[8,3,3,3,3,9],[8,3,3,9,9,9],[8,3,9,9,9,27]]
      @testset "v = $v" for v in 1:length(structures)
        g , d , disc_log = Hecke._multgrp_mod_pv(p,v)
        @test length(g) == length(d)
        for i in 1:length(g)
          @test verify_order(g[i],p^v,d[i])
        end
        G = DiagonalGroup(d)
        @test order(G) == prod(structures[v])
        H = DiagonalGroup(structures[v])
        @test Hecke.isisomorphic(G,H)
        # Test discrete logarithm
        Q = NfOrdQuoRing(O,p^v)
        for bas in g
          for exp in [-4,1,6]
            el = Q(bas)^exp
            dlog = disc_log(el.elem)
            prod = 1
            for j in 1:length(g)
              prod *= Q(g[j])^dlog[j]
            end
            @test el == prod
          end
        end
      end
    end

    @testset "f = x^6+6x^5-12*x^4-x^3-6*x^2+9*x+20" begin
      f = x^6+6x^5-12*x^4-x^3-6*x^2+9*x+20
      K, a = NumberField(f,"a");
      O = maximal_order(K)
      p = ideal(O,fmpz(3),O(2+2*a+a^2))
      p = collect(keys(factor(p)))[1]
      structures = Vector{fmpz}[[8],[8,3,3],[8,3,3,3,3],[8,9,9,3,3],[8,9,9,3,3,3,3],[8,9,9,9,9,3,3]]
      @testset "v = $v" for v in 1:length(structures)
        g , d , disc_log = Hecke._multgrp_mod_pv(p,v)
        @test length(g) == length(d)
        for i in 1:length(g)
          @test verify_order(g[i],p^v,d[i])
        end
        G = DiagonalGroup(d)
        @test order(G) == prod(structures[v])
        H = DiagonalGroup(structures[v])
        @test Hecke.isisomorphic(G,H)
        # Test discrete logarithm
        Q = NfOrdQuoRing(O,p^v)
        for bas in g
          for exp in [-1,0,1,6]
            el = Q(bas)^exp
            dlog = disc_log(el.elem)
            prod = 1
            for j in 1:length(g)
              prod *= Q(g[j])^dlog[j]
            end
            @test el == prod
          end
        end
      end
    end
  end

  @testset "_multgrp_mod_p" begin
    Qx, x = PolynomialRing(FlintQQ, "x");

    @testset "K = Q" begin
      K, a = NumberField(x,"a");
      O = maximal_order(K)

      @testset "p = <$(pnum)>" for pnum in [ x for x in 1:50 if isprime(fmpz(x))]
        p = ideal(O,O(pnum))
        #p = collect(keys(factor(p)))[1]
        g , n , dlog = Hecke._multgrp_mod_p(p)
        @test isa(g,NfOrdElem)
        @test isa(n,fmpz)
        @test order(p) == O
        @test n == pnum-1
        @test verify_order(g,p,n)
        for exp in [0,1,2,n-2,n-1,n,3234239]
          @test dlog(g^exp) == mod(fmpz(exp),n)
        end
      end
    end

    @testset "K = Q[√2]" begin
      K, a = NumberField(x^2-2,"a");
      O = maximal_order(K)

      primeideals = Vector{Hecke.NfOrdIdl}()
      for pnum in [ x for x in 1:40 if isprime(fmpz(x))]
        append!(primeideals,collect(keys(factor(ideal(O,O(pnum))))))
      end

      @testset "p = <$(p.gen_one), $(p.gen_two)>" for p in primeideals
        g , n , dlog = Hecke._multgrp_mod_p(p)
        @test isa(g,NfOrdElem)
        @test order(p) == O
        @test isa(n,fmpz)
        @test !iszero(g)
        @test !(g in p)
        @test verify_order(g,p,n)
        Q = NfOrdQuoRing(O,p)
        for exp in [0,1,2,n-2,n-1,n]
          @test dlog((Q(g)^exp).elem) == mod(fmpz(exp),n)
        end
      end
    end

    @testset "K = Q[x]/<x^6+6*x^5-12*x^4-x^3-6*x^2+9*x+20>" begin
      K, a = NumberField(x^6+6*x^5-12*x^4-x^3-6*x^2+9*x+20,"a");
      O = maximal_order(K)

      primeideals = Vector{Hecke.NfOrdIdl}()
      for pnum in [ x for x in 1:20 if isprime(fmpz(x)) ]
        append!(primeideals,collect(keys(factor(ideal(O,O(pnum))))))
      end

      @testset "p = <$(p.gen_one), $(p.gen_two)>" for p in primeideals
        g , n , dlog = Hecke._multgrp_mod_p(p)
        @test isa(g,NfOrdElem)
        @test order(p) == O
        @test isa(n,fmpz)
        @test !iszero(g)
        @test !(g in p)
        @test verify_order(g,p,n)
        Q = NfOrdQuoRing(O,p)
        for exp in [24,n-345]
          @test dlog((Q(g)^exp).elem) == mod(fmpz(exp),n)
        end
      end
    end

    @testset "K = Q[x]/<x^10-x^9+x^8-x^7+x^6-x^5+x^4-x^3+x^2-x+1>" begin
      K, a = NumberField(x^10-x^9+x^8-x^7+x^6-x^5+x^4-x^3+x^2-x+1,"a");
      O = maximal_order(K)

      primeideals = Vector{Hecke.NfOrdIdl}()
      for pnum in [2,3,5,11]
        append!(primeideals,collect(keys(factor(ideal(O,O(pnum))))))
      end

      @testset "p = <$(p.gen_one), $(p.gen_two)>" for p in primeideals
        g , n , dlog = Hecke._multgrp_mod_p(p)
        @test isa(g,NfOrdElem)
        @test isa(n,fmpz)
        @test order(p) == O
        @test !iszero(g)
        @test !(g in p)
        @test verify_order(g,p,n)
        Q = NfOrdQuoRing(O,p)
        for exp in [50,n-30]
          @test dlog((Q(g)^exp).elem) == mod(fmpz(exp),n)
        end
      end
    end
  end

  @testset "_1_plus_p_mod_1_plus_pv" begin
    Qx, x = PolynomialRing(FlintQQ, "x");

    @testset "Method: $method" for method in [:quadratic,:artin_hasse,:p_adic]

      @testset "K = Q" begin
        K, a = NumberField(x,"a");
        O = maximal_order(K)

        @testset "p = <$(pnum)>, v = $(v)" for pnum in [ x for x in 1:30 if isprime(fmpz(x))], v in [1,2,3,4,11,30]
          i = ideal(O,O(pnum))
          p = i
          #p = collect(keys(factor(i)))[1]
          g, D , disc_log = Hecke._1_plus_p_mod_1_plus_pv(p,v;method=method)
          @test typeof(D)==Array{fmpz,1}
          @test length(g) == length(D)
          # Test generators
          for i in 1:length(g)
            if(D[i] != 1)
              @test isa(g[i],NfOrdElem)
              @test parent(g[i]) == O
              @test g[i]-1 in p
              @test verify_order(g[i],p^v,D[i])
            end
          end
          # Test structure/relations
          G = DiagonalGroup(D)
          H = Hecke.multgrp_of_cyclic_grp(fmpz(pnum)^v)
          I = DiagonalGroup([Hecke._multgrp_mod_p(p)[2]])
          @test order(G) == div(order(H), order(I))
          check, J = Hecke._1_plus_pa_mod_1_plus_pb_structure(p,1,v)
          if check
            @test Hecke.isisomorphic(G,J)
          end
          # Test discrete logarithm
          Q = NfOrdQuoRing(O,p^v)
          for bas in g
            for exp in -1:2
              el = Q(bas)^exp
              dlog = disc_log(el.elem)
              prod = 1
              for j in 1:length(g)
                prod *= Q(g[j])^dlog[j]
              end
              @test el == prod
            end
          end
        end
      end

      @testset "K = Q[√2]" begin
        K, a = NumberField(x^2 - 2,"a")
        O = maximal_order(K)

        primeideals = Vector{Hecke.NfOrdIdl}()
        for pnum in [2,3,5,7,19]
          fac = factor(ideal(O,O(pnum)))
          ks = collect(keys(fac))
          vs = collect(values(fac))
          all(isone,vs) && append!(primeideals,ks)
        end

        @testset "p = <$(p.gen_one), $(p.gen_two)>, v = $(v)" for p in primeideals, v in [1,4,11]
          g, D, disc_log = Hecke._1_plus_p_mod_1_plus_pv(p,v;method=method)
          @test length(g) == length(D)
          # Test generators
          for i in 1:length(g)
            if(D[i] != 1)
              @test isa(g[i],NfOrdElem)
              @test parent(g[i]) == O
              @test g[i]-1 in p
              @test verify_order(g[i],p^v,D[i])
            end
          end
          # Test structure
          check, J = Hecke._1_plus_pa_mod_1_plus_pb_structure(p,1,v)
          if check
            G = DiagonalGroup(D)
            @test Hecke.isisomorphic(G,J)
          end
          # Test discrete logarithm
          Q = NfOrdQuoRing(O,p^v)
          for bas in g
            for exp in [-1,1,2]
              el = Q(bas)^exp
              dlog = disc_log(el.elem)
              prod = 1
              for j in 1:length(g)
                prod *= Q(g[j])^dlog[j]
              end
              @test el == prod
            end
          end
        end
      end

      @testset "K = Q[x]/<f = x^6+...>" begin
        f = x^6 + 6*x^5 - 12*x^4 - x^3 - 6*x^2 + 9*x + 20
        K, a = NumberField(f,"a")
        O = maximal_order(K)

        primeideals = Vector{Hecke.NfOrdIdl}()
        for pnum in [2,11,13]
          fac = factor(ideal(O,O(pnum)))
          ks = collect(keys(fac))
          vs = collect(values(fac))
          all(isone,vs) && append!(primeideals,ks)
        end

        @testset "p = <$(p.gen_one), $(p.gen_two)>, v = $(v)" for p in primeideals, v in [2,9]
          g, D, disc_log = Hecke._1_plus_p_mod_1_plus_pv(p,v;method=method)
          @test length(g) == length(D)
          # Test generators
          for i in 1:length(g)
            if(D[i] != 1)
              @test isa(g[i],NfOrdElem)
              @test parent(g[i]) == O
              @test g[i]-1 in p
              @test verify_order(g[i],p^v,D[i])
            end
          end
          # Test structure
          check, J = Hecke._1_plus_pa_mod_1_plus_pb_structure(p,1,v)
          if check
            G = DiagonalGroup(D)
            @test Hecke.isisomorphic(G,J)
          end
          # Test discrete logarithm
          Q = NfOrdQuoRing(O,p^v)
          for bas in g
            for exp in [-1,2]
              el = Q(bas)^exp
              dlog = disc_log(el.elem)
              prod = 1
              for j in 1:length(g)
                prod *= Q(g[j])^dlog[j]
              end
              @test el == prod
            end
          end
        end
      end
    end
  end

end


text:
# SPDX-License-Identifier: X11
# 2020-10-21
# Various distances (200pt)

N = parse(Int, readline())
X = (abs ∘ parse).(Int, split(readline()))

println(sum(X))
println((sqrt ∘ sum ∘ map)(x -> x^2, X))
println(maximum(X))

max_stars_repo_name: jthvai/competitive-programming

max_issues_repo_name: jthvai/competitive-programming

max_forks_repo_name: jthvai/competitive-programming


text:
srcdir = "/Users/anray/Desktop/TDGP/src"
any(srcdir .== LOAD_PATH) || push!(LOAD_PATH, srcdir)
any(pwd() .== LOAD_PATH) || push!(LOAD_PATH, pwd())
using PyPlot, Random, Statistics, Revise, LinearAlgebra, MCMC_Driver
import GP
##
Random.seed!(13)
xtest = collect(0:0.01:1)'
λ = [0.05]
δ = 0.02
δtry = 5δ
p = 2
demean = true

y = zeros(size(xtest,2))
change2 = round(Int, 3/4*length(y))
change1 = round(Int, 1/4*length(y))
y[change2:end] .= 2
y[change1:change2-1] .= 1
ynoisy = y + δ*randn(length(y))
dec = 4
ytrain = ynoisy[1:dec:end]
xtrain = xtest[1:dec:end]'

# f1 = figure()
# plot(xtest[:][:],y, "-k",label="true")
# plot(xtrain[:],ytrain, "+m", label="train",markersize=10)
#
ytest, σ2, σ_prior  = GP.GPfit(ytrain, xtrain, xtest, λ, δtry, p=p, demean=demean)
# plot(xtest[:],ytest,label="testing")
# plot(xtest[:],ytest+2sqrt.(diag(σ2)),"--r",alpha=0.2)
# plot(xtest[:],ytest-2sqrt.(diag(σ2)),"--b",alpha=0.2)

## non stationary
λ[1] = 0.05
λtest = λ[1]*ones(size(xtest))
λtrain = λ[1]*ones(Float64, size(xtrain))
ytest_ns, σ2_ns, σ_prior_ns  = GP.GPfit(ytrain, xtrain, xtest, λtest, λtrain, δtry, p=p, demean=demean)

f2, ax2 = plt.subplots(2,1, sharex=true, figsize=(7,7))
ax2[1].plot(xtest[:],y, "-k",label="true", linewidth=2)
ax2[1].plot(xtrain[:],ytrain, "xm", label="train",markersize=10, alpha=0.8)
ax2[1].plot(xtest[:],ytest_ns,label="fixed λ", linewidth=3)
# plot(xtest[:],ytest_ns+2sqrt.(diag(σ2_ns)),"--r",alpha=0.2)
# plot(xtest[:],ytest_ns-2sqrt.(diag(σ2_ns)),"--b",alpha=0.2)
s = gca()
##
lscale_max = 0.05
dropoff = 0.0025
λtest = lscale_max*(1 .+eps() .- (exp.(-1/dropoff*abs.(xtest.-xtest[change1+1]).^2) +
                    exp.(-1/dropoff*abs.(xtest.-xtest[change2-2]).^2)))
λtrain[:] = λtest[1:dec:end]
ytest_ns, σ2_ns, σ_prior_ns  = GP.GPfit(ytrain, xtrain,
    xtest, λtest, λtrain, δtry, p=2, demean=demean, nogetvars=false)

ax2[1].plot(xtest[:],ytest_ns,label="variable λ", linewidth=3, alpha=0.7, "-r")
ax2[2].plot(xtest[:],λtest[:], linewidth=3, color="r", alpha=0.7)
ax2[1].set_ylabel(L"f(y)")
ax2[2].set_xlabel(L"y")
ax2[2].set_ylabel(L"\lambda")
ax2[1].grid()
ax2[2].grid()
MCMC_Driver.nicenup(gcf(), fsize=14)
savefig("variable.png", dpi=300)
## timing
@info "meshkernel with broadcast"
@time for i = 1:1000
    GP.meshkernel(xtrain, xtest, λtest, λtrain, p)
end
@info "mapkernel with pairwise"
@time for i = 1:1000
    GP.mapkernel(xtrain, xtest, λtest, λtrain)
end
@info "loops and full expression"
@time for i = 1:1000
    GP.makekernel(xtrain, xtest, λtest, λtrain, p)
end


text:
# # Support Vector Machine
# We are showing how to solve a support vector machine problem with COSMO (and JuMP).
#
# ## Generating the Dataset
# We want to classify the points in this example dataset with $m = 100$ samples and $n = 2$ features:

using Distributions: MvNormal
using Plots, LinearAlgebra, SparseArrays, Random, Test

#-

## Generate dataset
rng = Random.MersenneTwister(123);
num_samples = 100;
Xpos = rand(rng, MvNormal([1.5, 1.5], 1.25), div(num_samples, 2))';
Xneg = rand(rng, MvNormal([-1.5, -1.5], 1.25), div(num_samples, 2))';
ypos = ones(div(num_samples, 2));
yneg = -ones(div(num_samples, 2));

#-
## Plot dataset
plot(Xpos[:, 1], Xpos[:, 2], color = :red, st=:scatter, markershape = :rect, label = "positive", xlabel = "x1", ylabel = "x2")
plot!(Xneg[:, 1], Xneg[:, 2], color = :blue, st=:scatter, markershape = :circle, label = "negative")

# with samples $(x_1, x_2, \ldots, x_m) \in \mathbb{R}^2$ and labels $y_i \in \{-1,1\}$.
#
# ## Solving SVM as a QP
# We want to compute the weights $w$ and bias term $b$ of the (soft-margin) SVM classifier:
#
# $$
# \begin{array}{ll}
#     \text{minimize}   & \|w\|^2 + \lambda \sum_{i=1}^m \text{max}(0, 1 - y_i(w^\top x_i  - b)),
# \end{array}
# $$
# where $\lambda$ is a hyperparameter. This problem can be solved as a quadratic program.
# We can rewrite above problem into an optimisation problem in primal form by introducing the auxiliary slack variables $t_i$:
#
# $$
# t_i = \text{max}(0, 1 - y_i(w^T x_i  - b)), \quad t_i \geq 0.
# $$

# This allows us to write the problems in standard QP format:
# $$
# \begin{array}{ll}
#     \text{minimize}   & \|w\|^2 + \lambda \sum_{i=1}^m t_i\\
#     \text{subject to} & y_i (w^\top x_i - b) \geq 1 - t_i, \quad \text{for } i = 1,\ldots, m\\
#                       & t_i \geq 0, \quad \text{for } i = 1,\ldots, m.
# \end{array}
# $$
# Next, we will remove the bias term $b$ by adding an initial feature $x_0 = -1$ to each sample (now: $n = 3$):

X = [-ones(num_samples) [Xpos; Xneg]];
y = [ypos; yneg];
m, n = size(X)

# ## Modelling in JuMP
# We can model this problem using `JuMP` and then hand it to `COSMO`:
using JuMP, COSMO
#-
λ = 1.0; # hyperparameter
model = JuMP.Model(optimizer_with_attributes(COSMO.Optimizer, "verbose" => true));

@variable(model, w[1:n]);
@variable(model, t[1:m] >= 0.);
@objective(model, Min, w' * w  + λ * ones(m)' * t);
@constraint(model, diagm(0 => y) * X * w .+ t .- 1 .>= 0);
status = JuMP.optimize!(model)
# The optimal weights $w = [w_0, w_1, w_2]^\top$ (where $w_0 = b$) are:
w_opt = JuMP.value.(w)

# ## Plotting the hyperplane
# The separating hyperplane is defined by $w^\top x - b = 0$. To plot the hyperplane, we calculate $x_2$ over a range of $x_1$ values:
# $$
# x_2 = (-w_1 x_1 - w_0) / w_2, \text{ where } w_0 = b.
# $$
x1 = -4:0.1:4;
x2 = (-w_opt[2] * x1  .- w_opt[1]) / w_opt[3]
plot!(x1, x2, label = "SVM separator", legend = :topleft)

# ## Modelling with COSMO
# The problem can also be solved by transforming it directly into `COSMO`'s problem format.
# Define `COSMO``s $x$-variable to be $x=[w, t]^\top$ and choose $P$, $q$, accordingly:
P = blockdiag(spdiagm(0 => ones(n)), spzeros(m, m));
q = [zeros(n); 0.5 * λ * ones(m)];

# Next we transform the first constraint $y_i (w^\top x_i - b) \geq 1 - t_i, \quad \text{for } i = 1,\ldots, m$ into
# `COSMO`'s constraint format: $Ax + b \in \mathcal{K}$.
A1 = [(spdiagm(0 => y) * X) spdiagm(0 => ones(m))];
b1 = -ones(m);
cs1 = COSMO.Constraint(A1, b1, COSMO.Nonnegatives);

# It remains to specify the constraint $t_i \geq 0, \quad \text{for } i = 1,\ldots, m$:
A2 = spdiagm(0 => ones(m));
b2 = zeros(m);
cs2 = COSMO.Constraint(A2, b2, COSMO.Nonnegatives, m+n, n+1:m+n);

# Create, assemble and solve the `COSMO.Model`:
model2 = COSMO.Model();
assemble!(model2, P, q, [cs1; cs2]);
result2 = COSMO.optimize!(model2);
w_opt2 = result2.x[1:3];
@test norm(w_opt2 - w_opt, Inf) < 1e-3


text:
if length(workers()) == 1 #src
    addprocs(4) #src
end #src
default(grid=false) #src
using ParallelDataTransfer #src
@everywhere using Parameters, JacProp, LTVModels, LTVModelsBase #src
@everywhere using Flux: params, jacobian #src
@everywhere using Flux.Optimise: Param, optimiser, expdecay #src
@everywhere begin #src

@with_kw struct LinearSys <: AbstractSystem
    A
    B
    N     = 1000
    nx    = size(A,1)
    nu    = size(B,2)
    h     = 0.02
    σ0    = 0
    sind  = 1:nx
    uind  = nx+1:(nx+nu)
    s1ind = (nx+nu+1):(nx+nu+nx)
end

function LinearSys(seed; nx = 10, nu = nx, h=0.02, kwargs...)
    srand(seed)
    A = randn(nx,nx)
    A = A-A'        # skew-symmetric = pure imaginary eigenvalues
    A = A - h*I     # Make 'slightly' stable
    A = expm(h*A)   # discrete time
    B = h*randn(nx,nu)
    LinearSys(;A=A, B=B, nx=nx, nu=nu, h=h, kwargs...)
end

function generate_data(sys::LinearSys, seed, validation=false)
    Parameters.@unpack A,B,N, nx, nu, h, σ0 = sys
    srand(seed)
    u      = filt(ones(5),[5], 10randn(N+2,nu))'
    t      = h:h:N*h+h
    x0     = randn(nx)
    x      = zeros(nx,N+1)
    x[:,1] = x0
    for i = 1:N-1
        x[:,i+1] = A*x[:,i] + B*u[:,i]
    end

    validation || (x .+= σ0 * randn(size(x)))
    u = u[:,1:N]
    @assert all(isfinite, u)
    x,u
end

function true_jacobian(sys::LinearSys, x, u)
    [sys.A sys.B]
end

function callbacker(epoch, loss,d,trace,model)
    printed = false
    info("New callbacker, epoch $epoch")
    function ()
        info("callback, epoch $epoch")
        l = sum(d->Flux.data(loss(d...)),d)
        increment!(trace,epoch,l)
        if !printed && epoch % 500 == 0
            println(@sprintf("Loss: %.4f", l))
            printed = true
        end
    end
end

num_params = 30
wdecay     = 0
stepsize   = 0.02
const sys  = LinearSys(1, N=200, h=0.02, σ0 = 0.01)
true_jacobian(x,u) = true_jacobian(sys,x,u)
nu         = sys.nu
nx         = sys.nx

end #src

# Generate validation data
function valdata()
    vx,vu,vy = Vector{Float64}[],Vector{Float64}[],Vector{Float64}[]
    for i = 20:60
        x,u = generate_data(sys,i, true)
        for j in 10:5:(sys.N-1)
            push!(vx, x[:,j])
            push!(vy, x[:,j+1])
            push!(vu, u[:,j])
        end
    end
    hcat(vx...),hcat(vu...),hcat(vy...)
end
vx,vu,vy = valdata()
vt = Trajectory(vx,vu,vy)

# Generate training trajectories
trajs = [Trajectory(generate_data(sys, i)...) for i = 1:3]

sendto(collect(2:5), trajs=trajs, vt=vt) #src

# # Without jacprop
f2 = @spawnat 2 begin #src
    srand(1)
    models     = [DiffSystem(nx,nu,num_params, a) for a in default_activations]
    opts       = ADAM.(params.(models), stepsize, decay=0.0005)#; [expdecay(Param(p), wdecay) for p in params(models[i]) if p isa AbstractMatrix]] for i = 1:length(models)]
    trainer  = ModelTrainer(models = models, opts = opts, losses = JacProp.loss.(models), cb=callbacker, P = 10, R2 = 10000I, σdivider = 20)
    for i = 1:3
        trainer(trajs[i], epochs=2000, jacprop=0, useprior=false)
    end
    trainer
end #src

# # With jacprop and prior
f3 = @spawnat 3 begin #src
    srand(1)
    models     = [DiffSystem(nx,nu,num_params, a) for a in default_activations]
    opts       = ADAM.(params.(models), stepsize, decay=0.0005)#; [expdecay(Param(p), wdecay) for p in params(models[i]) if p isa AbstractMatrix]] for i = 1:length(models)]
    trainerj  = ModelTrainer(models = models, opts = opts, losses = JacProp.loss.(models), cb=callbacker, P = 10, R2 = 1000I, σdivider = 20)
    for i = 1:3
        trainerj(trajs[i], epochs=1000, jacprop=1, useprior=true)
    end
    trainerj
end #src

# # With jacprop no prior
f4 = @spawnat 4 begin #src
    srand(1)
    models     = [DiffSystem(nx,nu,num_params, a) for a in default_activations]
    opts       = ADAM.(params.(models), stepsize, decay=0.0005)#; [expdecay(Param(p), wdecay) for p in params(models[i]) if p isa AbstractMatrix]] for i = 1:length(models)]
    trainerjn  = ModelTrainer(models = models, opts = opts, losses = JacProp.loss.(models), cb=callbacker, P = 10, R2 = 1000I, σdivider = 200)
    for i = 1:3
        trainerjn(trajs[i], epochs=1000, jacprop=1, useprior=false)
    end
    trainerjn
end #src

# # Visualize result
pyplot(reuse=false)
trainer,trainerj,trainerjn = fetch(f2), fetch(f3),fetch(f4)

## mutregplot(trainer, vt, true_jacobian, title="Witout jacprop", subplot=1, layout=(2,3), reuse=false, useprior=false, showltv=false)
## mutregplot!(trainerj, vt, true_jacobian, title="With jacprop and prior", subplot=2, link=:y, useprior=true, show=false, showltv=false)
## mutregplot!(trainerjn, vt, true_jacobian, title="With jacprop, no prior", subplot=3, link=:y, useprior=false, showltv=false)
## traceplot!(trainer, subplot=3)
## traceplot!(trainerj, subplot=5, show=false)
## traceplot!(trainerjn, subplot=4)

mutregplot(trainer, vt, true_jacobian, title="Witout jacprop", subplot=1, layout=(2,2), reuse=false, useprior=false, showltv=false, legend=false, xaxis=(1:3), xaxis=(1:3))
mutregplot!(trainerjn, vt, true_jacobian, title="With jacprop", subplot=2, link=:y, useprior=false, showltv=false, legend=false, xaxis=(1:3))
traceplot!(trainer, subplot=3, title="Training error", xlabel="Epoch", legend=false)
traceplot!(trainerjn, subplot=4, title="Training error", xlabel="Epoch", legend=false)

gui()
##

jacplot(trainer.models, vt, true_jacobian, label="Without", c=:red, reuse=false, fillalpha=0.2)
## jacplot!(KalmanModel(trainer, vt), vt, label="Without", c=:pink)
jacplot!(trainerj.models, vt, label="With", c=:blue, fillalpha=0.2)
## jacplot!(KalmanModel(trainerj, vt), vt, label="With", c=:cyan)
jacplot!(trainerjn.models, vt, label="With no", c=:green, fillalpha=0.2)
## jacplot!(KalmanModel(trainerjn, vt), vt, label="With no", c=:green)
gui()

## eigvalplot(trainer.models, vt, subplot=1, layout=3, title="g, Witout jacprop, wdecay=$wdecay")
## eigvalplot!(trainerj.models, vt, subplot=2, title="g, With jacprop and prior, wdecay=$wdecay")
## eigvalplot!(trainerjn.models, vt, subplot=3, title="g, With jacprop, no prior, wdecay=$wdecay")

eigvalplot(trainer.models, vt, true_jacobian, subplot=1, layout=2, title="g, Witout jacprop", markeralpha=0.05, link=:both, xlims=(-0.2, 1.2), ds=1, markerstrokealpha=0)
eigvalplot!(trainerjn.models, vt, true_jacobian, subplot=2, title="g, With jacprop", markeralpha=0.05, link=:both, xlims=(-0.2, 1.2), ds=1, markerstrokealpha=0)
plot!(ylims=(-0.8, 0.8))
gui()

ui = display_modeltrainer(trainerjn, size=(800,600))
jacplot(trainer.models, trainer.trajs[3], true_jacobian, ds=20)
@gif for i = 1:length(t)
    eigvalplot(trainer.models, trainer.trajs[1], true_jacobian, ds=20, onlyat=i)
end

## TODO: produce eigvalplots for system and diffsystem with and without wdecay
## TODO: make jacprop magnitude an option std()/10
## TODO: Validate vt in callback
## TODO: why can a network perform better but train with higher final loss? Overfitting?

wdecay = 0.0
f1 = @spawnat 1 begin #src
    srand(1)
    models     = [System(nx,nu,num_params, a) for a in default_activations]
    opts       = ADAM.(params.(models), stepsize, decay=0.0005)#; [expdecay(Param(p), wdecay) for p in params(models[i]) if p isa AbstractMatrix]] for i = 1:length(models)]
    trainers  = ModelTrainer(models = models, opts = opts, losses = JacProp.loss.(models), cb=callbacker, P = 10, R2 = 10000I, σdivider = 20)
    for i = 1:3
        trainers(trajs[i], epochs=2000, jacprop=1, useprior=false)
    end
    trainers
end #src

f2 = @spawnat 2 begin #src
    srand(1)
    models     = [DiffSystem(nx,nu,num_params, a) for a in default_activations]
    opts       = ADAM.(params.(models), stepsize, decay=0.0005)#; [expdecay(Param(p), wdecay) for p in params(models[i]) if p isa AbstractMatrix]] for i = 1:length(models)]
    trainerds  = ModelTrainer(models = models, opts = opts, losses = JacProp.loss.(models), cb=callbacker, P = 10, R2 = 10000I, σdivider = 20)
    for i = 1:3
        trainerds(trajs[i], epochs=2000, jacprop=1, useprior=false)
    end
    trainerds
end #src

wdecay = 0.1
f3 = @spawnat 3 begin #src
    srand(1)
    models     = [System(nx,nu,num_params, a) for a in default_activations]
    opts       = [[ADAM(params(models[i]), stepsize, decay=0.0005); [expdecay(Param(p), wdecay) for p in params(models[i]) if p isa AbstractMatrix]] for i = 1:length(models)]
    trainers  = ModelTrainer(models = models, opts = opts, losses = JacProp.loss.(models), cb=callbacker, P = 10, R2 = 10000I, σdivider = 20)
    for i = 1:3
        trainers(trajs[i], epochs=2000, jacprop=1, useprior=false)
    end
    trainers
end #src

f4 = @spawnat 4 begin #src
    srand(1)
    models     = [DiffSystem(nx,nu,num_params, a) for a in default_activations]
    opts       = [[ADAM(params(models[i]), stepsize, decay=0.0005); [expdecay(Param(p), wdecay) for p in params(models[i]) if p isa AbstractMatrix]] for i = 1:length(models)]
    trainerdswd  = ModelTrainer(models = models, opts = opts, losses = JacProp.loss.(models), cb=callbacker, P = 10, R2 = 10000I, σdivider = 20)
    for i = 1:3
        trainerdswd(trajs[i], epochs=2000, jacprop=1, useprior=false)
    end
    trainerdswd
end #src

trainers,trainerds = fetch(f1), fetch(f2)
trainerswd,trainerdswd = fetch(f3), fetch(f4)
eigvalplot(trainers.models, vt, true_jacobian, title="f, wdecay=0", layout=4, subplot=1, markeralpha=0.05, ds=1, markerstrokealpha=0)
eigvalplot!(trainerds.models, vt, true_jacobian,  title="g, wdecay=0", subplot=2, markeralpha=0.05, ds=1, markerstrokealpha=0)
eigvalplot!(trainerswd.models, vt, true_jacobian,  title="f, wdecay=$wdecay", subplot=3, markeralpha=0.05, ds=1, markerstrokealpha=0)
eigvalplot!(trainerdswd.models, vt, true_jacobian,  title="g, wdecay=$wdecay", subplot=4, markeralpha=0.05, ds=1, markerstrokealpha=0)
plot!(link=:both, xlims=(0,1.1), ylims=(-0.5, 0.5))
gui()

ui = display_modeltrainer(trainerdswd, size=(800,600))

for acti in 1:4
    plot(layout=4, link=:both)
    actstr = string(JacProp.default_activations[acti])
    actstr = actstr[1:2] == "NN" ? actstr[7:end] : actstr
    for tri in 1:4
        tr = [trainers, trainerds, trainerswd, trainerdswd][tri]
        eigvalplot!([tr.models[acti]], vt, true_jacobian, title="$actstr, $(tri%2==0 ? "g" : "f"), wdecay=$(tri>2 ? wdecay : 0)", subplot=tri, markeralpha=0.05, ds=3, markerstrokealpha=0)
    end
    gui()
end


text:
include("phi_ru.jl")
include("phi_ft.jl")

include("phi_cs.jl")
include("phi_pf.jl")

include("phi_ind.jl")
include("phi_res.jl")


text:
##### Sum wrapper

export Σfluxes, Σsources

"""
    flux

An individual flux.
See [`BalanceLaw`](@ref) for more info.
"""
function flux end

"""
    source

An individual source.
See [`BalanceLaw`](@ref) for more info.
"""
function source end

"""
    ntuple_sum(nt::NTuple{N,T}) where {N, T}

sum of `NTuple`, which requires more strict
type input than `sum`. This is added to better
synchronize the success/failure between CPU/GPU
runs to help improve debugging.
"""
ntuple_sum(nt::NTuple{N, T}) where {N, T} = sum(nt)

"""
    Σfluxes(fluxes::NTuple, args...)

Sum of the fluxes where
 - `fluxes` is an `NTuple{N, TendencyDef{Flux{O}, PV}} where {N, PV, O}`
"""
function Σfluxes(
    fluxes::NTuple{N, TendencyDef{Flux{O}, PV}},
    args...,
) where {N, PV, O}
    return ntuple_sum(ntuple(Val(N)) do i
        flux(fluxes[i], args...)
    end)
end
Σfluxes(fluxes::Tuple{}, args...) = 0

"""
    Σsources(sources::NTuple, args...)

Sum of the sources where
 - `sources` is an `NTuple{N, TendencyDef{Source, PV}} where {N, PV}`
"""
function Σsources(
    sources::NTuple{N, TendencyDef{Source, PV}},
    args...,
) where {N, PV}
    return ntuple_sum(ntuple(Val(N)) do i
        source(sources[i], args...)
    end)
end
Σsources(sources::Tuple{}, args...) = 0

max_forks_repo_name: jm-c/CLIMA


text:
using ScatteredArrays
using Test

@testset "ScatteredArrays.jl" begin
    # Write your own tests here.
end


text:
# Julia implementation of the multirate integration scheme used in the numerical experiments of the chapter
# Leyendecker S., Ober-Blöbaum S. (2013)
# A Variational Approach to Multirate Integration for Constrained Systems.
# In: Samin JC., Fisette P. (eds) Multibody Dynamics. Computational Methods in Applied Sciences, vol 28. Springer, Dordrecht.
# https://doi.org/10.1007/978-94-007-5404-1_5

using Flux: gradient
import Base.Iterators: flatten
using NLsolve
#using ProgressBars

function BuildDiscreteL(fastPotential, slowPotential, kineticTerm)
    # returns discrete multi-rate Lagrangian given continuous functions for the fast potential, the slow potential, the kinetic term

    function qsInterpolate(qs0,qs1,m,h,microsteps)
    1/microsteps*((microsteps-m)*qs0 + m*qs1)
    end

    function DiscreteSlowPotential(qs0,qs1,qf0,h,microsteps)                             # qf0 is matrix with dimension (microsteps,length(qf))
    	h_micro = h/microsteps  # length micro step
        h_micro*sum(map(m->slowPotential(qsInterpolate(qs0,qs1,m-1,h,microsteps),qf0[m,:]),1:microsteps))
    end

    function DiscreteFastPotential(qf0,h,microsteps)
        qfmp=(qf0[1:end-1,:]+qf0[2:end,:])/2

        # rewriting mapslices for Zygote
        s = 0.
        for j=1:microsteps
            s += fastPotential(qfmp[j,:])
        end

	h_micro = h/microsteps  # length micro step

        return h_micro*s

    end

    function DiscreteKineticTerm(qs0,qs1,qf0,h,microsteps)

        h_micro = h/microsteps  # length micro step

        qsdot = (qs1-qs0)/h
        qfdot = (qf0[2:end,:]-qf0[1:end-1,:])/h_micro

        #rewriting mapslices for Zygote
        s = 0.
        for j=1:microsteps
            s += kineticTerm(qsdot,qfdot[j,:])
        end

        return h_micro*s

    end

    function Ld(qs0,qs1,qf0,h,microsteps)
        DiscreteKineticTerm(qs0,qs1,qf0,h,microsteps) - DiscreteSlowPotential(qs0,qs1,qf0,h,microsteps) - DiscreteFastPotential(qf0,h,microsteps)
    end

    return Ld
end

function MultiRateIntegrator(Ld)
    # returns the function MultiRateIntegrate as well as InitializationValues, and MultiRateStep for a discrete multirate Lagrangian Ld.

    function InitializationValues(qs0,qf00,ps0,pf00,h,microsteps) # (qs0,qf00),(ps0,pf00) initial configuration
	# compute the initial step from initial configuration

        n = length(qf00)

        function InitHelper(qs1,qf0Tail)    # qf0Tail is matrix with dimension (microsteps,length(qf)) (leave out microstep at t_k^0 but include t_k^microstep)
	    # need to solve InitHelper(qs1,qf0tail)==0 to determine initialisation values from initial configuration
            qf0Inner = qf0Tail[1:end-1,:]
            qf0End   = qf0Tail[end,:]

            partial_Ld(qs0,qf00,qf0Inner) = Ld(qs0,qs1,vcat(transpose(qf00),qf0Inner,transpose(qf0End)),h,microsteps)

            dLd0 = gradient( partial_Ld, qs0,qf00,qf0Inner)
            dLd0[1] .+= ps0
            dLd0[2] .+= pf00

            return collect(flatten(dLd0))

        end

        function InitHelper(A)
            InitHelper(A[1:n],reshape(A[n+1:end],(microsteps,n)))
        end

        initsSol = nlsolve(InitHelper, zeros((microsteps+1)*n),autodiff = :forward)

        inits = initsSol.zero

        # distribute output to qs1 and inner microstep points qf0
        qs1 = inits[1:n]
        qf0Tail = reshape(inits[n+1:end],(microsteps,n))

        qf0=vcat(transpose(qf00),qf0Tail);

        return qs1, qf0,initsSol

    end

    function MultiRateStep(qs0,qs1,qf0,h,microsteps)
	# computes a step with the multi-rate integrator

        n = length(qs0)

        qf10 = qf0[end,:]

        function objective(qs2,qf1Tail)

            qf1Inner = qf1Tail[1:end-1,:]
            qf1End = qf1Tail[end,:]

            # attention: qf10 occurs in qf0 as last element and in qf1 as first element
            Ld_partial(qs1,qf10,qf1Inner) = Ld(qs1,qs2,vcat(transpose(qf10),qf1Inner,transpose(qf1End)),h,microsteps) + Ld(qs0,qs1,vcat(qf0[1:end-1,:],transpose(qf10)),h,microsteps)

            return gradient(Ld_partial, qs1,qf10,qf1Inner)
        end

        objective(A) = collect(flatten(objective(A[1:n],reshape(A[n+1:end],(microsteps,n)))))

        Aguess = vcat(qs1, reshape(qf0[2:end,:],microsteps*n))

        qnextSol = nlsolve(objective,Aguess,autodiff = :forward)

        # distribute to variables
        qs2     = qnextSol.zero[1:n]
        qf1Tail = reshape(qnextSol.zero[n+1:end],(microsteps,n))
        qf1 = vcat(transpose(qf0[end,:]),qf1Tail)

        return qs2,qf1,qnextSol

    end

    function MultiRateIntegrate(qs0,qf00,ps0,pf00,N,h,microsteps)

    	# returns slow and fast variables for N steps with the multi-rate integrator
    	# qs0,qf00,ps0,pf00 initial configuration
    	# N number of steps
    	# h lenggth macro step
    	# microsteps number of microsteps

        # compute 1st macro step
        qs1,qf0,initsSol=InitializationValues(qs0,qf00,ps0,pf00,h,microsteps)

        # pre-allocate memory
        QS = Array{Float64}(undef,(length(qs0),N+1))
        QF = Array{Float64}(undef,(length(qs0),N*microsteps+1))
        checkSolver = Array{Bool, 1}(undef, N-1)

        QS[:,1] = qs0
        QS[:,2] = qs1
        QF[:,1:microsteps+1] = transpose(qf0);

        for j = 2:N #ProgressBar(2:N)

            # computation of q_j

            qsk0 = QS[:,j-1]
            qsk1 = QS[:,j]
            qfk0 = QF[:,(j-2)*microsteps+1:(j-1)*microsteps+1]

            qsk2,qfk1, chkSolver = MultiRateStep(qsk0,qsk1,transpose(qfk0),h,microsteps)

            QS[:,j+1] = qsk2
            QF[:,(j-1)*microsteps+1:j*microsteps+1] = transpose(qfk1)
            checkSolver[j-1] = chkSolver.f_converged

        end

        return QS,QF,checkSolver

    end

    return MultiRateIntegrate, InitializationValues, MultiRateStep

end

function DiscreteConjugateMomenta(Ld,QS,QF,h,microsteps)

    # compute discrete conjugate momenta on macro-grid
    # Ld discrete Lagrangian
    # QS slow variables, Matrix of dimension (dimension of qs variable,number of macro-steps)
    # QF fast variables, Matrix of dimension (dimension of qf variable,number of micro-steps)

    dLd_qs0(qs0,qs1,qf0)  = gradient(qs0 -> Ld(qs0,qs1,qf0,h,microsteps),qs0)[1]
    dLd_qs1(qs0,qs1,qf0)  = gradient(qs1 -> Ld(qs0,qs1,qf0,h,microsteps),qs1)[1]
    dLd_qf00(qs0,qs1,qf0) = gradient(qf00 -> Ld(qs0,qs1,vcat(transpose(qf00),qf0[2:end,:]),h,microsteps),qf0[1,:])[1]
    dLd_qf0end(qs0,qs1,qf0) = gradient(qf0end -> Ld(qs0,qs1,vcat(qf0[1:end-1,:],transpose(qf0end)),h,microsteps),qf0[end,:])[1]

    microsteps = convert(Int64,(size(QF)[2]-1)/(size(QS)[2]-1))

    # pre-allocate memory
    PS = Array{Float64}(undef,size(QS))
    PF = Array{Float64}(undef,size(QS))

    # compute discrete conjugate momenta
    for j=1:(size(QS)[2]-1)

        qfj = transpose(QF[:,(j-1)*microsteps+1:j*microsteps+1])

        PS[:,j]   = -dLd_qs0(QS[:,j],QS[:,j+1],qfj)
        PS[:,j+1] = dLd_qs1(QS[:,j],QS[:,j+1],qfj)                    # only very last component will not get overwritten, excellent jit-compiler will notice
        PF[:,j]   = -transpose(dLd_qf00(QS[:,j],QS[:,j+1],qfj))
        PF[:,j+1] = transpose(dLd_qf0end(QS[:,j],QS[:,j+1],qfj))      # only very last component will not get overwritten, excellent jit-compiler will notice

    end

    return PS,PF

end

max_stars_repo_name: Christian-Offen/MultiRate

max_issues_repo_name: Christian-Offen/MultiRate

max_forks_repo_name: Christian-Offen/MultiRate


text:
### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# ╔═╡ e077ab16-22e5-11eb-15ed-5f0654bf7328
using Pkg, DrWatson

# ╔═╡ e75a9a06-22e5-11eb-2e36-8d4f62b830ed
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Statistics, DataFrames, CSV
	using Random, Distributions, Plots
	using HypothesisTests
	Random.seed!(0)
end;

# ╔═╡ ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
md"## Listing6.02"

# ╔═╡ 1aa41898-3afa-11eb-2f49-570853ce6435
begin
	data = CSV.read(datadir("machine1.csv"), DataFrame; header=false)[:,1]
	xBar, n = mean(data), length(data)
	s = std(data)
	alpha = 0.1
	t = quantile(TDist(n-1),1-alpha/2)

	Text("Calculating formula: $((xBar - t*s/sqrt(n), xBar + t*s/sqrt(n)))")
end

# ╔═╡ 88f0e5d6-3b0b-11eb-0af9-b964a1170ff3
Text("Using confint() function: $(confint(OneSampleTTest(xBar,s,n),alpha))")

# ╔═╡ 475ff888-22e6-11eb-2354-09f8f40a8e12
md"## End of listing6.02"

# ╔═╡ Cell order:
# ╟─ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
# ╠═e077ab16-22e5-11eb-15ed-5f0654bf7328
# ╠═e75a9a06-22e5-11eb-2e36-8d4f62b830ed
# ╠═1aa41898-3afa-11eb-2f49-570853ce6435
# ╠═88f0e5d6-3b0b-11eb-0af9-b964a1170ff3
# ╟─475ff888-22e6-11eb-2354-09f8f40a8e12


text:
## Higher criticism

"""
Higher criticism scores

# Examples

```jldoctest
julia> pvals = PValues([0.001, 0.01, 0.03, 0.5]);

julia> estimate(pvals, HigherCriticismScores())
4-element Array{Float64,1}:
 1.15008
 1.96
 3.32554
 2.3094

```

# References

Donoho, D., and Jin, J. (2008). Higher criticism thresholding: Optimal feature
selection when useful features are rare and weak. PNAS 105, 14790–14795.

Klaus, B., and Strimmer, K. (2013). Signal identification for rare and weak
features: higher criticism or false discovery rates? Biostatistics 14, 129–143.

"""
struct HigherCriticismScores
end

function estimate(pValues::PValues{T}, method::HigherCriticismScores) where T<:AbstractFloat
    n = length(pValues)
    F = (n+1 .- competerank(-pValues)) ./ n  # ECDF
    denom = F .* (one(T) .- F) ./ n
    # avoid denominator of 0 for last value
    idx0 = denom .== 0
    denom[idx0] .= minimum(denom[.!idx0]) .+ eps()  # conservative
    hcs = abs.(F .- pValues) ./ sqrt.(denom)
    return hcs
end

"""
Higher criticism threshold

# Examples

```jldoctest
julia> pvals = PValues([0.001, 0.01, 0.03, 0.5]);

julia> estimate(pvals, HigherCriticismThreshold())
0.03
```

# References

Donoho, D., and Jin, J. (2008). Higher criticism thresholding: Optimal feature
selection when useful features are rare and weak. PNAS 105, 14790–14795.

Klaus, B., and Strimmer, K. (2013). Signal identification for rare and weak
features: higher criticism or false discovery rates? Biostatistics 14, 129–143.

"""
struct HigherCriticismThreshold
end

function estimate(pValues::PValues{T}, method::HigherCriticismThreshold) where T<:AbstractFloat
    idx_hcv = argmax(estimate(pValues, HigherCriticismScores()))
    return pValues[idx_hcv]
end


text:
# This Map has non-trivial domain, thus we need the define testargs
"""
    Reindex(values) -> Map
"""
struct Reindex{A} <: Map
  values::A
end

function testargs(k::Reindex,i)
  @check length(k.values) !=0 "This map has empty domain"
  (one(i),)
end
function testargs(k::Reindex,i::Integer...)
  @check length(k.values) !=0 "This map has empty domain"
  map(one,i)
end
function return_value(k::Reindex,i...)
  length(k.values)!=0 ? evaluate(k,testargs(k,i...)...) : testitem(k.values)
end
return_cache(k::Reindex,i...) = array_cache(k.values)
evaluate!(cache,k::Reindex,i...) = getindex!(cache,k.values,i...)
evaluate(k::Reindex,i...) = k.values[i...]

#"""
#    reindex(i_to_v::AbstractArray, j_to_i::AbstractArray)
#"""
#function reindex(i_to_v::AbstractArray, j_to_i::AbstractArray)
#  lazy_map(Reindex(i_to_v),j_to_i)
#end

function lazy_map(k::Reindex{<:Fill},::Type{T}, j_to_i::AbstractArray) where T
  v = k.values.value
  Fill(v,size(j_to_i)...)
end

function lazy_map(k::Reindex{<:CompressedArray},::Type{T}, j_to_i::AbstractArray) where T
  i_to_v = k.values
  values = i_to_v.values
  ptrs = lazy_map(Reindex(i_to_v.ptrs),j_to_i)
  CompressedArray(values,ptrs)
end

function lazy_map(k::Reindex{<:LazyArray},::Type{T},j_to_i::AbstractArray) where T
  i_to_maps = k.values.maps
  i_to_args = k.values.args
  j_to_maps = lazy_map(Reindex(i_to_maps),eltype(i_to_maps),j_to_i)
  j_to_args = map(i_to_fk->lazy_map(Reindex(i_to_fk),eltype(i_to_fk),j_to_i), i_to_args)
  LazyArray(T,j_to_maps,j_to_args...)
end

# This optimization is important for surface-coupled problems
function lazy_map(k::Reindex{<:LazyArray{<:Fill{<:PosNegReindex}}},::Type{T},j_to_i::AbstractArray) where T
  i_to_iposneg = k.values.args[1]
  ipos_to_value = k.values.maps.value.values_pos
  ineg_to_value = k.values.maps.value.values_neg
  if aligned_with_pos(i_to_iposneg,j_to_i,length(ipos_to_value))
    ipos_to_value
  elseif aligned_with_neg(i_to_iposneg,j_to_i,length(ineg_to_value))
    ineg_to_value
  elseif all_in_pos(i_to_iposneg,j_to_i)
    j_to_ipos = lazy_map(Reindex(get_array(i_to_iposneg)),j_to_i)
    j_to_value = lazy_map(Reindex(ipos_to_value),j_to_ipos)
  elseif all_in_neg(i_to_iposneg,j_to_i)
    j_to_ineg = lazy_map(Reindex(get_array(i_to_iposneg)),j_to_i)
    j_to_value = lazy_map(Reindex(ineg_to_value),lazy_map(ineg->-ineg,j_to_ineg))
  else
    j_to_iposneg = lazy_map(Reindex(get_array(i_to_iposneg)),j_to_i)
    j_to_value = lazy_map(PosNegReindex(ipos_to_value,ineg_to_value),j_to_iposneg)
  end
end

function lazy_map(k::Reindex{<:LazyArray{<:Fill{<:PosNegReindex}}},::Type{T},indices::IdentityVector) where T
  @check length(k.values) == length(indices)
  k.values
end

function lazy_map(k::Reindex{<:LazyArray},::Type{T},indices::IdentityVector) where T
  @check length(k.values) == length(indices)
  k.values
end

function lazy_map(k::Reindex{<:AbstractArray},::Type{T},indices::IdentityVector) where T
  @check length(k.values) == length(indices)
  k.values
end

function lazy_map(k::Reindex{<:Fill},::Type{T},indices::IdentityVector) where T
  @check length(k.values) == length(indices)
  k.values
end

function lazy_map(k::Reindex{<:CompressedArray},::Type{T},indices::IdentityVector) where T
  @check length(k.values) == length(indices)
  k.values
end

@propagate_inbounds function Base.setindex!(a::LazyArray{<:Fill{<:Reindex}},v,j::Integer)
  k = a.maps.value
  i_to_v = k.values
  j_to_i, = a.args
  i = j_to_i[j]
  i_to_v[i]=v
end


text:
# !! auto-generated code, do not edit !!
module stereo_msgs
using RobotOSData.Messages
using RobotOSData.StdMsgs
using RobotOSData.CommonMsgs
struct DisparityImage <: Readable
    header::Header
    image::sensor_msgs.Image
    f::Float32
    T::Float32
    valid_window::sensor_msgs.RegionOfInterest
    min_disparity::Float32
    max_disparity::Float32
    delta_d::Float32
end
end


text:
using Test
using LinearAlgebra
using QuantumControlBase.Functionals
using QuantumControlBase.Functionals:
    chi_re!, chi_sm!, chi_ss!, make_zygote_gradient, make_zygote_chi, grad_J_T_sm!
using QuantumControlBase.TestUtils

N_HILBERT = 10
N = 4
L = 2
N_T = 50
PROBLEM = dummy_control_problem(; N=N_HILBERT, n_objectives=N, n_controls=L, n_steps=N_T)

@testset "zygote-gradients" begin

    # Test that Zygote gradients and analytical gradients give the same result

    objectives = PROBLEM.objectives
    G1 = zeros(L * N_T)
    G2 = zeros(L * N_T)
    τ = rand(N) .* exp.((1im * 2π) .* rand(N))
    ∇τ = [rand(L, N_T) .* exp.((1im * 2π) .* rand(L, N_T)) for k = 1:N]
    grad_J_T_sm_zyg! = make_gradient(J_T_sm, objectives; via=:tau, force_zygote=true)
    grad_J_T_sm!(G1, τ, ∇τ)
    grad_J_T_sm_zyg!(G2, τ, ∇τ)
    @test norm(G1 - G2) < 1e-15

end

@testset "chi-with-tau" begin

    # Test that the various chi routines give the same result whether they are
    # called with ϕ states or with τ values

    objectives = PROBLEM.objectives
    χ1 = [similar(obj.initial_state) for obj in objectives]
    χ2 = [similar(obj.initial_state) for obj in objectives]
    ϕ = [random_state_vector(N_HILBERT) for k = 1:N]
    τ = [obj.target_state ⋅ ϕ[k] for (k, obj) in enumerate(objectives)]

    chi_re!(χ1, ϕ, objectives)
    chi_re!(χ2, ϕ, objectives; τ=τ)
    @test maximum(norm.(χ1 .- χ2)) < 1e-15

    chi_sm!(χ1, ϕ, objectives)
    chi_sm!(χ2, ϕ, objectives; τ=τ)
    @test maximum(norm.(χ1 .- χ2)) < 1e-15

    chi_ss!(χ1, ϕ, objectives)
    chi_ss!(χ2, ϕ, objectives; τ=τ)
    @test maximum(norm.(χ1 .- χ2)) < 1e-15

end

@testset "zygote-chi" begin

    # Test that the zygote chi routines give the same results as the analytical
    # chi routines

    objectives = PROBLEM.objectives
    χ1 = [similar(obj.initial_state) for obj in objectives]
    χ2 = [similar(obj.initial_state) for obj in objectives]
    ϕ = [random_state_vector(N_HILBERT) for k = 1:N]
    τ = [obj.target_state ⋅ ϕ[k] for (k, obj) in enumerate(objectives)]

    chi_re!(χ1, ϕ, objectives)
    chi_re_zyg! = make_chi(J_T_re, objectives; force_zygote=true)
    chi_re_zyg!(χ2, ϕ, objectives; τ=τ)
    @test maximum(norm.(χ1 .- χ2)) < 1e-15

    chi_sm!(χ1, ϕ, objectives)
    chi_sm_zyg! = make_zygote_chi(J_T_sm, objectives)
    chi_sm_zyg!(χ2, ϕ, objectives; τ=τ)
    @test maximum(norm.(χ1 .- χ2)) < 1e-15

    chi_ss!(χ1, ϕ, objectives)
    chi_ss_zyg! = make_zygote_chi(J_T_ss, objectives)
    chi_ss_zyg!(χ2, ϕ, objectives; τ=τ)
    @test maximum(norm.(χ1 .- χ2)) < 1e-15

end

@testset "make-gradient" begin

    # Test that the routine returned by `make_gradient` gives the same result
    # as the Zygote gradient

    objectives = PROBLEM.objectives
    G1 = zeros(L * N_T)
    G2 = zeros(L * N_T)
    τ = rand(N) .* exp.((1im * 2π) .* rand(N))
    ∇τ = [rand(L, N_T) .* exp.((1im * 2π) .* rand(L, N_T)) for k = 1:N]

    for functional in (J_T_sm, J_T_re, J_T_ss)
        grad_auto! = make_gradient(functional, objectives; via=:tau)
        grad_zyg! = make_zygote_gradient(functional, objectives)
        grad_auto!(G1, τ, ∇τ)
        grad_zyg!(G2, τ, ∇τ)
        @test norm(G1 - G2) < 1e-15
    end

end

@testset "make-chi" begin

    # Test that the routine returned by `make_chi` gives the same result
    # as the Zygote chi

    objectives = PROBLEM.objectives
    χ1 = [similar(obj.initial_state) for obj in objectives]
    χ2 = [similar(obj.initial_state) for obj in objectives]
    ϕ = [random_state_vector(N_HILBERT) for k = 1:N]
    τ = [obj.target_state ⋅ ϕ[k] for (k, obj) in enumerate(objectives)]

    for functional in (J_T_sm, J_T_re, J_T_ss)
        chi_auto! = make_chi(functional, objectives)
        chi_zyg! = make_zygote_chi(functional, objectives)
        chi_auto!(χ1, ϕ, objectives; τ=τ)
        chi_zyg!(χ2, ϕ, objectives)
        @test maximum(norm.(χ1 .- χ2)) < 1e-15
    end

end

@testset "gradient-via-chi" begin

    # Test that gradients calculated via chi give the same results as gradients
    # calculated via tau

    problem = dummy_control_problem(; N=N_HILBERT, n_objectives=2, n_controls=1, n_steps=2)
    objectives = problem.objectives
    dt = problem.tlist[2] - problem.tlist[1]
    Ĥ₀ = Array(objectives[1].generator[1])
    Ĥ₁ = Array(objectives[1].generator[2][1])
    𝟘 = zeros(ComplexF64, N_HILBERT, N_HILBERT)
    ϵ = objectives[1].generator[2][2]
    Ψtgt = [obj.target_state for obj in objectives]
    χ = [similar(Ψ) for Ψ in Ψtgt]
    χ_from_τ = [similar(Ψ) for Ψ in Ψtgt]
    ∇τ_tgt = [zeros(ComplexF64, 1, 2) for k = 1:2]
    ∇τ_chi = [zeros(ComplexF64, 1, 2) for k = 1:2]
    G1 = zeros(2)
    G2 = zeros(2)

    # forward propagation of states
    Ψ0 = [obj.initial_state for obj in objectives]
    Ĥ1 = Ĥ₀ + ϵ[1] * Ĥ₁
    Ĥ2 = Ĥ₀ + ϵ[2] * Ĥ₁
    Ψ1 = [exp(-1im * Ĥ1 * dt) * Ψ for Ψ in Ψ0]
    Ψ2 = [exp(-1im * Ĥ2 * dt) * Ψ for Ψ in Ψ1]

    # overlap with target state
    τ_tgt = [Ψtgt[k] ⋅ Ψ2[k] for k = 1:2]

    # generators for gradient backward propagation
    G1⁺ = [Ĥ1'  Ĥ₁'; 𝟘  Ĥ1']
    G2⁺ = [Ĥ2'  Ĥ₁'; 𝟘  Ĥ2']

    # backward propagation of target states / gradients
    ξ̃2 = [[zeros(ComplexF64, N_HILBERT); Ψtgt[k]] for k = 1:2]
    ξ̃1 = [exp(1im * G2⁺ * dt) * ξ̃ for ξ̃ in ξ̃2]
    for k = 1:2
        ∇τ_tgt[k][1, 2] = ξ̃1[k][1:N_HILBERT] ⋅ Ψ1[k]
        ξ̃1[k][1:N_HILBERT] .= 0
    end
    ξ̃0 = [exp(1im * G1⁺ * dt) * ξ̃ for ξ̃ in ξ̃1]
    for k = 1:2
        ∇τ_tgt[k][1, 1] = ξ̃0[k][1:N_HILBERT] ⋅ Ψ0[k]
        ξ̃0[k][1:N_HILBERT] .= 0
    end

    for functional in (J_T_sm, J_T_re, J_T_ss)

        # boundary condition for χ-backward prop
        chi! = make_chi(functional, objectives)
        chi!(χ, Ψ2, objectives)
        chi!(χ_from_τ, Ψ2, objectives; τ=τ_tgt)
        @test maximum(norm.(χ .- χ_from_τ)) < 1e-15
        τ_chi = [χ[k] ⋅ Ψ2[k] for k = 1:2]

        # backward propagation of chi states / gradients
        χ̃2 = [[zeros(ComplexF64, N_HILBERT); χ[k]] for k = 1:2]
        χ̃1 = [exp(1im * G2⁺ * dt) * χ̃ for χ̃ in χ̃2]
        for k = 1:2
            ∇τ_chi[k][1, 2] = χ̃1[k][1:N_HILBERT] ⋅ Ψ1[k]
            χ̃1[k][1:N_HILBERT] .= 0
        end
        χ̃0 = [exp(1im * G1⁺ * dt) * χ̃ for χ̃ in χ̃1]
        for k = 1:2
            ∇τ_chi[k][1, 1] = χ̃0[k][1:N_HILBERT] ⋅ Ψ0[k]
            χ̃0[k][1:N_HILBERT] .= 0
        end

        grad_J_T_via_tau! = make_gradient(functional, objectives; via=:tau)
        grad_J_T_via_chi! = make_gradient(functional, objectives; via=:chi)
        @test grad_J_T_via_tau! ≢ grad_J_T_via_chi!
        grad_J_T_via_tau!(G1, τ_tgt, ∇τ_tgt)
        grad_J_T_via_chi!(G2, τ_chi, ∇τ_chi)
        @test norm(G1 - G2) < 1e-15

    end

end


text:
# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENCE in the project root.
# ------------------------------------------------------------------

"""
  log_gpd_pdf(arg, μ, σ, ξ)

Log density function of the generalized Pareto distribution,
with an expansion with ξ near zero.
"""
function log_gpd_pdf(_x, μ, σ, ξ)
  x = (_x-μ)/σ
  expn = if abs(ξ) < 1e-5
    # expansion for ξ near zero.
    -x*(ξ+1) + (x^2)*ξ*(ξ+1)/2 - (x^3)*(ξ^2)*(ξ+1)/3 + (x^4)*(ξ^3)*(ξ+1)/4
  else
    (-(ξ+1)/ξ)*log(max(0, 1 + x*ξ))
  end
  expn - log(σ)
end

"""
  log_gev_pdf(arg, μ, σ, ξ)

Log density function of the generalized extreme value distribution,
with an expansion with ξ near zero.
"""
function log_gev_pdf(_x, μ, σ, ξ)
  x = (_x-μ)/σ
  tx = if abs(ξ) < 1e-10 # this cutoff is _not_ the same as for log_gpd_pdf.
    # expansion near zero.
    -x + (x^2)*ξ/2 - (x^3)*(ξ^2)/3 + (x^4)*(ξ^3)/4
  else
    (-1/ξ)*log(max(0, 1 + x*ξ))
  end
  (ξ+1)*tx - exp(tx) - log(σ)
end

"""
    fit_mle(gev, bm)

Fit generalized extreme value distribution `gev` to block maxima
`bm` with constrained maximum likelihood estimation.
"""
function fit_mle(::Type{GeneralizedExtremeValue}, bm::BlockMaxima;
                 init=(μ=0.0, σ=1.0, ξ=0.05))
  # retrieve maxima values
  x = collect(bm)
  n = length(x)

  # set up the problem
  mle = Model(optimizer_with_attributes(Ipopt.Optimizer, "print_level"=>0,
                                        "sb"=>"yes", "max_iter"=>250))
  @variable(mle, μ, start=init.μ)
  @variable(mle, σ, start=init.σ)
  @variable(mle, ξ, start=init.ξ)
  JuMP.register(mle, :log_gev_pdf, 4, log_gev_pdf, autodiff=true)
  @NLobjective(mle, Max, sum(log_gev_pdf(z, μ, σ, ξ) for z in x))
  @NLconstraint(mle, [i=1:n], 1 + ξ*(x[i]-μ)/σ ≥ 0.0)
  @constraint(mle, σ ≥ 1e-10)
  @constraint(mle, ξ ≥ -1/2)

  # attempt to solve
  optimize!(mle)

  # retrieve solver status
  status = termination_status(mle)

  # acceptable statuses
  OK = (MOI.OPTIMAL, MOI.LOCALLY_SOLVED)

  if status ∈ OK
    GeneralizedExtremeValue(value(μ), value(σ), value(ξ))
  else
    error("could not fit distribution. Exited with status $status.")
  end
end

"""
    fit_mle(gp, pm)

Fit generalized Pareto distribution `gp` to peak over threshold
maxima `pm` with constrained maximum likelihood estimation.
"""
function fit_mle(::Type{GeneralizedPareto}, pm::PeakOverThreshold;
                 init=(σ=1.0, ξ=0.05))
  # retrieve maxima values
  x = collect(pm)
  n = length(x)

  # set up the problem
  mle = Model(optimizer_with_attributes(Ipopt.Optimizer, "print_level"=>0,
                                        "sb"=>"yes", "max_iter"=>250))
  @variable(mle, σ, start=init.σ)
  @variable(mle, ξ, start=init.ξ)
  JuMP.register(mle, :log_gpd_pdf, 4, log_gpd_pdf, autodiff=true)
  @NLobjective(mle, Max, sum(log_gpd_pdf(z, pm.u, σ, ξ) for z in x))
  @NLconstraint(mle, [i=1:n], 1 + ξ*(x[i]-pm.u)/σ ≥ 0.0)
  @constraint(mle, σ ≥ 1e-10)
  @constraint(mle, ξ ≥ -1/2)

  # attempt to solve both cases
  optimize!(mle)

  # retrieve solver status
  status  = termination_status(mle)

  # acceptable statuses
  OK = (MOI.OPTIMAL, MOI.LOCALLY_SOLVED)

  if status ∈ OK
    GeneralizedPareto(pm.u, value(σ), value(ξ))
  else
    error("could not fit distribution. Exited with status $status.")
  end
end


text:
include("matlab.jl")
include("edges.jl")
include("geometry.jl")
include("chunk_u.jl")
include("import.jl") #io outside module
include("io.jl") #io module


text:
#========================================================================================#
#	Laboratory 100
#
# Welcome to course 100: Foundations of Mathematics and Physics!
#
# Author: Niall Palfreyman, 24/04/2022
#========================================================================================#
[
	Activity(
		"""
		???
		""",
		"???",
		x -> true
	),
]


text:
using Test
using libcxxwrap_julia_jll

excluded = ["runtests.jl", "testcommon.jl"]

if Sys.iswindows() && Sys.WORD_SIZE == 32
  push!(excluded, "except.jl")
end

if get(ENV, "CXXWRAP_DISABLE_GC", "0") == "1"
  println("disabling GC for tests")
  GC.enable(false)
end

@testset "CxxWrap tests" begin
  for f in filter(fname -> fname ∉ excluded, readdir(@__DIR__))
    println("Running tests from $f...")
    include(f)
  end
end