"""
    ito_correction!(I, h=1)

Applies the Itô-correction for iterated integrals to `I`.
This amounts to subtracting ``\\frac{1}{2}h`` from every element on the diagonal.

# Example
```jldoctest; setup=:(using LevyArea)
julia> M = ones(5,5);

julia> LevyArea.ito_correction!(M, 0.5)

julia> M
5×5 Matrix{Float64}:
 0.75  1.0   1.0   1.0   1.0
 1.0   0.75  1.0   1.0   1.0
 1.0   1.0   0.75  1.0   1.0
 1.0   1.0   1.0   0.75  1.0
 1.0   1.0   1.0   1.0   0.75
```

"""
function ito_correction!(I, h=1)
    m,n = size(I)
    m == n || throw(DimensionMismatch("Matrix is not square: dimensions are $(size(I))"))
    @inbounds for i=1:m
        I[i,i] -= h/2
    end
end

"""
    iterated_integrals(W::AbstractVector, h, eps=h^(3/2);
        ito_correction=true,
        error_norm=MaxL2(),
        alg=optimal_algorithm(length(W),h,eps,error_norm)
    )

Simulates an approximation of the iterated stochastic integrals
``\\int_0^h\\int_0^sdW_i(t)dW_j(s)`` for all pairs ``1\\le i, j \\le m``
of the given ``m``-dimensional Brownian motion with step size h.

# Examples
```jldoctest; setup=:(using LinearAlgebra; using LevyArea)
julia> h = 1/2;

julia> W = [1.0, 0.5]
2-element Vector{Float64}:
 1.0
 0.5

julia> diag(iterated_integrals(W, h, h^(3/2))) ≈ 0.5*W.^2 .- 0.5h
true
```
"""
function iterated_integrals(W::AbstractVector, h::Real, eps::Real=h^(3/2);
    ito_correction=true,
    error_norm::AbstractErrorNorm=MaxL2(),
    alg::AbstractIteratedIntegralAlgorithm=optimal_algorithm(length(W),h,eps,error_norm)
)
    m = length(W)
    n = terms_needed(m, h, eps, alg, error_norm)
    I = levyarea(W/√h, n, alg)
    if ito_correction
        ito_correction!(I)
    end
    I .= 1//2 .* W.*W' .+ h.*I
    return I
end

"""
    iterated_integrals(W::AbstractVector, q_12::AbstractVector, h, eps;
        ito_correction=true,
        error_norm=FrobeniusL2(),
        alg=optimal_algorithm(length(W),q_12,h,eps,error_norm)
    )

Simulates an approximation of the iterated stochastic integrals for finite-dimensional approximations of
a Q-Wiener process with covariance matrix ``Q = Q^\\frac{1}{2}*Q^\\frac{1}{2}``.
Here `q_12` is a vector of the eigenvalues of ``Q^\\frac{1}{2}``; the square root of the covariance matrix.
Equivalently these are the square roots of the eigenvalues of ``Q``.

# Examples
```jldoctest; setup=:(using LinearAlgebra; using LevyArea)
julia> h = 0.01; dim=10; q = [1/k^2 for k=1:dim];

julia> W = √h * sqrt.(q) .* randn(dim);

julia> diag(iterated_integrals(W,sqrt.(q),h,h^(3/2))) ≈ 0.5*W.^2 .- 0.5*h*q
true
```
"""
function iterated_integrals(W::AbstractVector, q_12::AbstractVector, h::Real, eps::Real;
    ito_correction=true,
    error_norm::AbstractErrorNorm=FrobeniusL2(),
    alg::AbstractIteratedIntegralAlgorithm=optimal_algorithm(length(W),q_12,h,eps,error_norm)
)
    m = length(W)
    n = terms_needed(m, q_12, h, eps, alg, error_norm)
    I = levyarea(W./q_12./√h, n, alg)
    if ito_correction
        ito_correction!(I)
    end
    I .= 1//2 .* W.*W' .+ h.*q_12'.*I.*q_12 # scale correctly
    return I
end

"""
    iterated_integrals(W::Real, h::Real, eps::Real=0.0; ito_correction=true, kwargs...)

In the case of a scalar Brownian motion the integral can be explicitly
calculated as ``\\int_0^h\\int_0^sdW(t)dW(s) = \\frac{1}{2}W(h)^2 - \\frac{1}{2}h``.

The parameter `eps` (as well as all additional keyword arguments) has no effect but is available
to provide the same interface as the multidimensional version.
"""
iterated_integrals(W::Real, h::Real, eps::Real=0.0; ito_correction=true, kwargs...) = ito_correction ? (W^2-h)/2 : W^2/2

"""
    iterated_integrals(W::Real, q_12::Real, h::Real, eps::Real; ito_correction=true, kwargs...)

In the case of a scalar Q-Wiener process with (scalar) covariance Q the integral can be explicitly
calculated as ``\\int_0^h\\int_0^sdW(t)dW(s) = \\frac{1}{2}W(h)^2 - \\frac{1}{2}hQ``.

Note that, as in the multidimensional case, the parameter `q_12` denotes the square root of the covariance.

The parameter `eps` (as well as all additional keyword arguments) has no effect but is available
to provide the same interface as the multidimensional version.
"""
iterated_integrals(W::Real, q_12::Real, h::Real, eps::Real; ito_correction=true, kwargs...) = ito_correction ? (W^2-h*q_12^2)/2 : W^2/2


# # Stratified Kelvin-Helmholtz instability
#
# ## Install dependencies
#
# First let's make sure we have all required packages installed.

# ```julia
# using Pkg
# pkg"add Oceananigans, JLD2, CairoMakie"
# ```

# ## The physical domain
#
# We simulate Kelvin-Helmholtz instability in two-dimensions in ``x, z``
# and therefore assign `Flat` to the `y` direction,

using Oceananigans

grid = RectilinearGrid(size=(64, 64), x=(-5, 5), z=(-5, 5),
                       topology=(Periodic, Flat, Bounded))

# # The basic state
#
# We're simulating the instability of a sheared and stably-stratified basic state
# ``U(z)`` and ``B(z)``. Two parameters define our basic state: the Richardson number
# defined as,
#
# ```math
# Ri = \frac{∂_z B}{(∂_z U)^2} ,
# ```
#
# and the width of the stratification layer, ``h``.

shear_flow(x, y, z, t) = tanh(z)

stratification(x, y, z, t, p) = p.h * p.Ri * tanh(z / p.h)

U = BackgroundField(shear_flow)

B = BackgroundField(stratification, parameters=(Ri=0.1, h=1/4))

# Our basic state thus has a thin layer of stratification in the center of
# the channel, embedded within a thicker shear layer surrounded by unstratified fluid.

using CairoMakie

zF = znodes(Face, grid)
zC = znodes(Center, grid)

Ri, h = B.parameters

fig = Figure(resolution = (850, 450))

ax = Axis(fig[1, 1], xlabel = "U(z)", ylabel = "z")
lines!(ax, shear_flow.(0, 0, zC, 0), zC; linewidth = 3)

ax = Axis(fig[1, 2], xlabel = "B(z)")
lines!(ax, [stratification(0, 0, z, 0, (Ri=Ri, h=h)) for z in zC], zC; linewidth = 3, color = :red)

ax = Axis(fig[1, 3], xlabel = "Ri(z)")
lines!(ax, [Ri * sech(z / h)^2 / sech(z)^2 for z in zF], zF; linewidth = 3, color = :black) # Ri(z)= ∂_z B / (∂_z U)²; derivatives computed by hand

# In unstable flows it is often useful to determine the dominant spatial structure of the
# instability and the growth rate at which the instability grows.
# If the simulation idealizes a physical flow, this can be used to make
# predictions as to what should develop and how quickly.
# Since these instabilities are often attributed to a linear instability,
# we can determine information about the structure and the growth rate of the instability by analyzing the linear operator
# that governs small perturbations about a base state, or by solving for the linear dynamics.

# Here, we discuss first briefly linear instabilities and how one can obtain growth rates and structures of most unstable
# modes via eigenanalysis. Then we present an alternative method for approximating the eigenanalysis results when
# one does not have access to the linear dynamics or the linear operator about the base state.

# # Linear Instabilities
#
# The base state ``U(z)``, ``B(z)`` is a solution of the inviscid equations of motion. Whether the base state is
# stable or not is determined by whether small perturbations about this base state grow or decay. To formalize this,
# we study the linearized dynamics satisfied by perturbations about the base state:
# ```math
# \partial_t \Phi = L \Phi \, .
# ```
# where ``\Phi = (u, v, w, b)`` is a vector of the perturbation velocities ``u, v, w`` and perturbation buoyancy ``b``
# and ``L`` a linear operator that depends on the base state, ``L = L(U(z), B(z))`` (the `background_fields`).
# Eigenanalysis of the linear operator ``L`` determines the stability of the base state, such as the Kelvin-Helmholtz
# instability. That is, by using the ansantz
# ```math
# \Phi(x, y, z, t) = \phi(x, y, z) \, \exp(\lambda t) \, ,
# ```
# then ``\lambda`` and ``\phi`` are respectively eigenvalues and eigenmodes of ``L``, i.e., they obey
# ```math
# L \, \phi_j = \lambda_j \, \phi_j \quad j=1,2,\dots \, .
# ```
# From hereafter we'll use the convention that the eigenvalues are ordered according to their real part,
# ``\mathrm{Re}(\lambda_1) \ge \mathrm{Re}(\lambda_2) \ge \dotsb``.
#
# Remarks:
#
# As we touched upon briefly above, Oceananigans.jl, does not include a linearized version of the equations.
# Furthermore, Oceananigans.jl does not give us access to the linear operator ``L`` so that we can perform eigenanalysis.
# Below we discuss an alternative way of approximating the eigenanalysis results.
# The method boils down to solving the nonlinear equations while continually renormalize
# the magnitude of the perturbations to ensure that nonlinear terms
# (terms that are quadratic or higher in perturbations) remain negligibly small,
# i.e.,much smaller than the background flow.

# # The power method algorithm
#
# Successive application of ``L`` to a random initial state will eventually render it parallel
# with eigenmode ``\phi_1``:
# ```math
# \lim_{n \to \infty} L^n \Phi \propto \phi_1 \, .
# ```
# Of course, if ``\phi_1`` is an unstable mode (i.e., ``\sigma_1 = \mathrm{Re}(\lambda_1) > 0``), then successive application
# of ``L`` will lead to exponential amplification. (Similarly, if ``\sigma_1 < 0``, successive application of ``L`` will
# lead to exponential decay of ``\Phi`` down to machine precision.) Therefore, after each
# application of the linear operator ``L``, we rescale the output ``L \Phi`` back to a pre-selected amplitude.
#
# So, we initialize a `simulation` with random initial conditions with amplitude much less than those of
# the base state (which are ``O(1)``). Instead of "applying" ``L`` on our initial state, we evolve the
# (approximately) linear dynamics for interval ``\Delta \tau``. We measure how much the energy has grown
# during that interval, rescale the perturbations back to original energy amplitude and repeat.
# After some iterations the state will converge to the most unstable eigenmode.
#
# In summary, each iteration of the power method includes:
# - compute the perturbation energy, ``E_0``,
# - evolve the system for a time-interval ``\Delta \tau``,
# - compute the perturbation energy, ``E_1``,
# - determine the exponential growth of the most unstable mode during the interval
#   ``\Delta \tau`` as  ``\log(E_1 / E_0) / (2 \Delta \tau)``,
# - repeat the above until growth rate converges.
#
# By fiddling a bit with ``\Delta t`` we can get convergence after only a few iterations.
#
# Let's apply all these to our example.

# # The model

model = NonhydrostaticModel(timestepper = :RungeKutta3,
                              advection = UpwindBiasedFifthOrder(),
                                   grid = grid,
                               coriolis = nothing,
                      background_fields = (u=U, b=B),
                                closure = ScalarDiffusivity(ν=2e-4, κ=2e-4),
                               buoyancy = BuoyancyTracer(),
                                tracers = :b)

# We have included a "pinch" of viscosity and diffusivity in anticipation of what will follow furtherdown:
# viscosity and diffusivity will ensure numerical stability when we evolve the unstable mode to the point it becomes nonlinear.

# For this example, we take ``\Delta \tau = 15``.

simulation = Simulation(model, Δt=0.1, stop_iteration=150)

# Now some helper functions that will be used during for the power method algorithm.
#
# First a function that evolves the state for ``\Delta \tau`` and measure the energy growth
# over that period.

"""
    grow_instability!(simulation, energy)

Grow an instability by running `simulation`.

Estimates the growth rate ``σ`` of the instability
using the fractional change in volume-mean kinetic energy,
over the course of the `simulation`

``
energy(t₀ + Δτ) / energy(t₀) ≈ exp(2 σ Δτ)
``

where ``t₀`` is the starting time of the simulation and ``t₀ + Δτ``
the ending time of the simulation. We thus find that the growth rate
is measured by

``
σ = log(energy(t₀ + Δτ) / energy(t₀)) / (2 Δτ) .
``
"""
function grow_instability!(simulation, energy)
    ## Initialize
    simulation.model.clock.iteration = 0
    t₀ = simulation.model.clock.time = 0
    compute!(energy)
    energy₀ = energy[1, 1, 1]

    ## Grow
    run!(simulation)

    ## Analyze
    compute!(energy)
    energy₁ = energy[1, 1, 1]
    Δτ = simulation.model.clock.time - t₀

    ## ½(u² + v²) ~ exp(2 σ Δτ)
    σ = growth_rate = log(energy₁ / energy₀) / 2Δτ

    return growth_rate
end
nothing # hide

# Finally, we write a function that rescales the state. The rescaling is done via computing the
# kinetic energy and then rescaling all flow fields so that the kinetic energy assumes a targetted value.
#
# (Measuring the perturbation growth via the kinetic energy works fine _unless_ an unstable mode _only_ has
# buoyancy structure. In that case, the total perturbation energy is more adequate.)

"""
    rescale!(model, energy; target_kinetic_energy = 1e-3)

Rescales all model fields so that `energy = target_kinetic_energy`.
"""
function rescale!(model, energy; target_kinetic_energy = 1e-6)
    compute!(energy)
    rescale_factor = √(target_kinetic_energy / energy[1, 1, 1])

    for f in merge(model.velocities, model.tracers)
        f .*= rescale_factor
    end

    return nothing
