using JLD2, Plots, Printf, LinearAlgebra, Statistics
using LaTeXStrings
include(pwd() * "/analysis_scripts/" * "post_analysis.jl")

filename = "/Weno_20_zonal_averages.jld2"
file = jldopen( pwd() * filename)
data_dictionary = get_data(file)
fine_z  = file["grid"]["zC"][2:end-1]
fine_y  = file["grid"]["yC"][2:end-1]
fine_yF = file["grid"]["yF"][2:end-1]
fine_zF = file["grid"]["zF"][2:end-1]
fine_b = mean(data_dictionary["b"][2:end], dims = 1)[1]
fine_u = mean(data_dictionary["u"][2:end], dims = 1)[1]
fine_v = mean(data_dictionary["v"][2:end], dims = 1)[1]
fine_v = (fine_v[2:end,:] + fine_v[1:end-1,:]) ./ 2
data_dictionary = get_data(file)
close(file)

filename = "/Weno_CPU_20_zonal_averages.jld2"
file = jldopen( pwd() * filename)
data_dictionary = get_data(file)
coarse_z  = file["grid"]["zC"][2:end-1]
coarse_y  = file["grid"]["yC"][2:end-1]
coarse_yF = file["grid"]["yF"][2:end-1]
coarse_zF = file["grid"]["zF"][2:end-1]
coarse_b = mean(data_dictionary["b"][2:end], dims = 1)[1]
coarse_u = mean(data_dictionary["u"][2:end], dims = 1)[1]
coarse_v = mean(data_dictionary["v"][2:end], dims = 1)[1]
coarse_v = (coarse_v[2:end,:] + coarse_v[1:end-1,:]) ./ 2
close(file)

filename = "/Weno_CPU_Hack_20_zonal_averages.jld2"
file = jldopen( pwd() * filename)
data_dictionary = get_data(file)
hack_z  = file["grid"]["zC"][2:end-1]
hack_y  = file["grid"]["yC"][2:end-1]
hack_yF = file["grid"]["yF"][2:end-1]
hack_zF = file["grid"]["zF"][2:end-1]
hack_b = mean(data_dictionary["b"][50:end], dims = 1)[1]
hack_u = mean(data_dictionary["u"][50:end], dims = 1)[1]
hack_v = mean(data_dictionary["v"][50:end], dims = 1)[1]
hack_v = (hack_v[2:end,:] + hack_v[1:end-1,:]) ./ 2
close(file)

function plot_field(a, y, yF, z, zF; name = " ", clims = nothing)
    py = get_y(a, y, yF) # global scope
    pz = get_z(a, z, zF) # global scope
    if clims == nothing
        clims = (minimum(a), maximum(a))
    end
    further_label = "Zonal and Temporal Average of "
    p1 = contourf(py, pz, a',
    color = :thermometer, title = " " * name,
    xlabel = "South to North [m]", ylabel = "Depth [m]"
    , clims = clims, linewidth = 0)
    display(p1)
    return p1
end

coarse_plot = plot_field(coarse_b, coarse_y, coarse_yF, coarse_z, coarse_zF, name = "Coarse Buoyancy")
fine_plot = plot_field(fine_b, fine_y, fine_yF, fine_z, fine_zF, name = "Coarse Buoyancy")

function avgy(Φ, n)
    m = size(Φ)[1]
    scale = Int(floor(m/n))
    if ( abs(Int(floor(m/n)) - m/n) > eps(1.0))
        return error
    end
    ny, nz = size(Φ)
    Φ2 = zeros(n, nz)
    for i in 1:n
        Φ2[i,:] .= 0
            for j in 1:scale
                Φ2[i,:] .+= Φ[scale*(i-1) + j,:] / scale
            end
    end
    return Φ2
end

coarse_grained_b = avgy(fine_b, 12)
coarse_grained_plot = plot_field(coarse_grained_b, coarse_y, coarse_yF, coarse_z, coarse_zF, name = "Coarse Grained Buoyancy")

norm(coarse_grained_b - coarse_b) / norm(coarse_grained_b)
norm(coarse_grained_b - hack_b) / norm(coarse_grained_b)

residual_plot = plot_field(coarse_b-coarse_grained_b , coarse_y, coarse_yF, coarse_z, coarse_zF, name = "Slack Buoyancy")

## Zonal Velocity
coarse_plot = plot_field(coarse_v, coarse_y, coarse_yF, coarse_z, coarse_zF, name = "Coarse Meridional Velocity")
fine_plot = plot_field(fine_v, fine_y, fine_yF, fine_z, fine_zF, name = "Meridional Velocity")
coarse_grained_v = avgy(fine_v, 12)
clims = (minimum(coarse_v), maximum(coarse_v))
coarse_grained_plot = plot_field(coarse_grained_v - coarse_v, coarse_y, coarse_yF,
                                 coarse_z, coarse_zF,
                                 name = "Difference in Meridional Velocity")


module DynaUtils

include("DynaPlanning.jl")

using LightGraphs, LightGraphsFlows, SimpleWeightedGraphs
using GraphPlot, Plots
using Colors
using Compose, Cairo
using SparseArrays
using DataStructures
using Random
using Logging

import .DynaPlanning

"""
price_of_anarchy(r, c)

Compute price of anarchy ρ(G,r,c) for instances with linear cost functions. r and c are vectors of traffic rates and cost functions, respectively.

"""
function price_of_anarchy(r, c)

    C_f = sum([c[i](r[i]) * r[i] for i = 1:length(r)])

    # Optimal flow for the instance (G,r/2,c)
    r_star = 0.5 * r
    C_f_star = sum([c[i](r_star[i]) * r_star[i] for i = 1:length(r_star)])

    # Augmentation to (G,r,c)
    δ = 1
    C_f_star_augm = δ * (C_f - C_f_star)

    return C_f / C_f_star_augm

end

"""
    find_flow(G)

Compute the edge flow for graph `G`.

"""
function find_flow(SInstance)

    A = adjacency_matrix(SInstance.network)

    I, J, vals = findnz(A)

    # init f_e
    f_e = sparse(I, J, zeros(length(vals)))

    for a in SInstance.agents
        thisAgent = a
        thisPath = thisAgent.path
        edge_k = first(thisPath)
        f_e[edge_k.src, edge_k.dst] += 1.0
    end

    # normalize
    f_e = sparse(I, J, findnz(A)[3] / sum(findnz(A)[3]))

    return f_e
end

"""
    find_supply_demand(G)

Compute the supply and the demand for `SInstance`.

"""
function find_supply_demand(SInstance)

    # init d
    d = zeros(nv(SInstance.network), 1)

    for a in SInstance.agents
        thisAgent = a
        thisPath = thisAgent.path

        @debug thisPath

        edge_src = first(thisPath)
        d[edge_src.src] += 1.0

        edge_sink = last(thisPath)
        d[edge_sink.dst] -= 1.0
    end

    # normalize
    total_demand = sum([x] for x in d if x > zero(x))

    d *= 1 / total_demand

    @debug "Supply/demand:" d

    return d
end

"""
    recursive_DFS!(genericgraph, visited, sink)

Compute -recursively- the list of paths between `visited` (source) and `sink` (destination) pairs for `genericgraph`.

"""
function recursive_DFS!(pathList, genericgraph, visited, sink)

    nodes = LightGraphs.neighbors(genericgraph, visited[length(visited)])

    for i in nodes
        if i == sink
            push!(visited, i)
            push!(pathList, visited)
            pop!(visited)
        elseif !in(i, visited)
            push!(visited, i)
            recursive_DFS!(pathList, genericgraph, visited, sink)
            pop!(visited)
        end
    end

end

"""
    find_all_paths(G, src, dst)

Compute -recursively- all paths between `src` (source) and `dst` (destination) pairs for graph `G`.

"""
function find_all_paths(G, src, dst)

    A = adjacency_matrix(G)

    all_paths = Array{Array{Any},1}()

    DynaPlanning.recursive_DFS!(all_paths, SimpleDiGraph(A), [src], dst)

    return all_paths

end

function plot_graph(G, locs_x, locs_y)

    A = adjacency_matrix(G)

    nodelabel = collect(1:nv(SimpleDiGraph(A)))

    draw(
        PNG("graph.png", 10cm, 10cm),
        gplot(SimpleDiGraph(A), locs_x, locs_y, nodelabel = nodelabel),
    )

end

end


# Construct DiGraph
flow_graph = DiGraph(8)

# Load custom dataset
flow_edges = [
    (1,2,10),(1,3,5),(1,4,15),(2,3,4),(2,5,9),
    (2,6,15),(3,4,4),(3,6,8),(4,7,16),(5,6,15),
    (5,8,10),(6,7,15),(6,8,10),(7,3,6),(7,8,10)
]

capacity_matrix = zeros(Int ,nv(flow_graph) ,nv(flow_graph))

for e in flow_edges
    u,v,f = e
    add_edge!(flow_graph,u,v)
    capacity_matrix[u,v] = f
end

# Construct the residual graph
residual_graph = LightGraphs.residual(flow_graph)

# Test with default distances
@test LightGraphs.dinic_impl(residual_graph, 1, 8, LightGraphs.DefaultCapacity(residual_graph))[1] == 3

# Test with capacity matrix
@test LightGraphs.dinic_impl(residual_graph, 1, 8, capacity_matrix)[1] == 28

# Test on disconnected graphs
function test_blocking_flow(residual_graph, source, target, capacity_matrix, flow_matrix)
    #disconnect source
    h = copy(residual_graph)
    for dst in collect(neighbors(residual_graph, source))
        rem_edge!(h, source, dst)
    end
    @test LightGraphs.blocking_flow!(h, source, target, capacity_matrix, flow_matrix) == 0

    #disconnect target and add unreachable vertex
    h = copy(residual_graph)
    for src in collect(in_neighbors(residual_graph, target))
        rem_edge!(h, src, target)
    end
    @test LightGraphs.blocking_flow!(h, source, target, capacity_matrix, flow_matrix) == 0

    # unreachable vertex (covers the case where a vertex isn't reachable from the source)
    h = copy(residual_graph)
    add_vertex!(h)
    add_edge!(h, nv(residual_graph)+1, target)
    capacity_matrix_ = vcat(hcat(capacity_matrix, zeros(Int, nv(residual_graph))), zeros(Int, 1, nv(residual_graph)+1))
    flow_graph_  = vcat(hcat(flow_matrix, zeros(Int, nv(residual_graph))), zeros(Int, 1, nv(residual_graph)+1))

    @test LightGraphs.blocking_flow!(h, source, target, capacity_matrix_, flow_graph_ ) > 0

    #test with connected graph
    @test LightGraphs.blocking_flow!(residual_graph, source, target, capacity_matrix, flow_matrix) > 0
end

flow_matrix = zeros(Int, nv(residual_graph), nv(residual_graph))
test_blocking_flow(residual_graph, 1, 8, capacity_matrix, flow_matrix)


# Part 1 - parse passports
passports = strip.(split(read("day04_input.txt", String), "\n\n"))
passports = map(p -> Dict(split.(split(p), ':')), passports)
required = ["byr", "iyr", "eyr", "hgt", "hcl", "ecl", "pid"]
println("Part 1: There are $(sum(map(p -> issubset(required, keys(p)), passports))) valid passports")

# Part 2
function number_valid(p, key, re, mini, maxi)
    m = match(re, p[key])
    if isnothing(m) return false end
    mini <= parse(Int32, first(m.captures)) <= maxi
end

validators = [
    p -> number_valid(p, "byr", r"^(\d{4})$", 1920, 2002),
    p -> number_valid(p, "iyr", r"^(\d{4})$", 2010, 2020),
    p -> number_valid(p, "eyr", r"^(\d{4})$", 2020, 2030),
    p -> number_valid(p, "hgt", r"^(\d{3})cm$", 150, 193) || number_valid(p, "hgt", r"^(\d{2})in$", 59, 76),
    p -> !isnothing(match(r"^#[0-9a-f]{6}$", p["hcl"])),
    p -> p["ecl"] in ["amb", "blu", "brn", "gry", "grn", "hzl", "oth"],
    p -> !isnothing(match(r"^\d{9}$", p["pid"]))
]
is_valid(p) = issubset(required, keys(p)) && all(map(v -> v(p), validators))
println("Part 2: There are $(sum(is_valid.(passports))) valid passports")


immutable TensorMap{S,P,T}<:AbstractTensorMap{S,P,T}
    tensor::AbstractTensor{S,P,T}
    codom::P
    dom::P
    function TensorMap(t::AbstractTensor{S,P,T},codom::P,dom::P=codom)
        t in codom ⊗ dual(dom).' || throw(SpaceError("tensor cannot represent map from $dom to $codom"))
        new(t,codom,dom)
    end
end
tensormap{S,P,T}(t::AbstractTensor{S,P,T},codom::TensorSpace,dom::TensorSpace=codom)=TensorMap{S,P,T}(t,codom,dom)

# properties
codomain(A::TensorMap)=A.codom
domain(A::TensorMap)=A.dom

Base.isreal(A::TensorMap)=isreal(A.tensor)

