# The goal of this visualization is to see how `f(z) = eᶻ` can be
# approximated through `g(z) = (1 + z/n)ⁿ` as `n → ∞`. We will
# use a slider to change the value of `n` and also place a 'test set'
# in the domain to see how it gets transformed.
using Cairo
using ComplexVisual
@ComplexVisual.import_huge

const CVG = try using ComplexVisualGtk; ComplexVisualGtk; catch nothing; end

const fontface = cv_fontface("DejaVu Sans")

# We start by creating the function to draw the axes on the left and right
function setup_axis(setup::CV_SceneSetupChain)
    label_style = cv_black → fontface → cv_fontsize(20)
    # these are the rules used to label the x-axis ticks
                                # main ticks on left, get labeled
    ticks1, ticks1h, ticks2 = "%.0f" ⇒ -6:0,
                                # secondary ticks on left, don't get labeled
                                "" ⇒ -5.5:0.5,
                                # ticks on right, get labeled
                                "%.0f" ⇒ -3:3
    # rules used to label the y-axis ticks
    # uses a Tuple to label specific irrational points
    ticks3 = ("-π" ⇒ -pi, "-π/2" ⇒ -pi/2, "0" ⇒ 0.0, "π/2" ⇒ pi/2, "π" ⇒ pi,)

    # main ticks are drawn slightly longer than secondary ticks on left
    app_l = CV_TickLabelAppearance(; label_style, tick_length=10)
    app_s = CV_TickLabelAppearance(; label_style, tick_length=6)

    # use the standard axis builder from `lrdomains.jl`
    return cv_setup_lr_axis(setup,
        # x-axis labels                     # y-axis labels
        (app_l ↦ ticks1, app_s ↦ ticks1h), (app_l ↦ ticks3,),
        (app_l ↦ ticks2,),                  (app_l ↦ ticks3,)
    )
end

# Now we need to create labels + style for our custom slider. We wish
# to show more detail in the slider around the first few numbers,
# and have a slider position "at infinity"
function get_slider_rulers()  # {{{
    # points where the large ticks and medium ticks are drawn
    major_values, minor_values = Float64.(10:10:50), Float64.(15:5:45)
    # all leftover integers between 1 and 50, will be drawn with small ticks
    mini_values = setdiff(setdiff(Float64.(10:1:50), minor_values), major_values)
    # These values will all be labeled
    intro_values = Float64.(1:1:9)

    # label the major values between 10 and 50, and also ∞
    major_ticks = (("%.0f" ⇒  major_values)..., "∞" ⇒ 55.0)
    # minor values do not get labeled
    minor_ticks, mini_ticks = "" ⇒ minor_values, "" ⇒ mini_values
    intro_ticks = "%.0f" ⇒ intro_values

    # at this point the style is set, using `CV_TickLabelAppearance`
    major_ruler = CV_TickLabelAppearance(; tick_length=10,
            label_style=cv_black → fontface → cv_fontsize(20)) ↦ major_ticks
    minor_ruler = CV_TickLabelAppearance(; tick_length=7) ↦ minor_ticks
    mini_ruler = CV_TickLabelAppearance(; tick_length=4) ↦ mini_ticks
    intro_ruler = CV_TickLabelAppearance(; tick_length=4,
            label_style=cv_black → fontface → cv_fontsize(10)) ↦ intro_ticks
    return (major_ruler, minor_ruler, mini_ruler, intro_ruler)
end # }}}

# We have the labels, so we can build the interactive slider
function create_slider(setup, slider_pos)

    # setup will be a `CV_SceneSetupChain` which has a field `layout`.
    # For more info see `lrdomains.jl`
    layout = setup.layout

    # internal getter functions to access the layouts for the
    # left and right portraits. Referred to as the `seen rectangles`
    can_domain_l = cv_get_can_domain_l(layout)
    can_codomain_l = cv_get_can_codomain_l(layout)

    # will be used to position ('anchor') the slider,
    # based on the size of the portraits
    width = can_codomain_l.rectangle.right - can_domain_l.rectangle.left
    mid = cv_half(can_codomain_l.rectangle.right + can_domain_l.rectangle.left)
    bottom = min(can_domain_l.rectangle.bottom, can_codomain_l.rectangle.bottom)
    height = 15

    # now we create the basic slider object.
    # for more info, see `slider.jl`
    slider_data = cv_create_hslider(
        width, height, 0.5, 55.5, get_slider_rulers();
        decoraction_with_layout_and_position_callback=(inner_layout, pos) ->
            cv_border(inner_layout, pos, 1))

    # this is the bounding box ('container') that contains the drawing info
    cont_slider = slider_data.slider_container
    cc_cont_slider = slider_data.container_context

    # draw the slider in our setup.layout,
    # in the center (referring to x-position),
    # 20pts above the bottom of the left-right layout (y-position)
    cont_slider_l = cv_add_canvas!(layout,
        cont_slider,
        cv_anchor(cont_slider, :slider_south, cont_slider, :south),
        (mid, bottom - Int32(20)))

    # interactive part of the slider: if the user clicks
    # to the right of the max value, snap to the rightmost
    # slider position (labeled ∞)
    set_slider_value = z -> begin
        slider_pos.value = (z > 53) ? 55.0 : max(1.0, min(z, 50.0))
        return CV_Response(;redraw_flag=true)
    end

    # a few style decisions
    bg_painter = cv_color(.8,.8,.8) ↦ CV_FillPainter()
    mark_painter = (cv_op_source → cv_color(0,0,1) → cv_linewidth(2)) ↦
        CV_ValueMarkPainter(slider_pos,
            0.0, imag(cont_slider.can_slider.corner_ul), false)

    # finally the slider is added to the `CV_SceneSetupChain`
    setup = cv_setup_hslider(setup, slider_data, cont_slider_l,
        bg_painter → mark_painter, set_slider_value)
    return setup
end

# Putting everything together: almost identical to the
# 'standard' left-right setup process (see `cv_scene_lr_std` in 'lrdomains.jl')
function create_scene(trafo,
        domain, codomain, slider_pos; cut_test=nothing, gap=80,
        axis_label_style=cv_color(0,0,0) →
                  cv_fontface("sans-serif") → cv_fontsize(20), padding=30) # {{{
    # cyclic state counter counts right-clicks
    layout = CV_StateLayout(CV_2DLayout(), CV_CyclicValue(2))
    # creates the standard left-right layout
    layout = cv_do_lr_layout(cv_add(layout, trafo, domain, codomain), gap)

    # create a basic `CV_SceneSetupChain` with a default left-right setup
    setup = cv_setup_cycle_state(CV_LRSetupChain(layout))

    # place the slider in the setup chain
    setup = create_slider(setup, slider_pos)

    # standard left-right setup, using our custom axes
    # all the functions used here are in `lrdomains.jl`
    setup = cv_setup_lr_painters(setup; cut_test)
    setup = setup_axis(setup)
    setup = cv_setup_lr_border(setup)
    padding > 0 && cv_add_padding!(setup.layout, padding)
    setup = cv_setup_domain_codomain_scene(setup)
    cv_scene_lr_start(setup)
    return setup.layout
end # }}}

# `CV_TranslateByOffset(Float64)` lets the user interact with something that
# only needs to move left and right
slider_pos = CV_TranslateByOffset(Float64)
slider_pos.value = 1.0

# There is a discontinuity in our goal function due to our use of Log
# To get a better visualization, we add the cut test
cut_func = cv_create_angle_cross_test(π, 0.0, Inf; δ=1e-2)

cut_test = (z, w) -> (slider_pos.value == round(slider_pos.value)) ?
    false : cut_func(z + slider_pos.value, w + slider_pos.value)

trafo = z -> (slider_pos.value == 55.0) ? exp(z) :
    exp(slider_pos.value * log(1+z/slider_pos.value))
domain   = CV_Math2DCanvas(-6.5 + 3.5im, 0.5 - 3.5im, 100)
codomain = CV_Math2DCanvas(-3.5 + (pi+0.1)im, 3.5 - (pi+0.1)im, 100)

scene = create_scene(trafo, domain, codomain, slider_pos; cut_test)

# finally, run the scene!
if CVG !== nothing
    handler = CVG.cvg_visualize(scene); CVG.cvg_wait_for_destroy(handler.window)

    CVG.cvg_close(handler); CVG.cv_destroy(scene)
end

# vim:syn=julia:cc=79:fdm=marker:sw=4:ts=4:


Wigner  = PyCall.PyNULL()
Optics  = PyCall.PyNULL()
Spin = PyCall.PyNULL()

function init_physics()
    copy!(Wigner,  PyCall.pyimport_conda("sympy.physics.wigner",       "sympy"))
    copy!(Optics,  PyCall.pyimport_conda("sympy.physics.optics",       "sympy"))
    copy!(Spin ,   PyCall.pyimport_conda("sympy.physics.quantum.spin", "sympy"))
end

module Physics

using PyCall
using SymPy

physics = [(SymPy.Wigner, [:clebsch_gordan,
                           :dot_rot_grad_Ynm,
                           :gaunt,
                           :racah,
                           :wigner_3j,
                           :wigner_6j,
                           :wigner_9j]),
           (SymPy.Optics, [:RayTransferMatrix,
                           :FreeSpace,
                           :CurvedRefraction,
                           :FlatMirror]),
           (SymPy.Spin  , [:WignerD])
]

for (m, meths) in physics
    for meth in meths
        meth_name = string(meth)
        @eval begin
            # @doc """
            # `$($meth_name)`: a SymPy function.
            #     The SymPy documentation can be found through: http://docs.sympy.org/latest/search.html?q=$($meth_name)
            #     """ ->
            ($meth)(args...;kwargs...) = pycall(getproperty($m, $meth_name), PyAny, map(Sym, args)...,
                                                [k=>Sym(v) for (k,v) in pairs(kwargs)]...)
        end
        eval(Expr(:export, meth))
    end
end

end


module ParPoly

export @estrin, @horner_split, @horner_split_simd
using SIMD

include("estrin.jl")

macro horner_split_simd(t,p...)
    x1   = gensym("x1")
    x2   = gensym("x2")
    x1x2 = gensym("x1x2")
    x2x2 = gensym("x2x2")

    blk = quote
        T = typeof($(esc(t)))
        $x1 = $(esc(t))
        $x2 = $(esc(t)) * $(esc(t))
        $x1x2 = Vec{2,T}(($x1,$x2))
        $x2x2 = Vec{2,T}(($x2,$x2))
    end
    n = length(p)
    p0 = esc(p[1])
    m = n÷2
    c = Array(Symbol,m)
    for i = 1:m
        c[i] = gensym("c$i")
    end

    if isodd(n)
        for i = 1:m
            push!(blk.args, :( $(c[i]) = Vec{2,T}(( $(esc(p[2i])), $(esc(p[2i+1])) )) ))
        end
    elseif iseven(n)
        for i = 1:m-1
            push!(blk.args, :( $(c[i]) = Vec{2,T}(( $(esc(p[2i])), $(esc(p[2i+1])) )) ))
        end
        push!(blk.args, :( $(c[m]) = Vec{2,T}(( $(esc(p[2m])), zero(T) )) )) # pad with a zero
    end

    ex = c[end]
    for i = m-1:-1:1
        ex = :(muladd($(x2x2), $ex, $(c[i])))
    end
    cc = gensym("cc")
    push!(blk.args, :($p0 + sum($x1x2*$ex) ))

    return blk
end

macro horner_split(x,p...)
    t = gensym("x1")
    t2 = gensym("x2")
    blk = quote
        $t = $(esc(x))
        $t2 = $(esc(x)) * $(esc(x))
    end

    n = length(p)
    p0 = esc(p[1])

    if isodd(n)
        ex_o = esc(p[end-1])
        ex_e = esc(p[end])

        for i = n-3:-2:2
            ex_o = :(muladd($(t2), $ex_o, $(esc(p[i]))))
        end
        for i = n-2:-2:2
            ex_e = :(muladd($(t2), $ex_e, $(esc(p[i]))))
        end
    elseif iseven(n)
        ex_o = esc(p[end])
        ex_e = esc(p[end-1])

        for i = n-2:-2:2
            ex_o = :(muladd($(t2), $ex_o, $(esc(p[i]))))
        end
        for i = n-3:-2:2
            ex_e = :(muladd($(t2), $ex_e, $(esc(p[i]))))
        end
    end
    push!(blk.args,:($(p0) + $(t)*$(ex_o) + $(t2)*$(ex_e)) )
    return blk
end

end


# ---
# title: 1290. Convert Binary Number in a Linked List to Integer
# id: problem1290
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Linked List, Bit Manipulation
# link: <https://leetcode.com/problems/convert-binary-number-in-a-linked-list-to-integer/description/>
# hidden: true
# ---
#
# Given `head` which is a reference node to a singly-linked list. The value of
# each node in the linked list is either 0 or 1. The linked list holds the
# binary representation of a number.
#
# Return the _decimal value_ of the number in the linked list.
#
#
#
# **Example 1:**
#
# ![](https://assets.leetcode.com/uploads/2019/12/05/graph-1.png)
#
#
#
#     Input: head = [1,0,1]
#     Output: 5
#     Explanation: (101) in base 2 = (5) in base 10
#
#
# **Example 2:**
#
#
#
#     Input: head = [0]
#     Output: 0
#
#
# **Example 3:**
#
#
#
#     Input: head = [1]
#     Output: 1
#
#
# **Example 4:**
#
#
#
#     Input: head = [1,0,0,1,0,0,1,1,1,0,0,0,0,0,0]
#     Output: 18880
#
#
# **Example 5:**
#
#
#
#     Input: head = [0,0]
#     Output: 0
#
#
#
#
# **Constraints:**
#
#   * The Linked List is not empty.
#   * Number of nodes will not exceed `30`.
#   * Each node's value is either `0` or `1`.
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


@testset "Unwrap 1D" begin
    @test unwrap([0.1, 0.2, 0.3, 0.4]) ≈ [0.1, 0.2, 0.3, 0.4]
    @test unwrap([0.1, 0.2 + 2pi, 0.3, 0.4]) ≈ [0.1, 0.2, 0.3, 0.4]
    @test unwrap([0.1, 0.2 - 2pi, 0.3, 0.4]) ≈ [0.1, 0.2, 0.3, 0.4]
    @test unwrap([0.1, 0.2 - 2pi, 0.3 - 2pi, 0.4]) ≈ [0.1, 0.2, 0.3, 0.4]
    @test unwrap([0.1 + 2pi, 0.2, 0.3, 0.4]) ≈ [0.1, 0.2, 0.3, 0.4]
    @test unwrap([0.1, 0.2 + 2pi, 0.3, 0.4]) ≈ [0.1, 0.2, 0.3, 0.4]

    test_v = [0.1, 0.2, 0.3 + 2pi, 0.4]
    res_v = unwrap(test_v)
    @test test_v ≈ [0.1, 0.2, 0.3 + 2pi, 0.4]
    res_v .= 0
    unwrap!(test_v)
    @test test_v ≈ [0.1, 0.2, 0.3, 0.4]

    # test unwrapping within multi-dimensional array
    wrapped = [0.1, 0.2 + 2pi, 0.3, 0.4]
    unwrapped = [0.1, 0.2, 0.3, 0.4]
    wrapped = hcat(wrapped, wrapped)
    unwrapped = hcat(unwrapped, unwrapped)

    # test generically typed unwrapping
    types = (Float32, Float64)
    for T in types
        A_unwrapped = collect(range(0, stop=4convert(T, π), length=10))
        A_wrapped = A_unwrapped .% (2convert(T, π))

        test(I, J) = I ≈ J || I .+ 2π ≈ J || I .+ 4π ≈ J

        @test test(unwrap(A_wrapped), A_unwrapped)
        unwrap!(A_wrapped)
        @test test(A_wrapped, A_unwrapped)
    end