end

using Printf

# Another helper function for the power method,

"""
    convergence(σ)

Check if the growth rate has converged. If the array `σ` has at least 2 elements then returns the
relative difference between ``σ[end]`` and ``σ[end-1]``.
"""
convergence(σ) = length(σ) > 1 ? abs((σ[end] - σ[end-1]) / σ[end]) : 9.1e18 # pretty big (not Inf tho)
nothing # hide

# and the main function that performs the power method iteration.

"""
    estimate_growth_rate(simulation, energy, ω; convergence_criterion=1e-3)

Estimates the growth rate iteratively until the relative change
in the estimated growth rate ``σ`` falls below `convergence_criterion`.

Returns ``σ``.
"""
function estimate_growth_rate(simulation, energy, ω, b; convergence_criterion=1e-3)
    σ = []
    power_method_data = []
    compute!(ω)
    push!(power_method_data, (ω=collect(interior(ω)[:, 1, :]), b=collect(interior(b)[:, 1, :]), σ=deepcopy(σ)))

    while convergence(σ) > convergence_criterion
        compute!(energy)

        @info @sprintf("About to start power method iteration %d; kinetic energy: %.2e", length(σ)+1, energy[1, 1, 1])
        push!(σ, grow_instability!(simulation, energy))
        compute!(energy)

        @info @sprintf("Power method iteration %d, kinetic energy: %.2e, σⁿ: %.2e, relative Δσ: %.2e",
                       length(σ), energy[1, 1, 1], σ[end], convergence(σ))

        compute!(ω)
        rescale!(simulation.model, energy)
        push!(power_method_data, (ω=collect(interior(ω)[:, 1, :]), b=collect(interior(b)[:, 1, :]), σ=deepcopy(σ)))
    end

    return σ, power_method_data
end
nothing # hide

# # Eigenplotting
#
# A good algorithm wouldn't be complete without a good visualization,

u, v, w = model.velocities
b = model.tracers.b

perturbation_vorticity = Field(∂z(u) - ∂x(w))

xω, yω, zω = nodes(perturbation_vorticity)
xb, yb, zb = nodes(b)

# # Rev your engines...
#
# We initialize the power iteration with random noise and rescale to have a `target_kinetic_energy`

using Random, Statistics

mean_perturbation_kinetic_energy = Field(Average(1/2 * (u^2 + w^2)))
noise(x, y, z) = randn()
set!(model, u=noise, w=noise, b=noise)
rescale!(simulation.model, mean_perturbation_kinetic_energy, target_kinetic_energy=1e-6)
growth_rates, power_method_data = estimate_growth_rate(simulation, mean_perturbation_kinetic_energy, perturbation_vorticity, b)

@info "Power iterations converged! Estimated growth rate: $(growth_rates[end])"

# # Powerful convergence
#
# We animate the power method steps. A scatter plot illustrates how the growth rate converges
# as the power method iterates.

n = Observable(1)

fig = Figure(resolution=(800, 600))

kwargs = (xlabel="x", ylabel="z", limits = ((xω[1], xω[end]), (zω[1], zω[end])), aspect=1,)

ω_title(t) = t === nothing ? @sprintf("vorticity") : @sprintf("vorticity at t = %.2f", t)
b_title(t) = t === nothing ? @sprintf("buoyancy")  : @sprintf("buoyancy at t = %.2f", t)

ax_ω = Axis(fig[2, 1]; title = ω_title(nothing), kwargs...)

ax_b = Axis(fig[2, 3]; title = b_title(nothing), kwargs...)

ωₙ = @lift power_method_data[$n].ω
bₙ = @lift power_method_data[$n].b

σₙ = @lift [(i-1, i==1 ? NaN : growth_rates[i-1]) for i in 1:$n]

ω_lims = @lift (-maximum(abs, power_method_data[$n].ω) - 1e-16, maximum(abs, power_method_data[$n].ω) + 1e-16)
b_lims = @lift (-maximum(abs, power_method_data[$n].b) - 1e-16, maximum(abs, power_method_data[$n].b) + 1e-16)

hm_ω = heatmap!(ax_ω, xω, zω, ωₙ; colorrange = ω_lims, colormap = :balance)
Colorbar(fig[2, 2], hm_ω)

hm_b = heatmap!(ax_b, xb, zb, bₙ; colorrange = b_lims, colormap = :balance)
Colorbar(fig[2, 4], hm_b)

eigentitle(σ, t) = length(σ) > 0 ? @sprintf("Iteration #%i; growth rate %.2e", length(σ), σ[end]) : @sprintf("Initial perturbation fields")
σ_title = @lift eigentitle(power_method_data[$n].σ, nothing)

ax_σ = Axis(fig[1, :];
            xlabel = "Power iteration",
            ylabel = "Growth rate",
            title = σ_title,
            xticks = 1:length(power_method_data)-1,
            limits = ((0.5, length(power_method_data)-0.5), (-0.25, 0.25)))

scatter!(ax_σ, σₙ; color = :blue)

frames = 1:length(power_method_data)

record(fig, "powermethod.mp4", frames, framerate=1) do i
       @info "Plotting frame $i of $(frames[end])..."
       n[] = i
end

nothing #hide

# ![](powermethod.mp4)

# # Now for the fun part
#
# Now we simulate the nonlinear evolution of the eigenmode
# we've isolated for a few e-folding times ``1/\sigma``,

## Reset the clock
model.clock.iteration = 0
model.clock.time = 0

estimated_growth_rate = growth_rates[end]

simulation.stop_time = 5 / estimated_growth_rate
simulation.stop_iteration = 9.1e18 # pretty big (not Inf tho)

## Rescale the eigenmode
initial_eigenmode_energy = 5e-5
rescale!(simulation.model, mean_perturbation_kinetic_energy, target_kinetic_energy=initial_eigenmode_energy)

# Let's save and plot the perturbation vorticity and buoyancy and also the total vorticity and
# buoyancy (perturbation + basic state). It'll be also neat to plot the kinetic energy time-series
# and confirm it grows with the estimated growth rate.

total_vorticity = Field(∂z(u) + ∂z(model.background_fields.velocities.u) - ∂x(w))

total_b = Field(b + model.background_fields.tracers.b)

simulation.output_writers[:vorticity] =
    JLD2OutputWriter(model, (ω=perturbation_vorticity, Ω=total_vorticity, b=b, B=total_b, KE=mean_perturbation_kinetic_energy),
                     schedule = TimeInterval(0.10 / estimated_growth_rate),
                     filename = "kelvin_helmholtz_instability.jld2",
                     overwrite_existing = true)

# And now we...

@info "*** Running a simulation of Kelvin-Helmholtz instability..."
run!(simulation)

# ## Pretty things
#
# First we plot the nonlinear equilibration of the perturbation fields together
# with the evolution of the kinetic energy,

@info "Making a neat movie of stratified shear flow..."

filepath = simulation.output_writers[:vorticity].filepath

ω_timeseries = FieldTimeSeries(filepath, "ω")
b_timeseries = FieldTimeSeries(filepath, "b")
Ω_timeseries = FieldTimeSeries(filepath, "Ω")
B_timeseries = FieldTimeSeries(filepath, "B")
KE_timeseries = FieldTimeSeries(filepath, "KE")

times = ω_timeseries.times

t_final = times[end]

n = Observable(1)

ωₙ = @lift interior(ω_timeseries, :, 1, :, $n)
bₙ = @lift interior(b_timeseries, :, 1, :, $n)

fig = Figure(resolution=(800, 600))

kwargs = (xlabel="x", ylabel="z", limits = ((xω[1], xω[end]), (zω[1], zω[end])), aspect=1,)

title = @lift @sprintf("t = %.2f", times[$n])

ax_ω = Axis(fig[2, 1]; title = "perturbation vorticity", kwargs...)

ax_b = Axis(fig[2, 3]; title = "perturbation buoyancy", kwargs...)

ax_KE = Axis(fig[3, :];
             yscale = log10,
             limits = ((0, t_final), (initial_eigenmode_energy, 1e-1)),
             xlabel = "time")

fig[1, :] = Label(fig, title, textsize=24, tellwidth=false)

ω_lims = @lift (-maximum(abs, interior(ω_timeseries, :, 1, :, $n)) - 1e-16, maximum(abs, interior(ω_timeseries, :, 1, :, $n)) + 1e-16)
b_lims = @lift (-maximum(abs, interior(b_timeseries, :, 1, :, $n)) - 1e-16, maximum(abs, interior(b_timeseries, :, 1, :, $n)) + 1e-16)

hm_ω = heatmap!(ax_ω, xω, zω, ωₙ; colorrange = ω_lims, colormap = :balance)
Colorbar(fig[2, 2], hm_ω)

hm_b = heatmap!(ax_b, xb, zb, bₙ; colorrange = b_lims, colormap = :balance)
Colorbar(fig[2, 4], hm_b)

tₙ = @lift times[1:$n]
KEₙ = @lift KE_timeseries[1:$n]

lines!(ax_KE, [0, t_final], @. initial_eigenmode_energy * exp(2 * estimated_growth_rate * [0, t_final]);
       label = "~ exp(2 σ t)",
       linewidth = 2,
       color = :black)

lines!(ax_KE, times, KE_timeseries[:];
       label = "perturbation kinetic energy",
       linewidth = 4, color = :blue, alpha = 0.4)

KE_point = @lift Point2f[(times[$n], KE_timeseries[$n][1, 1, 1])]

scatter!(ax_KE, KE_point;
         marker = :circle, markersize = 16, color = :blue)

frames = 1:length(times)

record(fig, "kelvin_helmholtz_instability_perturbations.mp4", frames, framerate=8) do i
    @info "Plotting frame $i of $(frames[end])..."
    n[] = i
end
nothing #hide

# ![](kelvin_helmholtz_instability_perturbations.mp4)

# And then the same for total vorticity & buoyancy of the fluid.
n = Observable(1)

Ωₙ = @lift interior(Ω_timeseries, :, 1, :, $n)
Bₙ = @lift interior(B_timeseries, :, 1, :, $n)

fig = Figure(resolution=(800, 600))

kwargs = (xlabel="x", ylabel="z", limits = ((xω[1], xω[end]), (zω[1], zω[end])), aspect=1,)

title = @lift @sprintf("t = %.2f", times[$n])

ax_Ω = Axis(fig[2, 1]; title = "total vorticity", kwargs...)

ax_B = Axis(fig[2, 3]; title = "total buoyancy", kwargs...)

ax_KE = Axis(fig[3, :];
             yscale = log10,
             limits = ((0, t_final), (initial_eigenmode_energy, 1e-1)),
             xlabel = "time")

fig[1, :] = Label(fig, title, textsize=24, tellwidth=false)

hm_Ω = heatmap!(ax_Ω, xω, zω, Ωₙ; colorrange = (-1, 1), colormap = :balance)
Colorbar(fig[2, 2], hm_Ω)

hm_B = heatmap!(ax_B, xb, zb, Bₙ; colorrange = (-0.05, 0.05), colormap = :balance)
Colorbar(fig[2, 4], hm_B)

tₙ = @lift times[1:$n]
KEₙ = @lift KE_timeseries[1, 1, 1, 1:$n]

lines!(ax_KE, [0, t_final], @. initial_eigenmode_energy * exp(2 * estimated_growth_rate * [0, t_final]);
       label = "~ exp(2 σ t)",
       linewidth = 2,
       color = :black)

lines!(ax_KE, times, KE_timeseries[:];
       label = "perturbation kinetic energy",
       linewidth = 4, color = :blue, alpha = 0.4)

KE_point = @lift Point2f[(times[$n], KE_timeseries[$n][1, 1, 1])]

scatter!(ax_KE, KE_point;
         marker = :circle, markersize = 16, color = :blue)

axislegend(ax_KE; position = :rb)

record(fig, "kelvin_helmholtz_instability_total.mp4", frames, framerate=8) do i
       @info "Plotting frame $i of $(frames[end])..."
       n[] = i
end
nothing #hide

# ![](kelvin_helmholtz_instability_total.mp4)


using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("colorscheme.jl"))
using DynamicalSystems, InteractiveDynamics
import GLMakie
using AbstractPlotting
using LinearAlgebra
using Statistics

# Code for 3D animation
ds = Systems.lorenz()
u0 = trajectory(ds, 10; Ttr = 100)[end]
dt = 0.005
N = 100
framerate = 30
systemtitle = "Lorenz63"
azimuth = 2.24
elevation = 0.38

ds = Systems.towel()
u0 = trajectory(ds, 10; Ttr = 100)[end]
dt = 1
N = 15 # total steps to take
framerate = 2
systemtitle = "towelmap"

elevation = 0.5
azimuth = 6.02

# %%
tinteg = tangent_integrator(ds, 3; u0)
Y = get_deviations(tinteg)

# Now I define functions that make an ellipsoid given the point vectors
function ellipsoid_transform(Y)
    svddecomp = svd(Y)
    d = svddecomp.S
    s = d ./ maximum(d)
    Q = svddecomp.U * svddecomp.Vt # rotation matrix
    θx = atan(Q[3,2],Q[3,3])
    θy = atan(-Q[3,1], sqrt(Q[3,2]^2+Q[3,3]^2))
    θz = atan(Q[2,1], Q[1,1])
    M = AbstractPlotting.rotationmatrix_x(θx) *
        AbstractPlotting.rotationmatrix_y(θy) *
        AbstractPlotting.rotationmatrix_z(θz) *
        AbstractPlotting.scalematrix(Vec3f0(s...))
    return M, s
end
function transform_mesh(msh, mat4x4)
    pos_trans = Point3f0.(Ref(mat4x4) .* to_ndim.(Point4f0, spheremesh.position, 0))
    AbstractPlotting.GeometryBasics.Mesh(pos_trans, AbstractPlotting.GeometryBasics.faces(msh))
end
function ellipsoid_from_pointcloud(v, msh = spheremesh)
    M, s = ellipsoid_transform(v)
    return transform_mesh(msh, M), s
