include("tools.jl")
using Main.tools
using LinearAlgebra

# let the JIT Compiler do some work
# global n=50
# global dims = [n,n,n]
# global rhs = rhs_norm(n)
# global rhs = rhs_sin(n)
# global frames =  Array{Array{Float64,2},1}(undef, 3)
# global rhs_u = unfolding_fun(rhs, 1, dims)
# global dims_u = [dims[1], prod(dims[2:end]) ]
# global I,J = aca_fun(rhs_u, dims_u)
# global C,U,R = cur_fun(rhs_u, dims_u, I, J)
# frames[1] = C
# global dims[1] = length(I)
# global core = folding(U*R, 1, dims)
# for i = 2:3
#     # global core_u = unfolding(core, i)
#     # global I,J = aca(core_u)
#     # global C, U, R = cur(core_u, I, J)
#     # frames[i] = C
#     # dims[i] = length(I)
#     # global core = folding(U*R, i, dims)
#     # GC.gc()
#     I, J, C, U, R = do_it(core, dims, i)
# end

function do_calc(core_fun, dims)
    frames =  Array{Array{Float64,2},1}(undef, 3)
    rhs_u = unfolding_fun(rhs, 1, dims)
    dims_u = [dims[1], prod(dims[2:end]) ]
    I,J = aca_fun(rhs_u, dims_u)
    C,U,R = cur_fun(rhs_u, dims_u, I, J)
    frames[1] = C
    dims[1] = length(I)
    core = folding(U*R, 1, dims)
    for i = 2:3
        core_u = unfolding(core, i)
        I,J = aca(core_u)
        C,U,R = cur(core_u, I, J)
        frames[i] = C
        dims[i] = length(I)
        core = folding(U*R, i, dims)
        GC.gc()
    end
    return (core, frames, dims)
end

n=50
dims = [n,n,n]
rhs = rhs_norm(n)
rhs = rhs_sin(n)

core, frames, dims = do_calc(rhs, dims)
println(dims)

# n=50
# dims = [n,n,n]
# rhs = rhs_norm(n)
# rhs = rhs_sin(n)
# frames =  Array{Array{Float64,2},1}(undef, 3)
# rhs_u = unfolding_fun(rhs, 1, dims)
# dims_u = [dims[1], prod(dims[2:end]) ]
# I,J = aca_fun(rhs_u, dims_u)
# C,U,R = cur_fun(rhs_u, dims_u, I, J)
# frames[1] = C
# dims[1] = length(I)
# core = folding(U*R, 1, dims)
# for i = 2:3
#     # global core_u = unfolding(core, i)
#     # global I,J = aca(core_u)
#     # global C, U, R = cur(core_u, I, J)
#     # frames[i] = C
#     # dims[i] = length(I)
#     # global core = folding(U*R, i, dims)
#     # GC.gc()
#     I, J, C, U, R = do_it(core, dims, i)
# end

for n = [200,300,400]
    println("=====================")
    println("using n=", n)

    println()
    println("First example with sinus")
    global dims = [n,n,n]
    global rhs = rhs_sin(n)
    println("Starting approximation")

    # tic()
    start = time_ns()
    # calculate the approximation
    # global frames =  Array{Array{Float64,2},1}(undef, 3)
    # global rhs_u = unfolding_fun(rhs, 1, dims)
    # global dims_u = [dims[1], prod(dims[2:end]) ]
    # global I,J = aca_fun(rhs_u, dims_u)
    # global C,U,R = cur_fun(rhs_u, dims_u, I, J)
    # frames[1] = C
    # dims[1] = length(I)
    # global core = folding(U*R, 1, dims)
    # for i = 2:3
    #     global core_u = unfolding(core, i)
    #     global I,J = aca(core_u)
    #     global C, U, R = cur(core_u, I, J)
    #     frames[i] = C
    #     dims[i] = length(I)
    #     global core = folding(U*R, i, dims)
    #     GC.gc()
    # end
    core, frames, dims = do_calc(rhs, dims)
    # toc()
    elapsed_time = time_ns() - start
    println("elapsed time: ", elapsed_time*1e-9)
    ten = tten(core, frames)
    rhs_full = get_rhs_sin(n)
    err = norm( (rhs_full - fullten(ten))[:])

    println("Resulting ranks are: ", size(ten.core))
    println("Relative Error: ", err/norm(rhs_full[:]))

    println()
    println("Second example with norm")
    dims = [n,n,n]
    rhs = rhs_norm(n)
    println("Starting approximation")
    # tic()
    start = time_ns()
    # calculate the approximation
    # frames =  Array{Array{Float64,2},1}(undef, 3)
    # rhs_u = unfolding_fun(rhs, 1, dims)
    # dims_u = [dims[1], prod(dims[2:end]) ]
    # global I,J = aca_fun(rhs_u, dims_u)
    # global C,U,R = cur_fun(rhs_u, dims_u, I, J)
    # frames[1] = C
    # dims[1] = length(I)
    # global core = folding(U*R, 1, dims)
    # for i = 2:3
    #     global core_u = unfolding(core, i)
    #     global I,J = aca(core_u)
    #     global C, U, R = cur(core_u, I, J)
    #     frames[i] = C
    #     dims[i] = length(I)
    #     global core = folding(U*R, i, dims)
    #     GC.gc()
    # end
    core, frames, dims = do_calc(rhs, dims)
    # toc()
    elapsed_time = time_ns() - start
    println("elapsed time: ", elapsed_time*1e-9)
    ten = tten(core, frames)
    rhs_full = get_rhs_norm(n)
    err = norm( (rhs_full - fullten(ten))[:])

    println("Resulting ranks are: ", size(ten.core))
    println("Relative Error: ", err/norm(rhs_full[:]))
end

for n = [500,600,700]
    println("=====================")
    println("using n=", n)

    println()
    println("First example with sinus")
    dims = [n,n,n]
    rhs = rhs_sin(n)
    println("Starting approximation")

    # tic()
    start = time_ns()
    # calculate the approximation
    # frames =  Array{Array{Float64,2},1}(undef, 3)
    # rhs_u = unfolding_fun(rhs, 1, dims)
    # dims_u = [dims[1], prod(dims[2:end]) ]
    # global I,J = aca_fun(rhs_u, dims_u)
    # global C,U,R = cur_fun(rhs_u, dims_u, I, J)
    # frames[1] = C
    # dims[1] = length(I)
    # global core = folding(U*R, 1, dims)
    # for i = 2:3
    #     core_u = unfolding(core, i)
    #     global I,J = aca(core_u)
    #     global C, U, R = cur(core_u, I, J)
    #     frames[i] = C
    #     dims[i] = length(I)
    #     global core = folding(U*R, i, dims)
    #     GC.gc()
    # end
    core, frames, dims = do_calc(rhs, dims)
    # toc()
    elapsed_time = time_ns() - start
    println("elapsed time: ", elapsed_time*1e-9)
    ten = tten(core, frames)
    println("Resulting ranks are: ", size(ten.core))

    println()
    println("Second example with norm")
    dims = [n,n,n]
    rhs = rhs_norm(n)
    println("Starting approximation")
    # tic()
    start = time_ns()
    # calculate the approximation
    # frames =  Array{Array{Float64,2},1}(undef, 3)
    # rhs_u = unfolding_fun(rhs, 1, dims)
    # dims_u = [dims[1], prod(dims[2:end]) ]
    # global I,J = aca_fun(rhs_u, dims_u)
    # global C,U,R = cur_fun(rhs_u, dims_u, I, J)
    # frames[1] = C
    # dims[1] = length(I)
    # global core = folding(U*R, 1, dims)
    # for i = 2:3
    #     core_u = unfolding(core, i)
    #     global I,J = aca(core_u)
    #     global C, U, R = cur(core_u, I, J)
    #     frames[i] = C
    #     dims[i] = length(I)
    #     global core = folding(U*R, i, dims)
    #     GC.gc()
    # end
    core, frames, dims = do_calc(rhs, dims)
    # toc()
    elapsed_time = time_ns() - start
    println("elapsed time: ", elapsed_time*1e-9)
    ten = tten(core, frames)
    println("Resulting ranks are: ", size(ten.core))
end

GC.gc()


#
# test hierarchicalsolve
#
using ApproxFun, SingularIntegralEquations, Base.Test

M = 50
L = Array(LowRankMatrix{Float64},14)
[L[i] = LowRankMatrix(0.001rand(4M,2),0.001rand(4M,2)) for i=1:2]
[L[i] = LowRankMatrix(0.01rand(2M,2),0.01rand(2M,2)) for i=[3:4;9:10]]
[L[i] = LowRankMatrix(0.1rand(M,2),0.1rand(M,2)) for i=[5:8;11:14]]
D = Array(Diagonal{Float64},8)
[D[i] = Diagonal(1./collect(1:M)) for i=1:8]

H = HierarchicalMatrix(D,L)

@test rank(H) == 400
#@test blockrank(H) == fill(2,8,8)+diagm(fill(48,8))
@test isfactored(H) == false

B = Array(Vector{Float64},8)
[B[i] = rand(M) for i=1:8]
b = rand(size(H,1))#HierarchicalVector(B)

full(H)\b
@time x = full(H)\b
H\b # includes compile time
H = HierarchicalMatrix(D,L)
@time xw = H\b # includes precomputation
@time xw = H\b # H.A's precomputed and factored :)

@test isfactored(H) == true

CH = cond(H)

@test norm(x-xw) ≤ 10CH^2*eps()
@test norm(H*x-b) ≤ 10CH*eps()
@test norm(H*xw-b) ≤ 10CH*eps()

H+H-(H-H)
@test norm(full(H-H)) < 10norm(full(H))*eps()

# Hierarchical operators case

i = 2

dom = cantor(Interval(),i)
⨍ = DefiniteLineIntegral(dom)
f = Fun(x->logabs(x-5im),dom)
sp = Space(dom)

G = GreensFun((x,y)->1/2,CauchyWeight(sp⊗sp,0);method=:Cholesky)

@time u1 = ⨍[G]\transpose(f)
println("Adaptive QR  forward error norm is: ",norm(⨍[G]*u1-f))

@test norm(⨍[G]*u1-f) < 10eps()

hdom = clustertree(dom)
hsp = Space(hdom)

G1 = GreensFun((x,y)->1/2,CauchyWeight(hsp⊗hsp,0);method=:Cholesky)
H = ⨍[G1]

@time u2 = H\f
@time u2 = H\f

@test condest(H) ≤ 20.0

import ApproxFun: ∞
@test blockrank(H) == Number[∞ 8 9 9; 8 ∞ 9 9; 9 9 ∞ 8; 9 9 8 ∞]

println("The hierarchical forward error norm is: ",norm(⨍[G]*u2-f))

@test norm(⨍[G]*u2-f) < 10eps()

# Test three domains

dom = Interval(-2.0,-1.0+0im)∪Interval(1.0,2.0+0im)∪Interval(-0.5-2im,0.5-2im)
hdom = clustertree(dom)
dom = UnionDomain(SingularIntegralEquations.collectdata(hdom))
⨍ = DefiniteLineIntegral(dom)
f = Fun(x->logabs(x-5im),dom)
sp = Space(dom)

G = GreensFun((x,y)->1/2,CauchyWeight(sp⊗sp,0);method=:Cholesky)

@time u1 = ⨍[G]\transpose(f)

println("Adaptive QR  forward error norm is: ",norm(⨍[G]*u1-f))

@test norm(⨍[G]*u1-f) < 300eps()

if VERSION ≥ v"0.5pre"
    warn("Skipping test in 0.5 due to error in type inference due to #265")
else
    hsp = Space(hdom)

    G1 = GreensFun((x,y)->1/2,CauchyWeight(hsp⊗hsp,0);method=:Cholesky)
    H = ⨍[G1]

    @time u2 = H\f
    @time u2 = H\f

    println("The hierarchical forward error norm is: ",norm(⨍[G]*u2-f))

    @test norm(⨍[G]*u2-f) < 10000eps()
end


export ExpSweep, expsweep, analyze

# ALMOST COPY OF: https://github.com/ssfrr/MeasureIR.jl/blob/master/src/expsweep.jl

struct ExpSweep{AT, SR}
    sig::AT
    w1::Float64
    w2::Float64
    gain::Float64
    samplerate::SR
end

# the default constructor doesn't auto-convert the non-parametric arguments
ExpSweep(sig::T, w1, w2, g, sr::SR) where {T, SR} = ExpSweep{T, SR}(sig, w1, w2, g, sr)

"""
    expsweep(length; options...)
Create an impulse response measurement using an exponential sinusoid sweep from
`minfreq` to `maxfreq`, following the procedure from [1]. Durations and
frequencies can be given with or without Unitful units. With the default
`samplerate` of `1` durations can be interpreted as samples and frequencies are
in cycles/sample. The default minimum frequency of 0.0004 cycles/sample
corresponds to about 19Hz at 48kHz.
## Options
- `samplerate=1`
- `minfreq=0.0004`
- `maxfreq=0.5`
- `fadein`
- `fadeout`
- `fade`
- `optimize`
- `func`
## Examples
```julia
using MeasureIR
# generate a 10000-sample sweep
expsweep(10000)
# generate a 10s sweep at 48kHz samplerate
using Unitful: s, kHz
expsweep(10s; samplerate=48kHz)
```
[1]: Farina, Angelo, Simultaneous measurement of impulse response and distortion
with a swept-sine technique, Audio Engineering Society Convention 108 (2000).
"""
function expsweep(length;
        # TODO: consider making 2π the default samplerate, which makes the
        # frequencies interpretable as radians/sample
        samplerate = 1,
        minfreq=0.0004*samplerate,
        # TODO: this probably isn't the right default. for example for
        # meas = expsweep(30s; samplerate=48000Hz), it doesn't roll off on top
        maxfreq=0.5*samplerate,
        # default fades are computed below
        fadein=nothing,
        fadeout=nothing,
        fade=nothing,
        gain=1,
        optimize=true,
        func=sin)

    # TODO: review Novak: https://ant-novak.com/posts/research/2015-10-30_JAES_Swept/
    # it was recommended by Gordon

    w1 = uconvert(NoUnits, minfreq * 2π / samplerate)
    w2 = uconvert(NoUnits, maxfreq * 2π / samplerate)
    L = round(Int, uconvert(NoUnits, length*samplerate))

    if fade !== nothing
        if fadein !== nothing || fadeout !== nothing
            throw(ArgumentError("`fade` argument cannot be used with `fadein` or `fadeout`"))
        end
        fadein = fade
        fadeout = fade
    end

    # default fade times chosen heuristically to be pretty reasonable
    # we could definitely be smarter here, especially for small L
    if fadein === nothing
        f1 = round(Int, min(L/2, 200/w1))
    else
        f1 = round(Int, uconvert(NoUnits, fadein*samplerate))
    end

    if fadeout === nothing
        f2 = round(Int, min(L/4, 800/w2))
    else
        f2 = round(Int, uconvert(NoUnits, fadeout*samplerate))
    end

    if optimize
        w1 = _optimizew1(w1, w2, L)
    end
    sig = _expsweep(func, L, w1, w2)
    winin = 0.5 .- 0.5*cos.(range(0, π; length=f1))
    winout = 0.5 .- 0.5*cos.(range(π, 0; length=f2))
    sig[1:f1] .*= winin
    sig[(end-f2+1):end] .*= winout
    ExpSweep(sig, w1, w2, gain, samplerate)
