using ModernRoboticsBook
using Test

import LinearAlgebra
const linalg = LinearAlgebra

@testset "ModernRoboticsBook.jl" begin
    @testset "basic helper functions" begin
        @test NearZero(-1e-7)
        @test Normalize([1, 2, 3]) == [0.2672612419124244,
                                       0.5345224838248488,
                                       0.8017837257372732]
    end
    @testset "chapter 3: rigid-body motions" begin
        @test RotInv([0 0 1; 1 0 0; 0 1 0]) == [0 1 0; 0 0 1; 1 0 0]
        @test VecToso3([1, 2, 3]) == [0 -3 2; 3 0 -1; -2 1 0]
        @test so3ToVec([0 -3 2; 3 0 -1; -2 1 0]) == [1, 2, 3]
        @test AxisAng3([1, 2, 3]) == ([0.2672612419124244,
                                       0.5345224838248488,
                                       0.8017837257372732], 3.7416573867739413)

        @test isapprox(MatrixExp3([ 0 -3  2;
                                    3  0 -1;
                                   -2  1  0]),
                       [-0.694921  0.713521 0.0892929
                        -0.192007 -0.303785 0.933192
                         0.692978  0.63135  0.348107 ]; rtol=1e-6)
        @test MatrixExp3(zeros(3, 3)) == [1 0 0; 0 1 0; 0 0 1]

        @test isapprox(MatrixLog3([0 0 1;
                                   1 0 0;
                                   0 1 0]),
                       [ 0.0    -1.2092  1.2092
                         1.2092  0.0    -1.2092
                        -1.2092  1.2092  0.0   ]; rtol=1e-6)
        @test MatrixLog3(zeros(3, 3)) == zeros(3, 3)
        @test MatrixLog3([-3 0 0; 0 1 0; 0 0  1]) == [0 -π 0; π 0 0; 0 0 0]
        @test MatrixLog3([-2 0 0; 0 1 0; 0 0 -1]) == [0 0 π; 0 0 0; -π 0 0]
        @test MatrixLog3([1 0 0; 0 -1 0; 0 0 -1]) == [0 0 0; 0 0 -π; 0 π 0]

        @test RpToTrans([1 0  0;
                         0 0 -1;
                         0 1  0], [1, 2, 5]) == [1 0  0 1
                                                 0 0 -1 2
                                                 0 1  0 5
                                                 0 0  0 1]
        @test TransToRp([1 0  0 0;
                         0 0 -1 0;
                         0 1  0 3;
                         0 0  0 1]) == ([1 0  0;
                                         0 0 -1;
                                         0 1  0], [0, 0, 3])
        @test TransInv([1 0  0 0;
                        0 0 -1 0;
                        0 1  0 3;
                        0 0  0 1]) == [1  0 0  0
                                       0  0 1 -3
                                       0 -1 0  0
                                       0  0 0  1]
        @test VecTose3([1, 2, 3, 4, 5, 6]) == [ 0 -3  2  4
                                                3  0 -1  5
                                               -2  1  0  6
                                                0  0  0  0]
        @test se3ToVec([ 0 -3  2  4;
                         3  0 -1  5;
                        -2  1  0  6;
                         0  0  0  0]) == [1, 2, 3, 4, 5, 6]
        @test Adjoint([1  0  0  0;
                       0  0 -1  0;
                       0  1  0  3;
                       0  0  0  1]) == [1  0  0  0  0  0
                                        0  0 -1  0  0  0
                                        0  1  0  0  0  0
                                        0  0  3  1  0  0
                                        3  0  0  0  0 -1
                                        0  0  0  0  1  0]
        @test ScrewToAxis([3; 0; 0], [0; 0; 1], 2) == [0, 0, 1, 0, -3, 2]

        @test AxisAng6([1, 0, 0, 1, 2, 3]) == ([1, 0, 0, 1, 2, 3], 1)
        @test AxisAng6([0, 0, 0, 0, 0, 4]) == ([0, 0, 0, 0, 0, 1], 4)

        @test MatrixExp6([0    0    0    0   ;
                          0    0   -π/2  3π/4;
                          0    π/2  0    3π/4;
                          0    0    0    0   ]) ≈ [1  0  0  0;
                                                   0  0 -1  0;
                                                   0  1  0  3;
                                                   0  0  0  1]
        @test MatrixLog6([1  0  0  0;
                          0  0 -1  0;
                          0  1  0  3;
                          0  0  0  1]) ≈ [0    0    0    0   ;
                                          0    0   -π/2  3π/4;
                                          0    π/2  0    3π/4;
                                          0    0    0    0   ]
        @test MatrixLog6(Array(linalg.Diagonal(ones(4)))) == zeros(4, 4)

        @test isapprox(ProjectToSO3([ 0.675  0.150  0.720;
                                      0.370  0.771 -0.511;
                                     -0.630  0.619  0.472]),
                       [ 0.679011  0.148945   0.718859
                         0.373207  0.773196  -0.512723
                        -0.632187  0.616428   0.469421]; rtol=1e-6)
        @test ProjectToSO3([-1 0 0; 0 -1 0; 0 0 -1]) == [1 0 0; 0 -1 0; 0 0 -1]

        @test isapprox(ProjectToSE3([ 0.675  0.150  0.720  1.2;
                                      0.370  0.771 -0.511  5.4;
                                     -0.630  0.619  0.472  3.6;
                                      0.003  0.002  0.010  0.9]),
                                    [ 0.679011  0.148945   0.718859  1.2
                                      0.373207  0.773196  -0.512723  5.4
                                     -0.632187  0.616428   0.469421  3.6
                                      0.0       0.0        0.0       1.0]; rtol=1e-6)

        @test DistanceToSO3([1.0  0.0  0.0 ;
                             0.0  0.1 -0.95;
                             0.0  1.0  0.1 ]) == 0.08835298523536149

        @test DistanceToSE3([1.0  0.0  0.0   1.2 ;
                             0.0  0.1 -0.95  1.5 ;
                             0.0  1.0  0.1  -0.9 ;
                             0.0  0.0  0.1   0.98]) == 0.13493053768513638
        @test DistanceToSE3(Array(reshape(1:9, (3, 3)))) >= 1e+9

        @test !TestIfSO3([1.0  0.0  0.0 ;
                          0.0  0.1 -0.95;
                          0.0  1.0  0.1 ])
        @test !TestIfSE3([1.0  0.0  0.0   1.2 ;
                          0.0  0.1 -0.95  1.5 ;
                          0.0  1.0  0.1  -0.9 ;
                          0.0  0.0  0.1   0.98])
    end
    @testset "chapter 4: forward kinematics" begin
        M = [-1  0  0  0 ;
              0  1  0  6 ;
              0  0 -1  2 ;
              0  0  0  1 ]

        Blist = [ 0  0 -1  2  0  0   ;
                  0  0  0  0  1  0   ;
                  0  0  1  0  0  0.1 ]'

        Slist = [ 0  0  1  4  0  0   ;
                  0  0  0  0  1  0   ;
                  0  0 -1 -6  0 -0.1 ]'

        thetalist = [ π/2, 3, π ]

        @test isapprox(FKinBody(M, Blist, thetalist),
                       [-1.14424e-17  1.0           0.0  -5.0     ;
                         1.0          1.14424e-17   0.0   4.0     ;
                         0.0          0.0          -1.0   1.68584 ;
                         0.0          0.0           0.0   1.0     ]; rtol=1e-6)
        @test isapprox(FKinSpace(M, Slist, thetalist),
                       [-1.14424e-17  1.0           0.0  -5.0     ;
                         1.0          1.14424e-17   0.0   4.0     ;
                         0.0          0.0          -1.0   1.68584 ;
                         0.0          0.0           0.0   1.0     ]; rtol=1e-6)
    end
    @testset "chapter 5: velocity kinematics and statics" begin
        Blist = [0 0 1   0 0.2 0.2;
                 1 0 0   2   0   3;
                 0 1 0   0   2   1;
                 1 0 0 0.2 0.3 0.4]'

        Slist = [0 0 1   0 0.2 0.2;
                 1 0 0   2   0   3;
                 0 1 0   0   2   1;
                 1 0 0 0.2 0.3 0.4]'

        thetalist = [0.2, 1.1, 0.1, 1.2]

        @test isapprox(JacobianBody(Blist, thetalist),
                       [ -0.0452841  0.995004    0.0       1.0 ;
                          0.743593   0.0930486   0.362358  0.0 ;
                         -0.667097   0.0361754  -0.932039  0.0 ;
                          2.32586    1.66809     0.564108  0.2 ;
                         -1.44321    2.94561     1.43307   0.3 ;
                         -2.0664     1.82882    -1.58869   0.4 ]; rtol=1e-5)
        @test isapprox(JacobianSpace(Slist, thetalist),
                       [0.0  0.980067  -0.0901156   0.957494 ;
                        0.0  0.198669   0.444554    0.284876 ;
                        1.0  0.0        0.891207   -0.0452841;
                        0.0  1.95219   -2.21635    -0.511615 ;
                        0.2  0.436541  -2.43713     2.77536  ;
                        0.2  2.96027    3.23573     2.22512  ]; rtol=1e-5)
    end
    @testset "chapter 6: inverse kinematics" begin
        Blist = [ 0  0 -1  2  0  0   ;
                  0  0  0  0  1  0   ;
                  0  0  1  0  0  0.1 ]'

        Slist = [ 0  0  1  4  0  0   ;
                  0  0  0  0  1  0   ;
                  0  0 -1 -6  0 -0.1 ]'

        M = [ -1  0  0  0 ;
               0  1  0  6 ;
               0  0 -1  2 ;
               0  0  0  1 ]

        T = [ 0  1  0     -5 ;
              1  0  0      4 ;
              0  0 -1 1.6858 ;
              0  0  0      1 ]

        thetalist0 = [1.5, 2.5, 3]

        eomg, ev = 0.01, 0.001

        thetalist, success = IKinBody(Blist, M, T, thetalist0, eomg, ev)
        @test isapprox(thetalist, [1.57074, 2.99967, 3.14154]; rtol=1e-5)
        @test success

        thetalist, success = IKinSpace(Slist, M, T, thetalist0, eomg, ev)
        @test isapprox(thetalist, [1.57074, 2.99966, 3.14153]; rtol=1e-5)
        @test success
    end
    @testset "chapter 8: dynamics of open chains" begin
        @test ad([1, 2, 3, 4, 5, 6]) == [ 0.0  -3.0   2.0   0.0   0.0   0.0
                                          3.0   0.0  -1.0   0.0   0.0   0.0
                                         -2.0   1.0   0.0   0.0   0.0   0.0
                                          0.0  -6.0   5.0   0.0  -3.0   2.0
                                          6.0   0.0  -4.0   3.0   0.0  -1.0
                                         -5.0   4.0   0.0  -2.0   1.0   0.0 ]

        thetalist = [0.1, 0.1, 0.1]
        dthetalist = [0.1, 0.2, 0.3]
        ddthetalist = [2, 1.5, 1]
        g = [0, 0, -9.8]
        Ftip = [1, 1, 1, 1, 1, 1]
        M01 = [1  0  0         0 ;
               0  1  0         0 ;
               0  0  1  0.089159 ;
               0  0  0         1 ]
        M12 = [0  0  1     0.28 ;
               0  1  0  0.13585 ;
              -1  0  0        0 ;
               0  0  0        1 ]
        M23 = [1  0  0        0 ;
               0  1  0  -0.1197 ;
               0  0  1    0.395 ;
               0  0  0        1 ]
        M34 = [1  0  0        0 ;
               0  1  0        0 ;
               0  0  1  0.14225 ;
               0  0  0        1 ]
        G1 = linalg.Diagonal([0.010267, 0.010267, 0.00666, 3.7, 3.7, 3.7])
        G2 = linalg.Diagonal([0.22689, 0.22689, 0.0151074, 8.393, 8.393, 8.393])
        G3 = linalg.Diagonal([0.0494433, 0.0494433, 0.004095, 2.275, 2.275, 2.275])
        Glist = [G1, G2, G3]
        Mlist = [M01, M12, M23, M34]
        Slist = [ 1  0  1      0  1      0 ;
                  0  1  0 -0.089  0      0 ;
                  0  1  0 -0.089  0  0.425 ]'

        taulist_actual = InverseDynamics(thetalist, dthetalist, ddthetalist, g, Ftip, Mlist, Glist, Slist)

        @test taulist_actual ≈ [74.69616155287451, -33.06766015851458, -3.230573137901424]

        @test isapprox(MassMatrix(thetalist, Mlist, Glist, Slist),
                       [ 22.5433      -0.307147  -0.00718426;
                         -0.307147     1.96851    0.432157  ;
                         -0.00718426   0.432157   0.191631  ]; rtol=1e-5)

        @test VelQuadraticForces(thetalist, dthetalist, Mlist, Glist, Slist) ≈ [  0.26453118054501235 ;
                                                                                -0.0550515682891655  ;
                                                                                -0.006891320068248911]

        @test GravityForces(thetalist, g, Mlist, Glist, Slist) == [  28.40331261821983  ;
                                                                    -37.64094817177068  ;
                                                                     -5.4415891999683605]

        @test EndEffectorForces(thetalist, Ftip, Mlist, Glist, Slist) ≈ [ 1.4095460782639782;
                                                                          1.8577149723180628;
                                                                          1.392409          ]

        taulist = [0.5, 0.6, 0.7]
        @test ForwardDynamics(thetalist, dthetalist, taulist, g, Ftip, Mlist, Glist, Slist) ≈ [ -0.9739290670855626;
                                                                                25.584667840340558 ;
                                                                                -32.91499212478149  ]

        @test hcat(EulerStep([0.1, 0.1, 0.1], [0.1, 0.2, 0.3], [2, 1.5, 1], 0.1)...) ≈ hcat([0.11, 0.12, 0.13], [0.3, 0.35, 0.4])

        @testset "inverse dynamics trajectory" begin
            thetastart = [0, 0, 0]
            thetaend = [π/2, π/2, π/2]
            Tf = 3
            N = 10
            method = 5

            traj = JointTrajectory(thetastart, thetaend, Tf, N, method)

            thetamat = copy(traj)
            dthetamat = zeros(N, 3)
            ddthetamat = zeros(N, 3)

            dt = Tf / (N - 1)

            for i = 1:size(traj, 1) - 1
                dthetamat[i + 1, :] = (thetamat[i + 1, :] - thetamat[i, :]) / dt
                ddthetamat[i + 1, :] = (dthetamat[i + 1, :] - dthetamat[i, :]) / dt
            end

            Ftipmat = ones(N, 6)

            taumat_actual = InverseDynamicsTrajectory(thetamat, dthetamat, ddthetamat, g, Ftipmat, Mlist, Glist, Slist)

            taumat_expected = [ 1.32297079e+01 -3.62621080e+01 -4.18134100e+00 ;
                                1.96974217e+01 -3.59188701e+01 -4.07919728e+00 ;
                                5.11979532e+01 -3.44705050e+01 -3.59488765e+00 ;
                                9.41368122e+01 -3.14099606e+01 -2.41622731e+00 ;
                                1.25417579e+02 -2.45283212e+01  5.76295281e-02 ;
                                1.24948454e+02 -1.85038921e+01  1.94898550e+00 ;
                                1.01525941e+02 -1.88375820e+01  2.07369432e+00 ;
                                7.68134579e+01 -2.23610568e+01  1.96172027e+00 ;
                                6.44365965e+01 -2.46704062e+01  2.07890001e+00 ;
                                6.72354909e+01 -2.47008371e+01  2.19474783e+00 ]

            @test taumat_actual ≈ taumat_expected
        end
        @testset "forward dynamics trajectory" begin
            taumat = [ 3.63 -6.58 -5.57 ; 3.74 -5.55 -5.5  ;
                       4.31 -0.68 -5.19 ; 5.18  5.63 -4.31 ;
                       5.85  8.17 -2.59 ; 5.78  2.79 -1.7  ;
                       4.99 -5.3  -1.19 ; 4.08 -9.41  0.07 ;
                       3.56 -10.1  0.97 ; 3.49 -9.41  1.23 ]

            Ftipmat = ones(size(taumat, 1), 6)

            dt = 0.1
            intRes = 8

            thetamat_expected = [ 0.1         0.1         0.1        ;
                                  0.10643138  0.2625997  -0.22664947 ;
                                  0.10197954  0.71581297 -1.22521632 ;
                                  0.0801044   1.33930884 -2.28074132 ;
                                  0.0282165   2.11957376 -3.07544297 ;
                                 -0.07123855  2.87726666 -3.83289684 ;
                                 -0.20136466  3.397858   -4.83821609 ;
                                 -0.32380092  3.73338535 -5.98695747 ;
                                 -0.41523262  3.85883317 -7.01130559 ;
                                 -0.4638099   3.63178793 -7.63190052 ]

            dthetamat_expected = [ 0.1          0.2          0.3        ;
                                   0.01212502   3.42975773  -7.74792602 ;
                                  -0.13052771   5.55997471 -11.22722784 ;
                                  -0.35521041   7.11775879  -9.18173035 ;
                                  -0.77358795   8.17307573  -7.05744594 ;
                                  -1.2350231    6.35907497  -8.99784746 ;
                                  -1.31426299   4.07685875 -11.18480509 ;
                                  -1.06794821   2.49227786 -11.69748583 ;
                                  -0.70264871  -0.55925705  -8.16067131 ;
                                  -0.1455669   -4.57149985  -3.43135114 ]

            thetamat_actual, dthetamat_actual = ForwardDynamicsTrajectory(thetalist, dthetalist, taumat, g, Ftipmat, Mlist, Glist, Slist, dt, intRes)

            @test thetamat_actual ≈ thetamat_expected
            @test dthetamat_actual ≈ dthetamat_expected
        end
    end
    @testset "chapter 9: trajectory generation" begin
        @test CubicTimeScaling(2, 0.6) == 0.21600000000000003
        @test QuinticTimeScaling(2, 0.6) == 0.16308
        @testset "joint trajectory" begin
            thetastart = [1, 0, 0, 1, 1, 0.2, 0,1]
            thetaend = [1.2, 0.5, 0.6, 1.1, 2, 2, 0.9, 1]
            Tf = 4
            N = 6
            method = 3

            expected = [     1     0      0      1     1    0.2      0 1 ;
                        1.0208 0.052 0.0624 1.0104 1.104 0.3872 0.0936 1 ;
                        1.0704 0.176 0.2112 1.0352 1.352 0.8336 0.3168 1 ;
                        1.1296 0.324 0.3888 1.0648 1.648 1.3664 0.5832 1 ;
                        1.1792 0.448 0.5376 1.0896 1.896 1.8128 0.8064 1 ;
                           1.2   0.5    0.6    1.1     2      2    0.9 1 ]

            @test JointTrajectory(thetastart, thetaend, Tf, N, method) ≈ expected
        end
        @testset "screw trajectory" begin
            Xstart = [1 0 0 1 ;
                      0 1 0 0 ;
                      0 0 1 1 ;
                      0 0 0 1 ]
            Xend = [0 0 1 0.1 ;
                    1 0 0   0 ;
                    0 1 0 4.1 ;
                    0 0 0   1 ]

            Tf = 5
            N = 4
            method = 3

            expected =  [[ 1 0 0 1 ;
                           0 1 0 0 ;
                           0 0 1 1 ;
                           0 0 0 1 ],
                         [ 0.904 -0.25   0.346  0.441 ;
                           0.346  0.904 -0.25   0.529 ;
                          -0.25   0.346  0.904  1.601 ;
                           0      0      0      1     ],
                         [ 0.346 -0.25   0.904 -0.117 ;
                           0.904  0.346 -0.25   0.473 ;
                          -0.25   0.904  0.346  3.274 ;
                           0      0      0      1     ],
                         [ 0 0 1 0.1 ;
                           1 0 0 0   ;
                           0 1 0 4.1 ;
                           0 0 0 1   ]]

            actual = ScrewTrajectory(Xstart, Xend, Tf, N, method)

            @test isapprox(actual, expected; rtol=1e-3)
        end
        @testset "cartesian trajectory" begin
            Xstart = [ 1 0 0 1 ;
                       0 1 0 0 ;
                       0 0 1 1 ;
                       0 0 0 1 ]
            Xend = [ 0 0 1 0.1 ;
                     1 0 0 0   ;
                     0 1 0 4.1 ;
                     0 0 0 1   ]

            Tf = 5
            N = 4
            method = 5

            expected = [[ 1 0 0 1 ;
                          0 1 0 0 ;
                          0 0 1 1 ;
                          0 0 0 1 ],
                        [ 0.937 -0.214  0.277 0.811 ;
                          0.277  0.937 -0.214     0 ;
                         -0.214  0.277  0.937 1.651 ;
                          0      0      0     1     ],
                        [ 0.277 -0.214  0.937 0.289 ;
                          0.937  0.277 -0.214     0 ;
                         -0.214  0.937  0.277 3.449 ;
                          0      0      0     1     ],
                        [ 0 0 1 0.1 ;
                          1 0 0 0   ;
                          0 1 0 4.1 ;
                          0 0 0 1   ]]

            actual = CartesianTrajectory(Xstart, Xend, Tf, N, method)

            @test isapprox(actual, expected; rtol=1e-3)
        end
    end
    @testset "chapter 11: robot control" begin
        thetalist = [0.1, 0.1, 0.1]
        dthetalist = [0.1, 0.2, 0.3]

        # Initialise robot description (Example with 3 links)
        g = [0, 0, -9.8]

        M01 = [1  0  0         0 ;
               0  1  0         0 ;
               0  0  1  0.089159 ;
               0  0  0         1 ]
        M12 = [0  0  1     0.28 ;
               0  1  0  0.13585 ;
              -1  0  0        0 ;
               0  0  0        1 ]
        M23 = [1  0  0        0 ;
               0  1  0  -0.1197 ;
               0  0  1    0.395 ;
               0  0  0        1 ]
        M34 = [1  0  0        0 ;
               0  1  0        0 ;
               0  0  1  0.14225 ;
               0  0  0        1 ]

        G1 = linalg.Diagonal([0.010267, 0.010267, 0.00666, 3.7, 3.7, 3.7])
        G2 = linalg.Diagonal([0.22689, 0.22689, 0.0151074, 8.393, 8.393, 8.393])
        G3 = linalg.Diagonal([0.0494433, 0.0494433, 0.004095, 2.275, 2.275, 2.275])

        Mlist = [M01, M12, M23, M34]
        Glist = [G1, G2, G3]

        Slist = [ 1  0  1      0  1      0 ;
                0  1  0 -0.089  0      0 ;
                0  1  0 -0.089  0  0.425 ]'

        # Create a trajectory to follow
        thetaend = [π / 2, π, 1.5 * π]

        Tf = 1
        dt = 0.05
        N = Int(Tf / dt)
        method = 5

        traj = JointTrajectory(thetalist, thetaend, Tf, N, method)

        thetamatd = copy(traj)
        dthetamatd = zeros(N, 3)
        ddthetamatd = zeros(N, 3)
        dt = Tf / (N - 1)

        for i = 1:size(traj, 1) - 1
            dthetamatd[i + 1, :] = (thetamatd[i + 1, :] - thetamatd[i, :]) / dt
            ddthetamatd[i + 1, :] = (dthetamatd[i + 1, :] - dthetamatd[i, :]) / dt
        end

        # Possibly wrong robot description (Example with 3 links)
        gtilde = [0.8, 0.2, -8.8]

        Mhat01 = [1 0 0   0;
                  0 1 0   0;
                  0 0 1 0.1;
                  0 0 0   1]
        Mhat12 = [0 0 1 0.3;
                  0 1 0 0.2;
                 -1 0 0   0;
                  0 0 0   1]
        Mhat23 = [1 0 0    0;
                  0 1 0 -0.2;
                  0 0 1  0.4;
                  0 0 0    1]
        Mhat34 = [1 0 0   0;
                  0 1 0   0;
                  0 0 1 0.2;
                  0 0 0   1]

        Ghat1 = linalg.Diagonal([0.1, 0.1, 0.1, 4, 4, 4])
        Ghat2 = linalg.Diagonal([0.3, 0.3, 0.1, 9, 9, 9])
        Ghat3 = linalg.Diagonal([0.1, 0.1, 0.1, 3, 3, 3])

        Gtildelist = [Ghat1, Ghat2, Ghat3]
        Mtildelist = [Mhat01, Mhat12, Mhat23, Mhat34]

        # Other required arguments
        Ftipmat = ones(size(traj, 1), 6)

        Kp = 20
        Ki = 10
        Kd = 18
        intRes = 8

        taumat_expected = [ -14.2640765   -54.06797429  -11.265448   ;
                             71.7014572   -17.58330542    3.86417108 ;
                            208.80692807    6.94442209    8.4352746  ;
                            269.9223766    14.44412677   11.24081382 ;
                            316.48343344    6.4020598    10.60970699 ;
                            327.82241593   -3.98984379   14.31752441 ;
                            248.33306921  -16.39336633   21.61795095 ;
                             93.7564835   -28.5575642    28.0092122  ;
                             13.12918592  -44.38407547   19.04258057 ;
                             56.35246455   -8.56189073    1.69770764 ;
                             32.68030349   39.77791901   -5.94800597 ;
                            -49.85502041   37.95496258  -15.10367806 ;
                           -104.48630504   24.8129766   -16.25667052 ;
                           -123.14920836   -3.62727714  -14.4680164  ;
                            -84.15220471  -25.40152665  -12.85439272 ;
                            -50.09890916  -33.73575763  -12.38441089 ;
                            -30.41466046  -34.03362524  -11.30974711 ;
                            -13.90701987  -26.40700004   -9.50026445 ;
                              7.93416317  -12.95474009   -6.58132646 ;
                             44.38627151    4.53534718   -2.32611269 ]

        thetamat_expected = [ 0.1028237  0.10738308 0.07715206 ;
                              0.10475535 0.11170712 0.05271794 ;
                              0.11930448 0.13796752 0.12315113 ;
                              0.1582068  0.21615938 0.27654789 ;
                              0.22464764 0.35314707 0.51117109 ;
                              0.31872944 0.54551689 0.81897456 ;
                              0.4377715  0.78011155 1.21554584 ;
                              0.57501633 1.03974866 1.71548388 ;
                              0.72137128 1.30551854 2.2916328  ;
                              0.87221934 1.56564069 2.84693996 ;
                              1.0257972  1.841874   3.32092634 ;
                              1.17352729 2.14460368 3.72255083 ;
                              1.30459739 2.44287634 4.03953656 ;
                              1.41155584 2.71108778 4.28079221 ;
                              1.49166143 2.92224332 4.45724092 ;
                              1.54707169 3.06503575 4.57357234 ;
                              1.579692   3.14295613 4.62867261 ;
                              1.59184568 3.1665363  4.62924012 ;
                              1.58827951 3.15316968 4.59115503 ;
                              1.57746305 3.12631855 4.54394792 ]

        taumat_actual, thetamat_actual = SimulateControl(thetalist, dthetalist, g, Ftipmat, Mlist, Glist,
                                                         Slist, thetamatd, dthetamatd, ddthetamatd, gtilde,
                                                         Mtildelist, Gtildelist, Kp, Ki, Kd, dt, intRes)

        @test taumat_actual ≈ taumat_expected
        @test thetamat_actual ≈ thetamat_expected
    end
