# A form is something that ends up as geometry in the graphic.

abstract type FormPrimitive end

const empty_tag = Symbol("")

struct Form{P <: FormPrimitive} <: ComposeNode
    primitives::Vector{P}
    tag::Symbol

    Form{P}(prim, tag::Symbol=empty_tag) where P = new{P}(prim, tag)
end

Form(primitives::Vector{P}, tag::Symbol=empty_tag) where P <: FormPrimitive =
        Form{P}(primitives, tag)

isempty(f::Form) = isempty(f.primitives)

isscalar(f::Form) =
        length(f.primitives) == 1

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, form::Form) =
        Form([resolve(box, units, t, primitive) for primitive in form.primitives])

Base.similar(f::Form{T}) where T = Form{T}(T[])

form_string(::Form) = "FORM"  # fallback definition

# Polygon
# -------

struct SimplePolygonPrimitive{P <: Vec} <: FormPrimitive
    points::Vector{P}
end

const SimplePolygon{P<:SimplePolygonPrimitive} = Form{P}

const Polygon = SimplePolygon
const PolygonPrimitive = SimplePolygonPrimitive

polygon() = Form([PolygonPrimitive(Vec[])])

"""
    polygon(points)

Define a polygon. `points` is an array of `(x,y)` tuples
that specify the corners of the polygon.
"""
function polygon(points::AbstractArray{T}, tag=empty_tag) where T <: XYTupleOrVec
    XM, YM = narrow_polygon_point_types(Vector[points])
    if XM == Any
        XM = Measure
    end
    if YM == Any
        YM = Measure
    end
    VecType = Tuple{XM, YM}

    return Form([PolygonPrimitive(VecType[(x_measure(point[1]), y_measure(point[2]))
                    for point in points])], tag)
end

function polygon(point_arrays::AbstractArray, tag=empty_tag)
    XM, YM = narrow_polygon_point_types(point_arrays)
    VecType = XM == YM == Any ? Vec : Tuple{XM, YM}
    PrimType = XM == YM == Any ? PolygonPrimitive : PolygonPrimitive{VecType}

    polyprims = Array{PrimType}(length(point_arrays))
    for (i, point_array) in enumerate(point_arrays)
        polyprims[i] = PrimType(VecType[(x_measure(point[1]), y_measure(point[2]))
                                        for point in point_array])
    end
    return Form{PrimType}(polyprims, tag)
end

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::PolygonPrimitive) =
        PolygonPrimitive{AbsoluteVec2}( AbsoluteVec2[
            resolve(box, units, t, point) for point in p.points])

function boundingbox(form::PolygonPrimitive, linewidth::Measure,
                     font::AbstractString, fontsize::Measure)
    x0 = minimum([p[1] for p in form.points])
    x1 = maximum([p[1] for p in form.points])
    y0 = minimum([p[2] for p in form.points])
    y1 = maximum([p[2] for p in form.points])
    return BoundingBox(x0 - linewidth,
                       y0 - linewidth,
                       x1 - x0 + linewidth,
                       y1 - y0 + linewidth)
end

form_string(::SimplePolygon) = "SP"

struct ComplexPolygonPrimitive{P <: Vec} <: FormPrimitive
    rings::Vector{Vector{P}}
end

const ComplexPolygon{P<:ComplexPolygonPrimitive} = Form{P}

complexpolygon() = ComplexPolygon([ComplexPolygonPrimitive(Vec[])])

function complexpolygon(rings::Vector{Vector}, tag=empty_tag)
    XM, YM = narrow_polygon_point_types(rings)
    if XM == Any
        XM = Length{:cx, Float64}
    end
    if YM == Any
        YM = Length{:cy, Float64}
    end
    VecType = Tuple{XM, YM}

    return ComplexPolygon([
        ComplexPolygonPrimitive([VecType[(x_measure(point[1]), y_measure(point[2]))
                                         for point in points]])], tag)
end

function complexpolygon(ring_arrays::Vector{Vector{Vector}}, tag=empty_tag)
    XM, YM = narrow_polygon_point_types(coords)
    VecType = XM == YM == Any ? Vec : Tuple{XM, YM}
    PrimType = XM == YM == Any ? ComplexPolygonPrimitive : ComplexPolygonPrimitive{VecType}

    ComplexPolygon([PrimType[[(x_measure(x), y_measure(y)) for (x, y) in ring]
                    for ring in ring_array] for ring_array in ring_arrays], tag)
end

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::ComplexPolygonPrimitive) =
        ComplexPolygonPrimitive( [AbsoluteVec2[
            resolve(box, units, t, point, t) for point in ring] for ring in p.rings])

form_string(::ComplexPolygon) = "CP"

# Rectangle
# ---------

struct RectanglePrimitive{P <: Vec, M1 <: Measure, M2 <: Measure} <: FormPrimitive
    corner::P
    width::M1
    height::M2
end

const Rectangle{P<:RectanglePrimitive} = Form{P}

"""
    rectangle()

Define a rectangle that fills the current context completely.
"""
function rectangle()
    prim = RectanglePrimitive((0.0w, 0.0h), 1.0w, 1.0h)
    return Rectangle{typeof(prim)}([prim])
end

"""
    rectangle(x0, y0, width, height)

Define a rectangle of size `width`x`height` with its top left corner at the point (`x`, `y`).
"""
function rectangle(x0, y0, width, height, tag=empty_tag)
    corner = (x_measure(x0), y_measure(y0))
    width = x_measure(width)
    height = y_measure(height)
    prim = RectanglePrimitive(corner, width, height)
    return Rectangle{typeof(prim)}([prim], tag)
end

"""
    rectangle(x0s, y0s, widths, heights)

Arguments can be passed in arrays in order to perform multiple drawing operations at once.
"""
rectangle(x0s::AbstractArray, y0s::AbstractArray,
                   widths::AbstractArray, heights::AbstractArray, tag=empty_tag) =
        @makeform (x0 in x0s, y0 in y0s, width in widths, height in heights),
            RectanglePrimitive{Vec2, Measure, Measure}((x_measure(x0), y_measure(y0)),
                                                    x_measure(width), y_measure(height)) tag

function resolve(box::AbsoluteBox, units::UnitBox, t::Transform,
                 p::RectanglePrimitive)
    corner = resolve(box, units, t, p.corner)
    width = resolve(box, units, t, p.width)
    height = resolve(box, units, t, p.height)

    if isxflipped(units) && hasunits(Length{:cx}, p.corner[1])
        # if coordinates are flipped we end up with the corner on the other end
        # of the rectangle, which is fix here
        x = corner[1] - width
    else
        x = corner[1]
    end

    if isyflipped(units) && hasunits(Length{:cy}, p.corner[2])
        y = corner[2] - height
    else
        y = corner[2]
    end

    return RectanglePrimitive{AbsoluteVec2, AbsoluteLength, AbsoluteLength}(
        (x, y), width, height)
end

boundingbox(form::RectanglePrimitive, linewidth::Measure,
                     font::AbstractString, fontsize::Measure) =
        BoundingBox(form.corner.x - linewidth,
                    form.corner.y - linewidth,
                    form.width + 2*linewidth,
                    form.height + 2*linewidth)

form_string(::Rectangle) = "R"

# Circle
# ------

struct CirclePrimitive{P <: Vec, M <: Measure} <: FormPrimitive
    center::P
    radius::M
end

CirclePrimitive(center::P, radius::M) where {P, M} = CirclePrimitive{P, M}(center, radius)
CirclePrimitive(x, y, r) = CirclePrimitive((x_measure(x), y_measure(y)), x_measure(r))

const Circle{P<:CirclePrimitive} = Form{P}

"""
    circle()

Define a circle in the center of the current context with a diameter equal to the width of the context.
"""
function circle()
    prim = CirclePrimitive((0.5w, 0.5h), 0.5w)
    return Circle{typeof(prim)}([prim])
end

"""
    circle(x, y, r)

Define a circle with its center at (`x`,`y`) and a radius of `r`.
"""
function circle(x, y, r, tag=empty_tag)
    prim = CirclePrimitive(x, y, r)
    return Circle{typeof(prim)}([prim], tag)
end

"""
    circle(xs, ys, rs)

Arguments can be passed in arrays in order to perform multiple drawing operations.
"""
function circle(xs::AbstractArray, ys::AbstractArray, rs::AbstractArray, tag=empty_tag)
    if isempty(xs) || isempty(ys) || isempty(rs)
        prima = CirclePrimitive[]
        return Circle{eltype(prima)}(prima, tag)
    end

    return @makeform (x in xs, y in ys, r in rs),
        CirclePrimitive{Vec2, Measure}((x_measure(x), y_measure(y)), x_measure(r)) tag
end

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::CirclePrimitive) =
        CirclePrimitive{AbsoluteVec2, AbsoluteLength}(
            resolve(box, units, t, p.center),
            resolve(box, units, t, p.radius))

boundingbox(form::CirclePrimitive, linewidth::Measure, font::AbstractString, fontsize::Measure) =
        BoundingBox(form.center[1] - form.radius - linewidth,
            form.center[2] - form.radius - linewidth,
            2 * (form.radius + linewidth),
            2 * (form.radius + linewidth))

form_string(::Circle) = "C"

# Ellipse
# -------

struct EllipsePrimitive{P1<:Vec, P2<:Vec, P3<:Vec} <: FormPrimitive
    center::P1
    x_point::P2
    y_point::P3
end

const Ellipse{P<:EllipsePrimitive} = Form{P}

function ellipse()
    prim = EllipsePrimitive((0.5w, 0.5h),
                            (1.0w, 0.5h),
                            (0.5w, 1.0h))
    return Ellipse{typeof(prim)}([prim])
end

function ellipse(x, y, x_radius, y_radius, tag=empty_tag)
    xm = x_measure(x)
    ym = y_measure(y)
    prim = EllipsePrimitive((xm, ym),
                            (xm + x_measure(x_radius), ym),
                            (xm, ym + y_measure(y_radius)))
    return Ellipse{typeof(prim)}([prim], tag)
end

function ellipse(xs::AbstractArray, ys::AbstractArray,
                 x_radiuses::AbstractArray, y_radiuses::AbstractArray, tag=empty_tag)
        return @makeform (x in xs, y in ys, x_radius in x_radiuses, y_radius in y_radiuses),
            EllipsePrimitive((x_measure(x), y_measure(y)),
                     (x_measure(x) + x_measure(x_radius), y_measure(y)),
                     (x_measure(x), y_measure(y) + y_measure(y_radius))) tag
end

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::EllipsePrimitive) =
        EllipsePrimitive{AbsoluteVec2, AbsoluteVec2, AbsoluteVec2}(
            resolve(box, units, t, p.center),
            resolve(box, units, t, p.x_point),
            resolve(box, units, t, p.y_point))

function boundingbox(form::EllipsePrimitive, linewidth::Measure,
                     font::AbstractString, fontsize::Measure)
    x0 = min(form.x_point.x, form.y_point.x)
    x1 = max(form.x_point.x, form.y_point.x)
    y0 = min(form.x_point.y, form.y_point.y)
    y1 = max(form.x_point.y, form.y_point.y)
    xr = x1 - x0
    yr = y1 - y0
    return BoundingBox(x0 - linewidth - xr,
                       y0 - linewidth - yr,
                       2 * (xr + linewidth),
                       2 * (yr + linewidth))
end

form_string(::Ellipse) = "E"

# Text
# ----

abstract type HAlignment end
struct HLeft   <: HAlignment end
struct HCenter <: HAlignment end
struct HRight  <: HAlignment end

const hleft   = HLeft()
const hcenter = HCenter()
const hright  = HRight()

abstract type VAlignment end
struct VTop    <: VAlignment end
struct VCenter <: VAlignment end
struct VBottom <: VAlignment end

const vtop    = VTop()
const vcenter = VCenter()
const vbottom = VBottom()

struct TextPrimitive{P<:Vec, R<:Rotation, O<:Vec} <: FormPrimitive
    position::P
    value::AbstractString
    halign::HAlignment
    valign::VAlignment

    # Text forms need their own rotation field unfortunately, since there is no
    # way to give orientation with just a position point.
    rot::R

    offset::O
end

const Text{P<:TextPrimitive} = Form{P}

"""
    text(x, y, value [,halign::HAlignment [,valign::VAlignment [,rot::Rotation]]])

Draw the text `value` at the position (`x`,`y`) relative to the current context.

The default alignment of the text is `hleft` `vbottom`. The vertical and horizontal
alignment is specified by passing `hleft`, `hcenter` or `hright` and `vtop`,
`vcenter` or `vbottom` as values for `halgin` and `valgin` respectively.
"""
function text(x, y, value::AbstractString,
              halign::HAlignment=hleft, valign::VAlignment=vbottom,
              rot=Rotation(), offset::Vec2=(0mm,0mm);
              tag::Symbol=empty_tag)
    moffset = (x_measure(offset[1]), y_measure(offset[2]))
    prim = TextPrimitive((x_measure(x), y_measure(y)), value, halign, valign, rot, moffset)
    Text{typeof(prim)}([prim], tag)
end

function text(x, y, value,
              halign::HAlignment=hleft, valign::VAlignment=vbottom,
              rot=Rotation(), offset::Vec2=(0mm,0mm);
              tag::Symbol=empty_tag)
    moffset = (x_measure(offset[1]), y_measure(offset[2]))
    prim = TextPrimitive((x_measure(x), y_measure(y)), string(value), halign, valign, rot, moffset)
    Text{typeof(prim)}([prim], tag)
