@testset "Checking io-equations: single output" begin
    test_cases = []

    # 2-compartiment model
    push!(test_cases, Dict(
        "ode" => @ODEmodel(
            x0'(t) = -(a01 + a21) * x0(t) + a12 * x1(t) + u(t),
            x1'(t) = a21 * x0(t) - a12 * x1(t),
            y(t) = x0(t)
        ),
        "correct" => "y_2 + y_1*a01 + y_1*a21 + y_1*a12 + y_0*a01*a12 - u_0*a12 - u_1"
    ))

    #---------------------------------------
    # Chen-Lee model
    correct = "9*y1_3^2*y1_0^2 - 18*y1_3*y1_2*y1_1*y1_0 - 18*y1_3*y1_2*y1_0^2*a - 36*y1_3*y1_2*y1_0^2*b - 36*y1_3*y1_2*y1_0^2*c + 18*y1_3*y1_1^2*y1_0*a + 18*y1_3*y1_1^2*y1_0*b + 18*y1_3*y1_1^2*y1_0*c - 24*y1_3*y1_1*y1_0^4 + 18*y1_3*y1_1*y1_0^2*a*b + 18*y1_3*y1_1*y1_0^2*a*c + 18*y1_3*y1_1*y1_0^2*b^2 + 36*y1_3*y1_1*y1_0^2*b*c + 18*y1_3*y1_1*y1_0^2*c^2 + 24*y1_3*y1_0^5*a - 18*y1_3*y1_0^3*a*b^2 - 36*y1_3*y1_0^3*a*b*c - 18*y1_3*y1_0^3*a*c^2 + 9*y1_2^2*y1_1^2 + 18*y1_2^2*y1_1*y1_0*a + 36*y1_2^2*y1_1*y1_0*b + 36*y1_2^2*y1_1*y1_0*c + 9*y1_2^2*y1_0^2*a^2 + 36*y1_2^2*y1_0^2*a*b + 36*y1_2^2*y1_0^2*a*c + 27*y1_2^2*y1_0^2*b^2 + 90*y1_2^2*y1_0^2*b*c + 27*y1_2^2*y1_0^2*c^2 - 18*y1_2*y1_1^3*a - 18*y1_2*y1_1^3*b - 18*y1_2*y1_1^3*c + 24*y1_2*y1_1^2*y1_0^3 - 18*y1_2*y1_1^2*y1_0*a^2 - 72*y1_2*y1_1^2*y1_0*a*b - 72*y1_2*y1_1^2*y1_0*a*c - 54*y1_2*y1_1^2*y1_0*b^2 - 108*y1_2*y1_1^2*y1_0*b*c - 54*y1_2*y1_1^2*y1_0*c^2 + 48*y1_2*y1_1*y1_0^4*b + 48*y1_2*y1_1*y1_0^4*c - 18*y1_2*y1_1*y1_0^2*a^2*b - 18*y1_2*y1_1*y1_0^2*a^2*c - 18*y1_2*y1_1*y1_0^2*a*b^2 - 108*y1_2*y1_1*y1_0^2*a*b*c - 18*y1_2*y1_1*y1_0^2*a*c^2 - 18*y1_2*y1_1*y1_0^2*b^3 - 126*y1_2*y1_1*y1_0^2*b^2*c - 126*y1_2*y1_1*y1_0^2*b*c^2 - 18*y1_2*y1_1*y1_0^2*c^3 - 24*y1_2*y1_0^5*a^2 - 48*y1_2*y1_0^5*a*b - 48*y1_2*y1_0^5*a*c + 18*y1_2*y1_0^3*a^2*b^2 + 36*y1_2*y1_0^3*a^2*b*c + 18*y1_2*y1_0^3*a^2*c^2 + 18*y1_2*y1_0^3*a*b^3 + 126*y1_2*y1_0^3*a*b^2*c + 126*y1_2*y1_0^3*a*b*c^2 + 18*y1_2*y1_0^3*a*c^3 + 9*y1_1^4*a^2 + 18*y1_1^4*a*b + 18*y1_1^4*a*c + 9*y1_1^4*b^2 + 18*y1_1^4*b*c + 9*y1_1^4*c^2 - 24*y1_1^3*y1_0^3*a - 24*y1_1^3*y1_0^3*b - 24*y1_1^3*y1_0^3*c + 18*y1_1^3*y1_0*a^2*b + 18*y1_1^3*y1_0*a^2*c + 36*y1_1^3*y1_0*a*b^2 + 72*y1_1^3*y1_0*a*b*c + 36*y1_1^3*y1_0*a*c^2 + 18*y1_1^3*y1_0*b^3 + 54*y1_1^3*y1_0*b^2*c + 54*y1_1^3*y1_0*b*c^2 + 18*y1_1^3*y1_0*c^3 + 16*y1_1^2*y1_0^6 + 24*y1_1^2*y1_0^4*a^2 - 12*y1_1^2*y1_0^4*b^2 - 72*y1_1^2*y1_0^4*b*c - 12*y1_1^2*y1_0^4*c^2 - 18*y1_1^2*y1_0^2*a^2*b^2 - 18*y1_1^2*y1_0^2*a^2*c^2 - 18*y1_1^2*y1_0^2*a*b^3 + 18*y1_1^2*y1_0^2*a*b^2*c + 18*y1_1^2*y1_0^2*a*b*c^2 - 18*y1_1^2*y1_0^2*a*c^3 + 36*y1_1^2*y1_0^2*b^3*c + 72*y1_1^2*y1_0^2*b^2*c^2 + 36*y1_1^2*y1_0^2*b*c^3 - 32*y1_1*y1_0^7*a + 24*y1_1*y1_0^5*a^2*b + 24*y1_1*y1_0^5*a^2*c + 24*y1_1*y1_0^5*a*b^2 + 144*y1_1*y1_0^5*a*b*c + 24*y1_1*y1_0^5*a*c^2 - 72*y1_1*y1_0^3*a^2*b^2*c - 72*y1_1*y1_0^3*a^2*b*c^2 - 72*y1_1*y1_0^3*a*b^3*c - 144*y1_1*y1_0^3*a*b^2*c^2 - 72*y1_1*y1_0^3*a*b*c^3 + 16*y1_0^8*a^2 - 12*y1_0^6*a^2*b^2 - 72*y1_0^6*a^2*b*c - 12*y1_0^6*a^2*c^2 + 36*y1_0^4*a^2*b^3*c + 72*y1_0^4*a^2*b^2*c^2 + 36*y1_0^4*a^2*b*c^3"

    push!(test_cases, Dict(
        "ode" => @ODEmodel(
            x0'(t) = a * x0(t) - x1(t) * x2(t),
            x1'(t) = b * x1(t) + x0(t) * x2(t),
            x2'(t) = c * x2(t) + 1//3 * x0(t) * x1(t),
            y1(t) = x0(t)
        ),
        "correct" => correct
    ))

    #---------------------------------------

    # predator-prey model
    correct = "y1_2*y1_0 - y1_1^2 - y1_1*y1_0^2*d + y1_1*y1_0*c + y1_0^3*a*d - y1_0^2*a*c"

    push!(test_cases, Dict(
        "ode" => @ODEmodel(
            x0'(t) = a * x0(t) - b * x0(t) * x1(t),
            x1'(t) = -c * x1(t) + d * x0(t) * x1(t),
            y1(t) = x0(t)
        ),
        "correct" => correct
    ))

    #---------------------------------------

    for case in test_cases
        ode = case["ode"]
        io_eq = collect(values(find_ioequations(ode)))[1]
        gen_symbols = symbols(parent(io_eq))
        correct_expr = Meta.parse(case["correct"])
        for (i, s) in enumerate(gen_symbols)
            correct_expr = StructuralIdentifiability.MacroTools.postwalk(ex -> ex == s ? :(polyvars[$i]) : ex, correct_expr)
        end
        # dirty hack to eval within the local scope
        correct_function = @eval function(eq)
            polyvars = gens(parent(eq))
            $correct_expr
        end
        correct = correct_function(io_eq)
        divisibility, remainder = divides(io_eq, correct)
        @test divisibility
        @test total_degree(remainder) == 0
    end
end


