module ThreeOmegaMethod

export Layer, HeaterGeometry, Source, ThreeOmegaMatrix

abstract type LayerInformation end
struct Layer{T1} <: LayerInformation where {T1<:Float64}
    ky::T1; kxy::T1; d::T1; ρC::T1
end

abstract type HeaterInformation end
struct HeaterGeometry{T1} <: HeaterInformation where {T1<:Float64}
    b::T1; l::T1; ρh::Array{T1}
end

abstract type HeaterSource end
struct Source{T1} <: HeaterSource where {T1<:Float64}
    p::T1; f::Array{T1}
end

using QuadGK

function ThreeOmegaMatrix(layers::Array{T1,N1}, heater::T2, source::T3, thresistances::Array{T4,N4}, int_limit::T5=1.0e6) where {T1<:Layer, N1, T2<:HeaterGeometry, T3<:Source, T4<:Number, N4, T5<:Number}
    numLayers::Int64 = size(layers,2)
    # Check input
    numLayers > 1 || throw(DimensionMismatch("There should be at least one layer below the heater."))
    length(thresistances) == numLayers-1 || throw(DimensionMismatch("The number of interface-resistances should be lower than the number of layers: length(thresistances) == size(layers,2)-1."))
    # Power per unit length
    p_l::Float64 = source.p / heater.l
    # Reverse order of input for bottom-up calculation
    layers = reverse(layers, dims=2)
    thresistances = reverse(vec(thresistances), dims=1)
    # Lambda parameters
    λ = [lambda(layers[i].ky, layers[i].kxy, layers[i].ky / layers[i].ρC) for i = 1 : numLayers]
    # Construct B parameter for the system of layers
    B = concatenateB(λ, layers, thresistances, numLayers)
    # Construct integration term. The half term comes from the heater at the top of the stack
    F::Function = (η,ω) -> ( 1 ./ λ[end](η,ω) .* (B(η,ω)[1] + B(η,ω)[2]) ./ (0.5*B(η,ω)[2] - 0.5*B(η,ω)[1]) .* (sin.(heater.b*η)./(heater.b*η)).^2 )
    # return integration
    return integrateTemperature(source.f, int_limit, F, heater.ρh, p_l/π/2, p_l/2/heater.b, im*4*π .* source.f)
end # ThreeOmegaMatrix

"""Lambda parameter."""
lambda(ky::T1, kxy::T1, α::T1) where {T1<:Float64} = (η,ω) -> (ky*sqrt.(kxy*η^2 + im*2*ω/α))

"""Interface matrix."""
Λ(λ1::T0, λ2::T0, r::T1) where {T0<:Function, T1<:Float64} = (η,ω) -> (0.5 ./ λ2(η,ω) .* [λ2(η,ω)+λ1(η,ω)+λ2(η,ω).*λ1(η,ω).*r λ2(η,ω)-λ1(η,ω)-λ2(η,ω).*λ1(η,ω).*r; λ2(η,ω)-λ1(η,ω)+λ2(η,ω).*λ1(η,ω).*r λ2(η,ω)+λ1(η,ω)-λ2(η,ω).*λ1(η,ω).*r])

"""Transfer matrix."""
U(λ::T0, ky::T1, d::T1) where {T0<:Function, T1<:Float64} = (η,ω) -> ([exp.(-λ(η,ω)/ky*d) 0.; 0. exp.(λ(η,ω)/ky*d)])

"""Integration of the temperature term."""
function integrateTemperature(f::Array{T0,N0}, int_limit::T0, F::T1, ρh::Array{T0,N2}, plπ::T0, plb::T0, h_param::Array{T2,N0}) where {T0<:Float64, N0, T1<:Function, T2<:ComplexF64, N2}
    ΔTh = Array{ComplexF64,1}(undef, length(f))
    int_error = Array{Float64,1}(undef, length(f))
    for i in eachindex(f)
        temp0::Tuple{ComplexF64,Float64} = quadgk((η) -> F(η, 2*π*f[i]), 0, int_limit, rtol=sqrt.(eps()))
        temp1::ComplexF64 = temp0[1]
        int_error[i] = temp0[2]
        ΔTh[i] = ( plπ*temp1 + ρh[1]*plb ) ./ ( 1 + ρh[2]*h_param[i] * (ρh[1] + temp1/plb)  )
    end
    return ΔTh, int_error
end # function integrateTemperature()

"""Concatenate B terms depending on the number of layers. The boundary condition for the substrate is semi-infinite."""
function concatenateB(λ::Array{T0,N0}, layers::Array{T1,N1}, thresistances::Array{T2,N2}, numLayers::T3) where {T0<:Function, N0, T1<:Layer, N1, T2<:Number, N2, T3<:Int64}
    B = Λ(λ[end-1], λ[end], thresistances[end])
    for i = (numLayers-1) : -1 : 2
        tempU::Function = U(λ[i], layers[i].ky, layers[i].d)
        tempΛ::Function = Λ(λ[i-1], λ[i], thresistances[i])
        B = let B = B
            (η,ω) -> B(η,ω) * tempU(η,ω) * tempΛ(η,ω)
        end # B = let B = B
    end # i = numLayers-1 : -1 : 2
    # Add semi-infinite boundary condition
    B = let B = B
        (η,ω) -> B(η,ω) * [0.; 1.]
    end # B = let B = B
    return B
end # function concatenateB()

end # ThreeOmegaMethod


module BenchParticle

using Random
using BenchmarkTools
using HOODESolver

suite = BenchmarkGroup()

function fct(u, p, t)
    s1, c1 = sincos(u[1]/2)
    s2, c2 = sincos(u[2])
    s3, c3 = sincos(u[3])
    return [0, 0, u[6], c1*s2*s3/2, s1*c2*s3, s1*s2*c3]
end

Random.seed!(561909)
u0 = 2rand(BigFloat,6)-ones(BigFloat,6)
epsilon=big"1e-7"
ordmax=9
A = [0 0 0  1 0 0;
     0 0 0  0 1 0;
     0 0 0  0 0 0;
     0 0 0  0 1 0;
     0 0 0 -1 0 0;
     0 0 0  0 0 0]

t_0 = big"0.0"
t_max = big"1.0"
prob = HOODEProblem(fct, u0, (t_0, t_max), missing, A, epsilon)
nb = 10^3
n_tau = 32
suite["solve"] = @benchmarkable solve(prob, nb_t=nb, order=ordmax,
                                      getprecision=false, nb_tau=n_tau, dense=false)

end

BenchParticle.suite


function Black76(CallPutFlag, F, X, T, r, v)

  d1 = (log(F / X) + (v^2 / 2) * T) / (v * sqrt(T))
  d2 = d1 - v * sqrt(T)

  if CallPutFlag == "c"
          exp(-r * T) * (F * CND(d1) - X * CND(d2))
  elseif CallPutFlag == "p"
          exp(-r * T) * (X * CND(-d2) - F * CND(-d1))
  else
          println(STDERR, "Falsche Eingabe! Abbruch")
  end

end


function get_robust_path_two(start_node, end_node, p, q, c, d, Gamma_u, Gamma_v, origin, destination)
  @assert length(start_node)==length(end_node)
  @assert length(start_node)==length(p)
  @assert length(start_node)==length(q)
  @assert length(start_node)==length(c)
  @assert length(start_node)==length(d)

  no_arc = length(start_node)

  # Step 0
  diag = p.*q + q.*c + q.*d
  idx = sortperm(diag, rev=true)
  # Theta = Array((Float64, Float64),1)
  # Theta = Array{Tuple{Float64, Float64}}(1)
  Theta = [(0.0,0.0)]

  for t=1:no_arc
    ij = idx[t]

    # Step 1
    Theta = [Theta; (0, p[ij]*d[ij]+q[ij]*d[ij])]

    # Step 2
    Ay  = p[ij]*d[ij] + q[ij]*d[ij]
    Bx  = q[ij]*c[ij]
    Cxy = p[ij]*d[ij] + q[ij]*c[ij] + q[ij]*d[ij]
    Dx  = q[ij]*c[ij] + q[ij]*d[ij]
    Ey  = p[ij]*d[ij]

    for s=t+1:no_arc
      kl = idx[s]

      AAy  = p[kl]*d[kl] + q[kl]*d[kl]
      BBx  = q[kl]*c[kl]
      CCxy = p[kl]*d[kl] + q[kl]*c[kl] + q[kl]*d[kl]
      DDx  = q[kl]*c[kl] + q[kl]*d[kl]
      EEy  = p[kl]*d[kl]

      # Pattern 1
      if Ay>AAy && Bx>BBx && Ey <EEy
        Theta = [Theta; (BBx, Ay); (Cxy-EEy, EEy)]
      end

      # Pattern 2
      if Ay>AAy && Bx>BBx && Ey>EEy && Dx>DDx
        Theta = [Theta; (BBx, Ay); (Dx, EEy)]
      end

      # Pattern 3
      if Ay>AAy && Bx>BBx && Dx<DDx && Ey>EEy
        Theta = [Theta; (BBx, Ay); (Dx, CCxy-Dx)]
      end

      # Pattern 4
      if EEy<Ay && Ay<AAy && Bx>BBx && Ey<EEy
        Theta = [Theta; (CCxy-Ay,Ay); (Cxy-EEy, EEy)]
      end

      # Pattern 5
      if EEy<Ay && Ay<AAy && Bx>BBx && Ey>EEy && Dx > DDx
        Theta = [Theta; (CCxy-Ay,Ay); (Dx, EEy)]
      end

      # Pattern 6
      if EEy<Ay && Ay<AAy && Bx>BBx && Ey>EEy && Dx < DDx
        Theta = [Theta; (CCxy-Ay,Ay); (Dx, CCxy-Dx)]
      end

      # Pattern 7
      if Bx>DDx && Ay<EEy
        Theta = [Theta; (DDx, Ay); (Bx, EEy)]
      end

      # Pattern 8
      if Bx<BBx && Ey<EEy
        Theta = [Theta; (BBx, Cxy-BBx); (Cxy-EEy, EEy)]
      end

      # Pattern 9
      if Bx<BBx && Ey>EEy && Dx>DDx
        Theta = [Theta; (BBx, Cxy-BBx); (Dx, EEy)]
      end

      # Pattern 10
      if Bx<BBx && Ey>EEy && Dx<DDx
        Theta = [Theta; (BBx, Cxy-BBx); (Dx, CCxy-Dx)]
      end

      # Pattern 11
      if Ey>AAy && Dx<BBx
        Theta = [Theta; (Dx, AAy); (BBx, Ey)]
      end
    end
  end
  # The set Theta contains all necessary (theta_u, theta_v)

  theta_u = zeros(size(Theta))
  theta_v = zeros(size(Theta))
  z_front = zeros(size(Theta))
  for i=1:length(Theta)
    theta_u[i] = Theta[i][1]
    theta_v[i] = Theta[i][2]
    z_front[i] = Gamma_u*theta_u[i] + Gamma_v*theta_v[i]
  end

  z_idx = sortperm(z_front, rev=true)

  best_obj = Inf
  best_path = []
  best_x = []

  for i=1:length(Theta)
    k = z_idx[i]
    if Gamma_u*theta_u[k] + Gamma_v*theta_v[k] >= best_obj
        break
    end

    rho_mu = zeros(no_arc)
    for ij=1:no_arc
        qc = q[ij]*c[ij]
        pd = p[ij]*d[ij]
        qd = q[ij]*d[ij]

        Condition1 = theta_u[k]>=qc && theta_v[k]>=pd+qd
        Condition2 = theta_u[k]<=qc && theta_v[k]>=pd+qd
        Condition3 = pd<=theta_v[k] && theta_v[k]<=pd+qd && theta_u[k]+theta_v[k]<=pd+qc+qd
        Condition4 = pd<=theta_v[k] && theta_v[k]<=pd+qd && theta_u[k]+theta_v[k]>=pd+qc+qd
        Condition5 = theta_u[k]<=qc+qd && theta_v[k]<=pd
        Condition6 = theta_u[k]>=qc+qd && theta_v[k]<=pd

        if Condition1 || Condition4
          rho_mu[ij] = 0
        elseif Condition2
          rho_mu[ij] = qc-theta_u[k]
        elseif Condition3 || Condition5
          rho_mu[ij] = pd+qc+qd-theta_u[k]-theta_v[k]
        elseif Condition6
          rho_mu[ij] = pd-theta_v[k]
        end
    end

    link_cost = p.*c + rho_mu
    (pp, xx, ll) = get_shortest_path(start_node, end_node, link_cost, origin, destination)
    # current_obj = Gamma_u*theta_u[k] + Gamma_v*theta_v[k] + dot(link_cost, xx)
    current_obj = Gamma_u*theta_u[k] + Gamma_v*theta_v[k] + ll

    if current_obj < best_obj
        best_obj = current_obj
        best_path = pp
        best_x = xx
    end

  end

  return (best_path, best_x, best_obj)

end


# one-column version
function op_H( LF, Vpot, v::Array{Float64,1} )
  # Is this OK?
  Hv = -0.5*op_nabla2( LF, v ) + Vpot .* v
end

# one-column version, sparse
function op_H( LF, ∇2::SparseMatrixCSC{Float64,Int64}, Vpot, v::Array{Float64,1} )
  # Is this OK?
  Hv = -0.5 * ∇2 * v + Vpot .* v
end

# multi-column version, sparse
function op_H( LF, ∇2::SparseMatrixCSC{Float64,Int64}, Vpot, v::Array{Float64,2} )
  Hv = -0.5 * ∇2 * v + Vpot .* v
end

