yaoqasm.jl

using Yao
using YaoBlocks
using YaoBlocks.ConstGate: TGate, SGate, TdagGate, SdagGate

function Base.show(io::IO, ::MIME"qasm/open", blk::AbstractBlock)
    qcode = qasm(blk)
    println(io, "OPENQASM 2.0")
    maincode = pop!(qcode, "main")
    for k in keys
        println(io, qcode)
    end
    println(io, maincode)
end

######################## Basic Buiding Blocks ########################
const BASIC_1GATES = [:X, :Y, :Z, :H, :S, :Sdag, :T, :Tdag]

addrs2str(locs) = join(["reg[$(i-1)]" for i in locs], ", ")
floats2str(params) = join(["$p" for p in params], ", ")
args2str(pcounts) = join(["p$i" for i=pcounts], ", ")

qasm(blk::PutBlock{N}, args...) where N = qasm(blk.block, args...) * "  " * addrs2str(blk.locs)
qasm(blk::ControlBlock{N}, args...) where N = "C-" * qasm(blk.block, args...) * "  " * addrs2str(blk.locs)

# u3, u2,
# cu1, cu3
# if statement

# x, y, z, h, s, sdag, t, tdag
# u1 == z
for G in BASIC_1GATES
    GT = Symbol(G, :Gate)
    g = G |> String |> lowercase
    @eval qasm(blk::$GT, args...) = "$($g)"
    @eval qasm(blk::PutBlock{<:Any, <:Any, <:$GT}, args...) where N = "$($g)  " * addrs2str(blk.locs)
end

# cx, cy, cz, ch
for G in BASIC_1GATES[1:4]
    GT = Symbol(G, :Gate)
    g = G |> String |> lowercase
    @eval qasm(blk::ControlBlock{<:Any, <:$GT}, args...) where N = "c$($g)  " * addrs2str([blk.ctrl_locs[1], blk.locs[1]])
end

# rx, ry, rz
# crz, crx, cry
# ccx, ccy, ccz
for G in BASIC_1GATES[1:3]
    GT = Symbol(G, :Gate)
    g = G |> String |> lowercase
    @eval qasm(blk::ControlBlock{<:Any, <:RotationGate{<:Any, <:Any, <:$GT}}, args...) where N = "cr$($g)($(blk.content.theta))  " * addrs2str([blk.ctrl_locs[1], blk.locs[1]])
    @eval qasm(blk::PutBlock{<:Any, <:Any, <:RotationGate{<:Any, <:Any, <:$GT}}, args...) where N = "r$($g)($(blk.content.theta))  " * addrs2str(blk.locs)
    @eval qasm(blk::ControlBlock{<:Any, <:$GT, 2}, args...) where N = "cc$($g)  " * addrs2str([blk.ctrl_locs..., blk.locs[1]])
end

# id
qasm(blk::I2Gate, args...) = "id"
qasm(blk::PutBlock{<:Any, <:Any, <:I2Gate}, args...) = "id  " * addrs2str(blk.locs)

function qasm(blk::AbstractBlock, args...)
    typeof(blk).name
end

function qasm(blk::ChainBlock, funcs=Dict{String, String}("main"=>""), fcount::Int=0, pcount::Int=0, level::Int=0)
    GNAME = "G$pcount"
    # however, we don't support function parameters for the moment.
    nparams = nparameters(blk)
    code = "gate $GNAME reg
{
"
    code *= join(["    "^(level+1) * qasm(b, funcs, fcount+1, pcount+nparams, level+1) for b in subblocks(blk)], "
")
    code *= "
}"
    funcs[GNAME] = code
    "$GNAME reg"
end

using Test
@testset "qasm basic" begin
    @test control(3, (1,3), 2=>Z) |> qasm == "ccz  reg[0], reg[2], reg[1]"
    @test control(3, (1,), 2=>X) |> qasm == "cx  reg[0], reg[1]"
    @test put(5, 2=>rot(X, 0.3)) |> qasm == "rx(0.3)  reg[1]"
    @test control(3, (1,), 2=>rot(X, 0.3)) |> qasm == "crx(0.3)  reg[0], reg[1]"
    @test put(5, 3=>X) |> qasm == "x  reg[2]"
    @test put(5, 3=>X) |> qasm == "x  reg[2]"
    @test put(5, 3=>I2) |> qasm == "id  reg[2]"
end