end

function _expsweep(func, L, minfreq, maxfreq)
    # switch to notation from [1]
    w1 = minfreq
    w2 = maxfreq
    T = L
    t = 0:(T-1)
    @. func(w1*T /
           log(w2/w1) * (exp(t/T * log(w2/w1))-1))
end

"""
Compute a new starting freq w < w1 that gives us an integer number
of total cycles.
"""
function _optimizew1(w1, w2, L)
    # scale to [0-1]
    f1 = w1 * L
    f2 = w2 * L
    # compute the total (integer) number of cycles we want
    M = floor((f2 - f1) /
              (2π*log(f2/f1)))

    a(f) = (f2-f)/log(f2/f) - 2π*M
    fnew = find_zero(a, f1)

    fnew/L
end

"""
    function analyze(m::ExpSweep, response::AbstractArray; noncausal=false)
When the measurement is a sweep, you can pass the `noncausal=true` keyword
argument to include the noncausal part of the impulse response, which carries
energy from nonlinear distortion in the system. In this case the zero-lag part
of the impulse response starts at index `L÷2+1`, where `L` is the length of the
IR.
"""
function analyze(m::ExpSweep, response::AbstractArray; noncausal=false)
    # we don't truncate the response here because typically sweep measurements
    # are taken without repeating, so there's no real errors we'd catch by
    # truncating.
    T = length(m.sig)
    w1 = m.w1
    w2 = m.w2
    t = 0:(T-1)
    # compute the normalizing factor from the derivative of the frequency -
    # the faster the frequency is changing the less time the signal spends
    # at that frequency so we need to boost the amplitude to compensate
    amp = @. w1/T*log(w2/w1)*exp(t/T*log(w2/w1))
    invfilt = amp .* m.sig
    # calculate an extra scaling factor to get the center of the passband to
    # unity. At some point it might make sense to look at a range instead of
    # just the center frequency, in case we catch it in the middle of some
    # ripple. Generally things should be pretty flat if there's a fade in/out
    # TODO: seems kinda silly to compute the whole FFT just to look at the
    # center frequency. We could just do the dot product of that freq
    roundtrip = xcorr(m.sig * m.gain, invfilt; padmode=:longest)
    centeridx = round(Int, (w1 + (w2-w1)/2)/π*(length(roundtrip)/2))
    scale = abs(rfft(roundtrip)[centeridx])

    ir = mapslices(response, dims=1) do v
        xcorr(v, invfilt, padmode=:longest) ./ scale
    end

    # keep the full IR (including non-causal parts representing nonlinearities)
    # if m.nonlinear is true
    zeroidx = size(ir,1)÷2+1
    if noncausal
        SampleArray(reshape(ir[zeroidx-T:zeroidx+T], :, 1), m.samplerate)
    else
        # nonlinearity detection is not working right now, but it seems like a
        # useful feature to add in the future. One issue is that we need a way
        # to estimate the expected "noise floor" of the noncausal part of the
        # IR, which is _not_ the same as the noise floor of the response. It
        # definitely helps to deconvolve just the silence part and use that to
        # estimate the IR noise floor. The other issue is that the overall
        # energy in the noncausal part just doesn't change that much, because
        # all the extra energy is pretty concentrated and it gets averaged out.
        # probably to get this to work well we either need to explicitly look
        # for peaks, and possibly look at fixed delays from the main peak that
        # correspond to harmonics.
        # deconvsilence = mapslices(resp, 1) do v
        #     xcorr(v[1:prepadding(m)], invfilt)
        # end
        # nf = noisefloor(m, deconvsilence)

        # we don't check all the way up to the center because if the impulse is
        # bandlimited it'll have energy directly preceeding it. In reality
        # how far back we should go depends on the bandwidth of the response,
        # so this number may need to be tuned.
        # fortunately this is only to heuristically decide whether we print a
        # warning, so we're never going to get to 100% accuracy anyways...
        # preimpulse = 100
        # noncausalpower = sum(ir[zeroidx-T:zeroidx-preimpulse].^2) /
        #                  (2T+1-preimpulse)
        # if any(noncausalpower .> nf .* 1.1 .+ sqrt(eps()))
        #     @warn "Energy in noncausal IR is above noise floor. Check for nonlinearity"
        # end
        SampleArray(reshape(ir[zeroidx:zeroidx+T-1], :, 1), m.samplerate)
    end
end

function analyze(m::ExpSweep, response::SampleArray; kwargs...)
    @assert nchannels(response) == 1
    analyze(m, data(response)[:]; kwargs...)
end


### A Pluto.jl notebook ###
# v0.12.20

using Markdown
using InteractiveUtils

# ╔═╡ 5bf16f26-7191-11eb-1166-73a483cf120a
begin
	using Pkg
	Pkg.activate("MLJ_env", shared=true)
end

# ╔═╡ 3cfba9ce-7191-11eb-2295-3b4c85656c03
begin
	using MLJ
	using CSV
	using DataFrames
	using PlutoUI
end

# ╔═╡ f3283b68-7191-11eb-3290-33ba4d2aaece
using Random

# ╔═╡ 9dc7a434-7195-11eb-35a1-e9f79dd4ee79
using Test

# ╔═╡ 428fadd2-719a-11eb-3f4b-d98cb77cddbe
using Printf

# ╔═╡ 87aaecb0-7190-11eb-1572-d9110d9ffaaa
md"""
# Week 2: Multiple Regression (Interpretation)
"""

# ╔═╡ c5c92502-7190-11eb-3a75-87b7f156a7ce
md"""
The goal of this first notebook is to explore multiple regression and feature engineering with MLJ functions.

In this notebook we will use data on house sales in King County to predict prices using multiple regression. we will:

  - Use DataFrames to do some feature engineering
  - Use built-in MLJ functions to compute the regression weights (coefficients/parameters)
  - Given the regression weights, predictors and outcome write a function to compute the Residual Sum of Squares (RSS)
  - Look at coefficients and interpret their meanings
  - Evaluate multiple models via RSS

"""

# ╔═╡ 7a50bd00-7191-11eb-2ed6-071e6db68422
md"""
### Load in house sales data
"""

# ╔═╡ 890cb362-7191-11eb-251b-e155ed2d21f2
sales = CSV.File("../../ML_UW_Spec/C02/data/kc_house_test_data.csv";
	header=true) |> DataFrame;

# ╔═╡ b1ae4466-7191-11eb-0c07-29e30cce116a
first(sales, 1)

# ╔═╡ b43046d2-7191-11eb-0978-b1bfe1762756
md"""
### Split data into train/test sets
"""

# ╔═╡ cd45cf32-7191-11eb-1e68-53bccaa9d806
function train_test_split(df; split=0.8, seed=42)
	Random.seed!(seed)
	(nr, nc) = size(df)
	nrp = round(Int, nr * split)
	row_ixes = shuffle(1:nr)
	df_train = view(df[row_ixes, :], 1:nrp, 1:nc)
	df_test = view(df[row_ixes, :], nrp+1:nr, 1:nc)
	(df_train, df_test)
end

# ╔═╡ e054e324-7191-11eb-3e15-c1e0bf73bb0d
sales_train, sales_test = train_test_split(sales);

# ╔═╡ fa6c9bf8-7191-11eb-26b3-1f7f13b91971
size(sales_train), size(sales_test)

# ╔═╡ 0d2ca940-7192-11eb-34ea-ab594b994f41
md"""
### Learning a multiple regression model
"""

# ╔═╡ 14e8024c-7192-11eb-1c76-83a28b6120ee
@load LinearRegressor pkg=MLJLinearModels

# ╔═╡ 72a34b30-7192-11eb-3cda-15a4e9a5244c
describe(sales, :mean, :std, :eltype)

# ╔═╡ c264c3ec-7192-11eb-24ef-c797c42d3671
begin
	example_features = [:sqft_living, :bedrooms, :bathrooms]

	X = select(sales_train, example_features)
	y = sales_train.price

	first(X, 3)
end

# ╔═╡ 1f6d0112-7193-11eb-20a2-efd1550ca8e2
## 3 first house prices
y[1:3]

# ╔═╡ bd39bf8e-7193-11eb-23c0-257408e760a0
size(X), size(y)

# ╔═╡ 45479604-7193-11eb-068e-7ddd263e9be2
## Let's declare a simple multivariate linear regression model:
mdl = MLJLinearModels.LinearRegressor()

# ╔═╡ 89ce6cb2-7193-11eb-25bb-f161f4514a6d
## Let's wrap it in a machine which, in MLJ, is the composition of a model
## and data to apply the model on:

begin
	mach_mr = machine(mdl, X, y)
	fit!(mach_mr)
end

# ╔═╡ d34ab594-7193-11eb-1bed-01a763445a7e
begin
	fp = fitted_params(mach_mr)

	with_terminal() do
		for (name, val) in fp.coefs
    		println("$(rpad(name, 8)):  $(round(val, sigdigits=3))")
		end
		println("Intercept: $(round(fp.intercept, sigdigits=3))")
	end
end

# ╔═╡ 508db812-7194-11eb-3fde-ad654ba107d1
md"""
### Making Predictions
"""

# ╔═╡ 6732d8e0-7194-11eb-3506-93329a94dd84
begin
	ŷ = predict(mach_mr, X)
	(round(ŷ[1], sigdigits=3) , y[1])  # first house predicetd price vs actual price
end

# ╔═╡ c9a0a836-7194-11eb-027a-7bfbfe6cae3c
md"""
### Compute RSS

Now that we can make predictions given the model, let's write a function to compute the RSS [Residual Sum of Squares] of the model.

Let's write a function `get_rss`  to calculate RSS given the machine, data(X), and the outcome (y).
"""

# ╔═╡ 3deaf822-7195-11eb-2bfb-9f51b24982ed
function get_rss(mach, X, y)
    ŷ = predict(mach, X)     # First get the predictions
    diff = y .- ŷ            # Then compute the residuals/errors
    rss = sum(diff .* diff)  # Then square and add them up
    return rss
end

# ╔═╡ a5251374-7195-11eb-368e-5d4814c8a5e2
begin
	X_test = select(sales_test, example_features)
	y_test = sales_test.price

	rss_ex = get_rss(mach_mr, X_test, y_test)
	exp_value = 5.313821155174205e13  # should be 5.313821155174205e13
	@test rss_ex ≈ exp_value
    # println("rss_ex is: $(rss_ex)")
end

# ╔═╡ 83c35d98-7196-11eb-0a75-b952509c4632
md"""
### Create some new features

Let's consider transformations of existing features e.g. the log of the squarefeet or even "interaction" features such as the product of bedrooms and bathrooms.

Create the following 4 new features as column in both test and train data:

  - `bedrooms_squared = bedrooms * bedrooms`
  - `bed_bath_rooms = bedrooms * bathrooms`
  - `log_sqft_living = log(sqft_living)`
  - `lat_plus_long = lat + long`

"""

# ╔═╡ d1cc2b4a-7197-11eb-10fd-8fdadba255d2
begin
	Xtr, Xte = (sales_train, sales_test)
end

# ╔═╡ d0d2312a-7196-11eb-3192-2158f9b7af15
begin
 	nXtr = hcat(Xtr,
		Xtr.bedrooms .* Xtr.bedrooms,
		Xtr.bedrooms .* Xtr.bathrooms,
 		log.(Xtr.sqft_living),
		Xtr.lat .+ Xtr.long;
		makeunique=true);

	rename!(nXtr, :x1 => :bedrooms_squared);
	rename!(nXtr, :x1_1 => :bed_bath_rooms);
	rename!(nXtr, :x1_2 => :log_sqft_living);
	rename!(nXtr, :x1_3 => :lat_plus_long);
end

# ╔═╡ bbe4cfb8-7199-11eb-2c60-edaa2feab075
begin
	nXte = hcat(Xte, Xte.bedrooms .* Xte.bedrooms,
		Xte.bedrooms .* Xte.bathrooms,
		log.(Xte.sqft_living),
		Xte.lat .+ Xte.long;
		makeunique=true);

	rename!(nXte, :x1 => :bedrooms_squared);
	rename!(nXte, :x1_1 => :bed_bath_rooms);
	rename!(nXte, :x1_2 => :log_sqft_living);
	rename!(nXte, :x1_3 => :lat_plus_long);
end

# ╔═╡ fd7ca6fa-7199-11eb-0691-89ea26b44013
md"""

**Quiz Question: What is the mean (arithmetic average) value of your 4 new features on
test data? (round to 2 digits)**

"""

# ╔═╡ 0ac9bda2-719a-11eb-3036-d110c21a2fcc
with_terminal() do
	for attr ∈ [:bed_bath_rooms, :log_sqft_living, :lat_plus_long, :bedrooms_squared]
		@printf("mean of %s is: %3.2f\n", String(attr), mean(nXte[!, attr]))
	end
end

# ╔═╡ afde3d9c-719a-11eb-12cb-8397b9441a65
md"""
### Learning Multiple Models

Now we will learn the weights for three (nested) models for predicting house prices. The first model will have the fewest features the second model will add one more feature and the third will add a few more:
  - Model 1: `squarefeet`, `# bedrooms`, `# bathrooms`, `latitude & longitude`
  - Model 2: add `bedrooms *bathrooms`
  - Model 3: Add `log squarefeet`, `bedrooms squared`, and the (nonsensical) `latitude + longitude`
"""

# ╔═╡ fa0073b6-719a-11eb-3190-dd61c2541ae4
begin
	mod1_features = [:sqft_living, :bedrooms, :bathrooms, :lat, :long];
	mod2_features = [mod1_features..., :bed_bath_rooms];
	mod3_features = [mod2_features..., :bedrooms_squared, :log_sqft_living, :lat_plus_long];
end

# ╔═╡ 5caf09e6-719b-11eb-1652-8938ad992ea3
## model 1
begin
	X1 = select(nXtr, mod1_features)
	mach_mr1 = machine(mdl, X1, y)
	fit!(mach_mr1)

	fp1 = fitted_params(mach_mr1)

	with_terminal() do
		for (name, val) in fp1.coefs
    		println("$(rpad(name, 16)):  $(round(val, sigdigits=3))")
		end
		println("Intercept: $(round(fp1.intercept, sigdigits=3))")
	end