Base.transpose(A::TensorMap)=tensormap(A.tensor.',dual(A.dom),dual(A.codom))
Base.ctranspose(A::TensorMap)=tensormap(A.tensor',conj(dual(A.dom)),conj(dual(A.codom)))

# comparison
==(A::TensorMap,B::TensorMap)=A.map==B.map

# multiplication with vector
function Base.A_mul_B!{S,P}(y::AbstractTensor{S,P},A::TensorMap{S,P},x::AbstractTensor{S,P})
    x in domain(A) || throw(SpaceError("x in not in domain of map"))
    y in codomain(A) || throw(SpaceError("y in not in codomain of map"))
    N1=length(codomain(A))
    N2=length(domain(A))
    tensorcontract!(1,A.tensor,[-(1:N2),reverse(1:N1)],'N',x,1:N1,'N',0,y,-(1:N2))
    return y
end
function Base.Ac_mul_B!{S,P}(y::AbstractTensor{S,P},A::TensorMap{S,P},x::AbstractTensor{S,P})
    x in dual(conj(codomain(A))) || throw(SpaceError("x in not in domain of map"))
    y in dual(conj(domain(A))) || throw(SpaceError("y in not in domain of map"))
    N1=length(codomain(A))
    N2=length(domain(A))
    tensorcontract!(1,A.tensor,[(1:N2),-reverse(1:N1)],'C',x,1:N2,'N',0,y,-(1:N1))
    return y
end
function Base.At_mul_B!{S,P}(y::AbstractTensor{S,P},A::TensorMap{S,P},x::AbstractTensor{S,P})
    x in dual(codomain(A)) || throw(SpaceError("x in not in domain of map"))
    y in dual(domain(A)) || throw(SpaceError("y in not in domain of map"))
    N1=length(codomain(A))
    N2=length(domain(A))
    tensorcontract!(1,A.tensor,[(1:N2),-reverse(1:N1)],'N',x,1:N2,'N',0,y,-(1:N1))
    return y
end


# An implementation of the linesearch method

"""An implementation of the linesearch method.
Linesearch(model, algo ; ϵa, ϵr, itemax) solves a continuous optimization
problem 'model' with absolute and relative tolerances ϵa and ϵr.
Steps are calculated using the argument 'algo', an optimization method using linesearch
"""
function Linesearch(model, algo ; profil::Bool=false, ϵa::Float64=1e-6, ϵr::Float64=1e-6, itemax::Int=10000)

    @info(loggerLin, @sprintf("Linesearch: resolution of %s using %s", model.meta.name, string(algo)))
    x = model.meta.x0 # initial estimation from the model
    n = model.meta.nvar # size of the problem
    g = grad(model, x) # ∇f(x_0)
    H = hess_op(model, x) # ∇²f(x_0)

    normg0 = norm(g,2) # ‖∇f(x_0)‖
    normg = normg0 # ‖∇f(x_k)‖
    ϵ = ϵa + ϵr * normg # tolerance

    fx0 = obj(model, x) # f(x_0)
    fx = fx0 # f(x_k)

    fvalues = [fx] # improving values of the objective
    k = 0 # number of iterations

    @debug(loggerLin, @sprintf("%4s  %9s  %7s", "Iter", "f", "‖∇f‖"))
    @debug(loggerLin, @sprintf("%4d  %9.2e  %7.1e", k, fx, normg))

    while normg > ϵ && k < itemax # stopping criterion : ‖∇f(x_k)‖ <= ϵ or k >= itemax
        k += 1
        s = algo(H, -g, 0., min(0.1, sqrt(normg))) # descent direction
        X = Armijo(model, x, s, fx, g)
        x = X[1]
        fx = X[2]
        fvalues = push!(fvalues, fx)
        g = grad(model, x)
        normg = norm(g,2)
        @debug(loggerLin, @sprintf("%4d  %8.1e  %7.1e", k, fx, normg))
        H = hess_op(model, x)
    end

    @info(loggerLin, @sprintf("%30s %s %9s %9s %9s %9s %3s %3s %4s %4s","name", "nvar", "f(x*)", "/ f(x0)", "‖∇f(x*)‖", "/ ‖∇f(x0)‖", "#f", "#g", "#Hv", "#it"))
    @info(loggerLin, @sprintf("%30s %d  %8.1e  %8.1e    %7.1e  %7.1e     %d   %d   %d    %d", model.meta.name, n, fx, fx0, normg, normg0, neval_obj(model), neval_grad(model), neval_hprod(model), k))

    optimal = normg ≤ ϵ
    tired = k ≥ itemax
    status = optimal ? "optimal" : "tired"

    if profil
        return x, fx, normg, k, optimal, tired, status # to use two_solvers
    else
        fx = @sprintf("%8.3e", fx)
        fx0 = @sprintf("%8.3e", fx0)
        normg = @sprintf("%7.1e", normg)
        normg0 = @sprintf("%7.1e", normg0)
        return [model.meta.name[22:end] string(algo) n fx fx0 normg normg0 neval_obj(model) neval_grad(model) neval_hprod(model) k optimal] # to use launch_linesearch
    end
end


using FunManifolds
using Test
using StaticArrays

include("../utils.jl")

@testset "DiscretizedCurves" begin
    N = 50
    s2 = Sphere(2)
    r2 = Euclidean(2)

    dcss2 = DiscretizedCurves(s2, range(0.0, 1.0, length=N))
    test_manifold(
        dcss2,
        [
            discretize(dcss2, t -> project(s2, @SVector [t + 0.1, t^2, sin(t)])),
            discretize(dcss2, t -> project(s2, @SVector [t - 1, sin(t^2), 1])),
            discretize(dcss2, t -> project(s2, @SVector [t^2, t, cos(t)])),
        ];
        exp_log_atol_multiplier=10.0,
        is_tangent_atol_multiplier=10,
        test_default_vector_transport=true,
        test_project_point=true,
        test_project_tangent=true,
    )

    dcsr2 = DiscretizedCurves(r2, range(0.0, 1.0, length=N))
    test_manifold(
        dcsr2,
        [
            discretize(dcsr2, t -> project(r2, @SVector [t^2, sin(t)])),
            discretize(dcsr2, t -> project(r2, @SVector [sin(t^2), 1])),
            discretize(dcsr2, t -> project(r2, @SVector [t, cos(t)])),
        ];
        is_tangent_atol_multiplier=10,
        test_default_vector_transport=true,
    )

    @testset "SRVF" begin
        r3 = Euclidean(3)
        dcsr1 = DiscretizedCurves(r3, range(0.0, 1.0, length=9))
        f = t -> (@SVector [sin(t), t^2, 2 * t + 1])
        c = discretize(dcsr1, f)
        q = tsrvf(
            dcsr1,
            c,
            (@SVector [0.0, 0.0, 0.0]),
            FunManifolds.ProjectedDifferenceBackend(nothing),
        )
        c_rev = reverse_srvf(dcsr1, q, f(0.0))
        @test c ≈ c_rev atol = 1e-10
    end
end


export Vec3d, vec3d, norm2, norm, Ray, ray_at_dist

import Base: show

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

"A 3-D vector. Use `vec3d` to create instances of this object"
const Vec3d = StaticVector{3, Float64}

"""
    vec3d(x, y, z)

Create a `Vec3d` object
"""
vec3d(x, y, z) = SVector{3, Float64}(x, y, z)

"""
    norm2(v::Vec3d)

Squared length of the vector. It is faster to compute than the length, since
no squared-root operation is required.

# Example
```jldoctest
julia> norm2(vec3d(1, 0, 0))
1.0
julia> norm2(vec3d(3, 1, 2))
14.0
```
"""
norm2(v::Vec3d) = v[1]^2 + v[2]^2 + v[3]^2

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

"""
    struct Ray

A ray originating from some point and having a direction. The direction
is always normalized. It contains the following fields:
- `origin`
- `dir`

You can compute the point reached by the ray after some time `t` with the
function `ray_at_dist`. This can be short-handed using the ray as a function:
```jldoctest
julia> r = Ray(vec3d(0, 0, 0), vec3d(1, 0, 0));
julia> ray_at_dist(r, 10) ≈ r(10)
true
```
"""
struct Ray
    "Point where the ray originates"
    origin::Vec3d
    "Normalized direction of the ray"
    dir::Vec3d

    Ray(o, d) = new(o, normalize(d))
end

function show(io::IO, r::Ray)
    println(string("Ray([$(r.origin[1]), $(r.origin[2]), $(r.origin[3])] + ",
                   "t * [$(r.dir[1]), $(r.dir[2]), $(r.dir[3])]"))
end

"""
    ray_at_dist(ray::Ray, dist::Number)

Return the point (a `Vec3d` type) reached by the ray after having
travelled a distance `dist`.
"""
ray_at_dist(ray::Ray, dist::Number) = ray.origin + dist * ray.dir

(r::Ray)(dist::Number) = ray_at_dist(r, dist)


# example from http://openfst.org/twiki/bin/view/FST/ConcatDoc
sym=Dict("$(Char(i))"=>i for i=97:122)

A = WFST(sym)
add_arc!(A,1=>2,"<eps>"=>"<eps>")
add_arc!(A,2=>2,"a"=>"p", 2)
add_arc!(A,2=>3,"b"=>"q", 3)
add_arc!(A,3=>3,"c"=>"r", 4)
add_arc!(A,3=>2,"<eps>"=>"<eps>", 5)
initial!(A,1)
final!(A,1)
final!(A,3,5)
#println(A)

B = WFST(sym)
add_arc!(B,1=>1,"a"=>"p",2)
add_arc!(B,1=>2,"b"=>"q",3)
add_arc!(B,2=>2,"c"=>"r",4)
initial!(B,1)
final!(B,2,5)
#println(A)

W,D = typeofweight(A),Int

C = A*B
#println(C)
@test size(C) == (5,10)
@test get_initial(C; single=false) == Dict(1=>one(W))
@test get_final(C; single=false) == Dict(5=>W(5))
@test C[1][1] == Arc{W,D}(0,0,one(W),2)
@test C[1][2] == Arc{W,D}(0,0,one(W),4)
@test C[2][1] == Arc{W,D}(sym["a"],sym["p"],W(2),2)
@test C[2][2] == Arc{W,D}(sym["b"],sym["q"],W(3),3)
@test C[3][1] == Arc{W,D}(sym["c"],sym["r"],W(4),3)
@test C[3][2] == Arc{W,D}(0,0,W(5),2)
@test C[3][3] == Arc{W,D}(0,0,W(5),4)
@test C[4][1] == Arc{W,D}(sym["a"],sym["p"],W(2),4)
@test C[4][2] == Arc{W,D}(sym["b"],sym["q"],W(3),5)
@test C[5][1] == Arc{W,D}(sym["c"],sym["r"],W(4),5)

B = WFST(sym,["z"]); initial!(B,1); final!(B,2)
@test_throws ErrorException A*B
B = WFST(["z"],sym); initial!(B,1); final!(B,2)
@test_throws ErrorException A*B
B = WFST(sym; W = ProbabilityWeight{Float32}); initial!(B,1); final!(B,2)
@test_throws ErrorException A*B
###
### with deterministic acyclic WFSTs
W = ProbabilityWeight{Float64}
isym=Dict("$(Char(i))"=>i for i=97:122)
osym=Dict(Char(i)=>i for i=97:122)
A = WFST(isym,osym; W = W)
add_arc!(A,1=>2,"a"=>'a',1)
add_arc!(A,2=>4,"b"=>'b',2)
add_arc!(A,2=>3,"c"=>'c',3)
add_arc!(A,3=>4,"d"=>'d',4)
add_arc!(A,1=>4,"e"=>'e',5)
initial!(A,1)
final!(A,3,6)
final!(A,4,7)
println(A)

B = WFST(isym,osym; W = W)
add_arc!(B,1=>2,"x"=>'x',8)
add_arc!(B,1=>3,"y"=>'y',9)
initial!(B,1)
final!(B,2,10)
final!(B,3,11)
println(B)

pas = collectpaths(A)
pbs = collectpaths(B)
C = A*B
pcs = collectpaths(C)
@test all([pa*pb in pcs for pa in pas, pb in pbs])

C = B*A
pcs = collectpaths(C)
@test all([pb*pa in pcs for pa in pas, pb in pbs])

C = A*A
pcs = collectpaths(C)
@test all([pa*pa in pcs for pa in pas, pb in pbs])

C = B*B
pcs = collectpaths(C)
@test all([pb*pb in pcs for pa in pas, pb in pbs])


using SafeTestsets, Test

@safetestset "CanonicalSymplecticForms" begin include("test_CanonicalSymplecticForms.jl") end

@safetestset "Plan Unit Tests" begin
    @safetestset "FirstFinDiffPlans" begin include("test_FirstFinDiffPlans.jl") end
    @safetestset "FirstFourierPlans" begin include("test_FirstFourierPlans.jl") end

    @safetestset "SecondChebyshevPlans" begin include("test_SecondChebyshevPlans.jl") end
    @safetestset "SecondFinDiffPlans" begin include("test_SecondFinDiffPlans.jl") end
end

@safetestset "PoincareInvariants" begin
    include("test_PoincareInvariants.jl")
end


using Test
using TensorFactorizations
using TensorOperations
using Random
using ArgParse
using LinearAlgebra

function parse_pars()
    settings = ArgParseSettings(autofix_names=true)
    @add_arg_table(settings
        , "--iters", arg_type=Int, default=50
        , "--no_test_svd", action=:store_true
        , "--no_test_eig", action=:store_true
        , "--no_test_split", action=:store_true
    )
    pars = parse_args(ARGS, settings; as_symbols=true)
    pars[:test_svd] = !pars[:no_test_svd]
    pars[:test_eig] = !pars[:no_test_eig]
    pars[:test_split] = !pars[:no_test_split]
    return pars
end

""" Generate a random shape. """
function rshape(;n=nothing, chi=nothing, nlow=0, nhigh=5, chilow=1, chihigh=5)
    # n is the number of legs, chi is the bond dimension(s)
    n == nothing && (n = rand(nlow:nhigh))
    shape = (chi == nothing ? rand(chilow:chihigh, (n,))
             : collect(repeat([chi], n)))
    return shape
end

""" Generate a random tensor, either with a given or a random shape. """
function rtensor(;shape=nothing, n=nothing, chi=nothing,
                 nlow=0, nhigh=5, chilow=0, chihigh=6, dtype=nothing)
    if dtype == nothing
        dtype = rand([Int, Float64, Complex{Float64}])
    end
    if shape == nothing
        shape = rshape(n=n, chi=chi, nlow=nlow, nhigh=nhigh, chilow=chilow,
                       chihigh=chihigh)
    end
    A = rand(dtype, (shape...),)
    return A
end

function test_svd()
    shp = rshape(nlow=2)
    n = length(shp)
    perm = randperm(n)
    I = rand(1:n-1)
    J = n - I
    ins = perm[1:I]
    outs = perm[I+1:end]
    A = rtensor(shape=shp)

    # SVD and reconstruct without truncation
    U, S, Vt = tensorsvd(A, ins, outs)
    US = tensorcontract(U, [-1*(1:I); 1], Diagonal(S), [1,-I-1])
    Areco = tensorcontract(US, [-1*(1:I); 1], Vt, [1; -1*(I+1:n)])
    Areco = tensorcopy(Areco, collect(1:ndims(Areco)), invperm(perm))
    @test isapprox(A, Areco)

    # SVD and reconstruct with truncation
    eps = 10.0^(-rand(0:3))
    dim_in = prod(shp[perm][1:I])
    dim_out = prod(shp[perm][I+1:end])
    max_dim = min(dim_in, dim_out)
    max_chi = rand(1:max_dim)
    chis = 1:max_chi
    U, S, Vt, error = tensorsvd(A, ins, outs, eps=eps, chis=chis,
                                return_error=true)
    # This test could fail if there were degenerate singular values. We assume
    # this is not the case because the tensors are random.
    @test (error<eps || length(S) == max_chi)
    US = tensorcontract(U, [-1*(1:I); 1], Diagonal(S), [1,-I-1])
    Areco = tensorcontract(US, [-1*(1:I); 1], Vt, [1; -1*(I+1:n)])
    Areco = tensorcopy(Areco, collect(1:ndims(Areco)), invperm(perm))
    Anorm = norm(A)
    true_error = Anorm > 0 ? norm(Areco - A)/Anorm : norm(Areco)
    @test isapprox(error, true_error; atol=1e-8)
end

function test_eig()
    # Non-hermitian
    I = rand(1:3)
    n = 2*I
    shp = rshape(n=n)
    perm = randperm(n)
    ins = perm[1:I]
    outs = perm[I+1:end]
    # Making sure that the left and right legs match.
    shp = collect(shp)
    for (o,i) in zip(outs, ins)
        shp[o] = shp[i]
    end
    shp = (shp...,)
    A = rtensor(shape=shp)

    # Check that eigenvectors really are eigenvectors
    eps = 10.0^(-rand(0:3))
    dim_in = prod(shp[perm][1:I])
    dim_out = prod(shp[perm][I+1:end])
    max_dim = min(dim_in, dim_out)
    max_chi = rand(1:max_dim)
    chis = 1:max_chi
    E, U, error = tensoreig(A, ins, outs, eps=eps, chis=chis,
                            return_error=true, break_degenerate=true)
    @test (error<eps || length(E) == max_chi)
    A_contract_list = collect(repeat([0], n))
    for (k,i) in enumerate(ins)
        A_contract_list[i] = -k
    end
    for (k,i) in enumerate(outs)
        A_contract_list[i] = k
    end
    AU = tensorcontract(A, A_contract_list, U, [1:I; -(n+1)],
                        [-1*(1:I); -(n+1)])
    UE = tensorcontract(U, [-1*(1:I); 1], Diagonal(E), [1; -(n+1)],
                        [-1*(1:I); -(n+1)])
    @test isapprox(AU, UE)

    # Hermitian
    I = rand(1:3)
    n = 2*I
    shp = rshape(n=n)
    perm = randperm(n)
    ins = perm[1:I]
    outs = perm[I+1:end]
    # Making sure that the left and right legs match.
    shp = collect(shp)
    for (o,i) in zip(outs, ins)
        shp[o] = shp[i]
    end
    shp = (shp...,)
    A = rtensor(shape=shp)
    # Make Hermitian
    hperm = collect(1:n)
    for (o,i) in zip(outs, ins)
        hperm[o], hperm[i] = hperm[i], hperm[o]
    end
    At = conj!(tensorcopy(A, collect(1:ndims(A)), hperm))
    A = (A + At)/2

    # Decompose and reconstruct without truncation
    E, U = tensoreig(A, ins, outs, hermitian=true)
    UE = tensorcontract(U, [-1*(1:I); 1], Diagonal(E), [1,-I-1])
    Areco = tensorcontract(UE, [-1*(1:I); 1], conj(U), [-1*(I+1:n); 1])
    Areco = tensorcopy(Areco, collect(1:ndims(Areco)), invperm(perm))

    # Decompose and reconstruct with truncation
    eps = 10.0^(-rand(0:3))
    dim_in = prod(shp[perm][1:I])
    dim_out = prod(shp[perm][I+1:end])
    max_dim = min(dim_in, dim_out)
    max_chi = rand(1:max_dim)
    chis = 1:max_chi
    E, U, error = tensoreig(A, ins, outs, eps=eps, chis=chis,
                            return_error=true, hermitian=true)
    # This test could fail if there were degenerate singular values. We assume
    # this is not the case because the tensors are random.
    @test (error<eps || length(E) == max_chi)
    UE = tensorcontract(U, [-1*(1:I); 1], Diagonal(E), [1,-I-1])
    Areco = tensorcontract(UE, [-1*(1:I); 1], conj(U), [-1*(I+1:n); 1])
    Areco = tensorcopy(Areco, collect(1:ndims(Areco)), invperm(perm))
    Anorm = norm(A)
    true_error = Anorm > 0 ? norm(Areco - A)/Anorm : norm(Areco)
    @test isapprox(error, true_error; atol=1e-8)
end

function test_split()
    # Non-hermitian, SVD
    shp = rshape(nlow=2)
    n = length(shp)
    perm = randperm(n)
    I = rand(1:n-1)
    J = n - I
    ins = perm[1:I]
    outs = perm[I+1:end]
    A = rtensor(shape=shp)

    B1, B2 = tensorsplit(A, ins, outs)
    U, S, Vt = tensorsvd(A, ins, outs)
    Ssqrt = Diagonal(sqrt.(S))
    B1reco = tensorcontract(U, [-1*(1:I); 1], Ssqrt, [1, -n-1])
    B2reco = tensorcontract(Ssqrt, [-1,1], Vt, [1; -1*(2:J+1)])
    @test isapprox(B1reco, B1)
    @test isapprox(B2reco, B2)

    # Hermitian, eigendecomposition
    I = rand(1:3)
    n = 2*I
    shp = rshape(n=n)
    perm = randperm(n)
    ins = perm[1:I]
    outs = perm[I+1:end]
    # Making sure that the left and right legs match.
    shp = collect(shp)
    for (o,i) in zip(outs, ins)
        shp[o] = shp[i]
    end
    shp = (shp...,)
    A = rtensor(shape=shp)
    # Make Hermitian
    hperm = collect(1:n)
    for (o,i) in zip(outs, ins)
        hperm[o], hperm[i] = hperm[i], hperm[o]
    end
    At = conj!(tensorcopy(A, collect(1:ndims(A)), hperm))
    A = (A + At)/2

    B1, B2 = tensorsplit(A, ins, outs, hermitian=true)
    E, U = tensoreig(A, ins, outs, hermitian=true)
    Esqrt = Diagonal(sqrt.(complex.(E)))
    B1reco = tensorcontract(U, [-1*(1:I); 1], Esqrt, [1, -n-1])
    B2reco = tensorcontract(Esqrt, [-1,1], conj(U), [-1*(2:I+1); 1])
    @test isapprox(B1reco, B1)
    @test isapprox(B2reco, B2)
end

function main()
    pars = parse_pars()
    if pars[:test_svd]
        println("Testing SVD.")
        for iter_num in 1:pars[:iters]
            test_svd()
        end
        println("Done testing SVD.")
    end
    if pars[:test_eig]
        println("Testing eig.")
        for iter_num in 1:pars[:iters]
            test_eig()
        end
        println("Done testing eig.")
    end
    if pars[:test_split]
        println("Testing split.")
        for iter_num in 1:pars[:iters]
            test_split()
        end
        println("Done testing split.")
    end
end

main()


# ---
# title: 486. Predict the Winner
# id: problem486
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Dynamic Programming, Minimax
# link: <https://leetcode.com/problems/predict-the-winner/description/>
# hidden: true
# ---
#
# Given an array of scores that are non-negative integers. Player 1 picks one of
# the numbers from either end of the array followed by the player 2 and then
# player 1 and so on. Each time a player picks a number, that number will not be
# available for the next player. This continues until all the scores have been
# chosen. The player with the maximum score wins.
#
# Given an array of scores, predict whether player 1 is the winner. You can
# assume each player plays to maximize his score.
#
# **Example 1:**
#
#
#
#     Input: [1, 5, 2]
#     Output: False
#     Explanation: Initially, player 1 can choose between 1 and 2.
#     If he chooses 2 (or 1), then player 2 can choose from 1 (or 2) and 5. If player 2 chooses 5, then player 1 will be left with 1 (or 2).
#     So, final score of player 1 is 1 + 2 = 3, and player 2 is 5.
#     Hence, player 1 will never be the winner and you need to return False.
#
#
#
#
# **Example 2:**
#
#
#
#     Input: [1, 5, 233, 7]
#     Output: True
#     Explanation: Player 1 first chooses 1. Then player 2 have to choose between 5 and 7. No matter which number player 2 choose, player 1 can choose 233.
#     Finally, player 1 has more score (234) than player 2 (12), so you need to return True representing player1 can win.
#
#
#
#
# **Constraints:**
#
#   * 1 <= length of the array <= 20.
#   * Any scores in the given array are non-negative integers and will not exceed 10,000,000.
#   * If the scores of both players are equal, then player 1 is still the winner.
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


# ---
# title: 1416. Restore The Array
# id: problem1416
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/restore-the-array/description/>
# hidden: true
# ---
#
# A program was supposed to print an array of integers. The program forgot to
# print whitespaces and the array is printed as a string of digits and all we
# know is that all integers in the array were in the range `[1, k]` and there
# are no leading zeros in the array.
#
# Given the string `s` and the integer `k`. There can be multiple ways to
# restore the array.
#
# Return _the number of possible array_ that can be printed as a string `s`
# using the mentioned program.
#
# The number of ways could be very large so return it **modulo** `10^9 + 7`
#
#
#
# **Example 1:**
#
#
#
#     Input: s = "1000", k = 10000
#     Output: 1
#     Explanation: The only possible array is [1000]
#
#
# **Example 2:**
#
#
#
#     Input: s = "1000", k = 10
#     Output: 0
#     Explanation: There cannot be an array that was printed this way and has all integer >= 1 and <= 10.
#
#
# **Example 3:**
#
#
#
#     Input: s = "1317", k = 2000
#     Output: 8
#     Explanation: Possible arrays are [1317],[131,7],[13,17],[1,317],[13,1,7],[1,31,7],[1,3,17],[1,3,1,7]
#
#
# **Example 4:**
#
#
#
#     Input: s = "2020", k = 30
#     Output: 1
#     Explanation: The only possible array is [20,20]. [2020] is invalid because 2020 > 30. [2,020] is ivalid because 020 contains leading zeros.
#
#
# **Example 5:**
#
#
#
#     Input: s = "1234567890", k = 90
#     Output: 34
#
#
#
#
# **Constraints:**
#
#   * `1 <= s.length <= 10^5`.
#   * `s` consists of only digits and doesn't contain leading zeros.
#   * `1 <= k <= 10^9`.
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


"""
    UnimodalRosin <: AbstractImageBinarizationAlgorithm
    UnimodalRosin()

    binarize([T,] img, f::UnimodalRosin)
    binarize!([out,] img, f::UnimodalRosin)

Uses Rosin's Unimodal threshold algorithm to binarize the image.

# Output

Return the binarized image as an `Array{Gray{T}}` of size `size(img)`. If
`T` is not specified, it is inferred from `out` and `img`.

# Details

This algorithm first selects the bin in the image histogram with the highest
frequency. The algorithm then searches from the location of the maximum bin to
the last bin of the histogram for the first bin with a frequency of 0 (known as
the minimum bin.). A line is then drawn that passes through both the maximum and
minimum bins. The bin with the greatest orthogonal distance to the line is
chosen as the threshold value.

## Assumptions

This algorithm assumes that:

* The histogram is unimodal.
* There is always at least one bin that has a frequency of 0. If not, the algorithm will use the last bin as the minimum bin.

If the histogram includes multiple bins with a frequency of 0, the algorithm
will select the first zero bin as its minimum. If there are multiple bins with
the greatest orthogonal distance, the leftmost bin is selected as the threshold.

# Arguments

The function argument is described in more detail below.

##  `img::AbstractArray`

The image that needs to be binarized. The image is automatically converted
to `Gray` in order to construct the requisite graylevel histogram.

# Example

Compute the threshold for the "moonsurface" image in the `TestImages` package.

```julia
using TestImages, ImageBinarization

img = testimage("moonsurface")
img_binary = binarize(img, UnimodalRosin())
```

# Reference
1. P. L. Rosin, “Unimodal thresholding,” Pattern Recognition, vol. 34, no. 11, pp. 2083–2096, Nov. 2001.[doi:10.1016/s0031-3203(00)00136-9](https://doi.org/10.1016/s0031-3203%2800%2900136-9)
"""
struct UnimodalRosin <: AbstractImageBinarizationAlgorithm end

function (f::UnimodalRosin)(out::GenericGrayImage, img::GenericGrayImage)
    edges, counts = build_histogram(img,  256)
    t = find_threshold(HistogramThresholding.UnimodalRosin(), counts[1:end], edges)
    @simd for i in CartesianIndices(img)
      out[i] = img[i] < t ? 0 : 1
    end
    out
end

(f::UnimodalRosin)(out::GenericGrayImage, img::AbstractArray{<:Color3}) =
    f(out, of_eltype(Gray, img))


using Mamba, Stan

currdir = (@__FILE__() == nothing) ? "." : dirname(@__FILE__())
const BayesianSlopeModel = open(fh -> readall(fh), joinpath(currdir, "BayesianSlope.stan"))

# Create fake data to test
NumCoeffs = 100
nonzero_coeffs = [2, 3, -4, 7]
nx = length(nonzero_coeffs)
NumSpurious = NumCoeffs - nx
coeffs = zeros(NumCoeffs)
idxs_nonzero = shuffle(collect(1:NumCoeffs))[1:nx]
coeffs[idxs_nonzero] = nonzero_coeffs
p = length(coeffs)
s = 0.20
n = 60
X = randn(n, p)
y = X * coeffs .+ s * randn(n)

# Set lambdas according to Benjamini-Hochberg
using Distributions
const NormalDistr = Normal(0, 1)
function lambdas_bh(p, q)
    lambdas = zeros(Float64, p)
    for i in 1:p
        lambdas[p+1-i] = quantile(NormalDistr, 1-0.5*i*q/p)
    end
    lambdas
end

# Prep and then call stan
data = Dict("p" => p, "n" => n, "y" => y, "X" => X,
    "lambda" => lambdas_bh(p, 0.01), "a" => 1.0, "g" => 1.0);
bayes_slope_data = [data];

num_betas = length(coeffs);
beta_names = map(i -> "beta.$i", 1:num_betas);
convergence_stats = ["lp__", "accept_stat__"];

stanmodel = Stanmodel(name="bayesian_slope", model=BayesianSlopeModel,
    monitors = vcat(convergence_stats, beta_names), update=2000)
@time sim1 = stan(stanmodel, bayes_slope_data, currdir, CmdStanDir=CMDSTAN_HOME);

convstats_chain1 = sim1[1:size(sim1, 1), convergence_stats, 1];

betas_chain1 = sim1[1:size(sim1, 1), beta_names, 1].value;
est_betas_chain1 = median(betas_chain1, 1)[:];
#std_betas_chain1 = std(betas_chain1, 1)[:];
#includes_zero(val, stdest) = (val - stdest) < 0.0 < (val + stdest)
#nonzero_idxs = filter(i -> !includes_zero(est_betas_chain1[i], std_betas_chain1[i]), 1:length(est_betas_chain1));
includes_zeroes_hdr(vals, level=0.05) = (qs = quantile(vals, [level, 1.0-level]); qs[1] < 0.0 < qs[2])
nonzero_idxs = filter(i -> !includes_zeroes_hdr(betas_chain1[:,i,:][:]), 1:length(est_betas_chain1));
println("Non-zero coeffs:")
map(i -> println("  $i: ", join(map(v -> round(v,2), quantile(betas_chain1[:,i,:][:], [0.05, 0.50, 0.95])), " - ")), sort(nonzero_idxs));
println("Actual coeffs:")
map(i -> println("  $i: ", coeffs[i]), sort(idxs_nonzero));

# Now compare to Optim-based SLOPE regression
include("optim_based_slope.jl")
#betas_optim = optim_based_slope_regression(y, X; absmax = 1e2)
betas_optim = optim_based_slope_regression(y, X);
largest_abs = sortperm(abs(betas_optim), rev=true);
println("Largest 10: ", largest_abs[1:10])
println("Coeffs:     ", map(i -> round(betas_optim[i], 3), largest_abs[1:10]))

# Without extensive testing it seems that Bayesian SLOPE regression is more powerful:
#  - if the number of cases (n) is low and, in particular, if n<p
#  - in that it can give confidence intervals and not only point estimates
# However, it does take a lot longer time so if n > p it might not be the way to go...
# Optim-based SLOPE does not seem to work for n/p under 1.2 or so...


@testset "fdiff" begin
    @testset "API" begin
        # fdiff! works the same as fdiff
        mat_in = rand(3, 3, 3)
        mat_out = similar(mat_in)
        fdiff!(mat_out, mat_in, dims = 2)
        @test mat_out == fdiff(mat_in, dims = 2)

        mat_in = rand(3, 3, 3)
        mat_out = similar(mat_in)
        fdiff!(mat_out, mat_in, dims = 3, rev=true)
        @test mat_out == fdiff(mat_in, dims = 3, rev=true)

        # default `dims` are only available when input is a vector
        v = rand(9)
        @test fdiff(v) == fdiff(v; dims=1)
        m = rand(3, 3)
        @test_throws UndefKeywordError fdiff(m)
    end

    @testset "NumericalTests" begin
        a = reshape(collect(1:9), 3, 3)
        b_fd_1 = [1 1 1; 1 1 1; -2 -2 -2]
        b_fd_2 = [3 3 -6; 3 3 -6; 3 3 -6]
        b_bd_1 = [-2 -2 -2; 1 1 1; 1 1 1]
        b_bd_2 = [-6 3 3; -6 3 3; -6 3 3]
        out = similar(a)

        @test fdiff(a, dims = 1) == b_fd_1
        @test fdiff(a, dims = 2) == b_fd_2
        @test fdiff(a, dims = 1, rev=true) == b_bd_1
        @test fdiff(a, dims = 2, rev=true) == b_bd_2
        fdiff!(out, a, dims = 1)
        @test out == b_fd_1
        fdiff!(out, a, dims = 2)
        @test out == b_fd_2
        fdiff!(out, a, dims = 1, rev=true)
        @test out == b_bd_1
        fdiff!(out, a, dims = 2, rev=true)
        @test out == b_bd_2

        @testset "Boundary" begin
            # These originally lives in Images and we use it to check the compatibility
            forwarddiffy(u::AbstractMatrix) = [u[2:end,:]; u[end:end,:]] - u
            forwarddiffx(u::AbstractMatrix) = [u[:,2:end] u[:,end:end]] - u
            backdiffy(u::AbstractMatrix) = u - [u[1:1,:]; u[1:end-1,:]]
            backdiffx(u::AbstractMatrix) = u - [u[:,1:1] u[:,1:end-1]]

            X = rand(3, 3)
            @test forwarddiffx(X) == fdiff(X, dims=2, boundary=:zero)
            @test forwarddiffy(X) == fdiff(X, dims=1, boundary=:zero)
            @test backdiffx(X) == fdiff(X, dims=2, rev=true, boundary=:zero)
            @test backdiffy(X) == fdiff(X, dims=1, rev=true, boundary=:zero)
        end

        if VERSION >= v"1.1"
            # dims keyword for diff only exists after Julia 1.1
            # check numerical results with base implementation
            drop_last_slice(X, dims) = collect(StackView(collect(eachslice(X, dims=dims))[1:end-1]..., dims=dims))
            function drop_last_slice(X::AbstractVector, dims)
                @assert dims==1
                X[1:end-1]
            end
            drop_first_slice(X, dims) = collect(StackView(collect(eachslice(X, dims=dims))[2:end]..., dims=dims))
            function drop_first_slice(X::AbstractVector, dims)
                @assert dims==1
                X[2:end]
            end

            for N in 1:3
                sz = ntuple(_->5, N)
                A = rand(sz...)
                A_out = similar(A)

                for dims = 1:N
                    out_base = diff(A; dims=dims)
                    out = fdiff(A; dims=dims)
                    @test out_base == drop_last_slice(out, dims)

                    out_base = .-reverse(diff(reverse(A; dims=dims); dims=dims); dims=dims)
                    out = fdiff(A; dims=dims, rev=true)
                    @test out_base == drop_first_slice(out, dims)
                end
            end
        end
    end

    @testset "OffsetArrays" begin
        A = OffsetArray(rand(3, 3), -1, -1)
        A_out = fdiff(A, dims=1)
        @test axes(A_out) == (0:2, 0:2)
        @test A_out.parent == fdiff(parent(A), dims=1)

        A = OffsetArray(rand(3, 3), -1, -1)
        A_out = fdiff(A, dims=1, rev=true)
        @test axes(A_out) == (0:2, 0:2)
        @test A_out.parent == fdiff(parent(A), dims=1, rev=true)
    end
end


include("../../../../math/common/sparse/src/sparse_module.jl")

module mtx
using ..sparse

# mtrx is a file format for storing sparse matrices
# coming from https://sparse.tamu.edu/?per_page=100
#
# there is a heading which seems to be 13 lines
#
# the first non header row contains the number of cols, rows and nnz
#
# extracted files consist of 3 cols
# first col = column index in 1 based indexing
# second col = row index in 1 based indexing
# third col = data
#
# there is a space between the cols

struct MtxBody
    n_row::Int
    n_col::Int
    nnz::Int
    row_index::Vector{Int}
    col_index::Vector{Int}
    data::Vector{Float64}
end

struct MtxFileContent{T}
    header::Vector{T}
    body::MtxBody
end

function read_mtx(url::String)
    iostr = open(url, "r");
    iobuff = read(iostr, String);
    close(iostr);

    mtx_content = parse_mtx(iobuff)
    return mtx_content
end

function find_first_nonheader_index(lines)
    Ll = length(lines)
    kk = 0
    for outer kk = 1:Ll
        line = strip( lines[kk] )
        if line[1] == '%'
            continue
        else
            break
        end
    end
    first_non_header_index = kk
    return first_non_header_index
end

function extract_header(lines, first_non_header_index)
    header = lines[1:first_non_header_index-1]
    return header
end

function interpret_body(lines, first_non_header_index)

    L = length(lines)-1

    # handle first row separately
    line    = lines[first_non_header_index]
    strvec  = split(line, " ")
    n_row   = parse(Int, strvec[1])
    n_col   = parse(Int, strvec[2])
    nnz     = parse(Int, strvec[3])

    row_index = zeros(Int, nnz)
    col_index = zeros(Int, nnz)
    data      = zeros(Float64, nnz)

    ii = 0
    for kk = (first_non_header_index+1):L
        ii += 1

        strvec = split(lines[kk], " ")
        rr     = parse(Int, strvec[1])
        cc     = parse(Int, strvec[2])
        dd     = parse(Float64, strvec[3])

        row_index[ii]   = rr
        col_index[ii]   = cc
        data[ii]        = dd
    end

    body = MtxBody(n_row, n_col, nnz, row_index, col_index, data)
    return body
end

function parse_mtx(iobuff)

    mtx_content = iobuff
    lines       = split(iobuff, "\n")

    first_non_header_index = find_first_nonheader_index(lines)
    header = extract_header(lines, first_non_header_index)
    body   = interpret_body(lines, first_non_header_index)

    mtx_content = MtxFileContent(header, body)
    return mtx_content
end

function mtx_to_sparse_csc(mtx_content::MtxFileContent)

    return sparse.SparseMatrixCSC(
        mtx_content.body.n_row,
        mtx_content.body.n_col,
        mtx_content.body.data,
        mtx_content.body.row_index,
        mtx_content.body.col_index)

end

end # module


using BenchmarkTools
using LarSurf

# function grid_voxel_id_to_cart(grid_size::Array{Int, 1}, ind::Int64)
function grid_voxel_id_to_cart_calcul(grid_size, ind)
#     ind = (sz2 * sz3) * (i - 1)  + (j - 1) * sz3 + k
    sz1,sz2,sz3 = grid_size
    layer = sz2*sz3
    ir = div(ind - 1, layer) + 1
    rest1 = mod(ind - 1 , layer) + 1
#     println("rest1 ", rest1)
    row = sz3
    jr = div(rest1 -1, row) + 1
    rest2 = mod(rest1 - 1, row) + 1
    kr = rest2
#     println("rests: ", rest1, " ", rest2)
    return [kr, jr, ir]
end

function grid_voxel_id_to_cart_cartind(blocks_number_axis, block_i)
    a = Array{Int}(
        undef,
        blocks_number_axis[1],
        blocks_number_axis[2],
        blocks_number_axis[3]
    )
#     print("block_i: ", block_i)
    # TODO maybe the cartesian index is slow
    bsub = CartesianIndices(a)[block_i]
    return bsub
end

blocks_number_axis = [10,10,10]
_cart = Array{Int}(
    undef,
    blocks_number_axis[1],
    blocks_number_axis[2],
    blocks_number_axis[3]
)

function grid_voxel_id_to_cart_extern(blocks_number_axis, block_i)
#     print("block_i: ", block_i)
    # TODO maybe the cartesian index is slow
    bsub = CartesianIndices(_cart)[block_i]
    return bsub
end

bl1 = grid_voxel_id_to_cart_calcul([10,10,10], 125)
println(bl1)
bl2 = grid_voxel_id_to_cart_cartind([10,10,10], 125)
println(bl2)
# bl3 = LarSurf.grid_voxel_id_to_cart([10,10,10], 125)
bl3 = grid_voxel_id_to_cart_extern([10,10,10], 125)
println(bl3)

bm = @benchmark grid_voxel_id_to_cart_calcul([10,10,10], 472)
display(bm)

bm = @benchmark grid_voxel_id_to_cart_cartind([10,10,10], 472)
display(bm)

bm = @benchmark grid_voxel_id_to_cart_extern([10,10,10], 472)
display(bm)

# bm = @benchmark LarSurf.grid_voxel_id_to_cart([10,10,10], 472)
# display(bm)


"""
Round to the nearest integer.
If the number is halfway between two integers,
then round down. (See definition 2.8 (pg 27) of Bremner.)

Motivation: The built-in round function sends -0.5 to 0 instead of -1.
This function correctly handles negative numbers per the examples in Bremner.
"""
function round_ties_down{T<:AbstractFloat}(x::T)
    if abs(floor(x) - x) == 0.5
        round(Int, x, RoundDown)
    else
        round(Int, x)
    end
end

round_ties_down{T<:Integer}(x::T) = x


# Set up the environment to run this example. Make sure you're within the folder that this
# file lives in.
using Pkg
Pkg.activate(@__DIR__)
Pkg.instantiate()

using Stheno, Plots, Random
gr()

# Define model.
σ², ω, T = 1e0, 1.0, 25.0
gpc = GPC()
f₁ = periodic(GP(EQ(), gpc), ω)
f₂ = GP(0.1 * EQ(), gpc)
f₃ = f₁ + f₂

# Sample from marginal process to generate toy data.
rng = MersenneTwister(123456)
S = 25
x = range(0.0, T; length=100)
xp = range(-2.5, T + 2.5; length=500)
fx = f₃(x, σ²)
y = rand(rng, fx)

# Plots for stuff.
f₁′_plot, f₂′_plot, f₃′_plot = plot(), plot(), plot();

posterior_plot = plot();
approx_in_marginal_posterior_plot = plot();
approx_in_latents_posterior_plot = plot();

##################################### Exact Inference ######################################

# Compute the posterior processes, sample from them, and compute marginals.
@show logpdf(fx, y)
f₁′, f₂′, f₃′ = (f₁, f₂, f₃) | (fx ← y)
f₁′xp, f₂′xp, f₃′xp = rand(rng, [f₁′(xp, 1e-6), f₂′(xp, 1e-6), f₃′(xp, 1e-6)], S)

ms₁ = marginals(f₁′(xp))
ms₂ = marginals(f₂′(xp))
ms₃ = marginals(f₃′(xp))
μ₁′, σ₁′ = mean.(ms₁), std.(ms₁)
μ₂′, σ₂′ = mean.(ms₂), std.(ms₂)
μ₃′, σ₃′ = mean.(ms₃), std.(ms₃)

items = [
    (μ₁′, σ₁′, f₁′xp, :red, "exact", f₁′_plot),
    (μ₂′, σ₂′, f₂′xp, :red, "", f₂′_plot),
    (μ₃′, σ₃′, f₃′xp, :red, "", f₃′_plot),
]

# Posterior marginal variance.
for (μ, σ, _, colour, name, plt) in items
    plot!(plt, xp, [μ, μ];
        linewidth=0.0,
        fillrange=[μ .- 3 .* σ, μ .+ 3 * σ],
        fillalpha=0.3,
        fillcolor=colour,
        label="",
    );
end

# Posterior samples.
for (μ, σ, f, colour, name, plt) in items
    plot!(plt, xp, f,
        linecolor=colour,
        linealpha=0.2,
        label="",
    );
end

# Posterior mean.
for (μ, σ, f, colour, name, plt) in items
    plot!(plt, xp, μ;
        linecolor=colour,
        linewidth=2.0,
        label=name,
    );
end

# Plot observations.
scatter!(posterior_plot, x, y;
    markercolor=:blue,
    markershape=:circle,
    markerstrokewidth=0.0,
    markersize=4,
    markeralpha=0.7,
    label="",
);

##################### Approximate inference with pseudo-data in f₃ #########################

# Compute approximate posterior process.
M₃ = 15;
Z₃ = range(0.0, T; length=M₃);
Z₃ = 0:T
Z₃ = 0:0.5:T
# Z₃ = 0:0.25:
u₃ = f₃(Z₃);
# μ′u₃, Σ′u₃ = Stheno.optimal_q(f₃(X), ŷ, u₃, σ_noise);
# q_u₃ = Stheno.Titsias(u₃, μ′u₃, Σ′u₃, gpc);
pseudo_obs = Stheno.PseudoObs(f₃(x) ← y, u₃)
f₁′u₃ = f₁ | pseudo_obs
f₂′u₃ = f₂ | pseudo_obs
f₃′u₃ = f₃ | pseudo_obs
f₁′u₃xp, f₂′u₃xp, f₃′u₃xp = rand(rng, [f₁′u₃(xp), f₂′u₃(xp), f₃′u₃(xp)], S);

μ₁′u₃, σ₁′u₃ = mean.(marginals(f₁′u₃(Xp))), std.(marginals(f₁′u₃(Xp)))
μ₂′u₃, σ₂′u₃ = mean.(marginals(f₂′u₃(Xp))), std.(marginals(f₂′u₃(Xp)))
μ₃′u₃, σ₃′u₃ = mean.(marginals(f₃′u₃(Xp))), std.(marginals(f₃′u₃(Xp)))

@show elbo(fx, y, u₃);

items = [
    (μ₁′u₃, σ₁′u₃, f₁′u₃xp, :green, "Z in f3", f₁′_plot),
    (μ₂′u₃, σ₂′u₃, f₂′u₃xp, :green, "", f₂′_plot),
    (μ₃′u₃, σ₃′u₃, f₃′u₃xp, :green, "", f₃′_plot),
];

# Posterior marginal variance.
for (μ, σ, _, colour, name, plt) in items
    plot!(plt, Xp, [μ, μ];
        linewidth=0.0,
        fillrange=[μ .- 3 .* σ, μ .+ 3 * σ],
        fillalpha=0.3,
        fillcolor=colour,
        label="");
end

# # Posterior samples.
# for (μ, σ, f, colour, name, plt) in items
#     plot!(plt, Xp, f,
#         linecolor=colour,
#         linealpha=0.2,
#         label="");
# end

# Posterior mean.
for (μ, σ, f, colour, name, plt) in items
    plot!(plt, Xp, μ;
        linecolor=colour,
        linewidth=2.0,
        label=name,
    );
end

# Plot observations and pseudo-input locations.
scatter!(approx_in_marginal_posterior_plot, X, ŷ;
    markercolor=:blue,
    markershape=:circle,
    markerstrokewidth=0.0,
    markersize=4,
    markeralpha=0.7,
    label="",
);
scatter!(approx_in_marginal_posterior_plot, Z₃, zeros(M₃);
    markercolor=:black,
    markershape=:circle,
    markerstrokewidth=0.0,
    markersize=4,
    markeralpha=0.8,
    label="Z",
);

############# Perform approximate inference by placing pseudo-data in f₁ and f₂ ############

# Compute approximate posterior process.
M₁, M₂ = 10, 5
Z₁, Z₂ = linspace(0.0, 1 / ω, M₁), linspace(0.0, T, M₂)
u₁₂ = BlockGP([f₁(Z₁), f₂(Z₂)])
μ′u, Σ′u = Stheno.optimal_q(f₃(X), ŷ, u₁₂, σ_noise)
conditioner = Stheno.Titsias(u₁₂, μ′u, Σ′u, gpc)
f₁′u₁₂, f₂′u₁₂, f₃′u₁₂ = (f₁, f₂, f₃) | conditioner

f₁′u₁₂Xp, f₂′u₁₂Xp, f₃′u₁₂Xp = rand(rng, [f₁′u₁₂(Xp), f₂′u₁₂(Xp), f₃′u₁₂(Xp)], S)
μ₁′u₁₂, σ₁′u₁₂ = marginals(f₁′u₁₂(Xp))
μ₂′u₁₂, σ₂′u₁₂ = marginals(f₂′u₁₂(Xp))
μ₃′u₁₂, σ₃′u₁₂ = marginals(f₃′u₁₂(Xp))

@show elbo(f₃(X), ŷ, u₁₂, σ_noise)

items = [
    (μ₁′u₁₂, σ₁′u₁₂, f₁′u₁₂Xp, :blue, "Z in f1 and f2", f₁′_plot),
    (μ₂′u₁₂, σ₂′u₁₂, f₂′u₁₂Xp, :blue, "", f₂′_plot),
    (μ₃′u₁₂, σ₃′u₁₂, f₃′u₁₂Xp, :blue, "", f₃′_plot),
]

# Posterior marginal variance.
for (μ, σ, _, colour, name, plt) in items
    plot!(plt, Xp, [μ, μ];
        linewidth=0.0,
        fillrange=[μ .- 3 .* σ, μ .+ 3 * σ],
        fillalpha=0.3,
        fillcolor=colour,
        label="",
    )
end

# # Posterior samples.
# for (μ, σ, f, colour, name, plt) in items
#     plot!(plt, Xp, f,
#         linecolor=colour,
#         linealpha=0.2,
#         label="");
# end

# Posterior mean.
for (μ, σ, f, colour, name, plt) in items
    plot!(plt, Xp, μ;
        linecolor=colour,
        linewidth=2.0,
        label=name,
    )
end

# Plot observations and pseudo-input locations.
scatter!(approx_in_latents_posterior_plot, X, ŷ;
    markercolor=:blue,
    markershape=:circle,
    markerstrokewidth=0.0,
    markersize=4,
    markeralpha=0.7,
    label="",
);
scatter!(approx_in_latents_posterior_plot, Z₁, zeros(M₁);
    markercolor=:black,
    markershape=:circle,
    markerstrokewidth=0.0,
    markersize=4,
    markeralpha=0.8,
    label="Z₁",
);
scatter!(approx_in_latents_posterior_plot, Z₂, zeros(M₂);
    markercolor=:magenta,
    markershape=:circle,
    markerstrokewidth=0.0,
    markersize=4,
    markeralpha=0.8,
    label="Z₂",
);

# display(posterior_plot);
# display(approx_in_marginal_posterior_plot);
# display(approx_in_latents_posterior_plot);

plot(f₁′_plot, f₂′_plot, f₃′_plot; layout=(3, 1))


using DiffEqDevTools
using Base.Test

# write your own tests here
@time @testset "Benchmark Tests" begin include("benchmark_tests.jl") end
@time @testset "ODE AppxTrue Tests" begin include("ode_appxtrue_tests.jl") end
!is_windows() && @time @testset "ODE Tableau Convergence Tests" begin include("ode_tableau_convergence_tests.jl") end ## Windows 32-bit fails on Butcher62 convergence test


import Pkg; Pkg.add(Pkg.PackageSpec(url="https://github.com/JuliaComputing/JuliaAcademyData.jl"))
using JuliaAcademyData; activate("Deep learning with Flux")

#

# # Recognizing handwritten digits using a neural network

#-

# We have now reached the point where we can tackle a very interesting task: applying the knowledge we have gained with machine learning in general, and `Flux.jl` in particular, to create a neural network that can recognize handwritten digits! The data are from a data set called MNIST, which has become a classic in the machine learning world.
#
# [We could also try to apply the techniques to the original images of fruit instead. However, the fruit images are much larger than the MNIST images, which makes the learning a suitable neural network too slow.]

#-

# ## Data munging

#-

# As we know, the first difficulty with any new data set is locating it, understanding what format it is stored in, reading it in and decoding it into a useful data structure in Julia.
#
# The original MNIST data is available [here](http://yann.lecun.com/exdb/mnist); see also the [Wikipedia page](https://en.wikipedia.org/wiki/MNIST_database). However, the format that the data is stored in is rather obscure.
#
# Fortunately, various packages in Julia provide nicer interfaces to access it. We will use the one provided by `Flux.jl`.
#
# The data are images of handwritten digits, and the corresponding labels that were determined by hand (i.e. by humans). Our job will be to get the computer to **learn** to recognize digits by learning, as usual, the function that relates the input and output data.

#-

# ### Loading and examining the data

#-

# First we load the required packages:

using Flux, Flux.Data.MNIST, Images

# Now we read in the data:

labels = MNIST.labels();
images = MNIST.images();

# We see that `images` is a `Vector`, i.e. an `Array{T, 1}` with a complicated parameter `T`. It has length

length(labels)

# But we can just look at the first handful to get a sense of the contents

images[1:5]

#-

labels[1:5]' # transposed to match the above

# So the $i$th entry of the array is the data for the $i$th image.

typeof(images[1])

# As with the fruit images from the start of the course, the image is an array of color blocks, except that now each pixel just has a grey scale.
#
# To see the actual numeric content of the image, we can do, for example

Float64.(images[1])

# ### Setting up a neural network

#-

# In the previous notebooks, we arranged the input data for Flux as a `Vector` of `Vector`s.
# Now we will use an alternative arrangement, as a matrix, since that allows `Flux` to use matrix operations, which are more efficient.
#
# The column $i$ of the matrix is a vector consisting of the $i$th data point $\mathbf{x}^{(i)}$.  Similarly, the desired outputs are given as a matrix, with the $i$th column being the desired output $\mathbf{y}^{(i)}$.

n_inputs = unique(length.(images))[]

#-

n_outputs = length(unique(labels))

# #### Creating the features
#
# We want to create a vector of feature vectors, each with the floating point values of the 784 pixels.
#
# An image is a matrix of colours, but now we need a vector of floating point numbers instead. To do so, we just arrange all of the elements of the matrix in a certain way into a single list; fortunately, Julia already provides the function `vec` to do so!

#-

# Let's use a subset of $N=5,000$ of the total $60,000$ images available in order to hold out test images that our model hasn't been trained on.

preprocess(img) = vec(Float64.(img))

#-

xs = preprocess.(images[1:5000]);

# #### Creating the labels
#
# We can just use `Flux.onehot` to create them:

ys = [Flux.onehot(label, 0:9) for label in labels[1:5000]];

# #### Create the batched matrices for efficiency

#-

# Let's also create a function so we can easily create independent batches from arbitrary segments of the original dataset.

function create_batch(r)
    xs = [preprocess(img) for img in images[r]]
    ys = [Flux.onehot(label, 0:9) for label in labels[r]]
    return (Flux.batch(xs), Flux.batch(ys))
end

# We'll train our model on the first 5000 images.

trainbatch = create_batch(1:5000);

# ## Setting up the neural network

#-

# Now we must set up a neural network. Since the data is complicated, we may expect to need several layers.
# But we can start with a single layer.
#
# - The network will take as inputs the vectors $\mathbf{x}^{(i)}$, so the input layer has $n$ nodes.
#
# - The output will be a one-hot vector encoding the digit from 1 to 9 or 0 that is desired. There are 10 possible categories, so we need an output layer of size 10.
#
# It is then our task as neural network designers to insert layers between these input and output layers, whose weights will be tuned during the learning process. *This is an art, not a science*! But major advances have come from finding a good structure for the network.

model = Chain(Dense(n_inputs, n_outputs, identity), softmax)
L(x,y) = Flux.crossentropy(model(x), y)
opt = SGD(params(model))
@time Flux.train!(L, [trainbatch], opt)
@time Flux.train!(L, [trainbatch], opt)

# ## Training

#-

# Just as before, use `repeated` to create an **iterator**. It does not copy the data 100 times, which would be very wasteful; it just gives an object that repeatedly loops over the same data:

Iterators.repeated(trainbatch, 100);

# We can see what the total current loss is (after training just a handful to times above):

L(trainbatch...)

#-

@time Flux.train!(L, [trainbatch], opt)

#-

L(trainbatch...)

# ### Using callbacks

#-

# The `train!` function can take an optional keyword argument, `cb` (short for "**c**all**b**ack"). A callback function is a function that you provide as an argument to a function `f`, which "calls back" your function every so often.
#
# This provides the possibility to provide a function that is called at each step or every so often during the training process.
# A common use case is to provide a visual trace of the training process by printing out the current value of the `loss` function:

callback() = @show(L(trainbatch...))

Flux.train!(L, Iterators.repeated(trainbatch, 3), opt; cb = callback)

# However, it is expensive to calculate the complete `loss` function and it is not necessary to output it every step. So `Flux` also provides a function `throttle`, that provides a mechanism to call a given function at most once every certain number of seconds:

Flux.train!(L, Iterators.repeated(trainbatch, 40), opt; cb = Flux.throttle(callback, 1))

# Of course, that's just measuring the loss over the same data it's training on. It'd be more representative to test against novel data. In fact, let's track the performance of both as we continue training our model. In order to do so, we need to create a batch of test data.

testbatch = create_batch(5001:10000);

#-

using Printf
train_loss = Float64[]
test_loss = Float64[]
function update_loss!()
    push!(train_loss, L(trainbatch...).data)
    push!(test_loss, L(testbatch...).data)
    @printf("train loss = %.2f, test loss = %.2f\n", train_loss[end], test_loss[end])
end

#-

Flux.train!(L, Iterators.repeated(trainbatch, 1000), opt;
                cb = Flux.throttle(update_loss!, 1))

#-

using Plots
plot(1:length(train_loss), train_loss, xlabel="~seconds of training", ylabel="loss", label="train")
plot!(1:length(test_loss), test_loss, label="test")

# ## Testing phase

#-

# We now have trained a model, i.e. we have found the parameters `W` and `b` for the network layer(s). In order to **test** if the learning procedure was really successful, we check how well the resulting trained network performs when we test it with images that the network has not yet seen!
#
# Often, a dataset is split up into "training data" and "test (or validation) data" for this purpose, and indeed the MNIST data set has a separate pool of training data. We can instead use the images that we have not included in our reduced training process.

i = 5001
display(images[i])
labels[i], findmax(model(preprocess(images[i]))) .- (0, 1)

#-

model(preprocess(images[i]))

# ## Evaluation
#
# What percent of images are we correctly classifying if we take the highest element to be the chosen answer?

prediction(i) = findmax(model(preprocess(images[i])))[2]-1 # returns (max_value, index)

#-

sum(prediction(i) == labels[i] for i in 1:5000)/5000

#-

sum(prediction(i) == labels[i] for i in 5001:10000)/5000

# ## Improving the prediction

#-

# So far we have used a single layer. In order to improve the prediction, we probably need to use more layers. Try adding more layers yourself and see how the performance changes.

n_hidden = 20
model = Chain(Dense(n_inputs, n_hidden, relu),
              Dense(n_hidden, n_outputs, identity), softmax)
L(x,y) = Flux.crossentropy(model(x), y)
opt = ADAM(params(model))

#-

train_loss = Float64[]
test_loss = Float64[]
Flux.train!(L, Iterators.repeated(trainbatch, 1000), opt;
            cb = Flux.throttle(update_loss!, 1))

#-

plot(1:length(train_loss), train_loss, xlabel="~seconds of training", ylabel="loss", label="train")
plot!(1:length(test_loss), test_loss, label="test")

# ## What about image structure?

#-

# As a final note, notice that our model doesn't take into account any aspect of the image's connected-ness.

using Random
p = randperm(28)
images[1][p,p]

#-

show(preprocess(images[1]))

#-

#nb ?Conv
#jl @doc Conv

Courses/Deep learning with Flux/1900.Recognizing-handwriting-using-a-neural-network.jl


module TestRuns

using Test
using RLEVectors

@testset begin

    # numruns single
    @test numruns([1, 2, 2, 3, 3]) == 3
    @test numruns([1, 1, 2, 2, 3]) == 3

    # numruns pair
    @test numruns(["a", "b", "c", "d", "d", "e", "f", "g"], [3, 6, 9, 12, 15, 18, 21, 24]) == 7 # Run of 2 runvalues
    @test numruns(["a", "b", "c", "d", "d", "d", "e", "f"], [3, 6, 9, 12, 15, 18, 21, 24]) == 6 # Run of 3 runvalues
    @test numruns(["a", "b", "c", "d", "e", "f", "g"], [3, 6, 9, 12, 12, 15, 17]) == 6 # run of 2 runends
    @test numruns(["a", "b", "c", "d", "e", "f"], [3, 6, 12, 12, 12, 15]) == 4 # run of 3 runends
    @test numruns([1, 2, 3], [1, 2, 2]) == 2
    @test_throws ArgumentError numruns([1, 2, 3], [3, 2, 4])

    # one vector
    @test ree([1, 1, 2, 2, 3, 3, 4, 4, 4]) == ([1, 2, 3, 4], [2, 4, 6, 9])
    @test ree([1, 1, 2, 2, 3, 3, 4, 4, 7]) == ([1, 2, 3, 4, 7], [2, 4, 6, 8, 9])
    @test ree([0, 1, 2, 2, 3, 3, 4, 4, 7]) == ([0, 1, 2, 3, 4, 7], [1, 2, 4, 6, 8, 9])

    # Clean up an RLEVector
    @test ree(["a", "b", "c", "d", "d", "e", "f", "g"], [3, 6, 9, 12, 15, 18, 21, 24]) ==
          (["a", "b", "c", "d", "e", "f", "g"], [3, 6, 9, 15, 18, 21, 24]) # Run of 2 runvalues
    @test ree(["a", "b", "c", "d", "d", "d", "e", "f"], [3, 6, 9, 12, 15, 18, 21, 24]) ==
          (["a", "b", "c", "d", "e", "f"], [3, 6, 9, 18, 21, 24]) # Run of 3 runvalues
    @test ree(["a", "b", "c", "d", "e", "f", "g"], [3, 6, 9, 12, 12, 15, 17]) ==
          (["a", "b", "c", "d", "f", "g"], [3, 6, 9, 12, 15, 17]) # run of 2 runends
    @test ree(["a", "b", "c", "d", "e", "f"], [3, 6, 12, 12, 12, 15]) ==
          (["a", "b", "c", "f"], [3, 6, 12, 15]) # run of 3 runends
    @test ree([1, 2, 3], [1, 2, 2]) == ([1, 2], [1, 2])

    # two vectors
    @test ree([9], [1]) == ([9], [1])
    @test ree([4, 3], [2, 4]) == ([4, 3], [2, 4])
    @test ree([4, 3], [0, 4]) == ([3], [4])

    # scalar
    @test ree(5) == ([5], [1])
    @test ree([5]) == ([5], [1])

    # inverse_ree
    @test inverse_ree([4, 5, 6], [2, 5, 7]) == [4, 4, 5, 5, 5, 6, 6]
    @test inverse_ree(Real[], Int64[]) == Real[]
    @test_throws ArgumentError inverse_ree([], [1])

end # testset

end # module


# Load modules
using Optim

x = -10:0.01:10
len_x = length(x)

options = Optim.Options(g_abstol=1e-8, g_reltol=1e-8, iterations=10^5, show_trace=false, store_trace=false);

# Gaussian 1 peak
seed1 = [[0.0], [15.0, 0.0, sqrt(0.2)]]
# We generate some data using the Gaussian model
y_data = Utils.gaussian(x, seed1) .+ randn(len_x)
sol = fit_curve_model(:gaussian, x, y_data, seed1; options=options)
# Gaussian 2 peaks
seed2 = [[1.0], [15.0, -1.0, sqrt(0.2)], [10.0, 1.5, sqrt(0.1)]]
# We generate some data using the Gaussian model
y_data = Utils.gaussian(x, seed2) .+ randn(len_x)
sol = fit_curve_model(:gaussian, x, y_data, seed2; options=options)

# Lorentzian 1 peak
seed1 = [[0.0], [15.0, 0.0, 0.5]]
# We generate some data using the Lorentzian model
y_data = Utils.lorentzian(x, seed1) .+ randn(len_x)
sol = fit_curve_model(:lorentzian, x, y_data, seed1; options=options)
# Lorentzian 2 peaks
seed2 = [[0.0], [15.0, -1.0, 0.5], [30.0, 1.0, 0.5]]
# We generate some data using the Lorentzian model
y_data = Utils.lorentzian(x, seed2) .+ randn(len_x)
sol = fit_curve_model(:lorentzian, x, y_data, seed2; options=options)

# Voigtian 1 peak
seed1 = [[0.0], [18.0, 0.0, 0.0, 1.53]]
# We generate some data using the Voigtian model
y_data = Utils.voigtian(x, seed1) .+ randn(len_x)
sol = fit_curve_model(:voigtian, x, y_data, seed1; options=options)
# Voigtian 2 peaks
seed2 = [[0.0], [10.0, -1.6, 0.0, 1.53], [24.0, 5.0, 0.0, 1.53]]
# We generate some data using the Voigtian model
y_data = Utils.voigtian(x, seed2) .+ randn(len_x)
sol = fit_curve_model(:voigtian, x, y_data, seed2; options=options)


#!/usr/bin/env julia

export pe34

function pe34()
    intlength(n::Integer)=length(digits(n))
    fdsum(n::Integer)=sum(factorial.(digits(n)))
    max_digit=2
    while intlength(fdsum(10^max_digit-1))>max_digit
        max_digit +=1
    end
    filter(n->fdsum(n)==n,10:(10^max_digit-1))|>sum
end

if !haskey(ENV,"PROJECT_EULER_WITHOUT_RESULT")
    pe34()|>print
end


module EarthMoversDistance

import JuMP
import COSMO, MathOptInterface
import FluPredictibility.BioTools
using BioSequences: BioSequence

export EMD

"""
	EMD(f,P,Q)

Compute the earth mover's distance between distributions `P` and `Q`, using `f` as a metric.
Enumerating `P` (resp. `Q`) should yield pairs `(x,p)` (resp. `(y,q)`) representing state `x` and its weight `p`. Distance between elements is `f(x,y)`.
"""
function EMD(f::Function,P,Q)
	# Building the distance matrix
	d = Array{Float64, 2}(undef, length(P), length(Q))
	for (i,(x,p)) in enumerate(P)
		for (j,(y,q)) in enumerate(Q)
			d[i,j] = f(x,y)
		end
	end
	# Building the JuMP model
	model = JuMP.Model(COSMO.Optimizer)
	JuMP.set_silent(model)
	JuMP.@variable(model, w[i=1:length(P), j=1:length(Q)] >= 0)
	for (i,(x,p)) in enumerate(P)
		JuMP.@constraint(model, sum(w[i,j] for j in 1:length(Q)) ==  p)
	end
	for (j,(y,q)) in enumerate(Q)
		JuMP.@constraint(model, sum(w[i,j] for i in 1:length(P)) ==  q)
	end
	JuMP.@objective(model, Min, sum(sum(w.*d)))
	JuMP.optimize!(model)
	# Looking at the solution
	if JuMP.termination_status(model) != MathOptInterface.OPTIMAL
		@warn "Termination status of optimization: $(JuMP.termination_status(model))"
	end
	flow = JuMP.value.(w)
	emd = JuMP.objective_value(model)

	return (emd=emd, flow=flow)
end

##
## Specialized functions for BioSequences and BioTools
##
"""
	EMD(f::Function, X::Array{<:BioSequence,1},Y::Array{<:BioSequence,1})
	EMD(X::Array{<:BioSequence,1},Y::Array{<:BioSequence,1})
"""
function EMD(f::Function, X::Array{<:BioSequence,1},Y::Array{<:BioSequence,1})
	P = sequence_distribution(X)
	Q = sequence_distribution(Y)
	ret = EMD(f, P, Q)
	flow = Dict()
	for (i,(x,p)) in enumerate(P)
		for(j,(y,q)) in enumerate(Q)
			flow[x,y] = ret.flow[i,j]
		end
	end
	return (emd=ret.emd, flow=flow), P, Q
end
EMD(X::Array{<:BioSequence,1},Y::Array{<:BioSequence,1}) = EMD(BioTools.hamming, X, Y)
EMD(f::Function, X::Array{<:BioTools.AbstractStrain,1}, Y::Array{<:BioTools.AbstractStrain,1}) = EMD(f,
	[x.seq for x in X], [y.seq for y in Y])
EMD(X::Array{<:BioTools.AbstractStrain,1}, Y::Array{<:BioTools.AbstractStrain,1}) = EMD([x.seq for x in X], [y.seq for y in Y])
EMD(P::AbstractDict{<:BioSequence,Float64}, Q::AbstractDict{<:BioSequence,Float64}) = EMD(BioTools.hamming, P, Q)

"""
	sequence_distribution(X::Array{<:BioSequence})
	sequence_distribution(X::Array{<:AbstractStrain})

Convenience function. Transform an array `X` of sequences or strains into a distribution form usable by the EMD functions, *i.e.* a `Dict{BioSequence => weight::Float64}`.
"""
function sequence_distribution(X::Array{<:BioSequence})
	d = Dict{BioSequence, Float64}()
	for s in X
		d[s] = get(d, s, 0.) + 1. /length(X)
	end
	return d
end
function sequence_distribution(X::Array{<:BioTools.Strain})
	d = Dict{BioSequence, Float64}()
	for s in X
		d[s.seq] = get(d, s.seq, 0.) + 1. /length(X)
	end
	return d
end

end


@testset "bayesian.jl" begin
    @testset "fitbayes(model; niter, warmup)" begin

        # stationary GEV Bayesian fit
        n = 5000

        μ = 0.0
        σ = 1.0
        ξ = 0.1

        ϕ = log(σ)
        θ = [μ; ϕ; ξ]

        pd = GeneralizedExtremeValue(μ, σ, ξ)
        y = rand(pd, n)

        model = Extremes.BlockMaxima(Variable("y", y))

        fm = Extremes.fitbayes(model, niter=500, warmup=400)

        cinterval = cint(fm)

        @test [x[1] for x in cinterval] <= θ
        @test θ <= [x[2] for x in cinterval]

        # non-stationary GEV Bayesian fit
        n = 5000

        x₁ = randn(n)
        x₂ = randn(n)
        x₃ = randn(n) / 10

        μ = 1.0 .+ x₁
        ϕ = -0.5 .+ x₂
        ξ = x₃

        ϕ = log(σ)
        θ = [1.0; 1.0; -0.5; 1.0; 0.0; 1.0]

        pd = GeneralizedExtremeValue.(μ, σ, ξ)
        y = rand.(pd)

        model = Extremes.BlockMaxima(Variable("y", y), locationcov = [Variable("x₁", x₁)], logscalecov = [Variable("x₂", x₂)], shapecov = [Variable("x₃", x₃)])

        fm = Extremes.fitbayes(model, niter=500, warmup=400)

        cinterval = cint(fm)

        @test [x[1] for x in cinterval] <= θ
        @test θ <= [x[2] for x in cinterval]

        # stationary GP bayes fit
        n = 5000

        σ = 1.0
        ξ = 0.1

        ϕ = log(σ)
        θ = [ϕ; ξ]

        pd = GeneralizedPareto(σ, ξ)
        y = rand(pd, n)

        model = Extremes.ThresholdExceedance(Variable("y", y))

        fm = Extremes.fitbayes(model, niter=500, warmup=400)

        cinterval = cint(fm)

        @test [x[1] for x in cinterval] <= θ
        @test θ <= [x[2] for x in cinterval]

        # non-stationary GP Bayesian fit
        n = 5000

        x₁ = randn(n) / 3
        x₂ = randn(n) / 10

        ϕ = -.5 .+ x₁
        ξ = x₂

        σ = exp.(ϕ)
        θ = [-.5; 1.0; 0.0; 1.0]

        pd = GeneralizedPareto.(σ, ξ)
        y = rand.(pd)

        model = Extremes.ThresholdExceedance(Variable("y", y), logscalecov = [Variable("x₁", x₁)], shapecov = [Variable("x₂", x₂)])

        fm = Extremes.fitbayes(model, niter=500, warmup=400)

        cinterval = cint(fm)

        @test [x[1] for x in cinterval] <= θ
        @test θ <= [x[2] for x in cinterval]

end

    @testset "parametervar(model)" begin

        n = 1000

        μ = 0.0
        σ = 1.0
        ξ = 0.1

        ϕ = log(σ)
        θ = [μ; ϕ; ξ]

        pd = GeneralizedExtremeValue(μ, σ, ξ)
        y = rand(pd, n)

        fm = Extremes.gevfitbayes(y, niter=1000, warmup = 500)
        npar = 3 + Extremes.getcovariatenumber(fm.model)
        @test size(Extremes.parametervar(fm)) == (npar,npar)

    end

    include(joinpath("bayesian", "bayesian_gev_test.jl"))
    include(joinpath("bayesian", "bayesian_gp_test.jl"))

end


using Printf
using Distances

function check_for_random_data()
    nbspectra = 20000
    nbval = 1000
    @debug "check_for_random_data: dispatching nbspecta : " nbspectra
    spectra = Array{Array{Float64,1},1}(undef, nbspectra)
    for i in 1:nbspectra
        spectra[i] = rand(nbval)
    end
    D=Jaccard()
    depth = 8
    argument = RPTreeArg(D, 8, 1.5)
    rptree=RPTree(argument, spectra)
    @time leafCenters = randomProjection(rptree)
    leaves, diameters = analyzeSplittingInfo(rptree)
    #
    if length(leaves) != 2^depth
        @printf stdout "\n bad number of leaves , shoud be %d, got %d " 2^8  length(leaves)
        return false
    end
    #
    return true
end


################################################
struct VRPTW_Instance
    travel_time     ::Matrix{Float64}
    service_time    ::Vector{Float64}
    early_time      ::Vector{Float64}
    late_time       ::Vector{Float64}
    load            ::Vector{Float64}
    capacity        ::Float64
    max_travel_time ::Float64
end
################################################

################################################
# ESPPRC algorithm
################################################

mutable struct Label
    time        ::Float64
    load        ::Float64
    flag        ::Vector{Int}
    cost        ::Float64
    path        ::Vector{Int}
    function Label(time::Float64, load::Float64, flag::Vector{Int}, cost::Float64, path::Vector{Int})
        time = round(time, digits=DIGITS)
        load = round(load, digits=DIGITS)
        cost = round(cost, digits=DIGITS)
        new(time, load, flag, cost, path)
    end
end

mutable struct ESPPRC_Instance
    origin      :: Int64
    destination :: Int64
    capacity    :: Float64
    cost        :: Matrix{Float64}
    time        :: Matrix{Float64}
    load        :: Matrix{Float64}
    early_time  :: Vector{Float64}
    late_time   :: Vector{Float64}
    service_time:: Vector{Float64}
    forward_star    :: Vector{Vector{Int}}
    reverse_star    :: Vector{Vector{Int}}
    critical_nodes  :: Set{Int}
    max_T           :: Float64
    max_neg_routes  :: Int
end
function ESPPRC_Instance(origin, destination, capacity, cost, time, load, early_time, late_time, service_time)
    n_nodes = length(service_time)
    fs = save_forward_star(n_nodes, destination, cost; sorted=true)
    rs = save_reverse_star(n_nodes, origin, cost; sorted=true)
    critical_nodes = Set(1:n_nodes)
    max_T = calculate_max_T(destination, time, early_time, late_time, service_time)
    max_neg_routes = MAX_INT
    return ESPPRC_Instance(origin, destination, capacity, cost, time, load, early_time, late_time, service_time,
                            fs, rs, critical_nodes, max_T, max_neg_routes)
end
################################################

################################################
#  Branch and bounds
################################################
struct Branch
    history         :: OrderedDict
    vrptw_instance  :: VRPTW_Instance
    root_routes     :: Array
    branch_priority :: Array
    lower_bound     :: Float64
    veh_cond        :: Tuple{String, Int64}
end

mutable struct BestIncumbent
    y::Vector{Float64}
    routes::Vector{Vector{Int}}
    objective::Float64
end

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

################################################
# solomon datasets

struct Node
    id::Int
    # type::Int
    cx::Float64
    cy::Float64
end

struct Fleet
    # type::Int
    number::Int
    # departure_node::Int
    # arrival_node::Int
    capacity::Float64
    max_travel_time::Float64
end

struct Request
    id::Int
    start_time::Int
    end_time::Int
    quantity::Float64
    service_time::Float64
end

struct Solomon
    nodes::Vector{Node}
    fleet::Fleet
    requests::Vector{Request}
end
################################################


using LinearKernelModels
using LinearAlgebra, OffsetArrays
using Test

@testset "Single stimulus" begin
    S = rand(100)
    ktrue = OffsetArray(randn(5), -4:0)
    r = compute_r(S, ktrue)
    C, d = compute_Cd(S, r, 5)
    kcalc = reshape(C \ d, axes(ktrue))
    @test kcalc ≈ ktrue
    nfailed = 0
    for i = 1:100
        rnoisy = r .+ randn(size(r))
        C, d = compute_Cd(S, rnoisy, 5)
        knoisy = reshape(C \ d, axes(ktrue))
        kerr = compute_kerr(knoisy, C, S, rnoisy)
        nfailed += any((knoisy .- ktrue).^2 .> 3*kerr)
    end
    @test nfailed <= 5

    nfailed = 0
    for i = 1:100
        S = rand(10000)
        r = compute_r(S, ktrue)
        rnoisy = r .+ randn(size(r))
        C, d = compute_Cd(S, rnoisy, 5)
        knoisy = reshape(C \ d, axes(ktrue))
        kerr = compute_kerr(knoisy, C, S, rnoisy)
        nfailed += any((knoisy .- ktrue).^2 .> 3*kerr)
    end
    @test nfailed <= 5
end

@testset "Single acausal stimulus" begin
    S = rand(100)
    ktrue = OffsetArray(randn(5), -2:2)
    r = compute_r(S, ktrue)
    C, d = compute_Cd(S, r, axes(ktrue, 1))
    kcalc = reshape(C \ d, axes(ktrue))
    @test kcalc ≈ ktrue
end

@testset "Multiple stimuli and NaNs" begin
    S = rand(100, 3)
    ktrue = OffsetArray(randn(5, 3), -4:0, 1:3)
    r = compute_r(S, ktrue)
    C, d = compute_Cd(S, r, 5)
    kcalc = reshape(C \ d, axes(ktrue))
    @test kcalc ≈ ktrue

    r[5:10] .= NaN
    r[50:52] .= NaN
    C, d = compute_Cd(S, r, 5)
    kcalc = reshape(C \ d, axes(ktrue))
    @test kcalc ≈ ktrue
end

@testset "Multiple acausal stimuli and NaNs" begin
    S = rand(100,3)
    ktrue = OffsetArray(randn(5,3), -2:2, 1:3)
    r = compute_r(S, ktrue)
    C, d = compute_Cd(S, r, -2:2)
    kcalc = reshape(C \ d, axes(ktrue))
    @test kcalc ≈ ktrue

    r[5:10] .= NaN
    r[50:52] .= NaN
    C, d = compute_Cd(S, r, -2:2)
    kcalc = reshape(C \ d, axes(ktrue))
    @test kcalc ≈ ktrue
end

@testset "Single stimulus with constraints" begin
    S = rand(100)
    ktrue = OffsetArray([rand(3); 0; 0], -2:2)
    r = compute_r(S, ktrue)
    kcalc = solve_for_kernel(S, r, ktrue.!=0)
    @test kcalc ≈ ktrue
end

@testset "Multiple stimuli with constraints" begin
    S = rand(100,3)
    ktrue = OffsetArray(randn(5,3), -2:2, 1:3)
    ktrue[rand(1:15,5,1)] .= 0
    r = compute_r(S,ktrue)
    kcalc = solve_for_kernel(S, r, ktrue.!=0)
    @test kcalc ≈ ktrue

    # With a degeneracy that is broken by the constraint
    S = zeros(20, 2)
    S[10,:] .= 1
    r = zeros(20)
    r[11:12] .= 0.5
    kmask = OffsetArray([true false; false true], -2:-1, 1:2)
    kcalc = solve_for_kernel(S, r, kmask)
    @test kcalc ≈ OffsetArray([0.5 0; 0 0.5], -2:-1, 1:2)
end

@testset "Multiple stimuli with constraints and Nans" begin
    S = rand(100,7)
    ktrue = OffsetArray(rand(5,7), -2:2, 1:7)
    ktrue[rand(1:35, 10)] .= 0
    r = compute_r(S,ktrue)
    r[5:10] .= NaN
    r[50:52] .= NaN
    kcalc = solve_for_kernel(S, r, ktrue.!=0)
    @test kcalc ≈ ktrue
end


using ClimateMachine.Mesh.Grids
using ClimateMachine.Mesh.Topologies: BrickTopology
using Test
using MPI

MPI.Initialized() || MPI.Init()

@testset "2-D Mass matrix" begin
    topology = BrickTopology(MPI.COMM_WORLD, ntuple(_ -> (-1, 1), 2);)

    @testset for N in ((2, 2), (2, 3), (3, 2))
        Nq = N .+ 1
        Np = prod(Nq)
        Nfp = div.(Np, Nq)

        grid = DiscontinuousSpectralElementGrid(
            topology,
            FloatType = Float64,
            DeviceArray = Array,
            polynomialorder = N,
        )

        @views begin
            ω = grid.ω
            M = grid.vgeo[:, Grids._M, 1]
            MH = grid.vgeo[:, Grids._MH, 1]
            sM = grid.sgeo[Grids._sM, :, :, 1]

            M = reshape(M, Nq[1], Nq[2])
            @test M ≈ [ω[1][i] * ω[2][j] for i in 1:Nq[1], j in 1:Nq[2]]

            MH = reshape(MH, Nq[1], Nq[2])
            @test MH ≈ [ω[1][i] for i in 1:Nq[1], j in 1:Nq[2]]

            sM12 = sM[1:Nfp[1], 1:2]
            @test sM12[:, 1] ≈ [ω[2][j] for j in 1:Nq[2]]
            @test sM12[:, 2] ≈ [ω[2][j] for j in 1:Nq[2]]

            sM34 = sM[1:Nfp[2], 3:4]
            @test sM34[:, 1] ≈ [ω[1][j] for j in 1:Nq[1]]
            @test sM34[:, 2] ≈ [ω[1][j] for j in 1:Nq[1]]
        end
    end
end

@testset "3-D Mass matrix" begin
    topology = BrickTopology(MPI.COMM_WORLD, ntuple(_ -> (-1, 1), 3);)

    @testset for N in ((2, 2, 2), (2, 3, 4), (4, 3, 2), (2, 4, 3))
        Nq = N .+ 1
        Np = prod(Nq)
        Nfp = div.(Np, Nq)

        grid = DiscontinuousSpectralElementGrid(
            topology,
            FloatType = Float64,
            DeviceArray = Array,
            polynomialorder = N,
        )

        @views begin
            ω = grid.ω
            M = grid.vgeo[:, Grids._M, 1]
            MH = grid.vgeo[:, Grids._MH, 1]
            sM = grid.sgeo[Grids._sM, :, :, 1]

            M = reshape(M, Nq[1], Nq[2], Nq[3])
            @test M ≈ [
                ω[1][i] * ω[2][j] * ω[3][k]
                for i in 1:Nq[1], j in 1:Nq[2], k in 1:Nq[3]
            ]

            MH = reshape(MH, Nq[1], Nq[2], Nq[3])
            @test MH ≈ [
                ω[1][i] * ω[2][j] for i in 1:Nq[1], j in 1:Nq[2], k in 1:Nq[3]
            ]

            sM12 = reshape(sM[1:Nfp[1], 1:2], Nq[2], Nq[3], 2)
            @test sM12[:, :, 1] ≈
                  [ω[2][j] * ω[3][k] for j in 1:Nq[2], k in 1:Nq[3]]
            @test sM12[:, :, 2] ≈
                  [ω[2][j] * ω[3][k] for j in 1:Nq[2], k in 1:Nq[3]]

            sM34 = reshape(sM[1:Nfp[2], 3:4], Nq[1], Nq[3], 2)
            @test sM34[:, :, 1] ≈
                  [ω[1][i] * ω[3][k] for i in 1:Nq[1], k in 1:Nq[3]]
            @test sM34[:, :, 2] ≈
                  [ω[1][i] * ω[3][k] for i in 1:Nq[1], k in 1:Nq[3]]

            sM56 = reshape(sM[1:Nfp[3], 5:6], Nq[1], Nq[2], 2)
            @test sM56[:, :, 1] ≈
                  [ω[1][i] * ω[2][j] for i in 1:Nq[1], j in 1:Nq[2]]
            @test sM56[:, :, 2] ≈
                  [ω[1][i] * ω[2][j] for i in 1:Nq[1], j in 1:Nq[2]]
        end
    end
end


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

using Markdown
using InteractiveUtils

# ╔═╡ 3da9f1a8-f65a-11ea-2a34-c7eb4db8cd54
using Plots, Distributions, LinearAlgebra,Distributions, ProgressMeter

# ╔═╡ 316f1058-f65a-11ea-33d8-997377fba2b0
md"Preliminary: load packages, setup latex macro"

# ╔═╡ b37d84e0-f659-11ea-0e0e-1d14d3d5aba5
md"# Calculating the ELBO for PFVG
Here we detail the various calculation we need."

# ╔═╡ 5cabfa9e-f65a-11ea-0695-c73c26cf3342
md"## Model
$$\int \mathcal{N}(\mathbf{y}|\mathbf{A}\mathbf{f} - \mathbf{b}, \mathbf{S}) \mathcal{N}(\mathbf{f}|\mathbf{0},\mathbf{K}) \ \mathbf{df} \ \mathbf{dA}\ \mathbf{db}$$
where the matrix $$\mathbf{A}(\mathbf{\alpha})$$ and vector $$\mathbf{b}(\mathbf{\beta})$$ are function of the respective parameters.
"

# ╔═╡ d24a45be-f65b-11ea-2c07-d19988003b56
md"## Lower bound
$$\int q(\mathbf{f}) q(\mathbf{\alpha}) q(\mathbf{\beta}) \log\mathcal{N}(\mathbf{y}|\mathbf{A}\mathbf{f} - \mathbf{b}, \mathbf{S}) \mathcal{N}(\mathbf{f}|\mathbf{0},\mathbf{K}) \ \mathbf{df} \ \mathbf{dA}\ \mathbf{db} + \mathcal{H}[q_f] + \mathcal{H}[q_\alpha] + \mathcal{H}[q_\beta]$$
"

# ╔═╡ 3e54363e-f65c-11ea-0150-7d09f848a86e
md"## Inside the log
$$-\frac{1}{2}(\mathbf{y} - \mathbf{Af}-\mathbf{b})^T\mathbf{S}^{-1} (\mathbf{y} - \mathbf{Af}-\mathbf{b}) - \frac{1}{2}\log|\mathbf{K}| -\frac{1}{2}\mathbf{f}^T\mathbf{K}^{-1}\mathbf{f}$$
"

# ╔═╡ ab41e982-f65c-11ea-26f9-01f9de7053a3
md"## Work out posterior $$q(\mathbf{f})$$ of GP latent variables

Keep only relevant terms:

$$-\frac{1}{2}(-2\mathbf{y}^T\mathbf{S}^{-1}\mathbf{Af} + 2\mathbf{b}^T\mathbf{S}^{-1}\mathbf{Af} + \mathbf{f^T A^T S^{-1} A f})
-\frac{1}{2} \mathbf{f}^T\mathbf{K}^{-1}\mathbf{f}$$

We recognise the quadratic which leads to a Gaussian posterior:

$$-\frac{1}{2}(-2(\mathbf{-y+b})^T\mathbf{S}^{-1}\mathbf{Af}  + \mathbf{f^T A^T S^{-1} A f})
-\frac{1}{2} \mathbf{f}^T\mathbf{K}^{-1}\mathbf{f}$$

$$(\mathbf{y-b})^T\mathbf{S}^{-1}\mathbf{Af}  - \frac{1}{2} \mathbf{f^T (A^T S^{-1} A +K^{-1}) f}$$

The posterior reads:

$$q(\mathbf{f}) = \mathcal{N}(\mathbf{f} |  \mathbf{\Sigma_f} \ \mathbf{A^T S^{-1}(y-b)}, \ \underbrace{\mathbf{A^T S^{-1} A +K^{-1}}}_{\mathbf{\Sigma_f}} )$$

"

# ╔═╡ fe789bd2-f66e-11ea-3400-61f55bbf640b

# ╔═╡ Cell order:
# ╟─316f1058-f65a-11ea-33d8-997377fba2b0
# ╠═3da9f1a8-f65a-11ea-2a34-c7eb4db8cd54
# ╟─b37d84e0-f659-11ea-0e0e-1d14d3d5aba5
# ╟─5cabfa9e-f65a-11ea-0695-c73c26cf3342
# ╟─d24a45be-f65b-11ea-2c07-d19988003b56
# ╟─3e54363e-f65c-11ea-0150-7d09f848a86e
# ╟─ab41e982-f65c-11ea-26f9-01f9de7053a3
# ╠═fe789bd2-f66e-11ea-3400-61f55bbf640b


# ## Data

y = [-1, 1]

# ## Model

import Random

using Turing

Random.seed!(1)

@model function m8_2(y)
    α ~ Flat() ## improper prior with pobability one everywhere
    σ ~ FlatPos(0.0) ## improper prior with probability one everywhere above 0.0

    y ~ Normal(α, σ)
end;

# ## Output

chns = sample(m8_2(y), NUTS(), 1000)

# \defaultoutput{}


function orad(ZO)

  R=0.62; g1=0.6; g2=20;

  f=R*exp.(ZO/g1)+(1-R)*exp.(ZO/g2);
  f[ZO.<-200].=0;
  f=-diff(f,dims=1);

  return(f)

end


struct Rectangle <: Surface
    p1::AbstractVec
    p2::AbstractVec
    p3::AbstractVec
    v1::AbstractVec
    v2::AbstractVec
    normal::AbstractVec
    d::AbstractFloat
    u2::AbstractFloat
    u4::AbstractFloat

    function Rectangle(p1::AbstractVec, p2::AbstractVec, p3::AbstractVec)
        v1::AbstractVec = p2 - p1
        v2::AbstractVec = p3 - p2
        normal::AbstractVec = normalize(cross(v2, v1))
        d::AbstractFloat = -dot(normal, p1)
        u2::AbstractFloat = dot(v1, v1)
        u4::AbstractFloat = dot(v2, v2)
        new(p1, p2, p3, v1, v2, normal, d, u2, u4)
    end
end

function hit(rect::Rectangle, ray::Ray, range::NamedTuple{(:min, :max), Tuple{T, T}} where T <: AbstractFloat)::Tuple{Bool, AbstractFloat, Union{Nothing, HitPoint}}
    b::AbstractFloat = dot(rect.normal, ray.direction)

    if b != 0
        a::AbstractFloat = -dot(rect.normal, ray.origin) - rect.d
        t::AbstractFloat = a / b
        if range.min < t < range.max
            point::AbstractVec = point_at(ray, t)
            v::AbstractVec = point - rect.p1
            u1::AbstractFloat = dot(v, rect.v1)
            u3::AbstractFloat = dot(v, rect.v2)

            if 0 < u1 < rect.u2 && 0 < u3 < rect.u4
                if a < 0
                    return true, t, HitPoint(point, rect.normal, u3 / rect.u4, u1 / rect.u2, true)
                else
                    return true, t, HitPoint(point, -rect.normal, u3 / rect.u4, u1 / rect.u2, false)
                end
            end
        end
    end

    false, NaN, nothing
end


export fast_αscheme
function fast_αscheme(
    m::Int64, n::Int64, h::Float64,
    M::Union{SparseTensor, SparseMatrixCSC},
    C::Union{SparseTensor, SparseMatrixCSC},
    K::Union{SparseTensor, SparseMatrixCSC},
    body_force::Union{Array{Float64}, PyObject},
    d0::Union{Array{Float64, 1}, PyObject},
    v0::Union{Array{Float64, 1}, PyObject},
    a0::Union{Array{Float64, 1}, PyObject},
    Δt::Float64,
    bdnode::Union{Array{Int64, 1}, PyObject, Missing} = missing,
    bd_d::Union{Array{Float64, 2}, PyObject, Missing} = missing,
    bd_v::Union{Array{Float64, 2}, PyObject, Missing} = missing,
    bd_a::Union{Array{Float64, 2}, PyObject, Missing} = missing;
    ρ::Float64 = 1.0)
    nt = size(body_force, 1)
    if !ismissing(bdnode)
        if ismissing(bd_a)
            error("Boundary acceleration at Dirichlet nodes must be provided.")
        end
        for v in [bd_d, bd_v, bd_a]
            if !ismissing(v) && (size(v,1)!= nt || size(v,2)!= 2length(bdnode))
                error("Invalid shape for $v. Expected $(nt)×$(2length(bdnode)) but got $(size(v,1))×$(size(v,2))")
            end
        end
    end
    !ismissing(bd_d) && (bd_d = convert_to_tensor(bd_d, dtype=Float64))
    !ismissing(bd_v) && (bd_v = convert_to_tensor(bd_v, dtype=Float64))
    !ismissing(bd_a) && (bd_a = convert_to_tensor(bd_a, dtype=Float64))

    αm = (2ρ-1)/(ρ+1)
    αf = ρ/(1+ρ)
    γ = 1/2-αm+αf
    β = 0.25*(1-αm+αf)^2
    d = length(d0)

    M = isa(M, SparseMatrixCSC) ? constant(M) : M
    C = isa(C, SparseMatrixCSC) ? constant(C) : C
    K = isa(K, SparseMatrixCSC) ? constant(K) : K
    body_force, d0, v0, a0, Δt = convert_to_tensor([body_force, d0, v0, a0, Δt], [Float64, Float64, Float64, Float64, Float64])

    A = (1-αm)*M + (1-αf)*C*Δt*γ + (1-αf)*K*β*Δt^2
    if !ismissing(bdnode)
        A, Abd = fem_impose_Dirichlet_boundary_condition(A, bdnode, m, n, h)
    else
        Abd = missing
    end
    A = factorize(A)

    function equ(dc, vc, ac, dt, body_force, i)
        dn = dc + dt*vc + dt^2/2*(1-2β)*ac
        vn = vc + dt*((1-γ)*ac)

        df = (1-αf)*dn + αf*dc
        vf = (1-αf)*vn + αf*vc
        am = αm*ac

        rhs = body_force - (M*am + C*vf + K*df)
        if !ismissing(Abd)
            rhs -= Abd * bd_a[i]
            rhs = scatter_update(rhs, [bdnode; bdnode .+ (m+1)*(n+1)], bd_a[i])
        end

        A\rhs
    end

    function condition(i, tas...)
        return i<=nt
    end
    function body(i, tas...)
        dc_arr, vc_arr, ac_arr = tas
        dc = read(dc_arr, i)
        vc = read(vc_arr, i)
        ac = read(ac_arr, i)
        y = equ(dc, vc, ac, Δt, body_force[i], i)
        dn = dc + Δt*vc + Δt^2/2*((1-2β)*ac+2β*y)
        vn = vc + Δt*((1-γ)*ac+γ*y)

        if !ismissing(bd_v)
            vn = scatter_update(vn, [bdnode; bdnode .+ (m+1)*(n+1)], bd_v[i])
        end

        if !ismissing(bd_d)
            dn = scatter_update(dn, [bdnode; bdnode .+ (m+1)*(n+1)], bd_d[i])
        end

        i+1, write(dc_arr, i+1, dn), write(vc_arr, i+1, vn), write(ac_arr, i+1, y)
    end

    dM = TensorArray(nt+1); vM = TensorArray(nt+1); aM = TensorArray(nt+1)
    dM = write(dM, 1, d0)
    vM = write(vM, 1, v0)
    aM = write(aM, 1, a0)
    i = constant(1, dtype=Int32)
    _, d, v, a = while_loop(condition, body, [i,dM, vM, aM])
    set_shape(stack(d), (nt+1, length(d0))), set_shape(stack(v), (nt+1, length(v0))), set_shape(stack(a), (nt+1, length(a0)))
end


function adjoint_solve(prob, sol_init, u₀, tstops, p, rtol, atol)
    de = ModelingToolkit.modelingtoolkitize(prob)
    sym_jac = eval(ModelingToolkit.generate_jacobian(de)[1])
    function adjoint!(du, u, p, t)
        du = -Base.invokelatest(sym_jac, sol_init(t), p, t)' * u
    end
    rev_tstops = reverse(tstops) # reverse tstops for λ(T) = u₀?
    rev_tspan = (rev_tstops[1], rev_tstops[end])
    adj_prob = ODEProblem(adjoint!, u₀, rev_tspan, p)
    λ = solve(adj_prob, sol_init.alg, reltol=rtol, abstol=atol, tstops=rev_tstops)
    return λ
end

function condition_num(λₜ)
    K = LinearAlgebra.norm(λₜ, 1) + LinearAlgebra.norm(λₜ[1])
    return K
end

function adjoint_step_ctrl(λ₁ₜ, λ₂ₜ, E₂, Eₙ, C, gtol)
    K₁ = condition_num(λ₁ₜ)
    K₂ = condition_num(λ₂ₜ)
    K = (E₂ / Eₙ) * sqrt(K₁^2 + K₂^2)
    h = sqrt((gtol / (K * C))) # T in denominator?
    return h
end


"""
A wrapper for the interpolation methods of DataInterpolations.jl.

$(SIGNATURES)

Wraps the methods in such a way that they are callable as `f(u,t)` to
create and return an interpolation of `u` over `t`.
The first argument of the constructor always defines the interpolation method,
all following arguments will be used in the interpolation.

The additional keyword `crop = false` indicates to discard the first and last element of the time series.

# Example

```julia
# Create the wrapper struct
itp_method = InterpolationMethod(QuadraticSpline)
# Create a callable interpolation
itp = itp_method(u,t)
# Return u[2]
itp(t[2])
```
"""
struct InterpolationMethod{T} <: AbstractInterpolationMethod
  itp::T
  args

  function InterpolationMethod(itp, args...)
    return new{typeof(itp)}(itp, args)
  end

end

(x::InterpolationMethod)(u, t) = x.itp(u,t,x.args...)

# TODO Wrap all types
# Wrap the common itps
InterpolationMethod() = InterpolationMethod(LinearInterpolation)

# Taken from DiffEqFlux
# https://github.com/SciML/DiffEqFlux.jl/blob/master/src/collocation.jl
# On 3-11-2021

struct EpanechnikovKernel <: CollocationKernel end
struct UniformKernel <: CollocationKernel end
struct TriangularKernel <: CollocationKernel end
struct QuarticKernel <: CollocationKernel end
struct TriweightKernel <: CollocationKernel end
struct TricubeKernel <: CollocationKernel end
struct GaussianKernel <: CollocationKernel end
struct CosineKernel <: CollocationKernel end
struct LogisticKernel <: CollocationKernel end
struct SigmoidKernel <: CollocationKernel end
struct SilvermanKernel <: CollocationKernel end

function calckernel(::EpanechnikovKernel,t)
    if abs(t) > 1
        return 0
    else
        return 0.75*(1-t^2)
    end
end

function calckernel(::UniformKernel,t)
    if abs(t) > 1
        return 0
    else
        return 0.5
    end
end

function calckernel(::TriangularKernel,t)
    if abs(t) > 1
        return 0
    else
        return (1-abs(t))
    end
end

function calckernel(::QuarticKernel,t)
  if abs(t)>0
    return 0
  else
    return (15*(1-t^2)^2)/16
  end
end

function calckernel(::TriweightKernel,t)
  if abs(t)>0
    return 0
  else
    return (35*(1-t^2)^3)/32
  end
end

function calckernel(::TricubeKernel,t)
  if abs(t)>0
    return 0
  else
    return (70*(1-abs(t)^3)^3)/80
  end
end

function calckernel(::GaussianKernel,t)
  exp(-0.5*t^2)/(sqrt(2*π))
end

function calckernel(::CosineKernel,t)
  if abs(t)>0
    return 0
  else
    return (π*cos(π*t/2))/4
  end
end

function calckernel(::LogisticKernel,t)
  1/(exp(t)+2+exp(-t))
end

function calckernel(::SigmoidKernel,t)
  2/(π*(exp(t)+exp(-t)))
end

function calckernel(::SilvermanKernel,t)
  sin(abs(t)/2+π/4)*0.5*exp(-abs(t)/sqrt(2))
end

function construct_t1(t,tpoints)
    hcat(ones(eltype(tpoints),length(tpoints)),tpoints.-t)
end

function construct_t2(t,tpoints)
  hcat(ones(eltype(tpoints),length(tpoints)),tpoints.-t,(tpoints.-t).^2)
end

function construct_w(t,tpoints,h,kernel)
    W = @. calckernel((kernel,),(tpoints-t)/h)/h
    Diagonal(W)
end

"""
$(SIGNATURES)

Unified interface for collocation techniques. The input can either be
a `CollocationKernel` (see list below) or a wrapped `InterpolationMethod` from
[DataInterpolations.jl](https://github.com/PumasAI/DataInterpolations.jl).

Computes a non-parametrically smoothed estimate of `u'` and `u`
given the `data`, where each column is a snapshot of the timeseries at
`tpoints[i]`.

# Examples
```julia
u′,u,t = collocate_data(data,tpoints,kernel=SigmoidKernel())
u′,u,t = collocate_data(data,tpoints,tpoints_sample,interp,args...)
u′,u,t = collocate_data(data,tpoints,interp)
```

# Collocation Kernels
See [this paper](https://www.ncbi.nlm.nih.gov/pmc/articles/PMC2631937/) for more information.
+ EpanechnikovKernel
+ UniformKernel
+ TriangularKernel
+ QuarticKernel
+ TriweightKernel
+ TricubeKernel
+ GaussianKernel
+ CosineKernel
+ LogisticKernel
+ SigmoidKernel
+ SilvermanKernel

# Interpolation Methods
See [DataInterpolations.jl](https://github.com/PumasAI/DataInterpolations.jl) for more information.
+ ConstantInterpolation
+ LinearInterpolation
+ QuadraticInterpolation
+ LagrangeInterpolation
+ QuadraticSpline
+ CubicSpline
+ BSplineInterpolation
+ BSplineApprox
+ Curvefit
"""
function collocate_data(data,tpoints,kernel=TriangularKernel(); crop = false, kwargs...)
  _one = oneunit(first(data))
  _zero = zero(first(data))
  e1 = [_one;_zero]
  e2 = [_zero;_one;_zero]
  n = length(tpoints)
  h = (n^(-1/5))*(n^(-3/35))*((log(n))^(-1/16))

  Wd = similar(data, n, size(data,1))
  WT1 = similar(data, n, 2)
  WT2 = similar(data, n, 3)
  x = map(tpoints) do _t
      T1 = construct_t1(_t,tpoints)
      T2 = construct_t2(_t,tpoints)
      W = construct_w(_t,tpoints,h,kernel)
      mul!(Wd,W,data')
      mul!(WT1,W,T1)
      mul!(WT2,W,T2)
      (e2'*((T2'*WT2)\T2'))*Wd,(e1'*((T1'*WT1)\T1'))*Wd
  end
  estimated_derivative = reduce(hcat,transpose.(first.(x)))
  estimated_solution = reduce(hcat,transpose.(last.(x)))
  crop && return estimated_derivative[:,2:end-1], estimated_derivative[:, 2:end-1], tpoints[2:end-1]
  estimated_derivative,estimated_solution, tpoints
end

# Adapted to dispatch on InterpolationMethod
collocate_data(data, tpoints, interp::InterpolationMethod; kwargs...) = collocate_data(data, tpoints, tpoints, interp; kwargs...)

function collocate_data(data::AbstractVector,tpoints::AbstractVector,tpoints_sample::AbstractVector,
                        interp::InterpolationMethod; kwargs...)
  u, du, tpoints = collocate_data(reshape(data, 1, :),tpoints,tpoints_sample,interp; kwargs...)
  return du[1, :], u[1, :] , tpoints
end

# Adapted to dispatch on InterpolationMethod
function collocate_data(data::AbstractMatrix{T},tpoints::AbstractVector{T},
                        tpoints_sample::AbstractVector{T},interp::InterpolationMethod; crop = false, kwargs...) where T

  u = zeros(T,size(data, 1),length(tpoints_sample))
  du = zeros(T,size(data, 1),length(tpoints_sample))
  for d1 in 1:size(data,1)
    interpolation = interp(data[d1,:],tpoints)
    u[d1,:] .= interpolation.(tpoints_sample)
    du[d1,:] .= DataInterpolations.derivative.((interpolation,), tpoints_sample)
  end
  crop && return du[:, 2:end-1], u[:, 2:end-1], tpoints[2:end-1]
  return du, u, tpoints
end


using Test

new_floor = function(direction, x)
    if direction == '('
        return x+1
    elseif direction == ')'
        return x-1
    else
        return x
    end
end

day1_part1 = function(x)
    floor = 0
    for c in x
        floor = new_floor(c, floor)
    end
    return floor
end

@testset "day1_part1" begin
    @test 0 == day1_part1("")
    @test 1 == day1_part1("(")
    @test 2 == day1_part1("((")
    @test 1 == day1_part1("(()")
    @test 0 == day1_part1("(())")
    @test 0 == day1_part1("()()")
    @test 3 == day1_part1("(((")
    @test 3 == day1_part1("(()(()(")
    @test 3 == day1_part1("))(((((")
    @test -1 == day1_part1("())")
    @test -1 == day1_part1("))(")
    @test -3 == day1_part1(")))")
    @test -3 == day1_part1(")())())")
    @test 0 == day1_part1("()\n")
end

day1_part2 = function(x)
    str = ""
    for c in x
        str *= string(c)
        if -1 == day1_part1(str)
            break
        end
    end
    return length(str)
end

@testset "day1_part2" begin
    @test 1 == day1_part2(")")
    @test 5 == day1_part2("()())")
end

const day1_input = read("1.txt", String)
println("""Day 1, Part 1: $(day1_part1(day1_input))""")
println("""Day 1, Part 2: $(day1_part2(day1_input))""")