using SubstitutionTilings
using SubstitutionTilings.Penrose
import SubstitutionTilings.Penrose: ψ
using Test
using BenchmarkTools
using Plots

using Luxor

const L = Penrose.Qζ

setantialias(6)

#collars = SubstitutionTilings.CoreDefs.collars(penrose(), 6);
#length(collars)

width = 1280
height = 800
sc = 200
@draw begin
    colors = ["#DD93FC", "#E7977A", "#9B70AF", "#A0644F",]
    first_tile = hkite()
    tiling = substitute(penrose(), [first_tile], 1, Penrose.in_bounds, (w=width/sc, h=height/sc))
    println(typeof(tiling))
    setline(1)

    for tile in tiling
        origin()
        draw(tile, sc, colors[Penrose.color(tile)], :fill)
        origin()
        setcolor("black")
        scale(sc)
        #circle(embed_center(tile[1]), 0.1, :fill)
        sethue("black")
    end
    origin()
    setline(1)
    #draw(first_tile, sc/0.618, "black", :stroke)
end width height #"penrose.png"

width = 1280*2
height = 800*2
sc = 400
@draw begin
    colors = ["#DD93FC", "#E7977A", "#9B70AF", "#A0644F",]

    quadrants = Tiler(width, height, 1, 2, margin=5)

    begin
        pos = quadrants[1][1]
        first_tile = hkite()
        tiling = substitute(penrose(), [first_tile], 0, Penrose.in_bounds, (w=width/sc, h=height/sc))
        println(typeof(tiling))
        setline(1)

        for tile in tiling
            origin()
            Luxor.translate(pos)
            draw(tile, sc, colors[Penrose.color(tile)], :fill)
        end
        origin()
        Luxor.translate(pos)
        scale(sc)
        setline(2)
        setcolor("black")
        Luxor.line(Penrose.embed_nf_p(ζ^4),Penrose.embed_nf_p(L(1)), :stroke)
        Luxor.arrow(Penrose.embed_nf_p(ζ^4),Penrose.embed_nf_p(ζ^4+(L(1)-ζ^4)*(7//16)), arrowheadlength=0.1)
        Luxor.arrow(Penrose.embed_nf_p(ζ^4),Penrose.embed_nf_p(ζ^4+(L(1)-ζ^4)*(9//16)), arrowheadlength=0.1)
        Luxor.line(Penrose.embed_nf_p(ζ^6),Penrose.embed_nf_p(L(1)), :stroke)
        Luxor.arrow(Penrose.embed_nf_p(L(1)),Penrose.embed_nf_p(ζ^6+(L(1)-ζ^6)//2), arrowheadlength=0.1)
        Luxor.arrow(Penrose.embed_nf_p(L(1)),Penrose.embed_nf_p(ζ^6+(L(1)-ζ^6)*(3//8)), arrowheadlength=0.1)
        Luxor.arrow(Penrose.embed_nf_p(L(1)),Penrose.embed_nf_p(ζ^6+(L(1)-ζ^6)*(5//8)), arrowheadlength=0.1)
        Luxor.line(Penrose.embed_nf_p(ζ^4),Penrose.embed_nf_p(ζ^6), :stroke)
        Luxor.arrow(Penrose.embed_nf_p(ζ^6),Penrose.embed_nf_p(ζ^4+(ζ^6-ζ^4)*(1//2)), arrowheadlength=0.1)
    end

    begin
        pos = quadrants[2][1]
        first_tile = hdart()
        tiling = substitute(penrose(), [first_tile], 0, Penrose.in_bounds, (w=width/sc, h=height/sc))
        println(typeof(tiling))
        setline(1)

        for tile in tiling
            origin()
            Luxor.translate(pos)
            draw(tile, sc, colors[Penrose.color(tile)], :fill)
        end
        origin()
        Luxor.translate(pos)
        scale(sc)
        setline(2)
        setcolor("black")
        Luxor.line(Penrose.embed_nf_p(ζ^2-ζ^0//ϕ),Penrose.embed_nf_p(L(1-ζ^0//ϕ)), :stroke)
        Luxor.arrow(Penrose.embed_nf_p(ζ^2-ζ^0//ϕ),Penrose.embed_nf_p(ζ^2-ζ^0//ϕ+(L(1-ζ^0//ϕ)-(ζ^2-ζ^0//ϕ))//2), arrowheadlength=0.1)
        Luxor.line(Penrose.embed_nf_p(L(1-ζ^0//ϕ)),Penrose.embed_nf_p(ζ^8-ζ^0//ϕ), :stroke)
        Luxor.arrow(Penrose.embed_nf_p(L(1-ζ^0//ϕ)),Penrose.embed_nf_p(1-ζ^0//ϕ+((ζ^8-ζ^0//ϕ)-(1-ζ^0//ϕ))*(3//16)), arrowheadlength=0.1)
        Luxor.arrow(Penrose.embed_nf_p(L(1-ζ^0//ϕ)),Penrose.embed_nf_p(1-ζ^0//ϕ+((ζ^8-ζ^0//ϕ)-(1-ζ^0//ϕ))*(5//16)), arrowheadlength=0.1)
        Luxor.arrow(Penrose.embed_nf_p(L(1-ζ^0//ϕ)),Penrose.embed_nf_p(1-ζ^0//ϕ+((ζ^8-ζ^0//ϕ)-(1-ζ^0//ϕ))*(7//16)), arrowheadlength=0.1)
        Luxor.arrow(Penrose.embed_nf_p(L(1-ζ^0//ϕ)),Penrose.embed_nf_p(1-ζ^0//ϕ+((ζ^8-ζ^0//ϕ)-(1-ζ^0//ϕ))*(9//16)), arrowheadlength=0.1)
        Luxor.line(Penrose.embed_nf_p(ζ^8-ζ^0//ϕ),Penrose.embed_nf_p(ζ^2-ζ^0//ϕ), :stroke)
        Luxor.arrow(Penrose.embed_nf_p(ζ^8-ζ^0//ϕ),Penrose.embed_nf_p(ζ^8-ζ^0//ϕ-(ζ^8-ζ^2)*(7//16)), arrowheadlength=0.1)
        Luxor.arrow(Penrose.embed_nf_p(ζ^8-ζ^0//ϕ),Penrose.embed_nf_p(ζ^8-ζ^0//ϕ-(ζ^8-ζ^2)*(9//16)), arrowheadlength=0.1)

   end
end width height #"penrose.png"

pentagon = (([
    PenroseElem(mod(k+s,10),s,L(1))*PenroseElem(0,0,L(-1)) => Penrose.Hkite
    for k=0:2:10 for s=0:1
]));

width = 300
height = 300
sc = 50
@png begin
    colors = ["#DD93FC", "#E7977A", "#9B70AF", "#A0644F",]
    first_tile = hkite(0, false, L(0))
    tiling = pentagon

    for tile in tiling
        origin()
        draw(tile, sc, colors[Penrose.color(tile)], :fill)
    end
    origin()
    setline(1)
    draw(first_tile, sc, "black", :stroke)
end width height "pentagon.png"

trapezoid = Dict([
    hkite(0,0,L(0)),
    hdart(2,1, ζ^8//ϕ )])
width = 300
height = 300
sc = 50
@png begin
    colors = ["#DD93FC", "#E7977A", "#9B70AF", "#A0644F",]
    first_tile = hkite(0, false, L(0))
    tiling = trapezoid

    for tile in tiling
        origin()
        draw(tile, sc, colors[Penrose.color(tile)], :fill)
    end
    origin()
    setline(1)
    draw(first_tile, sc, "black", :stroke)
end width height "trapezoid.png"

# Comparison of Array and Dict operations
@assert substitute(penrose(), Dict([hkite(0,0,L(0))]), 20, Penrose.in_bounds, (w=50, h=50)) == Dict(substitute(penrose(), ([hkite(0,0,L(0))]), 20, Penrose.in_bounds, (w=50, h=50)))

t_dict = zeros(35)
n_dict = zeros(35)
for i in 1:35
    p = substitute(penrose(), Dict([hkite(0,0,L(0))]), 20, Penrose.in_bounds, (w=i*2, h=i*2));
    time = @timed empirical_frequency(pentagon, Dict(p))
    t_dict[i] = time.time
    n_dict[i] = length(p)
end # essentially linear in the amont of tiles: pretty good!

t_arr = zeros(35)
n_arr = zeros(35)
for i in 1:35
    p = substitute(penrose(), ([hkite(0,0,L(0))]), 20, Penrose.in_bounds, (w=i*2, h=i*2));
    time = @timed empirical_frequency(pentagon, (p))
    t_arr[i] = time.time
    n_arr[i] = length(p)
end
plot(n_dict,t_dict)
plot!(n_arr, t_arr)
plot((n_dict), (t_dict))
plot(log.(n_arr), log.(t_arr))
# Frequency counting with arrays is exponential in the amount of tiles
# with Dicts it is more like linear

tiling = substitute(penrose(), Dict([hkite(0,0,L(0))]), 20, Penrose.in_bounds, (w=50, h=50))
empirical_frequency(pentagon, tiling)

Penrose.frequency(pentagon, 4)
@testset "Frequencies" begin
    @test Penrose.frequency([hkite()], 4) == ψ
    @test Penrose.frequency([hdart()], 4) == 1-ψ
    @test Penrose.frequency(pentagon, 4) == 7*ψ - 4
    @test Penrose.frequency(pentagon, 5) == 7*ψ - 4
    @test Penrose.frequency(pentagon, 6) == 7*ψ - 4
end


using Dates
using Printf

struct Clock

    hr::Integer
    min::Integer

    function Clock(hr::Integer, min::Integer)

        # Add/subtract extra minutes
        if min % 60 == 0
            hr += div(min, 60)
            min = 0
        elseif min < 0
            hr += div(min, 60) - 1
            min = min % 60 + 60
        else
            hr += div(min, 60)
            min = min % 60
        end

        if hr % 24 == 0
            hr = 0
        elseif hr < 0
            hr = hr % 24 + 24
        else
            hr = hr % 24
        end

        new(hr, min)
    end
end

function Base.:+(clock::Clock, time::Dates.Minute)
    delta_mins = time.value
    return Clock(clock.hr, clock.min + delta_mins)
end

function Base.:-(clock::Clock, time::Dates.Minute)
    delta_mins = time.value
    return Clock(clock.hr, clock.min - delta_mins)
end

function Base.show(io::IO, clock::Clock)
    @printf(io, "%2d:%2d", clock.hr, clock.min)
end


using thesis, PGFPlotsX, DifferentialEquations, DifferentialEquations.EnsembleAnalysis, Statistics, ProgressMeter
setup_pgfplotsx()

## Helper functions
"""the expected firing rate of an lIF neuron with threshold `δ` and time-constant `α` in response to a constant input `x`"""
r_est_lif(x,δ,α)::Float64 = x ≤ α*δ ? 0.0 : -α/log(1-δ*α/x)

"""the expected firing rate of an LNP neuron with gain `λ` and time-constant `α` in response to a constant input `x`"""
r_est_lnp(x,λ,α)::Float64 = x ≤ 0.0 ? 0.0 : λ*x

"""the expected membrane potential of the decoding lIF neuron for an encoding lIF neuron with threshold `δ` and time-constant `α` in response to a constant input `x`"""
z_est_lif(x,δ,α) = δ*r_est_lif(x,δ,α)

"""the expected threshold of an lIF neuron with time-constant `α` that leads to a firing rate `r` in response to a constant input `x`"""
δ_est(x,r, α) = x/α*(1-exp(-α/r))

"""the expected gain of an LNP neuron with time-constant `α` that leads to a firing rate `r` in response to a constant input `x`"""
λ_est(x,r, α) = r/x

"""the RMSE expected for a ternary value in a finite `range` """
RMSE_ternary(range) = √(1/12)*range/3

"""the RMSE expected for a binary-encoded value in a finite `range` with fixed number of `bits` """
RMSE_digital(bits, range) = √(1/12)*range*2.0^(-bits)

"""the RMSE expected for an lIF neuron with threshold δ and time-constant α """
RMSE_lif(δ, α) = √(1/12)*α*δ

"""the threshold δ expected for an lIF neuron with time-constant α required for a fixed `RMSE`"""
RMSE_lif_inverse(RMSE, α) = √(12)*RMSE/α

"""the RMSE expected for an LNP neuron with gain `λ` and time-constant α """
RMSE_lnp(λ, α) = √(α/(2*λ))

"""the gain `λ` expected for an LNP neuron with time-constant α required for a fixed `RMSE`"""
RMSE_lnp_inverse(RMSE, α) = α/(2*RMSE^2)

"""the estimated rate at which the expected RMSE will equal that of an `nbits` encoding """
r_nbit_lif(s₀, nbits, α) = r_est_lif(s₀, RMSE_lif_inverse(RMSE_digital(nbits, 1.0), α), α)

"""the estimated rate at which the expected RMSE will equal that of a ternary encoding """
r_ternary_lif(α) = r_est_lif(s₀, RMSE_lif_inverse(RMSE_ternary(1.0), α), α)

## Configuration types

abstract type SNNConfig end;

struct lIFConfig{A,D,II,W,IW} <: SNNConfig
    num_neurons::Int
    α::A
    δ::D
    input_neurons::II
    input_weights::IW
    weights::W
    can_spike::Vector{Bool}
end

struct LNPConfig{A,L,II,W,IW} <: SNNConfig
    num_neurons::Int
    α::A
    λ₀::L
    input_neurons::II
    input_weights::IW
    weights::W
    can_spike::Vector{Bool}
end

## ODE and problem definition for lIF and LNP neurons

"""dynamics of leaky neurons (lIF and LNP) with exponential filter"""
function f′(du, u, p, t)
    du .= -p.cfg.α*u

    # work-around for jump problems (broadcasting doesnt work)
    tmp = p.cfg.input_weights .* p.input(t)
    for idx ∈ p.cfg.input_neurons
        du[idx] += tmp[idx]
    end
end

function lIF_spike_condition(out, u, t, integrator)
    out .= u .- integrator.p.cfg.δ
    # ignore neurons that cant spike
    out[(~).(can_spike)] .= Inf
end

function lIF_spike!(integrator, idx)
    integrator.u .+= integrator.p.cfg.weights[:,idx]

    has_spiked = zeros(Bool, integrator.p.cfg.num_neurons)
    has_spiked[idx] = true
    push!(integrator.p.spike_train, (integrator.t, idx))
    while true
        new_spikes = findall(integrator.p.cfg.can_spike .* (~).(has_spiked) .* (integrator.u .>= integrator.p.cfg.δ))
        if isempty(new_spikes)
            break
        end
        for i ∈ new_spikes
            integrator.u .+= integrator.p.cfg.weights[:,i]
            has_spiked[i] = true
            push!(integrator.p.spike_train, (integrator.t, i))
        end
    end
    integrator.u[has_spiked] .= 0.0
end

function make_LNP_jump(i)
    function LNP_rate(u,p,t)
        return u[i]*p.cfg.λ₀[i]
    end

    function LNP_spike!(integrator)
        integrator.u.u .+= integrator.p.cfg.weights[:,i]
        push!(integrator.p.spike_train, (integrator.t, i))
    end

    VariableRateJump(LNP_rate, LNP_spike!)
end

"""run an ensemble of lIF encoder-decoder-pair simulations over timespan `tspan` for a given `input` (function of time) and an lIFConfig `p`
Returns a solution object that tracks, for each trajectory, the decoding neuron's membrane potential (first trace) and spike rate (second trace)
"""
function runEnsemble(tspan, input, p::lIFConfig; trajectories=100, kwargs...)
    p.weights[2,1]=p.δ[1] * p.α

    spike_trains = [Tuple{Float64,Int}[] for i ∈ 1:trajectories]

    prob_func(prob,i,repeat) = ODEProblem(f′, rand(p.num_neurons).*p.δ, tspan, (spike_train=spike_trains[i], cfg=p, input=input))
    prob = prob_func(nothing,1,0)
    cb = VectorContinuousCallback(lIF_spike_condition, lIF_spike!, nothing, p.num_neurons)
    ensemble_prob = EnsembleProblem(
        prob,
        prob_func=prob_func;
        output_func=(sol,i)->(
            [
                sol(tspan[2])[2],
                length(filter(((t,n),)->(t>(tspan[2]+tspan[1])/2), spike_trains[i]))
            ],
            false
        )
    )
    sol = solve(ensemble_prob, Tsit5(), scheduler(); callback=cb, trajectories=trajectories, kwargs...)

    return hcat(sol...)'
end

"""run an ensemble of LNP encoder-decoder-pair simulations over timespan `tspan` for a given `input` (function of time) and an LNPConfig `p`
Returns a solution object that tracks, for each trajectory, the decoding neuron's membrane potential (first trace) and spike rate (second trace)
"""
function runEnsemble(tspan, input, p::LNPConfig; trajectories=100, kwargs...)
    p.weights[2,1]=p.α / p.λ₀[1]

    spike_trains = [Tuple{Float64,Int}[] for i ∈ 1:trajectories]

    prob_func(prob,i,repeat) = JumpProblem(
        ODEProblem(f′, rand(p.num_neurons), tspan, (spike_train=spike_trains[i], cfg=p, input=input)),
        Direct(),
        JumpSet((make_LNP_jump(j) for j ∈ 1:p.num_neurons if p.can_spike[j])...)
    )
    prob = prob_func(nothing,1,0)
    ensemble_prob = EnsembleProblem(
        prob,
        prob_func=prob_func;
        output_func=(sol,i)->(
            [
                sol(tspan[2])[2],
                length(filter(((t,n),)->(t>(tspan[2]+tspan[1])/2), spike_trains[i]))
            ],
            false
        )
    )
    sol = solve(ensemble_prob, Tsit5(), scheduler(); trajectories=trajectories, kwargs...)

    return hcat(sol...)'
end

"""simulates an ensemble of encoder-decoder pairs of neurons (lIF and LNP) for constant inputs with different `signal_levels`.
returns the `signal_levels`, means `μ` of lIF and LNP neurons and RMSE `err` of lIF and LNP neurons
"""
function calculate_transfer(tspan, signal_levels, lIF_cfg, LNP_cfg; kwargs...)
    μ = (lif=Float64[],lnp=Float64[])
    err = (lif=(neg=Float64[], pos=Float64[]),lnp=(neg=Float64[], pos=Float64[]))
    @showprogress 1.0 "Calculating transfer functions..." for level ∈ signal_levels
        for (i,cfg) ∈ enumerate((lIF_cfg, LNP_cfg))
            sol=runEnsemble(tspan, t->[level], cfg; kwargs...)

            quant = quantile(sol[:,1],[0.25,0.75])
            m = mean(sol[:,1])

            push!(err[i].neg,m-quant[1])
            push!(μ[i],m)
            push!(err[i].pos,quant[2]-m)
        end
    end
    return (signal_levels=signal_levels, μ=μ, err=err)
end

"""simulates an ensemble of encoder-decoder pairs of neurons (lIF and LNP) as well as a digital encoding for an `input` function.
returns digital encoding as well as decodings of lIF and LNP neurons
"""
function compare_digital_lif_lnp(times, input, (bits, rate), lIF_cfg, LNP_cfg; kwargs...)
    tspan = (minimum(times),maximum(times))
    spike_trains = (digital=Tuple{Float64,Int}[], lif=Tuple{Float64,Int}[], lnp=Tuple{Float64,Int}[])
    sample_points = minimum(times):1.0/rate:maximum(times)
    samples = input.(sample_points)

    digital = map(t->round(samples[searchsortedlast(sample_points, t)]*2^bits)/2^bits, times)

    prob = ODEProblem(f′, rand(lIF_cfg.num_neurons).*lIF_cfg.δ, tspan, (spike_train=spike_trains.lif, cfg=lIF_cfg, input=input))
    cb = VectorContinuousCallback(lIF_spike_condition, lIF_spike!, nothing, lIF_cfg.num_neurons)
    sol_lIF = solve(prob, Tsit5(); callback=cb, kwargs...)

    prob = JumpProblem(
        ODEProblem(f′, rand(LNP_cfg.num_neurons), tspan, (spike_train=spike_trains.lnp, cfg=LNP_cfg, input=input)),
        Direct(),
        JumpSet((make_LNP_jump(i) for i ∈ 1:LNP_cfg.num_neurons if LNP_cfg.can_spike[i])...)
    )
    sol_LNP = solve(prob, Tsit5(); kwargs...)

    return (solution=(digital=digital, lif=sol_lIF(times, idxs=2)[:].+0.5*lIF_cfg.δ[1]*lIF_cfg.α, lnp=sol_LNP(times, idxs=2)[:]), spikes=spike_trains)
end

"""simulates an ensemble of encoder-decoder pairs of neurons (lIF and LNP) for a fixed `input` function with different `scales` of firing rates (adjust δ or λ accordingly).
returns the `scales`, mean firing rates `rate` of lIF and LNP neurons and RMSE `rmse` of lIF and LNP neurons
"""
function calculate_error_vs_spikes!(tspan, scales, input, lIF_cfg, LNP_cfg; kwargs...)
    rate = (lif=Float64[],lnp=Float64[])
    rmse = (lif=Float64[],lnp=Float64[])
    @showprogress 1.0 "Calculating error/rate relation..." for s ∈ scales
        for (i,cfg) ∈ enumerate((lIF_cfg,LNP_cfg))
            if i==1
                cfg.δ[1] = δ_est(input(0)[1],s,cfg.α)
            else
                cfg.λ₀[1] = λ_est(input(0)[1],s,cfg.α)
            end

            sol=runEnsemble(tspan, input, cfg; kwargs...)

            push!(rate[i], mean(sol[:,2]./(tspan[2]-tspan[1])*2))

            v = sol[:,1]

            if i==1
                v.+= cfg.α.*cfg.δ[1]./2
            end
            push!(rmse[i], √(mean((v .- input(tspan[2])).^2)))
        end
    end

    return (scales=scales, rate=rate, rmse=rmse)
end

## Set up
use_multithreading = false                           # whether to use multithreading or not
num_neurons = 2                                     # we always simulate 2 neurons here
α=1000.0/25.0                                       # α = 40 corresponds to a time-scale of 25ms
δ=[δ_est(0.5,r_nbit_lif(0.5,4,α),α), 2.0]               # δ = 0.0015625 gives similar RMSE at a signal value 0.5 as a 4-bit encoding would
λ₀=1.0 ./ δ                                         # λ = 1/δ results in the same mean firing rate as for the lIF neuron
can_spike = [true,false]                            # only the encoding neuron can spike
input_neurons=[1]                                   # only the encoding neuron receives input
input_weights_lIF=[1.0]                             # the input-weights for the lIF neuron are 1
input_weights_LNP=input_weights_lIF.*α              # the input-weights for the LNP neuron are scaled up by α
weights=[0 0; 1 0].*δ'.*α                           # the weight from encoder to decode is scaled by αδ
trajectories = 50                                  # we simulate an ensemble of multiple trajectories to determine the distribution
signal_levels = LinRange(-0.2,1.0, 100)             # we evaluate the input-output function of the neurons for constant input signals with various levels
scales = 10.0 .^ LinRange(0, 4, 10)                 # we evaluate the RMSE for different parameterizations by sweeping over multiple scales (changes either δ or λ, respectively)
s₀ = 0.5                                            # we use a constant signal s(t)=0.5 to evaluate different parameterizations
equivalent_sampling_rate = α/2                      # a rectangular filter with equivalent delay as the exponential filter with time-constant α has width 2/α -> rate α/2
reference_bits = [2,4,6,8]                          # plot refrence lines for digital encodings using different numbers of bits
tspan = (0.0,10.0)                                  # we simulate for 10 seconds and take the last half of that for calculating averages
sine(t) = 0.4*sin(0.8*π*t)+0.5                      # we visualize the encoding for a sine-wave
plot_times = LinRange(-1,2.5,1000)                  # we plot the sine wave in the interval 0 to 10

inset_xmin,inset_xmax = 0.3,0.5
inset_ymin,inset_ymax = 0.7,1.0
plot_times_zoom = LinRange(inset_xmin,inset_xmax,200)

# the decoded signal will approximate the filtered version of the input signal
#∫ᵗsin(ωτ)α exp(-α(t-τ))dτ
#   = α exp(-αt)/(2i)∫ᵗ(exp((α+iω)τ)-exp((α-iω)τ))dτ
#   = α exp(-αt)/(2i)[1/(α+iω)exp((α+iω)t)-1/(α-iω)exp((α-iω)t)]
#   = α 1/(2iα-2ω)exp(iωt)-1/(2iα+2ω)exp(-iωt)
#   = α (α sin(ωt) - ω cos(ωt))/(α² + ω²)
sine_filtered(t) = 0.4*(α*sin(0.8*π*t)-0.8*π*cos(0.8*π*t)) * α/(α^2+(0.8π)^2)+0.5

##

scheduler = use_multithreading ? EnsembleThreads : EnsembleSerial
lIF_cfg = lIFConfig(num_neurons, α, copy(δ), input_neurons, copy(input_weights_lIF), copy(weights), can_spike)
LNP_cfg = LNPConfig(num_neurons, α, copy(λ₀), input_neurons, copy(input_weights_LNP), copy(weights), can_spike)

# Calculate the input-output transfer function
(s, μ, err) = calculate_transfer(tspan, signal_levels, lIF_cfg, LNP_cfg; trajectories=trajectories)

# Calculate input and decoding of a sine-wave
(decoded,spikes) = compare_digital_lif_lnp(plot_times, sine, (4, equivalent_sampling_rate), lIF_cfg, LNP_cfg; trajectories=trajectories)

# Calculate the RMSE/firing rate relationship as a function of varying δ or λ
(scales, rate, rmse) = calculate_error_vs_spikes!(tspan, scales, t->[s₀], lIF_cfg, LNP_cfg; trajectories=trajectories)

ss = 10 .^ LinRange(log10(minimum(scales)), log10(maximum(scales)), 200)

reference_rmse = [RMSE_digital(bits, 1.0) for bits in reference_bits]
reference_rate = [r_nbit_lif(s₀, bits,α) for bits in reference_bits]
reference_text = ["\\approx $(bits)-bit @ \$$(equivalent_sampling_rate)\$Hz " for bits in reference_bits]

t0=searchsortedfirst(plot_times,0)

## Plot the monstrosity
@pgf gp = TikzPicture(
    {remember_picture},
    GroupPlot(
    {
        group_style={
            group_size="2 by 3",
            horizontal_sep="2cm",
            vertical_sep="2cm",
        },
        height="6cm",
        width="8cm",
    },
    {xlabel=raw"input value $c$", ylabel=raw"rate$(c)$", title=raw"decoded rate", legend_pos="north west", xmin=signal_levels[1], xmax=signal_levels[end]},
    Plot(
        {"name path=lnplow", draw="none", forget_plot},
        Table(signal_levels, μ.lnp .- err.lnp.neg)
    ),
    Plot(
        {"name path=lnphigh", draw="none", forget_plot},
        Table(signal_levels, μ.lnp .+ err.lnp.pos)
    ),
    PlotInc(
        {color="Set1-B", opacity=0.5, forget_plot},
        raw"fill between [of=lnplow and lnphigh]"
    ),
    PlotInc(
        {color="Set1-B", very_thick},
        Table(signal_levels, μ.lnp)
    ),
    LegendEntry("LNP response"),
    Plot(
        {"name path=liflow", draw="none", forget_plot},
        Table(signal_levels, μ.lif .- err.lif.neg)
    ),
    Plot(
        {"name path=lifhigh", draw="none", forget_plot},
        Table(signal_levels, μ.lif .+ err.lif.pos)
    ),
    PlotInc(
        {color="Set1-A", forget_plot},
        raw"fill between [of=liflow and lifhigh]"
    ),
    PlotInc(
        {color="Set1-A", very_thick},
        Table(signal_levels, μ.lif)
    ),
    LegendEntry("lIF response"),
    PlotInc(
        {dashed, very_thick},
        Table(signal_levels, max.(0,signal_levels.- 0.5*α*δ[1]))
    ),
    LegendEntry(raw"shifted ReLU"),
    PlotInc(
        {dotted, very_thick},
        Table(signal_levels, z_est_lif.(signal_levels, δ[1], α))
    ),
    LegendEntry(raw"expected lIF response"),
    {xlabel=raw"firing rate $r$ [Hz]", ylabel=raw"RMSE$(r)$", title="RMSE for constant input \$x(t) = $(s₀)\$", ymode="log", xmode="log",xmin=ss[1],xmax=ss[end]},
    (Plot(
        {forget_plot, thick, color= i==2 ? "Set1-C" : "black!30"},
        Coordinates([(ss[1],e),(ss[end],e)])
    ) for (i,e) ∈ enumerate(reference_rmse))...,
    PlotInc(
        {forget_plot, very_thick, color="Set1-B"},
        Table(ss, RMSE_lnp.(λ_est.(s₀,ss,α),α))
    ),
    PlotInc(
        {only_marks, very_thick, color="Set1-B"},
        Table(rate.lnp, rmse.lnp)
    ),
    LegendEntry(raw"LNP neuron"),
    PlotInc(
        {forget_plot, very_thick, color="Set1-A"},
        Table(ss, RMSE_lif.(δ_est.(s₀,ss,α),α))
    ),
    PlotInc(
        {only_marks, very_thick, color="Set1-A"},
        Table(rate.lif, rmse.lif)
    ),
    LegendEntry(raw"lIF neuron"),
    Plot(
        {only_marks, mark="+", color="black"},
        Table(reference_rate, reference_rmse)
    ),
    ["\\node[anchor=north east, color=$(i==2 ? "Set1-C" : "black!50")] at (axis cs:$(r),$(e)) {$(t)};" for (i,(r,e,t)) ∈ enumerate(zip(reference_rate, reference_rmse, reference_text))],
    {width="16.5cm", xshift="4.25cm", title="digital code vs. lIF neuron vs. LNP neuron", ylabel="decoded signal", xticklabels="\\empty", xmin=0,xmax=2.5,ymin=0, ymax=2.0},
    "\\coordinate (boxsw) at (axis cs:$(inset_xmin),$(inset_ymin));",
    "\\coordinate (boxse) at (axis cs:$(inset_xmax),$(inset_ymin));",
    "\\coordinate (boxne) at (axis cs:$(inset_xmax),$(inset_ymax));",
    "\\coordinate (boxnw) at (axis cs:$(inset_xmin),$(inset_ymax));",
    "\\filldraw[color=black!90, very thick, fill=black!10] (boxsw) rectangle (boxne);",
    PlotInc(
        {very_thick, color="Set1-B"},
        Table(plot_times,decoded.lnp)
    ),
    LegendEntry("LNP decoding ($(length(filter(>(0), first.(spikes.lnp)))) spikes)"),
    PlotInc(
        {very_thick, color="Set1-A"},
        Table(plot_times,decoded.lif)
    ),
    LegendEntry("lIF decoding ($(length(filter(>(0), first.(spikes.lif)))) spikes)"),
    PlotInc(
        {very_thick},
        Table(plot_times,decoded.digital)
    ),
    LegendEntry("digital decoding ($(sum(sum.(thesis.discretize.(decoded.digital, Ref(LinRange(0,1,2^4+1)), Ref(digits.(0:2^4,base=2,pad=4)))[t0:end])))) bits)"),
    Plot(
        {ultra_thick, dashed},
        Table(plot_times,sine_filtered.(plot_times))
    ),
    LegendEntry(raw"target signal $(s*\kappa_\alpha)(t)$"),
    raw"\coordinate (insetpos) at (axis cs:0.1,1.1);",
    {"group/empty plot", "axis background/.style={fill=none}", xmin=0,xmax=1,ymin=0,ymax=1},
    {yshift="1.75cm", width="16.5cm", height="3cm", xlabel="time [s]", xmin=0,xmax=2.5,ymin=0, ymax=2.0, ylabel="spikes", ytick="\\empty"},
    (Plot(
        {color="Set1-A", forget_plot, opacity=0.5},
        Coordinates([(t,0.05),(t,0.95)])
    ) for t ∈ first.(spikes.lif))...,
    (Plot(
        {color="Set1-B", forget_plot, opacity=0.5},
        Coordinates([(t,1.05),(t,1.95)])
    ) for t ∈ first.(spikes.lnp))...,
    ),
    Axis(
        {very_thick, height="1.75cm", width="5cm", scale_only_axis, "axis background/.style"={fill="black!10", opacity=0.9},  xmin=inset_xmin,xmax=inset_xmax, ymin=inset_ymin, ymax=inset_ymax, at={"(insetpos)"}, anchor={"south west"}, ticks="none"},
        PlotInc(
            {ultra_thick, color="Set1-B"},
            Table(plot_times,decoded.lnp)
        ),
        PlotInc(
            {ultra_thick, color="Set1-A"},
            Table(plot_times,decoded.lif)
        ),
        PlotInc(
            {ultra_thick},
            Table(plot_times,decoded.digital)
        ),
        Plot(
            {ultra_thick, dashed},
            Table(plot_times,sine_filtered.(plot_times))
        ),
        "\\coordinate (insw) at (axis cs:$(inset_xmin),$(inset_ymin));",
        "\\coordinate (inse) at (axis cs:$(inset_xmax),$(inset_ymin));",
        "\\coordinate (inne) at (axis cs:$(inset_xmax),$(inset_ymax));",
        "\\coordinate (innw) at (axis cs:$(inset_xmin),$(inset_ymax));",
    ),
    raw"\draw[color=black!75, very thick] (boxsw) -- (insw);",
    raw"\draw[color=black!75, very thick] (boxse) -- (inse);",
)
#pgfsave("fig/encoding_comparison.pdf",gp)


module DynamicForestBehaviors

using AutomotiveDrivingModels

using RandomForestBehaviors.MvDecisionTrees
using Distributions.MvNormal

import AutomotiveDrivingModels:
    AbstractVehicleBehavior,
    AbstractVehicleBehaviorTrainParams,
    AbstractVehicleBehaviorPreallocatedData,

    preallocate_learning_data,
    select_action,
    calc_action_loglikelihood,
    train

export
    DynamicForestBehavior,
    DF_TrainParams,
    AutoregressiveMvNormLeaf,

    select_action,
    calc_action_loglikelihood,
    train

global GLOBAL_DIAG_SHRINKAGE_THRESOLD = 1e-6
DEFAULT_AUTOREGRESSION_CONSTANT = 0.5
DEFAULT_NUM_AUTOREGRESSION_PREDICTORS = 3
DEFAULT_NUM_PREDICTOR_SAMPLES = 10

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

immutable AutoregressiveMvNormLeaf <: MvDecisionTrees.Leaf
    A::Matrix{Float64} # regression matrix A = UΦᵀ(ΦΦᵀ + γI)⁻¹ where γ is a ridge regularization param, U is column-wise targets, Φ is column-wise features
    m::MvNormal        # μ = Aϕ must be overwritten on each use
    predictor_indeces::Vector{Int} # indeces of predictors
end

function _regress_on_predictor_subset{T<:AbstractFloat, S<:Real}(
    data::MvDecisionTrees.TreeData{T,S},
    assignment_id::Int,
    Γ::Matrix{T}, # ϕ+1 × ϕ+1 diagonal γ matrix (preallocated)
    U::Matrix{T}, #   o × m
    Φ::Matrix{T}, # ϕ+1 × m (is overwritten)
    A::Matrix{T}, #   o × ϕ+1 (is overwritten)
    predictor_indeces::Vector{Int}, # length ϕ
    )

    #=
    Returns the loss when using this predictor set
    Overwrites Φ
    =#

    n = size(data.X,1)
    o, m = size(U)
    ϕ = length(predictor_indeces)
    @assert(size(Φ,1) == ϕ+1)
    @assert(size(Φ,2) == m)
    @assert(size(Γ,1) == size(Γ,2) == ϕ+1 )

    u_index = 0
    for i in 1 : n
        if data.assignment[i] == assignment_id
            u_index += 1
            for (j,k) in enumerate(predictor_indeces)
                Φ[j,u_index] = data.X[i,k] # features : [nrows, nfeatures], Φ: [nfeatures, nrows]
            end
            Φ[ϕ+1,u_index] = 1.0 # constant 1-predictor
        end
    end

    # regress

    A[:] = U*Φ'/(Φ*Φ' + Γ)

    # calc residual loss (square deviation)

    loss = 0.0
    for i in 1 : n
        if data.assignment[i] == assignment_id
            for j in 1 : o
                r = data.Y[j,i] - A[j,ϕ+1]
                for k in 1 : ϕ
                    r -= A[j,k]*data.X[i,predictor_indeces[k]]
                end
                loss += r*r # squared error
            end
        end
    end

    loss
end
function MvDecisionTrees.build_leaf{T<:AbstractFloat, S<:Real}(::Type{AutoregressiveMvNormLeaf},
    data::MvDecisionTrees.TreeData{T,S},
    assignment_id::Int,
    γ::T=DEFAULT_AUTOREGRESSION_CONSTANT, # ridge-regularization parameter
    n_autoregression_predictors::Int = DEFAULT_NUM_AUTOREGRESSION_PREDICTORS,
    n_random_predictor_samples::Int = DEFAULT_NUM_PREDICTOR_SAMPLES,
    )

    #=
    Fit the regression matrix A = UΦᵀ(ΦΦᵀ + γI)⁻¹
       γ is a ridge regularization param ≥ 0
         0 → least squares
         large γ → increased penalty on regression coefficients
       U is column-wise targets
       Φ is column-wise features plus a constant 1-feature

    This function picks a set of 'n_autoregression_predictors' predictors used to regress on the targets
    This is done 'n_random_predictor_samples' times
    The best in terms of residual error is kept
    The covariance matrix is static and computed based on the final residual errors
    =#

    @assert(n_autoregression_predictors ≥ 0)
    @assert(n_random_predictor_samples > 0)

    n, o, p = size(data)
    id_count = MvDecisionTrees.get_id_count(data, assignment_id)
    @assert(id_count > 0)

    ##################################
    # pull target matrix

    U = Array(T, o, id_count) # column-wise contatenation of target variables

    u_index = 0
    for i in 1 : n
        if data.assignment[i] == assignment_id
            u_index += 1
            for j in 1 : o
                U[j,u_index] = data.Y[j,i]
            end
        end
    end

    #################################
    # pull predictor matrix

    ϕ = min(n_autoregression_predictors, p) # number of real features
    Φ = Array(T, ϕ+1, id_count) # column-wise contatenation of predictors (also including constant 1-predictor)
    Γ = diagm(fill(γ, (ϕ+1)))
    A = Array(T, o, ϕ+1)
    predictor_indeces = Array(Int, ϕ)

    if ϕ == 0
        _regress_on_predictor_subset(data, assignment_id, Γ, U, Φ, A, predictor_indeces) # just solve for A
    else
        # pick several samples of predictor_indeces
        temp_A = deepcopy(A)
        temp_predictor_indeces = deepcopy(predictor_indeces)
        best_loss = Inf

        for i = 1 : n_random_predictor_samples
            MvDecisionTrees._reservoir_sample!(temp_predictor_indeces, p, ϕ)
            loss = _regress_on_predictor_subset(data, assignment_id, Γ, U, Φ, temp_A, temp_predictor_indeces)
            if loss < best_loss
                best_loss = loss
                copy!(A, temp_A)
                copy!(predictor_indeces, temp_predictor_indeces)
            end
        end

        # in case we found nothing
        if isinf(best_loss)
            fill!(A, 0.0)
            fill!(predictor_indeces, 1)
        end
    end

    ###########################################################3
    # calc covariance

    Σ = zeros(T, o, o)
    r_vec = Array(Float64, o)

    # obtain upper triangle
    for i in 1 : n
        if data.assignment[i] == assignment_id

            for j in 1 : o
                r_vec[j] = data.Y[j,i] - A[j,ϕ+1]
                for k in 1 : ϕ
                    r_vec[j] -= A[j,k]*data.X[i,predictor_indeces[k]]
                end
            end

            for a in 1 : o
                for b in a : o
                    Σ[a,b] += r_vec[a] * r_vec[b]
                end
            end
        end
    end

    if id_count < 2
        warn("Not enough samples to build a covariance matrix")
        for i in 1 : o
            Σ[i,i] = 1.0
        end
    else
        for i in 1 : o*o
            @assert(!isnan(Σ[i]))
            Σ[i] /= (id_count-1)
        end

        ε = GLOBAL_DIAG_SHRINKAGE_THRESOLD
        MvDecisionTrees._diagonal_shrinkage!(Σ, ε)
        for a in 2:o
            for b in 1:a-1
                Σ[a,b] = Σ[b,a]
            end
        end
    end

    AutoregressiveMvNormLeaf(A, MvNormal(Array(T, o), Σ), predictor_indeces)
end

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

type DF_TrainParams <: AbstractVehicleBehaviorTrainParams

    targets::ModelTargets
    indicators::Vector{AbstractFeature}

    ntrees::Int
    max_tree_depth::Int
    n_split_tries::Int
    min_samples_split::Int
    min_samples_leaves::Int
    n_random_predictor_samples::Int
    n_autoregression_predictors::Int
    use_PCA::Bool

    diag_shrinkage_threshold::Float64
    min_split_improvement::Float64
    partial_sampling::Float64
    autogression_coef::Float64

    function DF_TrainParams(;
        targets::ModelTargets = ModelTargets(Features.FUTUREDESIREDANGLE, Features.FUTUREACCELERATION),
        indicators::Vector{AbstractFeature} = AbstractFeature[],

        ntrees::Integer=10,
        max_tree_depth::Integer=5,
        n_split_tries::Integer=50,
        min_samples_split::Integer=100,
        min_samples_leaves::Integer=20,
        n_random_predictor_samples::Integer=10,
        n_autoregression_predictors::Integer=2,
        use_PCA::Bool=false,

        diag_shrinkage_threshold::Float64=1e-6,
        min_split_improvement::Float64=0.0,
        partial_sampling::Float64=0.7,
        autogression_coef::Float64=0.1,
        )

        retval = new()
        retval.targets = targets
        retval.indicators = indicators
        retval.ntrees = ntrees
        retval.max_tree_depth = max_tree_depth
        retval.n_split_tries = n_split_tries
        retval.min_samples_split = min_samples_split
        retval.min_samples_leaves = min_samples_leaves
        retval.n_random_predictor_samples = n_random_predictor_samples
        retval.n_autoregression_predictors = n_autoregression_predictors
        retval.use_PCA = use_PCA

        retval.diag_shrinkage_threshold = diag_shrinkage_threshold
        retval.min_split_improvement = min_split_improvement
        retval.partial_sampling = partial_sampling
        retval.autogression_coef = autogression_coef
        retval
    end
end
type DF_PreallocatedData <: AbstractVehicleBehaviorPreallocatedData

    # the full dataset matrices pre-extracted
    X::Matrix{Float64} # [nfeatures × nframes]
    Y::Matrix{Float64} # [2 × nframes]
    Z::Matrix{Float64} # [n_pca_features × nframes]

    extractor::FeatureSubsetExtractor
    preprocess::DataPreprocessor
    preprocess_pca::DataPreprocessor
    action_clamper::DataClamper

    function DF_PreallocatedData(dset::ModelTrainingData2, params::DF_TrainParams)

        retval = new()

        targets = params.targets
        indicators = params.indicators
        trainingframes = dset.dataframe
        nframes = nrow(trainingframes)
        nindicators = length(indicators)

        X = Array(Float64, nindicators, nframes) # NOTE(tim): initially indicators × nframes, we later transpose it
        Y = Array(Float64, 2, nframes)

        pull_design_and_target_matrices!(X, Y, trainingframes, targets, indicators)

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

        extractor = FeatureSubsetExtractor(indicators)
        preprocess = ChainedDataProcessor(extractor)
        push!(preprocess, X, DataNaReplacer, extractor.indicators)
        process!(X, preprocess.processors[end]) # process in place
        push!(preprocess, X, DataScaler) # ensure data is standardized before running PCA
        process!(X, preprocess.processors[end]) # process in place
        retval.X = copy(X)

        # Add a PCA scaler
        preprocess_pca = deepcopy(preprocess)
        n_PCA_features = min(nindicators, 20)
        push!(preprocess_pca, X, DataLinearTransform, n_PCA_features)

        Z = Array(Float64, n_PCA_features, nframes)
        retval.Z = process!(Z, X, preprocess_pca.processors[end])

        retval.Y = Y
        retval.extractor = extractor
        retval.preprocess = preprocess
        retval.preprocess_pca = preprocess_pca
        retval.action_clamper = DataClamper(
                Array(Float64, 2),
                vec(minimum(Y, 2)),
                vec(maximum(Y, 2))
            )

        retval
    end
end
type DynamicForestBehavior <: AbstractVehicleBehavior
    forest::Ensemble

    targets::ModelTargets
    indicators::Vector{AbstractVehicleBehavior}
    extractor::FeatureSubsetExtractor
    processor::DataPreprocessor
    action_clamper::DataClamper # contains action vector and clamping range

    # preallocated memory
    ϕ::Vector{Float64} # [n_predictors+1] (last one is always 1.0)

    function DynamicForestBehavior(
        forest::Ensemble,
        n_predictors::Int,
        params::DF_TrainParams,
        pre::DF_PreallocatedData,
        )

        extractor = pre.extractor
        action_clamper = pre.action_clamper
        processor = params.use_PCA ? pre.preprocess_pca : pre.preprocess

        retval = new()
        retval.forest = forest
        retval.targets = params.targets

        # copy the processor and ensure that the copied io arrays match up
        retval.extractor = FeatureSubsetExtractor(deepcopy(extractor.x), extractor.indicators)
        retval.processor = deepcopy(processor, retval.extractor)
        retval.action_clamper = DataClamper(Array(Float64, 2),
                                   deepcopy(action_clamper.f_lo),
                                   deepcopy(action_clamper.f_hi))

        retval.ϕ = Array(Float64, n_predictors+1)

        retval
    end
end
AutomotiveDrivingModels.get_targets(behavior::DynamicForestBehavior) = behavior.targets
AutomotiveDrivingModels.get_indicators(behavior::DynamicForestBehavior) = behavior.extractor.indicators

function preallocate_learning_data(
    dset::ModelTrainingData2,
    params::DF_TrainParams)

    DF_PreallocatedData(dset, params)
end

function _condition_predictor_mean!(
    normal::MvNormal,
    predictor_indeces::Vector{Int},
    A::Matrix{Float64},
    behavior::DynamicForestBehavior,
    )

    ϕ = behavior.ϕ

    for (i,p) in enumerate(predictor_indeces)
        ϕ[i] = behavior.processor.z[p]
    end
    ϕ[end] = 1.0

    normal.μ[:] = A * ϕ

    nothing
end
function select_action(
    behavior::DynamicForestBehavior,
    runlog::RunLog,
    sn::StreetNetwork,
    colset::UInt,
    frame::Int
    )

    observe!(behavior.extractor, runlog, sn, colset, frame)
    process!(behavior.processor)

    # pick a tree at random
    tree_index = rand(1:length(behavior.forest.trees))

    # sample from the MvNorm for said tree
    leaf = apply_tree(behavior.forest.trees[tree_index], behavior.processor.z)::AutoregressiveMvNormLeaf
    _condition_predictor_mean!(leaf.m, leaf.predictor_indeces, leaf.A, behavior)

    Distributions._rand!(leaf.m, behavior.action_clamper.x)
    process!(behavior.action_clamper)

    action_lat = behavior.action_clamper.x[1]
    action_lon = behavior.action_clamper.x[2]

    (action_lat, action_lon)
end

function _set_and_process_action!(behavior::DynamicForestBehavior, action_lat::Float64, action_lon::Float64)
    behavior.action_clamper.x[1] = action_lat
    behavior.action_clamper.x[2] = action_lon
    process!(behavior.action_clamper)
    behavior.action_clamper
end
function _calc_action_loglikelihood(behavior::DynamicForestBehavior)

    #=
    Action should be in behavior.action_clamper and should already have been processed
    Observations should be in behavior.processor.z and should already have been processed
    =#

    ntrees = length(behavior.forest.trees)
    log_ntrees = log(ntrees)

    total_probability_density = 0.0
    logl = 0.0
    for tree in behavior.forest.trees
        leaf = apply_tree(tree, behavior.processor.z)::AutoregressiveMvNormLeaf
        _condition_predictor_mean!(leaf.m, leaf.predictor_indeces, leaf.A, behavior)
        normal  = leaf.m

        p = pdf(normal, behavior.action_clamper.x)/ntrees

        if total_probability_density == 0.0
            logl = logpdf(normal, behavior.action_clamper.x) - log_ntrees
        elseif p < total_probability_density
            # NOTE: log(a + b) = log a + log(1 + b/a)
            logl = logl + log1p(p/total_probability_density)
        else
            # NOTE: log(a + b) = log b + log(1 + a/b)
            logl = (logpdf(normal, behavior.action_clamper.x) - log_ntrees) + log1p(total_probability_density/p)
        end
        total_probability_density += p
    end

    logl
end

function calc_action_loglikelihood(
    behavior::DynamicForestBehavior,
    runlog::RunLog,
    sn::StreetNetwork,
    colset::UInt,
    frame::Int,
    action_lat::Float64,
    action_lon::Float64,
    )

    #=
    Compute the log-likelihood of the action taken during a single frame
    given the VehicleBehaviorGaussian.
    =#

    observe!(behavior.extractor, runlog, sn, colset, frame)
    process!(behavior.processor)

    _set_and_process_action!(behavior, action_lat, action_lon)

    _calc_action_loglikelihood(behavior)
end
function calc_action_loglikelihood(
    behavior::DynamicForestBehavior,
    features::DataFrame,
    frameind::Integer,
    )

    observe!(behavior.extractor, features, frameind)
    process!(behavior.processor)

    action_lat = features[frameind, symbol(behavior.targets.lat)]::Float64
    action_lon = features[frameind, symbol(behavior.targets.lon)]::Float64
    _set_and_process_action!(behavior, action_lat, action_lon)

    _calc_action_loglikelihood(behavior)
end

function train(
    training_data::ModelTrainingData2,
    preallocated_data::DF_PreallocatedData,
    params::DF_TrainParams,
    foldset::FoldSet,
    )

    extractor = preallocated_data.extractor

    ntrees = params.ntrees
    max_tree_depth = params.max_tree_depth
    n_split_tries = params.n_split_tries
    min_samples_split = params.min_samples_split
    min_samples_leaves = params.min_samples_leaves
    n_random_predictor_samples = params.n_random_predictor_samples
    n_autoregression_predictors = params.n_autoregression_predictors
    min_split_improvement = params.min_split_improvement
    partial_sampling = params.partial_sampling
    autogression_coef = params.autogression_coef

    X_orig = params.use_PCA ? preallocated_data.Z : preallocated_data.X
    Y_orig = preallocated_data.Y

    # TODO(tim): do this without globals
    global DEFAULT_AUTOREGRESSION_CONSTANT = autogression_coef
    global DEFAULT_NUM_AUTOREGRESSION_PREDICTORS = n_autoregression_predictors
    global DEFAULT_NUM_PREDICTOR_SAMPLES = n_random_predictor_samples
    global GLOBAL_DIAG_SHRINKAGE_THRESOLD = params.diag_shrinkage_threshold

    build_tree_params = BuildTreeParameters(
        round(Int, sqrt(length(extractor.indicators))),
        max_tree_depth,
        min_samples_split,
        min_samples_leaves,
        min_split_improvement,
        n_split_tries,
        LossFunction_MSE,
        AutoregressiveMvNormLeaf
        )

    nframes = length(foldset)
    X = Array(Float64, nframes, size(X_orig, 1)) # NOTE: this is transpose of X_orig
    y = Array(Float64, 2, nframes)

    i = 0
    for frame in foldset
        i += 1

        for j in 1 : size(X_orig, 1)
            X[i,j] = X_orig[j,frame]
        end

        y[1,i] = Y_orig[1,frame]
        y[2,i] = Y_orig[2,frame]
    end
    @assert(i == nframes)

    ensemble = build_forest(y, X, ntrees, build_tree_params, partial_sampling)
    DynamicForestBehavior(ensemble, n_autoregression_predictors, params, preallocated_data)
end

end # end module


using Profile, PProf
using Pkg;Pkg.activate(".")
using PairwiseVelocities
using StaticArrays
using BenchmarkTools
using DelimitedFiles

boxsize = 2000.
rbins = LinRange(0.,50,50)
rbins_c = 0.5*(rbins[2:end] + rbins[1:end-1])

filename = "moments.csv"
logn = -5.
run = 101
snapshot = 15

positions, velocities = read_data(run, snapshot, 10^logn, boxsize)
positions = convert(Array{Float64}, positions)
velocities = convert(Array{Float64}, velocities)

moments = PairwiseVelocities.compute_pairwise_velocity_moments(
    positions,
    velocities,
    rbins,
    boxsize,
)
open(filename; write=true) do f
    write(f, "# r_c v_r sigma_r sigma_t skewness_r skewness_rt kurtosis_r kurtosis_t kurtosis_rt \n")
    writedlm(f,
            zip(rbins_c, moments[2][:],
            moments[3][:], moments[6][:],
            moments[4][:], moments[7][:],
            moments[5][:], moments[8][:], moments[9][:]))
    end

println("Wrote file !")


using Test
using Gravityfield

let
    @test V(1.0m) == (4π / 3)m³
    @test ρ(m = (4π / 3)kg, r = 1m) == 1kg/m³
    @test ρ(r = 1m, m = (4π / 3)kg) == 1kg/m³
    @test ρ_earth |> ustrip |> typeof == Float64
end

@testset "Gravity acceleration" begin
    @test g(1, 0) == 0.0m^3*kg^-1*s^-2
    @test g(1., 1.) == γ
    @test g(m_earth, r_earth) ≈ 10.0391664m/s^2
    ω_earth_sun = 2π / yr
    a_earth_sun = ω_earth_sun^2 * d_sun_earth
    @test g(m_sun, d_sun_earth) / a_earth_sun ≈ 1.000 atol = 0.001
end

@testset "Force vector on 1 due to 2" begin
    p1 = [0.,0.,0.]m
    p2 = [0.,1.,0.]m
    m1 = 316227kg
    m2 = 129099kg
    f1_2 = m1* uvec(p1, p2) * g(m2, len(p2-p1))
    @test f1_2 |> length == 3
    @test ustrip.(f1_2) |> typeof == Vector{Float64}
    @test f1_2 ≈ [0.0, 2.7247474, 0.]N
end
@testset "Euclidean vector length" begin
    @test len(-1m) == 1m
    @test len([1]m) == 1m
    @test len([2,0]m) == 2m
    @test len([1m,1m]) == sqrt(2)m
end
@testset "Vector between bodies or points" begin
    p1 = [0.,0.,0.]m
    p2 = [0.,1.,0.]m
    s = Sphere(p1, m_earth, r_earth)
    @test rvec(s,s) == [0., 0., 0.]m
    @test len(s,s) == 0m
    @test uvec(s,s) == [0., 0., 0.]
    @test rvec(p2,s) == [0., -1.0, 0.]m
    @test len(p2,s) == 1m
    @test uvec(p2,s) == [0., -1.0, 0.]

    s2 = Sphere(p2, m_earth, r_earth)
    @test rvec(s,s2) == [0., 1.0, 0.]m
    @test len(s,s2) == 1m
    @test uvec(s,s2) == [0., 1.0, 0.]

    s3 = Sphere([-d_sun_earth], 1, 1)
    s4 = Sphere([0.0m], 1, 1)
    @test rvec(s3,s4) == [d_sun_earth]
    @test len(s4,s3) == d_sun_earth
    @test uvec(s3,s4) == [1]
    @test uvec(s4,s3) == [-1]
end

@testset "Gravity acceleration vector btw. bodies or point and body" begin
    p1 = [0.,0.,0.]m
    s1 = Sphere(p = p1, m = 1kg, r = 1m)
    @test gvec(s1,s1) == [0., 0., 0.]m/s²
    @test gvec([-1, 0., 0.]m, s1) == [γ*kg/m², 0m/s², 0m/s²]
    s2 = Sphere([0.0]m, m_earth, r_earth)
    @test gvec([r_earth], s2) == [-10.039166439909296]m/s²
    @test gvec([r_earth/2], s2) == [-10.039166439909296]m/s²/2
end

@testset "Volume" begin
    @test V(1m) == (4π/3)m³
    @test V(r_sun) * ρ_sun / m_sun ≈ 0.999 atol = 0.001
    @test V(r_sun) isa Quantity
end

@testset "Sphere constructors" begin
    s1 = Sphere(p = [0]m, r = r_earth)
    s2 = Sphere([0]m, m_earth, r_earth)
    @test s1.m == s2.m
    @test s1.r == s2.r
    s3 = Sphere()
    @test s3.m ≈ s2.m
    @test s3.r ≈ s2.r
    s4 = earthlikeSphere_r_known(r = 10.E6m)
    s5 = earthlikeSphere_m_known(m = ρ_earth * V(10.E6m))
    @test s4.m ≈ s5.m
    @test s4.r ≈ s5.r
end

@testset "Position vectors" begin
    @test [1, 2, 3]m isa Pos
    @test !isa([1 2; 3 4]m, Pos)
    @test !isa([1m, 2, 3], Pos)
    @test !isa([1, sin, 3], Pos)
    @test !isa([[1, 2, 3, 4,5,6]], Pos)
    @test [1, 2, 3, 4,5,6]m isa Pos
    @test !isa([[1, 2, 3, 4,5,6]], Pos)
    @test !isa([1, 2, 3, 4,5,6]kg, Pos)
end

@testset "Gravity acceleration vector in system of bodies" begin
    earth = Sphere()
    moon = earthlikeSphere_m_known(p = [d_moon_earth, 0m, 0m], m = m_moon)
    sun = Sphere([-d_sun_earth, 0m, 0m], m_sun, (3m_sun / (4π * ρ_sun))^(1/3))
    bodies = [sun, earth, moon]
    @test gvec(earth.p, bodies)[1] < 0m/s² # Sun stronger than moon in centre of earth
    @test gvec([r_earth/2, 0m, 0m], bodies)[1] < -5m/s² # Halfway to centre of earth, earth dominates
    xgrid, ygrid = xygrids( xmin = -d_moon_earth/2.,
        xmax = d_moon_earth * 3/2,
        ymin =-d_moon_earth/2,
        ymax = d_moon_earth * 3/2, n = 100)
    gmagn = gfieldvals(nest(xgrid, ygrid, moon.r),  bodies)
    @test foldr(+, gmagn)  ≈  222.5649149449058m/s²
    @test gvec(earth, bodies) == gvec(earth.p, bodies)
end

@testset "Nested position vectors" begin
    earth = Sphere()
    moon = earthlikeSphere_m_known(p = [d_moon_earth, 0m, 0m], m = m_moon)
    sun = Sphere([-d_sun_earth, 0m, 0m], m_sun, (3m_sun / (4π * ρ_sun))^(1/3))
    bodies = [sun, earth, moon]
    p1 = Pos([1., 2.]m)
    p2 = Pos([-1., -2.]m)
    vp = Npos([p1, p2])
    @test vp == [[1., 2.], [-1., -2.]]m
    mp = hcat(vp, vp)
    @test mp isa Npos
    pu1 = [1., 2.]
    pu2 = [-1., -2.]
    mup = hcat([pu1, pu2], [pu1, pu2])
    @test mp == mup*m
    @test nest(1m:1m:3m, 2m, 3m) == [[1, 2, 3], [2, 2, 3], [3, 2, 3]]m
    @test nest(1m,2m:1m:3m, 4m) == [[1, 2, 4], [1, 3, 4]]m
    @test nest(1m,2m, 3m:1m:4m) == [[1, 2, 3], [1, 2, 4]]m
    @test nest(1m:1m:2m, 1m:1m:2m) == [[1,1], [2,2]]m
    @test nest(1m,2m:1m:3m) == [[1,2], [1,3]]m
    @test nest(1m:1m:2m, 3m) == [[1,3], [2,3]]m
    @test nest(range(1.5m, 3.5m, length = 2)) == [[1.5], [3.5]]m
    @test gfieldvals(nest(1m,2m:1m:3m,4m), bodies)[1] ≈ 0.0059000661748m/s²
    @test nest([1m, 2m, 3m], [4, 5, 6]m, [7, 8, 9]m) ==  [[1, 4, 7], [2, 5, 8], [3, 6, 9]]m
    np = nest([1, 10]m, [2, 20]m, [3, 30]m, [4, 40]m, [5, 50]m, [6, 60]m)
    @test np ==  [[1, 2, 3, 4, 5, 6]m, [10, 20, 30, 40, 50, 60]m]
    @test slice_1(np) == [1, 10]m
    @test slice_2(np) == [2, 20]m
    @test slice_3(np) == [3, 30]m
    @test slice_4(np) == [4, 40]m
    @test slice_5(np) == [5, 50]m
    @test slice_6(np) == [6, 60]m
    @test gvec([[0m,0m,0m]], bodies)[1][1] ≈ -0.005898468352763m/s²
end

@testset "Bodies constructors" begin
    @test eltype(bodies_3d()) == Sphere
    @test eltype(bodies_2d()) == Sphere
    @test eltype(bodies_1d()) == Sphere
end

@testset "Pretty-printing" begin
    function sout(p)
        iob=IOBuffer()
        show(IOContext(iob, :color=>true), p)
        take!(iob)|>String
    end
    p = [1,2]m
    @test sout(p) == "[1.0, 2.0]\e[36mm\e[39m"
    function slongout(p)
        iob=IOBuffer()
        show(IOContext(iob, :color=>true), :"text/plain", p)
        take!(iob)|>String
    end
    @test slongout(p) == "Pos{Float64}([1.0, 2.0]\e[36mm\e[39m)"
    vp = [p, p]
    @test sout(vp) == "[[1.0, 2.0], [1.0, 2.0]]\e[36mm\e[39m"
    @test slongout(vp) == "2-element Npos{Pos{Float64}}:\n[[1.0, 2.0], [1.0, 2.0]]\e[36mm\e[39m"
    p1 = [1,1]m
    p2 = [2,2]m
    p3 = [3,3]m
    p4 = [4,4]m
    mp = hcat([p1, p2], [p3, p4])
    @test mp[2,1] == p2
    @test sout(mp) == "[[1.0, 1.0] [3.0, 3.0]; [2.0, 2.0] [4.0, 4.0]]\e[36mm\e[39m"
    @test slongout(mp) == "2x2 Npos{Pos{Float64}}:\n[[1.0, 1.0] [3.0, 3.0]; [2.0, 2.0] [4.0, 4.0]]\e[36mm\e[39m"
    @test sout(bodies_3d()) == "Sphere[Sphere(p = [-1.496e11, 0.0, 0.0]\e[36mm\e[39m, m = 1.989e30 kg, r = 6.95733e8 m), Sphere(p = [0.0, 0.0, 0.0]\e[36mm\e[39m, m = 5.97e24 kg, r = 6.3e6 m), Sphere(p = [3.844e8, 0.0, 0.0]\e[36mm\e[39m, m = 7.34767e22 kg, r = 1.45456e6 m)]"
    @test sout(bodies_2d()) == "Sphere[Sphere(p = [-1.496e11, 0.0]\e[36mm\e[39m, m = 1.989e30 kg, r = 6.95733e8 m), Sphere(p = [0.0, 0.0]\e[36mm\e[39m, m = 5.97e24 kg, r = 6.3e6 m), Sphere(p = [3.844e8, 0.0]\e[36mm\e[39m, m = 7.34767e22 kg, r = 1.45456e6 m)]"
    @test sout(bodies_1d()) == "Sphere[Sphere(p = [-1.496e11]\e[36mm\e[39m, m = 1.989e30 kg, r = 6.95733e8 m), Sphere(p = [0.0]\e[36mm\e[39m, m = 5.97e24 kg, r = 6.3e6 m), Sphere(p = [3.844e8]\e[36mm\e[39m, m = 7.34767e22 kg, r = 1.45456e6 m)]"
    @test slongout(bodies_3d()) == "3-element Array{Sphere,1}:\n  p = [-1.496e11, 0.0, 0.0]\e[36mm\e[39m m = 1.989e30 kg r = 6.95733e8 m\n  p = [0.0, 0.0, 0.0]\e[36mm\e[39m m = 5.97e24 kg r = 6.3e6 m\n  p = [3.844e8, 0.0, 0.0]\e[36mm\e[39m m = 7.34767e22 kg r = 1.45456e6 m"
    @test slongout(bodies_2d()) == "3-element Array{Sphere,1}:\n  p = [-1.496e11, 0.0]\e[36mm\e[39m m = 1.989e30 kg r = 6.95733e8 m\n  p = [0.0, 0.0]\e[36mm\e[39m m = 5.97e24 kg r = 6.3e6 m\n  p = [3.844e8, 0.0]\e[36mm\e[39m m = 7.34767e22 kg r = 1.45456e6 m"
    @test slongout(bodies_1d()) == "3-element Array{Sphere,1}:\n  p = [-1.496e11]\e[36mm\e[39m m = 1.989e30 kg r = 6.95733e8 m\n  p = [0.0]\e[36mm\e[39m m = 5.97e24 kg r = 6.3e6 m\n  p = [3.844e8]\e[36mm\e[39m m = 7.34767e22 kg r = 1.45456e6 m"
end

nothing


using TensorToolbox
include("LTR.jl")
include("methods/NTD.jl")
include("methods/lraSNTD.jl")
include("methods/LD/LD.jl")

function tucker_to_cprank(tucker_rank)
    if all(isequal(first(tucker_rank)),tucker_rank)
        cp = tucker_rank[1]
    else
        error("rank $tucker_rank cannot convert to CP rank")
    end
    return cp
end

function calc(X, method, reqrank)
    if method == "LT1R"
        @assert tucker_to_cprank(reqrank) == 1
        return LT1R(X)
    elseif method == "LTR"
        return LTR(X, reqrank)
    elseif method == "NTD_KL"
        _, _, Xr = NTD(X, reqrank, cost="KL", init_method=NTD_init, verbose=false)
        return Xr
    elseif method == "NTD_LS"
        _, _, Xr = NTD(X, reqrank, cost="LS", init_method=NTD_init, verbose=false)
        return Xr
    elseif method == "lraSNTD"
        return lraSNTD(X, reqrank)
    else
        error("calc method error")
    end
end


#Run all lines of code in PopGen.jl
#Then read in all functions within PairwiseRelatedness.jl

verbose = true
data = nancycats()

alleles = Dict()
for loc in loci(data)
    alleles[loc] = allele_freq(locus(data, loc))
end
alleles

#Working pair
ind1 = "N100"
ind2 = "N106"

#Not working pair ex.
ind1 = "N111"
ind2 = "N83"

L = all_loci_Pr_L_S(data, ind1, ind2, alleles)

Δ = Variable(8)
problem = maximize(sum(log(L * vcat(1 - sum(Δ), Δ))))
problem.constraints += 0 <= 1 - sum(Δ)
problem.constraints += 1 - sum(Δ) <= 1
problem.constraints += 0 <= Δ[1:8]
problem.constraints += Δ[1:8] <= 1

#shifted from actual relatedness calculations because the 1 - sum(Δ) goes at beginning
problem.constraints += 2 * ((1 - sum(Δ)) + 0.5 * (Δ[2] + Δ[4] + Δ[6]) + 0.25 * Δ[7]) <= 1
problem.constraints += 0 <= 2 * ((1 - sum(Δ)) + 0.5 * (Δ[2] + Δ[4] + Δ[6]) + 0.25 * Δ[7])

Convex.solve!(problem, ECOSSolver(verbose = verbose, maxit = 100), verbose = verbose)

n = Variable(9)
problem2 = minimize(sum(log(L * exp(n)) / sum(exp(n))))
Convex.solve!(problem2, ECOSSolver())


"""
    F, C = maxflow(G, s, t)

G is the adjacency matrix of an undirected graph, with edge weights equal to capacity.
Computes the maximum s-t flow, F, and the minimum cut, C.

F[a,b] is flow from a to b, and F[b,a] = -F[a,b].
C is the set reachable from s using unsaturated edges.

Uses Dinic's algorithm.  It would probably be easy to make this work for directed graphs.
"""
function maxflow(G::SparseMatrixCSC{Tv, Ti}, s::Int64, t::Int64) where {Tv, Ti}

  n = max(G.n, G.m)

  # initialize the capacity matrix
  C = copy(G)
  backInd = backIndices(G)

  Q = zeros(Ti, n)
  inQ = zeros(Bool, n)
  #=
   prev[i][1] = the vertex preceding vertex i in the BFS
   prev[i][2] = the edge index of prev[i][1]; nbri(G, prev[i][1], prev[i][2]) = i
  =#
  prev = [(0,0) for i in 1:n]
  totalflow = 0

  improve = true
  while improve
    improve = false

    # initialize the queue at the current step
    left = 0
    right = 1
    Q[right] = s
    inQ[s] = true

    # do bfs
    while left < right
      left = left + 1
      u = Q[left]
      if (u == t)
        continue
      end

      for i in 1:deg(G, u)
        v = nbri(G, u, i)

        if inQ[v] == false && weighti(C, u, i) > 0
          right = right + 1
          Q[right] = v
          inQ[v] = true
          prev[v] = (u, i)
        end
      end
    end

    # do block flow
    if inQ[t] == true
      improve = true

      for i in 1:deg(G, t)
        # ignore a neighbor that isn't part of the BF-tree
        if inQ[nbri(G, t, i)] == false
          continue
        end

        # go through every neighbor of t and see the increase in flow
        prev[t] = (nbri(G, t, i), weighti(backInd, t, i))
        u = t

        currentflow = typemax(Tv)
        while u != s
          prevU = prev[u][1]
          indexPrevU = prev[u][2]
          currentflow = min(currentflow, weighti(C, prevU, indexPrevU))
          u = prevU
        end

        if currentflow != 0
          #=
            prevU is the vertex immmediately before u on the path from s to t
            indexU is such that nbri(G, u, indexU) = prevU
            indexPrevU is such that nbri(G, prevU, indexPrevU) = u
            capPrevUU and capUPrevU have similar meanings
          =#
          u = t
          while u != s
            prevU = prev[u][1]
            indexPrevU = prev[u][2]
            indexU = weighti(backInd, prevU, indexPrevU)

            capPrevUU = weighti(C, prevU, indexPrevU)
            capUPrevU = weighti(C, u, indexU)

            setValue(C, prevU, indexPrevU, capPrevUU - currentflow)
            setValue(C, u, indexU, capUPrevU + currentflow)

            u = prevU
          end
        end

        totalflow = totalflow + currentflow
      end

    end

    # reset the queue
    for i in 1:right
      u = Q[i]

      Q[i] = 0
      inQ[u] = false
      prev[u] = (0,0)
    end
    left = right = 0

  end

  # compute the min-cut
  left = 0
  right = 1
  Q[right] = s
  inQ[s] = true

  while left < right
    left = left + 1

    u = Q[left]
    for i in 1:deg(G, u)
      v = nbri(G, u, i)

      if inQ[v] == 0 && weighti(C, u, i) > 0
        right = right + 1
        Q[right] = v
        inQ[v] = 1
      end
    end
  end

  cut = Q[1:right]
  flow = C - G

  return flow, cut
end


using NetCDF, LinearAlgebra
using Statistics, StatsBase
using PyCall
# xr = pyimport("xarray")
#pygui(:qt)
using PyPlot
using MultivariateStats
# eofs = pyimport("eofs")

path = "/Users/milan/Downloads/cams/ieee"
filelist = filter(x->endswith(x,"no2.nc"),readdir(path))

ncfile = NetCDF.open(joinpath(path,filelist[1]))
no2 = ncfile.vars["no2"][:,:,:,1]
level = ncfile.vars["level"][:]
lat = ncfile.vars["latitude"][:]
lon = ncfile.vars["longitude"][:]

no2m = mean(no2,dims=(1,2))[1,1,:]
no2mi = median(no2,dims=(1,2))[1,1,:]
no2min = minimum(no2,dims=(1,2))[1,1,:]
no2max = maximum(no2,dims=(1,2))[1,1,:]
no2p10 = [percentile(vec(no2[:,:,i]),10) for i in level]
no2p90 = [percentile(vec(no2[:,:,i]),90) for i in level]

## EOF in the vertical
no2f = Matrix(reshape(no2,(:,length(level)))')
M = fit(PCA,log10.(no2f),maxoutdim=20,pratio=0.999)
no2fhat = transform(M,log10.(no2f))
no2fr = 10f0 .^reconstruct(M,no2fhat)

## reconstruction statistics
no2rm = mean(no2fr,dims=2)[:,1]
no2rmi = median(no2fr,dims=2)[:,1]
no2rmin = minimum(no2fr,dims=2)[:,1]
no2rmax = maximum(no2fr,dims=2)[:,1]
no2rp10 = [percentile(no2fr[i,:],10) for i in level]
no2rp90 = [percentile(no2fr[i,:],90) for i in level]

## error
decimal_error = abs.(log10.(reshape(no2,(:,length(level)))' ./ no2fr))

dem = mean(decimal_error,dims=2)[:,1]
demi = median(decimal_error,dims=2)[:,1]
demin = minimum(decimal_error,dims=2)[:,1]
demax = maximum(decimal_error,dims=2)[:,1]
dep10 = [percentile(decimal_error[i,:],10) for i in level]
dep90 = [percentile(decimal_error[i,:],90) for i in level]

## highest loadings
P = projection(M)
Ps = zeros(Int64,size(P))
for i in 1:size(Ps)[2]
    Ps[:,i] = sortperm(abs.(P[:,i]))
end
L = zeros(Int64,size(P)[2])
L[1] = Ps[end,1]
for i in 2:length(L)
    j = 0
    while Ps[end-j,i] in L
        j += 1
    end
    L[i] = Ps[end-j,i]
end

##
fig,(ax,ax2,ax3) = subplots(1,3,figsize=(10,6),sharey=true)
ax.invert_yaxis()

ax.semilogx(no2m,level,"C0",lw=3,label="mean")
ax.plot(no2mi,level,"C0",lw=1.5,label="median")
ax.fill_betweenx(level,no2min,no2max,color="C0",alpha=0.2,label="Min-max range")
ax.fill_betweenx(level,no2p10,no2p90,color="C0",alpha=0.5,label="10-90% range")

ax2.semilogx(no2rm,level,"C0",lw=3,label="mean")
ax2.plot(no2rmi,level,"C0",lw=1.5,label="median")
ax2.fill_betweenx(level,no2rmin,no2rmax,color="C0",alpha=0.2,label="Min-max range")
ax2.fill_betweenx(level,no2rp10,no2rp90,color="C0",alpha=0.5,label="10-90% range")

for i in level
    xvals = no2f[i,(decimal_error[i,:] .>= 0.5f0) .* (decimal_error[i,:] .< 1f0)]
    ax.scatter(xvals,i*ones(length(xvals)),10,"yellow",alpha=0.05)

    xvals = no2f[i,(decimal_error[i,:] .>= 1f0) .* (decimal_error[i,:] .< 2f0)]
    ax.scatter(xvals,i*ones(length(xvals)),10,"C1",alpha=0.1)

    xvals = no2f[i,(decimal_error[i,:] .>= 2f0)]
    ax.scatter(xvals,i*ones(length(xvals)),10,"C3",alpha=0.1)
end

ax.scatter([1,2,3],[0,0,0],10,"yellow",alpha=0.5,label="error > 0.5")
ax.scatter([1,2,3],[0,0,0],10,"C1",alpha=0.5,label="error > 1")
ax.scatter([1,2,3],[0,0,0],10,"C3",alpha=0.5,label="error > 2")

ax3.semilogx(dem,level,"C0",lw=3,label="mean")
ax3.plot(demi,level,"C0",lw=1.5,label="median")
ax3.fill_betweenx(level,demin,demax,color="C0",alpha=0.2,label="Min-max range")
ax3.fill_betweenx(level,dep10,dep90,color="C0",alpha=0.5,label="10-90% range")

cmap = ColorMap("viridis")
for i in 1:length(L)
    ax2.plot([1e-15,1e-14],[L[i],L[i]].-0.5,color=cmap(i/length(L)),lw=2)
    ax2.text(2e-14,L[i],"$i")
end
ax2.text(1.5e-15,8,"modes")

ax.set_title(L"NO$_2$ vertical distribution")
ax2.set_title(L"NO$_2$ reconstructed from 20 PCs")
ax3.set_title("L1 log-error")
ax.set_ylabel("model level")
ax.set_xlabel("mixing ratio kg/kg")
ax2.set_xlabel("mixing ratio kg/kg")
ax.set_ylim(137,1)
ax.legend(loc=2,scatterpoints=3)

ax.set_xlim(1e-15,1e-7)
ax2.set_xlim(1e-15,1e-7)

tight_layout()


using LinearAlgebra
using StaticArrays

function position_in_world(x::Float64, y::Float64)
    target_width_px = 1123
    target_height_px = 791
    target_width_mm = 297
    target_height_mm = 210.025
    X = ((x + 0.5) / target_width_px) * target_width_mm - 0.5 * target_width_mm
    Y = ((y + 0.5) / target_height_px) * target_height_mm - 0.5 * target_height_mm
    Z = 0.0
    return SVector(X, Y, Z)
end

function compute_homograpy_DLT(M::StaticMatrix, P::StaticMatrix)
    T = promote_type(eltype(M), eltype(P))
    A = zeros(T, 8, 9)
    for i in 1:size(M, 1)
        Mi = M[i,:]
        A[2 * i - 1, 1:3] = Mi
        A[2 * i - 1, 7:9] = - P[i,1] * Mi
        A[2 * i, 4:6] = Mi
        A[2 * i, 7:9] = - P[i,2] * Mi
    end
    return nullspace(A)
end

function π_projection_function(K::StaticMatrix, R::StaticMatrix, t::StaticVector, point::StaticVector)
    # point_in_camera = R * point + t
    point_in_image = K * (R * point + t)
    return point_in_image / point_in_image[end]
end

function foo()
    K = @SMatrix [ 1169.19630    0.0     652.98743;
                0.0     1169.61014 528.83429;
                0.0        0.0       1.0]

    T = @SMatrix [0.961255 -0.275448 0.0108487  112.79;
            0.171961  0.629936 0.75737   -217.627;
        -0.21545  -0.72616  0.652895  1385.13]

    R = T[:, SOneTo(3)]
    t = T[:,end]

    tl = position_in_world(0.0, 0.0)
    tr = position_in_world(1123.0, 0.0)
    br = position_in_world(1123.0, 791.0)
    bl = position_in_world(0.0, 791.0)

    p1 = π_projection_function(K, R, t, tl)
    p2 = π_projection_function(K, R, t, tr)
    p3 = π_projection_function(K, R, t, br)
    p4 = π_projection_function(K, R, t, bl)
    P = [p1 p2 p3 p4]'

    M = @SMatrix [0.0    0.0 1.0;
                1123.0    0.0 1.0;
                1123.0 791.0 1.0;
                0.0 791.0 1.0]

    res = compute_homograpy_DLT(M, P)
    homography = reshape(res, (3,3))'
    return homography
end

using BenchmarkTools
@btime foo()


function set_volume(nodes, cellxmax, cellymax)
    volume = zeros(cellxmax, cellymax)
    for i in 1:cellxmax
        for j in 1:cellymax
            vec_r1x = nodes[i+1,j+1,1] - nodes[i,j,1]
            vec_r1y = nodes[i+1,j+1,2] - nodes[i,j,2]
            vec_r2x = nodes[i,j+1,1] - nodes[i+1,j,1]
            vec_r2y = nodes[i,j+1,2] - nodes[i+1,j,2]

            volume[i,j] = abs(vec_r1x*vec_r2y - vec_r1y*vec_r2x) /2
        end
    end
    return volume
end

function set_mu(Qbase,cellxmax,cellymax,specific_heat_ratio,Rd)
    mu = zeros(cellxmax,cellymax)

    # サザーランドの式
    # https://www.jstage.jst.go.jp/article/jsam1937/37/4/37_4_694/_pdf/-char/ja
    mu0 = 1.82e-5 # 基準粘度[Pa s]
    T0 = 293.15       # 基準温度[K]
    C = 117       # サザーランド定数[K]

    for i in 1:cellxmax
        for j in 1:cellymax
            # T = p/(rho Rd )
            T = Qbase[i,j,4]/(Qbase[i,j,1]*Rd)

            if T <0
                println("T is munus ( 'A' ) !!")
                println("i : "*string(i)*"  j : "*string(j))
                println(Qbase[i,j,:])
            end

            mu[i,j] = mu0 * (T/T0)^1.5 * (T0+C)/(T+C)
        end
    end

    return mu
end

function set_lambda(Qbase,cellxmax,cellymax,mu,specific_heat_ratio,Rd)

    lambda = zeros(cellxmax,cellymax)

    # サザーランドの式
    lam0 = 22.3*10^(-3)  # 基準熱伝導率　[W/mK]
    T0 = 273.15       # 基準温度[K]
    C = 125       # サザーランド定数[K]

    for i in 1:cellxmax
        for j in 1:cellymax
            T = Qbase[i,j,4]/(Qbase[i,j,1]*Rd)
            lambda[i,j] = lam0*((T0+C)/(T+C))*(T/T0)^1.5
        end
    end

    return lambda
end

function set_Minf(bdcon, specific_heat_ratio)
    rho = 0
    u   = 0
    v   = 0
    p   = 0

    for i in 1:4
        if Int(bdcon[i][1]) == 0 || Int(bdcon[i][1]) == 5
            rho = bdcon[i][2]
            u   = bdcon[i][3]
            v   = bdcon[i][4]
            p   = bdcon[i][5]
        end
    end

    a = (specific_heat_ratio * p / rho)^0.5
    u = (u^2 + v^2)^0.5
    M = u/a

    return M
end

function set_gasconst(Qbase,cellxmax,cellymax,nval,nch,R)
    Rhat = zeros(cellxmax, cellymax)
    mw   = [28e-3, 14e-3]
    npre = nval - nch

    for i in 1:cellxmax
        for j in 1:cellymax
            for s in 1:nch
                Rs = R/mw[s]
                Rhat[i,j] += Qbase[i,j,npre+s] * Rs
            end
        end
    end
    return Rhat
end


# small scrip to compute the solution of Lippmann-Schwinger equation
# We test the types introduced in FastConvolution.jl
# we test that the application is fast and that the construction
# is performed fast.

using PyPlot
using IterativeSolvers
using SpecialFunctions

include("../src/FastConvolution.jl")
include("../src/Preconditioner.jl")

# setting the number of threads for the FFT and BLAS
# libraries (please set them to match the number of
# physical cores in your system)
FFTW.set_num_threads(4);
BLAS.set_num_threads(4);

#Defining Omega
h = 0.005
k = 1/(h)

# size of box
a = 1
x = collect(-a/2:h:a/2)
y = collect(-a/2:h:a/2)
(n,m) = length(x), length(y)
N = n*m
X = repmat(x, 1, m)[:]
Y = repmat(y', n,1)[:]
# we solve \triangle u + k^2(1 + nu(x))u = 0

# We use the modified quadrature in Ruan and Rohklin
(ppw,D) = referenceValsTrapRule();
D0 = D[1];

# Defining the smooth perturbation of the slowness
nu(x,y) = 0.3*exp.(-40*(x.^2 + y.^2)).*(abs.(x).<0.48).*(abs.(y).<0.48);

## You can choose between Duan Rohklin trapezoidal quadrature
#fastconv = buildFastConvolution(x,y,h,k,nu)

# or Greengard Vico Quadrature (this is not optimized and is 2-3 times slower)
fastconv = buildFastConvolution(x,y,h,k,nu, quadRule = "Greengard_Vico");

# assembling the sparsifiying preconditioner
@time As = buildSparseA(k,X,Y,D0, n ,m);

# assembling As*( I + k^2G*nu)
@time Mapproxsp = As + k^2*(buildSparseAG(k,X,Y,D0, n ,m)*spdiagm(nu(X,Y)));

# defining the preconditioner
precond = SparsifyingPreconditioner(Mapproxsp, As)#, solverType="MKLPardiso")

# building the RHS from the incident field
u_inc = exp.(k*im*X);
rhs = -k^2*FFTconvolution(fastconv, nu(X,Y).*u_inc) ;

#rhs = u_inc;

# allocating the solution
u = zeros(Complex128,N);

# solving the system using GMRES
@time info =  gmres!(u, fastconv, rhs, Pl=precond, log = true)
#println(info[2].residuals[:])

# plotting the solution
figure(1)
clf()
imshow(real(reshape(u+u_inc,n,m)))


mview(c, inds, ::Val{1}) = view(c, :, inds[1], :)
mview(c, inds, ::Val{2}) = view(c, :, inds[1], inds[2], :)
mview(c, inds, ::Val{3}) = view(c, :, inds[1], inds[2], inds[3], :)
mview(c, inds, ::Val{4}) = view(c, :, inds[1], inds[2], inds[3], inds[4], :)

tensor_contraction(
    A,
    B,
    ::Val{1},
) = @tullio C[o, a, b] := A[i, a, o] * B[i, a, b]
tensor_contraction(A, B, ::Val{2}) =
    @tullio C[o, a₁, a₂, b] := A[i, a₁, a₂, o] * B[i, a₁, a₂, b]
tensor_contraction(A, B, ::Val{3}) =
    @tullio C[o, a₁, a₂, a₃, b] := A[i, a₁, a₂, a₃, o] * B[i, a₁, a₂, a₃, b]
tensor_contraction(A, B, ::Val{4}) = @tullio C[o, a₁, a₂, a₃, a₄, b] :=
    A[i, a₁, a₂, a₃, a₄, o] * B[i, a₁, a₂, a₃, a₄, b]

sparse_mean(w, c, ::Val{1}) = @tullio μ[o, a, b] := w[i, o] * c[i, a, b]
sparse_mean(
    w,
    c,
    ::Val{2},
) = @tullio μ[o, a₁, a₂, b] := w[i, o] * c[i, a₁, a₂, b]
sparse_mean(w, c, ::Val{3}) =
    @tullio μ[o, a₁, a₂, a₃, b] := w[i, o] * c[i, a₁, a₂, a₃, b]
sparse_mean(w, c, ::Val{4}) =
    @tullio μ[o, a₁, a₂, a₃, a₄, b] := w[i, o] * c[i, a₁, a₂, a₃, a₄, b]

sparse_covariance(
    w,
    c,
    ::Val{1},
) = @tullio μ[o, a, r, b] := w[i, r, o] * c[i, a, b]
sparse_covariance(w, c, ::Val{2}) =
    @tullio μ[o, a₁, a₂, r, b] := w[i, r, o] * c[i, a₁, a₂, b]
sparse_covariance(w, c, ::Val{3}) =
    @tullio μ[o, a₁, a₂, a₃, r, b] := w[i, r, o] * c[i, a₁, a₂, a₃, b]
sparse_covariance(w, c, ::Val{4}) =
    @tullio μ[o, a₁, a₂, a₃, a₄, r, b] := w[i, r, o] * c[i, a₁, a₂, a₃, a₄, b]


using SparseArrays
#------------------------------------------------------------------------------
struct SparseData
    i::Vector{Int}
    j::Vector{Int}
    e::Vector{Int}
    k0::Vector{Float64}
    k::Vector{Float64}
end

#------------------------------------------------------------------------------
struct FEM
    NNode::Int
    Node::Array{Float64,2}
    NElem::Int
    Elem::Array{Int,2}
    ElemLen::Vector{Float64}
    ElemNDof::Vector{Int}
    CumSumElemNDof2::Vector{Int}
    ElemDof::Vector{Vector{Int}}
    Supp::Support
    Load::NodalLoad
    LocalIntMap::Array{Float64,2}
    YoungModulus::Float64
    Force::Vector{Float64}
    FreeDofs::Vector{Int}
    Sparse::SparseData
end

#------------------------------------------------------------------------------
function InitFEA(Node::Array{Float64,2}, Bars::Array{Int,2}, Supp::Support, Load::NodalLoad, E::Float64)::FEM

    NNode, NElem = size(Node, 1), size(Bars, 1)

    D = Node[Bars[:,2],:] - Node[Bars[:,1],:]
    L = vec(sqrt.(sum(D.^2, dims=2))); N = D./L
    LocalInternalMap = [-N N]

    ElemNDof = 4*ones(Int, NElem); CumSumElemNDof2 = cumsum(ElemNDof.^2)
    Ke = zeros(4, 4); Nval::Int = sum(ElemNDof.^2)
    i = zeros(Int, Nval); j = zeros(Int, Nval); e = zeros(Int, Nval); k0 = zeros(Nval)

    idx = 0; ElemDof = Vector{Int}[]; sizehint!(ElemDof, NElem)
    for el = 1:NElem
        LocalK!(Ke, LocalInternalMap[el,:], L[el], E)

        NDof = ElemNDof[el]; Idx = idx+1:idx+NDof^2
        EDof = permutedims([2Bars[el,:].-1 2Bars[el,:]])[:]
        push!(ElemDof, EDof)

        IdxI = repeat(EDof, outer=(1, NDof))
        IdxJ = permutedims(IdxI)
        i[Idx] = IdxI[:]; j[Idx] = IdxJ[:]; k0[Idx] = Ke[:]; e[Idx] .= el
        idx += NDof^2
    end
    k = zeros(Nval)

    F = zeros(2NNode)
    for iL = 1:length(Load.node)
        F[2Load.node[iL] - 1] += Load.val[iL, 1]
        F[2Load.node[iL]] += Load.val[iL, 2]
    end

    NSupp = length(Supp.node)
    FixedDofs, idx = zeros(Int, 2NSupp), 0
    for iS = 1:NSupp
        if Supp.type[iS,1] == 1; idx += 1; FixedDofs[idx] = 2Supp.node[iS]-1; end
        if Supp.type[iS,2] == 1; idx += 1; FixedDofs[idx] = 2Supp.node[iS];   end
    end
    FreeDofs = setdiff(1:2NNode, FixedDofs[1:idx])

    return FEM(NNode, Node, NElem, Bars, L, ElemNDof, CumSumElemNDof2, ElemDof,
    Supp, Load, LocalInternalMap, E, F, FreeDofs, SparseData(i, j, e, k0, k))
end

#------------------------------------------------------------------------------
function LocalK!(Ke::Array{Float64,2}, LocalIntMap::Vector{Float64}, L::Float64, E::Float64)

    NDof = length(LocalIntMap)
    for j = 1:NDof, i = 1:NDof
        Ke[i, j] = E / L * (LocalIntMap[i] * LocalIntMap[j])
    end
end

#------------------------------------------------------------------------------
function FEAnalysis!(U::Vector{Float64}, x::Vector{Float64}, FEA::FEM)

    NDof, Free = 2FEA.NNode, FEA.FreeDofs

    for idx in eachindex(FEA.Sparse.i)
        @fastmath FEA.Sparse.k[idx] = FEA.Sparse.k0[idx] * x[FEA.Sparse.e[idx]]
    end
    @fastmath K = sparse(FEA.Sparse.i, FEA.Sparse.j, FEA.Sparse.k, NDof, NDof)

    @fastmath U[Free] = K[Free, Free] \ FEA.Force[Free]
end

#------------------------------------------------------------------------------
function Stress!(σ::Vector{Float64}, U::Vector{Float64}, FEA::FEM)

    for e = 1:FEA.NElem
        Δ = 0.
        for j = 1:FEA.ElemNDof[e]
            Δ += FEA.LocalIntMap[e, j] * U[FEA.ElemDof[e]][j]
        end
        σ[e] = FEA.YoungModulus / FEA.ElemLen[e] * Δ
    end
end


### A Pluto.jl notebook ###
# v0.14.8

using Markdown
using InteractiveUtils

# ╔═╡ 5ad5c202-20f8-11eb-23f1-4f38b687c285
using STMOZOO.Example

# ╔═╡ d2d007b8-20f8-11eb-0ddd-1181d4565a85
using Plots

# ╔═╡ 45189a82-20fa-11eb-0423-05ce1b84639d
using Zygote

# ╔═╡ 171fee18-20f6-11eb-37e5-2d04caea8c35
md"""
# Example: solving quadratic systems

By Michiel Stock

## Introduction

In this notebook, I have chosen to write code so minimize *quadratic functions* of the form:

$$f(\mathbf{x}) =  \frac{1}{2} \mathbf{x}^\intercal P\mathbf{x} + \mathbf{q} \cdot \mathbf{x} + r\,,$$

where $P \succ 0$.

Hence, solving

$$\min_\mathbf{x}\, f(\mathbf{x})\,.$$
"""

# ╔═╡ fb4aeb8c-20f7-11eb-0444-259de7b76883
md"""

## Derivation

The gradient of the quadratic function is

$$\nabla f(\mathbf{x})=P\mathbf{x} +\mathbf{q}\,.$$

Setting this to zero gives

$$\mathbf{x}^\star=-P^{-1}\mathbf{q}\,.$$

We also know that $\nabla^2f(\mathbf{x}) = P \succ 0$, so $\mathbf{x}^\star$ is a minimizer.
"""

# ╔═╡ 52e87238-20f8-11eb-2ea0-27ee1208d3c3
md"""
## Illustration

We will use a simple 2D example to illustrate this code.
"""

# ╔═╡ e9a9d69e-20f8-11eb-0d9d-330ee5e9cf25
P = [7 -1; -1 2]

# ╔═╡ fdd4e550-20f8-11eb-227b-25f36708484d
q = [-14.0, -2.0]

# ╔═╡ 025fd6e8-20f9-11eb-3e7d-3519f3c4b58f
r = 1

# ╔═╡ 096eff98-20f9-11eb-1e61-99d5714895ba
md"We have a function to generate the quadratic function."

# ╔═╡ 165509ca-20f9-11eb-107c-550cbba0f0e9
f_quadr = quadratic_function(P, q, r)

# ╔═╡ 1fffc82a-20f9-11eb-198c-c160d7dac87d
f_quadr([2, 1])

# ╔═╡ 26ab6ce2-20f9-11eb-1836-1756b290e5e3
md"No more need to remember the formulla for the minimizer! Just use `solve_quadratic_system`!"

# ╔═╡ 49832a8e-20f9-11eb-0841-19a40a12db18
x_star = solve_quadratic_system(P, q, r)

# ╔═╡ 55e0e274-20f9-11eb-36c0-753f228f7e9b
begin
	contourf(-20:0.1:20, -20:0.1:20, (x, y) -> f_quadr([x,y]), color=:speed)
	scatter!([x_star[1]], [x_star[2]], label="minimizer")
end

# ╔═╡ b1551758-20f9-11eb-3e8f-ff9a7127d7f8
md"""
## Approximating non-quadratic functions

We can approximate non-quadratic functions by a quadratic function: The second order Taylor approximation $\hat{f}$ of a function $f$ at $\mathbf{x}$ is
$$f(\mathbf{x}+\mathbf{v})\approx\hat{f}(\mathbf{x}+\mathbf{v}) = f(\mathbf{x}) + \nabla f(\mathbf{x})^\top \mathbf{v} + \frac{1}{2} \mathbf{v}^\top \nabla^2 f(\mathbf{x}) \mathbf{v}\,.$$

Let us use this idea for the Rosenbrock function.
"""

# ╔═╡ 41d8f1dc-20fa-11eb-3586-a989427c1fd6
f_nonquadr((x1, x2); a=1, b=5) = (a-x1)^2 + b * (x2 - x1^2)^2

# ╔═╡ 4ed4215e-20fa-11eb-11ee-f7741591163c
x = [0.0, 0.0]

# ╔═╡ 56af99ee-20fa-11eb-0240-69c675efb78c
fx = f_nonquadr(x)

# ╔═╡ 6c5473b4-20fa-11eb-327b-51ac560530eb
∇fx = f_nonquadr'(x)

# ╔═╡ 7518c2c0-20fa-11eb-32c0-a9db2a91cbc5
∇²fx = Zygote.hessian(f_nonquadr, x)

# ╔═╡ 34027942-20fb-11eb-261e-3b991ce4c9f8
v = solve_quadratic_system(∇²fx, ∇fx, fx)

# ╔═╡ 3bbeb85c-20fc-11eb-04d0-fb12d8ace50a
f̂(x′) = quadratic_function(∇²fx, ∇fx, fx)(x′ .- x)

# ╔═╡ 8623ac1a-20fa-11eb-2d45-49cce0fdac86
begin
	plot_nonquadr = contourf(-2:0.01:2, -2:0.01:2, (x, y) -> f_nonquadr([x,y]), color=:speed, title="non-quadratic function")
	scatter!(plot_nonquadr, [x[1]], [x[2]], label="x")
	scatter!(plot_nonquadr, [x[1]+v[1]], [x[2]+v[2]], label="x + v")

	plot_approx = contourf(-2:0.01:2, -2:0.01:2, (x, y) -> f̂([x,y]), color=:speed,
		title="quadratic approximation")
	scatter!(plot_approx, [x[1]], [x[2]], label="x")
	scatter!(plot_approx, [x[1]+v[1]], [x[2]+v[2]], label="x + v")

	plot(plot_nonquadr, plot_approx, layout=(2,1), size=(600, 800))

end

# ╔═╡ Cell order:
# ╟─171fee18-20f6-11eb-37e5-2d04caea8c35
# ╠═5ad5c202-20f8-11eb-23f1-4f38b687c285
# ╠═d2d007b8-20f8-11eb-0ddd-1181d4565a85
# ╟─fb4aeb8c-20f7-11eb-0444-259de7b76883
# ╟─52e87238-20f8-11eb-2ea0-27ee1208d3c3
# ╠═e9a9d69e-20f8-11eb-0d9d-330ee5e9cf25
# ╠═fdd4e550-20f8-11eb-227b-25f36708484d
# ╠═025fd6e8-20f9-11eb-3e7d-3519f3c4b58f
# ╟─096eff98-20f9-11eb-1e61-99d5714895ba
# ╠═165509ca-20f9-11eb-107c-550cbba0f0e9
# ╠═1fffc82a-20f9-11eb-198c-c160d7dac87d
# ╟─26ab6ce2-20f9-11eb-1836-1756b290e5e3
# ╠═49832a8e-20f9-11eb-0841-19a40a12db18
# ╠═55e0e274-20f9-11eb-36c0-753f228f7e9b
# ╟─b1551758-20f9-11eb-3e8f-ff9a7127d7f8
# ╠═41d8f1dc-20fa-11eb-3586-a989427c1fd6
# ╠═45189a82-20fa-11eb-0423-05ce1b84639d
# ╠═4ed4215e-20fa-11eb-11ee-f7741591163c
# ╠═56af99ee-20fa-11eb-0240-69c675efb78c
# ╠═6c5473b4-20fa-11eb-327b-51ac560530eb
# ╠═7518c2c0-20fa-11eb-32c0-a9db2a91cbc5
# ╠═34027942-20fb-11eb-261e-3b991ce4c9f8
# ╠═3bbeb85c-20fc-11eb-04d0-fb12d8ace50a
# ╟─8623ac1a-20fa-11eb-2d45-49cce0fdac86


module AdaMeanShift

    using Statistics, LinearAlgebra, SpecialFunctions, StaticArrays
    export atomic_meanshift, atomic_intensity, atomic_ratio, meanshift!, meanshift_nonadaptive!, extract_circ_region, extract_region
    @inline function Kern(x::AbstractVector{fT},h::AbstractVector{fT}) where fT
        max(zero(fT),one(fT) - norm(@inbounds x ./ max.(h,one(fT)))^2)
    end
    function clampdim(M,c::T,i)::T where T
        clamp(c,1,size(M,i))
    end
    function makerange(M,p,h,i)
        #round(Int,p[i]-h[i]-1):round(Int,p[i]+h[i]+1)
        @inbounds clampdim(M,round(Int,p[i]-h[i]-1),i):clampdim(M,round(Int,p[i]+h[i]+1),i)
    end

"""
        atomic_meanshift(M,p,h,isotropy,smoothing) -> NamedTuple(pnew,hnew,modeval,delta)

Performs a single iteration of meanshift over a single particle

# Arguments
- `M` density tensor
- `p` initial position for test particle
- `h` initial particle standard deviation
- `isotropy` scalar 0 -> 1, 0 is completely anisotropic, completely isotropic
- `smoothing` scalar 0 -> 1, regularization term for noisy density tensors

"""
    function atomic_meanshift(M,p,h,isotropy,smoothing)
        N=length(p)
        T=eltype(p)
        @assert N==length(h)
        @assert T==eltype(h)
        atomic_meanshift(M,SVector{length(p)}(p),SVector{length(h)}(h),T(isotropy),T(smoothing))
    end
    @generated function atomic_meanshift(M::Mty,p::K,h::K,isotropy::T,smoothing::T) where {T,N,K<:StaticArray,Mty<:AbstractArray{T,N}}
        D=(8//N^2 + 6//N +1)
        scale= sqrt(N * ((2//N + 1)*(4//N + 1)*(6//N + 1))//D)
        invvolume=gamma(N/2+1)*pi^(-N/2)
        modescale=invvolume*N^3*D
        Z=zero(T)
        O=one(T)
        quote
            @inbounds begin
                c=p.*$Z
                xsq=p.*$Z
                NW=IW=KW=W=$Z
                ls=$invvolume*inv(prod(max.(h,$O)))*smoothing

                Base.Cartesian.@nexprs $N i -> R_i = makerange(M,p,h,i)
                Base.Cartesian.@nloops $N i i->R_i begin
                    cp=$K(Base.Cartesian.@ntuple($N,k->i_k))
                    diff=p.-cp
                    keW=Kern(diff,h)
                    imW=Base.Cartesian.@nref($N,M,k->i_k) + ls
                    sW=keW*imW
                    isofactor = isotropy*norm(diff)^2/$N
                    xsq=xsq .+ (diff .^ 2 .*(1.0-isotropy) .+isofactor) .* sW
                    c=c.+cp.*sW
                    W+=sW
                    KW+=keW
                    IW+=imW
                    NW+=$O
                end
                if 2*NW<1
                    return (pnew=p,
                            hnew=h,
                            modeval=T(NaN),
                            delta=$Z)
                end
                xsq=xsq ./ W

                c=@. p+(c/W - p)/2
                return (pnew=K(Base.Cartesian.@ntuple($N,k->clampdim(M,c[k],k))),
                        hnew=max.(1,$scale .* sqrt.(xsq)),
                        modeval=W/prod(max.(h,$O))*$modescale,
                        delta=norm(c.-p))
            end
        end
    end

"""
        atomic_intensity(M,p,h) -> NamedTuple(intensity,numpoints,ndims)

Calculates the integrated intensity on the density tensor `M` over a ellipsoidal region caracterized by `p` region center, `h` semiaxis

# Arguments
- `M` density tensor (e.g. matrix N*N)
- `p` region center (e.g. 2D vector)
- `h` semiaxis lengths (e.g. 2D vector)

Note that both `p`,`h` should be of type Vector or StaticVector with the same length.
"""
    function atomic_intensity(M,p,h)
        N=length(p)
        T=eltype(p)
        @assert N==length(h)
        @assert T==eltype(h)
        atomic_intensity(M,SVector{length(p)}(p),SVector{length(h)}(h))
    end

    @generated function atomic_intensity(M::Mty,p::K,h::K) where {T,N,K<:StaticArray,Mty<:AbstractArray{T,N}}
        Z=zero(T)
        O=one(T)
        quote
            @inbounds begin
                NP=IP=$Z
                Base.Cartesian.@nexprs $N i -> R_i = makerange(M,p,h,i)
                Base.Cartesian.@nloops $N i i->R_i begin
                    cp=$K(Base.Cartesian.@ntuple($N,k->i_k))
                    norm((p.-cp)./h) > 1 && continue;
                    IP+=Base.Cartesian.@nref($N,M,k->i_k)
                    NP+=$O
                end
                return (intensity=IP,numpoints=NP,ndims=N)
            end
        end
    end

using ProgressMeter

    function atomic_expmeanlog(M,p,h)
        N=length(p)
        T=eltype(p)
        @assert N==length(h)
        @assert T==eltype(h)
        atomic_expmeanlog(M,SVector{length(p)}(p),SVector{length(h)}(h))
    end

    @generated function atomic_expmeanlog(M::Mty,p::K,h::K) where {T,N,K<:StaticArray,Mty<:AbstractArray{T,N}}
        Z=zero(T)
        O=one(T)
        quote
            @inbounds begin
                NP2=NP=IP=IP2=$Z
                maxnum=maximum(M)
                Base.Cartesian.@nexprs $N i -> R_i = makerange(M,p,h,i)
                Base.Cartesian.@nloops $N i i->R_i begin
                    cp=$K(Base.Cartesian.@ntuple($N,k->i_k))
                    norm((p.-cp)./h)>1 && continue;
                    I=Base.Cartesian.@nref($N,M,k->i_k)
                    I<=0 && continue;
                    wh=(I/maxnum)^2   # pow=2 true est, pow=1 seems good too
                    X=log(I)
                    IP+=X*wh
                    IP2+=X^2*wh
                    NP+=wh
                    NP2+=wh^2
                end
                #@show IP,IP2,NP,NP2
                IP/=NP
                IP2/=NP
                XV=exp(IP)
                XD=XV*sqrt(max(IP2-IP^2,$Z))*sqrt(NP2)/NP
                return (XV,XD)
            end
        end
    end

"""
        atomic_ratio(Mnum,Mden,p,h) -> (pointest=(ratio,stddev),ndims=N)

Calculates the pixel pair-wise ratio on for density tensor `Mnum/Mden` over a ellipsoidal region caracterized by `p` region center, `h` semiaxis

# Arguments
- `Mnum` density tensor (e.g. matrix N*N)
- `Mden` density tensor (e.g. matrix N*N)
- `p` region center (e.g. 2D vector)
- `h` semiaxis lengths (e.g. 2D vector)

Note that both `p`,`h` should be of type Vector or StaticVector with the same length.
"""
    function atomic_ratio(Mnum,Mden,p,h)
        n,σn=atomic_expmeanlog(Mnum,p,h)
        d,σd=atomic_expmeanlog(Mden,p,h)
        (n/d, sqrt(σn^2+(σd*n/d)^2)/d)
    end

    function meanshift_kernel!(::Val{isadaptive} ,M::Mty,P::AbstractVector{K},h::AbstractVector{K},w::AbstractVector{T}, hmax::T, isotropy::T,maxit::T,rtol::T,smoothing::T) where {isadaptive,T,N,K<:StaticArray,Mty<:AbstractArray{T,N}}
        pr = Progress(length(P));
        update!(pr,0)
        jj = Threads.Atomic{Int}(0)
        a=@elapsed Threads.@threads for i in eachindex(P)
            cnt=0
            delta=one(T)
            bufcnt=0
            while (cnt<maxit) & (bufcnt < 10)
                if norm(h[i])>hmax && isadaptive
                    h[i]=h[i] ./ norm(h[i]) .* hmax
                end
                data=atomic_meanshift(M,P[i],h[i],isotropy,smoothing)
                if isadaptive
                    @inbounds P[i],h[i],w[i],delta=data
                else
                    @inbounds P[i],_,w[i],delta=data
                end
                cnt+=1
                bufcnt=ifelse(delta<rtol,bufcnt+1,0)
            end
            Threads.atomic_add!(jj, 1)
            Threads.threadid() == 1 && update!(pr, jj[])
        end
        isadaptive && Threads.@threads for i in eachindex(P)
                if norm(h[i])>hmax
                    h[i]=h[i] ./ norm(h[i]) .* hmax
                end
        end
        finish!(pr)
        a
    end

"""
        meanshift!(M, P, h, w, hmax; isotropy=½, maxit=∞, rtol=√ϵ, smoothing=1)
        meanshift_nonadaptive!(M, P, h, w; maxit=∞, rtol=√ϵ, smoothing=1)

Performs MeanShift on a swarm of particles over a given density matrix locating all modes and their scale. The non adaptive method disables updating of the `h` vectors, and thus depends on fewer parameters.

# Arguments
- `M` is the density tensor over which particles evolve.
- `P` is a julia vector of StaticVectors describing all positions.
- `h` is a julia vector of StaticVectors describing all standard deviations.
- `w` is a julia vector which will contain the new modes intensity estimates.
- `hmax` is a real number for the maximum norm allowed to the `h` vectors.

# Keyword Arguments
- `isotropy` is a scalar [0,1] set to 1 for isotropic kernel, < 1 anisotropy along coordinate axis is allowed.
- `maxit` is a integer for the maximum number of meanshift iterations.
- `rtol` is the absolute tolerance to declare a particle as converged.
- `smoothing` is a regularization term for density tensors with noise.

"""
    function meanshift!(M,P,h,w, hmax; isotropy = 0.5,maxit=Inf,rtol=1e-8,smoothing=1.0)
        meanshift_kernel!(Val(true),M,P,h,w, promote(hmax, isotropy,maxit,rtol,smoothing)...)
    end
    function meanshift_nonadaptive!(M,P,h,w; maxit=Inf,rtol=1e-8,smoothing=1.0)
        meanshift_kernel!(Val(false),M,P,h,w, promote(0, 0,maxit,rtol,smoothing)...)
    end
"""
    extract_region(M,p,h) -> M',p',h
Extracts a rectangular section with center "p" and axes "h" from a multidimensional array and returns the region with the new center coordinates.
"""
    function extract_region(M::Mty,p::Vty,h::Vty) where {T,N,Mty<:AbstractArray{T,N},Vty<:AbstractVector{T}}
        rg=[makerange(M,p,h,i) for i in 1:N]
        redM=collect(view(M,rg...))
        startind=minimum.(rg)
        redp=p.-startind.+1
        redM,redp,h
    end
"""
    extract_circ_region(M,p,h) -> M',p',h
Extracts a spherical section with center "p" and axes "h" from a multidimensional array and returns the region with the new center coordinates.
"""
    function extract_circ_region(M::Mty,p::Vty,h::Vty) where {T,N,Mty<:AbstractArray{T,N},Vty<:AbstractVector{T}}
        redM,redp,redh=extract_region(M,p,h)
        for i in CartesianIndices(redM)
            if norm((Tuple(i).-redp)./redh) > 1
                redM[i] =0.0
            end
        end
        redM,redp,redh
    end
end


# ------------------------------------------------------------------------------------------
# ## 函数
#
# 话题：
# 1. 如何声明一个函数
# 2. Julia 中的鸭子类型（Duck-typing）
# 3. 原地修改与非原地修改函数（Mutating vs. non-mutating functions）
# 4. 高阶函数
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ### 如何声明一个函数
#
# Julia 提供了几种不同的声明函数的方法。
#
# #### 方法一：用 `function` 和 `end` 关键字
# ------------------------------------------------------------------------------------------

function sayhi(name)
    println("Hi $name, it's great to see you!")
end

function f(x)
    x^2
end

# ------------------------------------------------------------------------------------------
# 我们可以这样调用函数：
# ------------------------------------------------------------------------------------------

sayhi("C-3PO")

f(42)

# ------------------------------------------------------------------------------------------
# #### 方法二：用 `=`
# 等价的，我们可以只用一行声明函数，拿上面两个函数为例
# ------------------------------------------------------------------------------------------

sayhi2(name) = println("Hi $name, it's great to see you!")

f2(x) = x^2

sayhi2("R2D2")

f2(42)

# ------------------------------------------------------------------------------------------
# #### 方法三：匿名函数
# 最后，我们可以声明“匿名”函数
# ------------------------------------------------------------------------------------------

sayhi3 = name -> println("Hi $name, it's great to see you!")

f3 = x -> x^2

sayhi3("Chewbacca")

f3(42)

# ------------------------------------------------------------------------------------------
# ### Julia 中的鸭子类型（Duck-typing）
# > *“假如它嘎嘎像个鸭鸭，那它就是个鸭鸭”*
#
# 在未指定函数参数的类型时，不管输入参数是什么类型，只要函数体可以正常执行，
# 那么 Julia 函数就可以应用在这种类型的参数上
#
#
#
# 例如，函数 `sayhi` 可以接受这个写作整数的小电视角色的名字……
# ------------------------------------------------------------------------------------------

sayhi(55595472)

# ------------------------------------------------------------------------------------------
# 函数 `f` 可以接受一个矩阵
# ------------------------------------------------------------------------------------------

A = rand(3, 3)
A

f(A)

# ------------------------------------------------------------------------------------------
# 函数 `f` 也可以接受字符串比如"hi"，因为 `*` 在接受字符串时用作字符串拼接。
# ------------------------------------------------------------------------------------------

f("hi")

# ------------------------------------------------------------------------------------------
# 然而，`f` 函数无法接受一个向量（vector）。
# 不同于 `A^2` 是已经定义好的矩阵乘方，当 `v` 是向量时的 `v^2` 是未定义的，会导致错误。
# ------------------------------------------------------------------------------------------

v = rand(3)

f(v)

# ------------------------------------------------------------------------------------------
# ### 原地修改与非原地修改函数（Mutating vs. non-mutating functions）
#
# 约定：函数名以 `!` 结束的原地修改传入的变量，函数名非以 `!` 结束的则不会改变传入的变量。
#
# 例如下面的 `sort` 和 `sort!`
# ------------------------------------------------------------------------------------------

v = [3, 5, 2]

sort(v)

v

# ------------------------------------------------------------------------------------------
# `sort(v)`会返回一个新的、将 `v` 排序后的数组，`v` 本身不会改变
#
# 而当执行 `sort!(v)` 后，`v` 将会被原地排序修改
# ------------------------------------------------------------------------------------------

sort!(v)

v

# ------------------------------------------------------------------------------------------
# ### 一些高阶函数
#
# #### `map`
#
# `map` 是 Julia 中的高阶函数。“高阶”指的是它会*接受一个函数*作为它的传入参数。
# `map` 将传入的函数作用于传入的数据结构的每一个元素。例如执行
#
# ```julia
# map(f, [1, 2, 3])
# ```
# 将得到一个输出数组，数组的每一个元素为 `f` 作用在数组 `[1, 2, 3]` 的对应元素
# ```julia
# [f(1), f(2), f(3)]
# ```
# ------------------------------------------------------------------------------------------

map(f, [1, 2, 3])

# ------------------------------------------------------------------------------------------
# 这里我们把向量 `[1, 2, 3]` 的每一个元素都平方了，而不是对向量 `[1, 2, 3]` 本身进行平方。
#
# `map` 可以接受匿名函数作为参数，像这样
# ------------------------------------------------------------------------------------------

x -> x^3

# ------------------------------------------------------------------------------------------
# 通过
# ------------------------------------------------------------------------------------------

map(x -> x^3, [1, 2, 3])

# ------------------------------------------------------------------------------------------
# 这就完成了对数组 `[1, 2, 3]` 所有元素立方的计算！
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ### 广播（broadcast）
#
# 广播（`broadcast`）和 `map` 一样也是高阶函数。
# 广播是一种广义化的 `map`，所以 `map` 能做的它都能做。并不止如此，广播的调用方法和 `map` 一样。
# ------------------------------------------------------------------------------------------

broadcast(f, [1, 2, 3])

# ------------------------------------------------------------------------------------------
# 我们又一次通过函数 `f` 把向量 `[1, 2, 3]` 的每一个元素都平方了
# ——这次是用“广播”函数 `f` 的方式！
#
# 广播（`broadcast`）有个语法糖：调用要广播的函数时，在函数名和输入变量之间加一个 `.`。
#
# 比如：
# ```julia
# broadcast(f, [1, 2, 3])
# ```
# 就等价于
# ```julia
# f.([1, 2, 3])
# ```
# ------------------------------------------------------------------------------------------

f.([1, 2, 3])

# ------------------------------------------------------------------------------------------
# 再一次强调 `map` 和广播方式与如下直接调用函数的区别。
#
# ```julia
# f([1, 2, 3])
# ```
# 我们可以求向量中每个元素的平方，但是不能求向量的平方！
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# 为了再讲清楚这一点，我们再看一下以下调用的区别
#
# ```julia
# f(A)
# ```
# 与
# ```julia
# f.(A)
# ```
# 其中 `A` 是个矩阵：
# ------------------------------------------------------------------------------------------

A = [i + 3*j for j in 0:2, i in 1:3]

f(A)

# ------------------------------------------------------------------------------------------
# 和之前一样，对于矩阵 `A` 来说
# ```
# f(A) = A^2 = A * A
# ```
#
# 另一方面，
# ------------------------------------------------------------------------------------------

B = f.(A)

# ------------------------------------------------------------------------------------------
# 结果包含了矩阵 `A` 中每个元素的平方
#
# 广播的点语法让比较复杂的元素间运算的表示更自然、更贴近数学表示。
#
# 例如：
# ------------------------------------------------------------------------------------------

A .+ 2 .* f.(A) ./ A

# ------------------------------------------------------------------------------------------
# 而不是
# ------------------------------------------------------------------------------------------

broadcast(x -> x + 2 * f(x) / x, A)

# ------------------------------------------------------------------------------------------
# 这两种方法本质上完全一样
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ### 练习
#
# #### 6.1
# 声明一个函数 `add_one` 用来给它的输入参数加1。
# ------------------------------------------------------------------------------------------

function add_one(n)
    # 请在此编写你的代码
    n
end

@assert add_one(1) == 2

@assert add_one(11) == 12

# ------------------------------------------------------------------------------------------
# #### 6.2
# 使用函数`map`或者`broadcast`给矩阵`A`的每一个元素加1并把它赋值给`A1`。
# ------------------------------------------------------------------------------------------

# 初始化 A。请勿更改
A = [i + 3*j for j in 0:2, i in 1:3]

# 请在此编写你的代码
A1 = map(x->x, A)

@assert A1 == [2 3 4; 5 6 7; 8 9 10]

# ------------------------------------------------------------------------------------------
# #### 6.3
# 使用广播的点语法给矩阵`A`的每一个元素加1并把它赋值给`A2`。
# ------------------------------------------------------------------------------------------

# 请在此编写你的代码
A2 = A1

@assert A2 == [3 4 5; 6 7 8;9 10 11]

# ------------------------------------------------------------------------------------------
# 请在完成练习运行 `@assert` 块以检测结果是否正确，若无报错即为正确。
# ------------------------------------------------------------------------------------------


@testset "2d weights" begin
    simplex = SVector{3}(SVector{2, Float64}[[1., -1], [1., 1], [-1.5, 0]])
    weights = projection_weights(simplex)
    @test isapprox(weights, [0.3, 0.3, 0.4])
    @test isapprox(weights, projection_weights_reference(simplex))

    simplex = SVector{3}(SVector{2, Float64}[[1., -1], [1., 1], [1.5, 0]])
    weights = projection_weights(simplex)
    @test isapprox(weights, [0.5, 0.5, 0.0])
    @test isapprox(weights, projection_weights_reference(simplex))

    simplex = SVector{3}(SVector{2, Float64}[[1., -1], [1., 1], [0.5, 0]])
    weights = projection_weights(simplex)
    @test isapprox(weights, [0.0, 0.0, 1.0])
    @test isapprox(weights, projection_weights_reference(simplex))
end

@testset "3d weights" begin
    simplex = SVector{4}(SVector{3, Float64}[[1., -1, 0], [1., 1, 0], [0.4, 0, -0.1], [0.5, 0, 1]])
    weights = projection_weights(simplex)
    @test isapprox(weights, [0.0, 0.0,0.942623, 0.057377], atol=1e-6)
    @test isapprox(weights, projection_weights_reference(simplex))

    simplex = SVector{4}(SVector{3, Float64}[[1., -1, 0], [1., 1, 0], [1.4, 0, -0.1], [1.5, 0, 1]])
    weights = projection_weights(simplex)
    @test isapprox(weights, [0.5, 0.5, 0.0, 0.0])
    @test isapprox(weights, projection_weights_reference(simplex))
end

@testset "random simplex" begin
    Random.seed!(1)
    for i in 1:100
        simplex = SVector{4}([rand(SVector{3, Float64}) for i in 1:4])
        weights = projection_weights(simplex)
        @test all(weights .>= 0)
        @test isapprox(sum(weights), 1)
        @test isapprox(weights, projection_weights_reference(simplex))
    end
end

function test_no_penetration(simplex)
    weights = projection_weights(simplex)
    if all(x -> x > 0, weights)
        @show weights
        @test false
    end
end

@testset "numerical issues" begin
    # These simplices are degenerate in that all z-coordinates are the same.
    # Numerical issues due to floating point comparison can cause `projection_weights`
    # to return positive weights.
    let simplex = SVector{4, SVector{3, Float64}}([-3.33728783796428, 0.3321305518800686, 0.11004228580261734], [-3.33728783796428, -0.6678694481199314, 0.11004228580261734], [1.6627121620357201, -0.6678694481199314, 0.11004228580261734], [1.6627121620357201, 0.3321305518800686, 0.11004228580261734])
        test_no_penetration(simplex)
    end
    let simplex = SVector{4, SVector{3, Float64}}([-2.362029787803838, -0.4737843508044094, 0.11167470708499053], [2.637970212196162, 0.5262156491955906, 0.11167470708499053], [2.637970212196162, -0.4737843508044094, 0.11167470708499053], [-2.362029787803838, 0.5262156491955906, 0.11167470708499053])
        test_no_penetration(simplex)
    end
end