end

# ╔═╡ c56e31be-719b-11eb-022d-f796b95b6b32
## model 2
begin
	X2 = select(nXtr, mod2_features)
	mach_mr2 = machine(mdl, X2, y)
	fit!(mach_mr2)

	fp2 = fitted_params(mach_mr2)

	with_terminal() do
		for (name, val) in fp2.coefs
    		println("$(rpad(name, 16)):  $(round(val, sigdigits=3))")
		end
		println("Intercept: $(round(fp2.intercept, sigdigits=3))")
	end
end

# ╔═╡ d972ae42-719b-11eb-193e-675d543025c1
## model 3
begin
	X3 = select(nXtr, mod3_features)
	mach_mr3 = machine(mdl, X3, y)
	fit!(mach_mr3)

	fp3 = fitted_params(mach_mr3)

	with_terminal() do
		for (name, val) in fp3.coefs
    		println("$(rpad(name, 16)):  $(round(val, sigdigits=3))")
		end
		println("Intercept: $(round(fp3.intercept, sigdigits=3))")
	end
end

# ╔═╡ 266c0a0e-719c-11eb-1692-f74b2a6f5bc4
md"""

**Quiz Question: What is the sign (positive or negative) for the coefficient/weight for 'bathrooms' in model 1?**
  - [x] Positive (+)
  - [ ] Negative (-)

**Quiz Question: What is the sign (positive or negative) for the coefficient/weight for 'bathrooms' in model 2?**
  - [ ] Positive (+)
  - [x] Negative (-)

Think about what this means.

"""

# ╔═╡ 2dfa4aee-719c-11eb-1f89-1babfa6312d2
md"""
### Comparing multiple models

Now that we have learned three models and extracted the model weights we want to evaluate which model is best.

"""

# ╔═╡ 48377010-719c-11eb-19a7-233bf9bb4beb
# function calc_rss(;models=[(:model_1, mach_mr1),
# 			(:model_2, mach_mr2),
# 			(:model_3, mach_mr3)], tdata=sales_train)
# 	for (modn, mach) ∈ models
#     	rss_mod = get_rss(mach, tdata, tdata.price)
#     	@printf("model %10s rss is: %2.7e", String(modn), rss_mod)
# 	end
#   	return
# end

# ╔═╡ 9af9983a-719e-11eb-0a18-f5493f3df6a3
md"""
**Training Data**
"""

# ╔═╡ 06cdc690-719e-11eb-0819-1b5fce1e3a18
begin
	rss_mmr1 = get_rss(mach_mr1, X1, y)
	rss_mmr2 = get_rss(mach_mr2, X2, y)
	rss_mmr3 = get_rss(mach_mr3, X3, y)

	with_terminal() do
		for (ix, rss_) ∈ enumerate((rss_mmr1, rss_mmr2, rss_mmr3))
			@printf("model %2d rss is: %2.7e\n", ix, rss_)
		end
	end
end

# ╔═╡ e60ea370-719f-11eb-3d56-5dcfecf4ab3e
md"""
**Quiz Question: Which model (1, 2 or 3) has lowest RSS on *training* Data?** Is this what you expected?

  - [ ] Model 1
  - [ ] Model 2
  - [x] Model 3

  I would expect this (on training data)
"""

# ╔═╡ 90a5c372-719e-11eb-3483-6140d252b0c0
md"""
**Test Data**
"""

# ╔═╡ b2d5046c-719e-11eb-07db-454247474942
begin
	rss_mmr1_te = get_rss(mach_mr1, select(nXte, mod1_features), sales_test.price)
	rss_mmr2_te = get_rss(mach_mr2, select(nXte, mod2_features), sales_test.price)
	rss_mmr3_te = get_rss(mach_mr3, select(nXte, mod3_features), sales_test.price)

	with_terminal() do
		for (ix, rss_) ∈ enumerate((rss_mmr1_te, rss_mmr2_te, rss_mmr3_te))
			@printf("model %2d rss is: %2.7e\n", ix, rss_)
		end
	end
end

# ╔═╡ fb4038a6-719f-11eb-3e40-cd5aecfb149b
md"""

**Quiz Question: Which model (1, 2 or 3) has lowest RSS on *testing* Data?** Is this what you expected? Think about the features that were added to each model from the previous.

  - [x] Model 1
  - [ ] Model 2
  - [ ] Model 3

"""

# ╔═╡ 0483452a-71a0-11eb-30c0-7d56ed83e837

# ╔═╡ Cell order:
# ╟─87aaecb0-7190-11eb-1572-d9110d9ffaaa
# ╟─c5c92502-7190-11eb-3a75-87b7f156a7ce
# ╠═5bf16f26-7191-11eb-1166-73a483cf120a
# ╠═3cfba9ce-7191-11eb-2295-3b4c85656c03
# ╟─7a50bd00-7191-11eb-2ed6-071e6db68422
# ╠═890cb362-7191-11eb-251b-e155ed2d21f2
# ╠═b1ae4466-7191-11eb-0c07-29e30cce116a
# ╟─b43046d2-7191-11eb-0978-b1bfe1762756
# ╠═f3283b68-7191-11eb-3290-33ba4d2aaece
# ╠═cd45cf32-7191-11eb-1e68-53bccaa9d806
# ╠═e054e324-7191-11eb-3e15-c1e0bf73bb0d
# ╠═fa6c9bf8-7191-11eb-26b3-1f7f13b91971
# ╟─0d2ca940-7192-11eb-34ea-ab594b994f41
# ╠═14e8024c-7192-11eb-1c76-83a28b6120ee
# ╠═72a34b30-7192-11eb-3cda-15a4e9a5244c
# ╠═c264c3ec-7192-11eb-24ef-c797c42d3671
# ╠═1f6d0112-7193-11eb-20a2-efd1550ca8e2
# ╠═bd39bf8e-7193-11eb-23c0-257408e760a0
# ╠═45479604-7193-11eb-068e-7ddd263e9be2
# ╠═89ce6cb2-7193-11eb-25bb-f161f4514a6d
# ╠═d34ab594-7193-11eb-1bed-01a763445a7e
# ╟─508db812-7194-11eb-3fde-ad654ba107d1
# ╠═6732d8e0-7194-11eb-3506-93329a94dd84
# ╟─c9a0a836-7194-11eb-027a-7bfbfe6cae3c
# ╠═3deaf822-7195-11eb-2bfb-9f51b24982ed
# ╠═9dc7a434-7195-11eb-35a1-e9f79dd4ee79
# ╠═a5251374-7195-11eb-368e-5d4814c8a5e2
# ╟─83c35d98-7196-11eb-0a75-b952509c4632
# ╠═d1cc2b4a-7197-11eb-10fd-8fdadba255d2
# ╠═d0d2312a-7196-11eb-3192-2158f9b7af15
# ╠═bbe4cfb8-7199-11eb-2c60-edaa2feab075
# ╟─fd7ca6fa-7199-11eb-0691-89ea26b44013
# ╠═428fadd2-719a-11eb-3f4b-d98cb77cddbe
# ╠═0ac9bda2-719a-11eb-3036-d110c21a2fcc
# ╟─afde3d9c-719a-11eb-12cb-8397b9441a65
# ╠═fa0073b6-719a-11eb-3190-dd61c2541ae4
# ╠═5caf09e6-719b-11eb-1652-8938ad992ea3
# ╠═c56e31be-719b-11eb-022d-f796b95b6b32
# ╠═d972ae42-719b-11eb-193e-675d543025c1
# ╟─266c0a0e-719c-11eb-1692-f74b2a6f5bc4
# ╟─2dfa4aee-719c-11eb-1f89-1babfa6312d2
# ╠═48377010-719c-11eb-19a7-233bf9bb4beb
# ╟─9af9983a-719e-11eb-0a18-f5493f3df6a3
# ╠═06cdc690-719e-11eb-0819-1b5fce1e3a18
# ╟─e60ea370-719f-11eb-3d56-5dcfecf4ab3e
# ╟─90a5c372-719e-11eb-3483-6140d252b0c0
# ╠═b2d5046c-719e-11eb-07db-454247474942
# ╟─fb4038a6-719f-11eb-3e40-cd5aecfb149b
# ╠═0483452a-71a0-11eb-30c0-7d56ed83e837


include("../src/searching.jl")
using .Searching

function testbinarysearch(testcase, searchfunc)
    v, k, ans = testcase
    @test searchfunc(k, v) == ans
end

@testset "Binary search" begin
    testcases = [
        ([1, 2, 5], 5, 3),
        ([1, 2, 5], 4, -1),
        ([1], 1, 1),
        ([1], 5, -1),
        ([-1, 0, 3, 5, 10, 200], -5, -1),
        ([-1, 0, 3, 5, 10, 200], 10, 5),
    ]

    for testcase in testcases
        testbinarysearch(testcase, binarysearch)
        testbinarysearch(testcase,
            (k, v, algo=:iterative) -> binarysearch(k, v, algo=algo))
    end
end

function testactions(testcase, bst)
    actions, ans = testcase
    node = bst.root

    for action in actions
        if action == :left
            node = node.left
        elseif action == :right
            node = node.right
        end
    end

    if isnothing(ans)
        @test isnothing(node)
    else
        @test node.val == ans
    end
end

@testset "Binary search tree" begin
    v = [5, 3, 6, 7, 4, 10]

    testcases = [
        (actions=[],  ans=5),
        (actions=[:left, :left],  ans=nothing),
        (actions=[:left, :right], ans=4),
        (actions=[:right, :right, :right], ans=10)
    ]

    @testset "Basics" begin
        bst = BinarySearchTree(v)
        @test Searching.size(bst) == 6

        @test Searching.get(bst, 5).val == 5
        @test isnothing(Searching.get(bst, 100))

        @test Searching.maximum(bst).val == 10
        @test Searching.minimum(bst).val == 3

        for testcase in testcases
            testactions(testcase, bst)
        end
    end

    @testset "Test put!" begin
        bst = BinarySearchTree(v)
        result = Searching.put!(bst, 2)
        @test Searching.size(result) == 7
        @test result.root.left.left.val == 2
    end

    @testset "Test delete!" begin
        bst = BinarySearchTree(v)
        result = Searching.delete!(bst, 5)
        @test result.root.val == 6
        @test result.root.numnodes == 5

        bst = BinarySearchTree(v)
        result = Searching.delete!(bst, 6)
        @test result.root.val == 5
        @test result.root.numnodes == 5
        @test result.root.right.val == 7

        bst = BinarySearchTree(v)
        bst = Searching.delete!(bst, 5)
        bst = Searching.delete!(bst, 6)
        @test bst.root.val == 7
        @test Searching.size(bst.root) == 4
        @test bst.root.right.val == 10
    end

    @testset "Test delete! side effects" begin
        # suppose to change the node the same way as returned value
        node = BinarySearchTree(v).root
        Searching.delete!(node, 5)
        result = BinarySearchTree(node)
        @test Searching.size(result) == 5
        @test result.root.val == 6

        node = BinarySearchTree(v).root
        Searching.delete!(node, 6)
        result = BinarySearchTree(node)
        @test Searching.size(result) == 5
        @test result.root.val == 5
        @test result.root.right.val == 7
    end

    @testset "Test rank and select" begin
        bst = BinarySearchTree(v)
        @test Searching.rank(bst, 5) == 3
        @test Searching.rank(bst, 1) == 1
        @test Searching.rank(bst, 100) == 7

        bst = BinarySearchTree(v)
        @test Searching.select(bst, 5).val == 7
        @test Searching.select(bst, 1).val == 3
        @test isnothing(Searching.select(bst, 0))
        @test isnothing(Searching.select(bst, 100))
    end
end


function count_fixed_points(array)
    result = 0
    for (index, value) in enumerate(array)
        if index == value
            result += 1
        end
    end
    return result
end

function count_cycles(array)
    result = 0
    visited = trues(size(array, 1))
    for (index, value) in enumerate(array)
        if visited[index]
            j = index
            while true
                visited[j] = false
                j = array[j]
                if j == index
                    break
                end
            end
            result += 1
        end
    end
    return result
end


import LogDensityProblems
import ProgressMeter
import TransformVariables

const default_learning_rate = 1e-4
const default_max_iterations = 1_000_000
const default_show_progress_meter = true
const default_throw_convergence_exception = true
const default_tolerance = 1e-10

"""
    fit(transformed_gradient_problem, theta_hat_initial; kwargs...)

Find the maximum likelihood estimatator for the parameters `theta`.

# Arguments
- `transformed_gradient_problem`
- `theta_hat_initial`

# Keyword Arguments
- `learning_rate`. Default value: $(default_learning_rate)
- `max_iterations`. Default value: $(default_max_iterations)
- `show_progress_meter`. Default value: $(default_show_progress_meter)
- `throw_convergence_exception`. Default value: $(default_throw_convergence_exception)
- `tolerance`. Default value: $(default_tolerance)

See the documentation for fully worked-out examples.
"""
function fit(transformed_gradient_problem,
             theta_hat_initial;
             learning_rate = default_learning_rate,
             max_iterations = default_max_iterations,
             show_progress_meter = default_show_progress_meter,
             throw_convergence_exception = default_throw_convergence_exception,
             tolerance = default_tolerance)
    transformed_log_density = parent(transformed_gradient_problem)
    transformation = transformed_log_density.transformation
    theta_hat_initial_deepcopy = deepcopy(theta_hat_initial)
    theta_hat_initial_deepcopy_inversetransformed = TransformVariables.inverse(transformation,
                                                                               theta_hat_initial_deepcopy)
    theta_hat_initial_deepcopy_inversetransformed_deepcopy = deepcopy(theta_hat_initial_deepcopy_inversetransformed)
    theta_hat_new_inversetransformed = _fit(transformed_gradient_problem,
                                            theta_hat_initial_deepcopy_inversetransformed_deepcopy;
                                            learning_rate = learning_rate,
                                            max_iterations = max_iterations,
                                            show_progress_meter = show_progress_meter,
                                            throw_convergence_exception = throw_convergence_exception,
                                            tolerance = tolerance)
    theta_hat_new = TransformVariables.transform(transformation,
                                                 theta_hat_new_inversetransformed)
    return theta_hat_new
end

