# VectorAffineDecisionFunction #
# ========================== #
mutable struct VectorAffineDecisionConstraintBridge{T, S} <: MOIB.Constraint.AbstractBridge
    constraint::CI{MOI.VectorAffineFunction{T}, S}
    decision_function::VectorAffineDecisionFunction{T}
    set::S
end

function MOIB.Constraint.bridge_constraint(::Type{VectorAffineDecisionConstraintBridge{T, S}},
                                           model,
                                           f::VectorAffineDecisionFunction{T},
                                           set::S) where {T, S}
    # All decisions have been mapped to either the decision part constant or the variable part terms
    # at this point.
    F = MOI.VectorAffineFunction{T}
    # Calculate total constant
    constants = f.variable_part.constants +
        f.known_part.constants
    # Add the bridged constraint
    constraint = MOI.add_constraint(model,
                                    MOI.VectorAffineFunction(f.variable_part.terms, constants),
                                    set)
    # Save the constraint index, the decision function, and the set, to allow modifications
    return VectorAffineDecisionConstraintBridge{T, S}(constraint, f, set)
end

function MOI.supports_constraint(::Type{<:VectorAffineDecisionConstraintBridge{T}},
                                 ::Type{<:VectorAffineDecisionFunction{T}},
                                 ::Type{<:MOI.AbstractVectorSet}) where {T}
    return true
end
function MOIB.added_constrained_variable_types(::Type{<:VectorAffineDecisionConstraintBridge})
    return Tuple{DataType}[]
end
function MOIB.added_constraint_types(::Type{<:VectorAffineDecisionConstraintBridge{T, S}}) where {T, S}
    return [(MOI.VectorAffineFunction{T}, S)]
end
function MOIB.Constraint.concrete_bridge_type(::Type{<:VectorAffineDecisionConstraintBridge{T}},
                                              ::Type{<:VectorAffineDecisionFunction{T}},
                                              S::Type{<:MOI.AbstractVectorSet}) where T
    return VectorAffineDecisionConstraintBridge{T, S}
end

MOI.get(b::VectorAffineDecisionConstraintBridge{T, S}, ::MOI.NumberOfConstraints{MOI.VectorAffineFunction{T}, S}) where {T, S} = 1
MOI.get(b::VectorAffineDecisionConstraintBridge{T, S}, ::MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{T}, S}) where {T, S} = [b.constraint]

function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintFunction,
                 bridge::VectorAffineDecisionConstraintBridge{T}) where T
    f = bridge.decision_function
    # Remove mapped variables to properly unbridge
    from_decision(v) = begin
        result = any(t -> t.variable_index == v, f.decision_part.terms)
    end
    g = VectorAffineDecisionFunction(
        MOIU.filter_variables(v -> !from_decision(v), f.variable_part),
        copy(f.decision_part),
        copy(f.known_part))
    return g
end

function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintSet,
                 bridge::VectorAffineDecisionConstraintBridge{T}) where T
    return bridge.set
end

function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintPrimal,
                 bridge::VectorAffineDecisionConstraintBridge{T}) where T
    return MOI.get(model, attr, bridge.constraint)
end

function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintDual,
                 bridge::VectorAffineDecisionConstraintBridge{T}) where T
    return MOI.get(model, attr, bridge.constraint)
end

function MOI.delete(model::MOI.ModelLike, bridge::VectorAffineDecisionConstraintBridge)
    MOI.delete(model, bridge.constraint)
    return nothing
end

function MOI.set(model::MOI.ModelLike, ::MOI.ConstraintFunction,
                 bridge::VectorAffineDecisionConstraintBridge{T,S},
                 f::VectorAffineDecisionFunction{T}) where {T,S}
    # Update bridge functions and function constant
    bridge.decision_function = f
    # Recalculate total constants and modify constraint function
    constants = f.variable_part.constants +
        f.known_part.constants
    MOI.set(model, MOI.ConstraintFunction(), bridge.constraint,
            MOI.VectorAffineFunction(f.variable_part.terms, constants))
    return nothing
end

function MOI.set(model::MOI.ModelLike, ::MOI.ConstraintSet,
                 bridge::VectorAffineDecisionConstraintBridge{T,S}, change::S) where {T,S}
    f = bridge.decision_function
    bridge.set = change
    MOI.set(model, MOI.ConstraintSet(), bridge.constraint, change)
    return nothing
end

function MOI.modify(model::MOI.ModelLike, bridge::VectorAffineDecisionConstraintBridge{T,S}, change::MOI.VectorConstantChange) where {T,S}
    f = bridge.decision_function
    # Modify variable part of decision function
    f.variable_part.constants .= change.new_constant
    # Recalculate total constants and modify constraint function
    constants = f.variable_part.constants +
        f.known_part.constants
    MOI.set(model, MOI.ConstraintFunction(), bridge.constraint,
            MOI.VectorAffineFunction(f.variable_part.terms, constants))
    return nothing
end

function MOI.modify(model::MOI.ModelLike, bridge::VectorAffineDecisionConstraintBridge{T,S}, change::MOI.MultirowChange) where {T,S}
    f = bridge.decision_function
    # Modify variable part of decision function
    modify_coefficients!(f.variable_part.terms, change.decision, change.new_coefficients)
    # Modify decision part if variable originated from a decision
    if !isempty(findall(t -> t.scalar_term.variable_index == change.variable, f.decision_part.terms))
        modify_coefficients!(f.decision_part.terms, change.variable, change.new_coefficients)
    end
    # Modify the variable part of the constraint as usual
    MOI.modify(model, bridge.constraint, change)
    return nothing
end

function MOI.modify(model::MOI.ModelLike, bridge::VectorAffineDecisionConstraintBridge{T,S}, change::DecisionMultirowChange) where {T,S}
    f = bridge.decision_function
    # Update the decision coefficient
    modify_coefficients!(f.decision_part.terms, change.decision, change.new_coefficients)
    # Update mapped variable through MultiRowChange
    MOI.modify(model, bridge, MOI.MultiRowChange(change.decision, change.new_coefficients))
    return nothing
end

function MOI.modify(model::MOI.ModelLike, bridge::VectorAffineDecisionConstraintBridge{T,S}, change::KnownCoefficientChange) where {T,S}
    f = bridge.decision_function
    # Query known value
    known_value = MOI.get(model, MOI.VariablePrimal(), change.known)
    seen = Int[]
    for i in findall(t -> t.scalar_term.variable_index == change.decision, f.known_part.terms)
        # Update known part of constraint constant
        coefficient = f.known_part.terms[i].scalar_term.coefficient
        f.known_part.constants[i] +=
            (change.new_coefficient - coefficient) * known_value
        push!(seen, i)
    end
    for i in filter(i -> i ∉ seen, collect(1:MOI.output_dimension(f)))
        f.known_part.constants[i] += change.new_coefficient * known_value
    end
    # Update the known decision coefficient
    modify_coefficients!(f.known_part.terms, change.known, change.new_coefficients)
    # Recalculate total constants and modify constraint function
    constants = f.variable_part.constants +
        f.known_part.constants
    MOI.set(model, MOI.ConstraintFunction(), bridge.constraint,
            MOI.VectorAffineFunction(f.variable_part.terms, constants))
    return nothing
end

function MOI.modify(model::MOI.ModelLike, bridge::VectorAffineDecisionConstraintBridge{T,S}, change::KnownValueChange) where {T,S}
    f = bridge.decision_function
    seen = Int[]
    for i in findall(t -> t.scalar_term.variable_index == change.known, f.known_part.terms)
        # Update known part of constraint constant
        coefficient = f.decision_part.terms[i].scalar_term.coefficient
        f.known_part.constants[i] +=
            coefficient * change.value_difference
    end
    # Recalculate total constant and shift constraint set
    constant = f.variable_part.constant +
        f.known_part.constant
    MOI.set(model, MOI.ConstraintSet(), bridge.constraint,
            MOIU.shift_constant(bridge.set, -constant))
    return nothing
end

function MOI.modify(model::MOI.ModelLike, bridge::VectorAffineDecisionConstraintBridge{T,S}, change::KnownValuesChange) where {T,S}
    f = bridge.decision_function
    known_vals = zeros(T, MOI.output_dimension(f))
    for term in f.known_part.terms
        i = term.output_index
        scalar_term = term.scalar_term
        known_vals[i] += scalar_term.coefficient * MOI.get(model, MOI.VariablePrimal(), scalar_term.variable_index)
    end
    # Update known part of objective
    f.known_part.constants .= known_vals
    # Recalculate total constants and modify constraint function
    constants = f.variable_part.constants +
        f.known_part.constants
    MOI.set(model, MOI.ConstraintFunction(), bridge.constraint,
            MOI.VectorAffineFunction(f.variable_part.terms, constants))
    return nothing
end


module DistDict

using ComputeFramework
import ComputeFramework: split_range, PartitionScheme, alignfirst, DomainBranch, project, isempty, intersect, domain, cat_data

immutable DictDomain <: Domain
    hashrange::UnitRange
end

Base.isempty(d::DictDomain) = isempty(d.hashrange)
Base.intersect(d::DictDomain, d2::DictDomain) =
    DictDomain(intersect(d.hashrange, d2.hashrange))
project(d::DictDomain, d2::DictDomain) = d
Base.getindex(bigd::DictDomain, d::DictDomain) = d
alignfirst(a::DictDomain) = a

function Base.getindex(dict::Dict, d::DictDomain)
    res = Dict()
    for k in keys(dict)
        if hash(k) in d.hashrange
            res[k] = dict[k]
        end
    end
    res
end

immutable DictPartition <: PartitionScheme
    bins::Int
end

function ComputeFramework.partition(p::DictPartition, d::DictDomain)
    rs = collect(1:UInt(typemax(UInt)/p.bins+1):typemax(UInt))
    ranges = map(UnitRange, rs[1:end-1], rs[2:end])
    ComputeFramework.DomainBranch(d, map(DictDomain, ranges))
end

domain(x::Dict) = DictDomain(zero(UInt):(typemax(UInt)-1))

cat_data(x::DictPartition, dmn::DomainBranch, parts::AbstractArray) = merge(parts...)

Base.map(f, d::Dict) = Dict(Pair[k=>f(v) for (k,v) in d])

end # module


import RDates: WeekendCalendar, @rd_str, setcachedcalendar!
import Dates: Date

setcachedcalendar!(CALENDARS, "TARGET", WeekendCalendar() + RuleBasedCalendar([
    # New Years Day
    PeriodicCalendarRule(rd"1y", rd"1JAN"),
    # Labour Day
    WindowCalendarRule(PeriodicCalendarRule(rd"1y", rd"1MAY"); from=Date(2000,1,1)),
    # Good Friday
    WindowCalendarRule(PeriodicCalendarRule(rd"1y", rd"0E-1FRI"); from=Date(2000,1,1)),
    # Easter Monday
    WindowCalendarRule(PeriodicCalendarRule(rd"1y", rd"0E+1MON"); from=Date(2000,1,1)),
    # Christmas Day
    PeriodicCalendarRule(rd"1y", rd"25DEC"),
    # Day of Goodwill
    WindowCalendarRule(PeriodicCalendarRule(rd"1y", rd"26DEC"); from=Date(2000,1,1)),
    # End of Year
    ExplicitDateRule(inclusions=[Date(1998,12,31), Date(1999,12,31), Date(2001,12,31)]),
]))


