#=
By starting at the top of the triangle below and moving to adjacent numbers on the row below, the maximum total from top to bottom is 23.

3
7 4
2 4 6
8 5 9 3

That is, 3 + 7 + 4 + 9 = 23.

Find the maximum total from top to bottom in triangle.txt (right click and 'Save Link/Target As...'), a 15K text file containing a triangle with one-hundred rows.

NOTE: This is a much more difficult version of Problem 18. It is not possible to try every route to solve this problem, as there are 2^99 altogether! If you could check one trillion (1012) routes every second it would take over twenty billion years to check them all. There is an efficient algorithm to solve it. ;o)
=#
#a =

function calc()
  lines = readlines(open(dirname(@__FILE__()) * "/../Resources/p67.txt"))
  arr = map(x -> [parse(Int64, s) for s = split(strip(x), ',')], lines)
  result = Array[]
  push!(result, arr[length(arr)])
  for row in reverse(1:length(arr)-1)
    nextrow = Int64[]
    for column in 1:length(arr[row])
      lowleft = last(result)[column]
      lowright = last(result)[column+1]
      center = arr[row][column]
      push!(nextrow, max(center+lowleft, center+lowright))
    end
    push!(result, nextrow)
  end
  last(result)[1]
end
@time println(calc())


include("derivative.jl")

function quadratic(f)
    # Compute the first derivative of f
	f1 = derivative(f)

    # Get the y intercept (explictly passing a floating point number)
	c = f(0.0)

	# Get the y intercept of the first derivative
	b = f1(0.0)

	a = f(1.0) - b - c

    # Our old friend the quadratic formula
    # Notice how Julia lets you return multiple values
    # Multiple values should be separated by a comma and
    # are returned as a tuple
    # return (-b + sqrt(b^2 -4a*c)) / 2a, (-b - sqrt(b^2 -4a*c)) / 2a
	return (-b + sqrt(b^2 -4a*c + 0im)) / 2a, (-b - sqrt(b^2 -4a*c + 0im)) / 2a
end


module hw1_4

using LinearAlgebra, PlotlyJS, ORCA

export total_manuever, plot_maneuver

function oe_to_rθh(oe::Vector,μ::Real) :: Vector

  a,e,i,Ω,ω,ν = oe
  return [a*(1-e^2)/(1+e*cos(ν)),
  0,
  0,
  (μ/sqrt(μ*a*(1-e^2)))*e*sin(ν),
  (μ/sqrt(μ*a*(1-e^2)))*(1+e*cos(ν)),
  0]

end

function rθh_to_xyz(rθh_vec::Vector,oe::Vector)

  a,e,i,Ω,ω,ν = oe
  θ = ω+ν
  cΩ,sΩ,ci,si,cθ,sθ = cos(Ω),sin(Ω),cos(i),sin(i),cos(θ),sin(θ)
  DCM = [cΩ*cθ-sΩ*ci*sθ -cΩ*sθ-sΩ*ci*cθ sΩ*si;
  sΩ*cθ+cΩ*ci*sθ -sΩ*sθ+cΩ*ci*cθ -cΩ*si;
  si*sθ si*cθ ci]
  DCM = kron(Matrix(I,2,2),DCM)
  return DCM*rθh_vec

end

function oe_to_xyz(oe::Vector,μ::Real)

  return rθh_to_xyz(oe_to_rθh(oe,μ),oe)

end

function xyz_to_oe(cart_vec::Vector,μ::Real)

  r_xyz, v_xyz = cart_vec[1:3],cart_vec[4:6]
  r = norm(r_xyz)
  h_xyz = cross(r_xyz,v_xyz) #km^2/s
  h = norm(h_xyz) #km^2/s
  ξ = dot(v_xyz,v_xyz)/2 - μ/r #km^2 s^-2
  a = -μ/(2ξ) #km
  e = sqrt(1 + (2h^2*ξ)/μ^2)
  e_xyz = cross(v_xyz,h_xyz)/μ - r_xyz/r
  i = acos(h_xyz[3]/h) #rad
  n_xyz = cross([0,0,1],h_xyz)
  Ω = acos(dot(n_xyz,[1,0,0])/norm(n_xyz))
  ω = acos((dot(n_xyz,e_xyz)/(norm(n_xyz)*e)))
  ν = acos((dot(r_xyz,e_xyz))/(r*norm(e_xyz)))
  Ω = dot(n_xyz,[0,1,0]) > 0. ? Ω : -Ω
  ω = dot(e_xyz,[0,0,1]) > 0. ? ω : -ω
  ν = dot(r_xyz,v_xyz) > 0. ? ν : -ν
  return [a,e,i,Ω,ω,ν]

end

function total_manuever(r1::Float64,
                        r2::Float64,
                        i1::Float64,
                        i2::Float64,
                        i3::Float64,
                        i4::Float64,
                        μ::Float64)
  if i2 > i3
    return Inf
  else
    a1 = r1
    a2 = (r1+r2)/2
    e1 = 0
    e2 = (r2-r1)/(r2+r1)
    v1_pre = oe_to_xyz([a1, 0., i1, 0., 0., 0.], μ)[4:6]
    v1_post = oe_to_xyz([a2, e2, i2, 0., 0., 0.], μ)[4:6]
    v2_pre = oe_to_xyz([a2, e2, i2, 0., 0., π], μ)[4:6]
    v2_post = oe_to_xyz([a2, e2, i3, 0., 0., π], μ)[4:6]
    v3_pre = oe_to_xyz([a2, e2, i3, 0., 0., 0.], μ)[4:6]
    v3_post = oe_to_xyz([a1, e1, i4, 0., 0., 0.], μ)[4:6]
    Δv1 = norm(v1_post-v1_pre)
    Δv2 = norm(v2_post-v2_pre)
    Δv3 = norm(v3_post-v3_pre)
    return Δv1 + Δv2 + Δv3
  end
end

function plot_maneuver(r1::Float64, r2::Float64, Δ_i::Float64,  μ::Float64, lim::Float64)
  N = 200
  node1_inc = node3_inc = collect(range(0, length=N, stop=lim))
  z = zeros(Float64, N, N)
  base = zeros(Float64, N, N, 2)
  for i in 1:N
    for j in 1:N
      z[i,j] = total_manuever(r1, r2, 0., node1_inc[i], Δ_i-node3_inc[j], Δ_i, μ)
      base[i,j,:] = [node1_inc[i], node3_inc[j]]
    end
  end

  data = contour(;z=z, x=node1_inc, y=node3_inc,
                 name="ΔV", contours=attr(size=0.2),
                 colorbar=attr(title=attr(text="ΔV (km/s)")))
  if lim < Δ_i
    layout = Layout(
      ;title="Total ΔV vs. Δi₁ and Δi₃ for a 90° Bi-Elliptic Transfer (zoomed)",
      width=920,
      height=920,
      xaxis=attr(title=attr(text="Δi₁ (radians)")),
      yaxis=attr(title=attr(text="Δi₃ (radians)")))
    p = Plot(data, layout)
    savefig(p::Union{Plot,PlotlyJS.SyncPlot}, "hw1_4_zoomed.png")
  else
    layout = Layout(
      ;title="Total ΔV vs. Δi₁ and Δi₃ for a 90° Bi-Elliptic Transfer",
      width=920,
      height=920,
      xaxis=attr(title=attr(text="Δi₁ (radians)")),
      yaxis=attr(title=attr(text="Δi₃ (radians)")))
    p = Plot(data, layout)
    savefig(p::Union{Plot,PlotlyJS.SyncPlot}, "hw1_4.png")
  end
  display(plot(p))

  min, position = findmin(z)
  angles = base[position, :]
  println("Found min: ", min)
  println("At: ", 180/π*angles[1], "°, ", 180/π*angles[2], "°")
  while true
    println("Ctrl+C to close")
    sleep(10)
  end
end

end


"""
    hyperbola2(x, s1, s2, Δy, xoff=0.0, yoff=0.0)

One branch of a hyperbola as a function of x given by

- s1: asymptote slope for x<0
- s2: asymptote slope for x>0
- Δy: absolute value of abcissa at x=0
- optionally shift the origin by xoff, yoff

Ref https://maurow.bitbucket.io/docs/hyperbola.html
"""
function hyperbola(x, s1, s2, Δy, xoff=0.0, yoff=0.0)
    x = x-xoff
    # Deal with degenerate cases and select the right branch of the hyperbola
    if Δy==0 # return two lines
        i1 = x.<0
        i2 = x.>=0
        return vcat(s1*x[i1], s2*x[i2]) + yoff
    elseif s1==s2 # return one line
        branch = 0
    elseif s1>s2
        branch = -1
        s1,s2 = s2,s1
    else
        branch = +1
    end
    B = (s1+s2)/2
    C = -Δy^2
    D = - (s2-s1)^2/4
    return B*x + branch * sqrt(-D*x.^2-C) + yoff
end

"""
    hyperbola01(x, Δy, xoff=0.0, yoff=0.0)

A hyperbola with asymptotes y=0 for x<0 and y=x for x>0, and
abcissa Δy.  Its origin can be moved with xoff and yoff.
"""
function hyperbola01(x, Δy, xoff=0.0, yoff=0.0)
    x = x-xoff
    B = 1/2
    C = -Δy^2
    D = -1/4
    return B*x + sqrt(-D*x^2-C) + yoff
end

"""
    max_smooth(x0, x, delta)

to replace

    max(x0, x)

for a fixed x0 and varying x.  `delta` gives the difference:
delta = max_smooth(x0,x0,delta)-x0,
which is the maximal difference for any x.

Note, all returned values >= x0 (== for delta==0).
"""
max_smooth(x0, x, delta) =
    delta==0 ? max(x0,x) : hyperbola01(x, delta, x0, x0)
"""
    min_smooth(x0, x, delta)

see max_smooth
"""
min_smooth(x0, x, delta) = -max_smooth(-x0, -x, delta)

####
# Sigmoid smoothing

"""
    sigmoid(x, x0, w)

Smooth transition function from 0 to 1.  At x0+w its value is about 0.99,
at x0-w about 0.01.
"""
sigmoid(x, x0, w) = 1/(1+exp(-(x-x0)/w*5))

"""
    fn_cap(x0, width, fn, x, args...)

Instead of

    max(fn(x0,args...), fn(x, args...))

where at x0+/-width the value is to 1% correct.

There is also the not recommended:

    fn_cap_(f0, width, fn, args...)

where at f0+/-width the value is to 1% correct.
"""
function fn_cap(x0, width, fn, x, args...)
    ff = fn(x, args...)
    f0 = fn(x0, args...)
    si = sigmoid(x, x0, width)
    f0*(1-si) + ff*si
end
function fn_cap_(f0, width, fn, args...)
    ff = fn.(args...)
    si = sigmoid.(ff, f0, width)
    f0.*(1-si) + ff.*si
end

"""
    val_cap(x0,x,width)
instead of
    max(x0, x)

with the return value within 1% for x=x0+/-width.
"""
function val_cap(x0,x,width)
    si = sigmoid(x,x0,width)
    x0*(1-si)+x*si
end


cd(@__DIR__) # changes the directory to the current directory, the default I guess is the HOME
using Pkg; Pkg.activate("."); Pkg.instantiate()

using OrdinaryDiffEq, ModelingToolkit, DiffEqOperators, DomainSets

@parameters t, x
@variables u(..)
Dt = Differential(t)
Dx = Differential(x)
Dxx = Differential(x)^2

# PDE
eq  = Dt(u(t,x)) + u(t,x)*Dx(u(t,x)) ~ 0 #- (0.01/pi)*Dxx(u(t,x))

# Initial and boundary conditions
bcs = [u(0,x) ~ -sin(pi*x),
       u(t,-1) ~ 0.,
       u(t,1) ~ 0.,
       u(t,-1) ~ u(t,1)]

# Space and time domains
domains = [t ∈ Interval(0.0,1.0),
           x ∈ Interval(-1.0,1.0)]
# Discretization
dx = 0.1

# PDE system - ModelingToolkit
@named pdesys = PDESystem(eq,bcs,domains,[t,x],[u(t,x)])

order = 2
discretization = MOLFiniteDifference([x=>dx],t;centered_order=order)

# Convert the PDE problem into an ODE problem
prob = discretize(pdesys,discretization)

# Solve ODE problem
using OrdinaryDiffEq
sol = solve(prob,Vern9(),saveat=0.1)

#= # Plot results and compare with exact solution
x = (0:dx:1)[2:end-1]
t = sol.t

using Plots
plt = plot()

for i in 1:length(t)
    plot!(x,sol.u[i],label="Numerical, t=$(t[i])")
end
display(plt)
savefig("plot.png") =#


# indicator of a simplex

export IndSimplex

"""
**Indicator of a simplex**

    IndSimplex(a=1.0)

Returns the indicator of the set
```math
S = \\left\\{ x : x \\geq 0, \\sum_i x_i = a \\right\\}.
```
By default `a=1.0`, therefore ``S`` is the probability simplex.
"""
struct IndSimplex{R <: Real} <: ProximableFunction
    a::R
    function IndSimplex{R}(a::R) where {R <: Real}
        if a <= 0
            error("parameter a must be positive")
        else
            new(a)
        end
    end
end

is_convex(f::IndSimplex) = true
is_set(f::IndSimplex) = true

IndSimplex(a::R=1.0) where {R <: Real} = IndSimplex{R}(a)

function (f::IndSimplex)(x::AbstractArray{R}) where R <: Real
    if all(x .>= 0) && sum(x) ≈ f.a
        return R(0)
    end
    return R(Inf)
end

function simplex_proj_condat!(y::AbstractArray{R}, a, x::AbstractArray{R}) where R
    # Implements algorithm proposed in:
    # Condat, L. "Fast projection onto the simplex and the l1 ball",
    # Mathematical Programming, 158:575–585, 2016.
    v = [x[1]]
    v_tilde = R[]
    rho = x[1] - a
    N = length(x)
    for k in 2:N
        if x[k] > rho
            rho += (x[k] - rho) / (length(v) + 1)
            if rho > x[k] - a
                push!(v, x[k])
            else
                append!(v_tilde, v)
                v = [x[k]]
                rho = x[k] - a
            end
        end
    end
    for z in v_tilde
        if z > rho
            push!(v, z)
            rho += (z - rho) / length(v)
        end
    end
    v_changed = true
    while v_changed == true
        v_changed = false
        k = 1
        while k <= length(v)
            z = v[k]
            if z <= rho
                deleteat!(v, k)
                v_changed = true
                rho += (rho - z) / length(v)
            else
                k = k + 1
            end
        end
    end
    y .= max.(x .- rho, R(0))
end

function prox!(y::AbstractArray{R}, f::IndSimplex, x::AbstractArray{R}, _::R=R(1)) where R <: Real
    simplex_proj_condat!(y, f.a, x)
    return R(0)
end

fun_name(f::IndSimplex) = "indicator of the probability simplex"
fun_dom(f::IndSimplex) = "AbstractArray{Real}"
fun_expr(f::IndSimplex) = "x ↦ 0 if x ⩾ 0 and sum(x) = a, +∞ otherwise"
fun_params(f::IndSimplex) = "a = $(f.a)"

function prox_naive(f::IndSimplex, x::AbstractArray{R}, _::R=R(1)) where R <: Real
    low = minimum(x)
    upp = maximum(x)
    v = x
    s = Inf
    for i = 1:100
        if abs(s)/f.a ≈ 0
            break
        end
        alpha = (low+upp)/2
        v = max.(x .- alpha, R(0))
        s = sum(v) - f.a
        if s <= 0
            upp = alpha
        else
            low = alpha
        end
    end
    return v, R(0)
end


# Examples from Luc Jaulin, Michel Kieffer, Olivier Didrit and Eric Walter - Applied Interval Analysis

using IntervalArithmetic, IntervalRootFinding, StaticArrays

A = [4..5 -1..1 1.5..2.5; -0.5..0.5 -7.. -5 1..2; -1.5.. -0.5 -0.7.. -0.5 2..3]
sA = SMatrix{3}{3}(A)
mA = MMatrix{3}{3}(A)

b = [3..4, 0..2, 3..4]
sb = SVector{3}(b)
mb = MVector{3}(b)

