using Bridge, Distributions, StaticArrays
using Plots
using LinearAlgebra

#
t = 2.
n = 100
dt = t / n
x0 = 1.

# define a Float64 Wiener process
P = Wiener() #returns object representing standard Brownian motion

typeof(P)
#prints Bridge.Wiener{Float64}

supertype(typeof(P))
#prints Bridge.ContinuousTimeProcess{Float64}

# sample Brownian motion on an equidistant grid
W = sample(range(0., stop=t, length=n), P)
W = sample(0:dt:t, P) # similar way

println(W.tt)
# prints [0.0,0.03,0.06,..., 0.93.96,0.99]
plot(W.tt, W.yy)
# plots path

plot(W)
# Bridge defines a plot @recipe (via RecipesBase)

# sample complex Brownian motion on a nonequidistant grid

X = sample(sort(rand(1000)), Wiener{Complex{Float64}}())
plot(real(X.yy), imag(X.yy))

# sample a standard Brownian bridge ending in v at time s
s = 1.
v = 0.

B = sample(0:dt:s, WienerBridge(s, v))
plot(W.tt, W.yy, color="blue")
plot!(B.tt, B.yy, color="red")
# displays X and  Brownian bridge B in red

# Define a diffusion process
struct OrnsteinUhlenbeck  <: ContinuousTimeProcess{Float64}
    β::Float64 # drift parameter (also known as inverse relaxation time)
    σ::Float64 # diffusion parameter
    function OrnsteinUhlenbeck(β::Float64, σ::Float64)
        isnan(β) || β > 0. || error("Parameter λ must be positive.")
        isnan(σ) || σ > 0. || error("Parameter σ must be positive.")
        new(β, σ)
    end
end

# define drift and sigma of OrnsteinUhlenbeck
import Bridge: b, σ, a, transitionprob
Bridge.b(t,x, P::OrnsteinUhlenbeck) = -P.β*x
Bridge.σ(t, x, P::OrnsteinUhlenbeck) = P.σ
Bridge.a(t, x, P::OrnsteinUhlenbeck) = P.σ^2

# simulate OrnsteinUhlenbeck using Euler scheme
W = sample(0:0.01:10, Wiener{Float64}())
X = solve(EulerMaruyama(), 0.1, W, OrnsteinUhlenbeck(20., 1.))

plot(X.tt, X.yy)

# define transition density
transitionprob(s, x, t, P::OrnsteinUhlenbeck) = Normal(x*exp(-P.β*(t-s)), sqrt((0.5P.σ^2/P.β) *(1-exp(-2*P.β*(t-s)))))

# plot likelihood of β
LL = [(β, llikelihood(X, OrnsteinUhlenbeck(β, 1.))) for β in 1.:30.]
for (β, ll) in LL
    println("β $β loglikelihood ", ll )
end
plot(Float64[β for (β, ll) in LL], Float64[ll for (β, ll) in LL])

# sample OrnsteinUhlenbeck exactly. compare with euler scheme which degrates as dt = 0.07
X = solve(EulerMaruyama(), 0.1, sample(0:0.07:10, Wiener{Float64}()), OrnsteinUhlenbeck(20., 1.))
X2 = sample(0:0.07:10, OrnsteinUhlenbeck(20., 1.), 0.1)
plot(X.tt, 1 .+ X.yy)
plot!(X2.tt, X2.yy)

# sample vector Brownian motion
W2 = sample(0:0.1:10, Wiener{SVector{4,Float64}}())
llikelihood(W2,Wiener{SVector{4,Float64}}())

P = BridgeProp(OrnsteinUhlenbeck(3., 1.), 0:0.01:1,  (0., 1.), 1.)
X = solve(EulerMaruyama(), 0.1, sample(0:0.01:1, Wiener{Float64}()),P)

P2 = PBridgeProp(OrnsteinUhlenbeck(3., 1.), 0., 0., 1., 2., 2., 0., 1., 0.3, 1.)
Y = solve(EulerMaruyama(), 0.1, sample(0:0.01:2, Wiener{Float64}()),P2)
p = plot(Y.tt, Y.yy, xlim=(0, 2), ylim=(-3,3), legend=false)

for i in 1:100
    global Y = solve(EulerMaruyama(), 0.1, sample(0:0.01:2, Wiener{Float64}()),P2)
    plot!(p, Y.tt, Y.yy, xlim=(0, 2), ylim=(-3,3), linewidth=0.2)
end
display(p)

# Define a diffusion process
struct VOrnsteinUhlenbeck{d}  <: ContinuousTimeProcess{SVector{d,Float64}}
    β # drift parameter (also known as inverse relaxation time)
    σ # diffusion parameter
    function VOrnsteinUhlenbeck{d}(β, σ) where d
           new(β, σ)
    end
end

# define drift and sigma of VOrnsteinUhlenbeck

Bridge.b(t, x, P::VOrnsteinUhlenbeck) = -P.β*x
Bridge.σ(t, x, P::VOrnsteinUhlenbeck) = P.σ*I
Bridge.a(t, x, P::VOrnsteinUhlenbeck) = P.σ*P.σ'*I

# Simulate
tt = 0.:0.003:10.
X = solve(EulerMaruyama(), SVector(0., 0.), sample(tt, Wiener{SVector{2,Float64}}()), VOrnsteinUhlenbeck{2}(3., 1.))

yy = Bridge.mat(X.yy)
plot(yy[1,:], yy[2,:], xlim=(-2, 2), ylim=(-2,2), linewidth=0.5)


using DataArrays
using DataFrames
using Gadfly

function θestimation_()
    # 1970 batting averages for 18 major league players. first column = batting average for first 45 at bats 2nd column = batting everage for remainder of season
    data = [.400 .346
            .378 .298
            .356 .276
            .333 .222
            .311 .273
            .311 .270
            .289 .263
            .267 .210
            .244 .269
            .244 .230
            .222 .264
            .222 .256
            .222 .303
            .222 .264
            .222 .226
            .200 .285
            .178 .316
            .156 .200]

    # data transformation
    y = data[:,1]
    ytest = data[:,2]
    n = 45
    # arcsin transform
    x = sqrt(n).*asin(2.*y-1)

    # shrinkage estimate
    μBar = mean(x)
    V    = sumabs2(x - μBar)
    d    = length(x)
    s2   = V/d
    # by construction of the arcsin transform
    σ2   = 1
    B    = σ2/(σ2 + max(0, s2-σ2))
    μShrunk = μBar + (1-B).*(x-μBar)

    # back transform
    θShrunk = 0.5*(sin(μShrunk/sqrt(n))+1)
    θMLE    = y
    (ytest, θShrunk, θMLE)
end

"""
Baseball Batting Average Shrinkage Estimates
Reproduce example from "Data Analysis Using Stein's Estimator and its Generalizations" Bradley Efron; Carl Morris JASA Vol. 70, No. 350. (Jun., 1975), pp. 311-319.
"""
function shrinkageDemoBaseball()
    θtrue, θShrunk, θMLE = θestimation_()
    # plot Shrinkage Estimates
    p1 = plot([layer(x=[θMLE[i] θShrunk[i]], y=[1 0], Geom.line) for i=1:d]...,
              layer(x=θMLE, y=ones(1, d), Geom.point),
              layer(x=θShrunk, y=zeros(1, d), Geom.point),
              Theme(default_color=colorant"blue"))

    # histograms
    df = DataFrame(PN=[1:d, 1:d, 1:d], Theta=[ytest, θShrunk, θMLE], Type=[ones(d), 2*ones(d), 3*ones(d)])
    p2 = plot(df, y="Theta", x="PN", color="Type", Geom.bar(position=:dodge))

    [p1, p2]
end

function shrinkageMSEBaseball()
    θtrue, θShrunk, θMLE = θestimation_()
    mseMLE = mean((θtrue-θMLE).^2)
    mseShrink = mean((θtrue-θShrunk).^2)
    [mseMLE, mseShrink, mseMLE/mseShrink]
end


@doc doc"""
mlr_beta{Float64}(y::Vector{Float64}, n::Int64=10)

Moving linear regression intercept (column 1) and slope (column 2)
""" ->
function mlr_beta{Float64}(y::Vector{Float64}, n::Int64=10)
    @assert n<length(y) && n>0 "Argument n out of bounds."
    out = zeros(Float64, (length(y),2))
    out[1:n-1,:] = NaN
    xi = collect(1.0:n)
    xbar = mean(xi)
    ybar = runmean(y, n, false)
    @inbounds for i = n:length(y)
        yi = y[i-n+1:i]
        out[i,2] = cov(xi,yi) / var(xi)
        out[i,1] = ybar[i] - out[i,2]*xbar
    end
    return out
end

@doc doc"""
mlr_slope{Float64}(y::Vector{Float64}, n::Int64=10)

Moving linear regression slope
""" ->
function mlr_slope{Float64}(y::Vector{Float64}, n::Int64=10)
    @assert n<length(y) && n>0 "Argument n out of bounds."
    out = zeros(y)
    out[1:n-1] = NaN
    xi = collect(1.0:n)
    @inbounds for i = n:length(y)
        yi = y[i-n+1:i]
        out[i] = cov(xi,yi) / var(xi)
    end
    return out
end

@doc doc"""
mlr_intercept{Float64}(y::Vector{Float64}, n::Int64=10)

Moving linear regression y-intercept
""" ->
function mlr_intercept{Float64}(y::Vector{Float64}, n::Int64=10)
    @assert n<length(y) && n>0 "Argument n out of bounds."
    out = zeros(y)
    out[1:n-1] = NaN
    xi = collect(1.0:n)
    xbar = mean(xi)
    ybar = runmean(y, n, false)
    @inbounds for i = n:length(y)
        yi = y[i-n+1:i]
        out[i] = ybar[i] - xbar*(cov(xi,yi)/var(xi))
    end
    return out
end

@doc doc"""
mlr{Float64}(y::Vector{Float64}, n::Int64=10)

Moving linear regression estimates
""" ->
function mlr{Float64}(y::Vector{Float64}, n::Int64=10)
    b = mlr_beta(y, n)
    return b[:,1] + b[:,2]*float(n)
end

@doc doc"""
mlr_se{Float64}(y::Vector{Float64}, n::Int64=10)

Moving linear regression standard error of estimate
""" ->
function mlr_se{Float64}(y::Vector{Float64}, n::Int64=10)
    yhat = mlr(y, n)
    r = zeros(Float64, n)
    out = zeros(y)
    out[1:n-1] = NaN
    nf = float(n)
    @inbounds for i = n:length(y)
        r = y[i-n+1:i] - yhat[i]
        out[i] = sqrt(sum(r.^2)/nf)
    end
    return out
end

@doc doc"""
mlr_ub{Float64}(y::Vector{Float64}, n::Int64=10, se::Float64=2.0)

Moving linear regression upper bound
""" ->
function mlr_ub{Float64}(y::Vector{Float64}, n::Int64=10, se::Float64=2.0)
    return y + se*mlr_se(y,n)
end

@doc doc"""
mlr_lb{Float64}(y::Vector{Float64}, n::Int64=10, se::Float64=2.0)

Moving linear regression lower bound
""" ->
function mlr_lb{Float64}(y::Vector{Float64}, n::Int64=10, se::Float64=2.0)
    return y - se*mlr_se(y,n)
end

@doc doc"""
mlr_bands{Float64}(y::Vector{Float64}, n::Int64=10, se::Float64=2.0)

Moving linear regression bands

`Output:`

Column 1: Lower bound

Column 2: Regression estimate

Column 3: Upper bound
""" ->
function mlr_bands{Float64}(y::Vector{Float64}, n::Int64=10, se::Float64=2.0)
    out = zeros(Float64, (length(y),3))
    out[1:n-1,:] = NaN
    out[:,2] = mlr(y, n)
    out[:,1] = mlr_lb(y, n, se)
    out[:,1] = mlr_ub(y, n, se)
    return out
end

@doc doc"""
mlr_rsq{Float64}(y::Vector{Float64}, n::Int64=10; adjusted::Bool=false)

Moving linear regression R-squared (and adjusted R-squared)
""" ->
function mlr_rsq{Float64}(y::Vector{Float64}, n::Int64=10; adjusted::Bool=false)
    yhat = mlr(y, n)
    rsq = runcor(y, yhat, n, false) .^ 2
    if adjusted
        return rsq - (1.0-rsq)*(1.0/(float(n)-2.0))
    else
        return rsq
    end
end


#=
A palindromic number reads the same both ways. The largest palindrome made from the product of two 2-digit numbers is 9009 = 91 × 99.

Find the largest palindrome made from the product of two 3-digit numbers.
=#

function calc()
  min = 100; max = 999
  maxPalindrome = 0
  for i in min:max, j in min:max
      current = i * j
      if string(current) == reverse(string(current))
        maxPalindrome = current > maxPalindrome ? current : maxPalindrome
      end
  end
  maxPalindrome
end
@time println(calc())


using Zygote, StaticArrays

# Naive Matrix Multiply

function mul(a::AbstractMatrix, b::AbstractMatrix)
  [sum(a[i, k]*b[k,j] for k = 1:size(a, 2))
   for i = 1:size(a, 1), j = 1:size(b, 2)]
end

a = [1 2; 3 4]

mul(a, a)

# Static Matrix Multiple

_sum(xs) = reduce((a, b) -> :($a+$b), xs)

@generated function mul(a::SMatrix{M,N}, b::SMatrix{N,O}) where {M,N,O}
  :(SMatrix{M,O}(
      $([_sum(:(a[$i, $k]*b[$k,$j]) for k = 1:size(a, 2))
         for i = 1:size(a, 1), j = 1:size(b, 2)]...),))
end

@generated function mul(a::SMatrix{M,N}, b::SMatrix{N,O}) where {M,N,O}
  quote
    Base.@_inline_meta
    SMatrix{M,O}(
      $([_sum(:(a[$i, $k]*b[$k,$j]) for k = 1:size(a, 2))
        for i = 1:size(a, 1), j = 1:size(b, 2)]...),)
  end
end

a = @SArray [1 2; 3 4]

mul(a, a)

sqmul(a) = sum(mul(a, a))

@code_llvm sqmul(a)

# See the code output

function _mul(a::SMatrix{M,N}, b::SMatrix{N,O}) where {M,N,O}
  :(SMatrix{M,O}(
      $([_sum(:(a[$i, $k]*b[$k,$j]) for k = 1:size(a, 2))
         for i = 1:size(a, 1), j = 1:size(b, 2)]...),))
end

_mul(a, a)


export RandomizedKaczmarz, RandomizedBlockKaczmarz

"""
Randomized Kaczmarz

Section 3.3 in Gower, R. M., & Richtárik, P. (2015).
Randomized iterative methods for linear systems.
SIAM Journal on Matrix Analysis and Applications, 36(4), 1660-1690.

RK takes a step in the direction of the negative stochastic gradient.
This means that it is equivalent to the SGD method.
However, the stepsize choice is very special: RK chooses the stepsize which leads to the point
which is closest to x* in the Euclidean norm.
"""
function RandomizedKaczmarz(stp::AbstractStopping; is_zero_start::Bool = false, kwargs...)
  A, b = get_matrix(stp.pb), get_vector(stp.pb)
  state = stp.current_state
  m = size(A, 1)
  if is_zero_start
    state.res .= b
  else
    state.res .= b .- A * state.x
  end
  OK = start!(stp, no_opt_check = true)

  while !OK
    i = Int(floor(rand() * m) + 1)
    Ai = view(A, i, :)
    AiAi = dot(Ai, Ai)
    if AiAi != 0
      Aix = dot(Ai, state.x)
      state.x .-= (Aix - b[i]) / AiAi * Ai
    end
    state.res = b - A * state.x
    OK = cheap_stop!(stp)
  end

  return stp
