text:
"""
Module storing return codes and modes
"""
module Return

    """
    Module storing return codes enum

    Usage:
    - `Return.Code.CMD_FINISHED`
    - ...
    """
    module Code
        """
        Return code enum
        """
        @enum ReturnCode::UInt8 begin
            INVALID_CMD          = 0x00
            CMD_FINISHED         = 0x01
            INVALID_COMPONENT_ID = 0x02
            INVALID_PAGE_ID      = 0x03
            INVALID_PICTURE_ID   = 0x04
            INVALID_FONT_ID      = 0x05
            INVALID_BAUD         = 0x11
            INVALID_VARIABLE     = 0x1A
            INVALID_OPERATION    = 0x1B
            INVALID_ASSIGN       = 0x1C
            INVALID_EEPROM       = 0x1D
            INVALID_PARAMETER_QUANTITY = 0x1E
            INVALID_IO                 = 0x1F
            INVALID_ESC_CHAR           = 0x20
            INVALID_VAR_NAME_TOO_LONG  = 0x23
            EVENT_TOUCH_HEAD           = 0x65
            CURRENT_PAGE_ID_HEAD       = 0x66
            EVENT_POSITION_HEAD        = 0x67
            EVENT_SLEEP_POSITION_HEAD  = 0x68
            STRING_HEAD                = 0x70
            NUMBER_HEAD                = 0x71
            EVENT_ENTER_SLEEP_MODE     = 0x86
            EVENT_ENTER_WAKE_UP_MODE   = 0x87
            EVENT_LAUNCHED             = 0x88
            EVENT_UPGRADED             = 0x89
            EVENT_DATA_TR_FINISHED     = 0xFD
            EVENT_DATA_TR_READY        = 0xFE
        end
    end

    function code(val::UInt8)
        Code.ReturnCode(val)
    end

    """
    Module storing return modes enum

    Usage:
    - `Return.Mode.NO_RETURN`
    - ...
    """
    module Mode
        """
        Return mode enum
        """
        @enum ReturnMode::UInt8 begin
            NO_RETURN    = 0x00  # No return
            SUCCESS_ONLY = 0x01  # Only return the successful data
            FAIL_ONLY    = 0x02  # Only return the failed data
            ALWAYS       = 0x03  # Always return
        end
    end
end

"""
Module storing font alignment modes
(Horizontal or Vertical)
"""
module Alignment

    """
    Module storing horizontal alignment modes

    Usage:
    - `Alignment.Horizontal.LEFT`
    - ...
    """
    module Horizontal
        @enum AlignmentHorizontalCode::Int8 begin
            NONE = -1  # Undefined
            LEFT = 0   # Left (x)
            CENTRE = 1 # Centre
            RIGHT = 2  # Right (x)
        end
    end

    """
    Module storing vertical alignment modes

    Usage:
    - `Alignment.Vertical.DOWN`
    - ...
    """
    module Vertical
        @enum AlignmentVerticalCode::Int8 begin
            NONE = -1  # Undefined
            UP = 0     # Top (y)
            CENTRE = 1 # Centre
            DOWN = 2   # Bottom (y)
        end
    end
end

"""
Module storing background modes for widgets/drawables with backgrounds

Usage:
- `Background.SOLIDCOLOUR`
- ...
"""
module Background
    @enum BackgroundCode::Int8 begin
        NONE = -1        # Undefined
        CROPIMAGE = 0    # Cropped image
        SOLIDCOLOUR = 1  # Solid colour
        IMAGE = 2        # Full size image
        NOBACKCOLOUR = 3 # No backcolour
    end
end

"""
Module storing several default colour values

Usage:
- Colour.RED
- ...
"""
module Colour
    @enum ColourCode::Int32 begin
        NONE = -1
        BLACK = 0
        WHITE = 65535
        RED = 63488
        GREEN = 2016
        BLUE = 31
        GRAY = 33840
        BROWN = 48192
        YELLOW = 65504
    end
end

module Scroll
    """
    Module storing scrolling properties (such as directions for scrolling

    Usage:
    - Scroll.Direction.LEFT_TO_RIGHT
    """
    module Direction
        @enum ScrollDirectionCode::Int8 begin
            NONE = -1
            LEFT_TO_RIGHT = 0
            RIGHT_TO_LEFT = 1
            UP_TO_DOWN = 2
            DOWN_TO_UP = 3
        end
    end
end

"""
Module storing Nextion supported baudrates
"""
module Baudrate
    SUPPORTED = [2400, 4800, 9600, 19200, 38400, 57600, 115200]

    function _assert_supported(value)
        if !(value in SUPPORTED)
            error("baudrate=$value not in $SUPPORTED")
        end
    end

    function at(value)
        _assert_supported(value)
        value
    end

    function id(value)
        _assert_supported(value)
        findall(x -> x == value, SUPPORTED)[1]
    end
end

"""
Module storing Format for numeric values
"""
module Format
    @enum FormatCode::UInt8 begin
        Decimal = 0
        Currency
        Hex
    end
end


text:
include("iwp.jl")

env = IWPEnv()
x0 = [1*π/180, 0.5, 0, 0]
simulate!(env, x0)

A, B, K, S = lqr_gains(env; Q=env.policy.Q, R=env.policy.R);

function test_bo_vs_lqr(;iters::Int=1000, reset::Bool=false)
    res, opt = find_gains_BO(env; maxiters=100)

    unc_lqr_loss = simulate!(env, x0, K)
    bo_loss = simulate!(env, x0, res.observed_optimizer)

    while bo_loss > unc_lqr_loss && opt.iterations.i < iters
        if reset
            res, opt = find_gains_BO(env; maxiters=100)
        else
            res = boptimize!(opt)
        end
        bo_loss = simulate!(env, x0, res.observed_optimizer)
    end
    return bo_loss, unc_lqr_loss, res
end

bo_loss, unc_lqr_loss, res = test_bo_vs_lqr(iters=1000, reset=true)
@assert bo_loss <= unc_lqr_loss

ext: jl


text:
module CellValuesReindexTests

using Test
using Gridap
using Gridap.CellValuesGallery
using Gridap.CellMapApply: IndexCellMapFromKernel
using ..CellValuesMocks
using ..MapsMocks

a = [1,2,4,5,2,4,6,7,4,3,5,1,2,3]
ca = CellValueFromArray(a)
b = [10,11,3,1]
cb = CellValueFromArray(b)

cc = reindex(ca,cb)

d = a[b]
test_index_cell_value(cc,d)

b = [10,10,10,10]
cb = TestIterCellValue(10,4)

cc = reindex(ca,cb)
d = a[b]
test_iter_cell_value(cc,d)

l = 10
a = VectorValue(10,10)
p1 = VectorValue(1,1)
p2 = VectorValue(2,2)
p3 = VectorValue(3,3)
p = [p1,p2,p3]
m = MockMap(a)
cp = ConstantCellValue(p,l)
cm1 = ConstantCellValue(m,l)
cm2 = CompressedCellValue([m,],fill(1,l))
cm3 = cm1 + cm2
b = [10,11,3,1]
cb = CellValueFromArray(b)
rcm = reindex(cm3,cb)
rcp = reindex(cp,cb)
@test isa(rcm,IndexCellMapFromKernel)

end # module


text:
struct CartpoleCost{T} <: TO.CostFunction
    Q::Vector{T}
    R::Vector{T}
end
function TO.stage_cost(cost::CartpoleCost, x, u)
    return TO.stage_cost(cost, x) + 0.5 * cost.R[1] * u[1]^2
end
function TO.stage_cost(cost::CartpoleCost, x::AbstractVector)
    y = x[1]
    θ = x[2]
    ydot = x[3]
    θdot = x[4]
    J = cost.Q[2] * (1 .- sin(θ/2))
    J += 0.5* (cost.Q[1] * y^2 + cost.Q[3] * ydot^2 + cost.Q[4] * θdot^2)
    return J
end
TO.state_dim(::CartpoleCost) = 4
TO.control_dim(::CartpoleCost) = 1
Base.copy(cost::CartpoleCost) = CartpoleCost(copy(cost.Q), copy(cost.R))

@testset "Nonlinear Cartpole" begin
## Initialize model
model = RobotZoo.Cartpole()
N = 101
tf = 3.0

## Set up cost function
Q = [0.1,0.1,0.1,0.1]
R = [0.1]
Qf = [100,1000,100,100.]
costfun = CartpoleCost(Q,R)
costfun_term = CartpoleCost(Qf,R*0)
obj = Objective(costfun, costfun_term, N)

## Define problem
x0 = zeros(4)
xf = [0,pi,0,0]

# Solve w/ iLQR
prob = Problem(model, obj, xf, tf, x0=x0)
ilqr = Altro.iLQRSolver(prob,
    cost_tolerance=1e-3, gradient_tolerance=1e-2)
b = benchmark_solve!(ilqr)
err = states(ilqr)[end] - xf
@test err'err < 1e-3
if VERSION >= v"1.5"
    @test b.allocs == 0
end

# Add constraints
cons = ConstraintList(4,1,N)
add_constraint!(cons, GoalConstraint(xf), N)
prob = Problem(model, obj, xf, tf, x0=x0, constraints=cons)
solver = ALTROSolver(prob, cost_tolerance_intermediate=1e-2, show_summary=false)
solver.opts.ρ_primal = 1e-3
solve!(solver)
err = states(ilqr)[end] - xf
@test err'err < 1e-4

end


text:
# This file is a part of BAT.jl, licensed under the MIT License (MIT).

using BAT
using Test

using LinearAlgebra, Random
using Distributions, PDMats, StatsBase

@testset "distribution_functions" begin
    @testset "_check_rand_compat" begin
        @test BAT._check_rand_compat(MvNormal(ones(2)), ones(2,10)) == nothing
        @test_throws DimensionMismatch BAT._check_rand_compat(MvNormal(ones(3)), ones(2,10))
    end

    @testset "BAT.issymmetric_around_origin" begin
        @test BAT.issymmetric_around_origin(Normal(0.0, 3.2)) == true
        @test BAT.issymmetric_around_origin(Normal(1.7, 3.2)) == false
        @test BAT.issymmetric_around_origin(Gamma(4.2, 2.2)) == false
        @test BAT.issymmetric_around_origin(Chisq(20.3)) == false
        @test BAT.issymmetric_around_origin(MvNormal(zeros(2), ones(2))) == true
        @test BAT.issymmetric_around_origin(MvNormal(ones(2), ones(2))) == false
        @test BAT.issymmetric_around_origin(MvTDist(1.5, zeros(2), PDMat(Matrix{Float64}(I, 2, 2))))
        @test BAT.issymmetric_around_origin(MvTDist(1.5, ones(2), PDMat(Matrix{Float64}(I, 2, 2)))) == false
    end
end