function _fit(transformed_gradient_problem,
              theta_hat_initial;
              learning_rate,
              max_iterations,
              show_progress_meter,
              throw_convergence_exception,
              tolerance)
    theta_hat_new = deepcopy(theta_hat_initial)
    if show_progress_meter
        progress = ProgressMeter.ProgressThresh(tolerance, "Maximizing:")
    end
    for iter = 1:max_iterations
        theta_hat_old = theta_hat_new
        log_likelihood_value, gradient = LogDensityProblems.logdensity_and_gradient(transformed_gradient_problem,
                                                                                theta_hat_old)
        update = learning_rate * gradient
        theta_hat_new = theta_hat_old + update
        update_norm = sum(abs, update)
        if show_progress_meter
            showvalues = [(:iteration, iter),
                          (:loglikelihood, log_likelihood_value)]
            ProgressMeter.update!(progress,
                                  update_norm;
                                  showvalues = showvalues)
        end
        if update_norm < tolerance
            @info("Successfully converged after $(iter) iterations.")
            return theta_hat_new
        end
    end
    if show_progress_meter
        ProgressMeter.finish!(progress)
    end
    if throw_convergence_exception
        @error("Failed to converge after $(max_iterations) iterations.")
        throw(ConvergenceException())
    else
        @warn("Failed to converge after $(max_iterations) iterations.")
    end
    return theta_hat_new
end


using Test, RubiSymbolics
using RubiSymbolics: @variables, simplify, isequal, expand, Elliptic, HypergeometricFunctions, @rule, Fixpoint, Postwalk, Chain

const print_incorrect_tests = false

r1 = @rule (~base::(x -> x isa Number && x<0))^(~exp::(x -> !(x isa Integer))) --> Complex(~base)^~exp
r1 = @rule (-(~base::(x -> x isa Number && x>0)))^(~exp::(x -> !(x isa Integer))) --> Complex(-(~base))^~exp
exp_fixer = Fixpoint(Postwalk(Chain([r1])))
function preprocess(expr::Expr)
    return exp_fixer(expr)
end

macro test_int(expr)
    if (expr.head == :vect && length(expr.args) == 4)
        preprocessed = preprocess(expr)
        (integrand, symbol, steps, result) = preprocessed.args

        # return :(@test isequal(simplify(integrate($(esc(integrand)), $(esc(symbol)))), simplify($(esc(result)))))
        return :(@test isequal(integrate($(esc(integrand)), $(esc(symbol))), $(esc(result))))
        # quote
        #     isequal(simplify(integrate($(esc(integrand)), $(esc(symbol)))), simplify($(esc(result))))
        #     @test true
        # end
    else
        if print_incorrect_tests
            println("Incorrect test problem at ", __source__)
        end
        return :(@test_broken false)
    end
end

# include(".jl")
include("/home/main/.julia/dev/RubiSymbolics/test/1 Algebraic functions/1.1 Binomial products/1.1.1 Linear/.jl")


# JuMP unit tests for MOI Wrapper
using COSMO, Test, LinearAlgebra, JuMP, MathOptInterface

# ---------------------------
# Linear Program
# ---------------------------

c = [300.0; 500.0]

m = Model(with_optimizer(COSMO.Optimizer))
@variable(m, x1 >= 0)
@variable(m, x2 >= 0)

@objective(m, Max, c[1] * x1 + c[2] * x2)
@constraint(m, x1 + 2 * x2 <= 170)
@constraint(m, x1 + x2 <= 150)
@constraint(m, 3 * x2 <= 180)
status = JuMP.optimize!(m)

termination_status = JuMP.termination_status(m)
result_status = JuMP.primal_status(m)
optimal_x1 = JuMP.value(x1)
optimal_x2 = JuMP.value(x2)
objvalue = JuMP.objective_value(m)

@testset "LP - JuMP" begin
    @test isapprox(JuMP.objective_value(m), 49000, atol = 1)
    @test isapprox(JuMP.value(x1), 130, atol = 0.1)
    @test isapprox(JuMP.value(x2), 20, atol = 0.1)
    @test JuMP.termination_status(m) == MathOptInterface.Success
    @test JuMP.primal_status(m) == MathOptInterface.FeasiblePoint
end

# ---------------------------
# Quadratic Program
# ---------------------------

m = Model(with_optimizer(COSMO.Optimizer))
@variable(m, x[1:2])

P = [4. 1; 1 2]
q = [1; 1.]
@objective(m, Min, 0.5 * x' * P * x  + q' * x)
A = [1. 1; 1 0; 0 1];
l = [1.; 0; 0];
u = [1.; 0.7; 0.7]
@constraint(m, A * x .<= u)
@constraint(m, A * x .>= l)
status = JuMP.optimize!(m)

@testset "QP - JuMP" begin
    @test isapprox(JuMP.objective_value(m), 1.88, atol = 1e-3)
    @test isapprox(JuMP.value.(x)[1], 0.3, atol = 1e-3)
    @test isapprox(JuMP.value.(x)[2], 0.7, atol = 1e-3)
    @test JuMP.termination_status(m) == MathOptInterface.Success
    @test JuMP.primal_status(m) == MathOptInterface.FeasiblePoint
end

# ---------------------------
# Second order cone Program
# ---------------------------

m = Model(with_optimizer(COSMO.Optimizer));
@variable(m, x[1:3]);
@objective(m, Min, x[1]);
@constraint(m, x[2] >= 1);
@constraint(m, x[3] >= 1);
@constraint(m,  soc, [x[1];x[2:3]] in SecondOrderCone());
status = JuMP.optimize!(m);

@testset "SOCP - JuMP" begin
    @test isapprox(JuMP.objective_value(m), sqrt(2), atol = 1e-3)
    @test isapprox(JuMP.value.(x)[1], sqrt(2), atol = 1e-3)
    @test isapprox(JuMP.value.(x)[2], 1., atol = 1e-3)
    @test isapprox(JuMP.value.(x)[3], 1., atol = 1e-3)
    @test JuMP.termination_status(m) == MathOptInterface.Success
    @test JuMP.primal_status(m) == MathOptInterface.FeasiblePoint
end

# ---------------------------
# Semidefinite Program - Square
# ---------------------------
A1 = [1.0 0 1; 0 3 7; 1 7 5]
A2 = [0.0 2 8; 2 6 0; 8 0 4]
C = [1.0 2 3; 2 9 0; 3 0 7]
b1 = 11.0
b2 = 19.0

m = Model(with_optimizer(COSMO.Optimizer));
@variable(m, X[1:3,1:3]);
@objective(m, Min, tr(C * X));
@constraint(m, tr(A1 * X) == b1);
@constraint(m, tr(A2 * X) == b2);
@constraint(m, psdcon, X in JuMP.PSDCone()); # this is used to enforce a PsdConeSquare constraint
status = JuMP.optimize!(m);

@testset "SDP Square - JuMP" begin
    @test isapprox(JuMP.objective_value(m), 13.9022, atol = 1e-3)
    @test JuMP.termination_status(m) == MathOptInterface.Success
    @test JuMP.primal_status(m) == MathOptInterface.FeasiblePoint
end

# ---------------------------
# Semidefinite Program - Triangle
# ---------------------------
A1 = [1.0 0 1; 0 3 7; 1 7 5]
A2 = [0.0 2 8; 2 6 0; 8 0 4]
C = [1.0 2 3; 2 9 0; 3 0 7]
b1 = 11.0
b2 = 19.0

    m = Model(with_optimizer(COSMO.Optimizer));
    @variable(m, X[1:3,1:3], PSD);
    @objective(m, Min, tr(C * X));
    @constraint(m, tr(A1 * X) == b1);
    @constraint(m, tr(A2 * X) == b2);
    status = JuMP.optimize!(m);

@testset "SDP Triangle - JuMP" begin
    @test isapprox(JuMP.objective_value(m), 13.9022, atol = 1e-3)
    @test JuMP.termination_status(m) == MathOptInterface.Success
    @test JuMP.primal_status(m) == MathOptInterface.FeasiblePoint
end

nothing


# findnewton.jl
#
#   Faster search for sorted array
#
# Jeremy Rüffer
# Thünen Institut
# Institut für Agrarklimaschutz
# Junior Research Group NITROSPHERE
# Julia 0.7
# 09.12.2016
# Last Edit: 18.04.2019

" findnewton{T<:Any}(x,target::Array{T,1};max_iterations::Int,verbose::Bool)

In a sorted data set, search iteratively for the target value/s. Uses Newtonian method, look in the center for a match. If value is greater look halfway between end and current value. If less than look between current value and previous low value. Range will be halved each time until the same index or the sentinal value (iteration max) have been reached.

---

#### Keywords:\n
* max_iterations::Int = Maximum number of iterations to attempt, 1000 is default
* verbose::Bool = Display imformation as the function runs, FALSE is default\n\n

---

\n
#### Related Functions:\n
* searchsorted
* searchsortedfirst
* searchsortedlast"
function findnewton(x,target::Array{T,1};max_iterations::Int=1000,verbose::Bool=false) where T <: Any
	answer = Array{Int}(undef,length(target))
	answer[1] = findnewton(x,target[1];max_iterations=max_iterations,verbose=verbose)
	lastgood = Int(1)
	temp = []
	bad_vals = false
	for i=2:1:length(target)
		temp = findnewton(x,target[i];max_iterations=max_iterations,verbose=verbose)
		if temp <= answer[lastgood] + 1
			# If the answer is the same or just one more than the previous value
			# then a gap has been found, mark it for later deletion
			answer[i] = -1
			bad_vals = true
		else
			answer[i] = temp
			lastgood = i
		end
	end

	# Remove Bad Values
	if bad_vals
		answer = answer[findall(answer .!= -1)]
	end

	return answer
end

function findnewton(x,target;max_iterations::Int=1000,verbose::Bool=false)
	if verbose
		println("Min = " * string(findmin(x)[1]))
		println("Max = " * string(findmax(x)[1]))
		println("Target = " * string(target) * "\n")
	end

	############################
	##  Initialize Variables  ##
	############################
	lower = 1
	upper = length(x)
	test = false
	sentinal = 1
	answer = []

	########################
	##  Newtonian Search  ##
	########################
	while !test && sentinal < max_iterations
		mid = Int(lower + floor((upper-lower)/2))
		verbose ? println("  Iteration (" * string(sentinal) * ") = " * string(x[mid])) : []
		if x[mid] > target
			# Target is higher than guess
			if upper == mid
				test = true
				answer = mid
			else
				upper = mid
			end
		elseif x[mid] < target
			# Target is lower than guess
			if lower == mid
				test = true
				answer = mid
			else
				lower = mid
			end
		else
			# Exact Match!
			test = true
			answer = mid
		end

		sentinal += 1
	end

	# Sentinal Catch Case
	if sentinal >= max_iterations
		warn("Sentinal value reached, returned value may not be the best answer.")
		answer = mid
	end

	if verbose
		println("\nIndex = " * string(Int(answer)))
		println("Answer = " * string(x[answer]) * "\n")
	end
	return Int(answer)
end


"""
```
    MidwayEqualization <: AbstractHistogramAdjustmentAlgorithm
    MidwayEqualization(; nbins = 256, minval = 0, maxval = 1)

    adjust_histogram([T,] img_sequence, f::MidwayEqualization(nbins = 256, edges = nothing))
    adjust_histogram!([out_sequence,] img_sequence, f::MidwayEqualization(nbins = 256, edges = nothing))
```

Gives a pair of images the same histogram whilst maintaining as
much as possible their previous grey level dynamics.

# Details

The purpose of midway histogram equalization is to transform the intensities in
a pair of images so that the intensities distribute according to a common
"midway" distribution. The histogram representing the common distribution is
chosen so that the original  gray level dynamics of the images are preserved as
much as possible. If one interprets histograms as piecewise-constant models of
probability density functions (see [`build_histogram`](@ref
build_histogram(::AbstractArray, ::Integer, ::Union{Real,AbstractGray},
::Union{Real,AbstractGray}))), then the midway histogram equalization task can
be modeled as the problem of transforming one probability distribution into
another (see [`adjust_histogram`](@ref adjust_histogram(::Matching,::AbstractArray, ::AbstractArray, ::Integer))).
It turns out that the solution to this transformation problem involves the
cumulative and inverse cumulative distribution functions of the source and
"midway" probability density functions. In particular, let the random variables ``X_i \\thicksim p_{x_i} \\; (i = 1,2)``,
and ``Z \\thicksim p_{z}``  represent an intensity in the first, second and
"midway" image respectively, and let

```math
 S_{X_i}(x) = \\int_0^{x}p_{x_i}(w)\\mathrm{d} w \\; \\quad \\text{and} \\quad
 T_{Z}(x) = \\frac{2}{\\frac{1}{S_{X_1}(x)} + \\frac{1}{S_{X_2}(x)}}
```
represent the cumulative distribution functions of the two input images, and
their *harmonic mean*, respectively. Then the sought-after mapping
``Q_{X_i}(\\cdot)`` ``(i = 1,2)`` such that ``Q_{X_i}(x) \\thicksim p_{z} `` is
given by

```math
Q_{X_i}(x) =  T_{Z}^{-1}\\left( S_{X_i}(x) \\right),
```

where ``T_{Z}^{-1}(y) = \\operatorname{min} \\{ x \\in \\mathbb{R} : y \\leq
T_{Z}(x) \\}`` is the inverse cumulative distribution function of ``T_{Z}(x)``.

# Options

Various options for the parameters of the `adjust_histogram` function and
`MidwayEqualization` types are described in more detail below.

## Choices for `img_sequence`

The function `adjust_histogram` expects a length-2 `Vector` of images (the pair
of images) and returns a length-2 `Vector` of modified images.  The  function
can handle a variety of input types. The type of the returned image matches the
input type.

For colored images, the inputs are converted to
[YIQ](https://en.wikipedia.org/wiki/YIQ)  type and the distributions of the Y
channels are transformed according to a "midway" distribution. The modified Y
channel is then combined with the I and Q channels and the resulting image
converted to the same type as the input.

## Choices for `nbins`

You can specify the total number of bins in the histogram. If you do not
specify the number of bins then a default value of 256 bins is utilized.

## Choices for `edges`

If you do not designate the number of bins, then you have the option to directly
stipulate how the intervals will be divided by specifying a `AbstractRange`
type.

# Example

```julia
using Images, TestImages, ImageView, ImageContrastAdjustment

img = testimage("mandril_gray")

# The same image but with different intensitiy distributions
img1 = adjust_histogram(img, GammaCorrection(gamma = 2))
img2 = adjust_histogram(img, GammaCorrection(gamma = 1.2))

# Midway histogram equalization will transform these two images so that their
# intensity distributions are almost identical.
img_sequence = adjust_histogram([img1, img2], MidwayEqualization(nbins = 256))
img1o = first(img_sequence)
img2o = last(img_sequence)
```

# References
1. T. Guillemot and J. Delon, “*Implementation of the Midway Image Equalization*,” Image Processing On Line, vol. 5, pp. 114–129, Jun. 2016. [doi:10.5201/ipol.2016.140](https://doi.org/10.5201/ipol.2016.140)
"""
Base.@kwdef struct MidwayEqualization{T₁ <: Union{Integer, Nothing},
                					  T₂ <: Union{AbstractRange, Nothing}} <: AbstractHistogramAdjustmentAlgorithm
    nbins::T₁ = 256
    edges::T₂ = nothing
