"""
Given the number of rows and columns, return an NTuple{4,Float64} containing
`(width, height, vspace, hspace)`, where `width` and `height` are the
width and height of each subplot and `vspace` and `hspace` are the vertical
and horizonal spacing between subplots, respectively.
"""
function sizes(nr::Int, nc::Int, subplot_titles::Bool=false)
    # NOTE: the logic of this function was mostly borrowed from plotly.py
    dx = 0.2 / nc
    dy = subplot_titles ? 0.55 / nr : 0.3 / nr
    width = (1. - dx * (nc - 1)) / nc
    height = (1. - dy * (nr - 1)) / nr
vspace = nr == 1 ? 0.0 : (1 - height * nr) / (nr - 1)
    hspace = nc == 1 ? 0.0 : (1 - width * nc) / (nc - 1)
    width, height, vspace, hspace
end

function gen_layout(rows::Tuple{Vararg{Int}}, subplot_titles::Bool=false)

    x = 0.0  # start from left
    y = 1.0  # start from top

    out = Layout()

    subplot = 1
    nr = length(rows)
    for nc in rows
        w, h, dy, dx = sizes(nr, nc, subplot_titles)
        x = 0.0 # reset x as we start a new row
        for col in 1:nc

            out["xaxis$subplot"] = attr(domain=[x, x + w], anchor="y$subplot")
            out["yaxis$subplot"] = attr(domain=[y - h, y], anchor="x$subplot")

            x += nc == 1 ? 0.0 : w + dx
            subplot += 1
         end
         y -= nr == 1 ? 0.0 : h + dy
    end

    out

end

gen_layout(nr::Int, nc::Int, subplot_titles::Bool=false) = gen_layout(tuple(fill(nc, nr)), subplot_titles)

function handle_titles!(big_layout, sub_layout, ix::Int)
    # don't worry about it if the sub_layout doesn't have a title
    if _isempty(sub_layout.title)
        return big_layout
    end

    # check for title or title text
    text = sub_layout.title
    if text isa Dict
        text = get(text, :text, missing)
        if ismissing(text)
            return big_layout
        end
    end

    ann = attr(
        font_size=16,
        showarrow=false,
        text=text,
        x=mean(big_layout["xaxis$(ix).domain"]),
        xanchor="center",
        xref="paper",
        y=big_layout["yaxis$(ix).domain"][2],
        yanchor="bottom",
        yref="paper"
    )
    anns = get(big_layout.fields, :annotations, [])
    push!(anns, ann)
    big_layout[:annotations] = anns
    big_layout
end

# plots are 3d if any of their traces have a 3d type. This should flow down
# the methods as ordered here
_is3d(p::Plot) = any(_is3d, p.data)
_is3d(t::GenericTrace) = _is3d(t[:type])
_is3d(t::Symbol) = _is3d(string(t))
_is3d(s::AbstractString) = s in ["surface", "mesh3d", "scatter3d"]

# else (maybe if trace didn't have a :type field set)
_is3d(x::Any) = false

function _cat(rows::Tuple{Vararg{Int}}, ps::Plot...)
    copied_plots = Plot[copy(p) for p in ps]
    subplot_titles = any(map(x -> haskey(x.layout.fields, :title) ||
                                  haskey(x.layout.fields, "title"), ps))
    layout = gen_layout(rows, subplot_titles)

    for ix in 1:sum(rows)
        handle_titles!(layout, copied_plots[ix].layout, ix)
        layout["xaxis$ix"] = merge(copied_plots[ix].layout["xaxis"], layout["xaxis$ix"])
        layout["yaxis$ix"] = merge(copied_plots[ix].layout["yaxis"], layout["yaxis$ix"])

        if _is3d(copied_plots[ix])
            # need to remove (x|y)axis$ix and move their domains into scene$ix here
            xaxis = pop!(layout, "xaxis$(ix)")
            yaxis = pop!(layout, "yaxis$(ix)")
            layout["scene$ix"] = attr(
                domain = attr(
                    x = xaxis[:domain],
                    y = yaxis[:domain]
                )
            )
            for trace in copied_plots[ix].data
                trace["scene"] = "scene$ix"
            end
        else
            for trace in copied_plots[ix].data
                trace["xaxis"] = "x$ix"
                trace["yaxis"] = "y$ix"
            end
        end

    end

    Plot(vcat([p.data for p in copied_plots]...), layout)
end

_cat(nr::Int, nc::Int, ps::Plot...) = _cat(Tuple(fill(nc, nr)), ps...)

Base.hcat(ps::Plot...) = _cat(1, length(ps), ps...)
Base.vcat(ps::Plot...) = _cat(length(ps), 1,  ps...)

Base.hvcat(rows::Tuple{Vararg{Int}}, ps::Plot...) = _cat(rows, ps...)

# implementation of algorithms from plotly.py/packages/python/plotly/plotly/subplots.py

function _init_subplot_xy!(
        layout::Layout, secondary_y::Bool, domain::NamedTuple{(:x, :y)}, max_subplot_ids::Dict
    )
    # Get axis label and anchor
    x_count = max_subplot_ids["xaxis"] + 1
    y_count = max_subplot_ids["xaxis"] + 1

    # Compute x/y labels (the values of trace.xaxis/trace.yaxis
    x_label = "x$(x_count > 1 ? x_count : "")"
    y_label = "y$(y_count > 1 ? y_count : "")"

    # Anchor x and y axes to each other
    x_anchor, y_anchor = y_label, x_label

    # Build layout.xaxis/layout.yaxis containers
    xaxis_name = "xaxis$(x_count > 1 ? x_count : "")"
    yaxis_name = "yaxis$(y_count > 1 ? y_count : "")"

    #  copy over existing xaxis/yaxis objects, but delete titles
    x_axis = attr(;layout.xaxis...)
    y_axis = attr(;layout.yaxis...)
    x_axis.domain = domain.x
    x_axis.anchor = x_anchor
    y_axis.domain = domain.y
    y_axis.anchor = y_anchor

    pop!(x_axis, :title_text)
    pop!(y_axis, :title_text)
    layout[xaxis_name] = x_axis
    layout[yaxis_name] = y_axis

    subplot_refs = [
        SubplotRef(
            subplot_kind="xy",
            layout_keys=Symbol.([xaxis_name, yaxis_name]),
            trace_kwargs=attr(xaxis=x_label, yaxis=y_label)
        )
    ]

    if secondary_y
        y_count += 1
        secondary_yaxis_name = "yaxis$(y_count > 1 ? y_count : "")"
        secondary_y_label = "y$(y_count > 1 ? y_count : "")"
        push!(
            subplot_refs,
            SubplotRef(
                subplot_kind="xy",
                layout_keys=Symbol.([xaxis_name, secondary_yaxis_name]),
                trace_kwargs=attr(xaxis=x_label, yaxis=secondary_y_label),
            )
        )
        secondary_y_axis = attr(anchor=y_anchor, overlaying=y_label, side="right")
        layout[secondary_yaxis_name] = secondary_y_axis
    end

    max_subplot_ids["xaxis"] = x_count
    max_subplot_ids["yaxis"] = y_count

    subplot_refs
end

function _init_subplot_single!(
        layout::Layout, subplot_kind::String, domain::NamedTuple{(:x, :y)}, max_subplot_ids::Dict
    )
    count = max_subplot_ids[subplot_kind] + 1
    label = "$(subplot_kind)$(count > 1 ? count : "")"
scene = attr(domain=attr(;domain...))
    layout[label] = scene

    trace_key = subplot_kind in _subplot_prop_named_subplot ? "subplot" : subplot_kind
    subplot_ref = SubplotRef(
        subplot_kind=subplot_kind,
        layout_keys=[Symbol(label)],
        trace_kwargs=attr(;Dict(Symbol(trace_key) => label)...)
    )
    max_subplot_ids[subplot_kind] = count
    [subplot_ref]
end

function _init_subplot_domain!(domain::NamedTuple{(:x, :y)})
    [SubplotRef(
        subplot_kind="domain",
        layout_keys=[],
        trace_kwargs=attr(domain=attr(;domain...))
    )]
end

function get_subplotkind_from_trace_type(k::Symbol)
    schema = get_plotschema()

    trace_attrs = schema.traces[k][:attributes]
    trace_attr_names = keys(trace_attrs)
    if :xaxis in trace_attr_names && :yaxis in trace_attr_names
        return "xy"
    elseif :domain in trace_attr_names
        return "domain"
    elseif :geo in trace_attr_names
        return "geo"
    elseif :scene in trace_attr_names
        return "scene"
    elseif k === :splom
        return "xy"
    elseif :subplot in trace_attr_names
        return trace_attrs[:subplot][:dflt]
    else
        @error "Unknown subplot type for trace $k. Please open issue"
    end
end

function _init_subplot!(
        layout::Layout, subplot_kind::String, secondary_y::Bool,
        domain::NamedTuple{(:x, :y)}, max_subplot_ids::Dict
    )
    sksymbol = Symbol(subplot_kind)
    if sksymbol in _TRACE_TYPES
        subplot_kind = get_subplotkind_from_trace_type(sksymbol)
    end
    if subplot_kind == "xy"
        return _init_subplot_xy!(layout, secondary_y, domain, max_subplot_ids)
    elseif subplot_kind in _single_subplot_types
        return _init_subplot_single!(layout, subplot_kind, domain, max_subplot_ids)
    elseif subplot_kind == "domain"
        return _init_subplot_domain!(domain)
    end
    error("Unknown subplot_kind $subplot_kind")
end

function _configure_shared_axes!(layout::Layout, grid_ref::Matrix, specs::Matrix{<:_Maybe{Spec}}, x_or_y::String, shared::Union{Bool,String}, row_dir::Int)
    rows, cols = size(grid_ref)

    layout_key_ind = x_or_y == "x" ? 1 : 2

    rows_iter = collect(1:rows)
    row_dir < 0 && reverse!(rows_iter)

    function update_axis_matches!(
            first_axis_id, subplot_ref::SubplotRef, spec::Spec, remove_label::Bool
        )
        span = x_or_y == "x" ? spec.colspan : spec.rowspan

        if subplot_ref.subplot_kind == "xy" && span == 1
            if ismissing(first_axis_id)
                first_axis_name = subplot_ref.layout_keys[layout_key_ind]
                first_axis_id = replace(String(first_axis_name), "axis" => "")
            else
                axis_name = subplot_ref.layout_keys[layout_key_ind]
                axis_to_match = layout[axis_name]
                axis_to_match[:matches] = first_axis_id
                if remove_label
                    axis_to_match[:showticklabels] = false
                end
            end
        end
        return first_axis_id
    end

    if shared == "columns" || (x_or_y == "x" && shared == true)
        for c in 1:cols
            first_axis_id = missing
            ok_to_remove_label = x_or_y == "x"
            for r in rows_iter
                subplot_ref = grid_ref[r, c][1]
                spec = specs[r, c]
                first_axis_id = update_axis_matches!(
                    first_axis_id, subplot_ref, spec, ok_to_remove_label
                )
            end
        end
    elseif shared == "rows" || (x_or_y == "y" && shared == true)
        for r in rows_iter
            first_axis_id = missing
            ok_to_remove_label = x_or_y == "y"
            for c in 1:cols
                subplot_ref = grid_ref[r, c][1]
                spec = specs[r, c]
                first_axis_id = update_axis_matches!(
                    first_axis_id, subplot_ref, spec, ok_to_remove_label
                )
            end
        end
    elseif shared == "all"
        first_axis_id = missing
        for c in 1:cols, (ri, r) in enumerate(rows_iter)
            subplot_ref = grid_ref[r, c][1]
            spec = specs[r, c]
            ok_to_remove_label = let
                if x_or_y == "y"
                    c > 0
                else
                    row_dir > 0 ? ri > 0 : r < rows
                end
            end
            first_axis_id = update_axis_matches!(
                first_axis_id, subplot_ref, spec, ok_to_remove_label
            )
        end
    end