module MPSforQuantum
    export MPS, restore, OneQubitGate, TwoQubitGate, dstack, ddstack, init_rand_MPS
    export CX, inner_product, expectation, SVD_R, SVD_L, iterative_ground_state_search, j1j2_2D_Hamiltonian_MPO
    using LinearAlgebra
    using TensorOperations

    function SVD(C::Array{ComplexF64,2}, eps::Float64)
        #eps = 1e-2
        A = svd(C)
        filter!((x) -> x > eps, A.S)
        l = length(A.S)
        return A.U[:, 1:l], diagm(A.S), A.Vt[1:l, :]
        #return A.U[:, 1:l], diagm(A.S), transpose(A.Vt)[1:l, :]
    end

    function SVD(C::Array{ComplexF64,2}, D::Int64)
        A = svd(C)
        if length(A.S) > D
            return A.U[:, 1:D], diagm(A.S[1:D]), A.Vt[1:D, :]
            #return A.U[:, 1:D], diagm(A.S[1:D]), transpose(A.Vt)[1:D, :]
        else
            return A.U[:, :], diagm(A.S), A.Vt[:, :]
            #return A.U[:, :], diagm(A.S), transpose(A.Vt)[:, :]
        end
    end

    function SVD_L(C::Array{ComplexF64,2}, eps::Float64)
        U, S, Vt = SVD(C, eps)
        return U, S * Vt
    end

    function SVD_L(C::Array{ComplexF64,2}, D::Int64)
        U, S, Vt = SVD(C, D)
        return U, S * Vt
    end

    function SVD_R(C::Array{ComplexF64,2}, eps::Float64)
        U, S, Vt = SVD(C, eps)
        return U * S, Vt
    end

    function SVD_R(C::Array{ComplexF64,2}, D::Int64)
        U, S, Vt = SVD(C, D)
        return U * S, Vt
    end

    function MPS(C::Array{ComplexF64,1}, param, normalize::Char='l')
        d = 2
        N = Int64(log2(size(C)[1]))
        arrs = []
        r = 1
        if normalize == 'l'
            for i = 1:N-1
                C = reshape(C, d * r, d^(N - i))
                tmp_A, C = SVD_L(C, param)
                r = size(tmp_A, 2)
                col = convert(Int64, size(tmp_A, 1) / 2)
                push!(arrs, [tmp_A[1:col, :], tmp_A[(col+1):col*2, :]])
            end

            Ct = transpose(C)
            col2 = convert(Int64, size(C, 2) / 2)
            push!(arrs, [C[:, 1:col2], C[:, (col2+1):col2*2]])
            return arrs
        elseif normalize == 'r'
            for i = 1:N-1
                #C = reshape(C, d^(N - i), d * r)
                l = convert(Int64, size(C, 1) / 2)
                C = cat(C[1:l, :], C[l+1:l*2, :], dims=2)
                C, tmp_B = SVD_R(C, param)
                r = size(tmp_B, 1)
                col = convert(Int64, size(tmp_B, 2) / 2)
                push!(arrs, [tmp_B[:, 1:col], tmp_B[:, (col+1):col*2]])
            end

            Ct = transpose(C)
            col2 = convert(Int64, size(C, 1) / 2)
            push!(arrs, [C[1:col2, :], C[(col2+1):col2*2, :]])
            arrs_ = []
            for i = 1:N
                push!(arrs_, arrs[N - (i-1)])
            end
            return arrs_
        end
    end

    function OneQubitGate(arrs::Array{Any, 1}, O::Array{Complex{Float64},2}, n::Int64)
        arrs_ = similar(arrs[n])
        arrs__ = copy(arrs)
        arrs_[1] = arrs[n][1] * O[1, 1] + arrs[n][2] * O[1, 2]
        arrs_[2] = arrs[n][1] * O[2, 1] + arrs[n][2] * O[2, 2]
        arrs__[n] = arrs_
        return arrs__
    end

    function swap_ct(arr::Array{Complex{Float64},2})
        swap = [1 0 0 0; 0 0 1 0; 0 1 0 0; 0 0 0 1]
        _arr = copy(arr)
        _arr[:, 2] = arr[:, 3]
        _arr[:, 3] = arr[:, 2]
        return swap * _arr
    end

    function reduce(mps::Array{Any, 1}, t1::Int64, t2::Int64, param::Number)
        eps = 1e-3
        arr1 = cat(mps[t1][1], mps[t1][2], dims = 1)
        arr2 = cat(mps[t2][1], mps[t2][2], dims = 2)
        arr = arr1 * arr2

        # U, S, V = svd(arr)
        # filter!((x) -> x > eps, S)
        # sdim = length(S)
        U, S, V = SVD(arr, param)

        #U1 = U[:, 1:sdim]
        U1 = U
        U1 = [U1[1:convert(Int, size(U1)[1]/2), :], U1[convert(Int, size(U1)[1]/2)+1:size(U1)[1], :]]
        #U2 = diagm(S) * transpose(V)[1:sdim, :]
        U2 = S * V
        U2 = [U2[:, 1:convert(Int, size(U2)[2]/2)], U2[:, convert(Int, size(U2)[2]/2)+1:size(U2)[2]]]

        mps[t1] = U1
        mps[t2] = U2
        return mps
    end

    function TwoQubitGate(mps::Array{Any, 1}, arr::Array{Complex{Float64},2}, c::Int64, t::Int64, param::Number)
        eps = 1e-6 # constant cutoff for TwoQubitGate decomposition

        if c == t + 1
            arr = swap_ct(arr)
            (c, t) = (t, c)
        end

        arr_00 = transpose(reshape(arr[:, 1], 2, 2))
        arr_01 = transpose(reshape(arr[:, 2], 2, 2))
        arr_10 = transpose(reshape(arr[:, 3], 2, 2))
        arr_11 = transpose(reshape(arr[:, 4], 2, 2))

        tensor1 = cat(arr_00, arr_01, dims = 2)
        tensor2 = cat(arr_10, arr_11, dims = 2)
        tensor = cat(tensor1, tensor2, dims = 1)

        U, S, V = svd(tensor)
        filter!((x) -> x > eps, S)
        sdim = length(S)
        U1 = U[:, 1:sdim]
        U1 = [U1[1:convert(Int, size(U1)[1]/2), :], U1[convert(Int, size(U1)[1]/2)+1:size(U1)[1], :]]
        U2 = diagm(S) * transpose(V)[1:sdim, :]
        U2 = [U2[:, 1:convert(Int, size(U2)[2]/2)], U2[:, convert(Int, size(U2)[2]/2)+1:size(U2)[2]]]

        after_c = kron(U1[1], mps[c][1]) + kron(U1[2], mps[c][2])
        after_t = kron(U2[1], mps[t][1]) + kron(U2[2], mps[t][2])
        c_row = convert(Int, size(after_c)[1]/2)
        t_col = convert(Int, size(after_t)[2]/2)

        mps[t][1] = convert(Array{ComplexF64,2}, after_t[:, 1:t_col])
        mps[t][2] = convert(Array{ComplexF64,2}, after_t[:, t_col+1:size(after_t)[2]])
        mps[c][1] = convert(Array{ComplexF64,2}, after_c[1:c_row, :])
        mps[c][2] = convert(Array{ComplexF64,2}, after_c[c_row+1:size(after_c)[1], :])

        mps = reduce(mps, c, t, param)

        return mps
    end

    function dstack(operators)
        return cat(operators..., dims = 3)
    end

    function ddstack(operators)
        m = size(operators)[1]
        tmp2 = []
        for i=1:m
            push!(tmp2, cat(operators[i]..., dims=4))
        end
        tmp3 = Tuple(tmp2)
        mpo = cat(tmp3..., dims=3)
        return mpo
    end

    function inner_product(arrs1::Array{Any, 1}, arrs2::Array{Any, 1})
        N = Int64(size(arrs1)[1])
        ip = arrs1[1][1]' * arrs2[1][1] + arrs1[1][2]' * arrs2[1][2]
        for i = 2:N
            phys0 = arrs1[i][1]' * ip * arrs2[i][1]
            phys1 = arrs1[i][2]' * ip * arrs2[i][2]
            ip = phys0 + phys1
        end
        return ip
    end

    # function CX(arrs::Array{Any, 1}, n::Int64, t::Number)
    #     tmp = copy(arrs)
    #     arrs_ = [tmp[n + 1][1]*tmp[n + 2][1] tmp[n + 1][2]*tmp[n + 2][2]
    #             tmp[n + 1][1]*tmp[n + 2][2] tmp[n + 1][2]*tmp[n + 2][1]]
    #     arrs__ = MPSforQuantum.SVD_L(arrs_, t)
    #     col = convert(Int64, size(arrs__[1], 1) / 2)

    #     tmp[n + 1][1] = arrs__[1][1:col, :]
    #     tmp[n + 1][2] = arrs__[1][(col+1):(2*col), :]
    #     tmp[n + 2][1] = arrs__[2][:, 1:col]
    #     tmp[n + 2][2] = arrs__[2][:, (col+1):(2*col)]

    #     return tmp
    # end

    function pauli()
        pauliX = convert(Array{ComplexF64,2}, [0 1; 1 0])
        pauliY = convert(Array{ComplexF64,2}, [0 -1*im; 1*im 0])
        pauliZ = convert(Array{ComplexF64,2}, [1 0; 0 -1])
        pauliI = convert(Array{ComplexF64,2}, [1 0; 0 1])
        return (pauliX, pauliY, pauliZ, pauliI)
    end

    function pauli_zero()
        zero = convert(Array{ComplexF64,2}, [0 0; 0 0])
        return zero
    end

    function expectation(mps::Array{Any,1}, O::Array{Any,1})
        N_site = size(mps)[1]
        contracted_sites = []
        n_phys = 2 #ss
        for i=1:N_site
            if i==1
                # first step
                tmp_mps = cat(mps[i][1][:], mps[i][2][:], dims=2)
                tmp_mps_dag = cat(mps[i][1]'[:], mps[i][2]'[:], dims=2)
                tmp_O = O[i]
                a_len = size(mps[i][1])[2]
                b_len = size(O[i][1, 1, :])[1] # bb0
                arr_0 = zeros(ComplexF64, n_phys, b_len, a_len)
                arr = zeros(ComplexF64, a_len, b_len, a_len)
                @tensor begin
                    arr_0[sigma1, b, a_] = tmp_O[sigma1, sigma2, b] * tmp_mps[a_, sigma2]
                    arr[a, b, a_] = tmp_mps_dag[a, sigma1] * arr_0[sigma1, b, a_]
                end

            elseif i==N_site
                # last step
                tmp_mps = cat(mps[i][1][:], mps[i][2][:], dims=2)
                tmp_mps_dag = cat(mps[i][1]'[:], mps[i][2]'[:], dims=2)
                tmp_O = O[i]
                tmp_site = contracted_sites[i-1]
                a_len1 = size(mps[i-1][1])[2]
                b_len2 = size(O[i][1, 1, :])[1]
                arr_0 = zeros(ComplexF64, n_phys, a_len1, b_len2)
                arr_1 = zeros(ComplexF64, n_phys, a_len1)
                @tensor begin
                    arr_0[sigma2, a1, b2] = tmp_site[a1, b2, a1_0] * tmp_mps[a1_0, sigma2]
                    arr_1[sigma1, a1] = tmp_O[sigma1, sigma2, b2] * arr_0[sigma2, a1, b2]
                    arr = tmp_mps_dag[a1, sigma1] * arr_1[sigma1, a1]
                end

            else
                # Middle step
                tmp_mps = cat(mps[i][1], mps[i][2], dims=3)
                tmp_mps_dag = cat(mps[i][1]', mps[i][2]', dims=3)
                tmp_O = O[i]
                tmp_site = contracted_sites[i-1]
                a_len1 = size(mps[i-1][1])[2]
                a_len2 = size(mps[i][1])[2]
                b_len1, b_len2 = size(O[i][1, 1, :, :])
                arr_0 = zeros(ComplexF64, n_phys, a_len1, b_len1, a_len2)
                arr_1 = zeros(ComplexF64, n_phys, b_len2, a_len1, a_len2)
                arr = zeros(ComplexF64, a_len2, b_len2, a_len2)
                @tensor begin
                    arr_0[sigma2, a1, b1, a2] = tmp_site[a1, b1, a0_0] * tmp_mps[a0_0, a2, sigma2]
                    arr_1[sigma1, b2, a1, a2] = tmp_O[sigma1, sigma2, b1, b2] * arr_0[sigma2, a1, b1, a2]
                    arr[a2, b2, a2_] = tmp_mps_dag[a2, a1, sigma1] * arr_1[sigma1, b2, a1, a2_]
                end
            end
            push!(contracted_sites, arr)
        end
        return contracted_sites[N_site]
    end

    function R_expression(mps::Array{Any,1}, O::Array{Any,1}, t::Int64)
        # t: target site, sites t+1 to N are merged to R expression
        N_site = size(mps)[1]
        n_phys = 2
        contracted_sites = []
        for i = t+1:N_site
            if i==t+1
                if i==N_site
                    tmp_mps = cat(mps[i][1][:], mps[i][2][:], dims=2)
                    tmp_mps_dag = cat(mps[i][1]'[:], mps[i][2]'[:], dims=2)
                    tmp_O = O[i]
                    a_len1 = size(mps[i-1][1])[2]
                    a_len2 = size(mps[i][1])[2]
                    b_len = size(O[i][1, 1, :])[1]
                    arr = zeros(ComplexF64, a_len1, b_len, a_len1)
                    @tensor begin
                        arr[a0, b0, a0_] = tmp_mps_dag[a0, sigma1] * tmp_O[sigma1, sigma2, b0] * tmp_mps[a0_, sigma2]
                    end
                    push!(contracted_sites, arr)
                else
                    tmp_mps = cat(mps[i][1], mps[i][2], dims=3)
                    tmp_mps_dag = cat(mps[i][1]', mps[i][2]', dims=3)
                    tmp_O = O[i]
                    a_len1 = size(mps[i-1][1])[2]
                    a_len2 = size(mps[i][1])[2]
                    b_len = size(O[i][1, 1, :, :])[1]
                    arr = zeros(ComplexF64, a_len1, a_len2, b_len, b_len, a_len1, a_len2)
                    @tensor begin
                        arr[a0, a1, b0, b1, a0_, a1_] = tmp_mps_dag[a1, a0, sigma1] * tmp_O[sigma1, sigma2, b0, b1] * tmp_mps[a0_, a1_, sigma2]
                    end
                    push!(contracted_sites, arr)
                end

            elseif i==N_site
                # last step
                tmp_mps = cat(mps[i][1][:], mps[i][2][:], dims=2)
                tmp_mps_dag = cat(mps[i][1]'[:], mps[i][2]'[:], dims=2)
                tmp_O = O[i]
                tmp_site = contracted_sites[1]
                a_len0 = size(tmp_site)[1]
                a_len1 = size(mps[i-1][1])[2]
                a_len2 = size(mps[i][1])[2]
                b_len = size(O[i][1, 1, :])[1]
                arr_0 = zeros(ComplexF64, n_phys, a_len0, a_len1, b_len, b_len, a_len0)
                arr_1 = zeros(ComplexF64, n_phys, n_phys, a_len0, b_len, b_len, a_len0)
                arr = zeros(ComplexF64, a_len0, b_len, a_len0)
                @tensor begin
                    arr_0[sigma2, a0, a1, b0, b1, a0_] = tmp_site[a0, a1, b0, b1, a0_, a1_] * tmp_mps[a1_, sigma2]
                    arr_1[sigma1, sigma2, a0, b0, b1, a0_] = tmp_mps_dag[a1, sigma1] * arr_0[sigma2, a0, a1, b0, b1, a0_]
                    arr[a0, b0, a0_] = arr_1[sigma1, sigma2, a0, b0, b1, a0_] * tmp_O[sigma1, sigma2, b1]
                end
                contracted_sites[1] = arr
            else
                tmp_mps = cat(mps[i][1], mps[i][2], dims=3)
                tmp_mps_dag = cat(mps[i][1]', mps[i][2]', dims=3)
                tmp_O = O[i]
                tmp_site = contracted_sites[1]
                a_len0 = size(tmp_site)[1]
                a_len1 = size(mps[i-1][1])[2]
                a_len2 = size(mps[i][1])[2]
                b_len = size(O[i][1, 1, :, :])[1]
                arr_0 = zeros(ComplexF64, n_phys, a_len0, a_len1, b_len, b_len, a_len0, a_len2)
                arr_1 = zeros(ComplexF64, n_phys, n_phys, a_len0, a_len2, b_len, b_len, a_len0, a_len2)
                arr = zeros(ComplexF64, a_len0, a_len2, b_len, b_len, a_len0, a_len2)
                @tensor begin
                    arr_0[sigma2, a0, a1, b0, b1, a0_, a2_] = tmp_site[a0, a1, b0, b1, a0_, a1_] * tmp_mps[a1_, a2_, sigma2]
                    arr_1[sigma1, sigma2, a0, a2, b0, b1, a0_, a2_] = tmp_mps_dag[a2, a1, sigma1] * arr_0[sigma2, a0, a1, b0, b1, a0_, a2_]
                    arr[a0, a2, b0, b2, a0_, a2_] = arr_1[sigma1, sigma2, a0, a2, b0, b1, a0_, a2_] * tmp_O[sigma1, sigma2, b1, b2]
                end
                contracted_sites[1] = arr
            end
        end
        return contracted_sites[1]
    end

    function L_expression(mps::Array{Any,1}, O::Array{Any,1}, t::Int64)
        # t: target site, sites 1 to t-1 are merged to L expression
        N_site = size(mps)[1]
        n_phys = 2
        contracted_sites = []
        for i = 1:t-1
            if i==1
                tmp_mps = cat(mps[i][1][:], mps[i][2][:], dims=2)
                tmp_mps_dag = cat(mps[i][1]'[:], mps[i][2]'[:], dims=2)
                tmp_O = O[i]
                a_len = size(mps[i][1])[2]
                b_len = size(O[i][1, 1, :, :])[1]
                arr = zeros(ComplexF64, a_len, b_len, a_len)
                @tensor begin
                    arr[a0, b0, a0_] = tmp_mps_dag[a0, sigma1] * tmp_O[sigma1, sigma2, b0] * tmp_mps[a0_, sigma2]
                end
                push!(contracted_sites, arr)
            else
                tmp_mps = cat(mps[i][1], mps[i][2], dims=3)
                tmp_mps_dag = cat(mps[i][1]', mps[i][2]', dims=3)
                tmp_O = O[i]
                tmp_site = contracted_sites[1]
                a_len0 = size(tmp_site)[1]
                a_len1 = size(mps[i-1][1])[2]
                a_len2 = size(mps[i][1])[2]
                b_len = size(O[i][1, 1, :, :])[1]
                arr_0 = zeros(ComplexF64, n_phys, a_len1, b_len, a_len2)
                arr_1 = zeros(ComplexF64, n_phys, n_phys, a_len2, b_len, a_len2)
                arr = zeros(ComplexF64, a_len2, b_len, a_len2)
                @tensor begin
                    arr_0[sigma2, a1, b1, a2_] = tmp_site[a1, b1, a1_] * tmp_mps[a1_, a2_, sigma2]
                    arr_1[sigma1, sigma2, a2, b1, a2_] = tmp_mps_dag[a2, a1, sigma1] * arr_0[sigma2, a1, b1, a2_]
                    arr[a2, b2, a2_] = arr_1[sigma1, sigma2, a2, b1, a2_] * tmp_O[sigma1, sigma2, b1, b2]
                end
                contracted_sites[1] = arr
            end
        end
        return contracted_sites[1]
    end

    function left_norm_for_2_sites(mps::Array{Any,1}, t::Int64, D::Int64)
        # t: target, site to be left normalized
        mps_ = copy(mps)
        site_1 = cat(mps[t][1], mps[t][2], dims=1)
        site_2 = cat(mps[t+1][1], mps[t+1][2], dims=2)
        mixed_site = site_1 * site_2
        A, M = SVD_L(mixed_site, D)
        col = convert(Int64, size(A, 1) / 2)
        mps_[t] = [A[1:col, :], A[(col+1):col*2, :]]
        col2 = convert(Int64, size(M, 2) / 2)
        mps_[t+1] = [M[:, 1:col2], M[:, (col2+1):col2*2]]
        return mps_
    end

    function right_norm_for_2_sites(mps::Array{Any,1}, t::Int64, D::Int64)
        # t: target, site to be right normalized
        mps_ = copy(mps)
        site_1 = cat(mps[t-1][1], mps[t-1][2], dims=1)
        site_2 = cat(mps[t][1], mps[t][2], dims=2)
        mixed_site = site_1 * site_2
        M, B = SVD_R(mixed_site, D)
        col = convert(Int64, size(M, 1) / 2)
        mps_[t-1] = [M[1:col, :], M[(col+1):col*2, :]]
        col2 = convert(Int64, size(B, 2) / 2)
        mps_[t] = [B[:, 1:col2], B[:, (col2+1):col2*2]]
        return mps_
    end

    function left_most_site_update(mps_::Array{Any,1}, O::Array{Any,1}, t::Int64)
        R = R_expression(mps_, O, t)
        tmp_O = O[t]
        H = zeros(ComplexF64, size(R, 1), size(R, 3), size(tmp_O, 1), size(tmp_O, 2))
        @tensor begin
            H[sigma1, a0, sigma2, a0_] = tmp_O[sigma1, sigma2, b0] * R[a0, b0, a0_]
        end
        H_ = zeros(ComplexF64, size(tmp_O, 1) * size(R, 1), size(tmp_O, 2) * size(R, 3))
        for i=1:size(H, 1), j=1:size(H, 2), k=1:size(H, 3), l=1:size(H, 4)
            H_[(i-1)*size(H, 2) + j, (k-1)*size(H, 4) + l] = H[i,j,k,l]
        end
        v = eigvecs(H_)[:, 1]
        d = convert(Int64, size(v, 1) / 2)
        mps_[t][1][:] = v[1:d]
        mps_[t][2][:] = v[d+1:2*d]
        return mps_
    end

    function mid_site_update(mps_::Array{Any,1}, O::Array{Any,1}, t::Int64)
        R = R_expression(mps_, O, t)
        L = L_expression(mps_, O, t)
        tmp_O = O[t]
        H = zeros(ComplexF64, size(tmp_O, 1), size(L, 1), size(R, 1), size(tmp_O, 2), size(L, 3), size(R, 3))
        @tensor begin
            H[sigma1, a0, a1, sigma2, a0_, a1_] = L[a0, b0, a0_] * tmp_O[sigma1, sigma2, b0, b1] * R[a1, b1, a1_]
        end
        H_ = zeros(ComplexF64, size(tmp_O, 1) * size(L, 1) * size(R, 1), size(tmp_O, 2) * size(L, 3) * size(R, 3))
        for i=1:size(H, 1), j=1:size(H, 2), k=1:size(H, 3), l=1:size(H, 4), m=1:size(H, 5), n=1:size(H, 6)
            H_[(i-1)*size(H, 2)*size(H, 3) + (j-1)*size(H, 3) + k, (l-1)*size(H, 5)*size(H, 6) + (m-1)*size(H, 6) + n] = H[i,j,k,l,m,n]
        end
        v = eigvecs(H_)[:, 1]
        d = convert(Int64, size(v, 1) / 2)
        M_1 = transpose(reshape(v[1:d], size(transpose(mps_[t][1]))))
        M_2 = transpose(reshape(v[d+1:2*d], size(transpose(mps_[t][1]))))
        mps_[t][1][:, :] = M_1
        mps_[t][2][:, :] = M_2
        return mps_
    end

    function right_most_site_update(mps_::Array{Any,1}, O::Array{Any,1}, t::Int64)
        L = L_expression(mps_, O, t)
        tmp_O = O[t]
        H = zeros(ComplexF64, size(L, 1), size(L, 3), size(tmp_O, 1), size(tmp_O, 2))
        @tensor begin
            H[sigma1, a0, sigma2, a0_] = tmp_O[sigma1, sigma2, b0] * L[a0, b0, a0_]
        end
        H_ = zeros(ComplexF64, size(tmp_O, 1) * size(L, 1), size(tmp_O, 2) * size(L, 3))
        for i=1:size(H, 1), j=1:size(H, 2), k=1:size(H, 3), l=1:size(H, 4)
            H_[(i-1)*size(H, 2) + j, (k-1)*size(H, 4) + l] = H[i,j,k,l]
        end
        v = eigvecs(H_)[:, 1]
        d = convert(Int64, size(v, 1) / 2)
        mps_[t][1][:] = v[1:d]
        mps_[t][2][:] = v[d+1:2*d]
        return mps_
    end

    function iterative_ground_state_search(mps::Array{Any,1}, O::Array{Any,1}, D::Int64, ite::Int64)
        hist = []
        N = size(mps, 1)
        mps_ = copy(mps)
        push!(hist, expectation(mps_, O))
        for i = 1:ite
            for t = 1:N
                if t == 1
                    mps_ = left_most_site_update(mps_, O, t)
                    mps_ = left_norm_for_2_sites(mps_, t, D)
                elseif t == N
                    mps_ = right_most_site_update(mps_, O, t)
                else
                    mps_ = mid_site_update(mps_, O, t)
                    mps_ = left_norm_for_2_sites(mps_, t, D)
                end
                # push!(hist, expectation(mps_, O))
            end

            for t = N-1:-1:1
                mps_ = right_norm_for_2_sites(mps_, t+1, D)
                if t == 1
                    mps_ = left_most_site_update(mps_, O, t)
                else
                    mps_ = mid_site_update(mps_, O, t)
                end
                push!(hist, expectation(mps_, O))
            end
        end
        return (mps_, hist)
    end

    function init_rand_MPS(N::Int64, D::Int64, normalize::Char='l')
        mps = []
        mps2 = []
        half_N = convert(Int64, floor(N/2))
        for i in 1:half_N
            l = 2^(i-1)
            mps_1 = []
            mps_2 = []
            for j in 1:2
                arr1 = rand(ComplexF64, (min(l, D), min(l*2, D))) / min(l*2, D)
                arr2 = rand(ComplexF64, (min(l*2, D), min(l, D))) / min(l, D)
                push!(mps_1, arr1)
                push!(mps_2, arr2)
            end
            push!(mps, mps_1)
            push!(mps2, mps_2)
        end

        if N%2 == 1
            mps_ = []
            for j in 1:2
                arr = rand(ComplexF64, (min(2^half_N, D), min(2^half_N, D))) / min(2^half_N, D)
                push!(mps_, arr)
            end
            push!(mps, mps_)
        end

        for j in 1:half_N
            push!(mps, mps2[half_N - (j-1)])
        end

        for i in 1:N-1
            mps = left_norm_for_2_sites(mps, i, D)
        end
        norm = mps[N][1]' * mps[N][1] + mps[N][2]' * mps[N][2]
        mps[N][1] = mps[N][1]/sqrt(norm)
        mps[N][2] = mps[N][2]/sqrt(norm)
        if normalize == 'r'
            for i in N:-1:2
                mps = right_norm_for_2_sites(mps, i, D)
            end
        end
        return mps
    end

    function push_N(arr::Array{Any,1}, op::Array{Complex{Float64},2}, N::Int64)
        arr_ = copy(arr)
        for i in 1:N
            push!(arr_, op)
        end
        return arr_
    end

    function j1j2_connection(L::Int64)
        J1_term = Set()
        for i in 1:L
            for j in 1:(L-1)
                row = 1 + 4*(i-1)
                push!(J1_term, (row + j - 1, row + j))
                col = i
                push!(J1_term, (col + 4*(j - 1), col + 4 + 4*(j - 1)))
            end
        end

        J2_term = Set()
        for i in 1:(L-1)
            for j in 1:(L-1)
                row = 1 + 4*(i-1)
                push!(J2_term, (row + (j - 1), row + (j - 1) + 5))
                push!(J2_term, (row + j, row + j + 3))
            end
        end
        return (J1_term, J2_term)
    end

    function j1j2_2D_Hamiltonian_MPO(L::Int64, j1::Float64, j2::Float64)
        N = L^2
        max_distance = L + 1
        MPO_L = 3*(max_distance + 2)
        O = []
        first_MPO = []
        (pauliX, pauliY, pauliZ, pauliI) = pauli()
        zero = pauli_zero()
        ops = [pauliX, pauliY, pauliZ]
        (J1_term, J2_term) = j1j2_connection(L)
        for op in ops
            push!(first_MPO, zero)
            for k in 1:max_distance
                if (1, 1 + k) in J1_term
                    push!(first_MPO, op*j1)
                elseif (1, 1 + k) in J2_term
                    push!(first_MPO, op*j2)
                else
                    push!(first_MPO, zero)
                end
            end
            push!(first_MPO, zero)
        end
        first_MPO_tuple = Tuple(first_MPO)
        push!(O, dstack(first_MPO_tuple))

        for site in 1:N-2
            mid_MPO = []
            for i in 1:length(ops)
                for row in 1:(max_distance + 2)
                    tmp_row = []
                    for j in 1:(i-1) # fill zero
                        tmp_row = push_N(tmp_row, zero, max_distance + 2)
                    end

                    if row == 1
                        push!(tmp_row, pauliI)
                        tmp_row = push_N(tmp_row, zero, max_distance + 1)
                    elseif row == 2
                        push!(tmp_row, ops[i])
                        tmp_row = push_N(tmp_row, zero, max_distance + 1)
                    elseif row == (max_distance + 2)
                        push!(tmp_row, zero)
                        for k in 1:max_distance
                            if (site, site + k) in J1_term
                                push!(tmp_row, j1*ops[i])
                            elseif (site, site + k) in J1_term
                                push!(tmp_row, j2*ops[i])
                            else
                                push!(tmp_row, zero)
                            end
                        end
                        push!(tmp_row, zero)
                    else
                        tmp_row = push_N(tmp_row, zero, row - 2)
                        push!(tmp_row, pauliI)
                        tmp_row = push_N(tmp_row, zero, max_distance + 1 - (row - 2))
                    end
                    for j in 1:(3-i) # fill zero
                        tmp_row = push_N(tmp_row, zero, max_distance + 2)
                    end
                    tmp_row_tuple = Tuple(tmp_row)
                    push!(mid_MPO, tmp_row_tuple)
                end
            end
            push!(O, ddstack(mid_MPO))
        end

        last_MPO = []
        for op in ops
            push!(last_MPO, pauliI)
            push!(last_MPO, op)
            last_MPO = push_N(last_MPO, zero, max_distance)
        end
        last_MPO_tuple = Tuple(last_MPO)
        push!(O, dstack(last_MPO_tuple))
        return O
    end