end

"""
    text(xs, ys, values [,haligns::HAlignment [,valigns::VAlignment [,rots::Rotation]]])

Arguments can be passed in arrays in order to perform multiple drawing operations at once.
"""
text(xs::AbstractArray, ys::AbstractArray, values::AbstractArray{AbstractString},
              haligns::AbstractArray=[hleft], valigns::AbstractArray=[vbottom],
              rots::AbstractArray=[Rotation()], offsets::AbstractArray=[(0mm,0mm)];
              tag::Symbol=empty_tag) =
    @makeform (x in xs, y in ys, value in values, halign in haligns, valign in valigns, rot in rots, offset in offsets),
        TextPrimitive((x_measure(x), y_measure(y)), value, halign, valign, rot, (x_measure(offset[1]), y_measure(offset[2]))) tag

text(xs::AbstractArray, ys::AbstractArray, values::AbstractArray,
              haligns::AbstractArray=[hleft], valigns::AbstractArray=[vbottom],
              rots::AbstractArray=[Rotation()], offsets::AbstractArray=[(0mm,0mm)];
              tag::Symbol=empty_tag) =
    @makeform (x in xs, y in ys, value in values, halign in haligns, valign in valigns, rot in rots, offset in offsets),
        TextPrimitive((x_measure(x), y_measure(y)), value, halign, valign, rot, (x_measure(offset[1]), y_measure(offset[2]))) tag

function resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::TextPrimitive{P,R,O}) where {P,R,O}
    rot = resolve(box, units, t, p.rot)
    return TextPrimitive{AbsoluteVec2, typeof(rot), O}(
                resolve(box, units, t, p.position),
                p.value, p.halign, p.valign, rot, p.offset)
end

function boundingbox(form::TextPrimitive, linewidth::Measure,
                     font::AbstractString, fontsize::Measure)

    width, height = text_extents(font, fontsize, form.value)[1]

    if form.halign == hleft
        x0 = form.position.x
    elseif form.halign == hcenter
        x0 = form.position.x - width/2
    elseif form.halign == hright
        x0 = form.position.x - width
    end

    if form.valign == vbottom
        y0 = form.position.y - height
    elseif form.valign == vcenter
        y0 = form.position.y - height/2
    elseif form.valign == vtop
        y0 = form.position.y
    end

    return BoundingBox(x0 - linewidth + form.offset[1],
                       y0 - linewidth + form.offset[2],
                       width + linewidth,
                       height + linewidth)
end

form_string(::Text) = "T"

# Line
# ----

struct LinePrimitive{P<:Vec} <: FormPrimitive
    points::Vector{P}
end

const Line{P<:LinePrimitive} = Form{P}

function line()
    prim = LinePrimitive(Vec[])
    return Line{typeof(prim)}([prim])
end

function line(points::AbstractArray{T}, tag=empty_tag) where T <: XYTupleOrVec
    XM, YM = narrow_polygon_point_types(Vector[points])
    VecType = XM == YM == Any ? Vec2 : Tuple{XM, YM}
    prim = LinePrimitive(VecType[(x_measure(point[1]), y_measure(point[2])) for point in points])
    return Line{typeof(prim)}([prim], tag)
end

function line(point_arrays::AbstractArray, tag=empty_tag)
    XM, YM = narrow_polygon_point_types(point_arrays)
    VecType = XM == YM == Any ? Vec2 : Tuple{XM, YM}
    PrimType = XM == YM == Any ? LinePrimitive : LinePrimitive{VecType}

    lineprims = Array{PrimType}(length(point_arrays))
    for (i, point_array) in enumerate(point_arrays)
        p = PrimType(VecType[(x_measure(point[1]), y_measure(point[2]))
                             for point in point_array])
        lineprims[i] = p
    end
    return Form{PrimType}(lineprims, tag)
end

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::LinePrimitive) =
        LinePrimitive{AbsoluteVec2}(
            AbsoluteVec2[resolve(box, units, t, point) for point in p.points])

function boundingbox(form::LinePrimitive, linewidth::Measure,
                     font::AbstractString, fontsize::Measure)
    x0 = minimum([p.x for p in form.points])
    x1 = maximum([p.x for p in form.points])
    y0 = minimum([p.y for p in form.points])
    y1 = maximum([p.y for p in form.points])
    return BoundingBox(x0 - linewidth,
                       y0 - linewidth,
                       x1 - x0 + linewidth,
                       y1 - y0 + linewidth)
end

form_string(::Line) = "L"

# Curve
# -----

struct CurvePrimitive{P1<:Vec, P2<:Vec, P3<:Vec, P4<:Vec} <: FormPrimitive
    anchor0::P1
    ctrl0::P2
    ctrl1::P3
    anchor1::P4
end

const Curve{P<:CurvePrimitive} = Form{P}

function curve(anchor0::XYTupleOrVec, ctrl0::XYTupleOrVec,
               ctrl1::XYTupleOrVec, anchor1::XYTupleOrVec, tag=empty_tag)
    prim = CurvePrimitive((x_measure(anchor0[1]), y_measure(anchor0[2])),
                                 (x_measure(ctrl0[1]), y_measure(ctrl0[2])),
                                 (x_measure(ctrl1[1]), y_measure(ctrl1[2])),
                                 (x_measure(anchor1[1]), y_measure(anchor1[2])))
    return Curve{typeof(prim)}([prim], tag)
end

curve(anchor0s::AbstractArray, ctrl0s::AbstractArray,
               ctrl1s::AbstractArray, anchor1s::AbstractArray, tag=empty_tag) =
        @makeform (anchor0 in anchor0s, ctrl0 in ctrl0s, ctrl1 in ctrl1s, anchor1 in anchor1s),
            CurvePrimitive((x_measure(anchor0[1]), y_measure(anchor0[2])),
                (x_measure(ctrl0[1]), y_measure(ctrl0[2])),
                (x_measure(ctrl1[1]), y_measure(ctrl1[2])),
                (x_measure(anchor1[1]), y_measure(anchor1[2]))) tag

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::CurvePrimitive) =
        CurvePrimitive{AbsoluteVec2, AbsoluteVec2, AbsoluteVec2, AbsoluteVec2}(
            resolve(box, units, t, p.anchor0),
            resolve(box, units, t, p.ctrl0),
            resolve(box, units, t, p.ctrl1),
            resolve(box, units, t, p.anchor1))

form_string(::Curve) = "CV"

# Bitmap
# ------

struct BitmapPrimitive{P <: Vec, XM <: Measure, YM <: Measure} <: FormPrimitive
    mime::AbstractString
    data::Vector{UInt8}
    corner::P
    width::XM
    height::YM
end

const Bitmap{P<:BitmapPrimitive} = Form{P}

function bitmap(mime::AbstractString, data::Vector{UInt8}, x0, y0, width, height, tag=empty_tag)
    corner = (x_measure(x0), y_measure(y0))
    width = x_measure(width)
    height = y_measure(height)
    prim = BitmapPrimitive(mime, data, corner, width, height)
    return Bitmap{typeof(prim)}([prim], tag)
end

bitmap(mimes::AbstractArray, datas::AbstractArray,
                x0s::AbstractArray, y0s::AbstractArray,
                widths::AbstractArray, heights::AbstractArray, tag=empty_tag) =
        @makeform (mime in mimes, data in datas, x0 in x0s, y0 in y0s, width in widths, height in heigths),
            BitmapPrimitive(mime, data, x0, y0, x_measure(width), y_measure(height)) tag

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::BitmapPrimitive) =
        BitmapPrimitive{AbsoluteVec2, AbsoluteLength, AbsoluteLength}(
            p.mime, p.data,
            resolve(box, units, t, p.corner),
            resolve(box, units, t, p.width),
            resolve(box, units, t, p.height))

boundingbox(form::BitmapPrimitive, linewidth::Measure, font::AbstractString, fontsize::Measure) =
        BoundingBox(form.corner.x, form.corner.y, form.width, form.height)

form_string(::Bitmap) = "B"

# Path
# ----

# An implementation of the SVG path mini-language.

abstract type PathOp end

struct MoveAbsPathOp <: PathOp
    to::Vec
end

function assert_pathop_tokens_len(op_type, tokens, i, needed)
    provided = length(tokens) - i + 1
    provided < needed &&
            error("In path $(op_type) requires $(needed) argumens but only $(provided) provided.")
end

function parsepathop(::Type{MoveAbsPathOp}, tokens::AbstractArray, i)
    assert_pathop_tokens_len(MoveAbsPathOp, tokens, i, 2)
    op = MoveAbsPathOp((x_measure(tokens[i]), y_measure(tokens[i + 1])))
    return (op, i + 2)
end

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::MoveAbsPathOp) =
        MoveAbsPathOp(resolve(box, units, t, p.to))

struct MoveRelPathOp <: PathOp
    to::Vec
end

function parsepathop(::Type{MoveRelPathOp}, tokens::AbstractArray, i)
    assert_pathop_tokens_len(MoveRelPathOp, tokens, i, 2)
    op = MoveRelPathOp((tokens[i], tokens[i + 1]))
    return (op, i + 2)
end

function resolve_offset(box::AbsoluteBox, units::UnitBox, t::Transform, p::Vec)
    absp = resolve(box, units, t, p)
    zer0 = resolve(box, units, t, (0w, 0h))
    return (absp[1] - zer0[1], absp[2] - zer0[2])
end

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::MoveRelPathOp) =
        MoveRelPathOp(resolve_offset(box, units, t, p.to))

struct ClosePathOp <: PathOp
end

parsepathop(::Type{ClosePathOp}, tokens::AbstractArray, i) = (ClosePathOp(), i)

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::ClosePathOp) = p

struct LineAbsPathOp <: PathOp
    to::Vec
end

function parsepathop(::Type{LineAbsPathOp}, tokens::AbstractArray, i)
    assert_pathop_tokens_len(LineAbsPathOp, tokens, i, 2)
    op = LineAbsPathOp((x_measure(tokens[i]), y_measure(tokens[i + 1])))
    return (op, i + 2)
end

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::LineAbsPathOp) =
        LineAbsPathOp(resolve(box, units, t, p.to))

struct LineRelPathOp <: PathOp
    to::Vec
end

function parsepathop(::Type{LineRelPathOp}, tokens::AbstractArray, i)
    assert_pathop_tokens_len(LineRelPathOp, tokens, i, 2)
    op = LineRelPathOp((x_measure(tokens[i]), y_measure(tokens[i + 1])))
    return (op, i + 2)
end

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::LineRelPathOp) =
        LineRelPathOp(resolve(box, units, t, p.to))

struct HorLineAbsPathOp <: PathOp
    x::Measure
end

function parsepathop(::Type{HorLineAbsPathOp}, tokens::AbstractArray, i)
    assert_pathop_tokens_len(HorLineAbsPathOp, tokens, i, 1)
    op = HorLineAbsPathOp(x_measure(tokens[i]))
    return (op, i + 1)
end

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::HorLineAbsPathOp) =
        HorLineAbsPathOp(resolve(box, units, t, (p.x, 0mm))[1])

struct HorLineRelPathOp <: PathOp
    Δx::Measure
end

function parsepathop(::Type{HorLineRelPathOp}, tokens::AbstractArray, i)
    assert_pathop_tokens_len(HorLineRelPathOp, tokens, i, 1)
    op = HorLineRelPathOp(x_measure(tokens[i]))
    return (op, i + 1)
end

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::HorLineRelPathOp) =
        HorLineRelPathOp(resolve(box, units, t, p.Δx))

struct VertLineAbsPathOp <: PathOp
    y::Measure
end

function parsepathop(::Type{VertLineAbsPathOp}, tokens::AbstractArray, i)
    assert_pathop_tokens_len(VertLineAbsPathOp, tokens, i, 1)
    op = VertLineAbsPathOp(y_measure(tokens[i]))
    return (op, i + 1)
end

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::VertLineAbsPathOp) =
        VertLineAbsPathOp(resolve(box, units, t, (0mm, p.y))[2])

struct VertLineRelPathOp <: PathOp
    Δy::Measure
end

function parsepathop(::Type{VertLineRelPathOp}, tokens::AbstractArray, i)
    assert_pathop_tokens_len(VertLineRelPathOp, tokens, i, 1)
    op = VertLineRelPathOp(y_measure(tokens[i]))
    return (op, i + 1)
end

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::VertLineRelPathOp) =
        VertLineAbsPathOp(resolve(box, units, t, (0mmm, p.Δy))[2])

struct CubicCurveAbsPathOp <: PathOp
    ctrl1::Vec
    ctrl2::Vec
    to::Vec
end

function parsepathop(::Type{CubicCurveAbsPathOp}, tokens::AbstractArray, i)
    assert_pathop_tokens_len(CubicCurveAbsPathOp, tokens, i, 6)
    op = CubicCurveAbsPathOp((tokens[i],     tokens[i + 1]),
                             (tokens[i + 2], tokens[i + 3]),
                             (tokens[i + 4], tokens[i + 5]))
    return (op, i + 6)
end

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::CubicCurveAbsPathOp) =
        CubicCurveAbsPathOp(
            resolve(box, units, t, p.ctrl1),
            resolve(box, units, t, p.ctrl2),
            resolve(box, units, t, p.to))

