using Turing, Distributions, Test
using MacroTools

model_info = Dict(:main_body_names => Dict(:vi => :vi, :sampler => :sampler))
# unit test model macro
expr = Turing.generate_observe(:x, :y, model_info)
@test expr.head == :block
@test :(vi.logp += Turing.observe(sampler, y, x, vi)) in expr.args

@model testmodel_comp(x, y) = begin
    s ~ InverseGamma(2,3)
    m ~ Normal(0,sqrt(s))

    x ~ Normal(m, sqrt(s))
    y ~ Normal(m, sqrt(s))

    return x, y
end
testmodel_comp(1.0, 1.2)

# check if drawing from the prior works
@model testmodel0(x) = begin
    x ~ Normal()
    return x
end
f0_mm = testmodel0()
@test mean(f0_mm() for _ in 1:1000) ≈ 0. atol=0.1

# Test #544
@model testmodel0(x = Vector{Float64}(undef, 2)) = begin
    x[1] ~ Normal()
    x[2] ~ Normal()
    return x
end
f0_mm = testmodel0()
@test all(isapprox.(mean(f0_mm() for _ in 1:1000), 0., atol=0.1))

@model testmodel01(x) = begin
    x ~ Bernoulli(0.5)
    return x
end
f01_mm = testmodel01()
@test mean(f01_mm() for _ in 1:1000) ≈ 0.5 atol=0.1

# test if we get the correct return values
@model testmodel1(x1, x2) = begin
    s ~ InverseGamma(2,3)
    m ~ Normal(0,sqrt(s))

    x1 ~ Normal(m, sqrt(s))
    x2 ~ Normal(m, sqrt(s))

    return x1, x2
end
f1_mm = testmodel1(1., 10.)
@test f1_mm() == (1, 10)

# Test for assertions in observe statements.
@model brokentestmodel_observe1(x1, x2) = begin
    s ~ InverseGamma(2,3)
    m ~ Normal(0,sqrt(s))

    x1 ~ Normal(m, sqrt(s))
    x2 ~ x1 + 2

    return x1, x2
end

btest = brokentestmodel_observe1(1., 2.)
@test_throws ArgumentError btest()

@model brokentestmodel_observe2(x) = begin
    s ~ InverseGamma(2,3)
    m ~ Normal(0,sqrt(s))

    x = Vector{Float64}(undef, 2)
    x ~ [Normal(m, sqrt(s)), 2.0]

    return x
end

btest = brokentestmodel_observe2([1., 2.])
@test_throws ArgumentError btest()

# Test for assertions in assume statements.
@model brokentestmodel_assume1() = begin
    s ~ InverseGamma(2,3)
    m ~ Normal(0,sqrt(s))

    x1 ~ Normal(m, sqrt(s))
    x2 ~ x1 + 2

    return x1, x2
end

btest = brokentestmodel_assume1()
@test_throws ArgumentError btest()

@model brokentestmodel_assume2() = begin
    s ~ InverseGamma(2,3)
    m ~ Normal(0,sqrt(s))

    x = Vector{Float64}(undef, 2)
    x ~ [Normal(m, sqrt(s)), 2.0]

    return x
end

btest = brokentestmodel_assume2()
@test_throws ArgumentError btest()


using JuMP
using Clp
using DiffOpt
using Test
using ChainRulesCore

# This script creates the JuMP model for a small unit commitment instance
# represented in a solution map function taking parameters as arguments and returning
# the optimal solution as output.

# The derivatives of this solution map can then be expressed in ChainRules semantics
# and implemented using DiffOpt

ATOL=1e-4
RTOL=1e-4

"""
Solution map of the problem using parameters:
- `load1_demand, load2_demand` for the demand of two nodes
- `gen_costs` is the vector of generator costs
- `noload_costs` is the vector of fixed activation costs of the generators,
and returning the optimal output power `p`.
"""
function unit_commitment(load1_demand, load2_demand, gen_costs, noload_costs; model = Model(() -> diff_optimizer(Clp.Optimizer)))
    ## Problem data
    unit_codes = [1, 2] # Generator identifiers
    load_names = ["Load1", "Load2"] # Load identifiers
    n_periods = 4 # Number of time periods
    Pmin = Dict(1 => fill(0.5, n_periods), 2 => fill(0.5, n_periods)) # Minimum power output (pu)
    Pmax = Dict(1 => fill(3.0, n_periods), 2 => fill(3.0, n_periods)) # Maximum power output (pu)
    RR = Dict(1 => 0.25, 2 => 0.25) # Ramp rates (pu/min)
    P0 = Dict(1 => 0.0, 2 => 0.0) # Initial power output (pu)
    D = Dict("Load1" => load1_demand, "Load2" => load2_demand) # Demand (pu)
    Cp = Dict(1 => gen_costs[1], 2 => gen_costs[2]) # Generation cost coefficient ($/pu)
    Cnl = Dict(1 => noload_costs[1], 2 => noload_costs[2]) # No-load cost ($)

    ## Variables
    # Note: u represents the activation of generation units.
    # Would be binary in the typical UC problem, relaxed here to u ∈ [0,1]
    # for a linear relaxation.
    @variable(model, 0 <= u[g in unit_codes, t in 1:n_periods] <= 1) # Commitment
    @variable(model, p[g in unit_codes, t in 1:n_periods] >= 0) # Power output (pu)

    ## Constraints

    # Energy balance
    @constraint(
        model,
        energy_balance_cons[t in 1:n_periods],
        sum(p[g, t] for g in unit_codes) == sum(D[l][t] for l in load_names),
    )

    # Generation limits
    @constraint(model, [g in unit_codes, t in 1:n_periods], Pmin[g][t] * u[g, t] <= p[g, t])
    @constraint(model, [g in unit_codes, t in 1:n_periods], p[g, t] <= Pmax[g][t] * u[g, t])

    # Ramp rates
    @constraint(model, [g in unit_codes, t in 2:n_periods], p[g, t] - p[g, t - 1] <= 60 * RR[g])
    @constraint(model, [g in unit_codes], p[g, 1] - P0[g] <= 60 * RR[g])
    @constraint(model, [g in unit_codes, t in 2:n_periods], p[g, t - 1] - p[g, t] <= 60 * RR[g])
    @constraint(model, [g in unit_codes], P0[g] - p[g, 1] <= 60 * RR[g])

    # Objective
    @objective(
        model,
        Min,
        sum((Cp[g] * p[g, t]) + (Cnl[g] * u[g, t]) for g in unit_codes, t in 1:n_periods),
    )

    optimize!(model)
    @assert termination_status(model) == MOI.OPTIMAL
    # converting to dense matrix
    return JuMP.value.(p.data)
end

@show unit_commitment([1.0, 1.2, 1.4, 1.6], [1.0, 1.2, 1.4, 1.6], [1000.0, 1500.0], [500.0, 1000.0])

# Forward differentiation rule for the solution map of the unit commitment problem
# taking in input perturbations on the input parameters and returning perturbations propagated to the result
function ChainRulesCore.frule((_, Δload1_demand, Δload2_demand, Δgen_costs, Δnoload_costs), ::typeof(unit_commitment), load1_demand, load2_demand, gen_costs, noload_costs)
    # creating the UC model with a DiffOpt optimizer wrapper around Clp
    model = Model(() -> diff_optimizer(Clp.Optimizer))
    # building and solving the main model
    pv = unit_commitment(load1_demand, load2_demand, gen_costs, noload_costs, model=model)
    energy_balance_cons = model[:energy_balance_cons]

    # Setting some perturbation of the right-hand side of the energy balance constraints
    # the RHS is equal to the sum of load demands at each period.
    # the corresponding perturbation are set accordingly as the set of perturbations of the two loads
    MOI.set.(
        model,
        DiffOpt.ForwardInConstraint(), energy_balance_cons,
        AffExpr[d1 + d2 for (d1, d2) in zip(Δload1_demand, Δload2_demand)],
    )

    p = model[:p]
    u = model[:u]

    # setting the perturbation of the linear objective
    Δobj = sum(Δgen_costs ⋅ p[:,t] + Δnoload_costs ⋅ u[:,t] for t in size(p, 2))
    MOI.set(model, DiffOpt.ForwardInObjective(), Δobj)
    DiffOpt.forward(JuMP.backend(model))
    # querying the corresponding perturbation of the decision
    Δp = MOI.get.(model, DiffOpt.ForwardOutVariablePrimal(), p)
    return (pv, Δp.data)
end

load1_demand = [1.0, 1.2, 1.4, 1.6]
load2_demand = [1.0, 1.2, 1.4, 1.6]
gen_costs = [1000.0, 1500.0]
noload_costs = [500.0, 1000.0]

Δload1_demand = 0 * load1_demand .+ 0.1
Δload2_demand = 0 * load2_demand .+ 0.2
Δgen_costs = 0 * gen_costs .+ 0.1
Δnoload_costs = 0 * noload_costs .+ 0.4
@show (pv, Δpv) = ChainRulesCore.frule((nothing, Δload1_demand, Δload2_demand, Δgen_costs, Δnoload_costs), unit_commitment, load1_demand, load2_demand, gen_costs, noload_costs)

# Reverse-mode differentiation of the solution map
# The computed pullback takes a seed for the optimal solution `̄p` and returns
# derivatives wrt each input parameter.
function ChainRulesCore.rrule(::typeof(unit_commitment), load1_demand, load2_demand, gen_costs, noload_costs; model = Model(() -> diff_optimizer(Clp.Optimizer)))
    # solve the forward UC problem
    pv = unit_commitment(load1_demand, load2_demand, gen_costs, noload_costs, model=model)
    function pullback_unit_commitment(pb)
        p = model[:p]
        u = model[:u]
        energy_balance_cons = model[:energy_balance_cons]

        MOI.set.(model, DiffOpt.BackwardInVariablePrimal(), p, pb)
        DiffOpt.backward(JuMP.backend(model))

        obj = MOI.get(model, DiffOpt.BackwardOutObjective())

        # computing derivative wrt linear objective costs
        dgen_costs = similar(gen_costs)
        dgen_costs[1] = sum(JuMP.coefficient.(obj, p[1,:]))
        dgen_costs[2] = sum(JuMP.coefficient.(obj, p[2,:]))

        dnoload_costs = similar(noload_costs)
        dnoload_costs[1] = sum(JuMP.coefficient.(obj, u[1,:]))
        dnoload_costs[2] = sum(JuMP.coefficient.(obj, u[2,:]))

        # computing derivative wrt constraint constant
        dload1_demand = JuMP.constant.(MOI.get.(model, DiffOpt.BackwardOutConstraint(), energy_balance_cons))
        dload2_demand = copy(dload1_demand)
        return (dload1_demand, dload2_demand, dgen_costs, dnoload_costs)
    end
    return (pv, pullback_unit_commitment)
end

(pv, pullback_unit_commitment) = ChainRulesCore.rrule(unit_commitment, load1_demand, load2_demand, gen_costs, noload_costs; model = Model(() -> diff_optimizer(Clp.Optimizer)))
@show pullback_unit_commitment(ones(size(pv)))


using Revise
using ADCME
using PyCall
using LinearAlgebra
using AdFem
using PyPlot
using Random
Random.seed!(233)

# TODO: specify your input parameters
σ1 = rand(10, 3)
σ2 = zeros(10, 3)
ε1 = rand(10, 3)
ε2 = rand(10, 3)
μ = rand(10)
η = rand(10)
λ = rand(10)
Δt = rand()

u1 = update_stress_viscosity(ε2, ε1, σ1, η, μ, λ, Δt)

u2 = update_stress_viscosity(constant(ε2), ε1, σ1, η, μ, λ, Δt)

sess = Session(); init(sess)
@show run(sess, u2)-u1

# uncomment it for testing gradients
# error()

# TODO: change your test parameter to `m`
#       in the case of `multiple=true`, you also need to specify which component you are testings
# gradient check -- v
function scalar_function(x)
    return sum(update_stress_viscosity(ε2, ε1, σ1, η, μ, x, Δt)^2)
end

# TODO: change `m_` and `v_` to appropriate values
m_ = constant(rand(10))
v_ = rand(10)
y_ = scalar_function(m_)
dy_ = gradients(y_, m_)
ms_ = Array{Any}(undef, 5)
ys_ = Array{Any}(undef, 5)
s_ = Array{Any}(undef, 5)
w_ = Array{Any}(undef, 5)
gs_ =  @. 1 / 10^(1:5)

for i = 1:5
    g_ = gs_[i]
    ms_[i] = m_ + g_*v_
    ys_[i] = scalar_function(ms_[i])
    s_[i] = ys_[i] - y_
    w_[i] = s_[i] - g_*sum(v_.*dy_)
end

