"""
    normalize!(A, hasIntercept)

Centers and normalizes the columns of A in place

# Arguments

- A = 2d array of floats
- hasIntercept = boolean flag indicating whether the first column of A is the
  intercept

# Value

Centers and normalizes A in place and returns 2d arrays of the column means and L2
norms of A before standardization.

"""
function normalize!(A::AbstractArray{Float64,2}, hasIntercept::Bool)

    # If including intercept, subtract the column means from all other columns
    if hasIntercept == true
        means = mean(A, dims=1)
        A[:,2:end] = A[:,2:end].-transpose(means[2:end])
    else # Otherwise, subtract the column means from all columns
        means = Array{Float64}(undef, 1, size(A,2))
    end

    # Obtain and divide by the column L2 norms
    norms = Array{Float64}(undef, 1, size(A,2))
    for j in 1:size(A,2)
        norms[j] = norm(A[:,j])
        A[:,j] = A[:,j]/norms[j]
    end

    # Return the column means and norms
    return means, norms
end

"""
    backtransform!(B, meansX, meansZ, normsX, normsZ, Y, Xold, Zold)

Back-transform coefficient estimates B in place if X and Z were standardized
prior to the estimation-- when both X and Z include intercept columns.

# Arguments

- B = 2d array of coefficient estimates B
- meansX = 2d array of column means of X, obtained prior to standardizing X
- meansZ = 2d array of column means of Z, obtained prior to standardizing Z
- normsX = 2d array of column norms of X, obtained prior to standardizing X
- normsZ = 2d array of column norms of Z, obtained prior to standardizing Z
- Y = 2d array of response matrix Y
- Xold = 2d array row covariates X prior to standardization
- Zold = 2d array column covariates Z prior to standardization

# Value

None; back-transforms B in place

"""
function backtransform!(B::AbstractArray{Float64,2},
                        meansX::AbstractArray{Float64,2},
                        meansZ::AbstractArray{Float64,2},
                        normsX::AbstractArray{Float64,2},
                        normsZ::AbstractArray{Float64,2},
                        Y::AbstractArray{Float64,2},
                        Xold::AbstractArray{Float64,2},
                        Zold::AbstractArray{Float64,2})

    # Back transform the X intercepts (row main effects)
    prodX = (meansX[:,2:end]./normsX[:,2:end])*B[2:end, 2:end]
    B[1,2:end] = (B[1,2:end]-vec(prodX))./vec(normsZ[:,2:end])/normsX[1,1]

    # Back transform the Z intercepts (column main effects)
    prodZ = B[2:end, 2:end]*transpose(meansZ[:,2:end]./normsZ[:,2:end])
    B[2:end,1] = (B[2:end,1]-prodZ)./transpose(normsX[:,2:end])/normsZ[1,1]

    # Back transform the interactions
    B[2:end, 2:end] = B[2:end, 2:end]./transpose(normsX[:,2:end])./
                                       normsZ[:,2:end]

    # Re-estimate intercept
    B[1,1] = 0
    B[1,1] = mean(Y-Xold*B*transpose(Zold))
end

"""
    backtransform!(B, hasXIntercept, hasZIntercept,
	               meansX, meansZ, normsX, normsZ)

Back-transform coefficient estimates B in place if X and Z were standardized
prior to the estimation-- when not including intercept columns for either X
or Z.

# Arguments

- B = 2d array of coefficient estimates B
- hasXIntercept = boolean flag indicating whether or not to X has an
  intercept column
- hasZIntercept = boolean flag indicating whether or not to Z has an
  intercept column
- meansX = 2d array of column means of X, obtained prior to standardizing X
- meansZ = 2d array of column means of Z, obtained prior to standardizing Z
- normsX = 2d array of column norms of X, obtained prior to standardizing X
- normsZ = 2d array of column norms of Z, obtained prior to standardizing Z

# Value

None; back-transforms B in place

"""
function backtransform!(B::AbstractArray{Float64,2},
                        hasXIntercept::Bool, hasZIntercept::Bool,
                        meansX::AbstractArray{Float64,2},
                        meansZ::AbstractArray{Float64,2},
                        normsX::AbstractArray{Float64,2},
                        normsZ::AbstractArray{Float64,2})

    # Back transform the X intercepts (row main effects), if necessary
    if hasXIntercept == true
        prodX = (meansX[:,2:end]./normsX[:,2:end])*B[2:end, 2:end]
        B[1,2:end] = (B[1,2:end]-vec(prodX))./vec(normsZ[:,2:end])/normsX[1,1]
    end

    # Back transform the Z intercepts (column main effects), if necessary
    if hasZIntercept == true
        prodZ = B[2:end, 2:end]*transpose(meansZ[:,2:end]./normsZ[:,2:end])
        B[2:end,1] = (B[2:end,1]-prodZ)./transpose(normsX[:,2:end])/
                                         normsZ[1,1]
    end

    # Back transform the interactions, if necessary
    if (hasXIntercept == true) || (hasZIntercept == true)
        B[2:end, 2:end] = B[2:end, 2:end]./transpose(normsX[:,2:end])./
                                           normsZ[:,2:end]
    end

    # Back transform the interactions if not including any main effects
    if (hasXIntercept == false) && (hasZIntercept == false)
        B = B./transpose(normsX)./normsZ
    end
end

"""
    backtransform!(B, meansX, meansZ, normsX, normsZ, Y, Xold, Zold)

Back-transform coefficient estimates B in place if X and Z were standardized
prior to the estimation-- when both X and Z include intercept columns.

# Arguments

- B = 4d array of coefficient estimates B
- meansX = 2d array of column means of X, obtained prior to standardizing X
- meansZ = 2d array of column means of Z, obtained prior to standardizing Z
- normsX = 2d array of column norms of X, obtained prior to standardizing X
- normsZ = 2d array of column norms of Z, obtained prior to standardizing Z
- Y = 2d array of response matrix Y
- Xold = 2d array row covariates X prior to standardization
- Zold = 2d array column covariates Z prior to standardization

# Value

None; back-transforms B in place

# Some notes

B is a 4d array in which each coefficient matrix is stored along the third and fourth
dimension.

"""
function backtransform!(B::AbstractArray{Float64,4},
                        meansX::AbstractArray{Float64,2},
                        meansZ::AbstractArray{Float64,2},
                        normsX::AbstractArray{Float64,2},
                        normsZ::AbstractArray{Float64,2},
                        Y::AbstractArray{Float64,2},
                        Xold::AbstractArray{Float64,2},
                        Zold::AbstractArray{Float64,2})

    # Iterate through the first dimension of B to back-transform each
	# coefficient matrix.
    for j in 1:size(B,4)
        for i in 1:size(B,3)
            # Back transform the X intercepts (row main effects)
            prodX = (meansX[:,2:end]./normsX[:,2:end])*B[2:end, 2:end,i,j]
            B[1,2:end,i,j] = (B[1,2:end,i,j]-vec(prodX))./vec(normsZ[:,2:end])/
                                                    normsX[1,1]

            # Back transform the Z intercepts (column main effects)
            prodZ = B[2:end,2:end,i,j]*transpose(meansZ[:,2:end]./normsZ[:,2:end])
            B[2:end,1,i,j] = (B[2:end,1,i,j]-prodZ)./transpose(normsX[:,2:end])/
                                                normsZ[1,1]

            # Back transform the interactions
            B[2:end,2:end,i,j] = B[2:end,2:end,i,j]./transpose(normsX[:,2:end])./
                                                normsZ[:,2:end]

            # Re-estimate intercept
            B[1,1,i,j] = 0
            B[1,1,i,j] = mean(Y-Xold*B[:,:,i,j]*transpose(Zold))
        end
    end
end

"""
    backtransform!(B, hasXIntercept, hasZIntercept,
                   meansX, meansZ, normsX, normsZ)

Back-transform coefficient estimates B in place if X and Z were standardized
prior to the estimation-- when not including intercept columns for either X
or Z.

# Arguments

- B = 4d array of coefficient estimates
- hasXIntercept = boolean flag indicating whether or not to X has an
  intercept column
- hasZIntercept = boolean flag indicating whether or not to Z has an
  intercept column
- meansX = 2d array of column means of X, obtained prior to standardizing X
- meansZ = 2d array of column means of Z, obtained prior to standardizing Z
- normsX = 2d array of column norms of X, obtained prior to standardizing X
- normsZ = 2d array of column norms of Z, obtained prior to standardizing Z

# Value

None; back-transforms B in place.

# Some notes

B is a 4d array in which each coefficient matrix is stored along the third and fourth
dimension.

"""
function backtransform!(B::AbstractArray{Float64,4},
                        hasXIntercept::Bool, hasZIntercept::Bool,
                        meansX::AbstractArray{Float64,2},
                        meansZ::AbstractArray{Float64,2},
                        normsX::AbstractArray{Float64,2},
                        normsZ::AbstractArray{Float64,2})

    # Iterate through the first dimension of B to back-transform each
	# coefficient matrix:
    for j in 1:size(B,4)
        for i in 1:size(B,3)
            # Back transform the X intercepts (row main effects), if necessary
            if hasXIntercept == true
                prodX = (meansX[:,2:end]./normsX[:,2:end])*B[2:end,2:end,i,j]
                B[1,2:end,i,j] = (B[1,2:end,i,j]-vec(prodX))./vec(normsZ[:,2:end])/
                                                        normsX[1,1]
            end

            # Back transform the Z intercepts (column main effects), if necessary
            if hasZIntercept == true
                prodZ = B[2:end,2:end,i,j]*transpose(meansZ[:,2:end]./
                    normsZ[:,2:end])
                B[2:end,1,i,j] = (B[2:end,1,i,j]-prodZ)./transpose(normsX[:,2:end])/
                                                    normsZ[1,1]
            end

            # Back transform the interactions, if necessary
            if (hasXIntercept == true) || (hasZIntercept == true)
                B[2:end,2:end,i,j] = B[2:end,2:end,i,j]./transpose(normsX[:,2:end])./
                                                    normsZ[:,2:end]
            end

            # Back transform the interactions if not including any main effects
            if (hasXIntercept == false) && (hasZIntercept == false)
                B[:,:,i,j] = B[:,:,i,j]./transpose(normsX)./normsZ
            end
        end
    end
end


"""
    kDPP(m::Int, kernel::Kernel)

k-DPP (Determinantal Point Process) will return a subset of `X` of size `m`,
according to DPP probability
"""
struct kDPP{K<:Kernel} <: OffIPSA
    m::Int
    kernel::K
    function kDPP(m::Int, kernel::K) where {K<:Kernel}
        m > 0 || throw(ArgumentError("The number of inducing points m should be positive"))
        return new{K}(m, kernel)
    end
end

function inducingpoints(rng::AbstractRNG, alg::kDPP, X::AbstractVector; kwargs...)
    if alg.m >= length(X)
        return edge_case(alg.m, length(X), X)
    end
    return kdpp_ip(rng, X, alg.m, alg.kernel)
end

Base.show(io::IO, ::kDPP) = print(io, "k-DPP selection of inducing points")

function kdpp_ip(rng::AbstractRNG, X::AbstractVector, m::Int, kernel::Kernel)
    return X[rand(rng, DPP(kernel, X)(m))] # Sample m indices from a DPP
end


boltzmann(β, E::Number) = exp(-β*E)

function boltzmann(β, E)
    factors = boltzmann.(β, E)
    partition = sum(factors)
    factors ./ partition
end

function boltzmann(β, E, deg)
    factors = map(*, boltzmann.(β, E), deg)
    partition = sum(factors)
    convert(Array{Float64}, factors ./ partition)
end

boltzmann(β, space::StateSpace) = boltzmann(β, space.energy, space.deg)


### A Pluto.jl notebook ###
# v0.14.0

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 3a4957ec-8723-11eb-22a0-8b35322596e2
html"""

Section 2.2

 Sampling and Random Variables

Lecture Video

"""

# ╔═╡ 06d2666a-8723-11eb-1395-0febdf3dc2a4
begin

	import Pkg

	Pkg.activate(mktempdir())

	Pkg.add([

			Pkg.PackageSpec(name="Images", version="0.22.4"),

			Pkg.PackageSpec(name="ImageMagick", version="0.7"),

			Pkg.PackageSpec(name="PlutoUI", version="0.7"),

			Pkg.PackageSpec(name="Plots"),

			Pkg.PackageSpec(name="Colors"),

			Pkg.PackageSpec(name="StatsBase"),

			Pkg.PackageSpec(name="Distributions")

			])

	using Plots, PlutoUI, Colors, Images, StatsBase, Distributions
	using Statistics

end

# ╔═╡ 0a70bca4-8723-11eb-1bcf-e9abb9b1ab75
PlutoUI.TableOfContents(aside=true)

# ╔═╡ 472a41d2-8724-11eb-31b3-0b81612f0083
md"""
## Julia: Useful tidbits
"""

