using SymbolicUtils, Test
using SymbolicUtils: Term, Add, Mul, Div, metadata, BasicSymbolic, PolyCoeffT, ACDict
using SymbolicUtils: Fill, Mapper, Mapreducer, ShapeVecT, SymReal
import MultivariatePolynomials as MP
import TermInterface

isequal2(a, b) = SymbolicUtils.@manually_scope SymbolicUtils.COMPARE_FULL => true isequal(a, b)

const BSImpl = SymbolicUtils.BasicSymbolicImpl

struct Ctx1 end
struct Ctx2 end

struct ThrowingEqualityScalar
    tag::Symbol
    value::Int
end

Base.hash(x::ThrowingEqualityScalar, h::UInt) = hash(x.value, h)
function Base.isequal(a::ThrowingEqualityScalar, b::ThrowingEqualityScalar)
    a.tag === b.tag && return isequal(a.value, b.value)
    throw(ArgumentError("incomparable scalar tags"))
end

@testset "Sym" begin
    x1 = only(@syms x)
    x2 = only(@syms x)
    @test x1.id === x2.id
    x3 = only(@syms x::Float64)
    @test x1.id !== x3.id
    x4 = only(@syms x::Float64)
    @test x1.id !== x4.id
    @test x3.id === x4.id
    x5 = only(@syms x::Int)
    x6 = only(@syms x::Int)
    @test x1.id !== x5.id
    @test x3.id !== x5.id
    @test x5.id === x6.id

    xm1 = setmetadata(x1, Ctx1, "meta_1")
    xm2 = setmetadata(x1, Ctx1, "meta_1")
    @test xm1.id === xm2.id
    xm3 = setmetadata(x1, Ctx2, "meta_2")
    @test xm1.id !== xm3.id
end

@syms a b c

@testset "Term" begin
    t1 = sin(a)
    t2 = sin(a)
    @test t1.id === t2.id
    t3 = Term{SymReal}(identity, [a])
    t4 = Term{SymReal}(identity, [a])
    @test t3.id === t4.id
    t5 = Term{SymReal}(identity, [a]; type = Int)
    @test t3.id !== t5.id
    tm1 = setmetadata(t1, Ctx1, "meta_1")
    @test t1.id !== tm1.id
end

@testset "Add" begin
    d1 = a + b
    d2 = b + a
    @test d1.id === d2.id
    d3 = b - 2 + a
    d4 = a + b - 2
    @test d3.id === d4.id
    d5 = Add{SymReal}(0, ACDict{SymReal}(a => 1, b => 1); type = Int)
    @test d5.id !== d1.id

    dm1 = setmetadata(d1, Ctx1, "meta_1")
    @test d1.id !== dm1.id
end

@testset "Mul" begin
    m1 = a * b
    m2 = b * a
    @test m1.id === m2.id
    m3 = 6 * a * b
    m4 = 3 * a * 2 * b
    @test m3.id === m4.id
    m5 = Mul{SymReal}(1, ACDict{SymReal}(a => 1, b => 1); type = Int)
    @test m5.id !== m1.id

    mm1 = setmetadata(m1, Ctx1, "meta_1")
    @test m1.id !== mm1.id
end

@testset "Div" begin
    v1 = a / b
    v2 = a / b
    @test v1.id === v2.id
    v3 = -1 / a
    v4 = -1 / a
    @test v3.id === v4.id
    v5 = 3a / 6
    v6 = 2a / 4
    @test v5.id === v6.id
    v7 = Div{SymReal}(-1, a, false; type = Float64)
    @test v7.id !== v3.id

    vm1 = setmetadata(v1, Ctx1, "meta_1")
    @test vm1.id !== v1.id
end

@testset "Pow" begin
    p1 = a^b
    p2 = a^b
    @test p1.id === p2.id
    p3 = a^(2^-b)
    p4 = a^(2^-b)
    @test p3.id === p4.id

    pm1 = setmetadata(p1, Ctx1, "meta_1")
    @test pm1.id !== p1.id
end

@testset "Symbolics in metadata" begin
    @syms a b
    a1 = setmetadata(a, Int, b)
    b1 = setmetadata(b, Int, 3)
    a2 = setmetadata(a, Int, b1)
    @test a1.id !== a2.id
    @test !isequal2(a1, a2)
    @test metadata(metadata(a1)[Int]) === nothing
    @test metadata(metadata(a2)[Int])[Int] == 3
end