end

_build_subplot_title_annotations(::Missing, list_of_domains; kw...) = []

function _build_subplot_title_annotations(
        subplot_titles::Array{<:_Maybe{String}},
        list_of_domains::Vector{<:NamedTuple{(:x, :y)}};
        title_edge="top", offset=0
    )
    # If shared_axes is false (default) use list_of_domains
    # This is used for insets and irregular layouts
    # if not shared_xaxes and not shared_yaxes:
    x_dom = [dom.x for dom in list_of_domains]
    y_dom = [dom.y for dom in list_of_domains]
    subtitle_pos_x = []
    subtitle_pos_y = []

    if title_edge == "top"
        text_angle = 0
        xanchor = "center"
        yanchor = "bottom"

        for x_domains in x_dom
            push!(subtitle_pos_x, sum(x_domains) / 2)
        end
        for y_domains in y_dom
            push!(subtitle_pos_y, y_domains[2])
        end
        yshift = offset
        xshift = 0
    elseif title_edge == "bottom"
        text_angle = 0
        xanchor = "center"
        yanchor = "top"

        for x_domains in x_dom
            push!(subtitle_pos_x, sum(x_domains) / 2)
        end
        for y_domains in y_dom
            push!(subtitle_pos_y, y_domains[1])
        end
        yshift = -offset
        xshift = 0
    elseif title_edge == "right"
        text_angle = 90
        xanchor = "left"
        yanchor = "middle"

        for x_domains in x_dom
            push!(subtitle_pos_x, x_domains[2])
        end
        for y_domains in y_dom
            push!(subtitle_pos_y, sum(y_domains) / 2.0)
        end
        yshift = 0
        xshift = offset
    elseif title_edge == "left"
        text_angle = -90
        xanchor = "right"
        yanchor = "middle"

        for x_domains in x_dom
            push!(subtitle_pos_x, x_domains[1])
        end
        for y_domains in y_dom
            push!(subtitle_pos_y, sum(y_domains) / 2.0)
        end
        yshift = 0
        xshift = -offset
    else
        error("title_edge must be one of [top, bottom, left, right]")
    end

    plot_titles = []
    for index in 1:length(subplot_titles)
        title = subplot_titles[index]
        should_continue = (
            ismissing(title) || title == false || length(title) == 0 ||
            index > length(subtitle_pos_y)
        )
        if should_continue
            continue
        end
        annotation = attr(
            y=subtitle_pos_y[index], yref="paper", yanchor=yanchor,
            x=subtitle_pos_x[index], xref="paper", xanchor=xanchor,
            text=title, font_size=16,
            showarrow=false,
        )
        xshift != 0 && setindex!(annotation, xshift, :xshift)
        yshift != 0 && setindex!(annotation, yshift, :yshift)
        text_angle != 0 && setindex!(annotation, text_angle, :textangle)
        push!(plot_titles, annotation)
    end
    plot_titles
end

function Layout(sp::Subplots; kw...)
    layout = Layout(;kw...)

    @unpack_Subplots sp

    row_dir = start_cell == "top-left" ? -1 : 1

    col_seq = collect(1:cols) .- 1
    row_seq = collect(1:rows) .- 1
    row_dir < 0 && reverse!(row_seq)

    grid = [
        (
            sum(_widths[1:c]) + c * horizontal_spacing,
            sum(_heights[1:r]) + r * vertical_spacing
        )
        for r in row_seq, c in col_seq
    ]
    domains_grid = _Maybe{NamedTuple{(:x, :y)}}[missing for _ in 1:rows, __ in 1:cols]
    list_of_domains = NamedTuple{(:x, :y)}[]
    max_subplot_ids = _get_initial_max_subplot_ids()

    for r in 1:rows, c in 1:cols
        spec = specs[r, c]
        if ismissing(spec)
            continue
        end
        c_spanned = c + spec.colspan - 1
        r_spanned = r + spec.rowspan - 1

        c_spanned > cols && error("Some colspan value is too large for this subplot grid")
        r_spanned > rows && error("Some rowspan value is too large for this subplot grid")

        # grid x c_spanned -> x_domain
        x_s = grid[r, c][1]
        x_e = grid[r, c_spanned][1] + _widths[c_spanned] - spec.r
        x_domain = (max(0.0, x_s), min(1.0, x_e))

        # grid x r_spanned x row_dir -> yaxis_domain
        y_domain = let
            if row_dir > 1
                (grid[r, c][2] + spec.b, grid[r_spanned, c][2] + _heights[r_spanned] - spec.t)
            else
                (grid[r_spanned, c][2], grid[r, c][2] + _heights[end - r + 1] - spec.t)
            end
        end
        domain = (x = x_domain, y = y_domain)
        domains_grid[r, c] = domain

        push!(list_of_domains, domain)

        grid_ref[r, c] = _init_subplot!(layout, spec.kind, spec.secondary_y, domain, max_subplot_ids)

    end

    _configure_shared_axes!(layout, grid_ref, specs, "x", shared_xaxes, row_dir)
    _configure_shared_axes!(layout, grid_ref, specs, "y", shared_yaxes, row_dir)

    # handle insets

    insets_ref = ismissing(insets) ? missing : Any[missing for _ in 1:length(insets)]
    if !ismissing(insets)
        for (i_inset, inset) in enumerate(insets)
            r = inset.cell[1]
            c = inset.cell[2]

            0 <= r <= rows || error("Some `cell` out of range")
            0 <= c <= cols || error("Some `cell` out of range")

            x_s = grid[r, c][1] + inset.l * _widths[c]
            x_e = inset.w == "to_end" ? grid[r, c][1] + _widths[c] : (x_s + inset.w * _widths[c])
            x_domain = (x_s, x_e)

            y_s = grid[r, c][2] + inset.b * _heights[end - r + 1]
            y_e = inset.h == "to_end" ? (grid[r, c][2] + _heights[end - r + 1]) : (y_s + inset.h * _heights[end - r + 1])
            y_domain = (y_s, y_e)
            domain = (x = x_domain, y = y_domain)

            push!(list_of_domains, domain)

            insets_ref[i_inset] = _init_subplot!(
                layout, inset.kind, false, domain, max_subplot_ids
            )
        end
    end

    plot_title_annotations = _build_subplot_title_annotations(
        subplot_titles, list_of_domains
    )
    layout[:annotations] = plot_title_annotations

    if !ismissing(column_titles)
        domains_list = let
            if row_dir > 0
                [domains_grid[end, c] for c in 1:cols]
            else
                [domains_grid[1, c] for c in 1:cols]
            end
        end
        append!(
            layout[:annotations],
            _build_subplot_title_annotations(column_titles, domains_list)
        )
    end

    if !ismissing(row_titles)
        domains_list = [domains_grid[r, end] for r in 1:rows]
        append!(
            layout[:annotations],
            _build_subplot_title_annotations(row_titles, domains_list, title_edge="right")
        )
    end

    if !ismissing(x_title)
        domains_list = [(x = (0, max_width), y = (0, 1))]
        append!(
            layout[:annotations],
            _build_subplot_title_annotations(
                [x_title], domains_list, title_edge="bottom", offset=30
            )
        )
    end
    if !ismissing(y_title)
        domains_list = [(x = (0, 1), y = (0, 1))]
        append!(
            layout[:annotations],
            _build_subplot_title_annotations(
                [y_title], domains_list, title_edge="left", offset=40
            )
        )
    end
    layout.subplots = sp
    layout
end


using Test
using UnPack
using MOONS.Grids
using MOONS.Fields

@testset "Fields" begin
    FT = Float64
    a, b = FT(0.0), FT(1.0)
    n = 10
    β = FT(1.1)
    c1 = Coordinates(a,b,n; warpfun=Roberts_both, args=(β,))
    c2 = Coordinates(a,b,n; warpfun=Roberts_both, args=(β,))
    c3 = Coordinates(a,b,n; warpfun=Roberts_both, args=(β,))
    grid = Grid((c1,c2,c3))

    cent = CellCenter(grid)
    @test size(cent.data) == (n+2,n+2,n+2)

    corn = CellCorner(grid)
    @test size(corn.data) == (n+3,n+3,n+3)

    face = CellFace(grid, 1)
    @test size(face.data) == (n+3,n+2,n+2)
    face = CellFace(grid, 2)
    @test size(face.data) == (n+2,n+3,n+2)
    face = CellFace(grid, 3)
    @test size(face.data) == (n+2,n+2,n+3)

    edge = CellEdge(grid, 1)
    @test size(edge.data) == (n+2,n+3,n+3)
    edge = CellEdge(grid, 2)
    @test size(edge.data) == (n+3,n+2,n+3)
    edge = CellEdge(grid, 3)
    @test size(edge.data) == (n+3,n+3,n+2)

end

@testset "Fields iterator" begin

    FT = Float64
    a, b = FT(0.0), FT(1.0)
    n = 3
    β = FT(1.01)
    grid = Cube(a, b, n; warpfun=Roberts_both, args=(β,))

    # CellCenter field
    f = CellCenter(grid)
    fi = CellCenter(grid)
    @inbounds for local_grid in GridField(grid,fi)
        @unpack ijk,x,y,z = local_grid
        fi[ijk] = sin(3*π*x)* sin(3*π*y)* sin(3*π*z)
    end
    @inbounds begin
    for i in 1:grid.c[1].c.s, j in 1:grid.c[2].c.s, k in 1:grid.c[3].c.s
        x = grid.c[1].c.h[i]
        y = grid.c[2].c.h[j]
        z = grid.c[3].c.h[k]
        f[i,j,k] = sin(3*π*x)* sin(3*π*y)* sin(3*π*z)
    end; end
    @test all(fi .≈ f)

    # CellCorner field
    f = CellCorner(grid)
    fi = CellCorner(grid)
    @inbounds for local_grid in GridField(grid,fi)
        @unpack ijk,x,y,z = local_grid
        fi[ijk] = sin(3*π*x)* sin(3*π*y)* sin(3*π*z)
    end
    @inbounds begin
    for i in 1:grid.c[1].n.s, j in 1:grid.c[2].n.s, k in 1:grid.c[3].n.s
        x = grid.c[1].n.h[i]
        y = grid.c[2].n.h[j]
        z = grid.c[3].n.h[k]
        f[i,j,k] = sin(3*π*x)* sin(3*π*y)* sin(3*π*z)
    end; end
    @test all(fi .≈ f)

    # CellFace field
    for dir in 1:3
        f = CellFace(grid, dir)
        fi = CellFace(grid, dir)
        @inbounds for local_grid in GridField(grid,fi)
            @unpack ijk,x,y,z = local_grid
            fi[ijk] = sin(3*π*x)* sin(3*π*y)* sin(3*π*z)
        end
        nx = dir == 1 ? grid.c[1].n.s : grid.c[1].c.s
        ny = dir == 2 ? grid.c[2].n.s : grid.c[2].c.s
        nz = dir == 3 ? grid.c[3].n.s : grid.c[3].c.s
        @inbounds begin
        for i in 1:nx, j in 1:ny, k in 1:nz
            x = dir == 1 ? grid.c[1].n.h[i] : grid.c[1].c.h[i]
            y = dir == 2 ? grid.c[2].n.h[j] : grid.c[2].c.h[j]
            z = dir == 3 ? grid.c[3].n.h[k] : grid.c[3].c.h[k]
            f[i,j,k] = sin(3*π*x)* sin(3*π*y)* sin(3*π*z)
        end; end
        @test all(fi .≈ f)
    end

    # CellEdge field
    for dir in 1:3
        f = CellEdge(grid, dir)
        fi = CellEdge(grid, dir)
        @inbounds for local_grid in GridField(grid,fi)
            @unpack ijk,x,y,z = local_grid
            fi[ijk] = sin(3*π*x)* sin(3*π*y)* sin(3*π*z)
        end
        nx = dir == 1 ? grid.c[1].c.s : grid.c[1].n.s
        ny = dir == 2 ? grid.c[2].c.s : grid.c[2].n.s
        nz = dir == 3 ? grid.c[3].c.s : grid.c[3].n.s
        @inbounds begin
        for i in 1:nx, j in 1:ny, k in 1:nz
            x = dir == 1 ? grid.c[1].c.h[i] : grid.c[1].n.h[i]
            y = dir == 2 ? grid.c[2].c.h[j] : grid.c[2].n.h[j]
            z = dir == 3 ? grid.c[3].c.h[k] : grid.c[3].n.h[k]
            f[i,j,k] = sin(3*π*x)* sin(3*π*y)* sin(3*π*z)
        end; end
        @test all(fi .≈ f)
    end