# ╔═╡ aeb99f72-8725-11eb-2efd-d3e44686be03
md"""
The package that each function comes from is shown in brackets, unless it comes from `Base`.
"""

# ╔═╡ 4f9bd326-8724-11eb-2c9b-db1ac9464f1e
md"""

- Julia `Base` library (no `using` required):

  - `if...else...end`
  - `Dict`: Julia's dictionary type
  - `÷` or `div`:  integer division (type `\div` + <tab>)

  - `sum(S)`: sum of elements in the collection `S`, e.g. an array
  - `count(S)`: count the number of true elements of a Boolean collection

  - `rand(S)`: random sampling from a collection `S`

- `Statistics.jl` (pre-loaded standard library; just needs `using`)

  - `mean(S)`: calculate the mean of a collection `S`
  - `std(S)`:  calculate the standard deviation of a collection `S`

- `StatsBase.jl`:
  - `countmap`

- `Plots.jl`:

  - `histogram(x)`: Plot a histogram of data vector `x` [Plots]

  - `bar(d)`: Plot a bar graph of categorical data [Plots]

- `Colors.jl`:
  - `distinguishable_colors(n)`: Make `n` distinguishable colours [Colors]

"""

# ╔═╡ db2d25de-86b1-11eb-0c78-d1ee52e019ca
md"""
## Random sampling with `rand`
"""

# ╔═╡ e33fe4c8-86b1-11eb-1031-cf45717a3dc9
md"""
The `rand` function in Julia is quite versatile: it tries to generate, or **sample**, a random object from the argument that you pass in:
"""

# ╔═╡ f49191a2-86b1-11eb-3eab-b392ba058415
rand(1:6)

# ╔═╡ 1abda6c4-86b2-11eb-2aa3-4d1148bb52b7
rand([2, 3, 5, 7, 11])

# ╔═╡ 30b12f28-86b2-11eb-087b-8d50ec429b89
rand("MIT")

# ╔═╡ 4ce946c6-86b2-11eb-1820-0728798665ab
rand('a':'z')

# ╔═╡ fae3d138-8743-11eb-1014-b3a2a9b49aba
typeof('a':'z')

# ╔═╡ 6cdea3ae-86b2-11eb-107a-17bea3f54bc9
rand()   # random number between 0 and 1

# ╔═╡ 1c769d58-8744-11eb-3bd3-ab11ea1503ed
rand

# ╔═╡ 297fdfa0-8744-11eb-1934-9fe31e8be534
methods(rand);

# ╔═╡ 776ec3f2-86b3-11eb-0216-9b71d07e99f3
md"""
We can take random objects from a collection of objects of *any* type, for example:
"""

# ╔═╡ 5fcf8d4e-8744-11eb-080e-cba749004b08
distinguishable_colors(10)

# ╔═╡ 4898106a-8744-11eb-128a-35fec741e6b8
typeof(distinguishable_colors(10))

# ╔═╡ 0926366a-86b2-11eb-0f6d-31ae6981598c
rand(distinguishable_colors(3))   # from Colors.jl package

# ╔═╡ 7c8d7b72-86b2-11eb-2dd5-4f77bc5fb8ff
md"""
### Several random objects
"""

# ╔═╡ 2090b7f2-86b3-11eb-2a99-ed98800e1d63
md"""
To sample several random objects from the same collection, we could write an array comprehension:
"""

# ╔═╡ a7dff55c-86b2-11eb-330f-3d6279347095
[rand(1:6) for i in 1:10]

# ╔═╡ 2db33022-86b3-11eb-17dd-13c534ac9892
md"""
But in fact, just adding another argument to `rand` does the trick:
"""

# ╔═╡ 0de6f23e-86b2-11eb-39ff-318bbc4ecbcf
rand(1:6, 10)

# ╔═╡ 36c3da4a-86b3-11eb-0b2f-fffdde06fcd2
md"""
In fact, you can also generate not only random vectors, but also random matrices:
"""

# ╔═╡ 940c2bf6-86b2-11eb-0a5e-011abdd6352b
rand(1:6, 10, 12)

# ╔═╡ 5a4e7fc4-86b3-11eb-3376-0941b79574aa
rand(distinguishable_colors(5), 10, 10)

# ╔═╡ c433104e-86b3-11eb-20bb-af608bb281cc
md"""
We can also use random images:
"""

# ╔═╡ 78dc94e2-8723-11eb-1ff2-bb7104b62033
penny_image = load(download("https://www.usacoinbook.com/us-coins/lincoln-memorial-cent.jpg"))

# ╔═╡ bb1465c4-8723-11eb-1abc-bdb5a7028cf2
begin
	head = penny_image[:, 1:end÷2]
	tail = penny_image[:, end÷2:end]
end;

# ╔═╡ e04f3828-8723-11eb-3452-09f821391ad0
rand( [head, tail], 5, 5)

# ╔═╡ b7793f7a-8726-11eb-11d8-cd928f1a3645
md"""
## Uniform sampling
"""

# ╔═╡ ba80cc78-8726-11eb-2f33-e364f19295d8
md"""
So far, each use of `rand` has done **uniform sampling**, i.e. each possible object as the *same* probability.

Let's *count* heads and tails using the `countmap` function from the `StatsBase.jl` package:
"""

# ╔═╡ 8fe715e4-8727-11eb-2e7f-15b723bb8d9d
tosses = rand( ["head", "tail"], 10000)

# ╔═╡ 9da4e6f4-8727-11eb-08cb-d55e3bbff0e4
toss_counts = countmap(tosses)

# ╔═╡ 9647840c-8745-11eb-33fc-898ae351ddfe
"tail" => 3

# ╔═╡ 9d782aa6-8745-11eb-034e-0f8b01be987b
typeof("tail" => 3)

# ╔═╡ a693582c-8745-11eb-261b-ef79e503420e
toss_counts["tail"]

# ╔═╡ c3255d7a-8727-11eb-0421-d99faf27c7b0
md"""

We see that `countmap` returns a **dictionary** (`Dict`), which maps **keys** to **values**; we will say more about dictionaries in another chapter.
"""

# ╔═╡ e2037806-8727-11eb-3c36-1500f1dd545b
prob_tail = toss_counts["tail"] / length(tosses)

# ╔═╡ 28ff621c-8728-11eb-2ec0-2579cb313315
md"""
As we increase the number of tosses, we "expect" the probability to get closer to $1/2$.
"""

# ╔═╡ 57125768-8728-11eb-3c3b-1dd37e1ac189
md"""
## Tossing a weighted coin
"""

# ╔═╡ 6a439c78-8728-11eb-1969-27a19d254704
md"""
How could we model a coin that is **weighted**, so that it is more likely to come up heads? We want to assign a probability $p = 0.7$ to heads, and $q = 0.3$ to tails.
"""

# ╔═╡ 9f8ac08c-8728-11eb-10ad-f93ca225ce38
md"""
One way would be to generate random integers between 1 and 10 and assign heads to a subset of the possible results with the desired probability, e.g. 1:7 get heads, and 8:10 get tails:
"""

# ╔═╡ 062b400a-8729-11eb-16c5-235cef648edb
function simple_weighted_coin()
	outcome = if rand(1:10) ≤ 7
		"heads"
	else      # could have elseif
		"tails"
	end

	return outcome
end

# ╔═╡ 7c099606-8746-11eb-37fe-c3befde06e9d
function simple_weighted_coin2()
	if rand(1:10) ≤ 7
		"heads"
	else      # could have elseif
		"tails"
	end
end

# ╔═╡ 97cb3dde-8746-11eb-0b00-690f20cb26dc
result = for i in 1:10

end

# ╔═╡ 9d8cdc8e-8746-11eb-2b9a-b30a52026f09
result == nothing

# ╔═╡ 81a30c9e-8746-11eb-38c8-9be4f6ba2e80
simple_weighted_coin2()

# ╔═╡ 2a81ba88-8729-11eb-3dcb-1db26e468066
md"""
(Note that `if` statements have a return value in Julia.)
"""

# ╔═╡ 5ea5838a-8729-11eb-1749-030533fb0656
md"""
How could we generalise this to an arbitrary probability $p ∈ [0, 1]$?

We can generate a uniform floating-point number between 0 and 1, and check if it is less than $p$. This is called a **Bernoulli trial**.
"""

# ╔═╡ c9f21046-8746-11eb-27c6-910807240fd1
rand()

# ╔═╡ 806e6aae-8729-11eb-19ea-33722c60edf0
rand() < 0.314159

# ╔═╡ 90642564-8729-11eb-3cd7-b3d41a1553b4
md"""
Note that comparisons also return a value in Julia. Here we have switched from heads/tails to true/false as the output.
"""

# ╔═╡ 9a426a20-8729-11eb-0c0f-31e1d4dc91bc
md"""
Let's make that into a function:
"""

# ╔═╡ a4870b14-8729-11eb-20ee-e531d4a7108d
bernoulli(p) = rand() < p

# ╔═╡ 081e3796-8747-11eb-32ec-dfd998605737
bernoulli(0.7)

# ╔═╡ 008a40d2-872a-11eb-224d-5b3331f29c99
md"""
p = $(@bind p Slider(0.0:0.01:1.0, show_value=true, default=0.7))
"""

# ╔═╡ baed5908-8729-11eb-00e0-9f749406c30c
countmap( [bernoulli(p) for i in 1:1000] )

# ╔═╡ ed94eae8-86b3-11eb-3f1b-15c7a54903f5
md"""
## Bar charts and histograms
"""

# ╔═╡ f20504de-86b3-11eb-3125-3140e0e060b0
md"""
Once we have generated several random objects, it is natural to want to count **how many times** each one occurred.
"""

# ╔═╡ 7f1a3766-86b4-11eb-353f-b13acaf1503e
md"""
Let's roll a (fair) die 1000 times:
"""

# ╔═╡ 371838f8-86b4-11eb-1633-8d282e42a085
md"""
An obvious way to find the counts would be to run through the data looking for 1s, then run through again looking for 2s, etc.:
"""

# ╔═╡ 94688c1a-8747-11eb-13a3-eb36f731674c
rolls .== 1

# ╔═╡ ad701cdc-8747-11eb-3804-63a0fc881547
count(rolls .== 1)

# ╔═╡ 2405eb68-86b4-11eb-31b0-dff8e355d88e
counts = [count(rolls .== i) for i in 1:6]

# ╔═╡ 9e9d3556-86b5-11eb-3dfb-916e625da235
md"""
Note that this is *not* the most efficient algorithm!
"""

# ╔═╡ 90844738-8738-11eb-0604-3d23662152d9
md"""
We can plot **categorical data** using a **bar chart**, `bar` in Plots.jl. This counts each discrete item.
"""

# ╔═╡ 02d03642-86b4-11eb-365a-63ff61ddd3b5
rolls = rand(1:6, 100000)   # try modifying 100 by adding more zeros

# ╔═╡ 2d71fa88-86b5-11eb-0e55-35566c2246d7
begin
	bar(counts, alpha=0.5, leg=false, size=(500, 300))
	hline!([length(rolls) / 6], ls=:dash)
	title!("number of die rolls = $(length(rolls))")
	ylims!(0, length(rolls) / 3)
end

# ╔═╡ cb8a9762-86b1-11eb-0484-6b6cc8b1b14c
md"""
## Probability densities

### Rolling multiple dice
"""

# ╔═╡ d0c9814e-86b1-11eb-2f29-1d041bccc649
roll_dice(n) = sum( rand(1:12, n) )

# ╔═╡ b81b1090-8735-11eb-3a52-2dca4d4ed472
experiment() = roll_dice(n)

# experiment() = sum([randn()^2 for i in 1:n])

# ╔═╡ 7a16b674-86b7-11eb-3aa5-83712cdc8580
trials = 10^6

# ╔═╡ e8e811de-86b6-11eb-1cbf-6d4aeaee510a
data = [experiment() for t in 1:trials]

# ╔═╡ 2bfa712a-8738-11eb-3248-6f9bb93154e8
md"""
### Converging shape
"""

# ╔═╡ e4abcbf4-86b8-11eb-167a-d97c61e07837
data

# ╔═╡ 6c133ab6-86b7-11eb-15f6-7780da5afc31
md"""
n = $(@bind n Slider(1:50, show_value=true))
"""

# ╔═╡ 514f6be0-86b8-11eb-30c9-d1020f783afe
histogram(data, alpha=0.5, legend=false, bins=200, c=:lightsalmon1, title="n = $n")
# c = RGB(0.1, 0.2, 0.3))

# ╔═╡ a15fc456-8738-11eb-25bd-b15c2b16d461
md"""
Here we have switched from a bar chart to a **histogram**, which counts the number of items falling into a given range or **bin**. When $n$ is small this tends to look like a bar chart, but it looks like a "smooth bar chart" as $n$ gets larger, due to the **aggregation**.
"""

# ╔═╡ dd753568-8736-11eb-1f20-1b81110ae807
md"""
Does the above histogram look like a bell to you?
"""