end


module StateVectorM

export StateVector, GateIndices
export init_state!, getindex, setindex!, getIndicesAfterExclusion, generateBitPatterns
export applyPauliX!, applyHadamard!, applyRX!, applyRY!

import Base.getindex
import Base.setindex!

mutable struct StateVector
    num_qubits::Int64
    data::Vector{ComplexF64}
end

Base.@ccallable function create_StateVector(num_qubits::Int64)::StateVector
    return StateVector(num_qubits, zeros(2^num_qubits).+0im)
end

mutable struct StateVectorPtr
    num_qubits::Int64
    data::Ptr{ComplexF64}
end
mutable struct StateVectorRef
    num_qubits::Int64
    data::Ref{ComplexF64}
end

function init_state!(sv::StateVector)
    sv.data = zeros(2^sv.num_qubits).+0im
    sv[1] = 1 + 0im
    ;
end

struct GateIndices
    internal::Vector{Int64}
    external::Vector{Int64}
    GateIndices(wires, num_qubits) = new(generateBitPatterns(wires, num_qubits), generateBitPatterns(getIndicesAfterExclusion(wires, num_qubits), num_qubits))
end

function getindex(sv::StateVector, index::Int64)
    return getindex(sv.data, index)
end

function setindex!(sv::StateVector, value::ComplexF64, index::Int64)
    sv.data[index] = value