using Revise
using Random
using AutomotiveDrivingModels
using AutoViz
using AutomotivePOMDPs
using AutomotiveSensors
using POMDPs
using POMDPSimulators
using POMDPPolicies
using StatsBase
using StaticArrays
using ProgressMeter
using Parameters
using Printf
using JLD2
using Flux
using Flux: mse, batchseq, throttle, @epochs
using Base.Iterators: partition
includet("environment.jl")
includet("cooperative_IDM.jl")
includet("generative_mdp.jl")

const Obs = Vector{Float64}
const Traj = Vector{Vector{Float64}}

function generate_trajectory(mdp::GenerativeMergingMDP, policy::Policy, max_steps::Int64, rng::AbstractRNG)
    s0 = initialstate(mdp, rng)
    svec = convert_s(Vector{Float64}, s0, mdp)
    hr = HistoryRecorder(max_steps=max_steps, rng=rng)
    hist = simulate(hr, mdp, policy, s0)
    # extract data from the history
    X = Vector{MVector{length(svec), Float64}}(undef, max_steps+1)
    for i=1:max_steps+1
        X[i] = zeros(length(svec))
    end
    fill!(X, zeros(length(svec)))
    X[1:n_steps(hist)+1] = convert_s.(Ref(Vector{Float64}), state_hist(hist), Ref(mdp))
    # build labels
    Y = Vector{Vector{Float64}}(undef, max_steps + 1)
    for (i,x) in enumerate(X)
        Y[i] = [x[6], x[9], x[12], x[15]]
        x[6] = 0.5
        x[9] = 0.5
        x[12] = 0.5
        x[15] = 0.5
    end
    return X, Y
end

function collect_set(mdp::GenerativeMergingMDP, policy::Policy, max_steps::Int64, rng::AbstractRNG, n_set)
    X_batch = Vector{Traj}(undef, n_set)
    Y_batch = Vector{Traj}(undef, n_set)
    @showprogress for i=1:n_set
        X_batch[i], Y_batch[i] = generate_trajectory(mdp, policy, max_steps, rng)
    end
    return X_batch, Y_batch
end

rng = MersenneTwister(1)

mdp = GenerativeMergingMDP(random_n_cars = true, driver_type = :random, observe_cooperation = true)

# @load "mixed_policy.jld2" policy

policy = RandomPolicy(mdp, rng=rng)

max_steps = 40
n_train = 1000
n_val = 100

X_train, Y_train = collect_set(mdp, policy, max_steps, rng, n_train);
X_val, Y_val = collect_set(mdp, policy, max_steps, rng, n_val);

input_dims = length(X_train[1][1])
output_dims = length(Y_train[1][1])
model = Chain(LSTM(input_dims, 64), Dense(64, 32, relu), Dense(32, output_dims, σ))
model = Chain(Dense(input_dims, output_dims, σ)) # logistic regression

function loss(xs, ys)
    # l = sum(Flux.crossentropy.(model.(xs), ys))
    l = mean(Flux.mse.(model.(xs), ys))
    Flux.truncate!(model)
    Flux.reset!(model)
    return l
end

function accuracy(x, y)
   mean( (model(x) .> 0.5) .== y)
end

function sequence_accuracy(x, y)
    acc = mean(accuracy.(x, y))
    Flux.truncate!(model)
    Flux.reset!(model)
    return acc
end

batch_size = 32
xv = batchseq(X_val)
yv = batchseq(Y_val)
# xt = batchseq.(partition(X_train, batch_size))
# yt = batchseq.(partition(Y_train, batch_size))
xt = batchseq(X_train)
yt = batchseq(Y_train)

evalcb = () -> @printf("eval loss %2.2f | train loss %2.2f | eval acc %2.2f | train acc %2.2f \n ", loss(xv, yv), loss(xt, yt), sequence_accuracy(xv, yv), sequence_accuracy(xt, yt))

opt = ADAM(1e-2)
@epochs 1000 Flux.train!(loss, params(model), [(xt, yt)], opt, cb=evalcb)

loss(xt, yt)

sequence_accuracy(xt, yt)

for data in zip(xt, yt)
    @show loss(data...)
end

loss(xt, yt)

(model(xv[1]) .> 0.5) .== yv[1]

sequence_accuracy(xv, yv)

model.(xv)

@time loss(xv, yv)
# loss(X_val, Y_val)

size(X_train)
size(X_train[1])
size(X_train[1][1])

size(batchseq(X_train))
size(batchseq(X_train)[1])
size(batchseq(X_train)[1][1])

endind = findfirst(isequal(zeros(input_dims)), X_train[1])
state_hist = convert_s.(Ref(AugScene), X_train[1][1:endind], Ref(mdp))

using Blink
using Interact
ui = @manipulate for s=1:length(state_hist)
    AutoViz.render(state_hist[s].scene, mdp.env.roadway,
                   cam = StaticCamera(VecE2(-25.0, -10.0), 6.0))
end

w = Window()
body!(w, ui);

X_train[1][1]

s0 = initialstate(mdp, rng)
AutoViz.render(s0.scene, mdp.env.roadway, cam = StaticCamera(VecE2(-25.0, -10.0), 6.0))


using Printf

mutable struct GameLayer <: Ranger.AbstractNode
    base::Nodes.NodeData
    transform::Nodes.TransformProperties{Float64}

    # Collection of nodes.
    children::Array{Ranger.AbstractNode,1}

    min::Geometry.Point{Float64}
    max::Geometry.Point{Float64}
    buc_min::Geometry.Point{Float64}
    buc_max::Geometry.Point{Float64}

    host_node::HostNode

    function GameLayer(world::Ranger.World, name::String, parent::Ranger.AbstractNode)
        o = new()

        o.base = Nodes.NodeData(name, parent, world)
        o.children = Array{Ranger.AbstractNode,1}[]
        o.transform = Nodes.TransformProperties{Float64}()

        o.min = Geometry.Point{Float64}()
        o.max = Geometry.Point{Float64}()
        o.buc_min = Geometry.Point{Float64}()
        o.buc_max = Geometry.Point{Float64}()

        o
    end
end

function build(layer::GameLayer, world::Ranger.World)
    hw = Float64(world.view_width) / 2.0
    hh = Float64(world.view_height) / 2.0

    # top-left
    Geometry.set!(layer.min, -hw, -hh)

    # bottom-right
    Geometry.set!(layer.max, hw, hh)

    layer.host_node = HostNode(world, "HostNode", layer)
    Nodes.set_scale!(layer.host_node, 150.0)
    push!(layer.children, layer.host_node);

end

# --------------------------------------------------------
# Timing
# --------------------------------------------------------
function Nodes.update(layer::GameLayer, dt::Float64)
end

# --------------------------------------------------------
# Rendering
# --------------------------------------------------------
function Nodes.draw(layer::GameLayer, context::Rendering.RenderContext)
    if Nodes.is_dirty(layer)
        Rendering.transform!(context, layer.min, layer.max, layer.buc_min, layer.buc_max)
        Nodes.set_dirty!(layer, false)
    end

    # Render background (i.e. a solid gray color)
    Rendering.set_draw_color(context, RangerGame.darkgray)
    Rendering.render_aa_rectangle(context, layer.buc_min, layer.buc_max, Rendering.FILLED)

    # Draw layer name
    Rendering.set_draw_color(context, RangerGame.white)
    Rendering.draw_text(context, 10, 10, layer.base.name, 2, 2, false)
end

# --------------------------------------------------------
# Life cycle events
# --------------------------------------------------------
function Nodes.enter_node(layer::GameLayer, man::Nodes.NodeManager)
    println("enter ", layer)

    # Register node as a timing target in order to receive updates
    Nodes.register_target(man, layer)

    # Register host so it too can get updates
    Nodes.register_target(man, layer.host_node)

    # Register for io events such as keyboard or mouse
    Nodes.register_event_target(man, layer)
end

function Nodes.exit_node(layer::GameLayer, man::Nodes.NodeManager)
    println("exit ", layer)
    Nodes.unregister_target(man, layer)
    Nodes.unregister_target(man, layer.host_node)

    Nodes.unregister_event_target(man, layer)
end

# --------------------------------------------------------
# Events
# --------------------------------------------------------

# --------------------------------------------------------
# Grouping
# --------------------------------------------------------
function Nodes.get_children(node::GameLayer)
    node.children
end


result = [2.32551,
-0.155293,
0.469214,
-0.0870457,
0.271364,
-0.222398,
1.01423,
-0.179009,
1.01724,
-0.152972,
0.216039,
-0.251637,
1.05034,
-0.231289,
0.245388,
-0.231289,
1.03371,
-0.258872,
1.03256,
-0.262519,
1.04748,
-0.259143,
1.0535,
-0.266484,
1.04711,
-0.267177,
1.02131,
-0.280381]

include(joinpath(BASE_DIR, "test", "data_5bus_pu.jl"))
include(joinpath(BASE_DIR, "test", "data_14bus_pu.jl"))

nodes_5 = nodes5()
nodes_14 = nodes14()
c_sys14 = System(nodes_14, thermal_generators14(nodes_14), loads14(nodes_14), branches14(nodes_14), nothing, 100.0, nothing, nothing);

c_sys5_re = System(nodes_5, vcat(thermal_generators5(nodes_5), renewable_generators5(nodes_5)), loads5(nodes_5),
                nothing, nothing,  100.0, nothing, nothing)

@testset begin
    using NLsolve
    pf!, x0 = make_pf(c_sys14);
    res = NLsolve.nlsolve(pf!, x0)
    @test res.zero ≈ result rtol=1e-3

    solve_powerflow!(c_sys14, nlsolve, method = :newton)
    pf!, x0 = make_pf(c_sys5_re);
    res = NLsolve.nlsolve(pf!, x0)
    @test !res.f_converged
end


# For better tree printing and navigation
## Things we need to define
function AbstractTrees.children(node::TrapezoidalSearchNode)
    if node.left === nothing
        if node.right === nothing
            return ()
        else
            return (node.right, )
        end
    else
        if node.right === nothing
            return (node.left, )
        else
            return (node.left, node.right)
        end
    end
end

## Things that make printing prettier
AbstractTrees.printnode(io::IO, node::TrapezoidalSearchNode) = print(io, node)

function buildmega(table)
    s = 0
    for x in table.geometry
        s += length(x.points) - length(x.parts)
    end
    res = Vector{EData{typeof(table.geometry[1].points[1].x)}}(undef, s)
    offset = 0
    for (i, poly) in pairs(table.geometry)
        buildedges!(res, offset, poly, i)
        offset += length(poly.points) - length(poly.parts)
    end

    return res
end

function buildedges!(edges, offset, poly, id)
    points = poly.points
    parts = poly.parts

    @inbounds for k in eachindex(parts)
        i1 = parts[k] + 2
        i2 = k == length(parts) ? length(points) : parts[k + 1]
        prev = V(points[i1 - 1].x, points[i1 - 1].y)
        for i in i1:i2
            p = points[i]
            cur = V(p.x, p.y)
            if prev < cur
                edges[i - k + offset] = EData(E(prev, cur), -1, id)
            else
                edges[i - k + offset] = EData(E(cur, prev), id, -1)
            end
            prev = cur
        end
    end

    return edges
end

function buildedges(poly, id)
    points = poly.points
    parts = poly.parts
    res = Vector{EData{typeof(points[1].x)}}(undef, length(points) - length(parts))
    @inbounds for k in eachindex(parts)
        i1 = parts[k] + 2
        i2 = k == length(parts) ? length(points) : parts[k + 1]
        prev = V(points[i1 - 1].x, points[i1 - 1].y)
        for i in i1:i2
            p = points[i]
            cur = V(p.x, p.y)
            if prev < cur
                res[i - k] = EData(E(prev, cur), -1, id)
            else
                res[i - k] = EData(E(cur, prev), id, -1)
            end
            prev = cur
        end
    end

    return res