# ╔═╡ 8ab9001a-8737-11eb-1009-5717fbe83af7
begin
	bell = load(download("https://encrypted-tbn0.gstatic.com/images?q=tbn:ANd9GcRmxRAIQt_L-X99A_4FoP3vsC-l_WHlC3TtAw&usqp=CAU"))

	bell[1:end*9÷10, :]
end

# ╔═╡ f8b2dd20-8737-11eb-1593-43659c693109
md"""
### Normalising the $y$-axis
"""

# ╔═╡ 89bb366a-8737-11eb-10a6-e754ee817f9a
md"""
Notice that the *shape* of the curve seems to converge to a bell-shaped curve, but the axes do not. What can we do about this?

We need to **normalise** so that the total shaded *area* is 1. We can use the histogram option `norm=true` to do so.
"""

# ╔═╡ 14f0a090-8737-11eb-0ccf-391249267401
histogram(data, alpha=0.5, legend=false, bins=50, norm=true,
			c=:lightsalmon1, title="n = $n", ylims=(0, 0.05))

# ╔═╡ e305467e-8738-11eb-1213-eb11aaebe151
md"""
### Normalising the $x$ axis
"""

# ╔═╡ e8341288-8738-11eb-27ae-0795fa7e4a7e
md"""
As we changed $n$ above, the range of values on the $x$ axis also changed. We want to find a way to rescale $x$ at the same time so that both axes and shape stop changing as we increase $n$.

We need to make sure that the data is centred in the same place -- we will choose 0. And we need to make sure that the width is the same -- we will divide by the standard deviation.
"""

# ╔═╡ 07e8ae34-873b-11eb-1df2-175392ac4678
σ = std(data)

# ╔═╡ 77616afe-873a-11eb-3f11-1bc417f53138
mean(data), std(data)

# ╔═╡ 16023bea-8739-11eb-1e32-79f2d006b093
normalised_data = ( data .- mean(data) ) ./ std(data)

# ╔═╡ bfc52118-873b-11eb-0fbb-952c60bc7cc2
histogram(normalised_data, bins=-5 - (1/(2σ)):(1/σ):5, norm=true,
	alpha=0.5, leg=false, ylim=(0, 0.41), size=(500, 300))

# ╔═╡ aa6126f8-873b-11eb-3b4a-0f96fe07b7fb
plot!(x -> exp(-x^2 / 2) / √(2π), lw=2, c=:red, alpha=0.5)

# ╔═╡ 308547c6-873d-11eb-3a42-833f8bf496ae
md"""
Note that in the limit, the data becomes *continuous*, no longer discrete. The probability of any particular value is 0 ! We then talk about a **probability density function**, $f(x)$; the integral of this function over the interval $[a, b]$ gives the probability of being in that interval.
"""

# ╔═╡ cb2fb68e-8749-11eb-29ea-9729ac0c63b4
md"""
### Options for plotting functions
"""

# ╔═╡ e0a1863e-8735-11eb-1182-1b3c59b1e05a
md"""
Other options for the `histogram` function options:
- `legend=false` turns off the legend (key), i.e. the plot labels
- `bins=50` specifies the number of **bins**; can also be a vector of bin edges
- `linetype=:stephist`: use steps instead of bars

- `alpha`: a general plot option specifying transparency (0=invisible; 1=opaque)
- `c` or `color`: a general plot option for the colour
- `lw`: line width (default=1)

There are several different ways to specify colours. [Here](http://juliagraphics.github.io/Colors.jl/stable/namedcolors/) is a list of named colours, but you can also specify `RGB(0.1, 0.2, 0.3)`.
"""

# ╔═╡ 900af0c8-86ba-11eb-2270-71b1869b9a1a
md"""
Note that `linetype=:stephist` will give a stepped version of the histogram:
"""

# ╔═╡ be6e4c00-873c-11eb-1413-5326aba54216
md"""
## Sampling from other distributions
"""

# ╔═╡ 9a1136c2-873c-11eb-124f-c3939972ce4a
md"""
dof = $(@bind dof Slider(1:50, show_value=true))
"""

# ╔═╡ e01b6f70-873c-11eb-04a1-ad8e86578982
chisq_data = rand( Chisq(dof), 100000 )

# ╔═╡ b5251f76-873c-11eb-38cb-7db300c8fe3c
histogram( chisq_data, norm=true, bins=100, size=(500, 300), leg=false, alpha=0.5,
	xlims=(0, 10*√(dof)))

# ╔═╡ da62fd1c-873c-11eb-0758-e7cb48e964f1
histogram( [ sum( randn().^2 for _=1:dof )  for _ = 1:100000], norm=true,
	alpha=0.5, leg=false)

# ╔═╡ Cell order:
# ╟─3a4957ec-8723-11eb-22a0-8b35322596e2
# ╠═06d2666a-8723-11eb-1395-0febdf3dc2a4
# ╠═0a70bca4-8723-11eb-1bcf-e9abb9b1ab75
# ╟─472a41d2-8724-11eb-31b3-0b81612f0083
# ╟─aeb99f72-8725-11eb-2efd-d3e44686be03
# ╟─4f9bd326-8724-11eb-2c9b-db1ac9464f1e
# ╟─db2d25de-86b1-11eb-0c78-d1ee52e019ca
# ╟─e33fe4c8-86b1-11eb-1031-cf45717a3dc9
# ╠═f49191a2-86b1-11eb-3eab-b392ba058415
# ╠═1abda6c4-86b2-11eb-2aa3-4d1148bb52b7
# ╠═30b12f28-86b2-11eb-087b-8d50ec429b89
# ╠═4ce946c6-86b2-11eb-1820-0728798665ab
# ╠═fae3d138-8743-11eb-1014-b3a2a9b49aba
# ╠═6cdea3ae-86b2-11eb-107a-17bea3f54bc9
# ╠═1c769d58-8744-11eb-3bd3-ab11ea1503ed
# ╠═297fdfa0-8744-11eb-1934-9fe31e8be534
# ╟─776ec3f2-86b3-11eb-0216-9b71d07e99f3
# ╠═5fcf8d4e-8744-11eb-080e-cba749004b08
# ╠═4898106a-8744-11eb-128a-35fec741e6b8
# ╠═0926366a-86b2-11eb-0f6d-31ae6981598c
# ╟─7c8d7b72-86b2-11eb-2dd5-4f77bc5fb8ff
# ╟─2090b7f2-86b3-11eb-2a99-ed98800e1d63
# ╠═a7dff55c-86b2-11eb-330f-3d6279347095
# ╟─2db33022-86b3-11eb-17dd-13c534ac9892
# ╠═0de6f23e-86b2-11eb-39ff-318bbc4ecbcf
# ╟─36c3da4a-86b3-11eb-0b2f-fffdde06fcd2
# ╠═940c2bf6-86b2-11eb-0a5e-011abdd6352b
# ╠═5a4e7fc4-86b3-11eb-3376-0941b79574aa
# ╟─c433104e-86b3-11eb-20bb-af608bb281cc
# ╠═78dc94e2-8723-11eb-1ff2-bb7104b62033
# ╠═bb1465c4-8723-11eb-1abc-bdb5a7028cf2
# ╠═e04f3828-8723-11eb-3452-09f821391ad0
# ╟─b7793f7a-8726-11eb-11d8-cd928f1a3645
# ╟─ba80cc78-8726-11eb-2f33-e364f19295d8
# ╠═8fe715e4-8727-11eb-2e7f-15b723bb8d9d
# ╠═9da4e6f4-8727-11eb-08cb-d55e3bbff0e4
# ╠═9647840c-8745-11eb-33fc-898ae351ddfe
# ╠═9d782aa6-8745-11eb-034e-0f8b01be987b
# ╠═a693582c-8745-11eb-261b-ef79e503420e
# ╟─c3255d7a-8727-11eb-0421-d99faf27c7b0
# ╠═e2037806-8727-11eb-3c36-1500f1dd545b
# ╟─28ff621c-8728-11eb-2ec0-2579cb313315
# ╟─57125768-8728-11eb-3c3b-1dd37e1ac189
# ╟─6a439c78-8728-11eb-1969-27a19d254704
# ╟─9f8ac08c-8728-11eb-10ad-f93ca225ce38
# ╠═062b400a-8729-11eb-16c5-235cef648edb
# ╠═7c099606-8746-11eb-37fe-c3befde06e9d
# ╠═97cb3dde-8746-11eb-0b00-690f20cb26dc
# ╠═9d8cdc8e-8746-11eb-2b9a-b30a52026f09
# ╠═81a30c9e-8746-11eb-38c8-9be4f6ba2e80
# ╟─2a81ba88-8729-11eb-3dcb-1db26e468066
# ╟─5ea5838a-8729-11eb-1749-030533fb0656
# ╠═c9f21046-8746-11eb-27c6-910807240fd1
# ╠═806e6aae-8729-11eb-19ea-33722c60edf0
# ╟─90642564-8729-11eb-3cd7-b3d41a1553b4
# ╟─9a426a20-8729-11eb-0c0f-31e1d4dc91bc
# ╠═a4870b14-8729-11eb-20ee-e531d4a7108d
# ╠═081e3796-8747-11eb-32ec-dfd998605737
# ╟─008a40d2-872a-11eb-224d-5b3331f29c99
# ╠═baed5908-8729-11eb-00e0-9f749406c30c
# ╟─ed94eae8-86b3-11eb-3f1b-15c7a54903f5
# ╟─f20504de-86b3-11eb-3125-3140e0e060b0
# ╟─7f1a3766-86b4-11eb-353f-b13acaf1503e
# ╟─371838f8-86b4-11eb-1633-8d282e42a085
# ╠═94688c1a-8747-11eb-13a3-eb36f731674c
# ╠═ad701cdc-8747-11eb-3804-63a0fc881547
# ╠═2405eb68-86b4-11eb-31b0-dff8e355d88e
# ╟─9e9d3556-86b5-11eb-3dfb-916e625da235
# ╟─90844738-8738-11eb-0604-3d23662152d9
# ╠═02d03642-86b4-11eb-365a-63ff61ddd3b5
# ╠═2d71fa88-86b5-11eb-0e55-35566c2246d7
# ╟─cb8a9762-86b1-11eb-0484-6b6cc8b1b14c
# ╠═d0c9814e-86b1-11eb-2f29-1d041bccc649
# ╠═b81b1090-8735-11eb-3a52-2dca4d4ed472
# ╠═7a16b674-86b7-11eb-3aa5-83712cdc8580
# ╠═e8e811de-86b6-11eb-1cbf-6d4aeaee510a
# ╟─2bfa712a-8738-11eb-3248-6f9bb93154e8
# ╠═e4abcbf4-86b8-11eb-167a-d97c61e07837
# ╟─6c133ab6-86b7-11eb-15f6-7780da5afc31
# ╠═514f6be0-86b8-11eb-30c9-d1020f783afe
# ╟─a15fc456-8738-11eb-25bd-b15c2b16d461
# ╟─dd753568-8736-11eb-1f20-1b81110ae807
# ╠═8ab9001a-8737-11eb-1009-5717fbe83af7
# ╟─f8b2dd20-8737-11eb-1593-43659c693109
# ╟─89bb366a-8737-11eb-10a6-e754ee817f9a
# ╠═14f0a090-8737-11eb-0ccf-391249267401
# ╟─e305467e-8738-11eb-1213-eb11aaebe151
# ╟─e8341288-8738-11eb-27ae-0795fa7e4a7e
# ╠═07e8ae34-873b-11eb-1df2-175392ac4678
# ╠═77616afe-873a-11eb-3f11-1bc417f53138
# ╠═16023bea-8739-11eb-1e32-79f2d006b093
# ╠═bfc52118-873b-11eb-0fbb-952c60bc7cc2
# ╠═aa6126f8-873b-11eb-3b4a-0f96fe07b7fb
# ╟─308547c6-873d-11eb-3a42-833f8bf496ae
# ╟─cb2fb68e-8749-11eb-29ea-9729ac0c63b4
# ╟─e0a1863e-8735-11eb-1182-1b3c59b1e05a
# ╟─900af0c8-86ba-11eb-2270-71b1869b9a1a
# ╟─be6e4c00-873c-11eb-1413-5326aba54216
# ╟─9a1136c2-873c-11eb-124f-c3939972ce4a
# ╠═e01b6f70-873c-11eb-04a1-ad8e86578982
# ╠═b5251f76-873c-11eb-38cb-7db300c8fe3c
# ╠═da62fd1c-873c-11eb-0758-e7cb48e964f1


__precompile__()

module AnimalBreedingTools

using Printf
using Statistics
using LinearAlgebra
using Base.Threads
using StatsBase
using Jacobi

struct AnimalBreedingDataSetRaw
   pedlist::Matrix{Int}
   genid::Vector{Int}
   G::Matrix{Float64}
   y::Vector{Float64}
   X::Matrix{Float64}