struct CubicCurveRelPathOp <: PathOp
    ctrl1::Vec
    ctrl2::Vec
    to::Vec
end

function parsepathop(::Type{CubicCurveRelPathOp}, tokens::AbstractArray, i)
    assert_pathop_tokens_len(CubicCurveRelPathOp, tokens, i, 6)
    op = CubicCurveRelPathOp((tokens[i],     tokens[i + 1]),
                             (tokens[i + 2], tokens[i + 3]),
                             (tokens[i + 4], tokens[i + 5]))
    return (op, i + 6)
end

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::CubicCurveRelPathOp) =
        CubicCurveRelPathOp(
            resolve(box, units, t, p.ctrl1),
            resolve(box, units, t, p.ctrl2),
            resolve(box, units, t, p.to))

struct CubicCurveShortAbsPathOp <: PathOp
    ctrl2::Vec
    to::Vec
end

function parsepathop(::Type{CubicCurveShortAbsPathOp}, tokens::AbstractArray, i)
    assert_pathop_tokens_len(CubicCurveShortAbsPathOp, tokens, i, 4)
    op = CubicCurveShortAbsPathOp((x_measure(tokens[i]),     y_measure(tokens[i + 1])),
                                  (x_measure(tokens[i + 2]), y_measure(tokens[i + 3])))
    return (op, i + 4)
end

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::CubicCurveShortAbsPathOp) =
        CubicCurveShortAbsPathOp(
            resolve_offset(box, units, t, p.ctrl2),
            resolve_offset(box, units, t, p.to))

struct CubicCurveShortRelPathOp <: PathOp
    ctrl2::Vec
    to::Vec
end

function parsepathop(::Type{CubicCurveShortRelPathOp}, tokens::AbstractArray, i)
    assert_pathop_tokens_len(CubicCurveShortRelPathOp, tokens, i, 4)
    op = CubicCurveShortRelPathOp((x_measure(tokens[i]),     y_measure(tokens[i + 1])),
                                  (x_measure(tokens[i + 2]), y_measure(tokens[i + 3])))
    return (op, i + 4)
end

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::CubicCurveShortRelPathOp) =
        CubicCurveShortRelPathOp(
            resolve(box, units, t, p.ctrl2),
            resolve(box, units, t, p.to))

struct QuadCurveAbsPathOp <: PathOp
    ctrl1::Vec
    to::Vec
end

function parsepathop(::Type{QuadCurveAbsPathOp}, tokens::AbstractArray, i)
    assert_pathop_tokens_len(QuadCurveAbsPathOp, tokens, i, 4)
    op = QuadCurveAbsPathOp((tokens[i],     tokens[i + 1]),
                            (tokens[i + 2], tokens[i + 3]))
    return (op, i + 4)
end

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::QuadCurveAbsPathOp) =
        QuadCurveAbsPathOp(
            resolve(box, units, t, p.ctrl1),
            resolve(box, units, t, p.to))

struct QuadCurveRelPathOp <: PathOp
    ctrl1::Vec
    to::Vec
end

function parsepathop(::Type{QuadCurveRelPathOp}, tokens::AbstractArray, i)
    assert_pathop_tokens_len(QuadCurveRelPathOp, tokens, i, 4)
    op = QuadCurveRelPathOp((tokens[i],     tokens[i + 1]),
                            (tokens[i + 2], tokens[i + 3]))
    return (op, i + 4)
end

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::QuadCurveRelPathOp) =
        QuadCurveRelPathOp(
            (resolve(box, units, t, p.ctrl1[1]),
             resolve(box, units, t, p.ctrl1[2])),
            (resolve(box, units, t, p.to[1]),
             resolve(box, units, t, p.to[2])))

struct QuadCurveShortAbsPathOp <: PathOp
    to::Vec
end

function parsepathop(::Type{QuadCurveShortAbsPathOp}, tokens::AbstractArray, i)
    assert_pathop_tokens_len(QuadCurveShortAbsPathOp, tokens, i, 2)
    op = QuadCurveShortAbsPathOp((tokens[i], tokens[i + 1]))
    return (op, i + 2)
end

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::QuadCurveShortAbsPathOp) =
        QuadCurveShortAbsPathOp(resolve(box, units, t, p.to))

struct QuadCurveShortRelPathOp <: PathOp
    to::Vec
end

function parsepathop(::Type{QuadCurveShortRelPathOp}, tokens::AbstractArray, i)
    assert_pathop_tokens_len(QuadCurveShortRelPathOp, tokens, i, 2)
    op = QuadCurveShortRelPathOp((tokens[i], tokens[i + 1]))
    return (op, i + 2)
end

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::QuadCurveShortRelPathOp) =
        QuadCurveShortRelPathOp(
            (resolve(box, units, t, p.to[1]),
             resolve(box, units, t, p.to[2])))

struct ArcAbsPathOp <: PathOp
    rx::Measure
    ry::Measure
    rotation::Float64
    largearc::Bool
    sweep::Bool
    to::Vec
end

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::ArcAbsPathOp) =
        ArcAbsPathOp(
            resolve(box, units, t, p.rx),
            resolve(box, units, t, p.ry),
            p.rotation,
            p.largearc,
            p.sweep,
            resolve(box, units, t, p.to))

struct ArcRelPathOp <: PathOp
    rx::Measure
    ry::Measure
    rotation::Float64
    largearc::Bool
    sweep::Bool
    to::Vec
end

function parsepathop(::Type{T}, tokens::AbstractArray, i) where T <: Union{ArcAbsPathOp, ArcRelPathOp}
    assert_pathop_tokens_len(T, tokens, i, 7)

    if isa(tokens[i + 3], Bool)
        largearc = tokens[i + 3]
    elseif tokens[i + 3] == 0
        largearc = false
    elseif tokens[i + 3] == 1
        largearc = true
    else
        error("largearc argument to the arc path operation must be boolean")
    end

    if isa(tokens[i + 4], Bool)
        sweep = tokens[i + 4]
    elseif tokens[i + 4] == 0
        sweep = false
    elseif tokens[i + 4] == 1
        sweep = true
    else
        error("sweep argument to the arc path operation must be boolean")
    end

    isa(tokens[i + 2], Number) || error("path arc operation requires a numerical rotation")

    op = T(x_measure(tokens[i]),
           y_measure(tokens[i + 1]),
           convert(Float64, tokens[i + 2]),
           largearc, sweep,
           (x_measure(tokens[i + 5]), y_measure(tokens[i + 6])))

    return (op, i + 7)
end

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::ArcRelPathOp) =
        ArcRelPathOp(
            resolve(box, units, t, p.rx),
            resolve(box, units, t, p.ry),
            p.rotation,
            p.largearc,
            p.sweep,
            (resolve(box, units, t, p.to[1]),
             resolve(box, units, t, p.to[2])))

const path_ops = Dict(
     :M => MoveAbsPathOp,
     :m => MoveRelPathOp,
     :Z => ClosePathOp,
     :z => ClosePathOp,
     :L => LineAbsPathOp,
     :l => LineRelPathOp,
     :H => HorLineAbsPathOp,
     :h => HorLineRelPathOp,
     :V => VertLineAbsPathOp,
     :v => VertLineRelPathOp,
     :C => CubicCurveAbsPathOp,
     :c => CubicCurveRelPathOp,
     :S => CubicCurveShortAbsPathOp,
     :s => CubicCurveShortRelPathOp,
     :Q => QuadCurveAbsPathOp,
     :q => QuadCurveRelPathOp,
     :T => QuadCurveShortAbsPathOp,
     :t => QuadCurveShortRelPathOp,
     :A => ArcAbsPathOp,
     :a => ArcRelPathOp
)

# A path is an array of symbols, numbers, and measures following SVGs path
# mini-language.
function parsepath(tokens::AbstractArray)
    ops = PathOp[]
    last_op_type = nothing
    i = 1
    while i <= length(tokens)
        tok = tokens[i]
        strt = i
        if isa(tok, Symbol)
            if !haskey(path_ops, tok)
                error("$(tok) is not a valid path operation")
            else
                op_type = path_ops[tok]
                i += 1
                op, i = parsepathop(op_type, tokens, i)
                push!(ops, op)
                last_op_type = op_type
            end
        else
            op, i = parsepathop(last_op_type, tokens, i)
            push!(ops, op)
        end
    end

    return ops
end

struct PathPrimitive <: FormPrimitive
    ops::Vector{PathOp}
end

const Path = Form{PathPrimitive}

path(tokens::AbstractArray, tag=empty_tag) = Path([PathPrimitive(parsepath(tokens))], tag)

path(tokens::AbstractArray{T}, tag=empty_tag) where T <: AbstractArray =
        Path([PathPrimitive(parsepath(ts)) for ts in tokens], tag)

resolve(box::AbsoluteBox, units::UnitBox, t::Transform, p::PathPrimitive) =
        PathPrimitive([resolve(box, units, t, op) for op in p.ops])

# TODO: boundingbox


using OrderedCollections

include("graph.jl")

"""
    LazyBridgeOptimizer{OT<:MOI.ModelLike} <: AbstractBridgeOptimizer

The `LazyBridgeOptimizer` combines several bridges, which are added using the [`add_bridge`](@ref) function.
Whenever a constraint is added, it only attempts to bridge it if it is not supported by the internal model (hence its name `Lazy`).
When bridging a constraint, it selects the minimal number of bridges needed.
For instance, a constraint `F`-in-`S` can be bridged into a constraint `F1`-in-`S1` (supported by the internal model) using bridge 1 or
bridged into a constraint `F2`-in-`S2` (unsupported by the internal model) using bridge 2 which can then be
bridged into a constraint `F3`-in-`S3` (supported by the internal model) using bridge 3,
it will choose bridge 1 as it allows to bridge `F`-in-`S` using only one bridge instead of two if it uses bridge 2 and 3.
"""
mutable struct LazyBridgeOptimizer{OT<:MOI.ModelLike} <: AbstractBridgeOptimizer
    # Internal model
    model::OT
    # Bridged variables
    variable_map::Variable.Map
    var_to_name::Dict{MOI.VariableIndex, String}
    name_to_var::Union{Dict{String, MOI.VariableIndex}, Nothing}
    # Bridged constraints
    constraint_map::Constraint.Map
    con_to_name::Dict{MOI.ConstraintIndex, String}
    name_to_con::Union{Dict{String, MOI.ConstraintIndex}, Nothing}
    # Bridged objective
    objective_map::Objective.Map
    # Bellman-Ford graph
    graph::Graph
    # List of types of available bridges
    variable_free_bridge_type::Union{Nothing, DataType}
    variable_bridge_types::Vector{Any}
    variable_node::OrderedDict{Tuple{DataType}, VariableNode}
    variable_types::Vector{Tuple{DataType}}
    # List of types of available bridges
    constraint_bridge_types::Vector{Any}
    constraint_node::OrderedDict{Tuple{DataType, DataType}, ConstraintNode}
    constraint_types::Vector{Tuple{DataType, DataType}}
    # List of types of available bridges
    objective_bridge_types::Vector{Any}
    objective_node::OrderedDict{Tuple{DataType}, ObjectiveNode}
    objective_types::Vector{Tuple{DataType}}
end
function LazyBridgeOptimizer(model::MOI.ModelLike)
    return LazyBridgeOptimizer{typeof(model)}(
        model,
        Variable.Map(), Dict{MOI.VariableIndex, String}(), nothing,
        Constraint.Map(), Dict{MOI.ConstraintIndex, String}(), nothing,
        Objective.Map(),
        Graph(),
        nothing, Any[], OrderedDict{Tuple{DataType}, VariableNode}(), Tuple{DataType}[],
        Any[], OrderedDict{Tuple{DataType, DataType}, ConstraintNode}(), Tuple{DataType, DataType}[],
        Any[], OrderedDict{Tuple{DataType}, ObjectiveNode}(), Tuple{DataType}[],)
end

function Variable.bridges(bridge::LazyBridgeOptimizer)
    return bridge.variable_map
end
function Constraint.bridges(bridge::LazyBridgeOptimizer)
    return bridge.constraint_map
end
function Objective.bridges(b::LazyBridgeOptimizer)
    return b.objective_map
end

# After `add_bridge(b, BT)`, some constrained variables `(S,)` in
# `keys(b.variable_best)` or constraints `(F, S)` in `keys(b.constraint_best)`
# or `(F,)` in `keys(b.objective_best)` may be bridged
# with less bridges than `b.variable_dist[(S,)]`,
# `b.constraint_dist[(F, S)]` or `b.objective_dist[(F,)]` using `BT`.
# We could either recompute the distance from every node or clear the
# dictionary so that the distance is computed lazily at the next `supports_constraint`
# call. We prefer clearing the dictionaries so as this is called for each
# bridge added and recomputing the distance for each bridge would be a wase
# if several bridges are added consecutively.
function _reset_dist(b::LazyBridgeOptimizer)
    empty!(b.variable_node)
    empty!(b.variable_types)
    empty!(b.constraint_node)
    empty!(b.constraint_types)
    empty!(b.objective_node)
    empty!(b.objective_types)
    empty!(b.graph)
end

function edge(b::LazyBridgeOptimizer, bridge_index, BT::Type{<:AbstractBridge})
    return Edge(
        bridge_index,
        VariableNode[node(b, C[1]) for C in added_constrained_variable_types(BT)],
        ConstraintNode[node(b, C[1], C[2]) for C in added_constraint_types(BT)]
    )
