export Jacobi, Legendre, WeightedJacobi

immutable Jacobi{T,D<:Domain} <: PolynomialSpace{D}
    a::T
    b::T
    domain::D
end
Legendre(domain)=Jacobi(0.,0.,domain)
Legendre()=Legendre(Interval())
Jacobi(a,b,d::Domain)=Jacobi(promote(a,b)...,d)
Jacobi(a,b,d)=Jacobi(a,b,Domain(d))
Jacobi(a,b)=Jacobi(a,b,Interval())
Jacobi(A::Ultraspherical) = Jacobi(order(A)-0.5,order(A)-0.5,domain(A))
Jacobi(A::Chebyshev)=Jacobi(-0.5,-0.5,domain(A))

function Ultraspherical(J::Jacobi)
    if J.a == J.b
        Ultraspherical(J.a+0.5,domain(J))
    else
        error("Cannot construct Ultraspherical with a=$(J.a) and b=$(J.b)")
    end
end

Base.promote_rule{T,V,D}(::Type{Jacobi{T,D}},::Type{Jacobi{V,D}})=Jacobi{promote_type(T,V),D}
Base.convert{T,V,D}(::Type{Jacobi{T,D}},J::Jacobi{V,D})=Jacobi{T,D}(J.a,J.b,J.domain)

typealias WeightedJacobi{T,D} JacobiWeight{Jacobi{T,D},D}

@compat (::Type{WeightedJacobi})(α,β,d::Domain) = JacobiWeight(α,β,Jacobi(β,α,d))
@compat (::Type{WeightedJacobi})(α,β) = JacobiWeight(α,β,Jacobi(β,α))

spacescompatible(a::Jacobi,b::Jacobi) = a.a==b.a && a.b==b.b

function canonicalspace(S::Jacobi)
    if isapproxinteger(S.a+0.5) && isapproxinteger(S.b+0.5)
        Chebyshev(domain(S))
    else
        # return space with parameters in (-1,0.]
        Jacobi(mod(S.a,-1),mod(S.b,-1),domain(S))
    end
end

#####
# jacobirecA/B/C is from dlmf:
# p_{n+1} = (A_n x + B_n)p_n - C_n p_{n-1}
#####
jacobirecA{T}(::Type{T},α,β,k)=k==0&&((α+β==0)||(α+β==-1))?.5*(α+β)+one(T):(2k+α+β+one(T))*(2k+α+β+2one(T))/(2*(k+one(T))*(k+α+β+one(T)))
jacobirecB{T}(::Type{T},α,β,k)=k==0&&((α+β==0)||(α+β==-1))?.5*(α-β)*one(T):(α-β)*(α+β)*(2k+α+β+one(T))/(2*(k+one(T))*(k+α+β+one(T))*(2one(T)*k+α+β))
jacobirecC{T}(::Type{T},α,β,k)=(one(T)*k+α)*(one(T)*k+β)*(2k+α+β+2one(T))/((k+one(T))*(k+α+β+one(T))*(2one(T)*k+α+β))

#####
# jacobirecA/B/C is from dlmf:
# x p_{n-1} =γ_n p_{n-2} + α_n p_{n-1} +  p_n β_n
#####

jacobirecγ{T}(::Type{T},α,β,k)=jacobirecC(T,α,β,k-1)/jacobirecA(T,α,β,k-1)
jacobirecα{T}(::Type{T},α,β,k)=-jacobirecB(T,α,β,k-1)/jacobirecA(T,α,β,k-1)
jacobirecβ{T}(::Type{T},α,β,k)=1/jacobirecA(T,α,β,k-1)

for (REC,JREC) in ((:recα,:jacobirecα),(:recβ,:jacobirecβ),(:recγ,:jacobirecγ),
                   (:recA,:jacobirecA),(:recB,:jacobirecB),(:recC,:jacobirecC))
    @eval $REC{T}(::Type{T},sp::Jacobi,k)=$JREC(T,sp.a,sp.b,k)
end

function jacobip(r::Range,α,β,x)
    n=r[end]+1
    if n<=2
        v=[1.,.5*(α-β+(2+α+β)*x)]
    else
        T=promote_type(Float64,typeof(x))
        v=Vector{T}(n)  # x may be complex
        v[1]=1.
        v[2]=.5*(α-β+(2+α+β)*x)

        for k=2:n-1
            v[k+1]=((x-jacobirecα(T,α,β,k))*v[k] - jacobirecγ(T,α,β,k)*v[k-1])/jacobirecβ(T,α,β,k)
        end
    end
    v[r+1]
end

function jacobip(r::Range,α,β,x::Number)
    if x==1. && α==0.
        ones(length(r))
    elseif x==-1. && β==0.
        (-1.).^r
    else
        n=r[end]+1
        if n<=2
            v=[1.,.5*(α-β+(2+α+β)*x)]
        else
            T=promote_type(Float64,typeof(x))
            v=Vector{T}(n)  # x may be complex
            v[1]=1.
            v[2]=.5*(α-β+(2+α+β)*x)

            for k=2:n-1
                v[k+1]=((x-jacobirecα(T,α,β,k))*v[k] - jacobirecγ(T,α,β,k)*v[k-1])/jacobirecβ(T,α,β,k)
            end
        end
        v[r+1]
    end
end
jacobip(n::Integer,α,β,v)=jacobip(n:n,α,β,v)[1]
jacobip(n::Range,α,β,v::Vector)=transpose(hcat(map(x->jacobip(n,α,β,x),v)...))
jacobip(n::Integer,α,β,v::Vector)=map(x->jacobip(n,α,β,x),v)
jacobip(n,S::Jacobi,v)=jacobip(n,S.a,S.b,v)

include("jacobitransform.jl")
include("JacobiOperators.jl")

for op in (:(Base.ones),:(Base.zeros))
    @eval ($op){T<:Number}(::Type{T},S::Jacobi)=Fun(($op)(T,1),S)
    @eval ($op)(S::Jacobi)=Fun(($op)(1),S)
end

function identity_fun(J::Jacobi)
    if domain(J)==Interval()
        Fun([(J.b-J.a)/(2+J.a+J.b),2.0/(2+J.a+J.b)],J)
    else
        d=domain(J)
        complexlength(d)/2*(Fun([(J.b-J.a)/(2+J.a+J.b),2.0/(2+J.a+J.b)],J)+1.)+first(d)
    end
end

setdomain(S::Jacobi,d::Domain)=Jacobi(S.a,S.b,d)

# O(min(m,n)) Jacobi conjugated inner product

function conjugatedinnerproduct{S,V}(sp::Jacobi,u::Vector{S},v::Vector{V})
    T,mn = promote_type(S,V),min(length(u),length(v))
    α,β = sp.a,sp.b
    if mn > 1
        wi = 2^(α+β+1)*gamma(α+1)*gamma(β+1)/gamma(α+β+2)
        ret = u[1]*wi*v[1]
        for i=2:mn
            wi *= (α+i-1)*(β+i-1)/(i-1)/(i-1+α+β)*(2i+α+β-3)/(2i+α+β-1)
            ret += u[i]*wi*v[i]
        end
        return ret
    elseif mn > 0
        wi = 2^(α+β+1)*gamma(α+1)*gamma(β+1)/gamma(α+β+2)
        return u[1]*wi*v[1]
    else
        return zero(promote_type(eltype(u),eltype(v)))
    end
end

function bilinearform{J<:Jacobi}(f::Fun{J},g::Fun{J})
    @assert domain(f) == domain(g)
    if f.space.a == g.space.a == 0. && f.space.b == g.space.b == 0.
        return complexlength(domain(f))/2*conjugatedinnerproduct(g.space,f.coefficients,g.coefficients)
    else
        return defaultbilinearform(f,g)
    end
end

function bilinearform{J<:Jacobi,DD<:Interval}(f::Fun{JacobiWeight{J,DD}},g::Fun{J})
    @assert domain(f) == domain(g)
    if f.space.β == f.space.space.a == g.space.a && f.space.α == f.space.space.b == g.space.b
        return complexlength(domain(f))/2*conjugatedinnerproduct(g.space,f.coefficients,g.coefficients)
    else
        return defaultbilinearform(f,g)
    end
end

function bilinearform{J<:Jacobi,DD<:Interval}(f::Fun{J},g::Fun{JacobiWeight{J,DD}})
    @assert domain(f) == domain(g)
    if g.space.β == g.space.space.a == f.space.a && g.space.α == g.space.space.b == f.space.b
        return complexlength(domain(f))/2*conjugatedinnerproduct(f.space,f.coefficients,g.coefficients)
    else
        return defaultbilinearform(f,g)
    end
end

function bilinearform{J<:Jacobi,DD<:Interval}(f::Fun{JacobiWeight{J,DD}},g::Fun{JacobiWeight{J,DD}})
    @assert domain(f) == domain(g)
    if f.space.β + g.space.β == f.space.space.a == g.space.space.a && f.space.α + g.space.α == f.space.space.b == g.space.space.b
        return complexlength(domain(f))/2*conjugatedinnerproduct(f.space.space,f.coefficients,g.coefficients)
    else
        return defaultbilinearform(f,g)
    end
end

function linebilinearform{J<:Jacobi}(f::Fun{J},g::Fun{J})
    @assert domain(f) == domain(g)
    if f.space.a == g.space.a == 0. && f.space.b == g.space.b == 0.
        return arclength(domain(f))/2*conjugatedinnerproduct(g.space,f.coefficients,g.coefficients)
    else
        return defaultlinebilinearform(f,g)
    end
end

function linebilinearform{J<:Jacobi,DD<:Interval}(f::Fun{JacobiWeight{J,DD}},g::Fun{J})
    @assert domain(f) == domain(g)
    if f.space.β == f.space.space.a == g.space.a && f.space.α == f.space.space.b == g.space.b
        return arclength(domain(f))/2*conjugatedinnerproduct(g.space,f.coefficients,g.coefficients)
    else
        return defaultlinebilinearform(f,g)
    end
end

function linebilinearform{J<:Jacobi,DD<:Interval}(f::Fun{J},g::Fun{JacobiWeight{J,DD}})
    @assert domain(f) == domain(g)
    if g.space.β == g.space.space.a == f.space.a && g.space.α == g.space.space.b == f.space.b
        return arclength(domain(f))/2*conjugatedinnerproduct(f.space,f.coefficients,g.coefficients)
    else
        return defaultlinebilinearform(f,g)
    end
end

function linebilinearform{J<:Jacobi,DD<:Interval}(f::Fun{JacobiWeight{J,DD}},g::Fun{JacobiWeight{J,DD}})
    @assert domain(f) == domain(g)
    if f.space.β + g.space.β == f.space.space.a == g.space.space.a && f.space.α + g.space.α == f.space.space.b == g.space.space.b
        return arclength(domain(f))/2*conjugatedinnerproduct(f.space.space,f.coefficients,g.coefficients)
    else
        return defaultlinebilinearform(f,g)
    end
end


export goldensection, hyperband

"""

    goldensection(f,n;kwargs) => (fmin,xmin)

Find the minimum of `f` using concurrent golden section search in `n`
dimensions. See `Knet.goldensection_demo()` for an example.

`f` is a function from a `Vector{Float64}` of length `n` to a
`Number`.  It can return `NaN` for out of range inputs.  Goldensection
will always start with a zero vector as the initial input to `f`, and
the initial step size will be 1 in each dimension.  The user should
define `f` to scale and shift this input range into a vector
meaningful for their application. For positive inputs like learning
rate or hidden size, you can use a transformation such as `x0*exp(x)`
where `x` is a value `goldensection` passes to `f` and `x0` is your
initial guess for this value. This will effectively start the search
at `x0`, then move with multiplicative steps.

I designed this algorithm combining ideas from [Golden Section
Search](http://apps.nrbook.com/empanel/index.html?pg=492) and [Hill
Climbing Search](https://en.wikipedia.org/wiki/Hill_climbing). It
essentially runs golden section search concurrently in each dimension,
picking the next step based on estimated gain.

# Keyword arguments
* `dxmin=0.1`: smallest step size.
* `accel=φ`: acceleration rate. Golden ratio `φ=1.618...` is best.
* `verbose=false`: use `true` to print individual steps.
* `history=[]`: cache of `[(x,f(x)),...]` function evaluations.

"""
function goldensection(f,n; dxmin=0.1, accel=1.618033988749895, history=[], verbose=false)

    function feval(x)           # so we don't repeat function evaluations
        for (k,v) in history
            if isapprox(x,k)
                return v
            end
        end
        fx = f(x)
        push!(history, (x,fx))
        return fx
    end

    function setindex(x,v,d)    # non-mutating setindex
        y = copy(x)
        y[d] = v
        return y
    end

    x0 = zeros(n)               # initial point
    f0 = feval(x0)              # initial value
    dx = ones(n)                # step sizes
    df = Inf * ones(n)          # expected gains
    while maximum(abs.(dx)) > dxmin
        i = argmax(df)
        x1 = setindex(x0,x0[i]+dx[i],i)
        f1 = feval(x1)
        if verbose; println((; f0,x0,f1,x1,dx,df)); end
        isnan(f1) && (f1=f0+df[i])
        if f1 < f0
            dx[i] = accel * dx[i]
            df[i] = accel * (f0-f1)
            x0,f0 = x1,f1
            for j = 1:length(df)
                if abs(dx[j]) < dxmin * accel
                    dx[j] = sign(dx[j]) * dxmin * accel
                    df[j] = max(df[j],0) # max(df[j],-1-df[j])
                end
            end
        else
            dx[i] = -dx[i] / accel
            df[i] = (f1-f0) / accel
            if abs(dx[i]) < dxmin
                df[i] = -1 # -1-df[i]
            end
        end
    end
    return (f0,x0)
end

function goldensection_demo()
    include(Knet.dir("examples","mnist.jl"))
    neval = 0

    function f(x)
        neval += 1
        winit,lr,hidden = xform(x)
        if hidden < 10000
            w = MNIST.main("--winit $winit --lr $lr --hidden $hidden --seed 1 --epochs 1 --fast")
            corr,loss = MNIST.accuracy(w,MNIST.dtst)
        else
            corr,loss = NaN,NaN # prevent huge weights
        end
        println((:neval,neval,:loss,loss,:corr,corr,:winit,winit,:lr,lr,:hidden,hidden))
        return loss
    end

    function xform(x)
        winit,lr,hidden = exp(x) .* [ 0.01, 0.01, 100.0 ]
        hidden = ceil(Int, hidden)
        (winit,lr,hidden)
    end

    goldensection(f,3)
end

"""

    hyperband(getconfig, getloss, maxresource=27, reduction=3)

Hyperparameter optimization using the hyperband algorithm from ([Lisha
et al. 2016](https://arxiv.org/abs/1603.06560)).  You can try a simple
MNIST example using `Knet.hyperband_demo()`.

# Arguments
* `getconfig()` returns random configurations with a user defined type and distribution.
* `getloss(c,n)` returns loss for configuration `c` and number of resources (e.g. epochs) `n`.
* `maxresource` is the maximum number of resources any one configuration should be given.
* `reduction` is an algorithm parameter (see paper), 3 is a good value.

"""
function hyperband(getconfig, getloss, maxresource=27, reduction=3)
    smax = floor(Int, log(maxresource)/log(reduction))
    B = (smax + 1) * maxresource
    best = (Inf,)
    for s in smax:-1:0
        n = ceil(Int, (B/maxresource)*((reduction^s)/(s+1)))
        r = maxresource / (reduction^s)
        curr = halving(getconfig, getloss, n, r, reduction, s)
        if curr[1] < best[1]; (best=curr); end
    end
    return best
end