end


using SummationByParts
include("nodecalc3.jl")

function nodecalc(sbp::TriSBP, isDG::Bool)
  vtx = [0.0 0; 1 0; 0 1]  # why is this 0 to 1 but the other method is -1 to 1
  r1 = vtx[1, :]
  r2 = vtx[2, :]
  r3 = vtx[3, :]

  T = zeros(2,2)
  T[:, 1] = r2 - r1
  T[:, 2] = r3 - r1

  if isDG
    coords = SummationByParts.SymCubatures.calcnodes(sbp.cub, vtx)
  else
    coords = calcnodes(sbp, vtx)
  end

  xi = zeros(coords)

  for i=1:size(coords,2)
    xi[:, i] = T\(coords[:, i] - r1)
  end

  return xi, coords
end

function nodecalc(sbp::TetSBP, isDG::Bool)

  vtx = [-1.0 -1.0 -1.0
        1.0 -1.0 -1.0
       -1.0 1.0 -1.0
       -1.0 -1.0 1.0]

  r1 = vtx[1, :]
  r2 = vtx[2, :]
  r3 = vtx[3, :]
  r4 = vtx[4, :]
  T = zeros(3,3)
  T[:, 1] = r2 - r1
  T[:, 2] = r3 - r1
  T[:, 3] = r4 - r1

  if isDG
    coords = SummationByParts.SymCubatures.calcnodes(sbp.cub, vtx)
  else
    coords = calcnodes(sbp, vtx)
  end

  xi = zeros(coords)

  for i=1:size(coords,2)
    xi[:, i] = T\(coords[:, i] - r1)
  end

  return xi, coords