end
function edge(b::LazyBridgeOptimizer, bridge_index, BT::Type{<:Objective.AbstractBridge})
    return ObjectiveEdge(
        bridge_index,
        VariableNode[node(b, C[1]) for C in added_constrained_variable_types(BT)],
        ConstraintNode[node(b, C[1], C[2]) for C in added_constraint_types(BT)],
        node(b, set_objective_function_type(BT))
    )
end

functionized_type(::Nothing) = nothing
functionized_type(::Type{<:Constraint.ScalarFunctionizeBridge{T}}) where T = MOI.ScalarAffineFunction{T}
functionized_type(::Type{<:Constraint.VectorFunctionizeBridge{T}}) where T = MOI.VectorAffineFunction{T}
function functionized_type(b::LazyBridgeOptimizer, ::Type{MOI.SingleVariable})
    return functionized_type(_first_functionize_bridge(b.constraint_bridge_types, Constraint.ScalarFunctionizeBridge))
end
function functionized_type(b::LazyBridgeOptimizer, ::Type{MOI.VectorOfVariables})
    return functionized_type(_first_functionize_bridge(b.constraint_bridge_types, Constraint.VectorFunctionizeBridge))
end

function node(b::LazyBridgeOptimizer, S::Type{<:MOI.AbstractSet})
    F = MOIU.variable_function_type(S)
    if (MOI.supports_constraint(b.model, F, S)
        || (S <: MOI.AbstractScalarSet && MOI.supports_add_constrained_variable(b.model, S))
        || (S <: MOI.AbstractVectorSet && MOI.supports_add_constrained_variables(b.model, S)))
        return VariableNode(0)
    end
    variable_node = get(b.variable_node, (S,), nothing)
    if variable_node !== nothing
        return variable_node
    end
    variable_node = add_variable_node(b.graph)
    b.variable_node[(S,)] = variable_node
    push!(b.variable_types, (S,))
    if is_bridged(b, MOI.Reals)
        FF = functionized_type(b, F)
        if FF !== nothing
            # We assume the distance of the variable node `MOI.Reals` is `1`,
            # i.e. it is bridged by `Variable.FreeBridge` and then
            # the distance of `MOI.Nonnegatives` is zero.
            # We also use the functionize bridge which has cost 1.
            # And we add `+1` as we treat constrained variables as constraints.
            set_variable_constraint_node(b.graph, variable_node, node(b, FF, S), 3)
        end
    else
        # We add `+1` as we treat constrained variables as constraints.
        set_variable_constraint_node(b.graph, variable_node, node(b, F, S), 1)
    end
    for (i, BT) in enumerate(b.variable_bridge_types)
        if Variable.supports_constrained_variable(BT, S)
            add_edge(b.graph, variable_node,
                     edge(b, i, Variable.concrete_bridge_type(BT, S))::Edge)
        end
    end
    return variable_node
end

function node(b::LazyBridgeOptimizer, F::Type{<:MOI.AbstractFunction},
              S::Type{<:MOI.AbstractSet})
    if MOI.supports_constraint(b.model, F, S)
        return ConstraintNode(0)
    end
    constraint_node = get(b.constraint_node, (F, S), nothing)
    if constraint_node !== nothing
        return constraint_node
    end
    constraint_node = add_constraint_node(b.graph)
    b.constraint_node[(F, S)] = constraint_node
    push!(b.constraint_types, (F, S))
    for (i, BT) in enumerate(b.constraint_bridge_types)
        if MOI.supports_constraint(BT, F, S)
            add_edge(b.graph, constraint_node,
                     edge(b, i, Constraint.concrete_bridge_type(BT, F, S))::Edge)
        end
    end
    return constraint_node
end

function node(b::LazyBridgeOptimizer, F::Type{<:MOI.AbstractScalarFunction})
    if MOI.supports(b.model, MOI.ObjectiveFunction{F}())
        return ObjectiveNode(0)
    end
    objective_node = get(b.objective_node, (F,), nothing)
    if objective_node !== nothing
        return objective_node
    end
    objective_node = add_objective_node(b.graph)
    b.objective_node[(F,)] = objective_node
    push!(b.objective_types, (F,))
    # Add the added nodes to the graph.
    for (i, BT) in enumerate(b.objective_bridge_types)
        if Objective.supports_objective_function(BT, F)
            add_edge(b.graph, objective_node,
                     edge(b, i, Objective.concrete_bridge_type(BT, F))::ObjectiveEdge)
        end
    end
    return objective_node
end

function _bridge_types(b::LazyBridgeOptimizer, ::Type{<:Variable.AbstractBridge})
    return b.variable_bridge_types
end
function _bridge_types(b::LazyBridgeOptimizer, ::Type{<:Constraint.AbstractBridge})
    return b.constraint_bridge_types
end
function _bridge_types(b::LazyBridgeOptimizer, ::Type{<:Objective.AbstractBridge})
    return b.objective_bridge_types
end

"""
    add_bridge(b::LazyBridgeOptimizer, BT::Type{<:AbstractBridge})

Enable the use of the bridges of type `BT` by `b`.
"""
function add_bridge(b::LazyBridgeOptimizer, BT::Type{<:AbstractBridge})
    if !has_bridge(b, BT)
        _add_bridge(b, BT)
        _reset_dist(b)
    end
end
function _add_bridge(b::LazyBridgeOptimizer, BT::Type{<:AbstractBridge})
    push!(_bridge_types(b, BT), BT)
end
function _add_bridge(b::LazyBridgeOptimizer, BT::Type{<:Variable.FreeBridge})
    if b.variable_free_bridge_type !== nothing
        error("Bridge $BT cannot be added as $(b.variable_free_bridge_type) has already been added. Remove it first with `M̀OI.Bridges.remove_bridge`.")
    else
        b.variable_free_bridge_type = BT
    end
end

"""
    remove_bridge(b::LazyBridgeOptimizer, BT::Type{<:AbstractBridge})

Disable the use of the bridges of type `BT` by `b`.
"""
function remove_bridge(b::LazyBridgeOptimizer, BT::Type{<:AbstractBridge})
    if !_remove_bridge(b, BT)
        error("Cannot remove bridge `$BT` as it was never added or was already",
              " removed.")
    end
    _reset_dist(b)
end
function _remove_bridge(b::LazyBridgeOptimizer, BT::Type{<:AbstractBridge})
    bridge_types = _bridge_types(b, BT)
    i = findfirst(isequal(BT), bridge_types)
    if i === nothing
        return false
    end
    deleteat!(bridge_types, i)
    return true
end
function _remove_bridge(b::LazyBridgeOptimizer, BT::Type{<:Variable.FreeBridge})
    if b.variable_free_bridge_type === nothing || b.variable_free_bridge_type != BT
        return false
    end
    b.variable_free_bridge_type = nothing
    return true
end

"""
    has_bridge(b::LazyBridgeOptimizer, BT::Type{<:AbstractBridge})

Return a `Bool` indicating whether the bridges of type `BT` are used by `b`.
"""
function has_bridge(b::LazyBridgeOptimizer, BT::Type{<:AbstractBridge})
    return findfirst(isequal(BT), _bridge_types(b, BT)) !== nothing
end
function has_bridge(b::LazyBridgeOptimizer, BT::Type{<:Variable.FreeBridge})
    return b.variable_free_bridge_type == BT
end

# It only bridges when the constraint is not supporting, hence the name "Lazy"
function is_bridged(b::LazyBridgeOptimizer, S::Type{<:MOI.AbstractSet})
    return !MOI.supports_constraint(b.model, MOIU.variable_function_type(S), S)
end
function is_bridged(b::LazyBridgeOptimizer, F::Type{<:MOI.AbstractFunction}, S::Type{<:MOI.AbstractSet})
    return !MOI.supports_constraint(b.model, F, S)
end
function is_bridged(b::LazyBridgeOptimizer, F::Type{<:MOI.AbstractScalarFunction})
    return !MOI.supports(b.model, MOI.ObjectiveFunction{F}())
end

function bridge_index(b::LazyBridgeOptimizer, S::Type{<:MOI.AbstractSet})
    return bridge_index(b.graph, node(b, S))
end
function bridge_index(b::LazyBridgeOptimizer, F::Type{<:MOI.AbstractFunction},
                      S::Type{<:MOI.AbstractSet})
    n = node(b, F, S)
    index = bridge_index(b.graph, n)
    return index
end
function bridge_index(b::LazyBridgeOptimizer, F::Type{<:MOI.AbstractFunction})
    return bridge_index(b.graph, node(b, F))
end

supports_constraint_bridges(::LazyBridgeOptimizer) = true
function supports_bridging_constrained_variable(
    b::LazyBridgeOptimizer, ::Type{MOI.Reals}
)
    return b.variable_free_bridge_type !== nothing
end
function supports_bridging_constrained_variable(
    b::LazyBridgeOptimizer, S::Type{<:MOI.AbstractSet}
)
    return !iszero(bridge_index(b, S))
end
function supports_bridging_constraint(
    b::LazyBridgeOptimizer, F::Type{<:MOI.AbstractFunction},
    S::Type{<:MOI.AbstractSet}
)
    return !iszero(bridge_index(b, F, S))
end
function supports_bridging_objective_function(
    b::LazyBridgeOptimizer, F::Type{<:MOI.AbstractScalarFunction}
)
    return !iszero(bridge_index(b, F))
end
function is_variable_bridged(b::LazyBridgeOptimizer, S::Type{<:MOI.AbstractSet})
    return is_variable_edge_best(b.graph, node(b, S))
end
function bridge_type(b::LazyBridgeOptimizer, ::Type{MOI.Reals})
    if b.variable_free_bridge_type === nothing
        throw(MOI.UnsupportedConstraint{MOI.VectorOfVariables, MOI.Reals}())
    end
    return b.variable_free_bridge_type
end
function bridge_type(b::LazyBridgeOptimizer, S::Type{<:MOI.AbstractSet})
    index = bridge_index(b, S)
    if iszero(index)
        throw(MOI.UnsupportedConstraint{MOIU.variable_function_type(S), S}())
    end
    return Variable.concrete_bridge_type(b.variable_bridge_types[index], S)
end
function bridge_type(b::LazyBridgeOptimizer, F::Type{<:MOI.AbstractFunction}, S::Type{<:MOI.AbstractSet})
    index = bridge_index(b, F, S)
    if iszero(index)
        throw(MOI.UnsupportedConstraint{F, S}())
    end
    return Constraint.concrete_bridge_type(b.constraint_bridge_types[index], F, S)
end
function bridge_type(b::LazyBridgeOptimizer, F::Type{<:MOI.AbstractScalarFunction})
    index = bridge_index(b, F)
    if iszero(index)
        throw(MOI.UnsupportedAttribute(MOI.ObjectiveFunction{F}()))
    end
    return Objective.concrete_bridge_type(b.objective_bridge_types[index], F)
end

_func_name(::Type{Constraint.ScalarFunctionizeBridge}) = "SingleVariable", "constraint"
_func_name(::Type{Constraint.VectorFunctionizeBridge}) = "VectorOfVariables", "constraint"
_func_name(::Type{Objective.FunctionizeBridge}) = "SingleVariable", "objective"
function _first_functionize_bridge(bridge_types, target_type)
    index = findfirst(bridge_type -> bridge_type <: target_type,
                      bridge_types)
    if index === nothing
        return nothing
    end
    return bridge_types[index]
end
function _functionize_bridge(bridge_types, target_type)
    bridge_type = _first_functionize_bridge(bridge_types, target_type)
    if bridge_type === nothing
        func, name = _func_name(target_type)
        error("Need to apply a `$target_type` to a",
              " `$func` $name because the variable is",
              " bridged but no such $name bridge type was added. Add one",
              " with `add_bridge`.")
    end
    return bridge_type
end
function constraint_vector_functionize_bridge(b::LazyBridgeOptimizer)
    return _functionize_bridge(b.constraint_bridge_types, Constraint.VectorFunctionizeBridge)
end
function constraint_scalar_functionize_bridge(b::LazyBridgeOptimizer)
    return _functionize_bridge(b.constraint_bridge_types, Constraint.ScalarFunctionizeBridge)
end
function objective_functionize_bridge(b::LazyBridgeOptimizer)
    return _functionize_bridge(b.objective_bridge_types, Objective.FunctionizeBridge)
end


# Default barriers:
SyncBarriers.Barrier(n::Integer) = _Barrier(n, Threads.nthreads())
SyncBarriers.fuzzy_barrier(n::Integer) = _fuzzy_barrier(n, Threads.nthreads())
SyncBarriers.fuzzy_reduce_barrier(op, ::Type{T}, n::Integer) where {T} =
    _fuzzy_reduce_barrier(op, T, n, Threads.nthreads())
SyncBarriers.reduce_barrier(op, ::Type{T}, n::Integer) where {T} =
    _reduce_barrier(op, T, n, Threads.nthreads())

# TODO: find the best cutoff
# TODO: make it configurable?
function _Barrier(n::Integer, nthreads::Integer)
    nthreads = min(nthreads, n)
    if nthreads >= 32
        return DisseminationBarrier(n)
    else
        return CentralizedBarrier(n)
    end
end

