using FFTW, DSP, Test

!(@__DIR__() in LOAD_PATH) && push!(LOAD_PATH, @__DIR__)
using FilterTestHelpers

# not exported, but used for some internal testing
using DSP.Windows: makewindow
@testset "makewindow" begin
    # make sure we're evaluating the given function at the correct points for
    # all combinations of arguments
    @test makewindow(identity, 6, 0, false) ≈ [-0.5, -0.3, -0.1, 0.1, 0.3, 0.5]
    win = makewindow(identity, 6, 0, true)
    # doesn't matter whether the middle value is evaluated at +/- 0.5
    @test win ≈ [0.0, 1/6, 1/3, 1/2, -1/3, -1/6] || win ≈ [0.0, 1/6, 1/3, -1/2, -1/3, -1/6]
    # we actually only end up with one real zero here, because one of the
    # "padding" values is used to split the -1/2 from +1/2. For windows that go
    # to zero padding will actually add `padding` zeros, because evaluating at
    # +/- 0.5 will be zero.
    @test makewindow(identity, 6, 2, false) ≈ [-0.5, -0.3, -0.1, 0.1, 0.3, 0.5, 0.0, 0.0]
    @test makewindow(identity, 6, 2, true) ≈ [0.0, 1/6, 1/3, 1/2, 0.0, -1/2, -1/3, -1/6]
    @test makewindow(identity, 5, 0, false) ≈ [-0.5, -0.25, 0.0, 0.25, 0.5]
    @test makewindow(identity, 5, 0, true) ≈ [0.0, 0.2, 0.4, -0.4, -0.2]
    @test makewindow(identity, 5, 2, false) ≈ [-0.5, -0.25, 0.0, 0.25, 0.5, 0.0, 0.0]
    @test makewindow(identity, 5, 2, true) ≈ [0.0, 0.2, 0.4, 0.0, 0.0, -0.4, -0.2]

    @test makewindow(x->42.0, 1, 0, false) ≈ [42.0]
    @test makewindow(x->42.0, 1, 0, true) ≈ [42.0]
    @test makewindow(x->42.0, 1, 2, false) ≈ [42.0, 0.0, 0.0]
    @test makewindow(x->42.0, 1, 2, true) ≈ [42.0, 0.0, 0.0]
end

@testset "dspss" begin
    # Test dpss against dpss computed with MATLAB
    d1 = dpss(128, 4)
    d2 = read_reference_data("dpss128,4.txt")
    @test d1 ≈ d2

    # Test dpsseig against dpss from MATLAB
    lambda = [0.9999999997159923,0.9999999731146645,0.9999988168667646,
              0.9999680890685374,0.9994167543397652,0.9925560207018469,
              0.9368556668429153]
    @test dpsseig(d1, 4) ≈ lambda
end