end

# tests for multi-dimensional unwrapping
@testset "Unwrap 2D" begin
    types = (Float32, Float64)
    for T in types
        v_unwrapped = collect(range(0, stop=4convert(T, π), length=7))
        A_unwrapped = v_unwrapped .+ v_unwrapped'
        A_wrapped = A_unwrapped .% (2convert(T, π))

        test_unwrapped = unwrap(A_wrapped)
        d = first(A_unwrapped) - first(test_unwrapped)
        @test (test_unwrapped .+ d) ≈ A_unwrapped
        unwrap!(A_wrapped)
        d = first(A_unwrapped) - first(A_wrapped)
        @test (A_wrapped .+ d) ≈ A_unwrapped

        v_unwrapped_range = collect(range(0, stop=4, length=7))
        A_unwrapped_range = v_unwrapped_range .+ v_unwrapped_range'
        test_range = convert(T, 2)
        A_wrapped_range = A_unwrapped_range .% test_range

    end
end

@testset "Unwrap 3D" begin
    types = (Float32, Float64)
    f(x, y, z) = 0.1x^2 - 2y + 2z
    f_wraparound2(x, y, z) = 5*sin(x) + 2*cos(y) + z
    f_wraparound3(x, y, z) = 5*sin(x) + 2*cos(y) - 4*cos(z)
    for T in types
        grid = range(zero(T), stop=2convert(T, π), length=11)
        f_uw = f.(grid, grid', reshape(grid, 1, 1, :))
        f_wr = f_uw .% (2convert(T, π))
        uw_test = unwrap(f_wr)
        offset = first(f_uw) - first(uw_test)
        @test (uw_test.+offset) ≈ f_uw rtol=eps(T) #oop, nowrap
        # test in-place version
        unwrap!(f_wr)
        offset = first(f_uw) - first(f_wr)
        @test (f_wr.+offset) ≈ f_uw rtol=eps(T) #ip, nowrap

        f_uw = f_wraparound2.(grid, grid', reshape(grid, 1, 1, :))
        f_wr = f_uw .% (2convert(T, π))
        uw_test = unwrap(f_wr)
        offset = first(f_uw) - first(uw_test)
        @test (uw_test.+offset) ≈ f_uw #oop, 2wrap
        # test in-place version
        unwrap!(f_wr)
        offset = first(f_uw) - first(f_wr)
        @test (f_wr.+offset) ≈ f_uw #ip, 2wrap

        f_uw = f_wraparound3.(grid, grid', reshape(grid, 1, 1, :))
        f_wr = f_uw .% (2convert(T, π))
        uw_test = unwrap(f_wr)
        offset = first(f_uw) - first(uw_test)
        @test (uw_test.+offset) ≈ f_uw #oop, 3wrap
        # test in-place version
        unwrap!(f_wr)
        offset = first(f_uw) - first(f_wr)
        @test (f_wr.+offset) ≈ f_uw #oop, 3wrap
    end
end


### Hastings & Powell 1991
## Reproducing figure 4
using LaTeXStrings

### Figure 4 (Bifurcation diagram)

## Define initial parameters
p = a1, a2, b1, b2, d1, d2 = [5.0, 0.1, 2.0, 2.0, 0.4, 0.01]
u0 = [0.76, 0.16, 9.9]
tspan = (0.0, 10000.0)
# b1 values interval
bs = collect(2.0:0.01:6.2)
# number of values to collect
N = 1000

## Loop to select all values of z for all b1 values
Z = zeros(Float64, (N, length(bs)))
Zmax = zeros(Float64, length(bs))
Zmin = zeros(Float64, length(bs))
errors = zeros(Float64, length(bs))
for j in 1:length(bs)
  # Change b1 value
  local p[3] = local b1 = bs[j]
  # Solve for given b1 value
  local prob = ODEProblem(parameterized_model!, u0, tspan, p)
  local sol = solve(prob)
  # Save values
  if length(sol[3,:]) > N
    Z[:,j] = sol[3, end-(N-1):end]
    Zmax[j] = maximum(Z[:, j])
    Zmin[j] = minimum(Z[:, j])
  else
    # Safeguard for when true solution is unstable
    errors[j] = j
  end
end
Z
Zmax
Zmin
errors[errors .> 0]

## Select only local maxima
Zmaxloc = zeros(Float64, size(Z))
for i in 2:(N-1), j in 1:length(bs)
  # Select maxima values(greater than neighbors)
  if Z[i, j] > Z[i-1, j] && Z[i, j] > Z[i+1, j]
    # Select first local maxima only with arbitrary interval
    Zrange = Zmax[j] - Zmin[j]
    if Z[i, j] > (Zmin[j] + 0.66*(Zrange))
       Zmaxloc[i, j] = Z[i, j]
    end
  end
end
Zmaxloc

## Visualize results
# Vector of b1 values repeated N times
B = repeat(bs, inner = N)
# Match local maxima with respective b1 value
points_maxloc = unique(hcat(B, vec(Zmaxloc)); dims = 1)
# Remove null values
pts_maxloc = points_maxloc[points_maxloc[:,2] .> 0.00, :]
# Fig.4 A)
fig4a = scatter(pts_maxloc[:, 1], pts_maxloc[:, 2],
                markersize = 2, markeralpha = 0.5,
                xlim = (2.2, 3.2), ylim = (9.5, 13.0),
                xlabel = L"b_1", ylabel = L"z_{max}", legend = false,
                grid = :none, c = :black, msw = 0.0, msc = :black,
                title = "a", titleloc = :left
                )
# Fig.4 B)
fig4b = scatter(pts_maxloc[:, 1], pts_maxloc[:, 2],
                markersize = 2, markeralpha = 0.5,
                xlim = (3.0, 6.5), ylim = (3.0, 10.0),
                xticks = 3.0:0.5:6.5, yticks = (3.0:1.0:10.0, string.(3.0:1.0:10.0)),
                xlabel = L"b_1", ylabel = L"z_{max}", legend = false,
                grid = :none, c = :black, msw = 0.0, msc = :black,
                title = "b", titleloc = :left
                )
# Fig.4 C)
fig4c = scatter(pts_maxloc[:, 1], pts_maxloc[:, 2],
                markersize = 2, markeralpha = 0.5,
                xlim = (2.25, 2.6), ylim = (11.4, 12.8),
                xticks = 2.25:0.05:2.60, yticks = 11.1:0.2:12.8,
                xlabel = L"b_1", ylabel = L"z_{max}", legend = false,
                grid = :none, c = :black, msw = 0.0, msc = :black,
                title = "c", titleloc = :left
                )
# Fig.4 - Combine A,B,C
fig4 = plot(fig4a, fig4b, fig4c, layout = (3, 1),
            size = (500, 900), titlefontsize = 11,
            tickfontsize = 8, guidefontsize = 10,
            left_margin = 10Plots.mm, right_margin = 5Plots.mm,
            bottom_margin = 2Plots.mm,
            framestyle = :box
            )

## Export figures
savefig(fig4, joinpath("..", "article", "figures", "fig4.pdf"))


@recipe function f(::Type{Val{:tsaplot}}, x, y, z)

    N = length(x)

    grid := false
    legend := false
    linewidth --> 1

    @series begin
        seriestype := :stem
        seriescolor --> :black
        x := x
        y := y
    end

    bound = 1.96 / √N
    linestyle := :dash
    linecolor --> :blue

    @series begin
        seriestype := :hline
        y := -bound
    end

    @series begin
        seriestype := :hline
        y := bound
    end

    @series begin
        seriestype := :hline
        linestyle := :solid
        linecolor := :black
        linewidth := 1
        y := 0.0
    end

end

@userplot Correlogram
@recipe function f(c::Correlogram)
    if length(c.args) != 2 || !(typeof(c.args[1]) <: AbstractVector) ||
        !(typeof(c.args[2]) <: Integer)
        error("Correlograms should be given a vector and an integer.  Got: $(typeof(h.args))")
    end
    z, k = c.args
    seriestype := :tsaplot
    1:k, map(i -> autocorrelation(z, i), 1:k)
end

@userplot Partial_Correlogram
@recipe function f(c::Partial_Correlogram; recursive=true)
    if length(c.args) != 2 || !(typeof(c.args[1]) <: AbstractVector) ||
        !(typeof(c.args[2]) <: Integer)
        error("Correlograms should be given a vector and an integer.  Got: $(typeof(h.args))")
    end
    z, k = c.args
    seriestype := :tsaplot

    if recursive
        ρ = map(i -> autocorrelation(z, i), 1:k)
        T = typeof(zero(eltype(z)) / 1)
        ϕ = zeros(T, k)
        ϕpp = zeros(T, k)
        σ2 = Ref{T}(autocovariance(z, 0))
        for i in 1:k
            levinson_durbin!(ϕ, σ2, ρ, i)
            ϕpp[i] = ϕ[i]
        end
    else
        ϕpp = map(i -> partial_autocorrelation(z, i), 1:k)
    end

    1:k, ϕpp
end


function as_sub(x::AbstractMatrix)
    y = similar(x, eltype(x), tuple(([size(x)...]*2)...))
    y = subview(y, 2:2:size(y,1), 2:2:size(y,2))
    for j=1:size(x,2)
        for i=1:size(x,1)
            y[i,j] = x[i,j]
        end
    end
    y
end

as_sub(eye(2))


@testset "Testing discrete_rv.jl" begin

    # set up
    n = 10
    x = rand(n)
    x ./= sum(x)
    drv = DiscreteRV(x)

    # test Q sums to 1
    @test drv.Q[end] ≈ 1.0

    # test lln
    draws = rand(drv, 100_000)
    c = counter(draws)
    counts = Vector{Float64}(undef, n)
    for i=1:n
        counts[i] = c[i]
    end
    counts ./= sum(counts)

    @test isapprox(maximum(abs, counts - drv.q), 0.0; atol=1e-2)

    draws = Vector{Int}(undef, 100_000)
    rand!(draws, drv)
    c = counter(draws)
    counts = Vector{Float64}(undef, n)
    for i=1:n
        counts[i] = c[i]
    end
    counts ./= sum(counts)

    @test isapprox(maximum(abs, counts - drv.q), 0.0; atol=1e-2)

end  # testset


using Compat.Test
import Roots.fzero

# test robustness of derivative free algorithms

## tests:http://people.sc.fsu.edu/~jburkardt/cpp_src/test_zero/test_zero.html
function newton_baffler(x)
    if ( x - 0.0 ) < -0.25
        0.75 * ( x - 0 ) - 0.3125
    elseif  ( x - 0 ) < 0.25
        2.0 * ( x - 0 )
    else
        0.75 * ( x - 0 ) + 0.3125
    end
end

pathological = [
                (x ->sin( x ) - x / 2, .1),
                (x-> 2 * x - exp( - x ), 1),
                (x -> x * exp( - x ), 1),
                (x -> exp( x ) - 1 / ( 10 * x )^2, 1),
                (x -> ( x + 3 ) * ( x - 1 )^2, 1),

                (x -> exp( x ) - 2 - 1 / ( 10 * x )^2 + 2 / ( 100 * x )^3, 1),
                (x -> cbrt(x), 1),
                (x -> cos( x ) - x, 1),
                (x-> newton_baffler(x), 8),
                (x -> 20.0 * x / ( 100.0 * x^2 + 1.0), 0.095),

                (x ->  ( 4.0 + x^2) * ( 2.0 + x ) * ( 2.0 - x )  / ( 16.0 * x * x * x * x + 0.00001 ), 1),
                (x -> (x == 1.0) ? float(0) : sign(x-1.0) * exp(log(1e4) + log(abs(x - 1.0)) - 1.0/(x-1.0)^2), 1),
                (x -> 0.00000000001 * (x - 100.0), 1),
                (x -> 1.0 / ( ( x - 0.3 ) * ( x - 0.3 ) + 0.01 ) + 1.0 / ( ( x - 0.9 ) * ( x - 0.9 ) + 0.04 ) + 2.0 * x - 5.2, -1),
                (x -> ( 1 - 6x^2) * cbrt(x) * exp(-x^2) / (3*x), -0.25),

                (x -> ( pi * ( x - 5.0 ) / 180.0 ) - 0.8 * sin( pi * x / 180.0 ), 1),
                (x -> x^3 - 2*x - 5, 2),
                (x -> 1e6 * (x^7 -7x^6 +21x^5 -35x^4 +35x^3-21x^2+7x-1),  0.990),
                (x -> cos(100*x)-4*erf(30*x-10), 0.0)
                ]

for (f1, x0) in pathological

      try
          x = fzero(f1, x0)
          @assert (f1(x) == 0.0 || f1(prevfloat(x)) * f1(nextfloat(x)) <= 0 || abs(f1(x)) <= eps(float(x0))^(1/2))
      catch err
#          rethrow(err)
#          if !isa(err, Roots.PossibleExtremaReached)
#              throw(err)
#          end
      end
end

## make a graphic comparing values
function run_robustness_test()
    orders = [0,1,2,5,8,16]
    N = length(orders)
    k,n = length(pathological), 50
    m = ["" for i=1:(N + 2)*k, j=1:n]
    ## how to mark values. If error we also use "x"
    function howgood(f1, xstar)
        if f1(xstar) == 0       # perfect
            "."
        elseif f1(prevfloat(xstar)) * f1(nextfloat(xstar)) < 0 # perfect
        "-"
        elseif -1e-16 < f1(xstar) < 1e-16 # approximate
            "~"
        elseif abs(xstar) > 5   # ran away yet converged
            "!"
        else                    # not close?
            "*"
        end
    end

    for i in 1:k
        f1,x = pathological[i]
        xs = linspace(-5, 5, n)
        for k in 1:n
            for j in 1:length(orders)
                try
                    xstar = fzero(f1, xs[k], order=orders[j], maxeval=100)
                    m[(N+1)*(i-1)+j, k] = howgood(f1, xstar)
                catch e
                    m[(N+1)*(i-1)+j, k] = "x"
                end
            end
        end
    end
    println("\n----")
    print(join(mapslices(join, m, 2),"\n"))
end


function i_plus_x(i, x, N)
	i = (i+N+x) % N
	(i == 0) && (i = N)
	return i
end

function mean(vet::AbstractVector)
	(length(vet) == 0) && (return 0)
	return sum(vet) / length(vet)
end

function issymmetric(mat::AbstractMatrix)
	(size(mat, 1) != size(mat, 2)) && (return false)
	for i in 1:size(mat, 1), j in i+1:size(mat, 2)
		(mat[i,j] != mat[j,i]) && (return false)
	end
	return true
end

function hasnodeOrError(network::AbstractNetwork, idx_node::Integer)
	hasnode(network, idx_node) ||
		throw(ArgumentError("Invalid node index! Node $idx_node doesn't exist in the network!"))
end


using Test

#--- test meshes
m1a = Mesh( (xmin=-2*π,xmax=π,N=3) )
m1b = Mesh( [-2*π -π 0 ] )

m2a = Mesh( (xmin=-2*π,xmax=2π,N=2) )
m2b = Mesh( [-2*π 0 ] )

function the_mesh(m::Mesh)
    [m.N, m.xmin, m.xmax, m.dx, m.x, m.kmin, m.kmax, m.dk, m.k]
end

@testset "Tests on Mesh" begin
    @test the_mesh(m1a) ≈ [3,-2π,π,π,[-2π; -π; 0],-2/3,2/3,2/3,[0; 2/3; -2/3]]
    @test the_mesh(m1a) == the_mesh(m1b)

    @test the_mesh(m2a) ≈ [2,-2π,2π,2π,[-2π; 0],-1//2,0,1//2,[0; -1/2]]
    @test the_mesh(m2a) == the_mesh(m2b)

end

#--- test times
t1a = Times( 0:0.1:1; ns=3 )
t1b = Times( (dt=0.1 , T=1); ns=3 )

t2a = Times( 0:1//10:1 ; Ns=3 )
t2b = Times( (dt=0.1 , T=1); Ns=3 )

function the_times(t::Times)
    [t.Nc, t.Ns, t.ns, t.tfin, t.dt, t.tc, t.ts]
end

@testset "Tests on times" begin
    @test the_times(t1a) == [11,4,[3,3,3],1.,0.1,0:0.1:1,[0.0, 0.3, 0.6, 0.9]]
    @test the_times(t1a) == the_times(t1b)

    @test the_times(t2a) == [11,4,[3,4,3],1,0.1,0:0.1:1,[0.0, 0.3, 0.7, 1.0]]
    @test the_times(t2a) == the_times(t2b)
end

#--- test data
v1 = [1; 3; 5];v2 = [2; 4; 6];
M = [ v1 v2 ];
data=Data(M)
@testset "Tests on data" begin
    @test data.U==[[1 2 ; 3 4 ; 5 6 ]]
    @test data.datalength==length(v1)
    @test data.datasize==2
end

#--- test problems
@testset "Tests on problems" begin
    param = (L=10, N=20, ϵ = 0.1, μ = 0.1, T=1, dt=1/30)
    init = Random(param)
    solver=Euler_naive()
    #solver=RK4(model1)  # the tests will fail with this solver, since it uses pre-allocation
    # the tests also fail if we define  model=WWn(param;n=1), and use model to build the problems
    pb1  = Problem(WWn(param;n=1,mesh=Mesh(param)), init, param; solver=solver)
    pb1a = Problem(WWn(param;n=1), init, param; solver=solver)
    pb1b = Problem(WWn(param;n=1), init, param; solver=solver)
    pb2  = Problem(WWn(param;n=2), init, merge(param,(ns=2,)); solver=solver)
    pb2a = Problem(WWn(param;n=2), init, merge(param,(ns=2,)); solver=solver)
    pb2b = Problem(WWn(param;n=2), init, merge(param,(ns=2,)); solver=solver)

    solve!([pb1a pb2a pb1b pb2b];verbose=false);
    solve!(pb1;verbose=false);solve!(pb2;verbose=false);

    for i in 1:2
        @test solution(pb1)[i]==solution(pb1a)[i]
        @test solution(pb1)[i]==solution(pb1b)[i]
        @test solution(pb2)[i]==solution(pb2a)[i]
        @test solution(pb2)[i]==solution(pb2b)[i]
    end
end

#--- test init
@testset "Tests on init" begin
    param = (L=10, N=2^8)
    mesh=Mesh(param);x=mesh.x;
    f(x) = exp.(-x.^2);g(x) = x.*exp.(-(x.-1).^2);
    inits=InitialData[]
    push!(inits, Init(f , g) )
    push!(inits, Init( mesh, f(x) , g(x); fast = false ) )
    push!(inits, Init( x, f(x) , g(x); fast = false ) )
    init = inits[1]
    for i in 2:3
        @test init.η([π])≈inits[i].η([π])
        @test init.v([π])≈inits[i].v([π])
    end
end


using ColabINAOE2021

export dqc, dql, dqk
export wijct, wijlt, wijkt

"### `EExchanges.jl` gives the functions and method for the study of energy distribution and exchange within the evolution of initial amplitudes on masses arrays given in `MArrays.jl`"

# 1. ---------------------------------------------------------------------------

"### Circular array energy exchange "

"`dqc` is the derivative of the function `qc` for the evolution on circular masses arrays. It is called as `dqc(k, n, t, N, V::Vector)`"
function dqc(k, n, t, N, V::Vector)
    v = zeros(Float64, (N + 1, N + 1))
	for m in 0:N
		for l in 0:m
			v[l + 1, m + 1] = sin(2 * t * (√k) * sin(π * m / (N + 1))) * sin(π * m / (N + 1)) * cos(2 * π * m * (n - l) / (N + 1)) * V[l + 1]
		end
	end
	return ((-2 * (√k)) / (N + 1)) * sum(v)
end

"`wijc` is the energy exchange function between the position _i_ and _j_ at time _t_ of an initial amplitude vector _V_, such that, but not strictly, _i < j_. It is called `wijc(k, i, j, t, N, V::Vector)`."
wijc(k, i, j, t, N, V) = qc(k, i, t, N, V) * dqc(k, j, t, N, V) - dqc(k, i, t, N, V) * qc(k, j, t, N, V)

"`wijct` is an ad-hoc constructor (comprehension) to evalute the evolution of the energy exchange given a time span _t_ with an initial amplitudes vector _V_. It is called `wijct(k, i, j, t, N, V::Vector)`"
wijct(k, i, j, t, N, V) = [wijc(k, i, j, t[τ], N, V) for τ in 1:length(t)]

"`wijcm` is a function that calculates the energy exchange between the position _i_, and all other positions at time _t_. The item who return is a _N×`length(t)`_ matrix. It is called `wijcm(k, i, t, N, V::Vector)`"
function wijcm(k, i, t, N, V)
	v = zeros(Float64, (N, length(t)))
	for τ in 1:length(t)
			for j in 1:N
				v[j, τ] = wij(k, i, j, t[τ], N, V)
			end
		end
	return v
end

# 2. ---------------------------------------------------------------------------

"""### Lineal array energy exchange """

"`dql` is the derivative of the function `ql` for the evolution on lineal masses arrays. It is called as `dql(k, n, t, N, V::Vector)`"
function dql(k, n, t, N, V::Vector)
	v = zeros(Float64, (N + 1, N + 1))
	for i in 0:N
		for j in 0:N
			v[i + 1, j + 1] = (ColabINAOE2021.twocheb(j, n, i, N) / ColabINAOE2021.nc(j, N)) * sin(2 * t * (√k) * sin(ColabINAOE2021.ϕ(j + 1, N) / 2)) * sin(ColabINAOE2021.ϕ(j + 1, N) / 2) * V[i + 1]
		end
	end
	return (-2 * (√k)) * sum(v)
end

"`wijl` is the energy exchange function between the position _i_ and _j_ at time _t_ of an initial amplitude vector _V_, such that, but not strictly, _i < j_. It is called `wijl(k, i, j, t, N, V::Vector)`."
wijl(k, i, j, t, N, V) = ql(k, i, t, N, V) * dql(k, j, t, N, V) - dql(k, i, t, N, V) * ql(k, j, t, N, V)

"`wijlt` is an ad-hoc constructor (comprehension) to evalute the evolution of the energy exchange given a time span _t_ with an initial amplitudes vector _V_. It is called `wijlt(k, i, j, t, N, V::Vector)`"
wijlt(k, i, j, t, N, V) = [wijl(k, i, j, t[τ], N, V) for τ in 1:length(t)]

# 3. ---------------------------------------------------------------------------

"""### Kravchuk array energy exchange """

"`dqk` is the derivative of the function `ql` for the evolution on Kravchuk masses arrays. It is called as `dqk(n, t, N, V::Vector)`"
function dqk(n, t, N, V::Vector)
	v = zeros(Float64, (N + 1, N + 1))
	for m in 0:N
		for l in 0:N
			v[m + 1, l + 1] = (-√m) *ColabINAOE2021.k2(n, l, m, N) * sin(t * (√m)) * V[l + 1]
		end
	end
	return sum(v)
end

"`wijk` is the energy exchange function between the position _i_ and _j_ at time _t_ of an initial amplitude vector _V_, such that, but not strictly, _i < j_. It is called `wijk(i, j, t, N, V::Vector)`."
wijk(i, j, t, N, V) = qk(i, t, N, V) * dqk(j, t, N, V) - dqk(i, t, N, V) * qk(j, t, N, V)

"`wijkt` is an ad-hoc constructor (comprehension) to evalute the evolution of the energy exchange given a time span _t_ with an initial amplitudes vector _V_. It is called `wijkt(k, i, j, t, N, V::Vector)`"
wijkt(i, j, t, N, V) = [wijk(i, j, t[τ], N, V) for τ in 1:length(t)]


# Script

using Distributions, PyPlot, BayesianNonparametricStatistics

β=0.5
θ = sumoffunctions(vcat([faberschauderone],[faberschauder(j,k) for j in 0:4 for k in 1:2^j]),vcat([1.0],[(-1)^(j*k)*2^(-β*j) for j in 0:4 for k in 1:2^j]))

x = 0.0:0.001:1.0

y = θ.(x)

# Uncomment the following lines to plot θ.
# clf()
# plot(x,y)

sde = SDEWithConstantVariance(θ, 1.0, 0.0, 1000.0, 0.01)

X = rand(sde)

# Uncomment the following lines to plot a sample from sde.
# clf()
# plot(X)

M = SDEModel(1.0, 0.0)

Π = FaberSchauderExpansionWithGaussianCoefficients([2^(β*j) for j in 0:4])
postΠ = calculateposterior(Π, X, M )

for k in 1:100
    f = rand(postΠ)
    y = f.(x)
    plot(x,y)
end


@defcomp co2cycle begin
    MAT     = Variable(index=[time])    # Carbon concentration increase in atmosphere (GtC from 1750)
    ML      = Variable(index=[time])    # Carbon concentration increase in lower oceans (GtC from 1750)
    MU      = Variable(index=[time])    # Carbon concentration increase in shallow oceans (GtC from 1750)

    E       = Parameter(index=[time])   # Total CO2 emissions (GtCO2 per year)
    mat0    = Parameter()               # Initial Concentration in atmosphere 2010 (GtC)
    ml0     = Parameter()               # Initial Concentration in lower strata 2010 (GtC)
    mu0     = Parameter()               # Initial Concentration in upper strata 2010 (GtC)

    # Parameters for long-run consistency of carbon cycle
    b11     = Parameter()               # Carbon cycle transition matrix atmosphere to atmosphere
    b12     = Parameter()               # Carbon cycle transition matrix atmosphere to shallow ocean
    b21     = Parameter()               # Carbon cycle transition matrix biosphere/shallow oceans to atmosphere
    b22     = Parameter()               # Carbon cycle transition matrix shallow ocean to shallow oceans
    b23     = Parameter()               # Carbon cycle transition matrix shallow to deep ocean
    b32     = Parameter()               # Carbon cycle transition matrix deep ocean to shallow ocean
    b33     = Parameter()               # Carbon cycle transition matrix deep ocean to deep oceans

    function run_timestep(p, v, d, t)
        # Define function for MAT
        if is_first(t)
            v.MAT[t] = p.mat0
        else
            v.MAT[t] = v.MAT[t - 1] * p.b11 + v.MU[t - 1] * p.b21 + (p.E[t - 1] * (5 / 3.666))
        end

        # Define function for MU
        if is_first(t)
            v.MU[t] = p.mu0
        else
            v.MU[t] = v.MAT[t - 1] * p.b12 + v.MU[t - 1] * p.b22 + v.ML[t - 1] * p.b32
        end

        # Define function for ML
        if is_first(t)
            v.ML[t] = p.ml0
        else
            v.ML[t] = v.ML[t - 1] * p.b33 + v.MU[t - 1] * p.b23
        end
    end
end


using Statistics

X = rand(50, 20)

# True generating function has intercept, 2 factors and a lag range of 3-5.
genfn(r, X) = 2.56 - 7.8*mean(X[(r-5):(r-4), 2]) + 9.3*X[r-3, 11] + (0.01 * rand())

y = Float64[genfn(r, X) for r in 6:size(X, 1)]

# Now in reality we don't know the real model se we will have to train it
# from data. Let's do it and try to recover an interpretable model.
using InterpretableModels

# We don't know how many factors so let's aim for 3 and a lag range of 2-10
m = MovingAverageLITRM(3, size(X, 2); minlag=2, maxlag=10)

fit!(m, X, y; MaxTime = 10.0)

println(m) # The model should have been recovered quite well...


function triangle(n)
    n < 0 && throw(DomainError("Invalid number"))
    n == 0 && return []
    rows = [[1]]
    for i in 2:n
        previous_row = [rows[i-1]..., 0]
        push!(rows, previous_row + reverse(previous_row))
    end
    rows
end


using Plots
using Convex
using Mosek
using SCS
using Interpolations
using DelimitedFiles

# This is a julia implementaiton of an example
# from "ℓ1 Trend Filtering" by Kim et al (2009)

I_sets = [
[1, 2000],
[1, 968, 2000],
[1, 339, 897, 2000],
[1, 337, 983, 1145, 2000],
[1, 317, 791, 885, 1207, 2000],
[1, 350, 636, 755, 884, 1208, 2000],
[1, 349, 631, 788, 834, 991, 1150, 2000],
[1, 349, 631, 787, 835, 988, 1191, 1853, 2000],
[1, 350, 636, 755, 887, 922, 976, 1181, 1853, 2000],
[1, 365, 515, 545, 628, 769, 837, 987, 1190, 1853, 2000]
]

cost_l0 = [find_optimal_y_values(ℓ, I_sets[k])[2] for k=1:10]

#---

function is_local_max(y)
    v = zeros(y, Bool)
    for k=2:length(y)-1
        v[k] = (y[k] >= y[k-1]) && (y[k] >= y[k+1])
    end
    return v
end

#---
data = readdlm(joinpath(joinpath(dirname(@__FILE__),"data",,"snp500.txt"))

N = length(data)
H =  spdiagm((ones(N-2), -2*ones(N-2), ones(N-2)), (0,1,2))

x = Variable(N)

#---
r = 6
I = I_sets[r]
Y, _ = find_optimal_y_values(ℓ, I)
y_l0 = interpolate((I,), Y, Gridded(Linear()))[1.0:N]

λ = 100
problem = minimize(0.5*sumsquares(x - data) + λ*norm(H*x, 1))

solve!(problem, MosekSolver([("QUIET", true)]))
println("Opt. val: ",problem.optval, "   (problem status: ", problem.status, ")")

y_l1 = evaluate(x)[:]

v = sortperm(abs.(H*y_l1) + 1*is_local_max(abs.(H*y_l1)))

plot(H*y_l1)
plot!(v[end-r+2:end], zeros(r-1), m=:circle)

I_l1 = [1; sort(v[end-r+2:end]); N]
Y_l1, f_l1 = find_optimal_y_values(ℓ, I_l1)

y_l10 = interpolate((I_l1,), Y_l1, Gridded(Linear()))[1.0:N]

println("ℓ1: ", sum((evaluate(x) - data).^2))
println("ℓ0/ℓ1: ", sum((y_l10 - data).^2))
println("ℓ0: ", sum((y_l0 - data).^2))

#subplot(211)
plot(data)
#plot!(evaluate(x), color="green", linewidth=3)
plot!(I, Y, color="red", linewidth=3)
plot!(I_l1, Y_l1, l=(:cyan, :dash), linewidth=3)
#---

λ_vec = 10 .^ range(1, stop=4, length=60)
cost = zeros(length(λ_vec), 10)

for (k, λ)=enumerate(λ_vec)
    problem = minimize(0.5*sumsquares(x - data) + λ*norm(H*x, 1))

    solve!(problem, MosekSolver([("QUIET", true)]))
    println("Opt. val: ",problem.optval, "   (problem status: ", problem.status, ")")

    y_l1 = evaluate(x)[:]
    v = sortperm(abs.(H*y_l1) + 1*is_local_max(abs.(H*y_l1)))
    for r=1:10
        I_l1 = [1; sort(v[end-r+2:end]); N]
        _, f_l1 = find_optimal_y_values(ℓ, I_l1)
        cost[k, r] = f_l1
    end
end

#---

#---
plot(log10(λ_vec), cost[:,3:end], show=true, yaxis=((0,10), 0:0.5:10) )

plot!([1.5, 2.5], [1,1]*cost_l0[3:end]')

#---
[minimum(cost,1)' cost_l0]

#---
 costs = [42.8484   42.8484
  9.25796   9.25796
  4.46404   4.26137
  3.13439   3.02902
  3.09011   2.71251
  2.68868   2.27944
  2.16277   2.03848
  1.94739   1.8517
  1.89786   1.76981
  1.75176   1.59654]

plot(3:10,costs[3:end,:])


@doc raw"""
    Vector3

An object `v` of the type [`Vector3`](@ref) is a vector in the three-dimensional
vector space ``ℝ^3``.

Geometrically spoken, a vector is the difference of two points ``p_2, p_1`` and
and denotes the direction and the distance from ``p_1`` to ``p_2``.

CGAL defines a symbolic constant [`NULL_VECTOR`](@ref). We will explicitly
state where you can pass this constant as an argument instead of a vector
initialized with zeros.
"""
Vector3

"""
    Vector3(a::Point3, b::Point3)

Introduces the vector `b - a`.
"""
Vector3(a::Point3, b::Point3)

"""
    Vector3(s::Segment3)

Introduces the vector `target(s) - source(s)`.
"""
Vector3(s::Segment3)

"""
    Vector3(::NullVector = NULL_VECTOR)

Introduces a null vector `v`.
"""
Vector3(::NullVector = NULL_VECTOR)

"""
    Vector3(x::Real, y::Real, z::Real)

Introduces a vector `v` initialized to ``(x, y, z)``.
"""
Vector3(x::Real, y::Real, z::Real) =
    Vector3(convert(FT, x), convert(FT, y), convert(FT, z))

@doc raw"""
    Vector3(hx::Real, hy::Real, hz::Real, hw::Real = 1)

Introduces a vector `v` initialized to `(hx/hw,hy/hw,hz/hw)`.

!!! info "Precondition"

    `hw` ``≠`` `0`
"""
Vector3(hx::Real, hy::Real, hz::Real, hw::Real) =
    Vector3(convert(RT, x), convert(RT, y), convert(RT, z), convert(RT, hw))

"""
    hx(p::Vector3)

Returns the homogeneous ``x`` coordinate.
"""
hx(p::Vector3)

"""
    hy(p::Vector3)

Returns the homogeneous ``y`` coordinate.
"""
hy(p::Vector3)

"""
    hz(p::Vector3)

Returns the homogeneous ``z`` coordinate.
"""
hz(p::Vector3)

"""
    hw(p::Vector3)

Returns the homogenizing coordinate.
"""
hw(p::Vector3)

"""
    x(p::Vector3)

Returns the Cartesian ``x`` coordinate, that is [`hx()`](@ref)/[`hw()`](@ref).
"""
x(p::Vector3)

"""
    y(p::Vector3)

Returns the Cartesian ``y`` coordinate, that is [`hy()`](@ref)/[`hw()`](@ref).
"""
y(p::Vector3)

"""
    z(p::Vector3)

Returns the Cartesian ``z`` coordinate, that is [`hz()`](@ref)/[`hw()`](@ref).
"""
z(p::Vector3)

@doc raw"""
    homogeneous(v::Vector3, i::Integer)

Returns the `i`ᵗʰ homogeneous coordinate of `v`.

!!! info "Precondition"

    ``0 ≤ i ≤ 3``.
"""
homogeneous(v::Vector3, i::Integer)

@doc raw"""
    cartesian(v::Vector3, i::Integer)

Returns the `i`ᵗʰ Cartesian coordinate of `v`.

!!! info "Precondition"

    ``0 ≤ i ≤ 2``.
"""
cartesian(v::Vector3, i::Integer)

"""
    dimension(v::Vector3)

Returns the dimension (the constant 3).
"""
dimension(v::Vector3)

"""
    squared_length(v::Vector3)

Returns the squared length of `v`.
"""
squared_length(v::Vector3)

"""
    ==(v::Vector3, w::Vector3)

Test for equality: two vectors are equal, iff their ``x``, ``y`` and ``z``
coordinates are equal.

You can compare a vector with the `NULL_VECTOR`.
"""
==(v::Vector3, w::Vector3)

"""
    +(v::Vector3, w::Vector3)

Addition.
"""
+(v::Vector3, w::Vector3)

"""
    -(v::Vector3, w::Vector3)

Subtraction.
"""
-(v::Vector3, w::Vector3)

"""
    -(v::Vector3)

Returns the opposite vector.
"""
-(v::Vector3)

"""
    *(v::Vector3, w::Vector3)

Returns the scalar product (= inner product) of the two vectors.
"""
*(v::Vector3, w::Vector3)

"""
    /(v::Vector3, s::Real)

Division by a scalar.
"""
/(v::Vector3, s::Real)

@cxxdereference Base.:/(v::Vector3, s::Real) = v / convert(RT, s)


function im2col{T}(img::Array{T}, n::Int, col::Array{T}, width::Int, height::Int, channels::Int,
    kernel::NTuple{2,Int}, pad::NTuple{2,Int}, stride::NTuple{2,Int})

  im2col_impl(img[:,:,:,n], col, width, height, channels, kernel, pad, stride)
end

function im2col_impl{T}(img::Array{T}, col::Array{T}, width::Int, height::Int, channels::Int,
    kernel::NTuple{2,Int}, pad::NTuple{2,Int}, stride::NTuple{2,Int})

  kernel_w, kernel_h = kernel
  pad_w, pad_h = pad
  stride_w, stride_h = stride

  height_col = div(height + 2pad_h - kernel_h, stride_h) + 1
  width_col = div(width + 2pad_w - kernel_w, stride_w) + 1
  channels_col = channels * kernel_h * kernel_w

  for c = 0:channels_col-1
    w_offset = c % kernel_w
    h_offset = div(c, kernel_w) % kernel_h
    c_im = div(c, kernel_h * kernel_w) # channel
    for h = 0:height_col-1
      for w = 0:width_col-1
        h_pad = h*stride_h - pad_h + h_offset
        w_pad = w*stride_w - pad_w + w_offset
        if (h_pad >= 0 && h_pad < height && w_pad >= 0 && w_pad < width)
            @inbounds col[1 + (c*height_col+h) * width_col + w] =
                img[1 + (c_im * height + h_pad) * width + w_pad]
        else
          @inbounds col[1 + (c*height_col+h) * width_col + w] = 0
        end
      end
    end
  end
end

function col2im{T}(col::Array{T}, img::Array{T}, n::Int, img_buf::Array{T}, width::Int, height::Int, channels::Int,
    kernel::NTuple{2,Int}, pad::NTuple{2,Int}, stride::NTuple{2,Int})

  col2im_impl(col, img_buf, width, height, channels, kernel, pad, stride)
  img[:,:,:,n] = img_buf
end

function col2im_impl{T}(col::Array{T}, img::Array{T}, width::Int, height::Int, channels::Int,
    kernel::NTuple{2,Int}, pad::NTuple{2,Int}, stride::NTuple{2,Int})

  kernel_w, kernel_h = kernel
  pad_w, pad_h = pad
  stride_w, stride_h = stride

  height_col = div(height + 2pad_h - kernel_h, stride_h) + 1
  width_col = div(width + 2pad_w - kernel_w, stride_w) + 1
  channels_col = channels * kernel_h * kernel_w

  fill!(img, 0)
  for c = 0:channels_col-1
    w_offset = c % kernel_w
    h_offset = div(c, kernel_w) % kernel_h
    c_im = div(c, kernel_w * kernel_h)
    for h = 0:height_col-1
      for w = 0:width_col-1
        h_pad = h * stride_h - pad_h + h_offset
        w_pad = w * stride_w - pad_w + w_offset
        if h_pad >= 0 && h_pad < height && w_pad >= 0 && w_pad < width
          @inbounds img[1 + (c_im * height + h_pad) * width + w_pad] +=
              col[1 + (c * height_col + h) * width_col + w]
        end
      end
    end
  end
end


""" DSLs for defining categories, diagrams, and related structures.

Here "diagram" means diagram in the standard category-theoretic sense, not
string diagram or wiring diagram. DSLs for constructing wiring diagrams are
provided by other submodules.
"""
module DiagrammaticPrograms
export @graph, @fincat, @finfunctor, @diagram, @free_diagram,
  @migrate, @migration

using Base.Iterators: repeated
using MLStyle: @match

using ...GAT, ...Present, ...Graphs, ...CategoricalAlgebra
using ...Theories: munit
using ...CategoricalAlgebra.FinCats: mapvals, make_map
using ...CategoricalAlgebra.DataMigrations: ConjQuery, GlueQuery, GlucQuery
using ...Graphs.BasicGraphs: TheoryGraph
import ...CategoricalAlgebra.FinCats: vertex_name, vertex_named,
  edge_name, edge_named

# Graphs
########

@present TheoryNamedGraph <: TheoryGraph begin
  VName::AttrType
  EName::AttrType
  vname::Attr(V, VName)
  ename::Attr(E, EName)
end

""" Abstract type for graph with named vertices and edges.
"""
@abstract_acset_type AbstractNamedGraph <: AbstractGraph

""" Graph with named vertices and edges.

The default graph type used by [`@graph`](@ref), [`@fincat`](@ref),
[`@diagram`](@ref), and related macros.
"""
@acset_type NamedGraph(TheoryNamedGraph, index=[:src,:tgt,:ename],
                       unique_index=[:vname]) <: AbstractNamedGraph
# FIXME: The edge name should also be uniquely indexed, but this currently
# doesn't play nicely with nullable attributes.

vertex_name(g::AbstractNamedGraph, args...) = subpart(g, args..., :vname)
edge_name(g::AbstractNamedGraph, args...) = subpart(g, args..., :ename)

vertex_named(g::AbstractNamedGraph, name) = only(incident(g, name, :vname))
edge_named(g::AbstractNamedGraph, name)= only(incident(g, name, :ename))

""" Construct a graph in a simple, declarative style.

The syntax is reminiscent of Graphviz. Each line a declares a vertex or set of
vertices, or an edge. For example, the following defines a directed triangle:

```julia
@graph begin
  v0, v1, v2
  fst: v0 → v1
  snd: v1 → v2
  comp: v0 → v2
end
```

Vertices in the graph must be uniquely named, whereas edges names are optional.
"""
macro graph(graph_type, body)
  :(parse_graph($(esc(graph_type)), $(Meta.quot(body))))
end
macro graph(body)
  :(parse_graph(NamedGraph{Symbol,Union{Symbol,Nothing}}, $(Meta.quot(body))))
end

function parse_graph(G::Type, body::Expr; preprocess::Bool=true)
  g = G()
  if preprocess
    body = reparse_arrows(body)
  end
  for stmt in statements(body)
    @match stmt begin
      # v
      v::Symbol => add_vertex!(g, vname=v)
      # u, v, ...
      Expr(:tuple, vs...) => add_vertices!(g, length(vs), vname=vs)
      # u → v
      Expr(:call, :(→), u::Symbol, v::Symbol) => parse_edge!(g, u, v)
      # e : u → v
      Expr(:call, :(:), e::Symbol, Expr(:call, :(→), u::Symbol, v::Symbol)) =>
        parse_edge!(g, u, v, ename=e)
      # (e, f, ...) : u → v
      Expr(:call, (:), Expr(:tuple, es...), Expr(:call, :(→), u::Symbol, v::Symbol)) =>
        for e in es; parse_edge!(g, u, v, ename=e) end
      ::LineNumberNode => nothing
      _ => error("Cannot parse line in graph definition: $stmt")
    end
  end
  return g
end

function parse_edge!(g, s::Symbol, t::Symbol;
                     ename::Union{Symbol,Nothing}=nothing)
  e = add_edge!(g, vertex_named(g, s), vertex_named(g, t))
  if has_subpart(g, :ename)
    g[e,:ename] = ename
  end
  return e
end

function reparse_arrows(expr)
  # `f : x → y` is parsed by Julia as `(f : x) → y`, not `f : (x → y)`.
  @match expr begin
    Expr(:call, :(→), Expr(:call, :(:), f, x), y) =>
      Expr(:call, :(:), f, Expr(:call, :(→), x, y))
    Expr(head, args...) => Expr(head, (reparse_arrows(arg) for arg in args)...)
    _ => expr
  end
end

statements(expr) = (expr isa Expr && expr.head == :block) ? expr.args : [expr]

# Categories
############

""" Present a category by generators and relations.

The result is a finitely presented category (`FinCat`) represented by a graph,
possibly with path equations. For example, the simplex category truncated to one
dimension is:

```julia
@fincat begin
  V, E
  (δ₀, δ₁): V → E
  σ₀: E → V

  σ₀ ∘ δ₀ == id(V)
  σ₀ ∘ δ₁ == id(V)
end
```

The objects and morphisms must be uniquely named.
"""
macro fincat(body)
  parse_category(NamedGraph{Symbol,Symbol}, body)
end

function parse_category(G::Type, body::Expr; preprocess::Bool=true)
  g, eqs = G(), Pair[]
  if preprocess
    body = reparse_arrows(body)
  end
  for stmt in statements(body)
    @match stmt begin
      # v
      v::Symbol => add_vertex!(g, vname=v)
      # u, v, ...
      Expr(:tuple, vs...) => add_vertices!(g, length(vs), vname=vs)
      # e : u → v
      Expr(:call, :(:), e::Symbol, Expr(:call, :(→), u::Symbol, v::Symbol)) =>
        parse_edge!(g, u, v, ename=e)
      # (e, f, ...) : u → v
      Expr(:call, (:), Expr(:tuple, es...), Expr(:call, :(→), u::Symbol, v::Symbol)) =>
        for e in es; parse_edge!(g, u, v, ename=e) end
      # f == g
      Expr(:call, :(==), lhs, rhs) => push!(eqs, parse_path_equation(g, lhs, rhs))
      ::LineNumberNode => nothing
      _ => error("Cannot parse line in category definition: $stmt")
    end
  end
  isempty(eqs) ? FinCat(g) : FinCat(g, eqs)
end

""" Parse an equation between paths in a named graph.
"""
parse_path_equation(g, lhs, rhs) = parse_path(g, lhs) => parse_path(g, rhs)

""" Parse a path in a named graph.
"""
function parse_path(g, expr)
  function parse(expr)
    @match expr begin
      Expr(:call, :compose, args...) => mapreduce(parse, vcat, args)
      Expr(:call, :(⋅), f, g) ||
      Expr(:call, :(⨟), f, g) => vcat(parse(f), parse(g))
      Expr(:call, :(∘), f, g) => vcat(parse(g), parse(f))
      Expr(:call, :id, x::Symbol) => empty(Path, g, vertex_named(g, x))
      f::Symbol => Path(g, edge_named(g, f))
      _ => error("Invalid morphism expression $expr")
    end
  end
  parse(expr)
end

# Functors
##########

""" Define a functor between two finitely presented categories.

Such a functor is defined by sending the object and morphism generators of the
domain category to generic object and morphism expressions in the codomain
category. For example, the following functor embeds the schema for graphs into
the schema for circular port graphs by ignoring the ports:

```julia
@finfunctor TheoryGraph ThCPortGraph begin
  V => Box
  E => Wire
  src => src ⨟ box
  tgt => tgt ⨟ box
end
```
"""
macro finfunctor(dom_cat, codom_cat, body)
  :(parse_functor($(esc(dom_cat)), $(esc(codom_cat)), $(Meta.quot(body))))
end

function parse_functor(C::FinCat, D::FinCat, body::Expr)
  ob_rhs, hom_rhs = parse_ob_hom_maps(C, body)
  F = FinFunctor(mapvals(x -> parse_ob(D, x), ob_rhs),
                 mapvals(f -> parse_hom(D, f), hom_rhs), C, D)
  is_functorial(F, check_equations=false) ||
    error("Parsed functor is not functorial: $body")
  return F
end
function parse_functor(C::Presentation, D::Presentation, body::Expr; kw...)
  parse_functor(FinCat(C), FinCat(D), body; kw...)
end

function parse_ob_hom_maps(C::FinCat, body::Expr; allow_missing::Bool=false)
  assignments = Dict{Symbol,Union{Expr,Symbol}}()
  assign(lhs, rhs) = if haskey(assignments, lhs)
    error("Left-hand side $lhs assigned twice in $body")
  else
    assignments[lhs] = rhs
  end
  for stmt in statements(body)
    @match stmt begin
      # x => y
      Expr(:call, :(=>), lhs::Symbol, rhs) => assign(lhs, rhs)
      # (x, x′, ...) => y
      Expr(:call, :(=>), Expr(:tuple, args...), rhs) =>
        foreach(lhs -> assign(lhs, rhs), args)
      ::LineNumberNode => nothing
      _ => error("Invalid assignment statement $stmt")
    end
  end
  ob_rhs = make_map(ob_generators(C)) do x
    y = pop!(assignments, ob_generator_name(C, x), missing)
    (!ismissing(y) || allow_missing) ? y :
      error("Object $(ob_generator_name(C,x)) is not assigned")
  end
  hom_rhs = make_map(hom_generators(C)) do f
    g = pop!(assignments, hom_generator_name(C, f), missing)
    (!ismissing(g) || allow_missing) ? g :
      error("Morphism $(hom_generator_name(C,f)) is not assigned")
  end
  isempty(assignments) ||
    error(string("Unused assignment(s): ", join(keys(assignments), ", ")))
  (ob_rhs, hom_rhs)
end

""" Parse expression for object in a category.
"""
function parse_ob(C::FinCat{Ob,Hom}, expr) where {Ob,Hom}
  @match expr begin
    x::Symbol => ob_generator(C, x)
    Expr(:curly, _...) => parse_gat_expr(C, expr)::Ob
    _ => error("Invalid object expression $expr")
  end
end

""" Parse expression for morphism in a category.
"""
function parse_hom(C::FinCat{Ob,Hom}, expr) where {Ob,Hom}
  function parse(expr)
    @match expr begin
      Expr(:call, :compose, args...) =>
        mapreduce(parse, (fs...) -> compose(C, fs...), args)
      Expr(:call, :(⋅), f, g) ||
      Expr(:call, :(⨟), f, g) => compose(C, parse(f), parse(g))
      Expr(:call, :(∘), f, g) => compose(C, parse(g), parse(f))
      Expr(:call, :id, x::Symbol) => id(C, ob_generator(C, x))
      f::Symbol => hom_generator(C, f)
      Expr(:curly, _...) => parse_gat_expr(C, expr)::Hom
      _ => error("Invalid morphism expression $expr")
    end
  end
  parse(expr)
end

""" Parse GAT expression based on curly braces, rather than parentheses.
"""
function parse_gat_expr(C::FinCat, root_expr)
  pres = presentation(C)
  function parse(expr)
    @match expr begin
      Expr(:curly, head::Symbol, args...) =>
        invoke_term(pres.syntax, head, map(parse, args)...)
      x::Symbol => generator(pres, x)
      _ => error("Invalid GAT expression $root_expr")
    end
  end
  parse(root_expr)
end

# Diagrams
##########

""" Present a diagram in a given category.

Recall that a *diagram* in a category ``C`` is a functor ``F: J → C`` from a
small category ``J`` into ``C``. Given the category ``C``, this macro presents a
diagram in ``C``, i.e., constructs a finitely presented indexing category ``J``
together with a functor ``F: J → C``. This method of simultaneous definition is
often more convenient than defining ``J`` and ``F`` separately, as could be
accomplished by calling [`@fincat`](@ref) and then [`@finfunctor`](@ref).

As an example, the limit of the following diagram consists of the paths of
length two in a graph:

```julia
@diagram TheoryGraph begin
  v::V
  (e₁, e₂)::E
  (t: e₁ → v)::tgt
  (s: e₂ → v)::src
end
```

Morphisms in the indexing category can be left unnamed, which is convenient for
defining free diagrams (see also [`@free_diagram`](@ref)). For example, the
following diagram is isomorphic to the previous one:

```julia
@diagram TheoryGraph begin
  v::V
  (e₁, e₂)::E
  (e₁ → v)::tgt
  (e₂ → v)::src
end
```

Of course, unnamed morphisms cannot be referenced by name within the `@diagram`
call or in other settings, which can sometimes be problematic.
"""
macro diagram(cat, body)
  :(parse_diagram($(esc(cat)), $(Meta.quot(body))))
end

""" Present a free diagram in a given category.

Recall that a *free diagram* in a category ``C`` is a functor ``F: J → C`` where
``J`` is a free category on a graph, here assumed finite. This macro is
functionally a special case of [`@diagram`](@ref) but, for convenience, changes
the interpretation of equality expressions. Rather than interpreting them as
equations between morphisms in ``J``, equality expresions can be used to
introduce anonymous morphisms in a "pointful" style. For example, the limit of
the following diagram consists of the paths of length two in a graph:

```julia
@free_diagram TheoryGraph begin
  v::V
  (e₁, e₂)::E
  tgt(e₁) == v
  src(e₂) == v
end
```

Anonymous objects can also be introduced. For example, the previous diagram is
isomorphic to this one:

```julia
@free_diagram TheoryGraph begin
  (e₁, e₂)::E
  tgt(e₁) == src(e₂)
end
```

Some care must exercised when defining morphisms between diagrams with anonymous
objects, since they cannot be referred to by name.
"""
macro free_diagram(cat, body)
  :(parse_diagram($(esc(cat)), $(Meta.quot(body)), free=true))
end

function parse_diagram(C::FinCat, body::Expr; kw...)
  F_ob, F_hom, J = parse_diagram_data(
    C, x -> parse_ob(C,x), (f,x,y) -> parse_hom(C,f), body; kw...)
  F = FinFunctor(F_ob, F_hom, J, C)
  is_functorial(F, check_equations=false) ||
    error("Parsed diagram is not functorial: $body")
  return F
end
parse_diagram(pres::Presentation, body::Expr; kw...) =
  parse_diagram(FinCat(pres), body; kw...)

function parse_diagram_data(C::FinCat, parse_ob, parse_hom, body::Expr;
                            free::Bool=false, preprocess::Bool=true)
  g, eqs = NamedGraph{Symbol,Union{Symbol,Nothing}}(), Pair[]
  F_ob, F_hom = [], []
  function push_hom!(h, x, y; name=nothing)
    e = parse_edge!(g, x, y, ename=name)
    push!(F_hom, parse_hom(h, F_ob[src(g,e)], F_ob[tgt(g,e)]))
    return e
  end

  nanon = 0
  if preprocess
    body = reparse_arrows(body)
  end
  for stmt in statements(body)
    @match stmt begin
      # x => X
      # x::X
      Expr(:call, :(=>), x::Symbol, X) || Expr(:(::), x::Symbol, X) => begin
        add_vertex!(g, vname=x)
        push!(F_ob, parse_ob(X))
      end
      # (x, y, ...) => X
      # (x, y, ...)::X
      Expr(:call, :(=>), Expr(:tuple, xs...), X) ||
      Expr(:(::), Expr(:tuple, xs...), X) => begin
        add_vertices!(g, length(xs), vname=xs)
        append!(F_ob, repeated(parse_ob(X), length(xs)))
      end
      # (f: x → y) => h
      # (f: x → y)::h
      Expr(:call, :(=>), Expr(:call, :(:), f::Symbol,
                              Expr(:call, :(→), x::Symbol, y::Symbol)), h) ||
      Expr(:(::), Expr(:call, :(:), f::Symbol,
                       Expr(:call, :(→), x::Symbol, y::Symbol)), h) =>
        push_hom!(h, x, y, name=f)
      # (x → y) => h
      # (x → y)::h
      Expr(:call, :(=>), Expr(:call, :(→), x::Symbol, y::Symbol), h) ||
      Expr(:(::), Expr(:call, :(→), x::Symbol, y::Symbol), h) =>
        push_hom!(h, x, y)
      # h(x) == y
      # y == h(x)
      (Expr(:call, :(==), call::Expr, y::Symbol) ||
       Expr(:call, :(==), y::Symbol, call::Expr)) && if free end =>
         push_hom!(destructure_unary_call(call)..., y)
      # h(x) == k(y)
      Expr(:call, :(==), lhs::Expr, rhs::Expr) && if free end => begin
        (h, x), (k, y) = destructure_unary_call(lhs), destructure_unary_call(rhs)
        z = Symbol("##unnamed#$(nanon += 1)")
        add_vertex!(g, vname=z)
        push!(F_ob, nothing) # Assumes that codomain not needed to parse homs.
        e1, e2 = push_hom!(h, x, z), push_hom!(k, y, z)
        # Infer codomain from parsed homs.
        Z1, Z2 = codom(C, F_hom[e1]), codom(C, F_hom[e2])
        Z1 == Z2 || error("Codomain mismatch in $stmt: $Z1 != $Z2")
        F_ob[end] = Z1
      end
      # f == g
      Expr(:call, :(==), lhs, rhs) && if !free end =>
        push!(eqs, parse_path_equation(g, lhs, rhs))
      ::LineNumberNode => nothing
      _ => error("Cannot parse line in diagram definition: $stmt")
    end
  end
  J = isempty(eqs) ? FinCat(g) : FinCat(g, eqs)
  (F_ob, F_hom, J)
end

# Data migrations
#################

""" A diagram without a codomain category.

An intermediate data representation used internally by the parser for the
[`@migration`](@ref) macro.
"""
struct DiagramData{T,ObMap,HomMap,Shape<:FinCat}
  ob_map::ObMap
  hom_map::HomMap
  shape::Shape

  function DiagramData{T}(ob_map::ObMap, hom_map::HomMap, shape::Shape) where
      {T,ObMap,HomMap,Shape<:FinCat}
    new{T,ObMap,HomMap,Shape}(ob_map, hom_map, shape)
  end
end

Diagrams.ob_map(d::DiagramData, x) = d.ob_map[x]
Diagrams.hom_map(d::DiagramData, f) = d.hom_map[f]
Diagrams.shape(d::DiagramData) = d.shape

""" A diagram morphism without a domain or codomain.

Like [`DiagramData`](@ref), an intermediate data representation used internally
by the parser for the [`@migration`](@ref) macro.
"""
struct DiagramHomData{T,ObMap,HomMap}
  ob_map::ObMap
  hom_map::HomMap

  function DiagramHomData{T}(ob_map::ObMap, hom_map::HomMap) where {T,ObMap,HomMap}
    new{T,ObMap,HomMap}(ob_map, hom_map)
  end
end

""" Contravariantly migrate data from one acset to another.

This macro is shorthand for defining a data migration using the
[`@migration`](@ref) macro and then calling the `migrate` function. If the
migration will be used multiple times, it is more efficient to perform these
steps separately, reusing the functor defined by `@migration`.

For more about the syntax and supported features, see [`@migration`](@ref).
"""
macro migrate(tgt_type, src_acset, body)
  quote
    let T = $(esc(tgt_type)), X = $(esc(src_acset))
      migrate(T, X, parse_migration(Presentation(T), Presentation(X),
                                    $(Meta.quot(body))))
    end
  end
end

""" Define a contravariant data migration.

This macro provides a DSL to specify a contravariant data migration from
``C``-sets to ``D``-sets for given schemas ``C`` and ``D``. A data migration is
defined by a functor from ``D`` to a category of queries on ``C``. Thus, every
object of ``D`` is assigned a query on ``C`` and every morphism of ``D`` is
assigned a morphism of queries, in a compatible way. Example usages are in the
unit tests. What follows is a technical reference.

Several categories of queries are supported by this macro:

1. Trivial queries, specified by a single object of ``C``. In this case, the
   macro simply defines a functor ``D → C`` and is equivalent to
   [`@finfunctor`](@ref) or [`@diagram`](@ref).
2. *Conjunctive queries*, specified by a diagram in ``C`` and evaluated as a
   finite limit.
3. *Gluing queries*, specified by a diagram in ``C`` and evaluated as a finite
   colimit. An important special case is *linear queries*, evaluated as a
   finite coproduct.
4. *Gluc queries* (gluings of conjunctive queries), specified by a diagram of
   diagrams in ``C`` and evaluated as a colimit of limits. An important special
   case is *duc queries* (disjoint unions of conjunctive queries), evaluated as
   a coproduct of limits.

The query category of the data migration is not specified explicitly but is
inferred from the queries used in the macro call. Implicit conversion is
performed: trivial queries can be coerced to conjunctive queries or gluing
queries, and conjunctive queries and gluing queries can both be coerced to gluc
queries. Due to the implicit conversion, the resulting functor out of ``D`` has
a single query type and thus a well-defined codomain.

Syntax for the right-hand sides of object assignments is:

- a symbol, giving object of ``C`` (query type: trivial)
- `@product ...` (query type: conjunctive)
- `@join ...` (alias: `@limit`, query type: conjunctive)
- `@cases ...` (alias: `@coproduct`, query type: gluing or gluc)
- `@glue ...` (alias: `@colimit`, query type: gluing or gluc)

Thes query types supported by this macro generalize the kind of queries familiar
from relational databases. Less familiar is the concept of a morphism between
queries, derived from the concept of a morphism between diagrams in a category.
A query morphism is given by a functor between the diagrams' indexing categories
together with a natural transformation filling a triangle of the appropriate
shape. From a practical standpoint, the most important thing to remember is that
a morphism between conjunctive queries is contravariant with respect to the
diagram shapes, whereas a morphism between gluing queries is covariant. TODO:
Reference for more on this.
"""
macro migration(src_schema, body)
  :(parse_migration($(esc(src_schema)), $(Meta.quot(body))))
end
macro migration(tgt_schema, src_schema, body)
  :(parse_migration($(esc(tgt_schema)), $(esc(src_schema)), $(Meta.quot(body))))
end

""" Parse a contravariant data migration from a Julia expression.

The process kicked off by this internal function is somewhat complicated due to
the need to coerce queries and query morphisms to a common category. The
high-level steps of this process are:

1. Parse the queries and query morphisms into intermediate representations
   ([`DiagramData`](@ref) and [`DiagramHomData`](@ref)) whose final types are
   not yet determined.
2. Promote the query types to the tightest type encompassing all queries, an
   approach reminiscent of Julia's own type promotion system.
3. Convert all query and query morphisms to this common type, yielding `Diagram`
   and `DiagramHom` instances.
"""
function parse_migration(src_schema::Presentation, body::Expr;
                         preprocess::Bool=true)
  C = FinCat(src_schema)
  F_ob, F_hom, J = parse_query_diagram(C, body; free=false, preprocess=preprocess)
  make_migration_functor(F_ob, F_hom, J, C)
end
function parse_migration(tgt_schema::Presentation, src_schema::Presentation,
                         body::Expr; preprocess::Bool=true)
  D, C = FinCat(tgt_schema), FinCat(src_schema)
  if preprocess
    body = reparse_arrows(body)
  end
  ob_rhs, hom_rhs = parse_ob_hom_maps(D, body)
  F_ob = mapvals(expr -> parse_query(C, expr), ob_rhs)
  F_hom = mapvals(hom_rhs, keys=true) do f, expr
    parse_query_hom(C, expr, F_ob[dom(D,f)], F_ob[codom(D,f)])
  end
  make_migration_functor(F_ob, F_hom, D, C)
end

# Query parsing
#--------------

""" Parse expression defining a query.
"""
function parse_query(C::FinCat, expr)
  @match expr begin
    x::Symbol => ob_generator(C, x)
    Expr(:macrocall, form, args...) &&
        if form ∈ (Symbol("@limit"), Symbol("@join")) end => begin
      DiagramData{op}(parse_query_diagram(C, last(args))...)
    end
    Expr(:macrocall, form, args...) &&
        if form == Symbol("@product") end => begin
      d = DiagramData{op}(parse_query_diagram(C, last(args))...)
      is_discrete(shape(d)) ? d : error("Product query is not discrete: $expr")
    end
    Expr(:macrocall, form, args...) &&
        if form ∈ (Symbol("@colimit"), Symbol("@glue")) end => begin
      DiagramData{id}(parse_query_diagram(C, last(args))...)
    end
    Expr(:macrocall, form, args...) &&
        if form ∈ (Symbol("@coproduct"), Symbol("@cases")) end => begin
      d = DiagramData{id}(parse_query_diagram(C, last(args))...)
      is_discrete(shape(d)) ? d : error("Cases query is not discrete: $expr")
    end
    _ => error("Cannot parse query in migration definition: $expr")
  end
end
function parse_query_diagram(C::FinCat, expr::Expr;
                             free::Bool=true, preprocess::Bool=false)
  parse_diagram_data(C, X -> parse_query(C,X),
                     (f,x,y) -> parse_query_hom(C,f,x,y), expr;
                     free=free, preprocess=preprocess)
end

""" Parse expression defining a morphism of queries.
"""
parse_query_hom(C::FinCat{Ob}, expr, ::Ob, ::Union{Ob,Nothing}) where Ob =
  parse_hom(C, expr)

# Conjunctive fragment.

function parse_query_hom(C::FinCat, expr, d::DiagramData{op}, d′::DiagramData{op})
  ob_rhs, hom_rhs = parse_ob_hom_maps(shape(d′), expr)
  f_ob = mapvals(ob_rhs, keys=true) do j′, rhs
    parse_conj_query_ob_rhs(C, rhs, d, ob_map(d′, j′))
  end
  f_hom = mapvals(rhs -> parse_hom(shape(d), rhs), hom_rhs)
  DiagramHomData{op}(f_ob, f_hom)
end
function parse_query_hom(C::FinCat{Ob}, expr, c::Ob, d′::DiagramData{op}) where Ob
  ob_rhs, f_hom = parse_ob_hom_maps(shape(d′), expr, allow_missing=true)
  f_ob = mapvals(ob_rhs, keys=true) do j′, rhs
    ismissing(rhs) ? missing : (missing, parse_query_hom(C, rhs, c, ob_map(d′, j′)))
  end
  @assert all(ismissing, f_hom)
  DiagramHomData{op}(f_ob, f_hom)
end
function parse_query_hom(C::FinCat{Ob}, expr, d::DiagramData{op}, c′::Ob) where Ob
  DiagramHomData{op}([parse_conj_query_ob_rhs(C, expr, d, c′)], [])
end

# Gluing fragment.

function parse_query_hom(C::FinCat, expr, d::DiagramData{id}, d′::DiagramData{id})
  ob_rhs, hom_rhs = parse_ob_hom_maps(shape(d), expr)
  f_ob = mapvals(ob_rhs, keys=true) do j, rhs
    parse_glue_query_ob_rhs(C, rhs, ob_map(d, j), d′)
  end
  f_hom = mapvals(expr -> parse_hom(shape(d′), expr), hom_rhs)
  DiagramHomData{id}(f_ob, f_hom)
end
function parse_query_hom(C::FinCat{Ob}, expr, c::Union{Ob,DiagramData{op}},
                         d′::DiagramData{id}) where Ob
  DiagramHomData{id}([parse_glue_query_ob_rhs(C, expr, c, d′)], [])
end
function parse_query_hom(C::FinCat{Ob}, expr, d::DiagramData{id},
                         c′::Union{Ob,DiagramData{op}}) where Ob
  ob_rhs, f_hom = parse_ob_hom_maps(shape(d), expr, allow_missing=true)
  f_ob = mapvals(ob_rhs, keys=true) do j, rhs
    ismissing(rhs) ? missing : (missing, parse_query_hom(C, rhs, ob_map(d, j), c′))
  end
  @assert all(ismissing, f_hom)
  DiagramHomData{id}(f_ob, f_hom)
end

""" Parse RHS of object assignment in morphism out of conjunctive query.
"""
function parse_conj_query_ob_rhs(C::FinCat, expr, d::DiagramData{op}, c′)
  j_name, f_expr = @match expr begin
    x::Symbol => (x, nothing)
    Expr(:tuple, x::Symbol, f) => (x, f)
    Expr(:call, op, _...) && if op ∈ compose_ops end =>
      leftmost_arg(expr, (:(⋅), :(⨟)), all_ops=compose_ops)
    Expr(:call, name::Symbol, _) => reverse(destructure_unary_call(expr))
    _ => error("Cannot parse object assignment in migration: $expr")
  end
  j = ob_generator(shape(d), j_name)
  isnothing(f_expr) ? j :
    (j, parse_query_hom(C, f_expr, ob_map(d, j), c′))
end

""" Parse RHS of object assignment in morphism into gluing query.
"""
function parse_glue_query_ob_rhs(C::FinCat, expr, c, d′::DiagramData{id})
  j′_name, f_expr = @match expr begin
    x::Symbol => (x, nothing)
    Expr(:tuple, x::Symbol, f) => (x, f)
    Expr(:call, op, _...) && if op ∈ compose_ops end =>
      leftmost_arg(expr, (:(∘),), all_ops=compose_ops)
    _ => error("Cannot parse object assignment in migration: $expr")
  end
  j′ = ob_generator(shape(d′), j′_name)
  isnothing(f_expr) ? j′ :
    (j′, parse_query_hom(C, f_expr, c, ob_map(d′, j′)))
end

const compose_ops = (:(⋅), :(⨟), :(∘))

# Query construction
#-------------------

function make_migration_functor(F_ob, F_hom, D::FinCat, C::FinCat)
  diagram(make_query(C, DiagramData{id}(F_ob, F_hom, D)))
end

function make_query(C::FinCat{Ob}, d::DiagramData{T}) where {T, Ob}
  F_ob, F_hom, J = d.ob_map, d.hom_map, shape(d)
  F_ob = mapvals(x -> make_query(C, x), F_ob)
  query_type = mapreduce(typeof, promote_query_type, values(F_ob), init=Ob)
  @assert query_type != Any
  F_ob = mapvals(x -> convert_query(C, query_type, x), F_ob)
  F_hom = mapvals(F_hom, keys=true) do h, f
    make_query_hom(f, F_ob[dom(J,h)], F_ob[codom(J,h)])
  end
  Diagram{T}(if query_type <: Ob
    FinFunctor(F_ob, F_hom, shape(d), C)
  else
    # XXX: Why is the element type of `F_ob` sometimes too loose?
    D = TypeCat(typeintersect(query_type, eltype(values(F_ob))),
                eltype(values(F_hom)))
    FinDomFunctor(F_ob, F_hom, shape(d), D)

  end)
end

make_query(C::FinCat{Ob}, x::Ob) where Ob = x

function make_query_hom(f::DiagramHomData{op}, d::Diagram{op}, d′::Diagram{op})
  f_ob = mapvals(f.ob_map, keys=true) do j′, x
    x = @match x begin
      ::Missing => only_ob(shape(d))
      (::Missing, g) => (only_ob(shape(d)), g)
      _ => x
    end
    @match x begin
      (j, g) => Pair(j, make_query_hom(g, ob_map(d, j), ob_map(d′, j′)))
      j => j
    end
  end
  f_hom = mapvals(h -> ismissing(h) ? only_hom(shape(d)) : h, f.hom_map)
  DiagramHom{op}(f_ob, f_hom, d, d′)
end

function make_query_hom(f::DiagramHomData{id}, d::Diagram{id}, d′::Diagram{id})
  f_ob = mapvals(f.ob_map, keys=true) do j, x
    x = @match x begin
      ::Missing => only_ob(shape(d′))
      (::Missing, g) => (only_ob(shape(d′)), g)
      _ => x
    end
    @match x begin
      (j′, g) => Pair(j′, make_query_hom(g, ob_map(d, j), ob_map(d′, j′)))
      j′ => j′
    end
  end
  f_hom = mapvals(h -> ismissing(h) ? only_hom(shape(d′)) : h, f.hom_map)
  DiagramHom{id}(f_ob, f_hom, d, d′)
end

function make_query_hom(f::Hom, d::Diagram{T,C}, d′::Diagram{T,C}) where
    {T, Ob, Hom, C<:FinCat{Ob,Hom}}
  cat = codom(diagram(d))
  munit(DiagramHom{T}, cat, f, dom_shape=shape(d), codom_shape=shape(d′))
end
make_query_hom(f, x, y) = f

only_ob(C::FinCat) = only(ob_generators(C))
only_hom(C::FinCat) = (@assert is_discrete(C); id(C, only_ob(C)))

# Query promotion
#----------------

# Promotion of query types is modeled loosely on Julia's type promotion system:
# https://docs.julialang.org/en/v1/manual/conversion-and-promotion/

promote_query_rule(::Type, ::Type) = Union{}
promote_query_rule(::Type{<:ConjQuery{C}}, ::Type{<:Ob}) where {Ob,C<:FinCat{Ob}} =
  ConjQuery{C}
promote_query_rule(::Type{<:GlueQuery{C}}, ::Type{<:Ob}) where {Ob,C<:FinCat{Ob}} =
  GlueQuery{C}
promote_query_rule(::Type{<:GlucQuery{C}}, ::Type{<:Ob}) where {Ob,C<:FinCat{Ob}} =
  GlucQuery{C}
promote_query_rule(::Type{<:GlucQuery{C}}, ::Type{<:ConjQuery{C}}) where C =
  GlucQuery{C}
promote_query_rule(::Type{<:GlucQuery{C}}, ::Type{<:GlueQuery{C}}) where C =
  GlucQuery{C}

promote_query_type(T, S) = promote_query_result(
  T, S, Union{promote_query_rule(T,S), promote_query_rule(S,T)})
promote_query_result(T, S, ::Type{Union{}}) = typejoin(T, S)
promote_query_result(T, S, U) = U

convert_query(::FinCat, ::Type{T}, x::S) where {T, S<:T} = x

function convert_query(cat::C, ::Type{<:Diagram{T,C}}, x::Ob) where
  {T, Ob, C<:FinCat{Ob}}
  g = NamedGraph{Symbol,Symbol}(1, vname=nameof(x))
  munit(Diagram{T}, cat, x, shape=FinCat(g))
end
function convert_query(::C, ::Type{<:GlucQuery{C}}, d::ConjQuery{C}) where C
  munit(Diagram{id}, TypeCat(ConjQuery{C}, Any), d)
end
function convert_query(cat::C, ::Type{<:GlucQuery{C}}, d::GlueQuery{C}) where C
  J = shape(d)
  new_ob = make_map(ob_generators(J)) do j
    convert_query(cat, ConjQuery{C}, ob_map(d, j))
  end
  new_hom = make_map(hom_generators(J)) do h
    munit(Diagram{op}, cat, hom_map(d, h),
          dom_shape=new_ob[dom(J,h)], codom_shape=new_ob[codom(J,h)])
  end
  Diagram{id}(FinDomFunctor(new_ob, new_hom, J))
end
function convert_query(cat::C, ::Type{<:GlucQuery{C}}, x::Ob) where
    {Ob, C<:FinCat{Ob}}
  convert_query(cat, GlucQuery{C}, convert_query(cat, ConjQuery{C}, x))
end

# Utilities
###########

""" Left-most argument plus remainder of left-associated binary operations.
"""
function leftmost_arg(expr, ops; all_ops=nothing)
  isnothing(all_ops) && (all_ops = ops)
  function leftmost(expr)
    @match expr begin
      Expr(:call, op2, Expr(:call, op1, x, y), z) &&
          if op1 ∈ all_ops && op2 ∈ all_ops end => begin
        x, rest = leftmost(Expr(:call, op1, x, y))
        (x, Expr(:call, op2, rest, z))
      end
      Expr(:call, op, x, y) && if op ∈ ops end => (x, y)
      _ => (nothing, expr)
    end
  end
  leftmost(expr)
end

""" Destructure the expression `:(f(g(x)))` to `(:(f∘g), :x)`, for example.
"""
function destructure_unary_call(expr::Expr)
  @match expr begin
    Expr(:call, head, x::Symbol) => (head, x)
    Expr(:call, head, arg) => begin
      rest, x = destructure_unary_call(arg)
      (Expr(:call, :(∘), head, rest), x)
    end
  end
end

end


### A Pluto.jl notebook ###
# v0.12.7

using Markdown
using InteractiveUtils

# ╔═╡ de7a5e8a-1ee4-11eb-2bbf-0f0eb3e489ec
using Pkg, DrWatson

# ╔═╡ e4eeb13a-1ee4-11eb-0113-91f9a9c3b659
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Random, DataFrames
	Random.seed!(1)
end;

# ╔═╡ ac3ceabe-1ee4-11eb-3935-313e96aafd62
md"## Listing 2.12"

# ╔═╡ ab1ec376-1f63-11eb-3eb3-d7f282f2ad65
begin
	N = 10^5
	prob1 = 0.7
	eps0, eps1 = 0.1, 0.05
end

# ╔═╡ faad978c-1f63-11eb-189c-bf9767f58cae
flipWithProb(bit,prob) = rand() < prob ? xor(bit,1) : bit

# ╔═╡ faade49e-1f63-11eb-3c17-5f937cb796ca
begin
	TxData = rand(N) .< prob1
	RxData = [x == 0 ? flipWithProb(x,eps0) : flipWithProb(x,eps1) for x in TxData]

	numTx1 = 0
	totalRx1 = 0
	for i in 1:N
	   if RxData[i] == 1
			global totalRx1 += 1
			global numTx1 += TxData[i]
		end
	end
end

# ╔═╡ faae7008-1f63-11eb-2811-914bae0ac42d
monteCarlo = numTx1/totalRx1

# ╔═╡ fabb6902-1f63-11eb-1a72-c3d127ec17c2
analytic = ((1-eps1)*0.7)/((1-eps1)*0.7+0.3*eps0)

# ╔═╡ fabbef8a-1f63-11eb-2122-3791b9449669
Text("Monte Carlo: $monteCarlo Analytic: $analytic")

# ╔═╡ ea086874-1eeb-11eb-2cb0-4186697d4294
md"## End of listing 2.12"

# ╔═╡ Cell order:
# ╟─ac3ceabe-1ee4-11eb-3935-313e96aafd62
# ╠═de7a5e8a-1ee4-11eb-2bbf-0f0eb3e489ec
# ╠═e4eeb13a-1ee4-11eb-0113-91f9a9c3b659
# ╠═ab1ec376-1f63-11eb-3eb3-d7f282f2ad65
# ╠═faad978c-1f63-11eb-189c-bf9767f58cae
# ╠═faade49e-1f63-11eb-3c17-5f937cb796ca
# ╠═faae7008-1f63-11eb-2811-914bae0ac42d
# ╠═fabb6902-1f63-11eb-1a72-c3d127ec17c2
# ╠═fabbef8a-1f63-11eb-2122-3791b9449669
# ╟─ea086874-1eeb-11eb-2cb0-4186697d4294


check_prob_array(m :: AbstractArray{F1}) where F1 <: Real =
    check_float_array(m, zero(F1), one(F1));

"""
    $(SIGNATURES)

Given an array of probabilities: ensure that all are in [0, 1].
Error if bounds violation larger than rounding errors.
Make sure that sum does not exceed an upper bound.
"""
function make_valid_probs!(m :: AbstractArray{F1}; maxSum :: F1 = one(F1)) where F1 <: Real

    fSmall = F1(.0000001);
    pSum = sum(m);
    @assert (pSum < one(F1) + fSmall)  "Sum too large: $pSum"
    @assert all_at_least(m, -fSmall)  "Negative probabilities"
    @assert all_at_most(m, one(F1) + fSmall)  "Probabilities above 1"

    if any(x -> x < zero(F1), m)
        m[findall(x -> x < zero(F1), m)] .= zero(F1);
    end
    if pSum > maxSum
        m .*= ((maxSum - fSmall) / pSum);
    end
end

# Make a `Vector` into an `n x 1` Matrix.
# Don't use `reshape` b/c that produces a reshaped array instead.
matrix_from_vector(v :: AbstractVector{F1}) where F1 =
    repeat(v, outer = (1,1));

## ----------  Choice without pref shocks

# Given value work, value of colleges, vector of admitted colleges, return max value and college index (not among admitted ones)
function max_choice(vWork :: F1, vCollegeV :: AbstractVector{F1},
    admitV :: AbstractVector{Bool}) where {F1 <: Real}

    v, idx = findmax(vcat(vWork, vCollegeV .- (.!admitV) .* F1(1e8)));
    ic = idx - 1;
    return v, ic
end

# The same with an index vector as `admitV`
function max_choice(vWork :: F1, vCollegeV :: AbstractVector{F1},
    admitIdxV :: AbstractVector{I1}) where {F1 <: Real, I1 <: Integer}

    admitV = falses(size(vCollegeV));
    admitV[admitIdxV] .= true;
    return max_choice(vWork, vCollegeV, admitV)
end

# The same for many agents
function max_choices(vWork_jV :: AbstractVector{F1},
    vCollege_jcM :: AbstractMatrix{F1},
    admitV :: AbstractVector{I1}) where {F1 <: Real, I1}

    J = length(vWork_jV);
    valV = zeros(F1, J);
    icV = zeros(Int, J);
    for j = 1 : J
        valV[j], icV[j] = max_choice(vWork_jV[j], vec(vCollege_jcM[j,:]),
            admitV);
    end
    return valV, icV
end

function find_class(x :: F1, cutoffV) where F1 <: Real
    if x > last(cutoffV)
        j = length(cutoffV) + 1;
    elseif x < first(cutoffV)
        j = 1;
    else
        j = findfirst(cutoff -> x < cutoff, cutoffV);
    end
    return j
end

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


### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 82467e1e-97e0-11eb-35bc-b70fe47f7296
using Revise, FourierTools, ImageCore, TestImages, PlutoUI, ImageShow

# ╔═╡ 8f0854f4-97e0-11eb-14b8-7958b0ab243f
begin
	img = Float32.(testimage("fabio_512_gray"))
	z = zeros(Float32, (768, 768))
	FourierTools.center_set!(z, img)
end

# ╔═╡ e43962ba-97e0-11eb-274f-fd6ed85b03ff
@bind ϕ Slider(-180:1:180)

# ╔═╡ dc4917a8-97e0-11eb-341d-25ec7207b748
md"ϕ=$ϕ"

# ╔═╡ a9e968a8-97e0-11eb-0486-3f75f68be3c7
Gray.(FourierTools.rotate(z, ϕ))

# ╔═╡ f1fecc74-97e1-11eb-203e-2949bafae27d
Gray.(FourierTools.rotate(img, ϕ))

# ╔═╡ Cell order:
# ╠═82467e1e-97e0-11eb-35bc-b70fe47f7296
# ╠═8f0854f4-97e0-11eb-14b8-7958b0ab243f
# ╠═dc4917a8-97e0-11eb-341d-25ec7207b748
# ╠═e43962ba-97e0-11eb-274f-fd6ed85b03ff
# ╠═a9e968a8-97e0-11eb-0486-3f75f68be3c7
# ╠═f1fecc74-97e1-11eb-203e-2949bafae27d


########################
### TESTING ROUTINES ###
########################

# TEST THE EXCHANGE ALGORITHM ON RANDOM DATA
#
# This function simulates random uncorrelated Gaussian data to test the exchange algorithm.
#
# Arguments:
# -- n is the desired number of samples.
# -- p is the desired number of predictors.
# -- r is the desired number of true nonzero components, or the size of the true model.
# -- extra is the amount of extra model sizes to compute beyond r. Thus, the length of the computed
#    regularization path is r + extra.
#
# Optional Arguments:
# -- max_iter caps the number of outer iterations in the exchange algorithm. Defaults to 100.
# -- quiet is a Bool to control algorithm output to the console. Defaults to false (print output).
# -- window is the number of predictors to consider for swapping. Defaults to r + extra, which tells the algorithm
#    to consider all active predictors for swapping.
# -- seed sets the random seed for the data generation. Defaults to 2015.
# -- noise indicates the magnitude of the random noise added to the response y. Defaults to 0.5, which indicates
#    noise drawn from the N(0,0.25) distribution
#
# coded by Kevin L. Keys (2015)
# klkeys@g.ucla.edu
function test_exleastsq(n::Int, p::Int, r::Int, extra::Int; max_iter::Int = 100, quiet::Bool = false, window::Int = r+extra, seed::Int = 2015, noise::Float64 = 0.5)
	tic()

	# testing data
	x         = randn(n,p)
	y         = randn(n)
	b         = zeros(p)
	perm      = collect(1:p)
	inner     = Dict{Int,DenseArray{Float64,1}}()

	# declare all temporary arrays
	res        = zeros(n)	# Y - Xbeta
	df         = zeros(p)	# X'(Y - Xbeta)
	tempn      = zeros(n)	# temporary array of length n
	tempn2     = zeros(n)	# temporary array of length n
	tempp      = zeros(p)	# temporary array of length p
	dotprods   = zeros(p)	# hold in memory the dot products for current index

	# precompute sum of squares for each column of x
	const nrmsq = vec(sumabs2(x,1))

	# declare return array
	betas = zeros(p,r+extra)

	# use a random permutation of indices for true model
	true_model = sample(perm, r)
	b[true_model] = randn(r)
	b_true = copy(b)
	BLAS.gemv!('N', 1.0, x, b, 1.0, y)	# y = xb + noise
	fill!(b, 0.0)
	iter = 0

	# print time spent on parameter initialization
	quiet || println("Parameter initialization took ", toq(), " seconds.")

	# reset timer and compute path
	tic()
	@time begin
		for k = 1:(r+extra)
			b = exchange_leastsq!(b, x, y, perm, k, inner=inner, max_iter=max_iter, quiet=quiet, nrmsq=nrmsq, res=res, df=df, tempn=tempn, tempn2=tempn2, tempp=tempp, dotprods=dotprods, window=min(k,window))
			rss = sumabs2(res)

			# print results from this step of regularization path
			if !quiet
				println("#pred = ", k, ", rss = ", rss, ", #truepos = ", countnz(b[true_model]))
				print("\t")
				print("\n")
			end
			betas[:,k] = b
			update_col!(betas, b, k, n=p, p=r+extra, a=1.0)
		end
	end

	# print time spent on computation of regularization path
	quiet || println("Path from 1 to ", r+extra, " took ", toq(), " seconds to compute.")

	# return sparsified betas
	return b_true, sparse(betas)
end

# TEST THE EXCHANGE ALGORITHM FOR LEAST SQUARES REGRESSION
# This subroutine runs the IHT on a dataset of chromosome 1 data from the WTCCC.
function test_exchangeleastsq(x_path::String = "/Users/kkeys/Downloads/wtccc-n2000p32307.txt", y_path::String = "/Users/kkeys/Downloads/withnoiselevelsd0_1/Y.100.1", b_path::String = "/Users/kkeys/Downloads/withnoiselevelsd0_1/causal.100.1", r::Int = 100, tol::Float64 = 1e-6, max_iter::Int = 100, quiet::Bool = false, extra::Int = 100, window::Int = r)

    # repeat input options
    println("This function will test the exchange algorithm.")
    println("Given options:")
    println("\tPath to X  = ", x_path)
    println("\tPath to Y  = ", y_path)
    println("\tPath to B  = ", b_path)
    println("\tModel size = ", r)
    println("\ttolerance  = ", tol, "\n\tmax_iter   = ", max_iter)

    # precompile @time macro for later use
    println("\nCompiling @time macro...")
    @time 1+1;

    # now load data
    println("\nLoading data...")
    tic()

    # load design matrix
    x = readdlm(x_path)
    const (n,p) = size(x)

    # load response vector
    y = readdlm(y_path)
    y = vec(y)  # need y to be 1D

    # load model
    B = readdlm(b_path)

    # need indices for true model
    # add 1 since Julia uses 1-indexing
    bidx = convert(Array{Int,1}, B[:,2]) + 1

    # need components of true model
    be = convert(Array{Float64,1}, B[:,3])

    # how long did it take to load files?
	file_time = toq()
    quiet || println("Files took ", file_time, " seconds to load.")

	# load parameters
	println("Allocating temporary arrays...")
	tic()

	# declare all temporary arrays
    b          = zeros(p)	# parameter vector
	res        = zeros(n)	# Y - Xbeta
	df         = zeros(p)	# X'(Y - Xbeta)
	tempn      = zeros(n)	# temporary array of length n
	tempn2     = zeros(n)	# temporary array of length n
	tempp      = zeros(p)	# temporary array of length p
	dotprods   = zeros(p)	# hold in memory the dot products for current index
	perm       = collect(1:p) # hold permutation vector of b

	# declare associative array for storing inner products
	inner = Dict{Int,DenseArray{Float64,1}}()

	# how long did the allocation take?
	alloc_time = toq()
	quiet || println("Temporary arrays allocated in ", alloc_time, " seconds.")

	# precompute sum of squares for each column of x
	quiet || println("Precomputing squared Euclidean norms of matrix columns...")
	tic()
	const nrmsq = vec(sumabs2(x,1))
	norm_time = toq()
	quiet || println("Column norms took ", norm_time, " seconds to compute")

	# declare any return values
	iter = 0

	# print time spent on parameter initialization
	parameter_time = norm_time + alloc_time
	quiet || println("Parameter initialization took ", parameter_time, " seconds.")

    # run exchange
    println("\nRunning exchange algorithm...")

	# reset timer and compute path
	tic()
	@time begin
		for i = 1:(r+extra)
			b = exchange_leastsq!(b, x, y, perm, i, inner=inner, max_iter=max_iter, quiet=quiet, nrmsq=nrmsq, res=res, df=df, tempn=tempn, tempn2=tempn2, tempp=tempp, dotprods=dotprods, window=min(i,window))
			rss = sumabs2(res)

			# print results from this step of regularization path
			quiet || println("#pred = ", i, ", rss = ", rss, ", #truepos = ", countnz(b[bidx]))
		end
	end

	# print time spent on computation of regularization path
	quiet || println("Path from 1 to ", r+extra, " took ", toq(), " seconds to compute.")

    # recover vector from output
    bk = b[bidx]

    # evaluate model
    println("\nTrue positives: ", countnz(bk), "/", length(be), ".")
    println("Distances to true model:")
    println("Chebyshev (L-Inf) = ", chebyshev(bk, be), ".")
    println("Euclidean (L-2)   = ", euclidean(bk, be), ".")
end

# TEST THE EXCHANGE ALGORITHM FOR LEAST SQUARES REGRESSION OVER PLINK BINARY DATA
# This subroutine runs the exchange algorithm on a GWAS dataset from the WTCCC.
function test_exchangeleastsq_plink(x_path::String = "/Users/kkeys/Downloads/wtccc_full/wtccc-n2k_chr1_clean.bed", xt_path::String = "/Users/kkeys/Downloads/wtccc_full/wtccc-n2k_chr1_clean_t.bed", y_path::String = "/Users/kkeys/Downloads/withnoiselevelsd0_1/Y.100.1", b_path::String = "/Users/kkeys/Downloads/withnoiselevelsd0_1/causal.100.1", r::Int = 10, tol::Float64 = 1e-4, max_iter::Int = 100, quiet::Bool = false, extra::Int = 10, window::Int = 10)

    # repeat input options
	quiet || begin
		println("This function will test the IHT algorithm.")
		println("Given options:")
		println("\tPath to X  = ", x_path)
		println("\tPath to X' = ", xt_path)
		println("\tPath to Y  = ", y_path)
		println("\tPath to B  = ", b_path)
		println("\tModel size = ", r)
		println("\ttolerance  = ", tol, "\n\tmax_iter   = ", max_iter)
	end

    # precompile @time macro for later use
    quiet || println("\nCompiling @time macro...")
    @time 1+1;

    # now load data
    quiet || println("\nLoading data...")
    tic()

    # load design matrix
	x = BEDFile(x_path, xt_path)
	const (n,p) = size(x)

    # load response vector
    y = readdlm(y_path)
	y = vec(y) # need y to be 1D
	const y = convert(SharedArray{Float64,1}, y)

	const means   = PLINK.mean(Float64, x, shared=true)
	const invstds = PLINK.invstd(x, means, shared=true)

    # load model
    B = readdlm(b_path)

    # need indices for true model
    # add 1 since Julia uses 1-indexing
	const bidx = convert(Array{Int,1}, B[:,2]) + 1

    # need components of true model
	const be = convert(Array{Float64,1}, B[:,3])

    # how long did it take to load files?
	file_time = toq()
    quiet || println("\nFiles took ", file_time, " seconds to load.")

	# discard B to recover some memory
	quiet || println("Object B is no longer needed, recovering memory...")
	tic()
	B = false
	gc()
	B_time = toq()
	quiet || println("Discarding B took ", B_time, " seconds.")

	# load parameters
	quiet || println("\nAllocating temporary arrays...")
	tic()

	# declare all temporary arrays
	b        = SharedArray(Float64, p, init = S -> S[localindexes(S)] = 0.0)		# parameter vector
	Xb       = SharedArray(Float64, n, init = S -> S[localindexes(S)] = 0.0)		# Xbeta
	res      = SharedArray(Float64, n, init = S -> S[localindexes(S)] = 0.0)		# Y - Xbeta
	df       = SharedArray(Float64, p, init = S -> S[localindexes(S)] = 0.0)		# X'(Y - Xbeta)
	tempn    = SharedArray(Float64, n, init = S -> S[localindexes(S)] = 0.0)		# temporary array of length n
	tempn2   = SharedArray(Float64, n, init = S -> S[localindexes(S)] = 0.0)		# temporary array of length n
	tempp    = SharedArray(Float64, p, init = S -> S[localindexes(S)] = 0.0)		# temporary array of length p
	dotprods = SharedArray(Float64, p, init = S -> S[localindexes(S)] = 0.0)		# hold in memory the dot products for current index
	indices  = falses(p)

	# permutation vector of b has a more complicated initialization
	perm = SharedArray(Int, p, init = S -> S[localindexes(S)] = localindexes(S))

	# declare associative array for storing inner products
	inner = Dict{Int,SharedArray{Float64,1}}()

	# how long did the allocation take?
	alloc_time = toq()
	quiet || println("Temporary arrays allocated in ", alloc_time, " seconds.")

	# precompute sum of squares for each column of x
	quiet || println("\nPrecomputing squared Euclidean norms of matrix columns...")
	tic()
	const nrmsq = sumsq(Float64, x, shared=true, means=means, invstds=invstds)
	norm_time = toq()
	quiet || println("Column norms took ", norm_time, " seconds to compute")

	# print time spent on parameter initialization
	parameter_time = norm_time + alloc_time + file_time + B_time
	quiet || println("Parameter initialization took ", parameter_time, " seconds.")

    # run exchange
    quiet || println("\nRunning exchange algorithm...")

	# reset timer and compute path
	tic()
	@time begin
		for i = 2:(r+extra)
			exchange_leastsq!(b, x, y, perm, i, p=p, n=n, inner=inner, max_iter=max_iter, quiet=quiet, nrmsq=nrmsq, res=res, df=df, tempn=tempn, tempn2=tempn2, tempp=tempp, dotprods=dotprods, window=min(i,window), Xb=Xb, means=means, invstds=invstds, indices=indices)
			rss = sumabs2(res)

			# print results from this step of regularization path
			quiet || println("#pred = ", i, ", rss = ", rss, ", #truepos = ", countnz(b[bidx]))
		end
	end

	# print time spent on computation of regularization path
	quiet || println("Path from 1 to ", r+extra, " took ", toq(), " seconds to compute.")

    # recover vector from output
    bk = b[bidx]

    # evaluate model
	quiet || begin
		println("\nTrue positives: ", countnz(bk), "/", length(be), ".")
		println("Distances to true model:")
		println("Chebyshev (L-Inf) = ", chebyshev(bk, be), ".")
		println("Euclidean (L-2)   = ", euclidean(bk, be), ".")
	end

	return bk, bidx
end

test_exchangeleastsq_plink(r::Int, extra::Int, window::Int) = test_exchangeleastsq_plink("/Users/kkeys/Downloads/wtccc_full/wtccc-n2k_chr1_clean.bed", "/Users/kkeys/Downloads/wtccc_full/wtccc-n2k_chr1_clean_t.bed", "/Users/kkeys/Downloads/withnoiselevelsd0_1/Y.100.1", "/Users/kkeys/Downloads/withnoiselevelsd0_1/causal.100.1", r, 1e-4, 100, false, extra, window)

# CROSSVALIDATION ROUTINE FOR EXCHANGE ALGORITHM
function test_cv_exlstsq(kend,numfolds; tol::Float64 = 1e-4, max_iter::Int = 1000, quiet::Bool=true, compute_model::Bool = false)

	# notify of number of folds
	quiet || println("Testing ", numfolds, "-fold crossvalidation from k = 1 to k = ", kend, "...")

	# start timing
	tic()

	# file paths for data, response, model
	x_path = "/Users/kkeys/Downloads/wtccc-n2000p32307.txt"
	y_path = "/Users/kkeys/Downloads/withnoiselevelsd0_1/Y.100.1"
	b_path = "/Users/kkeys/Downloads/withnoiselevelsd0_1/causal.100.1"

	# load design matrix
	x = readdlm(x_path)
	const (n,p) = size(x)
	const x = convert(SharedArray, x)

	# load response vector
	y = readdlm(y_path)
	y = vec(y)	# need y to be 1D
	const y = convert(SharedArray,y)

	# load model
	B = readdlm(b_path)

	# need indices for true model
	# add 1 since Julia uses 1-indexing
	const bidx = convert(Array{Int,1}, B[:,2]) + 1

	# need components of true model
	const be = convert(Array{Float64,1}, B[:,3])

	parameter_time = toq()

	# inform console that the data are loaded
	quiet || println("Data loaded in ", parameter_time, " seconds.")

	# obtain a partition of the sample set
	myfolds = cv_get_folds(y,numfolds)

	# inform console that the folds are computed
	quiet || println("Folds computed.")

	# run cv_iht to crossvalidate dataset
	# be careful! mses will be a tuple if compute_model=true
	tic()
	if compute_model
		mses, b = cv_exlstsq(x,y,kend,numfolds, p=p, tol=tol, max_iter=max_iter, quiet=quiet, folds=myfolds, compute_model=compute_model)
	else
		mses = cv_exlstsq(x,y,kend,numfolds, p=p, tol=tol, max_iter=max_iter, quiet=quiet, folds=myfolds, compute_model=compute_model)
	end

	cv_time = toq()
	quiet || println("CV computed in ", cv_time, " seconds.")

	compute_model && return mses, b
	return mses
end

function compare_exlstsq(x::Matrix, kend, numfolds; tol::Float64 = 1e-4, max_iter::Int = 1000, quiet::Bool=true, compute_model::Bool = false, y_path::String = "/Users/kkeys/Desktop/Thesis/Exchange/Y_1_same_1.0_2.0.txt", b_path::String = "/Users/kkeys/Desktop/Thesis/Exchange/b_1_same_1.0_2.0.txt", bidx_path::String = "/Users/kkeys/Desktop/Thesis/Exchange/bidx_1_same_1.0_2.0.txt")

	# notify of number of folds
	quiet || println("Testing ", numfolds, "-fold crossvalidation from k = 1 to k = ", kend, "...")

	# start timing
	tic()

	# dimensions of design matrix
	const n,p = size(x)
	x = convert(SharedArray{Float64,2}, x)

	# load response vector
	y = readdlm(y_path)
	y = vec(y)	# need y to be 1D
#	y = convert(Array{Float64,1}, y)
	y = convert(SharedArray{Float64,1}, y)

	# load model
	be = readdlm(b_path)
	be = vec(be)
	const be = convert(Array{Float64,1}, be)

	# need indices for true model
	bidx = readdlm(bidx_path)
	bidx = vec(bidx)
	const bidx = convert(Array{Int,1}, bidx)

	parameter_time = toq()

	# inform console that the data are loaded
	quiet || println("Data loaded in ", parameter_time, " seconds.")

	# start with exchange algorithm
	quiet || println("Running crossvalidation with exchange algorithm...")

	# obtain a partition of the sample set
	myfolds = cv_get_folds(y,numfolds)

	# inform console that the folds are computed
	quiet || println("Folds computed.")

	# run cv_iht to crossvalidate dataset
	# be careful! mses will be a tuple if compute_model=true
	tic()
#	mses = cv_exlstsq(x,y,kend,numfolds, p=p, tol=tol, max_iter=max_iter, quiet=quiet, folds=myfolds, compute_model=compute_model)
	if compute_model
		mses, b, ex_model = cv_exlstsq(x,y,kend,numfolds, p=p, tol=tol, max_iter=max_iter, quiet=quiet, folds=myfolds, compute_model=compute_model)
	else
		mses = cv_exlstsq(x,y,kend,numfolds, p=p, tol=tol, max_iter=max_iter, quiet=quiet, folds=myfolds, compute_model=compute_model)
	end

	cv_time = toq()
	quiet || println("CV computed in ", cv_time, " seconds.")

	# now compute LASSO model
	quiet || println("\nComputing LASSO model...")

	# translate folds

	# run glmnet
	tic()
	cv = glmnetcv(sdata(x), sdata(y), dfmax=kend, nlambda=kend, nfolds=5, parallel=true, folds=my_folds, tol=tol)

	# extract best beta from lasso
	lasso_btemp = cv.path.betas[:,indmin(cv.meanloss)]

	# inform console that lasso computation is complete
	lasso_time = toq()
	quiet || println("LASSO model crossvalidated in ", lasso_time, " seconds.")

	if compute_model
#		lasso_bidx = lasso_btemp .!= 0.0
		lasso_bidx = find( function f(x) x .!= 0.0; end, lasso_btemp)
		x_lasso = x[:,lasso_bidx]
		lasso_b = (x_lasso'*x_lasso) \ (x_lasso'*y)
	end

	# return values based on compute_model
	compute_model && return mses, cv.meanloss, b, lasso_b, be, ex_model, lasso_bidx, bidx
	return mses, cv.meanloss
end

function compare_exlstsq(x::Matrix, y::Vector, be::Vector, bidx::Vector, kend::Int, numfolds::Int; tol::Float64 = 1e-4, max_iter::Int = 10000, quiet::Bool=true, compute_model::Bool = false, n::Int = length(y), p::Int = length(be), window::Int = 20, nrmsq::DenseArray{Float64,1} = vec(sumsq(x,1)))

	isequal((n,p),size(x)) || throw(DimensionMismatch("(n,p) = ($n,$p) != ($(size(x)))"))
	0 <= window <= p       || throw(error("For current X, must have window in (0,$p)"))
	0 <= kend <= p         || throw(error("Path legnth must be positive and cannot exceed the number of predictors"))

	# notify of number of folds
	quiet || println("Testing ", numfolds, "-fold crossvalidation from k = 1 to k = ", kend, "...")

	# start with exchange algorithm
	quiet || println("Running crossvalidation with exchange algorithm...")

	# obtain a partition of the sample set
	myfolds = cv_get_folds(y,numfolds)

	# inform console that the folds are computed
	quiet || println("Folds computed.")

	# run cv_iht to crossvalidate dataset
	tic()
	if compute_model
		mses, b, ex_model = cv_exlstsq(x,y,kend,numfolds, p=p, tol=tol, max_iter=max_iter, quiet=quiet, folds=myfolds, compute_model=compute_model, window=window)
	else
		mses = cv_exlstsq(x,y,kend,numfolds, p=p, tol=tol, max_iter=max_iter, quiet=quiet, folds=myfolds, compute_model=compute_model, window=window)
	end
	cv_time = toq()
	quiet || println("CV computed in ", cv_time, " seconds.")

	# now compute LASSO model
	quiet || println("\nComputing LASSO model...")

	# run glmnet
	tic()
#	cv = glmnetcv(x,y, dfmax=kend, nlambda=kend, nfolds=numfolds, parallel=true, folds=my_folds, tol=tol)

	# extract best beta from lasso
#	lasso_btemp = cv.path.betas[:,indmin(cv.meanloss)]
	lasso_btemp = zeros(p)
	lasso_btemp[1] = 1.0

	# inform console that lasso computation is complete
	lasso_time = toq()
	quiet || println("LASSO model crossvalidated in ", lasso_time, " seconds.")

	if compute_model
#		lasso_bidx = lasso_btemp .!= 0.0
		lasso_bidx = find( function f(x) x .!= 0.0; end, lasso_btemp)
		x_lasso = x[:,lasso_bidx]
		lasso_b = (x_lasso'*x_lasso) \ (x_lasso'*y)
	end

	# return values based on compute_model
	compute_model && return mses, zeros(kend), b, lasso_b, be, ex_model, lasso_bidx, bidx, lasso_time, cv_time
	return mses, cv.meanloss, cv_time, lasso_time
end