export center_extract, center_set!, get_indices_around_center, center_pos
export rr, rr_3D
export fftpos
export normabs2
export get_index_for_pos

"""
    get_index_for_pos(pos_arr, pos)

For a arr containing `range` like data,
return the `index` which is closest to the `pos` value

# Example
```julia-repl
julia> pos = collect(fftpos(1, 5))
5-element Vector{Float64}:
 -0.5
 -0.25
  0.0
  0.25
  0.5

julia> get_index_for_pos(pos, 0.124)
3

julia> get_index_for_pos(pos, 0.1251)
4

julia> pos = collect(fftpos(1, 6))
6-element Vector{Float64}:
 -0.5
 -0.3333333333333333
 -0.16666666666666666
  2.2204460492503052e-17
  0.16666666666666669
  0.33333333333333337

julia> get_index_for_pos(pos, 0)
4

julia> get_index_for_pos(pos, 0.1)
5
```
"""
function get_index_for_pos(pos_arr, pos)
    left = pos_arr[begin]
    right = pos_arr[end]
    Δ = (right - left) / (length(pos_arr) - 1)

    ind = 1 + round(Int, (pos - left) / Δ)
    return ind
end

"""
    normabs2(arr)

Takes abs2 and normalizes maximum to 1

 # Examples
```jldoctest
julia> normabs2([10.0, 12.1])
2-element Array{Float64,1}:
 0.6830134553650707
 1.0
```
"""
function normabs2(arr)
    arr2 = abs2.(arr)
    return arr2 ./ maximum(arr2)
end

"""
    fftpos(L, N)

Construct a range from -L/2 to L/2.
However, we ensure that everything is centered around the center
in a way that a FFT interpretes it correctly.
For odd sequences it is indeed in the real center.
For even sequences the center is at `N/2 + 1`.

 # Examples
```jldoctest
julia> collect(fftpos(1, 4))
4-element Array{Float64,1}:
 -0.5
 -0.25
  0.0
  0.25

julia> collect(fftpos(1, 5))
5-element Array{Float64,1}:
 -0.5
 -0.25
  0.0
  0.25
  0.5
```
"""
function fftpos(l, N)
    if iszero(l) || isone(N)
        return zero(l):zero(l)
    end
    if N % 2 == 0
        dx = l / N
        return range(-l/2, l/2-dx, length=N)
    else
        return range(-l/2, l/2, length=N)
    end
end

function fftpos_old(l, N)
    if iszero(l) && isone(N)
        return l:l
    end
    if N % 2 == 0
        dx = l / N
        return -l/2:dx:l/2-dx
    else
        dx = l / (N-1)
        return -l/2:dx:l/2
    end
end

"""
    get_indices_around_center(i_in, i_out)
A function which provides two output indices `i1` and `i2`
where `i2 - i1 = i_out`
The indices are chosen in a way that the set `i1:i2`
cuts the interval `1:i_in` in a way that the center frequency
stays at the center position.
Works for both odd and even indices
"""
function get_indices_around_center(i_in, i_out)
    if (mod(i_in, 2) == 0 && mod(i_out, 2) == 0
     || mod(i_in, 2) == 1 && mod(i_out, 2) == 1)
        x = (i_in - i_out) ÷ 2
        return 1 + x, i_in - x
    elseif mod(i_in, 2) == 1 && mod(i_out, 2) == 0
        x = (i_in - 1 - i_out) ÷ 2
        return 1 + x, i_in - x - 1
    elseif mod(i_in, 2) == 0 && mod(i_out, 2) == 1
        x = (i_in - (i_out - 1)) ÷ 2
        return 1 + x, i_in - (x - 1)
    end
end

"""
    center_extract(arr, new_size_array)
Extracts a center of an array.
`new_size_array` must be list of sizes indicating the output
size of each dimension. Centered means that a center frequency
stays at the center position. Works for even and uneven.
If `length(new_size_array) < length(ndims(arr))` the remaining dimensions
are untouched and copied.
# Examples
```jldoctest
julia> center_extract([1 2; 3 4], [1])
1×2 Array{Int64,2}:
 3  4
julia> center_extract([1 2; 3 4], [1, 1])
1×1 Array{Int64,2}:
 4
julia> center_extract([1 2 3; 3 4 5; 6 7 8], [2 2])
2×2 Array{Int64,2}:
 1  2
 3  4
```
"""
function center_extract(arr::AbstractArray, new_size_array)
    new_size_array = collect(new_size_array)

    # we construct two lists
    # the reason is, that we don't change higher dimensions which are not
    # specified in new_size_array
    out_indices1 = [get_indices_around_center(size(arr)[x], new_size_array[x])
                    for x = 1:length(new_size_array)]

    out_indices1 = [x[1]:x[2] for x = out_indices1]

    # out_indices2 contains just ranges covering the full size of each dimension
    out_indices2 = [1:size(arr)[i] for i = (1 + length(new_size_array)):ndims(arr)]
    return arr[out_indices1..., out_indices2...]
end

"""
    center_set!(arr_large, arr_small)
Puts the `arr_small` central into `arr_large`.
The convention, where the center is, is the same as the definition
as for FFT based centered.
Function works both for even and uneven arrays.
# Examples
```jldoctest
julia> center_set!([1, 1, 1, 1, 1, 1], [5, 5, 5])
6-element Array{Int64,1}:
 1
 1
 5
 5
 5
 1
```
"""
function center_set!(arr_large, arr_small)
    out_is = []
    for i = 1:ndims(arr_large)
        a, b = get_indices_around_center(size(arr_large)[i], size(arr_small)[i])
        push!(out_is, a:b)
    end

    #rest = ones(Int, ndims(arr_large) - 3)
    arr_large[out_is...] = arr_small

    return arr_large
end

"""
    center_pos(x)
Calculate the position of the center frequency.
Size of the array is `x`
# Examples
```jldoctest
julia> center_pos(3)
2
julia> center_pos(4)
3
```
"""
function center_pos(x::Integer)
    # integer division
    return div(x, 2) + 1
end

function rr_3D(s)
    rarr = zeros((s...))
    for k = 1:s[3]
        for j = 1:s[2]
            for i = 1:s[1]
                rarr[i, j, k] = sqrt( (i-center_pos(s[1]))^2 + (j-center_pos(s[2]))^2 + (k-center_pos(s[3]))^2)
            end
        end
    end
    return rarr
end

"""
    rr(s; L=nothing)

Generate a array with values being the distance to the center pixel.
`s` specifies the output size of the 2D array.
`L` specifies the length. Can be a single number indicating
the same length for both dimensions or a tuple specifying a length for each
dimension separately.

# Examples
```julia-repl
julia> PhysicalOptics.rr((4,4))
4×4 Matrix{Float64}:
 2.82843  2.23607  2.0  2.23607
 2.23607  1.41421  1.0  1.41421
 2.0      1.0      0.0  1.0
 2.23607  1.41421  1.0  1.41421

julia> PhysicalOptics.rr((4,4), L=1)
4×4 Matrix{Float64}:
 0.707107  0.559017  0.5   0.559017
 0.559017  0.353553  0.25  0.353553
 0.5       0.25      0.0   0.25
 0.559017  0.353553  0.25  0.353553

julia> PhysicalOptics.rr((4,4), L=(5, 1))
4×4 Matrix{Float64}:
 2.54951  1.34629  0.5   1.34629
 2.51247  1.27475  0.25  1.27475
 2.5      1.25     0.0   1.25
 2.51247  1.27475  0.25  1.27475

julia> PhysicalOptics.rr((5,5), L=(1, 1))
5×5 Matrix{Float64}:
 0.707107  0.559017  0.5   0.559017  0.707107
 0.559017  0.353553  0.25  0.353553  0.559017
 0.5       0.25      0.0   0.25      0.5
 0.559017  0.353553  0.25  0.353553  0.559017
 0.707107  0.559017  0.5   0.559017  0.707107
```
"""
function rr_old(s; L=nothing)
    # in case of odd/even array we need to fix the size
    s2 = s[2] %2 == 0 ? s[2] : s[2] -1
    s1 = s[1] %2 == 0 ? s[1] : s[1] -1
    x = fftpos(s2, s[2])'
    y = fftpos(s1, s[1])

    if isnothing(L) == false
        if typeof(L) <: Number
            x *= L ./ 2 ./ abs(x[1])
            y *= L ./ 2 ./  abs(y[1])
        else
            x *= L[1] ./ 2 ./ abs(x[1])
            y *= L[2] ./ 2 ./ abs(y[1])
        end
    end

    r = sqrt.(x.^2 .+ y .^ 2)
    return r
end

function rr(s, arr_type=Array; L=nothing)
    x, y = _get_rr_x_y(L, s, arr_type)
    r = sqrt.(x.^2 .+ y .^ 2)
    return r
end