end

include("interpolations.jl")

@testset "Interpolations - convergence" begin
    include("interpolations_convergence.jl")
end

include("extrapolations.jl")
include("discrete_ops.jl")


"""
    SetupPeriodicDomain(np::Integer=16)

Set up a periodic domain of size np x np

```
np=16 #domain size is np x np
Γ=SetPeriodicDomain(np)
```
"""
function SetupPeriodicDomain(np::Integer=16)
    γ,Γ=GridOfOnes("PeriodicDomain",1,np)
    Γ["XC"][1]=vec(0.5:1.0:np-0.5)*ones(1,np)
    Γ["XG"][1]=vec(0.0:1.0:np-1.0)*ones(1,np)
    Γ["YC"][1]=ones(np,1)*transpose(vec(0.5:1.0:np-0.5))
    Γ["YG"][1]=ones(np,1)*transpose(vec(0.0:1.0:np-1.0))
    return Γ
end

"""
    SetupRandomFlow(Γ::Dict)

Set up a random flow field over the domain specified by Γ

```
Γ=SetPeriodicDomain(16)
𝑃,ϕ=SetupRandomFlow(Γ)
```
"""
function SetupRandomFlow(Γ::Dict)
  (_,ϕ,_,_)=demo2(Γ);

  (u,v)=gradient(ϕ,Γ)
  u=u./Γ["DXC"]#normalization to grid units
  v=v./Γ["DYC"]

  (u,v)=exchange(u,v,1)
  u0=-v; u1=-v;
  v0=u; v1=u;

  𝑃 = Dict( "u0" => u0, "u1" => u1, "v0" => v0, "v1" => v1,
            "t0" => 0.0, "t1" => 400.0, "dt" => 0.1)
  𝑃=merge(𝑃,Γ)#add grid variables

  return 𝑃,ϕ
end


@testset "conic linear problems" begin
    @testset "conic_linear1_test" begin
    #=
    primal
        min -3x - 2y - 4z
    s.t.
        x +  y +  z == 3 :w_4
        y +  z == 2      :w_5
        x>=0 y>=0 z>=0   :w_1, w_2, w_3

    dual
        max 3w_4 + 2w_5
    s.t.
        w_1, w_2, w_3 \in Nonnegatives
        w_1 + w_4 == -3.0
        w_2 + w_4 + w_5 == -2.0
        w_3 + w_4 + w_5 == -4.0
    =#
        primal_model = conic_linear1_test()
        dual_model, primal_dual_map = dual_model_and_map(primal_model)

        @test MOI.get(dual_model, MOI.NumberOfVariables()) == 5
        list_of_cons =  MOI.get(dual_model, MOI.ListOfConstraints())
        @test Set(list_of_cons) == Set([
            (SAF{Float64}, MOI.EqualTo{Float64})
            (VVF, MOI.Nonnegatives)
        ])
        @test MOI.get(dual_model, MOI.NumberOfConstraints{SAF{Float64}, MOI.EqualTo{Float64}}()) == 3
        @test MOI.get(dual_model, MOI.NumberOfConstraints{VVF, MOI.Nonnegatives}()) == 1
        obj_type = MOI.get(dual_model, MOI.ObjectiveFunctionType())
        @test obj_type == SAF{Float64}
        @test MOI.get(dual_model, MOI.ObjectiveSense()) == MOI.MAX_SENSE
        obj = MOI.get(dual_model, MOI.ObjectiveFunction{obj_type}())
        @test MOI.constant(obj) == 0.0
        @test MOI.coefficient.(obj.terms) == [3.0; 2.0]

        eq_con1_fun = MOI.get(dual_model, MOI.ConstraintFunction(), CI{SAF{Float64}, MOI.EqualTo{Float64}}(2))
        eq_con1_set = MOI.get(dual_model, MOI.ConstraintSet(), CI{SAF{Float64}, MOI.EqualTo{Float64}}(2))
        @test MOI.coefficient.(eq_con1_fun.terms) == [1.0; 1.0]
        @test MOI.constant.(eq_con1_fun) == 0.0
        @test MOI.constant(eq_con1_set) == -3.0
        eq_con2_fun = MOI.get(dual_model, MOI.ConstraintFunction(), CI{SAF{Float64}, MOI.EqualTo{Float64}}(3))
        eq_con2_set = MOI.get(dual_model, MOI.ConstraintSet(), CI{SAF{Float64}, MOI.EqualTo{Float64}}(3))
        @test MOI.coefficient.(eq_con2_fun.terms) == [1.0; 1.0; 1.0]
        @test MOI.constant.(eq_con2_fun) == 0.0
        @test MOI.constant(eq_con2_set) == -2.0
        eq_con3_fun = MOI.get(dual_model, MOI.ConstraintFunction(), CI{SAF{Float64}, MOI.EqualTo{Float64}}(4))
        eq_con3_set = MOI.get(dual_model, MOI.ConstraintSet(), CI{SAF{Float64}, MOI.EqualTo{Float64}}(4))
        @test MOI.coefficient.(eq_con3_fun.terms) == [1.0; 1.0; 1.0]
        @test MOI.constant.(eq_con3_fun) == 0.0
        @test MOI.constant(eq_con3_set) == -4.0

        primal_con_dual_var = primal_dual_map.primal_con_dual_var
        @test primal_con_dual_var[CI{VAF{Float64}, MOI.Zeros}(2)] == VI.(4:5)
        @test primal_con_dual_var[CI{VVF, MOI.Nonnegatives}(1)] == VI.(1:3)

        primal_var_dual_con = primal_dual_map.primal_var_dual_con
        @test primal_var_dual_con[VI(1)] == CI{SAF{Float64}, MOI.EqualTo{Float64}}(2)
        @test primal_var_dual_con[VI(2)] == CI{SAF{Float64}, MOI.EqualTo{Float64}}(3)
        @test primal_var_dual_con[VI(3)] == CI{SAF{Float64}, MOI.EqualTo{Float64}}(4)
    end

    @testset "conic_linear3_test" begin
    #=
    primal
        min  3x + 2y - 4z + 0s
    s.t.
        x           -  s == -4    :w_4
            y            == -3    :w_5
        x      +  z      == 12    :w_6
        y <= 0 :w_3
        z >= 0 :w_2
        s zero :w_1

    dual
        max -4w_4 - 3w_5 + 12w_6
    s.t
        w_4 + w_6 == 3
        w_3 + w_5 == 2
        w_2 + w_6 == -4
        w_1 - w_4 == 0
        w_2 >= 0
        w_3 <= 0
        w_1 in Reals
    =#
        primal_model = conic_linear3_test()
        dual_model, primal_dual_map = dual_model_and_map(primal_model)

        @test MOI.get(dual_model, MOI.NumberOfVariables()) == 6
        list_of_cons =  MOI.get(dual_model, MOI.ListOfConstraints())
        @test Set(list_of_cons) == Set([
            (SAF{Float64}, MOI.EqualTo{Float64})
            (VVF, MOI.Nonnegatives)
            (VVF, MOI.Nonpositives)
        ])
        @test MOI.get(dual_model, MOI.NumberOfConstraints{SAF{Float64}, MOI.EqualTo{Float64}}()) == 4
        @test MOI.get(dual_model, MOI.NumberOfConstraints{VVF, MOI.Nonnegatives}()) == 1
        @test MOI.get(dual_model, MOI.NumberOfConstraints{VVF, MOI.Nonpositives}()) == 1
        obj_type = MOI.get(dual_model, MOI.ObjectiveFunctionType())
        @test obj_type == SAF{Float64}
        @test MOI.get(dual_model, MOI.ObjectiveSense()) == MOI.MAX_SENSE
        obj = MOI.get(dual_model, MOI.ObjectiveFunction{obj_type}())
        @test MOI.constant(obj) == 0.0
        @test MOI.coefficient.(obj.terms) == [-4.0; -3.0; 12.0]

        eq_con1_fun = MOI.get(dual_model, MOI.ConstraintFunction(), CI{SAF{Float64}, MOI.EqualTo{Float64}}(3))
        eq_con1_set = MOI.get(dual_model, MOI.ConstraintSet(), CI{SAF{Float64}, MOI.EqualTo{Float64}}(3))
        @test MOI.coefficient.(eq_con1_fun.terms) == [1.0; 1.0]
        @test MOI.constant.(eq_con1_fun) == 0.0
        @test MOI.constant(eq_con1_set) == 3.0
        eq_con2_fun = MOI.get(dual_model, MOI.ConstraintFunction(), CI{SAF{Float64}, MOI.EqualTo{Float64}}(4))
        eq_con2_set = MOI.get(dual_model, MOI.ConstraintSet(), CI{SAF{Float64}, MOI.EqualTo{Float64}}(4))
        @test MOI.coefficient.(eq_con2_fun.terms) == [1.0; 1.0]
        @test MOI.constant.(eq_con2_fun) == 0.0
        @test MOI.constant(eq_con2_set) == 2.0
        eq_con3_fun = MOI.get(dual_model, MOI.ConstraintFunction(), CI{SAF{Float64}, MOI.EqualTo{Float64}}(5))
        eq_con3_set = MOI.get(dual_model, MOI.ConstraintSet(), CI{SAF{Float64}, MOI.EqualTo{Float64}}(5))
        @test MOI.coefficient.(eq_con3_fun.terms) == [1.0; 1.0]
        @test MOI.constant.(eq_con3_fun) == 0.0
        @test MOI.constant(eq_con3_set) == -4.0
        eq_con6_fun = MOI.get(dual_model, MOI.ConstraintFunction(), CI{SAF{Float64}, MOI.EqualTo{Float64}}(6))
        eq_con6_set = MOI.get(dual_model, MOI.ConstraintSet(), CI{SAF{Float64}, MOI.EqualTo{Float64}}(6))
        @test MOI.coefficient.(eq_con6_fun.terms) == [1.0; -1.0]
        @test MOI.constant.(eq_con6_fun) == 0.0
        @test MOI.constant(eq_con6_set) == 0.0

        primal_con_dual_var = primal_dual_map.primal_con_dual_var
        @test primal_con_dual_var[CI{VVF, MOI.Zeros}(4)] == [VI(1)]
        @test primal_con_dual_var[CI{VVF, MOI.Nonnegatives}(3)] == [VI(2)]
        @test primal_con_dual_var[CI{VVF, MOI.Nonpositives}(2)] == [VI(3)]
        @test primal_con_dual_var[CI{VAF{Float64}, MOI.Zeros}(1)] == VI.(4:6)

        primal_var_dual_con = primal_dual_map.primal_var_dual_con
        @test primal_var_dual_con[VI(1)] == CI{SAF{Float64}, MOI.EqualTo{Float64}}(3)
        @test primal_var_dual_con[VI(2)] == CI{SAF{Float64}, MOI.EqualTo{Float64}}(4)
        @test primal_var_dual_con[VI(3)] == CI{SAF{Float64}, MOI.EqualTo{Float64}}(5)
        @test primal_var_dual_con[VI(4)] == CI{SAF{Float64}, MOI.EqualTo{Float64}}(6)
    end
