using PowerModels
using Ipopt
using JuMP
using DataFrames
using CSV

include("post_flex_model.jl")

filename = ARGS[1]
ADC_filename = ARGS[2]

# read data files
data = PowerModels.parse_file(filename)
flex_data = CSV.read(ADC_filename)

# number of loads
num_load = size(flex_data,1)
# dictionary contining the A and B matrices
A = Dict{Int64,Any}()
B = Dict{Int64,Any}()

# populating the A and B dictionaries
for i=1:num_load
    ca = split(flex_data[:A][i][2:end-1], '|')
    cb = split(flex_data[:b][i][2:end-1], '|')
    a = zeros(length(ca),2)
    b = zeros(length(ca))
    for j=1:length(ca)
        za = split(ca[j],' ')
        a[j,:] = parse.(Float64,za)
    end
    b = parse.(Float64,cb)
    A[i] = a
    B[i] = b
end

# build base optimization model with loads as variables
m = JuMP.Model(solver=IpoptSolver())
model, var_refs = post_ac_opf_withload(data, m, A, B)

# solver the model
status = solve(model)

# extract solution and write output.csv
gen_nums = [parse(Int64,i) for (i,gen) in data["gen"]]
pg = [getvalue(var_refs["pg"][i]) for i in gen_nums]
qg = [getvalue(var_refs["qg"][i]) for i in gen_nums]

df = DataFrame()
df[:gen_num] = gen_nums
df[:active_power] = pg
df[:reactive_power] = qg

CSV.write("output.csv",df)


using Snopt
using SparseArrays

# ---- unconstrained, no dervatives provided. -----

function matyas(g, df, dg, x, deriv)
    f = 0.26 * (x[1]^2 + x[2]^2) - 0.48 * x[1] * x[2]
    fail = false
    return f, fail
end
x0 = [5.0; 7]
lx = [-10.0; -10]
ux = [10.0; 10]
lg = []
ug = []
rows = []
cols = []
options = Dict(
    "Derivative option" => 0
)
xopt, fopt, info, out = snopta(matyas, x0, lx, ux, lg, ug, rows, cols, options)

# ---- unconstrained, providing gradient -----

function rosenbrock(g, df, dg, x, deriv)
    f = (1 - x[1])^2 + 100*(x[2] - x[1]^2)^2
    fail = false

    if deriv
        df[1] = -2*(1 - x[1]) + 200*(x[2] - x[1]^2)*-2*x[1]
        df[2] = 200*(x[2] - x[1]^2)
    end

    return f, fail
end

x0 = [4.0; 4.0]
lx = [-5.0; -5.0]
ux = [5.0; 5.0]
lg = []
ug = []
rows = []
cols = []

xopt, fopt, info, out = snopta(rosenbrock, x0, lx, ux, lg, ug, rows, cols)

# --------- constrained, not providing derivatives ------

function barnes(g, df, dg, x, deriv)

    a1 = 75.196
    a3 = 0.12694
    a5 = 1.0345e-5
    a7 = 0.030234
    a9 = 3.5256e-5
    a11 = 0.25645
    a13 = 1.3514e-5
    a15 = -5.2375e-6
    a17 = 7.0e-10
    a19 = -1.6638e-6
    a21 = 0.0005
    a2 = -3.8112
    a4 = -2.0567e-3
    a6 = -6.8306
    a8 = -1.28134e-3
    a10 = -2.266e-7
    a12 = -3.4604e-3
    a14 = -28.106
    a16 = -6.3e-8
    a18 = 3.4054e-4
    a20 = -2.8673

    x1 = x[1]
    x2 = x[2]
    y1 = x1*x2
    y2 = y1*x1
    y3 = x2^2
    y4 = x1^2

    # --- function value ---

    f = a1 + a2*x1 + a3*y4 + a4*y4*x1 + a5*y4^2 +
        a6*x2 + a7*y1 + a8*x1*y1 + a9*y1*y4 + a10*y2*y4 +
        a11*y3 + a12*x2*y3 + a13*y3^2 + a14/(x2+1) +
        a15*y3*y4 + a16*y1*y4*x2 + a17*y1*y3*y4 + a18*x1*y3 +
        a19*y1*y3 + a20*exp(a21*y1)
    fail = false

    # --- constraints ---

    g[1] = 1 - y1/700.0
    g[2] = y4/25.0^2 - x2/5.0
    g[3] = (x1/500.0- 0.11) - (x2/50.0-1)^2

    return f, fail

end

x0 = [10.0; 10.0]
lx = [0.0; 0.0]
ux = [65.0; 70.0]
lg = -Inf*ones(3)
ug = zeros(3)
rows = [1, 2, 3, 1, 2, 3]
cols = [1, 1, 1, 2, 2, 2]
options = Dict(
    "Derivative option" => 0
)

xopt, fopt, info, out = snopta(barnes, x0, lx, ux, lg, ug, rows, cols, options)

# ----- constrained, providing derivatives (dense jacobian) ---------

function barnesgrad(g, df, dg, x, deriv)

    a1 = 75.196
    a3 = 0.12694
    a5 = 1.0345e-5
    a7 = 0.030234
    a9 = 3.5256e-5
    a11 = 0.25645
    a13 = 1.3514e-5
    a15 = -5.2375e-6
    a17 = 7.0e-10
    a19 = -1.6638e-6
    a21 = 0.0005
    a2 = -3.8112
    a4 = -2.0567e-3
    a6 = -6.8306
    a8 = -1.28134e-3
    a10 = -2.266e-7
    a12 = -3.4604e-3
    a14 = -28.106
    a16 = -6.3e-8
    a18 = 3.4054e-4
    a20 = -2.8673

    x1 = x[1]
    x2 = x[2]
    y1 = x1*x2
    y2 = y1*x1
    y3 = x2^2
    y4 = x1^2

    # --- function value ---

    f = a1 + a2*x1 + a3*y4 + a4*y4*x1 + a5*y4^2 +
        a6*x2 + a7*y1 + a8*x1*y1 + a9*y1*y4 + a10*y2*y4 +
        a11*y3 + a12*x2*y3 + a13*y3^2 + a14/(x2+1) +
        a15*y3*y4 + a16*y1*y4*x2 + a17*y1*y3*y4 + a18*x1*y3 +
        a19*y1*y3 + a20*exp(a21*y1)
    fail = false

    # --- constraints ---

    g[1] = 1 - y1/700.0
    g[2] = y4/25.0^2 - x2/5.0
    g[3] = (x1/500.0- 0.11) - (x2/50.0-1)^2

    if deriv
        # --- derivatives of f ---

        dy1 = x2
        dy2 = y1 + x1*dy1
        dy4 = 2*x1
        dfdx1 = a2 + a3*dy4 + a4*y4 + a4*x1*dy4 + a5*2*y4*dy4 +
            a7*dy1 + a8*y1 + a8*x1*dy1 + a9*y1*dy4 + a9*y4*dy1 + a10*y2*dy4 + a10*y4*dy2 +
            a15*y3*dy4 + a16*x2*y1*dy4 + a16*x2*y4*dy1 + a17*y3*y1*dy4 + a17*y3*y4*dy1 + a18*y3 +
            a19*y3*dy1 + a20*exp(a21*y1)*a21*dy1

        dy1 = x1
        dy2 = x1*dy1
        dy3 = 2*x2
        dfdx2 = a6 + a7*dy1 + a8*x1*dy1 + a9*y4*dy1 + a10*y4*dy2 +
            a11*dy3 + a12*x2*dy3 + a12*y3 + a13*2*y3*dy3 + a14*-1/(x2+1)^2 +
            a15*y4*dy3 + a16*y4*y1 + a16*y4*x2*dy1 + a17*y4*y1*dy3 + a17*y4*y3*dy1 + a18*x1*dy3 +
            a19*y3*dy1 + a19*y1*dy3 + a20*exp(a21*y1)*a21*dy1

        df[1] = dfdx1
        df[2] = dfdx2

        # --- derivatives of g ---

        dg[1] = -x2/700.0  # 1, 1
        dg[2] = 2*x1/25^2   # 2, 1
        dg[3] = 1.0/500  # 3, 1
        dg[4] = -x1/700.0  # 1, 2
        dg[5] = -1.0/5   # 2, 2
        dg[6] = -2*(x2/50.0-1)/50.0  # 3, 2
    end

    return f, fail
end

x0 = [10.0; 10.0]
lx = [0.0; 0.0]
ux = [65.0; 70.0]
lg = -Inf*ones(3)
ug = zeros(3)
rows = [1, 2, 3, 1, 2, 3]
cols = [1, 1, 1, 2, 2, 2]
options = Dict(
    "Verify level" => 1
)

xopt, fopt, info, out = snopta(barnesgrad, x0, lx, ux, lg, ug, rows, cols, options)

# ----- constrained, providing derivatives (sparse jacobian) ---------

function sparsegrad(g, df, dg, x, deriv)

    f = x[1]^2 - x[2]
    fail = false

    # --- constraints ---

    g[1] = x[2] - 2*x[1]
    g[2] = -x[2]

    # --- derivatives of f ---

    df[1] = 2*x[1]
    df[2] = -1.0

    # --- derivatives of g ---

    dg[1] = -2.0  # 1, 1
    dg[2] = 1.0  # 1, 2
    dg[3] = -1.0   # 2, 2

    return f, fail

end

x0 = [0.0; 0.0]
lx = [-10.0, -10.0]
ux = [10.0, 10.0]
lg = -Inf*ones(2)
ug = zeros(2)
rows = [1, 1, 2]
cols = [1, 2, 2]
options = Dict(
    "Verify level" => 1
)

xopt, fopt, info, out = snopta(sparsegrad, x0, lx, ux, lg, ug, rows, cols, options)

# ------- adding names in output file ------

x0 = [12.0; 10.0]
lx = [0.0; 0.0]
ux = [65.0; 70.0]
lg = -Inf*ones(3)
ug = zeros(3)

rows = [1, 2, 3, 1, 2, 3]
cols = [1, 1, 1, 2, 2, 2]
names = Snopt.Names("howdy", ["x1", "x2"], ["obj", "g1", "g2", "g3"])
xopt, fopt, info, out = snopta(barnesgrad, x0, lx, ux, lg, ug, rows, cols, names=names)

# -------- warm start --------------

x0 = [12.0; 10.0]
lx = [0.0; 0.0]
ux = [65.0; 70.0]
lg = -Inf*ones(3)
ug = zeros(3)
rows = [1, 2, 3, 1, 2, 3]
cols = [1, 1, 1, 2, 2, 2]

# artificially limiting the major iterations so we can restart
options = Dict(
    "Major iterations limit" => 15
)

xopt, fopt, info, out = snopta(barnesgrad, x0, lx, ux, lg, ug, rows, cols, options)
println("major iter = ", out.major_iter)

# warm start from where we stopped
warmstart = out.warm  # a WarmStart type
xopt, fopt, info, out = snopta(barnesgrad, warmstart, lx, ux, lg, ug, rows, cols)
println("major iter = ", out.major_iter)

# -------- linear functions --------------

function lp(g, df, dg, x, deriv)
    f = 0.0  # no nonlinear component
    g .= 0.0
    df .= 0.0
    fail = false

    return f, fail
end

x0 = ones(10)
lx = zeros(10)
ux = 4*ones(10)
lg = [5.0, 7, 1]
ug = [8.0, Inf, 10.0]
rows = []
cols = []

A = [
    0.46 0.54 0.40 0.39 0.49 0.03 0.66 0.26 0.05 0.60  # obj
    0.56 0.84 0.23 0.48 0.05 0.69 0.87 0.85 0.88 0.62  # constraint 1
    0.29 0.98 0.36 0.14 0.26 0.41 0.87 0.97 0.13 0.69  # constraint 2
    0.48 0.55 0.78 0.59 0.79 0.84 0.01 0.77 0.13 0.10  # constraint 3
]

xopt, fopt, info, out = snopta(lp, x0, lx, ux, lg, ug, rows, cols, A=A)


# https://adventofcode.com/2021/day/3
using AdventOfCode
using Statistics
input = (x->parse.(Int,split(x,""))).(readlines("2021/data/day_3.txt"))
function parsebin(binvec::BitArray)
    sum(binvec.*2 .^((length(binvec)-1:-1:0)))
end
parsebin(vec::Vector{<:Int})=parsebin(Bool.(vec))
function part_1(input)
    gammavec= >(0.5).(mean(input))
    parsebin(gammavec)*
        parsebin(.!gammavec)
end

function sift(vectors, meancond, sumvec, digit=1)
    len=length(vectors)
    if len==1 return(parsebin(vectors[1])) end
    filter!(x->x[digit]==meancond(sumvec[digit]/len),vectors)
    sumvec.=sum(vectors)
    sift(vectors,meancond,sumvec,digit+1)
end
function part_2(input)
    sift(copy(input),≥(0.5),sum(input))*
        sift(copy(input),<(0.5),sum(input))
end


# Part 1: evolve until we get a repeated state, then report its "biodiversity"
dirs = CartesianIndex.([(0, -1), (1, 0), (0, 1), (-1, 0)])
biodiversity(plan) = sum([2^(i-1) for (i, v) in enumerate(plan) if v])
adjacent(plan, coord) = [try plan[coord + dir] catch BoundsError; false end for dir in dirs]

"""Evolve the plan state, returning new state. adjacency_func should return array of adjacent values"""
function evolve(plan, adjacency_func)
    new_plan = copy(plan)
    for coord in CartesianIndices(plan)
        adjacent_count = sum(adjacency_func(plan, coord))
        if plan[coord] && adjacent_count != 1
            new_plan[coord] = false
        elseif !plan[coord] && adjacent_count ∈ 1:2
            new_plan[coord] = true
        end
    end
    return new_plan
end

"""Find first repeated state, returning its biodiversity"""
function find_first_repeat(plan)
    states = Set()
    while true
        bd = biodiversity(plan)
        if bd ∈ states return bd end
        push!(states, bd)
        plan = evolve(plan, adjacent)
    end
end

plan = hcat(open(f -> collect.(readlines(f)), "day24_input.txt")...) .== '#'
println("Biodiversity of first repeated state: $(find_first_repeat(plan))")