function _get_rr_x_y(L::Number, s, arr_type)
    x = to_gpu_or_cpu(arr_type, fftpos(L, s[2])')
    y = to_gpu_or_cpu(arr_type, fftpos(L, s[1]))
    return x, y
end

function _get_rr_x_y(L, s, arr_type)
    x = to_gpu_or_cpu(arr_type, fftpos(L[2], s[2])')
    y = to_gpu_or_cpu(arr_type, fftpos(L[1], s[1]))
    return x, y
end

function _get_rr_x_y(L::Nothing, s, arr_type)
    s2 = s[2] %2 == 0 ? s[2] : s[2] -1
    s1 = s[1] %2 == 0 ? s[1] : s[1] -1
    x = to_gpu_or_cpu(arr_type, fftpos(s2, s[2])')
    y = to_gpu_or_cpu(arr_type, fftpos(s1, s[1]))
    return x, y
end

"""
    jinc(x)

Computes the jinc function which is \$\\text{jinc} = \\frac{J_1(x)}{x}\$
where \$J_1\$ being the first Bessel function.

In future, probably switch to: https://github.com/JuliaMath/SpecialFunctions.jl/issues/264

# Examples
```julia-repl
julia> jinc(0.0)
0.5

julia> jinc(0.0im)
0.5 + 0.0im

julia> jinc(0.0001im)
0.5000000006250005 - 3.0616170016954073e-17im

julia> jinc(1.0f0)
0.44005057f0

julia> jinc.([0.0 0.5 1.0])
1×3 Array{Float64,2}:
 0.5  0.484537  0.440051
```
"""
function jinc(x::T) where T
    if iszero(x)
		return convert(T, 0.5)
	else
		return besselj1(x) / x
	end
end

"""
    apply_rot_symmetry(arr, xs, xs_out, ys_out)

Takes a 1D array `arr` with rotational symmetry about the center point.
`arr` should have values evaluates at the positions `xs`.

`xs_out` and `ys_out` are the positions for the output array.
We interpolate `arr` and evaluate it at distances `sqrt(xs_out[j]^2+ys_out[i]^2)`.

"""
function apply_rot_symmetry(arr, xs, xs_out, ys_out)
    itp = Interpolations.scale(
            interpolate(arr,
                        BSpline(Cubic(Line(Interpolations.OnCell())))),
            xs)

    out = zeros(eltype(arr), (length(ys_out), length(xs_out)))

    for j in axes(out, 2)
        for i in axes(out, 1)
            r = sqrt(xs_out[j]^2+ys_out[i]^2)
            out[i, j] = itp(r)
        end
    end

    return out

end


using FractionalDiffEq

K = 1
fdorder = 1.9
dx = pi/20
dt = 0.01
n = 2
xStart = 0
xEnd = pi
x = collect(0:dx:xEnd)
t = collect(0:dt:n)
u0t = 0
uendt = 0
u0(x) = sin(x)

using Plots
plotlyjs()
U=solve(fdorder, dx, dt, xStart, xEnd, n, K, u0t, uendt, u0, FiniteDiffEx())
plt=plot(x, t, U, st=:surface)


using ConvOptDL
using ConvOptDL.Utils
using StatsBase
using LinearAlgebra: I
using ArgParse
using Serialization
using Pipe: @pipe
using Flux: update!, trainmode!
using Statistics: mean
using Logging: @info
using Dates: now
import Flux

function loss_log_softmax(logits, one_hot_vec)
    log_prob = log.(Flux.softmax(logits, dims = 1))
    loss = @pipe -(log_prob .* one_hot_vec) |> sum(_, dims = 1)
    mean(loss)
end

_format_batch(batch) = reshape(Float32.(batch), size(batch)[1:end-2]..., :)

function train!(loss, model, batch, opt)
    # https://fluxml.ai/Flux.jl/stable/training/training/#Custom-Training-loops-1
    ps = Flux.params(model)
    # smoothed onehot encoding
    # onehot is not differentiable, moving out of the do block
    query_onehot = @pipe ConvOptDL.Utils.onehot(batch.query_labels) |>
          reshape(_, batch.support_n_ways, :) |>
          _ .* (1 - 5f-2) .+ 5f-2 * (1 .- _) ./ (batch.support_n_ways - 1f0)
    support_onehot = ConvOptDL.Utils.onehot(batch.support_labels)

    local meta_loss
    gs = Flux.gradient(ps) do
        # (`feat_dim`, `n_ways`*`k_shots`, `num_tasks`)
        embed_query = reshape(
            model(_format_batch(batch.query_samples)),
            :,
            batch.n_query,
            size(batch),
        )
        embed_support = reshape(
            model(_format_batch(batch.support_samples)),
            :,
            batch.n_support,
            size(batch),
        )
        Q, p, G, h, A, b = crammer_svm(embed_support, support_onehot, batch)
        # solve QP
        X, λ, ν = solve_qp_batch(Q, p, G, h, A, b)
        α = @pipe X |>
              reshape(_, 1, batch.n_support, batch.support_n_ways, size(batch)) |>
              repeat(_, outer = (batch.n_query, 1, 1, 1)) |>
              permutedims(_, (3, 2, 1, 4))
        # logits
        logits = @pipe ConvOptDL.Utils.gram_matrix(embed_query, embed_support) |>
              reshape(_, 1, size(_)...) |>
              repeat(_, outer = (5, 1, 1, 1)) |>
              _ .* α |>
              sum(_, dims = 2) |>
              dropdims(_, dims = 2) |>
              reshape(_, batch.support_n_ways, :)
        meta_loss = loss(logits, query_onehot)
        return meta_loss
    end
    update!(opt, ps, gs)
    return meta_loss
end

function parse_opts()
    s = ArgParseSettings("Meta Convex Optimization Learning training")
    @add_arg_table! s begin
        "--batch-size"
        arg_type = Int64
        metavar = "INT"
        default = 8
        "--batches-per-episode"
        arg_type = Int64
        metavar = "INT"
        default = 200
        "--log-period"
        help = "print log for every given number of batches"
        arg_type = Int64
        metavar = "INT"
        default = 25
        "--num-episodes"
        arg_type = Int64
        metavar = "INT"
        default = 50
        "--learning-rate"
        arg_type = Float32
        default = 1f-1
        "-o"
        arg_type = String
        default = "model.jls"
        "data_file"
        arg_type = String
        default = "./test/test_data.jls"
    end
    args = parse_args(s)
    return args
end

if nameof(@__MODULE__) == :Main
    args = parse_opts()
    @info "Training Options: " args
    batch_size = args["batch-size"]
    batches_per_episode = args["batches-per-episode"]
    log_period = args["log-period"]
    num_episodes = args["num-episodes"]
    data_file = args["data_file"]
    out_model_file = args["o"]
    lr = args["learning-rate"]
    model_name, _ = splitext(out_model_file)
    dloader = FewShotDataLoader(data_file)
    opt = Flux.Optimise.Momentum(lr)
    model = resnet12()
    trainmode!(model)
    record = Dict()
    for episode = 1:num_episodes
        @info "episode $(episode) start" now()
        meta_losses = []
        total_time = 0
        for i = 1:batches_per_episode
            time_start = time()
            batch = sample(dloader, batch_size, support_n_ways = 5, support_k_shots = 2)
            meta_loss = train!(loss_log_softmax, model, batch, opt)
            time_end = time()
            push!(meta_losses, meta_loss)
            total_time += (time_end - time_start)
            if i % log_period == 0
                @info "[$(now())] last $(log_period) meta losses" meta_losses[end-log_period+1:end] mean(meta_losses)
                @info "[$(now())] average execution time" mean_exec_time = (total_time / i)
            end
        end
        @info "episode $(episode) end" now()
        record[episode] = meta_losses
        serialize("train_$(model_name)_meta_losses_$(episode).jls", meta_losses)
    end
    serialize("all_train_$(model_name)_meta_losses.jls", record)
    serialize(out_model_file, model)
end


abstract type AbstractEquationOfState{𝒯} end

struct BarotropicFluid{𝒯} <: AbstractEquationOfState{𝒯} end
struct IdealGas{𝒯} <: AbstractEquationOfState{𝒯} end

@inline function pressure(::BarotropicFluid{(:ρ, :ρu)}, state, aux, params)
    ρ = state.ρ
    cₛ = params.cₛ
    ρₒ = params.ρₒ

    return (cₛ * ρ)^2 / (2 * ρₒ)
end

@inline function pressure(::IdealGas{(:ρ, :ρu, :ρe)}, state, aux, params)
    ρ = state.ρ
    ρu = state.ρu
    ρe = state.ρe
    Φ = aux.Φ
    γ = params.γ

    return (γ - 1) * (ρe - dot(ρu, ρu) / 2ρ - ρ * Φ)
end

@inline function sound_speed(::BarotropicFluid{(:ρ, :ρu)}, state, aux, params)
    cₛ = params.cₛ
    ρₒ = params.ρₒ
    ρ = state.ρ

    return cₛ * sqrt(ρ / ρₒ)
end

@inline function sound_speed(eos::IdealGas{(:ρ, :ρu, :ρe)}, state, aux, params)
    ρ = state.ρ
    ρu = state.ρu
    ρe = state.ρe
    Φ = aux.Φ
    γ = params.γ

    p = calc_pressure(eos, state, aux, params)

    return sqrt(γ * p / ρ)
end


"""
```julia
gradient(O::Expression, vars::AbstractVector{<:Expression}; simplify = true)
```

A helper function for computing the gradient of an expression with respect to
an array of variable expressions.
"""
function gradient(O::Expression, vars::AbstractVector{<:Expression}; simplify = true)
    [expand_derivatives(Differential(v)(O),simplify) for v in vars]
end

"""
```julia
jacobian(ops::AbstractVector{<:Expression}, vars::AbstractVector{<:Expression}; simplify = true)
```

A helper function for computing the Jacobian of an array of expressions with respect to
an array of variable expressions.
"""
function jacobian(ops::AbstractVector{<:Expression}, vars::AbstractVector{<:Expression}; simplify = true)
    [expand_derivatives(Differential(v)(O),simplify) for O in ops, v in vars]
end

"""
```julia
sparsejacobian(ops::AbstractVector{<:Expression}, vars::AbstractVector{<:Expression}; simplify = true)
```

A helper function for computing the sparse Jacobian of an array of expressions with respect to
an array of variable expressions.
"""
function sparsejacobian(ops::AbstractVector{<:Expression}, vars::AbstractVector{<:Expression}; simplify = true)
    I = Int[]
    J = Int[]
    du = Expression[]

    sp = jacobian_sparsity(ops, vars)
    I,J,_ = findnz(sp)

    exprs = Expression[]

    for (i,j) in zip(I, J)
        push!(exprs, expand_derivatives(Differential(vars[j])(ops[i]), simplify))
    end
    sparse(I,J, exprs, length(ops), length(vars))
end

"""
```julia
jacobian_sparsity(ops::AbstractVector{<:Expression}, vars::AbstractVector{<:Expression})
```

Return the sparsity pattern of the Jacobian of an array of expressions with respect to
an array of variable expressions.
"""
function jacobian_sparsity(du, u)
    dict = Dict(zip(to_symbolic.(u), 1:length(u)))

    i = Ref(1)
    I = Int[]
    J = Int[]

    # This rewriter notes down which u's appear in a
    # given du (whose index is stored in the `i` Ref)
    r = [@rule ~x::(x->haskey(dict, x)) => begin
        push!(I, i[])
        push!(J, dict[~x])
        nothing
    end] |> Rewriters.Chain |> Rewriters.Postwalk

    for ii = 1:length(du)
        i[] = ii
        r(to_symbolic(du[ii]))
    end

    sparse(I, J, true, length(du), length(u))
end

"""
```julia
hessian(O::Expression, vars::AbstractVector{<:Expression}; simplify = true)
```

A helper function for computing the Hessian of an expression with respect to
an array of variable expressions.
"""
function hessian(O::Expression, vars::AbstractVector{<:Expression}; simplify = true)
    first_derivs = vec(jacobian([O], vars, simplify=simplify))
    n = length(vars)
    H = Array{Expression, 2}(undef,(n, n))
    fill!(H, 0)
    for i=1:n
        for j=1:i
            H[j, i] = H[i, j] = expand_derivatives(Differential(vars[i])(first_derivs[j]))
        end
    end
    H
end

isidx(x) = x isa TermCombination

let
    _scalar = one(TermCombination)

    linearity_propagator = [
          @rule +(~~xs) => reduce(+, filter(isidx, ~~xs), init=_scalar)
          @rule *(~~xs) => reduce(*, filter(isidx, ~~xs), init=_scalar)
          @rule (~f)(~x::(!isidx)) => _scalar
          @rule (~f)(~x::isidx) => if haslinearity(~f, Val{1}())
              combine_terms(linearity(~f, Val{1}()), ~x)
          else
              error("Function of unknown linearity used: ", ~f)
          end
          @rule (^)(~x::isidx, ~y) => ~y isa Number && isone(~y) ? ~x : (~x) * (~x)
          @rule (~f)(~x, ~y) => begin
              if haslinearity(~f, Val{2}())
                  a = isidx(~x) ? ~x : _scalar
                  b = isidx(~y) ? ~y : _scalar
                  combine_terms(linearity(~f, Val{2}()), a, b)
              else
                  error("Function of unknown linearity used: ", ~f)
              end
          end] |> Rewriters.Chain |> Rewriters.Postwalk |> Rewriters.Fixpoint

    # we do this in a let block so that Revise works on the list of rules
    global hessian_sparsity

    """
    ```julia
    hessian_sparsity(ops::AbstractVector{<:Expression}, vars::AbstractVector{<:Expression})
    ```

    Return the sparsity pattern of the Hessian of an array of expressions with respect to
    an array of variable expressions.
    """
    function hessian_sparsity(f, u)
        idx(i) = TermCombination(Set([Dict(i=>1)]))
        dict = Dict(SymbolicUtils.to_symbolic.(u) .=> idx.(1:length(u)))
        found = []
        f = Rewriters.Prewalk(x->haskey(dict, x) ? dict[x] : x)(to_symbolic(f))
        _sparse(linearity_propagator(to_symbolic(f)), length(u))
    end
end

"""
```julia
islinear(ex::Expression, u)
```
Check if an expression is linear with respect to a list of variable expressions.
"""
function islinear(ex::Expression, u)
    isempty(hessian_sparsity(ex, u).nzval)
end

"""
```julia
sparsehessian(O::Expression, vars::AbstractVector{<:Expression}; simplify = true)
```

A helper function for computing the sparse Hessian of an expression with respect to
an array of variable expressions.
"""
function sparsehessian(O::Expression, vars::AbstractVector{<:Expression}; simplify = true)
    S = hessian_sparsity(O, vars)
    I, J, _ = findnz(S)
    exprs = Array{Expression}(undef, length(I))
    fill!(exprs, 0)
    prev_j = 0
    d = nothing
    for (k, (i, j)) in enumerate(zip(I, J))
        j > i && continue
        if j != prev_j
            d = expand_derivatives(Differential(vars[j])(O), false)
        end
        expr = expand_derivatives(Differential(vars[i])(d), simplify)
        exprs[k] = expr
        prev_j = j
    end
    H = sparse(I, J, exprs, length(vars), length(vars))
    for (i, j) in zip(I, J)
        j > i && (H[i, j] = H[j, i])
    end
    return H
end

function simplified_expr(O::Operation)
  if O isa Constant
    return O.value
  elseif isa(O.op, Differential)
    return :(derivative($(simplified_expr(O.args[1])),$(simplified_expr(O.op.x))))
  elseif isa(O.op, Variable)
    isempty(O.args) && return O.op.name
    return Expr(:call, Symbol(O.op), simplified_expr.(O.args)...)
  end
  if O.op === (^)
      if length(O.args) > 1  && O.args[2] isa Constant && O.args[2].value < 0
          return Expr(:call, :^, Expr(:call, :inv, simplified_expr(O.args[1])), -(O.args[2].value))
      end
  end
  return Expr(:call, Symbol(O.op), simplified_expr.(O.args)...)
end

simplified_expr(c::Constant) = c.value

function simplified_expr(eq::Equation)
    Expr(:(=), simplified_expr(eq.lhs), simplified_expr(eq.rhs))
end

simplified_expr(eq::AbstractArray) = simplified_expr.(eq)
simplified_expr(x::Integer) = x
simplified_expr(x::AbstractFloat) = x


#-----------------------------------------------------------------------
#   FULL TEST OF BROYDEN CODE
#-----------------------------------------------------------------------

d = 12

r0 = [0.6027037403222177, 0.6823081703161471, 0.4969968192879606,
      0.9455425492979632, 0.09666423716951855, 0.5533229648934559,
      0.6257838591529667, 0.3219114195167574, 0.9465987429640241,
      0.9543249417929738, 0.1987862235853426, 0.050519233798469454]

r1 = [0.45667659772446867, 0.29604904127839915, 0.36276182662693257,
      0.05020123630607376, 0.5045105043961433, 0.21437469923464825,
      0.9063213634872238, 0.6585306236631778, 0.14703614205907733,
      0.5187343751664399, 0.137009564745886, 0.7477554965317557]

r2 = [0.8559101941413021, 0.8407795618867213, 0.9935645925564422,
      0.24371648277714364, 0.5445775437567715, 0.05095659932958019,
      0.788215773526395, 0.9845642037746567, 0.5623887100264071,
      0.11735622784087019, 0.07719560594089825, 0.5954214792591306]

x0 = 2*r0/1000
x1 = r1
x2 = -π*r2/10

f0 = x->2*π*(x-x0)
Jac0 = ForwardDiff.jacobian(f0,zeros(d));

@test (bsolve(f0,Jac0,f0(zeros(d)),zeros(d),2*d)-x0)./x0 ≈ zeros(d)

f1 = x->(x-x0).*(x-x1)
Jac1 = ForwardDiff.jacobian(f1,zeros(d));

@test (bsolve(f1,Jac1,f1(zeros(d)),zeros(d),2*d)-x0)./x0 ≈ zeros(d)

f2 = x->(x-x0).*(x-x1).*(x-x2)
Jac2 = ForwardDiff.jacobian(f2,zeros(d));

@test (bsolve(f2,Jac2,f2(zeros(d)),zeros(d),2*d)-x0)./x0 ≈ zeros(d)


# abstract type AbstractDensity end

# struct LineDensity{T} <: AbstractDensity
#     x::Vector{T}
#     N::Int
#     bins::Int
#     function LineDensity(x::Vector{<:Real}; bins::Int=100)
#         T = eltype(x)
#         N = length(x)
#         return new{T}(x, N, bins)
#     end
# end

"""
    polar1d(y::AbstractVector, theta::AbstractVector, grid::AbstractVector; normalize::Bool=True)
Compute polar order across channel. Normalized such that int n dy =1.
"""
function polar1d(y::AbstractVector, theta::AbstractVector, grid::AbstractVector; normalize::Bool=true)
    bins = length(grid) -1
    # initialize array to store number
    n = zeros(bins); np = zeros(bins);
    m = zeros(bins); mp = zeros(bins);
    # in a cell [a,b], particles that has a<=y<b is counted.
    # for the last cell, y==b is also counted.
    for i in 1:bins
        for (position, angle) in zip(y, theta)
            if grid[i] <=position <grid[i+1]
                n[i] += 1; m[i] += sin(angle);
                if angle<pi
                    np[i] += 1; mp[i] += sin(angle);
                end
            end
        end
    end
    # last cell: add particles on the right boundary
    ind=findall(y->y==grid[end], y)
    for angle in y[ind]
        m[end] += sin(angle);
        if angle<pi
            np[i] += 1; mp[i] += sin(angle);
        end
    end
    n[end] += length(ind)

    # use bin centers as representative locations.
    dy = diff(grid)
    bin_centers = grid[1:end-1] + dy/2
    ind = findall(n->n>0, n)
    m[ind] ./= n[ind]
    ind = findall(np->np>0, np)
    mp[ind] ./= np[ind]
    if normalize
        n ./= dy *length(y)
        np ./= dy*length(y)
    end
    return bin_centers, n, np, m, mp
end

"""
    polar1d(y::AbstractVector, theta::AbstractVector, ymin::Real, ymax::Real; bins::Int=100, normalize::Bool=true)
Compute polar order across channel. Normalized such that int n dy =1.
"""
function polar1d(y::AbstractVector, theta::AbstractVector, ymin::Real, ymax::Real; bins::Int=100, normalize::Bool=true)
    ymin < ymax ||throw(ArgumentError("ymin>=ymax"))
    grid = LinRange(ymin, ymax, bins+1)
    return polar1d(y, theta, grid, normalize=normalize)
end

"""
    polar1d(y::AbstractVector, theta::AbstractVector; bins::Int=100, normalize::Bool=true)
Compute polar order across channel.
"""
function polar1d(y::AbstractVector, theta::AbstractVector; bins::Int=100, normalize::Bool=true)
    ymin = minimum(y)
    ymax = maximum(y)
    return polar1d(y, theta, ymin, ymax, bins=bins, normalize=normalize)
end


# # Input PDF parametrisation and priors
#
# An important part of PDF fitting is defining a useful parametrisation for the
# PDF shapes, as well as meaningful prior distributions that encode our knowledge
# of the problem.
#
# In this notebook, we explore two different approaches:
# * *Full Dirichlet*
# * *Valence shape + Dirichlet*
#
# In the end, it seems like the latter option makes more sense for us and is
# therefore implemented elsewhere in the `PartonDensity` package.
# We demonstrate why below.

using Distributions, Plots, SpecialFunctions, Printf
const sf = SpecialFunctions;

# ## "Full Dirichlet" approach
#
# A clean way to ensure the momentum sum rule would be to sample different
# contributions of the momentrum density integral from a Dirichlet distribution,
# then use these weights to set the parameters on the individual Beta distributions.
# However, in practice this is non-trvial as we also want to fix the normalisation
# of the *number densities* of the valance contributions.

# 9 components of decreasing importance

dirichlet = Dirichlet([3., 2., 1, 0.5, 0.3, 0.2, 0.1, 0.1, 0.1])
data = rand(dirichlet, 1000);

# Have a look

plot()
for i in 1:9
    histogram!(data[i,:], bins=range(0, stop=1, length=20), alpha=0.7)
end
plot!(xlabel="I_i = A_i B_i")

# This would be great as the sum rule is automatically conserved

sum(data, dims=1)

# But, it is non-trival to define valence params from this

I = rand(dirichlet)

# Valance u component

λ_u = rand(Uniform(0, 1))
K_u = rand(Uniform(0, 10))

# Integral of number density must = 2

A_u = 2 / sf.beta(λ_u, K_u+1)

# integral of momentum density can be fixed by I[1]

I_1 = A_u * sf.beta(λ_u+1, K_u+1);

# Could use a root-finder to find K_u given I_1 and λ_u...
# Could be nasty to sample from though...

I_1 = 2 * (sf.beta(λ_u+1, K_u+1)/sf.beta(λ_u, K_u+1));

using Roots

function func_to_solve(K_u)
    return I_1 - 2 * (sf.beta(λ_u+1, K_u+1) / sf.beta(λ_u, K_u+1))
end

K_u ≈ find_zero(func_to_solve, (0, 10), Bisection())

# While this approach might be nice, there are two issues in practice:
# * It is difficult to set sensible priors on $\lambda_u$ that imply priors on $K_u$, and similarly for $\lambda_d$ and $K_d$
# * The problem is overconstrained and we hav to use a root finder. This is rather fragile, and could fail for certain parameter combinations, such as we might find in a fit.

# ## "Valence shape + Dirichlet" approach
#
# We can handle this more elegantly (maybe?) by specifying constraints on the
# valence params through the shape of their Beta distributions, then using a
# Dirichlet to specify the weights of the gluon and sea components. The problem
# here is it isn't clear how to specify that the d contribution must be less than
# the u contribution, but it is possible to do this indirectly through priors on
# the shape parameters. This will however require some further investigation.

x = range(0, stop=1, length=50)

# High-level priors
# Looks like we maybe want to change lambda and K priors to boost these components

λ_u = 0.7 #rand(Uniform(0, 1))
K_u = 4 #rand(Uniform(2, 10))
λ_d = 0.5 #rand(Uniform(0, 1))
K_d = 6 #rand(Uniform(2, 10))

u_V = Beta(λ_u, K_u+1)
A_u = 2 / sf.beta(λ_u, K_u+1)

d_V = Beta(λ_d, K_d+1)
A_d = 1 / sf.beta(λ_d, K_d+1)

# Integral contributions

I_u = A_u * sf.beta(λ_u+1, K_u+1)
I_d = A_d * sf.beta(λ_d+1, K_d+1)

plot(x, x .* A_u .* x.^λ_u .* (1 .- x).^K_u * 2, alpha=0.7, label="x u(x)", lw=3)
plot!(x, x .* A_d .* x.^λ_d .* (1 .- x).^K_d, alpha=0.7, label="x d(x)", lw=3)
plot!(xlabel="x", legend=:topright)

#

@printf("I_u = %.2f\n", I_u)
@printf("I_d = %.2f\n", I_d)

# The remaining 7 integrals can be dirichlet-sampled with decreasing importance

remaining = 1 - (I_u + I_d)
dirichlet = Dirichlet([3., 2., 1, 0.5, 0.3, 0.2, 0.1])
I = rand(dirichlet) * remaining;
sum(I) ≈ remaining

# Gluon contributions

λ_g1 = rand(Uniform(-1, 0))
λ_g2 = rand(Uniform(0, 1))
K_g = rand(Uniform(2, 10))
A_g2 = I[1] / sf.beta(λ_g2+1, K_g+1)
A_g1 = I[2] / sf.beta(λ_g1+1, 5+1);

# Sea quark contributions

λ_q = rand(Uniform(-1, 0))
A_ubar = I[3] / (2 * sf.beta(λ_q+1, 5+1))
A_dbar = I[4] / (2 * sf.beta(λ_q+1, 5+1))
A_s = I[5] / (2 * sf.beta(λ_q+1, 5+1))
A_c = I[6] / (2 * sf.beta(λ_q+1, 5+1))
A_b = I[7] / (2 * sf.beta(λ_q+1, 5+1));

total = A_u * sf.beta(λ_u+1, K_u+1) + A_d * sf.beta(λ_d+1, K_d+1)
total += A_g1 * sf.beta(λ_g1+1, 5+1) + A_g2 * sf.beta(λ_g2+1, K_g+1)
total += 2 * (A_ubar + A_dbar + A_s + A_c + A_b) * sf.beta(λ_q+1, 5+1)
total ≈ 1

#

x = 10 .^ range(-2, stop=0, length=500)

# How does it look?

xg2 = A_g2 * x.^λ_g2 .* (1 .- x).^K_g
xg1 = A_g1 * x.^λ_g1 .* (1 .-x).^5
plot(x, x .* A_u .* x.^λ_u .* (1 .- x).^K_u * 2, alpha=0.7, label="x u(x)", lw=3)
plot!(x, x .* A_d .* x.^λ_d .* (1 .- x).^K_d, alpha=0.7, label="x d(x)", lw=3)
plot!(x, xg1 + xg2, alpha=0.7, label="x g(x)", lw=3)
plot!(x, A_ubar * x.^λ_q .* (1.0 .- x).^5, alpha=0.7, label="x ubar(x)", lw=3)
plot!(x, A_dbar * x.^λ_q .* (1.0 .- x).^5, alpha=0.7, label="x dbar(x)", lw=3)
plot!(x, A_s * x.^λ_q .* (1.0 .- x).^5, alpha=0.7, label="x s(x)", lw=3)
plot!(x, A_c * x.^λ_q .* (1.0 .- x).^5, alpha=0.7, label="x c(x)", lw=3)
plot!(x, A_b * x.^λ_q .* (1.0 .-  x).^5, alpha=0.7, label="x b(x)", lw=3)
plot!(xlabel="x", legend=:bottomleft, xscale=:log, ylims=(1e-8, 10), yscale=:log)

# ### Prior predictive checks
#
# We can start to visualise the type of PDFs that are allowed by the
# combination of the choice of parametrisation and prior distributions
# with some simple prior predictive checks, as done below...

N = 100
alpha = 0.03
total = Array{Float64, 1}(undef, N)
first = true
leg = 0

plot()
for i in 1:N

    λ_u_i = rand(Uniform(0, 1))
    K_u_i = rand(Uniform(2, 10))
    λ_d_i = rand(Uniform(0, 1))
    K_d_i = rand(Uniform(2, 10))
    A_u_i = 2 / sf.beta(λ_u_i, K_u_i+1)
    A_d_i = 1 / sf.beta(λ_d_i, K_d_i+1)
    I_u_i = A_u * sf.beta(λ_u_i+1, K_u_i+1)
    I_d_i = A_d * sf.beta(λ_d_i+1, K_d_i+1)
    u_V_i = Beta(λ_u_i, K_u_i+1)
    d_V_i = Beta(λ_d_i, K_d_i+1)

    remaining_i = 1 - (I_u_i + I_d_i)
    dirichlet_i = Dirichlet([3., 2., 1, 0.5, 0.3, 0.2, 0.1])
    I_i = rand(dirichlet_i) * remaining_i

    λ_g1_i = rand(Uniform(-1, 0))
    λ_g2_i = rand(Uniform(0, 1))
    K_g_i = rand(Uniform(2, 10))
    A_g2_i = I_i[1] / sf.beta(λ_g2_i+1, K_g_i+1)
    A_g1_i = I_i[2] / sf.beta(λ_g1_i+1, 5+1)

    λ_q_i = rand(Uniform(-1, 0))
    A_ubar_i = I_i[3] / (2 * sf.beta(λ_q_i+1, 5+1))
    A_dbar_i = I_i[4] / (2 * sf.beta(λ_q_i+1, 5+1))
    A_s_i = I_i[5] / (2 * sf.beta(λ_q_i+1, 5+1))
    A_c_i = I_i[6] / (2 * sf.beta(λ_q_i+1, 5+1))
    A_b_i = I_i[7] / (2 * sf.beta(λ_q_i+1, 5+1))

    total[i] = A_u_i * sf.beta(λ_u_i+1, K_u_i+1) + A_d_i * sf.beta(λ_d_i+1, K_d_i+1)
    total[i] += A_g1_i * sf.beta(λ_g1_i+1, 5+1) + A_g2_i * sf.beta(λ_g2_i+1, K_g_i+1)
    total[i] += 2 * (A_ubar_i + A_dbar_i + A_s_i + A_c_i + A_b_i) * sf.beta(λ_q_i+1, 5+1)

    xg2_i = A_g2_i * x.^λ_g2_i .* (1 .- x).^K_g_i
    xg1_i = A_g1_i * x.^λ_g1_i .* (1 .- x).^5
    plot!(x, [x .* A_u_i .* x.^λ_u_i .* (1 .- x).^K_u_i * 2], alpha=alpha, color="blue", lw=3)
    plot!(x, x .* A_d_i .* x.^λ_d_i .* (1 .- x).^K_d_i, alpha=alpha, color="orange", lw=3)
    plot!(x, xg1_i + xg2_i, alpha=alpha, color="green", lw=3)
    plot!(x, A_ubar_i * x.^λ_q_i .* (1 .- x).^5, alpha=alpha, color="red", lw=3)
    plot!(x, A_dbar_i * x.^λ_q_i .* (1 .- x).^5, alpha=alpha, color="purple", lw=3)
    plot!(x, A_s_i * x.^λ_q_i .* (1 .- x).^5, alpha=alpha, color="brown", lw=3)
    plot!(x, A_c_i * x.^λ_q_i .* (1 .- x).^5, alpha=alpha, color="pink", lw=3)
    plot!(x, A_b_i * x.^λ_q_i .* (1 .- x).^5, alpha=alpha, color="grey", lw=3)
end

plot!(xlabel="x", ylabel="x f(x)", xscale=:log, legend=false,
      ylims=(1e-8, 10), yscale=:log)

# Looks like naive priors need some work...

# ## PDF Parametrisation interface
#
# `PartonDensity` provides a handy interface to the "Valence shape + Dirichlet" style parametrisation, as demonstrated here.

using PartonDensity

hyper_params = PDFParameters(λ_u=0.5, K_u=4.0, λ_d=0.6, K_d=6.0, λ_g1=-0.37, λ_g2=-0.7,
                             K_g=6.0, λ_q=0.5, seed=5, weights=[50., 0.5, 5., 5., 3., 2., 1.]);

plot_input_pdfs(hyper_params)

int_xtotx(hyper_params) ≈ 1


# See p. 280 of Murphy's Machine Learning
# x_k1 = x_k - alpha * gr + mu * (x - x_previous)

# L should be function or any other callable
immutable MomentumGradientDescent{L} <: Optimizer
    mu::Float64
    linesearch!::L
end

#= uncomment for v0.8.0
MomentumGradientDescent(; mu::Real = 0.01, linesearch = LineSearches.HagerZhang()) =
  MomentumGradientDescent(Float64(mu), linesearch)
=#

Base.summary(::MomentumGradientDescent) = "Momentum Gradient Descent"

function MomentumGradientDescent(; mu::Real = 0.01, linesearch = LineSearches.HagerZhang())
    MomentumGradientDescent(Float64(mu), linesearch)
end

type MomentumGradientDescentState{T,N}
    x::Array{T,N}
    x_previous::Array{T,N}
    f_x_previous::T
    s::Array{T,N}
    @add_linesearch_fields()
end

function initial_state{T}(method::MomentumGradientDescent, options, d, initial_x::Array{T})
    value_gradient!(d, initial_x)

    MomentumGradientDescentState(copy(initial_x), # Maintain current state in state.x
                         copy(initial_x), # Maintain previous state in state.x_previous
                         T(NaN), # Store previous f in state.f_x_previous
                         similar(initial_x), # Maintain current search direction in state.s
                         @initial_linesearch()...) # Maintain a cache for line search results in state.lsr
end

function update_state!{T}(d, state::MomentumGradientDescentState{T}, method::MomentumGradientDescent)
    n = length(state.x)
    # Search direction is always the negative gradient
    @simd for i in 1:n
        @inbounds state.s[i] = -gradient(d, i)
    end

    # Determine the distance of movement along the search line
    lssuccess = perform_linesearch!(state, method, d)

    # Update current position
    @simd for i in 1:n
        # Need to move x into x_previous while using x_previous and creating "x_new"
        @inbounds tmp = state.x_previous[i]
        @inbounds state.x_previous[i] = state.x[i]
        @inbounds state.x[i] = state.x[i] + state.alpha * state.s[i] + method.mu * (state.x[i] - tmp)
    end
    lssuccess == false # break on linesearch error
end


using Percolation
using Base.Test

# square
## nearest neighbor
@test try
	N = 50; p = 0.6
	sq = Square(N, p, "nn")
	label_components!(sq)
	heatmap!(sq)
	true
catch
	false
end

## next nearest neighbor
@test try
	N = 100; p = 0.45
	sq = Square(N, p, "nnn")
	label_components!(sq)
	heatmap!(sq)
	true
catch
	false
end

# triangular
@test try
	N = 100; p = 0.45
	tri = Triangular(N, p, "nnn")
	label_components!(tri)
	heatmap!(tri)
	true
catch
	false
end

# honeycomb
@test try
	N = 100; p = 0.45
	h = Honeycomb(N, p, "nnn")
	label_components!(h)
	heatmap!(h)
	true
catch
	false
end

# Percolating probability
## square lattice, nearest neighbor
@test try
	linsize = 50
	ps = 0.1
	pinc = 0.025
	pf = 1
	nsample = 50
	plot_percolation_prob("square", linsize, ps, pinc, pf, nsample)
	true
catch
	false
end

## triangular lattice
@test try
	linsize = 50
	ps = 0
	pinc = 0.025
	pf = 1
	nsample = 50
	plot_percolation_prob("triangular", linsize, ps, pinc, pf, nsample)
	true
catch
	false
end

## honeycomb lattice
@test try
	linsize = 50
	ps = 0
	pinc = 0.025
	pf = 1
	nsample = 50
	plot_percolation_prob("honeycomb", linsize, ps, pinc, pf, nsample)
	true
catch
	false
end

# Cluster property
lattice_types = [:Square, :Triangular, :Honeycomb]
@test try
    for lattice_type in lattice_types
        eval(quote
    	    N = 50; p = 0.6
    	    l = $lattice_type(N, p, "nn")
    	    heatmap!(l)
    	    clustersize!(l)
    	    clustersizefreq!(l)
    	    clusternumber!(l)
    	    average_clustersize!(l)
    	    strength!(l)
    end)
    end
    true
catch
	false
end

# Lifetime of forest fire
@test try
	linsize = 50
	ps = 0
	pinc = 0.025
	pf = 1
	nsample = 50
	plot_lifetime(linsize, ps, pinc, pf, nsample)
	true
catch
	false
end


using FullyThreadedTree
using Test

@testset "FullyThreadedTree.jl" begin
    # Verify that the cell_volume corresponding to all the cells cover the
    # domain exactly once
    function total_volume_via_cells(tree)
        vol = 0.0
        for cell ∈ cells(tree, filter = active)
            vol += volume(cell)
        end
        return vol
    end

    function total_firstMoment_via_cells(tree::Tree{N}) where N
        moment = zeros(N)
        for cell ∈ cells(tree, filter = active)
            moment += centroid(cell) * volume(cell)
        end
        return moment
    end

    # Verify that the cell_volume corresponding to all the faces cover the
    # domain exactly dim times (NB since boundary faces are inactive, we need
    # periodicity)
    function total_volume_via_faces(tree)
        vol = 0.0
        for face ∈ all_faces(tree, filter = active)
            vol += cell_volume(face)
        end
        return vol
    end

    for dim ∈ 1:3
        tree = Tree(rand(dim), periodic=ones(Bool, dim))
        @test total_volume_via_cells(tree) == 1.0
        @test total_firstMoment_via_cells(tree) ≈ tree.position
        @test total_volume_via_faces(tree) == dim

        refine!(tree)
        @test total_volume_via_cells(tree) == 1.0
        @test total_firstMoment_via_cells(tree) ≈ tree.position
        @test total_volume_via_faces(tree) == dim

        refine!(tree, recurse=true)
        @test total_volume_via_cells(tree) == 1.0
        @test total_firstMoment_via_cells(tree) ≈ tree.position
        @test total_volume_via_faces(tree) == dim

        refine!(tree.children[end].children[1])
        @test total_volume_via_cells(tree) == 1.0
        @test total_firstMoment_via_cells(tree) ≈ tree.position
        @test total_volume_via_faces(tree) == dim

        refine!(tree, recurse=true)
        @test total_volume_via_cells(tree) == 1.0
        @test total_firstMoment_via_cells(tree) ≈ tree.position
        @test total_volume_via_faces(tree) == dim
    end
end


"""
$(TYPEDEF)

A surreal number form with explicit left and right subsets.

A (valid) form represents a specific surreal number, and so instances of this
class behave as the number they represent.
"""
struct SurrealForm{L<:SurrealSet, R<:SurrealSet} <: Surreal
	l::L
	r::R

	SurrealForm(l::SurrealSet, r::SurrealSet) = new{typeof(l), typeof(r)}(l, r)
end

SurrealForm() = SurrealForm((), ())
SurrealForm(f::SurrealForm) = f
SurrealForm(l, r) = SurrealForm(SurrealSet(l), SurrealSet(r))  # Convert arguments
SurrealForm(x) = SurrealForm(left(x), right(x))

function SurrealForm(args...; check::Bool)
	f = SurrealForm(args...)
	check && check_numeric(f)
	return f
end

left(f::SurrealForm) = f.l
right(f::SurrealForm) = f.r
issurreal(f::SurrealForm) = true

Base.zero(::SurrealForm) = SurrealForm()
Base.one(::SurrealForm) = SurrealForm((0,), ())

"""
	isnumeric(f::SurrealForm)

Whether the form is *numeric* - i.e. represents a valid surreal number.
"""
isnumeric(f::SurrealForm) = all_lt(left(f), right(f))

checknumeric(f::SurrealForm) = isnumeric(f) || error("Form is not numeric")

Base.show(io::IO, mime::Union{FORM_TEXT_MIMES...}, f::SurrealForm) = show_form_outer(io, mime, f)


function fold(r)
# [rw] = fold(r)

# Fold left wing of vector in "FFT buffer format"
# onto right wing
# J.O. Smith, 1982-2002

   sz = size(r);
   n = maximum(sz)
   if prod(sz) != n
     error("fold.m: input must be a vector");
   end

   if n < 3; return r;
   elseif mod(n,2)==1
       nt = Int((n+1)/2);
       rw = [ r[1], r[2:nt] + conj(r[n:-1:nt+1]), 0*ones(n-nt) ];
   else
       nt = Int(n/2);
       rf = cat(1, r[2:nt] , 0);
       rf = rf + conj(r[n:-1:nt+1]);
       rw = cat(1, r[1] , rf , 0*ones(n-nt-1) );
   end

   rw = reshape(rw,sz)
   return rw
end


#!/usr/bin/env julia

using TimerOutputs, GLMakie, FileIO

N=4096
NPx = N
NPy = convert(Int64,round(3*NPx/4))		# 3/4 aspect ratio
const xmin	= -2.0
const xmax 	=  1.0
const ymin 	= -1.25
const ymax	=  1.25

dx = (xmax-xmin)/NPx
dy = (ymax-ymin)/NPy

xs = [ xmin + (i-1)*dx for i in 1:NPx]
ys = [ ymin + (i-1)*dy for i in 1:NPy]

to = TimerOutput()
out_array = zeros(Float64,NPx,NPy)

function mb_iterate(i,j,dx,dy,xmin,ymin,Niter)
	z = c =  xmin + (i-1) * dx + im*(ymin + (j-1) * dy)
	j = 0.0
	for i in 1:Niter
		j += 1.0
		z_prime = z*z + c
		if abs2(z_prime) <= 4.0
			z = z_prime
		else
			z = j + 0*im
			break
		end
	end
	return abs(z)
end

# compile mb_iterate with 1 call, then time it on the array.
x = mb_iterate(1,1,dx,dy,xmin,ymin,60)
print("x = $x\n") # to defeat being optimized away

@timeit to "run iterations for loops" @inbounds Threads.@threads for i in 1:NPx
	@simd for j in 1:NPy
		out_array[i,j] = mb_iterate(i,j,dx,dy,xmin,ymin,60)
	end
end

jrange = 1:NPy
irange = 1:NPx
@timeit to "run iterations for loop + broadcast" Threads.@threads for j in 1:NPy
		out_array[irange,j] = mb_iterate.(irange,j,dx,dy,xmin,ymin,30)
end

@timeit to "create figure and save" FileIO.save("mbs.png",
					heatmap(xs,ys,out_array,colormap = Reverse(:deep));
					resolution = size(out_array))

show(to)
print("\n")


using Luxor, Test, Random
Random.seed!(42)

function polygon_intersection_test(fname)
    Drawing(2000, 2000, fname)
    origin()
    background("white")
    sethue("magenta")
    setline(.2)
    fontsize(4)
    tiles = Tiler(2000, 2000, 5, 5)
    for (pos, n) in tiles
        @layer begin
            randomhue()
            translate(pos)
            setline(1)
            p1 = randompointarray(BoundingBox(box(O, tiles.tilewidth, tiles.tileheight, vertices=true))..., 3)
            p2 = randompointarray(BoundingBox(box(O, tiles.tilewidth, tiles.tileheight, vertices=true))..., 3)

            sethue("grey80")
            poly(p1, close=true, :stroke)
            poly(p2, close=true, :stroke)

            sethue("red")
            prettypoly(p1, close=false, :stroke, vertexlabels = (n, l) -> (text(string(n))))

            sethue("blue")
            prettypoly(p2, close=false, :stroke, vertexlabels = (n, l) -> (text(string(n))))

            pint = polyintersect(p1, p2)
            sethue("purple")
            setopacity(0.5)
            circle.(pint, 6, :fill)

            pint = polyintersect(p1, p2, closed=false)
            sethue("blue")
            setopacity(0.8)
            circle.(pint, 10, :stroke)
        end
    end
    @test finish() == true
end

fname = "polygon_intersection_test.png"
polygon_intersection_test(fname)
println("...finished test: output in $(fname)")


# utils
export variables, operator, affine
import Base: convert
import AbstractOperators: displacement

convert(::Type{Expression},x::Variable{T,N,A}) where {T,N,A} =
Expression{1}((x,),Eye(T,size(x)))

"""
    variables(ex::Expression)

Returns a tuple containing the `Variable`s of expression `ex`.

# Example

```julia
julia> x,y = Variable(2),Variable(2);

julia> ex = x+y;

julia> variables(ex)
(Variable(Float64, (2,)), Variable(Float64, (2,)))

```

"""
variables(A::Expression)    = A.x
variables(x::Variable)    = x

"""
    operator(ex::Expression)

Returns the `AbstractOperator` of expression `ex`.

# Example

```julia
julia> x = Variable(3)
Variable(Float64, (3,))

julia> ex = fft(x);

julia> operator(ex)
ℱ  ℝ^3 -> ℂ^3

```
"""
operator(A::Expression) = remove_displacement(A.L)
operator(x::Variable)   = Eye(~x)

"""
    affine(ex::Expression)

Returns the `AbstractOperator` of expression `ex` keeping any affine addition.

"""
affine(A::Expression) = A.L
affine(x::Variable)   = Eye(~x)

"""
    displacement(ex::Expression)

Returns the displacement of expression `ex`.

# Example

```julia
julia> x = Variable(3)
Variable(Float64, (3,))

julia> ex = fft(x)+[1.+im*2.;0.;3.+im*4];

julia> displacement(ex)
3-element Array{Complex{Float64},1}:
1.0+2.0im
0.0+0.0im
3.0+4.0im

```
"""
displacement(A::Variable{T}) where {T} = zero(T)
displacement(E::Expression) = displacement(affine(E))


### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

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

	using MLJ
	using CSV
	using DataFrames
	using PlutoUI
	using Test
	using Printf
	using Random
 	using Plots  # using PyPlot
end

# ╔═╡ 638e7b78-7544-11eb-2a45-9544fcbeb7bc
## include utils functions
include("./utils.jl")

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

# ╔═╡ c5c92502-7190-11eb-3a75-87b7f156a7ce
md"""
In this notebook, we will run ridge regression multiple times with different L2 penalties to see which one produces the best fit. We will revisit the example of polynomial regression as a means to see the effect of L2 regularization. In particular, we will:
* Use a pre-built implementation of regression (MLJLinearModels) to run polynomial regression, this time with L2 penalty
* Use plot to visualize polynomial regressions under L2 regularization
* Choose best L2 penalty using cross-validation.
* Assess the final fit using test data.

We will continue to use the House data from previous notebooks.  (In the next programming assignment for this module, you will implement your own ridge regression learning algorithm using gradient descent.)

"""

# ╔═╡ 14d8375a-734e-11eb-2518-2d4eee6b9e3e
md"""
### Polynomial regression, revisited

We build on the material from Week 3, where we wrote the function to produce an SFrame with columns containing the powers of a given input. Copy and paste the function `polynomial_sframe` from Week 3:
"""

# ╔═╡ 55fb5a1e-734e-11eb-2019-21b125a173d1
function polynomial_df(feature; degree=3)
	@assert degree ≥ 1 "Expect degree to be ≥ 1"

	hsh = Dict{Symbol, Vector{Float64}}(:power_1 => feature)
	for deg ∈ 2:degree
		hsh[Symbol("power_$(deg)")] = feature .^ deg
	end

	return DataFrame(hsh)
end

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

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

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

# ╔═╡ 76617496-734e-11eb-0177-ef53e0a61ebc
sort!(sales, [:sqft_living, :price], rev=[false, false]);

# ╔═╡ a626fab6-734e-11eb-0994-f153e8264aaf
md"""
Let us revisit the 15th-order polynomial model using the `sqft_living` input. Generate polynomial features up to degree 15 using `polynomial_df()` and fit a model with these features.

When fitting the model, use an L2 penalty of 1e-5:
"""

# ╔═╡ 9f4449e2-734e-11eb-0d61-19d14b53358a
begin
	poly_df_0 = polynomial_df(sales.sqft_living; degree=15)
	poly_df_0[!, :price] = sales.price

	first(poly_df_0, 3)
end

# ╔═╡ e985d03e-734e-11eb-26a7-effe0c92f0fc
@load RidgeRegressor pkg=MLJLinearModels

# ╔═╡ ed0179fc-734e-11eb-3309-03b0ef4cdc96
begin
	l2_small_penalty = 1e-5

	mdl0 = MLJLinearModels.RidgeRegressor(lambda=l2_small_penalty,
		penalize_intercept=false)

	X_0 = select(poly_df_0, :power_1)
	y_0 = poly_df_0.price

	mach0 = machine(mdl0, X_0, y_0)
	fit!(mach0)
	fp0 = fitted_params(mach0)

	with_terminal() do
		for (name, c) in fp0.coefs
   			println("$(rpad(name, 10)):  $(round(c, sigdigits=3))")
		end

		println("Intercept: $(round(fp0.intercept, sigdigits=3))")
	end
end

# ╔═╡ f4ac831a-7351-11eb-10b2-67efef23c436
md"""

**Quiz Qestion:  What's the learned value for the coefficient of feature `power_1`?**
  - Answer below:
"""

# ╔═╡ 4686d706-74a8-11eb-101d-c976b98872f3
## power_1
round(fp0.coefs[1][2], sigdigits=2)

# ╔═╡ 3c18cc2e-74a9-11eb-336b-1bdf4159d5f8
getfield(fp0, :coefs)[1][2]

# ╔═╡ 1fe9cdec-7352-11eb-1336-d98c99d2a002
md"""
### Observe overfitting

Recall from Week 3 that the polynomial fit of degree 15 changed wildly whenever the data changed. In particular, when we split the sales data into four subsets and fit the model of degree 15, the result came out to be very different for each subset. The model had a *high variance*. We will see in a moment that ridge regression reduces such variance.

But first, we must reproduce the experiment we did in Week 3.
"""

# ╔═╡ 4835efa8-7352-11eb-27dd-23dcf67ead55
begin
	(sales_a, sales_b) = train_test_split(sales; split=0.5, seed=42)

	(set1, set2) = train_test_split(sales_a; split=0.5, seed=42)
	(set3, set4) = train_test_split(sales_b; split=0.5, seed=42)

	(size(set1), size(set2), size(set3), size(set4))
end

# ╔═╡ b3b02ed6-7352-11eb-3a15-9f392deb56d8
begin

function make_poly(tset; degree=15, output=:price)
  poly_df = polynomial_df(tset.sqft_living; degree)
  features = names(poly_df)
  poly_df[!, output] = tset[!, output]
  (features, poly_df)
end

function fit_poly(tset; degree=15, output=:price, l2_penalty=l2_small_penalty)
  (features, poly_df) = make_poly(tset; degree, output)
  mdl = MLJLinearModels.RidgeRegressor(lambda=l2_penalty,
		penalize_intercept=false)
  X_ = select(poly_df, features)
  y_ = poly_df[!, output]

  mach = machine(mdl, X_, y_)
  fit!(mach)

  (mach, X_, y_)
end

function print_coeff(mach, ix)
  fp = fitted_params(mach)
  println("\n- model $(ix)")
  for (name, c) in fp.coefs[1:1]  # print only first coeff
	println("$(rpad(name, 10)):  $(round(c, sigdigits=3))")
  end

  println("Intercept: $(round(fp.intercept, sigdigits=3))")
end

function find(machs; key=:min)
	fps = map(m -> fitted_params(m), machs) |>
		fps -> map(fp -> getfield(fp, :coefs)[1][2], fps)

	ix = key == :min ? argmin(fps) : argmax(fps)
	(fps[ix], string("model_", ix))
end

end

# ╔═╡ 6431250e-7354-11eb-1446-37f9f51cb959
md"""
##### Set1, Set2, Set3 & Set4
"""

# ╔═╡ 64f0b460-7353-11eb-23b5-0b2a0c9006f4
begin
	hsh = Dict{Symbol, Vector{Any}}(:machine => [], :tset => [])

	for (ix, s) ∈ enumerate((set1, set2, set3, set4))
		(mach_, X_, y_) = fit_poly(s)
		push!(hsh[:machine], mach_)
		push!(hsh[:tset], (X_, y_))
	end

	ps = []

	for ix ∈ 1:length(hsh[:machine])
		(X_, y_) = hsh[:tset][ix]
		mach_ = hsh[:machine][ix]

		p1 = scatter(X_.power_1, y_, legend=false, color=[:lightblue], marker=".")
 		p2 = scatter!(X_.power_1, predict(mach_, X_), color=[:orange], marker="-")
		push!(ps, (p1, p2))
	end

end

# ╔═╡ c7c9f6dc-7495-11eb-2bd5-8ba7f20f997a
with_terminal() do
	for ix in 1:length(hsh[:machine])
		print_coeff(hsh[:machine][ix], ix)
	end
end

# ╔═╡ e22363b4-7499-11eb-3b50-5905e5eef045
begin
	lg1 = grid(2, 2, widths=[0.0, 0.9], heights=[0.5, 0.5])
	plot(ps[1]..., ps[2]..., layout=lg1)
end

# ╔═╡ f08611a8-749a-11eb-1014-e9e29e34a966
begin
	lg2 = grid(2, 2, widths=[0.0, 0.9], heights=[0.5, 0.5])
	plot(ps[3]..., ps[4]..., layout=lg2)
end

# ╔═╡ e4c6e384-7354-11eb-307d-579569191332
md"""
The four curves should differ from one another a lot, as should the coefficients you learned.

**Quiz Question:  For the models learned in each of these training sets, what are the smallest and largest values you learned for the coefficient of feature `power_1`?**

  - cf. below
"""

# ╔═╡ ac0d2d96-74a8-11eb-079e-036266a3b75e
(smallest=find(hsh[:machine]), largest=find(hsh[:machine]; key=:max))

# ╔═╡ 40bd2ee6-7355-11eb-00ba-1943adaa3fb9
md"""
### Ridge regression comes to rescue

Generally, whenever we see weights change so much in response to change in data, we believe the variance of our estimate to be large. Ridge regression aims to address this issue by penalizing "large" weights. (Weights of `model15` looked quite small, but they are not that small because 'sqft_living' input is in the order of thousands.)

With the argument `l2_penalty=1e5`, fit a 15th-order polynomial model on `set1`, `set2`, `set3`, and `set4`. Other than the change in the `l2_penalty` parameter, the code should be the same as the experiment above.
"""

# ╔═╡ 5aa44ab0-7355-11eb-17cb-7bd692a55a37
md"""
#### Set1, Set2, Set3 & Set4
"""

# ╔═╡ 5aa4d1cc-7355-11eb-09a9-dd185f7871e1
l2_penalty=1e5

# ╔═╡ 96c23cd0-749c-11eb-0c80-63c4f283c2e5
begin
	hsh_v1r = Dict{Symbol, Vector{Any}}(:machine => [], :tset => [])

	for (ix, s) ∈ enumerate((set1, set2, set3, set4))
		(mach_, X_, y_) = fit_poly(s; l2_penalty)
		push!(hsh_v1r[:machine], mach_)
		push!(hsh_v1r[:tset], (X_, y_))
	end

	## plot prep.
	ps2 = []
	for ix ∈ 1:length(hsh[:machine])
		(X_, y_) = hsh[:tset][ix]
		mach_ = hsh[:machine][ix]

		p1 = scatter(X_.power_1, y_, legend=false, color=[:lightblue], marker=".")
 		p2 = scatter!(X_.power_1, predict(mach_, X_), color=[:orange], marker="-")
		push!(ps2, (p1, p2)) # plot(p1, p2, layout=l)
	end

	with_terminal() do
		for ix in 1:length(hsh_v1r[:machine])
			print_coeff(hsh_v1r[:machine][ix], ix)
		end
	end
end

# ╔═╡ 03e1ed82-74a8-11eb-140c-e31183a2314b
begin
	lg11 = grid(2, 2, widths=[0.0, 0.9], heights=[0.5, 0.5])
	plot(ps2[1]..., ps2[2]..., layout=lg11)
end

# ╔═╡ 034e5e96-74a8-11eb-030d-fb9f7ea9af55
begin
	lg12 = grid(2, 2, widths=[0.0, 0.9], heights=[0.5, 0.5])
	plot(ps2[3]..., ps2[4]..., layout=lg12)
end

# ╔═╡ d32b7d62-7355-11eb-00f5-e15ead474ce5
md"""
These curves should vary a lot less, now that you applied a high degree of regularization.

**Quiz Question:  For the models learned with the high level of regularization in each of these training sets, what are the smallest and largest values you learned for the coefficient of feature `power_1`?**

  - cf. below
"""

# ╔═╡ a96763a2-74aa-11eb-2a92-87ee73d05221
(smallest=find(hsh_v1r[:machine]), largest=find(hsh_v1r[:machine]; key=:max))

# ╔═╡ 8e9ed252-7357-11eb-117c-91f7ff585a2d
md"""
### Selecting an L2 penalty via cross-validation

Just like the polynomial degree, the L2 penalty is a "magic" parameter we need to select. We could use the validation set approach as we did in the last module, but that approach has a major disadvantage: it leaves fewer observations available for training. **Cross-validation** seeks to overcome this issue by using all of the training set in a smart way.

We will implement a kind of cross-validation called **k-fold cross-validation**. The method gets its name because it involves dividing the training set into k segments of roughtly equal size. Similar to the validation set method, we measure the validation error with one of the segments designated as the validation set. The major difference is that we repeat the process k times as follows:

Set aside segment 0 as the validation set, and fit a model on rest of data, and evalutate it on this validation set

Set aside segment 1 as the validation set, and fit a model on rest of data, and evalutate it on this validation set

...

Set aside segment k-1 as the validation set, and fit a model on rest of data, and evalutate it on this validation set

After this process, we compute the average of the k validation errors, and use it as an estimate of the generalization error. Notice that  all observations are used for both training and validation, as we iterate over segments of data.

To estimate the generalization error well, it is crucial to shuffle the training data before dividing them into segments.

"""

# ╔═╡ 91a058b6-7359-11eb-1da5-eba5d339d100
function get_rss(mach, X, y)
    ŷ = predict(mach, X)     # First get the predictions
    diff = y .- ŷ            # Then compute the residuals/errors
    rss = sum(diff .* diff)  # Then square and add them up
    return rss
end

# ╔═╡ 957b90fe-7359-11eb-23fc-039055cd2914
function k_fold_cross_validation(k, l2_penalty, df, output, features;
	degree=15)
	n = size(df)[1]
	rss = []
	for ix ∈ 1:k
		start_ = (n * (ix - 1)) ÷ k + 1
		end_ = (n * ix) ÷ k

		## get the partition train/valid
		valid_set = df[start_:end_, :]
		train_set = vcat(df[1:start_ - 1, :], df[end_+1:end, :])

		##
		mdl = MLJLinearModels.RidgeRegressor(lambda=l2_penalty,
			penalize_intercept=false)
  		X_ = select(train_set, features)
  		y_ = train_set[!, output]

  		mach = machine(mdl, X_, y_)
  		fit!(mach)

		##
		push!(rss, get_rss(mach, valid_set, valid_set[!, output]))
	end
	sum(rss) / k
end

# ╔═╡ 9af5697a-73af-11eb-1558-dbc635b07638
function find_best_l2(train_valid; k=10, degree=15, output=:price)
	l2_rss, l2_p = [], []
	best_l2, min_rss = nothing, nothing
	(features, poly_set) = make_poly(train_valid; degree, output)

	for cur_l2p ∈ (10^ix for ix ∈ 1:0.5:7)
		push!(l2_p, cur_l2p)
		rss = k_fold_cross_validation(k, cur_l2p, poly_set, output, features)
		push!(l2_rss, rss)

		# @printf("current rss: %2.5e / current penalty: %2.5f\n", rss, cur_l2p)
		if isnothing(min_rss) || rss < min_rss
			min_rss = rss
			best_l2 = cur_l2p
		end
	end
	(best_l2, min_rss, l2_p, l2_rss)
end

# ╔═╡ 51150d2e-7358-11eb-2d17-c5c6898b7045
begin
	# train_valid is shuffled during the train-test_split
	(train_valid, test) = train_test_split(sales; split=0.9, seed=42)

	k = 10
	degree = 15
	output=:price

	(best_l2, min_rss, l2_p, l2_rss) = find_best_l2(train_valid; k, degree, output)
	with_terminal() do
		@printf("min rss: %2.5e / best penalty: %2.5e\n", min_rss, best_l2)
	end
end

# ╔═╡ daa7efec-73b0-11eb-3657-d520436222d5
plot(l2_p, l2_rss, marker="-", color=[:lightblue], xscale=:log)

# ╔═╡ b273d2d6-73b1-11eb-06df-1559778e561d
md"""
Once you found the best value for the L2 penalty using cross-validation, it is important to retrain a final model on all of the training data using this value of l2_penalty.

This way, your final model will be trained on the entire dataset.
"""

# ╔═╡ 199f8240-73b2-11eb-2753-ab74e806ab9f
begin
	(features, train_set) = make_poly(train_valid; degree, output)

	mdl = MLJLinearModels.RidgeRegressor(lambda=best_l2,
			penalize_intercept=false)
  		X_ = select(train_set, features)
  		y_ = train_set[!, output]

  		mach = machine(mdl, X_, y_)
  		fit!(mach)

	(_, test_set) = make_poly(test; degree, output)

	final_rss = get_rss(mach, test_set, test_set[!, output])

	with_terminal() do
		@printf("final rss: %2.5e\n", final_rss)
	end
end

# ╔═╡ 3db30566-73b3-11eb-1d4e-f919ad63641a
md"""
**Quiz Question: Using the best L2 penalty found above, train a model using all *training* data. What is the RSS on the *test* data of the model you learn with this L2 penalty?**

 - Final rss: 4.03826e+14
"""

# ╔═╡ Cell order:
# ╟─87aaecb0-7190-11eb-1572-d9110d9ffaaa
# ╟─c5c92502-7190-11eb-3a75-87b7f156a7ce
# ╠═5bf16f26-7191-11eb-1166-73a483cf120a
# ╟─14d8375a-734e-11eb-2518-2d4eee6b9e3e
# ╠═55fb5a1e-734e-11eb-2019-21b125a173d1
# ╟─7a50bd00-7191-11eb-2ed6-071e6db68422
# ╠═890cb362-7191-11eb-251b-e155ed2d21f2
# ╠═638e7b78-7544-11eb-2a45-9544fcbeb7bc
# ╟─b43046d2-7191-11eb-0978-b1bfe1762756
# ╠═76617496-734e-11eb-0177-ef53e0a61ebc
# ╟─a626fab6-734e-11eb-0994-f153e8264aaf
# ╠═9f4449e2-734e-11eb-0d61-19d14b53358a
# ╠═e985d03e-734e-11eb-26a7-effe0c92f0fc
# ╠═ed0179fc-734e-11eb-3309-03b0ef4cdc96
# ╟─f4ac831a-7351-11eb-10b2-67efef23c436
# ╠═4686d706-74a8-11eb-101d-c976b98872f3
# ╠═3c18cc2e-74a9-11eb-336b-1bdf4159d5f8
# ╟─1fe9cdec-7352-11eb-1336-d98c99d2a002
# ╠═4835efa8-7352-11eb-27dd-23dcf67ead55
# ╠═b3b02ed6-7352-11eb-3a15-9f392deb56d8
# ╟─6431250e-7354-11eb-1446-37f9f51cb959
# ╠═64f0b460-7353-11eb-23b5-0b2a0c9006f4
# ╠═c7c9f6dc-7495-11eb-2bd5-8ba7f20f997a
# ╠═e22363b4-7499-11eb-3b50-5905e5eef045
# ╠═f08611a8-749a-11eb-1014-e9e29e34a966
# ╟─e4c6e384-7354-11eb-307d-579569191332
# ╠═ac0d2d96-74a8-11eb-079e-036266a3b75e
# ╟─40bd2ee6-7355-11eb-00ba-1943adaa3fb9
# ╟─5aa44ab0-7355-11eb-17cb-7bd692a55a37
# ╠═5aa4d1cc-7355-11eb-09a9-dd185f7871e1
# ╠═96c23cd0-749c-11eb-0c80-63c4f283c2e5
# ╠═03e1ed82-74a8-11eb-140c-e31183a2314b
# ╠═034e5e96-74a8-11eb-030d-fb9f7ea9af55
# ╟─d32b7d62-7355-11eb-00f5-e15ead474ce5
# ╠═a96763a2-74aa-11eb-2a92-87ee73d05221
# ╟─8e9ed252-7357-11eb-117c-91f7ff585a2d
# ╠═91a058b6-7359-11eb-1da5-eba5d339d100
# ╠═957b90fe-7359-11eb-23fc-039055cd2914
# ╠═9af5697a-73af-11eb-1558-dbc635b07638
# ╠═51150d2e-7358-11eb-2d17-c5c6898b7045
# ╠═daa7efec-73b0-11eb-3657-d520436222d5
# ╟─b273d2d6-73b1-11eb-06df-1559778e561d
# ╠═199f8240-73b2-11eb-2753-ab74e806ab9f
# ╟─3db30566-73b3-11eb-1d4e-f919ad63641a


#=
image_geom.jl
Methods related to an image geometry for image reconstruction
2017-10-02 Samuel Rohrer, University of Michigan
2019-03-05 Jeff Fessler, Julia 1.1 + tests
2019-06-23 Jeff Fessler, overhaul
=#

export MIRT_image_geom, image_geom, cbct
export image_geom_circle
export image_geom_ellipse

#using MIRT: jim, downsample2, downsample3
using Test: @test, @test_throws, @inferred
using ImageTransformations: imresize

"""
`MIRT_image_geom`

Image geometry "struct" with essential parameters
"""
struct MIRT_image_geom
	# options for 2D image geometry
	nx::Int				# image dimension 1
	ny::Int				# image dimension 2
	dx::Float32			# pixel size
	dy::Float32			# pixel size
	offset_x::Float32	# unitless
	offset_y::Float32	# unitless

	# options for 3D image geometry
	nz::Int				# image dimension 3
	dz::Float32			# voxel size
	offset_z::Float32	# unitless

	mask::Array{Bool}	# logical mask
end

"""
`image_geom_help( ; io)`
"""
function image_geom_help( ; io::IO = isinteractive() ? stdout : IOBuffer())
	print(io, "propertynames:\n")
	print(io, propertynames(image_geom(nx=1, dx=1)))

	print(io,
	"
	Derived values for 2D (and 3D) images:

	is3	is it 3d (nz > 0?)
	fovs [|dx|*nx |dy|*ny ...]
	np	sum(mask) = # pixels to be estimated

	dim	(nx, ny, [nz])
	x	1D x coordinates of each pixel
	y	1D y coordiantes of each pixel
	wx	(nx - 1/2) * dx + offset_x
	wy	(ny - 1/2) * dy + offset_y
	wz	(nz - 1/2) * dz + offset_z

	xg	x coordinates of each pixel as a grid (2D or 3D)
	yg	y coordinates of each pixel as a grid (2D or 3D)

	u	1D frequency domain coordinates of each pixel (kx)
	v	1D frequency domain coordinates of each pixel (ky)
	w	1D frequency domain coordinates of each pixel (kz)

	ug	2D or 3D grid of frequency domain coordinates
	vg	2D or 3D grid of frequency domain coordinates
	fg	NamedTuple of 2D or 3D frequency coordinates
	ones	ones(dim)
	zeros	zeros(dim)

	mask_outline	binary image showing 2D outline of mask

	Derived values for 3D images:

	z	z coordinates of each pixel
	zg	z coordinates of each pixel as a grid (3D only)
	wg	3D grid of frequency domain coordinates

	mask_or	[nx ny] logical 'or' of mask

	Methods:

	embed(x)	turn short column(s) into array(s)
	maskit(x)	opposite of embed
	shape(x)	reshape long column x to [nx ny [nz]] array
	unitv(...)	j=j | i=(ix,iy) | c=[cx cy]
				j: single index from 1 to length(z)
				i: [ix,iy[,iz]] index from 1 to nx,ny
				c: [cx,cy[,cz]] index from +/- n/2 center at floor(n/2)+1
	circ(rx=,ry=,cx=,cy=)	circle of given radius and center (cylinder in 3D)
	plot()		plot the image geometry

	Methods that return a new MIRT_image_geom:

	down(down::Int)		down-sample geometry by given factor
	over(over::Int)		over-sample geometry by given factor
	expand_nz(nz_pad)	expand image geometry in z by nz_pad on both ends
	\n")
end

"""
`MIRT_cbct_ig`

Structure suitable for passing to C routines `cbct_*`
based on the struct `cbct_ig` found in `cbct,def.h`
"""
struct MIRT_cbct_ig
	nx::Cint
	ny::Cint
	nz::Cint
	dx::Cfloat
	dy::Cfloat
	dz::Cfloat
	offset_x::Cfloat
	offset_y::Cfloat
	offset_z::Cfloat
	mask2::Ptr{Cuchar}	# [nx,ny] 2D support mask: 0 or 1 .. nthread
	iy_start::Ptr{Cint} # [nthread] for mask2
	iy_end::Ptr{Cint}	# [nthread] for mask2
end

"""
`cbct(ig::MIRT_image_geom; nthread=1)`
constructor for `MIRT_cbct_ig`
"""
function cbct(ig::MIRT_image_geom; nthread::Int=1)
	iy_start = [0]
	iy_end = [ig.ny]
	nthread != 1 && throw("only nthread=1 for now due to iy_start")
	return MIRT_cbct_ig(
		Cint(ig.nx), Cint(ig.ny), Cint(ig.nz), Cfloat(ig.dx),
		Cfloat(ig.dy), Cfloat(ig.dz),
		Cfloat(ig.offset_x), Cfloat(ig.offset_y), Cfloat(ig.offset_z),
			pointer(UInt8.(ig.mask_or)),
			pointer(Cint.(iy_start)), pointer(Cint.(iy_end)),
		)
end

# see https://docs.julialang.org/en/stable/manual/constructors/
"""
`ig = image_geom(...)`

Constructor for `MIRT_image_geom`

# Arguments
- `nx::Int = 128`
- `ny::Int = nx`
- `dx::Real = ?` (must specify one of `dx` or `fov`)
- `dy::Real = -dx`
- `offset_x::Real = 0` (unitless)
- `offset_y::Real = 0` (unitless)
- `fov::Real = ?` (if specified, then `nx*dx=ny*dy`)
- `nz::Int = 0`
- `dz::Real = ?` (need one of `dz` or `zfov` if `nz > 0`)
- `zfov::Real = ?` (if specified, then `nz*dz`)
- `offset_z::Real = 0` (unitless)
- `offsets::Symbol = :none` or :dsp
- `mask::Union{Symbol,AbstractArray{Bool}} = :all` | `:circ` | `:all_but_edge`
"""
function image_geom( ;
		nx::Int = 128,
		ny::Int = nx,
		dx::Real = NaN,
		dy::Real = NaN,
		offset_x::Real = 0,
		offset_y::Real = 0,
		fov::Real = NaN,
		nz::Int = 0,
		dz::Real = NaN,
		zfov::Real = NaN,
		offset_z::Real = 0,
		offsets::Symbol = :none,
		mask::Union{Symbol,AbstractArray{Bool}}	= :all,
	)

	# handle optional arguments

	is3 = !isnan(dz) || !isnan(zfov) || (nz > 0)

	# offsets
	if offsets === :dsp
		# check that no nonzero offset was specified
		(offset_x != 0 || offset_y != 0) && throw("offsets usage incorrect")
		offset_x = 0.5
		offset_y = 0.5
		if is3
			offset_z != 0 && throw("offset_z usage incorrect")
			offset_z = 0.5
		end
	elseif offsets != :none
		throw("offsets $offsets")
	end

	# transverse plane (2d) distances
	if true
		if isnan(fov)
			isnan(dx) && throw("dx or fov required")
		#	fov = nx * dx
		else
			!isnan(dx) && throw("dx and fov?")
			dx = fov / nx
			dy = fov / ny
		end

		dy = isnan(dy) ? -dx : dy # default dy
	end

	# 3D geometry
	if is3
		if isnan(zfov)
			isnan(dz) && throw("dz or zfov required")
		#	zfov = nz * dz
		else
			!isnan(dz) && throw("dz and zfov?")
			dz = zfov / nz
		end
	end

	# mask
	if mask === :all
		mask = is3 ? trues(nx, ny, nz) : trues(nx, ny)
	elseif mask === :circ
		mask = image_geom_circle(nx, ny, dx, dy)
	elseif mask === :all_but_edge
		mask = trues(nx,ny,nz)
		mask[  1,   :, :] .= false
		mask[end,   :, :] .= false
		mask[  :,   1, :] .= false
		mask[  :, end, :] .= false
	elseif isa(mask, Symbol)
		throw("mask symbol $mask")
	elseif size(mask,1) != nx || size(mask,2) != ny || (is3 && size(mask,3) != nz)
		throw("mask size $(size(mask)), nx=$nx ny=$ny nz=$nz")
	end

	# return the object
	return MIRT_image_geom(nx, ny, dx, dy, offset_x, offset_y,
		nz, dz, offset_z, mask)
end

"""
`image_geom_down_round()`
helper function needed to downsample `image_geom`
"""
function _down_round(val::Real=1, dd::Real=1, down::Real=1)
	out = val / down
	if out != round(out)
		out = 2 * round(out / 2) # keep it even if not an integer
	end
	dd = dd * down
	return Int(out), dd
end

"""
`ig_down = downsample(ig, down::Union{Int,Vector{Int}})`
down sample an image geometry by the factor `down`
cf `image_geom_downsample`
"""
function downsample(ig::MIRT_image_geom, down::Union{Int,Vector{Int}})

	# check if it is a scalar by seeing if the collection of size is empty
	if isempty(size(down))
		downv = ig.is3 ? [down,down,down] : downv = [down,down]
	else
		downv = down
	end

	# call the down round function
	down_nx, down_dx = _down_round(ig.nx, ig.dx, downv[1])
	down_ny, down_dy = _down_round(ig.ny, ig.dy, downv[2])
	# adjust to "pixel" units
	down_offset_x = ig.offset_x / downv[1]
	down_offset_y = ig.offset_y / downv[2]

	if ig.is3 # 3d case
		down_nz, down_dz = _down_round(ig.nz, ig.dz, downv[3])
		down_offset_z = ig.offset_z / downv[3]
	else
		down_nz = 0; down_dz = 0; down_offset_z = 0
	end

	# carefully down-sample the mask
	mdim_vec = [size(ig.mask)...] # tuple to vector
	if ig.is3
		if [down_nx, down_ny, down_nz] .* downv == mdim_vec
			down_mask = downsample3(ig.mask, downv) .> 0
		else
			throw("bug: bad mask size. need to address mask downsampling")
		end
	else
		if [down_nx,down_ny] .* downv == mdim_vec
			down_mask = downsample2(ig.mask, downv) .> 0
		else
			down_mask = imresize(ig.mask, (down_nx,down_ny)...) .> 0
		end
	end

	return MIRT_image_geom(
		down_nx, down_ny, down_dx, down_dy, down_offset_x, down_offset_y,
		down_nz, down_dz, down_offset_z, down_mask)
end

"""
`ig_new = image_geom_expand_nz(ig::MIRT_image_geom, nz_pad::Int)`
pad both ends
"""
function image_geom_expand_nz(ig::MIRT_image_geom, nz_pad::Int)
	!ig.is3 && throw("expand_nz only valid for 3D")
	out_nz = ig.nz + 2*nz_pad
	out_mask = cat(dims=3, repeat(ig.mask[:,:,1], 1, 1, nz_pad), ig.mask,
			repeat(ig.mask[:,:,end], 1, 1, nz_pad))
	return MIRT_image_geom(
		ig.nx, ig.ny, ig.dx, ig.dy, ig.offset_x, ig.offset_y,
		out_nz, ig.dz, ig.offset_z, out_mask)
end

"""
`ig_over = image_geom_over(ig::MIRT_image_geom, over::Int)`
over-sample an image geometry by the factor `over`
"""
function image_geom_over(ig::MIRT_image_geom, over::Int)
	if all(ig.mask .== true)
		mask_over = trues(ig.dim)
	else
		mask_over = imresize(ig.mask, ig.dim .* over) .> 0
	end
	return MIRT_image_geom(
		ig.nx*over, ig.ny*over, ig.dx/over, ig.dy/over,
		ig.offset_x*over, ig.offset_y*over,
		ig.nz*over, ig.dz/over, ig.offset_z*over,
		mask_over)
end

# ellipse that just inscribes the rectangle
# but keeping a 1 pixel border due to ASPIRE regularization restriction
function image_geom_ellipse(nx::Int, ny::Int, dx::Real, dy::Real ;
		rx::Real = min(abs((nx/2-1)*dx), abs((ny/2-1)*dy)),
		ry::Real = min(abs((nx/2-1)*dx), abs((ny/2-1)*dy)),
		cx::Real = 0, cy::Real = 0, over::Int=2)
	ig = image_geom(nx=nx, ny=ny, dx=dx, dy=dy)
	circ = ellipse_im(ig, [cx cy rx ry 0 1], oversample=over) .> 0
	return circ
end

# default is a circle that just inscribes the square
# but keeping a 1 pixel border due to ASPIRE regularization restriction
function image_geom_circle(nx::Int, ny::Int, dx::Real, dy::Real ;
		rx::Real = min(abs((nx/2-1)*dx), abs((ny/2-1)*dy)),
		ry::Real = rx, cx::Real = 0, cy::Real = 0, nz::Int=0, over::Int=2)
	ig = image_geom(nx=nx, ny=ny, dx=dx, dy=dy)
	circ = ellipse_im(ig, [cx cy rx ry 0 1], oversample=over) .> 0
	if nz > 0
		circ = repeat(circ, 1, 1, nz)
	end
	return circ
end

"""
`out = image_geom_add_unitv(z::AbstractArray ; j=?, i=?, c=?)`

add a unit vector to an initial array `z` (typically of zeros)

# Arguments
- `j` single index from 1 to length(z)
- `i` [ix,iy[,iz]] index from 1 to nx,ny
- `c` [cx,cy[,cz]] index from +/- n/2 center at floor(n/2)+1

default with no arguments gives unit vector at center `c=[0,0]`
"""
function image_geom_add_unitv(
		z::AbstractArray{T} ; # starts with zeros()
		j::Int = 0,
		i::AbstractVector{Int} = zeros(Int, ndims(z)),
		c::AbstractVector{Int} = zeros(Int, ndims(z))
	) where {T <: Number}

	out = copy(z)

	if j > 0 && all(i .== 0)
		out[j] += one(T)
	elseif j == 0 && any(i .> 0)
		out[i...] += one(T)
	else
		tmp = c .+ Int.(floor.(size(out) ./ 2)) .+ 1
		out[tmp...] += one(T)
	end

	return out
end

"""
`image_geom_add_unitv_test()`
"""
function image_geom_add_unitv_test()
#	test != :test && throw(ArgumentError("test $test"))
	image_geom_add_unitv(zeros(3,4), j=2)
	image_geom_add_unitv(zeros(3,4), i=[2,3])
	image_geom_add_unitv(zeros(3,4))
	image_geom_add_unitv(zeros(3,4), c=[-1,-2])
	true
end

"""
`image_geom_plot(ig)`
"""
function image_geom_plot(ig::MIRT_image_geom; kwargs...)
	return ig.is3 ?
		jim(ig.x, ig.y, ig.mask_or,
			"(dx,dy,dz)=$(ig.dx),$(ig.dy),$(ig.dz)"; kwargs...) :
		jim(ig.x, ig.y, ig.mask, "(nx,ny)=$(ig.nx),$(ig.ny)"; kwargs...)
end

function Base.display(ig::MIRT_image_geom)
	ir_dump(ig)
end

# Extended properties

image_geom_fun0 = Dict([
	(:help, ig -> image_geom_help()),

	(:is3, ig -> ig.nz > 0),
	(:dim, ig -> ig.is3 ? (ig.nx, ig.ny, ig.nz) : (ig.nx, ig.ny)),
	(:fovs, ig -> ig.is3 ?
		[abs(ig.dx)*ig.nx, abs(ig.dy)*ig.ny, abs(ig.dz)*ig.nz] :
			[abs(ig.dx)*ig.nx, abs(ig.dy)*ig.ny]),

	(:zeros, ig -> zeros(Float32, ig.dim)), # (nx, ny, nz) : zeros(nx, ny)
	(:ones, ig -> ones(Float32, ig.dim)), # (nx, ny, nz) : zeros(nx, ny)

	(:x, ig -> ((0:ig.nx-1) .- ig.wx) * ig.dx),
	(:y, ig -> ((0:ig.ny-1) .- ig.wy) * ig.dy),
	(:z, ig -> ((0:ig.nz-1) .- ig.wz) * ig.dz),

	(:wx, ig -> (ig.nx - 1)/2 + ig.offset_x),
	(:wy, ig -> (ig.ny - 1)/2 + ig.offset_y),
	(:wz, ig -> (ig.nz - 1)/2 + ig.offset_z),

	(:xg, ig -> ig.is3 ? repeat(ig.x, 1, ig.ny, ig.nz) :
						repeat(ig.x, 1, ig.ny)),
	(:yg, ig -> ig.is3 ? repeat(ig.y', ig.nx, 1, ig.nz) :
						repeat(ig.y', ig.nx, 1)),
	(:zg, ig -> ig.is3 ? repeat(reshape(ig.z, 1, 1, ig.nz), ig.nx, ig.ny, 1) :
						zeros(ig.nx, ig.ny)),

	# DFT frequency sample grid
	(:u, ig -> (-ig.nx/2:(ig.nx/2-1)) / (ig.nx*ig.dx)),
	(:v, ig -> (-ig.ny/2:(ig.ny/2-1)) / (ig.ny*ig.dy)),
	(:w, ig -> (-ig.nz/2:(ig.nz/2-1)) / (ig.nz*ig.dz)),

	(:ug, ig -> ig.is3 ? repeat(ig.u, 1, ig.ny, ig.nz) :
						repeat(ig.u, 1, ig.ny)),
	(:vg, ig -> ig.is3 ? repeat(ig.v', ig.nx, 1, ig.nz) :
						repeat(ig.v', ig.nx, 1)),
	(:wg, ig -> ig.is3 ? repeat(reshape(ig.w, 1, 1, ig.nz), ig.nx, ig.ny, 1) :
						zeros(ig.nx, ig.ny)),
	(:fg, ig -> ig.is3 ? (ig.ug, ig.vg, ig.wg) : (ig.ug, ig.vg)),

	(:np, ig -> sum(ig.mask)),
	(:mask_or, ig -> mask_or(ig.mask)),
	(:mask_outline, ig -> mask_outline(ig.mask)),

	# simple functions

	(:plot, ig -> ((;kwargs...) -> image_geom_plot(ig; kwargs...))),
	(:embed, ig -> (x::AbstractArray{<:Number} -> embed(x, ig.mask))),
	(:maskit, ig -> (x::AbstractArray{<:Number} -> maskit(x, ig.mask))),
	(:shape, ig -> (x::AbstractArray{<:Number} -> reshape(x, ig.dim))),
	(:unitv, ig -> ((;kwargs...) -> image_geom_add_unitv(ig.zeros; kwargs...))),
	(:circ, ig -> ((;kwargs...) ->
		image_geom_circle(ig.nx, ig.ny, ig.dx, ig.dy, nz=ig.nz; kwargs...))),

	# functions that return new geometry:

	(:down, ig -> (down::Int -> downsample(ig, down))),
	(:over, ig -> (over::Int -> image_geom_over(ig, over))),
	(:expand_nz, ig -> (nz_pad::Int -> image_geom_expand_nz(ig, nz_pad))),

	])

# Tricky overloading here!

Base.getproperty(ig::MIRT_image_geom, s::Symbol) =
		haskey(image_geom_fun0, s) ? image_geom_fun0[s](ig) :
		getfield(ig, s)

Base.propertynames(ig::MIRT_image_geom) =
	(fieldnames(typeof(ig))..., keys(image_geom_fun0)...)

function image_geom_test2(ig::MIRT_image_geom)
	# test 2D functions provided by the constructor
	ig.dim
	ig.x
	ig.y
	ig.wx
	ig.wy
	ig.xg
	ig.yg
	ig.fovs
	ig.np
	ig.mask_outline
	ig.ones
	ig.zeros
	ig.u
	ig.v
	ig.ug
	ig.vg
	ig.fg
	@test ig.shape(ig.ones[:]) == ig.ones
	@test ig.embed(ig.ones[ig.mask]) == Float32.(ig.mask)
	@test ig.maskit(ig.ones) == Float32.(ones(ig.np))
	ig.unitv(j=4)
	ig.unitv(i=ones(Int, length(ig.dim)))
	ig.unitv(c=zeros(Int, length(ig.dim)))
#= todo-i: why do these fail?
	@inferred image_geom_ellipse(8, 10, 1, 2)
	@inferred ig.circ()
	@inferred ig.plot()
	@inferred ig.unitv()
=#
	ig.circ()
	ig.plot()
	ig.unitv()
	@inferred ig.down(2)
	@inferred ig.over(2)
	image_geom_ellipse(8, 10, 1, 2)
	true
end

function image_geom_test2()
	@inferred image_geom(nx=16, dx=2, offsets=:dsp, mask=:all_but_edge)
	@test_throws String image_geom(nx=16, dx=1, offsets=:bad)
	@test_throws String image_geom(nx=16, dx=1, mask=:bad) # mask type
	@test_throws String image_geom(nx=16, dx=1, mask=trues(2,2)) # mask size
	ig = image_geom(nx=16, dx=2)
	display(ig)
	image_geom_test2(ig)
	ig = image_geom(nx=16, dx=2, mask=:circ)
	ig.over(2)
	ig.down(3) # test both even and non-even factors
	ig.help
	true
end

function image_geom_test3(ig::MIRT_image_geom)
	@test image_geom_test2(ig)
	ig.wz
	ig.zg
	ig.mask_or
	@inferred ig.expand_nz(2)
	@inferred cbct(ig)
	true
end

function image_geom_test3()
	@inferred image_geom(nx=16, nz=4, dx=2, zfov=1)
	ig = image_geom(nx=16, nz=4, dx=2, dz=3)
	image_geom_test3(ig)
	true
end

"""
`image_geom(:test)`
self test
"""
function image_geom(test::Symbol)
	if test === :help
		image_geom_help()
		return true
	end
	test != :test && throw(ArgumentError("test $test"))
	@test image_geom_add_unitv_test()
	@test _down_round(4, 3, 2)[1] == 2
	@test image_geom_test2()
	@test image_geom_test3()
	@test image_geom(:help)
	true
end


##############################################################################
# predict

# if x happens to be a scalar, do we want this to work?
# predict(M::Model, x::Number,            theta=thetaopt(M))   = [x*theta]
# "predict" could be called "predict_y_from_x"
# for each record, x,y,u,v are always vectors
predict(M::Model, x::Array{Float64, 1}, theta=thetaopt(M))   = [dot(x, theta)]
# following could be defined using rowwise, but efficiency might dictate otherwise
predict(M::Model, X::Array{Float64, 2}, theta=thetaopt(M))   = X*theta

"""`predict_y_from_test(M [, theta])`

Allows you compute embedded predictions (i.e., y values) based on a trained ERM model M on test data. Option to
specify a choice of theta. It defaults to `theta=thetaopt(M)`
"""
predict_y_from_test(M::Model,                   theta=thetaopt(M))   = predict(M, Xtest(M), theta)

"""`predict_y_from_train(M [, theta])`

Allows you compute embedded predictions (i.e., y values)  based on a trained ERM model M on train data. Option to
specify a choice of theta. It defaults to `theta=thetaopt(M)`
"""
predict_y_from_train(M::Model,                  theta=thetaopt(M))   = predict(M, Xtrain(M), theta)

"""`predict_v_from_test(M [, theta])`

Allows you compute unembedded predictions (i.e., in V space)
based on a trained ERM model M on test data. Option to
specify a choice of theta. It defaults to `theta=thetaopt(M)`
"""
predict_v_from_test(M::Model,                   theta=thetaopt(M))   = unembedY(M.S, predict(M, Xtest(M), theta))

"""`predict_v_from_train(M [, theta])`

Allows you compute unembedded predictions (i.e., in V space)
based on a trained ERM model M on train data. Option to
specify a choice of theta. It defaults to `theta=thetaopt(M)`
"""
predict_v_from_train(M::Model,                  theta=thetaopt(M))   = unembedY(M.S, predict(M, Xtrain(M), theta))

"""`predict_y_from_u(M, U [, theta])`

Allows you compute embedded predictions (i.e., y values)
based on a trained ERM model M on one or many raw inputs, `U`. Option to
specify a choice of theta. It defaults to `theta=thetaopt(M)`
"""
predict_y_from_u(M::Model, u::Array{T,1}, theta=thetaopt(M)) where {T<:Any} =  predict(M::Model, embedU(M.S, u), theta)
predict_y_from_u(M::Model, U::Array{T,2}, theta=thetaopt(M)) where {T<:Any} =  rowwise(u -> predict_y_from_u(M, u, theta), U)

"""`predict_y_from_u(M, U [, theta])`

Allows you compute unembedded predictions (i.e., in V space)
based on a trained ERM model M on one or many raw inputs, `U`. Option to
specify a choice of theta. It defaults to `theta=thetaopt(M)`
"""
predict_v_from_u(M::Model, u::Array{T,1}, theta=thetaopt(M)) where {T<:Any} =  unembedY(M.S, predict_y_from_u(M, u, theta))
predict_v_from_u(M::Model, U::Array{T,2}, theta=thetaopt(M)) where {T<:Any} =  rowwise(u -> predict_v_from_u(M, u, theta), U)

function uniquevalues(x::Array{T,1}) where {T<:Any}
    n = length(x)
    valtonum = Dict()
    vals = Any[]
    num = 1
    possiblevalues = sort(collect(Set(x)))
    for val in possiblevalues
        valtonum[val] = num
        push!(vals, val)
        num += 1
    end
    return valtonum, vals
end

uniquevalues(x::Array{T,2}) where {T<:Any} = uniquevalues(x[:])

function confusionx(vhat, v)
    K = length(Set(v))
    C = zeros(K,K)
    n = size(v,1)
    valtonum, vals = uniquevalues(v)
    for k=1:n
        i = convert(Int64, vhat[k])
        j = convert(Int64, v[k])
        C[valtonum[i], valtonum[j]] += 1
    end
    return C
end

"""`confusion_train(M)`

Returns the confusion matrix based on training data
for a trained ERM model"""
function confusion_train(M::Model)
    vhat = predict_v_from_train(M)
    return confusionx(vhat, Vtrain(M))
end

"""`confusion_test(M)`

Returns the confusion matrix based on test data
for a trained ERM model"""
function confusion_test(M::Model)
    vhat = predict_v_from_test(M)
    return confusionx(vhat, Vtest(M))
end


@testset "Multiroute flow" begin
    #### Graphs for testing
    graphs = [
    # Graph with 8 vertices
    (8,
    [
    (1, 2, 10), (1, 3, 5),  (1, 4, 15), (2, 3, 4),  (2, 5, 9),
    (2, 6, 15), (3, 4, 4),  (3, 6, 8),  (4, 7, 16), (5, 6, 15),
    (5, 8, 10), (6, 7, 15), (6, 8, 10), (7, 3, 6),  (7, 8, 10),
    (8, 1, 8)                        # Reverse edge to test the slope in EMRF
    ],
    1, 8,                              # source/target
    [28, 28, 15, 0],                   # answer for 1 to 4 route(s) flows
    [(0., 0., 3), (5., 15., 2),        # breaking points
    (10., 25., 1), (13., 28., 0)],
    28.                                # value for 1.5 routes
    ),

    # Graph with 6 vertices
    (6,
    [
    (1, 2, 9), (1, 3, 9), (2, 3, 10), (2, 4, 8), (3, 4, 1),
    (3, 5, 3), (5, 4, 8), (4, 6, 10), (5, 6, 7)
    ],
    1, 6,                                      # source/target
    [12, 6, 0],                                # answer for 1 to 3 route(s) flows
    [(0., 0., 2), (3., 6., 1), (9., 12., 0)],  # breaking points
    9.                                         # value for 1.5 routes
    ),

    # Graph with 7 vertices
    (7,
    [
    (1, 2, 1), (1, 3, 2), (1, 4, 3), (1, 5, 4), (1, 6, 5),
    (2, 7, 1), (3, 7, 2), (4, 7, 3), (5, 7, 4), (6, 7, 5)
    ],
    1, 7,                                      # source/target
    [15, 15, 15, 12, 5, 0],                    # answer for 1 to 6 route(s) flows
    [(0., 0., 5), (1., 5., 4), (2., 9., 3),    # breaking points
    (3., 12., 2), (4., 14., 1), (5., 15., 0)],
    15.                                        # value for 1.5 routes
    ),

    # Graph with 6 vertices
    (6,
    [
    (1, 2, 1), (1, 3, 1), (1, 4, 2), (1, 5, 2),
    (2, 6, 1), (3, 6, 1), (4, 6, 2), (5, 6, 2),
    ],
    1, 6,                                      # source/target
    [6, 6, 6, 4, 0],                           # answer for 1 to 5 route(s) flows
    [(0., 0., 4), (1., 4., 2), (2., 6., 0)],   # breaking points
    6.                                         # value for 1.5 routes
    )
    ]

    for (nvertices, flow_edges, s, t, froutes, breakpts, ffloat) in graphs
        flow_graph = Graphs.DiGraph(nvertices)
        for g in testdigraphs(flow_graph)
            capacity_matrix = zeros(Int, nvertices, nvertices)
            for e in flow_edges
                u, v, f = e
                Graphs.add_edge!(g, u, v)
                capacity_matrix[u, v] = f
            end
            # Test ExtendedMultirouteFlowAlgorithm when the number of routes is either
            # Noninteger or 0 (the algorithm returns the breaking points)
            @test multiroute_flow(g, s, t, capacity_matrix) == breakpts
            @test multiroute_flow(g, s, t, capacity_matrix, routes=1.5)[1] ≈ ffloat
            @test multiroute_flow(breakpts, 1.5)[2] ≈ ffloat

            # Test all other algorithms - PR is unstable - see #553
            for AlgoFlow in [EdmondsKarpAlgorithm, DinicAlgorithm, BoykovKolmogorovAlgorithm, PushRelabelAlgorithm]
                # When the input are breaking points and routes number
                @test multiroute_flow(breakpts, 1.5, g, s, t, capacity_matrix)[1] ≈ ffloat
                for AlgoMrf in [ExtendedMultirouteFlowAlgorithm, KishimotoAlgorithm]
                    for (k, val) in enumerate(froutes)
                        @test multiroute_flow(g, s, t, capacity_matrix,
                            flow_algorithm = AlgoFlow(), mrf_algorithm = AlgoMrf(),
                            routes = k)[1] ≈ val
                    end
                end
            end
        end
    end
end


# Flux vNext
# ==========

# Optimisation Basics
# -------------------

using Flux
using Flux: step!

# Using Flux is very simple. You write a program, and we'll make tweaks to that
# program so that it gets gradually better.

# What does it mean to make a program "better"? That's up to you – your program
# returns a score, called the "loss", which determines how well the program
# is doing. Flux's job is to minimise that score. For example, let's take the
# simplest possible program, one that simply returns a constant.

w = 1
loss = () -> w

loss()

# This program doesn't look very interesting, but we can still do something
# interesting with it. The core function that optimises programs for us is
# `step!`. We have to pass `step!` an *optimiser* called `Descent`; this basically
# tells Flux how agressive to be, but we'll talk more about that later.

opt = Descent(1)
step!(loss, opt)

# `step!` returns the same loss as before, `1`. But something more interesting
# has happened; try running `loss()` again.

loss()

# It went down! And if we keep calling `step!` in a loop, it'll keep going down.

for i = 1:10
  @show step!(loss, opt)
end

loss()

# Of course, this case is pretty easy: we can always improve the parameter by
# making `w` less.

w

# Here's something harder; now our `loss` is always positive, so it can't keep
# improving indefinitely. Things will stop improving once we hit the *minimum*
# of this function (which we happen to know is at $w = 0$, where $loss = 0$.)

w = 1
loss = () -> w^2

opt = Descent(0.2)
for i = 1:10
  @show step!(loss, opt)
end

w

# You can see that our loss gradually tends towards $0$, and so does $w$. Note,
# however, that Flux will never say: "Ok, we're done here, here's the best value
# for $w$." Though there are tools that can do this in simple cases, Flux is
# designed to scale to extremely complex problems where this is no longer
# possible. So we only make tweaks and it's up to you when to finish.

# Let's put these ideas towards something a little more interesting. Say we want
# to solve $5x = 10$, to find an $x$ that makes this true. What's our program?
# Well, to start with we want to take $f(x) = 5x$. Then our loss should be something like
# $f(x) - 10$, so that it measures how far the $f(x)$ is from where we want it to
# be. This doesn't quite work, however, since the loss will be low (negative) if
# $f(x)$ is `-Inf`! So we can use our squaring trick again here, to make
# sure that $f(x) - 10$ tends to zero.

x = 1 # Our initial guess
f = x -> 5x

opt = Descent(0.01)

for i = 1:10
  l = step!(opt) do
    (f(x) - 10)^2
  end
  @show l
end

# Our loss ended up being pretty low. How's our function looking?

5x

# That looks pretty good. So we're beginning to be able to use Flux to solve
# problems where we know what the *output* should look like, but we're not
# sure what the *input* should be to get there.

# You now arguably understand everything you need to do productive ML. But let's
# look over a few more examples to see how it looks in practice.

# Optimising Colours
# ------------------

# Just like Julia more generally, Flux has good support for custom types.
# This means we can carry out optimisation on things like colours!

# This example uses the excellent Colors.jl. Colors contains, among other
# things, a `colordiff` function which uses fancy colour theory algorithms to
# estimate the *perceptual* difference between two colours. We can use this
# directly in our loss function.

using Colors

target = RGB(1, 0, 0)
colour = RGB(1, 1, 1)
[target, colour]
#-
opt = Descent(0.01)

for i = 1:10
  step!(opt, target) do y
    colordiff(colour, y)
  end
end

[colour, target]

# `colour` started out white and is now red. That makes sense, as we've
# minimised the distance between the two colours. But we can also *maximize*
# with a simple minus sign.

colour1 = RGB(1, 1, 1)

for i = 1:10
  step!(opt, target) do y
    -colordiff(colour1, y)
  end
end

[colour1, target]

# Now we have green, a colour that's arguably very different from red. However,
# there's a subtlety here; notice what happens if we use a different colour as
# our starting point.

colour2 = RGB(0, 0, 1)

for i = 1:10
  step!(opt, target) do y
    -colordiff(colour2, y)
  end
end

[colour2, target]

# Now we have a dark blue! If we look directly at `colourdiff` we'll see that
# green is better.

colordiff(target, colour1), colordiff(target, colour2)

# So why do we get blue here? This is another case where it's important that
# Flux optimises programs through a series of small tweaks. In this case, even
# though green is better overall, making our colour slightly more green actually
# makes our score worse temporarily.

colordiff(target, RGB(0, 0, 0.4)), colordiff(target, RGB(0, 0.1, 0.4))

# This is known as a *local optimimum*. It's important to understand how Flux
# optimises programs and what this means for you, so we'll cover this in more
# detail in future.

# Keras in 5 lines
# -----------------

# [Working on making this an MNIST demo, but here's the gist of it.]

# Dummy data.

x = rand(10)
y = [1, 0]

# Logistic regresion.

using Flux: crossentropy

W = randn(2, 10)
b = zeros(2)

predict = x -> softmax(W * x .+ b)

opt = Descent(0.1)

loss = (x, y) -> crossentropy(predict(x), y)

step!(loss, opt, x, y)

# Multi-layer perceptron.

function dense(in, out, σ = identity)
  W = randn(out, in)
  b = zeros(out)
  x -> σ.(W * x .+ b)
end

chain(fs...) = x -> foldl((x, m) -> m(x), fs, init = x)

model = chain(dense(10, 5, relu), dense(5, 2), softmax)

# Doesn't quite work yet.

## step!(opt, x, y) do x, y
##   crossentropy(model(x), y)
## end


include("splinesParallel.jl")

b1 = plarBezier(S1, [[0,1,0],[0,1,5]])
b2 = plarBezier(S1, [[0,0,0],[0,0,5]])
b3 = plarBezier(S1, [[1,0,0],[2,-1,2.5],[1,0,5]])
b4 = plarBezier(S1, [[1,1,0],[1,1,5]])
b5 = plarBezier(S1, [[0,1,0],[0,1,5]])
controls = [b1, b2, b3, b4, b5]
knots = [0,1,2,3,4,5,6,7]           # periodic B-spline
knots = [0,0,0,1,2,3,3,3]           # non-periodic B-spline
ptbspline = pt_bspline(S2, 2, knots, controls)
point = [0.7, 0.2]
result = ptbspline(point)


function macWave(nr,nx,b,Nmax,depth,period)

#   assembles the discretized equations and boundary conditions into
#   two matrices for a generalized eigenvalue problem. This version uses
#   the finite difference method on a wave equations for b_phi. The
#   solution does NOT include the diffusive magnetic boundary layer at the
#   based of the stratified layer as part of the solution. The equations
#   are expressed in non-dimensional form using diffusion-free scales
#
#   Scales
#   length      - radius
#   time        - 1/rotation rate  = 1 / omega
#   velocity    -  omega * L
#   mag. field  - sqrt(mu * rho) * omega * L
#
#   Inputs
#   nr      -  number of grid points in radius
#   nx      -  number of grid points in colatitude
#   b      -   radial (main) magnetic field (mT)
#   eta     -  magnetic diffusivity  (m^2/s)
#   Nmax   -   buoyancy frequency at CMB (Nmax/Omega)
#   depth   -  thickness of stratified layer (km)
#   period  -  target period for wave
#
#
#   Output
#   freq     -  eigenfrequency
#   q        -  quality factor
#   mode     -  eigenvector
#   r, x     -  r and x=cos(theta) coordinates
#
#
#   definition of eigenvalue problem
#   (A - d I) * v == 0
#

#   some constants
    omega = 0.7292e-4;       # rotation rate (1/s)
    sectoyear = 3.1557e7;     # seconds in a year
    mu = 4 * pi * 1.0e-7;    # permeability
    eta = 0.8;               # magnetic diffusivity  (m^2/s)
    rho = 1.0e4;             # density
    L = 3.480e6;             # core radius (length scale)
    bscale = sqrt(rho * mu) * omega * L;

#   dimensionless quantities
    B = b * 1.0e-3 / bscale;
    Eta = eta / (omega * L^2);

#     test for dipole and constant stratification
    eigenfreq1 = 0.5 * B * Nmax / 1.0;
    w1 = sqrt(1*2) * eigenfreq1
    w2 = sqrt(2*3) * eigenfreq1
    w3 = sqrt(3*4) * eigenfreq1
    w4 = sqrt(4*5) * eigenfreq1
    #println("w1 ",w1," w2 ", w2," w3 ",w3)

#   set up grid
    rbottom = (L - 1.0e3 * depth)/L;
    r,x,theta = grids(nr,nx,rbottom);

#   stratification
    iopt = 1;
    N2 = stratification(r,Nmax,rbottom,iopt);

#   magnetic field (a function of x - see code for details)
    br = magnetic_field(B,x);

#   set up linear system
     rsp = restoring(r,x,br,N2);
     dsp = diffusion(r,x,Eta);
     nsp = dbdt(r,x);
     asp = rsp + dsp + nsp;  # combine elements of A matrix

#  solve for eigenvalues ()
    freq = 2 * pi / (period * sectoyear * omega)
    shift = freq * 1im;
    d,v = eigs(asp;which=:LM,nev=8,sigma=shift,maxiter=600);

    return r,x,d,v

end


include("header.jl")

using SpecialFunctions
using AutoGrad: lgamma, loggamma # TODO: delete after everyone has SpecialFunctions 0.8

@testset "specialfunctions" begin
    o = (:delta=>0.0001,:rtol=>0.02,:atol=>0.01)
    ϵ = 0.1
    val_0_2(x)=rand() * (2-2ϵ) + ϵ
    val_gt_0(x)=abs(x) + ϵ
    val_gt_m1(x)=abs(x) - 1 + ϵ
    val_lt_1(x)=-abs(x) + 1 - ϵ
    abs_lt_1(x)=rand() * (2-2ϵ) - (1-ϵ)
    val_gamma(x)=(x < ϵ && abs(x-round(x)) < ϵ ? x+0.5 : x) # avoid <=0 integers

    # `airy(k,x)` is deprecated, use `airyai(x)`, `airyaiprime(x)`, `airybi(x)` or `airybiprime(x)` instead.
    @test randcheck(airyai; o...) # @primitive airyai(x),dy (dy.*(airyaiprime.(x)))
    @test randcheck(airyaiprime; o...) # @primitive airyaiprime(x),dy (dy.*(x .* airyai.(x)))
    # airyaiprimex
    # airyaix
    @test randcheck(airybi,val_lt_1; o...) # @primitive airybi(x),dy (dy.*(airybiprime.(x)))
    @test randcheck(airybiprime,val_lt_1; o...) # @primitive airybiprime(x),dy (dy.*(x .* airybi.(x)))
    # airybiprimex
    # airybix
    # `airyprime(z::Number)` is deprecated, use `airyaiprime(z)` instead.
    # `airyx(k,x)` is deprecated, use `airyaix(x)`, `airyaiprimex(x)`, `airybix(x)` or `airybiprimex(x)` instead.
    # besselh
    # besselhx
    # besseli
    # besselix
    # besselj
    @test randcheck(besselj0; o...) # @primitive besselj0(x),dy (dy.*(-(besselj1.(x))))
    @test randcheck(besselj1; o...) # @primitive besselj1(x),dy (dy.*((besselj0.(x) - besselj.(2,x)) / 2))
    # besseljx
    # besselk
    # besselkx
    # bessely
    @test randcheck(bessely0,val_gt_0; o...) # @primitive bessely0(x),dy (dy.*(-(bessely1.(x))))
    @test randcheck(bessely1,val_gt_0; o...) # @primitive bessely1(x),dy (dy.*((bessely0.(x) - bessely.(2,x)) / 2))
    # besselyx
    # beta
    # cosint
    @test randcheck(dawson; o...) # @primitive dawson(x),dy,y (dy.*((-2y) .* x + 1))
    @test randcheck(digamma,val_gamma; o...) # @primitive digamma(x),dy,y (dy.*(trigamma.(x))) ## avoid <=0 ints
    @test randcheck(erf; o...) # @primitive erf(x),dy,y (dy.*(exp.(-(abs2.(x))) * convert(eltype(x), 2 / √π)))
    @test randcheck(erfc; o...) # @primitive erfc(x),dy,y (dy.*(-(exp.(-(abs2.(x)))) * convert(eltype(x), 2 / √π)))
    @test randcheck(erfcinv,val_0_2; o...) # @primitive erfcinv(x),dy,y (dy.*(-(exp.(abs2.(y))) * convert(eltype(x), √π / 2)))
    @test randcheck(erfcx,val_gt_m1; o...) # @primitive erfcx(x),dy,y (dy.*((2y) .* x - convert(eltype(x), 2 / √π)))
    @test randcheck(erfi,abs_lt_1; o...) # @primitive erfi(x),dy,y (dy.*(exp.(abs2.(x)) * convert(eltype(x), 2 / √π)))
    @test randcheck(erfinv,abs_lt_1; o...) # @primitive erfinv(x),dy,y (dy.*(exp.(abs2.(y)) * convert(eltype(x), √π / 2)))
    # eta
    @test randcheck(gamma,val_gamma; o...) # @primitive gamma(x),dy,y (dy.*(y .* digamma.(x))) ## avoid <=0 ints
    # hankelh1
    # hankelh1x
    # hankelh2
    # hankelh2x
    @test randcheck(invdigamma, val_lt_1; o...) # @primitive invdigamma(x),dy,y (dy.*(1 ./ trigamma.(y)))
    # lbeta
    # `lfact` is deprecated, use `lfactorial` instead.
    # lfactorial
    @test randcheck(loggamma,val_gt_0; o...) # @primitive loggamma(x),dy,y (dy.*(digamma.(x))) ## avoid <=0
    # @test randcheck(polygamma) # @primitive polygamma(x1,x2),dy,y  nothing  unbroadcast(x2,dy.*polygamma(x1+1,x2))
    # sinint
    @test randcheck(trigamma,val_gamma; o...) # @primitive trigamma(x),dy,y (dy.*(polygamma.(2,x)))
    # zeta
end