end

Base.@ccallable function getIndicesAfterExclusion(indicesToExclude::Vector{Int64}, num_qubits::Int64)::Vector{Int64}
    indices = Set{Int64}(1:1:num_qubits)
    for excl_idx = indicesToExclude
        pop!(indices, excl_idx);
    end
    collect(indices)
end

Base.@ccallable function generateBitPatterns(qubitIndices::Vector{Int64}, num_qubits::Int64)::Vector{Int64}
    indices = zeros(1)
    for index_it = Iterators.reverse(qubitIndices)
        value = exp2(num_qubits - index_it)
        current_size = length(indices)
        for i = 1:1:current_size
            push!(indices, indices[i] + value)
        end
    end

    return indices
end

Base.@ccallable function applyPauliX!(sv::StateVector,
                    indices::Vector{Int64},
                    externalIndices::Vector{Int64},
                    inverse::Bool = false)::Cvoid
    for ext_idx = externalIndices
        tmp1 = sv[1 + ext_idx + indices[1]]
        sv[1 + ext_idx + indices[1]] = sv[1 + ext_idx + indices[2]]
        sv[1 + ext_idx + indices[2]] = tmp1
    end
    ;
end

Base.@ccallable function applyPauliX!(sv::StateVector, wire::Int64, inverse::Bool = false)::Cvoid
    gi = GateIndices([wire], sv.num_qubits);
    applyPauliX!(sv::StateVector, gi.internal, gi.external, inverse);
end

Base.@ccallable function applyHadamard!(sv::StateVector,
                        indices::Vector{Int64},
                        externalIndices::Vector{Int64},
                        inverse::Bool = false)::Cvoid
    for ext_idx = externalIndices
        v1 = sv[1 + ext_idx + indices[1]]
        v2 = sv[1 + ext_idx + indices[2]]

        sv[1 + ext_idx + indices[1]] = (1/sqrt(2))*(v1 + v2)
        sv[1 + ext_idx + indices[2]] = (1/sqrt(2))*(v1 - v2)
    end
    ;
end

Base.@ccallable function applyHadamard!(sv::StateVector, wire::Int64, inverse::Bool = false)::Cvoid
    gi = GateIndices([wire], sv.num_qubits);
    applyHadamard!(sv::StateVector, gi.internal, gi.external, inverse);
end

Base.@ccallable function applyRX!(sv::StateVector,
                    indices::Vector{Int64},
                    externalIndices::Vector{Int64},
                    param::Float64, inverse::Bool = false)::Cvoid
    c = cos(param / 2);
    js = (inverse == true) ?  -1im*sin(-param / 2) : 1im*sin(-param / 2);

    for ext_idx = externalIndices
        v1 = sv[1 + ext_idx + indices[1]]
        v2 = sv[1 + ext_idx + indices[2]]

        sv[1 + ext_idx + indices[1]] = c * v1 + js * v2;
        sv[1 + ext_idx + indices[2]] = js * v1 + c * v2;
    end
    ;
end

Base.@ccallable function applyRX!(sv::StateVector, wire::Int64, inverse::Bool = false)::Cvoid
    gi = GateIndices([wire], sv.num_qubits);
    applyRX!(sv::StateVector, gi.internal, gi.external, inverse);
end

Base.@ccallable function applyRY!(  sv::StateVector,
                    indices::Vector{Int64},
                    externalIndices::Vector{Int64},
                    param::Float64, inverse::Bool = false)::Cvoid
    c = cos(param / 2);
    s = (inverse == true) ? -sin(param / 2) : sin(param / 2);

    for ext_idx = externalIndices
        v1 = sv[1 + ext_idx + indices[1]]
        v2 = sv[1 + ext_idx + indices[2]]

        sv[1 + ext_idx + indices[1]] = c * v1 - s * v2;
        sv[1 + ext_idx + indices[2]] = s * v1 + c * v2;
    end
    ;
end

Base.@ccallable function applyRY!(sv::StateVector, wire::Int64, inverse::Bool)::Cvoid
    gi = GateIndices([wire], sv.num_qubits);
    applyRY!(sv::StateVector, gi.internal, gi.external, inverse);
end

end


using Revise
using UnivariateGaussianProcess

# An example usage of the package.
function run_gp()
    obj_func = sin
    bounds = (-π, π)
    cov_func = (x1,x2) -> matern1(x1, x2; σ=2)
    acq_func = ucb

    noise = 1e-2
    init_mean = (x) -> x
    x_points = [-π, 0., π]
    y_points = obj_func.(x_points)
    points = DATA_POINTS(x_points, y_points)

    return gaussian_process(
        obj_func,
        bounds,
        cov_func,
        acq_func;
        noise,
        init_mean,
        evals = 6,
        points,
        plot = true,
        plot_only_final = false,
    )
end


import jInv.ForwardShare.getData
function getData(m,pFor::ConvDiffParam)
	pFor.Ainv.doClear = 0 # PDE does not depend on m, so we shouldn't clear it

	U,pFor.Ainv    = solveLinearSystem(pFor.A, m + pFor.bc, pFor.Ainv)
	dobs = pFor.P*U
	pFor.Fields = U

	return dobs,pFor
end

import jInv.ForwardShare.getSensMatVec
function getSensMatVec(v::Vector{Float64},m::Vector{Float64},pFor::ConvDiffParam)
	return pFor.P*solveLinearSystem(pFor.A, v, pFor.Ainv)[1]
end