end

function names2indices(baddies, table)
    m = Dict(map(x -> x[2] => x[1], enumerate(table.tzid)))
    res = Tuple{Int, Int}[]
    for x in baddies
        push!(res, (m[x[1]], m[x[2]]))
    end
    return res
end

function mergeedges(edges, table; precedence = PRECEDENCE)
    precedence = names2indices(precedence, table)
    baddies = Set{Tuple{Int, Int}}()
    edges = sort(edges, by = x -> edge(x))
    res = similar(edges, 0)
    prev = edges[1]
    for i in 2:length(edges)
        cur = edges[i]
        if edge(cur) != edge(prev)
            push!(res, prev)
            prev = cur
        else
            if lhs(cur) != -1 && lhs(prev) == -1
                @set! prev.l = cur.l
            elseif cur.r != -1 && prev.r == -1
                @set! prev.r = cur.r
            else
                caught = false
                for prec in precedence
                    if (cur.l == prec[1] && prev.l == prec[2]) || (cur.l == prec[2] && prev.l == prec[1])
                        @set! prev.l = prec[1]
                        caught = true
                    elseif (cur.r == prec[1] && prev.r == prec[2]) || (cur.r == prec[2] && prev.r == prec[1])
                        @set! prev.r = prec[1]
                        caught = true
                    end
                    caught && break
                end
                if !caught
                    if cur.l != -1 && prev.l != -1
                        m1 = min(cur.l, prev.l)
                        m2 = max(cur.l, prev.l)
                        push!(baddies, (m1, m2))
                    elseif cur.r != -1 && prev.r != -1
                        m1 = min(cur.r, prev.r)
                        m2 = max(cur.r, prev.r)
                        push!(baddies, (m1, m2))
                    end
                end
            end
        end
    end
    push!(res, prev)

    return res, baddies
end

function boundingbox(poly)
    points = poly.points
    p, rest = Iterators.peel(points)
    p = V(p.x, p.y)
    leftp = p
    rightp = p
    topy = p.y
    bottomy = p.y
    for p in rest
        p = V(p.x, p.y)
        if p < leftp
            leftp = p
        end
        if p > rightp
            rightp = p
        end
        if p.y > topy
            topy = p.y
        end
        if p.y < bottomy
            bottomy = p.y
        end
    end
    top = E(V(leftp.x, topy), V(rightp.x, topy))
    bottom = E(V(leftp.x, bottomy), V(rightp.x, bottomy))

    return Trapezoid(leftp, rightp, top, bottom, -1)
end

function boundingbox(v::AbstractVector{T}) where {T <: EData}
    ed, rest = Iterators.peel(v)
    p = edge(ed).p
    leftp = p
    rightp = p
    topy = p.y
    bottomy = p.y
    for ed in rest
        p = edge(ed).p
        if p < leftp
            leftp = p
        end
        if p > rightp
            rightp = p
        end
        if p.y > topy
            topy = p.y
        end
        if p.y < bottomy
            bottomy = p.y
        end
    end
    top = E(V(leftp.x, topy), V(rightp.x, topy))
    bottom = E(V(leftp.x, bottomy), V(rightp.x, bottomy))

    return Trapezoid(leftp, rightp, top, bottom, -1)
end

########################################
# Angle sorting facilities
########################################

function pseudoangle1(v)
    atan(v.y, v.x)
end

function pseudoangle2(v)
    copysign(1. - v.x/(abs(v.x)+abs(v.y)), v.y)
end

# http://steve.hollasch.net/cgindex/math/fowler.html
# https://stackoverflow.com/a/27253590/1494135
function pseudoangle3(v)
    # 1 for above, 0 for below the diagonal/anti-diagonal
    diag  = v.x > v.y
    adiag = v.x > -v.y

    r = !adiag ? 4 : 0

    if (v.y == 0)
        return r
    end

    if (diag ⊻ adiag)
        r += 2 - v.x / v.y
    else
        r += v.y / v.x
    end

    return r
end

function isangleless(v1, v2)
    a1 = abs(v1.x) + abs(v1.y)
    a2 = abs(v2.x) + abs(v2.y)

    a2*copysign(a1 - v1.x, v1.y) < a1*copysign(a2 - v2.x, v2.y)
end


@testset "Cellular automata 1D" begin
  using Agents.CA1D
  rules = Dict("111"=>"0", "110"=>"0", "101"=>"0", "100"=>"1", "011"=>"0",
  "010"=>"1", "001"=>"1", "000"=>"0")  # rule 22
  ncols = 101
  model = CA1D.build_model(rules=rules, ncols=ncols)
  model.agents[51].status="1"
  runs = 10
  data = CA1D.ca_run(model, runs);
  @test size(data) == (ncols * (runs + 1), 4)
end

function dummyplot(data; nodesize=2.0, anim=2)
end

@testset "Cellular automata 2D" begin
  using Agents.CA2D
  rules = (2,3,3)
  dims=(10, 10)
  model = CA2D.build_model(rules=rules, dims=dims, Moore=true)
  for i in 1:nv(model)
    if rand() < 0.1
      model.agents[i].status="1"
    end
  end

  runs = 2
  anim = CA2D.ca_run(model, runs, dummyplot);

  @test anim == nothing
end


function AbstractRegion(arg0::Vector{Hyperplane}, arg1::jdouble)
    return AbstractRegion((Vector{Hyperplane}, jdouble), arg0, arg1)
end

function apply_transform(obj::AbstractRegion, arg0::Transform)
    return jcall(obj, "applyTransform", AbstractRegion, (Transform,), arg0)
end

function build_new(obj::AbstractRegion, arg0::BSPTree)
    return jcall(obj, "buildNew", AbstractRegion, (BSPTree,), arg0)
end

function check_point(obj::AbstractRegion, arg0::JVector)
    return jcall(obj, "checkPoint", Region_Location, (JVector,), arg0)
end

function check_point(obj::AbstractRegion, arg0::Point)
    return jcall(obj, "checkPoint", Region_Location, (Point,), arg0)
end

function contains(obj::AbstractRegion, arg0::Region)
    return jcall(obj, "contains", jboolean, (Region,), arg0)
end

function copy_self(obj::AbstractRegion)
    return jcall(obj, "copySelf", AbstractRegion, ())
end

function equals(obj::Object, arg0::Object)
    return jcall(obj, "equals", jboolean, (Object,), arg0)
end

function get_barycenter(obj::AbstractRegion)
    return jcall(obj, "getBarycenter", Point, ())
end

function get_boundary_size(obj::AbstractRegion)
    return jcall(obj, "getBoundarySize", jdouble, ())
end

function get_class(obj::Object)
    return jcall(obj, "getClass", Class, ())
end

function get_size(obj::AbstractRegion)
    return jcall(obj, "getSize", jdouble, ())
end

function get_tolerance(obj::AbstractRegion)
    return jcall(obj, "getTolerance", jdouble, ())
end

function get_tree(obj::AbstractRegion, arg0::jboolean)
    return jcall(obj, "getTree", BSPTree, (jboolean,), arg0)
end

function hash_code(obj::Object)
    return jcall(obj, "hashCode", jint, ())
end

function intersection(obj::AbstractRegion, arg0::SubHyperplane)
    return jcall(obj, "intersection", SubHyperplane, (SubHyperplane,), arg0)
end

function is_empty(obj::AbstractRegion)
    return jcall(obj, "isEmpty", jboolean, ())
end

function is_empty(obj::AbstractRegion, arg0::BSPTree)
    return jcall(obj, "isEmpty", jboolean, (BSPTree,), arg0)
end

function is_full(obj::AbstractRegion)
    return jcall(obj, "isFull", jboolean, ())
end

function is_full(obj::AbstractRegion, arg0::BSPTree)
    return jcall(obj, "isFull", jboolean, (BSPTree,), arg0)
end

function notify(obj::Object)
    return jcall(obj, "notify", void, ())
end

function notify_all(obj::Object)
    return jcall(obj, "notifyAll", void, ())
end

function project_to_boundary(obj::AbstractRegion, arg0::Point)
    return jcall(obj, "projectToBoundary", BoundaryProjection, (Point,), arg0)
end

function to_string(obj::Object)
    return jcall(obj, "toString", JString, ())
end

function wait(obj::Object)
    return jcall(obj, "wait", void, ())
end

function wait(obj::Object, arg0::jlong)
    return jcall(obj, "wait", void, (jlong,), arg0)
end

function wait(obj::Object, arg0::jlong, arg1::jint)
    return jcall(obj, "wait", void, (jlong, jint), arg0, arg1)
end


