# ---
# title: 1027. Longest Arithmetic Subsequence
# id: problem1027
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/longest-arithmetic-subsequence/description/>
# hidden: true
# ---
#
# Given an array `A` of integers, return the **length** of the longest
# arithmetic subsequence in `A`.
#
# Recall that a _subsequence_ of `A` is a list `A[i_1], A[i_2], ..., A[i_k]`
# with `0 <= i_1 < i_2 < ... < i_k <= A.length - 1`, and that a sequence `B` is
# _arithmetic_ if `B[i+1] - B[i]` are all the same value (for `0 <= i < B.length
# - 1`).
#
#
#
# **Example 1:**
#
#
#
#     Input: A = [3,6,9,12]
#     Output: 4
#     Explanation:
#     The whole array is an arithmetic sequence with steps of length = 3.
#
#
# **Example 2:**
#
#
#
#     Input: A = [9,4,7,2,10]
#     Output: 3
#     Explanation:
#     The longest arithmetic subsequence is [4,7,10].
#
#
# **Example 3:**
#
#
#
#     Input: A = [20,1,15,3,10,5,8]
#     Output: 4
#     Explanation:
#     The longest arithmetic subsequence is [20,15,10,5].
#
#
#
#
# **Constraints:**
#
#   * `2 <= A.length <= 1000`
#   * `0 <= A[i] <= 500`
#
#
## @lc code=start
using LeetCode

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


using LinearAlgebra

import ClimaCore:
    Domains, Fields, Geometry, Meshes, Operators, Spaces, Topologies

using OrdinaryDiffEq: ODEProblem, solve, SSPRK33

import Logging
import TerminalLoggers
Logging.global_logger(TerminalLoggers.TerminalLogger())

"""
    convergence_rate(err, Δh)

Estimate convergence rate given vectors `err` and `Δh`

    err = C Δh^p+ H.O.T
    err_k ≈ C Δh_k^p
    err_k/err_m ≈ Δh_k^p/Δh_m^p
    log(err_k/err_m) ≈ log((Δh_k/Δh_m)^p)
    log(err_k/err_m) ≈ p*log(Δh_k/Δh_m)
    log(err_k/err_m)/log(Δh_k/Δh_m) ≈ p

"""
convergence_rate(err, Δh) =
    [log(err[i] / err[i - 1]) / log(Δh[i] / Δh[i - 1]) for i in 2:length(Δh)]

# Advection problem on a sphere. The initial condition can be set via a command line
# argument. Possible test cases are: cosine_bell (default) and gaussian_bell

const R = 6.37122e6
const h0 = 1000.0
const r0 = R / 3
const u0 = 2 * pi * R / (86400 * 12)
const center = Geometry.LatLongPoint(0.0, 270.0)
const test_name = get(ARGS, 1, "cosine_bell") # default test case to run
const test_angle_name = get(ARGS, 2, "alpha0") # default test case to run
const cosine_test_name = "cosine_bell"
const gaussian_test_name = "gaussian_bell"
const alpha0_test_name = "alpha0"
const alpha45_test_name = "alpha45"

if test_angle_name == alpha45_test_name
    const α0 = 45.0
else # default test case, α0 = 0.0
    const α0 = 0.0
end

# Plot variables and auxiliary function
ENV["GKSwstype"] = "nul"
import ClimaCorePlots, Plots
Plots.GRBackend()
dirname = "cg_sphere_solidbody_$(test_name)"
dirname = "$(dirname)_$(test_angle_name)"
path = joinpath(@__DIR__, "output", dirname)
mkpath(path)

function linkfig(figpath, alt = "")
    # buildkite-agent upload figpath
    # link figure in logs if we are running on CI
    if get(ENV, "BUILDKITE", "") == "true"
        artifact_url = "artifact://$figpath"
        print("\033]1338;url='$(artifact_url)';alt='$(alt)'\a\n")
    end
end

FT = Float64
ne_seq = 2 .^ (2, 3, 4, 5)
Δh = zeros(FT, length(ne_seq))
L1err, L2err, Linferr = zeros(FT, length(ne_seq)),
zeros(FT, length(ne_seq)),
zeros(FT, length(ne_seq))
Nq = 4

# h-refinement study
for (k, ne) in enumerate(ne_seq)
    domain = Domains.SphereDomain(R)
    mesh = Meshes.EquiangularCubedSphere(domain, ne)
    grid_topology = Topologies.Topology2D(mesh)
    quad = Spaces.Quadratures.GLL{Nq}()
    space = Spaces.SpectralElementSpace2D(grid_topology, quad)

    coords = Fields.coordinate_field(space)

    Δh[k] = 2 * R / ne

    global_geom = space.global_geometry

    h_init = map(coords) do coord
        rd = Geometry.great_circle_distance(coord, center, global_geom)

        if test_name == gaussian_test_name
            h0 * exp(-(rd / r0)^2 / 2)
        else # default test case, cosine bell
            if rd < r0
                h0 / 2 * (1 + cospi(rd / r0))
            else
                0.0
            end
        end
    end

    u = map(coords) do coord
        ϕ = coord.lat
        λ = coord.long

        uu = u0 * (cosd(α0) * cosd(ϕ) + sind(α0) * cosd(λ) * sind(ϕ))
        uv = -u0 * sind(α0) * sind(λ)
        Geometry.UVVector(uu, uv)
    end

    function rhs!(dh, h, u, t)
        div = Operators.Divergence()

        @. dh = -div(h * u) # strong form of equation
        Spaces.weighted_dss!(dh)
    end

    # Set the RHS function
    rhs!(similar(h_init), h_init, u, 0.0)

    # Solve the ODE
    T = 86400 * 12
    dt = 20 * 60
    prob = ODEProblem(rhs!, h_init, (0.0, T), u)
    sol = solve(
        prob,
        SSPRK33(),
        dt = dt,
        saveat = dt,
        progress = true,
        adaptive = false,
        progress_message = (dt, u, p, t) -> t,
    )
    L1err[k] = norm(sol.u[end] .- h_init, 1)
    L2err[k] = norm(sol.u[end] .- h_init)
    Linferr[k] = norm(sol.u[end] .- h_init, Inf)

    @info "Test case: $(test_name) with α: $(test_angle_name)"
    @info "Number of elements per cube panel: $(ne) x $(ne)"
    @info "Solution norm at t = 0: ", norm(h_init)
    @info "Solution norm at t = $(T): ", norm(sol.u[end])
    @info "L₁ error at t = $(T): ", L1err[k]
    @info "L₂ error at t = $(T): ", L2err[k]
    @info "L∞ error at t = $(T): ", Linferr[k]
end

# Print convergence rate info
conv = convergence_rate(L2err, Δh)
@info "Converge rates for this test case are: ", conv

# Plot the errors
# L₁ error Vs number of elements
Plots.png(
    Plots.plot(
        collect(ne_seq),
        L1err,
        yscale = :log10,
        xlabel = "Nₑ",
        ylabel = "log₁₀(L₁ err)",
        label = "",
    ),
    joinpath(path, "L1error.png"),
)
linkfig(
    relpath(joinpath(path, "L1error.png"), joinpath(@__DIR__, "../..")),
    "L₁ error Vs Nₑ",
)

# L₂ error Vs number of elements
Plots.png(
    Plots.plot(
        collect(ne_seq),
        L2err,
        yscale = :log10,
        xlabel = "Nₑ",
        ylabel = "log₁₀(L₂ err)",
        label = "",
    ),
    joinpath(path, "L2error.png"),
)
linkfig(
    relpath(joinpath(path, "L2error.png"), joinpath(@__DIR__, "../..")),
    "L₂ error Vs Nₑ",
)

# L∞ error Vs number of elements
Plots.png(
    Plots.plot(
        collect(ne_seq),
        Linferr,
        yscale = :log10,
        xlabel = "Nₑ",
        ylabel = "log₁₀(L∞ err)",
        label = "",
    ),
    joinpath(path, "Linferror.png"),
)
linkfig(
    relpath(joinpath(path, "Linferror.png"), joinpath(@__DIR__, "../..")),
    "L∞ error Vs Nₑ",
)


# Most of this work is borrowed from
# https://github.com/stan-dev/rstan
# Copyright (C) 2012, 2013, 2014, 2015, 2016, 2017, 2018 Trustees of Columbia University
# Copyright (C) 2018, 2019 Aki Vehtari, Paul Bürkner

# References

# Andrew Gelman, John B. Carlin, Hal S. Stern, David B. Dunson, Aki
# Vehtari and Donald B. Rubin (2013). Bayesian Data Analysis, Third
# Edition. Chapman and Hall/CRC.

# Aki Vehtari, Andrew Gelman, Daniel Simpson, Bob Carpenter, and
# Paul-Christian Bürkner (2019). Rank-normalization, folding, and
# localization: An improved R-hat for assessing convergence of
# MCMC. arXiv preprint arXiv:1903.08008

using Statistics
using FFTW
using Distributions

function autocovariance(x)
    N = length(x)
    Mt2 = 2 * fft_nextgoodsize(N)
    yc = x .- mean(x)
    yc = append!(yc, repeat([0], Mt2 - N))
    t = bfft(yc)
    ac = bfft(conj(t) .* t)
    return real(ac)[1:N] ./ (N * N * 2)
end

function autocorrelation(x)
    ac = autocovariance(x)
    return ac / ac[1]
end

function fft_nextgoodsize(N)
    N <= 2 && return 2

    while true
        m = N
        while mod(m, 2) == 0
            m /= 2
        end
        while mod(m, 3) == 0
            m /= 3
        end
        while mod(m, 5) == 0
            m /= 5
        end
        if m <= 1
            return N
        end
        N += 1
    end
end

function ess(x)

    niterations, nchains = size(x)

    (niterations < 3 || any(isnan.(x))) && return NaN
    any(isinf.(x)) && return NaN
    isconstant(x) && return NaN

    acov = mapslices(autocovariance, x, dims=1)
    chain_mean = mean(x, dims=1)
    mean_var = mean(acov[1, :]) * niterations / (niterations - 1)
    var_plus = mean_var * (niterations - 1) / niterations

    if nchains > 1
        var_plus += var(chain_mean)
    end

    rhohat = zeros(niterations)
    t = 0
    rhohat_even = 1.0
    rhohat[t + 1] = rhohat_even
    rhohat_odd = 1 - (mean_var - mean(acov[t + 2, :])) / var_plus
    rhohat[t + 2] = rhohat_odd

    while t < niterations - 5 &&
        !isnan(rhohat_even + rhohat_odd) &&
        rhohat_even + rhohat_odd > 0
        t += 2
        rhohat_even = 1 - (mean_var - mean(acov[t + 1, :])) / var_plus
        rhohat_odd = 1 - (mean_var - mean(acov[t + 2, :])) / var_plus

        if rhohat_even + rhohat_odd >= 0
            rhohat[t + 1] = rhohat_even
            rhohat[t + 2] = rhohat_odd
        end
    end

    max_t = t
    # this is used in the improved estimate
    if rhohat_even > 0
        rhohat[max_t + 1] = rhohat_even
    end

    # Geyer's initial monotone sequence
    t = 0
    while t <= max_t - 4
        t += 2
        if rhohat[t + 1] + rhohat[t + 2] > rhohat[t - 1] + rhohat[t]
            rhohat[t + 1] = (rhohat[t - 1] + rhohat[t]) / 2
            rhohat[t + 2] = rhohat[t + 1]
        end
    end

    ess = nchains * niterations
    # Geyer's truncated estimate
    # it's possible max_t == 0; 1:0 does not behave like in R
    τ = -1 + 2.0 * sum(rhohat[1:max(1, max_t)]) + rhohat[max_t + 1]
    # Improved estimate reduces variance in antithetic case
    τ = max(τ, 1.0 / log10(ess))
    return ess / τ
end

function rhat_impl(x)
    any(isnan.(x)) && return NaN
    any(isinf.(x)) && return NaN
    isconstant(x) && return NaN

    niterations, nchains = size(x)

    chain_mean = mean(x, dims=1)
    chain_var = var(x, dims=1)

    var_between = niterations * var(chain_mean)
    var_within = mean(chain_var)

    return sqrt((var_between / var_within + niterations - 1) / niterations)
end

function rhat(x)
    rhatbulk = rhat_impl()
end

function isconstant(x, tol=sqrt(eps(0.0)))
    return isapprox(minimum(x), maximum(x), rtol=tol)
end

function splitchains(x)
    niterations = size(x, 1)
    niterations < 2 && return x

    half = niterations / 2
    ub_lowerhalf::Int = floor(half)
    lb_secondhalf::Int = ceil(half + 1)
    return hcat(x[1:ub_lowerhalf, :], x[lb_secondhalf:end, :])