text:
function load_dxf_line(filename::String)
    io = open(filename, "r")
    DXF = readlines(io)
    close(io)
    i = 1
    lines = Array{Float64,1}[]
    while DXF[i] != "EOF"
        if DXF[i] == "LINE"
            global line = Float64[]
        end
        if DXF[i] == "10"
            push!(line, parse(Float64,DXF[i+1]))
            i = i+1
        end
        if DXF[i] == "20"
            push!(line, parse(Float64,DXF[i+1]))
            i = i+1
        end
        if DXF[i] == "30"
            push!(line, parse(Float64,DXF[i+1]))
            i = i+1
        end
        if DXF[i] == "11"
            push!(line, parse(Float64,DXF[i+1]))
            i = i+1
        end
        if DXF[i] == "21"
            push!(line, parse(Float64,DXF[i+1]))
            i = i+1
        end
        if DXF[i] == "31"
            push!(line, parse(Float64,DXF[i+1]))
            i = i+1
            push!(lines,line)
        end
        i = i+1
    end

    V = hcat(reshape.(lines,3,2)...)
    EV = [[i,i+1] for i in range(1, size(V,2), step=2)]
    return V,EV
end

# V,EV = load_dxf_line("C:/Users/marte/Documents/GEOWEB/TEST/CONTEA_SEZIONE_z250/DXF/RAW/segment3D.dxf")
#
# using Visualization
# Visualization.VIEW([
#     Visualization.GLGrid(V,EV)
# ])


text:
using Revise
using Dice
using Dice: num_flips, num_nodes, to_dice_ir
using Distributions

code = @dice begin
    # triangle distribution

    function triangle(b::Int)
        s = false
        n = 2^b
        x = Vector(undef, b)
        y = Vector(undef, b)
        for i = b:-1:1
            x[i] = Dice.ifelse(s, flip(1/2), flip((3n - 2)/ (4n-4)))
            y[i] = flip((n-2)/(3n-2))
            s = s || (x[i] && !y[i])
            n = n/2
        end
        return DistInt(x)
    end

    function uniform(b::Int) # b is the bits for uniform, w is the bitwidth
        x = Vector(undef, b)
        for i = b:-1:1
            x[i] = flip(0.5)
        end
        return DistInt(x)
    end

    function discrete(p::Vector{Float64})
        mb = length(p)
        v = Vector(undef, mb)
        sum = 1
        for i=1:mb
            v[i] = p[i]/sum
            sum = sum - p[i]
        end

        # println(v)
        ans = DistInt(dicecontext(), mb-1)
        for i=mb-1:-1:1
            ans = if flip(v[i]) DistInt(dicecontext(), i-1) else ans end
        end
        return ans
    end

    function anyline(bits::Int, p::Float64)
        ans = Dice.ifelse(flip(p*2^bits), add_bits(uniform(bits), 3), add_bits(triangle(bits), 3))
        return ans
    end

    function gaussian(bits::Int, pieces::Int)
        d = Normal()
        start = quantile.(Normal(), 0.001)
        interval_sz = 2*abs(start)/pieces

        areas = Vector(undef, pieces)
        total_area = 0

        end_pts = Vector(undef, pieces)
        for i=1:pieces
            p1 = start + (i-1)*interval_sz
            p2 = p1 + interval_sz/2^bits
            p3 = start + (i)*interval_sz
            p4 = p3 - interval_sz/2^bits

            pts = [cdf.(d, p2) - cdf.(d, p1), cdf.(d, p3) - cdf.(d, p4)]
            end_pts[i] = pts

            areas[i] = (pts[1] + pts[2])*2^(bits - 1)
            total_area += areas[i]
        end

        # println(total_area)

        rel_prob = areas/total_area

        b = discrete(rel_prob)
        a = end_pts[pieces][1]/areas[pieces]
        l = a > 1/2^bits
        ans =  (if l
                    2^bits - 1 + (pieces - 1)*2^bits - add_bits(anyline(bits, 2/2^bits - a), 3)
                else
                    (pieces - 1)*2^bits + add_bits(anyline(bits, a), 3)
                end)[1]

        for i=pieces-1:-1:1
            a = end_pts[i][1]/areas[i]
            l = a > 1/2^bits
            ans = if prob_equals(b, i-1)
                    (if l
                        (2^bits - 1 + (i - 1)*2^bits - anyline(bits, 2/2^bits - a))
                    else
                        (i - 1)*2^bits +
                            anyline(bits, a)
                    end)[1]
                else
                    ans
                end
        end
        return ans
    end
    # println(max_bits(add_bits(anyline(2, 0.1), 3)))
    # println(max_bits(gaussian(2, 4)))
    g = gaussian(1, 108)
    rep = (g > 107) & (143 > g)
    rep
end

bdd = compile(code)
num_flips(bdd)
num_nodes(bdd)
println(infer(code, :bdd))
@assert infer(code, :bdd) ≈ 0.5

bdd = compile(code)
num_flips(bdd)
num_nodes(bdd)
infer(code, :bdd)


text:
using Base.Test
using suduko
p1=[0,7,0,9,6,4,0,0,0]
p2=[0,3,0,1,0,0,0,4,6]
p3=[0,0,0,8,0,0,0,0,5]
p4=[0,2,0,0,0,0,0,0,3]
p5=[6,0,0,0,0,0,0,0,4]
p6=[4,0,0,0,0,0,0,9,0]
p7=[1,0,0,0,0,3,0,0,0]
p8=[3,9,0,0,0,5,0,1,0]
p9=[0,0,0,7,1,6,0,5,0]

tester=hvcat(9,p1,p2,p3,p4,p5,p6,p7,p8,p9)
print ("in tester")
print(suduko.notConsistent (tester));
print ("done")

p1=[7,7,7,9,6,4,0,0,0] #this is not consistent
p2=[7,3,0,1,0,0,0,4,6]
p3=[0,0,0,8,0,0,0,0,5]
p4=[0,2,0,0,0,0,0,0,3]
p5=[6,0,0,0,0,0,0,0,4]
p6=[4,0,0,0,0,0,0,9,0]
p7=[1,0,0,0,0,3,0,0,0]
p8=[3,9,0,0,0,5,0,1,0]
p9=[0,0,0,7,1,6,0,5,0]

tester=hvcat(9,p1,p2,p3,p4,p5,p6,p7,p8,p9)

println("in tester")
print (suduko.notConsistent(tester))
print ("done")


text:
using Test, Mill, Flux
using Mill: p_map, inv_p_map

X = Matrix{Float64}(reshape(1:12, 2, 6))
BAGS = AlignedBags([1:1, 2:3, 4:6])
W = [1, 1/2, 1/2, 1/8, 1/3, 13/24]
W_mat = vcat(W, 2*W)
C = [1, 2]

@testset "basic aggregation functionality" begin
    @test SegmentedMean(2)(X, BAGS) ≈ [1.0 4.0 9.0; 2.0 5.0 10.0]
    @test SegmentedMax(2)(X, BAGS) ≈ [1.0 5.0 11.0; 2.0 6.0 12.0]
    @test SegmentedMeanMax(2)(X, BAGS) ≈ cat(SegmentedMean(2)(X, BAGS), SegmentedMax(2)(X, BAGS), dims=1)
    @test SegmentedMean(2)(X, BAGS, W) ≈ [1.0 4.0 236/24; 2.0 5.0 260/24]
    @test SegmentedMax(2)(X, BAGS, W) ≈ SegmentedMax(2)(X, BAGS)
    @test SegmentedMeanMax(2)(X, BAGS, W) ≈ cat(SegmentedMean(2)(X, BAGS, W), SegmentedMax(2)(X, BAGS, W), dims=1)
end

@testset "matrix functionality" begin
    @test SegmentedMean(2)(X, BAGS, W_mat) ≈ SegmentedMean(2)(X, BAGS, W)
end

@testset "pnorm functionality" begin
    for t = 1:10
        a, b, c, d, ρ1, ρ2, c1, c2, w1, w2 = randn(10)
        p1, p2 = p_map(ρ1), p_map(ρ2)
        w1 = abs(w1)
        w2 = abs(w2)
        bags = ScatteredBags([[1,2]])
        agg = SegmentedPNorm([ρ1, ρ2], [c1, c2], C)
        @test agg([a b; c d], bags) ≈ [
                                     (1/2*(abs(a-c1)^p1 + abs(b-c1)^p1))^(1/p1);
                                     (1/2*(abs(c-c2)^p2 + abs(d-c2)^p2))^(1/p2)
                                    ]
        @test agg([a b; c d], bags) ≈ [
                                     (1/2*(abs(a-c1)^p1 + abs(b-c1)^p1))^(1/p1);
                                     (1/2*(abs(c-c2)^p2 + abs(d-c2)^p2))^(1/p2)
                                    ]
        @test agg([a b; c d], bags, [w1, w2]) ≈ [
                                               (1/(w1+w2)*(w1*abs(a-c1)^p1 + w2*abs(b-c1)^p1))^(1/p1);
                                               (1/(w1+w2)*(w1*abs(c-c2)^p2 + w2*abs(d-c2)^p2))^(1/p2)
                                              ]
        x = randn(2, 6)
        agg = SegmentedPNorm(inv_p_map.([1+1e-16, 1+1e-16]), [0, 0], C)
        @test agg(x, BAGS) ≈ SegmentedMean(2)(abs.(x), BAGS)
        agg = SegmentedPNorm(inv_p_map.([2, 2]), [0, 0], C)
        @test agg(x, BAGS) ≈ hcat([sqrt.(sum(x[:, b] .^ 2, dims=2) ./ length(b)) for b in BAGS]...)
    end
end

@testset "lse functionality" begin
    for t = 1:10
        a, b, c, d, r1, r2 = randn(6)
        @test SegmentedLSE([r1, r2], C)([a b; c d], ScatteredBags([[1,2]])) ≈ [
                                                                               1/r1*log(1/2*(exp(a*r1)+exp(b*r1)));
                                                                               1/r2*log(1/2*(exp(c*r2)+exp(d*r2)))
                                                                              ]
        X = randn(2, 6)
        r1, r2 = randn(2)
        @test all(SegmentedLSE([r1, r2], C)(X, BAGS) .== SegmentedLSE([r1, r2], C)(X, BAGS, W))
        # the bigger value of r, the closer we are to the real maximum
        @test isapprox(SegmentedLSE([100, 100], C)(X, BAGS), SegmentedMax(2)(X, BAGS), atol=0.1)
    end
end

@testset "lse numerical stability" begin
    b1 = AlignedBags([1:2])
    b2 = ScatteredBags([[1,2]])
    @test SegmentedLSE([1,1], C)([1e15 1e15; 1e15 1e15], b1) ≈ [1e15; 1e15]
    @test SegmentedLSE([1,1], C)([-1e15 -1e15; -1e15 -1e15], b2) ≈ [-1e15; -1e15]
    @test SegmentedLSE([1,1], C)([1e15 1e15; 1e15 1e15], b1, [1, 1]) ≈ [1e15; 1e15]
    @test SegmentedLSE([1,1], C)([1e15 1e15; 1e15 1e15], b2, [2, 2]) ≈ [1e15; 1e15]
    @test SegmentedLSE([1,1], C)([-1e15 -1e15; -1e15 -1e15], b1, [1, 1]) ≈ [-1e15; -1e15]
    @test SegmentedLSE([1,1], C)([-1e15 -1e15; -1e15 -1e15], b2, [2, 2]) ≈ [-1e15; -1e15]
