text:
using SparseArrays, DiffEqOperators, LinearAlgebra, Random,
    Test, BandedMatrices, FillArrays

# Analytic solutions to higher order operators.
# Do not modify unless you are completely certain of the changes.

# 4th derivative, 4th order
function fourth_deriv_approx_stencil(N)
    A = zeros(N, N + 2)
    A[1, 1:8] = [3.5 -56 / 3 42.5 -54.0 251 / 6 -20.0 5.5 -2 / 3]
    A[2, 1:8] = [2 / 3 -11 / 6 0.0 31 / 6 -22 / 3 4.5 -4 / 3 1 / 6]

    A[N-1, N-5:end] = reverse([2 / 3 -11 / 6 0.0 31 / 6 -22 / 3 4.5 -4 / 3 1 / 6], dims=2)
    A[N, N-5:end] = reverse([3.5 -56 / 3 42.5 -54.0 251 / 6 -20.0 5.5 -2 / 3], dims=2)

    for i in 3:N-2
        A[i, i-2:i+4] = [-1 / 6 2.0 -13 / 2 28 / 3 -13 / 2 2.0 -1 / 6]
    end
    return A
end

function second_deriv_fourth_approx_stencil(N)
    A = zeros(N, N + 2)
    A[1, 1:6] = [5 / 6 -15 / 12 -1 / 3 7 / 6 -6 / 12 5 / 60]
    A[N, N-3:end] = [1 / 12 -6 / 12 14 / 12 -4 / 12 -15 / 12 10 / 12]
    for i in 2:N-1
        A[i, i-1:i+3] = [-1 / 12 4 / 3 -5 / 2 4 / 3 -1 / 12]
    end
    return A
end

function second_derivative_stencil(N)
    A = zeros(N, N + 2)
    for i in 1:N, j in 1:N+2
        (j - i == 0 || j - i == 2) && (A[i, j] = 1)
        j - i == 1 && (A[i, j] = -2)
    end
    A
end

# analytically derived central difference operators for the irregular grid:
# [0.0, 0.08, 0.1, 0.15, 0.19, 0.26, 0.29], and spacing:
# dx = [0.08, 0.02, 0.05, 0.04, 0.07, 0.03]
# with convention that the first number is the derivative order
# and the second is the approximation order

function analyticCtrOneTwoIrr()
    A = zeros(5, 7)
    A[1, 1:3] = [-5.0 / 2.0, -75.0 / 2.0, 40.0]
    A[2, 2:4] = [-250.0 / 7.0, 30.0, 40.0 / 7.0]
    A[3, 3:5] = [-80.0 / 9.0, -5.0, 125.0 / 9.0]
    A[4, 4:6] = [-1225.0 / 77.0, 825.0 / 77.0, 400.0 / 77.0]
    A[5, 5:7] = [-30.0 / 7.0, -400.0 / 21.0, 70.0 / 3.0]
    A
end

function analyticCtrTwoTwoIrr()
    A = zeros(5, 7)
    A[1, 1:3] = [250.0, -1250.0, 1000.0]
    A[2, 2:4] = [10000.0 / 7.0, -2000.0, 4000.0 / 7.0]
    A[3, 3:5] = [4000.0 / 9.0, -1000.0, 5000.0 / 9.0]
    A[4, 4:6] = [454.0 + 42.0 / 77.0, -5000.0 / 7.0, 20000.0 / 77.0]
    A[5, 5:7] = [6000.0 / 21.0, -20000.0 / 21.0, 2000.0 / 3.0]
    A
end

function analyticCtrTwoFourIrr()
    A = zeros(3, 7)
    A[1, 1:5] = [14.0 + 12012.0 / 13167.0, 1542.0 + 2736.0 / 13167.0,
        -2288.0 - 11704.0 / 13167.0, 838.0 + 1254.0 / 13167.0,
        -106.0 - 4298.0 / 13167.0]
    A[2, 2:6] = [-223.0 - 461.0 / 693.0, 847.0 + 154.0 / 693.0,
        -1311.0 - 477.0 / 693.0, 699.0 + 593.0 / 693.0,
        -11.0 - 502.0 / 693.0]
    A[3, 3:7] = [2.0 + 12166.0 / 13167.0, 538.0 + 12654.0 / 13167.0,
        -912.0 - 9196.0 / 13167.0, 508.0 + 8664.0 / 13167.0,
        -137.0 - 11121.0 / 13167.0]
    A
end

function analyticCtrFourTwoIrr()
    A = zeros(3, 7)
    A[1, 1:5] = [462000000.0 / 4389.0,
        -8550000000.0 / 4389.0,
        11704000000.0 / 4389.0,
        -5016000000.0 / 4389.0,
        1400000000.0 / 4389.0]
    A[2, 2:6] = [200000000.0 / 231.0,
        -385000000.0 / 231.0,
        360000000.0 / 231.0,
        -200000000.0 / 231.0,
        25000000.0 / 231.0]
    A[3, 3:7] = [770000000.0 / 4389.0,
        -3420000000.0 / 4389.0,
        4180000000.0 / 4389.0,
        -2850000000.0 / 4389.0,
        1320000000 / 4389.0]
    A
end

function convert_by_multiplication(::Type{Array}, A::AbstractDerivativeOperator{T}, N::Int=A.dimension) where {T}
    @assert N >= A.stencil_length # stencil must be able to fit in the matrix
    mat = zeros(T, (N, N + 2))
    v = zeros(T, N + 2)
    for i = 1:N+2
        v[i] = one(T)
        #=
            calculating the effect on a unit vector to get the matrix of transformation
            to get the vector in the new vector space.
        =#
        mul!(view(mat, :, i), A, v)
        v[i] = zero(T)
    end
    return mat
end

# Tests the corrrectness of stencils, along with concretization.
# Do not modify the following test-set unless you are completely certain of your changes.
@testset "Correctness of Uniform Stencils" begin
    N = 20
    L = CenteredDifference(4, 4, 1.0, N)
    correct = fourth_deriv_approx_stencil(N)

    # Check that stencils (according to convert_by_multiplication) agree with correct
    @test convert_by_multiplication(Array, L, N) ≈ correct

    # Check that concretization agrees correct
    @test Array(L) ≈ correct
    @test sparse(L) ≈ correct
    @test BandedMatrix(L) ≈ correct

    L = CenteredDifference(2, 4, 1.0, N)
    correct = second_deriv_fourth_approx_stencil(N)

    # Check that stencils (according to convert_by_multiplication) agree with correct
    @test convert_by_multiplication(Array, L, N) ≈ correct

    # Check that concretization agrees correct
    @test Array(L) ≈ correct
    @test sparse(L) ≈ correct
    @test BandedMatrix(L) ≈ correct
end

@testset "Correctness of Uniform Stencils, Complete" begin
    weights = []

    push!(weights, ([-0.5, 0, 0.5], [1.0, -2.0, 1.0], [-1 / 2, 1.0, 0.0, -1.0, 1 / 2]))
    push!(weights, ([1 / 12, -2 / 3, 0, 2 / 3, -1 / 12], [-1 / 12, 4 / 3, -5 / 2, 4 / 3, -1 / 12], [1 / 8, -1.0, 13 / 8, 0.0, -13 / 8, 1.0, -1 / 8]))

    for d in 1:3
        for (i, a) in enumerate([2, 4])
            D = CompleteCenteredDifference(d, a, 1.0)

            @test all(isapprox.(D.stencil_coefs, weights[i][d], atol=1e-10))
        end
    end
end

@testset "Correctness of Non-Uniform Stencils, Complete" begin
    weights = []

    push!(weights, ([-0.5, 0, 0.5], [1.0, -2.0, 1.0], [-1 / 2, 1.0, 0.0, -1.0, 1 / 2]))
    push!(weights, ([1 / 12, -2 / 3, 0, 2 / 3, -1 / 12], [-1 / 12, 4 / 3, -5 / 2, 4 / 3, -1 / 12], [1 / 8, -1.0, 13 / 8, 0.0, -13 / 8, 1.0, -1 / 8]))

    for d in 1:3
        for (i, a) in enumerate([2, 4])
            D = CompleteCenteredDifference(d, a, 0.0:1.0:10.0)
            @show D.stencil_coefs

            @test all(isapprox.(D.stencil_coefs[end], weights[i][d], atol=1e-10))
        end
    end
end

@testset "Correctness of Uniform Stencils, Complete Half" begin
    weights = (([0.5, 0.5], [-1 / 16, 9 / 16, 9 / 16, -1 / 16]),
        ([-1.0, 1.0], [1 / 24, -9 / 8, 9 / 8, -1 / 24]))
    for (i, a) in enumerate([2, 4])
        for d in 0:1
            D = CompleteHalfCenteredDifference(d, a, 1.0)
            @test all(D.stencil_coefs .≈ weights[d+1][i])
        end
    end

end

@testset "Correctness of Non-Uniform Stencils, Complete Half" begin
    weights = (([0.5, 0.5], [-1 / 16, 9 / 16, 9 / 16, -1 / 16]),
        ([-1.0, 1.0], [1 / 24, -9 / 8, 9 / 8, -1 / 24]))
    for (i, a) in enumerate([2, 4])
        for d in 0:1
            D = CompleteHalfCenteredDifference(d, a, 0.0:1.0:5.0)
            @test all(D.stencil_coefs[end] .≈ weights[d+1][i])
        end
    end

end

@testset "Correctness of Uniform Upwind Stencils" begin
    weights = (
        (
            [-1.0, 1.0],
            [-1.0, 3.0, -3.0, 1.0]
        ),
        (
            [-3 / 2, 2.0, -1 / 2],
            [-5 / 2, 9.0, -12.0, 7.0, -3 / 2]
        )
    )
    for (i, a) in enumerate(1:2)
        for (j, d) in enumerate([1, 3])
            D = CompleteUpwindDifference(d, a, 1.0, 0)
            @test all(isapprox.(D.stencil_coefs, weights[i][j], atol=1e-10))
        end
    end
end

@testset "Correctness of Non-Uniform Upwind Stencils" begin
    weights = (
        (
            [-1.0, 1.0],
            [-1.0, 3.0, -3.0, 1.0]
        ),
        (
            [-3 / 2, 2.0, -1 / 2],
            [-5 / 2, 9.0, -12.0, 7.0, -3 / 2]
        )
    )
    for (i, a) in enumerate(1:2)
        for (j, d) in enumerate([1, 3])
            D = CompleteUpwindDifference(d, a, 0.0:1.0:10.0, 0)
            @test all(isapprox.(D.stencil_coefs[end], weights[i][j], atol=1e-10))
        end
    end
end

@testset "Correctness of Non-Uniform Stencils" begin
    x = [0.0, 0.08, 0.1, 0.15, 0.19, 0.26, 0.29]
    nx = length(x)
    dx = diff(x)

    # Second-Order First Derivative
    L = CenteredDifference(1, 2, dx, nx - 2)
    correct = analyticCtrOneTwoIrr()

    # Check that stencils agree with correct
    for (i, coefs) in enumerate(L.stencil_coefs)
        @test Array(coefs) ≈ correct[i, correct[i, :].!=0.0]
    end
    @test Array(L) ≈ correct
    @test sparse(L) ≈ correct
    @test BandedMatrix(L) ≈ correct

    # Second-Order Second Derivative
    L = CenteredDifference(2, 2, dx, nx - 2)
    correct = analyticCtrTwoTwoIrr()

    # Check that stencils agree with correct
    for (i, coefs) in enumerate(L.stencil_coefs)
        @test Array(coefs) ≈ correct[i, correct[i, :].!=0.0]
    end
    @test Array(L) ≈ correct
    @test sparse(L) ≈ correct
    @test BandedMatrix(L) ≈ correct

    # Fourth-Order Second Derivative
    L = CenteredDifference(2, 4, dx, nx - 2)
    correct = analyticCtrTwoFourIrr()

    # Check that stencils agree with correct
    for (i, coefs) in enumerate(L.stencil_coefs)
        @test Array(coefs) ≈ correct[i, correct[i, :].!=0.0]
    end
    @test Array(L)[2:end-1, :] ≈ correct
    @test sparse(L)[2:end-1, :] ≈ correct
    @test BandedMatrix(L)[2:end-1, :] ≈ correct

    # Second-Order Fourth Derivative
    L = CenteredDifference(4, 2, dx, nx - 2)
    correct = analyticCtrFourTwoIrr()

    # Check that stencils agree with correct
    for (i, coefs) in enumerate(L.stencil_coefs)
        @test Array(coefs) ≈ correct[i, correct[i, :].!=0.0]
    end
    @test Array(L)[2:end-1, :] ≈ correct
    @test sparse(L)[2:end-1, :] ≈ correct
    @test BandedMatrix(L)[2:end-1, :] ≈ correct
end

