using SPHKernels, Test

@testset "SPH Kernels" begin

    @testset "Cubic Spline" begin

        @testset "1D value" begin
            k = Cubic()
            # < 0.5
            d = kernel_value_1D(k, 0.4, 0.5)
            @test d ≈ 0.2826666666666666
            # < 1.0
            d = kernel_value_1D(k, 0.5, 0.5)
            @test d ≈ 0.16666666666666666
            # > 1.0
            d = kernel_value_1D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "1D derivative" begin
            k = Cubic()
            # < 0.5
            d = kernel_deriv_1D(k, 0.4, 0.5)
            @test d ≈ -0.6399999999999999
            # < 1.0
            d = kernel_deriv_1D(k, 0.5, 0.5)
            @test d ≈ -0.5
            # > 1.0
            d = kernel_deriv_1D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "1D bias correction" begin
            k = Cubic()
            d = bias_correction_1D(k, 1.0, 0.0, 0.0)
            @test d ≈  1.0
        end

        @testset "2D value" begin
            k = Cubic()
            # < 0.5
            d = kernel_value_2D(k, 0.4, 0.5)
            @test d ≈ 0.19280484534561032
            # < 1.0
            d = kernel_value_2D(k, 0.5, 0.5)
            @test d ≈ 0.11368210220849667
            # > 1.0
            d = kernel_value_2D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "2D derivative" begin
            k = Cubic()
            # < 0.5
            d = kernel_deriv_2D(k, 0.4, 0.5)
            @test d ≈ -0.4365392724806272
            # < 1.0
            d = kernel_deriv_2D(k, 0.5, 0.5)
            @test d ≈ -0.34104630662549
            # > 1.0
            d = kernel_deriv_2D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "2D bias correction" begin
            k = Cubic()
            d = bias_correction_2D(k, 1.0, 0.0, 0.0)
            @test d ≈  1.0
        end

        @testset "3D value" begin
            k = Cubic()
            # < 0.5
            d = kernel_value_3D(k, 0.4, 0.5)
            @test d ≈ 0.13496339174192723
            # < 1.0
            d = kernel_value_3D(k, 0.5, 0.5)
            @test d ≈ 0.07957747154594767
            # > 1.0
            d = kernel_value_3D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "3D derivative" begin
            k = Cubic()
            # < 0.5
            d = kernel_deriv_3D(k, 0.4, 0.5)
            @test d ≈ -0.30557749073643903
            # < 1.0
            d = kernel_deriv_3D(k, 0.5, 0.5)
            @test d ≈ -0.238732414637843
            # > 1.0
            d = kernel_deriv_3D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "3D bias correction" begin
            k = Cubic()
            d = bias_correction_3D(k, 1.0, 0.0, 0.0)
            @test d ≈  1.0
        end
    end

    @testset "Quintic Spline" begin

        @testset "1D value" begin
            k = Quintic()
            # < 1/3
            d = kernel_value_1D(k, 0.3, 0.5)
            @test d ≈ 0.38972624999999994
            # < 2/3
            d = kernel_value_1D(k, 0.5, 0.5)
            @test d ≈ 0.092578125
            # < 1.0
            d = kernel_value_1D(k, 0.8, 0.5)
            @test d ≈ 0.0009719999999999988
            # > 1.0
            d = kernel_value_1D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "1D derivative" begin
            k = Quintic()
            # < 1/3
            d = kernel_deriv_1D(k, 0.3, 0.5)
            @test d ≈ -0.99984375
            # < 2/3
            d = kernel_deriv_1D(k, 0.5, 0.5)
            @test d ≈ -114.345703125
            # < 1.0
            d = kernel_deriv_1D(k, 0.8, 0.5)
            @test d ≈ -0.012149999999999987
            # > 1.0
            d = kernel_deriv_1D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "1D bias correction" begin
            k = Quintic()
            d = bias_correction_1D(k, 1.0, 0.0, 0.0)
            @test d ≈  1.0
        end

        @testset "2D value" begin
            k = Quintic()
            # < 1/3
            d = kernel_value_2D(k, 0.3, 0.5)
            @test d ≈ 0.32700352517410614
            # < 2/3
            d = kernel_value_2D(k, 0.5, 0.5)
            @test d ≈ 0.07767855829318415
            # < 1.0
            d = kernel_value_2D(k, 0.8, 0.5)
            @test d ≈ 0.0008155658657050455
            # > 1.0
            d = kernel_value_2D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "2D derivative" begin
            k = Quintic()
            # < 1/3
            d = kernel_deriv_2D(k, 0.3, 0.5)
            @test d ≈ -0.8389284295663888
            # < 2/3
            d = kernel_deriv_2D(k, 0.5, 0.5)
            @test d ≈ -95.94285222098028
            # < 1.0
            d = kernel_deriv_2D(k, 0.8, 0.5)
            @test d ≈ -0.010194573321313068
            # > 1.0
            d = kernel_deriv_2D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "2D bias correction" begin
            k = Quintic()
            d = bias_correction_2D(k, 1.0, 0.0, 0.0)
            @test d ≈  1.0
        end

        @testset "3D value" begin
            k = Quintic()
            # < 1/3
            d = kernel_value_3D(k, 0.3, 0.5)
            @test d ≈ 0.2791208661307549
            # < 2/3
            d = kernel_value_3D(k, 0.5, 0.5)
            @test d ≈ 0.06630419797168217
            # < 1.0
            d = kernel_value_3D(k, 0.8, 0.5)
            @test d ≈ 0.0006961437210839494
            # > 1.0
            d = kernel_value_3D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "3D derivative" begin
            k = Quintic()
            # < 1/3
            d = kernel_deriv_3D(k, 0.3, 0.5)
            @test d ≈ -0.7160853380941675
            # < 2/3
            d = kernel_deriv_3D(k, 0.5, 0.5)
            @test d ≈ -81.89407743147959
            # < 1.0
            d = kernel_deriv_3D(k, 0.8, 0.5)
            @test d ≈ -0.008701796513549367
            # > 1.0
            d = kernel_deriv_3D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "3D bias correction" begin
            k = Quintic()
            d = bias_correction_3D(k, 1.0, 0.0, 0.0)
            @test d ≈  1.0
        end
    end

    @testset "Wendland C2" begin

        @testset "1D value" begin
            k = WendlandC2()
            # < 1.0
            d = kernel_value_1D(k, 0.5, 0.5)
            @test d ≈ 0.1953125
            # > 1.0
            d = kernel_value_1D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "1D derivative" begin
            k = WendlandC2()
            # < 1.0
            d = kernel_deriv_1D(k, 0.5, 0.5)
            @test d ≈ -0.9375
            # > 1.0
            d = kernel_deriv_1D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "1D bias correction" begin
            k = WendlandC2()
            d = bias_correction_1D(k, 1.0, 1.0, 0.5)

            @test d ≈ 0.99540625
        end

        @testset "2D value" begin
            k = WendlandC2()
            # < 1.0
            d = kernel_value_2D(k, 0.5, 0.5)
            @test d ≈ 0.10444543140405632
            # > 1.0
            d = kernel_value_2D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "2D derivative" begin
            k = WendlandC2()
            # < 1.0
            d = kernel_deriv_2D(k, 0.5, 0.5)
            @test d ≈ -0.3481514380135211
            # > 1.0
            d = kernel_deriv_2D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "2D bias correction" begin
            k = WendlandC2()
            d = bias_correction_2D(k, 1.0, 1.0, 0.5)

            @test d ≈ 0.991811478177922
        end

        @testset "3D value" begin
            k = WendlandC2()
            # < 1.0
            d = kernel_value_3D(k, 0.5, 0.5)
            @test d ≈ 0.07833407355304224
            # > 1.0
            d = kernel_value_3D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "3D derivative" begin
            k = WendlandC2()
            # < 1.0
            d = kernel_deriv_3D(k, 0.5, 0.5)
            @test d ≈ -0.26111357851014083
            # > 1.0
            d = kernel_deriv_3D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "3D bias correction" begin
            k = WendlandC2()
            d = bias_correction_3D(k, 1.0, 1.0, 0.5)

            @test d ≈ 0.987717217266883
        end
    end

    @testset "Wendland C4" begin

        @testset "1D value" begin
            k = WendlandC4()
            # < 1.0
            d = kernel_value_1D(k, 0.5, 0.5)
            @test d ≈ 0.12890625
            # > 1.0
            d = kernel_value_1D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "1D derivative" begin
            k = WendlandC4()
            # < 1.0
            d = kernel_deriv_1D(k, 0.5, 0.5)
            @test d ≈ -0.4921875
            # > 1.0
            d = kernel_deriv_1D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "1D bias correction" begin
            k = WendlandC4()
            d = bias_correction_1D(k, 1.0, 1.0, 0.5)

            @test d ≈ 0.9992541503078153
        end

        @testset "2D value" begin
            k = WendlandC4()
            # < 1.0
            d = kernel_value_2D(k, 0.5, 0.5)
            @test d ≈ 0.07740152505836316
            # > 1.0
            d = kernel_value_2D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "2D derivative" begin
            k = WendlandC4()
            # < 1.0
            d = kernel_deriv_2D(k, 0.5, 0.5)
            @test d ≈ -0.37301939787162974
            # > 1.0
            d = kernel_deriv_2D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "2D bias correction" begin
            k = WendlandC4()
            d = bias_correction_2D(k, 1.0, 1.0, 0.5)

            @test d ≈ 0.9985755320162227
        end

        @testset "3D value" begin
            k = WendlandC4()
            # < 1.0
            d = kernel_value_3D(k, 0.5, 0.5)
            @test d ≈ 0.06651693559703084
            # > 1.0
            d = kernel_value_3D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "3D derivative" begin
            k = WendlandC4()
            # < 1.0
            d = kernel_deriv_3D(k, 0.5, 0.5)
            @test d ≈ -0.3205635450459318
            # > 1.0
            d = kernel_deriv_3D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "3D bias correction" begin
            k = WendlandC4()
            d = bias_correction_3D(k, 1.0, 1.0, 0.5)

            @test d ≈ 0.9975516956528828
        end
    end

    @testset "Wendland C6" begin

        @testset "1D value" begin
            k = WendlandC6()
            # < 1.0
            d = kernel_value_1D(k, 0.5, 0.5)
            @test d ≈ 0.0797271728515625
            # > 1.0
            d = kernel_value_1D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "1D derivative" begin
            k = WendlandC6()
            # < 1.0
            d = kernel_deriv_1D(k, 0.5, 0.5)
            @test d ≈ -0.417938232421875
            # > 1.0
            d = kernel_deriv_1D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "1D bias correction" begin
            k = WendlandC6()
            d = bias_correction_1D(k, 1.0, 1.0, 0.5)

            @test d ≈ 0.9997781503817778
        end

        @testset "2D value" begin
            k = WendlandC6()
            # < 1.0
            d = kernel_value_2D(k, 0.5, 0.5)
            @test d ≈ 0.052822211162893276
            # > 1.0
            d = kernel_value_2D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "2D derivative" begin
            k = WendlandC6()
            # < 1.0
            d = kernel_deriv_2D(k, 0.5, 0.5)
            @test d ≈ -0.3238607700806899
            # > 1.0
            d = kernel_deriv_2D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "2D bias correction" begin
            k = WendlandC6()
            d = bias_correction_2D(k, 1.0, 1.0, 0.5)

            @test d ≈ 0.9995421822101074
        end

        @testset "3D value" begin
            k = WendlandC6()
            # < 1.0
            d = kernel_value_3D(k, 0.5, 0.5)
            @test d ≈ 0.05055250677698771
            # > 1.0
            d = kernel_value_3D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "3D derivative" begin
            k = WendlandC6()
            # < 1.0
            d = kernel_deriv_3D(k, 0.5, 0.5)
            @test d ≈ -0.30994487761628525
            # > 1.0
            d = kernel_deriv_3D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "3D bias correction" begin
            k = WendlandC6()
            d = bias_correction_3D(k, 1.0, 1.0, 0.5)

            @test d ≈ 0.9991237081365336
        end
    end

    @testset "Wendland C8" begin

        @testset "2D value" begin
            k = WendlandC8()
            # < 1.0
            d = kernel_value_2D(k, 0.5, 0.5)
            @test d ≈ 0.034310013366734275
            # > 1.0
            d = kernel_value_2D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "2D derivative" begin
            k = WendlandC8()
            # < 1.0
            d = kernel_deriv_2D(k, 0.5, 0.5)
            @test d ≈ -0.25896353614740847
            # > 1.0
            d = kernel_deriv_2D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "2D bias correction" begin
            k = WendlandC8()
            d = bias_correction_2D(k, 1.0, 1.0, 0.5)

            @test d ≈ 1.0
        end

        @testset "3D value" begin
            k = WendlandC8()
            # < 1.0
            d = kernel_value_3D(k, 0.5, 0.5)
            @test d ≈ 0.017155006683367138
            # > 1.0
            d = kernel_value_3D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "3D derivative" begin
            k = WendlandC8()
            # < 1.0
            d = kernel_deriv_3D(k, 0.5, 0.5)
            @test d ≈ -0.12948176807370423
            # > 1.0
            d = kernel_deriv_3D(k, 1.5, 0.5)
            @test d == 0.0
        end

        @testset "3D bias correction" begin
            k = WendlandC8()
            d = bias_correction_3D(k, 1.0, 1.0, 0.5)

            @test d ≈ 1.0
        end
    end

    @testset "Multiple Dispatch" begin

        @testset "kernel value" begin
            k = WendlandC6()
            @test 𝒲₁(k, 0.5, 1.0) ≈ kernel_value_1D(k, 0.5, 1.0)
            @test 𝒲₂(k, 0.5, 1.0) ≈ kernel_value_2D(k, 0.5, 1.0)
            @test 𝒲₃(k, 0.5, 1.0) ≈ kernel_value_3D(k, 0.5, 1.0)
        end

        @testset "kernel derivative" begin
            k = WendlandC6()
            @test ∇𝒲₁(k, 0.5, 1.0) ≈ kernel_deriv_1D(k, 0.5, 1.0)
            @test ∇𝒲₂(k, 0.5, 1.0) ≈ kernel_deriv_2D(k, 0.5, 1.0)
            @test ∇𝒲₃(k, 0.5, 1.0) ≈ kernel_deriv_3D(k, 0.5, 1.0)
        end

        @testset "bias correction" begin
            k = WendlandC6()
            @test δρ₁(k, 1.0, 1.0, 0.5) ≈ bias_correction_1D(k, 1.0, 1.0, 0.5)
            @test δρ₂(k, 1.0, 1.0, 0.5) ≈ bias_correction_2D(k, 1.0, 1.0, 0.5)
            @test δρ₃(k, 1.0, 1.0, 0.5) ≈ bias_correction_3D(k, 1.0, 1.0, 0.5)
        end

    end