end

"""
Randomized block Kaczmarz

Section 3.5 in Gower, R. M., & Richtárik, P. (2015).
Randomized iterative methods for linear systems.
SIAM Journal on Matrix Analysis and Applications, 36(4), 1660-1690.
"""
function RandomizedBlockKaczmarz(
  stp::AbstractStopping;
  is_zero_start::Bool = false,
  r::Int = 15,
  rand_r::Bool = false,
  kwargs...,
)
  A, b = get_matrix(stp.pb), get_vector(stp.pb)
  state = stp.current_state
  m = size(A, 1)
  if is_zero_start
    state.res .= b
  else
    state.res .= b .- A * state.x
  end
  OK = start!(stp, no_opt_check = true)

  while !OK
    r = rand_r ? Int(floor(rand() * m) + 1) : min(r, m)
    sub = zeros(Int64, r)
    #sample!(1:m, sub) #x is a subset of [1,...,m] of size r #sample repeat entries
    sub = StatsBase.randperm(m)[1:r]

    Ai = view(A, sub, :)
    resk = state.res[sub]
    if !(Ai == 0)
      state.x .-= Ai' * pinv(Matrix(Ai * Ai')) * resk
    end
    state.res .= A * state.x .- b
    OK = cheap_stop!(stp)
  end

  return stp
end


# # An equation of second order in two dimmensions
#
# Let \\(\Omega = (0, 1)^3\\) the unit cube and denote \\( \Gamma = \partial \Omega\\) its boundary.
# Consider the following problem. Given \\(f \in C(\Omega)\\), find a function \\(u\\) satisfying
# -\Delta u + u = f in \Omega
# u = 0 on \Gamma

using Revise
using SymFEL
using SymPy
using LinearAlgebra
using SparseArrays
using JLD2
using WriteVTK

## discretization parameters
# we use the mesh square.msh (in gmsh format)
# obtained from square.geo using gmsh
# this mesh is formed by quad elements
# we use gmsh module for read the mesh
# there are some problems using Threads and gmsh (in linux)
# the file is square.jld2 is prepared separetely
@load "cube.jld2"

# elementary matrices -
elem_Mx = SymFEL.get_lagrange_em(1, 0, 0)
elem_Kx = SymFEL.get_lagrange_em(1, 1, 1)

nc = ([1, 2, 2, 1, 1, 2, 2, 1],
      [1, 1, 2, 2, 1, 1, 2, 2],
      [1, 1, 1, 1, 2, 2, 2, 2])
nr = nc

elem_Mxyz = SymFEL.get_cube_em(elem_Mx,
                               elem_Mx,
                               elem_Mx,
                               nc,
                               nr)

elem_Kxyz = SymFEL.get_cube_em(elem_Kx,
                               elem_Mx,
                               elem_Mx,
                               nc,
                               nr) +
            SymFEL.get_cube_em(elem_Mx,
                               elem_Kx,
                               elem_Mx,
                               nc,
                               nr) +
            SymFEL.get_cube_em(elem_Mx,
                               elem_Mx,
                               elem_Kx,
                               nc,
                               nr)

dx = norm(nodes_coordinate[:, elements_bound[1,1]] - nodes_coordinate[:, elements_bound[2,1]])

elem_Kxyz_dx = convert(Matrix{Float64}, elem_Kxyz.subs(h, dx))
elem_Mxyz_dx = convert(Matrix{Float64}, elem_Mxyz.subs(h, dx));

# global matrices
K = SymFEL.assemble_cubemesh_FE_matrix(elem_Kxyz_dx, elements_int, order1=1, order2=1)
M = SymFEL.assemble_cubemesh_FE_matrix(elem_Mxyz_dx, elements_int, order1=1, order2=1)

f = (3*pi^2 + 1) * (sin.(pi * nodes_coordinate[1,:])
                    .* sin.(pi * nodes_coordinate[2,:])
                    .* sin.(pi * nodes_coordinate[3,:]))

F = M * f

A = K + M

# boundary condition
tgv = 1e30
A[nodes_boundary, nodes_boundary] += tgv * sparse(Matrix{Float64}(I, nodes_boundary_N, nodes_boundary_N))

u = A \ F
u_exact = sin.(pi*nodes_coordinate[1,:]) .* sin.(pi*nodes_coordinate[2,:]) .* sin.(pi*nodes_coordinate[3,:])
err = u - u_exact