end


using Plots

include("rw_1d.jl")

nsteps = length(ARGS)>0 ? parse(Int, ARGS[1]) : 1000
prob = length(ARGS)>1 ? parse(Float64, ARGS[1]) : 0.5

result = randomwalk(nsteps,prob)

pl = plot(result, xaxis=("t"), yaxis=("x"), color=[:black], legend=false)
savefig(pl, "rw_1d.pdf")


export random_sequence,
       all_sequences

#################
### Functions ###
#################

"""
    random_sequence(automaton::Automaton, N::Integer)

The function takes an arbitrary walk of length `N` in `automaton`.
Returns a sequence that satisfy all weakly-hard constraints used to build the automaton.
"""
function random_sequence(automaton::Automaton{T}, N::S) where {T <: Integer, S <: Integer}

    # Initialise start vertex, character vector, and sequence
    IntType = BigInt
    if N < 8
        IntType = Int8
    elseif N < 16
        IntType = Int16
    elseif N < 32
        IntType = Int32
    elseif N < 64
        IntType = Int64
    end
    seq             = IntType(0)
    cmp_word        = IntType(1)
    current_vertex  = automaton.head

    # Random walk
    for _ in 1:N
        seq <<= T(1)
        seq |= current_vertex.w & cmp_word
        current_vertex = _rand_neighbour(current_vertex)
    end # for

    return seq
end # function

"""
    all_sequences(automaton::Automaton, N::Integer)

The function returns a set containing all sequences of length `N` satisfying the constraints used to build the automaton.
In other words, the function generates the satisfaction set of length `N` sequences.
"""
function all_sequences(automaton::Automaton{T}, N::S) where {T <: Integer, S <: Integer}
    # @description
    #   The function returns a set containing all sequences of length N
    #   satisfying the constraints used to build the automaton.
    # @param
    #   automaton::Automaton:     The weakly-hard automaton
    #   N::Integer:               The length of the sequence to be generated
    # @returns
    #   Set{<: Integer}: The resulting set of words (of length N) satisfying the
    #   weakly-hard automaton constraints.

    IntType = BigInt
    if N < 8
        IntType = Int8
    elseif N < 16
        IntType = Int16
    elseif N < 32
        IntType = Int32
    elseif N < 64
        IntType = Int64
    end

    seq_set = Set{IntType}()
    _recursive_sequence!(automaton.head, seq_set, N, IntType(0))
    return seq_set
end # function

#####################
### Aux functions ###
#####################

function _rand_neighbour(vertex::WordVertex{T}) where {T <: Integer}
    #= Picks a random neighbour =#
    if _childexists(vertex, :miss)
        return getfield(vertex, rand([:miss, :hit]))
    else
        return vertex.hit
    end
end # function

function _recursive_sequence!(vertex::WordVertex{T},
                              seq_set::Set{R},
                              N::S,
                              seq::R) where {T <: Integer, R <: Integer, S <: Integer}
    if N == 0 # base case
        push!(seq_set, seq)
    else
        if _childexists(vertex, :miss)
            _recursive_sequence!(vertex.miss, seq_set, N-S(1), seq << R(1))
        end
        _recursive_sequence!(vertex.hit, seq_set, N-S(1), (seq << R(1)) | R(1))
    end
end # function


const USE_GPU = false # ACHTUNG -> uncomment PT loop !!
const USE_DIRECT = true

import TriangleMesh
using Printf, LoopVectorization, LinearAlgebra, SparseArrays

include("../CreateMeshFCFV.jl")
include("../VisuFCFV.jl")
include("../DiscretisationFCFV.jl")
if USE_GPU
    using CUDA
    include("ResidualEvalFCFV_PT_GPU.jl")
else
    include("ResidualEvalFCFV_PT.jl")
end

#--------------------------------------------------------------------#

function SetUpProblem!( mesh, T, Tdir, Tneu, se, a, b, c, d, alp, bet )
    # Evaluate T analytic on cell faces
    @avx for in=1:mesh.nf
        x        = mesh.xf[in]
        y        = mesh.yf[in]
        Tdir[in] = exp(alp*sin(a*x + c*y) + bet*cos(b*x + d*y))
        dTdx     = Tdir[in] * (a*alp*cos(a*x + c*y) - b*bet*sin(b*x + d*y))
        dTdy     = Tdir[in] * (alp*c*cos(a*x + c*y) - bet*d*sin(b*x + d*y))
        Tneu[in] = -dTdy # nx*dTdx + nt*dTdy on SOUTH face
    end
    # Evaluate T analytic on barycentres
    @avx for iel=1:mesh.nel
        x       = mesh.xc[iel]
        y       = mesh.yc[iel]
        T       = exp(alp*sin(a*x + c*y) + bet*cos(b*x + d*y))
        T[iel]  = T
        se[iel] = T*(-a*alp*cos(a*x + c*y) + b*bet*sin(b*x + d*y))*(a*alp*cos(a*x + c*y) - b*bet*sin(b*x + d*y)) + T*(a^2*alp*sin(a*x + c*y) + b^2*bet*cos(b*x + d*y)) + T*(-alp*c*cos(a*x + c*y) + bet*d*sin(b*x + d*y))*(alp*c*cos(a*x + c*y) - bet*d*sin(b*x + d*y)) + T*(alp*c^2*sin(a*x + c*y) + bet*d^2*cos(b*x + d*y))
    end
    return
end

#--------------------------------------------------------------------#

function ComputeError( mesh, Te, qx, qy, a, b, c, d, alp, bet )
    eT  = zeros(mesh.nel)
    eqx = zeros(mesh.nel)
    eqy = zeros(mesh.nel)
    Ta  = zeros(mesh.nel)
    qxa = zeros(mesh.nel)
    qya = zeros(mesh.nel)
    @avx for iel=1:mesh.nel
        x        = mesh.xc[iel]
        y        = mesh.yc[iel]
        Ta[iel]       = exp(alp*sin(a*x + c*y) + bet*cos(b*x + d*y))
        qxa[iel]      = -Ta[iel] * (a*alp*cos(a*x + c*y) - b*bet*sin(b*x + d*y))
        qya[iel]      = -Ta[iel] * (alp*c*cos(a*x + c*y) - bet*d*sin(b*x + d*y))
        eT[iel]  = Te[iel] - Ta[iel]
        eqx[iel] = qx[iel] - qxa[iel]
        eqy[iel] = qy[iel] - qya[iel]
    end
    errT  = norm(eT)/norm(Ta)
    errqx = norm(eqx)/norm(qxa)
    errqy = norm(eqy)/norm(qya)
    return errT, errqx, errqy
end

#--------------------------------------------------------------------#

function StabParam(tau, dA, Vol, mesh_type) # ACHTUNG does not work on GPU (called from woithin kernel !)
    if mesh_type=="Quadrangles";        taui = tau;    end
    # if mesh_type=="UnstructTriangles";  taui = tau*dA; end
    if mesh_type=="UnstructTriangles";  taui = tau end
    return taui
end

#--------------------------------------------------------------------#

@views function main()

    println("\n******** FCFV POISSON ********")
    # Create sides of mesh
    xmin, xmax = 0, 1
    ymin, ymax = 0, 1
    nx, ny     = 20, 20
    R          = 0.5
    inclusion  = 0
    # mesh_type  = "Quadrangles"
    mesh_type  = "UnstructTriangles"

    # Generate mesh
    if mesh_type=="Quadrangles"
        tau  = 1
        mesh = MakeQuadMesh( nx, ny, xmin, xmax, ymin, ymax, inclusion, R )
    elseif mesh_type=="UnstructTriangles"
        tau  = 100
        mesh = MakeTriangleMesh( nx, ny, xmin, xmax, ymin, ymax, inclusion, R )
    end
    println("Number of elements: ", mesh.nel)

    # Source term and BCs etc...
    Tanal  = zeros(mesh.nel)
    se     = zeros(mesh.nel)
    Tdir   = zeros(mesh.nf)
    Tneu   = zeros(mesh.nf)
    alp = 0.1; bet = 0.3; a = 5.1; b = 4.3; c = -6.2; d = 3.4;
    println("Model configuration :")
    @time SetUpProblem!(mesh , Tanal, Tdir, Tneu, se, a, b, c, d, alp, bet)

    # Compute some mesh vectors
    println("Compute FCFV vectors:")
    @time ae, be, ze = ComputeFCFV(mesh, se, Tdir, tau)

    if USE_DIRECT
        # Assemble element matrices and RHS
        println("Compute element matrices:")
        @time Kv, fv = ElementAssemblyLoop(mesh, ae, be, ze, Tdir, Tneu, tau)

        # Assemble triplets and sparse
        println("Assemble triplets and sparse:")
        @time K, f = CreateTripletsSparse(mesh, Kv, fv)
        # display(UnicodePlots.spy(K))

        # Solve for hybrid variable
        println("Direct solve:")
        # @time Th = K\f
        PC  = 0.5 .* (K .+ transpose(K))
        PCc = cholesky(PC)
        Th  = zeros(mesh.nf)
        dTh = zeros(mesh.nf,1)
        r   = zeros(mesh.nf,1)
        r  .= f - K*Th
        for rit=1:5
            r    .= f - K*Th
            println("It. ", rit, " - Norm of residual: ", norm(r)/length(r))
            if norm(r)/length(r) < 1e-10
                break
            end
            dTh  .= PCc\r
            Th  .+= dTh[:]
        end
        # Compute residual on faces -  This is check
        @time Te, qx, qy = ComputeElementValues(mesh, Th, ae, be, ze, Tdir, tau)
        @time F = ResidualOnFaces(mesh, Th, Te, qx, qy, tau)
        # println(F)
        println("Norm of matrix-free residual: ", norm(F)/length(F))
        Te2 = copy(Te)
    end

    if USE_GPU
        @show maxthreads = mesh.nf
        BLOC = 256
        GRID = Int(ceil(maxthreads/BLOC))
        @show cuthreads = BLOC
        @show cublocks  = GRID

        Th    = CUDA.zeros(Float64, mesh.nf)
        Te    = CUDA.zeros(Float64, mesh.nel)
        qx    = CUDA.zeros(Float64, mesh.nel)
        qy    = CUDA.zeros(Float64, mesh.nel)
        F     = CUDA.zeros(Float64, mesh.nf)
        F0    = CUDA.zeros(Float64, mesh.nf)
        Th_PT = CUDA.zeros(Float64, mesh.nf)
        ae    = CuArray(ae)
        be    = CuArray(be)
        ze    = CuArray(ze)
        Tneu  = CuArray(Tneu)

        mesh_nf     = mesh.nf
        mesh_type   = mesh.type
        mesh_nf_el  = mesh.nf_el
        Mesh_bc     = CuArray(mesh.bc)
        Mesh_f2e    = CuArray(mesh.f2e)
        Mesh_dA_f   = CuArray(mesh.dA_f)
        Mesh_n_x_f  = CuArray(mesh.n_x_f)
        Mesh_n_y_f  = CuArray(mesh.n_y_f)
        Mesh_vole_f = CuArray(mesh.vole_f)
        Mesh_vole   = CuArray(mesh.vole)
        Mesh_e2f    = CuArray(mesh.e2f)
        Mesh_dA     = CuArray(mesh.dA)
        Mesh_n_x    = CuArray(mesh.n_x)
        Mesh_n_y    = CuArray(mesh.n_y)
    else
        Th    = zeros(mesh.nf)
        Te    = zeros(mesh.nel)
        qx    = zeros(mesh.nel)
        qy    = zeros(mesh.nel)
        F     = zeros(mesh.nf)
        F0    = zeros(mesh.nf)
        Th_PT = zeros(mesh.nf)
    end

    # Now a PT solve
    dmp     = 0.8
    dTdtau  = 0.05
    nout    = 1e3
    iterMax = 1e5
    ϵ_PT    = 1e-7

    # PT loop
    @time for iter=1:iterMax
         if USE_GPU
              # @cuda blocks=cublocks threads=cuthreads ResidualOnFaces_v2_GPU!(F, Mesh_bc, Mesh_f2e, Mesh_dA_f, Mesh_n_x_f, Mesh_n_y_f, Mesh_vole_f, Mesh_vole, Mesh_e2f, Mesh_dA, Mesh_n_x, Mesh_n_y, Th_PT, Te, Tneu, qx, qy, ae, be, ze, tau, mesh_nf, mesh_nf_el) #mesh_type not ok because string
              # synchronize()
              # @cuda blocks=cublocks threads=cuthreads Update_F_GPU!(F, Th_PT, F0, dTdtau, dmp, mesh_nf)
              # synchronize()
         else
            ResidualOnFaces_v2!(F, mesh, Th_PT, Te, qx, qy, ae, be, ze, tau, Tneu)
            F     .= (1 - dmp).*F0 .+ F
            Th_PT .= Th_PT .+ dTdtau.*F
            F0    .= F
         end
        if iter % nout == 0
            err = norm(F)/length(F)
            println("PT Iter. ", iter, " --- Norm of matrix-free residual: ", err)
            if err < ϵ_PT
                print("PT solve converged in")
                break
            elseif isnan(err)
                error("NaN !")
            end
        end
    end

    Th .= Th_PT

    # Reconstruct element values
    println("Compute element values:")
    @time Te, qx, qy = ComputeElementValues(mesh, Array(Th), Array(ae), Array(be), Array(ze), Array(Tdir), tau)
    @time Te1, qx1, qy1 = ComputeElementValuesFaces(mesh, Array(Th), Array(ae), Array(be), Array(ze), Array(Tdir), tau)
    if USE_DIRECT
        println(norm(Te.-Te2)/length(Te[:]))
    end
    println(norm(Te.-Te1)/length(Te[:]))
    println(norm(qx.-qx1)/length(qx[:]))
    println(norm(qy.-qy1)/length(qy[:]))

    # Compute discretisation errors
    err_T, err_qx, err_qy = ComputeError( mesh, Te, qx, qy, a, b, c, d, alp, bet )
    println("Error in T:  ", err_T )
    println("Error in qx: ", err_qx)
    println("Error in qy: ", err_qy)

    # Visualise
    print("Visualisation:")
    if USE_DIRECT
        @time PlotMakie(mesh, Err, xmin, xmax, ymin, ymax; cmap = :batlow)
    else
        @time PlotMakie(mesh, Te1, xmin, xmax, ymin, ymax; cmap = :batlow)
        # @time PlotMakie(mesh, Te1, xmin, xmax, ymin, ymax; cmap = :batlow, min_v = 0.9, max_v = 1.2)
    end
    print("Done! Total runtime:")
    return