end

function minNodeDist(sbp, isDG::Bool)
# get the minimum distance between nodes on a reference element of degree p

    xi, coords = nodecalc(sbp, isDG)
    min_dist = typemax(Float64)
#    println("coords = ", coords)
    for i=1:size(coords, 2)
      for j=(i+1):size(coords, 2)
	# calculate distance between node i and node j
	dist_j = norm(coords[:, i] - coords[:, j])

	if dist_j < min_dist
	  min_dist = dist_j
	end

      end  # end loop j
    end  # end loop i

#    println("for p=$p elements, min node distance = ", min_dist)

  return min_dist
end

# sbp = TriSBP{Float64}(degree=1, reorder=false, internal=false)
# sbp = getTriSBPWithDiagE(degree=4)
# xi, coords = nodecalc(sbp, true)
# printCaseStatement(xi)
# printBaryCoords(xi, coords)
# writedlm("coordsout.dat", coords, ' ')
#minNodeDist(2)


using ModelingToolkit, MethodOfLines, LinearAlgebra, OrdinaryDiffEq
using ModelingToolkit: operation, istree, arguments
using DomainSets
using NonlinearSolve

# # Define some variables
@testset "Heat Equation 1D 2 variables" begin
       @parameters t x
       @variables u(..) v(..)
       Dt = Differential(t)
       Dx = Differential(x)
       Dxx = Differential(x)^2
       eqs  = [Dt(u(t,x)) ~ Dxx(u(t,x)),
              Dt(v(t,x)) ~ Dxx(v(t,x))]
       bcs = [u(0,x) ~ - x * (x-1) * sin(x),
              v(0,x) ~ - x * (x-1) * sin(x),
              u(t,0) ~ 0.0, u(t,1) ~ 0.0,
              v(t,0) ~ 0.0, v(t,1) ~ 0.0]

       domains = [t ∈ Interval(0.0,1.0),
              x ∈ Interval(0.0,1.0)]

       @named pdesys = PDESystem(eqs,bcs,domains,[t,x],[u(t,x),v(t,x)])
       discretization = MOLFiniteDifference([x=>0.1],t;grid_align=edge_align)
       prob = discretize(pdesys,discretization) # This gives an ODEProblem since it's time-dependent
       sol = solve(prob,Tsit5())
end

@testset "Heat Equation 2D 1 variable" begin
       @parameters t x y
       @variables u(..)
       Dxx = Differential(x)^2
       Dyy = Differential(y)^2
       Dt = Differential(t)
       t_min= 0.
       t_max = 2.0
       x_min = 0.
       x_max = 2.
       y_min = 0.
       y_max = 2.

       # 3D PDE
       eq  = Dt(u(t,x,y)) ~ Dxx(u(t,x,y)) + Dyy(u(t,x,y))

       analytic_sol_func(t,x,y) = exp(x+y)*cos(x+y+4t)
       # Initial and boundary conditions
       bcs = [u(t_min,x,y) ~ analytic_sol_func(t_min,x,y),
              u(t,x_min,y) ~ analytic_sol_func(t,x_min,y),
              u(t,x_max,y) ~ analytic_sol_func(t,x_max,y),
              u(t,x,y_min) ~ analytic_sol_func(t,x,y_min),
              u(t,x,y_max) ~ analytic_sol_func(t,x,y_max)]

       # Space and time domains
       domains = [t ∈ Interval(t_min,t_max),
              x ∈ Interval(x_min,x_max),
              y ∈ Interval(y_min,y_max)]
       @named pdesys = PDESystem([eq],bcs,domains,[t,x,y],[u(t,x,y)])
       # Method of lines discretization
       dx = 0.1; dy = 0.2
       discretization = MOLFiniteDifference([x=>dx,y=>dy],t)
       prob = ModelingToolkit.discretize(pdesys,discretization)
       sol = solve(prob,Tsit5())
end

# Diffusion in a sphere
@testset "Spherical Diffusion" begin

       @parameters t r
       @variables u(..)
       Dt = Differential(t)
       Dr = Differential(r)
       Drr = Dr^2
       eq  = Dt(u(t,r)) ~ 1/r^2 * Dr(r^2 * Dr(u(t,r)))
       bcs = [u(0,r) ~ - r * (r-1) * sin(r),
              Dr(u(t,0)) ~ 0.0, u(t,1) ~ sin(1)]

       domains = [t ∈ Interval(0.0,1.0),
              r ∈ Interval(0.0,1.0)]

       @named pdesys = PDESystem(eq,bcs,domains,[t,r],[u(t,r)])
       discretization = MOLFiniteDifference([r=>0.1],t)
       prob = discretize(pdesys,discretization) # This gives an ODEProblem since it's time-dependent
       sol = solve(prob,Tsit5())
end

@testset "RHS = 0" begin
       @parameters x, t
       @variables u(..)
       Dt = Differential(t)
       Dx = Differential(x)
       Dx2 = Differential(x)^2
       Dx3 = Differential(x)^3
       Dx4 = Differential(x)^4

       α = 1
       β = 4
       γ = 1
       eq = Dt(u(x,t)) + u(x,t)*Dx(u(x,t)) + α*Dx2(u(x,t)) + β*Dx3(u(x,t)) + γ*Dx4(u(x,t)) ~ 0

       du(x,t;z = -x/2+t) = 15/2*(tanh(z) + 1)*(3*tanh(z) - 1)*sech(z)^2

       bcs = [u(x,0) ~ x^2,
              Dx(u(-10,t)) ~ du(-10,t),
              Dx(u(10,t)) ~ du(10,t)]

       # Space and time domains
       domains = [x ∈ Interval(-10.0,10.0),
              t ∈ Interval(0.0,0.5)]
       # Discretization
       dx = 0.4; dt = 0.2

       discretization = MOLFiniteDifference([x=>dx],t;approx_order=4)
       @named pdesys = PDESystem(eq,bcs,domains,[x,t],[u(x,t)])
       prob = discretize(pdesys,discretization)
end

@testset "Wave Equation" begin
       # Parameters, variables, and derivatives
       @parameters t x
       @variables u(..)
       Dt = Differential(t)             # Required in ICs and equation
       Dtt = Differential(t)^2           # required in equation
       Dxxxx = Differential(x)^4       # required in equation
       Dxx = Differential(x)^2         # required in BCs
       # some parameters
       EI  = 291.6667;
       m = 1.3850;
       c = 1.;
       p = 0.01;
       L = 2.0;

       # 1D PDE and boundaru conditions
       eq  = m*Dtt(u(t,x)) + c*Dt(u(t,x)) + EI*Dxxxx(u(t,x)) ~0

       ic_bc = [u(0,x) ~ (p*x*(x^3 + L^3 -2*L*x^2)/(24*EI)), #for all 0 < u < L
              Dt(u(0,x)) ~ 0.,        # for all 0 < u < L
              u(t,0) ~ 0.,            # for all t > 0,, displacement zero at u=0
              u(t,2) ~ 0.,            # for all t > 0,, displacement zero at u=L
              Dxx(u(t,0)) ~ 0.,       # for all t > 0,, curvature zero at u=0
              Dxx(u(t,2)) ~ 0.]       # for all t > 0,, curvature zero at u=L

       # Space and time domains
       domains = [t ∈ Interval(0.0,2.0),
              x ∈ Interval(0.0, L)]

       dt = 0.1   # dt related to saving the data.. not actual dt
       # PDE sustem
       @named pdesys = PDESystem(eq,ic_bc,domains,[t,x],[u(t,x)])

       # Method of lines discretization
       dx = 0.1
       order = 2
       discretization = MOLFiniteDifference([x=>dx, t=>dt], approx_order=order)

       # Convert the PDE problem into an ODE problem
       prob = discretize(pdesys,discretization)

       # Solve the ODE problem
       sol = NonlinearSolve.solve(prob, NewtonRaphson())
end

@testset "Rearranged Robin" begin
       @parameters x t

	@variables c(..)

	∂t  = Differential(t)
	∂x  = Differential(x)
	∂²x = Differential(x) ^ 2

       D₀ = 1.
       R = 0.5
       cₑ = 2.5
       ℓ = 2.0
       α = 1/3

	diff_eq = ∂t(c(x, t)) ~ ∂x((D₀ + α * c(x, t)) * ∂x(c(x, t)))

	bcs = [c(x, 0) 	 ~ 0.0,  	       # initial condition
                  R * (D₀ + α * c(0, t)) * ∂x(c(0, t)) ~ c(0, t) - cₑ, # Robin
                  ∂x(c(ℓ, t)) ~ 0.0]   # no flux

       domains = [t ∈ Interval(0.0,10.0),
                  x ∈ Interval(0.0, ℓ)]

       @named pdesys = PDESystem(diff_eq,bcs,domains,[t,x],[c(x,t)])

       # Method of lines discretization
       dx = 0.1
       order = 2
       discretization = MOLFiniteDifference([x=>dx],t,approx_order=order)

       # Convert the PDE problem into an ODE problem
       prob = discretize(pdesys,discretization)

       sol = solve(prob,Rodas4())
end

@testset "Array u" begin
	# Dependencies
	N = 6 # number of dependent variables

	# Variables, parameters, and derivatives
	@parameters x
	@variables u[1:N](..)
	Dx = Differential(x)
	Dxx = Differential(x)^2

	# Domain edges
	x_min= 0.
	x_max = 1.

	# Discretization parameters
	dx = 0.1
	order = 2

       #u = collect(u)

	# Equations
	eqs  = Vector{ModelingToolkit.Equation}(undef, N)
	for i = 1:N
		eqs[i] = Dxx(u[i](x)) ~ u[i](x)
	end

	# Initial and boundary conditions
	bcs = Vector{ModelingToolkit.Equation}(undef, 2*N)
	for i = 1:N
		bcs[i] = Dx(u[i](x_min)) ~ 0.
	end

	for i = 1:N
		bcs[i+N] = u[i](x_max) ~ rand()
	end

	# Space and time domains
	domains = [x ∈ Interval(x_min, x_max)]

	# PDE system
	@named pdesys = PDESystem(eqs, bcs, domains, [x], collect([u[i](x) for i = 1:N]))

	# Method of lines discretization
	discretization = MOLFiniteDifference([x=>dx], nothing, approx_order=order)
	prob = ModelingToolkit.discretize(pdesys,discretization)

	# # Solution of the ODE system
	sol = NonlinearSolve.solve(prob,NewtonRaphson())
end

@testset "2D variable connected to 1D variable at boundary #33" begin
       @parameters t x r
       @variables u(..) v(..)
       Dt = Differential(t)
       Dx = Differential(x)
       Dxx = Differential(x)^2
       Dr = Differential(r)
       Drr = Differential(r)^2

       s = u(t,x) + v(t,x,1)

       eqs  = [Dt(u(t,x)) ~ Dxx(u(t,x)) + s,
              Dt(v(t,x,r)) ~ Drr(v(t,x,r))]
       bcs = [u(0,x) ~ 1,
       v(0,x,r) ~ 1,
       Dx(u(t,0)) ~ 0.0, Dx(u(t,1)) ~ 0.0,
       Dr(v(t,x,0)) ~ 0.0, Dr(v(t,x,1)) ~ s]

       domains = [t ∈ Interval(0.0,1.0),
                  x ∈ Interval(0.0,1.0),
                  r ∈ Interval(0.0,1.0)]

       @named pdesys = PDESystem(eqs,bcs,domains,[t,x,r],[u(t,x),v(t,x,r)])

       # Method of lines discretization
       dx = 0.1
       dr = 0.1
       order = 2
       discretization = MOLFiniteDifference([x=>dx,r=>dr], t, approx_order=order)

       # Convert the PDE problem into an ODE problem
       prob = discretize(pdesys,discretization)

       sol = solve(prob, Tsit5())
end