end
spheremesh = AbstractPlotting.GeometryBasics.mesh(Sphere(Point3f0(0), 1))

ellipsoid, s = ellipsoid_from_pointcloud(Y, spheremesh)
colors = [RGBAf0((abs.(u)./norm(u))..., 1) for u in spheremesh.position]

ellipsobs = Observable(ellipsoid)

# Setup figure and plot everything

# fig = Figure(resolution = (1200, 600))
fig = Figure(resolution = (1200, 600)); display(fig)
# Plot trajectory
t = Observable(0.0)
tstr = lift(x -> systemtitle *", t = $(round(x; digits = 5))", t)
axtr = Axis3(fig[1,1]; title = tstr, titlealign = :left)
tr = trajectory(ds, 100N*dt, u0; dt)
dotobs = Observable([u0])
trobs = Observable([u0])
if ds isa ContinuousDynamicalSystem
    lines!(axtr, tr.data; linewidth = 0.2, color = RGBAf0(0,0,0,0.5), transparency = true )
    lines!(axtr, trobs; linewidth = 4.0, color = to_color(JULIADYNAMICS_COLORS[1]), overdraw = true)
else
    scatter!(axtr, tr.data; markersize = 1000, color = RGBAf0(0,0,0,0.5), strokewidth=0, transparency = true )
    scatter!(axtr, trobs; markersize = 2000, color = to_color(JULIADYNAMICS_COLORS[1]), overdraw = true)
end
scatter!(axtr, dotobs; markersize = 4000, marker=:diamond,
color = to_color(JULIADYNAMICS_COLORS[3]), overdraw = true)
axtr.azimuth = azimuth
axtr.elevation = elevation

# Plot ellipsoid
ax3D = Axis3(fig[1, 2]; aspect = :data)
mesh!(ellipsobs, color = colors)

ax3D.xticklabelsvisible = false
ax3D.yticklabelsvisible = false
ax3D.zticklabelsvisible = false
ax3D.xlabelvisible = false
ax3D.ylabelvisible = false
ax3D.zlabelvisible = false
ax3D.azimuth = azimuth
ax3D.elevation = elevation

ax2 = Axis(fig[1, 3]; width = 300) # this axis has values of principal lengths
ax2.title = "ellipsoid axes"
tightlimits!(ax2, Bottom())
sobs = Observable(s)
barplot!(ax2, 1:3, sobs; color = to_color.(COLORSCHEME[1:3]))

display(fig)

# %%
record(
        fig, joinpath(@__DIR__, "volume_growth_$(systemtitle).mp4"), 1:N;
        framerate
        ) do i
    DynamicalSystems.step!(tinteg, dt)
    Y = get_deviations(tinteg)
    ellipsobs[], sobs[] = ellipsoid_from_pointcloud(Y)
    # s[] = principals(newus)
    # autolimits!(ax2)
    # must set new limits centered on u
    # maxε = maximum(norm(u) for u in newus)
    # maxε = maximum(s)
    # ax3D.limits = map(j -> (-maxε, +maxε), (1,2,3))
    # Update trajectory and title and elevation
    u = get_state(tinteg)
    dotobs[] = [u]
    trobs[] = push!(trobs[], u)
    t[] = tinteg.t
end


export Partition, Partitions, YoungTableau, SkewDiagram
export dim, has_left_neighbor, has_bottom_neighbor, inskewdiag, isrimhook,
       hooklength, leglength, matrix_repr, partitionseq

##############################################################################
#
#   Partition type, AbstractVector interface
#
##############################################################################

length(p::Partition) = length(p.part)

size(p::Partition) = size(p.part)

Base.IndexStyle(::Type{T}) where {T <: Partition} = Base.IndexLinear()

getindex(p::Partition, i::Int) = p.part[i]

function setindex!(p::Partition, v::Int, i::Int)
   prev = Inf
   nex = 1
   if i == length(p)
      prev = p[i-1]
   elseif i == 1
      nex = p[2]
   elseif 1 < i < length(p)
      prev = p[i-1]
      nex = p[i+1]
   end
   nex <= v <= prev || throw("Partition must be positive and non-increasing")
   p.n += v - p.part[i]
   p.part[i] = v
   return p
end

==(p::Partition, m::Partition) = p.n ==m.n && p.part == m.part
hash(p::Partition, h::UInt) = hash(p.part, hash(Partition, h))

convert(::Type{Partition}, p::Vector{Int}) = Partition(p)

##############################################################################
#
#   IO for Partition
#
##############################################################################

function subscriptify(n::Int)
   subscript_0 = Int(0x2080) # Char(0x2080) -> subscript 0
   return join([Char(subscript_0 + i) for i in reverse(digits(n))])
end

function show(io::IO, p::Partition)
   uniq = unique(p.part)
   mults = [count(i -> i == u, p.part) for u in uniq]
   str = join((string(u)*subscriptify(m) for (u,m) in zip(uniq, mults)))
   print(io, str)
end

##############################################################################
#
#   Iterator interface for Integer Partitions
#
##############################################################################

# Implemented following RuleAsc (Algorithm 3.1) from
#    "Generating All Partitions: A Comparison Of Two Encodings"
# by Jerome Kelleher and Barry O’Sullivan, ArXiv:0909.2331

function Base.start(parts::Partitions)
    if parts.n < 1
        return (Int[], 0)
    elseif parts.n == 1
        return ([1], 0)
    else
        p = zeros(Int, parts.n)
        p[2] = parts.n
        return (p, 2)
    end
end

Base.next(parts::Partitions, state) = nextpart_asc(state...)
Base.done(parts::Partitions, state) = state[2] == 1
Base.eltype(::Type{Partitions}) = Partition
length(parts::Partitions) = BigInt(numpart(parts.n))

function nextpart_asc(part, k)
    if k == 0
        return Partition(part, false), (part, 1)
    end
    y = part[k] - 1
    k -= 1
    x = part[k] + 1
    while x <= y
        part[k] = x
        y -= x
        k += 1
    end
    part[k] = x + y
    return Partition(reverse(part[1:k]), false), (part, k)
end

doc"""
    conj(part::Partition)
> Returns the conjugated partition of `part`, i.e. the partition corresponding
> to the Young tableau of `part` reflected through the main diagonal.
"""
function conj(part::Partition)
    p = Int[]
    for i in 1:part.n
        n = sum(part .>= i)
        n == 0 && break
        push!(p, n)
    end
    return Partition(p, false)
end

##############################################################################
#
#   Partition sequences and Murnaghan-Nakayama formula
#
##############################################################################

doc"""
    partitionseq(lambda::Partition)
> Returns a sequence (as `BitVector`) of `false`s and `true`s constructed from
> `lambda`: tracing the lower contour of the Young Diagram associated to
> `lambda` from left to right a `true` is inserted for every horizontal and
> `false` for every vertical step. The sequence always starts with `true` and
> ends with `false`.
"""
function partitionseq(lambda::Partition)
   seq = trues(maximum(lambda) + length(lambda))
   j = lambda[end]
   for i in (length(lambda)-1):-1:1
      seq[j+1] = false
      j += lambda[i] - lambda[i+1] + 1
   end
   seq[j+1] = false
   return seq
end

partitionseq(v::Vector{Int}) = partitionseq(Partition(v))

doc"""
    partitionseq(seq::BitVector)
> Returns the essential part of the sequence `seq`, i.e. a subsequence starting
> at first `true` and ending at last `false`.
"""
partitionseq(seq::BitVector) = seq[findfirst(seq, true):findlast(seq, false)]

doc"""
    isrimhook(R::BitVector, idx::Int, len::Int)
> `R[idx:idx+len]` forms a rim hook in the Young Diagram of parition
> corresponding to `R` iff `R[idx] == true` and `R[idx+len] == false`.
"""
function isrimhook(R::BitVector, idx::Int, len::Int)
   return (R[idx+len] == false) && (R[idx] == true)
end

const _charvalsTable = Dict{Tuple{BitVector,Vector{Int}}, BigInt}()

doc"""
    MN1inner(R::BitVector, mu::Partition, t::Int, [charvals])
> Returns the value of $\lambda$-th irreducible character on conjugacy class of
> permutations represented by partition `mu`, where `R` is the (binary)
> partition sequence representing $\lambda$. Values already computed are stored
> in `charvals::Dict{Tuple{BitVector,Vector{Int}}, Int}`.
> This is an implementation (with slight modifications) of the
> Murnaghan-Nakayama formula as described in
>
>     Dan Bernstein,
>     "The computational complexity of rules for the character table of Sn"
>     _Journal of Symbolic Computation_, 37(6), 2004, p. 727-748.
"""
function MN1inner(R::BitVector, mu::Partition, t::Int,
        charvals=Dict{Tuple{BitVector,Vector{Int}}, BigInt}())
    if t > length(mu)
        chi = 1
    elseif mu[t] > length(R)
        chi = 0
    else
        chi = 0
        sgn = false

        for j in 1:mu[t]-1
            if R[j] == false
                sgn = !sgn
            end
        end
        for i in 1:length(R)-mu[t]
            if R[i] != R[i+mu[t]-1]
                sgn = !sgn
            end
            if isrimhook(R, i, mu[t])
                R[i], R[i+mu[t]] = R[i+mu[t]], R[i]
                essR = (partitionseq(R), mu[t+1:end])
                if !haskey(charvals, essR)
                    charvals[essR] = MN1inner(R, mu, t+1, charvals)
                end
                chi += (-1)^Int(sgn)*charvals[essR]
                R[i], R[i+mu[t]] = R[i+mu[t]], R[i]
            end
        end
    end
    return chi
end

##############################################################################
#
#   YoungTableau type, AbstractVector interface
#
##############################################################################

function YoungTableau(part::Partition, fill::Vector{Int}=collect(1:part.n))
   part.n == length(fill) || throw("Can't fill Young digaram of $part with $fill: different number of elemnets.")
   tab = zeros(Int, length(part), maximum(part))
   k=1
   for (idx, p) in enumerate(part)
      tab[idx, 1:p] = fill[k:k+p-1]
      k += p
   end
   return YoungTableau(part, tab)
end

YoungTableau(p::Vector{Int}) = YoungTableau(Partition(p))

size(Y::YoungTableau) = size(Y.tab)

Base.IndexStyle(::Type{T}) where {T <: YoungTableau} = Base.IndexLinear()

getindex(Y::YoungTableau, i::Int) = Y.tab[i]

function ==(Y1::YoungTableau,Y2::YoungTableau)
   Y1.part == Y2.part || return false
   Y1.tab == Y2.tab || return false
   return true
end

hash(Y::YoungTableau, h::UInt) = hash(Y.part, hash(Y.tab, hash(YoungTableau, h)))

doc"""
    conj(Y::YoungTableau)
> Returns the conjugated tableau, i.e. the tableau reflected through the main
> diagonal.
"""
conj(Y::YoungTableau) = YoungTableau(conj(Y.part), transpose(Y.tab))

##############################################################################
#
#   Misc functions for YoungTableaux
#
##############################################################################

rowlen(Y::YoungTableau, i, j) = sum(Y[i, j:end] .> 0)
collen(Y::YoungTableau, i, j) = sum(Y[i:end, j] .> 0)

doc"""
    hooklength(Y::YoungTableau, i, j)
> Returns the hooklength of an element in `Y` at position `(i,j)`. `hooklength`
> will return `0` for `(i,j)` not in the tableau `Y`.
"""
function hooklength(Y::YoungTableau, i, j)
   if Y[i,j] == 0
      return 0
   else
      return rowlen(Y, i, j) + collen(Y, i, j) - 1
   end
end

doc"""
    dim(Y::YoungTableau)
> Returns the dimension of the irreducible representation of
> `PermutationGroup(n)` associated to YoungTableau `Y` of a partition of `n`.
"""
function dim(Y::YoungTableau)
   n, m = size(Y)
   num = fac(maximum(Y))
   den = reduce(*, 1, BigInt(hooklength(Y,i,j)) for i in 1:n, j in 1:m if j <= Y.part[i])
   return divexact(num, den)
end

##############################################################################
#
#   SkewDiagrams
#
##############################################################################

SkewDiagram(lambda::Vector{Int}, mu::Vector{Int}) = SkewDiagram(Partition(lambda), Partition(mu))

/(lambda::Partition, mu::Partition) = SkewDiagram(lambda, mu)

==(xi::SkewDiagram, psi::SkewDiagram) = xi.lam == psi.lam && xi.mu == psi.mu
hash(xi::SkewDiagram, h::UInt) = hash(xi.lam, hash(xi.mu, hash(SkewDiagram, h)))

###############################################################################
#
#   String I/O
#
###############################################################################

doc"""
    matrix_repr(xi::SkewDiagram)
> Returns a binary representation of the diagram `xi`, i.e. a binary array `A`
> where `A[i,j] == 1` if and only if `(i,j)` is in `xi.lam` but not in `xi.mu`.
"""
function matrix_repr(xi::SkewDiagram)
   ydiag = zeros(Int,length(xi.lam), maximum(xi.lam))
   for i in 1:length(xi.mu)
      ydiag[i, xi.mu[i]+1:xi.lam[i]] .= 1
   end
   for i in length(xi.mu)+1:length(xi.lam)
      ydiag[i,1:xi.lam[i]] .= 1
   end
   return ydiag
end

show(io::IO, xi::SkewDiagram) = show(io, MIME("text/plain"), matrix_repr(xi))

##############################################################################
#
#   Misc functions for SkewDiagrams
#
##############################################################################

doc"""
    inskewdiag(xi::SkewDiagram, i::Int, j::Int)
> Checks if box at position `(i,j)` belongs to the skew diagram `xi`.
"""
function inskewdiag(xi::SkewDiagram, i::Int, j::Int)
   if i <= 0 || j <= 0
      return false
   elseif i > length(xi.lam) || j > maximum(xi.lam)
      return false
   elseif length(xi.mu) >= i
      return xi.mu[i] < j ≤ xi.lam[i]
   else
      return j ≤ xi.lam[i]
   end
end

doc"""
    has_left_neighbor(xi::SkewDiagram, i::Int, j::Int)
> Checks if box at position `(i,j)` has neighbour in `xi` to the left.
"""
function has_left_neighbor(xi::SkewDiagram, i::Int, j::Int)
   if j == 1
      return false
   else
      return inskewdiag(xi, i, j) && inskewdiag(xi, i, j-1)
   end