# TODO: document Successive/Sequential Halving:
# http://www.jmlr.org/proceedings/papers/v51/jamieson16.pdf,
# http://www.jmlr.org/proceedings/papers/v28/karnin13.pdf,
# Successive  Reject:
# http://certis.enpc.fr/~audibert/Mes%20articles/COLT10.pdf
function halving(getconfig, getloss, n, r=1, reduction=3, s=round(Int, log(n)/log(reduction)))
    best = (Inf,)
    T = [ getconfig() for i=1:n ]
    for i in 0:s
        ni = floor(Int,n/(reduction^i))
        ri = r*(reduction^i)
        # println((:s,s,:n,n,:r,r,:i,i,:ni,ni,:ri,ri,:T,length(T)))
        L = [ getloss(t, ri) for t in T ]
        l = sortperm(L); l1=l[1]
        L[l1] < best[1] && (best = (L[l1],ri,T[l1])) # ;println("best1: $best"))
        T = T[l[1:floor(Int,ni/reduction)]]
    end
    # println("best2: $best")
    return best
end

function hyperband_demo()
    include(Knet.dir("examples","mnist.jl"))
    best = (Inf,)
    neval = 0

    function getloss(config,epochs)
        neval += 1
        winit,lr,hidden = config
        epochs = round(Int,epochs)
        w = MNIST.main("--winit $winit --lr $lr --hidden $hidden --epochs $epochs --seed 1 --fast")
        corr,loss = MNIST.accuracy(w,MNIST.dtst)
        println((:epochs,epochs,:corr,corr,:loss,loss,:winit,winit,:lr,lr,:hidden,hidden))
        if loss < best[1]
            best = (loss, config, epochs)
        end
        return loss
    end

    function getconfig()
        winit = 0.001^rand()
        lr = 0.001^rand()
        hidden = 16 + floor(Int, 10000^rand())
        return (winit,lr,hidden)
    end

    hyperband(getconfig, getloss)
    println((:neval,neval,:minloss,best[1],:epochs,best[3],:winit,best[2][1],:lr,best[2][2],:hidden,best[2][3]))
end


### A Pluto.jl notebook ###
# v0.11.14

using Markdown
using InteractiveUtils

# ╔═╡ 81cdc1aa-04fd-11eb-3235-bff0e3f52535
using Pkg, DrWatson

# ╔═╡ 81cdee8c-04fd-11eb-29f5-85927396b6dc
begin
	@quickactivate "StatisticalRethinkingTuring"
	using Turing
	using StatisticalRethinking
end

# ╔═╡ 5e06a11c-04f6-11eb-0116-1fd51d614f99
md"## clip-04-69-71t"

# ╔═╡ b8b8de5e-0512-11eb-0932-830c3055e00a
#include(srcdir("tools.jl"))

md"### snippet 4.64"

# ╔═╡ 81ce6d32-04fd-11eb-2c4f-77bdf6a3f9e2
begin
	df = CSV.read(sr_datadir("Howell1.csv"), DataFrame)
	scale!(df, [:weight])
end;

# ╔═╡ 82407c7c-04fd-11eb-2bba-3783e79b0671
md"### snippet 4.69"

# ╔═╡ da70abc2-0527-11eb-31fa-d5b029193ba5
f_cube(weight_s, a, b1, b2, b3) = a + b1 * weight_s + b2 * weight_s^2 + b3 * weight_s^3

# ╔═╡ 82483e30-04fd-11eb-1290-b98ae5d06a45
@model function cube(weight_s, heights)
    a ~ Normal(178, 20)
    b1 ~ LogNormal(0, 1)
    b2 ~ Normal(0, 10)
    b3 ~ Normal(0, 10)
    σ ~ Uniform(0, 50)
    μ = f_cube.(weight_s, a, b1, b2, b3)
    heights ~ MvNormal(μ, σ)
end

# ╔═╡ 8250f05c-04fd-11eb-0370-899604fffdfd
begin
	quap4_6t = quap(cube(df.weight_s, df.height), NelderMead())
	quap4_6t.coef
end

# ╔═╡ 825e8410-04fd-11eb-3e36-6d3bac027215
md"### snippets 4.70, 4.71"