# Part 2: recursive bugs. Count total bugs after n=200 evolutions. Our plan is 3D this time
# Third dimension of the plan is 2n+1 long, 1:n are "outer", n+1 is initial, n+2:end are "inner"
function adjacent_3d(plan_3d, coord)
    adj = []
    sx, sy, sz = size(plan_3d)
    mx, my = (sx+1)÷2, (sy+1)÷2
    # Centre squares are meaningless now
    if (coord[1], coord[2]) == (mx, my) return [false] end
    for dir in dirs
        x, y, z = (coord + CartesianIndex((dir[1], dir[2], 0))).I
        try
            if x == 0 # Past left edge, go "out" a level
                push!(adj, plan_3d[CartesianIndex((mx-1, my, z-1))])
            elseif x == sx+1 # Past right edge, go "out" a level
                push!(adj, plan_3d[CartesianIndex((mx+1, my, z-1))])
            elseif y == 0 # Past top edge, go "out" a level
                push!(adj, plan_3d[CartesianIndex((mx, my-1, z-1))])
            elseif y == sy+1 # Past bottom edge, go "out" a level
                push!(adj, plan_3d[CartesianIndex((mx, my+1, z-1))])
            elseif x == mx && y == my # Centre square, go "in" a level, find sx/sy adjacent squares
                if coord[1] == mx-1 # Left edge of inner level
                    append!(adj, plan_3d[CartesianIndices((1, 1:sy, z+1))])
                elseif coord[1] == mx+1 # Right edge of inner level
                    append!(adj, plan_3d[CartesianIndices((sx, 1:sy, z+1))])
                elseif coord[2] == my-1 # Top edge of inner level
                    append!(adj, plan_3d[CartesianIndices((1:sx, 1, z+1))])
                elseif coord[2] == my+1 # Bottom edge of inner level
                    append!(adj, plan_3d[CartesianIndices((1:sx, sy, z+1))])
                end
            else # our old fashioned simple case
                push!(adj, plan_3d[CartesianIndex((x, y, z))])
            end
        catch BoundsError
            push!(adj, false)
        end
    end
    return adj
end

n = 200
plan_3d = zeros(Bool, size(plan)..., 2n+1)
plan_3d[:, :, n+1] = plan
evolved = foldl((x, y) -> evolve(x, adjacent_3d), 1:n, init=plan_3d)
println("Total bugs after $n evolutions: $(sum(evolved))")


import Distributions, StatsBase, Random
@everywhere using Distributions, StatsBase, Random
@everywhere function datagen(n::Int, Ctrue::Int, N::Int, J::Int, n_ij::Vector{Int64}; b::Int=round(Int,rand()*1000), adjustmean::Bool=false, showpower::Bool=false, Xstd::Real=1.0)

    Random.seed!(200*b)
    groupindex = StatsBase.inverse_rle(1:n, n_ij)
    X = rand(Distributions.Normal(), N, 2)
    betas_true=ones(J)
    if adjustmean
        mu_all = -log(1/0.81 - 1)
        if Ctrue == 1
            mu_true = [mu_all]
            wi_true = [1.0]
            sigmas_true = [0.3493]
        elseif Ctrue == 2
            mu_true = [mu_all - 1.0, mu_all + 1.2]
            wi_true = [.3, .7]
            sigmas_true = [.6, .5]
        elseif Ctrue == 3
            mu_true = [mu_all - 1.0, mu_all, mu_all + 2.]
            wi_true = [.2, .5, .3]
            sigmas_true = [.6, .4, .6]
        end
    elseif showpower
        mu_all = log(1/0.779 - 1)
        if Ctrue == 2
            mu_true = [mu_all - 1.0, mu_all + 0.8]
            wi_true =  [.6, .4]
            sigmas_true = [1.2, .8]
        elseif Ctrue == 3
            mu_true = [mu_all - 2.0, mu_all + 1.0, mu_all + 3.5]
            wi_true = [.3, .4, .3]
            sigmas_true = [1.2, .8, .9]
        end
    else
        mu_all = log(1/0.779 - 1)
        if Ctrue == 1
            mu_true = [mu_all]
            wi_true = [1.0]
            sigmas_true = [0.5]
        elseif Ctrue == 2
            mu_true = [mu_all - 2.0, mu_all + 2.0]
            wi_true =  [.5, .5]
            sigmas_true = [1.2, .8]
        elseif Ctrue == 3
            mu_true = [mu_all - 4.0, mu_all + 1.0, mu_all + 4.0]
            wi_true = [.3, .4, .3]
            sigmas_true = [1.2, .8, .9]
        end
    end

    #Randomly data generation based on the setting on datagen.jl
    Random.seed!(b * 100)
    m = Distributions.MixtureModel(map((u, v) -> Distributions.Normal(u, v), mu_true, sigmas_true), wi_true)
    gamma_true = rand(m, n)

    prob = exp.(gamma_true[groupindex] .+ X*betas_true)
    prob= prob ./ (1 .+ prob)
    Y = Bool[Distributions.rand(Distributions.Binomial(1, prob[i])) == 1 for i in 1:N];
    X = X .- mean(X, dims=1);
    return(X, Y, groupindex, wi_true, mu_true, sigmas_true, betas_true, m, gamma_true)
end


#!/usr/bin/env julia
"""
A small Julia module that defines a simple function, to be used in UCBjulia.py (with pyjulia).
"""

# Small Julia module to wrap the function that computes a UCB index
module UCBjulia
    function index(rewards, pulls, t, arm, alpha=4)
        if pulls[arm] < 1
            return Inf
        else
            return (rewards[arm] / pulls[arm]) + sqrt((alpha * log(t)) / (2 * pulls[arm]))
        end
    end
end

# Small Julia function that computes a UCB index
function index(rewards, pulls, t, arm, alpha=4)
    if pulls[arm] < 1
        return Inf
    else
        return (rewards[arm] / pulls[arm]) + sqrt((alpha * log(t)) / (2 * pulls[arm]))
    end
end


module ActuarialScience

export maleMort, femaleMort,
        MortalityTable,
        qx,px, tpx, tqx,
        tqxy,tpxy, tqx̅y̅, tpx̅y̅,
        ω,w, lx, dx, ex,
        ixVector,
        InterestRate,
        i,vx,tvx,v,
        LifeInsurance,
        Ax, Axn, äx,äxn,
        Dx, Mx

include("interest.jl")
include("decrement.jl")
include("mortality.jl")

# 'actuarial objects' that combine multiple forms of decrements (lapse, interest, death, etc)
abstract type AbstractActuarial end

mutable struct ActuarialHelper
    mort::MortalityTable
    int::InterestRate
end

mutable struct LifeInsurance <: AbstractActuarial
    ah::ActuarialHelper

    function LifeInsurance(mt::MortalityTable,int::InterestRate)
        return new(ActuarialHelper(mt,int))
    end
end

###################
## COMMUTATIONS ###
###################

Dx(ah::ActuarialHelper,x) = tvx(ah.int,x,1) * lx(ah.mort,x)

preMx(ah::ActuarialHelper,x) = tvx(ah.int,x+1,1) * dx(ah.mort,x)

function Nx(ah::ActuarialHelper,x)
range =x:min(ω(ah.mort),ω(ah.int))
    return sum(map((y) -> Dx(ah,y),range))
end

function Mx(ah::ActuarialHelper,x)
    range =  x:(min(ω(ah.mort),ω(ah.int))-1)
    return sum(map((x -> preMx(ah,x)),range))
end
tEx(ah::ActuarialHelper,t,x) = Dx(x+t) / Dx(x)

##################
### Insurances ###
##################

#term insurance on age x for n years
Axn(ins::LifeInsurance,x,n) = (Mx(ins.ah,x) - Mx(ins.ah,x + n) ) / Dx(ins.ah,x)

# whole life insurance
Ax(ins::LifeInsurance,x) = Mx(ins.ah,x) / Dx(ins.ah,x)

# life annuity due
äx(ins::LifeInsurance,x) = Nx(ins.ah,x) / Dx(ins.ah,x)

# finite duration life annuity due
äxn(ins::LifeInsurance,x,n) = (Nx(ins.ah,x) - Nx(ins.ah,x+n) )/ Dx(ins.ah,x)

Dx(ins::LifeInsurance,x) = Dx(ins.ah,x)
Mx(ins::LifeInsurance,x) = Mx(ins.ah,x)

end # module


# The SteinRectangularDomainMetaKernel is a "meta" kernel of the form:
#
# kR(x, y) = \prod_{j: l[j] > -inf} (x_j - l_j)(y_j - l_j) *
#            \prod_{k: u[k] <  inf} (x_k - u_k)(y_k - u_k) *
#            k(x,y)
#
# on the domain [l_1, u_1] x ... x [l_d, u_d].
# We assume subtypes of this define the fields:
#
# l, u, and basekernel.

mutable struct SteinRectangularDomainMetaKernel <: SteinKernel
    # the base kernel k(x,y)
    basekernel::SteinKernel
    # the lower bound vector (l_1, ..., l_d)
    l::Array{Float64,1}
    # the upper bound vector (u_1, ..., u_d)
    u::Array{Float64,1}
end

# domain is [l,u]^d
SteinRectangularDomainMetaKernel(basekernel::SteinKernel, l::Float64, u::Float64, d::Int64) =
SteinRectangularDomainMetaKernel(basekernel, l * ones(d), u * ones(d))

# HELPER FUNCTIONS

function w(ker::SteinRectangularDomainMetaKernel,
           x::Array{Float64,1})
    J = find(ker.l .> -Inf)
    K = find(ker.u .< Inf)
    prodl = prod(x[J] - ker.l[J])
    produ = prod(ker.u[K] - x[K])
    prodl * produ
end

function w(ker::SteinRectangularDomainMetaKernel,
           x::Array{Float64,1},
           y::Array{Float64,1})
    w(ker, x) * w(ker, y)
end

function alpha(ker::SteinRectangularDomainMetaKernel,
               x::Array{Float64,1})
    l = ker.l; u = ker.u
    d = length(x)
    fullprod = w(ker, x)
    result = zeros(d)

    for j in 1:d
        finitel = l[j] > -Inf
        finiteu = u[j] < Inf
        if !finitel && !finiteu
            continue
        end

        if !finitel && finiteu
            result[j] = fullprod / (x[j] - u[j])
        elseif finitel && !finiteu
            result[j] = fullprod / (x[j] - l[j])
        else
            result[j] = fullprod * (2*x[j] - l[j] - u[j]) / ((x[j] - l[j]) * (x[j] - u[j]))
        end
    end
    result
end

# Definition of kernel

function k(ker::SteinRectangularDomainMetaKernel,
           x::Array{Float64,1},
           y::Array{Float64,1})
    w(ker, x, y) * k(ker.basekernel, x, y)
end

function gradxk(ker::SteinRectangularDomainMetaKernel,
                x::Array{Float64,1},
                y::Array{Float64,1})
    wx = w(ker, x)
    wy = w(ker, y)
    wx * wy * gradxk(ker.basekernel, x, y) + wy * alpha(ker, x) * k(ker.basekernel, x, y)
end

function gradxyk(ker::SteinRectangularDomainMetaKernel,
                 x::Array{Float64,1},
                 y::Array{Float64,1})
    wx = w(ker, x)
    wy = w(ker, y)
    ax = alpha(ker, x)
    ay = alpha(ker, y)
    dot(ax, ay) * k(ker.basekernel, x, y) +
        wx * dot(ay, gradxk(ker.basekernel, x, y)) +
        wy * dot(ax, gradyk(ker.basekernel, x, y)) +
        wx * wy * gradxyk(ker.basekernel, x, y)
end


using Test, LinearMaps, LinearAlgebra, Quaternions