end

@time main()


#!/usr/bin/env julia

using Luxor

using Test

using Random
Random.seed!(42)

function test_epitrochoid(fname)
    Drawing(1400, 1400, fname)
    origin()
    background("grey20")
    sethue("antiquewhite")
    setopacity(0.8)
    tiles = Tiler(1200, 1200, 6, 6)
    for (pos, n) in tiles
        gsave()
        setline(.5)
        translate(pos)

        R = 20; r = n; d = 20
        sethue("antiquewhite")
        poly(epitrochoid(R, n, d, :stroke, stepby=pi/300, vertices=true), :stroke, close=true)
        text(string("R = $(R), r = $(r), d = $(d)"), halign=:center, O.x, O.y+tiles.tileheight/2)

        # test perimeterlength
        ptslist = epitrochoid(R, n, d, :stroke, stepby=pi/300, vertices=true)
        polyperimeter(ptslist)

        R = 20; r = 12; d = 5
        sethue("orange")
        (R == r) && (r += 1)
        p = epitrochoid(R, r, d, :none, stepby=pi/305, vertices=true)
        polydistances(p)
        text(string("R = $(R), r = $(r), d = $(d)"), halign=:center, O.x, O.y+tiles.tileheight/2 + 10)
        poly(p, :stroke, close=true)

        # testing portions of polys
        sethue("magenta")
        portion = rand()
        poly(polyportion(p, portion), :stroke, close=false)
        text(string("portion ", round(portion, digits=2)), halign=:center, O.x, O.y+tiles.tileheight/2 + 30)

        sethue("cyan")
        setline(0.3)
        poly(offsetpoly(p, 5), :stroke, close=true)
        text("offset 15", halign=:center, O.x, O.y+tiles.tileheight/2 + 20)

        grestore()
    end
    @test finish() == true
end

Random.seed!(42)

fname = "epitrochoid-test.pdf"
test_epitrochoid(fname)
println("...finished test: output in $(fname)")


# ----------------------------------------------------
# --- Import modules
# ----------------------------------------------------
using DigitalComm
using Test
# ----------------------------------------------------
# --- Tests
# ----------------------------------------------------
# Note --> The mapping system is described in bitMapping.jl

println("Tests for Max-Log soft  symbol demapper");
for alphabet = [4;16;64;256]
	@testset  "$alphabet-QAM" begin
		# Minimal Tx
		mcs	  = alphabet;		  # Constellation size
		n	  = Int(log2(mcs));		  # Bit per symbol
		N	  = 4096;	  # Numbver of test symbols
		bitSeq	= genBitSequence(N*n);
		qamSeq	= bitMappingQAM(mcs,bitSeq);
		# Assuming unitary channel
		channel = ones(Complex{Float64},length(qamSeq));
		# Calculating LLR
		llr		= symbolDemappingQAM(mcs,qamSeq,channel);
		hardB	= llrToHardBits(llr);
		# Checking size of decoding vector
		@test length(llr)	== N*n;
		@test length(hardB) == N*n;
		# Checking that bang and classic methods are the same
		buff  = zeros(Float64,N*n);
		symbolDemappingQAM!(buff,mcs,qamSeq,channel);
		buffB = zeros(UInt8,n*N);
		llrToHardBits!(buffB,buff);
		@test all(hardB .== buffB);
		# Checking that decoded method is initial method
		@test all( buffB .≈ bitSeq)
	end
end


function generate_autodiff_solver_fn(eval_obj,eval_cons,cs_eqs,cs_ineqs;hessian=true)
    function solver_fn(x)
        J = eval_obj(x)
        gJ = ForwardDiff.gradient(eval_obj, x)
        if hessian
            HJ = ForwardDiff.hessian(eval_obj, x)
        else
            HJ = zeros(length(x),length(x))
        end

        g = eval_cons(x)
        dgdx = ForwardDiff.jacobian(eval_cons, x)

        ceq = g[cs_eqs]
        c = g[cs_ineqs]
        gceq = dgdx[cs_eqs,:]
        gc = dgdx[cs_ineqs,:]

        J, ceq, c, gJ, gceq, gc, HJ
    end

    solver_fn
end

function generate_autodiff_solver_fn(eval_obj,fres,fcfg,eval_cons,gres,gcfg,cs_eqs,cs_ineqs;hessian=true)
    function solver_fn(x)
        ForwardDiff.hessian!(fres, eval_obj, x, fcfg)
        J = DiffResults.value(fres)
        gJ = DiffResults.gradient(fres)
        if hessian
            HJ = DiffResults.hessian(fres)
        else
            HJ = zeros(length(x),length(x))
        end

        ForwardDiff.jacobian!(gres, eval_cons, x, gcfg)
        g = DiffResults.value(gres)
        dgdx = DiffResults.jacobian(gres)

        ceq = g[cs_eqs]
        c = g[cs_ineqs]
        gceq = dgdx[cs_eqs,:]
        gc = dgdx[cs_ineqs,:]

        J, ceq, c, gJ, gceq, gc, HJ
    end

    solver_fn
end

function generate_autodiff_solver_fn(eval_obj,fres,eval_cons,gres,cs_eqs,cs_ineqs;hessian=true)
    function solver_fn(x)
        ForwardDiff.hessian!(fres, eval_obj, x)
        J = DiffResults.value(fres)
        gJ = DiffResults.gradient(fres)
        if hessian
            HJ = DiffResults.hessian(fres)
        else
            HJ = zeros(length(x),length(x))
        end

        ForwardDiff.jacobian!(gres, eval_cons, x)
        g = DiffResults.value(gres)
        dgdx = DiffResults.jacobian(gres)

        ceq = g[cs_eqs]
        c = g[cs_ineqs]
        gceq = dgdx[cs_eqs,:]
        gc = dgdx[cs_ineqs,:]

        J, ceq, c, gJ, gceq, gc, HJ
    end

    solver_fn
end

function generate_autodiff_solver_fn(eval_obj,eval_cons,cs_eqs,cs_ineqs,vs_num_vars;hessian=true)
    cs_num_cons = length(cs_eqs) + length(cs_ineqs)

    fres = DiffResults.HessianResult(zeros(vs_num_vars))
    fcfg = ForwardDiff.HessianConfig(eval_obj, fres, zeros(vs_num_vars))
    gres = DiffResults.JacobianResult(zeros(cs_num_cons), zeros(vs_num_vars))
    gcfg = ForwardDiff.JacobianConfig(eval_cons, zeros(vs_num_vars))

    generate_autodiff_solver_fn(eval_obj,fres,fcfg,eval_cons,gres,gcfg,cs_eqs,cs_ineqs,hessian=hessian)
end


function runExperimentAndWriteResults(kappaSquared::Array{Float64},h::Array{Float64,1},src::Array{Int64,1},n::Array{Int64,1},T_exact::Array{Float64},WU::Float64 )
Omega = zeros(2*length(n));
Omega[2:2:end] = (n.-1).*h;
Mesh = getRegularMesh(Omega,n.-1);
##########################################################

# if length(n)==2
	# matshow(sqrt(kappaSquared)); colorbar();
	# xlabel("y");
	# ylabel("x");
	# xticks(0:div(n[2],8):n[2],0:1:8);
	# yticks(0:div(n[1],4):n[1],0:1:4);
# end

pMem = getEikonalTempMemory(n);
pEik = getEikonalParam(Mesh,kappaSquared,src,false);
pEik.T1 = zeros(Float64,tuple(n...));

tictimes = 1000000;
times = div(tictimes,prod(n)) .+ 1;
println("Timing is based on averaged ",times," runs.");
tt = time_ns()
for k=1:times
	solveFastMarchingUpwindGrad(pEik,pMem);
end
t1 = (time_ns() - tt)/(10^9)
t1 = t1/times;

if Mesh.dim==2
	selfMultiplyWithAnalyticSolution2D(Mesh.n.+1,Mesh.h,src,pEik.T1);
