# The goal is to estimate the fractal dimension D of a curve by box counting
# algorith as N = eps^{-D}, where N is the number of boxes intersected
# by the curve at this level of refinement eps.

"""A hyper prism"""
struct BoundingBox{D, T}
    low::SVector{D, T}
    high::SVector{D, T}
end

"""Physical size of the bounding box; D vector of sizes"""
size{D, T}(box::BoundingBox{D, T}) = box.high - box.low

"""Length of bounding box in the i-th axis"""
size{D, T}(box::BoundingBox{D, T}, i::Int) = size(box)[i]

"""One refinement of a 2d box results in 4 boxes"""
function split{T}(box::BoundingBox{2, T})
    X, Y = box.low, box.high
    M = 0.5(X + Y)
    [BoundingBox(X, M),
     BoundingBox(SVector{2, T}(X[1], M[2]), SVector{2, T}(M[1], Y[2])),
     BoundingBox(SVector{2, T}(M[1], X[2]), SVector{2, T}(Y[1], M[2])),
     BoundingBox(M, Y)]
end

"""Bbox of segment"""
BoundingBox{D, T}(line::Segment{D, T}) = BoundingBox(SVector(Tuple(min(line.A[i], line.B[i]) for i in 1:D)),
                                                     SVector(Tuple(max(line.A[i], line.B[i]) for i in 1:D)))

"""Bbox of collection of segments"""
function BoundingBox{D, T}(segments::Vector{Segment{D, T}})

    low = SVector(Tuple(minimum(min(seg.A[i], seg.B[i]) for seg in segments) for i in 1:D))

    high = SVector(Tuple(maximum(max(seg.A[i], seg.B[i]) for seg in segments) for i in 1:D))

    return BoundingBox(low, high)
end

"""Bbox of a curve"""
function BoundingBox{D, T}(curve::Curve{D, T})
    segs = Vector{Segment{D, T}}(ne(curve))
    for (i, segment) in enumerate(iter_segments(curve))
        segs[i] = segment
    end
    BoundingBox(segs)
end

"""Does line collide with a box?"""
# NOTE: Liang-Barsky algorithm
function collides{D, T}(line::Segment{D, T}, box::BoundingBox{D, T})
    x, y = line.A
    dx = line.B[1] - x
    dy = line.B[2] - y

    p = [-dx, dx, -dy, dy]
    q = [x-box.low[1], box.high[1]-x, y-box.low[2], box.high[2]-y]

    u1 = -Inf
    u2 = Inf
    for i in 1:4
        (p[i] == 0 && q[i] < 0) && return false

        t = q[i]/p[i]

        if (p[i] < 0 && u1 < t)
            u1 = max(0, t)
        elseif (p[i] > 0 && u2 > t)
            u2 = min(1, t)
        end
    end

   (u1 > u2) && return false

    true
end

"""
Box counter holds after each refinement a pair of curve and box which
might intersect it. The curve is a collection if indices which refer
to segments of the unsplit curve.
"""
struct BoxCounter{D, T}
    segments::Vector{Segment{D, T}}
    seg_box::Deque{Pair{Vector{Int}, BoundingBox{D, T}}}
    stop_length::Float64
end

"""Box size at the current level"""
size{D, T}(bc::BoxCounter{D, T}) = size(first(bc.seg_box).second)

"""Construct BoxCounter for the curve"""
function BoxCounter{D, T}(curve::Curve{D, T})
    segs = Vector{Segment{D, T}}(ne(curve))
    for (i, segment) in enumerate(iter_segments(curve))
        segs[i] = segment
    end
    box = BoundingBox(segs)
    # Make the box slightly larger so that we don't hit corners
    X, Y = box.low, box.high
    center = (X+Y)/2.

    # Increaze the box a bit in an attempt/delay point isects
    # NOTE: as a result of this each run gives a slighly different
    # number - ideally avarage several runs
    shift1 = SVector(((Y - center).*(1. + rand(D)))...)
    shift2 = SVector(((Y - center).*(1. + rand(D)))...)

    box = BoundingBox(center - shift1, center + shift2)

    seg_box = Deque{Pair{Vector{Int}, BoundingBox{D, T}}}()
    push!(seg_box, Pair(collect(1:length(segs)), box))

    BoxCounter(segs, seg_box, min_segment(curve)/2)
end

# One box intersected. That box is bounding box
start{D, T}(bc::BoxCounter{D, T}) = (1, size(bc))

# Refine and count
function next{D, T}(bc::BoxCounter{D, T}, state)

    count = 0
    # Go over all the boxes on the current level
    for i in 1:length(bc.seg_box)

        index_segments, boxes = shift!(bc.seg_box)
        # Are children intersected
        for box in split(boxes)
            isected_segments = Vector{Int}()
            # Collect isects
            for index in index_segments
                collides(bc.segments[index], box) && push!(isected_segments, index)
            end
            # Increate count if there are and send for next round
            if !isempty(isected_segments)
                count += 1
                push!(bc.seg_box, Pair(isected_segments, box))
            end
        end
    end
    # We have halved the mesh
    (count, last(state)/2.)
end

# FIXME: It doesn't make sence to go much beyond the resolution of the curve
done{D, T}(c::BoxCounter{D, T}, state) = true #minimum(last(state)) < c.stop_length

"""Char size of the hyper box - best fit by cuve"""
volume{D, T}(p::SVector{D, T}) = (prod(p))^(1./D)

"""Remember volumes and counts as the refinement is done with box counter"""
mutable struct FractalDim{D, T}
    bc::BoxCounter{D, T}

    volumes::Vector{Float64}
    counts::Vector{Int}
    next_state::Tuple{Int, SVector{D, T}}
end

"""Estimate fd by lin fit of the stepped FractalDim"""
function estimate_fd{D, T}(fd::FractalDim{D, T}, slice=3)
    @assert !isempty(fd.volumes) && !isempty(fd.counts)

    if length(fd.counts) < slice
        x = log.(fd.volumes)
        y = log.(fd.counts)
    else
        x = log.(fd.volumes[end-slice:end])
        y = log.(fd.counts[end-slice:end])
    end
    -last(Polynomials.polyfit(x, y, 1).a)
end

"""Curve constructor"""
function FractalDim{D, T}(c::Curve{D, T})
    bc = BoxCounter(c)
    volumes = Vector{Float64}()
    counts = Vector{Int}()
    next_state = (1, size(bc))

    FractalDim(bc, volumes, counts, next_state)
end

"""Start stepping"""
function start{D, T}(fd::FractalDim{D, T})
    state = start(fd.bc)

    push!(fd.counts, first(state))
    push!(fd.volumes, volume(last(state)))
    fd.next_state = state

    state
end

"""Evolve"""
function next{D, T}(fd::FractalDim{D, T}, state)
    state = next(fd.bc, state)

    push!(fd.counts, first(state))
    push!(fd.volumes, volume(last(state)))
    fd.next_state = state

    state
end

"""Let box counter decide"""
done{D, T}(fd::FractalDim{D, T}, state) = done(fd.bc)

"""Evolve the started iterator nlevel times"""
function fractal_dim!{D, T}(fd::FractalDim{D, T}, nlevels::Int)
    state = fd.next_state
    for i in 1:nlevels
        state = next(fd, state)
    end
    (estimate_fd(fd), fd)
end

"""
Step the estimator for curve returning the counter which can be
further evolved if the precision is not okay
"""
fractal_dim{D, T}(c::Curve{D, T}, nlevels::Int) = fractal_dim!(FractalDim(c), nlevels)


typealias SubPaving{N,T} Vector{IntervalBox{N,T}}

type Paving{N,T}
    separator::Separator
    inner::SubPaving{N,T}
    boundary::SubPaving{N,T}
    ϵ::Float64
end

function setdiff{N,T}(x::IntervalBox{N,T}, subpaving::SubPaving{N,T})
    working = [x]
    new_working = IntervalBox{N,T}[]

    local have_split

    for y in subpaving
        for x in working
            have_split = false

            diff = setdiff(x, y)

            if diff != [x]
                have_split = true
                append!(new_working, diff)
            end
        end

        !have_split && push!(new_working, x)

        working = new_working
        new_working = IntervalBox{N,T}[]

    end

    return working

end

setdiff{N,T}(X::SubPaving{N,T}, Y::SubPaving{N,T}) = vcat([setdiff(x, Y) for x in X]...)

function setdiff{N,T}(x::IntervalBox{N,T}, paving::Paving{N,T})
    Y = setdiff(x, paving.inner)
    Z = setdiff(Y, paving.boundary)
    return Z
end

function show{N,T}(io::IO, p::Paving{N,T})
    print(io, """Paving:
                 - tolerance ϵ = $(p.ϵ)
                 - inner approx. of length $(length(p.inner))
                 - boundary approx. of length $(length(p.boundary))"""
              )
end


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

sm = @stochastic_model begin
    @stage 1 begin
        @decision(model, x¹ <= 2)
        @decision(model, w¹)
        @decision(model, y¹)
        @objective(model, Min, 100*x¹ + 3*w¹ + 0.5*y¹)
        @constraint(model, x¹ + w¹ - y¹ == 1)
    end
    @stage 2 begin
        @uncertain d
        @variable(model, x² <= 2)
        @variable(model, w²)
        @decision(model, y²)
        @objective(model, Min, x² + 3*w² + 0.5*y²)
        @constraint(model, y¹ + x² + w² - y² == d)
    end
    @stage 3 begin
        @uncertain d
        @variable(model, x³ <= 2)
        @variable(model, w³)
        @recourse(model, y³)
        @objective(model, Min, x³ + 3*w³)
        @constraint(model, y² + x³ + w³ - y³ == d)
    end
end
s₂₁ = @scenario d = 1.0 probability = 0.5
s₂₂ = @scenario d = 3.0 probability = 0.5

s₃₁ = @scenario d = 1.0 probability = 0.25
s₃₂ = @scenario d = 3.0 probability = 0.25
s₃₃ = @scenario d = 1.0 probability = 0.25
s₃₄ = @scenario d = 3.0 probability = 0.25

sp = instantiate(sm, ([s₂₁,s₂₂], [s₃₁,s₃₂,s₃₃,s₃₄]))


# It seems like the individual flashes aren't bright enough to navigate.
# However, you might have a better option: the flashes seem to be
# synchronizing!
#
# In the example above, the first time all octopuses flash simultaneously is
# step 195:
#
# After step 193:
#
# 5877777777
# 8877777777
# 7777777777
# 7777777777
# 7777777777
# 7777777777
# 7777777777
# 7777777777
# 7777777777
# 7777777777
#
# After step 194:
# 6988888888
# 9988888888
# 8888888888
# 8888888888
# 8888888888
# 8888888888
# 8888888888
# 8888888888
# 8888888888
# 8888888888
#
# After step 195:
# 0000000000
# 0000000000
# 0000000000
# 0000000000
# 0000000000
# 0000000000
# 0000000000
# 0000000000
# 0000000000
# 0000000000
#
# If you can calculate the exact moments when the octopuses will all flash
# simultaneously, you should be able to navigate through the cavern. What is
# the first step during which all octopuses flash?

# load grid
grid = hcat(map(eachline("input.txt")) do l
  parse.(Int8, split(l, ""))
end...)

# simulate until they all flash
for i in Iterators.countfrom()
  # increment the energy of all octopuses
  grid .+= 1

  # find flashes
  all_flashes = falses(10, 10)
  flashes = grid .> 9
  while any(flashes)
    # increment the energy of the octopuses around flashing octopuses
    @views begin
      grid[1:9, 1:9] .+= flashes[2:10, 2:10]
      grid[1:10, 1:9] .+= flashes[1:10, 2:10]
      grid[2:10, 1:9] .+= flashes[1:9, 2:10]
      grid[2:10, 1:10] .+= flashes[1:9, 1:10]
      grid[2:10, 2:10] .+= flashes[1:9, 1:9]
      grid[1:10, 2:10] .+= flashes[1:10, 1:9]
      grid[1:9, 2:10] .+= flashes[2:10, 1:9]
      grid[1:9, 1:10] .+= flashes[2:10, 1:10]
    end

    # update all_flashes and see if there are any cascading flashes
    all_flashes = all_flashes .|| flashes
    flashes = grid .> 9 .&& .!all_flashes
  end

  # if they've all flashed, we're done
  if all(all_flashes)
    println("All flash at: ", i)
    break
  end

  # set any flashing octopuses back to zero
  grid[all_flashes] .= 0
end


function summaryRanges(nums::Array{Number})
    length(nums) == 0 && return []

    start, endding = nums[1], nums[1]
    result::Array{String} = []

    for i in nums
        i > endding+1 &&
            begin
                start != endding ?
                    push!(result, string(start, "->", endding)) :
                    push!(result, string(start))
                start = i
            end
        endding = i
    end

    start != endding ?
        push!(result, string(start, "->", endding)) :
        push!(result, string(start))

    result
end

using Test
@test summaryRanges([0, 1, 2, 4, 5, 7]) == ["0->2", "4->5", "7"]
@test summaryRanges([0, 2, 3, 4, 6, 8, 9]) == ["0", "2->4", "6", "8->9"]


# Functions for representing graphs with GraphPlot package
# https://github.com/JuliaGraphs/GraphPlot.jl