p = fill(-1e16..1e16, 3)

rts = gauss_seidel_interval!(p, A, b, precondition=true) # Gauss-Seidel Method; precondition=true by default
rts = gauss_seidel_interval!(p, sA, sb, precondition=true) # Gauss-Seidel Method; precondition=true by default
rts = gauss_seidel_interval!(p, mA, mb, precondition=true) # Gauss-Seidel Method; precondition=true by default

rts = gauss_seidel_interval(A, b, precondition=true) # Gauss-Seidel Method; precondition=true by default
rts = gauss_seidel_interval(sA, sb, precondition=true) # Gauss-Seidel Method; precondition=true by default
rts = gauss_seidel_interval(mA, mb, precondition=true) # Gauss-Seidel Method; precondition=true by default

rts = gauss_seidel_contractor!(p, A, b, precondition=true) # Gauss-Seidel Method (Vectorized); precondition=true by default
rts = gauss_seidel_contractor!(p, sA, sb, precondition=true) # Gauss-Seidel Method (Vectorized); precondition=true by default
rts = gauss_seidel_contractor!(p, mA, mb, precondition=true) # Gauss-Seidel Method (Vectorized); precondition=true by default

rts = gauss_seidel_contractor(A, b, precondition=true) # Gauss-Seidel Method (Vectorized); precondition=true by default
rts = gauss_seidel_contractor(sA, sb, precondition=true) # Gauss-Seidel Method (Vectorized); precondition=true by default
rts = gauss_seidel_contractor(mA, mb, precondition=true) # Gauss-Seidel Method (Vectorized); precondition=true by default

rts = gauss_elimination_interval!(p, A, b, precondition=true) # Gaussian Elimination; precondition=true by default
rts = gauss_elimination_interval(A, b, precondition=true) # Gaussian Elimination; precondition=true by default


#=
Functions to compute quadratic sums

@author : Spencer Lyon <spencer.lyon@nyu.edu>

@date : 2014-08-19
=#