else
	selfMultiplyWithAnalyticSolution3D(Mesh.n.+1,Mesh.h,src,pEik.T1)
end

T1 = copy(pEik.T1)

pEik.HO = true;
ERR1 = T_exact .- pEik.T1;

tt = time_ns()
for k=1:times
	solveFastMarchingUpwindGrad(pEik,pMem);
end
t2 = (time_ns() - tt)/(10^9)
t2 = t2/times;

if Mesh.dim==2
	selfMultiplyWithAnalyticSolution2D(Mesh.n.+1,Mesh.h,src,pEik.T1);
else
	selfMultiplyWithAnalyticSolution3D(Mesh.n.+1,Mesh.h,src,pEik.T1)
end

sqrtN = sqrt(prod(n));
T2 = copy(pEik.T1)

ERR2 = T_exact - pEik.T1;

@printf("[%3.4f,%3.4f]\t[%d,%d]\t[%3.2e,%3.2e]\t%3.2fs(%3.2f)\t",h[1],h[2],n[1],n[2],norm(ERR1,Inf),norm(ERR1,2)/sqrtN,t1,t1/WU);

@printf("[%3.2e,%3.2e]\t%3.2fs(%3.2f)\n",norm(ERR2,Inf),norm(ERR2,2)/sqrtN,t2,t2/WU);

# if length(n)==2
	# figure()
	# contour(T_exact[end:-1:1,:],20,linestyles = "-",colors = "black",linewidths=1.5)
	# contour(T1[end:-1:1,:],20,linestyles = ":",colors = "red", linewidths=1.5)
	# contour(T2[end:-1:1,:],20,linestyles = "--",colors = "blue", linewidths=1.5)
	# xlabel("y");
	# ylabel("x");
	# xticks(0:div(n[2],8):n[2],0:1:8);
	# yticks(0:div(n[1],4):n[1],4:-1:0);

	# matshow(abs(T1 - T_exact)); colorbar();
	# xlabel("y");
	# ylabel("x");
	# xticks(0:div(n[2],8):n[2],0:1:8);
	# yticks(0:div(n[1],4):n[1],0:1:4);

	# matshow(abs(T2 - T_exact)); colorbar();
	# xlabel("y");
	# ylabel("x");
	# xticks(0:div(n[2],8):n[2],0:1:8);
	# yticks(0:div(n[1],4):n[1],0:1:4);
# end

# if length(n)==2
	# figure()
	# contour(T_exact[end:-1:1,:],100,linestyles = "-",colors = "black",linewidths=2.0)
	# contour(T1[end:-1:1,:],100,linestyles = ":",colors = "red", linewidths=2.0)
	# contour(T2[end:-1:1,:],100,linestyles = "--",colors = "blue", linewidths=2.0)
	# xlabel("y");
	# ylabel("x");
	# xticks(0:div(n[2],16):n[2],0:0.5:8);
	# yticks(0:div(n[1],40):n[1],4:-0.1:0);

# end

return;
end


"""
`optimalPLD_solveSAO_LP_Benders( n, revenues, numPermutations, lambda, orderings, MIP_timeLimit, x_initial, A_feas, b_feas)`

Solve the integer version of the Bertsimas and Mišić first-choice PLD formulation using Benders decomposition.

This function will formulate the Benders reformulation in Gurobi by specifying a lazy constraint callback to generate the Benders cuts. Cuts are generated using the closed-form expressions for the Benders cuts when the x_i's are integer (see Section 4.1 of Bertsimas and Mišić).

`n` = number of candidate products.

`revenues` = array of n numbers representing marginal revenue/profit of each of the n products. (These are the pi_i's in the Bertsimas and Mišić paper.)

`numPermutations` = number of rankings / customer types.

`lambda` = array specifying the probability of each ranking / customer type; should be nonnegative and sum to 1.

`orderings` = 2D array of integers encoding each ranking. orderings[k,j] contains the index of ranking k's jth top product; (n+1) encodes the no-purchase option. For example, if n = 5 and

	orderings[k,:] = [3,5,2,6,1,4]

then the customer prefers 3 < 5 < 2 < 0 (no-purchase option).

`MIP_timeLimit` = maximum time limit for Gurobi to solve the integer problem. If set to 0, no time limit is imposed.

`x_initial` = array with with n values (between 0 and 1) indicating the initial values of the x_i variables to use to warm-start the Benders decomposition method. If empty, no warm-start solution will

`A_feas`, `b_feas` = 2D array and 1D array specifying the system of linear inequalities Ax <= b for the PLD problem. This can be used to impose side constraints (this is the Cx <= d constraint system in the formulation in Bertsimas and Mišić 2018). If b_feas is empty, then no constraints will be added (i.e., it will solve the unconstrained problem).

`rankingid_list`, `alpha_list`, `beta_list`, `gamma_list` = 1D, 2D, 2D and 1D arrays containing information on cuts generated in Benders LP solution process. These cuts are added to the Benders integer formulation so as to warm-start it.

Returns:

`x_val` = array of optimal x_i variables for the integer problem.

`expRevenue` = expected per-customer revenue/profit of the integer solution.

`upperBound` = best upper bound on the optimal per-customer revenue/profit.

`solutiontime` = time required to solve the problem.

`MIPgap` = optimality gap of the best integer solution found.

`total_num_lazy_checked` = number of lazy constraints that are checked (incremented by 1 for each customer type and each invocation of the lazy constraint callback)

`total_num_lazy_generated` = number of lazy constraints added to the formulation (incremented by 1 every time a lazy constraint for any customer type is found to be violated and is added to the formulation)

"""
function optimalPLD_solveSAO_Benders( n, revenues, numPermutations, lambda, orderings, MIP_timeLimit, isRelaxation, x_initial,
								A_feas, b_feas,
								rankingid_list, alpha_list, beta_list, gamma_list )

	maxRev = maximum(revenues);

	if (MIP_timeLimit > 0)
		m = Model( solver = GurobiSolver(TimeLimit = MIP_timeLimit))
	else
		m = Model( solver = GurobiSolver())
	end

	if (!isRelaxation)
		@variable(m, x[1:n], Bin)
	else
		@variable(m, 0 <= x[1:n] <= 1)
	end

	if (!isRelaxation)
		if (!isempty(x_initial))
			for i=1:n
				setvalue(x[i], x_initial[i]);
			end
		end
	end

	@variable(m, 0 <= subobj[1:numPermutations] <= maxRev)

	for t=1:length(b_feas)
		@constraint(m, sum( A_feas[t,i]*x[i] for i=1:n) <= b_feas[t])
	end

	for t=1:length(rankingid_list)
		k = rankingid_list[t];
		@constraint(m, subobj[k] <= gamma_list[t] + sum( alpha_list[t,i]*x[i] for i=1:n) + sum(beta_list[t,i]*(1-x[i]) for i=1:n) )
	end

	@objective(m, Max, sum( lambda[k] * subobj[k] for k=1:numPermutations))

	const total_num_lazy_generated = 0;
	const total_num_lazy_checked = 0;

	orderings = convert( Array{Int64}, orderings);
	p2r_mat = zeros(Int64, numPermutations, n+1);
	temp_rev = [revenues;0];
	revenues_ordered_mat = zeros(numPermutations, n+1);
	sorted_rev_inds_mat = zeros(Int64, numPermutations, n+1);

	for k in 1:numPermutations
		single_ordering = convert(Array{Int64}, orderings[k,:][:]);
		revenues_ordered =  temp_rev[single_ordering[:]]
		revenues_ordered_mat[k,:] = revenues_ordered
		sorted_rev_inds_mat[k,:] = sortperm(revenues_ordered[:], rev= true);
		p2r_mat[k,:] = sortperm(single_ordering[:]); # single_ordering is r2p
	end

	tol_integrality = 1e-4;
	# Compute current lower bound.
	function benders_cb(cb)
		x_val = getvalue(x)[:]

		sol_is_integral = all( i->( (i>1-tol_integrality) &  (i<tol_integrality)), x_val )

		# println("candSet = $candSet")
		subobj_val = getvalue(subobj)[:]

		candSet_ponly = find(x_val .> 0.5)
		candSet = [candSet_ponly; n+1];
		for k=1:numPermutations
			total_num_lazy_checked = total_num_lazy_checked + 1;
			#p2r = sortperm(orderings[k,:][:])
			gamma_val = 0;
			alpha_val = zeros(n);
			beta_val = zeros(n+1);
			#temp = indmin(p2r_mat[k,candSet[:]])
			temp = indmin(p2r_mat[k,candSet])
			if (candSet[temp] <= n)
				gamma_val = revenues[candSet[temp]];
			else
				gamma_val = 0.0
			end
			# println("\tgamma_val is ")
			# println(gamma_val)

			if (!isempty(candSet_ponly))
				maxrev = maximum(revenues[candSet_ponly[:]]);
				beta_val[candSet[temp]] = maxrev - gamma_val
			end
			# println("\tbeta_val[temp] is (temp = $temp) ")
			# println(beta_val[candSet[temp]])

			#alpha_val = zeros(size(alpha_val))
			for i2=1:(n+1)
				if (orderings[k,i2] <= n)
					alpha_val[orderings[k,i2]] = max(revenues[orderings[k,i2]] - gamma_val - sum( beta_val[orderings[k,1:(i2-1)]]), 0 )
				end
			end

			if (gamma_val < subobj_val[k])
				# println("\tperm $k: gamma_val = $gamma_val, subobj_val[k] = $(subobj_val[k])")
				@lazyconstraint(cb, gamma_val + sum( alpha_val[i] * x[i] for i=1:n) + sum( beta_val[i] * (1 - x[i]) for i=1:n) >= subobj[k])
				total_num_lazy_generated = total_num_lazy_generated+1;
			end
		end
	end

	addlazycallback(m, benders_cb)

	solve(m)

	println(getvalue(x))

	x_val = getvalue(x)[:]

	solutiontime = Gurobi.get_runtime(m.internalModel.inner);
	if (!isRelaxation)
		MIPgap = Gurobi.get_dblattr(m.internalModel.inner, "MIPGap");
	else
		MIPgap = -1;
	end

	expRevenue = getobjectivevalue(m)

	upperBound = Gurobi.get_dblattr(m.internalModel.inner, "ObjBound");

	return x_val, expRevenue, upperBound, solutiontime, MIPgap, total_num_lazy_checked, total_num_lazy_generated
end


module FiniteFields

# taken from https://github.com/kalmarek/RamanujanGraphs.jl/blob/master/src/gf.jl

using Primes
export GF, int, characteristic
export generator, issqrt

struct GF{q} <: Number
    value::Int

    function GF{q}(n, check = true) where {q}
        if check
            @assert q > 1
            @assert isprime(q)
        end
        return new{q}(mod(n, q))
    end
end

GF{q}(n::GF{q}) where {q} = GF{q}(Int(n), false)

Base.Int(n::GF) = Int(n.value)
characteristic(::Type{GF{q}}) where {q} = q
characteristic(::GF{q}) where {q} = q

Base.:(==)(n::GF{q}, m::GF{q}) where {q} = Int(n) == Int(m)
# hash(GF) == 0x04fd9e474909f8bf
Base.hash(n::GF{q}, h::UInt) where {q} =
    xor(0x04fd9e474909f8bf, hash(q, hash(Int(n), h)))

