"""
Helper function for quadratic_interval that computes roots of a
real quadratic using interval arithmetic to bound rounding errors.
"""
function quadratic_helper!(a::Interval{T}, b::Interval{T}, c::Interval{T}, L::Array{Interval{T}}) where {T}

    Δ = b^2 - 4*a*c

    Δ.hi < 0 && return

    Δ = sqrt(Δ)

    if (b.lo >= 0)
        x0 = -0.5 * (b + Δ)

    else
        x0 = -0.5 * (b - Δ)
    end

    if (0 ∈ x0)
        push!(L, x0)

    else
        x1 = c / x0
        x0 = x0 / a
        push!(L, x0, x1)
    end

end

"""
Function to solve a quadratic equation where the coefficients are intervals.
Returns an array of intervals of the roots.
Arguments `a`, `b` and `c` are interval coefficients of `x²`, `x` and `1` respectively.
The interval case differs from the non-interval case in that
there might be three disjoint interval roots. In the third
case, one interval root extends to −∞ and another extends to +∞.
This algorithm finds the set of points where `F.lo(x) ≥ 0` and the set
of points where `F.hi(x) ≤ 0` and takes the intersection of these two sets.
Eldon Hansen and G. William Walster : Global Optimization Using Interval Analysis - Chapter 8
"""
function quadratic_roots(a::Interval{T}, b::Interval{T}, c::Interval{T}) where {T}

    L = Interval{T}[]
    R = Interval{T}[]

    quadratic_helper!(Interval(a.lo), Interval(b.lo), Interval(c.lo), L)
    quadratic_helper!(Interval(a.hi), Interval(b.hi), Interval(c.hi), L)
    quadratic_helper!(Interval(a.lo), Interval(b.hi), Interval(c.lo), L)
    quadratic_helper!(Interval(a.hi), Interval(b.lo), Interval(c.hi), L)

    if (length(L) == 8)
        resize!(L, 4)
    end

    if (a.lo < 0 || (a.lo == 0 && b.hi == 0) || (a.lo == 0 && b.hi == 0 && c.lo ≤ 0))
        push!(L, Interval(-∞))
    end

    if (a.lo < 0 || (a.lo == 0 && b.lo == 0) || (a.lo == 0 && b.lo == 0 && c.lo ≤ 0))
        push!(L, Interval(∞))
    end

    sort!(L, by = x -> x.lo)

    for i in 1:2:length(L)
        push!(R, Interval(L[i].lo, L[i+1].hi))
    end

    R
end


#=
   Locomotive problem.

   From Think Bayes, page 22ff
   """
   I found the locomotive problem in Frederick Mosteller's "Fifty Challenging
   Problems in Probability with Solutions" (Dover, 1987):
      'A railroad numbers its locomotives in order 1..N. One day you
       see a locomotive with the number 60. Estimate how many loco-
       motives the railroad has.'

   ...

   The mean of the posterior is 333, so that might be a good guess if you wanted to
   minimize error.
   """

   As the book (Think Bayes) mentions, the model is very sensitive to
   the value of maxInt. Here are some results for different maxInt:

   Here are the values from this model for different values of maxInt
   (using MH())
     maxInt:100  mean(n): 78.23232
     maxInt:200  mean(n): 115.87732
     maxInt:500  mean(n): 206.42599
     maxInt:1000  mean(n): 335.66271
     maxInt:2000  mean(n): 552.68673

   These are the exact values from WebPPL model using the enumerate sampler.
   maxInt  n
   -------------------------
    100    77.8166306057923
    200   115.84577808279282
    500   207.2393675826458
   1000   334.04414386751915
   2000   553.5237331955558

   Mosteller's solution: 2*(60-1)+1: 119.

  cf german_tank.jl
     german_tank_int.jl
     ~/blog/locomotive_problem.blog
     ~/webppl/locomotive_problem.wppl
=#

using Turing
include("jl_utils.jl")

@model function locomotive_problem(ys,maxInt=1000)
    len = length(ys)
    n ~ DiscreteUniform(1,maxInt)
    ys ~ filldist(DiscreteUniform(1,n),len)
end

# Run the model with different values of maxInt
function run_locomotive_problem(ys,maxInt)
   model = locomotive_problem(ys,maxInt)

    # chns = sample(model, Prior(), 100_000)
    chns = sample(model, MH(), 100_000)
    # chns = sample(model, PG(15), 10_000)
    # chns = sample(model, SMC(), 10_000)
    # chns = sample(model, IS(), 100_000)
    # display(chns)

    println("maxInt:$maxInt  mean(n): ", mean_val(chns,:n))
end

ys = [60]
for maxInt in [100,200,500,1000,2000]
    run_locomotive_problem(ys,maxInt)
end


export det
export minor
export cofactor

sub(M::Array{Float64}, i, j) = M[1:end .!= i, 1:end .!= j]
det(a::Array{Float64}) = begin
    if size(a) == (2, 2)
        a[1, 1] * a[2, 2] - a[1, 2] * a[2, 1]
    else
        sum([(-1)^(i+1) * a[i, 1] * det(sub(a, i, 1)) for i in 1:size(a, 1)])
        # performance tip, look for the best row/column choice for expansion.
    end
end
minor(M::Array{Float64}, i, j) = det(sub(M, i, j))
cofactor(M::Array{Float64}, i, j) = minor(M, i, j) * (-1)^(i+j)


@doc raw"""
# Lotka-Volterra model in 2D with symmetric Lagrangian

```math
\begin{aligned}
L (q, \dot{q}) &= \frac{1}{2} \frac{\log q_2}{q_1} \, \dot{q_1} - \frac{1}{2} \frac{\log q_1}{q_2} \, \dot{q_2} - H(q) , \\
H(q) &= a_1 \, q_1 + a_2 \, q_2 + b_1 \, \log q_1 + b_2 \, \log q_2
\end{aligned}
```

This Lagrangian is a slight generalization of Equation (5) in José Fernández-Núñez,
Lagrangian Structure of the Two-Dimensional Lotka-Volterra System, International
Journal of Theoretical Physics, Vol. 37, No. 9, pp. 2457-2462, 1998.

"""
module LotkaVolterra2dSymmetric

    ϑ₁(t, q) = + log(q[2]) / q[1] / 2
    ϑ₂(t, q) = - log(q[1]) / q[2] / 2

    dϑ₁dx₁(t, q) = - log(q[2]) / q[1]^2 / 2
    dϑ₁dx₂(t, q) = + 1 / (q[1] * q[2]) / 2

    dϑ₂dx₁(t, q) = - 1 / (q[2] * q[1]) / 2
    dϑ₂dx₂(t, q) = + log(q[1]) / q[2]^2 / 2

    include("lotka_volterra_2d_common.jl")
    include("lotka_volterra_2d_equations.jl")

end


using Plots

