using RobustAndOptimalControl
using ControlSystems
using Test
using LinearAlgebra
using Random

execute_tests = [true, true, true, true, true, true, true, true, true]

"""
Tests for the public and private methods of the hInfSynthesis function. This
function utilizes the preexisting ControlSystems toolbox, and performs a
H-infinity synthesis using the dual Riccati equation approach. As such,
the synthesis is done in a set of steps.

(1) Re-writing the specifications on an extended state-space form.
(2) Verifying that the resulting extended state-space object satisfies a set of
    assumptions required for proceeding with the synthesis.
(3) A coordinate transform to enable the synthesis.
(4) Synthesis using the γ-iterations, checking if a solution to the H-infinity
    problem exists in each iteration and applying a bisection method.
(5) Re-transforming the system to the original coordinates
(6) Verification that the computed solution is correct.

In addition to these six ponts, the code also enables

(7) A bilinear discretization with an inverse operation to move from continuous
to discrete time, thereby enabling approximate discrete-time synthesis.
(8) Plotting functionality to visualize the H-infinity synthesis.
(9) Three examples which can be used to demonstrate the tool.

Many of the tests are quite intuitive, and all points (1)-(9) are tested
extensively with detailed comments for each test-set, and below is a list
of all the functions tested in this unit test

  hinfsynthesize
  hinfassumptions
  hinfsignals
  hinfpartition
  bilineard2c
  bilinearc2d
  _detectable
  _stabilizable
  _synthesizecontroller
  _assertrealandpsd
  _checkfeasibility
  _solvehamiltonianare
  _solvematrixequations
  _γIterations
  _transformp2pbar
  _scalematrix
  _input2ss
  _coordinateTtansformqr
  _coordinateTtansformsvd

"""

if execute_tests[1]
    @testset "(1) Specifications" begin
        """
        These tests make sure that the specifications are written on a suitable form
        for any combination of none/empty, static gain, SISO, and MIMO
        specificaitons in the weighting functions. Essentially, the tests are made
        to verify that hInf_partition(G, WS, WU, WT) works as expected.
        """

        @testset "Conversion of user input" begin

            @testset "Empty and nothing" begin

                # Check that conversion of nothing is OK
                A, B, C, D = RobustAndOptimalControl._input2ss(nothing)
                @test isa(A, Array{Float64,2})
                @test isa(B, Array{Float64,2})
                @test isa(C, Array{Float64,2})
                @test isa(D, Array{Float64,2})

                @test size(A) == (0, 0)
                @test size(B) == (0, 0)
                @test size(C) == (0, 0)
                @test size(D) == (0, 0)

                # Check that conversion of empty objects are OK
                A, B, C, D = RobustAndOptimalControl._input2ss([])
                @test isa(A, Array{Float64,2})
                @test isa(B, Array{Float64,2})
                @test isa(C, Array{Float64,2})
                @test isa(D, Array{Float64,2})

                @test size(A) == (0, 0)
                @test size(B) == (0, 0)
                @test size(C) == (0, 0)
                @test size(D) == (0, 0)
            end

            @testset "Static gains" begin
                # Fixture
                number = 2.0

                # Check that conversion of numbers are OK
                A, B, C, D = RobustAndOptimalControl._input2ss(number)
                @test isa(A, Array{Float64,2})
                @test isa(B, Array{Float64,2})
                @test isa(C, Array{Float64,2})
                @test isa(D, Array{Float64,2})

                @test D[1, 1] == number

                @test size(A) == (0, 0)
                @test size(B) == (0, 1)
                @test size(C) == (1, 0)
                @test size(D) == (1, 1)
            end

            @testset "LTI models" begin

                # Fixture
                Random.seed!(0)
                M = 3
                N = 2
                SISO_tf = tf([1, 0], [1, 0, 1])
                SISO_ss = ss(SISO_tf)
                MIMO_ss = ss(
                    rand(Float64, (M, M)),
                    rand(Float64, (M, N)),
                    rand(Float64, (N, M)),
                    rand(Float64, (N, N)),
                )
                MIMO_tf = tf(MIMO_ss)

                # check that conversion of SISO tf data is OK
                A, B, C, D = RobustAndOptimalControl._input2ss(SISO_tf)
                @test isa(A, Array{Float64,2})
                @test isa(B, Array{Float64,2})
                @test isa(C, Array{Float64,2})
                @test isa(D, Array{Float64,2})

                @test ss(A, B, C, D) == ss(SISO_tf)

                # check that conversion of SISO tf data is OK
                A, B, C, D = RobustAndOptimalControl._input2ss(SISO_ss)
                @test isa(A, Array{Float64,2})
                @test isa(B, Array{Float64,2})
                @test isa(C, Array{Float64,2})
                @test isa(D, Array{Float64,2})

                @test ss(A, B, C, D) == SISO_ss

                # check that conversion of MIMO tf data is OK
                A, B, C, D = RobustAndOptimalControl._input2ss(MIMO_tf)
                @test isa(A, Array{Float64,2})
                @test isa(B, Array{Float64,2})
                @test isa(C, Array{Float64,2})
                @test isa(D, Array{Float64,2})

                @test ss(A, B, C, D) == ss(MIMO_tf)

                # check that conversion of MIMO tf data is OK
                A, B, C, D = RobustAndOptimalControl._input2ss(MIMO_ss)
                @test isa(A, Array{Float64,2})
                @test isa(B, Array{Float64,2})
                @test isa(C, Array{Float64,2})
                @test isa(D, Array{Float64,2})

                @test ss(A, B, C, D) == MIMO_ss
            end
        end

        #
        # @testset "Dimensionality checks"
        #
        #   # Check sensitivity function weight - must have the same number of inputs
        #   # as outputs, can have an arbitrarily large statespace, and must have the
        #   # same number of inputs as there are outputs in the process which is to be
        #   # controlled.
        #
        #   # Fixture
        #   M = 5; N = 3; L=4;
        #   G = ss(rand(Float64, (M,M)), rand(Float64, (M,L)),
        #          rand(Float64, (N,M)), rand(Float64, (N,L)))
        #
        #   # Create a sensitivity weighting function with ii inputs and ii outputs
        #   for ii = 2:6
        #     for jj = 2:6
        #       for kk = 2:6
        #         # Create some randome weighting functions
        #         WS = ss(rand(Float64, (M,M)), rand(Float64, (M,ii)),
        #                 rand(Float64, (ii,M)), rand(Float64, (ii,ii)))
        #         WU = ss(rand(Float64, (M,M)), rand(Float64, (M,jj)),
        #                 rand(Float64, (jj,M)), rand(Float64, (jj,jj)))
        #         WT = ss(rand(Float64, (M,M)), rand(Float64, (M,kk)),
        #                 rand(Float64, (kk,M)), rand(Float64, (kk,kk)))
        #
        #         println([ii,jj,kk])
        #         # Chech that the specifications can be re-written is possible
        #         if ii == N && jj == L && kk == N
        #           @test isa(hInf_partition(G, WS, WU, WT), ControlSystems.ExtendedStateSpace)
        #         else
        #           @test_throws ErrorException hInf_partition(G, WS, WU, WT)
        #         end
        #       end
        #     end
        #   end
        # end
    end
end