Base.:+(n::GF{q}, m::GF{q}) where {q} = GF{q}(Int(n) + Int(m), false)
Base.:-(n::GF{q}, m::GF{q}) where {q} = GF{q}(Int(n) - Int(m), false)
Base.:*(n::GF{q}, m::GF{q}) where {q} = GF{q}(Int(n) * Int(m), false)
Base.:*(n::Integer, m::GF{q}) where {q} = GF{q}(n * Int(m), false)
Base.:/(n::GF{q}, m::GF{q}) where {q} = n * inv(m)

Base.:-(n::GF{q}) where {q} = GF{q}(q - Int(n), false)
Base.inv(n::GF{q}) where {q} = GF{q}(invmod(Int(n), q), false)

function Base.:^(n::GF{q}, i::Integer) where {q}
    i < 0 && return inv(n)^-i
    return GF{q}(powermod(Int(n), i, q), false)
end

Base.zero(::Type{GF{q}}) where {q} = GF{q}(0, false)
Base.one(::Type{GF{q}}) where {q} = GF{q}(1, false)
Base.iszero(n::GF) = iszero(Int(n))
Base.isone(n::GF) = isone(Int(n))

Base.promote_rule(::Type{GF{q}}, ::Type{I}) where {q,I<:Integer} = GF{q}
Base.promote_rule(::Type{GF{p}}, ::Type{GF{q}}) where {p,q} = throw(
    DomainError(
        (GF{p}, GF{q}),
        "Cannot perform arithmetic on elements from different fields",
    ),
)

Base.div(n::GF{q}, m::Integer) where {q} = n / GF{q}(m, false)
Base.div(m::Integer, n::GF{q}) where {q} = GF{q}(m, false) / n

# taken from ValidatedNumerics, under under the MIT "Expat" License:
# https://github.com/JuliaIntervals/ValidatedNumerics.jl/blob/master/LICENSE.md
function subscriptify(n::Integer)
    subscript_0 = Int(0x2080) # Char(0x2080) -> subscript 0
    return join((Char(subscript_0 + d) for d in reverse(digits(n))), "")
end

Base.show(io::IO, n::GF{q}) where {q} = print(io, "$(Int(n))" * subscriptify(q))
Base.isless(n::GF{q}, m::GF{q}) where {q} = isless(Int(n), Int(m))

function legendresymbol(n, q)
    iszero(mod(n, q)) && return zero(n)
    isone(powermod(n, (q - 1) ÷ 2, q)) && return one(n)
    return -one(n)
end

function generator(::Type{GF{q}}) where {q}
    q == 2 && return one(GF{2})
    return rootofunity(GF{q}, q - 1)
end

function rootofunity(::Type{GF{q}}, ord::Integer) where {q}
    d, r = divrem(q - 1, ord)
    !iszero(r) &&
        throw(DomainError(GF{q}, "No root of unity of order $ord in $(GF{q})."))
    for i = 2:q-1
        g = GF{q}(i, false)
        acc = g
        for j = 2:ord-1
            acc *= g
            isone(acc) && @goto next_elt
        end
        isone(acc * g) && return g
        @label next_elt
    end
    return zero(GF{q}) # never hit, to keep compiler happy
end

Base.sqrt(n::GF{q}) where {q} = GF{q}(sqrtmod(Int(n), q), false)
issquare(n::GF{q}) where {q} = legendresymbol(Int(n), q) >= 0

function sqrtmod(n::Integer, q::Integer)
    l = legendresymbol(n, q)
    l == 0 && return zero(n)
    l == -1 && throw(DomainError(n, "$n is not a square modulo $q"))
    for i = 1:q # bruteforce loop
        y = powermod(i, 2, q)
        y == n && return oftype(n, i)
    end
    return zero(n) # never hit, to keep compiler happy
end

Base.iterate(::Type{GF{q}}, s = 0) where {q} =
    s >= q ? nothing : (GF{q}(s, false), s + 1)
Base.eltype(::Type{GF{q}}) where {q} = GF{q}
Base.size(gf::Type{<:GF}) = (characteristic(gf),)

# needed for basic linear algebra over GF{q}:
Base.abs(x::GF) = x
Base.conj(x::GF) = x

end # module


# `ql.jl` requires Turing, Distributions, Plots, and StatsPlots
using Turing, Distributions, Plots, StatsPlots

# definitions of type alias
Action = Int64
Reward = Real
Probability = Real
QValue = Real

# definition of the environment
struct Bandit
    p::Array{Probability, 1}
end

function step(bandit::Bandit, action::Action)
    p = bandit.p[action]
    if p >= rand()
        return 1.
    end
    return 0.
end

# definition of the model
mutable struct Agent
    q::Array{QValue, 1}
    alpha::Real
    beta::Real
end

function from_env(alpha::Real, beta::Real, env::Bandit)
    n = length(env.p)
    q = zeros(n)
    return Agent(q, alpha, beta)
end

function softmax(agent::Agent)
    qmax = maximum(agent.q)
    qexp = exp.(agent.beta * (agent.q .- qmax))
    return qexp / sum(qexp)
end

function choose_action(agent::Agent)
    p = softmax(agent)
    return rand(Categorical(p))
end

function update(agent::Agent, action::Action, reward::Reward)
    agent.q[action] += agent.alpha * (reward - agent.q[action])
end

# definition of steps in the simulation
function step(agent::Agent, env::Bandit)
    action = choose_action(agent)
    reward = step(env, action)
    update(agent, action, reward)
    return action, reward
end

function run(agent::Agent, env::Bandit, trial::Int64)
    result = Array{Tuple{Action, Reward}, 1}(undef, trial)
    for i in 1:trial
        ar = step(agent, env)
        result[i] = ar
    end
    return result
end

# run simulation and fit q-learning model to the simulated data
bandit = Bandit([0.1, 0.5])
agent = from_env(0.1, 2., bandit)
result = run(agent, bandit, 1000)

@model QLearning(env::Bandit, result::Array{Tuple{Action, Reward}}) = begin
    T = length(result)
    alpha ~ Beta(1, 1)
    beta ~ Gamma(1., 100.)
    agent = from_env(alpha, beta, env)

    for t in 1:(T - 1)
        a, r = result[t]
        p = softmax(agent)
        result[t][1] ~ Categorical(p)
        update(agent, a, r)
    end
end

chain = sample(QLearning(bandit, result), NUTS(), 500)
plot(chain)


using GaitSymmetry
using Test, InteractiveUtils

@testset "GaitSymmetry.jl" begin
    @testset "inverses" begin
        x = abs(randn()) + 1
        y = abs(randn()) + 1

        for F in subtypes(SymmetryFunc)
            s = F(x,y)
            @test F(inv(F)(s)...) ≈ s
        end

        @test (inv(Plo05)(0.6931471805599453; scaled=false) .≈ (1., 2.)) |> prod
        @test_throws DomainError inv(Plo05)(-10.)
    end

    @testset "conversions" begin
        x = abs(randn()) + 1
        y = abs(randn()) + 1

        for F1 in subtypes(SymmetryFunc), F2 in subtypes(SymmetryFunc)
            s1 = F1(x,y)
            # Plo05 and Allen2011 is non-directional, so any conversion to another
            # metric must only consider magnitude
            if F1 ∈ (Plo05,)
                # Sel86 needs inverses (not magnitude) to remove consideration of
                # directionality
                if F2 === Sel86
                    @test (convert(F2, F1)(s1) ≈ inv(F2(x,y)) ||
                           convert(F2, F1)(s1) ≈ F2(x,y))
                else
                    @test abs(convert(F2, F1)(s1)) ≈ abs(F2(x,y))
                end
            else
                @test convert(F2, F1)(s1) ≈ F2(x,y)
            end
        end
    end
end


# reimplements methods: hessfact, hessfact!
# and minor functions from LinAlg: full, A_mul_B! ..., reflectorApply!

import LinearAlgebra
import Base: A_mul_B!, Ac_mul_B!, A_mul_Bc!, copymutable, full
import LinearAlgebra: hessfact, hessfact!, Hessenberg, HessenbergQ
import LinearAlgebra: chkstride1, checksquare

import LinearAlgebra: reflector!, reflectorApply!
import LinearAlgebra.LAPACK: gehrd!

function Hessenberg{T<:BigFloatOrComplex}(A::StridedMatrix{T})
  Hessenberg(gehrd!(A)...)
end

hessfact!{T<:BigFloatOrComplex}(A::StridedMatrix{T}) = Hessenberg(A)
hessfact{T<:BigFloatOrComplex}(A::StridedMatrix{T}) = hessfact!(copy(A))

# Perfom Householder/Hessenberg
# The Householder matrices are unitary but not hermitian in the complex case
# to provide real subdiagonal of the Hessenberg matrix.
# mimic the corresponding LAPACK.gehrd! function.
function gehrd!{T<:BigFloatOrComplex}(ilo::Integer, ihi::Integer, A::StridedMatrix{T})
  chkstride1(A)
  n = checksquare(A)
  τ = zeros(T, max(0, n - 1))

  ilo = max(ilo, 1)
  ihi = min(ihi, n)

  for k = ilo:ihi-1
    x = view(A, k+1:ihi, k)
    τk = reflector!(x)
    τ[k] = τk
    reflectorApply!(x, τk, view(A, k+1:n, k+1:n))
    reflectorApply!(view(A, :, k+1:n), x, τk)
  end
  A, τ
end

"""

  `reflectorApply!(x::AbstractVector, τ::Number, A::StridedMatrix)`
  `reflectorApply!(A::StridedMatrix, x::AbstractVector, τ::Number)`

Apply reflector `x, τ` from left or right to matrix `A`, see `linalg/generic.jl`.
Matrix is modified by this operation.
"""
function reflectorApply!(A::StridedMatrix, x::AbstractVector, τ::Number)
  n, m = size(A)
  if length(x) != m
    throw(DimensionMismatch("reflector has length $(length(x)), which must match the second d    imension of matrix A, $m"))
  end
  @inbounds begin
  for j = 1:n
    # dot
    vAj = A[j,1]
    for i = 2:m
      vAj += x[i]*A[j,i]
    end

    vAj = τ*vAj

    # ger
    A[j, 1] -= vAj
    for i = 2:m
      A[j, i] -= x[i]'*vAj
    end
   end
  end
  return A
end

# various multiplications with HessenbergQ
# note the adjoint(τ[k]) - supports the case of complex τ!
function A_mul_B!{T<:BigFloatOrComplex}(HQ::HessenbergQ{T}, A::StridedVecOrMat{T})
  n = size(A, 1)
  τ = HQ.τ
  for k = length(τ):-1:1
    τk = τ[k]'
    if τk != 0
      x = view(HQ.factors, k+1:n, k)
      reflectorApply!(x, τk, view(A, k+1:n, :))
    end
  end
  A
end

function A_mul_B!{T<:BigFloatOrComplex}(A::StridedMatrix{T}, HQ::HessenbergQ{T})
  n = size(A, 2)
  τ = HQ.τ
  for k = 1:length(τ)
    τk = τ[k]
    if τk != 0
      x = view(HQ.factors, k+1:n, k)
      reflectorApply!(view(A, :, k+1:n), x, τk)
    end
  end
  A
end

function Ac_mul_B!{T<:BigFloatOrComplex}(HQ::HessenbergQ{T}, A::StridedVecOrMat{T})
  n = size(A, 1)
  τ = HQ.τ
  for k = 1:length(τ)
    τk = τ[k]
    if τk != 0
      x = view(HQ.factors, k+1:n, k)
      reflectorApply!(x, τk, view(A, k+1:n, :))
    end
  end
  A
end

function A_mul_Bc!{T<:BigFloatOrComplex}(A::StridedMatrix{T}, HQ::HessenbergQ{T})
  n = size(A, 2)
  τ = HQ.τ
  for k = length(τ):-1:1
    τk = τ[k]'
    if τk != 0
      x = view(HQ.factors, k+1:n, k)
      reflectorApply!(view(A, :, k+1:n), x, τk)
    end
  end
  A
