using Test
using ITensors
using ITensors.NDTensors

@testset "Allow general mixtures of collections of indices" begin
  d = 2
  is = Index.((d, d, d, d, d), ("i", "j", "k", "l", "m"))
  i, j, k, l, m = is
  is1 = ([i, j], k, (l, m))
  is2 = [[i, j], k, (l, m)]
  A = randn(dims(is))
  D = randn(minimum(dims(is)))
  x = randn()
  @test hassameinds(ITensor(i), (i,))
  @test hassameinds(ITensor(Float64, i), (i,))
  @test hassameinds(ITensor(is1), is)
  @test hassameinds(ITensor(is2), is)
  @test hassameinds(ITensor(is1...), is)
  @test hassameinds(ITensor(Float64, is1), is)
  @test hassameinds(ITensor(Float64, is2), is)
  @test hassameinds(ITensor(Float64, is1...), is)
  @test hassameinds(randomITensor(is1), is)
  @test hassameinds(randomITensor(is2), is)
  @test hassameinds(randomITensor(is1...), is)
  @test hassameinds(randomITensor(Float64, is1), is)
  @test hassameinds(randomITensor(Float64, is2), is)
  @test hassameinds(randomITensor(Float64, is1...), is)
  @test hassameinds(ITensor(x, is1), is)
  @test hassameinds(ITensor(x, is2), is)
  @test hassameinds(ITensor(x, is1...), is)
  @test hassameinds(ITensor(Float64, x, is1), is)
  @test hassameinds(ITensor(Float64, x, is2), is)
  @test hassameinds(ITensor(Float64, x, is1...), is)
  @test hassameinds(ITensor(Float64, undef, is1), is)
  @test hassameinds(ITensor(Float64, undef, is2), is)
  @test hassameinds(ITensor(Float64, undef, is1...), is)
  @test hassameinds(ITensor(undef, is1), is)
  @test hassameinds(ITensor(undef, is2), is)
  @test hassameinds(ITensor(undef, is1...), is)
  @test hassameinds(emptyITensor(is1), is)
  @test hassameinds(emptyITensor(is2), is)
  @test hassameinds(emptyITensor(is1...), is)
  @test hassameinds(emptyITensor(Float64, is1), is)
  @test hassameinds(emptyITensor(Float64, is2), is)
  @test hassameinds(emptyITensor(Float64, is1...), is)
  @test hassameinds(ITensor(A, is1), is)
  @test hassameinds(ITensor(A, is2), is)
  @test hassameinds(ITensor(A, is1...), is)
  @test hassameinds(itensor(A, is1), is)
  @test hassameinds(itensor(A, is2), is)
  @test hassameinds(itensor(A, is1...), is)
  @test hassameinds(ITensor(Float64, A, is1), is)
  @test hassameinds(ITensor(Float64, A, is2), is)
  @test hassameinds(ITensor(Float64, A, is1...), is)
  @test hassameinds(itensor(Float64, A, is1), is)
  @test hassameinds(itensor(Float64, A, is2), is)
  @test hassameinds(itensor(Float64, A, is1...), is)
  @test hassameinds(diagITensor(is1), is)
  @test hassameinds(diagITensor(is2), is)
  @test hassameinds(diagITensor(is1...), is)
  @test hassameinds(diagITensor(Float64, is1), is)
  @test hassameinds(diagITensor(Float64, is2), is)
  @test hassameinds(diagITensor(Float64, is1...), is)
  @test hassameinds(diagITensor(D, is1), is)
  @test hassameinds(diagITensor(D, is2), is)
  @test hassameinds(diagITensor(D, is1...), is)
  @test hassameinds(diagITensor(Float64, D, is1), is)
  @test hassameinds(diagITensor(Float64, D, is2), is)
  @test hassameinds(diagITensor(Float64, D, is1...), is)
  @test hassameinds(diagITensor(x, is1), is)
  @test hassameinds(diagITensor(x, is2), is)
  @test hassameinds(diagITensor(x, is1...), is)
  @test hassameinds(diagITensor(Float64, x, is1), is)
  @test hassameinds(diagITensor(Float64, x, is2), is)
  @test hassameinds(diagITensor(Float64, x, is1...), is)
  @test hassameinds(diagitensor(D, is1), is)
  @test hassameinds(diagitensor(D, is2), is)
  @test hassameinds(diagitensor(D, is1...), is)
  @test hassameinds(diagitensor(Float64, D, is1), is)
  @test hassameinds(diagitensor(Float64, D, is2), is)
  @test hassameinds(diagitensor(Float64, D, is1...), is)
  @test hassameinds(delta(is1), is)
  @test hassameinds(delta(is2), is)
  @test hassameinds(delta(is1...), is)
  @test hassameinds(delta(Float64, is1), is)
  @test hassameinds(delta(Float64, is2), is)
  @test hassameinds(delta(Float64, is1...), is)
  @test hasinds(combiner(is1), is)
  @test hasinds(combiner(is2), is)
  @test hasinds(combiner(is1...), is)
