using Test, YaoBlocks, YaoArrayRegister

@testset "test constructor" for T in [Float16, Float32, Float64]
    # NOTE: type should follow the axis
    @test RotationGate(X, 0.1) isa PrimitiveBlock{1}
    @test_throws TypeError RotationGate{1, Complex{T}, XGate} # will not accept non-real type

    @test Rx(T(0.1)) isa RotationGate{1, T, XGate}
    @test Ry(T(0.1)) isa RotationGate{1, T, YGate}
    @test Rz(T(0.1)) isa RotationGate{1, T, ZGate}
end

@testset "test matrix" begin
    theta = 2.0
    for (DIRECTION, MAT) in [
        (X, [cos(theta/2) -im*sin(theta/2); -im*sin(theta/2) cos(theta/2)]),
        (Y, [cos(theta/2) -sin(theta/2); sin(theta/2) cos(theta/2)]),
        (Z, [exp(-im*theta/2) 0;0 exp(im*theta/2)]),
        (CNOT, exp(-mat(CNOT)/2*theta*im |> Matrix)),
        (control(2, (1,), 2=>X), exp(-mat(CNOT)/2*theta*im |> Matrix))]

        @test mat(RotationGate(DIRECTION, theta)) ≈ MAT
    end
end

@testset "test apply" begin
    r = rand_state(1)
    @test state(apply!(copy(r), Rx(0.1))) ≈ mat(Rx(0.1)) * state(r)
end

@testset "test dispatch" begin
    @test dispatch!(Rx(0.1), 0.3) == Rx(0.3)
    @test nparameters(Rx(0.1)) == 1

    @testset "test $op" for op in [+, -, *, /]
        @test dispatch!(op, Rx(0.1), π) == Rx(op(0.1, π))
    end

    @test_throws AssertionError dispatch!(Rx(0.1), (0.2, 0.3))
end

@testset "adjoints" begin
    @test Rx(0.1)' == Rx(-0.1)
    @test Rx(0.2)' == Rx(-0.2)
    @test copy(Rx(0.1)) == Rx(0.1)
    
    g = Rx(0.1) # creates a new one
    @test copy(g) !== g
end
