using OrdinaryDiffEq
using DiffEqCallbacks
using Surrogates
using Random
using Optim
using Plots
using BenchmarkTools

function sir_ode!(du,u,p,t)
    (S,I,R) = u
    (β,γ) = p
    @inbounds begin
        du[1] = -β*S*I
        du[2] = β*S*I - γ*I
        du[3] = γ*I
    end
    nothing
end;

tmax = 10000.0
tspan = (0.0,tmax);

cb_ss = TerminateSteadyState();

n_train = 25 # Number of samples
n_test = 1000; # Number of samples

# Parameters are β
lb = 0.00005
ub = 0.001;

N = 1000.0
u0 = [990.0,10.0,0.0]
p = [0.0005,0.25]
prob_ode = ODEProblem(sir_ode!,u0,tspan,p);

Random.seed!(123)
sampler = LatinHypercubeSample();

β = Surrogates.sample(n_train,lb,ub,sampler)
β₀ = copy(β);

final_size = function(β)
  p = prob_ode.p
  p[1] = β
  prob = remake(prob_ode;p=p)
  sol = solve(prob, ROS34PW3(),callback=cb_ss)
  fsp = sol[end][3]/N
  fsp
end;

fs = final_size.(β);

β_grid = lb:0.00001:ub
plot(β_grid,final_size.(β_grid),label="True function",xlab="Infectivity parameter, β",ylab="Final size")
scatter!(β,fs,legend=:right,label="Sampled points")

fs_surrogate = RadialBasis(β, fs, lb, ub, rad=cubicRadial);

scatter(fs,
        fs_surrogate.(β),
        legend=false,
        xlabel="True final size",
        ylabel="Surrogate final size")
Plots.abline!(1.0,0.0)

surrogate_optimize(final_size,
                   SRBF(),
                   lb,
                   ub,
                   fs_surrogate,
                   sampler;
                   maxiters=1000,
                   num_new_samples=1000);

β₁ = copy(β)
length(β₀),length(β)

fs_opt = final_size.(β)
scatter(fs_opt,fs_surrogate.(β), xlabel="Final size", ylabel="Surrogate final size",legend=false)
Plots.abline!(1.0, 0.0)

β_test = sample(n_test,lb,ub,UniformSample())
fs_test = final_size.(β_test)
fs_test_pred = fs_surrogate.(β_test);

l = @layout [a b]
pl1 = plot(β_grid,final_size.(β_grid),color=:red,label="Model",xlabel="β",ylabel="Final size",legend=:right)
scatter!(pl1, β_test,fs_test_pred,color=:blue,label="Surrogate")
pl2 = scatter(fs_test,fs_test_pred,color=:blue,legend=false,xlabel="Final size",ylabel="Surrogate final size")
Plots.abline!(pl2,1.0,0.0)
plot(pl1,pl2,layout=l)

@benchmark final_size(0.0005)

@benchmark fs_surrogate(0.0005)

peak_infected = function(β)
  p = prob_ode.p
  p[1] = β
  prob = remake(prob_ode;p=p)
  sol = solve(prob, ROS34PW3(),callback=cb_ss)
  tss = sol.t[end]
  f = (t) -> -sol(t,idxs=2)
  opt = Optim.optimize(f,0.0,tss)
  pk = -opt.minimum
  pk
end

plot(β_grid, peak_infected,xlabel="β",ylabel="Peak infected",legend=false)

pk = peak_infected.(β)
pk_surrogate = RadialBasis(β,pk,lb,ub,rad = cubicRadial)
surrogate_optimize(peak_infected,
                   SRBF(),
                   lb,
                   ub,
                   pk_surrogate,
                   sampler;
                   maxiters=1000,
                   num_new_samples=1000)
pk_test = peak_infected.(β_test)
pk_test_pred = pk_surrogate.(β_test);

l = @layout [a b]
pl1 = plot(β_grid,peak_infected.(β_grid),color=:red,label="Model",xlabel="β",ylabel="Peak infected",legend=:right)
scatter!(pl1, β_test,pk_test_pred,color=:blue,label="Surrogate")
pl2 = scatter(pk_test,pk_test_pred,color=:blue,legend=false,xlabel="Peak infected",ylabel="Surrogate peak infected")
Plots.abline!(pl2,1.0,0.0)
plot(pl1,pl2,layout=l)

@benchmark peak_infected(0.0005)

@benchmark pk_surrogate(0.0005)

peak_time = function(β)
  p = prob_ode.p
  p[1] = β
  prob = remake(prob_ode;p=p)
  sol = solve(prob, ROS34PW3(),callback=cb_ss)
  tss = sol.t[end]
  f = (t) -> -sol(t,idxs=2)
  opt = Optim.optimize(f,0.0,tss)
  pt = opt.minimizer
  pt
end

plot(β_grid,peak_time,xlabel="β",ylabel="Peak time",legend=false)

pt = peak_time.(β)
pt_surrogate = RadialBasis(β,pt,lb,ub,rad = cubicRadial)
surrogate_optimize(peak_time,
                   SRBF(),
                   lb,
                   ub,
                   pt_surrogate,
                   sampler;
                   maxiters=1000,
                   num_new_samples=1000)
pt_test = peak_time.(β_test)
pt_test_pred = pt_surrogate.(β_test);

l = @layout [a b]
pl1 = plot(β_grid,peak_time.(β_grid),color=:red,label="Model",xlabel="β",ylabel="Peak time",legend=:right)
scatter!(pl1, β_test,pt_test_pred,color=:blue,label="Surrogate")
pl2 = scatter(pt_test,pt_test_pred,color=:blue,legend=false,xlabel="Peak time",ylabel="Surrogate peak time")
Plots.abline!(pl2,1.0,0.0)
plot(pl1,pl2,layout=l)

@benchmark peak_time(0.0005)

@benchmark pt_surrogate(0.0005)


export latin, check_latin

function funky_mod(i::Int, n::Int)
    v = mod(i,n)
    return v==0 ? n : v
end

"""
`latin(n)` returns a simple `n`-by-`n` Latin square. More generally,
`latin(n,a,b)` generates a Latin square whose `i,j`-entry is
`a*(i-1) + b*(j-1) + 1` (wrapping around `n`, of course).

*Note*: If parameters `n,a,b` do not generate a legitimate Latin square
an error is thrown.
"""
function latin(n::Int, a::Int=1, b::Int=1)::Matrix{Int}
    A = zeros(Int,n,n)
    for i=1:n
        for j=1:n
            A[i,j] = funky_mod(a*(i-1) + b*(j-1) + 1, n)
        end
    end
    @assert check_latin(A) "Parameters n=$n, a=$a, and b=$b do not generate a Latin square"
    return A
end

"""
`check_latin(A)` checks if `A` is a Latin square.
(The entries must be chosen from `1:n` when `A` is an `n`-by-`n`
matrix).
"""
function check_latin(A::Matrix{Int})::Bool
    r,c = size(A)
    if r != c
        return false
    end

    vals = collect(1:r)

    # check rows
    for i=1:r
        if sort(A[:,i]) != vals
            return false
        end
    end

    # check cols
    for i=1:r
        if sort(A[i,:]) != vals
            return false
        end
    end

    return true
end


mat = [
  75 00 00 00 00 00 00 00 00 00 00 00 00 00 00;
  95 64 00 00 00 00 00 00 00 00 00 00 00 00 00;
  17 47 82 00 00 00 00 00 00 00 00 00 00 00 00;
  18 35 87 10 00 00 00 00 00 00 00 00 00 00 00;
  20 04 82 47 65 00 00 00 00 00 00 00 00 00 00;
  19 01 23 75 03 34 00 00 00 00 00 00 00 00 00;
  88 02 77 73 07 63 67 00 00 00 00 00 00 00 00;
  99 65 04 28 06 16 70 92 00 00 00 00 00 00 00;
  41 41 26 56 83 40 80 70 33 00 00 00 00 00 00;
  41 48 72 33 47 32 37 16 94 29 00 00 00 00 00;
  53 71 44 65 25 43 91 52 97 51 14 00 00 00 00;
  70 11 33 28 77 73 17 78 39 68 17 57 00 00 00;
  91 71 52 38 17 14 91 43 58 50 27 29 48 00 00;
  63 66 04 68 89 53 67 30 73 16 69 87 40 31 00;
  04 62 98 27 23 09 70 98 73 93 38 53 60 04 23
]

n, m = size(mat)
for i = 2:n
  for j = 1:i
    mat[i,j] += max(mat[i-1,j], j>1 ? mat[i-1,j-1] : 0)
  end
end
println(maximum(mat[end,:]))


using CSV, Distributions, HypothesisTests

data1 = CSV.read("../data/machine1.csv", header=false)[:,1]
data2 = CSV.read("../data/machine2.csv", header=false)[:,1]
xBar1, n1 = mean(data1), length(data1)
xBar2, n2 = mean(data2), length(data2)
sig1, sig2 = 1.2, 1.6
delta0 = 0

testStatistic = ( xBar1-xBar2 - delta0 ) / ( sqrt( sig1^2 / n1 + sig2^2 / n2 ) )
pVal = 2*ccdf(Normal(), abs(testStatistic))

println("Sample mean machine 1: ",xBar1)
println("Sample mean machine 2: ",xBar2)
println("Manually calculated test statistc: ", testStatistic)
println("Manually calculated p-value: ", pVal)


# https://www.cs.cornell.edu/gries/Logic/Axioms.html
# The axioms of calculational propositional logic C are listed in the order in
# which they are usually presented and taught. Note that equivalence comes
# first. Note also that, after the first axiom, we take advantage of
# associativity of equivalence and write sequences of equivalences without
# parentheses. We use == for equivalence, | for disjunction, & for conjunction,
# ~ for negation (not), => for implication, and <= for consequence.
#
# Associativity of ==: ((p == q) == r) == (p == (q == r))
# Symmetry of ==: p == q == q == p
# Identity of ==: true == q == q
#
# Definition of false: false == ~true
# Distributivity of not: ~(p == q) == ~p == q
# Definition of =/=: (p =/= q) == ~(p == q)
#
# Associativity of |: (p | q) & r == p | (q | r)
# Symmetry of |: p | q == q | p
# Idempotency of |: p | p == p
# Distributivity of |: p |(q == r) == p | q == p | r
# Excluded Middle: p | ~p
#
# Golden rule: p & q == p == q == p | q
#
# Implication: p => q == p | q == q
# Consequence: p <= q == q => p

# Metatheory.options[:printiter] = true

using Metatheory

calc = @theory begin
    ((p == q) == r)     ==  (p == (q == r))
    (p == q)            ==  (q == p)
    (q == q)            =>  true

    ¬(p == q)           ==  (¬(p) == q)
    (p != q)            ==  ¬(p == q)

    ((p ∨ q) ∨ r)       ==  (p ∨ (q ∨ r))
    (p ∨ q)             ==  (q ∨ p)
    (p ∨ p)             =>  p
    (p ∨ (q == r))      ==  (p ∨ q == p ∨ r)
    (p ∨ ¬(p))          =>  true

    # DeMorgan
    ¬(p ∨ q)            ==  (¬p ∧ ¬q)
    ¬(p ∧ q)            ==  (¬p ∨ ¬q)

    (p ∧ q)             ==  ((p == q) == p ∨ q)

    (p => q)            ==  ((p ∨ q) == q)
    # (p => q)            ==  (¬p ∨ q)
    # (p <= q)            =>  (q => p)
end

fold = @theory begin
    (p::Bool == q::Bool)    |>     (p == q)
    (p::Bool ∨ q::Bool)     |>     (p || q)
    (p::Bool => q::Bool)    |>     ((p || q) == q)
    (p::Bool ∧ q::Bool)     |>     (p && q)
    ¬(p::Bool)              |>     (!p)
end

# t = or_alg ∪ and_alg ∪ neg_alg ∪ demorgan ∪ and_or_distrib ∪
#     absorption ∪ calc

t = calc ∪ fold

g = EGraph(:(((¬p == p) == false)))
saturate!(g, t)
extract!(g, astsize)

@test @areequal t true ((¬p == p) == false)
@test @areequal t true ((¬p == ¬p) == true)
@test @areequal t true ((¬p ∨ ¬p) == ¬p) (¬p ∨ p) ¬(¬p ∧ p)
@test @areequal t true ((p => (p ∨ p)) == true)
params = SaturationParams(timeout=12, eclasslimit=10000, schedulerparams=(1000, 5))

@test areequal(t, true, :(((p => (p ∨ p)) == ((¬(p) ∧ q) => q)) == true); params=params)

# Frege's theorem
# params = SaturationParams(timeout=12, eclasslimit=15000, scheduler=Schedulers.ScoredScheduler)
# params = SaturationParams(timeout=12, eclasslimit=15000, schedulerparams=(500, 2))
@test_skip areequal(t, true, :((p => (q => r)) => ((p => q) => (p => r))); params=params)

# Demorgan's
@test @areequal t true (¬(p ∨ q) == (¬p ∧ ¬q))

# Consensus theorem
# @test_skip
areequal(t, :((x ∧ y) ∨ (¬x ∧ z) ∨ (y ∧ z)), :((x ∧ y) ∨ (¬x ∧ z)); params=params)


using Mods, Primes, SimpleGraphs, LinearAlgebraX

