using Test
using QuantumSavory
using ResumableFunctions
using ConcurrentSim
using QuantumSavory.ProtocolZoo
using QuantumSavory.ProtocolZoo: EntanglementCounterpart, EntanglementHistory, EntanglementUpdateX, EntanglementUpdateZ
using Graphs
using Random
using Logging

@testset "ProtocolZoo Entanglement Tracker" begin

if isinteractive()
    logger = ConsoleLogger(Logging.Debug; meta_formatter=(args...)->(:black,"",""))
    global_logger(logger)
    println("Logger set to debug")
end

tagprefix(tag, n) = Tuple(tag[i] for i in 1:n)

##

# without an entanglement tracker

for i in 1:10

    net = RegisterNet([Register(3), Register(4), Register(2), Register(3)])
    sim = get_time_tracker(net)


    entangler1 = EntanglerProt(sim, net, 1, 2; rounds=1)
    @process entangler1()
    run(sim, 20)

    @test tagprefix.([net[1].tag_info[i].tag for i in net[1].guids], 3) == [(EntanglementCounterpart, 2, 1)]

    entangler2 = EntanglerProt(sim, net, 2, 3; rounds=1)
    @process entangler2()
    run(sim, 40)
    entangler3 = EntanglerProt(sim, net, 4, 3; rounds=1)
    @process entangler3()
    run(sim, 60)

    @test tagprefix.([net[1].tag_info[i].tag for i in net[1].guids], 3) == [(EntanglementCounterpart, 2, 1)]
    @test tagprefix.([net[2].tag_info[i].tag for i in net[2].guids], 3) == [(EntanglementCounterpart, 1, 1), (EntanglementCounterpart, 3, 1)]
    @test tagprefix.([net[3].tag_info[i].tag for i in net[3].guids], 3) == [(EntanglementCounterpart, 2, 2), (EntanglementCounterpart, 4, 1)]
    @test tagprefix.([net[4].tag_info[i].tag for i in net[4].guids], 3) == [(EntanglementCounterpart, 3, 2)]

    @test [islocked(ref) for i in vertices(net) for ref in net[i]] |> any == false


    swapper2 = SwapperProt(sim, net, 2; nodeL = <(2), nodeH = >(2), chooseL = argmin, chooseH = argmax, rounds = 1)
    swapper3 = SwapperProt(sim, net, 3; nodeL = <(3), nodeH = >(3), chooseL = argmin, chooseH = argmax, rounds = 1)
    @process swapper2()
    @process swapper3()
    run(sim, 80)

    # In the absence of an entanglement tracker the tags will not all be updated
    @test tagprefix.([net[1].tag_info[i].tag for i in net[1].guids], 3) == [(EntanglementCounterpart, 2, 1)]
    @test tagprefix.([net[2].tag_info[i].tag for i in net[2].guids], 6) == [(EntanglementHistory, 1, 1, 3, 1, 2), (EntanglementHistory, 3, 1, 1, 1, 1)]
    @test tagprefix.([net[3].tag_info[i].tag for i in net[3].guids], 6) == [(EntanglementHistory, 2, 2, 4, 1, 2), (EntanglementHistory, 4, 1, 2, 2, 1)]
    @test tagprefix.([net[4].tag_info[i].tag for i in net[4].guids], 3) == [(EntanglementCounterpart, 3, 2)]

    @test isassigned(net[1][1]) && isassigned(net[4][1])
    @test !isassigned(net[2][1]) && !isassigned(net[3][1])
    @test !isassigned(net[2][2]) && !isassigned(net[3][2])

    @test [islocked(ref) for i in vertices(net) for ref in net[i]] |> any == false

end

##

if isinteractive()
    logger = ConsoleLogger(Logging.Warn; meta_formatter=(args...)->(:black,"",""))
    global_logger(logger)
    println("Logger set to debug")
end

##

# same but this time with an entanglement tracker

for i in 1:30, n in 2:30
    net = RegisterNet([Register(j+3) for j in 1:n]; classical_delay=1e-9)
    sim = get_time_tracker(net)
    for j in vertices(net)
        tracker = EntanglementTracker(sim, net, j)
        @process tracker()
    end
    for e in edges(net)
        eprot = EntanglerProt(sim, net, e.src, e.dst; rounds=1, randomize=true)
        @process eprot()
    end
    for j in 2:n-1
        swapper = SwapperProt(sim, net, j; nodeL = <(j), nodeH = >(j), chooseL = argmin, chooseH = argmax, rounds = 1)
        @process swapper()
    end
    run(sim, 200)

    q1 = query(net[1], EntanglementCounterpart, n, ❓, ❓)
    q2 = isnothing(q1) ? nothing : query(net[n], EntanglementCounterpart, 1, q1.slot.idx, q1.tag[4])
    @test q1.tag[2] == n
    @test q2.tag[2] == 1
    @test observable((q1.slot, q2.slot), Z⊗Z) ≈ 1
    @test observable((q1.slot, q2.slot), X⊗X) ≈ 1
end
end
