# s = 2b-1
# b = (s+1)/2
# s = -1 <=> b = 0
# s = 1 <=> b = 1
function HOBO(ising::Ising{T}) where T<:Real
    variables = vars(ising)
    hobo = HOBO{T}(variables)
    for (key, val) = values(ising)
        if length(key) == 0
            hobo += HOBO(val, variables)
            continue
        end
        hobo += val*prod(one(T)-2*variables[HOBO{T},k...] for k=key)
    end
    hobo
end

function Ising(hobo::HOBO{T}) where T<:Real
    variables = vars(hobo)
    ising = Ising{T}(variables)
    for (key, val) = values(hobo)
        if length(key) == 0
            ising += Ising(val, variables)
            continue
        end
        ising += val*prod((T(1)-variables[Ising{T},k...])/T(2) for k=key)
    end
    ising
end


using AeroMDAO

## Fuselage definition
rads = range(0, 1, length = 10)
lens = 0.05:0.05:1.0
fuse = Fuselage(lens[1:end-1], [ rads; reverse(rads) ] )

## Property checks
projected_area(fuse)
length(fuse)

## Cosine interpolation
lens_rads = cosine_spacing(fuse, 40)

# Circles for plotting
n_pts          = 20
circle3D(r, n) = let arcs = 0:2π/n:2π; [ zeros(length(arcs)) r * cos.(arcs) r * sin.(arcs) ] end