sess = Session(); init(sess)
sval_ = run(sess, s_)
wval_ = run(sess, w_)
close("all")
loglog(gs_, abs.(sval_), "*-", label="finite difference")
loglog(gs_, abs.(wval_), "+-", label="automatic differentiation")
loglog(gs_, gs_.^2 * 0.5*abs(wval_[1])/gs_[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
loglog(gs_, gs_ * 0.5*abs(sval_[1])/gs_[1], "--",label="\$\\mathcal{O}(\\gamma)\$")

plt.gca().invert_xaxis()
legend()
xlabel("\$\\gamma\$")
ylabel("Error")
savefig("gradtest.png")


function solve()
    n, x = [parse(Int, x) for x in split(readline())]
    a = 0
    for i=1:n
        v, p = [parse(Int, x) for x in split(readline())]
        a += v * p
        if a > x * 100
            return i
        end
    end
    return -1
end

println(solve())


# TODO: add more target distributions and make them iteratable

# Dimension of testing distribution
const D = 5
# Tolerance ratio
const TRATIO = Int == Int64 ? 1 : 2
# Deterministic tolerance
const DETATOL = 1e-3 * D * TRATIO
# Random tolerance
const RNDATOL = 5e-2 * D * TRATIO * 2

# Hand-coded multivaite Gaussain

const gaussian_m = zeros(D)
const gaussian_s = ones(D)

function ℓπ(m, s, x::AbstractVecOrMat{T}) where {T}
    diff = x .- m
    v = s.^2
    return -(log(2 * T(pi)) .+ log.(v) .+ diff .* diff ./ v) / 2
end

function ℓπ(θ::AbstractVector)
    return sum(ℓπ(gaussian_m, gaussian_s, θ))
end

function ℓπ(θ::AbstractMatrix)
    return dropdims(sum(ℓπ(gaussian_m, gaussian_s, θ); dims=1); dims=1)
end

function ∂ℓπ∂θ(m, s, x::AbstractVecOrMat{T}) where {T}
    diff = x .- m
    v = s.^2
    v = -(log(2 * T(pi)) .+ log.(v) .+ diff .* diff ./ v) / 2
    g = -diff
    return v, g
end

function ∂ℓπ∂θ(θ::AbstractVector)
    v, g = ∂ℓπ∂θ(gaussian_m, gaussian_s, θ)
    return sum(v), g
end

function ∂ℓπ∂θ(θ::AbstractMatrix)
    v, g = ∂ℓπ∂θ(gaussian_m, gaussian_s, θ)
    return dropdims(sum(v; dims=1); dims=1), g
end

# For the Turing model
# @model gdemo() = begin
#     s ~ InverseGamma(2, 3)
#     m ~ Normal(0, sqrt(s))
#     1.5 ~ Normal(m, sqrt(s))
#     2.0 ~ Normal(m, sqrt(s))
#     return s, m
# end

using Distributions: logpdf, InverseGamma, Normal
using Bijectors: invlink, logpdf_with_trans

function invlink_gdemo(θ)
    s = invlink(InverseGamma(2, 3), θ[1])
    m = θ[2]
    return [s, m]
end

function ℓπ_gdemo(θ)
    s, m = invlink_gdemo(θ)
    logprior = logpdf_with_trans(InverseGamma(2, 3), s, true) + logpdf(Normal(0, sqrt(s)), m)
    loglikelihood = logpdf(Normal(m, sqrt(s)), 1.5) + logpdf(Normal(m, sqrt(s)), 2.0)
    return logprior + loglikelihood
end

using Distributions: MvNormal
import Turing

Turing.@model function mvntest(θ=missing, x=missing)
    θ ~ MvNormal(zeros(D), 2)
    x ~ Normal(sum(θ), 1)
    return θ, x
end

function get_primitives(x, modelgen)
    spl_prior = Turing.SampleFromPrior()
    function ℓπ(θ)
        vi = Turing.VarInfo(model)
        vi[spl_prior] = θ
        model(vi, spl_prior)
        Turing.getlogp(vi)
    end
    adbackend = Turing.Core.ForwardDiffAD{40}
    alg_ad = Turing.HMC{adbackend}(0.1, 1)
    model = modelgen(missing, x)
    vi = Turing.VarInfo(model)
    spl = Turing.Sampler(alg_ad, model)
    Turing.Core.link!(vi, spl)
    ∂ℓπ∂θ = θ -> Turing.Core.gradient_logp(adbackend(), θ, vi, model, spl)
    θ₀ = Turing.VarInfo(model)[Turing.SampleFromPrior()]
    return ℓπ, ∂ℓπ∂θ, θ₀
end

function rand_θ_given(x, modelgen, metric, κ; n_samples=20)
    ℓπ, ∂ℓπ∂θ, θ₀ = get_primitives(x, modelgen)
    h = Hamiltonian(metric, ℓπ, ∂ℓπ∂θ)
    samples, stats = sample(h, κ, θ₀, n_samples; verbose=false, progress=false)
    s = samples[end]
    return length(s) == 1 ? s[1] : s
end

# Test function
geweke_g(θ, x) = cat(θ, x; dims=1)


"""
    wshenk(object::Union{Pca, Plsr}, X; nlv = nothing)
Compute the Shenk et al. (1997) PLSR weights
* `object` : The fitted model.
* `X` : X-data on which the weights are computed.
* `nlv` : Nb. latent variables (LVs) to consider. If nothing,
    it is the maximum nb. of components.

For each observation (row) of `X`, the weights are returned
for the models with 1, ..., nlv LVs.

## References

Shenk, J., Westerhaus, M., Berzaghi, P., 1997. Investigation of a LOCAL calibration
procedure for near infrared instruments.
Journal of Near Infrared Spectroscopy 5, 223. https://doi.org/10.1255/jnirs.115

"""
function wshenk(object::Union{Plsr}, X; nlv = nothing, norm = false)
    X = ensure_mat(X)
    m = size(X, 1)
    p = size(X, 2)
    q = size(object.C, 1)
    a = size(object.T, 2)
    isnothing(nlv) ? nlv = a : nlv = min(nlv, a)
    E = similar(X)
    rms_r = similar(X, m, a)
    zB = similar(X, p + 1, q)
    rms_b = similar(X, a)
    W = similar(X, m, a)
    for j = 1:a
        E .= xresid(object, X, nlv = j)
        rms_r[:, j] .= vec(sqrt.(sum(E.^2, dims = 2)))
        z = coef(object; nlv = j)
        zB[1, :] .= vec(z.int)
        zB[2:end, :] .= z.B
        rms_b[j] = sqrt(mean(mean(zB.^2, dims = 2)))
        W[:, j] .= 1 ./ (vcol(rms_r, j) * rms_b[j])
    end
    for i = 1:m
        W[i, :] ./= sum(vrow(W, i))
    end
    (W = W, rms_r = rms_r, rms_b = rms_b)
end


### Statistics

function countmap(x::RLEVector{T1,T2}) where {T1,T2}
    tally = Dict{T1,T2}()
    for (v, w) in zip(values(x), widths(x))
        tally[v] = get(tally, v, 0) + w
    end
    tally
end

function mode(x::RLEVector)
    tally = countmap(x)
    which_max = argmax(collect(values(tally)))
    max_val = collect(keys(tally))[which_max]
    max_val
end


function pose_to_matrix(pose::Pose)::Matrix{Float32}
    R = Matrix{Float32}(pose.orientation)
    mat = zeros(Float32,4,4)
    mat[end,end] = 1
    mat[1:3,1:3] = R
    mat[1:3,4] = pose.pos
    return mat
end

function I4(t)
    x = zeros(t, 4,4)
    x[1,1] = 1.0
    x[2,2] = 1.0
    x[3,3] = 1.0
    x[4,4] = 1.0
    x
end

function compute_projection_matrix(fx, fy, cx, cy, near, far, skew=0f0)
    proj = I4(Float32)
    proj[1, 1] = fx
    proj[2, 2] = fy
    proj[1, 2] = skew
    proj[1, 3] = -cx
    proj[2, 3] = -cy
    proj[3, 3] = near + far
    proj[3, 4] = near * far
    proj[4, 4] = 0.0f0
    proj[4, 3] = -1f0
    return proj
end

function compute_ortho_matrix(left, right, bottom, top, near, far)
    ortho = I4(Float32)
    ortho[1, 1] = 2f0 / (right-left)
    ortho[2, 2] = 2f0 / (top-bottom)
    ortho[3, 3] = - 2f0 / (far - near)
    ortho[1, 4] = - (right + left) / (right - left)
    ortho[2, 4] = - (top + bottom) / (top - bottom)
    ortho[3, 4] = - (far + near) / (far - near)
    return ortho
end

function get_perspective_matrix(width, height, fx, fy, cx, cy, near, far)
    # (height-cy) is used instead of cy because of the difference between
    # image coordinate systems between OpenCV and OpenGL. In the former,
    # the origin is at the top-left of the image while in the latter the
    # origin is at the bottom-left.
    proj_matrix = compute_projection_matrix(
            fx, fy, cx, (height-cy),
            near, far, 0.f0)
    ndc_matrix = compute_ortho_matrix(0, width, 0, height, near, far)
    ndc_matrix * proj_matrix
end

function CameraIntrinsics_from_fov_aspect(width, height, fov_y, aspect_ratio, near, far)
    # Camera principal point is the center of the image.
    cx, cy = width / 2.0, height / 2.0

    # Vertical field of view is given.
    fov_y = deg2rad(fov_y)
    # Convert field of view to distance to scale by aspect ratio and
    # convert back to radians to recover the horizontal field of view.
    fov_x = 2 * atan(aspect_ratio * tan(fov_y / 2.0))

    # Use the following relation to recover the focal length:
    #   FOV = 2 * atan( (0.5 * IMAGE_PLANE_SIZE) / FOCAL_LENGTH )
    fx = cx / tan(fov_x / 2.0)
    fy = cy / tan(fov_y / 2.0)

    CameraIntrinsics(width, height,
        fx, fy, cx, cy,
        near, far)
end

function scale_down_camera(camera, factor)
    camera_modified = CameraIntrinsics(
        width=round(Int,camera.width/factor), height=round(Int,camera.height/factor),
        fx=camera.fx/factor, fy=camera.fy/factor, cx=camera.cx/factor, cy=camera.cy/factor,
        near=camera.near, far=camera.far)
    camera_modified
end


function reset!(alg::TrustRegion)
    alg.Δ[] = alg.Δ₀
    return nothing
end

function step!(x, alg::TrustRegion, f::Function, state)
    Δ = alg.Δ[]
    Δₘₐₓ = alg.Δₘₐₓ
    η = alg.η
    ηₛ = alg.ηₛ
    ηₑ = alg.ηₑ
    fx = state.f
    ∇fx = state.∇f
    B = approx_hessian(alg.hessian, state)
    model(p) = fx + ∇fx⋅p + (p'*B*p)/2
    p = alg.method(state, B, Δ)
    ρ = (fx - f(x + p)) / (model(zero(p)) - model(p))

    # Update radius
    if ρ ≥ ηₑ && norm(p) ≈ Δ # highly successful
        alg.Δ[] = min(alg.σₑ*Δ, Δₘₐₓ)
    elseif ρ < ηₛ
        alg.Δ[] = alg.σₛ*Δ
    end

    # Apply step
    if ρ ≥ η
        x .+= p
    end

    return nothing
end


        # isleft
        l = LineSegment2D(Point2D{T}(0, 0), Point2D{T}(1, 0))
        @test isleft(Point2D{T}(0, 1) , l)
        @test !isleft(Point2D{T}(0, -1) , l)
        @test !isleft(Point2D{T}(0, -1e-6) , l)
        @test isleft(Point2D{T}(0, 1e-6) , l)
        @test isleft(Point2D{T}(0.5, 0) , l)

        # sort_intersection_points
        l = LineSegment3D(Point3D{T}(0,0,0), Point3D{T}(10,0,0))
        p₁ = Point3D{T}(1, 0, 0)
        p₂ = Point3D{T}(2, 0, 0)
        p₃ = Point3D{T}(3, 0, 0)
        points = [p₃, p₁, p₂, Point3D{T}(1 + 1//1000000, 0, 0)]
        sort_intersection_points!(l, points)
        @test points[1] == p₁
        @test points[2] == p₂
        @test points[3] == p₃


using Plots; pyplot()

include("../OceanParameters.jl")

## Replicate Bottom Loss Curves

θ₁ = π/4
@show BL = OceanParameters.bottom_loss(θ₁)

l = @layout [a b; c d]
θ₁ = deg2rad.(range(0, 90, length = 100))

cₚ = [1550. 1600. 1800.]
Numθ₁ = length(θ₁)
Numcₚ = length(cₚ)
BL = zeros(AbstractFloat, (Numθ₁, Numcₚ))
for nθ₁ = 1:length(θ₁)
	for ncₚ = 1:length(cₚ)
		BL[nθ₁, ncₚ] = OceanParameters.bottom_loss(θ₁[nθ₁]; cₚ = cₚ[ncₚ])
	end
end

p1 = plot(rad2deg.(θ₁), BL,
	xaxis = 0:30:90,
	yaxis = ("Loss (dB)", 0:5:15),
	label = cₚ,
	legendtitle = "cₚ (m/s)",
	legend = :bottomright)

αₚ = [1. 0.5 0]
Numθ₁ = length(θ₁)
Numαₚ = length(αₚ)
BL = zeros(AbstractFloat, (Numθ₁, Numαₚ))
for nθ₁ = 1:length(θ₁)
	for nαₚ = 1:length(αₚ)
		BL[nθ₁, nαₚ] = OceanParameters.bottom_loss(θ₁[nθ₁]; αₚ = αₚ[nαₚ])
	end
end

p2 = plot(rad2deg.(θ₁), BL,
	xaxis = 0:30:90,
	yaxis = 0:5:15,
	label = αₚ,
	legendtitle = "αₚ (dB/λ)",
	legend = :bottomright)

ρ₂ = [1.5e3 2e3 2.5e3]
Numρ₂ = length(ρ₂)
BL = zeros(AbstractFloat, (Numθ₁, Numρ₂))
for nθ₁ = 1:length(θ₁)
	for nρ₂ = 1:length(ρ₂)
		BL[nθ₁, nρ₂] = OceanParameters.bottom_loss(θ₁[nθ₁]; ρ₂ = ρ₂[nρ₂])
	end
end

p3 = plot(rad2deg.(θ₁), BL,
	xaxis = ("Grazing angle θ₁ (deg)", 0:30:90),
	yaxis = ("Loss (dB)", 0:5:15),
	label = ρ₂,
	legendtitle = "ρ₂ (kg/m³)",
	legend = :bottomright)

cₛ = [6e2 4e2 2e2 0.0]
Numcₛ = length(cₛ)
BL = zeros(AbstractFloat, (Numθ₁, Numcₛ))
for nθ₁ = 1:length(θ₁)
	for ncₛ = 1:length(cₛ)
		BL[nθ₁, ncₛ] = OceanParameters.bottom_loss(θ₁[nθ₁]; cₛ = cₛ[ncₛ], αₚ = 0.)
	end
end

p4 = plot(rad2deg.(θ₁), BL,
	xaxis = ("Grazing angle θ₁ (deg)", 0:30:90),
	yaxis = 0:5:15,
	label = cₛ,
	legendtitle = "cₛ (m/s)",
	legend = :bottomright)

pt = plot(p1, p2, p3, p4,
	layout = l)

savefig(pt, "img/BottomLoss_Parameters.png")

## Bottom Loss Types
θ₁ = deg2rad.(range(0., 90., length = Int(1e3)))
BottomTypes = ["Clay" "Silt" "Sand" "Gravel" "Moraine" "Chalk" "Limestone" "Basalt"]
ρ₂ = [1.5e3 1.7e3 1.9e3 2e3 2.1e3 2.2e3 2.4e3 2.7e3]
cₚ = [1.5e3 1575 1650 1.8e3 1950 2.4e3 3e3 5250]
cₛ = [5e1 0 0 0 6e2 1e3 1.5e3 2.5e3]
αₚ = [0.2 1. 0.8 0.6 0.4 0.2 0.1 0.1]
αₛ = [1.0 1.5 2.5 1.5 1.0 0.5 0.2 0.2]

NumBot = length(BottomTypes)
Numθ = length(θ₁)
BL = zeros(AbstractFloat, (Numθ, NumBot))
for nBot = 1:NumBot
	for nθ = 1:Numθ
		BL[nθ, nBot] = OceanParameters.bottom_loss(θ₁[nθ]; cₚ = cₚ[nBot], αₚ = αₚ[nBot], ρ₂ = ρ₂[nBot], cₛ = cₛ[nBot], αₛ = αₛ[nBot])
	end
end

pt = plot(rad2deg.(θ₁), BL,
	title = "Bottom Loss Types",
	yaxis = ("Bottom Reflection Loss (dB)", 0:5:20),
	xaxis = ("Grazing Angle (deg)", 0:15:90),
	label = BottomTypes,
	ylims = (0, 20))

savefig(pt, "img/BottomLoss_Types.png")

##


using DelimitedFiles

function support_serial_qtest(mdp::Union{MDP,POMDP}, file::AbstractString; niter::Int64=100, res::Float64=1e-3)
    qt = readdlm(file)
    solver = ValueIterationSolver(max_iterations=niter, belres=res, verbose=true)
    policy = solve(solver, mdp)
    (q, u, p, am) = locals(policy)
    npolicy = ValueIterationPolicy(mdp, deepcopy(q))
    nnpolicy = ValueIterationPolicy(mdp, deepcopy(q), deepcopy(u), deepcopy(p))
    s = GridWorldState(1,1)
    a1 = action(policy, s)
    @inferred action(policy, s)
    v1 = value(policy, s)
    a2 = action(npolicy, s)
    v2 = value(npolicy, s)
    return (isapprox(qt, q, rtol=1e-5)) && (policy.policy == nnpolicy.policy)
end

function test_complex_gridworld()
    # Load correct policy from file and verify we can reconstruct it
    rstates = [GridWorldState(4,3), GridWorldState(4,6), GridWorldState(9,3), GridWorldState(8,8)]
    rvals = [-10.0, -5.0, 10.0, 3.0]
    xs = 10
    ys = 10
    mdp = LegacyGridWorld(sx=xs, sy=ys, rs = rstates, rv = rvals)
    file = "grid-world-10x10-Q-matrix.txt"
    niter = 100
    res = 1e-3

    return support_serial_qtest(mdp, file, niter=niter, res=res)
end

function test_simple_grid()
    # Simple test....
    # GridWorld(sx=2,sy=3) w reward at (2,3):
    # Here's our grid:
    # |state (x,y)____available actions__|
    # ----------------------------------------------
    # |5 (1,3)__u,d,l,r__|6 (2,3)__u,d,l,r+REWARD__|
    # |3 (1,2)__u,d,l,r__|4 (2,2)______u,d,l,r_____|
    # |1 (1,1)__u,d,l,r__|2 (2,1)______u,d,l,r_____|
    # ----------------------------------------------
    # 7 (0,0) is absorbing state
    mdp = LegacyGridWorld(sx=2, sy=3, rs = [GridWorldState(2,3)], rv = [10.0])

    solver = ValueIterationSolver(verbose=true)
    policy = solve(solver, mdp)

    # up: 1, down: 2, left: 3, right: 4
    correct_policy = [1,1,1,1,4,1,1] # alternative policies
    # are possible, but since they are tied & the first
    # action is always 1, we will always return 1 for tied
    # actions
    return policy.policy == correct_policy
end

function test_init_solution()
    # Initialize the value to the solution
    rstates = [GridWorldState(4,3), GridWorldState(4,6), GridWorldState(9,3), GridWorldState(8,8)]
    rvals = [-10.0, -5.0, 10.0, 3.0]
    xs = 10
    ys = 10
    mdp = LegacyGridWorld(sx=xs, sy=ys, rs = rstates, rv = rvals)
    qt = readdlm("grid-world-10x10-Q-matrix.txt")
    ut = maximum(qt, dims=2)[:]
    solver = ValueIterationSolver(verbose=true, init_util=ut, belres=1e-3)
    policy = solve(solver, mdp)
    return isapprox(ut, policy.util, rtol=1e-5)
end

function test_not_include_Q()
    # Load correct policy from file and verify we can reconstruct it
    rstates = [GridWorldState(4,3), GridWorldState(4,6), GridWorldState(9,3), GridWorldState(8,8)]
    rvals = [-10.0, -5.0, 10.0, 3.0]
    xs = 10
    ys = 10
    mdp = LegacyGridWorld(sx=xs, sy=ys, rs = rstates, rv = rvals)
    qt = readdlm("grid-world-10x10-Q-matrix.txt")
    ut = maximum(qt, dims=2)[:]
    niter = 100
    res = 1e-3
    solver = ValueIterationSolver(verbose=true, init_util=ut, belres=1e-3, include_Q=false)
    policy = solve(solver, mdp)
    return isapprox(ut, policy.util, rtol=1e-3)
end

function test_warning()
    mdp = LegacyGridWorld()
    solver = ValueIterationSolver()
    println("There should be a warning bellow: ")
    solve(solver, mdp, verbose=true)
end

function test_type_stability()
    mdp = LegacyGridWorld()
    solver = ValueIterationSolver()
    @inferred solve(solver, mdp)
    return true
end

@test test_complex_gridworld() == true
@test test_simple_grid() == true
@test test_init_solution() == true
@test test_not_include_Q() == true
test_warning()
@test_throws String solve(ValueIterationSolver(), TigerPOMDP())
@test test_type_stability() == true


# --- convolution

# http://developer.amd.com/resources/documentation-articles/articles-whitepapers/tiled-convolution-fast-image-filtering/
# has a detailed argument about why and how to do this better on a
# real gpu (although this case is simpler as it fits in a single tile)

"""convolve a 3x3 cross-shaped kernel."""
function convolve(input, border, kernel, output, n)
    @forall i j n begin
        a = 0
        # cannot use forneighbours because we need di, dj
        for (di, dj) in ((0, 1), (0, -1), (1, 0), (-1, 0))
            ii, jj = i+di, j+dj
            if ii > 0 && ii <= n && jj > 0 && jj <= n
                a = a + kernel[di+2, dj+2] * input[ii, jj]
            else
                a = a + kernel[di+2, dj+2] * border
            end
            output[i, j] = a
        end
    end
end

function blank!{N}(data, p::Position{N})
    @forall i j N begin
        if point(p, i, j) != empty
            data[i, j] = 0
        end
    end
end


using GaussianProcessODEs
using Test
using BenchmarkTools

using KernelFunctions
using QuadGK

rot(phi) = GaussianProcessODEs.Krot(phi)

struct altrotinvKernel{K <: KernelFunctions.Kernel, GKP }  <: GaussianProcessODEs.MatrixKernel
    kernel::K
    gkparams::GKP
end

# constructor
function altrotinvKernel(ker::K; N::Int) where K
    gkparams = NamedTuple{(:x, :weights)}(gauss(N, 0, 2pi))
    gkparams = zip(gkparams.x, gkparams.weights)
    gkparams = Iterators.product(gkparams, gkparams)
    altrotinvKernel{K, typeof(gkparams)}(ker, gkparams)
end

function (rk::altrotinvKernel)(z1, z2)
    gkp = rk.gkparams
    ker = rk.kernel

    mapreduce( g -> g[1][2]*ker(rot(g[1][1]) * z1, rot(g[2][1]) * z2)*g[2][2], +, gkp)
end

function (mker::altrotinvKernel)(w::Array{T, 1}) where {T <: Real}
    typeof(mker)(mker.kernel(w), mker.gkparams)
end

import KernelFunctions: kernelmatrix
function kernelmatrix(ker::altrotinvKernel, X)
    K = zeros(length(X), length(X))
    for i in 1:length(X), j in 1:i
        tmp = ker(X[i], X[j])
        K[i,j] = K[j,i] = tmp
    end
#     for (i, z1) in enumerate(X), (j,z2) in enumerate(X)
# #         K[i,j] = threadriker(ker, z1, z2)
#         K[i,j] = ker(z1,z2)
#     end
    return K
end

### setup

using LinearAlgebra
H(q,p) = norm(p)^2/2 - inv(norm(q))

borders = [[-2,2], [-2,2], [-2, 2], [-2, 2]]
ngp = 4
Z = rectanglegrid(borders, ngp) .+ [0.1*(rand(4) .- 0.5) for i in 1:ngp^length(borders)]
U = H.(getindex.(Z, Ref([1,2])), getindex.(Z, Ref([3,4])) );

### testing implementation above
@testset "invariance of reference" begin
    N = 55
    ker = pskernel(ones(5))
    riker = altrotinvKernel(ker; N);

    let phi = rand()*2pi
        @test (riker(Z[1], Z[2]) ≈ riker(rot(phi)*Z[1], Z[2]) ) atol = 1e-5
        @test (riker(Z[1], Z[2]) ≈ riker(Z[1], rot(phi)*Z[2]) ) atol = 1e-5
    end

end

### testing package implementation
@testset "invariance of package implementation" begin
    N = 55
    ker = pskernel(ones(5))
    tmpZ = Z[1:60];
    criker = crotinvKernel(ker; N)

    let phi = rand()*2pi
        @test (criker(Z[1], Z[2]) ≈ criker(rot(phi)*Z[1], Z[2]) ) atol = 1e-5
        @test (criker(Z[1], Z[2]) ≈ criker(Z[1], rot(phi)*Z[2]) ) atol = 1e-5
    end

end

### comparison
let
    N = 55
    ker = pskernel(ones(5))
    riker = altrotinvKernel(ker; N)
    criker = crotinvKernel(ker; N)
    tmpZ = Z[1:60]

    println("Comparison")
    println("reference")
    @time K1 = kernelmatrix(riker, tmpZ);
    K1 = kernelmatrix(riker, tmpZ);

    println("package")
    println("\r setup")
    @time tmpKx = Kx(criker, tmpZ);
    tmpKx = Kx(criker, tmpZ)
    println("\r kernel matrix")
    @time K2 = tmpKx();
    # @btime K2 = $tmpKx();
    K2 = tmpKx();

    @testset "accuracy" begin
        @test sum(K1 .- K2) < 1e-8
    end

    @testset "kernelmatrix()" begin
        @time K1 = kernelmatrix(criker, tmpZ)
        @time K2 = kernelmatrix(riker, tmpZ)
        @test K1 ≈ K2 atol=1e-6
    end
end

### test package implementation

## complete prediction
@testset "Kx" begin
    N = 55
    ker = pskernel(ones(5))
    riker = altrotinvKernel(ker; N)
    criker = crotinvKernel(ker; N)

    tmpZ = Z[1:60]
    tmpKx = Kx(criker, tmpZ)

    x = rand(4)
    @time Kx1 = tmpKx(x);
    @time Kx2 = riker.([x], tmpZ)';

    @test sum(Kx1[:] .- Kx2[:]) < 1e-6
end

@testset "Complete Prediction" begin
    N = 55
    ker = pskernel(ones(5))
    criker = crotinvKernel(ker; N)

    nZ = 6
    minZ = Z[1:nZ]
    minU = U[1:nZ]

    testker = pskernel([1., 12 .* ones(4)...])
    criker = crotinvKernel(testker; N)
    tmpKx2 = Kx(criker, minZ)
    K = tmpKx2()

    KinvU = K \ minU;
    # gpH(q,p) = (riker.([vcat(q,p)], permutedims(Zcl)) * KinvU)[1];
    gpf(z) = (tmpKx2(z) * KinvU)[1]

    @testset "recover data points" begin
        @test maximum(abs.(gpf.(minZ) .- minU)) < 1e-6
    end

    r1 = criker(Z[1], Z[2])
    phis = criker.gkparams.x
    rz1 = rot.(phis) .* [Z[1]]
    rz2 = rot.(phis) .* [Z[2]]

    r2 = criker(rz1, rz2)

    @testset "kernel itself" begin
        @test r1 ≈ r2 atol=1e-6
    end
end

### Derivative kernel

using Flux
@testset "Base derivative kernel" begin
    ker = pskernel(ones(5))
    dker = dKernel(ker)
    r1 = dker(Z[1], Z[2])
    r2 = gradient(x1 -> ker(x1, Z[2]), Z[1])[1]

    @test r1 ≈ r2 atol = 1e-6
end

### Invariant derivative kernel

@testset "Invariant derivative kernel" begin
    N = 60
    ker = pskernel(ones(5))
    dker = dKernel(ker)
    criker = crotinvKernel(ker; N)
    dcriker = dcrotinvKernel(dker; N)

    @time r1 = dcriker(Z[1], Z[2])
    @time r2 = gradient(x1 -> criker(x1, Z[2]), Z[1])[1]

    @test r1 ≈ r2 atol=1e-6
end

println("done")


using NetworkLayout, Roots
export MultigraphPlot, relabel, get_curves

function relabel(edgelist::Vector{Tuple{Int,Int}})
    edgelist_relab = Tuple{Int,Int}[]
    vmap = Dict{Int,Int}()
    i = 1
    for edge in edgelist
        for v in edge
            if v ∉ keys(vmap)
                vmap[v] = i
                i+=1
            end
        end
        push!(edgelist_relab, map(x->vmap[x], edge))
    end
    rev_vmap = Dict(v=>k for (k,v) in vmap)
    return edgelist_relab, rev_vmap
end

function relabel(edgelist::Vector{Tuple{String,String}})
    edgelist_relab = Tuple{Int,Int}[]
    vmap = Dict{String,Int}()
    i = 1
    for edge in edgelist
        for v in edge
            if v ∉ keys(vmap)
                vmap[v] = i
                i+=1
            end
        end
        push!(edgelist_relab, map(x->vmap[x], edge))
    end
    rev_vmap = Dict(v=>k for (k,v) in vmap)
    return edgelist_relab, rev_vmap
end

node_edge_intercept(x,w,h,r) = (4*h/(w^2) * x*(w-x))^2 + x^2 - r^2

function get_arc(
    src, dst;
    bendprop=0.1
    )
    # f(x) = 4h/w² x(w-x) (equation for arc from (0,0)->(w,0))
    # We then then translate and rotate this
    w = euclidean(src,dst)
    h = bendprop*euclidean(src,dst)
    θ = angle(complex(dst...)-complex(src...))
    cosθ, sinθ = (cos(θ), sin(θ))
    C = 4*h/w^2
    rawx = 0:0.01:w
    srcx,srcy = src
    xvals = [x*cosθ - C*x*(w-x)*sinθ + srcx for x in rawx]
    yvals = [x*sinθ + C*x*(w-x)*cosθ + srcy for x in rawx]
    return xvals, yvals
end

function get_arc_shorten(
    src, dst;
    bendprop=0.1,
    shorten=0.2
    )

    w = euclidean(src,dst)
    h = bendprop*euclidean(src,dst)
    x_shorten = find_zero(x->node_edge_intercept(x,w,h,shorten), (0.0,shorten))
    θ = angle(complex(dst...)-complex(src...))
    cosθ, sinθ = (cos(θ), sin(θ))
    C = 4*h/w^2
    rawx = x_shorten:0.01:(w-x_shorten)
    srcx,srcy = src
    xvals = [x*cosθ - C*x*(w-x)*sinθ + srcx for x in rawx]
    yvals = [x*sinθ + C*x*(w-x)*cosθ + srcy for x in rawx]
    return xvals, yvals
end

function get_curves(
    edges::AbstractVector,
    locs::AbstractVector;
    bendprop=0.1,
    shorten=0.2
    )

    xout, yout = (Vector{Float64}[],Vector{Float64}[])

    for (i,j) in edges
        xtmp, ytmp = get_arc_shorten(locs[i],locs[j],bendprop=bendprop,shorten=shorten)
        push!(xout,xtmp)
        push!(yout,ytmp)
    end

    return xout, yout

end

@userplot MultigraphPlot
@recipe function f(
    h::MultigraphPlot;
    relabelnodes=true,
    shorten=0.1,
    bendprop=0.1,
    seed=1
    )
    edgelist = h.args[1]

    do_relabel = (relabelnodes) || (eltype(edgelist)==Tuple{String,String})
    if do_relabel
        edgelist, rev_vmap = relabel(edgelist)
        V = length(keys(rev_vmap))
        adjmat = zeros(Int, V, V)
        for (i,j) in edgelist
            adjmat[i,j] += 1
        end
    end
    aspect_ratio --> 1
    locs = NetworkLayout.spring(adjmat,seed=seed)

    @series begin
        # seriestype := :line
        linecolor := :blue
        arrow --> true
        get_curves(edgelist, locs, shorten=shorten, bendprop=bendprop)
    end

    @series begin
        seriestype := :scatter
        markercolor := :blue
        locs
    end
end


using Formatting
using ArgParse, JSON
using NCDatasets

println("""
This program output the forcing file with QFLXT and QFLXS from QFLXT and QFLXS
variables in history files. It also needs the data folder generated by program
main_concat_and_convert_units.jl that has monthly TEMP and SALT profile.
""")

function runOneCmd(cmd)
    println(">> ", string(cmd))
    run(cmd)
end

function pleaseRun(cmd)
    if isa(cmd, Array)
        for i = 1:length(cmd)
            runOneCmd(cmd[i])
        end
    else
        runOneCmd(cmd)
    end
end

function parse_commandline()

    s = ArgParseSettings()

    @add_arg_table s begin

        "--input-file"
            help = "Casename"
            arg_type = String
            required = true

        "--output-file"
            help = "Casename"
            arg_type = String
            default = ""

    end

    return parse_args(s)
end

parsed = parse_commandline()

JSON.print(parsed,4)

output_file = parsed["output-file"]
input_file = parsed["input-file"]

Dataset(input_file, "r") do ds

    global Nx, Ny, Nz, Nt = size(ds["QFLXT"])

    ref = nomissing(ds["QFLXT"][:, :, :, 1], NaN)
    area = reshape(ds["area_sT"][:], Nx, Ny,  1)
    dz   = reshape(ds["dz_cT"][:],    1,  1, Nz)

    global valid_idx = isfinite.(ref)
    global mask = zeros(Float64, Nx, Ny, Nz)

    mask[valid_idx] .= 1.0

    tmp_QFLXT = zeros(Float64, Nx, Ny, Nz)
    tmp_QFLXS = copy(tmp_QFLXT)
    for t=1:Nt
        tmp_QFLXT .+= nomissing(ds["QFLXT"][:, :, :, t], 0.0)
        tmp_QFLXS .+= nomissing(ds["QFLXS"][:, :, :, t], 0.0)
    end
    tmp_QFLXT ./= Nt
    tmp_QFLXS ./= Nt

    vol = area .* dz
    vol .*= mask

    global sum_vol, total_QFLXT, total_QFLXS, mean_QFLXT, mean_QFLXS

    sum_vol = sum(vol)

    total_QFLXT = sum(vol .* tmp_QFLXT)
    total_QFLXS = sum(vol .* tmp_QFLXS)

    mean_QFLXT = total_QFLXT / sum_vol
    mean_QFLXS = total_QFLXS / sum_vol

end

println(format("Total volume : {:.2e}", sum_vol))
println(format("Total, Mean QFLXT   : {:.2e}, {:.2e}", total_QFLXT, mean_QFLXT))
println(format("Total, Mean QFLXS   : {:.2e}, {:.2e}", total_QFLXS, mean_QFLXS))

if output_file != ""
    println("Output file : $output_file")
    pleaseRun(`ncap2 -s "QFLXT=QFLXT-($(mean_QFLXT));QFLXS=QFLXS-($(mean_QFLXS));" $(input_file) $(output_file)`)
end


######################################################################
# general interface
######################################################################

export
    # general interface
    AbstractBasis,
    domain,
    collocation_points,
    degf,
    basis_matrix,
    InterpolatedFunction,
    collocation_values,
    # bases
    ChebyshevBasis,
    IntervalAB,
    LinearInterpolation,
    ExtrapolateLevel

abstract AbstractBasis

"""
Return the domain of the basis.
"""
function domain end

"A range of `n` linearly spaced elements in the domain of `basis`."
function linspace(basis::AbstractBasis, len)
    dom = domain(basis)
    linspace(dom.lo, dom.hi, len)
end

"""
Return the collocation points (if any) for a basis.
"""
function collocation_points end

"""
Degress of freedom (number of coefficients).
"""
degf(basis::AbstractBasis) = length(collocation_points(basis))

"""
Return the basis matrix, ie the functions in the basis evaluated at
the given points, one row for each point. Points default to
`collocation_points`.
"""
function basis_matrix end

"""
Interpolated function using `basis`, with coefficients `α`.
"""
immutable InterpolatedFunction{S, T}
    basis::S
    α::T
end

*(basis::AbstractBasis, α) = InterpolatedFunction(basis, α)

@forward InterpolatedFunction.basis domain, collocation_points, degf

function show(io::IO, ipf::InterpolatedFunction)
    print(io, "Interpolated function on ", ipf.basis)
end

"""
Evaluate the approximation on `basis` with coefficients `α` at at `x`.

Notes:
1. methods should be defined for various bases.
2. the callable interface is recommended for use outside this module.
"""
function _evaluate(basis, α, x)
    first(basis_matrix(basis, [x]) * α)
end

(f::InterpolatedFunction)(x) = _evaluate(f.basis, f.α, x)

"""
Evaluate the interpolated function at the collocation points.
"""
function collocation_values(ipf::InterpolatedFunction)
    ipf.(collocation_points(ipf.basis))
end

"""
Interpolate points in `y`, assumed to be values at the collocation
points, using `basis`.
"""
function \(basis::AbstractBasis, y::AbstractVector)
    # NOTE: this is just a fallback method, not necessarily efficient.
    @assert degf(basis) == length(y)
    InterpolatedFunction(basis, basis_matrix(basis) \ y)
end

"""
Interpolate `f`, evaluated at the collocation points, using `basis`.
"""
function \(basis::AbstractBasis, f)
    basis \ f.(collocation_points(basis))
end

zeros(basis::AbstractBasis) = basis \ zeros(degf(basis))

ones(basis::AbstractBasis) = basis \ ones(degf(basis))

# plot recipe for displaying interpolated functions
@recipe function f(ipf::InterpolatedFunction, label = "function", N = 100)
    @unpack basis, α = ipf
    dom = domain(basis)
    z = collect(linspace(dom.lo, dom.hi, N))
    @series begin
        seriestype := :path
        label := label
        z, ipf.(z)
    end
    @series begin
        primary := false
        seriestype := :scatter
        collocation_points(ipf), collocation_values(ipf)
    end
end

######################################################################
# Chebyshev basis on [-1,1]
######################################################################

immutable ChebyshevBasis <: AbstractBasis
    n::Int
    stretch::Bool
    extrapolate::Bool
    function ChebyshevBasis(n; stretch=false, extrapolate=false)
        new(n, stretch, extrapolate)
    end
end

size(b::ChebyshevBasis) = (b.n,)

domain(b::ChebyshevBasis, T=Float64) = -one(T)..one(T)

function collocation_points(b::ChebyshevBasis)
    zs = [cos((i-0.5)*pi/b.n) for i in b.n:-1:1]
    if b.stretch
        scale = 1/zs[end]
        zs .= zs .* scale
    end
    zs
end

function basis_matrix{T}(b::ChebyshevBasis,
                         x::AbstractVector{T}=collocation_points(b))
    if !b.extrapolate
        @assert all(abs(x) ≤ one(T) + √eps(T) for x in x)
    end
    A = Array{T}(length(x), b.n)
    for i in 1:b.n
        A[:,i] .= i == 1 ? one(T) : i == 2 ? x : (2.*x.*A[:, i-1]-A[:, i-2])
    end
    A
end

######################################################################
# IntervalAB
######################################################################

"""
Transform an inner basis to an interval.

Transforming to the inner basis is with `x ↦ fma(scale, shift)`.
"""
immutable IntervalAB{T,S <: AbstractBasis} <: AbstractBasis
    domain::Interval{T}
    inner_basis::S
    scale::T
    shift::T
    function IntervalAB(dom, inner_basis)
        @assert isfinite(dom)
        inner_dom = domain(inner_basis)
        scale = diam(inner_dom)/diam(dom)
        new(dom, inner_basis, scale, mid(inner_dom) - scale*mid(dom))
    end
end

function IntervalAB{T,S}(domain::Interval{T}, inner_basis::S)
    IntervalAB{T,S}(domain, inner_basis)
end

_map_to_inner(b::IntervalAB, x) = fma.(x, b.scale, b.shift)

_map_from_inner(b::IntervalAB, x) = fma.(x, 1/b.scale, -b.shift/b.scale)

domain(b::IntervalAB) = b.domain

function collocation_points(b::IntervalAB)
    _map_from_inner(b, collocation_points(b.inner_basis))
end

basis_matrix(b::IntervalAB) = basis_matrix(b.inner_basis)

basis_matrix(b::IntervalAB, x) = basis_matrix(b.inner_basis, _map_to_inner(b, x))

function show(io::IO, b::IntervalAB)
    print(io, b.inner_basis)
    print(io, " on ")
    print(io, b.domain)
end

######################################################################
# linear interpolation
######################################################################

immutable LinearInterpolation{T <: AbstractVector} <: AbstractBasis
    nodes::T
    function LinearInterpolation(nodes)
        @assert issorted(nodes) "Nodes need to be sorted"
        new(nodes)
    end
end

LinearInterpolation{T}(nodes::T) = LinearInterpolation{T}(nodes)

collocation_points(li::LinearInterpolation) = li.nodes

domain(li::LinearInterpolation) = li.nodes[1]..li.nodes[end]

basis_matrix(li::LinearInterpolation) = speye(length(li.nodes))

function \(li::LinearInterpolation, y::AbstractVector)
    @assert degf(li) == length(y)
    InterpolatedFunction(li, y)
end

function _evaluate(li::LinearInterpolation, α, x)
    @unpack nodes = li
    (nodes[1] ≤ x ≤ nodes[end]) || error(DomainError())
    index = searchsortedlast(nodes, x)
    @assert index ≠ 0 "internal error"
    # println("index = $index, node = $(nodes[index])")
    if x == nodes[index]
        α[index]
    else
        p = (x - nodes[index]) / (nodes[index+1] - nodes[index])
        α[index]*(1-p) + α[index+1]*p
    end
end

function collocation_values{S <: LinearInterpolation, T}(ipf::InterpolatedFunction{S,T})
    ipf.α
end

######################################################################
# ExtrapolateLevel
######################################################################

immutable ExtrapolateLevel{T} <: AbstractBasis
    inner_basis::T
end

@forward ExtrapolateLevel.inner_basis domain, collocation_points, basis_matrix, degf

function _evaluate(basis::ExtrapolateLevel, α, x)
    @unpack inner_basis = basis
    dom = domain(inner_basis)
    _evaluate(inner_basis, α, clamp(x, dom.lo, dom.hi))
end


using MAT
using Convex
using ECOS
using Plots
using LinearAlgebra
using LineSearches
using ProgressBars
using WoodburyMatrices
using Printf: @sprintf

basis = matread("../Curvelet_Basis_Construction/G_32_32.mat")
G = basis["G_mat"]
s = basis["scales"]
G = G.*(2 .^ (-2 .*s))
Gn = maximum(sqrt.(sum(G.^2, dims=2)))
G ./= Gn

function d_oed(X::Matrix{T}, B, ϵ, λ; maxiter=2500, ls=MoreThuente(), g_rtol::T = sqrt(eps(T)), g_atol::T = eps(T)) where T <: Number
    #initial vector
    n = size(X, 1)
    μ = ones(n, 1)
    μp = Variable(n, 1)
    gvec = similar(μ)
    s = similar(μ)
    #setup linesearch functions
    #linesearch - this needs to make sure that any interior functions actually have function barriers around global variables

    function Ml(μ, X, ϵ)
        println(Diagonal(μ))
        M = ϵ*I + X'*Diagonal(μ)*X
        Mc = cholesky(Hermitian(M))
        l = μ.*diag(X*(Mc\X'))
        return (M, Mc, l)
    end

    Ml(μ) = Ml(μ, X, ϵ)

    function f(μ, λ)
        M, Mc, l = Ml(μ)
        return -logdet(M) + λ*maximum(l)
    end

    f(μ) = f(μ, λ)

    function gi!(gvec, i, Mc, μ, u, lmind, X, λ, T)
        ld_term = @views -X[i,:]'*(Mc\X[i,:])
        tmp_A = SymWoodbury(Mc, X[i,:], -μ[i])
        tmp_b = @views tmp_A \ X[i,:]
        tmp_1 = (u'*tmp_b)^2
        tmp_d = @views X[i,:]'*tmp_b
        tmp_den = (1+μ[i]*tmp_d)^2
        lev_term_1 =  -μ[lmind]*((1+μ[i]*tmp_d)*tmp_1 - μ[i]*tmp_1*tmp_d) / tmp_den
        lev_term_2 = (i==lmind) ? u'*(Mc\u) : zero(T)
        gvec[i] = ld_term + λ*(lev_term_1 + lev_term_2)
    end

    gi!(gvec, i, Mc, μ, u, lmind) = gi!(gvec, i, Mc, μ, u, lmind, X, λ, T)

    function g!(gvec, μ)
        M, Mc, l = Ml(μ)
        lm, lmind = findmax(l)
        lmind=lmind[1]
        u = @views X[lmind, :]
        for i = 1:n
            gi!(gvec, i, Mc, μ, u, lmind)
        end
    end

    function fg!(gvec, μ)
        M, Mc, l = Ml(μ)
        lm, lmind = findmax(l)
        lmind=lmind[1]
        u = @views X[lmind, :]
        for i = 1:n
            gi!(gvec, i, Mc, μ, u, lmind)
        end
        return -logdet(M) + λ*maximum(l)
    end

    ϕ(α) = f(μ.+α.*s)

    function dϕ(α)
        g!(gvec, μ .+ α.*s)
        return dot(gvec, s)
    end

    function ϕdϕ(α)
        phi = fg!(gvec, μ .+ α.*s)
        dphi = dot(gvec, s)
        return (phi, dphi)
    end

    #get initial gradient and iterate
    g!(gvec, μ)
    fμ = f(μ)
    gnorm = norm(gvec)
    gtol = max(g_rtol*gnorm, g_atol)
    iter = ProgressBar(1:maxiter)
    set_description(iter, string(@sprintf("Starting Up")))

    for i in iter
        if gnorm < gtol
            break
        else
            #set step direction
            s .= -gvec
            dϕ_0 = dot(s, gvec)
            #linesearch gradient
            α, fμ = ls(ϕ, dϕ, ϕdϕ, 1.0, fμ, dϕ_0)
            # α = 0.001
            # fμ = f(μ)
            @. μ = μ + α*s
            #project into constrained space - decrease budget as we go along

            projection = minimize(sumsquares(μp-μ), [0<=μp, μp<=1, sum(μp)<=max(n-i, B)])
            solve!(projection, () -> ECOS.Optimizer(verbose=false))
            μ .= μp.value
            #update gradient
            g!(gvec, μ)
            gnorm = norm(gvec)

            set_description(iter, string(@sprintf("Objective: %.5f", fμ)))
        end
    end

    return μ, round.(μ)
end

resuround, res = d_oed(G, 23, 1.0, 10.0, maxiter=10)


### A Pluto.jl notebook ###
# v0.12.18

using Markdown
using InteractiveUtils

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

# ╔═╡ 87fefce0-02c5-11eb-3bc6-ad05fe23cf93
begin
	import Pkg
	Pkg.activate(mktempdir())
	Pkg.add("Plots")
	Pkg.add("PlutoUI")
	using Plots
	using PlutoUI
end

# ╔═╡ 3a7f37b0-02c9-11eb-28d5-1d8d43c6f44a
md"### Curva de Juros Interativa (Nelson-Siegel-Svensson)"

# ╔═╡ 5e857a22-02c9-11eb-2c2b-97f06fbcf1a8
md"Este pluto notebook permite interagir com os 6 parâmetros **β₀, β₁, β₂, β₃, λ₁, λ₂** do modelo de curva de juros de Nelson-Siegel-Svensson, gerando a visualização interativa da curva de juros.

Deste modo, é possível observar o impacto que cada parâmetro tem sobre o comportamento da curva de juros."

# ╔═╡ abc28efe-02db-11eb-12a1-e3aa3ea94d5f
md"O modelo segue a seguinte formulação:

$y_{t}(\tau)=\beta_{0, t}+\beta_{1, t} \left(\frac{1-e^{-\lambda_{1, x} \tau}}{\lambda_{1, t} \cdot \tau}\right)+\beta_{2, t} \left(\frac{1-e^{-\lambda_{1, x} \tau}}{\lambda_{1, t} \cdot \tau}-e^{-\lambda_{1, x} \cdot \tau}\right)+\beta_{3, t} \left(\frac{1-e^{-\lambda_{2, x} \cdot \tau}}{\lambda_{2, t} \cdot \tau}-e^{-\lambda_{2, x} \cdot \tau}\right)$"

# ╔═╡ c633a440-02c8-11eb-185a-8b3e594df121
begin
	n = @bind n Slider(5:100, default=25, show_value=true)
	β₀ = @bind β₀ Slider(0.01:0.001:1.0, default=0.05, show_value=true)
	β₁ = @bind β₁ Slider(-1:0.01:1.0, default=-0.05, show_value=true)
	β₂ = @bind β₂ Slider(-1:0.01:1.0, default=-0.1, show_value=true)
	β₃ = @bind β₃ Slider(-1:0.01:1.0, default=-0.05, show_value=true)
	λ₁ = @bind λ₁ Slider(0.01:0.01:2.0, default=0.1, show_value=true)
	λ₂ = @bind λ₂ Slider(0.01:0.01:1.0, default=0.08, show_value=true)

	md"""**Parâmetros da Curva de Juros**

	**N°**: $(n) 	Número de Obs. (21 du)

	**β₀**: $(β₀)	Fator de Nível (Longo Prazo)

	**β₁**: $(β₁)	Fator de Inclinação (Curto Prazo)

	**β₂**: $(β₂)	Fator de Curvatura (Médio Prazo)

	**β₃**: $(β₃)	Fator de Curvatura (Médio Prazo)

	**λ₁**: $(λ₁)   Velocidade de Decaimento (Médio Prazo)

	**λ₂**: $(λ₂)	Velocidade de Decaimento (Médio Prazo)"""
end

# ╔═╡ bdab38e0-02c5-11eb-09cf-43a08fb6bf8b
begin
	x = collect(21:21:(n*21))
	x[1] = 21
	@show x
end

# ╔═╡ b74b0b10-02ca-11eb-17cb-f5c8fdc4ee50
begin
	nss(β₀, β₁, β₂, β₃, λ₁, λ₂, t) =   β₀ +
									β₁*((1-exp(-λ₁*t))/(λ₁*t)) +
									β₂*((1-exp(-λ₁*t))/(λ₁*t) - exp(-λ₁*t)) +
									β₃*((1-exp(-λ₂*t))/(λ₂*t) - exp(-λ₂*t))
	y = [nss(β₀, β₁, β₂, β₃, λ₁, λ₂, i)*100 for i in x]
end

# ╔═╡ 3a9d9420-02d4-11eb-2df5-a12b2a9eecb8
begin
	gr()
	plot(x, y,
		leg=false,
		title="Curva de Juros (NSS)",
		xlabel="Vencimento em dias úteis", ylabel="% a.a",
		ylim=(-3,10),
		c=:blue, lw=3,
		shape=:hexagon,
		msize=4, # msize equivale a markersize
		markercolor=:white,
		markeralpha=1)
	plot!([-1, n*21], [0,0], c=:black, lab="", l=:dash)
end

# ╔═╡ Cell order:
# ╟─3a7f37b0-02c9-11eb-28d5-1d8d43c6f44a
# ╟─5e857a22-02c9-11eb-2c2b-97f06fbcf1a8
# ╟─abc28efe-02db-11eb-12a1-e3aa3ea94d5f
# ╟─87fefce0-02c5-11eb-3bc6-ad05fe23cf93
# ╟─c633a440-02c8-11eb-185a-8b3e594df121
# ╟─3a9d9420-02d4-11eb-2df5-a12b2a9eecb8
# ╟─bdab38e0-02c5-11eb-09cf-43a08fb6bf8b
# ╟─b74b0b10-02ca-11eb-17cb-f5c8fdc4ee50


using Scat
using FFTW
using Test

include("test_filterBank.jl")
include("test_filter.jl")
include("test_scattered.jl")
include("test_scatteringTransform.jl")
include("test_scatter1d.jl")


Random.seed!(123)

n = 2000
h = readdlm("h.txt")
A = Conv(Float64,(n,),h[:])
m = size(A,1)
#A = hcat([[zeros(i);h;zeros(n-1-i)] for i = 0:n-1]...) # Equivalent Full Matrix

x_star = Vector(sprandn(n,0.5))
b = A*x_star+20*randn(m)

f = Translate(SqrNormL2(), -b)
lam = 1e-1*norm(A'*b, Inf)
g = NormL1(lam)

## fast FBS/Adaptive
x0 = zeros(n)
@time it, xFBS, sol = ProximalAlgorithms.FBS(x0; fq=f, Aq=A, g=g, fast = true, tol = 1e-8)
@test it < 1000

# ZeroFPR/Adaptive
x0 = zeros(n)
@time it, xZ, sol = ProximalAlgorithms.ZeroFPR(x0; fq=f, Aq=A, g=g, tol = 1e-8)

@test norm(xFBS-xZ) < 1e-3
@test it < 70

# PANOC/Adaptive
x0 = zeros(n)
@time it, xP, sol = ProximalAlgorithms.PANOC(x0; fq=f, Aq=A, g=g, tol = 1e-8)

@test norm(xP-xZ) < 1e-3
@test it < 140

#####################
# Complex Variables #
#####################

Random.seed!(123)
m, n = 10,100

A = randn(m,n)+im*randn(m,n)
U,S,V = svd(A)
#create ill conditioned matrix
S[floor(Int,0.5*min(m,n)):end] .= 0.
A = U*diagm(0 => S)*V'

x_star = sprandn(n,0.5).+im.*sprandn(n,0.5)
b = A*x_star+20*(randn(m).+im*randn(m))

f = Translate(SqrNormL2(), -b)
f2 = LeastSquares(A, b)
lam = 1e-2*norm(A'*b, Inf)
g = NormL1(lam)

x0 = zeros(n)+im*zeros(n)
it, x_star, sol = ProximalAlgorithms.PANOC(x0; fq=f, Aq=A, g=g, verbose = 0, tol = 1e-9)

# fast FBS/Adaptive
x0 = zeros(n)+im*zeros(n)
@time it, x, sol = ProximalAlgorithms.FBS(x0; fq=f, Aq=A, g=g, fast = true, tol = 1e-6, maxit = 30000)

@test norm(x-x_star) < 1e-4
@test it < 25000

# ZeroFPR/Adaptive
x0 = zeros(n)+im*zeros(n)
@time it, x, sol = ProximalAlgorithms.ZeroFPR(x0; fq=f, Aq=A, g=g, tol = 1e-6)

@test norm(x-x_star) < 1e-4
@test it < 700

# PANOC/Adaptive
x0 = zeros(n)+im*zeros(n)
@time it, x, sol = ProximalAlgorithms.PANOC(x0; fq=f, Aq=A, g=g, tol = 1e-6)

@test norm(x-x_star) < 1e-4
@test it < 1500


using CompScienceMeshes

m1 = meshrectangle(1.0, 0.5, 0.1, 3)
m1 = CompScienceMeshes.translate(m1, [0.0, 0.5, 0.0])
m2 = meshrectangle(1.0, 0.5, 0.1, 3)

nc1 = numcells(m1)
nc2 = numcells(m2)

nv1 = numvertices(m1)
nv2 = numvertices(m2)

M1 = weld(m1,m2)
@assert numcells(M1) == nc1 + nc2
@assert numvertices(M1) == nv1 + nv2 - 11

M2 = weld(m1,-m2)
@assert numcells(M2) == nc1 + nc2
@assert numvertices(M2) == nv1 + nv2 - 11

@assert CompScienceMeshes.isoriented(M1)
@assert !CompScienceMeshes.isoriented(M2)

e1 = skeleton(m1,1)
e2 = skeleton(m2,1)
E1 = skeleton(M1,1)

S11 = CompScienceMeshes.embedding(e1,E1)
S21 = CompScienceMeshes.embedding(e2,E1)


# Density at a given lattice site
density(fᵢ::DenseVector) = sum(fᵢ)
function density(fᵢ::DenseArray)
    shape = size(fᵢ)[2:end]
    reshape([sum(fᵢ[:, i]) for i = 1:prod(shape)], shape)
end
density(sim::Simulation) = density(sim.populations)
# Momentum at given site
function momentum!(μ_out::DenseVector, fᵢ::DenseVector, cᵢ::DenseMatrix)
    @assert length(μ_out) == size(cᵢ, 1) && length(fᵢ) == size(cᵢ, 2)
    μ_out[:] = 0
    for i in 1:length(μ_out)
        for j in 1:length(fᵢ)
            if cᵢ[i, j] == 1
                μ_out[i] += fᵢ[j]
            elseif cᵢ[i, j] == -1
                μ_out[i] -= fᵢ[j]
            end
        end
    end
    μ_out
end
momentum(fᵢ::DenseVector, cᵢ::DenseMatrix) = cᵢ * fᵢ
# Velocities at a given lattice site
velocity(μ::DenseVector, ρ::Number) = μ / ρ
velocity(fᵢ::DenseVector, cᵢ::DenseMatrix, ρ::Number) = velocity(momentum(fᵢ, cᵢ), ρ)
velocity(fᵢ::DenseVector, cᵢ::DenseMatrix) = velocity(fᵢ, cᵢ, density(fᵢ))
function velocity!(ν_out::DenseVector, μ::DenseVector, ρ::Number)
    ν_out[:] = μ
    scale!(ν_out, 1 / ρ)
end

# Momentum and velocity functions that take the full grid, or a simulation object
for name in [:momentum, :velocity]
    @eval begin
        function $name(fᵢ::DenseArray, cᵢ::DenseMatrix)
            const shape = tuple(size(cᵢ, 1), size(fᵢ)[2:end]...)
            result = zeros(eltype(fᵢ), shape)
            for i = 1:prod(shape[2:end])
                result[1:shape[1], i] = $name(fᵢ[:, i], cᵢ)
            end
            result
        end
        $name(sim::Simulation) = $name(sim.populations, sim.lattice.celerities)
    end
end

#= Computes the equilibrium particle distributions $f^{eq}$:

    momentum: Macroscopic momentum μ at current lattice site
    celerities: d by n matrix of celerities ē for the lattice, with d the dimensionality of the
        lattice and n the number of particle distributions.
    weights: Weights associated with each celerity
    ρ: Density

    $f^{eq}$ = weights .* [ρ + 3ē⋅μ + \frac{9}{2ρ} (ē⋅μ)² - \frac{3}{2ρ} μ⋅μ]$
=#
function equilibrium{T, I}(
        ρ::T, momentum::DenseVector{T}, celerities::DenseMatrix{I}, weights::DenseVector{T})
    # computes momentum projected on each particle celerity first
    @assert length(momentum) == size(celerities, 1)
    @assert length(weights) == size(celerities, 2)
    μ_on_ē = celerities.'momentum
    weights .* (
        ρ
        + 3μ_on_ē
        + 9/(2ρ) * (μ_on_ē .* μ_on_ē)
        - 3/(2ρ) * dot(momentum, momentum)
    )
end
equilibrium{T, I}(lattice::Lattice{T, I}, fᵢ::DenseVector{T}) =
    equilibrium(density(fᵢ), momentum(fᵢ, lattice.celerities), lattice.celerities, lattice.weights)

equilibrium{T}(lattice::Lattice, ρ::T, momentum::DenseVector{T}) =
    equilibrium(ρ, momentum, lattice.celerities, lattice.weights)
equilibrium{T}(lattice::Symbol, ρ::T, momentum::DenseVector{T}) =
    equilibrium(getfield(LB, lattice), ρ, momentum)
# Computes equilibrium without memory-allocation
# Incoming must have the right size
function equilibrium!{T, I}(
        feq::DenseVector{T}, ρ::T, momentum::DenseVector{T}, celerities::DenseMatrix{I},
        weights::DenseVector{T})
    @assert length(feq) == size(celerities, 2)
    @assert length(momentum) == size(celerities, 1)
    @assert length(weights) == size(celerities, 2)
    const inv_tworho = 1 / (2ρ)
    # Interestingly enough, celerities[:, i] is fairly slower than celerities[j:j+d], where j is
    # incremented at each step of the loop
    const inc = length(momentum)
    const d = inc - 1
    j = 1
    for i in 1:length(feq)
        const μ_on_ē = 3dot(celerities[j:j+d], momentum)
        j += inc
        feq[i] = weights[i] * (
            ρ + μ_on_ē + (μ_on_ē * μ_on_ē - 3dot(momentum, momentum)) * inv_tworho)
    end
end


function test_simple_net_gradient(backend)
    println("-- Testing gradients on simple network (example for gradient checking code)")

    srand(12345678)
    ############################################################
    # Prepare Random Data
    ############################################################
    N = 5    # works with arbitrary minibatch size as long as
             # N == batch_size in MemoryDataLayer so it cycles through
             # and gets the same data during forward()
    M = 10
    P = 4

    X = rand(M, N)
    W = rand(M, P)
    B = rand(P, 1)

    Y = (W'*X .+ B)
    Y = Y + 0.01*randn(size(Y))

    ############################################################
    # Define network
    ############################################################
    data_layer = MemoryDataLayer(batch_size=N, data=Array[X,Y])

    w1 = InnerProductLayer(neuron=Neurons.Sigmoid(), name="ip1",output_dim=20, tops=[:a], bottoms=[:data])
    w2 = InnerProductLayer(neuron=Neurons.Identity(), name="ip2",output_dim=4, tops=[:b], bottoms=[:a])
    loss_layer = SquareLossLayer(name="loss", bottoms=[:b, :label] )

    net = Net("TEST", backend, [w1,w2, loss_layer, data_layer])

    # epsilon:     milage may vary 1e-4 - 1e-8
    # digit:       compare this many digits to check for 'identity'
    # visualize:   prints out correct and failed positions of gradients
    return test_gradients(net, epsilon=1e-8, digit=6, visual=false )
end

if test_cpu
    test_simple_net_gradient(backend_cpu)
end


# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Desription
# ==============================================================================
#
#   Tests related to conversion from direction cosine matrices to Euler angle
#   and axis.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

# File: ./src/conversions/dcm_to_angleaxis.jl
# ===========================================

# Functions: dcm_to_angleaxis
# ---------------------------

@testset "DCM => Euler angle and axis (Float64)" begin
    T = Float64

    # Create a random DCM.
    D = angle_to_dcm(_rand_ang(T), :X) *
        angle_to_dcm(_rand_ang(T), :Y) *
        angle_to_dcm(_rand_ang(T), :Z)

    # Convert to Euler angle and axis.
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T

    # Check if the rotation expressed by D is consistent, which is performed in
    # two steps:
    #
    #   1. A vector aligned with `v` does not change.
    #   2. A vector perpendicular to `v` is rotated by `a`.

    v = av.v
    vr = D * v
    @test vr ≈ v
    @test eltype(v) === eltype(vr) === T

    # Auxiliary vector to obtain a vector perpendicular to `v`.
    aux = @SVector randn(3)
    aux = aux / norm(aux)

    vp = av.v × aux
    vp = vp / norm(vp)
    vpr = D * vp
    @test eltype(vp) === eltype(vpr) === T

    # Compute the angle between vp and vpr in [0, 2π].
    a = acos(vp ⋅ vpr)
    @test a ≈ av.a
end

@testset "DCM => Euler angle and axis (Float32)" begin
    T = Float32

    # Create a random DCM.
    D = angle_to_dcm(_rand_ang(T), :X) *
        angle_to_dcm(_rand_ang(T), :Y) *
        angle_to_dcm(_rand_ang(T), :Z)

    # Convert to Euler angle and axis.
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T

    # Check if the rotation expressed by D is consistent, which is performed in
    # two steps:
    #
    #   1. A vector aligned with `v` does not change.
    #   2. A vector perpendicular to `v` is rotated by `a`.

    v = av.v
    vr = D * v
    @test vr ≈ v
    @test eltype(v) === eltype(vr) === T

    # Auxiliary vector to obtain a vector perpendicular to `v`.
    aux = @SVector randn(T, 3)
    aux = aux / norm(aux)

    vp = av.v × aux
    vp = vp / norm(vp)
    vpr = D * vp
    @test eltype(vp) === eltype(vpr) === T

    # Compute the angle between vp and vpr in [0, 2π].
    a = acos(vp ⋅ vpr)
    @test a ≈ av.a
end

@testset "DCM => Euler angle and axis (Special cases)" begin
    # Float64
    # ==========================================================================

    T = Float64

    D = DCM(T(1) * I)
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T
    @test av.a ≈ 0
    @test av.v ≈ [0, 0, 0]

    D = angle_to_dcm(T(π), :X)
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T
    @test av.a ≈ π
    @test av.v ≈ [1, 0, 0]

    D = angle_to_dcm(T(π), :Y)
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T
    @test av.a ≈ π
    @test av.v ≈ [0, 1, 0]

    D = angle_to_dcm(T(π), :Z)
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T
    @test av.a ≈ π
    @test av.v ≈ [0, 0, 1]

    D = angle_to_dcm(T(2π / 3), :X)
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T
    @test av.a ≈ T(2π / 3)
    @test av.v ≈ [1, 0, 0]

    D = angle_to_dcm(T(-2π / 3), :X)
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T
    @test av.a ≈ T(2π / 3)
    @test av.v ≈ [-1, 0, 0]

    D = angle_to_dcm(T(2π / 3), :Y)
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T
    @test av.a ≈ T(2π / 3)
    @test av.v ≈ [0, 1, 0]

    D = angle_to_dcm(T(-2π / 3), :Y)
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T
    @test av.a ≈ T(2π / 3)
    @test av.v ≈ [0, -1, 0]

    D = angle_to_dcm(T(2π / 3), :Z)
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T
    @test av.a ≈ T(2π / 3)
    @test av.v ≈ [0, 0, 1]

    D = angle_to_dcm(T(-2π / 3), :Z)
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T
    @test av.a ≈ T(2π / 3)
    @test av.v ≈ [0, 0, -1]

    D = angle_to_dcm(T(-2π / 3), :Z) * angle_to_dcm(T(-π / 2), :Y)
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T
    @test av.a ≈ 2.4188584057763776
    @test av.v ≈ [0.6546536707079772, -0.37796447300922736, -0.6546536707079772]

    D = DCM(
        T(-1 / 3), T(-2 / 3), T(-1 / 3),
        T(-1 / 3), T(-1 / 3), T(+2 / 3),
        T(-2 / 3), T(+2 / 3), T(-1 / 3)
    )'
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T
    @test av.a ≈ T(π)
    @test av.v ≈ [T(√3 / 3), T(-√3 / 3), T(-√3 / 3)]

    # Float32
    # ==========================================================================

    T = Float32

    D = DCM(T(1) * I)
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T
    @test av.a ≈ 0
    @test av.v ≈ [0, 0, 0]

    D = angle_to_dcm(T(π), :X)
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T
    @test av.a ≈ π
    @test av.v ≈ [1, 0, 0]

    D = angle_to_dcm(T(π), :Y)
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T
    @test av.a ≈ π
    @test av.v ≈ [0, 1, 0]

    D = angle_to_dcm(T(π), :Z)
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T
    @test av.a ≈ π
    @test av.v ≈ [0, 0, 1]

    D = angle_to_dcm(T(2π / 3), :X)
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T
    @test av.a ≈ T(2π / 3)
    @test av.v ≈ [1, 0, 0]

    D = angle_to_dcm(T(-2π / 3), :X)
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T
    @test av.a ≈ T(2π / 3)
    @test av.v ≈ [-1, 0, 0]

    D = angle_to_dcm(T(2π / 3), :Y)
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T
    @test av.a ≈ T(2π / 3)
    @test av.v ≈ [0, 1, 0]

    D = angle_to_dcm(T(-2π / 3), :Y)
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T
    @test av.a ≈ T(2π / 3)
    @test av.v ≈ [0, -1, 0]

    D = angle_to_dcm(T(2π / 3), :Z)
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T
    @test av.a ≈ T(2π / 3)
    @test av.v ≈ [0, 0, 1]

    D = angle_to_dcm(T(-2π / 3), :Z)
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T
    @test av.a ≈ T(2π / 3)
    @test av.v ≈ [0, 0, -1]

    D = angle_to_dcm(T(-2π / 3), :Z) * angle_to_dcm(T(-π / 2), :Y)
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T
    @test av.a ≈ 2.4188584057763776
    @test av.v ≈ [0.6546536707079772, -0.37796447300922736, -0.6546536707079772]

    D = DCM(
        T(-1 / 3), T(-2 / 3), T(-1 / 3),
        T(-1 / 3), T(-1 / 3), T(+2 / 3),
        T(-2 / 3), T(+2 / 3), T(-1 / 3)
    )'
    av = dcm_to_angleaxis(D)
    @test eltype(av) === T
    @test av.a ≈ T(π)
    @test av.v ≈ [T(√3 / 3), T(-√3 / 3), T(-√3 / 3)]
end


# Gradient boosting.
module GB

using Random
using GradientBoost.Util
using GradientBoost.LossFunctions

export GBAlgorithm,
       GBModel,
       stochastic_gradient_boost,
       fit,
       predict,
       build_base_func,
       create_sample_indices

# Gradient boost algorithm.
abstract type GBAlgorithm end

# Gradient boost model.
mutable struct GBModel{F <: AbstractFloat}
  learning_rate::F
  base_funcs::Vector{Function}
end

# Perform stochastic gradient boost.
#
# @param gb Gradient boosting algorithm.
# @param instances Instances.
# @param labels Labels.
# @return Gradient boost model.
function stochastic_gradient_boost(gb::GBAlgorithm, instances, labels)
  # Initialize base functions collection
  num_iterations = gb.num_iterations
  base_funcs = Array{Function, 1}(undef, num_iterations+1)

  # Create initial base function
  initial_val = minimizing_scalar(gb.loss_function, labels)
  initial_base_func = (instances) -> fill(initial_val, size(instances, 1))

  # Add initial base function to ensemble
  base_funcs[1] = initial_base_func

  # Build consecutive base functions
  prev_func_pred = initial_base_func(instances)
  for iter_ind = 2:num_iterations+1
    # Obtain current residuals
    psuedo = negative_gradient(
      gb.loss_function,
      labels,
      prev_func_pred
    )

    # Sample instances
    stage_sample_ind = create_sample_indices(gb, instances, labels)

    # Build current base function
    stage_base_func = build_base_func(
      gb,
      instances[stage_sample_ind, :],
      labels[stage_sample_ind],
      prev_func_pred[stage_sample_ind],
      psuedo[stage_sample_ind]
    )

    # Update previous function prediction
    prev_func_pred .+= gb.learning_rate .* stage_base_func(instances)

    # Add optimal base function to ensemble
    base_funcs[iter_ind] = stage_base_func
  end

  # Return model
  return GBModel(gb.learning_rate, base_funcs)
end

function fit(gb::GBAlgorithm, instances, labels)
  stochastic_gradient_boost(gb, instances, labels)
end
function predict(gb_model::GBModel, instances)
  outputs = gb_model.base_funcs[1](instances)
  for i = 2:length(gb_model.base_funcs)
    outputs .+= gb_model.learning_rate .* gb_model.base_funcs[i](instances)
  end
  return outputs
end

# Build base (basis) function for gradient boosting algorithm.
#
# @param gb Gradient boosting algorithm.
# @param instances Instances.
# @param labels Labels.
# @param prev_func_pred Previous base function's predictions.
# @param psuedo Psuedo-labels (psuedo-response).
# @return Function of form (instances) -> predictions.
function build_base_func(
  gb::GBAlgorithm,
  instances,
  labels,
  prev_func_pred,
  psuedo)

  err_must_be_overriden()
end

# Default sample method for gradient boosting algorithms.
# By default, it is sampling without replacement.
#
# @param gb Gradient boosting algorithm.
# @param instances Instances.
# @param labels Labels.
# @return Sample indices.
function create_sample_indices(gb::GBAlgorithm, instances, labels)
  n = size(instances, 1)
  prop = gb.sampling_rate

  ind = randperm(n)[1:Int(round(prop * n))]
end

end # module


# Computed via wolframalpha.com: SinIntegral[SetPrecision[Table[x,{x, 1,20,1}],20]] and CosIntegral[SetPrecision[Table[x,{x, 1,20,1}],20]]

@testset "trigonometric integrals" begin
    sinintvals = [  0.9460830703671830149, 1.6054129768026934849, 1.848652527999468256,
                    1.75820313894905306, 1.54993124494467414, 1.4246875512805065,
                    1.4545966142480936, 1.5741868217069421, 1.665040075829602,
                    1.658347594218874, 1.578306806945727416, 1.504971241526373371,
                    1.499361722862824564, 1.556211050077665054, 1.618194443708368739,
                    1.631302268270032886, 1.590136415870701122, 1.536608096861185462,
                    1.518630031769363932, 1.548241701043439840]
    cosintvals = [0.3374039229009681347, 0.4229808287748649957, 0.119629786008000328,
                    -0.14098169788693041, -0.19002974965664388, -0.06805724389324713,
                    0.07669527848218452, 0.122433882532010, 0.0553475313331336,
                    -0.045456433004455, -0.08956313549547997948, -0.04978000688411367560,
                    0.02676412556403455504, 0.06939635592758454727, 0.04627867767436043960,
                    -0.01420019012019002240, -0.05524268226081385053, -0.04347510299950100478,
                    0.00515037100842612857, 0.04441982084535331654]
    for x in 1:20
        @test sinint(x)     ≅ sinintvals[x]
        @test sinint(-x)    ≅ -sinintvals[x]
        @test cosint(x)     ≅ cosintvals[x]
    end

    @test sinint(Float16(1)) == Float16(sinint(1))
    @test cosint(Float16(1)) == Float16(cosint(1))
    @test sinint(Float32(1)) == Float32(sinint(1))
    @test cosint(Float32(1)) == Float32(cosint(1))

    @test sinint(1//2) == sinint(0.5)
    @test cosint(1//2) == cosint(0.5)

    @test sinint(1e300)     ≅ π/2
    @test cosint(1e300)     ≅ -8.17881912115908554103E-301
    @test sinint(1e-300)    ≅ 1.0E-300
    @test cosint(1e-300)    ≅ -690.1983122333121

    @test sinint(Inf) == π/2
    @test cosint(Inf) == 0.0
    @test isnan(sinint(NaN))
    @test isnan(cosint(NaN))

    @test_throws ErrorException sinint(big(1.0))
    @test_throws ErrorException cosint(big(1.0))
    @test_throws DomainError cosint(-1.0)
    @test_throws DomainError cosint(Float32(-1.0))
    @test_throws DomainError cosint(Float16(-1.0))
    @test_throws DomainError cosint(-1//2)
    @test_throws MethodError sinint(Complex(1))
    @test_throws MethodError cosint(Complex(1))
end


println("Loading Packges for BayesLM test...")
using BayesLM
using Base.Test
println("Starting Tests for BayesLM test...")

@testset "lm" begin
  srand(1);
  n = 100
  X = randn(n)
  b = [0,3]
  sig = .5
  y = [ones(n) X] * b + randn(n) * sig

  @time model = lm(y,X,B=10000);

  #= private tests
  using RCall
  @rput y X n b sig;
  R"summary(mod <- lm(y~X))"
  R"plot(X,y)"
  =#

  post_sum = summary(model,alpha=.05)
  println(post_sum)
  @test all(abs(b - post_sum.mean_beta) .< .2)
  @test abs(sig - post_sum.mean_sig) < .1
  @time CPO = cpo(model,y,[ones(n) X])

  #= visual tests
  @rput CPO
  R"plot(y,dnorm(y,cbind(1,X)%*%mod$coef,sig))"
  R"points(y,CPO,col='blue',cex=2)"

  # If I had CPO's from 2 different models, use this visual
  R"CPO2 <- dnorm(y,cbind(1,X)%*%b,sig) # this is not really a CPO, but will be used for demo";
  R"plot(y,CPO-CPO2,bty='n',fg='grey',
         main='CPO_M1 - CPO_M2',ylab='',xlab='Observed y')"
  R"abline(h=0,col='red')"

  # Note that since most of the points lie below 0,
  # M1 may be poorer than M2.
  =#
end

println()

@testset "glm" begin
  srand(1);
  N = 100
  b = [2,3,4,5.]
  X = [ones(N) randn(N,length(b)-1)]
  σ² = 2
  y = X*b + randn(N)*sqrt(σ²)

  function loglike(y::Vector{Float64}, mu::Vector{Float64}, θ::Hyper)
    -N/2 * log(2*pi*θ[:σ²]) - sum((y-mu).^2) / (2*θ[:σ²])
  end

  lp_θ(θ::Hyper) = (-2-1) * log(θ[:σ²]) - 1 / θ[:σ²]

  @time model = glm(y, X, eye(Float64,4)*.03, loglike,
                    Σ_θ=eye(Float64,1)*.1,
                    θ_bounds=[0. Inf],
                    θ_names=[:σ²],
                    θ_logprior=lp_θ);

  print(summary(model))
end

println()

@testset "glm logistic" begin
  using Distributions
  srand(1);
  N = 1000
  b = [.2, .4, .6, .8]
  J = length(b)
  X = [ones(N) randn(N,length(b)-1)]

  invlogit(xb::Float64) = 1 / (1 + exp(-xb))
  y = map(xb -> rand(Bernoulli(invlogit(xb)))*1.0, X*b)

  function loglike(y::Vector{Float64}, Xb::Vector{Float64}, θ::Hyper)
    const mu = invlogit.(Xb)
    sum([ logpdf(Bernoulli(mu[i]),y[i]) for i in 1:N ])
  end

  @time model = glm(y, X, eye(J)*1E-2, loglike, B=2000,burn=10000);
  println(summary(model))
  #=
  using RCall
  @rput y X
  println(R"glm(y~X-1,family='binomial')")
  =#
end


using CmdStan

ProjDir = @__DIR__
cd(ProjDir)

shared_file = "shared_funcs.stan"

bernoulli_model = "
  functions{

    #include $(shared_file)
    //#include shared_funcs.stan // a comment
    //#include shared_funcs.stan // a comment
    //#include /Users/rob/shared_funcs.stan // a comment

    void model_specific_function(){
        real x = 1.0;
        return;
    }

  }
  data {
    int<lower=1> N;
    int<lower=0,upper=1> y[N];
  }
  parameters {
    real<lower=0,upper=1> theta;
  }
  model {
    model_specific_function();
    theta ~ beta(my_function(),1);
    y ~ bernoulli(theta);
  }
";

observeddata = Dict("N" => 10, "y" => [0, 1, 0, 1, 0, 0, 0, 0, 0, 1])

tmpdir = ProjDir*"/tmp"

stanmodel = Stanmodel(name="bernoulli", model=bernoulli_model,
  printsummary=true, tmpdir=tmpdir);

rc, chn, cnames = stan(stanmodel, observeddata, ProjDir);

if rc == 0
  # Describe the results
  show(chn)
  println()

  # Ceate a ChainDataFrame
  summary_df = read_summary(stanmodel)
  summary_df[:theta, [:mean, :ess]]
end


# Mostly copy/pasted from StreamStats.jl

# v0.4 updates from
# https://github.com/PalladiumConsulting/StreamStats.jl/blob/master/src/hyper_log_log.jl

hash32(d::Any) = hash(d) % UInt32
maskadd32(x::UInt32, mask::UInt32, add::UInt32) = (x & mask) + add
ρ(s::UInt32) = UInt32(leading_zeros(s)) + 0x00000001
const toInt = Int
const toUInt = UInt

function α(m::UInt32)
    if m == 0x00000010
        return 0.673
    elseif m == 0x00000020
        return 0.697
    elseif m == 0x00000040
        return 0.709
    else # if m >= UInt32(128)
        return 0.7213 / (1 + 1.079 / m)
    end
end

"""
`HyperLogLog(b)`

Approximate count of distinct elements.  `HyperLogLog` differs from other OnlineStats
in that any input to `fit!(o::HyperLogLog, input)` is considered a singleton.  Thus,
a vector of inputs must be done by:

```julia
o = HyperLogLog(4)
for yi in y
    fit!(o, yi)
end
```
"""
type HyperLogLog <: OnlineStat
    m::UInt32
    M::Vector{UInt32}
    mask::UInt32
    altmask::UInt32
    n::Int
end

function HyperLogLog(b::Integer)
    if !(4 <= b <= 16)
        throw(ArgumentError("b must be an Integer between 4 and 16"))
    end

    m = 0x00000001 << b

    M = zeros(UInt32, m)

    mask = 0x00000000
    for i in 1:(b - 1)
        mask |= 0x00000001
        mask <<= 1
    end
    mask |= 0x00000001

    altmask = ~mask

    return HyperLogLog(m, M, mask, altmask, 0)
end

function Base.show(io::IO, counter::HyperLogLog)
    @printf(io, "A HyperLogLog counter w/ %d registers", Int(counter.m))
    return
end

function fit!(counter::HyperLogLog, v::Any)
    x = hash32(v)
    j = maskadd32(x, counter.mask, 0x00000001)
    w = x & counter.altmask
    counter.M[j] = max(counter.M[j], ρ(w))
    counter.n += 1
    return
end

nobs(o::HyperLogLog) = o.n

function value(counter::HyperLogLog)
    S = 0.0

    for j in 1:counter.m
        S += 1 / (2^counter.M[j])
    end

    Z = 1 / S

    E = α(counter.m) * toUInt(counter.m)^2 * Z

    if E <= 5//2 * counter.m
        V = 0
        for j in 1:counter.m
            V += toInt(counter.M[j] == 0x00000000)
        end
        if V != 0
            E_star = counter.m * log(counter.m / V)
        else
            E_star = E
        end
    elseif E <= 1//30 * 2^32
        E_star = E
    else
        E_star = -2^32 * log(1 - E / (2^32))
    end

    return E_star
end

# TODO: Figure out details here
# function confint(counter::HyperLogLog)
#     e = estimate(counter)
#     delta = e * 1.04 / sqrt(counter.m)
#     return e - delta, e + delta
# end


#=
@author : Spencer Lyon <spencer.lyon@nyu.edu>

References
----------
Based off the original python file binom_df.py
=#
using PyPlot
using Distributions

srand(42)  # reproducible results
fig, axes = plt.subplots(2, 2)
plt.subplots_adjust(hspace=0.4)
axes = [axes...]
ns = [1, 2, 4, 8]
dom = 1:9

for (ax, n) in zip(axes, ns)
    b = Binomial(n, 0.5)
    ax[:bar](dom, pdf(b, dom), alpha=0.6, align="center")
    ax[:set_xlim](-0.5, 8.5)
    ax[:set_ylim](0, 0.55)
    ax[:set_xticks](dom)
    ax[:set_yticks]((0, 0.2, 0.4))
    ax[:set_title](LaTeXString("\$n = $n\$"))
end


export CP

mutable struct CP <: PopulationModel
    k::Int
    cp::Union{CANDECOMP{Float64,3},Float64}
	function CP(k, cp=NaN)
		new(k,cp)
	end
end

function fit!(model::CP, Z::Array{Float64,3})
    model.cp = candecomp(Z, model.k, ([randn(x,model.k) for x in size(Z)]...,),
    compute_error=true, method=:ALS, verbose=false, maxiter=1000)
end

weights(model::CP) = model.cp.factors[1]

factors(model::CP, dim=0) = dim == 0 ? model.cp.factors : model.cp.factors[dim]


@testset "MembershipTest" begin
    @polyvar x y
    f = x^2 + y^2 + 1

    F = AmoebaFiber2D(f)

    result = membershiptest(F, (0.0, 0.0), PolynomialAmoebas.TorusStartValueGenerator{2}())
    @test result.successfull

    r2 = membershiptest(F, (0.1, 0.1), PolynomialAmoebas.TorusStartValueGenerator{2}(), (result.solution,), MembershipTestOptions())
    @test r2.successfull
    @test r2.tries == 1
end


spacing(::Periodic, n) = 1 / n
spacing(::NonPeriodic, n) = 1 / (n-1)

coordinate(top, n) = UnitRange(0, n-1) .* spacing(top, n)


using PihVfi890
using Test

function euler_test()
    @testset "Euler" begin
        m = init_test_model();
        t = 2;
        k = 5.0;
        kPrimeTomorrow(kPrime) = 0.5 .* kPrime;

        kPrime = solve_one_point(m, t, k, kPrimeTomorrow);
        dev = euler_dev_one_point(m, t, k, kPrime, kPrimeTomorrow);
        @test kPrime > kprime_min(m, t);
        @test kPrime < kprime_max(m, t);
        @test abs(dev) < 1e-4
        dev2 = euler_dev_one_point2(m, t, k, kPrime, kPrimeTomorrow);
        @test abs(dev2) < 1e-4
    end
end

function solve_test()
    @testset "Solve" begin
        m = init_test_model();
        kPrime_tV = solve(m);
    end
end

@testset "VFI" begin
    euler_test();
    solve_test();
end

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


function cosine_similarity(x, y)
    return dot(x, y) / (norm(x) * norm(y))
end

function train_and_find_support_vectors(model, data, solver)

    SVDD.set_data!(model, data)
    try
        SVDD.fit!(model, solver)
    catch e
        # For small data sets fall back to returning all points as SVs
        if size(data, 2) <= 5
            return collect(1:size(data, 2))
        end
        throw(SamplingException("Failed fitting model on data set of size $(size(data, 2)) due to $e."))
    end
    return SVDD.get_support_vectors(model)
end

function mcc(classification, ground_truth; pos_class = :outlier, neg_class = :inlier)::Float64
    @assert length(classification) == length(ground_truth)
    tp = sum((classification .== pos_class) .& (ground_truth .== pos_class))
    fp = sum((classification .== pos_class) .& (ground_truth .== neg_class))
    tn = sum((classification .== neg_class) .& (ground_truth .== neg_class))
    fn = sum((classification .== neg_class) .& (ground_truth .== pos_class))

    d = sqrt((tp + fp) * (tp + fn) * (tn + fp) * (tn + fn))
    return d == 0 ? 0.0 : (tp * tn - fp * fn) / d
end


function norm_matrix_induced(A::Array{Float64,2})
    N = size(A)[1]
    # FIXME no check for matrix form

    # unit-norm vector
    d = 1/sqrt(N)
    v1 = ones(N)*d
    #
    v = A*v1
    return norm(v)
end


# https://www.rp-photonics.com/brewsters_angle.html

# Load modules
using Plots, LaTeXStrings
pyplot()
using ThinFilmsTools

function main()
    # Define beam
    λ = [1064.]
    θ = LinRange(0.01,90,901)
    beam = PlaneWave(λ,θ)
    # Define layers: notice that lambda is outside the range
    layers = [
        LayerTMMO(RIdb.dummy(beam.λ,1.0,0.0)),
        LayerTMMO(RIdb.dummy(beam.λ,1.5,0.0), d=10.),
        LayerTMMO(RIdb.dummy(beam.λ,1.5,0.0)),
    ]
    # call main script
    sol = tmm_optics(beam, layers)
    return sol
end

sol = main()

# plot the R, T and A spectra
plot(SpectrumAngle1D(),
    sol.beam.θ,
    [sol.Spectra.Rp[1,:], sol.Spectra.Rs[1,:], sol.Spectra.Rs[1,:]./sol.Spectra.Rp[1,:]./1000.],
    label=["p-wave" "s-wave" "(Rs/Rp)/1000"],
    line=([:solid :dash :dashdot]),
    xlims=(sol.beam.θ[1], sol.beam.θ[end]),
    yaxis=("Reflectance", (0.,0.2)),
)
gui()

# plot the refractive index profile
plot(RIprofile(), sol)
gui()


# ---
# title: 1546. Maximum Number of Non-Overlapping Subarrays With Sum Equals Target
# id: problem1546
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/maximum-number-of-non-overlapping-subarrays-with-sum-equals-target/description/>
# hidden: true
# ---
#
# Given an array `nums` and an integer `target`.
#
# Return the maximum number of **non-empty**   **non-overlapping** subarrays
# such that the sum of values in each subarray is equal to `target`.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [1,1,1,1,1], target = 2
#     Output: 2
#     Explanation: There are 2 non-overlapping subarrays [ **1,1** ,1, **1,1** ] with sum equals to target(2).
#
#
# **Example 2:**
#
#
#
#     Input: nums = [-1,3,5,1,4,2,-9], target = 6
#     Output: 2
#     Explanation: There are 3 subarrays with sum equal to 6.
#     ([5,1], [4,2], [3,5,1,4,2,-9]) but only the first 2 are non-overlapping.
#
# **Example 3:**
#
#
#
#     Input: nums = [-2,6,6,3,5,4,1,2,8], target = 10
#     Output: 3
#
#
# **Example 4:**
#
#
#
#     Input: nums = [0,0,0], target = 0
#     Output: 3
#
#
#
#
# **Constraints:**
#
#   * `1 <= nums.length <= 10^5`
#   * `-10^4 <= nums[i] <= 10^4`
#   * `0 <= target <= 10^6`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


module MorePolynomials

using Polynomials
using LinearAlgebra
using Intervals
using Reduce # for genvandpoly macro
using FastGaussQuadrature
import Polynomials.fit
import Polynomials.domain
import Polynomials.coeffs
import Polynomials.showterm
import Polynomials.derivative
import Polynomials.integrate
import Base.convert
import Base.length
import Base.push!
import Base.length

const SymbolLike = Union{AbstractString,Char,Symbol}

include("legendre.jl")
include("IntervalsExtras.jl")
include("register.jl")
include("Lagrange.jl")
include("Global.jl")
include("vanderpoly.jl")
include("polyhelper.jl")

end # module


using RobustOptimization
using JuMP
using Test
using MathOptInterface
const MOI = MathOptInterface
@testset "RobustOptimization.jl" begin
    # Write your tests here.
    @testset "RobustOptimization.jl" begin
    # Write your tests here.

    m = RobustModel()

    @variable(m, 0 <= y <= 1,Uncertain(1))

    @variable(m, 1 <= x_interval <= 3.5,Decision(1,"test",m.uncertainVariables))
    #Add VariableTest
    @testset "variables" begin
        #init
        mod = RobustModel()
        #add vars
        @variable(mod, 0 <= y <= 1,Uncertain(1))
        @variable(mod, 0 <= z <= 1,Uncertain(1))
        @variable(mod, 1 <= x_interval <= 3.5,Decision(1,"test",mod.uncertainVariables))

        #test add vars
        @test mod.var_to_name[3] == "x_interval"
        @test mod.var_to_name[1] == "y"

        #test infos
        @test has_lower_bound(x_interval) == true
        @test lower_bound(x_interval) == 1.0
        @test has_upper_bound(x_interval) == true
        @test upper_bound(x_interval)== 3.5

        #test setter
        set_lower_bound(x_interval,0.5)
        @test lower_bound(x_interval) == 0.5
        set_upper_bound(x_interval,5.5)
        @test upper_bound(x_interval) == 5.5

        #test start value
        set_start_value(x_interval,2.0)
        @test start_value(x_interval) == 2.0

        #test delete
        delete_lower_bound(x_interval)
        @test has_lower_bound(x_interval) == false
        delete_upper_bound(x_interval)
        @test has_upper_bound(x_interval) == false

        #test fixing
        fix(x_interval, 4)
        @test is_fixed(x_interval) == true

        #test unfix
        unfix(x_interval)
        @test is_fixed(x_interval) == false

        #test binary
        set_binary(x_interval)
        @test is_binary(x_interval) == true
        unset_binary(x_interval)
        @test is_binary(x_interval) == false

        #test integer
        set_integer(x_interval)
        @test is_integer(x_interval) == true
        unset_integer(x_interval)
        @test is_integer(x_interval) == false

        #test delete
        decLen = length(mod.decisionVariables)
        delete(mod,x_interval)
        @test length(mod.decisionVariables)-decLen == -1

    end

    @testset "constraints" begin
        model = RobustModel()
        @variable(model, 0 <= y <= 1,Uncertain(1))
        @constraint(model, uncConst, 1.5y >= 1, UncertaintySetConstraint("1"))
        @constraint(model, modConst, 2y <= 2, UncertainConstraint("2"))
        @constraint(model, modConstt, 2y <= 2, UncertainConstraint("2"))
        @test name(uncConst) == "uncConst"
        print(model,"\n")

        @test string(model.uncertainConstraints[uncConst.type_idx].func)== "2 y - 2"
        @test typeof(model.uncertainConstraints[uncConst.type_idx].set) == MathOptInterface.LessThan{Float64}
        #JuMP.coefficient() #model.uncertainConstraints[uncConst.type_idx], model.uncertainVariables[y.type_idx], 4)
        print(typeof(model.uncertainConstraints[uncConst.type_idx].func),"\n")

    end

    #@constraint(m, c1, 2*x_interval + 2*y_interval <= 10)
    #@objective(m, Max, x_interval + y_interval)
    @test m.var_to_name[2] == "x_interval"
    #TODO get inner of Decision
    @test m.var_to_name[1] == "y"
    @test m.nextvaridx == length(m.uncertainVariables)+length(m.decisionVariables)
    #con =m.uncertainConstraints[RobustOptimization.Int(1)]
    #@test con.set == MathOptInterface.LessThan{Float64}(10.0)

    ########### TEST Variable Basics

    print(m.objectivesense)
end

end


@testset "disrelation.jl" begin
    x, A = demodata()
    Ã = A .- mean(A, dims=1)

    Λ = disrelation(Ã)
    @plottest begin
        heatmap(x, x, Λ)
        plot!(xflip=true, yflip=true)
        xlabel!("Peak position")
        ylabel!("Peak position")
        title!("Disrelation")
    end "disrelation.png"

    # Can this be prettier?
    @test .≈(diag(Λ), 0) |> all

    Ψ = asynchronous(Ã)
    @test maximum(sqrt.((Λ .- Ψ).^2)) ≈ 0 atol=2e-3
end


using ComputationalHomology
using LinearAlgebra
using SparseArrays
using Distances
using Test
using Random

include("cells.jl")
include("chains.jl")
include("simplicial.jl")
include("bitmap.jl")
include("cw.jl")
include("boundary.jl")
include("filtration.jl")
include("constructions.jl")
include("homology.jl")
include("persistence.jl")
include("landscape.jl")
include("pimage.jl")
include("distances.jl")
include("tropic.jl")


using LowRankModels, Random
import StatsBase: sample, Weights, mean
import LinearAlgebra: norm

# test MNL Ordinal loss

## generate data
Random.seed!(1);
m,n,k = 100,100,2;
kfit = k+1
nlevels = 5; # number of levels
d = nlevels-1 # embedding dimension
D = n*d;
# coordinates of covariates
X_real = randn(m,k)
# directions of observations
Y_real = randn(k,n)
# measurement thresholds
T_real = k*randn(d,n) # notice x^T y has variance k; so making this bigger makes the problem easier
for j=1:n # this scheme doesn't work to ensure uniform sampling
	T_real[:,j] = sort(T_real[:,j])
end

signedsums = Array{Float64}(undef, d, nlevels)
for i=1:d
    for j=1:nlevels
        signedsums[i,j] = i<j ? 1 : -1
    end
end

XY = X_real*Y_real;
XYplusT = zeros(Float64, (m,D))
A = zeros(Int, (m, n))
for i=1:m
	for j=1:n
		u = XY[i,j] .+ T_real[:,j]
		XYplusT[i,(j-1)*d .+ (1:d)] = u
		diffs = u'*signedsums
		wv = Weights(Float64[exp(-diffs[l]) for l in 1:nlevels])
		l = sample(wv)
		A[i,j] = l
	end
end
# loss is insensitive to shifts; regularizer should pick this shift
XYplusT = XYplusT .- mean(XYplusT, dims=2)

# and the model
losses = fill(MultinomialOrdinalLoss(nlevels),n)
rx, ry = lastentry1(QuadReg(.01)), MNLOrdinalReg(QuadReg(.01)) #lastentry_unpenalized(QuadReg(10));
Yord = randn(kfit,D)
yidxs = get_yidxs(losses)
for j=1:n
	prox!(ry, view(Yord,:,yidxs[j]), 1)
end
glrm = GLRM(A,losses,rx,ry,kfit, scale=false, offset=false, X=randn(kfit,m), Y=Yord);

# fit w/o initialization
@time X,Y,ch = fit!(glrm);
XYh = X'*Y#[:,1:d:D];
@show ch.objective
println("After fitting, parameters differ from true parameters by $(norm(XYplusT - XYh)/sqrt(prod(size(XYplusT)))) in RMSE")
A_imputed = impute(glrm)
println("After fitting, $(sum(A_imputed .!= A) / prod(size(A))*100)% of imputed entries are wrong")
println("After fitting, imputed entries are off by $(sum(abs.(A_imputed - A)) / prod(size(A))*100)% on average")
println("(Picking randomly, $((nlevels-1)/nlevels*100)% of entries would be wrong.)\n")

# initialize
init_svd!(glrm)
XYh = glrm.X' * glrm.Y
println("After initialization with the svd, parameters differ from true parameters by $(norm(XYplusT - XYh)/sqrt(prod(size(XYplusT)))) in RMSE")
A_imputed = impute(glrm)
println("After initialization with the svd, $(sum(A_imputed .!= A) / prod(size(A))*100)% of imputed entries are wrong")
println("After initialization with the svd, imputed entries are off by $(sum(abs,(A_imputed - A)) / prod(size(A))*100)% on average")
println("(Picking randomly, $((nlevels-1)/nlevels*100)% of entries would be wrong.)\n")

# fit w/ initialization
@time X,Y,ch = fit!(glrm);
XYh = X'*Y#[:,1:d:D];
@show ch.objective
println("After fitting, parameters differ from true parameters by $(norm(XYplusT - XYh)/sqrt(prod(size(XYplusT)))) in RMSE")
A_imputed = impute(glrm)
println("After fitting, $(sum(A_imputed .!= A) / prod(size(A))*100)% of imputed entries are wrong")
println("After fitting, imputed entries are off by $(sum(abs,(A_imputed - A)) / prod(size(A))*100)% on average")
println("(Picking randomly, $((nlevels-1)/nlevels*100)% of entries would be wrong.)\n")

# test scaling
mul!(glrm.ry, 3)


function tgear(thtl)  # Power command v. thtl. relationship
    if thtl <= 0.77
        cpow = 64.94 * thtl
    else
        cpow = 217.38 * thtl - 117.38
    end
    return cpow
end

function pdot(p3, p1) # PDOT= rate of change of power
    if p1 >= 50.0  # P3= actual power, P1= power command
        if p3 >= 50.0
            t = 5.0
            p2 = p1
        else
            p2 = 60.0
            t = rtau(p2 - p3)
        end
    else
        if p3 >= 50.0
            t = 5.0
            p2 = 40.0
        else
            p2 = p1
            t = rtau(p2 - p3)
        end
    end
    pd = t * (p2 - p3)
    return pd
end

function rtau(dp)
    if dp <= 25.0
        rt = 1.0  # Reciprocal time constant
    elseif dp >= 50.0
        rt = 0.1
    else
        rt = 1.9 - 0.036 * dp
    end
    return rt
end

idle_data =
    [
        1060.0 670.0 880.0 1140.0 1500.0 1860.0
        635.0 425.0 690.0 1010.0 1330.0 1700.0
        60.0 25.0 345.0 755.0 1130.0 1525.0
        -1020.0 -710.0 -300.0 350.0 910.0 1360.0
        -2700.0 -1900.0 -1300.0 -247.0 600.0 1100.0
        -3600.0 -1400.0 -595.0 -342.0 -200.0 700.0
    ]'

mil_data =
    [
        12680.0 9150.0 6200.0 3950.0 2450.0 1400.0
        12680.0 9150.0 6313.0 4040.0 2470.0 1400.0
        12610.0 9312.0 6610.0 4290.0 2600.0 1560.0
        12640.0 9839.0 7090.0 4660.0 2840.0 1660.0
        12390.0 10176.0 7750.0 5320.0 3250.0 1930.0
        11680.0 9848.0 8050.0 6100.0 3800.0 2310.0
    ]'

max_data =
    [
        20000.0 15000.0 10800.0 7000.0 4000.0 2500.0
        21420.0 15700.0 11225.0 7323.0 4435.0 2600.0
        22700.0 16860.0 12250.0 8154.0 5000.0 2835.0
        24240.0 18910.0 13760.0 9285.0 5700.0 3215.0
        26070.0 21075.0 15975.0 11115.0 6860.0 3950.0
        28886.0 23319.0 18300.0 13484.0 8642.0 5057.0
    ]'

function thrust(pow, alt, rmach)  # Engine thrust model

    alt = alt * M2FT

    if alt <= 0.
        alt = 0.0
    end

    h = 0.0001 * alt
    i = floor(Int, h)
    if i >= 5
        i = 4
    end

    dh = h - float(i)

    rm = 5.0 * rmach
    m = floor(Int, rm)
    if m >= 5
        m = 4
    end

    dm = rm - float(m)
    cdh = 1.0 - float(dh)

    i = i + 1
    m = m + 1

    s = mil_data[i, m] * cdh + mil_data[i+1, m] * dh
    t = mil_data[i, m+1] * cdh + mil_data[i+1, m+1] * dh
    tmil = s + (t - s) * dm

    if pow < 50.0
        s = idle_data[i, m] * cdh + idle_data[i+1, m] * dh
        t = idle_data[i, m+1] * cdh + idle_data[i+1, m+1] * dh
        tidl = s + (t - s) * dm
        thrst = tidl + (tmil - tidl) * pow * 0.02
    else
        s = max_data[i, m] * cdh + max_data[i+1, m] * dh
        t = max_data[i, m+1] * cdh + max_data[i+1, m+1] * dh
        tmax = s + (t - s) * dm
        thrst = tmil + (tmax - tmil) * (pow - 50.0) * 0.02
    end

    # Transform from lbf to N (using GD defined in Stevens)
    thrst = thrst * LB2KG * GD * FT2M
    return thrst
end

function calculate_prop_forces_moments(ac::F16, dss::DSState, mach, controls)

    # Assign state & control variables
    alt = get_height(dss)
    pow = get_engine_power(dss)

    thtl = controls[1]

    thrust_ = thrust(pow, alt, mach)

    return [thrust_, 0.0, 0.0, 0.0, 0.0, 0.0]
end

calculate_prop_gyro_effects(ac::F16) = [HX, 0, 0]

calculate_pdot(ac::F16, thtl, pow) = pdot(pow, tgear(thtl))

tgear(ac::F16, thtl) = tgear(thtl)


using LinearAlgebra

struct ValueOne; end
ValueOne()

# Compute X <- a X + b I.
function matfun_axpby!(X,a,b,Y::UniformScaling)
    m,n=size(X)
    if ~(a isa ValueOne)
        rmul!(X,a)
    end
    @inbounds for i=1:n
        X[i,i]+=(b isa ValueOne) ? 1 : b
    end
end

# Compute X <- a X + b Y.
function matfun_axpby!(X,a,b,Y)
    m,n=size(X)
    if ~(a isa ValueOne)
        rmul!(X,a)
    end
    @inbounds for i=1:m
        @inbounds for j=1:n
            if (b isa ValueOne)
                X[i,j]+=Y[i,j]
            else
                X[i,j]+=b*Y[i,j]
            end
        end
    end
end

@inline function exp_sastre_m6(A)
    T=promote_type(eltype(A),Float64)
    A_copy=similar(A,T); A_copy .= A;
    return exp_sastre_m6!(A_copy)
end

@inline function exp_sastre_m6!(A)
    T=promote_type(eltype(A),Float64) # Make it work for many 'bigger' types (matrices and scalars)
    # max_memslots=7
    n=size(A,1)
    # The first slots are precomputed nodes [:A]
    memslots2 = similar(A,T)
    memslots3 = similar(A,T)
    memslots4 = similar(A,T)
    memslots5 = similar(A,T)
    memslots6 = similar(A,T)
    memslots7 = similar(A,T)
    # Assign precomputed nodes memslots
    memslots1=A # overwrite A
    # Uniform scaling is exploited.
    # No matrix I explicitly allocated.
    value_one=ValueOne()
    # Computation order: B2 Bb3 B3 Ba4 Bb4 B4 Ba7_5 Bb5_3 B5 Ba6 Bb6 B6 Bb7 B7 T2k9
    # Computing B2 with operation: mult
    mul!(memslots2,memslots1,memslots1)
    # Computing Bb3 = x*A+x*B2
    coeff1=0.001052151783051235
    coeff2=0.0004675683454147702
    memslots3 .= coeff1.*memslots1 .+ coeff2.*memslots2
    # Computing B3 with operation: mult
    mul!(memslots4,memslots2,memslots3)
    # Deallocating Bb3 in slot 3
    # Computing Ba4 = x*A+x*B2+x*B3
    coeff1=0.2868706220817633
    coeff2=-0.03289442879547955
    coeff3=1.0
    memslots3 .= coeff1.*memslots1 .+ coeff2.*memslots2 .+ coeff3.*memslots4
    # Computing Bb4 = x*B2+x*B3
    coeff1=0.05317514832355802
    coeff2=1.0
    memslots5 .= coeff1.*memslots2 .+ coeff2.*memslots4
    # Computing B4 with operation: mult
    mul!(memslots6,memslots3,memslots5)
    # Deallocating Ba4 in slot 3
    # Deallocating Bb4 in slot 5
    # Computing Ba7_5 = x*I+x*A+x*B2+x*B3+x*B4
    coeff1=3.09646797193604
    coeff2=0.772360321294401
    coeff3=0.1673139636901279
    coeff4=7.922322450524197
    coeff5=1.0
    # Smart lincomb recycle B3
    memslots4 .= coeff2.*memslots1 .+ coeff3.*memslots2 .+ coeff4.*memslots4 .+ coeff5.*memslots6
    mul!(memslots4,true,I*coeff1,true,true)
    # Deallocating B4 in slot 6
    # Computing Bb5_3 = x*A+x*B2
    coeff1=0.002688394980266927
    coeff2=0.0004675683454147702
    memslots3 .= coeff1.*memslots1 .+ coeff2.*memslots2
    # Computing B5 with operation: mult
    mul!(memslots5,memslots2,memslots3)
    # Deallocating Bb5_3 in slot 3
    # Computing Ba6 = x*A+x*B2+x*B5
    coeff1=0.39689859154115
    coeff2=0.02219811707032801
    coeff3=1.0
    memslots3 .= coeff1.*memslots1 .+ coeff2.*memslots2 .+ coeff3.*memslots5
    # Computing Bb6 = x*B2+x*B5
    coeff1=0.0277140002806296
    coeff2=1.0
    memslots6 .= coeff1.*memslots2 .+ coeff2.*memslots5
    # Computing B6 with operation: mult
    mul!(memslots7,memslots3,memslots6)
    # Deallocating Ba6 in slot 3
    # Deallocating Bb6 in slot 6
    # Computing Bb7 = x*A+x*B2+x*B5+x*B6
    coeff1=0.3229486011362678
    coeff2=0.08092036376147299
    coeff3=1.930814505527068
    coeff4=1.0
    # Smart lincomb recycle A
    memslots1 .= coeff1.*memslots1 .+ coeff2.*memslots2 .+ coeff3.*memslots5 .+ coeff4.*memslots7
    # Deallocating B2 in slot 2
    # Deallocating B5 in slot 5
    # Deallocating B6 in slot 7
    # Computing B7 with operation: mult
    mul!(memslots2,memslots4,memslots1)
    # Deallocating Ba7_5 in slot 4
    # Deallocating Bb7 in slot 1
    # Computing T2k9 = x*I+x*B7
    coeff1=1.0
    coeff2=1.0
    # Smart lincomb recycle B7
    memslots2 .= coeff2.*memslots2
    mul!(memslots2,true,I*coeff1,true,true)
    return memslots2 # Returning T2k9
end


using OptimalTransport

using Distances
using PythonOT: PythonOT

using Random
using Test

const POT = PythonOT

Random.seed!(100)

@testset "quadratic.jl" begin
    @testset "quadreg" begin
        M = 250
        N = 200

        # create two uniform histograms
        μ = fill(1 / M, M)
        ν = fill(1 / N, N)

        # create random cost matrix
        C = pairwise(SqEuclidean(), rand(1, M), rand(1, N); dims=2)

        # compute optimal transport map (Julia implementation + POT)
        eps = 0.25
        γ = quadreg(μ, ν, C, eps, QuadraticOTNewton())

        γ_pot = POT.Smooth.smooth_ot_dual(μ, ν, C, eps; stopThr=1e-9)

        # need to use a larger tolerance here because of a quirk with the POT solver
        @test γ ≈ γ_pot rtol = 1e-1
    end
end