function plot_network(netw1::SparseMatrixCSC;
    layout = spring_layout,
    label = true,
    edgelinewidth = 1.0,
    color = Colors.parse(Colorant, "turquoise")
)

    directed = !is_symmetric(netw1)

    n = netw1.m
    swn = SimpleWeightedDiGraph(netw1')

    p = gplot(swn,
        layout = layout,
        NODESIZE = 0.20 / sqrt(n),
        nodefillc = color,
        nodelabel = 1:n,
        nodelabeldist = 0,
        nodelabelangleoffset = π / 4,
        NODELABELSIZE = (n < 50) ? 5 : 2,
        edgelabel = netw1.nzval,
        EDGELABELSIZE = label ? 3 : 0,
        edgestrokec = Colors.parse(Colorant, "gray74"),
        edgelinewidth = edgelinewidth,
        EDGELINEWIDTH = 1.7 / sqrt(n),
        arrowlengthfrac = (directed == false) ? 0.0 : maximum([0.028, 0.3 / length(netw1.nzval)]),
        arrowangleoffset = (π / 9),
    )
    return p
end

function plot_network(netw1::SimpleWeightedDiGraph;
    layout = spring_layout,
    label = true,
    edgelinewidth = 1.0,
    color = Colors.parse(Colorant, "turquoise"))
    return plot_network(netw1.weights; layout = layout, label = label,
        edgelinewidth = edgelinewidth, color = color)
end

function plot_network(netw1::SimpleWeightedGraph;
    layout = spring_layout,
    label = true,
    edgelinewidth = 1.0,
    color = Colors.parse(Colorant, "turquoise"))
    return plot_network(netw1.weights; layout = layout, label = label,
        edgelinewidth = edgelinewidth, color = color)
end

function plot_graph(g::SimpleGraph; color = false)
    # transform in a Simple Weighted Graph
    sparse_mat = adjacency_matrix(g)

    if color
        mat = SimpleWeightedGraph(sparse_mat)
        colors = palette(:default, sparse_mat.m)
        comp = dfs_find_components(mat)
        plot_network(sparse_mat; label = false, color = colors[comp])
    else
        plot_network(sparse_mat; label = false)
    end

end

function plot_graph(g::SimpleDiGraph; color = false)
    sparse_mat = adjacency_matrix(g)

    if color
        mat = SimpleWeightedDiGraph(sparse_mat')
        colors = palette(:default, sparse_mat.m)
        comp = dfs_find_components(mat)
        plot_network(sparse_mat; label = false, color = colors[comp])
    else
        plot_network(sparse_mat; label = false)
    end

end


#function show(io::IO, M::ResidueRingPolyMap)
#  println(io, "Map $(domain(M)) -> $(codomain(M)): gen -> $(M.gen_image)")
#  if isdefined(M, :coeff_map)
#    println(io, "using coeff map: ", M.coeff_map)
#  end
#end
#
#function ResidueRingPolyMap_image(M::ResidueRingPolyMap{D, C}, a::Generic.Res{T}) where {D, C, T}
#  #a should be in the domain of M...
#  if isdefined(M, :coeff_map)
#    I = codomain(M)(0)
#    for i=degree(a.data):-1:0
#      I = I*M.gen_image + M.coeff_map(coeff(a.data, i))
#    end
#  else
#    I = codomain(M)(0)
#    for i=degree(a.data):-1:0
#      I = I*M.gen_image + coeff(a.data, i)
#    end
#  end
#  return I
#end
#
#function ResidueRingPolyMap_preimage(M::ResidueRingPolyMap{D, C}, a::Generic.Res{T}) where {D, C, T}
#  #a should be in the codomain of M...
#  #
#  # a good 1st attempt. If many preimages are needed, the matrix Mt
#  # and the rref or whatever solve is used should be cached and reused.
#  #
#  R = codomain(M)
#  parent(a) == R || throw("mixed rings in preimage")
#  g = gens(domain(M))
#  im_gen = map(x->M(x), g) ## possibly should be cached and stored
#  ## need to write the elements in a matrix, solve the eqn for a
#  Mt = zero_matrix(base_ring(base_ring(R)), degree(R.modulus), degree(R.modulus))
#  for i=1:degree(R.modulus)
#    elem_to_mat_row!(Mt, i, im_gen[i])
#  end
#  b = zero_matrix(base_ring(base_ring(R)), 1, degree(R.modulus))
#  elem_to_mat_row!(b, 1, a)
#  s = solve_rational(Mt', b') # why, oh why is solve operating on columns????
#  if isa(s, Tuple) ## again, why, oh why is solve doing things differently
#                   ## over rings than fields?
#    s = s[1] * inv(s[2]) # all rings here (type) are actually fields (math)
#  end
#  ## need a version returning a bool to test is solution exists
#  ## without an error!
#  if isdefined(M, :coeff_map)
#    s = sum([preimage(M.coeff_map, s[i,1])*g[i] for i=1:length(g)])
#  else
#    s = sum([s[i,1]*g[i] for i=1:length(g)])
#  end
#  @assert parent(s) == domain(M)
#  return s
#end
#


using InducingPoints, KernelFunctions
using BenchmarkTools
using Plots
plotlyjs()

nv = [2, 4, 8, 16]
nnv = length(nv)
ker = SqExponentialKernel()

sizeug = zeros(nnv)
sizeexpgrid = zeros(nnv)

kermtimeug = zeros(nnv)
kermtimegrid = zeros(nnv)
kermmemug = zeros(nnv)
kermmemgrid = zeros(nnv)

sumtimeug = zeros(nnv)
sumtimegrid = zeros(nnv)
summemug = zeros(nnv)
summemgrid = zeros(nnv)

for (i, n) in enumerate(nv)
    iter = LinRange(0.0, 10.0, n)
    proditer = Iterators.product(iter, iter)
    ug = UniformGrid(proditer)
    expgrid = ug[:]
    println(length(expgrid))

    sizeug[i] = sizeof(ug)
    sizeexpgrid[i] = sizeof(expgrid)

    r1 = @benchmark kernelmatrix($ker, $ug)
    kermtimeug[i] = median(r1.times)
    kermmemug[i] = r1.memory
    r2 = @benchmark kernelmatrix($ker, $expgrid)
    kermtimegrid[i] = median(r2.times)
    kermmemgrid[i] = r2.memory

    rs1 = @benchmark sum.($ug)
    sumtimeug[i] = median(rs1.times)
    summemug[i] = rs1.memory
    rs2 = @benchmark sum.($expgrid)
    sumtimegrid[i] = median(rs2.times)
    summemgrid[i] = rs2.memory
end

elbs = ["" ""]
p1 = plot()
plot!(
    p1,
    nv,
    [sizeug, sizeexpgrid] ./ 1e3;
    label=["UniformGrid" "Explicit Grid"],
    xticks=nv,
    linewidth=3,
    legend=:topleft,
    xlabel="Nr Grid Points",
    ylabel="Object Size [kB]",
)

p2 = plot()
plot!(
    nv,
    [kermtimeug, kermtimegrid] ./ 1e3;
    yaxis=:log,
    xaxis=:log,
    label=elbs,
    title="Kernelmatrix",
    ylabel="Compute Time [μs]",
    left_margin=5.0Plots.Measures.mm,
    xticks=(nv, nv),
    linewidth=3,
)

p3 = plot()
plot!(
    nv,
    [sumtimeug, sumtimegrid] ./ 1e3;
    yaxis=:log,
    xaxis=:log,
    label=elbs,
    title="Broadcast",
    xticks=(nv, nv),
    linewidth=3,
    left_margin=-2.0Plots.Measures.mm,
)

p4 = plot()
plot!(
    p4,
    nv,
    [kermmemug, kermmemgrid] ./ 1e3;
    yaxis=:log,
    xaxis=:log,
    label=elbs,
    xlabel="Nr Grid Points",
    ylabel="Memory [kB]",
    left_margin=5.0Plots.Measures.mm,
    xticks=(nv, nv),
    linewidth=3,
)

p5 = plot()
plot!(
    p5,
    nv,
    [summemug, summemgrid] ./ 1e3;
    yaxis=:log,
    xaxis=:log,
    label=elbs,
    xlabel="Nr of Grid Points",
    left_margin=-2.0Plots.Measures.mm,
    xticks=(nv, nv),
    linewidth=3,
)

pt = plot(p2, p3)
pm = plot(p4, p5)
pR = plot(pt, pm; layout=(2, 1))

p = plot(
    p1, pR; size=(700, 400), margin=3Plots.Measures.mm, layout=grid(1, 2; widths=[0.3, 0.7])
)

savefig(p, "../src/assets/UniformGrid_bench.svg")


function parse_sequence(str::AbstractString)
    sequence = split.(split(str, ";"), ",")
    sequence = [uppercase(join(first.(s))) for s in sequence]
    return join(sequence, "#")
end

function parse_strokes(str::AbstractString)
    lines = filter(x -> x[1] != '#', split(str))
    sequence = parse_sequence(lines[1])
    data = map(x -> parse.(Float64, split(x, ",")), lines[2:end])
    return sequence, reduce(hcat, data)
end

function load_strokes(fn::AbstractString; transform=false, jitter=false)
    sequence, data = read(fn, String) |> parse_strokes
    transform && transform_strokes!(data)
    jitter && add_jitter!(data)
    return sequence, data
end

function transform_strokes!(data::AbstractMatrix;
                            invert_x::Bool=false, invert_y::Bool=true,
                            in_scale::Real=176.389, out_scale::Real=5.0)
    data .*= (out_scale / in_scale)
    if invert_x data[1,:] .*= -1 end
    if invert_y data[2,:] .*= -1 end
    min_bounds, max_bounds = minimum(data, dims=2), maximum(data, dims=2)
    mid_x, mid_y = 0.5*(min_bounds[1] + max_bounds[1]), 0.5*(min_bounds[2] + max_bounds[2])
    data[1,:] .+= (out_scale / 2 - mid_x)
    data[2,:] .+= (out_scale / 2 - mid_y)
    return data
end

function add_jitter!(data::AbstractMatrix; x_std::Real=0.05, y_std::Real=0.05)
    data[1,:] .+= (randn(size(data)[2]) .* x_std)
    data[2,:] .+= (randn(size(data)[2]) .* y_std)
    return data
end

function load_dataset(dir::AbstractString; transform=true, jitter=false)
    fns = readdir(dir)
    fns = filter(s -> match(r".*\.txt", s) !== nothing, fns)
    dataset = load_strokes.(joinpath.(dir, fns); transform=transform, jitter=jitter)
    return dataset
end


module TestEspy
using EspyInsideFunction,Test#,StaticArrays

## Test request

struct Element  end
requestable(el::Element)= (X=(3,2),  gp=forloop(2, (F=(2,2), material=(ε=(2,2),Σ=(2,2)))) )
el = Element()
r8 = @request X,gp[].(F,material.(Σ,ε))

ra    = requestable(el)

@testset "requestable" begin
    @test ra.X == (3,2)
    @test ra.gp.body == (F = (2, 2), material = (ε = (2, 2), Σ = (2, 2)))
end

@testset "Makekey" begin
    @test EspyInsideFunction.makekey_symbol(2,(2,3)) == ([3 5 7; 4 6 8], 8)
    @test EspyInsideFunction.makekey_tuple(2,:((a,b)),(a=(2,3),b=(2,2))) == ((a=[3 5 7; 4 6 8], b=[9 11; 10 12]), 12)
    @test EspyInsideFunction.makekey_tuple(2,:((a,b)),(a=(2,3),b=scalar)) == ((a=[3 5 7; 4 6 8], b=9), 9)
    @test EspyInsideFunction.makekey_loop(2,:((a,b)),forloop(3,(a=(2,3),b=(2,2)))) == ([(a=[3 5 7; 4 6 8], b=[9 11; 10 12]), (a=[13 15 17; 14 16 18], b=[19 21; 20 22]), (a=[23 25 27; 24 26 28], b=[29 31; 30 32])], 32)
    @test EspyInsideFunction.makekey_tuple(2,:((a[].(b,)),),(a=forloop(2,(b=(2,3),)),)) == (( a=[(b=[3 5 7; 4 6 8],), (b=[9 11 13; 10 12 14],)] ,), 14)
    @test EspyInsideFunction.makekey_tuple(2,:(a[].b),(a=forloop(2,(b=(2,3),)),)) == ((a=[(b=[3 5 7; 4 6 8],), (b=[9 11 13; 10 12 14],)],), 14)
    @test makekey(r8,ra) == ((X = [1 4; 2 5; 3 6], gp = [(F = [7 9; 8 10], material = (Σ = [11 13; 12 14], ε = [15 17; 16 18])), (F = [19 21; 20 22], material = (Σ = [23 25; 24 26], ε = [27 29; 28 30]))]), 30)
end

## annotated code (written by user)

struct El  end
requestable(el::El)= (gp=forloop(2, (z=scalar,s=scalar, material=(a=scalar,b=scalar))),)
el = El()

@espy function residual(x,y)
    ngp=2
    r = 0
    for igp=1:ngp
        :z = x[igp]+y[igp]
        :s,dum  = :material(z)
        r += s
    end
    return r
end
@espy function material(z)
    :a = z+1
    :b = a*z
    return b,3.
end

## macro'ed code (returned by @espy)
# function residual(out,key,x,y)
#     ngp = 2
#     r   = 0
#     for igp = 1:ngp
#         @espy_loop key gp                     # key_gp = key.gp[igp]
#         z = x[igp]+y[igp]
#         @espy_record out key_gp z             # out[key_gp.z] = z
#         s = @espy_call out key_gp material(z) # s = material(out,key_gp.material,z)
#         @espy_record out key_gp s             # out[key_gp.s] = s
#         r += s
#     end
#     return r
# end
# function material(out,key,z)
#     a = z+1
#     @espy_record out key a                    # out[key.a] = a
#     b = a*z
#     @espy_record out key b                    # out[key.b] = b
#     return b
# end
## final code (after @espy_call, @espy_loop and @espy_record are run)
# function residual(out,key,x,y) #
#     ngp = 2
#     r = 0
#     for igp = 1:ngp
#         key_gp = key.gp[igp] #
#         z = x[igp]+y[igp]
#         out[key_gp.z] = z #
#         s  = material(out,key_gp.material,z) #
#         out[key_gp.s] = s #
#         r += s
#     end
#     return r
# end
# function material(out,key,z) #
#     a = z+1
#     out[key.a] = a #
#     b = a*z
#     out[key.b] = b #
#     return b
# end

## Test result extraction
req       = @request gp[].(s,z,material.(a,b))  # generates an expression, not a variable
key,nkey  = makekey(req,requestable(el))
nstep,nel = 2,3
out       = Array{Float64,3}(undef,nkey,nel,nstep)
istep,iel = 1,2
x,y       = [1.,2.],[.5,.2]
r         = residual(@view(out[:,iel,istep]),key,x,y)

@testset "Espy" begin
    @test key == (gp=[(s=1, z=2, material=(a=3, b=4)), (s=5, z=6, material=(a=7, b=8))],)
    @test out[:,iel,istep] ≈ [3.75, 1.5, 2.5, 3.75, 7.04, 2.2, 3.2, 7.04]
end

out = Vector{Float64}(undef,5)
out[[1,2]] .= [1,2]
out[[3,4]] .= (3,4)
# out[[5,6]] .= @SVector [5,6] works, but requires StaticArrays in the test environment
out[[5]]   .= 5
@testset "Assign to out" begin
    @test out == [1.,2.,3.,4.,5.]
end

end


# Test of Depth-first visit

using Graphs
using Base.Test

type GraphTest
    graph_edges::Array{(Int,Int),1}
    dfs_path::Array{Int,1}
    is_cyclic::Bool
    topo_sort::Array{Int,1}

    GraphTest(gedges, dfspath, iscyclic, toposort) =
        new(gedges, dfspath, iscyclic, toposort)
end

dir_acyclic = GraphTest(
    [(1,2), (1,3), (1,6), (2,4), (2,5), (3,5), (3,6)],
    [1, 2, 4, 5, 3, 6],
    false,
    [1,3,6,2,5,4])
undir_acyclic = GraphTest(
    [(1,2), (1,3), (1,6), (2,4), (2,5)],
    [],
    false,
    [1,6,3,2,5,4])
cyclic  = GraphTest(
    [(1,2), (1,3), (1,6), (2,4), (2,5), (3,5), (3,6), (5,1)],
    [1, 2, 4, 5, 3, 6],
    true,
    [])

testsets = [
    (true, [dir_acyclic, cyclic]),
    (false, [undir_acyclic, cyclic])]

for tset in testsets
    is_dir, graphtests = tset

    for gtest in graphtests

        g = simple_inclist(6, is_directed = is_dir)
        map((edg) -> add_edge!(g, edg[1], edg[2]), gtest.graph_edges)

        gEx = graph(ExVertex[], ExEdge{ExVertex}[], is_directed = is_dir)
        map((x) -> add_vertex!(gEx, "edge:" * string(x)), 1:6)
        V = vertices(gEx)
        map((edg) -> add_edge!(gEx, V[edg[1]], V[edg[2]]), gtest.graph_edges)

        # DFS traversal
        if !isempty(gtest.dfs_path)
            vs1 = visited_vertices(g, DepthFirst(), 1)
            @assert vs1 == gtest.dfs_path

            vs2 = visited_vertices(gEx, DepthFirst(), V[1])
            @assert vs2 == collect(map((x) -> gEx.vertices[x], gtest.dfs_path))
        end

        # Cyclic test
        @assert test_cyclic_by_dfs(g) == gtest.is_cyclic
        @assert test_cyclic_by_dfs(gEx) == gtest.is_cyclic

        # Topological sort
        if gtest.is_cyclic
            @test_throws ArgumentError topological_sort_by_dfs(g)  # g2 contains a loop
            @test_throws ArgumentError topological_sort_by_dfs(gEx)  # g2 contains a loop

        elseif !isempty(gtest.topo_sort)
            ts = topological_sort_by_dfs(g)
            @assert ts == gtest.topo_sort

            ts = topological_sort_by_dfs(gEx)
            @assert [vertex_index(e,gEx) for e in ts] == gtest.topo_sort
        end
    end
end


module TestBootsampleNonParametric

using Bootstrap
using Test

using Statistics
using Random

using DataFrames
using RDatasets
using StatsBase

## 'city' dataset
const city = dataset("boot", "city")
const citya = convert(Matrix, city)

@testset "Non-parametric bootstraps" begin

    function test_bootsample(bs, ref, raw_data, n)

        show(IOBuffer(), bs)
        @test typeof(bs) <: NonParametricBootstrapSample
        t0 = original(bs)
        @test length(t0) == length(ref)
        [@test t ≈ r for (t, r) in zip(t0, ref)]

        t1 = straps(bs)
        @test length(t1) == length(t0)
        [@test length(t) == n for t in t1]
        [@test (minimum(t) <= tr && maximum(t) >= tr) for (t, tr) in zip(t1, t0)]
        [@test eltype(t) == eltype(tr) for (t, tr) in zip(t1, t0)]

        @test Bootstrap.data(bs) == raw_data

        @test nrun(sampling(bs)) == n
        @test nrun(bs) == n

        @test length(bias(bs)) == length(ref)
        [@test (b[1] > -Inf && b[1] < Inf) for b in bias(bs)]
        @test length(stderror(bs)) == length(ref)
        [@test s >= 0 for s in stderror(bs)]

        [@test original(bs, i) == original(bs)[i]  for i in 1:nvar(bs)]
        [@test straps(bs, i) == straps(bs)[i]  for i in 1:nvar(bs)]
        [@test bias(bs, i) == bias(bs)[i]  for i in 1:nvar(bs)]
        [@test stderror(bs, i) == stderror(bs)[i]  for i in 1:nvar(bs)]

        @test_throws MethodError model(bs)

        return Nothing
    end

    function test_confint(bs)

        cim_all = (BasicConfInt(), PercentileConfInt(), NormalConfInt(), BCaConfInt())
        for cim in cim_all
            c = confint(bs, cim)
            [@test (x[1] >= x[2]  && x[1] <= x[3]) for x in c]
            [@test x[1] ≈ t0 for (x, t0) in zip(c, original(bs))]
            @test level(cim) == 0.95
        end

        return Nothing
    end

    function test_confint_width0(bs)

        cim_all = (BasicConfInt(), PercentileConfInt(), NormalConfInt(), BCaConfInt())
        for cim in cim_all
            c = confint(bs, cim)
            # We expect a confidence interval of width 0
            [@test x[2] == x[1] == x[3] for x in c]
            [@test x[1] ≈ t0 for (x, t0) in zip(c, original(bs))]
            @test level(cim) == 0.95
        end

        return Nothing
    end

    n = 250

    city_ratio(x::AbstractArray) = mean(x[:,2]) ./ mean(x[:,1])
    city_ratio(x::AbstractDataFrame) = mean(x[!,:X]) ./ mean(x[!,:U])

    city_cor(x::AbstractArray) = cor(x[:,1], x[:,2])
    city_cor(x::AbstractDataFrame) = cor(x[!,:X], x[!,:U])

    @testset "Basic resampling" begin

        @testset "city_ratio with DataFrame input" begin
            ref = city_ratio(city)
            @test ref  ≈ 1.5203125
            bs = bootstrap(city_ratio, city, BasicSampling(n))
            test_bootsample(bs, ref, city, n)
            test_confint(bs)
        end

        @testset "city_cor with DataFrame input" begin
            ref = city_cor(city)
            bs = bootstrap(city_cor, city, BasicSampling(n))
            test_bootsample(bs, ref, city, n)
            test_confint(bs)
        end

        @testset "city_cor with DataArray input" begin
            ref = city_cor(citya)
            bs = bootstrap(city_cor, citya, BasicSampling(n))
            test_bootsample(bs, ref, citya, n)
            test_confint(bs)
        end

        @testset "mean_and_sd: Vector input, 2 output variables" begin
            r = randn(25)
            ref = mean_and_std(r)
            bs = bootstrap(mean_and_std, r, BasicSampling(n))
            test_bootsample(bs, ref, r, n)
            test_confint(bs)
        end

        @testset "mean_and_sd: Student CI" begin
            r = randn(25)
            bs = bootstrap(mean_and_std, r, BasicSampling(n))
            ## Student confint
            cim = StudentConfInt()
            c = confint(bs, straps(bs, 2), cim, 1)
            @test c[1] >= c[2]  && c[1] <= c[3]
            @test c[1] ≈ original(bs, 1)
            @test level(cim) == 0.95
        end
    end

    @testset "Antithetic resampling" begin

        @testset "mean_and_sd: Vector input, 2 output variables" begin
            r = randn(50)
            ref = mean_and_std(r)
            bs = bootstrap(mean_and_std, r, AntitheticSampling(n))
            test_bootsample(bs, ref, r, n)
            test_confint(bs)
        end

    end

    @testset "Balanced resampling" begin

        @testset "city_ratio with DataFrame input" begin
            ref = city_ratio(city)
            @test ref ≈ 1.5203125
            bs = bootstrap(city_ratio, city, BalancedSampling(n))
            test_bootsample(bs, ref, city, n)
            test_confint(bs)
        end

        @testset "city_cor with DataFrame input" begin
            ref = city_cor(city)
            bs = bootstrap(city_cor, city, BalancedSampling(n))
            test_bootsample(bs, ref, city, n)
            test_confint(bs)
        end

        @testset "city_cor with DataArray input" begin
            ref = city_cor(citya)
            bs = bootstrap(city_cor, citya, BalancedSampling(n))
            test_bootsample(bs, ref, citya, n)
            test_confint(bs)
        end

        @testset "mean_and_sd: Vector input, 2 output variables" begin
            r = randn(50)
            ref = mean_and_std(r)
            bs = bootstrap(mean_and_std, r, BalancedSampling(n))
            test_bootsample(bs, ref, r, n)
            test_confint(bs)
            ## mean should be unbiased
            @test isapprox(bias(bs)[1], 0.0, atol = 1e-8)
        end

    end

    @testset "Exact resampling" begin

        nc = Bootstrap.nrun_exact(nrow(city))

        @testset "city_ratio with DataFrame input" begin
            ref = city_ratio(city)
            @test ref ≈ 1.5203125
            bs = bootstrap(city_ratio, city, ExactSampling())
            test_bootsample(bs, ref, city, nc)
            test_confint(bs)
        end

        @testset "mean: Vector input, 1 output variables" begin
            r = randn(10)
            ref = mean(r)
            bs = bootstrap(mean, r, ExactSampling())
            test_bootsample(bs, ref, r, nc)
            test_confint(bs)
        end

    end

    @testset "Maximum Entropy resampling" begin

        # Simulated AR(1) data copied from
        # https://github.com/colintbowers/DependentBootstrap.jl/blob/master/test/runtests.jl#L8

        nobs = 100

        function test_obs(n, seed = 1234)
            Random.seed!(seed)
            e = randn(n)
            x = zeros(Float64, n)
            for i = 2:n
                x[i] = 0.8 * x[i - 1] + e[i]
            end
            return sin.(x)
        end

        r = test_obs(nobs)
        ref = mean(r)
        s = MaximumEntropySampling(n)
        bs = bootstrap(mean, r, s)
        test_bootsample(bs, ref, r, n)
        test_confint(bs)

        # Collect the samples
        samples = zeros(eltype(r), (nobs, n))
        for i in 1:n
            # For some reason the samples are only filled in if we have
            # an explicit assignment back into the matrix.
            samples[:, i] = draw!(s.cache, r, samples[:, i])
        end

        # Add some checks to ensure that our within sample variation is greater than our
        # across sample variation at any given "timestep".
        @test all(std(samples, dims = 2) .< std(r))
        @test mean(std(samples, dims = 2)) < 0.1  # NOTE: This is about 0.09 in julia and 0.08 in the R package
        @test std(r) > 0.5

    end

    @testset "ConfInt" begin
        @testset "All equal" begin
            r = rand(10)
            # In this case our count will always be the same because all values are positive
            bs = bootstrap(x->count(x .> 0.0), r, BasicSampling(n))
            test_confint_width0(bs)
        end
    end
end

end


using Random, Test
using TupleTools

using Base: tail, front

n = 10

p = randperm(n)
ip = invperm(p)

t = (p...,)

@test @inferred(TupleTools.argtail2(1, 2, 3, 4)) == (3, 4)
@test @inferred(TupleTools.tail2(t)) == t[3:n]
@test @inferred(TupleTools.unsafe_tail(t)) == t[2:n]
@test @inferred(TupleTools.unsafe_front(t)) == t[1:n-1]
@test @inferred(TupleTools.unsafe_tail(())) == ()
@test @inferred(TupleTools.unsafe_front(())) == ()
@test @inferred(TupleTools.vcat()) == ()
@test @inferred(TupleTools.vcat((1, 2))) == (1, 2)
@test @inferred(TupleTools.getindices(t, ())) == ()
@test @inferred(TupleTools.getindices(t, (1,2,3))) == t[1:3]

@test @inferred(TupleTools.deleteat((1, 2), (1, ))) == (2, )
for i = 1:n
    @test @inferred(TupleTools.deleteat(t, i)) == (deleteat!(copy(p), i)...,)
    @test @inferred(TupleTools.insertat(t, i, (1,2,3))) == (vcat(p[1:i-1], [1,2,3], p[i+1:n])...,)
end
@test @inferred(TupleTools.deleteat((1,2,3,4,5,6), (3,1,5))) == (2,4,6)
@test_throws BoundsError TupleTools.deleteat(t, 0)
@test_throws BoundsError TupleTools.deleteat(t, n+1)
@test_throws BoundsError TupleTools.insertat(t, 0, (1,2,3))
@test_throws BoundsError TupleTools.insertat(t, n+1, (1,2,3))
for i = 0:n
    @test @inferred(TupleTools.insertafter(t, i, (1,2,3))) == (vcat(p[1:i], [1,2,3], p[i+1:n])...,)
end
@test_throws BoundsError TupleTools.insertafter(t, -1, (1,2,3))
@test_throws BoundsError TupleTools.insertafter(t, n+1, (1,2,3))
@test @inferred(TupleTools.vcat((1,2,3),4,(5,),(),(6,7,8))) == (1,2,3,4,5,6,7,8)
@test @inferred(TupleTools.flatten((1,2,3),4,(5,),(),(6,7,8))) == (1,2,3,4,5,6,7,8)
@test @inferred(TupleTools.flatten((1,(2,3)),4,(5,),(),((6,),(7,(8,))))) == (1,2,3,4,5,6,7,8)

@test @inferred(TupleTools.sum(())) == 0
@test @inferred(TupleTools.sum((1, ))) == 1
@test @inferred(TupleTools.sum(t)) == sum(t)
@test @inferred(TupleTools.cumsum(())) == ()
@test @inferred(TupleTools.cumsum((1, ))) == (1, )
@test @inferred(TupleTools.cumsum(t)) == (cumsum(p)...,)
@test @inferred(TupleTools.prod(())) == 1
@test @inferred(TupleTools.prod((2, ))) == 2
@test @inferred(TupleTools.prod(t)) == prod(t)
@test @inferred(TupleTools.cumprod(())) == ()
@test @inferred(TupleTools.cumprod((1, ))) == (1, )
@test @inferred(TupleTools.cumprod(t)) == (cumprod(p)...,)

for a in (t, (1, ))
    @test @inferred(TupleTools.findmin(a)) == findmin(a)
    @test @inferred(TupleTools.findmax(a)) == findmax(a)
    @test @inferred(TupleTools.minimum(a)) == minimum(a)
    @test @inferred(TupleTools.maximum(a)) == maximum(a)
    @test @inferred(TupleTools.argmin(a)) == argmin(a)
    @test @inferred(TupleTools.argmax(a)) == argmax(a)
end

@test @inferred(TupleTools.sort((1, ))) == (1, )
@test @inferred(TupleTools.sort(())) == ()
@test @inferred(TupleTools.sort(t; rev = true)) == (sort(p; rev = true)...,)
@test @inferred(TupleTools.sort(t; rev = false)) == (sort(p; rev = false)..., )
@test @inferred(TupleTools.sortperm(t)) == (sortperm(p)...,)
@test @inferred(TupleTools.sortperm(t; rev=true)) == (sortperm(p; rev=true)..., )
@test @inferred(TupleTools.invperm(t)) == (ip...,)
@test @inferred(TupleTools.isperm(t)) == true
@test @inferred(TupleTools.isperm((1,2,1))) == false
@test @inferred(TupleTools.permute(t, t)) == (p[p]...,)

@test @inferred(TupleTools.vcat()) == ()
@test @inferred(TupleTools.diff(())) == ()
@test @inferred(TupleTools.diff((1, ))) == ()
@test @inferred(TupleTools.diff((1, 2, 3))) == (1, 1)


nbdir = ENV["HOME"] * "/.julia/dev/DeltaRCWA/notebooks/"
include("$(nbdir)NystromMethodQP.jl")
import FFTW: fftfreq
import LinearAlgebra: I, Diagonal

"""
    compute_BIE_method(η::Function, μ::Function, L::Float64, N::Int, k::Float64, i::Int)

Compute the fields using Carlos' Boundary Integral Equation (BIE) method.
The output fields are on a grid compatible with the RCWA method.

Arguments:
η:: A periodic function returning the metasurface magnetic conductance
μ:: A periodic function returning the metasurface electric resistance
L:: the length of the unit cell (some multiple of the period of η, μ)
N:: the number of discretization points along the unit cell
k:: the wavenumber/frequency of the incident planewave (assumed TM polarization)
i:: the index of the incident planewave in the Fourier basis

Keyword arguments:
method:: the Nystrom method to be used (choose default "MK" or "Alpert)
NC:: the number of cells in a windowed sum for Green's function (default 5)
NP:: the number of points in BIE discretization per wavelength (default 20)
s:: a scale factor to stretch the y coordinate of the output grid

Returns:
NamedTuple with fields:
I₁:: an array of the incident fields in the position basis for y<0
O₁:: an array of the outgoing fields in the position basis for y<0
I₂:: an array of the incident fields in the position basis for y>0
O₂:: an array of the outgoing fields in the position basis for y>0
"""
function compute_BIE_method(
    η::Function, μ::Function, L::Float64, N::Int, k::Float64, i::Int;
    method="MK", NC=5, NP=20, s=1
)
    dx = L/N
    k⃗ = 2π*fftfreq(N, 1/dx)
    β⃗ = @. sqrt(Complex(k^2 - k⃗^2))
    α = k⃗[i]
    β = β⃗[i]
    uInc(x,y)= exp(im*α*x+im*β*y)  # incident planewave
    # Defining geometry
    # Vertices that define the curve Γ
    ver = [
        L/2   0;   #1
        -L/2   0   #2
    ]
    # In this case, Γ is just a line segment from (L/2,0) to (-L/2,0).
    # This matrix defines the orientation of the curve
            #starting node  #ending node  #domain on the right #domain on the left
    info = [1               2             1                    0]
    # The parameter M defines the number of boundary points to used in the discretization of the BIE
    λ = 2π/k
    M  = Int(round(max(NP,NP/λ))) # ~20 pts per unit length
    # Γ contains all the information of the discretized curve that is needed for the numerical solution of the problem
    Γ  = SkeletonCctor(ver,info,[M])

    # Solving the boundary integral equations
    # method::Selection of the Nystrom method to be used
    # NC::Number of unit cells considered in the windowed sum (A = NC*L)
    mat = matricesSkeleton(1,Γ,k,L,α,NC,method) # Construction of the Nystrom matrices

    x1 = Γ.parts[1].x[:,1] # x-coordinate of the discretization points on Γ

    f1(x) = -2*im*k*η(x)*exp(im*α*x)
    A⁺ = -0.5*I+((im*k)*Diagonal(η.(x1)))*mat.S  # System matrix for Ω⁺
    b⁺ = f1.(x1) # boundary data for Ω⁺
    ψ⁺ = A⁺\b⁺

    f2(x) = -2*im*β*exp(im*α*x)
    A⁻ = -0.5*I+((im*k)*Diagonal(μ.(x1)))*mat.S # System matrix for Ω⁻
    b⁻ =  f2.(x1) # boundary data for Ω⁻
    ψ⁻ = A⁻\b⁻

    φ⁺ = 0.5*(ψ⁺+ψ⁻)
    φ⁻ = 0.5*(ψ⁺-ψ⁻)

    x = range(0, step=dx, length=N)
    y = range(-s*L, s*L, length=2N)
    y₁ = y[y .< 0] # y points in domain below sheet
    y₂ = y[y .> 0] # y points in domain above sheet

    ver⁺ = [L 0;-L 0;-L 2s*L;L 2s*L]
    Ω⁺ = isin(x,y,ver⁺)

    ver⁻ = [L 0;-L 0;-L -2s*L;L -2s*L]
    Ω⁻ = isin(x,y,ver⁻)

    UInc = reshape(uInc.(Ω⁺.pts[:,1],Ω⁺.pts[:,2]),Ω⁺.Ny,Ω⁺.Nx)

    pot⁺ = potentialsSkeleton(Ω⁺.pts,Γ,k,L,α,NC);
    U⁺scat = reshape(pot⁺.SL*φ⁺,Ω⁺.Ny,Ω⁺.Nx)
    U⁺   =  Ω⁺.In .* (U⁺scat .+ UInc)

    pot⁻= potentialsSkeleton(Ω⁻.pts,Γ,k,L,α,NC);
    U⁻scat = reshape(pot⁻.SL*φ⁻,Ω⁻.Ny,Ω⁻.Nx)
    U⁻   =  Ω⁻.In .* (U⁻scat .+ UInc)

    Ω⁺mask = isfinite.(Ω⁺.In)
    I₂ = reshape(UInc[Ω⁺mask] .- UInc[Ω⁺mask], length.((y₂, x)))
    O₂ = reshape(U⁺[Ω⁺mask], length.((y₂, x)))

    Ω⁻mask = isfinite.(Ω⁻.In)
    I₁ = reshape(UInc[Ω⁻mask], length.((y₁, x)))
    O₁ = reshape(U⁻[Ω⁻mask] .- UInc[Ω⁻mask], length.((y₁, x)))
    (x=x, y₁=y₁, y₂=y₂, I₁=I₁, O₁=O₁, I₂=I₂, O₂=O₂)
end


include("test_params.jl")
using SymbolicRegression, SymbolicUtils, Test, Random, ForwardDiff
using SymbolicRegression:
    Options, string_tree, eval_tree_array, Dataset, differentiable_eval_tree_array
using SymbolicRegression: print_tree, pow, eval_loss, score_func, Node
using SymbolicRegression:
    plus,
    sub,
    mult,
    square,
    cube,
    div,
    log_abs,
    log2_abs,
    log10_abs,
    sqrt_abs,
    acosh_abs,
    neg,
    greater,
    greater,
    relu,
    logical_or,
    logical_and,
    gamma
using SymbolicRegression: node_to_symbolic, symbolic_to_node
using SymbolicRegression: check_constraints, loss

x1 = 2.0
# Initialize functions in Base....
for unaop in [cos, exp, log_abs, log2_abs, log10_abs, relu, gamma, acosh_abs]
    for binop in [sub]
        function make_options(; kw...)
            return Options(;
                default_params...,
                binary_operators=(+, *, ^, /, binop),
                unary_operators=(unaop,),
                npopulations=4,
                verbosity=(unaop == gamma) ? 0 : Int(1e9),
                kw...,
            )
        end
        make_options()

        # for unaop in
        f_true = (x,) -> binop((3.0 * unaop(x))^2.0, -1.2)

        # binop at outside:
        tree = Node(5, (Node(3.0) * Node(1, Node("x1")))^2.0, -1.2)
        tree_bad = Node(5, (Node(3.0) * Node(1, Node("x1")))^2.1, -1.3)
        n = count_nodes(tree)

        true_result = f_true(x1)

        result = eval(Meta.parse(string_tree(tree, make_options())))

        # Test Basics
        @test n == 8
        @test result == true_result

        types_to_test = [Float32, Float64, BigFloat]
        if unaop == cos
            # Other unary operators produce numbers too large
            # to do meaningful tests
            types_to_test = [Float16, types_to_test...]
        end
        for T in types_to_test
            if T == Float16 || unaop == gamma
                zero_tolerance = 3e-2
            else
                zero_tolerance = 1e-6
            end

            Random.seed!(0)
            N = 100
            X = T.(randn(MersenneTwister(0), Float64, 5, N) / 3)
            X = X + sign.(X) * T(0.1)
            y = T.(f_true.(X[1, :]))
            dataset = Dataset(X, y)
            test_y, complete = eval_tree_array(tree, X, make_options())
            test_y2, complete2 = differentiable_eval_tree_array(tree, X, make_options())

            # Test Evaluation
            @test complete == true
            @test all(abs.(test_y .- y) / N .< zero_tolerance)
            @test complete2 == true
            @test all(abs.(test_y2 .- y) / N .< zero_tolerance)

            # Test loss:
            @test abs(eval_loss(tree, dataset, make_options())) < zero_tolerance
            @test eval_loss(tree, dataset, make_options()) ==
                score_func(dataset, one(T), tree, make_options())[2]

            #Test Scoring
            @test abs(score_func(dataset, one(T), tree, make_options(; parsimony=0.0))[1]) <
                zero_tolerance
            @test score_func(dataset, one(T), tree, make_options(; parsimony=1.0))[1] > 1.0
            @test score_func(dataset, one(T), tree, make_options())[1] <
                score_func(dataset, one(T), tree_bad, make_options())[1]
            @test score_func(dataset, one(T) * 10, tree_bad, make_options())[1] <
                score_func(dataset, one(T), tree_bad, make_options())[1]

            # Test gradients:
            df_true = x -> ForwardDiff.derivative(f_true, x)
            dy = T.(df_true.(X[1, :]))
            test_dy = ForwardDiff.gradient(
                _x -> sum(differentiable_eval_tree_array(tree, _x, make_options())[1]), X
            )
            test_dy = test_dy[1, 1:end]
            @test all(abs.(test_dy .- dy) / N .< zero_tolerance)
        end
    end
end

# Generic operator tests
types_to_test = [Float16, Float32, Float64, BigFloat]
for T in types_to_test
    val = T(0.5)
    val2 = T(3.2)
    @test sqrt_abs(val) == sqrt_abs(-val)
    @test abs(log_abs(-val) - log(val)) < 1e-6
    @test abs(log2_abs(-val) - log2(val)) < 1e-6
    @test abs(log10_abs(-val) - log10(val)) < 1e-6
    @test neg(-val) == val
    @test sqrt_abs(val) == sqrt(val)
    @test mult(val, val2) == val * val2
    @test plus(val, val2) == val + val2
    @test sub(val, val2) == val - val2
    @test square(val) == val * val
    @test cube(val) == val * val * val
    @test div(val, val2) == val / val2
    @test greater(val, val2) == T(0.0)
    @test greater(val2, val) == T(1.0)
    @test relu(-val) == T(0.0)
    @test relu(val) == val
    @test logical_or(val, val2) == T(1.0)
    @test logical_or(T(0.0), val2) == T(1.0)
    @test logical_and(T(0.0), val2) == T(0.0)
end

# Test SymbolicUtils interface
_inv(x) = 1 / x
options = Options(;
    default_params...,
    binary_operators=(+, *, ^, /, greater),
    unary_operators=(_inv,),
    constraints=(_inv => 4,),
    npopulations=4,
)
tree = Node(5, (Node(3.0) * Node(1, Node("x1")))^2.0, -1.2)

eqn = node_to_symbolic(tree, options; varMap=["energy"], index_functions=true)
tree2 = symbolic_to_node(eqn, options; varMap=["energy"])

@test string_tree(tree, options) == string_tree(tree2, options)

# Test constraint-checking interface
tree = Node(5, (Node(3.0) * Node(1, Node("x1")))^2.0, -1.2)
violating_tree = Node(1, tree)

@test check_constraints(tree, options) == true
@test check_constraints(violating_tree, options) == false

# Test different loss functions
customloss(x, y) = abs(x - y)^2.5
customloss(x, y, w) = w * (abs(x - y)^2.5)
testl1(x, y) = abs(x - y)
testl1(x, y, w) = abs(x - y) * w

for (loss_fnc, evaluator) in [(L1DistLoss(), testl1), (customloss, customloss)]
    local options = Options(;
        default_params...,
        binary_operators=(+, *, -, /),
        unary_operators=(cos, exp),
        npopulations=4,
        loss=loss_fnc,
    )
    x = randn(MersenneTwister(0), Float32, 100)
    y = randn(MersenneTwister(1), Float32, 100)
    w = abs.(randn(MersenneTwister(2), Float32, 100))
    @test abs(loss(x, y, options) - sum(evaluator.(x, y)) / length(x)) < 1e-6
    @test abs(loss(x, y, w, options) - sum(evaluator.(x, y, w)) / sum(w)) < 1e-6
end

# Test derivatives
include("test_derivatives.jl")

# Test simplification:
include("test_simplification.jl")

# Test `print`:
include("test_print.jl")

include("test_evaluation.jl")

include("test_prob_pick_first.jl")

include("test_crossover.jl")

include("test_nan_detection.jl")

include("test_constraints.jl")

include("test_complexity.jl")


export func_transform, euler_matrix,euler_matrix_update, nonlinear_transform, make_input_vector, make_hermitian
"""
    make_input_vector(x::Vector{T}) -> Vector{T}

Generates input vector for `QuNLDE` iterations.
"""
function make_input_vector(x::Vector{T}) where T
    if !(norm(x) ≈ 1)
        throw(ArgumentError("Input vector is not normalized"))
    end
    len = length(x)
    siz = nextpow(2,len + 1)
    z = zeros(T, siz)
    z[1] = one(T)
    z[2:len+1] = x
    normalize!(z)
    reg = [1,0] ⊗ z ⊗ z
    return reg
end

"""
    make_hermitian(A::Matrix) -> Matrix

Returns hermitian matrix containing A.
"""
function make_hermitian(A::Matrix)
    if ishermitian(A)
        return A
    end

    return [zero(A) im*A'; -im*A zero(A)]
end

function nonlinear_transform(H::Matrix, x::Vector, k::Int, ϵ::Real = 1e-4)
    r, N = taylorsolve(im*H,x,k,ϵ)
    n = log2i(length(x))
    nb = Int((n-1)/2)
    r = relax!(r) |> focus!(n)|> select!(1) |> focus!(1:nb...,) |> select!(0)
    return r,sqrt(2)*N/ϵ
end

"""
    func_transform(A::Matrix, x::Vector, k::Int) -> ArrayReg, <: Complex

Function transform sub-routine. Returns state register and inverse probability of finding  it.
"""
function func_transform(A::Matrix, x::Vector, k::Int,ϵ::Real = 1e-4)
    reg = make_input_vector(x)
    H = make_hermitian(A)
    r, N = nonlinear_transform(H,reg,k,ϵ)
    return r, N
end

"""
    euler_matrix(A::Matrix{CPType},b::Vector{CPType},h::Real) -> Matrix

Generates matrix for forward Euler iteration.
"""
function euler_matrix(A::Matrix{CPType},b::Vector{CPType},h::Real) where CPType
    n = length(b)
    A = CPType(h)*A
    A[1,1] = 1
    @inbounds for i in 1:n
        A[4*i+ 1, i+1] += 1
    end
    return A
end

"""
    euler_matrix(A::Matrix{CPType},b::Vector{CPType},h::Real) -> Matrix

Updates euler matrix for forward Euler iteration.
"""
function euler_matrix_update(A::Matrix{CPType}, b::Vector{CPType}, nrm::Real) where CPType
    n = length(b)
    A = CPType(nrm^2)*A
    A[1,1] = 1
    @inbounds for i in 1:n
        A[4*i+ 1, 1] = A[4*i+ 1, 1]/(nrm^2)
        @. A[4*i+ 1, 2:n+1] = A[4*i+ 1, 2:n+1]/nrm
        for j in 1:n
            A[4*i + 1, 4*j + 1] = A[4*i + 1, 4*j + 1]/nrm
        end
    end
    return A
end

function DiffEqBase.solve(prob::QuODEProblem,alg::QuNLDE; dt = (prob.tspan[2]-prob.tspan[1])/10, kwargs...)
    A = prob.A
    b = prob.b
    k = alg.k
    ϵ = alg.ϵ
    tspan = prob.tspan
    len = round(Int,(tspan[2] - tspan[1])/dt) + 1
    siz = length(b)
    res = Array{eltype(b),2}(undef,len,siz)
    res[1,:] = b
    A = euler_matrix(A,b,dt)
    H = make_hermitian(A)
    reg = make_input_vector(b)
    r, N = nonlinear_transform(H,reg,k,ϵ)
    ntem = 1
    @views for step in 2:len - 1
        tem = vec(state(r))*N*sqrt(2)
        res[step,:] = tem[2:siz+1]
        ntem = norm(res[step,:])
        C = euler_matrix_update(A,b,ntem)
        H = make_hermitian(C)
        reg = res[step,:]/ntem
        reg = make_input_vector(reg)
        r, N = nonlinear_transform(H,reg,k,ϵ)
    end
    tem = vec(state(r))*N*sqrt(2)
    @views res[len,:] = tem[2:siz+1]
    return res
end


 # Code to read simul.pdb file and return number of each species and its concentration in mol L-1 - viniciusp
 # first function

 using Statistics

 function conc( Nan::Float64, Ncat::Float64, Nh2o::Float64, Name::String, BOX::Vector)
   conc = round((Nan / (BOX[1]*BOX[2]*BOX[3])) * ( 1 / (6.02e-4)), digits = 2)
   conc_cat = round((Ncat / (BOX[1]*BOX[2]*BOX[3])) * ( 1 / (6.02e-4)), digits = 2)
   concw = round((Nh2o / (BOX[1]*BOX[2]*BOX[3])) * ( 1 / (6.02e-4)), digits = 2)
   println("System : Protein with $Name")
   println("Water -> molecules : $Nh2o / concentration (mol/L) : $concw   ")
   println("Ions of ionicliquids")
   println("Cations : $Ncat / concentration (mol/L) : $conc_cat ")
   println("Anions  : $Nan  / concentration (mol/L) : $conc")
   m = (BOX[1] + BOX[2] + BOX[3]) / 3
   println("Avarage cubic box length = $m")
 end

 function read_pdb(arquivo::String)

   nan = 0
   nwat = 0
   nan2 = 0
   ncat = 0

   f = open("system.pdb","r+")
   for line in eachline(f)
     if occursin("NC",line)
       nan = nan + 1
     elseif occursin("SOL",line)
       nwat = nwat + 1
     elseif occursin("BF4",line)
       nan2 = nan2 + 1
     elseif occursin("BMI",line) || occursin("EMI",line)  # for systems with one cation, search for EMI or BMI.
       ncat = ncat + 1
     end
   end
   close(f)

   nan = nan/5
   nwat = nwat / 3
   nan2 = nan2 / 5
   ncat = ncat / 20
   box = zeros(3)

   file = open(arquivo,"r")
   for line in eachline(file)
     if occursin("CRYST1",line)
       data = split(line)
       box[1] = parse(Float64,data[2])
       box[2] = parse(Float64,data[3])
       box[3] = parse(Float64,data[4])
     else
       continue
     end
   end

   box_n =  box   #(box[1] + box[2] + box[3]
   return nan, nwat, nan2, ncat, box_n
 end

 LIS = ["EMIMDCABF4" ]
 lista = ["0.50","1.00","1.50","2.00","2.50","3.00"]

  for i in lista
   cd("/home/viniciuspiccoli/Documents/DOUTORADO/EMIM_DCA_BF4/teste/$(LIS[1])/$i/00")

   nan,nwater,nan2,ncat,box = read_pdb("system.pdb")

   println("*****************************************")
   println("Main folder - initial concentration $i")
   conc(nan,ncat,nwater,"EMIMDCABF4",box)
   println("*****************************************")

 end


using Random
using LinearAlgebra
using GenericSVD

using ArbNumerics  # ArbFloat
using Decimals     # Decimal
using DecFP        # Decimal128
using DoubleFloats # Double64
using Quadmath     # Float128
using MultiFloats  # Float64x2

setextrabits(0)
setprecision(BigFloat, 106)
setprecision(ArbFloat, 106)

function test_qr(::Type{T}, n::Int) where {T}
    Random.seed!(0)
    A = T.(rand(Float64, n, n))
    result, elapsed_time, _, _, _ = @timed qr(A)
    num_correct_digits = -log10(Float64(sum(abs.(result.Q * result.R - A))))
    println(rpad(T, 30), " |  qr  | ", round(num_correct_digits, digits=1), " | ", elapsed_time)
end

function test_pinv(::Type{T}, m::Int, n::Int) where {T}
    Random.seed!(0)
    A = T.(rand(Float64, m, n))
    Apinv, elapsed_time, _, _, _ = @timed pinv(A)
    num_correct_digits = -log10(Float64(sum(abs.(A * Apinv * A - A))))
    println(rpad(T, 30), " | pinv | ", round(num_correct_digits, digits=1), " | ", elapsed_time)
end

const N = 40
const M = 25

function main()
    while true
        test_qr(Float64x2, N)
        test_qr(BigFloat, N)
        test_qr(ArbFloat, N)
        test_qr(Dec128, N)
        test_qr(Double64, N)
        test_qr(Float128, N)
        test_pinv(Float64x2, N, M)
        test_pinv(BigFloat, N, M)
        test_pinv(ArbFloat, N, M)
        test_pinv(Double64, N, M)
        test_pinv(Float128, N, M)
    end
end

main()


using Revise
using IsopycnalSurfaces, Test

@testset "IsopycnalSurfaces.jl" begin
    #############################
    # Test dedup! function
    @testset "column" begin
        ################################
        # Idealized mapping onto sigma1
        pz = collect(0.:500.:4000.) # pressure levels
        nz = length(pz)

        θz = collect(range(20,stop=10,length=nz))
        Sz = collect(range(36,stop=35,length=nz))
        ztest = sort(rand(2:8,2))
        σ₁true = sigma1column(θz[ztest],Sz[ztest],pz[ztest])

        σ₁grid = range(minimum(σ₁true),stop=maximum(σ₁true),length=20)
        σ₁=sigma1column(θz,Sz,pz)

        @testset "column_spline" begin
            splorder = 3
            sgood = findall(minimum(σ₁) .<= σ₁grid .<= maximum(σ₁))
            pσ = var2sigmacolumn(σ₁,pz,σ₁grid[sgood],splorder)

            @test isapprox(pσ[begin],pz[ztest[begin]])
            @test isapprox(pσ[end],pz[ztest[end]])
        end

        @testset "column_linear" begin
            splorder = 100
            sgood = findall(minimum(σ₁) .<= σ₁grid .<= maximum(σ₁))
            pσ = var2sigmacolumn(σ₁,pz,σ₁grid[sgood],splorder)

            @test isapprox(pσ[begin],pz[ztest[begin]])
            @test isapprox(pσ[end],pz[ztest[end]])
        end

    end
end


import PowerModelsACDC
const _PMACDC = PowerModelsACDC
import PowerModels
const _PM = PowerModels
import Ipopt
import Memento
# import CPLEX
import SCS
import Mosek
import MosekTools
import JuMP
import Gurobi  # needs startvalues for all variables!

file_pf_droop = "./test/data/case5_acdc_droop.m"

file_case5acdc = "./test/data/case5_acdc.m"
file_case5dc ="./test/data/case5_dc.m"
file_case5dcgrid ="./test/data/case5_dcgrid.m"
file_case5dcgridb0 ="./test/data/case5_dcgrid_b0.m"
file_case24 = "./test/data/case24_3zones_acdc.m"
file_case39 = "./test/data/case39_acdc.m"
file_case3120 = "./test/data/case3120sp_acdc.m"
file_case5dcgrid = "./test/data/case5_dcgrid.m"
file_case5_b2bdc = "./test/data/case5_b2bdc.m"
file_lcc = "./test/data/lcc_test.m"
file_588sdet_acdc = "./test/data/pglib_opf_case588_sdet_acdc.m"
file = file_case5acdc

data = _PM.parse_file(file)

_PMACDC.process_additional_data!(data)

ipopt = JuMP.optimizer_with_attributes(Ipopt.Optimizer, "tol" => 1e-6, "print_level" => 0)
mosek = JuMP.optimizer_with_attributes(Mosek.Optimizer)
scs = JuMP.optimizer_with_attributes(SCS.Optimizer)
gurobi = JuMP.optimizer_with_attributes(Gurobi.Optimizer, Presolve=0)
# cplex = JuMP.optimizer_with_attributes(CPLEX.Optimizer)

s = Dict("output" => Dict("branch_flows" => true), "conv_losses_mp" => true)

resultAC = _PMACDC.run_acdcopf(file, _PM.ACPPowerModel, ipopt; setting = s)

resultLPAC = _PMACDC.run_acdcopf(file, _PM.LPACCPowerModel, ipopt; setting = s)

resultQC = _PMACDC.run_acdcopf(file, _PM.QCRMPowerModel, ipopt; setting = s)

resultSOCBIM = _PMACDC.run_acdcopf(file, _PM.SOCWRPowerModel, ipopt; setting = s)
resultACSOCBIM = _PM.run_opf(file, _PM.SOCWRPowerModel, ipopt; setting = s)
# #
resultSOCBFM = _PMACDC.run_acdcopf_bf(file, _PM.SOCBFPowerModel, ipopt; setting = s)
resultSOCBFMConic = _PMACDC.run_acdcopf_bf(file, _PM.SOCBFConicPowerModel, mosek; setting = s)
resultSOCBFMConicSCS = _PMACDC.run_acdcopf_bf(file, _PM.SOCBFConicPowerModel, scs; setting = s)
# #
resultSDP = _PMACDC.run_acdcopf(file, _PM.SDPWRMPowerModel, mosek; setting = s)
# #
resultDC = _PMACDC.run_acdcopf(file, _PM.DCPPowerModel, gurobi; setting = s)
resultACPF5 = _PMACDC.run_acdcpf(file_case5acdc, _PM.ACPPowerModel, ipopt; setting = s)


function makeTensorT{T}(factorW::Array{T,2}, measureOperator::Array{T,1} = ones(T, size(factorW,1)))
	leng = size(factorW,1)
	tensorT = zeros(T,leng,leng,leng,leng)
	for a = 1:leng, x = 1:leng, xp = 1:leng, y = 1:leng, yp = 1:leng
		tensorT[x,xp,y,yp] += measureOperator[a] * factorW[a,x] * factorW[a,xp] * factorW[a,y] * factorW[a,yp]
	end
	return tensorT
end

function makeTensorP{T}(factorW::Array{T,2})
	leng = size(factorW,1)
	tensorP = zeros(T,leng,leng,leng)
	for a = 1:leng, x = 1:leng, xp = 1:leng, s = 1:leng
		tensorP[x,xp,s] += factorW[a,x] * factorW[a,xp] * factorW[a,s]
	end
	return tensorP
end

function makeTensorQ{T}(factorW::Array{T,2})
	leng = size(factorW,1)
	tensorQ = zeros(T,leng,leng)
	for a = 1:leng, x = 1:leng, y = 1:leng
		tensorQ[x,y] += factorW[a,x] * factorW[a,y]
	end
	return tensorQ
end


nt = 1000
seed = 99
ssm = acqr(0.8,1.5,2,3,0.5,2)
m = acqr_simula(ssm,nt,seed)
y = m["y"]
a = ssm.a
c = ssm.c
q = ssm.q
r = ssm.r
x10 = ssm.m1
P10 = ssm.P1

###########################
print("Testing acqr_kalman: ")
start = acqr(0.5,0.5,1,1,0,0)
me = acqr_em_s(y,start,max_iter=10)
mk = acqr_kalman(y,me.fit,s=true)
if sum(mk.prediction.^2) > 1e-6
    println("OK")
end


## TEAVAR TE
function TEAVAR(GRB_ENV, edges, capacity, flows, demand, beta, T, Tf, scenarios, scenario_probs; average=false)
    printstyled("\n** solving TEAVAR LP at beta $(beta)..\n", color=:yellow)
    nedges = length(edges)
    nflows = length(flows)
    ntunnels = length(T)
    nscenarios = length(scenarios)
    p = scenario_probs

    tunnel_num = 0
    for x in 1:length(Tf)
        if size(Tf[x],1) > tunnel_num
            tunnel_num = size(Tf[x],1)
        end
    end

    #CREATE TUNNEL SCENARIO MATRIX
    X  = ones(nscenarios,ntunnels)
    for s in 1:nscenarios
        for t in 1:ntunnels
            if size(T[t],1) == 0
                X[s,t] = 0
            else
                for e in 1:nedges
                    if scenarios[s][e] == 0
                        back_edge = findfirst(x -> x == (edges[e][2],edges[e][1], edges[e][3]), edges)
                        if in(e, T[t]) || in(back_edge, T[t])
                        # if in(e, T[t])
                            X[s,t] = 0
                        end
                    end
                end
            end
        end
    end

    #CREATE TUNNEL EDGE MATRIX
    L = zeros(ntunnels, nedges)
    for t in 1:ntunnels
        for e in 1:nedges
            if in(e, T[t])
                L[t,e] = 1
            end
        end
    end

    model = Model(() -> Gurobi.Optimizer(GRB_ENV))
    set_optimizer_attribute(model, "OutputFlag", 0)
    set_optimizer_attribute(model, "Threads", 32)

    model_def_time = @elapsed begin

    @variable(model, a[1:nflows, 1:tunnel_num] >= 0)
    @variable(model, alpha >= 0)
    @variable(model, umax[1:nscenarios] >= 0)
    @variable(model, u[1:nscenarios, 1:nflows] >= 0)

    for e in 1:nedges
        @constraint(model, sum(a[f,t] * L[Tf[f][t],e] for f in 1:nflows for t in 1:size(Tf[f],1)) <= capacity[e])   #overlapping flows cannot add up to the capacity of that link
    end

    # FLOW LEVEL LOSS
    @expression(model, satisfied[s=1:nscenarios, f=1:nflows], sum(a[f,t] * X[s,Tf[f][t]] for t in 1:size(Tf[f],1)) / demand[f])

    for s in 1:nscenarios
        for f in 1:nflows
            @constraint(model, u[s,f] == 1 - satisfied[s,f])
        end
    end

    for s in 1:nscenarios
        if average
            @constraint(model, umax[s] + alpha >= (sum(u[s,f] for f in 1:nflows)) / nflows)
        else
            for f in 1:nflows
                @constraint(model, umax[s] >= u[s,f] - alpha)
            end
        end
    end

    @objective(model, Min, alpha + (1 / (1 - beta)) * sum((p[s] * umax[s] for s in 1:nscenarios)))
    end  # model_def_time end

    optimize!(model)
    solve_runtime = solve_time(model)
    opt_runtime = solve_runtime + model_def_time

    return value.(alpha), objective_value(model), value.(a), value.(umax), value.(u), solve_runtime, opt_runtime
end


# Copyright (c) 2018: Matthew Wilhelm & Matthew Stuber.
# This code is licensed under MIT license (see LICENSE.md for full details)
#############################################################################
# EAGO
# A development environment for robust and global optimization
# See https://github.com/PSORLab/EAGO.jl
#############################################################################
# src/eago_optimizer/bound.jl
# Computes interval bounds of various functions.
#############################################################################

###
### AFFINE FUNCTIONS
###
function lower_interval_bound(m::GlobalOptimizer, f::AffineFunctionIneq)
    fL = f.constant
    for (c, j) in f.terms
        xL = _lower_bound(FullVar(), m, j)
        xU = _upper_bound(FullVar(), m, j)
        fL += (c > 0.0) ? c*xL : c*xU
    end
    return fL
end

function interval_bound(m::GlobalOptimizer, f::Union{AffineFunctionEq,AffineFunctionIneq})
    fL = fU = f.constant
    for (c, j) in f.terms
        xL = _lower_bound(FullVar(), m, j)
        xU = _upper_bound(FullVar(), m, j)
        if c > 0.0
            fL += c*xL
            fU += c*xU
        else
            fL += c*xU
            fU += c*xL
        end
    end
    return fL, fU
end

###
### QUADRATIC FUNCTIONS
###

function lower_interval_bound(m::GlobalOptimizer, f::BufferedQuadraticIneq)
    fval = Interval{Float64}(f.func.constant)
    for t in f.func.affine_terms
        c = t.coefficient
        j = t.variable.value
        xL = _lower_bound(FullVar(), m, j)
        xU = _upper_bound(FullVar(), m, j)
        fval += c > 0.0 ? c*xL : c*xU
    end
    for t in f.func.quadratic_terms
        c = t.coefficient
        i = t.variable_1.value
        j = t.variable_2.value
        xL = _lower_bound(FullVar(), m, i)
        xU = _upper_bound(FullVar(), m, i)
        if i == j
            if c > 0.0
                fval += (0.0 < xL) ? 0.5*c*xL*xL : ((xU <= 0.0) ? 0.5*c*xU*xU : 0.0)
            else
                fval += (xL < xU) ? 0.5*c*xU*xU : 0.5*c*xL*xL
            end
        else
            yL = _lower_bound(FullVar(), m, j)
            yU = _upper_bound(FullVar(), m, j)
            fval += c*Interval{Float64}(xL, xU)*Interval{Float64}(yL, yU)
        end
    end
    return fval.lo
end
function interval_bound(m::GlobalOptimizer, f::BufferedQuadraticIneq)

    fval = Interval{Float64}(f.func.constant)
    for t in f.func.affine_terms
        c = t.coefficient
        j = t.variable.value
        xL = _lower_bound(FullVar(), m, j)
        xU = _upper_bound(FullVar(), m, j)
        fval += c*Interval(xL, xU)
    end
    for t in f.func.quadratic_terms
        c = t.coefficient
        i = t.variable_1.value
        j = t.variable_2.value
        xL = _lower_bound(FullVar(), m, i)
        xU = _upper_bound(FullVar(), m, i)
        if i == j
            fval += 0.5*c*pow(Interval(xL, xU), 2)
        else
            yL = _lower_bound(FullVar(), m, j)
            yU = _upper_bound(FullVar(), m, j)
            fval += c*Interval(xL, xU)*Interval(yL, yU)
        end
    end
    return fval.lo, fval.hi
end

function interval_bound(m::GlobalOptimizer, f::BufferedQuadraticEq)

    fval = Interval{Float64}(f.func.constant)
    for t in f.func.affine_terms
        c = t.coefficient
        j = t.variable.value
        xL = _lower_bound(FullVar(), m, j)
        xU = _upper_bound(FullVar(), m, j)
        fval += c*Interval(xL, xU)
    end
    for t in f.func.quadratic_terms
        c = t.coefficient
        i = t.variable_1.value
        j = t.variable_2.value
        xL = _lower_bound(FullVar(), m, i)
        xU = _upper_bound(FullVar(), m, i)
        if i == j
            fval += 0.5*c*pow(Interval(xL, xU), 2)
        else
            yL = _lower_bound(FullVar(), m, j)
            yU = _upper_bound(FullVar(), m, j)
            fval += c*Interval(xL, xU)*Interval(yL, yU)
        end
    end
    return fval.lo, fval.hi
end

###
### SECOND-ORDER CONE
###
function lower_interval_bound(m::GlobalOptimizer, d::BufferedSOC)

    sol_branch_map = m._sol_to_branch_map
    lo_bnds = n.lower_variable_bounds
    up_bnds = n.upper_variable_bounds
    vec_of_vi = d.variables.variables

    norm_bound = Interval(0.0)
    for i = 2:length(vec_of_vi)
        mapped_vi = @inbounds sol_branch_map[vec_of_vi[i].value]
        x = Interval{Float64}(lo_bnds[mapped_vi], up_bnds[mapped_vi])
        norm_bound += pow(x, 2)
    end
    norm_bound = sqrt(norm_bound)

    mapped_vi = @inbounds sol_branch_map[vec_of_vi[1].value]
    lower_bound = norm_bound.lo -(@inbounds up_bnds[mapped_vi])

    return lower_bound
end

###
### NONLINEAR FUNCTIONS
###
function lower_interval_bound(m::GlobalOptimizer, d::BufferedNonlinearFunction{V,N,T}) where {V,N,T}
    !has_value(d) && forward_pass!(m._working_problem._relaxed_evaluator, d)
    is_num(d) ? num(d) : interval(d).lo
end
function interval_bound(m::GlobalOptimizer, d::BufferedNonlinearFunction{V,N,T}) where {V,N,T}
    !has_value(d) && forward_pass!(m._working_problem._relaxed_evaluator, d)
    v = is_num(d) ? Interval{Float64}(num(d)) : interval(d)
    return v.lo, v.hi
end

is_feasible(m::GlobalOptimizer, f::Union{AFI,BQI}) = lower_interval_bound(m, f) <= 0.0
function is_feasible(m::GlobalOptimizer, f::Union{AFE,BQE})
    l, u = interval_bound(m, f)
    l <= 0.0 <= u
end
function is_feasible(m::GlobalOptimizer, f::BufferedNonlinearFunction{V,N,T}) where {V,N,T}
    l, u = interval_bound(m, f)
    feasible_flag = (u >= lower_bound(f))
    feasible_flag && (l <= upper_bound(f))
end

bound_objective(m::GlobalOptimizer, f::BufferedNonlinearFunction) = interval_bound(m, f)
bound_objective(m::GlobalOptimizer, f::AffineFunctionIneq) = interval_bound(m, f)
bound_objective(m::GlobalOptimizer, f::BufferedQuadraticIneq) = interval_bound(m, f)
function bound_objective(m::GlobalOptimizer, f::VI)
    vval = f.value
    l = lower_bound(FullVar(), m, vval)
    u = upper_bound(FullVar(), m, vval)
    return l, u
end

bound_objective(t::ExtensionType, m::GlobalOptimizer) = bound_objective(m, m._working_problem._objective)
bound_objective(m::GlobalOptimizer{R,Q,S}) where {R,Q,S<:ExtensionType} = bound_objective(_ext(m), m)


nt = 1000
seed = 99
ssm = acqr(0.8,1.5,2,3,0.5,2)
m = acqr_simula(ssm,nt,seed)
y = m["y"]

#######################
print("Testing acqr_em_starting: ")
ssm = acqr_em_starting(y)
if sum(ssm.c^2) > 1e-6
    println("OK")
end


# Dynamically aware constraints manager

include("constraintManager.jl")
include("../objective/LQR_objective.jl")

"""
    constraintManager_Dynamics

Contains a list of constraints and the dual variables (lambda) that correspond
with each constraint.

Unlike [`constraintManager_Base`](@ref), this constraint manager also has the
dynamics constraints embedded such that the nearest dynamically feasible
trajectory can be determined at any given time.
"""
mutable struct constraintManager_Dynamics <: constraintManager
    cList::Array{constraint, 1}
    lambdaList
    dynamicsConstraint::AL_AffineEquality
    affineLambdaList
end

#############################################################
###            Evaluate the Affine Constraints            ###
#############################################################

function evalAffineEq(cM::constraintManager_Dynamics, yCurr)
    return getRaw(cM.dynamicsConstraint, yCurr)
end

function affineBlock(cM::constraintManager_Dynamics)
    return cM.dynamicsConstraint.A
end


function metropolis(T::Real, initial::Real, candidate::Real)
    if candidate < initial
        true
    else
        rand() < exp((initial - candidate) / T)
    end
end

function probabilistic_improvement(tuning_run::Run;
                                   threshold::AbstractFloat = 2.,
                                   acceptance_criterion::Function = metropolis)
    initial_cost  = tuning_run.starting_cost
    x             = deepcopy(tuning_run.starting_point)
    x_proposal    = deepcopy(tuning_run.starting_point)
    name          = "Probabilistic Improvement"
    cost_calls    = 0
    iteration     = 1
    neighbor!(x_proposal)
    proposal      = @fetch (tuning_run.measurement_method(tuning_run, x_proposal))
    cost_calls   += tuning_run.cost_evaluations
    if acceptance_criterion(threshold, initial_cost, proposal)
        update!(x, x_proposal.parameters)
        initial_cost = proposal
    end
    Result(name, tuning_run.starting_point, x, initial_cost, iteration,
           iteration, cost_calls, false)
end


## Connection and Linearization

## A connection between polynomial systems `P` and `Q` is a function `C(m,n)` satisfying
## p_n = ∑C(m,n) q_m or  pᵢ = ∑ Cʲᵢqⱼ (if unicode is helpful)
##
## From this a polynomial in P = ∑ aᵢ⋅pᵢ can be represented in a Q through
##
##  ∑ᵢ aᵢ⋅pᵢ = ∑ᵢ aᵢ  ∑ⱼ Cʲᵢqⱼ  = ∑ᵢ (∑ⱼ (aᵢ⋅Cʲᵢ)) qⱼ or ∑ⱼ (∑ᵢ aᵢCʲᵢ) qⱼ
##
## The choice of which depends on how the connection coefficients are presented.
##
## A connection between P and Q defines a `convert` method `Base.convert(Q,  p::P)`
## Conversion can also be achieved through polynomial evaluation. For example
## if `x=variable(Q)` then `Base.convert(Q,p) = p(x)` will also convert `p` into
## the basis Q. This is helpful when Q is a standard basis polynomial.
##
## However, if Q is not a standard basis polynomial, it isn't helpful unless
## one can multiply monomomials within Q. For standard basis polynomials this
## is known, as  xⁱ⋅xʲ = xⁱ⁺ʲ.
##
## For general Q though, one needs a linearization: a function α with
##
##  qᵢqⱼ = ∑_k α(k,i,j) q_k
##
## with a given α, multiplication can be defined. The linearization can be defined
## through composition: If P is the standard basis, connect Q to P, multiply ini P, then connect
## to Q.
##
## In short, if there is a linearization then multiplication can be defined and conversion defined through
## polynomial evaluation
## If there is a connection defined then conversion is defined and multiplication can be defined through composition
##

## For `AbstractCCOP` polynomials, Koepf and Schmersau present recursive formulas for the connection coefficients
## between Q and the standard basis, the standard and Q, and between P and Q when σ and σ̂ are the same. (This is the case when a polynomial system is parameterized, like Laguerre, Gegenbauer, Bessel, and  Jacobi.)

ConvertibleTypes = Union{AbstractCOP, Polynomials.StandardBasisPolynomial,  FallingFactorial}
kn(::Type{P},  n) where{P <: Union{Polynomials.StandardBasisPolynomial,  FallingFactorial}} = one(eltype(one(P)))
k1k0(::Type{P},  n) where{P <: Union{Polynomials.StandardBasisPolynomial,  FallingFactorial}} = one(eltype(one(P)))
knk_1(::Type{P},  n) where{P <: Union{Polynomials.StandardBasisPolynomial,  FallingFactorial}} = one(eltype(one(P)))

## If  the connection  coefficients for  (P,Q) satisfy c₀⋅ C̃ⁱⱼ + c₁⋅C̃ⁱ⁺¹ⱼ  + c₂⋅*C̃ⁱ⁺²ⱼ = 0 for some computable
##  values c₀,c₁,c₂ then this  will implement  `convert(Q,p::P)`
function _convert_cop(::Type{Q}, p::P) where {P <: ConvertibleTypes,
                                               Q <: ConvertibleTypes}

    d = degree(p)
    T,S = eltype(Q), eltype(p)  #
    R = typeof(one(T) * p[0]/1)
    as = zeros(R, 1+d)

    λⱼⱼ = one(R)  * k0(P) / k0(Q)

    for j in  0:d

        λ = λⱼⱼ

        λⱼⱼ *= k1k0(P,j) / k1k0(Q,j)

        pⱼ = p[j]
        iszero(pⱼ) && continue

        C̃ʲⱼ = one(R)

        as[1+j] += pⱼ * (λ * C̃ʲⱼ)

        C̃ⁱ⁺²ⱼ, C̃ⁱ⁺¹ⱼ = zero(R), C̃ʲⱼ  # for recursion starting  with C̃ʲ⁺¹ⱼ, C̃ʲⱼ
        for i in j-1:-1:0

            λ *= k1k0(Q,i)

            # c₀⋅ C̃ⁱⱼ + c₁⋅C̃ⁱ⁺¹ⱼ  + c₂⋅*C̃ⁱ⁺²ⱼ) = 0
            c₀,c₁,c₂ = connection_m(P, Q, i, j)
            C̃ⁱⱼ = iszero(c₀) ? zero(R) : -one(R)*(c₁*C̃ⁱ⁺¹ⱼ + c₂*C̃ⁱ⁺²ⱼ) / c₀
            as[1+i] += pⱼ * (λ * C̃ⁱⱼ)

            C̃ⁱ⁺²ⱼ, C̃ⁱ⁺¹ⱼ = C̃ⁱ⁺¹ⱼ, C̃ⁱⱼ
        end
    end
    X = Polynomials.indeterminate(Q,p)
    ⟒(Q){R,X}(as)
end

# ## Use FastTransform for conversion, as  possible
# ## FastTransforms.kind2string.(0:15)
# # Legendre<--Chebyshev
# function _convert(::Type{Q}, p::P) where {
#     T <: AbstractFloat,
#     Q <: Union{Legendre, OrthonormalLegendre},
#     P <: Union{Chebyshev{T}, OrthonormalChebyshev{T}}}
#     ps = coeffs(p)
#     Q(cheb2leg(ps, normcheb=isorthonormal(P), normleg=isorthonormal(Q)))
# end

# # Chebyshev<--Legendre
# function _convert(::Type{Q}, p::P) where {
#     T <: AbstractFloat,
#     Q <: Union{Chebyshev, OrthonormalChebyshev},
#     P <: Union{Legendre{T}, OrthonormalLegendre{T}}
# }
#     ps =  coeffs(p)
#     Q(leg2cheb(ps, normleg=isorthonormal(P), normcheb=isorthonormal(Q)) )
# end

# # ultraspherical<--ultraspherical
# function _convert(::Type{Q}, p::P) where {
#     β, α,  T <: AbstractFloat,
#     Q <: Union{Gegenbauer{β}, OrthonormalGegenbauer{β}},
#     P <: Union{Gegenbauer{α, T}, OrthonormalGegenbauer{α, T}}
# }

#     ps =  coeffs(p)
#     Q( ultra2ultra(ps, α, β, norm1=isorthonormal(P), norm2=isorthonormal(Q)) )

# end

# # Jacobi<--Jacobi
# function _convert(::Type{Q}, p::P) where {
#     γ, δ, α, β,  T <: AbstractFloat,
#     Q <: Union{Jacobi{γ, δ}, OrthonormalJacobi{γ, δ}},
#     P <: Union{Jacobi{α, β, T}, OrthonormalJacobi{α, β, T}}
# }

#     ps =  coeffs(p)
#     Q( jac2jac(ps, α, β, γ, δ,  norm1=isorthonormal(P), norm2=isorthonormal(Q)) )

# end

# # Laguerre<--Laguerre
# function _convert(::Type{Q}, p::P) where {
#     α, β,  T <: AbstractFloat,
#     Q <: Union{Laguerre{β}, OrthonormalLaguerre{β}},
#     P <: Union{Laguerre{α, T}, OrthonormalLaguerre{α, T}}
# }

#     ps =  coeffs(p)
#     Q( lag2lag(ps, α, β,  norm1=isorthonormal(P), norm2=isorthonormal(Q)) )

# end

# # Jacobi<--ultraspherical
# function _convert(::Type{Q}, p::P) where {
#     γ, δ, α,  T <: AbstractFloat,
#     Q <: Union{Jacobi{γ, δ}, OrthonormalJacobi{γ, δ}},
#     P <: Union{Gegenbauer{α, T}, OrthonormalGegenbauer{α, T}}
# }

#     ps =  coeffs(p)
#     Q( ultra2jac(ps, α, γ, δ,  normultra=isorthonormal(P), normjac=isorthonormal(Q)) )

# end

# # ultraspherical<--Jacobi
# function _convert(::Type{Q}, p::P) where {
#     α, γ, δ,  T <: AbstractFloat,
#     Q <: Union{Gegenbauer{α, T}, OrthonormalGegenbauer{α, T}},
#     P <: Union{Jacobi{γ, δ}, OrthonormalJacobi{γ, δ}}
# }

#     ps =  coeffs(p)
#     Q( jac2ultra(ps,γ, δ, α, normjac=isorthonormal(P), normultra=isorthonormal(Q)) )

# end

# # Jacobi<--Chebyshev
# function _convert(::Type{Q}, p::P) where {
#     γ, δ,  T <: AbstractFloat,
#     Q <: Union{Jacobi{γ, δ}, OrthonormalJacobi{γ, δ}},
#     P <: Union{Chebyshev{T}, OrthonormalChebyshev{T}}
# }

#     ps =  coeffs(p)
#     Q( cheb2jac(ps,γ, δ, normcheb=isorthonormal(P), normjac=isorthonormal(Q)) )

# end

# # Chebyshev<--Jacobi
# function _convert(::Type{Q}, p::P) where {
#     γ, δ,  T <: AbstractFloat,
#     Q <: Union{Chebyshev, OrthonormalChebyshev},
#     P <: Union{Jacobi{γ, δ,T}, OrthonormalJacobi{γ, δ,T}}
# }

#     ps =  coeffs(p)
#     Q( jac2cheb(ps,γ, δ, normjac = isorthonormal(P), normcheb=isorthonormal(Q)) )

# end

# # ultraspherical<--Chebyshev
# function _convert(::Type{Q}, p::P) where {
#     α,  T <: AbstractFloat,
#     Q <: Union{Gegenbauer{α}, OrthonormalGegenbauer{α}},
#     P <: Union{Chebyshev{T}, OrthonormalChebyshev{T}}
# }

#     ps =  coeffs(p)
#     Q( cheb2ultra(ps,α, normcheb=isorthonormal(P), normultra=isorthonormal(Q)) )

# end

# # Chebyshev<--ultraspherical
# function _convert(::Type{Q}, p::P) where {
#     α,  T <: AbstractFloat,
#     Q <: Union{Chebyshev, OrthonormalChebyshev},
#     P <: Union{Gegenbauer{α,T}, OrthonormalGegenbauer{α,T}}
# }
#     ps =  coeffs(p)
#     Q( ultra2cheb(ps,α, normultra=isorthonormal(P), normcheb=isorthonormal(Q)) )

# end

##  Koepf and Schmersa Thm 2, p11 connection for P to Q when σ = σ̂
function connection_m(::Type{P},::Type{Q},m,n) where {
    P <: AbstractCCOP,
    Q <: AbstractCCOP}

    a,b,c,d,e = abcde(P)
    ā,b̄,c̄,d̄,ē = abcde(Q)

    (a == ā && b == b̄ && c == c̄) || throw(ArgumentError("This connection assume σ = σ̄"))
    a²,b²,c²,d²,d̄²,e²,ē²,m²,n² = a^2, b^2, c^2, d^2, d̄^2, e^2, ē^2, m^2, n^2

    c0 = -(m-n) ⋅ (a⋅m + d - a + a⋅n) ⋅ (d̄ + 2a⋅m) ⋅ (d̄ + a + 2a⋅m) ⋅ (d̄ + 3a + 2a⋅m)
    c0 *= (d̄ + 2a⋅m  + 2a)^2

    c1 = -d⋅b⋅n⋅d̄ + 2d⋅a⋅m²⋅b + d⋅b⋅d̄ + 2d⋅a⋅m⋅b + 2d⋅ē⋅n⋅a
    c1 += d⋅d̄⋅ē + 2d⋅d̄⋅b⋅m - m⋅b⋅d̄² - e⋅d̄² - 4a²⋅m²⋅e - m²⋅a⋅b⋅d̄ + b⋅n⋅d̄⋅a - 2e⋅d̄⋅a
    c1 += -4a²⋅m⋅e - 4e⋅d̄⋅a⋅m + 2m²⋅a²⋅ē + 2ē⋅a²⋅n² - 2ē⋅a²⋅n - m⋅a⋅b⋅d̄ + 2m⋅d̄⋅ē⋅a
    c1 += 2m⋅ē⋅a² - b⋅n²⋅d̄⋅a
    c1 *= (d̄ + 2a⋅m + 2a) ⋅ (m + 1) ⋅  (d̄ + a + 2a⋅m) ⋅ (d̄ + 3a + 2a⋅m)

    c2 =  -(d̄ + 2a⋅m) ⋅ (m + 1) ⋅ (-a⋅m - 2a + a⋅n - d̄ + d) ⋅ (a⋅m + a⋅n + a + d̄)
    c2a = a⋅b²⋅m² - 4a²⋅m²⋅c - 8a²⋅m⋅c + 2a⋅m⋅b² - 4a⋅d̄⋅m⋅c + m⋅b²⋅d̄ - 4a⋅d̄⋅c - a⋅ē² + a⋅b² - c⋅d̄²
    c2a += b⋅ē⋅d̄ - 4⋅a²⋅c + b²⋅d̄
    c2 *= c2a ⋅ (m + 2)

    return (c0, c1, c2)

end

##  Koepf and Schmersa Thm 4, p14 connection for P, Q=x^n
function connection_m(::Type{P},::Type{Q},m,n) where {
    P <: AbstractCCOP,
    Q <: Polynomials.StandardBasisPolynomial}

    a,b,c,d,e = abcde(P)

    c0 = (m-n) ⋅ (a⋅n + d - a + a⋅m)

    c1 = (m + 1) ⋅ (b⋅m + e)

    c2 = c ⋅ (m + 1) ⋅ (m + 2)

    return (c0, c1, c2)

end

##  Koepf and Schmersa Thm 5, p16 connection for P=x^n, Q
function connection_m(::Type{P},::Type{Q},m,n) where {
    P <: Polynomials.StandardBasisPolynomial,
    Q <: AbstractCCOP}

    ā,b̄,c̄,d̄,ē = abcde(Q)
    ā²,b̄²,c̄²,d̄²,ē²,m²,n² = ā^2, b̄^2, c̄^2, d̄^2, ē^2, m^2, n^2

    c0 =  (n - m) ⋅ (d̄ + 2ā⋅m) ⋅ (d̄ + 3ā + 2ā⋅m) ⋅ (d̄ + ā + 2ā⋅m) ⋅ (d̄ + 2ā⋅m + 2ā)^2

    c1 =  (d̄⋅ē + b̄⋅d̄ +2d̄⋅b̄⋅m + 2ā⋅m²⋅b̄ + 2ā⋅m⋅b̄ +2ē⋅ā⋅n - d̄⋅b̄⋅n) ⋅ (d̄ + 2ā⋅m + 2ā)
    c1 *= (m+1) ⋅ (d̄ + 3ā + 2ā⋅m) ⋅  (d̄ + ā + 2ā⋅m)

    c2 = -(m+2)
    c2a = -4ā²⋅c̄⋅m²
    c2a +=  ā⋅b̄²⋅m² + 2ā⋅b̄²⋅m - 4ā⋅c̄⋅m⋅d̄ - 8ā²⋅c̄⋅m + m⋅b̄²⋅d̄ - ā⋅ē^2 - d̄²⋅c̄ + b̄⋅ē⋅d̄ -4ā²⋅c̄
    c2a += -4ā⋅c̄⋅d̄ + ā⋅b̄² +b̄²⋅d̄
    c2 *= c2a ⋅ (ā⋅m + ā⋅n + ā + d̄) ⋅ (m + 1) ⋅ (d̄ + 2ā⋅m)

    return (c0, c1, c2)

end

## With an explicit connectioni
function connection(::Type{P}, q::Q) where
    {P <: Polynomials.AbstractPolynomial,
     Q<:Polynomials.AbstractPolynomial}

    n = degree(q)
    T = eltype(q)
    T = eltype(one(eltype(one(P)))*one(q)/1)
    cs = zeros(T, n+1)

    for k in 0:n
        for (i,val) in Connection{P,Q}(n, k)
            cs[1+k] = muladd(q[i], val, cs[1+k])
        end
    end
    X = Polynomials.indeterminate(P,q)
    ⟒(P){T,X}(cs)

end


#---------------------------------------------------------
# # [Ideal age](@id ideal-age)
#---------------------------------------------------------

#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/notebooks/1_ideal_age.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/notebooks/1_ideal_age.ipynb)

#md # !!! note
#md #     All the AIBECS tutorials and how-to guides are available as Jupyter notebooks.
#md #     You can view them with [nbviewer](https://nbviewer.jupyter.org/)or execute them online with [binder](https://mybinder.org/) by clicking on the badges above!
#md #     (Note that binder can be slow to launch and its memory caps can be a problem when running.)

# The tracer equation for the ideal age is
#
# $$\left(\partial_t + \mathbf{T}\right) \boldsymbol{a} = 1 - \frac{\boldsymbol{a}}{τ} \, (\boldsymbol{z} \le z_0),$$
#
# where the sink term on the right clamps the age to $0$ at the surface (where $\boldsymbol{z} \le z_0$).
# The smaller the timescale $\tau$, the quicker $\boldsymbol{a}$ is restored to $0$ at the surface.

# AIBECS can interpret tracer equations as long as you arrange them under the generic form:
#
# $$\big(\partial_t + \mathbf{T}(\boldsymbol{p}) \big) \boldsymbol{x} = \boldsymbol{G}(\boldsymbol{x}, \boldsymbol{p}),$$
#
# where $\mathbf{T}(\boldsymbol{p})$ is the transport, $\boldsymbol{G}(\boldsymbol{x}, \boldsymbol{p})$ is the net local sources and sinks, and $\boldsymbol{p}$ is the vector of model parameters.
# We will then use the AIBECS to simulate the ideal age by finding the steady-state of the system, i.e., the solution of
#
# $$\partial_t \boldsymbol{x} = \boldsymbol{F}(\boldsymbol{x}, \boldsymbol{p}) = \boldsymbol{G}(\boldsymbol{x}, \boldsymbol{p}) - \mathbf{T}(\boldsymbol{p}) \, \boldsymbol{x} = 0.$$

# In this tutorial, we will simulate the ideal age by
# 1. defining functions for `T(p)` and `G(x,p)`,
# 1. defining the parameters `p`,
# 1. generating the state function `F(x,p)` and solving the associated steady-state problem,
# 1. and finally making a plot of our simulated ideal age.

# We start by telling Julia that we want to use the AIBECS package and the OCIM2 circulation
#nb # (the Ocean Circulation Inverse Model, see [1](https://doi.org/10.1029/2018JC014716) for details).
#md # (the Ocean Circulation Inverse Model[^1]).

#md # [^1]:
#md #     DeVries, T., & Holzer, M. (2019). Radiocarbon and helium isotope constraints on deep ocean ventilation and mantle‐³He sources. Journal of Geophysical Research: Oceans, 124, 3036–3057. doi:[10.1029/2018JC014716](https://doi.org/10.1029/2018JC014716)

using AIBECS
grd, TOCIM2 = OCIM2.load()

#md # !!! note
#md #     If it's your first time, Julia will ask you to download the OCIM2, in which case you should accept (i.e., type `y` and "return").
#md #     Once downloaded, AIBECS will remember where it downloaded the file and it will only load it from your laptop.
#nb # > **Note**
#nb # > If it's your first time, Julia will ask you to download the OCIM2, in which case you should accept (i.e., type `y` and "return").
#nb # > Once downloaded, AIBECS will remember where it downloaded the file and it will only load it from your laptop.

# `grd` is an `OceanGrid` object containing information about the 3D grid of the OCIM2 circulation and `TOCIM2` is the transport matrix representing advection and diffusion.

# The local sources and sinks for the age take the form

function G(x,p)
    @unpack τ, z₀ = p
    return @. 1 - x / τ * (z ≤ z₀)
end

# as per the tracer equation.
# The `@unpack` line unpacks the parameters `τ` and `z₀`.
# The `return` line returns the net sources and sinks.
# (The `@.` "macro" tells Julia that the operations apply to every element.)

# We can define the vector `z` of depths with `depthvec`.

z = depthvec(grd)

# Now we must construct a type for `p` the parameters.
# This type must contain our parameters `τ` and `z₀`.

struct IdealAgeParameters{U} <: AbstractParameters{U}
    τ::U
    z₀::U
end

# The type is now ready for us to generate an instance of the parameter `p`.
# Let's use `τ = 1.0` (s) and `z₀` the minimum depth of the model.

p = IdealAgeParameters(1.0, 30.0)

# We now use the AIBECS to generate the state function $\boldsymbol{F}$ (and its Jacobian) via

F = AIBECSFunction(TOCIM2, G)

# Now that `F` and `p` are defined, we are going to solve for the steady-state.
# But first, we must create a `SteadyStateProblem` object that contains `F`, `p`, and an initial guess `x_init` for the age.
# (`SteadyStateProblem` comes from [SciMLBase](https://github.com/JuliaDiffEq/SciMLBase.jl).)

# Let's make a vector of 0's for our initial guess.

nb = sum(iswet(grd))  # number of wet boxes
x_init = zeros(nb)    # Start with age = 0 everywhere

# Now we can create our `SteadyStateProblem` instance

prob = SteadyStateProblem(F, x_init, p)

# And finally, we can `solve` this problem, using the AIBECS `CTKAlg()` algorithm,

age = solve(prob, CTKAlg())

# This should take a few seconds.

# To conclude this tutorial, let's have a look at the age using AIBECS' plotting recipes and [Plots.jl](https://github.com/JuliaPlots/Plots.jl).

using Plots

# We first convert the age in years
# (because the default SI unit we used, i.e., seconds, is a bit small relative to global ocean timescales).

age_in_yrs = age * u"s" .|> u"yr"

# And we take a horizontal slice at about 2000m.

plothorizontalslice(age_in_yrs, grd, depth=2000u"m", color=:magma)

# Or look at the horiontal mean

plothorizontalmean(age_in_yrs, grd)

# That's it for this tutorial...
# Good job!


@doc raw"""
`SDE`: Stratonovich Stochastic Differential Equation

Defines a stochastic differential initial value problem
```math
\begin{aligned}
dq (t) &= v(t, q(t)) \, dt + B(t, q(t)) \circ dW , & q(t_{0}) &= q_{0} ,
\end{aligned}
```
with drift vector field ``v``, diffusion matrix ``B``,
initial conditions ``q_{0}``, the dynamical variable ``q``
taking values in ``\mathbb{R}^{d}``, and the m-dimensional Wiener process W

### Parameters

* `DT <: Number`: data type
* `TT <: Real`: time step type
* `AT <: AbstractArray{DT}`: array type
* `vType <: Function`: type of `v`
* `BType <: Function`: type of `B`
* `pType <: Union{NamedTuple,Nothing}`: parameters type

### Fields

* `d`:  dimension of dynamical variable ``q`` and the vector field ``v``
* `m`:  dimension of the Wiener process
* `ns`: number of sample paths
* `v`:  function computing the deterministic vector field
* `B`:  function computing the d x m diffusion matrix
* `t₀`: initial time
* `q₀`: initial condition for dynamical variable ``q`` (may be a random variable itself)
* `parameters`: either a `NamedTuple` containing the equations parameters or `nothing`
* `periodicity`: determines the periodicity of the state vector `q` for cutting periodic solutions

The functions `v` and `B`, providing the drift vector field and diffusion matrix,
`v(t, q, v)` and `B(t, q, B, col=0)`, where `t` is the current time, `q` is the
current solution vector, and `v` and `B` are the variables which hold the result
of evaluating the vector field ``v`` and the matrix ``B`` on `t` and `q` (if col==0),
or the column col of the matrix B (if col>0).

### Constructors

```julia
SDE(m, ns, v, B, t₀, q₀; parameters=NullParameters(), periodicity=zero(q₀[begin]))
SDE(m, ns, v, B, q₀::StateVector; kwargs...) = SDE(m, ns, v, B, 0.0, q₀; kwargs...)
SDE(m, ns, v, B, t₀, q₀::State; kwargs...) = SDE(m, ns, v, B, t₀, [q₀]; kwargs...)
SDE(m, ns, v, B, q₀::State; kwargs...) = SDE(m, ns, v, B, 0.0, q₀; kwargs...)
```

### Example

```julia
    function v(t, q, v, p)
        λ = p[:λ]
        v[1] = λ*q[1]
        v[2] = λ*q[2]
    end

    function B(t, q, B, p, col=0)
        μ = p[:μ]
        if col==0 #whole matrix
            B[1,1] = μ*q[1]
            B[2,1] = μ*q[2]
        elseif col==1
            #just first column
        end
    end

    t₀ = 0.
    q₀ = [1., 1.]
    λ  = 2.
    μ  = 1.
    p = (λ=λ, μ=μ)

    sde = SDE(v, B, t₀, q₀; parameters=p)
```
"""
struct SDE{dType <: Number, tType <: Real, arrayType <: AbstractArray{dType},
           vType <: Function, BType <: Function,
           pType <: OptionalParameters} <: AbstractEquationSDE{Nothing,Nothing,Nothing}

    d::Int
    m::Int
    ns::Int
    v::vType
    B::BType
    t₀::tType
    q₀::Vector{arrayType}
    parameters::pType
    periodicity::arrayType

    function SDE(m, ns, v::vType, B::BType, t₀::tType, q₀::Vector{arrayType};
                 parameters=NullParameters(), periodicity=zero(q₀[begin])) where {
                        dType <: Number, tType <: Real, arrayType <: AbstractArray{dType},
                        vType <: Function, BType <: Function}

        d  = length(q₀[begin])
        ni = length(q₀)

        @assert all(length(q) == d for q in q₀)

        @assert ns ≥ 1
        @assert ni == 1 || ns == 1
        # either multiple deterministic initial conditions and one sample path
        # or one deterministic initial condition and multiple sample paths

        new{dType,tType,arrayType,vType,BType,typeof(parameters)}(d, m, ns, v, B, t₀, q₀, parameters, periodicity)
    end
end

SDE(m, ns, v, B, q₀::StateVector; kwargs...) = SDE(m, ns, v, B, 0.0, q₀; kwargs...)
SDE(m, ns, v, B, t₀, q₀::State; kwargs...) = SDE(m, ns, v, B, t₀, [q₀]; kwargs...)
SDE(m, ns, v, B, q₀::State; kwargs...) = SDE(m, ns, v, B, 0.0, q₀; kwargs...)

# const SDEHT{HT,DT,TT,AT,VT,BT,PT} = SDE{DT,TT,AT,VT,BT,HT,PT} # type alias for dispatch on Hamiltonian type parameter
# const SDEPT{PT,DT,TT,AT,VT,BT,HT} = SDE{DT,TT,AT,VT,BT,HT,PT} # type alias for dispatch on parameters type parameter
const SDEPT{PT,DT,TT,AT,VT,BT} = SDE{DT,TT,AT,VT,BT,PT} # type alias for dispatch on parameters type parameter

Base.hash(sde::SDE, h::UInt) = hash(sde.d, hash(sde.m, hash(sde.ns,
                               hash(sde.v, hash(sde.B, hash(sde.t₀, hash(sde.q₀,
                               hash(sde.parameters, hash(sde.periodicity, h)))))))))

Base.:(==)(sde1::SDE, sde2::SDE) = (
                                sde1.d == sde2.d
                             && sde1.m == sde2.m
                             && sde1.ns == sde2.ns
                             && sde1.v == sde2.v
                             && sde1.B == sde2.B
                             && sde1.t₀ == sde2.t₀
                             && sde1.q₀ == sde2.q₀
                             && sde1.parameters == sde2.parameters
                             && sde1.periodicity == sde2.periodicity)

function Base.similar(equ::SDE, t₀::Real, q₀::StateVector, ns::Int=equ.ns)
    @assert all([length(q) == ndims(equ) for q in q₀])
    SDE(equ.m, ns, equ.v, equ.B, t₀, q₀; parameters=equ.parameters, periodicity=equ.periodicity)
end

Base.similar(equ::SDE, t₀::Real, q₀::State, ns::Int=equ.ns; kwargs...) = similar(equ, t₀, [q₀], ns; kwargs...)
Base.similar(equ::SDE, q₀::Union{State,StateVector}, ns::Int=equ.ns; kwargs...) = similar(equ, equ.t₀, q₀, ns; kwargs...)

Base.ndims(sde::SDE) = sde.d
Base.axes(equ::SDE) = axes(equ.q₀[begin])
GeometricBase.nsamples(equ::SDE) = length(equ.q₀)
GeometricBase.periodicity(equ::SDE) = equ.periodicity

initial_conditions(equ::SDE) = (equ.t₀, equ.q₀)

# hashamiltonian(::SDEHT{<:Nothing}) = false
# hashamiltonian(::SDEHT{<:Function}) = true

hasparameters(::SDEPT{<:NullParameters}) = false
hasparameters(::SDEPT{<:NamedTuple}) = true

_get_v(equ::SDE) = hasparameters(equ) ? (t,q,v) -> equ.v(t, q, v, equ.parameters) : equ.v
_get_B(equ::SDE) = hasparameters(equ) ? (t,q,B,col=0) -> equ.B(t, q, B, equ.parameters, col) : equ.B
# _get_h(equ::SDE) = hasparameters(equ) ? (t,q) -> equ.h(t, q, equ.parameters) : equ.h

function functions(equ::SDE)
    names = (:v,:B)
    equs  = (_get_v(equ), _get_B(equ))

    # if hashamiltonian(equ)
    #     names = (names..., :h)
    #     equs  = (equs..., _get_h(equ))
    # end

    NamedTuple{names}(equs)
end