#----------------------------------------------------------------------------
"""
    OlsPs(y,x,ExciseIt=false,UnExciseIt=false,SkipCovbIt=false)

Do OLS of y on x, for one dependent variable or SURE with same regressors

# Usage
`(b,res,yhat,Covb,R2a,T) = OlsPs(y,x,ExciseIt,UnExciseIt,SkipCovbIt)`

# Input
- `y::Array`:         Tx1 or Txn matrix of the dependent variables
- `x::Array`:         Txk matrix of regressors (including deterministic ones)
- `ExciseIt::Bool`:   (optional) bool, true to doing excise([y,x])
- `UnExciseIt::Bool`: (optional) bool, true to put residuals and fitted values in Txn matrix
- `SkipCovbIt::Bool`: (optional) bool, true to skip calculation of Covb (sensitive inversion)

# Output
- `b::Array`:     kxn matrix, regression coefficients
- `res::Array`:   Tx1 or Txn matrix, residuals y - yhat
- `yhat::Array`:  Tx1 or Txn matrix, fitted values x*b
- `Covb::Array`:  matrix, covariance matrix of vec(b) = [beq1;beq2;...]
- `R2a::Array`:   1xn vector, R2 values
- `Ty::Int`:      scalar, number of obs (after excise)

# Requires
- FindNoNaNPs

"""
function OlsPs(y,x,ExciseIt=false,UnExciseIt=false,SkipCovbIt=false)

  (T,n) = (size(y,1),size(y,2))

  if ExciseIt
     vvNoNaNR = FindNNPs(y,x)                #find rows with no NaNs
    (y,x)     = (y[vvNoNaNR,:],x[vvNoNaNR,:])
  end

  (Ty,Tx,k) = (size(y,1),size(x,1),size(x,2))

  if Tx != Ty
    error("y and x must have same number of observations")
  end
  if any(isnan.([y x]))
    error("NaN in x or y")
  end
  if ndims(x) == 1               #T vector to (T,1) matrix, could also do reshape(x,T,1)
    x = x[:,:]
  end

  if Ty >= k
    b = x\y
    #b = (x'x)\(x'y)
    #ndims(x) == 1 ? yhat2  = x*b[1] : yhat2  = x*b  #if? then: else. to handle x is vector
    yhat2  = x*b
    res2   = y - yhat2
    Covres = cov(res2)*(Ty-1)/Ty
    if SkipCovbIt
      Covb = NaN
    else
      Covb   = kron(Covres,inv(x'x))
    end
    R2a = 1 .- var(res2,dims=1)./var(y,dims=1)
  else
    (b,yhat2,res2,Covb,R2a) = (NaN,NaN,NaN,NaN,NaN)
  end

  if UnExciseIt && ExciseIt   #put fitted value and residual in Tyxn matrix
    (yhat,res)       = [fill(NaN,(T,n)) for i=1:2]
    yhat[vvNoNaNR,:] = yhat2
    res[vvNoNaNR,:]  = res2
  else
    yhat = yhat2
    res  = res2
  end

  if n == 1                #if only one regression
    R2a = R2a[1]
  end

  return b,res,yhat,Covb,R2a,Ty

end
#------------------------------------------------------------------------------


# Score-Based Generative Modeling
#
# Score-Based Generative Modeling through Stochastic Differential Equations.
# Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma,
# Abhishek Kumar, Stefano Ermon, and Ben Poole.
# https://arxiv.org/pdf/2011.13456.pdf

using MLDatasets
using Flux
using Flux: @functor, chunk, params
using Flux.Data: DataLoader
using Parameters: @with_kw
using BSON
using CUDA
using Images
using Logging: with_logger
using ProgressMeter: Progress, next!
using TensorBoardLogger: TBLogger, tb_overwrite
using Random
using Statistics

"""
Projection of Gaussian Noise onto a time vector.

# Notes
This layer will help embed our random times onto the frequency domain. \n
W is not trainable and is sampled once upon construction - see assertions below.

# References
paper-  https://arxiv.org/abs/2006.10739
"""
function GaussianFourierProjection(embed_dim, scale)
    # Instantiate W once
    W = randn(Float32, embed_dim ÷ 2) .* scale
    # Return a function that always references the same W
    function GaussFourierProject(t)
        t_proj = t' .* W * Float32(2π)
        [sin.(t_proj); cos.(t_proj)]
    end
end

"""
Helper function that computes the *standard deviation* of 𝒫₀ₜ(𝘹(𝘵)|𝘹(0)).

# Notes
Derived from the Stochastic Differential Equation (SDE):    \n
                𝘥𝘹 = σᵗ𝘥𝘸,      𝘵 ∈ [0, 1]                   \n

We use properties of SDEs to analytically solve for the stddev
at time t conditioned on the data distribution. \n

We will be using this all over the codebase for computing our model's loss,
scaling our network output, and even sampling new images!
"""
marginal_prob_std(t, sigma=25.0f0) = sqrt.((sigma .^ (2t) .- 1.0f0) ./ 2.0f0 ./ log(sigma))

"""
Create a UNet architecture as a backbone to a diffusion model. \n

# Notes
Images stored in WHCN (width, height, channels, batch) order. \n
In our case, MNIST comes in as (28, 28, 1, batch). \n

# References
paper-  https://arxiv.org/abs/1505.04597
"""
struct UNet
    layers::NamedTuple
end

"""
User Facing API for UNet architecture.
"""
function UNet(channels=[32, 64, 128, 256], embed_dim=256, scale=30.0f0)
    return UNet((
        gaussfourierproj=GaussianFourierProjection(embed_dim, scale),
        linear=Dense(embed_dim, embed_dim, swish),
        # Encoding
        conv1=Conv((3, 3), 1 => channels[1], stride=1, bias=false),
        dense1=Dense(embed_dim, channels[1]),
        gnorm1=GroupNorm(channels[1], 4, swish),
        conv2=Conv((3, 3), channels[1] => channels[2], stride=2, bias=false),
        dense2=Dense(embed_dim, channels[2]),
        gnorm2=GroupNorm(channels[2], 32, swish),
        conv3=Conv((3, 3), channels[2] => channels[3], stride=2, bias=false),
        dense3=Dense(embed_dim, channels[3]),
        gnorm3=GroupNorm(channels[3], 32, swish),
        conv4=Conv((3, 3), channels[3] => channels[4], stride=2, bias=false),
        dense4=Dense(embed_dim, channels[4]),
        gnorm4=GroupNorm(channels[4], 32, swish),
        # Decoding
        tconv4=ConvTranspose((3, 3), channels[4] => channels[3], stride=2, bias=false),
        dense5=Dense(embed_dim, channels[3]),
        tgnorm4=GroupNorm(channels[3], 32, swish),
        tconv3=ConvTranspose((3, 3), channels[3] + channels[3] => channels[2], pad=(0, -1, 0, -1), stride=2, bias=false),
        dense6=Dense(embed_dim, channels[2]),
        tgnorm3=GroupNorm(channels[2], 32, swish),
        tconv2=ConvTranspose((3, 3), channels[2] + channels[2] => channels[1], pad=(0, -1, 0, -1), stride=2, bias=false),
        dense7=Dense(embed_dim, channels[1]),
        tgnorm2=GroupNorm(channels[1], 32, swish),
        tconv1=ConvTranspose((3, 3), channels[1] + channels[1] => 1, stride=1, bias=false),
    ))
end

@functor UNet

"""
Helper function that adds `dims` dimensions to the front of a `AbstractVecOrMat`.
Similar in spirit to TensorFlow's `expand_dims` function.

# References:
https://www.tensorflow.org/api_docs/python/tf/expand_dims
"""
expand_dims(x::AbstractVecOrMat, dims::Int=2) = reshape(x, (ntuple(i -> 1, dims)..., size(x)...))

"""
Makes the UNet struct callable and shows an example of a "Functional" API for modeling in Flux. \n
"""
function (unet::UNet)(x, t)
    # Embedding
    embed = unet.layers.gaussfourierproj(t)
    embed = unet.layers.linear(embed)
    # Encoder
    h1 = unet.layers.conv1(x)
    h1 = h1 .+ expand_dims(unet.layers.dense1(embed), 2)
    h1 = unet.layers.gnorm1(h1)
    h2 = unet.layers.conv2(h1)
    h2 = h2 .+ expand_dims(unet.layers.dense2(embed), 2)
    h2 = unet.layers.gnorm2(h2)
    h3 = unet.layers.conv3(h2)
    h3 = h3 .+ expand_dims(unet.layers.dense3(embed), 2)
    h3 = unet.layers.gnorm3(h3)
    h4 = unet.layers.conv4(h3)
    h4 = h4 .+ expand_dims(unet.layers.dense4(embed), 2)
    h4 = unet.layers.gnorm4(h4)
    # Decoder
    h = unet.layers.tconv4(h4)
    h = h .+ expand_dims(unet.layers.dense5(embed), 2)
    h = unet.layers.tgnorm4(h)
    h = unet.layers.tconv3(cat(h, h3; dims=3))
    h = h .+ expand_dims(unet.layers.dense6(embed), 2)
    h = unet.layers.tgnorm3(h)
    h = unet.layers.tconv2(cat(h, h2, dims=3))
    h = h .+ expand_dims(unet.layers.dense7(embed), 2)
    h = unet.layers.tgnorm2(h)
    h = unet.layers.tconv1(cat(h, h1, dims=3))
    # Scaling Factor
    h ./ expand_dims(marginal_prob_std(t), 3)
end

"""
Model loss following the denoising score matching objectives:

# Notes
Denoising score matching objective:
```julia
min wrt. θ (
    𝔼 wrt. 𝘵 ∼ 𝒰(0, 𝘛)[
        λ(𝘵) * 𝔼 wrt. 𝘹(0) ∼ 𝒫₀(𝘹) [
            𝔼 wrt. 𝘹(t) ∼ 𝒫₀ₜ(𝘹(𝘵)|𝘹(0)) [
                (||𝘚₀(𝘹(𝘵), 𝘵) - ∇ log [𝒫₀ₜ(𝘹(𝘵) | 𝘹(0))] ||₂)²
            ]
        ]
    ]
)
```
Where 𝒫₀ₜ(𝘹(𝘵) | 𝘹(0)) and λ(𝘵), are available analytically and
𝘚₀(𝘹(𝘵), 𝘵) is estimated by a U-Net architecture.

# References:
http://www.iro.umontreal.ca/~vincentp/Publications/smdae_techreport.pdf \n
https://yang-song.github.io/blog/2021/score/#estimating-the-reverse-sde-with-score-based-models-and-score-matching \n
https://yang-song.github.io/blog/2019/ssm/
"""
function model_loss(model, x, ϵ=1.0f-5)
    batch_size = size(x)[end]
    # (batch) of random times to approximate 𝔼[⋅] wrt. 𝘪 ∼ 𝒰(0, 𝘛)
    random_t = rand!(similar(x, batch_size)) .* (1.0f0 - ϵ) .+ ϵ
    # (batch) of perturbations to approximate 𝔼[⋅] wrt. 𝘹(0) ∼ 𝒫₀(𝘹)
    z = randn!(similar(x))
    std = expand_dims(marginal_prob_std(random_t), 3)
    # (batch) of perturbed 𝘹(𝘵)'s to approximate 𝔼 wrt. 𝘹(t) ∼ 𝒫₀ₜ(𝘹(𝘵)|𝘹(0))
    perturbed_x = x + z .* std
    # 𝘚₀(𝘹(𝘵), 𝘵)
    score = model(perturbed_x, random_t)
    # mean over batches
    mean(
        # L₂ norm over WHC dimensions
        sum((score .* std + z) .^ 2; dims=1:(ndims(x) - 1))
    )
end

"""
Helper function that loads MNIST images and returns loader.
"""
function get_data(batch_size)
    xtrain, ytrain = MLDatasets.MNIST(:train)[:]
    xtrain = reshape(xtrain, 28, 28, 1, :)
    DataLoader((xtrain, ytrain), batchsize=batch_size, shuffle=true)
end

"""
Helper function from DrWatson.jl to convert a struct to a dict
"""
function struct2dict(::Type{DT}, s) where {DT<:AbstractDict}
    DT(x => getfield(s, x) for x in fieldnames(typeof(s)))
end
struct2dict(s) = struct2dict(Dict, s)

# arguments for the `train` function
@with_kw mutable struct Args
    η = 1e-4                                        # learning rate
    batch_size = 32                                 # batch size
    epochs = 50                                     # number of epochs
    seed = 1                                        # random seed
    cuda = false                                    # use CPU
    verbose_freq = 10                               # logging for every verbose_freq iterations
    tblogger = true                                 # log training with tensorboard
    save_path = "output"                            # results path
end

function train(; kws...)
    # load hyperparamters
    args = Args(; kws...)
    args.seed > 0 && Random.seed!(args.seed)

    # GPU config
    if args.cuda && CUDA.has_cuda()
        device = gpu
        @info "Training on GPU"
    else
        device = cpu
        @info "Training on CPU"
    end

    # load MNIST images
    loader = get_data(args.batch_size)

    # initialize UNet model
    unet = UNet() |> device

    # ADAM optimizer
    opt = ADAM(args.η)

    # parameters
    ps = Flux.params(unet)

    !ispath(args.save_path) && mkpath(args.save_path)

    # logging by TensorBoard.jl
    if args.tblogger
        tblogger = TBLogger(args.save_path, tb_overwrite)
    end

    # Training
    train_steps = 0
    @info "Start Training, total $(args.epochs) epochs"
    for epoch = 1:args.epochs
        @info "Epoch $(epoch)"
        progress = Progress(length(loader))

        for (x, _) in loader
            x = device(x)
            loss, grad = Flux.withgradient(ps) do
                model_loss(unet, x)
            end
            Flux.Optimise.update!(opt, ps, grad)
            # progress meter
            next!(progress; showvalues=[(:loss, loss)])

            # logging with TensorBoard
            if args.tblogger && train_steps % args.verbose_freq == 0
                with_logger(tblogger) do
                    @info "train" loss = loss
                end
            end
            train_steps += 1
        end
    end

    # save model
    model_path = joinpath(args.save_path, "model.bson")
    let unet = cpu(unet), args = struct2dict(args)
        BSON.@save model_path unet args
        @info "Model saved: $(model_path)"
    end
end

if abspath(PROGRAM_FILE) == @__FILE__
    train()
end


"""
    Faber(i, j)

change of indexation: from double indexation i,j to single indexation n
"""
function Faber(i::Int64, j::Int64)
    return Int(2<<(i-1) + j)
end

"""
    binlog(convert(UInt32, v))

efficient log_2 of integers
https://graphics.stanford.edu/~seander/bithacks.html#IntegerLog
"""
binlog(v::Integer) = binlog(convert(UInt32, v))
function binlog(v::UInt32)
    r =     (v > 0xFFFF) << 4; v >>= r;
    shift = (v > 0xFF  ) << 3; v >>= shift; r |= shift;
    shift = (v > 0xF   ) << 2; v >>= shift; r |= shift;
    shift = (v > 0x3   ) << 1; v >>= shift; r |= shift;
                                            r |= (v >> 1);
end

"""
    Faber(n)

change of indexation: from single indexation n to double indexation i,j
"""
function Faber(n)
    return (binlog(n)), Int(n - 1<<(binlog(n)))
end

"""
    OrderBasis(i, j)

compare two index (single indexation) and order it from the smaller
"""
function OrderBasis(i, j)
    if i <= j
        return (i, j)
    else
        return (j, i)
    end
end

"""
    OrderBasis(i1, j1, i2, j2)

compare two index (double indexation) and order it from the smaller
"""
function OrderBasis(i1 ,j1 ,i2 ,j2)
    if i1 <= i2
        return (i1, j1, i2, j2)
    else
        return (i2, j2, i1, j1)
    end
end
"""
    ShareDomain(i,j)

check if the basis relative to index i and index j (single indexation) share the same support
return 1 if they do
return 0 if they don't
"""
function ShareDomain(i,j)
    if i == j return 1
    else
        (i, j) = OrderBasis(i, j)
        i1, j1 = Faber(i)
        i2, j2 = Faber(j)
        if j1*2<<(i2 - i1 - 1) <= j2 < (j1 + 1)*2<<(i2 - i1-1)
            return 1
        else
            return 0
        end
    end
end

"""
    SharedSupport(n, L)

Collect all the indexes (single indexation) relative to the basis function sharing
sharing the same support of the function ϕ_n up to level L
"""
function SharedSupport(n,L)
    a = Int64[]
    for i in 1:(2<<(L)) - 1
        if ShareDomain(n,i) == 1 # && n != i
            push!(a, i)
        end
    end
    return a
end

"""
    SharedSupport(i,j, L)

Collect all the indexes (with single indexation) relative to the basis functions
sharing the same support of the function ϕ_{i,j} up to  level L
"""
function SharedSupport(i, j, L)
    n = Faber(i,j)
    a = SharedSupport(n, L)
    return a
end
"""
    NotSharedSupport(i,j, L)

Collect all the indexes (with single indexation) relative to the basis functions
which do the same support of the function ϕ_{i,j} up to  level L
"""
function NotSharedSupport(n, L)
    a = Int64[]
    for i in 1:(2<<(L)) - 1
        if ShareDomain(n,i) == 0
            push!(a, i)
        end
    end
    return a
end

"""
    FS

the strcut
```
struct Fs
    i::Int16    #idenx i (double indexation)
    j::Int16    #index j (double indexation)
    n::Int64    #idenx n (single indexation)
    lb::Float64 #pper bound domain
    ub::Float64 #lower Bound domain [not used?]
    range::Float64 #ub - lb
    nhb::Array{Int64} #neightbourhoods
    nhb::Array{Int64} #not neightbourhoods
    supr::Float64 #suprimum(ϕ)
    δ::Float64
end
```
"""
struct Fs
    i::Int16    #idenx i (double indexation)
    j::Int16    #index j (double indexation)
    n::Int64    #idenx n (single indexation)
    lb::Float64 #Upper bound domain
    ub::Float64 #Lower Bound domain [not used?]
    range::Float64
    nhb::Array{Int64} #neightbourhoods
    notnhb::Array{Int64}
    supr::Float64 #suprimum(ϕ)
    δ::Float64 # range*suprimum(ϕ)
    function Fs(i, j, L, T) #inizialization
        n = Faber(i, j)
        new(i, j, n, T*j/2^i, T*(j+1)/2^i, T/2^i, SharedSupport(n,L), NotSharedSupport(n,L), 2^(-i/2 - 1)*sqrt(T), T/2^i*2^(-i/2 -1)*sqrt(T))
    end
    function Fs(n, L, T) #inizialization
        i, j = Faber(n)
        Fs(i, j, L, T)
    end
end

"""
    generate(L, T)

generate all the set of basis used consistenly
"""
function generate(L, T)
    ϕ = Array{Fs}(undef, 2<<L - 1)
    for  n in 1:2<<L -1
        ϕ[n] = Fs(n, L, T)
    end
    ϕ
end

"""
    Λ(x, T)

function needed for construction of any Faber Schauder function i,j
see   Λ(x, l, k, T). This function should not be visible.
"""
function Λ(x, T)
    return max(0.0, sqrt(T)/2 - abs(x/sqrt(T) - sqrt(T)/2))
end

"""
    Λ(x, l, k, T)

Faber Schauder function `l`, `k` (double indexation) construct as in definition
rescaled by length `T` and evaluated at `x`
"""
function Λ(x, l, k, T)
    x = mod(x, T)   #ok
    2^(-l/2)*Λ((x)*(1<<l) - k*T, T)
end

"""
    IntQuad1(a, b, n1, n2)
integral of quadratic function1 with `a` quadratic term
`b` linear term and`n1` `n2` uppper and lower bound:

∫_n2^n1 a*t^2 + b*t dt
"""
function IntQuad1(a, b, n1, n2)
    a*(n1^3 - n2^3)/3 + b*(n1^2 - n2^2)/2
end

"""
    flip(i,j)

flip index j to the left to ease the integration:
int_0^1 psi_i (t) psi_j(t) dt

ex flip(1,1) = 0 flip(1,0) = 0 (do nothing), flip(3,5)= 2
"""
function flip(i,j)
    bound = ((1/2)*2^i - 1/2)
    if bound < j
        j = 2*bound - j
    end
    return Integer(j)
end
"""
    flip(i,j)

flip index j on the other side to ease the computation
of bound2(n,T)

ex. flip2(1,1) = 0, flip2(1,0) = 1, flip2(2,0) = 3, flip2(2,1) = 2.

"""
function flip2(i,j)
    bound = ((1/2)*2^i - 1/2)
    j = 2*bound - j
    return Integer(j)
end

"""
    int_prod_fs(i,j, T)

compute Φ_{ij} = ∫_0^T ϕ_i ϕ_j dt
ϕ_i being the ith (single indexation) Faber Schauder function
"""
#int_0^1 psi_i (t) psi_j(t) dt with i<j
function int_prod_fs(i,j, T)
    (i, j) = OrderBasis(i, j)
    i1, j1 = Faber(i) #ordering  i < j --> i1 < i2
    i2, j2 = Faber(j)
    i3 = i2 - i1
    j3 = j2 - 2^(i2-i1)*j1
    j3 = flip(i3,j3)
    A = j3/2^i3
    B = (j3 + 1/2)/2^i3
    C = (j3 + 1)/ 2^i3
    return T^2*2^((i2 - 5*i1)/2)*(IntQuad1(1,-j3/2^i3, B, A) - IntQuad1(1,-(j3+1)/2^i3, C, B))
end

"""
    intfs(i, T)

compute Φ_i = ∫_0^T ϕ_{ij} ds
"""
function intfs(n, T)
    i, j = Faber(n)
    return T^(1.5)*2^(-(3i + 4)/2 )
end

"""
    generate_matrix(L, T)
generate vector whose element is ∫_0^T ϕ(i1,j1)*ϕ(i2,j2) dt
for each i1,j1,i2,j2 until truncation level `L`
"""

function generate_matrix(L::Int64, T)
    dim = 2^(L+1) -1
    V = zeros(dim,dim)
    for i in 1:dim
        for j in 1:i
            if i == j
                i1, j1 = Faber(i)
                V[i,i] = T^2* 2.0^(-2*i1 - 2)/3
            else
                if ShareDomain(i,j)==1
                    res = int_prod_fs(i, j, T)
                    V[i,j], V[j,i] = (res, res)
                else
                    V[i,j], V[j,i] = (0 , 0)
                end
            end
        end
    end
    return V
end

"""
    generate_vector(L, T)
generate vector whose element is ∫_0^T ϕ(i,j) dt
for each i,j until truncation level `L`
"""

function generate_vector(L, T)
    [intfs(i, T) for i in 1:(2<<L -1)]
end

"""
    bound2(L, T)
computes ∫_0^T ϕ(n)(t/T) dt
N.B. `n` single inexation
"""
function bound1(n, T)
    i,j = Faber(n)
    A = j / (2^i)
    B = (0.5 + j)/ (2^i)
    C = (1.0 + j)/ (2^i)
    return T^(3/2)*2^(-i/2)*(IntQuad1(2^i,-j, B, A) + IntQuad1(-2^i, j + 1, C, B))
end

"""
    bound2(L, T)
computes ∫_0^T ϕ(n)(1 - t/T) dt
N.B. `n` single inexation
"""
function bound2(n, T)
    i,j = Faber(n)
    j = flip2(i,j)
    A = j / (2^i)
    B = (0.5 + j)/ (2^i)
    C = (1.0 + j)/ (2^i)
    return T^(3/2)*2^(-i/2)*(IntQuad1(2^i,-j, B, A) + IntQuad1(-2^i, j + 1, C, B))
end

"""
    generate_bound1(L, T)
generate vector whose element is ∫_0^T ϕ(i,j)(t/T) dt
for each i,j until truncation level `L`
"""
function generate_bound1(L, T)
    [bound1(n,T) for n in 1:1:2<<L-1]
end

"""
    generate_bound2(L, T)
generate vector whose element is ∫_0^T ϕ(i,j)(1 - t/T) dt
for each i,j until truncation level `L`
"""
function generate_bound2(L, T)
    [bound2(n,T) for n in 1:1:2<<L-1]
end

"""
    MCintegration(k::Fs)

take a random point inside the support of the basis function k
"""
function MCintegration(k::Fs)
    rand()*k.range + k.lb
end

"""
    Sol2E(a::Float64, b::Float64, c::Float64)

Solve a quadratic equation.
TODO add errors when the solution is not real.
solve quadratic equation ax^2 + bx + c = 0
"""
function Sol2E(a::Float64, b::Float64, c::Float64)
    return (-b + sqrt(b^2 - 4a*c))/2a
end


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

"""
    Polytope{N,Dim,T}

We say that a geometry is a N-polytope when it is a collection of "flat" sides
that constitue a `N`-dimensional subspace. They are called polygon and polyhedron
respectively for 2D (`N=2`) and 3D (`N=3`) subspaces, embedded in a `Dim`-dimensional
space. The parameter `N` is also known as the rank or parametric dimension of the
polytope: https://en.wikipedia.org/wiki/Abstract_polytope.

The term polytope expresses a particular combinatorial structure. A polyhedron,
for example, can be decomposed into faces. Each face can then be decomposed into
edges, and edges into vertices. Some conventions act as a mapping between vertices
and higher dimensional features (edges, faces, cells...), removing the need to
store all features. We follow the ordering conventions of the GMSH project:
https://gmsh.info/doc/texinfo/gmsh.html#Node-ordering

Additionally, the following property must hold in order for a geometry to be considered
a polytope: the boundary of a (N+1)-polytope is a collection of N-polytopes, which may
have (N-1)-polytopes in common. See https://en.wikipedia.org/wiki/Polytope.
"""
abstract type Polytope{N,Dim,T} <: Geometry{Dim,T} end

(::Type{PL})(vertices::Vararg{P}) where {PL<:Polytope,P<:Point} = PL(SVector(vertices))
(::Type{PL})(vertices::AbstractVector{TP}) where {PL<:Polytope,TP<:Tuple} = PL(Point.(vertices))
(::Type{PL})(vertices::Vararg{TP}) where {PL<:Polytope,TP<:Tuple} = PL(collect(vertices))

# type aliases for convenience
const Polygon = Polytope{2}
const Polyhedron = Polytope{3}
area(p::Polygon) = measure(p)
volume(p::Polyhedron) = measure(p)

"""
    paramdim(polytope)

Return the parametric dimension or rank of the polytope.
"""
paramdim(::Type{<:Polytope{N}}) where {N} = N

"""
    vertices(polytope)

Return the vertices of a `polytope`.
"""
vertices(p::Polytope) = p.vertices

"""
    nvertices(polytope)

Return the number of vertices in the `polytope`.
"""
nvertices(p::Polytope) = length(p.vertices)

"""
    facets(polytope, convention=GMSH)

Return the facets of a `polytope` according to some
ordering `convention`. Default convention is [`GMSH`](@ref).
See https://en.wikipedia.org/wiki/Facet_(geometry)
"""
function facets(p::Polytope{N}, convention=GMSH) where {N}
  faces(p, N-1, convention)
end

"""
    faces(polytope, rank, convention=GMSH)

Return the faces of the `polytope` of given `rank` according
to some ordering `convention`. Default convention is [`GMSH`](@ref).
"""
function faces(p::Polytope, rank, convention=GMSH)
  (materialize(ord, p.vertices) for ord in connectivities(typeof(p), rank, convention))
end

"""
    p1 == p2

Tells whether or not polytopes `p1` and `p2` are equal.
"""
==(p1::Polytope, p2::Polytope) = p1.vertices == p2.vertices

"""
    centroid(polytope)

Return the centroid of the `polytope`.
"""
centroid(p::Polytope) = Point(sum(coordinates.(p.vertices)) / length(p.vertices))

function Base.show(io::IO, p::Polytope)
  kind = nameof(typeof(p))
  vert = join(p.vertices, ", ")
  print(io, "$kind($vert)")
end

function Base.show(io::IO, ::MIME"text/plain", p::Polytope{N,Dim,T}) where {N,Dim,T}
  kind = nameof(typeof(p))
  lines = ["  └─$v" for v in p.vertices]
  println(io, "$kind{$Dim,$T}")
  print(io, join(lines, "\n"))
end

include("polytopes/segment.jl")
include("polytopes/triangle.jl")
include("polytopes/quadrangle.jl")
include("polytopes/pyramid.jl")
include("polytopes/tetrahedron.jl")
include("polytopes/hexahedron.jl")
include("polytopes/chain.jl")
include("polytopes/polyarea.jl")


# example to illustate how autoinit functions work, and used as a development script towards a standard unit test.

using Distributions
using IncrementalInference

import IncrementalInference: getSample

## MARKER START is a required block of code, but can be reviewed at the end of the tutorial,
# Run this but skip past these user defined functions for a quicker introduction======================
struct Prior{T} <: IncrementalInference.FunctorSingleton where T <: Distribution
  z::T
end
getSample(s::Prior, N::Int=1) = (rand(s.z,N), )
struct LinearOffset{T} <: IncrementalInference.FunctorPairwise where T <: Distribution
  z::T
end
getSample(s::LinearOffset, N::Int=1) = (rand(s.z,N), )
function (s::LinearOffset)(res::Array{Float64},
      idx::Int,
      meas::Tuple{Array{Float64, 1}},
      X1::Array{Float64,2},
      X2::Array{Float64,2}  )
  #
  res[1] = meas[1][idx] - (X2[1,idx] - X1[1,idx])
  nothing
end
# and a bi-modal conditional function
struct MultiModalOffset <: IncrementalInference.FunctorPairwise
  z::Vector{Distribution}
  c::Categorical
end
getSample(s::MultiModalOffset, N::Int=1) = (rand.(s.z, N)..., rand(s.c, N))
function (s::MultiModalOffset)(res::Array{Float64},
      idx::Int,
      meas::Tuple,
      X1::Array{Float64,2},
      X2::Array{Float64,2}
  )
  res[1] = meas[meas[end][idx]][idx] - (X2[1,idx] - X1[1,idx])
  nothing
end
## Define a model with these user defined functions ===============================================
## MARKER END

# Start with an empty graph
fg = emptyFactorGraph()

# add the first node
addNode!(fg, :x0, ContinuousScalar)

# this is unary (prior) factor and does not immediately trigger autoinit of :x0.
addFactor!(fg, [:x0], Prior(Normal(0,1)))

# writeGraphPdf(fg, file="/home/dehann/Downloads/fgx0.png")

# To visualize the factor graph structure, assuming GraphViz.jl is available.
# Graphs.plot(fg.g)
# Please find `writeGraphPdf` member definition at the end of this tutorial
# writeGraphPdf(fg, file="/home/dehann/Downloads/fgx01.png")

# Also consider automatic initialization of variables.
@show isInitialized(fg, :x0)
# Why is x0 not initialized?
# Since no other variable nodes have been 'connected to' (or depend) on :x0,
# and future intentions of the user are unknown, the initialization of :x0 is
# deferred until such criteria are met.

# Auto initialization of :x0 is triggered by connecting the next (probably uninitialized) variable node
addNode!(fg, :x1, ContinuousScalar)

# with a linear conditional belief to :x0
# P(Z | :x1 - :x0 ) where Z ~ Normal(10,1)
addFactor!(fg, [:x0, :x1], LinearOffset(Normal(10.0,1)))

# writeGraphPdf(fg, file="/home/dehann/Downloads/fgx01.png")

# x0 should not be initialized
@show isInitialized(fg, :x0)

# To plot the belief states of variables -- (assuming `RoMEPlotting` is available)
# Remember the first time executions are slow given required code compilation,
# and that future versions of these package will use more precompilation
# to reduce first execution running cost.
using RoMEPlotting

plotKDE(fg, :x0)

# Since no other variables 'are yet connected'/depend on :x1, it will not be initialized
@show isInitialized(fg, :x1)

# we can force all the variable nodes to initialize
ensureAllInitialized!(fg)

# now draw both :x0
plotKDE(fg, [:x0, :x1])

# add another node, but introduce more general beliefs
addNode!(fg, :x2, ContinuousScalar)

mmo = MultiModalOffset([Rayleigh(3); Uniform(30,55)], Categorical([0.4; 0.6]))
addFactor!(fg, [:x1, :x2], mmo)

# Graphs.plot(fg.g)
# writeGraphPdf(fg, file="/home/dehann/Downloads/fgx012.png")

# By again forcing the initialization of :x3 for illustration
ensureAllInitialized!(fg)

# the predicted marginal probability densities are
plotKDE(fg, [:x0, :x1, :x2])

# Now transmit this 'weird' multi-modal marginal belief through a another unimodal linear offset (conditional likelihood)
addNode!(fg, :x3, ContinuousScalar)

addFactor!(fg, [:x2, :x3], LinearOffset(Normal(-50, 1)))
# note, this addFactor step relies on :x2 being initialized and would have done so if we didn't call ensureAllInitialized! a few lines earlier.

# writeGraphPdf(fg, file="/home/dehann/Downloads/fgx0123.png")

ensureAllInitialized!(fg)
plotKDE(fg, [:x0, :x1, :x2, :x3])

lo3 = LinearOffset(Normal(40, 1))
addFactor!(fg, [:x3, :x0], lo3)

# writeGraphPdf(fg, file="/home/dehann/Downloads/fgx0123c.png")

plotKDE(fg, [:x0, :x1, :x2, :x3])

# Find global best likelihood solution (posterior belief)
# After defining the problem, we can find the 'minimum free energy' solution
tree = wipeBuildNewTree!(fg)
inferOverTree!(fg, tree)

# and look at the posterior belief, and notice which consensus modes stand out in the posterior
plotKDE(fg, [:x0, :x1, :x2, :x3])

# Helper function for graphing the factor graph structure (using GraphViz)

function writeGraphPdf(fgl::FactorGraph; file::String="fg.pdf")
  fgd = drawCopyFG(fgl)
  println("Writing factor graph file")
  fid = open("fg.dot","w+")
  write(fid,Graphs.to_dot(fgd.g))
  close(fid)
  fext = split(file, '.')[end]
  run(`dot fg.dot -T$fext -o $file`)
  Base.rm("fg.dot")
  nothing
end

# using Gadfly
# pl.guides[1] = Gadfly.Guide.xlabel("")
# push!(pl.guides, Gadfly.Guide.ylabel("density"))
#
# Gadfly.draw(PNG("/home/dehann/Downloads/plx012.png", 10cm, 7cm), pl)

## should complete and add to RoMEPlotting
# import KernelDensityEstimatePlotting: plot
# import Gadfly: plot
# import Graphs: plot
# import RoMEPlotting: plot
# function plot(fgl::FactorGraph, sym::Symbol; api::DataLayerAPI=IncrementalInference.dlapi)
#   PX = getKDE(getVert(fgl, sym, api=api))
#   plot(PX)
# end

#


# This file is a part of Julia. License is MIT: http://julialang.org/license

@testset "binomial" begin
    @test binomial(5,-1) == 0
    @test binomial(5,10) == 0
    @test binomial(5,3) == 10
    @test binomial(2,1) == 2
    @test binomial(1,2) == 0
    @test binomial(-2,1) == -2 # let's agree
    @test binomial(2,-1) == 0

    #Issue 6154
    @test binomial(Int32(34), Int32(15)) == binomial(BigInt(34), BigInt(15)) == 1855967520
    @test binomial(Int64(67), Int64(29)) == binomial(BigInt(67), BigInt(29)) == 7886597962249166160
    @test binomial(Int128(131), Int128(62)) == binomial(BigInt(131), BigInt(62)) == 157311720980559117816198361912717812000
    @test_throws InexactError binomial(Int64(67), Int64(30))
end

@testset "permutations" begin
    p = shuffle([1:1000;])
    @test isperm(p)
    @test all(invperm(invperm(p)) .== p)
    @test isperm(()) == true
    @test isperm((1,)) == true
    @test isperm((2,)) == false
    @test isperm((1,2)) == true
    @test isperm((2,1)) == true
    @test isperm((2,2)) == false
    @test isperm((1,3)) == false
    @test invperm(()) == ()
    @test invperm((1,)) == (1,)
    @test invperm((1,2)) == (1,2)
    @test invperm((2,1)) == (2,1)
    @test_throws ArgumentError invperm((1,3))

    push!(p, 1)
    @test !isperm(p)

    a = randcycle(10)
    @test ipermute!(permute!([1:10;], a),a) == [1:10;]

    # PR 12785
    let a = 2:-1:1
        @test ipermute!(permute!([1, 2], a), a) == [1, 2]
    end
end

@testset "factorial" begin
    @test factorial(7) == 5040
    @test factorial(Int8(7)) == 5040
    @test factorial(UInt8(7)) == 5040
    @test factorial(Int16(7)) == 5040
    @test factorial(UInt16(7)) == 5040
    @test factorial(Int32(7)) == 5040
    @test factorial(UInt32(7)) == 5040
    @test factorial(Int64(7)) == 5040
    @test factorial(UInt64(7)) == 5040
    @test factorial(Int128(7)) == 5040
    @test factorial(UInt128(7)) == 5040
    @test factorial(0) == 1
    @test_throws DomainError factorial(-1)
    @test factorial(Int64(20)) == 2432902008176640000
    # issue #6579
    @test_throws OverflowError factorial(Int64(21))
    @test typeof(factorial(Int8(2))) == typeof(factorial(Int8(1)))
    if Int === Int32
        @test factorial(Int32(12)) === Int32(479001600)
        @test_throws OverflowError factorial(Int32(13))
    end
end


using Zygote
include("chainrules_patch.jl")

import YaoExtensions, Random

c = YaoExtensions.variational_circuit(5)
h = YaoExtensions.heisenberg(5)

function loss(h, c, θ) where N
    # the assign is nessesary!
    c = dispatch(c, fill(θ, nparameters(c)))
    reg = apply(zero_state(nqubits(c)), c)
    real(expect(h, reg))
end

reg0 = zero_state(5)
zygote_grad = Zygote.gradient(θ->loss(h, c, θ), 0.5)[1]

# check gradients
using Test
dispatch!(c, fill(0.5, nparameters(c)))
greg, gparams = expect'(h, zero_state(5)=>c)
true_grad = sum(gparams)

@test zygote_grad ≈ true_grad

# the batched version
function loss2(h, c, θ) where N
    # the assign is nessesary!
    c = dispatch(c, fill(θ, nparameters(c)))
    reg = zero_state(nqubits(c),nbatch=2)
    reg = apply(reg, c)
    sum(real(expect(h, reg)))
end

reg0 = zero_state(5)
zygote_grad2 = Zygote.gradient(θ->loss2(h, c, θ), 0.5)[1]
true_grad2 = (loss2(h, c, 0.5+1e-5) - loss2(h, c, 0.5-1e-5)) / 2e-5
@test true_grad2 ≈ zygote_grad2


function linearindex(shape, idxs)
   stride = 1
   idx = 1
   @inbounds for d in 1:min(length(shape), length(idxs))
      i = idxs[d]
      n = shape[d]
      @assert 1<=i<=n "1<=$i<=$n   shape=$shape idxs=$idxs"
      idx += stride*(i-1)
      stride *= n
   end
   return idx
end


export Boussinesq

"""
    Boussinesq(params;kwargs)

Define an object of type `AbstractModel` in view of solving the initial-value problem for
`abcd`-Boussinesq models (with `b=d` and `c=0`).
(see Bona, J. L.; Chen, M. & Saut, J.-C., doi:10.1007/s00332-002-0466-4)

# Argument
`param` is of type `NamedTuple` and must contain

- dimensionless parameters `ϵ` (nonlinearity) and `μ` (dispersion);
- numerical parameters to construct the mesh of collocation points as `mesh = Mesh(param)`

## Optional keyword arguments
- two parameters `a` (default is `-1/3`) and `b` (default is `+1/3`) which determine the model solved. You need `a+2*b=1/3` for validity as a long wave model (without surface tension).
- `ktol`: tolerance of the low-pass Krasny filter (default is `0`, i.e. no filtering);
- `dealias`: dealiasing with Orlicz rule `1-dealias/(dealias+2)` (default is `0`, i.e. no dealiasing).
- `verbose`: prints information if `true` (default is `true`).

# Return values
Generate necessary ingredients for solving an initial-value problem via `solve!` and in particular
1. a function `Boussinesq.f!` to be called in the time-integration solver;
2. a function `Boussinesq.mapto` which from `(η,v)` of type `InitialData` provides the raw data matrix on which computations are to be executed;
3. a function `Boussinesq.mapfro` which from such data matrix returns the Tuple of real vectors `(η,v)`, where

    - `η` is the surface deformation;
    - `v` is the derivative of the trace of the velocity potential.

"""
mutable struct Boussinesq <: AbstractModel

	label   :: String
	f!		:: Function
	mapto	:: Function
	mapfro	:: Function
	param	:: NamedTuple
	kwargs	:: NamedTuple

    function Boussinesq(param::NamedTuple;
						a=-1/3,b=1/3,α=1, #α plays no role, but must be provided to WhithamBoussinesq
						dealias=0,ktol=0,verbose=true)

		m=WhithamBoussinesq(param;Boussinesq=true,a=a,b=b,
							dealias=dealias,ktol=ktol,verbose=verbose)

		new(m.label, m.f!, m.mapto, m.mapfro, m.param, m.kwargs)
    end
end


# ## Error bar

#md # ![FILE_NAME.png](images/FILE_NAME.png)

using Random
Random.seed!(2018)

f(x) = 2 * x + 1
x = 0:0.1:2
n = length(x)
y = f.(x) + randn(n)
plot(x, y,
    xerr=0.1 * rand(n),
    yerr=rand(n))


# Functions for efficiently evaluating all non-zero B-splines (and/or their
# derivatives) at a given point.

using Base.Cartesian: @ntuple
using Base: @propagate_inbounds
using StaticArrays: MVector

"""
    evaluate_all(
        B::BSplineBasis, x::Real,
        [op = Derivative(0)], [T = float(typeof(x))];
        [ileft = nothing],
    ) -> i, bs

Evaluate all B-splines which are non-zero at coordinate `x`.

Returns a tuple `(i, bs)`, where `i` is an index identifying the basis functions
that were computed, and `bs` is a tuple with the actual values.

More precisely:

- `i` is the index of the first B-spline knot ``t_{i}`` when going from ``x``
  towards the left.
  In other words, it is such that ``t_{i} ≤ x < t_{i + 1}``.

  It is effectively computed as `i = searchsortedlast(knots(B), x)`.
  If the correct value of `i` is already known, one can avoid this computation by
  manually passing this index via the optional `ileft` keyword argument.

- `bs` is a tuple of B-splines evaluated at ``x``:

  ```math
  (b_i(x), b_{i - 1}(x), …, b_{i - k + 1}(x)).
  ```

  It contains ``k`` values, where ``k`` is the order of the B-spline basis.
  Note that values are returned in backwards order starting from the ``i``-th
  B-spline.

## Computing derivatives

One can pass the optional `op` argument to compute B-spline derivatives instead
of the actual B-spline values.

## Examples

See [`AbstractBSplineBasis`](@ref) for some examples using the alternative
evaluation syntax `B(x, [op], [T]; [ileft])`, which calls this function.
"""
@propagate_inbounds function evaluate_all(
        B::BSplineBasis, x::Real, op::Derivative, ::Type{T}; kws...,
    ) where {T <: Number}
    _evaluate_all(knots(B), x, BSplineOrder(order(B)), op, T; kws...)
end

@propagate_inbounds evaluate_all(
    B, x, op::AbstractDifferentialOp = Derivative(0); kws...,
) = evaluate_all(B, x, op, float(typeof(x)); kws...)

@propagate_inbounds evaluate_all(B, x, ::Type{T}; kws...) where {T <: Number} =
    evaluate_all(B, x, Derivative(0), T; kws...)

@propagate_inbounds function _knotdiff(x, ts, i, n)
    @boundscheck checkbounds(ts, i:(i + n))
    @inbounds ti = ts[i]
    @inbounds tj = ts[i + n]
    # @assert ti ≠ tj
    (x - ti) / (tj - ti)
end

# TODO
# - this is redundant with Splines.knot_interval...
"""
    find_knot_interval(ts::AbstractVector, x::Real) -> (i, zone)

Finds the index ``i`` corresponding to the knot interval ``[t_i, t_{i + 1}]``
that should be used to evaluate B-splines at location ``x``.

The knot vector is assumed to be sorted in non-decreasing order.

It also returns a `zone` integer, which is:

- `0`  if `x` is within the knot domain (`ts[begin] ≤ x ≤ ts[end]`),
- `-1` if `x < ts[begin]`,
- `1`  if `x > ts[end]`.

This function is functionally equivalent to de Boor's `INTERV` routine (de Boor
2001, p. 74).
"""
function find_knot_interval(ts::AbstractVector, x::Real)
    if x < first(ts)
        return firstindex(ts), -1
    end
    i = searchsortedlast(ts, x)
    Nt = lastindex(ts)
    if i == Nt
        tlast = ts[Nt]
        while true
            i -= 1
            ts[i] ≠ tlast && break
        end
        zone = (x > tlast) ? 1 : 0
        return i, zone
    else
        return i, 0  # usual case
    end
end

function _evaluate_all(
        ts::AbstractVector, x::Real, ::BSplineOrder{k},
        op::Derivative{0}, ::Type{T};
        ileft = nothing,
    ) where {k, T}
    if @generated
        @assert k ≥ 1
        ex = quote
            i, zone = _find_knot_interval(ts, x, ileft)
            if zone ≠ 0
                return (i, @ntuple($k, j -> zero($T)))
            end
            # We assume that the value of `i` corresponds to the good interval,
            # even if `ileft` was passed by the user.
            # This is the same as in de Boor's BSPLVB routine.
            # In particular, this allows plotting the extension of a polynomial
            # piece outside of its knot interval (as in de Boor's Fig. 6, p. 114).
            bs_1 = (one(T),)
        end
        for q ∈ 2:k
            bp = Symbol(:bs_, q - 1)
            bq = Symbol(:bs_, q)
            ex = quote
                $ex
                Δs = @ntuple(
                    $(q - 1),
                    j -> @inbounds($T(_knotdiff(x, ts, i - j + 1, $q - 1))),
                )
                $bq = _evaluate_step(Δs, $bp, BSplineOrder($q))
            end
        end
        bk = Symbol(:bs_, k)
        quote
            $ex
            return i, $bk
        end
    else
        _evaluate_all_alt(ts, x, BSplineOrder(k), op, T; ileft = ileft)
    end
end

# Derivatives
function _evaluate_all(
        ts::AbstractVector, x::Real, ::BSplineOrder{k},
        op::Derivative{n}, ::Type{T};
        ileft = nothing,
    ) where {k, n, T}
    if @generated
        n::Int
        @assert n ≥ 1
        # We first need to evaluate the B-splines of order p.
        p = k - n
        if p < 1
            # Derivatives are zero. The returned index is arbitrary...
            return :( firstindex(ts), ntuple(_ -> zero($T), Val($k)) )
        end
        bp = Symbol(:bs_, p)
        ex = quote
            i, $bp = _evaluate_all(ts, x, BSplineOrder($p), Derivative(0), $T; ileft)
        end
        for q ∈ (p + 1):k
            bp = Symbol(:bs_, q - 1)
            bq = Symbol(:bs_, q)
            ex = quote
                $ex
                $bq = _evaluate_step_deriv(ts, i, $bp, BSplineOrder($q), $T)
            end
        end
        bk = Symbol(:bs_, k)
        quote
            $ex
            return i, $bk
        end
    else
        _evaluate_all_alt(ts, x, BSplineOrder(k), op, T; ileft = ileft)
    end
end

# Non-@generated version
@inline function _evaluate_all_alt(
        ts::AbstractVector, x::Real, ::BSplineOrder{k},
        ::Derivative{0}, ::Type{T};
        ileft = nothing,
    ) where {k, T}
    @assert k ≥ 1
    i, zone = _find_knot_interval(ts, x, ileft)
    if zone ≠ 0
        return (i, ntuple(j -> zero(T), Val(k)))
    end
    bq = MVector{k, T}(undef)
    Δs = MVector{k - 1, T}(undef)
    bq[1] = one(T)
    @inbounds for q ∈ 2:k
        for j ∈ 1:(q - 1)
            Δs[j] = _knotdiff(x, ts, i - j + 1, q - 1)
        end
        bp = bq[1]
        Δp = Δs[1]
        bq[1] = Δp * bp
        for j = 2:(q - 1)
            bpp, bp = bp, bq[j]
            Δpp, Δp = Δp, Δs[j]
            bq[j] = Δp * bp + (1 - Δpp) * bpp
        end
        bq[q] = (1 - Δp) * bp
    end
    i, Tuple(bq)
end

# Derivatives / non-@generated variant
function _evaluate_all_alt(
        ts::AbstractVector, x::Real, ::BSplineOrder{k},
        ::Derivative{n}, ::Type{T};
        ileft = nothing,
    ) where {k, n, T}
    n::Int
    @assert n ≥ 1
    # We first need to evaluate the B-splines of order m.
    m = k - n
    if m < 1
        # Derivatives are zero. The returned index is arbitrary...
        return firstindex(ts), ntuple(_ -> zero(T), Val(k))
    end
    i, bp = _evaluate_all(ts, x, BSplineOrder(m), Derivative(0), T; ileft)
    bq = MVector{k, T}(undef)
    us = MVector{k - 1, T}(undef)
    bq[1:m] .= bp
    for q = (m + 1):k
        p = q - 1
        for δj = 1:p
            @inbounds us[δj] = bq[δj] / (ts[i + q - δj] - ts[i + 1 - δj])
        end
        @inbounds bq[1] = p * us[1]
        for j = 2:(q - 1)
            # Note: adding @inbounds here can slow down stuff (from 25ns to
            # 80ns), which is very strange!! (Julia 1.8-beta2)
            bq[j] = p * (us[j] - us[j - 1])
        end
        @inbounds bq[q] = -p * us[p]
    end
    i, Tuple(bq)
end

function _find_knot_interval(ts, x, ileft)
    if isnothing(ileft)
        i, zone = find_knot_interval(ts, x)
    else
        i = ileft
        zone = (x < first(ts)) ? -1 : (x > last(ts)) ? 1 : 0
    end
    i, zone
end

@inline @generated function _evaluate_step(Δs, bp, ::BSplineOrder{k}) where {k}
    ex = quote
        @inbounds b_1 = Δs[1] * bp[1]
    end
    for j = 2:(k - 1)
        bj = Symbol(:b_, j)
        ex = quote
            $ex
            @inbounds $bj = (1 - Δs[$j - 1]) * bp[$j - 1] + Δs[$j] * bp[$j]
        end
    end
    b_last = Symbol(:b_, k)
    quote
        $ex
        @inbounds $b_last = (1 - Δs[$k - 1]) * bp[$k - 1]
        @ntuple $k b
    end
end

@inline @generated function _evaluate_step_deriv(
        ts, i, bp, ::BSplineOrder{k}, ::Type{T},
    ) where {k, T}
    p = k - 1
    ex = quote
        us = @ntuple(
            $p,
            δj -> @inbounds($T(bp[δj] / (ts[i + $k - δj] - ts[i + 1 - δj]))),
        )
        @inbounds b_1 = $p * us[1]
    end
    for j = 2:(k - 1)
        bj = Symbol(:b_, j)
        ex = quote
            $ex
            @inbounds $bj = $p * (-us[$j - 1] + us[$j])
        end
    end
    b_last = Symbol(:b_, k)
    quote
        $ex
        @inbounds $b_last = -$p * us[$p]
        @ntuple $k b
    end
end


### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# ╔═╡ ea9468a2-710f-11eb-2284-6d8af448d9eb
# dependencies
begin
	using PlutoUI
	using DataStructures
	using Logging
	Logging.global_logger(Logging.SimpleLogger(stdout, LogLevel(-5000)  ) )
	Logging.disable_logging(LogLevel(0))
end

# ╔═╡ 49d11f58-713c-11eb-20bb-e3a3dc889b6b
begin
	using LightGraphs, SimpleWeightedGraphs, Plots, LinearAlgebra

	"""
		mst(A::Array{<:Number, 2})

	From an adjacency matrix return a dictionary of possible
	destinations (including weights, returned a  tuple)
	"""
	function mst(A::Array{<:Number, 2})
		G = SimpleWeightedGraph(A) # initialise graph
		MST = kruskal_mst(G)       # calculate MST
		# translate MST to dict
		d = Dict{Int64, Array{Tuple{Int64, Float64},1}}()
		for edge in MST
		   push!(get!(d, edge.src, Array{Tuple{Int64, Float64},1}()), (edge.dst, edge.weight))
		end
		return d
	end
end;

# ╔═╡ 7782b3b2-710f-11eb-1a92-e12b8b1f3299
md"""
# A-star and the eight puzzle

Consider the eight puzzle, i.e. a grid of nine cases, with 8 integer values in them. The goal is to start with the empty square and then have all numbers in increasing order.

|initial   |       | state  |
|---|---|---|
| 7 | 2 | 4 |
| 5 |   | 6 |
| 8 | 3 | 1 |

|goal   |       | state  |
|---|---|---|
|   | 1 | 2 |
| 3 | 4 | 5 |
| 6 | 7 | 8 |

The objective of this session is to implement different heuristics and compare the performance of the A-star algorithm with the tree search methods implemented beforehand.

Below you find a number of generic and specific functions that allow you to identify the possible actions for a given state and to determine the next state from a previous one. For didactic purposes, there also is a function that allows you to visualize the current state.
"""

# ╔═╡ 58b8349c-7110-11eb-1aab-f7c716a12c57
md"""## Generic functions"""

# ╔═╡ 13d314dc-7110-11eb-3475-2126f5d414b7
begin
	"""
	Type used for representation of a node in the search tree.

	The node type will be inherited from the type of the state that is passed as an argument to the inner constructor

	Fields:
		- state: a node needs a state (of non-specified type T)
		- parent: each node has either a parent or nothing (in case of the root node)
		- action: the action that was taken to get from the previous node to the current one
		- pathcost: total cost of the current path (each contribution should be non-negative)
		- depth: to illustrate the depth of the tree and to be able to implement depth limited search.
	"""
	struct Node{T}
		state::T
		parent::Union{Nothing, Node, UInt}
		action::Union{Nothing, String, Int64, Tuple, UnitRange}
		pathcost::Float64
		depth::Int
		function Node(state::T, parent::Union{Nothing, Node, UInt}=nothing, action::Union{Nothing, String, Int64, Tuple, UnitRange}=nothing, pathcost=0.0, depth::Int=0) where {T}
			return new{T}(state, parent, action, pathcost, depth)
		end
	end

	"""
	Type used for representation of a problem with an initial and a goal state of type T
	"""
	struct Problem{T}
		initstate::T
		goalstate::T
		function Problem(initstate::T, goalstate::T) where {T}
			return new{T}(initstate, goalstate)
		end
	end

	"""
		path(n::Node)

	Return sequence of nodes to come to a solution (inluding the root node)
	"""
	function path(n::Node)
		track = Array{Node,1}()
		node = n
		pushfirst!(track, node)
		while node.parent ≠ nothing
			node = node.parent
			pushfirst!(track, node)
		end
		return track
	end

	"""
		solution(x::Array{Node,1})

	Return sequence of actions to come to a solution
	"""
	function solution(x::Array{Node,1})
		return [node.action for node in x[2:end]]
	end
end;

# ╔═╡ 423469c0-7110-11eb-0cef-03540836525b
md"""## Problem specific functions"""

# ╔═╡ 68def13a-7110-11eb-29cf-8da0df67b178
begin

	"""
		goaltest_puzzle(p::Problem, n::Node)

	Compare problem's endstate with current state of a node
	"""
	function goaltest_puzzle(p::Problem, n::Node)
		if p.goalstate == n.state
			return true
		else
			return false
		end
	end

	"""
		findactions_puzzle(state::NTuple{N,Int64}) where {N}

	Find all possible outcomes of actions, independent of gridsize
	"""
	function findactions_puzzle(state::NTuple{N,Int64}) where {N}
		d = Int(sqrt(N))
		# all possible actions (only four exist to move the blank case):
		possible_actions = ["UP","DOWN","LEFT","RIGHT"]
		# find location of empty case
		ind = findfirst(x->x==0, state)
		# remove options
		if ind <= d
			filter!(x -> x ≠ "UP", possible_actions)
		end
		if ind > N - d
			filter!(x -> x ≠ "DOWN", possible_actions)
		end
		if ind % d  == 1
			filter!(x -> x ≠ "LEFT", possible_actions)
		end
		if ind % d  == 0
			filter!(x -> x ≠ "RIGHT", possible_actions)
		end

		return possible_actions
	end

	"""
		applyaction_puzzle(state::NTuple{N,Int64}, action::String) where {N}

	Go from one state to another by applying an action.
	"""
	function applyaction_puzzle(state::NTuple{N,Int64}, action::String) where {N}
		d = Int(sqrt(N))
		# make copy of state
		state = collect(state)
		# find location of empty case
		ind = findfirst(x->x==0, state)
		# operations and index shifts   ## Maybe make this a constant for speed gains
		ops = Dict("UP"=>-d, "DOWN"=>d, "LEFT"=>-1, "RIGHT"=>1)
		state[ind], state[ind + ops[action]] = state[ind + ops[action]], state[ind]
		return Tuple(state)
	end

	"""
		prettyprint(state::Array{Int,1})

	Supporting function to visualize the playing field
	"""
	function prettyprint(state::NTuple{N,Int64}) where {N}
		d = Int(sqrt(N))
		function morph(x::Int)
			if x == 0
				return "  "
			elseif x < 10
				return " $(x)"
			elseif x >= 10
				return "$(x)"
			end
		end
		msg = "Game layout:\n\n"
		for i = 1:d
			msg = msg *"| " *  prod(["$(morph(state[(i-1)*d+j])) | " for j in 1:d]) * "\n"
		end
		println(msg)
	end
end;

# ╔═╡ b5309106-7110-11eb-3631-c1484cc8a5d9
md"""## Illustration of a small game"""

# ╔═╡ c28473ea-7110-11eb-2129-cdc0aab5444e
begin
	println("Illustration of an easy case (2 steps):")
	startstate = Tuple([3, 1, 2, 4, 0, 5, 6, 7, 8])
	endstate = Tuple(collect(0:8))
	p = Problem(startstate, endstate)
	prettyprint(p.initstate)
	actions = ["LEFT","UP"]
	state = p.initstate
	for action in actions
		state = applyaction_puzzle(state, action)
		prettyprint(state)
	end
end

# ╔═╡ ccfd8924-7110-11eb-3890-6d9d80cfb96b
md"""## Astar algorithm
Below you find the treesearch Astar algorithm implementation. Analyze the way it works ."""

# ╔═╡ d629933a-7110-11eb-39ad-7582f1db3d62
begin
	"""
		treesearchAstar(p::Problem; findactions::Function, applyaction::Function, goaltest::Function, g::Function, h::Function, kwargs...)

	Relatively generic implementation of A* tree search.
	Determines sequence of nodes and transformations.

	keyword arguments:
		- g(p, node, action)::Function: calculate the path cost from one node to the next one
		- h(state, goalstate)::Function: calculate the heuristic for a current state towards a goal state

		- to allow for extra output during the tree traversal, you need to enable a logger that can actually log this. e.g.
						Logging.global_logger(Logging.SimpleLogger(stdout, LogLevel(-5000)  ) )
						Logging.disable_logging(LogLevel(-5000))
	"""
	function treesearchAstar(p::Problem;
						 	findactions::Function,
							applyaction::Function,
							goaltest::Function,
							g::Function,
							h::Function, kwargs...)
		@debug "\n" * "#"^29 * "\n# Astar tree search # \n" * "#"^29
		node = Node(p.initstate)                                             # initialise root node
		fringe = DataStructures.PriorityQueue{Node, typeof(node.pathcost)}() # create Priority Queue
		enqueue!(fringe, node, node.pathcost + h(p, node; action=nothing, kwargs...))   # place root node in the Queue
		@debug "Current queue: $(fringe)"
		while !isempty(fringe)                                               # expand the fringe
			node = dequeue!(fringe)                                          # pop element from queue
			if goaltest(p, node; kwargs...)
				@debug "Solution found!"
				return path(node), solution(path(node))
			end
			@debug "Current node: $(node), current queue: $(fringe)"
			@debug "Possible action: $(findactions(node.state))"
			for action in findactions(node.state; kwargs...)
				child = Node(applyaction(node.state, action; kwargs...), node, action, g(p, node, action; kwargs...), node.depth + 1)
				@debug "new child: $(child)"
				enqueue!(fringe, child, child.pathcost + h(p, node; action=action, kwargs...))
				#enqueue!(fringe, child, child.pathcost + h(node.state, p.goalstate; kwargs...))
			end
		end
		@warn "failed to find a solution"
		return nothing
	end
end;

# ╔═╡ 902a810e-7111-11eb-1c4f-ef75bbcac693
md"""
* Implement both heuristics from the course (number of mismatches & manhattan mismatches)
* Compare the solution with previous methods (computation time + memory cost) on the case below:

```julia
    # Illustrations of a harder case:
    startstate = Tuple([1, 2, 3, 4, 0, 5, 6, 7, 8])
    endstate = Tuple(collect(0:8))
    p = Problem(startstate, endstate)
```

"""

# ╔═╡ 85073ccc-75e7-11eb-0d0f-c1c178e4e57e

# ╔═╡ 07f44e1c-713b-11eb-25b0-b9824cc5d251
md"""
## The Astar graph search method
Additional gains can be obtained by avoiding to visit the same state again (do you remember the requirements for the heuristic function in order for this to lead to the best solution?). Take the Astar algorithm from before and transform the tree search to graph search. One again, compare this to the previous methods.
"""

# ╔═╡ 369047da-713b-11eb-3f47-29e522f17fb5
begin
	"""
		Astargraphsearch(p::Problem; g::Function, h::Function; kwargs...)

	Relatively generic implementation of A* tree search.
	Determines sequence of nodes and transformations.

	keyword arguments:
		- g(p, node, action)::Function: calculate the path cost from one node to the next one
		- h(state, goalstate)::Function: calculate the heuristic for a current state towards a goal stat

		- to allow for extra output during the tree traversal, you need to enable a logger that can actually log this. e.g.
						Logging.global_logger(Logging.SimpleLogger(stdout, LogLevel(-5000)  ) )
						Logging.disable_logging(LogLevel(-5000))
	"""
	function Astargraphsearch(p::Problem;
							findactions::Function,
							applyaction::Function,
							goaltest::Function,
							g::Function,
							h::Function, kwargs...)
		@debug "\n" * "#"^29 * "\n# Astar tree search # \n" * "#"^29
		# initialise root node
		node = Node(p.initstate)
		# create Priority Queue
		fringe = DataStructures.PriorityQueue{Node, typeof(node.pathcost)}()
		explored = Set{typeof(p.initstate)}()
		# place root node in the Queue
		enqueue!(fringe, node, node.pathcost + h(p, node; action=nothing, kwargs...))
		@debug "Current queue: $(fringe)"
		# expand the fringe
		while !isempty(fringe)
			# pop element from queue
			node = dequeue!(fringe)
			if goaltest(p, node; kwargs...)
				@debug "Solution found!"
				return path(node), solution(path(node))
			end
			push!(explored, node.state)
			@debug "Current node: $(node), current queue: $(fringe)"
			@debug "Possible action: $(findactions(node.state))"
			for action in findactions(node.state)
				child = Node(applyaction(node.state, action), node, action, g(p, node, action; kwargs...), node.depth + 1)
				@debug "new child: $(child)"
				if child.state ∉ explored
					enqueue!(fringe, child, child.pathcost + h(p, node; action=action, kwargs...))
				end
			end
		end
		@warn "failed to find a solution"
		return nothing
	end

end;

# ╔═╡ 9b3d2832-75e7-11eb-0372-35f7b9ed402b

# ╔═╡ b35e5664-7143-11eb-1e75-17b00177d1f9
md"""
## Additional problems
We have seen several search algorithms, each time illustrated by one or more examples. The difficulty of this kind of problems is not so much in the implementation of an algorithm, but in formalising a problem and choosing a correct heuristic. The following problems will illustrate this.
"""

# ╔═╡ a653d9a8-7143-11eb-3812-396c3ec732bd

md"""
### 2D grid

$(PlutoUI.LocalResource("./img/fig3-9.png", :width => 500, :align => "middle"))
Consider the unbounded version of the regular 2D grid shown in Figure 3.9. The start state is at the origin, (0,0), and the goal state is at (x, y).
1. What is the branching factor b in this state space?
2. How many distinct states are there at depth k (for k > 0)?
3. What is the maximum number of nodes expanded by breadth-first tree search?
4. Is h = |u − x| + |v − y| an admissible heuristic for a state at (u, v)? Explain.
"""

# ╔═╡ 28f0121c-713c-11eb-2aff-e7a9c592fa53
md"""

### TSP
The [traveling salesperson problem (TSP)](https://en.wikipedia.org/wiki/Travelling_salesman_problem) can be solved with the [minimum-spanning-tree (MST)](https://en.wikipedia.org/wiki/Minimum_spanning_tree) heuristic, which estimates the cost of completing a tour, given that a partial tour has already been constructed. The MST cost of a set of cities is the smallest sum of the link costs of any tree that connects all the cities.

1. Show how this heuristic can be derived from a relaxed version of the TSP. What relaxation has been applied?
2. Think about how you can implement this problem in the general framework (what is a state, how do you check if a goalstate is realised, how do you define a transition and possible candidates...)
2. Show that the MST heuristic dominates straight-line distance. Compare the following heuristics:
    - MST
    - straight-line distance
    - nearest neighbor
3. Write a problem generator for instances of the TSP where cities are represented by random points in the unit square.
4. Compare the Astar graph search with what you have learned previously (e.g. solve it as a LP problem)

Below you can find a small example of the MST calculation for a tiny network that can be used for further three search.
"""

# ╔═╡ 99bbb244-713c-11eb-34be-d78f80834d46
# small illustration on a distance matrix
begin
	A = [0 30 5 7 0;
		30 0  1  0 0;
		 5 1  0  10 10;
		7 0  10  0 0;
		0 0 10 0 0]
	mst(A)
end

# ╔═╡ 8b9fb7a0-7146-11eb-089c-f5c25325324b
begin
	"""
		goaltest_TSP(p::Problem, n::Node; kwargs...)

	Compare problem's endstate with current state of a node.
	 (slightly modified the goaltest function to check all visited)
	"""
	function goaltest_TSP(p::Problem, n::Node; kwargs...)
		if p.goalstate[2] == union(n.state[2], n.state[1])
			return true
		else
			return false
		end
	end

	"""
		solution_TSP(x::Array{Node,1})

	Return sequence of action to come to a solution
	 (slightly modified solution so it returns a closed loop)
	"""
	function solution_TSP(x::Array{Node,1})
		return [x[1].state[1];[node.action for node in x[2:end]]...;x[1].state[1]]
	end

	"""
		findactions_TSP(state::Tuple{Int64,Set{Int64}})

	Find all possible actions for the travelling salesman
	i.e. exlude all visited locations and the current location
	"""
	function findactions_TSP(state::Tuple{Int64,Set{Int64},Int64}; kwargs...)
		return setdiff(Set(1:state[3]), state[2], state[1])
	end

	"""
		applyaction_TSP(state::Tuple{Int64,Set{Int64}}, action::Int64)

	Go from one state to another by applying an action.
	"""
	function applyaction_TSP(state::Tuple{Int64,Set{Int64},Int64}, action::Int64; kwargs...)
		# new location
		newpos = action
		# update visited list
		newvisited = copy(state[2])
		push!(newvisited, state[1])

		return (newpos, newvisited, state[3])
	end

	"""
		g_TSP(p::Problem, node::Node, action::Int; kwargs...)

	Path cost starting from a node and taking a specific action

	Required kwarg: "wm" = weight matrix describing the distance between each node in the global graph

	"""
	function g_TSP(p::Problem, node::Node, action::Int; kwargs...)
		cost = node.pathcost
		extra = kwargs[:wm][node.state[1], action]
		#@warn kwargs[:wm]
		#@warn "node depth: $(node.depth), cost: $(node.pathcost), extra: $(extra)"
		return cost + extra
	end

	"""
		h_TSP(p::Problem, node::Node; action::Int64, kwargs...)

	MST heuristic, which estimates the cost of completing the tour,
	given that a partial tour has already been constructed.

	Required kwarg: "wm" = weight matrix describing the distance between each node in the global graph

	"""
	function h_TSP(p::Problem, node::Node; action::Union{Nothing,Int64}, kwargs...)
		# in case of nothing as action
		if isnothing(action)
			return 0
		end

		current = node.state[1]
		vis = node.state[2]
		vcount = node.state[3]
		# subgraph weight matrix:
		# remove visited nodes, current state or action from global graph
		inds = [x for x in 1:vcount if x∉[current, action, vis...]]
		H_weights = kwargs[:wm][inds,inds]
		H = SimpleWeightedGraph(H_weights)
		try
			# calculate MST
			MST = prim_mst(H)
			# return estimate of remaining tour length
			return sum([weights(H)[e.src, e.dst] for e in MST])
		catch
			return 0
		end
	end
end;

# ╔═╡ c6fda06e-713c-11eb-3587-47c7faed0d07
begin
	# visualisations

	"""
		plotgraph(g::SimpleWeightedGraph, d::Dict,draw_MST::Bool=true,wf::Real=5)

	Visualize the network with its edges and their weights

	parameters:
		- g : SimpleWeightedGraph
		- d : Dict
			dictionary that links a node to a coordinate pair
		- draw_MST : Bool
			if true, also draws the MST on the figure
		- wf : Real
			multiplication factor for the edge weigth to obtain a nicer visualisation
	"""
	function plotgraph(g::SimpleWeightedGraph, d::Dict;wf::Real=5,draw_MST::Bool=true, draw_weights::Bool=false)
		p = plot(legend=:bottomright)
		for i in 1:length(d)
			annotate!([d[i][1]], [d[i][2]] .+0.05, "$(i)")
		end
		for e in edges(g)
			plot!([d[e.src][1], d[e.dst][1]],[d[e.src][2], d[e.dst][2]],
				markershape=:circle, linewidth=e.weight*wf, color=:black, markeralpha = 0.5,
				linealpha=0.2, label="", markercolor=:black, markersize=5)
			if draw_weights
				annotate!([(d[e.src][1] + d[e.dst][1])/2], [(d[e.src][2] + d[e.dst][2])/2],"$(e.weight)", color=:blue);
			end
		end
		plot!([],[],markershape=:circle, color=:black,
				linealpha=0.2, label="Full network", markercolor=:black, markersize=5, markeralpha = 0.5)
		if draw_MST
			MST = prim_mst(g, weights(g))
			for e in MST
				plot!([d[e.src][1], d[e.dst][1]],[d[e.src][2], d[e.dst][2]],
						markershape=:circle, linewidth=weights(g)[e.src,e.dst]*wf, color=:red,
						linealpha=0.8, label="", linestyle=:solid, markeralpha = 0)

			end
			plot!([],[],linealpha=0.8, label="MST", color=:red, linestyle=:solid)
		end
		xlims!(0,1)
		ylims!(0,1)
		plot!(legend=:best)
		return p
	end

	"""
		plotsol(fp::Array{Int64,1}, nodedict::Dict)

	Function that plots the final TSP solution NOK!
	"""
	function plotsol(fp::Array{Int64,1}, nodedict::Dict)
		#settings = Dict(:arrow => :arrow, :lw => 2, :color=>:blue, :label=>"")
		p = plot()
		startnode = pop!(filter(!in(fp),keys(nodedict)))
		plot!([nodedict[startnode][1]; nodedict[fp[1]][1]],
			  [nodedict[startnode][2]; nodedict[fp[1]][2]],
			  arrow=:arrow, color=:blue,label="")
		for i in 1:length(fp)-1
			x = [nodedict[fp[i]][1]; nodedict[fp[i+1]][1]]
			y = [nodedict[fp[i]][2]; nodedict[fp[i+1]][2]]
			plot!(x,y, arrow=:arrow, color=:blue,label="")
		end
		plot!([nodedict[fp[end]][1]; nodedict[startnode][1]],
			  [nodedict[fp[end]][2]; nodedict[startnode][2]],
			  arrow=:arrow, color=:blue,label="")
		xlims!(0,1)
		ylims!(0,1)
		return p
	end

end;

# ╔═╡ c6903b96-713c-11eb-3954-c7fda49d3e35
begin
	# Play and test case
	N = [(0.5, 0.9);(0.1,0.1); (0.9,0.1); (0.5,0.5)]
	nodedict = Dict( i => N[i] for i in 1:length(N)) # nodes dict
	dm = Float64[0 10 15 20;
		 10 0 35 25;
		 15 35 0 30;
		 20 25 30 0]
	G = SimpleWeightedGraph(dm)

	initstate = (1, Set(Array{Int64,1}()),4)
	goalstate = (1, Set(collect(1:length(N))),4)
	P = Problem(initstate, goalstate)

	let
		println("Astar graph search with for TSP with MST heuristic")
		@time (res, actions) = Astargraphsearch(P, findactions=findactions_TSP, 													applyaction=applyaction_TSP,
												goaltest=goaltest_TSP,
												g=g_TSP, h=h_TSP, wm=dm);

		plot(plotgraph(G,nodedict, wf=0.51, draw_weights=true), plotsol(actions, nodedict), size=(600,300))
		title!("TSP solution with graph\nsearch and MST heuristic")
end

end

# ╔═╡ 39578d8a-7167-11eb-0e98-314089370666
begin
	# specific generating functions
	"""
		problemgenerator(n::Int)

	generate n points in the unit square
	"""
	function problemgenerator(n::Int)
		return [(rand(), rand()) for _ in 1:n]
	end

	"""
		distancematrix(x::Array{Tuple{Float64,Float64},1}, p::Real=1)

	Generate distancematrix between all points.

	x is an array [n x 1] and p is the L_p-norm you want to use for the distancematrix
	"""
	function distancematrix(x::Array{Tuple{Float64,Float64},1}, p::Real=1)
		n = length(x)
		A = zeros(n,n)
		for (i,j) in collect(Base.Iterators.product(1:n, 1:n))
			A[i,j] = norm(x[i] .- x[j], p)
		end
		return A
	end
end;

# ╔═╡ 4cea4e00-7167-11eb-17de-abf6c558f9a9
let
	n = 15 # number of nodes
	N = problemgenerator(n) # nodes
	nodedict = Dict( i => N[i] for i in 1:length(N)) # nodes dict
	A = distancematrix(N,1) # distance matrix using manhattan norm
	G = SimpleWeightedGraph(A) # graph

	initstate = (1, Set(Array{Int64,1}()),n)
	goalstate = (1, Set(collect(1:length(N))),n)
	P = Problem(initstate, goalstate)

	println("Astar graph search with for TSP (n = $(n)) with MST heuristic")
		@time (res, actions) = Astargraphsearch(P, findactions=findactions_TSP, 													applyaction=applyaction_TSP,
												goaltest=goaltest_TSP,
												g=g_TSP, h=h_TSP, wm=A);

	plot(plotgraph(G,nodedict, wf=0.51), plotsol(actions, nodedict), size=(600,300))
	title!("TSP solution with graph\nsearch and MST heuristic")
end

# ╔═╡ Cell order:
# ╠═ea9468a2-710f-11eb-2284-6d8af448d9eb
# ╟─7782b3b2-710f-11eb-1a92-e12b8b1f3299
# ╟─58b8349c-7110-11eb-1aab-f7c716a12c57
# ╠═13d314dc-7110-11eb-3475-2126f5d414b7
# ╟─423469c0-7110-11eb-0cef-03540836525b
# ╠═68def13a-7110-11eb-29cf-8da0df67b178
# ╠═b5309106-7110-11eb-3631-c1484cc8a5d9
# ╠═c28473ea-7110-11eb-2129-cdc0aab5444e
# ╟─ccfd8924-7110-11eb-3890-6d9d80cfb96b
# ╠═d629933a-7110-11eb-39ad-7582f1db3d62
# ╟─902a810e-7111-11eb-1c4f-ef75bbcac693
# ╠═85073ccc-75e7-11eb-0d0f-c1c178e4e57e
# ╟─07f44e1c-713b-11eb-25b0-b9824cc5d251
# ╠═369047da-713b-11eb-3f47-29e522f17fb5
# ╠═9b3d2832-75e7-11eb-0372-35f7b9ed402b
# ╟─b35e5664-7143-11eb-1e75-17b00177d1f9
# ╟─a653d9a8-7143-11eb-3812-396c3ec732bd
# ╟─28f0121c-713c-11eb-2aff-e7a9c592fa53
# ╠═49d11f58-713c-11eb-20bb-e3a3dc889b6b
# ╠═99bbb244-713c-11eb-34be-d78f80834d46
# ╠═8b9fb7a0-7146-11eb-089c-f5c25325324b
# ╠═c6fda06e-713c-11eb-3587-47c7faed0d07
# ╠═c6903b96-713c-11eb-3954-c7fda49d3e35
# ╠═39578d8a-7167-11eb-0e98-314089370666
# ╠═4cea4e00-7167-11eb-17de-abf6c558f9a9


using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("colorscheme.jl"))

using InteractiveDynamics
using DynamicalSystems, OrdinaryDiffEq
import GLMakie

a = 0.2
b = 0.2
c = 5.7
amu = 0.02
ω1 = 1.1 + amu
ω2 = 1.1 - amu
k1 = 0.02
k2 = k1

y_init = [0.1, 0.2, 0., 0.11, 0.19, 0.1]

diffeq = (alg = Tsit5(), adaptive = false, dt = 0.02, reltol = 1e-9, abstol = 1e-9)

ds = Systems.coupled_roessler(y_init; ω1, ω2, a, b, c, k1, k2)

ps = Dict(
    1 => 1.0:0.01:1.5,
    2 => 1.0:0.01:1.5,
    5 => 0:0.01:6,
    6 => 0:0.001:0.2,
    7 => 0:0.001:0.2,
)
pnames = Dict(
    1 => "ωx",
    2 => "ωy",
    5 => "c",
    6 => "kx",
    7 => "ky",
)

u0s =  [ds.u0]

fig, obs = interactive_evolution_timeseries(
    ds, u0s, ps; tail = 1000, diffeq, idxs = [1,4], pnames,
    lims = ((-12,12), (-12,12))
)

ax = GLMakie.content(fig[1,1])
ax.xlabel = "x1"
ax.ylabel = "y1"

# %%
record_interaction(fig, string(@__FILE__)[1:end-2]*"mp4"; total_time = 30, sleep_time = 2)


"""A node which lags its input by a fixed number of knots."""
struct Lag{T} <: NodeOp{T}
    n::Int64
end

struct LagState{T} <: NodeEvaluationState
    # A buffer of the last number of values required.
    value_buffer::CircularBuffer{T}
end

function create_evaluation_state(::Tuple{Node}, op::Lag{T}) where {T}
    return LagState(CircularBuffer{T}(op.n))
end

function run_node!(
    op::Lag{T},
    state::LagState{T},
    ::DateTime,  # time_start
    ::DateTime,  # time_end
    input::Block{T},
) where {T}
    if isempty(input)
        # The input is empty, so regardless of the state we cannot lag onto anything.
        # The state is unmodified, and return the empty block.
        return input
    end

    w = op.n  # The window length.
    ni = length(input)  # The input length.
    ns = length(state.value_buffer)  # The state buffer length.

    # The output length.
    #   - at most the input size
    #   - at most total number of knots available in this batch: max((ns + ni) - w, 0)
    no = min(ni, max((ns + ni) - w, 0))

    # The number of elements to take from the state buffer.
    #   - at most the window size
    #   - at most the number of available elements in the buffer
    #   - at most the output size
    m = min(w, ns, no)

    times = @inbounds @view(input.times[(1 + ni - no):end])
    values = if m == 0
        @inbounds @view(input.values[1:no])
    else
        # Merge values from the state and in the Block.
        result = _allocate_values(T, no)
        for i in 1:m
            @inbounds result[i] = popfirst!(state.value_buffer)
        end
        @inbounds result[(m + 1):end] = @view(input.values[1:(no - m)])
        result
    end

    # Update the state with the remaining values in the input.
    @inbounds append!(state.value_buffer, @view(input.values[(1 + no - m):end]))

    return Block(unchecked, times, values)
end

"""
    lag(x::Node, n::Integer)

Construct a node which takes values from `x`, but lags them by `n` knots.

This means that we do not introduce any new timestamps that do not appear in `x`, however
we will not emit knots for the first `n` values that appear when evaluating `x`.
"""
function lag(x::Node, n::Integer)
    return if n == 0
        # Optimisation.
        x
    elseif n < 0
        throw(ArgumentError("Cannot lag by $n."))
    elseif _is_constant(x)
        # Constant nodes shouldn't be affected by lagging.
        x
    else
        obtain_node((x,), Lag{value_type(x)}(n))
    end
end

"""
    diff(x::Node[, n=1])

Compute the `n`-knot difference of `x`, i.e. `x - lag(x, n)`.
"""
function Base.diff(x::Node, n::Integer=1)
    n > 0 || throw(ArgumentError("n must be positive, got $n"))
    return x - lag(x, n)
end


using DataFrames, CSV, Statistics
data = CSV.read("data/purchaseData.csv", DataFrame)

rowsKeep = .!(ismissing.(data.Grade) .& ismissing.(data.Price))
data = data[rowsKeep, :]

replace!(x -> ismissing(x) ? "QQ" : x, data.Name)
replace!(x -> ismissing(x) ? "31/06/2008" : x, data.Date)

grPr = combine(groupby(dropmissing(data), :Grade),
    :Price=> (x -> round(mean(x), digits=-3)) => :AvgPrice)

d = Dict(grPr[:,1] .=> grPr[:,2])
nearIndx(v, x) = findmin(abs.(v.-x))[2]
for i in 1:nrow(data)
    if ismissing(data[i, :Price])
        data[i, :Price] = d[data[i, :Grade]]
    end
    if ismissing(data[i, :Grade])
        data[i, :Grade] = grPr[ nearIndx(grPr[:,2], data[i, :Price]), :Grade]
    end
end

println(first(data, 5), "\n")
println(describe(data))


#
# Parrot.jl/src/utils.jl
#
# By Elijah Mathews (me@elijahmathews.com)
# openpgp4fpr:ac1d3fb1e8a5eb7d14bd587b2932c725055a90d8
#

"""
    splitindices(frac::Real, N::Integer)

Generate two sets of unique indices of cumulative length `N`.
Returns a tuple:

    (indices1, indices2)

where `indices1` is an array of unique indices of length `frac × N` and
`indices2` is an array of unique indices of length `(1 - frac) × N`. Can
be used for generating indices to separate training and test datasets, for
example.
"""
function splitindices(frac, N::Integer)

    @assert 0 <= frac <= 1

    # Generate unique random indices.
    base = randperm(N)

    indices1 = base[1:ceil(Int, frac * N)]
    indices2 = base[(ceil(Int, frac * N) + 1):end]

    return indices1, indices2

end

"""
    fractionalerrorquantiles(fractionalerror; quants=[0.0005, 0.005, 0.025, 0.5, 0.975, 0.995, 0.9995])

Compute the 0.05%, 0.5%, 2.5%, 50% (median), 97.5%, 99.5%, and 99.95% quantiles
for a set of spectral energy distributions (SEDs) predicted by Parrot given an
array of the fractional errors `fractionalerror`.

The selected quantiles can be changed from their defaults by passing an array to
the `quants` keyword argument.
"""
function fractionalerrorquantiles(fractionalerror; quants=[0.0005, 0.005, 0.025, 0.5, 0.975, 0.995, 0.9995])

    returnval = []

    for i in 1:length(fractionalerror)

        result = zeros(length(quants), size(fractionalerror[i], 1))

        for j in 1:size(fractionalerror[i], 1)

            for k in 1:length(quants)

                result[k,j] = quantile(fractionalerror[i][j,:], quants[k])

            end

        end

        push!(returnval, result)

    end

    return returnval

end

"""
    convert(MultivariateStats.PCA{Float32}, pca::MultivariateStats.PCA{Float64})

Convert a `Float64` PCA model `pca` to a `Float32` PCA model.
"""
function Base.convert(::Type{MultivariateStats.PCA{Float32}}, pca::MultivariateStats.PCA{Float64})
    MultivariateStats.PCA(
        convert(Array{Float32,1}, pca.mean    ),
        convert(Array{Float32,2}, pca.proj    ),
        convert(Array{Float32,1}, pca.prinvars),
        convert(Float32,          pca.tprinvar),
        convert(Float32,          pca.tvar    ),
    )
end


"""
    BetaPlane{T} <: AbstractRotation

A parameter object for meridionally increasing Coriolis parameter (`f = f₀ + βy`)
that accounts for the variation of the locally vertical component of the rotation
vector with latitude.
"""
struct BetaPlane{T} <: AbstractRotation
    f₀ :: T
     β :: T
end

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

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

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

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

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

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

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

@inline x_f_cross_U(i, j, k, grid, coriolis::BetaPlane, U) =
    @inbounds - (coriolis.f₀ + coriolis.β * grid.yC[j]) * ℑxyᶠᶜᵃ(i, j, k, grid, U[2])

@inline y_f_cross_U(i, j, k, grid, coriolis::BetaPlane, U) =
    @inbounds   (coriolis.f₀ + coriolis.β * grid.yF[j]) * ℑxyᶜᶠᵃ(i, j, k, grid, U[1])

@inline z_f_cross_U(i, j, k, grid::AbstractGrid{FT}, coriolis::BetaPlane, U) where FT = zero(FT)

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


# This script shows stansummary ESS vakues vs. MCMCChains ESS values

using CmdStan, StatsPlots

ProjDir = @__DIR__
cd(ProjDir)

berstanmodel = "
data {
  int<lower=0> N;
  int<lower=0,upper=1> y[N];
}
parameters {
  real<lower=0,upper=1> theta;
}
model {
  theta ~ beta(1,1);
  for (n in 1:N)
    y[n] ~ bernoulli(theta);
}
"

data = Dict("y" => [0,1,0,0,0,0,0,0,0,1],
  "N" => 10)

stanmodel = Stanmodel(
   name = "berstanmodel", model = berstanmodel, nchains = 4,
   printsummary=false, tmpdir=mktempdir());

ess_array_mcmcchains = []
ess_array_cmdstan = []

for i in 1:100
  rc, chns, cnames = stan(stanmodel,
     data, summary=true, ProjDir);
  dfc = describe(chns)[1]
  append!(ess_array_mcmcchains, dfc[:theta, :ess])
  rs = read_summary(stanmodel)
  append!(ess_array_cmdstan, rs[:theta, :ess])
end

p = Array{Plots.Plot{Plots.GRBackend}}(undef, 1);
p[1] = plot(ess_array_cmdstan, lab="CmdStan", xlim=(0, 170))
p[1] = plot!(ess_array_mcmcchains, lab="MCMCChains")
plot(p..., layout=(1,1))
savefig("ess_comparison_plot.pdf")


"""
    baggr_oob(object::Baggr, X, Y; score = rmsep)
Compute the out-of-bag (OOB) error after bagging a regression model.
* `object` : Output of a bagging.
* `X` : X-data used in the bagging.
* `Y` : Y-data used in the bagging.
* `score` : Function computing the prediction error (default: msep).

## References

Breiman, L., 1996. Bagging predictors. Mach Learn 24, 123–140.
https://doi.org/10.1007/BF00058655

Breiman, L., 2001. Random Forests. Machine Learning 45, 5–32.
https://doi.org/10.1023/A:1010933404324

Genuer, R., 2010. Forêts aléatoires : aspects théoriques,
sélection de variables et applications. PhD Thesis. Université Paris Sud - Paris XI.

Gey, S., 2002. Bornes de risque, détection de ruptures, boosting :
trois thèmes statistiques autour de CART en régression (These de doctorat).
Paris 11. http://www.theses.fr/2002PA112245
"""
function baggr_oob(object::Baggr, X, Y; score = rmsep)
    X = ensure_mat(X)
    Y = ensure_mat(Y)
    n = size(X, 1)
    q = size(Y, 2)
    rep = length(object.fm)
    pred_oob = similar(X, n, q)
    k = zeros(n)
    @inbounds for i = 1:n
        zpred = fill(0., 1, q)
        zk = 0
        @inbounds for j = 1:rep
            soob = object.s_oob[j]
            if in(i, soob)
                zpred .+= predict(object.fm[j], X[i:i, vcol(object.s_col, j)]).pred
                zk += 1
            end
        end
        pred_oob[i, :] = zpred / zk
        k[i] = zk
    end
    scor = score(pred_oob, Y)
    (scor = scor, pred_oob, k)
end

"""
    baggr_vi(object::Baggr, X, Y; score = rmsep)
Compute variance importances (with permutation method) after bagging a regression model.
* `object` : Output of a bagging.
* `X` : X-data that was used in the model bagging.
* `Y` : Y-data that was used in the model bagging.

Variances importances are computed
by permuting sucessively each column of the out-of-bag (X_OOB),
and by looking at the effect on the MSEP.

## References

Breiman, L., 2001. Random Forests. Machine Learning 45, 5–32.
https://doi.org/10.1023/A:1010933404324

Genuer, R., 2010. Forêts aléatoires : aspects théoriques,
sélection de variables et applications. PhD Thesis. Université Paris Sud - Paris XI.

Gey, S., 2002. Bornes de risque, détection de ruptures, boosting :
trois thèmes statistiques autour de CART en régression. PhD Thesis.
Univ. Paris 11. http://www.theses.fr/2002PA112245
"""
function baggr_vi(object::Baggr, X, Y; score = rmsep)
    X = ensure_mat(X)
    Y = ensure_mat(Y)
    p = size(X, 2)
    q = size(Y, 2)
    rep = length(object.s_oob)
    ncol = size(object.s_col, 1)
    scol = similar(object.s_col, ncol)
    res = similar(X, p, rep, q)
    @inbounds for i = 1:rep
        soob = object.s_oob[i]  # soob has a variable length
        m = length(soob)
        scol .= vcol(object.s_col, i)
        zpred = predict(object.fm[i], X[soob, scol]).pred
        zY = @view(Y[soob, :]) ;
        zscore = score(zpred, zY)
        zX = similar(X[soob, :])
        @inbounds for j = 1:p
            zX .= X[soob, :]
            s = sample(1:m, m, replace = false)
            zX[:, j] .= zX[s, j]
            zpred .= predict(object.fm[i], zX[:, scol]).pred
            zscore_perm = score(zpred, zY)
            res[j, i, :] = zscore_perm .- zscore
        end
    end
    imp = reshape(mean(res, dims = 2), p, q)
    (imp = imp, res = res)
end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.7.0-DEV
#     language: julia
#     name: julia-1.7
# ---

# %%
@show VERSION
@time using DifferentialEquations
@time using Plots
@time using ProgressMeter

# %% [markdown]
# Wave equaiton: $\displaystyle
# \frac{\partial^2 u}{\partial t^2} =
# \frac{\partial^2 u}{\partial x^2} +
# \frac{\partial^2 u}{\partial y^2}
# $

# %%
# the wave equation on the square with free boundary condition

#f!(du, v, u, p, t) = @. du = v

function g!(dv, v, u, p, t)
    h⁻² = p[1]
    @. @views dv[2:end-1,2:end-1] = h⁻²*(
        u[1:end-2,2:end-1] + u[3:end,2:end-1] + u[2:end-1,1:end-2] + u[2:end-1,3:end] - 4u[2:end-1,2:end-1])
    @. @views dv[1  ,2:end-1] = h⁻²*(u[2,    2:end-1] + u[1,  1:end-2] + u[1,  3:end] - 3u[1,  2:end-1])
    @. @views dv[end,2:end-1] = h⁻²*(u[end-1,2:end-1] + u[end,1:end-2] + u[end,3:end] - 3u[end,2:end-1])
    @. @views dv[2:end-1,1  ] = h⁻²*(u[2:end-1,2    ] + u[1:end-2,1  ] + u[3:end,1  ] - 3u[2:end-1,1  ])
    @. @views dv[2:end-1,end] = h⁻²*(u[2:end-1,end-1] + u[1:end-2,end] + u[3:end,end] - 3u[2:end-1,end])
    dv[1  ,1  ] = h⁻²*(u[1  ,2    ] + u[2    ,1  ] - 2u[1  ,1  ])
    dv[end,1  ] = h⁻²*(u[end,2    ] + u[end-1,1  ] - 2u[end,1  ])
    dv[1  ,end] = h⁻²*(u[1  ,end-1] + u[2    ,end] - 2u[1  ,end])
    dv[end,end] = h⁻²*(u[end,end-1] + u[end-1,end] - 2u[end,end])
end

x = y = range(-1, 1; length=201)
h = step(x)
p = (1/h^2,)

f(x, y) = exp(-(x^2+y^2)/2)
u0 = f.(x', y)
v0 = zero(u0)
tspan = (0.0, 10.0)

#@time prob = DynamicalODEProblem(g!, f!, v0, u0, tspan, p)
@time prob = SecondOrderODEProblem(g!, v0, u0, tspan, p)
sol = nothing; GC.gc()
@time sol = solve(prob; saveat=0.02)
sol = nothing; GC.gc()
@time sol = solve(prob; sageat=0.02)
sol = nothing; GC.gc()
@time sol = solve(prob; saveat=0.02);

# %%
x = y = range(-1, 1; length=201)
h = step(x)
p = (1/h^2,)

f(x, y) = exp(-(x^2+y^2)/2)
u0 = f.(x', y)
v0 = zero(u0)
tspan = (0.0, 20.0)

#@time prob = DynamicalODEProblem(g!, f!, v0, u0, tspan, p)
@time prob = SecondOrderODEProblem(g!, v0, u0, tspan, p)
sol = nothing; GC.gc()
@time sol = solve(prob; saveat=0.02);

# %%
umin = minimum(minimum.(u.x[2] for u in sol.u))
umax = maximum(maximum.(u.x[2] for u in sol.u))
umin, umax

# %%
L = 201
ts = range(tspan...; length=L) |> r -> [fill(r[1], 20); r; fill(r[end], 20)]
prog = Progress(length(ts), 0)
@gif for t in ts
    heatmap(x, y, sol(t).x[2]'; clim=(umin, umax), colorbar=false, size=(300, 316), c=:CMRmap)
    title!("t = $(round(t, digits=1))"; titlefontsize=12, aspectratio=1)
    next!(prog)
end

# %%
L = 201
ts = range(tspan...; length=L) |> r -> [fill(r[1], 20); r; fill(r[end], 20)]
prog = Progress(length(ts), 0)
@gif for t in ts
    surface(x, y, sol(t).x[2]'; zlim=(umin, umax), colorbar=false, c=:CMRmap, camera=(30, 60))
    title!("t = $(round(t, digits=1))"; titlefontsize=12)
    next!(prog)
end

# %%
L = 201
ts = range(tspan...; length=L) |> r -> [fill(r[1], 20); r; fill(r[end], 20)]
prog = Progress(length(ts), 0)
@gif for t in ts
    wireframe(x[begin:4:end], y[begin:4:end],
        sol(t).x[2][begin:4:end, begin:4:end]';
        zlim=(umin, umax), colorbar=false, c=:CMRmap, camera=(30, 60))
    title!("t = $(round(t, digits=1))"; titlefontsize=12)
    next!(prog)
end

# %%
g(x, y) = exp(-x^2/2-y^2/4)
u0 = g.(x', y)
v0 = zero(u0)
tspan = (0.0, 20.0)

@time prob = SecondOrderODEProblem(g!, v0, u0, tspan, p)
sol = nothing; GC.gc()
@time sol = solve(prob; saveat=0.02)

umax = maximum(maximum.(u.x[2] for u in sol.u))
umin = minimum(minimum.(u.x[2] for u in sol.u))

L = 201
ts = range(tspan...; length=L) |> r -> [fill(r[1], 20); r; fill(r[end], 20)]
prog = Progress(length(ts), 0)
@gif for t in ts
    surface(x, y, sol(t).x[2]'; zlim=(umin, umax), colorbar=false, c=:CMRmap, camera=(30, 60))
    title!("t = $(round(t, digits=1))"; titlefontsize=12)
    next!(prog)
end

# %%
L = 201
ts = range(tspan...; length=L) |> r -> [fill(r[1], 20); r; fill(r[end], 20)]
prog = Progress(length(ts), 0)
@gif for t in ts
    wireframe(x[begin:4:end], y[begin:4:end],
        sol(t).x[2][begin:4:end, begin:4:end]';
        zlim=(umin, umax), colorbar=false, c=:CMRmap, camera=(30, 60))
    title!("t = $(round(t, digits=1))"; titlefontsize=12)
    next!(prog)
end

# %%

0003/Wave equation on the square 2021-06-09 Part 1.jl


using CSV, DataFrames, Dates

function load_board(filename)
    board = CSV.read(filename, DataFrame, header = 0) |> Array{Int8, 2}
    return board
end

function print_board(board::Array{Int8, 2})
    for i in 1:9
        if (i-1) % 3 == 0 && i != 0
            println("- - - - - - - - - - - -")
        end

        for j in 1:9
            if (j-1) % 3 == 0 && j != 1
                print(" | ")
            end
            if j == 9
                println(board[i, j])
            else
                print("$(board[i, j]) ")
            end
        end
    end
end

function solve(board::Array{Int8, 2})
    find = find_empty(board)
    if find == (10, 10)
        return true
    else
        row, col = find
        for i in 1:9
            if valid(board, i, (row, col))
                board[row, col] = i
                if solve(board)
                    return true
                else
                    board[row, col] = 0
                end
            end
        end
    end
    return false
end

function valid(board::Array{Int8, 2}, num::Int8, pos::Tuple{Int8, 1})
    # Check column
    num_in_col = findall( x -> x == num, board[:pos[2]] )
    if length(num_in_col) > 0
        return false
    end

    # Check row
    num_in_row = findall( x -> x == num, board[pos[1]:end] )
    if length(num_in_row) > 0
        return false
    end

    # Check box
    coordinates = Dict( 1 => 1, 2 => 1, 3 => 1,
                        4 => 4, 5 => 4, 6 => 4,
                        7 => 7, 8 => 7, 9 => 7)
    box = board[coordinates[pos[1]]:(coordinates[pos[1]]+2),
          coordinates[pos[2]]:(coordinates[pos[2]]+2)]
    num_in_box = findall( x -> x == num, box )

    if length(num_in_box) > 0
        return false
    end

    return true
end

function find_empty(board::Array{Int8, 2})
    for i in 1:9
        for j in 1:9
            if board[i, j] == 0
                return (i, j)  # row, col
            end
        end
    end
    return (10,10)
end

sudoku = load_board(raw"C:\Users\kaszt\github\Julia_Sudoku\9x9.csv")
println("\nOriginal:\n")
print_board(sudoku)
start_time = Time(Dates.now())
solve(sudoku)
end_time = Time(Dates.now())
println("\nSolved:\n")
print_board(sudoku)
time_taken = end_time - start_time
println(Dates.Millisecond(time_taken))


module NamedVecs

export NamedVec

"""
A `NamedVec` is a kind of hybrid between a `Vector` and a `NamedTuple`. This can be useful,
for example, in mathematical modelling where collections of parameters are some times treated
as a single vector but at other times acessed individually.

`NamedVec` supports vector operations, such as `+` and `*`, that are useful in linear algebra.

At the same time, a `NamedVec` can be indexed by field names to obtain reshaped views into
parts of the vector.

Example:
```
v = NamedVec((a=[1], b=[2, 3]))
v' * v # returns 14
f((;a, b)) = 2a .+ b # This notation requires Julia 1.7
f(v) # returns [4, 5]
```
"""
struct NamedVec{T,Names,D<:AbstractVector{T},M} <: AbstractVector{T}
    data::D
    maps::NamedTuple{Names, M}
end

# Constructors

function NamedVec(data::AbstractVector, names::NTuple{N, Symbol}, maps::Tuple{Vararg{Any,N}}) where {N}
    NamedVec(data, NamedTuple{names}(maps))
end

"""
A `NamedVec` can be created from a `NamedTuple` of arrays. The arrays are copied (not aliased) into the new object.
"""
function NamedVec(::Type{T}, xs::NamedTuple{Names, <:Tuple{Any, Vararg{Any}}}) where {T<:AbstractVector, Names}
    ns = map(length, Tuple(xs)) # TODO: should look at number of elements actually used.
    lastix = cumsum(ns)
    firstix = lastix .- ns .+ 1
    data = similar(T isa UnionAll ? T{promote_type(eltype.(Tuple(xs))...)} : T, (sum(ns),))
    maps = ntuple(length(xs)) do i
        (xs[i] isa Real ? firstix[i] : firstix[i]:lastix[i]) => viewfun(xs[i])
    end
    v = NamedVec(data, Names, maps)
    for n in Names
        setproperty!(v, n, getproperty(xs, n))
    end
    v
end

"NamedVec, if the vector type is not specified, uses `Vector`."
function NamedVec(xs::NamedTuple{Names, <:Tuple{Any, Vararg{Any}}}) where {Names}
    E = promote_type(eltype.(Tuple(xs))...)
    NamedVec(Vector{E}, xs)
end

NamedVec(::Type{T}; kwargs...) where {T<:AbstractVector} = NamedVec(T, NamedTuple(kwargs))
NamedVec(; kwargs...) = NamedVec(NamedTuple(kwargs))

"""
`viewfun(x)` - Returns a function that create an object that is like `x`, from a vector
with the correct number of elements.
"""
function viewfun(x::AbstractArray)
    let sz = size(x)
        v -> reshape(v, sz)
    end
end

function viewfun(x::NamedVec)
    let m = maps(x)
        v -> NamedVec(v, m)
    end
end

viewfun(::AbstractVector{<:Real}) = identity

viewfun(::Real) = getindex

# Conversion

"`vec(v::NamedVec)` returns the underlying data vector of a NamedVec"
Base.vec(v::NamedVec) = getfield(v, :data)

"oftype(x::NamedVec, y) converts a vector y into a NamedVec with the names and indices of x"
function Base.oftype(v::NamedVec{Names, <:Tuple{Vararg{AbstractArray}}}, w::AbstractVector) where {Names}
    @boundscheck eachindex(v) == eachindex(w) || error("Size missmatch.")
    NamedVec(w, maps(v))
end

"""
`Tuple(v::NamedVec)` returns a tuple of the views into `v`.
For a tuple of the elements of the vector, use `Tuple(vec(v))`.
"""
function Base.Tuple(v::NamedVec)
    let data=vec(v)
        map(m->m[2](view(data, m[1])), maps(v))
    end
end

"`NamedTuple(v::NamedVec)` returns a `NamedTuple` of the views into `v`."
Base.NamedTuple(v::NamedVec{<:Any,Names}) where {Names} = NamedTuple{Names}(Tuple(v))

# Display

function Base.show(io::IO, ::MIME"text/plain", v::NamedVec)
    println(io, length(v), "-element NamedVec{", eltype(v), "}:")
    show(io, NamedTuple(v))
end

# Vector-like

Base.eltype(::NamedVec{T}) where {T} = T
Base.length(v::NamedVec) = length(vec(v))
Base.eachindex(v::NamedVec) = eachindex(vec(v))
Base.size(v::NamedVec) = size(vec(v))
Base.size(v::NamedVec, dim::Integer) = size(vec(v), dim)
Base.getindex(v::NamedVec, ix::Integer) = getindex(vec(v), ix)
Base.setindex!(v::NamedVec, y, ix::Integer) = setindex!(vec(v), y, ix)

Base.:*(a::Number, v::NamedVec) = NamedVec(a * vec(v), maps(v))
Base.:+(a::NamedVec{Names}, b::NamedVec{Names}) where {Names} = NamedVec(vec(a) + vec(b), maps(a))

Base.similar(a::NamedVec) = NamedVec(similar(vec(a)), maps(a))
Base.copy(a::NamedVec) = NamedVec(copy(vec(a)), maps(a))

# Vector-like broadcasting

struct NamedVecStyle <: Broadcast.AbstractArrayStyle{1}
end

Base.Broadcast.BroadcastStyle(::Type{<:NamedVec}) = NamedVecStyle()
Base.Broadcast.BroadcastStyle(a::NamedVecs.NamedVecStyle, ::Base.Broadcast.DefaultArrayStyle{0}) = a
Base.Broadcast.BroadcastStyle(a::NamedVecs.NamedVecStyle, ::Base.Broadcast.DefaultArrayStyle{1}) = a
Base.Broadcast.BroadcastStyle(::NamedVecs.NamedVecStyle, b::Base.Broadcast.DefaultArrayStyle) = b

function Base.similar(bc::Broadcast.Broadcasted{NamedVecStyle}, ::Type{T}) where {T}
    # TOOD: We should maybe check that there are not conflicting sets of names.
    v = getnv(bc)
    NamedVec(similar(vec(v), T), maps(v))
end

"getnv extracts the first NamedVec it can find in a Broadcasted object"
getnv(::Any) = nothing
getnv(x::NamedVec) = x
function getnv(bc::Broadcast.Broadcasted{NamedVecStyle})
    for a in bc.args
        y = getnv(a)
        y isa NamedVec && return y
    end
    nothing
end

# NamedTuple-like

Base.propertynames(::NamedVec{<:Any, Names}, ::Bool=false) where {Names} = Names

function Base.getproperty(v::NamedVec, s::Symbol)
    ix, f = getproperty(maps(v), s)
    f(view(vec(v), ix))
end

function Base.setproperty!(v::NamedVec, s::Symbol, y)
    ix, f = getproperty(maps(v), s)
    if f === getindex
        vec(v)[ix] = y
    else
        f(view(vec(v), ix)) .= y
    end
end

# Utility functions
# (Unexported and mainly for internal use)

maps(v::NamedVec) = getfield(v, :maps)

end # modulev


module Multithread

using Base.Threads
import Random
import Future

# thread-local seeds for rand()
function seed_per_thread() :: Vector{Random.MersenneTwister}
    m = Random.MersenneTwister(1)
    return [m; accumulate(Future.randjump,
                          fill(big(10)^20,
                          nthreads()-1),
                          init=m)]
end

# returns the number of simulations
# for the current thread
function thread_num_sims(total_num_sims::Integer)
    n = nthreads()
    dv = div(total_num_sims, n)
    md = mod(total_num_sims, n)

    if dv == 0
        # degenerates to single-threaded
        if threadid() == 1
            return total_num_sims
        else
            return 0
        end
    else
        if threadid() != n
            return dv
        else
            return dv + md
        end
    end
end

# checks if we're running Julia in multithread mode
function check_julia_num_threads()
    if nthreads() == 1
        error("Run `export JULIA_NUM_THREADS=n`, n > 1, before launching Julia.")
    end
end

# with threading
function run_multithread(num_sims::T) :: Float64 where {T<:Integer}
    check_julia_num_threads()
    @assert num_sims > 0

    local acum = Threads.Atomic{Int}(0)
    local check_num_sims = Threads.Atomic{T}(0)

    seeds = seed_per_thread()

    @threads for tid in 1:nthreads()
        @assert tid == threadid()

        thread_seed = seeds[threadid()]
        thread_acum = 0
        local_num_sims = thread_num_sims(num_sims)

        for sim in 1:local_num_sims

            x = rand(thread_seed)
            y = rand(thread_seed)

            r = sqrt(x^2 + y^2)

            if r <= 1.0
                thread_acum += 1
            end
        end

        Threads.atomic_add!(acum, thread_acum)
        Threads.atomic_add!(check_num_sims, local_num_sims)
    end

    @assert check_num_sims[] == num_sims

    area = acum[] / num_sims
    return 4*area
end

end # module Threaded


export MyLinOp

"""
`MyLinOp(domainType::Type, dim_in::Tuple, [domainType::Type,] dim_out::Tuple, Fwd!::Function, Adj!::Function)`

Construct a user defined `LinearOperator` by specifing its linear mapping `Fwd!` and its adjoint `Adj!`. The functions `Fwd!` and `Adj` must be in-place functions consistent with the given dimensions `dim_in` and `dim_out` and the domain and codomain types.

```julia
julia> n,m = 5,4;

julia> A = randn(n,m);

julia> op = MyLinOp(Float64, (m,),(n,), (y,x) -> mul!(y,A,x), (y,x) -> mul!(y,A',x))
A  ℝ^4 -> ℝ^5

julia> op = MyLinOp(Float64, (m,), Float64, (n,), (y,x) -> mul!(y,A,x), (y,x) -> mul!(y,A',x))
A  ℝ^4 -> ℝ^5

```

"""
struct MyLinOp{N,M,C,D} <: LinearOperator
	dim_out::NTuple{N,Int}
	dim_in::NTuple{M,Int}
	Fwd!::Function
	Adj!::Function
end

# Constructors

MyLinOp(domainType::Type, dim_in::NTuple{N,Int}, dim_out::NTuple{M,Int},
	Fwd!::Function, Adj!::Function ) where {N,M} =
MyLinOp{N,M, domainType, domainType}(dim_out, dim_in, Fwd!, Adj! )

MyLinOp(domainType::Type, dim_in::NTuple{N,Int}, codomainType::Type, dim_out::NTuple{M,Int},
	Fwd!::Function, Adj!::Function ) where {N,M} =
MyLinOp{N,M, domainType, codomainType}(dim_out, dim_in, Fwd!, Adj! )

# Mappings

mul!(y::Array{C,N}, L::MyLinOp{N,M,C,D}, b::Array{D,M}) where {N,M,C,D} = L.Fwd!(y,b)
mul!(y::Array{C,N}, L::AdjointOperator{MyLinOp{N,M,C,D}}, b::Array{D,M}) where {N,M,C,D} = L.A.Adj!(y,b)

# Properties

size(L::MyLinOp) = (L.dim_out, L.dim_in)

codomainType(L::MyLinOp{N,M,C,D}) where {N,M,C,D} = C
  domainType(L::MyLinOp{N,M,C,D}) where {N,M,C,D} = D

fun_name(L::MyLinOp)  = "A"