end

function copymutable(HQ::HessenbergQ)
  Q = eye(HQ.factors)
  A_mul_B!(HQ, Q)
  Q
end

full{T<:BigFloatOrComplex,S<:AbstractMatrix}(HQ::HessenbergQ{T,S}) = copymutable(HQ)


include("sparsestcut.jl")
using Random

function flow_decomposition_with_fake_edges(F::SparseMatrixCSC,
     s::Int64, t::Int64, S::Vector{Int64}, flowvalue::Int64)
    N = size(F, 1) - 2
    S_bar = set_complement(N, S)
    matching_size = div(N, 2)

    M = zeros(Int, matching_size, 2)

    in_S = zeros(N + 2)
    in_S_bar = zeros(N + 2)
    in_S[S .+ 1] .= 1
    in_S_bar[S_bar .+ 1] .= 1

    Neighs, d = ConstructAdj(F, N + 2)
    Neighs_ptrs = ones(Int, N + 2)

    visited = zeros(N + 2)

    q = Queue{Pair{Int, Int}}()
    for _ = 1:length(S)
        enqueue!(q, s => s)
    end

    edges = 0
    while !isempty(q)
        x = dequeue!(q)
        u = x.first
        from = x.second
        uNeighs = Neighs[u]
        @assert Neighs_ptrs[u] <= d[u]

        v = uNeighs[Neighs_ptrs[u]]
        while F[u, v] <= 0
            Neighs_ptrs[u] += 1
            if Neighs_ptrs[u] > d[u]
                break
            end
            v = uNeighs[Neighs_ptrs[u]]
        end
        if Neighs_ptrs[u] > d[u]
            continue
        end
        if from == s
            @assert in_S[v] == 1
            from = v
        end
        F[u, v] -= 1
        if v == t
            edges += 1
            M[edges, 1] = from
            M[edges, 2] = u
            visited[from] = 1
            visited[u] = 1
        else
            enqueue!(q, v => from)
        end
    end
    @assert edges == flowvalue println(string(edges) * " " * string(matching_size))

    unvisited_S = []
    unvisited_S_bar = []
    for u = 2:N + 1
        if visited[u] == 0
            if in_S[u] == 1
                push!(unvisited_S, u)
            else
                push!(unvisited_S_bar, u)
            end
        end
    end
    shuffle(unvisited_S)
    shuffle(unvisited_S_bar)
    for i = 1:matching_size - flowvalue
        edges += 1
        M[edges, 1] = unvisited_S[i]
        M[edges, 2] = unvisited_S_bar[i]
    end
    return M
end

function flow_match_with_fake_edges(A::SparseMatrixCSC, S::Vector{Int}, beta::Float64)
    N = size(A, 1)
    target_flow = div(N, 2) - ceil(beta * N / log2(N) / log2(N))

    lb = 0
    ub = Int(ceil(log2(N)))
    while lb < ub
        mid = div(lb + ub, 2)
        B, s, t = build_graph(A, S, 2^mid)
        F = maxflow(B, s, t)
        if F.flowvalue >= target_flow
            ub = mid
        else
            lb = mid + 1
        end
    end
    B, s, t = build_graph(A, S, 2^ub)
    F = maxflow(B, s, t)
    M = flow_decomposition_with_fake_edges(F.F, s, t, S, Int(F.flowvalue)) .- 1

    power = max(0, ub - 1)
    B, s, t = build_graph(A, S, 2^power)
    F = maxflow(B, s, t)
    cut_S = source_nodes(F) .- 1
    cut_S = setdiff(cut_S, 0)
    return cut_S, M
end

function pseudo_balanced_cut_matching(A::SparseMatrixCSC, T::Int, b::Float64)
    Expender = expander(Vector{matching}())
    N = size(A, 1)
    min_expansion = N^2 # store the minimum expansion
    min_S = [] # store the set inducing the sparsest cut
    for i = 1:N
        S = [i]
        expansion = compute_expansion(A, S)
        if expansion < min_expansion
            min_expansion = expansion
            min_S = S
        end
    end
    for _ = 1:T
        r = gen_rand_unitv(N)
        bisec_S = find_bisection(Expender, r)
        beta = b / 2
        cut_S, M = flow_match_with_fake_edges(A, bisec_S, beta)
        if length(cut_S) > 0 && length(cut_S) < N
            expansion = compute_expansion(A, cut_S)
            if expansion < min_expansion
                min_expansion = expansion
                min_S = cut_S
            end
        end
        push!(Expender.matchings, matching(M))
    end
    return min_S, min_expansion
end


# Defines some commonly used states in quantum information

export bell, bellDiagState, isotropicState, rState, sState, sStateQutrit, rStatePhase, rStateCorrPhase, rStateCorrPhaseCopies;

""" `bell`

The Bell states (as vectors), bell[j,1:4] gives the j-th bell state.
"""
bell = zeros(4,4);
bell[1,1:4] = maxEntVec(2);
bell[2,1:4] = (kron(eVec(2,1),eVec(2,1)) -  kron(eVec(2,2),eVec(2,2)))/sqrt(2);
bell[3,1:4] = (kron(eVec(2,1),eVec(2,2)) +  kron(eVec(2,2),eVec(2,1)))/sqrt(2);
bell[4,1:4] = (kron(eVec(2,1),eVec(2,2)) -  kron(eVec(2,2),eVec(2,1)))/sqrt(2);

""" `epr` and `singlet`

Two of these have special names: we will use epr for the EPR pair, and singlet for the singlet.
"""
epr = bell[1,1:4];
singlet = bell[4,1:4];

""" `bellS`

An array of density matrices corresponding to the 4 Bell states.
"""
bS1 = epr*epr';
bS2 = bell[2,1:4]* bell[2,1:4]';
bS3 = bell[3,1:4]* bell[3,1:4]';
bS4 = bell[4,1:4]*bell[4,1:4]';
bellS = [bS1; bS2; bS3; bS4];

""" `PS`
Projector onto the symmetric subspace of 2 qubits.
"""
PS = bS1 + bS2 + bS3;

""" `PA`
Projector onto the antisymmetric subspace of 2 qubits.
"""
PA = bS4;

""" `rho = bellDiagState(p1,p2,p3)`

Returns a bell diagonal state that is a mixture between the 4 Bell states: p1 * phi^+ + p2 * psi^+ + p3 * phi^- + (1 - p1 - p2 - p3) * psi^-.
"""

function bellDiagState(p1::Number, p2::Number, p3::Number)

	@assert 0 <= p1 "Probilities must be positive."
	@assert 0 <= p2 "Probilities must be positive."
	@assert 0 <= p3 "Probilities must be positive."

	@assert p1 + p2 + p3 <= 1 "Probabilities cannot exceed 1."

	# Produce the desired mixture
	out = p1 * bS1 + p2 * bS3 + p3 * bS2 + (1-p1-p2-p3) * bS4;
	return out;
end

""" `rho = sState(p)`

Returns a state that is an a mixture between the EPR pair (with probability *p*), and the state |11><11|.
"""

function sState(p::Number)

	@assert 0 <= p "Probilities must be positive."
	@assert p <= 1 "Probabilities cannot exceed 1."

	# Generate the maximally entangled state
	epr = maxEnt(2);

	# Generate |11>
	e0 = [1 0];
	e1 = [0 1];
	v11 = kron(e1,e1);

	# Produce the desired mixture
	out = p * epr + (1-p) * v11'*v11;
	return out;
end

""" `rho = sStateQutrit(p)`

Returns a state that is a mixture between a 3 dimensional maximally entangled state (with probability *p*) and the state |00><00|
"""

function sStateQutrit(p::Number)

	@assert 0 <= p "Probilities must be positive."
	@assert p <= 1 "Probabilities cannot exceed 1."

	# Maximally entangled state
	epr = maxEnt(3);

	# Generate |00>
	e0 = [1 0 0];
	e1 = [0 0 1];
	v00 = kron(e0,e0);

	# Produce the desired mixture
	out = p * epr + (1-p) * v00'*v00;

  	return out;
end

""" `rho = rState(p)`

Returns a mixture between a state proportional to |01> + |10> (with probability *p*) and the state |11><11|.
"""

rState(p::Number) = rStatePhase(p, 0)

""" `rho = rStatePhase(p, phi)` or `rho = rStatePhase(p)`

Returns a mixture between a state proportional to |01> + e^(i *phi*) |10> (with probability *p*) and the state |11><11|. No value for *phi* defaults to *phi*=0.

"""

function rStatePhase(p::Number, phi::Number = 0.0)

	@assert 0 <= p "Probilities must be positive."
	@assert p <= 1 "Probabilities cannot exceed 1."

	# Produce the state |01> + e^(i phi) |10>
	e0 = eVec(2, 1);
	e1 = eVec(2, 2);
	vec = (kron(e0, e1) + e^(im*phi) * kron(e1,e0))/sqrt(2);

	if phi == 0 || phi == pi
		vec = real(vec);
	end

	# Produce a state orthogonal to the one above
	v11 = kron(e1,e1);

	# Construct the desire mixture
	out = p * vec*vec' + (1-p) * v11*v11';

  	return out
end

""" `rho = rStateCorrPhase(p, pd)`

Returns a state of the form integral phi [pd * r(p,phi) + (1-pd) * r(p,phi+pi)] tensor r(p,phi), where
r(p,phi) = rStatePhase(p,phi). Default for pd is 1.
"""
function rStateCorrPhase(p::Number, pd::Number = 1)

	@assert 0 <= p "Probilities must be positive."
	@assert p <= 1 "Probabilities cannot exceed 1."

    rStatePhaseDep(phi) = pd * rStatePhase(p, phi) + (1-pd) * rStatePhase(p, phi + pi)

	integrand(phi) = ( 1/(2*pi) ) * kron(rStatePhaseDep(phi), rStatePhase(p, phi));

	eps = 10.0^(-4)
	out = real( quadgk( integrand, 0, 2 * pi; reltol=sqrt(eps), abstol=0, maxevals=10^7, order=7, norm=vecnorm)[1])
  	return out
end

""" `rho = rStateCorrPhaseCopies(p)`

Returns a state of the form integral phi r(p,phi)^(tensor n) , where
r(p,phi) = rStatePhase(p,phi) and n is the number of desired copies.
"""
function rStateCorrPhaseCopies(p::Number, n::Int)

	@assert 0 <= p "Probilities must be positive."
	@assert p <= 1 "Probabilities cannot exceed 1."
	@assert n > 1 "Number of copies must be at least 1."

	integrand(phi) = ( 1/(2*pi) ) * copies(rStatePhase(p, phi), n);

	eps = 10.0^(-4)
	out = real( quadgk( integrand, 0, 2 * pi; reltol=sqrt(eps), abstol=0, maxevals=10^7, order=7, norm=vecnorm)[1])
  	return out
end

""" `rho = isotropicState(p)` or `rho = isotropicState(p,d)`

Returns an isotropic state, i.e., a mixture of a maximally entangled pair (with probability *p*) and the maximally mixed state in local dimension *d*. If no argument *d* is given the default is *d*=2, that is, we take the mixture of the EPR pair with the maximally mixed state of local dimension 2.
"""

function isotropicState(p::Number; d::Int = 2)

	@assert 0 <= p "Probilities must be positive."
	@assert p <= 1 "Probabilities cannot exceed 1."

        epr = maxEnt(d);
	out = p * epr*epr' + (1-p) * eye(d^2)/d^2;

	return out;
end