# "Escaping the swirling vortex"
#
# Nonautonomous ODE from pg 121 of Hale's ODE book, credited to
# Markus, L. and H. Yamabe "Global stability criteria for diferential systems",
# Osaka Math. J. 12 (1960), 305-317.
#
# A nonautonomous linear ODE of the form `du/dt = A(t)u`, with all the eigenvalues
# of `A(t)` being complex with negative real part.
#
# However, origin is an unstable fixed point!
#
# The vector field visualization was based on
# https://gist.github.com/Wikunia/4ca58b4c08da90978c0dcdee5b0c908a

using Javis
using Colors
using LinearAlgebra

"""
    ground(args...)

Define the background and the initial "foreground" color.
"""
function ground(args...)
    background("black")
    sethue("white")
    return nothing
end

"""
dudt(x, y, t)

Return the right hand side of the linear equation `u'=A(t)u`,
where u=[x,y]. In this case, the equation takes the from
    x' = (-1 + 1.5cos(t)^2)x + (1 - 1.5cos(t)sin(t))y
    y' = (-1 - 1.5sin(t)cos(t))x + (-1 + 1.5sin(t)^2)y
"""
function dudt(x, y, t)
    ct = cos(t)
    st = sin(t)
    dxdt = (-1.0 + 1.5 * ct^2) * x + (1.0 - 1.5 * ct * st) * y
    dydt = (-1.0 - 1.5 * st * ct) * x + (-1.0 + 1.5 * st^2) * y
    return dxdt, dydt
end

"""
    jacobian!(mat, t)

Return (and modify in place) the linear part `A(t)` of the equation `u'=A(t)u`
(see [`dudt`](@ref)). This is also seen as the Jacobian of `u -> f(t,u) = A(t)u`.
"""
function jacobian!(mat, t)
    ct = cos(t)
    st = sin(t)
    mat[1, 1] = -1.0 + 1.5 * ct^2
    mat[1, 2] = 1.0 - 1.5 * ct * st
    mat[2, 1] = -1.0 - 1.5 * st * ct
    mat[2, 2] = -1.0 + 1.5 * st^2
    return mat
end

"""
    solution!(u,t)

Return a particular solution of the equation `u'=A(t)u`, given by
`u = exp(t/2)  * [-cos(t), sin(t)] / 20.0` corresponding to the initial condition
`u(0) = [-0.05, 0.0]`
"""
function solution!(u, t)
    et2 = exp(t / 2) / 20.0
    u[1] = -et2 * cos(t)
    u[2] = et2 * sin(t)
    return u
end

"""
    space_to_frame_coordinates(u, frm_scl)

Transform a point in xy coordinates to the frame coordinates XY.
They are related according to `X = frm_scl * x` and `Y = - frm_scl * y`,
for a given scale factor `frm_scl`.
"""
space_to_frame_coordinates(u, frm_scl) = frm_scl * Point(u[1], -u[2])
space_to_frame_coordinates(x, y, frm_scl) = space_to_frame_coordinates([x, y], frm_scl)

"""
    ellipse_level_set!(u, t, mat, σ₁, σ₂, frm_scl)

Draw a "level-set"-like ellipse associated with the solution at time `t`.

For a given time `t`, the solutions `w=w(s)` of the autonomous equation
`dw(s)/ds = A(t)w(s)` spiral towards the origin, decreasing along the level
sets of an elliptic paraboloid. This function computes and draws the ellipse
level set containing the solution `u(t)`. Note, however, that `u=u(t)` is
the solution of the nonautonomous equation `du(t)/dt = A(t)u(t)` given
by [`solution!`](@ref), which spirals away from the origin.

The ellipse is computed by finding the complex eigenvector `v=v(t)` of the
complex eigenvalue `α + iβ` of `A(t)` and considering the real and invertible
matrix `P(t) = [Re(v(t)) Im(v(t))]` which transforms `A(t)` into its real
Jordan form `P(t)^-1 * A(t) P(t) = [α β, -β α]`. In this particular case, the
eigenvalues `α ± iβ` are time-independent. The ellipse is obtained from the
SVD decomposition `P = UΣVᵗ`, where `Σ = diag(σ₁, σ₂)` is diagonal with
nonnegative elements `σ₁ ≥ σ₂ ≥ 0`, and `U` and `V` are orthogonal matrices.
In general, all these terms may change in time, but in the particular
example we are considering, only `U=U(t)` and `V=V(t)` are time-dependent,
with `σ₁ > σ₂ > 0` being constant. The elements `σ₁` and `σ₂` are, precisely,
the semi-major and semi-minor axes of the ellipse associated with the level
set of the elliptic paraboloid at level one. It determines the eccentricity
of the ellipse. The orientation of the ellipse is given by the rotation
determined by the orthogonal matrix `U=U(t)`.

At each time `t`, we find and draw the level set that contains the solution `u(t)`
as follows. We rotate `u(t)` via `Uᵗ = inv(U)`, compute the semi-major
and semi-minor axis of the ellipse that has the same level as `Uᵗu(t)`, compute
the foci of this ellipse, rotate the foci back to the right orientation by applying
the orthogonal matrix `U`, and then find the ellipse, in frame coordinates, that
defined by the foci, in frame coordinates, and passing by the solution `u(t)`, in
frame coordinates.
"""
function ellipse_level_set!(u, t, mat, σ₁, σ₂, frm_scl)
    sethue("white")
    jacobian!(mat, t)
    v = eigen(mat).vectors[:, 1]
    U = svd([real(v) imag(v)]).U
    x, y = transpose(U) * solution!(u, t)
    γ = sqrt(x^2 / σ₁^2 + y^2 / σ₂^2) # level
    σ̃₁ = γ * σ₁ # semi-major
    σ̃₂ = γ * σ₂ # semi-minor
    fd = sqrt(abs(σ̃₁^2 - σ̃₂^2)) # focal distance to the origin/center
    f = U * [fd, 0.0] # one focus point; the other is -f
    poly(
        ellipse(
            space_to_frame_coordinates(f, frm_scl), # focus in frame coordinates
            -space_to_frame_coordinates(f, frm_scl), # the other focus
            space_to_frame_coordinates(u, frm_scl), # ellipse/level containing u(t)
        ),
        :stroke,
        close = true,
    )
    return nothing
end

"""
    ball(p=O, size=4, color="red")

Draw a filled circle of size `size` and color `color` around the point `p`
"""
function ball(p = O, size = 4, color = "red")
    sethue(color)
    circle(p, size, :fill)
    return p
end

"""
    path!(points, pos, color)

Update and draw the trail of the solution.

Add the point `pos` to the vector `points` containing the trail of the trajectory
and draw the trail.
"""
function path!(points, pos, color)
    sethue(color)
    push!(points, pos) # add pos to points
    prettypoly(points) # draw dotted path with circles
    return nothing
end

"""
    ode_arrow(x, y, t, frm_scl, color_min, color_max)

Create one arrow at the base point `(x,y)` and arrow vector given by `dudt(x, y, t)`.
The arrow is properly scaled by the factor `frm_scl` and the orientation takes into
account that the vertical coordinate increases downwards. Depending on the magnitude
of the vector field at that point, the color ranges from `color_min` to `color_max`.
"""
function ode_arrow(x, y, t, frm_scl, color_min, color_max)
    @JShape begin
        translate(space_to_frame_coordinates(x, y, frm_scl))
        dx, dy = dudt(x, y, t)
        delta_t = 0.1
        ξ = dx * delta_t
        η = dy * delta_t
        l = sqrt(ξ^2 + η^2)
        len_colors = 100
        color_range = range(color(color_min), stop = color(color_max), length = len_colors)
        idx = clamp(floor(Int, l * len_colors), 1, len_colors)
        sethue(color_range[idx])
        if ξ != 0 || η != 0
            arrow(
                O,
                space_to_frame_coordinates(ξ, η, frm_scl);
                linewidth = 2,
                arrowheadlength = 8,
            )
        end
    end x = x y = y t = t frm_scl = frm_scl
end

"""
    horizontal_grid_line(video, c, frm_scl)

Create an horizontal grid line at `y=c`
"""
function horizontal_grid_line(video, r, frm_scl)
    if r == 0
        setline(1.5)
    end
    sethue("gray")
    line(Point(-video.width / 2, r * frm_scl), Point(video.width / 2, r * frm_scl), :stroke)
    return nothing
end

"""
    vertical_grid_line(video, c, frm_scl)

Create a vertical grid line at `x=c`.
"""
function vertical_grid_line(video, c, frm_scl)
    if c == 0
        setline(1.5)
    end
    sethue("gray")
    line(
        Point(c * frm_scl, -video.height / 2),
        Point(c * frm_scl, video.height / 2),
        :stroke,
    )
    return nothing
end

"""
    animate()

Create the animation. 😃
"""
function animate()
    # initial setup
    vid = Video(500, 500)
    nframes = 200
    xgrid = -2.0:0.25:2.0 # grid for bars and arrow base points
    ygrid = -2.0:0.25:2.0 # grid for bars and arrow base points
    t₀, tf = 0.0, 5π / 2 # time interval for the simulation
    u = [0.0, 0.0] # dummy initizaliation of the solution vector
    solution!(u, t₀) # initial vector
    mat = zeros(2, 2) # dummy initialization of the jacobian
    jacobian!(mat, t₀) # jacobian at time t₀
    v = eigen(mat).vectors[:, 1] # first eigenvector
    svd_P = svd([real(v) imag(v)]) # P = [real(v) imag(v)] => inv(P) * mat * P = [α β; -β α]
    σ₁, σ₂ = svd_P.S # the semi-major and semi-minor axes are constant in this problem
    t = t₀ # set t to initial time
    frm_scl = 120 # frame scale
    solution_trail = Point[] # initialize history vector for the solution path

    # objects
    ## set background for all frames
    Background(1:nframes, ground)
    ## grids
    [
        Object(1:nframes, (args...) -> horizontal_grid_line(args[1], x, frm_scl)) for
        x in xgrid
    ]
    [Object(1:nframes, (args...) -> vertical_grid_line(args[1], y, frm_scl)) for y in ygrid]
    ## vector field
    arrows = [
        Object(1:nframes, ode_arrow(x, y, t, frm_scl, "yellow", "red")) for x in xgrid,
        y in ygrid
    ]
    ## solution at time t
    sol_pos = Object(
        1:nframes,
        (args...; t) -> ball(
            space_to_frame_coordinates(solution!(u, t), frm_scl),
            6,
            "lightskyblue",
        ),
    )
    ## solution trail
    Object(1:nframes, (args...) -> path!(solution_trail, pos(sol_pos), "lightskyblue"))
    ## ellipse
    obj_ellipse =
        Object(1:nframes, (args...; t) -> ellipse_level_set!(u, t, mat, σ₁, σ₂, frm_scl))

    # animation
    ## update solution, vector field and ellipse
    ## rmk: animation of the solution trail is a byproduct of updating solution
    act!(sol_pos, Action(1:nframes, change(:t, t₀ => tf)))
    act!(arrows, Action(1:nframes, change(:t, t₀ => tf)))
    act!(obj_ellipse, Action(1:nframes, change(:t, t₀ => tf)))

    # rendering
    ## render and generate gif
    render(
        vid;
        framerate = 15,
        pathname = joinpath(@__DIR__, "gifs/escaping_the_swirling_vortex.gif"),
    )

    return nothing
end

animate()


using LinearAlgebra
using Plots

function get_mass2_mat(nu1_mass::Float16, nu2_mass::Float16)::Array
    """Create a 2D mass-squared diagonal matrix"""
    ret_array::Array{Float16} = Diagonal([nu1_mass^2, nu2_mass^2])
    return ret_array
end

function get_mixing_matrix(mixing_angle::Float16, isdeg::Bool = true)::Array
    """Create a 2D mixing/rotation matrix with a given angle"""
    if isdeg
        mixing_angle = deg2rad(mixing_angle)
    end
    mixing_mat::Array{Float16} = [
        +cos(mixing_angle) -sin(mixing_angle)
        +sin(mixing_angle) +cos(mixing_angle)
    ]
    return mixing_mat
end

function time_propagator_mat(
        energy_nu::Float16,
        time_step::Float16,
        delta_mat::Array{Float16}
)::Array{ComplexF16}
    """Calculate the time propagator. Due to a bug in Julia, complex
    matrix arguments for the exponential function are not supported.
    So instead, I use the Euler relation instead.
    """
    exp_arg::Array{ComplexF16} = time_step / (2.0 * energy_nu) * delta_mat
    exp_mat::Array{ComplexF16} = cos(exp_arg) + (0 + 1im) * sin(exp_arg)
    return exp_mat
end

function schrodingerODE()
    """Run a 2D toy example of neutrino oscillations"""

    # Define scalars
    nu1_mass::Float16 = 1.0
    nu2_mass::Float16 = 2.0
    mixing_angle_deg::Float16 = 40
    time_steps::Int = 100
    time_step_size::Float16 = 0.01
    energy_nu::Float16 = 10.0

    # Define physics matrices
    mass2_mat::Array{Float16} = get_mass2_mat(nu1_mass, nu2_mass)
    mixing_mat::Array{Float16} = get_mixing_matrix(mixing_angle_deg, true)
    delta_mat::Array{Float16} = transpose(conj(mixing_mat)) * mass2_mat * mixing_mat

    # Store results in arrays
    # initial state is nu1 = 100%, nu2 = 0%
    initial_state_vector::Vector{ComplexF16} = Vector([1 + 0im, 0 + 0im])
    state_vectors::Array{ComplexF16} = zeros(ComplexF16, (2, time_steps))
    state_fractions::Array{Float16} = zeros(Float16, (2, time_steps))
    state_vectors[:, 1] = initial_state_vector
    state_fractions[:, 1] = Array([1.0; 0.0])

    # Evaluate the wave-function over time and save the state over each step
    for step_index::Int = 2:time_steps
        previous_step_index::Int = step_index - 1
        previous_state_vector::Array{ComplexF16} = state_vectors[:, previous_step_index]
        time_step::Float16 = (step_index - 1) * time_step_size
        time_propagator::Array{ComplexF16} =
            time_propagator_mat(energy_nu, time_step, delta_mat)
        new_state::Vector{ComplexF16} = time_propagator * previous_state_vector
        new_state_fraction::Vector{Float16} =
            Array([abs2(new_state[1]); abs2(new_state[2])])
        state_vectors[:, step_index] = new_state
        state_fractions[:, step_index] = new_state_fraction
    end

    # Plot the results
    Plots.pyplot()
    fig = plot(1:time_steps, state_fractions[1, :], label="\$\\nu_1\$")
    plot!(1:time_steps, state_fractions[2, :], label="\$\\nu_2\$")
    xlabel!("Time step")
    ylabel!("Wave Function Component")
    title!("Flavor Oscillations")
    @show fig