end

function (f::MidwayEqualization)(out_sequence::Vector{<:GenericGrayImage}, in_sequence::Vector{<:GenericGrayImage})
    length(out_sequence) == 2 || error("Please supply a length-2 output vector to store the pair of images.")
    length(in_sequence) == 2 || error("Please supply a length-2 input vector storing the image pair.")
    out1 = first(out_sequence)
    out2 = last(out_sequence)
    in1 = first(in_sequence)
    in2 = last(in_sequence)
    out1 .= in1
    out2 .= in2
    edges, pdf1, pdf2 = isnothing(f.edges) ? construct_pdfs(out1, out2, f.nbins) : construct_pdfs(out1, out2, f.edges)
    midway_pdf = zero(pdf1)
    cdf1 = cumsum(pdf1)
    cdf2 = cumsum(pdf2)
    # midway_cdf is the harmonic mean between cdf1 and cdf2.
    midway_cdf =  similar(cdf1)
    for i in eachindex(cdf1)
        if cdf1[i] == 0 || cdf2[i] == 0
            midway_cdf[i] = 0
        else
            midway_cdf[i] = 2 / (1/cdf1[i] + 1/cdf2[i])
        end
    end
    cdf2pdf!(midway_pdf, midway_cdf)
    index₁ = firstindex(out_sequence)
    index₂ = lastindex(out_sequence)
    out_sequence[index₁]  = match_pdf!(out1, edges, pdf1, midway_pdf)
    out_sequence[index₂]  = match_pdf!(out2, edges, pdf2, midway_pdf)
end

function (f::MidwayEqualization)(out_sequence::Vector{<:AbstractArray{<:Color3}}, in_sequence::Vector{<:AbstractArray{<:Color3}})
    length(out_sequence) == 2 || error("Please supply a length-2 output vector to store the pair of images.")
    length(in_sequence) == 2 || error("Please supply a length-2 input vector storing the image pair.")
    out1 = first(out_sequence)
    out2 = last(out_sequence)
    in1 = first(in_sequence)
    in2 = last(in_sequence)

    T₁ = eltype(out1)
    T₂ = eltype(out2)

    out1 .= in1
    out2 .= in2

    in_yiq1 = convert.(YIQ, in1)
    in_yiq1_view = channelview(in_yiq1)
    in_yiq2 = convert.(YIQ, in2)
    in_yiq2_view = channelview(in_yiq2)
    in_yiq_view_sequence = [view(in_yiq1_view,1,:,:), view(in_yiq2_view,1,:,:)]

    out_yiq1 = convert.(YIQ, out1)
    out_yiq1_view = channelview(out_yiq1)
    out_yiq2 = convert.(YIQ, out2)
    out_yiq2_view = channelview(out_yiq2)
    out_yiq_view_sequence = [view(out_yiq1_view,1,:,:), view(out_yiq2_view,1,:,:)]

    adjust_histogram!(out_yiq_view_sequence, in_yiq_view_sequence, f)

    index₁ = firstindex(out_sequence)
    index₂ = lastindex(out_sequence)
    out_sequence[index₁] .= convert.(T₁, out_yiq1)
    out_sequence[index₂] .= convert.(T₂, out_yiq2)
end

(f::MidwayEqualization)(out_sequence::Vector{<:GenericGrayImage}, img_sequence::Vector{<:AbstractArray{<:Color3}}) =
    f(out_sequence, map(img -> of_eltype(Gray, img), img_sequence))


# ---
# title: 137. Single Number II
# id: problem137
# author: Indigo
# date: 2021-06-03
# difficulty: Medium
# categories: Bit Manipulation
# link: <https://leetcode.com/problems/single-number-ii/description/>
# hidden: true
# ---
#
# Given an integer array `nums` where every element appears **three times**
# except for one, which appears **exactly once**. _Find the single element and
# return it_.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [2,2,3,2]
#     Output: 3
#
#
# **Example 2:**
#
#
#
#     Input: nums = [0,1,0,1,0,1,99]
#     Output: 99
#
#
#
#
# **Constraints:**
#
#   * `1 <= nums.length <= 3 * 104`
#   * `-231 <= nums[i] <= 231 - 1`
#   * Each element in `nums` appears exactly **three times** except for one element which appears **once**.
#
#
#
# **Follow up:**  Your algorithm should have a linear runtime complexity. Could
# you implement it without using extra memory?
#
#
## @lc code=start
using LeetCode

function single_number(nums::Vector{Int})::Int
    a, b = 0, 0
    for num in nums
        b = (b ⊻ num) & ~a
        a = (a ⊻ num) & ~b
    end
    return b
end
## @lc code=end


@static if VERSION >= v"1"
    using Distributed
end

function work(numsteps, np)
  sum = 0.0
  slice = 1.0/numsteps

  sum = @distributed (+) for i = 1:numsteps
    x = (i - 0.5) * slice
    (4.0/(1.0 + x^2))
  end

  mypi = sum * slice

  return mypi
end

function picalc(numsteps)

  # Get number of processes.
  np = nworkers()

  println("Calculating PI using:")
  println("  ", numsteps, " slices")
  println("  ", np, " worker(s)")
  temp = work(1,1)
  temp = work(numsteps,np)
  start = time()

  mypi = work(numsteps,np)

  stop = time()

  elapsed = (stop - start)

  println("Obtained value of PI: ", mypi)
  println("Time taken: ", elapsed, " seconds")

end

numsteps=1000000000

if length(ARGS) > 0
  numsteps = parse(Int, ARGS[1])
end

picalc(numsteps)


numbers = [0,14,1,3,7,9]
# numbers = [1,3,2]
# numbers = [2,1,3]

function get_next_number!(spoken)
    idx = findlast(==(spoken[end]), spoken[1:end-1])
    if idx === nothing
        push!(spoken, 0)
    else
        push!(spoken, length(spoken) - idx)
    end
end

spoken_history = [numbers...]
for i in length(numbers)+1:2020
    global spoken_history
    get_next_number!(spoken_history)
end

println(spoken_history[end])


using Mimi

@defcomp climatedynamics begin
    TATM = Variable(index=[time]) # Increase in temperature of atmosphere (degrees C from 1900)
    TOCEAN = Variable(index=[time]) # Increase in temperature of lower oceans (degrees C from 1900)

    FORC = Parameter(index=[time]) # Increase in radiative forcing (watts per m2 from 1900)
    fco22x = Parameter() # Forcings of equilibrium CO2 doubling (Wm-2)
    t2xco2 = Parameter() # Equilibrium temp impact (oC per doubling CO2)
    tatm0 = Parameter() # Initial atmospheric temp change (C from 1900)
    tatm1 = Parameter() # Period 2 atmospheric temp change (C from 1900)
    tocean0 = Parameter() # Initial lower stratum temp change (C from 1900)

    # Transient TSC Correction ("Speed of Adjustment Parameter")
    c1 = Parameter() # Climate equation coefficient for upper level
    c3 = Parameter() # Transfer coefficient upper to lower stratum
    c4 = Parameter() # Transfer coefficient for lower level

    function run_timestep(p, v, d, t)
        #Define function for TATM
        if is_first(t)
            v.TATM[t] = p.tatm0
        elseif t.t == 2
            v.TATM[t] = p.tatm1
        else
            v.TATM[t] = v.TATM[t-1] + p.c1 * ((p.FORC[t] - (p.fco22x/p.t2xco2) * v.TATM[t-1]) - (p.c3 * (v.TATM[t-1] - v.TOCEAN[t-1])))
        end

        #Define function for TOCEAN
        if is_first(t)
            v.TOCEAN[t] = p.tocean0
        else
            v.TOCEAN[t] = v.TOCEAN[t-1] + p.c4 * (v.TATM[t-1] - v.TOCEAN[t-1])
        end
    end
end


include("../Resistances/heat.jl")
include("../../Utilities/summary.jl")

using Measurements, Modia
#ENV["MODIA_PLOT"] = "PyPlot"
Modia.usePlotPackage("PyPlot")
@usingModiaPlot

# Boundary model(s)
T_a = 95.6u"K"
T_b = 273.15u"K"

# Parameters
r_inner = 0.1u"m"#measurement(0.1,0.01)u"m"
th_1 = 0.0025u"m"
th_ins = 0.01u"m"
th_2 = 0.0025u"m"

r_ins_inner = r_inner + th_1
r_ins_outer = r_ins_inner + th_ins
r_outer = r_ins_outer + th_2

model = Model(
    boundaryT_a = Temperature | Map(T=T_a),
    boundaryT_b = Temperature | Map(T=T_b),
    convectionInner = Convection | Map(surfaceArea=4*π*r_inner^2,alpha=150u"W/(m^2*K)"),
    linerInner = Sphere | Map(r_inner=r_inner,r_outer=r_ins_inner,lambda=15u"W/(m*K)"),
    contact_1 = Contact | Map(surfaceArea=4*π*r_ins_inner^2,Rc_pp=0.003u"m^2*K/W"),
    insulation = Sphere | Map(r_inner=r_ins_inner,r_outer=r_ins_outer,lambda=0.033u"W/(m*K)"),
    contact_2 = Contact | Map(surfaceArea=4*π*r_ins_outer^2,Rc_pp=0.003u"m^2*K/W"),
    linerOuter = Sphere | Map(r_inner=r_ins_outer,r_outer=r_outer,lambda=15u"W/(m*K)"),
    radiationOuter = Radiation | Map(surfaceArea=4*π*r_outer^2,epsilon=0.7u"m/m"), #nondim units?
    convectionOuter = Convection | Map(surfaceArea=4*π*r_outer^2,alpha=6.0u"W/(m^2*K)"),
    Q_total=Var(init=0.0u"J"),
    equations =:[
        der(Q_total) = convectionInner.port_a.Q_flow
    ],
    connect =:[
        (convectionInner.port_a,boundaryT_a.port)
        (convectionInner.port_b,linerInner.port_a)
        (linerInner.port_b,contact_1.port_a)
        (contact_1.port_b,insulation.port_a)
        (insulation.port_b,contact_2.port_a)
        (contact_2.port_b,linerOuter.port_a)
        (linerOuter.port_b,radiationOuter.port_a,convectionOuter.port_a)
        (convectionOuter.port_b,radiationOuter.port_b,boundaryT_b.port)
    ]
)

result = @instantiateModel(model, log=false,FloatType=Measurement{Float64})
@time simulate!(result, stopTime=1.0, tolerance=1e-4, log=false)
y_min, t_min, y_max, t_max = reportMinMax("Q_total","Thermal Resistances Model",result)
errorCheck(-69.43455u"J"±0.0u"J", y_min)
plot(result, ["Q_total"], figure=1)


using AutoGrad

"""
    svd(A::TrackedArray) -> TrackedTuple

Return tracked tuple of (U, S, V) that `A == USV'`.
"""
@primitive svd!(A), Δ, y svd_back(y..., Δ...)

A = randn(4,4) |> Param
O = randn(4,4)
function test_func(A)
    U, S, V = svd!(A)
    psi = U[:,1]
    psi'*O*psi
end
dU, dS, dV = randn(4,4), randn(4), randn(4,4)
y = @diff test_func(A)
@show value(y)
grad(y, A)

qr(A::TrackedArray) = track(qr, A)
function _forward(::typeof(qr), A)
    Q, R = qr(data(A))
    (Q, R), Δ -> (qr_back(data(A), Q, R, Δ...),)
end

lq(A::TrackedArray) = track(lq, A)
function _forward(::typeof(lq), A)
    L, Q = lq(data(A))
    (L, Q), Δ -> (lq_back(data(A), L, Q, Δ...),)
end


using Gen
import Random

include("static_model.jl")
include("dataset.jl")

@gen (static) function slope_proposal(trace)
    slope = trace[:slope]
    @trace(normal(slope, 0.5), :slope)
end

@gen (static) function intercept_proposal(trace)
    intercept = trace[:intercept]
    @trace(normal(intercept, 0.5), :intercept)
end

@gen (static) function inlier_std_proposal(trace)
    log_inlier_std = trace[:log_inlier_std]
    @trace(normal(log_inlier_std, 0.5), :log_inlier_std)
end

@gen (static) function outlier_std_proposal(trace)
    log_outlier_std = trace[:log_outlier_std]
    @trace(normal(log_outlier_std, 0.5), :log_outlier_std)
end

@gen (static) function flip_z(z::Bool)
    @trace(bernoulli(z ? 0.0 : 1.0), :z)
end

@gen (static) function is_outlier_proposal(trace, i::Int)
    prev_z = trace[:data => i => :z]
    @trace(bernoulli(prev_z ? 0.0 : 1.0), :data => i => :z)
end

@gen (static) function is_outlier_proposal(trace, i::Int)
    prev_z = trace[:data => i => :z]
    @trace(bernoulli(prev_z ? 0.0 : 1.0), :data => i => :z)
end

Gen.load_generated_functions()

function do_inference(xs, ys, num_iters)
    observations = choicemap()
    for (i, y) in enumerate(ys)
        observations[:data => i => :y] = y
    end

    # initial trace
    (trace, _) = generate(model, (xs,), observations)

    scores = Vector{Float64}(undef, num_iters)
    for i=1:num_iters

        # steps on the parameters
        for j=1:5
            (trace, _) = metropolis_hastings(trace, slope_proposal, ())
            (trace, _) = metropolis_hastings(trace, intercept_proposal, ())
            (trace, _) = metropolis_hastings(trace, inlier_std_proposal, ())
            (trace, _) = metropolis_hastings(trace, outlier_std_proposal, ())
        end

        # step on the outliers
        for j=1:length(xs)
            (trace, _) = metropolis_hastings(trace, is_outlier_proposal, (j,))
        end

        score = get_score(trace)
        scores[i] = score

        # print
        slope = trace[:slope]
        intercept = trace[:intercept]
        inlier_std = exp(trace[:log_inlier_std])
        outlier_std = exp(trace[:log_outlier_std])
        println("score: $score, slope: $slope, intercept: $intercept, inlier_std: $inlier_std, outlier_std: $outlier_std")
    end
    return scores
end

(xs, ys) = make_data_set(200)
do_inference(xs, ys, 10)
@time do_inference(xs, ys, 50)

# handcoded version

mutable struct State
    xs::Vector{Float64}
    log_inlier_std::Float64
    log_outlier_std::Float64
    slope::Float64
    intercept::Float64
    is_outlier::Vector{Bool}
    ys::Vector{Float64}