end

doc"""
    has_bottom_neighbor(xi::SkewDiagram, i::Int, j::Int)
> Checks if box at position `(i,j)` has neighbour in `xi` below.
"""
function has_bottom_neighbor(xi::SkewDiagram, i::Int, j::Int)
   if i == length(xi.lam)
      return false
   else
      return inskewdiag(xi, i, j) && inskewdiag(xi, i+1, j)
   end
end

doc"""
    isrimhook(xi::SkewDiagram)
> Checks if `xi` represents a rim-hook diagram, i.e. its diagram is
> edge-connected and contains no $2\times 2$ squares.
"""
function isrimhook(xi::SkewDiagram)
   i = 1
   j = xi.lam[1]
   while i != length(xi.lam) && j != 1
      left = has_left_neighbor(xi, i,j)
      bottom = has_bottom_neighbor(xi, i,j)
      if left && bottom # there is 2×2 square in xi
         return false
      elseif left
         j -= 1
      elseif bottom
         i += 1
      else
         lam_tail = xi.lam[i+1:end]
         mu_tail = zeros(Int, length(lam_tail))
         mu_tail[1:length(xi.mu)-i] = xi.mu[i+1:end]

         if any(lam_tail .- mu_tail .> 0)
            return false # xi is disconnected
         else
            return true # we arrived at the end of xi
         end
      end
   end
   return true
end

doc"""
    leglength(xi::SkewDiagram, check::Bool=true)
> Computes the leglength of a rim-hook `xi`, i.e. the number of rows with
> non-zero entries minus one. If `check` is `false` function will not check
> whether `xi` is actually a rim-hook.
"""
function leglength(xi::SkewDiagram, check::Bool=true)
   if check
      isrimhook(xi) || throw("$xi is not a rimhook. leglength is defined only for rim hooks")
   end
   m = zeros(length(xi.lam))
   m[1:length(xi.mu)] = xi.mu
   return sum((xi.lam .- m) .> 0) - 1
end


#=
  From BLOG example/bayesian-linear-regression.blog
  """
  A model for Bayesian linear regression

  @author yiwu
  @since 2015-07-02
  """

  Note: This is 0-based i.e.
     x y
     0 0
     1 1
     2 4
     3 6
     4 ?

  Summary Statistics
  parameters      mean       std   naive_se      mcse        ess      rhat   ess_per_sec
      Symbol   Float64   Float64    Float64   Float64    Float64   Float64       Float64

           m    1.9616    0.4171     0.0042    0.0218   198.2693    1.0229       25.6327
           b   -0.1645    0.7370     0.0074    0.0321   351.3579    1.0083       45.4244
       sigma    0.9349    0.4702     0.0047    0.0309    86.5685    1.0103       11.1918
          y4    7.6895    1.6275     0.0163    0.0742   229.2941    1.0250       29.6437
     post[1]   -0.1513    1.2702     0.0127    0.0441   765.0960    1.0014       98.9135
     post[2]    1.8138    1.1771     0.0118    0.0365   910.5267    1.0006      117.7152
     post[3]    3.7861    1.1241     0.0112    0.0386   667.4734    1.0052       86.2926
     post[4]    5.7280    1.3455     0.0135    0.0475   517.3522    1.0120       66.8846
     post[5]    7.7025    1.6070     0.0161    0.0705   311.8366    1.0193       40.3150
     post[6]    9.6650    1.9211     0.0192    0.0894   276.0979    1.0201       35.6946
     post[7]   11.5449    2.1770     0.0218    0.1019   276.8155    1.0169       35.7874
     post[8]   13.5644    2.6050     0.0260    0.1283   236.7469    1.0168       30.6072
     post[9]   15.5463    2.9901     0.0299    0.1471   210.4787    1.0206       27.2112
    post[10]   17.4897    3.3944     0.0339    0.1688   230.5202    1.0202       29.8022
    post[11]   19.4406    3.7283     0.0373    0.1901   214.3265    1.0238       27.7087

   Cf ~/webppl/bayesian_linear_regression.wppl

=#

using Turing
include("jl_utils.jl")

@model function bayesian_linear_regression(y=[0,1,4,6])
    m ~ Normal(0,2)
    b ~ Normal(0,2)
    sigma ~ Gamma(1,1)
    for i in 0:length(y)-1
        y[i+1] ~ Normal(m * i + b, sigma)
    end

    # Posterior
    y4 ~ Normal(m * (4) + b, sigma)
    post = Vector{Real}(undef, n)
    for i in 0:10
        post[i+1] ~ Normal(m * i + b, sigma)
    end
end

y = [0,1,4,6]
model = bayesian_linear_regression(y)

num_chains = 4
# chs = sample(model, Prior(), 1000)
# chs = sample(model, MH(), 10_000)
# chs = sample(model, PG(15), 10_000)
# chs = sample(model, IS(), 10_000)
#chs = sample(model, SMC(), 10_000)
# chs = sample(model, SMC(), MCMCThreads(), 10_000, num_chs)

# chs = sample(model, SGLD(), 10_000)
chs = sample(model,NUTS(), 10_000)
# chs = sample(model,HMC(0.01,5), 10_000)

display(chs)
# display(plot(chs))


using NumericalAnalysis
using Test

@testset "NumericalAnalysis.jl" begin
    #2x + 4y
    # Write your tests here.
end

@testset "basic test" begin
    @test 1<=ξ(1, 2)<=2
    @test absolute_error(1.0, 2) == 1.0
    @test relative_error(1, 2) == 1.0
    @test NthDerivative(exp, 4, 4)  == exp(4)
    @test NthDerivative(x->x^3, 4, 4) == 0
    @test 0.9<=TaylorPolynomials(cos, 0.1, 0, 6)<=1.0

end


# generic solve_hjb( ) Implicit method.
# TODO: ns states/nc controls
# TODO: add uncertainty, nx shocks...

# fcn NGM() returns struct NGM;
# HJB_options() returns struct HJB_options;
# solve_HJB: (m, opt, v) -> dist, v_err, v, dV_Upwind, c_Upwind, μ_Upwind
# Plot error/value/policy/
# Simulate w/ interp.

# optimal stopping w/ LCPsolve.jl

using LinearAlgebra, SparseArrays, Plots

_m(x) = max(x,1e-10)
_m(x) = max(x,1e-8)

# NGM ###############################
struct NGM
    ρ; σ;            # preference param
    δ; A; α;         # tech param
    #
    r                # return fcn. r(s,c)
    μ                # transition fcn.  ṡ=μ(s,c)
    dr               # derivative of return fcn.
    FOC              # FOC: solve dr(c)=V_s. c=u'^-1(V_s)
    μ_inv            # inv transition c=μ_inv(s,ṡ), solve ṡ=μ(s,c) for c
    #
    s_min; s_max; H; s; ds;
    dVf; dVb; dV_Upwind; c_Upwind;
end

function NGM(;
    ρ = 0.05, σ= 2.0,
    δ = 0.05, A = 1.0, α= 0.3,
    r = c -> (c^(1-σ))/(1-σ),               # return fcn.      make it r(s,c)
    μ = (s,c) -> A*(s^α) -δ*s -c,           # Transition fcn.  ṡ=μ(s,c)
    dr = c -> c^(-σ),                       # derivative of return fcn.
    FOC = v -> (max(v,1e-10))^(-1.0/σ),     # FOC: solve dr(c)=V_s. c=u'^-1(V_s)
    μ_inv = (s,ṡ) -> A*(s^α) - δ*s -ṡ,      # Inv Trans c=μ_inv(s,ṡ), solve ṡ=μ(s,c) for c
    s_min = 0.001*(α*A/(ρ+δ))^(1/(1-α)),
    s_max = 2.000*(α*A/(ρ+δ))^(1/(1-α)),
    H = 10_000,
    s = collect(LinRange(s_min, s_max, H)),
    ds = (s_max-s_min)/(H-1),
    dVf = zeros(H,1),
    dVb = zeros(H,1),
    dV_Upwind = zeros(H,1),
    c_Upwind = zeros(H,1),
    )
    return NGM(ρ,σ,δ,A,α,r,μ,dr,FOC,μ_inv,s_min,s_max,H,s,ds,dVf,dVb,dV_Upwind,c_Upwind)
end
mod = NGM()
mod = NGM(ρ = 0.05)
mod = NGM(A = 0.01)

# con-sav ###############################
struct CS
    ρ; σ;            # preference param
    A; w;            # tech param
    #
    r                # return fcn. r(s,c)
    μ                # transition fcn.  ṡ=μ(s,c)
    dr               # derivative of return fcn.
    FOC              # FOC: solve dr(c)=V_s. c=u'^-1(V_s)
    μ_inv            # inv transition c=μ_inv(s,ṡ), solve ṡ=μ(s,c) for c
    #
    s_min; s_max; H; s; ds;
    dVf; dVb; dV_Upwind; c_Upwind;
end

function CS(;
    ρ = 0.05, σ= 2.0,
    A = 0.045, w = 0.1,                     # A = interest rate
    r = c -> (_m(c)^(1-σ))/(1-σ),           # return fcn.      make it r(s,c)
    μ = (s,c) -> w + A*s - c,               # Transition fcn.  ṡ=μ(s,c)
    dr = c -> _m(c)^(-σ),                       # derivative of return fcn.
    FOC = v -> (_m(v))^(-1.0/σ),     # FOC: solve dr(c)=V_s. c=u'^-1(V_s)
    μ_inv = (s,ṡ) -> w + A*s -ṡ,      # Inv Trans c=μ_inv(s,ṡ), solve ṡ=μ(s,c) for c
    s_min = -0.02,
    s_max = 1.0,
    H = 10_000,
    s = collect(LinRange(s_min, s_max, H)),
    ds = (s_max-s_min)/(H-1),
    dVf = zeros(H,1),
    dVb = zeros(H,1),
    dV_Upwind = zeros(H,1),
    c_Upwind = zeros(H,1),
    )
    return CS(ρ,σ,A,w,r,μ,dr,FOC,μ_inv,s_min,s_max,H,s,ds,dVf,dVb,dV_Upwind,c_Upwind)
end

# NBC
# μ_inv(a_bar,0) > 0 ⟺ w + A*a_bar >0 ⟺ a_bar > -w/A
#a_bar = -w/A # Natural Borrowing Constraint
#a_bar = -w/A + eps()

mod = CS()
mod = CS(w = 0.1, s_min = -0.1/0.045 + eps())

# verify things are well defined @ corners
mod.μ_inv(mod.s_min,0) > 0
mod.s_max, mod.μ_inv(mod.s_max,0), mod.dr(mod.μ_inv(mod.s_max,0))
mod.s_min, mod.μ_inv(mod.s_min,0), mod.dr(mod.μ_inv(mod.s_min,0))

##################################
# options
struct HJB_options; Δ; maxit; ε; dist; end
HJB_options(;Δ=1_000,maxit=10_000,ε=1e-6,dist=[]) = HJB_options(Δ,maxit,ε,dist)
HJB_options()

##################################
# solve_HJB(m, opt, v)
function solve_HJB(m, opt, v)
    # unpack
    # σ,δ,A,α = m.σ, m.δ, m.A, m.α     # Do not need.
    ρ = m.ρ
    r,μ,dr,FOC,μ_inv = m.r, m.μ, m.dr, m.FOC, m.μ_inv
    s_min,s_max,H,s,ds = m.s_min, m.s_max, m.H, m.s, m.ds
    dVf,dVb,dV_Upwind,c_Upwind = m.dVf, m.dVb, m.dV_Upwind, m.c_Upwind
    #
    Δ,maxit,ε,dist = opt.Δ, opt.maxit, opt.ε, opt.dist

    # VFI
    for n=1:maxit
        #println(n)
        V=v
        dV = (V[2:end] - V[1:end-1])/ds

        # dVf = [dV; dr(μ_inv(s_max,0))]
        # dVb = [dr(μ_inv(s_min,0)); dV]

        # forward difference: if ṡ>0
        dVf[1:end-1] = dV
        dVf[end]     = dr(μ_inv(s_max,0))    # u'(c(s_max)) state const, for stab #dVf[H]=0
        cf           = FOC.(dVf)             # choice with forward difference
        μ_f          = μ.(s, cf)             # ṡ      with forward difference
        If = μ_f .> 0

        # backward difference: if ṡ<0
        dVb[2:end] = dV
        dVb[1]     = dr(μ_inv(s_min,0))      # state const, for stab #dVb[1] = (A*s_min)^(-σ)
        cb  = FOC.(dVb)                      # choice with backward difference
        μ_b = μ.(s, cb)
        Ib = μ_b .< 0

        # neither difference: if ṡ=0
        c0  = μ_inv.(s,0)        # c if ṡ=0
        dV0 = dr.(c0)
        # μ_0 = μ.(s, c0)          # μ_0 == zero(s)
        I0  = (1.0 .- If - Ib)   # choice betw forward & backward difference

        # I_concave = dVb .> dVf
        # scatter(I_concave) #1 everywhere EXCEPT @ last point H.

        dV_Upwind = dVf.*If + dVb.*Ib + dV0.*I0
        c_Upwind  = FOC.(dV_Upwind)
        μ_Upwind  = μ.(s, c_Upwind)
        u = r.(c_Upwind)

        # plot(dVf, lab="dVf")
        # plot!(dVb, lab="dVb")
        # plot!(dV0, lab="dV0")
        # plot!(dV_Upwind, lab="dV_Upwind")

        # plot(cf, lab="dVf")
        # plot!(cb, lab="dVb")
        # plot!(c0, lab="dV0")
        # plot!(c_Upwind, lab="dV_Upwind")

        # plot(μ_f, lab="dVf")
        # plot!(μ_b, lab="dVb")
        # plot!(zero(μ_f), lab="dV0")
        # plot!(μ_Upwind, lab="dV_Upwind")

        # create the transition matrix
        X = -min.(μ_b,0)/ds
        Y = -max.(μ_f,0)/ds + min.(μ_b,0)/ds
        Z = max.(μ_f,0)/ds

        a1 = sparse(Diagonal((Y[:])))
        a2 = [zeros(1,H); sparse(Diagonal((X[2:H]))) zeros(H-1,1)]
        a3 = [zeros(H-1,1) sparse(Diagonal((Z[1:H-1]))); zeros(1,H)]
        AA = a1 + a2 + a3

        B = (ρ + 1/Δ)*sparse(I,H,H) - AA
        b = u + V./Δ

        # Solve: B V = b
        # => (ρ + 1/Δ)*sparse(I,H,H)*V - AA*V = u + V./Δ
        # => (ρ + 1/Δ)*V = u + V./Δ +  AA*V
        # => (ρ)*V + V./Δ = u + V./Δ +  AA*V
        # => (ρ)*V = u  +  AA*V

        V = B \ b
        V_change = V-v
        v = V

        push!(dist,findmax(abs.(V_change))[1])
        println(n, " ", dist[n])
        if dist[n] .< ε
            println("Value Function Converged Iteration=")
            println(n)
            break
        end
    end
    μ_Upwind  = μ.(s, c_Upwind)
    v_err = r.(c_Upwind) + dV_Upwind.*μ_Upwind - ρ.*v # approx @ borrowing constraint

    return dist, v_err, v, dV_Upwind, c_Upwind, μ_Upwind