end

@testset "missing values" begin
    dummy = randn(2)
    for bags in [AlignedBags([0:-1]), AlignedBags([0:-1, 0:-1, 0:-1])]
        @test SegmentedMean(C)(missing, bags) == repeat(C, 1, length(bags))
        @test SegmentedMax(C)(missing, bags) == repeat(C, 1, length(bags))
        @test SegmentedLSE(dummy, C)(missing, bags) == repeat(C, 1, length(bags))
        @test SegmentedPNorm(dummy, dummy, C)(missing, bags) == repeat(C, 1, length(bags))
    end
end


text:
using Metatheory
using Test
using Metatheory.Library
using Metatheory.EGraphs
using Metatheory.Classic
using Metatheory.Util
using Metatheory.EGraphs.Schedulers

# =====================================================

# Zen Lineage Chart Example from Julog.jl https://github.com/ztangent/Julog.jl
# clauses = @julog [
#   ancestor(sakyamuni, bodhidharma) <<= true,
#   teacher(bodhidharma, huike) <<= true,
#   teacher(huike, sengcan) <<= true,
#   teacher(sengcan, daoxin) <<= true,
#   teacher(daoxin, hongren) <<= true,
#   teacher(hongren, huineng) <<= true,
#   ancestor(A, B) <<= teacher(A, B),
#   ancestor(A, C) <<= teacher(B, C) & ancestor(A, B),
#   grandteacher(A, C) <<= teacher(A, B) & teacher(B, C)
# ]

facts = [
    :(ancestor(sakyamuni, bodhidharma)),
    :(teacher(bodhidharma, huike)),
    :(teacher(huike, sengcan)),
    :(teacher(sengcan, daoxin)),
    :(teacher(daoxin, hongren)),
    :(teacher(hongren, huineng)),
]

function addfacts!(g::EGraph, facts)
    for fact ∈ facts
        fc, _ = addexpr!(g, fact)
        tc, _ = addexpr!(g, true)
        merge!(g, fc.id, tc.id)
    end
end

clauses = @theory begin
    teacher(a, b) => ancestor(a, b)
    # grandteacher(A, C) <<= teacher(A, B) & teacher(B, C)
end

# TODO syntax for MultiPatRewriteRule and PatEquiv
#   ancestor(A, C) <<= teacher(B, C) & ancestor(A, B),
lhs = @pat teacher(b, c)
rhs = @pat ancestor(a,c)
pat1 = PatEquiv((@pat ancestor(a, b)), @pat teacher(b,c))
q = MultiPatRewriteRule(lhs, rhs, [pat1])

push!(clauses, q)

# goals to prove: ancestor(sakyamuni, huineng)
g = EGraph()
addfacts!(g, facts)

query = :(ancestor(sakyamuni, huineng))
addexpr!(g, query)

params = SaturationParams(timeout=14)
saturate!(g, clauses, params)

# display(g.classes); println()

emptyt = @theory begin end
@test areequal(g, emptyt, true, query)


text:
module Protos

include("parsing.jl")

include("specs.jl")
using .Specs
export ProtoFile

include("utils.jl")

include("generation.jl")

include("serialization.jl")

end


text:
function test5()
  filename = "test5.asc"
  # exectuablepath = null string will not run LTspice.exe.  Test parsing only.
  sim = LTspiceSimulation(filename,executablepath="")
  @test LTspice.does_circuitfilearray_file_match(sim)
  show(IOBuffer(),sim)

  asteps = [1.0,2.0]

  @test stepvalues(sim) == (asteps,)
  @test measurementnames(sim) == ("sum","sump1000")
  @test stepnames(sim) == ("a",)
  @test logpath(sim) != ""

  @static if Sys.iswindows()
      @test circuitpath(sim) == filename
  end

  @test parametervalues(sim) == []
  @test measurementvalues(sim)[1,1] == 1.0
  @test executablepath(sim) == ""
  @test haskey(sim,"sum") == true
  @test length(sim) == 4

  for i in eachindex(asteps)
      @test measurementvalues(sim)[i,1] == asteps[i]
      @test sim["sum"][i] == asteps[i]
      @test measurementvalues(sim)[i,2] == asteps[i] + 1000
      @test sim["sump1000"][i] == asteps[i] + 1000
  end

  for line in perlineiterator(sim)
      @test (line[1] == line[2])
      @test (line[2] + 1000 == line[3])
  end

  show(IOBuffer(),sim)
end
test5()


text:
using OnlineStats, Test, LinearAlgebra, Random, StatsBase, Statistics, Dates
using LossFunctions, PenaltyFunctions

#-----------------------------------------------------------------------# utils
n = 1000
x,  y,  z  = rand(Bool, n), randn(n), rand(1:10, n)
x2, y2, z2 = rand(Bool, n), randn(n), rand(1:10, n)
xs, ys, zs = vcat(x, x2), vcat(y, y2), vcat(z, z2)

p = 5
xmat,  ymat,  zmat  = rand(Bool, n, p), randn(n, p), rand(1:10, n, p)
xmat2, ymat2, zmat2 = rand(Bool, n, p), randn(n, p), rand(1:10, n, p)

function mergestats(a::OnlineStat, y1, y2)
    b = copy(a)
    fit!(a, y1)             # fit a on y1
    fit!(b, y2)             # fit b on y2
    merge!(a, b)            # merge b into a
    fit!(b, y1)             # fit b on y1
    @test nobs(a) == nobs(b) == length(y1) + length(y2)
    a, b
end
mergevals(o1::OnlineStat, y1, y2) = map(value, mergestats(o1, y1, y2))

#-----------------------------------------------------------------------# AutoCov
@testset "AutoCov" begin
    o = fit!(AutoCov(10), y)
    @test autocov(o) ≈ autocov(y, 0:10)
    @test autocor(o) ≈ autocor(y, 0:10)
    @test nobs(o) == n
end
#-----------------------------------------------------------------------# Bootstrap
@testset "Bootstrap" begin
    o = fit!(Bootstrap(Mean(), 100, [1]), y)
    @test all(value.(o.replicates) .== value(o.stat))
    c = confint(o)
    @test length(c) == 2
    @test c[1] ≤ c[2]
end
#-----------------------------------------------------------------------# CallFun
@testset "CallFun" begin
    i = 0
    o = fit!(CallFun(Mean(), x -> i+=1), y)
    @test value(o) ≈ mean(y)
    @test i == n
    @test ≈(mergevals(CallFun(Mean(), x->nothing), y, y2)...)
end
#-----------------------------------------------------------------------# CovMatrix
@testset "CovMatrix" begin
    o = fit!(CovMatrix(), eachrow(ymat))
    @test value(o) ≈ cov(ymat)
    @test cov(o) ≈ cov(ymat)
    @test cor(o) ≈ cor(ymat)
    @test all(x -> ≈(x...), zip(var(o), var(ymat; dims=1)))
    @test all(x -> ≈(x...), zip(std(o), std(ymat; dims=1)))
    @test all(x -> ≈(x...), zip(mean(o), mean(ymat; dims=1)))

    @test ≈(mergevals(CovMatrix(), OnlineStatsBase.eachrow(ymat), OnlineStatsBase.eachrow(ymat2))...)
    @test ≈(mergevals(CovMatrix(), OnlineStatsBase.eachcol(ymat'), OnlineStatsBase.eachcol(ymat2'))...)
    @test ≈(mergevals(CovMatrix(Complex{Float64}), OnlineStatsBase.eachrow(ymat * im), OnlineStatsBase.eachrow(ymat2))...)
    @test ≈(mergevals(CovMatrix(Complex{Float64}), OnlineStatsBase.eachrow(ymat * im), OnlineStatsBase.eachrow(ymat2 * im))...)
end
#-----------------------------------------------------------------------# Diff
@testset "Diff" begin
    @test value(fit!(Diff(), y)) == y[end] - y[end-1]
    o = fit!(Diff(Int), 1:10)
    @test diff(o) == 1
    @test last(o) == 10
end
#-----------------------------------------------------------------------# Fit[Dist]
@testset "Fit[Dist]" begin
    @testset "FitBeta" begin
        @test value(FitBeta()) == (1.0, 1.0)
        a, b = mergevals(FitBeta(), rand(100), rand(100))
        @test a[1] ≈ b[1]
        @test a[2] ≈ b[2]
    end
    @testset "FitCauchy" begin
        @test value(FitCauchy()) == (0.0, 1.0)
        a, b = mergevals(FitCauchy(), y, y2)
        @test ≈(a[1], b[1]; atol=.5)
        @test ≈(a[2], b[2]; atol=.5)
    end
    @testset "FitGamma" begin
        @test value(FitGamma()) == (1.0, 1.0)
        a, b = mergevals(FitGamma(), rand(100), rand(100))
        @test a[1] ≈ b[1]
        @test a[2] ≈ b[2]
    end
    @testset "FitLogNormal" begin
        @test value(FitLogNormal()) == (0.0, 1.0)
        a, b = mergevals(FitLogNormal(), exp.(y), exp.(y2))
        @test a[1] ≈ b[1]
        @test a[2] ≈ b[2]
    end
    @testset "FitNormal" begin
        @test value(FitNormal()) == (0.0, 1.0)
        a, b = mergestats(FitNormal(), y, y2)
        @test value(a)[1] ≈ value(b)[1]
        @test value(a)[2] ≈ value(b)[2]

        @test mean(a) ≈ mean(ys)
        @test var(a) ≈ var(ys)
        @test std(a) ≈ std(ys)

        # pdf and cdf
        o = fit!(FitNormal(), [-1, 0, 1])
        @test OnlineStats.pdf(o, 0.0) ≈ 0.3989422804014327
        @test OnlineStats.pdf(o, -1.0) ≈ 0.24197072451914337
        @test OnlineStats.cdf(o, 0.0) ≈ 0.5
        @test ≈(OnlineStats.cdf(o, -1.0), 0.15865525393145702; atol=.001)
    end
    @testset "FitMultinomial" begin
        o = FitMultinomial(5)
        @test value(o)[2] == ones(5) / 5
        data = [1 2 3 4 5; 1 2 3 4 5]
        @test value(fit!(o, eachrow(data)))[2] == collect(2:2:10) ./ sum(data)

        data1 = OnlineStatsBase.eachrow(rand(1:4, 10, 3))
        data2 = OnlineStatsBase.eachrow(rand(2:7, 11, 3))
        a, b = mergevals(FitMultinomial(3), data1, data2)
        @test ≈(a[2], b[2])
    end
    @testset "FitMvNormal" begin
        @test value(FitMvNormal(2)) == (zeros(2), Matrix(I, 2, 2))
        a, b = mergevals(FitMvNormal(2), OnlineStatsBase.eachrow([y y2]), OnlineStatsBase.eachrow([y2 y]))
        @test a[1] ≈ b[1]
        @test a[2] ≈ b[2]
    end