function _fuzzy_barrier(n::Integer, nthreads::Integer)
    nthreads = min(nthreads, n)
    if nthreads >= 32
        return TreeBarrier{8}(n)
    else
        return CentralizedBarrier(n)
    end
end

function _fuzzy_reduce_barrier(op, ::Type{T}, n::Integer, nthreads::Integer) where {T}
    nthreads = min(nthreads, n)
    if nthreads >= 32
        return TreeBarrier{8,T}(op, n)
    elseif nthreads >= 16
        return FlatTreeBarrier{4,T}(op, n)
    else
        return FlatTreeBarrier{2,T}(op, n)
    end
end

function _reduce_barrier(op, ::Type{T}, n::Integer, nthreads::Integer) where {T}
    nthreads = min(nthreads, n)
    if nthreads >= 32
        return StaticTreeBarrier{8,8,T}(op, n)
    else
        return _fuzzy_reduce_barrier(op, T, n, nthreads)
    end
end

SyncBarriers.CentralizedBarrier(n) = CentralizedBarrier(n)
SyncBarriers.DisseminationBarrier(n) = DisseminationBarrier(n)
SyncBarriers.StaticTreeBarrier{NArrive,NDepart}(n) where {NArrive,NDepart} =
    StaticTreeBarrier{NArrive,NDepart}(n)
SyncBarriers.StaticTreeBarrier{NArrive,NDepart,T}(op, n) where {NArrive,NDepart,T} =
    StaticTreeBarrier{NArrive,NDepart,T}(op, n)
SyncBarriers.TreeBarrier{N}(n) where {N} = TreeBarrier{N}(n)
SyncBarriers.TreeBarrier{N,T}(op, n) where {N,T} = TreeBarrier{N,T}(op, n)
SyncBarriers.FlatTreeBarrier{N}(n) where {N} = FlatTreeBarrier{N}(n)
SyncBarriers.FlatTreeBarrier{N,T}(op, n) where {N,T} = FlatTreeBarrier{N,T}(op, n)

const FuzzyReduceBarrier{T,NBranches} =
    Union{TreeBarrier{NBranches,T},FlatTreeBarrier{NBranches,T}}
const FuzzyBarrier = Union{CentralizedBarrier,TreeBarrier,FlatTreeBarrier}

# Used only in testing ATM:
const ReduceBarrier{T} =
    Union{TreeBarrier{<:Any,T},FlatTreeBarrier{<:Any,T},StaticTreeBarrier{<:Any,<:Any,T}}
acctype(::Type{B}) where {T,B<:ReduceBarrier{T}} = T
acctype(b::SyncBarriers.Barrier) = acctype(typeof(b))

function SyncBarriers.cycle!(
    handle::BarrierHandle{<:FuzzyBarrier},
    spin::Union{Integer,Nothing} = nothing,
)
    islast = SyncBarriers.arrive!(handle)
    islast || SyncBarriers.depart!(handle, spin)
    return islast
end

SyncBarriers.arrive!(handle::BarrierHandle{<:FuzzyReduceBarrier{Nothing}}) =
    SyncBarriers.reduce_arrive!(handle, nothing) === Some(nothing)

Base.summary(io::IO, ::CentralizedBarrier) = show(io, SyncBarriers.CentralizedBarrier)
Base.summary(io::IO, ::DisseminationBarrier) = show(io, SyncBarriers.DisseminationBarrier)
Base.summary(io::IO, ::StaticTreeBarrier{NArrive,NDepart,T}) where {NArrive,NDepart,T} =
    show(io, SyncBarriers.StaticTreeBarrier{NArrive,NDepart,T})
Base.summary(io::IO, ::TreeBarrier{T,N}) where {T,N} = show(io, SyncBarriers.TreeBarrier{T,N})
Base.summary(io::IO, ::FlatTreeBarrier{T,N}) where {T,N} =
    show(io, SyncBarriers.FlatTreeBarrier{T,N})

function Base.show(io::IO, ::MIME"text/plain", barrier::SyncBarriers.Barrier)
    summary(io, barrier)
    print(io, " for ", length(barrier), " task(s)")
end


module Types
    import ...SSGrB: Abstract_GrB_Struct, AbstractUnaryOp, AbstractMonoid, AbstractSelectOp,
    AbstractSemiring, AbstractBinaryOp, AbstractDescriptor
    using ..libgb
end

mutable struct GBScalar{T} <: Abstract_GrB_Struct
    p::libgb.GxB_Scalar
    function GBScalar{T}(p::libgb.GxB_Scalar) where {T}
        s = new(p)
        function f(scalar)
            libgb.GxB_Scalar_free(Ref(scalar.p))
        end
        return finalizer(f, s)
    end
end

mutable struct GBVector{T} <: AbstractSparseArray{T, UInt64, 1}
    p::libgb.GrB_Vector
    function GBVector{T}(p::libgb.GrB_Vector) where {T}
        v = new(p)
        function f(vector)
            libgb.GrB_Vector_free(Ref(vector.p))
        end
        return finalizer(f, v)
    end
end

mutable struct GBMatrix{T} <: AbstractSparseArray{T, UInt64, 2}
    p::libgb.GrB_Matrix
    function GBMatrix{T}(p::libgb.GrB_Matrix) where {T}
        A = new(p)
        function f(matrix)
            libgb.GrB_Matrix_free(Ref(matrix.p))
        end
        return finalizer(f, A)
    end
end

const GBVecOrMat = Union{GBVector{T}, GBMatrix{T}} where T
const OpUnion = Union{
    libgb.GrB_Monoid,
    libgb.GrB_UnaryOp,
    libgb.GrB_Semiring,
    libgb.GrB_BinaryOp,
    libgb.GxB_SelectOp
}

const MonoidSemiringOrBinary = Union{
    libgb.GrB_Monoid,
    libgb.GrB_Semiring,
    libgb.GrB_BinaryOp,
    AbstractSemiring,
    AbstractBinaryOp,
    AbstractMonoid
}
const GBMatOrTrans = Union{<:GBMatrix, Transpose{<:Any, <:GBMatrix}}
const GBVecMatOrTrans = Union{<:GBVector, GBMatOrTrans}

function getoperator(op, t)
    #Default Semiring should be LOR_LAND for boolean
    if op == Semirings.PLUS_TIMES_SEMIRING
        if t == Bool
            op = Semirings.LOR_LAND_SEMIRING
        end
    end
    #Default BinaryOp should be LAND or LOR for boolean
    if op == BinaryOps.TIMES
        if t == Bool
            op = BinaryOps.LAND
        end
    elseif op == BinaryOps.PLUS
        if t == Bool
            op = BinaryOps.LOR
        end
    end
    #Default monoid should be LAND/LOR
    if op == Monoids.TIMES_MONOID
        if t == Bool
            op = Monoids.LAND_MONOID
        end
    elseif op == Monoids.PLUS_MONOID
        if t == Bool
            op = Monoids.LOR_MONOID
        end
    end

    if op isa AbstractOp
        return op[t]
    elseif op isa OpUnion
        return op
    else
        error("Not a valid GrB op/semiring")
    end
end


export QBasedPolicy

using MacroTools: @forward
using Flux
using Setfield

"""
    QBasedPolicy(;learner::Q, explorer::S)

Use a Q-`learner` to generate estimations of action values.
Then an `explorer` is applied on the estimations to select an action.
"""
Base.@kwdef mutable struct QBasedPolicy{Q<:AbstractLearner,E<:AbstractExplorer} <:
                           AbstractPolicy
    learner::Q
    explorer::E
end

Flux.functor(x::QBasedPolicy) = (learner = x.learner,), y -> @set x.learner = y.learner

(π::QBasedPolicy)(env) = π(env, ActionStyle(env))
(π::QBasedPolicy)(env, ::MinimalActionSet) = get_actions(env)[π.explorer(π.learner(env))]
(π::QBasedPolicy)(env, ::FullActionSet) =
    get_actions(env)[π.explorer(π.learner(env), get_legal_actions_mask(env))]

RLBase.get_prob(p::QBasedPolicy, env) = get_prob(p, env, ActionStyle(env))
RLBase.get_prob(p::QBasedPolicy, env, ::MinimalActionSet) =
    get_prob(p.explorer, p.learner(env))
RLBase.get_prob(p::QBasedPolicy, env, ::FullActionSet) =
    get_prob(p.explorer, p.learner(env), get_legal_actions_mask(env))

@forward QBasedPolicy.learner RLBase.get_priority, RLBase.update!

function Flux.testmode!(p::QBasedPolicy, mode = true)
    testmode!(p.learner, mode)
    testmode!(p.explorer, mode)
end


nelements(node::Leaf) = length(node)
nelements(node::Branch) = mapreduce(nelements, +, children(node), init=0)

# check the R-tree node (branch or leaf):
# * the parent
# * the level
# * the number of attached elements
# * element id uniqueness
# * the MBR (should fit the MBR of children)
function check(node::Node{T,N,V}, tree::RTree{T,N,V},
               nnodes_perlevel::AbstractVector{Int}, ids::Union{Set, Nothing}) where {T,N,V}
    nnodes_perlevel[level(node)] += 1
    for (i, child) in enumerate(children(node))
        if node isa Branch
            child.parent === node ||
                throw(SpatialIndexException("Node (lev=$(level(node)) len=$(length(node))): child #$i references different parent"))
            level(child) + 1 == level(node) ||
                throw(SpatialIndexException("Node (lev=$(level(node)) len=$(length(node))): child #$i level ($(level(child))) doesn't match parent's ($(level(node)))"))
            check(child, tree, nnodes_perlevel, ids)
        elseif idtrait(typeof(child)) !== HasNoID
            cid = id(child)
            (cid ∈ ids) && throw(SpatialIndexException("Duplicate data id=$cid"))
            push!(ids, cid)
        end
    end
    br = !isempty(node) ? mapreduce(mbr, combine, node.children) : empty(Rect{T,N})
    if isempty(node)
        isequal(br, mbr(node)) ||
            throw(SpatialIndexException("Node (lev=$(level(node)) len=$(length(node))) MBR ($(mbr(node))) should be empty"))
    elseif tree.tight_mbrs && br != mbr(node) || !in(br, mbr(node))
        throw(SpatialIndexException("Node (lev=$(level(node)) len=$(length(node))) MBR ($(mbr(node))) doesn't match its children MBR ($(br))"))
    end
    return true
end

# check the tree:
# * the root
# * the height
# * the number of elements
# * the number of nodes per level
# * element id uniqueness
# * MBRs of the nodes
function check(tree::RTree)
    tree.root.parent === nothing ||
        throw(SpatialIndexException("RTree: root has a parent"))
    level(tree.root) == height(tree) ||
        throw(SpatialIndexException("RTree: root level ($(level(tree.root))) doesn't match tree height ($(height(tree)))"))
    length(tree.nnodes_perlevel) == height(tree) ||
        throw(SpatialIndexException("RTree: nnodes_perlevel length ($(length(tree.nnodes_perlevel))) doesn't match tree height ($(height(tree)))"))
    ids = idtype(tree) !== Union{} ? Set{idtype(tree)}() : nothing
    nnodes_perlevel = zeros(Int, height(tree))
    check(tree.root, tree, nnodes_perlevel, ids)
    nelms = nelements(tree.root)
    length(tree) == nelms ||
        throw(SpatialIndexException("RTree: actual ($(nelms)) and reported ($(length(tree))) number of elements do not match"))
    ids === nothing || length(tree) == length(ids) ||
        throw(SpatialIndexException("RTree: the number of ids ($(length(ids))) doesn't match the reported number of elements ($(length(tree)))"))
    tree.nnodes_perlevel == nnodes_perlevel ||
        throw(SpatialIndexException("RTree: actual ($(nnodes_perlevel)) and reported ($(tree.nnodes_perlevel)) number of nodes per level do not match"))
    return true
end


using FileIO,JLD2,MAT,YAML,Mmap,LoopVectorization
import FileIO: save

include("SpikeGLX.jl")
include("Imager.jl")

"""
Drop array dims according to `MATLAB` convention.

- scalar: scalar instead of 1x1 array
- rvector: vector instead of 1xN array
- cvector: vector instead of Nx1 array
"""
dropmatdim!(x;scalar = true, rvector = true, cvector = true) = x
function dropmatdim!(x::Dict;scalar = true, rvector = true, cvector = true)
    foreach(k->x[k]=dropmatdim!(x[k],scalar=scalar,rvector=rvector,cvector=cvector),keys(x))
    return x
 end
function dropmatdim!(x::Array;scalar = true, rvector = true, cvector = true)
    if ndims(x) == 2
        nr,nc = size(x)
        if nr==1 && nc==1 && scalar
            return dropmatdim!(x[1,1],scalar=scalar,rvector=rvector,cvector=cvector)
        elseif nr==1 && nc>1 && rvector
            x = dropdims(x,dims=1)
        elseif nr>1 && nc==1 && cvector
            x = dropdims(x,dims=2)
        end
    end
    if x isa Array{Any}
        foreach(i->x[i]=dropmatdim!(x[i],scalar=scalar,rvector=rvector,cvector=cvector),eachindex(x))
    end
    return x
end