end


export clenshaw_curtis

"""
    clenshaw_curtis(N)

Compute the nodes `x` and weights `w` for integrating a
continuous functions from [-1,1] using the Clenshaw-Curtis
integration rule with order `N`.
"""
function clenshaw_curtis(N::Int64)

    @assert N>0 "Order must be greater than 0"

    x = zeros(N)
    w = ones(N)

    if N==1
        #x[1]=0.0
        w[1] = 2.0
        return x, w
    end

    # Compute x values
    @inbounds for i=1:N
        x[i] = cos((N-i)*π/(N-1))
    end

    n = N
    x[1] = -1.0

    if mod(n,2)==1
        x[ceil(Int64, (n+1)/2)] = 0.0
    end

    x[n] = 1.0

    @inbounds for i=1:n
        θ = (i-1)*π/(n-1)
        for j=1:ceil(Int64,(n-1)/2)
            if 2*j == n-1
                b = 1.0
            else
                b = 2.0
            end

            w[i] -= b*cos(2*j*θ)/(4*j*j-1)
        end
    end

    w[1] *= 1/(n-1)
    w[2:n] .*= 2.0/(n-1)
    w[n] *= 1/(n-1)

    return x, w
end


# Description:
# Dimensions: 2

# The plot on the left shows the three-hump Camel function on its recommended input domain, and the plot on the right shows only a portion of this domain, to allow for easier viewing of the function's key characteristics. The function has three local minima.

# Input Domain:
# The function is usually evaluated on the square xi ∈ [-5, 5], for all i = 1, 2.

# Global Minimum:
# https://www.sfu.ca/~ssurjano/camel33.png

function threeHumpCamel(xx)

    x1 = xx[1];
    x2 = xx[2];

    term1 = 2 * x1^2;
    term2 = -1.05 * x1^4;
    term3 = x1^6 / 6;
    term4 = x1 * x2;
    term5 = x2^2;

    y = term1 + term2 + term3 + term4 + term5;

    return [y]
end


require("bspline.jl")

function exercise_zero()
    controlPoints = [
                     [-0.2 2];
                     [-0.3 6.2];
                     [-1.2 4.8];
                     [-2.8 8.8];
                     [-0.7 14];
                     [1.4 14.7];
                     [3.6 10.2];
                     [3.2 5.1];
                     [1.5 6.2];
                     [1.4 2];
                     ]

    ## continuityVector = ones(length(controlPoints[:,1]))
    k = 4
    extendedPartition = extendedPartitionAttempt(k, ones(6), Clumped())
    M = sum(ones(6))
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    fittingPoints = 200
    bspline = drawCurve(extendedPartition, k, M, controlPoints, matrix, fittingPoints)

    writeArrayForGnuplot(controlPoints,
                         "exercise-zero-clumped-control-poly.coordinates")

    writeArrayForGnuplot(bspline,
                         "exercise-zero-clumped-bspline.coordinates")

    extendedPartition = extendedPartitionAttempt(k, ones(6), Uniformed())
    M = sum(ones(6))
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    fittingPoints = 200
    bspline = drawCurve(extendedPartition, k, M, controlPoints, matrix, fittingPoints)

    writeArrayForGnuplot(controlPoints,
                         "exercise-zero-uniformed-control-poly.coordinates")

    writeArrayForGnuplot(bspline,
                         "exercise-zero-uniformed-bspline.coordinates")

    closedControlPoints = [controlPoints; controlPoints[1,:];]
    closedBSpline = drawBSpline(k, closedControlPoints, 200)
    writeArrayForGnuplot(closedControlPoints,
                         "exercise-zero-closed-control-poly.coordinates")

    writeArrayForGnuplot(closedBSpline,
                         "exercise-zero-closed-bspline.coordinates")

end

function exercise_one()
    controlPoints = [
                     [-2.5 -1];
                     [-3 2];
                     [0.3 2.0];
                     [-0.3 -2.0];
                     [4 -2.0];
                     [3 1.0];
                     ]

    writeArrayForGnuplot(controlPoints,
                         "exercise-one-control-poly.coordinates")

    fittingPoints = 200

    k = 2
    continuityVector = ones(4)
    k_two_extendedPartition = extendedPartitionAttempt(k, continuityVector, Clumped())
    M = sum(continuityVector)
    matrix = buildFunctionsMatrix(k, M, k_two_extendedPartition)
    k_two_bspline = drawCurve(k_two_extendedPartition,
                              k,
                              M,
                              controlPoints,
                              matrix,
                              fittingPoints)
    writeArrayForGnuplot(k_two_bspline,
                         "exercise-one-clumped-bspline-k-two.coordinates")

    k = 3
    continuityVector = ones(3)
    k_three_extendedPartition = extendedPartitionAttempt(k, continuityVector, Clumped())
    M = sum(continuityVector)
    matrix = buildFunctionsMatrix(k, M, k_three_extendedPartition)
    k_three_bspline = drawCurve(k_three_extendedPartition,
                              k,
                              M,
                              controlPoints,
                              matrix,
                              fittingPoints)
    writeArrayForGnuplot(k_three_bspline,
                         "exercise-one-clumped-bspline-k-three.coordinates")

    k = 4
    continuityVector = ones(2)
    k_four_extendedPartition = extendedPartitionAttempt(k, continuityVector, Clumped())
    M = sum(continuityVector)
    matrix = buildFunctionsMatrix(k, M, k_four_extendedPartition)
    k_four_bspline = drawCurve(k_four_extendedPartition,
                              k,
                              M,
                              controlPoints,
                              matrix,
                              fittingPoints)
    writeArrayForGnuplot(k_four_bspline,
                         "exercise-one-clumped-bspline-k-four.coordinates")

    k = 5
    continuityVector = ones(1)
    k_five_extendedPartition = extendedPartitionAttempt(k, continuityVector, Clumped())
    M = sum(continuityVector)
    matrix = buildFunctionsMatrix(k, M, k_five_extendedPartition)
    k_five_bspline = drawCurve(k_five_extendedPartition,
                              k,
                              M,
                              controlPoints,
                              matrix,
                              fittingPoints)
    writeArrayForGnuplot(k_five_bspline,
                         "exercise-one-clumped-bspline-k-five.coordinates")

    k = 6
    continuityVector = ones(0)
    k_six_extendedPartition = extendedPartitionAttempt(k, continuityVector, Clumped())
    M = sum(continuityVector)
    matrix = buildFunctionsMatrix(k, M, k_six_extendedPartition)
    k_six_bspline = drawCurve(k_six_extendedPartition,
                              k,
                              M,
                              controlPoints,
                              matrix,
                              fittingPoints)
    writeArrayForGnuplot(k_six_bspline,
                         "exercise-one-clumped-bspline-k-six.coordinates")

end