function compare_power_flows(first_network, second_network, plot_dir)
    if split(lowercase(split(first_network, "/")[end]), '.')[end] == "json"
        first_mn_data = JSON.parsefile(first_network)
    else
        first_mn_data = PowerModels.parse_file(first_network)
    end
    # println("Finished parsing the first network")
    if split(lowercase(split(second_network, "/")[end]), '.')[end] == "json"
        second_mn_data = JSON.parsefile(second_network)
    else
        second_mn_data = PowerModels.parse_file(second_network)
    end
    # println("Finished parsing the second network")

    first_ref = PowerModels.build_ref(first_mn_data)[:sn]
    # println("Finished building ref for the first network")
    second_ref = PowerModels.build_ref(second_mn_data)[:sn]
    # println("Finished building ref for the second network")

    table_string = ""
    for (subnet_idx, ref_subnet) in first_ref
        println("subnet_idx: $subnet_idx")
        if (subnet_idx) == 0 && !(subnet_idx in keys(second_ref))
            subnet_idx = 1
        end
        if (subnet_idx) == 1 && !(subnet_idx in keys(second_ref))
            subnet_idx = 0
        end
        # total, max, min, plot all
        gen_compared_idx = Array{Number}(undef, length(ref_subnet[:gen]))
        gen_pg1 = Array{Float64}(undef, length(ref_subnet[:gen]))
        gen_qg1 = Array{Float64}(undef, length(ref_subnet[:gen]))
        gen_vm1 = Array{Float64}(undef, length(ref_subnet[:gen]))
        gen_pg2 = Array{Float64}(undef, length(ref_subnet[:gen]))
        gen_qg2 = Array{Float64}(undef, length(ref_subnet[:gen]))
        gen_vm2 = Array{Float64}(undef, length(ref_subnet[:gen]))
        for (i,(gen_i,gen)) in enumerate(ref_subnet[:gen])
            if ref_subnet[:gen][gen_i]["gen_status"] == 0
                gen_pg1[i] = NaN
                gen_qg1[i] = NaN
                gen_vm1[i] = NaN
                gen_pg2[i] = NaN
                gen_qg2[i] = NaN
                gen_vm2[i] = NaN
                gen_compared_idx[i] = NaN
            else
                gen_pg1[i] = ref_subnet[:gen][gen_i]["pg"]
                gen_qg1[i] = ref_subnet[:gen][gen_i]["qg"]
                gen_vm1[i] = ref_subnet[:gen][gen_i]["vg"]
                gen_pg2[i] = second_ref[subnet_idx][:gen][gen_i]["pg"]
                gen_qg2[i] = second_ref[subnet_idx][:gen][gen_i]["qg"]
                gen_vm2[i] = second_ref[subnet_idx][:gen][gen_i]["vg"]
                gen_compared_idx[i] = gen_i
            end
        end
        filter!(e->e≠NaN,gen_pg1)
        filter!(e->e≠NaN,gen_qg1)
        filter!(e->e≠NaN,gen_vm1)
        filter!(e->e≠NaN,gen_pg2)
        filter!(e->e≠NaN,gen_qg2)
        filter!(e->e≠NaN,gen_vm2)
        filter!(e->e≠NaN,gen_compared_idx)
        gen_pg_compared = gen_pg2 - gen_pg1
        gen_qg_compared = gen_qg2 - gen_qg1
        gen_vm_compared = gen_vm2 - gen_vm1

        # bus_vm_compared = Array{Float64}(undef, length(ref_subnet[:bus]))
        # bus_va_compared = Array{Float64}(undef, length(ref_subnet[:bus]))
        bus_compared_idx = Array{Number}(undef, length(ref_subnet[:bus]))
        bus_vm1 = Array{Float64}(undef, length(ref_subnet[:bus]))
        bus_va1 = Array{Float64}(undef, length(ref_subnet[:bus]))
        bus_vm2 = Array{Float64}(undef, length(ref_subnet[:bus]))
        bus_va2 = Array{Float64}(undef, length(ref_subnet[:bus]))
        for (i,(bus_i,bus)) in enumerate(ref_subnet[:bus])
            if bus["bus_type"] ∉ [1,2,3]
                bus_vm1[i] = NaN
                bus_va1[i] = NaN
                bus_vm2[i] = NaN
                bus_va2[i] = NaN
                bus_compared_idx[i] = NaN
            else
                bus_vm1[i] = ref_subnet[:bus][bus_i]["vm"]
                bus_va1[i] = ref_subnet[:bus][bus_i]["va"]
                bus_vm2[i] = second_ref[subnet_idx][:bus][bus_i]["vm"]
                bus_va2[i] = second_ref[subnet_idx][:bus][bus_i]["va"]
                bus_compared_idx[i] = bus_i
            end
        end
        filter!(e->e≠NaN,bus_vm1)
        filter!(e->e≠NaN,bus_va1)
        filter!(e->e≠NaN,bus_vm2)
        filter!(e->e≠NaN,bus_va2)
        filter!(e->e≠NaN,bus_compared_idx)
        bus_vm_compared = bus_vm2 - bus_vm1
        bus_va_compared = bus_va2 - bus_va1

        branch_compared_idx = Array{Number}(undef, length(ref_subnet[:branch]))
        branch_pt1 = Array{Float64}(undef, length(ref_subnet[:branch]))
        branch_pf1 = Array{Float64}(undef, length(ref_subnet[:branch]))
        branch_qt1 = Array{Float64}(undef, length(ref_subnet[:branch]))
        branch_qf1 = Array{Float64}(undef, length(ref_subnet[:branch]))
        branch_pt2 = Array{Float64}(undef, length(ref_subnet[:branch]))
        branch_pf2 = Array{Float64}(undef, length(ref_subnet[:branch]))
        branch_qt2 = Array{Float64}(undef, length(ref_subnet[:branch]))
        branch_qf2 = Array{Float64}(undef, length(ref_subnet[:branch]))
        for (i,(branch_i,branch)) in enumerate(ref_subnet[:branch])
            if ("pt" ∉ keys(branch)) || ("pt" ∉ keys(second_ref[subnet_idx][:branch][branch_i]))
                branch_pt1[i] = NaN
                branch_pf1[i] = NaN
                branch_qt1[i] = NaN
                branch_qf1[i] = NaN
                branch_pt2[i] = NaN
                branch_pf2[i] = NaN
                branch_qt2[i] = NaN
                branch_qf2[i] = NaN
                branch_compared_idx[i] = NaN
            else
                branch_pt1[i] = ref_subnet[:branch][branch_i]["pt"]
                branch_pf1[i] = ref_subnet[:branch][branch_i]["pf"]
                branch_qt1[i] = ref_subnet[:branch][branch_i]["qt"]
                branch_qf1[i] = ref_subnet[:branch][branch_i]["qf"]
                branch_pt2[i] = second_ref[subnet_idx][:branch][branch_i]["pt"]
                branch_pf2[i] = second_ref[subnet_idx][:branch][branch_i]["pf"]
                branch_qt2[i] = second_ref[subnet_idx][:branch][branch_i]["qt"]
                branch_qf2[i] = second_ref[subnet_idx][:branch][branch_i]["qf"]
                branch_compared_idx[i] = branch_i
            end
        end
        filter!(e->e≠NaN,branch_pt1)
        filter!(e->e≠NaN,branch_pf1)
        filter!(e->e≠NaN,branch_qt1)
        filter!(e->e≠NaN,branch_qf1)
        filter!(e->e≠NaN,branch_pt2)
        filter!(e->e≠NaN,branch_pf2)
        filter!(e->e≠NaN,branch_qt2)
        filter!(e->e≠NaN,branch_qf2)
        filter!(e->e≠NaN,branch_compared_idx)
        branch_pt_compared = branch_pt2 - branch_pt1
        branch_pf_compared = branch_pf2 - branch_pf1
        branch_qt_compared = branch_qt2 - branch_qt1
        branch_qf_compared = branch_qf2 - branch_qf1

        GR.inline("pdf")

        output_dir = "$(plot_dir)/subnet$(subnet_idx)"
        if !isdir(output_dir)
            mkpath(output_dir)
        end
        upscale = 2 #upscaling in resolution
        fntsm = font("serif", pointsize=round(8.0*upscale))
        fntlg = font("serif", pointsize=round(12.0*upscale))
        default(titlefont=fntlg, guidefont=fntlg, tickfont=fntsm, legendfont=fntsm)
        default(size=(800*upscale,600*upscale)) #Plot canvas size

        gr()

        plot_data = [gen_pg_compared, gen_qg_compared,
                     bus_vm_compared, bus_va_compared,
                     branch_pt_compared, branch_pf_compared, branch_qt_compared, branch_qf_compared
                     ]
        raw_data = [gen_pg1,
                    gen_qg1,
                    bus_vm1,
                    bus_va1,
                    branch_pt1,
                    branch_pf1,
                    branch_qt1,
                    branch_qf1
                    ]
        plot_idx = [gen_compared_idx, gen_compared_idx,
                    bus_compared_idx, bus_compared_idx,
                    branch_compared_idx, branch_compared_idx, branch_compared_idx, branch_compared_idx
                    ]
        plot_labels = ["gen_pg", "gen_qg",
                       "bus_vm", "bus_va",
                       "branch_pt", "branch_pf", "branch_qt", "branch_qt"
                       ]

        for (i,series) in enumerate(plot_data)

            series_idx = plot_idx[i]

            (max_val, max_i) = findmax(series)
            (min_val, min_i) = findmin(series)
            idx_max = series_idx[max_i]
            idx_min = series_idx[min_i]

            # tol = 1e-3
            # worst = series[abs.(series).>tol]
            # idx_worst = series_idx[abs.(series).>tol]
            perm_worst = sortperm(abs.(series), rev=true)
            ten_worst = series[perm_worst][1:10]
            idx_ten_worst = series_idx[perm_worst][1:10]
            percent = ten_worst ./ raw_data[i][perm_worst][1:10]

            result_string = "$(plot_labels[i])\n==========\nMaximum difference: $(maximum(series))\nMinimum difference: $(minimum(series))\nMean absolute difference: $(mean(abs.(series)))\n"*
                            "biggest differences: $ten_worst\nat indices: $idx_ten_worst\npercent difference = $percent\n"

            println(result_string)

            open("$output_dir/comparison_log.txt", "a") do f
                write(f, result_string)
            end

            table_string = "$table_string,$(max_val),$(min_val),$(mean(abs.(series)))"

            perm = sortperm(series_idx)
            x_sorted = series_idx[perm]
            y_sorted = series[perm]
            series_label = plot_labels[i]
            plt = plot(1:length(x_sorted), y_sorted,
                label=series_label,
                legend=true,
                left_margin=(10*upscale)*mm,
                right_margin=(10*upscale)*mm, top_margin=(10*upscale)*mm, bottom_margin=(10*upscale)*mm,
                linetype=:scatter,
                marker = (:dot, 1.5*upscale, 0.6, :green, stroke(0))
            )
            xlabel!(plt, "index")
            ylabel!(plt, series_label)

            savefig(plt,"$(output_dir)/$(series_label).pdf")

            bar_plot = false
            if bar_plot
                plt = bar(x_sorted, y_sorted,
                    label=series_label,
                    legend=false,
                    left_margin=(10*upscale)*mm,
                    right_margin=(10*upscale)*mm, top_margin=(10*upscale)*mm, bottom_margin=(10*upscale)*mm
                )
                xlabel!(plt, "index")
                ylabel!(plt, series_label)
                # savefig(p,"$(folder)/$(y_label) vs $(x_label).pdf")
                if !isdir(output_dir)
                    mkpath(output_dir)
                end
                # println("saving plot at $(output_folder)/plots/$plotdir/$(y_label) vs $(x_label).pdf")
                # display(plt)
                savefig(plt,"$(output_dir)/$(series_label)_bar.pdf")
            end
        end

        plot_raw_data = [[gen_pg1, gen_pg2],
                         [gen_qg1, gen_qg1],
                         [bus_vm1, bus_vm2],
                         [bus_va1, bus_va2]
                         ]
        plot_idx = [gen_compared_idx,
                    gen_compared_idx,
                    bus_compared_idx,
                    bus_compared_idx
                    ]
        plot_labels = ["gen_pg values",
                       "gen_qg values",
                       "bus_vm values",
                       "bus_va values"
                       ]
        for (i,series) in enumerate(plot_raw_data)
            series_idx = plot_idx[i]
            perm = sortperm(series_idx)
            x_sorted = series_idx[perm]
            series_label = plot_labels[i]
            y1_sorted = series[1][perm]
            plt = plot(1:length(x_sorted), y1_sorted,
                label="solution 1",
                legend=true,
                left_margin=(10*upscale)*mm,
                right_margin=(10*upscale)*mm, top_margin=(10*upscale)*mm, bottom_margin=(10*upscale)*mm,
                linetype=:scatter,
                marker = (:dot, 1.5*upscale, 0.6, :blue, stroke(0))
            )
            y2_sorted = series[2][perm]
            plt = plot!(plt,1:length(x_sorted), y2_sorted,
                label="solution 2",
                legend=true,
                left_margin=(10*upscale)*mm,
                right_margin=(10*upscale)*mm, top_margin=(10*upscale)*mm, bottom_margin=(10*upscale)*mm,
                linetype=:scatter,
                marker = (:dot, 1.5*upscale, 0.6, :red, stroke(0))
            )
            xlabel!(plt, "index")
            ylabel!(plt, series_label)

            savefig(plt,"$(output_dir)/$(series_label).pdf")
        end
    end
    return table_string
end


const IUPACnumDict = Dict{Char,Int8}(Dict())
IUPACnumDict['A']=1
IUPACnumDict['C']=2
IUPACnumDict['G']=3
IUPACnumDict['T']=4
IUPACnumDict['U']=5
IUPACnumDict['R']=6
IUPACnumDict['Y']=7
IUPACnumDict['S']=8
IUPACnumDict['W']=9
IUPACnumDict['K']=10
IUPACnumDict['M']=11
IUPACnumDict['B']=12
IUPACnumDict['D']=13
IUPACnumDict['H']=14
IUPACnumDict['V']=15
IUPACnumDict['N']=16