end
#-----------------------------------------------------------------------# FastNode
@testset "FastNode" begin
    X, Y = ymat, rand(1:3, 1000)
    X2, Y2 = ymat, rand(1:3, 1000)
    data  = zip(eachrow(X), Y)
    data2 = zip(eachrow(X2), Y2)
    o  = fit!(FastNode(5, 3), data)
    o2 = fit!(FastNode(5, 3), data2)
    merge!(o, o2)
    fit!(o2, data)
    for k in 1:3, j in 1:5
        @test value(o.stats[k][j])[1] ≈ value(o2.stats[k][j])[1]
        @test value(o.stats[k][j])[2] ≈ value(o2.stats[k][j])[2]
    end
    @test length(o[1]) == 3

    pvec = [mean(Y .== 1), mean(Y .== 2), mean(Y .== 3)]
    o = fit!(FastNode(5, 3), data)
    @test probs(o) == pvec
    @test nobs(o) == 1000
    @test OnlineStats.nkeys(o) == 3
    @test OnlineStats.nvars(o) == 5
    @test classify(o) ∈ [1, 2, 3]
end
#-----------------------------------------------------------------------# FastTree
@testset "FastTree" begin
    X, Y = OnlineStats.fakedata(FastNode, 10^4, 10)
    o = fit!(FastTree(10; splitsize=100), zip(eachrow(X),Y))
    @test classify(o, X[1,:]) ∈ [1, 2]
    @test all(0 .< classify(o, X) .< 3)
    @test OnlineStats.nkeys(o) == 2
    @test OnlineStats.nvars(o) == 10
    @test mean(classify(o, X) .== Y) > .5

    # Issue 116
    Random.seed!(218)
    X,Y = OnlineStats.fakedata(FastNode, 10^4, 1)
    fit!(FastTree(1, splitsize=100), zip(eachrow(X),Y))
end
#-----------------------------------------------------------------------# FastForest
@testset "FastForest" begin
    X, Y = OnlineStats.fakedata(FastNode, 10^4, 10)
    o = fit!(FastForest(10; splitsize=500, λ = .7), zip(eachrow(X), Y))
    @test classify(o, randn(10)) in 1:2
    @test mean(classify(o, X) .== Y) > .5
end

#-----------------------------------------------------------------------# HeatMap
@testset "HeatMap" begin
    data1 = OnlineStatsBase.eachrow(ymat[:, 1:2])
    data2 = OnlineStatsBase.eachrow(ymat2[:, 1:2])
    @test ==(mergevals(HeatMap(-5:.1:5, -5:.1:5), data1, data2)...)
end
#-----------------------------------------------------------------------# Hist
@testset "Hist" begin
    @test ==(mergevals(Hist(-5:.1:5), y, y2)...)
    @testset "Hist compared to StatsBase.Histogram" begin
        for edges in (-5:5, collect(-5:5), [-5, -3.5, 0, 1, 4, 5.5])
            for data in (y, -6:.75:6)
                w  = fit(Histogram, data, edges, closed = :left).weights
                w2 = fit(Histogram, data, edges, closed = :right).weights
                @test fit!(Hist(edges, Number; closed=false, left=true),  data).counts == w
                @test fit!(Hist(edges, Number; closed=false, left=false), data).counts == w2
            end
        end
    end
    o = fit!(Hist(-5:.1:5), y)
    for (v1, v2) in zip(extrema(o), extrema(y))
        @test ≈(v1, v2; atol=.1)
    end
    @test ≈(mean(o), mean(y); atol=.1)
    @test ≈(var(o), var(y); atol=.2)

    # merge unequal bins
    r1, r2 = -5:.2:5, -5:.1:5
    @test merge!(fit!(Hist(r1), y), fit!(Hist(r2), y2)) == fit!(Hist(r1), vcat(y, y2))
    @test OnlineStats.pdf(fit!(Hist(-5:.1:5), y), 0) > 0
    @test OnlineStats.pdf(fit!(Hist(-5:.1:5), y), 100) == 0
end
#-----------------------------------------------------------------------# KHist
@testset "KHist" begin
    @test KHist(10) == KHist(10)

    o = fit!(KHist(1000), y)
    @test mean(o) ≈ mean(y)
    @test var(o) ≈ var(y)
    @test median(o) ≈ median(y)
    @test quantile(o) ≈ quantile(y, [0, .25, .5, .75, 1])
    @test std(o) ≈ std(y)
    @test extrema(o) == extrema(y)

    ==(mergevals(KHist(2000), y, y2)...)
    ==(mergevals(KHist(1), y, y2)...)
    ==(mergevals(KHist(2000, Float32), Float32.(y), Float32.(y2))...)

    data = randn(10_000)
    o = fit!(KHist(50), data)
    @test OnlineStats.pdf(o, -10) == 0.0
    @test ≈(OnlineStats.pdf(o, 0.0), 0.3989422804014327, atol=.5)
    @test OnlineStats.pdf(o, 10) == 0.0
    @test OnlineStats.cdf(o, -10) == 0.0
    @test ≈(OnlineStats.cdf(o, 0.0), .5; atol=.1)
    @test OnlineStats.cdf(o, 10) == 1.0
end
#-----------------------------------------------------------------------# HyperLogLog
@testset "HyperLogLog" begin
    @test ≈(value(fit!(HyperLogLog(), y)), n; atol=20)
    @test ==(mergevals(HyperLogLog(), x, x2)...)
    @test ==(mergevals(HyperLogLog(), y, y2)...)
    @test ==(mergevals(HyperLogLog(), z, z2)...)
end
#-----------------------------------------------------------------------# IndexedPartition
@testset "IndexedPartition" begin
    o = IndexedPartition(Float64, Mean())
    fit!(o, zip(y, y2))
end
#-----------------------------------------------------------------------# KahanSum
@testset "KahanSum" begin
    @test value(fit!(KahanSum(), y)) ≈ sum(y)
    @test value(fit!(KahanSum(Int), x)) == sum(x)
    @test value(fit!(KahanSum(Int), z)) == sum(z)
    @test ≈(mergevals(KahanSum(), y, y2)...)
end
#-----------------------------------------------------------------------# KahanMean
@testset "KahanMean" begin
    @test value(fit!(KahanMean(), y)) ≈ mean(y)
    @test ≈(mergevals(KahanMean(), y, y2)...)
end
#-----------------------------------------------------------------------# KahanVariance
@testset "KahanVariance" begin
    o = fit!(KahanVariance(), y)
    @test mean(o) ≈ mean(y)
    @test var(o) ≈ var(y)
    @test std(o) ≈ std(y)
    @test ≈(mergevals(KahanVariance(), y, y2)...)

    # Issue 116
    @test std(KahanVariance()) == 1
    @test std(fit!(KahanVariance(), 1)) == 1
    @test std(fit!(KahanVariance(), [1, 2])) == sqrt(.5)
end
#-----------------------------------------------------------------------# KMeans
@testset "KMeans" begin
    o = fit!(KMeans(5,2), eachrow(ymat))
    sort!(o, rev=true)
    @test o.value[1].n ≥ o.value[2].n