end

function init_state(xs, ys)
    log_inlier_std = random(normal, 0, 2)
    log_outlier_std = random(normal, 0, 2)
    slope = random(normal, 0, 2)
    intercept = random(normal, 0, 2)
    is_outlier = Vector{Bool}(undef, length(xs))
    for i=1:length(xs)
        is_outlier[i] = (rand() < 0.5)
    end
    State(xs, log_inlier_std, log_outlier_std, slope, intercept, is_outlier, ys)
end

function log_likelihood(x, y, is_outlier, log_inlier_std, log_outlier_std, slope, intercept)
    std = is_outlier ? exp(log_outlier_std) : exp(log_inlier_std)
    logpdf(normal, y, slope * x + intercept, std)
end

function log_likelihood(state)
    ll = 0.
    for (x, is_outlier, y) in zip(state.xs, state.is_outlier, state.ys)
        ll += log_likelihood(x, y, is_outlier,
                    state.log_inlier_std, state.log_outlier_std, state.slope, state.intercept)
    end
    ll
end

macro parameter_mh_move!(state, addr)
    quote
        prev_ll = log_likelihood($(esc(state)))
        prev_val = $(esc(state)).$addr
        new_val = normal(prev_val, 0.5)
        prev_prior = logpdf(normal, prev_val, 0, 2)
        new_prior = logpdf(normal, new_val, 0, 2)
        $(esc(state)).$addr = new_val
        new_ll = log_likelihood($(esc(state)))
        if log(rand()) >= new_ll - prev_ll + new_prior - prev_prior
            $(esc(state)).$addr = prev_val # reject
        end
    end
end

function is_outlier_mh_move!(state, i::Int)
    prev_val = state.is_outlier[i]
    new_val = !prev_val
    x = state.xs[i]
    y = state.ys[i]
    prev_ll = log_likelihood(x, y, prev_val,
        state.log_inlier_std, state.log_outlier_std, state.slope, state.intercept)
    new_ll = log_likelihood(x, y, new_val,
        state.log_inlier_std, state.log_outlier_std, state.slope, state.intercept)
    if log(rand()) < new_ll - prev_ll
        state.is_outlier[i] = new_val
    end
end

function handcoded_inference(xs, ys, iters)
    state = init_state(xs, ys)
    for iter=1:iters
        for j=1:5
            @parameter_mh_move!(state, slope)
            @parameter_mh_move!(state, intercept)
            @parameter_mh_move!(state, log_inlier_std)
            @parameter_mh_move!(state, log_outlier_std)
        end
        for j=1:length(xs)
            is_outlier_mh_move!(state, j)
        end
    end
    println("slope: $(state.slope)")
    println("intercept: $(state.intercept)")
end

handcoded_inference(xs, ys, 50)
@time handcoded_inference(xs, ys, 50)


using Quadrature, Cuba, Cubature, Zygote, FiniteDiff, ForwardDiff
using Test

### One Dimensional

f(x,p) = sum(sin.(p[1] * x))
lb = 1.0
ub = 3.0
p = 2.0
prob = QuadratureProblem(f,lb,ub,p)
sol = solve(prob,QuadGKJL(),reltol=1e-3,abstol=1e-3)

function testf(lb,ub,p)
    prob = QuadratureProblem(f,lb,ub,p)
    sin(solve(prob,QuadGKJL(),reltol=1e-3,abstol=1e-3)[1])
end
dlb1,dub1,dp1 = Zygote.gradient(testf,lb,ub,p)
dlb2 = FiniteDiff.finite_difference_derivative(lb->testf(lb,ub,p),lb)
dub2 = FiniteDiff.finite_difference_derivative(ub->testf(lb,ub,p),ub)
dp2 = FiniteDiff.finite_difference_derivative(p->testf(lb,ub,p),p)

# dlb3 = ForwardDiff.derivative(lb->testf(lb,ub,p),lb)
# dub3 = ForwardDiff.derivative(ub->testf(lb,ub,p),ub)
dp3 = ForwardDiff.derivative(p->testf(lb,ub,p),p)

#@test dlb1 ≈ dlb3
#@test dub1 ≈ dub3
@test dp1 ≈ dp3

### N-dimensional

f(x,p) = sum(sin.(x .* p))
lb = ones(2)
ub = 3ones(2)
p = [1.5,2.0]
prob = QuadratureProblem(f,lb,ub,p)
sol = solve(prob,CubaCuhre(),reltol=1e-3,abstol=1e-3)

function testf(lb,ub,p)
    prob = QuadratureProblem(f,lb,ub,p)
    sin(solve(prob,CubaCuhre(),reltol=1e-6,abstol=1e-6)[1])
end

function testf2(lb,ub,p)
    prob = QuadratureProblem(f,lb,ub,p)
    sin(solve(prob,HCubatureJL(),reltol=1e-6,abstol=1e-6)[1])
end

dlb1,dub1,dp1 = Zygote.gradient(testf,lb,ub,p)
dlb2 = FiniteDiff.finite_difference_gradient(lb->testf(lb,ub,p),lb)
dub2 = FiniteDiff.finite_difference_gradient(ub->testf(lb,ub,p),ub)
dp2 = FiniteDiff.finite_difference_gradient(p->testf(lb,ub,p),p)

@test_broken dlb1 ≈ dlb2
@test_broken dub1 ≈ dub2
@test dp1 ≈ dp2

# dlb3 = ForwardDiff.gradient(lb->testf(lb,ub,p),lb)
# dub3 = ForwardDiff.gradient(ub->testf(lb,ub,p),ub)
dp3 = ForwardDiff.gradient(p->testf2(lb,ub,p),p)

#@test dlb1 ≈ dlb3
#@test dub1 ≈ dub3
@test dp1 ≈ dp3

# dlb4 = ForwardDiff.gradient(lb->testf(lb,ub,p),lb)
# dub4 = ForwardDiff.gradient(ub->testf(lb,ub,p),ub)
dp4 = ForwardDiff.gradient(p->testf(lb,ub,p),p)

#@test dlb1 ≈ dlb4
#@test dub1 ≈ dub4
@test dp1 ≈ dp4

### N-dimensional N-out

f(x,p) = sin.(x .* p)
lb = ones(2)
ub = 3ones(2)
p = [1.5,2.0]
prob = QuadratureProblem(f,lb,ub,p,nout=2)
sol = solve(prob,CubaCuhre(),reltol=1e-3,abstol=1e-3)

function testf(lb,ub,p)
    prob = QuadratureProblem(f,lb,ub,p,nout=2)
    sin(sum(solve(prob,CubaCuhre(),reltol=1e-6,abstol=1e-6)))
end

function testf2(lb,ub,p)
    prob = QuadratureProblem(f,lb,ub,p,nout=2)
    sin(sum(solve(prob,HCubatureJL(),reltol=1e-6,abstol=1e-6)))
end

dlb1,dub1,dp1 = Zygote.gradient(testf,lb,ub,p)
dlb2 = FiniteDiff.finite_difference_gradient(lb->testf(lb,ub,p),lb)
dub2 = FiniteDiff.finite_difference_gradient(ub->testf(lb,ub,p),ub)
dp2 = FiniteDiff.finite_difference_gradient(p->testf(lb,ub,p),p)

@test_broken dlb1 ≈ dlb2
@test_broken dub1 ≈ dub2
@test dp1 ≈ dp2

# dlb3 = ForwardDiff.gradient(lb->testf(lb,ub,p),lb)
# dub3 = ForwardDiff.gradient(ub->testf(lb,ub,p),ub)
dp3 = ForwardDiff.gradient(p->testf2(lb,ub,p),p)

#@test dlb1 ≈ dlb3
#@test dub1 ≈ dub3
@test dp1 ≈ dp3

# dlb4 = ForwardDiff.gradient(lb->testf(lb,ub,p),lb)
# dub4 = ForwardDiff.gradient(ub->testf(lb,ub,p),ub)
dp4 = ForwardDiff.gradient(p->testf(lb,ub,p),p)

#@test dlb1 ≈ dlb4
#@test dub1 ≈ dub4
@test dp1 ≈ dp4

### Batch Single dim
f(x,p) = x*p[1].+p[2]*p[3]

lb =1.0
ub = 3.0
p = [2.0, 3.0, 4.0]
prob = QuadratureProblem(f,lb,ub,p)

function testf3(lb,ub,p; f=f)
    prob = QuadratureProblem(f,lb,ub,p, batch = 10, nout=1)
    solve(prob, CubatureJLh(); reltol=1e-3,abstol=1e-3)[1]
end

dp1 = ForwardDiff.gradient(p->testf3(lb,ub,p),p)
dp2 = Zygote.gradient(p->testf3(lb,ub,p),p)[1]
dp3 = FiniteDiff.finite_difference_gradient(p->testf3(lb,ub,p),p)

@test dp1 ≈ dp3 #passes
@test dp2 ≈ dp3 #passes

### Batch single dim, nout
f(x,p) = (x*p[1].+p[2]*p[3]).*[1;2]

lb =1.0
ub = 3.0
p = [2.0, 3.0, 4.0]
prob = QuadratureProblem(f,lb,ub,p)

function testf3(lb,ub,p; f=f)
    prob = QuadratureProblem(f,lb,ub,p, batch = 10, nout=2)
    sum(solve(prob, CubatureJLh(); reltol=1e-3,abstol=1e-3))
end

dp1 = ForwardDiff.gradient(p->testf3(lb,ub,p),p)
# dp2 = Zygote.gradient(p->testf3(lb,ub,p),p)[1]
dp3 = FiniteDiff.finite_difference_gradient(p->testf3(lb,ub,p),p)

@test dp1 ≈ dp3 #passes
# @test dp2 ≈ dp3 #passes

### Batch multi dim
f(x,p) = x[1,:]*p[1].+p[2]*p[3]

lb =[1.0,1.0]
ub = [3.0,3.0]
p = [2.0, 3.0, 4.0]
prob = QuadratureProblem(f,lb,ub,p)

function testf3(lb,ub,p; f=f)
    prob = QuadratureProblem(f,lb,ub,p, batch = 10, nout=1)
    solve(prob, CubatureJLh(); reltol=1e-3,abstol=1e-3)[1]
end

dp1 = ForwardDiff.gradient(p->testf3(lb,ub,p),p)
dp2 = Zygote.gradient(p->testf3(lb,ub,p),p)[1]
dp3 = FiniteDiff.finite_difference_gradient(p->testf3(lb,ub,p),p)

@test dp1 ≈ dp3
@test dp2 ≈ dp3

### Batch multi dim nout
f(x,p) = x*p[1].+p[2]*p[3]

lb =[1.0,1.0]
ub = [3.0,3.0]
p = [2.0, 3.0, 4.0]
prob = QuadratureProblem(f,lb,ub,p)

function testf3(lb,ub,p; f=f)
    prob = QuadratureProblem(f,lb,ub,p, batch = 10, nout=2)
    sum(solve(prob, CubatureJLh(); reltol=1e-3,abstol=1e-3))
end

dp1 = ForwardDiff.gradient(p->testf3(lb,ub,p),p)
# dp2 = Zygote.gradient(p->testf3(lb,ub,p),p)[1]
dp3 = FiniteDiff.finite_difference_gradient(p->testf3(lb,ub,p),p)

@test dp1 ≈ dp3
# @test dp2 ≈ dp3

## iip Batch mulit dim
function g(dx, x,p)
    dx.= sum(x*p[1].+p[2]*p[3], dims=1)
end

lb =[1.0,1.0]
ub = [3.0,3.0]
p = [2.0, 3.0, 4.0]
prob = QuadratureProblem(g,lb,ub,p)

function testf3(lb,ub,p; f=g)
    prob = QuadratureProblem(f,lb,ub,p, batch = 10, nout=1)
    solve(prob, CubatureJLh(); reltol=1e-3,abstol=1e-3)[1]
end

testf3(lb,ub,p)

dp1 = ForwardDiff.gradient(p->testf3(lb,ub,p),p)
dp2 = Zygote.gradient(p->testf3(lb,ub,p),p)[1]
dp3 = FiniteDiff.finite_difference_gradient(p->testf3(lb,ub,p),p)

@test dp1 ≈ dp3
@test dp2 ≈ dp3

## iip Batch mulit dim nout
function g(dx, x,p)
    dx.= x*p[1].+p[2]*p[3]
end

lb =[1.0,1.0]
ub = [3.0,3.0]
p = [2.0, 3.0, 4.0]
prob = QuadratureProblem(g,lb,ub,p)

function testf3(lb,ub,p; f=g)
    prob = QuadratureProblem(f,lb,ub,p, batch = 10, nout=2)
    sum(solve(prob, CubatureJLh(); reltol=1e-3,abstol=1e-3))
end

dp1 = ForwardDiff.gradient(p->testf3(lb,ub,p),p)
# dp2 = Zygote.gradient(p->testf3(lb,ub,p),p)[1]
dp3 = FiniteDiff.finite_difference_gradient(p->testf3(lb,ub,p),p)

@test dp1 ≈ dp3
# @test dp2 ≈ dp3


module TrainingSetGenerator

#Declare linear function, which separates observations in 2D space
f(x) = 8*x+3.5

#Generating training data with labels for given N lenght
function generate_data(n)
  training_set = []
  for i in 1:n
    x = rand(0:200)
    y = rand(0:200)
    bias = rand(0:200)
    lineY = f(x)
    if y>lineY
      label = -1
    else
      label = 1
    end
    input_values = [x,y, bias]
    push!(training_set, [input_values, label])
  end
  return training_set
end

end


# ## Encoders

function encode_simpleArithmetic(x; encoder_size,in_max,in_min,buckets,range,w)
  b= floor(Int, (x-in_min) .* buckets./range) +1
  sdr= falses(encoder_size)
  sdr[b:b+w-1].= true
  sdr,b
end
function encode_powerDay(power,time,wkend; power_p,time_p,wkend_p)
  sdr_power,bucket= encode_simpleArithmetic(power; power_p...)
  sdr_time,_= encode_simpleArithmetic(time; time_p...)
  sdr_wkend,_= encode_simpleArithmetic(wkend; wkend_p...)
  [sdr_power;sdr_time;sdr_wkend], bucket
end

# ## Encoder parameter generators
function initenc_simpleArithmetic(ts; encoder_size,buckets=0,w=0)
  in_max= maximum(ts); in_min= minimum(ts)
  range= in_max - in_min +10*eps(in_max|>Float32)
  if buckets>0
    w= prod(encoder_size) - buckets +1
  elseif w>0
    buckets= prod(encoder_size) - w +1;
  else
    error("[initenc_simpleArithmetic]: either buckets or w must be provided")
  end
  (in_max=in_max, in_min=in_min, range=range, buckets=buckets, w=w, encoder_size=encoder_size)