IUPACnucLookup=[[true false false false false true false false true false true false true true true true]
[false true false false false false true true false false true true false true true true]
[false false true false false true false true false true false true true false true true]
[false false false true true false true false true true false true true true false true]
[false false false true true false true false true true false true true true false true]
[true false true false false true false true true true true true true true true true]
[false true false true true false true true true true true true true true true true]
[false true true false false true true true false true true true true true true true]
[true false false true true true true false true true true true true true true true]
[false false true true true true true true true true false true true true true true]
[true true false false false true true true true false true true true true true true]
[false true true true true true true true true true true true true true true true]
[true false true true true true true true true true true true true true true true]
[true true false true true true true true true true true true true true true true]
[true true true false false true true true true true true true true true true true]
[true true true true true true true true true true true true true true true true]];

function toIUPACnum(strin::String)
    return [IUPACnumDict[i] for i in strin]
end

function IUPAC_equals(a1,a2)
    return IUPACnucLookup[a1,a2]
end

"""Needleman-Wunch, with IUPAC codes, and end gaps penalized slightly less."""
function IUPAC_nuc_edit_dist(s1::String, s2::String; edge_reduction = 0.9999)
    # edge_reduction is a multiplicative score that gets multiplied to
    # the penalties along the edges, to prefer terminal gaps.
    ins_cost = -1.0  # -> : horizontal
    del_cost = -1.0  # V : vertical
    mismatch_cost = -1.0
    match_cost = 0.0

    s1numArr = toIUPACnum(s1)
    s2numArr = toIUPACnum(s2)
    s1arr = collect(s1)  # vertical
    s2arr = collect(s2)  # horizontal
    arr = zeros(length(s1arr)+1, length(s2arr)+1)
    traceArr = zeros(Int, length(s1arr), length(s2arr))

    # this will need to be generalized when we want to allow overhang
    arr[:, 1] = edge_reduction*del_cost*(0:length(s1arr))
    arr[1,:] = edge_reduction*ins_cost*(0:length(s2arr))

    for i in 2:length(s1arr)+1
        for j in 2:length(s2arr)+1
            if s1arr[i-1] == s2arr[j-1]
                diag = arr[i-1, j-1] + match_cost
            elseif IUPAC_equals(s1numArr[i-1],s2numArr[j-1])
                diag = arr[i-1, j-1] + match_cost-0.00001
            else
                diag = arr[i-1, j-1] + mismatch_cost
            end

            # to handle the lower edge penalties.
            delMult = 1
            if i == length(s1arr)+1
                delMult = edge_reduction
            end
            insMult=1
            if j == length(s2arr)+1
                insMult = edge_reduction
            end

            ins = arr[i-1, j]+(ins_cost*insMult)
            del = arr[i, j-1]+(del_cost*delMult)
            scores = [diag, del, ins]
            best = indmax(scores)
            if best == 1
                arr[i, j] = diag
            elseif best ==2
                arr[i, j] = del
            elseif best ==3
                arr[i, j] = ins
            end
            traceArr[i-1, j-1] = best
        end
    end
    alignedScore = arr[end, end]
    return alignedScore
end

"""Needleman-Wunch, with IUPAC codes, and end gaps penalized slightly less."""
function IUPAC_nuc_nw_align(s1::String, s2::String; edge_reduction = 0.9999)
    # edge_reduction is a multiplicative score that gets multiplied to
    # the penalties along the edges, to prefer terminal gaps.
    ins_cost = -1.0  # -> : horizontal
    del_cost = -1.0  # V : vertical
    mismatch_cost = -1.0
    match_cost = 0.0

    s1numArr = toIUPACnum(s1)
    s2numArr = toIUPACnum(s2)
    s1arr = collect(s1)  # vertical
    s2arr = collect(s2)  # horizontal
    arr = zeros(length(s1arr)+1, length(s2arr)+1)
    traceArr = zeros(Int, length(s1arr), length(s2arr))

    # this will need to be generalized when we want to allow overhang
    arr[:, 1] = edge_reduction*del_cost*(0:length(s1arr))
    arr[1,:] = edge_reduction*ins_cost*(0:length(s2arr))

    for i in 2:length(s1arr)+1
        for j in 2:length(s2arr)+1
            if s1arr[i-1] == s2arr[j-1]
                diag = arr[i-1, j-1] + match_cost
            elseif IUPAC_equals(s1numArr[i-1],s2numArr[j-1])
                diag = arr[i-1, j-1] + match_cost-0.00001
            else
                diag = arr[i-1, j-1] + mismatch_cost
            end

            # to handle the lower edge penalties.
            delMult = 1
            if i == length(s1arr)+1
                delMult = edge_reduction
            end
            insMult=1
            if j == length(s2arr)+1
                insMult = edge_reduction
            end

            ins = arr[i-1, j]+(ins_cost*insMult)
            del = arr[i, j-1]+(del_cost*delMult)
            scores = [diag, del, ins]
            best = indmax(scores)
            if best == 1
                arr[i, j] = diag
            elseif best ==2
                arr[i, j] = del
            elseif best ==3
                arr[i, j] = ins
            end
            traceArr[i-1, j-1] = best
        end
    end
    alignedScore = arr[end, end]

    # return arr

    # the trace endpoing will need to be generalized if we want to
    # allow overhang.
    trI, trJ = length(s1arr), length(s2arr)

    # First compute the trace running backwards. Initialized to the
    # maximum size. With unit penalties, is it possible to tell how
    # big this should be in advance?
    backtrace = Array{Int}(length(s1arr) + length(s2arr))
    # If you hit any boundary, you have to run all the way to the side.
    btInd = 1
    while (trI > 0) && (trJ > 0)
        backtrace[btInd] = traceArr[trI, trJ]
        if backtrace[btInd] == 1
            trI += -1
            trJ += -1
        elseif backtrace[btInd] == 2
            trJ += -1
        elseif backtrace[btInd] == 3
            trI += -1
        end
        btInd += 1
    end

    # If you hit the boundaries not at the top left corner.
    while trI > 0
        backtrace[btInd] = 3
        btInd += 1
        trI += -1
    end

    while trJ > 0
        backtrace[btInd] = 2
        btInd += 1
        trJ += -1
    end

    backtrace = backtrace[1:btInd-1]

    # This will need to be generalized to work on non-strings. Not
    # important.
    ali1arr = Array{Char}(length(backtrace))
    ali2arr = Array{Char}(length(backtrace))
    ind1 = 1
    ind2 = 1
    for i in 1:length(backtrace)
        tr = backtrace[length(backtrace)-(i-1)]
        if tr == 1
            ali1arr[i] = s1arr[ind1]
            ali2arr[i] = s2arr[ind2]
            ind1 += 1
            ind2 += 1
        elseif tr == 2
            ali2arr[i] = s2arr[ind2]
            ali1arr[i] = '-'
            ind2 += 1
        elseif tr == 3
            ali1arr[i] = s1arr[ind1]
            ali2arr[i] = '-'
            ind1 += 1
        end
    end
    return join(ali1arr), join(ali2arr), alignedScore
end

function banded_edit_dist(s1::String, s2::String; edge_reduction = 0.99, band_coeff = 1)
    # calculate band width
    avg_len = (length(s1) + length(s2)) / 2
    bandwidth = Int(ceil(band_coeff * sqrt(avg_len)))
    # edge_reduction is a multiplicative score that gets multiplied to
    # the penalties along the edges, to prefer terminal gaps.
    ins_cost = -1.0  # -> : horizontal
    del_cost = -1.0  # V : vertical
    mismatch_cost = -1.0
    match_cost = 1.0

    s1arr = collect(s1)  # vertical
    s2arr = collect(s2)  # horizontal

    # dimension difference added to band length on one side
    # to give padding for large sequence length differences
    dim_diff = length(s1arr) - length(s2arr)
    maxlen = max(length(s1), length(s2))
    arr = zeros(maxlen + 1, 2*bandwidth + 1 + abs(dim_diff))
    traceArr = zeros(Int, size(arr))
    for i in 1:bandwidth+1
        add_to_band!(arr, edge_reduction*del_cost*(i-1), i, 1, bandwidth, dim_diff)
        add_to_band!(arr, edge_reduction*ins_cost*(i-1), 1, i, bandwidth, dim_diff)
    end

    for i in 2:length(s1arr)+1
        # keep j within 2*bandwidth + 1 band of i, but keep bandwidth sized buffer on edges
        lo = 0
        # to keep a big enough padding around the edges when there's a big length mismatch
        if i <= bandwidth + 1 + max(0, dim_diff)
            lo = 2
        else
            lo = i - bandwidth - max(0, dim_diff)
        end
        for j in lo:length(s2arr)+1
            # keep j within band
            if !(in_band(i, j, bandwidth, dim_diff))
                break
            end

            if in_band(i-1, j-1, bandwidth, dim_diff) && s1arr[i-1] == s2arr[j-1]
                diag = get_band_val(arr, i-1, j-1, bandwidth, dim_diff) + match_cost
            else
                diag = get_band_val(arr, i-1, j-1, bandwidth, dim_diff) + mismatch_cost
            end

            # to handle the lower edge penalties.
            delMult = (i == length(s1arr)+1) ? edge_reduction : 1
            insMult = (j == length(s2arr)+1) ? edge_reduction : 1

            ins = in_band(i-1, j, bandwidth, dim_diff) ?
                    (get_band_val(arr, i-1, j, bandwidth, dim_diff)+(ins_cost*insMult)) : NaN
            del = in_band(i, j-1, bandwidth, dim_diff) ?
                    (get_band_val(arr, i, j-1, bandwidth, dim_diff)+(del_cost*delMult)) : NaN

            scores = [diag, del, ins]
            best = indmax(scores)
            add_to_band!(arr, scores[best], i, j, bandwidth, dim_diff)
            add_to_band!(traceArr, best, i-1, j-1, bandwidth, dim_diff)
        end
    end
    alignedScore = get_band_val(arr, length(s1arr)+1, length(s2arr)+1, bandwidth, dim_diff)

    return alignedScore
end

#Indmax in 2D, returning the coordinates
function maxcoord(mat)
    return ind2sub(mat,indmax(mat))
end

#THIS PRIMER MATCHING FUNCTION DOESN'T ASSUME WE KNOW PRIMER PAIRS
function matchPrimer(seq,forwardPrimers,revPrimers,forwardNames,revNames)
    fwdseq = seq;
    revseq = reverse_complement(fwdseq);
    fwdSeqFwdPrim = [IUPAC_nuc_edit_dist(fwdseq[1:length(prim)],prim) for prim in forwardPrimers];
    revSeqFwdPrim = [IUPAC_nuc_edit_dist(revseq[1:length(prim)],prim) for prim in forwardPrimers];
    fwdSeqRevPrim = [IUPAC_nuc_edit_dist(fwdseq[1:length(prim)],prim) for prim in revPrimers];
    revSeqRevPrim = [IUPAC_nuc_edit_dist(revseq[1:length(prim)],prim) for prim in revPrimers];

    fwdDirectionScores = zeros(length(forwardPrimers),length(revPrimers))
    revDirectionScores = zeros(length(forwardPrimers),length(revPrimers))

    for i in 1:length(forwardPrimers)
        for j in 1:length(revPrimers)
            fwdDirectionScores[i,j]=fwdSeqFwdPrim[i]+revSeqRevPrim[j];
            revDirectionScores[i,j]=revSeqFwdPrim[i]+fwdSeqRevPrim[j];
        end
    end

    fwd,back = maximum(fwdDirectionScores),maximum(revDirectionScores)

    #fwdDirectionScores = fwdSeqFwdPrim.+revSeqRevPrim;
    #revDirectionScores = revSeqFwdPrim.+fwdSeqRevPrim;
    fwd,back = maximum(fwdDirectionScores),maximum(revDirectionScores)
    fwdDir = fwd>back
    if fwdDir
        primInd = maxcoord(fwdDirectionScores)
        worstScore = minimum([fwdSeqFwdPrim[primInd[1]],revSeqRevPrim[primInd[2]]])
        retseq = fwdseq
    else
        primInd = maxcoord(revDirectionScores)
        worstScore = minimum([revSeqFwdPrim[primInd[1]],fwdSeqRevPrim[primInd[2]]])
        retseq = revseq
    end
    return primInd,forwardNames[primInd[1]]*"_"revNames[primInd[2]],retseq,round(worstScore),(!fwdDir)