import jInv.ForwardShare.getSensTMatVec
function getSensTMatVec(v::Vector{Float64},m::Vector{Float64},pFor::ConvDiffParam)
	return solveLinearSystem(pFor.A, pFor.P'*v, pFor.Ainv,1)[1]
end


using ManifoldLearning

include("nearestneighbors.jl")

k=13
X, L = ManifoldLearning.swiss_roll()

# Use default distance matrix based method to find nearest neighbors
M1 = fit(Isomap, X)
Y1 = transform(M1)

# Use NearestNeighbors package to find nearest neighbors
M2 = fit(Isomap, X, nntype=KDTree)
Y2 = transform(M2)

# Use FLANN package to find nearest neighbors
M3 = fit(Isomap, X, nntype=FLANNTree)
Y3 = transform(M3)

using Plots
plot(
    plot(X[1,:], X[2,:], X[3,:], zcolor=L, m=2, t=:scatter3d, leg=false, title="Swiss Roll"),
    plot(Y1[1,:], Y1[2,:], c=L, m=2, t=:scatter, title="Distance Matrix"),
    plot(Y2[1,:], Y2[2,:], c=L, m=2, t=:scatter, title="NearestNeighbors"),
    plot(Y3[1,:], Y3[2,:], c=L, m=2, t=:scatter, title="FLANN")
, leg=false)


#=
Associated YT tutorial: https://youtu.be/TEWy9vZcxW4
=#

using Random

Random.seed!(0)

X = [ 1    2    3    2.5;
      2.0  5.0 -1.0  2.0;
     -1.5  2.7  3.3 -0.8]

# Julia doesn't have classes, so instead we make a struct.
mutable struct Layer_Dense
    weights::Array{Float64, 2}
    biases::Array{Float64, 2}
    output::Array{Float64, 2}
end

# Make a constructor for this struct.
function Layer_Dense(n_inputs, n_neurons)
    weights = 0.10 * randn(n_inputs, n_neurons)
    biases = zeros(1, n_neurons)

    # Initialize with an undefined output.
    output = Array{Float64, 2}(undef, n_inputs, n_neurons)

    Layer_Dense(weights, biases, output)
end

# And define the forward function for this struct.
function forward(layer::Layer_Dense, inputs)
    layer.output = inputs * layer.weights .+ layer.biases
end

# Now do the example.
layer1 = Layer_Dense(4,5)
layer2 = Layer_Dense(5,2)

forward(layer1, X)
forward(layer2, layer1.output)
println(layer2.output)


using SimpleGraphs, SimpleGraphAlgorithms, ProgressMeter

"""
    Trees(n)
This type is used for iterating over all trees with vertex set `{1,2,...,n}`, like this:

```
for T ∈ Trees(6)
    some_action(T)
end
```

**WARNING**: There are `n^(n-2)` trees on `n` vertices. With `n=10`, that's
one hundred million.
"""
struct Trees
    n::Int
end

function Base.iterate(T::Trees, state = 0)
    N = T.n
    NN = Int128(N)^(N - 2)
    if state == NN
        return nothing
    end

    code = Int.(digits(state, base = N, pad = N - 2) .+ 1)
    G = prufer_restore(code)
    name(G, "Tree")
    return (G, state + 1)
end

Base.length(T::Trees) = Int128(T.n)^(T.n - 2)


"""
    Schur complement structure of the form:
       M = A B
           C D
       |A B| |x| = |u|
       |C D| |y| = |v|
    M ∈ R{n+m,n+m}
    A ∈ R{n,n}
    D ∈ R{m,m}
    The Schur complement is taken about A, the top-left block:
        SC = D - C * A^-1 * B
"""
mutable struct Schur{T,n,m,nn,nm,mm}
    A::SMatrix{n,n,T,nn}
    B::SMatrix{n,m,T,nm}
    C::SMatrix{m,n,T,nm}
    D::SMatrix{m,m,T,mm}
    Ai::SMatrix{n,n,T,nn} # Inverse of A
    CAi::SMatrix{m,n,T,nm} # C*A^{-1}
    CAiB::SMatrix{m,m,T,mm} # C*A^{-1}*B
    gs_data::SDMGSSolver{m,T} # Schur complement of the block A in M
    u::SVector{n,T}
    v::SVector{m,T}
    x::SVector{n,T}
    y::SVector{m,T}
end

"""
    Initialize a Schur Object from a matrix M. The matrix M is split in 4 blocks:
        - top left block = n x n
        - bottom right block = m x m
"""
function Schur(M::AbstractMatrix{T}; n::Int=1, m::Int=size(M)[1]-n) where {T}
    @assert all(size(M) .== n+m)
    A = SMatrix{n,n,T,n^2}(M[1:n,      1:n])
    B = SMatrix{n,m,T,n*m}(M[1:n,      n .+ (1:m)])
    C = SMatrix{m,n,T,n*m}(M[n .+ (1:m),1:n])
    D = SMatrix{m,m,T,m^2}(M[n .+ (1:m),n .+ (1:m)])
    Ai = inv(A)
    CAi = C*Ai
    CAiB = C*Ai*B
    gs_data = SDMGSSolver(m,T=T)
    factorize!(gs_data, D - CAiB)
    u = zeros(SVector{n,T})
    v = zeros(SVector{m,T})
    x = zeros(SVector{n,T})
    y = zeros(SVector{m,T})
    return Schur{T,n,m,n^2,n*m,m^2}(A,B,C,D,Ai,CAi,CAiB,gs_data,u,v,x,y)
end

"""
    Update the Schur complement structure to handle a change in the D matrix.
    It recomputes the Gram-Schmidt factorization of the Schur complement of the block A.
"""
function schur_factorize!(S::Schur{T,n,m,nn,nm,mm}, D::AbstractMatrix{T}) where {T,n,m,nn,nm,mm}
    # update D, gs_data and the shur complement
    # we may want to set S.D to D so that it's clean, even if it's not necessary
    # S.D = D
    CAiB = S.CAiB
    gs_data = S.gs_data
    factorize!(gs_data, D - CAiB)
    return nothing
end

"""
    Solve for [x;y] given [u,v] and the factorized matrix M.
"""
function schur_solve!(S::Schur{T,n,m,nn,nm,mm}, u::AbstractVector{T}, v::AbstractVector{T}) where {T,n,m,nn,nm,mm}
    B  = S.B
    Ai = S.Ai
    CAi = S.CAi
    gs_data = S.gs_data
    S.u = u
    S.v = v
    us = S.u # staticarrays
    vs = S.v # staticarrays
    # S.x = Ai*(us + B*(As*(C*(Ai*us) - vs)))
    # S.y = As*(- C*(Ai*us) + vs)

    qr_solve!(gs_data, CAi*us - vs)
    temp = gs_data.xs
    S.x = Ai*(us + B*temp)
    S.y = - temp
    return nothing
end


@testset "Test Scalar Function Spaces" begin

using HDiscontinuousGalerkin
using Tensors

root_file = "mesh/figure2.1"
mesh = parse_mesh_triangle(root_file)

finiteElement = ContinuousLagrange{2,RefTetrahedron,1}()
Wh = ScalarFunctionSpace(mesh, finiteElement)

# Basic Test
sq2 = sqrt(2)
sq3 = sqrt(3)
invs = (Tensor{2,2}([1.0 1.0;-2.0 0.0]),
        Tensor{2,2}([-1.0 -1.0;2.0 0.0]),
        Tensor{2,2}([1.0 1.0;-1.0 1.0]),
        Tensor{2,2}([1.0 -1.0;0.0 2.0]))
detsJf = ([sq2/2,sq2/2,1],
          [sq2/2,sq2/2,1],
          [1,sq2/2,sq2/2],
          [sq2/2,sq2/2,1])
@test getnlocaldofs(Wh) == 3
@inbounds for (cellcount, cell) in enumerate(CellIterator(mesh))
    reinit!(Wh, cell)
    @test Wh.detJ[] ≈ 0.5
    @test getdetJdV(Wh,1)/Wh.qr_weights[1] ≈ 0.5
    @test Wh.Jinv[] ≈ invs[cellcount]
    # Face Data
    @test Wh.detJf ≈ detsJf[cellcount]
    if cellcount == 1
        @test Wh.normals ≈ [Vec{2}([-sq2/2,sq2/2]), Vec{2}([-sq2/2,-sq2/2]), Vec{2}([1.0,0.0])]
    end
end
end


using JuMP, CPLEX

#Preparing optimization model
model = Model(CPLEX.Optimizer)
set_optimizer_attribute(model, "CPX_PARAM_EPINT", 1e-8)

const N_GEN = 3
const N_HRS = 4

fixedCost = [5 7 6];
variableCost = [0.100 0.125 0.150];
startUpCost = [20 18 5];
shutdownCost = [0.5 0.3 1.0];

minPowerOut = [50 80 40];
maxPowerOut = [350 200 140];
rampDownLimit = [300 150 100];
rampShutDownLimit = [300 150 100];
rampUpLimit = [200 100 100];
rampStartUpLimit = [200 100 100];

demand = [100; 160; 500; 400];
reserves = [0; 16; 50; 40];

@variable(model, genOnOff[1:N_GEN, 1:N_HRS], Bin)
@variable(model, genOut[1:N_GEN, 1:N_HRS] >= 0)
@variable(model, genStartUp[1:N_GEN, 1:N_HRS], Bin)
@variable(model, genShutDown[1:N_GEN, 1:N_HRS], Bin)

#Setting the objective function
@objective(model, Min, sum(sum( fixedCost[g]*genOnOff[g,t] + variableCost[g]*genOut[g,t] + startUpCost[g]*genStartUp[g,t] + shutdownCost[g]*genShutDown[g, t] for g in 1:N_GEN) for t in 1:N_HRS) )

#Generator #1 should be ON at t=0
@constraint(model, constInit[t=1, g=3], genOnOff[g,t] == 1)

#Initial constraints
@constraint(model, constShutDn[g=1:N_GEN], genShutDown[g,1] == 0)
@constraint(model, constShutDn[g=1:N_GEN], genStartUp[g,1] == 0)

#Demand constraint
#for t in 1:N_HRS
@constraint(model, constDemand[t=1:N_HRS], sum(genOut[g, t] for g=1:N_GEN) == demand[t])
#end

#Power bounds
@constraint(model, constMaxPow[t=1:N_HRS, g=1:N_GEN],  genOut[g,t] <= maxPowerOut[g]*genOnOff[g,t] )
@constraint(model, constMinPow[t=1:N_HRS, g=1:N_GEN],  genOut[g,t] >= minPowerOut[g]*genOnOff[g,t] )

#Logical Conditions
@constraint(model, constLog1[t=2:N_HRS, g=1:N_GEN], (genStartUp[g,t] - genShutDown[g,t]) == (genOnOff[g,t] - genOnOff[g,t-1]))
@constraint(model, constLog2[t=1:N_HRS, g=1:N_GEN], genStartUp[g,t] + genShutDown[g, t] <= 1)

#Ramp limits
@constraint(model, constRampUp[t=2:N_HRS, g=1:N_GEN], genOut[g,t] - genOut[g,t-1] <= rampUpLimit[g]*genOnOff[g, t-1] + rampStartUpLimit[g]*genStartUp[g,t])
@constraint(model, constRampDown[t=2:N_HRS, g=1:N_GEN], (genOut[g,t-1] - genOut[g,t]) <= (rampDownLimit[g]*genOnOff[g,t] + rampShutDownLimit[g]*genShutDown[g,t]))

#Reserves
@constraint(model, constReserves[t=1:N_HRS], sum(maxPowerOut[g]*genOnOff[g,t] for g=1:N_GEN) >= (demand[t] + reserves[t]))

optimize!(model)

println("------------------------------------")
println("------- Optimal Solutions -------")
for g=1:N_GEN
  for t=1:N_HRS
    print(JuMP.value(genOut[g,t]),"\t")
  end
  println("\t Generator: $g")
end

println("------------------------------------")
println("------- GENERATORS COMMITED --------")

for g=1:N_GEN
  for t=1:N_HRS
    print(JuMP.value(genOnOff[g,t]),"\t")
  end
  println("\t Generator: $g")
end


a = 1
c = -5
limit = 10

mf = SigmoidMF(a, c, limit)

# Evaluation tests
@assert mf.eval(c) == 0.5
@assert mf.eval(c - 5) < mf.eval(c + 5)

# Mean finding tests
@assert mf.mean_at(1) > mf.mean_at(0.6) > mf.mean_at(0.3) > mf.mean_at(0)


# --- L2 Cost

type L2Cost <: Layer
    output
    grad_input
    L2Cost() = new()
end
export L2Cost

# Defines forward & backward
forward!(cost::L2Cost, inputs) = cost.output = .5 * norm(inputs[1] - inputs[2])
backward!(cost::L2Cost, inputs, scale=1.) = cost.grad_input = scale * (inputs[1] - inputs[2])


struct Gaussian{K} <: WhittleLikelihoodInference.UnknownAcvTimeSeriesModel{1,Float64}
    α::Float64
    ωₚ::Float64
    κ::Float64

    norm::Float64
    halfκ⁻²::Float64
    function Gaussian{K}(α,ωₚ,κ) where {K}
        α > 0 || throw(ArgumentError("α must be > 0."))
        ωₚ > 0 || throw(ArgumentError("ωₚ must be > 0."))
        κ > 0 || throw(ArgumentError("κ must be > 0."))
        new(α,ωₚ,κ,α/(κ*sqrt(2π)),inv(2κ^2))
    end
    function Gaussian{K}(x) where {K}
        @boundscheck checkparameterlength(x,Gaussian{K})
        @inbounds Gaussian{K}(x[1],x[2],x[3])
    end
end

@fastmath function WhittleLikelihoodInference.sdf(model::Gaussian, ω::Real)
    return model.norm*exp(-(model.ωₚ-abs(ω))^2*model.halfκ⁻²)
end

WhittleLikelihoodInference.npars(::Type{Gaussian{K}}) where {K} = 3
WhittleLikelihoodInference.nalias(::Gaussian{K}) where {K} = K

WhittleLikelihoodInference.lowerbounds(::Type{Gaussian{K}}) where {K} = [0,  0, 0]
WhittleLikelihoodInference.upperbounds(::Type{Gaussian{K}}) where {K} = [Inf,Inf,Inf]

Gaussian(x::AbstractVector{Float64}) = Gaussian(1,1,1)
Gaussian(α,ωₚ,κ) = error("Gaussian process requires the ammount of aliasing specified as a type parameter. Use Gaussian{K}() where K ∈ N.")

Base.@propagate_inbounds @fastmath function WhittleLikelihoodInference.grad_add_sdf!(out, model::Gaussian, ω::Real)
    @boundscheck checkbounds(out,1:3)
    @inbounds begin
        ω = abs(ω)
        sdf = model.norm*exp(-(model.ωₚ-ω)^2 * model.halfκ⁻²)
        out[1] += sdf/model.α
        out[2] += 2(ω-model.ωₚ) * model.halfκ⁻² * sdf
        out[3] += sdf * ((ω-model.ωₚ)^2/model.κ^3-1/model.κ)
    end
    nothing
end


L1Loss(x, y) = mean(abs(x - y))

function NLLLoss(input::Array{Float64, 2}, weights::Array{Float64, 1}, size_average::Bool=true)
    if size_average
        return map((x, y) -> mean(-x * y), input, weights)
    else
        return map((x, y) -> sum(-x * y), input, weights)
    end
end

KLDivLoss(x, target) = mean(target .* (log(target) - x))

MSELoss(x, y) = mean((x - y) .* (x - y))

BCELoss(o, t, weights, size_average::Bool=true) = -mean(weights .* (t .* log(o) + (1 - t) * log(1 - o)))

HigheEmbeddingLoss(x, y, margin=1) = map((x, y) -> y == 1 ? x : max(0, margin - x), x, y)

SoftMarginLoss(x, y) = mean(log(1 + exp(- y .* x)))

CrossEntropyLoss(x, class, weights) = weights[class] * (-x[class] + log(sum(exp(x))))

CosineEmbeddingLoss(x1, x2, y, margin=0) = map((x1, x2, y) -> y == 1 ? 1 - cos(x1, x2) : max(0, cos(x1, x2) - margin), x1, x2, y)


function frankwolfe(network::AbstractNetwork,
              trips::AbstractMatrix{T},
              costfn::Function;
              basedon=:link,
              errtol=1e-4) where {T<:Real}
    # initialize
    flows = allornothing(network, trips, costfn; basedon=basedon)

    # start iteration
    err = 1.
    while err > errtol
        ## find target solution
        linkcosts, = costfn(flows)
        shortflows = allornothing(network, trips, linkcosts; basedon=basedon)

        ## find stepsize
        Δflows = shortflows - flows
        μ = bisection((μ) -> sum(costfn(flows + μ * Δflows).costs .* Δflows))

        ## calculate new link travel times
        flows += μ * Δflows

        ## calculate error
        err = sum(flows .* linkcosts)/sum(shortflows .* linkcosts) - 1.
    end
    return (flows, err)
end


function ellec_bulirsch(k2::Real)
ca = sqrt(eps(k2))
if k2 == 1.0
  return 1.0
end
m=1.0
a=1.0
b=1.0-k2
kc=sqrt(b)
c=a
a+=b

b=2.0*(c*kc+b)
c=a
m0=m
m+=kc
a+=b/m
while abs(m0-kc) > (ca*m0)
  kc=2.0*sqrt(kc*m0)
  b=2.0*(c*kc+b)
  c=a
  m0=m
  m+=kc
  a+=b/m
#  println(a,' ',b,' ',c,' ',m0,' ',kc)
end
return pi*0.25*a/m
end


@testset "eval scalar" begin
    mmesh = Mesh(
        [2.2 2
        3. 4.
        2.5 10.], [1 2 3]
    )
    edge = [1 2]
    out = eval_scalar_on_boundary_edge([5.0;6.0;10000.0], edge, mmesh)
    sess = Session(); init(sess)
    @test run(sess, out)≈[5.069431844202973
    5.330009478207572
    5.669990521792428
    5.930568155797027]
end


using MultivariateSeries

X = @ring x0 x1 x2
n = length(X)
d = 4
r = 4

Xi = rand(n,r)
w = fill(1.0,r)
T = tensor(w,Xi,X, d)

k = 2
H = hankel(T,k)
P = perp(T,k)

@assert length(P)==2


# Normalization: remove trailing zeros in coefficient
function normalize(x::Decimal; rounded::Bool=false)
    p = 0
    if x.c != 0
        while x.c % 10^(p+1) == 0
            p += 1
        end
    end
    c = BigInt(x.c / 10^p)
    q = (c == 0 && x.s == 0) ? 0 : x.q + p
    if rounded
        Decimal(x.s, abs(c), q)
    else
        round(Decimal(x.s, abs(c), q), digits=DIGITS, normal=true)
    end
end


#
# abstract interfaces
#

abstract type AbstractCurve end

# Required methods
"""
	point(C::AbstractCurve,t::Real)
Find the point on curve `C` at parameter value `t`, which should lie in the interval [0,1].
"""
point(c::AbstractCurve,t::Real) = @error "No point() method defined for type $(typeof(c))"

"""
	tangent(C::AbstractCurve,t::Real)
Find the complex number representing the tangent to curve `C` at parameter value `t` in [0,1].
"""
tangent(C::AbstractCurve,t::Real) = @error "No tangent() method defined for type $(typeof(C))"

reverse(C::AbstractCurve) = @error "No reverse() method defined for type $(typeof(C))"

"""
	isfinite(C::AbstractCurve)
Return `true` if the curve is bounded in the complex plane (i.e., does not pass through infinity).
"""
isfinite(C::AbstractCurve) = @error "No isfinite() method defined for type $(typeof(C))"

conj(C::AbstractCurve) = @error "No conj() method defined for type $(typeof(C))"
+(C::AbstractCurve,z::Number) = @error "No addition method defined for type $(typeof(C))"
-(C::AbstractCurve) = @error "No negation method defined for type $(typeof(C))"
*(C::AbstractCurve,z::Number) = @error "No multiplication method defined for type $(typeof(C))"
inv(C::AbstractCurve) = @error "No inversion method defined for type $(typeof(C))"

# Default implementations
"""
	point(C::AbstractCurve,t::AbstractArray)

Vectorize the `point` function for curve `C`.
"""
point(c::AbstractCurve,t::AbstractArray{T}) where T<:Real = [point(c,t) for t in t]

"""
	unittangent(C::AbstractCurve,t::Real)
Find the complex number representing the unit tangent to curve `C` at parameter value `t` in [0,1]. For Lines, Segments, and Rays, the `t` argument is optional.
"""
unittangent(C::AbstractCurve,t::Real) = sign(tangent(C,t))

"""
	normal(C::AbstractCurve,t::Real)
Find the unit complex number in the direction of the leftward-pointing normal to curve `C` at parameter value `t` in [0,1].
"""
normal(c::AbstractCurve,t::Real) = 1im*unittangent(c,t)

+(z::Number,C::AbstractCurve) = +(C,z)
-(C::AbstractCurve,z::Number) = C + (-z)
-(z::Number,C::AbstractCurve) = z + (-C)
*(z::Number,C::AbstractCurve) = *(C,z)
/(C::AbstractCurve,z::Number) = C*(1/z)
/(z::Number,C::AbstractCurve) = z*inv(C)

function arclength(C::AbstractCurve,part=[0,1])
	f = t -> abs(tangent(C,t))
	intadapt(f,part...,DEFAULT[:tol])
end

"""
	plotdata(C::AbstractCurve,n=501)

Compute `n` points along the curve `C` suitable to make a plot of it.
"""
plotdata(C::AbstractCurve) = adaptpoints(t->point(C,t),t->unittangent(C,t),0,1)

show(io::IO,C::AbstractCurve) = print(io,"Complex-valued $(typeof(C))")
show(io::IO,::MIME"text/plain",C::AbstractCurve) = print(io,"Complex-valued $(typeof(C))")

# AbstractClosedCurve
abstract type AbstractClosedCurve <: AbstractCurve end

# Default implementations
function winding(C::AbstractClosedCurve,z::Number)
	# Integrate around the curve
	f = t -> imag(tangent(C,t)/(point(C,t)-z))
	w = intadapt(f,0,1,1e-4)
	return round(Int,w/(2π))
end

isinside(z::Number,C::AbstractClosedCurve) = winding(C,z) != 0
isoutside(z::Number,C::AbstractClosedCurve) = winding(C,z) == 0

#
# generic Curve
#

"""
(type) Smooth curve defined by an explicit function of a real paramerter in [0,1].
"""
struct Curve <: AbstractCurve
	point
	tangent
end

"""
	Curve(f)
	Curve(f,a,b)
Construct a `Curve` object from the complex-valued function `f` accepting an argument in the interval [0,1]. If `a` and `b` are given, they are the limits of the parameter in the call to the supplied `f`. However, the resulting object will be defined on [0,1], which is internally scaled to [a,b].

	Curve(f,df[,a,b])
Construct a curve with point location and tangent given by the complex-valued functions `f` and `df`, respectively, optionally with given limits on the parameter.
"""
Curve(f) = Curve(f,t->fdtangent(f,t))
Curve(f,df,a::Real,b::Real) = Curve(t->f(scaleto(a,b,t)),t->df(scaleto(a,b,t)))
Curve(f,a::Real,b::Real) = Curve(t->f(scaleto(a,b,t)))

# Required methods
point(C::Curve,t::Real) = C.point(t)
(C::Curve)(t::Real) = point(C,t)
tangent(C::Curve,t::Real) = C.tangent(t)
conj(C::Curve) = Curve(t->conj(C.point(t)))
reverse(C::Curve) = Curve(t->C.point(1-t))
isfinite(C::Curve) = true

+(C::Curve,z::Number) = Curve(t->C.point(t)+z,C.tangent)
-(C::Curve) = Curve(t->-C.point(t),t->-C.tangent(t))
*(C::Curve,z::Number) = Curve(t->C.point(t)*z,t->C.tangent(t)*z)
inv(C::Curve) = Curve(t->1/C.point(t),t->-C.tangent(t)/C.point(t)^2)

#
# generic ClosedCurve
#

"""
(type) Smooth closed curve defined by an explicit function of a real paramerter in [0,1].
"""
struct ClosedCurve <: AbstractClosedCurve
	curve::Curve
	function ClosedCurve(c::Curve;tol=DEFAULT[:tol])
		@assert isapprox(point(c,0),point(c,1);rtol=tol,atol=tol) "Curve does not close"
		new(c)
	end
end

"""
	ClosedCurve(f; tol=<default>)
	ClosedCurve(f,a,b; tol=<default)
Construct a `ClosedCurve` object from the complex-valued function `point` accepting an argument in the interval [0,1]. The constructor checks whether `f(0)≈f(1)` to tolerance `tol`. If `a` and `b` are given, they are the limits of the parameter in the call to the supplied `f`. However, the resulting object will be defined on [0,1], which is internally scaled to [a,b].

	ClosedCurve(f,df[,a,b]; tol=<default>)
Construct a closed curve with point location and tangent given by the complex-valued functions `f` and `df`, respectively, optionally with given limits on the parameter.
"""
ClosedCurve(f,df=t->fdtangent(f,t);kw...) = ClosedCurve(Curve(f,df;kw...))
ClosedCurve(f,a::Real,b::Real;kw...) = ClosedCurve(Curve(f,a,b;kw...))
ClosedCurve(f,df,a::Real,b::Real;kw...) = ClosedCurve(Curve(f,df,a,b;kw...))

point(C::ClosedCurve,t::Real) = point(C.curve,t)
(C::ClosedCurve)(t::Real) = point(C.curve,t)
tangent(C::ClosedCurve,t::Real) = tangent(C.curve,t)
conj(C::ClosedCurve) = ClosedCurve(conj(C.curve))
reverse(C::ClosedCurve) = ClosedCurve(reverse(C.curve))
isfinite(C::ClosedCurve) = isfinite(C.curve)
+(C::ClosedCurve,z::Number) = ClosedCurve(C.curve+z)
-(C::ClosedCurve) = ClosedCurve(-C.curve)
*(C::ClosedCurve,z::Number) = ClosedCurve(C.curve*z)
inv(C::ClosedCurve) = ClosedCurve(inv(C.curve))

include("lines.jl")
include("rays.jl")
include("segments.jl")
include("circles.jl")
include("arcs.jl")
include("intersections.jl")


function freqresp(sys::DelayLtiSystem, ω::AbstractVector{T}) where {T <: Real}
    ny = noutputs(sys)
    nu = ninputs(sys)

    P_fr = freqresp(sys.P.P, ω);

    G_fr = zeros(eltype(P_fr), length(ω), ny, nu)

    for ω_idx=1:length(ω)
        P11_fr = P_fr[ω_idx, 1:ny, 1:nu]
        P12_fr = P_fr[ω_idx, 1:ny, nu+1:end]
        P21_fr = P_fr[ω_idx, ny+1:end, 1:nu]
        P22_fr = P_fr[ω_idx, ny+1:end, nu+1:end]

        delay_matrix_inv_fr = Diagonal(exp.(im*ω[ω_idx]*sys.Tau)) # Frequency response of the diagonal matrix with delays
        # Inverse of the delay matrix, so there should not be any minus signs in the exponents

        G_fr[ω_idx,:,:] .= P11_fr + P12_fr/(delay_matrix_inv_fr - P22_fr)*P21_fr # The matrix is invertible (?!)
    end

    return G_fr
end

function evalfr(sys::DelayLtiSystem, s)
    (ny, nu) = size(sys)

    P_fr = evalfr(sys.P.P, s)

    P11_fr = P_fr[1:ny, 1:nu]
    P12_fr = P_fr[1:ny, nu+1:end]
    P21_fr = P_fr[ny+1:end, 1:nu]
    P22_fr = P_fr[ny+1:end, nu+1:end]

    delay_matrix_inv_fr = Diagonal(exp.(s*sys.Tau))

    return P11_fr + P12_fr/(delay_matrix_inv_fr - P22_fr)*P21_fr
end

"""
    delayd_ss(τ, Ts)
Discrete-time statespace realization of a delay τ sampled with period Ts,
i.e. of z^-N where N = τ / Ts.

τ must be a multiple of Ts.
"""
function delayd_ss(τ::Number, Ts::Number)
    n = Int(round(τ / Ts))
    if !(τ - n*Ts ≈ 0)
        error("The delay τ must be a multiple of the sample time Ts")
    end
    ss(diagm(1 => ones(n-1)), [zeros(n-1,1); 1], [1 zeros(1,n-1)], 0, Ts)
end

"""
    c2d(G::DelayLtiSystem, Ts, method=:zoh)
"""
function c2d(G::DelayLtiSystem, Ts::Real, method=:zoh)
    if !(method === :zoh)
        error("c2d for DelayLtiSystems only supports zero-order hold")
    end
    X = append([delayd_ss(τ, Ts) for τ in G.Tau]...)
    Pd = c2d(G.P.P, Ts)
    return lft(Pd, X)
end

"""
    `y, t, x = lsim(sys::DelayLtiSystem, u, t::AbstractArray{<:Real}; x0=fill(0.0, nstates(sys)), alg=MethodOfSteps(Tsit5()), abstol=1e-6, reltol=1e-6, kwargs...)`

    Simulate system `sys`, over time `t`, using input signal `u`, with initial state `x0`, using method `alg` .

    Arguments:

    `t`: Has to be an `AbstractVector` with equidistant time samples (`t[i] - t[i-1]` constant)
    `u`: Function to determine control signal `ut` at a time `t`, on any of the following forms:
        Can be a constant `Number` or `Vector`, interpreted as `ut .= u` , or
        Function `ut .= u(t)`, or
        In-place function `u(ut, t)`. (Slightly more effienct)
    `alg, abstol, reltol` and `kwargs...`: are sent to `DelayDiffEq.solve`.

    Returns: times `t`, and `y` and `x` at those times.
"""
function lsim(sys::DelayLtiSystem{T,S}, u, t::AbstractArray{<:Real};
        x0=fill(zero(T), nstates(sys)),
        alg=DelayDiffEq.MethodOfSteps(Tsit5()),
        abstol=1e-6, reltol=1e-6,
        kwargs...) where {T,S}

    # Make u! in-place function of u
    u! = if isa(u, Number) || isa(u,AbstractVector) # Allow for u to be a constant number or vector
        (uout, t) -> uout .= u
    elseif DiffEqBase.isinplace(u, 2)               # If u is an inplace (more than 1 argument function)
        u
    else                                            # If u is a regular u(t) function
        (out, t) -> (out .= u(t))
    end

    _lsim(sys, u!, t, x0, alg; abstol=abstol, reltol=reltol, kwargs...)
end

# Generic parametrized dde used for simulating DelayLtiSystem
# We simulate the integral of the states
# u, du are from the notation of variables in DifferentialEquations.jl
# The state from the control system is x
# u!, uout is the control law and its output
function dde_param(du, u, h, p, t)
    A, B1, B2, C1, C2, D11, D12, D21, D22, y, Tau, u!, uout, hout, tmpy, tsave = p

    nx = size(A,1)
    nd = length(Tau)
    ny = size(C1,1)

    dx = view(du, 1:nx)
    dY = view(du, (nx+1):(nx+ny))
    dD = view(du, (nx+ny+1):(nx+ny+nd))
    x = view(u, 1:nx)

    # uout = u(t)
    u!(uout, t)

    # hout = d(t-Tau)
    for k=1:length(Tau)
        # Get the derivative of the history for d at indices corresponding to Tau
        hout[k] = h(p, t-Tau[k], Val{1}, idxs=(nx+ny+k))
    end

    #dx(t) .= A*x(t) + B1*u(t) +B2*d(t-tau)
    mul!(dx, A, x)
    mul!(dx, B1, uout, true, true)
    mul!(dx, B2, hout, true, true)

    # dY = y(t) = C1*x + D11*u(t) + D12*d(t-Tau)
    mul!(dY, C1, x)
    mul!(dY, D11, uout, true, true)
    mul!(dY, D12, hout, true, true)

    # dD = d(t) = C2*x + D21*u(t) + D22*d(t-Tau)
    mul!(dD, C2, x)
    mul!(dD, D21, uout, true, true)
    mul!(dD, D22, hout, true, true)

    # Save y value in tmpy to be used by dde_saver
    if t in tsave
        # The value of y at this time is given by the derivative
        tmpy .= dY
    end
    return
end

# Save x(t) and y(t) to output
function dde_saver(u,t,integrator)
    A, B1, B2, C1, C2, D11, D12, D21, D22, y, Tau, u!, uout, hout, tmpy, tsave = integrator.p
    nx = size(A,1)
    # y is already saved in tmpy
    u[1:nx], copy(tmpy)
end

function _lsim(sys::DelayLtiSystem{T,S}, Base.@nospecialize(u!), t::AbstractArray{<:Real}, x0::Vector{T}, alg; kwargs...) where {T,S}

    P = sys.P

    t0 = first(t)
    dt = t[2] - t[1]

    # Get all matrices to save on allocations
    A, B1, B2, C1, C2, D11, D12, D21, D22 = P.A, P.B1, P.B2, P.C1, P.C2, P.D11, P.D12, P.D21, P.D22
    Tau = sys.Tau
    nx = size(A,1)
    nd = length(Tau)
    ny = size(C1,1)
    nt = length(t)

    hout = fill(zero(T), nd)            # in place storage for delays
    uout = fill(zero(T), ninputs(sys))  # in place storage for u
    tmpy = similar(x0, ny)              # in place storage for output
    y = Matrix{T}(undef, ny, nt)        # Output matrix
    x = Matrix{T}(undef, nx, nt)        # State matrix

    p = (A, B1, B2, C1, C2, D11, D12, D21, D22, y, Tau, u!, uout, hout, tmpy, t)

    # This callback computes and stores the delay term
    sv = SavedValues(Float64, Tuple{Vector{Float64},Vector{Float64}})
    cb = SavingCallback(dde_saver, sv, saveat = t)

    # History function, only used for d
    # The true d(t) is stored as a derivative
    h!(p, t, deriv::Type{Val{1}}; idxs=0) = zero(T)

    # The states are x(t), Y(t), D(t), where Y, D are integrals of y(t), d(t)
    u0 = [x0;zeros(T,ny);zeros(T,nd)]
    prob = DDEProblem{true}(dde_param, u0, h!,
                (T(t[1]), T(t[end])),
                p,
                constant_lags=sort(Tau),# Not sure if sort needed
                neutral=true,           # We have derivatives on RHS (d(t))
                callback=cb)
    # Important to stop at t since we can not access derivatives in SavingCallback
    sol = DelayDiffEq.solve(prob, alg; tstops=t, saveat=t, kwargs...)

    # Retrive the saved values
    for k = 1:nt
        x[:,k] .= sv.saveval[k][1]
        y[:,k] .= sv.saveval[k][2]
    end

    return y', t, x'
end

# We have to default to something, look at the sys.P.P and delays
function _bounds_and_features(sys::DelayLtiSystem, plot::Val)
    ws, pz =  _bounds_and_features(sys.P.P, plot)
    logtau = log10.(abs.(sys.Tau))
    logtau = filter(x->x>4, logtau) # Ignore low frequency
    if isempty(logtau)
        return ws, pz
    end
    extreme = extrema(logtau)
    return [min(ws[1], floor(extreme[1]-0.2)), max(ws[2], ceil(extreme[2]+0.2))], pz
end

# Againm we have to do something for default vectors, more or less a copy from timeresp.jl
function _default_dt(sys::DelayLtiSystem)
    if !isstable(sys.P.P)
        return 0.05   # Something small
    else
        ps = pole(sys.P.P)
        r = minimum([abs.(real.(ps));0]) # Find the fastest pole of sys.P.P
        r = min(r, minimum([sys.Tau;0])) # Find the fastest delay
        if r == 0.0
            r = 1.0
        end
        return 0.07/r
    end
end

function Base.step(sys::DelayLtiSystem{T}, t::AbstractVector; kwargs...) where T
    nu = ninputs(sys)
    if t[1] != 0
        throw(ArgumentError("First time point must be 0 in step"))
    end
    u = (out, t) -> (t < 0 ? out .= 0 : out .= 1)
    x0=fill(zero(T), nstates(sys))
    if nu == 1
        y, tout, x = lsim(sys, u, t; x0=x0, kwargs...)
    else
        x = Array{T}(undef, length(t), nstates(sys), nu)
        y = Array{T}(undef, length(t), noutputs(sys), nu)
        for i=1:nu
            y[:,:,i], tout, x[:,:,i] = lsim(sys[:,i], u, t; x0=x0, kwargs...)
        end
    end
    return y, tout, x
end

function impulse(sys::DelayLtiSystem{T}, t::AbstractVector; alg=MethodOfSteps(BS3()), kwargs...) where T
    nu = ninputs(sys)
    iszero(sys.P.D12) || @warn("Impulse with a direct term from input to delay vector leads to poor accuracy.")
    if t[1] != 0
        throw(ArgumentError("First time point must be 0 in impulse"))
    end
    u = (out, t) -> (out .= 0)
    if nu == 1
        y, tout, x = lsim(sys, u, t; alg=alg, x0=sys.P.B[:,1], kwargs...)
    else
        x = Array{T}(undef, length(t), nstates(sys), nu)
        y = Array{T}(undef, length(t), noutputs(sys), nu)
        for i=1:nu
            y[:,:,i], tout, x[:,:,i] = lsim(sys[:,i], u, t; alg=alg, x0=sys.P.B[:,i], kwargs...)
        end
    end
    return y, tout, x
end

# Used for pade approximation
"""
`p2 = _linscale(p::Polynomial, a)`

Given a polynomial `p` and a number `a, returns the polynomials `p2` such that
`p2(s) == p(a*s)`.
"""
function _linscale(p::Polynomial, a)
    # This function should perhaps be implemented in Polynomials.jl
    coeffs_scaled = similar(p.coeffs, promote_type(eltype(p), typeof(a)))
    a_pow = 1
    coeffs_scaled[1] = p.coeffs[1]
    for k=2:length(p.coeffs)
        a_pow *= a
        coeffs_scaled[k] = p.coeffs[k]*a_pow
    end
    return Polynomial(coeffs_scaled)
end

# Coefficeints for Padé approximations
# PADE_Q_COEFFS = [Polynomial([binomial(N,i)*prod(N+1:2*N-i) for i=0:N]) for N=1:10]
const PADE_Q_COEFFS =  [[2, 1],
 [12, 6, 1],
 [120, 60, 12, 1],
 [1680, 840, 180, 20, 1],
 [30240, 15120, 3360, 420, 30, 1],
 [665280, 332640, 75600, 10080, 840, 42, 1],
 [17297280, 8648640, 1995840, 277200, 25200, 1512, 56, 1],
 [518918400, 259459200, 60540480, 8648640, 831600, 55440, 2520, 72, 1],
 [17643225600, 8821612800, 2075673600, 302702400, 30270240, 2162160, 110880, 3960, 90, 1],
 [670442572800, 335221286400, 79394515200, 11762150400, 1210809600, 90810720, 5045040, 205920, 5940, 110, 1]]

"""
    pade(τ::Real, N::Int)

Compute the `N`th order Padé approximation of a time-delay of length `τ`.
"""
function pade(τ::Real, N::Int)
    if !(1 <= N <= 10); error("Order of Padé approximation must be between 1 and 10. Got $N."); end

    Q = Polynomials.Polynomial(PADE_Q_COEFFS[N])

    return tf(_linscale(Q, -τ), _linscale(Q, τ)) # return Q(-τs)/Q(τs)
end

"""
    pade(G::DelayLtiSystem, N)

Approximate all time-delays in `G` by Padé approximations of degree `N`.
"""
function pade(G::DelayLtiSystem, N)
    ny, nu = size(G)
    nTau = length(G.Tau)
    X = append([ss(pade(τ,N)) for τ in G.Tau]...) # Perhaps append should be renamed blockdiag
    return lft(G.P.P, X)
end


using ModelConstructors, FileIO, Random, SMC, JLD2

### Construct a generic model and populate it with parameters
reg = GenericModel()
fn = dirname(@__FILE__)
reg <= Setting(:dataroot, "$(fn)/save/input_data/")
reg <= Setting(:saveroot, "$(fn)/save/")

reg <= parameter(:α1, 0., (-1e5, 1e5), (-1e5, 1e5), Untransformed(), Normal(0, 10), fixed = false)
reg <= parameter(:β1, 0., (-1e5, 1e5), (-1e5, 1e5), Untransformed(), Normal(0, 10), fixed = false)

### Estimate with SMC

## Make data
X = rand(100)
β = 1.
α = 1.
y = β*X + α
N = 1
data = hcat(y, X)
jldopen(dataroot(reg)*"reg_data.jld2", "w") do file
    file["data"] = data
end

## Construct likelihood function:
function likelihood_fnct(p, d)
    α = p[1]
    β = p[2]
    Σ = 1 #ones(N,N)
    det_Σ = det(Σ)
    inv_Σ = inv(Σ)
    term1 = -N / 2 * log(2 * π) - 1 /2 * log(det_Σ)
    logprob = 0.
    errors = d[:, 1] .- α .- β .* d[:, 2]
    for t in 1:size(d,1)
        logprob += term1 - 1/2 * dot(errors, inv_Σ * errors)
    end
    return logprob
end

Random.seed!(1793)
println("Starting to estimate Regression with SMC . . .")
@everywhere using SMC
smc(likelihood_fnct, reg.parameters, data, n_parts = 10, use_fixed_schedule = false, tempering_target = 0.97)


function GenerateExamples(example::Char)
    # Generate damping examples of dimension n=801, 1601, 2001
    # with example=='S','L'or 'H', for Small, Large and Homogeneous, respectively.

    if example=='S'
        n₁=400
    elseif example=='L'
        n₁=800
    else
        n₁=1000
    end

    # Dimension of the system
    n=2*n₁+1
    m₁=1
    m₂=1

    if example=='S'
        m₀=1200.0
    elseif example=='L'
        m₀=1800.0
    else
        m₀=2000.0
    end

    k₁=100
    k₂=150
    k₃=200

    first_part=Array{Int,1}(undef,n₁)
    second_part=Array{Int,1}(undef,n₁)

    if example=='S'
        for ii=1:400
            first_part[ii]=1000-4*ii
        end
        for ii=201:400
            first_part[ii]=3*ii-400
        end
        for ii=1:400
            second_part[ii]=500+ii
        end
    elseif example=='L'
        for ii=1:div(n₁,2)
            first_part[ii]=2000-4*ii
        end
        for ii=div(n₁,2)+1:n₁
            first_part[ii]=3*ii-800
        end
        for ii=1:n₁
            second_part[ii]=500+ii
        end
    else
        for ii=1:n₁
           first_part[ii]=1000
        end
        for ii=1:n₁
          second_part[ii]=1500
        end
    end

    M=Diagonal([first_part; second_part; m₀])

    K₁=SymTridiagonal(2*ones(n₁),-ones(n₁-1))
    α=0.02
    K₁=Matrix(K₁)

    K=zeros(2*n₁+1,2*n₁+1)
    K[1:n₁,1:n₁]=k₁*K₁
    K[n₁+1:2*n₁,n₁+1:2*n₁]=k₂*K₁

    K[1:n₁,end]=-[zeros(1,n₁-1) k₁]'
    K[n₁+1:2*n₁,end]=-[zeros(1,n₁-1) k₂]'

    K[end,1:n₁]=-[zeros(1,n₁-1) k₁]
    K[end,n₁+1:2*n₁]=-[zeros(1,n₁-1) k₂]
    K[2*n₁+1, 2*n₁+1]=k₁+k₂+k₃

    # Return stiffness and mass matrix
    return K, M
end


#TODO: Some kind of generator/iterator/resumable function would be prefered here...
function read_return(episode, i, γ, non_termination_penalty)
    s, a, r = episode[i]
    g = r
    w = 1
    for j = (i+1):length(episode)-1
        w = w * γ
        _, _, r1 = episode[j]
        g += w * r1
    end
    s_fin, _, _ = episode[end]
    if s_fin != TERMINAL_STATE_CODE
        g -= non_termination_penalty
    end
    return s, a, g
end

function update_state_action_values_repository!(R, episode, γ::Real, non_termination_penalty)
    for i = 1:length(episode)-1
        s, a, g = read_return(episode, i, γ, non_termination_penalty)
        ∑r, n = R[s, a]
        R[s, a] = (∑r + g, n + 1)
    end
end

function Q_from_state_action_values_repository!(Q, R)
    for s = 1:size(Q, 1)
        for a = 1:size(Q, 2)
            ∑r, n = R[s, a]
            if n != 0
                Q[s, a] = ∑r / n
            else
                Q[s, a] = NaN
            end
        end
    end
end

"""
    mk_evaluate_policy!(V, Q, mdp, policy, γ; tol, maxiter)

Evaluates (deterministic or stochastic) decision policy `policy` for the given Markov
decision process `mdp` and given discount factor `γ`.

By saying that a policy has been "evaluated", we imply that state values `V` and
state-action values `Q` are computed for this policy. This function evaluates policies in-place,
approximately, in an iterrative manner, until maximal difference between state values in consequtive
iterations is less than the specified `tol`, or until a maximal number of iterations `maxiter`
is performed.

The function returns a 2-tuple containing: the number of iterations performed, and a boolean
indicator of convergence (which is `true` if the vector of state values is evaluated with the
required precission, and `false` otherwise).
"""
function mk_evaluate_policy!(Q::AbstractMatrix{<:Real}, policy_simulator, γ::Real; maxiter = 100, non_termination_penalty = 1000)
    R = fill((0.0, 0), size(Q))
    for _ = 1:maxiter
        s0 = rand(1:size(Q, 1))
        a0 = rand(1:size(Q, 2))
        episode = policy_simulator(s0, a0)
        update_state_action_values_repository!(R, episode, γ, non_termination_penalty)
    end # for: iterations
    Q_from_state_action_values_repository!(Q, R)
end

function mk_update_Q!(Q, ω, policy_simulator, γ; non_termination_penalty = 1000)
    s0 = rand(1:size(Q, 1))
    a0 = rand(1:size(Q, 2))
    episode = policy_simulator(s0, a0)
    Δ = -Inf
    for i = 1:length(episode)-1
        s, a, g = read_return(episode, i, γ, non_termination_penalty)
        new_value = ω * Q[s, a] + (1 - ω) * g
        Δ = max(Δ, abs(new_value - Q[s, a]))
        Q[s, a] = new_value
    end
    return Δ
end


"""
Handles the non-relativistic NBody problem for planets,
and other celestial bodies.
"""
module NBody

using Reexport
@reexport using ..AstrodynamicsCore

include("../Misc/DocStringExtensions.jl")
include("../Misc/UnitfulAliases.jl")

using StaticArrays: SVector, @SVector, SMatrix, @SMatrix

export Body, NBodySystem, system_energy,
       system_angular_momentum, promote, convert,
       Float16, Float32, Float64, BigFloat,
       length, getindex

include("NBodyStates.jl")
include("NBodyCalculations.jl")

end


export linpath
export momentumpath

"""
    momentumpath

Generate a list of momenta

# Arguments
* `anchorpoints`
* (Optional) `nseg` - number of points in each segment
"""
function linpath(anchors::AbstractVector...; nseg::Integer=100)
    n = length(anchors)
    if n <= 1
        throw(ArgumentError("Number of anchor points should be more than one"))
    end
    d = length(anchors[1])
    for i in 2:n
        if length(anchors[i]) != d
            throw(DimensionMismatch("All anchor points should have the same dimension"))
        end
    end
    out = Matrix{Float64}(undef, (d, (n-1) * nseg + 1))
    for i in 1:n-1
        from = anchors[i]
        to = anchors[i+1]
        dvec = (to - from) ./ nseg
        for j in 0:(nseg-1)
            out[:, (i-1)*nseg + j + 1] = from + dvec * j
        end
    end
    out[:, end] = anchors[end]
    return out
end

"""
    momentumpath

The anchorpoints are given in units of the reciprocal lattice vectors.
"""
function momentumpath(
    unitcell::UnitCell,
    anchorpoints::AbstractVector{<:AbstractVector{<:Number}},
)
    real_anchorpoints = [unitcell.reciprocallatticevectors * ap for ap in anchorpoints]
    return linpath(real_anchorpoints...)
end


export Star,
       center,
       basis,
       predicate,
       intersection!,
       dim

"""
    Star{N, VN<:AbstractVector{N}, MN<:AbstractMatrix{N}, PT<:AbstractPolyhedron{N}} <: AbstractPolyhedron{N}

Generalized star set with a polyhedral predicate, i.e.

```math
X = \\{x ∈ \\mathbb{R}^n : x = x₀ + \\sum_{i=1}^m α_i v_i,~~\\textrm{s.t. } P(α) = ⊤ \\},
```
where ``x₀ ∈ \\mathbb{R}^n`` is the center, the ``m`` vectors ``v₁, …, vₘ`` form
the basis of the star set, and the combination factors
``α = (α₁, …, αₘ) ∈ \\mathbb{R}^m`` are the predicates' decision variables,
i.e. ``P : α ∈ \\mathbb{R}^m → \\{⊤, ⊥\\}`` where the polyhedral predicate satisfies
``P(α) = ⊤`` if and only if ``Aα ≤ b`` for some fixed ``A ∈ \\mathbb{R}^{p × m}`` and
``b ∈ \\mathbb{R}^p``.

### Fields

- `c` -- vector that represents the center
- `V` -- matrix where each column corresponds to a basis vector
- `P` -- polyhedral set that represents the predicate

### Notes

The predicate function is implemented as a conjunction of linear constraints, i.e.
a subtype of `AbstractPolyhedron`. By a slight abuse of notation, the *predicate*
is also used to denote the subset of ``\\mathbb{R}^n`` such that ``P(α) = ⊤`` holds.

The ``m`` basis vectors (each one ``n``-dimensional) are stored as the columns
of an ``n × m`` matrix.

We remark that a `Star` is mathematically equivalent to the lazy affine map of
the polyhedral set `P`, with the transformation matrix and translation vector
being `V` and `c` respectively.

### Examples

This example is drawn from Example 1 in [2]. Consider the two-dimensional plane
``\\mathbb{R}^2``. Let

```jldoctest star_constructor
julia> V = [[1.0, 0.0], [0.0, 1.0]];
```
be the basis vectors and take

```jldoctest star_constructor
julia> c = [3.0, 3.0];
```
as the center of the star set. Let the predicate be the infinity-norm ball of radius 1,

```jldoctest star_constructor
julia> P = BallInf(zeros(2), 1.0);
```
Finally, the star set ``X = ⟨c, V, P⟩`` defines the set:

```jldoctest star_constructor
julia> S = Star(c, V, P)
Star{Float64, Vector{Float64}, Matrix{Float64}, BallInf{Float64, Vector{Float64}}}([3.0, 3.0], [1.0 0.0; 0.0 1.0], BallInf{Float64, Vector{Float64}}([0.0, 0.0], 1.0))
```

We can use getter functions for each component field:

```jldoctest star_constructor
julia> center(S)
2-element Vector{Float64}:
 3.0
 3.0

julia> basis(S)
2×2 Matrix{Float64}:
 1.0  0.0
 0.0  1.0

julia> predicate(S)
BallInf{Float64, Vector{Float64}}([0.0, 0.0], 1.0)
```
In this case, we know calculating by hand that the generalized star ``S`` is
defined by the rectangular set

```math
    T = \\{(x, y) ∈ \\mathbb{R}^2 : (2 ≤ x ≤ 4) ∧ (2 ≤ y ≤ 4) \\}
```
It holds that ``T`` and ``S`` are equivalent.

### References

Star sets as defined here were introduced in [1]; see also [2] for a preliminary
definition. For applications in reachability analysis of neural networks, see
[3].

- [1] Duggirala, P. S., and Mahesh V. *Parsimonious, simulation based verification of linear systems.*
      International Conference on Computer Aided Verification. Springer, Cham, 2016.

- [2] Bak S, Duggirala PS. *Simulation-equivalent reachability of large linear systems with inputs.*
      In International Conference on Computer Aided Verification
      2017 Jul 24 (pp. 401-420). Springer, Cham.

- [3] Tran, H. D., Lopez, D. M., Musau, P., Yang, X., Nguyen, L. V., Xiang, W., & Johnson, T. T. (2019, October).
      *Star-based reachability analysis of deep neural networks.*
      In International Symposium on Formal Methods (pp. 670-686). Springer, Cham.
"""
struct Star{N, VN<:AbstractVector{N}, MN<:AbstractMatrix{N}, PT<:AbstractPolyhedron{N}} <: AbstractPolyhedron{N}
    c::VN # center
    V::MN # basis
    P::PT # predicate

    # default constructor with size checks
    function Star(c::VN, V::MN, P::PT) where {N, VN<:AbstractVector{N}, MN<:AbstractMatrix{N}, PT<:AbstractPolyhedron{N}}
        @assert length(c) == size(V, 1) "the center of the basis vectors should be compatible, " *
                                        "but they are of length $(length(c)) and $(size(V, 1)) respectively"

        @assert dim(P) == size(V, 2) "the number of basis vectors should be compatible " *
                                     "with the predicates' dimension, but they are $(size(V, 2)) and $(dim(P)) respectively"

        return new{N, VN, MN, PT}(c, V, P)
    end
end

# constructor from center and list of generators
function Star(c::VN, vlist::AbstractVector{VN}, P::PT) where {N, VN<:AbstractVector{N}, PT<:AbstractPolyhedron{N}}
    V = to_matrix(vlist, length(c))
    return Star(c, V, P)
end

# type used for dispatch
const STAR{N, VN<:AbstractVector{N},
              MN<:AbstractMatrix{N},
              PT<:AbstractPolyhedron{N}} = AffineMap{N, PT, N, MN, VN}

isoperationtype(::Type{<:Star}) = false
isconvextype(::Type{<:Star}) = true

# ============================
# Star set getter functions
# ============================

"""
    center(X::Star)

Return the center of a star.

### Input

- `X` -- star

### Output

The center of the star.
"""
center(X::Star) = X.c

"""
    basis(X::Star)

Return the basis vectors of a star.

### Input

- `X` -- star

### Output

A matrix where each column is a basis vector of the star.
"""
basis(X::Star) = X.V

"""
    predicate(X::Star)

Return the predicate of a star.

### Input

- `X` -- star

### Output

A polyhedral set representing the predicate of the star.
"""
predicate(X::Star) = X.P

# =============================
# LazySets interface functions
# =============================

"""
    dim(X::Star)

Return the dimension of a star.

### Input

- `X` -- star

### Output

The dimension of a star.
"""
function dim(X::Star)
    return length(X.c)
end

"""
    σ(d::AbstractVector, X::Star)

Return the support vector of a star.

### Input

- `d` -- direction
- `X` -- star

### Output

The support vector in the given direction.
"""
function σ(d::AbstractVector, X::Star)
    A = basis(X)
    return A * σ(_At_mul_B(A, d), predicate(X)) + center(X)
end

"""
    ρ(d::AbstractVector, X::Star)

Return the support function of a star.

### Input

- `d` -- direction
- `X` -- star

### Output

The support function in the given direction.
"""
function ρ(d::AbstractVector, X::Star)
    return ρ(_At_mul_B(basis(X), d), predicate(X)) + dot(d, center(X))
end

"""
    an_element(X::Star)

Return some element of an affine map.

### Input

- `X` -- star

### Output

An element of the star. It relies on the `an_element` function of the
wrapped set.
"""
function an_element(X::Star)
    return basis(X) * an_element(predicate(X)) + center(X)
end

"""
    isempty(X::Star)

Return whether an affine map is empty or not.

### Input

- `am` -- affine map

### Output

`true` iff the predicate is empty.
"""
function isempty(X::Star)
    return isempty(predicate(X))
end

"""
    isbounded(am::Star; cond_tol::Number=DEFAULT_COND_TOL)

Determine whether a star is bounded.

### Input

- `X`        -- star
- `cond_tol` -- (optional) tolerance of matrix condition (used to check whether
                the basis matrix is invertible)

### Output

`true` iff the star is bounded.

### Algorithm

See [`isbounded(::AbstractAffineMap)`](@ref).
"""
function isbounded(X::Star; cond_tol::Number=DEFAULT_COND_TOL)
    am = convert(STAR, X)
    return isbounded(am, cond_tol=cond_tol)
end

"""
    ∈(v::AbstractVector, X::Star)

Check whether a given point is contained in a star set.

### Input

- `v`  -- point/vector
- `X` -- star

### Output

`true` iff ``v ∈ X``.

### Algorithm

See [`∈(::AbstractVector, ::AbstractAffineMap)`](@ref).
"""
function ∈(x::AbstractVector, X::Star)
    return basis(X) \ (x - center(X)) ∈ predicate(X)
end

"""
    vertices_list(X::Star; apply_convex_hull::Bool=true)

Return the list of vertices of a star set.

### Input

- `X`                 -- star
- `apply_convex_hull` -- (optional, default: `true`) if `true`, apply the convex
                         hull operation to the list of vertices of the star

### Output

A list of vertices.

### Algorithm

See [`vertices_list(::AbstractAffineMap)`](@ref).
"""
function vertices_list(X::Star; apply_convex_hull::Bool=true)
    am = convert(STAR, X)
    return vertices_list(am, apply_convex_hull=apply_convex_hull)
end

"""
    constraints_list(X::Star)

Return the list of constraints of a star.

### Input

- `X` -- star

### Output

The list of constraints of the star.

### Algorithm

See [`constraints_list(::AbstractAffineMap)`](@ref).
"""
function constraints_list(X::Star)
    am = convert(STAR, X)
    return constraints_list(am)
end

"""
    linear_map(M::AbstractMatrix, X::Star)

Return the linear map of a star.

### Input

- `M` -- matrix
- `X` -- star

### Output

The star obtained by applying `M` to `X`.
"""
function linear_map(M::AbstractMatrix, X::Star)
    c′ = M * X.c
    V′ = M * X.V
    P′ = X.P
    return Star(c′, V′, P′)
end

"""
    affine_map(M::AbstractMatrix, X::Star, v::AbstractVector)

Return the affine map of a star.

### Input

- `M` -- matrix
- `X` -- star
- `v` -- vector

### Output

The star obtained by the affine map with matrix `M` and displacement `v` to `X`.
"""
function affine_map(M::AbstractMatrix, X::Star, v::AbstractVector)
    c′ = M * X.c + v
    V′ = M * X.V
    P′ = X.P
    return Star(c′, V′, P′)
end

"""
    rand(::Type{Star}; [N]::Type{<:Real}=Float64, [dim]::Int=2,
         [rng]::AbstractRNG=GLOBAL_RNG, [seed]::Union{Int, Nothing}=nothing)

Create a random star.

### Input

- `Star`       -- type for dispatch
- `N`              -- (optional, default: `Float64`) numeric type
- `dim`            -- (optional, default: 2) dimension
- `rng`            -- (optional, default: `GLOBAL_RNG`) random number generator
- `seed`           -- (optional, default: `nothing`) seed for reseeding

### Output

A random star.

### Algorithm

This functions calls the function to generate a random zonotope,
because every zonotope is a star.
"""
function rand(::Type{Star};
              N::Type{<:Real}=Float64,
              dim::Int=2,
              rng::AbstractRNG=GLOBAL_RNG,
              seed::Union{Int, Nothing}=nothing)
    rng = reseed(rng, seed)
    Z = rand(Zonotope, N=N, dim=dim, rng=rng, seed=seed)
    X = convert(Star, Z)
    return X
end