end

schrodingerODE()


"""
Approximates the area under the curve using the trapezoidal rule
Arguments:
 - f: function for the
 - x_start: starting value for x
 - x_end: ending value for x
 - steps: steps taken while integrating.
"""

function trapazoidal_area(f, x_start, x_end, steps)
	x1 = x_start
	fx1 = f(x_start)
	area = 0.0

	for _ in 1:steps
		x2 = ((x_end - x_start) / steps) + x1
		fx2 = f(x2)
		area += abs(fx2 + fx1) * (x2 - x1) / 2
		x1 = x2
		fx1 = fx2
	end

	return area
end


#
#--------------------------------------#
export derivMat
#--------------------------------------#
"""
 Compute derivative matrix for lagrange
 interpolants on points [x]
"""
function derivMat(x)

n = length(x)

a = ones(1,n)
for i=1:n
    for j=1:(i-1) a[i]=a[i]*(x[i]-x[j]) end
    for j=(i+1):n a[i]=a[i]*(x[i]-x[j]) end
end
a = 1 ./ a # Barycentric weights

# diagonal elements
D = x .- x'
for i=1:n D[i,i] = 1. end
D = 1 ./ D
for i=1:n
    D[i,i] = 0.
    D[i,i] = sum(D[i,:])
end

# off-diagonal elements
for j=1:n for i=1:n
    if(i!=j) D[i,j] = a[j] / (a[i]*(x[i]-x[j])) end
end end

return D
end
#--------------------------------------#


using LinearAlgebra

# Parameters in regression model
beta = [148.73, -18.85]; v = 54.40;
y_hat(x) = x'*beta + v;
# Evaluate regression model prediction
x = [0.846, 1]; y = 115;
y_hat(x), y                     # (161.37557999999999, 115)
x = [1.324,2];  y = 234.50;
y_hat(x), y

using VMLS                      # see below to install package
D = house_sales_data();
price = D["price"]
area = D["area"]
beds = D["beds"]

v = 54.4017;
beta = [ 147.7251, -18.8534 ];
predicted = v .+ beta[1] * area + beta[2] * beds;

using Plots
scatter(price, predicted, lims = (0,800));
plot!([0, 800], [0, 800], linestyle = :dash);
# make axes equal and add labels
plot!(xlims = (0,800), ylims = (0,800), size = (500,500));
plot!(xlabel = "Actual price", ylabel = "Predicted price");

# Install VMLS to JuliaBox
# Go to git, enter 'https:////github.com/VMLS-book/VMLS.jl.git'
# Go to package manager, add 'VMLS'
# Wait up to 15 minutes (for package build)

# --OR--

# Install VMLS to Julia runtime (local machine)
# Go to Julia prompt, type ']' to open package manager
# Get library from 'pkg>' prompt
#    type 'add https:////github.com/VMLS-book/VMLS.jl'
#    'ctrl + C' to exit pakage manager prompt


### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# ╔═╡ 32504772-aaef-11eb-0bb0-0b145c5102db
begin
	using Base.Threads
	using Random
	using Distributions
	using SpecialFunctions
	using Plots
	using StatsPlots
	using Query
	using DataFrames
	using CSV
end

# ╔═╡ ec00bfbc-4889-4fbc-9f5b-ce19b361f1fe
md"""
## Estimating π in Buffon's Needle Problem
First we estimate π many times by simulating Buffon's needle problem. This is repeated for a variety of experimental setups, varying the number of needles thrown (`n_throws`) and the needle length (`needle_length`). For simplicity, we'll just assume that the strip height (`strip_height`) is one - so you can think of the needle length as the ratio of the needle length to the strip height.
"""

# ╔═╡ 7e54380a-fada-40c0-ad20-5795075f2c41
function needle_crossing_fraction(strip_height, needle_length, n_throws)
	D = rand(Uniform(0, strip_height/2), n_throws)
	θ = rand(Uniform(0, π), n_throws)

	crosses_line = (needle_length / 2) * sin.(θ) .> D
	n_crossings = sum(crosses_line)

	return n_crossings / n_throws
end

# ╔═╡ 3de8036e-723b-4805-9384-f87a6aaae5ba
function buffon_pi_estimate(strip_height, needle_length, n_throws)
	p_hat = needle_crossing_fraction(strip_height, needle_length, n_throws)

	return 2 * (needle_length / strip_height) / p_hat
end

# ╔═╡ 9e97dbc2-7d87-40cd-92a8-e892e485d9e6
buffon_pi_estimate(1, 0.1, 10^7)

# ╔═╡ 7838810d-63c0-4de6-8ec7-97a4b65750e2
function estimate_pi_repeatedly(
		n_needle_lengths = 10,
		n_repetitions = 100,
		n_throws_max_exp = 6
	)

	# Set up experiment, arrange data so that (when multi-threading) the
	# processors will be running similar problems at similar times
	df = DataFrame(
		[(needle_length,  n_throws, repetition)
			for n_throws = [10^k for k = 2:n_throws_max_exp]
			for needle_length = [k/n_needle_lengths for k = 1:n_needle_lengths]
			for repetition = 1:n_repetitions]
	)
	rename!(df, ["needle_length", "n_throws", "repetition"])

	# Number of experiments to run
	n_experiments = nrow(df)

	# Run multi-threaded experiments to estimate π
	pi_estimates = zeros(n_experiments)

	Threads.@threads for i ∈ 1:n_experiments
		needle_length = df[i, :needle_length]
		n_throws = df[i, :n_throws]
		pi_estimates[i] = buffon_pi_estimate(1, needle_length, n_throws)
	end

	df.pi_estimate = pi_estimates

	# Evaluate error
	df.error = df.pi_estimate .- π
	df.abs_error = abs.(df.error)

	return df
end

# ╔═╡ bbd5cc5f-80f2-405f-aa71-da0022d74863
pi_estimate_results = estimate_pi_repeatedly(10, 1000, 6)

# ╔═╡ 62e8e764-516a-452a-8f7f-806bf696b126
CSV.write("pi_estimate_results.csv", pi_estimate_results)

# ╔═╡ 27eabc0d-b40f-448f-90ce-301871f3b63b
md"""
## Analyzing the Error in Estimating π
In this section, we take a look at the error committed in estimating π in the various experiments performed above. In particular, we take a look at how the error depends upon the number of needles thrown and how long our needles are. In particular, we take a look at the median error and 95ᵗʰ percentile of our experiments.

In addition, I have computed the theoretical error "bounds" (note this bound is only asymptotically correct / hence more of a guide than a hard and fast "bound"), for 50% and 95% confidence bounds. I didn't bother to plot the actual vs theoretical comparison however, since they're not hard and fast bounds.

I have (however) plotted the observed median and 95ᵗʰ percentile observed error across the experimental set-ups.
"""

# ╔═╡ 6c17402b-8905-463c-aa0f-3a4905877944
md"""

### Analytical Asymptotic Error Bound
Elsewhere, I computed that the following bound holds asymptotically with probability $$P = 1-\alpha$$ (here, $$\Phi^{-1}$$ denotes the probit function, $$h$$ the strip width, $$\ell$$ the needle length, and $$\hat{\pi}$$ the estimate from throwing needles):

$$\vert \hat{\pi} - \pi \vert \lessapprox \pi  \sqrt{\frac{\pi}{2} - \frac{\ell}{h}} \cdot \Phi^{-1}\left(1 - \frac{\alpha}{2}\right)\cdot \sqrt{\frac{h}{\ell}} \cdot \frac{1}{\sqrt{n_\text{throws}}}.$$

We will compare against this below.
"""

# ╔═╡ d75c194b-20df-4ea8-a415-7c44167271d4
function probit(p)
	sqrt(2) * erfinv(2*p - 1)
end

# ╔═╡ 60f54651-6926-4248-aa40-f135d01e4347
function error_bound(needle_length, strip_width, prob_in_interval, n_throws)
	return π * √(π/2 - (needle_length/ strip_width)) *
	                    probit(1 - prob_in_interval/2) *
		                √(strip_width / needle_length) *
	                    1/√(n_throws)
end

# ╔═╡ 3c1ade5a-7a15-4585-8cf2-48b6c84813f7
md"""
The following table shows empirical error bounds, compared against the theoretical approximate bounds.
"""

# ╔═╡ a3304dce-9e87-444d-ab5d-c0f3fffea134
begin
	pi_error_analysis = combine(
		groupby(pi_estimate_results, [:needle_length, :n_throws]),
		:abs_error => median,
		:abs_error => (t -> quantile(t, 0.95))
	)

	rename!(pi_error_analysis, "abs_error_function" => "abs_error_95_pctile")

	transform!(pi_error_analysis, :, [:needle_length, :n_throws] =>
                                     ((a,b) -> error_bound.(a, 1, 0.5, b)) =>
		                             :abs_error_median_theoretical)

	transform!(pi_error_analysis, :, [:needle_length, :n_throws] =>
                                     ((a,b) -> error_bound.(a, 1, 0.05, b)) =>
		                             :abs_error_95_pctile_theoretical)

	pi_error_analysis
end

# ╔═╡ 902e5796-6741-4c4a-9856-a43006c326ab
CSV.write("pi_error_analysis.csv", pi_error_analysis)

# ╔═╡ 7014ffed-1bed-41aa-95ef-ecdcd366b510
md"""
The following plot shows the empirical median absolute error, compared against the theoretical approximate 50% percent confident bound on error.
"""

# ╔═╡ ab93cbb3-6b82-48d7-8fdd-4c91f300d33c
begin
	plot_ref_50th = pi_error_analysis |>
	@df scatter(
		:n_throws,
		:abs_error_median,
		group=:needle_length,
		color=reshape(palette(:auto)[1:10], (1, 10)),
		label=reshape(["" for i =1:10], (1,10)),
		# ugly hack to hide junk entry:
		xaxis=((75, 10^6+5*10^5), :log),
		yaxis=:log
	)

	pi_error_analysis |>
	@df plot!(
		:n_throws,
		:abs_error_median_theoretical,
		group=:needle_length,
		#legend=false
		color=reshape(palette(:auto)[1:10], (1, 10)),
		ylabel="Median of Absolute Error",
		xlabel="Number of Needles Thrown",
		title="Median Absolute Error vs Number of Needles Thrown\nby Needle Length",
		titlefontsize=10,
		legendtitle="Needle Length",
		legendtitlefontsize=8,
		legend = :outertopright

	)

	# junk data to get legend entries
    plot!([], [], label = " ", color = "white")
    plot!([], [], label = "Theoretical", color = "black")
    scatter!([10], [1], label = "Observed", color="black")

	# save figure
	savefig(plot_ref_50th, "error_plot_50pctile.svg")

	# display result in Pluto
	plot_ref_50th
end

# ╔═╡ ed32a816-05d8-46f2-b3fe-6a9c5a2f2636
md"""
The following plot shows the empirical 95ᵗʰ percentile of absolute error, compared against the theoretical 95% confident bound on error.
"""

# ╔═╡ a5f76835-73ce-4fa9-a0b0-e63f1f8bff69
begin
	plot_ref_95th = pi_error_analysis |>
	@df scatter(
		:n_throws,
		:abs_error_95_pctile,
		group=:needle_length,
		color=reshape(palette(:auto)[1:10], (1, 10)),
		label=reshape(["" for i =1:10], (1,10)),
		# ugly hack to hide junk entry:
		xaxis=((75, 10^6+5*10^5), :log),
		yaxis=:log
	)

	pi_error_analysis |>
	@df plot!(
		:n_throws,
		:abs_error_95_pctile_theoretical,
		group=:needle_length,
		#legend=false
		color=reshape(palette(:auto)[1:10], (1, 10)),
		ylabel="95ᵗʰ Percentile of Absolute Error",
		xlabel="Number of Needles Thrown",
		title="95ᵗʰ Percentile of Absolute Error vs Number of Needles Thrown\nby Needle Length",
		titlefontsize=10,
		legendtitle="Needle Length",
		legendtitlefontsize=8,
		legend = :outertopright

	)

	# junk data to get legend entries
    plot!([], [], label = " ", color = "white")
    plot!([], [], label = "Theoretical", color = "black")
    scatter!([10], [1], label = "Observed", color="black")

	# save figure
	savefig(plot_ref_95th, "error_plot_95pctile.svg")

	# display result in Pluto
	plot_ref_95th
end

# ╔═╡ bde8330a-fb98-4950-ad92-c82e2e5e65c6
md"""
### Plotting π Estimates when `needle_length = 1.0`

Since we have a lot of estimates for π, but have only looked at the error we made when estimating them (as opposed to the estimates themselves), I figure I should at least take a look at the estimates in one of the experimental settings.
"""