"""
Read variables of a `MATLAB` MAT file into a Dictionary

1. f: MAT file path
2. vars: variable names in the varset to read

- varset: variable names to select in vars
- scalar: scalar instead of 1x1 matrix
- rvector: vector instead of 1xN matrix
- cvector: vector instead of Nx1 matrix
"""
function readmat(f::AbstractString,vars...;varset=["spike","lfp","digital","analog1k","image"],scalar = true, rvector = true, cvector = true)
    if isempty(vars)
        d=matread(f)
    else
        d=Dict()
        matopen(f) do file
            fvs = names(file)
            nonselectvars=setdiff(fvs,varset)
            selectvars = intersect(varset,vars)
            for i in setdiff(selectvars,fvs)
                @warn """Selected variable "$i" not found in "$f"."""
            end
            readvars=union(nonselectvars,intersect(fvs,selectvars))
            if !isempty(readvars)
                for v in readvars
                    d[v]=read(file,v)
                end
            end
        end
    end
    if scalar || rvector || cvector
        d = dropmatdim!(d,scalar=scalar,rvector=rvector,cvector=cvector)
    end
    return d
end

"Read `DataExport` Metadata MAT File"
readmeta(f::AbstractString) = DataFrame(readmat(f)["Tests"])

"Load images in directory"
function loadimageset(dir;name=[],n=typemax(Int),alpha=false)
    isdir(dir) || error("Invalid Directory")
    for (root, dirs, files) in walkdir(dir)
        if isempty(name)
            n = min(n,length(files))
            n<=0 && error("n <= 0")
            imgs=load.(joinpath.([root],string.(1:n).*splitext(files[1])[2]))
        else
            imgs=load.(joinpath.([root],string.(Int.(name)).*splitext(files[1])[2]))
        end
        if alpha
            imgs = map(i->coloralpha.(i),imgs)
        end
        return imgs
    end
end

"Default digital input channels"
const CondDCh=1
const MarkDCh=2
const StartDCh=3
const StopDCh=4
const Bits16DCh=5
"Get digital channel time and value."
function digitaldata(dataset::Dict,ch)
    chidx = findfirst(dataset["digital"]["channel"].==ch)
    if !isempty(chidx)
        return dataset["digital"]["time"][chidx],dataset["digital"]["data"][chidx]
    else
        return [],[]
    end
end
"Parse bit event time and value."
function digitalbit(dt,dv,bits...)
    n = length(dt)
    maxbit = maximum(bits)
    bn = length(bits)
    bt = [zeros(Float64,n+1) for _ in 1:bn]
    bv = [zeros(UInt8,n+1) for _ in 1:bn]
    j=ones(Int,bn)
    for i in 1:n
        ds = digits(UInt8,dv[i],2,maxbit)
        for b in 1:bn
            v=ds[bits[b]]
            if v != bv[b][j[b]]
                j[b]=j[b]+1
                bt[b][j[b]]=dt[i]
                bv[b][j[b]]=v
            end
        end
    end
    for b in 1:bn
        bt[b]=bt[b][2:j[b]]
        bv[b]=bv[b][2:j[b]]
    end
    if bn==1
        return bt[1],bv[1]
    end
    return bt,bv
end

"Read and Prepare Dataset in `MATLAB` MAT File"
prepare(f::AbstractString,vars...)=prepare!(readmat(f,vars...))
function prepare!(d::Dict)
    if haskey(d,"secondperunit")
        settimeunit(d["secondperunit"])
    end
    if haskey(d,"sourceformat")
        fmt = d["sourceformat"]
        if fmt=="Ripple"
            d=prepare_ripple!(d)
        elseif fmt=="OI"
            d=prepare_oi!(d)
        elseif fmt=="Experica"
            d=prepare_experica!(d)
        end
    end
    if haskey(d,"imecindex")
        imecspike=Dict()
        for i in d["imecindex"]
            s = "spike$(i)_kilosort"
            if haskey(d,s)
                imecspike[i] = unitspike_kilosort(d[s],syncindex=i,sortspike=true)
            end
        end
        d = mergeimecspike!(imecspike,d)
    end
    return d
end
"Prepare `Ripple` Data"
function prepare_ripple!(d::Dict)
    if haskey(d,"spike")
        st=d["spike"]["time"]
        su=d["spike"]["unitid"]
        for i in 1:length(st)
            if ndims(st[i])==0
                st[i]=[st[i]]
                su[i]=[su[i]]
            end
        end
    end
    return d
end
"Prepare `Optical Imaging` Block Data"
function prepare_oi!(d::Dict)
    return d
end
"Prepare `Experica` Experiment Data"
function prepare_experica!(d::Dict)
    return d
end

"Extract state time of condtest"
function statetime(ct::Dict;statetype::AbstractString="CONDSTATE",state::AbstractString="COND")
    if haskey(ct,statetype)
        filter!(l->!isempty(l),map(i->begin
        t = filter!(k->!isempty(k),map(j->haskey(j,state) ? j[state] : [],i))
        length(t)==1 ? t[1] : t
    end,ct[statetype]))
else
    []
end
end

function getparam(d::Dict{String,Any},name,fromobject="")
    for k in keys(d)
        startswith(k,name) && endswith(k,fromobject) && return d[k]
    end
    nothing
end

"Get `CondTestCond` DataFrame"
function condtestcond(ctcd::Dict)
    ctc = DataFrame(ctcd)
    return ctc
end
"Get `CondTest` DataFrame"
function condtest(ctd::Dict)
    ct = DataFrame(ctd)
    return ct
end
"Get `CondTest` and `CondTestCond` DataFrame"
function ctctc(ctd::Dict,ctcd::Dict)
    ct = condtest(ctd)
    ctc = condtestcond(ctcd)
    return ct,ctc
end
function ctctc(ex::Dict)
    return ctctc(ex["CondTest"],ex["CondTestCond"])
end

"Map time to data timeline, optionally add latency."
function maptodatatime(x,ex::Dict;addlatency=true)
    t=x*(1+ex["TimerDriftSpeed"])+ex["t0"]
    if addlatency
        t+=ex["Latency"]
    end
    return t
end

"Map time between ref and target, according to the timing diff between target and ref of the same sync signal"
function ref2sync(t,dataset::Dict,si="0")
    haskey(dataset,"syncdt") || return t
    tt = "syncdiff$si"
    if haskey(dataset,tt)
        ref2sync(t,dataset[tt],dataset["syncdt"])
    else
        @warn "No syncdiff for Sync$si, skip timing correction.";return t
    end
end
function sync2ref(t,dataset::Dict,si="0")
    haskey(dataset,"syncdt") || return t
    tt = "syncdiff$si"
    if haskey(dataset,tt)
        sync2ref(t,dataset[tt],dataset["syncdt"])
    else
        @warn "No syncdiff for Sync$si, skip timing correction.";return t
    end
end
ref2sync(t,syncdiff,syncdt=500) = t .+ @view syncdiff[clamp!(round.(Int,t./syncdt),1,length(syncdiff))]
sync2ref(t,syncdiff,syncdt=500) = t .- @view syncdiff[clamp!(round.(Int,t./syncdt),1,length(syncdiff))]

"Epochs of `Neuropixels` Channel Sample `x`, optionally gain corrected(voltage), line noise(60,120,180Hz) removed and bandpass filtered"
function epochsamplenp(x,fs,epochs,chs;meta=[],bandpass=[1,100])
    f = nothing
    if !isempty(meta)
        f = i -> gaincorrectnp(i,meta)
        if !isempty(bandpass)
            f = i -> hlpass(rmline!(gaincorrectnp(i,meta),fs),fs,high=bandpass[1],low=bandpass[2])
        end
    end
    epochsample(x,fs,epochs,chs,fun=f)
end

"""
Organize each spiking unit from `Kilosort` result.
"""
function unitspike_kilosort(data::Dict;syncindex="0",sortspike::Bool=true)
    # each unit
    unitid = data["clusterid"]
    unitgood = data["clustergood"].==1
    unitindex = [data["cluster"].==i for i in unitid]
    unitspike = map(i->data["time"][i],unitindex)
    unittemplate = map(i->data["template"][i],unitindex)
    unitamplitude = map(i->data["amplitude"][i],unitindex)

    chmap = data["chanmap"]
    chposition = data["channelposition"]
    unittemplatesindex = unique.(unittemplate)
    unitposition = vcat(map(i->mean(data["templatesposition"][i,:],dims=1),unittemplatesindex)...)
    unittemplateamplitude = map(i->mean(data["templatesamplitude"][i]),unittemplatesindex)
    # first found template feature as unit template feature
    unittemplatefeature = Dict(k=>data["templateswaveformfeature"][k][first(unittemplatesindex)] for k in keys(data["templateswaveformfeature"]))
    unitfeature = haskey(data,"clusterwaveformfeature") ? data["clusterwaveformfeature"] : unittemplatefeature

    sortspike && foreach(sort!,unitspike)
    return Dict("unitid"=>unitid,"unitgood"=>unitgood,"unitspike"=>unitspike,"chposition"=>chposition,"unitposition"=>unitposition,
                "unittemplateamplitude"=>unittemplateamplitude,"unitfeature"=>unitfeature,"isspikesorted"=>sortspike,"unitsync"=>fill(syncindex,size(unitspike)))
end

function mergeimecspike!(imecspike,d)
    n = length(imecspike)
    if n == 0
        @warn "No IMEC spike data"
    elseif n==1
        d["spike"] = first(values(imecspike))
    else
        @info "Merging spike from multiple IMEC, not implemented"
    end
    return d
end

"convert `Matlab` struct of array to `DataFrame`"
matdictarray2df(d) = DataFrame(Any[squeeze(v,2) for v in values(d)],[Symbol(k) for k in keys(d)])

"Regular expression to match `VLab` data file names"
function vlabregex(testtype;subject="[0-9]",maxch=1,cell="[a-z]",maxrepeat=3,format="mat")
    mr = lpad(maxrepeat,2,0)
    Regex("^$subject+[A-Za-z]+[1-$maxch]$cell$testtype[0-$(mr[1])][0-$(mr[2])]\\.$format")
end

"Get matched file names in path"
function getdatafile(testtype;subject="[0-9]",path="./data")
    matchfile(vlabregex(testtype,subject=subject),path=path)
end

"Regular Expression to match Optical Imaging `VDAQ` block file name"
function oifileregex(;subject="[A-Za-z0-9]",session="[A-Za-z0-9]",experimentid="[0-9]",format="mat")
    Regex("^$subject+_$session*_?E$experimentid+B[0-9]+[.]$format")
end

"Get Optical Imaging `VDAQ` matched files in directory"
function getoifile(;subject="[A-Za-z0-9]",session="[A-Za-z0-9]",experimentid="[0-9]",format="mat",dir="",adddir=false)
    matchfile(oifileregex(subject=subject,session=session,experimentid=experimentid,format=format),dir=dir,adddir=adddir)
end

"Regular Expression to match `Experica` data file name"
function expericafileregex(;subject="[A-Za-z0-9]",session="[A-Za-z0-9]",site="[A-Za-z0-9]",test="[A-Za-z0-9]",maxrepeat=10,format="mat")
    mr = lpad(maxrepeat,2,"0");d2=mr[1];d1=parse(Int,d2)>0 ? 9 : mr[2]
    Regex("^$subject+_$session*_?$site*_?$test+_[0-$d2]?[0-$d1][.]$format")
end

"Get matched `Experica` files in directory"
function getexpericafile(;subject="[A-Za-z0-9]",session="[A-Za-z0-9]",site="[A-Za-z0-9]",test="[A-Za-z0-9]",maxrepeat=10,format="mat",dir="",adddir=false)
    matchfile(expericafileregex(subject=subject,session=session,site=site,test=test,maxrepeat=maxrepeat,format=format),dir=dir,adddir=adddir)
end

"Get matched file names in directory, optionally join directory to get file path"
function matchfile(pattern::Regex;dir="",join::Bool=false)
    fs = filter!(f->occursin(pattern,f),readdir(dir))
    if join
        fs=joinpath.(dir,fs)
    end
    return fs
end