end

function deMux(seqs,forwardPrimers,reversePrimers, forwardNames,reverseNames)
    matched = [matchPrimer(seq,forwardPrimers,reversePrimers,forwardNames,reverseNames) for seq in seqs];
    return [[i[2],i[3], i[5], i[4]] for i in matched]
end

function demux_fastx(fwdPrimers, revPrimers; minLen = 2300, errorRate = 0.01,
    path = "/Users/vrkumar/Downloads/",
    filename = "LP97.fastq",
    mismatchThresh = 2
)
    forwardNames = [i[1] for i in fwdPrimers];
    reverseNames = [i[1] for i in revPrimers];
    forwardPrimers = [i[2] for i in fwdPrimers];
    reversePrimers = [i[2] for i in revPrimers];

    newName = path*filename*"_"*string(minLen)*"_"*string(errorRate)*".fastq"
    usearch_filter(path*filename,newName, errorRate=errorRate,minLength=minLen,labelPrefix="seq",errorOut = true)

    if filename[end] == 'a'
        names, seqs = read_fasta_with_names(newName)
        phreds = nothing;
    else
        seqs, phreds, names = read_fastq(newName);
    end

    deMuxed = deMux(seqs,forwardPrimers,reversePrimers,forwardNames,reverseNames);

    primerPairs = [i[1] for i in deMuxed];
    primerMatchedSeqs = [i[2] for i in deMuxed];
    reversed = [i[3] for i in deMuxed];
    scores = [i[4] for i in deMuxed]

    if phreds != nothing
        out_phreds = [reversed[i] ? reverse(phreds[i]) : phreds[i] for i in 1:length(reversed)]
    else
        out_phreds = [0 for i in 1:length(reversed)]
    end

    out = Dict()
    for i in 1:length(primerPairs)
        if scores[i] > mismatchThresh
            continue
        end
        if !(haskey(out, primerPairs[i]))
            out[primerPairs[i]] = [[],[],[]]
        end
        push!(out[primerPairs[i]][1], primerMatchedSeqs[i])
        push!(out[primerPairs[i]][2], out_phreds[i])
        push!(out[primerPairs[i]][3], names[i])
    end

    for p in keys(out)
        if phreds != nothing
            write_fastq(path*p*".fastq", out[p][1], Array{Array{Int8,1},1}(out[p][2]), names = out[p][3])
        else
            write_fasta(path*p*".fastq", out[p][1], names = out[p][3])
        end
    end

end


using Gtk.GConstants.GtkAlign: GTK_ALIGN_START, GTK_ALIGN_END, GTK_ALIGN_FILL

function change_red(col::CLim{T}, chan::CLim{T2}) where {T<:AbstractRGB, T2<:GrayLike}
    cmin = T(chan.min, green(col.min), blue(col.min))
    cmax = T(chan.max, green(col.max), blue(col.max))
    return CLim(cmin, cmax)
end
function change_green(col::CLim{T}, chan::CLim{T2}) where {T<:AbstractRGB, T2<:GrayLike}
    cmin = T(red(col.min), chan.min, blue(col.min))
    cmax = T(red(col.max), chan.max, blue(col.max))
    return CLim(cmin, cmax)
end
function change_blue(col::CLim{T}, chan::CLim{T2}) where {T<:AbstractRGB, T2<:GrayLike}
    cmin = T(red(col.min), green(col.min), chan.min)
    cmax = T(red(col.max), green(col.max), chan.max)
    return CLim(cmin, cmax)
end
change_red(col::Signal, chan::CLim) = change_red(value(col), chan)
change_green(col::Signal, chan::CLim) = change_green(value(col), chan)
change_blue(col::Signal, chan::CLim) = change_blue(value(col), chan)

function contrast_gui(enabled::Signal{Bool}, hists::Vector, clim::Signal{CLim{T}}) where {T<:AbstractRGB}
    @assert length(hists) == 3 #one signal per color channel
    chanlims = channel_clims(value(clim))
    rsig = Signal(chanlims[1])
    gsig = Signal(chanlims[2])
    bsig = Signal(chanlims[3])
    #make sure that changes to individual channels update the color clim signal and vice versa
    bindmap!(clim, x->change_red(clim, x), rsig, x->channel_clim(red, x); initial = false)
    bindmap!(clim, x->change_green(clim, x), gsig, x->channel_clim(green, x); initial = false)
    bindmap!(clim, x->change_blue(clim, x), bsig, x->channel_clim(blue, x); initial = false)
    names = ["Red Contrast"; "Green Contrast"; "Blue Contrast"]
    csigs = [rsig; gsig; bsig]
    cguis = []
    for i=1:length(hists)
        push!(cguis, contrast_gui(enabled, hists[i], csigs[i]; wname = names[i]))
    end
    return cguis
end

contrast_gui(enabled, hist::Vector, clim) = contrast_gui(enabled, hist[1], clim)

function contrast_gui(enabled::Signal{Bool}, hist::Signal, clim::Signal; wname="Contrast")
    vhist, vclim = value(hist), value(clim)
    T = eltype(vclim)
    Δ = T <: Integer ? T(1) : eps(T)
    rng = vhist.edges[1]
    cmin, cmax = vclim.min, vclim.max
    if !(cmin < cmax)
        cmin, cmax = first(rng), last(rng)
        if !(cmin < cmax)
            cmin, cmax = zero(cmin), one(cmax)
        end
    end
    smin = Signal(convert(eltype(rng), cmin))
    smax = Signal(convert(eltype(rng), cmax))
    cgui = contrast_gui_layout(smin, smax, rng; wname=wname)
    signal_connect(cgui["window"], :destroy) do widget
        push!(enabled, false)
    end
    updateclim = map(smin, smax) do cmin, cmax
        # if min/max is outside the current range, update the sliders
        adj = Gtk.Adjustment(widget(cgui["slider_min"]))
        rmin, rmax = Gtk.G_.lower(adj), Gtk.G_.upper(adj)
        if cmin < rmin || cmax > rmax || cmax-cmin < Δ
            # Also, don't cross the sliders
            bigmax = max(cmin,cmax,rmin,rmax)
            bigmin = min(cmin,cmax,rmin,rmax)
            thismax = min(typemax(T), max(cmin, cmax, rmax))
            thismin = max(typemin(T), min(cmin, cmax, rmin))
            rng = range(thismin, stop=thismax, length=255)
            cminT, cmaxT = T(min(cmin, cmax)), T(max(cmin, cmax))
            if cminT == cmaxT
                cminT = min(cminT, cminT-Δ)
                cmaxT = max(cmaxT, cmaxT+Δ)
            end
            mn, mx = minimum(rng), maximum(rng)
            cmin, cmax = clamp(cminT, mn, mx), clamp(cmaxT, mn, mx)
            push!(cgui["slider_min"], rng, cmin)
            push!(cgui["slider_max"], rng, cmax)
        end
        # Update the image contrast
        push!(clim, CLim(cmin, cmax))
    end
    # TODO: we might want to throttle this?
    redraw = draw(cgui["canvas"], hist) do cnvs, hst
        if get_gtk_property(cgui["window"], :visible, Bool) # protects against window destruction
            rng, cl = hst.edges[1], value(clim)
            mn, mx = minimum(rng), maximum(rng)
            push!(cgui["slider_min"], rng, clamp(cl.min, mn, mx))
            push!(cgui["slider_max"], rng, clamp(cl.max, mn, mx))
            drawhist(cnvs, hst)
        end
    end
    GtkReactive.gc_preserve(cgui["window"], (cgui, redraw, updateclim))
    cgui
end

function contrast_gui_layout(smin::Signal, smax::Signal, rng; wname="Contrast")
    win = Window(wname) |> (g = Grid())
    slmax = slider(rng; signal=smax)
    slmin = slider(rng; signal=smin)
    for sl in (slmax, slmin)
        set_gtk_property!(sl, :draw_value, false)
    end
    g[1,1] = widget(slmax)
    g[1,3] = widget(slmin)
    cnvs = canvas(UserUnit)
    g[1,2] = widget(cnvs)
    set_gtk_property!(cnvs, :expand, true)
    emax_w = Entry(; width_chars=5, hexpand=false, halign=GTK_ALIGN_END, valign=GTK_ALIGN_START)
    emin_w = Entry(; width_chars=5, hexpand=false, halign=GTK_ALIGN_END, valign=GTK_ALIGN_END)
    g[2,1] = emax_w
    g[2,3] = emin_w
    # By not specifying the range on the textbox, we let the user
    # enter something out-of-range, which can be handy in some
    # circumstances.
    emax = textbox(eltype(smax); widget=emax_w, signal=smax) # , range=rng)
    emin = textbox(eltype(smin); widget=emin_w, signal=smin) #, range=rng)

    Gtk.showall(win)
    Dict("window"=>win, "canvas"=>cnvs, "slider_min"=>slmin, "slider_max"=>slmax, "textbox_min"=>emin, "textbox_max"=>emax)
end

# We could use one of the plotting toolkits, but most are pretty slow
# to load and/or produce the first plot. So let's just do it manually.
@guarded function drawhist(canvas, hist)
    ctx = getgc(canvas)
    fill!(canvas, colorant"white")
    edges, counts = hist.edges[1], hist.weights
    xmin, xmax = first(edges), last(edges)
    cmax = maximum(counts)
    if cmax <= 0 || !(xmin < xmax)
        return nothing
    end
    set_coordinates(ctx, BoundingBox(xmin, xmax, log10(cmax+1), 0))
    set_source(ctx, colorant"black")
    move_to(ctx, xmax, 0)
    line_to(ctx, xmin, 0)
    for (i, c) in enumerate(counts)
        line_to(ctx, edges[i], log10(c+1))
        line_to(ctx, edges[i+1], log10(c+1))
    end
    line_to(ctx, xmax, 0)
    fill(ctx)
    nothing
end


insts = OrderedDict()
insts["minimal"] = [
    ((1, 2, 2, 2),),
    ]
insts["fast"] = [
    ((2, 2, 3, 2),),
    ((3, 3, 3, 3),),
    ((3, 3, 5, 4),),
    ((5, 2, 5, 3),),
    ((1, 30, 2, 30),),
    ((10, 1, 3, 1),),
    ]
insts["various"] = [
    ((3, 3, 3, 3),),
    ((3, 3, 5, 4),),
    ((5, 2, 5, 3),),
    ((2, 30, 3, 30),),
    ((2, 15, 6, 15),),
    ((2, 30, 2, 40),),
    ((8, 3, 3, 3),),
    ]
return (PolyEnvelopeJuMP, insts)