end
initenc_powerDay(power,time,wkend; encoder_size, w)= (
  power_p= initenc_simpleArithmetic(power,encoder_size=encoder_size[1], w=w[1]),
  time_p=  initenc_simpleArithmetic(time,encoder_size=encoder_size[2], w=w[2]),
  wkend_p= initenc_simpleArithmetic(wkend,encoder_size=encoder_size[3], w=w[3])
)


### Edge and Gradient related tests ###

EPS = 1e-14

## Checkerboard array, used to test image gradients

let
    white{T}(::Type{T}) = one(T)
    black{T}(::Type{T}) = zero(T)
    white{T<:Unsigned}(::Type{T}) = typemax(T)
    black{T<:Unsigned}(::Type{T}) = typemin(T)

    global checkerboard
    function checkerboard{T}(::Type{T}, sq_width::Integer, count::Integer)
        wh = fill(white(T), (sq_width,sq_width))
        bk = fill(black(T), (sq_width,sq_width))
        bw = [wh bk; bk wh]
        vert = repmat(bw, (count>>1), 1)
        isodd(count) && (vert = vcat(vert, [wh bk]))
        cb = repmat(vert, 1, (count>>1))
        isodd(count) && (cb = hcat(cb, vert[:,1:sq_width]))
        cb
    end

    checkerboard(sq_width::Integer, count::Integer) = checkerboard(Uint8, sq_width, count)
end

SZ=5

cb_array    = checkerboard(SZ,3)
cb_image_xy = grayim(cb_array)
cb_image_yx = grayim(cb_array)
cb_image_yx["spatialorder"] = ["y","x"]

for method in ["sobel", "prewitt", "ando3", "ando4", "ando5", "ando4_sep", "ando5_sep"]
    ## Checkerboard array

    (agx, agy) = imgradients(cb_array, method)
    amag = magnitude(agx, agy)
    agphase = phase(agx, agy)
    @assert (amag, agphase) == magnitude_phase(agx, agy)

    @assert agx[1,SZ]   < 0.0   # white to black transition
    @assert agx[1,2*SZ] > 0.0   # black to white transition
    @assert agy[SZ,1]   < 0.0   # white to black transition
    @assert agy[2*SZ,1] > 0.0   # black to white transition

    # Test direction of increasing gradient
    @assert cos(agphase[1,SZ])   - (-1.0) < EPS   # increasing left  (=  pi   radians)
    @assert cos(agphase[1,2*SZ]) -   1.0  < EPS   # increasing right (=   0   radians)
    @assert sin(agphase[SZ,1])   -   1.0  < EPS   # increasing up    (=  pi/2 radians)
    @assert sin(agphase[2*SZ,1]) - (-1.0) < EPS   # increasing down  (= -pi/2 radians)

    # Test that orientation is perpendicular to gradient
    aorient = orientation(agx, agy)
    @assert all((cos(agphase).*cos(aorient) .+ sin(agphase).*sin(aorient) .< EPS) |
                ((agphase .== 0.0) & (aorient .== 0.0)))  # this part is where both are
                                                          # zero because there is no gradient

    ## Checkerboard Image with row major order

    (gx, gy) = imgradients(cb_image_xy, method)
    mag = magnitude(gx, gy)
    gphase = phase(gx, gy)
    @assert (mag, gphase) == magnitude_phase(gx, gy)

    @assert gx[SZ,1]   < 0.0   # white to black transition
    @assert gx[2*SZ,1] > 0.0   # black to white transition
    @assert gy[1,SZ]   < 0.0   # white to black transition
    @assert gy[1,2*SZ] > 0.0   # black to white transition

    @assert cos(gphase[SZ,1])   - (-1.0) < EPS   # increasing left  (=  pi   radians)
    @assert cos(gphase[2*SZ,1]) -   1.0  < EPS   # increasing right (=   0   radians)
    @assert sin(gphase[1,SZ])   -   1.0  < EPS   # increasing up    (=  pi/2 radians)
    @assert sin(gphase[1,2*SZ]) - (-1.0) < EPS   # increasing down  (= -pi/2 radians)

    # Test that orientation is perpendicular to gradient
    orient = orientation(gx, gy)
    @assert all((cos(gphase).*cos(orient) .+ sin(gphase).*sin(orient) .< EPS) |
                ((gphase .== 0.0) & (orient .== 0.0)))  # this part is where both are
                                                        # zero because there is no gradient

    ## Checkerboard Image with column-major order

    (gx, gy) = imgradients(cb_image_yx, method)
    mag = magnitude(gx, gy)
    gphase = phase(gx, gy)
    @assert (mag, gphase) == magnitude_phase(gx, gy)

    @assert gx[1,SZ]   < 0.0   # white to black transition
    @assert gx[1,2*SZ] > 0.0   # black to white transition
    @assert gy[SZ,1]   < 0.0   # white to black transition
    @assert gy[2*SZ,1] > 0.0   # black to white transition

    # Test direction of increasing gradient
    @assert cos(gphase[1,SZ])   - (-1.0) < EPS   # increasing left  (=  pi   radians)
    @assert cos(gphase[1,2*SZ]) -   1.0  < EPS   # increasing right (=   0   radians)
    @assert sin(gphase[SZ,1])   -   1.0  < EPS   # increasing up    (=  pi/2 radians)
    @assert sin(gphase[2*SZ,1]) - (-1.0) < EPS   # increasing down  (= -pi/2 radians)

    # Test that orientation is perpendicular to gradient
    orient = orientation(gx, gy)
    @assert all((cos(gphase).*cos(orient) .+ sin(gphase).*sin(orient) .< EPS) |
                ((gphase .== 0.0) & (orient .== 0.0)))  # this part is where both are
                                                        # zero because there is no gradient

end

# Create an image with white along diagonals -2:2 and black elsewhere
m = zeros(Uint8, 20,20)
for i = -2:2; m[diagind(m,i)] = 0xff; end