end

@testset "Allow general mixtures of collections of QN indices" begin
  d = [QN() => 2]
  is = Index.((d, d, d, d, d), ("i", "j", "k", "l", "m"))
  i, j, k, l, m = is
  is1 = ([i, j], k, (l, m))
  is2 = [[i, j], k, (l, m)]
  A = randn(dims(is))
  D = randn(minimum(dims(is)))
  x = randn()
  @test hassameinds(ITensor(i), (i,))
  @test hassameinds(ITensor(Float64, i), (i,))
  @test hassameinds(ITensor(is1), is)
  @test hassameinds(ITensor(is2), is)
  @test hassameinds(ITensor(is1...), is)
  @test hassameinds(ITensor(Float64, is1), is)
  @test hassameinds(ITensor(Float64, is2), is)
  @test hassameinds(ITensor(Float64, is1...), is)
  @test hassameinds(randomITensor(is1), is)
  @test hassameinds(randomITensor(is2), is)
  @test hassameinds(randomITensor(is1...), is)
  @test hassameinds(randomITensor(Float64, is1), is)
  @test hassameinds(randomITensor(Float64, is2), is)
  @test hassameinds(randomITensor(Float64, is1...), is)
  @test hassameinds(ITensor(x, is1), is)
  @test hassameinds(ITensor(x, is2), is)
  @test hassameinds(ITensor(x, is1...), is)
  @test hassameinds(ITensor(Float64, x, is1), is)
  @test hassameinds(ITensor(Float64, x, is2), is)
  @test hassameinds(ITensor(Float64, x, is1...), is)
  @test hassameinds(ITensor(Float64, undef, is1), is)
  @test hassameinds(ITensor(Float64, undef, is2), is)
  @test hassameinds(ITensor(Float64, undef, is1...), is)
  @test hassameinds(ITensor(undef, is1), is)
  @test hassameinds(ITensor(undef, is2), is)
  @test hassameinds(ITensor(undef, is1...), is)
  @test hassameinds(emptyITensor(is1), is)
  @test hassameinds(emptyITensor(is2), is)
  @test hassameinds(emptyITensor(is1...), is)
  @test hassameinds(emptyITensor(Float64, is1), is)
  @test hassameinds(emptyITensor(Float64, is2), is)
  @test hassameinds(emptyITensor(Float64, is1...), is)
  @test hassameinds(ITensor(A, is1), is)
  @test hassameinds(ITensor(A, is2), is)
  @test hassameinds(ITensor(A, is1...), is)
  @test hassameinds(itensor(A, is1), is)
  @test hassameinds(itensor(A, is2), is)
  @test hassameinds(itensor(A, is1...), is)
  @test hassameinds(ITensor(Float64, A, is1), is)
  @test hassameinds(ITensor(Float64, A, is2), is)
  @test hassameinds(ITensor(Float64, A, is1...), is)
  @test hassameinds(itensor(Float64, A, is1), is)
  @test hassameinds(itensor(Float64, A, is2), is)
  @test hassameinds(itensor(Float64, A, is1...), is)
  @test hassameinds(diagITensor(is1), is)
  @test hassameinds(diagITensor(is2), is)
  @test hassameinds(diagITensor(is1...), is)
  @test hassameinds(diagITensor(Float64, is1), is)
  @test hassameinds(diagITensor(Float64, is2), is)
  @test hassameinds(diagITensor(Float64, is1...), is)
  @test hassameinds(diagITensor(D, is1), is)
  @test hassameinds(diagITensor(D, is2), is)
  @test hassameinds(diagITensor(D, is1...), is)
  @test hassameinds(diagITensor(Float64, D, is1), is)
  @test hassameinds(diagITensor(Float64, D, is2), is)
  @test hassameinds(diagITensor(Float64, D, is1...), is)
  @test hassameinds(diagITensor(x, is1), is)
  @test hassameinds(diagITensor(x, is2), is)
  @test hassameinds(diagITensor(x, is1...), is)
  @test hassameinds(diagITensor(Float64, x, is1), is)
  @test hassameinds(diagITensor(Float64, x, is2), is)
  @test hassameinds(diagITensor(Float64, x, is1...), is)
  @test hassameinds(diagitensor(D, is1), is)
  @test hassameinds(diagitensor(D, is2), is)
  @test hassameinds(diagitensor(D, is1...), is)
  @test hassameinds(diagitensor(Float64, D, is1), is)
  @test hassameinds(diagitensor(Float64, D, is2), is)
  @test hassameinds(diagitensor(Float64, D, is1...), is)
  @test hassameinds(delta(is1), is)
  @test hassameinds(delta(is2), is)
  @test hassameinds(delta(is1...), is)
  @test hassameinds(delta(Float64, is1), is)
  @test hassameinds(delta(Float64, is2), is)
  @test hassameinds(delta(Float64, is1...), is)
  @test hasinds(combiner(is1), is)
  @test hasinds(combiner(is2), is)
  @test hasinds(combiner(is1...), is)

  # With flux
  @test hassameinds(ITensor(QN(), i), (i,))
  @test hassameinds(ITensor(Float64, QN(), i), (i,))
  @test hassameinds(ITensor(QN(), is1), is)
  @test hassameinds(ITensor(QN(), is2), is)
  @test hassameinds(ITensor(QN(), is1...), is)
  @test hassameinds(ITensor(Float64, QN(), is1), is)
  @test hassameinds(ITensor(Float64, QN(), is2), is)
  @test hassameinds(ITensor(Float64, QN(), is1...), is)
  @test hassameinds(randomITensor(QN(), is1), is)
  @test hassameinds(randomITensor(QN(), is2), is)
  @test hassameinds(randomITensor(QN(), is1...), is)
  @test hassameinds(randomITensor(Float64, QN(), is1), is)
  @test hassameinds(randomITensor(Float64, QN(), is2), is)
  @test hassameinds(randomITensor(Float64, QN(), is1...), is)
  @test hassameinds(ITensor(x, QN(), is1), is)
  @test hassameinds(ITensor(x, QN(), is2), is)
  @test hassameinds(ITensor(x, QN(), is1...), is)
  @test hassameinds(ITensor(Float64, x, QN(), is1), is)
  @test hassameinds(ITensor(Float64, x, QN(), is2), is)
  @test hassameinds(ITensor(Float64, x, QN(), is1...), is)
  @test hassameinds(ITensor(Float64, undef, QN(), is1), is)
  @test hassameinds(ITensor(Float64, undef, QN(), is2), is)
  @test hassameinds(ITensor(Float64, undef, QN(), is1...), is)
  @test hassameinds(ITensor(undef, QN(), is1), is)
  @test hassameinds(ITensor(undef, QN(), is2), is)
  @test hassameinds(ITensor(undef, QN(), is1...), is)
  @test_throws ErrorException hassameinds(emptyITensor(QN(), is1), is)
  @test_throws ErrorException hassameinds(emptyITensor(QN(), is2), is)
  @test_throws ErrorException hassameinds(emptyITensor(QN(), is1...), is)
  @test_throws ErrorException hassameinds(emptyITensor(Float64, QN(), is1), is)
  @test_throws ErrorException hassameinds(emptyITensor(Float64, QN(), is2), is)
  @test_throws ErrorException hassameinds(emptyITensor(Float64, QN(), is1...), is)