# multicolumn version
function H_psi( LF, Vpot, psi::Array{Float64,2} )
  Ncol = size(psi)[2]
  Hpsi = zeros( Float64, Npoints, Ncol )
  for ic = 1:Ncol
    Hpsi[:,ic] = -0.5*op_nabla2( LF, psi[:,ic] ) + Vpot .* psi[:,ic]
  end
  return H_psi
end


using EventSimulation
using Distributions

mutable struct Customer
    tin::Float64
    tout::Float64
end

mutable struct Server
    queue::SimQueue{Customer}
    customer::Union{Customer, Nothing}
end

mutable struct State <: AbstractState
    servers::Vector{Server}
    left_customers::Vector{Customer}
    shortest::Bool
    ad::Exponential{Float64}
    sd::Exponential{Float64}
end

function arrival(s)
    if s.state.shortest
        qlen = [length(s.queue) for s in s.state.servers]
        idleserver = isnothing.(s.customer for s in s.state.servers)
        i = argmin(qlen .- idleserver)
        q = s.state.servers[i].queue
    else
        q = rand(s.state.servers).queue
    end
    provide!(s, q, Customer(s.now, NaN))
end

function request_service(server)
    function start_service(s, customer)
        @assert isnothing(server.customer)
        server.customer = customer
        register!(s, finish_service, rand(s.state.sd))
    end

    function finish_service(s)
        server.customer.tout = s.now
        push!(s.state.left_customers, server.customer)
        server.customer = nothing
        request!(s, server.queue, request_service(server))
    end
    start_service
end

function runsim(ar::Number, sr::Number, n::Integer, shortest::Bool)
    servers = [Server(SimQueue{Customer}(), nothing) for i in 1:n]
    ss = State(servers, Customer[],
               shortest, Exponential(1/ar), Exponential(1/sr))
    s = Scheduler(ss)
    repeat_register!(s, arrival, x -> rand(ss.ad))
    for server in ss.servers
        request!(s, server.queue, request_service(server))
    end
    go!(s, 1_000_000)
    mean(c.tout - c.tin for c in s.state.left_customers)
end

@show runsim(1, 0.4, 3, true)
@show runsim(1, 0.4, 3, false)


@testset "linear combination" begin

    lc = LinearCombination([1, -1])
    @test lc isa Transform
    @test cardinality(lc) == ManyToOne()

    @testset "default reduces over columns" begin
        M = [1 1; 2 2; 3 5]
        lc = LinearCombination([1, -1, 1])
        @test FeatureTransforms.apply(M, lc) == [2, 4]
        @test lc(M) == [2, 4]
    end

    @testset "dims = :" begin
        M = [1 1; 2 2; 3 5]
        lc = LinearCombination([1, -1, 1])
        @test_deprecated FeatureTransforms.apply(M, lc; dims=:)
    end

    @testset "dimension mismatch" begin
        M = [1 1 1; 2 2 2]
        lc = LinearCombination([1, -1, 1])  # there are only 2 rows
        @test_throws DimensionMismatch FeatureTransforms.apply(M, lc)
    end

    @testset "N-dim Array" begin
        A = reshape(1:27, 3, 3, 3)
        lc = LinearCombination([1, -1, 1])
        @test FeatureTransforms.apply(A, lc) == [2 11 20; 5 14 23; 8 17 26]
        @test lc(A) == [2 11 20; 5 14 23; 8 17 26]
    end
end


using BaseTestNext
using GraphLayout
srand(1)

@testset "GraphLayout.jl" begin

@testset "Render a pentagon" begin
    adj_matrix = ones(5,5) - eye(5,5)
    @testset "layout_spring_adj" begin
        loc_x, loc_y = layout_spring_adj(adj_matrix)

        draw_layout_adj(adj_matrix, loc_x, loc_y, filename="pentagon_spring.svg")

        draw_layout_adj(adj_matrix, loc_x, loc_y, filename="pentagon_labeled.svg",
        	labels=collect(1:5), labelsize=2.0)
        draw_layout_adj(adj_matrix, loc_x, loc_y, filename="pentagon_noarrows.svg",
        	arrowlengthfrac=0.0)
        draw_layout_adj(adj_matrix, loc_x, loc_y, filename="pentagon_longarrows.svg",
        	arrowlengthfrac=0.5)
    end  # "layout_spring_adj"

    @testset "layout_stressmajorize_adj" begin
        X = layout_stressmajorize_adj(adj_matrix)
        draw_layout_adj(adj_matrix, X[:,1], X[:,2], filename="pentagon_stress.svg")
    end  # "layout_stressmajorize_adj"
end  # "Render a pentagon"

@testset "Render a random graph" begin

    adj_matrix = full(sprand(100,100,0.02))
    @testset "layout_spring_adj" begin
        loc_x, loc_y = layout_spring_adj(adj_matrix)

        draw_layout_adj(adj_matrix, loc_x, loc_y, filename="random_spring.svg")

        draw_layout_adj(adj_matrix, loc_x, loc_y, filename="random_spring_color.svg",
            labelc="#000000",
            labelsize=4.0,
            nodefillc="#21AAFF",
            nodestrokec="#7BB1B1",
            edgestrokec="#B11B1B")
    end

    @testset "layout_stressmajorize_adj" begin
        X = layout_stressmajorize_adj(adj_matrix)
        draw_layout_adj(adj_matrix, X[:,1], X[:,2], filename="random_stress.svg")
    end
end  # "Render a random graph"

# Trees
@testset "Render a tree" begin
    adj_list = Vector{Int}[
        [2,3,4],
        [4],
        [4],
        []
    ]
    labels = ["Apple", "Banana", "Carrot", "Durian"]

    layout_tree(adj_list,labels,filename="tree_1.svg",
                cycles = false, ordering = :optimal, coord = :optimal,
                xsep = 5, ysep = 10, scale = 0.2, labelpad = 1.2,
                background = nothing)

    layout_tree(adj_list,labels,filename="tree_2.svg",
                cycles = false, ordering = :optimal, coord = :optimal,
                xsep = 1, ysep = 20, scale = 0.1, labelpad = 2,
                background = "#DDDDFF")

    rm("tree_1.svg")
    rm("tree_2.svg")
end

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

@testset "Check that output == cached data" begin

#Compare with cached output
cachedout = joinpath(Pkg.dir("GraphLayout"), "test", "examples")
differentfiles = AbstractString[]
for output in readdir(".")
    endswith(output, ".svg") || continue
    contains(output, "tree") && continue  # don't test trees right now
    cached = open(readall, joinpath(cachedout, output))
    genned = open(readall, joinpath(output))
    if cached != genned
        push!(differentfiles, output)
    else #Delete generated file
        rm(output)
    end
end

#Print out which files differ and their diffs
if length(differentfiles)>0
    #Capture diffs
    diffs = map(
        output -> output * ":\n" *
            readall(ignorestatus(`diff $(joinpath(cachedout, output)) $(joinpath(output))`)) *
            "\n\n",
        differentfiles)
    error(string("Generated output differs from cached test output:\n",
        join(differentfiles, "\n"), "\n\n", join(diffs, "\n")))
else
    println("All files matched!")
end

end  # "Check that output == cached data"

end  # "GraphLayout.jl"


# Copyright 2017, Iain Dunning, Joey Huchette, Miles Lubin, and contributors    #src
# This Source Code Form is subject to the terms of the Mozilla Public License   #src
# v.2.0. If a copy of the MPL was not distributed with this file, You can       #src
# obtain one at https://mozilla.org/MPL/2.0/.                                   #src

# # The cannery problem

# A JuMP implementation of the cannery problem from:
#
# Dantzig, Linear Programming and Extensions, Princeton University Press,
# Princeton, NJ, 1963.
#
# It was originally contributed by Louis Luangkesorn, January 30, 2015.

using JuMP
import GLPK
import Test  #src

function example_cannery()
    ## Origin plants.
    plants = ["Seattle", "San-Diego"]
    num_plants = length(plants)
    ## Destination markets.
    markets = ["New-York", "Chicago", "Topeka"]
    num_markets = length(markets)
    ## Capacity and demand in cases.
    capacity = [350, 600]
    demand = [300, 300, 300]
    ## Distance in thousand miles.
    distance = [2.5 1.7 1.8; 2.5 1.8 1.4]
    ## Cost per case per thousand miles.
    freight = 90
    cannery = Model()
    set_optimizer(cannery, GLPK.Optimizer)
    ## Create decision variables.
    @variable(cannery, ship[1:num_plants, 1:num_markets] >= 0)
    ## Ship no more than plant capacity
    @constraint(
        cannery, capacity_con[i = 1:num_plants], sum(ship[i, :]) <= capacity[i]
    )
    ## Ship at least market demand
    @constraint(
        cannery, demand_con[j = 1:num_markets], sum(ship[:, j]) >= demand[j]
    )
    ## Minimize transporatation cost
    @objective(
        cannery,
        Min,
        sum(
            distance[i, j] * freight * ship[i, j]
            for i = 1:num_plants, j = 1:num_markets
        )
    )
    optimize!(cannery)
    println("RESULTS:")
    for i = 1:num_plants
        for j = 1:num_markets
            println("  $(plants[i]) $(markets[j]) = $(value(ship[i, j]))")
        end
    end
    Test.@test termination_status(cannery) == MOI.OPTIMAL    #src
    Test.@test primal_status(cannery) == MOI.FEASIBLE_POINT  #src
    Test.@test objective_value(cannery) == 151_200.0         #src
    return
end

example_cannery()


using RobotDynamics
using StaticArrays
using LinearAlgebra
using BenchmarkTools
using Test

## Setup
include("random_linear.jl")

N = 11
n,m = 15,10
dt = 0.1
times = range(0,step=dt,length=N)
t = times[3]
Ad,Bd = gencontrollable(n,m)
@test isstabled(Ad)
x = @SVector rand(n)
u = @SVector rand(m)
z = KnotPoint(x,u,dt,t)

## Discrete Time Invariant
struct DLTI{TA,TB} <: DiscreteLTI
    A::TA
    B::TB
end
RobotDynamics.state_dim(::DLTI) = 15
RobotDynamics.control_dim(::DLTI) = 10

model = DLTI(SMatrix{n,n}(Ad),SMatrix{n,m}(Bd))
@test RobotDynamics.is_affine(model) == Val(false)
@test RobotDynamics.is_time_varying(model) == false

@test_throws ErrorException discrete_dynamics(DiscreteSystemQuadrature, model, x, u, t, dt)
RobotDynamics.get_A(model::DLTI) = model.A
RobotDynamics.get_B(model::DLTI) = model.B
x2 = discrete_dynamics(DiscreteSystemQuadrature, model, x, u, t, dt)
@test x2 ≈ Ad*x + Bd*u
@test (@ballocated discrete_dynamics(DiscreteSystemQuadrature, $model, $x, $u, $t, $dt) samples=2 evals=2) == 0

F = RobotDynamics.DynamicsJacobian(model)
discrete_jacobian!(DiscreteSystemQuadrature, F, model, z)
@test F ≈ [Ad Bd]

## DLTV
struct DLTV{TA,TB} <: DiscreteLTV
    A::Vector{TA}
    B::Vector{TB}
    times::Vector{Float64}
end
RobotDynamics.state_dim(::DLTV) = 15
RobotDynamics.control_dim(::DLTV) = 10
RobotDynamics.get_A(model::DLTV, k::Int) = model.A[k]
RobotDynamics.get_B(model::DLTV, k::Int) = model.B[k]

N = 11
systems = [gencontrollable(n,m) for k = 1:N]
A = [SMatrix{n,n}(s[1]) for s in systems]
B = [SMatrix{n,m}(s[2]) for s in systems]

model = DLTV(A,B,collect(times))

@test RobotDynamics.is_affine(model) == Val(false)
@test RobotDynamics.is_time_varying(model) == true

@test_throws ErrorException discrete_dynamics(DiscreteSystemQuadrature, model, x, u, t, dt)
RobotDynamics.get_times(model::DLTV) = model.times
x2 = discrete_dynamics(DiscreteSystemQuadrature, model, x, u, t, dt)
@test x2 ≈ A[3]*x + B[3]*u
@test !(x2 ≈ A[1]*x + B[1]*u)
@test (@ballocated discrete_dynamics(DiscreteSystemQuadrature, $model, $x, $u, $t, $dt)) == 0

F = RobotDynamics.DynamicsJacobian(model)
discrete_jacobian!(DiscreteSystemQuadrature, F, model, z)
@test F ≈ [A[3] B[3]]
@test !(F ≈ [A[1] B[1]])

## CLTI
struct CLTI{TA,TB} <: ContinuousLTI
    A::TA
    B::TB
end
RobotDynamics.state_dim(::CLTI) = 15
RobotDynamics.control_dim(::CLTI) = 10

model = CLTI(SMatrix{n,n}(Ad),SMatrix{n,m}(Bd))
@test_throws ErrorException dynamics(model, x, u, t)
RobotDynamics.get_A(model::CLTI) = model.A
RobotDynamics.get_B(model::CLTI) = model.B
xdot = dynamics(model, x, u, t)
@test xdot ≈ Ad*x + Bd*u
@test (@ballocated dynamics($model, $x, $u, $t)) == 0

jacobian!(F, model, z)
@test F ≈ [Ad Bd]

x2 = discrete_dynamics(RK3, model, z)
begin
    local k1_ = (Ad*x + Bd*u)*dt
    local k2_ = (Ad*(x + k1_/2) + Bd*u)*dt
    local k3_ = (Ad*(x - k1_ + 2k2_) + Bd*u)*dt
    @test x2 ≈ x + (k1_+4k2_+k3_)/6
end