end

##################################
# DO IT.
mod = NGM() # mod.s_min =0.0048
mod = NGM(s_max=3.0)

# NGM Closed form. Need: ρ = δ*(α*σ - 1)
mod = NGM(σ= 4.0, δ = 0.05, A = 1.0, α= 0.3, ρ = 0.05*(0.3*4 - 1))

mod = CS()
# verify TVC: ρ-(1-σ)*A >0
# mod.ρ - (1-mod.σ)*mod.A   # AZ: always true if σ>1

mod.w, mod.s_min, -mod.w/mod.A
# AZ: lower s_min, the longer it takes to hit constraint. Bigger T_sim
mod = CS(s_min=-.5)
mod = CS(s_min=-1.5)
mod = CS(s_min=-2)
mod = CS(s_min=-0.1/0.045 + eps())
mod = CS(s_min=-0.1/0.045 -.25)

mod = CS(w=0.0, s_min=0.015) # binding
mod = CS(w=0.0, s_min=0.0)   # NBC
mod = CS(w=0.0, s_min=-0.01)  # loose
#mod = CS(w=0.0, s_min=-0.10)  # loose. Inaccurate when s close to 0

# AZ: got errors before bc w=0 & s_min =-.02

#
opt = HJB_options(ε=1e-7)
v0 = @. mod.r(mod.μ_inv(mod.s,0))/mod.ρ #v0 = @. r(μ_inv(s,0))/ρ #initial guess for V

dist, v_err, v, dV_Upwind, c_Upwind, μ_Upwind = solve_HJB(mod, opt, v0)

##################################
# Plot: dist, v_err, v, c_Upwind, μ_Upwind
plot(dist,
		xlabel="Iteration", ylabel="||V^{n+1} - V^n||",
		legend=false, title="")
plot(mod.s, v_err,
		xlabel="s", ylabel="Error in the HJB equation",
		legend=false, title="")
plot(mod.s, v,
		xlabel="s", ylabel="V(s)",
		legend=false, title="")
plot(mod.s, c_Upwind,
		xlabel="s", ylabel="c(s)",
		legend=false, title="")
plot(mod.s, μ_Upwind,
		xlabel="s", ylabel="ṡ(s)",
		title="", label="ṡ(s)", legend=false);
plot!([0.0],  seriestype = :hline, lab="", color="grey", l=:dash)

##################################
# Simulate: s, c, ṡ
# problem: c(s) only on finite grid s => interpolate
using Interpolations
ĉ = LinearInterpolation(mod.s, c_Upwind[:], extrapolation_bc = Line())

# NGM: converges quickly
Δt = 0.01;
T_sim = 200;

# CS: converges slowly
Δt = 0.01;
T_sim = 3_500;

t_sim = 0.0:Δt:T_sim
s_sim, c_sim, ṡ_sim = [zeros(length(t_sim),1) for i in 1:3]

s_0=mod.s[1] # start below ss
s_0=0.5
#s_0=0.75

s_sim[1] = s_0
for i in 2:length(t_sim)
    c_sim[i-1] = ĉ(s_sim[i-1])
    ṡ_sim[i-1] = mod.μ(s_sim[i-1], c_sim[i-1])
    s_sim[i]   = s_sim[i-1] + Δt * ṡ_sim[i-1]
    # s_sim[i] = s_sim[i-1] + Δt * μ(s_sim[i-1], c_sim[i-1])
end
#
ix = 1:(length(t_sim)-1)
plot();
plot!(t_sim[ix], s_sim[ix], lab = "s");
plot!(t_sim[ix], c_sim[ix], lab = "c");
plot!(t_sim[ix], ṡ_sim[ix], lab = "ṡ");
plot!([mod.s_min],  seriestype = :hline, lab="", color="grey", l=:dash)

# Verify: HACT Fig 2 Cara/F3 CRRA, hit borrowing constraint in finite time

########################################################
# Special case of NGM has closed-form sol. Compare Policy & Sim
# Policy: c = (1-s)(k^α)
cp(s) = (1-((α*δ)/(ρ+δ)))*s^α
s_d(s) = (s^α) - δ*s - cp(s)
plot(legend=:topleft)
plot!(s, cp,      lab="c: closed form")
plot!(s, c_Upwind,lab="c: Upwind")
plot!(s, s_d,     lab="ṡ: closed form")
plot!(s, s_dot,   lab="ṡ: Upwind")

#
s_ss = (α*A/(ρ+δ))^(1/(1-α))
z_ss = (s_ss)^(1-α)
c_ss = (s_ss)^(α)  - δ*s_ss

#s_0 = 5;
z_0 = (s_0)^(1-α)

λ=(1-α)*δ
z(t) = z_ss + exp(-λ*t)*(z_0 - z_ss)
kk(t) = z(t)^(1/(1-α))
cc(t) = (1-((α*δ)/(ρ+δ)))*kk(t)^α

#t_sim = 0.01:0.01:100
t_sim = time[ix]

plot(legend=:topleft)
#plot!(t_sim, z, lab="z")
plot!(t_sim, kk, lab="s: closed form")
plot!(t_sim, s_sim[ix], lab="s: Upwind")
plot!(t_sim, cc, lab="c: closed form")
plot!(t_sim, c_sim[ix], lab="c: Upwind")
#plot!([z_ss s_ss c_ss],  seriestype = :hline, lab="", color="grey")
plot!([s_ss c_ss],  seriestype = :hline, lab="", color="grey")


"""
    Lasso(; alpha = 0.1, tol = 1e-4, mi = 1e+8)
Lasso Regression structure. eEach parameters are as follows:
- `alpha` : leaarning rate.
- `tol` : Allowable error.
- `mi` : Maximum number of learning.

# Example
```jldoctest regression
julia> model = Lasso()
Lasso(Float64[], 0.1, 0.0001, 100000000)

julia> fit!(model, x, t)
3-element Vector{Float64}:
   0.0
   0.5022766549841176
 154.43624186616267

julia> model(x)
20-element Vector{Float64}:
 188.64541774549468
 187.30088075704523
 181.5191726873298
 184.15748544986084
 185.72158909964372
 176.33798923891868
 185.80014722707335
 184.71565381947883
 189.20524796663838
 189.67502263476888
 189.50409373058318
 188.39535519538825
 188.481083670683
 188.88872347085172
 182.8477136378307
 188.64156231429416
 181.43996475587224
 188.9400571253936
 179.39836073711297
 185.6065850765288
```
"""
mutable struct Lasso
    w::Array{Float64, 1}
    α::Float64
    tol::Float64
    mi::Int64
    Lasso(; alpha = 0.1, tol = 1e-4, mi = 1e+8) = new(Array{Float64}(undef, 0), alpha, tol, mi)
end

sfvf(x, y) = sign(x) * max(abs(x) - y, 0)