end

function fold(x)
    return abs.(x .- median(x))
end

function tiedrank(x)
    # Borrowed with slight modification from StatsBase
    # https://github.com/JuliaStats/StatsBase.jl/blob/master/src/ranking.jl

    # Copyright (c) 2012-2016: Dahua Lin, Simon Byrne, Andreas Noack,
    # Douglas Bates, John Myles White, Simon Kornblith, and other contributors.

    # Permission is hereby granted, free of charge, to any person obtaining
    # a copy of this software and associated documentation files (the
    # "Software"), to deal in the Software without restriction, including
    # without limitation the rights to use, copy, modify, merge, publish,
    # distribute, sublicense, and/or sell copies of the Software, and to
    # permit persons to whom the Software is furnished to do so, subject to
    # the following conditions:
    #
    # The above copyright notice and this permission notice shall be
    # included in all copies or substantial portions of the Software.
    #
    # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
    # EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
    # MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
    # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
    # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
    # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
    # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

    if length(size(x)) > 1
        x = x[:]
    end

    n = length(x)
    p = sortperm(x)
    rks = zeros(n)

    if n > 0
        v = x[p[1]]

        s = 1  # starting index of current range
        e = 2  # pass-by-end index of current range
        while e <= n
            cx = x[p[e]]
            if cx != v
                # fill average rank to s : e-1
                ar = (s + e - 1) / 2
                for i = s : e-1
                    rks[p[i]] = ar
                end
                # switch to next range
                s = e
                v = cx
            end
            e += 1
        end

        # the last range (e == n+1)
        ar = (s + n) / 2
        for i = s : n
            rks[p[i]] = ar
        end
    end

    return rks
end

function zscale(x)
    r = tiedrank(x[:])
    N = Normal()
    z = quantile.(N, (r .- 0.5) ./ length(x[:]))
    if length(size(x)) > 1
        z = reshape(z, size(x))
    end
    return z
end

function rhat_basic(x, split = true)
    split && (x = splitchains(x))
    return rhat_impl(x)
end

function ess_basic(x, split = true)
    split && (x = splitchains(x))
    return ess(x)
end

function rhat(x)
    rhat_bulk = rhat_impl(zscale(splitchains(x)))
    rhat_tail = rhat_impl(zscale(splitchains(fold(x))))
    return max(rhat_bulk, rhat_tail)
end

function ess_bulk(x)
    return ess(zscale(splitchains(x)))
end

function ess_tail(x)
    I05 = x .<= quantile(x[:], 0.05)
    q05_ess = ess(splitchains(I05))
    I95 = x .<= quantile(x[:], 0.95)
    q95_ess = ess(splitchains(I95))
    return min(q05_ess, q95_ess)
end

function ess_quantile(x, prob::Real)
    @assert prob >= 0 && prob <= 1
    I = x .<= quantile(x[:], prob)
    return ess(splitchains(I))
end

function ess_mean(x)
    return ess(splitchains(x))
end

function ess_sq(x)
    return ess(splitchains(x .^ 2))
end

function ess_std(x)
    return min(ess_mean(x), ess_sq(x))
end

function mcse_quantile(x, prob::Real)
    ess = ess_quantile(x, prob)
    p = [0.1586553; 0.8413447]
    B = Beta(ess * prob + 1, ess * (1 - prob) + 1)
    a = quantile.(B, p)
    ssims = sort(x[:])
    S = length(ssims)
    th1 = ssims[convert(Int64, max(floor(a[1] * S), 1))]
    th2 = ssims[convert(Int64, min(ceil(a[2] * S), S))]
    return (th2 - th1) / 2
end

function mcse_mean(x)
    return std(x[:]) / sqrt(ess(x))
end

function mcse_std(x)
    ess_sd = ess_std(x)
    return std(x[:]) * sqrt(exp(1) * (1 - 1 / ess_sd) ^ (ess_sd - 1) - 1)