# ╔═╡ d82dd81c-bf96-443e-8b5a-aa75b21e037e
begin
	df_tmp = filter(row -> row[:needle_length] == 1.0, pi_estimate_results)

	plot_ref_len_1 = df_tmp |>
	@df dotplot(
		string.(:n_throws),
		:pi_estimate,
		ylabel="Estimate for π",
		xlabel="Number of Needles Thrown",
		title="Distribution of π Estimates by Number of Needles Thrown\nFor Needle Length=1.0",
		titlefontsize=11,
		marker=(:lightgrey, stroke(0))
	)

	plot_ref_len_1 = df_tmp |>
	@df boxplot!(
		string.(:n_throws),
		:pi_estimate,
		marker=(:black, stroke(0)),
		fillcolor=:lightblue,
		fillalpha=0.75,
		legend=false
	)

	# save figure
	savefig(plot_ref_len_1, "pi_estimate_boxplot.svg")

	# show figure
	plot_ref_len_1
end

# ╔═╡ b29276d8-e541-4221-b81f-fde2600a8979
md"""
The following table summarizes the results:
"""

# ╔═╡ 312ba44a-dcad-4b18-b130-3439625f78f3
begin
	pi_rslt = combine(
		groupby(pi_estimate_results, [:needle_length, :n_throws]),
		:pi_estimate => mean => :mean,
		:pi_estimate => minimum => :min,
		:pi_estimate => (t -> quantile(t, 0.25)) => :pctile_25,
		:pi_estimate => median => :median,
		:pi_estimate => (t -> quantile(t, 0.75)) => :pctile_75,
		:pi_estimate => maximum => :max,
	)
end

# ╔═╡ b847e348-4afd-4e63-b91b-84152cbb89b7
md"""
Let's take a look at the `needle_length = 1` case.
"""

# ╔═╡ e2042def-bf39-4b95-bdc9-1d46db79c88e
# show an interesting subset
filter(row -> row[:needle_length] == 1.0, pi_rslt)

# ╔═╡ 9b3f7234-9978-43bb-9af9-55bdf967ed62
# save to file
CSV.write("pi_rslt.csv", pi_rslt)

# ╔═╡ 489c446d-f942-4703-b0d7-7a29c98ce4d5
md"""
## How Many Needles Do I Need to Throw?
Finally, I will use the theoretical 95ᵗʰ percentile error bound to estimate how many needles one needs to throw in order to achieve a specified decimal digit accuracy (with 95% confidence) in estimating π.
"""

# ╔═╡ aeb3c67c-5248-4312-9930-852259f43de9
function n_throws_required(desired_accuracy, needle_length, strip_width,
		                   prob_in_interval)

	p = 1 - prob_in_interval

	n_throws =  π^2 * (π/2 - (needle_length / strip_width)) *
	                  probit(1 - p/2)^2 *
		              (strip_width / needle_length) *
	                  (1 / desired_accuracy^2)

	return n_throws
end

# ╔═╡ b988a3e4-1d15-4abf-8b5a-fc8a2e119c7e
begin
	df_required_analytical = DataFrame([
	 (decimal_digits, needle_length, n_throws_required(1/10^decimal_digits,                                                                  needle_length, 1, 0.95))
			for decimal_digits = [1, 2, 3, 4, 5, 6]
			for needle_length = [k/10 for k = 1:10]
	])

	rename!(df_required_analytical, ["decimal_digits", "needle_length",                                                    "throws_required"])
end

# ╔═╡ 55234b47-1b22-4fab-a449-c43b2976b7fa
md"""
Let's take a look at the `needle_length = 1` case.
"""

# ╔═╡ 9668af8a-ecea-4c11-9c3a-f32b416f02a9
# show an interesting subset
filter(row -> row[:needle_length] == 1.0, df_required_analytical)

# ╔═╡ fd88ca22-a191-4e92-900f-c86c7f516641
CSV.write(
	"required_throws_analytical_solution.csv",
	df_required_analytical
)

# ╔═╡ 70b92add-4a58-4181-8409-c896b57665f9
begin
	plot_ref_throws_req = df_required_analytical |>
	@df plot(
		:decimal_digits,
		:throws_required,
		group=:needle_length,
		yaxis=:log,
		ylabel="Number of Throws Required",
		xlabel="Number of Decimal Digits",
		title="Number of Throws Required to Achieve Specified\nDecimal Digit Accuracy with 95% Confidence\nby Needle Length",
		titlefontsize=11,
		legendtitle="Needle Length",
		legendtitlefontsize=8,
		legend = :outertopright
	)

	# save figure
	savefig(plot_ref_throws_req, "throws_required_95pctile.svg")

	# display plot
	plot_ref_throws_req
end

# ╔═╡ Cell order:
# ╟─ec00bfbc-4889-4fbc-9f5b-ce19b361f1fe
# ╠═7e54380a-fada-40c0-ad20-5795075f2c41
# ╠═3de8036e-723b-4805-9384-f87a6aaae5ba
# ╠═9e97dbc2-7d87-40cd-92a8-e892e485d9e6
# ╠═7838810d-63c0-4de6-8ec7-97a4b65750e2
# ╠═bbd5cc5f-80f2-405f-aa71-da0022d74863
# ╠═62e8e764-516a-452a-8f7f-806bf696b126
# ╟─27eabc0d-b40f-448f-90ce-301871f3b63b
# ╟─6c17402b-8905-463c-aa0f-3a4905877944
# ╠═d75c194b-20df-4ea8-a415-7c44167271d4
# ╠═60f54651-6926-4248-aa40-f135d01e4347
# ╟─3c1ade5a-7a15-4585-8cf2-48b6c84813f7
# ╠═a3304dce-9e87-444d-ab5d-c0f3fffea134
# ╠═902e5796-6741-4c4a-9856-a43006c326ab
# ╟─7014ffed-1bed-41aa-95ef-ecdcd366b510
# ╠═ab93cbb3-6b82-48d7-8fdd-4c91f300d33c
# ╟─ed32a816-05d8-46f2-b3fe-6a9c5a2f2636
# ╠═a5f76835-73ce-4fa9-a0b0-e63f1f8bff69
# ╟─bde8330a-fb98-4950-ad92-c82e2e5e65c6
# ╠═d82dd81c-bf96-443e-8b5a-aa75b21e037e
# ╟─b29276d8-e541-4221-b81f-fde2600a8979
# ╠═312ba44a-dcad-4b18-b130-3439625f78f3
# ╟─b847e348-4afd-4e63-b91b-84152cbb89b7
# ╠═e2042def-bf39-4b95-bdc9-1d46db79c88e
# ╠═9b3f7234-9978-43bb-9af9-55bdf967ed62
# ╟─489c446d-f942-4703-b0d7-7a29c98ce4d5
# ╠═aeb3c67c-5248-4312-9930-852259f43de9
# ╠═b988a3e4-1d15-4abf-8b5a-fc8a2e119c7e
# ╟─55234b47-1b22-4fab-a449-c43b2976b7fa
# ╠═9668af8a-ecea-4c11-9c3a-f32b416f02a9
# ╠═fd88ca22-a191-4e92-900f-c86c7f516641
# ╠═70b92add-4a58-4181-8409-c896b57665f9
# ╠═32504772-aaef-11eb-0bb0-0b145c5102db


multsum(n, lim) = (occ = div(lim-1, n); div(n*occ*(occ+1), 2))
multsum(n, m, lim) = multsum(n, lim) + multsum(m, lim) - multsum(lcm(n,m), lim)


using LinearAlgebra
using PyPlot
function linesegment_intersect(p1, p2, q1, q2, δ)
    A = [p2-p1 q1-q2]
    b = q1-p1
    if rank(A) == 2
        st = A \ b
        if all(-δ .≤ st .≤ 1 + δ)
            pintersect = p1 + st[1] * (p2 - p1)
            return true, pintersect
        else
            return false, nothing
        end
    else
        # Warning: parallel lines, should check if they actually overlap
        return false, nothing
    end
end

# Example: n random lines, find and plot all intersections
n = 10
lines = [ [rand(2), rand(2)] for i = 1:n ]

# Plot all lines
for line in lines
    plot([line[1][1], line[2][1]], [line[1][2], line[2][2]])
end

# Find all intersections - note that we only consider j>i to avoid checking
# the same 2 line-segments twice
for i = 1:n
    for j = i+1:n
        isect, pq = linesegment_intersect(lines[i][1], lines[i][2], lines[j][1], lines[j][2], 0)
        if isect
            plot(pq[1], pq[2], "o")
        end
    end
end

axis("equal"); grid(true);


export ladder_graph

"""
    ladder_adj(w1, w2)
return true if words `w1` and `w2` have the same length
and differ in exactly one location.
"""
function ladder_adj(w1::String, w2::String)::Bool
    if length(w1) != length(w2)
        return false
    end
    c = collect(w1) .!= collect(w2)
    return sum(c) == 1
end

"""
    ladder_graph(S::Set{String}, len::Int=0, trim::Bool=false)

Create a graph whose vertices are words in which two words are adjacent if they differ
in exactly one letter.

We use only words of length `len`, unless this is set to `0`
in which case all words are used.

If `trim` is set to `true` then vertices of degree zero are deleted.
"""
function ladder_graph(S::Set{String}, len::Int = 0, trim::Bool = false)::SimpleGraph{String}
    G = _bare_graph(S, len)

    PM = Progress(NV(G))

    for v ∈ G.V
        next!(PM)
        nv = length(v)
        for i ∈ 1:nv
            for c ∈ _A2Z
                w = v[1:i-1] * c * v[i+1:nv]
                if w ∈ G.V
                    add!(G, v, w)
                end
            end
        end
    end

    # remove vertices of degree zero (if trim is true)
    if trim
        for w in G.V
            if deg(G, w) == 0
                delete!(G, w)
            end
        end
    end

    if len < 1
        name(G, "Ladder Graph on all words")
    else
        name(G, "Ladder Graph on $len-letter words")
    end

    return G

end


"""
    gauss_lobatto_quad(α, β, N)

Computes Legendre-Gauss-Lobatto quadrature points and weights with Jacobi weights α,β.
"""
function gauss_lobatto_quad(α::Int, β::Int, N)

    if !(iszero(α) && iszero(β))
        error("alpha/beta must be zero for Lobatto")
    end

    x = zeros(N+1)
    w = zeros(N+1)
    if N == 0
        x[1] = 0.
        w[1] = 2.
    elseif N == 1
        x[1] = -1.0
        x[2] = 1.0
        w[1] = 1.0
        w[2] = 1.0
    elseif N==2
        x[1] = -1.
        x[2] = 0.
        x[3] = 1.
        w[1] = 0.333333333333333
        w[2] = 1.333333333333333
        w[3] = 0.333333333333333
    else
        xint, wint = gauss_quad(α+1, β+1, N-2)
        x = vcat(-1,xint,1)
        wend = 2/(N*(N+1))
        w = vcat(wend, (@. wint / (1-xint^2)) ,wend)
    end
    return x[:],w[:]
end