@testset "common windows" begin
    # Checking Hanning, Hamming, Triangular, Bartlett, Bartlett-Hann, Tukey,
    # Blackman, Blackman-Harris-4term and Nuttall-4term windows against values
    # computed with MATLAB. Lanczos and cosine are checked against values
    # generated with DSP.jl v0.4.0 and Blackman-Harris-3term as well as
    # Nuttall-3term with DSP.jl v0.8.5 to test for regressions, as there's no
    # reference MATLAB implementation. Gaussian is compared against DSP.jl from
    # commit da1b195, when the implementation was corrected (see GH issue #204)
    @test rect(128) == ones(128)

    hanning_jl = hanning(128)
    hann_jl = hann(128)
    hanning_ml = read_reference_data("hanning128.txt")
    @test hanning_jl ≈ hanning_ml
    @test hann_jl ≈ hanning_ml

    hamming_jl = hamming(128)
    hamming_ml = read_reference_data("hamming128.txt")
    @test hamming_jl ≈ hamming_ml

    triang_jl = triang(128)
    triang_ml = read_reference_data("triang128.txt")
    @test triang_jl ≈ triang_ml

    # for odd `n` the `triang` window should be the middle n-2 samples of the
    # bartlett
    @test triang(5) ≈ bartlett(7)[2:6]

    bartlett_jl = bartlett(128)
    bartlett_ml = read_reference_data("bartlett128.txt")
    @test bartlett_jl ≈ bartlett_ml

    barthann_jl = bartlett_hann(128)
    barthann_ml = read_reference_data("bartlett_hann128.txt")
    @test bartlett_jl ≈ bartlett_ml

    blackman_jl = blackman(128)
    blackman_ml = read_reference_data("blackman128.txt")
    @test blackman_jl ≈ blackman_ml
    @test minimum(blackman_jl) == 0.0

    blackmanharris_3term_jl = blackmanharris(128, 3)
    blackmanharris_3term_ref = read_reference_data("blackmanharris_3term_128.txt")
    @test blackmanharris_3term_jl ≈ blackmanharris_3term_ref

    blackmanharris_4term_jl = blackmanharris(128)
    blackmanharris_4term_ml = read_reference_data("blackmanharris_4term_128.txt")
    @test blackmanharris_4term_jl ≈ blackmanharris_4term_ml

    nuttall_3term_jl = nuttall(128, 3)
    nuttall_3term_ref = read_reference_data("nuttall_3term_128.txt")
    @test nuttall_3term_jl ≈ nuttall_3term_ref

    nuttall_4term_jl = nuttall(128)
    nuttall_4term_ml = read_reference_data("nuttall_4term_128.txt")
    @test nuttall_4term_jl ≈ nuttall_4term_ml

    kaiser_jl = kaiser(128, 0.4/π)
    kaiser_ml = read_reference_data("kaiser128,0.4.txt")
    @test kaiser_jl ≈ kaiser_ml

    flattop_jl = flattop(128)
    flattop_ml = read_reference_data("flattop.txt")
    @test flattop_jl ≈ flattop_ml

    gaussian_jl = gaussian(128, 0.2)
    gaussian_ref = read_reference_data("gaussian128,0.2.txt")
    @test gaussian_jl ≈ gaussian_ref

    tukey_jl = tukey(128, 0.4)
    tukey_ml = read_reference_data("tukey128,0.4.txt")
    @test tukey_jl ≈ tukey_ml

    @test tukey(128, 0) == rect(128)

    lanczos_jl = lanczos(128)
    lanczos_ref = read_reference_data("lanczos128.txt")
    @test lanczos_jl ≈ lanczos_ref

    cosine_jl = cosine(128)
    cosine_ref = read_reference_data("cosine128.txt")
    @test cosine_jl ≈ cosine_ref

    @test_throws ArgumentError blackmanharris(128, 2)
    @test_throws ArgumentError nuttall(128, 2)
end

zeroarg_wins = [rect, hanning, hamming, cosine, lanczos,
    bartlett, bartlett_hann, blackman, triang, flattop]
onearg_wins = [gaussian, kaiser, tukey]
termarg_wins = [blackmanharris, nuttall]
@testset "zero-phase windows" begin
    for winf in zeroarg_wins
        if winf == triang
            # triang needs to be special-cased here because it has different
            # definitions for odd and even `n` (and also the underlying
            # continuous function changes with `n`), so it doesn't match the
            # zerophase assumption below
            @test triang(6, zerophase=true) ≈ [1.0, 0.75, 0.5, 0.25, 0.5, 0.75]
            @test triang(7, zerophase=true) ≈ [1.0, 0.75, 0.5, 0.25, 0.25, 0.5, 0.75]
        else
            @test winf(8, zerophase=true) ≈ ifftshift(winf(9)[1:8])
            # the `zerophase=false` version doesn't hit the center point, but the
            # zerophase one does, so this ends up introducing a 1/2-sample shift.
            @test winf(9, zerophase=true) ≈ ifftshift(winf(19)[2:2:end])
        end
    end

    # test the window functions that need extra args
    for winf in onearg_wins
        @test winf(8, 0.5, zerophase=true) == ifftshift(winf(9, 0.5)[1:8])
        @test winf(9, 0.5, zerophase=true) == ifftshift(winf(19, 0.5)[2:2:end])
    end

    # Test the window functions that accept the `term` arg (number of coefficients)
    for winf in termarg_wins
        @test winf(8, 4, zerophase=true) == ifftshift(winf(9, 4)[1:8])
        @test winf(9, 4, zerophase=true) == ifftshift(winf(19, 4)[2:2:end])
    end
    # 3-term blackmanharris and nuttall are special cases because of their odd
    # number of coefficients. These coefficients are generated with DSP.jl
    # v0.8.5 to test for regressions, as there's no reference MATLAB
    # implementation.
    @test blackmanharris(6, 3, zerophase=true) ≈ [1.0, 0.632395, 0.134845, 0.0049, 0.134845, 0.632395]
    @test nuttall(6, 3, zerophase=true) ≈ [1.0, 0.63391075, 0.13657015, 0.0053188, 0.13657015, 0.63391075]

    @test dpss(8, 2, 1, zerophase=true)[:] == ifftshift(dpss(9, 2, 1)[1:8])
    # odd-length zerophase dpss windows not currently supported
    @test_throws ArgumentError dpss(9, 2, 1, zerophase=true)