@testset "noncommutative number type" begin
    x = Quaternion.(rand(10), rand(10), rand(10), rand(10))
    v = rand(10)
    A = Quaternion.(rand(10,10), rand(10,10), rand(10,10), rand(10,10))
    α = UniformScaling(Quaternion.(rand(4)...))
    β = UniformScaling(Quaternion.(rand(4)...))
    L = LinearMap(A)
    @test Array(L) == A
    @test Array(L') == A'
    @test Array(transpose(L)) == transpose(A)
    @test Array(α * L) == α * A
    @test Array(L * α) == A * α
    @test Array(α * L) == α * A
    @test Array(L * α ) == A * α
    @test Array(α * L') == α * A'
    @test Array((α * L')') ≈ (α * A')' ≈ A * conj(α)
    @test L * x ≈ A * x
    @test L' * x ≈ A' * x
    @test α * (L * x) ≈ α * (A * x)
    @test α * L * x ≈ α * A * x
    @test (α * L') * x ≈ (α * A') * x
    @test (α * L')' * x ≈ (α * A')' * x
    @test (α * L')' * v ≈ (α * A')' * v
    @test Array(@inferred adjoint(α * L * β)) ≈ conj(β) * A' * conj(α)
    @test Array(@inferred transpose(α * L * β)) ≈ β * transpose(A) * α
end

@testset "nonassociative number type" begin
    x = Octonion.(rand(10), rand(10), rand(10), rand(10),rand(10), rand(10), rand(10), rand(10))
    v = rand(10)
    A = Octonion.(rand(10,10), rand(10,10), rand(10,10), rand(10,10),rand(10,10), rand(10,10), rand(10,10), rand(10,10))
    α = UniformScaling(Octonion.(rand(8)...))
    β = UniformScaling(Octonion.(rand(8)...))
    L = LinearMap(A)
    @test Array(L) == A
    @test Array(α * L) == α * A
    @test Array(L * α) == A * α
    @test Array(α * L) == α * A
    @test Array(L * α) == A * α
    @test (α * L')' * v ≈ (α * A')' * v
end


###
#   transforms concs and dirality
###
import NumericExtensions.negate!, NumericExtensions.divide!, NumericExtensions.add!, NumericExtensions.exp!

function flow(Pot :: Matrix)
    d1, d2 = size(Pot)

    Flow = Array(Number8, d1, d2)

    for j in 1:d2
        for i in 1:d1

            west  = j == 1  ? d2 : j-1
            east  = j == d2 ? 1  : j+1
            north = i == d1 ? 1  : i+1
            south = i == 1  ? d1 : i-1

            ###
            # Calculate the probability of flow out of a cell based on the pot differnece

            p   = Pot[i    ,   j]    ;
            psw = Pot[south,west] - p;
            psj = Pot[south,j   ] - p;
            pse = Pot[south,east] - p;
            pie = Pot[i    ,east] - p;
            pne = Pot[north,east] - p;
            pnj = Pot[north,j   ] - p;
            pnw = Pot[north,west] - p;
            piw = Pot[i    ,west] - p;

            ge = [psw, psj, pse, pie, pne, pnj, pnw, piw]
            ge2 = [psw, psj, pse, pie, pne, pnj, pnw, piw]
            ge = negate!(divide!(ge2, add!(negate!(exp!(ge)), 1))) #

            ###
            # Remove NaN
            ###

            ge_wo_na = Number8([isnan(x) ? 1.0 / 8.0 : x/8.0 for x in ge])

            ###
            # Normalize by eight.
            ###

            Flow[i,j] = ge_wo_na
        end
    end
    return Flow
end

function diffusionJl!(Conc :: Matrix, pot :: Matrix, p_move :: Matrix) #concentration and direction
    d1, d2 = size(Conc)

    Flow = flow(pot)

    for j in 1:d2
        for i in 1:d1

            west  = j == 1  ? d2 : j-1
            east  = j == d2 ? 1  : j+1
            north = i == d1 ? 1  : i+1
            south = i == 1  ? d1 : i-1

            ###
            # Get the flow out of cell ij
            ###

            outflow =  Conc[i,j] * sum(Flow[i,j])

            ###
            # Calculate the inflow based on the outflow from other cells into this on.
            ###

            inflow  =   Conc[south, west] * Flow[south, west].s4 +
                        Conc[south, j   ] * Flow[south, j   ].s5 +
                        Conc[south, east] * Flow[south, east].s6 +
                        Conc[i    , east] * Flow[i    , east].s7 +
                        Conc[north, east] * Flow[north, east].s0 +
                        Conc[north, j   ] * Flow[north, j   ].s1 +
                        Conc[north, west] * Flow[north, west].s2 +
                        Conc[i    , west] * Flow[i    , west].s3 ;

            ###
            # Inflow - outflow = change
            ###

            p_move[i,j] = inflow - outflow
        end
    end
end


module LocalMethods

using GSL
using LinearAlgebra
using Statistics
import JLD2

export cov_n1_dot_n2, cov_angle

const data_dir = "../data/"
const cls = JLD2.FileIO.load(data_dir*"cls.jld2", "cls")
const w_on_Pl = cls[:cl_len_scalar][:,1] .* (2 .* cls[:ell] .+ 1) ./ (4π) ./ cls[:factor_on_cl_cmb]
const lmax = maximum(cls[:ell])

function cov_cosθ(cosθ,lmax::Int)
	@assert lmax >= 2 "lmax too small, must be greater than 1"
	dot(sf_legendre_Pl_array(lmax, cosθ)[3:end], w_on_Pl[3:lmax+1])
end

cov_n1_dot_n2(n1_dot_n2) = cov_cosθ(n1_dot_n2,lmax)
cov_angle(angle) = cov_cosθ(cos(angle),lmax)

function bin_mean(fk; bin=0)
	fkrtn = copy(fk)
	if bin > 1
		Nm1    = length(fk)
		subNm1 = bin * (Nm1÷bin)
		fmk    = reshape(fk[1:subNm1], bin, Nm1÷bin)
		fmk   .= mean(fmk, dims=1)
		fkrtn[1:subNm1] .= vec(fmk)
		fkrtn[subNm1+1:end] .= mean(fk[subNm1+1:end])
	end
	return fkrtn
end

end # module


# randperm
n = 100
r = [1:n;]
randperm!(r)
@test isperm(r)
@test r != [1:n;]

# randchoice
n = 10
r = [1:n;]
x = zeros(Int, 1000)
for i=1:length(x)
    x[i] = randchoice(r)
    @test 1 <= x[i] <= n
end
@test sort(union(x)) == r

a = randchoice(MersenneTwister(17),r)
@test 1 <= a <= n

r = (i for i in 1:20)
a = randchoice(r)
@test 1<= a <= 20

#getRNG
@test getRNG(-1) === Base.Random.GLOBAL_RNG
@test getRNG(0) === Base.Random.GLOBAL_RNG
@test getRNG(17) == MersenneTwister(17)

# randchoice!

s = randchoice(1:10, 3)
@test length(s) == 3
for  e in s
    @test 1 <= e <= 10
end

s = randchoice!([1:10;], 3)
@test length(s) == 3
for  e in s
    @test 1 <= e <= 10
end

s = randchoice!([1:10;], 6, exclude=[1,2])
@test length(s) == 6
for  e in s
    @test 3 <= e <= 10
end


using FinancialToolbox
if !(VERSION.major==0&&VERSION.minor<=6)
	using Test,Dates
else
	using Base.Test
end
function yearFractionTester(StartDate,EndDate,MatlabResults,TestToll=1e-14)
for i=0:FinancialToolbox.currMaxImplemented
	@test(abs(MatlabResults[i+1]-yearfrac(StartDate,EndDate,i))<TestToll)
end
return ;
end

StartDate=Date(1992,12,14);
EndDate=Date(1996,2,28);
@test(-yearfrac(StartDate,StartDate,0)==0)
@test(-yearfrac(StartDate,EndDate,0)==yearfrac(EndDate,StartDate,0))
MatlabResults=[3.20821917808219;3.20555555555556;3.25277777777778;3.20821917808219;3.20555555555556;3.20555555555556;3.20555555555556;3.20821917808219;3.20821917808219;3.25277777777778;3.20821917808219;3.20555555555556;3.20765027322409];
yearFractionTester(StartDate,EndDate,MatlabResults);

EndDate2=Date(1996,2,29);
MatlabResults2=[3.21095890410959;3.20833333333333;3.25555555555556;3.21095890410959;3.20833333333333;3.20833333333333;3.20833333333333;3.21095890410959;3.21095890410959;3.25555555555556;3.21095890410959;3.20833333333333;3.21038251366122];
yearFractionTester(StartDate,EndDate2,MatlabResults2);

EndDate3=Date(2026,2,28);
MatlabResults3=[33.2301369863014;33.2055555555556;33.6916666666667;33.2301369863014;33.2055555555556;33.2055555555556;33.2055555555556;33.2082191780822;33.2301369863014;33.6916666666667;33.2301369863014;33.2055555555556;33.2080844374580];
TestToll=1e-13;
yearFractionTester(StartDate,EndDate3,MatlabResults3,TestToll);

StartDate2=Date(1992,2,29);
EndDate4=Date(1996,2,29);
MatlabResults4=[3.99180327868852;4;4.05833333333333;4.00273972602740;3.99722222222222;4;4;4;3.99180327868852;4.05833333333333;4.00273972602740;4;4]
yearFractionTester(StartDate2,EndDate4,MatlabResults4);

StartDate3=Date(1992,3,30);
EndDate5=Date(1996,12,31);
MatlabResults5=[4.75890410958904;4.75000000000000;4.82500000000000;4.75890410958904;4.75000000000000;4.75000000000000;4.75000000000000;4.75616438356164;4.75890410958904;4.82500000000000;4.75890410958904;4.75000000000000;4.75409836065569]
yearFractionTester(StartDate3,EndDate5,MatlabResults5);

StartDate4=Date(1992,3,31);
EndDate6=Date(1996,12,30);
MatlabResults6=[4.75342465753425;4.75000000000000;4.81944444444445;4.75342465753425;4.75000000000000;4.75000000000000;4.75000000000000;4.75068493150685;4.75342465753425;4.81944444444445;4.75342465753425;4.75000000000000;4.74863387978144]
yearFractionTester(StartDate4,EndDate6,MatlabResults6);

@test_throws(ErrorException,yearfrac(StartDate,EndDate3,FinancialToolbox.currMaxImplemented+1));
@test_throws(ErrorException,yearfrac(StartDate,EndDate3,-1));

@test(daysact(Date(1992,12,14),Date(1992,12,13))==-1)

##Excel test
@test(fromExcelNumberToDate(33952)==Date(1992,12,14))

println("Test Dates Passed")


# This file tests if the number of f and g calls are correctly counted
@testset "f_g_counts" begin

    fcalls = Ref(0)
    gcalls = Ref(0)

    function f_counts_ref!(x::Vector, fvec::Vector, fcalls)
        fcalls[] += 1
        fvec[1] = 1 - x[1]
        fvec[2] = 10(x[2]-x[1]^2)
    end
    f_counts!(fvec, x) = f_counts_ref!(x, fvec, fcalls)

    function g_counts_ref!(x::Vector, fjac::Matrix, gcalls)
        gcalls[] += 1
        fjac[1,1] = -1
        fjac[1,2] = 0
        fjac[2,1] = -20x[1]
        fjac[2,2] = 10
    end
    g_counts!(fjac, x) = g_counts_ref!(x, fjac, gcalls)

    x0 = [-1.2; 1.]

    df = OnceDifferentiable(f_counts!, g_counts!, similar(x0), similar(x0))

    fcalls[] = 0
    gcalls[] = 0

    r = nlsolve(df, x0, method = :trust_region)
    @test r.f_calls == fcalls[]
    @test r.g_calls == gcalls[]

    df = OnceDifferentiable(f_counts!, g_counts!, similar(x0), similar(x0))

    fcalls[] = 0
    gcalls[] = 0

    r = nlsolve(df, x0, method = :newton)
    @test r.f_calls == fcalls[]
    @test r.g_calls == gcalls[]
end


# # Brusselator
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/models/Brusselator.ipynb)
#
#md # !!! note "Overview"
#md #     System type: polynomial continuous system\
#md #     State dimension: 2\
#md #     Application domain: Chemical kinetics
#
# ## Model description
#
# A chemical reaction is said to be *autocatalytic* if one of the reaction products is
# also a catalyst for the same or a coupled reaction, and such a reaction is called an autocatalytic reaction.
# We refer to the wikipedia article [Autocatalysis](https://en.wikipedia.org/wiki/Autocatalysis) for details.

# The Brusselator is a mathematical model for a class of autocatalytic reactions.
# The dynamics of the Brusselator is given by the two-dimensional ODE
#
# ```math
#   \left\{ \begin{array}{lcl} \dot{x} & = & A + x^2\cdot y - B\cdot x - x \\
#    \dot{y} & = & B\cdot x - x^2\cdot y \end{array} \right.
# ```

# The numerical values for the model's constants (in their respective units) are
# given in the following table.
#
# |Quantity|Value|
# |-----|-------|
# |A    | 1     |
# |B    | 1.5   |

using ReachabilityAnalysis #!jl

const A = 1.0
const B = 1.5
const B1 = B + 1

@taylorize function brusselator!(du, u, p, t)
    x, y = u
    x² = x * x
    aux = x² * y
    du[1] = A + aux - B1*x
    du[2] = B*x - aux
    return du
end

#md # !!! tip "Performance tip"
#md #     The auxiliary variables `B1`, `x²` and `aux` have been defined to make
#md #     better use of `@taylorize` and help to reduce allocations.

# ## Reachability settings
#
# The initial set is defined by ``x \in [0.8, 1]``, ``y \in [0, 0.2]``.
# These settings are taken from [^CAS13].

U₀ = (0.8 .. 1.0) × (0.0 .. 0.2); #!jl
prob = @ivp(u' = brusselator!(U), u(0) ∈ U₀, dim: 2); #!jl

# ## Results

# We use `TMJets` algorithm with sixth-order expansion in time and second order expansion
# in the spatial variables.

sol = solve(prob, T=18.0, alg=TMJets20(orderT=6, orderQ=2)); #!jl

#-

using Plots #!jl

plot(sol, vars=(1, 2), xlab="x", ylab="y", lw=0.2, color=:blue, lab="Flowpipe", legend=:bottomright) #!jl
plot!(U₀, color=:orange, lab="Uo") #!jl

# We observe that the system converges to the equilibrium point `(1.0, 1.5)`.

# Below we plot the flowpipes projected into the time domain.

plot(sol, vars=(0, 1), xlab="t", lw=0.2, color=:blue, lab="x(t)", legend=:bottomright)  #!jl
plot!(sol, vars=(0, 2), xlab="t", lw=0.2, color=:red, lab="y(t)")  #!jl

# ## Changing the initial volume

# The model was considered in [^GCLASG20] but using a different set of initial conditions. Let us parametrize
# the initial states as a ball centered at ``x = y = 1`` and radius ``r > 0``:

U0(r) = Singleton([1.0, 1.0]) ⊕ BallInf(zeros(2), r)

#-

# The parametric initial-value problem is defined accordingly.

bruss(r) = @ivp(u' = brusselator!(u), u(0) ∈ U0(r), dim: 2)

# First we solve for ``r = 0.01``:

sol_01 = solve(bruss(0.01), T=30.0, alg=TMJets20(orderT=6, orderQ=2))  #!jl

LazySets.set_ztol(Float64, 1e-15)  #!jl

plot(sol_01, vars=(1, 2), xlab="x", ylab="y", lw=0.2, color=:blue, lab="Flowpipe (r = 0.01)", legend=:bottomright)  #!jl

plot!(U0(0.01), color=:orange, lab="Uo", xlims=(0.6, 1.3))  #!jl

# We observe that the wrapping effect is controlled and the flowpipe doesn't blow up even for the large time horizon ``T = 30.0``.
# Next we plot the flowpipe zoomed to the last portion and compare ``r = 0.01`` with a set of larger initial states, ``r = 0.1``.

sol_1 = solve(bruss(0.1), T=30.0, alg=TMJets20(orderT=6, orderQ=2))

fig = plot(xlab="x", ylab="y", xlims=(0.9, 1.05), ylims=(1.43, 1.57), legend=:bottomright)  #!jl

plot(sol_1, vars=(1, 2), lw=0.2, color=:red, lab="r = 0.1", alpha=.4)  #!jl

plot!(sol_01, vars=(1, 2), lw=0.2, color=:blue, lab="r = 0.01")  #!jl

# The volume at time ``T = 9.0`` can be obtained by evaluating the flowpipe and computing the volume of the hyperrectangular overapproximation:

vol_01 = overapproximate(sol_01(9.0), Hyperrectangle) |> set |> volume  #!jl

#-

vol_1 = overapproximate(sol_1(9.0), Hyperrectangle) |> set |> volume  #!jl

# ## References

# [^CAS13]: X. Chen, E. Abraham, S. Sankaranarayanan. *Flow*: An Analyzer for Non-Linear Hybrid Systems.* In Proceedings of the 25th International Conference on Computer Aided Verification (CAV’13). Volume 8044 of LNCS, pages 258-263, Springer, 2013.

# [^GCLASG20]: S. Gruenbacher, J. Cyranka, M. Lechner, Md. Ariful Islam,, Scott A. Smolka, R. Grosu *Lagrangian Reachtubes: The Next Generation*. arXiv: 2012.07458


###################################################
# METODO DO GRADIENTE COM BUSCA LINEAR ARMIJO + INTERPOLAÇÃO QUADRÁTICA + SALVAGUARDAS
# Veja
#   leonardosecchin.github.io/files/otim1/4.1.Metodos_descida_gerais.pdf
#
# Aqui, a busca linear segue o proposto na implementação de referência do SPG:
#   www.ime.usp.br/~egbirgin/tango/codes.php
#
# Autor: Leonardo D. Secchin
# Data : 06/03/2021
#
# Exemplos de uso:
#   gradiente_interp(nlp);
#   gradiente_interp(nlp, x0=[1;1]);
#   gradiente_interp(nlp, eps=1e-8);
#   x, f, gsupn, iter, status = gradiente_interp(nlp, x0=[1;1], eps=1e-8, maxiter=2000, saidas=true);
# onde 'nlp' é a estrutura MathOptNLPModel do problema.
###################################################

using Printf, LinearAlgebra, NLPModels

# FUNÇÃO PRINCIPAL
function gradiente_interp(nlp; x0=nothing, eps=1.0e-6, maxiter=10000, saidas=true)

    # DADOS DE ENTRADA
    # nlp     : estrutura MathOptNLPModel do problema
    # x0      : ponto inicial (opcional)
    # eps     : precisão para convergencia (opcional, padrão=1.0e-6)
    # maxiter : número máximo de iterações (opcional, padrão=1000)
    # saidas  : mensagens na tela?

    # DADOS DE SAÍDA
    # x       : último iterando
    # f       : f(x)
    # gsupn   : norma do supremo do gradiente de f
    # iter    : número de iterações
    # status  : 0=sucesso, 1=falha

    iter   = 0
    status = 1

    # parâmetro busca linear inexata (Armijo)
    eta = 1e-4

    # define ponto inicial caso não fornecido
    if x0 == nothing
        # captura ponto inicial da estrutura nlp
        if nlp.meta.x0 != nothing
            x0 = nlp.meta.x0
        else
            # caso não exista, seta x0=origem
            x0 = zeros(Float64,nlp.meta.nvar)
        end
    end

    # f, gradiente e norma
    x     = float(x0)
    f     = obj(nlp, x)
    g     = grad(nlp, x)
    gsupn = norm(g, Inf)

    # imprimi cabeçalho saídas
    if saidas
        @printf("\nit     \tf         |grad|\n==========================")
        @printf("\n%d\t%8.2e  %8.2e", iter, f, gsupn)
    end

    while gsupn > eps && iter < maxiter
        # direção
        d = -g

        # retorna novo iterando após busca linear (Armijo + interpolação quadrática)
        x, f = buscalinear(nlp, x, f, f, g, d, eta)

        # atualiza dados do iterando
        g     = grad(nlp, x)
        gsupn = norm(g, Inf)

        iter += 1

        # imprimi iteração corrente
        if saidas
            if mod(iter,20) == 0
                @printf("\n\nit     \tf         |grad|\n==========================")
            end
            @printf("\n%d\t%8.2e  %8.2e", iter, f, gsupn)
        end
    end

    # status de sucesso
    if gsupn <= eps
        status = 0
    end

    # saída
    if saidas
        println("\n\n*******************************")
        if status == 0
            println("PROBLEMA RESOLVIDO COM SUCESSO!")
        else
            println("FALHA NA RESOLUÇÃO.")
        end
        println("*******************************\n")
    end

    return x, f, gsupn, iter, status
end

# BUSCA LINEAR COM INTERPOLAÇÃO QUADRÁTICA E SALVAGUARDAS
#
# Código implementado como no SPG disponível em
#   www.ime.usp.br/~egbirgin/tango/codes.php
#
# Aqui, a interpolação quadrática é tentada sucessivas vezes nos pontos x+t*d.
# Salvaguardas são empregadas, isto é, caso o passo da
# interpolação seja próximo de 0 ou 1, usa backtracking (divide por 2).
# Essa estratégia une a aceleração da interpolação quadrática
# com um melhor controle de passos extremos.
#
# Ao calcular o passo da interpolação no ponto x+t*d, aparece um t^2.
# Em
#   leonardosecchin.github.io/files/otim1/4.1.Metodos_descida_gerais.pdf
# só calculamos a interpolação referente ao ponto x+d. Você pode
# refazer as contas dessa aula e verificar que aparece o t^2 quando o ponto
# de referência é x+t*d.
#
# A variável "fmax" neste código será igual à "f". No SPG, "fmax" receberá o
# máximo dos últimos valores de f, enquanto "f" receberá o valor de f no ponto
# corrente.
function buscalinear(nlp, x, f, fmax, g, d, eta)

    # calcula g' * d
    gtd = g' * d

    # passo inicial
    t = 1.0

    xnew = x + d
    fnew = obj(nlp, xnew)

    # enquanto Armijo não é satisfeito, atualiza t...
    while fnew > fmax + t*eta*gtd
        # se t for próximo de 0, usa backtracking
        if t <= 0.1
            t /= 2.0
        else
            # passo da interpolação quadrática
            tquad = - 0.5*gtd*(t^2) / (fnew - f - gtd)

            # se tquad for próximo de 1, usa backtracking
            if (tquad < 0.1) || (tquad > 0.9)
                tquad = t / 2.0
            end

            t = tquad
        end

        # nova tentativa
        xnew = x + t*d
        fnew = obj(nlp, xnew)
    end

    # retorna novo iterando
    return xnew, fnew
end


### 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

# ╔═╡ 250252f0-6e61-11ec-2ffe-ebddaa9f0fb3
begin
	using Pkg
	Pkg.activate("../../")
	using Revise
end

# ╔═╡ 0eeb5bb1-bb41-43cc-8c14-eb9c1ebf6fad
using ImgProcMic, FourierTools, IndexFunArrays, ImageShow, FFTW, TestImages, Colors

# ╔═╡ 1b2ebe1a-f3a7-494c-bcca-18df5403ee65
using PlutoTest, Noise, Tullio, PlutoUI, Images, Plots

# ╔═╡ 4aeceb65-f226-41cf-988b-273f42c75059
using SpecialFunctions

# ╔═╡ c11d25d8-8fe9-476e-9a6d-66f625f29348
md"# 1 Convolution with Small Kernels

When the kernel (e.g. the PSF) of a convolution is spatially very small (e.g. $3 \times 3$), it is more efficient to perform the convolution with a sum approach (_for loops_) in real space.

In this part, you implement the convolution via the sum approach.

You will notice that the question arises, how to treat the boundary pixels. We follow the approach, that the output image is smaller than the input image. This is due to the reason that we only convolve inside the image where the kernel fits completely in.

If you need a graphical visualization, see [here](https://kharshit.github.io/img/edge_detection.png).
"

# ╔═╡ 6f87bb34-ae2b-4c8e-9c21-621644445df3
"""
	conv_s(img, kernel)

Convolution of `img` with `kernel`.

## Procedure
* Determine the correct output size
* Move the kernel over the image like [here](https://kharshit.github.io/img/edge_detection.png).
"""
function conv_s(img, kernel)
	out = similar(img, size(img) .- size(kernel) .+ 1)

	# runtime complexity is:
	for i in 1:size(out, 1)
		for j in 1:size(out, 2)
			# value we are going to write into out
			# that value accumulates the elementwise multiplication
			# and summation
			v = zero(eltype(img))
			for ki in 1:size(kernel, 1)
				for kj in 1:size(kernel, 2)
					# -1 because of Julia's one-based indexing style
					v += kernel[ki, kj] * img[i + ki - 1, j + kj - 1]
				end
			end
			out[i, j] = v
		end
	end
	return out
end

# ╔═╡ b030ccf9-9a48-4410-8ac1-f3a2a02b5b5c
md"runtime complexity is:

$\mathcal{O}(L^2 \cdot K^2)$

where $L$ is the image size, and $K$ the kernel size.
"

# ╔═╡ a879425f-a152-4d4d-b6b8-f2a03b7a452f
img = Float32.(testimage("mandril_gray"));

# ╔═╡ 012dd588-350b-4c2f-99c9-6371e443e32c
similar(img, (2,2))

# ╔═╡ 61ccfb00-513e-4bca-80eb-133aa58a1def
begin
	# todo create a kernel of size (5,5) filled with 1s
	# but afterwards normalize the sum to 1
	kernel = ones(Float32, (5,5))
	kernel ./= sum(kernel)
end

# ╔═╡ d51fadc0-3227-4e80-96f2-269a8fd2bac2
img_b = conv_s(img, kernel);

# ╔═╡ caab13b5-12a6-4058-8931-1f8a9abe57e8
[gray_show(FourierTools.select_region(img, new_size=size(img_b))) gray_show(img_b)]

# ╔═╡ cc090365-274f-4e56-9be5-cbc1a32295ad
kernel_h_filter = repeat([1 0 -1], inner=(3, 1))

# ╔═╡ b63022e3-3197-4dfe-9b7f-3ee8babce73a
md"
Those examples show the vertical and horizontal edges of the mandril
"

# ╔═╡ 7ec68244-fe43-4418-b9ee-0cb7e67572c9
[gray_show(conv_s(img, kernel_h_filter)) gray_show(conv_s(img, kernel_h_filter'))]

# ╔═╡ 2bd635be-6f07-40df-a25e-cfdf2dd22d6f
md"## Benchmark with FFT based
Is your method faster than FFT based conv?
Rerun below cell more often!

I guess, you won't beat Tullio!
"

# ╔═╡ 3765b001-deb1-49b2-87a1-0100d36d3087
function conv_tullio(img, k)
	@tullio out_t[i, j] := (img[i + ik - 1, j + jk - 1] *
							kernel[ik, jk]) (i in 1:508) (j in 1:508)
end

# ╔═╡ 57427d81-961c-488e-9a9d-e625a2e536c3
# run me multiple times!
with_terminal() do
	print("Real space conv\t\t")
	@time conv_s(img, kernel)
	print("FFT based conv\t\t")
	kernel_big = ifftshift(FourierTools.select_region(kernel, new_size=size(img)))
	@time conv(img, kernel_big)
	print("Tullio based conv\t")
	@time conv_tullio(img, kernel)
end

# ╔═╡ 6b4028de-e605-455f-947d-62649f2e26fa
md"## 1 Test"

# ╔═╡ 53ae1016-55d7-4545-a29d-b9a4fc90b854
PlutoTest.@test size(conv_s(ones((3,3)), ones((2,2)))) == ((2,2))

# ╔═╡ 01428578-4258-413d-8b8b-27203527fd2b
img_tt = zeros((6,6));

# ╔═╡ 81605617-d091-47b2-bd85-1cd8ae0fef18
img_tt[:, 1:3] .= 10;

# ╔═╡ f3fb21e1-990a-42ac-8a5d-52154be23cfb
# error disappears if size is correct
PlutoTest.@test conv_s(img_tt, kernel_h_filter) ≈ [0.0 30.0 30.0 0.0; 0.0 30.0 30.0 0.0; 0.0 30.0 30.0 0.0; 0.0 30.0 30.0 0.0]

# ╔═╡ 94694b2a-715a-450b-a024-fc474ce19975
md"# 2 Cell Detection

Based on correlation, we can come up with a simple cell detection algorithm. For the detection we need a template cell.

However, such correlation based methods are not very robust with respect to noise and distortions.

Btw, what you implemented above is mathematically not a convolution but a correlation (we didn't reverse the kernel).

The equation we want to implement in this part, is called _normalized cross-correlation_ since it divides by the local standard deviation of intensities:

$\text{NCC}\{I, K\}[i, j] = \frac{\sum_{m, n} I[i + m - 1, j + n - 1] \cdot K[m, n]}{\sqrt{\sum_{m, n} I[i + m - 1, j + n - 1]^2} \cdot \sqrt{\sum_{m, n} K[m, n]^2}}$
"

# ╔═╡ f43bd606-d202-4f74-8e46-28bab5af6d45
# load cells
cells = 1 .- Float32.(Gray.(testimage("blobs")));

# ╔═╡ 4ed29eeb-0cbc-4810-8a6c-1b7aaa0014be
gray_show(cells)

# ╔═╡ 1664be9b-08bc-4144-b191-b82bc9e2cee2
"""
	get_template(cells[, t_size=(20, 20)])

Extracts a template cell at position `pos` with template size `t_size`.

# Procedure
* Start the extract at `pos`.
* End the extraction at `pos .+ t_size`.
* Return the extracted part
"""
function get_template(cells, pos, t_size=(20, 20))
	# we need the -1 because of the one based indexing in julia
	return cells[pos[1]:pos[1] + t_size[1] - 1, pos[2]:pos[2] + t_size[2] - 1]
end

# ╔═╡ d27183d2-0d2b-4868-b34b-a90b2b5a0636
# you can change to different cells as well
cell_pos = (74, 53)

# ╔═╡ 0be4d313-6a6e-4eb6-bdc9-006ef6a726e6
# you can also change size
cell_size = (23, 23)

# ╔═╡ fe0632ef-e1b2-4a6f-ba53-44865b7971d8
cell = get_template(cells, cell_pos, cell_size);

# ╔═╡ 5f9744fa-a818-46bf-b83e-942ad1baac12
gray_show(cell)

# ╔═╡ cd861f28-8923-4ff7-ba24-6875cfa99f39
"""
	NCC(img, kernel)

Normalized cross-correlation of `img` with `kernel`.

## Procedure
* Determine the correct output size
* Move the kernel over the image like [here](https://kharshit.github.io/img/edge_detection.png).
* Calculate the local standard deivations and divide by them
"""
function NCC(img, kernel)
	out = similar(img, size(img) .- size(kernel) .+ 1)

	# runtime complexity is:
	for i in 1:size(out, 1)
		for j in 1:size(out, 2)
			# value we are going to write into out
			# that value accumulates the elementwise multiplication
			# and summation
			numerator = zero(eltype(img))
			denominator1 = zero(eltype(img))
			denominator2 = zero(eltype(img))

			for m in 1:size(kernel, 1)
				for n in 1:size(kernel, 2)
					# -1 because of Julia's one-based indexing style
					numerator += kernel[m, n] * img[i + m - 1, j + n - 1]
					denominator1 += img[i + m - 1, j + n - 1]^2
					denominator2 += kernel[m, n]^2
				end
			end
			out[i, j] = numerator / (√(denominator1) * √(denominator2))
		end
	end
	return out
end

# ╔═╡ 3fa412f2-d107-4caa-abf6-bf47c8333027
corr = NCC(cells, cell);

# ╔═╡ ff16b0cb-b8ba-4d3a-bcc8-da5c120ea898
gray_show(corr)

# ╔═╡ 0c108a6b-599a-4581-8d8d-68b9ae5865df
"""
	find_max(arr)

Returns the value and the index of the maximum intensity.

### Procedure
* Iterate over the array
* Update the current value if you find a larger one

### Examples
```julia-repl
julia> find_max([1 2 3; 4 5 6; 100 8 9])
(100, (3, 1))
```
"""
function find_max(arr)
	# TODO: initate inds and value meaningfully
	inds = (0, 0)
	value = eltype(arr)(-Inf)

	#for cind in CartesianIndices(arr)
	#	v = arr[cind]
	#	if v > value
	#		value = v
	#		inds = Tuple(cind)
	#	end
	#end

	for i in 1:size(arr, 1)
		for j in 1:size(arr, 2)
			v = arr[i, j]
			if v > value
				value = v
				inds = (i, j)
			end
		end
	end
	# TODO: return value and inds
	return value, inds
end

# ╔═╡ 2d7d48ac-c2e1-4acd-8e21-3a9314e19390
collect(eachindex(img))[2]

# ╔═╡ d1660777-a3a2-4abb-a394-4872da699c4e
Float64(-Inf)

# ╔═╡ 49a583a5-ef6b-4e5b-8ea2-0e9a5088b3cc
md"Now check whether `find_max` finds your template in the image again"

# ╔═╡ 64af243c-77f9-4c87-8f59-a4e506e3cb46
find_max(corr)

# ╔═╡ 5610563c-ea7f-40c6-ba8c-5a816bc71a2e
PlutoTest.@test find_max(corr)[2] == cell_pos

# ╔═╡ 37f43e35-015a-4596-96b0-13ccb720c767
md"## 2 Test"

# ╔═╡ bef2f5e8-e6ef-49df-9125-1d814feb31d7
kernel_ncc = [9 10; 13 14];

# ╔═╡ cf18b86d-6c82-4335-b46c-60cd161504b5
img_ncc = [Float64((i - 1) * 4 + j) for i = 1:4, j = 1:4];

# ╔═╡ 2cdd244d-33cb-487e-9fa2-069f8a845ef9
arr_rand = randn((32,32));

# ╔═╡ 8e5f5bf2-5b93-4c52-9960-ea1fb60888ce
PlutoTest.@test size(get_template(cells, (10, 10), (3, 7))) == (3,7)

# ╔═╡ 15626656-af92-44d3-83dd-9775172caed8
PlutoTest.@test get_template(cells, (10, 10), (3, 3)) ≈ Float32[0.15686274 0.18823528 0.31372547; 0.18823528 0.25098038 0.4078431; 0.25098038 0.34509802 0.43921566]

# ╔═╡ a9db1ffb-cc86-4d78-8f24-8b847443aae8
PlutoTest.@test Tuple.(find_max(arr_rand)) == Tuple.(findmax(arr_rand))

# ╔═╡ e144fe4f-cb8e-4158-a59a-8a93faca1f0c
PlutoTest.@test NCC(img_ncc, kernel_ncc) ≈ [0.9376736528408373 0.9683640522700839 0.9836212432229419; 0.9958743946641203 0.9981668178901743 0.9993408297073572; 0.9999999999999999 0.9999029998713893 0.9996660511437935]

# ╔═╡ b34bb590-0fd4-491f-a392-a538d16a86b6

# ╔═╡ 1fc3e23d-e283-43a1-a645-cfbaa3821870
md"# 3 Watershed
Based on the correlation map, we can use that as seed for a watershed algorithm.

What we need is an image with _seeds_. By that we mean an image full of 0s.
At the rough positions of the cells we put a 1.
"

# ╔═╡ 2fc20959-426a-410f-900b-bdde889426fd
"""
	get_seeds(cells, cell)

The `cells` you want to segment with the template `cell`.

# Procedure
* calculate the `NCC` arr (called `corr`)
* use `findlocalmaxima(corr)` to get the local peaks
* return an array which is 0 everywhere, except at the local peaks, where it is 1
"""
function get_seeds(cells, cell)
	# TODO: calc corr (call method from above)
	corr = NCC(cells, cell)

	# TODO: create an arr_n filled with zeros
	arr_n = zeros(eltype(corr), size(corr))
	# TODO: use findlocalmaxima
	maxs = findlocalmaxima(corr)

	# TODO: iterate over local maxima and set the corresponding
	# values of the arr_n to 1
	for cinds in maxs
		arr_n[cinds] = 1
	end

	return arr_n
end

# ╔═╡ e0f00df4-7c63-448c-b005-eda95cb4cdb8
findlocalmaxima(corr)

# ╔═╡ 7d95732e-94f4-448c-8e19-7969271059a9
10 * gray_show(get_seeds(cells, cell)) + gray_show(corr)

# ╔═╡ 86a8bf70-3c82-4100-8b63-5d010523a196
# we call watershed algorithm
segments = watershed(.- corr, label_components(get_seeds(cells, cell)));

# ╔═╡ 3641d445-1fff-416d-849b-f9d5715fef0f
heatmap(segments.image_indexmap .* (FourierTools.select_region(cells, new_size=size(segments.image_indexmap)) .> 0.3), c=:glasbey_hv_n256)

# ╔═╡ cc688fb0-0da2-4d5c-a0e8-bcff8af05ad1
gray_show(reverse(cells, dims=(1)))

# ╔═╡ e559b4b6-e26c-446e-b67e-b196e8ac3021
md"# 3 Test"

# ╔═╡ 6b966b4c-4b5c-4d62-a842-efe204b1c0fb
PlutoTest.@test get_seeds(Float32[1 2 3 4; 5 6 7 8; 9 10 11 12; 13 14 15 16], [3 4; 5 6]) ≈ Float32[0.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 0.0]

# ╔═╡ 37c5f515-74a3-4bd0-b7e6-6e498b375e23

# ╔═╡ cfe15c1b-7b29-4514-82c4-e6ffe86456b0
md"# 4 NCC starts to fail
For which values does the NCC start to fail?
"

# ╔═╡ 814d4261-c433-47a9-a35a-16d83be3d4f4
"""
	draw(img, template, pos)

Draw the `template` into the `img`.
The top left corner aligns each with `pos`.

So we basically set the `template` into the `img`.
"""
function draw(img, template, pos)
	img[pos[1]:pos[1] + size(template, 1) - 1,
		pos[2]:pos[2] + size(template, 2) - 1] .= template
	return img
end

# ╔═╡ 3e30b081-e4df-4518-98aa-96e6e0880a4a
md"σ = $(@bind σ Slider(0:0.01:0.4, show_value=true))"

# ╔═╡ d45bcd94-6adf-40eb-baee-e8edab966f07
md"N = $(@bind N Slider(1:300, show_value=true))"

# ╔═╡ eb005f06-1b2f-45fe-b6b2-abfc6b4dbb1b
md"
First, we shear the image, then we add Poisson noise and finally Gauss noise.
"

# ╔═╡ a01fab74-b77b-4b6f-8415-b0f60530593b
cells_n = add_gauss(poisson(shear(cells, 20), N), σ);

# ╔═╡ de0edf52-cf26-4265-8f10-d8720dfee53f
corr_n = NCC(cells_n, cell);

# ╔═╡ 15dc3591-8141-4e8d-9778-8290de7627ad
[ gray_show(cells_n) gray_show(draw(cells_n, cell .* 3, find_max(corr_n)[2]))]

# ╔═╡ 5ef8d15b-b407-4f18-8ed3-b0d4c69b6345
gray_show(corr_n)

# ╔═╡ 5e32a3b3-30a9-4cfe-8005-b2eebfd5a4b4
md"Those values should roughly agree. Remember we sheared the array, so they won't exactly be the same"

# ╔═╡ b94655ba-382e-402e-a0c0-89d9d4d91bb6
find_max(corr_n)

# ╔═╡ 3ee0c65f-02b2-41b9-a95b-9932203223c8
find_max(corr)

# ╔═╡ 84a06e2e-13b9-44ea-a1e3-809791017656
md"## 4 Test"

# ╔═╡ 0d4763cc-f47a-4de0-9cde-12a73556b5e3
PlutoTest.@test draw(zeros((5, 5)), ones((2, 2)), (2, 2)) ≈ [0.0 0.0 0.0 0.0 0.0; 0.0 1.0 1.0 0.0 0.0; 0.0 1.0 1.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0]

# ╔═╡ 15b6196d-0b87-45fd-8303-770969386d10
md"# 5 Simulating a Scalar High NA PSF
We often discussed how you can simulate a simple PSF:
* start with an array full of 1s in Fourier space
* low pass filter it
* take the Fourier transform
* take `abs2`

However, that procedure suffers from hard edge artifacts in Fourier space.
Fortunately, in 2D we know that the solution of a perfect PSF is a

$\text{jinc}^2(x) = \left(\frac{J_1(x)}{x}\right)^2$

where $J_1$ is the Bessel function of the first kind.

We now want to simulate PSF but with an apodization function multiplied in Fourier space (e.g. the aplantic factor).

"

# ╔═╡ 75852dd4-6cc9-4040-84e2-cdb0ce812ec4
"""
	simulate_PSF(s, r, θ, [T=Float64])

Simulate the incoherent 2D PSF with radius `r` and total output size `s`.
This functions returns the PSF centered around the center pixel.
`θ` is the angle over the field size needed for the aplanatic factor.

The procedure is as following:
* create a disc with values of 1 and radius `r` (use `rr2`)
* multiply `cos.(θ)`
* take the `ift` of it.
* normalize the sum to 1
* center pixcel should be located in the center
"""
function simulate_PSF(s, r, θ, T=Float64)
	disc = T.(rr2(T, s) .< r^2)

	# a modification of the perfect pupil
	disc .*= cos.(θ)
	apsf = ift(disc)
	psf = abs2.(apsf)
	psf ./= sum(psf)

    return psf
end

# ╔═╡ bfa26eb6-d9ea-4a5d-80a5-84f048ed3ea0
"""
	simulate_PSF_jinc(s, scale, θ, [T=Float64])

`s` is the size. `scale` is the argument passed to `rr`.
`θ` is the angle over the field size needed for the aplanatic factor.
This functions returns the PSF centered around the center pixel.

## Procedure
* Call `jinc.` of the `rr` function which is scaled by `scale`
* Take `ft`.
* multiply aplanatic factor
* take `ift`.
* take `abs2.` and normalize sum
"""
function simulate_PSF_jinc(s, scale, θ, T=Float64)
	# ideal APSF
	apsf = jinc.(rr(s, scale=scale))

	# go to Fourier space and apply cos factor
	pupil = ft(apsf)
	pupil .*= cos.(θ)

	# go to real space again
	apsf_n = ift(pupil)

	# calculate psf
	psf = abs2.(apsf_n)
	psf ./= sum(psf)

	return psf
end

# ╔═╡ 23cc65b9-76d8-43c9-a398-22212b00c76a
θ = abs.(asin.(min.(1, rr(Float32, (100, 100), scale=0.05))));

# ╔═╡ 88b9c633-01ef-4bf2-9180-0bb2b59f4766
# todo
# What is the correct scaling such that both methods return a same scaled PSF
# as a hint, look at the cell below
# factor of 5 comes from: 10 / 100 * 2
# 10 is the radius of the PSF in Fourier space

scale = 1 / 5 # fix value

# ╔═╡ 170b5e91-dfaa-4f13-ac45-01b875ba02d8
# hint, look at this output
gray_show(abs.(ft(jinc.(scale .* rr((100, 100))))))

# ╔═╡ 19e2a53c-7716-4fd3-8716-87586c60ec73
gray_show(rr((100, 100)) .< 10)

# ╔═╡ 497f6e35-4b9a-4349-8366-9f54aa1476c5
md"### Can you spot differences?"

# ╔═╡ 71510434-8685-49ee-b514-31d1c7c411d8
psf_rr = simulate_PSF((100, 100), 10, θ);

# ╔═╡ 195bb3ce-5ade-4c31-9385-12cc4c119762
psf_jinc = simulate_PSF_jinc((100, 100), scale, θ);

# ╔═╡ 8d2c5280-d9c1-414d-a64c-8d8ef5a88b10
gray_show(log.(psf_rr), set_zero=true)

# ╔═╡ abf9ed24-ea46-4067-b786-3a6b23be4c80
gray_show(log.(psf_jinc), set_zero=true)

# ╔═╡ e3d0bb6e-e9c9-4a2d-97bf-73041958771b
md"## 5 Test"

# ╔═╡ f4b0f15f-016e-41ba-a999-e9f792a534fc
PlutoTest.@test simulate_PSF((3, 3), 10, 1) ≈ [0.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 0.0]

# ╔═╡ db517950-7031-4495-8198-f054eb833e16
PlutoTest.@test simulate_PSF_jinc((3, 3), 1000000, 1) ≈ [0.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 0.0]

# ╔═╡ 9a09fa1a-8045-4e93-9140-7f57915146ac
PlutoTest.@test simulate_PSF((10, 10), 3, 0.3) ≈ Float32[0.00040000005 0.0 0.0006111457 0.0 0.0041888547 0.01 0.0041888547 0.0 0.0006111457 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0006111457 0.0 0.0009337476 0.0 0.006400001 0.015278643 0.006400001 0.0 0.0009337476 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0041888547 0.0 0.006400001 0.0 0.043866254 0.10472137 0.043866254 0.0 0.006400001 0.0; 0.01 0.0 0.015278643 0.0 0.10472137 0.25 0.10472137 0.0 0.015278643 0.0; 0.0041888547 0.0 0.006400001 0.0 0.043866254 0.10472137 0.043866254 0.0 0.006400001 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0006111457 0.0 0.0009337476 0.0 0.006400001 0.015278643 0.006400001 0.0 0.0009337476 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0]

# ╔═╡ ea3b65bb-2e41-4c17-9df8-35eaf69975f6
PlutoTest.@test simulate_PSF_jinc((4, 4), 0.1, 0.3, Float64) ≈ [0.05502252115995306 0.059372295904787806 0.060884063254034754 0.059372295904787806; 0.059372295904787806 0.0640036567471649 0.0656126102678531 0.0640036567471649; 0.06088406325403477 0.06561261026785312 0.06725510065275418 0.06561261026785312; 0.059372295904787806 0.0640036567471649 0.0656126102678531 0.0640036567471649]

# ╔═╡ Cell order:
# ╠═250252f0-6e61-11ec-2ffe-ebddaa9f0fb3
# ╠═0eeb5bb1-bb41-43cc-8c14-eb9c1ebf6fad
# ╠═1b2ebe1a-f3a7-494c-bcca-18df5403ee65
# ╟─c11d25d8-8fe9-476e-9a6d-66f625f29348
# ╠═6f87bb34-ae2b-4c8e-9c21-621644445df3
# ╟─b030ccf9-9a48-4410-8ac1-f3a2a02b5b5c
# ╠═012dd588-350b-4c2f-99c9-6371e443e32c
# ╠═a879425f-a152-4d4d-b6b8-f2a03b7a452f
# ╠═61ccfb00-513e-4bca-80eb-133aa58a1def
# ╠═d51fadc0-3227-4e80-96f2-269a8fd2bac2
# ╠═caab13b5-12a6-4058-8931-1f8a9abe57e8
# ╠═cc090365-274f-4e56-9be5-cbc1a32295ad
# ╟─b63022e3-3197-4dfe-9b7f-3ee8babce73a
# ╠═7ec68244-fe43-4418-b9ee-0cb7e67572c9
# ╟─2bd635be-6f07-40df-a25e-cfdf2dd22d6f
# ╠═3765b001-deb1-49b2-87a1-0100d36d3087
# ╠═57427d81-961c-488e-9a9d-e625a2e536c3
# ╟─6b4028de-e605-455f-947d-62649f2e26fa
# ╠═53ae1016-55d7-4545-a29d-b9a4fc90b854
# ╠═01428578-4258-413d-8b8b-27203527fd2b
# ╠═81605617-d091-47b2-bd85-1cd8ae0fef18
# ╠═f3fb21e1-990a-42ac-8a5d-52154be23cfb
# ╟─94694b2a-715a-450b-a024-fc474ce19975
# ╠═f43bd606-d202-4f74-8e46-28bab5af6d45
# ╠═4ed29eeb-0cbc-4810-8a6c-1b7aaa0014be
# ╠═1664be9b-08bc-4144-b191-b82bc9e2cee2
# ╠═d27183d2-0d2b-4868-b34b-a90b2b5a0636
# ╠═0be4d313-6a6e-4eb6-bdc9-006ef6a726e6
# ╠═fe0632ef-e1b2-4a6f-ba53-44865b7971d8
# ╠═5f9744fa-a818-46bf-b83e-942ad1baac12
# ╠═cd861f28-8923-4ff7-ba24-6875cfa99f39
# ╠═3fa412f2-d107-4caa-abf6-bf47c8333027
# ╠═ff16b0cb-b8ba-4d3a-bcc8-da5c120ea898
# ╠═0c108a6b-599a-4581-8d8d-68b9ae5865df
# ╠═2d7d48ac-c2e1-4acd-8e21-3a9314e19390
# ╠═d1660777-a3a2-4abb-a394-4872da699c4e
# ╟─49a583a5-ef6b-4e5b-8ea2-0e9a5088b3cc
# ╠═64af243c-77f9-4c87-8f59-a4e506e3cb46
# ╠═5610563c-ea7f-40c6-ba8c-5a816bc71a2e
# ╟─37f43e35-015a-4596-96b0-13ccb720c767
# ╠═bef2f5e8-e6ef-49df-9125-1d814feb31d7
# ╠═cf18b86d-6c82-4335-b46c-60cd161504b5
# ╠═2cdd244d-33cb-487e-9fa2-069f8a845ef9
# ╠═8e5f5bf2-5b93-4c52-9960-ea1fb60888ce
# ╠═15626656-af92-44d3-83dd-9775172caed8
# ╠═a9db1ffb-cc86-4d78-8f24-8b847443aae8
# ╠═e144fe4f-cb8e-4158-a59a-8a93faca1f0c
# ╠═b34bb590-0fd4-491f-a392-a538d16a86b6
# ╟─1fc3e23d-e283-43a1-a645-cfbaa3821870
# ╠═2fc20959-426a-410f-900b-bdde889426fd
# ╠═e0f00df4-7c63-448c-b005-eda95cb4cdb8
# ╠═7d95732e-94f4-448c-8e19-7969271059a9
# ╠═86a8bf70-3c82-4100-8b63-5d010523a196
# ╠═3641d445-1fff-416d-849b-f9d5715fef0f
# ╠═cc688fb0-0da2-4d5c-a0e8-bcff8af05ad1
# ╟─e559b4b6-e26c-446e-b67e-b196e8ac3021
# ╠═6b966b4c-4b5c-4d62-a842-efe204b1c0fb
# ╠═37c5f515-74a3-4bd0-b7e6-6e498b375e23
# ╟─cfe15c1b-7b29-4514-82c4-e6ffe86456b0
# ╠═814d4261-c433-47a9-a35a-16d83be3d4f4
# ╠═3e30b081-e4df-4518-98aa-96e6e0880a4a
# ╠═d45bcd94-6adf-40eb-baee-e8edab966f07
# ╟─eb005f06-1b2f-45fe-b6b2-abfc6b4dbb1b
# ╠═a01fab74-b77b-4b6f-8415-b0f60530593b
# ╠═15dc3591-8141-4e8d-9778-8290de7627ad
# ╠═5ef8d15b-b407-4f18-8ed3-b0d4c69b6345
# ╠═de0edf52-cf26-4265-8f10-d8720dfee53f
# ╟─5e32a3b3-30a9-4cfe-8005-b2eebfd5a4b4
# ╠═b94655ba-382e-402e-a0c0-89d9d4d91bb6
# ╠═3ee0c65f-02b2-41b9-a95b-9932203223c8
# ╟─84a06e2e-13b9-44ea-a1e3-809791017656
# ╠═0d4763cc-f47a-4de0-9cde-12a73556b5e3
# ╟─15b6196d-0b87-45fd-8303-770969386d10
# ╠═4aeceb65-f226-41cf-988b-273f42c75059
# ╠═75852dd4-6cc9-4040-84e2-cdb0ce812ec4
# ╠═bfa26eb6-d9ea-4a5d-80a5-84f048ed3ea0
# ╠═23cc65b9-76d8-43c9-a398-22212b00c76a
# ╠═88b9c633-01ef-4bf2-9180-0bb2b59f4766
# ╠═170b5e91-dfaa-4f13-ac45-01b875ba02d8
# ╠═19e2a53c-7716-4fd3-8716-87586c60ec73
# ╟─497f6e35-4b9a-4349-8366-9f54aa1476c5
# ╠═71510434-8685-49ee-b514-31d1c7c411d8
# ╠═195bb3ce-5ade-4c31-9385-12cc4c119762
# ╠═8d2c5280-d9c1-414d-a64c-8d8ef5a88b10
# ╠═abf9ed24-ea46-4067-b786-3a6b23be4c80
# ╟─e3d0bb6e-e9c9-4a2d-97bf-73041958771b
# ╠═f4b0f15f-016e-41ba-a999-e9f792a534fc
# ╠═db517950-7031-4495-8198-f054eb833e16
# ╠═9a09fa1a-8045-4e93-9140-7f57915146ac
# ╠═ea3b65bb-2e41-4c17-9df8-35eaf69975f6


include("header.jl")

#=
# Manometry
In this notebook, we will learn
* How to use manometers to measure pressure
* How to combine several working fluids in a manometer
* How to calculate for tilted manometers
=#

# ### Set up the module
using MAE103
#-
using Plots

#=
In fluid dynamics, we are often interested in measuring pressure differences
between two points, often because those pressure differences allow us
to determine flow speed.

We'll discuss how we get flow speed from pressure difference later. For now, our
goal is to show how we can measure pressure difference between two
points. The key notion here is that we only need to measure pressure *difference*,
not absolute pressure. Also, one of those "points" might simply be the atmosphere,
where the pressure has some ambient value.

Manometers enable this by turning pressure difference into a measurable height of a
*working fluid* in the manometer. The manometer itself consists of a narrow tube,
bent into a U-shaped form and arranged vertically so that gravity and fluid statics can
be exploited. Sometimes the manometer has additional feature, e.g.,

- one of the channels in the tube might be at an angle rather than vertical
- one or more of the channels might be connected to a "cistern", with a
much wider diameter than the tube.

The working fluid is generally a **liquid**, and we can assume that it does not
change its volume during operation. This fact will be important in the analysis.
=#
#-
#=
### Example 1
We will demonstrate here how we use the fluid statics equation to determine pressure
by reading a measured height. Consider the example setup shown in the figure below:

#

We seek to measure the gauge pressure (i.e., relative to ambient) in the air in the tank
by measuring the heights of liquids. Suppose that the oil has specific gravity 0.9. Also,
mercury (Hg) has a very high specific gravity, 13.6: it is a liquid metal.

The heights shown are $h_1 = 36$ inches, $h_2 = 6$ inches, and $h_3 = 9$ inches.
Let's find the gauge pressure in the air, in psi.

First set the values:
=#
h1 = Height(36u"inch")
h2 = Height(6u"inch")
h3 = Height(9u"inch")
γoil = SpecificWeight(SpecificGravity(0.9))
γHg = SpecificWeight(SpecificGravity(13.6))
patm = Pressure(1u"atm");
#=
We work backward from the open end, where pressure is atmospheric. Just to be careful,
we will add this atmospheric pressure, so we are working with absolute pressures,
but then subtract it at the end to find our answer. At point 2, the pressure
increases by $γ_{\mathrm{Hg}}h_3$ from atmospheric:
=#
p2 = Pressure(patm + γHg*h3)
#=
Point 1 on the mercury side is at the same elevation as point 2, and in the same
fluid, so its pressures are equal.
=#
p1 = p2
#=
Also $p_1$ in oil is the same as in mercury. And this pressure is larger than
the pressure in the air by $\gamma_{\mathrm{oil}}(h_1 + h_2)$, so $p_{\mathrm{air}}$
=#
p_air = Pressure(p1 - γoil*(h1+h2))
#=
Remember, this was the absolute pressure. The gauge pressure is obtained by
subtracting the atmospheric pressure:
=#
p_airg = Pressure(p_air-patm)

#=
Finally, subtract the convert to psi
=#
value(p_airg,u"psi")

#=
In summary, the answer was

$$p_{\mathrm{air},g} = γ_{\mathrm{Hg}}h_3 - \gamma_{\mathrm{oil}}(h_1 + h_2)$$
=#

#=
In this example, the manometer used mercury. What if it had used water instead?
What would $h_3$ be (assuming the same gauge pressure in the air)? This is
given by

$$h_3 = \left(p_{\mathrm{air},g} + \gamma_{\mathrm{oil}}(h_1 + h_2)\right)/γ_w$$

where $γ_w$ is the specific weight of water. Let's find out what $h_3$ would be
(in feet).
=#
h3 = Height((p_airg + γoil*(h1+h2))/SpecificWeight(Water))
value(h3,u"ft")

#=
10.2 feet instead of 9 inches! There is big advantage in using a heavy working
liquid in a manometer, because it does not rise as high. It's easier to
get readings.
=#

#=
### Example 2
In the previous example, we pointed out that mercury is useful for translating
a large pressure into a relatively small height change. Sometimes we need the
opposite: if we want to measure a relatively small pressure difference,
we want it to be easy to read on the manometer. For this purpose, we can use an
inclined tube manometer, as in the setup shown in the figure below:

#

We read the length $l_2$ along the manometer tube. How does this length come into the
manometer equation? Only the vertical part of $l_2$ enters this equation, i.e.,
$l_2 \sin \theta$

The equation between $A$ and $B$ is

$$p_A + \gamma_1 h_1 - \gamma_2 l_2 \sin \theta - \gamma_3 h_3 = p_B$$

(From $A$, we drop down by $h_1$ with fluid 1, rise by $l_2 \sin \theta$ with fluid 2,
 then rise by $h_3$ with fluid 3.)

Let's assume that fluids 1 and 3 are gases, so their specific weights are negligible.
Then

$$p_B - p_A = \gamma_2 l_2 \sin \theta$$

or, in other words,

$$l_2 = \dfrac{p_B - p_A}{\gamma_2\sin\theta}$$

Let's try this out for a few different angles of inclined tube. We will assume that the
pressure difference is 100 Pa --- a pretty small pressure difference --- and the
working fluid (fluid 2) is water.
=#
dp = Pressure(100)
γ = SpecificWeight(Water)
#=
We'll start with a vertical tube (90 degrees)
=#
Θ = 90π/180
l = Length(dp/γ/sin(Θ))
value(l,u"cm")
#=
Now a tube at 45 degrees
=#
Θ = 45π/180
l = Length(dp/γ/sin(Θ))
value(l,u"cm")
#=
and finally a tube at 10 degrees
=#
Θ = 10π/180
l = Length(dp/γ/sin(Θ))
value(l,u"cm")
#=
Notice how the length goes up as the tube becomes more and more horizontal.
This makes it much easier to read the manometer without being prone to error.
=#


include("../examples/bubblesort.jl")
a = rand(1:100,100)
@test issorted(sort!(a;alg=BubbleSort))

include("../examples/enum.jl")
@enum TestEnum TestEnum1 TestEnum2 TestEnum3
@test [TestEnum1.n,TestEnum2.n,TestEnum3.n] == [0,1,2]

include("../examples/lru.jl")
include("../examples/lru_test.jl")

include("../examples/modint.jl")
b = ModInts.ModInt{10}(2)
c = ModInts.ModInt{10}(4)
@test b + c == ModInts.ModInt{10}(6)
@test c - b == ModInts.ModInt{10}(2)

include("../examples/ndgrid.jl")
r = repmat(1:10,1,10)
r1, r2 = ndgrid(1:10, 1:10)
@test r1 == r
@test r2 == r'
r3, r4 = meshgrid(1:10,1:10)
@test r3 == r'
@test r4 == r

include("../examples/quaternion.jl")
q = Quaternions.Quaternion(1,0,0,0)
x = Quaternions.Quaternion(0,1,1,1)
@test q*2.0+2 == Quaternions.Quaternion(4,0,0,0)
@test abs((-q+x*2)/4) == 0.9013878188659973

include("../examples/queens.jl")
@test solve(8, 8, 1) == Array{Int,1}[[1,1]]
@test solve(8, 8, 7) == Array{Int,1}[[1,1],[2,3],[3,5],[4,2],[5,8],[7,4],[8,7]]

# At least make sure code laods
include("../examples/plife.jl")

include("../examples/preduce.jl")

include("../examples/wordcount.jl")


"""
    update_partial_residuals!(r, y, x, perm, b, k)

A subroutine to compute the partial residuals `r = Y - X*b` in-place based on a permutation vector `perm` that indexes the nonzeroes in `b`.
"""
function update_partial_residuals!{T <: Float}(
    r    :: DenseVector{T},
    y    :: DenseVector{T},
    x    :: DenseMatrix{T},
    perm :: DenseArray{Int,1},
    b    :: DenseVector{T},
    k    :: Int
)
    copy!(r, y)
    @inbounds for j = 1:k
        idx = perm[j]
        @inbounds @simd for i in eachindex(y)
            r[i] += -b[idx]*x[i,idx]
        end
    end
    return nothing
end

"""
    exchange_leastsq!(v, x, y, k) -> b

This function minimizes the residual sum of squares

    0.5*sumabs2( y - x*b )

subject to `b` having no more than `k` nonzero components. The function will compute a `b` for a given value of `k`.
For optimal accuracy, this function should be run for multiple values of `k` over a path of model sizes.
For optimal performance, allocate `v` once and then reuse it between different runs of `exchange_leastsq!`.

Arguments:

- `v` is the `ELSQVariables` object housing all temporary arrays, including `b`
- `x` is the n x p statistical design matrix.
- `y` is the n-dimensional response vector.
- `k` is the desired number of nonzero components in `b`.

Optional Arguments:

- `n` and `p` are the dimensions of `x`; the former defaults to `length(y)` while the latter defaults to `size(x,2)`.
- `window` is an `Int` to dictate the dimension of the search window for potentially exchanging predictors.
   Defaults to `k`, which tells `exchange_leastsq!` to search all active predictors.
   Decreasing this quantity tells the algorithm to search through fewer current active predictors,
   which can decrease compute time but can also degrade model recovery performance.
- `max_iter` is the maximum permissible number of iterations. Defaults to `100`.
- `tol` is the convergence tolerance. Defaults to `1e-6`.
- `quiet` is a `Bool` to control output. Defaults to `false` (full output).
"""
function exchange_leastsq!{T <: Float}(
    v        :: ELSQVariables{T},
    x        :: DenseMatrix{T},
    y        :: DenseVector{T},
    k        :: Int;
    n        :: Int  = length(y),
    p        :: Int  = size(x,2),
    window   :: Int  = k,         # <--- search *all* active predictors
    max_iter :: Int  = 100,
    tol      :: T    = convert(T, 1e-6),
    quiet    :: Bool = true
)
    # initial value for previous residual sum of squares
    old_rss = oftype(tol,Inf)

    # obtain top r components of bvec in magnitude
    selectperm!(v.perm, v.b, k, by=abs, rev=true, initialized=true)

    # compute partial residuals based on top r components of perm vector
    update_partial_residuals!(v.r, y, x, v.perm, v.b, k)

    # save value of RSS before starting algorithm
    rss = sumabs2(v.r) / 2

    # compute inner products of X and residuals
    # this is basically the negative gradient
    BLAS.gemv!('T', one(T), x, v.r, zero(T), v.df)

    # outer loop controls number of total iterations for algorithm run on one r
    for iter = 1:(max_iter)

        # output algorithm progress to console
        quiet || println("\titer = ", iter, ", RSS = ", rss)

        # middle loop tests each of top r parameters (by magnitude?)
        for i = abs(k-window+1):k

            # save information for current value of i
            l     = v.perm[i]
            betal = v.b[l]
            #update_col!(v.tempn, x, l)  # tempn now holds X[:,l]
            copy!(v.tempn, view(x, :, l))

            # if necessary, compute inner products of current predictor against all other predictors
            # store this information in Dict inner
            # for current index, hold dot products in memory for duration of inner loop
            get_inner_product!(v.dotprods, v.tempn, v, x, l)

            # subroutine compares current predictor i against all predictors k+1, k+2, ..., p
            # these predictors are candidates for inclusion in set
            # _exlstsq_innerloop! find best new predictor r
            a, b, r, adb = _exlstsq_innerloop!(v, k, i, p, tol)

            # now want to update residuals with current best predictor
            m = update_current_best_predictor!(v, x, betal, adb, r)

            # if necessary, compute inner product of current predictor against all other predictors
            # save in our Dict for future reference
            get_inner_product!(v.tempp, v.tempn2, v, x, m)

            # also update df
            axpymbz!(v.df, betal, v.dotprods, adb, v.tempp)

            # now swap best predictor with current predictor
            _swap_predictors!(v, i, r, m, adb)

        end # end middle loop over predictors

        # update residual sum of squares
        rss = sumabs2(v.r) / 2

        # test for descent failure
        # if no descent failure, then test for convergence
        # if not converged, then save RSS and check finiteness
        # if not converged and still finite, then save RSS and continue
        ascent    = rss > old_rss + tol
        converged = abs(old_rss - rss) / abs(old_rss + 1) < tol
        old_rss   = rss
        ascent && print_descent_error(iter, old_rss, rss)
        converged && return nothing

        check_finiteness(rss)

    end # end outer iteration loop

    # at this point, maximum iterations reached
    # warn and return
    print_maxiter(max_iter, rss)
    return nothing

end # end exchange_leastsq!

"""
    exlstsq(x, y) -> ELSQResults

For a statistical model `b`, `exlstsq` minimizes the residual sum of squares

    0.5*sumabs2( y - x*b )

subject to `b` having no more than `k` nonzero components. `exlstsq` will compute a `b` for several model sizes `k`.

Arguments:

- `x` is the n x p statistical design matrix.
- `y` is the n-dimensional response vector.

Optional Arguments:

- `v` is the `ELSQVariables` object housing all temporary arrays, including `b`.
- `models` is the integer vector of model sizes to test. It defaults to `collect(1:p)`, where `p = min(20, size(x,2))`.
- `window` is an `Int` to dictate the maximum size of the search window for potentially exchanging predictors.
   Defaults to `max(20, min(maximum(models), size(x,2)))` or the greater of `20` and `r`, where `r` is the lesser of the maximum model size and the number of predictors.
   Decreasing this quantity tells the algorithm to search through fewer current active predictors,
   which can decrease compute time but can also degrade model recovery performance.
- `max_iter` is the maximum permissible number of iterations. Defaults to `100`.
- `tol` is the convergence tolerance. Defaults to `1e-6`.
- `quiet` is a `Bool` to control output. Defaults to `false` (full output).
"""
function exlstsq{T <: Float}(
    x        :: DenseMatrix{T},
    y        :: DenseVector{T};
    v        :: ELSQVariables{T} = ELSQVariables(x, y),
    models   :: DenseVector{Int} = collect(1:min(20, size(x,2))),
    window   :: Int  = max(20, min(maximum(models), size(x,2))),
    max_iter :: Int  = 100,
    tol      :: T    = convert(T, 1e-6),
    quiet    :: Bool = true
)
    # dimensions of problem
    nmodels = length(models)
    n,p = size(x)

    # error checking
    errorcheck(x, y, tol, max_iter, window, p)

    # initialize sparse matrix to return
    betas = spzeros(T, p, nmodels)

    # loop through models
    for i in models

        # monitor output
        quiet || println("Testing model size $i.")
        exchange_leastsq!(v, x, y, i, window=min(window,i), max_iter=max_iter, tol=tol, quiet=quiet, n=n, p=p)
        betas[:,i] = sparse(sdata(v.b))
    end

    # return matrix of betas
    return betas
end


# Tests for Von-Mises distribution

using Distributions
using Test

function test_vonmises(μ::Float64, κ::Float64)
    d = VonMises(μ, κ)
    @test length(d) == 1
    @test mean(d) == μ
    @test median(d) == μ
    @test mode(d) == μ
    @test d == typeof(d)(params(d)...,d.I0κx)
    @test d == deepcopy(d)
    @test partype(d) == Float64
    # println(d)

    # conversions
    @test typeof(convert(VonMises{Float32}, d)) == VonMises{Float32}

    # Support
    @test support(d) == RealInterval(d.μ-π,d.μ+π)
    @test pdf(d, d.μ-2π) == 0.0
    @test pdf(d, d.μ+2π) == 0.0

end

## General testing

for (μ, κ) in [(2.0, 1.0),
               (2.0, 5.0),
               (3.0, 1.0),
               (3.0, 5.0),
               (5.0, 2.0)]

    test_vonmises(μ, κ)
end


#=
   Ruin problem.

   Cf ruin_problem.jl

   Compared to ruin_problem.jl this version uses a recursive approach
   with concatenated arrays instead.

   The arrays always ends with a 0 (ruin).
   Note: We have to have an maxLen since otherwise the sequence goes to
   infinity...

   Summary Statistics
   parameters      mean       std   naive_se      mcse          ess      rhat   ess_per_sec
       Symbol   Float64   Float64    Float64   Float64      Float64   Float64       Float64

          len   11.0188   15.7735     0.1577    0.1505    9765.3342    1.0001    30233.2330
         ruin    0.8925    0.3098     0.0031    0.0028   10015.3143    0.9999    31007.1649

   Distributions of variable len (num:0)
   2.00000 =>    4999  (0.499900)
   4.00000 =>    1320  (0.132000)
   50.00000 =>    1097  (0.109700)
   6.00000 =>     603  (0.060300)
   8.00000 =>     393  (0.039300)
   10.00000 =>     265  (0.026500)
   12.00000 =>     204  (0.020400)
   14.00000 =>     148  (0.014800)
   16.00000 =>     126  (0.012600)
   18.00000 =>     126  (0.012600)
   20.00000 =>      99  (0.009900)
   22.00000 =>      87  (0.008700)
   24.00000 =>      65  (0.006500)
   30.00000 =>      58  (0.005800)
   26.00000 =>      56  (0.005600)
   28.00000 =>      51  (0.005100)
   36.00000 =>      48  (0.004800)
   32.00000 =>      42  (0.004200)
   38.00000 =>      41  (0.004100)
   34.00000 =>      39  (0.003900)
   40.00000 =>      30  (0.003000)
   46.00000 =>      29  (0.002900)
   42.00000 =>      28  (0.002800)
   44.00000 =>      27  (0.002700)
   48.00000 =>      19  (0.001900)

   Distributions of variable ruin (num:0)
   1.00000 =>    8925  (0.892500)
   0.00000 =>    1075  (0.107500)

.  Cf ~/webppl/ruin_problem2.wppl

=#
using Turing, StatsPlots
include("jl_utils.jl")

@model function ruin_problem2(maxLen=50,start=1,win=1,loose=1)
    head = 1
    tail = 2
    coins = [head, tail]

    function draw(arr)
        if length(arr) == 0
            return draw([start])
        elseif arr[end] == 0 || length(arr) >= maxLen
            return arr
        else
            lastVal = arr[end]
            c = rand(coins)
            if c == head
                return draw(push!(arr,lastVal+win))
            else
                if lastVal - loose <= 0
                    return push!(arr,0)
                else
                    return draw(push!(arr,lastVal-loose))
                end
            end
        end
    end

    a = draw([])
    len ~ Dirac(length(a))
    ruin ~ Dirac(a[end] == 0)

    # Use PG(5) when using these. Otherwise Prior() works and is faster.
    # true ~ Dirac(len == 6)
    # true ~ Dirac(a[end] ==8)

end

maxLen=50
start=1
win=1
loose=1
model = ruin_problem2(maxLen,start,win,loose)

chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
# chns = sample(model, PG(5), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 10_000)

display(chns)
# display(plot(chns))

show_var_dist_pct(chns, :len)
show_var_dist_pct(chns, :ruin)

# chains_params = Turing.MCMCChains.get_sections(chns, :parameters)
# genq = generated_quantities(model, chains_params)
# show_var_dist_pct(genq)


# using Revise
using Test, BifurcationKit, LinearAlgebra, SparseArrays, Setfield, Parameters, ForwardDiff
const BK = BifurcationKit

f1(u, v) = u^2 * v
norminf = x -> norm(x, Inf)

function Fbru!(f, x, p)
	@unpack α, β, D1, D2, l = p
	n = div(length(x), 2)
	h = 1.0 / n; h2 = h*h
	c1 = D1 / l^2 / h2
	c2 = D2 / l^2 / h2

	u = @view x[1:n]
	v = @view x[n+1:2n]

	# Dirichlet boundary conditions
	f[1]   = c1 * (α	  - 2u[1] + u[2] ) + α - (β + 1) * u[1] + f1(u[1], v[1])
	f[end] = c2 * (v[n-1] - 2v[n] + β / α)			 + β * u[n] - f1(u[n], v[n])

	f[n]   = c1 * (u[n-1] - 2u[n] +  α  )  + α - (β + 1) * u[n] + f1(u[n], v[n])
	f[n+1] = c2 * (β / α  - 2v[1] + v[2])			 + β * u[1] - f1(u[1], v[1])

	for i=2:n-1
		  f[i] = c1 * (u[i-1] - 2u[i] + u[i+1]) + α - (β + 1) * u[i] + f1(u[i], v[i])
		f[n+i] = c2 * (v[i-1] - 2v[i] + v[i+1])			  + β * u[i] - f1(u[i], v[i])
	end
	return f
end

Fbru(x, p) = Fbru!(similar(x), x, p)

function Jbru_sp(x, p)
	@unpack α, β, D1, D2, l = p
	# compute the Jacobian using a sparse representation
	n = div(length(x), 2)
	h = 1.0 / n; h2 = h*h

	c1 = D1 / p.l^2 / h2
	c2 = D2 / p.l^2 / h2

	u = @view x[1:n]
	v = @view x[n+1:2n]

	diag   = zeros(eltype(x), 2n)
	diagp1 = zeros(eltype(x), 2n-1)
	diagm1 = zeros(eltype(x), 2n-1)

	diagpn = zeros(eltype(x), n)
	diagmn = zeros(eltype(x), n)

	@. diagmn = β - 2 * u * v
	@. diagm1[1:n-1] = c1
	@. diagm1[n+1:end] = c2

	@. diag[1:n]    = -2c1 - (β + 1) + 2 * u * v
	@. diag[n+1:2n] = -2c2 - u * u

	@. diagp1[1:n-1] = c1
	@. diagp1[n+1:end] = c2

	@. diagpn = u * u
	J = spdiagm(0 => diag, 1 => diagp1, -1 => diagm1, n => diagpn, -n => diagmn)
	return J
end

Jbru_ana(x, p) = ForwardDiff.jacobian(z->Fbru(z,p),x)
jet = BK.get3Jet(Fbru, Jbru_ana)

n = 100
par_bru = (α = 2., β = 5.45, D1 = 0.008, D2 = 0.004, l = 0.3)
	sol0 = vcat(par_bru.α * ones(n), par_bru.β/par_bru.α * ones(n))

# test that the jacobian is well computed
@test Jbru_sp(sol0, par_bru) - Jbru_ana(sol0, par_bru) |> sparse |> nnz == 0

opt_newton = NewtonPar(tol = 1e-11, verbose = false)
	out, hist, flag = newton(jet[1], jet[2], sol0 .* (1 .+ 0.01rand(2n)), par_bru, opt_newton)

opts_br0 = ContinuationPar(dsmin = 0.001, dsmax = 0.1, ds= 0.01, pMax = 1.8, newtonOptions = opt_newton, detectBifurcation = 3, nev = 16, nInversion = 4)
	br, = continuation(jet[1], jet[2], out, (@set par_bru.l = 0.3), (@lens _.l), opts_br0, printSolution = (x, p) -> norm(x, Inf64))
###################################################################################################
# Hopf continuation with automatic procedure
outhopf, = newton(jet[1], jet[2], br, 1; startWithEigen = true, d2F = jet[3])
outhopf, = newtonHopf(jet[1], jet[2], br, 1; startWithEigen = true)
optconthopf = ContinuationPar(dsmin = 0.001, dsmax = 0.15, ds= 0.01, pMax = 6.8, pMin = 0., newtonOptions = opt_newton, maxSteps = 50, detectBifurcation = 2)
outhopfco, = continuationHopf(jet[1], jet[2], br, 1, (@lens _.β), optconthopf; startWithEigen = true, updateMinAugEveryStep = 1, plot = false)

# Continuation of the Hopf Point using Dense method
ind_hopf = 1
# av = randn(Complex{Float64},2n); av = av./norm(av)
# bv = randn(Complex{Float64},2n); bv = bv./norm(bv)
hopfpt = BK.HopfPoint(br, ind_hopf)
bifpt = br.specialpoint[ind_hopf]
hopfvariable = HopfProblemMinimallyAugmented(
					jet[1], jet[2], nothing, nothing,
					(@lens _.l),
					conj.(br.eig[bifpt.idx].eigenvec[:, bifpt.ind_ev]),
					(br.eig[bifpt.idx].eigenvec[:, bifpt.ind_ev]),
					# av,
					# bv,
					opts_br0.newtonOptions.linsolver)

hopfvariable(hopfpt, par_bru) |> norm

Bd2Vec(x) = vcat(x.u, x.p)
Vec2Bd(x) = BorderedArray(x[1:end-2], x[end-1:end])
hopfpbVec(x, p) = Bd2Vec(hopfvariable(Vec2Bd(x),p))

# finite differences Jacobian
Jac_hopf_fdMA(u0, p) = ForwardDiff.jacobian( u-> hopfpbVec(u, p), u0)
# ``analytical'' jacobian
Jac_hopf_MA(u0, p, pb::HopfProblemMinimallyAugmented) = (return (x=u0,params=p ,hopfpb=pb))

rhs = rand(length(hopfpt))
jac_hopf_fd = Jac_hopf_fdMA(Bd2Vec(hopfpt), par_bru)
sol_fd = jac_hopf_fd \ rhs

# create a linear solver
hopfls = BK.HopfLinearSolverMinAug()
tmpVecforσ = zeros(ComplexF64, 2+2n)
sol_ma,  = hopfls(Jac_hopf_MA(hopfpt, par_bru, hopfvariable), BorderedArray(rhs[1:end-2],rhs[end-1:end]), debugArray = tmpVecforσ)

# we test the expression for σp
σp_fd = Complex(jac_hopf_fd[end-1,end-1], jac_hopf_fd[end,end-1])
σp_fd_ana = tmpVecforσ[1]
@test σp_fd ≈ σp_fd_ana rtol = 1e-5

# we test the expression for σω
σω_fd = Complex(jac_hopf_fd[end-1,end], jac_hopf_fd[end,end])
σω_fd_ana = tmpVecforσ[2]
@test σω_fd ≈ σω_fd_ana rtol = 1e-4

# we test the expression for σx
σx_fd = jac_hopf_fd[end-1, 1:end-2] + Complex(0,1) * jac_hopf_fd[end, 1:end-2]
σx_fd_ana = tmpVecforσ[3:end]
@test σx_fd ≈ σx_fd_ana rtol = 1e-3

outhopf, hist, flag = newton(
		jet[1], jet[2], br, 1, Jᵗ = (x, p) -> transpose(Jbru_sp(x, p)))
		# flag && printstyled(color=:red, "--> We found a Hopf Point at l = ", outhopf.p[1], ", ω = ", outhopf.p[end], ", from l = ",hopfpt.p[1],"\n")

outhopf, _, flag, _ = newton((u, p) -> hopfvariable(u, p),
							(x, p) -> Jac_hopf_MA(x, p, hopfvariable),
							hopfpt, par_bru, NewtonPar(verbose = false, linsolver = BK.HopfLinearSolverMinAug()))
	# flag && printstyled(color=:red, "--> We found a Hopf Point at l = ", outhopf.p[1], ", ω = ", outhopf.p[2], " from ", hopfpt.p, "\n")

# version with analytical Hessian = 2 P(du2) P(du1) QU + 2 PU P(du1) Q(du2) + 2PU P(du2) Q(du1)
function d2F(x, p1, du1, du2)
	n = div(length(x),2)
	out = similar(du1)
	@views out[1:n] .= 2 .* x[n+1:end] .* du1[1:n] .* du2[1:n] .+
				2 .* x[1:n] .* du1[1:n] .* du2[n+1:end] .+
				2 .* x[1:n] .* du2[1:n] .* du1[1:n]
	@inbounds for ii=1:n
		out[ii+n] = -out[ii]
	end
	return out
end

outhopf, hist, flag = newton(jet[1], jet[2], br, 1,
		Jᵗ = (x, p) -> transpose(Jbru_sp(x, p)),
		d2F = (x, p1, v1, v2) -> d2F(x, 0., v1, v2))
		# flag && printstyled(color=:red, "--> We found a Hopf Point at l = ", outhopf.p[1], ", ω = ", outhopf.p[end], ", from l = ",hopfpt.p[1],"\n")

br_hopf, u1_hopf = continuation(
			jet[1], jet[2], br, ind_hopf, (@lens _.β),
			ContinuationPar(dsmin = 0.001, dsmax = 0.05, ds= 0.01, pMax = 6.5, pMin = 0.0, a = 2., theta = 0.4, maxSteps = 3, newtonOptions = NewtonPar(verbose = false)), verbosity = 0, plot = false)

br_hopf, u1_hopf = continuation(jet[1], jet[2], br, ind_hopf, (@lens _.β), ContinuationPar(dsmin = 0.001, dsmax = 0.05, ds= 0.01, pMax = 6.5, pMin = 0.0, a = 2., theta = 0.4, maxSteps = 3, newtonOptions = NewtonPar(verbose = false)), Jᵗ = (x, p) ->  transpose(Jbru_sp(x, p)), d2F = (x, p1, v1, v2) -> d2F(x, 0., v1, v2), verbosity = 0, plot = false)
#################################################################################################### Continuation of Periodic Orbit
ind_hopf = 1
hopfpt = BK.HopfPoint(br, ind_hopf)

l_hopf = hopfpt.p[1]
ωH	   = hopfpt.p[2] |> abs
M = 20

orbitguess = zeros(2n, M)
phase = []; scalphase = []
vec_hopf = geteigenvector(opt_newton.eigsolver, br.eig[br.specialpoint[ind_hopf].idx][2], br.specialpoint[ind_hopf].ind_ev-1)
for ii=1:M
	t = (ii-1)/(M-1)
	orbitguess[:, ii] .= real.(hopfpt.u + 26*0.1 * vec_hopf * exp(-2pi * complex(0, 1) * (t - .252)))
end

orbitguess_f = vcat(vec(orbitguess), 2pi/ωH) |> vec

# test guess using function
l_hopf, Th, orbitguess2, hopfpt, vec_hopf = BK.guessFromHopf(br, ind_hopf, opt_newton.eigsolver, M, 2.6; phase = 0.252)

poTrap = PeriodicOrbitTrapProblem(jet[1], Jbru_sp, real.(vec_hopf), hopfpt.u, M, 2n	)

jac_PO_fd = BK.finiteDifferences(x -> poTrap(x, (@set par_bru.l = l_hopf + 0.01)), orbitguess_f)
jac_PO_sp = poTrap(Val(:JacFullSparse), orbitguess_f, (@set par_bru.l = l_hopf + 0.01))

# test of the Jacobian for PeriodicOrbit via Finite differences VS the FD associated jacobian
# test jacobian expression for Periodic Orbit solve problem
@test norm(jac_PO_fd - jac_PO_sp, Inf64) < 1e-4

# test various jacobians and methods
jac_PO_sp =  poTrap(Val(:BlockDiagSparse), orbitguess_f, (@set par_bru.l = l_hopf + 0.01))
BK.getTimeDiff(poTrap, orbitguess_f)
# BK.Jc(poTrap, reshape(orbitguess_f[1:end-1], 2n, M), par_bru, reshape(orbitguess_f[1:end-1], 2n, M))
# BK.Jc(poTrap, orbitguess_f, par_bru, orbitguess_f)

# newton to find Periodic orbit
opt_po = NewtonPar(tol = 1e-8, verbose = false, maxIter = 150)
	outpo_f, _, flag = @time newton(
		(x, p) ->  poTrap(x, p),
		(x, p) ->  poTrap(Val(:JacFullSparse),x,p),
		copy(orbitguess_f), (@set par_bru.l = l_hopf + 0.01), opt_po)
	# println("--> T = ", outpo_f[end])
# flag && printstyled(color=:red, "--> T = ", outpo_f[end], ", amplitude = ", BK.amplitude(outpo_f, n, M; ratio = 2),"\n")

newton(poTrap, orbitguess_f, (@set par_bru.l = l_hopf + 0.01), opt_po; linearPO = :FullLU)

# jacobian of the functional
Jpo2 = poTrap(Val(:JacCyclicSparse), orbitguess_f, (@set par_bru.l = l_hopf + 0.01))

# calcul des exposants de Floquet, extract full vector
BK.MonodromyQaD(Val(:ExtractEigenVector), poTrap, orbitguess_f, par_bru, orbitguess_f[1:2n])

# calcul des exposants de Floquet
floquetES = FloquetQaD(DefaultEig())

# continuation of periodic orbits using :BorderedLU linear algorithm
opts_po_cont = ContinuationPar(dsmin = 0.0001, dsmax = 0.05, ds= 0.001, pMax = 2.3, maxSteps = 3, theta = 0.1, newtonOptions = NewtonPar(verbose = false), detectBifurcation = 1)
	br_pok2, = continuation(
		poTrap, orbitguess_f, (@set par_bru.l = l_hopf + 0.01), (@lens _.l), opts_po_cont; linearPO = :BorderedLU,
		plot = false, verbosity = 0)

# test of simple calls to newton / continuation
deflationOp = DeflationOperator(2.0, (x,y) -> dot(x[1:end-1], y[1:end-1]),1.0, [zero(orbitguess_f)])
# opt_po = NewtonPar(tol = 1e-8, verbose = false, maxIter = 15)
opts_po_cont = ContinuationPar(dsmin = 0.001, dsmax = 0.03, ds= 0.01, pMax = 3.0, maxSteps = 3, newtonOptions = (@set opt_po.verbose = false), nev = 2, precisionStability = 1e-8, detectBifurcation = 1)
for linalgo in [:FullLU, :BorderedLU, :FullSparseInplace]
	@show linalgo
	# with deflation
	@time newton(poTrap,
			copy(orbitguess_f), (@set par_bru.l = l_hopf + 0.01), opt_po, deflationOp; linalgo = linalgo, normN = norminf)
	# classic Newton-Krylov
	outpo_f, hist, flag = @time newton(poTrap,
			copy(orbitguess_f), (@set par_bru.l = l_hopf + 0.01), opt_po; linalgo = linalgo, normN = norminf)
	# continuation
	br_pok2, = @time continuation(poTrap,
			copy(orbitguess_f), (@set par_bru.l = l_hopf + 0.01), (@lens _.l),
			opts_po_cont; linearPO = linalgo, verbosity = 0,
			plot = false, normC = norminf)
end

# test of Matrix free computation of Floquet exponents
eil = EigKrylovKit(x₀ = rand(2n))
ls = GMRESKrylovKit()
ls = DefaultLS()
opt_po = NewtonPar(tol = 1e-8, verbose = false, maxIter = 10, linsolver = ls, eigsolver = eil)
opts_po_cont = ContinuationPar(dsmin = 0.001, dsmax = 0.03, ds= 0.01, pMax = 3.0, maxSteps = 3, newtonOptions = (@set opt_po.verbose = false), nev = 2, precisionStability = 1e-8, detectBifurcation = 2)
br_pok2, upo, = continuation(poTrap, outpo_f, (@set par_bru.l = l_hopf + 0.01), (@lens _.l), opts_po_cont; linearPO = :FullLU, normC = norminf, verbosity = 0)


## Exercise 4-12
## Read about spirals at https://en.wikipedia.org/wiki/Spiral; then write a program that draws an Archimedan spiral as in Archimedan spiral.
## https://benlauwens.github.io/ThinkJulia.jl/latest/book.html#fig04-4
println("Ans : ")

println("End.")


mutable struct Board
	wonAt::Int32
	rows::Vector{Vector{Int32}}
	columns::Vector{Vector{Int32}}
end

input = readlines("input.txt")

boards = Vector{Board}()
for boardDefinition in Iterators.partition(input[3:end], 6)
	board = Board(0, Vector{Vector{Int32}}(), Vector{Vector{Int32}}())
	for row in boardDefinition[1:end-1]
		push!(board.rows, map(x -> parse(Int32, x), split(strip(row), r"\s+")))
	end
	board.columns = [x[:] for x in eachrow(hcat(board.rows...))]
	push!(boards, board)
end

winners = Vector{Board}()
for number in map(x -> parse(Int32, x), split(input[1], ","))
	for board in boards
		f(x) = x != number
		board.rows = [filter(f, x) for x in board.rows]
		board.columns = [filter(f, x) for x in board.columns]

		e(x) = size(x, 1) == 0
		if findfirst(e, board.rows) != nothing || findfirst(e, board.columns) != nothing
			board.wonAt = number
			push!(winners, board)
		end
	end
	deleteat!(boards, findall(x -> x in winners, boards))
end

function score(board)
	return sum([sum(x) for x in board.rows]) * board.wonAt
end

println("The answer to the first part is: ", score(winners[1]))
println("The answer to the second part is: ", score(winners[end]))