@testset "Compare metadata of expression tree" begin
    @syms a b
    aa = setmetadata(a, Int, b)
    @test aa.id !== a.id
    @test isequal(a, aa)
    @test !isequal2(a, aa)
    @test !isequal2(2a, 2aa)
end

@testset "Hashconsing can be toggled" begin
    SymbolicUtils.ENABLE_HASHCONSING[] = false
    @syms a b
    x1 = a + b
    x2 = a + b
    @test x1.id === nothing === x2.id
    SymbolicUtils.ENABLE_HASHCONSING[] = true
end

@testset "`hash` is cached" begin
    @syms a b f(..)
    for ex in [a + b, a * b, f(a)]
        h = hash(ex)
        @test h == ex.hash[]
        ex2 = setmetadata(ex, Int, 3)
        # the hash is metadata-free
        @test hash(ex2) == h
    end
end

@testset "Metadata type is considered in hashconsing" begin
    @syms x
    x1 = setmetadata(x, Int, 1)
    x2 = setmetadata(x, Int, 1.0)
    @test x1 !== x2
    @test !isequal2(x1, x2)

    x1 = setmetadata(x, Int, [1])
    x2 = setmetadata(x, Int, [1])
    @test x1 === x2
    @test hash(x1) == hash(x2)
end

@testset "Throwing scalar equality is not a hashconsing match" begin
    @test_throws ArgumentError isequal(
        ThrowingEqualityScalar(:tag1, 1), ThrowingEqualityScalar(:tag2, 1)
    )

    x1 = BSImpl.Const{SymReal}(ThrowingEqualityScalar(:tag1, 1))
    x2 = BSImpl.Const{SymReal}(ThrowingEqualityScalar(:tag2, 1))
    x3 = BSImpl.Const{SymReal}(ThrowingEqualityScalar(:tag1, 1))

    @test x1 !== x2
    @test x1 === x3
    @test !isequal2(x1, x2)
end

@testset "Operations hash by content" begin
    @syms p q

    @testset "Fill" begin
        f1 = Fill(ShapeVecT([1:3, 1:2]))
        f2 = Fill(ShapeVecT([1:3, 1:2]))
        f3 = Fill(ShapeVecT([1:3, 1:4]))
        @test f1 == f2
        @test isequal(f1, f2)
        @test hash(f1) == hash(f2)
        @test isequal(f1, deepcopy(f1)) && hash(f1) == hash(deepcopy(f1))
        @test f1 != f3
        @test hash(f1) != hash(f3)
        # As an `ArrayOp` operation, content-identical fills hash-cons together.
        a1 = f1(p)
        a2 = f2(p)
        @test isequal(a1, a2)
        @test hash(a1) == hash(a2)
    end

    @testset "Mapper" begin
        m1 = Mapper(sin)
        m2 = Mapper(sin)
        m3 = Mapper(cos)
        @test m1 == m2
        @test isequal(m1, m2)
        @test hash(m1) == hash(m2)
        @test isequal(m1, deepcopy(m1)) && hash(m1) == hash(deepcopy(m1))
        @test m1 != m3
        @test hash(m1) != hash(m3)
        # As a `Term` operation, content-identical mappers hash-cons to the same node.
        t1 = Term{SymReal}(Mapper(sin), [p])
        t2 = Term{SymReal}(Mapper(sin), [p])
        @test t1.id === t2.id
        @test hash(t1) == hash(t2)
    end

    @testset "Mapreducer" begin
        r1 = Mapreducer(sin, +, 1, nothing)
        r2 = Mapreducer(sin, +, 1, nothing)
        r3 = Mapreducer(sin, +, 2, nothing)   # different dims
        r4 = Mapreducer(cos, +, 1, nothing)   # different f
        r5 = Mapreducer(sin, +, 1, 0.0)       # different init
        @test r1 == r2
        @test isequal(r1, r2)
        @test hash(r1) == hash(r2)
        @test isequal(r1, deepcopy(r1)) && hash(r1) == hash(deepcopy(r1))
        @test r1 != r3 && hash(r1) != hash(r3)
        @test r1 != r4 && hash(r1) != hash(r4)
        @test r1 != r5 && hash(r1) != hash(r5)
        t1 = Term{SymReal}(Mapreducer(sin, +, 1, nothing), [p])
        t2 = Term{SymReal}(Mapreducer(sin, +, 1, nothing), [p])
        @test t1.id === t2.id
        @test hash(t1) == hash(t2)
    end
end