function vfsolvebasic(vnew, kgrid, tolerance, imax)
    α = 0.3
    δ = 0.1
    β = 0.9

    n = length(kgrid)
    v = vnew .+ 2*tolerance
    cartesianindex = Array{CartesianIndex{2}, n}
    i = 1
    while maximum(abs.(v - vnew)) > tolerance && i<=imax
        v = vnew;
        c = zeros(n, n);
        for i in 1:n
            for j in 1:n
                c[i,j] = kgrid[i]^α + (1-δ)*kgrid[i] - kgrid[j];
                if c[i,j] < 0
                    c[i,j] = 0
                end
            end
            (vnew, cartesianindex) = findmax(log.(c) .+ β*v', dims = 2);
        end
        i += 1
    end
    kprimeindex = getindex.(cartesianindex, 2)
    return (v = vnew, kprime = kgrid[kprimeindex], kprimindex = kprimeindex)
end

kupper = 2
klower = 0.001
n = 10
tolerance = 0.001
imax = 1000
kgrid = collect(range(klower, stop = kupper, length = n))
(vbasic, kprimebasic, kprimeindex) = vfsolvebasic(zeros(n), kgrid, tolerance, imax)

# Check it agrees with previous version
scatter(kgrid, vbasic, label = "v")
scatter!(kgrid, vnew, label = "v")

scatter(kgrid, kprime, label = "k'")
scatter!(kgrid, kprimebasic, label = "k'")

vbasic == vnew
kprimebasic == kprime


using DifferentialEquations
using NLsolve
using Plots

# The range of values α = M₂/(M₁+M₂) where the Lagrange points will be found
αvals = 0.001:0.001:0.5
# The number of values of α
nαvals = length(αvals)
# Will store the positions of the Lagrange points
lagrangepoints = zeros(Float64, nαvals, 10);

for i ∈ 1:nαvals
    # Set the value of α = M₂/(M₁+M₂)
	α = αvals[i]

    # Set up the function for NLsolve to use
	function evalF!(F, X)
		x, y = X
		F[1] = x -
			(1 - α)*(x + α)/√((x + α)^2 + y^2)^3 -
			α*(x + α - 1)/√((x + α - 1)^2 + y^2)^3
		F[2] = y -
			(1 - α)*y/√((x + α)^2 + y^2)^3 -
			α*y/√((x + α - 1)^2 + y^2)^3
	end

    # Run nlsolve at each of the five predetermined starting positions to find the five Lagrange points
	lagrangepoints[i,1:2 ] = nlsolve(evalF!, [ 0.0, 0.0], autodiff=:forward).zero
	lagrangepoints[i,3:4 ] = nlsolve(evalF!, [ 1.0, 0.0], autodiff=:forward).zero
	lagrangepoints[i,5:6 ] = nlsolve(evalF!, [-1.0, 0.0], autodiff=:forward).zero
	lagrangepoints[i,7:8 ] = nlsolve(evalF!, [ 0.0, 1.0], autodiff=:forward).zero
	lagrangepoints[i,9:10] = nlsolve(evalF!, [ 0.0,-1.0], autodiff=:forward).zero

end

# Plot the values of α against the coordinates of the Lagrange points,
#  both as found numerically, and as approximated theoretically
function locationplots()
    return plot(
        plot(αvals, [lagrangepoints[:,1 ], α -> 1 - ∛(α/3)]),
        plot(αvals, [lagrangepoints[:,2 ], α -> 0]),
        plot(αvals, [lagrangepoints[:,3 ], α -> 1 + ∛(α/3)]),
        plot(αvals, [lagrangepoints[:,4 ], α -> 0]),
        plot(αvals, [lagrangepoints[:,5 ], α -> - 1 - 5α/12]),
        plot(αvals, [lagrangepoints[:,6 ], α -> 0]),
        plot(αvals, [lagrangepoints[:,7 ], α -> 1/2 - α]),
        plot(αvals, [lagrangepoints[:,8 ], α -> (√3)/2]),
        plot(αvals, [lagrangepoints[:,9 ], α -> 1/2 - α]),
        plot(αvals, [lagrangepoints[:,10], α -> -(√3)/2]),
        linestyle = [:solid :dash],
        layout = (5,2),
        size = (600, 1000),
        label = ["numerical" "theoretical"],
        yticks = false,
        title = ["L₁ x" "L₁ y" "L₂ x" "L₂ y" "L₃ x" "L₃ y" "L₄ x" "L₄ y" "L₅ x" "L₅ y"]
    )
end

# Calculates the acceleration of m, given position r and velocity v, with p a vector of parameters
function acceleration(r, v, p)
    # Parameters for the system
    G, M₁, M₂, r₁, r₂, Ω = p

    # The four forces (including fictitious) acting on m
    gravitational₁ = - ( G * M₁ / hypot((r - r₁)...)^3 ) .* (r - r₁)
    gravitational₂ = - ( G * M₂ / hypot((r - r₂)...)^3 ) .* (r - r₂)
    coriolis = - 2Ω .* [[0, -1] [1, 0]] * v
    centrifugal = (Ω^2) .* r
    return gravitational₁ + gravitational₂ + coriolis + centrifugal
end

# Sets up the derivative function for DifferentialEquations to use
function derivative!(du, u, p, t)
    du[1] = u[3]
    du[2] = u[4]
    du[3], du[4] = acceleration(u[1:2], u[3:4], p)
end

# Uses DifferentialEquations to simulate the trajectory of m for time T with initial conditions
#  r₀ and v₀, and parameters α and M₁
function simulate(α, G, r₀, v₀, T)
    # M₂ is the mass of the other large stationary body
    M₂ = M₁ * α/(1-α)

    # r₁ and r₂ are the positions of M₁ and M₂
    r₁ = [-α, 0]
    r₂ = [1-α, 0]

	# G is the gravitational constant, which is 1 for an appropriate system of units
	G = 1
    # Ω is the magnitude of the angular velocity of the rotating frame of reference
    Ω = √(G * (M₁ + M₂))

    # DifferentialEquations simulates the mass for time T
    return solve(ODEProblem(derivative!, vcat(r₀, v₀), (0.0, float(T)), [G, M₁, M₂, r₁, r₂, Ω]))
end

# Creates a diagram of the Sun/Earth/Moon system to test the simulation
function moondiagram()
    diagram = plot(
		simulate(3e-6, 39, [1.003, 0.0], [0.0, 0.215], 1/12),
		vars = (1,2),
		linecolor = :white,
		linewidth = 1,
		arrow = true,
		bg = :black,
		ticks = false,
		showaxis = false,
		legend = false,
		xlims = [0.995, 1.005],
		ylims = [-0.005, 0.005]
	)

	scatter!(
		diagram,
		[(1-(3e-6),0)],
		markercolor = :aqua,
		markersize = 4
	)

    return diagram
end

# Plots the trajectory of m starting near Lₙ up to time T, with parameters α and M₁
function trajectory(α, M₁, n, T, lims = [2,1])
	i = findfirst(≥(α), αvals)
    simulatedtrajectory = simulate(αvals[i], M₁, lagrangepoints[i, (2n-1):2n], [0.0, 0.0], T)

	diagram = plot(
		simulatedtrajectory,
		vars = (1,2),
		linecolor = :white,
		linewidth = 1,
		bg = :black,
		ticks = false,
		showaxis = false,
		legend = false,
        xlims = [-lims[1],lims[1]],
        ylims = [-lims[2],lims[2]],
        size = Tuple(lims) .* round(Int64, 1200/sum(lims))
	)

	scatter!(
		diagram,
		[Tuple(simulatedtrajectory.u[1][1:2]), (-α,0), (1-α, 0)],
		markercolor = [:white, :yellow, :aqua],
		markersize = [3, 10, 4]
	)

	return diagram
end


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

"""
    spheredir(θ, φ)

Return the 3D direction given polar angle `θ` and
azimuthal angle `φ` in degrees according to the ISO
convention.
"""
function spheredir(theta, phi)
  θ, φ = deg2rad(theta), deg2rad(phi)
  Vec(sin(θ)*cos(φ), sin(θ)*sin(φ), cos(θ))
end

"""
    planebasis(normal)

Return 2D basis vectors in the plane with given 3D `normal`.
"""
function planebasis(normal::Vec{3,T}) where {T}
  # normalize input
  n = normal ./ sqrt(sum(normal[i]^2 for i in 1:3))

  # find last non-zero component
  idx = -1
  for (i, c) in enumerate(reverse(n))
    if c != 0
      idx = length(n) - i + 1
      break
    end
  end

  @assert idx > 0 "invalid normal vector"

  # first basis vector (perturb and subtract projection)
  u = ntuple(i -> i == idx%3 + 1 ? n[i] + one(T) : n[i], 3)
  l = sum(u[i]*n[i] for i in 1:3)
  u = ntuple(i -> u[i] - l*n[i], 3)

  # second basis vector (cross product)
  nx, ny, nz = n
  ux, uy, uz = u
  v = (ny*uz - nz*uz, nz*ux - nx*uz, nx*uy - ny*ux)

  # normalize output
  u = u ./ sqrt(sum(u[i]^2 for i in 1:3))
  v = v ./ sqrt(sum(v[i]^2 for i in 1:3))

  Vec(u), Vec(v)
end

planebasis(normal::NTuple{3,T}) where {T} = planebasis(Vec(normal))


doc"""
    Chisq(ν)

The *Chi squared distribution* (typically written χ²) with `ν` degrees of freedom has the
probability density function

$f(x; k) = \frac{x^{k/2 - 1} e^{-x/2}}{2^{k/2} \Gamma(k/2)}, \quad x > 0.$

If `ν` is an integer, then it is the distribution of the sum of squares of `ν` independent standard [`Normal`](:func:`Normal`) variates.

```julia
Chisq(k)     # Chi-squared distribution with k degrees of freedom

params(d)    # Get the parameters, i.e. (k,)
dof(d)       # Get the degrees of freedom, i.e. k
```

External links

* [Chi-squared distribution on Wikipedia](http://en.wikipedia.org/wiki/Chi-squared_distribution)

"""
immutable Chisq <: ContinuousUnivariateDistribution
    ν::Float64

    Chisq(ν::Real) = (@check_args(Chisq, ν > zero(ν)); new(ν))
end

@distr_support Chisq 0.0 Inf

#### Parameters

dof(d::Chisq) = d.ν
params(d::Chisq) = (d.ν,)

#### Statistics

mean(d::Chisq) = d.ν

var(d::Chisq) = 2.0 * d.ν

skewness(d::Chisq) = sqrt(8.0 / d.ν)

kurtosis(d::Chisq) = 12.0 / d.ν

mode(d::Chisq) = d.ν > 2.0 ? d.ν - 2.0 : 0.0

function median(d::Chisq; approx::Bool=false)
    if approx
        return d.ν * (1.0 - 2.0 / (9.0 * d.ν))^3
    else
        return quantile(d, 0.5)
    end
end

function entropy(d::Chisq)
    hν = 0.5 * d.ν
    hν + logtwo + lgamma(hν) + (1.0 - hν) * digamma(hν)
end

#### Evaluation

@_delegate_statsfuns Chisq chisq ν

mgf(d::Chisq, t::Real) = (1.0 - 2.0 * t)^(-d.ν * 0.5)

cf(d::Chisq, t::Real) = (1.0 - 2.0 * im * t)^(-d.ν * 0.5)

gradlogpdf(d::Chisq, x::Float64) =  x > 0.0 ? (d.ν * 0.5 - 1) / x - 0.5 : 0.0

#### Sampling

_chisq_rand(ν::Float64) = StatsFuns.Rmath.chisqrand(ν)
rand(d::Chisq) = _chisq_rand(d.ν)


"""
    fixpoint(g,p0)

adalah fungsi untuk mencari solusi dari `x=g(x)` dengan nilai tebakan awal `p0`

# Examples
```jl
julia> g(x) = exp(-x);

julia> pn, flag, M = fixpoint(g,0.5);

julia> pn
0.5671432633594872

julia> flag
0

julia> M
27×3 Array{Float64,2}:
  0.0  0.5       NaN
  1.0  0.606531    0.106531
  2.0  0.545239    0.0612914
  3.0  0.579703    0.0344639
  ⋮
 23.0  0.567143    4.09741e-7
 24.0  0.567143    2.32382e-7
 25.0  0.567143    1.31794e-7
 26.0  0.567143    7.4746e-8
```
return solusi `pn` dengan `flag` bernilai 0 jika metode bisection berhasil menemukan
solusi dan gagal jika tidak 0. Serta, matriks `M` yang berisi catatan proses tiap
iterasi `[k, pk, f(pk)]`

"""
function fixpoint(g,p0)
    # Definisikan nilai toleransi, maksimum iterasi dan tebakan awal
    delta = 10^-7;
    maxi = 100;
    flag = 1;
    pn = p0
    M = [0 pn NaN];
    for n = 2:maxi
        pn1 = pn;
        pn = g(pn1);
        err = abs(pn-pn1);
        relerr = err/(abs(pn)+eps());
        M = [M; [n-1 pn err]]
        if (err<delta) || (relerr<delta)
            flag = 0; break
        end
    end
    return pn, flag, M
end


"""Julia program to implement Ternary Search algorithm.
Ternary Search is a Divide and Conquer strategy based algorithm similar to binary search.
Here we divide the entire array into three parts and recursively search for the required element
"""

function ternary_search_recursive(arr, low, high, ele)
    while (high >= low)
        # FInd the two required midpoints to divide the array into three.
        midA = low + (high - low) ÷ 3
        midB = high - (high - low) ÷ 3
        # Search, if the element is present in the midpoint indices
        if (arr[midA] == ele)
            return true
        elseif (arr[midB] == ele)
            return true
        end

        # Identify the part in which the required element belongs to and recursively call the function with limited length
        if (ele < arr[midA])
            return ternary_search_recursive(arr, low, midA - 1, ele)
        elseif (ele > arr[midB])
            return ternary_search_recursive(arr, midB + 1, high, ele)
        else
            return ternary_search_recursive(arr, midA + 1, midB - 1, ele)
        end
    end
    return false
end

print("How many numbers are present in the array? ")
n = readline()
n = parse(Int, n)
if (n <= 0)
    println("Array is Empty!!!")
    exit()
end
arr = Int[]
print("Enter the numbers: ")
arr = [parse(Int, num) for num in split(readline())]
print("Which number do you want to search in the array? ")
ele = readline()
ele = parse(Int, ele)
# Sort the array in ascending order
arr = sort(arr)
res = ternary_search_recursive(arr, 1, n, ele)
if (res == 0)
    print("The number $ele is not present in the array")
else
    print("The number $ele is present in the array.")
end

"""
Time Complexity - O(log(n)), where 'n' is the size of the array
Space Complexity - O(n)

SAMPLE INPUT AND OUTPUT

SAMPLE I

How many numbers are present in the array? 5
Enter the numbers: 1 2 3 4 5
Which number do you want to search in the array? 6
The number 6 is not present in the array

SAMPLE II

How many numbers are present in the array? 3
Enter the numbers: 3 1 2
Which number do you want to search in the array? 2
The number 2 is present in the array.

"""


# linreg.jl
#
# File with methods related to linear regression.
#
# Author: John Duchi (jduchi@stanford.edu)

# (A, b, x) = GenerateLinearRegressionData(m, n; condition_number, noise)
#
# Generates data for a linear regression problem
function GenerateLinearRegressionData(mm::Int64, nn::Int64;
                                      condition_number::Float64 = 1.0,
                                      noise::Float64 = 0.0)
  A = randn(mm, nn);
  x = randn(nn);
  if (condition_number > 1)
    D = LinRange(1, condition_number, nn);
    A = A * Diagonal(D);
  end
  b = A * x + noise * randn(mm);
  return (A, b, x);
end

# o = LinearRegressionObj(A::Matrix{Float64}, b::Vector{Float64},
#                         x::Vector{Float64})
#
# Computes and returns linear regression objective (mean squared error).
function LinearRegressionObj(A::Matrix{Float64}, b::Vector{Float64},
                             x::Vector{Float64})
  margins = A * x - b;
  return mean(margins.^2) / 2;
end

# x = ProximalUpdateLinreg(a::Vector{Float64}, b::Float64,
#                              x_init::Vector, stepsize::Float64)
#
# Sets x to minimize
#
#  .5 * (a' * x - b)^2 + norm(x - x_init)^2 / (2 * stepsize)
#
# Returns the minimizing x.
function ProximalUpdateLinreg(a::Vector{Float64}, b::Float64,
                              x_init::Vector, stepsize::Float64)
  # Replace this code to return the correct update.
  return x_init;
end

# x = SGUpdateLinreg(a::Vector{Float64}, b::Float64,
#                    x_init::Vector, stepsize::Float64)
#
# Sets x to minimize the linear approximation to the standard squared
# loss for linear regression, or, if g is the gradient of the loss
#
#   F(x; (a, b)) = .5 * (a' * x - b)^2
#
# at the point x_init, updates
#
#   x = x_init - stepsize * g;
function SGUpdateLinreg(a::Vector{Float64}, b::Float64,
                        x_init::Vector, stepsize::Float64)
  # Replace this code to return the correct update.
  return x_init;
end

# x = TruncatedUpdateLinreg(a::Vector{Float64}, b::Float64,
#                              x_init::Vector, stepsize::Float64)
#
# Let F(x; (a, b)) = .5 * (a' * x - b)^2. Sets x to minimize the
# positive approximation to F at x_init, that is, for
#
#  F_lin(x) = (F(x_init) + F'(x_init) * (x - x_init))_+
#
# sets x to minimize
#
#  F_lin + norm(x - x_init)^2 / (2 * stepsize).
#
# Returns the minimizing x.
function TruncatedUpdateLinreg(a::Vector{Float64}, b::Float64,
                               x_init::Vector, stepsize::Float64)
  # Replace this code to return the correct update.
  return x_init;
end

Machine Learning Summer School 2019 (London, UK)/tutorials/optimization/linreg.jl


#=
## Control System: Spaceship
As an example we want to build an IOSystem controlling the altitude of a spacecraft.
The spacecraft has mass ``m`` and can be controlled with thrusters which apply the force ``F(t)`` to the spacecraft. The altitude ``x(t)``
```math
\dot v(t) = \frac{ F(t) }{m}\\
\dot x(t) =  v(t)
```
in our model this system has the input ``F(t)``, the internal state ``v(t)`` (vertical velocity) and the output ``x(t)``.

```
       +------------+
F(t) --| spacecraft |-- x(t)
       | m, v(t)    |
       +------------+

```
=#

using BlockSystems
using ModelingToolkit
@parameters t M F(t)
@variables x(t) v(t)
D = Differential(t)

spacecraft = IOBlock([D(v) ~ F/M, D(x) ~ v], # define the equation
                     [F], # inputs of the system
                     [x], # outputs of the system
                     name = :spacecraft)

#=
We want to model a controller which takes a desired altitude as an input parameter and outputs the force for thrusters.

## Simple proportional controller
A proportional controller takes an input `i` and calculates the output proportional to the input.
```math
 o(t) = K\cdot i(t)
```
```
       +--------+
i(t) --| prop K |-- o(t)
       +--------+
```
=#
@parameters K i(t)
@variables o(t)

prop = IOBlock([o ~ K*i], [i], [o], name = :prop)
nothing # hide
#=
In order to make this useful as an controller, the input has to be the difference
between the reference and the system variable (negative feedback). We can model this
as an IOSystem where
```math
Δ = p - m\,.
```
```
       +--------+
p(t) --|  diff  |-- Δ(t)
m(t) --|        |
       +--------+
```
=#
@parameters p(t) m(t)
@variables Δ(t)
diff = IOBlock([Δ ~ p - m], [p, m], [Δ], name=:diff)
nothing # hide
#=
Now we can connect both of the defined models to create an proportional controller

```
             +----------------------------------------+
             | propc                                  |
             |         +--------+   +--------+        |
  target(t)--|--p(t) --|  diff  |---| prop K |--o(t)--|--o(t)
feedback(t)--|--m(t) --|        |   +--------+        |
             |         +--------+                     |
             +----------------------------------------+
```

If we don't provide additional information the system will try to promote all of the enclosed
variables to the new systemwide namespace.
=#
prop_c = IOSystem([diff.Δ => prop.i], # connect output of diff to input of prop
                  [diff, prop], # subsystems
                  name=:propc)

#=
For finer control, it is often preferred to give new names manually, this is
done with the `namespace_map` argument. Per default, all of the outputs of the
subsystems will become outputs of the connected system (in this case also the
output `diff.Δ`). We can prevent this by supplying the `outputs` argument
manually. Sub outputs which are not referenced here will become internal states
of the connected system.

The rhs of the namespace map can be given as a Variable/Parameter type from MTK.
For simple renaming one can also give the rhs as a `Symbol` type.
=#
prop_c = IOSystem([diff.Δ => prop.i], [diff, prop],
                  namespace_map = [prop.o => o,
                                   diff.p => :target,
                                   diff.m => :feedback],
                  outputs = [o],
                  name=:propc)

#=
Right now, the created object is a container for the two included systems. However,
it is possible to transform the object into a new `IOBlock` by calling the `connect_system`
function. The resulting is equivalent to
```
             +----------------------------------+
  target(t)--| prop_c_block                     |--o(t)
feedback(t)--| o(t)=K*(target(t) - feedback(t)) |
             +----------------------------------+
```
=#
prop_c_block = connect_system(prop_c)
nothing #hide

#=
Now we can hook our spaceship to this controller. It does not matter whether we use the
connected `IOBlock` version `prop_c` or the `IOSystem` version `prop_c_block`. We want to build
the connected system

```
           +--------------------------------------------+
           |  control system                            |
           |       +--------+   +------------+          |
target(t)--|-------| prop_c |---| spacecraft |-x(t)--+--|--altitude(t)
           |  +-fb-|        |   | m, v(t)    |       |  |
           |  |    +--------+   +------------+       |  |
           |  +--------------------------------------+  |
           +--------------------------------------------+
```
=#
@variables altitude(t)
space_controller = IOSystem([prop_c.o => spacecraft.F, spacecraft.x => prop_c.feedback],
                            [prop_c, spacecraft],
                            namespace_map = [spacecraft.x => altitude],
                            outputs = [altitude])
## we want to reduce the space_controller to a block
space_controller = connect_system(space_controller)
@info "Variables of space_controller" space_controller equations(space_controller.system)

#=
## Simulate System
In order to simulate the system we can have to build the Julia functions.
=#
gen = generate_io_function(space_controller)
nothing # hide
#=
By doing so we get access to a named tuple with the fileds
- `gen.f_ip` in-place function
- `gen.f_oop` out-of-place function
- `gen.massm` mass matrix of the system
- `gen.states` symbols of states (in order)
- `gen.inputs` symbols of inputs (in order)
- `gen.params` symbols of parameters (in order)
- (see docstring for full list)

The functions have the form
`f_ip(du, u, inputs, params, t)`
where `u` are all the states (outputs stacked on top of internal states) and `t` is the independent variable of the system.
The order of the inputs and states can be controlled.
=#
gen = generate_io_function(space_controller, f_states=[altitude, v], f_params=[K, M])
@info "Generated function" gen.massm gen.states gen.inputs gen.params
nothing # hide

#=
Well, let's see how our model is doing.
=#
using Plots
using OrdinaryDiffEq
targetfun(t) = t>1.0 ? 1.0 : 0
odefun(du, u, p, t) = gen.f_ip(du, u, [targetfun(t)], p, t)
p = [0.5, 1.0] # K, m
u0 = [0.0, 0.0] # altitude, v
tspan = (0.0, 30.0)
prob = ODEProblem(odefun, u0, tspan, p)
sol = solve(prob, Tsit5())
plot(t->sol(t)[1],tspan..., label="altitude", title="proportional control")
plot!(t->targetfun(t),tspan..., label="target")

#=
Well who could have thought, proportional control looks like an harmonic oscillator 🤷♂

## Defining a better controller
We might just add a damping term (a force proportional to the velocity of the spaceship).
If it works for a harmonic oscillator, it should work for our spaceship.
```
           +--------------------------------------------------------+
           |  control system                                        |
           |  +--------------------------------------------------+  |
           |  |    +--------+     +---+                          |  |
           |  +-v--| prop_v |-(-)-| d |     +------------+       |  |
           |       +--------+     | i |--F--| spacecraft |-v(t)--+  |
           |       +--------+     | f |     | m          |-x(t)--+--|--altitude(t)
target(t)--|-------| prop_c |-(+)-| f |     +------------+       |  |
           |  +-fb-|        |     +---+                          |  |
           |  |    +--------+                                    |  |
           |  +--------------------------------------------------+  |
           +--------------------------------------------------------+
```
In order to do so we have to slightly redefine the spaceship system: now the velocity `v(t)` is also an output and not and internal state.
=#

spacecraft = IOBlock([D(v) ~ F/M, D(x) ~ v],
                     [F],
                     [x,v],
                     name = :spacecraft)

# One can define new blocks based on previously defined blocks.
prop_v = IOBlock(prop, name=:prop_v)
fdiff = IOBlock(diff, name=:fdiff)

space_controller = IOSystem([spacecraft.v => prop_v.i,
                             spacecraft.x => prop_c.feedback,
                             prop_c.o => fdiff.p,
                             prop_v.o => fdiff.m,
                             fdiff.Δ => spacecraft.F],
                            [prop_v, prop_c, fdiff, spacecraft],
                            namespace_map = [spacecraft.x => altitude],
                            outputs = [altitude])

space_controller = connect_system(space_controller)
gen = generate_io_function(space_controller, f_states=[altitude, v], f_params=[prop_c.K, M, prop_v.K])

odefun(du, u, p, t) = gen.f_ip(du, u, [targetfun(t)], p, t)
p = [1.0, 1.0, 0.5] # propc, M, propv
u0 = [0.0, 0.0] # altitude, v
tspan = (0.0, 30.0)
prob = ODEProblem(odefun, u0, tspan, p)
sol = solve(prob, Tsit5())
plot(sol, vars=(0,1), label="altitude", title="better control")
plot!(t->targetfun(t),tspan..., label="target")

#=
## Defining an PT1 controller
```
             +-----------------------------------------------+
             | pi_c                                   +---+  |
             |                           +------------|   |  |
             |        +----+  +--------+ | +-------+  |sum|--|--o(t)
  target(t)--|--p(t)--|diff|--| prop K |-+-| int T |--|   |  |
feedback(t)--|--m(t)--|    |  +--------+   +-------+  +---+  |
             |        +----+                                 |
             +-----------------------------------------------+
```
=#

@parameters T a1(t) a2(t)
@variables Σ(t) altitude(t)
int = IOBlock([D(o) ~ 1/T * i - o], [i], [o], name=:int)
adder = IOBlock([Σ ~ a1 + a2], [a1, a2], [Σ], name=:add)

pi_c = IOSystem([diff.Δ => prop.i,
                 prop.o => int.i,
                 prop.o => adder.a1,
                 int.o => adder.a2],
                [diff, prop, int, adder],
                namespace_map = [diff.p => :target,
                                 diff.m => :feedback,
                                 adder.Σ => o],
                outputs = [o],
                name=:pi_c)
nothing # hide

# as before we can close the loop and build the control circuit

space_controller = IOSystem([pi_c.o => spacecraft.F , spacecraft.x => pi_c.feedback],
                            [pi_c, spacecraft],
                            namespace_map = [spacecraft.x => altitude],
                            outputs = [altitude])
space_controller = connect_system(space_controller, verbose=false)
@info "Variables of space_controller" space_controller equations(space_controller.system)

# and we can simulate and plot the system
gen = generate_io_function(space_controller, f_states=[altitude], f_params=[K, T, M])

odefun(du, u, p, t) = gen.f_ip(du, u, [targetfun(t)], p, t)
p = [0.5, -1.5, 1.0] # K, T, m
u0 = [0.0, 0.0, 0.0] # altitude, int.o, v
tspan = (0.0, 50.0)
prob = ODEProblem(odefun, u0, tspan, p)
sol = solve(prob, Tsit5())
plot(sol, vars=(0,[ 1,2 ]), label=["altitude" "integrator"], title="PT1 controller")
plot!(t->targetfun(t),tspan..., label="target")

# thank you for flying with us :)


# # Linear regression

import MaximumLikelihoodProblems

import Distributions
import ForwardDiff
import LogDensityProblems
import TransformVariables

struct LinearRegression{Ty, TX}
    y::Ty
    X::TX
end

function (problem::LinearRegression)(θ)
    y = problem.y
    X = problem.X

    β = θ.β
    σ = θ.σ

    η = X * β
    μ = η
    ε = y - μ

    ## these two lines are equivalent:
    ## log_likelihood = Distributions.loglikelihood(Distributions.Normal(0, σ), ε)
    ## log_likelihood = sum(Distributions.logpdf.(Distributions.Normal(0, σ), ε))

    log_likelihood = sum(Distributions.logpdf.(Distributions.Normal(0, σ), ε))
    return log_likelihood
end

N = 10_000

## the first column (the column of all ones) is the intercept
X = hcat(ones(N), randn(N), randn(N))

size_β = (3,)
β_true = [1.0, 2.0, -1.0]
σ_true = 0.5
η_true = X * β_true
μ_true = η_true
ε_true = randn(N) .* σ_true
y = μ_true + ε_true

function generate_problem_transformation(p::LinearRegression)
    return TransformVariables.as((β = TransformVariables.as(Array, size(p.X, 2)),
                                  σ = TransformVariables.asℝ₊))
end

problem = LinearRegression(y, X)
transformation = generate_problem_transformation(problem)
transformed_problem = LogDensityProblems.TransformedLogDensity(transformation,
                                                               problem)
transformed_gradient_problem = LogDensityProblems.ADgradient(:ForwardDiff,
                                                             transformed_problem)

β_hat_initial_guess = zeros(size_β)
σ_hat_initial_guess = 1.0

θ_hat_initial = (; β = β_hat_initial_guess,
                   σ = σ_hat_initial_guess)

# θ_hat:

θ_hat = MaximumLikelihoodProblems.fit(transformed_gradient_problem,
                                      θ_hat_initial;
                                      learning_rate = 1e-6)

# β_hat:

β_hat = θ_hat[:β]

# σ_hat:

σ_hat = θ_hat[:σ]

# Value of the log likelihood function evaluated at θ_hat:

MaximumLikelihoodProblems.loglikelihood(transformed_gradient_problem, θ_hat)


"""
    update(x, measurement, H)

UPDATE step in Kalman filtering for linear dynamics models:
```math
K = P_{n+1}^P * H^T * S^{-1}
m_{n+1} = m_{n+1}^P + K * (0 - z)
P_{n+1} = P_{n+1}^P - K*S*K^T
```

This function provides a very simple UPDATE implementation.
In the solvers, we recommend to use the non-allocating [`update!`](@ref).
"""
function update(x::Gaussian, measurement::Gaussian, H::AbstractMatrix)
    m, C = x
    z, S = measurement

    K = C * H' * inv(S)
    m_new = m - K * z
    C_new = C - K * S * K'

    return Gaussian(m_new, C_new)
end
"""UPDATE step in Joseph-form, with square-root matrix inputs"""
function update(x::SRGaussian, measurement::SRGaussian, H::AbstractMatrix)
    m, C = x
    z, S = measurement

    K = C * H' * inv(S)
    m_new = m - K * z
    C_new = X_A_Xt(C, (I - K * H))

    return Gaussian(m_new, C_new)
end

"""
    update!(x_out, x_pred, measurement, H, R=0)

UPDATE step in Kalman filtering for linear dynamics models, given a measurement `Z=N(z, S)`.
In-place implementation of [`update`](@ref), saving the result in `x_out`.

```math
K = P_{n+1}^P * H^T * S^{-1}
m_{n+1} = m_{n+1}^P + K * (0 - z)
P_{n+1} = P_{n+1}^P - K*S*K^T
```

Implemented in Joseph Form.

See also: [`predict`](@ref)
"""
function update!(
    x_out::Gaussian,
    x_pred::Gaussian,
    measurement::Gaussian,
    H::AbstractMatrix,
    K_cache::AbstractMatrix,
    M_cache::AbstractMatrix,
    m_tmp,
)
    z, S = measurement.μ, copy!(m_tmp.Σ, measurement.Σ)
    m_p, P_p = x_pred.μ, x_pred.Σ
    D = length(m_p)

    # K = P_p * H' / S
    S_chol = cholesky!(S)
    K = _matmul!(K_cache, Matrix(P_p), H')
    rdiv!(K, S_chol)

    # x_out.μ .= m_p .+ K * (0 .- z)
    x_out.μ .= m_p .- _matmul!(x_out.μ, K, z)

    # M_cache .= I(D) .- mul!(M_cache, K, H)
    _matmul!(M_cache, K, H, -1.0, 0.0)
    @inbounds @simd ivdep for i in 1:D
        M_cache[i, i] += 1
    end

    X_A_Xt!(x_out.Σ, P_p, M_cache)

    return x_out
end


### A Pluto.jl notebook ###
# v0.12.4

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 10b19f7a-1f8b-11eb-128e-11743ce72ef5
using PlutoUI, Plots

# ╔═╡ e5a688ac-1f91-11eb-104a-2753c9263f14
md"""
 # Floating Point Arithmetic
"""

# ╔═╡ 2df31082-1f8b-11eb-145a-93fa6e01e345
@bind f Slider(1:100)

# ╔═╡ 37c0afe8-1f8b-11eb-1154-31f5954ca4ff
f

# ╔═╡ 8d2203e8-1f8a-11eb-0332-7d512d831ec1
bitstring(f)

# ╔═╡ 4aa7bc3c-1f8b-11eb-30dd-89d64f42fc27
length(bitstring(f))

# ╔═╡ c90f7b06-1f8a-11eb-17d7-35a09af858f1
float(f)

# ╔═╡ d4cd33f2-1f8a-11eb-1af7-bdbc1fb236da
bitstring(float(f))

# ╔═╡ 4f354440-1f8b-11eb-225a-f55552bc9e29
length(bitstring(float(f)))

# ╔═╡ b7eb6210-1f9e-11eb-2a57-95b4fab4187f
md"""
IEEE 64 bit representation (double precision) (64 bits0
"""

# ╔═╡ 2ee2a8ec-1f8e-11eb-23e6-9d5e2bb00d1e
ieee0(x::Float64) = [bitstring(x)[1:1], bitstring(x)[2:12], bitstring(x)[13:64]]

# ╔═╡ d6b5d1f6-1f9e-11eb-0c32-9b9cb51352a0
md"""
## Let me work with numbers from 1 to 2
"""

# ╔═╡ e0cd119a-1f9e-11eb-1181-3d3fa6e0653f
 nextfloat(1.0)-1.0

# ╔═╡ f3357c5a-1f9e-11eb-270d-d3b5fd6bd542
2.0 ^(-52)

# ╔═╡ fbe68df8-1f9e-11eb-1278-a102df99c88c

# ╔═╡ 0bf5f86a-1f8f-11eb-17fd-75f4b45f16d5
@bind j Slider(1: 2^(52), show_value = true)

# ╔═╡ 2ff48c68-1f8f-11eb-0bd0-cde93f2cc18b
g = 1 + j/2^52

# ╔═╡ 316b7dd0-1f8e-11eb-0da9-9fc5b63a4e80
ieee0(g)

# ╔═╡ 56e8258a-1f8f-11eb-3249-6592feed84b3
ieee0(1.0)

# ╔═╡ 5d7dcbe0-1f8f-11eb-05e9-1b2ae35cf681
ieee0(-1.0)

# ╔═╡ bf3bd13c-1f95-11eb-3187-a7311125be18
nextfloat(1.0)-1

# ╔═╡ d3a6bbaa-1f95-11eb-27de-e99b0e230d21
eps(1.0)

# ╔═╡ c6d813ba-1f95-11eb-1be6-55cb3ffeb358
nextfloat(1.0f0)-1

# ╔═╡ d66f2a16-1f95-11eb-3edc-4182c1a4b8ad
eps(1.0f0)

# ╔═╡ d31baed4-1f95-11eb-19df-15ff5c378487
eps(1.0f0) * 2^23

# ╔═╡ 81a7921c-1f8f-11eb-33c6-397d38b57a75
ieee(x) = [ parse.(Int,ieee0(x),base=2)      ieee0(x)]

# ╔═╡ 85c8f822-1f8f-11eb-15a1-7956d269779a
ieee(g)

# ╔═╡ dd3031d8-1f9f-11eb-39b2-cb7ac57c6bdf
ieee(1.0)

# ╔═╡ e4b2733a-1f9f-11eb-30ac-0dc5305b8938
ieee(-1.0)

# ╔═╡ ff15233a-1f9f-11eb-0012-f198e05abaf9
md"""
# s, e, m  (sign, exponent, mantissa)  =  (-1)^s  * 2^(e-1023) * (1 + m/2^52)
"""

# ╔═╡ 3f36e4a8-1fa0-11eb-1c46-97ff95905b41
ieee(prevfloat(1.0))

# ╔═╡ f8612fa6-1f91-11eb-36d6-cbd0f9b5b0ea
float2ieee(x::Float64) =  parse.(Int,ieee0(x),base=2)

# ╔═╡ 66f85a94-1fa0-11eb-32bd-178184d8aef9
2^52-1

# ╔═╡ 0a44c782-1f92-11eb-1b67-eb93f1dfe1ff
float2ieee(1.0)

# ╔═╡ 37039e06-1f92-11eb-192e-e59ffb46a368
begin
	i =-2:2
	x =  2.0 .^i
	plot(yticks=[],xticks=round.(x,digits=2),size=(500,100))
	for xx∈x

		z = (1+1/16):(1/16):(2-1/16)
		for zz∈ xx.*z

	  plot!( [zz, zz], [-.1, 0.0], color=:red, ratio=10, legend=false, ylims=(-.1,.1),tickfontsize=1)

		end
		  plot!( [xx, xx], [-.1, .1], color=:blue, ratio=5, legend=false, ylims=(-.1,.1),tickfontsize=4)
	end

	plot!()
end

# ╔═╡ c51ad4de-1fa1-11eb-3e01-fb46ed72f493
(1 + 1.000001*eps(1.0)/2)-1

# ╔═╡ 79e5e5ca-1fa0-11eb-087f-9fad526c1b0e
ieee(0.0) ## does not follow the same rule

# ╔═╡ 8fd89a82-1fa0-11eb-2921-0d3c31c93cbc
ieee(-0.0) ## is this the same as 0.0?

# ╔═╡ a207d416-1fa0-11eb-3bf0-bdc40925736e
1.0 / 0.0

# ╔═╡ c7473096-1fa0-11eb-39fd-31c4c4b1eb04
ieee(Inf)

# ╔═╡ d08a02dc-1fa0-11eb-3369-a7c84dad581a
ieee(1.0/-0.0)

# ╔═╡ f072cf66-1fa0-11eb-0789-733f904d41d0
0.0 / 0.0

# ╔═╡ ff8975cc-1fa0-11eb-1fa6-514db7898e52
md"""
# Not a Number
"""

# ╔═╡ 0b340f86-1fa1-11eb-37bd-839bd7525dc9
NaN == NaN

# ╔═╡ 0d1c1902-1f96-11eb-0734-1bb3dda2aa7d
md"""
Why do we need this floating point representation?
"""

# ╔═╡ a40c4352-1f94-11eb-263e-632828cb6079
float2ieee(0.0)

# ╔═╡ b7b6a198-1f94-11eb-0093-8f3ab007d770
float2ieee(-0.0)

# ╔═╡ bb04eb2c-1f94-11eb-0a4a-3d2765531e94
float2ieee(Inf)

# ╔═╡ c0e4f8f2-1f94-11eb-1305-d9028395ab2f
float2ieee(-Inf)

# ╔═╡ c500ec66-1f94-11eb-079a-1b5b4befade2
1/-0.0

# ╔═╡ 1e37c20a-1f95-11eb-37f9-11c5a6041e1b
Inf - Inf

# ╔═╡ 25be9fbc-1f95-11eb-294a-1dd62a135bfa
0.0/0.0

# ╔═╡ 7345c08e-1f96-11eb-0bf1-c50d0108facd
md"""
$$\begin{matrix}   \text{exponent\\mantissa} & 0 & ≠0  \\ \hline 0 & ±0  &  \text{denormal} \\
\text{regular} &  2^k & \text{general} \\ \text{max} &  \pm\infty &  \text{NaN} \end{matrix}$$
"""

# ╔═╡ b22dc266-1f9a-11eb-199b-170ef50b702c
md"""
Rules of IEEE arithmetic
"""

# ╔═╡ 3f2943c6-1f95-11eb-0342-6569169f32ca
Float32(1.0)

# ╔═╡ 47153ebe-1f95-11eb-2224-716d2e152db0
bitstring(Float32(1.0))

# ╔═╡ 60135b2e-1f95-11eb-293d-4f8a96a677b2
float32ieee0(x::Float32) = [bitstring(x)[1:1], bitstring(x)[2:9], bitstring(x)[10:32]]

# ╔═╡ 77de76f8-1f95-11eb-00bc-915dd7977c75
float32ieee0(1.0f0)

# ╔═╡ 65ef756c-1f96-11eb-2546-35b37cb7f54c
## The rules of +, -, / , * and √

# ╔═╡ f707de2e-1fa1-11eb-0147-3fca8c03ab5f

# ╔═╡ 2b0265d4-1f96-11eb-1553-ef2ecb7d451b
md"""
What precision do you need?
 - Small absolute precision?
 - Small relative precision?
 - Neither?
"""

# ╔═╡ c5f5c97c-1f97-11eb-3683-55c16be8274e
md"""
What is displayed? What can you represent?
"""

# ╔═╡ b86c366c-1f95-11eb-0a18-a52cdf00da4c
1/5

# ╔═╡ d400e394-1f97-11eb-0f74-476c32bdd9d0
bitstring(1/5)

# ╔═╡ f5e3b162-1f97-11eb-2879-f7c7954143a0
big(1/5)

# ╔═╡ 000c656c-1f98-11eb-0f0e-657b9651f0bd
big(1)/5

# ╔═╡ 149b4c9e-1f98-11eb-19b7-eb9988c9333a
BigFloat(1//5,precision=1000)

# ╔═╡ 3aae6bac-1f98-11eb-22a5-bdc5e5de3f4c
BigFloat(π,precision=1000)

# ╔═╡ 43bed31c-1f98-11eb-1131-f5dd8dd27086
BigFloat(float(pi),precision=1000)

# ╔═╡ 56fb9c94-1f98-11eb-216e-f9e6e04d6f45
md"""
Note that every IEEE number terminates (is not a repeating expansion)
"""

# ╔═╡ 32cf8afa-1f99-11eb-0193-e969cca12f86
md"""
Find the first floating point number bigger than 1 such that x * (1/x) ≠ 1
"""

# ╔═╡ 51baa858-1f99-11eb-3329-4d9bf850b1c0
begin
	t = 1+rand()
	t,ieee(t*1/t)
end

# ╔═╡ 64c9235e-1f99-11eb-24f3-d940be06c7c8
let
	 α = 1.0
     while( α * (1/α) == 1.0)
		α = nextfloat(α)
	 end
	 Int((α-1)/eps(1.0))
end

# ╔═╡ ecc5c55a-1f99-11eb-3f44-4161850426cc
α= 1 + 257736490*eps()

# ╔═╡ 40575ec2-1f9a-11eb-2d29-5b517d1f3b18
α*(1/α)

# ╔═╡ 0b0719c2-1f9c-11eb-3ed6-6927c927145d

# ╔═╡ 11ebb920-1f9c-11eb-061a-8df3899e8945

# ╔═╡ Cell order:
# ╟─e5a688ac-1f91-11eb-104a-2753c9263f14
# ╠═2df31082-1f8b-11eb-145a-93fa6e01e345
# ╠═37c0afe8-1f8b-11eb-1154-31f5954ca4ff
# ╠═8d2203e8-1f8a-11eb-0332-7d512d831ec1
# ╠═4aa7bc3c-1f8b-11eb-30dd-89d64f42fc27
# ╠═c90f7b06-1f8a-11eb-17d7-35a09af858f1
# ╠═d4cd33f2-1f8a-11eb-1af7-bdbc1fb236da
# ╠═4f354440-1f8b-11eb-225a-f55552bc9e29
# ╟─b7eb6210-1f9e-11eb-2a57-95b4fab4187f
# ╠═2ee2a8ec-1f8e-11eb-23e6-9d5e2bb00d1e
# ╟─d6b5d1f6-1f9e-11eb-0c32-9b9cb51352a0
# ╠═e0cd119a-1f9e-11eb-1181-3d3fa6e0653f
# ╠═f3357c5a-1f9e-11eb-270d-d3b5fd6bd542
# ╠═fbe68df8-1f9e-11eb-1278-a102df99c88c
# ╠═0bf5f86a-1f8f-11eb-17fd-75f4b45f16d5
# ╠═2ff48c68-1f8f-11eb-0bd0-cde93f2cc18b
# ╠═316b7dd0-1f8e-11eb-0da9-9fc5b63a4e80
# ╠═56e8258a-1f8f-11eb-3249-6592feed84b3
# ╠═5d7dcbe0-1f8f-11eb-05e9-1b2ae35cf681
# ╠═bf3bd13c-1f95-11eb-3187-a7311125be18
# ╠═d3a6bbaa-1f95-11eb-27de-e99b0e230d21
# ╠═c6d813ba-1f95-11eb-1be6-55cb3ffeb358
# ╠═d66f2a16-1f95-11eb-3edc-4182c1a4b8ad
# ╠═d31baed4-1f95-11eb-19df-15ff5c378487
# ╠═81a7921c-1f8f-11eb-33c6-397d38b57a75
# ╠═85c8f822-1f8f-11eb-15a1-7956d269779a
# ╠═dd3031d8-1f9f-11eb-39b2-cb7ac57c6bdf
# ╠═e4b2733a-1f9f-11eb-30ac-0dc5305b8938
# ╠═ff15233a-1f9f-11eb-0012-f198e05abaf9
# ╠═3f36e4a8-1fa0-11eb-1c46-97ff95905b41
# ╠═f8612fa6-1f91-11eb-36d6-cbd0f9b5b0ea
# ╠═66f85a94-1fa0-11eb-32bd-178184d8aef9
# ╠═0a44c782-1f92-11eb-1b67-eb93f1dfe1ff
# ╠═37039e06-1f92-11eb-192e-e59ffb46a368
# ╠═c51ad4de-1fa1-11eb-3e01-fb46ed72f493
# ╠═79e5e5ca-1fa0-11eb-087f-9fad526c1b0e
# ╠═8fd89a82-1fa0-11eb-2921-0d3c31c93cbc
# ╠═a207d416-1fa0-11eb-3bf0-bdc40925736e
# ╠═c7473096-1fa0-11eb-39fd-31c4c4b1eb04
# ╠═d08a02dc-1fa0-11eb-3369-a7c84dad581a
# ╠═f072cf66-1fa0-11eb-0789-733f904d41d0
# ╟─ff8975cc-1fa0-11eb-1fa6-514db7898e52
# ╠═0b340f86-1fa1-11eb-37bd-839bd7525dc9
# ╟─0d1c1902-1f96-11eb-0734-1bb3dda2aa7d
# ╠═a40c4352-1f94-11eb-263e-632828cb6079
# ╠═b7b6a198-1f94-11eb-0093-8f3ab007d770
# ╠═bb04eb2c-1f94-11eb-0a4a-3d2765531e94
# ╠═c0e4f8f2-1f94-11eb-1305-d9028395ab2f
# ╠═c500ec66-1f94-11eb-079a-1b5b4befade2
# ╠═1e37c20a-1f95-11eb-37f9-11c5a6041e1b
# ╠═25be9fbc-1f95-11eb-294a-1dd62a135bfa
# ╠═7345c08e-1f96-11eb-0bf1-c50d0108facd
# ╟─b22dc266-1f9a-11eb-199b-170ef50b702c
# ╠═3f2943c6-1f95-11eb-0342-6569169f32ca
# ╠═47153ebe-1f95-11eb-2224-716d2e152db0
# ╠═60135b2e-1f95-11eb-293d-4f8a96a677b2
# ╠═77de76f8-1f95-11eb-00bc-915dd7977c75
# ╠═65ef756c-1f96-11eb-2546-35b37cb7f54c
# ╠═f707de2e-1fa1-11eb-0147-3fca8c03ab5f
# ╠═2b0265d4-1f96-11eb-1553-ef2ecb7d451b
# ╟─c5f5c97c-1f97-11eb-3683-55c16be8274e
# ╠═b86c366c-1f95-11eb-0a18-a52cdf00da4c
# ╠═d400e394-1f97-11eb-0f74-476c32bdd9d0
# ╠═f5e3b162-1f97-11eb-2879-f7c7954143a0
# ╠═000c656c-1f98-11eb-0f0e-657b9651f0bd
# ╠═149b4c9e-1f98-11eb-19b7-eb9988c9333a
# ╠═3aae6bac-1f98-11eb-22a5-bdc5e5de3f4c
# ╠═43bed31c-1f98-11eb-1131-f5dd8dd27086
# ╟─56fb9c94-1f98-11eb-216e-f9e6e04d6f45
# ╟─32cf8afa-1f99-11eb-0193-e969cca12f86
# ╠═51baa858-1f99-11eb-3329-4d9bf850b1c0
# ╠═64c9235e-1f99-11eb-24f3-d940be06c7c8
# ╠═ecc5c55a-1f99-11eb-3f44-4161850426cc
# ╠═40575ec2-1f9a-11eb-2d29-5b517d1f3b18
# ╠═0b0719c2-1f9c-11eb-3ed6-6927c927145d
# ╠═11ebb920-1f9c-11eb-061a-8df3899e8945
# ╠═10b19f7a-1f8b-11eb-128e-11743ce72ef5


#############
# This demonstrates the Padua transform and inverse transform,
# explaining precisely the normalization and points
#############

using FastTransforms

N = 15
pts = paduapoints(N)
x = pts[:,1]; y = pts[:,2]

f = (x,y) -> exp(x + cos(y))
f̌ = paduatransform(f.(x , y))
f̃ = (x,y) -> begin
    j = 1
    ret = 0.0
    for n in 0:N, k in 0:n
        ret += f̌[j]*cos((n-k)*acos(x)) * cos(k*acos(y))
        j += 1
    end
    ret
end

f̃(0.1,0.2) ≈ f(0.1,0.2)

ipaduatransform(f̌) ≈ f̃.(x,y)


function fractran(n::Integer, ratios::Vector{<:Rational}, steplim::Integer)
    rst = zeros(BigInt, steplim)
    for i in 1:steplim
        rst[i] = n
        if (pos = findfirst(x -> isinteger(n * x), ratios)) > 0
            n *= ratios[pos]
        else
            break
        end
    end
    return rst
end

using IterTools
macro ratio_str(s)
    a = split(s, r"[\s,/]+")
    return collect(parse(BigInt, n) // parse(BigInt, d) for (n, d) in partition(a, 2))
end

fracs = ratio"""17 / 91, 78 / 85, 19 / 51, 23 / 38, 29 / 33, 77 / 29, 95 / 23,
                77 / 19, 1 / 17, 11 / 13, 13 / 11, 15 / 14, 15 / 2, 55 / 1"""
println("The first 20 in the series are ", fractran(2, fracs, 20))

prmfound = 0
n = big(2)
while prmfound < 20
    if isinteger(log2(n))
        prmfound += 1
        println("Prime $prmfound found: $n is 2 ^ $(Int(log2(n)))")
    end
    n = fractran(n, fracs, 2)[2]
end


using Statistics;
using Random;

function distance(from::Vector{T}, to::Vector{V}) where {T,V}
    sum((from .- to) .^2) |> sqrt
end

function KMEANS(
    matrix::Vector{Vector{T}},
    k::Int;
    max::Int,
)::Vector{Vector{Vector{T}}} where {T}
    len = matrix |> length
    μ = shuffle(MersenneTwister(888888), 1:len)[1:k] .|> x -> matrix[x]
    C::Vector{Vector{Vector{T}}} = 1:k .|> _ -> []
    for _ ∈ 1:max
        for x ∈ matrix
            distances = μ .|> ϵ -> distance(x, ϵ)
            min = argmin(distances)
            C[min] = C[min] ∪ [x]
        end
        δ = C .|> x -> mean(x)
        μ == δ && break
        μ = δ
    end
    C
end

KMEANS(
    [
        [0, 2],
        [0, 0],
        [1.5, 0],
        [5, 0],
        [5, 2]
    ],
    2,
    max = 20,
) .|> println


@doc raw"""
    gausschebyshev(n::Integer) -> Tuple{Vector{Float64},Vector{Float64}}
    gausschebyshev(n::Integer, 1) -> Tuple{Vector{Float64},Vector{Float64}}

Return nodes and weights of [Gauss-Chebyshev quadrature](https://en.wikipedia.org/wiki/Chebyshev%E2%80%93Gauss_quadrature) of the 1st kind.

```math
\int_{-1}^{1} \frac{f(x)}{\sqrt{1-x^2}} dx \approx \sum_{i=1}^{n} w_i f(x_i)
```

# Examples
```jldoctest
julia> x, w = gausschebyshev(3);

julia> f(x) = x^4;

julia> I = dot(w, f.(x));

julia> I ≈ 3π/8
true
```

---

    gausschebyshev(n::Integer, 2) -> Tuple{Vector{Float64},Vector{Float64}}

Return nodes and weights of [Gauss-Chebyshev quadrature](https://en.wikipedia.org/wiki/Chebyshev%E2%80%93Gauss_quadrature) of the 2nd kind.

```math
\int_{-1}^{1} f(x)\sqrt{1-x^2} dx \approx \sum_{i=1}^{n} w_i f(x_i)
```

# Examples
```jldoctest
julia> x, w = gausschebyshev(3, 2);

julia> f(x) = x^4;

julia> I = dot(w, f.(x));

julia> I ≈ π/16
true
```

---

    gausschebyshev(n::Integer, 3) -> Tuple{Vector{Float64},Vector{Float64}}

Return nodes and weights of [Gauss-Chebyshev quadrature](https://en.wikipedia.org/wiki/Chebyshev%E2%80%93Gauss_quadrature) of the 3rd kind.

```math
\int_{-1}^{1} f(x)\sqrt{1-x^2} dx \approx \sum_{i=1}^{n} w_i f(x_i)
```

# Examples
```jldoctest
julia> x, w = gausschebyshev(3, 3);

julia> f(x) = x^4;

julia> I = dot(w, f.(x));

julia> I ≈ 3π/8
true
```

---

    gausschebyshev(n::Integer, 4) -> Tuple{Vector{Float64},Vector{Float64}}

Return nodes and weights of [Gauss-Chebyshev quadrature](https://en.wikipedia.org/wiki/Chebyshev%E2%80%93Gauss_quadrature) of the 4th kind.

```math
\int_{-1}^{1} f(x)\sqrt{1-x^2} dx \approx \sum_{i=1}^{n} w_i f(x_i)
```

# Examples
```jldoctest
julia> x, w = gausschebyshev(3, 4);

julia> f(x) = x^4;

julia> I = dot(w, f.(x));

julia> I ≈ 3π/8
true
```
"""
function gausschebyshev(n::Integer, kind::Integer=1)
    # GAUSS-CHEBYSHEV NODES AND WEIGHTS.

    if n < 0
        throw(DomainError(n, "Input n must be a non-negative integer"))
    end

    # Use known explicit formulas. Complexity O(n).
    if kind == 1
        # Gauss-ChebyshevT quadrature, i.e., w(x) = 1/sqrt(1-x^2)
        return ([cos((2 * k - 1) * π / (2 * n)) for k = n:-1:1], fill(π / n, n))
    elseif kind == 2
        # Gauss-ChebyshevU quadrature, i.e., w(x) = sqrt(1-x^2)
        return ([cos(k * π / (n + 1)) for k = n:-1:1],
                [π/(n + 1) * sin(k / (n + 1) * π)^2 for k = n:-1:1])
    elseif kind == 3
        # Gauss-ChebyshevV quadrature, i.e., w(x) = sqrt((1+x)/(1-x))
        return ([cos((k - .5) * π / (n + .5)) for k = n:-1:1],
                [2π / (n + .5) * cos((k - .5) * π / (2 * (n + .5)))^2 for k = n:-1:1])
    elseif kind == 4
        # Gauss-ChebyshevW quadrature, i.e., w(x) = sqrt((1-x)/(1+x))
        return ([cos(k * π / (n + .5)) for k = n:-1:1],
                [2π / (n + .5) * sin(k * π / (2 * (n + .5)))^2 for k = n:-1:1])
    else
        throw(ArgumentError("Chebyshev kind should be 1, 2, 3, or 4"))
    end
end


# # [Stochastic Gradient Descent](@id SGDTutorial)
#
# This tutorial illustrates how to use the [`stochastic_gradient_descent`](@ref)
# solver and different [`DirectionUpdateRule`](@ref)s in order to introduce
# the average or momentum variant, see [Stochastic Gradient Descent](https://en.wikipedia.org/wiki/Stochastic_gradient_descent).
#
# Computationally we look at a very simple but large scale problem,
# the Riemannian Center of Mass or [Fréchet mean](https://en.wikipedia.org/wiki/Fréchet_mean):
# For given points ``p_i ∈\mathcal M``, ``i=1,…,N`` this optimization problem reads
#
# ```math
# \operatorname*{arg\,min}_{x∈\mathcal M} \frac{1}{2}\sum_{i=1}^{N}
#   \operatorname{d}^2_{\mathcal M}(x,p_i),
# ```
# which of course can be (and is) solved by a gradient descent, see the [introductionary tutorial](@ref Optimize).
# If ``N`` is very large it might be quite expensive to evaluate the complete gradient.
# A remedy is, to evaluate only one of the terms at a time and choose a random order for these.
#
# We first initialize the manifold (see [])
export_folder = joinpath( #src
    @__DIR__, #src
    "..", #src
    "..", #src
    "docs", #src
    "src", #src
    "assets", #src
    "images", #src
    "tutorials", #src
) #src
using Manopt, Manifolds, Random, Colors
# and we define some colors from [Paul Tol](https://personal.sron.nl/~pault/)
black = RGBA{Float64}(colorant"#000000")
TolVibrantOrange = RGBA{Float64}(colorant"#EE7733") # Start
TolVibrantBlue = RGBA{Float64}(colorant"#0077BB") # a path
TolVibrantTeal = RGBA{Float64}(colorant"#009988") # points
nothing #hide
#
# And optain a large data set
n = 5000
σ = π / 12
M = Sphere(2)
x = 1 / sqrt(2) * [1.0, 0.0, 1.0]
Random.seed!(42)
data = [exp(M, x, random_tangent(M, x, Val(:Gaussian), σ)) for i in 1:n]
nothing #hide
# which looks like
#
asymptote_export_S2_signals( #src
    export_folder * "/centerAndLargeData.asy"; #src
    points=[[x], data], #src
    colors=Dict(:points => [TolVibrantOrange, TolVibrantTeal]), #src
    dot_sizes=[2.5, 1.25], #src
    camera_position=(1.0, 0.5, 0.5), #src
) #src
render_asymptote(export_folder * "/centerAndLargeData.asy"; render=2) #src
#md # ```julia
#md # asymptote_export_S2_signals("centerAndLargeData.asy";
#md #     points = [ [x], data],
#md #     colors=Dict(:points => [TolVibrantBlue, TolVibrantTeal]),
#md #     dot_sizes = [2.5, 1.0], camera_position = (1.,.5,.5)
#md # )
#md # render_asymptote("centerAndLargeData.asy"; render = 2)
#md # ```
#md #
#md # ![The data of noisy versions of $x$](../assets/images/tutorials/centerAndLargeData.png)
#
# Note that due to the construction of the points as zero mean tangent vectors, the mean should
# be very close to our initial point `x`.
#
# In order to use the stochastic gradient, we now need a function that returns the vector of gradients.
# There are two ways to define it in `Manopt.jl`: as one function, that returns a vector or a vector of funtions.
#
# The first variant is of course easier to define, but the second is more efficient when only evaluating one of the gradients.
# For the mean we have as a gradient
#
# ```math
#  gradF(x) = \sum_{i=1}^N \operatorname{grad}f_i(x) \quad \text{where} \operatorname{grad}f_i(x) = -\log_x p_i
# ```
#
# Which we define as
F(M, x) = 1 / (2 * n) * sum(map(p -> distance(M, x, p)^2, data))
gradF(M, x) = [grad_distance(M, p, x) for p in data]
gradf = [(M, x) -> grad_distance(M, p, x) for p in data];
# The calls are only slightly different, but notice that accessing the 2nd gradient element
# requires evaluating all logs in the first function.
# So while you can use both `gradF` and `gradf` in the following call, the second one is faster:
@time x_opt1 = stochastic_gradient_descent(M, gradF, x);
# versus
@time x_opt2 = stochastic_gradient_descent(M, gradf, x);
# This result is reasonably close. But we can improve it by using a [`DirectionUpdateRule`](@ref),
# namely:
# On the one hand [`MomentumGradient`](@ref), which requires both the manifold and the initial value,
#    in order to keep track of the iterate and parallel transport the last direction to the current iterate.
#    you can also set a `vector_transport_method`, if `ParallelTransport()` is not
#    available on your manifold. Here we simply do
@time x_opt3 = stochastic_gradient_descent(
    M, gradf, x; direction=MomentumGradient(M, x, StochasticGradient(zero_vector(M, x)))
);
# And on the other hand the [`AverageGradient`](@ref) computes an average of the last `n` gradients, i.e.
@time x_opt4 = stochastic_gradient_descent(
    M, gradf, x; direction=AverageGradient(M, x, 10, StochasticGradient(zero_vector(M, x)))
);
# note that the default [`StoppingCriterion`](@ref) is a fixed number of iterations.
#
# Note that since you can apply both also in case of [`gradient_descent`](@ref),
# i.e. to use [`IdentityUpdateRule`](@ref) and evaluate the classical gradient,
# both constructors have to know that internally the default evaluation of the Stochastic gradient
# (choosing one gradient $\operatorname{grad}f_k$ at random) has to be specified.
#
# For this small example you can of course also use a gradient descent with [`ArmijoLinesearch`](@ref),
# but it will be a little slower usually
@time x_opt5 = gradient_descent(
    M, F, (M, x) -> sum(gradF(M, x)), x; stepsize=ArmijoLinesearch()
);
# but it is for sure faster than the variant above that evaluates the full gradient on every iteration,
# since stochastic gradient descent takes more iterations.
#
# Note that all 5 of couse yield the same result
[distance(M, x, y) for y in [x_opt1, x_opt2, x_opt3, x_opt4, x_opt5]]


# ---
# title: 396. Rotate Function
# id: problem396
# author: Tian Jun
# date: 2020-11-14
# difficulty: Medium
# categories: Math
# link: <https://leetcode.com/problems/rotate-function/description/>
# hidden: true
# ---
#
# Given an array of integers `A` and let _n_ to be its length.
#
# Assume `Bk` to be an array obtained by rotating the array `A` _k_ positions
# clock-wise, we define a "rotation function" `F` on `A` as follow:
#
# `F(k) = 0 * Bk[0] + 1 * Bk[1] + ... + (n-1) * Bk[n-1]`.
#
# Calculate the maximum value of `F(0), F(1), ..., F(n-1)`.
#
# **Note:**
# _n_ is guaranteed to be less than 105.
#
# **Example:**
#
#
#
#     A = [4, 3, 2, 6]
#
#     F(0) = (0 * 4) + (1 * 3) + (2 * 2) + (3 * 6) = 0 + 3 + 4 + 18 = 25
#     F(1) = (0 * 6) + (1 * 4) + (2 * 3) + (3 * 2) = 0 + 4 + 6 + 6 = 16
#     F(2) = (0 * 2) + (1 * 6) + (2 * 4) + (3 * 3) = 0 + 6 + 8 + 9 = 23
#     F(3) = (0 * 3) + (1 * 2) + (2 * 6) + (3 * 4) = 0 + 2 + 12 + 12 = 26
#
#     So the maximum value of F(0), F(1), F(2), F(3) is F(3) = 26.
#
#
#
## @lc code=start
using LeetCode

function max_rotate_function(A::Vector{Int})::Int
    ## f(i+1) = f(i) + ∑(A) - A[end] * length(A)
    s = sum(A)
    n = length(A)
    cur = sum(k -> (k - 1) * A[k], 1:n)
    res = cur
    for i in n:-1:2
        cur += s - n * A[i]
        res = max(res, cur)
    end
    return res
end
## @lc code=end


export ode45, rk4, αscheme, αscheme_time,
    αscheme_atime, TR_BDF2, ExplicitNewmark, Newmark

function runge_kutta_one_step(f::Function, t::PyObject, y::PyObject, Δt::PyObject, θ::Union{PyObject, Missing})
    k1 = Δt*f(t, y, θ)
    k2 = Δt*f(t+Δt/2, y+k1/2, θ)
    k3 = Δt*f(t+Δt/2, y+k2/2, θ)
    k4 = Δt*f(t+Δt, y+k3, θ)
    y = y + k1/6 + k2/3 + k3/3 + k4/6
end

function ode45_one_step(f::Function, t::PyObject, y::PyObject, h::PyObject, θ::Union{PyObject, Missing})
    k1 = h * f(t, y, θ);
    k2 = h * f(t + (1/5)*h, y + (1/5)*k1, θ);
    k3 = h * f(t + (3/10)*h, y + (3/40)*k1 + (9/40)*k2, θ);
    k4 = h * f(t + (4/5)*h, y + (44/45)*k1 + (-56/15)*k2 + (32/9)*k3, θ);
    k5 = h * f(t + (8/9)*h, y + (19372/6561)*k1 + (-25360/2187)*k2 + (64448/6561)*k3 + (-212/729)*k4, θ);
    k6 = h * f(t + h, y + (9017/3168)*k1 + (-355/33)*k2 + (46732/5247 )*k3 + (49/176)*k4 + (-5103/18656)*k5, θ);
    k7 = h * f(t + h, y + (35/384)*k1 + (500/1113)*k3 + (125/192)*k4 + (-2187/6784)*k5 + (11/84)*k6, θ);

    y_new = y + (35/384)*k1 + (500/1113)*k3 + (125/192)*k4 + (-2187/6784)*k5 + (11/84)*k6;
end

@doc raw"""
    runge_kutta(f::Function, T::Union{PyObject, Float64},
                NT::Union{PyObject,Int64}, y::Union{PyObject, Float64, Array{Float64}}, θ::Union{PyObject, Missing}=missing; method::String="rk4")

Solves
```math
\frac{dy}{dt} = f(y, t, \theta)
```
with Runge-Kutta method.

For example, the default solver, `RK4`, has the following numerical scheme per time step
```math
\begin{aligned}
k_1 &= \Delta t f(t_n, y_n, \theta)\\
k_2 &= \Delta t f(t_n+\Delta t/2, y_n + k_1/2, \theta)\\
k_3 &= \Delta t f(t_n+\Delta t/2, y_n + k_2/2, \theta)\\
k_4 &= \Delta t f(t_n+\Delta t, y_n + k_3, \theta)\\
y_{n+1} &= y_n + \frac{k_1}{6} +\frac{k_2}{3} +\frac{k_3}{3} +\frac{k_4}{6}
\end{aligned}
```
"""
function runge_kutta(f::Function, T::Union{PyObject, Float64},
                NT::Union{PyObject,Int64}, y::Union{PyObject, Float64, Array{Float64}}, θ::Union{PyObject, Missing}=missing; method::String="rk4")

    local one_step
    if lowercase(method)=="rk4"
        one_step = runge_kutta_one_step
    elseif lowercase(method)=="rk45"
        one_step = ode45_one_step
    else
        error("Method $method not implemented yet")
    end
    y = convert_to_tensor(y)
    Δt = convert_to_tensor(T/NT)

    ta = TensorArray(NT+1) # storing y
    function condition(i, ta)
        i <= NT+1
    end
    function body(i, ta)
        y = read(ta, i-1)
        y_ = one_step(f, (cast(eltype(Δt), i)-1)*Δt, y, Δt, θ)
        ta = write(ta, i, y_)
        i+1, ta
    end
    ta = write(ta, 1, y)
    i = constant(2, dtype=Int32)
    _, out = while_loop(condition, body, [i, ta])
    res = stack(out)
end

@doc raw"""
    rk4(y::Union{PyObject, Float64, Array{Float64}}, T::Union{PyObject, Float64},
                NT::Union{PyObject,Int64}, f::Function, θ::Union{PyObject, Missing}=missing)

Solves
```math
\frac{dy}{dt} = f(y, t, \theta)
```
with Runge-Kutta (order 4) method.
"""
rk4(args...;kwargs...) = runge_kutta(args...;method="rk4", kwargs...)

@doc raw"""
    ode45(y::Union{PyObject, Float64, Array{Float64}}, T::Union{PyObject, Float64},
                NT::Union{PyObject,Int64}, f::Function, θ::Union{PyObject, Missing}=missing)

Solves
```math
\frac{dy}{dt} = f(y, t, \theta)
```
with six-stage, fifth-order, Runge-Kutta method.
"""
ode45(args...;kwargs...) = runge_kutta(args...;method="rk45", kwargs...)

@doc raw"""
    αscheme(M::Union{SparseTensor, SparseMatrixCSC},
        C::Union{SparseTensor, SparseMatrixCSC},
        K::Union{SparseTensor, SparseMatrixCSC},
        Force::Union{Array{Float64}, PyObject},
        d0::Union{Array{Float64, 1}, PyObject},
        v0::Union{Array{Float64, 1}, PyObject},
        a0::Union{Array{Float64, 1}, PyObject},
        Δt::Array{Float64};
        solve::Union{Missing, Function} = missing,
        extsolve::Union{Missing, Function} = missing,
        ρ::Float64 = 1.0)

Generalized α-scheme.
$$M u_{tt} + C u_{t} + K u = F$$

`Force` must be an array of size `n`×`p`, where `d0`, `v0`, and `a0` have a size `p`
`Δt` is an array (variable time step).

The generalized α scheme solves the equation by the time stepping
```math
\begin{aligned}
\bf d_{n+1} &= \bf d_n + h\bf v_n + h^2 \left(\left(\frac{1}{2}-\beta_2 \right)\bf a_n + \beta_2 \bf a_{n+1}  \right)\\
\bf v_{n+1} &= \bf v_n + h((1-\gamma_2)\bf a_n + \gamma_2 \bf a_{n+1})\\
\bf F(t_{n+1-\alpha_{f_2}}) &= M \bf a _{n+1-\alpha_{m_2}} + C \bf v_{n+1-\alpha_{f_2}} + K \bf{d}_{n+1-\alpha_{f_2}}
\end{aligned}
```
where
```math
\begin{aligned}
\bf d_{n+1-\alpha_{f_2}} &= (1-\alpha_{f_2})\bf d_{n+1} + \alpha_{f_2} \bf d_n\\
\bf v_{n+1-\alpha_{f_2}} &= (1-\alpha_{f_2}) \bf v_{n+1} + \alpha_{f_2} \bf v_n \\
\bf a_{n+1-\alpha_{m_2} } &= (1-\alpha_{m_2}) \bf a_{n+1} + \alpha_{m_2} \bf a_n\\
t_{n+1-\alpha_{f_2}} & = (1-\alpha_{f_2}) t_{n+1 + \alpha_{f_2}} + \alpha_{f_2}t_n
\end{aligned}
```

Here the parameters are computed using
```math
\begin{aligned}
\gamma_2 &= \frac{1}{2} - \alpha_{m_2} + \alpha_{f_2}\\
\beta_2 &= \frac{1}{4} (1-\alpha_{m_2}+\alpha_{f_2})^2 \\
\alpha_{m_2} &= \frac{2\rho_\infty-1}{\rho_\infty+1}\\
\alpha_{f_2} &= \frac{\rho_\infty}{\rho_\infty+1}
\end{aligned}
```

∘ `solve`: users can provide a solver function, `solve(A, rhs)` for solving `Ax = rhs`
∘ `extsolve`: similar to `solve`, but the signature has the form
```julia
extsolve(A, rhs, i)
```
This provides the users with more control, e.g., (time-dependent) Dirichlet boundary conditions.
See [Generalized α Scheme](https://kailaix.github.io/ADCME.jl/dev/alphascheme/) for details.

!!! note
    In the case $u$ has a nonzero essential boundary condition $u_b$, we let $\tilde u=u-u_b$, then
    $$M \tilde u_{tt} + C \tilde u_t + K u = F - K u_b - C \dot u_b$$
"""
function αscheme(M::Union{SparseTensor, SparseMatrixCSC},
                      C::Union{SparseTensor, SparseMatrixCSC},
                      K::Union{SparseTensor, SparseMatrixCSC},
                      Force::Union{Array{Float64, 2}, PyObject},
                      d0::Union{Array{Float64, 1}, PyObject},
                      v0::Union{Array{Float64, 1}, PyObject},
                      a0::Union{Array{Float64, 1}, PyObject},
                      Δt::Array{Float64, 1};
                      solve::Union{Missing, Function} = missing,
                      extsolve::Union{Missing, Function} = missing,
                      ρ::Float64 = 1.0)
    if !ismissing(solve) && !ismissing(extsolve)
        error("You cannot provide `solve` and `extsolve` at the same time.")
    end
    nt = length(Δt)
    αm = (2ρ-1)/(ρ+1)
    αf = ρ/(1+ρ)
    γ = 1/2-αm+αf
    β = 0.25*(1-αm+αf)^2
    d = length(d0)

    M = isa(M, SparseMatrixCSC) ? constant(M) : M
    C = isa(C, SparseMatrixCSC) ? constant(C) : C
    K = isa(K, SparseMatrixCSC) ? constant(K) : K
    Force, d0, v0, a0, Δt = convert_to_tensor([Force, d0, v0, a0, Δt], [Float64, Float64, Float64, Float64, Float64])

    function equ(dc, vc, ac, dt, Force, i)
        dn = dc + dt*vc + dt^2/2*(1-2β)*ac
        vn = vc + dt*((1-γ)*ac)

        df = (1-αf)*dn + αf*dc
        vf = (1-αf)*vn + αf*vc
        am = αm*ac

        rhs = Force - (M*am + C*vf + K*df)
        A = (1-αm)*M + (1-αf)*C*dt*γ + (1-αf)*K*β*dt^2

        if !ismissing(solve)
            return solve(A, rhs)
        elseif !ismissing(extsolve)
            return extsolve(A, rhs, i)
        else
            return A\rhs
        end
    end

    function condition(i, tas...)
        return i<=nt
    end
    function body(i, tas...)
        dc_arr, vc_arr, ac_arr = tas
        dc = read(dc_arr, i)
        vc = read(vc_arr, i)
        ac = read(ac_arr, i)
        y = equ(dc, vc, ac, Δt[i], Force[i], i)
        dn = dc + Δt[i]*vc + Δt[i]^2/2*((1-2β)*ac+2β*y)
        vn = vc + Δt[i]*((1-γ)*ac+γ*y)
        i+1, write(dc_arr, i+1, dn), write(vc_arr, i+1, vn), write(ac_arr, i+1, y)
    end

    dM = TensorArray(nt+1); vM = TensorArray(nt+1); aM = TensorArray(nt+1)
    dM = write(dM, 1, d0)
    vM = write(vM, 1, v0)
    aM = write(aM, 1, a0)
    i = constant(1, dtype=Int32)
    _, d, v, a = while_loop(condition, body, [i,dM, vM, aM])
    set_shape(stack(d), (nt+1, length(d0))), set_shape(stack(v), (nt+1, length(v0))), set_shape(stack(a), (nt+1, length(a0)))
end

@doc raw"""
    αscheme_time(Δt::Array{Float64}; ρ::Float64 = 1.0)

Returns the integration time $t_{i+1-\alpha_{f_2}}$ between $[t_i, t_{i+1}]$ using the alpha scheme.
If $\Delta t$ has length $n$, the output will also have length $n$.
"""
function αscheme_time(Δt::Array{Float64}; ρ::Float64 = 1.0)
    n = length(Δt)
    αm = (2ρ-1)/(ρ+1)
    αf = ρ/(1+ρ)
    γ = 1/2-αm+αf
    β = 0.25*(1-αm+αf)^2
    function equ(tc, dt)
        tf1 = (1-αf)*(tc+dt) + αf*tc
        return tf1
    end

    tcf = Float64[]
    tc = 0.0
    for i = 1:n
        t1 = equ(tc, Δt[i])
        push!(tcf, t1)
        tc += Δt[i]
    end
    return tcf
end

@doc raw"""
    TR_BDF2(D0::Union{SparseTensor, SparseMatrixCSC},
        D1::Union{SparseTensor, SparseMatrixCSC},
        Δt::Float64)

Constructs a TR-BDF2 (the Trapezoidal Rule with Second Order Backward Difference Formula) handler for
the DAE

$$D_1 \dot y + D_0 y = f$$

The struct is a functor, which performs one step simulation

```
(tr::TR_BDF2)(y::Union{PyObject, Array{Float64, 1}},
    f1::Union{PyObject, Array{Float64, 1}},
    f2::Union{PyObject, Array{Float64, 1}},
    f3::Union{PyObject, Array{Float64, 1}})
```
Here `f1`, `f2`, and `f3` correspond to the right hand side at time step $n$, $n+\frac12$, and $n+1$.

Or we can pass a batched `F` defined as a `(2NT+1) × DOF` array

```
(tr::TR_BDF2)(y0::Union{PyObject, Array{Float64, 1}},
    F::Union{PyObject, Array{Float64, 2}})
```

The output will be the entire solution of size `(NT+1) × DOF`.

!!! info
    The scheme takes the following form for n = 0, 1, ...
    $$\begin{aligned} D_1(y^{n+\frac12}-y^n) = \frac12\frac{\Delta t}{2}\left(f^{n+\frac12} + f^n - D_0 \left(y^{n+\frac12} + y^n\right)\right)\\ \left(\frac{\Delta t}{2}\right)^{-1} D_1 \left(\frac32y^{n+1} - 2y^{n+\frac12} + \frac12 y^n\right) + D_0 y^{n+1} = f^{n+1}\end{aligned}$$
"""
mutable struct TR_BDF2
    D0::Union{SparseTensor, SparseMatrixCSC}
    D1::Union{SparseTensor, SparseMatrixCSC}
    Δt::Float64
    _D0::Union{SparseTensor, SparseMatrixCSC}
    _D1::Union{SparseTensor, SparseMatrixCSC}
    symbolic::Bool
    function TR_BDF2(D0::Union{SparseTensor, SparseMatrixCSC},
            D1::Union{SparseTensor, SparseMatrixCSC},
            Δt::Float64)
        if isa(D0, SparseTensor) || isa(D1, SparseTensor)
            symbolic = true
            D0 = constant(D0)
            D1 = constant(D1)
        else
            symbolic = false
        end
        @assert size(D0, 1)==size(D0,2)==size(D1,1)==size(D1,2)
        _D0 = D1 + Δt/4 * D0
        _D1 = 1/(Δt/2) * 3/2 * D1  + D0
        new(D0, D1, Δt, _D0, _D1, symbolic)
    end
end

"""
    constant(tr::TR_BDF2)

Converts `tr` to a symbolic solver.
"""
function constant(tr::TR_BDF2)
    TR_BDF2(constant(tr.D0), constant(tr.D1), tr.Δt)
end

function Base.:show(io::IO, tr::TR_BDF2)
    print("""TR_BDF2 (DOF = $(size(tr.D0, 1)), Δt = $(tr.Δt)$(tr.symbolic ? ", symbolic" : ""))""")
end

function (tr::TR_BDF2)(y::Union{PyObject, Array{Float64, 1}},
    f1::Union{PyObject, Array{Float64, 1}},
    f2::Union{PyObject, Array{Float64, 1}},
    f3::Union{PyObject, Array{Float64, 1}})
    y, f1, f2, f3 = convert_to_tensor([y, f1, f2, f3], [Float64, Float64, Float64, Float64])
    r1 = tr.Δt/4 * (f2 + f1) - tr.Δt/4 * (tr.D0 * y) + tr.D1 * y
    yn = tr._D0\r1
    r2 = f3 + 1/(tr.Δt/2)*(tr.D1* (2*yn - 0.5*y))
    tr._D1\r2
end

function (tr::TR_BDF2)(y0::Union{PyObject, Array{Float64, 1}},
    F::Union{PyObject, Array{Float64, 2}})
    @assert size(F, 2)==size(tr.D0, 1) && mod(size(F, 1), 2)==1
    y0, F = convert_to_tensor([y0, F], [Float64, Float64])
    NT = size(F, 1)÷2
    y_arr = TensorArray(NT+1)
    y_arr = write(y_arr, 1, y0)
    function condition(i, y_arr)
        i<=NT
    end
    function body(i, y_arr)
        y = read(y_arr, i)
        f1, f2, f3 = F[2*i-1], F[2*i], F[2*i+1]
        yn = tr(y, f1, f2, f3)
        i+1, write(y_arr, i+1, yn)
    end
    i = constant(1, dtype = Int32)
    _, ya = while_loop(condition, body, [i, y_arr])
    set_shape(stack(ya), (NT+1, length(y0)))
end

function (tr::TR_BDF2)(y::Array{Float64, 1},
    f1::Array{Float64, 1},
    f2::Array{Float64, 1},
    f3::Array{Float64, 1})
    if tr.symbolic
        return tr(constant(y), f1, f2, f3)
    end
    r1 = tr.Δt/4 * (f2 + f1) - tr.Δt/4 * (tr.D0 * y) + tr.D1 * y
    yn = tr._D0\r1
    r2 = f3 + 1/(tr.Δt/2)*(tr.D1* (2*yn - 0.5*y))
    tr._D1\r2
end

function (tr::TR_BDF2)(y0::Array{Float64,1}, F::Array{Float64, 2})
    if tr.symbolic
        return tr(constant(y0), F)
    end
    @assert size(F, 2)==size(tr.D0, 1) && mod(size(F, 1), 2)==1
    y = zeros(size(F, 1)÷2+1, 2)
    y[1,:] = y0
    for i = 1:size(F, 1)÷2
        y[i+1,:] = tr(y[i,:], F[2*i-1,:], F[2*i,:], F[2*i+1,:])
    end
    y
end

# design principle: proceed one step
# (states) u0, u1, ..., (coefficients) M0, M1, ..., (step size) Δt

@doc raw"""
    ExplicitNewmark(M::Union{SparseTensor, SparseMatrixCSC}, Z1::Union{Missing, SparseTensor, SparseMatrixCSC}, Z2::Union{Missing, SparseTensor, SparseMatrixCSC}, Δt::Float64)

An explicit Newmark integrator for

$$M \ddot{\mathbf{d}} + Z_1 \dot{\mathbf{d}} + Z_2 \mathbf{d} + f = 0$$

The numerical scheme is

$$\left(\frac{1}{\Delta t^2} M + \frac{1}{2\Delta t}Z_1\right)d^{n+1} = \left(\frac{2}{\Delta t^2} M - Z_2\right)d^n - \left(\frac{1}{\Delta t^2} M - \frac{1}{2\Delta t}Z_1\right) d^{n-1} - f$$

To use this integrator,

```julia
en = ExplicitNewmark(M, Z1, Z2, Δt)
d2 = step(en, d0, d1, f)
```
"""
mutable struct ExplicitNewmark
    A::Union{PyObject, Tuple{SparseTensor, PyObject}, SparseTensor}
    B::Union{PyObject, SparseTensor}
    C::Union{PyObject, SparseTensor}
    function ExplicitNewmark(M::Union{SparseTensor, PyObject, Array{Float64, 2}, SparseMatrixCSC},
            Z1::Union{Missing, PyObject, Array{Float64, 2}, SparseTensor, SparseMatrixCSC},
            Z2::Union{Missing, PyObject, Array{Float64, 2}, SparseTensor, SparseMatrixCSC}, Δt::Float64;
            factorize_A::Bool = true)
        M = constant(M)
        if ismissing(Z1)
            A = 1/Δt^2 * M
            C = -1/Δt^2 * M
        else
            Z1 = constant(Z1)
            A = (1/Δt^2 * M + 1/(2Δt) * Z1)
            C = -(1/Δt^2 * M - 1/(2Δt) * Z1)
        end
        if isa(A, SparseTensor) && factorize_A
            A = factorize(A)
        end
        if ismissing(Z2)
            B = 2/Δt^2 * M
        else
            Z2 = constant(Z2)
            B = (2/Δt^2 * M - Z2)
        end
        new(A, B, C)
    end
end

function Base.:show(io::IO, en::ExplicitNewmark)
    print("ExplicitNewmark(DOF=$(size(en.B, 1)))")
end

function Base.:step(en::ExplicitNewmark,
    d0::Union{Array{Float64, 1}, PyObject},
    d1::Union{Array{Float64, 1}, PyObject},
    f::Union{Array{Float64, 1}, PyObject})
    d0, d1, f = convert_to_tensor([d0, d1, f], [Float64, Float64, Float64])
    en.A \ (en.B * d1 + en.C * d0 - f)
end

@doc raw"""
    Newmark(M::Union{SparseTensor, PyObject, Array{Float64, 2}, SparseMatrixCSC},
                        C::Union{SparseTensor, PyObject, Array{Float64, 2}, SparseMatrixCSC},
                        K::Union{SparseTensor, PyObject, Array{Float64, 2}, SparseMatrixCSC},
                        Δt::Float64; β::Float64 = 1/4, γ::Float64 = 1/2, factorize_S::Bool = true)

Creates a Newmark algorithm object. This function uses "a-form" implementation for

$$M\ddot \mathbf{d} + C \dot\mathbf{d} + K \mathbf{d} + f = 0$$

The numerical scheme is

$$\begin{aligned}
\tilde d_{n+1} &= d_n + \Delta t v_n + \Delta^2/2(1-2\beta)a_n \\
\tilde v_{n+1} &= v_n + (1-\gamma)\Delta a_n \\
(M+\gamma\Delta C+\beta \Delta^2 K)a_{n+1} &= F_{n+1} - C\tilde v_{n+1} - K\tilde d_{n+1}\\
d_{n+1} &= \tilde d_{n+1} + \beta \Delta^2 a_{n+1}\\
v_{n+1} &= \tilde v_{n+1} + \gamma \Delta ta_n
\end{aligned}$$

For unconditional stability
$$2\beta \geq \gamma \geq \frac12$$

For an explicit form, see [`ExplicitNewmark`](@ref), which corresponds to $\beta = 0, \gamma = \frac12$.

A common choice for Newmark algorithm is $\beta = \frac14, \gamma = \frac12$ (a.k.a., average acceleration, or
Trapezoidal rule), which is unconditionally stable,
has second order of accuracy.

!!! info
    Internally, $M+\gamma\Delta C+\beta \Delta^2 K$ is represented by `S`. Users can elect to factorize $S$ (default $=$ true).

To perform time marching, use
```
step(nm::Newmark, d::PyObject, v::PyObject, a::PyObject,
        F::Union{Array{Float64, 1}, PyObject})
```
"""
mutable struct Newmark
    M::SparseTensor
    C::SparseTensor
    K::SparseTensor
    S::Union{Tuple{SparseTensor, PyCall.PyObject}, SparseTensor}
    β::Float64
    γ::Float64
    Δt::Float64
    function Newmark(M::Union{SparseTensor, PyObject, Array{Float64, 2}, SparseMatrixCSC},
            C::Union{SparseTensor, PyObject, Array{Float64, 2}, SparseMatrixCSC},
            K::Union{SparseTensor, PyObject, Array{Float64, 2}, SparseMatrixCSC},
            Δt::Float64; β::Float64 = 1/4, γ::Float64 = 1/2, factorize_S::Bool = true)
        M = SparseTensor(M)
        C = SparseTensor(C)
        K = SparseTensor(K)
        S = M + γ*Δt*C + β*Δt^2*K
        if factorize_S
            S = factorize(S)
        end
        new(M, C, K, S, β, γ, Δt)
    end
end

function Base.:show(io::IO, en::Newmark)
    print("Newmark(DOF=$(size(en.S, 1)), γ=$(en.γ), β=$(en.β))")
end

function Base.:step(nm::Newmark, d::PyObject, v::PyObject, a::PyObject,
        F::Union{Array{Float64, 1}, PyObject})
    β, γ, Δt = nm.β, nm.γ, nm.Δt
    td = d + Δt * v + Δt^2/2*(1-2β)*a
    tv = v + (1-γ)*Δt*a
    a_ = nm.S\(-F - nm.C*tv - nm.K*td)
    d_ = td + β*Δt^2*a_
    v_ = tv + γ*Δt*a_
    d_, v_, a_
end


export combinations,
       CoolLexCombinations,
       multiset_combinations,
       with_replacement_combinations,
       powerset

#The Combinations iterator

struct Combinations{T}
    a::T
    t::Int
end

function Base.iterate(c::Combinations, s = collect(1:c.t))
    (!isempty(s) && s[1] > length(c.a) - c.t + 1) && return

    comb = [c.a[si] for si in s]
    if c.t == 0
        # special case to generate 1 result for t==0
        return (comb, [length(c.a)+2])
    end
    s = copy(s)
    for i = length(s):-1:1
        s[i] += 1
        if s[i] > (length(c.a) - (length(s) - i))
            continue
        end
        for j = i+1:lastindex(s)
            s[j] = s[j-1]+1
        end
        break
    end
    (comb, s)
end

Base.length(c::Combinations) = binomial(length(c.a), c.t)

Base.eltype(::Type{Combinations{T}}) where {T} = Vector{eltype(T)}

"""
    combinations(a, n)

Generate all combinations of `n` elements from an indexable object `a`. Because the number
of combinations can be very large, this function returns an iterator object.
Use `collect(combinations(a, n))` to get an array of all combinations.
"""
function combinations(a, t::Integer)
    if t < 0
        # generate 0 combinations for negative argument
        t = length(a) + 1
    end
    Combinations(a, t)
end

"""
    combinations(a)

Generate combinations of the elements of `a` of all orders. Chaining of order iterators
is eager, but the sequence at each order is lazy.
"""
combinations(a) = Iterators.flatten([combinations(a, k) for k = 1:length(a)])

# cool-lex combinations iterator

"""
    CoolLexCombinations

Produce ``(n,k)``-combinations in cool-lex order.

# Reference

Ruskey, F., & Williams, A. (2009). The coolest way to generate combinations.
*Discrete Mathematics*, 309(17), 5305-5320.
"""
struct CoolLexCombinations
    n::Int
    t::Int
end

struct CoolLexIterState{T<:Integer}
    R0::T
    R1::T
    R2::T
    R3::T
end

function Base.iterate(C::CoolLexCombinations)
    if C.n < 0
        throw(DomainError(C.n))
    end
    if C.t ≤ 0
        throw(DomainError(C.t))
    end

    #What integer size should I use?
    if C.n < 8sizeof(Int)
        T = Int
    else
        T = BigInt
    end

    state = CoolLexIterState{T}(0, 0, T(1) << C.n, (T(1) << C.t) - 1)
    iterate(C, state)
end

function Base.iterate(C::CoolLexCombinations, S::CoolLexIterState)
    (S.R3 & S.R2 != 0) && return

    R0 = S.R0
    R1 = S.R1
    R2 = S.R2
    R3 = S.R3

    R0 = R3 & (R3 + 1)
    R1 = xor(R0, R0 - 1)
    R0 = R1 + 1
    R1 &= R3
    R0 = max((R0 & R3) - 1, 0)
    R3 += R1 - R0

    _cool_lex_visit(S.R3), CoolLexIterState(R0, R1, R2, R3)
end

#Converts an integer bit pattern X into a subset
#If X & 2^k == 1, then k is in the subset
function _cool_lex_visit(X::Integer)
    subset = Int[]
    n = 1
    while X != 0
        X & 1 == 1 && push!(subset, n)
        X >>= 1
        n += 1
    end
    subset
end

Base.length(C::CoolLexCombinations) = max(0, binomial(C.n, C.t))

struct MultiSetCombinations{T}
    m::T
    f::Vector{Int}
    t::Int
    ref::Vector{Int}
end

Base.eltype(::Type{MultiSetCombinations{T}}) where {T} = Vector{eltype(T)}

function Base.length(c::MultiSetCombinations)
    t = c.t
    if t > length(c.ref)
        return 0
    end
    p = [1; zeros(Int, t)]
    for i in 1:length(c.f)
        f = c.f[i]
        if i == 1
            for j in 1:min(f, t)
                p[j+1] = 1
            end
        else
            for j in t:-1:1
                p[j+1] = sum(p[max(1,j+1-f):(j+1)])
            end
        end
    end
    return p[t+1]
end

function multiset_combinations(m, f::Vector{<:Integer}, t::Integer)
    length(m) == length(f) || error("Lengths of m and f are not the same.")
    ref = length(f) > 0 ? vcat([[i for j in 1:f[i] ] for i in 1:length(f)]...) : Int[]
    if t < 0
        t = length(ref) + 1
    end
    MultiSetCombinations(m, f, t, ref)
end

"""
    multiset_combinations(a, t)

Generate all combinations of size `t` from an array `a` with possibly duplicated elements.
"""
function multiset_combinations(a, t::Integer)
    m = unique(collect(a))
    f = Int[sum([c == x for c in a]) for x in m]
    multiset_combinations(m, f, t)
end

function Base.iterate(c::MultiSetCombinations, s = c.ref)
    ((!isempty(s) && max(s[1], c.t) > length(c.ref)) || (isempty(s) && c.t > 0)) && return

    ref = c.ref
    n = length(ref)
    t = c.t
    changed = false
    comb = [c.m[s[i]] for i in 1:t]
    if t > 0
        s = copy(s)
        for i in t:-1:1
            if s[i] < ref[i + (n - t)]
                j = 1
                while ref[j] <= s[i]
                    j += 1
                end
                s[i] = ref[j]
                for l in (i+1):t
                    s[l] = ref[j+=1]
                end
                changed = true
                break
            end
        end
        !changed && (s[1] = n+1)
    else
        s = [n+1]
    end
    (comb, s)
end

struct WithReplacementCombinations{T}
    a::T
    t::Int
end

Base.eltype(::Type{WithReplacementCombinations{T}}) where {T} = Vector{eltype(T)}

Base.length(c::WithReplacementCombinations) = binomial(length(c.a) + c.t - 1, c.t)

"""
    with_replacement_combinations(a, t)

Generate all combinations with replacement of size `t` from an array `a`.
"""
with_replacement_combinations(a, t::Integer) = WithReplacementCombinations(a, t)

function Base.iterate(c::WithReplacementCombinations, s = [1 for i in 1:c.t])
    (!isempty(s) && s[1] > length(c.a) || c.t < 0) && return

    n = length(c.a)
    t = c.t
    comb = [c.a[si] for si in s]
    if t > 0
        s = copy(s)
        changed = false
        for i in t:-1:1
            if s[i] < n
                s[i] += 1
                for j in (i+1):t
                    s[j] = s[i]
                end
                changed = true
                break
            end
        end
        !changed && (s[1] = n+1)
    else
        s = [n+1]
    end
    (comb, s)
end

## Power set

"""
    powerset(a, min=0, max=length(a))

Generate all subsets of an indexable object `a` including the empty set, with cardinality
bounded by `min` and `max`. Because the number of subsets can be very large, this function
returns an iterator object. Use `collect(powerset(a, min, max))` to get an array of all
subsets.
"""
function powerset(a, min::Integer=0, max::Integer=length(a))
    itrs = [combinations(a, k) for k = min:max]
    min < 1 && append!(itrs, eltype(a)[])
    Iterators.flatten(itrs)
end


### A Pluto.jl notebook ###
# v0.12.18

using Markdown
using InteractiveUtils

# ╔═╡ ceff6126-51fc-11eb-3cf5-dd30d637e267
begin
	using Dates
	using MortalityTables
	using DifferentialEquations
	using Plots
	using ActuaryUtilities
end

# ╔═╡ cd5e7080-5348-11eb-2c22-455a3556cab0
md"
# Universal Life Policy Account Mechanics as a Differential Equation
"

# ╔═╡ de0ace1a-5348-11eb-35c9-5f0ff4009d97
md"""
## Introduction

This demonstrates an example of defining an universal life policy value roll like a discrete differential equation.

It then uses the [SciML DifferentialEquations](https://sciml.ai/) package to "solve" the policy projection given a single point, but also to see how the policy projection behaves under different premium and interest rate conditions./
"""

# ╔═╡ ef26ffb6-5348-11eb-04a8-c5cbfdcab72a
md"""
Let's use the 2001 CSO table as the basis for cost of insurance charges:
"""

# ╔═╡ 45dab59c-528d-11eb-357d-af3664466910
begin
	tables = MortalityTables.tables()
	cso = tables["2001 CSO Super Preferred Select and Ultimate - Male Nonsmoker, ANB"]
end

# ╔═╡ ffe9b6e0-5348-11eb-2d38-a5613eb24383
md"""
Next, policy mechanics are coded. It's essentially a discrete differential equation, so it leverages [`DifferentialEquations.jl`](https://diffeq.sciml.ai/dev/)

The projection is coded in the [Discrete DifferentialEquation format](https://diffeq.sciml.ai/dev/types/discrete_types/):

``u_{n+1} = f(u,p,t_{n+1})``

In the code below, this translates to:

- `u` is the *state* of the system. We  will track three variables to represent the `state`:
   - `state[1]` is the account value
   - `state[2]` is the premium paid
   - `state[3]` is the policy duration
- `p` are the parameters of the system.
- `t` is the time, which will represent days since policy issuance
"""

# ╔═╡ da199644-51fc-11eb-0aab-dd68b1fcddc6
function policy_projection(state,p,t)
    # grab the state from the inputs
    av = state[1]

    # calculated variables
    cur_date = p.issue_date + Day(t)
    dur = duration(p.issue_date,cur_date)
    att_age = p.issue_age + dur - 1

	# lapse if AV <= 0
    lapsed = (av <= 0.0 ) & (t > 1)

    if !lapsed

        monthly_coi_rate = (1 - (1-p.mort_assump[att_age]) ^ (1/12))

		## Periodic Policy elements

        # annual events
        if Dates.monthday(cur_date) == Dates.monthday(p.issue_date) ||
			cur_date ==p.issue_date + Day(1) # OR first issue date

            premium = p.annual_prem
        else
            premium = 0.0
        end

        # monthly_events
        if Dates.day(cur_date) == Dates.day(p.issue_date)
            coi = max((p.face - av) * monthly_coi_rate,0.0)
        else
            coi = 0.0
        end

        # daily events
        int(av) = av * ((1 + p.int_rate) ^ (1 / 360) - 1.0)

        # av
        new_av = max(0.0,av - coi + premium + int(av-coi))

				# new state
		return [new_av, premium, dur] # AV, Prem, Dur

    else
		# new state
		return [0.0, 0.0, dur] # AV, Prem, Dur

    end

end

# ╔═╡ 8a227236-5349-11eb-353d-1d3f94a968d0
md"The following function will create a named tuple of parameters given a varying `prem` (premium) and `int` (credit rate)."

# ╔═╡ 956811e4-5349-11eb-2b34-b5441def32a1
params(prem,int) = (
    int_rate = int,
    issue_date = Date(2010,1,1),
    face = 1e6,
    issue_age = 25,
    mort_assump = tables["2001 CSO Super Preferred Select and Ultimate - Male Nonsmoker, ANB"].ultimate,
    projection_years = 75,
    annual_prem = prem,
)

# ╔═╡ 9ee35710-5349-11eb-0d45-23bc292d8a3d
md"""
## Runing the system

This results in the following plot. The tracked output variables u1 and u2 represent the two vars that we tracked above: account value and cumulative premium.
"""

# ╔═╡ b65ef28c-5349-11eb-1baa-6575ad7fb946
begin
	p = params(
	        8000.0, # 8,000 annual premium
	        0.08    # 8% interest
	        )

	# calculate the number of days to project
	projection_end_date = p.issue_date + Year(p.projection_years)
	days_to_project = Dates.value(projection_end_date - p.issue_date)

	# the [0.0,..] are the initial conditions for the tracked variables
	 prob = DiscreteProblem(policy_projection,[0.0,0.0,0],(0,days_to_project),p)
	 proj = solve(prob,FunctionMap())

	 plot(proj)
end

# ╔═╡ c4637664-5349-11eb-0e81-59e39974f02e
md"""
## Moving up the ladder of abstraction

An excellent way to understand the behavior of a model is to [move up the ladder of abstraction](http://worrydream.com/LadderOfAbstraction/). Below, we will see what happens to the projection at varying levels of credit rates and annual premiums.
"""

# ╔═╡ 02cc352a-5409-11eb-0fda-336d5fa42fa9
begin
	prem_range = 1000.0:100.0:9000.0
	int_range = 0.02:0.0025:0.08

	function ending_av(ann_prem,int,days_to_project)
	    p = params(ann_prem,int)
	    prob = DiscreteProblem(policy_projection,[0.0,0.0,0],(0,days_to_project),p)
	    proj = solve(prob,FunctionMap())
	    end_av = proj[end][1]
	    if end_av == 0.0
	        lapse_time = findfirst(isequal(0.0),proj[1,2:end])
	    else
	        lapse_time = length(proj)
	    end
	    duration = proj[3,lapse_time]
	    end_age = p.issue_age + duration - 1.0

	    return end_av,end_age

	end

	end_age = zeros(length(prem_range),length(int_range))
	end_av = zeros(length(prem_range),length(int_range))

	# loop through each projection we did and fill our ranges with the ending AV and ending age
	for (i,vp) in enumerate(prem_range)
	    for (j,vi) in enumerate(int_range)
	        (end_av[i,j],end_age[i,j]) = ending_av(vp,vi,days_to_project)
	    end
	end
end

# ╔═╡ 280d25a6-5409-11eb-17d9-674f218f4351
begin
	using ColorSchemes # for Turbo colors, which emphasize readability

	viz = plot(layout=2) # side by side plot

	#
	contour!(viz[1],int_range,
	    prem_range,
	    end_av ./ 1e6, # scale to millions for readability

	    contour_labels=true,
	    c=cgrad(ColorSchemes.turbo.colors),
	    fill=true,
	    title="AV at age 100 (\$M)",
	    ylabel="Annual Premium (\$)",

	)

	contour!(viz[2],int_range,
	    prem_range,
	    end_age,

	    contour_labels=true,
	    c=cgrad(ColorSchemes.turbo.colors),
	    fill=true,
	    yaxis=false,
	    title="Age at Lapse",
	)

	annotate!(viz[2],0.055,7000,Plots.text("Doesn't lapse \nbefore age 100", 8, :white, :center))

end

# ╔═╡ 21e53f60-5409-11eb-291d-d903da8e16b4
md"
Now let's plot the result. Not surprising, interest has a huge effect on the policy projection. Premium is also a major influence.

One thing that's remarkable is how going from 2000 premium to just ~2200 of premium results in about a $5m difference at 8% interest. The power of compound interest!
"

# ╔═╡ 684f7efc-5409-11eb-0757-19b40c5d418e
md"
## Conclusion

This shows how universal life mechanics are a dynamic system. the growth/decay is governed by two competing feedback loops:

- Growth: the force of interest lets the balance grow exponentially over long periods of time

- Decay: low balances increase the net amount at risk and the resulting COI charges.

"

# ╔═╡ 7e80faf0-5409-11eb-33f6-4f6be6fea835
md"
## Endnotes

This is not meant to represent any particular insurance product, nor fully replicate typical account mechanics.
"

# ╔═╡ Cell order:
# ╟─cd5e7080-5348-11eb-2c22-455a3556cab0
# ╟─de0ace1a-5348-11eb-35c9-5f0ff4009d97
# ╠═ceff6126-51fc-11eb-3cf5-dd30d637e267
# ╟─ef26ffb6-5348-11eb-04a8-c5cbfdcab72a
# ╠═45dab59c-528d-11eb-357d-af3664466910
# ╟─ffe9b6e0-5348-11eb-2d38-a5613eb24383
# ╠═da199644-51fc-11eb-0aab-dd68b1fcddc6
# ╟─8a227236-5349-11eb-353d-1d3f94a968d0
# ╠═956811e4-5349-11eb-2b34-b5441def32a1
# ╟─9ee35710-5349-11eb-0d45-23bc292d8a3d
# ╠═b65ef28c-5349-11eb-1baa-6575ad7fb946
# ╟─c4637664-5349-11eb-0e81-59e39974f02e
# ╠═02cc352a-5409-11eb-0fda-336d5fa42fa9
# ╟─21e53f60-5409-11eb-291d-d903da8e16b4
# ╠═280d25a6-5409-11eb-17d9-674f218f4351
# ╟─684f7efc-5409-11eb-0757-19b40c5d418e
# ╟─7e80faf0-5409-11eb-33f6-4f6be6fea835


using Cubature
using Distributions

"""
Demo of robust Bayesian analysis, from Berger 1984
"""
function robustPriorDemo()
    x = 5
    obsVar = 1

    # First let us verify that the Gaussian prior indeed satisfies the prior quartiles
    priorMu = 0
    priorVar = 2.19

    nprior = Normal(priorMu, sqrt(priorVar))
    @assert isapprox(cdf(nprior, -1), .25; atol=.001)
    @assert isapprox(cdf(nprior, 0) - cdf(nprior, -1), .25; atol=.001)
    @assert isapprox(cdf(nprior, 1) - cdf(nprior, 0), .25; atol=.001)
    @assert isapprox(1-cdf(nprior, 1), .25; atol=.001)
    @assert isapprox(median(nprior), 0; atol=.001)

    # Now compute posterior mean using Gaussian prior
    postVar = 1/(1/obsVar + 1/priorVar)
    postMeanNormal = postVar*(priorMu/priorVar + x/obsVar)
    @assert isapprox(postMeanNormal, 3.43; atol=.01)

    # Now let us do the same thing for the Cauchy
    cprior = Cauchy(0, 1)
    @assert isapprox(cdf(cprior, -1), .25; atol=.001)
    @assert isapprox(cdf(cprior, 0) - cdf(cprior, -1), .25; atol=.001)
    @assert isapprox(cdf(cprior, 1) - cdf(cprior, 0), .25; atol=.001)
    @assert isapprox(1-cdf(cprior, 1), .25; atol=.001)
    @assert isapprox(median(cprior), 0; atol=.001)

    # Now let us compute posterior mean using Cauchy

    lik(θ) = pdf(Normal(θ, sqrt(obsVar)), x)
    prior(θ) = pdf(cprior, θ)
    post(θ) = lik(θ).*prior(θ)

    ∞=10

    Z, err = hquadrature(post, -∞, ∞; reltol=1e-4)

    f(θ) = θ.*post(θ)/Z

    postMeanCauchy, err = hquadrature(f, -∞, ∞; reltol=1e-4)
    @assert isapprox(postMeanCauchy, 4.56; atol=.01)

    (postMeanNormal, postMeanCauchy)
end


function solve_poisson_pcg( Lmat::SparseMatrixCSC{Float64,Int64},
                            prec::SparseMatrixCSC{Float64,Int64},
                            rho::Array{Float64,1}, NiterMax::Int64;
                            verbose=false, TOL=5.e-10 )
  #
  Npoints = size(rho)[1]
  phi = zeros( Float64, Npoints ) # XXX or use some starting guess
  #
  r = zeros( Float64, Npoints )
  p = zeros( Float64, Npoints )
  z = zeros( Float64, Npoints )
  #
  nabla2_phi = Lmat*phi
  r = rho - nabla2_phi
  z = apply_prec_ilu0( prec, r)
  p = copy(z)

  rsold = dot( r, z )

  for iter = 1 : NiterMax
    #
    nabla2_phi = Lmat*p
    #
    alpha = rsold/dot( p, nabla2_phi )
    #
    phi = phi + alpha * p
    r = r - alpha * nabla2_phi
    z = apply_prec_ilu0( prec, r)
    #
    rsnew = dot( z, r )
    deltars = rsold - rsnew
    #println(typeof(rsnew))
    #println(rsnew)
    if verbose
      @printf("%8d %18.10e\n", iter, sqrt(abs(rsnew)))
    end
    #
    if sqrt(abs(rsnew)) < TOL
      if verbose
        @printf("Convergence achieved in solve_poison_pcg: N, iter: %d %d\n", Npoints, iter)
      end
      break
    end
    #
    p = z + (rsnew/rsold) * p
    #
    rsold = rsnew
  end
  #
  return phi
  #
end # of function


function besselj(m, x)
    sn = 0
    an = 1/factorial(m)
    k = 0

    while abs(an) > eps()
        sn += an
        an = -an * x * x/4/(k+1)/(k+m+1)
        k += 1
    end

    sn *= (x/2)^m
    return sn
end

function besselj_plots()
    x = 0:0.02:20
    y1, y2, y3, y4, y5, y6 =[], [], [], [], [], []

    for i in x
        push!(y1, besselj(0, i))
        push!(y2, besselj(1, i))
        push!(y3, besselj(2, i))
        push!(y4, besselj(3, i))
        push!(y5, besselj(4, i))
        push!(y6, besselj(5, i))
    end

    p = plot(x, y1)
    plot!(p, x, y2)
    plot!(p, x, y3)
    plot!(p, x, y4)
    plot!(p, x, y5)
    plot!(p, x, y6)
    display(p)
end


"""
    dct2!(a, fft_type::FFTType=Radix2Type())

Discrete cosine transform of type II

Y = DCT2(X) returns the discrete cosine transform of X,
based on the staggered-grid definition
           N
   Y(k) =  ∑  X(j) cos (π*(k-1)*(j-1/2)/N)
          j=1
The vector Y is the same size as X and contains the
discrete cosine transform coefficients.

Y = DCT2(X,N) pads or truncates the vector X to length N
before transforming.

If X is a matrix, the DCT2 operation is applied to each
column. This transform can be inverted using IDCT2.

Verified correct within machine precision for 2D array input 2/28/2021

Traditional dct2:
          N-1
   Y(k) =  ∑  X(j) cos (π*(j+1/2)*k/N)
          j=0

"""
function dct2!(x, fft_type=Radix2Type())
    x .= dct2(x, fft_type)
end
function dct2(a, fft_type::FFTType=Radix2Type())
    FT = eltype(a)
    if min(size(a)...)==1
        if size(a,2)>1
            do_trans = true;
        else
            do_trans = false;
        end
        a = a[:];
    else
        do_trans = false;
    end
    n = size(a,1);
    m = size(a,2);
    # Pad or truncate a if necessary
    if size(a,1)<n
      aa = zeros(n,m);
      aa[1:size(a,1),:] = a;
    else
      aa = a[1:n,:];
    end
    y = zeros(2*n, m);
    y[1:n,:] .= aa;
    i = Complex(0,1);
    e = 0.5 * exp.(-i*FT(0.5) * π * (0:n-1) ./ n);
    y[n+1:2*n,:] .= reverse(aa[1:n,:]; dims=1);
    yy = fft(y, 1);
    b = zeros(n,m);
    for l = 1:m
        b[1:n,l] .= real.(yy[1:n,l] .* e);
    end
    if isreal(a)
        b = real.(b);
    end
    if do_trans
        b = b'
    end
    return b
end


"""
# Divide-and-conquer algorithms

- Multiplication
- Recurrence relations
    - Master theorem
    - Binary search
- Mergesort
    - An n log n lower bound for sorting
- Medians
    - A randomized divide-and-conquer algorithm for selection
    - Efficiency analysis
    - The Unix sort command
- Matrix multiplication
- The fast Fourier transform
    - Why multiply polynomials?
    - An alternative representation of polynomials
    - Evaluation by divide-and-conquer
    - Interpolation
    - A matrix reformulation
    - Interpolation resolved
    - Inversion formula
    - A closer look at the fast Fourier transform
    - The definitive FFT algorithm
    - The fast Fourier transform unraveled
    - The slow spread of a fast algorithm
"""
module Chapter2

export
    karatsuba,
    karatsuba2,
    mergesort,
    iterativemergesort,
    polynomialmultiplication

"""
    karatsuba(x, y)

Using the Karatsuba algorithm, compute the product of two integers `x` and `y`.

# Examples

```jldoctest
julia> karatsuba(2412, 3231)
7793172
```
"""
function karatsuba(x::Int, y::Int)
    if ndigits(x) == 1 || ndigits(x) == 1
        return x * y
    end

    n = max(ndigits(x), ndigits(y))
    m = n ÷ 2

    xL, xR = x ÷ 10^m, x % 10^m
    yL, yR = y ÷ 10^m, y % 10^m

    P1 = karatsuba(xL, yL)
    P2 = karatsuba(xR, yR)
    P3 = karatsuba((xL + xR), (yL + yR))

    return (P1 * 10^(2m)) + ((P3 - P1 - P2) * 10^m) + P2
end

"""
    karatsuba2(x, y)

Using the Karatsuba algorithm, compute the product of two integers `x` and `y`.

# Examples

```jldoctest
julia> karatsuba(2412, 3231)
7793172
```
"""
function karatsuba2(x::Int, y::Int)
    if ndigits(x, base=2) == 1 || ndigits(y, base=2) == 1
        return x * y
    end

    n = max(ndigits(x, base=2), ndigits(y, base=2))
    m = n ÷ 2

    xL, xR = x ÷ 2^m, x % 2^m
    yL, yR = y ÷ 2^m, y % 2^m

    P1 = karatsuba(xL, yL)
    P2 = karatsuba(xR, yR)
    P3 = karatsuba((xL + xR), (yL + yR))

    return (P1 * 2^n) + ((P3 - P1 - P2) * 2^m) + P2
end

"""
    mergesort(a)

Sort an array of numbers `a` using the merge sort algorithm.

# Examples

```jldoctest
julia> mergesort([10, 2, 5, 3, 7, 13, 1, 6])
8-element Vector{Int64}:
  1
  2
  3
  5
  6
  7
 10
 13

julia> mergesort([1, 4.54312, 3.4321, π])
4-element Vector{Float64}:
 1.0
 3.141592653589793
 3.4321
 4.54312
```
"""
function mergesort(a::AbstractVector)
    if length(a) > 1
        m = length(a) ÷ 2
        return merge(mergesort(a[1:m]), mergesort(a[(m + 1):end]))
    else
        return a
    end
end

function merge(x::AbstractVector, y::AbstractVector)
    if length(x) == 0
        return y
    end
    if length(y) == 0
        return x
    end
    if x[1] <= y[1]
        return pushfirst!(merge(x[2:end], y), x[1])
    else
        return pushfirst!(merge(x, y[2:end]), y[1])
    end
end

function iterativemergesort(a::AbstractVector)
    Q = []
    for i ∈ 1:length(a)
        push!(Q, [a[i]])
    end
    while length(Q) > 1
        push!(Q, merge(eject(Q), eject(Q)))
    end
    return eject(Q)
end

function polynomialmultiplication(A::AbstractVector{T}, B::AbstractVector{T}) where {T <: Real}
    product = zeros(length(A) + length(B))

    for i ∈ 1:length(A)
        for j ∈ 1:length(B)
            product[i + j] += A[i] * B[j]
        end
    end

    return product
end

# TODO: FFT
# function FFT(A, ω)
#
# end

end # module


"""
    maxwell_boltzmann(f)

Maxwell Boltzmann entropy
"""
maxwell_boltzmann(f) = f * log(f) - f

"""
    maxwell_boltzmann_prime(x)

Prim of Maxwell Boltzmann entropy
"""
maxwell_boltzmann_prime(x) = log(x)

"""
    maxwell_boltzmann_dual(f)

Dual of Maxwell Boltzmann entropy
"""
maxwell_boltzmann_dual(f) = exp(f)

"""
    maxwell_boltzmann_dual_prime(f)

Dual prim of Maxwell Boltzmann entropy
"""
maxwell_boltzmann_dual_prime(f) = exp(f)

"""
    kinetic_entropy(α, m, weights)

Reconstruct mathematical entropy from Legendre dual
"""
function kinetic_entropy(α::AbstractArray, m::AbstractArray, weights::AbstractVector)
    B = KitBase.maxwell_boltzmann_dual_prime.(α' * m)[:]
    return sum(maxwell_boltzmann.(B) .* weights)
end