end

export get_design_matrix, get_interaction_matrix, directsum, solve_pcg, approximate_trace_of_inverse
export v2r, c2r, r2v, r2c, vech, bending, bending2
export AnimalBreedingDataSetRaw, load_data_set
export chol_symtrid, ldlt_symtrid, takahashi_ldlt_symtrid!
export normalized_legendre, normalized_legendre_matrix, hrv_class
export disp

include("designmatrix.jl")
include("solver.jl")
include("tools.jl")
include("data.jl")
include("polynomials.jl")
include("io.jl")

end


"""
Module for performing a Doppler-constrained PCA analysis

For algorithm information, see:
  - Jones, Stenning, Ford et al. https://arxiv.org/abs/1711.01318
  - See also: Gilbertson, Ford, Jones & Stenning 2020 https://arxiv.org/abs/2009.01085

Author: Eric Ford
Date:   September 2020
"""
module DCPCA

using Statistics, MultivariateStats

#import .RvSpectML : speed_of_light_mps
const speed_of_light_mps = 299792458.0 # TODO: Update value

export doppler_constrained_pca, compute_spectra_perp_doppler_shift

function compute_spectra_perp_doppler_shift(spectra::AA, deriv::V1, rvs::V2) where {
            T1<:Real, AA<:AbstractArray{T1,2}, T2<:Real, V1<:AbstractVector{T2}, T3<:Real, V2<:AbstractVector{T3} }
   @assert size(spectra,1) == length(deriv)
   @assert size(spectra,2) == length(rvs)
   fm_perp = spectra .- rvs' .* deriv./speed_of_light_mps
end

function doppler_constrained_pca(flux::AbstractArray{T1,2}, deriv::AbstractVector{T2}, rvs::AbstractVector{T3} )  where { T1<:Real, T2<:Real, T3<:Real  }#, times::AbstractVector, rvs::AbstractVector)
  rv_shift = rvs .- mean(rvs)
  fm_perp = compute_spectra_perp_doppler_shift(flux,deriv, rv_shift )
  idx_good_obs = 1:length(rvs)
  M = fit(PCA, fm_perp[:,idx_good_obs]; maxoutdim=10)
  pca_out = MultivariateStats.transform(M,fm_perp[:,idx_good_obs])
  return pca_out, M
end

end # DCPCA


using MPI
using StaticArrays
using ClimateMachine
using ClimateMachine.VariableTemplates
using ClimateMachine.Mesh.Topologies
using ClimateMachine.Mesh.Grids
using ClimateMachine.MPIStateArrays
using ClimateMachine.DGMethods

using ClimateMachine.BalanceLaws

import ClimateMachine.BalanceLaws: vars_state, nodal_init_state_auxiliary!

using ClimateMachine.Mesh.Geometry: LocalGeometry

struct GradTestModel{dim, dir} <: BalanceLaw end

vars_state(m::GradTestModel, ::Auxiliary, T) = @vars begin
    a::T
    ∇a::SVector{3, T}
    ∇a_exact::SVector{3, T}
end
vars_state(::GradTestModel, ::Prognostic, T) = @vars()

function nodal_init_state_auxiliary!(
    ::GradTestModel{dim, dir},
    aux::Vars,
    tmp::Vars,
    g::LocalGeometry,
) where {dim, dir}
    x, y, z = g.coord
    if dim == 2
        aux.a = x^2 + y^3 - x * y
        if dir isa EveryDirection
            aux.∇a_exact = SVector(2 * x - y, 3 * y^2 - x, 0)
        elseif dir isa HorizontalDirection
            aux.∇a_exact = SVector(2 * x - y, 0, 0)
        elseif dir isa VerticalDirection
            aux.∇a_exact = SVector(0, 3 * y^2 - x, 0)
        end
    else
        aux.a = x^2 + y^3 + z^2 * y^2 - x * y * z
        if dir isa EveryDirection
            aux.∇a_exact = SVector(
                2 * x - y * z,
                3 * y^2 + 2 * z^2 * y - x * z,
                2 * z * y^2 - x * y,
            )
        elseif dir isa HorizontalDirection
            aux.∇a_exact =
                SVector(2 * x - y * z, 3 * y^2 + 2 * z^2 * y - x * z, 0)
        elseif dir isa VerticalDirection
            aux.∇a_exact = SVector(0, 0, 2 * z * y^2 - x * y)
        end
    end
end

using Test
function test_run(mpicomm, dim, direction, Ne, N, FT, ArrayType)

    brickrange = ntuple(j -> range(FT(0); length = Ne[j] + 1, stop = 3), dim)
    topl = StackedBrickTopology(
        mpicomm,
        brickrange,
        periodicity = ntuple(j -> true, dim),
    )

    grid = DiscontinuousSpectralElementGrid(
        topl,
        FloatType = FT,
        DeviceArray = ArrayType,
        polynomialorder = N,
    )

    model = GradTestModel{dim, direction}()
    dg = DGModel(
        model,
        grid,
        nothing,
        nothing,
        nothing;
        state_gradient_flux = nothing,
    )

    continuous_field_gradient!(
        model,
        dg.state_auxiliary,
        (:∇a,),
        dg.state_auxiliary,
        (:a,),
        grid,
        direction,
    )

    # Wrapping in Array ensure both GPU and CPU code use same approx
    @test Array(dg.state_auxiliary.∇a) ≈ Array(dg.state_auxiliary.∇a_exact)
end

let
    ClimateMachine.init()
    ArrayType = ClimateMachine.array_type()

    mpicomm = MPI.COMM_WORLD

    numelem = (5, 5, 1)
    lvls = 1
    polynomialorder = 4

    @testset for FT in (Float64, Float32)
        @testset for dim in 2:3
            @testset for direction in (
                EveryDirection(),
                HorizontalDirection(),
                VerticalDirection(),
            )
                for l in 1:lvls
                    test_run(
                        mpicomm,
                        dim,
                        direction,
                        ntuple(j -> 2^(l - 1) * numelem[j], dim),
                        polynomialorder,
                        FT,
                        ArrayType,
                    )
                end
            end
        end
    end
end

nothing


include("common.jl")

# bkamins: this function mostly reuses what is below; to be discussed
function calibrate(rs::RideShare, initial_solution::Vector{Vector{Int}})
    temperature = 1000.0
    iterations_per_temp = 1000
    cooling = 0.6
    accept_prob = 1.0
    current_solution = deepcopy(initial_solution)
    i, accepted, total, k = 0, 0.0, 0, 1

    current_score = rideshare_optimality(rs, initial_solution)

    # bkamins: I have change changed the logic of calculation of accept_prob
    # to make it more consistent
    while accept_prob > 0.6
        if i > iterations_per_temp
            accept_prob = total == 0 ? 1.0 : accepted / total
            temperature *= cooling
            accepted, total, i = 0.0, 0, 0
        end
        s_new = iterate_solution(current_solution)
        score_new = rideshare_optimality(rs, s_new)
        if score_new > current_score
            current_score = score_new
            current_solution = s_new
        else
            delta_score = score_new - current_score
            prob = 1 - exp(delta_score / (k * temperature))
            accepted += prob
            total = total + 1
            if prob < rand()
                current_score = score_new
                current_solution = s_new
            end
        end
        i = i + 1
    end
    println("Calibrated Temperature $temperature")
    temperature
end

function simulated_annealing(rs::RideShare;
                             final_temp=0.001, max_iterations=1000,
                             max_accept=2, alpha=0.90, threshold=1000)
    progress = Float64[]
    temp_progress = Float64[]
    current_progress = Float64[]
    solution = get_initial_solution(rs)
    current_temp = calibrate(rs, solution)
    current = rideshare_optimality(rs, solution)
    best = current
    best_schedule = solution
    i, accept, k = 0, 0.0, 1

    println("Initial Solution:")
    foreach(((x, s),) -> println("Driver #$x : $s"), enumerate(solution))
    println("Initial revenue: $current")
    waitbar = Progress(max_iterations, 1)
    while current_temp > final_temp && i < max_iterations && current < threshold
        if accept >= max_accept
            accept = 0
            current_temp *= alpha
        end
        next_schedule = iterate_solution(solution)
        next = rideshare_optimality(rs,next_schedule)
        if next > current
            accept += 1
            solution, current = next_schedule, next
        else
            delta = next - current
            if exp(delta / (k * current_temp)) > rand()
                accept += 1
                solution, current = next_schedule, next
            end
        end
        if current > best
            best, best_schedule = current, solution
        end
        i += 1
        push!(progress, best)
        push!(temp_progress, current_temp)
        push!(current_progress, current)
        next!(waitbar)
    end

    println("Iterations: $i")
    println("Best Solution")
    foreach(((x, s),) -> println("Driver #$x : $s"), enumerate(best_schedule))
    println("Best revenue: $best")
    return (current_progress, best_schedule)
end

using Random
Random.seed!(1234)

res = simulated_annealing(rs_samples[1]; max_iterations=10000, alpha=0.95)
res = simulated_annealing(rs_samples[2]; max_iterations=10000, alpha=0.95)


include("./exportData.jl")
include("./fenkar_ode.jl")
using .exportData, .fenkar_ode, DifferentialEquations, DiffEqParamEstim, PyPlot, BlackBoxOptim, Optim, DelimitedFiles

#solve and take data
prob = ODEProblem(fenkar!, u0, tspan, p)
sol = solve(prob, Tsit5(); saveat=tspan[begin]:5.0:tspan[end])

# getExpData
ind = 78
t, data, p[15] = getExpData(ind, tInds=1:500)
p[14] = 0.0 #t[findfirst(data[1,:] .> 0.5)]-p[15]/2

# remake problem
prob = remake(prob, tspan=(0.0, t[end]), p=p)

# use the save_idxs to select the mid-point of the u-variable
sol = solve(prob, Tsit5(); saveat=t, save_idxs=1:1)

fig = plt.figure(constrained_layout=true)
plt.plot(t, data[:], "-", label="Exp. Sample")
plt.plot(sol.t, sol[1,:], "--", label="Initial Parameters")
plt.xlabel("\$ t \$")

# define the cost function
cost_function = build_loss_objective(prob, Tsit5(), L2Loss(t,data), maxiters=10000, verbose_opt=true, verbose_steps=50; save_idxs=1:1);

# define limits for parameters
#          tsi,  tv1m,   tv2m,   tvp,   twm,    twp,  td,    to,    tr,   xk,  uc,    uv, ucsi,        t0,     TI,   IA
#   p = [ 29.0,  19.6, 1250.0,   3.33, 41.0,  870.0, 0.25, 12.5,  33.3, 10.0, 0.13, 0.04,  0.85,      1.0,   BCL, 0.05]
lower = [  1.0,  10.0,  400.0,   2.33, 10.0,  100.0, 0.10, 10.0,  10.0,  9.0, 0.10, 0.01,  0.10,      0.0,   40.0, 0.00]
upper = [ 50.0,  50.0, 2000.0,   4.33, 50.0, 1000.0, 0.50, 50.0,  50.0, 11.0, 0.90, 0.90,  0.90,    p[15], 1000.0, 1.00]

@assert all(p .>= lower) && all(p .<= upper)

bound = [(lower[n], upper[n]) for n in 1:length(p)]

# fit
result = bboptimize(cost_function; SearchRange = bound, MaxSteps = 1e5)#, NThreads=Threads.nthreads()-1)
@info "Initial p = $(p) \n"
@info "Minimizer = $(result.archive_output.best_candidate) \n"
prob = remake(prob, p=result.archive_output.best_candidate)
sol = solve(prob, Tsit5(); saveat=t, save_idxs=1:1)
plt.plot(sol.t, sol[1,:], "--", label="BBO Parameters")

result = optimize(cost_function, result.archive_output.best_candidate, NelderMead())
@info "Initial p = $(p) \n"
@info "Minimizer = $(result.minimizer) \n"
prob = remake(prob, p=result.minimizer)
sol = solve(prob, Tsit5(); saveat=t, save_idxs=1:1)
plt.plot(sol.t, sol[1,:], "--", label="Nelder-Mead Parameters")

plt.savefig("./ode/bbo/$(ind)_fittings.pdf")
plt.close()

open("./ode/bbo/$(ind)_params.txt", "w") do io
	writedlm(io, ["tsi" "tv1m" "tv2m" "tvp" "twm" "twp" "td" "to" "tr" "xk" "uc" "uv" "ucsi" "t0" "TI" "IA"])
	writedlm(io, "\n")
	writedlm(io, transpose(p))
	writedlm(io, "\n")
	#writedlm(io, transpose(result.archive_output.best_candidate))
	writedlm(io, transpose(result.minimizer))
end