function save(filepath::Union{AbstractString,IO},cm::Dict)
    ext = splitext(filepath)[2]
    if ext == ".yaml"
        data = deepcopy(cm)
        foreach(n->data[n]=Dict(k=>cm isa String ? cm : vec.(cm) for (k,cm) in pairs(data[n])),keys(data))
        YAML.write_file(filepath,data)
    elseif ext == ".mat"
        data = deepcopy(cm)
        foreach(n->data[n]=Dict(k=>cm isa String ? cm : mapreduce(c->vec(c)',vcat,cm) for (k,cm) in pairs(data[n])),keys(data))
        matwrite(filepath,data)
    end
end


const TOKEN_PREFIX = '/'
"""
    JSONPointer(token)

Follows https://tools.ietf.org/html/rfc6901 standard

### Differences are
- Index numbers starts from '1' instead of '0'
- User can declare type with '::T' notation at the end

### Example
"""
struct JSONPointer{T}
    token::Tuple
end
function JSONPointer(token::AbstractString)
    if !startswith(token, TOKEN_PREFIX)
        throw(ArgumentError("JSONPointer must starts with '/' prefix"))
    end

    T = Any
    jk = convert(Array{Any, 1}, split(chop(token; head=1, tail=0), TOKEN_PREFIX))
    if occursin("::", jk[end])
        x = split(jk[end], "::")
        jk[end] = x[1]
        T = (x[2] == "Vector" ? "Vector{Any}" : x[2]) |> Meta.parse |> eval
    end
    @inbounds for i in 1:length(jk)
        if occursin(r"^\d+$", jk[i]) # index of a array
            jk[i] = parse(Int, string(jk[i]))
            if iszero(jk[i])
                throw(AssertionError("Julia uses 1-based indexing"))
            end
        end
    end

    JSONPointer{T}(tuple(jk...))
end

"""
    null_value(p::JSONPointer{T}) where T

provide appropriate value for 'T'
'Real' return 'zero(T)' and 'AbstractString' returns '""'

If user wants different null value for 'T' override 'null_value(p::JSONPointer{T})' method

"""
null_value(p::JSONPointer{T}) where T = missing
function null_value(p::JSONPointer{T}) where T <: Array
    eltype(T) <: Real ? eltype(T)[] :
    eltype(T) <: AbstractString ? eltype(T)[] :
    Any[]
end

Base.Dict(p::JSONPointer) = create_by_pointer(Dict, p)
DataStructures.OrderedDict(p::JSONPointer) = create_by_pointer(OrderedDict, p)

function create_by_pointer(::Type{T}, p::JSONPointer) where T <: AbstractDict
    val = missing

    @inbounds for i in length(p):-1:1
        k = p.token[i]
        if isa(k, Integer)
            tmp = Array{Any, 1}(missing, k)
            tmp[k] = if i == length(p)
                        null_value(p)
                    else
                        val
                    end
            val = tmp
        elseif isa(k, AbstractString)
            val = if i == length(p)
                    T{String, Any}(k => null_value(p))
                else
                    T{String, Any}(k => val)
                end
        end
    end
    return val
end

function create_by_pointer(::Type{T}, arr::Array) where T <: AbstractDict
    template = T(arr[1])
    if length(arr) > 1
        @inbounds for p in arr
            p2 = JSONPointer{Any}(p.token)
            template[p2] = null_value(p)
        end
    end
    return template
end

for T in (Dict, OrderedDict)
    @eval Base.haskey(dict::$T{K,V}, p::JSONPointer) where {K, V} = haskey_by_pointer(dict, p)
    @eval Base.getindex(dict::$T{K,V}, p::JSONPointer) where {K, V} = getindex_by_pointer(dict, p)
    @eval Base.setindex!(dict::$T{K,V}, v, p::JSONPointer) where {K <: AbstractString, V} = setindex_by_pointer!(dict, v, p)
    @eval Base.setindex!(dict::$T{K,V}, v, p::JSONPointer) where {K <: Integer, V} = setindex_by_pointer!(dict, v, p)
end
function haskey_by_pointer(collection, p::JSONPointer)::Bool
    b = true
    val = collection
    @inbounds for (i, k) in enumerate(p.token)
        val = begin
            if isa(val, Array)
                if !isa(k, Integer)
                    missing
                else
                    length(val) >= k ? val[k] : missing
                end
            else
                if isa(k, Integer)
                    missing
                else
                    haskey(val, k) ? val[k] : missing
                end
            end
        end

        if ismissing(val)
            b = false
            break
        end
    end
    return b
end
function getindex_by_pointer(collection, p::JSONPointer, i = 1)
    val = getindex(collection, p.token[i])
    if i < length(p)
        val = getindex_by_pointer(val, p, i+1)
    end
    return val
end

function setindex_by_pointer!(collection::T, v, p::JSONPointer{U}) where {T <: AbstractDict, U}
    v = ismissing(v) ? null_value(p) : v
    if !isa(v, U) &&
        try
            v = convert(eltype(p), v)
        catch e
            msg = isa(v, Array) ? "Vector" : "Any"
            error("$v is not valid value for $p use '::$msg' if you don't need static type")
            throw(e)
        end
    end
    prev = collection
    DT = OrderedDict{String, Any}

    @inbounds for (i, k) in enumerate(p.token)
        if isa(prev, AbstractDict)
            DT = typeof(prev)
        end
        if isa(prev, Array)
            if !isa(k, Integer)
                throw(MethodError(setindex!, k))
            end
            grow_array!(prev, k)
        else
            if isa(k, Integer)
                throw(MethodError(setindex!, k))
            end
            if !haskey(prev, k)
                setindex!(prev, missing, k)
            end
        end

        if i < length(p)
            tmp = getindex(prev, k)
            if ismissing(tmp)
                next_key = p.token[i+1]
                if isa(next_key, Integer)
                    new_data = Array{Any,1}(missing, next_key)
                else
                    new_data = DT(next_key => missing)
                end
                setindex!(prev, new_data, k)
            end
            prev = getindex(prev, k)
        end
    end
    setindex!(prev, v, p.token[end])
end

function grow_array!(arr, target_size)
    x = target_size - length(arr)
    if x > 0
        T = eltype(arr)
        new_arr = Array{T, 1}(undef, x)
        new_arr .= T <: Real ? zero(T) : missing
        append!(arr, new_arr)
    end
    return arr
end

Base.length(x::JSONPointer) = length(x.token)
Base.eltype(x::JSONPointer{T}) where T = T

macro j_str(token)
    JSONPointer(token)
end

function Base.show(io::IO, x::JSONPointer{T}) where T
    print(io,
    "JSONPointer{", T, "}(\"/", join(x.token, "/"), "\")")
end


module BKMaxflow

include("bitmask.jl")
export BKStatusBits, BK_EMPTY, BK_FREE, BK_S, BK_T, BK_ACTIVE, BK_ORPHAN
export BK_S_ACTIVE, BK_T_ACTIVE, BK_S_ORPHAN, BK_T_ORPHAN, BK_S_ACTIVE_ORPHAN, BK_T_ACTIVE_ORPHAN

include("boykov_kolmogorov.jl")
export boykov_kolmogorov

abstract type AbstractBKMaxflow{Tv,Ti} end
struct JuliaImpl{Tv<:Real,Ti<:Integer} <: AbstractBKMaxflow{Tv,Ti} end
struct CImpl{Tv<:Real} <: AbstractBKMaxflow{Tv,Integer} end

function create_graph(::Type{JuliaImpl{Tv,Ti}}, vertexNum::Integer) where {Tv<:Real,Ti<:Integer}
    weights = Vector{Tv}()
    neighbors = [Vector{Tuple{Ti,Ti}}() for i = 1:vertexNum]
    return weights, neighbors
end

function add_edge!(weights::AbstractVector{Tv}, neighbors::AbstractVector{Vector{Tuple{Ti,Ti}}},
                   p::Ti, q::Ti, weightp2q::Tv, weightq2p::Tv) where {Tv<:Real,Ti<:Integer}
    idx = length(weights) ÷ 2
    # store forward and reverse weights compactly for cache efficiency,
    # assume a forward weight is the weight from a smaller index to a bigger index
    if p < q
        push!(weights, weightp2q, weightq2p)
        push!(neighbors[p], (q,idx+1))
        push!(neighbors[q], (p,idx+1))
    else
        push!(weights, weightq2p, weightp2q)
        push!(neighbors[q], (p,idx+1))
        push!(neighbors[p], (q,idx+1))
    end
end

export AbstractBKMaxflow, JuliaImpl, CImpl
export create_graph, add_edge!

include("c_wrapper.jl")

end # module


# Command line args
!isempty(ARGS) || error("No config supplied.")
isfile(ARGS[1]) || error("Cannot read '$(ARGS[1])'")
isabspath(ARGS[1]) || error("Please use an absolute path for the config.")
println("Config supplied: '$(ARGS[1])'")
config_file = ARGS[1]

include(joinpath(@__DIR__, "gmm_utils.jl"))
using SVDD, OneClassActiveLearning
using MLKernels
using Memento
using Serialization, DelimitedFiles, Logging, Random

Random.seed!(0)
include(config_file)

function generate_experiments(data_dir, data_gen_settings::DataGenGMM, data_split_strategy, initial_pool_strategy, solver, num_resamples_initial_pool=5)
    exp_dir = joinpath(data_output_root, exp_name)
    println(exp_dir)
    if isdir(exp_dir)
        print("Type 'yes' or 'y' to delete and overwrite experiment $(exp_dir): ")
        argin = readline()
        if argin == "yes" || argin == "y"
            rm(exp_dir, recursive=true)
        else
            error("Overwriting anyways... Just kidding, nothing happened.")
        end
    else
        !isdir(exp_dir) || error("The experiment directory $(exp_dir) already exists.")
    end

    isdir(exp_dir) || mkpath(exp_dir)
    mkpath(joinpath(exp_dir, "data"))
    mkpath(joinpath(exp_dir, "log", "experiment"))
    mkpath(joinpath(exp_dir, "log", "worker"))
    mkpath(joinpath(exp_dir, "results"))
    @info "Generating experiment directory with name: $exp_dir and config: $(config_file). ($num_resamples_initial_pool resamples of the initial pool)"

    experiments = []

    for x in 1:data_gen_settings.num_repetitions
        for dim in data_gen_settings.dimensions
            for num_g in data_gen_settings.num_gaussians
                @info "[Repetition $x] Starting data generation with $dim dimensions and $num_g gaussians."
                Random.seed!(x)
                gmm = rand_gmm(num_gaussians=num_g, d=dim)
                Random.seed!(x)
                data_test, labels_test = generate_gmm_test_data(gmm, data_gen_settings.outlier_threshold, data_gen_settings.num_data)
                oracle = QuerySynthesisGMMOracle(gmm, data_gen_settings.outlier_threshold)
                @info "[Repetition $x] Data generation done."

                for num_g_train in data_gen_settings.num_gaussians_train
                    gmm_train = deconstruct_gmm(gmm, collect(1:num_g_train))
                    data_train, labels_train = generate_gmm_test_data(gmm_train, gmm, data_gen_settings.outlier_threshold, data_gen_settings.num_data, 0)
                    data, labels = hcat(data_train, data_test), vcat(labels_train, labels_test)
                    train_mask = vcat(trues(length(labels_train)), falses(length(labels_test)))
                    split_strategy = DataSplits(train_mask, .~train_mask, LabeledSplitStrat(), FullSplitStrat())

                    @assert size(data_test, 2) == length(labels_test)
                    @assert size(data_train, 2) == length(labels_train)
                    @assert size(data_train, 1) == size(data_test, 1)
                    @assert size(data, 2) == length(labels)
                    @assert length(labels) == length(train_mask)

                    # write data file
                    data_file = joinpath(exp_dir, "data", "gmm_holdout_$(x)_seed_$(dim)_dim_$(num_g)_gaussians_$(num_g_train)_num_gaussians_train.csv")
                    writedlm(data_file, hcat(data', labels), ',')

                    for ip in initial_pool_strategy
                        ip[1] == "Pnin" || error("Cannot generate experiment with initial pools '$(ip[1])'.")
                        haskey(ip[2], :n) || error("Initial pools '$(ip[1])' has no parameter 'n'.")
                        num_train = ip[2][:n]
                        for n in 1:num_resamples_initial_pool
                            Random.seed!(n)
                            initial_pools = get_initial_pools(data, labels, split_strategy, ip[1], n=num_train)
                            @assert length(labels) == length(initial_pools)

                            param = Dict(:num_al_iterations => num_al_iterations,
                                         :solver => solver,
                                         :initial_pools => initial_pools,
                                         :adjust_K => true,
                                         :initial_pool_resample_version => n,
                                         :classify_precision => classify_precision)

                            for model in models
                                for init_strategy in init_strategies
                                    for query_strategy in query_strategies
                                        out_dir = split(data_file, '/')[end-1]
                                        output_path = joinpath(exp_dir, "results", out_dir)
                                        isdir(output_path) || mkdir(output_path)
                                        experiment = Dict{Symbol, Any}(
                                                :data_file => data_file,
                                                :data_set_name => out_dir,
                                                :split_strategy_name => "Shl",
                                                :initial_pool_strategy_name => ip,
                                                :model => merge(model, Dict(:init_strategy => init_strategy)),
                                                :query_strategy => Dict(:type => query_strategy[:type],
                                                                        :param => query_strategy[:param]),
                                                :split_strategy => split_strategy,
                                                :oracle => oracle,
                                                :param => param)

                                        exp_hash = hash(sprint(print, experiment))
                                        @show data_file, ip, n, exp_hash
                                        @assert exp_hash == hash(sprint(print, deepcopy(experiment)))
                                        experiment[:hash] = "$exp_hash"

                                        out_name = splitext(splitdir(data_file)[2])[1]
                                        out_name = joinpath(output_path, "$(out_name)_$(query_strategy[:type])_$(model[:type])_$(exp_hash).json")
                                        experiment[:output_file] = out_name
                                        experiment[:log_dir] = joinpath(exp_dir, "log")
                                        push!(experiments, deepcopy(experiment))
                                    end
                                end
                            end
                        end
                    end
                end
            end
        end
    end

    # save the experimental setup
    cp(@__FILE__, joinpath(exp_dir, splitdir(@__FILE__)[2]), follow_symlinks=true)
    cp(config_file, joinpath(exp_dir, splitdir(config_file)[2]), follow_symlinks=true)

    @info "Created $exp_dir with $(length(experiments)) instances."

    open(joinpath(exp_dir, "experiment_hashes"), "a") do f
        for e in experiments
            write(f, "$(e[:hash])\n")
        end
    end
    serialize(open(joinpath(exp_dir, "experiments.jser"), "w"), experiments)
    return experiments
end

generate_experiments(data_dir, data_gen_settings, split_strategy, initial_pool_strategy, solver, num_resamples_initial_pool)


export NDMSymm

struct NDMSymm{T,T2,T4} <: MatrixNeuralNetwork
    bare_net::T
    symm_net::T

    ∇b_mat::T2
    ∇h_mat::T2
    ∇d_mat::T2
    ∇w_mat::T2
    ∇u_mat::T2

    symm_map::T4
end
@functor NDMSymm

NDMSymm(n_in::Int, args...) = NDMSymm(STD_REAL_PREC, args...)
NDMSymm(T::Type, hilb::AbstractHilbert, args...) =
    NDMSymm(T, nsites(hilb), args...)
function NDMSymm(T::Type{<:Real}, n_in::Int, αh, αa, permutations, σ::Function=logℒ)
    n_symm = length(permutations)
    @assert length(first(permutations)) == n_in

    symm_net = NDM(T, n_in, αh//n_in, αa//n_in, σ)
    bare_net = NDM(T, n_in, αh//n_in*n_symm, αa//n_in*n_symm, σ)
    set_bare_params!(bare_net, symm_net, permutations)
    ∇b_mat, ∇h_mat, ∇d_mat, ∇w_mat, ∇u_mat = construct_∇matrices(bare_net, symm_net, n_in, αh, αa, permutations)

    NDMSymm(bare_net, symm_net, ∇b_mat, ∇h_mat, ∇d_mat, ∇w_mat, ∇u_mat, permutations)
end

out_type(net::NDMSymm)           = out_type(net.bare_net)
is_analytic(net::NDMSymm)        = is_analytic(net.bare_net)
#weights(cnet::NDMSymm)           = cnet.symm_net

grad_cache(net::NDMSymm)         = grad_cache(net.symm_net)
trainable(net::NDMSymm)          = trainable(net.symm_net)
update!(opt, cnet::NDMSymm, Δ, state=nothing) = (res = update!(opt, weights(cnet), weights(Δ), state);
                                            set_bare_params!(cnet.bare_net, cnet.symm_net, cnet.symm_map);
                                            res)

Base.show(io::IO, m::NDMSymm) = print(io,
    "NDMSymm{$(eltype(m.bare_net.b_μ))}, n=$(length(m.symm_net.b_μ)), αₕ=$(length(m.symm_net.h_μ)), αₐ=$(length(m.symm_net.d_λ)), f=$(m.symm_net.f))")

@inline (net::NDMSymm)(args...) = net.bare_net(args...)
function logψ_and_∇logψ!(∇lnψ_symm, net::NDMSymm, args...)
    lnψ, ∇lnψ = logψ_and_∇logψ(net.bare_net, args...)

    symmetrize_∇logψ_NDM!(∇lnψ_symm, ∇lnψ, net)
    lnψ, ∇lnψ_symm
end

struct NDMSymmCache{BC,BD} <: NNCache{NDMSymm}
    bare_cache::BC
    bare_∇logψ::BD
end

cache(net::NDMSymm)              =
    NDMSymmCache(cache(net.bare_net),
                 grad_cache(net.bare_net))

@inline (net::NDMSymm)(c::NDMSymmCache, args...) = net.bare_net(c.bare_cache, args...)

function logψ_and_∇logψ!(∇lnψ_symm, net::NDMSymm, c::NDMSymmCache, args...)
    lnψ = logψ_and_∇logψ!(c.bare_∇logψ, net.bare_net, c.bare_cache, args...)

    symmetrize_∇logψ_NDM!(∇lnψ_symm, c.bare_∇logψ, net)
    return lnψ
end

#
function symmetrize_∇logψ_NDM!(∇lnψ_symm, ∇lnψ, net)
    mul!(∇lnψ_symm.b_μ, net.∇b_mat, ∇lnψ.b_μ)
    mul!(∇lnψ_symm.b_λ, net.∇b_mat, ∇lnψ.b_λ)
    mul!(∇lnψ_symm.h_μ, net.∇h_mat, ∇lnψ.h_μ)
    mul!(∇lnψ_symm.h_λ, net.∇h_mat, ∇lnψ.h_λ)
    mul!(∇lnψ_symm.d_λ, net.∇d_mat, ∇lnψ.d_λ)
    mul!(vec(∇lnψ_symm.w_μ), net.∇w_mat, vec(∇lnψ.w_μ))
    mul!(vec(∇lnψ_symm.w_λ), net.∇w_mat, vec(∇lnψ.w_λ))
    mul!(vec(∇lnψ_symm.u_μ), net.∇u_mat, vec(∇lnψ.u_μ))
    mul!(vec(∇lnψ_symm.u_λ), net.∇u_mat, vec(∇lnψ.u_λ))
    ∇lnψ_symm
end

function set_bare_params!(bare_net, symm_net, permutations)
    # Load a bit
    n_symm = length(permutations)
    n_in   = length(symm_net.b_μ)
    αh   = size(symm_net.w_μ,1)
    αa   = size(symm_net.u_μ,1)

    # Local bias
    symm_net.b_μ .= sum(symm_net.b_μ)/length(symm_net.b_μ)
    symm_net.b_λ .= sum(symm_net.b_λ)/length(symm_net.b_λ)
    bare_net.b_μ .=  symm_net.b_μ
    bare_net.b_λ .=  symm_net.b_λ

    # Hidden biases
    for (i, (h_μ_i, h_λ_i)) = enumerate(zip(symm_net.h_μ, symm_net.h_λ))
        bare_net.h_μ[((i-1)*n_symm+1):(i*n_symm)] .=  h_μ_i
        bare_net.h_λ[((i-1)*n_symm+1):(i*n_symm)] .=  h_λ_i
    end

    # Ancilla biases
    for (i, d_λ_i) = enumerate(symm_net.d_λ)
        bare_net.d_λ[((i-1)*n_symm+1):(i*n_symm)] .=  d_λ_i
    end

    # Hidden connections
    for f=1:αh
        for (j,perm)=enumerate(permutations)
            for (i, i_p)=enumerate(perm)
                bare_net.w_μ[j + (f-1)*n_symm, i_p] = symm_net.w_μ[f,i]
                bare_net.w_λ[j + (f-1)*n_symm, i_p] = symm_net.w_λ[f,i]
            end
        end
    end

    # Ancilla connections
    for f=1:αa
        for (j,perm)=enumerate(permutations)
            for (i, i_p)=enumerate(perm)
                bare_net.u_μ[j + (f-1)*n_symm, i_p] = symm_net.u_μ[f,i]
                bare_net.u_λ[j + (f-1)*n_symm, i_p] = symm_net.u_λ[f,i]
            end
        end
    end
    bare_net
end

function construct_∇matrices(bare_net, symm_net, n_in, αh, αa, permutations)
    # Load a bit
    n_symm = length(permutations)

    # Local bias matrix
    ∇b_mat   = ones(eltype(bare_net.u_λ), n_in, n_in)./n_in

    # Hidden bias matrix
    ∇h_mat   = zeros(eltype(bare_net.h_λ), αh, αh*n_symm)
    for s=1:αh
        ∇h_mat[s, (s-1)*n_symm+1:s*n_symm] .= 1
    end

    # Ancillary bias matrix
    ∇d_mat   = zeros(eltype(bare_net.d_λ), αa, αa*n_symm)
    for s=1:αa
        ∇d_mat[s, (s-1)*n_symm+1:s*n_symm] .= 1
    end

    ∇w_mat   = zeros(eltype(bare_net.w_λ), αh*n_in, αh*n_symm*n_in)
    lin_ind_bare  = LinearIndices(bare_net.w_λ)
    lin_ind_symm  = LinearIndices(symm_net.w_λ)
    for f=1:αh
        for (j,perm)=enumerate(permutations)
            for (i, i_p)=enumerate(perm)
                row = lin_ind_symm[f,i]
                col = lin_ind_bare[j + (f-1)*n_symm, i_p]
                ∇w_mat[row,col] = 1
            end
        end
    end

    ∇u_mat   = zeros(eltype(bare_net.u_λ), αa*n_in, αa*n_symm*n_in)
    lin_ind_bare  = LinearIndices(bare_net.u_λ)
    lin_ind_symm  = LinearIndices(symm_net.u_λ)
    for f=1:αa
        for (j,perm)=enumerate(permutations)
            for (i, i_p)=enumerate(perm)
                row = lin_ind_symm[f,i]
                col = lin_ind_bare[j + (f-1)*n_symm, i_p]
                ∇u_mat[row,col] = 1
            end
        end
    end
    ∇b_mat, ∇h_mat, ∇d_mat, ∇w_mat, ∇u_mat
end

######


function build_op_model!(op_model::OperationModel{M}; kwargs...) where M<:AbstractOperationModel
    sys = get_system(op_model)
    _build_canonical!(op_model.canonical, op_model.model_ref, sys; kwargs...)
    return
end

function _build_canonical!(canonical::CanonicalModel, ref::ModelReference, sys::PSY.System; kwargs...)

    verbose = get(kwargs, :verbose, true)
    transmission = ref.transmission

    # Build Injection devices
    for (_, device_model) in ref.devices
        verbose && @info "Building $(device_model.device) with $(device_model.formulation) formulation"
        construct_device!(canonical, sys, device_model, transmission; kwargs...)
    end

    # Build Network
    verbose && @info "Building $(transmission) network formulation"
    construct_network!(canonical, sys, transmission; kwargs...)

    # Build Branches
    for (_, branch_model) in ref.branches
        verbose && @info "Building $(branch_model.device) with $(branch_model.formulation) formulation"
        construct_device!(canonical, sys, branch_model, transmission; kwargs...)
    end

    #Build Service
    for (_, service_model) in ref.services
        #construct_service!(canonical, sys, service_model, transmission; kwargs...)
    end

    # Objective Function
    verbose && @info "Building Objective"
    JuMP.@objective(canonical.JuMPmodel,
                    MOI.MIN_SENSE,
                    canonical.cost_function)

    return

end


# Utility function which returns a Dual number of the index matches the target.
@inline function context_getindex(ctx, am, addr)
    if (addr, ) == ctx.target
        Dual(getindex(am, addr), 1.0)
    else
        getindex(am, addr)
    end
end

function forward(addr, params, cl::DynamicCallSite, seed)
    ctx = ForwardMode(addr, params, cl, seed)
    ret = ctx(cl.fn, cl.args...)
    ret, ctx.weight
end

# ------------ Choice sites ------------ #

@inline function (ctx::ForwardModeContext)(call::typeof(trace),
                                           addr::A,
                                           d::Distribution{K}) where {A <: Address, K}
    visit!(ctx, addr)
    v = context_getindex(ctx, ctx.map, addr)
    ctx.weight += logpdf(d, v)
    v
end

# ------------ Learnable ------------ #

@inline function (ctx::ForwardModeContext)(fn::typeof(learnable), addr::Address)
    visit!(ctx, addr)
    haskey(ctx.params, addr) && return context_getindex(ctx, ctx.params, addr)
    error("Parameter not provided at address $addr.")
end

# ------------ Call sites ------------ #

@inline function (ctx::ForwardModeContext)(c::typeof(trace),
                                           addr::A,
                                           call::Function,
                                           args...) where A <: Address
    visit!(ctx, addr)
    ps = get_sub(ctx.params, addr)
    ret, w = forward(ctx.target[2 : end], ps, get_sub(ctx.map, addr), Dual(1.0, 0.0))
    ctx.weight += w
    return ret
end

# ------------ Convenience ------------ #

function get_choice_gradient(addr::T, cl::C) where {T <: Tuple, C <: CallSite}
    fn = seed -> begin
        ret, w = forward(addr, Empty(), cl, seed)
        w
    end
    d = fn(Dual(1.0, 0.0))
    cl[addr], d.partials.values[1]
end

function get_choice_gradient(addr::T, ps::P, cl::C) where {P <: AddressMap, T <: Tuple, C <: CallSite}
    fn = seed -> begin
        ret, w = forward(addr, ps, cl, seed)
        w
    end
    d = fn(Dual(1.0, 0.0))
    cl[addr], d.partials.values[1]
end

function get_learnable_gradient(addr::T, ps::P, cl::C) where {P <: AddressMap, T <: Tuple, C <: CallSite}
    fn = seed -> begin
        ret, w = forward(addr, ps, cl, seed)
        w
    end
    d = fn(Dual(1.0, 0.0))
    ps[addr], d.partials.values[1]
end

function get_choice_gradient(addr::T, cl::DynamicCallSite) where T <: Tuple
    fn = seed -> begin
        ret, w = forward(addr, Empty(), cl, seed)
        w
    end
    d = fn(Dual(1.0, 0.0))
    cl[addr], d.partials.values[1]
end

function get_choice_gradient(addr::T, ps::P, cl::DynamicCallSite) where {P <: AddressMap, T <: Tuple}
    fn = seed -> begin
        ret, w = forward(addr, ps, cl, seed)
        w
    end
    d = fn(Dual(1.0, 0.0))
    cl[addr], d.partials.values[1]
end

function get_learnable_gradient(addr::T, ps::P, cl::DynamicCallSite) where {P <: AddressMap, T <: Tuple}
    fn = seed -> begin
        ret, w = forward(addr, ps, cl, seed)
        w
    end
    d = fn(Dual(1.0, 0.0))
    ps[addr], d.partials.values[1]
end

# ------------ Documentation ------------ #