##
xs = lens_rads[:,1]
circs = [ reduce(hcat, eachrow(circ) .+ Ref([x; 0; 0]))' for (x, circ) in zip(xs, circle3D.(lens_rads[:,2], n_pts)) ]

##
using Plots
pyplot()
plt = plot()
[ plot!(circ[:,1], circ[:,2], circ[:,3], color = :blue, label = :none) for circ in circs ]
plot!()


"""
    asymptotic_monopole_wavenumbers(ω, medium::Acoustic, species; num_wavenumbers = 2)

Calculates the asymptotic effective wavenumbers for monopole scatterers by assuming a large number of wavenumbers. The 2D results are deduced in [Section 5](https://arxiv.org/pdf/1905.06996.pdf).
"""
function asymptotic_monopole_wavenumbers(ω::T, medium::Acoustic{T,2}, species::Species{T,2};
        num_wavenumbers = 2
    ) where T

    P = Int(round(num_wavenumbers / 2) + 1)
    k = ω / medium.c

    b = mean(s1.exclusion_distance * outer_radius(s1) + s2.exclusion_distance * outer_radius(s2) for s1 in species, s2 in species)
    num_density = mean(number_density(s) for s in species)

    T0 = mean(t_matrix(s.particle, medium, ω, 0) for s in species)[1]

    c = sqrt(2pi) * num_density * b^2 * T0 * hankelh1(0,k * b) * exp(-im * pi / 4.0)

    rc = abs(c)
    θc = angle(c)

    kps = map(-ceil(θc / (2pi)):P) do p
        σ = θc + 2pi * p
        kp = (σ + im * log(abs(σ)^(3/2) / rc)) / b
    end

    kns = map(ceil(θc / (2pi) - 3 / 4):P) do p
        σ = θc - 3pi / 2 - 2pi * p
        kn = (σ + im * log(abs(σ)^(3/2) / rc)) / b
    end

   return [kns; kps]
end

function asymptotic_monopole_wavenumbers(ω::T, medium::Acoustic{T,3}, species::Species{T,3};
        num_wavenumbers = 2
    ) where T

   @warn "asymptotic_monopole_wavenumbers has not yet been implemented for 3D acoustics, will use 2D acoustics instead."

   medium2D = Acoustic(2; ρ=medium.ρ, c=medium.c)

   species2D = [
       Specie(Acoustic(2; ρ = s.particle.medium.ρ, c = s.particle.medium.c), Circle(outer_radius(s));
           volume_fraction = volume_fraction(s))
   for s in species]

   return asymptotic_monopole_wavenumbers(ω, medium2D, species2D; num_wavenumbers = num_wavenumbers)
end


module Atmosphere

using StaticArrays, StructArrays, ComponentArrays

using Flight.Geodesy

import Flight.Modeling: SystemDescriptor, System, init_x0, init_y0, init_u0, init_d0, f_cont!, f_disc!

export TunableISA
export TunableWind
export AtmosphereDescriptor, AtmosphericData, AtmosphericSystem

const R = 287.05287 #gas constant for dry air
const γ = 1.40 #heat capacity ratio for dry air
const βs = 1.458e-6 #Sutherland's empirical constant for dynamic viscosity
const S = 110.4 #Sutherland's empirical constant for dynamic viscosity

const T_std = 288.15
const p_std = 101325.0
const ρ_std = p_std / (R * T_std)
const g_std = 9.80665

const ISA_layers = StructArray(
    β =      SVector{7,Float64}([-6.5e-3, 0, 1e-3, 2.8e-3, 0, -2.8e-3, -2e-3]),
    h_ceil = SVector{7,Float64}([11000, 20000, 32000, 47000, 51000, 71000, 84852]))

@inline density(p,T) = p/(R*T)
@inline speed_of_sound(T) = √(γ*R*T)
@inline dynamic_viscosity(T) = (βs * T^1.5) / (T + S)

@inline ISA_temperature_law(h::Real, T_b, h_b, β)::Float64 = T_b + β * (h - h_b)

@inline function ISA_pressure_law(h::Real, g0, p_b, T_b, h_b, β)::Float64
    if β != 0.0
        p_b * (1 + β/T_b * (h - h_b)) ^ (-g0/(β*R))
    else
        p_b * exp(-g0/(R*T_b) * (h - h_b))
    end
end

Base.@kwdef struct SLConditions
    p::Float64 = p_std
    T::Float64 = T_std
    g::Float64 = g_std
end

struct ISAData
    p::Float64
    T::Float64
    ρ::Float64
    a::Float64
    μ::Float64
end

ISAData() = ISAData(AltGeop(0))

@inline function ISAData(h_geo::AltGeop, sl::SLConditions = SLConditions())

    h = Float64(h_geo)
    h_base = 0; T_base = sl.T; p_base = sl.p; g_base = sl.g

    for i in 1:length(ISA_layers)
        β, h_ceil = ISA_layers[i]
        if h < h_ceil
            T = ISA_temperature_law(h, T_base, h_base, β)
            p = ISA_pressure_law(h, g_base, p_base, T_base, h_base, β)
            return ISAData(p, T, density(p, T), speed_of_sound(T), dynamic_viscosity(T) )
        end
        T_ceil = ISA_temperature_law(h_ceil, T_base, h_base, β)
        p_ceil = ISA_pressure_law(h_ceil, g_base, p_base, T_base, h_base, β)
        h_base = h_ceil; T_base = T_ceil; p_base = p_ceil
    end

    throw(ArgumentError("Altitude out of bounds"))

end

##################### AbstractISA ###########################

#a System{<:AbstractISA} may have an output type of its own, as any other
#System. however, this output generally will not be an ISAData instance. the
#ISAData returned by a ISA System depends on the location specified in the
#query. instead, the output from an ISA System may hold quantities of interest
#related to its own internal state, if it has one due to it being a dynamic ISA
#implementation.

#AbstractISA subtypes: ConstantUniformISA (TunableISA), ConstantFieldISA,
#DynamicUniformISA, DynamicFieldISA.

abstract type AbstractISA <: SystemDescriptor end

function SLConditions(::T, ::Abstract2DLocation) where {T<:System{<:AbstractISA}}
    error("SLConditions constructor not implemented for $T")
end

@inline function ISAData(sys::System{<:AbstractISA}, p::Geographic)

    h_geop = Altitude{Geopotential}(p.alt, p.l2d)
    sl = SLConditions(sys, p.l2d)
    ISAData(h_geop, sl)

end

##################### TunableISA ###########################

#a TunableISA System does not have a state and therefore cannot evolve on its
#own. but its input vector can still be used to manually tune the SL conditions
#during simulation.

struct TunableISA <: AbstractISA end #Constant, Uniform ISA

Base.@kwdef mutable struct UTunableISA #only allocates upon System instantiation
    T_sl::Float64 = T_std
    p_sl::Float64 = p_std
end

init_u0(::TunableISA) = UTunableISA()
f_cont!(::System{<:TunableISA}, args...) = nothing
f_disc!(::System{<:TunableISA}, args...) = false

function SLConditions(s::System{<:TunableISA}, ::Abstract2DLocation)
    SLConditions(T = s.u.T_sl, p = s.u.p_sl, g = g_std)
    #alternative using actual local SL gravity:
    # return (T = s.u.T_sl, p = s.u.p_sl, g = gravity(Geographic(l2d, AltOrth(0.0))))
end

######################## AbstractWind ############################

abstract type AbstractWind <: SystemDescriptor end

Base.@kwdef struct WindData
    v_ew_n::SVector{3,Float64} = zeros(SVector{3})
end

function WindData(::T, ::Abstract3DLocation) where {T<:System{<:AbstractWind}}
    error("WindData constructor not implemented for $T")
end

######################## TunableWind ############################

struct TunableWind <: AbstractWind end

Base.@kwdef mutable struct USimpleWind
    v_ew_n::MVector{3,Float64} = zeros(MVector{3}) #MVector allows changing single components
end

init_u0(::TunableWind) = USimpleWind()
f_cont!(::System{<:TunableWind}, args...) = nothing
f_disc!(::System{<:TunableWind}, args...) = false

function WindData(wind::System{<:TunableWind}, ::Abstract3DLocation)
    wind.u.v_ew_n |> SVector{3,Float64} |> WindData
end

#################### AtmosphereDescriptor ############################

Base.@kwdef struct AtmosphereDescriptor{I <: AbstractISA, W <: AbstractWind} <: SystemDescriptor
    isa_::I = TunableISA()
    wind::W = TunableWind()
end

#this one's convenient for dispatching on plots() methods, but for U and D we
#can use simply NamedTuples
struct AtmosphereCmpY{I, W}; isa_::I; wind::W; end

init_x0(atm::AtmosphereDescriptor) = ComponentVector(isa_ = init_x0(atm.isa_), wind = init_x0(atm.wind))
init_y0(atm::AtmosphereDescriptor) = AtmosphereCmpY(init_y0(atm.isa_), init_y0(atm.wind))
init_u0(atm::AtmosphereDescriptor) = (isa_ = init_u0(atm.isa_), wind = init_u0(atm.wind))
init_d0(atm::AtmosphereDescriptor) = (isa_ = init_d0(atm.isa_), wind = init_d0(atm.wind))

function System(atm::AtmosphereDescriptor, ẋ = init_x0(atm), x = init_x0(atm),
                    y = init_y0(atm), u = init_u0(atm), d = init_d0(atm), t = Ref(0.0))

    isa_sys = System(atm.isa_, ẋ.isa_, x.isa_, y.isa_, u.isa_, d.isa_, t)
    wind_sys = System(atm.wind, ẋ.wind, x.wind, y.wind, u.wind, d.wind, t)
    params = nothing
    subsystems = (isa_ = isa_sys, wind = wind_sys,)
    System{map(typeof, (atm, x, y, u, d, params, subsystems))...}(
                                ẋ, x, y, u, d, t, params, subsystems)
end

const AtmosphericSystem = System{<:AtmosphereDescriptor}

Base.@kwdef struct AtmosphericData
    isa_::ISAData = ISAData()
    wind::WindData = WindData()
end

function AtmosphericData(a::AtmosphericSystem, pos::Geographic)
    AtmosphericData(
        ISAData(a.subsystems.isa_, pos),
        WindData(a.subsystems.wind, pos))
end

function f_cont!(sys::AtmosphericSystem)
    f_cont!(sys.isa_)
    f_cont!(sys.wind)
end

#we can create a World <: SystemDescriptor to compose Aircraft and
#AtmosphericSystem, and step them together in a single Model. if the
#AtmosphericSystem is dynamic but still we don't want to compose it with the
#Aircraft, we can always put it in a separate model, step it separately, and
#simply pass a reference to its underlying AtmosphericSystem to the
#AircraftSystem, which will be wrapped in its own Model

# #top-down / recursive implementation
# @inline function get_tp(h::Real, T0::Real = T0_std, p0::Real = p0_std, g0::Real = g0_std)

#     h == 0 && return (T0, p0)
#     (h_b, β) = layer_parameters(h)
#     (T_b, p_b) = get_tp(h_b, T0, p0, g0) #get pt at the layer base
#     T = ISA_temperature_law(h, T_b, h_b, β)
#     p = ISA_pressure_law(h, g0, p_b, T_b, h_b, β)
#     return (T, p)
# end

end


"""
```
  genlasso(
      X    :: AbstractMatrix{T},
      y    :: AbstractVector{T};
      path :: Bool = true,
      ρ    :: Union{AbstractVector, Number} = zero(T),
      D    :: AbstractMatrix{T} = eye(size(X, 2)),
      solver = ECOSSolver(maxit=10e8, verbose=0)
    )
```
Solve generalized lasso problem by reformulating it as constrained lasso problem.
Note genralized lasso minimizes
    `0.5sumabs2(√obswt .* (y - X * β)) + ρ * sumabs(D * β)`.

### Arguments
- `X`       : predictor matrix.
- `y`       : response vector.

### Optional arguments
- `path`    : set `path=false` if user wishes to supply parameter value(s).
              Default is true.
- `ρ`       : tuning parameter value(s). Default is 0.
- `D`       : penalty matrix. Default is identity matrix.
- `solver`  : a solver Convex.jl supports. Default is ECOS.
              Note that Mosek and Gurobi are more robust than ECOS. Unlike ECOS or
              SCS, both Mosek and Gurobi require a license (free for academic
              use). <http://convexjl.readthedocs.io/en/latest/solvers.html>

### Returns
- `β`       : estimated coefficents.
- `objval`  : optimal objective value.
- `problem` : Convex.jl problem.
"""
function genlasso(
    X::AbstractMatrix{T},
    y::AbstractArray{T};
    path::Bool = true,
    ρ::Union{AbstractVector, Number} = zero(T),
    D::AbstractMatrix{T} = eye(size(X, 2)),
    solver = ECOSSolver(maxit=10e8, verbose=0)
    ) where T

  # singular value decomposition on D
  m = size(D, 1)
  F = svd!(D, full = true)
  # extract singular values and submatrices
  singvals = F.S
  rankD = count(!iszero,F.S .> abs(F.S[1]) * eps(F.S[1]) * maximum(size(D)))
  V1 = F.V[:, 1:rankD]
  V2 = F.V[:, rankD+1:end]
  U1 = F.U[:, 1:rankD]
  U2 = F.U[:, rankD+1:end]
  # calculate the MP-inverse of D
  Dplus = V1 * broadcast(*, U1', 1 ./ F.S)
  # transform the design matrix
  XDplus = X * Dplus
  XV2 = X * V2
  # projection matrix onto C(XV2)
  Pxv2 = (1 / dot(XV2, XV2)) * XV2 * transpose(XV2)
  # orthogonal projection matrix
  Mxv2 = Matrix{Float64}(I, size(XV2, 1),size(XV2, 1)) - Pxv2
  # obtain the new design matrix and response vector
  ỹ = vec(Mxv2 * y)
  X̃ = Mxv2 * XDplus
  # solve the constrained lasso problem
  if path
    α̂path, ρpath, = lsq_classopath(X̃, ỹ; solver=solver, Aeq = U2',
          beq = zeros(size(U2, 2)))
  else
    α̂path, = lsq_constrsparsereg(X̃, ỹ, ρ; solver=solver, Aeq = U2',
          beq = zeros(size(U, 2)))
    ρpath = ρ
  end
  # transform back to beta
  β̂path = Base.LinAlg.BLAS.ger!(1.0, vec(V2 * ((1 / dot(XV2, XV2)) *
    		At_mul_B(XV2, y))), ones(size(ρpath)), (Matrix{Float64}(I, size(V2, 1),size(V2, 1)) -
    		V2 * ((1 / dot(XV2, XV2)) * At_mul_B(XV2, X))) * Dplus * α̂path)

  return β̂path, ρpath

end # end of function


using Test
using MiniKalman, GaussianDistributions, StaticArrays, FillArrays, Statistics, Random,
    Unitful, Parameters

@with_kw struct RandomWalk <: MiniKalman.Model
    step_size
    obs_noise
end
MiniKalman.observation_noise(rw::RandomWalk, _, _) = white_noise1(rw.obs_noise)
MiniKalman.transition_noise(rw::RandomWalk, _, _) = white_noise1(rw.step_size)
MiniKalman.initial_state(rw::RandomWalk) = white_noise1(0.5) # centered around 0
MiniKalman.nobservations(rw::RandomWalk, _) = 1

# Check that sampling from the true model, then inferring the parameters of RandomWalk
# through maximum likelihood applied to that sample, yields the true model.
true_model = RandomWalk(step_size=2.0, obs_noise=0.5)
sr = sample_and_recover(true_model, nothing, 100000,
                        start_model=RandomWalk(step_size=1.0, obs_noise=1.0),
                        rng=MersenneTwister(1))

@test sr.estimated_model.step_size ≈ sr.true_model.step_size rtol=0.02
@test sr.estimated_model.obs_noise ≈ sr.true_model.obs_noise rtol=0.02

# Needs revision:
# N = 100
# rng = MersenneTwister(1)

# # Need to generate it as regular vectors/matrices first to use `\` below.
# xx = rand(rng, 1, 1,N)
# XX = hcat(ones(1,1,N), xx)
# XX_mat = [convert(SMatrix{1,2}, XX[:,:,i]) for i in 1:N]
# YY = 7*xx[:] + randn(rng, N) .* sqrt(1) .+ 15
# YY_vec = map(SVector, YY)

# xxf2, ll2 = kalman_filter(Gaussian([20., 20.0], [1000.0 0; 0 1000.0]), YY_vec,
#                           Fill(MiniKalman.white_noise1(10.0), N),
#                           observation_mats=XX_mat)

# xf = mean(xxf2[end])
# @test xf ≈ XX[1, :, :]' \ YY rtol=0.01
# @assert N==100 # otherwise the following test is not exact
# @test xf ≈ [15.0589, 6.73069] rtol=0.001

# @test sum(ll2) ≈ -328.39288090529027
# rand(white_noise2(1.0u"m^2"))


### WarpedArray
"""
A `WarpedArray` `W` is an AbstractArray for which `W[x] = A[ϕ(x)]` for
some parent array `A` and some deformation `ϕ`.  The object is created
lazily, meaning that computation of the displaced values occurs only
when you ask for them explicitly.

Create a `WarpedArray` like this:

```
W = WarpedArray(A, ϕ)
```
where

- The first argument `A` is an `AbstractExtrapolation` that can be
  evaluated anywhere.  See the Interpolations package.
- ϕ is an `AbstractDeformation`
"""
struct WarpedArray{T,N,A<:Extrapolatable,D<:AbstractDeformation} <: AbstractArray{T,N}
    data::A
    ϕ::D
end

# User already supplied an interpolatable ϕ
function WarpedArray(data::Extrapolatable{T,N},
                     ϕ::GridDeformation{S,N,A}) where {T,N,S,A<:AbstractInterpolation}
    WarpedArray{T,N,typeof(data),typeof(ϕ)}(data, ϕ)
end

# Create an interpolatable ϕ
function WarpedArray(data::Extrapolatable{T,N}, ϕ::GridDeformation) where {T,N}
    itp = scale(interpolate(ϕ.u, BSpline(Quadratic(Flat(OnCell())))), ϕ.nodes...)
    ϕ′ = GridDeformation(itp, ϕ.nodes)
    WarpedArray{T,N,typeof(data),typeof(ϕ′)}(data, ϕ′)
end

WarpedArray(data, ϕ::GridDeformation) = WarpedArray(to_etp(data), ϕ)

Base.size(A::WarpedArray) = size(A.data)
Base.size(A::WarpedArray, i::Integer) = size(A.data, i)
Base.axes(A::WarpedArray) = axes(A.data)
Base.axes(A::WarpedArray, i::Integer) = axes(A.data, i)

@inline function Base.getindex(W::WarpedArray{T,N}, I::Vararg{Number,N}) where {T,N}
    ϕx = W.ϕ(I...)
    return W.data(ϕx...)
end

ImageAxes.getindex!(dest, W::WarpedArray{T,N}, coords::Vararg{Any,N}) where {T,N} =
    Base._unsafe_getindex!(dest, W, coords...)


@testset "36.valid-sudoku.jl" begin
    board = [
        ['5', '3', '.', '.', '7', '.', '.', '.', '.'],
        ['6', '.', '.', '1', '9', '5', '.', '.', '.'],
        ['.', '9', '8', '.', '.', '.', '.', '6', '.'],
        ['8', '.', '.', '.', '6', '.', '.', '.', '3'],
        ['4', '.', '.', '8', '.', '3', '.', '.', '1'],
        ['7', '.', '.', '.', '2', '.', '.', '.', '6'],
        ['.', '6', '.', '.', '.', '.', '2', '8', '.'],
        ['.', '.', '.', '4', '1', '9', '.', '.', '5'],
        ['.', '.', '.', '.', '8', '.', '.', '7', '9'],
    ]
    @test is_valid_sudoku(board) == true

    board = [
        ['.', '.', '4', '.', '.', '.', '6', '3', '.'],
        ['.', '.', '.', '.', '.', '.', '.', '.', '.'],
        ['5', '.', '.', '.', '.', '.', '.', '9', '.'],
        ['.', '.', '.', '5', '6', '.', '.', '.', '.'],
        ['4', '.', '3', '.', '.', '.', '.', '.', '1'],
        ['.', '.', '.', '7', '.', '.', '.', '.', '.'],
        ['.', '.', '.', '5', '.', '.', '.', '.', '.'],
        ['.', '.', '.', '.', '.', '.', '.', '.', '.'],
        ['.', '.', '.', '.', '.', '.', '.', '.', '.'],
    ]
    @test is_valid_sudoku(board) == false

    board = [
        ['8', '3', '.', '.', '7', '.', '.', '.', '.'],
        ['6', '.', '.', '1', '9', '5', '.', '.', '.'],
        ['.', '9', '8', '.', '.', '.', '.', '6', '.'],
        ['8', '.', '.', '.', '6', '.', '.', '.', '3'],
        ['4', '.', '.', '8', '.', '3', '.', '.', '1'],
        ['7', '.', '.', '.', '2', '.', '.', '.', '6'],
        ['.', '6', '.', '.', '.', '.', '2', '8', '.'],
        ['.', '.', '.', '4', '1', '9', '.', '.', '5'],
        ['.', '.', '.', '.', '8', '.', '.', '7', '9'],
    ]
    @test is_valid_sudoku(board) == false
end


sqr(x) = x ^ 2


function test_sparse_vanilla_same(m)
    vanilla_solver = ValueIterationSolver(max_iterations=500, belres=1e-8)
    sparse_solver = SparseValueIterationSolver(max_iterations=500, belres=1e-8)
    vanilla_policy = solve(vanilla_solver, m)
    sparse_policy = solve(sparse_solver, m)
    @test vanilla_policy.policy == sparse_policy.policy
    @test isapprox(vanilla_policy.util, sparse_policy.util; atol=1e-3)
    for s in states(m)
        si = stateindex(m, s)
        for a in actions(m, s)
            ai = actionindex(m, a)
            @test isapprox(vanilla_policy.qmat[si, ai], sparse_policy.qmat[si, ai]; atol=1e-3)
        end
    end
end

m1 = SimpleGridWorld(size=(10,10))
test_sparse_vanilla_same(m1)

m2 = SpecialGridWorld()
test_sparse_vanilla_same(m2)

## Two state MDP to test for terminal states

struct TwoStatesMDP <: MDP{Int, Int} end

POMDPs.states(mdp::TwoStatesMDP) = 1:2
POMDPs.stateindex(mdp::TwoStatesMDP, s) = s
POMDPs.actions(mdp::TwoStatesMDP) = 1:2
POMDPs.actionindex(mdp::TwoStatesMDP, a) = a
POMDPs.discount(mdp::TwoStatesMDP) = 0.95
POMDPs.transition(mdp::TwoStatesMDP, s, a) = SparseCat([a], [1.0])
POMDPs.reward(mdp::TwoStatesMDP, s, a, sp) = float(sp == 2)
POMDPs.isterminal(mdp::TwoStatesMDP, s) = s == 2

mdp = TwoStatesMDP()
solver = ValueIterationSolver(verbose = true)
policy = solve(solver, mdp)

sparsesolver = SparseValueIterationSolver(verbose=true)
sparsepolicy = solve(sparsesolver, mdp)

@test sparsepolicy.qmat == policy.qmat
@test value(sparsepolicy, 2) ≈ 0.0
@test_throws String solve(SparseValueIterationSolver(), TigerPOMDP())

@inferred solve(SparseValueIterationSolver(verbose=false), mdp)


@testset "intermodes" begin
    original_image = testimage("lena")
    for T in (Gray{N0f8}, Gray{N0f16}, Gray{Float32}, Gray{Float64})
        img = T.(original_image)
        img₀₁ = binarize(Intermodes(), img)

        # Check original image is unchanged.
        @test img == T.(testimage("lena"))

        # Check that the image only has ones or zeros.
        non_zeros = findall(x -> x != 0.0 && x != 1.0, img₀₁)
        @test length(non_zeros) == 0

        # Check type of binarized image.
        @test typeof(img₀₁) == Array{Gray{Bool},2}

        # Check that ones and zeros have been assigned to the correct side of the threshold.
        maxval, maxpos = findmax(Gray.(img))
        @test img₀₁[maxpos] == 1
        minval, minpos = findmin(Gray.(img))
        @test img₀₁[minpos] == 0
    end
end


include("configure.jl")

corrTs = [2.1, 2.2, 2.25, 2.26, 2.265, 2.269]
μs = Array(Float64,mcd,length(corrTs))   # means of correlator
μas = Array(Float64,mcd,length(corrTs))  # sample means of expected radius
σs = Array(Float64,mcd,length(corrTs))   # standard devs of correlator
σas = Array(Float64,mcd,length(corrTs))  # sample standard devs of expected radius
estimators = Array(Float64, length(corrTs))
corrSteps = 2000

# find spatial correlators for multiple T values
for t in 1:length(corrTs)

    @time mags, acc, naivecorrelators =    # throw in correlators later on
            cluster(Nx, Ny, corrSteps, J, 0.0, corrTs[t], mcd)

    # need to find μs and σs for the spatial correlators at diff separations
    μs[:,t] = mean(naivecorrelators, 2)
    μas[:,t] = mean(aoft(naivecorrelators), 2)
    σs[:,t] = std(naivecorrelators, 2)
    σas[:,t] = std(aoft(naivecorrelators), 2)
    ### estimators[t] = mean(correlators)

    writedlm("cormeans", μs)
    writedlm("corstddevs", σs)
    ### writedlm("estimatormeans", estimatormeans) # DOESNT WORK ANYWAY
end


@doc raw"""
    SubGradientProblem <: Problem

A structure to store information about a subgradient based optimization problem

# Fields
* `manifold` – a [Manifold](https://juliamanifolds.github.io/Manifolds.jl/stable/interface.html#ManifoldsBase.Manifold)
* `cost` – the function $F$ to be minimized
* `subgradient` – a function returning a subgradient $\partial F$ of $F$

# Constructor

    SubGradientProblem(M, f, ∂f)

Generate the [`Problem`] for a subgradient problem, i.e. a function `f` on the
manifold `M` and a function `∂f` that returns an element from the subdifferential
at a point.
"""
struct SubGradientProblem{mT<:Manifold,TCost,TSubgradient} <: Problem
    M::mT
    cost::TCost
    subgradient::TSubgradient
end
"""
    get_subgradient(p,x)

Evaluate the (sub)gradient of a [`SubGradientProblem`](@ref)` p` at the point `x`.
"""
function get_subgradient(p::SubGradientProblem, x)
    return p.subgradient(x)
end
"""
    SubGradientMethodOptions <: Options
stories option values for a [`subgradient_method`](@ref) solver

# Fields
* `retraction_method` – the retration to use within
* `stepsize` – a [`Stepsize`](@ref)
* `stop` – a [`StoppingCriterion`](@ref)
* `x` – (initial or current) value the algorithm is at
* `x_optimal` – optimal value
* `∂` the current element from the possivle subgradients at `x` that is used
"""
mutable struct SubGradientMethodOptions{TRetract<:AbstractRetractionMethod,TStepsize,P,T} <:
               Options where {P,T}
    retraction_method::TRetract
    stepsize::TStepsize
    stop::StoppingCriterion
    x::P
    x_optimal::P
    ∂::T
    function SubGradientMethodOptions(
        M::TM,
        x::P,
        sC::StoppingCriterion,
        s::Stepsize,
        retraction_method=ExponentialRetraction(),
    ) where {TM<:Manifold,P}
        return new{typeof(retraction_method),typeof(s),P,typeof(zero_tangent_vector(M, x))}(
            retraction_method, s, sC, x, x
        )
    end
end


module TestPITMatrix

using Copulas
using Base.Test

## single observation
kk = PITMatrix([0.1, 0.3])

kk2 = [0.1 NaN;
       NaN 0.3]
expOut = PITMatrix(reshape(kk2, 2, 2, 1))

@test isequal(kk, expOut)

## multiple observations
kk = PITMatrix([0.1 0.3;
                0.4 0.7;
                0.6 0.4])

@test size(dg(kk)) == (2, 2, 3)
@test kk[1, 1] == [0.1, 0.4, 0.6]
@test isnan(dg(kk)[1, 2])

## throw error during construction
@test_throws Exception PITMatrix([0.1 0.3;
                                  0.4 1.2;
                                  0.6 0.4])

@test_throws Exception PITMatrix([0.1 0.3;
                                  0.4 -0.2;
                                  0.6 0.4])

@test_throws Exception PITMatrix([0.1 0.3;
                                  0.4 -0.2;
                                  0.6 0.4])

kk = reshape([0.1 0.2;
              0.4 0.3;
              0.5 0.1], 3, 2, 1)

@test_throws Exception PITMatrix(kk)

end


s = open("input.txt") do file
    readlines(file);
end

startTime = parse(Int, s[1])

nums = split(s[2], ',');
getnums = function(x)
    if x == "x"
        return -1;
    else
        return parse(Int, x);
    end
end
busNums = map(getnums, nums);

notX = busNums.!=-1;
routeNums = busNums[notX];

cur = zeros(length(routeNums));

while all(cur .< startTime)
    cur .+= routeNums;
end

final = maximum(cur);
index = findall(x->x>=startTime, cur)

routeNums[3] * (final-startTime)


using DoubleFloats, GenericLinearAlgebra, LinearAlgebra
using Test

using Random
phi = Base.MathConstants.golden

include("specialvalues.jl")
include("construct.jl")
include("compare.jl")
include("promote.jl")

include("arithmetic.jl")
include("functions.jl")
include("matmul.jl")
include("complex.jl")

include("concrete_accuracy.jl")
include("function_accuracy.jl")

include("coverage.jl")

include("string.jl")
include("parse.jl")


if !isdefined(Main, :CustomStencil)
    include("../module/CustomStencil.jl")
    using Main.CustomStencil
end
include("misc.jl")
include("cpu_stencils.jl")

nx = 5
ny = 5
nz = 5
dx = 1<<(nx)
dy = 1<<(ny)
dz = 1<<(nz)

original_data = CreateData(dx,dy,dz)
## Stencil Definition

i = 1
radius = 1
t_steps = 2
coefs = [Float64(round(1/i, digits=2)) for i = 1:2]
star_stencil = @def_stencil_expression c[0]D[x,y,z] + @sum(i, 1, 1, c[i]*(
            D[x+i,y,z] + D[x,y+i,z] + D[x,y,z+i] +
            D[x-i,y,z] + D[x,y-i,z] + D[x,y,z-i]))
st_def = CreateStencilDefinition(star_stencil, coefs)

st_v = @def_stencil_expression v*c[0]D[x,y,z] + @sum(i, 1, 1, v*c[i]*(
            D[x+i,y,z] + D[x,y+i,z] + D[x,y,z+i] +
            D[x-i,y,z] + D[x,y-i,z] + D[x,y,z-i]))
st_v_def = CreateStencilDefinition(st_v, coefs)
st_inst = NewStencilInstance(st_def,st_v_def)

## Input Data size Definition
data = PadData(radius, original_data)
@time dout = ApplyStencil(st_inst, data, t_steps)
gpu_out = Array(dout)

using SymEngine

d = st_inst.stencil_sym
total_max_radius = st_inst.max_radius
radius = total_max_radius
max_radius = radius
Dsym = [symbols("w$(i)$(j)$(k)") for i = 1:(2total_max_radius+1),
                                    j = 1:(2total_max_radius+1),
                                    k = 1:(2total_max_radius+1)]

Doff = [[i,j,k] for i = -radius:radius, j=-radius:radius, k=-radius:radius]
new_Dsym = [symbols("w$(i)$(j)$(k)") for i = 1:(4total_max_radius+1),
                                    j = 1:(4total_max_radius+1),
                                    k = 1:(4total_max_radius+1)]
new_Doff =  [[i,j,k] for i = -2radius:2radius, j=-2radius:2radius, k=-2radius:2radius]

new_eq = 0
for i = 1:(2*max_radius+1), j = 1:(2*max_radius+1),k = 1:(2*max_radius+1)
    global d, Doff, Dsym, new_Dsym, new_Doff,new_eq,radius
    if coeff(d,Dsym[i,j,k]) != 0
        @show Dsym[i,j,k]
        @show c1 = coeff(d,Dsym[i,j,k])
        for ii = 1:(2*max_radius+1), jj = 1:(2*max_radius+1),kk = 1:(2*max_radius+1)
            if coeff(d,Dsym[ii,jj,kk]) != 0
                @show c = coeff(d,Dsym[ii,jj,kk])
                @show off2 = Doff[ii,jj,kk] + [i;j;k] .+ radius
                new_eq += c1*c*new_Dsym[off2...]
            end
        end
    end

end
new_eq

d1 = [0 0.5 0; 0.5 1 0.5; 0 0.5 0]
c = PaddedView(0, d1, (1:7,1:7), (3:5,3:5) )
c = Array(c)
res = zeros(Float64, 7 , 7)
for i = 1:3
    global res
    for j = 1:3
        res[i:i+2,j:j+2] += c[i+2,j+2]*d1

    end
end

@show res

SymEngine.

eq = a^2 + a + 2

for i in enumerate(eq)
    println(i)
end

SymEngine.subs(eq, b, symbols(c))


using GeoRDD
using Test

function sim_data(n::Int, τ::Real)
    X = rand(2, n)
    f_X = sin.(X[1,:] * 2.0) .+ 3.0 .* X[2,:].^2.0
    m = -0.8 # mean offset
    treatment_radius = 0.8
    covarA = randn(n)
    covarB = randn(n)
    categ = rand(["A", "B", "C"], n) # random cateogry
    treat = sqrt.(X[1,:].^2 + X[2,:].^2) .< treatment_radius

    # outcome is: smooth surface + noise + treatment effect
    Y = m .+ # mean offset
        f_X .+ # smooth surface
        0.3*randn(n) .+ # noise
        τ .* treat .+ # treatment effect
        covarA .* 0.1 .+ # effect of real-valued covariate
        (categ.=="B") .* 0.2 # raise category B a little bit
    border_X1 = treatment_radius .* cos.(range(0,stop=π/2,length=1000))
    border_X2 = treatment_radius .* sin.(range(0,stop=π/2,length=1000))
    border_geo = LibGEOS.LineString([[border_X1[i], border_X2[i]] for i in 1:n])
    geordd_df = DataFrame(
        X1 = X[1, :],
        X2 = X[2, :],
        outcome = Y,
        region = treat, # treatment indicator
        covarA = covarA,
        covarB = covarB,
        categ = categ
        )
    categorical!(geordd_df, :categ)
    categorical!(geordd_df, :region)
    return geordd_df, border_X1, border_X2, border_geo
end
include("test_geometry.jl")
include("test_gprealisations.jl")
include("test_simulated.jl")


# Requires LightGraphs and SimpleWeightedGraphs
module GraphUtils
using LightGraphs, SimpleWeightedGraphs
using Base.Cartesian
export CompleteBipartiteDiGraph, CompleteMultipartiteDiGraph, ComplexGraph, mesh_graph, grid_layout, add_inputs!

"""
    CompleteBipartiteDiGraph(n1, n2)

Create a directed [complete bipartite graph](https://en.wikipedia.org/wiki/Complete_bipartite_graph)
with `n1 + n2` vertices.
"""
function CompleteBipartiteDiGraph(n1::T, n2::T) where {T <: Integer}
    (n1 < 0 || n2 < 0) && return SimpleDiGraph{T}(0)
    Tw = widen(T)
    n = T(Tw(n1) + Tw(n2))  # checks if T is large enough for n1 + n2

    adjmat = falses(n, n)
    range1 = 1:n1
    range2 = (n1 + 1):n
    @inbounds for u in range1, v in range2
        adjmat[u, v] = true
    end
    return SimpleDiGraph{T}(adjmat)
end

"""
    CompleteMultipartiteDiGraph(partitions)

Create a directed [complete multipartite graph](https://en.wikipedia.org/wiki/Multipartite_graph)
with `sum(partitions)` vertices. A partition with `0` vertices is skipped.
### Implementation Notes
Preserves the eltype of the partitions vector. Will error if the required number of vertices
exceeds the eltype.
"""
function CompleteMultipartiteDiGraph(partitions::AbstractVector{T}) where {T<:Integer}
    any(x -> x < 0, partitions) && return SimpleDiGraph{T}(0)
    length(partitions) == 1 && return SimpleDiGraph{T}(partitions[1])
    length(partitions) == 2 && return CompleteBipartiteDiGraph(partitions[1], partitions[2])

    n = sum(partitions)
    adjmat = falses(n, n)
    cur = 1
    for i in 1:length(partitions)
        curp = partitions[i]
        nextp = i < length(partitions) ? partitions[i+1] : 0
        currange = cur:(cur+curp-1) # all vertices in the current partition
        upperrange = (cur+curp):(cur+curp+nextp-1)   # all vertices higher than the current partition

        @inbounds for u in currange, v in upperrange
            adjmat[u,v] = true
        end
        cur += curp
    end

    return SimpleDiGraph{T}(adjmat)
end

"""
    ComplexGraph(metagraph, layers)

Construct a complex graph from a metagraph, a graph where each vertex vi represents a set of layers[vi] vertices in the complex graph which have identical edges. If selfloops is true, a self-loop in the metagraph will produce self-loops in all of the vertices in the resulting set. If completeloops is true, a self-loop in the metagraph will produce full connectivity between all vertices in the resulting set.
"""
function ComplexGraph(metagraph::G, layers::AbstractVector{<:Integer}, probabilityGraph::WG=SimpleWeightedDiGraph(SimpleDiGraph(weights(metagraph)), 1.0); selfloops=true, completeloops=true) where {G <: AbstractGraph, WG <: AbstractSimpleWeightedGraph{T,U} where {T, U}}
    nlayers = length(layers)
    if nlayers != nv(metagraph)
        error("$layers must have the same length as the number of vertices in $metagraph.")
    elseif nv(metagraph) != nv(probabilityGraph)
        error("$metagraph must have the same number of vertices as $probabilityGraph.")
    elseif !iszero(ne(symmetric_difference(metagraph, probabilityGraph)))
        error("$metagraph must have the same enges as $probabilityGraph.")
    end

    # construct ranges for each layer
    cur = 1
    layerrange = Vector{Vector{Int}}(undef, nlayers)
    for i in 1:nlayers
        layerrange[i] = cur:(cur+layers[i] - 1)
        cur += layers[i]
    end

    # fill in the adjacency matrix
    n = sum(layers)
    adjmat = falses(n, n)
    for i in 1:nlayers
        for j in outneighbors(metagraph, i)
            for u in layerrange[i]
                for v in layerrange[j]
                    # fill in the edge with probability weights(probabilityGraph)[i, j]
                    if (i != j || (u == v && selfloops) || (u != v && completeloops)) && rand() <= weights(probabilityGraph)[i, j]
                        adjmat[u,v] = true
                    end
                end
            end
        end

    end

    return G(adjmat)
end

function ComplexGraph(metagraph::G, layers::AbstractVector{<:Integer}, probabilityGraph::G=SimpleWeightedDiGraph(SimpleDiGraph(weights(metagraph)), 1.0); selfloops=true, completeloops=true) where G<:AbstractSimpleWeightedGraph{T, U} where {T,U}
    nlayers = length(layers)
    if nlayers != nv(metagraph)
        error("$layers must have the same length as the number of vertices in $metagraph.")
    end

    # construct ranges for each layer
    cur = 1
    layerrange = Vector{Vector{Int}}(undef, nlayers)
    for i in 1:nlayers
        layerrange[i] = cur:(cur+layers[i] - 1)
        cur += layers[i]
    end

    # fill in the adjacency matrix
    n = sum(layers)
    w = weights(metagraph)
    adjmat = zeros(U, n, n)
    for i in 1:nlayers
        for j in outneighbors(metagraph, i)
            for u in layerrange[i]
                for v in layerrange[j]
                    if (i != j || (u == v && selfloops) || (u != v && completeloops)) && rand() <= weights(probabilityGraph)[i, j]
                        adjmat[u,v] = w[i, j]
                    end
                end
            end
        end
    end

    return G(adjmat)
end

"""
    ComplexGraph(Ns::Vector{<:Integer}, edges::Vararg{NTuple{4, Integer}}, NE}; kwargs...) where NE

Construct a graph from a metagraph with `length(Ns)` nodes.
`edges` is a variable number of tuples `(u, v, w, p)`, where the edge on the
metagraph is from vertex `u` to vertex `v` with weight `w` and probability `p`.
"""
function ComplexGraph(Ns::Vector{<:Integer}, edges::Vararg{Tuple{<:Integer, <:Integer, <:Real, <:Real}}; kwargs...)
    # Construct a weighted meta-graph with length(N) populations
    metagraph = SimpleWeightedDiGraph(length(Ns))
    probabilityGraph = SimpleWeightedDiGraph(length(Ns))

    for (u, v, w, p) in edges
        add_edge!(metagraph, u, v, w)
        add_edge!(probabilityGraph, u, v, p)
    end
    return ComplexGraph(metagraph, Ns, probabilityGraph; kwargs...)
end

"""
    _neighbors(ind::CartesianIndex{N}; forward::Vector{<:Integer}=ones(Int, N), backward::Vector{<:Integer}=ones(Int, N),
               self::Bool=true, diagonals::Bool=true, horizontals::Bool=true)::Vector{CartesianIndex{N}} where N

For any `N`-dimensional index `ind`, generate a vector of all indices neighboring
`ind`, including all indices within `backward` indices behind ind to `forward` in
front of ind. `ind` is included in the return value iff `self` is true. Indices
diagonal to `ind` are included iff `diagonals` is true, and indices horizontal to
`ind` are included iff `horizontals` is true.
"""
function _neighbors(ind::CartesianIndex{N}; forward::Vector{<:Integer}=ones(Int, N), backward::Vector{<:Integer}=ones(Int, N), self::Bool=true, diagonals::Bool=true, horizontals::Bool=true)::Vector{CartesianIndex{N}} where N
    if length(forward) != N || length(backward) != N
        error("Forward and backward strides must have length $N.")
    end

    return [ind+i for i in -CartesianIndex(backward...):CartesianIndex(forward...) if (self && i == zero(i)) || (horizontals && any(iszero.(Tuple(i)))&& any(.!iszero.(Tuple(i)))) || (diagonals && !any(iszero.(Tuple(i))))]
end

"""
    _mesh(G::Type{<:AbstractGraph}, dims::Tuple{Vararg{<:Integer, N}}, args...; kwargs...) where N

Generate an `N` dimensional mesh graph of type `G` with size `dims` in each dimension.
Optional arguments are passed to LightGraphs.add_edge!, and keyword arguments are
passed to `SpikingNeuralNets.GraphUtils_neighbors`.
"""
function _mesh(G::Type{<:AbstractGraph}, dims::Tuple{Vararg{<:Integer, N}}, args...; kwargs...) where N
    if iszero(N)
        return G()
    end

    g = G(reduce(*, dims))
    verts = reshape(vertices(g), dims)    # store the linear index of each vertex
    inds = CartesianIndices(verts)
    Ifirst, Ilast = Tuple(first(inds)), Tuple(last(inds))
    for u in inds
        for v in _neighbors(u; self=false, kwargs...)
            if all(Tuple(v) .>= Ifirst) && all(Tuple(v) .<= Ilast)
                add_edge!(g, verts[u], verts[v], args...)
            end
        end
    end

    return g
end

"""
    mesh_graph(dims::Tuple{Vararg{<:Integer, N}}; directed=true, weight=nothing, kwargs...) where N

Generate an `N`-dimensional mesh graph of size `dims` in each dimension.
Keyword arguments:
  `directed`: if true (default), return a directed graph. Else return an undirected graph.
  `weight`: if nothing (default), return an `AbstractSimpleGraph`. Else return an
            `AbstractSimpleWeightedGraph` where each edge has weight `weight`.
  `backward`: the number of vertices (in each dimension) behind each vertex `v`
              with incoming edges from `v`. By default, `ones(Int, N)`.
  `forward`: the number of vertices (in each dimension) in front of each vertex
             `v` with incoming edges from `v`. By default, `ones(Int, N)`.
  `horizontals`: if true (default), vertices are connected horizontally (where vertices
                 are aligned on any single dimension). Else these edges are omitted.
  `diagonals`: if true (default), vertices are connected diagonally, where vertices are not
               aligned on any dimension. Else these edges are omitted.
  `self`: if true (default), edges `(v, v)` are included for all vertices `v`.
          Else these edges are omitted.
"""
function mesh_graph(dims::Tuple{Vararg{<:Integer, N}}; directed=true, weight=nothing, kwargs...) where N
    if isnothing(weight)
        if directed
            return _mesh(SimpleDiGraph, dims; kwargs...)
        else
            return _mesh(SimpleGraph, dims; kwargs...)
        end
    else
        if directed
            return _mesh(SimpleWeightedDiGraph, dims, weight; kwargs...)
        else
            return _mesh(SimpleWeightedGraph, dims, weight; kwargs...)
        end
    end
end

"""
    mesh_graph(dims::Vararg{<:Integer, N}; directed=true, weight=nothing, kwargs...) where N

Generate an `N`-dimensional mesh graph of size `dims` in each dimension.
Keyword arguments:
  `directed`: if true (default), return a directed graph. Else return an undirected graph.
  `weight`: if nothing (default), return an `AbstractSimpleGraph`. Else return an
            `AbstractSimpleWeightedGraph` where each edge has weight `weight`.
  `backward`: the number of vertices (in each dimension) behind each vertex `v`
              with incoming edges from `v`. By default, `ones(Int, N)`.
  `forward`: the number of vertices (in each dimension) in front of each vertex
             `v` with incoming edges from `v`. By default, `ones(Int, N)`.
  `horizontals`: if true (default), vertices are connected horizontally (where vertices
                 are aligned on any single dimension). Else these edges are omitted.
  `diagonals`: if true (default), vertices are connected diagonally, where vertices are not
               aligned on any dimension. Else these edges are omitted.
  `self`: if true (default), edges `(v, v)` are included for all vertices `v`.
          Else these edges are omitted.
"""
function mesh_graph(dims::Vararg{<:Integer, N}; kwargs...) where N
    return mesh_graph(dims; kwargs...)
end

"""
    grid_layout(g::AbstractGraph; reverse_x::Bool=false, reverse_y::Bool=false, columnwise=false)

Return a tuple of x locations and y locations for a graph, laying out
each vertex in a single row (if `columnwise == false`, default) or column
(if `columnwise == true`). The row/column order can be reversed using the boolean
flags `reverse_x` and `reverse_y`, respectively.
"""
function grid_layout(g::AbstractGraph; reverse_x::Bool=false, reverse_y::Bool=false, columnwise=false)
    if columnwise
        locs_x, locs_y = zeros(nv(g)), vertices(g)
    else
        locs_x, locs_y = vertices(g), zeros(nv(g))
    end

    if reverse_x
        reverse!(locs_x)
    end
    if reverse_y
        reverse!(locs_y)
    end
    return locs_x, locs_y
end

"""
    grid_layout(g::AbstractGraph, dims::Tuple{<:Integer, <:Integer};
                reverse_x::Bool=false, reverse_y::Bool=false, columnwise::Bool=true)

Return a tuple of x locations and y locations for all vertices in a graph, laying out
each vertex an a grid of size `dims` in row-major order (if `columnwise == false`)
or column-major order (if `columnwise == true`, default). The row/column order can
be reversed using the boolean flags `reverse_x` and `reverse_y`, respectively.
"""
function grid_layout(g::AbstractGraph, dims::Tuple{<:Integer, <:Integer};
                     reverse_x::Bool=false, reverse_y::Bool=false, columnwise::Bool=true)
    if reduce(*, dims) != nv(g)
        error("dims ($(dims)) does not match the number of vertices in g ($(nv(g))).")
    end

    inds = vec([(Float64(x), Float64(y)) for x in Base.OneTo(dims[1]), y in Base.OneTo(dims[2])])
    if columnwise
        locs_x = getindex.(inds, 2)
        locs_y = getindex.(inds, 1)
    else
        locs_x = getindex.(inds, 1)
        locs_y = getindex.(inds, 2)
    end

    if reverse_x
        reverse!(locs_x)
    end
    if reverse_y
        reverse!(locs_y)
    end

    return locs_x, locs_y
end

"""
    add_inputs!(g::AbstractGraph{T}, vertices::AbstractVector{<:T}=vertices(g))

For every vertex `vi` in `vertices` a new vertex `vinput` having a single edge `vinput->vi`.
Return the indices of all new input vertices.
"""
function add_inputs!(g::AbstractGraph{T}, vertices::AbstractVector{<:T}=vertices(g)) where {T<:Integer}
    vstart = nv(g)+1
    vstop = vstart + length(vertices)
    add_vertices!(g, length(vertices))
    for (vi, vinput) in zip(vertices, vstart:vstop)
        add_edge!(g, vinput, vi)
    end
    return vstart:vstop
end

"""
    add_inputs!(g::AbstractSimpleWeightedGraph{T,U}, vertices::AbstractVector{<:T}=vertices(g), w::AbstractVector{<:U}=ones(U, length(vertices)))

For every vertex `vi` in `vertices` a new vertex `vinput` having a single edge `vinput->vi` with weight `weights[i]`.
Return the indices of all new input vertices.
"""
function add_inputs!(g::AbstractSimpleWeightedGraph{T,U}, vertices::AbstractVector{<:T}=vertices(g), weights::AbstractVector{<:U}=ones(U, length(vertices))) where {T<:Integer, U<:Real}
    vstart = nv(g)+1
    vstop = vstart + length(vertices)
    add_vertices!(g, length(vertices))
    for (vi, vinput, w) in zip(vertices, vstart:vstop, weights)
        add_edge!(g, vinput, vi, w)
    end
    return vstart:vstop
end

end


"""
    StartPairSystem(F)

A system which consider the system `F(x, p)` as a system `F([x p])`, i.e., variable and
parameter space are considered together.
"""
struct StartPairSystem{S<:AbstractSystem,C<:FiniteDiff.JacobianCache} <: AbstractSystem
    F::S
    # it's a little bit unnecessary to have these here, since views are theoretically
    # sufficient. However, we want to avoid any unnecessary recompilation
    # of the potential expensive evaluate! and evaluate_and_jacobian! functions.
    # So we give exactly the same types as in the tracking
    x::Vector{ComplexF64}
    p::Vector{ComplexF64}
    J::Matrix{ComplexF64}
    fd_cache::C
end

StartPairSystem(F::System; compile::Union{Bool,Symbol} = COMPILE_DEFAULT[]) =
    StartPairSystem(fixed(F; compile = compile))
function StartPairSystem(F::AbstractSystem)
    x = zeros(ComplexF64, size(F, 2))
    p = zeros(ComplexF64, nparameters(F))
    J = zeros(ComplexF64, size(F))
    fd_cache = FiniteDiff.JacobianCache(similar(p), zeros(ComplexF64, size(F, 1)))
    StartPairSystem(F, x, p, J, fd_cache)
end
Base.size(F::StartPairSystem) = (size(F.F, 1), size(F.F, 2) + nparameters(F.F))

function ModelKit.evaluate!(u, F::StartPairSystem, xp, ::Nothing = nothing)
    N, m = length(F.x), length(F.p)
    F.x .= view(xp, 1:N)
    F.p .= view(xp, N+1:N+m)
    evaluate!(u, F.F, F.x, F.p)
end

function ModelKit.evaluate_and_jacobian!(u, U, F::StartPairSystem, xp, ::Nothing = nothing)
    n, N, m = size(F.F, 1), length(F.x), length(F.p)
    F.x .= view(xp, 1:N)
    F.p .= view(xp, N+1:N+m)

    evaluate_and_jacobian!(u, F.J, F.F, F.x, F.p)
    view(U, :, 1:N) .= F.J

    fx = (v, q) -> evaluate!(v, F.F, F.x, q)
    FiniteDiff.finite_difference_jacobian!(view(U, :, N+1:N+m), fx, F.p, F.fd_cache)
end


#Pseudo-Random Bit Sequence Tests

@testset "PRBS Tests" begin #Scope for test data

_DEBUG = false

#Dumps contents of a bit sequence vector for debugging purposes:
function dbg_dumpseq(prefix, seq)
	if !_DEBUG; return; end
	print(prefix)
	for bit in seq
		if bit > 0
			print("1")
		else
			print("0")
		end
	end
	println()
end

@testset "Cyclical tests: MaxLFSR" begin
	for prbslen in 3:15 #Don't perform tests on higher-order patterns (too much time/memory)
		patlen = (2^prbslen) - 1
		genlen = patlen + prbslen #Full pattern length + "internal register length"
		regmask = (Int64(1)<<prbslen) -1 #Mask out "non-register" bits
		seed = rand(Int64) & regmask
		pattern = collect(sequence(MaxLFSR(prbslen), seed=seed, len=genlen, output=Int))

		#Make sure pattern repeats for at least as many bits as are in internal register:
		mismatch = pattern[(1:prbslen)] .!= pattern[patlen .+ (1:prbslen)]
		@test sum(convert(Vector{Int}, mismatch)) == 0

		#@show prbslen, pattern[1:prbslen], pattern[patlen .+ (1:prbslen)]
	end
end

@testset "Error detection: MaxLFSR" begin
	#TODO: Compare algorithm against known good pattern (if one can be found).
	prbslen = 15; seed = 11; len = 50
	pattern = collect(sequence(MaxLFSR(prbslen), seed=seed, len=len, output=Int))
	_errors = sequence_detecterrors(MaxLFSR(prbslen), pattern)

	dbg_dumpseq("_errors = ", _errors) #DEBUG
	@test sum(_errors) == 0

	#Test with error injected:
	pattern[prbslen+5] = 1 - pattern[prbslen+5] #Flip "bit"
	_errors = sequence_detecterrors(MaxLFSR(prbslen), pattern)

	dbg_dumpseq("_errors = ", _errors) #DEBUG
	@test sum(_errors) == 1 #Should have a single error
end

end #PRBS Tests

#Last line


@testset "Heuristic tests" begin
    # Load a graph to test on.
    g = graph_from_gr("../examples/test.gr")

    # Test the min fill heuristic.
    order, tw = min_fill(g)
    @test length(order) == lg.nv(g)

    # Test min fill sampling
    order, tw = min_fill(g, 10)
    @test length(order) == lg.nv(g)

    # # Test the min width heuristic.
    # order, tw = min_width(g)
    # @test length(order) == lg.nv(g)

    # # Test sampling version of min width.
    # order, tw = min_width(g, 10)
    # @test length(order) == lg.nv(g)

    # Test maximum minimum degree heuristic for a lower bound.
    lb = mmd(g)
    @test lb <= tw
end


using StatisticalTests
using Test, RDatasets, Statistics, StatsModels, Distributions

test_show(x) = show(IOBuffer(), x)

include("test_t.jl")
include("test_f.jl")
include("test_ks.jl")
include("test_wilcox.jl")


module YaoToEinsum

using Yao, OMEinsum
using LinearAlgebra

export yao2einsum, add_gate!

struct EinBuilder{T}
    slots::Vector{Int}
    labels::Vector{Vector{Int}}
    tensors::Vector{AbstractArray{T}}
    maxlabel::Base.RefValue{Int}
end

Yao.nqubits(eb::EinBuilder) = length(eb.slots)
function add_tensor!(eb::EinBuilder{T}, tensor::AbstractArray{T,N}, labels::Vector{Int}) where {N,T}
    @assert N == length(labels)
    push!(eb.tensors, tensor)
    push!(eb.labels, labels)
end

function EinBuilder(::Type{T}, n::Int) where T
    EinBuilder(collect(1:n), Vector{Int}[], AbstractArray{T}[], Ref(n))
end
newlabel!(eb::EinBuilder) = (eb.maxlabel[] += 1; eb.maxlabel[])

function add_gate!(eb::EinBuilder{T}, b::PutBlock{D,C}) where {T,D,C}
    return add_matrix!(eb, C, mat(T, b.content), collect(b.locs))
end
# general and diagonal gates
function add_matrix!(eb::EinBuilder{T}, k::Int, m::AbstractMatrix, locs::Vector) where T
    if !isdiag(m)
        nlabels = [newlabel!(eb) for _=1:k]
        add_tensor!(eb, reshape(Matrix{T}(m), fill(2, 2k)...), [nlabels..., eb.slots[locs]...])
        eb.slots[locs] .= nlabels
    else
        add_tensor!(eb, reshape(Vector{T}(diag(m)), fill(2, k)...), eb.slots[locs])
    end
    return eb
end
# swap gate
function add_gate!(eb::EinBuilder{T}, b::PutBlock{2,2,ConstGate.SWAPGate}) where {T,N}
    lj = eb.slots[b.locs[2]]
    eb.slots[b.locs[2]] = eb.slots[b.locs[1]]
    eb.slots[b.locs[1]] = lj
    return eb
end

# control gates
function add_gate!(eb::EinBuilder{T}, b::ControlBlock{BT,C,M}) where {T, BT,C,M}
    return add_controlled_matrix!(eb, M, mat(T, b.content), collect(b.locs), collect(b.ctrl_locs), collect(b.ctrl_config))
end
function add_controlled_matrix!(eb::EinBuilder{T}, k::Int, m::AbstractMatrix, locs::Vector, control_locs, control_vals) where T
    if length(control_locs) == 0
        return add_matrix!(eb, k, m, locs)
    end
    sig = eb.slots[control_locs[1]]
    val = control_vals[1]
    for i=1:length(control_locs)-1
        newsig = newlabel!(eb)
        add_tensor!(eb, and_gate(T, control_vals[i+1], val), [newsig,eb.slots[control_locs[i+1]],sig])
        sig = newsig
        val = 1
    end
    if !isdiag(m)
        t1 = reshape(Matrix{T}(m), fill(2, 2k)...)
        t2 = reshape(Matrix{T}(I, 1<<k, 1<<k), fill(2, 2k)...)
        if val == 1
            t1, t2 = t2, t1
        end
        nlabels = [newlabel!(eb) for _=1:k]
        add_tensor!(eb, cat(t1, t2; dims=2k+1), [nlabels..., eb.slots[locs]..., sig])
        eb.slots[locs] .= nlabels
    else
        t1 = reshape(Vector{T}(diag(m)), fill(2, k)...)
        t2 = reshape(ones(T, 1<<k), fill(2, k)...)
        if val == 1
            t1, t2 = t2, t1
        end
        add_tensor!(eb, cat(t1, t2; dims=k+1), [eb.slots[locs]..., sig])
    end
    return eb
end

function and_gate(::Type{T}, a::Int, b::Int) where T
    m = zeros(T, 2, 2, 2)
    for v1 in (0, 1)
        for v2 in (0, 1)
            # the first is output
            m[(v1==a && v2==b)+1, v1+1,v2+1] = 1
        end
    end
    return m
end

function add_gate!(eb::EinBuilder, b::ChainBlock)
    for ib in subblocks(b)
        add_gate!(eb, ib)
    end
    return eb
end

function add_gate!(eb::EinBuilder, b::AbstractBlock)
    B = Optimise.to_basictypes(b)
    if typeof(B) == typeof(b)
        throw("block of type `$(typeof(b))` can not be converted to tensor network representation!")
    else
        add_gate!(eb, B)
    end
    return eb
end

"""
    yao2einsum(circuit, initial_state::Dict, final_state::Dict)
    yao2einsum(circuit; initial_state=Dict(), final_state=Dict())

Transform a Yao `circuit` to a generalized tensor network (einsum) notation.
This function returns a 2-tuple of (einsum code, input tensors).
`initial_state` and `final_state` are dictionaries that specify the initial state and final state as product states,
e.g. a vector `Dict(1=>1, 2=>1, 3=>0, 4=>1)` specifies a product state `|1⟩⊗|1⟩⊗|0⟩⊗|1⟩`.
If an qubit in initial state or final state is not specified, it will be treated as an open edge.

```jldoctest
julia> using YaoToEinsum, Yao

julia> c = chain(3, put(3, 2=>X), put(3, 1=>Y), control(3, 1, 3=>Y))
nqubits: 3
chain
├─ put on (2)
│  └─ X
├─ put on (1)
│  └─ Y
└─ control(1)
   └─ (3,) Y

julia> yao2einsum(c; initial_state=Dict(1=>0, 2=>1), final_state=Dict(1=>ArrayReg([0.6, 0.8im]), 2=>1))
(1, 2, 4∘2, 5∘1, 6∘3∘5, 5, 4 -> 6∘3, AbstractArray{ComplexF64}[[1.0 + 0.0im, 0.0 + 0.0im], [0.0 + 0.0im, 1.0 + 0.0im], [0.0 + 0.0im 1.0 + 0.0im; 1.0 + 0.0im 0.0 + 0.0im], [0.0 + 0.0im 0.0 - 1.0im; 0.0 + 1.0im 0.0 + 0.0im], [1.0 + 0.0im 0.0 + 0.0im; 0.0 + 0.0im 1.0 + 0.0im;;; 0.0 + 0.0im 0.0 - 1.0im; 0.0 + 1.0im 0.0 + 0.0im], [0.6 + 0.0im, 0.0 + 0.8im], [0.0 + 0.0im, 1.0 + 0.0im]])
```
"""
function yao2einsum(circuit::AbstractBlock{D}; initial_state::Dict=Dict{Int,Int}(), final_state::Dict=Dict{Int,Int}()) where {D}
    T = promote_type(ComplexF64, dict_regtype(initial_state), dict_regtype(final_state), Yao.parameters_eltype(circuit))
    vec_initial_state = Dict{Int,ArrayReg{D,T}}([k=>render_single_qubit_state(T, v) for (k, v) in initial_state])
    vec_final_state = Dict{Int,ArrayReg{D,T}}([k=>render_single_qubit_state(T, v) for (k, v) in final_state])
    yao2einsum(circuit, vec_initial_state, vec_final_state)
end
dict_regtype(d::Dict) = promote_type(_regtype.(values(d))...)
_regtype(::ArrayReg{D,VT}) where {D,VT} = VT
_regtype(::Int) = ComplexF64
render_single_qubit_state(::Type{T}, x::Int) where T = x == 0 ? zero_state(T, 1) : product_state(T, bit"1")
render_single_qubit_state(::Type{T}, x::ArrayReg) where T = ArrayReg(collect(T, statevec(x)))

function yao2einsum(circuit::AbstractBlock{D}, initial_state::Dict{Int,<:ArrayReg{D,T}}, final_state::Dict{Int,<:ArrayReg{D,T}}) where {D,T}
    n = nqubits(circuit)
    eb = EinBuilder(T, n)
    openindices = Int[]
    for k=1:n
        if haskey(initial_state, k)
            add_tensor!(eb, statevec(initial_state[k]), [eb.slots[k]])
        else
            push!(openindices, eb.slots[k])
        end
    end
    add_gate!(eb, circuit)
    openindices2 = Int[]
    for k=1:n
        if haskey(final_state, k)
            add_tensor!(eb, statevec(final_state[k]), [eb.slots[k]])
        else
            push!(openindices2, eb.slots[k])
        end
    end
    return build_einsum(eb, vcat(openindices2, openindices))
end

function build_einsum(eb::EinBuilder, openindices)
    return EinCode(eb.labels, openindices), eb.tensors
end

end


using SymbolicUtils, SymbolicUtils.Code, Test
@testset "CSE" begin
    @syms x
    t = cse(hypot(hypot(cos(x), sin(x)), atan(cos(x), sin(x))))

    @test t isa Let
    @test length(t.pairs) == 2
    @test occursin(t.pairs[1].lhs, t.body)
    @test occursin(t.pairs[2].lhs, t.body)
end


include("divide_singularity.jl")
include("JacobiWeight.jl")
include("JacobiWeightOperators.jl")
include("JacobiWeightChebyshev.jl")
include("LogWeight.jl")
include("ExpWeight.jl")


"""
    dualBinaryTree(ptr, N)
Draw the dual domain's binary partition tree.
"""
function dualBinaryTree(ptr, N)
    n = 0 # number of nodes in the tree
    jmax = length(ptr)+1
    for j = 1:jmax-1
        n += length(ptr[j])
    end
    n += N

    node_ind = (Array{Int64,1})[]
    ind = 1
    for j = 1:jmax-1
        lvl_nodes = (Int64)[]
        for i = 1:length(ptr[j])
            push!(lvl_nodes, ind)
            ind += 1
        end
        push!(node_ind, lvl_nodes)
    end
    push!(node_ind, [ind-1+i for i=1:N])

    BT = Graph(n)
    X = zeros(n,2)
    hv = ["(0,0)" for i = 1:n]
    for j = 1:jmax-1
        for i in 1:length(ptr[j])
            c1 = ptr[j][i][1]
            add_edge!(BT, Edge(node_ind[j][i], node_ind[j+1][c1]))
            X[node_ind[j+1][c1],:] = X[node_ind[j][i],:] + [0,1]
            hv[node_ind[j+1][c1]] = "($(j),$(c1-1))"
            if length(ptr[j][i]) == 2
                c2 = ptr[j][i][2]
                add_edge!(BT, Edge(node_ind[j][i], node_ind[j+1][c2]))
                X[node_ind[j+1][c1],:] = X[node_ind[j][i],:] + [-2^(jmax-j),1]
                X[node_ind[j+1][c2],:] = X[node_ind[j][i],:] + [2^(jmax-j),1]
                hv[node_ind[j+1][c2]] = "($(j),$(c2-1))"
            end
        end
    end
    return X, BT, node_ind, hv
end

function find_deepest_leaf(node, ptr)
    jmax = length(ptr)
    j, k = node
    while j <= jmax
        k = ptr[j][k][1]
        j += 1
    end
    return [j, k]
end


"""
    Percentile(x)

Indicate that `x` should be interpreted as a [percentile](https://en.wikipedia.org/wiki/Percentile) rather than an absolute value. For example,

- `canny(img, 1.4, (80, 20))` uses absolute thresholds on the edge magnitude image
- `canny(img, 1.4, (Percentile(80), Percentile(20)))` uses percentiles of the edge magnitude image as threshold
"""
struct Percentile{T} <: Real p::T end

"""
HomogeneousPoint(x::NTuple{N, T})

In projective geometry [homogeneous coordinates](https://en.wikipedia.org/wiki/Homogeneous_coordinates) are the
natural coordinates for describing points and lines.

For instance, the homogeneous coordinates for a planar point are a triplet of real numbers ``(u, v ,w)``, with ``w \\neq 0``.
This triple can be associated with a point ``P = (x,y)`` in Cartesian coordinates, where ``x = \\frac{u}{w}`` and ``y = \\frac{v}{w}``
[(more details)](http://www.geom.uiuc.edu/docs/reference/CRC-formulas/node6.html#SECTION01140000000000000000).

In particular, the `HomogeneousPoint((10.0,5.0,1.0))` is the standardised projective representation of the Cartesian
point `(10.0,5.0)`.
"""
struct HomogeneousPoint{T <: AbstractFloat,N}
    coords::NTuple{N, T}
end

# By overwriting Base.to_indices we can define how to index into an N-dimensional array
# given an (N+1)-dimensional [`HomogeneousPoint`](@ref) type.
# We do this by converting the homogeneous coordinates to Cartesian coordinates
# and rounding to nearest integer.
#
# For homogeneous coordinates of a planar point we return
# a tuple of permuted Cartesian coordinates, (y,x), since matrices
# are indexed  according to row and then column.
# For homogeneous coordinates of other dimensions we do not permute
# the corresponding Cartesian coordinates.
Base.to_indices(A::AbstractArray, p::Tuple{<: HomogeneousPoint}) = homogeneous_point_to_indices(p[1])

function homogeneous_point_to_indices(p::HomogeneousPoint{T,3}) where T
    if  p.coords[end] == 1
        return round(Int,  p.coords[2]), round(Int, p.coords[1])
    else
        return round(Int,  p.coords[2] / p.coords[end]), round(Int, p.coords[1] / p.coords[end])
    end
end

function homogeneous_point_to_indices(p::HomogeneousPoint)
    if  p.coords[end] == 1
        return round.(Int, p.coords)
    else
        return round.(Int, p.coords ./ p.coords[end])
    end
end


abstract type ReproducingKernel end
abstract type ReproducingKernel_0 <: ReproducingKernel end
abstract type ReproducingKernel_1 <: ReproducingKernel_0 end
abstract type ReproducingKernel_2 <: ReproducingKernel_1 end

@doc raw"
`struct RK_H0{T} <: ReproducingKernel_0`

Defines a type of reproducing kernel of Bessel Potential space ``H^{n/2 + 1/2}_ε (R^n)`` ('Basic Matérn kernel'):
```math
V(\eta , \xi, \varepsilon) = \exp (-\varepsilon |\xi - \eta|) \, .
```
# Fields
- `ε::T`: 'scaling parameter' from the Bessel Potential space definition,
           it may be omitted in the struct constructor otherwise it must be greater than zero
"
struct RK_H0{T} <: ReproducingKernel_0
    ε::T
    RK_H0() = new{Float64}(0.0)
    RK_H0(ε) = (@assert ε > 0; new{typeof(float(ε))}(float(ε)))
    RK_H0{T}(ε) where {T} = (@assert ε > 0; new{T}(T(ε)))
end
Base.eltype(::RK_H0{T}) where {T} = T

@doc raw"
`struct RK_H1{T} <: ReproducingKernel_1`

Defines a type of reproducing kernel of Bessel Potential space ``H^{n/2 + 3/2}_ε (R^n)`` ('Linear Matérn kernel'):
```math
V(\eta , \xi, \varepsilon) = \exp (-\varepsilon |\xi - \eta|)
             (1 + \varepsilon |\xi  - \eta|) \, .
```
# Fields
- `ε::T`: 'scaling parameter' from the Bessel Potential space definition,
           it may be omitted in the struct constructor otherwise it must be greater than zero
"
struct RK_H1{T} <: ReproducingKernel_1
    ε::T
    RK_H1() = new{Float64}(0.0)
    RK_H1(ε) = (@assert ε > 0; new{typeof(float(ε))}(float(ε)))
    RK_H1{T}(ε) where {T} = (@assert ε > 0; new{T}(T(ε)))
end
Base.eltype(::RK_H1{T}) where {T} = T

@doc raw"
`struct RK_H2{T} <: ReproducingKernel_2`

Defines a type of reproducing kernel of Bessel Potential space ``H^{n/2 + 5/2}_ε (R^n)`` ('Quadratic Matérn kernel'):
```math
V(\eta , \xi, \varepsilon) = \exp (-\varepsilon |\xi - \eta|)
             (3 + 3\varepsilon |\xi  - \eta| + \varepsilon ^2 |\xi - \eta| ^2 ) \, .
```
# Fields
- `ε::T`: 'scaling parameter' from the Bessel Potential space definition,
           it may be omitted in the struct constructor otherwise it must be greater than zero
"
struct RK_H2{T} <: ReproducingKernel_2
    ε::T
    RK_H2() = new{Float64}(0.0)
    RK_H2(ε) = (@assert ε > 0; new{typeof(float(ε))}(float(ε)))
    RK_H2{T}(ε) where {T} = (@assert ε > 0; new{T}(T(ε)))
end
Base.eltype(::RK_H2{T}) where {T} = T

@inline function _rk(kernel::RK_H2, η::SVector, ξ::SVector)
    x = kernel.ε * norm(η - ξ)
    return (3 + x * (3 + x)) * exp(-x)
end

@inline function _rk(kernel::RK_H1, η::SVector, ξ::SVector)
    x = kernel.ε * norm(η - ξ)
    return (1 + x) * exp(-x)
end

@inline function _rk(kernel::RK_H0, η::SVector, ξ::SVector)
    x = kernel.ε * norm(η - ξ)
    return exp(-x)
end

@inline function _∂rk_∂e(kernel::RK_H2, η::SVector, ξ::SVector, e::SVector)
    t = η - ξ
    x = kernel.ε * norm(t)
    return kernel.ε^2 * exp(-x) * (1 + x) * (t ⋅ e)
end

@inline function _∂rk_∂e(kernel::RK_H1, η::SVector, ξ::SVector, e::SVector)
    t = η - ξ
    x = kernel.ε * norm(t)
    return kernel.ε^2 * exp(-x) * (t ⋅ e)
end

@inline function _∂rk_∂η(kernel::RK_H2, η::SVector, ξ::SVector)
    t = η - ξ
    x = kernel.ε * norm(t)
    return -kernel.ε^2 * exp(-x) * (1 + x) * t
end

@inline function _∂rk_∂η(kernel::RK_H1, η::SVector, ξ::SVector)
    t = η - ξ
    x = kernel.ε * norm(t)
    return -kernel.ε^2 * exp(-x) * t
end

@inline function _∂rk_∂η(kernel::RK_H0, η::SVector, ξ::SVector)
    # Note: Derivative of spline built with reproducing kernel RK_H0 does not exist at the spline nodes
    t    = η - ξ
    tnrm = norm(t)
    x    = kernel.ε * tnrm
    ∇    = kernel.ε * exp(-x)
    t    = ifelse(x > eps(typeof(x)), t, zeros(t))
    ∇   *= ifelse(x > eps(typeof(x)), inv(tnrm), one(tnrm))
    ∇   *= t
end

@inline function _∂²rk_∂²e(kernel::RK_H2, η::SVector{n}, ξ::SVector{n}, êη::SVector{n}, êξ::SVector{n}) where {n}
    ε     = kernel.ε
    ε²    = ε * ε
    t     = η - ξ
    tnrm  = norm(t)
    x     = ε * tnrm
    ε²e⁻ˣ = ε² * exp(-x)
    ∇²    = ((1 + x) * ε²e⁻ˣ) * (êξ ⋅ êη)
    ∇²   -= (ε² * ε²e⁻ˣ) * (êξ ⋅ t) * (t ⋅ êη)
end

@inline function _∂²rk_∂²e(kernel::RK_H1, η::SVector{n}, ξ::SVector{n}, êη::SVector{n}, êξ::SVector{n}) where {n}
    # Note: Second derivative of spline built with reproducing kernel RK_H1 does not exist at the spline nodes
    ε     = kernel.ε
    ε²    = ε * ε
    t     = η - ξ
    tnrm  = norm(t)
    x     = ε * tnrm
    ε²e⁻ˣ = ε² * exp(-x)
    ∇²    = ε²e⁻ˣ * (êξ ⋅ êη)
    ∇²   -= ifelse(x > eps(typeof(x)), (ε * ε²e⁻ˣ / tnrm) * (êξ ⋅ t) * (t ⋅ êη), zero(∇²))
end

@inline function _∂²rk_∂η∂ξ(kernel::RK_H2, η::SVector{n}, ξ::SVector{n}) where {n}
    ε     = kernel.ε
    ε²    = ε * ε
    t     = η - ξ
    tnrm  = norm(t)
    x     = ε * tnrm
    ε²e⁻ˣ = ε² * exp(-x)
    S     = SMatrix{n,n,typeof(x)}
    ∇²    = S(((1 + x) * ε²e⁻ˣ) * I)
    ∇²   -= ((ε² * ε²e⁻ˣ) * t) * t'
end

@inline function _∂²rk_∂η∂ξ(kernel::RK_H1, η::SVector{n}, ξ::SVector{n}) where {n}
    # Note: Second derivative of spline built with reproducing kernel RK_H1 does not exist at the spline nodes
    ε     = kernel.ε
    ε²    = ε * ε
    t     = η - ξ
    tnrm  = norm(t)
    x     = ε * tnrm
    ε²e⁻ˣ = ε² * exp(-x)
    S     = SMatrix{n,n,typeof(x)}
    ∇²    = S(ε²e⁻ˣ * I)
    ∇²   -= ifelse(x > eps(typeof(x)), ((ε * ε²e⁻ˣ / tnrm) * t) * t', zeros(S))
end


@doc raw"""
    UnweightedParticleEnsemble{T}

An ensemble of `N` particles, each of dimension `n`.
"""
struct MultivariateGaussian{S, TM, TP} <: ParametricRepresentation{Vector{S}, S}
    mean::TM
    cov::TP
    MultivariateGaussian(mean::AbstractVector{S}, cov::AbstractMatrix{S}) where S = new{S, typeof(mean), typeof(cov)}(mean, cov)
end

mean(gauss::MultivariateGaussian) = gauss.mean
cov(gauss::MultivariateGaussian)  = gauss.cov
var(gauss::MultivariateGaussian)  = LinearAlgebra.diag(gauss.cov)

parameter_vector(gauss::MultivariateGaussian) = vcat(mean(gauss), vec(cov(gauss)))


# Test Interpolation

@testset "GaussLegendre Quadrature order 1" begin
    order = 1
    dim = 1
    point_type = :legendre
    qr_C = QuadratureRule{dim, RefLine}(point_type, order);
    qr_F = QuadratureRule{dim, RefCube}(point_type, order);
    gp_C = getpoints(qr_C)
    gp_F = getpoints(qr_F)
    gw_C = getweights(qr_C)
    println(gp_C)
    println(gw_C)
    @test typeof(gp_C) == typeof(gp_F)
    @test size(gp_C) == (order,)
    @test size(gw_C) == (order,)
end

@testset "GaussLegendre Quadrature order 2" begin
    order = 2
    dim = 1
    point_type = :legendre
    qr_C = QuadratureRule{dim, RefLine}(point_type, order);
    qr_F = QuadratureRule{dim, RefCube}(point_type, order);
    gp_C = getpoints(qr_C)
    gp_F = getpoints(qr_F)
    gw_C = getweights(qr_C)
    println(gp_C)
    println(gw_C)
    @test typeof(gp_C) == typeof(gp_F)
    @test size(gp_C) == (order,)
    @test size(gw_C) == (order,)
end

@testset "GaussLegendre Quadrature order 3" begin
    order = 3
    dim = 1
    point_type = :legendre
    qr_C = QuadratureRule{dim, RefLine}(point_type, order);
    qr_F = QuadratureRule{dim, RefCube}(point_type, order);
    gp_C = getpoints(qr_C)
    gp_F = getpoints(qr_F)
    gw_C = getweights(qr_C)
    println(gp_C)
    println(gw_C)
    @test typeof(gp_C) == typeof(gp_F)
    @test size(gp_C) == (order,)
    @test size(gw_C) == (order,)
end

@testset "GaussLegendre Quadrature order 10" begin
    order = 10
    dim = 1
    point_type = :legendre
    qr_C = QuadratureRule{dim, RefLine}(point_type, order);
    qr_F = QuadratureRule{dim, RefCube}(point_type, order);
    gp_C = getpoints(qr_C)
    gp_F = getpoints(qr_F)
    gw_C = getweights(qr_C)
    println(gp_C)
    println(gw_C)
    @test typeof(gp_C) == typeof(gp_F)
    @test size(gp_C) == (order,)
    @test size(gw_C) == (order,)
end

@testset "GaussLobbato Quadrature order 1" begin
    order = 1
    dim = 1
    point_type = :lobatto
    qr_C = QuadratureRule{dim, RefLine}(point_type, order);
    gp_C = getpoints(qr_C)
    gw_C = getweights(qr_C)
    println(gp_C)
    println(gw_C)
    @test size(gp_C) == (order,)
    @test size(gw_C) == (order,)
end

@testset "GaussLobbato  Quadrature order 2" begin
    order = 2
    dim = 1
    point_type = :lobatto
    qr_C = QuadratureRule{dim, RefLine}(point_type, order);
    qr_F = QuadratureRule{dim, RefCube}(point_type, order);
    gp_C = getpoints(qr_C)
    gp_F = getpoints(qr_F)
    gw_C = getweights(qr_C)
    println(gp_C)
    println(gw_C)
    @test typeof(gp_C) == typeof(gp_F)
    @test size(gp_C) == (order,)
    @test size(gw_C) == (order,)
end

@testset "GaussLobbato Quadrature order 3" begin
    order = 3
    dim = 1
    point_type = :lobatto
    qr_C = QuadratureRule{dim, RefLine}(point_type, order);
    qr_F = QuadratureRule{dim, RefCube}(point_type, order);
    gp_C = getpoints(qr_C)
    gp_F = getpoints(qr_F)
    gw_C = getweights(qr_C)
    println(gp_C)
    println(gw_C)
    @test typeof(gp_C) == typeof(gp_F)
    @test size(gp_C) == (order,)
    @test size(gw_C) == (order,)
end

@testset "GaussLegendre Quadrature order 3" begin
    order = 10
    dim = 1
    point_type = :lobatto
    qr_C = QuadratureRule{dim, RefLine}(point_type, order);
    qr_F = QuadratureRule{dim, RefCube}(point_type, order);
    gp_C = getpoints(qr_C)
    gp_F = getpoints(qr_F)
    gw_C = getweights(qr_C)
    println(gp_C)
    println(gw_C)
    @test typeof(gp_C) == typeof(gp_F)
    @test size(gp_C) == (order,)
    @test size(gw_C) == (order,)
end


geom_sigdigits(::Type{Int64})::Int = 12
geom_sigdigits(::Type{Float32})::Int = 6
geom_sigdigits(::Type{Float64})::Int = 12

geom_atol_zero(::Type{Int64})::Float64 = 1e-12
geom_atol_zero(::Type{Float32})::Float32 = 1e-9
geom_atol_zero(::Type{Float64})::Float64 = 1e-12

function geom_round(v::T)::T where {T}
    return isapprox(v, 0, atol = geom_atol_zero(T)) ? T(0) : round(v, sigdigits = geom_sigdigits(T))
end
# function geom_round(cyl::CylindricalPoint{T})::CylindricalPoint{T} where T
#     return Cylindrical{T}(geom_round(cyl.r),geom_round(cyl.φ),geom_round(cyl.z))
# end

function geom_round(pt::CylindricalPoint{T})::CylindricalPoint{T} where {T <: SSDFloat}
    return CylindricalPoint{T}( geom_round(pt.r), geom_round(pt.φ), geom_round(pt.z)  )
end

function geom_round(pt::CartesianPoint{T})::CartesianPoint{T} where {T <: SSDFloat}
    return CartesianPoint{T}( geom_round(pt.x), geom_round(pt.y), geom_round(pt.z)  )
end


export getDivergenceMatrix

function getDivergenceMatrix(mesh::TetraMesh)

if isempty(mesh.Div)

  nf = mesh.nf
	nc = mesh.nc
	T  = mesh.Tetras

  ii = zeros(Int64,4*nc)
  jj = zeros(Int64,4*nc)
  kk = zeros(4*nc)

	ii = [1:nc
			  1:nc
			  1:nc
			  1:nc]

	jj = [findRow(mesh.faces,[T[:,1] T[:,2] T[:,3]])
	      findRow(mesh.faces,[T[:,1] T[:,2] T[:,4]])
			  findRow(mesh.faces,[T[:,1] T[:,3] T[:,4]])
			  findRow(mesh.faces,[T[:,2] T[:,3] T[:,4]])]
	kk = [ones(nc)
			 -ones(nc)
			  ones(nc)
			 -ones(nc)]

	#    cnt = 1
	#    for i=1:nc
	#    ii[cnt] = nc
	# 	jj[cnt] = findRow(mesh.faces,[T[i,1] T[i,2] T[i,3]])[1]
	# 	kk[cnt] = 1
	# 	cnt += 1
	#
	#    ii[cnt] = nc
	# 	jj[cnt] = findRow(mesh.faces,[T[i,1] T[i,2] T[i,4]])[1]
	# 	kk[cnt] = -1
	# 	cnt += 1
	#
	#    ii[cnt] = nc
	# 	jj[cnt] = findRow(mesh.faces,[T[i,1] T[i,3] T[i,4]])[1]
	# 	kk[cnt] = 1
	# 	cnt += 1
	#
	#    ii[cnt] = nc
	# 	jj[cnt] = findRow(mesh.faces,[T[i,2] T[i,3] T[i,4]])[1]
	# 	kk[cnt] = -1
	# 	cnt += 1
	# end

  F = getFaceArea(mesh)

  for k = 1:4*nc
    kk[k] *= F[jj[k]]      #divergence scaled using face sizes
  end

	mesh.Div = sparse(ii, jj, kk, nc, nf)

end

return mesh.Div

end


module SparseGrids

import FastGaussQuadrature: gausshermite
import IterTools: collect, product

include("grids.jl")
include("nested.jl")

# package code goes here
export
	sparsegrid,
	tensorgrid,
	listNdq,
	combvec,
	kpn

end # module


"""
    FullGenerator{T} <: Generator{T}

An Matrix or SparseMatrixCSC representation of the generator of a generator of a DiscretisedFluidQueue.

Higher memory requirements than LazyGenerator (as blocks are duplicated) but much faster matrix arithmetic.
"""
struct FullGenerator{T} <: Generator{T}
    B::Union{Array{Float64,2}, SparseArrays.SparseMatrixCSC{Float64,Int}}
    dq::DiscretisedFluidQueue{T}
end
FullGenerator(B,dq) = FullGenerator{typeof(dq.mesh)}(B,dq)

size(B::FullGenerator) = size(B.B)
# getindex(B::FullGenerator,i) = B.B[i]
getindex(B::FullGenerator,i,j) = B.B[i,j]
getindex(B::FullGenerator,t::Tuple,s::Tuple) = _getindex_from_tuple(B,t,s)
setindex!(B::FullGenerator,x::Float64,i,j) = (B.B[i,j]=x)
sum(B::FullGenerator; kwargs...) = sum(B.B; kwargs...)

fast_mul(A::AbstractArray{<:Real,2}, B::FullGenerator) = A*B.B
fast_mul(B::FullGenerator, A::AbstractArray{<:Real,2}) = fast_mul(A,B)
fast_mul(A::FullGenerator, B::FullGenerator) = A.B*B.B
fast_mul(A::FullGenerator, x::Real) = A.B*x
fast_mul(x::Real,A::FullGenerator) = fast_mul(A, x)

function show(io::IO, mime::MIME"text/plain", B::FullGenerator)
    if VERSION >= v"1.6"
        show(io, mime, B.B)
    else
        show(io, mime, Matrix(B.B))
    end
end

"""
    build_full_generator(dq::DiscretisedFluidQueue; v::Bool = false)

Returns a SparseMatrixCSC generator of a DiscretisedFluidQueue.
"""
function build_full_generator(dq::DiscretisedFluidQueue; v::Bool=false)
    lazy = build_lazy_generator(dq; v=v)
    return build_full_generator(lazy)
end

"""
    build_full_generator(lzB::LazyGenerator)

Returns a SparseMatrixCSC representation of a LazyGenerator.
"""
function build_full_generator(lzB::LazyGenerator)
    B = fast_mul(SparseArrays.SparseMatrixCSC{Float64,Int}(LinearAlgebra.I(size(lzB,1))),lzB)
    return FullGenerator(B,lzB.dq)
end


function stderrs{T <: Model}(::Type{T}, inputs::Vector, weights::Vector)
    se = sqrt.(diag(cov(T, inputs, weights)))
    return se
end


import DSGE: Measurement, get_setting, n_observables, n_shocks_exogenous

"""
```
measurement(m::KrusellSmith{T}, TTT::Matrix{T}, RRR::Matrix{T},
            CCC::Vector{T}) where {T<:AbstractFloat}
```

Assign measurement equation

```
y_t = ZZ*s_t + DD + u_t
```

where

```
Var(ϵ_t) = QQ
Var(u_t) = EE
Cov(ϵ_t, u_t) = 0
```
"""
function measurement(m::KrusellSmith{T},
                     TTT::Matrix{T}, RRR::Matrix{T},
                     CCC::Vector{T}) where {T<:AbstractFloat}
    endo      = m.endogenous_states
    exo       = m.exogenous_shocks
    obs       = m.observables

    _n_model_states = n_model_states(m)

    _n_observables = n_observables(m)
    _n_shocks_exogenous = n_shocks_exogenous(m)

    ZZ = zeros(_n_observables, _n_model_states)
    DD = zeros(_n_observables)
    EE = zeros(_n_observables, _n_observables)
    QQ = zeros(_n_shocks_exogenous, _n_shocks_exogenous)

    # Output in log levels
    ZZ[obs[:obs_gdp], endo[:z′_t]]   .= (m[:Kstar]^m[:α])*(m[:Lstar]^(1.0-m[:α]))
    ZZ[obs[:obs_gdp], endo[:K′_t]]   .= (m[:α]/m[:Kstar])*(m[:Kstar]^m[:α])*(m[:Lstar]^(1.0-m[:α]))

    # Measurement error
    EE[obs[:obs_gdp], obs[:obs_gdp]] = m[:e_y]

    # Variance of innovations
    QQ[exo[:z_sh], exo[:z_sh]] = m[:σ_z]^2

    # Adjustment to DD because measurement equation assumes CCC is the zero vector
    if any(CCC .!= 0)
        DD += ZZ*((UniformScaling(1) - TTT)\CCC)
    end

    return Measurement(ZZ, DD, QQ, EE)
end


using DrWatson
@quickactivate "StatReth"

# %%
using CSV
using DataFrames
using Turing
using Plots
using StatsPlots

include(srcdir("quap.jl"))
include(srcdir("tools.jl"))

# %% 4E1
"""
The first line: yᵢ ~ Normal(μ, σ).
"""

# %% 4E2
"""
There are two parameters, μ and σ.
"""

# %% 4E3
"""
Bayes Theorem: P(A|B) = P(B|A) * P(A) / P(B)
P(μ,σ | y) = P(y | μ,σ) * P(μ,σ) / P(y)
           = pdf(Normal(μ, σ), yᵢ) * pdf(Normal(0, 10), μ) * pdf(Exponential(1)), μ) / P(y)
"""

# %% 4E4
"""
The line μᵢ = α + βxᵢ
"""

# %% 4E5
"""
There are three parameters, α, β, and σ.
"""


include("building_tree.jl")
include("building_weighted_tree.jl")
include("utilities.jl")
using LinearAlgebra

function fusion_points(X, Y, proportion)
  n = length(Y)
  m = n*proportion
  m ÷= 1
  for i in 1:m
    # trouver points d'indice a et b les plus proches de meme classe
    distance = 1
    a, b = 0, 0
    for k in 1:size(X, 1)
      for l in 1:size(X, 1)
        if Y[k] == Y[l] && k != l
          distance_provisoire = norm(X[k,:] - X[l,:]) # distance euclidienne
          if distance_provisoire < distance && distance_provisoire > 0
            distance = distance_provisoire
            a = k
            b = l
          end
        end
      end
    end
    # les fusionner dans X et dans Y
    X[a,:] = 0.5*(X[a,:] + X[b,:])
    # supprimer X[b] et Y[b]
    X = X[1:end .!= b,:]
    splice!(Y, b)
  end
  return X,Y
end

function main()

    # Pour chaque jeu de données
    for dataSetName in ["iris", "seeds", "wine", "haberman", "banknote"]

        print("=== Dataset ", dataSetName)

        # Préparation des données
        include("../data/" * dataSetName * ".txt")
        train, test = train_test_indexes(length(Y))
        X_train = X[train, :] # modifier X_train pour fusionner gamma pourcent des points de meme classe
        Y_train = Y[train] # idem sur Y_train
        X_train, Y_train = fusion_points(X_train, Y_train, 0.2)
        X_test = X[test, :] # mais pas sur les donnees de test !
        Y_test = Y[test]

        println(" (train size ", size(X_train, 1), ", test size ", size(X_test, 1), ", ", size(X_train, 2), ", features count: ", size(X_train, 2), ")")

        # Temps limite de la méthode de résolution en secondes
        time_limit = 30

        # On ajoute ou pas des poids aléatoires aux points/aux clusters
        weigthed = true
        if weigthed
          println("avec poids : ")
        end

        # Pour chaque profondeur considérée
        for D in 2:4

            println("  D = ", D)

            ## 1 - Univarié (séparation sur une seule variable à la fois)
            # Création de l'arbre
            print("    Univarié...  \t")
            if weigthed
              T, obj, resolution_time, gap = build_weighted_tree(X_train, Y_train, D,  multivariate = false, time_limit = time_limit)
            else
              T, obj, resolution_time, gap = build_tree(X_train, Y_train, D,  multivariate = false, time_limit = time_limit)
            end

            # Test de la performance de l'arbre
            print(round(resolution_time, digits = 1), "s\t")
            print("gap ", round(gap, digits = 1), "%\t")
            if T != nothing
                print("Erreurs train/test ", prediction_errors(T,X_train,Y_train))
                print("/", prediction_errors(T,X_test,Y_test), "\t")
            end
            println()

            ## 2 - Multivarié
            print("    Multivarié...\t")
            if weigthed
              T, obj, resolution_time, gap = build_weighted_tree(X_train, Y_train, D, multivariate = true, time_limit = time_limit)
            else
              T, obj, resolution_time, gap = build_tree(X_train, Y_train, D, multivariate = true, time_limit = time_limit)
            end
            print(round(resolution_time, digits = 1), "s\t")
            print("gap ", round(gap, digits = 1), "%\t")
            if T != nothing
                print("Erreurs train/test ", prediction_errors(T,X_train,Y_train))
                print("/", prediction_errors(T,X_test,Y_test), "\t")
            end
            println("\n")
        end
    end
end


using AlgebraicMatrices
using Random
using Test

# Set reproducible random number seed
Random.seed!(0)

# Random rationals
function Base.rand(rng::AbstractRNG,
                   ::Random.SamplerType{Rational{T}}) where {T}
    l = lcm(1:10)
    return Rational{T}(T(rand(rng, (-l):l)) // l)
end

@testset "Algebraic matrices as vector spaces D=$D" for D in 0:3
    T = Rational{BigInt}
    niters = 20
    for iter in 1:niters
        print("\r$iter/$niters...")

        sz = rand(0:10, D)
        x = arand(AlgebraicArray{T,D}, sz)
        y = arand(AlgebraicArray{T,D}, sz)
        z = arand(AlgebraicArray{T,D}, sz)
        n = zero(AlgebraicArray{T,D}, sz)
        a = rand(T)
        b = rand(T)

        @test x isa AbstractArray{T,D}
        @test y isa AbstractArray{T,D}
        @test z isa AbstractArray{T,D}
        @test n isa AbstractArray{T,D}
        @test a isa T
        @test b isa T

        @test size(x) == size(y) == size(z) == size(n)

        @test x == x
        @test x == n || x + x ≠ x

        if D == 1
            @test (x < y) + (x == y) + (x > y) == 1
        end

        @test collect(x) isa Array{T,D}
        @test evaluate(x) isa AbstractArray{T,D}
        @test collect(x) == evaluate(x)

        x′ = map(x -> 2x + 1, x)
        @test x′ == map(x -> 2x + 1, evaluate(x))

        @test (x + y) + z == x + (y + z)
        @test n + x == x
        @test x + n == x
        @test x + y == y + x
        @test one(T) * x == x
        @test x * one(T) == x
        @test a * n == n
        @test n * a == n
        if a ≠ zero(T)
            @test inv(a) * (a * x) == x
        end
        @test +x == x
        @test -x == n - x
        @test x + (-x) == n
        @test (-x) + x == n
        @test a * (x + y) == a * x + a * y
        @test (a + b) * x == a * x + b * x

        @test (x == y) == (evaluate(x) == evaluate(y))
        if D == 1
            @test (x < y) == (evaluate(x) < evaluate(y))
            @test (x > y) == (evaluate(x) > evaluate(y))
        end
        @test iszero(evaluate(n))
        @test evaluate(x + y) == evaluate(x) + evaluate(y)
        @test evaluate(a * x) == a * evaluate(x)
        @test evaluate(-x) == -evaluate(x)

        @test simplify(x) == x
    end
    println("\r", " "^40, "\r")
end

@testset "Algebraic matrices as categories D=$D1" for D1 in 1:3
    T = Rational{BigInt}
    niters = 10
    for iter in 1:niters
        print("\r$iter/$niters...")

        D2 = rand(1:3)
        D3 = rand(1:3)
        sz1 = rand(0:10, D1)
        sz2 = (sz1[end], rand(0:10, D2 - 1)...)
        sz3 = (sz2[end], rand(0:10, D3 - 1)...)
        x = arand(AlgebraicArray{T,D1}, sz1)
        x′ = arand(AlgebraicArray{T,D1}, sz1)
        y = arand(AlgebraicArray{T,D2}, sz2)
        z = arand(AlgebraicArray{T,D3}, sz3)
        e1 = one(AlgebraicArray{T,2}, sz1[1], sz1[1])
        e2 = one(AlgebraicArray{T,2}, sz1[end], sz1[end])
        nx = zero(AlgebraicArray{T,D1}, sz1)
        ny = zero(AlgebraicArray{T,D2}, sz2)
        D12 = D1 + D2 - 2
        sz12 = (sz1[1:(end - 1)]..., sz2[2:end]...)
        nxy = zero(AlgebraicArray{T,D12}, sz12)
        a = rand(T)

        if ndims(y) ≥ 2
            @test (x * y) * z == x * (y * z)
        end
        @test e1 * x == x
        @test x * e2 == x
        @test nx * y == nxy
        @test x * ny == nxy
        @test (a * x) * y == a * (x * y)
        @test x * (a * y) == a * (x * y)
        @test (x + x′) * y == x * y + x′ * y

        @test isone(evaluate(e1))

        x′ = evaluate(x)
        y′ = evaluate(y)
        sxy = (size(x)[1:(end - 1)]..., size(y)[2:end]...)
        xy′ = zeros(T, sxy)
        for i in CartesianIndices(size(x)[1:(end - 1)])
            for j in CartesianIndices(size(y)[2:end])
                for k in 1:size(y, 1)
                    xy′[i, j] += x′[i, k] * y′[k, j]
                end
            end
        end
        # @test evaluate(x * y) == evaluate(x) * evaluate(y)
        @test evaluate(x * y) == xy′
    end
    println("\r", " "^40, "\r")
end


### A Pluto.jl notebook ###
# v0.12.3

using Markdown
using InteractiveUtils

# ╔═╡ 92be606a-0998-11eb-053e-cb42246e0598
using Pkg, DrWatson

# ╔═╡ 92be94e0-0998-11eb-03fc-4b2293b5e32a
begin
	@quickactivate "StatisticalRethinkingTuring"
	using Turing
	using StatisticalRethinking
end

# ╔═╡ 8cc034b8-0998-11eb-2c71-578b14536d5d
md"## Clip-05-28-44.jl"

# ╔═╡ 92c90f38-0998-11eb-086f-4b4900891b2f
md"### snippets 5.28 - 5.32"

# ╔═╡ 94042790-099b-11eb-1a63-6dc1af913def
md"##### Notice that `missing` values are dropped!"

# ╔═╡ 92d87388-0998-11eb-1881-11abc4118c34
begin
	df = CSV.read(sr_datadir("milk.csv"), DataFrame; missingstring = "NA")
	dropmissing!(df, [:kcal_per_g, :neocortex_perc, :mass])
	df.K = zscore(df.kcal_per_g)
	df.N = zscore(df.neocortex_perc)
	df.M = zscore(log.(df.mass))
end;

# ╔═╡ 55142898-0999-11eb-1cff-6d82e287b64a
md"### snippet 5.33"

# ╔═╡ fdde27d4-099a-11eb-2156-9b15bba277f2
@model function m5_5_draft(N, K)
    a ~ Normal(0, 1)
    bN ~ Normal(0, 1)
    σ ~ Exponential(1)
    μ = lin(a, bN, N)
    K ~ MvNormal(μ, σ)
end

# ╔═╡ fdde6780-099a-11eb-0fd8-b597b7b34d15
begin
	m5_5_draftt = m5_5_draft(df.N, df.K)
	quap5_5_draftt = quap(m5_5_draftt)
end

# ╔═╡ e17b5820-099d-11eb-242b-c1f8c9e2fc14
md"### snippet 5.34"

# ╔═╡ d03a085c-09a2-11eb-3ba3-d18810ca8ab5
xseq = -2:0.1:2

# ╔═╡ e4d38704-099d-11eb-0c30-9d7f0556f5b1
begin
	prior = sample(m5_5_draftt, Prior(), 50) |> DataFrame
	μ = lin(prior.a', prior.bN', xseq)

	fig1 = plot(
		xlab="neocortex percent (std)", ylab="kcal per g (std)",
		title="a ~ Normal(0, 1)\nbN ~ Normal(0, 1)",
		legend = false, ylims = extrema(xseq))
	for c in eachcol(μ)
		plot!(xseq, c, color = :black, alpha = 0.3)
	end
	plot!()
end

# ╔═╡ e4d3b814-099d-11eb-2d2d-bbf106ef33b6
md"### snippet 5.35"

# ╔═╡ e4d468b6-099d-11eb-2b83-3b2d265a4fad
@model function m5_5(N, K)
    a ~ Normal(0, 0.2)
    bN ~ Normal(0, 0.5)
    σ ~ Exponential(1)
    μ = lin(a, bN, N)
    K ~ MvNormal(μ, σ)
end

# ╔═╡ e4ea98f4-099d-11eb-0ca5-bd416c74dbab
begin
	m5_5t = m5_5(df.N, df.K)
	chns5_5t = sample(m5_5t, NUTS(0.65), 2000)
end

# ╔═╡ 82cf498e-09a3-11eb-1edf-4f32d126c20e
begin
	post = sample(m5_5_draftt, NUTS(0.65), 200) |> DataFrame
	μ_post = lin(post.a', post.bN', xseq)

	fig2 = plot(
		xlab="neocortex percent (std)", ylab="kcal per g (std)",
		title="a ~ Normal(0, 0.2)\nbN ~ Normal(0, 0.5)",
		legend = false, ylims = extrema(xseq))
	for c in eachcol(μ_post)
		plot!(xseq, c, color = :black, alpha = 0.3)
	end
	plot(fig1, fig2, layout=(1, 2))
end

# ╔═╡ e4ebf03e-099d-11eb-16e3-93fea8ed8f8a
md"### snippet 5.36"

# ╔═╡ e4f60d74-099d-11eb-324e-31bef9d5181f
begin
	post5_5t = DataFrame(Array(chns5_5t), names(chns5_5t, [:parameters]))
	Text(precis(post5_5t; io=String))
end

# ╔═╡ 26000e2c-099f-11eb-05cc-bb0bede678d1
begin
	quap5_5t = quap(m5_5t)
	Text(precis(quap5_5t; io=String))
end

# ╔═╡ 4010b8b4-09a1-11eb-173a-2bb20afc882d
md"### snippet 5.37"

# ╔═╡ 44a91bfa-09a1-11eb-3bd9-af4110eab915
begin
	post1 = DataFrame(rand(quap5_5t.distr, 1000)', quap5_5t.params)
	μ1 = lin(post1.a', post1.bN', xseq) |> meanlowerupper

	fig3 = plot(xlab="neocortex percent (std)", ylab="kcal per gram (std)")
	scatter!(df.N, df.K, legend = false)
	plot!(xseq, μ1.mean, ribbon = (μ1.mean .- μ1.lower, μ1.upper .- μ1.mean))
end

# ╔═╡ 44a9500c-09a1-11eb-36e3-750c47541694
md"### snippet 5.38"

# ╔═╡ 44aa2266-09a1-11eb-3d17-5d3f2da14256
@model function m5_6(M, K)
    a ~ Normal(0, 0.2)
    bM ~ Normal(0, 0.5)
    σ ~ Exponential(1)
    μ = lin(a, bM, M)
    K ~ MvNormal(μ, σ)
end

# ╔═╡ 44bf965a-09a1-11eb-2d3b-95b5e58e1edd
begin
	m5_6t = m5_6(df.M, df.K)
	Text(precis(quap(m5_6t); io=String))
end

# ╔═╡ 753a9688-09a4-11eb-1bce-654aeb39c454
begin
	quap5_6t = quap(m5_6t)
	post2 = DataFrame(rand(quap5_6t.distr, 1000)', quap5_6t.params)
	μ2 = lin(post2.a', post2.bM', xseq) |> meanlowerupper

	fig4 = plot(xlab="log body mass (std)", ylab="kcal per gram (std)")
	scatter!(df.N, df.K, legend = false)
	plot!(xseq, μ2.mean, ribbon = (μ2.mean .- μ2.lower, μ2.upper .- μ2.mean))
end

# ╔═╡ 44c06a94-09a1-11eb-243e-e94c4ea7b657
md"### snippet 5.39"

# ╔═╡ 44d48ee8-09a1-11eb-0dd6-c5aa78fc84ec
@model function m5_7(M, N, K)
    a ~ Normal(0, 0.2)
    bM ~ Normal(0, 0.5)
    bN ~ Normal(0, 0.5)
    σ ~ Exponential(1)
    μ = lin(a, bM, M, bN, N)
    K ~ MvNormal(μ, σ)
end

# ╔═╡ 44d79f14-09a1-11eb-2122-916596c22454
begin
	m5_7t = m5_7(df.M, df.N, df.K)
	quap5_7t = quap(m5_7t)
	post5_7t = DataFrame(rand(quap5_7t.distr, 1000)', quap5_7t.params)
	Text(precis(post5_7t; io=String))
end

# ╔═╡ 2aa59598-09a2-11eb-0854-5b7b1a9aa1b2
md"### snippet 5.40"

# ╔═╡ 99bfbc84-09a1-11eb-01a6-35c26e589d7b
begin
	quap_res, fig = plotcoef([m5_5t, m5_6t, m5_7t], [:a, :bN, :bM, :σ])
	plot(fig)
end

# ╔═╡ 1982c47a-09a2-11eb-357c-1b16ba7aacb1
quap_res

# ╔═╡ 59c2eef4-09a5-11eb-2a39-b75a6de102ad
md"### snippet 5.41"

# ╔═╡ 5e7f6ba2-09a5-11eb-1168-7bbd0f34e93b
begin
	mu2 = lin(post5_7t.a', post5_7t.bN', xseq, post5_7t.bM', zeros(length(xseq))) |> meanlowerupper
	fig5 = plot(xlab="neocortex percent (std)", ylab="kcal per gram (std)",
		title="Counterfactual (M=0)", leg=false)
	plot!(xseq, mu2.mean, ribbon = (mu2.upper .- mu2.mean, mu2.mean .- mu2.lower))

	mu3 = lin(post5_7t.a', post5_7t.bM', xseq, post5_7t.bN', zeros(length(xseq))) |> meanlowerupper
	fig6 = plot(xlab="neocortex percent (std)", ylab="log body mass (std)",
		title="Counterfactual (N=0)", leg=false)
	plot!(xseq, mu3.mean, ribbon = (mu3.upper .- mu3.mean, mu3.mean .- mu3.lower))

	plot(fig3, fig4, fig5, fig6, layout=(2,2))
end

# ╔═╡ 405a080e-0995-11eb-01c0-3ff35fe8c925
md"### snippet 5.42"

# ╔═╡ 99aa57b4-09a7-11eb-2080-01a888bd0999
n = 100

# ╔═╡ 99aa96fa-09a7-11eb-17c6-bb516e055371
begin
	# M -> K <- N
	# M -> N
	M1 = randn(n)
	N1 = rand.(Normal.(M1))
	K1 = rand.(Normal.(N1 .- M1))
	d_sim1 = DataFrame((; K1, N1, M1))
	Text(precis(d_sim1; io=String))
end

# ╔═╡ 99b7614e-09a7-11eb-21b4-716994e5df9d
md"### snippet 5.43"

# ╔═╡ 99ba1a76-09a7-11eb-1729-339be9f12454
begin
	# M -> K <- N
	# N -> M
	N2 = randn(n)
	M2 = rand.(Normal.(N2))
	K2 = rand.(Normal.(N2 .- M2))
	d_sim2 = DataFrame((; K2, N2, M2))
	Text(precis(d_sim2; io=String))
end

# ╔═╡ 99d1457a-09a7-11eb-18c8-bb41b89014f3
begin
	# M -> K <- N
	# M -> U <- N
	U3 = randn(n)
	N3 = rand.(Normal.(U3))
	M3 = rand.(Normal.(U3))
	K3 = rand.(Normal.(N3 .- M3))
	d_sim3 = DataFrame((; K3, N3, M3))
	Text(precis(d_sim3; io=String))
end

# ╔═╡ 99d510ba-09a7-11eb-0fcf-97d4242bb925
md"### snippet 5.44"

# ╔═╡ 99e1a37a-09a7-11eb-3839-87157105995e
md"##### TODO working with DAGs."

# ╔═╡ 99eecdfe-09a7-11eb-12c8-fd7d822403c1
md"## End of clip-05-28-44.jl"

# ╔═╡ Cell order:
# ╟─8cc034b8-0998-11eb-2c71-578b14536d5d
# ╠═92be606a-0998-11eb-053e-cb42246e0598
# ╠═92be94e0-0998-11eb-03fc-4b2293b5e32a
# ╟─92c90f38-0998-11eb-086f-4b4900891b2f
# ╟─94042790-099b-11eb-1a63-6dc1af913def
# ╠═92d87388-0998-11eb-1881-11abc4118c34
# ╟─55142898-0999-11eb-1cff-6d82e287b64a
# ╠═fdde27d4-099a-11eb-2156-9b15bba277f2
# ╠═fdde6780-099a-11eb-0fd8-b597b7b34d15
# ╟─e17b5820-099d-11eb-242b-c1f8c9e2fc14
# ╠═d03a085c-09a2-11eb-3ba3-d18810ca8ab5
# ╠═e4d38704-099d-11eb-0c30-9d7f0556f5b1
# ╟─e4d3b814-099d-11eb-2d2d-bbf106ef33b6
# ╠═e4d468b6-099d-11eb-2b83-3b2d265a4fad
# ╠═e4ea98f4-099d-11eb-0ca5-bd416c74dbab
# ╠═82cf498e-09a3-11eb-1edf-4f32d126c20e
# ╟─e4ebf03e-099d-11eb-16e3-93fea8ed8f8a
# ╠═e4f60d74-099d-11eb-324e-31bef9d5181f
# ╠═26000e2c-099f-11eb-05cc-bb0bede678d1
# ╟─4010b8b4-09a1-11eb-173a-2bb20afc882d
# ╠═44a91bfa-09a1-11eb-3bd9-af4110eab915
# ╟─44a9500c-09a1-11eb-36e3-750c47541694
# ╠═44aa2266-09a1-11eb-3d17-5d3f2da14256
# ╠═44bf965a-09a1-11eb-2d3b-95b5e58e1edd
# ╠═753a9688-09a4-11eb-1bce-654aeb39c454
# ╟─44c06a94-09a1-11eb-243e-e94c4ea7b657
# ╠═44d48ee8-09a1-11eb-0dd6-c5aa78fc84ec
# ╠═44d79f14-09a1-11eb-2122-916596c22454
# ╟─2aa59598-09a2-11eb-0854-5b7b1a9aa1b2
# ╠═99bfbc84-09a1-11eb-01a6-35c26e589d7b
# ╠═1982c47a-09a2-11eb-357c-1b16ba7aacb1
# ╟─59c2eef4-09a5-11eb-2a39-b75a6de102ad
# ╠═5e7f6ba2-09a5-11eb-1168-7bbd0f34e93b
# ╟─405a080e-0995-11eb-01c0-3ff35fe8c925
# ╠═99aa57b4-09a7-11eb-2080-01a888bd0999
# ╠═99aa96fa-09a7-11eb-17c6-bb516e055371
# ╟─99b7614e-09a7-11eb-21b4-716994e5df9d
# ╠═99ba1a76-09a7-11eb-1729-339be9f12454
# ╠═99d1457a-09a7-11eb-18c8-bb41b89014f3
# ╟─99d510ba-09a7-11eb-0fcf-97d4242bb925
# ╟─99e1a37a-09a7-11eb-3839-87157105995e
# ╟─99eecdfe-09a7-11eb-12c8-fd7d822403c1


__precompile__()

module Baseline
using Flux
abstract type AbstractBaseline end

struct Linear{DT<:AbstractFloat} <: AbstractBaseline
   nfeat::Int
   T::Int
   numT::Int
   reg_coeff::DT
   coeff::Vector{DT} # states x N

   # workspace
   featnd::Array{DT,3}
   A::Matrix{DT}
   bline::Vector{DT}
   function Linear{DT}(reg::AbstractFloat, ns::Integer, T::Integer, numT::Integer) where DT<:AbstractFloat
      nfeat = 2*ns+3+1
      featnd=zeros(DT, nfeat, T, numT)
      A = zeros(DT, nfeat, nfeat)
      bline = zeros(DT, T*numT)
      return new(nfeat, T, numT, reg, zeros(DT, nfeat),
                 featnd, A, bline)
   end
end

struct Quadratic{DT<:AbstractFloat} <: AbstractBaseline
   nfeat::Int
   T::Int
   numT::Int
   reg_coeff::DT
   coeff::Vector{DT} # states x N

   # workspace
   featnd::Array{DT,3}
   A::Matrix{DT}
   bline::Vector{DT}
   function Quadratic{DT}(reg::AbstractFloat, ns::Integer, T::Integer, numT::Integer) where DT<:AbstractFloat
      nfeat = ns+convert(Int,ns*(ns+1)/2)+ns+ns+4+1
      featnd=zeros(DT, nfeat, T, numT)
      gettimefeatures!(featnd)

      A = zeros(DT, nfeat, nfeat)
      bline = zeros(DT, T*numT)
      return new(nfeat, T, numT, reg, zeros(DT, nfeat),
                 featnd, A, bline)
   end
end

struct NN{DT<:AbstractFloat} <: AbstractBaseline
   nhidden::Int
   nfeat::Int
   T::Int
   numT::Int

   loss::Function
   opt::Function

   #workspace
   featnd::Array{DT,2}
   m::Any
   #w::Array{Any}
   bline::Vector{DT}
   function NN{DT}(ns::Integer, T::Integer, numT::Integer, nhidden::Integer) where DT<:AbstractFloat
      nfeat = 1*ns+3+1
      featnd=zeros(DT, nfeat, T*numT)
      m = Chain(Dense(nfeat, nhidden, tanh),
                Dense(nhidden, nhidden, tanh),
                Dense(nhidden, 1)) |> gpu
      bline = zeros(DT, T*numT)
      return new(nhidden, nfeat, T, numT,
                 (a, b) -> sum((m(a) - b).^2)/length(b),
                 ADAM(params(m)), #, 0.0001),
                 featnd, m, bline)
   end
end

######################################################################## common

function gettimefeatures!{DT<:AbstractFloat}(feat::Array{DT,3})
   T    = size(feat, 2)
   numT = size(feat, 3)
   al = convert(Array{DT}, linspace(1, T, T)') / DT(T)
   for t=1:numT
      feat[end-4,:,t] = al
      feat[end-3,:,t] = al.^2
      feat[end-2,:,t] = al.^3
      feat[end-1,:,t] = al.^4
      feat[end,:,t]   = DT(1.0)
   end
end

function getlinearfeatures(obs::Matrix{Float64})
   T = size(obs,2)
   # TODO add clipping
   o = max.(obs, -10.0)
   o = min.(o, 10.0)
   al = linspace(1, T, T)' / T
   feat = vcat(o, o.^2, al, al.^2, al.^3, ones(1, T))
end

function getquadfeatures!{DT<:AbstractFloat}(feat::Array{DT,3},
                                             obs::Matrix{DT},
                                             t::Integer)
   n = size(obs, 1)
   T = size(feat, 2)

   feat[1:n,:,t] = obs[:,(t-1)*T+1:t*T] #implicite conversion between 32, 64
   feat[1:n,:,t] = max.(feat[1:n,:,t], DT(-10.0))
   feat[1:n,:,t] = min.(feat[1:n,:,t], DT(10.0))
   feat[1:n,:,t] /= DT(10.0)

   qfsize = convert(Int, n*(n+1)/2)
   qf = view(feat, (n+1):(n+qfsize) , :, t) #Matrix{Float64}(qfsize, T)
   k = 1
   @inbounds for i=1:n
      for j=i:n
         for x=1:T
            qf[k,x] = feat[i,x,t].*feat[j,x,t]
         end
         k += 1
      end
   end

   i0 = n+qfsize
   i1 = i0+n
   @inbounds for j = 1:n
      for i=1:T
         feat[i0+j,i,t] = sin(pi*feat[j,i,t])
         feat[i1+j,i,t] = cos(pi*feat[j,i,t])
      end
   end
end

function getNNfeatures(obs::Matrix{Float64})
   T = size(obs,2)
   # TODO add clipping
   o = max.(obs, -10.0)
   o = min.(o, 10.0)
   # al = linspace(1, T, T)' / 100.0
   al = linspace(1, T, T)' / 1000.0
   # feat = vcat(o, o.^2, al, al.^2, al.^3, ones(1, T))
   feat = vcat(o, al, al.^2, al.^3, ones(1, T))
end

function prefit!{DT<:AbstractFloat}(b::AbstractBaseline, returns::Vector{DT})
   featmat = reshape(b.featnd, b.nfeat, b.numT*b.T)
   #b.A[:] = (featmat*featmat')
   A_mul_Bt!(b.A, featmat, featmat)

   return featmat*returns
end

# returns: T*numT
function fit!{DT<:AbstractFloat}(b::AbstractBaseline, returns::Vector{DT})

   target = prefit!(b, returns) # sets A matrix, returns target

   for i=1:b.nfeat # set diag
      b.A[i,i] += b.reg_coeff
   end
   #b.coeff[:] = b.A\(target)
   if isposdef(b.A)
      A_ldiv_B!(b.coeff, cholfact!(b.A), target) # needs A to be factorized
   else
      A_ldiv_B!(b.coeff, lufact!(convert(Array{Float64}, b.A)), target) # needs A to be factorized
   end
end

function fit!{DT<:AbstractFloat}(b::NN, returns::Vector{DT})
   nsamples = length(returns)
   const nbatch = 50

   X = b.featnd |> gpu
   Y = returns |> gpu

   rndidx = Flux.chunk(shuffle(1:nsamples), nbatch)

   datas = [ (X[:,i], Y[i]') for i in rndidx ]

   opt = ADAM(params(b.m))
   evalcb() = @show(b.loss(b.featnd, returns'))

   BLAS.set_num_threads(Threads.nthreads())
   for i=1:3
      Flux.train!(b.loss, datas, b.opt)#, cb = Flux.throttle(evalcb, 1))
   end

end

######################################################################## linear

function predict!{T<:AbstractFloat}(b::Linear, obs::Matrix{T})
   nthread = min.(b.numT, Threads.nthreads()) # cant have more threads than data
   Threads.@threads for tid=1:nthread
      #for tid=1:nthread
      thread_range = Distributed.splitrange(b.numT, nthread)[tid]

      for t=thread_range
         b.featnd[:,:,t] = getlinearfeatures(obs[:,(t-1)*b.T+1:t*b.T])
      end
   end
   b.bline[:] = reshape(b.featnd, b.nfeat, b.numT*b.T)'*b.coeff
end

##################################################################### quadratic
# TODO this is hacky and this whole module is a disgrace
function predict{DT<:AbstractFloat}(b::Quadratic, obs::Matrix{DT})
   nsamples = size(obs, 2)
   featnd=zeros(DT, b.nfeat, 1, nsamples)
   gettimefeatures!(featnd)
   for t=1:nsamples
      getquadfeatures!(featnd, obs, t)
   end
   values = reshape(featnd, b.nfeat, nsamples)'*b.coeff
   return values
end

function predict!{DT<:AbstractFloat}(b::Quadratic, obs::Matrix{DT})
   nthread = min.(b.numT, Threads.nthreads()) # cant have more threads than data
   Threads.@threads for tid=1:nthread
      #for tid=1:nthread
      thread_range = Distributed.splitrange(b.numT, nthread)[tid]

      for t=thread_range
         getquadfeatures!(b.featnd, obs, t)
      end
   end
   b.bline[:] = reshape(b.featnd, b.nfeat, b.numT*b.T)'*b.coeff
end

##################################################################### Neural Network

function predict!{T<:AbstractFloat}(b::NN, obs::Matrix{T})
   b.featnd[:,:] = getNNfeatures(obs)
   b.bline .= b.m(b.featnd).data[1,:]
end

end # module