"""
    relief(data::Array{<:Real,2}, target::Array{<:Integer,1}, m::Signed=-1,
                dist_func::Any=(e1, e2) -> sum(abs.(e1 .- e2), dims=2); f_type::String="continuous")::Array{Float64,1}

Compute feature weights using Relief algorithm. The f_type argument specifies whether the features are continuous or discrete
and can either have the value of "continuous" or "discrete".

---
# Reference:
- Kenji Kira and Larry A. Rendell. The feature selection problem: Tra-
ditional methods and a new algorithm. In Proceedings of the Tenth
National Conference on Artificial Intelligence, AAAI’92, pages 129–134.
AAAI Press, 1992.
"""
function relief(data::Array{<:Real,2}, target::Array{<:Integer,1}, m::Signed=-1,
                dist_func::Any=(e1, e2) -> sum(abs.(e1 .- e2), dims=2); f_type::String="continuous")::Array{Float64,1}

    # Initialize feature weights vector.
    weights = zeros(Float64, size(data, 2))

    # Compute vectors of maximum and minimum feature values.
    max_f_vals = vec(maximum(data, dims=1))
    min_f_vals = vec(minimum(data, dims=1))

    # Sample m samples without replacement. If m has signal value -1, use all samples.
    sample_idxs = StatsBase.sample(1:size(data, 1), if (m==-1) size(data,1) else m end, replace=false)
    if (m == -1) m = size(data, 1) end # If m has signal value -1, set m to total number of examples.

    # Compute pairwise distances between samples (vector form).
    # Note that the number of elements in the distance vector is {n \choose 2} = n!/(2!*(n-2)!) = n*(n-1)/2.
    dists = Array{Float64}(undef, Int64((size(data, 1)^2 - size(data, 1))/2))

    # Construct pairwise distances vector using vectorized distance function.
    top_ptr = 1
    @inbounds for idx = 1:size(data,1)-1
        upper_lim = top_ptr + size(data, 1) - idx - 1
        dists[top_ptr:upper_lim] = dist_func(data[idx:idx, :], data[idx+1:end, :])
        top_ptr = upper_lim + 1
    end

    # Go over sampled indices.
    @inbounds for idx = sample_idxs

        # Row and column indices for querying pairwise distance vector.
        # Use zero-based numbering.
        row_idxs = repeat([idx - 1], size(data, 1))
        col_idxs = collect(0:size(data, 1)-1)

        # Get indices of neighbors with same class in distances vector and find nearest hit.
        sel_neigh_same = (target .== target[idx]) .& (1:length(target) .!= idx)
        neigh_idx_hit = square_to_vec(row_idxs[sel_neigh_same],
                                        col_idxs[sel_neigh_same], size(data, 1)) .+ 1
        idx_nearest_hit = argmin(dists[neigh_idx_hit])
        nearest_hit = vec(data[sel_neigh_same, :][idx_nearest_hit, :])

        # Get indices of neighbors with different class in distances vector and find nearest miss.
        sel_neigh_miss = target .!= target[idx]
        neigh_idx_miss = square_to_vec(row_idxs[sel_neigh_miss], col_idxs[sel_neigh_miss], size(data, 1)) .+ 1
        idx_nearest_miss = argmin(dists[neigh_idx_miss])
        nearest_miss = vec(data[sel_neigh_miss, :][idx_nearest_miss, :])

        ### Weights Update ###

        if f_type == "continuous"
            # If features continuous.
            weights = weights .- (abs.(data[idx, :] .- nearest_hit)./(max_f_vals .- min_f_vals .+ eps(Float64)))./m .+
                (abs.(data[idx, :] .- nearest_miss)./(max_f_vals .- min_f_vals .+ eps(Float64)))./m
        elseif f_type == "discrete"
            # If features discrete.
            weights = weights .- Int64.(data[idx, :] .!= nearest_hit)./m .+ Int64.(data[idx, :] .!= nearest_miss)./m
        else
            throw(DomainError(f_type, "f_type can only be equal to \"continuous\" or \"discrete\"."))
        end

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

    end

    # Return computed feature weights.
    return weights
end


using Indicators
using Temporal
using Test
using Random
using Statistics

const global N = 1_000
const global X0 = 50.0
const global SEED = 1

@testset "Utilities" begin
    Random.seed!(SEED)
    @testset "Array" begin
        x = cumsum(randn(N)) .+ X0
        y = x + randn(N)
        cxo = crossover(x, y)
        cxu = crossunder(x, y)
        @test any(cxo)
        @test any(cxu)
        @test !any(cxo .* cxu)  # ensure crossovers and crossunders never coincide
    end
end