"""
Computes the expected discounted quadratic sum

    q(x_0) = E sum_{t=0}^{infty} beta^t x_t' H x_t

Here {x_t} is the VAR process x_{t+1} = A x_t + C w_t with {w_t}
standard normal and x_0 the initial condition.

##### Arguments
- `A::Union{Float64, Matrix{Float64}}` The n x n matrix described above (scalar)
if n = 1
- `C::Union{Float64, Matrix{Float64}}` The n x n matrix described above (scalar)
if n = 1
- `H::Union{Float64, Matrix{Float64}}` The n x n matrix described above (scalar)
if n = 1
- `beta::Float64`: Discount factor in (0, 1)
- `x_0::Union{Float64, Vector{Float64}}` The initial condtion. A conformable
array (of length n) or a scalar if n=1

##### Returns

- `q0::Float64` : Represents the value q(x_0)

##### Notes

The formula for computing q(x_0) is q(x_0) = x_0' Q x_0 + v where

- Q is the solution to Q = H + beta A' Q A and
- v = \trace(C' Q C) \beta / (1 - \beta)

"""
function var_quadratic_sum(A::ScalarOrArray, C::ScalarOrArray, H::ScalarOrArray,
                           bet::Real, x0::ScalarOrArray)
    n = size(A, 1)

    # coerce shapes
    A = reshape([A;], n, n)
    C = reshape([C;], n, n)
    H = reshape([H;], n, n)
    x0 = reshape([x0;], n)

    # solve system
    Q = solve_discrete_lyapunov(sqrt(bet) .* A', H)
    cq = C'*Q*C
    v = trace(cq) * bet / (1 - bet)
    q0 = x0'*Q*x0 + v
    return q0[1]
end

"""
Computes the quadratic sum

    V = sum_{j=0}^{infty} A^j B A^{j'}

V is computed by solving the corresponding discrete lyapunov equation using the
doubling algorithm.  See the documentation of `solve_discrete_lyapunov` for
more information.

##### Arguments

- `A::Matrix{Float64}` : An n x n matrix as described above.  We assume in order
for convergence that the eigenvalues of A have moduli bounded by unity
- `B::Matrix{Float64}` : An n x n matrix as described above.  We assume in order
for convergence that the eigenvalues of B have moduli bounded by unity
- `max_it::Int(50)` : Maximum number of iterations

##### Returns

- `gamma1::Matrix{Float64}` : Represents the value V

"""
function m_quadratic_sum(A::Matrix, B::Matrix; max_it=50)
    solve_discrete_lyapunov(A, B, max_it)
end


"""
    prlyap(A, C) -> X

Solve the reverse-time periodic discrete-time Lyapunov equation

    A'σXA + C = X

where `σ` is the forward shift operator `σX(i) = X(i+1)`.

For the periodic matrices `A` and `C` with the same sampling period and commensurate periods,
the periodic matrix `X` is determined, whose period is automatically set to the
the least common period of `A` and `C`.
"""
function prlyap(A::PeriodicArray, C::PeriodicArray)
   A.Ts ≈ C.Ts || error("A and C must have the same sampling time")
   period = promote_period(A, C)
   nta = numerator(rationalize(period/A.period))
   K = nta*A.nperiod*A.dperiod
   X = pslyapd(A.M, C.M; adj = true)
   return PeriodicArray(X, period; nperiod = div(K,size(X,3)))
end
prlyap(A::PeriodicArray, C::AbstractMatrix) = prlyap(A, PeriodicArray(C, A.Ts; nperiod = 1))
"""
    pflyap(A, C) -> X

Solve the forward-time periodic discrete-time Lyapunov equation

    A'XA + C = σX

where `σ` is the forward shift operator `σX(i) = X(i+1)`.

For the periodic matrices `A` and `C` with the same sampling period and commensurate periods,
the periodic matrix `X` is determined, whose period is automatically set to the
the least common period of `A` and `C`.
"""
function pflyap(A::PeriodicArray, C::PeriodicArray)
   A.Ts ≈ C.Ts || error("A and C must have the same sampling time")
   period = promote_period(A, C)
   nta = numerator(rationalize(period/A.period))
   K = nta*A.nperiod*A.dperiod
   X = pslyapd(A.M, C.M; adj = false)
   return PeriodicArray(X, period; nperiod = div(K,size(X,3)))
end
pflyap(A::PeriodicArray, C::AbstractMatrix) = pflyap(A, PeriodicArray(C, A.Ts; nperiod = 1))
"""
    pslyapd(A, C; adj = true) -> X

Solve the periodic discrete-time Lyapunov equation.

For the square `n`-th order periodic matrices `A(i)`, `i = 1, ..., pa` and
`C(i)`, `i = 1, ..., pc`  of periods `pa` and `pc`, respectively,
the periodic solution `X(i)`, `i = 1, ..., p` of period `p = lcm(pa,pc)` of the
periodic Lyapunov equation is solved:  for `adj = true`,

    A(i)'*X(i+1)*A(i) + C(i) = X(i), i = 1, ..., p

and for `adj = false`

    A(i)*X(i)*A(i)' + C(i) = X(i+1), i = 1, ..., p.

The periodic matrices `A` and `C` are stored in the `n×n×pa` and `n×n×pc` 3-dimensional
arrays `A` and `C`, respectively, and `X` results as a `n×n×p` 3-dimensional array.

The periodic discrete analog of the Bartels-Steward method based on the periodic Schur form
of the periodic matrix `A` is employed [1].

_Reference:_

[1] A. Varga. Periodic Lyapunov equations: some applications and new algorithms.
              Int. J. Control, vol, 67, pp, 69-87, 1997.
"""
function pslyapd(A::AbstractArray{T1, 3}, C::AbstractArray{T2, 3}; adj::Bool = true) where {T1, T2}
   n = LinearAlgebra.checksquare(A[:,:,1])
   pa = size(A,3)
   pc = size(C,3)
   (LinearAlgebra.checksquare(C[:,:,1]) == n && all([issymmetric(C[:,:,i]) for i in 1:pc])) ||
      throw(DimensionMismatch("all C[:,:,i] must be $n x $n symmetric matrices"))
   p = lcm(pa,pc)

   T = promote_type(T1, T2)
   T <: BlasFloat  || (T = promote_type(Float64,T))
   A1 = T1 == T ? A : A1 = convert(Array{T,3},A)
   C1 = T2 == T ? C : C1 = convert(Array{T,3},C)

   # Reduce A to Schur form and transform C
   AS, Q, _, KSCHUR = pschur(A1)

   #X = Q'*C*Q
   X = Array{T,3}(undef, n, n, p)

   for i = 1:p
       ia = mod(i-1,pa)+1
       ic = mod(i-1,pc)+1
       ia1 = mod(i,pa)+1

       X[:,:,i] = adj ? utqu(view(C1,:,:,ic),view(Q,:,:,ia)) :
                        utqu(view(C1,:,:,ic),view(Q,:,:,ia1))
   end
   # solve A'σXA - X + C = 0
   pdlyaps!(KSCHUR, AS, X; adj)

   #X <- Q*X*Q'
   for i = 1:p
       ia = mod(i-1,pa)+1
       utqu!(view(X,:,:,i),view(Q,:,:,ia)')
   end
   return X
end
"""
     pslyapdkr(A, C; adj = true) -> X

Solve the periodic discrete-time Lyapunov matrix equation

      A'σXA + C = X, if adj = true,

or

      A*X*A' + C =  σX, if adj = false,

where `σ` is the forward shift operator `σX(i) = X(i+1)`.
The Kronecker product expansion of equations is employed. `A` and `C` are
periodic square matrices, and `A` must not have characteristic multipliers on the unit circle.
This function is not recommended for large order matrices or large periods.
"""
function pslyapdkr(A::AbstractArray{T, 3}, C::AbstractArray{T, 3}; adj = true) where {T}
    m, n, pc = size(C)
    n == LinearAlgebra.checksquare(A[:,:,1])
    m == LinearAlgebra.checksquare(C[:,:,1])
    m == n  || throw(DimensionMismatch("A and C have incompatible dimensions"))
    pa = size(A,3)
    n2 = n*n
    p = lcm(pa,pc)
    N = p*n2
    R = zeros(T, N, N)
    adj ? copyto!(view(R,N-n2+1:N,1:n2),kron(A[:,:,pa]',A[:,:,pa]')) :
          copyto!(view(R,1:n2,N-n2+1:N),kron(A[:,:,pa],A[:,:,pa]))
    (i2, j2) = adj ? (n2, n2+n2) : (n2+n2, n2)
    for i = 1:p-1
        i1 = i2-n2+1
        j1 = j2-n2+1
        ia = mod(i-1,pa)+1
        adj ? copyto!(view(R,i1:i2,j1:j2),kron(A[:,:,ia]',A[:,:,ia]')) :
              copyto!(view(R,i1:i2,j1:j2),kron(A[:,:,ia],A[:,:,ia]))
        i2 += n2
        j2 += n2
    end
    indc = mod.(0:p-1,pc).+1
    return adj ? reshape((I-R) \ (C[:,:,indc][:]), n, n, p) :
    reshape((I-R) \ (C[:,:,circshift(indc,1)][:]), n, n, p)
end
function pdlyaps!(KSCHUR::Int, A::StridedArray{T1,3}, C::StridedArray{T1,3}; adj = true) where {T1<:BlasReal}
   n = LinearAlgebra.checksquare(A[:,:,1])
   pa = size(A,3)
   pc = size(C,3)
   (LinearAlgebra.checksquare(C[:,:,1]) == n && all([issymmetric(C[:,:,i]) for i in 1:pc])) ||
      throw(DimensionMismatch("all C[:,:,i] must be $n x $n symmetric matrices"))
   rem(pc,pa) == 0 || error("the period of C must be an integer multiple of A")
   (KSCHUR <= 0 || KSCHUR > pa ) &&
         error("KSCHUR has a value $KSCHUR, which is inconsistent with A ")

   if pa == 1 && pc == 1
      lyapds!(view(A,:,:,1), view(C,:,:,1); adj)
      return C[:,:,:]
   end
   ONE = one(T1)

   # determine the dimensions of the diagonal blocks of real Schur form

   G = Array{T1,3}(undef,2,2,pc)
   W = Matrix{Float64}(undef,2,14)
   WX = Matrix{Float64}(undef,4,5)
   ba, p = MatrixEquations.sfstruct(A[:,:,KSCHUR])
   if adj
      #
      # Solve    A(j)'*X(j+1)*A(j) + C(j) = X(j) .
      #
      # The (K,L)th blocks of X(j), j = 1, ..., p are determined
      # starting from upper-left corner column by column by
      #
      #   A(j)(K,K)'*X(j+1)(K,L)*A(j)(L,L) - X(j)(K,L) = -C(j)(K,L) - R(j)(K,L)
      #
      # where
      #                K              L-1
      #   R(j)(K,L) = SUM {A(j)(I,K)'*SUM [X(j+1)(I,J)*A(j)(J,L)]}
      #               I=1             J=1
      #
      #                 K-1
      #             +  {SUM [A(j)(I,K)'*X(j+1)(I,L)]}*A(j)(L,L)
      #                 I=1
      i = 1
      @inbounds  for kk = 1:p
          dk = ba[kk]
          k = i:i+dk-1
          j = 1
          ir = 1:i-1
          for ll = 1:kk
              dl = ba[ll]
              j1 = j+dl-1
              l = j:j1
              Ckl = view(C,k,l,1:pc)
              y = view(G,1:dk,1:dl,1:pc)
              copyto!(y,Ckl)
              if kk > 1
                 # C(j+1)[l,k] = C(j+1)[l,ir]*A(j)[ir,k]
                 ic = 1:j1
                 for ii = 1:pc
                     ia = mod(ii-1,pa)+1
                     ii1 = mod(ii,pc)+1
                     mul!(view(C,l,k,ii1),view(C,l,ir,ii1),view(A,ir,k,ia))
                     #y += C(j+1)[ic,k]'*A(j)[ic,l]
                     mul!(view(y,:,:,ii),transpose(view(C,ic,k,ii1)),view(A,ic,l,ia),ONE,ONE)
                 end
              end
              Ckl[:,:,:] .= dpsylv2(adj, dk, dl, KSCHUR, view(A,k,k,1:pa), view(A,l,l,1:pa), y, W, WX)
              if ll == kk && dl == 2
                 for ii = 1:pc
                     temp = 0.5*(Ckl[1,2,ii]+Ckl[2,1,ii])
                     Ckl[1,2,ii] = temp; Ckl[2,1,ii] = temp
                 end
              end
              j += dl
              if ll < kk
                 # C(j+1)[l,k] += C(j+1)[k,l]'*A(j)[k,k]
                 for ii = 1:pc
                     ia = mod(ii-1,pa)+1
                     ii1 = mod(ii,pc)+1
                     mul!(view(C,l,k,ii1),transpose(view(C,k,l,ii1)),view(A,k,k,ia),ONE,ONE)
                 end
              end
          end
          if kk > 1
             # C(j)[ir,k] = C(j)[k,ir]'
             for ii = 1:pc
                 transpose!(view(C,ir,k,ii),view(C,k,ir,ii))
             end
          end
          i += dk
      end
   else
      #
      # Solve    A(j)*X(j)*A(j)' + C(j) = X(j+1) .
      #
      # The (K,L)th block of X(j) is determined starting from
      # bottom-right corner column by column by
      #
      #    A(j)(K,K)*X(j)(K,L)*A(j)(L,L)' - X(j+1)(K,L) = -C(j)(K,L) - R(j)(K,L)
      #
      # Where
      #
      #                 N               N
      #    R(j)(K,L) = SUM {A(j)(K,I)* SUM [X(j)(I,J)*A(j)(L,J)']} +
      #                I=K            J=L+1
      #
      #                N
      #             { SUM [A(j)(K,J)*X(j)(J,L)]}*A(j)(L,L)'
      #              J=K+1
      j = n
      for ll = p:-1:1
          dl = ba[ll]
          l = j-dl+1:j
          i = n
          ir = j+1:n
          for kk = p:-1:ll
              dk = ba[kk]
              i1 = i-dk+1
              k = i1:i
              Clk = view(C,l,k,1:pc)
              y = view(G,1:dl,1:dk,1:pc)
              copyto!(y,Clk)
              if ll < p
                 ic = i1:n
                 for ii = 1:pc
                     ia = mod(ii-1,pa)+1
                     # C(j)[k,l] = C(j)[k,ir]*A(j)[l,ir]'
                     mul!(view(C,k,l,ii),view(C,k,ir,ii),transpose(view(A,l,ir,ia)))
                     # y += (A(j)[k,ic]*C(j)[ic,l])'
                     mul!(view(y,:,:,ii),transpose(view(C,ic,l,ii)),transpose(view(A,k,ic,ia)),ONE,ONE)
                 end
              end
              Clk[:,:,:] .= dpsylv2(adj, dl, dk, KSCHUR, view(A,l,l,1:pa), view(A,k,k,1:pa), y, W, WX)
              i -= dk
              if i >= j
                 for ii = 1:pc
                     ia = mod(ii-1,pa)+1
                     # C(j)[k,l] += (A(j)[l,l]*C(j)[l,k])'
                     mul!(view(C,k,l,ii),transpose(view(C,l,k,ii)),transpose(view(A,l,l,ia)),ONE,ONE)
                 end
              else
                 break
              end
          end
          if ll < p
             ir = i+2:n
             for ii = 1:pc
                 # C(j)[ir,l] = C(j)[l,ir]'
                 transpose!(view(C,ir,l,ii),view(C,l,ir,ii))
             end
          end
          j -= dl
      end
   end
   return C[:,:,:]
end
function dpsylv2(REV::Bool, N1::Int, N2::Int, KSCHUR::Int, TL::StridedArray{T,3}, TR::StridedArray{T,3},
                 B::StridedArray{T,3}, W::AbstractMatrix{T}, WX::AbstractMatrix{T}) where {T}
#     To solve for the N1-by-N2 matrices X_j, j = 1, ..., P,
#     1 <= N1,N2 <= 2, in the P simultaneous equations:

#     if REV = true

#       TL_j'*X_(j+1)*TR_j - X_j = B_j, X_(P+1) = X_1  (1)

#     or if REV = false

#       TL_j*X_j*TR_j' - X_(j+1) = B_j, X_(P+1) = X_1  (2)

#     where TL_j is N1 by N1, TR_j is N2 by N2, B_j is N1 by N2,
#     and ISGN = 1 or -1.

#     NOTE: This routine is primarily intended to be used in conjuntion
#           with solvers for periodic Lyapunov equations. Thus, both
#           TL and TR are formed from the diagonal blocks of the same
#           matrix in periodic real Schur form. W and WX are working matrices
#           allocated only once in the caller routine.
#           TL and TR are assumed to have the same period, but
#           B may have different period than TL and TR. The period
#           of B must be an integer multiple of that of TL and TR.
#           In the interests of speed, this routine does not
#                   check the inputs for errors.

#     METHOD

#     An initial approximation X_j, j=1, ..., P, is computed by reducing
#     the system (1) or (2) to an equivalent single Lyapunov equation.
#     Then, the accuracy of the solution is iteratively improved by
#     performing forward or backward sweeps depending on the stability
#     of eigenvalues. A maximum of 30 sweeps are performed.

#     REFERENCES

#     [1] A. Varga.
#         Periodic Lyapunov equations: some applications and new algorithms.
#         Int. J. Control, vol, 67, pp, 69-87, 1997.

	IND(J,P) = mod( J-1, P ) + 1
   P = size(TL,3)
   PB = size(B,3)
   i1 = 1:N1; i2 = 1:N2;
   Xw = view(WX,:,1:4)
   Yw = view(WX,:,5)
   Z = view(W,:,7:8)
   X = Array{T,3}(undef, N1, N2, PB)

   # Quick return if possible.
   (N1 == 0 || N2 == 0) && (return X)
   if P == 1 && PB == 1
      copyto!(view(X,i1,i2,1), view(B,i1,i2,1))
      MatrixEquations.lyapdsylv2!(REV, view(X, i1, i2, 1),  N1, N2, view(TL, i1, i1, 1), view(TR, i2, i2, 1), Xw, Yw)
      return X[:,:,:]
   end
   # partition working space
   AL = view(W,:,1:2)
   AR = view(W,:,3:4)
   Q = view(W,:,5:6)
   ZOLD = view(W,:,9:10)
   ATMP = view(W,:,11:12)
   BTMP = view(W,:,13:14)

   EPSM = 2*eps(T)

#     Define

#       AL(i,j) := TL_{j+i-1}*...*TL_{j+1}*TL_{j}, AL(j,j) := I;
#       AR(i,j) := TR_{j+i-1}*...*TR_{j+1}*TR_{j}, AR(j,j) := I;

#     If REV = true, compute an initial approximation for
#     X_{KSCHUR+1} = Z by solving

#       AL(KSCHUR+P+1,KSCHUR+1)'*Z*AR(KSCHUR+P+1,KSCHUR+1)  - Z = Qr(KSCHUR+1)

#     where

#       Qr(j)  = B_{j} + TL_{j}'*Qr(j+1)*TR_{j},  Qr(KSCHUR) = B(KSCHUR),
#                for j = KSCHUR+P-1, ..., KSCHUR+1

#     If REV = false, compute an initial approximation for
#     X_{KSCHUR} = Z by solving

#       AL(KSCHUR+P,KSCHUR)*Z*AR(KSCHUR+P,KSCHUR)'  - Z = Qf(KSCHUR)

#     where

#       Qf(j)  = B_{j} + TL_{j}*Qf(j-1)*TR_{j}',  Qf(KSCHUR) = B(KSCHUR),
#                for j = KSCHUR+1, ..., KSCHUR+P-1
   x = 1
   if REV
      L1 = KSCHUR+PB-1
      L2 = KSCHUR+1
      LSTEP = -1
   else
      L1 = KSCHUR+1
      L2 = KSCHUR+PB-1
      LSTEP = 1
   end
   AL[ i1, i1 ] = TL[ i1, i1, KSCHUR ]
   AR[ i2, i2 ] = TR[ i2, i2, KSCHUR ]
   Q[i1,i2] = B[i1,i2,KSCHUR]

   if  N1 == 1 && N2 == 1
       for JJ = L1:LSTEP:L2
           J = IND( JJ, P )
           JB = IND( JJ, PB )
           X11 =  TL[ 1, 1, J ]
           Y11 = TR[ 1, 1, J ]
           AL[ 1, 1 ] = AL[ 1, 1] * X11
           AR[ 1, 1 ] = AR[ 1, 1] * Y11
           Q[ 1, 1 ] = B[ 1, 1, JB] + X11 * Q[ 1, 1 ] * Y11
       end
   elseif N1 == 1 && N2 == 2
       for JJ = L1:LSTEP:L2
           J = IND( JJ, P )
           JB = IND( JJ, PB )
           X11 =  TL[ 1, 1, J ]
           AL[ 1, 1 ] = AL[ 1, 1] * X11
           Q[1,1] = X11*Q[1,1]
           Q[1,2] = X11*Q[1,2]

           Y11 = TR[ 1, 1, J ]
           Y22 = TR[ 2, 2, J ]
           Y12 = TR[ 1, 2, J ]
           if REV
               AR[1,2] = AR[1,1]*Y12 + AR[1,2]*Y22
               AR[1,1] = AR[1,1]*Y11
               AR[2,2] = AR[2,1]*Y12 + AR[2,2]*Y22
               AR[2,1] = AR[2,1]*Y11
               Q[1,2] = Q[1,1]*Y12 + Q[1,2]*Y22 + B[ 1, 2, JB ]
               Q[1,1] = Q[1,1]*Y11 + B[ 1, 1, JB ]
           else
               AR[1,1] = Y11*AR[1,1] + Y12*AR[2,1]
               AR[2,1] = Y22*AR[2,1]
               AR[1,2] = Y11*AR[1,2] + Y12*AR[2,2]
               AR[2,2] = Y22*AR[2,2]
               Q[1,1] = Q[1,1]*Y11 + Q[1,2]*Y12 + B[ 1, 1, JB ]
               Q[1,2] = Q[1,2]*Y22 + B[ 1, 2, JB ]
           end
       end
   elseif N1 == 2 && N2 == 1
       for JJ = L1:LSTEP:L2
           J = IND( JJ, P )
           JB = IND( JJ, PB )
           X11 = TL[ 1, 1, J ]
           X12 = TL[ 1, 2, J ]
           X22 = TL[ 2, 2, J ]
           if REV
              AL[1,2] = AL[1,1]*X12 + AL[1,2]*X22
              AL[1,1] = AL[1,1]*X11
              AL[2,2] = AL[2,1]*X12 + AL[2,2]*X22
              AL[2,1] = AL[2,1]*X11
              Q[2,1] = X12*Q[1,1] + X22*Q[2,1]
              Q[1,1] = X11*Q[1,1]
           else
              AL[1,1] = X11*AL[1,1] + X12*AL[2,1]
              AL[2,1] = X22*AL[2,1]
              AL[1,2] = X11*AL[1,2] + X12*AL[2,2]
              AL[2,2] = X22*AL[2,2]
              Q[1,1] = X11*Q[1,1] + X12*Q[2,1]
              Q[2,1] = X22*Q[2,1]
            end

            Y11 = TR[ 1, 1, J ]
            AR[1,1] = AR[1,1]*Y11
            Q[1,1] = Q[1,1]*Y11 + B[ 1, 1, JB ]
            Q[2,1] = Q[2,1]*Y11 + B[ 2, 1, JB ]
       end
   elseif N1 == 2 && N2 == 2
       for JJ = L1:LSTEP:L2
          J = IND( JJ, P )
          JB = IND( JJ, PB )
          X11 = TL[ 1, 1, J ]
          X12 = TL[ 1, 2, J ]
          X22 = TL[ 2, 2, J ]
          if REV
             AL[1,2] = AL[1,1]*X12 + AL[1,2]*X22
             AL[1,1] = AL[1,1]*X11
             AL[2,2] = AL[2,1]*X12 + AL[2,2]*X22
             AL[2,1] = AL[2,1]*X11
             Q[2,1] = X12*Q[1,1] + X22*Q[2,1]
             Q[1,1] = X11*Q[1,1]
             Q[2,2] = X12*Q[1,2] + X22*Q[2,2]
             Q[1,2] = X11*Q[1,2]
          else
             AL[1,1] = X11*AL[1,1] + X12*AL[2,1]
             AL[2,1] = X22*AL[2,1]
             AL[1,2] = X11*AL[1,2] + X12*AL[2,2]
             AL[2,2] = X22*AL[2,2]
             Q[1,1] = X11*Q[1,1] + X12*Q[2,1]
             Q[2,1] = X22*Q[2,1]
             Q[1,2] = X11*Q[1,2] + X12*Q[2,2]
             Q[2,2] = X22*Q[2,2]
          end

          Y11 = TR[ 1, 1, J ]
          Y12 = TR[ 1, 2, J ]
          Y22 = TR[ 2, 2, J ]
          if REV
             AR[1,2] = AR[1,1]*Y12 + AR[1,2]*Y22
             AR[1,1] = AR[1,1]*Y11
             AR[2,2] = AR[2,1]*Y12 + AR[2,2]*Y22
             AR[2,1] = AR[2,1]*Y11
             Q[1,2] = Q[1,1]*Y12 + Q[1,2]*Y22 + B[ 1, 2, JB ]
             Q[1,1] = Q[1,1]*Y11 + B[ 1, 1, JB ]
             Q[2,2] = Q[2,1]*Y12 + Q[2,2]*Y22 + B[ 2, 2, JB ]
             Q[2,1] = Q[2,1]*Y11 + B[ 2, 1, JB ]
          else
             AR[1,1] = Y11*AR[1,1] + Y12*AR[2,1]
             AR[2,1] = Y22*AR[2,1]
             AR[1,2] = Y11*AR[1,2] + Y12*AR[2,2]
             AR[2,2] = Y22*AR[2,2]
             Q[1,1] = Q[1,1]*Y11 + Q[1,2]*Y12 + B[ 1, 1, JB ]
             Q[1,2] = Q[1,2]*Y22 + B[ 1, 2, JB ]
             Q[2,1] = Q[2,1]*Y11 + Q[2,2]*Y12 + B[ 2, 1, JB ]
             Q[2,2] = Q[2,2]*Y22 + B[ 2, 2, JB ]
          end
       end
   end
   Z[i1,i2] = -Q[i1,i2]
#     Compute X_[KSCHUR+1] (if REV=.TRUE.) or X_KSCHUR (if REV=.FALSE.).
   MatrixEquations.lyapdsylv2!(REV, view(Z,i1,i2), N1, N2, view(AL,i1,i1), view(AR,i2,i2), Xw, Yw)
   XNORM = norm(view(Z,i1,i2), Inf)
   XNORM == 0 && (XNORM = one(T))

#     Determine the type of iteration to use

   if N1 == 1
      EL = abs( AL[1,1] )
   else
      EL = sqrt( abs( AL[1,1]*AL[2,2] - AL[1,2]*AL[2,1] ) )
   end
   if N2 == 1
      ER = abs( AR[1,1] )
   else
      ER = sqrt( abs( AR[1,1]*AR[2,2] - AR[1,2]*AR[2,1] ) )
   end
   STAB = EL*ER < one(T)

#     Save initial X_[KSCHUR+1].

   ZOLD[i1,i2] = Z[i1,i2]

#     Set iteration indices.

   if REV
      if STAB
          L1 = KSCHUR + PB
          L2 = KSCHUR + 1
          LSTEP = -1
          ITIND = -1
      else
          L1 = KSCHUR + 1
          L2 = KSCHUR + PB
          LSTEP = 1
          ITIND = 0
      end
   else
      if STAB
          L1 = KSCHUR
          L2 = KSCHUR + PB - 1
          LSTEP = 1
          ITIND = 0
      else
          L1 = KSCHUR + PB - 1
          L2 = KSCHUR
          LSTEP = -1
          ITIND = -1
      end
   end
   for ITER = 1:30
       if N1 == 1 && N2 == 1
          if STAB

#             Use direct recursion.

             for JJ = L1:LSTEP:L2
               J = IND( JJ, P )
               JB = IND( JJ, PB )
               J1 = IND( JJ+ITIND+1, PB )
               Z[1,1] = TL[ 1, 1, J ] * Z[1,1] * TR[ 1, 1, J ] - B[ 1, 1, JB ]
               X[ 1, 1, J1 ] = Z[ 1, 1 ]
             end
          else

#             Use inverse recursion.

             for JJ = L1:LSTEP:L2
                  J = IND( JJ, P )
                  JB = IND( JJ, PB )
                  J1 = IND( JJ+ITIND+1, PB )
                  Z[1,1] = (Z[1,1] + B[ 1, 1, JB ] ) / TL[ 1, 1, J ] / TR[ 1, 1, J ]
                  X[ 1, 1, J1 ] = Z[ 1, 1 ]
             end
          end
       elseif N1 == 1 && N2 == 2
          if STAB

#              Use direct recursion.

             for JJ = L1:LSTEP:L2
                 J = IND( JJ, P )
                 JB = IND( JJ, PB )
                 J1 = IND( JJ+ITIND+1, PB )
                 X11 = TL[ 1, 1, J ]
                 Y11 = TR[ 1, 1, J ]
                 Y22 = TR[ 2, 2, J ]
                 if REV
                    Y12 = TR[ 1, 2, J ]
                    Y21 = TR[ 2, 1, J ]
                 else
                    Y12 = TR[ 2, 1, J ]
                    Y21 = TR[ 1, 2, J ]
                 end

                 Z[1,1] = X11*Z[1,1]
                 Z[1,2] = X11*Z[1,2]
                 TEMP = Z[1,1]*Y11 + Z[1,2]*Y21
                 Z[1,2] = Z[1,1]*Y12 + Z[1,2]*Y22 - B[ 1, 2, JB ]
                 Z[1,1] = TEMP - B[ 1, 1, JB ]
                 X[ 1, 1, J1 ] = Z[ 1, 1 ]
                 X[ 1, 2, J1 ] = Z[ 1, 2 ]
              end
          else

#              Use inverse recursion.

             for JJ = L1:LSTEP:L2
                 J = IND( JJ, P )
                 JB = IND( JJ, PB )
                 J1 = IND( JJ+ITIND+1, PB )
                 BTMP[ 1, 1 ] = ( Z[1,1] + B[1,1,JB] ) / TL[ 1, 1, J ]
                 BTMP[ 2, 1 ] = ( Z[1,2] + B[1,2,JB] ) / TL[ 1, 1, J ]
                 ATMP[ 1, 1 ] = TR[ 1, 1, J ]
                 ATMP[ 2, 2 ] = TR[ 2, 2, J ]
                 if REV
                    ATMP[ 1, 2 ] = TR[ 2, 1, J ]
                    ATMP[ 2, 1 ] = TR[ 1, 2, J ]
                 else
                    ATMP[ 1, 2 ] = TR[ 1, 2, J ]
                    ATMP[ 2, 1 ] = TR[ 2, 1, J ]
                 end

                 #CALL DGESV( 2, 1, ATMP, 2, JPIV, BTMP, 2, INFO )
                 luslv!(ATMP,view(BTMP,1:2,1:1))
                 #BTMP = ATMP\BTMP
                 Z[ 1, 1 ] = BTMP[ 1, 1 ]
                 Z[ 1, 2 ] = BTMP[ 2, 1 ]
                 X[ 1, 1, J1 ] = Z[ 1, 1 ]
                 X[ 1, 2, J1 ] = Z[ 1, 2 ]
              end
          end
       elseif N1 == 2 && N2 == 1
          if STAB

#              Use direct recursion.

             for JJ = L1:LSTEP:L2
                 J = IND( JJ, P )
                 JB = IND( JJ, PB )
                 J1 = IND( JJ+ITIND+1, PB )
                 X11 = TL[ 1, 1, J ]
                 X22 = TL[ 2, 2, J ]
                 if REV
                    X12 = TL[ 2, 1, J ]
                    X21 = TL[ 1, 2, J ]
                 else
                    X12 = TL[ 1, 2, J ]
                    X21 = TL[ 2, 1, J ]
                 end
                 Y11 = TR[ 1, 1, J ]

                 TEMP = X11*Z[1,1] + X12*Z[2,1]
                 Z[2,1] = X21*Z[1,1] + X22*Z[2,1]
                 Z[1,1] = TEMP
                 Z[1,1] = Z[1,1]*Y11 - B[ 1, 1, JB ]
                 Z[2,1] = Z[2,1]*Y11 - B[ 2, 1, JB ]
                 X[ 1, 1, J1 ] = Z[ 1, 1 ]
                 X[ 2, 1, J1 ] = Z[ 2, 1 ]
              end
          else

#              Use inverse recursion.

             for JJ = L1:LSTEP:L2
                 J = IND( JJ, P )
                 JB = IND( JJ, PB )
                 J1 = IND( JJ+ITIND+1, PB )
                 Z[ 1, 1 ] = ( Z[1,1] + B[1,1,JB] ) / TR[ 1, 1, J ]
                 Z[ 2, 1 ] = ( Z[2,1] + B[2,1,JB] ) / TR[ 1, 1, J ]
                 ATMP[ 1, 1 ] = TL[ 1, 1, J ]
                 ATMP[ 2, 2 ] = TL[ 2, 2, J ]
                 if REV
                    ATMP[ 1, 2 ] = TL[ 2, 1, J ]
                    ATMP[ 2, 1 ] = TL[ 1, 2, J ]
                 else
                    ATMP[ 1, 2 ] = TL[ 1, 2, J ]
                    ATMP[ 2, 1 ] = TL[ 2, 1, J ]
                 end

                 #CALL DGESV( 2, 1, ATMP, 2, JPIV, Z, 2, INFO )
                 luslv!(ATMP,view(Z,1:2,1:1))
                 #Z = ATMP\Z
                 X[ 1, 1, J1 ] = Z[ 1, 1 ]
                 X[ 2, 1, J1 ] = Z[ 2, 1 ]
              end
          end
       elseif N1 == 2 && N2 == 2
          if STAB

#              Use direct recursion.

             for JJ = L1:LSTEP:L2
                 J = IND( JJ, P )
                 JB = IND( JJ, PB )
                 J1 = IND( JJ+ITIND+1, PB )
                 X11 = TL[ 1, 1, J ]
                 X22 = TL[ 2, 2, J ]
                 Y11 = TR[ 1, 1, J ]
                 Y22 = TR[ 2, 2, J ]
                 if REV
                    X12 = TL[ 2, 1, J ]
                    X21 = TL[ 1, 2, J ]
                    Y12 = TR[ 1, 2, J ]
                    Y21 = TR[ 2, 1, J ]
                 else
                    X12 = TL[ 1, 2, J ]
                    X21 = TL[ 2, 1, J ]
                    Y12 = TR[ 2, 1, J ]
                    Y21 = TR[ 1, 2, J ]
                 end

                 TEMP = X11*Z[1,1] + X12*Z[2,1]
                 Z[2,1] = X21*Z[1,1] + X22*Z[2,1]
                 Z[1,1] = TEMP
                 TEMP = X11*Z[1,2] + X12*Z[2,2]
                 Z[2,2] = X21*Z[1,2] + X22*Z[2,2]
                 Z[1,2] = TEMP
                 TEMP = Z[1,1]*Y11 + Z[1,2]*Y21
                 Z[1,2] = Z[1,1]*Y12 + Z[1,2]*Y22 - B[ 1, 2, JB ]
                 Z[1,1] = TEMP - B[ 1, 1, JB ]
                 TEMP = Z[2,1]*Y11 + Z[2,2]*Y21
                 Z[2,2] = Z[2,1]*Y12 + Z[2,2]*Y22 - B[ 2, 2, JB ]
                 Z[2,1] = TEMP - B[ 2, 1, JB ]
                 X[ 1, 1, J1 ] = Z[ 1, 1 ]
                 X[ 2, 1, J1 ] = Z[ 2, 1 ]
                 X[ 1, 2, J1 ] = Z[ 1, 2 ]
                 X[ 2, 2, J1 ] = Z[ 2, 2 ]
              end
          else

#              Use inverse recursion.

             for JJ = L1:LSTEP:L2
                 J = IND( JJ, P )
                 JB = IND( JJ, PB )
                 J1 = IND( JJ+ITIND+1, PB )
                 BTMP[ 1, 1 ] = Z[ 1, 1 ] + B[ 1, 1, JB ]
                 BTMP[ 2, 1 ] = Z[ 1, 2 ] + B[ 1, 2, JB ]
                 BTMP[ 1, 2 ] = Z[ 2, 1 ] + B[ 2, 1, JB ]
                 BTMP[ 2, 2 ] = Z[ 2, 2 ] + B[ 2, 2, JB ]
                 ATMP[ 1, 1 ] = TR[ 1, 1, J ]
                 ATMP[ 2, 2 ] = TR[ 2, 2, J ]
                 if REV
                    ATMP[ 1, 2 ] = TR[ 2, 1, J ]
                    ATMP[ 2, 1 ] = TR[ 1, 2, J ]
                 else
                    ATMP[ 1, 2 ] = TR[ 1, 2, J ]
                    ATMP[ 2, 1 ] = TR[ 2, 1, J ]
                 end

                 #CALL DGESV( 2, 2, ATMP, 2, JPIV, BTMP, 2, INFO )
                 luslv!(ATMP,view(BTMP,1:2,1:2))
                 #BTMP = ATMP\BTMP
                 Z[ 1, 1 ] = BTMP[ 1, 1 ]
                 Z[ 1, 2 ] = BTMP[ 2, 1 ]
                 Z[ 2, 1 ] = BTMP[ 1, 2 ]
                 Z[ 2, 2 ] = BTMP[ 2, 2 ]

                 ATMP[ 1, 1 ] = TL[ 1, 1, J ]
                 ATMP[ 2, 2 ] = TL[ 2, 2, J ]
                 if REV
                    ATMP[ 1, 2 ] = TL[ 2, 1, J ]
                    ATMP[ 2, 1 ] = TL[ 1, 2, J ]
                 else
                    ATMP[ 1, 2 ] = TL[ 1, 2, J ]
                    ATMP[ 2, 1 ] = TL[ 2, 1, J ]
                 end

                 #CALL DGESV( 2, 2, ATMP, 2, JPIV, Z, 2, INFO )
                 luslv!(ATMP,view(Z,1:2,1:2))
                 #Z = ATMP\Z
                 X[ 1, 1, J1 ] = Z[ 1, 1 ]
                 X[ 1, 2, J1 ] = Z[ 1, 2 ]
                 X[ 2, 1, J1 ] = Z[ 2, 1 ]
                 X[ 2, 2, J1 ] = Z[ 2, 2 ]
              end
          end
       end
       DNORM = zero(T)
       for J = 1:N2
          for I = 1:N1
             DNORM = max( DNORM, abs( ZOLD[I,J] - Z[I,J] ) )
             ZOLD[ I, J ] = Z[ I, J ]
          end
       end
       #println("XNORM = $XNORM DNORM = $DNORM")
       DNORM <= EPSM*XNORM && (return -X)
   end
   @warn "iterative process not converging: solution may be inaccurate"
   return -X
#   END of DPSYLV2
end
@inline function luslv!(A::AbstractMatrix{T}, B::AbstractMatrix{T}) where T
   #
   #  fail = luslv!(A,B)
   #
   # This function is a speed-oriented implementation of a Gaussion-elimination based
   # solver of small order linear equations of the form A*X = B. The computed solution X
   # overwrites the vector B, while the resulting A contains in its upper triangular part,
   # the upper triangular factor U of its LU decomposition.
   # The diagnostic output parameter fail, of type Bool, is set to false in the case
   # of normal return or is set to true if the exact singularity of A is detected
   # or if the resulting B has non-finite components.
   #
   n, m = size(B)
   @inbounds begin
         for k = 1:n
            # find index max
            kp = k
            if k < n
                amax = abs(A[k, k])
                for i = k+1:n
                    absi = abs(A[i,k])
                    if absi > amax
                        kp = i
                        amax = absi
                    end
                end
            end
            iszero(A[kp,k]) && return true
            if k != kp
               # Interchange
               for i = 1:n
                   tmp = A[k,i]
                   A[k,i] = A[kp,i]
                   A[kp,i] = tmp
               end
               for j = 1:m
                   tmp = B[k,j]
                   B[k,j] = B[kp,j]
                   B[kp,j] = tmp
               end
            end
            # Scale first column
            Akkinv = inv(A[k,k])
            i1 = k+1:n
            Ak = view(A,i1,k)
            rmul!(Ak,Akkinv)
            # Update the rest of A and B
            for j = k+1:n
                axpy!(-A[k,j],Ak,view(A,i1,j))
            end
            for j = 1:m
                axpy!(-B[k,j],Ak,view(B,i1,j))
            end
         end
         ldiv!(UpperTriangular(A), view(B,:,:))
         return any(!isfinite, B)
   end
end


#=
Given a array of numbers representing the stock prices of a company in chronological order, write a function that calculates the maximum profit you could have made from buying and selling that stock once. You must buy before you can sell it.

For example, given [9, 11, 8, 5, 7, 10], you should return 5, since you could buy the stock at 5 dollars and sell it at 10 dollars.
=#

using Test
include("Solutions/problem47_find_maximum_profit.jl")

@test maximumprofit([9, 11, 8, 5, 7, 10]) == 5


module BS

using ..Distributions
import ..Roots

#
# types
#

struct D₁Parts{T}
    numerator::T
    denominator::T
end

abstract type BSOption end

"European call option."
struct EuropeanCall <: BSOption end

"European put option."
struct EuropeanPut <: BSOption end

@inline N(x) = cdf(Normal(), x)

@inline D₁Parts(s, k, t, σ) = D₁Parts(log(s/k) + (σ^2/2)*t, σ*sqrt(t))

@inline d₁(parts::D₁Parts) = parts.numerator / parts.denominator
@inline d₂(parts::D₁Parts) = d₁(parts) - parts.denominator

@inline d₁(s, k, t, σ) = d₁(D₁Parts(s, k, t, σ))
@inline d₂(s, k, t, σ) = d₂(D₁Parts(s, k, t, σ))

@inline function bscall(s, k, t, σ)
    parts = D₁Parts(s, k, t, σ)
    return s*N(d₁(parts)) - k*N(d₂(parts))
end

@inline function bsput(s, k, t, σ)
    parts = D₁Parts(s, k, t, σ)
    return k*N(-d₂(parts)) - s*N(-d₁(parts))
end

@inline price(::EuropeanCall, s, k, t, σ) = bscall(s, k, t, σ)
@inline price(::EuropeanPut, s, k, t, σ) = bsput(s, k, t, σ)

function impvol(opt::BSOption, observed_price, s, k, t, interval::Tuple=infer_impvol_interval(opt, observed_price, s, k, t))
    f(σ) = observed_price - price(opt, s, k, t, σ)
    return Roots.find_zero(f, interval, Roots.Bisection())
end

function infer_impvol_interval(opt::BSOption, observed_price, s, k, t)
    # most instruments have volatility values below 100%
    min_vol = 0.0
    max_vol = 1.0
    if observed_price <= price(opt, s, k, t, max_vol)
        return (min_vol, max_vol)
    end

    # will increment max_vol 10x until we reach a final interval
    while true
        min_vol = max_vol
        max_vol = min_vol * 10
        if observed_price <= price(opt, s, k, t, max_vol)
            return (min_vol, max_vol)
        end
    end
end

#
# Greeks
#

@inline delta(::EuropeanCall, s, k, t, σ) = N(d₁(s, k, t, σ))
@inline delta(::EuropeanPut, s, k, t, σ) = N(d₁(s, k, t, σ)) - 1

@inline function theta(::EuropeanCall, s, k, t, r, σ)
    parts = D₁Parts(s, k, t, σ)
    return -(s * pdf(Normal(), d₁(parts)) * σ) / ( 2*sqrt(t) ) - r*k*N(d₂(parts))
end

@inline function theta(::EuropeanPut, s, k, t, r, σ)
    parts = D₁Parts(s, k, t, σ)
    return -(s * pdf(Normal(), d₁(parts)) * σ) / ( 2*sqrt(t) ) + r*k*N(-d₂(parts))
end

@inline function gamma(s, k, t, σ)
    parts = D₁Parts(s, k, t, σ)
    return pdf(Normal(), d₁(parts)) / (s * parts.denominator)
end
@inline gamma(::BSOption, s, k, t, σ) = gamma(s, k, t, σ)

@inline vega(s, k, t, σ) = s*sqrt(t)*pdf(Normal(), d₁(s, k, t, σ))
@inline vega(::BSOption, s, k, t, σ) = vega(s, k, t, σ)

@inline rho(::EuropeanCall, s, k, t, σ) = k*t*N(d₂(s, k, t, σ))
@inline rho(::EuropeanPut, s, k, t, σ) = -k*t*N(-d₂(s, k, t, σ))

end # module


#Test start to finish

using SimpleNeuralNet
using LinearAlgebra #for norm()
using Printf

#Build the neural net
nNeurons = [4; 5; 2]; #number of neurons by layer
nn = neuralNet(nNeurons);

#Define some important functions
SimpleNeuralNet.nnAct(z) = sigmoid.(z); #activation function (assume the same everywhere)
SimpleNeuralNet.nnDAct(z) = dsigmoid.(z); #derivative of activation function (assume the same everywhere)
SimpleNeuralNet.nnCost(datay,y) = 0.5*norm(y-datay,2)^2;
SimpleNeuralNet.nnDCost(datay,y) = y - datay;

#Initialize the weights and biases to random values
nn.w = rand(length(nn.w));
nn.b = rand(length(nn.b));

#Define the training parameters
lr=1.0;      #learning rate
epochs=1000; #epochs

#Define the dataset
nData = 10;
datax = rand(nn.nNeurons[1],nData);   #draw input dataset at random
datay = rand(nn.nNeurons[end],nData); #draw output dataset at random

#Compute the initial costs
c0 = nnComputeCosts(nn,datax,datay;print=false);

#Train
nnTrain(nn,datax,datay,epochs,lr;verbose=true);

#Compute the final costs
cf = nnComputeCosts(nn,datax,datay;print=false);

#Print a comparison of initial/final costs
for d=1:length(cf)
    @printf("dataset %3d: initial cost=%10.8f final cost=%10.8f\n",d,c0[d],cf[d]);
end


"""
    geometric_factors(x, y, Dr, Ds)
    geometric_factors(x, y, z, Dr, Ds, Dt)

Compute metrics of mappings between "real" elements and reference elements,
outward pointing normals on faces of every elements, and Jacobian.

x,y,z are arrays of coordinates, and Dr, Ds, Dt are nodal differentiation matrices

Geometric terms in 3D are constructed to ensure satisfaction of free-stream
preservation using the curl-based construction of David Kopriva (2001).

"""

# 2D version
function geometric_factors(x, y, Dr, Ds)
    "Transformation and Jacobian"

    xr = Dr*x;   xs = Ds*x;
    yr = Dr*y;   ys = Ds*y;

    J = -xs.*yr + xr.*ys;
    rxJ =  ys;  sxJ = -yr;
    ryJ = -xs;  syJ =  xr;

    return rxJ, sxJ, ryJ, syJ, J
end

# 3D version. Filters = tuple of filtering matrices.
function geometric_factors(x, y, z, Dr, Ds, Dt, Filters=(I,I,I))

    xr = Dr*x;  xs = Ds*x;  xt = Dt*x
    yr = Dr*y;  ys = Ds*y;  yt = Dt*y
    zr = Dr*z;  zs = Ds*z;  zt = Dt*z

    Fr = (Dr*y).*z
    Fs = (Ds*y).*z
    Ft = (Dt*y).*z
    Fr,Fs,Ft = ((A,x)->A*x).(Filters,(Fr,Fs,Ft))
    rxJ = Dt*(Fs) - Ds*(Ft)
    sxJ = Dr*(Ft) - Dt*(Fr)
    txJ = Ds*(Fr) - Dr*(Fs)

    Fr = (Dr*x).*z
    Fs = (Ds*x).*z
    Ft = (Dt*x).*z
    Fr,Fs,Ft = ((A,x)->A*x).(Filters,(Fr,Fs,Ft))
    ryJ = -(Dt*(Fs) - Ds*(Ft))
    syJ = -(Dr*(Ft) - Dt*(Fr))
    tyJ = -(Ds*(Fr) - Dr*(Fs))

    Fr = (Dr*y).*x
    Fs = (Ds*y).*x
    Ft = (Dt*y).*x
    Fr,Fs,Ft = ((A,x)->A*x).(Filters,(Fr,Fs,Ft))
    rzJ = -(Dt*(Fs) - Ds*(Ft))
    szJ = -(Dr*(Ft) - Dt*(Fr))
    tzJ = -(Ds*(Fr) - Dr*(Fs))

    J = @. xr*(ys*zt-zs*yt) - yr*(xs*zt-zs*xt) + zr*(xs*yt-ys*xt)

    return rxJ, sxJ, txJ, ryJ, syJ, tyJ, rzJ, szJ, tzJ, J
end


module PolyharmonicInterpolation
using LinearAlgebra

export PolyharmonicInterpolator,polyharmonicSpline

struct PolyharmonicInterpolator
	knots::AbstractArray
	values::Vector{Float64}
	nPts::Int
	dim::Int
	order::Int
	w::Vector{Float64}
	v::Vector{Float64}
end

PolyharmonicInterpolator(knots::AbstractArray{T},values::Vector{T},
order::Int=2) where T <: Real = begin

	# Check to make sure knots is N x d and values is N x 1
	if size(knots)[1] != length(values)
		error("knots and values length must agree")
	end

	if length(size(knots)) == 1
		nPts = length(knots)
		dim = 1
	else
		nPts,dim = size(knots)
	end

	# initialize w and v
	w = Vector{Float64}(undef,nPts)
	v = Vector{Float64}(undef,dim+1)

	# Fill npts and dim using the dimension of knots
	computeWeights(PolyharmonicInterpolator(knots,values,nPts,dim,order,w,v))

end

function polyharmonicRBF(r::Float64,k::Int)

	if r >= 1
		val = (k % 2 == 0) ? r^k*log(r) : r^k
	else
		val = (k % 2 == 0) ? r^(k-1)*log(r^r) : r^k
	end

	return val

end

function createBlockMatrix(p::PolyharmonicInterpolator)

	# Create the matrix of differences between the knots
	dist = [norm(p.knots[i,:] - p.knots[j,:]) for i in 1:p.nPts, j in 1:p.nPts]

	# Create the matrix A
	A = map(r->polyharmonicRBF(r,p.order),dist)

	# Create the matrix B
	B = [ones(p.nPts) p.knots];

	return A,B

end

function computeWeights(p::PolyharmonicInterpolator)

	# Create block matrices used to solve for weights
	A,B = createBlockMatrix(p)

	# Make block matrix for the symmetric linear system
	blockMat = [A B; transpose(B) zeros(p.dim+1,p.dim+1)]

	# Get weights nPts and number of dimensions
	weights = blockMat \ [p.values; zeros(p.dim+1,1)]

	p.w .= weights[1:p.nPts]
	p.v .= weights[p.nPts+1:end]

	return p

end

function polyharmonicSpline(p::PolyharmonicInterpolator,interpPoints::AbstractArray)

	# Check to see the dimensions of the interpPoints array agree
	let
		nInterpDims = nothing
		try
			_, nInterpDims = size(interpPoints)
		catch
			nInterpDims = 1
		end

		if p.dim != nInterpDims
			error("dimension of knots and points to be interpolated must agree")
		end
	end

	# Initialize array
	interpValues = zeros(size(interpPoints)[1])

	for (i,point) in enumerate(eachrow(interpPoints))

		tmppoint =[1.0;point]

		# Get distance to all scattered points
		dist = sqrt.(sum((p.knots .- point').^2,dims=2))

		# reduce to normal size (Julia returns a 2 dimensional row array for some reason...)
		dist = reduce(vcat,dist)

		# Get distance between basis and point of interest and compute the
		# polyharmonic basis
		phi = map(r->polyharmonicRBF(r,p.order),dist)

		interpValues[i] = p.w'phi + p.v'tmppoint

	end

	return interpValues

end

end


using LinearAlgebraicRepresentation, ViewerGL, SparseArrays
Lar = LinearAlgebraicRepresentation; GL = ViewerGL
using Base.union

# 3D Boolean example generation
#-------------------------------------------------------------------------------
n,m,p = 1,1,1
V,(VV,EV,FV,CV) = Lar.cuboidGrid([n,m,p],true)
cube = V,FV,EV

V,FV = Lar.sphere()()
EV = Lar.simplexFacets(FV)
sphere = V,FV,EV

# three cubes in "assembly"
assembly = Lar.Struct([ cube,
    Lar.t(.3,.4,.25), Lar.r(pi/5,0,0), Lar.r(0,0,pi/12), cube,
    Lar.t(-.2,.4,-.2), Lar.r(0,pi/5,0), Lar.r(0,pi/12,0), cube ])

V,FV,EV = Lar.struct2lar(assembly)
GL.VIEW([ GL.GLGrid(V,FV), GL.GLFrame ]);

W, (copEV, copFE, copCF), boolmatrix = Lar.bool3d(assembly)
Matrix(boolmatrix)
#three-chains = [ for k = 1:3]

A = boolmatrix[:,2]
B = boolmatrix[:,3]
C = boolmatrix[:,4]
AorB = A .| B
AandB = A .& B
AxorB = AorB .& (.! AandB)
AorBorC = A .| B .| C
AorBorC = .|(A,B,C)
AandBandC = A .& B .& C
AandBandC = .&(A,B,C)
AminusBminusC = .&(A, .!(B .| C)) # A - B - C

unione = .|(A,B,C)
unione = Matrix(copCF)' * Int.(unione) # coord vector of Faces
intersection = Matrix(copCF)' * Int.(AandBandC) # coord vector of Faces
difference = Matrix(copCF)' * Int.(AminusBminusC) # coord vector of Faces

V,CVs,FVs,EVs = Lar.pols2tria(W, copEV, copFE, copCF) # whole assembly
Fs = unione
V,CVs,FVs,EVs = Lar.pols2tria(W, copEV, copFE, copCF, Fs) # part of assembly

# EV = Lar.cop2lar(copEV)
# EVor = [ev for (k,ev) in enumerate(EV) if abs(union[k])==1 ]
# EVand = [ev for (k,ev) in enumerate(EV) if abs(intersection[k])==1 ]
# EVxor = [ev for (k,ev) in enumerate(EV) if abs(xor[k])==1 ]

GL.VIEW(GL.GLExplode(V,FVs,1.,1.,1.,99,1));
GL.VIEW(GL.GLExplode(V,EVs,1.,1.,1.,99,1));
meshes = GL.GLExplode(V,CVs[2:end],1.5,1.5,1.5,99,1);
GL.VIEW( push!( meshes, GL.GLFrame) );

# GL.VIEW(GL.GLExplode(V,[EVor],1.,1.,1.,99,1));

GL.VIEW([ GL.GLGrid(V,EV, GL.COLORS[1],0.5), GL.GLFrame ]);


immutable Point{T<:FloatingPoint}
    x::T
    y::T
end

immutable Circle{T<:FloatingPoint}
    c::Point{T}
    r::T
end
Circle{T<:FloatingPoint}(a::Point{T}) = Circle(a, zero(T))

using AffineTransforms

function circlepoints{T<:FloatingPoint}(a::Point{T}, b::Point{T}, r::T)
    cp = Circle{T}[]
    r >= 0 || return (cp, "No Solution, Negative Radius")
    if a == b
        if abs(r) < 2eps(zero(T))
            return (push!(cp, Circle(a)), "Point Solution, Zero Radius")
        else
            return (cp, "Infinite Solutions, Indefinite Center")
        end
    end
    ca = Complex(a.x, a.y)
    cb = Complex(b.x, b.y)
    d = (ca + cb)/2
    tfd = tformtranslate([real(d), imag(d)])
    tfr = tformrotate(angle(cb-ca))
    tfm = tfd*tfr
    u = abs(cb-ca)/2
    r-u > -5eps(r) || return(cp, "No Solution, Radius Too Small")
    if r-u < 5eps(r)
        push!(cp, Circle(apply(Point, tfm*[0.0, 0.0]), r))
        return return (cp, "Single Solution, Degenerate Centers")
    end
    v = sqrt(r^2 - u^2)
    for w in [v, -v]
        push!(cp, Circle(apply(Point, tfm*[0.0, w]), r))
    end
    return (cp, "Two Solutions")
end


using Measurements
using Printf
include("/home/jefter66/Projects/Lablib/src/Stats.jl")
include("/home/jefter66/Projects/Lablib/src/LeastSquares.jl")
γ(h1,h3) = return (h1) / (h1 - h3);

Δh = 0.1;

h1 = [18.1 ± Δh, 16.0 ± Δh, 17.6 ± Δh];
h3 = [5.7, 4,5.1];

ɣs = [γ(h1[1],h3[1]), γ(h1[2],h3[2]), γ(h1[3],h3[3])];
deviation(ɣs);

A(d) = return π * (d^2 / 4);
# ΔA = 0.00001;

P0 = 691 * 133.322 ± 0.01; # ΔP = 0.01;
g = 9.81;

d = 0.016 ± 0.001; # Δd = 0.001;
m = 0.0167 ± 0.0001; # Δm = 0.0001;
V = 0.0104  ± 0.0001; # ΔV = 0.0001;
P = (m * g)/A(d) + P0;
T = 1.2  ± 0.01; # ΔT = 0.01;

𝛄 = 4 * π^2 * ((m * V)/(P * A(d)^2 * T^2));
#x, Δx = m * V, Δm * V + ΔV * m;
#y, Δy = P*A(d)^2 * T^2, 2*T*ΔT*(P*A(d)^2)  +  (ΔP * A(d)^2 + 2 * A(d) * ΔA * P);
#Δ𝛄 = 4 * π^2 * (Δx * y + Δy * x) / (y^2);

@printf("𝛄 = %.6f ± %.6f", 𝛄.val, 𝛄.err);

x = [24.7,1.0,-196.0,97.0];
y = [69.2,62.8,17.7,82.3];

leastSquares(x,y);


using CompEcon
using Dierckx
using Statistics

linspace(a,b,n) = range(a, stop=b, length=n)

# Create functions to test against
@. f1(x, y) = sin(x) * cos(y)
@. f2(x, y) = exp(x) * cos(y)
@. f3(x, y) = 1.5 + log(x * y)
fs = [f1, f2, f3]
nf = length(fs)

x = collect(linspace(.5, 2, 35))
y = collect(linspace(.2, 4, 35))

function compare_D_CE_levels(f::Function, x::Array{Float64, 1}, y::Array{Float64, 1}, order::Int)

    # Create Dierckx Spline (and data)
    n = length(x)
    z = f(x, y')
    dspl = Spline2D(x, y, z; s=0.0, kx=order, ky=order)

    # Create CompEcon Spline (and data)
    cebasis_x = Basis(SplineParams(x, 0, order))
    cebasis_y = Basis(SplineParams(y, 0, order))
    cebasis = Basis(cebasis_x, cebasis_y)
    xx = nodes(cebasis)[1]
    yy = f(xx[:, 1], xx[:, 2])
    c, bs = funfitxy(cebasis, xx, yy)

    # Evaluate Splines on finer grid
    xfine = collect(linspace(x[1], x[end], 2*n + 1))
    yfine = collect(linspace(y[1], y[end], 2*n + 1))
    zfine = f(xfine, yfine')

    deval = Dierckx.evalgrid(dspl, xfine, yfine)
    derr = abs.(zfine .- deval)

    ceeval = CompEcon.funeval(c, cebasis, gridmake(xfine, yfine))
    ceerr = abs.(zfine .- reshape(ceeval, 2*n+1, 2*n+1))

    return derr, ceerr
end

for i=1:nf

    curr_f = fs[i]

    derr1_levels, ceerr1_levels = compare_D_CE_levels(curr_f, x, y, 1)
    derr2_levels, ceerr2_levels = compare_D_CE_levels(curr_f, x, y, 2)
    derr3_levels, ceerr3_levels = compare_D_CE_levels(curr_f, x, y, 3)

    println("Linear Approximation to Function")
    println("Dierckx Level Errors \n\t Max: $(maximum(derr1_levels)) \n\t Min: $(minimum(derr1_levels)) \n\t Mean: $(mean(derr1_levels))")
    println("CompEcon Level Errors \n\t Max: $(maximum(ceerr1_levels)) \n\t Min: $(minimum(ceerr1_levels)) \n\t Mean: $(mean(ceerr1_levels))")

    println("Quadratic Approximation to Function")
    println("Dierckx Level Errors \n\t Max: $(maximum(derr2_levels)) \n\t Min: $(minimum(derr2_levels)) \n\t Mean: $(mean(derr2_levels))")
    println("CompEcon Level Errors \n\t Max: $(maximum(ceerr2_levels)) \n\t Min: $(minimum(ceerr2_levels)) \n\t Mean: $(mean(ceerr2_levels))")

    println("Cubic Approximation to Function")
    println("Dierckx Level Errors \n\t Max: $(maximum(derr3_levels)) \n\t Min: $(minimum(derr3_levels)) \n\t Mean: $(mean(derr3_levels))")
    println("CompEcon Level Errors \n\t Max: $(maximum(ceerr3_levels)) \n\t Min: $(minimum(ceerr3_levels)) \n\t Mean: $(mean(ceerr3_levels))")

end


using .FFTW
laplacianunwrap(ϕ) = laplacianunwrap!(copy(ϕ))
function laplacianunwrap!(ϕ::AbstractArray)
    FFTW.set_num_threads(Threads.nthreads())
    ϕ .+= 2π .* k(ϕ)
end

# Schofield and Zhu 2003, https://doi.org/10.1364/OL.28.001194
k(ϕw) = 1 / 2π .* ∇⁻²(∇²_nw(ϕw) - ∇²(ϕw))  # (1)

∇²(x) = -(2π)^ndims(x) / length(x) .* idct(pqterm(size(x)) .* dct(x))  # (2)

∇⁻²(x) = -length(x) / (2π)^ndims(x) .* idct(dct(x) ./ pqterm(size(x)))  # (3)

∇²_nw(ϕw) = cos.(ϕw) .* ∇²(sin.(ϕw)) .- sin.(ϕw) .* ∇²(cos.(ϕw))  # (in text)

pqterm(sz::NTuple{1}) = (1:sz[1]).^2  # 1D case
pqterm(sz::NTuple{2}) = [p^2 + q^2 for p in 1:sz[1], q in 1:sz[2]]  # 2D case
pqterm(sz::NTuple{3}) = [p^2 + q^2 + t^2 for p in 1:sz[1], q in 1:sz[2], t in 1:sz[3]]  # 3D case
pqterm(sz::NTuple{4}) = [p^2 + q^2 + t^2 + r^2 for p in 1:sz[1], q in 1:sz[2], t in 1:sz[3], r in 1:sz[4]]  # 4D case

"""
    laplacianunwrap(ϕ::AbstractArray)

Performs laplacian unwrapping on the input phase. (1D - 4D)
The phase has to be scaled to radians.
The implementation is close to the original publication: Schofield and Zhu 2003, https://doi.org/10.1364/OL.28.001194.
It is not the fastest implementation of laplacian unwrapping (doesn't use discrete laplacian).
"""
laplacianunwrap, laplacianunwrap!


using ChebyshevQuantum
using Test

tol_val = 1e-11

@testset "test C interp" begin

    function f(x)
        return exp(-0.2*sin(3*x))
    end

    c = ChebyshevQuantum.Chebyshev.make_Cheb(f , a=2, b = 3)

    @test isapprox(c(2.3), f(2.3), atol=tol_val)

    function f2(x)
        if x < 0.1
            return exp(-0.2*sin(3*x))
        else
            return 1+x^2
        end
    end

    c = ChebyshevQuantum.Chebyshev.make_Cheb(f2 , a=-1, b = 1)

    x=0.0
    @test isapprox(c(x), f2(x), atol=tol_val)
    x=0.2
    @test isapprox(c(x), f2(x), atol=tol_val)

end


function generate_window_functions_expr()
    # x_exprW = :(clamp.(- border_in .+ abs.(scale .* (x .- offset))./(border_out-border_in),0,1))
    x_exprW = :(clamp.(1 .-(abs.(scale .* (x .- offset)).-border_in)./(border_out .- border_in),0,1))
    x_exprW2 = :(clamp.((abs.(scale .* (x .- offset)).-border_in)./(border_out .- border_in),0,Inf))  # no outer border and starting from 0 at the inner border
    x_exprRW = :(clamp.(1 .-(sqrt.(sum((scale .* (x .- offset)).^2)).-border_in)./(border_out .- border_in),0,1))
    # Attention: The function below returns the square distance to avoid sqrt followed by squaring in the gaussian
    x_exprRW2 = :(clamp.((sum((scale .* (x .- offset)).^2).-border_in)./(border_out .- border_in),0,Inf))   # no outer border and starting from 0 at the inner border

    functions = [  # see https://en.wikipedia.org/wiki/Window_function
        (:(window_linear),  :(x -> T(prod(($x_exprW))))),
        (:(window_edge),  :(x -> T(prod(($x_exprW).>0.5)))),
        (:(window_hanning),  :(x -> T(prod(sinpi.(0.5 .* ($x_exprW)).^2)))),
        (:(window_half_cos),  :(x -> T(prod(sinpi.(0.5 .* ($x_exprW)))))),
        (:(window_hamming),  :(x -> T(prod(0.54 .-0.46.*cospi.(($x_exprW)))))),
        (:(window_blackman_harris),  :(x -> T(prod(0.35875 .- 0.48829.*cospi.($x_exprW).+0.14128.*cospi.(2 .*$x_exprW).-0.01168.*cospi.(3 .*$x_exprW))))),
        (:(window_gaussian),  :(x -> T(prod(exp.(-2 .* abs2.($x_exprW2)))))),
        (:(window_radial_linear),  :(x -> T($x_exprRW))),
        (:(window_radial_edge),  :(x -> T(($x_exprRW).>0.5))),
        (:(window_radial_hanning),  :(x -> T(sinpi.(0.5 .* ($x_exprRW)).^2))),
        (:(window_radial_hamming),  :(x -> T((0.54 .-0.46.*cospi.(($x_exprRW)))))),
        (:(window_radial_blackman_harris),  :(x -> T((0.35875 .-0.48829.*cospi.($x_exprRW).+0.14128.*cospi.(2 .*$x_exprRW).-0.01168.*cospi.(3 .*$x_exprRW))))),
        (:(window_radial_gaussian),  :(x -> T(exp.(-2 .*($x_exprRW2))))),
    ]
    x_expr8 = :(exp(.- sum(abs2.(x .- offset).*scale))) # scale is 1/(2 sigma^2)
    return functions
end

# we automatically generate the functions for different windows like hanning
for F in generate_window_functions_expr()
    # default functions with certain offset and scaling behavior
    @eval function $(F[1])(::Type{T}, size::NTuple{N, Int};
                           offset=CtrFT,
                           scale=ScaFTEdge, border_in=0.8, border_out=1.0, dims=ntuple(+, N)) where{N, T}
        scale_init = get_scale(size, scale)
        scale = ntuple(i -> i ∈ dims ? scale_init[i] : zero(scale_init[1]), N)
        offset = get_offset(size, offset)
        IndexFunArray(T, $(F[2]), size)
    end

    # change order of offset and scale
    @eval function $(F[1])(size::NTuple{N, Int};
                           offset=CtrFT,
                           scale=ScaFTEdge, border_in=0.8, border_out=1.0, dims=ntuple(+, N)) where{N}
        T = $DEFAULT_T
        $(F[1])(T, size, scale=scale, offset=offset, border_in=border_in, border_out=border_out, dims=dims)
    end

    # convenient wrapper to provide an array as input
    @eval function $(F[1])(arr::AbstractArray{T, N};
                           offset=CtrFT,
                           scale=ScaFTEdge, border_in=0.8, border_out=1.0, dims=ntuple(+, N)) where{N, T}
        $(F[1])(default_type(T, $DEFAULT_T), size(arr), scale=scale, offset=offset, border_in=border_in, border_out=border_out, dims=dims)
    end

    @eval export $(F[1])
end


using DifferentialEquations
using SimpleDiffEq
using Random
using DataFrames
using StatsPlots
using BenchmarkTools

function infection_rate(u,p,t)
    (S,I,R) = u
    (β,c,γ) = p
    N = S+I+R
    β*c*I/N*S
end
function infection!(integrator)
  integrator.u[1] -= 1
  integrator.u[2] += 1
end
infection_jump = ConstantRateJump(infection_rate,infection!);

function recovery_rate(u,p,t)
    (S,I,R) = u
    (β,c,γ) = p
    γ*I
end
function recovery!(integrator)
  integrator.u[2] -= 1
  integrator.u[3] += 1
end
recovery_jump = ConstantRateJump(recovery_rate,recovery!);

tmax = 40.0
tspan = (0.0,tmax);

δt = 0.1
t = 0:δt:tmax;

u0 = [990,10,0]; # S,I,R

p = [0.05,10.0,0.25]; # β,c,γ

Random.seed!(1234);

prob = DiscreteProblem(u0,tspan,p)

prob_jump = JumpProblem(prob,Direct(),infection_jump,recovery_jump)

sol_jump = solve(prob_jump,SSAStepper());

out_jump = sol_jump(t);

df_jump = DataFrame(out_jump')
df_jump[!,:t] = out_jump.t;

@df df_jump plot(:t,
    [:x1 :x2 :x3],
    label=["S" "I" "R"],
    xlabel="Time",
    ylabel="Number")

@benchmark solve(prob_jump,FunctionMap())

include(joinpath(@__DIR__,"tutorials","appendix.jl"))
appendix()


# All square roots are periodic when written as continued fractions and can be
# written in the form:
#
# √N = a0 + 1 / (a1 + 1 / (a2 + 1 / (a3 + ...)))
#
# For example, let us consider √23:
# √23 = 4 + √23 − 4 = 4 + 1 / (1 / (√23 − 4)) = 4 + 1 / (1 + (√23 − 3) / 7)
#
# If we continue we would get the following expansion:
# √23 = 4 + 1 / (1 + 1 / (3 + 1 / (1 + 1 / (8 + ...))))
#
# The process can be summarised as follows:
#
# a0 = 4, 1 / (23√ - 4) = (√23 + 4) / 7 = 1 + (√23 - 3) / 7
# a1 = 1, 7 / (23√ - 3) = 7 * (23√+3)14 = 3 + (√23 - 3) / 2
# a2 = 3, 2 / (23√ - 3) = 2 * (23√+3)14 = 1 + (√23 - 4) / 7
# a3 = 1, 7 / (23√ - 4) = 7 * (23√+4)7 = 8 + √23 - 4
# a4 = 8, 1 / (23√ - 4) = (√23 + 4) / 7 = 1 + (√23 - 3) / 7
# a5 = 1, 7 / (23√ - 3) = 7 * (√23 + 3) / 14 = 3 + (√23 - 3) / 2
# a6 = 3, 2 / (23√ - 3) = 2 * (√23 + 3) / 14 = 1 + (√23 - 4) / 7
# a7 = 1, 7 / (23√ - 4) = 7 * (√23 + 4) / 7 = 8 + √23 - 4
#
# It can be seen that the sequence is repeating. For conciseness, we use the
# notation √23 = [4;(1,3,1,8)], to indicate that the block (1,3,1,8) repeats
# indefinitely.
#
# The first ten continued fraction representations of (irrational) square roots
# are:
#
# √2 = [1;(2)], period=1
# √3 = [1;(1,2)], period=2
# √5 = [2;(4)], period=1
# √6 = [2;(2,4)], period=2
# √7 = [2;(1,1,1,4)], period=4
# √8 = [2;(1,4)], period=2
# √10 = [3;(6)], period=1
# √11 = [3;(3,6)], period=2
# √12 = [3;(2,6)], period=2
# √13 = [3;(1,1,1,1,6)], period=5
#
# Exactly four continued fractions, for N ≤ 13, have an odd period.
#
# How many continued fractions for N ≤ 10000 have an odd period?

using ProjectEulerSolutions

# A continued fraction period ends when the a,m,d tuple matches a previously
# seen one.  Just maintain the list of tuples until one matches, measuring the
# length of the continued fraction period to find the odd numbered ones.
function p064solution(n_max::Integer=10)::Integer

    odd_count = 0
    for n in 2:n_max
        a0 = floor(sqrt(n))
        if sqrt(n) == a0
            continue
        end
        a0 = Integer(a0)
        a, m, d = a0, 0, 1
        s = Array{Tuple{Integer, Integer, Integer}, 1}()
        while true
            m = d * a - m
            d = div(n - m * m, d)
            a = div(a0 + m, d)
            key = (a, m, d)
            if length(s) > 0 && key == s[1]
                break
            else
                push!(s, (a, m, d))
            end
        end
        if length(s) % 2 == 1
            odd_count += 1
        end
    end
    return odd_count
end

p064 = Problems.Problem(p064solution)

Problems.benchmark(p064, 10_000)


# Example of logistic map from Keith Briggs's 2006 paper

using ExactReals

"Iterate the function `f` for `n` steps, starting at `x0`."
function iterates(f, n, x0)

    x = x0

    for i in 1:n
        x = f(x)
    end

    return x
end

"Logistic map"
logistic(a, x) = a * x * (1 - x)
logistic(a) = x -> logistic(a, x)

a = rationalize(3.999)
x0 = rationalize(0.9)
n = 54

y = iterates(logistic(ExactReal(a)), n, ExactReal(x0))

@show y > ExactReal(1//2)

# Floating point:

z = iterates(logistic(float(a)), n, float(x0))

@show z > 0.5

nothing

#
# using IntervalArithmetic
#
# x = interval(big"0.9")
#
# iterates(logistic(@biginterval(3.999)), 53, @biginterval(0.9)) |> showfull
#
# iterates(logistic(I"3.999"), 53, I"0.9")
#
# showfull(I"3.999")
#
# # example where double precision is on wrong side of 1/2:
# iterates(logistic(3.999), 54, 0.9)
#
# iterates(logistic(@biginterval(3.999)), 54, @biginterval(0.9)) |> showfull
#
#
# @time y = iterates(logistic(ExactReal(3999//1000)), 54, ExactReal(9//10))
#
# y > ExactReal(1//2)


function Sch_solve_Emin_cg( LF::LF3dGrid, Vpot, Ncol::Int64;
                            v0 = nothing,
                            α_t=3e-5, Niter=1000,
                            verbose=false )
  #
  Npoints = LF.Nx * LF.Ny * LF.Nz
  ΔV = LF.LFx.h * LF.LFy.h * LF.LFz.h
  #
  # Setup trial solution
  #
  if v0 == nothing
    # random trial solution
    srand(1234)
    v = rand( Npoints, Ncol )
    v = orthonormalize( LF, v )
  else
    v = copy(v0)
    # v0 is assumed to be ortonormalized
  end
  #
  Energies = calc_Energies( LF, Vpot, v )
  #
  Etot_old = Energies.Total
  #
  g     = zeros(Float64,Npoints,Ncol)
  g_old = zeros(Float64,Npoints,Ncol)
  d     = zeros(Float64,Npoints,Ncol)
  d_old = zeros(Float64,Npoints,Ncol)
  #
  β = 0
  α = 0
  for iter = 1:Niter
    g = calc_grad( LF, Vpot, v )
    #
    if iter != 1
      β = trace( g' * g)/trace( g_old'*g_old )
      #β = trace( (g-g_old)'*g ) / trace( g_old'*g_old )
      #β = trace( (g-g_old)'*g )/ trace( (g-g_old)'*d )
    end
    d = -g + β*d_old
    #
    # compute gradient at trial step
    #
    v2 = orthonormalize(LF, v + α_t*d)
    g_t = calc_grad( LF, Vpot, v2 )
    #
    # compute estimate of best step and update current trial vectors
    #
    denum = trace( (g - g_t)' * d )
    if denum != 0.0
      α = abs( α_t * trace( g'*d )/denum )
    else
      α = 0.0
    end
    v = v + α*d
    #
    v = orthonormalize( LF, v )
    Energies = calc_Energies( LF, Vpot, v )
    Etot = Energies.Total
    #
    if verbose
      @printf("%8d %16.9e %16.9e\n", iter, Etot, abs(Etot-Etot_old))
    end
    if abs(Etot-Etot_old) < 1.e-7
      if verbose
        @printf("Emin CG converges in %8d iterations\n", iter)
      end
      break
    end
    Etot_old = Etot
    #
    g_old = g[:,:]
    d_old = d[:,:]
  end
  return Energies, v
end


"""
    P(n, τ)

Returns the value of the Legendre polynomial of degree `n` at `τ`.
(Recursive function used in `../test`).
"""
function P(n::Integer, τ::T) where T <: AbstractFloat
    if n == 0
        return one(T)
    elseif n == 1
        return τ
    else
        return ( (2n-1) * τ * P(n-1, τ) - (n-1) * P(n-2, τ) ) / n
    end
end

"""
    legendre_polys!(P, τ)

If `τ` is a scalar, then `P` is a vector with `P[n+1]` equal to the value
of the Legendre polynomial of degree `n` at `τ`.

If `τ` is a vector, then `P` is a matrix with `P[n+1,j]` equal to the value
of the Legendre polynomial of degree `n` at `τ[j]`.
"""
function legendre_polys!(P::AbstractVector{T}, τ::T) where T <: AbstractFloat
    I = length(P)
    if I ≥ 1
        P[1] = one(T)
    end
    if I ≥ 2
        P[2] = τ
    end
    for j = 1:I, i = 1:I-2
        P[i+2] = ((2i+1) * τ * P[i+1] - i * P[i] ) / (i+1)
    end
end

function legendre_polys!(P::AbstractMatrix{T}, τ::AbstractVector{T}
                        ) where T <: AbstractFloat
    I = size(P, 1)
    J = size(P, 2)
    @argcheck length(τ) == J
    if I ≥ 1
        for j = 1:J
            P[1,j] = one(T)
        end
    end
    if I ≥ 2
        for j = 1:J
            P[2,j] = τ[j]
        end
    end
    for j = 1:J, i = 1:I-2
        P[i+2,j] = ((2i+1) * τ[j] * P[i+1,j] - i * P[i,j] ) / (i+1)
    end
end

"""
    dP(n, τ)

Returns the value of the derivative of the Legendre polynomial of degree `n`
at `τ`.  (Recursive function used in `../test`).
"""
function dP(n::Integer, τ::T) where T <: AbstractFloat
    if n == 0
        return zero(T)
    elseif n == 1
        return one(T)
    else
        return ( (2n-1) * τ * dP(n-1, τ) - n * dP(n-2, τ) ) / ( n - 1 )
    end
end

function deriv_legendre_polys!(dP::AbstractVector{T}, τ::T
                              ) where T <: AbstractFloat
    I = length(dP)
    if I ≥ 1
        dP[1] = zero(T)
    end
    if I ≥ 2
        dP[2] = one(T)
    end
    for i = 1:I-2
        dP[i+2] = ( (2i+1) * τ * dP[i+1] - (i+1) * dP[i] ) / i
    end
end

function deriv_legendre_polys!(dP::AbstractMatrix{T}, τ::AbstractVector{T}
                               ) where T <: AbstractFloat
    I = size(dP, 1)
    J = size(dP, 2)
    @argcheck length(τ) == J
    if I ≥ 1
        for j = 1:J
            dP[1,j] = zero(T)
        end
    end
    if I ≥ 2
        for j = 1:J
            dP[2,j] = one(T)
        end
    end
    for j = 1:J, i = 1:I-2
        dP[i+2,j] = ((2i+1) * τ[j] * dP[i+1,j] - (i+1) * dP[i,j] ) / i
    end
end

"""
    pcwise_t, pcwise_U = evaluate_pcwise_poly(U, t, ppI, store)
"""
function evaluate_pcwise_poly!(U::Vector{Vector{T}}, t::OffsetVector{T},
                               ppI::Integer, store::Store{T}
                              ) where { T <: AbstractFloat }
    τ = range(-one(T), stop=one(T), length=ppI)
    pcwise_t, pcwise_U = evaluate_pcwise_poly!(U, t, τ, store)
    return pcwise_t, pcwise_U
end

function evaluate_pcwise_poly!(U::Vector{Vector{T}}, t::OffsetVector{T},
                               τ::AbstractVector{T}, store::Store{T}
                              ) where T <: AbstractFloat
    N = length(U)
    ppI = length(τ)
    @argcheck ppI ≤ store.ppImax
    pcwise_t = Array{T}(undef, ppI, N)
    pcwise_U = similar(pcwise_t)
    evaluate_pcwise_poly!(pcwise_t, pcwise_U, U, t, τ, store)
    return pcwise_t, pcwise_U
end

function evaluate_pcwise_poly!(pcwise_t::Matrix{T}, pcwise_U::Matrix{T},
                               U::Vector{Vector{T}}, t::OffsetVector{T},
                               τ::AbstractVector{T}, store::Store{T}
                              ) where { T <: AbstractFloat }
    N = length(U)
    pts_per_interval = length(τ)
    @argcheck size(pcwise_t) == (pts_per_interval, N)
    @argcheck size(pcwise_U) == (pts_per_interval, N)
    @argcheck length(t) == N+1
    rmax = store.rmax
    Ψ = view(store.Ψ, 1:rmax, 1:pts_per_interval)
    legendre_polys!(Ψ, τ)
    for n = 1:N
        rn = length(U[n])
        for m = 1:pts_per_interval
            pcwise_t[m,n] = ( (1-τ[m])*t[n-1] + (1+τ[m])*t[n] ) / 2
            s = zero(T)
            for j = 1:rn
                s += U[n][j] * Ψ[j,m]
            end
            pcwise_U[m,n] = s
        end
    end
end

function Fourier_Legendre_coefs(u::Function, r::Integer, t::OffsetArray{T},
                                store::Store{T}) where T <: AbstractFloat
    N = length(t) - 1
    rmax = store.rmax
    a = Vector{Vector{T}}(undef, N)
    M = store.Mmax
    τ, wτ = rule(store.legendre[M])
    Ψ = view(store.Ψ, 1:r, 1:M)
    legendre_polys!(Ψ, τ)
    for n = 1:N
        a[n] = Vector{T}(undef, r)
        for j = 1:r
            s = zero(T)
            for m = 1:M
                tnm = ( (1-τ[m])*t[n-1] + (1+τ[m])*t[n] ) / 2
                s += wτ[m] * u(tnm) * Ψ[j,m]
            end
            a[n][j] = (2j-1) * s / 2
        end
    end
    return a
end

function reconstruction_pts(In::Tuple{T,T}, r::Integer) where T <: AbstractFloat
    tnm1, tn = In
    tstar = reconstruction_pts(T, r)
    tstar[0] = tnm1
    for m = 1:r-1
        tstar[m] = ( ( 1 - tstar[m] ) * tnm1 + ( tstar[m] + 1 ) * tn ) / 2
    end
    tstar[r] = tn
    return tstar
end

function reconstruction_pts(::Type{T}, r::Integer) where T <: AbstractFloat
    right = GaussQuadrature.right
    τ, wτ = GaussQuadrature.legendre(T, r, right)
    return OffsetArray([ -one(0); τ ], 0:r)
end

"""
    Uhat = reconstruction(U, u0, store)
"""
function reconstruction(U::Vector{Vector{T}}, u0::T,
                        store::Store{T}) where T <: AbstractFloat
    Uhat = reconstruction(U, u0, store.rmax)
    return Uhat
end

"""
    Uhat = reconstruction(U, u0, rmax)
"""
function reconstruction(U::Vector{Vector{T}}, u0::T,
	                rmax::Integer) where T <: AbstractFloat
    N = length(U)
    Uhat = Vector{Vector{T}}(undef, N)
    pow = OffsetArray{T}(undef, 0:rmax)
    pow[0] = one(T)
    for n = 1:rmax
        pow[n] = -pow[n-1]
    end
    r1 = length(U[1])
    @argcheck r1+1 ≤ rmax
    U_left = zero(T)   # U_left  = U(t[0] + 0) = value at the left endpoint.
    U_right = zero(T)  # U_right = U(t[1] - 0) = value at the right endpoint.
    for j = 1:r1
        U_left += pow[j-1] * U[1][j]
        U_right += U[1][j]
    end
    jumpU0 = U_left - u0
    Uhat[1] = Vector{T}(undef, r1+1)
    Uhat[1][1:r1] .= U[1][1:r1]
    Uhat[1][r1] += pow[r1] * jumpU0 / 2
    Uhat[1][r1+1] = - pow[r1] * jumpU0 / 2
    for n = 2:N
        rn = length(U[n])
        @argcheck rn + 1 ≤ rmax
        U_left = zero(T) # U_left = U(t[n-1] + 0 )
        for j = 1:rn
            U_left += pow[j-1] * U[n][j]
        end
        jumpUnm1 = U_left - U_right
        Uhat[n] = Vector{T}(undef, rn+1)
        Uhat[n][1:rn] .= U[n][1:rn]
        Uhat[n][rn] += pow[rn] * jumpUnm1 / 2
        Uhat[n][rn+1]  = -pow[rn] * jumpUnm1 / 2
        U_right = zero(T) # U_right = U(t[n] - 0 )
        for j = 1:rn
            U_right += U[n][j]
        end
    end
    return Uhat
end

function max_order(U::Vector{Vector{T}}) where T <: AbstractFloat
    N = length(U)
    r = length(U[1])
    for n = 2:N
        rn = length(U[n])
        r = max(r, rn)
    end
    return r
end

function jumps(U::Vector{Vector{T}}, t::OffsetVector{T},
	u0::T) where T <: AbstractFloat
    N = length(t) - 1
    r = max_order(U)
    JU = OffsetVector{T}(undef, 0:N-1)
    U_left = u0
    U_right = zero(T)
    pow = ones(T, r)
    pow[2:2:r] .= -one(T)
    for n = 1:N
	U_left = zero(T)
	for j = 1:r
	    U_left += pow[j] * U[n][j]
	end
	JU[n-1] = U_left
    end
    JU[0] -= u0
    for n = 2:N
	U_right = sum(U[n-1])
	JU[n-1] -= U_right
    end
    return JU
end

function dG_error_estimator(JU::OffsetVector{T}, t::OffsetVector{T},
        r::Integer, τ::AbstractVector{T}) where T <: AbstractFloat
    N = length(t) - 1
    pts_per_interval = length(τ)
    pcwise_t = Matrix{T}(undef, pts_per_interval, N)
    pcwise_approx_err = similar(pcwise_t)
    Ψ = Matrix{T}(undef, r+1, pts_per_interval)
    legendre_polys!(Ψ, τ)
    if r % 2 == 0
        sign = 1
    else
        sign = -1
    end
    for n = 1:N
        for m = 1:pts_per_interval
            pcwise_t[m,n] = ( (1-τ[m])*t[n-1] + (1+τ[m])*t[n] ) / 2
            pcwise_approx_err[m,n] = (sign/2) * JU[n-1] * ( Ψ[r+1,m] - Ψ[r,m] )
        end
    end
    return pcwise_t, pcwise_approx_err
end

function dG_error_estimator(JU::OffsetVector{T}, t::OffsetVector{T},
        r::Integer, pts_per_interval::Integer) where T <: AbstractFloat
    τ = range(-1.0, 1.0, length=pts_per_interval)
    return dG_error_estimator(JU, t, r, τ)
end


function step_size(x0, zero_tol = 1e-15)
    #=
    if abs(x0) <= zero_tol
        x0 = copysign(zero_tol, x0)
    end
    =#
    h = sqrt( eps(x0) )
    return h
end

function step_size_cbrt(x0, zero_tol = 1e-15)
    #=
    if abs(x0) <= zero_tol
        x0 = copysign(zero_tol, x0)
    end
    =#
    h = cbrt( eps(x0) )
    return h
end

function numdiff_fw(f, x0)
    # f: R -> R
    #=
    h = step_size(x0)
    x1 = x0 + h
    f0 = f(x0)
    f1 = f(x1)
    g = (f1 - f0) / h
    return g
    =#
    h = step_size(x0)
    numdiff_fw(f, x0, h)
end

function numdiff_fw(f, x0, h)
    # f: R -> R
    f0 = f(x0)
    numdiff_fw(f, x0, h, f0)
end

function numdiff_fw(f, x0, h, f0)
    # f: R -> R
    x1 = x0 + h
    f1 = f(x1)
    g = (f1 - f0) / h
end

function numdiff_bw(f, x0)
    # f: R -> R
     #=
    h = step_size(x0)
    x1 = x0 - h
    f0 = f(x0)
    f1 = f(x1)
    g = (f0 - f1) / h
    =#
    h = step_size(x0)
    numdiff_bw(f, x0, h)
end

function numdiff_bw(f, x0, h)
    f0 = f(x0)
    numdiff_bw(f, x0, h, f0)
end

function numdiff_bw(f, x0, h, f0)
    x1 = x0 - h
    f1 = f(x1)
    g = (f0 - f1) / h
end

function numdiff_central(f, x0)
    # f: R -> R
    h = step_size_cbrt(x0)
    xm1 = x0 - h
    xp1 = x0 + h
    fm1 = f(xm1)
    fp1 = f(xp1)
    g = (fp1 - fm1) / (2*h)
    return g
end

function gradient_fw(f, x0)
    # f: R^n -> R
    # g_i = df / dx_i
    h = step_size.(x0)
    return gradient_fw(f, x0, h)
end

function gradient_fw(f, x0, h)
    # f: R^n -> R
    # g_i = df / dx_i
    f0 = f(x0)
    return gradient_fw(f, x0, h, f0)
end

function gradient_fw(f, x0, h, f0)
    # f: R^n -> R
    # g_i = df / dx_i
    g = similar(x0)
    x1 = copy(x0)
    for kk = 1:length(x0)
        x1[kk] += h[kk]
        f1 = f(x1)
        g[kk] = (f1 - f0) / h[kk]
        x1[kk] -= h[kk]
    end
    return g
end

function directional_diff_fw(f, x0, d)
    h = step_size( norm(d) )
    return directional_diff_fw(f, x0, d, h)
end

function directional_diff_fw(f, x0, d, h)
    f0 = f(x0)
    return directional_diff_fw(f, x0, d, h, f0)
end

function directional_diff_fw(f, x0, d, h, f0)

    g = zero(eltype(x0))

    x1 = x0 + d * h
    f1 = f(x1)

    g = (f1 - f0) / h

    return g
end

function jacobian_fw(f, x0)
    # f: R^n -> R^m
    # g_ij = df_i / dx_j

    h = step_size.(x0)  # n-by-1
    return jacobian_fw(f, x0, h)
end

function jacobian_fw(f, x0, h)
    # f: R^n -> R^m
    # g_ij = df_i / dx_j

    f0 = f(x0)                  # m-by-1
    return jacobian_fw(f, x0, h, f0)   # m-by-n
end

function jacobian_fw(f, x0, h, f0)
    # f: R^n -> R^m
    # g_ij = df_i / dx_j

    jac = zeros(eltype(f0), length(f0), length(x0)) # m-by-n
    x1 = copy(x0)
    for kk = 1:length(x0)
        x1[kk] += h[kk]
        f1 = f(x1)
        jac[:, kk] = (f1 .- f0) ./ h[kk]
        x1[kk] -= h[kk]
    end
    return jac
end

function jacobian_fw!(jac, f, x0)
    # f: R^n -> R^m
    # g_ij = df_i / dx_j

    h = step_size.(x0)  # n-by-1
    return jacobian_fw!(jac, f, x0, h)
end

function jacobian_fw!(jac, f, x0, h)
    # f: R^n -> R^m
    # g_ij = df_i / dx_j

    f0 = f(x0)                  # m-by-1
    return jacobian_fw!(jac, f, x0, h, f0)   # m-by-n
end

function jacobian_fw!(jac, f, x0, h, f0)
    # f: R^n -> R^m
    # g_ij = df_i / dx_j

    f1 = similar(f0)
    x1 = copy(x0)
    for kk = 1:length(x0)
        x1[kk]     += h[kk]
        f1         .= f(x1)
        jac[:, kk] .= (f1 .- f0) ./ h[kk]
        x1[kk]     -= h[kk]
    end
    return jac
end


# This file includes the MonteCarlo simulation of PSK modulation scheme and compares
# the numerical results with the theoretical results.

using DigitalCommunications
using Plots

# Simulation parameters
k = 3
M = 2^k
nsymbols = Int(1e6)
nbits = k * nsymbols
ebno = collect(0 : 10)
esno = ebno .+ 10 * log10(k)

# Communcation system components
gen = SymbolGenerator(nsymbols, M)
modulator = Modulator(PSK(M))
channel = AWGNChannel()
detector = Detector(modulator(1:M))

# Monte Carlo simulation
message = gen.symbols
symerr = zeros(length(esno))
for i in 1 : length(symerr)
    channel.esno = esno[i]  # Update channel snr
    mbar = message |> modulator |> channel |> detector  # Extracted message signal
    symerr[i] = sum(mbar .!= message) / length(message)  # Symbol error rate
end

# Plots
plt = plot(title="$M-PSK", xlabel="ebno [dB]", ylabel="Pe")
plot!(ebno, berpsk.(esno, M), marker=:circle, yscale=:log10, label="theoretical")
plot!(ebno, symerr, marker=:circle, yscale=:log10, label="montecarlo")


using MultivariateStats
using LinearAlgebra, StatsBase, SparseArrays
using Test
import Statistics: mean, cov
import Random

@testset "Whitening" begin

    Random.seed!(34568)

    ## data

    d = 3
    n = 5

    X = rand(d, n)
    mval = vec(mean(X, dims=2))
    C = cov(X, dims=2)
    C0 = copy(C)

    emax = maximum(eigvals(C))
    rc = 0.01
    Cr = C + (emax * rc) * Matrix(I, d, d)

    ## cov_whitening

    cf = cholesky(Hermitian(C, :U))
    W = cov_whitening(cf)
    @test istriu(W)
    @test W'C * W ≈ Matrix(I, d, d)

    cf = cholesky(Hermitian(C, :L))
    W = cov_whitening(cf)
    @test istriu(W)
    @test W'C * W ≈ Matrix(I, d, d)

    W = cov_whitening(C)
    @test C == C0
    @test istriu(W)
    @test W'C * W ≈ Matrix(I, d, d)

    W = cov_whitening(C, 0.0)
    @test C == C0
    @test istriu(W)
    @test W'C * W ≈ Matrix(I, d, d)

    W = cov_whitening(C, rc)
    @test istriu(W)
    @test W'Cr * W ≈ Matrix(I, d, d)

    # Whitening from data

    f = fit(Whitening, X)
    W = f.W
    @test isa(f, Whitening{Float64})
    @test mean(f) === f.mean
    @test length(f) == d
    @test size(f) == (d,d)
    @test istriu(W)
    @test W'C * W ≈ Matrix(I, d, d)
    @test transform(f, X) ≈ W' * (X .- f.mean)

    f = fit(Whitening, X; regcoef=rc)
    W = f.W
    @test W'Cr * W ≈ Matrix(I, d, d)

    f = fit(Whitening, X; mean=mval)
    W = f.W
    @test W'C * W ≈ Matrix(I, d, d)

    f = fit(Whitening, X; mean=0)
    Cx = (X * X') / (n - 1)
    W = f.W
    @test W'Cx * W ≈ Matrix(I, d, d)

    # invsqrtm

    R = invsqrtm(C)
    @test C == C0
    @test R ≈ inv(sqrt(C))

    # mixing types
    X = rand(Float64, 5, 10)
    XX = convert.(Float32, X)

    M = fit(Whitening, X)
    MM = fit(Whitening, XX)

    # mixing types should not error
    transform(M, XX)
    transform(MM, X)

    # type consistency
    @test eltype(mean(M)) == Float64
    @test eltype(mean(MM)) == Float32

    # sparse arrays
    SX = sprand(Float32, d, n, 0.75)
    SM = fit(Whitening, SX; mean=sprand(Float32, 3, 0.75))
    Y = transform(SM, SX)
    @test eltype(Y) == Float32

    # different dimensions
    @test_throws DomainError fit(Whitening, X'; dims=3)
    M1 = fit(Whitening, X'; dims=1)
    M2 = fit(Whitening, X; dims=2)
    @test M1.W == M2.W
    @test_throws DimensionMismatch transform(M1, rand(6,4))
    @test_throws DimensionMismatch transform(M2, rand(4,6))
    Y1 = transform(M1,X')
    Y2 = transform(M2,X)
    @test Y1' == Y2
    @test_throws DimensionMismatch transform(M1, rand(7))
    V1 = transform(M1,X[:,1])
    V2 = transform(M2,X[:,1])
    @test V1 == V2
end


using HigherOrderKernels
using GaussQuadrature
using Test

import HigherOrderKernels: density_kernel, density_bandwidth_constant, roughness, firstnzmoment

for s = 0:20, ν = 2:2:20
    N = s + div(ν, 2)
    x, w = legendre(N)
    @test density_kernel.(PolynomialKernel{s,ν}, x)'*w ≈ 1
end

for ν = 2:2:20
    N = div(ν, 2)
    k = x -> density_kernel(GaussianKernel{ν}, x)
    x, w = hermite(N)
    @test sum(w[i]*k(sqrt(2)*x[i])*exp(x[i]^2) for i=1:N)*sqrt(2) ≈ 1
end

# table 1
@test roughness(     UniformKernel{2}) ≈ 1/2
@test roughness(EpanechnikovKernel{2}) ≈ 3/5
@test roughness(    BiweightKernel{2}) ≈ 5/7
@test roughness(   TriweightKernel{2}) ≈ 350/429
@test roughness(    GaussianKernel{2}) ≈ 1/(2sqrt(pi))
@test firstnzmoment(     UniformKernel{2}) ≈ 1/3
@test firstnzmoment(EpanechnikovKernel{2}) ≈ 1/5
@test firstnzmoment(    BiweightKernel{2}) ≈ 1/7
@test firstnzmoment(   TriweightKernel{2}) ≈ 1/9
@test firstnzmoment(    GaussianKernel{2}) ≈ 1

# table 2
@test roughness(EpanechnikovKernel{4}) ≈ 5/4
@test roughness(    BiweightKernel{4}) ≈ 805/572
@test roughness(   TriweightKernel{4}) ≈ 3780/2431
@test roughness(    GaussianKernel{4}) ≈ 27/(32sqrt(pi))
@test firstnzmoment(EpanechnikovKernel{4}) ≈ -1/21
@test firstnzmoment(    BiweightKernel{4}) ≈ -1/33
@test firstnzmoment(   TriweightKernel{4}) ≈ -3/143
@test firstnzmoment(    GaussianKernel{4}) ≈ -3

# table 3
@test roughness(EpanechnikovKernel{6}) ≈ 1575/832
@test roughness(    BiweightKernel{6}) ≈ 29295/14144
@test roughness(   TriweightKernel{6}) ≈ 301455/134368
@test roughness(    GaussianKernel{6}) ≈ 2265/(2048sqrt(pi))
@test firstnzmoment(EpanechnikovKernel{6}) ≈ 5/429
@test firstnzmoment(    BiweightKernel{6}) ≈ 1/143
@test firstnzmoment(   TriweightKernel{6}) ≈ 1/221
@test firstnzmoment(    GaussianKernel{6}) ≈ 15

# table 4
@test round(density_bandwidth_constant(PolynomialKernel{1,2}), digits=2) == 2.34
@test round(density_bandwidth_constant(PolynomialKernel{1,4}), digits=2) == 3.03
@test round(density_bandwidth_constant(PolynomialKernel{1,6}), digits=2) == 3.53
@test round(density_bandwidth_constant(PolynomialKernel{2,2}), digits=2) == 2.78
@test round(density_bandwidth_constant(PolynomialKernel{2,4}), digits=2) == 3.39
@test round(density_bandwidth_constant(PolynomialKernel{2,6}), digits=2) == 3.84
@test round(density_bandwidth_constant(PolynomialKernel{3,2}), digits=2) == 3.15
@test round(density_bandwidth_constant(PolynomialKernel{3,4}), digits=2) == 3.72
@test round(density_bandwidth_constant(PolynomialKernel{3,6}), digits=2) == 4.13
@test round(density_bandwidth_constant(GaussianKernel{2}), digits=2) == 1.06
@test round(density_bandwidth_constant(GaussianKernel{4}), digits=2) == 1.08
@test round(density_bandwidth_constant(GaussianKernel{6}), digits=2) == 1.08


module Spirograph

using Plots

function lastly()
    plot!(aspectratio = 1, legend = false, axis = false, grid = false, ticks = false)
end

@inline inner_angle(a::Int, b::Int, t::Real)::Float64 = b * t / a

@inline function spot(a::Int, b::Int, offset::Real, t::Real)::Complex{Float64}
    θ = inner_angle(a, b, t)
    return im * ((a + b) * exp(θ * im) + offset * exp(t * im))
end

@inline stop_t(a::Int, b::Int) = 2 * π * (lcm(a, b) ÷ abs(b))

function spiro_points(a::Int, b::Int, offset::Real, step::Real)
    if a <= 0 || b == 0
        error(
            "Radius of fixed circle must be positive and radius of the rolling disk must be nonzero",
        )
    end
    T = stop_t(a, b) #+ step
    return [spot(a, b, offset, t) for t = 0:step:T]
end

_DEFAULT_STEP = 0.05

"""
    spirograph(a,b,offset;args...)
Draw a spirograph picture where
+ `a` is the radius of the fixed wheel,
+ `b` is the radius of the moving wheel, and
+ `offset` is the distance from the center of the moving wheel to the pen.

The optional `args` are passed to `plot`.

## Example
`spirograph(20,-9,12,linecolor=:green)`
"""
function spirograph(a::Int, b::Int, offset::Real, step::Real = _DEFAULT_STEP; args...)
    pts = spiro_points(a, b, offset, step)
    plot(real(pts), imag(pts); args...)
    lastly()
end

"""
    spirograph!
Same as `spirograph` but does not clear the drawing screen first. This way spirograph
images can be combined.
"""
function spirograph!(a::Int, b::Int, offset::Real, step::Real = _DEFAULT_STEP; args...)
    pts = spiro_points(a, b, offset, step)
    plot!(real(pts), imag(pts); args...)
    lastly()
end

export spirograph, spirograph!
end # module


# Gaussian mixture model
struct GMM{T<:AbstractVector{<:Real}, M<:Normal, S<:InverseGamma, W<:Dirichlet} <: MCMC.Model
  y::T
  K::Int
  mu::M
  sigmasq::S
  eta::W
end
function GMM(y::AbstractVector{<:Real}, K::Int; mu=Normal(0, 1),
             sigmasq=InverseGamma(3, 2), eta=Dirichlet(K, 1/K))
  return GMM(y, K, mu, sigmasq, eta)
end
function MCMC.make_init_state(m::GMM)
  nobs = length(m.y)
  eta = rand(m.eta)
  return (mu=rand(m.mu, m.K), sigmasq=rand(m.sigmasq, m.K),
          eta=eta, lambda=rand(Categorical(eta), nobs))
end
function MCMC.update(m::GMM, s::T) where T
  s = setproperty!!(s, :mu, update_mu(m, s))
  s = setproperty!!(s, :sigmasq, update_sigmasq(m, s))
  s = setproperty!!(s, :eta, update_eta(m, s))
  s = setproperty!!(s, :lambda, update_lambda(m, s))

  return s
end
function update_mu(m::GMM, s::T) where T
  m_mu, s_mu = params(m.mu)

  vkernels = zero.(s.mu)
  mkernels = zero.(s.mu)

  for n in eachindex(m.y)
    k = s.lambda[n]
    vkernels[k] += 1
    mkernels[k] += m.y[n]
  end

  vnew = 1 ./ (s_mu^-2 .+ vkernels ./ s.sigmasq)
  mnew = vnew .* (m_mu/(s_mu^2) .+ mkernels ./ s.sigmasq)

  return randn(m.K) .* sqrt.(vnew) + mnew
end
function update_sigmasq(m::GMM, s::T) where T
  a, b = params(m.sigmasq)
  akernel = zero.(s.sigmasq)
  bkernel = zero.(s.sigmasq)

  for n in eachindex(m.y)
    k = s.lambda[n]
    akernel[k] += 1
    bkernel[k] += (m.y[n] - s.mu[k]) ^ 2
  end

  anew = a .+ akernel / 2
  bnew = b .+ bkernel / 2
  s.sigmasq .= rand.(InverseGamma.(anew, bnew))

  return s.sigmasq
end
function update_eta(m::GMM, s::T) where T
  anew = copy(collect(m.eta.alpha))
  for lam in s.lambda
    anew[lam] += 1
  end
  return rand(Dirichlet(anew))
end
function update_lambda(m::GMM, s::T) where T
  logeta = s.eta
  lambda = [let
              logmix = normlogpdf.(s.mu, sqrt.(s.sigmasq), yn) + logeta
              wsample_logprob(logmix)
            end for yn in m.y]
  return lambda
end
@testset "GMM" begin
  true_model = MixtureModel([Normal(3, .6), Normal(-2, 1)], [.3, .7])

  Random.seed!(0)

  nobs = 1000
  y = rand(true_model, nobs)
  model = GMM(y, 2)
  nsamps = 1000
  chain, _ = mcmc(model, nsamps, nburn=1000)

  mu = hcat(map(c -> c.mu, chain)...)
  sigma = hcat(map(c -> sqrt.(c.sigmasq), chain)...)
  eta = hcat(map(c -> c.eta, chain)...)
  lambda = hcat(map(c -> c.lambda, chain)...)

  @test size(mu) == (model.K, nsamps)
  @test size(sigma) == (model.K, nsamps)
  @test size(eta) == (model.K, nsamps)
  @test size(lambda) == (nobs, nsamps)

  function match_components(postmean::AbstractVector{<:Real}, truth::AbstractVector{<:Real}, atol)
    cond1 = abs2(postmean[1] - truth[1]) < atol
    cond2 = abs2(postmean[2] - truth[2]) < atol
    cond3 = abs2(postmean[1] - truth[2]) < atol
    cond4 = abs2(postmean[2] - truth[1]) < atol
    return (cond1 && cond2) || (cond3 && cond4)
  end
  function match_components(postmean, sym::Symbol, atol)
    truth = [getfield(true_model.components[1], sym),
             getfield(true_model.components[2], sym)]
    return match_components(postmean, truth, atol)
  end
  mean_mu = vec(mean(mu, dims=2))
  mean_sigma = vec(mean(sigma, dims=2))
  mean_eta = vec(mean(eta, dims=2))
  # println(mean_mu)
  # println(mean_sigma)
  # println(mean_eta)
  @test match_components(mean_mu, :μ, 1e-1)
  @test match_components(mean_sigma, :σ, 1e-2)
  @test match_components(mean_eta, true_model.prior.p, 1e-2)
end


#=

   The Euro coin problem.

   From Think Bayes, page 33ff
   """
   A statistical statement appeared in "The Guardian" on Friday January 4, 2002:
      When spun on edge 250 times, a Belgian one-euro coin
      came up heads 140 times and tails 110. 'It looks very
      suspicious to me,' said Barry Blight, a statistics lecturer
      at the London School of Economics. 'If the coin were
      unbiased, the chance of getting a result as extreme as
      that would be less than 7%.'

   But do these data give evidence that the coin is biased rather than fair?
   """

  Model 1
  mean(prob): 0.5639764325445531
  Distributions of variable probLt0_5 (num:0)
  0.00000 =>    9896  (0.989600)
  1.00000 =>     104  (0.010400)

  Distributions of variable probGt0_5 (num:0)
  1.00000 =>    9896  (0.989600)
  0.00000 =>     104  (0.010400)

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

        prob     0.5496    0.0296     0.0003    0.0026     52.7368    1.0154        6.8901
       heads   139.9668    1.8488     0.0185    0.0332   2384.8100    1.0002      311.5770
   probLt0_5     0.0006    0.0245     0.0002    0.0006   1544.1256    1.0005      201.7410
   probGt0_5     0.9994    0.0245     0.0002    0.0006   1544.1256    1.0005      201.7410

   Cf ~/blog/euro_coin_problem.blog
      ~/psi/euro_coin_problem.blog
      ~/webppl/euro_coin_problem.jl

=#

using Turing, StatsPlots, DataFrames
include("jl_utils.jl")

@model function euro_coin_problem1()
    n = 250

    # Probability of throwing head
    prob ~ Beta(2,2) # Uniform(0,1)
    coin ~ filldist(flip(prob),n)

    sum250 ~ Dirac(sum([coin[i] == true ? 1 : 0 for i in 1:n]))
    true ~ Dirac(sum250 == 140)

    probLt0_5 ~ Dirac(prob < 0.5)
    probGt0_5 ~ Dirac(prob > 0.5)

end

# Simpler and faster
@model function euro_coin_problem2()
    prob ~ Beta(2,2)
    heads ~ Binomial(250,prob)

    true ~ Dirac(heads==140)

    probLt0_5 ~ Dirac(prob < 0.5)
    probGt0_5 ~ Dirac(prob > 0.5)

end

println("Model 1")
model = euro_coin_problem1()

num_chains = 4

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(5), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 10_000)

# display(chns)
# display(plot(chns))

println("mean(prob): ", mean(chns[:prob]))
show_var_dist_pct(chns, :probLt0_5)
show_var_dist_pct(chns, :probGt0_5)

println("\nModel 2")
model = euro_coin_problem2()

num_chains = 4

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(5), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 10_000)

display(chns)
# display(plot(chns))

show_var_dist_pct(chns, :probLt0_5)
show_var_dist_pct(chns, :probGt0_5)