end

@testset "Test Index collection as Vector of abstract type" begin
  d = 2
  i = Index(d)
  A = randn(d, d)
  T = itensor(A, Index[i', dag(i)])
  @test storage(T) isa NDTensors.Dense{Float64}
  T = itensor(A, Any[i', dag(i)])
  @test storage(T) isa NDTensors.Dense{Float64}

  i = Index([QN() => d])
  A = randn(d, d)
  T = itensor(A, Index[i', dag(i)])
  @test storage(T) isa NDTensors.BlockSparse{Float64}
  T = itensor(A, Any[i', dag(i)])
  @test storage(T) isa NDTensors.BlockSparse{Float64}
end

@testset "Test output types of ITensors.indices" begin
  i = Index(2)
  @test ITensors.indices([i'', i', i]) == Index{Int}[i'', i', i]
  @test ITensors.indices((i'', i', i)) == (i'', i', i)
  @test ITensors.indices(((i'',), (i',), i)) == (i'', i', i)
  @test ITensors.indices(((i'', i'), (i,))) == (i'', i', i)
  @test ITensors.indices([(i'',), (i',), (i,)]) == Index{Int}[i'', i', i]
  @test ITensors.indices(Any[(i'',), (i',), (i,)]) == Index{Int}[i'', i', i]
  @test ITensors.indices([(i'',), (i',), [i]]) == Index{Int}[i'', i', i]
  @test ITensors.indices([(i'',), i', [i]]) == Index{Int}[i'', i', i]
  @test ITensors.indices(Any[(i'',), i', [i]]) == Index{Int}[i'', i', i]
  @test ITensors.indices(((i'',), i', [i])) == Index{Int}[i'', i', i]
end