abstract type GradientMode end
abstract type SemiGradient <: GradientMode end
abstract type FullGradient <: GradientMode end

export GradientMode, SemiGradient, FullGradient

include("accumulating.jl")
include("trueonline.jl")
include("tidbd.jl")

export AbstractAccumulatingTraceOptimizer, AccumulatingTraceOptimizer, Parl2AccumulatingTraceOptimizer
export TIDBDOptimizer, AutoTIDBDOptimizer
export AbstractTrueOnlineTraceOptimizer, TrueOnlineTraceOptimizer


module Milkbot

using  CSV, HTTP, StructTypes,  Plots, DataFrames, JSON3, TensorCast, Setfield, Printf
export MilkPoint, Lactation, FittedLactation, LactationSet, MBParams, Fit, Priors, FitOptions, FittingJob, Credentials  #, StructTypes.StructType
export getPriorsTree, fitLactation, fitLactationSet, mbParamsArray, milkbot, lactationPlot, paramsPlot, refitPlot, fitPlot

struct Credentials
    server::String
    apikey::String
end

struct MilkPoint
    dim::Integer
    milk::Float64
end
MilkPoint((dim,milk)) = MilkPoint(dim,milk)

struct Lactation
    lacKey::String
    breed::String
    parity::Integer
    points
end

struct MBParams
    scale::Float16
    ramp::Float16
    decay::Float16
    offset::Float16
end
MBParams(js ::JSON3.Object{Base.CodeUnits{UInt8, String}, SubArray{UInt64, 1, Vector{UInt64}, Tuple{UnitRange{Int64}}, true}}) = MBParams(js.scale,js.ramp,js.decay,js.offset)
MBParams(d::Dict) = MBParams(d["scale"],d["ramp"],d["decay"],d["offset"])

const mbParamsArray(p::MBParams) = [p.scale, p.ramp,p.decay,p.offset]
struct Fit
    lacKey::String
    params::MBParams
    n
    sdResiduals
    path
end
struct FittedLactation
    lac:: Lactation
    fit:: Fit
end
struct Priors
    means::MBParams
    sd::MBParams
    seMilk::Float16
    milkUnit::String
end
struct FitOptions
    steppedFit::Bool
    returnInputData::Bool
    returnPriors::Bool
    returnPath::Bool
    fitMethod::String  #: String = "MilkBot@2.01",
    preferredMilkUnit::String #: MilkUnit = Kg
end
const trackedOptions = FitOptions(false,false,false,true,"MilkBot@2.01","Kg")
struct LactationSet
    name #: Option[String],
    lactations #: Seq[Lactation],
    milkUnit #: Option[MilkUnit.Value]
end
struct FittingJob
    lactationSet::LactationSet
    priorsTree
    options
end
const TrackedFittingJob(l::Lactation) = FittingJob(LactationSet([l]),nothing,trackedOptions)

StructTypes.StructType(::Type{FittingJob}) = StructTypes.Struct()
StructTypes.StructType(::Type{LactationSet}) = StructTypes.Struct()
StructTypes.StructType(::Type{Lactation}) = StructTypes.Struct()
StructTypes.StructType(::Type{MilkPoint}) = StructTypes.Struct()
StructTypes.StructType(::Type{FitOptions}) = StructTypes.Struct()
StructTypes.StructType(::Type{Priors}) = StructTypes.Struct()
StructTypes.StructType(::Type{MBParams}) = StructTypes.Struct()
StructTypes.StructType(::Type{Fit}) = StructTypes.Struct()

const getPriorsTree(cred::Credentials) =
begin
    try
        response = HTTP.get("$(cred.server)/priorstree?milkUnit=Kg", Dict("X-API-KEY"=>cred.apikey))
        return String(response.body)
    catch e
        return "Error occurred : $e"
    end
end

const fitLactation(lac::Lactation, cred::Credentials) =
begin
    jsbody = JSON3.write("lactation" => lac)
    try
        response = HTTP.post("$(cred.server)/fitlactation", Dict("X-API-KEY"=>cred.apikey),jsbody)
        return String(response.body)
    catch e
        return "Error occurred : $e"
    end
end

const fitLactation(lac::Lactation, priors::Priors, cred::Credentials) =
begin
    jsbody = JSON3.write(Dict("lactation" => lac, "priors" => priors))
    try
        response = HTTP.post("$(cred.server)/fitlactation", Dict("X-API-KEY"=>cred.apikey),jsbody)
        return String(response.body)
    catch e
        return "Error occurred : $e"
    end
end

const fitLactation(lac::Lactation, priors::Priors, cred::Credentials, withPath::Bool) =
begin
    jsbody = JSON3.write(Dict("lactation" => lac, "priors" => priors))
    try
        response = HTTP.post("$(cred.server)/fitlactation?includePath=$(withPath)", Dict("X-API-KEY"=>cred.apikey),jsbody)
        return String(response.body)
    catch e
        return "Error occurred : $e"
    end
end

const fitLactationSet(fj::FittingJob, cred::Credentials) =
begin
    jsbody = JSON3.write("fittingJob" => fj)
    try
        response = HTTP.post("$(cred.server)/fitlactations", Dict("X-API-KEY"=>cred.apikey),jsbody)
        return String(response.body)
    catch e
        return "Error occurred : $e"
    end
end

#const milkbot(r::RowTypeSitko) = t ->  ((1 - ℯ^((r.offset - t)/r.ramp)/2)*r.scale)/ℯ^(r.decay * t)
const milkbot(f::Fit) = t ->  ((1 - ℯ^((f.params.offset - t)/f.params.ramp)/2)*f.params.scale)/ℯ^(f.params.decay * t)
const milkbot(p:: MBParams) = t ->  ((1 - ℯ^((p.offset - t)/p.ramp)/2)*p.scale)/ℯ^(p.decay * t)

const lactationPlot(fl::FittedLactation) =
 begin
    dim = map(m -> m.dim, fl.lac.points)
    milk = map(m -> m.milk, fl.lac.points)
    plot(dim  ,milk, seriestype = :scatter, color = :black, label = fl.lac.lacKey)
    plot!(milkbot(fl.fit),0,maximum(dim), label = "fit")
 end

 const paramsPlot(fit::Fit, priors::Priors) =
 begin
     data = [
        map(d->(d["scale"]-priors.means.scale)/priors.sd.scale,fit.path),
        map(d->(d["ramp"]-priors.means.ramp +.004)/priors.sd.ramp,fit.path),
        map(d->(d["decay"]-priors.means.decay)/priors.sd.decay,fit.path),
        map(d->(d["offset"]-priors.means.offset-.004)/priors.sd.offset-.01,fit.path)
     ]
      plot(1:length(fit.path),data , labels = ["scale" "ramp" "decay" "offset"], legend=:topleft)
 end

 fitPlot(l::Lactation, priors::Priors, cred:: Credentials) =
 begin
    dim = map(m -> m.dim, l.points)
    milk = map(m -> m.milk, l.points)
    fit = JSON3.read(fitLactation(l, priors, cred, true), Fit) #call to API to re-fit
    plot(dim,milk, seriestype = :scatter, label = l.lacKey)
    s = map(p ->milkbot(MBParams(p)) , collect(fit.path)) #vector of mb functions
    plot!(pop!(s),0,maximum(dim), color = :black, linewidth =2, label = "final",
    annotate = (((.4, .15), (@sprintf("SE: %.2f", fit.sdResiduals), 10,:right))))
    p1 = plot!(s[1:min(12, length(s))],0,maximum(dim), label = "")
    p2 =   paramsPlot(fit, priors)
    df = vcat(DataFrame.(fit.path)...)
    plot(p1,p2, layout = (2,1), size = (600,600))
    #vcat(DataFrame.(reFit.path)...)
 end

 #version including API call
 const refitPlot(fl::FittedLactation, priors::Priors, cred:: Credentials) =
 begin
    dim = map(m -> m.dim, fl.lac.points)
    milk = map(m -> m.milk, fl.lac.points)
    reFit = JSON3.read(fitLactation(fl.lac, priors, cred, true), Fit) #call to API to re-fit
    plot(dim, milk, seriestype = :scatter, label = fl.lac.lacKey)
    plot!(milkbot(fl.fit),0,maximum(dim), label = "original fit", linestyle = :dash, color = :black, linewidth =2,
    annotate = (((.4, .25), ("original SE: $(fl.fit.sdResiduals)", 10, :right))))
    s = map(p ->milkbot(MBParams(p)) , collect(reFit.path)) #vector of mb functions
    plot!(pop!(s),0,maximum(dim), color = :black, linewidth =2, label = "re-fit",
    annotate = (((.4, .15), (@sprintf("refit SE: %.2f", reFit.sdResiduals), 10,:right))))
    p1 = plot!(s[1:min(12, length(s))],0,maximum(dim), label = "")
    p2 =   paramsPlot(reFit, priors)
    df = vcat(DataFrame.(reFit.path)...)
    plot(p1,p2, layout = (2,1), size = (600,600))
    #vcat(DataFrame.(reFit.path)...)
 end

  #no API call
 const refitPlot(fl::FittedLactation, reFit:: Fit) =
 begin
    dim = map(m -> m.dim, fl.lac.points)
    milk = map(m -> m.milk, fl.lac.points)
    plot(dim  , milk, seriestype = :scatter, label = fl.lac.lacKey)
    plot!(milkbot(fl.fit),0,maximum(dim), label = "original fit", linestyle = :dash, color = :black, linewidth =2)
    s = map(p ->milkbot(MBParams(p)) , collect(reFit.path)) #vector of mb functions
    plot!(pop!(s),0,maximum(dim), color = :black, linewidth =2, label = "re-fit")
    p1 = plot!(s[1:min(12, length(s))],0,maximum(dim), label = "")
    p2 =   paramsPlot(reFit, priors)
    df = vcat(DataFrame.(reFit.path)...)
    plot(p1,p2, layout = (2,1), size = (600,600))
 end

end #Milkbot module


using Plots
using Plots.Measures
using DataFrames
using LinearAlgebra
using JSON

using ACEtb
using ACEtb.FitTools
using ACEtb.AnalysisTools
using LaTeXStrings

##

const DFTB_PATH = "$(@__DIR__)/../dftbplus/build/install/bin/"
const OUTPUT_DIR = joinpath(@__DIR__, "Results_fits/")
const PARAM_FILE = "Al-Al.acetb.json"

fc = FitContext(@__DIR__, DFTB_PATH, OUTPUT_DIR, PARAM_FILE)

results = tabulate_results(fc)
filter!(x -> !startswith(x.param_hash, "order"), results)

subresults = dropmissing(results)
subresults = subresults[(subresults.rcut .== 10.0) .& (subresults.N_train .<= 2) .& (subresults.rtol .== 1e-8), :]

gdf = groupby(subresults, :config)

degplots = [plot(; yscale=:log10, legend=:bottomleft, xlabel="Degree",
                ylabel="Training error")
            plot(; yscale=:log10, legend=:false, xlabel="Degree",
                ylabel="Band error FCC", ylim=(10^0.2, 10^0.4))
            plot(; yscale=:log10, legend=:false, xlabel="Degree",
                ylabel="Band error BCC")] #ylim=(10^0.2, 10^0.4))]