# ╔═╡ 423cd608-0529-11eb-126e-7fde3de2165b
begin
	weight_seq_2 = range(-2.2, 2, length = 30)
	post_2 = DataFrame(rand(quap4_6t.distr, 1_000)', quap4_6t.params)
	mu_2 = f_cube.(weight_seq_2, post_2.a', post_2.b1', post_2.b2', post_2.b3') |> meanlowerupper
	sim_2 = rand.(Normal.(mu_2.raw, post_2.σ')) |> meanlowerupper
	weight_seq_rescaled = weight_seq_2 .* std(df.weight) .+ mean(df.weight)
	scatter(df.weight, df.height, ms = 3, alpha = 0.7, legend = false)
	plot!(weight_seq_rescaled, mu_2.mean, ribbon = (mu_2.mean .- mu_2.lower, mu_2.upper .- mu_2.mean))
	plot!(weight_seq_rescaled, sim_2.lower, fillrange = sim_2.upper, alpha = 0.3, la = 0.0, c = 2)
end

# ╔═╡ 4f386d5c-0573-11eb-084c-1758217d06f9
md"## End of clip-04-69-71t.jl"

# ╔═╡ Cell order:
# ╟─5e06a11c-04f6-11eb-0116-1fd51d614f99
# ╠═81cdc1aa-04fd-11eb-3235-bff0e3f52535
# ╠═81cdee8c-04fd-11eb-29f5-85927396b6dc
# ╟─b8b8de5e-0512-11eb-0932-830c3055e00a
# ╠═81ce6d32-04fd-11eb-2c4f-77bdf6a3f9e2
# ╟─82407c7c-04fd-11eb-2bba-3783e79b0671
# ╠═da70abc2-0527-11eb-31fa-d5b029193ba5
# ╠═82483e30-04fd-11eb-1290-b98ae5d06a45
# ╠═8250f05c-04fd-11eb-0370-899604fffdfd
# ╟─825e8410-04fd-11eb-3e36-6d3bac027215
# ╠═423cd608-0529-11eb-126e-7fde3de2165b
# ╟─4f386d5c-0573-11eb-084c-1758217d06f9


using ScalarRadau
using DifferentialEquations
using BenchmarkTools
using Statistics
using PyPlot

##

#stiff, scalar ODE for testing
f_sr(x, y, _) = 50*(cos(x) - y)
f_de(u, _, t) = 50*(cos(t) - u)
#known solution
s(x) = (50*(-50*exp(-50*x) + sin(x) + 50*cos(x)))/2501

##

#number of tolerances to sample
N = 10
#tolerance values
tols = 10 .^ LinRange(-6, -2, N)
#times and errors for ScalarRadau evaluations
t_sr = zeros(N)
e_sr = zeros(N)
#times and errors for DifferentialEquations evaluations
t_de = zeros(N)
e_de = zeros(N)
#starting point
x₀ = 0.0
y₀ = 0.0
#finishing coordinate
xₙ = 3.0
#exact solution end values
sₙ = s(xₙ)
#the problem for DifferentialEquations
prob = ODEProblem(f_de, x₀, (x₀,xₙ))

##

for (i,tol) ∈ enumerate(tols)

    #ScalarRadau error
    yₙ = radau!(y, x, f_sr, y₀, x₀, xₙ, rtol=tol, atol=tol)
    e_sr[i] = abs(yₙ - sₙ)/abs(sₙ)
    #ScalarRadau timing
    b = @benchmark radau($f_sr, $y₀, $x₀, $xₙ, rtol=$tol, atol=$tol)
    t_sr[i] = mean(b.times)

    #DifferentialEquations error
    sol = solve(prob, RadauIIA5(), reltol=tol, abstol=tol, save_start=false, saveat=xₙ)
    e_de[i] = abs(sol.u[end] - sₙ)/abs(sₙ)
    #DifferentialEquations timing
    b = @benchmark solve($prob, $RadauIIA5(), reltol=$tol, abstol=$tol, save_start=false, saveat=$xₙ)
    t_de[i] = mean(b.times)

end


using Printf

using Oceananigans.Operators: ℑxyᶜᶠᵃ, ℑxyᶠᶜᵃ

#####
##### Physical constants
#####

const Ω_Earth = 7.292115e-5 # [s⁻¹] https://en.wikipedia.org/wiki/Earth%27s_rotation#Angular_speed
const R_Earth = 6371.0e3    # Mean radius of the Earth [m] https://en.wikipedia.org/wiki/Earth

#####
##### Functions for non-rotating models
#####

@inline x_f_cross_U(i, j, k, grid::AbstractGrid{FT}, ::Nothing, U) where FT = zero(FT)
@inline y_f_cross_U(i, j, k, grid::AbstractGrid{FT}, ::Nothing, U) where FT = zero(FT)
@inline z_f_cross_U(i, j, k, grid::AbstractGrid{FT}, ::Nothing, U) where FT = zero(FT)

#####
##### The 'FPlane' approximation. This is equivalent to a model with a constant
##### rotation rate around its vertical axis.
#####

"""
    FPlane{FT} <: AbstractRotation

A parameter object for constant rotation around a vertical axis.
"""
struct FPlane{FT} <: AbstractRotation
    f :: FT
end

"""
    FPlane([FT=Float64;] f=nothing, rotation_rate=Ω_Earth, latitude=nothing)

Returns a parameter object for constant rotation at the angular frequency
`f/2`, and therefore with background vorticity `f`, around a vertical axis.
If `f` is not specified, it is calculated from `rotation_rate` and
`latitude` according to the relation `f = 2*rotation_rate*sind(latitude).

By default, `rotation_rate` is assumed to be Earth's.

Also called `FPlane`, after the "f-plane" approximation for the local effect of
a planet's rotation in a planar coordinate system tangent to the planet's surface.
"""
function FPlane(FT::DataType=Float64; f=nothing, rotation_rate=Ω_Earth, latitude=nothing)

    use_f = !isnothing(f)
    use_planet_parameters = !isnothing(latitude)

    if !xor(use_f, use_planet_parameters)
        throw(ArgumentError("Either both keywords rotation_rate and latitude must be " *
                            "specified, *or* only f must be specified."))
    end

    if use_f
        return FPlane{FT}(f)
    elseif use_planet_parameters
        return FPlane{FT}(2rotation_rate*sind(latitude))
    end
end

@inline x_f_cross_U(i, j, k, grid, coriolis::FPlane, U) = - coriolis.f * ℑxyᶠᶜᵃ(i, j, k, grid, U.v)
@inline y_f_cross_U(i, j, k, grid, coriolis::FPlane, U) =   coriolis.f * ℑxyᶜᶠᵃ(i, j, k, grid, U.u)
@inline z_f_cross_U(i, j, k, grid::AbstractGrid{FT}, coriolis::FPlane, U) where FT = zero(FT)

Base.show(io::IO, f_plane::FPlane{FT}) where FT =
    println(io, "FPlane{$FT}: f = ", @sprintf("%.2e", f_plane.f))

#####
##### βeta Plane
#####

"""
    BetaPlane{T} <: AbstractRotation

A parameter object for meridionally increasing Coriolis parameter (`f = f₀ + βy`).
"""
struct BetaPlane{T} <: AbstractRotation
    f₀ :: T
     β :: T
end

"""
    BetaPlane([T=Float64;] f₀=nothing, β=nothing,
                           rotation_rate=Ω_Earth, latitude=nothing, radius=R_Earth)

A parameter object for meridionally increasing Coriolis parameter (`f = f₀ + βy`).

The user may specify both `f₀` and `β`, or the three parameters `rotation_rate`,
`latitude`, and `radius` that specify the rotation rate and radius of a planet, and
the central latitude at which the `β`-plane approximation is to be made.

By default, the `rotation_rate` and planet `radius` is assumed to be Earth's.
"""
function BetaPlane(T=Float64; f₀=nothing, β=nothing,
                              rotation_rate=Ω_Earth, latitude=nothing, radius=R_Earth)

    use_f_and_β = !isnothing(f₀) && !isnothing(β)
    use_planet_parameters = !isnothing(latitude)

    if !xor(use_f_and_β, use_planet_parameters)
        throw(ArgumentError("Either both keywords f₀ and β must be specified, " *
                            "*or* all of rotation_rate, latitude, and radius."))
    end

    if use_planet_parameters
        f₀ = 2rotation_rate * sind(latitude)
         β = 2rotation_rate * cosd(latitude) / radius
     end

    return BetaPlane{T}(f₀, β)
end

@inline x_f_cross_U(i, j, k, grid, coriolis::BetaPlane, U) =
    @inbounds - (coriolis.f₀ + coriolis.β * grid.yC[j]) * ℑxyᶠᶜᵃ(i, j, k, grid, U.v)
@inline y_f_cross_U(i, j, k, grid, coriolis::BetaPlane, U) =
    @inbounds   (coriolis.f₀ + coriolis.β * grid.yF[j]) * ℑxyᶜᶠᵃ(i, j, k, grid, U.u)
@inline z_f_cross_U(i, j, k, grid::AbstractGrid{FT}, coriolis::BetaPlane, U) where FT = zero(FT)

Base.show(io::IO, β_plane::BetaPlane{FT}) where FT =
    println(io, "BetaPlane{$FT}: ", @sprintf("f₀ = %.2e, β = %.2e", β_plane.f₀, β_plane.β))


if VERSION >= v"0.7"
  using Test
  using LinearAlgebra
else
  using Base.Test
  include("random.jl")
end
include("../src/define_constants.jl")
include("loglinspace.jl")

include("test_cel_vec.jl")
include("test_s2_gradient.jl")
include("test_transit_poly.jl")

#include("test_IJv_derivative.jl")
#include("test_Mn.jl")
#include("test_sn_jacobian2.jl")

skip_plots = false
#skip_plots = true
include("test_transit_poly_gradient2.jl")
include("test_integrate_transit_simpson_vec.jl")


include("7.1.1 (a+b arcsinh(c x))^n.jl")
include("7.1.2 (d x)^m (a+b arcsinh(c x))^n.jl")
include("7.1.3 (d+e x^2)^p (a+b arcsinh(c x))^n.jl")
include("7.1.4 (f x)^m (d+e x^2)^p (a+b arcsinh(c x))^n.jl")
include("7.1.5 u (a+b arcsinh(c x))^n.jl")
include("7.1.6 Miscellaneous inverse hyperbolic sine.jl")

src/rules/7 Inverse hyperbolic functions/7.1 Inverse hyperbolic sine/.jl


function encrypt(msg::AbstractString, key::AbstractString)
    msg    = uppercase(join(filter(isalpha, collect(msg))))
    key    = uppercase(join(filter(isalpha, collect(key))))
    msglen = length(msg)
    keylen = length(key)

    if keylen < msglen
        key = repeat(key, div(msglen - keylen, keylen) + 2)[1:msglen]
    end

    enc = Vector{Char}(msglen)

    @inbounds for i in 1:length(msg)
        enc[i] = Char((Int(msg[i]) + Int(key[i]) - 130) % 26 + 65)
    end

    return join(enc)
end

function decrypt(enc::AbstractString, key::AbstractString)
    enc    = uppercase(join(filter(isalpha, collect(enc))))
    key    = uppercase(join(filter(isalpha, collect(key))))
    msglen = length(enc)
    keylen = length(key)

    if keylen < msglen
        key = repeat(key, div(msglen - keylen, keylen) + 2)[1:msglen]
    end

    msg = Vector{Char}(msglen)

    @inbounds for i in 1:length(enc)
        msg[i] = Char((Int(enc[i]) - Int(key[i]) + 26) % 26 + 65)
    end

    return join(msg)
end

const messages = ("Attack at dawn.", "Don't attack.", "The war is over.", "Happy new year!")
const key = "LEMON"

for msg in messages
    enc = encrypt(msg, key)
    dec = decrypt(enc, key)
    println("Original: $msg\n -> encrypted: $enc\n -> decrypted: $dec")
end


using Flux, NNlib, XLA, Test

xresult(f, args...) = collect(xla(f)(args...))

x, y = randn(2), randn(1, 2)
@test collect(xla(*)(x, y)) ≈ x*y

f(x) = gradient(x -> sum(x.*x), x)

@test collect(xla(f)([1, 2, 3])[1]) == [2, 4, 6]

f(x) = gradient(x -> sum(x.+x), x)

@test collect(xla(f)([1, 2, 3])[1]) == [2, 2, 2]

W = rand(2, 3)

f(x) = gradient(x -> sum(W*x), x)

@test collect(xla(f)([1.0, 2, 3])[1]) == f([1.0, 2, 3])[1]

@test xla(sin'')(0.5) == -sin(0.5)

x = rand(10)
m = Chain(Dense(10, 5), Dense(5, 2)) |> f64

f(x) = gradient(m -> sum(m(x)), m)[1]

xf = xla(f)

@test collect(xf(x).layers[1].W) ≈ f(x).layers[1].W

x = rand(10, 10, 5, 1)
w = rand(2, 2, 5, 3)
y = rand(size(conv(x, w))...)

@test xresult(conv, x, w) ≈ conv(x, w)

@test ∇conv_data(y, w, DenseConvDims(x, w)) ≈ xresult(y, w) do y, w
  ∇conv_data(y, w, DenseConvDims(x, w))
end

@test ∇conv_filter(x, y, DenseConvDims(x, w)) ≈ xresult(x, y) do x, y
  ∇conv_filter(x, y, DenseConvDims(x, w))
end

@test xresult(x -> maxpool(x, (2, 2)), x) ≈ maxpool(x, (2, 2))

y = maxpool(x, (2, 2))

@test xresult(x -> ∇maxpool(y, y, x, PoolDims(x, (2, 2))), x) ≈ ∇maxpool(y, y, x, PoolDims(x, (2, 2)))


#=
Part of this work is derived from astropy/photutils. The relevant derivations
are considered under a BSD 3-clause license. =#

"""
    EllipticalAperture(x, y, a, b, θ=0)
    EllipticalAperture(position, a, b, θ=0)

An elliptical aperture with semi-major axis `a`, semi-minor axis `b`, and position angle `θ`. `a` and `b` must be ≥ 0, `θ` is measured in degrees counter-clockwise the standard x-axis.

# Examples
```jldoctest
julia> ap = EllipticalAperture(0, 0, 4, 2, 35)
5×7 EllipticalAperture{Int64} with indices -2:2×-3:3:
 0.873382  1.0       1.0       0.796137  0.23968   0.0       0.0
 0.844185  1.0       1.0       1.0       0.990119  0.435284  0.0
 0.324917  0.997821  1.0       1.0       1.0       0.997821  0.324917
 0.0       0.435284  0.990119  1.0       1.0       1.0       0.844185
 0.0       0.0       0.23968   0.796137  1.0       1.0       0.873382
```
"""
struct EllipticalAperture{T <: Number} <: AbstractAperture{T}
    x::T
    y::T
    a::T
    b::T
    theta::T
end

EllipticalAperture(x::Number, y::Number, a, b, theta=0) = EllipticalAperture(promote(x, y, a, b, theta)...)
EllipticalAperture(center, a, b, theta=0) = EllipticalAperture(center..., a, b, theta)

function Base.show(io::IO, e::EllipticalAperture)
    print(io, "EllipticalAperture($(e.x), $(e.y), a=$(e.a), b=$(e.b), θ=$(e.theta)°)")
end

oblique_coefficients(ap::EllipticalAperture) = oblique_coefficients(ap.a, ap.b, ap.theta)
oblique_coefficients(ap::Subpixel{T,<:EllipticalAperture}) where {T} = oblique_coefficients(ap.ap)

function oblique_coefficients(a, b, theta)
    sintheta, costheta = sincosd(theta)
    a2 = a^2
    b2 = b^2
    cxx = costheta^2 / a2 + sintheta^2 / b2
    cyy = sintheta^2 / a2 + costheta^2 / b2
    cxy = 2 * costheta * sintheta * (1 / a2 - 1 / b2)
    return cxx, cyy, cxy
end

function overlap(ap::EllipticalAperture, i, j)
    cxx, cyy, cxy = oblique_coefficients(ap)
    flags = (
        inside_ellipse(j - 0.5, i - 0.5, ap.x, ap.y, cxx, cyy, cxy),
        inside_ellipse(j - 0.5, i + 0.5, ap.x, ap.y, cxx, cyy, cxy),
        inside_ellipse(j + 0.5, i - 0.5, ap.x, ap.y, cxx, cyy, cxy),
        inside_ellipse(j + 0.5, i + 0.5, ap.x, ap.y, cxx, cyy, cxy)
    )
    all(flags) && return Inside
    all(!, flags) && return Outside

    return Partial
end

function elliptical_bounds(cx, cy, a, b, theta)
    iszero(a) && return cx, cx, cy, cy
    sintheta, costheta = sincosd(theta)

    t = atan(-b * tand(theta), a)

    sint, cost = sincos(t)
    xmin = cx + a * cost * costheta - b * sint * sintheta
    xmax = xmin

    for n in -2:2
        sint, cost = sincos(t + n * pi)
        xmin = min(xmin, cx + a * cost * costheta - b * sint * sintheta)
        xmax = max(xmax, cx + a * cost * costheta - b * sint * sintheta)
    end

    t2 = atan(b * cotd(theta),  a)

    sint, cost = sincos(t2)
    ymin = cy + b * sint * costheta + a * cost * sintheta
    ymax = ymin

    for n in -2:2
        sint, cost = sincos(t2 + n * pi)
        ymin = min(ymin, cy + b * sint * costheta + a * cost * sintheta)
        ymax = max(ymax, cy + b * sint * costheta + a * cost * sintheta)
    end

    xmin = ceil(Int, xmin)
    xmax = floor(Int, xmax)
    ymin = ceil(Int, ymin)
    ymax = floor(Int, ymax)

    return xmin, xmax, ymin, ymax
end

bounds(e::EllipticalAperture) = elliptical_bounds(e.x, e.y, e.a, e.b, e.theta)

partial(ap::EllipticalAperture, x, y) = elliptical_overlap_exact(x - 0.5, y - 0.5, x + 0.5, y + 0.5, ap.a, ap.b, ap.theta)
partial(sub_ap::Subpixel{T,<:EllipticalAperture}, x, y) where {T} = elliptical_overlap_single_subpixel(x - 0.5, y - 0.5, x + 0.5, y + 0.5, oblique_coefficients(sub_ap.ap)..., sub_ap.N)

#######################################################

"""
    EllipticalAnnulus(x, y, a_in, a_out, b_out, θ=0)
    EllipticalAnnulus(position, a_in, a_out, b_out, θ=0)
An elliptical annulus with inner semi-major axis `a_in`, outer semi-major axis `a_out`, outer semi-minor axis `b_out`, and position angle `θ`.
`a_out` ≥ `a_in` ≥ 0 and `b_out` must be ≥ 0, `θ` is measured in degrees counter-clockwise the standard x-axis.

`b_in` will automatically be calculated from `(a_in / a_out) * b_out`. Note this may cause a type instability.
# Examples
```jldoctest
julia> ap = EllipticalAnnulus(0, 0, 4, 10, 5, 45)
15×15 EllipticalAnnulus{Float64} with indices -7:7×-7:7:
 0.594853   1.0       1.0       1.0         …  0.0       0.0       0.0
 1.0        1.0       1.0       1.0            0.0       0.0       0.0
 1.0        1.0       1.0       1.0            0.0       0.0       0.0
 1.0        1.0       1.0       1.0            0.0       0.0       0.0
 1.0        1.0       1.0       1.0            0.0       0.0       0.0
 0.814163   1.0       1.0       1.0         …  0.414163  0.0       0.0
 0.369432   1.0       1.0       1.0            0.975704  0.193728  0.0
 0.0112571  0.809079  1.0       1.0            1.0       0.809079  0.0112571
 0.0        0.193728  0.975704  1.0            1.0       1.0       0.369432
 0.0        0.0       0.414163  1.0            1.0       1.0       0.814163
 0.0        0.0       0.0       0.546165    …  1.0       1.0       1.0
 0.0        0.0       0.0       0.00252321     1.0       1.0       1.0
 0.0        0.0       0.0       0.0            1.0       1.0       1.0
 0.0        0.0       0.0       0.0            1.0       1.0       1.0
 0.0        0.0       0.0       0.0            1.0       1.0       0.594853
```
"""
struct EllipticalAnnulus{T <: Number} <: AbstractAperture{T}
    x::T
    y::T
    a_in::T
    b_in::T
    a_out::T
    b_out::T
    theta::T
end

EllipticalAnnulus(x::Number, y::Number, a_in, a_out, b_out, theta=0) = EllipticalAnnulus(promote(x, y, a_in, a_in / a_out * b_out, a_out, b_out, theta)...)
EllipticalAnnulus(center, a_in, a_out, b_out, theta=0) = EllipticalAnnulus(center..., a_in, a_out, b_out, theta)

function Base.show(io::IO, e::EllipticalAnnulus)
    print(io, "EllipticalAnnulus($(e.x), $(e.y), a_in=$(e.a_in), a_out=$(e.a_out), b_in=$(e.b_in), b_out=$(e.b_out), θ=$(e.theta)°)")
end

function overlap(ap::EllipticalAnnulus, i, j)
    coeffs_out = oblique_coefficients(ap.a_out, ap.b_out, ap.theta)
    flags_out = (
        inside_ellipse(j - 0.5, i - 0.5, ap.x, ap.y, coeffs_out...),
        inside_ellipse(j - 0.5, i + 0.5, ap.x, ap.y, coeffs_out...),
        inside_ellipse(j + 0.5, i - 0.5, ap.x, ap.y, coeffs_out...),
        inside_ellipse(j + 0.5, i + 0.5, ap.x, ap.y, coeffs_out...)
    )

    coeffs_in = oblique_coefficients(ap.a_in, ap.b_in, ap.theta)
    flags_in = (
        inside_ellipse(j - 0.5, i - 0.5, ap.x, ap.y, coeffs_in...),
        inside_ellipse(j - 0.5, i + 0.5, ap.x, ap.y, coeffs_in...),
        inside_ellipse(j + 0.5, i - 0.5, ap.x, ap.y, coeffs_in...),
        inside_ellipse(j + 0.5, i + 0.5, ap.x, ap.y, coeffs_in...)
    )

   all(flags_out) && all(!, flags_in) && return Inside
   all(flags_in) || all(!, flags_out) && return Outside

    return Partial
end

bounds(e::EllipticalAnnulus) =  elliptical_bounds(e.x, e.y, e.a_out, e.b_out, e.theta)

function partial(ap::EllipticalAnnulus, x, y)
    f1 = elliptical_overlap_exact(x - 0.5, y - 0.5, x + 0.5, y + 0.5, ap.a_out, ap.b_out, ap.theta)
    f2 = elliptical_overlap_exact(x - 0.5, y - 0.5, x + 0.5, y + 0.5, ap.a_in, ap.b_in, ap.theta)
    return f1 - f2
end

function partial(sub_ap::Subpixel{T,<:EllipticalAnnulus}, x, y) where T
    ap = sub_ap.ap
    coeffs_out = oblique_coefficients(ap.a_out, ap.b_out, ap.theta)
    f1 = elliptical_overlap_single_subpixel(x - 0.5, y - 0.5, x + 0.5, y + 0.5, coeffs_out..., sub_ap.N)
    coeffs_in = oblique_coefficients(ap.a_in, ap.b_in, ap.theta)
    f2 = elliptical_overlap_single_subpixel(x - 0.5, y - 0.5, x + 0.5, y + 0.5, coeffs_in..., sub_ap.N)
    return f1 - f2
end


#=
Tools for working with Markov Chains

@author : Spencer Lyon, Zac Cranko

@date: 07/10/2014

References
----------

https://lectures.quantecon.org/jl/finite_markov.html

=#
import LightGraphs: DiGraph, period, attracting_components,
                    strongly_connected_components, is_strongly_connected

@inline check_stochastic_matrix(P) = maximum(abs, sum(P, dims = 2) .- 1) < 5e-15 ? true : false

"""
Finite-state discrete-time Markov chain.

Methods are available that provide useful information such as the stationary
distributions, and communication and recurrent classes, and allow simulation of
state transitions.

##### Fields

- `p::AbstractMatrix` : The transition matrix. Must be square, all elements must be nonnegative, and all rows must sum to unity.
- `state_values::AbstractVector` : Vector containing the values associated with the states.
"""
mutable struct MarkovChain{T, TM<:AbstractMatrix{T}, TV<:AbstractVector}
    p::TM # valid stochastic matrix
    state_values::TV

    function MarkovChain{T,TM,TV}(p::AbstractMatrix, state_values) where {T,TM,TV}
        n, m = size(p)

        n != m &&
            throw(DimensionMismatch("stochastic matrix must be square"))

        minimum(p) <0 &&
            throw(ArgumentError("stochastic matrix must have nonnegative elements"))

        !check_stochastic_matrix(p) &&
            throw(ArgumentError("stochastic matrix rows must sum to 1"))

        length(state_values) != n &&
            throw(DimensionMismatch("state_values should have $n elements"))

        return new{T,TM,TV}(p, state_values)
    end
end

# Provide constructor that infers T from eltype of matrix
MarkovChain(p::AbstractMatrix, state_values=1:size(p, 1)) =
    MarkovChain{eltype(p), typeof(p), typeof(state_values)}(p, state_values)

Base.eltype(mc::MarkovChain{T,TM,TV}) where {T,TM,TV} = eltype(TV)

"Number of states in the Markov chain `mc`"
n_states(mc::MarkovChain) = size(mc.p, 1)

function Base.show(io::IO, mc::MarkovChain{T,TM}) where {T,TM}
    println(io, "Discrete Markov Chain")
    println(io, "stochastic matrix of type $TM:")
    print(io, mc.p)
end

@doc doc"""
This routine computes the stationary distribution of an irreducible Markov
transition matrix (stochastic matrix) or transition rate matrix (generator
matrix) ``A``.

More generally, given a Metzler matrix (square matrix whose off-diagonal
entries are all nonnegative) ``A``, this routine solves for a nonzero solution
``x`` to ``x (A - D) = 0``, where ``D`` is the diagonal matrix for which the rows of
``A - D`` sum to zero (i.e., ``D_{ii} = \sum_j A_{ij}`` for all ``i``). One (and only
one, up to normalization) nonzero solution exists corresponding to each
reccurent class of ``A``, and in particular, if ``A`` is irreducible, there is a
unique solution; when there are more than one solution, the routine returns the
solution that contains in its support the first index ``i`` such that no path
connects ``i`` to any index larger than ``i``. The solution is normalized so that
its 1-norm equals one. This routine implements the Grassmann-Taksar-Heyman
(GTH) algorithm (Grassmann, Taksar, and Heyman 1985), a numerically stable
variant of Gaussian elimination, where only the off-diagonal entries of ``A`` are
used as the input data. For a nice exposition of the algorithm, see Stewart
(2009), Chapter 10.

##### Arguments

- `A::Matrix{T}` : Stochastic matrix or generator matrix. Must be of shape n x
  n.

##### Returns

- `x::Vector{T}` : Stationary distribution of ``A``.

##### References

- W. K. Grassmann, M. I. Taksar and D. P. Heyman, "Regenerative Analysis and
  Steady State Distributions for Markov Chains, " Operations Research (1985),
  1107-1116.
- W. J. Stewart, Probability, Markov Chains, Queues, and Simulation, Princeton
  University Press, 2009.

"""
gth_solve(A::Matrix{T}) where {T<:Real} = gth_solve!(copy(A))

# convert already makes a copy, hence gth_solve!
gth_solve(A::Matrix{T}) where {T<:Integer} =
    gth_solve!(convert(Matrix{Rational{T}}, A))

"""
Same as `gth_solve`, but overwrite the input `A`, instead of creating a copy.

"""
function gth_solve!(A::Matrix{T}) where T<:Real
    n = size(A, 1)
    x = zeros(T, n)

    @inbounds for k in 1:n-1
        scale = sum(A[k, k+1:n])
        if scale <= zero(T)
            # There is one (and only one) recurrent class contained in
            # {1, ..., k};
            # compute the solution associated with that recurrent class.
            n = k
            break
        end
        A[k+1:n, k] /= scale

        for j in k+1:n, i in k+1:n
            A[i, j] += A[i, k] * A[k, j]
        end
    end

    # backsubstitution
    x[n] = 1
    @inbounds for k in n-1:-1:1, i in k+1:n
        x[k] += x[i] * A[i, k]
    end

    # normalisation
    x /= sum(x)

    return x
end

"""
Find the recurrent classes of the Markov chain `mc`.

##### Arguments

- `mc::MarkovChain` : MarkovChain instance.

##### Returns

- `::Vector{Vector{Int}}` : Vector of vectors that describe the recurrent
  classes of `mc`.

"""
recurrent_classes(mc::MarkovChain) = attracting_components(DiGraph(mc.p))

"""
Find the communication classes of the Markov chain `mc`.

#### Arguments

- `mc::MarkovChain` : MarkovChain instance.

### Returns

- `::Vector{Vector{Int}}` : Vector of vectors that describe the communication
  classes of `mc`.

"""
communication_classes(mc::MarkovChain) = strongly_connected_components(DiGraph(mc.p))

"""
Indicate whether the Markov chain `mc` is irreducible.

##### Arguments

- `mc::MarkovChain` : MarkovChain instance.

##### Returns

- `::Bool`

"""
is_irreducible(mc::MarkovChain) =  is_strongly_connected(DiGraph(mc.p))

"""
Indicate whether the Markov chain `mc` is aperiodic.

##### Arguments

- `mc::MarkovChain` : MarkovChain instance.

##### Returns

- `::Bool`

"""
is_aperiodic(mc::MarkovChain) = period(mc) == 1

"""
Return the period of the Markov chain `mc`.

##### Arguments

- `mc::MarkovChain` : MarkovChain instance.

##### Returns

- `::Int` : Period of `mc`.

"""
function period(mc::MarkovChain)
    g = DiGraph(mc.p)
    recurrent = attracting_components(g)

    d = 1
    for r in recurrent
        pd = period(g[r])
        d *= div(pd, gcd(pd, d))
    end

    return d
end

for (S, ex_T, ex_gth) in ((Real, :(T), :(gth_solve!)),
                          (Integer, :(Rational{T}), :(gth_solve)))
    @eval function stationary_distributions(mc::MarkovChain{T}) where T<:$S
        n = n_states(mc)
        rec_classes = recurrent_classes(mc)
        T1 = $ex_T
        stationary_dists = Vector{Vector{T1}}(undef, length(rec_classes))

        for (i, rec_class) in enumerate(rec_classes)
            dist = zeros(T1, n)
            # * mc.p[rec_class, rec_class] is a copy, hence gth_solve!
            #   while gth_solve for Int matrix
            # * todense is to convert a sparse matrix to a dense matrix
            #   with eltype T1
            A = todense(T1, mc.p[rec_class, rec_class])
            dist[rec_class] = $ex_gth(A)
            stationary_dists[i] = dist
        end

        return stationary_dists
    end
end

@doc doc"""
Compute stationary distributions of the Markov chain `mc`, one for each
recurrent class.

##### Arguments

- `mc::MarkovChain{T}` : MarkovChain instance.

##### Returns

- `stationary_dists::Vector{Vector{T1}}` : Vector of vectors that represent
  stationary distributions, where the element type `T1` is `Rational` if `T` is
  `Int` (and equal to `T` otherwise).

"""
stationary_distributions

# From base/sparse/sparsematrix.jl
"""Custom version of `full`, which allows convertion to type `T`"""
function todense(T::Type, S::SparseMatrixCSC)
    A = zeros(T, S.m, S.n)
    for Sj in 1:S.n
        for Sk in nzrange(S, Sj)
            Si = S.rowval[Sk]
            Sv = S.nzval[Sk]
            A[Si, Sj] = Sv
        end
    end
    return A
end

"""If `A` is already dense, return `A` as is"""
todense(::Type, A::Array) = A

mutable struct MCIndSimulator{T<:MarkovChain,S}
    mc::T
    len::Int
    init::Int
    drvs::S
end

function MCIndSimulator(mc::MarkovChain, len::Int, init::Int)
    # NOTE: ensure dense array and transpose before slicing the array. Then
    #       when passing to DiscreteRV use `sub` to avoid allocating again
    p = Matrix(mc.p)'
    drvs = [DiscreteRV(view(p, :, i)) for i in 1:size(mc.p, 1)]
    MCIndSimulator(mc, len, init, drvs)
end

# Base.start(mcis::MCIndSimulator) = (mcis.init, 0)
# function Base.next(mcis::MCIndSimulator, state::Tuple{Int,Int})
#     ix, t = state
#     (ix, (rand(mcis.drvs[ix]), t+1))
# end
# Base.done(mcis::MCIndSimulator, s::Tuple{Int,Int}) = s[2] >= mcis.len

function Base.iterate(mcis::MCIndSimulator, state::Tuple{Int,Int}=(mcis.init, 0))
    ix, t = state
    if t >= mcis.len
        return nothing
    end
    (ix, (rand(mcis.drvs[ix]), t+1))
end

Base.length(mcis::MCIndSimulator) = mcis.len
Base.eltype(mcis::MCIndSimulator) = Int
Base.IteratorSize(mcis::MCIndSimulator) = Base.HasLength()

mutable struct MCSimulator{T<:MCIndSimulator}
    mcis::T
end

function MCSimulator(mc::MarkovChain, len::Int, init::Int)
    MCSimulator(MCIndSimulator(mc, len, init))
end

# only need to implement next and eltype differently...
function Base.iterate(mcs::MCSimulator, state::Tuple{Int,Int}=(mcs.mcis.init, 0))
    output = iterate(mcs.mcis, state)
    if output === nothing
        return output
    end
    ix, new_state = output
    (mcs.mcis.mc.state_values[ix], new_state)
end
Base.eltype(mcs::MCSimulator) = eltype(mcs.mcis.mc)

# ...the rest of the interface can derive from mcis
Base.length(mcs::MCSimulator) = length(mcs.mcis)
Base.IteratorSize(mcs::MCSimulator) = Base.IteratorSize(mcs.mcis)

"""
Simulate one sample path of the Markov chain `mc`.
The resulting vector has the state values of `mc` as elements.

### Arguments

- `mc::MarkovChain` : MarkovChain instance.
- `ts_length::Int` : Length of simulation
- `;init::Int=rand(1:n_states(mc))` : Initial state

### Returns

- `X::Vector` : Vector containing the sample path, with length
  ts_length
"""
function simulate(mc::MarkovChain, ts_length::Int;
                  init::Int=rand(1:n_states(mc)))
    X = Vector{eltype(mc)}(undef, ts_length)
    simulate!(X, mc; init=init)
end

"""
Fill `X` with sample paths of the Markov chain `mc` as columns.
The resulting matrix has the state values of `mc` as elements.

### Arguments

- `X::Matrix` : Preallocated matrix to be filled with sample paths
of the Markov chain `mc`. The element types in `X` should be the
same as the type of the state values of `mc`
- `mc::MarkovChain` : MarkovChain instance.
- `;init=rand(1:n_states(mc))` : Can be one of the following
    - blank: random initial condition for each chain
    - scalar: same initial condition for each chain
    - vector: cycle through the elements, applying each as an
      initial condition until all columns have an initial condition
      (allows for more columns than initial conditions)
"""
function simulate!(X::Union{AbstractVector,AbstractMatrix},
                   mc::MarkovChain; init=rand(1:n_states(mc), size(X, 2)))
    mcs = MCSimulator(mc, size(X, 1), init[1])

    for (i, init) in enumerate(take(cycle(init), size(X, 2)))
        mcs.mcis.init = init
        copyto!(view(X, :, i), mcs)
    end
    X
end

# ------------------------ #
# simulate_indices methods #
# ------------------------ #

"""
Simulate one sample path of the Markov chain `mc`.
The resulting vector has the indices of the state values of `mc` as elements.

### Arguments

- `mc::MarkovChain` : MarkovChain instance.
- `ts_length::Int` : Length of simulation
- `;init::Int=rand(1:n_states(mc))` : Initial state

### Returns

- `X::Vector{Int}` : Vector containing the sample path, with length
  ts_length
"""
function simulate_indices(mc::MarkovChain, ts_length::Int;
                          init::Int=rand(1:n_states(mc)))
    X = Vector{Int}(undef, ts_length)
    simulate_indices!(X, mc; init=init)
end

"""
Fill `X` with sample paths of the Markov chain `mc` as columns.
The resulting matrix has the indices of the state values of `mc` as elements.

### Arguments

- `X::Matrix{Int}` : Preallocated matrix to be filled with indices
of the sample paths of the Markov chain `mc`.
- `mc::MarkovChain` : MarkovChain instance.
- `;init=rand(1:n_states(mc))` : Can be one of the following
    - blank: random initial condition for each chain
    - scalar: same initial condition for each chain
    - vector: cycle through the elements, applying each as an
      initial condition until all columns have an initial condition
      (allows for more columns than initial conditions)
"""
function simulate_indices!(X::Union{AbstractVector{T},AbstractMatrix{T}},
               mc::MarkovChain; init=rand(1:n_states(mc), size(X, 2))) where T<:Integer
    mcis = MCIndSimulator(mc, size(X, 1), init[1])

    for (i, init) in enumerate(take(cycle(init), size(X, 2)))
        mcis.init = init
        copyto!(view(X, :, i), mcis)
    end
    X
end

@doc doc"""
Simulate one sample path of the non-stationary vector of Markov chains `mcs`.
The resulting vector has the state values of `mcs` as elements.

### Arguments

- `mcs::Vector{MarkovChain}` : Vector of MarkovChains.
- `;init::Int` : Index of initial state

### Returns

- `X::Vector` : Vector containing the sample path, with length
  length(mcs)+1 (includes initial state)
"""
function simulate_ns(mcs::Vector, init::Int)
    ind_0 = init
    ind = zeros(Int, length(mcs))
    X = zeros(Real, length(mcs))

    for t in eachindex(mcs)
        X[t] = simulate(mcs[t], 2, init = ind_0)[2]
        ind[t] = findmin( abs.(mcs[t].state_values .- X[t]) )[2]
        if t < length(mcs)
            ind_0 = ind[t] #findmin( abs.(mcs[t].state_values .- X[t]) )[2]
        end
    end
    ind = [init; ind]
    X = [mcs[1].state_values[init]; X]
end


# irreducible unitary representations of the groups SO(3) and SE(3) and representations of
# their Lie algebras.
# For SO(3), see
# Varshalovich, D. A., Moskalev, A. N. & Khersonskii, V. K. Quantum Theory of Angular
# Momentum. 1988. World Scientific, Singapore. ISBN: 9971-5-0107-4.
# For SO(3) and SE(3), see §12.9 and 12.12 of
# Chirikjian, G. S. Stochastic Models, Information Theory, and Lie Groups, Volume 2. 2012.
# ISBN: 978-0-8176-4943-2. doi: 10.1007/978-0-8176-4944-9.

# diagonals of infinitesimal rotation elements of Varshalovich §4.18.3.
function representation_diag(::typeof(so3), ::typeof(E1), ℓ, ::Val{i}) where {i}
    ℓ2 = Int(2ℓ)
    absi = abs(i)
    isone(absi) && return @. im * (-sqrt((ℓ2:-1:1) * (1:ℓ2)) / 2)
    return zeros(complex(float(eltype(ℓ))), ℓ2 + 1 - absi)
end
function representation_diag(::typeof(so3), ::typeof(E2), ℓ, ::Val{i}) where {i}
    ℓ2 = Int(2ℓ)
    absi = abs(i)
    isone(absi) && return @. i * sqrt((ℓ2:-1:1) * (1:ℓ2)) / 2
    return zeros(float(eltype(ℓ)), ℓ2 + 1 - absi)
end
function representation_diag(::typeof(so3), ::typeof(E3), ℓ, ::Val{i}) where {i}
    iszero(i) && return Vector(im .* (ℓ:-1:(-ℓ)))
    return zeros(complex(eltype(ℓ)), Int(2ℓ) + 1 - abs(i))
end

"""
    representation_block(::typeof(so3), X, ℓ)

Compute the size `(2ℓ+1, 2ℓ+1)` block ``ℓ`` of the representation of the element ``X`` of
the element of the Lie algebra ``𝔰𝔬(3)``.

The representation is defined as

```math
u^ℓ(X) = \\lim_{t → 0} \\frac{d}{dt} U^ℓ(\\exp(t X)),
```

where ``U^ℓ(R)`` is the representation of the rotation ``R``, and ``\\exp(R)`` is the group
exponential map on ``\\mathrm{SO}(3)``.
"""
representation_block(::typeof(so3), X, ℓ)

function representation_block(::typeof(so3), Ei::BasisVector, ℓ)
    dl = representation_diag(so3, Ei, ℓ, Val(-1))
    d = representation_diag(so3, Ei, ℓ, Val(0))
    du = representation_diag(so3, Ei, ℓ, Val(1))
    return Tridiagonal(dl, d, du)
end
function representation_block(::typeof(so3), ::typeof(E3), ℓ)
    return Diagonal(representation_diag(so3, E3, ℓ, Val(0)))
end
function representation_block(::typeof(so3), Xⁱ::AbstractVector{<:Real}, ℓ)
    # because representation_element returns a result with real eltype Float64, we use `T`
    # to avoid promoting if the eltype of Xⁱ is lower precision
    T = complex(float(eltype(Xⁱ)))
    uE1 = representation_block(so3, E1, ℓ)
    uE2 = representation_block(so3, E2, ℓ)
    uE3 = representation_block(so3, E3, ℓ)
    return Tridiagonal(
        Xⁱ[1] .* T.(uE1.dl) .+ Xⁱ[2] .* T.(uE2.dl),
        Xⁱ[3] .* T.(uE3.diag),
        Xⁱ[1] .* T.(uE1.du) .+ Xⁱ[2] .* T.(uE2.du),
    )
end
function representation_block(::typeof(so3), X::AbstractMatrix, ℓ)
    # represent the vector X as an actual Vector using the basis
    Xⁱ = Manifolds.vee(SO3, Manifolds.Identity(SO3, one(X)), X)
    return representation_block(so3, Xⁱ, ℓ)
end

"""
    representation_block(::typeof(SO3), p, ℓ) -> Matrix{Complex}

Compute the size `(2ℓ+1, 2ℓ+1)` block ``ℓ`` of the representation of the rotation ``p``.

This block is identical to the ``D^ℓ`` block of the Wigner ``D``-matrix using the
conventions of [^Varshalovich1988].

[^Varshalovich1988]: Varshalovich, D. A., Moskalev, A. N. & Khersonskii, V. K. Quantum Theory of Angular
    Momentum. 1988. World Scientific, Singapore. ISBN: 9971-5-0107-4.
# Examples

To create the "little-``d``" block ``d^ℓ(\\frac{π}{4}) = D^ℓ(0, \\frac{π}{4}, 0)``,

```julia-repl
julia> using EnsembleInference, Rotations

julia> EnsembleInference.representation_block(SO3, RotZXZ(0, π/2, 0), 1)
3×3 Matrix{ComplexF64}:
  0.5+0.0im              0.0-0.707107im  -0.5+0.0im
  0.0-0.707107im  1.2175e-16+0.0im        0.0-0.707107im
 -0.5+0.0im              0.0-0.707107im   0.5+0.0im

julia> EnsembleInference.representation_block(SO3, RotZXZ(0, π/2, 0), 2)
5×5 Matrix{ComplexF64}:
      0.25+0.0im   0.0-0.5im          -0.612372+0.0im           0.0+0.5im               0.25+0.0im
       0.0-0.5im  -0.5-0.0im                0.0-3.24998e-16im  -0.5-0.0im                0.0+0.5im
 -0.612372-0.0im   0.0-2.41689e-16im       -0.5-0.0im           0.0-3.62534e-16im  -0.612372-0.0im
       0.0+0.5im  -0.5+0.0im                0.0-2.03875e-16im  -0.5+0.0im                0.0-0.5im
      0.25+0.0im   0.0+0.5im          -0.612372+0.0im           0.0-0.5im               0.25+0.0im
```
"""
function representation_block(::typeof(SO3), p, ℓ)
    X = Manifolds.group_log(SO3, p)
    u = representation_block(so3, X, ℓ)
    return exp(Matrix(u))
end


# File: fVAI.jl
# Author:
# Date: 29-June-2020

# VALUE ADDED INDEX (VAI)
function fVAI(x)

# --------------------------------------------------------------------------
# x:       returns vector
# vaidx:   Value Added Index (VAI)
# --------------------------------------------------------------------------

    n = size(x,1); # data points
    fVAI = ones(1,n); # pre-allocate memory()
    fVAI(1) = 100; # initialise to $100
    for i = 1:i:end
        fVAI(i+1) = fVAI(i) * (1+x[i]./100); # calculate VAI
    end
    vaidx = fVAI;
    return vaidx
end


# Functional roles

A = [
     false true false false false false;
     false false false true false false;
     false true false false false false;
     false false false false true false;
     false false false false false true;
     false false false false false false
    ]
U = UnipartiteNetwork(A)
L = [1,1,1,2,2,2]
P = Partition(U, L)
n = networkroles(P)

# Values of z
@test n[1,1] ≈ -0.57 atol = 0.02
@test n[2,1] ≈  1.15 atol = 0.02
@test n[3,1] ≈ -0.57 atol = 0.02
@test n[4,1] ≈ -0.57 atol = 0.02
@test n[5,1] ≈  1.15 atol = 0.02
@test n[6,1] ≈ -0.57 atol = 0.02

# Values of c
@test n[1,2] ≈ 0.00 atol = 0.02
@test n[2,2] ≈ 0.45 atol = 0.02
@test n[3,2] ≈ 0.00 atol = 0.02
@test n[4,2] ≈ 0.50 atol = 0.02
@test n[5,2] ≈ 0.00 atol = 0.02
@test n[6,2] ≈ 0.00 atol = 0.02


import Base: size, getindex, setindex!, similar, ndims
#' in the future: Base.show(io::IO, tenmat::Tenmat) = print(io, "...")

abstract AbstractTenmat{T} # <: AbstractArray
##### functions to be implemented to make it sub of AbstractArray
# AbstractArray{T,N} where T is the type of object returned by integer indexing (A[1,...,1], when A is not empty)
# 	and N should be the length of the tuple returned by size()

"""
	Tenmat{T}(matrix::Array{T,2}, rowindex::Vector{Int}, colindex::Vector{Int}, tensorsize::Tuple{Vararg{Int}})

Matrix which stores the indices to convert it back to tensor.

## Example

```jldoctest
matrix = Array(reshape(1:48, 6, 8)) # some 6 x 8 matrix
rowindex = [3,1]
colindex = [2,4]
tensorsize = (2,4,3,2)
tenmat = Tenmat(matrix, rowindex, colindex, tensorsize)

# output

TensorMatrices_lemon.Tenmat{Int64}([1 7 … 37 43; 2 8 … 38 44; … ; 5 11 … 41 47; 6 12 … 42 48],[3,1],[2,4],(2,4,3,2))
```
"""
type Tenmat{T} <: AbstractTenmat{T}
	# e.g. Tenmat(A,[1,2],[3,4],(4,3,2,4)), where A is 12x8 matrix
	"converted matrix from tensor"
	matrix::Array{T,2}
	"Tensor indices to map to row"
	rowindex::Vector{Int}
	"Tensor indices to map to coloumn"
	colindex::Vector{Int}
	tensorsize::Tuple{Vararg{Int}} # tensor size

	function Tenmat{T}(matrix::Array{T,2}, rowindex::Vector{Int}, colindex::Vector{Int}, tensorsize::Tuple{Vararg{Int}})
		# the condition for the arguments:
		# the dimension of the mat matches to the tensize
		# -> the product of tensize[rowindex] == mat's row dimension
		# -> also for the column
		(length(rowindex) + length(colindex) == length(tensorsize)) &&
		(productWithIndex(tensorsize, rowindex) == size(matrix, 1)) &&
		(productWithIndex(tensorsize, colindex) == size(matrix, 2)) ?
 		new(matrix, rowindex, colindex, tensorsize)   :
		error("dimension does not match")
	end
end

Tenmat{T}(matrix::Matrix{T}, rowindex::Vector{Int}, colindex::Vector{Int}, tensorsize::Tuple{Vararg{Int}})  = Tenmat{T}(matrix, rowindex, colindex, tensorsize)

"""
	Tenmat{T}(matrix::Matrix{T})

Set a matrix as a Tenmat.
"""
Tenmat{T}(matrix::Matrix{T}) = Tenmat{T}(matrix, [1], [2], size(matrix))

"""
	Tenmat{T,N}(tensor::Array{T,N}, rowindex::Vector{Int}, colindex::Vector{Int})

Set a tensor as a Tenmat.

`Tenmat.matrix` will return the converted matrix.

## Example

```jldoctest basicMethods
A = Array(reshape(1:3024, 4,7,2,9,6)) # tensor, size of [4,7,2,9,6]
rowindex = [5,2] # 5th and 2nd indices maps to the row
colindex = [3,4,1] # 3rd, 4th and 1st indices maps to the column
tenmat = Tenmat(A,rowindex,colindex)

# output

TensorMatrices_lemon.Tenmat{Int64}([1 29 … 452 480; 505 533 … 956 984; … ; 2041 2069 … 2492 2520; 2545 2573 … 2996 3024],[5,2],[3,4,1],(4,7,2,9,6))
```
"""
Tenmat{T,N}(tensor::Array{T,N}, rowindex::Vector{Int}, colindex::Vector{Int}) = tensor2tenmat(tensor, rowindex, colindex)

##### define functions with macro
# tenmat_with_other_args_non_preserving = [:size, :getindex, : setindex!

"""
	size(tenmat::Tenmat, args...)

Return the size of the tenmat.matrix.

```jldoctest basicMethods
julia> size(tenmat)
(42,72)

julia> size(tenmat,2)
72
```
"""
size(tenmat::Tenmat, args...) = size(tenmat.matrix, args...)

"""
	setindex!(tenmat::Tenmat, args...)

Expands `setindex!` by tenmat.matrix.

```jldoctest basicMethods
julia> setindex!(tenmat, 100, 2, 6);
```
"""
setindex!(tenmat::Tenmat, args...) = setindex!(tenmat.matrix, args...)

"""
	getindex(tenmat::Tenmat, args...)

Expands `getindex` by tenmat.matrix.

```jldoctest basicMethods
julia> getindex(tenmat, 2, 6)
100
```
"""
getindex(tenmat::Tenmat, args...) = getindex(tenmat.matrix, args...)

"""
	similar

similar

## Examples

```jldoctest
A = rand(4,7,6)
tenmat = Tenmat(A,[3,1],[2])
similar(tenmat, Int)
similar(tenmat, Int, (5,2,7))
similar(tenmat, Int, [1,3],[2])
similar(tenmat, Int, [4,2,5], [1,3], (2,4,3,1,6))
similar(tenmat)
similar(tenmat, (5,2,7))
similar(tenmat, [1,3],[2])
B = similar(tenmat, [4,2,5], [1,3], (2,4,3,1,6))
size(B)

# output

(24,6)

```
"""
similar{T}(tenmat::Tenmat{T}) = Tenmat{T}(similar(tenmat.matrix), tenmat.rowindex, tenmat.colindex, tenmat.tensorsize)
similar{T}(tenmat::Tenmat{T}, S::Type) =
	similar(tenmat, S, tensorSize(tenmat))
similar{N}(tenmat::Tenmat, T::Type, tensorsize::Tuple{Vararg{Int64,N}}) =
	similar(tenmat, T, tenmat.rowindex, tenmat.colindex, tensorsize)
similar{T}(tenmat::Tenmat{T}, S::Type, rowindex::Vector{Int}, colindex::Vector{Int}) =
	similar(tenmat, S, rowindex, colindex, tensorSize(tenmat))
similar{T}(tenmat::Tenmat{T}, S::Type, rowindex::Vector{Int}, colindex::Vector{Int}, tensorsize::Tuple{Vararg{Int}}) =
	Tenmat(similar(tenmat2tensor(tenmat), T, tensorsize), rowindex, colindex)

similar{N}(tenmat::Tenmat, tensorsize::Tuple{Vararg{Int64,N}}) =
	similar(tenmat, tenmat.rowindex, tenmat.colindex, tensorsize)
similar{T}(tenmat::Tenmat{T}, rowindex::Vector{Int}, colindex::Vector{Int}) = similar(tenmat, rowindex, colindex, tensorSize(tenmat))
similar{T}(tenmat::Tenmat{T}, rowindex::Vector{Int}, colindex::Vector{Int}, tensorsize::Tuple{Vararg{Int}}) =
	Tenmat(similar(tenmat2tensor(tenmat), tensorsize), rowindex, colindex)

"""
	ndims(tenmat::Tenmat)

Return the original dimension of tenmat,
from `tensorsize`.

```jldoctest tenmat2
A = Array(reshape(1:3024, 4,7,2,9,6))
tenmat = Tenmat(A,[5,2],[3,4,1])
ndims(tenmat)

# output

5
```
"""
ndims(tenmat::Tenmat) = length(tensorSize(tenmat))

"""
	tensorSize(tenmat::Tenmat)

Return the tensorsize.

```jldoctest tenmat2
julia> tensorSize(tenmat)
(4,7,2,9,6)
```
"""
function tensorSize(tenmat::Tenmat)
	return tenmat.tensorsize
end

"""
	getMatrix(tenmat::Tenmat)

Return the matrix.

```jldoctest
A = Array(reshape(1:18, 3, 2, 3))
tenmat = Tenmat(A,[3,1],[2])
getMatrix(tenmat)

# output

9×2 Array{Int64,2}:
  1   4
  7  10
 13  16
  2   5
  8  11
 14  17
  3   6
  9  12
 15  18
```
"""
function getMatrix{T}(tenmat::Tenmat{T})
	return tenmat.matrix
end

"""
	productWithIndex{T}(tuple::Tuple{Vararg{T}}, index = collect(1:length(tuple))::Vector{Int})

Get a product of tuple elements with index number in index vector
"""
function productWithIndex{T}(tuple::Tuple{Vararg{T}}, index = collect(1:length(tuple))::Vector{Int})
	prod = 1
	for i in index
		prod *= tuple[i]
	end
	return prod
end

"""
	tensor2tenmat{T,N}(tensor::Array{T,N}, rowindex::Vector{Int}, colindex::Vector{Int})

Convert tensor to tenmat.

## examples

```jldoctest
julia> A = rand(3,5,3,2,6);

julia> tenmatA = tensor2tenmat(A,[3,5,2],[4,1]);

julia> size(getMatrix(tenmatA))
(90,6)

julia> tensorSize(tenmatA)
(3,5,3,2,6)

julia> B = Array(reshape(1:10,2,5));

julia> tenmatB = tensor2tenmat(B,[2],[1]);

julia> tenmatB.matrix == transpose(B)
true

julia> tensorSize(tenmatB)
(2,5)

```
"""
function tensor2tenmat{T,N}(tensor::Array{T,N}, rowindex::Vector{Int}, colindex::Vector{Int})
#=
julia> a = 5; b = 7; c = 8;
julia> D = reshape(1:a*b*c,a,b,c)
julia> tenmatD = tensor2tenmat(D,[3],[1,2])
julia> tenmatD.tensorsize == (a,b,c)
julia> size(tenmatD.matrix) == (c,a*b)
julia> for i = 1:a, j = 1:b, k = 1:c
julia> 	println(D[i,j,k] == tenmatD.matrix[k,(j-1)*a+i])
julia> end
julia>
julia> a = 5; b = 7; c = 8; d = 11
julia> D = reshape(1:a*b*c*d,a,b,c,d)
julia> tenmatD = tensor2tenmat(D,[3],[1,4,2])
julia> tenmatD.tensorsize == (a,b,c,d)
julia> size(tenmatD.matrix) == (c,a*d*b)
julia> for i = 1:a, j = 1:b, k = 1:c, l = 1:d
julia> 	println(D[i,j,k,l] == tenmatD.matrix[k,(j-1)*a*d+(l-1)*a+i])
julia> end

=#
	tensorsize = size(tensor)
	(N == length(rowindex) + length(colindex)) || error("index does not match")

	permten = permutedims(tensor,vcat(rowindex,colindex))
	rowdim = productWithIndex(size(tensor),rowindex)
	coldim = productWithIndex(size(tensor),colindex)

	matrix = reshape(permten,rowdim,coldim)
	tenmat = Tenmat(matrix,rowindex,colindex,tensorsize)

	return tenmat
end

"""
	tenmat2tensor{T}(tenmat::Tenmat{T})

Convert tenmat back to the tensor.

## examples
```jldoctest
julia> A = rand(3,5,3,2,8);

julia> tenmatA = tensor2tenmat(A,[3,5,2],[4,1]);

julia> back2A = tenmat2tensor(tenmatA);

julia> A == back2A
true
```

```jldoctest
B = rand(8,6,7,5,8,5,3,8)
tenmatB = tensor2tenmat(B,[1,7,3,2,6],[5,8,4])
back2B = tenmat2tensor(tenmatB)
check = true
for a=1:8, b=1:6, c=1:7, d=1:5, e=1:8, f=1:5, g=1:3, h=1:8
	back2B[a,b,c,d,e,f,g,h] == B[a,b,c,d,e,f,g,h] || check * false
end
check

# output

true
```
"""
function tenmat2tensor{T}(tenmat::Tenmat{T})
	indices = vcat(tenmat.rowindex, tenmat.colindex)
	tensor = reshape(tenmat.matrix, [tenmat.tensorsize[i] for i in indices]...)
	tensor = ipermutedims(tensor, indices)
	return tensor
end


using Memoize

@memoize function pattern(i::Int)
    vcat(zeros(Int,i), ones(Int,i), zeros(Int,i), fill(-1, i))
end

function pattern(position::Int, element::Int)
    p = pattern(position)
    return p[(element % length(p)) + 1]
end

function fft(input::String, phases::Int)
    m = [pattern(r, c) for r=1:length(input),c=1:length(input)]
    v = parse.(Int,split(input,""))

    for i=1:phases
        v = abs.(rem.(m * v, 10))
    end

    return v
end

examples = ["80871224585914546619083218645595", "19617804207202209144916044189917",
    "69317163492948606335995924319873"]

println("Day 16 Examples: ")
for example in examples
    println("$(example) becomes $(foldl(*, map(string, fft(example,100)[1:8])))")
end

input_file = open("input.txt")
input = readline(input_file)
close(input_file)

println("Day 16 Part 1: $(foldl(*, map(string, fft(input,100)[1:8])))")

# I never really understood the problem statement for part 2. From Reddit,
# it looks like I was not the only one.
# The real signal is the puzzle input repeated 10000 times.
signal = input ^ 10000

# The message offset is the first seven digits.
offset = parse(Int, signal[1:7])


##
cd(@__DIR__)
using DelimitedFiles
input = readdlm("input.txt", Int) |> vec;

##
count(>(0),diff(input))

## PARTE 2
using DataStructures: nil, cons

##
res = foldl(1:length(input)-2; init=nil()) do l, w
    cons(sum(input[w:w+2]), l)
end;

##
count(>(0), diff(collect(reverse(res))))


struct Singleton{T}
  value::T
end

Base.show(io::IO, x::Singleton) = print(io, "Singleton(", x.value, ")")

collapse(x::Bernoulli) =
  mean(x) == 1.0 ? Singleton(true) :
  mean(x) == 0.0 ? Singleton(false) :
  x

# TODO make this a distribution, approximate logpdf etc

struct Empirical{T}
  samples::Vector{T}
end

@forward Empirical.samples Statistics.mean, Statistics.var, Statistics.std, Base.length

rounded(x) = @sprintf("%.2g", x)

function Base.show(io::IO, d::Empirical{Bool})
  print(io, "Empirical Bernoulli(p=$(rounded(mean(d)))), N=$(length(d))")
end

function Base.show(io::IO, d::Empirical{<:Real})
  print(io, "Empirical(mean = $(rounded(mean(d))), std = $(rounded(std(d))), N=$(length(d)))")
end


function solve_lp()
    glpk_optimizer =
        optimizer_with_attributes(GLPK.Optimizer, "msg_lev" => GLPK.GLP_MSG_OFF)
    model = Model(optimizer_with_attributes(
        CS.Optimizer,
        "lp_optimizer" => glpk_optimizer,
        "logging" => [],
        "seed" => 1,
    ))

    # Variables
    @variable(model, inclusion[h = 1:3], Bin)
    @variable(model, 0 <= allocations[h = 1:3, a = 1:3] <= 1, Int)
    @variable(model, 0 <= days[h = 1:3, a = 1:3] <= 5, Int)

    # Constraints
    @constraint(
        model,
        must_include[h = 1:3],
        sum(allocations[h, a] for a in 1:3) <= inclusion[h]
    )
    # at least n
    @constraint(model, min_hospitals, sum(inclusion[h] for h in 1:3) >= 3)
    # every h must be allocated at most one a
    @constraint(model, must_visit[h = 1:3], sum(allocations[h, a] for a in 1:3) <= 1)
    # every allocated h must have fewer than 5 days of visits per week
    @constraint(
        model,
        max_visits[h = 1:3],
        sum(days[h, a] for a in 1:3) <= 5 * inclusion[h]
    )

    @objective(model, Max, sum(days[h, a] * 5 for h in 1:3, a in 1:3))
    optimize!(model)
end


"""
    Code for a Gaussian mask shape for use with CCF
Author: Eric Ford
Created: August 2020
"""

import SpecialFunctions: erf

"""   GaussianCCFMask
A truncated Gaussian mask with two parameters, it's stdandard deviation and where to truncate it, both as a velocity in m/s.
Mask weights are stored separately in a line list.
"""
struct GaussianCCFMask <: AbstractCCFMaskShape
    σ_sqrt2::Float64
    half_width_truncation::Float64
    normalization::Float64
    cdf_normalization::Float64

    """ GaussianCCFMask( σ ; half_truncation_width_in_σ=2 ) """
    function GaussianCCFMask(σ::Real, w::Real=2 )
        @assert 0 < σ <= 300000   # 300km/s is arbitrary choice for an upper limit
        @assert 0 < w <= 4
        norm = 1.0/(sqrt(2π)*σ*erf(w/(sqrt(2.0))))
        cdf_norm = 0.5/(erf(w/(sqrt(2.0))))
        new(σ*sqrt(2.0),σ*w,norm,cdf_norm)
    end

end

""" `GaussianCCFMask( inst ; scale_factor )` """
function GaussianCCFMask(inst::AbstractInstrument; σ_scale_factor::Real = 1, truncation_scale_factor::Real = default_gaussian_ccf_truncation_scale_factor )
    #=
    default_σ = 5000
    σ_sqrt2 = scale_factor * default_σ * sqrt(2.0)
    w = scale_factor * default_ccf_mask_v_width(inst)
    norm = 1.0/(σ_sqrt2*sqrt(π)*erf(w/σ_sqrt2))
    GaussianCCFMask(σ_sqrt2,w/2) # ,norm)
    =#
    σ = σ_scale_factor * default_ccf_mask_v_width(inst)
    w = truncation_scale_factor
    GaussianCCFMask(σ,w/2)
end

λ_min(m::GaussianCCFMask,λ::Real) = λ/calc_doppler_factor(m.half_width_truncation)
λ_max(m::GaussianCCFMask,λ::Real) = λ*calc_doppler_factor(m.half_width_truncation)

function integrate(m::GaussianCCFMask, v_lo::Real,v_hi::Real)
    #quadgk(m, v_lo, v_hi)[1]
    m.cdf_normalization*(erf(v_hi/m.σ_sqrt2)-erf(v_lo/m.σ_sqrt2))
end

""" Functor for returning PSF for Δv <= half_width.  """
function (m::GaussianCCFMask)(Δv::Real)
    if abs2(Δv) > abs2(m.half_width_truncation)
        return zero(Δv)
    else
        return m.normalization*exp(-(Δv/m.σ_sqrt2)^2)
    end
end

function mask_with_increased_fwhm(m::GaussianCCFMask, Δfwhm::Real )
    σ = m.σ_sqrt2/sqrt(2)
    w = m.half_width_truncation/σ
    #fwhm_orig = sqrt(8*log(2)) * σ
    fwhm_orig = 2*sqrt(log(2)) * m.σ_sqrt2
    fwhm_new = sqrt(fwhm_orig^2 + Δfwhm^2)
    σ_new = 0.5*fwhm_new/sqrt(2*log(2))
    return GaussianCCFMask(σ_new, w)
end


using Zygote, LinearAlgebra
using IterativeSolvers: lobpcg

function test_eigmin(A::AbstractMatrix)
    results = lobpcg(A, false, 1)
    ev = minimum(results.λ)
    return ev
end

A = rand(4, 4); A = A + A';

@show eigmin(A)
@show test_eigmin(A)
@show Zygote.gradient(test_eigmin, A)


module Miletus

using Dates, LinearAlgebra
using Dates: days
using Statistics: mean

export CurrencyUnit, CurrencyQuantity, USD, GBP,
 value, process, ivol, fit, fit_ivol,
 SingleStock, CoreModel, CoreForwardModel, GeomBMModel,
 montecarlo, MonteCarloModel, LeastSquaresMonteCarlo

import Optim
import StatsBase: fit
import StatsFuns: normcdf, norminvcdf, normpdf, invsqrt2, sqrthalfπ, invsqrt2π, logistic, logit
import Base.Math: @horner
import ForwardDiff: Dual

extract_derivative(x::Dual) = ForwardDiff.partials(x,1)
extract_derivative(x::Real) = zero(x)

import Reexport: @reexport
@reexport using BusinessDays

const Φ = normcdf
const invΦ = norminvcdf
const ϕ = normpdf

include("daycounts.jl")
include("termstructure.jl")
using .TermStructure

include("currency.jl")
using .Currency

include("utils/math.jl")
include("utils/black.jl")
include("utils/ivol.jl")

include("observables.jl")
include("contracts.jl")
include("process.jl")

include("models/abstractmodel.jl")
include("models/core.jl")
include("models/coreforward.jl")
include("models/geombm.jl")
include("models/binomial.jl")
include("models/binomial_ivol.jl")

include("models/montecarlo.jl")
include("models/lsmc.jl")

include("print.jl")

include("greeks.jl")
export vega, delta, rho, greek

end # module


import ACE.Transforms: GetVal

@doc raw"""
`Ylm1pBasis`

One-particle basis component of the form
```math
\phi_{lm}({\bm r}) = Y_l^m(\hat{\br r})
```
Fundamental building block of ACE basis sets of the form
```math
   R_{nl}^{\mu_i, \mu_j}(r_{ij}) Y_l^m(\hat{\bm r})
```
This type basically just translates the `SHBasis` into a valid one-particle
basis.
"""
function Ylm1pBasis(L::Integer; varsym = :rr, lsym = :l, msym = :m,
                                  label = "Y$lsym$msym")
   SH = SphericalHarmonics.SHBasis(L)
   NTLM = NamedTuple{(lsym, msym)}
   spec = [ NTLM(tuple(SphericalHarmonics.idx2lm(i)...))
            for i = 1:length(SH) ]
   degrees = [ b[lsym] for b in spec ]
   return B1pComponent(SH, GetVal{varsym}(), spec, degrees, label)
end


using Distributions, DataFrames

const θ_dist = Uniform(-1, 1)
const ϕ_dist = Uniform(0, π)

struct ChemicalShift
    σᵢₛₒ::Array{Distribution}
    Δσ::Array{Distribution}
    ησ::Array{Distribution}
    weights::Array{Float64}
end

function ChemicalShift(p::Array{Float64})
    sites = length(p) ÷ 7
    σᵢₛₒ = Array{Distribution}(undef, sites)
    Δσ = Array{Distribution}(undef, sites)
    ησ = Array{Distribution}(undef, sites)
    weights = zeros(sites)

    for i in 1:length(p)
        i % 7 in [0, 2, 4, 5, 6] && p[i] < 0 && return missing
        (isnan(p[i]) || isinf(p[i])) && return missing
        (isnothing(p[i]) || ismissing(p[i])) && return missing
    end

    for i = 1:sites
        σᵢₛₒ[i] = Normal(p[7*(i-1)+1], p[7*(i-1)+2])
        Δσ[i] = Normal(p[7*(i-1)+3], p[7*(i-1)+4])
        ησ[i] = truncated(Normal(p[7*(i-1)+5], p[7*(i-1)+6]), 0.0, 1.0)
        weights[i] = p[7*i]
    end
    weights /= sum(weights)

    return ChemicalShift(σᵢₛₒ, Δσ, ησ, weights)
end

function get_ν(
    θ::Float64,
    ϕ::Float64,
    σᵢₛₒ::Distribution,
    Δσ::Distribution,
    ησ::Distribution,
)
    rand(σᵢₛₒ) + (rand(Δσ) / 2) * (3 * θ^2 - 1 + rand(ησ) * (1-θ^2) * cos(2 * ϕ));
end

function estimate_powder_pattern(p::ChemicalShift, N::Int64)
    powder_pattern = zeros(N)
    i = 1
    for j = 1:(length(p.weights) - 1)
        to_add = floor(Int, p.weights[j] * N)
        θ = rand(θ_dist, to_add)
        ϕ = rand(ϕ_dist, to_add)
        powder_pattern[i:(i + to_add - 1)] = get_ν.(θ, ϕ, p.σᵢₛₒ[j], p.Δσ[j], p.ησ[j])
        i += to_add
    end
    θ = rand(θ_dist, N - i + 1)
    ϕ = rand(ϕ_dist, N - i + 1)
    powder_pattern[i:end] = get_ν.(θ, ϕ, p.σᵢₛₒ[end], p.Δσ[end], p.ησ[end])
    return powder_pattern
end

function get_chemical_shift_starting_values(sites::Int64)
    starting_values = zeros(7 * sites)
    starting_values[1:7:end] = rand(Uniform(-4000, 4000), sites)
    starting_values[2:7:end] = rand(Uniform(0.000001, sqrt(800)), sites)
    starting_values[3:7:end] = rand(Uniform(-4000, 4000), sites)
    starting_values[4:7:end] = rand(Uniform(0.000001, sqrt(400)), sites)
    starting_values[5:7:end] = rand(Uniform(0, 1), sites)
    starting_values[6:7:end] = rand(Uniform(0.000001, 1), sites)
    starting_values[7:7:end] = map(x -> 1 / sites, 1:sites)
    return starting_values
end

function get_output_table(minimizer::Array{Float64})
    minimizer = transform_params(push!(minimizer,1.0), ChemicalShift)
    df = DataFrame(["σᵢₛₒ" minimizer[1] minimizer[2]
                    "Δσ" minimizer[3] minimizer[4]
                    "ησ"  minimizer[5]  minimizer[6]])
    rename!(df, Symbol.(["Parameter","Mean","St. Dev."]))
    return df
end


import MeasureTheory.Kernel, MeasureTheory.kernel

MeasureTheory.mapcall(t) = map(func -> func(), t)
(k::Kernel{M,<:NamedTuple})() where {M} = M(; MeasureTheory.mapcall(k.ops)...)

const GaussKernel = MeasureTheory.Kernel{<:Gaussian}
#const Copy = MeasureTheory.Kernel{<:Dirac}

"""
    AffineMap(B, β)

Represents a function `f = AffineMap(B, β)`
such that `f(x) == B*x + β`.
"""
struct AffineMap{S,T}
    B::S
    β::T
end
(a::AffineMap)(x) = a.B*x + a.β
(a::AffineMap)(p::Gaussian) = Gaussian(μ = a.B*mean(p) + a.β, Σ = a.B*cov(p)*a.B')
(a::AffineMap)(p::WGaussian) = WGaussian(μ = a.B*mean(p) + a.β, Σ = a.B*cov(p)*a.B', c=p.c)

"""
    LinearMap(B)

Represents a function `f = LinearMap(B)`
such that `f(x) == B*x`.
"""
struct LinearMap{T}
    B::T
end
(a::LinearMap)(x) = a.B*x
(a::LinearMap)(p::Gaussian) = Gaussian(μ = a.B*mean(p), Σ = a.B*cov(p)*a.B')
(a::LinearMap)(p::WGaussian) = WGaussian(μ = a.B*mean(p), Σ = a.B*cov(p)*a.B', c=p.c)

"""
    ConstantMap(β)

Represents a function `f = ConstantMap(β)`
such that `f(x) == β`.
"""
struct ConstantMap{T}
    x::T
end
(a::ConstantMap)(x) = a.x
(a::ConstantMap)() = a.x

const LinearGaussianKernel = Kernel{T,NamedTuple{(:μ, :Σ),Tuple{A, C}}} where {T<:Gaussian, A<:LinearMap, C<:ConstantMap}
params(k::LinearGaussianKernel) = k.ops.μ.B, Zero(), k.ops.Σ.x
const AffineGaussianKernel = Kernel{T,NamedTuple{(:μ, :Σ),Tuple{A, C}}} where {T<:Gaussian, A<:AffineMap, C<:ConstantMap}
params(k::AffineGaussianKernel) = k.ops.μ.B, k.ops.μ.β, k.ops.Σ.x
const ConstantGaussianKernel =  Kernel{T,NamedTuple{(:μ, :Σ),Tuple{A, C}}} where {T<:Gaussian, A<:ConstantMap, C<:ConstantMap}
params(k::ConstantGaussianKernel) = k.ops.μ.x, k.ops.Σ.x

"""
    correct(prior, obskernel, obs) = u, yres, S

Joseph form correction step of a Kalman filter with `prior` state
and `obs` the observation with observation kernel
`obskernel = kernel(Gaussian; μ=LinearMap(H), Σ=ConstantMap(R))`
`H` is the observation operator and `R` the observation covariance. Returns corrected/conditional
distribution `u`, the residual and the innovation covariance.
See https://en.wikipedia.org/wiki/Kalman_filter#Update.
"""
function correct(u::Gaussian{T}, k::LinearGaussianKernel, y) where {T}
    x, Ppred = meancov(u)
    H = k.ops.μ.B
    R = k.ops.Σ.x
    (x, P), yres, S = correct_joseph((x, Ppred), H, R, y)
    Gaussian{T}(μ=x, Σ=P), yres, S
end

function correct_joseph((x, Ppred), H, R, y)
    yres = y - H*x # innovation residual
    S = (H*Ppred*H' + R) # innovation covariance

    K = Ppred*H'/S # Kalman gain
    x = x + K*yres
    P = (I - K*H)*Ppred*(I - K*H)' + K*R*K' #  Joseph form
    (x, P), yres, S
end

function correct_kalman((x, Ppred), H, R, y)
    yres = y - H*x # innovation residual
    S = (H*Ppred*H' + R) # innovation covariance

    K = Ppred*H'/S # Kalman gain
    x = x + K*yres
    P = (I - K*H)*Ppred # Kalman form
    (x, P), yres, S
end

function smooth(Gs::T, Gf, Gpred, Φ) where {T}
    xs, Ps = meancov(Gs)
    xf, Pf = meancov(Gf)
    xpred, Ppred = meancov(Gpred)

    J = Pf*Φ'/Ppred # C/(C+w)
    xs = xf + J*(xs - xpred)
    Ps = Pf + J*(Ps - Ppred)*J'

    (xs, Ps)
end


# This file was generated, do not modify it. # hide
using DelimitedFiles

function cnt(ReZ, ImZ, ReC, ImC)
    k::UInt8 = 0
    while (k < 255)
        new_ReZ = ReZ * ReZ - ImZ * ImZ + ReC
        new_ImZ = 2 * ReZ * ImZ + ImC
        if new_ReZ * new_ReZ + new_ImZ * new_ImZ > 4.
            return k
        end
        ReZ = new_ReZ
        ImZ = new_ImZ
        k += 1
    end
    return k
end

function mand(M,N)
    grid = Array{UInt8}(undef, M, N)
    init_Re = 0.0
    init_Im = 0.0
    for j = 1:N
        for i = 1:M
            ReC = 4.0 / (M - 1) * (j - 1) - 2.0
            ImC = 4.0 / (N - 1) * (i - 1) - 2.0
            ReZ = init_Re
            ImZ = init_Im
            grid[i, j] = cnt(ReZ, ImZ, ReC, ImC)
        end
    end
    return grid
end

function main()
    s = time()
    M = N = 4200
    grid = mand(M,N)
    t = time()
    @show(t - s)
    writedlm("/tmp/resultJ.txt",grid,",")
end

main()

using BenchmarkTools
@btime mand(5000, 5000);


"""
    covfn(x1, x2, gp_params)
Calculate GP covariance between x1 and x2 based on a squared exponential covariance function

Returns GP covariance matrix
"""
function covfn(x1, x2, gp_params::SVGP_params)
    rho = exp.(gp_params.log_rho)
    R   = exp.(-Distances.pairwise(SqEuclidean(), x1 ./ rho, x2 ./ rho, dims=1))

    return exp(gp_params.log_kappa[1]) * R
end

# function covfn(x1, x2, params::SVGP_params)
#     rho = exp.(params.log_rho)
#     x2    = x2 ./ rho
#     R     = exp.(-pairwise(haversine(), x1', x2')) # WORK ON BEING ABLE TO USE THIS KERNEL FOR EARTH SURFACE
#     return exp(params.log_kappa[1]) * R .+ (2.0 * params.mean_factor)^2
# end

"""
    covfn_lonlat(x1, x2, gp_params)
Calculate GP covariance between x1 and x2 based on the Haversine distance on a sphere

Warning: Currently does not work due to lack of auto-differentiation through Haversine distance call

Returns GP covariance matrix
"""
function covfn_lonlat(x1, x2, gp_params::SVGP_params)
    rho = exp.(gp_params.log_rho)
    R   = exp.(-Distances.pairwise(Haversine(), x1 ./ rho, x2 ./ rho, dims=1))

    return exp(gp_params.log_kappa[1]) * R
end


#--------------------------------------------------------------------
# Spacetime Discretization methods in Julia
# Soham 01-2018
#--------------------------------------------------------------------

function restrictmodes!(coeffs::Array{Float64,1}, M::Int)::Array{Float64,1}
    if M+1 > size(coeffs)[1]
        error("Cannot restrict to larger number of modes")
    else
        for index in CartesianRange(size(coeffs))
            index[1] > M+1 ? coeffs[index] = 0 : coeffs[index] = coeffs[index]
        end
    end
    return coeffs
end

function restrictmodes!(coeffs::Array{Float64,2}, M::Int, N::Int)::Array{Float64,2}
    if M+1 > size(coeffs)[1] || N+1 > size(coeffs)[2]
        error("Cannot restrict to larger number of modes")
    else
        for index in CartesianRange(size(coeffs))
            index[1] > M + 1 || index[2] > N + 1 ? coeffs[index] = 0 : coeffs[index] = coeffs[index]
        end
    end
    return coeffs
end

function prolongatemodes(coeffs::Array{Float64,1}, M::Int)::Array{Float64,1}
    if M+1 < size(coeffs)[1]
        error("Cannot prolongate to smaller  number of modes")
    else
        ncoeffs = zeros(M+1)
        oldM = size(coeffs)[1] - 1
        for index in CartesianRange(size(ncoeffs))
            index[1] < oldM+1 ? ncoeffs[index] = coeffs[index] : ncoeffs[index] = ncoeffs[index]
        end
    end
    return ncoeffs
end

function prolongatemodes(coeffs::Array{Float64,2}, M::Int, N::Int)::Array{Float64,2}
    if M+1 < size(coeffs)[1] || N+1 < size(coeffs)[2]
        error("Cannot prolongate to smaller number of modes")
    else
        ncoeffs = zeros(M+1,N+1)
        oldM = size(coeffs)[1] - 1
        oldN = size(coeffs)[2] - 1
        for index in CartesianRange(size(ncoeffs))
            index[1] < oldM+1 && index[2] < oldN+1 ? ncoeffs[index] = coeffs[index] : ncoeffs[index] = ncoeffs[index]
        end
    end
    return ncoeffs
end

function prolongateOP(Nx::Int, M::Int)::Array{Float64,2}
    invndmx = inv(vandermonde(Nx))
    Px = zeros((Nx+1)*M, Nx+1)
    for k in 1:M
        pvndmx = pseudovandermonde(Nx, chebgrid(Nx, M, k))
        Px[1+(k-1)*(Nx+1):k*(Nx+1), :] = pvndmx*invndmx
    end
    return Px
end

function restrictOP(Nx::Int, M::Int)::Array{Float64,2}
    return pinv(prolongateOP(Nx, M))
end

function prolongatePatch(patch::Patch, M::Int)::Dict{Array{Int, 1}, Patch}
    Nx = size(patch.value)[1] - 1
    Ny = size(patch.value)[2] - 1
    Px = prolongateOP(Nx, M)
    Py = prolongateOP(Ny, M)
    mPatch = Px*patch.value*Py'
    return array2dict(mPatch, Nx, Ny, M)
end

function restrictPatch(dbase::Dict{Array{Int, 1}, Patch})::Patch
    Nx = size(dbase[[1,1]].value)[1] - 1
    Ny = size(dbase[[1,1]].value)[2] - 1
    M  = convert(Int, sqrt(length(dbase)))
    Rx = restrictOP(Nx, M)
    Ry = restrictOP(Ny, M)
    sPatch = dict2array(dbase)
    return Patch([1,1], Rx*sPatch*Ry')
end


module Kabsch
export calc_centroid, kabsch, rotate, rmsd, translate_points, kabsch_rmsd
	# Calculate root mean square deviation of two matrices A, B
	# http://en.wikipedia.org/wiki/Root-mean-square_deviation_of_atomic_positions
	function rmsd(A::Array{Float64,2}, B::Array{Float64,2})

		RMSD::Float64 = 0.0

		# D pairs of equivalent atoms
		D::Int = size(A)[1]::Int # <- oddly _only_ changing this to Int makes it work on 32-bit systems.
		# N coordinates
		N::Int = length(A)::Int

		for i::Int64 = 1:N
			RMSD += (A[i]::Float64 - B[i]::Float64)^2
		end
		return sqrt(RMSD / D)
	end

	# calculate a centroid of a matrix
	function calc_centroid(m::Array{Float64,2})

		sum_m::Array{Float64,2} = sum(m,1)
		size_m::Int64 = size(m)[1]

		return map(x -> x/size_m, sum_m)
	end

	# Translate P, Q so centroids are equal to the origin of the coordinate system
	# Translation der Massenzentren, so dass beide Zentren im Ursprung des Koordinatensystems liegen
	function translate_points(P::Array{Float64,2}, Q::Array{Float64,2})
		# Calculate centroids P, Q
		# Die Massenzentren der Proteine
		centroid_p::Array{Float64,2} = calc_centroid(P)
		centroid_q::Array{Float64,2} = calc_centroid(Q)

		P = broadcast(-,P, centroid_p)

		Q = broadcast(-,Q, centroid_q)

		return P, Q, centroid_p, centroid_q
	end

	# Input: Two sets of points: reference, coords as Nx3 Matrices (so)
	# returns optimally rotated matrix
	function kabsch(reference::Array{Float64,2}, coords::Array{Float64,2})

		centered_reference::Array{Float64,2}, centered_coords::Array{Float64,2}, centroid_p::Array{Float64,2}, centroid_q::Array{Float64,2}  = translate_points(reference, coords)
		# Compute covariance matrix A
		A::Array{Float64,2} = *(centered_coords', centered_reference)

		# Calculate Singular Value Decomposition (SVD) of A
		u::Array{Float64,2}, d::Array{Float64,1}, vt::Array{Float64,2} = svd(A)

		# check for reflection
		f::Int64 = sign(det(vt) * det(u))
		m::Array{Int64,2} = [1 0 0; 0 1 0; 0 0 f]

		# Calculate the optimal rotation matrix _and_ superimpose it
		return broadcast(+, *(centered_coords, u, m, vt'), centroid_p)

	end

	# directly return RMSD for matrices P, Q for convenience
	function kabsch_rmsd(P::Array{Float64,2}, Q::Array{Float64,2})
		return rmsd(P,kabsch(P,Q))
	end

end


import Base.in

export SetConstraint
export in

### Set constraint
mutable struct SetConstraint{T <: ContiSet} <: Constraint
  head::Symbol
  id_hash::UInt64
  lhs::AbstractExpr
  rhs::T
  size::Tuple{Int, Int}
  dual::ValOrNothing
end

function SetConstraint(lhs::Variable, rhs::ContiSet)
  id_hash = hash((lhs, rhs, :(==)))
  return SetConstraint(:(set), id_hash, lhs, rhs, (1, 1), nothing)
end

function get_cv(constraint::SetConstraint)
  return get_cv(constraint.lhs)
end

function get_sv(constraint::SetConstraint)
  return get_sv(constraint.lhs)
end

in(lhs::Variable, rhs::ContiSet) = SetConstraint(lhs, rhs)

# the fenchel conjugate of a set indicator function
fenchel(setconstraint::SetConstraint, w::AbstractArray) = fenchel(setconstraint.rhs, w)

# the prox of a set indicator function
prox(setconstraint::SetConstraint, w::AbstractArray) = prox(setconstraint.rhs, w)


# * Envelopes
#   The envelopes implemented below are all /vector potential/
#   envelopes, since that is what is being used in
#   calculations. However, they may be specified using
#   intensity-related quantities, e.g. Gaussian pulses are most often
#   specified using the FWHM duration of their /intensity/ envelopes.
#

envelope_types = Dict{Symbol,Any}()

complex_squared(t) = t*t

# ** Gaussian
@doc raw"""
    GaussianEnvelope

A Gaussian pulse is given by

```math
f(t) = \exp\left(-\frac{t^2}{2\sigma^2}\right),
```

where the standard deviation ``σ`` is related to the FWHM duration τ
of the intensity envelope as

```math
\sigma = \frac{\tau}{2\sqrt{2\ln 2}}.
```

Since we define all fields in terms of the vector potential ``A(t)``,
and the [`instantaneous_intensity`](@ref) is given by
``\abs{-\partial_t A(t)}^2``, we have to find a coefficient ``α`` such
that the field amplitude

```math
F(t) \sim -\partial_t \exp(-\alpha t^2) \sin(\omega t + \phi)
```

has an _intensity envelope_ with the desired FWHM; we do this
iteratively in [`gaussian_common!`](@ref). This is mainly important
for ultrashort pulses, since for long pulse durations we can
approximate ``\exp(-\alpha t^2) \sim 1`` such that only the carrier
contributes to the time derivative.

Since a Gaussian never ends, we specify how many ``σ`` we require; the
resulting time window will be rounded up to an integer amount of
cycles of the fundamental.

Required parameters:
- `λ|T|f|ν|ω|ħω`,
- `τ|σ`,
- `σmax|tmax|Tmax`,
- `env=:gauss` (optional, since [`GaussianEnvelope`](@ref) is the default envelope).
"""
struct GaussianEnvelope{T} <: AbstractEnvelope
    τ::T # Intensity FWHM
    σ::T # Intensity std.dev.
    α::T # Vector potential exponential coefficient
    σmax::T
    tmax::T # Maximum time. Time window: [-tmax,tmax]
end
envelope_types[:gauss] = GaussianEnvelope

(env::GaussianEnvelope)(t) = exp(-env.α*complex_squared(t))

show(io::IO, env::GaussianEnvelope) =
    printfmt(io, "Gaussian envelope of duration {1:s} (intensity FWHM; ±{2:0.2f}σ)",
             au2si_round(env.τ, u"s"), env.σmax)

@doc raw"""
    gaussian_common!(field_params, carrier[; Tmax_rounder, verbosity])

Compute parameters common to Gaussian envelopes,
i.e. [`GaussianEnvelope`](@ref) and
[`TruncatedGaussianEnvelope`](@ref); most importantly, given an
intensity FWHM or σ, we need to figure out the coefficient ``\alpha``
for the envelope of the vector potential ``\exp(-\alpha t^2)``, such
that the electric field amplitude and intensity have the desired
durations. The optional function `Tmax_rounder` determines if the
envelop should be extended to encompass an integer amount of cycles of
the `carrier` (default).
"""
function gaussian_common!(field_params, carrier;
                          Tmax_rounder = Base.Fix1(ceil, Int), verbosity=0)
    @namespace!(field_params) do
        if τ
            σ = τ/(2*√(2log(2)))
        else
            τ = 2*√(2log(2))*σ
        end

        if σmax
            Tmax = Tmax_rounder(σmax*σ/T)
            tmax = Tmax*T
        else
            if tmax
                Tmax = Tmax_rounder(tmax/T)
            elseif Tmax
                tmax = Tmax*T
            end
            tmax = Tmax*T
        end
        σmax = tmax/σ
    end
    τ = austrip(field_params[:τ])
    σ = austrip(field_params[:σ])
    α = 2log(2)/τ^2
    period = austrip(field_params[:T])
    ω = austrip(field_params[:ω])

    σmax = austrip(field_params[:σmax])
    tmax = austrip(field_params[:tmax])

    τ < 0.5period &&
        @warn "Pulse durations smaller than half a period not reliably supported"

    # Find the exponential coefficient for the vector potential that
    # will give the desired FWHM for the intensity profile.
    f = α -> begin
        env = GaussianEnvelope(τ, σ, α[1], σmax, tmax)
        field = make_temp_field(carrier, env, field_params)
        abs2(intensity(field, τ/2) - ω^2/2)
    end
    verbosity > 0 && @info "Finding optimal vector potential coefficient"
    res = optimize(f, [α], BFGS())
    verbosity > 0 && display(res)

    field_params[:α] = res.minimizer[1]
end

function GaussianEnvelope(field_params::Dict{Symbol,Any}, carrier)
    test_field_parameters(field_params, [:T]) # Period time required to round time window up
    test_field_parameters(field_params, [:τ, :σ])
    test_field_parameters(field_params, [:σmax, :tmax, :Tmax])

    gaussian_common!(field_params, carrier)

    @unpack τ, σ, α, σmax, tmax = field_params
    GaussianEnvelope(austrip(τ), austrip(σ), α, austrip(σmax), austrip(tmax))
end

continuity(::GaussianEnvelope) = Inf
span(env::GaussianEnvelope) = -env.tmax..env.tmax

time_integral(env::GaussianEnvelope) = env.σ*√(2π)

# *** Spectrum
@doc raw"""
    spectrum(env::GaussianEnvelope)

Gaussians belong to the [Schwartz
class](https://en.wikipedia.org/wiki/Schwartz_space), i.e. functions
who, under Fourier transform, are mapped back to the same space. That
is to say, the Fourier transform of a Gaussian is a Gaussian:

```math
\exp(-\alpha t^2) \leftrightarrow
\frac{1}{\sqrt{2\alpha}}
\exp\left(-\frac{\omega^2}{4\alpha}\right).
```

Comparing with the above, we find that the spectral standard
deviation

```math
\Omega = \sqrt{2\alpha} = \frac{2\sqrt{\ln 2}}{\tau},
```

and the Gaussian function in the spectral domain is thus

```math
E(\omega) =
\frac{E_0\tau}{2\sqrt{\ln 2}}
\exp\left[-\frac{(\omega\tau)^2}{8\ln2}\right].
```
"""
function spectrum(env::GaussianEnvelope)
    N = env.E₀*env.τ/(2*√(log(2)))
    ω -> N*exp(-(ω*env.τ)^2/8log(2))
end

# ** Truncated Gaussian

@doc raw"""
    TruncatedGaussianEnvelope

Since a Gaussian function never ends, this envelope adds a soft truncation over a time interval:

```math
f(t)=
\begin{cases}
\exp(-\alpha t^2), & \abs{t} \leq t_{\textrm{off}},\\
\exp\left\{
-\alpha\left[
t_{\textrm{off}} + \frac{2}{\pi}(t_{\textrm{max}} - t_{\textrm{off}})
\tan\left(
\frac{\pi}{2}
\frac{\abs{t} - t_{\textrm{off}}}{t_{\textrm{max}} - t_{\textrm{off}}}
\right)
\right]^2
\right\},
& t_{\textrm{off}} < \abs{t} \leq t_{\textrm{max}}, \\
0, & t_{\textrm{max}} < \abs{t}.
\end{cases}
```
This is Eq. (72) of

- Patchkovskii, S., & Muller, H. (2016). Simple, accurate, and
  efficient implementation of 1-electron atomic time-dependent
  Schrödinger equation in spherical coordinates. Computer Physics
  Communications, 199,
  153–169. [10.1016/j.cpc.2015.10.014](http://dx.doi.org/10.1016/j.cpc.2015.10.014)

Required parameters:
- `λ|T|f|ν|ω|ħω`,
- `τ|σ`,
- `toff`,
- `σmax|tmax|Tmax`,
`env=:trunc_gauss`.

Beyond this, everything else is the same as for
[`GaussianEnvelope`](@ref).
"""
struct TruncatedGaussianEnvelope{T} <: AbstractEnvelope
    τ::T # Intensity FWHM
    σ::T # Intensity std.dev.
    α::T # Vector potential exponential coefficent
    toff::T # Start of hard turn-off
    tmax::T # Maximum time; end of hard turn-off. Time window: [-tmax,tmax]
    Tmax::T # Maximum time, in cycles of the fundamental.
end
envelope_types[:trunc_gauss] = TruncatedGaussianEnvelope

function (env::TruncatedGaussianEnvelope{T})(t) where T
    at = abs(real(t))
    at > env.tmax && return zero(T)
    @unpack α, toff, tmax = env
    if at ≤ env.toff
        exp(-α*complex_squared(t))
    else
        exp(-α*complex_squared(toff + 2/π*(tmax-toff)*tan(π/2*(at-toff)/(tmax-toff))))
    end
end

show(io::IO, env::TruncatedGaussianEnvelope) =
    printfmt(io, "Truncated Gaussian envelope of duration {1:.4f} jiffies = {2:s} (intensity FWHM; turn-off from {3:s} to {4:s})",
             env.τ, au2si_round(env.τ, u"s"),
             au2si_round(env.toff, u"s"), au2si_round(env.tmax, u"s"))

function TruncatedGaussianEnvelope(field_params::Dict{Symbol,Any}, carrier)
    test_field_parameters(field_params, [:T]) # Period time required to round time window up
    test_field_parameters(field_params, [:τ, :σ])
    test_field_parameters(field_params, [:toff, :σoff])
    test_field_parameters(field_params, [:σmax, :tmax, :Tmax])

    gaussian_common!(field_params, carrier, Tmax_rounder=NoUnits)
    @namespace!(field_params) do
        if σoff
            toff = σ*σoff
        end
    end

    @unpack τ, σ, α, toff, tmax, Tmax = field_params
    toff < tmax || throw(ArgumentError("Hard turn-off must occur before end of pulse"))
    TruncatedGaussianEnvelope(austrip(τ), austrip(σ), α, austrip(toff), austrip(tmax), Tmax)
end

continuity(::TruncatedGaussianEnvelope) = Inf # This is not exactly true
span(env::TruncatedGaussianEnvelope) = -env.tmax..env.tmax

# TODO: Take truncation into account
time_integral(env::TruncatedGaussianEnvelope) = env.σ*√(2π)

# ** Trapezoidal

@doc raw"""
    TrapezoidalEnvelope

This is a very simple piecewise linear function:

```math
f(t)=
\begin{cases}
r/r_{\textrm{up}},
& 0 \leq r < r_{\textrm{up}},\\
1,
& r_{\textrm{up}} \leq r < r_{\textrm{up}} + r_{\textrm{flat}}, \\
1 - \frac{r-r_{\textrm{up}}-r_{\textrm{flat}}}{r_{\textrm{down}}},
& r_{\textrm{up}} + r_{\textrm{flat}} \leq r < r_{\textrm{up}} + r_{\textrm{flat}} + r_{\textrm{down}}, \\
0,
& \textrm{else},
\end{cases}
\quad
r \defd t/T,
```
where ``T`` is the period time of the carrier.

Required parameters:
- `λ|T|f|ν|ω|ħω`,
- `ramp | (ramp_up & ramp_down)`,
- `flat`,
- `env=:trapezoidal | env=:tophat`.

Beware that this envelope can introduce artifacts at the ends of the
pulse, such that the electric field is non-vanishing, depending of
e.g. the phase of the carrier.
"""
struct TrapezoidalEnvelope{T} <: AbstractEnvelope
    ramp_up::T
    flat::T
    ramp_down::T
    period::T

    function TrapezoidalEnvelope(ramp_up::T, flat::T, ramp_down::T, period::T) where T
        ramp_up >= 0 || error("Negative up-ramp not supported")
        flat >= 0 || error("Negative flat region not supported")
        ramp_down >= 0 || error("Negative down-ramp not supported")
        ramp_up + flat + ramp_down > 0 || error("Pulse length must be non-zero")

        new{T}(ramp_up, flat, ramp_down, period)
    end
end
envelope_types[:trapezoidal] = TrapezoidalEnvelope
# Common alias
envelope_types[:tophat] = TrapezoidalEnvelope

function (env::TrapezoidalEnvelope{T})(t) where T
    t = real(t)/env.period
    if t < 0
        zero(T)
    elseif t < env.ramp_up
        t/env.ramp_up
    elseif t < env.ramp_up + env.flat
        one(T)
    elseif t < env.ramp_up + env.flat + env.ramp_down
        one(T) - (t-env.ramp_up-env.flat)/env.ramp_down
    else
        zero(T)
    end
end

show(io::IO, env::TrapezoidalEnvelope) =
    printfmt(io, "/{1:d}‾{2:d}‾{3:d}\\ cycles trapezoidal envelope",
             env.ramp_up, env.flat, env.ramp_down)

function TrapezoidalEnvelope(field_params::Dict{Symbol,Any}, args...)
    test_field_parameters(field_params, [:T]) # Period time required to relate ramps/flat to cycles
    test_field_parameters(field_params, [:ramp, :ramp_up])
    test_field_parameters(field_params, [:ramp, :ramp_down])
    test_field_parameters(field_params, [:flat])

    @namespace!(field_params) do
        if ramp
            ramp_up = ramp
            ramp_down = ramp
        end
    end

    @unpack ramp_up, flat, ramp_down, T = field_params
    TrapezoidalEnvelope(ramp_up, flat, ramp_down, austrip(T))
end

continuity(::TrapezoidalEnvelope) = 0
span(env::TrapezoidalEnvelope{T}) where T =
    zero(T)..((env.ramp_up + env.flat + env.ramp_down)*env.period)

# ** Cos²

@doc raw"""
    Cos²Envelope

```math
f(t) = \begin{cases}
\cos^2 \left(\frac{\pi t}{cT}\right),
& -1 \leq t/(cT) \leq 1,\\
0, & \textrm{else},
\end{cases}
```
where ``c`` is the number of cycles from zero to zero of the
``\cos^2`` envelope and ``T`` the period time.

Required parameters:
- `λ|T|f|ν|ω|ħω`,
- `cycles`,
- `env=:cos² | env=:cos2`.
"""
struct Cos²Envelope{T} <: AbstractEnvelope
    cycles::T
    period::T

    function Cos²Envelope(cycles::T, period::T) where T
        cycles >= 0 || error("Negative duration not supported")
        new{T}(cycles, period)
    end
end
envelope_types[:cos²] = Cos²Envelope
envelope_types[:cos2] = Cos²Envelope

function (env::Cos²Envelope{T})(t) where T
    t /= (env.cycles*env.period)
    if -1 ≤ 2real(t) ≤ 1
        complex_squared(cospi(t))
    else
        zero(T)
    end
end

show(io::IO, env::Cos²Envelope) =
    printfmt(io, "{1:0.2f} cycles cos² envelope", env.cycles)

function Cos²Envelope(field_params::Dict{Symbol,Any}, args...)
    test_field_parameters(field_params, [:T]) # Period time required to relate ramps/flat to cycles
    test_field_parameters(field_params, [:cycles])

    @unpack cycles, T = field_params
    Cos²Envelope(cycles, austrip(T))
end

continuity(::Cos²Envelope) = 0
function span(env::Cos²Envelope)
    s = env.cycles*env.period/2
    -s..s
end

# ** Exports

export continuity


# Suit of functions to create LCModels of Catalecticant matrices.
# For example, for ternary quartics:
#
# Σ = Catalecticant(2, 3)
#
# To compute the ML-degree of Σ
#
# ml_degree_witness(Σ; dual=true)
#

using LinearCovarianceModels
import HomotopyContinuation
const HC = HomotopyContinuation

function nvectors_sumingupm(n, m)
    ntuples = Iterators.product(ntuple(_ -> 0:m, n)...)
    return filter(I -> sum(I) == m, collect(ntuples))
end

function Catalecticantindices(r, s)
    rowindeces = collect.(nvectors_sumingupm(s, r))
    m = binomial(r + s - 1, s - 1)
    M = reduce(hcat, fill(rowindeces, m))
    return M + permutedims(M)
end

function Catalecticantmatrix(r, s)
    return [HC.Variable(:θ, I...) for I in Catalecticantindices(r, s)]
end

Catalecticant = LCModel ∘ Catalecticantmatrix


using LinearAlgebra
using Plots

struct parameters
    gamma::Float64
    x_range::Vector{Float64}
    y_range::Vector{Float64}
    max_tree_nodes::Int64
    nominal_step::Float64
    expansion_time::Float64
    RewireTime::Float64
    obstacle_grid_n::Int64
    cell_size_x::Float64
    cell_size_y::Float64
    function parameters(gamma::Float64, x_range::Vector{Float64}, y_range::Vector{Float64},
        max_tree_nodes::Int64, nominal_step::Float64, expansion_time::Float64,
        RewireTime::Float64, obstacle_grid_n::Int64)
        cell_size_x = diff(x_range)[1] / obstacle_grid_n
        cell_size_y = diff(y_range)[1] / obstacle_grid_n
        new(gamma, x_range, y_range, max_tree_nodes, nominal_step, expansion_time,
            RewireTime, obstacle_grid_n, cell_size_x, cell_size_y)
    end
end

function calc_edge_cost(p1::Vector{Float64}, p2::Vector{Float64})
    "Euclidean distance edge cost"
    return norm(p2[1:2] - p1[1:2])
end

function sample_point(parameters)
    "Sample uniformly in workspace"
    x = parameters.x_range[1] + (parameters.x_range[2] - parameters.x_range[1]) * rand()
    y = parameters.y_range[1] + (parameters.y_range[2] - parameters.y_range[1]) * rand()
    return [x, y]
end

function plot_tree(tree)
    x_data = []
    y_data = []

    for node in tree
        if node.parent_id != 0
            push!(x_data, [tree[node.parent_id].pose[1], node.pose[1]])
            push!(y_data, [tree[node.parent_id].pose[2], node.pose[2]])
        end
    end
    plot!(x_data, y_data, marker=(:circle,2), legend=false, color="grey")
    plot!([tree[1].pose[1]], [tree[1].pose[2]], marker=(:circle,5), legend=false, color="red")
end

function intersected_cells(points::Vector{Vector{Float64}}, params::parameters)
    # Min point p1, max point x2
    min_x = params.x_range[2]
    max_x = params.x_range[1]
    min_y = params.y_range[2]
    max_y = params.y_range[1]

    for pt in points
        pt[1] < min_x ? min_x = pt[1] : nothing
        pt[1] > max_x ? max_x = pt[1] : nothing
        pt[2] < min_y ? min_y = pt[2] : nothing
        pt[2] > max_y ? max_y = pt[2] : nothing
    end

    # Get indices of occ cells in obstacle grid
    id_x_min = 1 + clamp(floor(Int64, (min_x - params.x_range[1]) / params.cell_size_x), 0, params.obstacle_grid_n - 1);
    id_y_min = 1 + clamp(floor(Int64, (min_y - params.y_range[1]) / params.cell_size_y), 0, params.obstacle_grid_n - 1);
    id_x_max = 1 + clamp(floor(Int64, (max_x - params.x_range[1]) / params.cell_size_x), 0, params.obstacle_grid_n - 1);
    id_y_max = 1 + clamp(floor(Int64, (max_y - params.y_range[1]) / params.cell_size_y), 0, params.obstacle_grid_n - 1);

    occ_cells = Vector{Vector{Int64}}()
    for i = id_x_min:id_x_max, j = id_y_min:id_y_max
        push!(occ_cells, [i, j])
    end
    return occ_cells
end