# tests for full and sparse function
@testset "Full and Sparse functions:" begin
    N = 10
    d_order = 2
    approx_order = 2
    correct = second_derivative_stencil(N)
    A = CenteredDifference(d_order, approx_order, 1.0, N)

    @test convert_by_multiplication(Array, A, N) == correct
    @test Array(A) == second_derivative_stencil(N)
    @test sparse(A) == second_derivative_stencil(N)
    @test BandedMatrix(A) == second_derivative_stencil(N)
    @test opnorm(A, Inf) == opnorm(correct, Inf)

    # testing higher derivative and approximation concretization
    N = 20
    d_order = 4
    approx_order = 4
    A = CenteredDifference(d_order, approx_order, 1.0, N)
    correct = convert_by_multiplication(Array, A, N)

    @test Array(A) ≈ correct
    @test sparse(A) ≈ correct
    @test BandedMatrix(A) ≈ correct

    N = 26
    d_order = 8
    approx_order = 8
    A = CenteredDifference(d_order, approx_order, 1.0, N)
    correct = convert_by_multiplication(Array, A, N)

    @test Array(A) ≈ correct
    @test sparse(A) ≈ correct
    @test BandedMatrix(A) ≈ correct

    # testing correctness of multiplication
    N = 1000
    d_order = 4
    approx_order = 10
    y = collect(1:1.0:N+2) .^ 4 - 2 * collect(1:1.0:N+2) .^ 3 + collect(1:1.0:N+2) .^ 2
    y = convert(Array{BigFloat,1}, y)

    A = CenteredDifference(d_order, approx_order, one(BigFloat), N)
    correct = convert_by_multiplication(Array, A, N)
    @test Array(A) ≈ correct
    @test sparse(A) ≈ correct
    @test BandedMatrix(A) ≈ correct
    @test A * y ≈ Array(A) * y
end

@testset "Indexing tests" begin
    N = 1000
    d_order = 4
    approx_order = 10

    A = CenteredDifference(d_order, approx_order, 1.0, N)
    @test A[1, 1] == Array(A)[1, 1]
    @test A[10, 20] == 0

    correct = Array(A)
    for i in 1:N
        @test A[i, i] == correct[i, i]
    end

    # Indexing Tests
    N = 1000
    d_order = 2
    approx_order = 2

    A = CenteredDifference(d_order, approx_order, 1.0, N)
    M = Array(A, 1000)
    @test A[1, 1] == M[1, 1]
    @test A[1:4, 1] == M[1:4, 1]
    @test A[5, 2:10] == M[5, 2:10]
    @test A[60:100, 500:600] == M[60:100, 500:600]
end