m_xy = grayim(m')
m_yx = grayim(m)
m_yx["spatialorder"] = ["y","x"]

for method in ["sobel", "prewitt", "ando3", "ando4", "ando5", "ando4_sep", "ando5_sep"]
    ## Diagonal array

    (agx, agy) = imgradients(m, method)
    amag = magnitude(agx, agy)
    agphase = phase(agx, agy)
    @assert (amag, agphase) == magnitude_phase(agx, agy)

    @assert agx[7,9]  < 0.0   # white to black transition
    @assert agx[10,8] > 0.0   # black to white transition
    @assert agy[10,8] < 0.0   # white to black transition
    @assert agy[7,9]  > 0.0   # black to white transition

    # Test direction of increasing gradient
    @assert abs(agphase[10,8] -    pi/4 ) < EPS   # lower edge (increasing up-right  =   pi/4 radians)
    @assert abs(agphase[7,9]  - (-3pi/4)) < EPS   # upper edge (increasing down-left = -3pi/4 radians)

    # Test that orientation is perpendicular to gradient
    aorient = orientation(agx, agy)
    @assert all((cos(agphase).*cos(aorient) .+ sin(agphase).*sin(aorient) .< EPS) |
                ((agphase .== 0.0) & (aorient .== 0.0)))  # this part is where both are
                                                          # zero because there is no gradient

    ## Diagonal Image, row-major order

    (gx, gy) = imgradients(m_xy, method)
    mag = magnitude(gx, gy)
    gphase = phase(gx, gy)
    @assert (mag, gphase) == magnitude_phase(gx, gy)

    @assert gx[9,7]  < 0.0   # white to black transition
    @assert gx[8,10] > 0.0   # black to white transition
    @assert gy[8,10] < 0.0   # white to black transition
    @assert gy[9,7]  > 0.0   # black to white transition

    # Test direction of increasing gradient
    @assert abs(gphase[8,10] -    pi/4 ) < EPS   # lower edge (increasing up-right  =   pi/4 radians)
    @assert abs(gphase[9,7]  - (-3pi/4)) < EPS   # upper edge (increasing down-left = -3pi/4 radians)

    # Test that orientation is perpendicular to gradient
    orient = orientation(gx, gy)
    @assert all((cos(gphase).*cos(orient) .+ sin(gphase).*sin(orient) .< EPS) |
                ((gphase .== 0.0) & (orient .== 0.0)))  # this part is where both are
                                                        # zero because there is no gradient

    ## Diagonal Image, column-major order

    (gx, gy) = imgradients(m_yx, method)
    mag = magnitude(gx, gy)
    gphase = phase(gx, gy)
    @assert (mag, gphase) == magnitude_phase(gx, gy)

    @assert gx[7,9]  < 0.0   # white to black transition
    @assert gx[10,8] > 0.0   # black to white transition
    @assert gy[10,8] < 0.0   # white to black transition
    @assert gy[7,9]  > 0.0   # black to white transition

    # Test direction of increasing gradient
    @assert abs(gphase[10,8] -    pi/4 ) < EPS   # lower edge (increasing up-right  =   pi/4 radians)
    @assert abs(gphase[7,9]  - (-3pi/4)) < EPS   # upper edge (increasing down-left = -3pi/4 radians)

    # Test that orientation is perpendicular to gradient
    orient = orientation(gx, gy)
    @assert all((cos(gphase).*cos(orient) .+ sin(gphase).*sin(orient) .< EPS) |
                ((gphase .== 0.0) & (orient .== 0.0)))  # this part is where both are
                                                        # zero because there is no gradient
end

# Nonmaximal suppression

function thin_edges(img)
    # Get orientation
    gx,gy = imgradients(img)
    orient = phase(gx,gy)

    # Do NMS thinning
    thin_edges_nonmaxsup_subpix(img, orient, radius=1.35)
end

function nms_test_horiz_vert(img, which)
    ## which = :horizontal or :vertical

    # Do NMS thinning
    t,s = thin_edges(img)

    # Calc peak location by hand

    # Interpolate values 1.35 pixels left and right
    # Orientation is zero radians -> to the right
    v1 = 6 - 0.35   # slope on right is 1
    v2 = 5 - 0.35*2 # slope on left is 2
    c = 7.0         # peak value

    # solve v = a*r^2 + b*r + c
    a = (v1 + v2)/2 - c
    b = a + c - v2
    r = -b/2a

    @assert abs(r - 1/6) < EPS

    # Location and value at peak
    peakloc = r*1.35 + 3
    peakval = a*r^2 + b*r + c

    transposed = spatialorder(img)[1] == "x"
    horizontal = which == :horizontal

    test_axis1 = transposed $ !horizontal

    @assert test_axis1 ? all(t[:,[1,2,4,5]] .== 0) : all(t[[1,2,4,5],:] .== 0)
    @assert test_axis1 ? all(t[:,3]   .== peakval) : all(t[3,:]   .== peakval)
    @assert test_axis1 ? all(s[:,[1,2,4,5]] .== zero(Base.Graphics.Point)) :
                         all(s[[1,2,4,5],:] .== zero(Base.Graphics.Point))

    if transposed
        if which == :horizontal
            @assert     [pt.x for pt in s[:,3]]  == [1:5]
            @assert all([pt.y for pt in s[:,3]] .== peakloc)
        else
            @assert all([pt.x for pt in s[3,:]] .== peakloc)
            @assert     [pt.y for pt in s[3,:]]  == [1:5]
        end
    else
        if which == :horizontal
            @assert     [pt.x for pt in s[3,:]]  == [1:5]
            @assert all([pt.y for pt in s[3,:]] .== peakloc)
        else
            @assert all([pt.x for pt in s[:,3]] .== peakloc)
            @assert     [pt.y for pt in s[:,3]]  == [1:5]
        end
    end
end

# Test image: vertical edge
m = [3.0  5.0  7.0  6.0  5.0
     3.0  5.0  7.0  6.0  5.0
     3.0  5.0  7.0  6.0  5.0
     3.0  5.0  7.0  6.0  5.0
     3.0  5.0  7.0  6.0  5.0]

m_xy = grayim(m')
m_yx = grayim(m)
m_yx["spatialorder"] = ["y","x"]

nms_test_horiz_vert(m, :vertical)
nms_test_horiz_vert(m_xy, :vertical)
nms_test_horiz_vert(m_yx, :vertical)

# Test image: horizontal edge
m = m'
m_xy = grayim(m')
m_yx = grayim(m)
m_yx["spatialorder"] = ["y","x"]

nms_test_horiz_vert(m, :horizontal)
nms_test_horiz_vert(m_xy, :horizontal)
nms_test_horiz_vert(m_yx, :horizontal)

function nms_test_diagonal(img)
    # Do NMS thinning
    t,s = thin_edges(img)

    # Calc peak location by hand

    # Interpolate values 1.35 pixels up and left, down and right
    # using bilinear interpolation
    # Orientation is π/4 radians -> 45 degrees up
    fr = 1.35*cos(π/4)
    lower = (7 + fr*(6-7))
    upper = (6 + fr*(5-6))
    v1 = lower + fr*(upper-lower)

    lower = (7 + fr*(5-7))
    upper = (5 + fr*(3-5))
    v2 = lower + fr*(upper-lower)

    c = 7.0         # peak value

    # solve v = a*r^2 + b*r + c
    a = (v1 + v2)/2 - c
    b = a + c - v2
    r = -b/2a

    @assert (r - 1/6) < EPS

    transposed = spatialorder(img)[1] == "x"

    # Location and value at peak

    x_peak_offset, y_peak_offset = r*fr, -r*fr
    peakval = a*r^2 + b*r + c

    @assert all(diag(data(t))[2:4] .== peakval)  # Edge pixels aren't interpolated here
    @assert all(t - diagm(diag(data(t))) .== 0)

    diag_s = copy(s, diagm(diag(data(s))))
    @assert s == diag_s

    @assert all([pt.x for pt in diag(data(s))[2:4]] - ([2:4] + x_peak_offset) .< EPS)
    @assert all([pt.y for pt in diag(data(s))[2:4]] - ([2:4] + y_peak_offset) .< EPS)

end

# Test image: diagonal edge
m = [7.0  6.0  5.0  0.0  0.0
     5.0  7.0  6.0  5.0  0.0
     3.0  5.0  7.0  6.0  5.0
     0.0  3.0  5.0  7.0  6.0
     0.0  0.0  3.0  5.0  7.0]

m_xy = grayim(m')
m_yx = grayim(m)
m_yx["spatialorder"] = ["y","x"]

nms_test_diagonal(m)
nms_test_diagonal(m_xy)
nms_test_diagonal(m_yx)


# ==============================================================================
using PyPlot, AcousticWaveLab
nz = 201; nx = 373; ext = 30; iflag = 2; dz = 10.; dx = 10.; dt = 1e-3;
v  = 3000. * ones(nz, nx);
v[101:end, :] = 4000.
# initialize source term
pos = [1 180]; f0 = 30.0; ot = [0.0];
src = InitSources(pos, f0, ot, dt);

# discretize spatial derivative operator
vmax = maximum(v); vmin = minimum(v);
fidMtx = InitFidMtx(nz, nx, ext, iflag, dz, dx, dt, vmax, vmin, f0, v);

ix = collect(1:5:nx); iz = 1 * ones(Int64, length(ix));
pos = hcat(iz, ix)
shot = MultiStepForward(pos, src[1], fidMtx, print_flag=true)
SeisPlot(shot.d, pclip=98)
# ==============================================================================
# using PyPlot, AcousticWaveLab
# nz = 201; nx = 201; ext = 20; iflag = 1; dz = 10.; dx = 10.; dt = 2e-3;
# v  = 3000. * ones(nz, nx);
# # initialize source term
# pos = [101 101]; f0 = 30.0; ot = [0.0];
# src = InitSources(pos, f0, ot, dt);
# # discretize spatial derivative operator
# vmax = maximum(v); vmin = minimum(v);
# fidMtx = InitFidMtx(nz, nx, ext, iflag, dz, dx, dt, vmax, vmin, f0, v);
# path = "/Users/wenlei/Desktop/acoustic.bin"
# MultiStepForward(path, src, fidMtx, tmax=1.0);
#
# iz = 5 * ones(Int64, 201); ix = collect(1:201);
# pos = hcat(iz, ix)
# shot = fullWfd2Shot(pos, path);
# shot1 = MultiStepForward(pos, src[1], fidMtx, tmax=1.0)
#
# path = "/Users/wenlei/Desktop/adjoint.bin"
# MultiStepAdjoint!(path, shot, fidMtx)
#
#
# using PyPlot
# imshow(read(path, 500))
#
# iz = 5 * ones(Int64, 201); ix = collect(1:201);
# pos = hcat(iz, ix)
# shot = wfd2shot(pos, path);
# figure()
# plot(shot.d[:,71])

# ==============================================================================
# dot product test for one step
# spt1 = InitSnapShot(101, 101, nz, nx, ext, 1, dt, 1);
# spt1.vz = randn(length(spt1.vz));
# spt1.vx = randn(length(spt1.vx));
# spt1.pz = randn(length(spt1.pz));
# spt1.px = randn(length(spt1.px));
# spt2 = InitSnapShot(101, 101, nz, nx, ext, 1, dt, 1);
# OneStepForward!(spt2, spt1, fidMtx);
#
# tmp2 = InitSnapShot(101, 101, nz, nx, ext, 1, dt, 2)
# tmp2.vz = randn(length(tmp2.vz))
# tmp2.vx = randn(length(tmp2.vx))
# tmp2.pz = randn(length(tmp2.pz))
# tmp2.px = randn(length(tmp2.px))
# tmp1 = InitSnapShot(101, 101, nz, nx, ext, 1, dt, 2)
# OneStepAdjoint!(tmp1, tmp2, fidMtx)
#
# x  = vcat(spt1.vz, spt1.vx, spt1.pz, spt1.px)
# x1 = vcat(tmp1.vz, tmp1.vx, tmp1.pz, tmp1.px)
# y  = vcat(spt2.vz, spt2.vx, spt2.pz, spt2.px)
# y1 = vcat(tmp2.vz, tmp2.vx, tmp2.pz, tmp2.px)
#
# x'*x1 - y'*y1

# ==============================================================================
# nz = 201; nx = 201; dt = 2e-3; iflag = 1; ext = 20
# if iflag ==1
#    zupper = ext
#    Nz = nz + 2*ext
# elseif iflag == 2
#    Nz = nx +   ext
#    zupper = ext
# end
# Nx = nx + 2*ext
# spt1 = InitSnapShot(101, 101, nz, nx, ext, iflag, dt, 1);
# spt1.vz = randn(length(spt1.vz));
# spt1.vx = randn(length(spt1.vx));
# spt1.pz = randn(length(spt1.pz));
# spt1.px = randn(length(spt1.px));
#
# iz = collect(1:5:201); ix = collect(1:3:201);
# nrz= length(iz); nrx = length(ix);
# iz = vec(repmat(iz, nrx, 1));
# ix = vec(repmat(ix', (nrz), 1));
# pos = hcat(iz, ix);
# shot1 = InitShot(101, 101, pos, 0.0, 1, dt)
# spt2shot!(shot1, spt1, Nz, Nx, ext, zupper)
#
#
# shot2 = InitShot(101, 101, pos, 0.0, 1, dt)
# shot2.d = randn(1, size(shot2.d, 2))
# spt2 = InitSnapShot(101, 101, nz, nx, ext, iflag, dt, 1);
# AddShot2SnapShot!(spt2, shot2, Nz, Nx, ext, zupper)
#
# x  = vcat(spt1.vz, spt1.vx, spt1.pz, spt1.px)
# x1 = vcat(spt2.vz, spt2.vx, spt2.pz, spt2.px)
# y  = vec(shot1.d)
# y1 = vec(shot2.d)
# x'*x1 - y'*y1

# ==============================================================================
# using AcousticWave
# nz = 201; nx = 201; ext = 20; iflag = 1; dz = 10.; dx = 10.; dt = 2e-3; f0 = 30.0
# v  = 3000. * ones(nz, nx);
# vmax = maximum(v); vmin = minimum(v);
# fidMtx = InitFidMtx(nz, nx, ext, iflag, dz, dx, dt, vmax, vmin, f0, v);
# # initialize source term
# ix  = collect(1:2:201); iz = 101*ones(Int64, length(ix));
# pos = hcat(iz, ix)
# f0 = 30.0; ot = zeros(size(pos,1));
# src = InitSources(pos, f0, ot, dt);
# path = "/Users/wenlei/Desktop/src2spts.bin"
# src2SnapShots(path, src, nz, nx, ext, iflag);
#
# # position of receiver
# irz = []; irx = []
# for ix = 1 : nx
#     for iz = 1 : nz
#         if rand() <= 0.1
#            irz = push!(irz, iz)
#            irx = push!(irx, ix)
#         end
#     end
# end
# pos_rec = convert(Array{Int64,2}, hcat(irz, irx))
# shot  = MultiStepForward(pos_rec, path, fidMtx, tmax=1.0)
# path_wfd = "/Users/wenlei/Desktop/acoustic.bin"
# MultiStepForward(path_wfd, src , fidMtx, tmax=1.0)
# shot1 = fullWfd2Shot(pos_rec, path_wfd)

# ==============================================================================
# using AcousticWave
# nz = 201; nx = 201; ext = 20; iflag = 1; dz = 10.; dx = 10.; dt = 2e-3; f0 = 30.0
# v  = 3000. * ones(nz, nx);
# vmax = maximum(v); vmin = minimum(v);
# fidMtx = InitFidMtx(nz, nx, ext, iflag, dz, dx, dt, vmax, vmin, f0, v);
# # initialize source term
# isz = []; isx = []
# for ix = 1 : nx
#     for iz = 1 : nz
#         if rand() <= 0.1
#            isz = push!(isz, iz)
#            isx = push!(isx, ix)
#         end
#     end
# end
# pos_src = convert(Array{Int64,2}, hcat(isz, isx))
# f0 = 30.0; ot = zeros(size(pos_src,1));
# src = InitSources(pos_src, f0, ot, dt);
# tmax = 1; nt = round(Int64, tmax/dt+1);
# w = Ricker(f0, dt); ns = nt + 1 - length(w)
# for isrc = 1 : length(src)
#     src[isrc].nt= nt
#     src[isrc].p = conv(randn(ns),w)
# end
# path = "/Users/wenlei/Desktop/src2spts.bin"
# src2SnapShots(path, src, nz, nx, ext, iflag);
# # position of receiver
# irx = collect(1:2:nx); irz = 5*ones(Int64,length(irx));
# pos_rec = hcat(irz, irx)
# shot  = MultiStepForward(pos_rec, path, fidMtx, tmax=1.0)
#
#
# path_adj = "/Users/wenlei/Desktop/adjoint.bin"
# shot1 = InitShot(0, 0, pos_rec, 0.0, nt, dt)
# for irec = 1 : length(shot1.irz)
#     shot1.d[:,irec] = conv(randn(ns), w)
# end
# MultiStepAdjoint_test(path_adj, shot1, fidMtx)
#
# L = 0.0
# for ir = 1 : length(shot.irx)
#     L = L + dot(shot.d[:,ir], shot1.d[:,ir])
# end
#
# T = 0.0
# for it = 1 : nt
#     spt1 = readSnapShot(path    , it)
#     spt2 = readSnapShot(path_adj, it)
#     T = T + dot(spt1.vz, spt2.vz) + dot(spt1.vx, spt2.vx) + dot(spt1.pz, spt2.pz) + dot(spt1.px, spt2.px)
# end

# ==============================================================================
# using AcousticWave
# path = "/Users/wenlei/Desktop/acoustic.bin"
# nz = 201; nx = 201; ext = 20; iflag = 1; dz = 10.; dx = 10.; dt = 2e-3;
# v  = 3000. * ones(nz, nx);
# # initialize source term
# pos = [101 101]; f0 = 30.0; ot = [0.0];
# src = InitSources(pos, f0, ot, dt);
# # discretize spatial derivative operator
# vmax = maximum(v); vmin = minimum(v);
# fidMtx = InitFidMtx(nz, nx, ext, iflag, dz, dx, dt, vmax, vmin, f0, v);
# path = "/Users/wenlei/Desktop/acoustic.bin"
# MultiStepForward(path, src, fidMtx, tmax=1.0);
#
# using PyPlot, AcousticWave
# path = "/Users/wenlei/Desktop/acoustic.bin"
# pathout = "/Users/wenlei/Desktop/homo"
# waveAnim(path, pathout)
# sptPlot(path, 100)
# irx = collect(1:2:nx); irz = 5*ones(Int64,length(irx));
# pos_rec = hcat(irz, irx)
# shot = fullWfd2Shot(pos_rec, path)


using Test: @test, @testset

using ODE

@testset "ODE Tests" begin
    @testset "stepsize_adjust" begin
        (hnew, retry) = ODE.stepsize_adjust([1.0], [1e-8], 0.1, 1)
        @test isapprox(hnew, 0.1/(1e-8/(1e-8 + 1e-8))^0.5)
        @test retry == false

        (hnew, retry) = ODE.stepsize_adjust([2.234], [1e-3], 0.01, 3)
        @test isapprox(hnew, 0.5*0.01)
        @test retry == true

        (hnew, retry) = ODE.stepsize_adjust([2.234], [1e-12], 0.01, 3)
        @test isapprox(hnew, 2*0.01)
        @test retry == false
    end

    @testset "evolve tests" begin
        x0 = 0.0
        y0 = [1.0]
        x1 = 1.0

        function exp_rhs(x, y)
            return [-y[1]]
        end

        y1 = evolve_ode(Midpoint, exp_rhs, x0, y0, x1; atol=1e-4)

        @test abs(y1[1] - exp(-1)) < 1e-3
    end
end


using Test
using ReversibleSeismic
using NiLang

"""
the reversible loss
"""
@i function i_loss!(out::T, param, src, srcv::AbstractVector{T}, c::AbstractMatrix{T},
          tu::AbstractArray{T,3}, tφ::AbstractArray{T,3}, tψ::AbstractArray{T,3}) where T
     i_solve!(param, src, srcv, c, tu, tφ, tψ)
     for i=1:size(tu, 1)
        for j=1:size(tu, 2)
            out += tu[i,j,end] ^ 2
        end
     end
end

"""
obtain gradients with NiLang.AD
"""
function getgrad(c::AbstractMatrix{T}; nstep::Int) where T
     param = AcousticPropagatorParams(nx=size(c,1)-2, ny=size(c,2)-2,
          Rcoef=0.2, dx=20.0, dy=20.0, dt=0.05, nstep=nstep)

     c = copy(c)
     tu = zeros(T, size(c)..., nstep+1)
     tφ = zeros(T, size(c)..., nstep+1)
     tψ = zeros(T, size(c)..., nstep+1)

     src = size(c) .÷ 2 .- 1
     srcv = Ricker(param, 100.0, 500.0)
     NiLang.AD.gradient(Val(1), i_loss!, (0.0, param, src, srcv, c, tu, tφ, tψ))[end-3]
end

"""
obtain gradients numerically, for gradient checking.
"""
function getngrad(c::AbstractMatrix{T}, i, j; nstep::Int, δ=1e-4) where T
     param = AcousticPropagatorParams(nx=size(c,1)-2, ny=size(c,2)-2,
          Rcoef=0.2, dx=20.0, dy=20.0, dt=0.05, nstep=nstep)

     c = copy(c)
     tu = zeros(T, size(c)..., nstep+1)
     tφ = zeros(T, size(c)..., nstep+1)
     tψ = zeros(T, size(c)..., nstep+1)

     src = size(c) .÷ 2 .- 1
     srcv = Ricker(param, 100.0, 500.0)
     c[i,j] += δ
     fpos = i_loss!(0.0, param, src, srcv, copy(c), copy(tu), copy(tφ), copy(tψ))[1]
     c[i,j] -= 2δ
     fneg = i_loss!(0.0, param, src, srcv, copy(c), copy(tu), copy(tφ), copy(tψ))[1]
     return (fpos - fneg)/2δ
end

@testset "loss" begin
     nx = ny = 100
     nstep = 2000
     param = AcousticPropagatorParams(nx=nx, ny=ny,
          Rcoef=0.2, dx=20.0, dy=20.0, dt=0.05, nstep=nstep)

     tu = zeros(nx+2, ny+2, nstep+1)
     tφ = zeros(nx+2, ny+2, nstep+1)
     tψ = zeros(nx+2, ny+2, nstep+1)

     c = 1000*ones(nx+2, ny+2)
     src = (nx ÷ 2, ny ÷ 2)
     srcv = Ricker(param, 100.0, 500.0)

     loss = i_loss!(0.0, param, src, srcv, c, tu, tφ, tψ)[1]
     @test check_inv(i_loss!, (0.0, param, src, srcv, c, tu, tφ, tψ); atol=1e-6)
     #@test loss ≈ 10.931466822080788
     @test loss ≈ 6.234873084873294e-5
end

@testset "gradient" begin
     nx = ny = 100
     nstep = 2000
     c = 1000*ones(nx+2, ny+2)
     g4545 = getgrad(c; nstep=nstep)[45,45]
     ng4545 = getngrad(c, 45, 45; nstep=nstep, δ=1e-4)
     @test isapprox(g4545, ng4545; rtol=1e-2)
end