end


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:4] <= 1)

                     @variable(m, -1.3634297333714107 <= q <= -0.20981772159177886)

                     add_NL_constraint(m, :(1/(1 + exp(-(-0.25086622252953283 + -0.4158912166699471*1/(1 + exp(-(0.18862560740304835 + -0.15744636994779526*1/(1 + exp(-(0.5575683410422836 + 0.3104874388021095*1/(1 + exp(-(0.8764207742833086 + -0.10444718155116828*$(x[1]) + 0.8852151451027264*$(x[2]) + -0.8060184127084176*$(x[3]) + -0.617200197422199*$(x[4])))) + 0.12330243724213741*1/(1 + exp(-(0.934307327106453 + 0.477402449923221*$(x[1]) + -0.9131747528764835*$(x[2]) + 0.07837215353953164*$(x[3]) + 0.2564992344946111*$(x[4]))))))) + 0.6894873557418655*1/(1 + exp(-(0.5818927234969289 + 0.4060546195732435*1/(1 + exp(-(0.8764207742833086 + -0.10444718155116828*$(x[1]) + 0.8852151451027264*$(x[2]) + -0.8060184127084176*$(x[3]) + -0.617200197422199*$(x[4])))) + 0.22410021012329473*1/(1 + exp(-(0.934307327106453 + 0.477402449923221*$(x[1]) + -0.9131747528764835*$(x[2]) + 0.07837215353953164*$(x[3]) + 0.2564992344946111*$(x[4])))))))))) + 0.15556574286983027*1/(1 + exp(-(-0.6464546506721769 + -0.3052437329487665*1/(1 + exp(-(0.5575683410422836 + 0.3104874388021095*1/(1 + exp(-(0.8764207742833086 + -0.10444718155116828*$(x[1]) + 0.8852151451027264*$(x[2]) + -0.8060184127084176*$(x[3]) + -0.617200197422199*$(x[4])))) + 0.12330243724213741*1/(1 + exp(-(0.934307327106453 + 0.477402449923221*$(x[1]) + -0.9131747528764835*$(x[2]) + 0.07837215353953164*$(x[3]) + 0.2564992344946111*$(x[4]))))))) + 0.4287563119683666*1/(1 + exp(-(0.5818927234969289 + 0.4060546195732435*1/(1 + exp(-(0.8764207742833086 + -0.10444718155116828*$(x[1]) + 0.8852151451027264*$(x[2]) + -0.8060184127084176*$(x[3]) + -0.617200197422199*$(x[4])))) + 0.22410021012329473*1/(1 + exp(-(0.934307327106453 + 0.477402449923221*$(x[1]) + -0.9131747528764835*$(x[2]) + 0.07837215353953164*$(x[3]) + 0.2564992344946111*$(x[4]))))))))))))) + 1/(1 + exp(-(-0.9805026843034819 + -0.19393497198753096*1/(1 + exp(-(0.18862560740304835 + -0.15744636994779526*1/(1 + exp(-(0.5575683410422836 + 0.3104874388021095*1/(1 + exp(-(0.8764207742833086 + -0.10444718155116828*$(x[1]) + 0.8852151451027264*$(x[2]) + -0.8060184127084176*$(x[3]) + -0.617200197422199*$(x[4])))) + 0.12330243724213741*1/(1 + exp(-(0.934307327106453 + 0.477402449923221*$(x[1]) + -0.9131747528764835*$(x[2]) + 0.07837215353953164*$(x[3]) + 0.2564992344946111*$(x[4]))))))) + 0.6894873557418655*1/(1 + exp(-(0.5818927234969289 + 0.4060546195732435*1/(1 + exp(-(0.8764207742833086 + -0.10444718155116828*$(x[1]) + 0.8852151451027264*$(x[2]) + -0.8060184127084176*$(x[3]) + -0.617200197422199*$(x[4])))) + 0.22410021012329473*1/(1 + exp(-(0.934307327106453 + 0.477402449923221*$(x[1]) + -0.9131747528764835*$(x[2]) + 0.07837215353953164*$(x[3]) + 0.2564992344946111*$(x[4])))))))))) + -0.01642706539433636*1/(1 + exp(-(-0.6464546506721769 + -0.3052437329487665*1/(1 + exp(-(0.5575683410422836 + 0.3104874388021095*1/(1 + exp(-(0.8764207742833086 + -0.10444718155116828*$(x[1]) + 0.8852151451027264*$(x[2]) + -0.8060184127084176*$(x[3]) + -0.617200197422199*$(x[4])))) + 0.12330243724213741*1/(1 + exp(-(0.934307327106453 + 0.477402449923221*$(x[1]) + -0.9131747528764835*$(x[2]) + 0.07837215353953164*$(x[3]) + 0.2564992344946111*$(x[4]))))))) + 0.4287563119683666*1/(1 + exp(-(0.5818927234969289 + 0.4060546195732435*1/(1 + exp(-(0.8764207742833086 + -0.10444718155116828*$(x[1]) + 0.8852151451027264*$(x[2]) + -0.8060184127084176*$(x[3]) + -0.617200197422199*$(x[4])))) + 0.22410021012329473*1/(1 + exp(-(0.934307327106453 + 0.477402449923221*$(x[1]) + -0.9131747528764835*$(x[2]) + 0.07837215353953164*$(x[3]) + 0.2564992344946111*$(x[4]))))))))))))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


function main()
	N = 100
	s = string( factorial(big(N)) )
	reduce(+, map(parseint,collect(s)))
end

main()


using Test
using TestSetExtensions

using JSON
using LinearAlgebra: SingularException

using Nemo

include("../src/ExactODEReduction.jl")
using .ExactODEReduction: from_dense, to_dense, from_COO, zero_sparsik, random_sparsik,
                          construct_jacobians, check_invariance!, @sparse, linear_span!,
                          check_inclusion!, density, is_thorough, transpose!, tr, dim, scale,
                          to_cartesian, to_plain, load_COO_set, load_MTX, find_radical,
                          rational_reconstruction, modular_reduction, Sparsik, inner,
                          scale!, reduce!, first_nonzero, find_basis_1_β, find_basis_1,
                          minimal_polynomial_wiedemann, unit_sparsik, evaluate, reconstruct!,
                          square_nonsingular_deterministic_wiedemann, square_nonsingular_randomized_wiedemann

@info "Testing started"

@testset "All the tests" begin

    @includetests ARGS

end

@info "All tests OK"


module LLW2d

# Linear Long Wave (LLW) tsunami in 2D Cartesian Coordinate

struct Matrices{T,M<:AbstractMatrix{T}}
    absorbing_boundary::M
    ocean_depth::M
    x_averaged_depth::M
    y_averaged_depth::M
    land_filter_m::M
    land_filter_n::M
    land_filter_e::M
end

const g_n = 9.80665

# Set station locations.
function set_stations!(ist::AbstractVector{Int},
                       jst::AbstractVector{Int},
                       station_distance_x::T,
                       station_distance_y::T,
                       station_boundary_x::T,
                       station_boundary_y::T,
                       dx::T,
                       dy::T) where T
    # synthetic station locations
    nst = 0

    # let's check ist and jst have a square number of elements before computing n
    @assert mod(sqrt(length(ist)),1.0) ≈ 0
    @assert mod(sqrt(length(jst)),1.0) ≈ 0
    n = floor(Int, sqrt(length(ist)))

    @inbounds for i in 1:n, j in 1:n
        nst += 1
        ist[nst] = floor(Int, (station_boundary_x + (i - 1) * station_distance_x) / dx)
        jst[nst] = floor(Int, (station_boundary_y + (j - 1) * station_distance_y) / dy)
    end
end

function timestep!(dx_buffer::AbstractMatrix{T},
                   dy_buffer::AbstractMatrix{T},
                   eta1::AbstractMatrix{T},
                   mm1::AbstractMatrix{T},
                   nn1::AbstractMatrix{T},
                   eta0::AbstractMatrix{T},
                   mm0::AbstractMatrix{T},
                   nn0::AbstractMatrix{T},
                   absorbing_boundary::AbstractMatrix{T},
                   x_averaged_depth::AbstractMatrix{T},
                   y_averaged_depth::AbstractMatrix{T},
                   land_filter_m::AbstractMatrix{T},
                   land_filter_n::AbstractMatrix{T},
                   land_filter_e::AbstractMatrix{T},
                   dx::Real,dy::Real,dt::Real) where T
    nx, ny = size(eta1)
    @assert (nx, ny) == size(mm1) == size(nn1) == size(eta0) == size(mm0) ==
        size(nn0) == size(x_averaged_depth) == size(y_averaged_depth) ==
        size(land_filter_m) == size(land_filter_n) == size(land_filter_e) ==
        size(absorbing_boundary)

    # diffs
    for j in 1:ny
        for i in 2:nx
            @inbounds dx_buffer[i,j] = (eta0[i,j] - eta0[i - 1,j]) / dx
        end
        @inbounds dx_buffer[1,j] = (eta0[1,j] - 0) / dx
    end
    for i in 1:nx
        for j in 2:ny
            @inbounds dy_buffer[i,j] = (eta0[i,j] - eta0[i, j - 1]) / dy
        end
        @inbounds dy_buffer[i,1] = (eta0[i,1] - 0) / dy
    end

    # Update Velocity
    for j in 1:ny, i in 1:nx
        @inbounds mm1[i,j] = mm0[i, j] - g_n * x_averaged_depth[i, j] * dx_buffer[i, j] * dt
        @inbounds nn1[i,j] = nn0[i, j] - g_n * y_averaged_depth[i, j] * dy_buffer[i, j] * dt
    end

    # boundary condition
    for j in 1:ny, i in 1:nx
        @inbounds mm1[i, j] = mm1[i, j] * land_filter_m[i, j] * absorbing_boundary[i, j]
        @inbounds nn1[i, j] = nn1[i, j] * land_filter_n[i, j] * absorbing_boundary[i, j]
    end

    # diffs
    for j in 1:ny
        @inbounds dx_buffer[nx, j] = (-mm1[nx, j]) / dx
        for i in 1:(nx-1)
            @inbounds dx_buffer[i, j] = (mm1[i + 1, j] - mm1[i, j]) / dx
        end
    end
    for i in 1:nx
        @inbounds dy_buffer[i, ny] = (-nn1[i, ny]) / dy
        for j in 1:(ny-1)
            @inbounds dy_buffer[i, j] = (nn1[i,j + 1] - nn1[i, j]) / dy
        end
    end

    # Update Wave Heigt
    for j in 1:ny, i in 1:nx
        @inbounds eta1[i, j] = eta0[i, j] - (dx_buffer[i, j] + dy_buffer[i, j]) * dt
    end

    # boundary condition
    for j in 1:ny, i in 1:nx
        @inbounds eta1[i, j] = eta1[i, j] * land_filter_e[i, j] * absorbing_boundary[i, j]
    end
    return eta1, mm1, nn1
end

function timestep!(dx_buffer::AbstractMatrix{T},
                   dy_buffer::AbstractMatrix{T},
                   eta1::AbstractMatrix{T},
                   mm1::AbstractMatrix{T},
                   nn1::AbstractMatrix{T},
                   eta0::AbstractMatrix{T},
                   mm0::AbstractMatrix{T},
                   nn0::AbstractMatrix{T},
                   matrices::Matrices{T},
                   dx::Real,dy::Real,dt::Real) where T
    # Unpack the relevant fields of `matrices`
    return timestep!(dx_buffer, dy_buffer, eta1, mm1, nn1, eta0, mm0, nn0,
                     matrices.absorbing_boundary, matrices.x_averaged_depth,
                     matrices.y_averaged_depth, matrices.land_filter_m,
                     matrices.land_filter_n, matrices.land_filter_e, dx, dy, dt)
end

function setup(nx::Int,
               ny::Int,
               bathymetry_val::Real,
               absorber_thickness_fraction::Real,
               apara::Real,
               cutoff_depth::Real,
               T::DataType = Float64)
    # Memory allocation
    ocean_depth = Matrix{T}(undef, nx, ny)
    x_averaged_depth = Matrix{T}(undef, nx, ny)
    y_averaged_depth = Matrix{T}(undef, nx, ny)
    absorbing_boundary = ones(T, nx, ny)
    land_filter_m = ones(T, nx, ny) # land filters
    land_filter_n = ones(T, nx, ny) # "
    land_filter_e = ones(T, nx, ny) # "

    nxa = floor(Int, nx * absorber_thickness_fraction)
    nya = floor(Int, nx * absorber_thickness_fraction)

    # Bathymetry set-up. Users may need to modify it
    fill!(ocean_depth, bathymetry_val)
    @inbounds for j in 1:ny, i in 1:nx
        if ocean_depth[i,j] < 0
            ocean_depth[i,j] = 0
        elseif ocean_depth[i,j] < cutoff_depth
            ocean_depth[i,j] = cutoff_depth
        end
    end

    # average bathymetry for staabsorbing_boundaryered-grid computation
    for j in 1:ny
        for i in 2:nx
            x_averaged_depth[i, j] = (ocean_depth[i, j] + ocean_depth[i - 1, j]) / 2
            if ocean_depth[i, j] <= 0 || ocean_depth[i - 1, j] <= 0
                x_averaged_depth[i, j] = 0
            end
        end
        x_averaged_depth[1, j] = ocean_depth[1, j]
    end
    for i in 1:nx
        for j in 2:ny
            y_averaged_depth[i, j] = (ocean_depth[i, j] + ocean_depth[i, j - 1]) / 2
            if ocean_depth[i, j] <= 0 || ocean_depth[i, j - 1] <= 0
                y_averaged_depth[i, j] = 0
            end
        end
        y_averaged_depth[i, 1] = ocean_depth[i, 1]
    end

    # Land filter
    @inbounds for j in 1:ny,i in 1:nx
        (x_averaged_depth[i, j] < 0) && (land_filter_m[i, j] = 0)
        (y_averaged_depth[i, j] < 0) && (land_filter_n[i, j] = 0)
        (ocean_depth[i, j] < 0) && (land_filter_e[i, j] = 0)
    end

    # Sponge absorbing boundary condition by Cerjan (1985)
    @inbounds for j in 1:ny, i in 1:nx
        if i <= nxa
            absorbing_boundary[i, j] *= exp(-((apara * (nxa - i)) ^ 2))
        end
        if i >= nx - nxa + 1
            absorbing_boundary[i, j] *= exp(-((apara * (i - nx + nxa - 1)) ^ 2))
        end
        if j <= nya
            absorbing_boundary[i, j] *= exp(-((apara * (nya - j)) ^ 2))
        end
        if j >= ny - nya + 1
            absorbing_boundary[i, j] *= exp(-((apara * (j - ny + nya - 1)) ^ 2))
        end
    end
    return Matrices(absorbing_boundary, ocean_depth, x_averaged_depth, y_averaged_depth,
                    land_filter_m, land_filter_n, land_filter_e)
end

# Calculates the initial surface height at point x,y
function get_initial_height(x::T, y::T, peak_x::T, peak_y::T, peak_height::T, cutoff::T) where T

    dx = peak_x - x
    dy = peak_y - y
    distance_to_peak = sqrt(dx^2 + dy^2)

    if distance_to_peak <= cutoff
        return peak_height / 4.0 * ((1 + cospi(dx / cutoff)) * (1 + cospi(dy / cutoff)))
    else
        return 0.0
    end

end

# Initializes the surface height on all points of the grid
function initheight!(height::AbstractMatrix{T},
                     ocean_depth::AbstractMatrix{T},
                     dx::Real,dy::Real,cutoff_distance::Real,
                     peak_height::Real, peak_position::AbstractVector{T}) where T

    @assert size(height) == size(ocean_depth)

    nx, ny = size(height)

    x = (1:nx) .* dx
    y = (1:ny) .* dy

    height .= get_initial_height.(x', y, peak_position[1], peak_position[2], peak_height, cutoff_distance)
    height[ocean_depth .< eps(T)] .= 0.0

end

# Initializes the surface height on all points of the grid
function initheight!(height::AbstractMatrix{T},
                     matrices::Matrices{T},
                     dx::Real,dy::Real,cutoff_distance::Real,
                     peak_height::Real, peak_position::AbstractVector{T}) where T
    # Unpack the relevant field of `matrices`
    initheight!(height, matrices.ocean_depth, dx, dy, cutoff_distance, peak_height, peak_position)
end

end # module


function dematel(comparisonMat::DataFrame; threshold=nothing)

    n, _ = size(comparisonMat)

    csums =  colsums(comparisonMat)
    rsums =  rowsums(comparisonMat)

    largest = maximum(vcat(csums, rsums))

    K = convert(Array{Float64,2}, comparisonMat)
    ND = K ./ largest

    T = ND * inv(I(n) - ND)

    c = rowsums(T)
    r = colsums(T)

    c_plus_r  = c .+ r
    c_minus_r = c .- r

    if threshold === nothing
        threshold = sum(T) / length(T)
    end

    influence_matrix = zeros(Float64, n, n)
    for i in 1:n
        for j in 1:n
            if T[i, j] > threshold
                influence_matrix[i, j] = 1
            end
        end
    end

    w_row = sqrt.((c_plus_r.^2.0) .+ (c_minus_r.^2.0))
    w = w_row ./ sum(w_row)

    result = DematelResult(
        comparisonMat,
        ND,
        T,
        c,
        r,
        c_plus_r,
        c_minus_r,
        threshold,
        influence_matrix,
        w
    )

    return result
end


using Sundials

## f routine. Compute function f(t,y).

function f(t, y, ydot)
    ydot[1] = -0.04*y[1] + 1.0e4*y[2]*y[3]
    ydot[3] = 3.0e7*y[2]*y[2]
    ydot[2] = -ydot[1] - ydot[3]
    return Sundials.CV_SUCCESS
end

t = [0.0; 4 * exp10.(range(-1., stop=7., length=9))]
y0 = [1.0, 0.0, 0.0]
res = Sundials.cvode(f, y0, t)


using NLsolve, BenchmarkTools

# number of sectors
numSectors = 3

# test initial capital
k = ones(numSectors)

# test initial output statespace
y = ones(numSectors)

# linear production function parameter
ϕ = 0.5
# time discount factor
β = 0.9
# elasticity parameter
ρ = -10
# depreciation
δ = 1
# savings scaling parameters
σ = ones(numSectors, numSectors)

# utility function of consumption vector
u(c) = sum(log.(c))
# derivative of utility w.r.t consumption in one element
u′(c) = 1 / c

# production function of capital
f(k) = ϕ * k
# derivative of production w.r.t capital
f′(k) = ϕ

# future capital from current parameters
kp(s,k) = (1 - δ) * k + s

# this would be replaced by our interpolators in each sector in the iterations
# a g function of the current below form is equivalent to our "first guess" of an interpolator
g(y) = ones(2*numSectors-1)

# consumption as function of S_ij ∀ i,j
c(S,y) = y .- sum(S, 2)

# S_ij as function of S39
function S(S39)
    S = ones(numSectors, numSectors)

    # fills in the known values first (the diagonal and end collumn)
    for i in 1:numSectors
        for j in 1:numSectors
            if i == j
                S[i,j] = S39[numSectors - 1 + i]
            elseif j == numSectors
                S[i,j] = S39[i]
            end
        end
    end
    # make partial Sb from the know values of S above
    Sb = S ./ σ

    # fills in the empty values in S using our partial Sb
    for i in 1:numSectors
        for j in 1:numSectors
            if i != j && j != numSectors
                S[i,j] = sigma[i,j] * Sb[jj] * Sb[i,numSectors] / Sb[j,numSectors]
            end
        end
    end

    # make full Sbb
    for i in 1:numSectors
        for j in 1:numSectors
            Sbb[i,j] = Sb[i,j] / Sb[j,j]
        end
    end

    # return savings matrix
    return S
end

Sb(S) = S ./ σ

# s_j from full S_ij
s(S) = (sum(σ[l,:] .* S[l,:].^ρ for l in 1:numSectors)).^(1/ρ)

# NLsolve approach
function f!(F, S39)
    for i in 1:2*numSectors-1
        if i < numSectors
            # Nth (last) collumn of double blackboard S matrix, excluding diagonal corner
            F[i] = u′(c(S, y)[numSectors]) / u′(c(S, y)[i]) - Sb(S(S39))[i,numSectors]/Sb(S(S39))[numSectors,numSectors]
        else
            # Diagonals of S matrix
            j = i - numSectors + 1
            F[i] = (f′(β) * β * u′(c(S(g(y)), kp(s(S), k))[j])) / (u′(c(S, y)))^(-rho/(1 - rho)) - S39[j] / (σ[j,j] * s(S))
        end
    end
end

# warm start for savings values
initial_x = 1.5*ones(2*numSectors-1)

# solve
nlsolve(f!, initial_x)


# Partial trace

using Convex
using SCS

import Base.sign
import Convex.monotonicity
import Convex.conic_form!
import Convex.curvature

export ptrace, partialtrace

"""
Type to be used internally by Convex.jl
"""

type PartialTraceAtom <: AbstractExpr
	head::Symbol
	id_hash::UInt64
	children::Tuple{AbstractExpr}
	size::Tuple{Integer, Integer}
	sys::Integer
	dims::Vector

	function PartialTraceAtom(x::AbstractExpr, sys::Integer, dims::Vector)
		if x.size[1] ≠ x.size[2]
			error("Only square matrices are supported")
		end

		if ! (1 ≤ sys ≤ length(dims))
			error("Invalid system, should between 1 and ", length(dims))
		end

		if x.size[1] ≠ prod(dims)
			error("Dimension of system doesnt correspond to dimension of subsystems")
		end

		children = (x, )
		newsize = (round(Int, x.size[2]/dims[sys]), round(Int, x.size[1]/dims[sys]))

		return new(:partialTrace,
		hash(children),
		children,
		newsize,
		sys,
		dims)
	end
end

function sign(x::PartialTraceAtom)
	 return NoSign()
	 return sign(x.children[0])
 end

function curvature(x::PartialTraceAtom)
  return ConstVexity()
end

function monotonicity(x::PartialTraceAtom)
	return (NoMonotonicity(),)
end

function conic_form!(x::PartialTraceAtom, unique_conic_forms::UniqueConicForms)
	if !has_conic_form(unique_conic_forms, x)
		sys = x.sys
		dims = x.dims
		function entry(ρ, j::Integer)
			bra = speye(1)
			ket = speye(1)
			i_sys = 1
			for dim in dims
				if i_sys == sys
					vO = sparsevec([j], [1], dim);
					bra = kron(bra, vO')
					ket = kron(ket, vO)
				else
					bra = kron(bra, speye(dim))
					ket = kron(ket, speye(dim))
				end
				i_sys += 1
			end
			return bra * ρ * ket
		end

		objective = conic_form!(sum([entry(x.children[1], j) for j in 1:dims[sys]]), unique_conic_forms)
		cache_conic_form!(unique_conic_forms, x, objective)
	end
	return get_conic_form(unique_conic_forms, x)
end

function partialtrace(x, systems::Vector, dims::Vector)
	dims = copy(dims)
	out = x
	for sys in systems
		out = partialtrace(out, sys, dims)
		dims[sys] = 1
	end
	return out
end

# General function, works in Convex.jl and with normal matrices

""" `rhoOut = partialtrace(rhoIn, sysNum, dimVec)`

Computes the partial trace of *rhoIn*, tracing out system number *sysNum*, where the dimensions of the systems are specified in the vector *dimVec*.

Example: rhoOut = partialtrace(epr, 2, [2 2])

Traces out the second qubit of a 2x2 state, here the EPR pair.
"""

function partialtrace(x, sys::Integer, dims::Vector)

	# Define a helper function returning matrix entries
	function entry(rho, j::Integer)
		bra = speye(1)
		ket = speye(1)
		i_sys = 1
		for dim in dims
			if i_sys == sys
				vO = sparsevec([j], [1], dim)
				bra = kron(bra, vO')
				ket = kron(ket, vO)
			else
				bra = kron(bra, speye(dim))
				ket = kron(ket, speye(dim))
			end
			i_sys += 1
		end
		return bra * rho * ket
	end

	# Compute the partial trace by picking the appropriate entries
	return sum([entry(x, j) for j in 1:dims[sys]])
end

ptrace(x::AbstractExpr, sys::Integer, dim::Vector) = partialtrace(x, sys, dim)
ptrace(x::AbstractArray, sys::Integer, dim::Vector) = partialtrace(x, sys, dim)


"""
Calculate water critical properties
T --- temperature in Kelvin
p --- pressure in bar
ro --- density in mol/l
doi: 10.1063/1.555786
"""
function methanolCritical()
	sub=newCritical()
	sub.T=512.50    #K
	sub.p=80.9464   #bar
	sub.ro= 8.40    #mol/l #322 kg/m^3
	return sub
end

"""
Calculate water properties at triple point
T --- temperature in Kelvin
p --- pressure in bar
vRo --- vapor density in mol/l
lRo --- liquid density in mol/l
doi: 10.1063/1.555786
"""
function methanolTriple()
	sub=newTriple()
	sub.T=175.59
	sub.p=1.835*10.0^-6	#bar
	sub.lRo=28.226	#kg/m^3 to mol/l
	sub.vRo=1.271*10.0^(-7)	#kg/m^3s to mol/l
	return sub
end

"""
calculate methanol properties at saturation line
T --- temperature in Kelvin
p --- pressure in bar
vRo --- vapor density in mol/l
lRo --- liquid density in mol/l
doi: 10.1063/1.555786
"""
function waterSaturate(T)
a=[-10.75284879, 16.758206642, -3.603424623, 4.373231941, -2.381377449, 4.572198698]    #pressure
b=[2.51730906, -2.46669419, 3.06681768, -1.32507680]    #liquid density
c=[10.61966850, -2.55668202, 3.81845421, 4.79556752]    #vapor density
sub=newSaturate()
crit=methanolCritical()

if (T<crit.T)
	v=1.0-T/crit.T
	sub.p=exp(a[1]/v + a[2] + a[3]*v + a[4]*v^2 + a[5]*v^3 + a[6]*(1-v)^1.7)
	sub.lRo=crit.ro*(1+b[1]*(1-v)^0.35 + b[2]*(v-1) + b[3]*(v^2-1) + b[4]*(v^3-1))
	sub.vRo=crit.ro*exp(c[1]*(1 - 1/v) + c[2]*(1-v)^0.35 + c[3]*(1-v) + c[4]*(1-v)^2)
	sub.T=T
else
	sub.p=-1
	sub.lRo=-1
	sub.vRo=-1
	sub.T=-1
end
return sub
end


using LinearAlgebra

export calcmoit, eigenmoit

function calcmoit(atoms)
    # We center the atoms at their center of mass
    Atoms = Molecules.translate(atoms, Molecules.center_of_mass(atoms))
    I = zeros(Float64, (3,3))
    len = size(Atoms)[1]
    for i = 1:3, j = 1:3
        if i == j
            a = i % 3 + 1
            b = (i+1) % 3 + 1
            suum = 0.0
            for k = 1:len
                suum += Atoms[k].mass * (Atoms[k].xyz[a]^2 + Atoms[k].xyz[b]^2)
            end
            I[i,i] = suum
        else
            suum = 0.0
            for k = 1:len
                suum -= Atoms[k].mass * Atoms[k].xyz[i] * Atoms[k].xyz[j]
            end
            I[i,j] = suum
            I[j,i] = I[i,j]
        end
    end
    return I
end

function eigenmoit(moit)
    A = LinearAlgebra.Symmetric(moit)
    eval, evec = LinearAlgebra.eigen(A)
    return (eval, evec)
end


"""
    VaR(lat::Lattice,alpha::Float32)
implements the nested version of the Value-at-Risk at level ``\\alpha`` for the lattice ``lat``.
See [here](https://rubsc.github.io/riskMeasures.jl/dev/) for the univariate version.
"""
function VaR(lat::Lattice,alpha::Float32)
    # look for smallest x such that  P(-states <= x) >alpha , i.e.
    T0 = deepcopy(lat);

    for i=(length(T0.state)-1):-1:1
		for j=1:length(T0.state[i])
			T0.state[i][j] = RiskMeasures.VaR(T0.state[i+1],T0.probability[i+1][j,:,1],alpha)
		end
	end
    return(T0.state[1],T0)
end

"""
    CTE(lat::Lattice,alpha::Float32)
implements the nested version of the Conditional Tail Expectation level ``\\alpha`` for the lattice ``lat``.
See [here](https://rubsc.github.io/riskMeasures.jl/dev/) for the univariate version.
"""
function CTE(lat::Lattice,alpha::Float32)
    T0 = deepcopy(lat);

	for i=(length(T0.state)-1):-1:1
		for j=1:length(T0.state[i])
			T0.state[i][j] = RiskMeasures.CTE(T0.state[i+1],T0.probability[i+1][j,:,1],alpha)
		end
	end

    return(T0.state[1],T0)
end

"""
    EVaR2(lat::Lattice,beta::Float32)
implements the nested version of the Entropic Value-at-Risk at level ``\\beta`` for the lattice ``lat``.
See [here](https://rubsc.github.io/riskMeasures.jl/dev/) for the univariate version. Here the optimization is done with a simple
goldenSearch algorithm.
"""
function EVaR2(lat::Lattice,beta::Float32)
	T0 = deepcopy(lat);

	for i=(length(T0.state)-1):-1:1
		for j=1:length(T0.state[i])
			T0.state[i][j] = RiskMeasures.EVaR2(T0.state[i+1],T0.probability[i+1][j,:,1],beta)[1]
		end
	end

    return(T0.state[1],T0)
end

"""
    EVaR(lat::Lattice,beta::Float32)
implements the nested version of the Entropic Value-at-Risk at level ``\\beta`` for the lattice ``lat``.
See [here](https://rubsc.github.io/riskMeasures.jl/dev/) for the univariate version. The optimization routine uses JuMP and Ipopt.
"""
function EVaR(lat::Lattice,beta::Float32)
	T0 = deepcopy(lat);

    for i=(length(T0.state)-1):-1:1
		for j=1:length(T0.state[i])
			T0.state[i][j] = RiskMeasures.EVaR(T0.state[i+1],T0.probability[i+1][j,:,1],beta)[1]
		end
	end
    return(T0.state[1],T0)
end

"""
    AVaR(lat::Lattice,alpha::Float32)
implements the nested version of the Average Value-at-Risk at level ``\\alpha`` for the lattice ``lat``.
See [here](https://rubsc.github.io/riskMeasures.jl/dev/) for the univariate version.
"""
function AVaR(lat::Lattice, alpha::Float32)
    T0 = deepcopy(lat);

    for i=(length(T0.state)-1):-1:1
		for j=1:length(T0.state[i])
			T0.state[i][j] = RiskMeasures.AVaR(T0.state[i+1],T0.probability[i+1][j,:,1],alpha)[1]
		end
	end
    return(T0.state[1],T0)
end


using PyPlot
using Statistics

using SubmodularMaximization

num_trials = 50

num_events = 50
num_agents = 50
num_sensors = 10
max_success_probability = 1.0

struct ProbabilisticCoverageParams
  nominal_area
  radius_factor
end

area_range = collect(0.2:0.4:3.0)
radius_factor_range = collect(1:1.0:8.0)

param_set = [ProbabilisticCoverageParams(nominal_area, radius_factor)
             for nominal_area in area_range,
             radius_factor in radius_factor_range]

function performance_ratio(params::ProbabilisticCoverageParams)
  println("Testing: area ($(params.nominal_area)), radius-factor ($(params.radius_factor))")

  sensor_radius = sqrt(params.nominal_area / (num_agents * pi))
  station_radius = params.radius_factor * sensor_radius

  agent_specification = ProbabilisticAgentSpecification(max_success_probability,
                                                        sensor_radius,
                                                        station_radius,
                                                        num_sensors)

  problems = map(1:num_trials) do unused
    events = generate_events(num_events)
    f(x) = mean_detection_probability(x, events)

    agents = generate_agents(agent_specification, num_agents)

    ExplicitPartitionProblem(f, agents)
  end

  myopic_mean = mean(map(x->solve_myopic(x).value, problems))
  sequential_mean = mean(map(x->solve_sequential(x).value, problems))

  myopic_mean / sequential_mean
end

performance_ratios = map(performance_ratio, param_set)

_, best_index = findmin(performance_ratios)
best_params = param_set[best_index]

println("Best ratio: $(performance_ratios[best_index])")
println("Params: area ($(best_params.nominal_area)), radius-factor ($(best_params.radius_factor))")

figure()
PyPlot.plt.hist(performance_ratios[:], 10)

figure()
pcolormesh(area_range, radius_factor_range, performance_ratios', cmap="Blues")
colorbar()
xlabel("nominal area")
ylabel("radius factor")

nothing


using MLDatasets
using Flux
using BSON
using Random
using Statistics
using Base.Iterators: partition
using Flux: crossentropy, onehotbatch, onecold

accuracy(x, y) = mean(onecold(cpu(m(x))) .== onecold(cpu(y)))

function reshape_data(X::AbstractArray{T, 3}, y::AbstractVector) where T
    s = size(X)
    return reshape(X, s[1], s[2], 1, s[3]), reshape(y, 1, :)
end

function train_model!(m, X, y;
    opt=ADAM(0.001),
    batch_size=128,
    n_epochs=10,
    file_name="")

    loss(x, y) = crossentropy(m(x), y)

    batches_train = map(partition(randperm(size(y, 2)), batch_size)) do inds
        return (gpu(X[:, :, :, inds]), gpu(y[:, inds]))
    end

    for i in 1:n_epochs
        println("Iteration " * string(i))
        Flux.train!(loss, params(m), batches_train, opt)
    end

    !isempty(file_name) && BSON.bson(file_name, m=m|>cpu)

    return
end

function load_data(dataset; onehot=false, T=Float32)
    classes = 0:9
    X_train, y_train = reshape_data(dataset.traindata(T)...)
    X_test, y_test = reshape_data(dataset.testdata(T)...)
    y_train = T.(y_train)
    y_test = T.(y_test)

    if onehot
        y_train = onehotbatch(y_train[:], classes)
        y_test = onehotbatch(y_test[:], classes)
    end

    return X_train, y_train, X_test, y_test
end

dataset = MLDatasets.MNIST
T = Float32
X_train, y_train, X_test, y_test = load_data(dataset; T=T, onehot=true)

m = Chain(
    Conv((2,2), 1=>16, sigmoid),
    MaxPool((2,2)),
    Conv((2,2), 16=>8, sigmoid),
    MaxPool((2,2)),
    flatten,
    Dense(288, size(y_train,1)), softmax) |> gpu

file_name = joinpath("data", "mnist_sigmoid.bson")
train_model!(m, X_train, y_train; file_name=file_name, n_epochs=100)

accuracy(X_test |> gpu, y_test |> gpu)


using VortexDistributions, Test, SafeTestsets

@safetestset "Point Vortex" begin include("point_vortex_test.jl") end
@safetestset "Single Vortex Accuracy" begin include("single_vort_acc_test.jl") end
@safetestset "Multiple Vortices" begin include("multivort_test.jl") end
@safetestset "Periodic Dipole" begin include("periodic_dipole_test.jl") end
@safetestset "Creation functions" begin include("creations_test.jl") end
@safetestset "Multiple Vortices 2" begin include("multivort_test2.jl") end
@safetestset "Unwrap" begin include("test_unwrap.jl") end
@safetestset "Unwrap2D" begin include("test_unwrap2D.jl") end
@safetestset "Phase Jumps" begin include("phase_jumps.jl") end
@safetestset "Periodic Dipole" begin include("periodic_dipole_test.jl") end
@safetestset "Scalar Vortex Arguments" begin include("vortex_types.jl") end
@safetestset "Single Vortex Sphere Field" begin include("single_vort_sphere.jl") end
@safetestset "Keep Vortices" begin include("keep_vortices_test.jl") end


# Inverse CDF method
#
# copied from Distributions.jl, Feb 14, 2015
function randtn_invcdf(d::Truncated)
    d0 = d.untruncated
    if d.tp > 0.25
        while true
            r = rand(d0)
            if d.lower <= r <= d.upper
                return r
            end
        end
    else
        return quantile(d0, d.lcdf + rand() * d.tp)
    end
end

# Rejection sampler based on algorithm from Robert (1995)
#
#  - Available at http://arxiv.org/abs/0907.4010
#
# copied from Distributions.jl, Feb 14, 2015
function randtn_robert(d::Truncated{Normal})
    d0 = d.untruncated
    μ = mean(d0)
    σ = std(d0)
    a = (d.lower - μ) / σ
    b = (d.upper - μ) / σ
    z = randtn_robert_base(a, b, d.tp)
    return μ + σ * z
end

function randtn_robert_base(lb::Float64, ub::Float64, tp::Float64)
    r::Float64
    if tp > 0.3   # has considerable chance of falling in [lb, ub]
        r = randn()
        while r < lb || r > ub
            r = randn()
        end
        return r

    else
        span = ub - lb
        if lb > 0 && span > 2.0 / (lb + sqrt(lb^2 + 4.0)) * exp((lb^2 - lb * sqrt(lb^2 + 4.0)) / 4.0)
            a = (lb + sqrt(lb^2 + 4.0))/2.0
            while true
                r = rand(Exponential(1.0 / a)) + lb
                u = rand()
                if u < exp(-0.5 * (r - a)^2) && r < ub
                    return r
                end
            end
        elseif ub < 0 && ub - lb > 2.0 / (-ub + sqrt(ub^2 + 4.0)) * exp((ub^2 + ub * sqrt(ub^2 + 4.0)) / 4.0)
            a = (-ub + sqrt(ub^2 + 4.0)) / 2.0
            while true
                r = rand(Exponential(1.0 / a)) - ub
                u = rand()
                if u < exp(-0.5 * (r - a)^2) && r < -lb
                    return -r
                end
            end
        else
            while true
                r = lb + rand() * (ub - lb)
                u = rand()
                if lb > 0
                    rho = exp((lb^2 - r^2) * 0.5)
                elseif ub < 0
                    rho = exp((ub^2 - r^2) * 0.5)
                else
                    rho = exp(-r^2 * 0.5)
                end
                if u < rho
                    return r
                end
            end
        end
    end
end


###################################################################################################
# FUNCTIONS
###################################################################################################
"""
    `generate_xcal(N)`

Generate state space 𝒳 for the methylation vector with N CpG sites.

# Examples
```julia-repl
julia> xcal=CpelAsm.generate_xcal(2)
4-element Array{Vector{Int64},1}:
 [-1, -1]
 [1, -1]
 [-1, 1]
 [1, 1]
```
"""
function generate_xcal(n::Int64)::Array{Vector{Int64},1}

    # Generate iterative object
    xcal = Vector{Int64}[]
    [push!(xcal,2 .* digits(i, base=2, pad=n) - ones(n)) for i in 0:(2^n-1)]

    # Return 𝒳
    return xcal

end # end generate_xcal
"""
    `gen_ising_full_data(M,N;α,β)`

Generate M FULLY observed reads with [N1,...,NK] CpG sites from energy function
U(X) parametrized by [α1,...,αK] and β.

# Examples
```julia-repl
julia> Random.seed!(1234);
julia> xobs=CpelAsm.gen_ising_full_data(10,[4])
10-element Array{Array{Int64,1},1}:
 [-1, 1, -1, -1]
 [1, -1, 1, -1]
 [-1, 1, 1, -1]
 [1, 1, 1, -1]
 [1, 1, 1, -1]
 [-1, 1, -1, 1]
 [-1, 1, -1, 1]
 [-1, 1, -1, 1]
 [-1, -1, 1, 1]
 [1, 1, 1, 1]
```
"""
function gen_ising_full_data(m::Int64,n::Vector{Int64};a=zeros(length(n)),b=0.0)::Array{Vector{Int64},1}

    # Generate iterative object
    xcal = generate_xcal(sum(n))

    # Create function for energy computation
    Ux_fun = create_Ux(n,a,b)

    # Compute probability over state space
    p = zeros(2^sum(n))
    p = [exp(-Ux_fun(xcal[k])) for k in 1:length(xcal)]
    p = p ./ sum(p)

    # Sample full observations from model
    indices = rand(Distributions.Multinomial(m, p))
    indices = vcat([repeat([i],indices[i]) for i in 1:length(indices)]...)

    # Return observations
    return xcal[indices]

end # end gen_ising_full_data
"""
    `gen_ising_part_data(M,N;α,β)`

Generate M partially observed reads with [N1,...,NK] CpG sites from energy function U(X)
parametrized by [α1,...,αK] and β. The values code for:
    1:  Methylated CpG site.
    -1: Unmethylated CpG site.
    0:  Unobserved CpG site.

# Examples
```julia-repl
julia> Random.seed!(1234);
julia> xobs=gen_ising_part_data(10,[4])
10-element Array{Vector{Int64},1}:
[-1, -1, -1, -1]
[1, -1, -1, -1]
[-1, 1, -1, -1]
[-1, -1, 1, -1]
[1, -1, 1, -1]
[1, -1, -1, 1]
[0, 1, -1, 1]
[1, -1, 1, 1]
[1, -1, 1, 1]
[-1, 1, 1, 1]
```
"""
function gen_ising_part_data(m::Int64,n::Vector{Int64};a=zeros(length(n)),b=0.0)::Array{Vector{Int64},1}

    # Get full data
    full_data = gen_ising_full_data(m,n;a=a,b=b)

    # Add missing values on sides
    missleft = rand(Distributions.Binomial(sum(n), 0.01),m)
    missright = rand(Distributions.Binomial(sum(n), 0.01),m)
    for i in 1:m
        full_data[i][1:missleft[i]] = fill(0, missleft[i])
        full_data[i][(sum(n)-missright[i]+1):end] = fill(0, missright[i])
    end

    # Return partial observations
    return full_data

end # end gen_ising_part_data
"""
    `gen_mult_full_data(M;N,p)`

Generate M FULLY observed reads from a multinomial model with vector of probabilities
`p=[p_1,...,p_{2^N}]`.

# Examples
```julia-repl
julia> Random.seed!(1234);
julia> xobs=CpelAsm.gen_mult_full_data(2)
2-element Array{Array{Int64,1},1}:
 [1, 1, -1, 1]
 [1, 1, -1, 1]
```
"""
function gen_mult_full_data(m::Int64;n=4,p=1.0/(2.0^n)*ones(2^n))::Array{Vector{Int64},1}

    # Generate iterative object
    xcal = generate_xcal(n)

    # Sample full observations from model
    indices = rand(Distributions.Multinomial(m, p))
    indices = vcat([repeat([i],indices[i]) for i in 1:length(indices)]...)

    # Return observations
    return xcal[indices]

end # end gen_mult_full_data
"""
    `gen_mult_part_data(M;N,p)`

Generate M partially observed reads with N CpG sites from a multinomial model with vector
probabilities `p=[p_1,...,p_{2^N}]`. The values code for:
    1:  Methylated CpG site.
    -1: Unmethylated CpG site.
    0:  Unobserved CpG site.

# Examples
```julia-repl
julia> Random.seed!(1234);
julia> xobs=CpelAsm.gen_mult_part_data(10)
10-element Array{Array{Int64,1},1}:
 [0, 0, 0, 0]
 [1, -1, 1, -1]
 [0, 1, 1, -1]
 [1, 1, 1, 0]
 [1, 1, 1, 0]
 [0, 1, -1, 1]
 [0, 1, -1, 1]
 [0, 1, -1, 1]
 [0, -1, 1, 1]
 [1, 1, 1, 1]
```
"""
function gen_mult_part_data(m::Int64;n=4,p=1.0/(2.0^n)*ones(2^n))::Array{Vector{Int64},1}

    # Get full data
    full_data = gen_mult_full_data(m;n=n,p=p)

    # Add missing values on sides
    missleft = rand(Distributions.Binomial(n+1, 0.05),m)
    missright = rand(Distributions.Binomial(n+1, 0.05),m)
    for i in 1:m
        full_data[i][1:missleft[i]] = fill(0, missleft[i])
        full_data[i][(n-missright[i]+1):end] = fill(0, missright[i])
    end

    # Return partial observations
    return full_data

end # end gen_mult_part_data


using Plots

include_implicit_dynamics()
include_ddp()
include_model("rocket3D")

h = 0.05
n = model.n
m = model.m

# control limits
ul = [-1.0; -1.0; 0.0]
uu = [1.0; 1.0; 11.5]

# control projection problem
m = 3
nz = m + 1 + 1 + 1 + 1 + m
nθ = m + 1

idx = [3; 1; 2]
function residual(z, θ, κ)
    u = z[1:m]
    s = z[m .+ (1:1)]
    y = z[m + 1 .+ (1:1)]
    w = z[m + 1 + 1 .+ (1:1)]
    p = z[m + 1 + 1 + 1 .+ (1:1)]
    v = z[m + 1 + 1 + 1 + 1 .+ (1:m)]

    ū = θ[1:m]
    T = θ[m .+ (1:1)]

    [
     u - ū - v - [0.0; 0.0; y[1] + p[1]];
     -y[1] - w[1];
     T[1] - u[3] - s[1];
     w .* s .- κ
     p .* u[3] .- κ
     second_order_cone_product(v[idx], u[idx]) .- κ .* [1.0; 0.0; 0.0]
    ]
end

@variables z[1:nz], θ[1:nθ], κ[1:1]

r = residual(z, θ, κ)
r .= simplify(r)
r_func = eval(Symbolics.build_function(r, z, θ, κ)[2])

rz = Symbolics.jacobian(r, z)
rz = simplify.(rz)
rz_func = eval(Symbolics.build_function(rz, z, θ)[2])

rθ = Symbolics.jacobian(r, θ)
rθ = simplify.(rθ)
rθ_func = eval(Symbolics.build_function(rθ, z, θ)[2])

idx_ineq = collect([3, 4, 6, 7])
idx_soc = [collect([3, 1, 2]), collect([10, 8, 9])]#collect([3, 1, 2]), collect([10, 8, 9])]

# ul = -1.0 * ones(m)
# uu = 1.0 * ones(m)
ū = [0.0, 0.0, 0.0]

z0 = 0.1 * ones(nz)
z0[1:m] = copy(ū)
z0[3] += 1.0
z0[10] += 1.0

θ0 = [copy(ū); uu[3]]

# solver
opts_con = InteriorPointOptions(
  κ_init = 0.1,
  κ_tol = 1.0e-4,
  r_tol = 1.0e-6,
  diff_sol = false)

ip_con = interior_point(z0, θ0,
  r! = r_func, rz! = rz_func,
  rz = similar(rz, Float64),
  rθ! = rθ_func,
  rθ = similar(rθ, Float64),
  idx_ineq = idx_ineq,
  idx_soc = idx_soc,
  opts = opts_con)

opts_jac = InteriorPointOptions(
	κ_init = 0.1,
	κ_tol = 1.0e-2,
	r_tol = 1.0e-6,
	diff_sol = true)

ip_jac = interior_point(z0, θ0,
	r! = r_func, rz! = rz_func,
	rz = similar(rz, Float64),
	rθ! = rθ_func,
	rθ = similar(rθ, Float64),
	idx_ineq = idx_ineq,
    idx_soc = idx_soc,
	opts = opts_jac)

interior_point_solve!(ip_con)
interior_point_solve!(ip_jac)

function soc_projection(x)
    proj = second_order_cone_projection(x[[3;1;2]])[1]
	ip_con.z .= [proj[2:3]; proj[3]; 0.1 * ones(7)]
    ip_con.z[3] += 1.0
    ip_con.z[10] += 1.0
	ip_con.θ .= [x; uu[3]]

	interior_point_solve!(ip_con)

	return ip_con.z[1:m]
end

soc_projection([100.0, 0.0, 0.0])

function soc_projection_jacobian(x)
    proj = second_order_cone_projection(x[[3;1;2]])[1]
	ip_con.z .= [proj[2:3]; proj[3]; 0.1 * ones(7)]
    ip_con.z[3] += 1.0
    ip_con.z[10] += 1.0
	ip_con.θ .= [x; uu[3]]

	interior_point_solve!(ip_jac)

	return ip_jac.δz[1:m, 1:m]
end

soc_projection_jacobian([10.0, 0.0, 100.0])

function fd(model::Rocket3D{Midpoint, FixedTime}, x, u, w, h, t)
	u_proj = soc_projection(u)
	rz(z) = fd(model, z, x, u_proj, w, h, t) # implicit midpoint integration
	x⁺ = newton(rz, copy(x))
    # x⁺ = levenberg_marquardt(rz, copy(x))
	return x⁺
	# return view(x, 1:model.n) + h * f(model, view(x, 1:model.n) + 0.5 * h * f(model, view(x, 1:model.n), view(u, 1:model.m), w), view(u, 1:model.m), w)
end

fd(model, ones(model.n), zeros(model.m), zeros(model.d), h, 1)

function fdx(model::Rocket3D{Midpoint, FixedTime}, x, u, w, h, t)
	u_proj = soc_projection(u)
	rz(z) = fd(model, z, x, u_proj, w, h, t) # implicit midpoint integration
    x⁺ = newton(rz, copy(x))
    # x⁺ = levenberg_marquardt(rz, copy(x))
	∇rz = ForwardDiff.jacobian(rz, x⁺)
	rx(z) = fd(model, x⁺, z, u_proj, w, h, t) # implicit midpoint integration
	∇rx = ForwardDiff.jacobian(rx, x)
	return -1.0 * ∇rz \ ∇rx
end

fdx(model, zeros(model.n), zeros(model.m), zeros(model.d), h, 1)

function fdu(model::Rocket3D{Midpoint, FixedTime}, x, u, w, h, t)
	u_proj = soc_projection(u)
	u_proj_jac = soc_projection_jacobian(u)

	rz(z) = fd(model, z, x, u_proj, w, h, t) # implicit midpoint integration
    x⁺ = newton(rz, copy(x))
    # x⁺ = levenberg_marquardt(rz, copy(x))
	∇rz = ForwardDiff.jacobian(rz, x⁺)
	ru(z) = fd(model, x⁺, x, z, w, h, t) # implicit midpoint integration
	∇ru = ForwardDiff.jacobian(ru, u_proj)
	return -1.0 * ∇rz \ (∇ru * u_proj_jac)
    # return -1.0 * ∇rz \ ∇ru

	# f(z) = fd(model, x, z, w, h, t)
	# return ForwardDiff.jacobian(f, u)
end

fdu(model, zeros(model.n), zeros(model.m), zeros(model.d), h, 1)

# Time
T = 101
# h = 0.05

# Initial conditions, controls, disturbances
x1 = zeros(model.n)
x1[1] = 5.0
x1[2] = -2.5
x1[3] = 10.0
mrp = MRP(RotZ(0.0 * π) * RotY(-0.0 * π) * RotX(0.0 * π))
x1[4:6] = [mrp.x; mrp.y; mrp.z]
x1[9] = -5.0

# visualize!(vis, model, [x1], Δt = h)

xT = zeros(model.n)
# xT[1] = 2.5
# xT[2] = 0.0
xT[3] = model.length
mrpT = MRP(RotZ(0.0 * π) * RotY(0.0))
xT[4:6] = [mrpT.x; mrpT.y; mrpT.z]

u_ref = [0.0; 0.0; 0.0]#model.mass * 9.81]
ū = [u_ref + [1.0e-3; 1.0e-3; 1.0e-3] .* randn(model.m) for t = 1:T-1]
w = [zeros(model.d) for t = 1:T-1]

# Rollout
x̄ = rollout(model, x1, ū, w, h, T)

# visualize!(vis, model, x̄, Δt = h)
# x̄ = linear_interpolation(x1, xT, T)
# plot(hcat(x̄...)')

# Objective
Q = h * [(t < T ? 1.0 * Diagonal([1.0e-1 * ones(3); 0.0 * ones(3); 1.0e-1 * ones(3); 1000.0 * ones(3)])
        : 1000.0 * Diagonal(1.0 * ones(model.n))) for t = 1:T]
q = h * [-2.0 * Q[t] * xT for t = 1:T]

R = h * [Diagonal([100000.0; 100000.0; 1000.0]) for t = 1:T-1]
r = h * [-2.0 * R[t] * u_ref  for t = 1:T-1]

obj = StageCosts([QuadraticCost(Q[t], q[t],
	t < T ? R[t] : nothing, t < T ? r[t] : nothing) for t = 1:T], T)

function g(obj::StageCosts, x, u, t)
	T = obj.T
    if t < T
		Q = obj.cost[t].Q
		q = obj.cost[t].q
	    R = obj.cost[t].R
		r = obj.cost[t].r
        return x' * Q * x + q' * x + u' * R * u + r' * u
    elseif t == T
		Q = obj.cost[T].Q
		q = obj.cost[T].q
        return x' * Q * x + q' * x
    else
        return 0.0
    end
end

# Constraints
# x_con = [-0.5; 0.5]
# y_con = [-0.75; 0.75]
p = [t < T ? 2 * m + 1 : 2 * n for t = 1:T]
info_t = Dict(:ul => ul, :uu => uu, :inequality => (1:2 * m + 1))
info_T = Dict(:xT => xT, :inequality => (1:(2 * n)))
con_set = [StageConstraint(p[t], t < T ? info_t : info_T) for t = 1:T]
# idx_T = collect([3, 7, 8, 9, 10, 11, 12])
# xT[4:6]
function c!(c, cons::StageConstraints, x, u, t)
	T = cons.T
	p = cons.con[t].p

	if t < T
		ul = cons.con[t].info[:ul]
		uu = cons.con[t].info[:uu]
		# c[1:2 * m] .= [ul - u; u - uu]
		c[2 * m + 1] = model.length - x[3]
	elseif t == T
		xT = cons.con[T].info[:xT]
        c[(1:n)] .= (xT - [0.25, 0.45, 0.0, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001]) - x
        c[n .+ (1:n)] .= x - (xT + [0.25, 0.45, 0.01, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001])
		# c[4 .+ (1:(n - 5))] .= (x - xT)[idx_T]
        # c[4 + (n - 5) .+ (1:3)] .= [-0.05; -0.05; -0.1] - x[4:6]
        # c[4 + (n - 5) + 3 .+ (1:3)] = x[4:6] - [0.05; 0.05; 0.3]
	else
		nothing
	end
end

prob = problem_data(model, obj, con_set, copy(x̄), copy(ū), w, h, T,
	analytical_dynamics_derivatives = true)

u_ref = [0.0; 0.0; 0.0] # model.mass * 9.81]
ū = [u_ref + [1.0e-2; 1.0e-2; 1.0e-2] .* randn(model.m) for t = 1:T-1]

# Solve
@time constrained_ddp_solve!(prob,
    max_iter = 1000, max_al_iter = 10,
	con_tol = 0.001,
	ρ_init = 1.0, ρ_scale = 10.0)

x, u = current_trajectory(prob)
x̄, ū = nominal_trajectory(prob)

x̄_soc = x̄
ū_soc = ū

# @save "/home/taylor/Research/motion_planning/examples/differential_dynamic_programming/implicit_dynamics/rocket_landing_soc.jld2" x̄_soc ū_soc
# @load "/home/taylor/Research/motion_planning/examples/differential_dynamic_programming/implicit_dynamics/rocket_landing_soc.jld2"

all([second_order_cone_projection(ū[t][idx])[2] for t = 1:T-1])
[second_order_cone_projection(ū[t][idx])[2] for t = 1:T-1]
t = 4
fd(model, x̄[t], ū[t], zeros(model.d), h, t) - x̄[t+1]
soc_projection(ū[t])
t = 1
norm(ū[t][1:2]) <= ū[t][3]

# Trajectories
maximum(hcat(ū...))
t = range(0, stop = h * (T-1), length = T)
plt = plot(t, hcat(ū..., ū[end])',
	width = 2.0,
	color = [:magenta :orange :cyan],
	title = "rocket soft landing",
	xlabel = "time (s)",
	ylabel = "control",
	labels = ["f1" "f2" "f3"],
	legend = :right,
	linetype = :steppost)

# savefig(plt, "/home/taylor/Research/implicit_dynamics_manuscript/figures/rocket_control.png")

plot(hcat(x̄...)[1:3, :]', linetype = :steppost)

# Visualize
include(joinpath(pwd(), "models/visualize.jl"))

vis = Visualizer()
render(vis)
# open(vis)
x_anim = [[x̄[1] for i = 1:50]..., x̄..., [x̄[end] for i = 1:50]...]
visualize!(vis, model,
	x_anim,
	Δt = h, T_off = length(x_anim)-75)
obj_platform = joinpath(pwd(), "models/rocket/space_x_platform.obj")
mtl_platform = joinpath(pwd(), "models/rocket/space_x_platform.mtl")
ctm_platform = ModifiedMeshFileObject(obj_platform,mtl_platform,scale=1.0)

setobject!(vis["platform"],ctm_platform)
settransform!(vis["platform"], compose(Translation(0.0,2.0,-0.6), LinearMap(0.75 * RotZ(pi)*RotX(pi/2))))

settransform!(vis["/Cameras/default"], compose(Translation(0.0, 0.0, 0.0),
	LinearMap(RotZ(1.5 * π + 0.0 * π))))
setvisible!(vis["/Grid"], false)
setprop!(vis["/Cameras/default/rotated/<object>"], "zoom", 1.0)

# line_mat = LineBasicMaterial(color=color=RGBA(1.0, 153.0 / 255.0, 51.0 / 255.0, 1.0), linewidth=10.0)
# points = Vector{Point{3,Float64}}()
# for xt in x̄
# 	push!(points, Point(xt[1:3]...))
# end
# setobject!(vis[:traj], MeshCat.Line(points, line_mat))
#
# t = 35
#
# setobject!(vis["rocket2"]["starship"], ctm)
#
# settransform!(vis["rocket2"]["starship"],
# 	compose(Translation(0.0, 0.0, -model.length),
# 		LinearMap(0.25 * RotY(0.0) * RotZ(0.5 * π) * RotX(0.5 * π))))
#
# body = Cylinder(Point3f0(0.0, 0.0, -1.25),
#   Point3f0(0.0, 0.0, 0.5),
#   convert(Float32, 0.125))
#
# setobject!(vis["rocket2"]["body"], body,
#   MeshPhongMaterial(color = RGBA(1.0, 0.0, 0.0, 1.0)))
#
# settransform!(vis["rocket2"],
# 	  compose(Translation(x̄[t][1:3]),
# 			LinearMap(MRP(x̄[t][4:6]...) * RotX(0.0))))
#
# t = 85
#
# setobject!(vis["rocket3"]["starship"], ctm)
#
# settransform!(vis["rocket3"]["starship"],
# 	compose(Translation(0.0, 0.0, -model.length),
# 		LinearMap(0.25 * RotY(0.0) * RotZ(0.5 * π) * RotX(0.5 * π))))
#
# body = Cylinder(Point3f0(0.0, 0.0, -1.25),
#   Point3f0(0.0, 0.0, 0.5),
#   convert(Float32, 0.125))
#
# setobject!(vis["rocket3"]["body"], body,
#   MeshPhongMaterial(color = RGBA(1.0, 0.0, 0.0, 1.0)))
#
# settransform!(vis["rocket3"],
# 	  compose(Translation(x̄[t][1:3]),
# 			LinearMap(MRP(x̄[t][4:6]...) * RotX(0.0))))
#
#
# t = T
#
# setobject!(vis["rocket4"]["starship"], ctm)
#
# settransform!(vis["rocket4"]["starship"],
# 	compose(Translation(0.0, 0.0, -model.length),
# 		LinearMap(0.25 * RotY(0.0) * RotZ(0.5 * π) * RotX(0.5 * π))))
#
# body = Cylinder(Point3f0(0.0, 0.0, -1.25),
#   Point3f0(0.0, 0.0, 0.5),
#   convert(Float32, 0.125))
#
# setobject!(vis["rocket4"]["body"], body,
#   MeshPhongMaterial(color = RGBA(1.0, 0.0, 0.0, 1.0)))
#
# settransform!(vis["rocket4"],
# 	  compose(Translation(x̄[t][1:3]),
# 			LinearMap(MRP(x̄[t][4:6]...) * RotX(0.0))))
#
#
# # # Visualize
# # obj_rocket = joinpath(pwd(), "models/starship/Starship.obj")
# # mtl_rocket = joinpath(pwd(), "models/starship/Starship.mtl")
# # ctm = ModifiedMeshFileObject(obj_rocket, mtl_rocket, scale=1.0)
# # setobject!(vis["rocket"]["starship"], ctm)
# #
# # settransform!(vis["rocket"]["starship"],
# # 	compose(Translation(0.0, 0.0, -model.length),
# # 		LinearMap(0.25 * RotY(0.0) * RotZ(0.5 * π) * RotX(0.5 * π))))
# #
# # default_background!(vis)
# # settransform!(vis["rocket"],
# # 	compose(Translation(0.0, 0.0, 0.0),
# # 	LinearMap(RotY(0.0))))
#
#
# using PGFPlots
# const PGF = PGFPlots
#
# plt_F1_smooth = PGF.Plots.Linear(t, hcat(ū_nominal..., ū_nominal[end])[3,:],
# 	mark="none",style="const plot, color=cyan, line width = 2pt, dashed",legendentry="F1")
#
# plt_F2_smooth = PGF.Plots.Linear(t, hcat(ū_nominal..., ū_nominal[end])[1,:],
# 	mark="none",style="const plot, color=orange, line width = 2pt, dashed",legendentry="F2")
#
# plt_F3_smooth = PGF.Plots.Linear(t, hcat(ū_nominal..., ū_nominal[end])[2,:],
# 	mark="none",style="const plot, color=magenta, line width = 2pt, dashed",legendentry="F2")
#
# plt_F1_soc = PGF.Plots.Linear(t, hcat(ū_soc..., ū_soc[end])[3,:],
# 	mark="none",style="const plot, color=cyan, line width = 2pt",legendentry="F1 (soc)")
#
# plt_F2_soc = PGF.Plots.Linear(t, hcat(ū_soc..., ū_soc[end])[1,:],
# 	mark="none",style="const plot, color=orange, line width = 2pt",legendentry="F2 (soc)")
#
# plt_F3_soc = PGF.Plots.Linear(t, hcat(ū_soc..., ū_soc[end])[2,:],
# 	mark="none",style="const plot, color=magenta, line width = 2pt",legendentry="F2 (soc)")
#
# a = Axis([plt_F1_soc; plt_F2_soc; plt_F3_soc; plt_F1_smooth; plt_F2_smooth; plt_F3_smooth],
#     axisEqualImage=false,
#     hideAxis=false,
# 	ylabel="control",
# 	xlabel="time (s)",
# 	# xlims=(0.0, 2.0),
# 	legendStyle="{at={(0.5,0.5)},anchor=west}")
#
# PGF.save("/home/taylor/Research/implicit_dynamics_manuscript/figures/rocket_control.tikz", a, include_preamble=false)


# This code presents an example to illustrate how NLboxsolve.jl can be used

using NLboxsolve

function test_nlboxsolve_outplace()

    function test3(x)

        f = similar(x)

        f[1] = 0.5*x[1] + x[2] - x[3] + x[1]*x[2] - 1.5
        f[2] = x[1] - 0.5*x[2] + x[3] + x[2]^2 - 2.5
        f[3] = -x[1] + x[2] + 0.5*x[3] + x[2]*x[3] - 1.5
        f[4] = x[1]^2 - x[2]^2 + x[3]^2 + x[4]^2 - 2.0

        return f
    end

    # Solve test3 model

    x0 = [1.5, 1.5, 1.5, 1.5]
    l = [-Inf,-Inf,-Inf,-Inf]
    u = [Inf,Inf,Inf,Inf]

    soln_newton     = constrained_newton(test3,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_lm         = constrained_levenberg_marquardt(test3,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_lm_kyf     = constrained_levenberg_marquardt_kyf(test3,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_lm_ar      = constrained_levenberg_marquardt_ar(test3,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_dogleg     = constrained_dogleg_solver(test3,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_dogleg_bmp = constrained_dogleg_bmp_solver(test3,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_nk         = constrained_newton_krylov(test3,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_jfnk       = constrained_jacobian_free_newton_krylov(test3,x0,l,u,xtol=1e-15,ftol=1e-15)

    soln_newton_sparse     = constrained_newton_sparse(test3,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_sparse         = constrained_levenberg_marquardt_sparse(test3,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_kyf_sparse     = constrained_levenberg_marquardt_kyf_sparse(test3,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_ar_sparse      = constrained_levenberg_marquardt_ar_sparse(test3,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_dogleg_sparse     = constrained_dogleg_solver_sparse(test3,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_dogleg_bmp_sparse = constrained_dogleg_solver_sparse(test3,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_nk_sparse         = constrained_newton_krylov_sparse(test3,x0,l,u,ftol=1e-15,xtol=1e-15)

    soln1 = nlboxsolve(test3,x0,l,u,method = :dogleg,xtol=1e-15,ftol=1e-15)
    soln2 = nlboxsolve(test3,x0,l,u,xtol=1e-15,ftol=1e-15)

    test_one   = maximum(abs,soln_newton.zero       - soln2.zero) < 1e-14
    test_two   = maximum(abs,soln_lm.zero           - soln2.zero) < 1e-14
    test_three = maximum(abs,soln_lm_kyf.zero       - soln2.zero) < 1e-14
    test_four  = maximum(abs,soln_lm_ar.zero        - soln2.zero) < 1e-14
    test_five  = maximum(abs,soln_dogleg.zero       - soln2.zero) < 1e-14
    test_six   = maximum(abs,soln_dogleg_bmp.zero   - soln2.zero) < 1e-14
    test_seven = maximum(abs,soln_nk.zero           - soln2.zero) < 1e-14
    test_eight = maximum(abs,soln_jfnk.zero         - soln2.zero) < 1e-14

    test_nine     = maximum(abs,soln_newton_sparse.zero       - soln2.zero) < 1e-14
    test_ten      = maximum(abs,soln_lm_sparse.zero           - soln2.zero) < 1e-14
    test_eleven   = maximum(abs,soln_lm_kyf_sparse.zero       - soln2.zero) < 1e-14
    test_twelve   = maximum(abs,soln_lm_ar_sparse.zero        - soln2.zero) < 1e-14
    test_thirteen = maximum(abs,soln_dogleg_sparse.zero       - soln2.zero) < 1e-14
    test_fourteen = maximum(abs,soln_dogleg_bmp_sparse.zero   - soln2.zero) < 1e-14
    test_fifteen  = maximum(abs,soln_nk_sparse.zero           - soln2.zero) < 1e-14

    return test_one, test_two, test_three, test_four, test_five, test_six, test_seven, test_eight, test_nine, test_ten, test_eleven,
           test_twelve, test_thirteen, test_fourteen, test_fifteen

end

function test_nlboxsolve_inplace()

    function test3(f,x)

        f[1] = 0.5*x[1] + x[2] - x[3] + x[1]*x[2] - 1.5
        f[2] = x[1] - 0.5*x[2] + x[3] + x[2]^2 - 2.5
        f[3] = -x[1] + x[2] + 0.5*x[3] + x[2]*x[3] - 1.5
        f[4] = x[1]^2 - x[2]^2 + x[3]^2 + x[4]^2 - 2.0

    end

    # Solve test3 model

    x0 = [1.5, 1.5, 1.5, 1.5]
    l = [-Inf,-Inf,-Inf,-Inf]
    u = [Inf,Inf,Inf,Inf]

    soln_newton     = constrained_newton(test3,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_lm         = constrained_levenberg_marquardt(test3,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_lm_kyf     = constrained_levenberg_marquardt_kyf(test3,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_lm_ar      = constrained_levenberg_marquardt_ar(test3,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_dogleg     = constrained_dogleg_solver(test3,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_dogleg_bmp = constrained_dogleg_bmp_solver(test3,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_nk         = constrained_newton_krylov(test3,x0,l,u,xtol=1e-15,ftol=1e-15)
    soln_jfnk       = constrained_jacobian_free_newton_krylov(test3,x0,l,u,xtol=1e-15,ftol=1e-15)

    soln_newton_sparse     = constrained_newton_sparse(test3,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_sparse         = constrained_levenberg_marquardt_sparse(test3,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_kyf_sparse     = constrained_levenberg_marquardt_kyf_sparse(test3,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_lm_ar_sparse      = constrained_levenberg_marquardt_ar_sparse(test3,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_dogleg_sparse     = constrained_dogleg_solver_sparse(test3,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_dogleg_bmp_sparse = constrained_dogleg_solver_sparse(test3,x0,l,u,ftol=1e-15,xtol=1e-15)
    soln_nk_sparse         = constrained_newton_krylov_sparse(test3,x0,l,u,ftol=1e-15,xtol=1e-15)

    soln1 = nlboxsolve(test3,x0,l,u,method = :dogleg,xtol=1e-15,ftol=1e-15)
    soln2 = nlboxsolve(test3,x0,l,u,xtol=1e-15,ftol=1e-15)

    test_one   = maximum(abs,soln_newton.zero       - soln2.zero) < 1e-14
    test_two   = maximum(abs,soln_lm.zero           - soln2.zero) < 1e-14
    test_three = maximum(abs,soln_lm_kyf.zero       - soln2.zero) < 1e-14
    test_four  = maximum(abs,soln_lm_ar.zero        - soln2.zero) < 1e-14
    test_five  = maximum(abs,soln_dogleg.zero       - soln2.zero) < 1e-14
    test_six   = maximum(abs,soln_dogleg_bmp.zero   - soln2.zero) < 1e-14
    test_seven = maximum(abs,soln_nk.zero           - soln2.zero) < 1e-14
    test_eight = maximum(abs,soln_jfnk.zero         - soln2.zero) < 1e-14

    test_nine     = maximum(abs,soln_newton_sparse.zero       - soln2.zero) < 1e-14
    test_ten      = maximum(abs,soln_lm_sparse.zero           - soln2.zero) < 1e-14
    test_eleven   = maximum(abs,soln_lm_kyf_sparse.zero       - soln2.zero) < 1e-14
    test_twelve   = maximum(abs,soln_lm_ar_sparse.zero        - soln2.zero) < 1e-14
    test_thirteen = maximum(abs,soln_dogleg_sparse.zero       - soln2.zero) < 1e-14
    test_fourteen = maximum(abs,soln_dogleg_bmp_sparse.zero   - soln2.zero) < 1e-14
    test_fifteen  = maximum(abs,soln_nk_sparse.zero           - soln2.zero) < 1e-14

    return test_one, test_two, test_three, test_four, test_five, test_six, test_seven, test_eight, test_nine, test_ten, test_eleven,
           test_twelve, test_thirteen, test_fourteen, test_fifteen

end

test_nlboxsolve_outplace()
test_nlboxsolve_inplace()


# Test MonteCarlo functions

using TrajOptSOCPs
using Test

@testset "Feasible" begin

    # Define the polyhedron
    A = [4 5; -4 5; 0 -1]
    b = [20; 30; 1]

    # Test a point in the feasible region (true)
    @test TrajOptSOCPs.isFeasiblePolyHedron(A, b, [-1; 0.5])

    # Test three points outside the feasible region (false)
    @test !TrajOptSOCPs.isFeasiblePolyHedron(A, b, [-10; 0.5])
    @test !TrajOptSOCPs.isFeasiblePolyHedron(A, b, [-1; 10])
    @test !TrajOptSOCPs.isFeasiblePolyHedron(A, b, [-1; -10])
end

@testset "MonteCarlo" begin

    # Generating a 3D point
    pt1 = TrajOptSOCPs.getPoint([3; 2; -5], [10; 100; 0.5])

    # Checking that the point is in the bounds
    @test (10 > pt1[1] > 3)
    @test (100 > pt1[2] > 2)
    @test (0.5 > pt1[3] > -5)

    # Generating 3 Random 4D Points
    numpts = 3
    pts = TrajOptSOCPs.montecarlo([3; 2; -5; -0.01],
                                  [10; 100; 0.5; 0.01],
                                  numpts)

    @test size(pts, 1) == numpts

    # Check that the points are indeed 4D
    @test size(pts[1], 1) == 4
    @test size(pts[numpts], 1) == 4

    # Generating 2 Random Points in a smaller [0, 1] box:
    pts2 = TrajOptSOCPs.montecarlo([-20; -20], [20; 20], 2, TrajOptSOCPs.inBox)

    @test size(pts2, 1) == 2

    # Check that the points are indeed 2D
    @test size(pts2[1], 1) == 2
    @test size(pts2[2], 1) == 2

end


using DistributedFactorGraphs
using IncrementalInference
using Manifolds
using StaticArrays
using Test

##

@testset "Test SpecialOrthogonal(2) prior" begin

##

Base.convert(::Type{<:Tuple}, M::SpecialOrthogonal{2}) = (:Circular,)
Base.convert(::Type{<:Tuple}, ::IIF.InstanceType{SpecialOrthogonal{2}})  = (:Circular,)

# @defVariable SpecialOrthogonal2 SpecialOrthogonal(2) @MMatrix([1.0 0.0; 0.0 1.0])
@defVariable SpecialOrthogonal2 SpecialOrthogonal(2) [1.0 0.0; 0.0 1.0]

M = getManifold(SpecialOrthogonal2)
@test M == SpecialOrthogonal(2)
pT = getPointType(SpecialOrthogonal2)
# @test pT == MMatrix{2, 2, Float64, 4}
@test pT == Matrix{Float64}
pϵ = getPointIdentity(SpecialOrthogonal2)
@test pϵ == [1.0 0.0; 0.0 1.0]

@test is_point(getManifold(SpecialOrthogonal2), getPointIdentity(SpecialOrthogonal2))

fg = initfg()

v0 = addVariable!(fg, :x0, SpecialOrthogonal2)

mp = ManifoldPrior(SpecialOrthogonal(2), SA[1.0 0.0; 0.0 1.0], MvNormal([0.01]))
p = addFactor!(fg, [:x0], mp)

##

doautoinit!(fg, :x0)

vnd = getVariableSolverData(fg, :x0)
@test all(isapprox.(mean(vnd.val), [1 0; 0 1], atol=0.1))
@test all(is_point.(Ref(M), vnd.val))

##
v1 = addVariable!(fg, :x1, SpecialOrthogonal2)
mf = ManifoldFactor(SpecialOrthogonal(2), MvNormal([pi], [0.01]))
f = addFactor!(fg, [:x0, :x1], mf)

doautoinit!(fg, :x1)

##
# smtasks = Task[]
solveTree!(fg) #; smtasks, verbose=true, recordcliqs=ls(fg))
# hists = fetchCliqHistoryAll!(smtasks);
# SArray 0.763317 seconds (2.36 M allocations: 160.488 MiB, 4.16% gc time)
# Vector 0.786390 seconds (2.41 M allocations: 174.334 MiB, 3.97% gc time)
# Vector 0.858993 seconds (2.42 M allocations: 176.613 MiB, 3.43% gc time) sample not tuple
##

end

@testset "Test SpecialOrthogonal(3) prior" begin

##

Base.convert(::Type{<:Tuple}, M::SpecialOrthogonal{3}) = (:Euclid, :Euclid, :Euclid)
Base.convert(::Type{<:Tuple}, ::IIF.InstanceType{SpecialOrthogonal{3}})  =  (:Euclid, :Euclid, :Euclid)

# @defVariable SO3 SpecialOrthogonal(3) @MMatrix([1.0 0.0; 0.0 1.0])
@defVariable SO3 SpecialOrthogonal(3) [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 1.0]

M = getManifold(SO3)
@test M == SpecialOrthogonal(3)
pT = getPointType(SO3)
# @test pT == MMatrix{2, 2, Float64, 4}
@test pT == Matrix{Float64}
pϵ = getPointIdentity(SO3)
@test pϵ == [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 1.0]

@test is_point(getManifold(SO3), getPointIdentity(SO3))

fg = initfg()

v0 = addVariable!(fg, :x0, SO3)

mp = ManifoldPrior(SpecialOrthogonal(3), SA[1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 1.0], MvNormal([0.01, 0.01, 0.01]))
p = addFactor!(fg, [:x0], mp)

doautoinit!(fg, :x0)

vnd = getVariableSolverData(fg, :x0)
@test all(isapprox.( mean(SpecialOrthogonal(3),vnd.val), [1 0 0; 0 1 0; 0 0 1], atol=0.01))
@test all(is_point.(Ref(M), vnd.val))

points = sampleFactor(fg, :x0f1, 100)
IIF.calcCovarianceBasic(SpecialOrthogonal(3), points)

##
v1 = addVariable!(fg, :x1, SO3)
mf = ManifoldFactor(SpecialOrthogonal(3), MvNormal([0.01,0.01,0.01], [0.01,0.01,0.01]))
f = addFactor!(fg, [:x0, :x1], mf)

doautoinit!(fg, :x1)

vnd = getVariableSolverData(fg, :x1)
@test all(isapprox.( mean(SpecialOrthogonal(3),vnd.val), [0.9999 -0.00995 0.01005; 0.01005 0.9999 -0.00995; -0.00995 0.01005 0.9999], atol=0.01))
@test all(is_point.(Ref(M), vnd.val))

# smtasks = Task[]
solveTree!(fg) # ; smtasks, verbose=true, recordcliqs=ls(fg))

# test them again after solve
vnd = getVariableSolverData(fg, :x0)
@test all(isapprox.( mean(SpecialOrthogonal(3),vnd.val), [1 0 0; 0 1 0; 0 0 1], atol=0.01))
@test all(is_point.(Ref(M), vnd.val))

vnd = getVariableSolverData(fg, :x1)
@test all(isapprox.( mean(SpecialOrthogonal(3),vnd.val), [0.9999 -0.00995 0.01005; 0.01005 0.9999 -0.00995; -0.00995 0.01005 0.9999], atol=0.01))
@test all(is_point.(Ref(M), vnd.val))

##
end


import CLBLAS
import OpenCL.cl

const clblas = CLBLAS

const COUNT = int(ARGS[1])

function seq_mat_mul_sdot{T}(Mdim::Int, Ndim::Int, Pdim::Int,
                             A::Array{T}, B::Array{T}, C::Array{T})
    for i in 1:Ndim
        for j in 1:Mdim
            tmp = zero(Float32)
            for k in 1:Pdim
                @inbounds tmp += A[(i-1)*Ndim+k] * B[(k-1)*Pdim+j]
            end
            @inbounds C[(i-1)*Ndim+j] = tmp
        end
    end
end

A = rand(Float32, 10, 10)
B = rand(Float32, 10, 10)
C = rand(Float32, 10, 10)
# Force compilation
seq_mat_mul_sdot(10, 10, 10, A, B, C)

clblas.setup(true)
device, ctx, queue = clblas.get_next_compute_context()

matSizes = Array(Tuple, COUNT)
for i in 1:COUNT
   matSizes[i] = (rand(20:100) * 32, rand(20:100) * 32, rand(20:100) * 32)
end

const alpha = 5.0
const beta  = 2.0

cpu_times        =  zeros(Float32, COUNT)
cpu_seq_times    =  zeros(Float32, COUNT)
gpu_times        =  zeros(Float32, COUNT)
gpu_times_block  =  zeros(Float32, COUNT)
gpu_times_clblas =  zeros(Float32, COUNT)

kernel_source = open(readall, "GEMM-NVIDIA.cl")
prg  = cl.Program(ctx, source=kernel_source) |> cl.build!
mmul = cl.Kernel(prg, "mmul")

kernel_source = open(readall, "block_form.cl")
prg2  = cl.Program(ctx, source=kernel_source) |> cl.build!
block_mmul = cl.Kernel(prg2, "mmul")

blocksize = 16
localmem1 = cl.LocalMem(Float32, blocksize^2)
localmem2 = cl.LocalMem(Float32, blocksize^2)

for id in 1:COUNT
   m, k, n = matSizes[id]
   A = rand(Float32, m, k)
   B = rand(Float32, k, n)
   C = rand(Float32, m, n)

   info("-----CPU------")
   startT = time()
     C_RES = alpha*A*B + beta*C
   runtime = time() - startT
   cpu_times[id] = runtime

   info("------ CPU SEQ ---------")
   startT = time()
   #seq_mat_mul_sdot(n, m, k, A, B, C)
   runtime = time() - startT
   cpu_seq_times[id] = runtime

   info("-----GPU------")
   A_buf = cl.Buffer(Float32, ctx, (:r, :copy), hostbuf=A)
   B_buf = cl.Buffer(Float32, ctx, (:r, :copy), hostbuf=B)
   C_buf = cl.Buffer(Float32, ctx, (:rw, :copy), hostbuf=C)

   evt = mmul[queue, (m, n), (blocksize, blocksize)](convert(Csize_t, m), convert(Csize_t, k),
                                                     alpha, beta, C_buf, A_buf, B_buf, localmem1, localmem2)
   runtime = evt[:profile_duration] / 1e9
   gpu_times[id] = runtime

   info("-----GPU BLOCK------")
   evt = cl.call(queue, block_mmul, (m, n), (16,16),
                 int32(m), int32(k),
                 C_buf, A_buf, B_buf, localmem1, localmem2)
   runtime = evt[:profile_duration] / 1e9
   gpu_times_block[id] = runtime

   info("-----GPU CLBLAS------")
   m = convert(Csize_t, m)
   k = convert(Csize_t, k)
   n = convert(Csize_t, n)

   lda = m
   ldb = k
   ldc = m

   event = cl.UserEvent(ctx)
   ptr_event = [pointer(event)]

   startT = time()
   clblas.clblasSgemm(1, 0, 0, m, n, k,
                           alpha, pointer(A_buf), 0, lda,
                           pointer(B_buf), 0, ldb, beta,
                           pointer(C_buf), 0, ldc,
                           1, [pointer(queue)],
                           0, C_NULL, pointer(ptr_event))

   cl.finish(queue)
   endT = time()
#   startT = cl.CL_ulong[0]
#   cl.api.clGetEventProfilingInfo( event.id, cl.CL_PROFILING_COMMAND_START, sizeof(cl.CL_long), startT, C_NULL)
#   println("Start:", startT[1])

#   endT = cl.CL_ulong[0]
#   cl.api.clGetEventProfilingInfo(event.id, cl.CL_PROFILING_COMMAND_END, sizeof(cl.CL_long), endT, C_NULL)
#   println("End:", endT[1])

#   runtime = (endT[1] - startT[1] ) / 1e9
   gpu_times_clblas[id] = endT - startT
end

using PyCall
@pyimport matplotlib.pyplot as plt
#plt.plot([1:COUNT], cpu_seq_times, color="blue", linewidth=2.0, linestyle="--", label="CPU")
plt.plot([1:COUNT], cpu_times, color="red", linewidth=2.0, linestyle="--", label="CPU")
plt.plot([1:COUNT], gpu_times, color="orange", linewidth=2.0, linestyle="--", label="GPU")
plt.plot([1:COUNT], gpu_times_block, color="green", linewidth=2.0, linestyle="--", label="GPU Block")
plt.plot([1:COUNT], gpu_times_clblas, color="black", linewidth=2.0, linestyle="--", label="GPU BLAS")

plt.ylabel("time in seconds")
plt.title("CPU vs GPU times")
plt.show()