for df in gdf
    for order in (2, 3)
        nrow(df) <= 1 && continue
        degree = df[df.order .== order, :degree]
        train_err = df[df.order .== order, :error_train]
        band_err_fcc = df[df.order .== order, :band_error_FCC]
        band_err_bcc = df[df.order .== order, :band_error_BCC]
        config = df.config[df.order .== order][1]
        sortorder = sortperm(degree)
        plot!(degplots[1], degree[sortorder],
              train_err[sortorder], linewidth=2,
              label="Order $order $config")
        plot!(degplots[2], degree[sortorder],
              band_err_fcc[sortorder], linewidth=2,
              label="Order $order $config")
        plot!(degplots[3], degree[sortorder],
              band_err_bcc[sortorder], linewidth=2,
              label="Order $order $config")
    end
end

rcut_plot = [plot(; xlabel="Cutoff radius", ylabel="Training Error", legend=:bottomleft, yscale=:log10),
             plot(; xlabel="Cutoff radius", ylabel="FCC Band Error", legend=false, yscale=:log10),
             plot(; xlabel="Cutoff radius", ylabel="BCC Band Error", legend=false, yscale=:log10)]

order = 2
for (color, config_type) in enumerate(["highT", "lowT", "mixT"])
    rcut_mask =  ((results.config .== config_type) .&
                    (results.rcut .<= 11.0) .&
                    (results.order .== order) .&
                    (results.rcut .> 6.0) .&
                    (results.N_train .<= 2) .&
                    (results.pcut .== 1))
    if config_type == "highT"
        rcut_mask .&= results.degree .== 10
    else
        rcut_mask .&= results.degree .== 10
    end
    rcut_mask = findall(rcut_mask)
    rcut_mask = rcut_mask[sortperm(results[rcut_mask, :rcut])]
    @show order config_type rcut_mask, results[rcut_mask, [:rcut, :pot_file]]

    plot!(rcut_plot[1], results[rcut_mask, :rcut], marker=:dot,
            markersize=3, markerstrokecolor=color, color=color,
            results[rcut_mask, :error_train], linewidth=2,
            label="Order $order $config_type")

    plot!(rcut_plot[2], results[rcut_mask, :rcut], marker=:dot,
            markersize=3, markerstrokecolor=color, color=color,
            results[rcut_mask, :band_error_FCC], linewidth=2,
            label="Order $order $config_type")

    plot!(rcut_plot[3], results[rcut_mask, :rcut], marker=:dot,
            markersize=3, markerstrokecolor=color, color=color,
            results[rcut_mask, :band_error_BCC], linewidth=2,
            label="Order $order $config_type")

end

rp = plot(rcut_plot..., layout=(1, 3), size=(1200, 400))
savefig("rcut.pdf")

dp = plot(degplots..., layout=(1, 3), size=(1200, 400))
savefig("degree.pdf")

# bp = comp_plot(fc, "order2_deg10_envdeg4_rcut11_renv3.0_conf15_pcut1_Ntrain1",
#                label="mixT Order 2 Deg 10 rcut 11")
# savefig("bands.pdf")

# bp

##

N_train = 100 * 321

p1 = scatter(results.N_bond, results.error_train ./ sqrt.(results.N_bond), xscale=:log10, yscale=:log10, label="Train", legend=:bottomright)
xlabel!(L"N_\mathrm{bonds}")
ylabel!("RMSE Hamiltonian")

scatter!(results.N_bond, results.error_test ./ sqrt(N_train), xscale=:log10, yscale=:log10, label="Test")
# xlabel!(L"N_\mathrm{bonds}")
# ylabel!("RMSE Hamiltonian")

p3 = scatter(results.N_bond, results.band_error_FCC, xscale=:log10, yscale=:log10, label="FCC")
scatter!(results.N_bond, results.band_error_BCC, xscale=:log10, yscale=:log10, label="BCC")

xlabel!(L"N_\mathrm{bonds}")
ylabel!(L"\mathrm{Band\; error}")

p = plot(p1, p3, size=(600, 300), layout=(1, 2), leftmargin=5Plots.mm, bottommargin=5Plots.mm)

##

println("Loading test data...")
param_hash = results[argmin(results.N_train), :param_hash] # pick smallest model
pot_dict, _, Bint_table, cutf_func = load_model(fc, param_hash)
test_dataset = pot_dict["test_datasets"]
test_data = ACEtb.Utils.get_data(test_dataset, cutf_func, ACEtb.Bonds.get_env)
println("done")

##

test_errors = Dict()
train_errors = Dict()

for param_hash in sort(results, :N_train)[1:5, :param_hash]
    @show param_hash
    pot_dict, _, Bint_table, cutf_func = load_model(fc, param_hash)
    test_errors[param_hash] = hcat(test_error(test_data, Bint_table)...)
    println("Test error: 2-norm $(norm(test_errors[param_hash])), inf-norm $(norm(test_errors[param_hash],Inf)), rms $(mean(sqrt(test_errors[param_hash].^2)))")

    train_dataset = pot_dict["training_datasets"]
    train_data = ACEtb.Utils.get_data(test_dataset, cutf_func, ACEtb.Bonds.get_env)
    train_errors[param_hash] = hcat(test_error(train_data, Bint_table)...)
    println("Test error: 2-norm $(norm(train_errors[param_hash])), inf-norm $(norm(train_errors[param_hash],Inf)), rms $(mean(sqrt(train_errors[param_hash].^2)))")
end


# generateur d'Instances pour la RODD projet 4

P = 1
m = 15
n = 15
w1 = 1
w2 = 5
L = 3
g = 1.26157

# Ouvrir le fichier "output.txt" dans lequel on pourra écrire
fout = open("output30.txt", "w")

println(fout, "P = " * string(P))
println(fout, "m = " * string(m))
println(fout, "n = " * string(n))
println(fout, "w1 = " * string(w1))
println(fout, "w2 = " * string(w2))
println(fout, "L = " * string(L))
println(fout, "g = " * string(g))
println(fout, "title")

for i in 1:n
  myLine = string(i)
  for j in 1:(m)
    myLine = myLine * " " * string(rand(70:100))
  end
  println(fout, myLine)
end

close(fout)


using OpenQASM.Types
using OpenQASM.RBNF: Token

#todo support custom yao gates to qasm

"""
    convert_to_qasm(qc, ncreg)

Parses an `AbstractBlock` into code based on the `OpenQASM` spec

- `qc`: A `ChainBlock`(circuit that is to be run).
- `ncreg` (optional) : Number of classical registers
While performing operations like measuring, one can input desired number of classical regs(each size equal to number of qubits). Defaults to 1.
"""
function convert_to_qasm(qc::AbstractBlock{N}, ncreg::Int = 1) where {N}
    prog = generate_prog(qc, ncreg)
    MainProgram(v"2", prog)
end

"""
    generate_prog(qc)

Parses the YaoIR into a list of `QASM` instructions

- `qc`: A `ChainBlock`(circuit that is to be run).
"""
function generate_prog(qc::AbstractBlock{N}, ncreg::Int) where {N}
    prog = []
    generate_defs(prog, N, ncreg)
    generate_prog!(prog, basicstyle(qc), [0:N-1...], Int[])
    return prog
end

function generate_defs(prog, nq::Int, ncreg::Int)
    qregs = RegDecl(Token{:reserved}("qreg"), Token{:type}("q"), Token{:int}("$nq"))
    cregs = collect(
        RegDecl(Token{:reserved}("creg"), Token{:type}("c$i"), Token{:int}("$nq")) for
        i = 1:ncreg
    )
    push!(prog, Include(Bit("\"qelib1.inc\"")), qregs, cregs...)
end

function generate_prog!(prog, qc_simpl::ChainBlock, locs, controls)
    for block in subblocks(qc_simpl)
        generate_prog!(prog, block, locs, controls)
    end
end

function generate_prog!(prog, blk::PutBlock{N,M}, locs, controls) where {N,M}
    generate_prog!(prog, blk.content, sublocs(blk.locs, locs), controls)
end

function generate_prog!(prog, blk::ControlBlock{N,GT,C}, locs, controls) where {N,GT,C}
    any(==(0), blk.ctrl_config) && error("Inverse Control used in Control gate context")
    generate_prog!(
        prog,
        blk.content,
        sublocs(blk.locs, locs),
        [controls..., sublocs(blk.ctrl_locs, locs)...],
    )
end

function generate_prog!(prog, m::YaoBlocks.Measure{N}, locs, controls) where {N}
    mlocs = sublocs(m.locations isa AllLocs ? [1:N...] : [m.locations...], locs)
    (length(controls) == 0) || error("controlled measure is not supported")

    if !(m.operator isa ComputationalBasis) && m.operator isa AbstractBlock
        generate_prog!(prog, m.operator, mlocs, ())
    end

    # can be improved
    for i in mlocs
        cargs = Bit("c1", i)
        qargs = Bit("q", i)
        inst = Types.Measure(qargs, cargs)
        push!(prog, inst)
    end
end

# IBMQ Chip only supports ["id", "u1", "u2", "u3", "cx"]
# x, y, z and control x, y, z, id, t, swap and other primitive gates
for (GT, NAME, MAXC) in [
    (:XGate, "x", 2),
    (:YGate, "y", 2),
    (:ZGate, "z", 2),
    (:I2Gate, "id", 0),
    (:TGate, "t", 0),
    (:SWAPGate, "swap", 0),
]
    @eval function generate_prog!(prog, ::$GT, locs, controls)
        if length(controls) <= $MAXC
            # push!(prog, )
            name = "c"^(length(controls)) * $NAME
            if name != "cx" && name != "swap"
                cargs = [] #empty for now
                qargs =
                    isempty(controls) ? [Bit("q", locs...)] :
                    [Bit("q", controls...), Bit("q", locs...)]
                inst = Instruction(name, cargs, qargs)
            elseif name == "swap"
                cargs = [] #empty for now
                qargs = [Bit("q", i) for i in locs]
                inst = Instruction(name, cargs, qargs)
            else
                inst = CXGate(Bit("q", controls...), Bit("q", locs...))
            end
            push!(prog, inst)
        else
            error("too many control bits!")
        end
    end
end

for (GT, NAME, PARAMS, MAXC) in [
    (:(RotationGate{1,T,XGate} where {T}), "rx", :(b.theta), 0),
    (:(RotationGate{1,T,YGate} where {T}), "ry", :(b.theta), 0),
    (:(RotationGate{1,T,ZGate} where {T}), "rz", :(b.theta), 0),
    (:(ShiftGate), "p", :(b.theta), 1),
    (:(HGate), "h", :(nothing), 0),
]
    @eval function generate_prog!(prog, b::$GT, locs, controls)
        if length(controls) <= $MAXC
            name = "c"^(length(controls)) * $NAME
            if $PARAMS === nothing
                cargs = []
                qargs = [Bit("q", locs...)]
            else
                params = $PARAMS
                cargs = [Token{:float64}("$params")]
                qargs =
                    isempty(controls) ? [Bit("q", locs...)] :
                    [Bit("q", controls...), Bit("q", locs...)]
            end
            push!(prog, Instruction(name, cargs, qargs))
        else
            error("too many control bits! got $controls (length > $($(MAXC)))")
        end
    end
end

sublocs(subs, locs) = [locs[i] for i in subs]

function basicstyle(blk::AbstractBlock)
    YaoBlocks.Optimise.simplify(blk, rules = [YaoBlocks.Optimise.to_basictypes])
end


using DiffEqProblemLibrary, Sundials, Base.Test

prob = prob_ode_2Dlinear
integrator = init(prob,CVODE_BDF())
step!(integrator)
integrator(integrator.t,Val{1})

for i in integrator
    @show i.t
end

integrator.t == 1.0
step!(integrator)
integrator.t > 1.0