function exercise_two()
    controlPoints = [
                     [-0.7 -0.4];
                     [-1 -1];
                     [-1.3 -0.3];
                     [-1 1.2];
                     [-0.3 1.4];
                     [-0.15 0];
                     [0.15 0];
                     [0.3 1.6];
                     [2 0.6];
                     [2.6 0.2];
                     ]

    writeArrayForGnuplot(controlPoints,
                         "exercise-two-control-poly.coordinates")

    k = 4
    extendedPartition = expandPartition(0:10, [4 [1 for _ in 1:10]'])
    extendedPartition = normalizeRespect(extendedPartition, Maximum(), k)
    M = sum(ones(6))
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    fittingPoints = 200
    bspline = drawCurve(extendedPartition,
                        k,
                        M,
                        controlPoints,
                        matrix,
                        fittingPoints)

    writeArrayForGnuplot(bspline,
                         "exercise-two-bspline-first.coordinates")

    k = 4
    extendedPartition = expandPartition(0:4, [4 2 2 2 4])
    extendedPartition = normalizeRespect(extendedPartition, Maximum())
    M = sum(ones(6))
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    fittingPoints = 200
    bspline = drawCurve(extendedPartition, k, M, controlPoints, matrix, fittingPoints)

    writeArrayForGnuplot(bspline,
                         "exercise-two-bspline-second.coordinates")

    k = 6
    extendedPartition = expandPartition(0:5, [6 1 1 1 1 6])
    extendedPartition = normalizeRespect(extendedPartition, Maximum())
    M = sum(ones(4))
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    fittingPoints = 200
    bspline = drawCurve(extendedPartition, k, M, controlPoints, matrix, fittingPoints)

    writeArrayForGnuplot(bspline,
                         "exercise-two-bspline-three.coordinates")

    k = 6
    extendedPartition = expandPartition(0:4, [6 1 1 2 6])
    extendedPartition = normalizeRespect(extendedPartition, Maximum())
    M = sum(ones(4))
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    fittingPoints = 200
    bspline = drawCurve(extendedPartition, k, M, controlPoints, matrix, fittingPoints)

    writeArrayForGnuplot(bspline,
                         "exercise-two-bspline-four.coordinates")

    k = 8
    extendedPartition = expandPartition(0:3, [8 1 1 8])
    extendedPartition = normalizeRespect(extendedPartition, Maximum())
    M = sum(ones(2))
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    fittingPoints = 200
    bspline = drawCurve(extendedPartition, k, M, controlPoints, matrix, fittingPoints)

    writeArrayForGnuplot(bspline,
                         "exercise-two-bspline-five.coordinates")

end

function exercise_three()
    controlPoints = [
                     [-3.0 2];
                     [-1.0 -2];
                     [1 1.0];
                     [1 1.0];
                     [3.0 -2];
                     [5 3.0];
                     ]

    writeArrayForGnuplot(controlPoints,
                         "exercise-three-control-poly.coordinates")

    fittingPoints = 200

    k = 3
    continuityVector = ones(3)
    extendedPartition = extendedPartitionAttempt(k, continuityVector, Clumped())
    M = sum(continuityVector)
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    bspline = drawCurve(extendedPartition,
                        k,
                        M,
                        controlPoints,
                        matrix,
                        fittingPoints)
    writeArrayForGnuplot(bspline,
                         "exercise-three-bspline.coordinates")
end

function exercise_four()
    controlPoints = [
                     [-2.0 2];
                     [-3.0 4];
                     [0 6];
                     [0 6];
                     [3 4];
                     [2.0 2];
                     ]

    writeArrayForGnuplot(controlPoints,
                         "exercise-four-control-poly.coordinates")

    fittingPoints = 200

    k = 4
    ## continuityVector = ones(3)
    ## extendedPartition = extendedPartitionAttempt(k, continuityVector, Clumped())
    extendedPartition = [0 0 0 0 1/4 3/4 1 1 1 1]
    M = 2
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    bspline = drawCurve(extendedPartition,
                        k,
                        M,
                        controlPoints,
                        matrix,
                        fittingPoints)
    writeArrayForGnuplot(bspline,
                         "exercise-four-bspline-first-partition.coordinates")

    extendedPartition = [0 0 0 0 1/2 1/2 1 1 1 1]
    M = 2
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    bspline = drawCurve(extendedPartition,
                        k,
                        M,
                        controlPoints,
                        matrix,
                        fittingPoints)
    writeArrayForGnuplot(bspline,
                         "exercise-four-bspline-second-partition.coordinates")

end

function exercise_five()
    controlPoints = [
                     [-2.0 1];
                     [-1.0 -1];
                     [0 1];
                     [1 -1];
                     [2.0 1];
                     ]

    writeArrayForGnuplot(controlPoints,
                         "exercise-five-control-poly.coordinates")

    fittingPoints = 200

    k = 4
    continuityVector = ones(1)
    extendedPartition = extendedPartitionAttempt(k, continuityVector, Clumped())
    M = sum(continuityVector)
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    bspline = drawCurve(extendedPartition,
                        k,
                        M,
                        controlPoints,
                        matrix,
                        fittingPoints)
    writeArrayForGnuplot(bspline,
                         "exercise-five-bspline-single.coordinates")

    controlPoints = [
                     [-2.0 1];
                     [-1.0 -1];
                     [0 1];
                     [0 1];
                     [1 -1];
                     [2.0 1];
                     ]

    fittingPoints = 200

    k = 4
    continuityVector = ones(2)
    extendedPartition = extendedPartitionAttempt(k, continuityVector, Clumped())
    M = sum(continuityVector)
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    bspline = drawCurve(extendedPartition,
                        k,
                        M,
                        controlPoints,
                        matrix,
                        fittingPoints)
    writeArrayForGnuplot(bspline,
                         "exercise-five-bspline-doubled.coordinates")

    controlPoints = [
                     [-2.0 1];
                     [-1.0 -1];
                     [0 1];
                     [0 1];
                     [0 1];
                     [1 -1];
                     [2.0 1];
                     ]

    fittingPoints = 200

    k = 4
    continuityVector = ones(3)
    extendedPartition = extendedPartitionAttempt(k, continuityVector, Clumped())
    M = sum(continuityVector)
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    bspline = drawCurve(extendedPartition,
                        k,
                        M,
                        controlPoints,
                        matrix,
                        fittingPoints)
    writeArrayForGnuplot(bspline,
                         "exercise-five-bspline-tripled.coordinates")

end

function exercise_six()

    k = 4

    firstClosedControlPoints = [
                                [1.0 4];
                                [-0.8 4];
                                [-1.9 8.2];
                                [0.3 13.2];
                                [2.8 4];
                                [1.0 4];
                                ]

    firstClosedBSpline = drawBSpline(k, firstClosedControlPoints, 200)
    writeArrayForGnuplot(firstClosedControlPoints,
                         "exercise-six-first-closed-control-poly.coordinates")

    writeArrayForGnuplot(firstClosedBSpline,
                         "exercise-six-first-closed-bspline.coordinates")

    secondClosedControlPoints = [
                                 [0.4 6];
                                 [-1.8 3];
                                 [-2 14.8];
                                 [0.4 11];
                                 [3.0 14];
                                 [3.4 2.2];
                                 [0.4 6];
                                 ]

    secondClosedBSpline = drawBSpline(k, secondClosedControlPoints, 200)
    writeArrayForGnuplot(secondClosedControlPoints,
                         "exercise-six-second-closed-control-poly.coordinates")

    writeArrayForGnuplot(secondClosedBSpline,
                         "exercise-six-second-closed-bspline.coordinates")

end

function exercise_knotsInsertion()
    controlPoints = [
                     [0.0 0.0];
                     [1.0 2.0];
                     [2.0 2.5];
                     [3.0 2.2];
                     [4.0 0.5];
                     ]

    partition = [-1; 0.5; 1.8; 2.8; 3.2; 4;]

    k = 3
    refined_controlPoints =
        knotsInsertion_head(controlPoints, partition, k, 1)

    writeArrayForGnuplot(controlPoints,
                         "exercise-knots-insertion-original-control-points.coordinates")

    writeArrayForGnuplot(refined_controlPoints,
                         "exercise-knots-insertion.coordinates")

    ## deBoorBSpline = drawBSpline(k, controlPoints, 200)
    ## writeArrayForGnuplot(deBoorBSpline,
    ##                      "exercise-knots-insertion-deBoor.coordinates")

end


using QuadGK

function main()

    f(x) = x^2
    g(x) = x^4
    h(x) = (x+1)^(-1)

    a = 0
    b = 2
    n = 10

    println(simpsons(f, a, b, n))
    println(quadgk(f, a, b, rtol=1e-5))

    println(simpsons(g, a, b, n))
    println(quadgk(g, a, b, rtol=1e-5))

    println(simpsons(h, a, b, n))
    println(quadgk(h, a, b, rtol=1e-5))
end

function simpsons(df, a, b, n)

    # n should be even
    h = (b - a) / n

    xi0 = df(a) + df(b)         # end points
    xi1 = 0                     # odd points
    xi2 = 0                     # even points

    # performing integration
    for i in 1:n-1

        x = a + i*h

        # if even add to xi2
        if i%2 == 0
            xi2 += df(x)
        else
            # else odd add to xi1
            xi1 += df(x)
        end
    end

    # add everything togeteher
    xi = h*(xi0 + 2*xi2 + 4*xi1)/3

    return xi

end

main()


function eval(membership_function::TriangularMF, x::T) where {T <: Number}
    return maximum([minimum([((x - membership_function.l_vertex) / (membership_function.center - membership_function.l_vertex)), ((membership_function.r_vertex - x) / (membership_function.r_vertex - membership_function.center))]), 0])
end

function mean_at(membership_function::TriangularMF, firing_strength::T) where {T <: Number}
    if firing_strength != 1
        p1 = (membership_function.center - membership_function.l_vertex) * firing_strength + membership_function.l_vertex
        p2 = (membership_function.center - membership_function.r_vertex) * firing_strength + membership_function.r_vertex
        (p1 + p2) / 2
    elseif firing_strength == 1
        return membership_function.center
    end
end

function eval(menbership_function::GaussianMF, x::T) where {T <: Number}
    return exp(- 0.5 * ((x - menbership_function.center) / menbership_function.sigma)^2)
end

function mean_at(membership_function::GaussianMF, firing_strength::T) where {T <: Number}
    return membership_function.center
end

function eval(membership_function::BellMF, x::T) where {T <: Number}
    return (1 / (1 + abs((x - membership_function.c) / membership_function.a)^(2 * membership_function.b)))
end

function mean_at(membership_function::BellMF, firing_strength::T) where {T <: Number}
    return membership_function.c
end

function eval(membership_function::TrapezoidalMF, x::T) where {T <: Number}
    return  maximum([minimum([((x - membership_function.l_bottom_vertex) / (membership_function.l_top_vertex - membership_function.l_bottom_vertex)), 1, ((membership_function.r_bottom_vertex - x) / (membership_function.r_bottom_vertex - membership_function.r_top_vertex))]), 0])
end

function mean_at(membership_function::TrapezoidalMF, firing_strength::T) where {T <: Number}
    p1 = (membership_function.l_top_vertex - membership_function.l_bottom_vertex) * firing_strength + membership_function.l_bottom_vertex
    p2 = (membership_function.r_top_vertex - membership_function.r_bottom_vertex) * firing_strength + membership_function.r_bottom_vertex
    (p1 + p2) / 2
end

function eval(membership_function::SigmoidMF, x::T) where {T <: Number}
    return  1 / (1 + exp(-membership_function.a * (x - membership_function.c)))
end

function mean_at(membership_function::SigmoidMF, firing_strength::T) where {T <: Number}
    if firing_strength == 1
        p_firing_strength = 0.999
    elseif firing_strength == 0
        p_firing_strength = 0.001
    else
        p_firing_strength = firing_strength
    end

    p1 = -log((1 / p_firing_strength) - 1) / membership_function.a + membership_function.c
    p2 = membership_function.limit
    (p1 + p2) / 2
end


"""
   function lagdiff(x::AbstractVector, d=1)

Returns `out[i] = x[i] - x[i-d]` with out padded with missings so that `length(out)==length(x)`.
"""
function lagdiff(x::AbstractVector, d=1)
  if d>0
    dx = vcat(fill(missing, d), x[(d+1):end] .- x[1:(end-d)])
  elseif d<0
    dx = vcat(-x[(d+1):end] .+ x[1:(end-d)], fill(missing,d))
  else
    dx = zero(x)
  end
  return(dx)
end

"""
    function smooth(x::AbstractVector; w=pdf(Normal(), range(-3, 3, length=7)))

Returns the running weighted mean of x*w. When `n=length(w)`, then
`out[i] = sum(x[i - n÷2 + j-1] * w[j] for j in 1:n)/sum(w)`.
"""
function smooth(x::AbstractVector;  w=pdf(Normal(), range(-3, 3, length=7)))
  sx = Vector{Union{Missing, Float64}}(undef, length(x))
  n = length(w)
  if n % 2 != 1
    error("only odd length windows allowed")
  end
  shift = -(n÷2):(n÷2)
  for i in 1:length(sx)
    s = findfirst((i .+ shift) .> 0)
    l = findlast((i .+ shift) .< length(x))
    sx[i] = sum(w[s:l].*x[i .+ shift[s:l]])./sum(w[s:l])
    if isnan(sx[i])
      sx[i] = 0
    end
  end
  return(sx)
end


## Julia program for Atkinson's Poisson Simulation
## author: weiya <szcfweiya@gmail.com>
## date: 2018-08-19

function AtkinsonPois(lambda)
    # parameters
    beta = pi/sqrt(3*lambda)
    alpha = lambda*beta
    c = 0.767 - 3.36/lambda
    k = log(c) - lambda - log(beta)
    # step 1: propose new x
    u1 = rand()
    while true
        global x
        x = (alpha - log((1-u1)/u1))/beta
        x > -0.5 && break
    end
    while true
        # step 2: transform to N
        N = floor(Int, x)
        u2 = rand()
        # step 3: accept or not
        lhs = alpha - beta*x + log(u2/(1+exp(alpha-beta*x))^2)
        rhs = k + N*log(lambda) - log(factorial(lambda))
        if lhs <= rhs
            return(N)
        end
    end
end

## example
N = 100;
res = ones(Int, N);
for i = 1:N
    res[i] = AtkinsonPois(10)
end
# ans: 8 9 13 10 12 .......

# alternative method
function SimplePois(lambda)
    s = 0
    k = 0
    while true
        u = rand()
        x = -log(u)/lambda
        s = s + x
        if s > 1
            return(k)
        end
        k = k + 1
    end
end

## example for simple poisson
res2 = ones(Int, N);
for i = 1:N
    res2[i] = SimplePois(10)
end
# ans: 5 7 16 7 10 .......


using Plots;gr()

f(x) = x*x - 2

function bisection(a, b, max_iter)
    c = 1e+9
    error = zeros(max_iter)
    for i in 1:max_iter
        c = (a + b) / 2
        error[i] = abs(c - sqrt(2))
        f_c = f(c)

        if (f_c == 0)
            break
        elseif (f_c < 0)
            a = c
        else
            b = c
        end
    end

    (c, error)
end

a = 1.0
b = 2.0
max_iter = 20
c, error = bisection(a, b, max_iter)
t = 1:max_iter

plot(t, error, color=:cornflowerblue, label="bisection",xlabel="t",ylabel="Error")
savefig("bisection.png")


#==============================================================================

          Solving the HJB with a state constraint using
                the explicit method

               Translated from matlab code on Ben Moll's website:
                              http://www.princeton.edu/~moll/HACTproject.htm

==============================================================================#

using LinearAlgebra, SparseArrays, Plots

γ = 2 # parameter from CRRA utility
r = 0.03 # the interest rate
ρ = 0.05 # the discount rate

z1=.1
z2 =.2
z= [z1 z2]

λ1 = 0.02
λ2 = 0.03

λ = [λ1 λ2]

H = 500 # number of points in the grid space
amin=-0.02
amax=2

a = LinRange(amin,amax,H)
da = (amax-amin)/(H-1)

aa = [a a]
zz = ones(H,1)*z

maxit = 100 # the maximum number of iterations we allow the finite differencing algorithm
crit = 10^(-6) # our critical value
Δ = 1000

dVf=zeros(H,2)
dVb = zeros(H,2)
c= zeros(H,2)

# Create the matrix that captures the systems dynamic
Aswitch = [-sparse(I,H,H)*λ[1] sparse(I,H,H)*λ[1]; sparse(I,H,H)*λ[2] -sparse(I,H,H)*λ[2]]

# Inital Guess
v0=zeros(H,2)
v0[:,1] = (z[1] .+ r.*a).^(1-γ)/(1-γ)/ρ
v0[:,2] = (z[2] .+ r.*a).^(1-γ)/(1-γ)/ρ
global v=v0

dist=[]
V_n=[]

# The finite differeing loop

for n in 1:maxit
    global V=v
    push!(V_n,V)
    # Forward differencing
    dVf[1:H-1,:]=(V[2:H,:]-V[1:H-1,:])/da
    dVf[H,:] = (z .+ r.*amax).^(-γ) # impose state constraint at the max, just in case
    # Backward differencing
    dVb[2:H,:]=(V[2:H,:]-V[1:H-1,:])/da
    dVb[1,:]= (z .+r.*amin).^(-γ)

    # Find consumption and savings with the forward difference
    cf = dVf.^(-1/γ)
    sf = zz +r.*aa -cf

    # Find consumption and savings with the backward difference
    cb= dVb.^(-1/γ)
    sb = zz +r.*aa -cb

    # Find consumption and savings at steady state
    c0 = zz + r.*aa
    dV0 = c0.^(-γ)

    # Now implement the upwind scheme in order to select the best differencing method
    If = sf .> 0
    Ib = sb .< 0
    I0 = (1 .- If - Ib)

    # State constraint at amin is automatically implemented

    global dV_Upwind = dVf.*If + dVb.*Ib + dV0.*I0

    global c = dV_Upwind.^(-1/γ)
    global u = c.^(1-γ)/(1-γ)
    # Construct matrix for the evolution of the system
    X = -min.(sb,0)/da
    Y= -max.(sf,0)/da + min.(sb,0)/da
    Z = max.(sf,0)/da

    A1 = spdiagm(0=>Y[:,1], -1 => X[2:H,1], 1=> Z[1:H-1,1])
    A2 = spdiagm(0=>Y[:,2], -1 => X[2:H,2], 1=> Z[1:H-1,2])
    AA = [A1 spzeros(H,H); spzeros(H,H) A2]

    A = AA + Aswitch

    B = (ρ +1/Δ)*sparse(I,2*H,2*H) -A
    u_stacked = [u[:,1]; u[:,2]]
    V_stacked = [V[:,1]; V[:,2]]

    b = u_stacked + V_stacked/Δ
    V_stacked = B\b # Solves the system of equations

    V = [V_stacked[1:H] V_stacked[H+1:2*H]]

    V_change = V - v

    global v = V

    push!(dist, findmax(abs.(V_change))[1])

    if dist[n] < crit
        println("Value Function converged, Iteration=")
        println(n)
        break
    end
end

# Graphs
adot = zz + r.*aa -c

plot(dist[:], xlabel="Iteration", ylabel="\$\\lvert V^{n+1}-V^{n}\\rvert\$")

plot(a,V, xlabel="a", ylabel="\$V_{i}(a)\$", xlims=(amin,amax),legend=false)

plot(a,c,xlabel="a", ylabel="\$c_{i}(a)\$", xlims=(amin,amax), legend=false)

plot(a, adot, xlabel="a", ylabel="\$s_{i}(a)\$", xlims=(amin,amax))


# # [Sparse linear regression](@id sparse_linreg)
#
# Let's look at a least squares regression problem with L1 regularization:
# we will use the "diabetes dataset" (see [here](https://www4.stat.ncsu.edu/~boos/var.select/)), so let's start by loading the data.

using HTTP

splitlines(s) = split(s, "\n")
splitfields(s) = split(s, "\t")
parsefloat64(s) = parse(Float64, s)

function load_diabetes_dataset()
    res = HTTP.request("GET", "https://www4.stat.ncsu.edu/~boos/var.select/diabetes.tab.txt")
    lines = res.body |> String |> strip |> splitlines
    return hcat((line |> splitfields .|> parsefloat64 for line in lines[2:end])...)'
end

data = load_diabetes_dataset()

training_input = data[1:end-100, 1:end-1]
training_label = data[1:end-100, end]

test_input = data[end-99:end, 1:end-1]
test_label = data[end-99:end, end]

n_training, n_features = size(training_input)

# Now we can set up the optimization problem we want to solve: we will minimize the mean squared error
# for a linear model, appropriately scaled so that the features in the training data are normally distributed
# ("standardization", this is known to help the optimization process).
#
# After some simple manipulation, this standardized linear model can be implemented as follows:

using LinearAlgebra
using Statistics

input_loc = mean(training_input, dims=1) |> vec
input_scale = std(training_input, dims=1) |> vec

linear_model(wb, input) = input * wb[1:end-1] .+ wb[end]

function standardized_linear_model(wb, input)
    w_scaled = wb[1:end-1] ./ input_scale
    wb_scaled = vcat(w_scaled, wb[end] - dot(w_scaled, input_loc))
    return linear_model(wb_scaled, input)
end

# The loss term in the cost is then the following. Note that this is a regular Julia function:
# since the algorithm we will apply requires its gradient, automatic differentiation will
# do the work for us.

mean_squared_error(label, output) = mean((output .- label) .^ 2) / 2

training_loss(wb) = mean_squared_error(training_label, standardized_linear_model(wb, training_input))

# As regularization we will use the L1 norm, implemented in [ProximalOperators](https://github.com/JuliaFirstOrder/ProximalOperators.jl):

using ProximalOperators

reg = ProximalOperators.NormL1(1)

# We want to minimize the sum of `training_loss` and `reg`, and for this task we can use `FastForwardBackward`,
# which implements the fast proximal gradient method (also known as fast forward-backward splitting, or FISTA).
# Therefore we construct the algorithm, then apply it to our problem by providing a starting point,
# and the objective terms `f=training_loss` (smooth) and `g=reg` (non smooth).

using ProximalAlgorithms

ffb = ProximalAlgorithms.FastForwardBackward()
solution, iterations = ffb(x0=zeros(n_features + 1), f=training_loss, g=reg)

# We can now check how well the trained model performs on the test portion of our data.

test_output = standardized_linear_model(solution, test_input)
mean_squared_error(test_label, test_output)


#=

Plot the price-dividend ratio in the risk neutral case, for the Markov asset pricing
lecture.

=#
using QuantEcon
using Plots
using LaTeXStrings

pyplot()

n = 25  # size of state space
beta = 0.9
mc = tauchen(n, 0.96, 0.02)

K = mc.p .* reshape(exp(mc.state_values), 1, n)

I = eye(n)
v = (I - beta * K) \  (beta * K * ones(n, 1))

plot(mc.state_values,
    v,
    lw=2,
    ylabel="price-dividend ratio",
    xlabel="state",
    alpha=0.7,
    label=L"$v$")


function agm(x, y, e::Real = 5)
    (x ≤ 0 || y ≤ 0 || e ≤ 0) && throw(DomainError("x, y must be strictly positive"))
    g, a = minmax(x, y)
    while e * eps(x) < a - g
        a, g = (a + g) / 2, sqrt(a * g)
    end
    a
end

x, y = 1.0, 1 / √2
println("# Using literal-precision float numbers:")
@show agm(x, y)

println("# Using half-precision float numbers:")
x, y = Float32(x), Float32(y)
@show agm(x, y)

println("# Using ", precision(BigFloat), "-bit float numbers:")
x, y = big(1.0), 1 / √big(2.0)
@show agm(x, y)


"""
$(SIGNATURES)
"""
average(m::AbstractMatrix) = (sum(m) / length(m))

"""
$(SIGNATURES)
"""
average_absolute(m::AbstractMatrix) = (sum(abs.(m)) / length(m))

"""
$(SIGNATURES)
"""
test_average_absolute(m::AbstractMatrix) = (average_absolute(m)) < 1e6

"""
$(SIGNATURES)
"""
function smallest_positive(x1::Float64, x2::Float64)
    (x1 > 0.0 && x2 < 0.0) && return x1
    (x1 < 0.0 && x2 > 0.0) && return x2
    (x1 > 0.0 && x2 > 0.0) && return min(x1, x2)
    return nothing
end

"""
$(SIGNATURES)

Approximates the gradient of a scalar field.
"""
function ∇(ϕ::AbstractMatrix{Float64})
    height, width = size(ϕ)

    # x, y only represents the first and the second variables. No reference to graphical representations.
    ∇ϕx = zeros(Float64, height, width)   # divergence on the right edge will be filled with zeros
    ∇ϕy = zeros(Float64, height, width)   # divergence on bottom edge will be filled with zeros

    ∇ϕx[begin:end-1, :] = ϕ[begin+1:end, :] - ϕ[begin:end-1, :]
    ∇ϕy[:, begin:end-1] = ϕ[:, begin+1:end] - ϕ[:, begin:end-1]

    fill_borders!(∇ϕx, 0.0)
    fill_borders!(∇ϕy, 0.0)

    return ∇ϕx, ∇ϕy
end

"""
$(SIGNATURES)

Calculate the second-order Laplace operator by convolution of a kernel.
"""
function laplacian(ϕ::AbstractMatrix{Float64})
    height, width = size(ϕ)

    # Embed matrix within a larger matrix for better vectorization and avoid duplicated code
    # Padded matrix adds 1 row/col after the size of ϕ.
    # ϕ is inserted in padded matrix within  1:1+height+1 x 1:1+width+1.
    # The rest of the padded matrix (its borders) are set at 0.0.
    padded_ϕ = zeros(Float64, 1 + height + 1, 1 + width + 1)
    padded_ϕ[begin+1:end-1, begin+1:end-1] .= ϕ[1:end, 1:end]

    # Convolution = up + down + left + right
    ∇²ϕ =
        padded_ϕ[begin+1-1:end-1-1, begin+1:end-1] +
        padded_ϕ[begin+1+1:end-1+1, begin+1:end-1] +
        padded_ϕ[begin+1:end-1, begin+1-1:end-1-1] +
        padded_ϕ[begin+1:end-1, begin+1+1:end-1+1] -
        4.0 * padded_ϕ[begin+1:end-1, begin+1:end-1]

    return ∇²ϕ[1:end-1, 1:end-1]
end


using Downhill

"""
    rosenbrock(x; b=2)

Return the value of Rosenbrock function
    ```math
    \\sum_{i=1}^n \\right[(1 - x_{i-1})^2 + b (x_i - x_{i-1})^2 \\left]
    ```
    in `n = length(x)` dimensions.
"""
function rosenbrock(x::AbstractVector{T}; b=2) where {T}
    f = zero(T) * zero(b) # type-generic zero
    for i in firstindex(x)+1:lastindex(x)
        f += (one(T) - x[i-1])^2 + b * (x[i] - x[i-1]^2)^2
    end
    return f
end

"""
    drosenbrock(x; b=2)

Return the gradient of Rosenbrock function
    ```math
    \\sum_{i=1}^n \\right[(1 - x_{i-1})^2 + b (x_i - x_{i-1})^2 \\left]
    ```
    in `n = length(x)` dimensions.
"""
function drosenbrock(x::AbstractVector{T}; b=2) where {T}
    g = fill!(similar(x, promote_type(T, typeof(b))), 0)
    for i in firstindex(x)+1:lastindex(x)
        g[i-1] += 2 * (x[i-1] - 1) + 4 * b * x[i-1] * (x[i-1]^2 - x[i])
        g[i] += 2 * b * (x[i] - x[i-1]^2)
    end
    return g
end

"""
    rosenbrock!(x, g; b=2)

Return the value of Rosenbrock function in `length(x)` dimensions.
"""
function rosenbrock!(x::AbstractVector, g::AbstractVector; b=2)
    f = zero(eltype(g))
    fill!(g, 0)
    inds = eachindex(x, g)
    for i in 2:last(inds)
        f += (1 - x[i-1])^2 + b * (x[i] - x[i-1]^2)^2
        g[i-1] += 2 * (x[i-1] - 1) + 4 * b * x[i-1] * (x[i-1]^2 - x[i])
        g[i] += 2 * b * (x[i] - x[i-1]^2)
    end
    return f, g
end

ros(x, g) = (rosenbrock(x), drosenbrock(x))

ros!(x, g) = rosenbrock!(x, g)

ans_nonmutating = let x0 = zeros(2)
    opt = BFGS(x0)
    optresult = optimize!(ros, opt, x0; maxiter=1000)
    println("""
        ==Nonmutating gradient evaluation==

        Optimization converged: $(optresult.converged)
        Final argument: $(optresult.argument)
        Final gradient: $(optresult.gradient)
        Iteration count: $(optresult.iterations)
        Call count: $(optresult.calls)
    """)
    optresult.argument
end

ans_mutating = let x0 = zeros(2)
    opt = BFGS(x0)
    optresult = optimize!(ros!, opt, x0; maxiter=1000)
    println("""
        ==Mutating gradient evaluation==

        Optimization converged: $(optresult.converged)
        Final argument: $(optresult.argument)
        Final gradient: $(optresult.gradient)
        Iteration count: $(optresult.iterations)
        Call count: $(optresult.calls)
    """)
    optresult.argument
end

@assert ans_nonmutating ≈ ans_mutating


"""
   HermiteWeight()

is a quasi-vector representing `exp(-x^2)` on ℝ.
"""
struct HermiteWeight{T} <: Weight{T} end

HermiteWeight() = HermiteWeight{Float64}()
axes(::HermiteWeight{T}) where T = (Inclusion(ℝ),)
function getindex(w::HermiteWeight, x::Number)
    x ∈ axes(w,1) || throw(BoundsError())
    exp(-x^2)
end

sum(::HermiteWeight{T}) where T = sqrt(convert(T, π))

struct Hermite{T} <: OrthogonalPolynomial{T} end
Hermite() = Hermite{Float64}()
orthogonalityweight(::Hermite{T}) where T = HermiteWeight{T}()

==(::Hermite, ::Hermite) = true
axes(::Hermite{T}) where T = (Inclusion(ℝ), oneto(∞))

"""
     hermiteh(n, z)

computes the `n`-th Hermite polynomial, orthogonal with
respec to `exp(-x^2)`, at `z`.
"""
hermiteh(n::Integer, z::Number) = Base.unsafe_getindex(Hermite{typeof(z)}(), z, n+1)

broadcasted(::LazyQuasiArrayStyle{2}, ::typeof(*), ::HermiteWeight{T}, ::Hermite{V}) where {T,V} = Weighted(Hermite{promote_type(T,V)}())

# H_{n+1} = 2x H_n - 2n H_{n-1}
# 1/2 * H_{n+1} + n H_{n-1} = x H_n
# x*[H_0 H_1 H_2 …] = [H_0 H_1 H_2 …] * [0    1; 1/2  0     2; 1/2   0  3; …]
jacobimatrix(H::Hermite{T}) where T = Tridiagonal(Fill(one(T)/2,∞), Zeros{T}(∞), one(T):∞)
recurrencecoefficients(H::Hermite{T}) where T = Fill{T}(2,∞), Zeros{T}(∞), zero(T):2:∞

massmatrix(::Hermite{T}) where T = Diagonal(sqrt(convert(T,π)) .* convert(T,2) .^ (0:∞) .* factorial.(convert(T,0):∞))

@simplify *(A::QuasiAdjoint{<:Any,<:Weighted{<:Any,<:Hermite}}, B::Hermite) = massmatrix(Hermite{promote_type(eltype(A),eltype(B))}())

##########
# Derivatives
##########

@simplify function *(D::Derivative, H::Hermite)
    T = promote_type(eltype(D),eltype(H))
    D = _BandedMatrix((zero(T):2:∞)', ℵ₀, -1,1)
    H*D
end

@simplify function *(D::Derivative, Q::OrthonormalWeighted{<:Any,<:Hermite})
    X = jacobimatrix(Q.P)
    Q * Tridiagonal(-X.ev, X.dv, X.ev)
end


#nb # %% A slide [markdown] {"slideshow": {"slide_type": "subslide"}}
# ## Symmetric Strip Grating
# This example consists of a symmetric strip grating, i.e. a grating where the strip width
# is half the unit cell period ``P``:
#
# ![diagram](./assets/symmetric_strip_diagram.png)
#
# Only three of the infinite number of strips in the grating are shown, and they extend infinitely to the left and right.
# The grating lies in the ``z=0`` plane with free space on both sides. The shaded areas represent metalization.
# The dashed lines show two possible choices for the unit cell location: "J" for a formulation in terms of electric
# surface currents, and "M" for magnetic surface currents.
#
# For normal incidence there is a closed-form solution due to Weinstein,
# but for a more recent reference one can find the solution in Problem 10.6 of R. E. Collin,
# *Field Theory of Guided Waves, Second Ed.*,
# IEEE Press, 1991.  Here is the code for computing the exact solution:
"""
    grating(kP, nterms=30) -> (Γ, T)

Compute the normal incidence reflecton and transmission coefficients of a symmetric grid of
zero-thickness conducting strips.  The product of the period of the strips and the incident
electric field wavenumber is `kP` (dimensionless).  The incident electric field is
perpendicular to the direction along the axis of the strips.  The series have been
accelerated by applying a Kummer's transformation, using the first two terms in the Maclaurin
series for the inverse sin function.  `kP` must be in the half-open interval [0,1). The
default number of summed terms `nterms` yields better than 10 digits of accuracy over the
interval [0.01,0.99].
"""
function grating(kP; nterms=30)
    sum1 = 1.3862943611198906 # \sum_{n=1}^{\infty} 1/(n-1/2) - 1/n = log(4)
    sum3 = 7.2123414189575710 # \sum_{n=1}^{\infty} (n-1/2)^{-3} - n^{-3} = 6 * \zeta(3)
    x = kP/(4π)
    θ = x*sum1 + x^3/6 * sum3
    for n = 1:nterms
        xonmhalf = x/(n - 0.5)
        xon = x/n
        term = asin(xonmhalf) - (xonmhalf + (xonmhalf)^3/6) -
              (asin(xon) - (xon + xon^3/6))
        θ += term
    end
    Γ = sin(θ) * cis(-π/2 - θ)
    T = 1 + Γ
    return (Γ, T)
end

#  Note that using the extension of
# [Babinet's Principle for electromagnetic fields](http://kirkmcd.princeton.edu/examples/babinet.pdf)
# this also provides the solution (upon appropriate interchange and sign change of the coefficients) for
# the case where the incident wave polarization is parallel to the direction of the strips.

# Here is the PSSFSS code to analyze this structure using electric currents as the unknowns.  We will
# scale the geometry so that the frequency in GHz is numerically equal to the period of the strips
# measured in wavelengths.

using Plots, PSSFSS
c = 11.802852677165355 # light speed [inch*GHz]
period = c  # so the period/wavelength = freq in GHz
Py = period
Ly = period/2
Px = Lx = Ly/10 # Infinite in x direction so this can be anything
Ny = 60
Nx = round(Int, Ny*Lx/Ly)
sheet = rectstrip(;Px, Py, Lx, Ly, Nx, Ny, units=inch)
flist = 0.02:0.02:0.98
steering = (θ=0, ϕ=0)
strata = [Layer()
          sheet
          Layer()]
results_j = analyze(strata, flist, steering, showprogress=false,
                    resultfile=devnull, logfile=devnull);
p1 = plot(sheet)
p2 = plot(sheet, unitcell=true)
ptitle = plot(title = "Symmetric Strip Triangulation",
             grid = false, showaxis = false, xtick=[], ytick=[],
             bottom_margin = -50Plots.px)
plot(ptitle, p1, p2, layout = @layout([A{0.09h}; [B C]]))
#md savefig("symstrip1.png"); nothing  # hide
#-
#md # ![](symstrip1.png)

# Note that setting `Lx = Px` causes the strip to fully occupy the x-extent
# of the unit cell.  PSSFSS automatically ensures that the triangle edges at these unit
# cell boundaries define basis functions that satisfy the Floquet (phase shift) boundary
# conditions, so that currents are free to flow across these unit cell boundaries.

# We can also analyze the same structure using magnetic currents in the areas free of
# metalization as the unknowns:
sheet = rectstrip(;class='M', Px, Py, Lx, Ly, Nx, Ny, units=inch)
strata = [Layer()
          sheet
          Layer()]
results_m = analyze(strata, flist, steering, showprogress=false,
                    resultfile=devnull, logfile=devnull);

#
# Each 50-frequency run of `analyze` takes about 14 seconds
# for this geometry of 720 triangles on my machine.
# More detailed timing information is available in the log file
# (which is omitted for generating this documentation).

# We will compare the PSSFSS results to the analytic solution:
## Generate exact results:
rt = grating.(2π*flist)
rperp_exact = first.(rt)
tperp_exact = last.(rt)
rpar_exact = -tperp_exact
tpar_exact = -rperp_exact;

# Obtain PSSFSS results for electric and magnetic currents:
outrequest = @outputs s11(v,v) s21(v,v) s11(h,h) s21(h,h)
rperp_j, tperp_j, rpar_j, tpar_j =
      collect.(eachcol(extract_result(results_j, outrequest)))
rperp_m, tperp_m, rpar_m, tpar_m =
      collect.(eachcol(extract_result(results_m, outrequest)));

# Generate the comparison plots:
angdeg(z) = rad2deg(angle(z)) # Convenience function

p1 = plot(title = "Perpendicular Reflection Magnitude",
          xlabel = "Period (wavelengths)",
          ylabel = "Coefficient Magnitude",
          legend=:topleft)
plot!(p1, flist, abs.(rperp_exact), ls=:dash, label="Exact")
plot!(p1, flist, abs.(rperp_j), label="PSSFSS J")
plot!(p1, flist, abs.(rperp_m), label="PSSFSS M")
#md savefig("symstrip2.png"); nothing  # hide
#-
#md # ![](symstrip2.png)
#-
p2 = plot(title = "Perpendicular Reflection Phase",
          xlabel = "Period (wavelengths)",
          ylabel = "Phase (deg)")
plot!(p2, flist, angdeg.(rperp_exact), ls=:dash, label="Exact")
plot!(p2, flist, angdeg.(rperp_j), label="PSSFSS J")
plot!(p2, flist, angdeg.(rperp_m), label="PSSFSS M")
#md savefig("symstrip3.png"); nothing  # hide
#-
#md # ![](symstrip3.png)
#-
p1 = plot(title = "Parallel Reflection Magnitude",
          xlabel = "Period (wavelengths)",
          ylabel = "Coefficient Magnitude")
plot!(p1, flist, abs.(rpar_exact), ls=:dash, label="Exact")
plot!(p1, flist, abs.(rpar_j), label="PSSFSS J")
plot!(p1, flist, abs.(rpar_m), label="PSSFSS M")
#md savefig("symstrip4.png"); nothing  # hide
#-
#md # ![](symstrip4.png)
#-
p2 = plot(title = "Parallel Reflection Phase",
          xlabel = "Period (wavelengths)",
          ylabel = "Phase (deg)")
plot!(p2, flist, angdeg.(rpar_exact), ls=:dash, label="Exact")
plot!(p2, flist, angdeg.(rpar_j), label="PSSFSS J")
plot!(p2, flist, angdeg.(rpar_m), label="PSSFSS M")
#md savefig("symstrip5.png"); nothing  # hide
#-
#md # ![](symstrip5.png)

# Now look at the transmission coefficients:
p1 = plot(title = "Perpendicular Transmission Magnitude",
          xlabel = "Period (wavelengths)",
          ylabel = "Coefficient Magnitude")
plot!(p1, flist, abs.(tperp_exact), ls=:dash, label="Exact")
plot!(p1, flist, abs.(tperp_j), label="PSSFSS J")
plot!(p1, flist, abs.(tperp_m), label="PSSFSS M")
#md savefig("symstrip6.png"); nothing  # hide
#-
#md # ![](symstrip6.png)
#-
p2 = plot(title = "Perpendicular Transmission Phase",
          xlabel = "Period (wavelengths)",
          ylabel = "Phase (deg)")
plot!(p2, flist, angdeg.(tperp_exact), ls=:dash, label="Exact")
plot!(p2, flist, angdeg.(tperp_j), label="PSSFSS J")
plot!(p2, flist, angdeg.(tperp_m), label="PSSFSS M")
#md savefig("symstrip7.png"); nothing  # hide
#-
#md # ![](symstrip7.png)
#-
p1 = plot(title = "Parallel Transmission Magnitude",
          xlabel = "Period (wavelengths)",
          ylabel = "Coefficient Magnitude", legend=:topleft)
plot!(p1, flist, abs.(tpar_exact), ls=:dash, label="Exact")
plot!(p1, flist, abs.(tpar_j), label="PSSFSS J")
plot!(p1, flist, abs.(tpar_m), label="PSSFSS M")
#md savefig("symstrip8.png"); nothing  # hide
#-
#md # ![](symstrip8.png)
#-
p2 = plot(title = "Parallel Transmission Phase",
          xlabel = "Period (wavelengths)",
          ylabel = "Phase (deg)")
plot!(p2, flist, angdeg.(tpar_exact), ls=:dash, label="Exact")
plot!(p2, flist, angdeg.(tpar_j), label="PSSFSS J")
plot!(p2, flist, angdeg.(tpar_m), label="PSSFSS M")
#md savefig("symstrip9.png"); nothing  # hide
#-
#md # ![](symstrip9.png)

# ### Conclusion
# Although good agreement is obtained, as expected the best agreement between
# all three results occurs for the lowest frequencies, where the triangles are
# smallest in terms of wavelength.  This emphasizes the fact that it is necessary for the
# user to check that enough triangles have been requested for good convergence
# over the frequency band of interest.  This example is an extremely demanding case
# in terms of bandwidth, as the ratio of maximum to minimum frequency here
# is ``0.98/0.02 = 49:1``


@testset "vector projections" begin
    n̂  = SVector{3,Float64}(0, 0, 1)
    v1 = SVector{3,Float64}(1, 0, 0)
    v2 = SVector{3,Float64}(0, 1, 1)

    @test vec_proj(v1, n̂) == 0.0 * n̂
    @test vec_proj(n̂, n̂) == 1.0 * n̂
    @test vec_proj(v2, n̂) == 1.0 * n̂

    @test vec_sub_vec_proj(v1, n̂) == v1
    @test vec_sub_vec_proj(n̂, n̂) == 0.0 * n̂
    @test vec_sub_vec_proj(v2, n̂) == SVector{3,Float64}(0, 1, 0)
end

@testset "a_dot_one_pad_b" begin
	a = SVector{4,Float64}(1.0, 2.0, 3.0, 4.0)
	b1 = SVector{3,Float64}(2.0, 0.0, 0.0)
	b2 = SVector{3,Float64}(1.0, 2.0, 0.0)
	b3 = SVector{3,Float64}(1.0, 2.0, 3.0)

	@test a_dot_one_pad_b(a, b1) == 6.0
	@test a_dot_one_pad_b(a, b2) == 9.0
	@test a_dot_one_pad_b(a, b3) == 18.0
end


using LabelledArrays

export verdes

"""
    eom_verdes(u, p, t) → SVector{3}

Equations of motion for a 3D system where the response X is a highly
nonlinear combination of Y and Z [1]. The forcings Y and Z involve
sines and cosines, respectively, and have different periods.

The implementation here allows for tuning the period of the
forcing signals.

# References
Verdes, P. F. "Assessing causality from multivariate time series." Physical Review E 72.2 (2005): 026222.
"""
function eom_verdes(u, p, t)
    x, y, z = (u...,)
    ωy, ωz, σx, σy, σz = (p...,)

    ηx = σx == 0 ? 0 : rand(Normal(0, σx))
    ηy = σy == 0 ? 0 : rand(Normal(0, σy))
    ηz = σz == 0 ? 0 : rand(Normal(0, σz))

    dx = y*(18y - 27y^2 + 10)/2 + z*(1-z) + ηx
    dy = (1 - cos((2*pi/ωy) * t))/2 + ηy
    dz = (1 - sin((2*pi/ωz) * t))/2 + ηz
    return SVector{3}(dx, dy, dz)
end

function verdes(u₀, ωy, ωz, σx, σy, σz)
    p = LVector(ωy = ωy, ωz = ωz, σx = σx, σy = σy, σz = σz)

    DiscreteDynamicalSystem(eom_verdes, u₀, p)
end

"""
    verdes(;u₀ = rand(3), ωy = 315, ωz = 80,
        σx = 0.0, σy = 0.0, σz = 0.0) → DiscreteDynamicalSystem

Intitialise a 3D system where the response X is a highly nonlinear combination
of Y and Z. The forcings Y and Z involve sines and cosines, and
have different periods, which controlled by `ωy` and `ωz`.

The equations of motion are

```math
\\begin{aligned}
x(t+1) &= \\dfrac{y(t)(18y(t) - 27y(t)^2 + 10)}{2} + z(t)(1-z(t)) + ηx \\
y(t+1) &= \\dfrac{(1 - \\dfrac{\\cos(2\\pi)}{\\omega y}t)}{2} + ηy \\
z(t+1) &= \\dfrac{(1 - \\dfrac{\\sin(2\\pi)}{\\omega z}t)}{2} + ηz
\\end{aligned}
```
where ηx, ηy, ηz is gaussian noise with mean 0 and standard deviation `σx`, `σy`
and `σz`.

# References
Verdes, P. F. "Assessing causality from multivariate time series." Physical
Review E 72.2 (2005): 026222.
"""
verdes(;u₀ = rand(3),
    ωy = 315, ωz = 80,
    σx = 0.01, σy = 0.01, σz = 0.01) =
    verdes(u₀, ωy, ωz, σx, σy, σz)


using Plots
using DifferentialEquations

@userplot LorenzPlot
@recipe function f(lp::LorenzPlot)
    x, y, z, i = lp.args
    n = length(x)
    inds = circshift(1:n, 1-i)
    seriesalpha --> range(0, 1, length=n)
    aspect_ratio --> 1
    label --> false
    x[inds], y[inds], z[inds]
end

function anim()

    function lorenz(du, u, p, t)

        du[1] = p[1]*(u[2] - u[1])
        du[2] = u[1]*(p[2] - u[3]) - u[2]
        du[3] = u[1]*u[2] - p[3]*u[3]

    end

    u0 = [1, 5, 10]
    tspan = (0, 100)
    p = (10, 28, 8/3)

    prob = ODEProblem(lorenz, u0, tspan, p, reltol=1e-10, abstol=1e-10)
    sol = solve(prob)
    n = length(sol.t)

    plt = plot3d(1, xaxis=("x", (-30, 30)), yaxis=("y", (-30, 30)), zaxis=("z", (0, 60)),
                 title="Lorenz Attractor", marker=2, markercolor="black", dpi=300)

    # anim = @animate for i in 1:n
    #     push!(plt, sol[1,i], sol[2,i], sol[3,i])
    # end
    # gif(anim, "de_anim.gif", fps=30)

    anim = @animate for i in 1:10:n
        lorenzplot(sol[1,:], sol[2,:], sol[3,:], i, c = :thermal)
    end
    gif(anim, "de_anim_60.gif", fps=60)

end

anim()

    # anim = @animate for i in 1:10:n
    #     lorenzplot(sol[1,:], sol[2,:], sol[3,:], i)
    # end
    # gif(anim, "de_anim60.gif", fps=60)


module test

using Roots,Interpolations,PyPlot,Optim
using CompEcon, NLsolve

alpha     = 0.65
beta      = 0.95
grid_max  = 2  # upper bound of capital grid
n         = 150  # number of grid points
N_iter    = 3000  # number of iterations
kgrid     = 1e-6:(grid_max-1e-6)/(n-1):grid_max  # equispaced grid
f(x) = x.^alpha  # defines the production function f(k)
tol = 1e-9

ab        = alpha * beta
c1        = (log(1 - ab) + log(ab) * ab / (1 - ab)) / (1 - beta)
c2        = alpha / (1 - ab)
v_star(k) = c1 .+ c2 .* log(k)
k_star(k) = ab * k.^alpha
c_star(k) = (1-ab) * k.^alpha
ufun(x) = log(x)

# Bellman Operator
# inputs
# `grid`: grid of values of state variable
# `v0`: current guess of value function

# output
# `v1`: next guess of value function
# `pol`: corresponding policy function

#takes a grid of state variables and computes the next iterate of the value function.
function bellman_operator(grid,v0)

    v1  = zeros(n)     # next guess
    pol = zeros(Int,n)     # policy function
    w   = zeros(n)   # temporary vector

    # loop over current states
    # current capital
    for (i,k) in enumerate(grid)

        # loop over all possible kprime choices
        for (iprime,kprime) in enumerate(grid)
            if f(k) - kprime < 0   #check for negative consumption
                w[iprime] = -Inf
            else
                w[iprime] = ufun(f(k) - kprime) + beta * v0[iprime]
            end
        end
        # find maximal choice
        v1[i], pol[i] = findmax(w)     # stores Value und policy (index of optimal choice)
    end
    return (v1,pol)   # return both value and policy function
end

# VFI iterator
#
## input
# `n`: number of grid points
# output
# `v_next`: tuple with value and policy functions after `n` iterations.
function VFI()
    v_init = zeros(n)     # initial guess
    for iter in 1:N_iter
        v_next = bellman_operator(kgrid,v_init)  # returns a tuple: (v1,pol)
        # check convergence
        if maxabs(v_init.-v_next[1]) < tol
            verrors = maxabs(v_next[1].-v_star(kgrid))
            perrors = maxabs(kgrid[v_next[2]].-k_star(kgrid))
            println("discrete VFI:")
            println("Found solution after $iter iterations")
            println("maximal value function error = $verrors")
            println("maximal policy function error = $perrors")
            return v_next
        elseif iter==N_iter
            warn("No solution found after $iter iterations")
            return v_next
        end
        v_init = v_next[1]  # update guess
    end
end

# plot
function plotVFI()
    v = VFI()
    figure("discrete VFI",figsize=(10,5))
    subplot(131)
    plot(kgrid,v[1],color="blue")
    plot(kgrid,v_star(kgrid),color="black")
    xlim(-0.1,grid_max)
    ylim(-50,-30)
    xlabel("k")
    ylabel("value")
    title("value function")

    subplot(132)
    plot(kgrid,kgrid[v[2]])
    plot(kgrid,k_star(kgrid),color="black")
    xlabel("k")
    title("policy function")

    subplot(133)
    plot(kgrid,kgrid[v[2]].-k_star(kgrid))
    title("policy function error")
end

function bellman_operator2(grid,v0)

    v1  = zeros(n)     # next guess
    pol = zeros(n)     # consumption policy function

    Interp = interpolate((collect(grid),), v0, Gridded(Linear()) )

    # loop over current states
    # of current capital
    for (i,k) in enumerate(grid)

        objective(c) = - (log(c) + beta * Interp[f(k) - c])
        # find max of ojbective between [0,k^alpha]
        res = optimize(objective, 1e-6, f(k))
        pol[i] = f(k) - res.minimizer
        v1[i] = -res.minimum
    end
    return (v1,pol)   # return both value and policy function
end

function VFI2()
    v_init = zeros(n)     # initial guess
    for iter in 1:N_iter
        v_next = bellman_operator2(kgrid,v_init)  # returns a tuple: (v1,pol)
        # check convergence
        if maxabs(v_init.-v_next[1]) < tol
            verrors = maxabs(v_next[1].-v_star(kgrid))
            perrors = maxabs(v_next[2].-k_star(kgrid))
            println("continuous VFI:")
            println("Found solution after $iter iterations")
            println("maximal value function error = $verrors")
            println("maximal policy function error = $perrors")
            return v_next
        elseif iter==N_iter
            warn("No solution found after $iter iterations")
            return v_next
        end
        v_init = v_next[1]  # update guess
    end
    return nothing
end

function plotVFI2()
    v = VFI2()
    figure("discrete VFI - continuous control",figsize=(10,5))
    subplot(131)
    plot(kgrid,v[1],color="blue")
    plot(kgrid,v_star(kgrid),color="black")
    xlim(-0.1,grid_max)
    ylim(-50,-30)
    xlabel("k")
    ylabel("value")
    title("value function")

    subplot(132)
    plot(kgrid,v[2])
    plot(kgrid,k_star(kgrid),color="black")
    xlabel("k")
    title("policy function")

    subplot(133)
    plot(kgrid,v[2].-k_star(kgrid))
    xlabel("k")
    title("policy function error")
end

function policy_iter(grid,c0,u_prime,f_prime)

    c1  = zeros(length(grid))     # next guess
    pol_fun = interpolate((collect(grid),), c0, Gridded(Linear()) )

    # loop over current states
    # of current capital
    for (i,k) in enumerate(grid)
        objective(c) = u_prime(c) - beta * u_prime(pol_fun[f(k)-c]) * f_prime(f(k)-c)
        c1[i] = fzero(objective, 1e-10, f(k)-1e-10)
    end
    return c1
end

uprime(x::Float64) = 1.0 / x
fprime(x::Float64) = alpha* x ^ (alpha-1)
function fprime(x::Vector)
    r = similar(x)
    for ix in eachindex(x)
        if x[ix] > 0
            r[ix] = alpha* x[ix]^(alpha-1)
        else
            r[ix] = 10000.0
        end
    end
    return r
end

function PFI()
    c_init = kgrid
    for iter in 1:N_iter
        c_next = policy_iter(kgrid,c_init,uprime,fprime)
        # check convergence
        if maxabs(c_init.-c_next) < tol
            perrors = maxabs(c_next.-c_star(kgrid))
            println("PFI:")
            println("Found solution after $iter iterations")
            println("max policy function error = $perrors")
            return c_next
        elseif iter==N_iter
            warn("No solution found after $iter iterations")
            return c_next
        end
        c_init = c_next  # update guess
    end
end
function plotPFI()
    v = PFI()
    figure("PFI")
    subplot(121)
    plot(kgrid,v)
    plot(kgrid,c_star(kgrid),color="black")
    xlabel("k")
    title("policy function")

    subplot(122)
    plot(kgrid,v.-c_star(kgrid))
    xlabel("k")
    title("policy function error")
end
#plotVFI()
#plotVFI2()
#plotPFI()

function proj2(n=25)

    a = 1e-6
    aa = 1e-2
    b = 1.0
    bb = 0.95
    basis = fundefn(:cheb,n,a,b)
    k = funnode(basis)[1]   # collocation points

    fs = f.(k)

    # system of equations: for each k, one line
    function ff(coef::Vector,result::Vector,b,coll_points,fs)
        # get cons function
        cons = funeval(coef,b,coll_points)[1]
        # println("coef= $coef")
        # println("cons = $cons")

        # put resulting residuals into result[:]
        # put into euler equation
        nextk = fs.-cons
        for i in eachindex(nextk)
            if nextk[i] < 0
                result[i] = -8000.0
            else
                result[i] = uprime(cons[i]) .- beta * uprime(nextk[i]) .* fprime(nextk[i])
            end
        end
        # println(result)
    end

    f_closure(x::Vector,r::Vector) = ff(x,r,basis,k,fs)

    res = nlsolve(f_closure,ones(n)*0.4)
    x = linspace(aa,bb,501)
    y = funeval(res.zero,basis,x)[1]
    plot(x,y)
    return res
end

function proj(n=25)

    alpha = 1.0
    eta   = 1.5
    a     = 0.1
    b     = 3.0
    basis = fundefn(:cheb,n,a,b)
    p     = funnode(basis)[1]   # collocation points

    c0 = ones(n)*0.3
    function resid!(c::Vector,result::Vector,p,basis,alpha,eta)
        # your turn!
        q = funeval(c,basis,p)[1]
        q2 = zeros(q)
        for i in eachindex(q2)
            if q[i] < 0
                q2[i] = -20.0
            else
                q2[i] = sqrt(q[i])
            end
        end
        result[:] = p.+ q .*((-1/eta)*p.^(eta+1)) .- alpha*q2 .- q.^2
    end
    f_closure(r::Vector,x::Vector) = resid!(x,r,p,basis,alpha,eta)
    res = nlsolve(f_closure,c0)
    println(res)

    # plot residual function
    x = collect(linspace(a,b,501))
    y = similar(x)
    resid!(res.zero,y,x,basis,alpha,eta);
    y = funeval(res.zero,basis,x)[1]
    pl = Any[]
    push!(pl,plot(x,y,title="residual function"))

    # plot supply functions at levels 1,10,20

    # plot demand function

    y = funeval(res.zero,basis,x)[1]
    p2 = plot(y,x,label="supply 1")
    plot!(10*y,x,label="supply 10")
    plot!(20*y,x,label="supply 20")
    d = x.^(-eta)
    plot!(d,x,label="Demand")

    push!(pl,p2)

    plot(pl...,layout=2)
end

end # module


# Refined Löwner-John ellipsoid example with Mosek, MosekTools, and JuMP

# Computes the inner and outer Löwner-John ellipsoids around the feasible
# space of the JuMP model.

# Uses Polyhedra and CDDLib packages to express the polyhedral set of the JuMP
# model by inequalities and a convex hull of a finite set of points

using JuMP, Mosek, MosekTools
using LinearAlgebra, Polyhedra, CDDLib
using Plots

function lowner_john_inner_ellipsoid(A,b)
    """
    Löwner-John inner ellipsoid:
        the function takes H-representation (A,b) of a polytope and
        outputs the parameters of the maximum volume inscribed ellipsoid.

        polyhedral set S:
        S = { x ∈ Rn | ai' x ⩽ bi, i= 1,...,m }

        Proposition 3.7.1 [BenTalN01]
            ... Then the largest volume ellipsoid contained in S is
            E = { x = Z∗u + z∗ | uTu ≤1 }, where (Z∗,z∗) are from

            maximize    t                                           (1.0)
                s.t.    t ⩽ (DetZ)^1/n                              (1.1)
                        Z - PSD                                     (1.2)
                        || Z * ai ||_2 ≤ bi −ai' * z, i = 1,...,m   (1.3)
                        Z ∈ S^n, z ∈ Rn, t ∈ R                      (1.4)

        To model (1.1), refer to [Mosek, 6.2.3 Log-determinant]

            t ⩽ (DetZ)^1/n <==> [Z B; B' diag(B)] - PSD             (2.1)
                                B is lower triangular               (2.2)
                                t ⩽ (Π_i={1,...,n} Bii)^1/n         (2.3)

            where (2.3) maximizes the geometric mean of the Bii variables. It
            can be modelled using power cone [Mosek, 4.2.4 Geometric mean] and
            MOI.GeometricMeanCone available with the MathOptInterface.jl

    [BenTalN01]	    A. Ben-Tal and A. Nemirovski. Lectures on Modern Convex
                    Optimization: Analysis, Algorithms, and Engineering Applications.
                    MPS/SIAM Series on Optimization. SIAM, 2001.
    [Mosek]         https://docs.mosek.com/modeling-cookbook/index.html
    """
    # model size
    m,n = size(A)
    # model declaraton
    model = Model(optimizer_with_attributes(Mosek.Optimizer, "LOG" => 0))
    # variables
    @variable(model, t)
    @variable(model, Z[1:n,1:n], PSD)
    @variable(model, z[1:n])
    @variable(model, B[1:n,1:n])
    # maximize the volume
    @objective(model, Max, t)
    # SOC constraint (1.3)
    @constraint(model, SOC[i=1:m], [ b[i] - (A[i,:]' * z) ; (A[i,:]'*Z)'] in SecondOrderCone())
    # PSD constraint (2.1)
    @constraint(model, [Z LowerTriangular(B); LowerTriangular(B)' diagm(diag(B))] in PSDCone())
    # Power cone costraint (2.3)
    @constraint(model, [t; vec([B[i,i] for i in 1:n])] in MOI.GeometricMeanCone(n+1))
    # optimize
    optimize!(model)
    status=termination_status(model)
    CPU_time = MOI.get(model, MOI.SolveTime())
    # return solution
    return JuMP.value.(Z), JuMP.value.(z), CPU_time
end

function lownerjohn_outer(x)
    """
    Löwner-John outer ellipsoid:
        the function takes V-representation (x) of a polytope and
        outputs the parameters of the maximum volume inscribed ellipsoid.

        Polyhedral set S = Conv{x_1,...,x_n}

        Uses the same Mosek features as the inner ellipsoid function,
        refer to Proposition 3.7.2 [BenTalN01] for further details.
    """

    m,n = size(x)
    model = Model(optimizer_with_attributes(Mosek.Optimizer, "LOG" => 0))

    @variable(model,t >= 0)
    @variable(model,P[1:n,1:n],PSD)
    @variable(model,B[1:n,1:n])
    @variable(model,c[1:n])

    # (1, Px-c) in cone
    for i in 1:m
        @constraint(model,[1 ; (x[i,:]' * P)' - c] in MOI.SecondOrderCone(n+1))
    end

    @constraint(model, [P LowerTriangular(B); LowerTriangular(B)' diagm(diag(B))] in PSDCone())
    @constraint(model, [t; vec([B[i,i] for i in 1:n])] in MOI.GeometricMeanCone(n+1))

    @objective(model,Max,t)

    optimize!(model)
    status=termination_status(model)
    CPU_time = MOI.get(model, MOI.SolveTime())

    JuMP.value.(P), JuMP.value.(c), CPU_time
end

# create a JuMP model
function JuMP_model(n)
    model = Model(optimizer_with_attributes(Mosek.Optimizer, "LOG" => 0))
    @variable(model, x[1:n])
    @objective(model, Min, x'diagm(ones(n))'x)
    @constraint(model, con1[i=1:n], 0 <= x[i] <= 1)
    @constraint(model, ones(n)'x >= n/3)
    optimize!(model)
    return model
end

# model dimension
n = 2
# get JuMP model
model = JuMP_model(n)
@info("done exporting JuMP model")
# create a polyhedron from the feasible set of a JuMP model
poly = polyhedron(model, CDDLib.Library(:exact))
@info("done exporting the polyhedron")

# obtain H-representation of the polyhedron
hr = MixedMatHRep(hrep(poly))
A = hr.A
b = hr.b
@info("done obtaining H-representation")
# solve ellipsoid problem
Ci, di, CPU_time = lowner_john_inner_ellipsoid(A,b)
@info("done solving inner ellipsoid problem in $(CPU_time) seconds")

# obtain V-representation of the polyhedron
vr = vrep(poly)
vr = MixedMatVRep(vr)
p = vr.V
@info("done obtaining V-representation")
Po, co, CPU_time = lownerjohn_outer(p)
@info("done solving outter ellipsoid problem in $(CPU_time) seconds")
Poinv = Po^-1

# make a plot
if n == 2
    println("Plotting solution...")
    Plots.gr()
    Plots.plot([p[:,1] ; p[1,1]], [p[:,2] ; p[1,2]])
    Plots.plot!(t -> Ci[1,:]' * [ cos(t),sin(t) ] + di[1],
                t -> Ci[2,:]' * [ cos(t),sin(t) ] + di[2],
                0,2*pi)
    Plots.plot!(t -> Poinv[1,:]' * ( [cos(t),sin(t)] .+ co[1] ),
                t -> Poinv[2,:]' * ( [cos(t),sin(t)] .+ co[2] ),
                0,2*pi)
end


# ---
# title: 172. Factorial Trailing Zeroes
# id: problem172
# author: zhwang
# date: 2022-03-25
# difficulty: Easy
# categories: Math
# link: <https://leetcode.com/problems/factorial-trailing-zeroes/description/>
# hidden: true
# ---
#
# Given an integer `n`, return _the number of trailing zeroes in`n!`_.
#
# **Follow up:** Could you write a solution that works in logarithmic time
# complexity?
#
#
#
# **Example 1:**
#
#
#
#     Input: n = 3
#     Output: 0
#     Explanation:  3! = 6, no trailing zero.
#
#
# **Example 2:**
#
#
#
#     Input: n = 5
#     Output: 1
#     Explanation:  5! = 120, one trailing zero.
#
#
# **Example 3:**
#
#
#
#     Input: n = 0
#     Output: 0
#
#
#
#
# **Constraints:**
#
#   * `1 <= n <= 104`
#
#
## @lc code=start
using LeetCode

## simulate Mathematica's function NestWhileList
function nest_while_list(f::Function, val::T, chk::Function)::Vector{T} where {T}
    res = [val]
    while chk(val)
        val = f(val)
        push!(res, val)
    end
    return res
end

trailing_zeroes(n::Int) = sum(nest_while_list(i -> i ÷ 5, n ÷ 5, >(1)))
## equivalent to the following
## function trailing_zeroes(n::Int)::Int
##     res = 0
##     while n >=5
##         n ÷= 5
##         res += n
##     end
##     res
## end

## @lc code=end


function LidDrivenCavity(;
                         tfinal = Inf,
                         Lx = 1, Ly = 1, CFL = 0.5, Re = 100,
                         Nx = 65, Ny = 65,
                         u_n = 1, u_s = 0, v_w = 0, v_e = 0,
                         printfreq = 10)
    t0 = time()
    println("------------------Ny = $(Ny), Nx = $(Nx) ---------------")
    Δy  = Ly/(Ny-1)
    Δx  = Lx/(Nx-1)
    x = 0:Δx:Lx
    y = 0:Δy:Ly
    Δt = CFL * Δx

    # Construct matrix for Poisson equation
    A_poisson = BuildPoissonMatrix(Ny,Nx,Δx,Δy) # for coNxgrad
    # Construct matrix for advection-diffusion equation
    ap,an,as,ae,aw = BuildAdvectionDiffusionCoefficients(Re,Δt,Δx,Δy)
    # allocate empty matrices for Gauss-Siedel solver
    Rp = zeros(Ny,Nx); res = zeros(Ny,Nx)

    # initialize ω and ψ
    ω = zeros(Ny,Nx)
    ψ = zeros(Ny,Nx)

    # keep track of changes with Δ's
    ω_old = zeros(Ny,Nx)
    ψ_old = zeros(Ny,Nx)
    ω_hist = []
    ψ_hist = []
    residual = 1

    ###########################
    k0,t = 0,0
    while t < tfinal && maximum(residual) > 1e-8
        t += Δt
        k0 += 1
        # Solve Poisson equation for ψ:
        LinearSolve!(A_poisson,ψ,-ω)

        VorticityBoundaryConditions!(ω,ψ,Δx,Δy,u_n,u_s,v_e,v_w)
        # Modify the explicit part of advection-diffusion equation
        BuildAdvectionDiffusionRHS!(Rp,ω,ψ,Δt,Δx,Δy,Ny,Nx,Re)

        # Solve advection-diffusion equation for ω:
        GaussSiedel!(ω,ap,an,as,ae,aw,Rp,res)

        residual = RecordHistory!.([ω,ψ],[ω_old,ψ_old],[ω_hist,ψ_hist])

        if (k0 % printfreq == 0)
            println("Step: $k0 \t Time: $(round(t,digits=3))\t",
                    "|Δω|: $(round((residual[1]),digits=8)) \t",
                    "|Δψ|: $(round((residual[2]),digits=8)) \t")
        end
    end
    tt = round(time() - t0, digits=3)
    println("This took $tt seconds.")
    println("--------------------------------------------------------")
    Results(ψ,ω,hcat(ω_hist,ψ_hist),x,y,t,k0,Re)
end


# Sine Dictionary based example
using PyPlot
using Statistics
using Random

include("../src/LocallyCompetitive.jl")
Random.seed!(0xcafebeef)

# First we create 'dim' basis functions, each consisting of 'n_elm' samples.
# Using the dictionary 'D' a set of output samples can be created.
# We will create a 'k' sparse signal 'y' and attempt to reconstruct it.
# For a basis function we will use a set of sine functions with the norm of each
# function set to 1.
# The sparse signal 'y' will have 'k' randomly chosen signals.
#
# LCA can effectively reconstruct 'y' original representation even with changes
# to the dictionary used to create it and if LCA attempts to reconstruct
# 'y'+noise it will attempt to impose sparsity on the solution

dim     = 64   #number of basis functions
n_elms  = 1024 #samples per function
D       = zeros(n_elems,dim);
k       = 3    #Sparsity factor
s       = zeros(dim)

for i=1:dim,
    j=1:n_elems
    D[j,i] = sin(2pi*j*i/n_elems)
end
LCA.unitary_columns!(D)

# The original sparse representation
s = rand(dim)*4
s[sortperm(randn(dim))[k+1:end]] .= 0

# The signal that is observed
y = D*s

noise   = randn(n_elems)

sparse_clean  = LCA.lca(y+0.0*noise, D, T_soft=0.20)
sparse_noise  = LCA.lca(y+0.2*noise, D, T_soft=0.25)

figure(1);PyPlot.clf()
subplot(3,1,1)
title("The original sparse elements")
ylabel("Magnitude")
stem(s)
subplot(3,1,2)
title("The reconstructed sparse elements")
ylabel("Magnitude")
stem(sparse_clean)
subplot(3,1,3)
title("The reconstructed elements under noisy observation")
xlabel("Dictionary Element")
ylabel("Magnitude")
stem(sparse_noise)

figure(2);PyPlot.clf()
title("The original sparse elements")
ylabel("Magnitude")
plot(D*s)
plot(D*sparse_clean)
figure(43);PyPlot.clf()
subplot(1,3,1)
plot(y+0.2*noise)
axis([0,n_elms, -0.6, 0.6])
subplot(1,3,2)
plot(y)
axis([0,n_elms, -0.6, 0.6])
subplot(1,3,3)
plot(D*sparse_noise)
axis([0,n_elms, -0.6, 0.6])

println("If LCA is working as expected, then the sparsity of the original")
println("signal input should be identical. For processing the signal with noise")
println("the sparsity should be similar though higher (sparsity will go down")
println("with higher soft_threshold values)")

println()

println("Sparsity of original signal = ", sum(s.!=0))
println("Sparsity of LCA on the clean observation = ", sum(sparse_clean.!=0))
println("Sparsity of LCA on the dirty observation = ", sum(sparse_noise.!=0))

println()
println("Although the sparsity of LCA on the dirty observation is higher,")
println("let's see what happens to the reconstruction error")
println()
println("Original     = ", norm(D*s-y), " (Defined to be zero)")
println("LCA clean    = ", norm(D*sparse_clean-y), " (Should be small)")
println("noisy signal = ", norm(y+noise-y),
        " (Defined to be the norm of the noise")
println("LCA dirty    = ", norm(D*sparse_noise-y),
        " (Much lower than the noisy observation)")


# 第一章
"""
線形回帰分析の係数と切片
```julia
using Random,Distributions
Random.seed!(123)
#データの生成
N=100
a=rand(Normal(2,1),N)
b=randn(1)
x = randn(N)
y = a.*x .+ b + randn(N)
a1, b1 = min_sq(x,y)
```
"""
function min_sq(x::Vector{T},y::Vector{T}) where {T<:Number}
    x̄ = mean(x)
    ȳ = mean(y)
    β₁ = (x .- x̄)'*(y .- ȳ) / norm(x .- x̄)^2 #傾き
    β₀ = ȳ - β₁ * x̄ #切片
    return  β₁,β₀
end

function expand_matrix(X)
    N = typeof(X) <: Matrix  ?  size(X)[1] :  length(X)
    T = eltype(X)
    return hcat(ones(T,N),X)
end
function insert_ones(X)
    N = typeof(X) <: Matrix  ?  size(X)[1] :  length(X)
    T = eltype(X)
    return hcat(ones(T,N),X)
end

"""
重回帰分析の係数
```julia
n = 100; p =2
β = Float64[1,2,3]
x = randn(n,2)
y = @. β[1] + β[2] * x[:,1] + β[3]*x[:,2] + $randn(n)
MultipleRegression(x,y)

3-element Array{Float64,1}:
 0.8957078029816884
 1.904721221695989
 3.1212560196461494
```
"""
function multiple_regression(x::Matrix{T}, y::Vector{T}) where {T<:Number}
    N,_ = size(x)
    @assert N==length(y)
    X = insert_ones(x)
    return (X'X)\X'y
end
function multiple_regression(x::Vector{T}, y::Vector{T}) where {T<:Number}
    return multiple_regression(x[:,:],y)
end

"""
線形回帰分析でのRSS値
"""
function RSS(x::Vector,y)
    β̂ = min_sq(x,y)
    ŷ = β̂[1]*x .+ β̂[2]
    return (y-ŷ)'*(y-ŷ)
end
function RSS(x::Matrix,y)
    ŷ = insert_ones(x)*multiple_regression(x,y)
    return (y-ŷ)'*(y-ŷ)
end

function TSS(y)
    length(y)==0 && return 0
    Y = y .- mean(y)
    return Y'*Y
end

ESS(x,y) = TSS(y) - RSS(x,y)

R2(x,y) = ESS(x,y)/TSS(y)

function RSE(x::Matrix,y)
    n,p = size(x)
    RSS(x,y) / (n-p-1)
end

function RSE(x::Vector,y)
    RSS(x,y) / (length(x)-2)
end

function Bdiag(x)
    X = insert_ones(x)
    return X'X |> inv |> diag
end

"""
VIF(variance inflation factor)
"""
function VIF(x::Matrix)
    _,ncol= size(x)
    v = Vector{Float64}(undef,ncol)
    for j in 1:ncol
        # index = delete!(Set(1:ncol),j) |> collect
        index = setdiff(1:ncol,j)
        v[j] = 1/(1-R2(x[:,index],x[:,j]))
    end
    return v
end

"""
confident_interval(xp,x,y;α=0.01)

線形回帰における信頼区間(予測したい場所での標準偏差)を返す。
xp:説明変数の予測データ
x:説明変数の学習データ
y:目的変数の学習データ
α:0.01 は99%信頼区間
"""
function confident_interval(xp,x,y;α=0.01)
    N,p = size(x)
    X = insert_ones(x)
    XP = insert_ones(xp)
    typeof(xp) <:Number && (xp =[xp])
    yerror = quantile(TDist(N-p-1),1-α/2) * sqrt.(diag(XP * ((X'X) \ XP')))
end

function prediction_interval(xp,x,y;α=0.01)
    N,p = size(x)
    X = insert_ones(x)
    XP = insert_ones(xp)
    typeof(xp) <:Number && (xp =[xp])
    yerror = quantile(TDist(N-p-1),1-α/2) * sqrt.( 1 .+ diag(XP * ((X'X) \ XP')))
end


#RSDeltaSigmaPort: Compute power values.
#-------------------------------------------------------------------------------

"""`dbm(v,R=50) = 10*log10(v^2/R*1000)`

The equivalent in dBm of an RMS voltage v
"""
dbm(v; R::Float64=50.0) = 10*log10.(abs.(v .^ 2)/R) .+ 30;

"""`dbp(x) = 10*log10(x)`

The dB equivalent of the power x
"""
dbp(x) = 10*log10.(abs.(x))

"""`dbv(x) = 20*log10(abs(x))`

The dB equivalent of the voltage x
"""
dbv(x) = 20*log10.(abs.(x))

"`y = rms(x; no_dc=false)`"
function rms(x; no_dc::Bool=false)
	if no_dc
		x = x - mean(x)
	end
	return norm(x)/sqrt(length(x))
end

"""`v = undbm(p, z=50) = sqrt(z*10^(p/10-3))`

RMS voltage equivalent to a power p in dBm
"""
undbm(p; z::Float64=50.0) = sqrt.(z*10 .^ (p/10 .- 3))

"""`y = undbp(x)`

Convert x from dB to a power
"""
undbp(x) = 10 .^ (x/10)

"""`y = undbv(x)`

Convert x from dB to a voltage
"""
undbv(x) = 10 .^ (x/20)

#Last line


# using PyPlot
using ConvDiffMIPDECO
using jInv.Mesh
# using jInvVis
using jInv.LinearSolvers
using MUMPSjInv

domain      = [0 2 0 1]
nc		    = [128 64]*4
pNoise      = .05
noBoreHoles = 30
Mc          = getRegularMesh(domain,nc)
P,mf        = getRandomBoreholes2D(noBoreHoles,nc,nc) 		# Randomly selecting borehols
pForFEM     = getConvDiffFEMParam(Mc,sig=0.01,Ainv=getMUMPSsolver());
v = kron([1;0.0],ones(prod(Mc.n+2)))
pFor  = getConvDiffParam(Mc,v,sig=pForFEM.sig,Ainv=getMUMPSsolver(),bc=(:dir,:neu,:neu,:neu));
w = zeros(Mc.n[1],Mc.n[2]);
w[30:40,30:40] = 1;
viewImage2D(w,Mc)
title("source term")
colorbar()
pFor.sig

dobsFEM, = getData(vec(w),pForFEM);
dobs,    = getData(vec(w),pFor);

subplot(3,1,1)
viewImage2D(pForFEM.Fields,getPaddedMesh(Mc))
ylabel("FEM")
colorbar()
title("Fields")

subplot(3,1,2)
viewImage2D(pFor.Fields,Mc)
ylabel("FD")
colorbar()

subplot(3,1,3)
viewImage2D(getNodalAverageMatrix(Mc)*pForFEM.Fields-pFor.Fields,Mc)
colorbar()
ylabel("Difference")

println("relative error (PDE solve FEM vs FD): $(norm(vec(getNodalAverageMatrix(Mc)*pForFEM.Fields-pFor.Fields))/norm(vec(pFor.Fields)))")

"""
fdAnisoTV(W,Mc)

prototype implementation of regularizer using for loops
"""
function fdAnisoTV(W,Mc)
    W = reshape(W,Mc.n[1], Mc.n[2])
    L = Mc.h

    Reg = 0.0;
    # compute first term of regularizer (derivatives in first dimension)
    for i=2:Mc.n[1]
        for j=1:Mc.n[2]
            Reg += Mc.h[2]*abs(W[i,j] - W[i-1,j])
        end
    end
    # compute first term of regularizer (derivatives in second dimension)
    for i=1:Mc.n[1]
        for j=2:Mc.n[2]
            Reg += Mc.h[1]*abs(W[i,j] - W[i,j-1])
        end
    end
    return Reg
end

Rfd = fdAnisoTV(w,Mc)

using jInv.InverseSolve
Rjinv = anisoTVReg(vec(w),0*vec(w),Mc,eps=0)[1]
println("Relative error (regularizer for loop): $(abs(Rfd-Rjinv)/abs(Rjinv))")

"""
fdPDE(U,Mc)

prototype implementation of PDE using for loops

here, assume that v = [1;0]
"""
function fdPDE(U,Mc;c=0.01)
    Up = zeros(Mc.n[1]+2, Mc.n[2]+2)
    Up[2:end-1,2:end-1] = reshape(U,Mc.n[1], Mc.n[2])

    # left boundary: 0-Dirichlet
    Up[1,2:end-1] = -U[1,:]
    # right boundary: 0-Neumann
    Up[end,2:end-1] = U[end,:]
    # top boundary: 0-Neumann
    Up[2:end-1,end] = U[:,end]
    # bottom boundary: 0-Neumann
    Up[2:end-1,1] = U[:,1]

    L = Mc.h

    # allocate output (with ghostpoints for simpler indexing)
    V = zeros(Mc.n[1]+2, Mc.n[2]+2)
    for i=2:Mc.n[1]+1
        for j=2:Mc.n[2]+1
        V[i,j] = (c/(L[1]^2))*(2*Up[i,j]-Up[i-1,j]-Up[i+1,j]) +
                 (c/(L[2]^2))*(2*Up[i,j]-Up[i,j-1]-Up[i,j+1]) +
                 (Up[i,j]-Up[i-1,j])/L[1]
        end
    end
    V = V[2:Mc.n[1]+1,2:Mc.n[2]+1]
    return V
end
U = randn(Mc.n[1],Mc.n[2])
# U = copy(w)
Vfd = fdPDE(copy(U),Mc,c=pFor.sig);

Vjinv = reshape(pFor.A*vec(U),Mc.n[1],Mc.n[2]);

subplot(3,1,1);
viewImage2D(Vfd,Mc)
colorbar()

subplot(3,1,2);
viewImage2D(Vjinv,Mc)
colorbar()

subplot(3,1,3);
viewImage2D(abs.(Vfd-Vjinv),Mc)
colorbar()

println("relative error (conv diff for loop): $(norm(vec(Vjinv-Vfd))/norm(vec(Vjinv)))")


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.5.0
#   kernelspec:
#     display_name: Julia 1.4.2
#     language: julia
#     name: julia-1.4
# ---

# # 主成分分析をする

# - PCA の例として MNIST 画像 dataset を主成分分析を行う

# ## Load modules

using MLDatasets
using Images
using Statistics
using LinearAlgebra
using Plots

# ## Download MNIST dataset

MNIST.download(;i_accept_the_terms_of_use=true)

train_x, train_y = MNIST.traindata()
test_x, test_y = MNIST.testdata()

# ## Study data

?MNIST.convert2image

MNIST.traintensor(1)|>MNIST.convert2image

train_x |> size # W, H, Num_data

N = size(train_x)[end]

train_x_nhw = permuteddimsview(train_x, (3, 2, 1)); # transpose so that it has NxHxW layout

id = 4
colorview(Gray, train_x_nhw[id, :, :]) |> display
println("label = $(train_y[id])")

# # Apply PCA to 7

sevens = train_x[:, :, train_y .== 7];

μ_sevens = dropdims(mean(sevens, dims=3),dims=3);

colorview(Gray, μ_sevens')

colorview(Gray,(sevens[:,:,1] .- μ_sevens)')

X_ = sum(sevens .- μ_sevens, dims=3)/√N
X = reshape(X_, 784)
X |> size

A = X * X'; # variance-covariance matrix w.r.t sample

F = eigen(A);
λ = F.values[end]
v = F.vectors[:, end];

# visualize the first principal vector
normalize_img(v) = (v .- minimum(v)) ./ (maximum(v)-minimum(v))
colorview(Gray, reshape(normalize_img(v),28,28)') |> display
v

# +
λs = F.values[end:-1:end-5]
vs = F.vectors[:, end:-1:end-5]
normalize(v) = (v .- minimum(v)) ./ (maximum(v)-minimum(v))
gs = []
for i in 1:size(vs)[2]
    v = reshape(vs[:, i] |> normalize,28,28)'
    push!(gs, colorview(Gray, v))
end

colorview.(Gray,gs)
# -

# # Apply PCA to whole dataset

train_x, train_y = MNIST.traindata()
μ_train_x = mean(train_x,dims=3)
X_ = sum(train_x .- μ_train_x, dims=3)/√N
X = reshape(X_, 784)
A = X * X';

F=eigen(A)
λ1 = F.values[end]
λ2 = F.values[end-1]
λ3 = F.values[end-2]
w1 = F.vectors[:, end]
w2 = F.vectors[:, end-1]
w3 = F.vectors[:, end-2];

# ## Project 2D space

p = plot()
data = reshape(train_x,(28^2, N))
for i in 1:10
    d = data[:,train_y .== (i-1)][:, 1:100]
    x1 = w1' * d
    x2 = w2' * d
    scatter!(p,x1,x2, label=false, color=palette(:tab10)[i])
end
p

# ## Project 3D space

# +
#=
If we load `ORCA` package or use plotting operation based on PlotlyJS backend over the sysimage generated by PackageCompiler using my method, it will cause the following error.
"""
InitError: IOError: could not spawn `'/usr/bin/xvfb-run /root/.julia/conda/3/envs/_ORCA_jl_/bin/orca' server --port=7982 --graph-only`: no such file or directory (ENOENT)
during initialization of module ORCA
"""
We will ignore the error at the moment, since there is no harm with respect to our purpose.
Here is my workaround "Just load `ORCA` and ignore the error above or plot a small dummy data in advance"
=#

using ORCA

# +
plotlyjs()

p = plot(xlabel="x", ylabel="y" , zlabel="z")
data = reshape(train_x,(28^2, N))
target = [0,1,2,3,4,5,6,7,8,9]
for i in target
    d = data[:,train_y .== i][:, 1:100]
    x1 = w1' * d
    x2 = w2' * d
    x3 = w3' * d
    color = palette(:tab10)[i+1]
    scatter!(p,
        x1, x2, x3,
        label=false, color=color,
        markersize = 2,
    )
end

display(palette(:tab10)) # show color palette
p


using DataStructures
using StatsBase

# Problem 1
nWashers = 5
nTrials = 1000000
success = 0
for i in 1:nTrials
    rn = rand(1:nWashers,nWashers)
    if get(countmap(rn),1,0) >= 4
        success+=1
    end
end
success/nTrials

# Problem 2
wait1 = 5
wait2 = 7
nTrials = 1000000
success = 0
for i in 1:nTrials
    rn = rand(2)
    if (0 <= (rn[2]-rn[1])*30 <= wait1) ||
       (0 <= (rn[1]-rn[2])*30 <= wait2)
        success+=1
    end
end
success/nTrials

# Problem 3
# nCars=3;  0.6666666666644826
# nCars=4;  0.6667395
# nCars=5;  0.6663580000064863
# nCars=6;  0.6665606666653675
# nCars=7;  0.6665282857209788
# nCars=8;  0.66674625
# nCars=9;  0.6667986666646687
# nCars=10; 0.6665119999960215
# nCars=11; 0.6665550909138054
# nCars=12; 0.6666493333327062
# nCars=20; 0.6666518999954775
# nCars=30; 0.6665245333273143

nCars = 7
nTrials = 1000000
success=0
for i in 1:nTrials
    rn = rand(nCars)
    rn = sort(rn)
    cn = zeros(nCars)
    cn[1]=2
    cn[nCars]=nCars-1
    for j in 2:nCars-1
        if rn[j+1]-rn[j] >= rn[j]-rn[j-1]
            cn[j]=j-1
        else
            cn[j]=j+1
        end
    end
    cnTotal = 0
    j = 1
    while j <= nCars-1
        if (cn[j]==j+1 && cn[j+1]==j)
            j+=2
            cnTotal+=2
        else
            j+=1
        end
    end
    success+=cnTotal/nCars
end
success/nTrials


# ---
# title: 754. Reach a Number
# id: problem754
# author: Indigo
# date: 2021-07-11
# difficulty: Medium
# categories: Math
# link: <https://leetcode.com/problems/reach-a-number/description/>
# hidden: true
# ---
#
# You are standing at position `0` on an infinite number line. There is a goal
# at position `target`.
#
# On each move, you can either go left or right. During the _n_ -th move
# (starting from 1), you take _n_ steps.
#
# Return the minimum number of steps required to reach the destination.
#
# **Example 1:**
#
#
#
#     Input: target = 3
#     Output: 2
#     Explanation:
#     On the first move we step from 0 to 1.
#     On the second step we step from 1 to 3.
#
#
# **Example 2:**
#
#
#
#     Input: target = 2
#     Output: 3
#     Explanation:
#     On the first move we step from 0 to 1.
#     On the second move we step  from 1 to -1.
#     On the third move we step from -1 to 2.
#
#
# **Note:**
#
# * `target` will be a non-zero integer in the range `[-10^9, 10^9]`.
#
#
## @lc code=start
using LeetCode

function reach_number(target::Integer)
    target = abs(target)
    k = 0
    while target > 0
        k += 1
        target = target - k
    end
    return iseven(target) ? k : k + 1 + k % 2
end
## @lc code=end