# trendy
@testset "Trendlines" begin
    Random.seed!(SEED)
    @testset "Array" begin
        x = cumsum(randn(N))
        tmp = resistance(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = support(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = minima(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = maxima(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
    end
end

# analytical
@testset "Chaos" begin
    Random.seed!(SEED)
    @testset "Array" begin
        x = randn(252)
        # helpers
        a, b = Indicators.divide(x)
        @test [a; b] == x
        x = randn(101)
        a, b = Indicators.divide(x)
        @test [a; b] == x
        # workhorses
        h = hurst(x, n=100)
        @test size(h) == size(x)
        rs = rsrange(x)
        @test size(rs) == size(x)
        x = randn(100)
    end
    @testset "Temporal" begin
        # chaos indicators
        x = TS(randn(N))
        tmp = hurst(x)
        @test size(tmp,1) == size(x,1)
        @test size(tmp,2) == 1
        tmp = rsrange(x)
        @test size(tmp,1) == size(x,1)
        @test size(tmp,2) == 1
    end
end

# moving regressions
@testset "Regressions" begin
    Random.seed!(SEED)
    @testset "Array" begin
        x = cumsum(randn(N))
        tmp = mlr_beta(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 2
        @test sum(isnan.(tmp)) != N
        tmp = mlr_slope(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = mlr_intercept(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = mlr(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = mlr_se(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = mlr_ub(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = mlr_lb(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = mlr_bands(tmp)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 3
        @test sum(isnan.(tmp)) != N
        tmp = mlr_rsq(x, adjusted=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = mlr_rsq(x, adjusted=false)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
    end
    @testset "Temporal" begin
        x = TS(cumsum(randn(N)))
        # moving regressions
        tmp = mlr_beta(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 2
        tmp = mlr_slope(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = mlr_intercept(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = mlr(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = mlr_se(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = mlr_ub(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = mlr_lb(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = mlr_bands(tmp)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 3
        tmp = mlr_rsq(x, adjusted=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = mlr_rsq(x, adjusted=false)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
    end
end

@testset "Running Calculations" begin
    Random.seed!(SEED)
    @testset "Array" begin
        x = cumsum(randn(N))
        X = cumsum(randn(N, 4), dims=1)
        tmp = diffn(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = diffn(X)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == size(X,2)
        @test sum(isnan.(tmp)) != N
        tmp = runmean(x, cumulative=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = runmean(x, cumulative=false)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = runsum(x, cumulative=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = wilder_sum(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = runmad(x, cumulative=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = runmad(x, cumulative=false)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = runvar(x, cumulative=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = runvar(x, cumulative=false)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = runsd(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = runcov(x, x.*rand(N), cumulative=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = runcov(x, x.*rand(N), cumulative=false)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = runcor(x, x.*rand(N), cumulative=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = runcor(x, x.*rand(N), cumulative=false)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = runmin(x, cumulative=true, inclusive=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = runmin(x, cumulative=true, inclusive=false)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = runmin(x, cumulative=false, inclusive=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = runmin(x, cumulative=false, inclusive=false)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = runmax(x, cumulative=true, inclusive=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = runmax(x, cumulative=true, inclusive=false)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = runmax(x, cumulative=false, inclusive=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = runmax(x, cumulative=false, inclusive=false)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = mode(map(xi->round(xi), x))
        @test size(tmp, 1) == 1
        @test size(tmp, 2) == 1
        @test !isnan(tmp)
        tmp = runquantile(x, cumulative=true)
        @test !isnan(tmp[2]) && isnan(tmp[1])
        @test tmp[10] == quantile(x[1:10], 0.05)
        tmp = runquantile(x, cumulative=false)
        @test tmp[10] == quantile(x[1:10], 0.05)
        n = 20
        tmp = runacf(x, n=n, maxlag=15, cumulative=true)
        @test all(tmp[n:end,1] .== 1.0)
        tmp = runacf(x, n=n, maxlag=15, cumulative=false)
        @test all(tmp[n:end,1] .== 1.0)
    end
    @testset "Temporal" begin
        # running calculations
        x = TS(cumsum(randn(N)))
        X = TS(cumsum(randn(N, 4), dims=1))
        tmp = runmean(x, cumulative=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = runmean(x, cumulative=false)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = runsum(x, cumulative=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = wilder_sum(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = runmad(x, cumulative=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = runmad(x, cumulative=false)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = runvar(x, cumulative=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = runvar(x, cumulative=false)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = runsd(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = runcov(X[:,1], X[:,4], cumulative=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = runcov(X[:,1], X[:,4], cumulative=false)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = runcor(X[:,1], X[:,4], cumulative=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = runcor(X[:,1], X[:,4], cumulative=false)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = runmin(x, cumulative=true, inclusive=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = runmin(x, cumulative=true, inclusive=false)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = runmin(x, cumulative=false, inclusive=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = runmin(x, cumulative=false, inclusive=false)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = runmax(x, cumulative=true, inclusive=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = runmax(x, cumulative=true, inclusive=false)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = runmax(x, cumulative=false, inclusive=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = runmax(x, cumulative=false, inclusive=false)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = runquantile(x, cumulative=true)
        @test !isnan(tmp.values[2,1]) && isnan(tmp.values[1,1])
        @test tmp.values[10,1] == quantile(x.values[1:10,1], 0.05)
        tmp = runquantile(x, cumulative=false)
        @test tmp.values[10,1] == quantile(x.values[1:10,1], 0.05)
        n = 20
        tmp = runacf(x, n=n, maxlag=15, cumulative=true)
        @test all(tmp.values[n:end,1] .== 1.0)
        tmp = runacf(x, n=n, maxlag=15, cumulative=false)
        @test all(tmp.values[n:end,1] .== 1.0)
    end
end

# moving average functions
@testset "Moving Averages" begin
    Random.seed!(SEED)
    @testset "Array" begin
        x = cumsum(randn(N))
        tmp = sma(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = mama(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 2
        @test sum(isnan.(tmp)) != N
        tmp = ema(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = wma(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = hma(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = trima(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = mma(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = tema(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = dema(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = swma(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = kama(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = alma(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = zlema(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
    end
    @testset "Temporal" begin
        x = TS(cumsum(randn(N)))
        # moving average functions
        tmp = sma(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = mama(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 2
        tmp = ema(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = wma(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = hma(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = trima(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = mma(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = tema(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = dema(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = swma(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = kama(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = alma(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = zlema(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
    end
end

# momentum function
@testset "Momentum" begin
    Random.seed!(SEED)
    @testset "Array" begin
        x = cumsum(randn(N))     # close
        Y = cumsum(randn(N, 2), dims=1)  # high-low
        Z = cumsum(randn(N, 3), dims=1)  # high-low-close
        tmp = aroon(Y)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 3
        @test sum(isnan.(tmp)) != N
        tmp = donch(Y)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 3
        @test sum(isnan.(tmp)) != N
        tmp = momentum(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = roc(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = macd(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 3
        @test sum(isnan.(tmp)) != N
        tmp = rsi(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = adx(Z)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 3
        @test sum(isnan.(tmp)) != N
        tmp = adx(Z, wilder=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 3
        @test sum(isnan.(tmp)) != N
        tmp = psar(Y)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = kst(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = wpr(Z)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = cci(Z)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = stoch(Z, kind=:fast)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 2
        @test sum(isnan.(tmp)) != N
        tmp = stoch(Z, kind=:slow)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 2
        @test sum(isnan.(tmp)) != N
        tmp = smi(Z)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 2
        @test sum(isnan.(tmp)) != N
    end
    @testset "Temporal" begin
        x = TS(cumsum(randn(N)))
        Y = TS(cumsum(randn(N, 2), dims=1))
        Z = TS(cumsum(randn(N, 3), dims=1))
        # momentum function
        tmp = aroon(Y)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 3
        tmp = donch(Y)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 3
        tmp = momentum(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = roc(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = macd(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 3
        tmp = rsi(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = adx(Z)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 3
        tmp = adx(Z, wilder=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 3
        tmp = psar(Y)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = kst(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = wpr(Z)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = cci(Z)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = stoch(Z, kind=:fast)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 2
        tmp = stoch(Z, kind=:slow)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 2
        tmp = smi(Z)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 2
    end
end

# volatility functions
@testset "Volatility" begin
    Random.seed!(SEED)
    @testset "Array" begin
        x = cumsum(randn(N))
        X = cumsum(randn(N, 3), dims=1)
        tmp = bbands(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 3
        @test sum(isnan.(tmp)) != N
        tmp = tr(X)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = atr(X)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = keltner(X)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 3
        @test sum(isnan.(tmp)) != N
    end
    @testset "Temporal" begin
        x = TS(cumsum(randn(N)))
        X = TS(cumsum(randn(N, 3), dims=1))
        tmp = bbands(x)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 3
        tmp = tr(X)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = atr(X)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        tmp = keltner(X)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 3
    end
end

# chart patterns functions
@testset "Charting" begin
    Random.seed!(SEED)
    @testset "Array" begin
        X = cumsum(randn(N, 3), dims=1)
        tmp = renko(X, use_atr=true)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
        tmp = renko(X, use_atr=false)
        @test size(tmp, 1) == N
        @test size(tmp, 2) == 1
        @test sum(isnan.(tmp)) != N
    end
end


import Pkg
Pkg.activate("..")

using CrimsonSkyline
using Distributions: Beta, Geometric, Binomial, truncated
using StatsBase: mean, std

function bayes_net(t::Trace, data::Int64)
    prob = sample(t, "prob", Beta(2.0, 2.0))
    n = observe(t, "n", truncated(Geometric(prob), 1, Inf), data)
    sample(t, "count", Binomial(n, prob))
end

function main()
    @warn "~~~ generating kernels for closed-world inference (model is identifiable) ~~~"
    data = 7
    kernels = make_kernels(bayes_net; params = (data,))
    results = inference(
        bayes_net, Metropolis();
        params = (data,),
        inference_params = Dict(
            "kernels" => kernels,
            "num_iterations" => 55000,
            "burn" => 10000,
            "thin" => 500,
            "inverse_verbosity" => 5000
        )
    )
    prob_stats = map(f -> f(results, "prob"), [mean, std])
    count_stats = map(f -> f(results, "count"), [mean, std])
    @info "Posterior stats: prob (mean, std) = $prob_stats, count (mean, std) = $count_stats"
end

if abspath(PROGRAM_FILE) == @__FILE__
    main()
end


function markov_blanket(n,m₀,graph)
    sg = size(graph,1)
    mb1::Array{UInt16,1} = m₀
    mb2::Array{UInt16,1} = findall(graph[n,:])

    mb12 = intersect(mb1,mb2,sg)
    if length(mb12) == 0
        mb = bfs(n,graph)
        if length(mb) > 0
            mb2 = union(mb2,mb,sg)
        end
    end

    mb2 = [1;n;mb2]
    union(mb1,mb2,sg)
end

function Base.intersect(a::AbstractArray{UInt16,1},b::AbstractArray{UInt16,1},s::Int)
    x = zeros(Int,s)
    y = zeros(Int,s)
    x[a] .= 1
    y[b] .= 1
    findall((x .* y) .== 1)
end

function Base.union(a::AbstractArray{UInt16,1},b::AbstractArray{UInt16,1},s::Int)
    x = zeros(Int,s)
    y = zeros(Int,s)
    x[a] .= 1
    y[b] .= 1
    findall((x .+ y) .> 0)
end

function bfs(n,graph)
    sg = size(graph,1)
    d = -ones(Int,sg)
    d[1] = 0
    q = zeros(Int,sg)
    i = 1
    j = 2
    q[i] = 1
    found = false
    p = zeros(Int,0)
    while  !found && i < j
        curr = q[i]
        i += 1
        next = findall(graph[n,:] .* (d.==-1))
        s1 = length(next)
        q[j:j+s1-1] = next
        j += s1
        d[next] .= curr
        found=d[n]==curr
    end

    if found
        p = curr
        prev = d[curr]
        while prev ≠ 1
            p = [p prev]
            prev = d[prev]
        end
    end
    return p
end


export
    default_grid_size, plot_function_around, plot_local_slices

"""
    $SIGNATURES

Find integers `w` and `h` such that `w * h ≥ n` and `abs2(w/h-aspect_ratio)` is
minimized. Return `(w, h)`.

Useful for determining grid size for `n` subplots.

The algorithm uses a heuristic, making an initial guess as the closest integer
to `√(aspect_ratio/r)`, and minimizing the objective within the nearest `± d`
integers.
"""
function default_grid_size(n; aspect_ratio = 4//3, d = 2)
    hs = filter(h -> h ∈ 1:n, @. round(√(n/aspect_ratio) + (-d:d)))
    ws = cld.(n, hs)
    _, ix = findmin(@. abs2((ws / hs) - aspect_ratio))
    ws[ix], hs[ix]
end

"""
    $SIGNATURES

Return `PGFPlotsX.Plot` elements visualizing the function ``f`` on ``x ± Δ``,
using `n` points. For plotting directly, use

```julia
Axis(plot_function_around(f, x, Δ)...)
```
"""
function plot_function_around(f, x::Real, Δ::Real; n = 100)
    xs = linspace(x - Δ, x + Δ, n)
    @pgf (Plot(Table(xs, f.(xs))),
          Plot({ only_marks }, Table([x], [f(x)])))
end

"""
    $SIGNATURES

Plot slices of the ``ℝⁿ→ℝ`` around `x ± Δs`. Return a `GroupPlot`, with subplots
arranged using `default_grid_size` using `aspect_ratio`. `xlabelfun` generates
labels for coordinates from the coordinate index, while `ylabel` provides the
`y` label.
"""
function plot_local_slices(f, x::AbstractVector, Δs::AbstractVector;
                           aspect_ratio = 4//3,
                           grid_size = default_grid_size(length(x)),
                           xlabelfun = i -> "\$x_{$(i)}\$",
                           ylabel = raw"$f$")
    @argcheck length(x) == length(Δs)
    w, h = grid_size
    plots = []
    for (i, Δ) in enumerate(Δs)
        push!(plots, @pgf { xlabel = xlabelfun(i) })
        append!(plots, plot_function_around(function_slice(f, x, i),
                                            zero(eltype(x)), Δ))
    end
    @pgf GroupPlot({ group_style = { group_size = "$w by $h"},
                     ylabel = ylabel},
                   plots...)
end


"""
    AbstractSpecialPolynomial{T,X}

An abstract type to distinguish the different polynomial types in this package.

The concrete types specify different bases for the space of polynomials of degree `n` or less.

This package includes:

* several classic orthogonal polynomials.
* Newton and Lagrange interpolating polynomials
* Bernstein polynomials

As many of the methods for the base `Polynomials` class are directly coded if possible, but quite a few
depend on conversion to the base `Polynomial` type (which uses the standard polynomial basis).

"""
abstract type AbstractSpecialPolynomial{T,X} <: Polynomials.AbstractPolynomial{T,X} end

# polynomial like Vector{T} with variable
Base.eltype(::Type{<:AbstractSpecialPolynomial{T}}) where {T} = T
Base.eltype(::Type{<:AbstractSpecialPolynomial}) = Float64

# to strip off type parameters. See Polynomials.constructorof
# We want ⟒(P{α,T,N}) = P{α} where {T, N}x
@generated function constructorof(::Type{T}) where {T}
    getfield(parentmodule(T), nameof(T))
end
⟒(P::Type{<:AbstractSpecialPolynomial}) = constructorof(P)
⟒(::Type{<:Polynomial}) = Polynomial

## Display

function unicode_subscript(io, j)
    a = ("⁻", "", "", "₀", "₁", "₂", "₃", "₄", "₅", "₆", "₇", "₈", "₉")
    for i in string(j)
        print(io, a[Int(i) - 44])
    end
end

function Polynomials.showterm(
    io::IO,
    ::Type{P},
    pj::T,
    var,
    j,
    first::Bool,
    mimetype,
) where {N,T,P<:AbstractSpecialPolynomial}
    iszero(pj) && return false
    !first && print(io, " ")

    if Polynomials.hasneg(T) && Polynomials.isneg(pj)
        print(io, "- ")
        pj = -pj
    else
        !first && print(io, "+ ")
    end

    Polynomials.printcoefficient(io, pj, 1, mimetype) # j != 0; want to have ()
    print(io, "⋅")
    print(io, basis_symbol(P))

    unicode_subscript(io, j)
    print(io, "($(var))")
    return true
end

# Conversion is *possible* when multiplication  in `P` is defined
# but otherwise,  it must be handled specially
# function Base.convert(::Type{P}, q::Q) where {P <: AbstractSpecialPolynomial, Q <: AbstractSpecialPolynomial}
#     convert(P, convert(Polynomial, q))
# end

# return the domain as a tuple, not an interval object
function Base.extrema(::Type{P}) where {P<:AbstractSpecialPolynomial}
    dom = domain(P)
    first(dom), last(dom)
end
Base.extrema(p::P) where {P<:AbstractSpecialPolynomial} = extrema(P)

## Polynomial operations

##
## --------------------------------------------------
##

function Polynomials.derivative(p::P, order::Integer=1) where {P<:AbstractSpecialPolynomial}
    T = eltype(one(P))
    q = convert(Polynomial{T}, p)
    convert(⟒(P), derivative(q, order))
end

function Polynomials.integrate(p::P) where {P<:AbstractSpecialPolynomial}
    q = convert(Polynomial, p)
    integrate(q)
end

function Base.divrem(num::P, den::P) where {P<:AbstractSpecialPolynomial}
    p1 = convert(Polynomial, num)
    p2 = convert(Polynomial, den)
    q, r = divrem(p1, p2)
    convert.(P, (q, r))
end

function Polynomials.companion(p::P) where {P<:AbstractSpecialPolynomial}
    companion(convert(Polynomial, p))
end

function Polynomials.vander(
    ::Type{P},
    x::AbstractVector{T},
    n::Integer,
) where {P<:AbstractSpecialPolynomial,T}
    N = length(x) - 1
    R = typeof(one(T) / one(T))
    V = zeros(R, N + 1, n + 1)

    for j in 0:n
        bj = Polynomials.basis(P, j)
        for i in 0:N
            V[i + 1, j + 1] = bj(x[i + 1])
        end
    end

    V
end


# MIT License

# Copyright (c) 2021 Fran Urbano
# Copyright (c) 2021 Val Lyashov
# Copyright (c) 2020 Bryan Palmer

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""
Package: Forecast

    hma(x, n)

Applies the Henderson moving average filter to dataset `x` with `n`-term.

"Henderson moving averages are filters which were derived by Robert Henderson in 1916
for use in actuarial applications. They are trend filters, commonly used in time series
analysis to smooth seasonally adjusted estimates in order to generate a trend estimate.

They are used in preference to simpler moving averages because they can reproduce
polynomials of up to degree 3, thereby capturing trend turning points.

The ABS uses Henderson moving averages to produce trend estimates from a seasonally
adjusted series. The trend estimates published by the ABS are typically derived using
a 13 term Henderson filter for monthly series, and a 7 term Henderson filter for quarterly series.

Henderson filters can be either symmetric or asymmetric. Symmetric moving averages can be applied
at points which are sufficiently far away from the ends of a time series. In this case, the
smoothed value for a given point in the time series is calculated from an equal number of values
on either side of the data point." - Australian Bureau of Statistics (www.abs.gov.au)

# Arguments
- `x`: Observations' support.
- `n`: Observation values, tipically 13 or 7 for quarterly data.

# Returns
An array of Henderson filter smoothed values provided in `s`.

# Examples
```julia-repl
julia> hma(rand(1000), 303)
1000-element Vector{Float64}}:
[...]
```
"""
@inline function hma(x::AbstractVector{<:T}, n::Integer) where T<:Real

    if !isodd(n)
        @warn "Henderson moving averages only accepts odds values for 'n';\nreturning instead the average of hma(x,"*string(n)*"-1) and hma(x,"*string(n)*"+1)"

        hma1 = hma(x,n-1)
        hma2 = hma(x,n+1)
        return (hma1 + hma2) / 2
    end

    @assert n >= 5 "n must be greater or equal 5"
    @assert (lx = length(x)) >= n "length(s) must be greater or equal to n"

    # Promotion to Float type
    P = Base.promote_op(/,T,T)
    x,_ = Base.promote(collect(x),[P(1.0)])

    w = hmaSymmetricWeights(n,P)

    m = (n-1) ÷ 2

    ic = n < 13 ? P(1.0) : (13 <= n < 15 ? P(3.5) : P(4.5))
    b2s2 = P(4.0)/P(pi)/ic^2

    function hmai(i)
        if i - 1 < m
            u = hmaAsymmetricWeights(m + i, w, b2s2)[end:-1:1]
            sum(x[1:i + m] .* u)
        elseif i - 1 + m >= lx
            u = hmaAsymmetricWeights(m + lx - i + 1, w, b2s2)
            sum(x[i-m:lx] .* u)
        else
            sum(x[i-m:i+m] .* w)
        end
    end

    hmav = similar(x)
    for i in 1:lx
        @inbounds hmav[i] = hmai(i)
    end
    hmav

end

"""
package: Smoothers

    hmaSymmetricWeights(n,P)

Caluclate the hma symmetric weights for 'n'

# Arguments
- `n`: Number of symmetric weights

# Returns

 Vector of symmetrical hma weights with type related to the bit size of `n`

# Refenrences
- "A Guide to Interpreting Time Series" ABS (2003), page 41.
"""
@inline function hmaSymmetricWeights(n::Integer,P::DataType)

    m = (P(n)-1)÷2
    m1 = (m+1)^2
    m2 = (m+2)^2
    m3 = (m+3)^2
    logd = log(m+2)+log(m2-1)+log(4*m2-1)+log(4*m2-9)+log(4*m2-25)
    d = 315/(8*exp(logd))

    w = Vector{P}(undef,Int(m)+2)
    for (i,v) in enumerate(0:m+1)
         @inbounds w[i] = real(exp(log(m1-v^2)+log(m2-v^2)+log(m3-v^2)+
                               log(complex(P(3.0)*m2-P(11.0)*v^2-P(16.0)))+log(d)))
    end

    u = vcat(w[end-1:-1:1],w[2:end-1])

    return mod(n, 2) != 0 ? u : vcat(u, missing)

end

"""
package: Smoothers

    hmaAsymmetricWeights(m,w)

Calculate the hma asymmetric end-weights for 'm' given 'w'

# Arguments
- `m`: Number of asymmetric weights (m < length(w))
- `w`: Vector of symmetrical hma weights

# Returns

Vector of asymmetrical hma weights

# References

- Mike Doherty (2001), "The Surrogate Henderson Filters in X-11",Aust, NZ J of Stat. 43(4), 2001, 385–392
"""
@inline function hmaAsymmetricWeights(m::Integer, w::AbstractVector{<:T}, b2s2::T) where T<:Real

    n = length(w)
    # @assert m <= n "The m argument must be less than w"
    # @assert m >= (n-1)÷2 "The m argument must be greater than (n-1)/2"

    sumResidual = sum(w[range(m + 1, n, step = 1)]) / T(m)

    sumEnd = T(0.0)
    m12 = (m+T(1.0))/T(2.0)
    for i in m+1:n
        @inbounds sumEnd += (i-m12)*w[i]
    end

    denominator = T(1.0) + ((m*(m-T(1.0))*(m+T(1.0)) / T(12.0) ) * b2s2)

    aw = Vector{T}(undef,m)
    for r in 1:m
        numerator = (T(r) - (m+T(1.0)) / T(2.0)) * b2s2
        @inbounds aw[r] = w[r] + sumResidual +  numerator / denominator * sumEnd
    end
    aw
end


##################################################
## SymbolicObject types have field x::PyCall.PyObject

## Symbol class for controlling dispatch
abstract type SymbolicObject <: Number end
struct Sym <: SymbolicObject
    __pyobject__::PyCall.PyObject
end

## Matrices
## We use this class for `ImmutableMatrices`
## Mutable matrices are mapped to `AbstractArray{Sym,N}`
## cf. matrix.jl
"""
    SymMatrix

Type to store a SymPy matrix, as created by `sympy.ImmutableMatrix`.

These have 0-based indexing defined for them to match SymPy

The traditional infix mathmatical operations are defined, but no dot broadcasting.

The `convert(Matrix{Sym}, M)` call is useful to covert to a Julia matrix

"""
mutable struct SymMatrix <: SymbolicObject
    __pyobject__::PyCall.PyObject
end

## Permutations
## A permutation of {0, 1, 2, ..., n} -- 0-based
struct SymPermutation <: SymbolicObject
  __pyobject__::PyCall.PyObject
end
export SymPermutation
Base.convert(::Type{SymPermutation}, o::PyCall.PyObject) = SymPermutation(o)

## A permutation of {0, 1, 2, ..., n} -- 0-based
struct SymPermutationGroup <: SymbolicObject
  __pyobject__::PyCall.PyObject
end
export SymPermutationGroup
Base.convert(::Type{SymPermutationGroup}, o::PyCall.PyObject) = SymPermutationGroup(o)

# a Lambda function
struct Lambda <: SymbolicObject
    __pyobject__::PyCall.PyObject
end
Lambda(args, expression) = Lambda(sympy.Lambda(args, expression).__pyobject__)
(λ::Lambda)(args...; kwargs...) = λ.__pyobject__(args...; kwargs...)
export Lambda

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

## important override
## this allows most things to flow though PyCall
PyCall.PyObject(x::SymbolicObject) = x.__pyobject__
## Override this so that using symbols as keys in a dict works
hash(x::SymbolicObject) = hash(PyObject(x))
==(x::SymbolicObject, y::SymbolicObject) = PyObject(x) == PyObject(y)

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

## Show methods
"create basic printed output"
function jprint(x::SymbolicObject)
    out = PyCall.pycall(pybuiltin("str"), String, PyObject(x))
    out = replace(out, r"\*\*" => "^")
    out
end
jprint(x::AbstractArray) = map(jprint, x)

## text/plain
Base.show(io::IO, s::Sym) = print(io, jprint(s))
Base.show(io::IO, ::MIME"text/plain", s::SymbolicObject) =  print(io, sympy.pretty(s))

## latex enhancements: Sym, array, Dict
#Base.show(io::IO, ::MIME"text/latex", x::SymbolicObject) = print(io, sympy.latex(x, mode="equation*"))

as_markdown(x) = Markdown.parse("``$x``")
function Base.show(io::IO, ::MIME"text/latex", x::SymbolicObject)
    print(io, sympy.latex(x, mode="inline",fold_short_frac=false)) # plain? equation*?
#    #print(io, as_markdown(sympy.latex(x, mode="equation*")))
#    #print(io, as_markdown(sympy.latex(x, mode="plain",fold_short_frac=false))) # inline?
end
function  show(io::IO, ::MIME"text/latex", x::AbstractArray{Sym})
    function toeqnarray(x::Vector{Sym})
        a = join([sympy.latex(x[i]) for i in 1:length(x)], "\\\\")
        """\\left[ \\begin{array}{r}$a\\end{array} \\right]"""
#        "\\begin{bmatrix}$a\\end{bmatrix}"
    end
    function toeqnarray(x::AbstractArray{Sym,2})
        sz = size(x)
        a = join([join(map(sympy.latex, x[i,:]), "&") for i in 1:sz[1]], "\\\\")
        "\\left[ \\begin{array}{" * repeat("r",sz[2]) * "}" * a * "\\end{array}\\right]"
#        "\\begin{bmatrix}$a\\end{bmatrix}"
    end
    print(io, as_markdown(toeqnarray(x)))
end
function show(io::IO, ::MIME"text/latex", d::Dict{T,S}) where {T<:SymbolicObject, S<:Any}
    Latex(x::Sym) = sympy.latex(x)
    Latex(x) = sprint(io -> show(IOContext(io, :compact => true), x))

    out = "\\begin{equation*}\\begin{cases}"
    for (k,v) in d
        out = out * Latex(k) * " & \\text{=>} &" * Latex(v) * "\\\\"
    end
    out = out * "\\end{cases}\\end{equation*}"
    print(io, as_markdown(out))
end

latex(x::Sym) = sympy.latex(x)

# (this is used by Polynomials.jl)
function Base.show_unquoted(io::IO, pj::SymbolicObject, indent::Int, prec::Int)
    if Base.operator_precedence(:+) <= prec
        print(io, "(")
        show(io, pj)
        print(io, ")")
    else
        show(io, pj)
    end
end

## Following recent changes to PyCall where:
# For o::PyObject, make o["foo"], o[:foo], and o.foo equivalent to o.foo in Python,
# with the former returning an raw PyObject and the latter giving the PyAny
# conversion.
# We do something similar to SymPy
#
# We only implement for symbols here, not strings
function Base.getproperty(o::T, s::Symbol) where {T <: SymbolicObject}
    if (s in fieldnames(T))
        getfield(o, s)
    else
        getproperty(PyCall.PyObject(o), s)
    end
end

# XXX Needs version v1.2+
#function Base.hasproperty(o::T, s::Symbol) where {T <: SymbolicObject}
#    s ∈ fieldnames(T) && return true
#    hasproperty(PyCall.PyObject(o), s)
#end


# Before running this, please make sure to activate and instantiate the environment
# corresponding to [this `Project.toml`](https://raw.githubusercontent.com/alan-turing-institute/DataScienceTutorials.jl/master/Project.toml) and [this `Manifest.toml`](https://raw.githubusercontent.com/alan-turing-institute/DataScienceTutorials.jl/master/Manifest.toml)
# so that you get an environment which matches the one used to generate the tutorials:
#
# ```julia
# cd("DataScienceTutorials") # cd to folder with the *.toml
# using Pkg; Pkg.activate("."); Pkg.instantiate()
# ```

using MLJ
using StableRNGs
import DataFrames

@load RidgeRegressor pkg=MultivariateStats

rng = StableRNG(6616) # for reproducibility
x1 = rand(rng, 300)
x2 = rand(rng, 300)
x3 = rand(rng, 300)
y = exp.(x1 - x2 -2x3 + 0.1*rand(rng, 300))
X = DataFrames.DataFrame(x1=x1, x2=x2, x3=x3)

test, train = partition(eachindex(y), 0.8);

Xs = source(X)
ys = source(y)

std_model = Standardizer()
stand = machine(std_model, Xs)
W = MLJ.transform(stand, Xs)

box_model = UnivariateBoxCoxTransformer()
box_mach = machine(box_model, ys)
z = MLJ.transform(box_mach, ys)

ridge_model = RidgeRegressor(lambda=0.1)
ridge = machine(ridge_model, W, z)
ẑ = predict(ridge, W)

ŷ = inverse_transform(box_mach, ẑ)

surrogate = Deterministic()
mach = machine(surrogate, Xs, ys; predict=ŷ)

fit!(mach)
predict(mach, X[test[1:5], :])

@from_network mach begin
    mutable struct CompositeModel
        regressor=ridge_model
    end
end

cm = machine(CompositeModel(), X, y)
res = evaluate!(cm, resampling=Holdout(fraction_train=0.8, rng=51),
                measure=rms)
round(res.measurement[1], sigdigits=3)

mutable struct CompositeModel2 <: DeterministicNetwork
    std_model::Standardizer
    box_model::UnivariateBoxCoxTransformer
    ridge_model::RidgeRegressor
end

function MLJ.fit(m::CompositeModel2, verbosity::Int, X, y)
    Xs = source(X)
    ys = source(y)
    W = MLJ.transform(machine(m.std_model, Xs), Xs)
    box = machine(m.box_model, ys)
    z = MLJ.transform(box, ys)
    ẑ = predict(machine(m.ridge_model, W, z), W)
    ŷ = inverse_transform(box, ẑ)
    mach = machine(Deterministic(), Xs, ys; predict=ŷ)
    fit!(mach, verbosity=verbosity - 1)
    return mach()
end

mdl = CompositeModel2(Standardizer(), UnivariateBoxCoxTransformer(),
                      RidgeRegressor(lambda=0.1))
cm = machine(mdl, X, y)
res = evaluate!(cm, resampling=Holdout(fraction_train=0.8), measure=rms)
round(res.measurement[1], sigdigits=3)

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl


@testset "modes" begin

@testset "center_mode" begin
    offsets = ((-1, -1), (0, -1), (1, -1), (-1, 0), (1, 0), (-1, 1), (0, 1), (1, 1))
    X = [6  7  9
         5  6  2
         2  1  6]
    @test 6 == center_mode(X, CartesianIndex(2, 2), offsets)

    iX = IntegralArray(X)
    @test 6 == center_mode(iX, CartesianIndex(2, 2), CartesianIndex(0, 0), offsets)
    @test mean(X[2:3,2:3]) == center_mode(iX, CartesianIndex(2, 2), CartesianIndex(1, 1), offsets)
end

@testset "average_mode" begin
    offsets = ((-1, -1), (0, -1), (1, -1), (-1, 0), (1, 0), (-1, 1), (0, 1), (1, 1))
    X = [6  7  9
         5  6  2
         2  1  6]
    @test mean(X) ≈ average_mode(X, CartesianIndex(2, 2), offsets)

    iX = IntegralArray(X)
    @test mean(X) ≈ average_mode(iX, CartesianIndex(2, 2), CartesianIndex(0, 0), offsets)

    ref = (mean(X[1:2, 1:2]) + mean(X[2:3, 1:2]) + mean(X[3, 1:2]) +
          mean(X[1:2, 2:3]) + mean(X[2:3, 2:3]) + mean(X[3, 2:3]) +
          mean(X[1:2, 3  ]) + mean(X[2:3, 3  ]) + mean(X[3, 3]))/9
    @test ref ≈ average_mode(iX, CartesianIndex(2, 2), CartesianIndex(1, 1), offsets)
end

@testset "LBP" begin
    lbp = LBP((1, 1))
    @test lbp.mode == center_mode
    @test lbp.block_size == (1, 1)

    @test_throws MethodError LBP()
    @test_throws ArgumentError LBP((0, 3))
end

end


# utility functions

function get_first_nz_idx(x::Array{T}, mask = nothing) where {T<:Real}
    first_nz_idx = 0
    for i in 1:length(x)
        mask != nothing && !mask[i] && continue
        first_nz_idx = i
        abs(x[i]) > eps() && break
    end
    first_nz_idx
end

"""Compute elementary symmetric polynomials for given Λ and k.
"""
function elem_symm_poly(Λ::Array{Float64}, k::Int)
    N = length(Λ)
    poly = zeros(k + 1, N + 1)
    poly[1, :] .= 1
    for l in (1:k) .+ 1
        for n in (1:N) .+ 1
            poly[l, n] = poly[l, n - 1] + Λ[n - 1] * poly[l - 1, n - 1]
        end
    end
    poly
end


# f1 = 3*x1-cos(x2*x3)-0.5
# f2 = x1^2-81*(x2+0.1)^2+sin(x3)+1.06
# f3 = e^(-x1*x2)+20*x3+(10*π-3)/3

x1 = 0.01
x2 = 0.01
x3 = -0.01
f1 = 3*x1-cos(x2*x3)-0.5
f2 = x1^2-81*(x2+0.1)^2+sin(x3)+1.06
f3 = e^(-x1*x2)+20*x3+(10*π-3)/3

x10 = x1
x20 = x2
x30 = x3

dx1 = 0.01
dx2 = 0.01
dx3 = 0.01
x1 = x1 + dx1
x2 = x2 + dx2
x3 = x3 + dx3

while (abs(f1)>1e-6)|(abs(f2)>1e-6)|(abs(f3)>1e-6)
  f10 = f1
  f20 = f2
  f30 = f3

  f1 = 3*x1-cos(x2*x3)-0.5
  f2 = x1^2-81*(x2+0.1)^2+sin(x3)+1.06
  f3 = e^(-x1*x2)+20*x3+(10*π-3)/3

  x11 = x1
  x22 = x2
  x33 = x3

  x1 = x1 + -f1/(f1-f10)
  x2 = x2 + -f1/(f1-f10)
  x3 = x3 + -f1/(f1-f10)

  x10 = x11
  x20 = x22
  x30 = x33
end

print("x1 is : $(x1)\n")
print("x2 is : $(x2)\n")
print("x3 is : $(x3)\n")


module _Inplace

using Brainstorm._DataStructure: PreAllocatedStack
using Brainstorm._Math._GMP: set!, add!, mul!, div!, neg!, pow!, isqrt!

include("binarysplitting.jl")

const c3_over_24 = Int64(640320)^3 ÷ 24
const digits_per_term = log10(c3_over_24 / (6 * 2 * 6))
const guard_terms = 5

preallocatedstack(size) = PreAllocatedStack(0, [(big(0), big(0), big(0)) for i in 1:size])

# Compute int(pi * 10^digits)
# This is done using Chudnovsky's series with binary splitting
pi(digits::Integer) = begin
    # how many terms to compute
    n = ceil(Int64, digits / digits_per_term) + guard_terms

    depth = 1
    while (1 << depth) <= n; depth += 1; end
    s = preallocatedstack(depth + 1)

    # calculate P(0, N), Q(0, N) and T(0, N)
    binarysplitting!(s, zero(n), n)
    p, q, t = s[1]

    sqrtC = p
    set!(sqrtC, 10)
    pow!(sqrtC, 2*digits) # one_squared
    mul!(sqrtC, 10005)
    isqrt!(sqrtC)

    mul!(q, 426880)
    mul!(q, sqrtC)
    div!(q, t)
end

#
# Source: http://www.craig-wood.com/nick/articles/pi-chudnovsky/
#
# Computes the terms for binary splitting the Chudnovsky infinite series
#
# a(a) = +/- (13591409 + 545140134*a)
# p(a) = (6*a-5)*(2*a-1)*(6*a-1)
# b(a) = 1
# q(a) = a*a*a*C3_OVER_24
#
compute!(s, a, b) = begin
    p, q, t = s[1]

    if a == 0
        set!(p, 1)
        set!(q, 1)
    else
        # p = (6a - 5) * (2a - 1) * (6a - 1)
        set!(p, 6*a - 5)
        mul!(p, 2*a - 1)
        mul!(p, 6*a - 1)

        # q = a^3 * c3_over_24
        set!(q, a)
        pow!(q, 3)
        mul!(q, c3_over_24)
    end

    # t(a) = p(a) * a(a)
    # t = p * (13591409 + 545140134 * a)
    set!(t, a)
    mul!(t, 545140134)
    add!(t, 13591409)
    mul!(t, p)
    isodd(a) && neg!(t)
end

combine!(s) = begin
    p₁, q₁, t₁ = s[1]
    p₂, q₂, t₂ = s[2]

    # t = q₂ * t₁ + p₁ * t₂
    mul!(t₁, q₂)
    mul!(t₂, p₁)
    add!(t₁, t₂)

    # p = p₁ * p₂
    mul!(p₁, p₂)

    # q = q₁ * q₂
    mul!(q₁, q₂)
end

end


"""
ドットプロットでは，データ点が重ならないようにランダムな数値（Jitter）を加えているので，その並びに何らかの意味がありそうに見えてしまう。Python の seaborn にある swarmplot() を Julia から使えるようになっているが，これも同じ特質を持っている。

そのようなことを避けるために
https://blog.goo.ne.jp/r-de-r/e/237f8bdfda925c5d724604aa9eabc8e3
にプログラムを呈示している。
"""

using FreqTables, Plots
function simple_dotplot(x, y; accu=0, stp=0, logflag=false, simple=false,
            symmetrical=true, sizew=600, sizeh=400)
    function logaxis(z)
        pos = zeros(100)
        minval = minimum(z)
        maxval = maximum(z)
        logmaxval = log10(maxval)
        logminval = log10(minval)
        width = logmaxval - logminval
        factor = logmaxval < 3.5 ? 15 : 8
        w = width / factor
        st = floor(logminval)
        delta = 10^st
        ic = 0
        while true
            for i = 1:9
                v = i * delta
                if minval - delta <= v <= maxval + delta
                    if (i == 1 || ic == 0) || (ic > 0 && log10(v) - pos[ic] > w)
                        if i == 1 && ic > 0 && log10(v) - pos[ic] <= w
                            ic -= 1
                        end
                        ic += 1
                        pos[ic] = log10(v)
                        prev = v
                    end
                end
            end
            delta *= 10
            if delta > maxval
                break
            end
        end
        return ic > 3 ? pos[2:ic-1] : [logminval, median([logminval, logmaxval]), logmaxval]
    end
    if logflag == true
        y0 = y
        y = log10.(y);
    end
    minx = minimum(y)
    maxx = maximum(y)
    rng = maxx - minx
    minx = minx - rng * 20 / sizeh
    maxx = maxx + rng * 20 / sizeh
    if accu == 0
        accu = (maxx - minx) / 50
    end
    if stp == 0
        stp = 30 / sizew
    end
    y = round.(y ./ accu, digits=0) * accu
    table = freqtable(y, x)
    yname, xname = names(table)
    nvalue, ngroup = size(table)
    # pyplot()
    p = plot(tick_direction=:out, xticks=false, grid=false,
                size=(sizew, sizeh), xlims = (0, length(xname)+1),
                ylims=(minx, maxx), label="")
    xticks!(1:ngroup, xname)
    for i = 1:ngroup
        for j = 1:nvalue
            if table[j, i] != 0
                n = table[j, i]
                if n == 1
                    height, position = [yname[j]], [i]
                else
                    height = repeat([yname[j]], n)
                    position = collect(range(i-stp*(n-1)/2, i+stp*(n-1)/2, length=n))
                end
                p = scatter!(position, height, labels="", alpha=0.7,
                                markercolor=i, markerstrokecolor=i)
            end
        end
    end
    if logflag
        pos = logaxis(y0)
        yticks!(pos, string.(round.(10 .^ pos, digits = 1)))
    end
    println("accu = $accu, stp = $stp")
    # display(p)
    p
end;

"""
使用法

julia> using RDatasets

julia> airquality = RDatasets.dataset("datasets", "airquality");

julia> x = string.(airquality.Month); # 型変換が必要

julia> y = airquality.Wind;

julia> simple_dotplot(x, y, accu=0.105, stp=0.05)
accu = 0.105, stp = 0.05

"""


function showextremes()
  values = [0.0, -0.0, Inf, -Inf, NaN]
  println(1 ./ values)
end

showextremes()

@show Inf + 2.0
@show Inf + Inf
@show Inf - Inf
@show Inf * Inf
@show Inf / Inf
@show Inf * 0
@show 0 == -0
@show NaN == NaN
@show NaN === NaN


#: Optimization functions and solvers

export OptFISTA, opt_fista

# Options type for FISTA

mutable struct OptFISTA{T<:Real}<:OptimOptions
    steplength::T
    niter::Int64
    tol_x::Union{Nothing,T}
    nesterov::Bool
end

opt_fista(steplength::T; niter::Int64=1000, tol_x::Union{Nothing,T}=nothing, nesterov::Bool=true) where {T<:Real} = OptFISTA{T}(steplength, niter, tol_x, nesterov)

# Generic FISTA solver

"""
Solver for the regularized problem via FISTA-like gradient-projections:
```math
min_x f(x)+g(x)
```
Reference: Beck, A., and Teboulle, M., 2009, A Fast Iterative Shrinkage-Thresholding Algorithm for Linear Inverse Problems
https://www.ceremade.dauphine.fr/~carlier/FISTA
"""
function minimize_fista!(fun::DiffPlusProxFun{CT,N}, initial_estimate::AbstractArray{CT,N}, steplength::T, niter::Int64, nesterov::Bool, tol_x::Union{Nothing,T}, x0::AbstractArray{CT,N}) where {T<:Real,N,CT<:RealOrComplex{T}}

    # Initialization
    x0   .= initial_estimate
    xtmp  = similar(x0)
    x     = similar(x0)
    t0 = T(1)

    # Optimization loop
    for i = 1:niter

        grad!(fun.f, x0, xtmp; eval=false)         # Compute gradient
        xtmp .= x0-steplength*xtmp               # Gradient update
        # xtmp .*= -steplength; xtmp .+= x0          # Gradient update
        proxy!(xtmp, steplength, fun.g, x)         # Compute proxy
        if nesterov                                # > Nesterov two-step update:
            t = T(0.5)*(T(1)+sqrt(T(1)+T(4)*t0^2)) # - compute dynamic step size
            x .+= (t0-T(1))/t*(x-x0)               # - update momentum
            # x .*= (t+t0-T(1))/t; x .-= (t0-T(1))/t*x0 # - update momentum
            t0 = t                                 # - update step size
        end                                        # <

        # Update unknowns
        ~isnothing(tol_x) && norm(x-x0)/norm(x0) < tol_x && break
        x0 .= x

    end

    # Return output
    return x0

end

function minimize_fista_debug!(fun::DiffPlusProxFun{CT,N}, initial_estimate::AbstractArray{CT,N}, steplength::T, niter::Int64, nesterov::Bool, tol_x::Union{Nothing,T}, x0::AbstractArray{CT,N}) where {T<:Real,N,CT<:RealOrComplex{T}}

    # Initialization
    x0 .= initial_estimate
    xtmp = similar(x0)
    x = similar(x0)
    t0 = T(1)
    fval_hist = Array{T,1}(undef, niter)
    err_rel = Array{T,1}(undef, niter)

    # Optimization loop
    for i = 1:niter

        fval = grad!(fun.f, x0, xtmp; eval=true)   # Compute gradient
        xtmp .= x0-steplength*xtmp                 # Gradient update
        proxy!(xtmp, steplength, fun.g, x)         # Compute proxy
        if nesterov                                # > Nesterov two-step update:
            t = T(0.5)*(T(1)+sqrt(T(1)+T(4)*t0^2)) # - compute dynamic step size
            x .= x+(t0-T(1))/t*(x-x0)              # - update momentum
            t0 = t                                 # - update step size
        end                                        # <

        # Convergence check
        fval_hist[i] = fval
        err_rel[i] = norm(x-x0)/norm(x0)
        ~isnothing(tol_x) && err_rel[i] < tol_x && break

        # Update unknowns
        x0 .= x

    end

    # Return output
    return x0, fval_hist, err_rel

end

minimize!(fun::DiffPlusProxFun{CT,N}, x0::AbstractArray{CT,N}, opt::OptFISTA{T}, x::AbstractArray{CT,N}) where {T<:Real,N,CT<:RealOrComplex{T}} = minimize_fista!(fun, x0, opt.steplength, opt.niter, opt.nesterov, opt.tol_x, x)

minimize_debug!(fun::DiffPlusProxFun{CT,N}, x0::AbstractArray{CT,N}, opt::OptFISTA{T}, x::AbstractArray{CT,N}) where {T<:Real,N,CT<:RealOrComplex{T}} = minimize_fista_debug!(fun, x0, opt.steplength, opt.niter, opt.nesterov, opt.tol_x, x)