"""
    gauss_quad(α, β, N)

Compute nodes and weights for Gaussian quadrature with Jacobi weights (α,β).
"""
function gauss_quad(α, β, N)
    if N == 0
        x = [-(α-β)/(α+β+2)]
        w = [2.]
        return x, w
    end

    J = zeros(N+1, N+1)
    h₁ = @. 2*(0:N)+α+β
    J = diagm(0 => @. -1/2*(α^2-β^2)/(h₁+2)/h₁) +
        diagm(1 => @. 2/(h₁[1:N]+2) * sqrt((1:N) * ((1:N)+α+β) * ((1:N)+α) * ((1:N)+β) / (h₁[1:N]+1) / (h₁[1:N]+3)))
    if α + β < 10*eps()
        J[1,1] = 0.0
    end

    x, V = eigen(Symmetric(J+J'))
    w = @. transpose(V[1,:])^2 * 2^(α+β+1) / (α+β+1) * gamma(α+1) * gamma(β+1) / gamma(α+β+1)
    return x[:], w[:]
end

"""
    jacobiP(x, α, β, N, tmp_array=zeros(N+1))
    jacobiP(x::AbstractArray, α, β, N)

Evaluate the Jacobi Polynomial (α, β) of order N at points x.
Can optionally specify `tmp_array` of length > N+1 to avoid allocations.
"""
function jacobiP(x, α, β, N, tmp_array=zeros(N+1))

    PL = tmp_array
    γ₀ = 2^(α+β+1) / (α+β+1) * gamma(α+1) * gamma(β+1) / gamma(α+β+1)
    PL[1] = 1.0 / sqrt(γ₀)
    if N == 0
        return PL[1]
    end

    γ₁ = (α+1) * (β+1) / (α+β+3) * γ₀
    PL[2] = ((α+β+2) * x/2 + (α-β)/2) / sqrt(γ₁)
    if N == 1
        return PL[N+1]
    end

    aold = 2 / (2+α+β) * sqrt((α+1) * (β+1) / (α + β + 3))

    for i = 1:N-1
        h₁ = 2i + α + β
        anew = 2 / (h₁+2) * sqrt((i+1) * (i+1+α+β) * (i+1+α) * (i+1+β) / (h₁+1) / (h₁+3))
        bnew = -(α^2-β^2) / h₁ / (h₁+2)
        PL[i+2] = 1 / anew * (-aold * PL[i] + (x - bnew) * PL[i+1])
        aold = anew
    end

    return PL[N+1]
end

# vectorized version.
# TODO: deprecate?
function jacobiP(x::AbstractArray, α, β, N)

    PL = zeros(N+1, length(x))
    γ₀ = 2^(α+β+1) / (α+β+1) * gamma(α+1) * gamma(β+1) / gamma(α+β+1)
    @. PL[1,:] = 1.0/sqrt(γ₀)
    if N == 0
        return vec(PL)
    end

    γ₁ = (α+1) * (β+1) / (α+β+3) * γ₀
    @. PL[2, :] = ((α+β+2) * x/2 + (α-β)/2) / sqrt(γ₁)
    if N == 1
        return PL[N+1,:]
    end

    aold = 2 / (2+α+β) * sqrt((α+1) * (β+1) / (α + β + 3))

    for i = 1:N-1
        h₁ = 2i+α+β
        anew = 2 / (h₁+2) * sqrt((i+1) * (i+1+α+β) * (i+1+α) * (i+1+β) / (h₁+1) / (h₁+3))
        bnew = -(α^2-β^2) / h₁ / (h₁+2)
        @. PL[i+2,:] = 1 / anew * (-aold * PL[i,:] .+ (x .- bnew) .* PL[i+1,:])
        aold = anew
    end

    return PL[N+1,:]
end

"""
    grad_jacobiP(r, α, β, N, tmp_array=())

Evaluate derivative of Jacobi polynomial (α, β) of order N at points r.
Pass in `tmp_array = tuple(zeros(N+1))` to avoid allocations.
"""

function grad_jacobiP(r, α, β, N, tmp_array=())
    if N != 0
        return sqrt(N * (N+α+β+1)) * jacobiP(r, α+1, β+1, N-1, tmp_array...)
    else
        return zero(r)
    end
end


using OceanTurb, Printf

@use_pyplot_utils

# # Operations on fields
#
# Generate a grid with 100 points and height H=10. By default the
# base of the grid is at z=-H and the top of the grid is at z=0.
grid = UniformGrid(100, 10)

# Define a function that will generate our field's data.
c0(z) = exp(1.2z)

# Create a field of zeros on the grid.
c = CellField(grid)

# Set c to the function c0.
OceanTurb.set!(c, c0)

# Calculate the derivative of c and set it to cz.
cz = ∂z(c)

# Plot the results
plot(c, label=L"c = e^{1.2z}")
plot(cz, label=L"\partial_z c = 1.2 e^{1.2z}")
legend()

title("A field and it's derivative")
xlabel("\$ c \$ and \$ \\partial_z c \$")
ylabel(L"z")

# # Interpolation
#
# Next we demonstrate how to interpolate from one grid to another,
# respecting the integral budget of a quantity.

grid1 = UniformGrid(30, 10)
grid2 = UniformGrid(10, 10)
grid3 = UniformGrid(3, 10)

c1 = CellField(grid1)
c2 = CellField(grid2)
c3 = CellField(grid3)

OceanTurb.set!(c1, c)
OceanTurb.set!(c2, c)
OceanTurb.set!(c3, c)

@show integral(c)
@show integral(c1)
@show integral(c2)
@show integral(c3)

plot(c, ".", label="high-res \$ c \$")
plot(c1, "^", label="medium-res \$ c \$")
plot(c2, "x", label="low-res \$ c \$")
plot(c3, "+", label="crazy low-res \$ c \$")

legend()

title("\$ e^{1.2z} \$ at different resolutions")
xlabel(L"c")
ylabel(L"z")


N = Int(5e1)  # number of data
M = Int(1e5)  # number of realisations

mu = [7; 10] * 1.0  # mean of data
c = [5 3; 3 7] * 1.0  # covarianza en float

# generate data
using Distributions
probMVgaus = MvNormal(mu,c)
x = rand(probMVgaus, N*M)
x = reshape(x, N, M, 2)

# estimo los mu y c de cada relizacion
muest = mean(x, 1)
dif = reshape(x .- muest,N,M,2,1) # ver como hacer esta resta facilmente
cest = sum(dif .* permutedims(dif,[1,2,4,3]), 1) / (N - 1)

# saco la media y varianza de los estimadores entre todas las realizaciones
# media y var de mu
muExp = mean(muest, 2)
difmu = reshape(muest .- muExp, M, 2, 1)
muVar = sum(difmu .* permutedims(difmu,[1,3,2]), 1) / (M - 1)
# media y var de c
cExp = mean(cest, 2)
difc = reshape(cest .- cExp, M,2,2,1,1)
cVar = sum(difc .* permutedims(difc, [1,5,4,3,2]), 1) / (M - 1)
cVar = cVar[1,:,:,:,:]

# saco las cuentas analiticas
expMu = mu
VarMu = c / N
expC = c
# no es necesario trasponer porque c es simetrica
cShaped = reshape(c,2,2,1,1)
VarCnn = cShaped .* permutedims(cShaped, [4,3,2,1])
nCons = (2 * N - 1) / (N - 1)^2
VarC = VarCnn * nCons

print("numerico: \n", cVar, "\n analitico \n", VarC)
print(cVar / nCons)
print(VarCnn)


discount(t, κ) = 1/(1 + κ*t)

function prob(κ, τ, t_ss, t_ll, r_ss, r_ll)
    u_ll = r_ll * discount(t_ll, κ)
    u_ss = r_ss * discount(t_ss, κ)
    return 1/(1 + exp(-τ * (u_ll - u_ss)))
end

function loglike(κ, τ, t_ss, t_ll, r_ss, r_ll, data)
    p = prob(κ, τ, t_ss, t_ll, r_ss, r_ll)
    p = p == 1 ? 1 - eps() : p
    p = p == 0 ? eps() : p
    LL = data ? log(p) : log(1 - p)
    # println(" choice ", data, " kappa ", κ, " tau ", τ, " t_ss ", t_ss,
    #     " t_ll ", t_ll, " r_ss ", r_ss, " r_ll ", r_ll, " LL ", LL)
    return LL
end

function simulate(κ, τ, t_ss, t_ll, r_ss, r_ll)
    p = prob(κ, τ, t_ss, t_ll, r_ss, r_ll)
    return rand() ≤ p ? true : false
end


using Plots

function booth(x)
    return (x[1] + 2x[2] - 7)^2 + (2x[1] + x[2] - 5)^2;
end

# plot

contour(-30:0.1:30, -30:0.1:30, (x, y)->booth([x, y]))


# This file investigates mean square error (MSE) with respect to number of interpolation points in 1D interpolation

using FractalTools
using Makie

# Construct interpolation data
f(x) = sin(2π * x)
xi = 0.
xf = 1.
freevar = 0.001
npts = 10 : 2 : 50
ntpts = 10 * npts[end]

# Compute errors
mse = map(npts) do npt
    # Construct interpolation data
    x = collect(range(xi, xf, length=npt))
    y = f.(x)
    pts = collect.(zip(x, y))

    # Construct interpolation
    interp = interpolate(pts, Interp1D(freevar))

    # Construct test data
    xt = collect(range(xi, xf, length=ntpts))

    # Compute error
    sum((f.(xt) - interp.(xt) ).^2) / length(xt)
end

# Plot mse
fig = Figure()
ax = fig[1, 1] = Axis(fig, xlabel="Number of Points", ylabel="MSE", title="1D Interpolation MSE")
stem!(ax, npts, mse, color=:black)
save(joinpath(@__DIR__, "interp1d_mse.png"), fig)
display(fig)


set_theory = [
    Axiom("Definition of the universal set", P ⊂ Ω, TRUE),
    Axiom("Definition of the universal set", x ∈ Ω, TRUE),
    Axiom("Definition of the empty set", ∅ ⊂ A, TRUE),
    Axiom("Definition of the empty set", x ∈ ∅, FALSE),
    Axiom("Equal sets are subsets", A ⊂ A, TRUE),

    Axiom("Definition of Element Of", x ∈ (A ∩ B), (x ∈ A) ∧ (x ∈ B)),
    Axiom("Definition of Element Of", (x ∈ A) ∧ (x ∈ B), x ∈ (A ∩ B)),
    Axiom("Definition of Element Of", x ∈ (A ∪ B), (x ∈ A) ∨ (x ∈ B)),
    Axiom("Definition of Element Of", (x ∈ A) ∨ (x ∈ B), x ∈ (A ∪ B)),

    Axiom("Flatten Condition Set", ConditionSet(y, y ∈ A), A),  # todo: change to f(x)
    Axiom("Expand Condition Set", A, ConditionSet(_y, _y ∈ A)),  # todo: change to f(x)
    Axiom("Definition of Subset", A ⊂ B, (_y ∈ A) ⟹ (_y ∈ B)),
    Axiom("Definition of Subset", (_y ∈ A) ⟹ (_y ∈ B), A ⊂ B),

    Axiom("Definition of the universal set", x ∈ (P ∪ Ω), x ∈ Ω),

    #Theorem("Intersections are subsets of Unions", x ∈ (A ∩ B), x ∈ (A ∪ B), Proof([])),
    #Theorem("Intersections are subsets of Unions", (A ∩ B) ⊂ (A ∪ B), TRUE, Proof([])),
]


#  Negative Binomial type 1 and type II log likelihood
#  The parameterization follows
#  Deb and Trivedi, J. Appl. Econometrics,
#  V. 12, 1997, pp. 313 - 336.
using SpecialFunctions, Distributions
function  negbin(θ, y, x, nbtype)
    n, k = size(x)
    β = θ[1:k]
    eps = 1e-8
    λ = eps .+ exp.(x*β)
    α = eps .+ exp(θ[end])
    nbtype == 1 ? r = λ./α : r = ones(n)/α
    p = r ./ (r + λ)
    all(r .> 0.0) & all(p .> 0.0) & all(p .< 1.0) ? log.(pdf.(NegativeBinomial.(r, p),y)) : -Inf
end


using LinearAlgebra

# Define and use function
f(x) = x[1] + x[2] - x[4]^2     # f(x) = x(1) + x(2) − x(4)^2
f([-1,0,1,2])                   # -5

# Superposition
# eval f at a linear combination of 2 vectors
# == forming the linear combination of f eval at the 2 vectors

a = [-2, 0, 1, -3];
f(x) = a'*x

x=[2,2,-1,1]; y=[0,1,-1,0];
alpha = 1.5;  beta = -3.7;
lhs = f(alpha * x + beta * y)

rhs = alpha * f(x) + beta * f(y)

# Verify superposition for f(e(3)) = a(3)
e3 = [0, 0, 1, 0];
f(e3)                           # 1.0

using Statistics
avg(x) = (ones(length(x)) / length(x))'*x;
x = [1, -3, 2, -1];
avg(x)


function conv(signal1::Vector{Complex}, signal2::Vector{Complex})
    n = length(signal1) + length(signal2)
    out = Vector{Complex}(n)
    sum = 0

    for i = 0:n
        for j = 0:i
            if(j < length(signal1))
                sum += signal1[j] * signal2[i-j]
            end
        end
        out[i] = sum
        sum = 0
    end

    return out
end

function conv_fft(signal1::Vector{Complex}, signal2::Vector{Complex})
    return ifft(fft(signal1).*fft(signal2))
end


struct CronbachAlpha{T <: Real}
    alpha::T
    dropped::Vector{T}
end

function Base.show(io::IO, x::CronbachAlpha)
    @printf(io, "Cronbach's alpha for all items: %.4f\n", x.alpha)
    isempty(x.dropped) && return
    println(io, "\nCronbach's alpha if an item is dropped:")
    for (idx, val) in enumerate(x.dropped)
        @printf(io, "item %i: %.4f\n", idx, val)
    end
end

"""
    cronbachalpha(covmatrix::AbstractMatrix{<:Real})

Calculate Cronbach's alpha (1951) from a covariance matrix `covmatrix` according to
the [formula](https://en.wikipedia.org/wiki/Cronbach%27s_alpha):

```math
\\rho = \\frac{k}{k-1} (1 - \\frac{\\sum^k_{i=1} \\sigma^2_i}{\\sum_{i=1}^k \\sum_{j=1}^k \\sigma_{ij}})
```

where ``k`` is the number of items, i.e. columns, ``\\sigma_i^2`` the item variance,
and ``\\sigma_{ij}`` the inter-item covariance.

Returns a `CronbachAlpha` object that holds:

* `alpha`: the Cronbach's alpha score for all items, i.e. columns, in `covmatrix`; and
* `dropped`: a vector giving Cronbach's alpha scores if a specific item,
  i.e. column, is dropped from `covmatrix`.

# Example
```jldoctest
julia> using StatsBase

julia> cov_X = [10 6 6 6;
                6 11 6 6;
                6 6 12 6;
                6 6 6 13];

julia> cronbachalpha(cov_X)
Cronbach's alpha for all items: 0.8136

Cronbach's alpha if an item is dropped:
item 1: 0.7500
item 2: 0.7606
item 3: 0.7714
item 4: 0.7826
```
"""
function cronbachalpha(covmatrix::AbstractMatrix{<:Real})
    isposdef(covmatrix) || throw(ArgumentError("Covariance matrix must be positive definite."))
    k = size(covmatrix, 2)
    k > 1  || throw(ArgumentError("Covariance matrix must have more than one column."))
    v = vec(sum(covmatrix, dims=1))
    σ = sum(v)
    for i in axes(v, 1)
        v[i] -= covmatrix[i, i]
    end
    σ_diag = sum(i -> covmatrix[i, i], 1:k)

    alpha = k * (1 - σ_diag / σ) / (k - 1)
    if k > 2
        dropped = typeof(alpha)[(k - 1) * (1 - (σ_diag - covmatrix[i, i]) / (σ - 2*v[i] - covmatrix[i, i])) / (k - 2)
                                for i in 1:k]
    else
        # if k = 2 do not produce dropped; this has to be also
        # correctly handled in show
        dropped = Vector{typeof(alpha)}()
    end
    return CronbachAlpha(alpha, dropped)
end


using Distributions, Random, Plots, LaTeXStrings; pyplot()

seed = 1
N = 100
lamGrid = 0.01:0.01:0.99

theorM(lam) = mean(Uniform(0,2*lam*(1-lam)))
estM(lam) = mean(rand(Uniform(0,2*lam*(1-lam)),N))

function estM(lam,seed)
    Random.seed!(seed)
    estM(lam)
end

trueM = theorM.(lamGrid)
estM0 = estM.(lamGrid)
estMCRN = estM.(lamGrid,seed)

plot(lamGrid,trueM,
	c=:black, label="Expected curve")
plot!(lamGrid,estM0,
	c=:blue, label="No CRN estiamte")
plot!(lamGrid,estMCRN,
	c=:red, label="CRN estimate",
	xlims=(0,1), ylims=(0,0.4), xlabel=L"\lambda", ylabel = "Mean")


# 10/24/2020 © Shuowen Chen
# This script simulates CAPM data using Tauchen & Hussey (1991) method
# as in Wright (2003) and explores sensitivity of Euler equation estimation
# to number of discretization points.

# load packages
using Random, Distributions, LinearAlgebra, Statistics, FastGaussQuadrature, Optim
# set seed
Random.seed!(123)

#======= Defining Functions =======#
# univariate tauchen
function tauchen(m, ρ, N, σ²)
    # INPUTS:
    # 1. m: multiple of unconditional standard deviation
    # 2. ρ: AR(1) coefficient
    # 3. N: number of grid points
    # 4. σ²: variance of the white noise

    # OUTPUTS:
    # 1. Transition probability matrix
    # 2. Sequence of grid points

    # unconditional standard deviation of the AR(1)
    std_un = sqrt(σ²/(1 - ρ^2));
    # max and min of the grid
    ymax = m*std_un;
    ymin = -ymax;
    # length of adjacent grid point
    w = (ymax - ymin)/(N - 1);
    # lay out the grids
    ygrid = collect(ymin:w:ymax);
    # Transition probability matrix (N by N)
    P = zeros(N, N);
    P[:, 1] = cdf.(Normal(0, 1), ((ymin .- ρ*ygrid .+ w/2)./sqrt(σ²)));
    P[:, end] = 1 .- cdf.(Normal(0, 1), ((ymax .- ρ*ygrid .- w/2)./sqrt(σ²)));
    for k = 2:(N - 1)
        P[:, k] = cdf.(Normal(0, 1), (ygrid[k] .- ρ*ygrid .+ w/2)./sqrt(σ²)) -
        cdf.(Normal(0, 1), (ygrid[k] .- ρ*ygrid .- w/2)./sqrt(σ²));
    end
    return(transition = P, grid = ygrid);
end

# Univariate Tauchen and Hussey (1991)
function tauchenhussey(ρ, a, σ², N)
    # PURPOSE:
    # Discretize the following AR(1) process:
    # z_{t+1} = a + ρ*z_{t} + epsilon_{t+1}
    # INPUTS:
    # 1. ρ: AR(1) coefficient
    # 2. a: intercept
    # 3. σ²: variance of epsilon
    # 4. N: number of grid points

    # OUTPUTS:
    # 1. P: transition probability matrix
    # 2. grid: vector of grid points

    # A function from FastGaussQuadrature
    (x, w) = gausshermite(N);
    # Gaussian nodes and weights for normal distribution
    # Acknowledgement: I benefit from reading slides by Karen Kopecky
    # http://www.karenkopecky.net/Teaching/eco613614/Notes_DiscretizingAR1s.pdf

    # unconditional mean of the process
    μ = a / (1 - ρ);
    # Establish the grid using Gaussian-Hermite node
    grid = x .* sqrt(2*σ²) .+ μ;
    # Rescale the Gaussian-Hermite weight
    w = w ./ sqrt(pi);
    # Compute transition probability matrix
    P = zeros(N, N);
    for i = 1:N
        for j = 1:N
            EZprime = a + ρ*grid[i];
            # Gauchen-Hussey weight
            P[i, j] = w[j] * pdf(Normal(EZprime, sqrt(σ²)), grid[j]) /
            pdf(Normal(μ, sqrt(σ²)), grid[j]);
        end
    end
    # Normalization
    for i = 1:N
        P[i, :] = P[i, :] ./ sum(P[i, :])
    end
    return(P, grid);
end

# An auxiliary function to be called in tauchenhusseyvar
function productrule(N, EZ, μ, Σ, grid_j, weight_j, k)
    # Purpose:
    # Compute individual transition prob from state i to j and aggregate
    # by Gauss product rule (tensor product)

    # Inputs:
    # 1. N: number of grid points
    # 2. EZ: expected value in state i (k by 1)
    # 3. μ: unconditional mean of VAR process (k by 1)
    # 4. Σ: var-cov matrix of VAR(1) process
    # 5. grid_j: grid points for the next state
    # 6. weight_j: Gauss-Hermite weight for the next state
    # 7. k: number of variables in the VAR

    # Outputs:
    # 1. prob: aggregate transition probability from state i to j

    den_ind = zeros(1, k);
    for i = 1:k
        den_ind[i] = weight_j/sqrt(pi) * pdf(Normal(EZ[i], diag(sqrt(Σ))[i]),
        grid_j[i]) / pdf(Normal(μ[i], diag(sqrt(Σ))[i]), grid_j[i])
    end
    den_agg = prod(den_ind);
    return(den_agg);
end

# Bivariate Tauchen-Hussey (1991)
function tauchenhusseyvar(k, A, B, Σ, N)
    # PURPOSE: Discretize the following VAR(1) process
    # Y_t = A + B * Y_t-1 + e_t,
    # where e_t ∼ N(0, Σ)
    # Acknowledgement: I benefit from reading the quadrature code by Tauchen,
    # which is available in his website.

    # INPUTS:
    # 1. k: number of variables in the VAR
    # 2. A: constant of VAR(1) (k by 1)
    # 3. B: VAR(1) coefficient matrix (k by k)
    # 4. Σ: iid var-cov matrix
    # 5. N: number of discrete points for each variable k in the VAR

    # OUTPUTS:
    # 1. P: transition matrix (Nᵏ by Nᵏ)
    # 2. Grid: grid points (k by Nᵏ).
    # Note: number of states in the system is equal to number of grids * number of lags,
    # since we consider VAR(1), number of states = number of grid points

    # unconditional mean of the process
    μ = inv(I(k) - B) * A;
    # obtain Gaussian-Hermite nodes and weights
    (x, w) = gausshermite(N);
    # Grid for each variable
    grid_ind = repeat(x, outer = [1, k])' .* sqrt.(2 * diag(Σ)) .+ μ;
    # total grid (k by N^k) and GH weights (N^k) using product rule
    grid = zeros(k, N^k);
    w_tot = zeros(k, N^k);
    for i = 1:k
        grid[i, :] = repeat(grid_ind[i, :], inner = [N^(k-i), 1],
                            outer = [N^(i-1), 1]);
        w_tot[i, :] = repeat(w, inner = [N^(k-i), 1], outer = [N^(i-1), 1]);
    end
    w_tot = prod(w_tot, dims = 1);
    # Compute transition probability matrix
    P = zeros(N^k, N^k);
    for i = 1:N^k
        for j = 1:N^k
            # expected value in state i at time t-1
            EZprime = A + B*grid[:, i]; # k by 1
            # Gauchen-Hussey
            P[i, j] = productrule(N, EZprime, μ, Σ, grid[:, j], w_tot[j], k);
        end
    end
    # Normalization
    for i = 1:N^k
        P[i, :] = P[i, :] ./ sum(P[i, :])
    end
    return(grid, P);
end

# Simulate data using Markov transition probability and grid points
function simulationMC(Pr_mat, grid, T, sim)
    # INPUTS:
    # 1. Pr_mat: transition matrix
    # 2. grid: grid points
    # 3. T: number of observations in each simulation
    # 4. sim: number of simulations
    # 5. k: number of variables

    # OUTPUT:
    # 1. Z: simulated states using MC (k by T for # sim times)
    # Economic meaning: log consumption and dividend growth data
    # 2. state: state index for stock return simulation

    # first obs of the MC (middle state)
    initial_state = convert(Int, floor((length(grid)/size(grid)[1]+1)/2));
    # placeholder for simulated states
    Z = zeros(k, T, sim);
    # initial state
    Z[:, 1, :] .= grid[:, initial_state]; # need to broadcast

    # placeholder for state index
    state = zeros(T, sim);
    # initial state
    state[1, :] .= initial_state;

    # This follows from Adda and Cooper (2003)
    # standard uniformly distributed rv for
    # determining the new state
    u = rand(Uniform(), (T-1, sim));
    # store ∑_{l=1}^{j}π_{i,l} for j = 1,...,N
    Pr_sum = cumsum(Pr_mat, dims = 2);
    old_state = initial_state;
    for i = 1:sim
        for t = 2:T
            new_state = findfirst(u[t-1, i] .<= Pr_sum[old_state, :])
            Z[:, t, i] = grid[:, new_state];
            # store the state index
            state[t, i] = new_state;
            # move to the next obs
            old_state = new_state;
        end
    end
    state = convert.(Int, state);
    return(Z, state);
end

# a function that directly simulates the VAR
function simulationVAR(T, sim, A, B, Σ, k)
    # INPUTS:
    # 1. T: number of observations in each simulation
    # 2. sim: number of simulations
    # 3. A: constant of VAR(1) (k by 1)
    # 4. B: VAR(1) coefficient matrix (k by k)
    # 5. Σ: iid var-cov matrix
    # 6. K: number of variables

    # OUTPUTS:
    # 1. TS: simulated time series (k by 1)

    # normal errors e_t ∼ N(0, Σ) for t = 2,...,T and each simulation
    err = zeros(k, T-1, sim);
    for i = 1:sim
        err[:, :, i] = rand(MvNormal([0; 0], Σ), T-1);
    end
    # placeholder for simulated VAR
    TS = zeros(k, T, sim);
    # initial observation (steady state)
    for i = 1:sim
        TS[:, 1, i] = inv(I(k) - B) * A;
    end
    # Then use the relationship Y_t = A + B*Y_t-1 + e_t
    for i = 1:sim
        for j = 2:T
            # Note the index for err is due to specification of err
            TS[:, j, i] = A + B*TS[:, j-1, i] + err[:, j-1, i];
        end
    end
    return(TS)
end

# Simulate stock return
function simulateReturn(β, γ, P, CG, DG, N, stateindex)
    # INPUTS:
    # 1. β: discount factor
    # 2. γ: CRRA preference
    # 3. P: probability transition matrix
    # 4. CG: consumption growth
    # 5. DG: dividend growth
    # 6. N: number of grid points for each variable

    # OUTPUTS:
    # 1. smreturn: simulated return data

    # Define transition matrix for return
    B = β*P;
    for j = 1:N^2
        B[:, j] = B[:, j] .* CG[j].^(-γ) .* DG[j];
    end
    # State contingent price-dividend ratio (N^2 by 1)
    psm = inv(I(N^2) - B) * B * ones(N^2);
    # Realized stock return from state i to j (N^2 by N^2)
    Rmat = zeros(N^2, N^2);
    for i = 1:N^2
        for j = 1:N^2
            Rmat[i, j] = DG[j]*(1+psm[j])/psm[i];
        end
    end
    # simulate the return data
    T = size(stateindex)[1] - 1;
    smreturn = zeros(T, 1);
    for t = 1:T
        smreturn[t] = Rmat[stateindex[t], stateindex[t+1]];
    end
    return(smreturn);
end

# A function to conduct 2 step GMM
function GMM2step(moments, gradients, k, initial, samplesize)
    # INPUTS:
    # 1. moments: an anonymous function that maps theta to moment conditions
    # 2. gradients: an anoymous function that mapes theta to gradient conditions
    # 3. k: number of moments
    # 4. initial: initial values for estimation (must be floating)
    # 5. samplesize: sample size of the moments
    # OUTPUTS:
    # 1. est: parameter estimate
    # 2. se: standard errors

    # sample average moments
    sample_m = theta -> vec(mean(moments(theta), dims = 1));
    # first step weighting matrix
    W = I(k);
    # objective function in the first step
    obj = theta -> sample_m(theta)'*W*sample_m(theta);
    # minimization using limited BFGS
    firststep = Optim.optimize(obj, initial, LBFGS());
    # second step weighting matrix
    W2 = inv(cov(moments(firststep.minimizer)));
    # objective function in the second step
    obj2 = theta -> sample_m(theta)'*W2*sample_m(theta);
    secondstep = Optim.optimize(obj2, initial, LBFGS());
    # compute the gradient
    # dimension of D: k (# of pars) by p (# of moments)
    D = reshape(mean(gradients(secondstep.minimizer), dims = 1), k, length(initial));
    var = inv(D'*inv(cov(moments(secondstep.minimizer)))*D);
    # compute standard errors
    se = sqrt.(diag(var)./samplesize);
    return(firststep.minimizer, secondstep.minimizer, se);
end

#======= Main Program Starts Here =======#
# Setting parameters
β = 1.139; # utility discount factor
γ = 13.7; # CRRA parameter
k = 2; # number of variables
A = [0.021; 0.04]; # Calibrated values as in Wright (2003)
B = [-0.161 0.017; 0.414 0.117];
Σ = [0.0012 0.00177; 0.00177 0.014];
N = 4; # number of grids for each variable
T = 400; # number of time periods
sim = 400; # number of simulations

# Generate grid points and transition probability
(grid, P) = tauchenhusseyvar(k, A, B, Σ, N);
# Simulate log consumption and dividend growth
(timeseries, stateindex) = simulationMC(P, grid, T, sim);

# Consumption and dividend growth
CG_tot = exp.(timeseries[1, :, :]);
DG_tot = exp.(timeseries[2, :, :]);
# use state index to generate return stock market return
R = zeros(T-1, sim);
for i = 1:sim
    R[:, i] = simulateReturn(β, γ, P, CG_tot[:, i], DG_tot[:, i], N, stateindex[:, i]);
end

# conduct estimation
container = zeros(sim, 2);
for i = 1:sim
    # intercept, lagged return and lagged consumption growth
    ins = [ones(T-2) R[1:end-1, i] CG_tot[1:end-2, i]];
    # sample size of the moments
    t_mom = size(instruments)[1];
    # Construct Euler equation
    euler = theta -> (theta[1] .* R[2:end, i] .* CG_tot[2:end-1, i] .^ (-theta[2]) .-1) .* ins;

    grad = theta -> hcat(R[2:end, i] .* CG_tot[2:end-1, i] .^ (-theta[2]) .* ins,
                             -(theta[1] .* R[2:end, i] .* CG_tot[2:end-1, i] .^ (-theta[2]) .* log.(CG_tot[2:end-1, i])) .* ins);

    gest = GMM2step(euler, grad, 3, [0.0, 13.0], t_mom);
    container[i,:] = gest[2];
end

mean(container, dims =1)


#' ---
#' title: Network Flows
#' ---

#' **Originally Contributed by**: Arpit Bhatia

#' In graph theory, a flow network (also known as a transportation network) is a directed graph where
#' each edge has a capacity and each edge receives a flow. The amount of flow on an edge cannot exceed the capacity of the edge.
#' Often in operations research, a directed graph is called a network, the vertices are called nodes and the edges are called arcs.
#' A flow must satisfy the restriction that the amount of flow into a node equals the amount of flow out of it,
#' unless it is a source, which has only outgoing flow, or sink, which has only incoming flow.
#' A network can be used to model traffic in a computer network, circulation with demands, fluids in pipes,
#' currents in an electrical circuit, or anything similar in which something travels through a network of nodes.

using JuMP
using GLPK

#' ## The Shortest Path Problem
#' Suppose that each arc $(i, j)$ of a graph is assigned a scalar cost $a_{i,j}$, and
#' suppose that we define the cost of a forward path to be the sum of the costs of its arcs.
#' Given a pair of nodes, the shortest path problem is to find a forward path that connects these nodes and has minimum cost.

#' $$
#' \begin{align*}
#' \min && \sum_{\forall e(i,j) \in E} a_{i,j} \times x_{i,j} \\
#' s.t. && b(i) = \sum_j x_{ij} - \sum_k x_{ki} = \begin{cases}
#' 1 &\mbox{if $i$ is the starting node,} \\
#' -1 &\mbox{if $i$ is the ending node,} \\
#' 0 &\mbox{otherwise.} \end{cases} \\
#' && x_{e} \in \{0,1\} && \forall e \in E
#' \end{align*}
#' $$

#'

G = [
0 100 30  0  0;
0   0 20  0  0;
0   0  0 10 60;
0  15  0  0 50;
0   0  0  0  0
]

n = size(G)[1]

shortest_path = Model(with_optimizer(GLPK.Optimizer))

@variable(shortest_path, x[1:n,1:n], Bin)
@constraint(shortest_path, [i = 1:n, j = 1:n; G[i,j] == 0], x[i,j] == 0) # Arcs with zero cost are not a part of the path as they do no exist
@constraint(shortest_path, [i = 1:n; i != 1 && i != 2], sum(x[i,:]) == sum(x[:,i])) # Flow conservation constraint
@constraint(shortest_path, sum(x[1,:]) - sum(x[:,1]) == 1) # Flow coming out of source = 1
@constraint(shortest_path, sum(x[2,:]) - sum(x[:,2]) == -1) # Flowing coming out of destination = -1 i.e. Flow entering destination = 1
@objective(shortest_path, Min, sum(G .* x))

optimize!(shortest_path)
@show objective_value(shortest_path);
@show value.(x);

#' ## The Assignment Problem
#' Suppose that there are $n$ persons and $n$ objects that we have to match on a one-to-one basis.
#' There is a benefit or value $a_{i,j}$ for matching person $i$ with object $j$, and
#' we want to assign persons to objects so as to maximize the total benefit.
#' There is also a restriction that person $i$ can be assigned to object $j$ only if $(i, j)$ belongs to a given set of pairs $A$.
#' Mathematically, we want to find a set of person-object pairs $(1, j_{1}),..., (n, j_{n})$ from $A$ such that
#' the objects $j_{1},...,j_{n}$ are all distinct, and the total benefit $\sum_{i=1}^{y} a_{ij_{i}}$ is maximized.

#' $$
#' \begin{align*}
#' \max && \sum_{(i,j) \in A} a_{i,j} \times y_{i,j} \\
#' s.t. && \sum_{\{j|(i,j) \in A\}} y_{i,j} = 1 && \forall i = \{1,2....n\} \\
#' && \sum_{\{i|(i,j) \in A\}} y_{i,j} = 1 && \forall j = \{1,2....n\} \\
#' && y_{i,j} \in \{0,1\} && \forall (i,j) \in \{1,2...k\}
#' \end{align*}
#' $$

#'

G = [
6 4 5 0;
0 3 6 0;
5 0 4 3;
7 5 5 5;
]

n = size(G)[1]

assignment = Model(with_optimizer(GLPK.Optimizer))
@variable(assignment, y[1:n,1:n], Bin)
@constraint(assignment, [i = 1:n], sum(y[:,i]) == 1) # One person can only be assigned to one object
@constraint(assignment, [j = 1:n], sum(y[j,:]) == 1) # One object can only be assigned to one person
@objective(assignment, Max, sum(G .* y))

optimize!(assignment)
@show objective_value(assignment);
@show value.(y);

#' ## The Max-Flow Problem
#' In the max-flow problem, we have a graph with two special nodes: the $source$, denoted by $s$, and the $sink$, denoted by $t$.
#' The objective is to move as much flow as possible from $s$ into $t$ while observing the capacity constraints.

#' $$
#' \begin{align*}
#' \max && \sum_{v:(s,v) \in E} f(s,v) \\
#' s.t. && \sum_{u:(u,v) \in E} f(u,v)  = \sum_{w:(v,w) \in E} f(v,w) && \forall v \in V - \{s,t\} \\
#' && f(u,v) \leq c(u,v) && \forall (u,v) \in E \\
#' && f(u,v) \geq 0 && \forall (u,v) \in E
#' \end{align*}
#' $$

#'

G = [
0 3 2 2 0 0 0 0
0 0 0 0 5 1 0 0
0 0 0 0 1 3 1 0
0 0 0 0 0 1 0 0
0 0 0 0 0 0 0 4
0 0 0 0 0 0 0 2
0 0 0 0 0 0 0 4
0 0 0 0 0 0 0 0
]

n = size(G)[1]

max_flow = Model(with_optimizer(GLPK.Optimizer))

@variable(max_flow, f[1:n,1:n] >= 0)
@constraint(max_flow, [i = 1:n, j = 1:n], f[i,j] <= G[i,j]) # Capacity constraints
@constraint(max_flow, [i = 1:n; i != 1 && i != 8], sum(f[i,:]) == sum(f[:,i])) # Flow conservation contraints
@objective(max_flow, Max, sum(f[1, :]))

optimize!(max_flow)
@show objective_value(max_flow);
@show value.(f);


"""Returns a matrix of water depth for the whole domain that contains a
Gaussian seamount in the middle. Water depth, heigth and width of the
seamount are adjusted with the constants water_depth, topofeat_height and topofeat_width."""
function seamount()
    xx_T,yy_T = meshgrid(x_T_halo,y_T_halo)
    bumpx = exp.(-((xx_T .- Lx/2).^2)/(2*topofeat_width^2))
    bumpy = exp.(-((yy_T .- Ly/2).^2)/(2*topofeat_width^2))

    H = water_depth .- topofeat_height*bumpx.*bumpy
    return Numtype.(H)
end

"""Returns a matrix of water depth for the whole domain that contains a
meridional Gaussian ridge in the middle. Water depth, heigth and width of the
ridge are adjusted with the constants water_depth, topofeat_height and topofeat_width."""
function ridge()
    xx_T,yy_T = meshgrid(x_T_halo,y_T_halo)
    bumpx = exp.(-((xx_T .- Lx/2).^2)/(2*topofeat_width^2))

    H = water_depth .- topofeat_height*bumpx
    return Numtype.(H)
end

"""Same as ridge() but for 3 ridges at 1/4,1/2,3/4 of the domain."""
function ridges()
    xx_T,yy_T = meshgrid(x_T_halo,y_T_halo)

    # bumps in x direction
    # shift slightly left/right to avoid a symmetric solution
    b0x = exp.(-(xx_T.^2)/(2*topofeat_width^2))
    b1x = exp.(-((xx_T .- 0.99*Lx/4).^2)/(2*topofeat_width^2))
    b2x = exp.(-((xx_T .- 1.01*Lx/2).^2)/(2*topofeat_width^2))
    b3x = exp.(-((xx_T .- 0.99*3*Lx/4).^2)/(2*topofeat_width^2))
    b4x = exp.(-((xx_T .- Lx).^2)/(2*topofeat_width^2))

    th = topofeat_height    # for convenience
    H = water_depth .- th*b0x .- th*b1x .- th*b2x .- th*b3x .- th*b4x
    return Numtype.(H)
end

"""Returns a matrix of constant water depth specified by the constant water_depth."""
function flat_bottom()
    H = fill(water_depth,(nx+2*haloη,ny+2*haloη))
    return Numtype.(H)
end

# rename for convenience
if topography_feature == "ridge"
    topography = ridge
elseif topography_feature == "ridges"
    topography = ridges
elseif topography_feature == "seamount"
    topography = seamount
elseif topography_feature == "flat"
    topography = flat_bottom
else
    throw(error("Topography feature not correctly declared."))
end


"""
    (uvels, vvels, wvels) = streamfunction_xz()

Define a two-dimensional flow field in the x-z directions from a simple streamfunction.
"""
function streamfunction_xz()
    scal = 2e-1
    f(x, y, z) = scal*sin(x*2π/128)*sin(z*π/128) #stream function

    ϕcorners=[f(x,0.,z) for x in 0:128, z in -128:0]
    ϕcenters=[f(x,0.,z) for x in 0.5:128, z in -128:-0.5]

    uu=-diff(ϕcorners,dims=2)[1:end-1,:]
    ww=diff(ϕcorners,dims=1)
    uu=reshape(uu,(128,1,128))
    ww=reshape(ww,(128,1,129))
    uu = reverse(uu, dims=3)
    ww = reverse(ww, dims=3)

    uvels = fill(uu, 2)
    vvels = fill(0*uu, 2)
    wvels = fill(ww, 2)
    uvels = cat(uvels..., dims=4)
    vvels = cat(vvels..., dims=4)
    wvels = cat(wvels..., dims=4)

    return uvels, vvels, wvels, ϕcenters
end

"""
    (uvels, vvels, wvels) = streamfunction_xy()

Define a two-dimensional flow field in the x-y directions from a simple streamfunction
as explained at https://shaddenlab.berkeley.edu/uploads/LCS-tutorial/examples.html#Sec7.1
"""
function streamfunction_xy()
    scal = 3e-1
    f(x, y, z) = scal*(0.3*sin(x*1π/128)*sin(y*2π/128)+0.7*sin(x*2π/128)*sin(y*1π/128)) #stream function

    ϕcorners=[f(x,y,0.) for x in 0:128, y in 0:128]
    ϕcenters=[f(x,y,0.) for x in 0.5:128, y in 0.5:128]

    uu=-diff(ϕcorners,dims=2)[1:end-1,:]
    vv=diff(ϕcorners,dims=1)[:,1:end-1]
    uu=reshape(uu,(128,128,1))
    vv=reshape(vv,(128,128,1))
    ww=zeros(128,128,2)

    uvels = fill(uu, 2)
    vvels = fill(vv, 2)
    wvels = fill(ww, 2)
    uvels = cat(uvels..., dims=4)
    vvels = cat(vvels..., dims=4)
    wvels = cat(wvels..., dims=4)

    return uvels, vvels, wvels, ϕcenters
end


function init_CN(model)
    nx, ny = model.grid.nx, model.grid.ny
    Δx, Δy = model.grid.Δx, model.grid.Δy
    ν, Δt  = model.parameters.ν, model.Δt

    Nx = (nx-2)
    Ny = (ny-2)
    Nxy = Nx*Ny
    βx = 0.5*ν*Δt/(Δx*Δx)
    βy = 0.5*ν*Δt/(Δy*Δy)

    # New a1
    a1 = spzeros(Nxy,Nxy)
    ij_ind = LinearIndices((1:Ny, 1:Nx))
    for ij in ij_ind
        a1[ij,ij] = 1 + 2*βx + 2*βy
    end
    for ij in 1:Nx*(Ny-1)
        a1[view(ij_ind,2:Ny,:)[ij],view(ij_ind,1:Ny-1,:)[ij]] = -βy # Up
        a1[view(ij_ind,1:Ny-1,:)[ij],view(ij_ind,2:Ny,:)[ij]] = -βy # Down
    end
    for ij in 1:(Nx-1)*Ny
        a1[view(ij_ind,:,2:Nx)[ij],view(ij_ind,:,1:Nx-1)[ij]] = -βx # Right
        a1[view(ij_ind,:,1:Nx-1)[ij],view(ij_ind,:,2:Nx)[ij]] = -βx # Left
    end

    # New a2
    a2 = spzeros(Nxy,Nxy)
    ij_ind = LinearIndices((1:Ny, 1:Nx))
    for ij in ij_ind
        a2[ij,ij] = 1 - 2*βx - 2*βy
    end
    for ij in 1:Nx*(Ny-1)
        a2[view(ij_ind,2:Ny,:)[ij],view(ij_ind,1:Ny-1,:)[ij]] = βy # Up
        a2[view(ij_ind,1:Ny-1,:)[ij],view(ij_ind,2:Ny,:)[ij]] = βy # Down
    end
    for ij in 1:(Nx-1)*Ny
        a2[view(ij_ind,:,2:Nx)[ij],view(ij_ind,:,1:Nx-1)[ij]] = βx # Right
        a2[view(ij_ind,:,1:Nx-1)[ij],view(ij_ind,:,2:Nx)[ij]] = βx # Left
    end

    a1_lu = factorize(a1)
    return a1_lu, a2
end
function init_ADI(model)
    nx, ny = model.grid.nx, model.grid.ny
    Δx, Δy = model.grid.Δx, model.grid.Δy
    ν, Δt  = model.parameters.ν, model.Δt

    βx = 0.5*ν*Δt/(Δx*Δx)
    βy = 0.5*ν*Δt/(Δy*Δy)
    d1x = ones(nx-2) .+ 2 * βx
    e1x = -βx * ones(nx-3)
    d1y = ones(ny-2) .+ 2 * βy
    e1y = -βy * ones(ny-3)
    # matrix ( I - 0.5 nu dt delta_x^2 )
    a1x = SymTridiagonal(d1x,e1x)
    a1y = SymTridiagonal(d1y,e1y)

    d2x = ones(nx-2) .- 2 * βx
    e2x = -βx * ones(nx-3)
    d2y = ones(ny-2) .- 2 * βy
    e2y = -βy * ones(ny-3)
    # matrix ( I + 0.5 nu dt delta_x^2)
    a2x = SymTridiagonal(d2x,-e2x)
    a2y = SymTridiagonal(d2y,-e2y)

    #LDLt factorize
    a1x_ldlt = factorize(a1x)
    a1y_ldlt = factorize(a1y)
    return a1x_ldlt, a1y_ldlt, a2x, a2y
end


"""
    poisson_eigenvalues(N, L, dim, ::Periodic)

Return the eigenvalues satisfying the discrete form of Poisson's equation
with periodic boundary conditions along the dimension `dim` with `N` grid
points and domain extent `L`.
"""
function poisson_eigenvalues(N, L, dim, ::Periodic)
    inds = reshape(1:N, reshaped_size(N, dim)...)
    return @. (2sin((inds - 1) * π / N) / (L / N))^2
end

"""
    poisson_eigenvalues(N, L, dim, ::Bounded)

Return the eigenvalues satisfying the discrete form of Poisson's equation
with staggered Neumann boundary conditions along the dimension `dim` with
`N` grid points and domain extent `L`.
"""
function poisson_eigenvalues(N, L, dim, ::Bounded)
    inds = reshape(1:N, reshaped_size(N, dim)...)
    return @. (2sin((inds - 1) * π / 2N) / (L / N))^2
end

"""
    poisson_eigenvalues(N, L, dim, ::Flat)

No need to do any transforms along `Flat` dimensions so just return zeros
since they will be added to the eigenvalues for other dimensions.
"""
poisson_eigenvalues(N, L, dim, ::Flat) = zeros(1, 1, 1)


## In this notebook, we'll look at stochastic gradient descent on a least squares
## system

##
using Plots
pyplot()
theme(:dark)
##
# https://vincentarelbundock.github.io/Rdatasets/csv/carData/Davis.csv
using CSV
df = CSV.read(download("https://vincentarelbundock.github.io/Rdatasets/csv/carData/Davis.csv"))
X = Float64.([df[:weight] df[:height]])
# Filter the data
goodpts = X[:,1] .<= 115
Xf = X[goodpts,:]
##
scatter(Xf[:,1],Xf[:,2],label="",xlabel="weight (kg)", ylabel="height (cm)")
##
x = Xf[:,1]
A = [x ones(length(x))]
b = Xf[:,2]
ab = A\b # solve the least squares problem

##
using Printf
using LinearAlgebra
""" Run stochastic gradient descent to solve a least squares problem. """
function stochastic_gradient_descent_leastsquares(A,b;
            step=0.0001, maxiter=10*size(A,1),tol=1e-6)
  x = zeros(size(A,2)) # start at 0
  for i=1:1000000
    si = rand(1:size(A,1))
    ai = A[si,:]
    bi = b[si]
    g = 2.0.*(ai'*x - bi)*ai
    x -= step*g
  end
  return x
end
stochastic_gradient_descent_leastsquares(A,b)
## Animate the algorithm
quadratic_leastsq(A,b) = x-> (0.5*norm(A*x - b)^2)/size(A,1)
ezsurf(x, y, f) = begin
    X = repeat(x', length(y), 1)
    Y = repeat(y, 1, length(x))
    # Evaluate each f(x, y)
    Z = map((x,y) -> f([x,y]), X, Y)
    plot(x, y, Z, st=:surface)
end
ezsurf(-3:0.1:3,-50:0.5:250,
  quadratic_leastsq(A,b))
plot!(colorbar=false,size=(450,450),dpi=300)
plot!(camera=(170,30))
x = zeros(size(A,2)) # initialize
p = plot3d!(1, label="", marker = 2)
step = 0.0001
anim = Animation()
for i=0:1000000
  global x
  g = (A'*A*x - A'*b)/size(A,1)
  x -= step*g
  if ((i <= 10) || ((i <= 30) & (rem(i,2) == 0)) || ((i <= 100) && rem(i,5) == 0) ||
      ((i <= 1000) && rem(i,10) == 0) || ((i <=10000) && rem(i,100)==0) ||
      ((i <= 100000) && rem(i,1000)==0) || rem(i,10000)==0)
    push!(p[1][2], x[1], x[2], quadratic_leastsq(A,b)(x))
    title!(@sprintf("Iteration %3i  Residual = %.3e", i, norm(g)))
    @printf("Iteration %3i  Residual = %.3e\n", i, norm(g))
    frame(anim)
  end
end
gif(anim, "stochatic-grad-descent-1.gif",fps=30)

##
## Animate the stochastic algorithm
using Random
Random.seed!(0)
quadratic_leastsq(A,b) = x-> (0.5*norm(A*x - b)^2)/size(A,1)
ezsurf(x, y, f) = begin
    X = repeat(x', length(y), 1)
    Y = repeat(y, 1, length(x))
    # Evaluate each f(x, y)
    Z = map((x,y) -> f([x,y]), X, Y)
    plot(x, y, Z, st=:surface)
end
ezsurf(-3:0.1:3,-50:0.5:250,
  quadratic_leastsq(A,b))
plot!(colorbar=false,size=(450,450),dpi=300)
plot!(camera=(170,30))
x = zeros(size(A,2)) # initialize
xgd = zeros(size(A,2)) # initialize
p = plot3d!(1, label="", color=3)
p2 = plot3d!(1, label="", color=2)
step = 0.0001
anim = Animation()
for i=0:1000000
  global x, xgd
  si = rand(1:size(A,1))
  ai = A[si,:]
  bi = b[si]
  g = (ai'*x - bi)*ai
  x -= step*g

  xgd -= step*(A'*A*xgd - A'*b)/size(A,1)
  if ((i <= 10) || ((i <= 30) & (rem(i,2) == 0)) || ((i <= 100) && rem(i,5) == 0) ||
      ((i <= 1000) && rem(i,10) == 0) || ((i <=10000) && rem(i,100)==0) ||
      ((i <= 100000) && rem(i,1000)==0) || rem(i,10000)==0)
    push!(p[1][2], x[1], x[2], quadratic_leastsq(A,b)(x))
    push!(p[1][3], xgd[1], xgd[2], quadratic_leastsq(A,b)(xgd))
    title!(@sprintf("Iteration %3i  Residual = %.3e", i, norm(g)))
    @printf("Iteration %3i  Residual = %.3e\n", i, norm(g))
    frame(anim)
  end
end
gif(anim, "stochatic-grad-descent-2.gif",fps=30)

## Show the quadratic surface
quadratic_leastsq(A,b) = x-> (0.5*norm(A*x - b)^2)
ezsurf(x, y, f) = begin
    X = repeat(x', length(y), 1)
    Y = repeat(y, 1, length(x))
    # Evaluate each f(x, y)
    Z = map((x,y) -> f([x,y]), X, Y)
    plot(x, y, Z, st=:surface)
end
ezsurf(-3:0.1:3,-50:0.5:250,
  quadratic_leastsq(A,b))
##
plot!(colorbar=false,size=(300,300),dpi=300)
anim = @animate for i=1:360
  plot!(camera=(i,30))
end
gif(anim, "leastsq-quadratic.gif")

##
## Show the quadratic surface with alpha-blending
quadratic_leastsq(A,b) = x-> (0.5*norm(A*x - b)^2)
ezsurf(x, y, f) = begin
    X = repeat(x', length(y), 1)
    Y = repeat(y, 1, length(x))
    # Evaluate each f(x, y)
    Z = map((x,y) -> f([x,y]), X, Y)
    plot()
    plot!(x, y, Z, st=:surface)
end
ezsurf(-3:0.1:3,-50:0.5:250,
  quadratic_leastsq(A,b))

##
## Show k random samples of the quadratic surface
plotlyjs()
quadratic_leastsq(A,b) = x-> (0.5*norm(A*x - b)^2)/size(A,1)
ezsurf!(x, y, f) = begin
    X = repeat(x', length(y), 1)
    Y = repeat(y, 1, length(x))
    # Evaluate each f(x, y)
    Z = map((x,y) -> f([x,y]), X, Y)
    plot!(x, y, Z, st=:surface, alpha=0.1)
end
orig = ezsurf(-3:0.25:3,-50:1:250,
  quadratic_leastsq(A,b))
nsamp = 25
k = 25
for i=1:k
  subset = randperm(size(A,1))[1:25] # choose 25 of 200 samples.
  ezsurf!(-3:0.1:3,-50:0.5:250,
    quadratic_leastsq(A[subset,:],b[subset,:]))
end
plot!() # get the plot to show
## Show k random samples of the quadratic surface
using Random
using LinearAlgebra
plotlyjs()
quadratic_leastsq(A,b) = x-> (0.5*norm(A*x - b)^2)/size(A,1)
ezsurf!(x, y, f) = begin
    X = repeat(x', length(y), 1)
    Y = repeat(y, 1, length(x))
    # Evaluate each f(x, y)
    Z = map((x,y) -> f([x,y]), X, Y)
    plot!(x, y, Z, st=:surface, alpha=0.4)
end
plot(framestyle=:none,colorbar=false,size=(300,300),dpi=300)
nsamp = 10
k = 25
for i=1:k
  subset = randperm(size(A,1))[1:nsamp] # choose 25 of 200 samples.
  ezsurf!(-3:0.25:3,-50:2:250,
    quadratic_leastsq(A[subset,:],b[subset,:]))
end
plot!() # get the plot to show
##
#plot!(colorbar=false,size=(300,300),dpi=300)
plot!(colorbar=false,size=(900,900))
anim = @animate for i=2:2:360
  plot!(camera=(i,30))
end
gif(anim, "leastsq-quadratic-samples.gif")

## Show the same thing with a log10
using Random
using LinearAlgebra
plotlyjs()
quadratic_leastsq(A,b) = x-> log10(0.5*norm(A*x - b)^2)/size(A,1)
plot(framestyle=:none,colorbar=false,size=(300,300),dpi=300)
nsamp = 10
k = 25
for i=1:k
  subset = randperm(size(A,1))[1:nsamp] # choose 25 of 200 samples.
  ezsurf!(-3:0.1:3,-50:0.5:250,
    quadratic_leastsq(A[subset,:],b[subset,:]))
end
plot!() # get the plot to show
##
#plot!(colorbar=false,size=(300,300),dpi=300)
plot!(colorbar=false,size=(900,900))
anim = @animate for i=2:2:360
  plot!(camera=(i,30))
end
gif(anim, "leastsq-quadratic-samples-log.gif")
##


"""
    Exponential(lnℓ::AbstractFloat) <: RadialKernel{Euclidean}

The exponential kernel with length scale
``\\exp(\\ln \\ell) = \\ell > 0``.

```math
k_{\\ell}(x, y) = k_{\\ell}(\\lVert x - y\\rVert_2) =
\\exp\\bigg\\lbrace -\\frac{\\lVert x - y\\rVert_2}{\\ell}\\bigg\\rbrace
```
"""
mutable struct Exponential{F<:AbstractFloat} <: RadialKernel{Euclidean}
    dist::Euclidean
    lnℓ::F
    ℓ::F
    function Exponential(dist::Euclidean, lnℓ::AbstractFloat)
        new{typeof(lnℓ)}(dist, lnℓ, exp(lnℓ))
    end
end

Exponential(lnℓ::AbstractFloat) = Exponential(Euclidean(), lnℓ)

@inline (k::Exponential)(τ::AbstractFloat) = exp(-τ / k.ℓ)

numparams(::Exponential) = (1,)
paramtypes(::Exponential{F}) where F = (F,)
params(k::Exponential) = (lnℓ = k.lnℓ,)
function setparams!(k::Exponential{F}, lnℓ::F) where F
    k.lnℓ = lnℓ
    k.ℓ = exp(lnℓ)
end


function hypot(x::DoubleFloat{T}, y::DoubleFloat{T}) where {T<:IEEEFloat}
    ax = abs(x)
    ay = abs(y)
    ay, ax = minmax(ax, ay)

    r = ay
    if !iszero(ax)
        r /= ax
    end
    r *= r

    rr = ax * sqrt(one(DoubleFloat{T}) + r)

    # from Base
    # use type of rr to make sure that return type
    #    is the same for all branches
    if isnan(ay)
        isinf(ax) && return oftype(rr, Inf)
        isinf(ay) && return oftype(rr, Inf)
        return oftype(rr, r)
    end
    return rr
 end

"""
    normalize(x,y)

    x_normalized^2 + y_normalized^2 == one(promote_type(typeof(x),typeof(y)))
"""
function normalize(x::DoubleFloat{T}, y::DoubleFloat{T}) where {T<:IEEEFloat}
    hypotenuse = hypot(x, y)
    if !iszero(hypotenuse) && isfinite(hypotenuse)
        xnorm = x / hypotenuse
        ynorm = y / hypotenuse
    else
        # try for xnorm, ynorm such that xnorm + ynorm == 1
        xnorm = x/(x+y)
        ynorm = y/(x+y)
    end
    return xnorm, ynorm
end


"""
    create_ols_approximation(y::Vector{<:Real}, x::Vector{<:Real}, base_x::Real = 0.0, degree::Integer = 1, intercept::Bool = true)
    create_ols_approximation(y::Vector{<:Real}, x::Union{Vector{DateTime},Vector{Date},Vector{Union{Date,DateTime}}}, base_x::Union{Date,DateTime} = global_base_date, degree::Integer = 1, intercept::Bool = true)

An approximation function calculated via OLS.

### Inputs
* `y` - A `Vector` with the y coordinates
* `x` - A `Vector` with the x coordinates
* `base_x` - A real that offsets the x. So a coordinate with x value of 2.0 will be converted to 1.8 if base_x is 0.2.
* `degree` - What the highest power of x should be. So if this is 3 then the equation will have x, x^2, x^3 as predictors.
* `intercept` - Should there be an x intercept.
### Returns
* A `Sum_Of_Functions` containing the approximation function.
"""
function create_ols_approximation(y::Vector{<:Real}, x::Vector{<:Real}, base_x::Real = 0.0, degree::Integer = 1, intercept::Bool = true)
    obs = length(y)
    if degree < 0
        error("Cannot approximate with OLS with a degree that is negative")
    end
    x = x .- base_x
    if intercept
        X = ones(obs)
        for i in 1:degree
            X = hcat(X, (x .^ i))
        end
    else
        X = x
        for i in 2:degree
            X = hcat(X, (x .^ i))
        end
    end

    lm1 = fit(LinearModel,  hcat(X), y)
    beta = lm1.pp.beta0
    func_array = Array{PE_Function,1}(undef,convert(Int, intercept) + degree)
    if intercept
        func_array[1] = PE_Function(beta[1], 0.0, base_x, 0)
    end
    for d in 1:degree
        func_array[d+Integer(intercept)] = PE_Function(beta[d+Integer(intercept)], 0.0, base_x, d)
    end
    return Sum_Of_Functions(func_array)
end

function create_ols_approximation(y::Vector{<:Real}, x::Union{Vector{DateTime},Vector{Date},Vector{Union{Date,DateTime}}}, base_x::Union{Date,DateTime} = global_base_date, degree::Integer = 1, intercept::Bool = true)
    base = years_from_global_base.(base_x)
    xx   = years_from_global_base.(x)
    return create_ols_approximation(y, xx, base, degree, intercept)
end

function get_cholesky_coefficients(chebyshev::Sum_Of_Functions, y::Vector{<:Real}, unnormalised_nodes::Vector{<:Real})
    chebyshev_on_nodes = evaluate.(Ref(chebyshev), unnormalised_nodes)
    a = sum(y .* chebyshev_on_nodes) / sum(chebyshev_on_nodes .^ 2)
    return a
end

"""
    create_chebyshev_approximation(func::Function, nodes::Integer, degree::Integer, left::Real, right::Real)

An function that will approximate another function via Chebyshev polynomials.

### Inputs
* `func` - A function that you want to approximation
* `nodes` - The number of approximation nodes
* `degree` - The degree of the Chebyshev polynomials.
* `left` - The left limit of the approximation
* `right` - The right limit of the approximation.
### Returns
* A `Sum_Of_Functions` containing the approximation function.
"""
function create_chebyshev_approximation(func::Function, nodes::Integer, degree::Integer, left::Real, right::Real)
    # This is all after Algorithm 6.2 from Judd (1998) Numerical Methods in Economics.
    if nodes <= degree
        error("Need to have more nodes than degree to use a chebyshev approximation")
    end
    k = 1:nodes
    unnormalised_nodes = -cos.( (((2 .* k) .- 1) ./ (2 * nodes)) .* pi    )
    normalised_nodes = ((unnormalised_nodes .+ 1) .* ((right-left)/2)) .+ left
    y = func.(normalised_nodes)
    chebyshevs = get_chevyshevs_up_to(degree, true)
    a = get_cholesky_coefficients.(chebyshevs, Ref(y), Ref(unnormalised_nodes))
    transformed_chebyshevs = convert_to_linearly_rescale_inputs.(chebyshevs, (right-left)/2, +(right+left)/2)
    # Note that these alpha and beta parameters in the convert_to_linearly_rescale_inputs function differ from those in Judd because those did not work in this context.
    all_terms = a .* transformed_chebyshevs
    final_func = Sum_Of_Functions(all_terms)
    return final_func
end


"""
    psMF(lnν)

Press-Schechter halo multiplicity function.

*Reference*: Press & Schechter, 187, 425 (1974)
- See also Bond, Cole, Efstathiou & Kaiser, ApJ, 379, 440 (1991)
- See Sheth & Tormen, MNRAS, 329, 61 (2002) for the notation used here. `psMF(lnν)` corresponds to νf(ν) in this reference.

# Arguments
- `lnν::Real`: natural logarithm of a threshold, ν, i.e., `lnν` = log(ν), defined by ν ≡ [δc/σ(R)]^2. Here, δc = 1.6865 and σ(R) is the r.m.s. mass fluctuation within a top-hat smoothing of scale R.

Press & Schechter's multiplicity function is normalized, i.e,

``∫_-∞^∞ dlnν psMF(lnν) = 1``

This function is based on [Cosmology Routine Library (CRL)](https://wwwmpa.mpa-garching.mpg.de/~komatsu/crl/).
"""
function psMF(lnν::Real)
    ν = exp(lnν)
    mf = exp(-ν / 2) * √(ν / 2π)
end

"""
    stMF(lnν)

Sheth-Tormen halo multiplicity function.

*Reference*: Equation (2) of Sheth & Tormen, MNRAS, 329, 61 (2002)
- See also Sheth, Mo & Tormen, MNRAS, 323, 1 (2001)
- `stMF(lnν)` corresponds to νf(ν) in these references

# Arguments
- `lnν::Real`: natural logarithm of a threshold, ν, i.e., `lnν` = log(ν), defined by ν ≡ [δc/σ(R)]^2. Here, δc = 1.6865 and σ(R) is the r.m.s. mass fluctuation within a top-hat smoothing of scale R.

Sheth & Tormen's multiplicity function is normalized, i.e,

``∫_-∞^∞ dlnν stMF(lnν) = 1``

This function is based on [Cosmology Routine Library (CRL)](https://wwwmpa.mpa-garching.mpg.de/~komatsu/crl/).
"""
function stMF(lnν::Real)
    A, p, q = 0.3222, 0.3, 0.707
    ν = exp(lnν)
    mf = A * (1 + (q * ν)^-p) * exp(-q * ν / 2) * √(q * ν / 2π)
end

"""
    jenkinsMF(lnν)

Jenkins et al.'s halo multiplicity function.

*Reference*: Equation (B3) of Jenkins et al., MNRAS, 321, 372 (2001)
- `jenkinsMF(lnν)` corresponds to f(σ)/2 in this reference, where σ=√(δc/ν).
- We use the fitting function for Δm = 180. Here, Δm is an overdensity within a spherical region of radius R, whose mean density is equal to Δm times the mean **mass** density of the Universe.

# Arguments
- `lnν::Real`: natural logarithm of a threshold, ν, i.e., `lnν` = log(ν), defined by ν ≡ [δc/σ(R)]^2. Here, δc = 1.6865 and σ(R) is the r.m.s. mass fluctuation within a top-hat smoothing of scale R.

Jenkins et al.'s multiplicity function is NOT normalized, i.e,

``∫_-∞^∞ dlnν jenkinsMF(lnν) ≠ 1``

This function is based on [Cosmology Routine Library (CRL)](https://wwwmpa.mpa-garching.mpg.de/~komatsu/crl/).
"""
function jenkinsMF(lnν::Real)
    ν = exp(lnν)
    σ = 1.6865 / √ν
    mf = 0.5 * 0.301 * exp(-abs(log(1 / σ) + 0.64)^3.82)
end


using Plots

function bracket_minimum(f, x=0, s=1e-2, k=2.0)
    a, ya = x, f(x);
    b, yb = a + s, f(a + s);

    if yb > ya
        a, b = b, a;
        ya, yb = yb, ya;
        s = -s;
    end

    while true
        c, yc = b + s, f(b + s);
        if yc > yb
            return a < c ? (a, c) : (c, a);
        end

        a, ya, b, yb = b, yb, c, yc;
        s *= k;
    end
end

# test

f1 = x -> (1 .+ x).^2;
f2 = x -> (1 + x)^2;
x = -3:0.1:3;
y = f1(x);

a, b = bracket_minimum(f2, -2, 1e-2, 1.01);

plot(x, y)
vline!([a, b])


"""
    angle(v::AbstractVector{<:Real}) -> Float

Angle the vector `v` makes with the x axis, throws error if vector is not of size 2
"""
function angle(v::AbstractVector{<:Real})
    if length(v)!=2 error("Vector must be of length 2") end
    atan(v[2],v[1])
end

angle(v::SVector{2, <:Real})  = atan(v[2],v[1])

"""
    sqrt(z, branchcutangle::Real=π) -> typeof(z)

Return \\sqrt{x}, analytically continued from real line with the branch cut from 0 in a straight line along `branchcutangle`
"""
function sqrt(z, branchcutangle::Real)

    if branchcutangle > 0
        if angle(z) > branchcutangle return -sqrt(z) end
    else
        if angle(z) < branchcutangle return -sqrt(z) end
    end

    return sqrt(z)
end

"""
    round(num::Real, base::Real) -> Real

Nearest multiple of `base` to `num`, result will be promoted to num and base
"""
round(num::Real, base::Real) = round(num/base)*base

"""
    ceil(num::Real, base::Real) -> Real

Smallest multiple of `base` greater than `num`, result will be promoted to num and base
"""
ceil(num::Real, base::Real) = ceil(num/base)*base

"""
    ceil(num::Real, base::Real) -> Real

Largest multiple of `base` smaller than `num`, result will be promoted to num and base
"""
floor(num::Real, base::Real) = floor(num/base)*base


# Generates the dataset to be used in GMM example for EM course

using Plots, Distributions, Random, DataFrames, CSV
Random.seed!(123)

# Set GMM parameters
n = 200
α = 0.4
μ₀ = [-1, 2]
μ₁ = [3.0, 7.0]
Σ₀ = [2.0 1.0; 1.0 2.0]
Σ₁ = [1.5 0; 0 0.5]
𝒩₀ = MvNormal(μ₀, Σ₀)
𝒩₁ = MvNormal(μ₁, Σ₁)
X₁ = zeros(n)
X₂ = zeros(n)

for i=1:n
    X₁[i],X₂[i] = rand() < α ? rand(𝒩₁) : rand(𝒩₀)
end

# Store in data frame
df = DataFrame(X1 = X₁, X2 = X₂)

# Save to csv
CSV.write("gmm_observations.csv", df)


# basisfunctions.jl

"""
    fourier(k::Int)

Implements the Fourier basis of functions ϕ_k, defined, if k is odd, by
ϕ_k(x)=sqrt(2)sin((k+1)π*x) and if k is even, by ϕ_k(x)=sqrt(2)cos(k*x).
This function is defined for k≥0 and ϕ_0≡1.

# Examples

```julia
x = -1.0:0.01:2.0
y = fourier(3).(x)
```
"""
function fourier(k::Int)
  k ≥ 0 || throw(AssertionError("k should be non-negative"))
  sqrttwo = sqrt(2.0)
  if k == 0
    return x -> 1.0
  elseif isodd(k)
    return x -> sqrttwo * sinpi(float(k+1)*x)
  else
    return x -> sqrttwo * cospi(float(k)*x)
  end
end

"""
    faberschauderone(x::Float64)

Implements the first Faber-Schauder function defined by 1-2x for 0≤x≤1/2 and
-1+2x for 0.5≤x≤1, and is 1-periodically extended to all x∈R.

#See also: faberschauder

#Warning

Note the difference between faberschauderone and faberschauder. The first is a
function that takes a Float64 and returns a Float64, the second takes (j,k) and
returns a anonymous function that takes a Float64 and returns a Float64.

#Examples

```julia
x=-2.0:0.001:2.0
y=faberschauderone.(x)
```
"""
function faberschauderone(x::Float64)
    y = mod(x, 1.0)
    if 0≤y≤0.5
        return 1.0 - 2*y
    else
        return -1.0 + 2*y
    end
end

"""
    faberschauder(j::Int, k::Int)

Implements the k-th Faber-Schauder function of level j. Here, j≥0 and 1≤k≤2^j.
It is a one periodic function and defined on [0,1] by 2^(j+1)(x-(k-1)2^(-j)) on
(k-1)2^(-j)≤x≤(k-1/2)2^(-j) and 1 - 2^(j+1)(x-(k-1/2)2^(-j)) on
[(k-1/2)2^(-j), k2^(-j)] and zero outside these intervals.

# See also: faberschauderone.

#Warning

Note the difference between faberschauderone and faberschauder. The first is a
function that takes a Float64 and returns a Float64, the second takes (j,k) and
returns a anonymous function that takes a Float64 and returns a Float64.

#Example using PyPlot
```julia
using PyPlot
clf()
j=3
x=-2.0:0.001:2.0
for k in 1:2^j
    y = faberschauder(j,k).(x)
    plot(x,y)
end
```
"""
function faberschauder(j::Int,k::Int)
  j ≥ 0 || throw(AssertionError("j should be a nonnegative integer."))
  1≤k≤2^j || throw(AssertionError("k should be an integer between 1 and 2^j."))
  twotothepowerminusj = 2.0^-j
  k_float = float(k)
  j_float = float(j)
  return function(x::Float64)
    y = mod(x,1.0)
    if y ≤ (k_float-1.0)*twotothepowerminusj || y ≥ k_float*twotothepowerminusj
      return 0.0
    elseif (k_float-1.0)*twotothepowerminusj < y ≤ (k_float-0.5)*twotothepowerminusj
      return 2.0^(j_float+1.0)*y-2*k_float+2.0
    else
      return 2.0*k_float - 2.0^(j_float+1.0) * y
    end
  end
end