@testset "Testing discretization of varied systems" begin
	@parameters x t

	@variables c(..)

	∂t  = Differential(t)
	∂x  = Differential(x)
	∂²x = Differential(x) ^ 2

	D₀ = 1.5
	α = 0.15
	χ = 1.2
	R = 0.1
	cₑ = 2.0
	ℓ = 1.0
	Δx = 0.1

	bcs = [
              # initial condition
		c(x, 0) ~ 0.0,
		# Robin BC
		∂x(c(0.0, t)) / (1 + exp(α * (c(0.0, t) - χ))) * R * D₀ + cₑ - c(0.0, t) ~ 0.0,
		# no flux BC
		∂x(c(ℓ, t)) ~ 0.0]

       # define space-time plane
       domains = [x ∈ Interval(0.0, ℓ), t ∈ Interval(0.0, 5.0)]

       @testset "Test 01: ∂t(c(x, t)) ~ ∂x(D * ∂x(c(x, t)))" begin
              D = D₀ / (1.0 + exp(α * (c(x, t) - χ)))
              diff_eq = ∂t(c(x, t)) ~ ∂x(D * ∂x(c(x, t)))
              @named pdesys = PDESystem(diff_eq, bcs, domains, [x, t], [c(x, t)]);
              discretization = MOLFiniteDifference([x=>Δx], t)
       end

       @testset "Test 02: ∂t(c(x, t)) ~ ∂x(D * ∂x(c(x, t)))" begin
              D = 1.0 / (1.0 + exp(α * (c(x, t) - χ)))
              diff_eq = ∂t(c(x, t)) ~ ∂x(D * ∂x(c(x, t)))
              @named pdesys = PDESystem(diff_eq, bcs, domains, [x, t], [c(x, t)]);
              discretization = MOLFiniteDifference([x=>Δx], t)
       end

       @testset "Test 03: ∂t(c(x, t)) ~ ∂x(1.0 / (1.0/D₀ + exp(α * (c(x, t) - χ))/D₀) * ∂x(c(x, t)))" begin
              diff_eq = ∂t(c(x, t)) ~ ∂x(1.0 / (1.0/D₀ + exp(α * (c(x, t) - χ))/D₀) * ∂x(c(x, t)))
              @named pdesys = PDESystem(diff_eq, bcs, domains, [x, t], [c(x, t)]);
              discretization = MOLFiniteDifference([x=>Δx], t)
       end

       @testset "Test 04: ∂t(c(x, t)) ~ ∂x(D₀ / (1.0 + exp(α * (c(x, t) - χ))) * ∂x(c(x, t)))" begin
              diff_eq = ∂t(c(x, t)) ~ ∂x(D₀ / (1.0 + exp(α * (c(x, t) - χ))) * ∂x(c(x, t)))
              @named pdesys = PDESystem(diff_eq, bcs, domains, [x, t], [c(x, t)]);
              discretization = MOLFiniteDifference([x=>Δx], t)
       end

       @testset "Test 05: ∂t(c(x, t)) ~ ∂x(1/x * ∂x(c(x, t)))" begin
              diff_eq = ∂t(c(x, t)) ~ ∂x(1/x * ∂x(c(x, t)))
              @named pdesys = PDESystem(diff_eq, bcs, domains, [x, t], [c(x, t)]);
              discretization = MOLFiniteDifference([x=>Δx], t)
       end

       @testset "Test 06: ∂t(c(x, t)) ~ ∂x(x*∂x(c(x, t)))/c(x,t)" begin
              diff_eq = ∂t(c(x, t)) ~ ∂x(x*∂x(c(x, t)))/c(x,t)
              @named pdesys = PDESystem(diff_eq, bcs, domains, [x, t], [c(x, t)]);
              discretization = MOLFiniteDifference([x=>Δx], t)
       end

       @testset "Test 07: ∂t(c(x, t)) ~ ∂x(1/(1+c(x,t)) ∂x(c(x, t)))" begin
              diff_eq = ∂t(c(x, t)) ~ ∂x(1/(1+c(x,t)) * ∂x(c(x, t)))
              @named pdesys = PDESystem(diff_eq, bcs, domains, [x, t], [c(x, t)]);
              discretization = MOLFiniteDifference([x=>Δx], t)
       end

       @testset "Test 08: ∂t(c(x, t)) ~ c(x, t) * ∂x(c(x,t) * ∂x(c(x, t)))" begin
              diff_eq = ∂t(c(x, t)) ~ c(x, t)*∂x(c(x,t) * ∂x(c(x, t)))
              @named pdesys = PDESystem(diff_eq, bcs, domains, [x, t], [c(x, t)]);
              discretization = MOLFiniteDifference([x=>Δx], t)
       end

       @testset "Test 09: ∂t(c(x, t)) ~ c(x, t) * ∂x(c(x,t) * ∂x(c(x, t)))/(1+c(x,t))" begin
              diff_eq = c(x, t) * ∂x(c(x,t) * ∂x(c(x, t)))/(1+c(x,t))
              @named pdesys = PDESystem(diff_eq, bcs, domains, [x, t], [c(x, t)]);
              discretization = MOLFiniteDifference([x=>Δx], t)
       end

       @testset "Test 10: ∂t(c(x, t)) ~ c(x, t) * ∂x(c(x,t) * ∂x(c(x, t)))/(1+c(x,t))" begin
              diff_eq = c(x, t) * ∂x(c(x,t)^(-1) * ∂x(c(x, t)))
              @named pdesys = PDESystem(diff_eq, bcs, domains, [x, t], [c(x, t)]);
              discretization = MOLFiniteDifference([x=>Δx], t)
       end

       @testset "Test 11: ∂t(c(x, t)) ~ ∂x(1/(1+c(x,t)^2) ∂x(c(x, t)))" begin
              diff_eq = ∂t(c(x, t)) ~ ∂x(1/(1+c(x,t)^2) * ∂x(c(x, t)))
              @named pdesys = PDESystem(diff_eq, bcs, domains, [x, t], [c(x, t)]);
              discretization = MOLFiniteDifference([x=>Δx], t)
       end

end


using CSV, GLM, Plots, Random; pyplot()
Random.seed!(0)

df = CSV.read("../data/weightHeight.csv", copycols = true)
n = size(df)[1]
df[shuffle(1:n),:] = df
df[[10,40,60,130,140,175,190,200],:Sex] .= "O1"
df[[9,44,63,132,138,172,192,199],:Sex] .= "O2"

model = lm(@formula(Height ~ Weight + Sex), df,
        contrasts=Dict(:Sex=>DummyCoding(base="F",levels=["M","O1","O2","F"])))
b0, b1, b2, b3, b4  = coef(model)
pred(weight,sex) = b0+b1*weight+b2*(sex=="M")+b3*(sex=="O1")+b3*(sex=="O2")
println(model)

males = df[df.Sex .== "M",:]
females = df[df.Sex .== "F",:]
other = df[(df.Sex .!= "M") .& (df.Sex .!= "F"),:]

xlim = [minimum(df.Weight), maximum(df.Weight)]
scatter(males.Weight, males.Height, c=:blue, msw=0, label="Males")
plot!(xlim, pred.(xlim,"M"), c=:blue, label="Male model")

scatter!(females.Weight, females.Height, c=:red, msw=0, label="Females")
plot!(xlim, pred.(xlim,"F"),
	c=:red, label="Female model", xlims=(xlim),
	xlabel="Weight (kg)", ylabel="Height (cm)", legend=:topleft)

scatter!(other.Weight, other.Height, c=:green, msw=0, label="Other")


@testset "polynomials" begin
    @test DiscretisedFluidQueues.gauss_lobatto_points(2.0,2.5,1)≈[2.25]
    @test DiscretisedFluidQueues.gauss_lobatto_points(2.0,3.0,2)≈[2.0;3.0]
    @test DiscretisedFluidQueues.gauss_lobatto_points(2.0,4.0,3)≈[2.0;3.0;4.0]
    nnodes = 5
    nodes = DiscretisedFluidQueues.gauss_lobatto_points(-1.0,1.0,nnodes)
    for n in 1:nnodes
        en = zeros(nnodes)
        en[n] = 1.0
        @test DiscretisedFluidQueues.lagrange_polynomials(nodes,nodes[n])≈en atol=sqrt(eps())
    end
    @test DiscretisedFluidQueues.lagrange_polynomials([1.0],1.0)≈[1.0] atol=sqrt(eps())
    @test DiscretisedFluidQueues.gauss_lobatto_weights(-1.0,1.0,1)≈[2.0]
    nnodes = 3
    nodes = DiscretisedFluidQueues.gauss_lobatto_points(-1.0,1.0,nnodes)
    @test DiscretisedFluidQueues.gauss_lobatto_weights(-1.0,1.0,nnodes)≈[1.0,4.0,1.0]./3
    x4 = x->x^4
    x4_interp = DiscretisedFluidQueues.lagrange_interpolation(x4,-1.0,1.0,5)
    x = collect(-3:0.02:3)
    @test x4_interp.(x)≈x4.(x)
    x4_quad = DiscretisedFluidQueues.gauss_lobatto_quadrature(x4,-1.0,1.0,4)
    @test x4_quad≈2/5

    x4_interp_p0 = DiscretisedFluidQueues.lagrange_interpolation(x->1.0-x4(x),-1.0,1.0,1)
    x = collect(-3:0.02:3)
    @test x4_interp_p0.(x)≈ones(size(x))
    x4_quad_p0 = DiscretisedFluidQueues.gauss_lobatto_quadrature(x->1.0-x4(x),-1.0,1.0,1)
    @test x4_quad_p0≈2.0
end


# Methods from Moukalled et al. FVM - OpenFOAM, Matlab

mutable struct Cell
    faceIndices::Vector{Int64}
    pointIndices::Vector{Int64}
end

######################### Mesh/Cell Geometry ###########################
function triangleCentroid(points)
    center = [ 0.0, 0.0, 0.0 ]
    nPts = size(points, 1)
    for pt in 1:nPts
        center = center .+ points[pt]
    end
    center /= nPts

    return center
end

#Alternate name for same function
geometricCenter = triangleCentroid

function triangleArea(points::Array{Array{Float64, 1}})
    side1 = points[2] .- points[1]
    side2 = points[3] .- points[1]
    fAVec = cross(side1, side2) ./ 2
    return fAVec
end

#=
    Calculates face area vector and centroid from the points that make up the face
    Points must be ordered sequentially

    How it works:
        1. Splits face into subtriangles
        2. Area and centroid is computed for each subtriangle
        3. Areas vectors are summed and returned
        4. The centroid returned is obtained from an area-weighted of sum of the subtriangle centroids
=#
function faceAreaCentroid(points::Array{Array{Float64, 1}})
    gC = geometricCenter(points)
    nPts = size(points, 1)

    fAVec = [ 0.0, 0.0, 0.0 ]
    centroid = [ 0.0, 0.0, 0.0 ]

    for i in 1:nPts
        if i < nPts
            subTriPts = [ gC, points[i], points[i+1] ]
        else
            subTriPts = [ gC, points[i], points[1] ]
        end

        triCentroid = triangleCentroid(subTriPts)
        subFAVec = triangleArea(subTriPts)

        fAVec += subFAVec
        centroid += triCentroid .* mag(subFAVec)
    end

    centroid /= mag(fAVec)

    return fAVec, centroid
end

#=
    Calculates cell volume (scalar) and centroid (vector) from the points and face area vectors (fAVecs) that make up the cell
        fAVecs can be computed using the faceAreaCentroids function

    How it works:
        1. Splits cell into polygonal pyramids, each incorporating a single face and the geometric center of the cell
        2. Computes volume and centroid of each sub-pyramid
        3. Resulting volume is sum of sub-pyramid volumes, centroid is the volume-weighted sum of sub-pyramid centroids
=#
function cellVolCentroid(points, fAVecs, faceCentroids)
    gC = geometricCenter(points)
    nFaces = size(fAVecs,1)

    vol = 0.0
    centroid = [ 0.0, 0.0, 0.0 ]

    for f in 1:nFaces
        cellCenterVec = faceCentroids[f] .- gC
        subPyrVol = abs(sum(fAVecs[f] .* cellCenterVec) / 3)
        subPyrCentroid = 0.75.*faceCentroids[f] .+ 0.25.*gC

        vol += subPyrVol
        centroid += subPyrCentroid .* subPyrVol
    end

    centroid /= vol

    return vol, centroid
end

# For each face, calculates the vectors from its owner and neighbour cell centers to its own center
function cellCentroidToFaceVec(faceCentroids::Array{Array{Float64, 1}}, cellCentroids::Array{Array{Float64, 1}})
    #TODO: Understand ordering to make sure we're accepting faces in the order we need to pass them back in??
    nFaces = size(faceCentroids, 1)

    cellToFaceVec = zeros(nFaces, 3)

    for f in 1:nFaces
        cellToFaceVec[f,:] = faceCentroids[f] - cellCentroids
    end

    return cellToFaceVec
end

######################### Utility Functions ###########################
# Returns basic info about the mesh: nCells, nFaces, nBoundaries, nBdryFaces
function unstructuredMeshInfo(mesh::Mesh)
    nCells = size(mesh.cells, 1)
    nFaces = size(mesh.faces, 1)
    nBoundaries = size(mesh.boundaryFaces, 1)

    # Count boundary faces
    nBdryFaces = 0
    for bdry in 1:nBoundaries
        nBdryFaces += size(mesh.boundaryFaces[bdry], 1)
    end

    return nCells, nFaces, nBoundaries, nBdryFaces
end

# Checks whether a string represents a number (using a regular expression)
function isNumber(str)
    re = r"^[+-]?([0-9]+([.][0-9]*)?|[.][0-9]+)$"
    return occursin(re, str)
end