function fit!(model::Lasso, x, t)
    function update!(x, t, w, α)
        n, d = size(x)
        w[1] = mean(t - x' * w[2:end])
        wvec = fill!(Array{Float64}(undef, d), w[1])
        for k in 1 : n
            ww = w[2:end]
            ww[k] = 0
            q = (t - wvec - x' * ww) ⋅ x[k, :]
            r = x[k, :] ⋅ x[k, :]
            w[k+1] = sfvf(q / r, α)
        end
    end
    α, tol, mi = model.α, model.tol, model.mi
    check_size(x, t)
    if ndims(x) == 1 x = x[:, :] end
    w = zeros(size(x, 1) + 1)
    e = 0.0
    for _ in 1 : mi
        eb = e
        update!(x, t, w, α)
        e = sum(abs.(w)) / length(w)
        abs(e - eb) <= tol && break
    end
    model.w = w[end:-1:1]
end

(model::Lasso)(x) = expand(x)' * model.w


using ReverseDiff: GradientTape, GradientConfig, gradient, gradient!, compile
using Turing: _hmc_step
using Turing
using HDF5, JLD
sv_data = load(TPATH*"/example-models/nips-2017/sv-data.jld.data")["data"]

@model sv_model(T, y) = begin
    ϕ ~ Uniform(-1, 1)
    σ ~ Truncated(Cauchy(0,5), 0, +Inf)
    μ ~ Cauchy(0, 10)
    # h = tzeros(Real, T)
    h = Vector{Real}(T)
    h[1] ~ Normal(μ, σ / sqrt(1 - ϕ^2))
    y[1] ~ Normal(0, exp.(h[1] / 2))
    for t = 2:T
      h[t] ~ Normal(μ + ϕ * (h[t-1] - μ) , σ)
      y[t] ~ Normal(0, exp.(h[t] / 2))
    end
  end

mf = sv_model(data=sv_data[1])
chain_nuts = sample(mf, HMC(0.05, 10), 2000)

println("mean of m: $(mean(chn[1000:end, :μ, :].value))")

# θ_dim = 1
# function lj_func(θ)
#   _lj = zero(Real)

#   s = 1

#   m = θ[1]
#   _lj += logpdf(Normal(0, sqrt(s)), m)

#   _lj += logpdf(Normal(m, sqrt(s)), 2.0)
#   _lj += logpdf(Normal(m, sqrt(s)), 2.5)

#   return _lj
# end

# neg_lj_func(θ) = -lj_func(θ)
# const f_tape = GradientTape(neg_lj_func, randn(θ_dim))
# const compiled_f_tape = compile(f_tape)

# function grad_func(θ)

#   inputs = θ
#   results = similar(θ)
#   all_results = DiffResults.GradientResult(results)

#   gradient!(all_results, compiled_f_tape, inputs)

#   neg_lj = all_results.value
#   grad, = all_results.derivs

#   return -neg_lj, grad

# end

# std = ones(θ_dim)
# θ = randn(θ_dim)
# lj = lj_func(θ)

# chn = []
# accept_num = 1

# totla_num = 5000
# for iter = 1:totla_num
#   push!(chn, θ)
#   θ, lj, is_accept, τ_valid, α = _hmc_step(θ, lj, lj_func, grad_func, 5, 0.05, std)
#   accept_num += is_accept
#   # if (iter % 50 == 0) println(θ) end
# end

# @show mean(chn), lj
# @show accept_num / totla_num


# ---
# title: 72. Edit Distance
# id: problem72
# author: Qling
# date: 2021-03-17
# difficulty: Hard
# categories: String, Dynamic Programming
# link: <https://leetcode.com/problems/edit-distance/description/>
# hidden: true
# ---
#
# Given two strings `word1` and `word2`, return _the minimum number of
# operations required to convert`word1` to `word2`_.
#
# You have the following three operations permitted on a word:
#
#   * Insert a character
#   * Delete a character
#   * Replace a character
#
#
#
# **Example 1:**
#
#
#
#     Input: word1 = "horse", word2 = "ros"
#     Output: 3
#     Explanation:
#     horse -> rorse (replace 'h' with 'r')
#     rorse -> rose (remove 'r')
#     rose -> ros (remove 'e')
#
#
# **Example 2:**
#
#
#
#     Input: word1 = "intention", word2 = "execution"
#     Output: 5
#     Explanation:
#     intention -> inention (remove 't')
#     inention -> enention (replace 'i' with 'e')
#     enention -> exention (replace 'n' with 'x')
#     exention -> exection (replace 'n' with 'c')
#     exection -> execution (insert 'u')
#
#
#
#
# **Constraints:**
#
#   * `0 <= word1.length, word2.length <= 500`
#   * `word1` and `word2` consist of lowercase English letters.
#
#
## @lc code=start
using LeetCode

function min_distance(word1::String, word2::String)::Int32
    m, n = length(word1) + 1, length(word2) + 1
    dp = fill(0, m, n)

    for i = 1:m
        dp[i, 1] = i - 1
    end

    for j = 1:n
        dp[1, j] = j - 1
    end

    for i = 2:m, j = 2:n
        if word1[i-1] == word2[j-1]
            dp[i, j] = dp[i-1, j-1]
        else
            dp[i, j] = min(dp[i-1, j-1], dp[i, j-1], dp[i-1, j]) + 1
        end
    end

    return dp[m, n]
end

## @lc code=end


#####
##### ordering of faces in terms of vertices
#####
function face_vertices(elem::Line)
    return 1,2
end

function find_face_nodes(elem::Tri,r,s,tol=50*eps())
    e1 = findall(@. abs(s+1)<tol)
    e2 = findall(@. abs(r+s)<tol)
    e3 = findall(@. abs(r+1)<tol)
    return e1,reverse(e2),reverse(e3)
end

function find_face_nodes(elem::Quad,r,s,tol=50*eps())
    e1 = findall(@. abs(s+1)<tol)
    e2 = findall(@. abs(r-1)<tol)
    e3 = findall(@. abs(s-1)<tol)
    e4 = findall(@. abs(r+1)<tol)
    return e1,e2,reverse(e3),reverse(e4)
end

function find_face_nodes(elem::Hex,r,s,t,tol=50*eps())
    fv1 = findall(@. abs(r+1) < tol)
    fv2 = findall(@. abs(r-1) < tol)
    fv3 = findall(@. abs(s+1) < tol)
    fv4 = findall(@. abs(s-1) < tol)
    fv5 = findall(@. abs(t+1) < tol)
    fv6 = findall(@. abs(t-1) < tol)
    return fv1,fv2,fv3,fv4,fv5,fv6
end

function find_face_nodes(elem::Tet,r,s,t,tol=50*eps())
    fv1 = findall(@. abs(s+1) < tol)
    fv2 = findall(@. abs(r+s+t+1) < tol)
    fv3 = findall(@. abs(r+1) < tol)
    fv4 = findall(@. abs(t+1) < tol)
    return fv1,fv2,fv3,fv4
end

# face vertices = face nodes of degree 1
face_vertices(elem) = find_face_nodes(elem,nodes(elem,1)...)

#####
##### face data for diff elements
#####

function map_face_nodes(elem::Tri, face_nodes)
    r1D = face_nodes
    e = ones(size(r1D)) # vector of all ones
    rf = [r1D; -r1D; -e];
    sf = [-e; r1D; -r1D];
    return rf,sf
end
function map_face_nodes(elem::Quad, face_nodes)
    r1D = face_nodes
    e = ones(size(r1D))
    rf = [r1D; e; -r1D; -e]
    sf = [-e; r1D; e; -r1D]
    return rf,sf
end
function map_face_nodes(elem::Hex, face_nodes...)
    r,s = face_nodes
    e = ones(size(r))
    rf = [-e; e; r; r; r; r]
    sf = [r; r; -e; e; s; s]
    tf = [s; s; s; s; -e; e]
    return rf,sf,tf
end
function map_face_nodes(elem::Tet, face_nodes...)
    r,s = face_nodes
    e = ones(size(r))
    rf = [r; r; -e; r]
    sf = [-e; s; r; s]
    tf = [s; -(e + r + s); s; -e]
    return rf,sf,tf
end

"""
    function inverse_trace_constant(rd::RefElemData)

Returns the degree-dependent constant in the inverse trace equality over the reference element (as
reported in ["GPU-accelerated dG methods on hybrid meshes"](https://doi.org/10.1016/j.jcp.2016.04.003)
by Chan, Wang, Modave, Remacle, Warburton 2016).

Can be used to estimate dependence of maximum stable timestep on degree of approximation.
"""
inverse_trace_constant(rd::RefElemData{1}) = (rd.N+1) * (rd.N+2) / 2
inverse_trace_constant(rd::RefElemData{2, Quad}) = (rd.N+1) * (rd.N + 2)
inverse_trace_constant(rd::RefElemData{3, Hex}) = 3 * (rd.N + 1) * (rd.N + 2) / 2
inverse_trace_constant(rd::RefElemData{1, Line, SBP{TensorProductLobatto}}) = rd.N * (rd.N + 1) / 2
inverse_trace_constant(rd::RefElemData{2, Quad, SBP{TensorProductLobatto}}) = rd.N * (rd.N + 1)
inverse_trace_constant(rd::RefElemData{3, Hex, SBP{TensorProductLobatto}}) = 3 * rd.N * (rd.N + 1) / 2

# precomputed
_inverse_trace_constants(rd::RefElemData{2, Tri, Polynomial}) = (6.0, 10.898979485566365, 16.292060161853993, 23.999999999999808, 31.884512140579055, 42.42373503225737, 52.88579066878113, 66.25284319164409, 79.3535377715693, 95.53911875636945)
_inverse_trace_constants(rd::RefElemData{3, Tet, Polynomial}) = (10., 16.892024376045097, 23.58210016200093, 33.828424659883034, 43.40423356477473, 56.98869932201791, 69.68035962892684)
inverse_trace_constant(rd::RefElemData{2, Tri, Polynomial}) where {Dim} = _inverse_trace_constants(rd)[rd.N]
inverse_trace_constant(rd::RefElemData{3, Tet, Polynomial}) where {Dim} = _inverse_trace_constants(rd)[rd.N]

# generic fallback
function inverse_trace_constant(rd::RefElemData)
    return maximum(eigvals(Matrix(rd.Vf' * diagm(rd.wf) * rd.Vf), Matrix(rd.Vq' * diagm(rd.wq) * rd.Vq)))
end


# Calculator in Julia

function calculator()
    println("------- Welcome to the Calculator app written in Julia -------")
    while true
        # Print the Calculator app Menu|
        println("==== Menu ====")
        println("1. Addition ➕")
        println("2. Subtraction ➖")
        println("3  Multiplication ✖")
        println("4. Division ➗")
        println("5. Exit the Calculator app ↗")
        operation = parse(Float64, readline())
        # Convert the string value input to a float

        # Depending on the user's choice, do the corresponding operation
        if operation == 1
            print("Enter the first number: ")
            num1 = parse(Float64, readline())
            print("Enter the second number: ")
            num2 = parse(Float64, readline())
            println("Output: $(num1 + num2)")
        elseif operation == 2
            print("Enter the first number: ")
            num1 = parse(Float64, readline())
            print("Enter the second number: ")
            num2 = parse(Float64, readline())
            println("Output: $(num1 - num2)")
        elseif operation == 3
            print("Enter the first number: ")
            num1 = parse(Float64, readline())
            print("Enter the second number: ")
            num2 = parse(Float64, readline())
            println("Output: $(num1 * num2)")
        elseif operation == 4
            print("Enter the first number: ")
            num1 = parse(Float64, readline())
            print("Enter the second number: ")
            num2 = parse(Float64, readline())
            println("Output: $(num1 / num2)")
        elseif operation == 5
            println("Exiting calculator")
            break # Stop the while loop if the user decides to exit
        else
            # If the user enters another number, then ask to try again
            println("Invalid input, please try again.")
        end

    end
end


using Test, SegmentIntersections

@testset "Test random lines" begin
    random_data = 100 .* rand(100, 4);
    segments = Segment{Float64}[]
    for datap in eachrow(random_data)
        push!(segments, Segment(datap...))
    end
    bo_intersections = find_intersections(segments)
    brute_intersections = find_intersections_brute(segments)
    @test length(bo_intersections) == length(brute_intersections)
    for expected in brute_intersections
        match = false
        for intersection in bo_intersections
            if isapprox(intersection, expected, rtol = 1e-2)
                match = true
            end
        end
        @test match == true
    end
end

# This test does currently not pass
#@testset "Test complete graph" begin
#    n = 4
#    r = 100
#    theta_values = range(0, 2π, length=n+1)[1:end-1]
#    coordinates = [[r*cos(theta), r*sin(theta)] for theta in theta_values]
#    segments = Segment{Float64}[]
#    for i in 1:n
#        for j in (i+1):n
#            push!(segments, Segment(Point(coordinates[i]...), Point(coordinates[j]...)))
#        end
#    end
#    bo_intersections = find_intersections(segments)
#    brute_intersections = find_intersections_brute(segments)
#    @test length(bo_intersections) == length(brute_intersections)
#    matches = 0
#    println("checking intersections...")
#    for expected in brute_intersections
#        for intersection in bo_intersections
#            if isapprox(intersection, expected, rtol = 1e-2)
#                matches += 1
#            end
#        end
#    end
#    @test matches == length(bo_intersections)
#end


using Flux
using Flux: gradient
using Flux.Optimise: update!
using DelimitedFiles, Statistics
using Parameters: @with_kw

# This replicates the housing data example from the Knet.jl readme. Although we
# could have reused more of Flux (see the mnist example), the library's
# abstractions are very lightweight and don't force you into any particular
# strategy.

# Struct to define hyperparameters
@with_kw mutable struct Hyperparams
    lr::Float64 = 0.1		# learning rate
    split_ratio::Float64 = 0.1	# Train Test split ratio, define percentage of data to be used as Test data
end

function get_processed_data(args)
    isfile("housing.data") ||
        download("https://raw.githubusercontent.com/MikeInnes/notebooks/master/housing.data",
            "housing.data")

    rawdata = readdlm("housing.data")'

    # The last feature is our target -- the price of the house.
    split_ratio = args.split_ratio # For the train test split

    x = rawdata[1:13,:]
    y = rawdata[14:14,:]

    # Normalise the data
    x = (x .- mean(x, dims = 2)) ./ std(x, dims = 2)

    # Split into train and test sets
    split_index = floor(Int,size(x,2)*split_ratio)
    x_train = x[:,1:split_index]
    y_train = y[:,1:split_index]
    x_test = x[:,split_index+1:size(x,2)]
    y_test = y[:,split_index+1:size(x,2)]

    train_data = (x_train, y_train)
    test_data = (x_test, y_test)

    return train_data,test_data
end

# Struct to define model
mutable struct model
    W::AbstractArray
    b::AbstractVector
end

# Function to predict output from given parameters
predict(x, m) = m.W*x .+ m.b

# Define the mean squared error function to be used in the loss
# function. An implementation is also available in the Flux package
# (https://fluxml.ai/Flux.jl/stable/models/losses/#Flux.Losses.mse).
meansquarederror(ŷ, y) = sum((ŷ .- y).^2)/size(y, 2)

function train(; kws...)
    # Initialize the Hyperparamters
    args = Hyperparams(; kws...)

    # Load the data
    (x_train,y_train),(x_test,y_test) = get_processed_data(args)

    # The model
    m = model((randn(1,13)),[0.])

    loss(x, y) = meansquarederror(predict(x, m), y)

    ## Training
    η = args.lr
    θ = params(m.W, m.b)

    for i = 1:500
        g = gradient(() -> loss(x_train, y_train), θ)
        for x in θ
            update!(x, g[x]*η)
        end
        if i%100==0
            @show loss(x_train, y_train)
        end
    end

    # Predict the RMSE on the test set
    err = meansquarederror(predict(x_test, m),y_test)
    println(err)
end
##===============mmy
using DifferentialEquations, Plots
## Setup ODE
function lotka_volterra(du,u,p,t)
  x, y = u
  α, β, δ, γ = p
  du[1] = dx = α*x - β*x*y
  du[2] = dy = -δ*y + γ*x*y
end
u0 = [1.0,1.0]
tspan = (0.0,10.0)
p = [1.5,1.0,3.0,1.0]
prob = ODEProblem(lotka_volterra,u0,tspan,p)
# Verify ODE solution
sol = solve(prob,Tsit5())
plot(sol)

# Generate data from the ODE
sol = solve(prob,Tsit5(),saveat=0.1)
A = sol[1,:] # length 101 vector
t = 0:0.1:10.0
scatter!(t,A)

using Flux, DiffEqFlux
# Build a neural network that sets the cost as the difference
# from the generated data and 1
p = [2.2, 1.0, 2.0, 0.4]
p = Flux.params(p) # Initial Parameter Vector

function predict_rd() # Our 1-layer neural network
  diffeq_rd(p,prob,Tsit5(),saveat=0.1)[1,:]
end
loss_rd() = sum(abs2,x-1 for x in predict_rd()) # loss function

# Optimize the parameters so the ODE's solution stays near 1
data = Iterators.repeated((), 100)
opt = ADAM(0.1)
Flux.train!(loss_rd, [p], data, opt, cb = cb)

##=====Poisson
using NeuralPDE, Flux, ModelingToolkit, GalacticOptim, Optim, DiffEqFlux
@parameters x y
@variables u(..)
Dxx = Differential(x)^2
Dyy = Differential(y)^2
# 2D PDE
eq  = Dxx(u(x,y)) + Dyy(u(x,y)) ~ -sin(pi*x)*sin(pi*y)
# Boundary conditions
bcs = [u(0,y) ~ 0.f0, u(1,y) ~ -sin(pi*1)*sin(pi*y),
       u(x,0) ~ 0.f0, u(x,1) ~ -sin(pi*x)*sin(pi*1)]
# Space and time domains
domains = [x ∈ IntervalDomain(0.0,1.0), y ∈ IntervalDomain(0.0,1.0)]
# Discretization
dx = 0.1
# Neural network
dim = 2 # number of dimensions
chain = FastChain(FastDense(dim,16,Flux.σ),FastDense(16,16,Flux.σ),
                  FastDense(16,1))
discret = PhysicsInformedNN(chain, GridTraining(dx)) # PINNs model
pde_system = PDESystem(eq,bcs,domains,[x,y],[u])
prob = discretize(pde_system,discret)
res = GalacticOptim.solve(prob, Optim.BFGS(); maxiters=1000)
phi = discret.phi

cd(@__DIR__)
train()


# This file is part of the TaylorIntegration.jl package; MIT licensed

using TaylorIntegration
using Test

const _order = 28
const _abstol = 1.0E-20
const tT = Taylor1(_order)

@testset "Tests: dot{x}=x^2, x(0) = 1" begin
    eqs_mov(x, p, t) = x^2
    t0 = 0.0
    x0 = 1.0
    x0T = Taylor1(x0, _order)
    tT[1] = t0
    TaylorIntegration.jetcoeffs!(eqs_mov, tT, x0T, nothing)
    @test x0T.coeffs[end] == 1.0
    δt = _abstol^inv(_order-1)
    @test TaylorIntegration.stepsize(x0T, _abstol) == δt

    tv, xv = taylorinteg(eqs_mov, 1, 0.0, 1.0, _order, _abstol)
    @test length(tv) == 501
    @test length(xv) == 501
    @test xv[1] == x0
    @test tv[end] < 1.0

    tv, xv = taylorinteg(eqs_mov, x0, 0.0, 1.0, _order, _abstol, nothing)
    @test length(tv) == 501
    @test length(xv) == 501
    @test xv[1] == x0
    @test tv[end] < 1.0

    trange = 0.0:1/8:1.0
    xv = taylorinteg(eqs_mov, 1, trange, _order, _abstol)
    @test length(xv) == length(trange)
    @test typeof(xv) == Array{typeof(x0),1}
    @test xv[1] == x0
    @test isnan(xv[end])
    @test abs(xv[5] - 2.0) ≤ eps(2.0)
    tvr, xvr = taylorinteg(eqs_mov, x0, trange[1], trange[end-1], _order, _abstol)
    @test tvr[1] == trange[1]
    @test tvr[end] == trange[end-1]
    @test xvr[1] == xv[1]
    @test xvr[end] == xv[end-1]

    trange = 0.0:1/8:1.0
    xv = taylorinteg(eqs_mov, x0, trange, _order, _abstol, nothing)
    @test length(xv) == length(trange)
    @test typeof(xv) == Array{typeof(x0),1}
    @test xv[1] == x0
    @test isnan(xv[end])
    @test abs(xv[5] - 2.0) ≤ eps(2.0)
    tvr, xvr = taylorinteg(eqs_mov, x0, trange[1], trange[end-1], _order, _abstol, nothing)
    @test tvr[1] == trange[1]
    @test tvr[end] == trange[end-1]
    @test xvr[1] == xv[1]
    @test xvr[end] == xv[end-1]

    tarray = vec(trange)
    xv2 = taylorinteg(eqs_mov, x0, tarray, _order, _abstol)
    @test xv[1:end-1] == xv2[1:end-1]
    @test xv2[1:end-1] ≈ xv[1:end-1] atol=eps() rtol=0.0
    @test length(xv2) == length(tarray)
    @test typeof(xv2) == Array{typeof(x0),1}
    @test xv2[1] == x0
    @test isnan(xv2[end])
    @test abs(xv2[5] - 2.0) ≤ eps(2.0)
end

@testset "Tests: dot{x}=x^2, x(0) = 3; nsteps <= maxsteps" begin
    eqs_mov(x, p, t) = x.^2 #the ODE (i.e., the equations of motion)
    exactsol(t, x0) = x0/(1.0-x0*t) #the analytical solution
    t0 = 0.0
    tmax = 0.3
    x0 = 3.0
    x0T = Taylor1(x0, _order)
    tT[1] = t0
    TaylorIntegration.jetcoeffs!(eqs_mov, tT, x0T, nothing)
    @test x0T.coeffs[end] == 3.0^(_order+1)
    δt = (_abstol/x0T.coeffs[end-1])^inv(_order-1)
    @test TaylorIntegration.stepsize(x0T, _abstol) == δt

    tv, xv = taylorinteg(eqs_mov, x0, 0, tmax, _order, _abstol)
    @test length(tv) < 501
    @test length(xv) < 501
    @test length(tv) == 14
    @test length(xv) == 14
    @test xv[1] == x0
    @test tv[end] < 1/3
    @test tv[end] == tmax
    @test abs(xv[end]-exactsol(tv[end], xv[1])) < 2e-14

    tmax = 0.33
    tv, xv = taylorinteg(eqs_mov, x0, t0, tmax, _order, _abstol)
    @test length(tv) < 501
    @test length(xv) < 501
    @test length(tv) == 28
    @test length(xv) == 28
    @test xv[1] == x0
    @test tv[end] < 1/3
    @test tv[end] == tmax
    @test abs(xv[end]-exactsol(tv[end], xv[1])) < 5e-12
end

@testset "Test non-autonomous ODE (1): dot{x}=cos(t)" begin
    let f(x, p, t) = cos(t)
        t0 = 0//1
        tmax = 10.25*(2pi)
        abstol = 1e-20
        order = 25
        x0 = 0.0 #initial conditions such that x(t)=sin(t)
        tv, xv = taylorinteg(f, x0, t0, tmax, order, abstol)
        @test length(tv) < 501
        @test length(xv) < 501
        # @test length(xT[:,2]) < 501
        @test xv[1] == x0
        @test tv[1] == t0
        @test abs(sin(tmax)-xv[end]) < 1e-14

        tmax = 15*(2pi)
        tv, xv = taylorinteg(f, x0, t0, tmax, order, abstol)
        @test length(tv) < 501
        @test length(xv) < 501
        @test xv[1] == x0
        @test abs(sin(tmax)-xv[end]) < 1e-14
    end
end


@testset "BinaryHeap" begin
    data = [16, 2, 6, 5, 7, 1, 3, 4, 11]
    @testset "MinHeap" begin
        @testset "init heap/buildheap!" begin
            h = MinHeap(data)
            @test peek(h) == 1
            @test length(h) == 9
            @test !isempty(h)
            @test isequal(h.tree.data, [1, 2, 3, 4, 7, 6, 16, 5, 11])
        end
        @testset "insert! and remove!" begin
            h = MinHeap(Vector{Int}())
            @test_throws ArgumentError peek(h)
            @test length(h) == 0
            @test isempty(h)
            expected = [
                 [16],
                 [2, 16],
                 [2, 16, 6],
                 [2, 5, 6, 16],
                 [2, 5, 6, 16, 7],
                 [1, 5, 2, 16, 7, 6],
                 [1, 5, 2, 16, 7, 6, 3],
                 [1, 4, 2, 5, 7, 6, 3, 16],
                 [1, 4, 2, 5, 7, 6, 3, 16, 11]]
            for i = 1:length(data)
                insert!(h, data[i])
                @test length(h) == i
                @test !isempty(h)
                @test isequal(h.tree.data, expected[i])
            end
            removed = [1, 2, 3, 4, 5, 6, 7, 11, 16]
            while !isempty(h)
                @test remove!(h) == popfirst!(removed)
            end
        end
        @testset "remove! at position" begin
            h = MinHeap(data)
            @test_throws BoundsError(h, 10) remove!(h, 10)
            @test remove!(h, length(h)) == 11
            @test remove!(h, 3) == 3
            @test length(h) == 7
        end
    end
    @testset "MaxHeap" begin
        @testset "init heap/buildheap!" begin
            h = MaxHeap(data)
            @test peek(h) == 16
            @test length(h) == 9
            @test !isempty(h)
            @test isequal(h.tree.data, [16, 11, 6, 5, 7, 1, 3, 4, 2])
        end
        @testset "insert! and remove!" begin
            h = MaxHeap(Vector{Int}())
            @test_throws ArgumentError peek(h)
            @test length(h) == 0
            @test isempty(h)
            expected = [
                 [16],
                 [16, 2],
                 [16, 2, 6],
                 [16, 5, 6, 2],
                 [16, 7, 6, 2, 5],
                 [16, 7, 6, 2, 5, 1],
                 [16, 7, 6, 2, 5, 1, 3],
                 [16, 7, 6, 4, 5, 1, 3, 2],
                 [16, 11, 6, 7, 5, 1, 3, 2, 4]]
            for i = 1:length(data)
                insert!(h, data[i])
                @test length(h) == i
                @test !isempty(h)
                @test isequal(h.tree.data, expected[i])
            end
            removed = [1, 2, 3, 4, 5, 6, 7, 11, 16]
            while !isempty(h)
                @test remove!(h) == pop!(removed)
            end
        end
        @testset "remove! at position" begin
            h = MaxHeap(data)
            @test_throws BoundsError(h, 10) remove!(h, 10)
            @test remove!(h, length(h)) == 2
            @test remove!(h, 2) == 11
            @test length(h) == 7
        end
    end
end


"""
    Point(x)

represents a single point at `x`.
"""
struct Point{T} <: Domain{T}
    x::T
end

similardomain(d::Point, ::Type{T}) where {T} = Point{T}(d.x)

convert(::Type{Number}, d::Point{<:Number}) = d.x
convert(::Type{N}, d::Point{<:Number}) where N<:Number = convert(N, convert(Number, d.x))
Number(d::Point) = convert(Number, d)

convert(::Type{Domain}, c::Number) = Point(c)
convert(::Type{Domain{T}}, c::Number) where T = Point{T}(c)

==(a::Point,b::Point) = a.x == b.x
indomain(x, d::Point) = x == d.x
isempty(::Point) = false

approx_indomain(x, d::Point, tolerance) = norm(x-d.x) <= tolerance

dimension(d::Point{Vector{T}}) where {T} = length(d.x)

canonicaldomain(d::Point{T}) where {T<:StaticTypes} = Point(zero(T))
canonicaldomain(d::Point{T}) where {T<:AbstractVector} =
    Point(zeros(eltype(T),dimension(d)))

mapfrom_canonical(d::Point) = Translation(d.x)

isopenset(d::Point) = false
isclosedset(d::Point) = true

boundary(d::Point) = d
boundingbox(d::Point) = d.x..d.x

interior(d::Point{T}) where {T} = EmptySpace{T}()
closure(d::Point) = d

point_in_domain(d::Point) = d.x

function map_domain(map, p::Point)
    x = applymap(map, p.x)
    Point(x)
end

mapped_domain(invmap, p::Point) = map_domain(inverse(invmap), p)

for op in (:+,:-)
    @eval $op(a::Point, b::Point) = Point($op(a.x,b.x))
end

# Interval minus a point:
setdiffdomain(d::Interval, x::Number) = setdiffdomain(d, Point(x))
setdiffdomain(d::Interval, p::Point) = setdiffdomain(promote_domains((d,p))...)
function setdiffdomain(d::Interval{L,R,T}, p::Point{T}) where {L,R,T}
    a = leftendpoint(d)
    b = rightendpoint(d)
    x = p.x

    a == x && return Interval{:open,R,T}(a,b)
    a < x < b && return UnionDomain(Interval{L,:open,T}(a,p.x), Interval{:open,R,T}(p.x,b))
    b == x && return Interval{L,:open,T}(a,b)
    return d
end

issubset1(d1::Point, d2) = d1.x ∈ d2

setdiffdomain1(p::Point, d2) = issubset(p, d2) ? EmptySpace{eltype(p)}() : p


mutable struct BernoulliDist
    _mu::Float64
    _beta::BetaDist
    # if use BetaDist, _bayes == true
    _bayes::Bool
    function BernoulliDist(mu::AbstractArray{Float64,1})
        _mu::Float64 = sum(mu) / size(mu)[1]
        new(_mu, BetaDist(0, 0), false)
    end
    function BernoulliDist(beta::BetaDist)
        new(0, beta, true)
    end
end

function pdf(dist::BernoulliDist, x::Int64)
    return dist._mu^x * (1.0 - dist._mu)^(1 - x)
end

function fitting(dist::BernoulliDist, X::AbstractArray{Float64,1})
    if dist._bayes == false
        return fitting_ml(dist, X)
    else
        return fitting_map(dist, X)
    end
end

function fitting_ml(dist::BernoulliDist, X::AbstractArray{Float64,1})
    # X is the AbstractArray of 0 or 1
    # X .== 0 is the AbstractArray indicating if X[i] == 0
    n_zeros = sum(X .== 0.0)
    # X .== 1 is the AbstractArray indicating if X[i] == 1
    n_ones = sum(X .== 1.0)
    dist._mu = n_ones * 1.0 / (n_ones + n_zeros)
end

function fitting_map(dist::BernoulliDist, X::AbstractArray{Float64,1})
    # X is the AbstractArray of 0 or 1
    # X .== 0 is the AbstractArray indicating if X[i] == 0
    n_zeros = sum(X .== 0.0)
    # X .== 1 is the AbstractArray indicating if X[i] == 1
    n_ones = sum(X .== 1.0)
    dist._beta._n_ones += n_ones
    dist._beta._n_zeros += n_zeros
end

function draw(bern::BernoulliDist, n_trials::Int64)
    trials = rand(Uniform(0, 1), n_trials)

    mu = 0
    if (bern._bayes == true)
        # sample from beta
        mu = bern._beta._n_ones / (bern._beta._n_ones + bern._beta._n_zeros)
    else
        # sample ml
        mu = bern._mu
    end

    return sum(mu .> trials)
end


import Base.*, Base./, Base.+, Base.-, Base.==, Base.^, Base.exp, Base.isapprox

@doc raw"""
    tf = TransferFunction([1, 2], [3, 5, 8])
    tf = TransferFunction([1, 2], [3, 5, 8], 3.0)

construct a transfer function representing a linear, time-invariant system.

# example
to construct the transfer function

```math
G(s) = \frac{4e^{-2.2s}}{2s+1}
```

in Julia:

```jldoctest
tf = TransferFunction([4], [2, 1], 2.2)
# output
    4.0
----------- e^(-2.2*s)
2.0*s + 1.0
```

# attributes
* `numerator::Polynomial{Float64, :s}`: the polynomial in the numerator of the transfer function
* `denominator::Polynomial{Float64, :s}`: the polynomial in the denominator of the transfer function
* `time_delay::Float64`: the associated time delay
"""
struct TransferFunction
    numerator::Polynomial{Float64, :s}
    denominator::Polynomial{Float64, :s}
    time_delay::Float64
end

TransferFunction(num::Union{Vector{Int}, Vector{Float64}},
                 den::Union{Vector{Int}, Vector{Float64}},
                 td::Union{Int, Float64}) =
    TransferFunction(Polynomial(reverse(convert(Vector{Float64}, num)), :s),
                     Polynomial(reverse(convert(Vector{Float64}, den)), :s),
                     convert(Float64, td))

TransferFunction(num::Union{Vector{Int}, Vector{Float64}},
                 den::Union{Vector{Int}, Vector{Float64}}) =
    TransferFunction(num, den, 0.0)

const s = TransferFunction([1, 0], [1])

function *(tf1::TransferFunction, tf2::TransferFunction)
    g = TransferFunction(tf1.numerator * tf2.numerator,
                         tf1.denominator * tf2.denominator,
                         tf1.time_delay + tf2.time_delay)
    return pole_zero_cancellation(g)
end

*(k::Number, tf::TransferFunction) = TransferFunction(k * tf.numerator, tf.denominator, tf.time_delay)
*(tf::TransferFunction, k::Number) = *(k, tf)
function ^(tf::TransferFunction, x::Int)
    # x< 0 case handled by MethodError
    if x == 0
        return TransferFunction([1], [1])
    else
        tfp = tf
        for i = 2:x
            tfp = tfp * tf
        end
        return tfp
    end
end

"""
    tf_time_delay = exp(- 3 * s)

Conveniently create a time delay by exp(- θ * s).

# Example
g = 1 / (s + 1) * exp(-2.0 * s) # introduce time delay of 2.0
"""
function exp(tf::TransferFunction)
    # if numerator is zero, then the value is 1.0 and no time delay is introduced
    if degree(tf.numerator) == -1
        return TransferFunction([1], [1])
    end
    if ! ((degree(tf.numerator) == 1) && (degree(tf.denominator) == 0) && (tf.time_delay == 0.0))
        error("exp only functions with exp(-3.0 * s) for example, to introduce time delays")
    end
    return TransferFunction([1.0], [1.0], -tf.numerator[1])
end

function +(tf1::TransferFunction, tf2::TransferFunction)
    if (tf1.time_delay != tf2.time_delay)
        error("""we cannot add two transfer functions with two different time delays.

        if you are trying to construct a closed-loop transfer function with a time delay,
        instead use `ClosedLoopTransferFunction`.""")
    end
    g = TransferFunction(tf1.numerator * tf2.denominator + tf1.denominator * tf2.numerator,
                            tf1.denominator * tf2.denominator,
                            tf1.time_delay)
    return pole_zero_cancellation(g)
end

+(tf::TransferFunction, x::Number) = tf + TransferFunction([x], [1.0])
+(x::Number, tf::TransferFunction) = +(tf, x)
-(tf::TransferFunction, x::Number) = +(tf, -x)
-(x::Number, tf::TransferFunction) = +(-1 * tf, x)
-(tf::TransferFunction) = -1.0 * tf
-(tf1::TransferFunction, tf2::TransferFunction) = +(tf1, -1*tf2)

function /(tf1::TransferFunction, tf2::TransferFunction)
    g = TransferFunction(tf1.numerator * tf2.denominator,
                        tf1.denominator * tf2.numerator,
                        tf1.time_delay - tf2.time_delay)
    return pole_zero_cancellation(g)
end
/(tf::TransferFunction, x::Number) = TransferFunction(tf.numerator / x, tf.denominator, tf.time_delay)
/(x::Number, tf::TransferFunction) = TransferFunction([x], [1.0], 0.0) / tf

# note: must have polynomails written in same form
==(tf1::TransferFunction, tf2::TransferFunction) = (tf1.numerator == tf2.numerator) && (tf1.denominator == tf2.denominator) && (tf1.time_delay == tf2.time_delay)

@doc raw"""
    tf = zpk_form(tf)

write transfer function `tf` in zeros, poles, k-factor form:

$$g(s)=k\dfrac{\Pi_j (s-z_j)}{\Pi_j (s-p_j)}$$
where $z_j$ is zero $j$, $p_j$ is pole $j$, and $k$ is a constant factor (not equal to the zero-frequency gain) that uniquely specifies the transfer function.

this is achieved by multiplying by 1.0 in a fancy way such that the highest power of $s$ in the denominator is associated with a coefficient of $1$.

# Example

```jldoctest
g = 8.0 / (2 * s^2 + 3 * s + 4)
g_zpk = zpk_form(g)
# output
         4.0
---------------------
1.0*s^2 + 1.5*s + 2.0
```
"""
function zpk_form(tf::TransferFunction)
    coeff_highest_power_denom = tf.denominator[end]

    return TransferFunction(
        tf.numerator / coeff_highest_power_denom,
        tf.denominator / coeff_highest_power_denom,
        tf.time_delay)
end

function isapprox(tf1::TransferFunction, tf2::TransferFunction; atol=0.0)
    # to directly compare numerators and denominators, put in zpk form
    tf1s = zpk_form(tf1)
    tf2s = zpk_form(tf2)
    return isapprox(tf1s.time_delay, tf2s.time_delay, atol=atol) && isapprox(
	    tf1s.numerator.coeffs, tf2s.numerator.coeffs, atol=atol) && isapprox(tf1s.denominator.coeffs, tf2s.denominator.coeffs, atol=atol)
end

_zeros(tf::TransferFunction) = roots(tf.numerator)

_poles(tf::TransferFunction) = roots(tf.denominator)

_k(tf::TransferFunction) = tf.numerator[end] / tf.denominator[end]

@doc raw"""
    K = zero_frequency_gain(tf)

compute the (signed) zero frequency gain of a transfer function $g(s)$, which is:

```math
K := \lim_{s\rightarrow 0} G(s)
```

the zero-frequency gain "represents the ratio of the steady state value of the output with respect to a step input" [source](http://www.cds.caltech.edu/~murray/books/AM05/pdf/am06-xferfcns_16Sep06.pdf)

# example

```jldoctest
g = 5 / (3 * s + 1)
K = zero_frequency_gain(g)
# output
5.0
```

# arguments
* `tf::TransferFunction`: the transfer function

# returns
* `K::Float64`: the zero-frequency gain of the transfer function
"""
function zero_frequency_gain(tf::TransferFunction)
    tf_simplified = pole_zero_cancellation(tf)
    return evaluate(tf_simplified, 0.0)
end

@doc raw"""
    # compute the zeros, poles, and k-factor of a transfer function
    z, p, k = zeros_poles_k(tf)
    # construct a transfer function from its zeros, poles, and k-factor
    tf = zeros_poles_k(z, p, k, time_delay=0.0)

the representation of a transfer function in this context is:

$$g(s)=k\dfrac{\Pi_j (s-z_j)}{\Pi_j (s-p_j)}$$
where $z_j$ is zero $j$, $p_j$ is pole $j$, and $k$ is a constant factor (not equal to the zero-frequency gain) that uniquely specifies the transfer function.

* the zeros are the zeros of the numerator of the transfer function.
* the poles are the zeros of the denominator of the transfer function.
"""
zeros_poles_k(tf::TransferFunction) = _zeros(tf), _poles(tf), _k(tf)

# docstring above. this overloaded function is for *constructing* tfs
function zeros_poles_k(zeros::Array, poles::Array, k::Union{Float64, Int};
                       time_delay::Union{Int64, Float64}=0)
    # construct numerator polynomial
    top = length(zeros) == 0 ? Polynomial(1.0, :s) : fromroots(zeros, var=:s)

    # construct denominator polynomial
    bottom = length(poles) == 0 ? Polynomial(1.0, :s) : fromroots(poles, var=:s)

    # owing to numerical errors, the coeff's could be imaginary with tiny imaginary parts
    #  let's check that the imaginary parts are tiny, then convert the coefficients to real.
    @assert maximum(abs.(imag.(top.coeffs))) < 1e-6
    @assert maximum(abs.(imag.(bottom.coeffs))) < 1e-6
    top = Polynomial(real.(top.coeffs), :s) # reconstruct a polynomial with real coeffs
    bottom = Polynomial(real.(bottom.coeffs), :s)

    return TransferFunction(k * top, bottom, time_delay)
end

@doc raw"""
    z, p, gain = zeros_poles_gain(tf)

Compute the zeros, poles, and zero-frequency gain of a transfer function.

* the zeros are the zeros of the numerator of the transfer function.
* the poles are the zeros of the denominator of the transfer function.
* the zero-frequency gain is the transfer function evaluated at $s=0$
"""
zeros_poles_gain(tf::TransferFunction) = _zeros(tf), _poles(tf), zero_frequency_gain(tf)

"""
    evaluate(tf, z)

evaluate a `TransferFunction`, `tf`, at a particular number `z` (could be complex).

# example
```jldoctest eval
tf = TransferFunction([1], [3, 1])
evaluate(tf, 1.0)
# output
0.25
```

```jldoctest eval
```
"""
function evaluate(tf::TransferFunction, z::Number)
    return tf.numerator(z) / tf.denominator(z) * exp(-tf.time_delay * z)
end

"""
    proper(tf)

Return `true` if transfer function `tf` is proper and `false` otherwise.
"""
proper(tf::TransferFunction) = degree(tf.numerator) <= degree(tf.denominator)

"""
    strictly_proper(tf)

Return `true` if transfer function `tf` is strictly proper and `false` otherwise.
"""
strictly_proper(tf::TransferFunction) = degree(tf.numerator) < degree(tf.denominator)

"""
    tf = pole_zero_cancellation(tf, verbose=false, digits=8)

find (pole, zero) pairs such that pole = zero and return a new transfer function with those pairs cancelled.
this is achieved by comparing the poles and zeros with `isapprox`, with poles and zeros rounded to
`digits` digits (also applies to reconstruction).

# arguments
* `tf::TransferFunction`: the transfer function
* `verbose::Bool=false`: print off which poles, zeros are cancelled.
* `digits::Int`: number of digits to round poles and zeros to, for (i) cancelling and (ii) reconstruction.

# example
```jldoctest
pole_zero_cancellation(s * (s - 1) / (s * (s + 1)))
# output
1.0*s - 1.0
-----------
1.0*s + 1.0
```
"""
function pole_zero_cancellation(tf::TransferFunction; verbose::Bool=false, digits::Int=6)
    # compute poles, zeros, and k-factor of the transfer function
    zs, ps, k = zeros_poles_k(tf)
    zs = round.(zs, digits=digits)
    ps = round.(ps, digits=digits)

    # store in these boolean arrays whether we will cancel them or not
    canceled_zeros = [false for i = 1:length(zs)]
    canceled_poles = [false for i = 1:length(ps)]

    # loop through poles and zeros. if they are equal, cancel!
    for (i_z, z) in enumerate(zs)
        for (i_p, p) in enumerate(ps)
            # the pole and zero are equal...
            # *and* if this pole has not already been canceled...
	    if isapprox(p, z) && (! canceled_poles[i_p])
                canceled_zeros[i_z] = true
                canceled_poles[i_p] = true
                break
            end
        end
    end
    @assert sum(canceled_zeros) == sum(canceled_poles)
    if verbose && sum(canceled_poles) != 0
        println("canceling the following poles and zeros: ", ps[canceled_poles])
    end

    if sum(canceled_zeros) == 0
        return tf
    else
        # reconstruct
        return zeros_poles_k(zs[.! canceled_zeros], ps[.! canceled_poles], k, time_delay=tf.time_delay)
    end
end

"""
    o = system_order(tf::TransferFunction)

return the order of the numerator and denominator of the transfer function `tf`.

use [`pole_zero_cancellation`](@ref) first if you wish to cancel poles and zeros that are equal before determining the order.

# returns
`o::Tuple{Int, Int}`: (order of numerator, order of denominator)

# examples
```jldoctest
g = 1 / (s + 1)
system_order(g)
# output
(0, 1)

g = (s + 1) / ((s + 2) * (s + 3))
system_order(g)
# output
(1, 2)
```
"""
function system_order(tf::TransferFunction)
    return (degree(tf.numerator), degree(tf.denominator))
end


export
    ecl_to_equ,
    equ_to_ecl,
    equ_to_horiz,
    equ_to_horiz1,
    geocentric_to_topocentric

"""
    ecl_to_equ(long, lat, obliquity)

Convert ecliptic to equatorial coordinates.

[Meeus-1998: equations 13.3, 13.4]

Parameters:
    long : ecliptic longitude in radians
    lat : ecliptic latitude in radians
    obliquity : obliquity of the ecliptic in radians

Returns:
    Right ascension in radians (2 pi is 24 hours)
    Declination in radians
"""
function ecl_to_equ(long, lat, obliquity)
    cose = cos(obliquity)
    sine = sin(obliquity)
    sinl = sin(long)
    x = cos(long)
    y = sinl * cose - tan(lat) * sine
    alpha1 = atan(y, x)
    ra = mod2pi(alpha1)
    # convert to hours
    declin = asin(sin(lat) * cose + cos(lat) * sine * sinl)
    return (ra, declin)
end

"""
    equ_to_horiz(H, decl)

Convert equatorial to horizontal coordinates.

[Meeus-1998: equations 13.5, 13.6]

Note that azimuth is measured westward starting from the south.

This is not a good formula for using near the poles.

Parameters:
    H : hour angle in radians
    decl : declination in radians

Returns:
    azimuth in radians
    altitude in radians
"""
function equ_to_horiz(H, decl)
    global latitude
    cosH = cos(H)
    sinLat = sin(latitude)
    cosLat = cos(latitude)
    alt = asin((sinLat * sin(decl)) + (cosLat * cos(decl) * cosH))

    # atan(y,x)... not the other way round
    azi = atan((-1 * cos(decl)) * cosLat * sin(H), sin(decl) - (sinLat * sin(alt)))

    # Duffett-Smith says if azi positive, true azimuth is (360 - azi)
    if sin(H) > 0
    azi = mod2pi(azi)
    end
    return (azi, alt)
end

"""
    equ_to_ecl(ra, declin, obliquity)

Convert equatorial to ecliptic coordinates.

[Meeus-1998: equations 13.1, 13.2]

Parameters:
    ra : right accension radians
    declin : declination in radians
    obliquity : obliquity of the ecliptic in radians

Returns:
    ecliptic longitude in radians
    ecliptic latitude in radians
"""
function equ_to_ecl(ra, declin, obliquity)
    cose = cos(obliquity)
    sine = sin(obliquity)
    sina = sin(ra)
    long = mod2pi(atan(sina * cose + tan(declin) * sine, cos(ra)))
    lat = mod2pi(asin(sin(declin) * cose - cos(declin) * sine * sina))
    return (long, lat)
end

"""
    geocentric_to_topocentric(phi, H, L, ra, decl, d, jd)

Convert geocentric to topocentric coordinates.

[Meeus-1998: equations 40.2, 40.3]

Parameters:
    phi: geocentric radius in radians
    H : observer's altitude (above sea level) in meters
    L ; observer longitude (in radians)
    ra : body right ascension
    decl : body declination in radians
    d : body distance (in AU)
    jd : Julian date

Returns:
    topocentric right ascension in radians
    topocentric declination in radians
"""
function geocentric_to_topocentric(phi, H, L, ra, decl, d, jd)
     sinparallax = sin(deg2rad(8.794/3600))/d
     ha    = mod2pi(apparent_sidereal_time_greenwich(jd) - L - deg2rad(ra* 15))
     p1, psinp1, pcosp1 = geographical_to_geocentric_lat(phi, H)
     tmp   = cos(decl) - pcosp1 * sinparallax * cos(ha)
     da    = atan(-pcosp1 * sinparallax * sin(ha), tmp)
     tdecl = atan((sin(decl) - psinp1 * sinparallax) * cos(da), tmp)
    return (ra+da, tdecl)
end


using NumericalMethodsforEngineers

ProjDir = dirname(@__FILE__)
cd(ProjDir) do

  function f(x::Float64, y::Vector{Float64})
    [y[2], -2y[2]/y[1]]
  end

  res = shootingmethod(f, [0.0, 1.0, 1.0, 2.0], [1.0, 3.0], 10, 0.00001, 40)
  @show round.(res, 4)

end