discrete_jacobian!(RK3, F, model, z)
@test F.A ≈ I(n) + (Ad*dt + 4*Ad*(I + Ad*dt/2)*dt + Ad*(I - Ad*dt + 2*Ad*(I + Ad*dt/2)*dt)*dt)/6
@test F.B ≈ (Bd*dt + 4*(Ad'Bd*dt/2 + Bd)*dt + (Ad'*(-Bd*dt + 2*(Ad'Bd*dt/2 + Bd)*dt) + Bd)*dt)/6
@test (@ballocated discrete_jacobian!(RK3, $F, $model, $z)) == 0

## CLTV
struct CLTV{TA,TB} <: ContinuousLTV
    A::Vector{TA}
    B::Vector{TB}
    times::Vector{Float64}
end
RobotDynamics.state_dim(::CLTV) = 15
RobotDynamics.control_dim(::CLTV) = 10
RobotDynamics.get_A(model::CLTV, k::Int) = model.A[k]
RobotDynamics.get_B(model::CLTV, k::Int) = model.B[k]

model = CLTV(A,B,collect(times))
@test is_time_varying(model)
@test is_affine(model) == Val(false)

@test_throws ErrorException dynamics(model, x, u, t)
RobotDynamics.get_times(model::CLTV) = model.times
xdot = dynamics(model, x, u, t)
@test xdot ≈ dynamics(model, z)
@test xdot ≈ A[3]*x + B[3]*u

jacobian!(F, model, z)
@test F ≈ [A[3] B[3]]

x2 = discrete_dynamics(RK2, model, z)
A_,B_ = A[3], B[3]
begin
    local k1_ = (A_*x + B_*u)*dt
    local k2_ = (A_*(x + k1_/2) + B_*u)*dt
    @test x2 ≈ x + k2_
end

discrete_jacobian!(RK2, F, model, z)
@test F.A ≈ I(n) + (A_'*(I + A_*dt/2))*dt
@test F.B ≈ (A_'B_*dt/2 + B_)*dt

## Test an affine system
struct DLTI2{TA,TB,TD} <: DiscreteLTI
    A::TA
    B::TB
    d::TD
end
RobotDynamics.state_dim(::DLTI2) = 15
RobotDynamics.control_dim(::DLTI2) = 10
RobotDynamics.get_A(model::DLTI2) = model.A
RobotDynamics.get_B(model::DLTI2) = model.B
RobotDynamics.get_d(model::DLTI2) = model.d
RobotDynamics.is_affine(::DLTI2) = Val(true)

d = @SVector rand(n)
model = DLTI2(SMatrix{n,n}(Ad), SMatrix{n,m}(Bd), d)
x2 = discrete_dynamics(DiscreteSystemQuadrature, model, z)
@test x2 ≈ Ad*x + Bd*u + d
@test (@ballocated discrete_dynamics(DiscreteSystemQuadrature, $model, $z)) == 0

F = RobotDynamics.DynamicsJacobian(model)
discrete_jacobian!(DiscreteSystemQuadrature, F, model, z)
@test F ≈ [Ad Bd]


################################################################################
##
##  This file contains boxcounting methods to analyse phase space volume
##  ratios in 2D billiard maps and 3D billiard flows.
##
################################################################################
export boundarymap_portion, phasespace_portion

#######################################################################################
## Boundary Map Portion
#######################################################################################
"""
    boundarymap_portion(bd::Billiard, t, p::AbstractParticle, δξ, δφ = δξ)
Calculate the portion of the boundary map of the billiard `bd` covered by the
particle `p` when it is evolved for time `t` (float or integer). Notice that the

The boundary map is partitioned into boxes of size `(δξ, δφ)` and as the particle
evolves visited boxes are counted. The returned ratio is this count divided
by the total boxes of size `(δξ, δφ)` needed to cover the boundary map.

**Important:** This portion **does not** equate the portion the particle's orbit covers
on the full, three dimensional phase space. Use the function
[`phasespace_portion`](@ref) for that!
"""
function boundarymap_portion(bd::Billiard{T}, t,
                             par::AbstractParticle{T}, δξ, δφ = δξ;
                             intervals = arcintervals(bd)) where {T}

    p = copy(par)

    count = zero(T)
    t_to_write = zero(T)

    d = Dict{SV{Int}, Int}()

    while count < t
        i, tmin = bounce!(p,bd)
        t_to_write += tmin

        if typeof(bd[i]) <: PeriodicWall
            continue # do not write output if collision with with PeriodicWall
        else
            ξ, sφ = to_bcoords(p, bd[i])
            ξ += intervals[i]

            ind = SV{Int}(floor(Int, ξ/δξ), floor(Int, (sφ + 1)/δφ))
            d[ind] = get(d, ind, 0) + 1

            count += increment_counter(t, t_to_write)
            t_to_write = zero(T)
        end
    end #time or collision number loop

    total_boxes = ceil(Int, totallength(bd)/δξ) * ceil(Int, 2/δφ)
    ratio = length(keys(d))/total_boxes

    return ratio, d
end

#######################################################################################
## Phase Space Portion
#######################################################################################
"""
    phasespace_portion(bd::Billiard, t, p::AbstractParticle, δξ, δφ = δξ)
Calculate the portion of the phase space of the billiard `bd` covered by the
particle `p` when it is evolved for time `t` (float or integer).

This function extends [`boundarymap_portion`](@ref) using a novel approach. For
each visited box of the boundary map, [`bounce!`](@ref) attributes a third dimension
(the collision time, equal to collision distance) which expands the two dimensions
of the boundary map to the three dimensions of the phase space.

The true phase space portion is then the weighted portion of boxes visited by the
particle, divided by the total weighted sum of boxes. The weights of the boxes are
the collision times.
"""
function phasespace_portion(bd::Billiard{T}, t,
                            par::AbstractParticle{T}, δξ, δφ = δξ) where {T}

    r, dict = boundarymap_portion(bd, t, par, δξ, δφ,
                                  intervals = arcintervals(bd))

    ints =  arcintervals(bd)
    maxξ = ceil(Int, totallength(bd)/δξ)
    maxφ = ceil(Int, 2/δφ)

    dummy = copy(par)
    total = zero(T); visited = zero(T)

    for ξcell ∈ 1:maxξ-1, φcell ∈ 1:maxφ-1
        #get center position
        ξc = (ξcell - 0.5)*δξ
        φc = (φcell - 0.5)*δφ - 1

        pos, vel, i = from_bcoords(ξc, φc, bd, ints)

        dummy.pos = pos
        dummy.vel = vel
        _, τ = next_collision(dummy,bd)
        total += τ
        (haskey(dict, SV{Int64}(ξcell, φcell))) && (visited += τ)
    end
    return visited/total
end


# You finally decode the Elves' message. HI, the message says. You continue
# searching for the sleigh keys.
#
# Ahead of you is what appears to be a large ocean trench. Could the keys have
# fallen into it? You'd better send a probe to investigate.
#
# The probe launcher on your submarine can fire the probe with any integer
# velocity in the x (forward) and y (upward, or downward if negative)
# directions. For example, an initial x,y velocity like 0,10 would fire the
# probe straight up, while an initial velocity like 10,-1 would fire the probe
# forward at a slight downward angle.
#
# The probe's x,y position starts at 0,0. Then, it will follow some trajectory
# by moving in steps. On each step, these changes occur in the following order:
#
# - The probe's x position increases by its x velocity.
# - The probe's y position increases by its y velocity.
# - Due to drag, the probe's x velocity changes by 1 toward the value 0; that
#   is, it decreases by 1 if it is greater than 0, increases by 1 if it is less
#   than 0, or does not change if it is already 0.
# - Due to gravity, the probe's y velocity decreases by 1.
#
# For the probe to successfully make it into the trench, the probe must be on
# some trajectory that causes it to be within a target area after any step. The
# submarine computer has already calculated this target area (your puzzle
# input). For example:
#
# target area: x=20..30, y=-10..-5
#
# This target area means that you need to find initial x,y velocity values such
# that after any step, the probe's x position is at least 20 and at most 30,
# and the probe's y position is at least -10 and at most -5.
#
# Given this target area, one initial velocity that causes the probe to be
# within the target area after any step is 7,2:
#
# .............#....#............
# .......#..............#........
# ...............................
# S........................#.....
# ...............................
# ...............................
# ...........................#...
# ...............................
# ....................TTTTTTTTTTT
# ....................TTTTTTTTTTT
# ....................TTTTTTTT#TT
# ....................TTTTTTTTTTT
# ....................TTTTTTTTTTT
# ....................TTTTTTTTTTT
#
# In this diagram, S is the probe's initial position, 0,0. The x coordinate
# increases to the right, and the y coordinate increases upward. In the bottom
# right, positions that are within the target area are shown as T. After each
# step (until the target area is reached), the position of the probe is marked
# with #. (The bottom-right # is both a position the probe reaches and a
# position in the target area.)
#
# Another initial velocity that causes the probe to be within the target area
# after any step is 6,3:
#
# ...............#..#............
# ...........#........#..........
# ...............................
# ......#..............#.........
# ...............................
# ...............................
# S....................#.........
# ...............................
# ...............................
# ...............................
# .....................#.........
# ....................TTTTTTTTTTT
# ....................TTTTTTTTTTT
# ....................TTTTTTTTTTT
# ....................TTTTTTTTTTT
# ....................T#TTTTTTTTT
# ....................TTTTTTTTTTT
#
# Another one is 9,0:
#
# S........#.....................
# .................#.............
# ...............................
# ........................#......
# ...............................
# ....................TTTTTTTTTTT
# ....................TTTTTTTTTT#
# ....................TTTTTTTTTTT
# ....................TTTTTTTTTTT
# ....................TTTTTTTTTTT
# ....................TTTTTTTTTTT
#
# One initial velocity that doesn't cause the probe to be within the target
# area after any step is 17,-4:
#
# S..............................................................
# ...............................................................
# ...............................................................
# ...............................................................
# .................#.............................................
# ....................TTTTTTTTTTT................................
# ....................TTTTTTTTTTT................................
# ....................TTTTTTTTTTT................................
# ....................TTTTTTTTTTT................................
# ....................TTTTTTTTTTT..#.............................
# ....................TTTTTTTTTTT................................
# ...............................................................
# ...............................................................
# ...............................................................
# ...............................................................
# ................................................#..............
# ...............................................................
# ...............................................................
# ...............................................................
# ...............................................................
# ...............................................................
# ...............................................................
# ..............................................................#
#
# The probe appears to pass through the target area, but is never within it
# after any step. Instead, it continues down and to the right - only the first
# few steps are shown.
#
# If you're going to fire a highly scientific probe out of a super cool probe
# launcher, you might as well do it with style. How high can you make the probe
# go while still reaching the target area?
#
# In the above example, using an initial velocity of 6,9 is the best you can
# do, causing the probe to reach a maximum y position of 45. (Any higher
# initial y velocity causes the probe to overshoot the target area entirely.)
#
# Find the initial velocity that causes the probe to reach the highest y
# position and still eventually be within the target area after any step. What
# is the highest y position it reaches on this trajectory?

# read input
(_, _, targety1, _) = parse.(
  Int,
  match(
    r"x=(-?\d+)..(-?\d+), y=(-?\d+)..(-?\d+)",
    readline("input.txt")
  )
)

# If our probe starts at a y velocity of Y, it will reach its maximum height
# and then fall back to 0 with a velocity of -Y because, at each step, the
# velocity will be Y, Y-1, Y-2, ..., 0 (at its maximum height), -1, -2, -3,
# ..., -Y (at a y position of 0), and so on. Assuming the minimum target y is
# negative, the velocity that produces the maximum height will be
# abs(min(targety))-1. Then we can just find the maximum height as the sum of
# (Y, Y-1, Y-2, ..., 0).
inity = abs(targety1) - 1
maxheight = inity * (inity + 1) / 2
println("Maximum Height: ", maxheight)


"""
    Flat()

The *flat distribution* is the improper distribution of real numbers that has the improper
probability density function

```math
f(x) = 1.
```
"""
struct Flat <: ContinuousUnivariateDistribution end

Base.minimum(::Flat) = -Inf
Base.maximum(::Flat) = Inf

Base.rand(rng::Random.AbstractRNG, d::Flat) = rand(rng)
Distributions.logpdf(::Flat, x::Real) = zero(x)

# TODO: only implement `logpdf(d, ::Real)` if support for Distributions < 0.24 is dropped
Distributions.pdf(d::Flat, x::Real) = exp(logpdf(d, x))

# For vec support
Distributions.logpdf(::Flat, x::AbstractVector{<:Real}) = zero(x)
Distributions.loglikelihood(::Flat, x::AbstractVector{<:Real}) = zero(eltype(x))

"""
    FlatPos(l::Real)

The *positive flat distribution* with real-valued parameter `l` is the improper distribution
of real numbers that has the improper probability density function

```math
f(x) = \\begin{cases}
0 & \\text{if } x \\leq l, \\\\
1 & \\text{otherwise}.
\\end{cases}
```
"""
struct FlatPos{T<:Real} <: ContinuousUnivariateDistribution
    l::T
end

Base.minimum(d::FlatPos) = d.l
Base.maximum(d::FlatPos) = Inf

Base.rand(rng::Random.AbstractRNG, d::FlatPos) = rand(rng) + d.l
function Distributions.logpdf(d::FlatPos, x::Real)
    z = float(zero(x))
    return x <= d.l ? oftype(z, -Inf) : z
end

# TODO: only implement `logpdf(d, ::Real)` if support for Distributions < 0.24 is dropped
Distributions.pdf(d::FlatPos, x::Real) = exp(logpdf(d, x))

# For vec support
function Distributions.loglikelihood(d::FlatPos, x::AbstractVector{<:Real})
    lower = d.l
    T = float(eltype(x))
    return any(xi <= lower for xi in x) ? T(-Inf) : zero(T)
end

"""
    BinomialLogit(n, logitp)

The *Binomial distribution* with logit parameterization characterizes the number of
successes in a sequence of independent trials.

It has two parameters: `n`, the number of trials, and `logitp`, the logit of the probability
of success in an individual trial, with the distribution

```math
P(X = k) = {n \\choose k}{(\\text{logistic}(logitp))}^k (1 - \\text{logistic}(logitp))^{n-k}, \\quad \\text{ for } k = 0,1,2, \\ldots, n.
```

See also: [`Binomial`](@ref)
"""
struct BinomialLogit{T<:Real,S<:Real} <: DiscreteUnivariateDistribution
    n::Int
    logitp::T
    logconstant::S

    function BinomialLogit{T}(n::Int, logitp::T) where T
        n >= 0 || error("parameter `n` has to be non-negative")
        logconstant = - (log1p(n) + n * StatsFuns.log1pexp(logitp))
        return new{T,typeof(logconstant)}(n, logitp, logconstant)
    end
end

BinomialLogit(n::Int, logitp::Real) = BinomialLogit{typeof(logitp)}(n, logitp)

Base.minimum(::BinomialLogit) = 0
Base.maximum(d::BinomialLogit) = d.n

# TODO: only implement `logpdf(d, k::Real)` if support for Distributions < 0.24 is dropped
Distributions.pdf(d::BinomialLogit, k::Real) = exp(logpdf(d, k))
Distributions.logpdf(d::BinomialLogit, k::Real) = _logpdf(d, k)
Distributions.logpdf(d::BinomialLogit, k::Integer) = _logpdf(d, k)

function _logpdf(d::BinomialLogit, k::Real)
    n, logitp, logconstant = d.n, d.logitp, d.logconstant
    _insupport = insupport(d, k)
    _k = _insupport ? round(Int, k) : 0
    result = logconstant + _k * logitp - SpecialFunctions.logbeta(n - _k + 1, _k + 1)
    return _insupport ? result : oftype(result, -Inf)
end

function Base.rand(rng::Random.AbstractRNG, d::BinomialLogit)
    return rand(rng, Binomial(d.n, logistic(d.logitp)))
end
Distributions.sampler(d::BinomialLogit) = sampler(Binomial(d.n, logistic(d.logitp)))

"""
    BernoulliLogit(logitp::Real)

Create a univariate logit-parameterised Bernoulli distribution.
"""
BernoulliLogit(logitp::Real) = BinomialLogit(1, logitp)

"""
    OrderedLogistic(η, c::AbstractVector)

The *ordered logistic distribution* with real-valued parameter `η` and cutpoints `c` has the
probability mass function

```math
P(X = k) = \\begin{cases}
    1 - \\text{logistic}(\\eta - c_1) & \\text{if } k = 1, \\\\
    \\text{logistic}(\\eta - c_{k-1}) - \\text{logistic}(\\eta - c_k) & \\text{if } 1 < k < K, \\\\
    \\text{logistic}(\\eta - c_{K-1}) & \\text{if } k = K,
\\end{cases}
```
where `K = length(c) + 1`.
"""
struct OrderedLogistic{T1, T2<:AbstractVector} <: DiscreteUnivariateDistribution
    η::T1
    cutpoints::T2

    function OrderedLogistic{T1,T2}(η::T1, cutpoints::T2) where {T1,T2}
        issorted(cutpoints) || error("cutpoints are not sorted")
        return new{typeof(η), typeof(cutpoints)}(η, cutpoints)
    end
end

function OrderedLogistic(η, cutpoints::AbstractVector)
    return OrderedLogistic{typeof(η),typeof(cutpoints)}(η, cutpoints)
end

Base.minimum(d::OrderedLogistic) = 0
Base.maximum(d::OrderedLogistic) = length(d.cutpoints) + 1

# TODO: only implement `logpdf(d, k::Real)` if support for Distributions < 0.24 is dropped
Distributions.pdf(d::OrderedLogistic, k::Real) = exp(logpdf(d, k))
Distributions.logpdf(d::OrderedLogistic, k::Real) = _logpdf(d, k)
Distributions.logpdf(d::OrderedLogistic, k::Integer) = _logpdf(d, k)

function _logpdf(d::OrderedLogistic, k::Real)
    η, cutpoints = d.η, d.cutpoints
    K = length(cutpoints) + 1

    _insupport = insupport(d, k)
    _k = _insupport ? round(Int, k) : 1
    logp = unsafe_logpdf_ordered_logistic(η, cutpoints, K, _k)

    return _insupport ? logp : oftype(logp, -Inf)
end

function Base.rand(rng::Random.AbstractRNG, d::OrderedLogistic)
    η, cutpoints = d.η, d.cutpoints
    K = length(cutpoints) + 1
    # evaluate probability mass function
    ps = map(1:K) do k
        exp(unsafe_logpdf_ordered_logistic(η, cutpoints, K, k))
    end
    k = rand(rng, Categorical(ps))
    return k
end
function Distributions.sampler(d::OrderedLogistic)
    η, cutpoints = d.η, d.cutpoints
    K = length(cutpoints) + 1
    # evaluate probability mass function
    ps = map(1:K) do k
        exp(unsafe_logpdf_ordered_logistic(η, cutpoints, K, k))
    end
    return sampler(Categorical(ps))
end

# unsafe version without bounds checking
function unsafe_logpdf_ordered_logistic(η, cutpoints, K, k::Int)
    @inbounds begin
        logp = if k == 1
            -StatsFuns.log1pexp(η - cutpoints[k])
        elseif k < K
            tmp = StatsFuns.log1pexp(cutpoints[k-1] - η)
            -tmp + StatsFuns.log1mexp(tmp - StatsFuns.log1pexp(cutpoints[k] - η))
        else
            -StatsFuns.log1pexp(cutpoints[k-1] - η)
        end
    end
    return logp
end

"""
    LogPoisson(logλ)

The *Poisson distribution* with logarithmic parameterization of the rate parameter
descibes the number of independent events occurring within a unit time interval, given the
average rate of occurrence ``exp(logλ)``.

The distribution has the probability mass function

```math
P(X = k) = \\frac{e^{k \\cdot logλ}{k!} e^{-e^{logλ}}, \\quad \\text{ for } k = 0,1,2,\\ldots.
```

See also: [`Poisson`](@ref)
"""
struct LogPoisson{T<:Real,S} <: DiscreteUnivariateDistribution
    logλ::T
    λ::S

    function LogPoisson{T}(logλ::T) where T
        λ = exp(logλ)
        return new{T,typeof(λ)}(logλ, λ)
    end
end

LogPoisson(logλ::Real) = LogPoisson{typeof(logλ)}(logλ)

Base.minimum(d::LogPoisson) = 0
Base.maximum(d::LogPoisson) = Inf

function _logpdf(d::LogPoisson, k::Real)
    _insupport = insupport(d, k)
    _k = _insupport ? round(Int, k) : 0
    logp = _k * d.logλ - d.λ - SpecialFunctions.loggamma(_k + 1)

    return _insupport ? logp : oftype(logp, -Inf)
end

# TODO: only implement `logpdf(d, ::Real)` if support for Distributions < 0.24 is dropped
Distributions.pdf(d::LogPoisson, k::Real) = exp(logpdf(d, k))
Distributions.logpdf(d::LogPoisson, k::Integer) = _logpdf(d, k)
Distributions.logpdf(d::LogPoisson, k::Real) = _logpdf(d, k)

Base.rand(rng::Random.AbstractRNG, d::LogPoisson) = rand(rng, Poisson(d.λ))
Distributions.sampler(d::LogPoisson) = sampler(Poisson(d.λ))

Bijectors.logpdf_with_trans(d::NoDist{<:Univariate}, ::Real, ::Bool) = 0
Bijectors.logpdf_with_trans(d::NoDist{<:Multivariate}, ::AbstractVector{<:Real}, ::Bool) = 0
function Bijectors.logpdf_with_trans(d::NoDist{<:Multivariate}, x::AbstractMatrix{<:Real}, ::Bool)
    return zeros(Int, size(x, 2))
end
Bijectors.logpdf_with_trans(d::NoDist{<:Matrixvariate}, ::AbstractMatrix{<:Real}, ::Bool) = 0


points(S::Jacobi, n) = points(Chebyshev(domain(S)), n)

struct JacobiTransformPlan{T,CPLAN,CJT} <: AbstractTransformPlan{T}
    chebplan::CPLAN
    cjtplan::CJT
end

JacobiTransformPlan(chebplan::CPLAN, cjtplan::CJT) where {CPLAN,CJT} =
    JacobiTransformPlan{eltype(chebplan),CPLAN,CJT}(chebplan, cjtplan)

plan_transform(S::Jacobi, v::AbstractVector) =
    JacobiTransformPlan(plan_transform(Chebyshev(), v), plan_icjt(v, S.a, S.b))
*(P::JacobiTransformPlan, vals::AbstractVector) = P.cjtplan*(P.chebplan*vals)

struct JacobiITransformPlan{T,CPLAN,CJT} <: AbstractTransformPlan{T}
    ichebplan::CPLAN
    icjtplan::CJT
end

JacobiITransformPlan(chebplan::CPLAN, cjtplan::CJT) where {CPLAN,CJT} =
    JacobiITransformPlan{eltype(chebplan),CPLAN,CJT}(chebplan, cjtplan)

plan_itransform(S::Jacobi, v::AbstractVector) =
    JacobiITransformPlan(plan_itransform(Chebyshev(), v), plan_cjt(v, S.a, S.b))
*(P::JacobiITransformPlan, cfs::AbstractVector) = P.ichebplan*(P.icjtplan*cfs)

function coefficients(f::AbstractVector,a::Jacobi,b::Chebyshev)
    if domain(a) == domain(b) && (!isapproxinteger(a.a-0.5) || !isapproxinteger(a.b-0.5))
        cjt(f,a.a,a.b)
    else
        defaultcoefficients(f,a,b)
    end
end
function coefficients(f::AbstractVector,a::Chebyshev,b::Jacobi)
    isempty(f) && return f
    if domain(a) == domain(b) && (!isapproxinteger(b.a-0.5) || !isapproxinteger(b.b-0.5))
        icjt(f,b.a,b.b)
    else
        defaultcoefficients(f,a,b)
    end
end

function coefficients(f::AbstractVector,a::Jacobi,b::Jacobi)
    if domain(a) == domain(b) && (!isapproxinteger(a.a-b.a) || !isapproxinteger(a.b-b.b))
        jjt(f,a.a,a.b,b.a,b.b)
    else
        defaultcoefficients(f,a,b)
    end
end


function findallmax(A)
    max = A[firstindex(A)]
    max_num = []
    for i in (firstindex(A):lastindex(A))
        if A[i] > max
            max = A[i]
            max_num = []
        end
        if A[i] == max
            append!(max_num, i)
        end
    end
    return max_num
end

print(findallmax([3,4,9,2,7,0,1,9,2,3,9,4,5,6]))


using FactCheck

import CDLasso

##############################################
#
#  Groups Lasso
#
##############################################

facts("Group Active set") do

  numGroups = 3
  groupToIndex = Array(typeof(1:2), numGroups)
  for i=1:numGroups
    groupToIndex[i] = UnitRange((i-1)*2+1, i*2)
  end

  as = CDLasso.GroupLassoData(zeros(6), groupToIndex)
  @fact as.numActive --> 0

  x = [1., 0., 0., 0., 1., 2.]
  as = CDLasso.GroupLassoData(x, groupToIndex)
  @fact as.numActive --> 2

  @fact CDLasso._group_norm(x, as, 3) --> roughly(sqrt(5))
  @fact CDLasso._group_norm(x, as, 2) --> roughly(0.)

  CDLasso._fill_zero!(x, as, 1)
  @fact x[1:2] --> zeros(2)

  A = eye(6)
  b = [1., 0., 0., 0., 1., 2.]
  x = zeros(6)
  as = CDLasso.GroupLassoData(x, groupToIndex)
  λ = [0.9, 0., 3]

  CDLasso._add_violator!(as, x, A, b, λ)
  @fact as.numActive --> 1
  @fact as.groups[1] --> 1

  A = eye(6)
  b = [1., 0., 0., 0., 1., 2.]
  x = zeros(6)
  as = CDLasso.GroupLassoData(x, groupToIndex)
  λ = [1.9, 0., 1.]

  CDLasso._add_violator!(as, x, A, b, λ)
  @fact as.numActive --> 1
  @fact as.groups[1] --> 3
end

facts("Minimize one group") do
  numGroups = 1
  groupToIndex = Array(typeof(1:2), numGroups)
  for i=1:numGroups
    groupToIndex[i] = UnitRange((i-1)*3+1, i*3)
  end

  b = [1., 1., 1.]
  A = eye(3)
  Asvd = svdfact(A)
  λ = 1.
  x = zeros(3)

  as = CDLasso.GroupLassoData(x, groupToIndex)
  CDLasso._add_violator!(as, x, A, b, fill(λ, 3))
  CDLasso._min_one_group!(x, Asvd, b, as, 1, λ)

  @fact A*x+b+(λ/norm(x))*x --> roughly(zeros(3); atol=1e-5)

  ##

  A = randn(100, 3)
  A = A'A/100
  Asvd = svdfact(A)
  b = [1., 1., 1.]
  λ = 1.
  x = zeros(3)

  as = CDLasso.GroupLassoData(x, groupToIndex)
  CDLasso._add_violator!(as, x, A, b, fill(λ, 3))
  CDLasso._min_one_group!(x, Asvd, b, as, 1, λ)

  @fact A*x+b+(λ/norm(x))*x --> roughly(zeros(3); atol=1e-5)

  ##
  numGroups = 1
  groupToIndex = Array(typeof(1:2), numGroups)
  for i=1:numGroups
    groupToIndex[i] = UnitRange((i-1)*10+1, i*10)
  end

  A = randn(100, 10)
  A = A'A/100
  Asvd = svdfact(A)
  b = randn(10)
  λ = 0.5
  x = zeros(10)

  as = CDLasso.GroupLassoData(x, groupToIndex)
  CDLasso._add_violator!(as, x, A, b, fill(λ, 10))
  CDLasso._min_one_group!(x, Asvd, b, as, 1, λ)

  @fact A*x+b+(λ/norm(x))*x --> roughly(zeros(10); atol=1e-5)

end

facts("Minimize active set") do
  numGroups = 2
  groupToIndex = Array(typeof(1:2), numGroups)
  for i=1:numGroups
    groupToIndex[i] = UnitRange((i-1)*3+1, i*3)
  end

  b = ones(6)
  A = eye(6)
  λ = [1., 0.5]

  x = ones(6)
  as = CDLasso.GroupLassoData(x, groupToIndex)
  CDLasso._add_violator!(as, x, A, b, λ)

  CDLasso._group_lasso!(x, A, b, λ, as)

  for i=1:numGroups
    @fact A[groupToIndex[i], :]*x + b[groupToIndex[i]] + (λ[i]/norm(x[groupToIndex[i]]))*x[groupToIndex[i]]  --> roughly(zeros(length(groupToIndex[i])); atol=1e-5)
  end

  ##

  numGroups = 2
  groupToIndex = Array(typeof(1:2), numGroups)
  for i=1:numGroups
    groupToIndex[i] = UnitRange((i-1)*3+1, i*3)
  end

  b = randn(6)
  A = randn(100, 6)
  A = A'A/100
  λ = [1., 0.5]

  x = ones(6)
  as = CDLasso.GroupLassoData(x, groupToIndex)
  CDLasso._add_violator!(as, x, A, b, λ)

  CDLasso._group_lasso!(x, A, b, λ, as)

  for i=1:numGroups
    @fact A[groupToIndex[i], :]*x + b[groupToIndex[i]] + (λ[i]/norm(x[groupToIndex[i]]))*x[groupToIndex[i]]  --> roughly(zeros(length(groupToIndex[i])); atol=1e-5)
  end

  ##

  numGroups = 2
  groupToIndex = Array(typeof(1:2), numGroups)
  for i=1:numGroups
    groupToIndex[i] = UnitRange((i-1)*3+1, i*3)
  end

  b = randn(6)
  A = randn(100, 6)
  A = A'A/100
  λ = [100., 0.1]

  x = ones(6)
  as = CDLasso.GroupLassoData(x, groupToIndex)
  CDLasso._add_violator!(as, x, A, b, λ)

  CDLasso._group_lasso!(x, A, b, λ, as)

  i = 1
  @fact norm(A[groupToIndex[i], :]*x + b[groupToIndex[i]]) <= λ[i]  --> true
  i = 2
  @fact A[groupToIndex[i], :]*x + b[groupToIndex[i]] + (λ[i]/norm(x[groupToIndex[i]]))*x[groupToIndex[i]]  --> roughly(zeros(length(groupToIndex[i])); atol=1e-4)
end

facts("Group Lasso") do
  numGroups = 200
  groupToIndex = Array(typeof(1:2), numGroups)
  for i=1:numGroups
    groupToIndex[i] = UnitRange((i-1)*5+1, i*5)
  end

  n = 400
  p = 1000
  X = randn(n, p)
  Y = X[:,1:50] * ones(50) + 0.1 * randn(n)
  XX = X' * X / n
  Xyn = -X' * Y / n

  λ = 0.4 .* ones(numGroups)
  x = zeros(p)
  CDLasso.group_lasso!(x, XX, Xyn, λ, groupToIndex)

  for i=1:200
    if norm(x[groupToIndex[i]]) < 1e-6
      @fact norm(XX[groupToIndex[i], :]*x + Xyn[groupToIndex[i]]) <= λ[i] + 1e-4  --> true
    else
      @fact XX[groupToIndex[i], :]*x + Xyn[groupToIndex[i]] + (λ[i]/norm(x[groupToIndex[i]]))*x[groupToIndex[i]] --> roughly(zeros(length(groupToIndex[i])); atol=1e-4)
    end
  end

  λ = 0.01 .* ones(numGroups)
  x = zeros(p)
  CDLasso.group_lasso!(x, XX, Xyn, λ, groupToIndex)

  for i=1:200
    if norm(x[groupToIndex[i]]) < 1e-6
      @fact norm(XX[groupToIndex[i], :]*x + Xyn[groupToIndex[i]]) <= λ[i] + 1e-4  --> true
    else
      @fact XX[groupToIndex[i], :]*x + Xyn[groupToIndex[i]] + (λ[i]/norm(x[groupToIndex[i]]))*x[groupToIndex[i]] --> roughly(zeros(length(groupToIndex[i])); atol=1e-4)
    end
  end

  λ = 40 .* ones(numGroups)
  x = zeros(p)
  CDLasso.group_lasso!(x, XX, Xyn, λ, groupToIndex)

  for i=1:200
    if norm(x[groupToIndex[i]]) < 1e-6
      @fact norm(XX[groupToIndex[i], :]*x + Xyn[groupToIndex[i]]) <= λ[i] + 1e-4  --> true
    else
      @fact XX[groupToIndex[i], :]*x + Xyn[groupToIndex[i]] + (λ[i]/norm(x[groupToIndex[i]]))*x[groupToIndex[i]] --> roughly(zeros(length(groupToIndex[i])); atol=1e-4)
    end
  end

end

##############################################
#
#  Lasso
#
##############################################

facts("Orthogonal design") do
  b = [1.0, 2.0, 3.0]
  A = eye(3)

  λ = 1.5
  x = zeros(3)
  CDLasso.lasso!(x, A, b, λ)
  @fact x --> roughly([0, -0.5, -1.5])

  λ = [0.9, 2.1, 3.1]
  x = zeros(3)
  CDLasso.lasso!(x, A, b, λ)
  @fact x --> roughly([-0.1, 0, 0])
end

facts("Lasso shrink to zero") do
  n = 100
  p = 10

  X = randn(n, p)
  Y = X * ones(p) + 0.1 * randn(n)
  A = X' * X / n
  b = -X' * Y / n

  beta = ones(p)
  λ = maximum(abs(b)) + 0.1
  CDLasso.lasso!(beta, A, b, λ)
  @fact beta --> roughly(zeros(p))
end

##############################################
#
#  Other
#
##############################################

facts("SoftThreshold") do

  @fact CDLasso.shrink(5., 3.) --> 2.
  @fact CDLasso.shrink(-2., 3.) --> 0.
  @fact CDLasso.shrink(-1.,1.) --> 0.
  @fact CDLasso.shrink(-1.,0.9) --> roughly(-0.1, 1e-5)

end

FactCheck.exitstatus()


export Flip

"""
    Flip(p)

Constructs a very simple *sfunc* corresponding to a Bernoulli distribution,
represented by a `Cat`.  The output is `true` with probability `p`, and `false`
with probability `1-p`.

See also: [`Cat`](@ref)
"""
Flip(p) = Cat([false, true], [1-p, p])


# # Vertical Collision Avoidance (VCAS)
#
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/models/VertCAS.ipynb)
#
# This example [^JK] [^ABKL] considers commercial aircraft, which are required to
# operate with a collision avoidance system that gives vertical climbrate advisories
# to pilots to prevent near midair collisions (NMACs). An NMAC occurs when the
# aircraft are separated by less than 100 ft vertically and 500 ft horizontally.

module VertCAS  #jl

using NeuralNetworkAnalysis, LinearAlgebra

# ## Model
#
# This benchmark is a closed-loop variant of aircraft collision avoidance
# system ACAS X. The scenario involves two aircraft, the ownship and the
# intruder, where the ownship is equipped with a collision avoidance system
# referred to as VerticalCAS [3]. VerticalCAS once every second issues vertical
# climbrate advisories to the ownship pilot to avoid a near mid-air collision
# (NMAC), a region where the ownship and intruder are separated by less than
# 100ft vertically and 500ft horizontally. The ownship (black) is assumed
# to have a constant horizontal speed, and the intruder (red) is assumed to
# follow a constant horizontal trajectory towards ownship, see the figure below.
#
# ![](https://raw.githubusercontent.com/JuliaReach/NeuralNetworkAnalysis.jl/master/models/VertCAS/nmac.png)
#
# The current geometry of the system is described by:
#
# 1) Intruder altitude relative to ownship $h(ft)$
# 2) Ownship vertical climbrate $\dot h_0 (ft/min)$
# 3) The seconds until the ownship (black) and intruder (red) are no longer horizontally separated $τ(s)$
#
# We can, therefore, assume that the intruder is static and the horizontal
# separation $\tau$ decreases by one each second.
# There are 9 advisories and each of them instructs the pilot to accelerate
# until the vertical climbrate of the ownship complies with the advisory:
#
# 1) COC: Clear of Conflict
# 2) DNC: Do Not Climb
# 3) DND: Do Not Descend
# 4) DES1500: Descend at least 1500 ft/min
# 5) CL1500: Climb at least 1500 ft/min
# 6) SDES1500: Strengthen Descent to at least 1500 ft/min
# 7) SCL1500: Strengthen Climb to at least 1500 ft/min
# 8) SDES2500: Strengthen Descent to at least 2500 ft/min
# 9) SCL2500: Strengthen Climb to at least 2500 ft/min
#
# In addition to the parameters describing the geometry of the encounter, the
# current state of the system stores the advisory $adv$ issued to the ownship
# at the previous time step. VerticalCAS is implemented as nine ReLU networks
# $N_i$, one for each (previous) advisory, with three inputs
# $(h,\dot{h}_0,\tau)$, five fully-connected hidden layers of 20 units each,
# and nine outputs representing the score of each possible advisory. Therefore,
# given a current state $(h,\dot{h}_0,\tau,\text{adv})$, the new advisory
# $adv$ is obtained by computing the argmax of the output of $N_{\text{adv}}$
# on $(h,\dot{h}_0,\tau)$.
# Given the new advisory, if the current climbrate does not comply with it, the
# pilot can choose acceleration $\ddot{h}_0$ from the given set:
#
# 1) COC: $\{-\frac{g}{8}, 0, \frac{g}{8}\}$
# 2) DNC: $\{-\frac{g}{3}, -\frac{7g}{24}, -\frac{g}{4}\}$
# 3) DND: $\{\frac{g}{4}, \frac{7g}{24}, \frac{g}{3}\}$
# 4) DES1500: $\{-\frac{g}{3}, -\frac{7g}{24}, -\frac{g}{4}\}$
# 5) CL1500: $\{\frac{g}{4}, \frac{7g}{24}, \frac{g}{3}\}$
# 6) SDES1500: $\{-\frac{g}{3}\}$
# 7) SCL1500: $\{\frac{g}{3}\}$
# 8) SDES2500: $\{-\frac{g}{3}\}$
# 9) SCL2500: $\{\frac{g}{3}\}$
#
# where $g$ represents the gravitational constant $32.2 \ \text{ft/s}^2$.
# If the new advisory is COC(1), then it can be any acceleration from the set ${−g/8, 0, g/8}$.
# For all remaining advisories, if the previous advisory coincides with the new one and the
# current climb rate complies with the new advisory (e.g., $\dot{h}_0$ is non-positive for DNC and
# $\dot{h}_0 ≥ 1500$ for CL1500) the acceleration `\ddot{h}_0$ is $0$.
#
# Given the current system state $(h,\dot{h}_0,\tau,\text{adv})$, the new
# advisory $\text{adv}'$ and the acceleration $\ddot{h}_0$, the new state
# of the system $(h(t+1),\dot{h}_0(t+1),\tau(t+1),\text{adv}(t+1))$ can
# be computed as follows:
#
# ```math
# \begin{aligned}
# h(t+1) &=& h - \dot{h}_0 \Delta\tau - 0.5 \ddot{h}_0 \Delta\tau^2 \\
# \dot{h}_0(t+1) &=& \dot{h}_0 + \ddot{h}_0 \Delta\tau \\
# \tau(t+1) &=& \tau - \Delta\tau \\
# \text{adv}(t+1) &=& \text{adv}'
# \end{aligned}
# ```
# where $\Delta\tau=1$.

const g = 32.2  # gravitational constant

## accelerations (middle)
const ACC_MIDDLE = Dict(:COC => 0.0, :DNC => -7g/24, :DND => 7g/24,
                        :DES1500 => -7g/24, :CL1500 => 7g/24, :SDES1500 => -g/3,
                        :SCL1500 => g/3, :SDES2500 => -g/3, :SCL2500 => g/3)

## continuous dynamics matrix (h, hdot0)
const Δτ = 1.0
const A = [1  -Δτ; 0  1.];

# We load the controllers in a dictionary with the keys being the advisories.

CONTROLLERS = Dict{Symbol, Any}()

CTRL_IDX = [:COC, :DNC, :DND, :DES1500,
            :CL1500, :SDES1500, :SCL1500,
            :SDES2500, :SCL2500]

path = @modelpath("VertCAS", "")
for i = 1:9
    file = joinpath(path, "VertCAS_noResp_pra0$(i)_v9_20HU_200.nnet")
    adv = CTRL_IDX[i]
    CONTROLLERS[adv] = read_nnet(file)
end

struct State{T, N}
    state::T # state for (h, hdot0) variables
    τ::N
    adv::Symbol
end

ADVISORIES = Dict{Symbol, Any}()

## for every advisory, return a function to check
## whether the current climbrate complies with the advisory or not
## if the advisory *does not* comply => it is changed according to the ACC dictionary
## otherwise, the new acceleration is zero
using Symbolics
@variables x

ADVISORIES[:COC] = EmptySet(1)
ADVISORIES[:DNC] = HalfSpace(x <= 0)
ADVISORIES[:DND] = HalfSpace(x >= 0)
ADVISORIES[:DES1500] = HalfSpace(x <= -1500)
ADVISORIES[:CL1500] = HalfSpace(x >= 1500)
ADVISORIES[:SDES1500] = HalfSpace(x <= -1500)
ADVISORIES[:SCL1500] = HalfSpace(x >= 1500)
ADVISORIES[:SDES2500] = HalfSpace(x <= -2500)
ADVISORIES[:SCL2500] = HalfSpace(x >= 2500)

## this function receives X = [h, hdot0, τ, adv′] and the
## *previous* advisory adv
function get_acceleration(X::State, adv; ACC=ACC_MIDDLE)

    ## obtain projection on hdot
    hdot = _interval(X.state, 2)

    ## transform units from ft/s to ft/min
    hdot = 60 * hdot

    ## new advisory
    adv′ = X.adv

    ## check whether the current state complies with the advisory
    comply = hdot ⊆ ADVISORIES[adv′]

    if adv == adv′ && comply
        return 0.0
    else
        return ACC[adv′]
    end
end

const normalization_additive = -[0.0, 0, 20]
const normalization_multiplicative = 1.0 ./ [16000.0, 5000, 40]

function normalize(x::AbstractVector)
    y = x .+ normalization_additive
    z = y .* normalization_multiplicative
    return z
end

function normalize(X::AbstractZonotope)
    Y = minkowski_sum(X, Singleton(normalization_additive))
    Z = linear_map(diagm(normalization_multiplicative), Y)
    return Z
end

## scalar case; alg is ignored
function forward_adv(X::Singleton, τ, adv; alg=nothing)
    v = vcat(element(X), τ)
    v = normalize(v)
    u = forward(CONTROLLERS[adv], v)
    imax = argmax(u)
    return CTRL_IDX[imax]
end

## set-based case
function forward_adv(X::AbstractZonotope, τ, adv; alg=Ai2())
    Y = cartesian_product(X, Singleton([τ]))

    Y = normalize(Y)
    out = forward_network(alg, CONTROLLERS[adv], Y)

    imax = argmax(high(out))
    return CTRL_IDX[imax]
end

function VCAS!(out::Vector{State{T, N}}, KMAX; ACC=ACC_MIDDLE, alg_nn=Ai2()) where {T, N}

    ## unpack initial state
    X0 = first(out)
    S = X0.state
    τ = X0.τ
    adv = X0.adv

    ## get initial acceleration
    hddot = ACC[adv]

    for i in 1:KMAX
        ## compute next state
        b = [-hddot*Δτ^2 / 2, hddot * Δτ]
        S′ = affine_map(A, S, b)
        τ′ = τ - 1
        adv′ = forward_adv(S′, τ′, adv, alg=alg_nn)

        ## store new state
        X′ = State(S′, τ′, adv′)
        push!(out, X′)

        ## get acceleration from network
        ## this logic only works for ACC_MIDDLE
        hddot = get_acceleration(X′, adv; ACC=ACC)

        ## update current state
        S = S′
        τ = τ′
        adv = adv′
    end
    return out
end;

# ## Specifications
#
# For this benchmark the aim is to verify that the ownship avoids entering the
# NMAC zone after $k \in \{1, \dots, 10\}$ time steps, i.e., $h(k) > 100$ or
# $h(k) < -100$, for all possible choices of acceleration by the pilot. The
# set of initial states considered is as follows: $h(0) \in [-133, -129]$,
# $\dot{h}_0(0) \in \{-19.5, -22.5, -25.5, -28.5\}$, $\tau(0) = 25$ and
# $\text{adv}(0) = \text{COC}$.

bad_states = HalfSpace([1.0, 0.0], 100.) ∩ HalfSpace([-1.0, 0.0], 100.)

## property for guaranteed violation
predicate = X -> X ⊆ bad_states
predicate_sol = sol -> any(predicate(R) for F in sol for R in F);

# ## Results

# ### Simulation

const h0 = Interval(-133, -129)
const hdot0 = [-19.5,-22.5, -25.5, -28.5]
const τ0 = 25.0
const adv0 = :COC

function _random_states(k=1, include_vertices::Bool=false, rand_h0::Bool=true)
    N = Float64
    T = Singleton{N, Vector{N}}
    states = Vector{State{T, N}}()
    xs = sample(h0, k, include_vertices=include_vertices)
    for x in xs
        if rand_h0
            ## use a random value for y
            y = hdot0[rand(1:4)]
            S0 = State(Singleton([x[1], y]), τ0, adv0)
            push!(states, S0)
            continue
        end
        ## use all possible values for y
        for i in 1:4
            y = hdot0[i]
            S0 = State(Singleton([x[1], y]), τ0, adv0)
            push!(states, S0)
        end
    end
    return states
end

function _all_states()
    S0 = [convert(Zonotope, concretize(h0 × Singleton([hdot0[i]]))) for i in 1:4]
    return [State(S0i, τ0, adv0) for S0i in S0]
end

function simulate_VCAS(X0::State; KMAX=10)
    out = [X0]
    sizehint!(out, KMAX+1)

    VCAS!(out, KMAX, ACC=ACC_MIDDLE)
    return out
end

## project onto the h variable
function _project(X::Vector{State{T, N}}) where {T<:Singleton, N}
    return [Singleton([Xi.state.element[1], Xi.τ]) for Xi in X]
end

_interval(X::LazySet, i) = overapproximate(Projection(X, (i,)), Interval)

function _project(X::Vector{State{T, N}}) where {T<:Zonotope, N}
    Xint = [_interval(Xi.state, 1) × Singleton([Xi.τ]) for Xi in X]
end

function run(X0)
    ensemble = [simulate_VCAS(X0i) for X0i in X0]
    res = _project.(ensemble)
    return res
end

function check(sol)
    println("property checking")
    res_pred = @timed predicate_sol(sol)
    print_timed(res_pred)
    if res_pred.value
        println("The property is violated.")
    else
        println("The property may be satisfied.")
    end
end;

# Simulation result for random choice of velocity:

X0 = _random_states(10, true, false)  # randomly sampled points (incl. vertices)
println("$(length(X0)) simulations with central advisories")
res = @timed begin
    res1 = run(X0)
    check(res1)
end
println("total analysis time")
print_timed(res);

# Simulation result for all choices of velocity:

println("flowpipe construction (unsound) with central advisories")
res = @timed begin
    res2 = run(_all_states())
    check(res1)
end
println("total analysis time")
print_timed(res);

# Finally we plot the results:

using Plots
import DisplayAs

function plot_helper()
    fig = plot(xlims=(-200, 200), ylims=(14, 26), xlab="h (vertical distance)",
               ylab="τ (time to reach horizontally)")
    plot!(fig, bad_states, alpha=0.2, c=:red, lab="unsafe states")
    return fig
end

fig = plot_helper()
for o in res1
    plot!(fig, o, alpha=1.0)
end
fig = DisplayAs.Text(DisplayAs.PNG(fig))
## savefig("VertCAS-rand.png")
fig

#-

fig = plot_helper()
for (i, c) in [(1, :blue), (2, :orange), (3, :green), (4, :cyan)]
    [plot!(fig, o, lw=2.0, alpha=1., markershape=:none, seriestype=:shape, c=c) for o in res2[i]]
end
fig = DisplayAs.Text(DisplayAs.PNG(fig))
## savefig("VertCAS-sets.png")
fig

# ## References

# [^JK]: Julian, K. D., & Kochenderfer, M. J. (2019). *A reachability method for
#        verifying dynamical systems with deep neural network controllers.*
#        [arXiv preprint arXiv:1903.00520.](https://arxiv.org/pdf/1903.00520.pdf).
#
# [^ABKL]: Akintunde, M. E., Botoeva, E., Kouvaros, P., & Lomuscio, A. (2020).
#          *Formal Verification of Neural Agents in Non-deterministic Environments.*
#          In [Proceedings of the 19th International Conference on Autonomous Agents and
#          Multiagent Systems (pp. 25-33)](http://ifaamas.org/Proceedings/aamas2020/pdfs/p25.pdf).

end  #jl
nothing  #jl


using Test

using MathOptInterface
const MOI = MathOptInterface
const MOIT = MathOptInterface.Test
const MOIU = MathOptInterface.Utilities
const MOIB = MathOptInterface.Bridges

include("../utilities.jl")

mock = MOIU.MockOptimizer(MOIU.Model{Float64}())
config = MOIT.TestConfig()

bridged_mock = MOIB.Variable.SOCtoRSOC{Float64}(mock)

@testset "soc1v" begin
    mock.optimize! = (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(mock, [1/√2 + 1/2, 1/√2 - 1/2, 1/√2],
        (MOI.VectorAffineFunction{Float64}, MOI.Zeros) => [[-√2]])
    MOIT.soc1vtest(bridged_mock, config)

    ceqs = MOI.get(mock, MOI.ListOfConstraintIndices{
        MOI.VectorAffineFunction{Float64}, MOI.Zeros}())
    @test length(ceqs) == 1
    MOI.set(bridged_mock, MOI.ConstraintName(), ceqs[1], "ceq")

    @testset "Test mock model" begin
        var_names = ["a", "b", "c"]
        MOI.set(
            mock, MOI.VariableName(),
            MOI.get(mock, MOI.ListOfVariableIndices()), var_names)
        rsocs = MOI.get(mock, MOI.ListOfConstraintIndices{
            MOI.VectorOfVariables, MOI.RotatedSecondOrderCone}())
        @test length(rsocs) == 1
        MOI.set(mock, MOI.ConstraintName(), rsocs[1], "rsoc")

        invs2 = 1 / √2
        s = """
        variables: a, b, c
        rsoc: [a, b, c] in MathOptInterface.RotatedSecondOrderCone(3)
        ceq: [$invs2*a + $invs2*b + -1.0] in MathOptInterface.Zeros(1)
        maxobjective: $invs2*a + -$invs2*b + c
        """
        model = MOIU.Model{Float64}()
        MOIU.loadfromstring!(model, s)
        MOIU.test_models_equal(mock, model, var_names, ["rsoc", "ceq"])
    end

    @testset "Test bridged model" begin
        var_names = ["x", "y", "z"]
        MOI.set(
            bridged_mock, MOI.VariableName(),
            MOI.get(bridged_mock, MOI.ListOfVariableIndices()), var_names)
        socs = MOI.get(bridged_mock, MOI.ListOfConstraintIndices{
            MOI.VectorOfVariables, MOI.SecondOrderCone}())
        @test length(socs) == 1
        MOI.set(bridged_mock, MOI.ConstraintName(), socs[1], "soc")

        s = """
        variables: x, y, z
        soc: [x, y, z] in MathOptInterface.SecondOrderCone(3)
        ceq: [x + -1.0] in MathOptInterface.Zeros(1)
        maxobjective: y + z
        """
        model = MOIU.Model{Float64}()
        MOIU.loadfromstring!(model, s)
        MOIU.test_models_equal(bridged_mock, model, var_names, ["soc", "ceq"])
    end

    @testset "Delete" begin
        xyz = MOI.get(bridged_mock, MOI.ListOfVariableIndices())

        message = string("Cannot delete variable as it is constrained with other",
                         " variables in a `MOI.VectorOfVariables`.")
        for i in eachindex(xyz)
            err = MOI.DeleteNotAllowed(xyz[i], message)
            @test_throws err MOI.delete(bridged_mock, xyz[i])
        end

        test_delete_bridged_variables(bridged_mock, xyz, MOI.SecondOrderCone, 3, (
            (MOI.VectorOfVariables, MOI.RotatedSecondOrderCone, 0),
        ))
    end
end


using Optimize
using NLPModels

using Plots
# pyplot()
plotly()

n_min = 2
n_max = 200

# Math Prog Base collection
include("MPBProblems.jl")
test_probs = mpb_probs

# Select solvers
using LineSearch
using Stopping
using LSDescentMethods

#stop_norm = ARCTR.stop_norm
atol = 1.0e-5
rtol = 1.0e-10

solvers = [Newton, Newton, Newton, Newton, Newton]
labels = ["Newton Armijo backtracking";"Newton HZ"; "Newton MT"; "Newton TR_SecA"; "Newton TR_Nwt"]

stp = TStopping(atol = atol, rtol = rtol, max_iter = 1000000, max_eval = 1000000, max_time = 100.0)

options = [Dict{Symbol,Any}(:verbose=>false, :stp => stp, :linesearch => Newarmijo_wolfe)
           Dict{Symbol,Any}(:verbose=>false, :stp => stp, :linesearch => _hagerzhang2!)#, :τ₀ => 0.01, :mayterminate => true, :τ₁ => 0.1)
           Dict{Symbol,Any}(:verbose=>false, :stp => stp, :linesearch => _morethuente2!)#, :τ₀ => 0.01, :τ₁ => 0.1)
           Dict{Symbol,Any}(:verbose=>false, :stp => stp, :linesearch => TR_SecA_ls, :τ₀ => 0.000932010000000, :τ₁ => 0.206998804000000)
           Dict{Symbol,Any}(:verbose=>false, :stp => stp, :linesearch => TR_Nwt_ls, :τ₀ => 0.200018755000000, :τ₁ => 0.999996368310000)
           ]

include("compare_solvers.jl")
P5 = compare_solvers_with_options2(solvers, options, labels, test_probs, n_min, n_max, printskip = false)[2]


struct BBK{TGV, TF<:AbstractFloat, TM} <: SecondOrderNonMetropolisSampler
    ∇V!::TGV
    β::TF
    γ::TF
    M::TM
    Δt::TF
    σ::TF
end

"""
    BBK(∇V!, β, γ, M, Δt)

Set up the BBK integrator for inertial Langevin.

### Fields

* ∇V!   - In place gradient of the potential
* β     - Inverse temperature
* γ     - Damping Coefficient
* M     - Mass (either scalar or vector)
* Δt    - Time step
"""
function BBK(∇V!::TGV, β::TF, γ::TF, M::TM, Δt::TF) where {TGV, TF<:AbstractFloat,TM}
    σ = sqrt(γ * Δt / β);
    return BBK(∇V!, β, γ, M, Δt, σ)
end

mutable struct BBKState{Tq, Tx<:AbstractVector{Tq}} <:SecondOrderNonMetropolisSamplerState
    x::Tx
    p_mid::Tq
    ∇V::Tq
end

function InitState!(x₀, sampler::BBK)
    ∇V = similar(x₀[1]);
    sampler.∇V!(∇V, x₀[1]);
    return BBKState(x₀, similar(x₀[1]), copy(∇V));
end

function InitState(x₀, sampler::BBK)

    ∇V = deepcopy(x₀[1]);
    sampler.∇V!(∇V, x₀[1]);
    return BBKState(deepcopy(x₀), similar(x₀[1]), copy(∇V));
end

function UpdateState!(state::BBKState, sampler::BBK)

    @. state.p_mid = state.x[2] - 0.5 * sampler.Δt * state.∇V - 0.5 * sampler.Δt * sampler.γ/sampler.M * state.x[2] + sampler.σ * randn();
    @. state.x[1] = state.x[1] + sampler.Δt/sampler.M * state.p_mid
    sampler.∇V!(state.∇V, state.x[1]);
    @. state.x[2] = (state.p_mid - 0.5 * sampler.Δt * state.∇V + sampler.σ * randn())/(1 + 0.5 * sampler.Δt * sampler.γ/sampler.M);

    state
end


using FountainCodes, Test, Distributions

"""return an LT code object and a vector of source symbols"""
function init(K; Tv=GF256, M=K, δ=1e-9)
    dd = Soliton(K, M, δ)
    LT{Tv}(K, dd)
end

function test_decoding(K::Integer, N::Integer; Tv=GF256)
    rng = MersenneTwister(123)
    code = init(K; Tv)
    src = rand_nonzero(rng, Tv, K)
    G = generator_matrix(code, 1:N)
    enc = G'*src
    dec = decode(G, enc)
    @test dec == src
end
test_decoding(100, 120)


using SpikingNetworks

parseParameters"""
τ = 10
σ=80
Vₜ=0.8
Vᵣ=0
N=100
W=0.1*(ones(N,N)-diagm(ones(N)))"""

parseEquations"""
dv/dt = (Iₑ(t)-v+σ*ξ(t)+W⊙sigmoid(v-0.5))/τ
Iₑ(t)=0.5"""

net=init"""
N=N
v=rand(Vᵣ:0.01:Vₜ,N)"""

spike"v.>Vₜ"
reset"v[i]=Vᵣ"

net.dt=0.1
@time run!(net,100/net.dt)
time,rate=activity(net)
using MatlabPlot
mplot(time,rate)


using LinearAlgebra, CUDA, SparseArrays

export MatFacModel

mutable struct MatFacModel

    X::AbstractMatrix            # KxM "instance factor" matrix
    Y::AbstractMatrix            # KxN "feature factor" matrix

    instance_reg_mats::AbstractVector{AbstractMatrix}  # K x (M x M) (usually sparse)
    feature_reg_mats::AbstractVector{AbstractMatrix}  # K x (N x N) (usually sparse)

    instance_offset::AbstractVector # instance-wise "intercept" terms
    instance_offset_reg::AbstractMatrix # regularizer matrix (usually sparse)

    feature_offset::AbstractVector  # feature-wise "intercept" terms
    feature_offset_reg::AbstractMatrix # regularizer matrix (usually sparse)

    feature_precision::AbstractVector # feature-specific precisions

    instance_covariate_coeff::Union{Nothing,AbstractMatrix} # Regression coefficients
                                                          # for sample covariates
    instance_covariate_coeff_reg::AbstractMatrix

    losses::AbstractVector       # N-dim vector of feature-specific losses

    factor_importances::AbstractVector # K-dim vector of factor importances
end

function MatFacModel(instance_reg_mats::AbstractVector,
                     feature_reg_mats::AbstractVector,
                     losses::AbstractVector;
                     instance_offset_reg::Union{Nothing,AbstractMatrix}=nothing,
                     feature_offset_reg::Union{Nothing,AbstractMatrix}=nothing,
                     instance_covariate_coeff_reg::Union{Nothing,AbstractMatrix}=nothing,
                     K::Union{Nothing,Integer}=nothing)

    M = size(instance_reg_mats[1],1)
    N = size(feature_reg_mats[1],1)

    if K == nothing
        K = max(length(instance_reg_mats),
                length(feature_reg_mats))
    else
        @assert K >= max(length(instance_reg_mats),
                         length(feature_reg_mats))
    end

    X = 0.01 .* randn(K, M) ./ sqrt(K)
    Y = 0.01 .* randn(K, N)

    instance_offset = 0.01 .* randn(M)
    feature_offset = 0.01 .* randn(N)

    if instance_offset_reg == nothing
        instance_offset_reg = spzeros(M,M)
    end
    if feature_offset_reg == nothing
        feature_offset_reg = spzeros(N,N)
    end
    if instance_covariate_coeff_reg == nothing
        instance_covariate_coeff_reg = spzeros(N,N)
    end

    feature_precision = ones(N)

    instance_covariate_coeff = nothing

    factor_importances = zeros(K)

    return MatFacModel(X, Y, instance_reg_mats,
                             feature_reg_mats,
                             instance_offset,
                             instance_offset_reg,
                             feature_offset,
                             feature_offset_reg,
                             feature_precision,
                             instance_covariate_coeff,
                             instance_covariate_coeff_reg,
                             losses,
                             factor_importances
                      )
end


using GeometricFlux
using Flux
using Flux: onehotbatch, onecold, crossentropy, throttle
using JLD2
using StatsBase
using SparseArrays
using LightGraphs.SimpleGraphs

@load "data/cora_features.jld2" features
@load "data/cora_labels.jld2" labels
@load "data/cora_graph.jld2" g

num_nodes = 2708
num_features = 1433
target_catg = 7

## Preprocessing data
train_X = features  # dim: num_features * num_nodes
train_y = labels  # dim: target_catg * num_nodes

## Model
model = Chain(GCNConv(g, num_features=>1000, relu),
              GCNConv(g, 1000=>500, relu),
              GCNConv(g, 500=>7),
              softmax)
# test model
# model(train_X)

## Loss
loss(x, y) = crossentropy(model(x), y)
accuracy(x, y) = mean(onecold(model(x)) .== onecold(y))

## Training
ps = Flux.params(model)
train_data = [(train_X, train_y)]
opt = ADAM(0.01)
evalcb() = @show(accuracy(train_X, train_y))

Flux.train!(loss, ps, train_data, opt, cb=throttle(evalcb, 10))


using JuMP
using Gurobi

polyEval(coeffs, pt) = sum([coeffs[i] * pt^(i-1) for i = 1:length(coeffs)])

function facto(m::Int64)
    factori = 1
    for j = 1:m
        factori = factori * j
    end
    return factori
end

function comb(n::Int64, m::Int64)
    combi = facto(n) / (facto(m) * facto(n-m))
    return convert(Int, combi)
end

function setUpFitting(deg::Int64, c::Float64)

    normCoeffs = Array(Float64, deg+1)
    for i=1:deg + 1
        normCoeffs[i] = comb(deg, i-1) * c^(deg-i+1)
    end

    m = Model(solver=GurobiSolver(OutputFlag=false))

	@variable(m, coeffs[1:deg+1])

	return m, coeffs, normCoeffs

end

function addResid(m, coeffs, ys, demands, arcs, scaling)
    	@variable(m, resid)
	@variable(m, dual_cost)
	@variable(m, primal_cost)

	@constraint(m, dual_cost == sum(demands[(s,t)] * (ys[(s,t), t] - ys[(s,t), s]) for (s,t)=keys(demands)))
	@constraint(m, primal_cost == sum(a.flow * a.freeflowtime * polyEval(coeffs, a.flow/a.capacity) for a=values(arcs)))
	@constraint(m, resid >= (primal_cost - dual_cost) / scaling)
	@constraint(m, resid >= 0)

	return resid
end

function addIncreasingCnsts(m, coeffs, arcs)
	sorted_flows = sort([a.flow / a.capacity for a in values(arcs)])
	for i = 2:length(sorted_flows)
		@constraint(m, polyEval(coeffs, sorted_flows[i-1]) <= polyEval(coeffs, sorted_flows[i]))
	end
end

function normalize(m, coeffs)
    @constraint(m, coeffs[1] == 1)
end

function addNetworkCnsts(m, coeffs, demands, arcs, numNodes)
	@variable(m, ys[keys(demands), 1:numNodes])
	for k = keys(arcs)
		a = arcs[k]
		rhs = a.freeflowtime * polyEval(coeffs, a.flow/a.capacity)
		for od in keys(demands)
			@constraint(m, ys[od, k[2]] - ys[od, k[1]] <= rhs)
		end
	end
	return ys
end

##########
#Fitting Funcs
##########

function train(lam::Float64, deg::Int, c::Float64, demands, arcs; fcoeffs=nothing)
    numNodes = maximum(map(pair->pair[1], keys(arcs)))
    m, coeffs, normCoeffs = setUpFitting(deg, c)

    addIncreasingCnsts(m, coeffs, arcs)  #uses the original obs flows

    normalize(m, coeffs)

    resids = Variable[]

    vArcs = Arc[]
    for arc in values(arcs)
        push!(vArcs, arc)
    end

    flow_data = [a.flow::Float64 for a in vArcs]

    #copy the flow data over to the arcs
    for (ix, a) in enumerate(vArcs)
        a.flow = flow_data[ix]
    end

    #Dual Feasibility
    ys = addNetworkCnsts(m, coeffs, demands, arcs, numNodes)

    #add the residual for this data point
    push!(resids, addResid(m, coeffs, ys, demands, arcs, 1e6))

    if fcoeffs != nothing
        fixCoeffs(m, fcoeffs, coeffs)
    end

    @objective(m, Min, sum(resids[i] for i = 1:length(resids))
                            + lam * sum(coeffs[i] * coeffs[i] / normCoeffs[i] for i=1:deg + 1))
    solve(m)

    return [getvalue(coeffs[i]) for i =1:length(coeffs)], getobjectivevalue(m)
end


"""
    load_gid_mesh(filename) -> mesh
"""
function load_gid_mesh(fn::AbstractString)
    open(fn, "r") do io
        load_gid_mesh(io)
    end
end

"""
    load_gid_mesh(stream) ->mesh
"""
function load_gid_mesh(file)

    lineCount = countlines(file);
    seekstart(file);
    readline(file); # MESH dimension 3 ElemType Triangle Nnode 3
    readline(file); # Coordinates

    # Read the vertices
    vertices = Vector{Pt{3,Float64}}(undef,lineCount)
    vertexCount = 0;
    while !eof(file)
            line = split(readline(file));
            size(line, 1) == 4 || break; # End Coordinates
            vertexCount += 1;
            #vertices[vertexCount] = Pt{3,Float64}(map(Meta.parse, line[2:4]));
            vertices[vertexCount] = Pt{3,Float64}(parse.(Float64,line[2:4]))
    end
    vertices= vertices[1:vertexCount]

    readline(file); # "\n"
    readline(file); # Elements

    # Read the triangles
    triangles = Vector{Pt{3,Int}}(undef,lineCount-vertexCount);
    triangleCount = 0;
    while !eof(file)
            line = split(readline(file));
            size(line, 1) == 4 || break; # End Elements
            triangleCount += 1;
            triangles[triangleCount] = Pt{3,Int}(map(Meta.parse, line[2:4]));
    end
    triangles = triangles[1:triangleCount]
    # Mesh(vertices[1:vertexCount], triangles[1:triangleCount])

    # Place them on a space filling curve
    Q = Pt{3,Float64}
    ctrs = Q[]
    for tr in triangles
            ctr = sum(vertices[tr])/3
            push!(ctrs,ctr)
    end
    sorted = sort_sfc(ctrs)
    triangles = triangles[sorted]

    Mesh(vertices, triangles)
end


function Graphs.weights(g::AbstractDataGraph{T,VL,VD,ED}) where {T,VL,VD,ED<:Nothing}
    return Graphs.DefaultDistance(nv(g))
end

function get_edgeweight(
    g::AbstractDataGraph{T,VL,VD,ED}, s::Integer, d::Integer
) where {T,VL,VD,ED<:Real}
    return get_data(g, s, d)
end

function get_edgeweight(
    g::AbstractDataGraph{T,VL,VD,ED}, s::Integer, d::Integer
) where {T,VL,VD,ED}
    return get_data(g, s, d).weight
end

@doc raw"""
    weights(g)

Compute the edge weights for an [`AbstractDataGraph`](@ref).

- If the edge data is of type `ED = Nothing`, return a lazy matrix full of ones.
- If the edge data is of type `ED <: Real`, fill a sparse matrix with it.
- If the edge data is of another type `ED`, try to use the attribute `edge_data.weight` for each edge.
"""
function Graphs.weights(g::AbstractDataGraph{T}) where {T}
    n = nv(g)
    eds = edges(g)
    I = [src(edge) for edge in eds]
    J = [dst(edge) for edge in eds]
    W = [get_edgeweight(g, src(edge), dst(edge)) for edge in eds]
    if is_directed(g)
        return sparse(I, J, W, n, n)
    else
        return sparse(vcat(I, J), vcat(J, I), vcat(W, W), n, n)
    end
end


@testset "Combined Manifolds" begin
using Random: seed!

M1 = Euclidean(3)
x1 = RnPoint([1.,2.,0.])
y1 = RnPoint([2.,3.,0.])
ξ1 = log(M1,x1,y1)
η1 = RnTVector([1.,0.,0.])

M2 = Sphere(2)
x2 = SnPoint([1.,0.,0.])
y2 = SnPoint([0.,0.,1.])
ξ2 = log(M2,x2,y2)
η2 = SnTVector(1/sqrt(2)*[0., 1.,1.])

#
# Prod
M = Product([M1, M2])
x = ProdPoint([x1,x2])
y = ProdPoint([y1,y2])
ξ = ProdTVector([ξ1,ξ2])
η = ProdTVector([η1, η2])

@test getValue(x) == [x1,x2]
@test getValue(ξ) == [ξ1, ξ2]

@test distance(M,x,y) == sqrt(distance(M1,x1,y1)^2 + distance(M2,x2,y2).^2)
@test dot(M,x,ξ,η) == dot(M1,x1,ξ1,η1) + dot(M2,x2,ξ2,η2)
@test distance( M, exp(M,x,ξ), ProdPoint([exp(M1,x1,ξ1),exp(M2,x2,ξ2)]) ) == 0
@test norm( M, x, log(M,x,y) - ProdTVector([log(M1,x1,y1), log(M2,x2,y2)]) ) ≈ 0
@test manifoldDimension(x) == 5
@test manifoldDimension(M) == 5
@test norm(M,y, parallelTransport(M,x,y,ξ) - ProdTVector([
    parallelTransport(M1,x1,y1,ξ1), parallelTransport(M2,x2,y2,ξ2) ]) ) ≈ 0

@test typeofMPoint(ξ) == ProdPoint{Array{MPoint,1}}
@test typeofTVector(x) == ProdTVector{Array{TVector,1}}
@test typeofMPoint(ProdTVector{Array{TVector,1}}) == ProdPoint{Array{MPoint,1}}
@test typeofTVector(ProdPoint{Array{MPoint,1}}) == ProdTVector{Array{TVector,1}}
@test validateMPoint(M, randomMPoint(M) )
@test validateTVector(M,x, randomTVector(M,x) )

@test_throws DomainError validateMPoint(Product([M1]),x)
@test_throws DomainError validateTVector(M,x,ProdTVector([ξ1]))

@test typicalDistance(M) ≈ sqrt( 2* (typicalDistance(M1).^2 + typicalDistance(M2).^2) )

@test zeroTVector(M,x) == ProdTVector([zeroTVector(M1,x1), zeroTVector(M2,x2)])

@test "$M" == "The Product Manifold of [ $(M1), $(M2) ]"
@test "$x" == "Prod[ $(x1), $(x2) ]"
@test "$ξ" == "ProdT[ $(ξ1), $(ξ2) ]"
#
# Power
M = Power(M1,2)
@test M.powerSize == (2,)
@test Power(M1,(3,2)).powerSize == (3,2)

x = PowPoint([x1,x1])
y = PowPoint([y1,y1])
ξ = PowTVector([ξ1,ξ1])
η = PowTVector([η1,η1])
#
@test size(x) == size(x.value)
@test distance(M, copy(x), x) == 0
@test norm(M,x, ξ - copy(ξ)) == 0
@test ndims(x) == ndims(x.value)
@test getValue(x) == [x1,x1]
@test getValue(ξ) == [ξ1,ξ1]
#
x[2] = RnPoint([1., 2., 0.])
@test getValue(x[2]) == [1., 2., 0.]
ξ[2] = RnTVector([1., 2., 0.])
@test getValue(ξ[2]) == [1., 2., 0.]
@test PowPoint([x1 x1; x1 x1])[1,2] == x1
x[2] = x1
ξ[2] = ξ1
xG = PowPoint([x1 x1;y1 y1])
@test xG[1:2,2].value == [x1,y1]
@test ndims(xG) == 2
ξG = PowTVector([ξ1 ξ1;η1 η1])
@test ξG[1:2,2].value == [ξ1,η1]
@test ndims(ξG) == 2
@test size(repeat(x, outer=[2,1])) == (4,1)
@test size(repeat(x, 2,1)) == (4,1)
@test size(repeat(x, inner=[1,2])) == (2,2)
@test size(repeat(ξ, outer=[2,1])) == (4,1)
@test size(repeat(ξ, 2,1)) == (4,1)
@test size(repeat(ξ, inner=[1,2])) == (2,2)
#
@test getValue.(getValue( adjointJacobiField(M, x, y, 0.5, ξ))) == getValue.(
        adjointJacobiField.(Ref(M1), [x1,x1], [y1,y1], 0.5, [ξ1,ξ1]) )
@test getValue.(getValue( jacobiField(M, x, y, 0.5, ξ))) ==
        getValue.(jacobiField.(Ref(M1), [x1,x1], [y1,y1], 0.5, [ξ1,ξ1]) )
#
@test distance(M,x,y) ≈ sqrt(2)*distance(M1,x1,y1)
@test dot(M,x,ξ,η) ≈ 2*dot(M1,x1,ξ1,η1)
#
@test getValue(exp(M, x, ξ)) == [exp(M1, x1, ξ1), exp(M1, x1, ξ1)]
@test getValue(log(M, x, y)) == [log(M1, x1, y1), log(M1, x1, y1)]
#
@test manifoldDimension(x) == 2*manifoldDimension(x1)
@test manifoldDimension(M) == 2*manifoldDimension(M1)
@test norm(M,x,ξ) ≈ sqrt(2)*norm(M1,x1,ξ1)
@test ndims(x) == 1
@test ndims(ξ) == 1
@test getValue( parallelTransport(M,x,y,ξ) ) == [parallelTransport(M1, x1, y1, ξ1),parallelTransport(M1, x1, y1, ξ1)]
@test validateMPoint(M,randomMPoint(M))
@test validateTVector(M,x,randomTVector(M,x))
@test validateTVector(M,x,randomTVector(M,x,:Gaussian))

@test typeofMPoint(ξ) == PowPoint{RnPoint{Float64},1}
@test typeofTVector(x) == PowTVector{RnTVector{Float64},1}
@test typeofTVector(PowPoint{RnPoint{Float64},1}) == PowTVector{RnTVector{Float64},1}
@test typeofMPoint(PowTVector{RnTVector{Float64},1}) == PowPoint{RnPoint{Float64},1}

@test_throws DomainError validateMPoint(Power(M,3),x)
@test_throws DomainError validateTVector(M,x,PowTVector([ξ1]))

@test typicalDistance(M) ≈ sqrt( 2 ) * typicalDistance(M1)

@test zeroTVector(M,x) == PowTVector([zeroTVector(M1,x1), zeroTVector(M1,x1)])

a0 = RnTVector([0.,0.,0.]); a1 = RnTVector([1.,0.,0.]);
a2 = RnTVector([0.,1.,0.]); a3 = RnTVector([0.,0.,1.])
ΞT = [PowTVector([a1,a0]), PowTVector([a2,a0]), PowTVector([a3,a0]),
    PowTVector([a0,a1]), PowTVector([a0,a2]), PowTVector([a0,a3]) ]
κT = zeros(6)
Ξ,κ = tangentONB(M,x,y)

@test sum( abs.(κ-κT)) == 0
@test sum( norm.(Ref(M), Ref(x),  ΞT .- Ξ ) ) == 0
# show
@test "$x" == "Pow[Rn([1.0, 2.0, 0.0]), Rn([1.0, 2.0, 0.0])]"
@test "$ξ" == "PowT[RnT([1.0, 1.0, 0.0]), RnT([1.0, 1.0, 0.0])]"
@test "$M" == "The Power Manifold of The 3-dimensional Euclidean space of size (2,)."
#
# Tangent Bundle
M = TangentBundle(M1)
@test getBase(M) == M1
x = TBPoint(x1,ξ1)
y = TBPoint(y1,η1)
xT = TBPoint( (x1,ξ1) )
@test x == xT
@test getValue(x) == (x1,ξ1)
xE = MPointE(x)
@test getBase(x) == x1
@test getTangent(x) == ξ1
@test getBase(xE) == MPointE(x1)
@test getTangent(xE) == TVectorE(ξ1,x1)

ξ = TBTVector(ξ1,η1)
ξT = TBTVector( (ξ1,η1) )
@test ξ == ξT
@test getValue(ξ) == (ξ1,η1)
ξE = TVectorE(ξ,x)
@test getBase(ξ) == ξ1
@test getTangent(ξ) == η1
@test getBase(ξE) == TVectorE(ξ1,x1)
@test getTangent(ξE) == TVectorE(η1,x1)

@test 2*ξ == TBTVector(2*ξ1, 2*η1)
@test ξ+ξ == TBTVector(ξ1+ξ1,η1+η1)
@test ξ-ξ == TBTVector(ξ1-ξ1,η1-η1)
@test -ξ == TBTVector(-ξ1,-η1)
@test +ξ == TBTVector(ξ1,η1)

@test distance(M,x,y) ≈ sqrt(3)
@test norm(M,x,ξ) ≈ sqrt(dot(M,x,ξ,ξ))
@test norm(M,x,ξ) ≈ sqrt(norm(M1,x1,ξ1)^2 + norm(M1,x1,η1))
@test exp(M,x,ξ) == TBPoint(exp(M1,x1,ξ1), ξ1+η1)
@test log(M,x,y) ==  TBTVector(log(M1,x1,y1),ξ1-η1)
@test manifoldDimension(M) == 2*manifoldDimension(M1)
@test manifoldDimension(x) == 2*manifoldDimension(x1)
@test parallelTransport(M,x,y,ξ) == ξ
@test validateMPoint(M,randomMPoint(M))
@test validateTVector(M,x,randomTVector(M,x))

@test typeofMPoint(ξ) == TBPoint{RnPoint{Float64},RnTVector{Float64}}
@test typeofTVector(x) == TBTVector{RnTVector{Float64}}

Ξ, κ = tangentONB(M,x,y)
@test κ == zeros(6)
@test Ξ == [TBTVector(a1,a0), TBTVector(a2,a0), TBTVector(a3,a0), TBTVector(a0,a1), TBTVector(a0,a2), TBTVector(a0,a3)]

@test typicalDistance(M) == sqrt( typicalDistance(M1)^2 + manifoldDimension(M1))
@test zeroTVector(M,x) == TBTVector(a0,a0)

@test "$M" == "The Tangent bundle of <$(M.manifold)>"
@test "$x" == "TB(Rn([1.0, 2.0, 0.0]), RnT([1.0, 1.0, 0.0]))"
@test "$ξ" == "TBT(RnT([1.0, 1.0, 0.0]), RnT([1.0, 0.0, 0.0]))"
end


f(a, b) = a + b

@code_lowered f(1, 1)
@code_typed f(1, 1)
@code_llvm f(1, 1)
@code_native f(1, 1)

@code_native f(1.0, 3.0)
@code_native f(1.2, 4)

using BenchmarkTools, LinearAlgebra

arr = rand(Float32, 1_000, 10)
arr = rand(Float32, 100_000, 10)

function f(a)
    b = a * a'
    b[diagind(b)] .= 0
    sum(b, dims=2)
end

f(arr)
@benchmark f(arr)

function g(a::Array{Float32, 2})::Array{Float32}
    """
    # vectorized version of the following operations with M (n, m) => NM (n, n) => n
    out = a * a' => creates a a we cannot store in RAM
    out[diagind(out)] .= 0
    out = sum(out, dims=2)
    """
    N, M = size(a)

    out = Array{Float32, 1}(undef, N)

    Threads.@threads for i in 1:N
        total = zero(Float32)
        @inbounds for ii in 1:N
            if (i == ii) continue end
            @inbounds @simd for j in 1:M
                total += a[i, j] * a[ii, j]
            end
        end
        @inbounds out[i] = total
    end
    return out
end

g(arr)
@benchmark g(arr)

using CUDA
function g_kernel(matrix::CuDeviceMatrix{Float32}, out::CuDeviceVector{Float32}, len::Int, N::Int, M::Int)::Nothing
    index = (blockIdx().x - 1) * blockDim().x + threadIdx().x

    if (index) < len
        for i in 1:N
            if index == i continue end
            for j in 1:M
                @inbounds out[index] += matrix[index, j] * matrix[i, j]
            end
        end
    end
    return nothing
end

function g(matrix::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})::Array{Float32}
    len = length(matrix)
    N, M = size(matrix)
    threads_per_block = 256
    blocks = Int(ceil(N / threads_per_block))

    out = CUDA.zeros(Float32, N)

    @cuda threads=threads_per_block blocks=blocks g_kernel(matrix, out, len, N, M)
    return Array(out)
end

arr = CUDA.rand(Float32, 100_000, 10)
g(arr)

@CUDA.time g(arr)