println("L2 error : ", sqrt(err' * M * err))
println("H1 error : ", sqrt(err' * K * err))

# export to vtk
cells = [MeshCell(VTKCellTypes.VTK_HEXAHEDRON, elements_int[1:8, i]) for i = 1:elements_int_N]

points_x = nodes_coordinate[1, :]
points_y = nodes_coordinate[2, :]
points_z = nodes_coordinate[3, :]
vtkfile = vtk_grid("ex5-output", points_x, points_y, points_z, cells)

vtkfile["u", VTKPointData()] = u
outfiles = vtk_save(vtkfile)


"""
    kernel = Skernel(k,σ=1.0)

Construct the scattering kernel functions in momentum space.
The kernel function is used to construct both the energy-damping term, and noise in the energy-damped
stochastic projected Gross-Pitaevskii equation.

`k`: momentum-space grid in where scattering is to be evaluated.

`σ`: width parameter for `k`-space degrees of freedom in the case of N=1,2 spatial dimensions.

External links

[Low-Dimensional Stochastic Projected Gross-Pitaevskii Equation, Bradley, Rooney, MacDonald, Physical Review A 92, 033631 (2015)](https://arxiv.org/abs/1507.02023)

"""
function Skernel1(k,σ=1.0)
    erfcx(abs(k)*σ/sqrt(2))/sqrt(8*π*σ^2)
end
function Skernel2(k,σ=1.0)
    k==0.0 ? 0.0 : exp(abs(k*σ/2)^2)*besselk(0,abs(k*σ/2)^2)/(2*π)
end
function Skernel3(k)
    k==0.0 ? 0.0 : 1/abs(k)
end

Skernel(k::Array{Float64,1},σ=1.0) = Skernel1.(k,σ)
Skernel(k::Array{Float64,2},σ=1.0) = Skernel2.(k,σ)
Skernel(k::Array{Float64,3},σ=1.0) = Skernel3.(k)


#functions for computing mutual informations (in bits)
#po is the distribution where the KL is averaged over
#the KL measures going from px to pxgy
function mutualinformation(py::Vector, px::Vector, pxgy::Matrix)
    card_y = size(py,1)
    card_x = size(px,1)
    if (size(pxgy,1)!=card_x) || (size(pxgy,2)!=card_y)
        error("Dimensionality of p(x|y) does not match p(x), p(y)!")
    end

    MI = 0
    for i in 1:card_y
        MI += py[i] * kl_divergence_bits(vec(pxgy[:,i]),px)
    end

    return MI
end

#functions for computing the conditional mutual information
#I(X;Y|Z)(in bits)
function conditional_mutualinformation(py::Vector, pzgy::Matrix, pxgz::Matrix, pxgzy)
    #I(X;Y|Z) = ∑_y p(y) DKL( p(x|z,y)||p(x|z) )
    #I(X;Y|Z) = ∑_y p(y) ∑_z p(z|y) ∑_x p(x|z,y) log( p(x|z,y)/p(x|z) )

    card_y = length(py)
    card_z = size(pzgy,1)

    MI = 0
    for j in 1:card_y
        for k in 1:card_z
            MI += py[j]*pzgy[k,j] * kl_divergence_bits( vec(pxgzy[:,k,j]), vec(pxgz[:,k]))
        end
    end

    return MI
end

#function for computing the expected utility
#pagw and U_pre must have the same dimensionality
function expectedutility(pw::Vector, pagw::Matrix, U_pre::Matrix)
    #E[U] = ∑_a,w p(w)p(a|w) U(a,w)
    card_w = size(pw,1)
    card_a = size(pagw,1)
    if (size(pagw,2)!=card_w) || (size(U_pre,1)!=card_a) || (size(U_pre,2)!=card_w)
        error("Dimensionality of p(a|w), U_pre(a,w) and p(w) does not match!")
    end

    EU = 0
    for i in 1:card_w
            EU += pw[i] * sum(pagw[:,i] .* U_pre[:,i])
    end

    return EU
end

#This function assumes that the utility function U(a,w) is not a function of the percept o
function expectedutility(pw::Vector, pogw::Matrix, pagow, U_pre::Matrix)
    #E[U] = ∑_a,o,w p(w)p(o|w)p(a|o,w) U(a,w)

    #reuse the function that computes E[U] for ∑_a,w p(w)p(a|w) U(a,w)
    #to do so, compute p(a|w)
    pagw = marginalizeo(pogw,pagow)

    return expectedutility(pw,pagw,U_pre)
end

#Entropy in bits (using entropy from Distributions.jl)
function entropybits(d::Distribution)
    return Distributions.entropy(d,2) #in bits
end

#Entropy in bits for an discrete distribution represented as a vector
function entropybits(p::Vector)
    return entropybits(Categorical(p))
end

#Explicitly provide a function for the log in bits here.
#The rationale behind this is that it makes it easily possible to spot all the
#places (in the code) to change if one wants to use the nats (natural logarithm) intead of bits.
function log_bits(x)
    return log2(x)
end

#Kullback-Leibler divergence in bits
function kl_divergence_bits(p_x::Vector, p0_x::Vector)
    #D_KL_bits = ∑_x p(x) log2 (p(x)/p0(x))

    if sum(p0_x.==0) > 0
        error("Zeros in denominator before kl_divergence computation!")
    end

    kl_div = kl_divergence(p_x, p0_x)/log(2) #using function from Distances.jl

    return kl_div
end

#compute value of rate-distortion objective (avg ΔF)
function RDobjective(EU,I,β)
    return EU-I/β
end

#Objective-value for the three-variable general case
function ThreeVArRDobjective(EU, I_ow, I_ao, I_awgo, β1, β2, β3)
    if(β3==0)
        #sequential case - I(A;W|O) should be zero and since β3=Inf → ∞ * 0 = NaN
        #handle this by using a tiny β3 - if I(A;W|O) is nonzero, this should be easily detectable through the huge
        #value of the objective
        β3 = eps()
    end
    return EU - (1/β1)*I_ow - (1/β2)*I_ao - (1/β3)*I_awgo
end

#compute I(A;W), H(A), H(A|W), E[U] and E[U]-I(A;W)/β
function analyzeBAsolution(pw::Vector, pa::Vector, pagw::Matrix, U_pre::Matrix, β)
    #compute I(A;W)
    I = mutualinformation(pw,pa,pagw)
    #compute H(A)
    Ha = entropybits(pa)
    #compute H(A|W)
    Hagw = Ha-I
    #compute EU
    EU = expectedutility(pw,pagw,U_pre)
    #compute value of objective
    RDobj = RDobjective(EU,I,β)

    return I, Ha, Hagw, EU, RDobj
end

#compute mutual informations, entropies and value of objective for three-variable general case
function analyze_three_var_BAsolution(pw::Vector, po::Vector, pa::Vector, pogw::Matrix,
    pago::Matrix, pagow, pagw::Matrix, U_pre::Matrix, β1, β2, β3)

    #compute I(O;W)
    I_ow = mutualinformation(pw,po,pogw)

    #compute I(A;O)
    I_ao = mutualinformation(po,pa,pago)

    #compute I(A;W|O)
    I_awgo = conditional_mutualinformation(pw, pogw, pago, pagow)

    #compute I(A;W)
    I_aw = mutualinformation(pw,pa,pagw)

    #compute H(O)
    #println("1: ∑p(o) = $(sum(po))")
    Ho = entropybits(po)

    #compute H(A)
    #println("2: ∑p(a) =  $(sum(pa))")
    Ha = entropybits(pa)

    #compute H(O|W)
    Hogw = Ho - I_ow

    #compute H(A|O)
    Hago = Ha - I_ao

    #compute H(A|O,W)
    Hagow = Hago - I_awgo

    #compute H(A|W)
    Hagw = Ha - I_aw

    #compute EU
    EU = expectedutility(pw,pagw,U_pre)

    #compute value of objective
    ThreeVarRDobj = ThreeVArRDobjective(EU, I_ow, I_ao, I_awgo, β1, β2, β3)

    return I_ow, I_ao, I_awgo, I_aw, Ho, Ha, Hogw, Hago, Hagow, Hagw, EU, ThreeVarRDobj
end


module Basis

export basis, nodes

import FastGaussQuadrature: gausslegendre
import Polynomials: Poly

function lagrange(x, w)
  """
  Taken from SciPy 0.18.1:

  Return a Lagrange interpolating polynomial.
  Given two 1-D arrays `x` and `w,` returns the Lagrange interpolating
  polynomial through the points ``(x, w)``.
  Warning: This implementation is numerically unstable. Do not expect to
  be able to use more than about 20 points even if they are chosen optimally.
  Parameters
  ----------
  x : array_like
      `x` represents the x-coordinates of a set of datapoints.
  w : array_like
      `w` represents the y-coordinates of a set of datapoints, i.e. f(`x`).
  Returns
  -------
      The Lagrange interpolating polynomial.
  """
  M = size(x)[1]
  p = Poly([0])
  for j=1:M
    pt = Poly([w[j]])
    for k=1:M
      if k != j
        fac = x[j]-x[k]
        pt *= Poly([-x[k], 1.0]) / fac
      end
    end
    p += pt
  end
  return p
end

function nodes(N)
  # Returns Legendre-Gauss nodes, scaled to [0,1]
  return 0.5 * (1 + gausslegendre(N+1)[1])
end

function basis(N)
  # Returns basis polynomials
  nodeArray = nodes(N)
  return [lagrange(nodeArray, eye(N+1)[:,i]) for i in 1:N+1]
end

end   # module Basis


doc"""
    Gumbel(μ, θ)

The *Gumbel distribution*  with location `μ` and scale `θ` has probability density function

$f(x; \mu, \theta) = \frac{1}{\theta} e^{-(z + e^z)},
\quad \text{ with } z = \frac{x - \mu}{\theta}$

```julia
Gumbel()            # Gumbel distribution with zero location and unit scale, i.e. Gumbel(0, 1)
Gumbel(u)           # Gumbel distribution with location u and unit scale, i.e. Gumbel(u, 1)
Gumbel(u, b)        # Gumbel distribution with location u and scale b

params(d)        # Get the parameters, i.e. (u, b)
location(d)      # Get the location parameter, i.e. u
scale(d)         # Get the scale parameter, i.e. b
```

External links

* [Gumbel distribution on Wikipedia](http://en.wikipedia.org/wiki/Gumbel_distribution)
"""
immutable Gumbel{T<:Real} <: ContinuousUnivariateDistribution
    μ::T  # location
    θ::T  # scale

    Gumbel(μ::T, θ::T) = (@check_args(Gumbel, θ > zero(θ)); new(μ, θ))
end

Gumbel{T<:Real}(μ::T, θ::T) = Gumbel{T}(μ, θ)
Gumbel(μ::Real, θ::Real) = Gumbel(promote(μ, θ)...)
Gumbel(μ::Integer, θ::Integer) = Gumbel(Float64(μ), Float64(θ))
Gumbel(μ::Real) = Gumbel(μ, 1.0)
Gumbel() = Gumbel(0.0, 1.0)

@distr_support Gumbel -Inf Inf

const DoubleExponential = Gumbel

#### Conversions

convert{T <: Real, S <: Real}(::Type{Gumbel{T}}, μ::S, θ::S) = Gumbel(T(μ), T(θ))
convert{T <: Real, S <: Real}(::Type{Gumbel{T}}, d::Gumbel{S}) = Gumbel(T(d.μ), T(d.θ))

#### Parameters

location(d::Gumbel) = d.μ
scale(d::Gumbel) = d.θ
params(d::Gumbel) = (d.μ, d.θ)
@inline partype{T<:Real}(d::Gumbel{T}) = T

#### Statistics

mean(d::Gumbel) = d.μ + d.θ * γ

median{T<:Real}(d::Gumbel{T}) = d.μ - d.θ * log(T(logtwo))

mode(d::Gumbel) = d.μ

var{T<:Real}(d::Gumbel{T}) = T(π)^2/6 * d.θ^2

skewness{T<:Real}(d::Gumbel{T}) = 112*sqrt(T(6))*zeta(T(3))/π/π/π

kurtosis{T<:Real}(d::Gumbel{T}) = T(12)/5

entropy(d::Gumbel) = log(d.θ) + 1 + γ

#### Evaluation

zval(d::Gumbel, x::Real) = (x - d.μ) / d.θ
xval(d::Gumbel, z::Real) = x * d.θ + d.μ

function pdf(d::Gumbel, x::Real)
    z = zval(d, x)
    exp(-z - exp(-z)) / d.θ
end

function logpdf(d::Gumbel, x::Real)
    z = zval(d, x)
    - (z + exp(-z) + log(d.θ))
end

cdf(d::Gumbel, x::Real) = exp(-exp(-zval(d, x)))
logcdf(d::Gumbel, x::Real) = -exp(-zval(d, x))

quantile(d::Gumbel, p::Real) = d.μ - d.θ * log(-log(p))

gradlogpdf(d::Gumbel, x::Real) = - (1 + exp((d.μ - x) / d.θ)) / d.θ

#### Sampling

rand(d::Gumbel) = quantile(d, rand())


# function Mvx(vx_pml::SparseMatrixCSC{Float64,Int64}, rho::Array{Float64,2}, dx::Float64, dt::Float64, ext::Int64, iflag::Int64)
#     rho = modExpand(rho, ext, iflag)
#     (m,n) = size(rho)
#     a1 = 9/8;   a2 = -1/24;
#     c1 = a1/dx; c2 = a2/dx;
#     C3 = zeros(m*n)
#     denum = zeros(m*n)
#     for ix = 1 : n
#         for iz = 1 : m
#             if ix < n
#                a = (rho[iz,ix+1]+rho[iz,ix]) / (2*dt)
#             else
#                a = rho[iz,ix] / dt
#             end
#             b = vx_pml[iz,ix] / 2
#             denum[(ix-1)*m+iz] = 1 / (a+b)
#             C3[(ix-1)*m+iz]= (a-b) / (a+b)
#         end
#     end
#     tmp = spzeros(n, n)
#     tmp[1,1] = -c1; tmp[1,2] = c1; tmp[1,3] = c2;
#     for ix = 2: n-2
#         tmp[ix,ix  ] = -c1; tmp[ix, ix+1] = c1;
#         tmp[ix,ix-1] = -c2; tmp[ix, ix+2] = c2;
#     end
#     tmp[n-1,n-1] = -c1; tmp[n-1, n  ] = c1; tmp[n-1,n-2] = -c2;
#     tmp[n  ,n  ] = -c1; tmp[n  , n-1] =-c2;
#     MvxBvx = spdiagm(C3)
#     MvxBp = spdiagm(denum) * kron(tmp, speye(m))
#     return MvxBvx, MvxBp
# end

function Mvx(vx_pml::SparseMatrixCSC{Float64,Int64}, dx::Float64, dt::Float64, ext::Int64, iflag::Int64)
    (m,n) = size(vx_pml)
    a1 = 9/8;   a2 = -1/24;
    c1 = a1/dx; c2 = a2/dx;
    C3 = zeros(m*n)
    denum = zeros(m*n)
    for ix = 1 : n
        for iz = 1 : m
            a = 1/dt + vx_pml[iz,ix]/2
            b = 1/dt - vx_pml[iz,ix]/2
            denum[(ix-1)*m+iz] = 1 / a
            C3[(ix-1)*m+iz]= b / a
        end
    end
    tmp = spzeros(n, n)
    tmp[1,1] = -c1; tmp[1,2] = c1; tmp[1,3] = c2;
    for ix = 2: n-2
        tmp[ix,ix  ] = -c1; tmp[ix, ix+1] = c1;
        tmp[ix,ix-1] = -c2; tmp[ix, ix+2] = c2;
    end
    tmp[n-1,n-1] = -c1; tmp[n-1, n  ] = c1; tmp[n-1,n-2] = -c2;
    tmp[n  ,n  ] = -c1; tmp[n  , n-1] =-c2;
    MvxBvx = spdiagm(C3)
    MvxBp  = spdiagm(denum) * kron(tmp, speye(m))
    return MvxBvx, MvxBp
end

function Mvx_back(vx_pml::SparseMatrixCSC{Float64,Int64}, dx::Float64, dt::Float64, ext::Int64, iflag::Int64)
    (m,n) = size(vx_pml)
    a1 = 9/8;   a2 = -1/24;
    c1 = a1/dx; c2 = a2/dx;
    C3 = zeros(m*n)
    denum = zeros(m*n)
    for ix = 1 : n
        for iz = 1 : m
            a = 1/dt - vx_pml[iz,ix]/2
            b = 1/dt + vx_pml[iz,ix]/2
            denum[(ix-1)*m+iz] = 1 / a
            C3[(ix-1)*m+iz]= b / a
        end
    end
    tmp = spzeros(n, n)
    tmp[1,1] = -c1; tmp[1,2] = c1; tmp[1,3] = c2;
    for ix = 2: n-2
        tmp[ix,ix  ] = -c1; tmp[ix, ix+1] = c1;
        tmp[ix,ix-1] = -c2; tmp[ix, ix+2] = c2;
    end
    tmp[n-1,n-1] = -c1; tmp[n-1, n  ] = c1; tmp[n-1,n-2] = -c2;
    tmp[n  ,n  ] = -c1; tmp[n  , n-1] =-c2;
    MvxBvx = spdiagm(C3)
    MvxBp  = -spdiagm(denum) * kron(tmp, speye(m))
    return MvxBvx, MvxBp
end


#!/usr/bin/env julia

function main()
    accum = 0
    fib_n  = 1
    fib_nn = 2
    fib_max = 4_000_000
    while fib_n <= fib_max
        if fib_n % 2 == 0
            accum += fib_n
        end
        tmp = fib_n + fib_nn
        fib_n = fib_nn
        fib_nn = tmp
    end
    @printf("The sum of even Fibonacci numbers up to %d is %d\n",
            fib_max, accum)
end

main()


# implement QZ-40 algorithm
# Dellnitz, M., Schütze, O., & Zheng, Q. (2002). Locating all the zeros of an analytic function in one complex variable. Journal of Computational and Applied Mathematics, 138(2), 325–333. https://doi.org/10.1016/S0377-0427(01)00371-5
# currently assumes all roots have multiplicity one

using QuadGK

function winding_number(dlogfdz::Function, rectangle::Array{Float64,1}, maxevals::Int = 10^4)::Int
    (x_min, x_max, y_min, y_max) = rectangle
    (bl, br, tr, tl) = (x_min + y_min*im, x_max + y_min*im, x_max + y_max*im, x_min + y_max*im)
    round(Int,real(1 ./(2*pi*im)*quadgk(dlogfdz,bl,br,tr,tl,bl,rtol=0, atol=0.45*2*pi,maxevals=maxevals)[1]))
end

function qz40(dlogfdz::Function, rectangle::Array{Float64,1}, verbose::Bool=false, maxevals::Int=10^4)

    rectangles = Array{Float64,1}[rectangle]
    zeroes = Complex[]
    wn = -1
    qz40_recursive(dlogfdz,rectangles,zeroes,wn,verbose,maxevals)
end

function locate_root(dlogfdz::Function, rectangle::Array{Float64,1}, maxevals::Int = 10^4)
    (x_min, x_max, y_min, y_max) = rectangle
    (bl, br, tr, tl) = (x_min + y_min*im, x_max + y_min*im, x_max + y_max*im, x_min + y_max*im)
    integrand(z::Complex) = z*dlogfdz(z)
    location = 1 ./(2*pi*im)*quadgk(integrand,bl,br,tr,tl,bl,maxevals=maxevals)[1]
    if (real(location) < x_min || real(location) > x_max || imag(location) < y_min || imag(location) > y_max)
        error("Root location $(location) outside of rectangle (x_min, x_max, y_min, y_max) $(x_min),$(x_max),$(y_min),$(y_max).")
    end
    return location
end

function random_subdivide(rectangle::Array{Float64,1})

    (x_min, x_max, y_min, y_max) = rectangle
    x_mid = x_min + rand()*(x_max - x_min)
    y_mid = y_min + rand()*(y_max - y_min)
    # rectangles = Array{Array{Float64,1},1}[]
    rectangles = Array{Array{Float64,1},1}(undef,0)
    push!(rectangles,[x_min, x_mid, y_min, y_mid]);
    push!(rectangles,[x_mid, x_max, y_min, y_mid]);
    push!(rectangles,[x_mid, x_max, y_mid, y_max]);
    push!(rectangles,[x_min, x_mid, y_mid, y_max]);
    return rectangles
end

function qz40_recursive(dlogfdz::Function, rectangles::Array{Array{Float64,1},1},zeroes::Vector{Complex},nroots::Int,verbose::Bool,maxevals::Int)
    # the rows of rectangles are the delimiters (x_min, x_max, y_min, y_max) of the rectangles
    new_rectangles = Array{Float64,1}[]
    n_rectangles = size(rectangles)[1]
    if (verbose)
        println("Number of rectangles: $(n_rectangles)")
    end
    nroots_new = length(zeroes)
    for i in 1:n_rectangles
        rect = rectangles[i]
        wn = winding_number(dlogfdz,rect,maxevals)
        nroots_new += wn
        if (wn == 1)
            push!(zeroes,locate_root(dlogfdz,rect,maxevals))
        elseif (wn > 1)
            append!(new_rectangles,random_subdivide(rect))
        elseif (wn < 0)
            error("Negative winding numbers")
        end
    end
    if (nroots > 0 && nroots_new != nroots)
        error("Expected number of roots: $(nroots), new number of roots: $(nroots_new).")
    end
    if (nroots == -1)
        println("Number of roots: $(nroots_new)")
    end
    if size(new_rectangles)[1] == 0
        return zeroes
    else
        return qz40_recursive(dlogfdz, new_rectangles, zeroes, nroots_new, verbose,maxevals)
    end
end

function newton(f, dfdz, z_init, n_iter)

    z_current = z_init
    for i in 1:n_iter
        z_current = z_current - f(z_current)/dfdz(z_current)
    end
    z_current
end


export real_matrix, complex_matrix

"""
`real_matrix(x::SimpleQuaternion)` returns a 4-by-4 real matrix that represents
`x`.
"""
function real_matrix(x::SimpleQuaternion)
    a, b, c, d = get_parts(x)
    M = [
        a -b -c -d
        b a -d c
        c d a -b
        d -c b a
    ]
    return M
end

"""
`complex_matrix(x::SimpleQuaternion)` returns a 2-by-2 complex matrix
that represents `x`.
"""
function complex_matrix(x::SimpleQuaternion)
    a, b, c, d = get_parts(x)
    M = [
         a+b*im     c+d*im
        -c+d*im     a-b*im
    ]
    return M
end

function SimpleQuaternion(A::AbstractMatrix{T}) where T
    r,c = size(A)
    if r!=c || !in(r,[2,4])
        error("Matrix must be 2-by-2 or 4-by-4; got $r-by-$c.")
    end

    if r==4
        v = A[:,1]
        x = SimpleQuaternion(v...)
        if real_matrix(x) != A
            error("Matrix not derived from a quaternion")
        end
        return x
    end

    a,b = reim(A[1,1])
    c,d = reim(A[1,2])
    x = SimpleQuaternion(a,b,c,d)
    if complex_matrix(x) != A
        error("Matrix not derived from a quaternion")
    end
    return x
end


abstract type AbstractTrajectory end

"""
Bezier curve
θ: a vector whose length is (N+1); each element can also be vector,
e.g., θ = [[1, 2, 3], [2, 3, 4]]
"""
struct Bezier <: AbstractTrajectory
    θ::AbstractVector
    t0::Real
    tf::Real
end

function (bezier::Bezier)(t::Real)
    @unpack t0, tf, θ = bezier
    N = length(θ) -1
    @assert t <= tf && t >= t0
    s = (t-t0)/(tf-t0)
    basis = 0:N |> Map(i -> binomial(N, i) * s^i * (1-s)^(N-i)) |> collect
    basis .* θ |> sum  # sum of N_C_i * s^(N-i) * (1-s)^i * θ_i
end


### A Pluto.jl notebook ###
# v0.12.7

using Markdown
using InteractiveUtils

# ╔═╡ 6fd616f6-1ea1-11eb-3814-8bfb4a096c49
using Pkg, DrWatson

# ╔═╡ 6ffe2628-1ea1-11eb-24ea-57f985146a72
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Distributions, Random, StatsBase, DataFrames, Plots
	Random.seed!(1)
end;

# ╔═╡ ed174bc4-1ea0-11eb-1e2f-a32874cec549
md"## Listing 3.12"

# ╔═╡ 70179b62-1ea1-11eb-13ce-7fb4d54bad4b
begin
	triangDist = TriangularDist(0,2,1)
	xGrid = 0:0.1:2
		N = 10^6
	inverseSampledData = quantile.(triangDist,rand(N))
end

# ╔═╡ b690890a-2135-11eb-1abc-933c757a93f8
begin
	histogram( inverseSampledData, bins=30, normed=true,
		ylims=(0,1.1), label="Inverse transform data")
	plot!( xGrid, pdf.(triangDist,xGrid), c=:red, lw=4,
		xlabel="x", label="PDF", ylabel = "Density", legend=:topright)
end

# ╔═╡ 70182cd8-1ea1-11eb-094d-8d0d49cf15f3
md"## End of listing 3.12"

# ╔═╡ Cell order:
# ╟─ed174bc4-1ea0-11eb-1e2f-a32874cec549
# ╠═6fd616f6-1ea1-11eb-3814-8bfb4a096c49
# ╠═6ffe2628-1ea1-11eb-24ea-57f985146a72
# ╠═70179b62-1ea1-11eb-13ce-7fb4d54bad4b
# ╠═b690890a-2135-11eb-1abc-933c757a93f8
# ╟─70182cd8-1ea1-11eb-094d-8d0d49cf15f3


"""
    ndgrid(x,y,...)

For d vector inputs, return d matrices representing the coordinate functions on the
tensor product grid.
"""
function ndgrid(x...)
    I = CartesianIndices( fill(undef,length.(x)) )
    [ [ x[d][i[d]] for i in I]  for d = 1:length(x) ]
end

"""
    rectdisc(m,xspan,n,yspan)

Create matrices and helpers for finite-difference discretization of a rectangle that is
the tensor  product of intervals `xspan` and `yspan`, using `m`+1 and `n`+1 points in
the two coordinates.
"""
function rectdisc(m,xspan,n,yspan)
    # Initialize grid and finite differences.
    x,Dx,Dxx = diffmat2(m,xspan)
    y,Dy,Dyy = diffmat2(n,yspan)
    X = repeat(x,outer=(1,n+1))
    Y = repeat(y',outer=(m+1,1))

    # Locate boundary points.
    isbndy = fill(true,m+1,n+1)
    isbndy[2:m,2:n] .= false

    # Get the diff. matrices recognized as sparse. Also include reshaping functions.
    disc = (
        Dx=sparse(Dx), Dxx=sparse(Dxx),
        Dy=sparse(Dy), Dyy=sparse(Dyy),
        Ix=Diagonal(ones(m+1)), Iy=Diagonal(ones(n+1)),
        isbndy=isbndy,
        vec=vec,
        unvec=u -> reshape(u,m+1,n+1)
        )
    return X,Y,disc
end

"""
    poissonfd(f,g,m,xspan,n,yspan)

Solve Poisson's equation on a rectangle by finite differences. Function `f` is the
forcing function and function `g` gives the  Dirichlet boundary condition. The rectangle
is the tensor product of intervals `xspan` and `yspan`,  and the discretization uses
`m`+1 and `n`+1 points in the two coordinates.

Return matrices of the solution values, and the coordinate functions, on the grid.
"""
function poissonfd(f,g,m,xspan,n,yspan)
    # Initialize the rectangle discretization.
    X,Y,d = rectdisc(m,xspan,n,yspan)

    # Form the collocated PDE as a linear system.
    A = kron(d.Iy,d.Dxx) + kron(d.Dyy,d.Ix)  # Laplacian matrix
    b = d.vec(f.(X,Y))

    # Replace collocation equations on the boundary.
    scale = maximum(abs.(A[n+2,:]))
    I = kron(d.Iy,d.Ix)
    A[d.isbndy[:],:] = scale*I[d.isbndy[:],:]                 # Dirichet assignment
    b[d.isbndy[:]] = scale*g.( X[d.isbndy],Y[d.isbndy] )  # assigned values

    # Solve the linear sytem and reshape the output.
    u = A\b
    U = d.unvec(u)
    return U,X,Y
end

"""
    newtonpde(f,g,m,xspan,n,yspan)

Newton's method with finite differences to solve the PDE `f`(u,x,y,disc)=0 on the
rectangle `xspan` ``\times`` `yspan`, subject to `g`(x,y)=0 on the boundary. Use `m`+1
points in x by `n`+1 points in y.

Return matrices of the solution values, and the coordinate functions, on the grid.
"""

function newtonpde(f,g,m,xspan,n,yspan)
    # Discretization.
    X,Y,d = rectdisc(n,xspan,n,yspan)

    # This evaluates the discretized PDE and its Jacobian, with all the
    # boundary condition modifications applied.
    function residual(U)
        R,J = f(U,X,Y,d)
        Ixy = Diagonal(ones(size(J,1)))   # used for row replacements
        scale = maximum(abs.(J))
        J[vec(d.isbndy),:] = scale*Ixy[d.isbndy[:],:]
        XB = X[d.isbndy];  YB = Y[d.isbndy];
        @. R[d.isbndy] = scale*(U[d.isbndy] - g(XB,YB))
        r = d.vec(R)
        return r,J
    end

    # Intialize the Newton iteration.
    U = zeros(size(X))
    r,J = residual(U)
    tol = 1e-10;  itermax = 20;
    s = Inf;  normr = norm(r);  k = 1;

    lambda = 1
    while (norm(s) > tol) && (normr > tol)
        s = -(J'*J + lambda*I) \ Vector(J'*r)  # damped step
        Unew = U + d.unvec(s)
        rnew,Jnew = residual(Unew)

        if norm(rnew) < normr
            # Accept and update.
            lambda = lambda/6;   # dampen the Newton step less
            U = Unew;  r = rnew;  J = Jnew;
            normr = norm(r)
            k = k+1
            println("Norm of residual = $normr")
        else
            # Reject.
            lambda = lambda*4;   # dampen the Newton step more
        end

        if k==itermax
            @warn "Maximum number of Newton iterations reached."
            break
        end
    end
    return U,X,Y
end


# Tools for initial digits stuff

char2digit(c::Char) = c - '0'

function first_digit(x::Real, alert::Bool = false)
    if x == 0
        if alert
            warn("Received 0, returning 0")
        end
        return 0
    end

    x = abs(x) # make sure it's positive

    while x < 1
        x *= 10
    end

    c = first(string(x))
    return char2digit(c)
end

function first_counts(x::Array{T,1}) where {T<:Real}
    digs = map(first_digit, x)
    n = length(digs)
    counts = zeros(Int, 9)
    for d in digs
        if d != 0
            counts[d] += 1
        end
    end
    return counts
end

function first_hists(x::Array{T,1}) where {T<:Real}
    counts = first_counts(x)
    hist = zeros(9)
    S = sum(counts)
    for k = 1:9
        hist[k] = counts[k] / S
    end
    return counts, hist
end

function report(x::Array{T,1}) where {T<:Real}
    counts, hist = first_hists(x)
    for k = 1:9
        println(
            k,
            "\t",
            counts[k],
            "\t",
            round(hist[k] * 100, digits = 1),
            "%\t",
            log(10, k + 1) - log(10, k),
        )
    end
end

function report_array(x::Array{T,1}) where {T<:Real}
    counts, hist = first_hists(x)
    A = Array(Any, (9, 3))
    for k = 1:9
        A[k, 1] = k
        A[k, 2] = counts[k]
        A[k, 3] = round(hist[k] * 100, digits = 2)
    end
    return A
end

function experiment()
    n = 1000
    x = rand(n)
    println("Initial distribution")
    report(x)
    println("\n\n")
    sleep(1)

    count = 0
    while count <= 1000
        m = minimum(x)
        count += 1
        if m < 1
            x *= 100.0
        end
        x = x .* rand(n)
        if count % 25 == 0
            println("step = ", count)
            report(x)
            println("-------------------------------")
            sleep(0.1)
        end
    end
end

function digit_split(n::Int)
    map(char2digit, collect(string(n)))
end

# create a 9^d multiplication table
function mult_table_old(d::Int)
    vals = Int[]
    for n = 10^(d-1):10^d-1
        digs = digit_split(n)
        v = prod(digs)
        # println(digs, " --> ", v)
        if v != 0
            push!(vals, v)
        end
    end
    return vals
end

function mult_table(d::Int)
    tic()
    counts = zeros(Int, 9)
    for n = 10^(d-1):10^d-1
        digs = digit_split(n)
        v = prod(digs)
        if v != 0
            d = first_digit(v)
            counts[d] += 1
        end
    end
    toc()
    return counts
end

mult_report_old(d::Int) = report(mult_table(d))

function mult_report(d::Int)
    counts = mult_table(d)
    N = sum(counts)

    A = Array(Any, (9, 3))
    for k = 1:9
        A[k, 1] = k
        A[k, 2] = counts[k]
        A[k, 3] = round(100 * counts[k] / N, 2)
    end
    return A
end


# Probability computations

"""
	$(SIGNATURES)

Validate a probability matrix. Elements are expected to sum to 1.
"""
function validate_prob_matrix(probM :: AbstractMatrix{T1}) where T1 <: AbstractFloat
    isValid = all_at_least(probM, 0.0)  &&  all_at_most(probM, 1.0);
    isValid = isValid  &&  isapprox(sum(probM), 1.0);
    return isValid
end

"""
	$(SIGNATURES)

Validate a probability vector. Elements are expected to sum to 1.
"""
validate_prob_vector(probV :: AbstractVector{T1}) where T1 <: AbstractFloat =
    validate_prob_matrix(probV');

"""
	$(SIGNATURES)

Conditional probabilities of rows or columns given a probability matrix.
"""
prob_j(m :: AbstractMatrix{T1}) where T1 <: AbstractFloat = vec(sum(m, dims = 2));
prob_j(m :: AbstractMatrix{T1}, j :: Integer) where T1 <: AbstractFloat =
    sum(m[j,:]);
prob_j(m :: AbstractMatrix{T1}, j) where T1 <: AbstractFloat =
    sum(m[j, :], dims = 2);

"""
	$(SIGNATURES)

Conditional probabilities of rows or columns given a probability matrix.
"""
prob_k(m :: AbstractMatrix{T1}) where T1 <: AbstractFloat = vec(sum(m, dims = 1));
prob_k(m :: AbstractMatrix{T1}, k :: Integer) where T1 <: AbstractFloat =
    sum(m[:, k]);
prob_k(m :: AbstractMatrix{T1}, k) where T1 <: AbstractFloat =
    vec(sum(m[:, k], dims = 1));

"""
    $(SIGNATURES)

Conditional probability (j | k).
"""
prob_j_k(m :: AbstractMatrix, j, k :: Integer) =
    m[j, k] ./ prob_k(m, k);

"""
    $(SIGNATURES)

Conditional probability (k | j).
"""
prob_k_j(m :: AbstractMatrix, k, j :: Integer) =
    m[j, k] ./ prob_j(m, j);

## -----------  Expected values

"""
	$(SIGNATURES)

Compute expected value of `x`, given row index `j`. Based on matrix of values `xM[j,k]` and matrix of probabilities (sum to 1) `prM[j, k]`.
"""
ev_given_j(x :: AbstractMatrix, prM :: AbstractMatrix, j :: Integer) =
    sum([(prob_k_j(prM, k, j) * x[j, k])  for k = 1 : size(x, 2) ]);
ev_given_j(x :: AbstractMatrix, prM :: AbstractMatrix) =
    [ev_given_j(x, prM, j)  for j = 1 : size(x, 1)];

"""
	$(SIGNATURES)

Compute expected value of `x`, given column index `k`. Based on matrix of values `xM[j,k]` and matrix of probabilities (sum to 1) `prM[j, k]`.
"""
ev_given_k(x :: AbstractMatrix, prM :: AbstractMatrix, k :: Integer) =
    sum([(prob_j_k(prM, j, k) * x[j, k])  for j = 1 : size(x, 1) ]);
ev_given_k(x :: AbstractMatrix, prM :: AbstractMatrix) =
    [ev_given_k(x, prM, k)  for k = 1 : size(x, 2)];

    """
    $(SIGNATURES)

Given an array of probabilities: ensure that all are in [0, 1].
Error if bounds violation larger than rounding errors.
Make sure that sum does not exceed an upper bound.
"""
function scale_prob_array!(m :: AbstractArray{F1}; maxSum :: F1 = one(F1)) where F1 <: AbstractFloat

    fSmall = F1(.0000001);
    pSum = sum(m);
    @assert (pSum < one(F1) + fSmall)  "Sum too large: $pSum"
    @assert all_at_least(m, -fSmall)  "Negative probabilities"
    @assert all_at_most(m, one(F1) + fSmall)  "Probabilities above 1"

    if any(x -> x < zero(F1), m)
        m[findall(x -> x < zero(F1), m)] .= zero(F1);
    end
    if pSum > maxSum
        m .*= ((maxSum - fSmall) / pSum);
    end
end

# ---------


## GeomUtils.jl --- constructs for geometry

export distance, anglespan, dihedral
export centroid, RMSD

"""
    distance(a, b) -> Real

Calculate Euclidean distance from `a` to `b`.
"""
distance(a, b) = norm(a-b)

"""
    anglespan(a, b, c) ∈ [0, π]

Calculate smaller angle between two vectors ``AB`` and ``BC`` meeting at ``B``.
"""
function anglespan(a, b, c)
    Δ1 = a - b; n1 = norm(Δ1)
    Δ2 = c - b; n2 = norm(Δ2)
    (n1 == 0 || n2 == 0) && error("two positions overlap")
    cos_ = Δ1⋅Δ2 / norm(Δ1) / norm(Δ2)
    cos_ >= 1 && return 0        ## handle "overflow?" cases where cos_ ∉ [-1,1]
    cos_ <= -1 && return π
    acos(cos_)
end

"""
    dihedral(a, b, c, d) ∈ [-π, π]

Calculate right handed dihedral (torsion) angle between the plane ``ABC`` and
``BCD`` intersecting along ``BC``. ``ABC`` is determined by the first three
arguments as points, while ``BCD`` is determined by the last three.
"""
function dihedral(a, b, c, d)
    Δ1 = a - b
    Δ2 = b - c; n2 = norm(Δ2)
    Δ3 = d - c
    n2 == 0 && error("torque axis is undefined; `b`, `c` overlap")

    Δ23 = Δ2×Δ3
    Δ21 = Δ2×Δ1
    (norm(Δ23) == 0 || norm(Δ21) == 0) && error("some points overlap/colinear")

    atan(Δ23⋅(Δ21×Δ2) / n2, Δ23⋅Δ21)
end

"""
    centroid(collection) -> eltype(collection)

Get the centroid of a collection of points. The preferred point representation
is a subtype of `StaticVector`. This centroid is an _unweighted_ average of
positions.
"""
function centroid(collection::Vector{StaticVector})
    sm_ = sum(collection)
    sm_ / length(collection)
end

"""
    centered(collection) -> centered positions

Get a copy of the centered positions in collection. The preferred point
representation is a subtype of `StaticVector`.

Equivalent to `collection .- centroid(collection)`.
"""
centered(collection) = collection .- centroid(collection)

"""
    RMSD(cA, cB) -> Real

Calculate root-mean-square deviation of points in `cA` from those in `cB`.
Preferred point representation is a subtype of `StaticVector`.
"""
function RMSD(cA::Vector{M}, cB::Vector{M}) where M <: StaticVector
    len = length(cA)
    @assert length(cB) == len "arguments don't have the same length"
    msd = sum(cA, cB; init=zero(eltype(M))) do (a, b)
        distance(a, b)
    end / len

    sqrt(msd)
end


using Base: IEEEFloat

#=

  - error-free transformations (`two_sum`, `two_diff`, `two_square`, `two_prod`,
                                `two_hilo_sum`, `two_lohi_sum`, `two_hilo_diff`, `two_lohi_diff`)
  - least-error transformations (`two_sqrt`, `two_inv`, `two_div`)

 nomenclature:
 - `two_<op>`: the "two" refers to the number of values returned
    two and three argument versions of `two_[sum,diff,prod]` are given
 - `two_hilo_<op>` the argments are ordered by non-increasing absolute value
    - this is `fast_two_sum` in the literature
 - `two_lohi_<op>` the argments are ordered by non-decreasing absolute value
    - this is not available in the literature

=#

"""
    two_sum(a, b)

Computes `hi = fl(a+b)` and `lo = err(a+b)`.
- Unchecked Precondition: !(isinf(a) | isinf(b))
"""
@inline function two_sum(a::T, b::T) where {T}
    hi = a + b
    v  = hi - a
    lo = (a - (hi - v)) + (b - v)
    return hi, lo
end

"""
   two_sum(a, b, c)

Computes `hi = fl(a+b+c)` and `lo = err(a+b+c)`.
- Unchecked Precondition: !(isinf(a) | isinf(b) | isinf(c))
"""
function two_sum(a::T, b::T, c::T) where {T}
    a, b, c = magnitude_mintomax(a, b, c)
    md, lo = two_sum(b, c)
    hi, md = two_sum(a, md)
    hi, lo = two_hilo_sum(hi, md+lo)
    return hi, lo
end

"""
    two_diff(a, b)

Computes `s = fl(a-b)` and `e = err(a-b)`.
- Unchecked Precondition: !(isinf(a) | isinf(b))
"""
@inline function two_diff(a::T, b::T) where {T}
    hi = a - b
    v  = hi - a
    lo = (a - (hi - v)) - (b + v)
    return hi, lo
end

"""
    two_diff(a, b, c)

Computes `s = fl(a-b-c)` and `e1 = err(a-b-c), e2 = err(e1)`.
- Unchecked Precondition: !(isinf(a) | isinf(b) | isinf(c))
"""
function two_diff(a::T,b::T,c::T) where {T}
    a, b, c = magnitude_maxtomin(a, b, c)
    s, t = two_diff(-b, c)
    hi, u = two_sum(a, s)
    hi, lo = two_hilo_sum(hi, u+t)
    return hi, lo
end

"""
    two_hilo_sum(a, b)

*unchecked* requirement `|a| ≥ |b|`
Computes `hi = fl(a+b)` and `lo = err(a+b)`.
- Unchecked Precondition: !(isinf(a) | isinf(b))
"""
@inline function two_hilo_sum(a::T, b::T) where {T}
    hi = a + b
    lo = b - (hi - a)
    return hi, lo
end

"""
    two_lohi_sum(a, b)

*unchecked* requirement `|b| ≥ |a|`
Computes `hi = fl(a+b)` and `lo = err(a+b)`.
- Unchecked Precondition: !(isinf(a) | isinf(b))
"""
@inline function two_lohi_sum(a::T, b::T) where {T}
    hi = b + a
    lo = a - (hi - b)
    return hi, lo
end

"""
    two_hilo_diff(a, b)

*unchecked* requirement `|a| ≥ |b|`
Computes `hi = fl(a-b)` and `lo = err(a-b)`.
- Unchecked Precondition: !(isinf(a) | isinf(b))
"""
@inline function two_hilo_diff(a::T, b::T) where {T}
    hi = a - b
    lo = (a - hi) - b
    hi, lo
end

"""
    two_lohi_diff(a, b)

*unchecked* requirement `|b| ≥ |a|`
Computes `hi = fl(a-b)` and `lo = err(a-b)`.
- Unchecked Precondition: !(isinf(a) | isinf(b))
"""
@inline function two_lohi_diff(a::T, b::T) where {T}
    hi = b - a
    lo = (b - hi) - a
    hi, lo
end

"""
    two_square(a)

Computes `hi = fl(a*a)` and `lo = fl(err(a*a))`.
- Unchecked Precondition: !(isinf(a))
"""
@inline function two_square(a::T) where {T}
    hi = a * a
    lo = fma(a, a, -hi)
    hi, lo
end

"""
    two_prod(a, b)

Computes `hi = fl(a*b)` and `lo = fl(err(a*b))`.
- Unchecked Precondition: !(isinf(a) | isinf(b))
"""
@inline function two_prod(a::T, b::T) where {T}
    hi = a * b
    lo = fma(a, b, -hi)
    hi, lo
end

"""
    two_prod(a, b, c)

Computes `hi = fl(a*b*c)` and `lo = err(a*b*c)`.
- Unchecked Precondition: !(isinf(a) | isinf(b) | isinf(c))
"""
function two_prod(a::T, b::T, c::T) where {T}
    abhi, ablo = two_prod(a, b)
    hi, abhiclo = two_prod(abhi, c)
    ablochi, abloclo = two_prod(ablo, c)
    lo = ablochi + (abhiclo + abloclo)
    return hi, lo
end

"""
    two_inv(a)

Computes `hi = fl(inv(a))` and `lo = err(inv(a))`.
- Unchecked Precondition: !(isinf(a))
"""
@inline function two_inv(a::T) where {T}
     hi = inv(a)
     lo = fma(-hi, a, one(T))
     lo /= a
     return hi, lo
end

"""
    two_div(a, b)

Computes `hi = fl(a/b)` and `lo = err(a/b)`.
- Unchecked Precondition: !(isinf(a) | isinf(b))
"""
@inline function two_div(a::T, b::T) where {T}
     hi = a / b
     lo = fma(-hi, b, a)
     lo /= b
     return hi, lo
end

"""
    two_sqrt(a)

Computes `hi = fl(sqrt(a))` and `lo = err(sqrt(a))`.
- Unchecked Precondition: !(isinf(a))
"""
@inline function two_sqrt(a::T) where {T}
    hi = sqrt(a)
    lo = fma(-hi, hi, a)
    lo /= 2
    lo /= hi
    return hi, lo
end

Base.:(+)(xhi::T, xlo::T, y::T) where {T<:IEEEFloat} = two_sum(xhi, y, xlo)
Base.:(-)(xhi::T, xlo::T, y::T) where {T<:IEEEFloat} = two_sum(xhi, xlo, -y)

Base.:(+)(x::Tuple{T,T}, y::T) where {T<:IEEEFloat} = (+)(x[1], x[2], y)
Base.:(-)(x::Tuple{T,T}, y::T) where {T<:IEEEFloat} = (+)(x[1], x[2], y)

# Algorithm 12 from Tight and rigourous error bounds.  relative error <= 5u²
function Base.:(*)(xhi::T, xlo::T, y::T) where {T<:IEEEFloat}
    hi, lo = two_prod(xhi, y)
    t = fma(xlo, y, lo)
    hi, lo = two_hilo_sum(hi, t)
    return hi, lo
end

function Base.:(/)(xhi::T, xlo::T, y::T) where {T<:IEEEFloat}
    hi = xhi / y
    uh, ul = two_prod(hi, y)
    lo = ((((xhi - uh) - ul) + xlo))/y
    hi,lo = two_hilo_sum(hi, lo)
    return hi, lo
end

Base.:(*)(x::Tuple{T,T}, y::T) where {T<:IEEEFloat} = (*)(x[1], x[2], y)
Base.:(/)(x::Tuple{T,T}, y::T) where {T<:IEEEFloat} = (/)(x[1], x[2], y)

# exchange sort for magnitudes, 2 or 3 values

@inline magnitude_minmax(a, b) = abs(a) < abs(b) ? (a, b) : (b, a)

@inline function magnitude_mintomax(a, b, c)
    b, c = magnitude_minmax(b, c)
    a, c = magnitude_minmax(a, c)
    a, b = magnitude_minmax(a, b)
    return a, b, c
end


using PyPlot, SparseArrays
rows = [2,3,4,4,5,5,5,6,6,7,7,8,8,9,9,10,11]
cols = [3,2,1,2,2,4,6,2,5,2,5,2,5,2,5,5,5]
A = sparse(rows, cols, 1, 11, 11)
A[sum(A,dims=2)[:] .== 0,:] .= 1  # If a vertex has no outbound edges, all edges to all other vertices
spy(A, marker=".");

function pagerank(A, d=0.85, tol=1e-6)
    N = size(A,2)
    L = sum(A, dims=2)
    M = (A ./ L)'
    R = ones(N) / N
    while true
        newR = d * (M * R) .+ (1 - d) / N
        if maximum(abs.(newR - R)) < tol
            return newR
        end
        R = newR
    end
end

R = pagerank(A, 0.85)
round.(100R, digits=1) # Show percentages


export substellarlatitude, hourangle
export diurnalfluxfactor, diurnalfluxfactors
export annualfluxfactor, annualfluxfactors

"""
    substellarlatitude(f, γ)

Compute the latitude of the substellar point for a given solar longitude `f` (true anomaly) and obliquity `γ`
"""
substellarlatitude(f, γ) = asin(cos(f)*sin(γ))

"""
    hourangle(θ, θₛ)

Compute the [hour angle](https://en.wikipedia.org/wiki/Hour_angle)
"""
function hourangle(θ, θₛ)
    x = -sin(θ)*sin(θₛ)/(cos(θ)*cos(θₛ))
    if x <= -1
        return π
    elseif x >= 1
        return 0.0
    end
    return acos(x)
end

#θ - latitude
#θₛ - substellar latitude
"""
    diurnalfluxfactor(θ, θₛ)

Compute the diurnally averaged fraction of incoming stellar flux received by a point at latitude `θ` when the substellar latitude is `θₛ`
"""
function diurnalfluxfactor(θ, θₛ)
    h = hourangle(θ, θₛ)
    return (sin(h)*cos(θ)*cos(θₛ) + h*sin(θ)*sin(θₛ))/π
end

#θ - latitude
#f - solar longitude
#γ - obliquity
"""
    diurnalfluxfactor(θ, f, γ)

Compute the diurnally averaged fraction of incoming stellar flux received by a point at latitude `θ` when the planet is at solar longitude (true anomaly) `f`, with obliquity `γ`
"""
diurnalfluxfactor(θ, f, γ) = diurnalfluxfactor(θ, substellarlatitude(f, γ))

"""
    diurnalfluxfactor(t, a, m, e, θ, γ, p)

Compute the diurnally averaged fraction of incoming stellar flux received by a point at latitude `θ` for a general elliptical orbit
"""
function diurnalfluxfactor(t, a, m, e, θ, γ, p)
    f = trueanomaly(t, a, m, e)
    r = orbitaldistance(a, f, e)
    return diurnalfluxfactor(θ, f - p, γ)*(a/r)^2
end

"""
    diurnalfluxfactors(γ; nf=251, nθ=181)

Compute a grid of diurnally averaged fractions of incoming stellar flux received by a point at latitude `θ` for a planet with obliquity `γ` in a circular orbit. Returns  a solar longitude vector (column values), latitude vector (row values), and the grid of flux factors. `nf` indicates the number of points around the orbit and `nθ` indicates the number of latitudes.
"""
function diurnalfluxfactors(γ; nf::Int=251, nθ::Int=181)
    θ = LinRange(-π/2, π/2, nθ)
    f = LinRange(0, 2π, nf)
    F, Θ = meshgrid(f, θ)
    return (f, θ, diurnalfluxfactor.(Θ, F, γ))
end

"""
    diurnalfluxfactors(a, m, e, γ, p; nt=251, nθ=181)

Compute a grid of diurnally averaged fractions of incoming stellar flux for a planet in a general elliptical orbit. Returns a time vector (column values) over one orbital period, latitude vector (row values), and the grid of flux factors. `nt` indicates the number of time samples around the orbit and `nθ` indicates the number of latitudes.
"""
function diurnalfluxfactors(a, m, e, γ, p; nt::Int=251, nθ::Int=181)
    t = LinRange(0, orbitalperiod(a, m), nt)
    θ = LinRange(-π/2, π/2, nθ)
    T, Θ = meshgrid(t, θ)
    return (t, θ, diurnalfluxfactor.(T, a, m, e, Θ, γ, p))
end

"""
    annualfluxfactor(e, θ, γ, p)

Compute the annually averaged flux factor for a latitude `θ` on a planet in a general elliptical orbit.
"""
function annualfluxfactor(e, θ, γ, p; tol::Float64=1e-4)
    T = orbitalperiod(1.0, 1.0)
    f(t) = diurnalfluxfactor(t, 1.0, 1.0, e, θ, γ, p)
    F, _ = hquadrature(f, 0, T, reltol=tol, abstol=tol)
    return F/T
end

"""
    annualfluxfactors(e, γ, p; nθ=181)

Compute a range of annually averaged flux factors for a planet in a general elliptical orbit. Returns a latitude vector (row values) and a vector of flux factors. `nθ` indicates the number of latitude samples.
"""
function annualfluxfactors(e, γ, p; nθ::Int=181)
    θ = LinRange(-π/2, π/2, nθ)
    F = annualfluxfactor.(e, θ, γ, p)
    return θ, F
end


abstract type AbstractODEProblem end

struct ODEProblem{F,T,U,P} <: AbstractODEProblem
    f::F
    tspan::T
    u0::U
    θ::P
end

abstract type ODESolver end
struct Euler{T} <: ODESolver
    dt::T
end
struct RK2{T} <: ODESolver
    dt::T
end

function f(x,θ)
    α, β, γ, δ = θ
    x₁, x₂ = x

    dx₁ = α*x₁ - β*x₁*x₂
    dx₂ = δ*x₁*x₂ - γ*x₂

    [dx₁, dx₂]
end

function solve(prob::AbstractODEProblem, solver::ODESolver)
    t = prob.tspan[1]; u = prob.u0
    us = [u]; ts = [t]
    while t < prob.tspan[2]
        (u,t) = solver(prob, u, t)
        push!(us,u)
        push!(ts,t)
    end
    ts, reduce(hcat,us)
end

function (solver::Euler)(prob::ODEProblem, u, t)
    f, θ, dt  = prob.f, prob.θ, solver.dt
    (u + dt*f(u,θ), t+dt)
end

function (solver::RK2)(prob::ODEProblem, u, t)
    f, θ, dt  = prob.f, prob.θ, solver.dt
    uh = u + f(u,θ)*dt
    u + dt/2*(f(u,θ) + f(uh,θ)), t+dt
end

θ = [0.1,0.2,0.3,0.2]
u0 = [1.0,1.0]
tspan = (0.,100.)
dt = 0.1
prob = ODEProblem(f,tspan,u0,θ)

t,X=solve(prob, RK2(0.2))

# using Plots
# p1 = plot(t, X[1,:], label="x", lw=3)
# plot!(p1, t, X[2,:], label="y", lw=3)

# display(p1)

#
θ = [0.2,0.2,0.3,0.2]
u0 = [1.0,1.0]
tspan = (0.,100.)
dt = 0.1
prob2 = ODEProblem(f,tspan,u0,θ)

t,X2=solve(prob2, RK2(0.2))

using Optim

function loss(θin,prob::ODEProblem,Y)
    prob.θ.=θin
    t,Xn=solve(prob,RK2(0.2))
    sum((Y.-Xn).^2)
end
θopt = copy(θ)
O=Optim.optimize(θ->loss(θ,prob,X),θopt)
O=Optim.optimize(θ->loss(θ,prob,X),θopt,LBFGS())

using DiffEqFlux
nn=FastDense(2,2)
p = initial_params(nn)
nn([1,2],p)

function fy(x,θ)
    α, β, γ, δ, ω = θ
    x₁, x₂ = x

    dx₁ = α*x₁ - β*x₁*x₂ + ω*x₂
    dx₂ = δ*x₁*x₂ - γ*x₂

    [dx₁, dx₂]
end

#
θy = [0.2,0.2,0.3,0.2,0.1]
u0 = [1.0,1.0]
tspan = (0.,100.)
dt = 0.1
proby = ODEProblem(fy,tspan,u0,θy)

t,Xy=solve(proby, RK2(0.2))

# py = plot(t, Xy[1,:], label="x", lw=3)
# plot!(py, t, Xy[2,:], label="y", lw=3)
# savefig("LV_omega.svg")

function fnn(x,θ)
    α, β, γ, δ = θ[1:4]
    x₁, x₂ = x

    dx₁ = α*x₁ - β*x₁*x₂
    dx₂ = δ*x₁*x₂ - γ*x₂

    [dx₁, dx₂]+nn(x,@view θ[5:end])
end

θnn = [0.2,0.2,0.3,0.2,0.01*initial_params(nn)...]
probnn = ODEProblem(fnn,tspan,u0,θnn)

θopt = copy(θnn)
O=Optim.optimize(θ->loss(θ,probnn,Xy),θopt,Optim.Options(iterations=10000))


"""`care(A, B, Q, R)`

Compute 'X', the solution to the continuous-time algebraic Riccati equation,
defined as A'X + XA - (XB)R^-1(B'X) + Q = 0, where R is non-singular.

Algorithm taken from:
Laub, "A Schur Method for Solving Algebraic Riccati Equations."
http://dspace.mit.edu/bitstream/handle/1721.1/1301/R-0859-05666488.pdf
"""
function care(A, B, Q, R)
    G = try
        B*inv(R)*B'
    catch y
        if y isa SingularException
            error("R must be non-singular in care.")
        else
            throw(y)
        end
    end

    Z = [A  -G;
        -Q  -A']

    S = schur(Z)
    S = ordschur(S, real(S.values).<0)
    U = S.Z

    (m, n) = size(U)
    U11 = U[1:div(m, 2), 1:div(n,2)]
    U21 = U[div(m,2)+1:m, 1:div(n,2)]
    return U21/U11
end

"""`dare(A, B, Q, R)`

Compute `X`, the solution to the discrete-time algebraic Riccati equation,
defined as A'XA - X - (A'XB)(B'XB + R)^-1(B'XA) + Q = 0, where Q>=0 and R>0

Algorithm taken from:
Laub, "A Schur Method for Solving Algebraic Riccati Equations."
http://dspace.mit.edu/bitstream/handle/1721.1/1301/R-0859-05666488.pdf
"""
function dare(A, B, Q, R)
    if (!ishermitian(Q) || minimum(eigvals(real(Q))) < 0)
        error("Q must be positive-semidefinite.");
    end
    if (!isposdef(R))
        error("R must be positive definite.");
    end

    n = size(A, 1);

    E = [
        Matrix{Float64}(I, n, n) B/R*B';
        zeros(size(A)) A'
    ];
    F = [
        A zeros(size(A));
        -Q Matrix{Float64}(I, n, n)
    ];

    QZ = schur(F, E);
    QZ = ordschur(QZ, abs.(QZ.alpha./QZ.beta) .< 1);

    return QZ.Z[(n+1):end, 1:n]/QZ.Z[1:n, 1:n];
end

"""`dlyap(A, Q)`

Compute the solution `X` to the discrete Lyapunov equation
`AXA' - X + Q = 0`.
"""
function dlyap(A::AbstractMatrix, Q)
    lhs = kron(A, conj(A))
    lhs = I - lhs
    x = lhs\reshape(Q, prod(size(Q)), 1)
    return reshape(x, size(Q))
end

"""`gram(sys, opt)`

Compute the grammian of system `sys`. If `opt` is `:c`, computes the
controllability grammian. If `opt` is `:o`, computes the observability
grammian."""
function gram(sys::AbstractStateSpace, opt::Symbol)
    if !isstable(sys)
        error("gram only valid for stable A")
    end
    func = iscontinuous(sys) ? lyap : dlyap
    if opt == :c
        # TODO probably remove type check in julia 0.7.0
        return func(sys.A, sys.B*sys.B')#::Array{numeric_type(sys),2} # lyap is type-unstable
    elseif opt == :o
        return func(Matrix(sys.A'), sys.C'*sys.C)#::Array{numeric_type(sys),2} # lyap is type-unstable
    else
        error("opt must be either :c for controllability grammian, or :o for
                observability grammian")
    end
end

"""`obsv(A, C)` or `obsv(sys)`

Compute the observability matrix for the system described by `(A, C)` or `sys`.

Note that checking for observability by computing the rank from `obsv` is
not the most numerically accurate way, a better method is checking if
`gram(sys, :o)` is positive definite."""
function obsv(A::AbstractMatrix, C::AbstractMatrix)
    T = promote_type(eltype(A), eltype(C))
    n = size(A, 1)
    ny = size(C, 1)
    if n != size(C, 2)
        throw(ArgumentError("C must have the same number of columns as A"))
    end
    res = fill(zero(T), n*ny, n)
    res[1:ny, :] = C
    for i=1:n-1
        res[(1 + i*ny):(1 + i)*ny, :] = res[((i - 1)*ny + 1):i*ny, :] * A
    end
    return res
end
obsv(sys::StateSpace) = obsv(sys.A, sys.C)

"""`ctrb(A, B)` or `ctrb(sys)`

Compute the controllability matrix for the system described by `(A, B)` or
`sys`.

Note that checking for controllability by computing the rank from
`ctrb` is not the most numerically accurate way, a better method is
checking if `gram(sys, :c)` is positive definite."""
function ctrb(A::AbstractMatrix, B::AbstractMatrix)
    T = promote_type(eltype(A), eltype(B))
    n = size(A, 1)
    nu = size(B, 2)
    if n != size(B, 1)
        throw(ArgumentError("B must have the same number of rows as A"))
    end
    res = fill(zero(T), n, n*nu)
    res[:, 1:nu] = B
    for i=1:n-1
        res[:, (1 + i*nu):(1 + i)*nu] = A * res[:, ((i - 1)*nu + 1):i*nu]
    end
    return res
end
ctrb(sys::StateSpace) = ctrb(sys.A, sys.B)

"""`P = covar(sys, W)`

Calculate the stationary covariance `P = E[y(t)y(t)']` of the output `y` of a
`StateSpace` model `sys` driven by white Gaussian noise `w` with covariance
`E[w(t)w(τ)]=W*δ(t-τ)` (δ is the Dirac delta).

Remark: If `sys` is unstable then the resulting covariance is a matrix of `Inf`s.
Entries corresponding to direct feedthrough (D*W*D' .!= 0) will equal `Inf`
for continuous-time systems."""
function covar(sys::AbstractStateSpace, W)
    (A, B, C, D) = ssdata(sys)
    if !isa(W, UniformScaling) && (size(B,2) != size(W, 1) || size(W, 1) != size(W, 2))
        error("W must be a square matrix the same size as `sys.B` columns")
    end
    if !isstable(sys)
        return fill(Inf,(size(C,1),size(C,1)))
    end
    func = iscontinuous(sys) ? lyap : dlyap
    Q = try
        func(A, B*W*B')
    catch
        error("No solution to the Lyapunov equation was found in covar")
    end
    P = C*Q*C'
    if iscontinuous(sys)
        #Variance and covariance infinite for direct terms
        direct_noise = D*W*D'
        for i in 1:size(C,1)
            if direct_noise[i,i] != 0
                P[i,:] .= Inf
                P[:,i] .= Inf
            end
        end
    else
        P += D*W*D'
    end
    return P
end

covar(sys::TransferFunction, W) = covar(ss(sys), W)

# Note: the H∞ norm computation is probably not as accurate as with SLICOT,
# but this seems to be still reasonably ok as a first step
"""
`..  norm(sys, p=2; tol=1e-6)`

`norm(sys)` or `norm(sys,2)` computes the H2 norm of the LTI system `sys`.

`norm(sys, Inf)` computes the H∞ norm of the LTI system `sys`.
The H∞ norm is the same as the H∞ for stable systems, and Inf for unstable systems.
If the peak gain frequency is required as well, use the function `hinfnorm` instead.
See [`hinfnorm`](@ref) for further documentation.

`tol` is an optional keyword argument, used only for the computation of L∞ norms.
It represents the desired relative accuracy for the computed L∞ norm
(this is not an absolute certificate however).

`sys` is first converted to a `StateSpace` model if needed.
"""
function LinearAlgebra.norm(sys::AbstractStateSpace, p::Real=2; tol=1e-6)
    if p == 2
        return sqrt(tr(covar(sys, I)))
    elseif p == Inf
        return hinfnorm(sys; tol=tol)[1]
    else
        error("`p` must be either `2` or `Inf`")
    end
end
LinearAlgebra.norm(sys::TransferFunction, p::Real=2; tol=1e-6) = norm(ss(sys), p, tol=tol)

"""
`   (Ninf, ω_peak) = hinfnorm(sys; tol=1e-6)`

Compute the H∞ norm `Ninf` of the LTI system `sys`, together with a frequency
`ω_peak` at which the gain Ninf is achieved.

`Ninf := sup_ω σ_max[sys(iω)]`  if `G` is stable (σ_max = largest singular value)
      :=        `Inf'           if `G` is unstable

`tol` is an optional keyword argument for the desired relative accuracy for
the computed H∞ norm (not an absolute certificate).

`sys` is first converted to a state space model if needed.

The continuous-time L∞ norm computation implements the 'two-step algorithm' in:\\
**N.A. Bruinsma and M. Steinbuch**, 'A fast algorithm to compute the H∞-norm of
a transfer function matrix', Systems and Control Letters (1990), pp. 287-293.

For the discrete-time version, see:\\
**P. Bongers, O. Bosgra, M. Steinbuch**, 'L∞-norm calculation for generalized
state space systems in continuous and discrete time', American Control Conference, 1991.

See also [`linfnorm`](@ref).
"""
function hinfnorm(sys::AbstractStateSpace; tol=1e-6)
    if iscontinuous(sys)
        return _infnorm_two_steps_ct(sys, :hinf, tol)
    else
        return _infnorm_two_steps_dt(sys, :hinf, tol)
    end
end
hinfnorm(sys::TransferFunction; tol=1e-6) = hinfnorm(ss(sys); tol=tol)

"""
`   (Ninf, ω_peak) = linfnorm(sys; tol=1e-6)`

Compute the L∞ norm `Ninf` of the LTI system `sys`, together with a frequency
`ω_peak` at which the gain `Ninf` is achieved.

`Ninf := sup_ω σ_max[sys(iω)]` (σ_max denotes the largest singular value)

`tol` is an optional keyword argument representing the desired relative accuracy for
the computed L∞ norm (this is not an absolute certificate however).

`sys` is first converted to a state space model if needed.

The continuous-time L∞ norm computation implements the 'two-step algorithm' in:\\
**N.A. Bruinsma and M. Steinbuch**, 'A fast algorithm to compute the H∞-norm of
a transfer function matrix', Systems and Control Letters (1990), pp. 287-293.

For the discrete-time version, see:\\
**P. Bongers, O. Bosgra, M. Steinbuch**, 'L∞-norm calculation for generalized
state space systems in continuous and discrete time', American Control Conference, 1991.

See also [`hinfnorm`](@ref).
"""
function linfnorm(sys::AbstractStateSpace; tol=1e-6)
    if iscontinuous(sys)
        return _infnorm_two_steps_ct(sys, :linf, tol)
    else
        return _infnorm_two_steps_dt(sys, :linf, tol)
    end
end
linfnorm(sys::TransferFunction; tol=1e-6) = linfnorm(ss(sys); tol=tol)

function _infnorm_two_steps_ct(sys::AbstractStateSpace, normtype::Symbol, tol=1e-6, maxIters=250, approximag=1e-10)
    # norm type :hinf or :linf the reason that to not use `hinfnorm(sys) = isstable(sys) : linfnorm ? (Inf, Nan)`
    # is to avoid re computing the poles and return the peak frequencies for, e.g., 1/(s^2 + 1)
    # `maxIters`: the maximum  number of iterations allowed in the algorithm (default 1000)
    # approximag is a tuning parameter: what does it mean for a number to be on the imaginary axis
    # Because of this tuning for example, the relative precision that we provide on the norm computation
    # is not a true guarantee, more an order of magnitude
    # outputs: An approximatation of the L∞ norm and the frequency ω_peak at which it is achieved
    # QUESTION: The tolerance for determining if there are poles on the imaginary axis
    # would not be very appropriate for systems with slow dynamics?
    T = promote_type(real(numeric_type(sys)), Float64)

    on_imag_axis = z -> abs(real(z)) < approximag # Helper fcn for readability

    if sys.nx == 0  # static gain
        return (opnorm(sys.D), T(0))
    end

    pole_vec = pole(sys)

    # Check if there is a pole on the imaginary axis
    pidx = findfirst(on_imag_axis, pole_vec)
    if !(pidx isa Nothing)
        return (T(Inf), imag(pole_vec[pidx]))
        # note: in case of cancellation, for s/s for example, we return Inf, whereas Matlab returns 1
    end

    if normtype == :hinf && any(z -> real(z) > 0, pole_vec)
        return T(Inf), T(NaN) # The system is unstable
    end

    # Initialization: computation of a lower bound from 3 terms
    if isreal(pole_vec)  # only real poles
        ω_p = minimum(abs.(pole_vec))
    else  # at least one pair of complex poles
        maxidx = argmax([abs(imag(p)/real(p))/abs(p) for p in pole_vec])
        ω_p = abs(pole_vec[maxidx])
    end

    m_vec_init = [0, ω_p, Inf]

    (lb, idx) = findmax([opnorm(evalfr(sys, im*m_vec_init[1]));
                         opnorm(evalfr(sys, im*m_vec_init[2]));
                         opnorm(sys.D)])
    ω_peak = m_vec_init[idx]

    # Iterations
    for iter=1:maxIters
        gamma = (1+2*T(tol))*lb
        R = sys.D'*sys.D - gamma^2*I
        S = sys.D*sys.D' - gamma^2*I
        M = sys.A-sys.B*(R\sys.D')*sys.C
        H = [         M              -gamma*sys.B*(R\sys.B') ;
               gamma*sys.C'*(S\sys.C)            -M'            ]

        Λ = complex(eigvals(H)) # To make type stable

        if numeric_type(sys) <: Real
            # Only need to consider one eigenvalue in each complex-conjugate pairs
            filter!(z -> imag(z) >= 0, Λ)
        end

        # Find eigenvalues on the imaginary axis
        Λ_on_imag_axis = filter(on_imag_axis, Λ)

        ω_vec = imag.(Λ_on_imag_axis)

        sort!(ω_vec)

        if isempty(ω_vec)
            return (1+T(tol))*lb, ω_peak
        end

        # Improve the lower bound
        # if not empty, ω_vec contains at least two values
        for k=1:length(ω_vec)-1
            mk = (ω_vec[k] + ω_vec[k+1])/2
            sigmamax_mk = opnorm(evalfr(sys,mk*1im))
            if sigmamax_mk > lb
                lb = sigmamax_mk
                ω_peak = mk
            end
        end
    end
    error("In _infnorm_two_steps_dt: The computation of the H∞/L∞ norm did not converge in $maxIters iterations")
end

function _infnorm_two_steps_dt(sys::AbstractStateSpace, normtype::Symbol, tol=1e-6, maxIters=250, approxcirc=1e-8)
    # Discrete-time version of linfnorm_two_steps_ct above
    # Compuations are done in normalized frequency θ

    on_unit_circle = z -> abs(abs(z) - 1) < approxcirc # Helper fcn for readability

    T = promote_type(real(numeric_type(sys)), Float64)

    if sys.nx == 0  # static gain
        return (opnorm(sys.D), T(0))
    end

    pole_vec = pole(sys)

    # Check if there is a pole on the unit circle
    pidx = findfirst(on_unit_circle, pole_vec)
    if !(pidx isa Nothing)
        return (T(Inf), angle(pole_vec[pidx])/T(sys.Ts))
    end

    if normtype == :hinf && any(z -> abs(z) > 1, pole_vec)
        return T(Inf), T(NaN) # The system is unstable
    end

    # Initialization: computation of a lower bound from 3 terms

    if isreal(pole_vec)  # not just real poles
        # find frequency of pôle closest to unit circle
        θ_p = angle(pole_vec[argmin(abs.(abs.(pole_vec).-1))])
    else
        θ_p = T(pi)/2
    end

    if isreal(pole_vec)  # only real poles
        ω_p = minimum(abs.(pole_vec))
    else  # at least one pair of complex poles
        maxidx = argmax([abs(imag(p)/real(p))/abs(p) for p in pole_vec])
        ω_p = abs(pole_vec[maxidx])
    end

    m_vec_init = [0, θ_p, pi]

    (lb, idx) = findmax([opnorm(evalfr(sys, exp(im*m))) for m in m_vec_init])
    θ_peak = m_vec_init[idx]

    # Iterations
    for iter=1:maxIters
        gamma = (1+2*T(tol))*lb
        R = gamma^2*I - sys.D'*sys.D
        RinvDt = R\sys.D'
        L = [ sys.A+sys.B*RinvDt*sys.C  sys.B*(R\sys.B');
              zeros(T, sys.nx,sys.nx)      I]
        M = [ I                                 zeros(T, sys.nx,sys.nx);
              sys.C'*(I+sys.D*RinvDt)*sys.C     L[1:sys.nx,1:sys.nx]']

        Λ = complex(eigvals(L,M)) # complex is to ensure type stability

        if numeric_type(sys) <: Real
            # Only need to consider one eigenvalue in each complex-conjugate pairs
            filter!(z -> imag(z) >= 0, Λ)
        end

        # Find eigenvalues on the unit circle
        Λ_on_unit_cirlce = filter(on_unit_circle, Λ)

        θ_vec = angle.(Λ_on_unit_cirlce)

        sort!(θ_vec)

        if isempty(θ_vec)
            return (1+T(tol))*lb, θ_peak/T(sys.Ts)
        end

        # Improve the lower bound
        # if not empty, θ_vec contains at least two values
        for k=1:length(θ_vec)-1
            mk = (θ_vec[k] + θ_vec[k+1])/2
            sigmamax_mk = opnorm(evalfr(sys,exp(mk*1im)))
            if sigmamax_mk > lb
                lb = sigmamax_mk
                θ_peak = mk
            end
        end
    end
    error("In _infnorm_two_steps_dt: The computation of the L∞ norm did not converge in $maxIters iterations")
end

"""`S, P, B = balance(A[, perm=true])`

Compute a similarity transform `T` resulting in `B = T\\A*T` such that the row
and column norms of `B` are approximately equivalent. If `perm=false`, the
transformation will only scale `A` using diagonal `S`, and not permute `A` (i.e., set `P=I`)."""
function balance(A, perm::Bool=true)
    n = LinearAlgebra.checksquare(A)
    B = copy(A)
    job = perm ? 'B' : 'S'
    ilo, ihi, scaling = LAPACK.gebal!(job, B)

    S = diagm(0 => scaling)
    for j = 1:(ilo-1)   S[j,j] = 1 end
    for j = (ihi+1):n   S[j,j] = 1 end

    P = Matrix{Int}(I,n,n)
    if perm
        if ilo > 1
            for j = (ilo-1):-1:1 cswap!(j, round(Int, scaling[j]), P) end
        end
        if ihi < n
            for j = (ihi+1):n    cswap!(j, round(Int, scaling[j]), P) end
        end
    end
    return S, P, B
end

function cswap!(i::Integer, j::Integer, X::StridedMatrix)
    for k = 1:size(X,1)
        X[i, k], X[j, k] = X[j, k], X[i, k]
    end
end

"""
`sysr, G = balreal(sys::StateSpace)`

Calculates a balanced realization of the system sys, such that the observability and reachability gramians of the balanced system are equal and diagonal `G`

See also `gram`, `baltrunc`

Glad, Ljung, Reglerteori: Flervariabla och Olinjära metoder
"""
function balreal(sys::ST) where ST <: AbstractStateSpace
    P = gram(sys, :c)
    Q = gram(sys, :o)

    Q1 = try
        cholesky(Hermitian(Q)).U
    catch
        throw(ArgumentError("Balanced realization failed: Observability grammian not positive definite, system needs to be observable"))
    end
    U,Σ,V = svd(Q1*P*Q1')
    Σ .= sqrt.(Σ)
    Σ1 = diagm(0 => sqrt.(Σ))
    T = Σ1\(U'Q1)

    Pz = T*P*T'
    Qz = inv(T')*Q*inv(T)
    if norm(Pz-Qz) > sqrt(eps())
        @warn("balreal: Result may be inaccurate")
        println("Controllability gramian before transform")
        display(P)
        println("Controllability gramian after transform")
        display(Pz)
        println("Observability gramian before transform")
        display(Q)
        println("Observability gramian after transform")
        display(Qz)
        println("Singular values of PQ")
        display(Σ)
    end

    sysr = ST(T*sys.A/T, T*sys.B, sys.C/T, sys.D, sys.Ts), diagm(0 => Σ)
end

"""
`sysr, G = baltrunc(sys::StateSpace, atol = √ϵ, rtol=1e-3, unitgain=true)`

Reduces the state dimension by calculating a balanced realization of the system sys, such that the observability and reachability gramians of the balanced system are equal and diagonal `G`, and truncating it such that all states corresponding to singular values less than `atol` and less that `rtol σmax` are removed. If `unitgain=true`, the matrix `D` is chosen such that unit static gain is achieved.

See also `gram`, `balreal`

Glad, Ljung, Reglerteori: Flervariabla och Olinjära metoder
"""
function baltrunc(sys::ST; atol = sqrt(eps()), rtol = 1e-3, unitgain = true) where ST <: AbstractStateSpace
    sysbal, S = balreal(sys)
    S = diag(S)
    S = S[S .>= atol]
    S = S[S .>= S[1]*rtol]
    n = length(S)
    A = sysbal.A[1:n,1:n]
    B = sysbal.B[1:n,:]
    C = sysbal.C[:,1:n]
    D = sysbal.D
    if unitgain
        D = D/(C*inv(-A)*B)
    end

    return ST(A,B,C,D,sys.Ts), diagm(0 => S)
end

"""
    syst = similarity_transform(sys, T)
Perform a similarity transform `T : Tx̃ = x` on `sys` such that
```
Ã = T⁻¹AT
B̃ = T⁻¹ B
C̃ = CT
D̃ = D
```
"""
function similarity_transform(sys::ST, T) where ST <: AbstractStateSpace
    Tf = factorize(T)
    A = Tf\sys.A*T
    B = Tf\sys.B
    C = sys.C*T
    D = sys.D
    ST(A,B,C,D,sys.Ts)
end

"""
sysi = innovation_form(sys, R1, R2)
sysi = innovation_form(sys; sysw=I, syse=I, R1=I, R2=I)

Takes a system
```
x' = Ax + Bu + w ~ R1
y  = Cx + e ~ R2
```
and returns the system
```
x' = Ax + Kv
y  = Cx + v
```
where `v` is the innovation sequence.

If `sysw` (`syse`) is given, the covariance resulting in filtering noise with `R1` (`R2`) through `sysw` (`syse`) is used as covariance.

See Stochastic Control, Chapter 4, Åström
"""
function innovation_form(sys::ST, R1, R2) where ST <: AbstractStateSpace
    K = kalman(sys, R1, R2)
    ST(sys.A, K, sys.C, Matrix{eltype(sys.A)}(I, sys.ny, sys.ny), sys.Ts)
end
# Set D = I to get transfer function H = I + C(sI-A)\ K
function innovation_form(sys::ST; sysw=I, syse=I, R1=I, R2=I) where ST <: AbstractStateSpace
	K = kalman(sys, covar(sysw,R1), covar(syse, R2))
	ST(sys.A, K, sys.C, Matrix{eltype(sys.A)}(I, sys.ny, sys.ny), sys.Ts)
end


"""
    horner(c,x)

Evaluate a polynomial whose coefficients are given in descending order in `c`,
at the point `x`, using Horner's rule.
"""
function horner(c,x)

    n = length(c)
    y = c[1]
    for k = 2:n
        y = x*y + c[k]
    end

    return y

end


# Bi-objective unidimensionnal 01 knapsack problem (biukp)
#
# Exercise 10.2 page 290 of
# Multicriteria Optimization (2nd edt), M. Ehrgott, Springer 2005.

# ---- Packages to use
using vOptGeneric, JuMP, GLPK

# ---- Values of the instance to solve
p1 = [10, 3,  6, 8, 2]  # coefficients's vector of the objective 1
p2 = [12, 9, 11, 5, 6]  # coefficients's vector of the objective 2
w  = [ 4, 5,  2, 5, 6]  # coefficients's vector of weights
c  = 17                 # nominal capacity
n  = length( p1 )       # number of items

# ---- setting the model
biukp = vModel( GLPK.Optimizer ) #; JuMP.set_silent( biukp )
@variable( biukp, x[1:n], Bin )
@addobjective( biukp, Max, sum( p1[j]*x[j] for j=1:n ) )
@addobjective( biukp, Max, sum( p2[j]*x[j] for j=1:n ) )
@constraint( biukp, sum( w[j]*x[j] for j=1:n ) <= c )

# ---- Invoking the solver (dichotomic method)
vSolve( biukp, method=:dichotomy, verbose=false )

# ---- Querying the results
Y_N = getY_N( biukp )

# ---- Displaying the results (X_{SE} and Y_{SN})
for i = 1:length(Y_N)
    X = value.(x, i)
    print("X = ", findall(elt -> elt ≈ 1, X))
    println(" | Z = ",Y_N[i])
end


## MIT license (C) 2019 by Andrew Lyasoff

# Julia 1.1.0 code to illustrate the construction of univariate histograms
# and the Monte Carlo simulation technique for multivariate Gaussian samples with a given covariance matrix.
#See (3.67), p.91, in "Stochastic Methods in Asset Pricing."

begin
    using LinearAlgebra
    using SpecialFunctions
    using StatsBase
    using Random
    using Plots
    pyplot()
end

#We write our own histogram function. It takes as an input a single 1-dimensional array of data.
#The number of bins in the histogram is determined automatically by using the Diaconis-Friedman rule.
#The function returns two arrays: the mid-points of the bins and the (unnormalized) heights of the bars.

function hstgram(data_sample::Array{Float64,1})
    data_sorted=sort(data_sample)
    first=data_sorted[1]
    last=data_sorted[end]
    nmb=length(data_sorted)
    IQR=percentile(data_sorted,75)-percentile(data_sorted,25)
    bin_size_loc = 2*IQR*(nmb^(-1.0/3))
    num_bins=Int(floor((last-first)/bin_size_loc))
    bin_size=(last-first)/(num_bins)
    bin_end_points=[first+(i-1)*bin_size for i=1:(num_bins+1)]
    ahist_val=[length(data_sorted[data_sorted .< u]) for u in bin_end_points]
    hist_val=[ahist_val[i+1]-ahist_val[i] for i=1:num_bins]
    mid_bins=[first-bin_size/2+i*bin_size for i=1:num_bins]
    return mid_bins, hist_val
end

## First, create data sampled from the standard univariate normal density.
val=(x->((2*π)^(-1/2)*exp(-x^2/2))).(-3.3:0.05:3.3);

Random.seed!(0xabcdef12); # if needed to generate the same samples

#simulate 10,000 standard normals and generate the histogram
begin
    nval=randn!(zeros(10000));
    U,V=hstgram(nval);
    VV=V/(sum(V)*(U[2]-U[1]));
end

#check the length of the bin
length(VV)

begin
    plot(U.+(U[2]-U[1])/2,VV,line=(:steppre,1),linewidth=0.05,label="histogram")
    xlabel!("samples")
    ylabel!("frequency")
    plot!(-3.3:0.05:3.3,val,label="normal density")
end

#generate another sample
begin
    nval=randn(10000);
    U,V=hstgram(nval);
    VV=V/(sum(V)*(U[2]-U[1]));
    plot(U.+(U[2]-U[1])/2,VV,line=(:steppre,1),linewidth=0.05,label="histogram")
    xlabel!("samples")
    ylabel!("frequency")
    plot!(-3.3:0.05:3.3,val,label="normal density")
end

# standard normals can be generated by sampling from the uniform distribution
begin
    uval=rand(10000);
    nval=(x->sqrt(2)*erfinv(2*x-1)).(uval);
    U,V=hstgram(nval);
    VV=V/(sum(V)*(U[2]-U[1]));
    plot(U.+(U[2]-U[1])/2,VV,line=(:steppre,1),linewidth=0.05,label="histogram")
    xlabel!("samples")
    ylabel!("frequency")
    plot!(-3.3:0.05:3.3,val,label="normal density")
end

#simulate 10,000 standard bi-variate normals
begin
    nval=randn!(zeros(10000));
    nnval=randn!(zeros(10000));
end

scatter(nval,nnval,ratio=1,markersize=1,label="")

# to simulate bi-variate normals with a given covariance matrix
#   first create a fictitios covariance matrix

begin
    A=rand(2,2);
    #Cov=A'A; # another alternative that always yields a positive definite matrix
    Cov=Symmetric(A) # may not produce a positive definite matrix
end

#check if Cov is positive definite; if not, repeat the last step
eigvals(Cov)

begin
    eigdCov=Diagonal(eigvals(Cov).^0.5)
    eigCov=eigvecs(Cov); # matrix of eigen vectors
    chlCov=cholesky(Cov); # Cholesky "square root" of Cov
    MM=eigdCov*eigCov; # Spectral "square root" of Cov
end

# check that the factorizations give what is expected
Cov-MM'*MM
Cov-chlCov.L*chlCov.U
Cov-(chlCov.U)'*chlCov.U
Cov-(chlCov.L)*(chlCov.L)'

#this should be an orthogonal matrix
eigCov'*eigCov

#Transform the randomly generated standard bi-variate sample through the "square root" of the covariance matrix.
#method 1
begin
    NN=hcat(nval,nnval)';
    data_2_dim=MM'NN;
    scatter(data_2_dim[1,:],data_2_dim[2,:],ratio=1,markersize=1,label="")
end

#method 2
begin
    data_2_dim=(chlCov.L)*NN;
    scatter(data_2_dim[1,:],data_2_dim[2,:],ratio=1,markersize=1,label="")
end


## Environment and packages
cd(@__DIR__)
using Pkg; Pkg.activate("."); Pkg.instantiate()

using OrdinaryDiffEq
using ModelingToolkit
using DataDrivenDiffEq
using LinearAlgebra, Optim
using DiffEqFlux, Flux
using Plots
gr()
using JLD2, FileIO
using Statistics
using DelimitedFiles
# Set a random seed for reproduceable behaviour
using Random
Random.seed!(5443)

#### NOTE
# Since the recent release of DataDrivenDiffEq v0.6.0 where a complete overhaul of the optimizers took
# place, SR3 has been used. Right now, STLSQ performs better and has been changed.
# Additionally, the behaviour of the optimization has changed slightly. This has been adjusted
# by decreasing the tolerance of the gradient.

svname = "HudsonBay"
## Data Preprocessing
# The data has been taken from https://jmahaffy.sdsu.edu/courses/f00/math122/labs/labj/q3v1.htm
# Originally published in E. P. Odum (1953), Fundamentals of Ecology, Philadelphia, W. B. Saunders
hudson_bay_data = readdlm("hudson_bay_data.dat", '\t', Float32, '\n')
# Measurements of prey and predator
Xₙ = Matrix(transpose(hudson_bay_data[:, 2:3]))
t = hudson_bay_data[:, 1] .- hudson_bay_data[1, 1]
# Normalize the data; since the data domain is strictly positive
# we just need to divide by the maximum
xscale = maximum(Xₙ, dims =2)
Xₙ .= 1f0 ./ xscale .* Xₙ
# Time from 0 -> n
tspan = (t[1], t[end])

# Plot the data
scatter(t, transpose(Xₙ), xlabel = "t", ylabel = "x(t), y(t)")
plot!(t, transpose(Xₙ), xlabel = "t", ylabel = "x(t), y(t)")

## Direct Identification via SINDy + Collocation

# Create the problem using a gaussian kernel for collocation
full_problem = ContinuousDataDrivenProblem(Xₙ, t, DataDrivenDiffEq.GaussianKernel())
# Look at the collocation
plot(full_problem.t, full_problem.X')
plot(full_problem.t, full_problem.DX')

# Create a Basis
@variables u[1:2]

# Generate the basis functions, multivariate polynomials up to deg 5
# and sine
b = [polynomial_basis(u, 5); sin.(u)]
basis = Basis(b, u)

# Create the thresholds which should be used in the search process
λ = Float32.(exp10.(-7:0.1:5))
# Create an optimizer for the SINDy problem
opt = STLSQ(λ)

# Best result so far
full_res = solve(full_problem, basis, opt, maxiter = 10000, progress = true, denoise = true, normalize = true)

println(full_res)
println(result(full_res))

## Define the network
# Gaussian RBF as activation
rbf(x) = exp.(-(x.^2))

# Define the network 2->5->5->5->2
U = FastChain(
    FastDense(2,5,rbf), FastDense(5,5, rbf), FastDense(5,5, tanh), FastDense(5,2)
)

# Get the initial parameters, first two is linear birth / decay of prey and predator
p = [rand(Float32,2); initial_params(U)]

# Define the hybrid model
function ude_dynamics!(du,u, p, t)
    û = U(u, p[3:end]) # Network prediction
    # We assume a linear birth rate for the prey
    du[1] = p[1]*u[1] + û[1]
    # We assume a linear decay rate for the predator
    du[2] = -p[2]*u[2] + û[2]
end

# Define the problem
prob_nn = ODEProblem(ude_dynamics!,Xₙ[:, 1], tspan, p)

## Function to train the network
# Define a predictor
function predict(θ, X = Xₙ[:,1], T = t)
    Array(solve(prob_nn, Vern7(), u0 = X, p=θ,
                tspan = (T[1], T[end]), saveat = T,
                abstol=1e-6, reltol=1e-6,
                sensealg = ForwardDiffSensitivity()
                ))
end

# Define parameters for Multiple Shooting
group_size = 5
continuity_term = 200.0f0

function loss(data, pred)
	return sum(abs2, data - pred)
end

function shooting_loss(p)
    return multiple_shoot(p, Xₙ, t, prob_nn, loss, Vern7(),
                          group_size; continuity_term)
end

function loss(θ)
    X̂ = predict(θ)
    sum(abs2, Xₙ - X̂) / size(Xₙ, 2) + convert(eltype(θ), 1e-3)*sum(abs2, θ[3:end]) ./ length(θ[3:end])
end

# Container to track the losses
losses = Float32[]

# Callback to show the loss during training
callback(θ,args...) = begin
	l = loss(θ) # Equivalent L2 loss
    push!(losses, l)
    if length(losses)%5==0
        println("Current loss after $(length(losses)) iterations: $(losses[end])")
    end
    false
end

## Training -> First shooting / batching to get a rough estimate

# First train with ADAM for better convergence -> move the parameters into a
# favourable starting positing for BFGS
res1 = DiffEqFlux.sciml_train(shooting_loss, p, ADAM(0.1f0), cb=callback, maxiters = 100)
println("Training loss after $(length(losses)) iterations: $(losses[end])")
# Train with BFGS to achieve partial fit of the data
res2 = DiffEqFlux.sciml_train(shooting_loss, res1.minimizer, BFGS(initial_stepnorm=0.01f0), cb=callback, maxiters = 500)
println("Training loss after $(length(losses)) iterations: $(losses[end])")
# Full L2-Loss for full prediction
res3 = DiffEqFlux.sciml_train(loss, res2.minimizer, BFGS(initial_stepnorm=0.01f0), cb=callback, maxiters = 10000)
println("Final training loss after $(length(losses)) iterations: $(losses[end])")

pl_losses = plot(1:101, losses[1:101], yaxis = :log10, xaxis = :log10, xlabel = "Iterations", ylabel = "Loss", label = "ADAM (Shooting)", color = :blue)
plot!(102:302, losses[102:302], yaxis = :log10, xaxis = :log10, xlabel = "Iterations", ylabel = "Loss", label = "BFGS (Shooting)", color = :red)
plot!(302:length(losses), losses[302:end], color = :black, label = "BFGS (L2)")
savefig(pl_losses, joinpath(pwd(), "plots", "$(svname)_losses.pdf"))

# Rename the best candidate
p_trained = res3.minimizer

## Analysis of the trained network
# Interpolate the solution
tsample = t[1]:0.5:t[end]
X̂ = predict(p_trained, Xₙ[:,1], tsample)
# Trained on noisy data vs real solution
pl_trajectory = scatter(t, transpose(Xₙ), color = :black, label = ["Measurements" nothing], xlabel = "t", ylabel = "x(t), y(t)")
plot!(tsample, transpose(X̂), color = :red, label = ["UDE Approximation" nothing])
savefig(pl_trajectory, joinpath(pwd(), "plots", "$(svname)_trajectory_reconstruction.pdf"))

# Neural network guess
Ŷ = U(X̂,p_trained[3:end])

pl_reconstruction = scatter(tsample, transpose(Ŷ), xlabel = "t", ylabel ="U(x,y)", color = :red, label = ["UDE Approximation" nothing])
plot!(tsample, transpose(Ŷ), color = :red, lw = 2, style = :dash, label = [nothing nothing])
savefig(pl_reconstruction, joinpath(pwd(), "plots", "$(svname)_missingterm_reconstruction.pdf"))
pl_missing = plot(pl_trajectory, pl_reconstruction, layout = (2,1))
savefig(pl_missing, joinpath(pwd(), "plots", "$(svname)_reconstruction.pdf"))
## Symbolic regression via sparse regression (SINDy based)
# We reuse the basis and optimizer defined at the beginning

nn_problem = ContinuousDataDrivenProblem(X̂, tsample, DX = Ŷ)
nn_res = solve(nn_problem, basis, opt, maxiter = 10000, progress = true, normalize = false, denoise = true)
println(nn_res)
println(result(nn_res))

# Define the recovered, hyrid model with the rescaled dynamics
function recovered_dynamics!(du,u, p, t)
    û = nn_res(u, p[3:end]) # Network prediction
    du[1] = p[1]*u[1] + û[1]
    du[2] = -p[2]*u[2] + û[2]
end

p_model = [p_trained[1:2];parameters(nn_res)]

estimation_prob = ODEProblem(recovered_dynamics!, Xₙ[:, 1], tspan, p_model)
# Convert for reuse
sys = modelingtoolkitize(estimation_prob);
dudt = ODEFunction(sys);
estimation_prob = ODEProblem(dudt,Xₙ[:, 1], tspan, p_model)
estimate = solve(estimation_prob, Tsit5(), saveat = t)

##  Fit the found model
function loss_fit(θ)
    X̂ = Array(solve(estimation_prob, Tsit5(), p = θ, saveat = t))
    sum(abs2, X̂ .- Xₙ)
end

# Post-fit the model
res_fit = DiffEqFlux.sciml_train(loss_fit, p_model, BFGS(initial_stepnorm = 0.1f0), maxiters = 1000)
p_fitted = res_fit.minimizer

# Estimate
estimate_rough = solve(estimation_prob, Tsit5(), saveat = 0.1*mean(diff(t)), p = p_model)
estimate = solve(estimation_prob, Tsit5(), saveat = 0.1*mean(diff(t)), p = p_fitted)

# Plot
pl_fitted = plot(t, transpose(Xₙ), style = :dash, lw = 2,color = :black, label = ["Measurements" nothing], xlabel = "t", ylabel = "x(t), y(t)")
plot!(estimate_rough, color = :red, label = ["Recovered" nothing])
plot!(estimate, color = :blue, label = ["Recovered + Fitted" nothing])
savefig(pl_fitted,joinpath(pwd(),"plots","$(svname)recovery_fitting.pdf"))

## Simulation

# Look at long term prediction
t_long = (0.0f0, 50.0f0)
estimate_long = solve(estimation_prob, Tsit5(), saveat = 0.25f0, tspan = t_long,p = p_fitted)
plot(estimate_long.t, transpose(xscale .* estimate_long[:,:]), xlabel = "t", ylabel = "x(t),y(t)")

## Save the results
save(joinpath(pwd(),"results","Hudson_Bay_recovery.jld2"),
    "X", Xₙ, "t" , t, "neural_network" , U, "initial_parameters", p, "trained_parameters" , p_trained, # Training
    "losses", losses, "result", nn_res, "recovered_parameters", parameters(nn_res), # Recovery
    "model", recovered_dynamics!, "model_parameter", p_model, "fitted_parameter", p_fitted,
    "long_estimate", estimate_long) # Estimation

## Post Processing and Plots

c1 = 3 # RGBA(174/255,192/255,201/255,1) # Maroon
c2 = :orange # RGBA(132/255,159/255,173/255,1) # Red
c3 = :blue # RGBA(255/255,90/255,0,1) # Orange
c4 = :purple # RGBA(153/255,50/255,204/255,1) # Purple

p3 = scatter(t, transpose(Xₙ), color = [c1 c2], label = ["x data" "y data"],
             title = "Recovered Model from Hudson Bay Data",
             titlefont = "Helvetica", legendfont = "Helvetica",
             markersize = 5)

plot!(p3,estimate_long, color = [c3 c4], lw=1, label = ["Estimated x(t)" "Estimated y(t)"])
plot!(p3,[19.99,20.01],[0.0,maximum(Xₙ)*1.25],lw=1,color=:black, label = nothing)
annotate!([(10.0,maximum(Xₙ)*1.25,text("Training \nData",12 , :center, :top, :black, "Helvetica"))])
savefig(p3,joinpath(pwd(),"plots","$(svname)full_plot.pdf"))


# ------------------------------------------------------------------------------------------
# # Julia for Data Science - Plotting
#
# ### Data visualization: generating nice looking plots in Julia is straight forward
# In what's next, we will see some of the tools that Julia plotting provides to produce high
# quality figures for your data. In particular we'll look at
#
# 1. Plotting mathematical functions
# 1. Visualizing statistics
# 1. Subplotting
#
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## Part 1: plot math functions (specifically latex equations) in our plots
# ------------------------------------------------------------------------------------------

using LaTeXStrings
using Plots
pyplot(leg=false)
x = 1:0.2:4

# ------------------------------------------------------------------------------------------
# Create three functions and plot them all!
# ------------------------------------------------------------------------------------------

y1 = sqrt.(x)
y2 = log.(x)
y3 = x.^2

f1 = plot(x,y1)
plot!(f1,x,y2) # "plot!" means "plot on the same canvas we just plot on"
plot!(f1,x,y3)

# ------------------------------------------------------------------------------------------
# Now we can annotate each of these plots! using either text, or latex strings
# ------------------------------------------------------------------------------------------

annotate!(f1,[(x[6],y1[6],text(L"\sqrt{x}",16,:center)),
          (x[11],y2[11],text(L"log(x)",:right,16)),
          (x[6],y3[6],text(L"x^2",16))])

# ------------------------------------------------------------------------------------------
# ## Part 2: Visualizing statistics
#
# **2D histograms** are really easy!
# ------------------------------------------------------------------------------------------

n = 1000
set1 = randn(n)
set2 = randn(n)
histogram2d(set1,set2,nbins=20,colorbar=true)

# ------------------------------------------------------------------------------------------
# **Let's go back to our houses dataset and learn even more things about it!**
# ------------------------------------------------------------------------------------------

using DataFrames
houses = readtable("houses.csv")
filter_houses = houses[houses[!, :sq__ft].>0,:]
x = filter_houses[!, :sq__ft]
y = filter_houses[!, :price]

gh = histogram2d(x,y,nbins=20,colorbar=true)
xaxis!(gh,"square feet")
yaxis!(gh,"price")

# ------------------------------------------------------------------------------------------
# Interesting!
#
# Most houses sold are in the range 1000-1500 and they cost approximately 150,000 dollars
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# *Let's see more stats plots.*
#
# We can convince ourselves that random distrubutions are indeed very similar.
#
# Let's do that through **box plots** and **violin plots**.
# ------------------------------------------------------------------------------------------

using StatsPlots
y = rand(10000,6) # generate 6 random samples of size 1000 each
f2 = violin(["Series 1" "Series 2" "Series 3" "Series 4" "Series 5"],y,leg=false,color=:red)

boxplot!(["Series 1" "Series 2" "Series 3" "Series 4" "Series 5"],y,leg=false,color=:green)

# ------------------------------------------------------------------------------------------
# These plots look almost identical, so we do have the same distribution indeed.
#
# Let's study the price distributions in different cities in the houses dataset.
# ------------------------------------------------------------------------------------------

some_cities = ["SACRAMENTO","RANCHO CORDOVA","RIO LINDA","CITRUS HEIGHTS","NORTH HIGHLANDS","ANTELOPE","ELK GROVE","ELVERTA" ] # try picking pther cities!

fh = plot(xrotation=90)
for ucity in some_cities
    subs = filter_houses[filter_houses[!, :city].==ucity,:]
    city_prices = subs[!, :price]
    violin!(fh,[ucity],city_prices,leg=false)
end
display(fh)

# ------------------------------------------------------------------------------------------
# ## Part 3: Subplots are very easy!
#
# To create a plot with subplots, all we need to do is throw the variables bound to
# individual plots inside another call to `plot`!
# ------------------------------------------------------------------------------------------

x = -10:.1:10
p1 = plot(x, x)
p2 = plot(x, x.^2)
p3 = plot(x, x.^3)
p4 = plot(x, x.^4)
plot(p1,p2,p3,p4,layout=(2,2),legend=false)

# ------------------------------------------------------------------------------------------
# You can create your own layout as follows.
# ------------------------------------------------------------------------------------------

mylayout = @layout([a{0.5h};[b{0.7w} c]])
plot(fh,f2,gh,layout=mylayout,legend=false)

# this layout:
# 1
# 2 3

# ------------------------------------------------------------------------------------------
# ### Please let us know how we're doing!
#
# https://tinyurl.com/JuliaDataScience
# ------------------------------------------------------------------------------------------