@testset "Operations on matrices" begin
    N = 51
    M = 101
    d_order = 2
    approx_order = 2

    xarrs = Vector{Union{Array,StepRangeLen}}(undef, 0)
    yarrs = Vector{Union{Array,StepRangeLen}}(undef, 0)
    push!(xarrs, range(0, stop=1, length=N))
    push!(yarrs, range(0, stop=1, length=M))
    push!(xarrs, vcat(xarrs[1][1:floor(Int, N / 2)] .^ 0.2020, xarrs[1][ceil(Int, N / 2):end] .^ 2.015))
    push!(yarrs, vcat(yarrs[1][1:floor(Int, M / 2)] .^ 1.793, yarrs[1][ceil(Int, M / 2):end] .^ 2.019))

    for (i, xarr, yarr) in zip(1:2, xarrs, yarrs)
        if i == 1
            dx = xarr[2] - xarr[1]
            dy = yarr[2] - yarr[1]
        elseif i == 2
            dx = diff(xarr)
            dy = diff(yarr)
        end
        F = [x^2 + y for x = xarr, y = yarr]

        A = CenteredDifference(d_order, approx_order, dx, length(xarr) - 2)
        B = i == 1 ? CenteredDifference(d_order, approx_order, dy, length(yarr)) :
            CenteredDifference(d_order, approx_order, dy, length(yarr) - 2)

        @test A * F ≈ 2 * ones(N - 2, M) atol = 1e-9

        # Operators are defined such that
        # B must be applied to F from the left
        @test all(abs.(B * F') .<= 1e-8)
        @test all(abs.(A * (B * F')') .<= 1e-4)

        G = [x^2 + y^2 for x = xarr, y = yarr]

        @test A * G ≈ 2 * ones(N - 2, M) atol = 1e-9
        @test B * G' ≈ 2 * ones(M - 2, N) atol = 1e-8
        @test all(abs.(A * (B * G')') .<= 1e-4)
    end
end

@testset "Scalar multiplication with operators" begin
    A = CenteredDifference(2, 2, 10.0, 3)

    scalar_ans = [0.033 -0.066 0.033 0.0 0.0
        0.0 0.033 -0.066 0.033 0.0
        0.0 0.0 0.033 -0.066 0.033]
    @test Array(3.3 * A) ≈ scalar_ans
end

@testset "Left-multiplying operators with a vector of coefficients" begin
    A = CenteredDifference(2, 2, 1.0, 3)
    B = UpwindDifference(1, 1, 1.0, 3, 1.0)
    C = UpwindDifference(1, 1, 1.0, 3, [2.0, 3.0, 4.0])
    c = [1.0, 1.0, 3.0]
    cA = c * A
    cB = c * B
    cC = c * C

    @test cA.coefficients == c
    @test cB.coefficients == c
    @test cC.coefficients == [2.0, 3.0, 12.0]

    x = [1.0, 2.0, 3.0, 4.0, 5.0]
    @test (A * x) .* c ≈ cA * x
    @test (B * x) .* c ≈ cB * x
    @test (C * x) .* c ≈ cC * x

    @test_throws DimensionMismatch ones(10) * A
    @test_throws DimensionMismatch ones(10) * cA
    @test_throws DimensionMismatch ones(10) * B
    @test_throws DimensionMismatch ones(10) * C
end

# These tests are broken due to the implementation 2.2*LD creating a DerivativeOperator
# rather than an Array
#=
@testset "Linear combinations of operators" begin
    N = 10
    Random.seed!(0); LA = DiffEqArrayOperator(rand(N,N+2))
    LD = CenteredDifference(2,2,1.0,N)
    L = 1.1*LA - 2.2*LD + 3.3*Eye(N,N+2)
    # Builds convert(L) the brute-force way
    fullL = zeros(N,N+2)
    v = zeros(N+2)
    for i = 1:N+2
        v[i] = 1.0
        fullL[:,i] = L*v
        v[i] = 0.0
    end
    @test_broken convert(AbstractMatrix,L) ≈ fullL
    for p in [1,2,Inf]
        @test_broken opnorm(L,p) ≈ opnorm(fullL,p)
    end
end
=#


text:
"""
    module Mesh

Mesh data structures and the interface expected from an `AbstractMesh`.

Exports: $(EXPORTS)
"""

module Mesh

using DocStringExtensions
using StaticArrays
using LinearAlgebra
using ForwardDiff # for computing derivatives of parametric elements
using OrderedCollections

using WaveProp.Utils
using WaveProp.Geometry
using WaveProp.Integration

import WaveProp.Geometry: ambient_dimension, geometric_dimension, domain, range, boundary, parametrization
import WaveProp.Integration: qnodes, qweights, qnormals, integrate, normal, jacobian, singular_quadrature

export GenericMesh, ElementIterator, nodes, etypes, elements, ent2tags, el2qnodes, elements, near_interaction_list, nodes,
compute_quadrature!,SubMesh, AbstractMesh, CartesianMesh, meshgen, NystromMesh, dof, AbstractElement, LagrangeLine, LagrangeTriangle, LagrangeRectangle, LagrangeTetrahedron, LagrangeElement, ParametricElement, derivative, derivative2, measure, convert_to_2d, decompose

include("element.jl")
include("meshes.jl")
include("submesh.jl")
include("cartesianmesh.jl")
include("meshgen.jl")
include("decompose.jl")
include("elementiterator.jl")

end

ext: jl

max_stars_repo_name: WaveProp/WaveProp

max_issues_repo_name: WaveProp/WaveProp

max_forks_repo_name: WaveProp/WaveProp


text:
# TODO: set up our own tfuncs

function CC.builtin_tfunction(interp::JETInterpreter, @nospecialize(f), argtypes::Array{Any,1},
                              sv::InferenceState) # `JETInterpreter` isn't overloaded on `return_type`
    ret = @invoke builtin_tfunction(interp::AbstractInterpreter, f, argtypes::Array{Any,1},
                                    sv::Union{InferenceState,Nothing})

    if f === throw
        # here we only report a selection of "serious" exceptions, i.e. those should be
        # reported even if they may be caught in actual execution;
        # other general `throw` calls will be reported within `_typeinf(interp::JETInterpreter, frame::InferenceState)`
        # only when they are not caught by control flow, which is judged by whether if the
        # final return type of `sv` is annotated as `Bottom` or not
        if length(argtypes) ≥ 1
            a = first(argtypes)
            if isa(a, Const)
                v = a.val
                if isa(v, UndefKeywordError)
                    report!(interp, UndefKeywordErrorReport(interp, sv, v, get_lin(sv)))
                end
            end
        end
        return ret
    elseif f === getfield
        # `getfield` is so common, let's special case it
        if 2 ≤ length(argtypes) ≤ 3
            obj, fld = argtypes
            if isa(fld, Const)
                name = fld.val
                if isa(name, Symbol)
                    if isa(obj, Const) && (mod = obj.val; isa(mod, Module))
                        if isdefined(mod, name)
                            if istoplevel_globalref(interp, sv)
                                # when accessing to a global variable in a module concretized by `interp`,
                                # take a risk and eagerly propagate its type
                                # NOTE logic here should be synced with that of `abstract_eval_special_value(::JETInterpreter, ::Any, ::VarTable, ::InferenceState)`
                                val = getfield(mod, name)
                                return isa(val, AbstractGlobal) ? val.t : Const(val)
                            end
                            # TODO; `ret` should be `Any` here, add report pass here (for performance linting)
                        else
                            # report access to undefined global variable
                            report!(interp, GlobalUndefVarErrorReport(interp, sv, mod, name))
                            # return Bottom
                        end
                    elseif ret === Bottom
                        # general case when an error is detected by the native `getfield_tfunc`
                        typ = widenconst(obj)
                        report!(interp, NoFieldErrorReport(interp, sv, typ, name))
                        return ret
                    end
                end
            end
        end
    elseif f === fieldtype
        # the valid widest possible return type of `fieldtype_tfunc` is `Union{Type,TypeVar}`
        # because fields of unwrapped `DataType`s can legally be `TypeVar`s,
        # but this will cause lots of false positive `NoMethodErrorReport`s for inference
        # with accessing to abstract fields since most methods don't expect `TypeVar`
        # (e.g. `@report_call readuntil(stdin, 'c')`)
        # JET.jl further widens this case to `Any` and give up further analysis rather than
        # trying hard to do sound and noisy analysis
        # xref: https://github.com/JuliaLang/julia/pull/38148
        if ret === Union{Type, TypeVar}
            return Any
        end
    elseif length(argtypes) == 2 && begin
            f === Intrinsics.checked_sdiv_int ||
            f === Intrinsics.checked_srem_int ||
            f === Intrinsics.checked_udiv_int ||
            f === Intrinsics.checked_urem_int
        end
        # `DivideError` for these intrinsics are handled in C and should be special cased
        a = argtypes[2]
        t = widenconst(a)
        if t <: Base.BitSigned64 || t <: Base.BitUnsigned64
            if isa(a, Const) && a.val === zero(t)
                report!(interp, DivideErrorReport(interp, sv))
                return Bottom
            end
        end
    end

    if ret === Bottom
        # XXX: for general case, JET just relies on the (maybe too persmissive) return type
        # from native tfuncs to report invalid builtin calls and probably there're lots of
        # false negatives
        report!(interp, InvalidBuiltinCallErrorReport(interp, sv, argtypes))
    end

    return ret
end

# check if this frame is for `getproperty(::Module, ::Symbol)`, which accesses to a global
# variable traced by `interp`
function istoplevel_globalref(interp::JETInterpreter, sv::InferenceState)
    def = sv.linfo.def
    def.name === :getproperty || return false
    def.sig === Tuple{typeof(getproperty), Module, Symbol} || return false
    parent = sv.parent
    return !isnothing(parent) && istoplevel(interp, parent)
end

# `return_type_tfunc` internally uses `abstract_call` to model `return_type` function and
# here we shouldn't pass `JETInterpreter` to it; otherwise we may get false error reports from
# the  `abstract_call`, which isn't the abstraction of actual execution, thus here we just
# check if the call of `return_type` is valid or not
function CC.return_type_tfunc(interp::JETInterpreter, argtypes::Vector{Any}, sv::InferenceState)
    if length(argtypes) ≠ 3
        # invalid argument number, let's report and return error result (i.e. `Bottom`)
        report!(interp, NoMethodErrorReport(interp,
                                            sv,
                                            false,
                                            # this is not necessary to be computed correctly, though
                                            argtypes_to_type(argtypes),
                                            ))
        @static if isdefined(CC, :ReturnTypeCallInfo)
            return CallMeta(Bottom, nothing)
        else
            return Bottom
        end
    else
        # don't recursively pass on `JETInterpreter` via `@invoke`
        return return_type_tfunc(interp.native, argtypes, sv)
    end
end


text:
module Audio

# Audio modules (to load and monitor data)
# import LibSndFile
using PortAudio		  # Handle audio
using FileIO
using Suppressor	  # Disable warning from PortAudio
using DSP

# --- Calling audio functions
#include("functions_play.jl");
export  play, play44,  play2, playbis
export createAudioBuffer, play

""" play
---
Play a sound based on input buffer
# --- Syntax
  play(buffer)
# --- Input parameters
- buffer  : Vector to be played [Array{Int}]
# --- Output parameters
- []
# ---
# v 1.0 - Robin Gerzaguet.
"""
function play(buffer)
	# --- Due to macro, need to create Any stream
	devices = PortAudio.devices()
	stream = Any;
	# --- Suprress to avoid underflow warnings
	@suppress let
		# --- Create default streamer,
		stream = PortAudioStream();
	end
	# --- Playing file through device
	@suppress write(stream,buffer);
	close(stream);
	# We need to close the streamer.
end

function createAudioBuffer()
	# --- Due to macro, need to create Any stream
	devices = PortAudio.devices()
	stream = Any;
	# --- Suprress to avoid underflow warnings
	@suppress let
		# --- Create default streamer,
		stream = PortAudioStream();
	end
	return stream;
end
function play(stream,buffer)
	 write(stream,buffer);
end

""" play
---
Play a sound based on input buffer with a 2kHz-> 44kHz upscaling
# --- Syntax
  play(buffer)
# --- Input parameters
- buffer  : Vector to be played [Array{Int}]
# --- Output parameters
- []
# ---
# v 1.0 - Corentin Lavaud
"""
function play44(buffer,fe)
	# --- Due to macro, need to create Any stream
	devices = PortAudio.devices()
	stream = Any;
	buf2 = to44(buffer,fe)
	# --- Suprress to avoid underflow warnings
	@suppress let
		# --- Create default streamer,
		stream = PortAudioStream();
	end
	# --- Playing file through device
	@suppress write(stream,buf2);
	close(stream);
	# We need to close the streamer.
end

""" playbis
---
Play a sound based on input buffer
# --- Syntax
  play2(buffer)
# --- Input parameters
- buffer  : Vector to be played [Array{Int}]
# --- Output parameters
- []
# ---
# v 1.0 - Corentin Lavaud.
"""
function playbis(buffer)
	# config
	buf_size = 256
	sample_rate = 44100;
	# simply play a single buffer
	#play(buffer, sample_rate);
	# use default device
	devID = -1
	# or retrieve a specific device by name
	#devID = PortAudio.find_device("default")

	# open a stream and write audio to the output
	stream = open(devID, (0, 2), sample_rate, buf_size)
	write(stream, buffer)
	close(stream)
end

"""
---
Resample input audio buffer of frequency fe to 44100 Hz for play() routine
--- Syntax
y = to44(x,fe)
# --- Input parameters
- x	  : Input signal @fe
- fe  : Audio sampling frequency of x
# --- Output parameters
- y	  : Audio buffer from x @44100 Hz
# ---
# v 1.0 - Robin Gerzaguet.
# """
function to44(sig,fe)

	rate = (44100 / fe);
	return resample(sig,rate);
end

end


text:
############################################################################################
"""
$(SIGNATURES)
Callable struct to make initializing PMCMC sampler easier in sampling library.

# Examples
```julia
```

"""
struct PMCMCConstructor{P,F<:ParticleFilterConstructor,M<:MCMCConstructor,D<:PMCMCDefault} <:
       AbstractConstructor
    "Valid PCMC Kernel."
    kernel::P
    "Particle Filter Constructor."
    filter::F
    "MCMC Constructor"
    mcmc::M
    "All other relevant keywords."
    default::D
    function PMCMCConstructor(
        kernel::Type{P}, filter::F, mcmc::M, default::D
    ) where {P<:PMCMCKernel,F<:ParticleFilterConstructor,M<:MCMCConstructor,D<:PMCMCDefault}
        return new{typeof(kernel),F,M,D}(kernel, filter, mcmc, default)
    end
end
function (constructor::PMCMCConstructor)(
    _rng::Random.AbstractRNG,
    model::ModelWrapper,
    data::D,
    temperature::F,
    info::BaytesCore.SampleDefault
) where {D, F<:AbstractFloat}
    return PMCMC(
        _rng,
        constructor.kernel,
        constructor.filter(_rng, model, data, temperature, info),
        constructor.mcmc(_rng, model, data, temperature, info),
        constructor.default,
        info
    )
end
function PMCMC(
    kernel::Type{P}, filter::F, mcmc::M; kwargs...
) where {P<:PMCMCKernel,F<:ParticleFilterConstructor,M<:MCMCConstructor}
    return PMCMCConstructor(kernel, filter, mcmc, PMCMCDefault(; kwargs...))
end

function get_sym(constructor::PMCMCConstructor)
    #return get_sym(constructor.mcmc)
    ## Obtain symbols from mcmc and particle filter
    sym1 = get_sym(constructor.mcmc)
    sym2 = get_sym(constructor.filter)
    ## Return unique symbols as tuple
    sym = unique((sym1..., sym2...))
    return Tuple(sym)
end

############################################################################################
"""
$(SIGNATURES)
Infer PMCMC diagnostics type.

# Examples
```julia
```

"""
function infer(
    _rng::Random.AbstractRNG,
    diagnostics::Type{AbstractDiagnostics},
    pmcmc::PMCMC,
    model::ModelWrapper,
    data::D,
) where {D}
    TPrediction = infer(_rng, pmcmc.kernel.pf, model, data)
    pfdiag = infer(_rng, diagnostics, pmcmc.kernel.pf, model, data)
    mcmcdiag = infer(_rng, diagnostics, pmcmc.kernel.mcmc, model, data)
    return PMCMCDiagnostics{TPrediction, pfdiag, mcmcdiag}
end

"""
$(SIGNATURES)
Infer type of predictions of kernel.

# Examples
```julia
```

"""
function infer(
    _rng::Random.AbstractRNG, pmcmc::PMCMC, model::ModelWrapper, data::D
) where {D}
    return infer(_rng, pmcmc.kernel.pf, model, data)
end

############################################################################################
function results(
    diagnosticsᵛ::AbstractVector{M}, pmcmc::PMCMC, Ndigits::Integer, quantiles::Vector{T}
) where {T<:Real,M<:PMCMCDiagnostics}
    ## Print MCMC diagnostics
    results(
        [diagnosticsᵛ[iter].mcmc for iter in eachindex(diagnosticsᵛ)],
        pmcmc.kernel.mcmc,
        Ndigits,
        quantiles,
    )
    ## Print PF diagnostics
    results(
        [diagnosticsᵛ[iter].pf for iter in eachindex(diagnosticsᵛ)],
        pmcmc.kernel.pf,
        Ndigits,
        quantiles,
    )
    return nothing
end

############################################################################################
function result!(pmcmc::PMCMC, result::L) where {L<:ℓObjectiveResult}
    result!(pmcmc.kernel.mcmc, result)
    return nothing
end
function get_result(pmcmc::PMCMC)
    return get_result(pmcmc.kernel.mcmc)
end
function predict(_rng::Random.AbstractRNG, kernel::PMCMCKernel, objective::Objective)
    return predict(_rng, kernel.pf, objective)
end
function predict(_rng::Random.AbstractRNG, pmcmc::PMCMC, objective::Objective)
    return predict(_rng, pmcmc.kernel, objective)
end

############################################################################################
#export
export PMCMCConstructor, infer


text:
struct MDBMcontainer{RTf,RTc,AT}
    funval::RTf
    cval::RTc
    callargs::AT
end

#TODO memoization for multiple functions Vector{Function}
struct MemF{fT,cT,RTf,RTc,AT} <:Function
    f::fT
    c::cT
    fvalarg::Vector{MDBMcontainer{RTf,RTc,AT}}#funval,callargs
    memoryacc::Vector{Int64} #number of function value call for already evaluated parameters
    MemF(f::fT,c::cT,cont::Vector{MDBMcontainer{RTf,RTc,AT}}) where {fT,cT,RTf,RTc,AT}=new{fT,cT,RTf,RTc,AT}(f,c,cont,[Int64(0)])
end

(memfun::MemF{fT,cT,RTf,RTc,AT})(::Type{RTf},::Type{RTc},args...,) where {fT,cT,RTf,RTc,AT} =( memfun.f(args...,)::RTf, memfun.c(args...,)::RTc)

function (memfun::MemF{fT,cT,RTf,RTc,AT})(args...,) where {fT,cT,RTf,RTc,AT}
    location=searchsortedfirst(memfun.fvalarg,args,lt=(x,y)->isless(x.callargs,y));
    if length(memfun.fvalarg)<location
        x=memfun(RTf,RTc,args...,);
        push!(memfun.fvalarg,MDBMcontainer{RTf,RTc,AT}(x...,args))
        return x
    elseif  memfun.fvalarg[location].callargs!=args
        x=memfun(RTf,RTc,args...,);
        insert!(memfun.fvalarg, location, MDBMcontainer{RTf,RTc,AT}(x...,args))
        return x
    else
        memfun.memoryacc[1]+=1;
        return (memfun.fvalarg[location].funval,memfun.fvalarg[location].cval);
    end
end

function (memfun::MemF{fT,cT,RTf,RTc,AT})(args::Tuple) where {fT,cT,RTf,RTc,AT}
    memfun(args...,)
end

"""
    Axis{T}

Represent the parameter space in a discretized grid (vector) in `ticks`.
"""
struct Axis{T} <: AbstractVector{T}
    ticks::Vector{T}
    name
    function Axis(T::Type,a::AbstractVector,name=:unknown)
        new{T}(T.(a), name)
    end
end
Base.getindex(ax::Axis{T}, ind) where T = ax.ticks[ind]::T
Base.setindex!(ax::Axis, X, inds...) = setindex!(ax.ticks, X, inds...)
Base.size(ax::Axis) = size(ax.ticks)

function Axis(a::AbstractVector{T}, name=:unknown) where T<:Real
    Axis(Float64, a, name)
end
function Axis(a::AbstractVector{T}, name=:unknown) where T
    Axis(T, a, name)
end
function Axis(a, name=:unknown) where T
    Axis([a...], name)
end
function Axis(a::Axis)
    a
end

struct Axes{N,aT}
    axs::aT
end
function Axes(axs...)
    axstosave=Axis.(axs)
    Axes{length(axstosave),typeof(axstosave)}(Axis.(axstosave))
end
(axs::Axes)(ind...) = getindex.(axs.axs,ind)
# (axs::Axes)(ind::AbstractVector{<:Integer}) = axs(ind...)
Base.getindex(axs::Axes,inds...) = axs.axs[inds...]
Base.getindex(axs::Axes,inds::Vector{Integer}) = axs.axs[inds]
Base.iterate(axs::Axes) = iterate(axs.axs)
Base.iterate(axs::Axes,i) = iterate(axs.axs,i)
Base.size(::Axes{N,aT}) where {N,aT} = (N,)
Base.size(::Axes{N,aT},::Integer) where {N,aT} = N
Base.length(::Axes{N,aT}) where {N,aT} = N

struct NCube{IT,FT,N}
    corner::MVector{N,IT} #"bottom-left" #Integer index of the axis
    size::MVector{N,IT}#Integer index of the axis
    posinterp::MVector{N,FT}#relative coordinate within the cube "(-1:1)" range
    bracketingncube::Bool
    # gradient ::MVector{MVector{T}}
    # curvnorm::Vector{T}
end

Base.isless(a::NCube{IT,FT,N},b::NCube{IT,FT,N}) where IT where FT where N = Base.isless([a.corner,a.size],[b.corner,b.size])
Base.isequal(a::NCube{IT,FT,N},b::NCube{IT,FT,N}) where IT where FT where N = all([a.corner==b.corner,a.size==b.size])
import Base.==
==(a::NCube{IT,FT,N},b::NCube{IT,FT,N}) where IT where FT where N = all([a.corner==b.corner,a.size==b.size])

"""
    struct MDBM_Problem{fcT,N,Nf,Nc,t01T,t11T,IT,FT}

Store the main data for the Multi-Dimensional Bisection Method.

# Examples
```julia
include("MDBM.jl")
using Reexport
@reexport using .MDBM

function foo(x,y)
    x^2.0+y^2.0-2.0^2.0
end
function c(x,y) #only the c>0 domain is analysed
    x-y
end

ax1=Axis([-5,-2.5,0,2.5,5],"x") # initial grid in x direction
ax2=Axis(-5:2:5.0,"y") # initial grid in y direction

mymdbm=MDBM_Problem(foo,[ax1,ax2],constraint=c)
```
"""
struct MDBM_Problem{fcT,N,Nf,Nc,t01T,t11T,IT,FT,aT}
    "(memoized) function and constraint in a Tuple"
    fc::fcT
    "The vector of discretized parameters space"
    axes::Axes{N,aT}
    "the bracketing n-cubes (which contains a solution)"
    ncubes::Vector{NCube{IT,FT,N}}
    T01::t01T#
    T11pinv::t11T
end

function MDBM_Problem(fc::fcT,axes,ncubes::Vector{NCube{IT,FT,N}},Nf,Nc) where fcT<:Function where IT<:Integer where FT<:AbstractFloat where N
    T01=T01maker(Val(N))
    T11pinv=T11pinvmaker(Val(N))
    MDBM_Problem{fcT,N,Nf,Nc,typeof(T01),typeof(T11pinv),IT,FT,typeof((axes...,))}(fc,Axes(axes...),
    [NCube{IT,FT,N}(SVector{N,IT}([x...]),SVector{N,IT}(ones(IT,length(x))),SVector{N,FT}(zeros(IT,length(x))),true) for x in Iterators.product((x->1:(length(x.ticks)-1)).(axes)...,)][:]
    ,T01,T11pinv)
end

function MDBM_Problem(f::Function, axes0::AbstractVector{<:Axis};constraint::Function=(x...,)->nothing, memoization::Bool=true,
    #Nf=length(f(getindex.(axes0,1)...)),
    Nf=f(getindex.(axes0,1)...) === nothing ? 0 : length(f(getindex.(axes0,1)...)),
    Nc=constraint(getindex.(axes0,1)...) === nothing ? 0 : length(constraint(getindex.(axes0,1)...)))#Float16(1.), nothing
    axes=deepcopy.(axes0);
    argtypesofmyfunc=map(x->typeof(x).parameters[1], axes);#Argument Type
    AT=Tuple{argtypesofmyfunc...};
    type_f=Base.return_types(f,AT)
    if length(type_f)==0
        error("input of the function is not compatible with the provided axes")
    else
        RTf=type_f[1];#Return Type of f
    end

    type_con=Base.return_types(constraint, AT)
    if length(type_con)==0
        error("input of the constraint function is not compatible with the provided axes")
    else
        RTc=type_con[1];#Return Type of the constraint function
    end

    if memoization
        fun=MemF(f, constraint,Array{MDBMcontainer{RTf,RTc,AT}}(undef, 0));
    else
        fun=(x)->(f(x...), constraint(x...));
    end
    Ndim=length(axes)
    MDBM_Problem(fun, axes, Vector{NCube{Int64,Float64,Ndim}}(undef, 0),  Nf, Nc)
end

function MDBM_Problem(f::Function, a::AbstractVector{<:AbstractVector};constraint::Function=(x...,)->nothing, memoization::Bool=true,
    #Nf=length(f(getindex.(axes0,1)...)),
    Nf=f(getindex.(a,1)...) === nothing ? 0 : length(f(getindex.(a,1)...)),
    Nc=constraint(getindex.(a,1)...) === nothing ? 0 : length(constraint(getindex.(a,1)...)))
    axes=[Axis(ax) for ax in a]
    MDBM_Problem(f, axes, constraint = onstraint, memoization=memoization, Nf=Nf, Nc=Nc)#,Vector{NCube{Int64,Float64,Val(Ndim)}}(undef, 0))
end

function Base.show(io::IO, mdbm::MDBM_Problem{fcT,N,Nf,Nc,t01T,t11T,IT,FT,aT}) where {fcT,N,Nf,Nc,t01T,t11T,IT,FT,aT}
    println(io,"Multi-Dimensional Bisection Method Problem")
    println(io,"  parameter dimension:   ", N)#typeof(mdbm).parameters[2])#N
    println(io,"  co-dimension:          ", Nf)#typeof(mdbm).parameters[3])#Nf
    println(io,"  number of constraints: ", Nc)#typeof(mdbm).parameters[4])#Nc
    println(io,"Axes:")
    for k in 1:length(mdbm.axes)
        println(io,"  axis #",k,": elements: ",length(mdbm.axes[k]),"; elementtype: ", typeof(mdbm.axes[k][1]),"; first: ", mdbm.axes[k][1],"; last: ",mdbm.axes[k][end])
    end
        println(io,"number of bracketing n-cubes: ", length(mdbm.ncubes))
        println()

    if (typeof(mdbm.fc) <: MemF)
        println(io,"number of function evaluation: ", length(mdbm.fc.fvalarg))
        println(io,"number of memoized function call: ", mdbm.fc.memoryacc[1])
        #println(io,"Ration of function evaluation compared to 'brute-force method': ", length(mdbm.fc.fvalarg)/prod(length.(mdbm.axes)))
    else
     println(io,"non-memoized version")
    end
end


text:
"""
    MMatrix(A::AbstractMatrix)

MMatrix is a type concentrating on the mathematical properties of a Matrix.
That is in opposition to wrappers like Symmetric, Adjoint, etc., which have a bias
on the internal representation.
MMatrix(Symmetric(A)) has type `MMatrix{T,SparseMatrixCSC{T},:upper,:symm,:ident,:ident}`,
which allows to specialize methods by the storage type and the symmetry property.

The type prameters after the storage type are in detail:
  - Bands (:all, :upper, :lower, :diagonal)
  - Moves (:ident, :transpose, :symmetry)
  - Transform (:ident, :conjugate)
  - Diag (:ident, :hermitian, :ones)

Some are restricted by element type of the matrix, so :conjugate and :hermitian are not
supported for Real arguments. :diagonal implies Moves(:ident).
Other combinations of type parameters describe differnt combinations of standard wrappers.
TODO: We will see how the open semigroup generated by the standard wrappers looks like.
"""
struct MMatrix{Tv,N,S<:AbstractArray{<:Any,N},B,M,T,D} <:AbstractArray{Tv,N}
    parent::S
end

const StorageMatrix{Tv} = Union{StridedArray{Tv},SparseMatrixCSC{Tv}}
const StorageArray{Tv} = Union{StorageMatrix{Tv},SparseVector{Tv}}

Base.parent(ma::MMatrix) = wparent(ma)
Base.show(io::IO, ma::MMatrix) = show(io, wparent(ma))

# MMatrix of storage types
MMatrix(a::S) where {Tv,S<:StorageMatrix{Tv}} = MMatrix{Tv,2,S,:all,:ident,:ident,:ident}(a)
MMatrix(a::S) where {Tv,S<:SparseVector{Tv}} = MMatrix{Tv,1,S,:all,:ident,:ident,:ident}(a)
function MMatrix(a::S) where {Tv,X<:StorageArray{Tv},S<:Transpose{Tv,X}}
    T<:Real || throw(ArgumentError("transpose of complex not supported"))
    return MMatrix{T,2,X,:all,:transpose,:ident,:ident}(parent(a))
end
function MMatrix(a::S) where {Tv<:Real,N,X<:StorageArray{Tv},S<:Adjoint{Tv,X}}
    MMatrix{T,N,X,:all,:transpose,:ident,:ident}(parent(a))
end
function MMatrix(a::S) where {Tv,N,X<:StorageArray{Tv},S<:Adjoint{Tv,X}}
    MMatrix{T,N,X,:all,:transpose,:conjugate,:hermitian}(parent(a))
end
function MMatrix(a::S) where {Tv,X<:StorageMatrix,S<:Symmetric{Tv,X}}
    T<:Real || throw(ArgumentError("symmetric of complex not supported"))
    uplo = a.uplo == 'U' ? :upper : lower
    MMatrix{Tv,2,X,uplo,:symmetric,:ident,:ident}(parent(a))
end
function MMatrix(a::S) where {Tv,X<:StorageMatrix,S<:Hermitian{Tv,X}}
    if ! (Tv<:Real)
        isreal(diag(a)) || throw(ArgumentError("Hermitian needs real diagonal"))
        uplo = a.uplo == 'U' ? :upper : lower
        MMatrix{Tv,2,X,uplo,:symmetric,:conjugate,:hermitian}(parent(a))
    else
        uplo = a.uplo == 'U' ? :upper : lower
        MMatrix{Tv,2,X,uplo,:symmetric,:ident,:ident}(parent(a))
    end
end
function MMatrix(a::S) where {Tv,X<:StorageMatrix,S<:AbstractTriangular{Tv,X}}
    uplo = S isa Union{UpperTriangular,UnitUpperTriangular} ? :upper : :lower
    unit = S isa Union{UnitUpperTriangular,UnitLowerTriangular} :ones : :ident
    MMatrix{Tv,2,X,uplo,:ident,:ident,unit}(parent(a))
end

# MMatrix of wrapped types
for wr in (Transpose,Adjoint,Symmetric,Hermitian,
           UpperTriangular,LowerTriangular,UnitUpperTriangular,UnitLowerTriangular)

    @eval MMatrix(a::S) where {Tv,X,S<:$wr{Tv,X}} = MMatrix(MMatrix(parent(a)))
end
# conversions to base types
# Matrix uses AbstractMatrix
SparseMatrixCSC{Tv,Ti}(a::MMatrix) where {Tv,Ti} = SparseMatrixCSC{Tv,Ti}(wparent(a))
unwrap(a::MMatrix) = unwrap(wparent(a))
iswrsparse(a::MMatrix) = iswrsparse(a.parent)

Base.convert(::Type{MMatrix}, a::AbstractMatrix) = MMatrix(a)
Base.eltype(::MMatrix{T}) where T = T
mmtype(::MMatrix{<:Any,<:Any,B,M,T,D}) where {B,M,T,D} = (B,M,T,D)
mmstorage_type(::MMatrix{<:Any,X}) where X = X

wparent(a::MMatrix{<:Any,<:Any,<:Any,:all,:ident,:ident,:ident}) = a.parent
wparent(a::MMatrix{<:Any,<:Any,<:Any,:all,:transpose,:ident,:ident}) = transpose(a.parent)
wparent(a::MMatrix{<:Any,<:Any,<:Any,:all,:transpose,:conjugate,:ident}) = adjoint(a.parent)
wparent(a::MMatrix{<:Any,<:Any,<:Any,:upper,:symmetry,:ident,:ident}) = symmetric(a.parent)
wparent(a::MMatrix{<:Any,<:Any,<:Any,:lower,:symmetry,:ident,:ident}) = symmetric(a.parent, :L)
wparent(a::MMatrix{<:Any,<:Any,<:Any,:upper,:symmetry,:conjugate,:hermitian}) = hermitian(a.parent)
wparent(a::MMatrix{<:Any,<:Any,<:Any,:lower,:symmetry,:conjugate,:hermitian}) = hermitian(a.parent, :L)
wparent(a::MMatrix{<:Any,<:Any,<:Any,:upper,:ident,:ident,:ident}) = UpperTriangular(a.parent)
wparent(a::MMatrix{<:Any,<:Any,<:Any,:lower,:ident,:ident,:ident}) = LowerTriangular(a.parent)
wparent(a::MMatrix{<:Any,<:Any,<:Any,:upper,:ident,:ident,:ones}) = UnitUpperTriangular(a.parent)
wparent(a::MMatrix{<:Any,<:Any,<:Any,:lower,:ident,:ident,:ones}) = UnitLowerTriangular(a.parent)

# inversion
adjoint(a::MMatrix) = MMatrix(adjoint(parent(a)))
transpose(a::MMatrix) = MMatrix(transpose(parent(a)))
symmetric(a::MMatrix, s::Symbol=:U) = MMatrix(symmetric(parent(a), s))
hermitian(a::MMatrix, s::Symbol=:U) = MMatrix(hermitian(parent(a), s))

# AbstractArray interface
Base.getindex(a::MMatrix, i...) = getindex(parent(a), i...)
Base.setindex!(a::MMatrix, v, i...) = setindex!(parent(a), v, i...)
Base.size(a::MMatrix) = size(parent(a))

# fallback to Matrix or SparseMatrixCSC
import Base: +, -, *, /, \
for op in (:+, :-, :*, :/, :\)
    for ty in (Number, AbstractVector, AbstractMatrix)
        @eval ($op)(a::MMatrix, b::$ty) = MMatrix(($op)(unwrap(a), unwrap(b)))
    end
    for ty in (Number, Adjoint{<:Any,<:AbstractVector}, Transpose{<:Any,<:AbstractVector}, AbstractMatrix)
        @eval ($op)(a::$ty, b::MMatrix) = MMatrix(($op)(unwrap(a), unwrap(b)))
    end
    @eval ($op)(a::MMatrix, b::MMatrix) = MMatrix(($op)(unwrap(a), unwrap(b)))
end


text:
# Power spectra for three orbits of Henon-Heiles
using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using DynamicalSystems, PyPlot

u0s = (
    [0.0, -0.25, 0.42, 0.0], # chaotic
    [0.0, 0.1, 0.5, 0.0], # quasiperiodic
    [0.0, 0.30266571044921875, 0.4205654433900762, 0.0], # periodic
)

labels = (
    "chaotic",
    "quasiperiodic",
    "periodic",
)

hh = Systems.henonheiles()
fig, axs = subplots(3, 1; figsize = (figx/2, 2figy), sharex = true)

δt = 0.05

using FFTW, Statistics
for (i, u) in enumerate(u0s)
   # r = trajectory(hh, 1000.0, u; dt = 0.1)[:, 1]
    r = trajectory(hh, 30000.0, u; dt = 0.01, Δt = δt, Ttr = 10)[:, 1]
    P = abs2.(rfft(r .- mean(r)))
    P[1] = P[2]
    ν = rfftfreq(length(r))/δt
    # axs[i].plot(ν, P ./ maximum(P),
    axs[i].semilogy(ν, P ./ maximum(P),
    label = labels[i], linewidth = i == 1 ? 1.0 : 2.0, color = "C$(i-1)")
    axs[i].text(0.99, 0.8, labels[i]; ha = "right", transform = axs[i].transAxes,
    color = "C$(i-1)")
    # @show std(r)
    # r .+= 0.5randn(length(r))
    # P = abs2.(rfft(r .- mean(r)))
    # PyPlot.plot(10rfftfreq(length(r)), P ./ maximum(P),
    # lw = 1.0, alpha = 0.5, color = "C$(i-1)")
    # PyPlot.plot(r)
    # axs[i].set_yticks([])
    axs[i].set_ylim(10.0^(-4.5), 1.0)
    axs[i].set_yticks(10.0 .^ (0:-2:-5))
    axs[i].set_xlim(0, 0.4)
end
axs[2].set_ylabel("\$P / \\mathrm{max}(P)\$")
axs[3].set_xlabel("frequency \$\\nu\$")
fig.tight_layout(pad=0.3)
wsave(plotsdir("2", "spectra"), fig)

max_stars_repo_name: JuliaDynamics/NonlinearDynamicsTextbook

max_issues_repo_name: JuliaDynamics/NonlinearDynamicsTextbook

max_forks_repo_name: JuliaDynamics/NonlinearDynamicsTextbook


text:
"""
    inject(::T, v; kwargs...) where {T <: AbstractInjectLocation}

Inject the value `v` to the specific location in the log record.  If `v` is callable
(e.g. function) then it will be evalated at runtime before injecting into the log.
"""
function inject end

function inject(log, loc::T, v; sep = " ") where {T <: MessageLocation}
    value = v isa Base.Callable ? v() : v
    msg = if T <: BeginningMessageLocation && length(log.message) > 0
              "$value$sep$(log.message)"
          elseif T <: EndingMessageLocation && length(log.message) > 0
              "$(log.message)$sep$value"
          else
              "$value"
          end
    merge(log, (message = msg,))
end

function inject(log, ::KwargsLocation, v)
    value = v isa Base.Callable ? v() : v
    updated_kwargs = (value..., log.kwargs...)
    merge(log, (kwargs = updated_kwargs,))
end

function inject(log, ::LevelLocation, v)
    value = v isa Base.Callable ? v() : v
    merge(log, (level = value,))
end

# Removing stuffs

"""
    remove(prop::T) where {T <: AbstractLogProperty}

Remove either `message` or `kwargs` data from the log record.  The `prop`
argument can be either `MessageProperty()` or `KwargsProperty`.
"""
function remove end

remove(log, ::KwargsProperty) = merge(log, (kwargs = (),))
remove(log, ::MessageProperty) = merge(log, (message = "",))
remove(log, ::LevelProperty) = error("Level property cannot be removed")

# Migrating stuffs

kv_string(kwargs, sep, divider) = join(["$k$sep$v" for (k,v) in kwargs], divider)

"""
    migrate(::MessageProperty, ::KwargsProperty; label = :message)

Migrate the message string to kwargs location with key `label`.
"""
function migrate(log, ::MessageProperty, ::KwargsProperty; label = :message)
    return @pipe log |>
                 inject(_, KwargsLocation(), (label => _.message,)) |>
                 remove(_, MessageProperty())
end

"""
    migrate(::KwargsProperty, ::MessageProperty; sep = "=", divider = " ", prepend = " ", transform::Function)

Migrate all kwargs to the end of message string.  By default kwargs will be formatted as key=value
pairs separated by a space (specified as `divider`). However, a custom `transform` function may be
passed for custom formatting.
"""
function migrate(log, ::KwargsProperty, ::MessageProperty;
                 sep = "=", divider = " ", prepend = " ",
                 transform = (kwargs) -> kv_string(kwargs, sep, divider))
    return @pipe log |>
                 inject(_, EndingMessageLocation(), transform(_.kwargs); sep = prepend) |>
                 remove(_, KwargsProperty())
end

"""
    migrate(::LevelProperty, ::MessageProperty; loc = EndingMessageLocation(), transform::Function = string)

Migrate the log level to the message field at the `loc` location.  A custom transform function
may be specified should a different format is desired.
"""
function migrate(log, ::LevelProperty, ::MessageProperty;
                 location = BeginningMessageLocation(),
                 transform::T = string) where {T <: Function}
    return inject(log, location, transform(log.level))
end

"""
    migrate(::LevelProperty, ::KwargsProperty; label = :level, transform::Function = string)

Migrate the log level to the kwargs field with key `label`.  A custom transform function
may be specified should a different format is desired.
"""
function migrate(log, ::LevelProperty, ::KwargsProperty;
                 label = :level,
                 transform::T = identity) where {T <: Function}
    return inject(log, KwargsLocation(), (label => transform(log.level),))
end

"""
    mutate(::AbstractLogProperty; transform::Function)

Mutate the log property by applying `transform` function over it.  The transformation
function will be passed with an immutable log record.
"""
function mutate end

function mutate(log, ::MessageProperty;
                transform::T) where {T <: Function}
    return @pipe log |>
                 remove(_, MessageProperty()) |>
                 inject(_, BeginningMessageLocation(), transform(log))
end

function mutate(log, ::LevelProperty;
                transform::T) where {T <: Function}
    inject(log, LevelLocation(), () -> transform(log))
end

function mutate(log, ::KwargsProperty;
                transform::T) where {T <: Function}
    return @pipe log |>
                 remove(_, KwargsProperty()) |>
                 inject(_, KwargsLocation(), () -> transform(log))
end


text:
"""
IteratorStyle

Construction and reading of a ChunkedStream is facilitated by the IteratorStyle.
* DefinedIterator: assumes all iterators are defined at construction
* ConstructOnIterator: assumes that each chunk is defined upon iteration which
    entails appending the `chunks` field with the

If ConstructOnIterator is specified then the user must define an iterator function
that constructs subsequent iterations (i.e., appends the `chunks` field with each
subsequent `ImageStream`.
"""
abstract type IteratorStyle end
struct DefinedIterator <: IteratorStyle end  # all iterators are defined before streaming begins
struct ConstructOnIterator <: IteratorStyle end  # iteration leads to construction of next ImageStream

mutable struct ChunkedStream{S,T,N,I} <: AbstractArrayStream{S,T,N,I}
    chunks::AbstractVector{<:ArrayStream}
    nitr::Int
    itr::Int
end

getchunk(s::ChunkedStream) = s.chunks[s.itr]
stream(s::ChunkedStream) = stream(getchunk(s))
streamindices(s::ChunkedStream) = streamindices(getchunk(s))
needswap(s::ChunkedStream) = needswap(getchunk(s))
ownstream(s::ChunkedStream) = ownstream(getchunk(s))

function chunk(s::ImageStream{S,T,N,Ax}, dim::Int) where {S,T,N,Ax}
    ncolon = ntuple(i->Colon(), dim-1)
    chunk_indices = CartesianIndices(map(i->length(axes(s,i)), dim:N))
    ChunkedStream([s[ncolon..., i.I...] for i in chunk_indices], length(chunk_indices), 1)
end

########
# Read #
########
function Base.read!(s::ChunkedStream, sink::A; kwargs...) where {A<:AbstractArray}
    read!(getchunk(s), A)
    iterate(s)
end

function Base.read(s::ChunkedStream)
    ret = read(getchunk(s))
    iterate(s)
    return ret
end

#########
# Write #
#########
Base.write(s::ChunkedStream, sink::A) where {A<:AbstractArray} = write(getchunk(s), sink)


text:
using TermWin

trun( ()->begin
    n = 2000
    for i=1:n
        TermWin.progressUpdate( i/n )
        TermWin.progressMessage( ( @sprintf( "%d/%d ", i, n ) ) * strftime( "%H:%M:%S", time() ) )
        if mod(i,10)==0
            sleep(50.0/n)
        end
    end
    n
end )


text:
module equations

using StatsModels, MixedModels, CategoricalArrays, CSV, StatsBase, DataStructures, DataFrames

include("misc.jl")

function make_ran_matrix(x1::AbstractVector,x2::AbstractVector)
        isa(x1, CategoricalArray) ||
                       throw(ArgumentError("ran() only works with CategoricalArrays (got $(typeof(2)))"))
        isa(x2, CategoricalArray) ||
                       throw(ArgumentError("ran() only works with CategoricalArrays (got $(typeof(2)))"))

        u = unique(x2);
        filter!(x->x≠0,u)
        Z = Matrix{Bool}(undef, length(x1), length(u))
        for i in eachindex(u)
        	@. Z[:, i] = x1 .== u[i]
        end
           return unique(x1),Z
       end

ranMat(arg1,arg2,data1,data2) = make_ran_matrix(data1[!,Symbol(arg1)],data2[!,Symbol(arg2)])

function mme(f, userHints, userData; blocks,path2ped,paths2geno)
        terms4StatsModels = String.(split(repr(f.rhs), ('+')))
        terms4StatsModels = replace.(terms4StatsModels, ":" => "")
        terms4StatsModels = [filter(x -> !isspace(x), trm) for trm in terms4StatsModels]

        yVec = StatsModels.modelmatrix(f.lhs, userData)

        FE = OrderedDict{Any,Any}() #any to block work

        RE = OrderedDict{Any,Any}()

	ME = OrderedDict{Any,Any}()
	regionSizes = OrderedDict{String,Int64}()

        #read pedigree
	if isempty(path2ped)
		A = []
	else
		pedigree = CSV.read(path2ped,DataFrame)

		pedigree.ID  = CategoricalArray(pedigree.ID)
		pedigree.Dam = CategoricalArray(pedigree.Dam)
		pedigree.Sire = CategoricalArray(pedigree.Sire)

		A = makeA(pedigree[!,:Sire],pedigree[!,:Dam])
	end

	#column id within pedigree
	idRE = []

        for i in 1:length(f.rhs)
		if (f.rhs[i] isa FunctionTerm) && (String(nameof(f.rhs[i].forig)) == "PR")
			println("$(terms4StatsModels[i]) is BayesPR Type")
			arg1 = repr((f.rhs[i].args_parsed)[1])
			arg2 = parse(Int64,repr((f.rhs[i].args_parsed)[2]))
			path = paths2geno[Symbol(arg1)]
			thisM = CSV.read(path,CSV.Tables.matrix)
			#centering
			thisM .-= mean(thisM,dims=1)
			println("size of $arg1 data: $(size(thisM))")
			println("region size for $arg1: $arg2")
			ME[arg1] = thisM
                        thisM = 0 #I can directly merge to dict above
			regionSizes[arg1] = arg2
                elseif (f.rhs[i] isa FunctionTerm) && (String(nameof(f.rhs[i].forig)) == "ran")
                        println("$(terms4StatsModels[i]) is ran Type")
                        sym1 = repr((f.rhs[i].args_parsed)[1]) #now it is Symbol
                        sym2 = repr((f.rhs[i].args_parsed)[2]) #now it is from string
                        println("sym1: $sym1 sym2: $sym2")

			IDs,thisZ = ranMat(sym1, sym2, userData, pedigree)
			RE[(sym1,sym2)] = thisZ
			thisZ = 0
			push!(idRE,IDs)
                elseif (f.rhs[i] isa FunctionTerm) && (String(nameof(f.rhs[i].forig)) == "|")
                        println("$(terms4StatsModels[i]) is | Type")
                        my_sch = schema(userData, userHints) #work on userData and userHints
                        my_ApplySch = apply_schema(terms(f.rhs[i]), my_sch, MixedModels.MixedModel)
			#####NO IDs for this effect!!! Will be added later!!!!#####################################################
                       	thisZ = modelcols(my_ApplySch, userData)
			RE[terms4StatsModels[i]] = thisZ
			thisZ = 0
			push!(RE,modelcols(my_ApplySch, userData))

                else
                println("$(terms4StatsModels[i]) is $(typeof(f.rhs[i])) type")
		thisX = StatsModels.modelmatrix(f.rhs[i], userData,hints= userHints)
		FE[terms4StatsModels[i]] = thisX
		thisX = 0
                end
        end

	#BLOCK FIXED EFFECTS
	for b in blocks
		getThese = intersect(collect(keys(FE)), b)
		FE[Tuple(getThese)] = hcat(getindex.(Ref(FE), getThese)...)
		for d in getThese
			delete!(FE,d)
		end
	end

        return idRE, A, vec(yVec), FE, RE, ME, regionSizes
        end

end


text:
# Autogenerated wrapper script for RadeonProRender_jll for aarch64-apple-darwin
export libNorthstar64, libProRenderGLTF, libRadeonProRender64, libRprLoadStore64, libTahoe64

using CompilerSupportLibraries_jll
JLLWrappers.@generate_wrapper_header("RadeonProRender")
JLLWrappers.@declare_library_product(libNorthstar64, "libNorthstar64.dylib")
JLLWrappers.@declare_library_product(libProRenderGLTF, "@rpath/libProRenderGLTF.dylib")
JLLWrappers.@declare_library_product(libRadeonProRender64, "libRadeonProRender64.dylib")
JLLWrappers.@declare_library_product(libRprLoadStore64, "libRprLoadStore64.dylib")
JLLWrappers.@declare_library_product(libTahoe64, "@rpath/libTahoe64.dylib")
function __init__()
    JLLWrappers.@generate_init_header(CompilerSupportLibraries_jll)
    JLLWrappers.@init_library_product(
        libNorthstar64,
        "lib/libNorthstar64.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libProRenderGLTF,
        "lib/libProRenderGLTF.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libRadeonProRender64,
        "lib/libRadeonProRender64.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libRprLoadStore64,
        "lib/libRprLoadStore64.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libTahoe64,
        "lib/libTahoe64.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
### A Pluto.jl notebook ###
# v0.16.1

using Markdown
using InteractiveUtils

# ╔═╡ 926e4270-ee52-11eb-15b9-0b59dac4c681
using Pkg; Pkg.develop(path=".."); Pkg.activate("..")

# ╔═╡ 1d009b9d-dcf2-4bdf-8dd8-fef135cb8524
begin
	using SqState
	using QuantumStateBase
	using QuantumStatePlots
	using DataDeps
	using Plots
	using MAT
end;

# ╔═╡ 05865de0-1458-4c59-880c-8619d4c7dd83
md"
# Real time system

JingYu Ning

The real time quantum state tomogrophy system plays an important role in the homodyme experiment. It provides the features of data visualization without delay and get rid of the extremely complex data analasis works.

In this demonstration, we will show you how the data flows and how the machine learning model works.

"

# ╔═╡ d372a693-4a33-45b1-898c-72339f9e910f
md"
## A small yet elegant model

Apart from the intergration of the model implemented by HsenYi, the improvement of perforement is another essential part of real time system.

Therefore I implemented a small yet elegant CNN model that will infer the arguments `r`, `θ` and `n̄` for a squeezed thermal state to boost the inference process.
"

# ╔═╡ 7a22ff61-a32f-4def-9fd9-23a540791609
m = get_model("model")

# ╔═╡ 6f6ad264-cac8-4a04-906a-e83ebac85136
md"
## Experiment data from the lab

The experiments of the following data are accomplish by YiRu in NTHU
"

# ╔═╡ 1b4271a4-8540-4881-b6cb-7a3dd8e9d0b1
begin
	data_path = joinpath(SqState.data_path(), "Flow")
	files = readdir(data_path)
end

# ╔═╡ 22a0af68-9012-4cc1-aebf-2f91bac387b4
md"
## About the data

The data is generated by the homodyme detector and capture by the oscilloscope.
"

# ╔═╡ 55d14429-f6e6-4db3-a7ea-6aec92fccd7d
begin
	data = SqState.get_data(joinpath("Flow", files[end-1]))
	x_data = SqState.sample(data, 4096)
end

# ╔═╡ 49fe2f66-5f90-4ad3-8ab3-846b6684b9e7
scatter(LinRange(0, 2π, 4096), x_data, size=(800, 400), title="Quadrature data")

# ╔═╡ 8bb50bb4-dbae-4adb-acd6-46dc45d8c109
md"
## Let the model do the magic

the model will infere the arguments for the state via the given quaduture data.
"

# ╔═╡ b015f3e0-89d8-4ecd-a3e5-23d9cc83a6e5
r, θ, n̄, c1, c2, c3 = SqState.infer_arg(data, 100)

# ╔═╡ d02b6dbc-a45d-4900-881a-30b550564f6d
md"
After the inference, we can construct the quantum state and plot the Wigner function, or apply some operator to it.
"

# ╔═╡ 90a4e898-f95a-4bd1-8825-9c7d8561a683
state =
	c1 * SqueezedState(ξ(r, θ), rep=StateMatrix) +
	c2 * SqueezedThermalState(ξ(r, θ), n̄) +
	c3 * ThermalState(n̄)

# ╔═╡ d7756447-61ae-49f7-b0cb-ea1ff0432d29
begin
	wf = WignerFunction(LinRange(-3, 3, 101), LinRange(-3, 3, 101))
	plot_wigner(wf(state), Contour)
end

# ╔═╡ 8ff17755-21eb-45c2-adbc-27a5d021f79f
md"
Calculate all args for all data
"

# ╔═╡ 870e4dd3-f8e6-45a1-87ba-721237e00066
begin
	argv = Matrix{Float64}(undef, 6, length(files))
	for i in 1:size(argv, 2)
		dataᵢ = SqState.get_data(joinpath("Flow", files[i]))
		argv[:, i:i] .= SqState.infer_arg(dataᵢ, 10)
	end
end

# ╔═╡ 90e1807a-0779-4156-9b19-44cab50e2efc
scatter(argv[1, :], size=(800, 200), legend=false, title="r")

# ╔═╡ b24b71fc-2f37-408d-9888-318eeba180dd
scatter(argv[2, :], size=(800, 200), legend=false, title="θ")

# ╔═╡ c730f57d-1b41-4165-b50c-e564fbda8861
scatter(argv[3, :], size=(800, 200), legend=false, title="n̄")

# ╔═╡ 203ee767-c38b-4123-be29-0d237dc57c40
begin
	plot(size=(800, 200), legend=:left)
	scatter!(argv[4, :], label="c₁")
	scatter!(argv[5, :], label="c₂")
	scatter!(argv[6, :], label="c₃")
end

# ╔═╡ 5bb4f85c-af93-456c-af73-2dc069d0237a
md"
## Wigner flow
"

# ╔═╡ 3fdf7f29-b8b6-478b-9b4d-fb96407e99ae
begin
	function get_w(data_name::String; n_sample=10, fix_θ=true, dim=70)
		data = SqState.get_data(joinpath("Flow", data_name))
    	state, w = SqState.calc_w(
			SqState.infer_arg(data, n_sample)..., dim, fix_θ,
			wf=wf
		)

		return w
	end

	anim = @animate for f in files
		plot_wigner(get_w(f, fix_θ=true), QuantumStatePlots.Contour)
		annotate!(-2.5, 2.5, text("$f", :left))
	end

	gif(anim, fps=2)
end

# ╔═╡ c2fcff08-cbab-47f6-b5b3-5a2989b11c91
# begin
# 	𝐰s = Array{Float64}(undef, 101, 101, 21)
# 	for (i, f) in enumerate(files)
# 		𝐰s[:, :, i] .= get_w(f, fix_θ=true).𝐰_surface
# 	end

# 	wfile = matopen("w.mat", "w")
# 	write(wfile, "ws", 𝐰s)
# 	close(wfile)
# end

# ╔═╡ 87ca3223-89cf-4dc3-aa25-f66e9bb22a2b
# begin
# 	argvfile = matopen("argv.mat", "w")
# 	write(argvfile, "argv", argv)
# 	close(argvfile)
# end

# ╔═╡ e9ea5eba-70fc-40d0-8bb4-62a289305a7c
# begin
# 	rs = [0.0952 0.0809 0.1271 0.1572 0.2652 0.2753 0.3454 0.3492]
# 	n̄s = [0.2549 0.2417 0.2283 0.2355 0.2402 0.2761 0.2793 0.2595]

# 	𝐰s_sqth = Array{Float64}(undef, 101, 101, 8)
# 	𝐰s_th = Array{Float64}(undef, 101, 101, 8)
# 	for (i, (r, n̄)) in enumerate(zip(rs, n̄s))
# 		𝐰s_sqth[:, :, i] .= wf(SqueezedThermalState(ξ(r, 0.), n̄)).𝐰_surface
# 		𝐰s_th[:, :, i] .= wf(ThermalState(n̄)).𝐰_surface
# 	end

# 	w_sqth_th_file = matopen("w_sqth_th.mat", "w")
# 	write(w_sqth_th_file, "w_sqth", 𝐰s_sqth)
# 	write(w_sqth_th_file, "w_th", 𝐰s_th)
# 	close(w_sqth_th_file)
#   heatmap(𝐰s_sqth[:, :, 8])
# end

# ╔═╡ 16f28d60-1fd4-4f2b-88c3-5f82ae48fd75
# begin
# 	rs = [0.0609 0.1382 0.3428 0.3692 0.4676 0.4820 0.5393 0.5445]

# 	𝐰s_sq = Array{Float64}(undef, 101, 101, 8)
# 	for (i, r) in enumerate(rs)
# 		𝐰s_sq[:, :, i] .= wf(SqueezedState(ξ(r, 0.))).𝐰_surface
# 	end

# 	w_sq_file = matopen("w_sq.mat", "w")
# 	write(w_sq_file, "w_sq", 𝐰s_sq)
# 	close(w_sq_file)
# 	heatmap(𝐰s_sq[:, :, 8])
# end

# ╔═╡ Cell order:
# ╟─05865de0-1458-4c59-880c-8619d4c7dd83
# ╠═926e4270-ee52-11eb-15b9-0b59dac4c681
# ╠═1d009b9d-dcf2-4bdf-8dd8-fef135cb8524
# ╟─d372a693-4a33-45b1-898c-72339f9e910f
# ╠═7a22ff61-a32f-4def-9fd9-23a540791609
# ╟─6f6ad264-cac8-4a04-906a-e83ebac85136
# ╠═1b4271a4-8540-4881-b6cb-7a3dd8e9d0b1
# ╟─22a0af68-9012-4cc1-aebf-2f91bac387b4
# ╠═55d14429-f6e6-4db3-a7ea-6aec92fccd7d
# ╠═49fe2f66-5f90-4ad3-8ab3-846b6684b9e7
# ╟─8bb50bb4-dbae-4adb-acd6-46dc45d8c109
# ╠═b015f3e0-89d8-4ecd-a3e5-23d9cc83a6e5
# ╟─d02b6dbc-a45d-4900-881a-30b550564f6d
# ╠═90a4e898-f95a-4bd1-8825-9c7d8561a683
# ╠═d7756447-61ae-49f7-b0cb-ea1ff0432d29
# ╟─8ff17755-21eb-45c2-adbc-27a5d021f79f
# ╠═870e4dd3-f8e6-45a1-87ba-721237e00066
# ╟─90e1807a-0779-4156-9b19-44cab50e2efc
# ╟─b24b71fc-2f37-408d-9888-318eeba180dd
# ╟─c730f57d-1b41-4165-b50c-e564fbda8861
# ╟─203ee767-c38b-4123-be29-0d237dc57c40
# ╟─5bb4f85c-af93-456c-af73-2dc069d0237a
# ╠═3fdf7f29-b8b6-478b-9b4d-fb96407e99ae
# ╠═c2fcff08-cbab-47f6-b5b3-5a2989b11c91
# ╠═87ca3223-89cf-4dc3-aa25-f66e9bb22a2b
# ╠═e9ea5eba-70fc-40d0-8bb4-62a289305a7c
# ╠═16f28d60-1fd4-4f2b-88c3-5f82ae48fd75


text:
import Base.show
import Base.display

include("display.jl")
include("data.jl")

using FileIO
using FileIO: @format_str

using Flux: onehotbatch, cpu
using Random

using MLDatasets

struct MyImage
    arr::AbstractArray
end

function Base.show(io::IO, ::MIME"image/png", x::MyImage)
    img = cpu(x.arr)
    size(img)[3] in [1, 3] || error("Unsupported channel size: ", size(img)[3])
    if size(img)[3] == 3
        getRGB(X) = colorview(RGB, permutedims(X, (3,1,2)))
        img = getRGB(img)
    end
    FileIO.save(FileIO.Stream(format"PNG", io), img)
end

function Base.display(d::EmacsDisplay, mime::MIME"image/png", x::MyImage)
    # path, io = Base.Filesystem.mktemp()
    # path * ".png"
    path = tempname() * ".png"
    open(path, "w") do io
        show(io, mime, x)
    end

    println("$(path)")
    println("#<Image: $(path)>")
end

# FIXME use Zygote.jl
# using Tracker: TrackedArray

function sample_and_view(x, y=nothing, model=nothing)
    if length(size(x)) < 4
        x = x[:,:,:,:]
    end
    # if typeof(x) <: Flux.Tracker.TrackedArray
    #     x = x.data
    # end
    size(x)[1] in [28,32,56] ||
        error("Image size $(size(x)[1]) not correct size. Currently support 28 or 32.")
    num = min(size(x)[4], 10)
    @info "Showing $num images .."
    if num == 0 return end
    imgs = cpu(hcat([x[:,:,:,i] for i in 1:num]...))
    # viewrepl(imgs)
    display(MyImage(imgs))
    if y != nothing
        labels = onecold(cpu(y[:,1:num]), 0:9)
        @show labels
    end
    if model != nothing
        preds = onecold(cpu(model(gpu(x))[:,1:num]), 0:9)
        @show preds
    end
    nothing
end

function load_MNIST_ds(;batch_size)
    # Ideally, it takes a batch size, and should return an iterator that repeats
    # infinitely, and shuffle before each epoch. The data should be moved to GPU
    # on demand. I prefer to just move online during training.
    train_x, train_y = MLDatasets.MNIST.traindata();
    test_x,  test_y  = MLDatasets.MNIST.testdata();
    # normalize into -1, 1
    train_x = (train_x .- 0.5) .* 2
    test_x = (test_x .- 0.5) .* 2

    # reshape to add channel, and onehot y encoding
    #
    # I'll just permute dims to make it column major
    train_ds = DataSetIterator(reshape(permutedims(train_x, [2,1,3]), 28,28,1,:),
                               onehotbatch(train_y, 0:9), batch_size);
    # FIXME test_ds should repeat only once
    # TODO add iterator interface, for x,y in ds end
    test_ds = DataSetIterator(reshape(permutedims(test_x, [2,1,3]), 28,28,1,:),
                              onehotbatch(test_y, 0:9), batch_size);
    # x, y = next_batch!(train_ds);
    return train_ds, test_ds
end

max_stars_repo_name: lihebi/DAG-EQ

max_issues_repo_name: lihebi/DAG-EQ

max_forks_repo_name: lihebi/DAG-EQ


text:
### Integrated autocorrelation time (IACT)

iact(mcvariance::N, iidvariance::N) where {N<:Real} = mcvariance/iidvariance

iact(v::AbstractVector, vtype::Symbol, args...) = iact(mcvar(v, Val{vtype}, args...), mcvar(v, Val{:iid}))

iact(v::AbstractArray, vtype::Symbol, region, args...) = mapslices(x -> iact(vec(x), vtype, args...), v, region)

iact(s::VariableNState{Univariate}, vtype::Symbol, args...) = iact(s.value, vtype, args...)

iact(s::VariableNState{Multivariate}, vtype::Symbol, i::Integer, args...) = iact(s.value[i, :], vtype, args...)

iact(s::VariableNState{Multivariate}, vtype::Symbol, r::AbstractVector=1:s.size, args...) =
  eltype(s)[iact(s, vtype, i, args...) for i in r]

iact(v::AbstractVector, args...) = iact(v, :imse, args...)

iact(v::AbstractArray, region, args...) = iact(v, :imse, region, args...)

iact(s::VariableNState{Univariate}, args...) = iact(s.value, :imse, args...)

iact(s::VariableNState{Multivariate}, i::Integer, args...) = iact(s, :imse, i, args...)

iact(s::VariableNState{Multivariate}, r::AbstractVector=1:s.size, args...) = iact(s, :imse, r, args...)


text:
is_constraint(::typeof(reactive_power_constraints)) = true

# Here, the power constraints are ordered as:
# g = [qg_gen]
function _reactive_power_constraints(
    qg, v_m, v_a, pinj, qinj, qload,
    ybus_re, ybus_im, pv, pq, ref, pv_to_gen, ref_to_gen, nbus
)
    if isa(qg, Array)
        device = KA.CPU()
        kernel! = reactive_power_kernel!(KA.CPU())
    else
        device = KA.CUDADevice()
        kernel! = reactive_power_kernel!(KA.CUDADevice())
    end
    range_ = length(pv) + length(ref)
    ev = kernel!(
        qg,
        v_m, v_a, pinj,
        pv, ref, pv_to_gen, ref_to_gen,
        ybus_re.nzval, ybus_re.colptr, ybus_re.rowval,
        ybus_im.nzval, qload,
        ndrange=range_,
        dependencies=Event(device)
    )
    wait(ev)
end

function reactive_power_constraints(polar::PolarForm, cons, buffer)
    # Refresh reactive power generation in buffer
    update!(polar, PS.Generators(), PS.ReactivePower(), buffer)
    # Constraint on Q_ref (generator) (Q_inj = Q_g - Q_load)
    copy!(cons, buffer.qg)
    return
end

# Function for AD with ForwardDiff
function reactive_power_constraints(polar::PolarForm, cons, vm, va, pbus, qbus)
    nbus = length(vm)
    pv = polar.indexing.index_pv
    pq = polar.indexing.index_pq
    ref = polar.indexing.index_ref
    pv_to_gen = polar.indexing.index_pv_to_gen
    ref_to_gen = polar.indexing.index_ref_to_gen
    ybus_re, ybus_im = get(polar.topology, PS.BusAdmittanceMatrix())
    _reactive_power_constraints(
        cons, vm, va, pbus, qbus, polar.reactive_load,
        ybus_re, ybus_im, pv, pq, ref, pv_to_gen, ref_to_gen, nbus
    )
end

function size_constraint(polar::PolarForm, ::typeof(reactive_power_constraints))
    return PS.get(polar.network, PS.NumberOfGenerators())
end

function bounds(polar::PolarForm{T, IT, VT, MT}, ::typeof(reactive_power_constraints)) where {T, IT, VT, MT}
    q_min, q_max = PS.bounds(polar.network, PS.Generators(), PS.ReactivePower())
    return convert(VT, q_min), convert(VT, q_max)
end

# Adjoint
function adjoint!(
    polar::PolarForm,
    pbm::AutoDiff.TapeMemory{F, S, I},
    cons, ∂cons,
    vm, ∂vm,
    va, ∂va,
    pinj, ∂pinj,
) where {F<:typeof(reactive_power_constraints), S, I}
    nbus = get(polar, PS.NumberOfBuses())
    ref = polar.indexing.index_ref
    pv = polar.indexing.index_pv
    pq = polar.indexing.index_pq
    pv_to_gen = polar.indexing.index_pv_to_gen
    ref_to_gen = polar.indexing.index_ref_to_gen
    ybus_re, ybus_im = get(polar.topology, PS.BusAdmittanceMatrix())

    fill!(pbm.intermediate.∂edge_vm_fr , 0.0)
    fill!(pbm.intermediate.∂edge_vm_to , 0.0)
    fill!(pbm.intermediate.∂edge_va_fr , 0.0)
    fill!(pbm.intermediate.∂edge_va_to , 0.0)

    adj_reactive_power!(
        cons, ∂cons,
        vm, ∂vm,
        va, ∂va,
        ybus_re, ybus_im, polar.topology.sortperm,
        pinj, ∂pinj,
        pbm.intermediate.∂edge_vm_fr,
        pbm.intermediate.∂edge_vm_to,
        pbm.intermediate.∂edge_va_fr,
        pbm.intermediate.∂edge_va_to,
        polar.reactive_load,
        pv, pq, ref, pv_to_gen, ref_to_gen, nbus,
    )
end

function matpower_jacobian(polar::PolarForm, X::Union{State,Control}, ::typeof(reactive_power_constraints), V)
    nbus = get(polar, PS.NumberOfBuses())
    pf = polar.network
    ref = pf.ref
    pv = pf.pv
    pq = pf.pq
    gen2bus = polar.indexing.index_generators
    Ybus = pf.Ybus

    dSbus_dVm, dSbus_dVa = PS.matpower_residual_jacobian(V, Ybus)

    if isa(X, State)
        j11 = imag(dSbus_dVa[gen2bus, [pv; pq]])
        j12 = imag(dSbus_dVm[gen2bus, pq])
        return [j11 j12]
    elseif isa(X, Control)
        j11 = imag(dSbus_dVm[gen2bus, [ref; pv]])
        j12 = spzeros(length(gen2bus), length(pv))
        return [j11 j12]
    end
end

function matpower_hessian(polar::PolarForm, ::typeof(reactive_power_constraints), buffer, λ)
    nbus = get(polar, PS.NumberOfBuses())
    ref = polar.indexing.index_ref
    pv = polar.indexing.index_pv
    pq = polar.indexing.index_pq
    gen2bus = polar.indexing.index_generators
    # Check consistency
    @assert length(λ) == length(gen2bus)

    λq = zeros(nbus)
    # Select only buses with generators
    λq[gen2bus] .= λ

    V = buffer.vmag .* exp.(im .* buffer.vang)
    hxx, hxu, huu = PS.reactive_power_hessian(V, polar.network.Ybus, λq, pv, pq, ref)
    return FullSpaceHessian(
        hxx, hxu, huu,
    )
end


text:
using IndexedTables

import IndexedTables.NDSparse

function NDSparse(n::NamedArray)
	L = length(n) # elements in array
	cols = Dict{Symbol, Array}()
	factor = 1
	for d in 1:ndims(n)
		nlevels = size(n, d)
		nrep = L ÷ (nlevels * factor)
		data = repeat(vcat([fill(x, factor) for x in names(n, d)]...), nrep)
		cols[Symbol(dimnames(n, d))] = data
		factor *= nlevels
	end
	return NDSparse(Columns(;cols...), array(n)[:])
end


text:
module AbstractSDRsSpectrum
# ----------------------------------------------------
## --- Loading Modules
# ----------------------------------------------------
# --- Core Julia modules
using Plots
gr()
using Printf
using FFTW
using Dates
using AbstractSDRs
using Statistics
using LinearAlgebra

# --- Package for GUI
using Blink;        # For GUI in electron
using Interact;      # For Widget and interactions
using PlotlyJS;

# using Suppressor ## To avoid spam when rate is high ??
# # export UHD_LOG_FASTPATH_DISABLE=1

# --- Define custom macros
export @updateCarrierFreq
export @updateGain
export @updateBand
export @updateMean
export @updateLim

# ----------------------------------------------------
# --- GUI functions
# ----------------------------------------------------
include("gui.jl");

# ----------------------------------------------------
## --- Processing functions
# ----------------------------------------------------
""" hostSpectrum
---
Acquire signal from USRP and compute the spectrum. Plot the obtained PSD  on a GR plot figure
# --- Syntax
hostSpectrum(nFFT)
# --- Input parameters
- radio         : Radio object
- nFFT			: FFT size
# --- Output parameters
- []
# ---
# v 1.0 - Robin Gerzaguet.
"""
function hostSpectrum(radioAll,nFFT,nbSegMean;maxHold=false);
    # --- Update USRP config.
    # We define global variable to update USRP config
    global doTask	= 1;
    global changed  = 0;
    global newMean  = nbSegMean;
    global newLims  = (-60,20);
    nbSamples = nFFT;
    # global newSamplingRate  = getSamplingRate(radio);
    global newSamplingRate  = getSamplingRate(radio)
    # --- Define local parameters
    yLim	        = newLims;
    sF 	            = zeros(Cfloat,nFFT);
    y	            = zeros(Cfloat,nFFT);
    maxH            = zeros(Cfloat,nFFT);
    xAx		        = ((0:nFFT-1) ./ nFFT .-0.5) .*  round(newSamplingRate,digits=2);
    fMHz            = round(getCarrierFreq(radio) / 1e6, digits=3) ;
    buffer          = zeros(Complex{Cfloat},nbSamples);
    sig				= zeros(Complex{Cfloat},nFFT);
    P               = plan_fft(sig;flags=FFTW.PATIENT);
    α 				= 1 / nbSegMean;
    cnt				= 0;
    try
        while(true)
            # --- Getting signal
            recv!(buffer,radioAll);
            if nbSamples < nFFT
                sig[1:nbSamples] .= buffer;
            else
                sig[1:nFFT] .= buffer[1:nFFT];
            end
            # --- Apply | FFT(.) | ^2
            y  .= abs2.(fftshift(fft(sig)));
            # --- Averaging
            sF .= (1-α) .* sF .+ α .* y;
            # --- Configuration axis
            if changed == 1
                nbSegMean   = newMean;
                yLim        = newLims;
                # --- Recalculate x axis based on the new configuration
                sleep(0.001);
                xAx     = ((0:nFFT-1) ./ nFFT .-0.5) .* round(newSamplingRate,digits=2);
                fMHz	= round(getCarrierFreq(radio) / 1e6, digits=3) ;
                # --- Relock the process until a new change is detected
                changed = 0;
                α 		= 1/nbSegMean;
                sF 		.= 0;
                maxH 	.= 0;
            end
            # --- Update plot
            plt		= Plots.plot(xAx/1e6,10.0*log10.(sF),title="Spectrum of $(round(getSamplingRate(radio)/1e6,digits=2)) MHz @ $(round(getCarrierFreq(radio) / 1e6, digits=3)) MHz ",xlabel="Frequency [MHz]",ylabel="Power",label="",ylims=yLim,reuse=false);
            if maxHold
                maxH 	= max.(maxH,sF);
                Plots.plot!(plt,xAx/1e6,10.0*log10.(maxH),label="");
            end
            if mod(cnt,20) == 0
                plt.attr[:size]=(1200,800)
                display(plt);
            end
            cnt+=1;
            yield();
            # --- Interruption manager
            if doTask != 1
                break;
            end
        end
    catch exception
        rethrow(exception);
    end
    # --- Release USRP
    close(radioAll);
end

"""  hostWaterfall
---
Acquire signal from USRP and compute the spectrum. Plot the obtained PSD  on a GR plot figure
# --- Syntax
hostSpectrum(nFFT)
# --- Input parameters
- radio         : Radio object
- nFFT			: FFT size
# --- Output parameters
- []
# ---
# v 1.0 - Robin Gerzaguet.
"""
function hostWaterfall(radio,nFFT,nbSegMean);
    # --- Update USRP config.
    # We define global variable to update USRP config
    global doTask	= 1;
    global changed  = 0;
    global newMean  = nbSegMean;
    global newLims  = (-60,20);
    global newSamplingRate  = getSamplingRate(radio);
    # --- Define local parameters
    yLim	        = newLims;
    sF 	            = zeros(Cfloat,nFFT,nbSegMean);
    y	            = zeros(Cfloat,nFFT);
    xAx		        = ((0:nFFT-1) ./ nFFT .-0.5) .*  round(getSamplingRate(radio),digits=2);
    yAy 			= round.((0:nbSegMean-1)./ getSamplingRate(radio) .* 1e3,digits=3)
    fMHz            = round(getCarrierFreq(radio) / 1e6, digits=3) ;
    sig				= zeros(Complex{Cfloat},nFFT);
    P               = plan_fft(sig;flags=FFTW.PATIENT);
    tSeg 			= 1/ newSamplingRate * nFFT * 1000;
    yAy 	= (0:nbSegMean-1).*tSeg;
    try
        while(true)
            # --- Getting signal
            recv!(sig,radio);
            # --- Apply | FFT(.) | ^2
            y  .= abs2.(fftshift(fft(sig)));
            # --- Averaging
            sF	   	= circshift(sF,(0,-1));
            sF[:,end] = y;
            # --- Configuration axis
            if changed == 1
                # # nbSegMean   = newMean;
                yLim        = newLims;
                # --- Recalculate x axis based on the new configuration
                xAx     = ((0:nFFT-1) ./ nFFT .-0.5) .* round(newSamplingRate,digits=2);
                # yAy 	= (0:nbSegMean*nFFT-1) ./ getSamplingRate(radio) .* 1e3;
                fMHz	= round(getCarrierFreq(radio) / 1e6, digits=3) ;
                tSeg 	= 1/newSamplingRate * nFFT*1000;
                yAy 	= (0:nbSegMean-1).*tSeg;
                # --- Relock the process until a new change is detected
                changed = 0;
                sF 		.= 0;
            end
            # --- Update plot
            # plt		= heatmap(xAx/1e6,(0:nbSegMean-1),10.0*log10.(sF)',title="Spectrum of $(round(getSamplingRate(radio)/1e6,digits=2)) MHz @ $(round(getCarrierFreq(radio) / 1e6, digits=3)) MHz ",xlabel="Frequency [MHz]",ylabel="Time [burst]",clim=yLim,label="",reuse=false);
            plt		= heatmap(xAx/1e6,yAy,10.0*log10.(sF)',title="Spectrum of $(round(getSamplingRate(radio)/1e6,digits=2)) MHz @ $(round(getCarrierFreq(radio) / 1e6, digits=3)) MHz ",xlabel="Frequency [MHz]",ylabel="Time [ms]",clim=yLim,label="",reuse=false);

            # plt = heatmap(transpose(sF));
            # end
            plt.attr[:size]=(1200,800)
            gui(plt);
            yield();
            # gui(plt);
            # --- Interruption manager
            if doTask != 1
                break;
            end
        end
    catch exception ;
        rethrow(exception);
    end
    # --- Release USRP
    close(radio);
end

function processing!(out,sig,internal,P)
    # --- Plan FFT
    mul!(internal,P,sig);
    # Shift
    sig .= fftshift(internal)
    # --- |.|^2
    out .= abs2.(sig);
end

""" killSpectrum
---
Interrupt hostSpectrum function when run in asynchronous mode
# --- Syntax
killSpectrum()
# --- Input parameters
- []
# --- Output parameters
- []
# ---
# v 1.0 - Robin Gerzaguet.
"""
function stop()
    global doTask = 0;
end
function killSpectrum()
    global doTask			  = 0;
end

""" @updateCarrierFreq
---
Update carrier frequency during T/F analysis
# --- Syntax
@updateCarrierFreq xxx
# --- Input parameters
- xxx	: Value of new desired carrier frequency
# --- Output parameters
- []
# ---
# v 1.0 - Robin Gerzaguet.
"""
macro updateCarrierFreq(param)
    # --- Updating carrier frequency
    global changed = 1;
    global newCarrierFreq = param;
    # --- Calling routine to update radio
    updateCarrierFreq!(radio,param);
end

""" @updateGain
---
Macro to dynamically update Rx gain during T/F grid
# --- Syntax
@updateGain xxx
# --- Input parameters
- xxx	: New value of Rx gain
# --- Output parameters
- []
# ---
# v 1.0 - Robin Gerzaguet.
"""
macro updateGain(param)
    # --- Updating carrier frequency
    global newGain = param;
    global changed = 1;
    # --- Calling routine to update radio
    updateGain!(radio,param);
end

""" @updateBand
---
Dynamically update Rx sample rate
# --- Syntax
@updateBand xxx
# --- Input parameters
- xxx	: New desired value for sample rate
# --- Output parameters
- []
# ---
# v 1.0 - Robin Gerzaguet.
"""
macro updateBand(param)
    # --- Updating carrier frequency
    global newSamplingRate = param;
    global changed = 1;
    # --- Calling routine to update radio
    updateSamplingRate!(radio,newSamplingRate);
end

""" @updateLim
---
Update the y axis limit view dynamically
# --- Syntax
@updateLim (xxx1,xxx2)
# --- Input parameters
- xxx1	  : New minimal y value
- xxx2	  : New maximal  y value
# --- Output parameters
- []
# ---
# v 1.0 - Robin Gerzaguet.
"""
macro updateLim(param)
    global changed = 1;
    global newLims=eval(param);
end

"""	@updateMean
---
Update the averaging system to compute spectrum
# --- Syntax
@updateMean xxx
# --- Input parameters
- xxx	: New number of mean
# --- Output parameters
- []
# ---
# v 1.0 - Robin Gerzaguet.
"""
macro updateMean(param)
    global changed = 1;
    global newMean = param;
end

# ----------------------------------------------------
## --- Main call
# ----------------------------------------------------
""" main
---
Main routine to display Power spectral density
# --- Syntax
main()
# --- Input parameters
- []
# --- Output parameters
- []
# ---
# v 1.0 - Robin Gerzaguet.
"""
function start(sdr;kwargs...)
    main(sdr;kwargs...);
end
function main(sdr::Union{Symbol,String};carrierFreq=770e6,gain=20,samplingRate=80e6,maxHold::Bool=false,waterfall::Bool=false,kwargs...)
    @info "Spectrum display -- "
    # --- Simulation parameters
    global newCarrierFreq	  = carrierFreq;		  # --- Starting frequency
    global newGain		      = gain;			  # --- Analog Rx gain
    global newSamplingRate	  = samplingRate;		  # --- Sampling
    nFFT 		              = 1024;
    # --- Update radio configuration
    global radio			= openSDR(sdr,newCarrierFreq,newSamplingRate,newGain;kwargs...);
    if !waterfall
        @async hostSpectrum(radio,nFFT,32;maxHold=maxHold);
    else
        @async hostWaterfall(radio,nFFT,32);
    end
end

end # module


text:
using DSGE, DataFrames, JLD

path = dirname(@__FILE__())

# Set up arguments
m = AnSchorfheide(testing = true)
m <= Setting(:date_forecast_start, quartertodate("2015-Q4"))

df, system = jldopen("$path/../reference/forecast_args.jld","r") do file
    read(file, "df"), read(file, "system")
end

# Read expected output
exp_states, exp_shocks, exp_pseudo = jldopen("$path/../reference/smooth_out.jld", "r") do file
    read(file, "exp_states"),
    read(file, "exp_shocks"),
    read(file, "exp_pseudo")
end

# Smooth without drawing states
states = Dict{Symbol, Matrix{Float64}}()
shocks = Dict{Symbol, Matrix{Float64}}()
pseudo = Dict{Symbol, Matrix{Float64}}()

@testset "Test smoother without drawing states" begin
    for smoother in [:hamilton, :koopman, :carter_kohn, :durbin_koopman]
        m <= Setting(:forecast_smoother, smoother)

        states[smoother], shocks[smoother], pseudo[smoother] =
            smooth(m, df, system; draw_states = false)

        @test @test_matrix_approx_eq exp_states states[smoother]
        @test @test_matrix_approx_eq exp_shocks shocks[smoother]
        @test @test_matrix_approx_eq exp_pseudo pseudo[smoother]
    end
end

# Smooth, drawing states
for smoother in [:carter_kohn, :durbin_koopman]
    m <= Setting(:forecast_smoother, smoother)
    smooth(m, df, system; draw_states = true)
end

nothing


text:
module TrySequence
#import PatternDispatch.Patterns.depsof, PatternDispatch.Patterns.subs
#using PatternDispatch.Immutable,
import PatternDispatch
using PatternDispatch.Patterns, PatternDispatch.Recode, Toivo
#import PatternDispatch.Patterns.TupleRef

# ---- Stuff for Patterns.jl --------------------------------------------------

#@in PatternDispatch.Patterns begin
#end

# ---- sequence! -------------------------------------------------------------

type Sequence
    intent::Intension
    results::Dict{Node,Node}
    seq::Vector{Node}
end
Sequence(intent::Intension) = Sequence(intent,   Dict{Node,Node}(), Node[])
Sequence(s::Sequence)       = Sequence(s.intent, copy(s.results),     Node[])

wrap(node::Guard) = node
wrap(node::Node)  = Result(node)

function sequence!(s::Sequence, node::Node)
    @show node
    if has(s.results, node)
        newnode = s.results[node]
        if isa(newnode, Result); newnode.nrefs += 1; end
        return
    end

    for dep in depsof(s.intent, node);  sequence!(s, dep)  end
    newnode = s.results[node] = wrap(subs(s.results, node))
    push(s.seq, newnode)
end

# ---- code_match etc --------------------------------------------------------

function code_match(p::Pattern)
    sg = Sequence(p.intent)
    for g in guardsof(p.intent);    sequence!(sg, Guard(g));     end
    sb = Sequence(sg)
    for node in values(p.bindings); sequence!(sb, node);  end

    pred = code_predicate(sg.seq)
    prebind = encoded(sb.seq)
    binds = { :( $name = $(sb.results[node].ex)) for (name,node) in p.bindings }
    bind = vcat(prebind, binds)

    pred, bind
end

function code_predicate(seq::Vector{Node})
    code = {}
    pred = nothing
    for node in seq
        if isa(node, Guard)
            factor = isempty(code) ? node.pred.ex :  quote; $(code...); $(node.pred.ex) end
            pred = pred === nothing ? factor : (:($pred && $factor))
            code = {}
        else
            encode!(code, node)
        end
    end
    pred
end

function encoded(seq::Vector{Node})
    code = {}
    for node in seq;  encode!(code, node);  end
    code
end

encode!(code::Vector, ::Guard) = error("Undefined!")
function encode!(code::Vector, node::Result)
    if node.ex != nothing;  return  end
    ex = encode(node.node)

    if isa(ex, Symbol) || node.nrefs == 1
        node.ex = ex
    else
        var = gensym("t")
        node.ex = var
        push(code, :( $var = $ex ))
    end
end

# ---- test code --------------------------------------------------------------

function code_method(p::Pattern, body)
    pred, bind = code_match(p)
    fdef = :($argsym->begin
        if $pred
            $(bind...)
            (true, $body)
        else
            (false, nothing)
        end
    end)
    @show p
    @show pred
    @show bind
    @show fdef
    fdef
end

code_method((@qpat (x::Int,)), :(x^2))

end # module