"""
`generate_points(p::Int)` generates the homogeneous coordinates of the
points in a finite projective plane of prime order `p`.
"""
function generate_points(p::Integer)::Vector{HVector}
    if ~isprime(p)
        error("Modulus $p is not prime")
    end

    n = p * p + p + 1
    A = Array{HVector}(undef, n)
    count = 1

    # case c=1
    for a = 0:p-1
        for b = 0:p-1
            x = Mod{p}.([a, b, 1])
            A[count] = HVector(x)
            count += 1
        end
    end

    # case c=0, b=1
    for a = 0:p-1
        x = Mod{p}.([a, 1, 0])
        A[count] = HVector(x)
        count += 1
    end

    A[count] = HVector(Mod{p}(1), 0, 0)

    return A
end

"""
`incidence_matrix(p::Int)` creates the point/line incidence matrix
of a finite projective plane of prime order `p`.
"""
function incidence_matrix(p::Int)::Matrix
    pts = generate_points(p)
    n = p*p + p + 1
    A = zeros(Int,n,n)
    for u = 1:n
        P = pts[u]
        for v = 1:n
            Q = pts[v]
            if dot(P,Q)==0
                A[u,v] = 1
            end
        end
    end
    return A
end

"""
`incidence_graph(p::Int)` creates the bipartite point-line
incidence graph of a finite projective plane of prime order `p`.
"""
function incidence_graph(p::Int)::SimpleGraph{Int}
    M = incidence_matrix(p)
    r,c = size(M)
    A = [ zeros(r,c)  M; copy(M')  zeros(r,c) ]
    G = SimpleGraph(A)
end

function incidence_hypergraph(p::Int)::SimpleHypergraph
    VV = generate_points(p)
    H = SimpleHypergraph{HVector{Mod{p}}}()
    for v in VV
        add!(H,v)
    end

    for v in VV
        S = [ w for w in VV if dot(v,w)==0 ]
        add!(H,Set(S))
    end

    return H
end


"""
`assemble_system!` - assemble the linear system for solving
"""
function assemble_system!(problem::FiniteDifferenceBVPProblem{TF, TBCL, TBCR}) where {TF<:AbstractArray, TBCL<:DirichletBC, TBCR<:DirichletBC}

    # build first row
    problem.A[1,1] = 2/problem.Δx^2;
    problem.A[1,2] = -1/problem.Δx^2;
    # build rows 2-n-1
    for i in 2:problem.n-1
        problem.A[i,i-1] = -1/problem.Δx^2;
        problem.A[i,i] = 2/problem.Δx^2;
        problem.A[i,i+1] = -1/problem.Δx^2;
    end
    # build row n
    problem.A[problem.n,problem.n-1] =  -1/problem.Δx^2;
    problem.A[problem.n,problem.n] =  2/problem.Δx^2;

    # build rhs
    @. problem.rhs = problem.f;
    problem.rhs[1] += problem.left_bc.u_bc/problem.Δx^2
    problem.rhs[problem.n] += problem.right_bc.u_bc/problem.Δx^2

    problem
end

function assemble_system!(problem::FiniteDifferenceBVPProblem{TF, TBCL, TBCR}) where {TF<:Function, TBCL<:DirichletBC, TBCR<:DirichletBC}

    # build first row
    problem.A[1,1] = 2/problem.Δx^2;
    problem.A[1,2] = -1/problem.Δx^2;
    # build rows 2-n-1
    for i in 2:problem.n-1
        problem.A[i,i-1] = -1/problem.Δx^2;
        problem.A[i,i] = 2/problem.Δx^2;
        problem.A[i,i+1] = -1/problem.Δx^2;
    end
    # build row n
    problem.A[problem.n,problem.n-1] =  -1/problem.Δx^2;
    problem.A[problem.n,problem.n] =  2/problem.Δx^2;

    # build rhs
    @. problem.rhs = problem.f(problem.x);
    problem.rhs[1] += problem.left_bc.u_bc/problem.Δx^2
    problem.rhs[problem.n] += problem.right_bc.u_bc/problem.Δx^2

    problem
end

function assemble_system!(problem::FiniteDifferenceBVPProblem{TF, TBCL, TBCR}) where {TF<:AbstractArray, TBCL<:NeumannBC, TBCR<:NeumannBC}

    # build first row
    problem.A[1,1] = 1/problem.Δx^2;
    problem.A[1,2] = -1/problem.Δx^2;
    # build rows 2-n+1
    for i in 2:problem.n+1
        problem.A[i,i-1] = -1/problem.Δx^2;
        problem.A[i,i] = 2/problem.Δx^2;
        problem.A[i,i+1] = -1/problem.Δx^2;
    end
    # build row n+2
    problem.A[end,end-1] =  -1/problem.Δx^2;
    problem.A[end,end] =  1/problem.Δx^2;

    # build rhs
    @. problem.rhs[2:end-1] = problem.f;
    problem.rhs[1] = -problem.left_bc.v_bc/problem.Δx
    problem.rhs[end] = problem.right_bc.v_bc/problem.Δx

    problem
end

function assemble_system!(problem::FiniteDifferenceBVPProblem{TF, TBCL, TBCR}) where {TF<:Function, TBCL<:NeumannBC, TBCR<:NeumannBC}

    # build first row
    problem.A[1,1] = 1/problem.Δx^2;
    problem.A[1,2] = -1/problem.Δx^2;
    # build rows 2-n+1
    for i in 2:problem.n+1
        problem.A[i,i-1] = -1/problem.Δx^2;
        problem.A[i,i] = 2/problem.Δx^2;
        problem.A[i,i+1] = -1/problem.Δx^2;
    end
    # build row n+2
    problem.A[end,end-1] =  -1/problem.Δx^2;
    problem.A[end,end] =  1/problem.Δx^2;

    # build rhs
    @. problem.rhs[2:end-1] = problem.f(problem.x[2:end-1]);
    # problem.rhs[1] = 0.5 * problem.f(problem.x[1])-problem.left_bc.v_bc/problem.Δx
    # problem.rhs[end] = 0.5 * problem.f(problem.x[end]) + problem.right_bc.v_bc/problem.Δx
    problem.rhs[1] = -problem.left_bc.v_bc/problem.Δx
    problem.rhs[end] =problem.right_bc.v_bc/problem.Δx

    problem
end

function assemble_system!(problem::FiniteDifferenceBVPProblem{TF, TBCL, TBCR}) where {TF<:AbstractArray, TBCL<:PeriodicBC, TBCR<:PeriodicBC}

    # build first row
    problem.A[1,1] = 2/problem.Δx^2;
    problem.A[1,2] = -1/problem.Δx^2;
    problem.A[1,end] = -1/problem.Δx^2;
    # build rows 2-n
    for i in 2:problem.n
        problem.A[i,i-1] = -1/problem.Δx^2;
        problem.A[i,i] = 2/problem.Δx^2;
        problem.A[i,i+1] = -1/problem.Δx^2;
    end
    # build row n+1
    problem.A[end,1] =  -1/problem.Δx^2;
    problem.A[end,end-1] =  -1/problem.Δx^2;
    problem.A[end,end] =  2/problem.Δx^2;

    # build rhs
    @. problem.rhs = problem.f;

    problem
end

function assemble_system!(problem::FiniteDifferenceBVPProblem{TF, TBCL, TBCR}) where {TF<:Function, TBCL<:PeriodicBC, TBCR<:PeriodicBC}

    # build first row
    problem.A[1,1] = 2/problem.Δx^2;
    problem.A[1,2] = -1/problem.Δx^2;
    problem.A[1,end] = -1/problem.Δx^2;
    # build rows 2-n
    for i in 2:problem.n
        problem.A[i,i-1] = -1/problem.Δx^2;
        problem.A[i,i] = 2/problem.Δx^2;
        problem.A[i,i+1] = -1/problem.Δx^2;
    end
    # build row n+1
    problem.A[end,1] =  -1/problem.Δx^2;
    problem.A[end,end-1] =  -1/problem.Δx^2;
    problem.A[end,end] =  2/problem.Δx^2;

    # build rhs
    @. problem.rhs = problem.f(problem.x);

    problem
end


function LUSolve(A::Array,B::Vector)

      #A = the matrix of coefficients 'A' must be squared
  #B = The Vector 'B' right hand side vector Vertical
    m,n =size(A)
    v = size(B,1)
    if m != n
        error("Matrix must be squared")
    elseif m != v
        error("Vector Length must be equal to Matrix")
    end

L,U,P,p=DolittleLU(A)
    A=P*A
    B=P*B
    X=D=zeros(m)

    #Fordward Substitution
    D[1]=B[1]
    for i=2:m
       D[i]=(B[i] - (L[i,1:i-1]'*D[1:i-1]))

    end

         #Back Substitution
    X[m]=D[end]/U[end,end]

    for i=m-1:-1:1
       X[i]=(D[i] - (U[i,i+1:m]'*X[i+1:m]))/U[i,i]
    end

    return X
end


### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# ╔═╡ e077ab16-22e5-11eb-15ed-5f0654bf7328
using Pkg, DrWatson

# ╔═╡ e75a9a06-22e5-11eb-2e36-8d4f62b830ed
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Random, Distributions, Plots, LaTeXStrings
	Random.seed!(0)
end

# ╔═╡ ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
md"## Listing5.19"

# ╔═╡ fa68607e-22e5-11eb-0558-c9a4d9f77426
begin
	alpha, beta = 8, 2
	prior(lam) = pdf(Gamma(alpha, 1/beta), lam)
	data = [2,1,0,0,1,0,2,2,5,2,4,0,3,2,5,0]

	like(lam) = *([pdf(Poisson(lam),x) for x in data]...)
	posteriorUpToK(lam) = like(lam)*prior(lam)
end

# ╔═╡ 9ed0aafc-3afb-11eb-1f46-5bc8b9a49938
begin
	delta = 10^-4.
	lamRange = 0:delta:10
	K = sum([posteriorUpToK(lam)*delta for lam in lamRange])
	posterior(lam) = posteriorUpToK(lam)/K
end

# ╔═╡ 9ed0def2-3afb-11eb-2bd0-5b8a32cba3e1
bayesEstimate = sum([lam*posterior(lam)*delta for lam in lamRange])

# ╔═╡ 9ed16ca0-3afb-11eb-332d-3ba7b882417f
newAlpha, newBeta = alpha + sum(data), beta + length(data)

# ╔═╡ 9ed86364-3afb-11eb-0bb4-a586cc54fb14
closedFormBayesEstimate = mean(Gamma(newAlpha, 1/newBeta))

# ╔═╡ 9ee4be2c-3afb-11eb-02df-0f03e0f586fc
Text("Computational Bayes Estimate: $(bayesEstimate)")

# ╔═╡ 9ee54e44-3afb-11eb-19ef-511b9b6134a4
Text("Closed form Bayes Estimate: $(closedFormBayesEstimate)")

# ╔═╡ 9ef09b16-3afb-11eb-1cc0-a7d27312610f
begin
	plot(lamRange, prior.(lamRange),
		c=:blue, label="Prior distribution")
	plot!(lamRange, posterior.(lamRange),
		c=:red, label="Posterior distribution",
		xlims=(0, 10), ylims=(0, 1.2),
			xlabel=L"\lambda",ylabel="Density")
end

# ╔═╡ 475ff888-22e6-11eb-2354-09f8f40a8e12
md"## End of listing5.19"

# ╔═╡ Cell order:
# ╟─ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
# ╠═e077ab16-22e5-11eb-15ed-5f0654bf7328
# ╠═e75a9a06-22e5-11eb-2e36-8d4f62b830ed
# ╠═fa68607e-22e5-11eb-0558-c9a4d9f77426
# ╠═9ed0aafc-3afb-11eb-1f46-5bc8b9a49938
# ╠═9ed0def2-3afb-11eb-2bd0-5b8a32cba3e1
# ╠═9ed16ca0-3afb-11eb-332d-3ba7b882417f
# ╠═9ed86364-3afb-11eb-0bb4-a586cc54fb14
# ╠═9ee4be2c-3afb-11eb-02df-0f03e0f586fc
# ╠═9ee54e44-3afb-11eb-19ef-511b9b6134a4
# ╠═9ef09b16-3afb-11eb-1cc0-a7d27312610f
# ╟─475ff888-22e6-11eb-2354-09f8f40a8e12


"""
    closeness_centrality(g, distmx=weights(g); normalize=true)

Calculate the [closeness centrality](https://en.wikipedia.org/wiki/Centrality#Closeness_centrality)
of the graph `g`. Return a vector representing the centrality calculated for each node in `g`.

### Optional Arguments
- `normalize=true`: If true, normalize the centrality value of each
node `n` by ``\\frac{|δ_n|}{|V|-1}``, where ``δ_n`` is the set of vertices reachable
from node `n`.

# Examples
```jldoctest
julia> using LightGraphs

julia> closeness_centrality(StarGraph(5))
5-element Array{Float64,1}:
 1.0
 0.5714285714285714
 0.5714285714285714
 0.5714285714285714
 0.5714285714285714

julia> closeness_centrality(PathGraph(4))
4-element Array{Float64,1}:
 0.5
 0.75
 0.75
 0.5
```
"""
function closeness_centrality(g::AbstractGraph,
    distmx::AbstractMatrix=weights(g);
    normalize=true)

    n_v = nv(g)
    closeness = zeros(n_v)

    for u in vertices(g)
        if degree(g, u) == 0     # no need to do Dijkstra here
            closeness[u] = 0.0
        else
            d = dijkstra_shortest_paths(g, u, distmx).dists
            δ = filter(x -> x != typemax(x), d)
            σ = sum(δ)
            l = length(δ) - 1
            if σ > 0
                closeness[u] = l / σ
                if normalize
                    n = l / (n_v - 1)
                    closeness[u] *= n
                end
            end
        end
    end
    return closeness
end


using CSV, Distributions, HypothesisTests

data = CSV.read("../data/machine1.csv", header=false)[:,1]
n = length(data)
mu0A, mu0B = 52.2, 53

xPositiveA = sum(data .> mu0A)
testStatisticA = max(xPositiveA, n-xPositiveA)

xPositiveB = sum(data .> mu0B)
testStatisticB = max(xPositiveB, n-xPositiveB)

binom = Binomial(n,0.5)
pValA = 2*ccdf(binom, testStatisticA)
pValB = 2*ccdf(binom, testStatisticB)

println("Binomial mean: ", mean(binom))

println("Case A: mu0: ", mu0A)
println("\tTest statistc: ", testStatisticA)
println("\tP-value: ", pValA)

println("Case B: mu0: ", mu0B)
println("\tTest statistc: ", testStatisticB)
println("\tP-value: ", pValB)


### A Pluto.jl notebook ###
# v0.17.1

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local iv = try Base.loaded_modules[Base.PkgId(Base.UUID("6e696c72-6542-2067-7265-42206c756150"), "AbstractPlutoDingetjes")].Bonds.initial_value catch; b -> missing; end
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : iv(el)
        el
    end
end

# ╔═╡ 9112b09c-5a21-11ec-0847-1971b3c9a6cd
begin
	using Pkg
	Pkg.activate("../../")
	using Revise
end

# ╔═╡ 4b66d9ba-eb36-4f53-8e87-28b6eadbd828
using FFTW, Colors, ImageShow, TestImages, Tullio, Zygote, Noise, IndexFunArrays, FourierTools, ImgProcMic, PlutoUI, DeconvOptim, PlutoTest, Optim, Statistics, Plots

# ╔═╡ 9bc5b039-5c13-46bc-8d3b-9edae93d164a
using ComponentArrays, Images, Random, MicroscopyTools

# ╔═╡ 606ad4de-213f-4aac-a0a9-9a6827c05c29
begin
	"""
		circ(size, radius)

	`size` the size of the resulting array
	and `radius` the radius of the circle
	"""
	circ(size, radius) = rr2(size) .<= radius^2
end

# ╔═╡ e7a503cf-c59c-4737-95f5-b8f38548833a
"""
	calc_psf(arr_size, radius)

Calculate a simple PSF. `arr_size` is the output size and `radius` is the radius
in Fourier space.
The output array is normalized by its sum to 1.
"""
function calc_psf(arr_size, radius)
	# TODO
	rr2 = circ(arr_size, radius)
	h = abs2.(ift(rr2))
	h ./= sum(h)
	return h
end

# ╔═╡ 3194824d-ed15-479a-af10-a1803f9f20bc
md"# 1 Lucy Richardson Deconvolution
In the last exercise we implemented the Wiener filter which is suited for Gaussian noise.
However, in microscopy the most dominant noise is usually Poisson noise.
The LR deconvolution algorithm is derived on a Poisson noise assumption.

Based on the lecture notes, implement the LR algorithm.
"

# ╔═╡ 7ef1aaf2-081c-4885-9ab3-4606cf0656b4
"""
	LR(measured, psf, N)

`measured` is the measured image. `psf` is the point spread function.
`N` is the number of iterations.
"""
function LR(measured, psf, N)
	otf = rfft(psf)
	otf_c = conj.(otf)

	rec = one.(measured)

	conv(object, o) = irfft(rfft(object) .* o, size(measured, 1))

	for i = 1:N
		rec .*= conv(measured ./ conv(rec, otf_c), otf)
	end
	return rec
end

# ╔═╡ 400aeebf-21bd-4b58-9b64-ccec3071df63
# testimage
img = Float32.(testimage("resolution_test_512"));

# ╔═╡ d7741c07-0982-4e57-bf95-a3ee0bc6f0b1
# calculate the PSF with a radius of 30
psf = ifftshift(Float32.(calc_psf(size(img), 30)));

# ╔═╡ aa15a3ed-0f97-4ef9-9704-0ede98426b6c
# our measured image, degraded with poisson
measured = 0.0001f0 .+ poisson(FourierTools.conv(img, psf), 100);

# ╔═╡ f4e2cc11-e86e-4afc-a021-aed27283a5f1
gray_show(measured)

# ╔═╡ dc9fa103-6a00-4129-b5d1-c1b36216d13d
md"
After

iterations=$(@bind iterations Slider(1:300, show_value=true))

the image should look _better_ than the original.
Take that as test!
"

# ╔═╡ 6ba1619b-036a-480d-b3bb-88ea6f5b259e
rec = LR(measured, psf, iterations);

# ╔═╡ f9bb5a74-067d-4afd-ae61-7f41f2e1dbca
gray_show(rec)

# ╔═╡ 0d7e91be-e03e-454c-9ebd-c0b68ca93848
md"# 2 (Smoothed) Total Variation Regularizer
In this part we want to implement a Total Variation Regularizer.
It's purpose is to impose a constraint of _piecewise smoothness_ to the image.
Formally, it is defined as

$\text{TV}(\text{img}) = \sum_{i,j} \sqrt{\epsilon + \left|\text{img}[i, j] - \text{img}[i+1, j]\right|^2 + \left|\text{img}[i, j] - \text{img}[i, j+1]\right|^2}$

where $\epsilon$ is a small value prevent a singularity in the derivative.

The output of this equation is a scalar value and can be understood as a loss function.
If we add this term to our loss function of the LR algorithm, the optimizer has to minimize not only the LR loss term but also this additional regularizer.
Often this can introduce nicer results.

As you learnt in the last homework, we also need the gradient of the `TV` function.
Fortunately, for that purpose we can use the automatic differentiation package [Zygote.jl](https://github.com/FluxML/Zygote.jl).
"

# ╔═╡ 72cdab83-0162-4951-8751-b91cb6fd0e49
md"For the implementation of the regularizer itself, we use the package [Tullio.jl](https://github.com/mcabbott/Tullio.jl). _Tullio.jl_ speeds up the calculation and is conventient to use.
For simple usage, see below:"

# ╔═╡ f8c665e4-e07f-4995-ad84-5a31658be5f4
begin
	# example how @tullio works
	arr = [1.0, 2.0, 3.0]
	@tullio arr_ex = abs2(arr[i] - arr[i-1])
end

# ╔═╡ bfb88c5e-0e45-428b-a66c-b9c37f633c6b
"""
	TV(img::AbstractArray{T, N}, ϵ=T(1e-6)) where {T, N}

Calculates the TV regularizer. `ϵ` is a small positive value to make it differentiable.
"""
function TV(img::AbstractArray{T, N}, ϵ=T(1e-6)) where {T, N}
	# insert tullio expression here!
	@tullio r = sqrt(ϵ + abs2(img[i, j] - img[i, j+1]) + abs2(img[i, j] - img[i + 1, j]))
end

# ╔═╡ 49143f59-2f8f-46e9-ac32-c4681164ef82
TV(img)

# ╔═╡ e33d9b1f-dc9b-49c0-a6e3-b1e2266df999
# the gradient can be calculated likes this!
# it describes how the value of the TV change for each pixel value
gradient(TV, img)[1]

# ╔═╡ 8a9da6e2-528a-41e5-9bc5-750885d94ca0
md"## 2 Test"

# ╔═╡ 53364e6f-b672-4c60-8fe1-845ba85399ba
PlutoTest.@test typeof(TV([1.0 2; 3 4])) <: Number

# ╔═╡ 30e313ee-aca2-4469-97cf-08acd489aaef
PlutoTest.@test TV([1.0 2; 3 4]) ≈ 2.2360682011065762

# ╔═╡ 4c91d0da-201a-4281-91a2-2de626ab4a30
PlutoTest.@test Zygote.gradient(TV, [1.0 2; 3 4])[1] ≈ [-1.34164   0.447214; 0.894427  0.0] rtol=1e-5

# ╔═╡ 340851e8-a66d-4643-83d1-48fb775dde8e
md"# 3 Lucy Richardson with TV
Additionally to the plain LR we want to incorporate the TV regularizer in this task.

The gradient of `TV` can be calculated with `gradient(TV, arr)[1]`.

In the lectures it is sketched how to add the regularizer loss to the normal loss function.
Copy your `LR` from above and try to incorporate the gradient of the `TV` loss value.
The `TV` gradient value is weighted with `λ`.
"

# ╔═╡ cfc01fbb-5ee4-4909-ba67-2034506e55e8
function LR_TV(measured, psf, λ, N)
	otf = rfft(psf)
	otf_c = conj.(otf)

	buffer = rfft(img)
	rec = one.(measured)

	conv(object, o) = irfft(rfft(object) .* o, size(measured, 1))

	for i = 1:N
		rec .*= (conv(measured ./ conv(rec, otf_c), otf) .- λ .* gradient(TV, rec)[1])
	end
	return rec
end

# ╔═╡ ae7c45b6-c8a7-4fe0-b922-9ab40fcce372
# that can easily take a minute. Zygote is very sloooooooow on the first call
rec2 = LR_TV(measured, psf, 0.02, 400)

# ╔═╡ 8c91142a-5b2d-4dfc-9764-87a66f90f0cb
gray_show(rec2)

# ╔═╡ dcba043b-db09-4adc-8239-a6439a9c074e
md" ## 4 Bonus
Anyone who can beat [DeconvOptim.jl](https://github.com/roflmaostc/DeconvOptim.jl) in terms of speed (with equal reconstruction quality), immediatly gets the bonus for the exam.
"

# ╔═╡ bbb7dda6-84c2-4172-a888-85a3658c1556
reg = DeconvOptim.TV(num_dims=2);

# ╔═╡ 231698fd-bb16-4895-9278-0d41dfafdb7e
rec_DO,r = deconvolution(measured, psf, regularizer=reg, iterations=20, λ=0.02);

# ╔═╡ 45f089ec-a7cb-4941-ad04-7a6c6c55dbef
r

# ╔═╡ d725c173-0123-40ed-aaa0-feec1dcd7a25
gray_show(rec_DO)

# ╔═╡ dfc9d8f3-6b16-476a-bba9-1c7f55229f57
md"# 5 Single Molecule Localization

In this part we fit the position, intensity and the standard deviation of a point source emitter to a measured image.
"

# ╔═╡ 954f06a6-ac27-4596-a978-e1274b2532d8
md"### Hint
Use the following pattern:
"

# ╔═╡ 46185370-aaf3-48c3-b297-b3a2843a0435
1:100

# ╔═╡ 7c51a261-a188-469d-8020-876b1b1dcbf0
(1:100)'

# ╔═╡ fae6889b-65f5-4b90-972f-56187b69fee5
"""
	gauss(I, σ, x_offset, y_offset)

Returns an image of size `(256, 256)` with a Gauss peak
centered at `(x_offset, y_offset)` with standard deviation `σ` and intensity
`I` (which is multiplied to the ideal Gauss)
"""
function gauss(I, σ::T, x_offset::T, y_offset::T) where T
	# TODO
	s = (256, 256)
	x = 1:s[1]
	y = (1:s[2])'
	@. I * exp(-T(0.5) * (abs2(x - x_offset) + abs2(y-y_offset)) / σ^2)
end

# ╔═╡ 23b9ea69-99d5-4a49-8e36-cd25323c8858
md"## 5.1 Test Gauss"

# ╔═╡ e9c17987-029e-4bd9-824e-47175f4d6924
z = zeros((256, 256));

# ╔═╡ 046ec177-9394-4b7d-bead-68590e169f1b
z[10, 10] = 42.0;

# ╔═╡ eb9b2246-c9d9-442e-bde4-48ed14bb3fa6
PlutoTest.@test gauss(42.0, 0.000001, 10.0, 10.0) ≈ z

# ╔═╡ 9cdf2e65-af36-4421-a15b-fdc4c87569b4
PlutoTest.@test gauss(11.0, 100000000000000000000.0, 0.0, 0.0) ≈ 11.0 .* ones((256, 256))

# ╔═╡ 7ecb1d8d-334c-4043-9378-adcbe0b0bf35
PlutoTest.@test (gauss(10.0, 2.0, 3.0, 3.0))[1:5, 1:5] ≈ [3.6787944117144233 5.352614285189903 6.065306597126334 5.352614285189903 3.6787944117144233; 5.352614285189903 7.788007830714049 8.824969025845954 7.788007830714049 5.352614285189903; 6.065306597126334 8.824969025845954 10.0 8.824969025845954 6.065306597126334; 5.352614285189903 7.788007830714049 8.824969025845954 7.788007830714049 5.352614285189903; 3.6787944117144233 5.352614285189903 6.065306597126334 5.352614285189903 3.6787944117144233]

# ╔═╡ 9b5b3495-0d3b-4f7c-9f51-04245ebd59c2
md"## 5.2 Forward Model
In the next part, we assemble a forward model mapping a set of parameters to an image.
"

# ╔═╡ 26950064-0127-41fd-9973-6c170bd6b8fe
"""
	draw_parameters(N)

Randomly draws `N` parameters for the Gauss peak.
"""
function draw_parameters(N)
	# make the random numbers reproducible
	Random.seed!(40)
	σ = [rand(1:0.01:5.0) for i = 1:N]
	x_offset = [rand(10:0.1:240) for i = 1:N]
	y_offset = [rand(10:0.1:240) for i = 1:N]
	I = [rand(0.5:0.1:1.0) for i = 1:N]
	return ComponentVector(σ=σ, x_offset=x_offset, y_offset=y_offset, I=I)
end

# ╔═╡ ef9c0737-ffbc-498b-80d2-2e93558c508f
# draws randomly a few parameters
# the output is a ComponentVector which is needed for Optim
params = draw_parameters(5);

# ╔═╡ 70a9e03f-d44f-4cce-a82d-898e63f789fd
# all fields can be access with params.FIELD, see below

# ╔═╡ 70453553-7dbc-46f1-a141-f4d19f2ae3d4
params.σ

# ╔═╡ 6ac55320-b94f-44c2-9669-1d4e7eb3ddee
params.x_offset

# ╔═╡ 73416e0d-ed2b-4b95-aa43-ce58de7f29aa
"""
	forward(params)

Maps the `params` to an output image.
"""
function forward(params)
	sum(gauss.(params.I, params.σ, params.x_offset, params.y_offset))
end

# ╔═╡ 56d566e3-632c-4852-839b-adf505122ad2
measurement = randn() * 0.3 .+ poisson(forward(params), 50);

# ╔═╡ 26f1a0c2-3787-4335-958b-12719e2901bf
# noisy measurement of the point sources
gray_show(measurement)

# ╔═╡ 2b3c151a-d7bd-40e4-97d7-f121c9da23a6
md"## 5.3 Loss Function"

# ╔═╡ 32bd26f1-0d04-4589-a2d5-1e01abce4a7d
"""
	loss(measurement, pred)

Compares the prediction `pred` and the `measurement` under the L2 norm.
"""
function loss(measurement, pred)
	# TODO
	sum(abs2.(measurement .- pred))
end

# ╔═╡ dc66401d-2581-4d6f-8b36-d15a2d21391a
md"## 5.4 Reconstruction"

# ╔═╡ daa77619-e7a8-4c6a-9854-d8ef21042277
"""
	find_peaks(arr, bin_factor)

Finds the peaks in the measurement `arr`.
To get a better SND we bin the measurement by a binning factor `bin_factor`.
"""
function find_peaks(arr, bin_factor)
	arr_s = bin(arr, (bin_factor,bin_factor))
	map(x -> x .* bin_factor, Tuple.(findlocalmaxima(arr_s)))
end

# ╔═╡ 1a505ff5-4745-42c5-bc40-bc1075fcd4d6
"""
	init_parameters(N, measurement)

Initializes a good set of parameters for the reconstruction.
`N` is the number of PSFs and `measurement` the measured image.

The critical part is to find good initial positions for the peaks.
For that we use `find_peaks`.

"""
function init_params(N, measurement, bin_factor, σ_guess)
	# find a σ_guess which results in a good estimate
	σ = [σ_guess for i = 1:N]

	peaks = find_peaks(measurement, bin_factor)
	x_offset = [peaks[i][1] for i =1:N]
	y_offset = [peaks[i][2] for i =1:N]

	# you can change that as well! Is it critical?
	I = [1.0 for i = 1:N]

	return ComponentVector(σ=σ, x_offset=x_offset, y_offset=y_offset, I=I)
end

# ╔═╡ 74f8a126-7e5c-4e1d-90bd-fe66065d5faf
begin
	# initial guess
	# try to find an binning factor which results in good peaks
	bin_factor = 5 # TODO
	σ_guess = 5.30 # TODO
	params_guess = init_params(5, measurement, bin_factor, σ_guess);
end

# ╔═╡ 3c38cdec-99d8-4c11-bedb-87137ea119ba
md"
The left image is the measured image.
The center image the initial guess.
The right image is the difference map suited to judge the position of the peaks.
"

# ╔═╡ 0a3cfe9b-888c-420a-9465-bef61d058c17
gray_show([measurement forward(params_guess) abs.(measurement .- forward(params_guess))])

# ╔═╡ 22d4842a-908f-485d-9bf4-96c053b106c6
md"### 5.5 Optim Specific Parts for the Reconstruction"

# ╔═╡ c3996433-e27d-4227-8e18-9531a30ea389
begin
	# the function maps the params to a single scalar value
	function f(params)
		loss(measurement, forward(params))
	end
	# this function calculates the
	g!(G, params) = isnothing(G) ? G : G .= gradient(f, params)[1]
end

# ╔═╡ d30c5e75-8a58-4231-8164-c7d783b9bc02
res = optimize(f, g!, params_guess, LBFGS(), Optim.Options(iterations=50))

# ╔═╡ 90b98a76-1bc7-43e0-927c-1cf94d35ffd0
output_optim = forward(res.minimizer);

# ╔═╡ b5880735-c615-440c-a0a0-56a6a024494a
gray_show([measurement output_optim measurement .- output_optim])

# ╔═╡ efaf6cfb-bf37-4cc2-be30-abbc6706e905
heatmap(measurement .- output_optim, title="Difference map")

# ╔═╡ 19a1ff11-1276-46ab-b84b-a66898522068
md"## Check if the values look plausible"

# ╔═╡ 060ef38d-312b-4487-b59d-0ba288b16daf
md"The operator `|>` in an expression like `f(x) |> g` is the same as `g(f(x))`. It's only a convenvient wrapper"

# ╔═╡ e60c3150-7658-4a3d-819f-b1e6293a27fb
params.σ |> sort

# ╔═╡ 17ed7c8d-4d12-4008-8745-e96f36ce85b1
res.minimizer.σ |> sort

# ╔═╡ e677b9f8-38ee-41cd-b319-ef1cc48b17e2
params.I |> sort

# ╔═╡ 854f732f-441b-45b1-9a00-2b7ae854dc47
res.minimizer.I |> sort

# ╔═╡ 45abec82-b841-4d8a-b2ae-bba738352da3
res.minimizer.y_offset |> sort

# ╔═╡ ff1c6ed2-e4b1-4f6a-87df-44f350bb50f7
params.y_offset |> sort

# ╔═╡ 8ffb693f-f5ca-45a5-8d1b-5c96885118d0
PlutoTest.@test (res.minimizer.x_offset |> sort) ≈ (params.x_offset |> sort) rtol=0.01

# ╔═╡ 67db2159-c2f0-4130-8e82-6b75907db82f
PlutoTest.@test (res.minimizer.y_offset |> sort) ≈ (params.y_offset |> sort) rtol=0.01

# ╔═╡ Cell order:
# ╠═9112b09c-5a21-11ec-0847-1971b3c9a6cd
# ╠═4b66d9ba-eb36-4f53-8e87-28b6eadbd828
# ╟─606ad4de-213f-4aac-a0a9-9a6827c05c29
# ╟─e7a503cf-c59c-4737-95f5-b8f38548833a
# ╟─3194824d-ed15-479a-af10-a1803f9f20bc
# ╠═7ef1aaf2-081c-4885-9ab3-4606cf0656b4
# ╠═400aeebf-21bd-4b58-9b64-ccec3071df63
# ╠═d7741c07-0982-4e57-bf95-a3ee0bc6f0b1
# ╠═aa15a3ed-0f97-4ef9-9704-0ede98426b6c
# ╠═f4e2cc11-e86e-4afc-a021-aed27283a5f1
# ╟─dc9fa103-6a00-4129-b5d1-c1b36216d13d
# ╠═6ba1619b-036a-480d-b3bb-88ea6f5b259e
# ╠═f9bb5a74-067d-4afd-ae61-7f41f2e1dbca
# ╟─0d7e91be-e03e-454c-9ebd-c0b68ca93848
# ╟─72cdab83-0162-4951-8751-b91cb6fd0e49
# ╠═f8c665e4-e07f-4995-ad84-5a31658be5f4
# ╠═bfb88c5e-0e45-428b-a66c-b9c37f633c6b
# ╠═49143f59-2f8f-46e9-ac32-c4681164ef82
# ╠═e33d9b1f-dc9b-49c0-a6e3-b1e2266df999
# ╟─8a9da6e2-528a-41e5-9bc5-750885d94ca0
# ╠═53364e6f-b672-4c60-8fe1-845ba85399ba
# ╠═30e313ee-aca2-4469-97cf-08acd489aaef
# ╠═4c91d0da-201a-4281-91a2-2de626ab4a30
# ╟─340851e8-a66d-4643-83d1-48fb775dde8e
# ╠═cfc01fbb-5ee4-4909-ba67-2034506e55e8
# ╠═ae7c45b6-c8a7-4fe0-b922-9ab40fcce372
# ╠═8c91142a-5b2d-4dfc-9764-87a66f90f0cb
# ╟─dcba043b-db09-4adc-8239-a6439a9c074e
# ╠═bbb7dda6-84c2-4172-a888-85a3658c1556
# ╠═231698fd-bb16-4895-9278-0d41dfafdb7e
# ╠═45f089ec-a7cb-4941-ad04-7a6c6c55dbef
# ╠═d725c173-0123-40ed-aaa0-feec1dcd7a25
# ╟─dfc9d8f3-6b16-476a-bba9-1c7f55229f57
# ╠═9bc5b039-5c13-46bc-8d3b-9edae93d164a
# ╟─954f06a6-ac27-4596-a978-e1274b2532d8
# ╠═46185370-aaf3-48c3-b297-b3a2843a0435
# ╠═7c51a261-a188-469d-8020-876b1b1dcbf0
# ╠═fae6889b-65f5-4b90-972f-56187b69fee5
# ╟─23b9ea69-99d5-4a49-8e36-cd25323c8858
# ╠═e9c17987-029e-4bd9-824e-47175f4d6924
# ╠═046ec177-9394-4b7d-bead-68590e169f1b
# ╠═eb9b2246-c9d9-442e-bde4-48ed14bb3fa6
# ╠═9cdf2e65-af36-4421-a15b-fdc4c87569b4
# ╠═7ecb1d8d-334c-4043-9378-adcbe0b0bf35
# ╟─9b5b3495-0d3b-4f7c-9f51-04245ebd59c2
# ╠═26950064-0127-41fd-9973-6c170bd6b8fe
# ╠═ef9c0737-ffbc-498b-80d2-2e93558c508f
# ╠═70a9e03f-d44f-4cce-a82d-898e63f789fd
# ╠═70453553-7dbc-46f1-a141-f4d19f2ae3d4
# ╠═6ac55320-b94f-44c2-9669-1d4e7eb3ddee
# ╠═73416e0d-ed2b-4b95-aa43-ce58de7f29aa
# ╠═56d566e3-632c-4852-839b-adf505122ad2
# ╠═26f1a0c2-3787-4335-958b-12719e2901bf
# ╟─2b3c151a-d7bd-40e4-97d7-f121c9da23a6
# ╠═32bd26f1-0d04-4589-a2d5-1e01abce4a7d
# ╟─dc66401d-2581-4d6f-8b36-d15a2d21391a
# ╠═daa77619-e7a8-4c6a-9854-d8ef21042277
# ╠═1a505ff5-4745-42c5-bc40-bc1075fcd4d6
# ╠═74f8a126-7e5c-4e1d-90bd-fe66065d5faf
# ╟─3c38cdec-99d8-4c11-bedb-87137ea119ba
# ╠═0a3cfe9b-888c-420a-9465-bef61d058c17
# ╟─22d4842a-908f-485d-9bf4-96c053b106c6
# ╠═c3996433-e27d-4227-8e18-9531a30ea389
# ╠═d30c5e75-8a58-4231-8164-c7d783b9bc02
# ╠═90b98a76-1bc7-43e0-927c-1cf94d35ffd0
# ╠═b5880735-c615-440c-a0a0-56a6a024494a
# ╠═efaf6cfb-bf37-4cc2-be30-abbc6706e905
# ╟─19a1ff11-1276-46ab-b84b-a66898522068
# ╟─060ef38d-312b-4487-b59d-0ba288b16daf
# ╠═e60c3150-7658-4a3d-819f-b1e6293a27fb
# ╠═17ed7c8d-4d12-4008-8745-e96f36ce85b1
# ╠═e677b9f8-38ee-41cd-b319-ef1cc48b17e2
# ╠═854f732f-441b-45b1-9a00-2b7ae854dc47
# ╠═45abec82-b841-4d8a-b2ae-bba738352da3
# ╠═ff1c6ed2-e4b1-4f6a-87df-44f350bb50f7
# ╠═8ffb693f-f5ca-45a5-8d1b-5c96885118d0
# ╠═67db2159-c2f0-4130-8e82-6b75907db82f


#############################
##  Numerical experiments associated with section 5.4.2
#############################
## In general, VIDEs are implemented as a combination of the approaches for nonlinear and integro-differential equations.
## Whether this is more efficient than standard methods once again depends on the sparsity of the resulting operator
## and the polynomial order required to resolve the solution.
##

using NLsolve, ApproxFun, MultivariateOrthogonalPolynomials, BandedMatrices, BlockBandedMatrices, Plots
using SparseVolterraExamples

#####################################################
## Problem in Equation (28), u_2(x) = tan(x)
#####################################################

####
## Some initialization
DD = 1                      # Highest derivative order of the problem
    f(u) = u^2              # This is the nonlinearity
    lincomb(x) = sin(x)     # In lincomb we store potentially appearing additional linear terms. In this case this is sin(x) as seen in the example.
    g(x) = x+cos(x)-tan(x)+tan(x)^2   # the g function from the problem statement
    K(x,y) = 1                        # the kernel is constant in this example
    conditions = [0]                  # Vector of length DD containing the conditions that uniquely solve the differential equation
    n = 28                            # Set some desired polynomial degree of approximation.
####
## Now we can set up the appropriate objective function
objective(x) = triNonLinearIntegroDiffVolterraObjective(x,f,K,g,n,DD,conditions,lincomb)

####
## Solve using basic Newton iteration with a zero vector guess
sol = nlsolve(objective,zeros(n),method=:newton,ftol=1e-15)

####
## We can compare the numerical solution against the analytic one.
plot(Fun(Jacobi(0,1,0..1),sol.zero),label="sparse method solution",grid=false,legend=:topleft,ylabel="u(x)",xlabel="x")
plot!(x->tan(x),0,1,label="analytic tan(x)",grid=false,legend=:topleft,ylabel="u(x)",xlabel="x")

####
## Or just plot the error on (0,1). A straightforward loop over the polynomial order and a maximum function over this difference then leads to Figure 8(b).
plot(Fun(Jacobi(0,1,0..1),sol.zero)-Fun(x->tan(x),0..1),label=false,grid=false,legend=:topleft,ylabel="error",xlabel="x")

examples/Section 5-4-2 - Nonlinear VIDEs.jl


# Quadrilateral in 3D defined by its 4 vertices.

# NOTE: Quadrilaterals are assumed to be convex and planar (all points in some plane, not
# necessarily xy, yx, etc.)!
# Quadrilaterals must be convex to be valid finite elements. See link below
# https://math.stackexchange.com/questions/2430691/jacobian-determinant-for-bi-linear-quadrilaterals
struct Quadrilateral_3D{T <: AbstractFloat}
    # Counter clockwise order
    points::NTuple{4, Point_3D{T}}
end

# Constructors
# -------------------------------------------------------------------------------------------------
Quadrilateral_3D(p₁::Point_3D{T},
                 p₂::Point_3D{T},
                 p₃::Point_3D{T},
                 p₄::Point_3D{T}) where {T <: AbstractFloat} = Quadrilateral_3D((p₁, p₂, p₃, p₄))

# Methods
# -------------------------------------------------------------------------------------------------
# Interpolation
function (quad::Quadrilateral_3D{T})(r::R, s::S) where {T <: AbstractFloat,
                                                        R <: Real,
                                                        S <: Real}
    # See The Visualization Toolkit: An Object-Oriented Approach to 3D Graphics, 4th Edition
    # Chapter 8, Advanced Data Representation, in the interpolation functions section
    rₜ = T(r)
    sₜ = T(s)
    return (1 - rₜ)*(1 - sₜ)*quad.points[1] +
                 rₜ*(1 - sₜ)*quad.points[2] +
                       rₜ*sₜ*quad.points[3] +
                 (1 - rₜ)*sₜ*quad.points[4]
end

function triangulate(quad::Quadrilateral_3D{T}) where {T <: AbstractFloat}
    # Return the two triangles that partition the domain
    A, B, C, D = quad.points
    return (Triangle_3D(A, B, C), Triangle_3D(C, D, A))
end

function area(quad::Quadrilateral_3D{T}) where {T <: AbstractFloat}
    # Using the convex quadrilateral assumption, just return the sum of the areas of the two
    # triangles that partition the quadrilateral. If the convex assumption ever changes, you
    # need to verify that the triangle pair partitions the quadrilateral. Choosing the wrong
    # pair overestimates the area, so just get the areas of both pairs of valid triangles and use
    # the smaller area.
    return sum(area.(triangulate(quad)))
end

#function intersect(l::LineSegment_3D{T}, quad::Quadrilateral_3D{T}) where {T <: AbstractFloat}
#    # Triangulate the quadrilateral, intersect the triangles.
#    tri = triangulate(quad)
#    intersection1 = l ∩ tri[1]
#    if intersection1[1]
#        return true, intersection1[2]
#    end
#    intersection2 = l ∩ tri[2]
#    if intersection2[1]
#        return true, intersection2[2]
#    end
#    return false, Point_3D(T, 0)
#end


module HaltonSequences

using Primes
export haltonvalue,Halton,HaltonPoint

const MAXLEN = typemax(UInt32)

"""
    haltonvalue(n, base[, T])

Return the `n`th element of the Halton sequence with prime base `base` (defaults to 2). If
a `<:Real` type `T` is given, compute the value in that type.

# Examples
```jldoctest
julia> haltonvalue(4,3,Rational)
4//9
```
"""
function haltonvalue(n::Integer,base::Integer=2,T::Type=float(Int))
    haltonvalue!(zero(T),n,base)
end

function haltonvalue!(x::Real,n::Integer,base::Integer=2)
    x = zero(x)
    s = one(x)
    while n > 0
        s /= base
        n,r = divrem(n,base)
        x += s*r
    end
    return x
end

const val = haltonvalue  # for terser qualified uses

########################################################
# Vector of Halton sequence values
########################################################
"""
    Halton(base=2; <keyword arguments>)
    Halton{T<:Real}(base=2; <keyword arguments>)

Create a `Vector`-like object for a Halton sequence. Optionally, use a `<:Real` type
`T` for the elements.

# Arguments
- `base::Integer`: prime number base for the sequence
- `start=1` (keyword): starting/offset index
- `length=typemax(UInt32)` (keyword): length of the sequence vector

# Examples
```jldoctest
julia> Halton(3,length=4)
4-element Halton{Float64}:
 0.3333333333333333
 0.6666666666666666
 0.1111111111111111
 0.4444444444444444
```
```jldoctest
julia> Halton{Rational}(2,length=7)
7-element Halton{Rational}:
 1//2
 1//4
 3//4
 1//8
 5//8
 3//8
 7//8
```
```jldoctest
julia> sum( Halton()[1:10000] )/10000
0.4998324768066406
```

 See also: [`HaltonPoint`](@ref), [`haltonvalue`](@ref)
"""
struct Halton{T<:Real} <: AbstractArray{T, 1}
    base::Int
    start::Int
    length::Int
    # Constructor for when the type is explicit.
    function Halton{T}(base::Integer=2;start=1,length=MAXLEN) where {T<:Real}
        new(base,start,length)
    end
end

# Constructor when no type given. Also includes a check for prime base.
function Halton(base::Integer=2;kw...)
    isprime(base) || error("Base must be a prime number.")
    Halton{float(Int)}(base;kw...)
end

# Serve the AbstractArray interface
Base.size(H::Halton) = (H.length,)
Base.IndexStyle(::Type{<:Halton}) = IndexLinear()
function Base.getindex(H::Halton{T}, n::Integer) where T
    n <= H.length || throw(BoundsError(H,n))
    haltonvalue(n+H.start-1,H.base,T)
end

########################################################
# Vector of multidimensional points based on different bases
########################################################
"""
    HaltonPoint(ndim::Integer; <keyword arguments>)
    HaltonPoint(base::Vector{<:Integer}; <keyword arguments>)
    HaltonPoint{T<:Real}(base::Vector{<:Integer}; <keyword arguments>)

Create a `Vector`-like object for a sequence of `ndim`-dimensional vectors constructed
from Halton sequences. If given, `base` should be a vector of primes and then
`ndim=length(base)`. If `T` is given, it is the element type for the vectors.

# Keyword arguments
- `start=1`: starting/offset index
- `length=typemax(UInt32)`: length of the sequence vector

# Examples
```jldoctest
julia> HaltonPoint(4,length=5)
5-element HaltonPoint{Float64}:
 [0.5, 0.333333, 0.2, 0.142857]
 [0.25, 0.666667, 0.4, 0.285714]
 [0.75, 0.111111, 0.6, 0.428571]
 [0.125, 0.444444, 0.8, 0.571429]
 [0.625, 0.777778, 0.04, 0.714286]
```
```jldoctest
julia> h = HaltonPoint([5,11,17]);  sum(h[10001:20000])/10000
3-element Array{Float64,1}:
 0.49996162048000015
 0.5000650725546567
 0.49992980926952507
```

 See also: [`Halton`](@ref), [`haltonvalue`](@ref)
"""
struct HaltonPoint{T<:Real} <: AbstractArray{T, 1}
    base::Vector{Int}
    start::Int
    length::Int
    # Constructor when the type and base are explicit.
    function HaltonPoint{T}(base::Vector{S};start=1,length=MAXLEN) where {T<:Real} where {S<:Integer}
        new(base,start,length)
    end
    # Constructor when the type is explicit, but bases are not given.
    function HaltonPoint{T}(ndim=Integer;kw...) where {T<:Real}
        function firstprimes(k::Integer)
            n = zeros(UInt16,k)
            n[1] = 2
            for i = 2:k
                n[i] = nextprime(n[i-1]+1)
            end
            return n
        end
        HaltonPoint{T}(firstprimes(ndim);kw...)
    end
end

# Constructor when the type is not given.
HaltonPoint(b_or_n;kw...) = HaltonPoint{float(Int)}(b_or_n;kw...)

# Serve the AbstractArray interface
Base.size(P::HaltonPoint) = (P.length,)
Base.IndexStyle(::Type{<:HaltonPoint}) = IndexLinear()
Base.eltype(::HaltonPoint{T}) where T = Vector{T}
function Base.getindex(P::HaltonPoint{T}, n::Integer) where T
    [ haltonvalue(n+P.start-1,P.base[i],T) for i in eachindex(P.base) ]
end
# This seems to be needed for range/vector indexing, since the eltype is not scalar.
function Base.getindex(P::HaltonPoint{T}, n::Union{AbstractVector,AbstractRange}) where T
    [ P[n[i]] for i in eachindex(n) ]
end

end # module


using OrdinaryDiffEq
using DiffEqSensitivity
using Zygote
using Plots

function sir_ode!(du,u,p,t)
    (S,I,R) = u
    (β,c,γ) = p
    N = S+I+R
    @inbounds begin
        du[1] = -β*c*I/N*S
        du[2] = β*c*I/N*S - γ*I
        du[3] = γ*I
    end
    nothing
end;

δt = 1.0
tmax = 40.0
tspan = (0.0,tmax)
t = 0.0:δt:tmax
num_timepoints = length(t);

u0 = [990.0,10.0,0.0] # S,I,R
num_states = length(u0);

p = [0.05,10.0,0.25]; # β,c,γ
num_params = length(p);

prob_ode = ODEProblem(sir_ode!,u0,tspan,p);

sim_ode = (u0,p)-> solve(prob_ode,Tsit5(),u0=u0,p=p,saveat=t,sensealg=QuadratureAdjoint());

sol_ode = sim_ode(u0,p);

du0,dp = Zygote.jacobian(sim_ode,u0,p);

dβ = reshape(dp[:,1],(num_states,:))' # as β is the first parameter
dc = reshape(dp[:,2],(num_states,:))' # c is 2nd parameter
dγ = reshape(dp[:,3],(num_states,:))' # γ is 3rd parameter
dI₀ = reshape(du0[:,2],(num_states,:))'; # I₀ is the 2nd initial condition

plot(sol_ode.t,
     Array(sol_ode(t))',
     labels = ["S" "I" "R"],
     xlabel = "Time",
     ylabel = "Number")

l = @layout [a b; c d]
pl1 = plot(t,dβ,xlabel="Time",ylabel="dp",label=["S" "I" "R"],title="Sensitivity to β")
pl2 = plot(t,dc,xlabel="Time",ylabel="dp",label=["S" "I" "R"],title="Sensitivity to c")
pl3 = plot(t,dγ,xlabel="Time",ylabel="dp",label=["S" "I" "R"],title="Sensitivity to γ")
pl4 = plot(t,dI₀,xlabel="Time",ylabel="dp",label=["S" "I"  "R"],title="Sensitivity to I₀")
plot(pl1,pl2,pl3,pl4,layout=l)


"""
The prime factors of 13195 are 5, 7, 13 and 29.

What is the largest prime factor of the number 600851475143 ?
"""
# Running isprime will have to test prime each time it is called
# It might be better to generate a prime prime_sieve till the sqrt and check from reverse if the number is divisible.

function prime_sieve(n::Int)
    isprime = trues(n)
    isprime[1] = false

    for i in 2:isqrt(n)
        if isprime[i]
            for j in i^2:i:n
                isprime[j] = false
            end
        end
    end
    return filter(x -> isprime[x],1:n)
end

function largest_prime_factor(n)
    primes = prime_sieve(Int(trunc(√(n))))
    for prime in Iterators.reverse(primes)
        if n%prime == 0
            return prime
        end
    end
    return -1
end

@assert largest_prime_factor(13195) == 29
@show largest_prime_factor(600851475143)


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

"""
    EigenAnalysis(proj; maxdim=nothing, pratio=1.0)

The eigenanalysis of the covariance with a given projection `proj`.
Optionally specify the maximum number of dimensions in the output `maxdim`
and the percentage of variance to retain `pratio`. Default to all dimensions of
the input.

## Projections

* `:V` - Uncorrelated variables (PCA transform)
* `:VD` - Uncorrelated variables and variance one (DRS transform)
* `:VDV` - Uncorrelated variables and variance one (SDS transformation)

The `:V` projection used in the PCA transform projects the data on the eigenvectors
V of the covariance matrix.

The `:VD` projection used in the DRS transform. Similar to the `:V` projection,
but the eigenvectors are multiplied by the squared inverse of the eigenvalues D.

The `:VDV` projection used in the SDS transform. Similar to the `:VD` transform,
but the data is projected back to the basis of the original variables using the Vᵀ matrix.

See [https://geostatisticslessons.com/lessons/sphereingmaf](https://geostatisticslessons.com/lessons/sphereingmaf)
for more details about these three variants of eigenanalysis.

# Examples

```julia
EigenAnalysis(:V)
EigenAnalysis(:VD)
EigenAnalysis(:VDV)
EigenAnalysis(:V, maxdim=3)
EigenAnalysis(:VD, pratio=0.99)
EigenAnalysis(:VDV, maxdim=3, pratio=0.99)
```
"""
struct EigenAnalysis <: Transform
  proj::Symbol
  maxdim::Union{Int,Nothing}
  pratio::Float64

  function EigenAnalysis(proj, maxdim, pratio)
    @assert proj ∈ (:V, :VD, :VDV) "Invalid projection."
    @assert 0 ≤ pratio ≤ 1 "Invalid pratio."
    new(proj, maxdim, pratio)
  end
end

EigenAnalysis(proj; maxdim=nothing, pratio=1.0) =
  EigenAnalysis(proj, maxdim, pratio)

assertions(::Type{EigenAnalysis}) = [assert_continuous]

isrevertible(::Type{EigenAnalysis}) = true

function apply(transform::EigenAnalysis, table)
  # basic checks
  for assertion in assertions(transform)
    assertion(table)
  end

  # original columns names
  cols = Tables.columns(table)
  onames = Tables.columnnames(cols)

  # table as matrix
  X = Tables.matrix(table)

  # center the data
  μ = mean(X, dims=1)
  Y = X .- μ

  # eigenanalysis of covariance
  S, S⁻¹ = eigenmatrices(transform, Y)

  # project the data
  Z = Y * S

  # column names
  names = Symbol.(:PC, 1:size(Z, 2))

  # table with transformed columns
  𝒯 = (; zip(names, eachcol(Z))...)
  newtable = 𝒯 |> Tables.materializer(table)

  newtable, (μ, S, S⁻¹, onames)
end

function revert(::EigenAnalysis, newtable, cache)
  # table as matrix
  Z = Tables.matrix(newtable)

  # retrieve cache
  μ, S, S⁻¹, onames = cache

  # undo projection
  Y = Z * S⁻¹

  # undo centering
  X = Y .+ μ

  # table with original columns
  𝒯 = (; zip(onames, eachcol(X))...)
  𝒯 |> Tables.materializer(newtable)
end

function reapply(transform::EigenAnalysis, table, cache)
  # basic checks
  for assertion in assertions(transform)
    assertion(table)
  end

  # table as matrix
  X = Tables.matrix(table)

  # retrieve cache
  μ, S, S⁻¹, onames = cache

  # center the data
  Y = X .- μ

  # project the data
  Z = Y * S

  # column names
  names = Symbol.(:PC, 1:size(Z, 2))

  # table with transformed columns
  𝒯 = (; zip(names, eachcol(Z))...)
  𝒯 |> Tables.materializer(table)
end

_maxdim(maxdim::Int, Y) = maxdim
_maxdim(::Nothing, Y) = size(Y, 2)

function outdim(transform, Y, λ)
  pratio = transform.pratio
  csums  = cumsum(λ)
  ratios = csums ./ last(csums)
  mdim   = _maxdim(transform.maxdim, Y)
  pdim   = findfirst(≥(pratio), ratios)
  min(mdim, pdim)
end

function eigenmatrices(transform, Y)
  proj = transform.proj

  Σ = cov(Y)
  λ, V = eigen(Σ, sortby=λ -> -real(λ))

  if proj == :V
    S   = V
    S⁻¹ = transpose(V)
  elseif proj == :VD
    Λ   = Diagonal(sqrt.(λ))
    S   = V * inv(Λ)
    S⁻¹ = Λ * transpose(V)
  elseif proj == :VDV
    Λ   = Diagonal(sqrt.(λ))
    S   = V * inv(Λ) * transpose(V)
    S⁻¹ = V * Λ * transpose(V)
  end

  d = outdim(transform, Y, λ)

  S[:, 1:d], S⁻¹[1:d, :]
end

"""
    PCA(; maxdim=nothing, pratio=1.0)

The PCA transform is a shortcut for
`ZScore() → EigenAnalysis(:V; maxdim, pratio)`.

See also: [`ZScore`](@ref), [`EigenAnalysis`](@ref).

# Examples

```julia
PCA(maxdim=2)
PCA(pratio=0.86)
PCA(maxdim=2, pratio=0.86)
```
"""
PCA(; maxdim=nothing, pratio=1.0) =
  ZScore() → EigenAnalysis(:V, maxdim, pratio)

"""
    DRS(; maxdim=nothing, pratio=1.0)

The DRS transform is a shortcut for
`ZScore() → EigenAnalysis(:VD; maxdim, pratio)`.

See also: [`ZScore`](@ref), [`EigenAnalysis`](@ref).

# Examples

```julia
DRS(maxdim=3)
DRS(pratio=0.87)
DRS(maxdim=3, pratio=0.87)
```
"""
DRS(; maxdim=nothing, pratio=1.0) =
  ZScore() → EigenAnalysis(:VD, maxdim, pratio)

"""
    SDS(; maxdim=nothing, pratio=1.0)

The SDS transform is a shortcut for
`ZScore() → EigenAnalysis(:VDV; maxdim, pratio)`.

See also: [`ZScore`](@ref), [`EigenAnalysis`](@ref).

# Examples

```julia
SDS()
SDS(maxdim=4)
SDS(pratio=0.88)
SDS(maxdim=4, pratio=0.88)
```
"""
SDS(; maxdim=nothing, pratio=1.0) =
  ZScore() → EigenAnalysis(:VDV, maxdim, pratio)


#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# Program: pi_monte_carlo.jl
#          Monte-Carlo calculation of PI
#
# Usage: julia pi_monte_carlo.jl
#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
function montepi(n::Int)
   R = 1.0
   s = 0
   for i = 1: n
      x = R * rand()
      y = R * rand()
      if x^2 + y^2 <= R^2
         s = s + 1
      end
   end
   return 4.0*s/n
end

# Main program
for i in 3: 8
    n = 10^i
    p = montepi(n)
    println("N = $n: PI = $p")
end


module ch1

import StatsBase; sb = StatsBase

# Page 4
count = 600_000
sample = [1:6...]

x = sb.sample(sample, count, replace=true)
P = sum(x.==3)/count

@show P

# Page 5
wins = zeros(Int32, 1, count)
wins .= 0
for ind = 1:length(wins)
    d = sb.sample(sample, 2)
    if sum(d) == 7 || sum(d) == 11
        wins[ind] = 1
    end
end
wins_count = sum(wins)

@show wins_count
@show wins_count/count

# Page 7
# PDF: Probability Density Function
using Plotly
y = [0,1]
p = [1,1]
line = scatter(x=y, y=p, mode="lines") # "lines+markers"
layout = Layout(
    title="PDF",
    # mirror for the bounding box
    xaxis=attr(title="y", linecolor="black", mirror=true),
    yaxis=attr(title="p(y)", range=[0,1.1], linecolor="black", mirror=true)
)
pt = plot([line], layout)
#display(pt)

# Page 10
import RDatasets; rdata = RDatasets
import CSV
import DataFrames; DF = DataFrames
import Distributions; Dist=Distributions
#using StringEncodings

if !isfile("discoveries.csv")
    download("https://raw.githubusercontent.com/vincentarelbundock/Rdatasets/master/csv/datasets/discoveries.csv",
    "discoveries.csv")
end
#med = rdata.dataset("datasets", "discoveries")
# frame = df.DataFrame()
#frame = CSV.read("discoveries.csv", df)
df = CSV.File(open(read, "discoveries.csv")) |> DF.DataFrame
@show names(df)
#@show frame

# Histogram
hist = histogram(x=df.value, histnorm="probability")
# Poisson distribution
d = Dist.Poisson(3.1)
x = range(0,12,length=13)
line = scatter(x=x,y=Dist.pdf.(d,x))
# Plot
layout = Layout(title="Discoveries")
pt = plot([hist, line], layout)
#display(pt)

# 1.3.1 The Binomial Distribution

# Page 14
# θ is a parameter
# Θ is set of θ
# Binomial

# Page 16
d = Dist.Binomial(3,0.1)
x = range(0,3,length=4)
scat = scatter(x=x,y=Dist.pdf.(d,x), mode="markers",
    #opacity=0.1
    marker=attr(size=15, opacity=1, color="transparent",
        line=attr(color="black", width=2)
    )
)
layout = Layout(title="N=3, p=0.1")
pt = plot([scat], layout)
#display(pt)
savefig(pt, "binomial.png", format="png")
savefig(pt, "binomial.html", format="html")

scat_list = Vector{AbstractTrace}()
d = Dist.Binomial(3,0.1)
x = range(0,3,length=4)
scat = scatter(x=x,y=Dist.pdf.(d,x), mode="markers",
    #opacity=0.1
    marker=attr(size=15, opacity=1, color="transparent",
        line=attr(color="black", width=2)
    )
)
push!(scat_list,scat)

d = Dist.Binomial(3,0.5)
scat = scatter(x=x,y=Dist.pdf.(d,x), mode="markers",
    xaxis="x2", yaxis="y2",
    #opacity=0.1
    marker=attr(size=15, opacity=1, color="transparent",
        line=attr(color="black", width=2)
    )
)
push!(scat_list, scat)
# show how to plot multiple traces in one subplot
line = scatter(x=x, y=Dist.pdf.(d,x), mode="lines",
    xaxis="x2", yaxis="y2",
    line=attr(color="black", width=2)
)
push!(scat_list, line)

d = Dist.Binomial(3,0.9)
scat = scatter(x=x,y=Dist.pdf.(d,x), mode="markers",
    xaxis="x3", yaxis="y3",
    #opacity=0.1
    marker=attr(size=15, opacity=1, color="transparent",
        line=attr(color="black", width=2)
    )
)
push!(scat_list, scat)

x = range(0,30,length=31)
d = Dist.Binomial(30,0.1)
scat = scatter(x=x,y=Dist.pdf.(d,x), mode="markers",
    xaxis="x4", yaxis="y4",
    #opacity=0.1
    marker=attr(size=10, opacity=1, color="transparent",
        line=attr(color="black", width=2)
    )
)
push!(scat_list, scat)

x = range(0,30,length=31)
d = Dist.Binomial(30,0.5)
scat = scatter(x=x,y=Dist.pdf.(d,x), mode="markers",
    xaxis="x5", yaxis="y5",
    #opacity=0.1
    marker=attr(size=10, opacity=1, color="transparent",
        line=attr(color="black", width=2)
    )
)
push!(scat_list, scat)

x = range(0,30,length=31)
d = Dist.Binomial(30,0.9)
scat = scatter(x=x,y=Dist.pdf.(d,x), mode="markers",
    xaxis="x6", yaxis="y6",
    #opacity=0.1
    marker=attr(size=10, opacity=1, color="transparent",
        line=attr(color="black", width=2)
    )
)
push!(scat_list, scat)

layout = Layout(title="N=3, p=0.1",
    grid=attr(rows=2, columns=3, pattern="independent")
    #yaxis2=attr(range=[0,0.4])
    )
pt = plot(scat_list, layout)
#display(pt)

# Example 1.3, page 15
# The hard way
# P[X=k | theta] = binomial(N,k) * theta^k * (1-theta)^(N-k)
# P[X >= 1] =
Ans = 0.0
theta = 2/9
N = 4
for ind = 1:N
    global Ans
    P = binomial(N,ind) * theta^ind * (1-theta)^(N-ind)
    Ans += P
end
println("Example 1.3, Ans = $Ans")

# The easy way
dist = Dist.Binomial(4,2/9)
Ans = 1 - Dist.cdf(dist,0)
println("Example 1.3, Ans = $Ans")

# The R way
using RCall
R"x <- 1 - dbinom(0,4,2/9)"
Ans = @rget x
println("Example 1.3, the R-way, Ans=$Ans")

# 1.3.2 The Poisson Distribution, page 17
x = range(0,100,length=101)
d = Dist.Poisson(60)
scat = scatter(x=x,y=Dist.pdf.(d,x), mode="lines")
cdf = scatter(x=x,y=Dist.cdf.(d,x), mode="lines", yaxis="y2")
layout = Layout(title="lambda=3",
    yaxis=attr(title="PDF"),
    yaxis2=attr(side="right", overlaying="y")
)
pt = plot([scat,cdf], layout)
display(pt)

# Page 19
# Need to run from terminal to pop R-plot
R"lam <- seq(0, 10, length=50)"
R"y <- dpois(3, lam)"
R"plot(lam, y, xlab=\"lambda\", ylab=\"P[x=3]\", type=\"l\")"

end
nothing #suppress printing of Main...


# ### m11.5t.jl

using Pkg, DrWatson

@quickactivate "StatisticalRethinkingTuring"
using Turing
using StatisticalRethinking
Turing.setprogress!(false)

delim=';'
df = CSV.read(sr_datadir("UCBadmit.csv"), DataFrame; delim);

# Turing model

@model ppl11_5(admit, applications) = begin
    N=length(applications)
    θ ~ truncated(Exponential(1), 0, Inf)
    α ~ Normal(0,2)

    for i ∈ 1:N
        prob = logistic(α)

        # alpha and beta for the BetaBinomial must be provided.
        # The two parameterizations are related by
        # alpha = prob * theta, and beta = (1-prob) * theta.
        # See https://github.com/rmcelreath/rethinking/blob/master/man/dbetabinom.Rd
        alpha = prob * θ
        beta = (1 - prob) * θ

        admit[i] ~ BetaBinomial(applications[i], alpha, beta)
    end
end

# Sample

m11_5t = ppl11_5(df.admit, df.applications)
nchains = 4; sampler = NUTS(0.65); nsamples=2000
chns11_5t = mapreduce(c -> sample(m11_5t, sampler, nsamples), chainscat, 1:nchains)

# Result rethinking

m11_5s = "
         mean   sd  5.5% 94.5% n_eff Rhat
θ        2.74 0.96  1.43  4.37  3583    1
a       -0.37 0.31 -0.87  0.12  3210    1
";

# End of `11/m811.5t.jl`


"""
    expand_poly(x, [degree = 5], [obsdim]) -> Matrix

Perform a polynomial basis expansion of the given `degree` for
the vector `x`.

The optional parameter `obsdim` can be used to specify which
dimension of the resulting `Matrix` should denote the
observations. By default the each column will denote an
observation, which means that the returning `Matrix` will have
the size `(degree, length(x))`.

```julia
julia> expand_poly(1:5, degree = 3)
3×5 Array{Float64,2}:
 1.0  2.0   3.0   4.0    5.0
 1.0  4.0   9.0  16.0   25.0
 1.0  8.0  27.0  64.0  125.0

julia> expand_poly(1:5, 3); # same but type-stable
```

If you want each row to denote an observation you can set
`obsdim = 1`.

```julia
julia> expand_poly(1:5, degree = 3, obsdim = 1)
5×3 Array{Float64,2}:
 1.0   1.0    1.0
 2.0   4.0    8.0
 3.0   9.0   27.0
 4.0  16.0   64.0
 5.0  25.0  125.0

julia> expand_poly(1:5, 3, ObsDim.First()); # same but type-stable
```
"""
function expand_poly(x::AbstractVector; degree::Integer = 5, obsdim = ObsDim.Last())
    expand_poly(x, degree, convert(LearnBase.ObsDimension, obsdim))
end

function expand_poly(x::AbstractVector, degree::Integer, ::ObsDim.Last)
    expand_poly(x, degree, ObsDim.Constant{2}())
end

function expand_poly{T<:Number}(x::AbstractVector{T}, degree::Integer, ::ObsDim.Constant{2} = ObsDim.Constant{2}())
    n = length(x)
    # FIXME: consider support for colorants (i.e. don't hardcode Float64)
    x_vec = convert(Vector{Float64}, x)
    X = zeros(Float64, (degree, n))
    @inbounds for i = 1:n
        for d = 1:degree
            X[d, i] += x_vec[i]^(d)
        end
    end
    X
end

function expand_poly{T<:Number}(x::AbstractVector{T}, degree::Integer, ::ObsDim.Constant{1})
    n = length(x)
    # FIXME: consider support for colorants (i.e. don't hardcode Float64)
    x_vec = convert(Vector{Float64}, x)
    X = zeros(Float64, (n, degree))
    @inbounds for d = 1:degree
        for i = 1:n
            X[i, d] += x_vec[i]^(d)
        end
    end
    X
end


# ------------------------------------------------------------------------------------------
#
#
# This notebook accompanies the [workshop paper](http://jiahao.github.io/parallel-prefix/)
# by Jiahao Chen and Alan Edelman entitled "Parallel Prefix Polymorphism Permits
# Parallelization, Presentation & Proof" and will appear in the proceedings of the [First
# Workshop for High Performance Technical Computing in Dynamic
# Languages](http://jiahao.github.io/hptcdl-sc14/), held in conjunction with [SC14: The
# International Conference on High Performance Computing, Networking, Storage and
# Analysis](http://sc14.supercomputing.org/).
#
#
# ------------------------------------------------------------------------------------------

# using Pkg; Pkg.add(["Compose", "Gadfly", "Interact"])
using Compose, Gadfly, Interact

# ------------------------------------------------------------------------------------------
# # `reduce()`
#
# Reduction applies a binary operator to a vector repeatedly to return a scalar. Thus +
# becomes sum, and * becomes prod.
#
# It is considered a basic parallel computing primitive.
#
#
# ------------------------------------------------------------------------------------------

reduce(+, 1:8), sum(1:8)  # triangular numbers

reduce(*, 1:8), prod(1:8) # factorials

boring(a,b)=a
@show reduce(boring, 1:8)
boring2(a,b)=b
@show reduce(boring2, 1:8)

# ------------------------------------------------------------------------------------------
# You can also use reduce to compute Fibonacci numbers using their recurrences.
#
# $$\begin{pmatrix} f_2 \\f_1 \end{pmatrix} = \begin{pmatrix} f_1 + f_0 \\ f_1 \end{pmatrix}
# = \begin{pmatrix} 1 & 1 \\ 1 & 0 \end{pmatrix} \begin{pmatrix} f_1 \\ f_0 \end{pmatrix} $$
#
# $$\begin{pmatrix} f_{n+1} \\ f_n \end{pmatrix} = \dots
# = \begin{pmatrix} 1 & 1 \\ 1 & 0 \end{pmatrix}^n \begin{pmatrix} f_1 \\ f_0 \end{pmatrix}
# $$
#
# From this, you can show that
#
# $$\begin{pmatrix} 1 & 1 \\ 1 & 0 \end{pmatrix}^n  = \begin{pmatrix} f_{n+1} & f_n \\ f_n &
# f_{n-1} \end{pmatrix} $$
#
# (this applies reduce to the same argument over and over again -- there are of course other
# ways)
# ------------------------------------------------------------------------------------------

M=[1 1; 1 0]
reduce(*,fill(M,3))
prod(fill(M,3))

@manipulate for n=1:100
    prod(fill(big.(M),n))
end

fib(j)=reduce(*, fill(M,j))
fib.([4,7])

# ------------------------------------------------------------------------------------------
# You can solve recurrences of any complexity using `reduce`. For example, `reduce` can
# compute a Hadamard matrix from its definition in terms of its submatrices:
#
# $$H_2 = \begin{pmatrix} H_1 & H_1 \\ H_1 & -H_1 \end{pmatrix} = \begin{pmatrix} 1 & 1 \\ 1
# & -1 \end{pmatrix} \otimes H_1$$
#
# and so on.
#
# (Note: this is just using reduce to compute a matrix power.
# One can think of alternative ways for sure.)
# ------------------------------------------------------------------------------------------

# [A B]
# If A is m x n
# If B is p x q
# then kron(A,B) is mp x nq and has all the elements of A times all of the elements of B

A=[1 2;3 4]
B=[10 100; 1 -10]
⊗(A,B)=kron(A,B)

M=[ 1 1;1 -1]
H=⊗(⊗(⊗(M,M),M),M)
H*H'

Hadamard(n)=reduce(⊗, fill(M,n))
H=Hadamard(3)

H'H

H*H' #This is a legitimate Hadamard matrix

using StatsBase

rolldice() = rand(1:6)+rand(1:6)
rolls(n)=fit(Histogram,[rolldice() for i=1:n ],2:12,closed=:left)

rolls(1000)

[rolls(100) for i=1:10]

reduce(merge,[rolls(100) for i=1:10])

#Reduction of matrix multiplications
M=[rand(-3:3,2,2) for i=1:4]

display(M[4]*M[3]*M[2]*M[1])
display(reduce((A,B)->B*A, M)) #backward multiply
display(reduce(*, M))          #forward multiply

# ------------------------------------------------------------------------------------------
# In the following example we apply `reduce()` to  function composition:
# ------------------------------------------------------------------------------------------

(sin ∘ cos)(1), sin(cos(1))

((-) ∘ sin)(1), -sin(1)

( (!) ∘ )

# h=reduce( (f,g) -> (x->f(g(x))), [sin cos tan])
h=reduce(∘, [sin cos tan])
[h(π) sin(cos(tan(π)))]

cumsum(1:8)  # It is useful to know that cumsum is a linear operator
# You can use power method! Below is the underlying matrix
A=tril(ones(Int,8,8))

[A*(1:8),cumsum(1:8)]

# ------------------------------------------------------------------------------------------
# # `prefix`
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# Having discussed `reduce`, we are now ready for the idea behind prefix sum.
# Prefix or scan is long considered an important parallel
# primitive as well.
#
# Suppose you wanted to compute the partial sums of a vector, i.e. given
# `y[1:n]`, we want to overwrite the vector `y` with the vector of partial sums
#
# ```julia
# new_y[1] = y[1]
# new_y[2] = y[1] + y[2]
# new_y[3] = y[1] + y[2] + y[3]
# ...
# ```
#
# At first blush, it seems impossible to parallelize this, since
#
# ```julia
# new_y[1] = y[1]
# new_y[2] = new_y[1] + y[2]
# new_y[3] = new_y[2] + y[3]
# ...
# ```
#
# which appears to be an intrinsically serial process.
# ------------------------------------------------------------------------------------------

function prefix_serial!(y, ⊕)
    for i=2:length(y)
        y[i] = y[i-1] ⊕ y[i]
        end
    y
end

prefix_serial!([1:8;],*)

prefix_serial!([rand(1:5,2,2) for i=1:4],*)

# ------------------------------------------------------------------------------------------
# However, it turns out that because addition (`+`) is associative, we can regroup the
# _order_ of how these sums are carried out. (This of course extends to other associative
# operations such as multiplication.) Another ordering of 8 associative operations is
# provided by `prefix8!`:
# ------------------------------------------------------------------------------------------

# Eight only for pedagogy
function prefix8!(y, ⋅)
    length(y)==8 || error("length 8 only")
    for i in [2,4,6,8]; y[i] = y[i-1] ⋅ y[i]; end
    for i in [  4,  8]; y[i] = y[i-2] ⋅ y[i]; end
    for i in [      8]; y[i] = y[i-4] ⋅ y[i]; end
    for i in [    6  ]; y[i] = y[i-2] ⋅ y[i]; end
    for i in [ 3,5,7 ]; y[i] = y[i-1] ⋅ y[i]; end
    y
end

# Generalization to any n
function prefix!(y, ⋅)
    l=length(y)
    k=ceil(Int, log2(l))
    @inbounds for j=1:k, i=2^j:2^j:min(l, 2^k)              #"reduce"
        y[i] = y[i-2^(j-1)] ⋅ y[i]
    end
    @inbounds for j=(k-1):-1:1, i=3*2^(j-1):2^j:min(l, 2^k) #"broadcast"
        y[i] = y[i-2^(j-1)] ⋅ y[i]
    end
    y
end

prefix!([1:12;],*)

# ------------------------------------------------------------------------------------------
# # Polymorphism for visualization
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# We can visualize the operations with a little bit of trickery. In Julia, arrays are simply
# types that expose the array protocol. In particular, they need to implement  methods for
# the generic functions `length`, `getindex` and `setindex!`. The last two are used in
# indexing operations, since statements
#
#     y = A[1]
#     A[3] = y
#
# get desugared to
#
#     y = getindex(A, 1)
#     setindex!(A, y, 3)
#
# respectively.
#
# We can trace through the iterable by introduce a dummy array type, `DummyArray`, which
# stores no useful information but records every access to `getindex` and `setindex!`.
#
# Specifically:
#
# - `length(A::DummyArray)` returns an integer that is stored internally in the `A.length`
# field
# - `getindex(A::DummyArray, i)` records read access to the index `i` in the `A.read` field
# and always retuns `nothing`.
# - `setindex!(A::DummyArray, x, i)` records write access to the index `i`. The `A.history`
# field is appended with a new tuple consisting of the current `A.read` field and the index
# `i`.
#
# The way `DummyArray` works, it assumes an association between a single `setindex!` call
# and and all the preceding `getindex` calls since the previous `setindex!` call, which is
# sufficient for the purposes of tracing through prefix calls.
# ------------------------------------------------------------------------------------------

import Base: getindex, setindex!, length

isdefined(:DummyArray) || type DummyArray
    length :: Int
    read :: Vector
    history :: Vector{Any}
    DummyArray(length, read=[], history=[])=new(length, read, history)
end

length(A::DummyArray)=A.length

function getindex(A::DummyArray, i)
    push!(A.read, i)
    nothing
end

function setindex!(A::DummyArray, x, i)
    push!(A.history, (A.read, [i]))
    A.read = []
end

M = DummyArray(4)

M[7] = M[3] + M[2]
M

import Base.+
+(a::Void, b::Void)=a

@manipulate for i=1:100
    M = DummyArray(4)
    M[i] = M[i-2] + M[2*i]
    M
end

# ------------------------------------------------------------------------------------------
# We also need a dummy associative operator to pass to the prefix call.
# ------------------------------------------------------------------------------------------

A=prefix8!(DummyArray(8),+)

A.history

# ------------------------------------------------------------------------------------------
# Now let's visualize this! Each entry in `A.history` is rendered by a gate object:
# ------------------------------------------------------------------------------------------

mutable struct gate
    ins :: Vector
    outs:: Vector
end

import Gadfly.render

function render(G::gate, x₁, y₁, y₀; rᵢ=0.1, rₒ=0.25)
    ipoints = [(i, y₀+rᵢ) for i in G.ins]
    opoints = [(i, y₀+0.5) for i in G.outs]
    igates  = [Compose.circle(i..., rᵢ) for i in ipoints]
    ogates  = [Compose.circle(i..., rₒ) for i in opoints]
    lines = [line([i, j]) for i in ipoints, j in opoints]
    compose(context(units=UnitBox(0.5,0,x₁,y₁+1)),
    compose(context(), stroke(colorant"black"), fill(colorant"white"),
            igates..., ogates...),
    compose(context(), linewidth(0.3mm), stroke(colorant"black"),
            lines...))
end

A=gate([1,2],[2])
render(A,2,0,0)

# ------------------------------------------------------------------------------------------
# Now we render the whole algorithm. We have to scan through the trace twice; the first time
# merely calculates the maximum depth that needs to be drawn and the second time actually
# generates the objects.
# ------------------------------------------------------------------------------------------

function render(A::DummyArray)
    #Scan to find maximum depth
    olast = depth = 0
    for y in A.history
        (any(y[1] .≤ olast)) && (depth += 1)
        olast = maximum(y[2])
    end
    maxdepth = depth

    olast = depth = 0
    C = []
    for y in A.history
        (any(y[1] .≤ olast)) && (depth += 1)
        push!(C, render(gate(y...), A.length, maxdepth, depth))
        olast = maximum(y[2])
    end

    push!(C, compose(context(units=UnitBox(0.5,0,A.length,1)),
      [line([(i,0), (i,1)]) for i=1:A.length]...,
    linewidth(0.1mm), stroke(colorant"grey")))
    compose(context(), C...)
end

render(prefix!(DummyArray(8), +))

# ------------------------------------------------------------------------------------------
# Now we can see that `prefix!` rearranges the operations to form two spanning trees:
# ------------------------------------------------------------------------------------------

render(prefix!(DummyArray(120),+))

function prefix!(y, +)
    l=length(y)
    k=ceil(Int, log2(l))
    @inbounds for j=1:k, i=2^j:2^j:min(l, 2^k)              #"reduce"
        y[i] = y[i-2^(j-1)] + y[i]
    end
    @inbounds for j=(k-1):-1:1, i=3*2^(j-1):2^j:min(l, 2^k) #"broadcast"
        y[i] = y[i-2^(j-1)] + y[i]
    end
    y
end
@show prefix!(DummyArray(8),+)
render(prefix!(DummyArray(9),+))

#Look at the compose tree
#render(prefix8!(AccessArray(8),+)) |> introspect

# ------------------------------------------------------------------------------------------
# as contrasted with the serial code:
# ------------------------------------------------------------------------------------------

render(prefix_serial!(DummyArray(8),+))

@manipulate for npp=1:180
    render(prefix!(DummyArray(npp),+))
end

@manipulate for npp=1:180
    render(prefix_serial!(DummyArray(npp),+))
end

draw(PDF("serial.pdf", 3.5inch, 3.5inch), render(prefix_serial!(AccessArray(8),+)))
draw(PDF("tree.pdf", 3.5inch, 3.5inch), render(prefix8!(AccessArray(8),+)))

# ------------------------------------------------------------------------------------------
# # Polymorphism for parallel operations
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# Now let's run `prefix` in parallel on every core on the CPU. Use `addprocs` to attach
# additional worker processes.
# ------------------------------------------------------------------------------------------

import Base: fetch, length, +, *
fetch(t::Vector) = map(fetch, t) #Vectorize fetch

#Define elementary operations on remote data
length(r1::Future)=length(fetch(r1))
+(r1::Future,r2::Future)=@spawnat r2.where fetch(r1)+fetch(r2)
*(r1::Future,r2::Future)=@spawnat r2.where fetch(r1)*fetch(r2)

# ------------------------------------------------------------------------------------------
# For 8 processes, the serial version requires 7 operations. The parallel version uses 11
# operations, but they are grouped into 5 simultaneous chunks of operations. Hence we expect
# that the parallel version runs in 5/7 the time needed for the naïve serial code.
# ------------------------------------------------------------------------------------------

#Before timing, force Julia to JIT compile the functions
[f([randn(3,3) for i=1:2], *) for f in (prefix_serial!, prefix!)]

# addprocs(max(0, Sys.CPU_CORES-nprocs())) #Add all possible processors
using JuliaRunClient
# initializeCluster(2)
initializeCluster(16) # more workers can be added with a larger subscription

timings=Dict{Int,Tuple{Float64,Float64}}()
for np in 2:nprocs()
    gc_enable(false) #Disable garbage collector for timing purposes
    n=1024#4096#3072

    #Create RemoteRefs on np processes
    r=[@spawnat p randn(n,n) for p in procs()[1:np]]
    s=fetch(r) #These are in-memory matrices

    #Repeat timings a couple of times and take smallest
    t_ser = minimum([(tic(); prefix_serial!(s, *); toc()) for i=1:2])
    t_par = minimum([(tic(); @sync prefix!(r, *) ; toc()) for i=1:2])
    timings[np] = (t_ser, t_par)
    gc_enable(true)
end
timings

#timings = #Data from julia.mit.edu
#Dict(
#     1 => ( 15.585,  15.721),
#     2 => ( 16.325,  15.893),
#     3 => ( 32.079,  31.678),
#     4 => ( 46.037,  33.257),
#     5 => ( 61.947,  51.412),
#     6 => ( 77.247,  50.769),
#     7 => ( 92.588,  67.796),
#     8 => (108.057,  65.185),
#     9 => (123.556,  68.197),
#    10 => (138.408,  66.658),
#    11 => (153.392,  82.580),
#    12 => (168.874,  83.373),
#    13 => (184.227,  82.875),
#    14 => (199.741,  83.098),
#    15 => (215.881, 101.271),
#    16 => (230.531, 104.860),
#    17 => (246.576, 103.229),
#    18 => (261.412, 102.497),
#    19 => (276.926, 103.367),
#    20 => (294.404, 104.162),
#    21 => (308.995, 104.944),
#    22 => (323.287, 104.838),
#    23 => (340.173, 121.495),
#    24 => (353.717, 119.729),
#    25 => (369.113, 120.281),
#    26 => (384.638, 118.710),
#    27 => (402.484, 119.237),
#    28 => (417.036, 119.980),
#    29 => (431.273, 120.463),
#    30 => (446.288, 120.560),
#    31 => (461.756, 137.209),
#    32 => (476.653, 135.873),
#    33 => (493.989, 138.150),
#    34 => (508.089, 137.998),
#    35 => (523.703, 139.059),
#    36 => (539.305, 135.201),
#    37 => (554.201, 138.242),
#    38 => (569.276, 137.470),
#    39 => (585.268, 137.261),
#    40 => (599.921, 137.871),
#    41 => (615.718, 137.887),
#    42 => (632.463, 137.579),
#    43 => (647.355, 138.181),
#    44 => (664.344, 139.917),
#    45 => (678.899, 138.382),
#    46 => (695.815, 138.354),
#    47 => (717.068, 154.826),
#    48 => (730.902, 155.603),
#    49 => (745.608, 154.549),
#    50 => (756.012, 154.244),
#    51 => (769.611, 154.000),
#    52 => (793.350, 152.893),
#    53 => (801.568, 155.276),
#    54 => (818.174, 153.436),
#    55 => (832.994, 155.923),
#    56 => (849.526, 155.721),
#    57 => (863.337, 155.289),
#    58 => (881.577, 156.188),
#    59 => (893.735, 157.889),
#    60 => (911.597, 155.557),
#    61 => (925.117, 157.048),
#    62 => (940.897, 153.338),
#    63 => (954.825, 173.942),
#    64 => (990.757, 173.236),
#    65 => (987.457, 172.005),
#    66 => (1009.688, 173.869),
#    67 => (1022.357, 173.303),
#    68 => (1043.536, 175.044),
#    69 => (1052.156, 173.361),
#    70 => (1067.174, 172.230),
#    71 => (1079.543, 173.491),
#    72 => (1101.516, 174.311),
#    73 => (1119.993, 174.057),
#    74 => (1136.092, 175.804),
#    75 => (1154.758, 175.075),
#    76 => (1170.732, 173.453),
#    77 => (1184.598, 171.936),
#    78 => (1190.922, 172.809),
#    79 => (1204.645, 172.878),
#)

speedup(p::Integer) = ((p-1)/(floor(Int, log2(p)) + 1 + floor(Int, log2(p/3))))
#for (np, (t_ser, t_par)) in sort!([x for x in timings])
#    @printf("np: %3d Serial: %.3f sec  Parallel: %.3f sec  speedup: %.3fx (theory=%.3fx)\n",
#        np, t_ser, t_par, t_ser/t_par, speedup(np))
#end

xend = 1+maximum(keys(timings))

theplot=plot(layer(x=2:xend, y=map(speedup, 2:xend), Geom.line(), Theme(default_color=colorant"darkkhaki")),
layer(x=[np+1 for (np, (t_ser, t_par)) in timings],
      y=[t_ser/t_par for (np, (t_ser, t_par)) in timings],
    Geom.point, Theme(default_color=colorant"firebrick",
        #=default_point_size=1.25pt, =#highlight_width=0pt)),
    Guide.xlabel("Number of processors, p"), Guide.ylabel("Speedup, r"))

labeledplot = compose(render(theplot),
compose(context(), stroke(colorant"purple"), #Powers of 2
text(0.21, 0.73, "2"), text(0.24, 0.72, "4"), text(0.27, 0.7, "8"), text(0.33, 0.63, "16"),
text(0.47, 0.53, "32"), text(0.76, 0.35, "64")),
compose(context(), stroke(colorant"green"), #3x powers of 2
text(0.25, 0.62, "6"), text(0.28, 0.57, "12"), text(0.39, 0.49, "24"), text(0.63, 0.32, "48"),
)
)
draw(SVG(3.5inch, 3.2inch), labeledplot)

draw(SVG("scaling.svg", 3.5inch, 3.2inch), labeledplot)
#draw(PDF("scaling.pdf", 3.5inch, 3inch), theplot)

# ------------------------------------------------------------------------------------------
# The exact same serial code now runs in parallel thanks to multiple dispatch!
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# # Formal verification
#
# Julia allows us to implement very easily the interval of summations monoid technique for
# verifying the correctness of prefix sums
# ([doi:10.1145/2535838.2535882](http://dx.doi.org/10.1145/2535838.2535882)).
# ------------------------------------------------------------------------------------------

#Definition 4.3 of Chong et al., 2014
abstract type ⊤ end
abstract type Id end
const 𝕊 = Union{UnitRange, Type{⊤}, Type{Id}}

⊕(I::UnitRange, J::UnitRange) = I.stop+1==J.start ? (I.start:J.stop) : ⊤
⊕(::Type{Id}, ::Type{Id}) = Id
⊕(I::𝕊, ::Type{Id}) = I
⊕(::Type{Id}, I::𝕊) = I
⊕(I::𝕊, J::𝕊) = ⊤

for kernel in (prefix_serial!, prefix!), n=1:1000
    @assert kernel([k:k for k=1:n], ⊕) == [1:k for k=1:n]
end

# ------------------------------------------------------------------------------------------
# Wrong kernels produce type errors!
# ------------------------------------------------------------------------------------------

function prefix_wrong!(y, .+)
    l = length(y)
    offset = 1
    z = copy(y)
    while offset < l
        for i=offset+1:l
            z[i] = y[i-offset] .+ z[i]
        end
        offset += 2
        y = copy(z)
    end
    y
end

prefix_wrong!(Any[k:k for k=1:8], ⊕) #Should produce conversion error

@manipulate for npp=1:10
    render(prefix_wrong!(AccessArray(npp),+))
end

# ------------------------------------------------------------------------------------------
# # Experimental code which doesn't work
#
# In this section I play with other parallel prefix algorithms. All the ones here suffer
# from read antidependencies and they break the logic I used to draw the gates.
# ------------------------------------------------------------------------------------------

#Kogge-Stone variant
#Works but doesn't render properly
function prefix_ks!(y, .+)
    l = length(y)
    offset = 1
    z = copy(y)
    while offset < l
        for i=offset+1:l
            z[i] = y[i-offset] .+ z[i]
        end
        offset *= 2
        y = copy(z)
    end
    y
end

@assert prefix_ks!(ones(8), +) == [1:8;]

Base.copy(A::AccessArray) = A

render(prefix_ks!(AccessArray(8),+))

#Hillis-Steele = Kogge-Stone
function prefix_hs!(y, +)
    l = length(y)
    lk = floor(Int, log2(l))
    for d=1:lk
        for k=2d:l
            y[k] = y[k-2^(d-1)] + y[k]
        end
    end
    y
end
prefix_hs!(ones(4), +)

render(prefix_hs!(AccessArray(4), +))

#Sklansky
function prefix_s!(y, +)
    n = length(y)
    for d=1:log2(n)
        for k=1:n÷2+1
            block = 2 * (k - (mod(k, 2^d)))
            me = block + mod(k, 2^d) + 2^d
            spine = block + 2^d - 1
            if spine ≤ n && me ≤ n
                y[me] = y[me] + y[spine]
            end
        end
    end
    y
end
prefix_s!(ones(4), +)

function prefix(v,⊕)
    @show v
    k=length(v)
    k≤1 && return v
    v[2:2:k] = @show v[1:2:k-1]⊕v[2:2:k]  # Line 1
    v[2:2:k] = @show prefix(v[2:2:k],⊕)   # Line 2
    v[3:2:k] = @show v[2:2:k-1]⊕v[3:2:k]  # Line 3
    v
end

prefix(ones(8),+)

function ⊕{T<:Range}(v1::Vector{T}, v2::Vector{T})
    for v in (v1, v2), i in 1:length(v)-1
        @assert v[i].stop + 1 == v[i+1].start
    end
    if length(v1)>0 && length(v2)>0
        @assert v1[end].stop + 1 == v2[1].start
        v1[1].start:v2[end].stop
    elseif length(v1)>0 #&& length(v2)==0
        v1[1].start:v1[end].stop
    elseif length(v2)>0 #&& length(v1)==0
        v2[1].start:v2[end].stop
    else #Both are empty
        v1
    end
end

for kernel in (prefix,), n=1:100#(prefix_serial!, prefix!, prefix_ks!), n=1:100
    @assert kernel([k:k for k=1:n], ⊕) == [1:k for k=1:n]
end

.nbexports/introductory-tutorials/intro-to-julia/Reduce and Parallel Prefix - JuliaBox.jl