end

@testset "window return types" begin
    # make sure return types are correct
    n = 12
    ft = typeof(1.0)
    for W in(:rect, :hanning, :hamming, :cosine, :lanczos, :triang, :bartlett, :bartlett_hann, :blackman)
        @eval @test Vector{$ft} == typeof($W($n)) && length($W($n)) == $n
    end

    @test Vector{ft} == typeof(tukey(n, 0.4)) && length(tukey(n, 0.4)) == n
    @test Vector{ft} == typeof(gaussian(n, 0.4)) && length(gaussian(n, 0.4)) == n
    @test Vector{ft} == typeof(kaiser(n, 0.4)) && length(kaiser(n, 0.4)) == n
    @test Vector{ft} == typeof(blackmanharris(n, 4)) && length(blackmanharris(n, 4)) == n
    @test Vector{ft} == typeof(blackmanharris(n, 3)) && length(blackmanharris(n, 3)) == n
    @test Vector{ft} == typeof(nuttall(n, 4)) && length(nuttall(n, 4)) == n
    @test Vector{ft} == typeof(nuttall(n, 3)) && length(nuttall(n, 3)) == n

    @test Matrix{ft} == typeof(dpss(n, 1.5)) && size(dpss(n, 1.5),1) == n   # size(,2) depends on the parameters
end

@testset "tensor product windows" begin
    # test all combinations of arguments. Each arg and kwarg can be not present,
    # a single value, or a 2-tuple
    function push_args!((w1e, w2e, w3e), arg, ::Type{T}, f) where {T}
        if arg isa T
            push!(w1e.args, f(arg))
            push!(w2e.args, f(arg))
            push!(w3e.args, f(arg))
        elseif arg isa Tuple
            push!(w1e.args, f(arg[1]))
            push!(w2e.args, f(arg[2]))
            push!(w3e.args, f(arg))
        end
    end

    function test_matrix_window(winfs, args)
        for winf in winfs,
             arg in args
            for padding in (nothing, 4, (4,5)),
              zerophase in (nothing, true, (true,false))
                w1_expr = :($winf(15))
                w2_expr = :($winf(20))
                w3_expr = :($winf((15,20)))
                w_all = (w1_expr, w2_expr, w3_expr)

                push_args!(w_all, arg, Real, identity)
                push_args!(w_all, padding, Integer, s -> Expr(:kw, :padding, s))
                push_args!(w_all, zerophase, Bool, s -> Expr(:kw, :zerophase, s))

                w1 = eval(w1_expr)
                w2 = eval(w2_expr)
                w3 = eval(w3_expr)
                @test w3 ≈ w1 * w2'
            end
        end
    end

    test_matrix_window(zeroarg_wins, (nothing,))
    test_matrix_window(onearg_wins, (0.4, (0.4, 0.5)))
    test_matrix_window(termarg_wins, (3, 4))
end