if execute_tests[2]
    @testset "(2) Assumptions" begin
        """
        Tests the methods used to check that the assumptions are satisfied,
        incorporating two separate algorithms to determine detectability and
        observability, and a rectangular pseudoinverse to check the rank conditions.
        """

        @testset "Stabilizability check" begin
            """
            Test the check for stabilizability using the Hautus Lemma
            """
            ### Fixture
            Random.seed!(0)
            N = 10
            M = 5
            R = rand(Float64, (N, N))
            Q = eigvecs(R + R')
            L = rand(Float64, N) .- 0.5
            A = Q * Diagonal(L) * Q'
            E = eigvals(A)

            # Here we should return false if for any positive eigenvalue λ of A, we get
            # that rank(A-λI, B) < N. We define B as conlums which are linearly
            # dependent with columns in A-λI, and check that the system returns false
            # in every case where λ > 0 and never when λ≦0.
            for ii = 1:N
                for jj = 1:(N-M)
                    Ahat = A - E[ii] * Matrix{Float64}(I, N, N)
                    B = Ahat[:, jj:(jj+M)]
                    if E[ii] >= 0
                        @test !RobustAndOptimalControl._stabilizable(A, B)
                    else
                        @test RobustAndOptimalControl._stabilizable(A, B)
                    end
                end
            end

            # Check common input types and ensure that a method error is thrown when
            # not using two abstract matrices, but some other common type
            N = 10
            M = 5
            A = rand(Float64, (N, N))
            B = rand(Float64, (N, M))
            @test_throws MethodError RobustAndOptimalControl._stabilizable(A, nothing)
            @test_throws MethodError RobustAndOptimalControl._stabilizable(nothing, B)
            @test_throws MethodError RobustAndOptimalControl._stabilizable(A, [])
            @test_throws MethodError RobustAndOptimalControl._stabilizable([], B)
            @test_throws MethodError RobustAndOptimalControl._stabilizable(A, ss(1))
            @test_throws MethodError RobustAndOptimalControl._stabilizable(ss(1), B)
            @test_throws MethodError RobustAndOptimalControl._stabilizable(A, tf(1))
            @test_throws MethodError RobustAndOptimalControl._stabilizable(tf(1), B)
        end

        @testset "Detectability check" begin
            """
            Test the check for detectability using the Hautus Lemma
            """
            ### Fixture
            Random.seed!(0)
            N = 10
            M = 5
            R = rand(Float64, (N, N))
            Q = eigvecs(R + R')
            L = rand(Float64, N) .- 0.5
            A = Q * Diagonal(L) * Q'
            E = eigvals(A)

            # Here we should return false if for any positive eigenvalue λ of A, we get
            # that rank(A-λI; C) < N. We define C as rows which are linearly
            # dependent with rows in A-λI, and check that the system returns false
            # in every case where λ > 0 and never when λ≦0.
            for ii = 1:N
                for jj = 1:(N-M)
                    Ahat = A - E[ii] * Matrix{Float64}(I, N, N)
                    C = Ahat[jj:(jj+M), :]
                    if E[ii] >= 0
                        @test !RobustAndOptimalControl._detectable(A, C)
                    else
                        @test RobustAndOptimalControl._detectable(A, C)
                    end
                end
            end

            # Check common input types and ensure that a method error is thrown when
            # not using two abstract matrices, but some other common type
            N = 10
            M = 5
            A = rand(Float64, (M, M))
            C = rand(Float64, (N, M))
            @test_throws MethodError RobustAndOptimalControl._detectable(A, nothing)
            @test_throws MethodError RobustAndOptimalControl._detectable(nothing, C)
            @test_throws MethodError RobustAndOptimalControl._detectable(A, [])
            @test_throws MethodError RobustAndOptimalControl._detectable([], C)
            @test_throws MethodError RobustAndOptimalControl._detectable(A, ss(1))
            @test_throws MethodError RobustAndOptimalControl._detectable(ss(1), C)
            @test_throws MethodError RobustAndOptimalControl._detectable(A, tf(1))
            @test_throws MethodError RobustAndOptimalControl._detectable(tf(1), C)
        end

        # TODO: write tests using the above submethods directly in hInf_assumptions
    end
end

if execute_tests[3]
    @testset "(3) Coordinate transform" begin
        """
        Computes the coordinate transfomration to write the system on the form of
        assumption A4 in [1], i.e. take an extended statespace model, P, satisfying
        assumption A2, and transform is to a system P_bar, where the matrices
        D12_bar = [0,I] and D21_bar = [0;I].
        """

        @testset "Find transformation for full rank matrix with QR" begin
            """
            Test computation of the transformation in square and rectangular cases.
            """
            # Fixture
            tolerance = 1e-10

            for N = 1:5
                for M = 1:5
                    A = rand(Float64, (M, N))
                    I_mat = Matrix{Float64}(I, min(M, N), min(M, N))

                    if M == N
                        # Square case
                        A_bar_true = I_mat
                    elseif M > N
                        # Rectangular case (more rows than columns)
                        Z_mat = zeros(Float64, (max(M, N) - min(M, N), min(M, N)))
                        A_bar_true = [Z_mat; I_mat]
                    else
                        # Rectangular case (more columns than rows)
                        Z_mat = zeros(Float64, (min(M, N), max(M, N) - min(M, N)))
                        A_bar_true = [Z_mat I_mat]
                    end
                    Tl, Tr = RobustAndOptimalControl._coordinatetransformqr(A)
                    @test opnorm(Tl * A * Tr - A_bar_true) < tolerance
                end
            end
        end

        @testset "Find transformation for full rank matrix with SVD" begin
            """
            Test computation of the transformation in square and rectangular cases.
            """
            # Fixture
            tolerance = 1e-10

            for N = 1:5
                for M = 1:5
                    A = rand(Float64, (M, N))
                    I_mat = Matrix{Float64}(I, min(M, N), min(M, N))

                    if M == N
                        # Square case
                        A_bar_true = I_mat
                    elseif M > N
                        # Rectangular case (more rows than columns)
                        Z_mat = zeros(Float64, (max(M, N) - min(M, N), min(M, N)))
                        A_bar_true = [Z_mat; I_mat]
                    else
                        # Rectangular case (more columns than rows)
                        Z_mat = zeros(Float64, (min(M, N), max(M, N) - min(M, N)))
                        A_bar_true = [Z_mat I_mat]
                    end
                    Tl, Tr = RobustAndOptimalControl._coordinatetransformsvd(A)
                    @test opnorm(Tl * A * Tr - A_bar_true) < tolerance
                end
            end
        end

        @testset "Find transformation for arbitrary matrix inputs" begin
            """
            Test computation of the transformation in square and rectangular cases.
            """
            # Fixture
            test_matrices_full = [
                rand(Float64, (1, 1)),
                rand(Float64, (5, 5)),
                rand(Float64, (2, 5)),
                rand(Float64, (5, 2)),
            ]

            # Compute rank deficient matrices
            test_matrices_rank_deficient = test_matrices_full
            for ii = 1:length(test_matrices_full)
                U, S, V = svd(test_matrices_full[ii])
                S[1] = 0
                test_matrices_rank_deficient[ii] = U * Diagonal(S) * V'
            end

            # Make sure that an exception is raised if assumption A2 is violated,
            # any of the matrices which are to be decomposed are rank deficient
            for A in test_matrices_rank_deficient
                println(A)
                @test_throws ErrorException RobustAndOptimalControl._scalematrix(A)
            end

            # Various bad inputs
            @test_throws MethodError RobustAndOptimalControl._scalematrix(tf(1))
            @test_throws MethodError RobustAndOptimalControl._scalematrix(ss(1))
            @test_throws MethodError RobustAndOptimalControl._scalematrix([])
            @test_throws MethodError RobustAndOptimalControl._scalematrix(nothing)
            @test_throws MethodError RobustAndOptimalControl._scalematrix(1)

            # Check that errors are thrown if the matric has zise zero
            @test_throws ErrorException RobustAndOptimalControl._scalematrix(zeros(Float64, (0, 1)))

            # Various bad methods
            A = test_matrices_full[1]
            @test_throws ErrorException RobustAndOptimalControl._scalematrix(
                A;
                method = "bad method",
            )
            @test_throws ErrorException RobustAndOptimalControl._scalematrix(A; method = 1)
        end

        @testset "Test application of the transformation to an ESS object" begin
            """
            The wa
            """
        end
    end
end

if execute_tests[4]
    @testset "(4) Gamma iterations" begin
        """
        Tests the core methods of the γ-iteration bisection method, including
        the ARE hamiltonian solver, the eigenvalue solver, the feasibility checks.
        """

        @testset "Solution feasibility check" begin
            """
            Check that a solution to the dual riccati equations is correctly reported
            as being feasible if satisfying the conditions of positive definiteness on
            the X and Y solutions are met, and the spectral radius condition of X*Y is
            also met
            """
            # Fixture
            tolerance = 1e-10
            iteration = 1
            Random.seed!(0)
            N = 10
            M = 5
            R = rand(Float64, N, N)
            Q = eigvecs(R + R')
            ρX = 0.2
            ρY = 0.3
            LX = rand(Float64, N)
            LX = ρX * sort(LX / maximum(LX))
            LY = rand(Float64, N)
            LY = ρY * sort(LY / maximum(LY))
            Xinf = Q * Diagonal(LX) * Q'
            Yinf = Q * Diagonal(LY) * Q'
            γ = 1

            # Test that the fesibility is true if ρ(Xinf*Yinf) < γ^2, that is the
            # check should be true for any
            #
            #    γ = sqrt(ρX*ρY) + ϵ
            #
            # with any epsilon greater than or equal to zero.
            @test !RobustAndOptimalControl._checkfeasibility(
                Xinf,
                Yinf,
                sqrt(ρX * ρY) - tolerance,
                tolerance,
                iteration;
                verbose = false,
            )
            @test RobustAndOptimalControl._checkfeasibility(
                Xinf,
                Yinf,
                sqrt(ρX * ρY) + tolerance,
                tolerance,
                iteration;
                verbose = false,
            )

            # Test that errors are thrown if the matrix Xinf and Yinf are not PSD down
            # to the numerical tolerance.
            L = LX
            L[1] += -L[1] + 2 * tolerance
            Xpos = Q * Diagonal(L) * Q' # slightly positive eigenvalue
            L[1] += -L[1]
            Xzero = Q * Diagonal(LX) * Q'               # exactly one zero eigenvalue
            L[1] += -L[1] - 2 * tolerance
            Xneg = Q * Diagonal(L) * Q'   # slightly negative eigenvalue
            @test RobustAndOptimalControl._checkfeasibility(
                Xpos,
                Yinf,
                sqrt(ρX * ρY) + tolerance,
                tolerance,
                iteration;
                verbose = false,
            )
            @test RobustAndOptimalControl._checkfeasibility(
                Xzero,
                Yinf,
                sqrt(ρX * ρY) + tolerance,
                tolerance,
                iteration;
                verbose = false,
            )
            @test !RobustAndOptimalControl._checkfeasibility(
                Xneg,
                Yinf,
                sqrt(ρX * ρY) + tolerance,
                tolerance,
                iteration;
                verbose = false,
            )

            L = LY
            L[1] += -L[1] + 2 * tolerance
            Ypos = Q * Diagonal(L) * Q' # slightly positive eigenvalue
            L[1] += -L[1]
            Yzero = Q * Diagonal(L) * Q'               # exactly one zero eigenvalue
            L[1] += -L[1] - 2 * tolerance
            Yneg = Q * Diagonal(L) * Q'   # slightly negative eigenvalue
            @test RobustAndOptimalControl._checkfeasibility(
                Xinf,
                Ypos,
                sqrt(ρX * ρY) + tolerance,
                tolerance,
                iteration;
                verbose = false,
            )
            @test RobustAndOptimalControl._checkfeasibility(
                Xinf,
                Yzero,
                sqrt(ρX * ρY) + tolerance,
                tolerance,
                iteration;
                verbose = false,
            )
            @test !RobustAndOptimalControl._checkfeasibility(
                Xinf,
                Yneg,
                sqrt(ρX * ρY) + tolerance,
                tolerance,
                iteration;
                verbose = false,
            )
        end

        # TODO: Include a check to verify that the bisection works as intended.
        # TODO: Check to verify that the Hamiltonian Shur-solver is working
        # TODO: Check to verify that the Hamiltonian Eigenvalue-solver is working
    end
end

if execute_tests[7]
    @testset "(7) Bilinear discretization" begin
        """
        This tests the bilinear method of discretizing a continuous time StateSpace
        or ExtendedStateSpace object, moving from the Laplace-domain to the Z-domain.
        However, importantly, the method also enables approximating discrete-time
        systems with continuous-time state-space domain, moving from the Z-domain back
        to the laplace-domain. Since the L∞-norm is invariant over the bilinear
        discretization and "continuization", this effectively allows us to approximate
        discrete-time systems with a continuous-time equivalent, design an H∞-optimal
        controller for the continuous-time plant, and then re-discretizing this
        controller.

        The structure of these tests is to simply discretize a set of plants with
        various dimensions, transport them back into continuous-time domain, and then
        to discrete-time again. By then comparing the resulting two pairs of discrete-
        time and continuous-time systems in the frequency domain, we can test if the
        bilinear discretization operates as expected.
        """

        @testset "SS data" begin
            # Fixture for the SS-type data
            Random.seed!(0)
            N = [1, 5, 10]
            M = [1, 5, 10]
            H = [0.1, 0.01, 0.001]
            tolerance = 1e-7

            # Tests for the SS-type data
            for (ii, m) in enumerate(N)
                for (jj, n) in enumerate(M)
                    for (kk, h) in enumerate(H)
                        freq = [10^i for i in range(-6, stop = log10(pi / h), length = 101)]

                        AcTrue, BcTrue, CcTrue, DcTrue =
                            rand(m, m), rand(m, n), rand(n, m), rand(n, n)

                        ev = eigvals(AcTrue)
                        if !isempty(ev[real(ev).<0])
                            AcTrue =
                                AcTrue -
                                one(AcTrue) * 2 * maximum(abs.(real(ev[real(ev).<0])))
                        end

                        Gtrue = ss(AcTrue, BcTrue, CcTrue, DcTrue)
                        valA, fA = sigma(Gtrue, freq)
                        sysB = bilinearc2d(Gtrue, h)
                        valB, fB = sigma(sysB, freq)
                        sysC = bilineard2c(bilinearc2d(Gtrue, h))
                        valC, fC = sigma(sysC, freq)
                        sysD = bilinearc2d(bilineard2c(bilinearc2d(Gtrue, h)), h)
                        valD, fD = sigma(sysD, freq)

                        @test abs(maximum(svd(sysB.B).S) - maximum(svd(sysB.C).S)) <
                              tolerance
                        @test abs(maximum(svd(sysC.B).S) - maximum(svd(sysC.C).S)) <
                              tolerance
                        @test abs(maximum(svd(sysD.B).S) - maximum(svd(sysD.C).S)) <
                              tolerance

                        # Test that the C->D->C and D->C->D results in the same
                        @test norm(valA - valC, Inf) < tolerance
                        @test norm(valB - valD, Inf) < tolerance
                    end
                end
            end
        end

        @testset "Extended SS data" begin
            # Fixture for the extended SS-type data
            Random.seed!(0)
            N = [1, 4]
            P1 = [1, 4]
            P2 = [1, 4]
            M1 = [1, 4]
            M2 = [1, 4]
            H = [0.1, 0.01, 0.001]
            tolerance = 1e-5

            for (in1, n) in enumerate(N)
                for (im1, m1) in enumerate(M1)
                    for (im2, m2) in enumerate(M2)
                        for (ip1, p1) in enumerate(P1)
                            for (ip2, p2) in enumerate(P2)
                                for (ih, h) in enumerate(H)
                                    freq = [
                                        10^i for i in
                                        range(-4, stop = log10(pi / h), length = 101)
                                    ]

                                    A, B1, B2 = 0.01 * rand(n, n), rand(n, m1), rand(n, m2)
                                    C1, D11, D12 = rand(p1, n), rand(p1, m1), rand(p1, m2)
                                    C2, D21, D22 = rand(p2, n), rand(p2, m1), rand(p2, m2)

                                    ev = eigvals(A)
                                    if !isempty(ev[real(ev).>0])
                                        A =
                                            A -
                                            one(A) *
                                            2 *
                                            maximum(abs.(real(ev[real(ev).>0])))
                                    end

                                    EsysA = ss(A, B1, B2, C1, C2, D11, D12, D21, D22)

                                    valA, fA = sigma(
                                        ss(
                                            EsysA.A,
                                            [EsysA.B1 EsysA.B2],
                                            [EsysA.C1; EsysA.C2],
                                            [EsysA.D11 EsysA.D12; EsysA.D21 EsysA.D22],
                                        ),
                                        freq,
                                    )

                                    # To discrete time
                                    EsysB = bilinearc2d(EsysA, h)
                                    valB, fB = sigma(
                                        ss(
                                            EsysB.A,
                                            [EsysB.B1 EsysB.B2],
                                            [EsysB.C1; EsysB.C2],
                                            [EsysB.D11 EsysB.D12; EsysB.D21 EsysB.D22],
                                            h,
                                        ),
                                        freq,
                                    )

                                    # To continuous time
                                    EsysC = bilineard2c(bilinearc2d(EsysA, h))
                                    valC, fC = sigma(
                                        ss(
                                            EsysC.A,
                                            [EsysC.B1 EsysC.B2],
                                            [EsysC.C1; EsysC.C2],
                                            [EsysC.D11 EsysC.D12; EsysC.D21 EsysC.D22],
                                        ),
                                        freq,
                                    )

                                    # To discrete time
                                    EsysD =
                                        bilinearc2d(bilineard2c(bilinearc2d(EsysA, h)), h)
                                    valD, fD = sigma(
                                        ss(
                                            EsysD.A,
                                            [EsysD.B1 EsysD.B2],
                                            [EsysD.C1; EsysD.C2],
                                            [EsysD.D11 EsysD.D12; EsysD.D21 EsysD.D22],
                                            h,
                                        ),
                                        freq,
                                    )

                                    # Test that the C->D->C and D->C->D results in the same
                                    @test norm(valA - valC, Inf) < tolerance
                                    @test norm(valB - valD, Inf) < tolerance
                                end
                            end
                        end
                    end
                end
            end
        end
    end
end

if execute_tests[9]
    @testset "(9) Synthesis Examples" begin

        @testset "DC motor example" begin
            # Fixture
            tolerance = 1e-4

            # Make sure that the code runs
            @test_nowarn include("../examples/hinf_example_DC.jl")
            Ω = [10^i for i in range(-3, stop = 3, length = 201)]

            # Check that the optimal gain is correct
            @test abs(γ - 4.463211059) < tolerance

            # Check that the closed loop satisfies ||F_l(P(jω), C(jω)||_∞ < γ,  ∀ω ∈ Ω
            valPcl = sigma(Pcl, Ω)[1]
            @test all(valPcl .< (γ + tolerance))

            # Check that ||S(jω)/WS(jω)||_∞ < γ,  ∀ω ∈ Ω
            if isa(WS, LTISystem) || isa(WS, Number)
                valSWS = sigma(S * WS, Ω)[1]
                @test all(valSWS .< (γ + tolerance))
            end

            # Check that ||C(jω)S(jω)/W_U(jω)||_∞ < γ,  ∀ω ∈ Ω
            if isa(WU, LTISystem) || isa(WU, Number)
                valKSWU = sigma(CS * WU, Ω)[1]
                @test all(valKSWU .< (γ + tolerance))
            end

            # Check that ||T(jω)/W_T(jω)||_∞ < γ,  ∀ω ∈ Ω
            if isa(WT, LTISystem) || isa(WT, Number)
                valTWT = sigma(T * WT, Ω)[1]
                @test all(valTWT .< (γ + tolerance))
            end

            # C_test = let
            #     _A = [-1.99970880438311e-7 -3.28465977959578e-12 -1.46200072004469e-10; 5213.86459190004 -599.125731164477 -26175.7449027424; 0.000132853900209503 0.249985014202232 -0.000667020485484034]
            #     _B = [3.9999999999998; -5.48401518999795e-14; -9.03638143773611e-7]
            #     _C = [651.733073987504 -74.8757163955595 -3271.9681128428]
            #     _D = [0.0]
            #     ss(_A, _B, _C, _D)
            # end
            # @test C ≈ C_test

            # Pcl_test = let
            #     _A = [-2.0e-7 0.0 -22.4 0.0 0.0 0.0; 0.0 -0.12 0.0 5213.86459190003 -599.005731164476 -26175.7449027424; 0.0 0.25 0.0 0.0 0.0 0.0; 0.0 0.0 -22.3999999999989 -1.99970880438311e-7 -3.28465977959578e-12 -1.46200072004469e-10; 0.0 0.0 3.07104850639885e-13 5213.86459190004 -599.125731164477 -26175.7449027424; 0.0 0.0 5.06037360513222e-6 0.000132853900209503 0.249985014202232 -0.000667020485484034]
            #     _B = [4.0; 0.0; 0.0; 3.9999999999998; -5.48401518999795e-14; -9.03638143773611e-7]
            #     _C = [4.99999996666667 0.0 -3.73333333333333 0.0 0.0 0.0; 0.0 0.0 0.0 651.733073987504 -74.8757163955595 -3271.9681128428]
            #     _D = [0.666666666666667; 0.0]
            #     ss(_A, _B, _C, _D)
            # end
            # @test Pcl ≈ Pcl_test

            @test_nowarn include("../examples/hinf_example_DC_discrete.jl")
        end

        @testset "MIT open courseware example" begin

            # Fixture
            tolerance = 1e-4

            # Make sure that the code runs
            @test_nowarn include("../examples/hinf_example_MIT.jl")
            Ω = [10^i for i in range(-7, stop = 7, length = 201)]

            # Check that the optimal gain is correct
            @test abs(γ - 0.923430124918) < tolerance

            # Check that the closed loop satisfies ||F_l(P(jω), C(jω)||_∞ < γ,  ∀ω ∈ Ω
            valPcl = sigma(Pcl, Ω)[1]
            @test all(valPcl .< (γ + tolerance))

            # Check that ||S(jω)/WS(jω)||_∞ < γ,  ∀ω ∈ Ω
            if isa(WS, LTISystem) || isa(WS, Number)
                valSWS = sigma(S * WS, Ω)[1]
                @test all(valSWS .< (γ + tolerance))
            end

            # Check that ||C(jω)S(jω)/W_U(jω)||_∞ < γ,  ∀ω ∈ Ω
            if isa(WU, LTISystem) || isa(WU, Number)
                valKSWU = sigma(CS * WU, Ω)[1]
                @test all(valKSWU .< (γ + tolerance))
            end

            # Check that ||T(jω)/W_T(jω)||_∞ < γ,  ∀ω ∈ Ω
            if isa(WT, LTISystem) || isa(WT, Number)
                valTWT = sigma(T * WT, Ω)[1]
                @test all(valTWT .< (γ + tolerance))
            end
            @test_nowarn include("../examples/hinf_example_MIT_discrete.jl")
        end

        @testset "Quad tank example" begin

            # Fixture
            tolerance = 1e-2

            # Make sure that the code runs
            @test_nowarn include("../examples/hinf_example_tank.jl")
            Ω = [10^i for i in range(-7, stop = 7, length = 201)]

            # Check that the optimal gain is correct
            @test abs(γ - 0.9534467) < tolerance

            # Check that the closed loop satisfies ||F_l(P(jω), C(jω)||_∞ < γ,  ∀ω ∈ Ω
            valPcl = sigma(Pcl, Ω)[1]
            @test all(valPcl[:, 1] .< (γ + tolerance))

            # Check that ||S(jω)/WS(jω)||_∞ < γ,  ∀ω ∈ Ω
            if isa(WS, LTISystem) || isa(WS, Number)
                valSWS = sigma(S * WS, Ω)[1]
                @test all(valSWS[:, 1] .< (γ + tolerance))
            end

            # Check that ||C(jω)S(jω)/W_U(jω)||_∞ < γ,  ∀ω ∈ Ω
            if isa(WU, LTISystem) || isa(WU, Number)
                valKSWU = sigma(CS * WU, Ω)[1]
                @test all(valKSWU[:, 1] .< (γ + tolerance))
            end

            # Check that ||T(jω)/W_T(jω)||_∞ < γ,  ∀ω ∈ Ω
            if isa(WT, LTISystem) || isa(WT, Number)
                valTWT = sigma(T * WT, Ω)[1]
                @test all(valTWT[:, 1] .< (γ + tolerance))
            end
            @test_nowarn include("../examples/hinf_example_tank_discrete.jl")
        end
    end
end


@doc raw"""
    energy_balance(canonical::Canonical,
                        initial_conditions::Vector{InitialCondition},
                        efficiency_data::Tuple{Vector{String}, Vector{InOut}},
                        cons_name::Symbol,
                        var_names::Tuple{Symbol, Symbol, Symbol})

Constructs multi-timestep constraint from initial condition, efficiency data, and variable tuple

# Constraints

If t = 1:

``` varenergy[name, 1] == initial_conditions[ix].value + varin[name, 1]*eff_in*fraction_of_hour - varout[name, 1]*fraction_of_hour/eff_out ```

If t > 1:

``` varenergy[name, t] == varenergy[name, t-1] + varin[name, t]*eff_in*fraction_of_hour - varout[name, t]*fraction_of_hour/eff_out ```

# LaTeX

`` x^{energy}_1 == x^{energy}_{init} + frhr \eta^{in} x^{in}_1 - \frac{frhr}{\eta^{out}} x^{out}_1, \text{ for } t = 1 ``

`` x^{energy}_t == x^{energy}_{t-1} + frhr \eta^{in} x^{in}_t - \frac{frhr}{\eta^{out}} x^{out}_t, \forall t \geq 2 ``

# Arguments
* canonical::Canonical : the canonical model built in PowerSimulations
* initial_conditions::Vector{InitialCondition} : for time zero 'varenergy'
* efficiency_data::Tuple{Vector{String}, Vector{InOut}} :: charging/discharging efficiencies
* cons_name::Symbol : name of the constraint
* var_names::Tuple{Symbol, Symbol, Symbol} : the names of the variables
- : var_names[1] : varin
- : var_names[2] : varout
- : var_names[3] : varenergy

"""
function energy_balance(canonical::Canonical,
                        initial_conditions::Vector{InitialCondition},
                        efficiency_data::Tuple{Vector{String}, Vector{InOut}},
                        cons_name::Symbol,
                        var_names::Tuple{Symbol, Symbol, Symbol})

    time_steps = model_time_steps(canonical)
    resolution = model_resolution(canonical)
    fraction_of_hour = Dates.value(Dates.Minute(resolution))/60
    name_index = efficiency_data[1]

    varin = get_variable(canonical, var_names[1])
    varout = get_variable(canonical, var_names[2])
    varenergy = get_variable(canonical, var_names[3])

    constraint = add_cons_container!(canonical, cons_name, name_index, time_steps)

    for (ix, name) in enumerate(name_index)
        eff_in = efficiency_data[2][ix].in
        eff_out = efficiency_data[2][ix].out

        constraint[name, 1] = JuMP.@constraint(canonical.JuMPmodel,
                                   varenergy[name, 1] == initial_conditions[ix].value + varin[name, 1]*eff_in*fraction_of_hour
                                                    - (varout[name, 1])*fraction_of_hour/eff_out)

    end

    for t in time_steps[2:end], (ix, name) in enumerate(name_index)
        eff_in = efficiency_data[2][ix].in
        eff_out = efficiency_data[2][ix].out

        constraint[name, t] = JuMP.@constraint(canonical.JuMPmodel,
                                   varenergy[name, t] == varenergy[name, t-1] + varin[name, t]*eff_in*fraction_of_hour
                                                    - (varout[name, t])*fraction_of_hour/eff_out)
    end

    return

end


#
#--------------------------------------#
export advect
#--------------------------------------#
"""
 for v,u,T in H¹₀(Ω)

 (v,(u⃗⋅∇)T) = (v,ux*∂xT + uy*∂yT)\n
            = v' *B*(ux*∂xT + uy*∂yT)\n

 implemented as

 R'R * QQ' * B * (ux*∂xT + uy*∂yT)

 (u⃗⋅∇)T = ux*∂xT + uy*∂yT

        = [ux uy] * [Dx] T
                    [Dx]

 ux,uy, ∇T are interpolated to
 a grid with higher polynomial order
 for dealiasing (over-integration)
 so we don't commit any
 "variational crimes"

"""
function advect(T  ::Array
               ,ux ::Array
               ,uy ::Array
               ,msh::Mesh)

    @unpack B = msh

    Tx,Ty = grad(T,msh)

    Cu   = @. ux * Tx + uy * Ty
    Cu .*= B

#   Cu = gatherScatter(Cu,msh)
#   Cu = mask(Cu,M)

    return Cu
end
#--------------------------------------#
function advect(T   ::Array
               ,ux  ::Array
               ,uy  ::Array
               ,mshV::Mesh
               ,mshD::Mesh  # dealias
               ,Jr,Js)

    Tx,Ty = grad(T,mshV)

    JTx = ABu(Js,Jr,Tx)
    JTy = ABu(Js,Jr,Ty)
    Jux = ABu(Js,Jr,ux)
    Juy = ABu(Js,Jr,uy)

    JCu   = @. Jux*JTx + Juy*JTy
    JCu .*= mshD.B
    Cu    = ABu(Js',Jr',JCu)

    return Cu
end
#--------------------------------------#
function advect(T   ::Array
               ,ux  ::Array
               ,uy  ::Array
               ,mshV::Mesh
               ,mshD::Mesh) # dealias

    Jr = interpMat(mshD.zr,mshV.zr)
    Js = interpMat(mshD.zs,mshV.zs)

    Cu = advect(T,ux,uy,mshV,mshD,Jr,Js)

    return Cu
end
#--------------------------------------#


@testset "TestSystems" begin
    @testset "cyclic5 validate solutions" begin
        cyclic5 = TestSystems.cyclic5()

        solutions = TestSystems.cyclic5Solutions()

        @test all(sol -> norm(evaluate(cyclic5, sol)) < 1e-8, solutions) == true
    end

    @testset "cyclic7 validate solutions" begin
        cyclic7 = TestSystems.cyclic7()

        solutions = TestSystems.cyclic7Solutions()

        @test all(sol -> norm(evaluate(cyclic7, sol)) < 1e-8, solutions) == true
    end
end


@testset "attention.jl" begin
    @testset "Attention" begin
        dim_x = 2
        dim_y = 3
        dim_embedding = 20
        num_heads = 5
        batch_size = 6
        n = 7
        m = 8

        layer = attention(
            dim_x=dim_x,
            dim_y=dim_y,
            dim_embedding=dim_embedding,
            num_heads=num_heads
        )

        xc = randn(Float32, n, dim_x, batch_size)
        yc = randn(Float32, n, dim_y, batch_size)
        xt = randn(Float32, m, dim_x, batch_size)

        # Perform encodings.
        keys = layer.encoder_x(xc)
        queries = layer.encoder_x(xt)
        values = layer.encoder_xy(cat(xc, yc, dims=2))

        dim_head = div(dim_embedding, num_heads)

        # Compute with layer.
        _, out = code(layer, xc, yc, xt)

        # Brute-force the attention computation.
        embeddings = zeros(Float32, m, dim_head, num_heads, batch_size)
        for c = 1:num_heads
            for b = 1:batch_size
                # Calculate weights.
                weights = Array{Float32}(undef, n, m)
                for i = 1:n, j = 1:m
                    key = keys[i, :, c, b]
                    query = queries[j, :, c, b]
                    # Normalise to control variance of product.
                    product = dot(key, query) / Float32(sqrt(length(key)))
                    weights[i, j] = exp(product)
                end
                for j = 1:m
                    weights[:, j] ./= sum(weights[:, j])
                end

                # Calculate embeddings.
                for i = 1:n, j = 1:m
                    embeddings[j, :, c, b] .+= values[i, :, c, b] .* weights[i, j]
                end
            end
        end
        ref = layer.transformer(layer.mixer(embeddings), queries)

        # Check that the layer lines up with the brute-force reference.
        @test out ≈ ref
    end

    @testset "BatchedMLP" begin
        layer = NeuralProcesses.batched_mlp(
            dim_in=2,
            dim_out=3,
            dim_hidden=10,
            num_layers=3
        )
        x = randn(10, 2, 4, 5)
        @test size(layer(x)) == (10, 3, 4, 5)
    end
end


using Brownian
using Base.Test

@test convert(FGN, FBM(0:0.5:10, 0.4)) == FGN(0.757858283255199,0.4)

p = FBM(0:1/2^10:1, 0.4)

rand(p)
rand(p, fbm=false)
rand([p, p])
rand(p, rtype=:chol)
rand(p, fbm=false, rtype=:chol)
rand([p, p], rtype=:chol)


function MinMaxModel(A,b,C,d,XLB,XUB,VariableDef,YR)

    m,n = size(A);
    nObj = size(C)[1];
    MyModel = Model();
    @variable(MyModel,XLB[i] <= X[i=1:n] <= XUB[i],category = VariableType(VariableDef,i))#
    @variable(MyModel,Z)

    @objective(MyModel,Min,Z)
    # @constraintref Ucons[1:m];
    # @constraintref LAMBDAcons[1:nObj];
    for i = 1:m
         @constraint(MyModel,sum(A[i,j]*X[j] for j = 1:n) >= b[i] )
    end
    for j = 1:nObj
         @constraint(MyModel,sum(C[j,i]*X[i] for i = 1:n) + d[j] - Z <= YR[j] )
    end
    for j = 1:nObj
         @constraint(MyModel,sum(C[j,i]*X[i] for i = 1:n) + d[j] >= YR[j])
    end
    return MyModel;

end

function MinMaxModel(A,b,C,d,XLB,XUB,VariableDef,YR; Cuts = [])

    m,n = size(A);
    nObj = size(C)[1];
    MyModel = Model();
    @variable(MyModel,XLB[i] <= X[i=1:n] <= XUB[i],category = VariableType(VariableDef,i))#
    @variable(MyModel,Z)

    @objective(MyModel,Min,Z)
    # @constraintref Ucons[1:m];
    # @constraintref LAMBDAcons[1:nObj];
    for i = 1:m
         @constraint(MyModel,sum(A[i,j]*X[j] for j = 1:n) >= b[i] )
    end
    for j = 1:nObj
         @constraint(MyModel,sum(C[j,i]*X[i] for i = 1:n) + d[j] - Z <= YR[j] )
    end
    for j = 1:nObj
         @constraint(MyModel,sum(C[j,i]*X[i] for i = 1:n) + d[j] >= YR[j])
    end
    if !isempty(Cuts)
        for i in Cuts
            ZeroVals = setdiff(VariableDef[1]+VariableDef[2]+1:n,i.OneVals)
            @constraint(MyModel, sum(1-X[Int(j+VariableDef[1]+VariableDef[2])] for j in i.OneVals)
            + sum(X[Int(j)] for j = ZeroVals) >= 1)
        end
    end
    return MyModel;

end


#=
JuMP helpers for constructing formulations for spectral/nuclear norm cone constraints
=#

function add_symm_specnuc(
    use_nat::Bool,
    is_nuc::Bool,
    u_aff::JuMPScalar,
    W_aff::AbstractMatrix{<:JuMPScalar},
    model::JuMP.Model,
)
    nrow = size(W_aff, 1)
    @assert issymmetric(W_aff)
    if use_nat
        len = Cones.svec_length(nrow)
        K = Hypatia.EpiNormSpectralTriCone{Float64, Float64}(1 + len, is_nuc)
        W_vec = Vector{JuMP.AffExpr}(undef, len)
        Cones.smat_to_svec!(W_vec, W_aff, rt2)
        JuMP.@constraint(model, vcat(u_aff, W_vec) in K)
    else
        if is_nuc
            # EF for symmetric nuclear norm (L1 of eigvals)
            JuMP.@variable(model, X1[1:nrow, 1:nrow], PSD)
            JuMP.@variable(model, X2[1:nrow, 1:nrow], PSD)
            eq = [W_aff[i, j] - X1[i, j] + X2[i, j] for j in 1:nrow for i in 1:j]
            JuMP.@constraint(model, eq .== 0)
            JuMP.@constraint(model, u_aff >= tr(X1) + tr(X2))
        else
            # EF for symmetric spectral norm (Linf of eigvals)
            uI = u_aff * Matrix(I, nrow, nrow)
            JuMP.@constraint(model, Symmetric(uI - W_aff) in JuMP.PSDCone())
            JuMP.@constraint(model, Symmetric(uI + W_aff) in JuMP.PSDCone())
        end
    end
    return
end

function add_nonsymm_specnuc(
    use_nat::Bool,
    is_nuc::Bool,
    u_aff::JuMPScalar,
    W_aff::AbstractMatrix{<:JuMPScalar},
    model::JuMP.Model,
)
    (nrow, ncol) = size(W_aff)
    if use_nat
        K = Hypatia.EpiNormSpectralCone{Float64, Float64}(nrow, ncol, is_nuc)
        JuMP.@constraint(model, vcat(u_aff, vec(W_aff)) in K)
    else
        if is_nuc
            # EF for nuclear norm
            X1 = JuMP.@variable(model, [1:nrow, 1:nrow], Symmetric)
            X2 = JuMP.@variable(model, [1:ncol, 1:ncol], Symmetric)
            JuMP.@constraint(model, u_aff >= (tr(X1) + tr(X2)) / 2)
        else
            # EF for spectral norm
            X1 = u_aff * Matrix(I, nrow, nrow)
            X2 = u_aff * Matrix(I, ncol, ncol)
        end
        mat = hvcat((2, 2), X1, W_aff, W_aff', X2)
        JuMP.@constraint(model, Symmetric(mat) in JuMP.PSDCone())
    end
    return
end


using ApproxFun, SingularIntegralEquations, Test
    import ApproxFun: testbandedoperator, testraggedbelowoperator, testblockbandedoperator

include("runtests.jl")

@testset "Full" begin
    ## Memory CurveTests

    d = exp(im*Interval(0.1,0.2))
    x = Fun(d)
    w = 1/(sqrt(abs(first(d)-x))*sqrt(abs(last(d)-x)))
    testbandedoperator(Hilbert(space(w)))

    ## 3 domain ideal fluid flow

    Γ = Segment(-im,1.0-im) ∪ Curve(Fun(x->exp(0.8im)*(x+x^2-1+im*(x-4x^3+x^4)/6))) ∪ Circle(2.0,0.2)
        z = Fun(Γ)

    S = PiecewiseSpace(map(d->isa(d,Circle) ? Fourier(d) : JacobiWeight(0.5,0.5,Ultraspherical(1,d)),components(Γ)))
    H = Hilbert(S)
    testblockbandedoperator(H)

    #  TODO: fix testraggedbelowoperator(H)

    Ai = [Operator(Fun(ones(component(Γ,1)),Γ)) Fun(ones(component(Γ,2)),Γ) Fun(ones(component(Γ,3)),Γ) real(H)]

    @test ApproxFun.israggedbelow(Ai)
    @test ApproxFun.israggedbelow(Ai.ops[4])
    @test ApproxFun.israggedbelow(Ai.ops[4].op)

    B=ApproxFun.SpaceOperator(ApproxFun.BasisFunctional(3),S,ApproxFun.ConstantSpace(Float64))

    Ai=[Operator(0)                 0                 0                 B;
        Fun(ones(component(Γ,1)),Γ) Fun(ones(component(Γ,2)),Γ) Fun(ones(component(Γ,3)),Γ) real(H)]

    @time testblockbandedoperator(Ai)

    k=114;
        α=exp(k/50*im)
        @time a,b,c,ui= Ai \ [0; imag(α*z)]

    u =(x,y)->α*(x+im*y)+2cauchy(ui,x+im*y)
    @test u(1.1,0.2) ≈ (-0.8290718508107162+0.511097153754im)
end

println("Example Tests")
include("ExamplesTest.jl")

include("WienerHopfTest.jl")

include("IdealFluidFlowTest.jl")


using MittagLeffler
using Test

@testset "mittleff" begin
    @test isapprox(mittleff(0.5, 0.5, 0.5),  1.5403698281390346)
    @test isapprox(mittleff(1.5, 0.5, 0.5), 1.1448466286155243)
    @test isapprox(mittleff(2.3, 0.7 + 2.0 * im), 1.201890136368392 + 0.7895394560075035 * im)
    @test isapprox(mittleff(2.3, 0.7 + 0.2 * im) , (1.268233154873853 + 0.07914994421659409im))
    @test mittleff(big".3",100.0) ==
        big"8.721285946907744692995882256235296113802695745418015206361825134909144332670706e+2015816"
    # allow alpha == beta == 1  --> exp(z)
    @test isapprox(mittleff(1, -2), 0.1353352832366127)
end

@testset "branches" begin
    # convert z to float or complex float to avoid integer to negative power.
    @test isapprox(mittleff(0.9, 0.5, 22 + 22im), -2.7808021618204008e13 - 2.8561425165239754e13im)
    # Test branch that calls `Pint`
    @test isapprox(mittleff(0.1, 1.05, 0.9 + 0.5im), 0.17617901349590603 + 2.063981943021305im)
    @test isapprox(mittleff(4.1,1), 1.0358176744122032)
end

@testset "derivative" begin
    myder(f,x,h) = (f(x+h/2)-f(x-h/2))/h
    @test isapprox(myder(z -> mittleff(.4,z),.4,1e-5), mittleffderiv(.4,.4))
    @test abs(myder(z -> mittleff(big".4",z),big".4",BigFloat(1//10^16)) - mittleffderiv(big".4",big".4")) < 1e-32
end


## Julia program for ARMS Metropolis-Hastings
## author: weiya <szcfweiya@gmail.com>
## date: 2018-08-22

include("../GenRV/ars.jl")
using Main.corears

# data
x = rand(10)
y = rand(10)
# parameters
tau = 1

# function of h
function h(b)
    return(b * sum(x .* y) - sum(log.(1 .+ exp.(b*x))) - sum(1 .- 1 ./ (1 .+ exp.(b*x))) - b^2/(2*tau^2))
end

# derivative of h
function dh(b)
    return(sum(x .* y) - sum(x .* (1 .- 1 ./ (1 .+ exp.(b*x)) )) + sum(x .* exp.(b*x) ./ (1 .+ exp.(b*x)).^2 ) - b/tau^2)
end

function phi(x, yfixed::Array)
    if exp(h(x)) <= exp(hplus(x, yfixed))
        return(exp(h(x)))
    else
        return(exp(hplus(x, yfixed)))
    end
end

# ARMS
function arms(T, yfixed::Array)
    x = ones(T+1)
    for t = 1:T
        # generate Y
        while true
            global y
            y = gplus_sample(yfixed)
            u = rand()
            u <= exp(h(y)-hplus(y, yfixed)) && break
        end
        # accept or not
        v = rand()
        r = exp(h(y))*phi(x[t], yfixed)/(exp(h(x[t]))*phi(y, yfixed))
        println(r)
        if r >= 1
            x[t+1] = y
        else
            if v <= r
                x[t+1] = y
            else
                x[t+1] = x[t]
            end
        end
    end
    return(x)
end

## example
x = arms(10000, [-1.5, -0.7, 0.7, 1.5])


function budget_constraint(ind_vec::NTuple{N,<:Integer}, mdl::Model, add_args = nothing) where N

    R,Rᴸ = mdl.Prices[1:2]
    ω, τ = mdl.AddPrices
    v = get_vars(ind_vec, mdl.GridDef)
    #l₁, l₂, d₁, d₂
    # function c = ... . If negative it means that this selection is invalid
    #TODO:fix below (lambdas are state (emp/unemp) variables)
    Δₗ = v.l₂ - (1 - δᶠ) * v.l₁ * ( (Rᴸ) / 𝜋 )
    dᴿ = v.d₂ - v.l₂ * δᶠ * δˡ ≥ 0

    𝔠 = v.d₁ * ( ρ * R / 𝜋) + ( 1 - τ ) * ω * (v.λᵢ==1) + ϐ * (v.λᵢ==2) + Δₗ -
    δᶠ * v.l₁ * ( (Rᴸ) / 𝜋 ) - v.d₂

    # C_ deposit requirement:
    c_min = max(0.0, ( ( 1 - τ ) * ω * (v.λᵢ==1) + ϐ * (v.λᵢ==2) + v.d₁ * ( ρ * R / 𝜋) -
                    - δᶠ * v.l₁ * ( (Rᴸ) / 𝜋 ) ) )
    c = min(c_min,𝔠) + ( 𝔠 - min(c_min,𝔠) ) * ξ

    tozero_l = (ind_vec[2] - 1) == ind_vec[4]
    # tozero_d = (ind_vec[1] - 1) == ind_vec[3]
    return(c,Δₗ≥0,dᴿ,tozero_l)

end


# compute ols coefficients, fitted values, and errors
function lsfit(Y, X)
    beta = X\Y
    fit = X*beta
    errors = Y - fit
    return beta, fit, errors
end


# For defining the ODEs for describing all the molecular dynamics.
# Designed for use with the Sundials ODE solver.
# Args: xdot = vector, [dCa2+/dt, dBCa/dt, dECa/dt, dImCa/dt, dBCa*/dt]
#       x = current state of continuous variables, [Ca2+, BCa, ECa, ImCa, BCa*]
#       t = time
#       params =  Dict, parameters [calcium_baseline, ca_influx, ca_outflux, b_i, f_i, indicator,
#                 b_e, f_e, endogeneous, b_im, f_im, immobile,
#                 excitation_rate, release_rate, previous_BCa*]
# Returns: nothing

function molecularODE(t::Float64, x::Array{Float64}, xdot::Array{Float64}, baseline::Int64,
                      calcium_rate::Float64, indicator::Int64, endogeneous::Int64,
                      immobile::Int64, b_i::Float64, f_i::Float64, b_e::Float64, f_e::Float64,
                      b_im::Float64, f_im::Float64, excitation::Float64, release::Float64)
  B, E, Im = indicator - (x[2] + x[5]), endogeneous - x[3], immobile - x[4]
  # d[Ca2+]/dt = b_i*[BCa] + b_e*[ECa] + b_im*x[ImCa] - f_i*[Ca2+]*[B] - f_e*[Ca2+]*[E] + f_im*x[Ca2+]*[Im] + \beta_{Ca}([Ca2+]_0 - [Ca2+])
  xdot[1] = b_i*x[2] + b_e*x[3] + b_im*x[4] - f_i*x[1]*B - f_e*x[1]*E - f_im*x[1]*Im + calcium_rate * (baseline - x[1])
  # d[BCa]/dt = f_i*[Ca2+]*[B] - b_i*[BCa] + release*[BCa*] - excitation*[BCa]
  xdot[2] = f_i*x[1]*B - b_i*x[2] + release*x[5] - excitation*x[2]
  # d[ECa]/dt = f_e*[Ca2+]*[E] - b_e*[ECa]
  xdot[3] = f_e*x[1]*E - b_e*x[3]
  # d[ImCa]/dt = f_im*[Ca2+]*[Im] - b_im*[ImCa]
  xdot[4] = f_im*x[1]*Im - b_im*x[4]
  # d[BCa*]/dt = excitation*[BCa] - release*[BCa*]
  xdot[5] = excitation*x[2] - release*x[5]
  return nothing
end


#
# Compute Polar Decomposition via SVD
#

struct SVDAlg <: PolarAlg end

function solve!(alg::SVDAlg,
                X::Matrix{T}, U::Matrix{T}, H::Matrix{T}) where {T}

    F = svd(X, full = false)
    PQt = Array{T}(undef, size(U))
    mul!(PQt, F.U, transpose(F.V))
    copyto!(U, PQt)
    copyto!(H, F.V * Diagonal(F.S) * F.Vt)
    H = (1/T(2)) * (H + H')

    return SVDResult(U, H)
end


# Alias sampling
immutable FastSampler{Tv,Ti}
    F::Array{Tv,1}    # array of p[s.A[i]] / (1 / n)
    A::Array{Ti,1}    # array of indices of elements < (1 / n)
    V::Array{Ti,1}    # arary of indices of elements > (1 / n)
    n::Ti
end

# sample a random number
function sample{Tv,Ti}(s::FastSampler{Tv,Ti})
    #i = rand(1:s.n)
    i = ceil(Ti,rand()*s.n)
    f = rand()
    if f < s.F[i]
        return s.A[i]
    else
        return s.V[i]
    end
end

# sample a random number
function sampleMany{Tv,Ti}(s::FastSampler{Tv,Ti},sampCount::Ti)
    #i = rand(1:s.n)
    samples = Array{Ti,1}(sampCount)
    for j = 1:sampCount
        i = ceil(Ti,rand()*s.n)
        f = rand()
        if f < s.F[i]
            samples[j] = s.A[i]
        else
            samples[j] = s.V[i]
        end
    end
    return samples
end

# initialize the sampler. To get the residual error after building the sampler, set residual to true
function FastSampler{Tv}(p::Array{Tv,1}; residual::Bool = false)

    @assert(minimum(p) > 0, "The probability vector has a negative entry")

    n = length(p)

    # generate the Sampler
    F = Array{Tv}(n)
    A = Array{Int64}(n)
    V = Array{Int64}(n)

    # generate a normalized p, so that it sums up to n
    newp = n * copy(p) / sum(p)

    # fill up the arays A and V
    pos = 0                      # number of filled or half filled buckets
    posFilled = 0                # number of filled buckets

    for i in 1:n
        if newp[i] < 1 #&& !epsequal(newp[i], 1)
            pos = pos + 1
            A[pos] = i
            F[pos] = newp[i]
        end
    end

    err = 0
    for i in 1:n
        if newp[i] >= 1
            val = newp[i]

            while val >= 1
                if posFilled < pos
                    # fill up a bucket already containing a small element, and subtract that value from the current element
                    posFilled = posFilled + 1
                    V[posFilled] = i
                    val = val - (1 - F[posFilled])
                else
                    # create a bucket with only one large element. note posFilled = pos
                    pos = pos + 1
                    posFilled = posFilled + 1
                    A[pos] = V[pos] = i
                    F[pos] = 1
                    val = val - 1
                end
            end

            if val > 0 && pos == n
            	err = err + val
            end

            # treat the case in which val becomes smaller than 1. Have to do the pos < n check because of precision errors
            if val > 0 && pos < n
                pos = pos + 1
                A[pos] = i
                F[pos] = val
            end
        end
    end

    while posFilled < n
        # if we are in this case, posFilled should be equal to pos - 1
        posFilled = posFilled + 1
        V[posFilled] = A[posFilled]
        err += 1-F[posFilled]
    end

    if residual
        return FastSampler(F, A, V, n), err
    else
        return FastSampler(F, A, V, n)
    end
end


# using TuringModels
using CSV, DataFrames, Query, Turing, Optim, StatsBase, StatsPlots

# define precis from StatisticalRethinking module
using NamedArrays
function precis(df::DataFrame; digits=3, depth=Inf, alpha=0.11)
  m = zeros(length(names(df)), 5)
  for (indx, col) in enumerate(names(df))
    m[indx, 1] = mean(df[:, col])
    m[indx, 2] = std(df[:, col])
    q = quantile(df[:, col], [alpha/2, 0.5, 1-alpha/2])
    m[indx, 3] = q[1]
    m[indx, 4] = q[2]
    m[indx, 5] = q[3]
  end
  NamedArray(
    round.(m, digits=digits),
    (names(df), ["mean", "sd", "5.5%", "50%", "94.5%"]),
    ("Rows", "Cols"))
end

# ### snippet 5.3

# TODO: use population standard deviation to be consistent with Rethinking?
scale(x) = (x .- mean(x)) ./ std(x)

df = CSV.read(joinpath(@__DIR__, "data", "milk.csv"))
categorical!(df, :clade)

df = df |> @mutate(clade_id = levelcode(_.clade)) |> DataFrame

df.K = scale(df.kcal_per_g)

K = df.K
clade = df.clade_id

@model model5_9(K, clade) = begin
    n_clades = length(unique(clade))
    σ ~ Exponential(1)
    α ~ filldist(Normal(0, 0.5), n_clades)
    μ = α[clade]
    @. K ~ Normal(μ, σ)
end

model = model5_9(K, clade)

m5_9 = sample(model, NUTS(), 2000)
loglik = logprob"K = K | chain = m5_9, clade = clade, model = model5_9"

describe(m5_9)
plot(m5_9)

params = ["α["] .* map(string, 1:4) .* "]"
dict = Dict(zip(params, levels(df.clade)))

post = DataFrame(m5_9)
rename!(post, dict)
post_long = DataFrames.stack(post)

using VegaLite

post_long |>
  @filter(_.variable != "σ") |>

  @vlplot(y={"variable:o", title = ""}) +
  @vlplot(
    mark={
        :point,
        filled=true
    },
    x={
        "mean(value)",
        scale={zero=false},
        title="expected kcal (std)"},
    color={value=:black}) +
  @vlplot(
    mark={
        :errorbar,
        extent=:stdev
     },
     x={:value, title="expected kcal (std)"})

using Random
Random.seed!(63);

df.house = sample(repeat(1:4, inner = 8), nrow(df))

@model model5_10(K, clade, house) = begin
    n_clades = length(unique(clade))
    n_houses = length(unique(house))
    σ ~ Exponential(1)
    α ~ filldist(Normal(0, 0.5), n_clades)
    h ~ filldist(Normal(0, 0.5), n_houses)
    μ = α[clade] + h[house]
    @. K ~ Normal(μ, σ)
end

K, clade, house = df.K, df.clade_id, df.house
model = model5_10(K, clade, house)

m5_10 = sample(model, NUTS(), 2000)

post = DataFrame(m5_10)

function recover_levels(var::String, factor::CategoricalArray)
    flevels = levels(factor)
    nlevels = length(flevels)
    params = [var * "["] .* map(string, 1:nlevels) .* "]"
    dict = Dict(zip(params, flevels))
    return dict
end

clades = recover_levels("α", df.clade)
rename!(post, clades)

houses = categorical(["Gryffindor", "Ravenclaw", "Hufflepuff", "Slytherin"])
houses = recover_levels("h", houses)
rename!(post, houses)

post_long = DataFrames.stack(post)

post_long |>
  @filter(_.variable != "σ") |>

  @vlplot(y={"variable:o", title = ""}) +
  @vlplot(
    mark={
        :point,
        filled=true
    },
    x={
        "mean(value)",
        scale={zero=false},
        title="expected kcal (std)"},
    color={value=:black}) +
  @vlplot(
    mark={
        :errorbar,
        extent=:stdev
     },
     x={:value, title="expected kcal (std)"})


### Quaternion type heavily based off of Base lib Complex.jl

"""
    Quaternion{T<:Real} <: Number
Quaternion number type with real and imaginary parts of type `T`.
`QuatF64` and `QuatF32` or `QuaternionF64` and `QuaternionF32` are aliases for
`Quaternion{Float64}` and `Quaternion{Float32}` respectively.
"""
struct Quaternion{T<:Real} <: Number
    re::T
    im::T
    jm::T
    km::T
end

### Constructors
Quaternion(a::Real, b::Real, c::Real, d::Real) = Quaternion(promote(a,b,c,d)...)
Quaternion(x::Real) = Quaternion(x, zero(x), zero(x), zero(x))
Quaternion(z::Complex) = Quaternion(real(z), imag(z), zero(real(z)), zero(real(z)))
Quaternion(q::Quaternion) = q

(::Type{Quaternion{T}})(x::Real) where {T} = Quaternion{T}(x, 0, 0, 0)
(::Type{Quaternion{T}})(z::Complex) where {T} = Quaternion{T}(z.re, z.im, 0, 0)
(::Type{Quaternion{T}})(q::Quaternion) where {T} = Quaternion{T}(q.re, q.im, q.jm, q.km)

### For representation as q = z + c*j for complex z and c
Quaternion(z1::Complex, z2::Complex) = Quaternion(real(z1), imag(z1), real(z2), imag(z2))

# Doesn't seem like a good design, fix this ?
Quaternion(v::Vector{<:Real}) = Quaternion(v[1], v[2], v[3], v[4])

"""
    jm
Quaternion imaginary unit.
"""
const jm = Quaternion(false,false,true,false)
"""
    km
Quaternion imaginary unit.
"""
const km = Quaternion(false,false,false,true)

const QuaternionF16 = Quaternion{Float16}
const QuaternionF32 = Quaternion{Float32}
const QuaternionF64 = Quaternion{Float64}
const QuatF16 = Quaternion{Float16}
const QuatF32 = Quaternion{Float32}
const QuatF64 = Quaternion{Float64}

convert(::Type{Quaternion{T}}, x::Real) where {T<:Real} = Quaternion{T}(x,0,0,0)
convert(::Type{Quaternion{T}}, z::Complex) where {T<:Real} = Quaternion{T}(real(z),imag(z),0,0)
convert(::Type{Quaternion{T}}, q::Quaternion) where {T<:Real} = Quaternion{T}(q.re, q.im, q.jm, q.km)

promote_rule(::Type{Quaternion{T}}, ::Type{S}) where {T<:Real,S<:Real} =
    Quaternion{promote_type(T,S)}
promote_rule(::Type{Quaternion{T}}, ::Type{Complex{S}}) where {T<:Real,S<:Real} =
    Quaternion{promote_type(T,S)}
promote_rule(::Type{Quaternion{T}}, ::Type{Quaternion{S}}) where {T<:Real,S<:Real} =
    Quaternion{promote_type(T,S)}

widen(::Type{Quaternion{T}}) where {T} = Quaternion{widen(T)}

bswap(q::Quaternion) = Quaternion(bswap(q.re), bswap(q.im), bswap(q.jm), bswap(q.km))

### Redefine/overload methods for Complex
real(q::Quaternion) = q.re
imag(q::Quaternion) = q.im

"""
    jmag(q)
Return the j-imaginary part of the quaternion number `q`.
# Example
```jldoctest
julia> jmag(1 + 3im + 2jm + 5km)
2
```
"""
jmag(q::Quaternion) = q.jm

"""
    kmag(q)
Return the k-imaginary part of the quaternion number `q`.
# Example
```jldoctest
julia> kmag(1 + 3im + 2jm + 5km)
5
```
"""
kmag(q::Quaternion) = q.km

real(::Type{Quaternion{T}}) where {T<:Real} = T

isreal(q::Quaternion) = iszero(q.im) && iszero(q.jm) && iszero(q.km)
isinteger(q::Quaternion) = isreal(q) && isinteger(q.re)
isfinite(q::Quaternion) = isfinite(q.re) && isfinite(q.im) && isfinite(q.jm) && isfinite(q.km)
isnan(q::Quaternion) = isnan(q.re) || isnan(q.im) || isnan(q.jm) || isnan(q.km)
isinf(q::Quaternion) = isinf(q.re) || isinf(q.im) || isinf(q.jm) || isinf(q.km)
iszero(q::Quaternion) = iszero(q.re) && iszero(q.im) && iszero(q.jm) && iszero(q.km)
iscomplex(q::Quaternion) = iszero(q.jm) && iszero(q.km)
isone(q::Quaternion) = isreal(q) && isone(q.re)

zero(::Type{Quaternion{T}}) where {T<:Real} = Quaternion{T}(zero(T), zero(T), zero(T), zero(T))

"""
    quat(r::Real [, i::Real, j::Real, k::Real])
    quat(c::Complex [, z::Complex])
Convert real numbers or arrays to quaternion. `i`, `j`, `k`, default to zero.
Convert a complex number to a quaternion, with `j` and `k` parts zero.
Convert a pair of complex numbers to a quaternion `c + z*j`.
Mostly just a more convienient way to call the constructor Quaternion(...).
"""
quat(a::Real, b::Real, c::Real, d::Real) = Quaternion(a, b, c, d)
quat(x::Real) = Quaternion(x)
quat(z::Complex) = Quaternion(z)
quat(q::Quaternion) = q
quat(z1::Complex, z2::Complex) = Quaternion(z1, z2)
quat(v::Vector{<:Real}) = Quaternion(v)

"""
    complex(q::Quaternion)
Returns the complex part of a quaternion, equivalent to `complex(q.re, q.im)`.
"""
complex(::Type{Quaternion{T}}) where {T<:Real} = Complex{T}
complex(q::Quaternion) = complex(q.re, q.im)

"""
    quat(T::Type)
Returns an appropriate type which can represent a value of type `T` as a quaternion.
Equivalent to `typeof(quat(zero(T)))`.
"""
quat(::Type{T}) where {T<:Real} = Quaternion{T}
quat(::Type{Complex{T}}) where {T<:Real} = Quaternion{T}
quat(::Type{Quaternion{T}}) where {T<:Real} = Quaternion{T}

vec(q::Quaternion) = vcat(q.re, q.im, q.jm, q.km)

function show(io::IO, q::Quaternion)
    r, i, j, k = vec(q)
    compact = get(io, :compact, false)
    show(io, r)
    for (x, xm) in [(i, "im"), (j, "jm"), (k, "km")]
        if signbit(x) && !isnan(x)
            x = -x
            print(io, compact ? "-" : " - ")
        else
            print(io, compact ? "+" : " + ")
        end
        show(io, x)
        if !(isa(x,Integer) && !isa(x,Bool) || isa(x,AbstractFloat) && isfinite(x))
            print(io, "*")
        end
        print(io, xm)
    end
end

function show(io::IO, q::Quaternion{Bool})
    if q == im
        print(io, "im")
    elseif q == jm
        print(io, "jm")
    elseif q == km
        print(io, "km")
    else
        print(io, "Quaternion($(q.re),$(q.im),$(q.jm),$(q.km))")
    end
end

function show(io::IO, ::MIME"text/html", q::Quaternion)
    r, i, j, k = vec(q)
    compact = get(io, :compact, false)
    show(io, r)
    for (x, xm) in [(i, "i"), (j, "j"), (k, "k")]
        if signbit(x) && !isnan(x)
            x = -x
            print(io, compact ? "-" : " - ")
        else
            print(io, compact ? "+" : " + ")
        end
        show(io, x)
        if !(isa(x,Integer) && !isa(x,Bool) || isa(x,AbstractFloat) && isfinite(x))
            print(io, "*")
        end
        print(io, "" * xm * "")
    end
end

function show(io::IO, ::MIME"text/html", q::Quaternion{Bool})
    if q == im
        print(io, "i")
    elseif q == jm
        print(io, "j")
    elseif q == km
        print(io, "k")
    else
        print(io, "Quaternion($(q.re),$(q.im),$(q.jm),$(q.km))")
    end
end

function read(s::IO, ::Type{Quaternion{T}}) where {T}
    r = read(s,T)
    i = read(s,T)
    j = read(s,T)
    k = read(s,T)
    Quaternion{T}(r,i,j,k)
end
function write(s::IO, q::Quaternion)
    write(s,q.re,q.im,q.jm,q.km)
end

==(q::Quaternion, w::Quaternion) = q.re == w.re && q.im == w.im && q.jm == w.jm && q.km == w.km
==(q::Quaternion, z::Complex) = iscomplex(q) && q.im == imag(z) && q.re == real(z)
==(q::Quaternion, x::Real) = isreal(q) && q.re == x

isequal(q::Quaternion, w::Quaternion) = isequal(q.re,w.re) & isequal(q.im,w.im) &
                                        isequal(q.jm,w.jm) & isequal(q.km,w.km)

#TODO: hash

"""
    conj(q::Quaternion)
Compute the complex conjugate of a quaternion `z`.
"""
conj(q::Quaternion) = Quaternion(q.re,-q.im,-q.jm,-q.km)

"""
    abs(q::Quaternion)
Compute the absolute value (euclidean norm) of a quaternion.
"""
abs(q::Quaternion)  = sqrt(abs2(q))

"""
    abs2(q::Quaternion)
Compute the squared euclidean norm of a quaternion.
Equivalent to `q * conj(q)` but computationally faster and returns a real.
"""
abs2(q::Quaternion) = q.re*q.re + q.im*q.im + q.jm*q.jm + q.km*q.km

inv(q::Quaternion)  = conj(q)/abs2(q)
inv(q::Quaternion{<:Integer}) = inv(float(q))

-(q::Quaternion) = Quaternion(-q.re, -q.im, -q.jm, -q.km)
+(q::Quaternion, w::Quaternion) = Quaternion(q.re + w.re, q.im + w.im,
                                             q.jm + w.jm, q.km + w.km)
-(q::Quaternion, w::Quaternion) = Quaternion(q.re - w.re, q.im - w.im,
                                             q.jm - w.jm, q.km - w.km)
*(q::Quaternion, w::Quaternion) = Quaternion(q.re*w.re - q.im*w.im - q.jm*w.jm - q.km*w.km,
                                             q.re*w.im + q.im*w.re + q.jm*w.km - q.km*w.jm,
                                             q.re*w.jm - q.im*w.km + q.jm*w.re + q.km*w.im,
                                             q.re*w.km + q.im*w.jm - q.jm*w.im + q.km*w.re)

+(x::Real, q::Quaternion) = Quaternion(x + q.re, q.im, q.jm, q.km)
+(q::Quaternion, x::Real) = Quaternion(q.re + x, q.im, q.jm, q.km)
+(z::Complex, q::Quaternion) = Quaternion(real(z) + q.re, imag(z) + q.im, q.jm, q.km)
+(q::Quaternion, z::Complex) = Quaternion(q.re + real(z), q.im + imag(z), q.jm, q.km)
function -(x::Real, q::Quaternion)
    # we don't want the default type for -(Bool)
    re = x - q.re
    Quaternion(re, - oftype(re, q.im), - oftype(re, q.jm), - oftype(re, q.km))
end
function -(z::Complex, q::Quaternion)
    # we don't want the default type for -(Bool)
    re = real(z) - q.re
    Quaternion(re, imag(z) - q.im, - oftype(re, q.jm), - oftype(re, q.km))
end
-(q::Quaternion, x::Real) = Quaternion(q.re - x, q.im, q.jm, q.km)
-(q::Quaternion, z::Complex) = Quaternion(q.re - real(z), q.im - imag(z), q.jm, q.km)
*(x::Real, q::Quaternion) = Quaternion(x * q.re, x * q.im, x * q.jm, x * q.km)
*(q::Quaternion, x::Real) = Quaternion(q.re * x, q.im * x, q.jm * x, q.km * x)

*(q::Quaternion, z::Complex) = Quaternion(q.re*z.re - q.im*z.im, q.re*z.im + q.im*z.re,
                                                      q.jm*z.re + q.km*z.im, q.km*z.re - q.jm*z.im)
*(z::Complex, q::Quaternion) = Quaternion(z.re*q.re - z.im*q.im, z.re*q.im + z.im*q.re,
                                                      z.re*q.jm - z.im*q.km, z.re*q.km + z.im*q.jm)

/(a::R, q::S) where {R<:Real,S<:Quaternion} = (T = promote_type(R,S); a*inv(T(q)))
/(q::Quaternion, x::Real) = Quaternion(q.re/x, q.im/x, q.jm/x, q.km/x)
/(q::Quaternion, w::Quaternion) = q * inv(w)

rand(r::AbstractRNG, ::SamplerType{Quaternion{T}}) where {T<:Real} =
    Quaternion(rand(r,T), rand(r,T), rand(r,T), rand(r,T))

"""
When the type argument is quaternion, the values are drawn
from the circularly symmetric quaternion normal distribution.
This is std normal in each component but with variance scaled by 1/4.
"""
randn(r::AbstractRNG, ::Type{Quaternion{T}}) where {T<:AbstractFloat} =
    Quaternion(T(0.5)*randn(r,T), T(0.5)*randn(r,T), T(0.5)*randn(r,T), T(0.5)*randn(r,T))

"""
    norm(q::Quaternion, p::Real=2)
Computes the `p` norm of the quaternion components of `q` as a vector.
For the default `p=2` euclidean norm, this is same as the quaternion `abs()`.
"""
norm(q::Quaternion{T}, p::Real=2) where {T<:Real} = norm(vec(q), p)

"""
    normalize(q::Quaternion, p::Real=2)
Computes the normalized unit quaternion. Optional choice of `p` norm.
"""
normalize(q::Quaternion{T}, p::Real=2) where {T<:Real} = Quaternion(normalize(vec(q), p))

exp(q::Quaternion{<:Integer}) = exp(float(q))
log(q::Quaternion{<:Integer}) = log(float(q))

### See Neil Dantam - Quaternion Computation - for exp and log
function exp(q::Quaternion{<:AbstractFloat})
    V = norm([q.im q.jm q.km])
    s = if V < 1e-8
        1 - V*V/6 + V*V*V*V/120
    else
        sin(V)/V
    end
    exp(q.re)*(cos(V) + quat(zero(q.re), q.im*s, q.jm*s, q.km*s))
end

function log(q::Quaternion{<:AbstractFloat})
    V = norm([q.im q.jm q.km])
    Q = abs(q)
    ϕ = atan(V, q.re)
    s = if V < 1e-8
        (1+ϕ*ϕ/6 +7*ϕ*ϕ*ϕ*ϕ/360)/Q
    else
        ϕ/V
    end
    quat(log(Q), q.im*s, q.jm*s, q.km*s)
end

### NOTE: sqrt(q^2) == q is NOT true for pure quaternion for this computation
sqrt(q::Quaternion) = exp(0.5*log(q))
^(q::Quaternion, w::Quaternion) = exp(w*log(q))

function round(q::Quaternion, rr::RoundingMode=RoundNearest, ri::RoundingMode=rr; kwargs...)
    Quaternion(round(real(q), rr; kwargs...),
               round(imag(q), ri; kwargs...),
               round(jmag(q), ri; kwargs...),
               round(kmag(q), ri; kwargs...))
end

float(q::Quaternion{<:AbstractFloat}) = q
float(q::Quaternion) = Quaternion(float(q.re), float(q.im), float(q.jm), float(q.km))

big(::Type{Quaternion{T}}) where {T<:Real} = Quaternion{big(T)}
big(q::Quaternion{T}) where {T<:Real} = Quaternion{big(T)}(q)

## Matrix representations of Quaternions

"""
    cmatrix(q::Quaternion)
    cmatrix(q::Matrix{Quaternion})
Returns the complex matrix representation of a quaternion.
If we write `q = x + y*j` then this is `[x y; -conj(y) conj(x)]`.
This works if `q` is a quaternion of matrix of quaternions.
For a matrix of size `(n,m)` this gives a complex matrix of size `(2n, 2m)`.
"""
function cmatrix(q::Quaternion)
    [complex( q.re, q.im) complex(q.jm,  q.km);
     complex(-q.jm, q.km) complex(q.re, -q.im)]
end

function cmatrix(Q::Matrix{Quaternion{T}}) where {T}
    [complex.( real.(Q),  imag.(Q)) complex.( jmag.(Q),  kmag.(Q));
     complex.(-jmag.(Q),  kmag.(Q)) complex.( real.(Q), -imag.(Q))]
end

"""
    qmatrix(c::Matrix{Complex})
Given a complex matrix representation of a quaternion of quaternion matrix,
returns the quaternion representation of the matrix of size `(n÷2, m÷2)`.
This gives an inverse of the `cmatrix` function.
This function does no error checking, make sure your matrix is valid beforehand.
"""
function qmatrix(C::Matrix{Complex{T}}) where {T}
    n, m = size(C)
    quat.(C[1:n÷2, 1:m÷2], C[1:n÷2, m÷2+1:m])
end

## Array operations on complex numbers ##

quat(A::AbstractArray{<:Quaternion}) = A


"""
  genGaussian(mean_in, cov_in, class_label, n)

Draw from a gaussian distribution
# Arguments
- `mean_in` : means
- `cov_in` : covariances
- `class_label` : class_label
- `n` : number of samples to draw
"""
function genGaussian(mean_in, cov_in, class_label, n)
  nv = Distributions.MvNormal(mean_in, cov_in)
  X = rand(nv, n)
  y = ones(n) .* class_label
  return nv, transpose(X), y
end

"""
  genZafarData(n = 10000, d = pi/4)

Generate synthetic data from Zafar et al., 2017 Fairness Constraints: Mechanisms for Fair Classification
# Arguments
- `n=10000` : number of samples
- `d=pi/4` : discrimination factor
# Returns
- `X` : DataFrame containing features and protected attribute z
- `y` : Binary Target variable
"""
function genZafarData(n = 10000, d = pi/4)
    mu1, sigma1 = [2., 2.],  [[5., 1.] [1., 5.]]
    mu2, sigma2 = [-2.,-2.], [[10., 1.]  [1., 3.]]
    nv1, X1, y1 = genGaussian(mu1, sigma1, 1, Int64(floor(n/2))) # positive class
    nv2, X2, y2 = genGaussian(mu2, sigma2, -1, Int64(floor(n/2))) # negative class
    perm = shuffle(1:(2*Int64(floor(n/2))))
    X = vcat(X1, X2)[perm,:]
    y = vcat(y1, y2)[perm]
    rotation_mult = [[cos(d), -sin(d)] [sin(d), cos(d)]]
    X_aux = X * rotation_mult
    """ Generate the sensitive feature here """
    x_control = [] # this array holds the sensitive feature value
    for i in 1:(2*Int64(floor(n/2)))
        x = X_aux[i,:]
        # probability for each cluster that the point belongs to it
        p1 = Distributions.pdf(nv1, x)
        p2 = Distributions.pdf(nv2, x)
        # normalize the probabilities from 0 to 1
        s = p1+p2
        p1 = p1/s
        p2 = p2/s
        if rand(1)[1] < p1 # the first cluster is the positive class
            append!(x_control, 1.0)
        else
            append!(x_control, 0.0)
        end
    end
    X = DataFrame(X)
    X.z = x_control
    coerce!(X, :z => Multiclass)
    y = categorical(y)
    return X, y
end

"""
  genSubgroupData(n=10000, setting="B00")

Generate synthetic data from Loh et al., 2019 : Subgroup identification for precision medicine: A comparative review of 13 methods
# Arguments
- `n=10000` : number of samples
- `setting="B00"` : Simulation data setting: one of "B00", ..., "B02", "B1", ... , "B8"
For "B00", ..., "B02" there is no "bias" in the data, i.e. group membership has no effect on y.
whereas for  "B1", ... , "B8", there is a direct effect of group membership z on y, usually mediated by one or more features.
# Returns
- `X` : DataFrame containing features and protected attribute z
- `y` : Binary Target variable
"""
function genSubgroupData(n = 10000, setting = "B00")
    xdists = [
      Distributions.Normal(0.,1.), # x1 is N(0,1)
      Distributions.MvNormal([0., 0.], [[1, 0.5] [0.5, 1]]), # x2 and x3 are N(0,1) with cor 0.5
      Distributions.Exponential(1.), # x4 is exp(1)
      Distributions.Bernoulli(0.5), # x5 is Ber(0.5)
      Distributions.Categorical(repeat([.1], 10)), # x6 is Multinom. w. equal
      Distributions.MvNormal([0., 0., 0., 0.], [[1, 0.5, 0.5, 0.5] [0.5, 1, 0.5, 0.5] [0.5, 0.5, 1, 0.5] [0.5, 0.5, 0.5, 1]]
      ) # x7 to 10 are N(0,1) with cor 0.5
    ]
    z = rand(Distributions.Bernoulli(0.5), n)
    X = []
    for d in xdists
      if length(X) > 0
        xn = rand(d, n)
        if size(xn)[1] != n
          xn = transpose(xn)
        end
        X = hcat(X, xn)
      else
        X = rand(d, n)
      end
    end
    y = logit_fun(X, z, setting)
    X = DataFrame(X)
    X.z = z
    coerce!(X, :z => Multiclass)
    y = categorical(y)
    return X, y
end

# Logistic Link function
function log_link(x)
  return exp(x)/(1+exp(x))
end

"""
  logit_fun(X, z, setting)

Compute y from X and z according to a setting provided in Loh et al., 2019: Subgroup identification for precision medicine: A comparative review of 13 methods
# Arguments
- `X` : matrix of features
- `z` : vector of group assignments
- `setting` : Simulation data setting: one of "B00", ..., "B02", "B1", ... , "B8"
"""
function logit_fun(X, z, setting)
  if setting == "B00"
    logit = repeat([0.], length(z))
  elseif setting == "B01"
    logit = 0.5*(X[:,1] + X[:,2])
  elseif setting == "B02"
    logit = 0.5*(X[:,1] + X[:,2]) .^ 2
  elseif setting == "B1"
    logit = 0.5*(X[:,1] + X[:,2] - X[:,5]) + 2. * z .* map(x -> ifelse(isodd(Int64(x)), 1., 0.), X[:,6])
  elseif setting == "B2"
    logit = 0.5*X[:,2] + 2. * z .* map(x -> ifelse(x > 0, 1., 0.), X[:,1])
  elseif setting == "B3"
    logit = 0.3*(X[:,1] + X[:,2]) + 2. * z .* map(x -> ifelse(x > 0, 1., 0.), X[:,1])
  elseif setting == "B4"
    logit = 0.2*(X[:,2] + X[:,3] .- 2.) + 2. * z .* X[:,4]
  elseif setting == "B5"
    logit = 0.2*(X[:,1] + X[:,2] .- 3) + 2. * z .* map(x -> ifelse(x < 0, 1., 0.), X[:,1]) .* map(x -> ifelse(isodd(Int64(x)), 1., 0.), X[:,6])
  elseif setting == "B6"
    logit = 0.5*(X[:,2] .- 1.) + 2. * z .* map(x -> ifelse(abs(x) < 0.8, 1., 0.), X[:,1])
  elseif setting == "B7"
    logit = 0.2*(X[:,2] + X[:,2] .^ 2 .- 6.)  + 2. * z .* map(x -> ifelse(x < 0, 1., 0.), X[:,1])
  elseif setting == "B8"
    logit = 0.5 * X[:,2] + 2. * z .* X[:,5]
  end
  yprob = log_link.(logit)
  # Compute a random y with probability yprob
  return ifelse.(yprob .> rand(1)[1], 1, 0)
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    Segment(p1, p2)

A line segment with points `p1`, `p2`.
"""
struct Segment{Dim,T,V<:AbstractVector{Point{Dim,T}}} <: Polytope{1,Dim,T}
  vertices::V
end

measure(s::Segment) = norm(s.vertices[2] - s.vertices[1])

facets(s::Segment) = (v for v in s.vertices)

"""
    s1 ∩ s2

Compute the intersection of two line segments `s1` and `s2`.
"""
Base.intersect(s1::Segment, s2::Segment) = get(intersecttype(s1, s2))

abstract type Intersection end

Base.get(I::Intersection) = I.value

struct CrossingSegments{P<:Point} <: Intersection
  value::P
end

struct MidTouchingSegments{P<:Point} <: Intersection
  value::P
end

struct CornerTouchingSegments{P<:Point} <: Intersection
  value::P
end

struct OverlappingSegments{S<:Segment} <: Intersection
  value::S
end

struct NonIntersectingSegments <: Intersection end

Base.get(::NonIntersectingSegments) = nothing

"""
    intersecttype(s1, s2)

Compute the intersection type of two line segments `s1` and `s2`.

The intersection type can be one of five types according to
Balbes, R. and Siegel, J. 1990:

1. intersect at point interior to both segments
2. intersect at end point of one and an interior point of the other
3. intersect at end point of both segments
4. overlap at more than one point
5. do not overlap nor intersect

## References

* Balbes, R. and Siegel, J. 1990. [A robust method for calculating
  the simplicity and orientation of planar polygons]
  (https://www.sciencedirect.com/science/article/abs/pii/0167839691900198)
"""
function intersecttype(s1::Segment{2,T}, s2::Segment{2,T}) where {T}
  x̄ = center(s1)
  ȳ = center(s2)
  x1, x2 = vertices(s1)
  y1, y2 = vertices(s2)
  Q = [x1, y1, x2, y2, x1]

  # winding number of x̄
  windingx = zero(T)
  determinatex = true
  for j in 1:length(Q)-1
    α = ∠(Q[j], x̄, Q[j+1])
    if α ≈ 0 || α ≈ π || α ≈ -π
      determinatex = false
    end
    windingx += α
  end

  # winding number of ȳ
  windingy = zero(T)
  determinatey = true
  for j in 1:length(Q)-1
    β = ∠(Q[j], ȳ, Q[j+1])
    if β ≈ 0 || β ≈ π || β ≈ -π
      determinatey = false
    end
    windingy += β
  end

  # winding number ≈ 0?
  wxzero = isapprox(windingx, 0, atol=atol(T))
  wyzero = isapprox(windingy, 0, atol=atol(T))

  # there are three cases to consider
  # Balbes, R. and Siegel, J. 1990.
  if determinatex && determinatey # CASE (I)
    if !wxzero && !wyzero
      # configuration (1)
      CrossingSegments(intersectpoint(s1, s2))
    else
      # configuration (5)
      NonIntersectingSegments()
    end
  elseif determinatex || determinatey # CASE (II)
    if !(determinatex ? wxzero : wyzero)
      if x1 ≈ y1 || x1 ≈ y2 || x2 ≈ y1 || x2 ≈ y2
        # configuration (3)
        CornerTouchingSegments(intersectpoint(s1, s2))
      else
        # configuration (2)
        MidTouchingSegments(intersectpoint(s1, s2))
      end
    else
      # configuration (5)
      NonIntersectingSegments()
    end
  elseif !determinatex && !determinatey # CASE (III)
    if !isapprox((x2 - x1) × (y2 - y1), 0, atol=atol(T))
      # configuration (3)
      CornerTouchingSegments(intersectpoint(s1, s2))
    else
      # configuration (3), (4) or (5)
      intersectcolinear(s1, s2)
    end
  end
end

# compute the intersection of two line segments assuming that it is a point
function intersectpoint(s1::Segment{Dim,T}, s2::Segment{Dim,T}) where {Dim,T}
  x1, x2 = vertices(s1)
  y1, y2 = vertices(s2)

  v1 = x1 - x2
  v2 = y1 - y2
  v12 = x1 - y1

  # the intersection point lies in between x1 and x2 at a fraction s
  # (https://en.wikipedia.org/wiki/Line-line_intersection#Formulas)
  s = (v12[1] * v2[2] - v12[2] * v2[1]) /
      (v1[1] * v2[2] - v1[2] * v2[1])

  x1 - s*v1
end

# intersection of two line segments assuming that they are colinear
function intersectcolinear(s1::Segment{Dim,T}, s2::Segment{Dim,T}) where {Dim,T}
  # make sure the first segment is larger than the second
  sa, sb = measure(s1) < measure(s2) ? (s2, s1) : (s1, s2)
  x1, x2 = vertices(sa)
  y1, y2 = vertices(sb)

  # given that the segments are colinear we can simply
  # operate on the scalar parameters of the vertices
  # along the line defined by the vector --x1--x2->
  Δ = x2 - x1
  i = findfirst(!iszero, Δ)
  t1 = (y1 - x1)[i] / Δ[i]
  t2 = (y2 - x1)[i] / Δ[i]

  # fix direction of second segment to match the first
  t1 > t2 && ((t1, t2) = (t2, t1))

  if t1 < zero(T) && t2 < zero(T)
    NonIntersectingSegments()
  elseif t1 < zero(T) && t2 ≈ zero(T)
    CornerTouchingSegments(x1)
  elseif t1 < zero(T) && t2 > zero(T)
    OverlappingSegments(Segment(x1, x1 + t2*Δ))
  elseif t1 ≈ zero(T) && t2 > zero(T)
    OverlappingSegments(Segment(x1, x1 + t2*Δ))
  elseif t1 > zero(T) && t2 < one(T)
    OverlappingSegments(Segment(x1 + t1*Δ, x1 + t2*Δ))
  elseif t1 < one(T) && t2 ≈ one(T)
    OverlappingSegments(Segment(x1 + t1*Δ, x2))
  elseif t1 < one(T) && t2 > one(T)
    OverlappingSegments(Segment(x1 + t1*Δ, x2))
  elseif t1 ≈ one(T) && t2 > one(T)
    CornerTouchingSegments(x2)
  elseif t1 > one(T) && t2 > one(T)
    NonIntersectingSegments()
  else
    @error "please report bug"
  end
end


module UtilityFunctionsLH

using ArgCheck, DocStringExtensions
export CRRA, utility, marginal_utility, cons_age1, cons_growth, cons_path,
    euler_deviation, lifetime_utility, mu_wealth,
    pv_consumption
export discounted_sum

# -------------  Types

abstract type UtilityFunction end
abstract type UtilityOneArg <: UtilityFunction end

"""
	$(SIGNATURES)

Utility = c^(1-sigma) / (1-sigma) - 1.0.

Log utility is special case when `sigma == 1.0`.
"""
mutable struct CRRA <: UtilityOneArg
    sigma  ::  Float64
    dbg  ::  Bool
 end

include("crra.jl")

## ---------  Helpers

"""
	$(SIGNATURES)

Discounted sum: `xV[1] + ... + pBeta ^ (T-1) * xV[T]`
"""
function discounted_sum(pBeta :: T1, xV) where T1 <: AbstractFloat
    T = length(xV);
    s = zero(eltype(xV));
    for t = T : -1 : 1
        s = pBeta * s + xV[t];
    end
    return s
end

# Vector with constant growth factor
function const_growth_vector(c1 :: T1,
    growthFactor :: T1, T :: Integer) where T1 <: AbstractFloat

    cV = Vector{T1}(undef, T);
    ct = c1;
    for j in eachindex(cV)
        cV[j] = ct;
        ct *= growthFactor;
    end
    return cV
end

# Sum of a constant growth vector: 1 + growthFactor + ... + growthFactor ^ (T-1)
function const_growth_sum(growthFactor :: T1, T :: Integer) where T1
    if isapprox(growthFactor, one(T1))
        gSum = T1(T);
    else
        gSum = (one(T1) - growthFactor ^ T) / (one(T1) - growthFactor);
    end
    return gSum
end

end # module


using Sobol
using LatinHypercubeSampling

abstract type SamplingAlgorithm end

"""
GridSample{T}

T is the step dx for lb:dx:ub
"""
struct GridSample{T} <: SamplingAlgorithm
    dx::T
end

struct UniformSample <: SamplingAlgorithm end
struct SobolSample <: SamplingAlgorithm end
struct LatinHypercubeSample <: SamplingAlgorithm end

"""
LowDiscrepancySample{T}

T is the base for the sequence
"""
struct LowDiscrepancySample{T} <: SamplingAlgorithm
    base::T
end

struct RandomSample <: SamplingAlgorithm end

struct KroneckerSample{A,B} <: SamplingAlgorithm
    alpha::A
    s0::B
end

struct GoldenSample <: SamplingAlgorithm end

struct SectionSample{T} <: SamplingAlgorithm
    x0::Vector{T}
    sa::SamplingAlgorithm
end

"""
sample(n,lb,ub,S::GridSample)

Returns a tuple containing numbers in a grid.

"""
function sample(n,lb,ub,S::GridSample)
    dx = S.dx
    if lb isa Number
        return vec(rand(lb:S.dx:ub,n))
    else
        d = length(lb)
        x = [[rand(lb[j]:dx[j]:ub[j]) for j = 1:d] for i in 1:n]
        return Tuple.(x)
    end
end

"""
sample(n,lb,ub,::UniformRandom)

Returns a Tuple containing uniform random numbers.
"""
function sample(n,lb,ub,::UniformSample)
    if lb isa Number
        return rand(Uniform(lb,ub),n)
    else
        d = length(lb)
        x = [[rand(Uniform(lb[j],ub[j])) for j in 1:d] for i in 1:n]
        return Tuple.(x)
    end
end

"""
sample(n,lb,ub,::SobolSampling)

Returns a Tuple containing Sobol sequences.
"""
function sample(n,lb,ub,::SobolSample)
    s = SobolSeq(lb,ub)
    skip(s,n)
    if lb isa Number
        return [next!(s)[1] for i = 1:n]
    else
        return Tuple.([next!(s) for i = 1:n])
    end
end

"""
sample(n,lb,ub,::LatinHypercube)

Returns a Tuple containing LatinHypercube sequences.
"""
function sample(n,lb,ub,::LatinHypercubeSample)
    d = length(lb)
    if lb isa Number
        x = vec(LHCoptim(n,d,1)[1])
        # x∈[0,n], so affine transform
        return @. (ub-lb) * x/(n) + lb
    else
        lib_out = float(LHCoptim(n,d,1)[1])
        # x∈[0,n], so affine transform column-wise
        @inbounds for c = 1:d
            lib_out[:,c] = (ub[c]-lb[c])*lib_out[:,c]/n .+ lb[c]
        end
        x = [lib_out[i,:] for i = 1:n]
        return Tuple.(x)
    end
end

"""
sample(n,lb,ub,S::LowDiscrepancySample)

Low discrepancy sample:
- Dimension 1: Van der Corput sequence
- Dimension > 1: Halton sequence
If dimension d > 1, all bases must be coprime with each other.
"""
function sample(n,lb,ub,S::LowDiscrepancySample)
    d = length(lb)
    if d == 1
        #Van der Corput
        b = S.base
        x = zeros(Float32,n)
        for i = 1:n
            expansion = digits(i,base = b)
            L = length(expansion)
            val = zero(Float32)
            for k = 1:L
                val += expansion[k]*float(b)^(-(k-1)-1)
            end
            x[i] = val
        end
        # It is always defined on the unit interval, resizing:
        return @. (ub-lb) * x + lb
    else
        #Halton sequence
        x = zeros(Float32,n,d)
        for j = 1:d
            b = S.base[j]
            for i = 1:n
                val = zero(Float32)
                expansion = digits(i, base = b)
                L = length(expansion)
                val = zero(Float32)
                for k = 1:L
                    val += expansion[k]*float(b)^(-(k-1)-1)
                end
                x[i,j] = val
            end
        end
        #Resizing
        # x∈[0,1], so affine transform column-wise
        @inbounds for c = 1:d
            x[:,c] = (ub[c]-lb[c])*x[:,c] .+ lb[c]
        end

        y = [x[i,:] for i = 1:n]
        return Tuple.(y)
    end
end

"""
sample(n,d,D::Distribution)

Returns a Tuple containing numbers distributed as D
"""
function sample(n,d,D::Distribution)
    if d == 1
        return rand(D,n)
    else
        x = [[rand(D) for j in 1:d] for i in 1:n]
        return Tuple.(x)
    end
end

"""
sample(n,d,K::KroneckerSample)

Returns a Tuple containing numbers following the Kronecker sample
"""
function sample(n,lb,ub,K::KroneckerSample)
    d = length(lb)
    alpha = K.alpha
    s0 = K.s0
    if d == 1
        x = zeros(n)
        @inbounds for i = 1:n
            x[i] = (s0+i*alpha)%1
        end
        return @. (ub-lb) * x + lb
    else
        x = zeros(n,d)
        @inbounds for j = 1:d
            for i = 1:n
                x[i,j] = (s0[j] + i*alpha[j])%i
            end
        end
        #Resizing
        # x∈[0,1], so affine transform column-wise
        @inbounds for c = 1:d
            x[:,c] = (ub[c]-lb[c])*x[:,c] .+ lb[c]
        end

        y = [x[i,:] for i = 1:n]
        return Tuple.(y)
    end
end

function sample(n,lb,ub,G::GoldenSample)
    d = length(lb)
    if d == 1
        x = zeros(n)
        g = (sqrt(5)+1)/2
        a = 1.0/g
        for i = 1:n
            x[i] = (0.5+a*i)%1
        end
        return @. (ub-lb) * x + lb
    else
        x = zeros(n,d)
        for j = 1:d
            #Approximate solution of x^(d+1) = x + 1, a simple newton is good enough
            y = 2.0
            for s = 1:10
                g = (1+y)^(1/(j+1))
            end
            a = 1.0/g
            for i = 1:n
                x[i,j] = (0.5+a*i)%1
            end
        end
        @inbounds for c = 1:d
            x[:,c] = (ub[c]-lb[c])*x[:,c] .+ lb[c]
        end
        y = [x[i,:] for i = 1:n]
        return Tuple.(y)
    end
end

fixed_dimensions(
        section_sampler::SectionSample)::Vector{Int64} = findall(
    x->x == false, isnan.(section_sampler.x0))

free_dimensions(
        section_sampler::SectionSample)::Vector{Int64} = findall(
    x->x == true, isnan.(section_sampler.x0))

"""
sample(n,d,K::SectionSample)

Returns Tuples constrained to a section.

In surrogate-based identification and control, optimization can alternate between unconstrained sampling in the full-dimensional parameter space, and sampling constrained on specific sections (e.g. a planes in a 3D volume),

A SectionSampler allows sampling and optimizing on a subset of 'free' dimensions while keeping 'fixed' ones constrained.
The sampler is defined as in e.g.

section_sampler_y_is_10 = SectionSample([NaN64, NaN64, 10.0, 10.0], Surrogates.UniformSample())

where the first argument is a Vector{T} in which numbers are fixed coordinates and `NaN`s correspond to free dimensions, and the second argument is a SamplingAlgorithm which is used to sample in the free dimensions.
"""
function sample(n,lb,ub,section_sampler::SectionSample)
    if lb isa Number
        return rand(Uniform(lb,ub),n)
    else
        d_free = Surrogates.free_dimensions(section_sampler)
        new_samples = sample(n, lb[d_free], ub[d_free], section_sampler.sa)
        out_as_vec = repeat(section_sampler.x0', n, 1)
        for y in 1:size(out_as_vec,1)
            for xi in 1:length(d_free)
                out_as_vec[y,xi] = new_samples[y][xi]
            end
        end
        [Tuple(out_as_vec[y,:]) for y in 1:size(out_as_vec,1)]
    end
end


module OpenModels
import Base: ==
using SemanticModels.ModelTools
import SemanticModels.ModelTools: model
using Catlab.Graphics.Graphviz
import Catlab.Graphics.Graphviz: Graph, Edge

export OpenModel

struct OpenModel{V,M} <: AbstractModel
    dom::V
    model::M
    codom::V
end

"""    Graph(f::OpenModel)

convert an OpenModel into a GraphViz Graph. This calls Graph(::Model) and then adds vertices (and edges) for the domain and codomain of the open model.
"""
function Graph(f::OpenModel)
    g = Graph(f.model)
    A, M, B = f.dom, f.model, f.codom
    stmts_dom = map(enumerate(A)) do (i,a)
        m = M.S[a]
        Edge(["I$i", "$m"], Attributes(:style=>"dashed"))
    end
    stmts_codom = map(enumerate(B)) do (i,a)
        m = M.S[a]
        Edge(["$m", "O$i"], Attributes(:style=>"dashed"))
    end
    append!(g.stmts, append!(stmts_dom, stmts_codom))
    return g
end

function ==(f::OpenModel,g::OpenModel)
    all(isequal.(f.dom, g.dom)) && all(isequal.(f.codom, g.codom)) && f.model == g.model
end

# otimes and compose general that use disjoin union functions

end


# Experiment 4.1 from:
# C.-H. Chen, L.H. Lee: Stochastic Simulation Optimization, World Scientific, 2011

using Random
using EventSimulation
using Distributions

mutable struct SeqQueue <: AbstractState
    q1::SimResource{Int}
    q2::SimResource{Int}
    arrive::Vector{Float64}
    depart::Vector{Float64}
end

arrival(s) = (provide!(s, s.state.q1, 1); push!(s.state.arrive, s.now))
stage1(s) = register!(s, leave1, rand(Uniform(1, 39)))
leave1(s) = (provide!(s, s.state.q2, 1); request!(s, s.state.q1, 1, stage1))
stage2(s) = register!(s, leave2, rand(Uniform(5, 45)))
leave2(s) = (push!(s.state.depart, s.now); request!(s, s.state.q2, 1, stage2))

function exec(n1, n2)
    state = SeqQueue(SimResource{Int}(), SimResource{Int}(), Float64[], Float64[])
    s = Scheduler(state)
    repeat_register!(s, arrival, (x -> length(s.state.arrive) < 100 ? randexp() : nothing))
    foreach(i -> request!(s, state.q1, 1, stage1), 1:n1)
    foreach(i -> request!(s, state.q2, 1, stage2), 1:n2)
    go!(s)
    @assert length(state.depart) == length(state.arrive)
    return mean(state.depart) - mean(state.arrive)
end

for i in 1:10
    v = mean(exec(10+i, 21-i) for j in 1:2048)
    println((10+i, 21-i),":\t", v)
end


function coarsen_grid(grd::Grid)

    _, ks, _ = findnz(grd.Proj)

    _Nxc = Int64(ceil(grd._Nx/2))
    _Nyc = Int64(ceil(grd._Ny/2))
    Nbuffer = Int64(ceil(grd._Nbuffer/2))

    dr = 2*grd.dr
    if iseven(grd._Nbuffer)
        r0 = grd.r0
        nan_outside_boundaries = grd._nan_outside_boundaries[1:2:end,1:2:end]
    else
        r0 = grd.r0 + [1,1]*grd.dr
        nan_outside_boundaries = grd._nan_outside_boundaries[2:2:end,2:2:end]
    end
    Nxc, Nyc = size(nan_outside_boundaries)
    r1 = r0 + dr*[Nxc, Nyc]

    points = zeros(2, grd.Nk)
    projc_js = zeros(grd.Nk)
    Nkc = 0
    for k=1:grd.Nk
        kc = ks[k]
        j = Int(ceil(kc / grd._Nx))
        i = kc - (j - 1) * grd._Nx
        if isodd(i) && isodd(j)
            Nkc += 1
            points[:,Nkc] = grd.points[:,k]

            ic = (i + 1)/2
            jc = (j + 1)/2
            projc_js[Nkc] = ic + (jc-1)*_Nxc
        end
    end

    Projc = sparse([1:Nkc...], projc_js[1:Nkc], ones(Nkc), Nkc, _Nxc*_Nyc)
    return Grid(r0, r1, points[:,1:Nkc], Projc, dr, Nkc, _Nxc, _Nyc, Nbuffer, nan_outside_boundaries)
end

function restriction_and_interpolation(xb, finegrid, coarsegrid, extr_c, extr_f)

    Interp_cart, Restr_cart = cartesian_2d_crs(coarsegrid._Nx, coarsegrid._Ny, 4, 4)
    Interp_i = finegrid.Proj*Interp_cart*transpose(coarsegrid.Proj)
    Restr_i = coarsegrid.Proj*Restr_cart*transpose(finegrid.Proj)
    xb_c = Restr_i*xb

    Interp_mat = extr_f.Proj*Interp_i*extr_c.R*transpose(extr_c.Proj)
    Interp_f = extr_f.Proj*Interp_i*(I-extr_c.R)*xb_c

    Restr_mat = extr_c.Proj*Restr_i*extr_f.R*transpose(extr_f.Proj)
    Restr_f = extr_c.Proj*Restr_i*(I-extr_f.R)*xb

    Interp = AffineMap{Float64}(Interp_mat, Interp_f)
    Restr = AffineMap{Float64}(Restr_mat, Restr_f)
    return Interp, Restr, xb_c
end


using DiffEqFlux, OrdinaryDiffEq, DiffEqSensitivity
using CUDA, Test, Zygote
CUDA.allowscalar(false)

H = CuArray(rand(Float32, 2, 2))
ann = FastChain(FastDense(1, 4, tanh))
p = initial_params(ann)

function func(x, p, t)
    ann([t],p)[1]*H*x
end

x0 = CuArray(rand(Float32, 2))
x1 = CuArray(rand(Float32, 2))

prob = ODEProblem(func, x0, (0.0f0, 1.0f0))

function evolve(p)
    solve(prob, Tsit5(), p=p, save_start=false,
          save_everystep=false, abstol=1e-4, reltol=1e-4,
          sensealg=QuadratureAdjoint(autojacvec=ZygoteVJP())).u[1]
end

function cost(p)
    x = evolve(p)
    c = sum(abs,x - x1)
    #println(c)
    c
end

grad = Zygote.gradient(cost,p)[1]
@test !iszero(grad[1])
@test iszero(grad[2:4])
@test !iszero(grad[5])
@test iszero(grad[6:end])