######################### Parse OpenFOAM Meshes ###########################
#=
    In the points, faces, owner, and neighbour files, the beginning of useful information in the file looks like this:

        ...
        object      faces;
        }
        // * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * //

        5135
        (
        4(79 1516 1380 37)
        4(1381 1422 1249 39)
        ...

    The 5135 indicates how many lines of information follow it.
    This function would return 5135, and the line number of the first piece of information in the file ((line number of 5135) + 2)
=#
function OFFile_FindNItems(fileLines)
    nLines = size(fileLines, 1)
    lineCounter = 1
    itemCount = 0
    startLine = 0
    for line in fileLines
        if isNumber(strip(line))
            itemCount = parse(Int64, line)
            startLine = lineCounter+2
            break
        end
        lineCounter += 1
    end
    return startLine, itemCount
end

#=
    Takes the file path to an OpenFOAM mesh 'points' file, returns a 2D array of point coordinates:
            X   Y   Z
        P1  x1  y1  z1
        P2  x2  y2  z2
        ...
=#
function readOFPointsFile(filePath)
    f = open(filePath)
    pointsLines = readlines(f)
    startLine, pCount = OFFile_FindNItems(pointsLines)

    points = zeros(pCount, 3)
    for line in startLine:(startLine+pCount-1)
        pLine = pointsLines[line]

        bracketsRemoved = pLine[2:end-1]
        coords = split(bracketsRemoved)
        for d in 1:3
            points[line+1-startLine, d] = parse(Float64, coords[d])
        end
    end

    return points
end

#=
    Takes the file path to an OpenFOAM mesh 'faces' file, returns a 1D array of 1D arrays.
    Each entry in the upper level array represents a single face
    Each face-array contains the indices of the points that make up that face
    Note that all indices from the file are incremented by 1, to switch from OpenFOAM's 0-based indexing to Julia's 1-based indexing

    Example return value:
        [ [face1_pt1Index, face1_pt2Index, face1_pt3Index], [face2_pt1Index, face2_pt2Index, face2_pt3Index, face2_pt4Index], ... ]
=#
function readOFFacesFile(filePath)
    f = open(filePath)
    facesLines = readlines(f)
    startLine, fCount = OFFile_FindNItems(facesLines)

    faces = []
    for line in startLine:(startLine+fCount-1)
        fLine = facesLines[line]
        facePts = []

        bracketIndex = findfirst("(", fLine)[1]
        bracketsRemoved = fLine[bracketIndex+1:end-1]
        ptNumbers = split(bracketsRemoved)
        nPts = length(ptNumbers)
        for p in 1:nPts
            # Add one to convert to 1-based indexing
            push!(facePts, parse(Int64, ptNumbers[p])+1)
        end
        push!(faces, facePts)
    end

    return faces
end

#=
    Takes the file path to an OpenFOAM 'owner' file, returns a 1D array of owner cell indices
    Example return value:
        [ face1_ownerCellIndex, face2_ownerCellIndex, ...  ]
=#
function readOFOwnerFile(filePath)
    f = open(filePath)
    ownerLines = readlines(f)
    startLine, oCount = OFFile_FindNItems(ownerLines)

    faceOwnerCells = zeros(Int64, oCount)
    for line in startLine:(startLine+oCount-1)
        # Add one to convert to 1-based indexing
        faceOwnerCells[line-startLine+1] = parse(Int64, ownerLines[line])+1
    end

    return faceOwnerCells
end

#=
    Takes the file path to an OpenFOAM 'neighbour' file, returns a 1D array of neighbour cell indices
    Note that all indices from the file are incremented by 1, to switch from OpenFOAM's 0-based indexing to Julia's 1-based indexing

    Example return value:
        [ face1_neighbourCellIndex, face2_neighbourCellIndex, ...  ]
=#
function readOFNeighbourFile(filePath)
    f = open(filePath)
    neighbourLines = readlines(f)
    startLine, nCount = OFFile_FindNItems(neighbourLines)

    faceNeighbourCells = zeros(Int64, nCount)
    for line in startLine:(startLine+nCount-1)
        # Add one to convert to 1-based indexing
        faceNeighbourCells[line-startLine+1] = parse(Int64, neighbourLines[line])+1
    end

    return faceNeighbourCells
end

#=
    Returns index of first line containing str
    Otherwise returns -1
=#
function findInLines(str, lines, startLine)
    nLines = size(lines, 1)
    for i in startLine:nLines
        if occursin(str, lines[i])
            return i
        end
    end
    return -1
end

#=
    Takes the file path to an OpenFOAM 'boundary' file, returns 3 arrays:
    boundaryNames:          Array of boundary names (strings)
    boundaryNumFaces:       Array containing the number of faces in each boundary
    boundaryStartFaces:     Array containing the index of the first face in the boundary (boundaries always occupy a continuous sequence of face indices, and are usually numbered last)
=#
function readOFBoundaryFile(filePath)
    f = open(filePath)
    bLines = readlines(f)
    startLine, bCount = OFFile_FindNItems(bLines)
    nLines = size(bLines, 1)

    boundaryNames = Array{String, 1}(undef, bCount)
    boundaryStartFaces = Array{Int64, 1}(undef, bCount)
    boundaryNumFaces = Array{Int64, 1}(undef, bCount)
    for i in 1:bCount
        bNameLine = findInLines("{", bLines, startLine)-1
        boundaryNames[i] = strip(bLines[bNameLine])

        bNFacesLine = findInLines("nFaces", bLines, startLine)
        boundaryNumFaces[i] = parse(Int64, split(bLines[bNFacesLine])[2][1:end-1])

        bStartFaceLine = findInLines("startFace", bLines, startLine)
        boundaryStartFaces[i] = parse(Int64, split(bLines[bStartFaceLine])[2][1:end-1])+1

        startLine = findInLines("}", bLines, startLine)+1
    end

    return boundaryNames, boundaryNumFaces, boundaryStartFaces
end

#=
    Input: Path to an OpenFOAM mesh FOLDER.
    Output: Calls readOFPoints/Faces/Owner/NeighbourFile and returns all of their results (reads all OpenFOAM's mesh file data into arrays)
=#
function readOpenFOAMMesh(polyMeshPath)
    println("Reading mesh: $polyMeshPath")
    pointsFilePath = "$polyMeshPath/points"
    points = readOFPointsFile(pointsFilePath)
    facesFilePath = "$polyMeshPath/faces"
    faces = readOFFacesFile(facesFilePath)
    ownerFilePath = "$polyMeshPath/owner"
    owner = readOFOwnerFile(ownerFilePath)
    neighbourFilePath = "$polyMeshPath/neighbour"
    neighbour = readOFNeighbourFile(neighbourFilePath)
    boundaryFilePath = "$polyMeshPath/boundary"
    boundaryNames, boundaryNumFaces, boundaryStartFaces = readOFBoundaryFile(boundaryFilePath)

    return points, faces, owner, neighbour, boundaryNames, boundaryNumFaces, boundaryStartFaces
end

#=
    Function used to find all the points in each cell.
    Not required for CFD, which is face-based, but required for cell-based .vtk file output.

    Input: Path to an OpenFOAM mesh FOLDER.
    Returns:
        points:         Array of points obtained from readOFPoints()
        cellPtIndices:  Array of arrays, where each subarrary represents a cell, and each entry in a cell's array is the index of one of its points
=#
function OpenFOAMMesh_findCellPts(polyMeshPath)
    pointLocations, pointIndicesByFace, faceOwnerCellIndices, faceNeighborCellIndices, boundaryNames, boundaryNumFaces, boundaryStartFaces = readOpenFOAMMesh(polyMeshPath)
    nCells = maximum(faceOwnerCellIndices)
    nFaces = size(pointIndicesByFace, 1)
    nBoundaries = size(boundaryNames, 1)

    # Output
    cells = Array{Cell, 1}(undef, nCells)
    for i in eachindex(cells)
        cells[i] = Cell(Vector{Int64}(undef, 0), Vector{Int64}(undef, 0))
    end

    ### Populate the cells array ###
    function addCellFaceIndices(adjacentCellList)
        nFaces = length(adjacentCellList)
        for f in 1:nFaces
            ownerCellIndex = adjacentCellList[f]
            push!(cells[ownerCellIndex].faceIndices, f)
        end
    end

    # Start by a just adding face indices, want to order the points appropriately before adding them
    addCellFaceIndices(faceOwnerCellIndices)
    addCellFaceIndices(faceNeighborCellIndices)

    function addAllNewPoints!(cell::Cell, faceIndex)
        for pointIndex in pointIndicesByFace[faceIndex]
            if !any(x->x==pointIndex, cell.pointIndices)
                push!(cell.pointIndices, pointIndex)
            end
        end
    end

    function disjoint(f1Points, f2Points)
        allPoints = vcat(f1Points, f2Points)
        pointSet = Set(allPoints)
        return length(allPoints) == length(pointSet)
    end

    function intersection(points1, points2)
        return intersect(Set(points1), Set(points2))
    end

    function index(element, array)
        for i in eachindex(array)
            if element == array[i]
                return i
            end
        end
        return -1
    end

    function addEdges!(cellPoints, facePoints, endFacePoints, endFacePointsSet, unusedFaces, pointOffset=4)
        oppositeFaceIndex = -1
        for i in 1:length(unusedFaces)
            fi = unusedFaces[i]
            fiPoints = pointIndicesByFace[fi]
            commonPoints = intersection(fiPoints, facePoints)

            if length(commonPoints) == 2
                # These two points form an edge connecting faces one and two
                points = collect(commonPoints)

                if points[1] in endFacePointsSet
                    p1 = points[1]
                    p2 = points[2]
                else
                    p1 = points[2]
                    p2 = points[1]
                end

                position = index(p1, endFacePoints) # Find position over the point on the face one side
                if position == -1
                    throw(ErrorException("Point $p1 not found in $endFacePoints"))
                end
                cellPoints[ position + pointOffset ] = p2 # The point on the face two side goes in the matching spot
            else
                oppositeFaceIndex = i # This is the face opposite to the one containing facePoints
            end
        end
        return oppositeFaceIndex
    end

    function populatePointIndices_Tet!(cell::Cell)
        # For a tetrahedron, the order of points is unimportant
        # And all points from an arbitrary face
        addAllNewPoints!(cell, cell.faceIndices[1])

        # Choose another face, it will contain the fourth point we need to complete the tetrahedron
        addAllNewPoints!(cell, cell.faceIndices[2])
    end

    function populatePointIndices_Pyramid!(cell::Cell)
        # Pyramids need to have the points that make up their square base ordered 1-4, with the tip of the pyramid coming last
        # Find the quad face
        for faceIndex in cell.faceIndices
            if length(pointIndicesByFace[faceIndex]) == 4
                addAllNewPoints!(cell, faceIndex)

                # There is only one square face, and all other faces contain the pyramid tip
                # Find any other face and add its points to make the pyramid tip the final point
                if faceIndex != cell.faceIndices[1]
                    addAllNewPoints!(cell, cell.faceIndices[1])
                else
                    addAllNewPoints!(cell, cell.faceIndices[2])
                end

                # All done
                break
            end
        end
    end

    function populatePointIndices_Wedge!(cell::Cell)
        # Wedges need the triangular faces numbered 1-3 and 4-6 respectively, where 1 is aligned with 4, 2 with 5, and 3 with 6
        # Using similar strategy it to that used for hex cells: Identify two end faces, then use the faces that connect them to ensure their points are aligned

        # First identify the two triangular faces, keep track of one and get rid of the other
        unusedFaces = deepcopy(cell.faceIndices)
        t1 = -1 # Triangular face one, we will base the ordering off this one

        for j in 1:2
            for i in eachindex(unusedFaces)
                faceIndex = unusedFaces[i]

                if length(pointIndicesByFace[faceIndex]) == 3
                    # Found a triangular face
                    if t1 == -1
                        t1 = faceIndex
                    end
                    deleteat!(unusedFaces, i)
                    break
                end
            end
        end

        if length(unusedFaces) != 3
            throw(ErrorException("Expected three quadrilateral faces remaining after removal of two triangular end faces, got $(length(unusedFaces))"))
        end

        t1Points = pointIndicesByFace[t1]
        t1PointsSet = Set(t1Points)
        cell.pointIndices = vcat(t1Points, [0, 0, 0]) # Still have to determine proper ordering of points in the other triangular end face, 0's will be replaced below

        # Now all the remaining faces (3) will be quadrilaterals connecting the two triangular end faces
        q1 = pop!(unusedFaces) # Select one at random
        q1Points = pointIndicesByFace[q1]
        addEdges!(cell.pointIndices, q1Points, t1Points, t1PointsSet, unusedFaces, 3) # Use it to populate 2/3 of the unknown points

        q2 = pop!(unusedFaces) # Select another at random
        q2Points = pointIndicesByFace[q2]
        addEdges!(cell.pointIndices, q2Points, t1Points, t1PointsSet, unusedFaces, 3) # Use it to populate the last unknown point
    end

    function populatePointIndices_Hex!(cell::Cell)
        # Need to order point such that the points that make up a single quadrilateral face are ordered 1-4
            # The face opposite to it must have its points ordered 5-8, with point 1 connected to/aligned with point 5, 2 with 6, etc...
            # We will accomplish this by first picking a starting face at random, this face's points will be numbered 1-4
            # Then we will identify its opposite face and remove it
            # Then all the remaining faces form the connections between the first face we picked and its opposite Face
            # Pairs of these faces will contain 0 or 2 points in common
                # When zero, they are opposite each other
                # When two, they are adjacent and the two points they have in common form an edge connecting the two end faces
                    # This indicates that these two points should be aligned, and is used to properly order the remaining points

        # Start by choosing an arbitrary face
        unusedFaces = deepcopy(cell.faceIndices)
        f1 = pop!(unusedFaces)
        f1Points = pointIndicesByFace[f1]
        f1PointsSet = Set(f1Points)

        # Find the opposite face (no points in common) and remove it
        for i in 1:length(unusedFaces)
            fi = unusedFaces[i]
            fiPoints = pointIndicesByFace[fi]

            if disjoint(f1Points, fiPoints)
                deleteat!(unusedFaces, i)
                break
            end
        end

        ### Now find the correct orientation
            # Point f1_i needs to be aligned (spatially) with Point f2_i
            # We can check for this by using one of the other faces as a guide, since it forms part of the connection between f1 and f2
        cell.pointIndices = vcat(f1Points, [0, 0, 0, 0])

        # Pick an arbitrary other face
        f3 = pop!(unusedFaces)
        f3Points = pointIndicesByFace[f3]
        lastFace = addEdges!(cell.pointIndices, f3Points, f1Points, f1PointsSet, unusedFaces)

        # Add points from the edges of the face opposite f3
        lastFacePoints = pointIndicesByFace[unusedFaces[lastFace]]
        deleteat!(unusedFaces, lastFace)
        noResult = addEdges!(cell.pointIndices, lastFacePoints, f1Points, f1PointsSet, unusedFaces)

        if noResult != -1
            throw(ErrorException("Failure to appropriately order hexahedral cell points for writing to .vtk"))
        end
    end

    # Now use the face indices to gather and order each cell's points appropriately
        # Correct point ordering is important for .vtk output: (See figure 2) https://vtk.org/wp-content/uploads/2015/04/file-formats.pdf
    for cell in cells
        nFaces = length(cell.faceIndices)
        if nFaces == 4
            populatePointIndices_Tet!(cell)
        elseif nFaces == 5
            quadFaceCount = 0

            for f in cell.faceIndices
                if length(pointIndicesByFace[f]) == 4
                    quadFaceCount += 1
                end
            end

            if quadFaceCount == 1
                populatePointIndices_Pyramid!(cell)
            elseif quadFaceCount == 3
                populatePointIndices_Wedge!(cell)
            else
                throw(ErrorException("Unrecognized cell type, cell: $cell. Expecting only hex, tet, wedge, or pyramid cells for vtk output."))
            end
        elseif nFaces == 6
            populatePointIndices_Hex!(cell)
        end
    end

    return pointLocations, cells
end

#=
    Reads an OpenFOAM mesh and returns a Mesh object suitable for calculations in FvCFD.jl.

    Mesh defined in dataStructures.jl, documented in dataStructureDefintions.md
=#
function OpenFOAMMesh(polyMeshPath)
    points, OFfaces, owner, neighbour, boundaryNames, boundaryNumFaces, boundaryStartFaces = readOpenFOAMMesh(polyMeshPath)
    nCells = maximum(owner)
    nFaces = size(OFfaces, 1)
    nBoundaries = size(boundaryNames, 1)

    cells = Array{Array{Int64, 1}, 1}(undef, nCells)                # Indices of faces that make up a cell
    cVols = zeros(nCells)                                           # Cell Volumes
    cCenters = Array{Array{Float64, 1}, 1}(undef, nCells)           # Cell Centroids
    cellSizes = Matrix{Float64}(undef, nCells, 3)                   # Cell sizes (x,y,z) direction

    faces = Array{Array{Int64, 1}, 1}(undef, nFaces)                # Indices of cells adjacent to each face always owner first, then neighbour
    fAVecs = Array{Array{Float64, 1}, 1}(undef, nFaces)             # Face Area Vectors
    fCenters = Array{Array{Float64, 1}, 1}(undef, nFaces)           # Face Centroids
    boundaryFaces = Array{Array{Int64, 1}, 1}(undef, nBoundaries)   # Indices of faces that make up each boundary

    # Calculate face area vectors and face centroids from list of points in each face
    for f in 1:nFaces
        fPts = [ points[pt,:] for pt in OFfaces[f] ]
        fAVecs[f], fCenters[f] = faceAreaCentroid(fPts)
    end
    #fAVecs and fCenters now complete

    # Build 'cells' and 'faces' arrays from info from 'owner' file
    nOwners = size(owner, 1)
    for f in 1:nOwners
        ownerCell = owner[f]
        if isassigned(cells, ownerCell)
            push!(cells[ownerCell], f)
        else
            cells[ownerCell] = [f,]
        end

        faces[f] = [ownerCell, -1]
    end

    # Finish building 'cells' and 'faces' arrays, adding info from 'neighbour' file
    nNeighbours = size(neighbour, 1)
    for f in 1:nNeighbours
        neighbourCell = neighbour[f]
        if isassigned(cells, neighbourCell)
            push!(cells[neighbourCell], f)
        else
            cells[neighbourCell] = [f,]
        end

        faces[f][2] = neighbourCell
    end
    # fAVecs, fCenters, faces, cells now complete

    # For each cell, use the face area vectors and face centroids for all of its faces, to calculate it's volume and centroid location
    for c in 1:nCells
        pts = []
        for f in cells[c]
            for pt in OFfaces[f]
                if !(points[pt,:] in pts)
                    push!(pts, points[pt,:])
                end
            end
        end
        fCs = [ fCenters[f] for f in cells[c] ]
        cell_fAVecs = [ fAVecs[f] for f in cells[c] ]

        cVols[c], cCenters[c] = cellVolCentroid(pts, cell_fAVecs, fCs)
    end
    # fAVecs, fCenters, faces, cells, cVols, cCenters now complete

    # Create boundaryFaces array
    for b in 1:nBoundaries
        startF = boundaryStartFaces[b]
        endF = startF + boundaryNumFaces[b] - 1
        boundaryFaces[b] = Array(startF:endF)
    end

    # Compute cell sizes in the x, y, z directions
    # In the past, cell sizes were used for a shitty CFL calculation, may be able to get rid of them now
    maxCoords = [ -1000000.0, -1000000.0, -1000000.0 ]
    minCoords = [ 1000000.0, 1000000.0, 1000000.0 ]
    for c in 1:nCells
        fill!(maxCoords, -10000000)
        fill!(minCoords, 10000000)

        pts = Vector{Vector{Float64}}()

        # Add points
        for f in cells[c]
            for pt in OFfaces[f]
                if !any(x->x==pt, pts)
                    push!(pts, points[pt, :])
                end
            end
        end

        for pt in pts
            for d in 1:3
                maxCoords[d] = max(maxCoords[d], pt[d])
                minCoords[d] = min(minCoords[d], pt[d])
            end
        end

        for d in 1:3
            cellSizes[c,d] = maxCoords[d] - minCoords[d]
        end
    end

    return Mesh(cells, cVols, cCenters, cellSizes, faces, fAVecs, fCenters, boundaryFaces)
end


"""
    Cylinder(x, y, zmin, zmax, radius)

Create a `Cylinder` that can be used for spatial lookup.

See also `AbstractRegion`
"""
struct Cylinder{T <: Real} <: AbstractRegion{T}
    x::T
    y::T
    zmin::T
    zmax::T
    radius::T
end

"""
    volume(cylinder::Cylinder)

Returns the volume of the `cylinder`.
"""
@inline function volume(cylinder::Cylinder)
    π * cylinder.radius * cylinder.radius * (cylinder.zmax - cylinder.zmin)
end

function boundingbox(cylinder::Cylinder{T}) where T
    xmin = cylinder.x - cylinder.radius
    xmax = cylinder.x + cylinder.radius
    ymin = cylinder.y - cylinder.radius
    ymax = cylinder.y + cylinder.radius
    return BoundingBox{T}(xmin, ymin, cylinder.zmin, xmax, ymax, cylinder.zmax)
end

function in(p::StaticVector{3, <:Real}, cylinder::Cylinder)
    r² = cylinder.radius * cylinder.radius
    Δx = p[1] - cylinder.x
    Δy = p[2] - cylinder.y
    return Δx*Δx + Δy*Δy <= r² && p[3] >= cylinder.zmin && p[3] <= cylinder.zmax
end


#
# asian-option.jl
#
function run-asian(N = 100000, PutCall = 'C';)
# European Asian option.
# Uses geometric or arithmetic average.
# Euler and Milstein discretization for Black-Scholes.
# Option features.
  println("Setting option parameters");
  S0  = 100;      # Spot price
  K   = 100;      # Strike price
  r   = 0.05;     # Risk free rate
  q   = 0.0;      # Dividend yield
  v   = 0.2;      # Volatility
  tma = 0.25;     # Time to maturity

  Averaging = 'A';  # 'A'rithmetic or 'G'eometric
  OptType = (PutCall == 'C' ? "CALL" : "PUT");
  println("Option type is $OptType");
# Simulation settings.
  println("Setting simulation parameters");
  T = 100;         # Number of time steps
  dt = tma/T;      # Time increment

# Initialize the terminal stock price matrices
# for the Euler and Milstein discretization schemes.
 S = zeros(Float64,N,T);
  for n=1:N
    	S[n,1] = S0;
  end

# Simulate the stock price under the Euler and Milstein schemes.
# Take average of terminal stock price.
  println("Looping $N times.");
  A = zeros(Float64,N);
  for n=1:N
    	for t=2:T
      		dW = (randn(1)[1])*sqrt(dt);
     		 z0 = (r - q - 0.5*v*v)*S[n,t-1]*dt;
     		 z1 = v*S[n,t-1]*dW;
      		z2 = 0.5*v*v*S[n,t-1]*dW*dW;
      		S[n,t] = S[n,t-1] + z0 + z1 + z2;
   	 end
   	 if cmp(Averaging,'A') == 0
      		A[n] = mean(S[n,:]);
   	 elseif cmp(Averaging,'G') == 0
      		A[n] = exp(mean(log(S[n,:])));
   	 end
  end

# Define the payoff
  P = zeros(Float64,N);
  if cmp(PutCall,'C') == 0
    	for n = 1:N
      		P[n] = max(A[n] - K, 0);
  	  end
  elseif cmp(PutCall,'P') == 0
    	for n = 1:N
     		 P[n] = max(K - A[n], 0);
   	 end
  end
# Calculate the price of the Asian option
 AsianPrice = exp(-r*tma)*mean(P);
 @printf "Price: %10.4f\n" AsianPrice;
end

run-asian()


@doc raw"""
    ℛ(z::Number, tableau::AbstractButcherTableau) :: Number
    ℛ(z::Number, solver::AbstractRungeKuttaSolver) :: Number

returns the stability function of an [`AbstractRungeKuttaSolver`](@ref):
```math
    R(z) = \frac{\det(I - z(A - \mathbb{1}b^\intercal))}{\det(I - zA)}.
```
"""
function ℛ(z::Number, tableau::AbstractButcherTableau)
    @↓ A, b, s = tableau
    e = ones(s)
    return det(I - z * (A - e * b')) / det(I - z * A)
end
ℛ(z::Number, solver::AbstractRungeKuttaSolver) = ℛ(z, solver.tableau)

"""
    ℛ(Z::AbstractMatrix, tableau::AbstractButcherTableau) :: AbstractMatrix
    ℛ(Z::AbstractMatrix, solver::AbstractRungeKuttaSolver) :: AbstractMatrix

returns the stability function of an [`AbstractRungeKuttaSolver`](@ref).
"""
function ℛ(Z::AbstractMatrix, tableau::AbstractButcherTableau)
    @↓ A, b, s = tableau
    e = ones(s)
    tmp = kron(e, Z)
    tmp = (I - kron(A, Z)) \ tmp
    tmp = kron(b', Matrix(1.0I, size(Z)...)) * tmp
    return I + tmp
end
ℛ(Z::AbstractMatrix, solver::AbstractRungeKuttaSolver) = ℛ(Z, solver.tableau)


### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# ╔═╡ e077ab16-22e5-11eb-15ed-5f0654bf7328
using Pkg, DrWatson

# ╔═╡ e75a9a06-22e5-11eb-2e36-8d4f62b830ed
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Random, Distributions, DataFrames, Plots
	Random.seed!(0)
end;

# ╔═╡ ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
md"## Listing6.14"

# ╔═╡ 1aa41898-3afa-11eb-2f49-570853ce6435
begin
	lambda = 0.1
	dist = Exponential(1/lambda)
	actualMedian = median(dist)

	M = 10^3
	N = 10^4
	nRange = 2:2:10
	alpha = 0.05
end

# ╔═╡ fba0be84-3b2f-11eb-1e0f-5f8a1f264dba
begin
	df = DataFrame()
	for n in nRange
		coverageCount = 0
		for _ in 1:M
			sampleData = rand(dist, n)
			bootstrapSampleMeans = [median(rand(sampleData, n)) for _ in 1:N]
			L = quantile(bootstrapSampleMeans, alpha/2)
			U = quantile(bootstrapSampleMeans, 1-alpha/2)
			coverageCount += L < actualMedian && actualMedian < U
		end
		append!(df, DataFrame(:n => n, :coverage => coverageCount/M))
	end
	df
end

# ╔═╡ 475ff888-22e6-11eb-2354-09f8f40a8e12
md"## End of listing6.14"

# ╔═╡ Cell order:
# ╟─ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
# ╠═e077ab16-22e5-11eb-15ed-5f0654bf7328
# ╠═e75a9a06-22e5-11eb-2e36-8d4f62b830ed
# ╠═1aa41898-3afa-11eb-2f49-570853ce6435
# ╠═fba0be84-3b2f-11eb-1e0f-5f8a1f264dba
# ╟─475ff888-22e6-11eb-2354-09f8f40a8e12


### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# ╔═╡ bca1c92c-8561-11eb-0420-6f0123bacb7f
begin
	# dependencies
	using Plots
	using LightGraphs, SimpleWeightedGraphs, LinearAlgebra, Random
	import StatsBase:sample
end

# ╔═╡ 774fa124-8560-11eb-2317-53b0d106b67a
md"""
# Local Search
"""

# ╔═╡ 2b5e6876-8561-11eb-0685-dfe1cca39c45
md"""
## Simulated annealing
Simulated annealing is similar to hill climbing, but it allows downhill moves. There are several aspects to consider.
* What is considered a neighbor? And how to chose them?
* How to evaluate the energy (i.e. the fitness)?
* How to let the temperature evolve. Several options exist here:
    - exponential decrease: $T = T_0(1-\alpha)^k$, where $\alpha$ is the ‘cooling rate’ and $k$ is the iteration
    - fast decrease: $T = T_0/k$
    - boltzmann decrease: $T = T_0/log(k)$
"""

# ╔═╡ 4e190a2e-8561-11eb-3759-795e3bb4d8d2
md"""
To illustrate the principle, let's try to find the maximum of
\\[ f(x) = -\dfrac{1}{4}x^4 +  \dfrac{7}{3} x^3 - 7x^2 + 8x \\]
in the interval $[0, 5]$. You can observe that there is a global maximum in $x=4$, however, there is a local maximum in $x=1$.
"""

# ╔═╡ b49f0528-8561-11eb-36a7-fd42574edc81
begin
	f = x-> -1/4*x^4 + 7/3 *x^3 - 7*x^2 + 8*x
	x = range(0,5, length=100);
	plot(x,f.(x),label="",size=(300,200))
end

# ╔═╡ d5019d8a-8561-11eb-2a43-a9ce616532f5
begin
	"""
	Evaluate the fitness of a specific point
	"""
	function fitness(x::Float64,f::Function)
		return f(x)
	end

	"""
	Determine successors for a given location (with a fixed step size)
	You could include automatic step size modification if you want.
	"""
	function successors(x::Float64; step=0.1)
		return x .+ step*[-1; 1]
	end

	"""
	mapping of time to temperature. You might make this a wrapper function for
	different cooling schemes. Currently only has exponential cooling.
	"""
	function scheduletemp(t::Int64;T0::Float64=273.0, α::Float64=0.005)
		return T0*(1-α)^t
	end

	"""
		simulated_annealing(args; kwargs)

	Simulated annealing implementation for a 1D problem.

	Since this is stochastic (sort of a random walk), it will not always reach the optimal value, therefore, one typically runs several iterations and the best one is retained.

	The probabiliy of accepting a value that lead to a decrease in energy is given by exp(ΔE/T), which is between 0 (temp = 0) and one (ΔE for any T ≠ 0).

	In the literature you might also find  1/(1 + exp(ΔE/T)) ∈ [1/2, 1] as a treshold.

	arguments:
	- x0: starting point
	keywords:
	- tmax: maximum iterations
	- T0: initial temperature
	- α: cooling rate
	- fitness: fitness function
	- optimfun: the function we want to optimize
	"""
	function simulated_annealing(x0;
				tmax=1000, T0::Float64=273.0, α::Float64=0.005,
				fitness::Function=fitness, optimfun::Function=f)
		current = x0
		next = x0
		T = T0
		best = x0
		for t = 1:tmax
			T = scheduletemp(t,T0=T0,α=α)
			if T == 0
				return current
			else
				next = rand(successors(current))
				ΔE = fitness(next, optimfun) - fitness(current, optimfun)
				if ΔE > 0
					current = next
				else
					if rand() < exp(ΔE/T)
						current = next
					end
				end
				if fitness(current, optimfun) > fitness(best, optimfun)
					best = current
				end
			end
		end
		return best
	end

	"""
		SAsearch

	Wrapper function for multiple SA iterations. Returns the best n found value for N runs of simulated annealing.
	"""
	function SAsearch(x0, n::Int64=1, N::Int64=10;kwargs...)
		return sort!([simulated_annealing(x0; kwargs...) for _ in 1:N],
					 by=x->kwargs[:fitness](x,kwargs[:optimfun]),
					 rev=true)[1:n]
	end
end

# ╔═╡ 9c2e2032-8568-11eb-103f-fdc57c11ab1f
md"""
### Solving the TSP problem with simulated annealing
Use simulated annealing to solve the travelling saleman problem again. You will need to think in particalar about:
* an appropriate inital solution (e.g. a tour by using a nearest neigbor approach or purely random)
* generating appropriate neigbors (e.g. by flipping a subpart of a tour)
"""

# ╔═╡ c84f6af2-8568-11eb-06c4-53f7664294ac
begin
	"""
		plotgraph(g::SimpleWeightedGraph, d::Dict,draw_MST::Bool=true,wf::Real=5)

	Visualize the network with its edges and their weights

	parameters:
		- g : SimpleWeightedGraph
		- d : Dict
			dictionary that links a node to a coordinate pair
		- draw_MST : Bool
			if true, also draws the MST on the figure
		- wf : Real
			multiplication factor for the edge weigth to obtain a nicer visualisation
	"""
	function plotgraph(g::SimpleWeightedGraph, d::Dict;wf::Real=5,draw_MST::Bool=true)
		p = plot(legend=:bottomright)
		for e in edges(g)
			plot!([d[e.src][1], d[e.dst][1]],[d[e.src][2], d[e.dst][2]],
				markershape=:circle, linewidth=e.weight*wf, color=:black, markeralpha = 0.5,
				linealpha=0.2, label="", markercolor=:black, markersize=5)
		end
		plot!([],[],markershape=:circle, color=:black,
				linealpha=0.2, label="Full network", markercolor=:black, markersize=5, markeralpha = 0.5)
		if draw_MST
			MST = prim_mst(g, weights(g))
			for e in MST
				plot!([d[e.src][1], d[e.dst][1]],[d[e.src][2], d[e.dst][2]],
						markershape=:circle, linewidth=weights(g)[e.src,e.dst]*wf, color=:red,
						linealpha=0.8, label="", linestyle=:solid, markeralpha = 0)

			end
			plot!([],[],linealpha=0.8, label="MST", color=:red, linestyle=:solid)
		end
		xlims!(0,1)
		ylims!(0,1)
		plot!(legend=:best)
		return p
	end

	"""
		plotsol(fp::Array{Int64,1}, nodedict::Dict)

	Function that plots the final TSP solution
	"""
	function plotsol(fp::Array{Int64,1}, nodedict::Dict)
		p = plot()
		for i in 1:length(fp)-1
			x = [nodedict[fp[i]][1]; nodedict[fp[i+1]][1]]
			y = [nodedict[fp[i]][2]; nodedict[fp[i+1]][2]]
			plot!(x,y,color=:blue,label="", marker=:circle)
		end
		# close loop
		x = [nodedict[fp[1]][1]; nodedict[fp[end]][1]]
		y = [nodedict[fp[1]][2]; nodedict[fp[end]][2]]
		plot!(x,y,color=:red,label="", marker=:circle)
		xlims!(0,1)
		ylims!(0,1)
		return p
	end

	"plot functions"
end

# ╔═╡ 46e468e0-8569-11eb-0a30-d11e814ccf7e
begin
	"""
		problemgenerator(n::Int)

	generate n points in the unit square
	"""
	function problemgenerator(n::Int)
		return [(rand(), rand()) for _ in 1:n]
	end

	"""
		distancematrix(x::Array{Tuple{Float64,Float64},1}, p::Real=1)

	Generate distancematrix between all points.

	x is an array [n x 1] and p is the L_p-norm you want to use for the distancematrix
	"""
	function distancematrix(x::Array{Tuple{Float64,Float64},1}, p::Real=1)
		n = length(x)
		A = zeros(n,n)
		for (i,j) in collect(Base.Iterators.product(1:n, 1:n))
			A[i,j] = norm(x[i] .- x[j], p)
		end
		return A
	end
end

# ╔═╡ 7c5282fa-8569-11eb-0dcf-abc7b3b795ef
begin
	N = problemgenerator(15) # nodes
	nodedict = Dict( i => N[i] for i in 1:length(N)) # nodes dict
	A = distancematrix(N,2) # distance matrix using manhattan norm
	G = SimpleWeightedGraph(A) # graph
	p = plotgraph(G,nodedict;draw_MST=false) # layout of the current graph
end

# ╔═╡ eb27eada-8569-11eb-0739-7717087a5240
begin
	"""
		finit(A::Array{Float64,2})

	Generate an initial distribution based and a random order
	"""
	function finit(A::Array{Float64,2})
		return shuffle(1:size(A,1))
	end

	finit(A)
end

# ╔═╡ 1692c654-856a-11eb-1b02-77b0ab027c29
begin
	"""
		fitness(T::Array{Int64,1},A::Array{Float64,2})

	Given a tour, compute the tour distance based on the distance matrix
	"""
	function fitness(T::Array{Int64,1},A::Array{Float64,2})
		return -sum([A[T[i],T[i+1]] for i in 1:length(T)-1]) - A[T[1],T[end]]
	end

	fitness(finit(A),A)
end

# ╔═╡ 3c58cc1a-856a-11eb-17bd-a31c244f6229
begin
	"""
	Determine successors for a given tour (simple random permutation of partial tour)
	"""
	function successors(T::Array{Int64,1}; nsuc::Int64=5)
		res = Array{Array{Int64,1},1}()
		for _ in 1:nsuc
			i,j = sample(1:length(T),2,replace=false,ordered=true)
			cand = copy(T)
			cand[i:j] = shuffle(T[i:j])
			push!(res, cand)
		end

		return res
	end
	successors(finit(A))
end

# ╔═╡ 7ab96168-856c-11eb-166a-b7531078ac90
begin
	"""
		simulated_annealing(args; kwargs...)

	Simulated annealing implementation for TSP

	Since this is stochastic (sort of a random walk), it will not always reach the optimal value, therefore, one typically runs several iterations and the best one is retained.

	The probabiliy of accepting a value that lead to a decrease in energy is given by exp(ΔE/T). In the literature you might also find  1/(1 + exp(ΔE/T)) ∈ [1/2, 1] as a treshold.
	"""
	function simulated_annealing(A::Array{Float64,2}, init::Array{Int64,1};
								tmax=1000, T0::Float64=273.0, α::Float64=0.005,
								fitness::Function=fitness,nsuc::Int64=10, kwargs...)
		current = init
		next = init
		T = T0
		best = init
		for t = 1:tmax
			T = scheduletemp(t,T0=T0)
			if T == 0
				return current
			else
				next = rand(successors(current, nsuc=nsuc))
				ΔE = fitness(next, A) - fitness(current, A)
				if ΔE > 0
					current = next
				else
					if rand() < exp(ΔE/T)
						current = next
					end
				end
				if fitness(current, A) > fitness(best, A)
					best = current
				end
			end
		end
		return best
	end

	"""
	wrapper function for multiple SA iterations

	Returns the best n found value for N runs of simulated annealing.
	"""
	function SAsearch(A::Array{Float64,2}, n::Int64=1, N::Int64=10;kwargs...)
		sort!([simulated_annealing(A, finit(A); kwargs...) for _ in 1:N],
				by=x->kwargs[:fitness](x,A),rev=false)[1:n]
	end

	simulated_annealing(A, finit(A))
end

# ╔═╡ 934aa9e4-856e-11eb-1924-fd1df8fd3fe2
md"""
## Genetic algorithms
a genetic algorithm (GA) is inspired by the process of natural selection. In order to implement this you need the following items:
* an initial population (e.g. random)
* a fitness function that allows you to rank the population members
* a selection method (e.g. roulette based proportions, ‘elite selection’)
* a cross-over method (cutting; 50/50 per gene etc.)
* mutation rate and method
* \\(\dots\\)
"""

# ╔═╡ 8753fd9a-856f-11eb-2d91-13657e302031
md"""
### Toy example
Let's try write a message by using a genetic algorithm. E.g "172 Pol evolved!" starting from a random set of characters (a-zA-Z0-9. !)
"""

# ╔═╡ 48c05e06-8570-11eb-0367-1f3db0aa20d0
begin
	chars = vcat(65:90,97:122, 48:57,[Int(' ');Int('.'); Int('!')])
	const genes = [Char(i) for i = chars]
	const pop_size = 100
	const goal = collect("172 Pol evolved!");
end

# ╔═╡ 6faf987e-8570-11eb-1861-7de0f661bd63
begin
	function fitness(p::Array{Char,1}; goal=goal)
		sum(p.==goal)
	end

	testcase = rand(genes,length(goal))
	testcase, fitness(testcase)
end

# ╔═╡ 728b94da-8566-11eb-1b02-e94c64424845
begin
	options = Dict(	:tmax=>1000, :T0=>273.0, :α=>0.005,
               		:fitness=>fitness, :optimfun=>f)
	x0 = 0.0;
	simulated_annealing(x0; options...)
end

# ╔═╡ a72d7c3c-8569-11eb-36a3-cb68ab2622cf
let
	n = 10;
	res = SAsearch(x0,n;options...)
end

# ╔═╡ 0ce6432c-8567-11eb-1465-432d66051360
begin
	n = 10;
	res = SAsearch(x0,n;options...)
	plot(x,f.(x),label="")
	scatter!([x0],[f.(x0)],label="start",marker=(:circle, :green),size=(300,200), legend=:bottom)
	scatter!([res],[f.(res)],label="optimal values",marker=(:circle, :red),size=(300,200), legend=:bottom)
end

# ╔═╡ ef64ca86-856c-11eb-0cd0-178565e85576
let
	options = Dict(:tmax=>1000, :T0=>273.0, :α=>0.005, :fitness=>fitness)
	@show res = SAsearch(A; options...)[1];
	init = finit(A)
	println("Cost of a random trip: $(-fitness(init,A))")
	println("Cost of a optimised trip: $(-fitness(res,A))")
	plotsol(res, nodedict)
	title!("TSP solution using simulated annealing")
end

# ╔═╡ 57726a64-856d-11eb-2b37-877202c47ede
let
	# Larger example
	n = 50 # number of nodes
	N = problemgenerator(n) # nodes
	options = Dict(:tmax=>1000, :T0=>273.0, :α=>0.005, :fitness=>fitness)
	nodedict = Dict( i => N[i] for i in 1:length(N)) # nodes dict
	A = distancematrix(N,2) # distance matrix using manhattan norm
	init = finit(A)

	@time res = SAsearch(A; options...)[1]
	println("Cost of a random trip: $(-fitness(init,A))")
	println("Cost of a optimised trip: $(-fitness(res,A))")
	plotsol(res, nodedict)
end

# ╔═╡ d70478c8-856d-11eb-3f27-0b565955b6ed
let
	# trajectory cost distribution
	res = Array{Float64,1}()
	for _ in 1:50
		push!(res, -fitness(SAsearch(A;options...)[1], A))
	end
	histogram(res, normalize=:pdf, alpha=0.5,label="PDF optimised tour cost")
	scatter!([-fitness(finit(A),A)],[0],label="initial non-optimised cost")
	xlabel!("Tour cost")
	ylabel!("PDF")
	title!("Histogram for tour cost ($(size(A,1)) cities)")
end

# ╔═╡ 9cc2fa40-8570-11eb-2eba-c15433dcb90e
begin
	"""
		mate(p1,p2; pmut, genes)

	mating function to obtain children. includes both crossover and mutation.
	"""
	function mate(p1::Array{Char,1}, p2::Array{Char,1}; pmut=0.1, genes=genes)
		child = typeof(p1)() # empty child
		# walk over genome sequence
		for i in 1:length(p1)
			p = rand()
			if p < (1-pmut)/2
				# parent 1 wins
				push!(child, p1[i])
			elseif p < 1-pmut
				# parent 2 wins
				push!(child, p2[i])
			else
				# mutation
				push!(child, rand(genes))
			end
		end
		return child
	end

	# demo
	p1 = ['a';'b';'c']
	p2 = ['d';'e';'f']
	mate(p1,p2, pmut=1)
end

# ╔═╡ 47153f8e-8572-11eb-1310-c932bd478721
md"""
### Some ideas for crossover:
Below you can find some algorithms that have proven to work quite well for this type of application. Both of these are described and an example and an implementation is provided. You are also invited to come up with a method yourself if you think you have a good idea.

#### The partially mapped crossover
After choosing two random cut points on parents to build offspring, the portion between cut points, one parent’s string is mapped onto the other parent’s string and the remaining information is exchanged.

E.g. for parents $P_1$ and $P_2$ with cutting positions $c_1$ and $c_2$:
* ``P_1 = [3, 4, 8 | 2, 7, 1 | 6, 5 ]``
* ``P_` = [4, 2, 5 | 1, 6, 8 | 3, 7 ]``
We get the following mappings:
* ``m_{1 \rightarrow 2}: \{2 \rightarrow 1, 7 \rightarrow 6, 1 \rightarrow 8\}``
* ``m_{2 \rightarrow 1}: \{1 \rightarrow 2, 6 \rightarrow 7, 8 \rightarrow 1\}``
which leads to the following initial offspring:
* ``O_1 = [0, 0, 0 | 1, 6, 8 | 0, 0 ] ``
* ``O_2 = [0, 0, 0 | 2, 7, 1 | 0, 0 ] ``
we then add the non-conflicting values (i.e. the values not already present the offspring due do the crossover) from the original parents:
* ``O_1 = [3, 4, 0 | 1, 6, 8 | 0, 5 ]``
* ``O_2 = [4, 0, 5 | 2, 7, 1 | 3, 0 ] ``
For the remaining values, we (recursively) make use of the mappings using $m_{2 \rightarrow 1}$ for $P_1$ and $m_{1 \rightarrow 2}$ for $P_2$. This finally leads to the following offspring:
* ``O_1 = [3, 4, 2 | 1, 6, 8 | 7, 5 ] ``
* ``O_2 = [4, 8, 5 | 2, 7, 1 | 3, 6 ] ``

#### The cycle crossover
The Cycle Crossover operator identifies a number of so-called cycles between two parent chromosomes. To form Child 1, cycle 1 is copied from parent 1, cycle 2 from parent 2, cycle 3 again from parent 1, and so on. An example might help to better understand the principle of a cycle. E.g.
* ``P_1 = [1,2,3,4,5,6,7,8] ``
* ``P_2 = [8,5,2,1,3,6,4,7] ``
For the first cycle, we start at index 1 in $P_1$. Following the links from $P_1$ to $P_2$ we find the following mappings:
* ``1 \rightarrow 8`` with index 8
* ``8 \rightarrow 7`` with index 7
* ``7 \rightarrow 4`` with index 4
* ``4 \rightarrow 1`` with index 1
after this last mapping, we are back at the starting point, so our first cycle is complete and is $\{1,4,7,8 \}$.

The second cycle will start at the first index that is not yet present in a cycle. In this case the index is equal to 2. Following the links from $P_1$ to $P_2$ we find the following mappings
* ``2 \rightarrow 5`` with index 2
* ``5 \rightarrow 3`` with index 5
* ``3 \rightarrow 2`` with index 3

So our second cycle is equal to $\{2,3,5\}$.

The third an final cycle is equal to the only remaining part, i.e. $\{6 \}$

We have now determined all cycles. The children aren obtained as follows:
* ``O_1[1,4,7,8] =  P_1[1,4,7,8] ``
* ``O_1[2,3,5] = P_2[2,3,5] ``
* ``O_1[6]  = P_1[6] ``

This leads to $O_1 = [1,5,2,4,3,6,7,8]$. In a similar fashion, but inverting the role of $P_1$ and $P_2$ one finds $O_2 = [8,2,3,1,5,6,4,7]$.

A downside of this type of crossover is that it can sometimes lead to the offspring being identical to the parents. You can verify this for $[3,4,8,2,7,1,6,5]$ & $[4,2,5,1,6,8,3,7]$.

#### Your own method
Come up with something that seems adequate...

"""

# ╔═╡ 1f5d255c-8576-11eb-1463-89f94258168d
begin
	"""
		mate(p1,p2;method=partialcrossover, kwargs...)

	Function that handles the crossover for two different tours. The method keyword argument can be any of the functions you define as long as 1. the function exists 2. the function works in a similar way as the other ones (i.e. return two children from two parents.)

	Additional options for your own crossover implementation can be passed via the kwargs. You should document the options for each algorithm sperately.

	"""
	function mate(p1::Array{Int64,1},p2::Array{Int64,1}; method::Function=partialcrossover, kwargs...)
			return method(p1,p2; kwargs...)
	end

	"""
		partialcrossover(p1, p2; kwargs...)

	Implementation of the partial crossover function for TSP routing problems.
	"""
	function partialcrossover(p1::Array{Int64,1},p2::Array{Int64,1}; kwargs...)
		"""
			deconflicter!(O::Array{Int64,1}, mapping::Dict)

		Function that deals with the conflicts after the initial assignments.
		"""
		function deconflicter(O::Array{Int64,1},p::Array{Int64,1}, mapping::Dict)
			for i in 1:length(O)
				if O[i] == 0
					cand = p[i]
					while cand in O
						cand = mapping[cand]
					end
					O[i] = cand
				end
			end
			return O
		end

		# choose cut points
		c1, c2 = sort(sample(1:length(p1),2, replace=false))
		# initiate offspring
		O1, O2 = [zeros(Int64,length(p1)) for _ in 1:2]
		# determine mapping
		map12 = Dict(p1[i]=> p2[i] for i in c1:c2)
		map21 = Dict(value => key for (key, value) in map12)
		# set offspring initial values
		O1[c1:c2] = p2[c1:c2]
		O2[c1:c2] = p1[c1:c2]
		# fill non-conflicting positions
		for i in setdiff(1:length(p1),c1:c2)
			O1[i] = p1[i] in O1 ? 0 : p1[i]
			O2[i] = p2[i] in O2 ? 0 : p2[i]
		end
		# fill conflicting positions using the mapping
		O1, O2 = deconflicter(O1, p1, map21), deconflicter(O2, p2, map12)

		return O1, O2
	end

	"""
		cyclecrossover(p1, p2)

	Implementation of the cycle crossover function for TSP routing problems.
	"""
	function cyclecrossover(p1::Array{Int64,1},p2::Array{Int64,1})
		# initiate children
		o1=zeros(Int64,length(p1))
		o2=zeros(Int64,length(p2))

		# Pt. 1 - find all cycles

		# make hashtable: value => index
		h1 = Dict{Int64, Int64}(p1[i]=>i for i in 1:length(p1))
		h2 = Dict{Int64, Int64}(p2[i]=>i for i in 1:length(p2))
		# make hashtable to track usage in a cycle: index => bool
		u1 = Dict{Int64, Bool}(i=>false for i in 1:length(p1))
		# loop over parent creating cycles on the go
		# a cycle is stored in a Set because this is hashable. This is (much) faster than an array for ∉ check
		# the values stores in the cycles are the indices to be used
		cycles = Array{Set{Int64},1}()
		for i in 1:length(p1)
			if !u1[h1[p1[i]]] # if not used yet, start a new cycle
				cycle = Set{Int64}()
				next = h1[p1[i]] # add initial values
				while next ∉ cycle
					push!(cycle, next)  # update cycle
					u1[next] = true     # update tracking
					next = h1[p2[next]] # next location
				end
				push!(cycles, cycle)
			else
				continue
			end
		end

		# Pt. 2 - make children from cycles
		for i in 1:length(cycles)
			inds = collect(cycles[i])
			if isodd(i)
				o1[inds] = p1[inds]
				o2[inds] = p2[inds]
			else
				o1[inds] = p2[inds]
				o2[inds] = p1[inds]
			end
		end
		return o1, o2
	end

	"""
		mycrossover(p1, p2; myoption::Float64=0.5, kwargs...)

	Your own method that takes predefined options and also accepts others.
	"""
	function mycrossover(p1::Array{Int64,1},p2::Array{Int64,1}; myoption::Float64=0.5, kwargs...)
		error("Function <mycrossover> not defined yet... Get creative!")
	end
end

# ╔═╡ fd95d464-8570-11eb-1a38-879d8079d127
begin
	function textga(;goal::Array{Char,1}=goal, popsize::Int64=pop_size,
                	 ngen::Int64=500, next::Float64=0.4, pmut=0.1, kwargs...)
		# initiate population
		population = [rand(genes, length(goal)) for _ in 1:popsize]
		gen = 0
		for _ in 1:ngen
			# sort by fitness
			sort!(population; by=x->fitness(x), rev=true)
			if population[1] == goal
				break
			end
			gen += 1
			# select adequate parents
			goodparents = population[1: round(Int,popsize*next)]
			# make children
			new_generation = typeof(population)()
			for _ in 1:popsize
				push!(new_generation, mate(rand(goodparents), rand(goodparents),
											pmut=pmut ))
			end
			population = new_generation
		end
		@info "stopped after $(gen) generations"

		return (gen, prod(population[1]), fitness(population[1])/fitness(goal))
	end

	textga(ngen=500, popsize=100, pmut=0.1)
end

# ╔═╡ 7c4ef658-8576-11eb-3f6d-d3587a1af7e8
begin
	let
		# small demo matching the explanation of the algorithm above.
		# Matches the detailed example if you force c1, c2 = 4, 6 in the partialcrossover algorithm
		p1 = [3,4,8,2,7,1,6,5]
		p2 = [4,2,5,1,6,8,3,7]
		@time mate(p1,p2)
	end

	let
		# Longer route does not necessarily imply greater time for offspring generation
		# with the partial crossover algorithm (why could that be?)
		@info "100 items"
		N = 100
		p1 = collect(1:N)
		p2 = sample(p1,N,replace=false)
		@time mate(p1,p2);
	end

	let
		# demo from the method description
		p1 = [1,2,3,4,5,6,7,8]
		p2 = [8,5,2,1,3,6,4,7]
		@time res = mate(p1,p2, method=cyclecrossover)
		@assert res ==  ([1,5,2,4,3,6,7,8],  [8,2,3,1,5,6,4,7])
	end

	let
		# demo for children equal to parents
		p1 = [3,4,8,2,7,1,6,5]
		p2 = [4,2,5,1,6,8,3,7]
		@time res = mate(p1,p2, method=cyclecrossover)
		@assert res == (p1 ,p2)
	end

	let
		# Longer routes scales nicely for cyclecrossover as well (≈10x longer route and only ≈2x slower)
		N = 100
		p1 = collect(1:N)
		p2 = sample(p1,N,replace=false)
		@time res = mate(p1,p2, method=cyclecrossover);
	end
end

# ╔═╡ 8b77cf1c-8571-11eb-2ea7-531abac106cc
md"""
### Genetic algorithm - TSP
Now that you are familiar with a genetic algorithm, let's try to solve the TSP problem again, but this time, using a genetic algorithm.

We should think about:
* representation
* an initial configuration
* a fitness function
* a selection method (e.g. roulette based proportions, ‘elite selection’ etc.)
* cross-over method: (what is the difference with the text application?)
* mutation rate and method (is this required? why (not)?)
* when do you stop making new generations?
* \\( \dots \\)

Tasks:
* Come up with your own suggestions for cross-over/mutation
* Implement the genetic algorithm to solve the TSP
* Compare the performance of the tree cross-over algorithms to deal with eachother. Give that there is some randomness involved, a statistical approached might be more appropriate. Since the initial population will most likely be chosen at random, you should impose the random seed to make sure you have the same starting point for the comparisons.
* Compare the scalability with what you know from previous practical sessions (tree search etc.) What is better and what is worse?
"""

# ╔═╡ 712cb76c-8576-11eb-37c3-87beed1c3fa1
begin
	function TSP_ga(d::Array{Float64,2};
					popsize::Int64=30, ngen::Int64=500, next::Float64=0.4,
					method::Function=partialcrossover, kwargs...)
		# initiate population
		n = size(d,1)
		population = [sample(1:n,n,replace=false) for _ in 1:popsize]
		gen = 0
		# go over the generations
		for _ in 1:ngen
			# sort by fitness
			sort!(population; by=x->fitness(x,d), rev=true)
			#@show "current tour cost: $(fitness(population[1],d))"
			# include a stopcondition?
			gen += 1
			# select adequate parents
			goodparents = population[1: round(Int,popsize*next)]
			# make children
			new_generation = typeof(population)()
			for _ in 1:2:popsize
				p1,p2 = rand(goodparents,2)
				offspring = mate(p1,p2;method=method, kwargs...)
				append!(new_generation, offspring)
			end
			population = new_generation
		end

		return (gen, population[1], fitness(population[1],d))
	end
end

# ╔═╡ fa780328-8576-11eb-3a16-b7109ffbad63
let
	n = 50 # number of nodes
	N = problemgenerator(n) # nodes
	nodedict = Dict( i => N[i] for i in 1:length(N)) # nodes dict
	A = distancematrix(N,1) # distance matrix using manhattan norm
	@info fitness(finit(A), A)
	G = SimpleWeightedGraph(A) # graph
	@time begin
    	D = Dict()
    	methods = [cyclecrossover, partialcrossover]
		for method in methods
			res = []
			for _ in 1:100
				(_, tour, cost) = TSP_ga(A, ngen=20, popsize=n,method=method)
				push!(res,(tour,cost))
			end
			D[method] = [x[2] for x in res]
		end
	end

	# illustration
	histogram()
	for (key, val) in D
		histogram!(-val, label=string(key), normalize=:pdf, alpha=0.5)
	end
	scatter!([-fitness(finit(A),A)],[0],label="initial non-optimised cost")
	xlabel!("Tour cost")
	ylabel!("PDF")
	title!("Histogram for tour cost ($(size(A,1)) cities)")
end

# ╔═╡ 1a79eb5a-8577-11eb-168e-8fa9230f26fb

# ╔═╡ Cell order:
# ╟─774fa124-8560-11eb-2317-53b0d106b67a
# ╠═bca1c92c-8561-11eb-0420-6f0123bacb7f
# ╠═2b5e6876-8561-11eb-0685-dfe1cca39c45
# ╠═4e190a2e-8561-11eb-3759-795e3bb4d8d2
# ╠═b49f0528-8561-11eb-36a7-fd42574edc81
# ╠═d5019d8a-8561-11eb-2a43-a9ce616532f5
# ╠═728b94da-8566-11eb-1b02-e94c64424845
# ╠═a72d7c3c-8569-11eb-36a3-cb68ab2622cf
# ╠═0ce6432c-8567-11eb-1465-432d66051360
# ╟─9c2e2032-8568-11eb-103f-fdc57c11ab1f
# ╟─c84f6af2-8568-11eb-06c4-53f7664294ac
# ╠═46e468e0-8569-11eb-0a30-d11e814ccf7e
# ╠═7c5282fa-8569-11eb-0dcf-abc7b3b795ef
# ╠═eb27eada-8569-11eb-0739-7717087a5240
# ╠═1692c654-856a-11eb-1b02-77b0ab027c29
# ╠═3c58cc1a-856a-11eb-17bd-a31c244f6229
# ╠═7ab96168-856c-11eb-166a-b7531078ac90
# ╠═ef64ca86-856c-11eb-0cd0-178565e85576
# ╠═57726a64-856d-11eb-2b37-877202c47ede
# ╠═d70478c8-856d-11eb-3f27-0b565955b6ed
# ╟─934aa9e4-856e-11eb-1924-fd1df8fd3fe2
# ╟─8753fd9a-856f-11eb-2d91-13657e302031
# ╠═48c05e06-8570-11eb-0367-1f3db0aa20d0
# ╠═6faf987e-8570-11eb-1861-7de0f661bd63
# ╠═9cc2fa40-8570-11eb-2eba-c15433dcb90e
# ╠═fd95d464-8570-11eb-1a38-879d8079d127
# ╟─47153f8e-8572-11eb-1310-c932bd478721
# ╠═1f5d255c-8576-11eb-1463-89f94258168d
# ╠═7c4ef658-8576-11eb-3f6d-d3587a1af7e8
# ╟─8b77cf1c-8571-11eb-2ea7-531abac106cc
# ╠═712cb76c-8576-11eb-37c3-87beed1c3fa1
# ╠═fa780328-8576-11eb-3a16-b7109ffbad63
# ╠═1a79eb5a-8577-11eb-168e-8fa9230f26fb