end
#-----------------------------------------------------------------------# LinReg
@testset "LinReg" begin
    ≈(mergevals(LinReg(), zip(eachrow(ymat), y), zip(eachrow(ymat2), y2))...)

    o = fit!(LinReg(), zip(eachrow(ymat), y))
    @test coef(o) ≈ ymat \ y
    @test coef(o, .1) ≈ (ymat'ymat ./ n + .1I) \ ymat'y ./ n
    @test coef(o, .1:.1:.5) ≈ (ymat'ymat ./ n + Diagonal(.1:.1:.5)) \ ymat'y ./ n
    @test predict(o, ymat) == ymat * o.β
    @test predict(o, ymat[1,:]) == dot(ymat[1,:], o.β)
end
#-----------------------------------------------------------------------# LinRegBuilder
@testset "LinRegBuilder" begin
    @test ≈(mergevals(LinRegBuilder(), OnlineStatsBase.eachrow(ymat), OnlineStatsBase.eachrow(ymat2))...)

    o = fit!(LinRegBuilder(), eachrow(ymat))
    for i in 1:5
        data = ymat[:, setdiff(1:5, i)]
        @test coef(o; y=i) ≈ [data ones(n)] \ ymat[:,i]
        @test coef(o, .1; y=i, bias=false) ≈ (data'data ./ n + .1*I) \ data'ymat[:,i] ./ n
    end

    o2 = fit!(LinReg(), zip(eachrow(ymat[:,[4,1]]), ymat[:,3]))
    @test coef(o, [.2,.4]; y=3, x = [4,1], bias=false) ≈ coef(o2, [.2, .4])
end
#-----------------------------------------------------------------------# ML
@testset "ML" begin
    o = OnlineStats.preprocess(OnlineStatsBase.eachrow(ymat))
    for i in 1:5
        @test o.group[i] isa OnlineStats.Numerical
    end
    o = OnlineStats.preprocess(OnlineStatsBase.eachrow(xmat))
    for i in 1:2
        @test o.group[i] isa OnlineStats.Categorical
    end
    o = OnlineStats.preprocess(OnlineStatsBase.eachrow(zmat), 3 => OnlineStats.Categorical(Int))
    @test o.group[3] isa OnlineStats.Categorical
    @test o.group[1] isa OnlineStats.Numerical
end

#-----------------------------------------------------------------------# Mosaic
@testset "Mosaic" begin
    @test ==(mergevals(Mosaic(Int,Int), zip(z, z2), zip(z2, z))...)
end
#-----------------------------------------------------------------------# MovingTimeWindow
@testset "MovingTimeWindow" begin
    dates = Date(2010):Day(1):Date(2011)
    data = Int.(1:length(dates))
    o = fit!(MovingTimeWindow(Day(4); timetype=Date, valtype=Int), zip(dates, data))
    @test value(o) == collect(Pair(a,b) for (a,b) in zip(dates[end-4:end], data[end-4:end]))

    d1 = zip(dates[1:2], data[1:2])
    d2 = zip(dates[3:4], data[3:4])
    @test ==(mergevals(MovingTimeWindow(Day(4); timetype=Date, valtype=Int), d1, d2)...)
end
#-----------------------------------------------------------------------# MovingWindow
@testset "MovingWindow" begin
    o = fit!(MovingWindow(10, Int), 1:12)
    for i in 1:10
        @test o[i] == (1:12)[i + 2]
    end
end
#-----------------------------------------------------------------------# NBClassifier
@testset "NBClassifier" begin
    o = fit!(NBClassifier(5, Bool), zip(eachrow(ymat),x))
    merge!(o, fit!(NBClassifier(5, Bool), zip(eachrow(ymat2), x2)))
    @test nobs(o) == 2000
    @test length(probs(o)) == 2
    @test sum(predict(o, ymat[1,:])) ≈ 1
    @test classify(o, ymat[1, :]) || !classify(o, ymat[1, :])
    @test OnlineStats.nvars(o) == 5
    @test OnlineStats.nkeys(o) == 2
    @test length(o[2]) == 2
end
#-----------------------------------------------------------------------# OrderStats
@testset "OrderStats" begin
    @test ≈(mergevals(OrderStats(100), y, y2)...)
    o = fit!(OrderStats(n), y)
    @test value(o) == sort(y)
    @test quantile(o, 0:.25:1) == quantile(y, 0:.25:1)
end
#-----------------------------------------------------------------------# Partition
@testset "Partition" begin
    # merging
    o = fit!(Partition(Mean(), 1000), y)
    o2 = fit!(Partition(Mean(), 1000), y2)
    merge!(o, o2)
    fit!(o2, y)
    @test nobs(o) == nobs(o2)
    @test all(nobs.(o.parts) .== nobs.(o2.parts))
    for i in 1:5
        @test value(o.parts[i]) ≈ value(o2.parts[500 + i])
    end
end
#-----------------------------------------------------------------------# ProbMap
@testset "ProbMap" begin
    @test sort(collect(keys(fit!(ProbMap(Float64), y).value))) == sort(y)
    # merge
    data, data2 = rand(1:4, 100), rand(1:4, 100)
    o = fit!(ProbMap(Int), data)
    o2 = fit!(ProbMap(Int), data2)
    merge!(o, o2)
    fit!(o2, data)
    @test sort(collect(keys(o.value))) == sort(collect(keys(o2.value)))
    @test probs(fit!(ProbMap(Int), [1,1,2,2,3,3,4,4])) ≈ fill(.25, 4)
    @test probs(fit!(ProbMap(Int), [1,1,2,2,3,3,4,4]), [1,2,9]) ≈ [.5, .5, 0]
end
#-----------------------------------------------------------------------# Quantile
@testset "Quantile/P2Quantile" begin
    data = randn(10_000)
    data2 = randn(10_000)
    τ = .1:.1:.9
    for o in [
            Quantile(τ; alg=SGD()),
            Quantile(τ; alg=MSPI()),
            Quantile(τ; alg=OMAS()),
            Quantile(τ; alg=ADAGRAD())
            ]
        @test ≈(value(fit!(copy(o), data)), quantile(data, τ), atol=.5)
        @test ≈(mergevals(o, data, data2)...; atol=.5)

    end
    for τi in τ
        @test ≈(value(fit!(P2Quantile(τi),data)), quantile(data, τi), atol=.2)
    end
end
#-----------------------------------------------------------------------# ReservoirSample
@testset "ReservoirSample" begin
    @test value(fit!(ReservoirSample(n), y)) == y
    a, b = mergestats(ReservoirSample(9), y, y2)
    for yi in value(a)
        @test (yi ∈ y) || (yi ∈ y2)
    end
end
#-----------------------------------------------------------------------# StatLag
@testset "StatLag" begin
    o = fit!(StatLag(Mean(), 10), 1:20)
    @test length(o.lag.value) == 10
    for (i, m) in enumerate(reverse(o.lag.value))
        @test nobs(m) == 10 + i
    end
    @test nobs(value(o)[1]) == 20
end
#-----------------------------------------------------------------------# StatLearn
@testset "StatLearn" begin
    X = randn(10_000, 5)
    β = collect(-1:.5:1)
    Y = X * β + randn(10_000)
    Y2 = 2.0 .* [rand()< 1 /(1 + exp(-η)) for η in X*β] .- 1.0
    for A in [SGD(),ADAGRAD(),ADAM(),ADAMAX(),ADADELTA(),RMSPROP(),OMAS(),OMAP(),MSPI()]
        print("  > $A")
        print(": ")
        for L in [.5 * L2DistLoss()]
            print(" | $L")
            # sanity checks
            o = fit!(StatLearn(5, A, L; rate=LearningRate(.7)), zip(eachrow(X),Y))
            @test o.loss isa typeof(L)
            @test o.alg isa typeof(A)
            any(isnan.(o.β)) && @info((L, A))
            merge!(o, copy(o))
            @test coef(o) == o.β
            @test predict(o, X) == X * o.β
            @test ≈(coef(o), β; atol=1.5)
            @test OnlineStats.objective(o, X, Y) ≈ value(o.loss, Y, predict(o, X), AggMode.Mean()) + value(o.penalty, o.β, o.λ)
        end
        for L in [LogitMarginLoss(), DWDMarginLoss(1.0)]
            print(" | $L")
            o = fit!(StatLearn(5, A, L), zip(eachrow(X),Y2))
            @test mean(Y2 .== classify(o, X)) > .5
        end
        println()
    end
end
#-----------------------------------------------------------------------# Kahan

include("test_kahan.jl")

#-----------------------------------------------------------------------# Show methods
@testset "Show methods" begin
    for stat in [BiasVec([1,2,3]), Bootstrap(Mean()), CallFun(Mean(), println), FastNode(5),
                 FastTree(5), FastForest(5), FTSeries(Variance()), Group(Mean(), Mean()),
                 HyperLogLog{10}(), LinRegBuilder(4), NBClassifier(5, Float64), ProbMap(Int),
                 P2Quantile(.5), Series(Mean()), StatLearn(5)]
        println("  > ", stat)
    end
end


text:
"""
Experiment in SubModules

$(DocStringExtensions.IMPORTS)
$(DocStringExtensions.EXPORTS)
"""
module PnmlBase

using DocStringExtensions
using Reexport
#using requires

# Could be many include/reexport pairs.
# These are each a submodule.
include("XmlUtils.jl")

@reexport using .XmlUtils

#function __init__() end

end


text:
# Invertible conditional HINT layer from Kruse et al. (2020)
# Author: Philipp Witte, pwitte3@gatech.edu
# Date: January 2020

export ConditionalLayerSLIM

"""
    CI = ConditionalLayerSLIM(nx1, nx2, nx_in, nx_hidden, ny1, ny2, ny_in, ny_hidden, batchsize, Op;
        type="affine", k1=3, k2=3, p1=1, p2=1, s1=1, s2=1)

 Create a conditional SLIM layer based on the HINT architecture.

 *Input*:

 - `nx1`, `nx2`: spatial dimensions of `X`

 - `nx_in`, `nx_hidden`: number of input and hidden channels of `X`

 - `ny1`, `ny2`: spatial dimensions of `Y`

 - `ny_in`, `ny_hidden`: number of input and hidden channels of `Y`

 - `Op`: Linear forward modeling operator

 - `type`: string to indicate which type of data coupling layer to use (`"additive"`, `"affine"`, `"learned"`)

 - `k1`, `k2`: kernel size of convolutions in residual block. `k1` is the kernel of the first and third
    operator, `k2` is the kernel size of the second operator.

 - `p1`, `p2`: padding for the first and third convolution (`p1`) and the second convolution (`p2`)

 - `s1`, `s2`: stride for the first and third convolution (`s1`) and the second convolution (`s2`)

 *Output*:

 - `CI`: Conditional SLIM coupling layer.

 *Usage:*

 - Forward mode: `Zx, Zy, logdet = CI.forward_X(X, Y, Op)`

 - Inverse mode: `X, Y = CI.inverse(Zx, Zy, Op)`

 - Backward mode: `ΔX, ΔY, X, Y = CI.backward(ΔZx, ΔZy, Zx, Zy, Op)`

 - Forward mode Y: `Zy = CI.forward_Y(Y)`

 - Inverse mode Y: `Y = CI.inverse(Zy)`

 *Trainable parameters:*

 - None in `CI` itself

 - Trainable parameters in coupling layers `CI.CL_X`, `CI.CL_Y`, `CI.CL_XY` and in
   permutation layers `CI.C_X` and `CI.C_Y`.

 See also: [`CouplingLayerHINT`](@ref), [`AffineCouplingLayerSLIM`](@ref), [`get_params`](@ref), [`clear_grad!`](@ref)
"""
struct ConditionalLayerSLIM <: NeuralNetLayer
    CL_X::CouplingLayerHINT
    CL_Y::CouplingLayerHINT
    CL_XY::Union{AdditiveCouplingLayerSLIM, AffineCouplingLayerSLIM, LearnedCouplingLayerSLIM}
    C_X::Conv1x1
    C_Y::Conv1x1
end

@Flux.functor ConditionalLayerSLIM

# Constructor from input dimensions
function ConditionalLayerSLIM(nx1::Int64, nx2::Int64, nx_in::Int64, nx_hidden::Int64, ny1::Int64, ny2::Int64, ny_in::Int64, ny_hidden::Int64,
    batchsize::Int64; type="affine", k1=3, k2=3, p1=1, p2=1, s1=1, s2=1)

    # Create basic coupling layers
    CL_X = CouplingLayerHINT(nx1, nx2, nx_in, nx_hidden, batchsize;
        k1=k1, k2=k2, p1=p1, p2=p2, s1=s1, s2=s2, logdet=true)
    CL_Y = CouplingLayerHINT(Int(ny1/2), Int(ny2/2), Int(ny_in*4), ny_hidden, batchsize;
        k1=k1, k2=k2, p1=p1, p2=p2, s1=s1, s2=s2, logdet=true)

    if type == "affine"
        CL_XY = AffineCouplingLayerSLIM(nx1, nx2, nx_in, nx_hidden, batchsize, identity;
            k1=k1, k2=k2, p1=p1, p2=p2, s1=s1, s2=s2, logdet=true, permute=false)
    elseif type == "additive"
        CL_XY = AdditiveCouplingLayerSLIM(nx1, nx2, nx_in, nx_hidden, batchsize, identity;
            k1=k1, k2=k2, p1=p1, p2=p2, s1=s1, s2=s2, logdet=true, permute=false)
    elseif type == "learned"
        CL_XY = LearnedCouplingLayerSLIM( nx1, nx2, nx_in, ny1, ny2, ny_in, nx_hidden, batchsize;
            k1=k1, k2=k2, p1=p1, p2=p2, s1=s1, s2=s2, logdet=true, permute=false)
    else
        throw("Specified layer type not defined.")
    end

    # Permutation using 1x1 convolution
    C_X = Conv1x1(nx_in)
    C_Y = Conv1x1(Int(ny_in*4))

    return ConditionalLayerSLIM(CL_X, CL_Y, CL_XY, C_X, C_Y)
end

function forward(X, Y, Op, CI::ConditionalLayerSLIM)

    # Y-lane: coupling
    Ys = wavelet_squeeze(Y)
    Yp = CI.C_Y.forward(Ys)
    Zy, logdet2 = CI.CL_Y.forward(Yp)
    Zy = wavelet_unsqueeze(Zy)

    # X-lane
    Xp = CI.C_X.forward(X)
    X, logdet1 = CI.CL_X.forward(Xp)
    if typeof(CI.CL_XY) == LearnedCouplingLayerSLIM
        Zx, logdet3 = CI.CL_XY.forward(X, reshape(Y, :, size(Y, 4)))   # Learn operator
    else
        Zx, logdet3 = CI.CL_XY.forward(X, reshape(Y, :, size(Y, 4)), Op)   # Use provided operator
    end
    logdet = logdet1 + logdet2 + logdet3
    return Zx, Zy, logdet
end

function inverse(Zx, Zy, Op, CI::ConditionalLayerSLIM)

    # Y-lane
    Zy = wavelet_squeeze(Zy)
    Yp = CI.CL_Y.inverse(Zy)
    Ys = CI.C_Y.inverse(Yp)
    Y = wavelet_unsqueeze(Ys)

    # X-lane
    if typeof(CI.CL_XY) == LearnedCouplingLayerSLIM
        X = CI.CL_XY.inverse(Zx, reshape(Y, :, size(Y, 4)))
    else
        X = CI.CL_XY.inverse(Zx, reshape(Y, :, size(Y, 4)), Op)
    end
    Xp = CI.CL_X.inverse(X)
    X = CI.C_X.inverse(Xp)

    return X, Y
end

function backward(ΔZx, ΔZy, Zx, Zy, Op, CI::ConditionalLayerSLIM)

    # Y-lane
    ΔZy = wavelet_squeeze(ΔZy)
    Zy = wavelet_squeeze(Zy)
    ΔYp, Yp = CI.CL_Y.backward(ΔZy, Zy)
    ΔYs, Ys = CI.C_Y.inverse((ΔYp, Yp))
    Y = wavelet_unsqueeze(Ys)
    ΔY = wavelet_unsqueeze(ΔYs)

    # X-lane
    if typeof(CI.CL_XY) == LearnedCouplingLayerSLIM
        ΔX, ΔY_, X = CI.CL_XY.backward(ΔZx, Zx, reshape(Y, :, size(Y, 4)))
    else
        ΔX, ΔY_, X = CI.CL_XY.backward(ΔZx, Zx, reshape(Y, :, size(Y, 4)), Op)
    end
    ΔY += reshape(ΔY_, size(ΔY))
    ΔXp, Xp = CI.CL_X.backward(ΔX, X)
    ΔX, X = CI.C_X.inverse((ΔXp, Xp))

    return ΔX, ΔY, X, Y
end

function forward_Y(Y, CI::ConditionalLayerSLIM)
    Ys = wavelet_squeeze(Y)
    Yp = CI.C_Y.forward(Ys)
    Zy, logdet2 = CI.CL_Y.forward(Yp)
    Zy = wavelet_unsqueeze(Zy)
    return Zy
end

function inverse_Y(Zy, CI::ConditionalLayerSLIM)
    Zy = wavelet_squeeze(Zy)
    Yp = CI.CL_Y.inverse(Zy)
    Ys = CI.C_Y.inverse(Yp)
    Y = wavelet_unsqueeze(Ys)
    return Y
end

# Clear gradients
function clear_grad!(CI::ConditionalLayerSLIM)
    clear_grad!(CI.CL_X)
    clear_grad!(CI.CL_Y)
    clear_grad!(CI.CL_XY)
    clear_grad!(CI.C_X)
    clear_grad!(CI.C_Y)
end

# Get parameters
function get_params(CI::ConditionalLayerSLIM)
    p = get_params(CI.CL_X)
    p = cat(p, get_params(CI.CL_Y); dims=1)
    p = cat(p, get_params(CI.CL_XY); dims=1)
    p = cat(p, get_params(CI.C_X); dims=1)
    p = cat(p, get_params(CI.C_Y); dims=1)
    return p
end

# Put parameters
function put_params!(CI::ConditionalLayerSLIM, Params::Array{Any,1})
    idx_s = 1
    counter = 0

    nlayers_x = length(CI.CL_X.CL)
    isnothing(CI.CL_X.C) ? counter += 5*nlayers_x : counter += 5*nlayers_x + 3
    put_params!(CI.CL_X, Params[idx_s:idx_s+counter-1])

    idx_s += counter
    counter = 0
    nlayers_y = length(CI.CL_Y.CL)
    isnothing(CI.CL_Y.C) ? counter += 5*nlayers_y : counter += 5*nlayers_y + 3
    put_params!(CI.CL_Y, Params[idx_s:idx_s+counter-1])

    idx_s += counter
    counter = 5
    isnothing(CI.CL_XY.C) ? counter += 0 : counter += 3
    put_params!(CI.CL_XY, Params[idx_s:idx_s+counter-1])

    idx_s += counter
    counter = 3
    ~isnothing(CI.C_X) && put_params!(CI.C_X, Params[idx_s:idx_s+counter-1]); idx_s += counter
    ~isnothing(CI.C_Y) && put_params!(CI.C_Y, Params[idx_s:idx_s+counter-1]); idx_s += counter
end


text:
using BesselK, BenchmarkTools, Printf

include("shared.jl")

const PAIRS = ((0.5,   1.0,  "half-integer"),
               (1.0,   1.0,  "whole integer"),
               (3.001, 1.0,  "near-integer order"),
               (3.001, 8.0,  "near-integer order, borderline arg"),
               (1.85,  1.0,  "small order"),
               (1.85,  8.0,  "small order, borderline arg"),
               (1.85,  14.0, "intermediate arg"),
               (1.85,  29.0, "large intermediate arg"),
               (1.85,  35.0, "large argument"))

for (j, (v, x, descriptor)) in enumerate(PAIRS)
  t_us    = (@belapsed BesselK._besselk($v, $x)   samples=1_000)*1e9 # our code
  t_am    = (@belapsed BesselK.besselk($v, $x)    samples=1_000)*1e9 # AMOS
  t_us_d  = (@belapsed adbesselkdv($v, $x)        samples=1_000)*1e9 # our code
  t_am_d  = (@belapsed fdbesselkdv($v, $x)        samples=1_000)*1e9 # AMOS
  t_us_d2 = (@belapsed adbesselkdvdv($v, $x)      samples=1_000)*1e9 # our code
  t_am_d2 = (@belapsed fdbesselkdvdv($v, $x)      samples=1_000)*1e9 # AMOS
  if j != length(PAIRS)
    @printf "(%1.3f, %1.0f) & %1.0f & %1.0f & %1.0f & %1.0f & %1.0f & %1.0f & %s \\\\\n" v x t_us t_am t_us_d t_am_d t_us_d2 t_am_d2 descriptor
  else
    @printf "(%1.3f, %1.0f) & %1.0f & %1.0f & %1.0f & %1.0f & %1.0f & %1.0f & %s\n" v x t_us t_am t_us_d t_am_d t_us_d2 t_am_d2 descriptor
  end
end


text:
module PandaMPC

using LinearAlgebra, Plots, ECOS, Convex, DelimitedFiles

import LinearAlgebra, Plots, ECOS, Convex, DelimitedFiles

###############################################################################
# This block contains the ball's trajectory simulation functions
"""
Given an array with the initial position and velocity of the ball [x;y;z;vx;vy;vz],\n
this function calculates the trajectory of the ball (with and without drag force)\n
during the simulation time defined by the user
"""
function ball(cond0, simtime)
    x0 = cond0[1];
    y0 = cond0[2];
    z0 = cond0[3];
    v0x = cond0[4];
    v0y = cond0[5];
    v0z = cond0[6];
    #Required ball data
    m = 0.05; # mass of the ball [kg]
    r = 0.05; # radius of the ball [meters]
    g = 9.81; # gravity
    t = range(0,stop=simtime,length=100)

    C = 0.96211*r^2/m;
    x = zeros(length(t),1); # x initialization
    y = zeros(length(t),1); # y initialization
    z = zeros(length(t),1); # y initialization
    xnodrag = zeros(length(t),1); # x initialization
    ynodrag = zeros(length(t),1); # y initialization
    znodrag = zeros(length(t),1); # y initialization

    #ODE coefficients
    # y direction (vertical)
    c1 = atan((-v0z*sqrt(C*g))/g)/sqrt(C*g);
    c2 = z0 - log(cos(c1*sqrt(C*g)))/C;
    # x direction
    c3 = 1/v0x;
    c4 = x0 - log(c3)/C;
    # z direction
    c5 = 1/v0y;
    c6 = y0 - log(c5)/C;

    for i = 1:length(t)
        x[i] = log(c3 + C*t[i])/C + c4;
        y[i] = log(c5 + C*t[i])/C + c6;
        z[i] = log(cos((c1+t[i])*sqrt(C*g)))/C + c2;

        xnodrag[i] = x0 + v0x*t[i];
        ynodrag[i] = y0 + v0y*t[i];
        znodrag[i] = z0 + v0z*t[i] - 0.5*g*t[i]^2;
    end

    posdrag = [x y z]
    posnodrag = [xnodrag ynodrag znodrag]
    return t, posdrag, posnodrag
end
"""
Given an array with the trajectory of the ball, this function calculates the\n
time remaining until the ball crosses a plane perpendicular to the x-direction\n
of the base frame of the robot. This plane will be at a distance xplane to the\n
origin of the frame. Cartesian position of this point is also provided
"""
function intercplane(t, posdrag, xplane = 0.4)
    xcapture = 0;
    ycapture = 0;
    zcapture = 0;
    tcapture = 0;
    for i in 1:length(t)
        if abs(posdrag[i,1]+xplane) < 0.01
            xcapture = -posdrag[i,1];
            ycapture = posdrag[i,2];
            zcapture = posdrag[i,3];
            tcapture = t[i];
        end
    end
    return xcapture, ycapture, zcapture, tcapture
end

################################################################################
# This block contains the inverse kinematics functions for the Franka Panda Robot
"""
Calculates the inverse kinematics of the Panda Robot. Providing the desired\n
cartesian position of the end-effector and the starting configuration, the function\n
will calculate the angles of each of the 7 joints that will achieve this.
"""
function getjointangles(posicCartes, theta0)
    theta = theta0;
    xd = posicCartes;
    M = [1 0 0 0.088; 0 -1 0 0; 0 0 -1 0.8226; 0 0 0 1];
    S = [0 0 0 0 0 0 0; 0 1 0 -1 0 -1 0; 1 0 1 0 1 0 -1; 0 -0.333 0 0.649 0 1.033 0; 0 0 0 0 0 0 0.088; 0 0 0 -0.0825 0 0 0];

    k = 1;
    while k<800
        res =  zeros(4,4,length(theta))
        for i in 1 : length(theta)
            Sw = S[1:3, i];
            Sv = S[4:6, i];
            Smat = [0 -Sw[3] Sw[2] Sv[1]; Sw[3] 0 -Sw[1] Sv[2]; -Sw[2] Sw[1] 0 Sv[3]; 0 0 0 0];
            res[:,:,i] = exp(Smat*theta[i]);
        end

        T = res[:,:,1]*res[:,:,2]*res[:,:,3]*res[:,:,4]*res[:,:,5]*res[:,:,6]*res[:,:,7]*M;

        Js = zeros(6,7)
        Adtotal = zeros(4,4,6)
        # Joint 1
        Js[:,1] = S[:,1];

        # All the other joints
        Adtotal[:,:,1] = res[:,:,1];
        Adtotal[:,:,2] = res[:,:,1]*res[:,:,2];
        Adtotal[:,:,3] = res[:,:,1]*res[:,:,2]*res[:,:,3];
        Adtotal[:,:,4] = res[:,:,1]*res[:,:,2]*res[:,:,3]*res[:,:,4];
        Adtotal[:,:,5] = res[:,:,1]*res[:,:,2]*res[:,:,3]*res[:,:,4]*res[:,:,5];
        Adtotal[:,:,6] = res[:,:,1]*res[:,:,2]*res[:,:,3]*res[:,:,4]*res[:,:,5]*res[:,:,6];

        for j = 1:size(Adtotal,3)
            Ad = Adtotal[:,:,j];
            R = Ad[1:3,1:3];
            p  = Ad[1:3,4];
            p2 = [0 -p[3] p[2]; p[3] 0 -p[1]; -p[2] p[1] 0];
            Js[:,j+1] = [R zeros(3,3); p2*R R]*S[:,j+1];
        end

        Jss = [Js[4:6,:]; Js[1:3,:]];
        Jbb = adi(T)*Jss;
        Jsb = [T[1:3,1:3] zeros(3,3); zeros(3,3) T[1:3,1:3]] * Jbb; #La buena

        e = [xd[1:3]-T[1:3,4]; 0;0;0];
        theta = theta + pinv(Jsb)*e;
        if abs(e[1])+abs(e[2])+abs(e[3]) < 0.005
            k = 5000;
        end
        k = k+1;
    end

    for i in 1:length(theta)
        theta[i] = theta[i] % (2*pi); #para evitar dar varias vueltas al eje (no afecta a T)
    end
    return theta
end

"""Calculates the adjoint of the inverse of a tranformation matrix T, which is also the inverse of the adjoint of T"""
function adi(T)
    # Computes the adjoint of transformation matrix inv(T)
    R = T[1:3,1:3]'
    t = skew(-R*T[1:3,4])
    Z = zeros(3,3)
    A = [R t*R; Z R]
    return A
end

skew(s)::AbstractMatrix{eltype(s)} = [0 -s[3] s[2];s[3] 0 -s[1]; -s[2] s[1] 0]

"""
Calculates the forward kinematics of the Panda Robot. An array with the 7 joint\n
angles must be provided so that the function calculates the cartesian position\n
of the end-effector
"""
function forwardkin(theta)
    M = [1 0 0 0.088; 0 -1 0 0; 0 0 -1 0.8226; 0 0 0 1];
    S = [0 0 0 0 0 0 0; 0 1 0 -1 0 -1 0; 1 0 1 0 1 0 -1; 0 -0.333 0 0.649 0 1.033 0; 0 0 0 0 0 0 0.088; 0 0 0 -0.0825 0 0 0];

    k = 1;

    res =  zeros(4,4,length(theta))
    for i in 1 : length(theta)
        Sw = S[1:3, i];
        Sv = S[4:6, i];
        Smat = [0 -Sw[3] Sw[2] Sv[1]; Sw[3] 0 -Sw[1] Sv[2]; -Sw[2] Sw[1] 0 Sv[3]; 0 0 0 0];
        res[:,:,i] = exp(Smat*theta[i]);
    end
    T = res[:,:,1]*res[:,:,2]*res[:,:,3]*res[:,:,4]*res[:,:,5]*res[:,:,6]*res[:,:,7]*M;

    return T
end

################################################################################
# This block contains the MPC functions
"""
Generates the matrixes Q and R for solving the problem: quadform(z,Q) + quadform(u,R)\n
    m is the number of joints considered in the robot (1, 2 or 3)
"""
function weights(weightQ, weightR, m)
    if m == 1
        Q = zeros(3,3)
        for i = [2,3]
            Q[i,i] = weightQ
        end
        R = weightR
    elseif m == 2
        Q = zeros(6,6)
        for i = [2,3,5,6]
            Q[i,i] = weightQ
        end
        R = zeros(2,2)
        for j = 1:2
            R[j,j] = weightR
        end
    elseif m == 3
        Q = zeros(9,9)
        for i = [2,3,5,6,8,9]
            Q[i,i] = weightQ
        end
        R = zeros(3,3)
        for j = 1:3
            R[j,j] = weightR
        end
        return Q, R
    elseif m == 7
        Q = zeros(21,21)
        for i = [2,3,5,6,8,9,11,12,14,15,17,18,20,21]
            Q[i,i] = weightQ
        end
        R = zeros(7,7)
        for j = 1:7
            R[j,j] = weightR
        end
        return Q, R
    else
        error("Number of joints should be 1, 2, 3 or 7")
    end
    return Q, R
end

"""
 Generates the model matrixes A, B and Dz\n
    h is the sampling period\n
    n is the number of degrees of freedom (3 per joint; 1, 2 or 3 joints)
"""
function model(h, n)
    Fitilde = [1 h h^2/2; 0 1 h; 0 0 1]
    Gamma1tilde = [h^4/24; h^3/6; h^2/2];
    Gammatilde = [h^3/6; h^2; h];
    if n == 3
        Fi = Fitilde
        Gamma1 = Gamma1tilde
        Gamma = Gammatilde
    elseif n == 6
        Fi = [Fitilde zeros(3,3); zeros(3,3) Fitilde]
        Gamma1 = [Gamma1tilde zeros(3,1); zeros(3,1) Gamma1tilde]
        Gamma = [Gammatilde zeros(3,1); zeros(3,1) Gammatilde]
    elseif n == 9
        Fi = [Fitilde zeros(3,3) zeros(3,3); zeros(3,3) Fitilde zeros(3,3); zeros(3,3) zeros(3,3) Fitilde]
        Gamma1 = [Gamma1tilde zeros(3,1) zeros(3,1); zeros(3,1) Gamma1tilde zeros(3,1); zeros(3,1) zeros(3,1) Gamma1tilde]
        Gamma = [Gammatilde zeros(3,1) zeros(3,1); zeros(3,1) Gammatilde zeros(3,1); zeros(3,1) zeros(3,1) Gammatilde]
    elseif n == 21
        Fi = [Fitilde zeros(3,3) zeros(3,3) zeros(3,3) zeros(3,3) zeros(3,3) zeros(3,3);
        zeros(3,3) Fitilde zeros(3,3) zeros(3,3) zeros(3,3) zeros(3,3) zeros(3,3);
        zeros(3,3) zeros(3,3) Fitilde zeros(3,3) zeros(3,3) zeros(3,3) zeros(3,3);
        zeros(3,3) zeros(3,3) zeros(3,3) Fitilde zeros(3,3) zeros(3,3) zeros(3,3);
        zeros(3,3) zeros(3,3) zeros(3,3) zeros(3,3) Fitilde zeros(3,3) zeros(3,3);
        zeros(3,3) zeros(3,3) zeros(3,3) zeros(3,3) zeros(3,3) Fitilde zeros(3,3);
        zeros(3,3) zeros(3,3) zeros(3,3) zeros(3,3) zeros(3,3) zeros(3,3) Fitilde]

        Gamma1 = [Gamma1tilde zeros(3,1) zeros(3,1) zeros(3,1) zeros(3,1) zeros(3,1) zeros(3,1);
        zeros(3,1) Gamma1tilde zeros(3,1) zeros(3,1) zeros(3,1) zeros(3,1) zeros(3,1);
        zeros(3,1) zeros(3,1) Gamma1tilde zeros(3,1) zeros(3,1) zeros(3,1) zeros(3,1);
        zeros(3,1) zeros(3,1) zeros(3,1) Gamma1tilde zeros(3,1) zeros(3,1) zeros(3,1);
        zeros(3,1) zeros(3,1) zeros(3,1) zeros(3,1) Gamma1tilde zeros(3,1) zeros(3,1);
        zeros(3,1) zeros(3,1) zeros(3,1) zeros(3,1) zeros(3,1) Gamma1tilde zeros(3,1);
        zeros(3,1) zeros(3,1) zeros(3,1) zeros(3,1) zeros(3,1) zeros(3,1) Gamma1tilde]

        Gamma = [Gammatilde zeros(3,1) zeros(3,1) zeros(3,1) zeros(3,1) zeros(3,1) zeros(3,1);
        zeros(3,1) Gammatilde zeros(3,1) zeros(3,1) zeros(3,1) zeros(3,1) zeros(3,1);
        zeros(3,1) zeros(3,1) Gammatilde zeros(3,1) zeros(3,1) zeros(3,1) zeros(3,1);
        zeros(3,1) zeros(3,1) zeros(3,1) Gammatilde zeros(3,1) zeros(3,1) zeros(3,1);
        zeros(3,1) zeros(3,1) zeros(3,1) zeros(3,1) Gammatilde zeros(3,1) zeros(3,1);
        zeros(3,1) zeros(3,1) zeros(3,1) zeros(3,1) zeros(3,1) Gammatilde zeros(3,1);
        zeros(3,1) zeros(3,1) zeros(3,1) zeros(3,1) zeros(3,1) zeros(3,1) Gammatilde]
    else
        error("Number of joints should be 1, 2, 3 or 7")
    end

    eye9 = zeros(n,n)
    for i = 1:n
        eye9[i,i] = 1
    end

    A = Fi
    B = Gamma + (Fi - eye9)*Gamma1/h
    Dz = Gamma1/h

    return A,B,Dz
end

"""
Robot Timescale\n
    You introduce a list of velocities (and times) with a time spacing greater than 1 milisecond\n
     and this function returns a list of velocities with a exact spacing of 1 milisecond\n
     (the gaps between the original time samples are filled by means of a linear interpolation) \n
"""
function robot_timescale(tfin, xfin2)
    for k = 1:length(tfin)
        tfin[k] = round(tfin[k]; digits=3)
        xfin2[k] = round(xfin2[k]; digits=5)
    end

    trobot = collect(tfin[1]:0.001:tfin[end])
    vrobot = [0;1]

    for i = 1:length(tfin)-1
        samples = tfin[i]:0.001:tfin[i+1]
        for j = 1:length(samples)-1
            vrobot = [vrobot; ((xfin2[i+1]-xfin2[i])/(tfin[i+1]-tfin[i]))*(samples[j] - tfin[i]) + xfin2[i]]
        end
    end
    vrobot = [vrobot[3:end]; xfin2[end]]

    return trobot, vrobot
end

"""
Simulation for 2 joints:\n\n

The idea behind this is that the robot first waits for an instruction,\n
then it executes that, waits in that position, goes back to the origin\n
and then waits for the next order.\n\n

pos1 is the array of positions of the first joint (where its first value is the starting point)\n

pos2 is the array of positions of the second joint (where its first value is the starting point)\n

span0 is the initial waiting time

span1 is the time span for getting to the goal state

wait_time is the waiting time at the goal state

span2 is the time span for getting back to the goal state from to the goal state

total_time is the total time of the iteration and in this context determines\n
the waiting time until the next iteration

"""
function iteraciones2(pos1, pos2, span0, span1, wait_time, span2, total_time)
    #Initial Pause - Awaiting instructions
    #span0 = 1
    final_z =  [pos1[1]; 0; 0; pos2[1]; 0; 0]
    initial_x = [pos1[1]; 0; 0; pos2[1]; 0; 0]
    initial_u = [0; 0]
    initial_z = [pos1[1]; 0; 0; pos2[1]; 0; 0]
    (tfin, xfin, ufin) = go(span0, final_z,initial_x, initial_u, initial_z, resample)
    timecalc = tfin
    x = xfin
    u = ufin
    for i = 2:length(pos1)
        #Instruction
        final_z =  [pos1[i]; 0; 0; pos2[i]; 0; 0]
        initial_x = [pos1[1]; 0; 0; pos2[1]; 0; 0]
        initial_u = [0; 0]
        initial_z = [pos1[1]; 0; 0; pos2[1]; 0; 0]
        (tfin, xfin, ufin) = go(span1, final_z,initial_x, initial_u, initial_z, resample)
        for i = 1:length(tfin)
            tfin[i] = tfin[i] + timecalc[end] + 0.0001
        end
        timecalc = [timecalc; tfin]
        x = [x xfin]
        u = [u ufin]

        #Waiting time in goal
        final_z =  [pos1[i]; 0; 0; pos2[i]; 0; 0]
        initial_x = [pos1[i]; 0; 0; pos2[i]; 0; 0]
        initial_u = [0; 0]
        initial_z = [pos1[i]; 0; 0; pos2[i]; 0; 0]
        (tfin, xfin, ufin) = go(wait_time, final_z,initial_x, initial_u, initial_z, resample)
        for i = 1:length(tfin)
            tfin[i] = tfin[i] + timecalc[end] + 0.0001
        end
        timecalc = [timecalc; tfin]
        x = [x xfin]
        u = [u ufin]

        #Back to origin
        final_z =  [pos1[1]; 0; 0; pos2[1]; 0; 0]
        initial_x = [pos1[i]; 0; 0; pos2[i]; 0; 0]
        initial_u = [0; 0]
        initial_z = [pos1[i]; 0; 0; pos2[i]; 0; 0]
        (tfin, xfin, ufin) = go(span2, final_z,initial_x, initial_u, initial_z, resample)
        for i = 1:length(tfin)
            tfin[i] = tfin[i] + timecalc[end] + 0.0001
        end
        timecalc = [timecalc; tfin]
        x = [x xfin]
        u = [u ufin]

        #Waiting next instruction
        span3 = total_time - span1 - span2
        final_z =  [pos1[1]; 0; 0; pos2[1]; 0; 0]
        initial_x = [pos1[1]; 0; 0; pos2[1]; 0; 0]
        initial_u = [0; 0]
        initial_z = [pos1[1]; 0; 0; pos2[1]; 0; 0]
        (tfin, xfin, ufin) = go(span3, final_z,initial_x, initial_u, initial_z, resample)
        for i = 1:length(tfin)
            tfin[i] = tfin[i] + timecalc[end] + 0.0001
        end
        timecalc = [timecalc; tfin]
        x = [x xfin]
        u = [u ufin]
    end
    return timecalc, x, u
end

"""
Plotting for position, velocity, acceleration and input (jerk) of one of the two controlled joints\n
    joint should be 1 or 2

"""
function plotjoint(timecalc, u, x, joint)
    if joint == 1
        plot(timecalc,[x[1,:],x[2,:],x[3,:],u[1,:]], layout = 4, legend = false)
    elseif joint == 2
        plot(timecalc,[x[4,:],x[5,:],x[6,:],u[2,:]], layout = 4, legend = false)
    elseif joint == 3
        plot(timecalc,[x[7,:],x[8,:],x[9,:],u[3,:]], layout = 4, legend = false)
    elseif joint == 4
        plot(timecalc,[x[10,:],x[11,:],x[12,:],u[4,:]], layout = 4, legend = false)
    elseif joint == 5
        plot(timecalc,[x[13,:],x[14,:],x[15,:],u[5,:]], layout = 4, legend = false)
    elseif joint == 6
        plot(timecalc,[x[16,:],x[17,:],x[18,:],u[6,:]], layout = 4, legend = false)
    elseif joint == 7
        plot(timecalc,[x[19,:],x[20,:],x[21,:],u[7,:]], layout = 4, legend = false)
    else
        error("Choose a valid joint")
    end
end

"""
MPC generation\n
    span is the calculation time\n
    resample lets you set how often does the MPC recalculate\n
    final_z is the final state\n
    initial_x, initial_u, initial_z are the initial conditions\n\n

    Example - 2 joints, 6 degrees of freedom and time span of 1 second and recalculated every 20ms:\n
    span = 1\n
    resample = 0.02\n
    final_z =  [1; 0; 0; 1; 0; 0]\n
    initial_x = [0; 0; 0; 0; 0; 0]\n
    initial_u = [0; 0]\n
    initial_z = [0; 0; 0; 0; 0; 0]
"""
function go7(span, final_z, initial_x, initial_u, initial_z, resample, T, n, l, m, weightQ, weightR)
    recalc_times = collect(0:resample:span-resample)
    tfin = [0;1] #only for creating these variables
    zfin = [0 1; 0 1; 0 1; 0 1; 0 1; 0 1; 0 1; 0 1; 0 1; 0 1; 0 1; 0 1; 0 1; 0 1; 0 1 ; 0 1; 0 1; 0 1; 0 1; 0 1; 0 1]
    ufin = [0 1; 0 1; 0 1; 0 1; 0 1; 0 1; 0 1]
    for k = 1:length(recalc_times)

        tsim = span - recalc_times[k] #Simulation time
        h = tsim/(T-1)
        (A, B, Dz) = model(h, n)

        # Create a optimization variables of size n x 1.
        x = Variable(n, T)
        u = Variable(m, T)
        z = Variable(l, T)

        # Create a problem instance
        (Q,R) = weights(weightQ, weightR, m);
        problem = minimize(quadform(z,Q) + quadform(u,R))

        # Equality constraints
        problem.constraints += x[:,1] == initial_x
        problem.constraints += u[:,1] == initial_u
        problem.constraints += z[:,1] == initial_z
        problem.constraints += z[:,T] == final_z

        for i in 1 : T -1
          problem.constraints += u[:,i+1] <= [5000; 3000; 4000; 4000; 5000; 7000; 7000]
          problem.constraints += u[:,i+1] >= [-5000; -3000; -4000; -4000; -5000; -7000; -7000]
          problem.constraints += z[:,i+1] <= [2.5; 2; 13; 1.5; 2; 6; 2.5; 2; 8; -0.07; 2; 10; 2.5; 2.5; 13; 3.5; 2.5; 18; 2.5; 2.5; 18]
          problem.constraints += z[:,i+1] >= [-2.5; -2; -13; -1.5; -2; -6; -2.5; -2; -8; -3; -2; -10; -2.5; -2.5; -13; 0; -2.5; -18; -2.5; -2.5; -18]
          problem.constraints += x[:,i+1] == A*x[:, i] + B* u[:, i]
          problem.constraints += z[:,i+1] == x[:,i+1] + Dz* u[:,i+1]
        end

        # Solve the problem by calling solve!
        solve!(problem, ECOSSolver())

        # Check the status of the problem
       problem.status # :Optimal, :Infeasible, :Unbounded etc.

        # Get the optimum value
        problem.optval

        # global itera, tbien
        itera = 0
        t3 = collect(recalc_times[k]:h:span)

        for i = 1:length(t3)
            if k < length(recalc_times)
                if t3[i] <= resample*k
                    tfin = [tfin; t3[i]]
                    zfin = [zfin z.value[:,i]]
                    ufin = [ufin u.value[:,i]]
                    itera = i
                end
                tbien = t3[1:itera]
                if tbien[end] == resample*k
                    initial_x = x.value[:,itera]
                    initial_u = u.value[:,itera]
                    initial_z = z.value[:,itera]
                else #Lineal interpolation
                    initial_x = ((x.value[:,itera+1]-x.value[:,itera])/(t3[itera+1]-t3[itera]))*(resample*k-t3[itera]) + x.value[:,itera]
                    initial_u = ((u.value[:,itera+1]-u.value[:,itera])/(t3[itera+1]-t3[itera]))*(resample*k-t3[itera]) + u.value[:,itera]
                    initial_z = ((z.value[:,itera+1]-z.value[:,itera])/(t3[itera+1]-t3[itera]))*(resample*k-t3[itera]) + z.value[:,itera]
                end
            else
                tfin = [tfin; t3[i]]
                zfin = [zfin z.value[:,i]]
                ufin = [ufin u.value[:,i]]
                itera = i
            end
        end

    end

    tfin = tfin[3:end]
    zfin = zfin[:,3:end]
    ufin = ufin[:,3:end]
    return tfin, zfin, ufin
end

export ball, intercplane, getjointangles, adi, forwardkin, skew, weights, model, robot_timescale, go, go7, iteraciones2, plotjoint
end


text:
struct Actuator
    name::String
    jointid::JointID
end

function parse_actuators(mechanism::Mechanism{T}, urdffile::String) where T
    xdoc = parse_file(urdffile)
    xroot = LightXML.root(xdoc)
    @assert LightXML.name(xroot) == "robot"

    xml_transmissions = get_elements_by_tagname(xroot, "transmission")
    actuators = Actuator[]
    for joint in tree_joints(mechanism) # order matters
        num_velocities(joint) == 0 && continue
        for xml_transmission in xml_transmissions
            xml_jointref = find_element(xml_transmission, "joint")
            if string(joint) == attribute(xml_jointref, "name")
                xml_actuator = find_element(xml_transmission, "actuator")
                xml_reduction = find_element(xml_actuator, "mechanicalReduction")
                xml_reduction == nothing || @assert parse(Int, content(xml_reduction)) == 1
                push!(actuators, Actuator(attribute(xml_actuator, "name"), JointID(joint)))
            end
        end
    end

    actuators
end

struct HumanoidRobotInfo{T}
    mechanism::Mechanism{T}
    feet::Dict{Side, RigidBody{T}}
    hands::Dict{Side, RigidBody{T}}
    floatingjoint::Joint{T, QuaternionFloating{T}}
    floatingbody::RigidBody{T}
    world_aligned_floating_body_frame::CartesianFrame3D
    actuators::Vector{Actuator}
    revolutejoints::Vector{Joint{T, Revolute{T}}}

    function HumanoidRobotInfo(
            mechanism::Mechanism{T},
            feet::AbstractDict{Side, RigidBody{T}},
            hands::AbstractDict{Side, RigidBody{T}},
            actuators::Vector{Actuator}) where {T}
        sides = instances(Side)
        @assert isempty(setdiff(keys(feet), sides))
        @assert isempty(setdiff(keys(hands), sides))
        floating_joints = filter(isfloating, tree_joints(mechanism))
        @assert length(floating_joints) == 1
        floatingjoint = first(floating_joints)
        floatingbody = successor(floatingjoint, mechanism)
        world_aligned_floating_body_frame = CartesianFrame3D("world_aligned_floating_body_frame")
        revolutejoints = Joint{T, Revolute{T}}[j for j in tree_joints(mechanism) if joint_type(j) isa Revolute{T}]
        new{T}(mechanism, feet, hands, floatingjoint, floatingbody, world_aligned_floating_body_frame, actuators, revolutejoints)
    end
end

num_actuators(info::HumanoidRobotInfo) = length(info.actuators)
actuators(info::HumanoidRobotInfo) = info.actuators

function findactuator(info::HumanoidRobotInfo, name::String)
    for actuator in info.actuators
        actuator.name == name && return actuator
    end
    throw(KeyError(name))
end