using LoopVectorization, OffsetArrays
using Test

@testset "dot" begin
  dotq = :(
    for i ∈ eachindex(a, b)
      s += a[i] * b[i]
    end
  )
  lsdot = LoopVectorization.loopset(dotq)
  @test LoopVectorization.choose_order(lsdot) ==
        (Symbol[:i], :i, Symbol("##undefined##"), :i, 4, -1)
  function mydot(a::AbstractVector, b::AbstractVector)
    s = zero(eltype(a))
    za = OffsetArray(a, OffsetArrays.Origin(0))
    zb = OffsetArray(b, OffsetArrays.Origin(0))
    @inbounds @simd for i ∈ LoopVectorization.CloseOpen(min(length(a), length(b)))
      s += za[i] * zb[i]
    end
    s
  end
  function mydotavx(a::AbstractVector, b::AbstractVector)
    s = zero(eltype(a))
    za = OffsetArray(a, OffsetArrays.Origin(0))
    zb = OffsetArray(b, OffsetArrays.Origin(0))
    @turbo for i ∈ LoopVectorization.CloseOpen(min(length(a), length(b)))
      s += za[i] * zb[i]
    end
    s
  end
  @test LoopVectorization.ArrayInterface.static_step(
    LoopVectorization.CloseOpen(-5, 10),
  ) === LoopVectorization.One()
  function mydot(a, b)
    s = zero(eltype(a))
    @inbounds @simd for i ∈ eachindex(a, b)
      s += a[i] * b[i]
    end
    s
  end
  function mydotavx(a, b)
    s = zero(eltype(a))
    @turbo for i ∈ eachindex(a, b)
      s += a[i] * b[i]
    end
    s
  end

  function mydot_avx(a, b)
    s = zero(eltype(a))
    @_avx for i ∈ eachindex(a, b)
      aᵢ = getindex(a, i)
      s += aᵢ * b[i]
    end
    s
  end


  selfdotq = :(
    for i ∈ indices(a, 1)
      s += a[i] * a[i]
    end
  )
  lsselfdot = LoopVectorization.loopset(selfdotq)
  @test LoopVectorization.choose_order(lsselfdot) ==
        (Symbol[:i], :i, Symbol("##undefined##"), :i, 8, -1)

  function myselfdot(a)
    s = zero(eltype(a))
    @inbounds @simd for i ∈ eachindex(a)
      s += getindex(a, i) * a[i]
    end
    s
  end
  function myselfdotavx(a)
    s = zero(eltype(a))
    @turbo for i ∈ eachindex(a)
      s += a[i] * a[i]
    end
    s
  end
  function myselfdotavx_v2(a)
    s = zero(eltype(a))
    @turbo for i ∈ 1:length(a)
      s += a[i] * a[i]
    end
    s
  end
  function myselfdotavx_range(a)
    s = zero(eltype(a))
    rng = axes(a, 1)
    @turbo for i ∈ rng
      s += a[i] * a[i]
    end
    s
  end
  function myselfdot_avx(a)
    s = zero(eltype(a))
    @_avx for i ∈ eachindex(a)
      s += a[i] * a[i]
    end
    s
  end
  function dot_unroll2avx(x::Vector{T}, y::Vector{T}) where {T<:Number}
    z = zero(T)
    o = (3,)
    b = (3,)
    @turbo unroll = 2 for i ∈ 1:length(x)
      z += x[i] * y[i+o[1]-b[1]]
    end
    z
  end
  function dot_unroll3avx(x::Vector{T}, y::Vector{T}) where {T<:Number}
    z = zero(T)
    @turbo unroll = 3 for i ∈ 1:length(x)
      z += x[i] * y[i]
    end
    z
  end
  # @macroexpand @turbo inline=false unroll=2 for i ∈ 1:length(x)
  #         z += x[i]*y[i]
  #     end

  function dot_unroll2avx_noinline(x::Vector{T}, y::Vector{T}) where {T<:Number}
    z = zero(T)
    @turbo inline = true unroll = 2 for i ∈ 1:length(x)
      z += x[i] * y[i]
    end
    z
  end
  function dot_unroll3avx_inline(x::Vector{T}, y::Vector{T}) where {T<:Number}
    z = zero(T)
    @turbo unroll = 3 inline = true check_empty = true for i ∈ 1:length(x)
      z += x[i] * y[i]
    end
    z
  end
  function dot_unroll2_avx(x::Vector{T}, y::Vector{T}) where {T<:Number}
    z = zero(T)
    @_avx unroll = 2 for i ∈ 1:length(x)
      z += x[i] * y[i]
    end
    z
  end
  function dot_unroll3_avx(x::Vector{T}, y::Vector{T}) where {T<:Number}
    z = zero(T)
    @_avx unroll = 3 for i ∈ 1:length(x)
      z += x[i] * y[i]
    end
    z
  end
  function complex_dot_soa(
    xre::AbstractVector{T},
    xim::AbstractVector{T},
    yre::AbstractVector{T},
    yim::AbstractVector{T},
  ) where {T}
    zre = zero(T)
    zim = zero(T)
    @turbo for i ∈ 1:length(xre)
      zre += xre[i] * yre[i] - xim[i] * yim[i]
      zim += xre[i] * yim[i] + xim[i] * yre[i]
    end
    Complex{T}(zre, zim)
  end
  qc = :(
    for i ∈ 1:length(xre)
      zre += xre[i] * yre[i] - xim[i] * yim[i]
      zim += xre[i] * yim[i] + xim[i] * yre[i]
    end
  )
  lsc = LoopVectorization.loopset(qc)
  function complex_mul_with_index_offset!(c_re, c_im, a_re, a_im, b_re, b_im)
    @inbounds @simd ivdep for i = 1:length(a_re)-1
      c_re[i] = b_re[i] * a_re[i+1] - b_im[i] * a_im[i+1]
      c_im[i] = b_re[i] * a_im[i+1] + b_im[i] * a_re[i+1]
    end
  end
  function complex_mul_with_index_offsetavx!(c_re, c_im, a_re, a_im, b_re, b_im)
    @turbo for i = 1:length(a_re)-1
      c_re[i] = b_re[i] * a_re[i+1] - b_im[i] * a_im[i+1]
      c_im[i] = b_re[i] * a_im[i+1] + b_im[i] * a_re[i+1]
    end
  end
  function complex_mul_with_index_offset_avx!(c_re, c_im, a_re, a_im, b_re, b_im)
    @_avx for i = 1:length(a_re)-1
      c_re[i] = b_re[i] * a_re[i+1] - b_im[i] * a_im[i+1]
      setindex!(c_im, b_re[i] * a_im[i+1] + b_im[i] * a_re[i+1], i)
    end
  end
  # q = :(for i = 1:length(a_re) - 1
  #         c_re[i] = b_re[i] * a_re[i + 1] - b_im[i] * a_im[i + 1]
  #         c_im[i] = b_re[i] * a_im[i + 1] + b_im[i] * a_re[i + 1]
  #       end);
  # ls = LoopVectorization.loopset(q)

  function mcpi(x, y)
    acc = 0
    @inbounds @simd for i ∈ eachindex(x)
      acc += (x[i] * x[i] + y[i] * y[i]) < 1.0
    end
    4acc / length(x)
  end
  function mcpiavx(x, y)
    acc = 0
    @turbo for i ∈ eachindex(x)
      acc += (x[i] * x[i] + y[i] * y[i]) < 1.0
    end
    4acc / length(x)
  end
  function mcpiavx_u4(x, y)
    acc = 0
    @turbo unroll = 4 for i ∈ eachindex(x)
      acc += (x[i] * x[i] + y[i] * y[i]) < 1.0
    end
    4acc / length(x)
  end
  function mcpi_avx(x, y)
    acc = 0
    @_avx for i ∈ eachindex(x)
      acc += (x[i] * x[i] + y[i] * y[i]) < 1.0
    end
    4acc / length(x)
  end
  function mcpi_avx_u4(x, y)
    acc = 0
    @_avx unroll = 4 for i ∈ 1:length(x)
      acc += (x[i] * x[i] + y[i] * y[i]) < 1.0
    end
    4acc / length(x)
  end

  function dotloopinductvarpow(x::AbstractArray{T}) where {T}
    s = zero(T)
    for i ∈ eachindex(x)
      s += x[i] * T(i)^3
    end
    s
  end
  function dotloopinductvarpowavx(x)
    s = zero(eltype(x))
    @turbo for i ∈ eachindex(x)
      s += x[i] * i^3
    end
    s
  end
  function dot_from_n_to_100(a, b, n)
    s = zero(eltype(a))
    @turbo for i ∈ n:100
      s += a[i] * b[i]
    end
    s
  end
  function dot33(a, b)
    s = zero(eltype(a))
    @turbo for i ∈ 1:33
      s += a[i] * b[i]
    end
    s
  end
  function dot17(a, b)
    s = zero(eltype(a))
    @turbo for i ∈ 1:17
      s += a[i] * b[i]
    end
    s
  end
  # @macroexpand @_avx for i = 1:length(a_re) - 1
  #     c_re[i] = b_re[i] * a_re[i + 1] - b_im[i] * a_im[i + 1]
  #     c_im[i] = b_re[i] * a_im[i + 1] + b_im[i] * a_re[i + 1]
  # end

  # a = rand(400);
  for T ∈ (Float32, Float64, Int32, Int64)
    @show T, @__LINE__
    N = 143
    R = T <: Integer ? (T(-100):T(100)) : T
    a = rand(R, N)
    b = rand(R, N)
    ao = OffsetArray(a, -60:N-61)
    bo = OffsetArray(b, -60:N-61)
    s = mydot(a, b)
    @test mydotavx(a, b) ≈ s
    @test mydot_avx(a, b) ≈ s
    @test mydotavx(ao, bo) ≈ s
    @test mydot_avx(ao, bo) ≈ s
    @test dot_unroll2avx(a, b) ≈ s
    @test dot_unroll3avx(a, b) ≈ s
    @test dot_unroll2_avx(a, b) ≈ s
    @test dot_unroll3_avx(a, b) ≈ s
    @test dot_unroll2avx_noinline(a, b) ≈ s
    @test dot_unroll3avx_inline(a, b) ≈ s
    s = myselfdot(a)
    @test myselfdotavx(a) ≈ s
    @test myselfdotavx_v2(a) ≈ s
    @test myselfdotavx_range(a) ≈ s
    @test myselfdot_avx(a) ≈ s
    @test myselfdotavx(a) ≈ s

    A = OffsetArray(rand(T, 37, 61), -5, 10)
    s = myselfdot(A)
    @test myselfdotavx(A) ≈ myselfdotavx(A') ≈ s
    @test myselfdotavx_v2(A) ≈ myselfdotavx_v2(A') ≈ s
    # @test myselfdot_avx(A) ≈ myselfdot_avx(A') ≈ s

    @test dot17(a, b) ≈ @view(a[1:17])' * @view(b[1:17])
    @test dot33(a, b) ≈ @view(a[1:33])' * @view(b[1:33])

    if T <: Union{Float32,Float64}
      πest = mcpi(a, b)
      @test πest ≈ mcpiavx(a, b)
      @test πest ≈ mcpiavx_u4(a, b)
      @test πest ≈ mcpi_avx(a, b)
      @test πest ≈ mcpi_avx_u4(a, b)
    end

    if !(
      !Bool(LoopVectorization.VectorizationBase.has_feature(Val(:x86_64_avx2))) &&
      T === Int32
    )
      @test dotloopinductvarpow(a) ≈ dotloopinductvarpowavx(a)
    end
    @test dot_from_n_to_100(a, b, 33) == @views mydotavx(a[33:100], b[33:100])

    a_re = rand(R, N)
    a_im = rand(R, N)
    b_re = rand(R, N)
    b_im = rand(R, N)
    ac = Complex.(a_re, a_im)
    bc = Complex.(b_re, b_im)

    @test mydot(ac, bc) ≈ complex_dot_soa(a_re, a_im, b_re, b_im)

    c_re1 = similar(a_re)
    c_im1 = similar(a_im)
    c_re2 = similar(a_re)
    c_im2 = similar(a_im)
    # b_re = rand(R, length(a_re) + 1); b_im = rand(R, length(a_im) + 1);
    complex_mul_with_index_offset!(c_re1, c_im1, a_re, a_im, b_re, b_im)
    complex_mul_with_index_offsetavx!(c_re2, c_im2, a_re, a_im, b_re, b_im)
    c_re1v, c_im1v, c_re2v, c_im2v =
      @views c_re1[1:end-1], c_im1[1:end-1], c_re2[1:end-1], c_im2[1:end-1]
    @test c_re1v ≈ c_re2v
    @test c_im1v ≈ c_im2v
    fill!(c_re2, -999999)
    fill!(c_im2, 99999999)
    complex_mul_with_index_offset_avx!(c_re2, c_im2, a_re, a_im, b_re, b_im)
    @test c_re1v ≈ c_re2v
    @test c_im1v ≈ c_im2v

    # Float32 is not accurate enough
    # Alternatively, loosen approx requirement?
    R == Float32 && continue
    A = rand(R, N, N, N)
    B = rand(R, N, N, N)
    @test mydot(A, B) ≈ mydotavx(A, B)
    # test CartesianIndices
    for i ∈ [3, :, 1:N-1], j ∈ [5, :, 1:N-2], k ∈ [:, 1:N-3]
      Av = view(A, i, j, k)
      Bv = view(B, i, j, k)
      # @show i, j, k
      @test mydot(Av, Bv) ≈ mydotavx(Av, Bv)
    end
  end
end
