# ******************************************************************************
# PACKAGES AND INCLUDES
# ******************************************************************************
using Plots
using ImplicitEquations

include("project2.jl")
include("helpers.jl")
include("simple.jl")

# ******************************************************************************
# PATH STORING FUNCTIONS
# ******************************************************************************
function path_hooke_jeeves(f, g, c, x, num_evals, ρ1, ρ2, γ, δ)
    x1_arr = Float64[]; x2_arr = Float64[]; y_arr = Float64[]; max_c_arr = Float64[];
    push!(x1_arr, x[1]); push!(x2_arr, x[2]); push!(y_arr, f(x))
    evals = c(x); max_c = maximum([evals[1], evals[2], 0]); push!(max_c_arr, max_c)
    α, y, n, ϵ = 1, f(x) + penalty(c, x, ρ1, ρ2, δ), length(x), 1e-9
    while α > ϵ && count(f, g, c) < (num_evals-8)
        evals = c(x); max_c = maximum([evals[1], evals[2], 0]); push!(max_c_arr, max_c)
        x, y, α = hooke_jeeves(x -> f(x) + penalty(c, x, ρ1, ρ2, δ), x, y, α, n, 0.5)
        push!(x1_arr, x[1]);push!(x2_arr, x[2]); push!(y_arr, f(x))
        ρ1 *= γ; ρ2 *= γ
    end
    return x, x1_arr, x2_arr, y_arr, max_c_arr
end

function path_mesh_adaptive(f, g, c, x, num_evals, ρ1, ρ2, γ, δ)
    x1_arr = Float64[]; x2_arr = Float64[]; y_arr = Float64[]; max_c_arr = Float64[];
    push!(x1_arr, x[1]); push!(x2_arr, x[2]); push!(y_arr, f(x))
    evals = c(x); max_c = maximum([evals[1], evals[2], 0]); push!(max_c_arr, max_c)
    α, y, n, ϵ = 1, f(x) + penalty(c, x, ρ1, ρ2, δ), length(x), 1e-9
    while α > ϵ && count(f, g, c) < (num_evals-8)
        evals = c(x); max_c = maximum([evals[1], evals[2], 0]); push!(max_c_arr, max_c)
        x, y, α = mesh_adaptive_direct_search(x -> f(x) + penalty(c, x, ρ1, ρ2, δ), x, y, α, n)
        push!(x1_arr, x[1]); push!(x2_arr, x[2]); push!(y_arr, f(x))
        ρ1 *= γ; ρ2 *= γ
    end
    return x, x1_arr, x2_arr, y_arr, max_c_arr
end

# ******************************************************************************
# PLOTTING HELPER FUNCTIONS
# ******************************************************************************
function scatter_helper(x1_arr, x2_arr, x0, x_best, color)
    scatter!([x1_arr], [x2_arr], color=color, label = "", markersize = 3, outline = color)
    scatter!([x0[1]], [x0[2]], color=color,  marker = :rect,label = "", markersize = 5)
    scatter!([x_best[1]], [x_best[2]], color=color, marker = :star4, label = "", markersize = 7)
end

function plot_simple1_contour()
    x1 = -3.0:0.01:3.0
    x2 = -3.0:0.01:3.0
    plot_f(x1, x2) = begin
        -x1 * x2 + 2.0 / (3.0 * sqrt(3.0))
    end
    X = repeat(reshape(x1, 1, :), length(x2), 1)
    Y = repeat(x2, 1, length(x1))
    Z = map(plot_f, X, Y)
    contour(x1, x2, Z, levels=31, fill=false, c=cgrad(:viridis, rev = true), linewidth=2.0, colorbar=true)

    c1(x1,x2) = -x1 - x2
    c2(x1,x2) = x1 + x2^2 - 1

    plot!((Lt(c1,0)) & (Lt(c2,0)), fill=(0.65,"#495C6F"), lims=[-3,3],
        label="Feasible Region")
end

function plot_simple2_contour()
    x1 = -3.0:0.01:3.0
    x2 = -3.0:0.01:3.0
    plot_f(x1, x2) = begin
        (1.0 - x1)^2 + 100.0 * (x2 - x1^2)^2
    end
    X = repeat(reshape(x1, 1, :), length(x2), 1)
    Y = repeat(x2, 1, length(x1))
    Z = map(plot_f, X, Y)
    h = [minimum(Z), 2.5, 10, 25, 100, 250, 500, 750,1000, 1500, 2500, 3500, 5000, 6500, 8000, 10000,maximum(Z)]
    contour(x1, x2, Z, levels=h, fill=false, c=cgrad(:viridis, rev = true), linewidth=2.0, colorbar=true)

    c1(x1,x2) = (x1-1)^3 - x2 + 1
    c2(x1,x2) = x1 + x2 - 2

    plot!((Lt(c1,0)) & (Lt(c2,0)), fill=(0.65,"#495C6F"), lims=[-3,3],
        label="Feasible Region")
end

# ******************************************************************************
# PATH PLOTTING FUNCTIONS
# ******************************************************************************
function plot_simple1_hooke_jeeves(f, g, c, n)

    plot_simple1_contour()

    x01 = [2.5, -2.5]; x02 = [2.5, 2.5]; x03 = [-2.0, 2.0]
    x1_best, x11_arr, x21_arr, _, _ =
        path_hooke_jeeves(f, g, c, x01, n, 1.01, 0, 1.001, 0.5);
    x2_best, x12_arr, x22_arr, _, _ =
        path_hooke_jeeves(f, g, c, x02, n, 1.01, 0, 1.001, 0.5);
    x3_best, x13_arr, x23_arr, _, _ =
        path_hooke_jeeves(f, g, c, x03, n, 1.01, 0, 1.001, 0.5);

    # Plot first path
    plot!(x11_arr,x21_arr,color=:magenta, label = "x0: (2.5, -2.5)",legend = :bottomleft,
        title = "Hooke-Jeeves Method to Minimize Simple1", titlefontsize = 12,
        xlabel = "x1", ylabel = "x2", dpi=300, colorbar = true)
    scatter_helper(x11_arr, x21_arr, x01, x1_best, :magenta)
    # Plot second path
    plot!(x12_arr,x22_arr,color=:red, label = "x0: (2.5, 2.5)")
    scatter_helper(x12_arr, x22_arr, x02, x2_best, :red)
    # Plot third path
    plot!(x13_arr,x23_arr,color=:darkorange, label = "x0: (-2.0, 2.0)")
    scatter_helper(x13_arr, x23_arr, x03, x3_best, :darkorange)

    png("simple1_hooke_jeeves.png")
end

function plot_simple1_mesh_adaptive(f, g, c, n)
    # Plot contour
    plot_simple1_contour()

    # Store paths
    x01 = [2.5, -2.5]; x02 = [2.5, 2.5]; x03 = [-2.0, 2.0]
    x1_best, x11_arr, x21_arr, _, _ =
        path_mesh_adaptive(f, g, c, x01, n, 1.01, 0, 1.001, 0.5);
    x2_best, x12_arr, x22_arr, _, _ =
        path_mesh_adaptive(f, g, c, x02, n, 1.01, 0, 1.001, 0.5);
    x3_best, x13_arr, x23_arr, _, _ =
        path_mesh_adaptive(f, g, c, x03, n, 1.01, 0, 1.001, 0.5);

    # Plot first path
    plot!(x11_arr,x21_arr,color=:magenta, label = "x0: (2.5, -2.5)",legend = :bottomleft,
        title = "Mesh Adaptive Direct Search Method to Minimize Simple1", titlefontsize = 12,
        xlabel = "x1", ylabel = "x2", dpi=300, colorbar = true)
    scatter_helper(x11_arr, x21_arr, x01, x1_best, :magenta)
    # Plot second path
    plot!(x12_arr,x22_arr,color=:red, label = "x0: (2.5, 2.5)")
    scatter_helper(x12_arr, x22_arr, x02, x2_best, :red)
    # Plot third path
    plot!(x13_arr,x23_arr,color=:darkorange, label = "x0: (-2.0, 2.0)")
    scatter_helper(x13_arr, x23_arr, x03, x3_best, :darkorange)

    png("simple1_mesh_adaptive.png")
end

function plot_simple2_hooke_jeeves(f, g, c, n)
    # Plot contour
    plot_simple2_contour()

    # Store paths
    x01 = [2.5, -2.5]; x02 = [2.5, 2.5]; x03 = [-2.0, 2.0]
    x1_best, x11_arr, x21_arr, _, _ =
        path_hooke_jeeves(f, g, c, x01, n, 100, 1, 1.1, 0.75);
    x2_best, x12_arr, x22_arr, _, _ =
        path_hooke_jeeves(f, g, c, x02, n, 100, 1, 1.1, 0.75);
    x3_best, x13_arr, x23_arr, _, _ =
        path_hooke_jeeves(f, g, c, x03, n, 100, 1, 1.1, 0.75);

    # Plot first path
    plot!(x11_arr,x21_arr,color=:magenta, label = "x0: (2.5, -2.5)",legend = :bottomleft,
        title = "Hooke-Jeeves Method to Minimize Simple2", titlefontsize = 12,
        xlabel = "x1", ylabel = "x2", dpi=300, colorbar = true)
    scatter_helper(x11_arr, x21_arr, x01, x1_best, :magenta)
    # Plot second path
    plot!(x12_arr,x22_arr,color=:red, label = "x0: (2.5, 2.5)")
    scatter_helper(x12_arr, x22_arr, x02, x2_best, :red)
    # Plot third path
    plot!(x13_arr,x23_arr,color=:darkorange, label = "x0: (-2.0, 2.0)")
    scatter_helper(x13_arr, x23_arr, x03, x3_best, :darkorange)

    png("simple2_hooke_jeeves.png")
end

function plot_simple2_mesh_adaptive(f, g, c, n)
    # Plot contour
    plot_simple2_contour()

    # Store paths
    x01 = [0.0, -2.0]; x02 = [1.0, -1.0]; x03 = [2.0, 0.0]
    x1_best, x11_arr, x21_arr, _, _ =
        path_mesh_adaptive(f, g, c, x01, n, 100, 1, 1.01, 1.0);
    x2_best, x12_arr, x22_arr, _, _ =
        path_mesh_adaptive(f, g, c, x02, n, 100, 1, 1.01, 1.0);
    x3_best, x13_arr, x23_arr, _, _ =
        path_mesh_adaptive(f, g, c, x03, n, 100, 1, 1.01, 1.0);

    # Plot first path
    plot!(x11_arr,x21_arr,color=:magenta, label = "x0: (0.0, -2.0)",legend = :bottomleft,
        title = "Mesh Adaptive Direct Search Method to Minimize Simple2", titlefontsize = 12,
        xlabel = "x1", ylabel = "x2", dpi=300, colorbar = true)
    scatter_helper(x11_arr, x21_arr, x01, x1_best, :magenta)
    # Plot second path
    plot!(x12_arr,x22_arr,color=:red, label = "x0: (1.0, -1.0)")
    scatter_helper(x12_arr, x22_arr, x02, x2_best, :red)
    # Plot third path
    plot!(x13_arr,x23_arr,color=:darkorange, label = "x0: (2.0, 0.0)")
    scatter_helper(x13_arr, x23_arr, x03, x3_best, :darkorange)

    png("simple2_mesh_adaptive.png")
end

function plot_obj_fun_max_con_hooke_jeeves(f, g, c, n)

    x01 = [2.5, -2.5]; x02 = [2.5, 2.5]; x03 = [-2.0, 2.0]
    _, _, _, y1_arr, max_c1_arr =
        path_hooke_jeeves(f, g, c, x01, n, 100, 1, 1.1, 0.75);
    _, _, _, y2_arr, max_c2_arr =
        path_hooke_jeeves(f, g, c, x02, n, 100, 1, 1.1, 0.75);
    _, _, _, y3_arr, max_c3_arr =
        path_hooke_jeeves(f, g, c, x03, n, 100, 1, 1.1, 0.75);

    x_lim = minimum([length(y1_arr), length(y2_arr), length(y3_arr), 20])
    x_vals = 1:x_lim

    plot(x_vals, y1_arr[1:x_lim], lw=3, color = "#B2BfCC", label="x0: (2.5, -2.5)",legend = :topright,
        title="Objective Function vs Iteration for Hooke-Jeeves",
        xlabel="Iteration",ylabel="Maximum Constraint Violation",dpi=300)
    plot!(x_vals, y2_arr[1:x_lim], lw = 3, color = "#495C6F", label = "x0: (2.5, 2.5)")
    plot!(x_vals, y3_arr[1:x_lim], lw = 3, color = "#1E262E", label = "x0: (-2.0, 2.0)")

    png("objective_function_hooke_jeeves.png")

    plot(x_vals, max_c1_arr[1:x_lim], lw=3, color = "#B2BfCC", label="x0: (2.5, -2.5)",legend = :topright,
        title="Maximum Constraint Violation for Hooke-Jeeves",
        xlabel="Iteration",ylabel="Maximum Constraint Violation",dpi=300)
    plot!(x_vals, max_c2_arr[1:x_lim], lw = 3, color = "#495C6F", label = "x0: (2.5, 2.5)")
    plot!(x_vals, max_c3_arr[1:x_lim], lw = 3, color = "#1E262E", label = "x0: (-2.0, 2.0)")

    png("constraint_violation_hooke_jeeves.png")
end

function plot_obj_fun_max_con_mesh_adaptive(f, g, c, n)

    x01 = [0.0, -2.0]; x02 = [1.0, -1.0]; x03 = [2.0, 0.0]
    _, _, _, y1_arr, max_c1_arr =
        path_mesh_adaptive(f, g, c, x01, n, 100, 1, 1.01, 1.0);
    _, _, _, y2_arr, max_c2_arr =
        path_mesh_adaptive(f, g, c, x02, n, 100, 1, 1.01, 1.0);
    _, _, _, y3_arr, max_c3_arr =
        path_mesh_adaptive(f, g, c, x03, n, 100, 1, 1.01, 1.0);

    x_lim = minimum([length(y1_arr), length(y2_arr), length(y3_arr), 20])
    x_vals = 1:x_lim

    plot(x_vals, y1_arr[1:x_lim], lw=3, color = "#B2BfCC", label="x0: (0.0, -2.0)",legend = :topright,
        title="Objective Function for Mesh Adaptive",
        xlabel="Iteration",ylabel="Maximum Constraint Violation",dpi=300)
    plot!(x_vals, y2_arr[1:x_lim], lw = 3, color = "#495C6F", label = "x0: (1.0, -1.0)")
    plot!(x_vals, y3_arr[1:x_lim], lw = 3, color = "#1E262E", label = "x0: (2.0, 0.0)")

    png("objective_function_mesh_adaptive.png")

    plot(x_vals, max_c1_arr[1:x_lim], lw=3, color = "#B2BfCC", label="x0: (0.0, -2.0)",legend = :topright,
        title="Maximum Constraint Violation for Mesh Adaptive",
        xlabel="Iteration", ylabel="Maximum Constraint Violation",dpi=300)
    plot!(x_vals, max_c2_arr[1:x_lim], lw = 3, color = "#495C6F", label = "x0: (1.0, -1.0)")
    plot!(x_vals, max_c3_arr[1:x_lim], lw = 3, color = "#1E262E", label = "x0: (2.0, 0.0)")

    png("constraint_violation_mesh_adaptive.png")
end

n = 100000
# Path plots
plot_simple1_hooke_jeeves(simple1, simple1_gradient, simple1_constraints, n)
plot_simple1_mesh_adaptive(simple1, simple1_gradient, simple1_constraints, n)
plot_simple2_hooke_jeeves(simple2, simple2_gradient, simple2_constraints, n)
plot_simple2_mesh_adaptive(simple2, simple2_gradient, simple2_constraints, n)

# Objective function and maximum constraint plots
plot_obj_fun_max_con_hooke_jeeves(simple2, simple2_gradient, simple2_constraints, n)
plot_obj_fun_max_con_mesh_adaptive(simple2, simple2_gradient, simple2_constraints, n)


"""
    inhomogeneous_poisson_process(lam, dt; kwargs...)

Returns a sequence of spike times produced by the inhomogeneous poisson process
with rate `lambda`. `dt` is 1/sampling rate.

To produce a more reliable spike train, increase `sampling_rate_factor`. If
`sampling_rate_factor = k` then the function first computes `kron(lam, ones(k))`
and uses sampling rate `dt/k` (For instance, with the CRCNS data, `k` between 10
and 20 seems to work).

Will not produce spikes with repeated spike times (this could come up if, e.g.
`lam[t]` is very large)

Options:
 * `sampling_rate_factor::Int = 1` Increases sampling rate

Stopping criteria:
 * `Tmax = length(lam) * dt` Will produces spikes until first spike later than this time
 * `max_real_time` Will stop this process early if it's taking forever (i.e. if `lam` is very low). Default value `0` ignores this option.
 * `max_loops` Same as above, counting number of loops instead of amount of real time. Default `0` to ignore.
 * `max_spikes` Same as above, counting number of spikes produced (e.g. if you want the first-so-many spieks produced by a process.)

These additional kwargs can be passed arrays. The first entry will be modified
in place. This allows additional information to be collected about the process
without requiring the function to return multiple values.
 * `total_time` The total run time of the function
 * `total_loops` The total number of times the process looped
 * `total_spikes` The same as `length(times)` where `times` is the output of this function.

This implementation is based on the first algorithm outlined at
http://freakonometrics.hypotheses.org/724

"""
function inhomogeneous_poisson_process(lam::Vector{Float64}, dt::Float64;
    Tmax::Float64=length(lam) * dt, sampling_rate_factor::Int=1,
    max_real_time::Float64=0.0, max_loops::Int=0, max_spikes::Int=0,
    total_time=[], total_loops=[], total_spikes=[],
    exit_status=Symbol[])
    # based on the first algorithm outlined at http://freakonometrics.hypotheses.org/724
    dt = dt/sampling_rate_factor
    # Lam = cumsum_kbn(kron(lam, ones(sampling_rate_factor)) * dt)
    Lam = cumsum_kbn(kron(lam, fill(dt, sampling_rate_factor)))
    times = Vector{Float64}()
    t_last = 0.0
    s = 0.0; t = 0.0
    real_time = 0.0
    n_loops = 0
    n_spikes = 0
    while t < Tmax
        tic()
        u = 1 - rand()
        s = s - log(u)
        idx = searchsortedlast(Lam, s)
        t = idx * dt
        if Tmax > t > t_last
            # Avoid repeated spikes by comparing to t_last
            push!(times, t)
            t_last = t
            n_spikes += 1
        end
        real_time += toq()
        n_loops += 1
        if idx == length(Lam)
            # if we've hit the end of Lam, then we should quit while we're ahead
            push!(exit_status,:end_of_lambda_reached)
            break
        end
        if 0.0 < max_real_time < real_time
            push!(exit_status,:max_real_time_reached)
            break
        end
        if 0 < max_loops < n_loops
            push!(exit_status,:max_loops_reached)
            break
        end
        if 0 < max_spikes < n_spikes
            push!(exit_status,:max_number_of_spikes)
            break
        end
    end
    if t >= Tmax
        push!(exit_status,:max_spike_time_reached)
    end
    if length(total_time) > 0
        total_time[1] = real_time
    end
    if length(total_loops) > 0
        total_loops[1] = n_loops
    end
    if length(total_spikes) > 0
        total_spikes[1] = length(times)
    end
    return times
end

# MAYBEDO maybe implement a version of this function that returns a proper
# SpikeTrains object. Probably not necessary, but keeping this flag here just in
# case.


@show parse(Int, "123459")
@show parse(Int, "0xabcf123")
@show parse(Int, "0o7651")
@show parse(Int, "0b101011001")


"""
    axes3D(n=100)

Draw labelled 3D axes at `(0, 0, 0)` with length `n`.
"""
function axes3D(n=100)
    projection = CURRENTPROJECTION[1]
    @layer begin
        fontsize(16)
        setline(2)
        xaxis1 = project(Point3D(0.1,  0.1,  0.1))
        xaxis2 = project(Point3D(n,    0.1,  0.1))
        sethue("red")
        if (xaxis1 != nothing) && (xaxis2 != nothing) && !isapprox(xaxis1, xaxis2)
            arrow(xaxis1, xaxis2)
            label("X", :N, xaxis2)
        end
        yaxis1 = project(Point3D(0.1,  0.1,  0.1))
        yaxis2 = project(Point3D(0.1,  n,  0.1))
        sethue("green")

        if (yaxis1 != nothing) &&  (yaxis2 != nothing) && !isapprox(yaxis1, yaxis2)
            arrow(yaxis1, yaxis2)
            label("Y", :N, yaxis2)
        end
        zaxis1 = project(Point3D(0.1,  0.1,  0.1))
        zaxis2 = project(Point3D(0.1,  0.1,  n))
        sethue("blue")
        if (zaxis1 != nothing) &&  (zaxis2 != nothing) && !isapprox(zaxis1, zaxis2)
            arrow(zaxis1, zaxis2)
            label("Z", :N, zaxis2)
        end
    end
end

"""
    drawcube(n=10, action=:stroke)

Draw a cube. `drawcube(1)` draws a wireframe unit cube.
"""
function drawcube(n=10, action=:stroke)
    @layer begin
        setline(2)
        fontsize(10)
        setlinecap("butt")
        setlinejoin("bevel")
        p1 = project(Point3D(0,   0,  0))
        p2 = project(Point3D(n,  0,  0))
        p3 = project(Point3D(n, n,  0))
        p4 = project(Point3D(0,  n,  0))
        p5 = project(Point3D(0,   0,  n))
        p6 = project(Point3D(n,  0,  n))
        p7 = project(Point3D(n,  n,  n))
        p8 = project(Point3D(0,  n,  n))

        if all(i -> i != nothing, [p1, p2, p3, p4]) # , p5, p6, p7, p8])
            sethue("red")
            prettypoly([p1, p2, p3, p4], action, close=true, () -> begin
                circle(O, n/100)
            end)
        end

        if all(i -> i != nothing, [p1, p2, p6, p5])
            sethue("purple")
            prettypoly([p1, p2, p6, p5], action, close=true, () -> begin
                circle(O, n/100)
            end)
        end

        if all(i -> i != nothing, [p3, p7, p8, p4])
            sethue("magenta")
            prettypoly([p3, p7, p8, p4], action, close=true, () -> begin
                circle(O, n/100)
            end)
        end

        if all(i -> i != nothing, [p5, p6, p7, p8])
            sethue("green")
            prettypoly([p5, p6, p7, p8], action, close=true, () -> begin
                circle(O, n/100)
            end)
        end
    end
end

"""
    carpet(n; kind=:circular)

Draw a circular carpet centered at the origin, using current Luxor parameters.

If `kind` is not `:circular`, the carpet will be a square.

Points that can't be rendered are not included in the final shape.
"""
function carpet(n=100; kind=:circular)
    projection = CURRENTPROJECTION[1]
    @layer begin
        if kind != :circular
            p1 = project(Point3D(-n/2, -n/2,  0))
            p2 = project(Point3D(n/2,  -n/2,  0))
            p3 = project(Point3D(n/2, n/2,  0))
            p4 = project(Point3D(-n/2,  n/2,  0))
            if all(i -> i != nothing, [p1, p2, p3,  p4])
                poly([p1, p2, p3, p4], :fill, close=true)
            end
        else
            shape3D = [Point3D(n * cos(theta), n * sin(theta), 0) for theta in 0:0.1:2pi]
            shape2D = Point[]
            for i in shape3D
                pt = project(i)
                if pt != nothing
                    push!(shape2D, pt)
                end
            end
            poly(shape2D, :fill, close=true)
        end
    end
end

"""
    crossproduct3D(A::Point3D, B::Point3D)

Find one of these.
"""
function crossproduct3D(A::Point3D, B::Point3D)
    x = (A.y * B.z) - (A.z * B.y)
    y = (A.z * B.x) - (A.x * B.z)
    z = (A.x * B.y) - (A.y * B.x)
    return Point3D(x, y, z)
end

"""
    pointsperpendicular(p1::Point3D, p2::Point3D, radius, angles = [0, π])

Find points perpendicular to a line joining p1 and p2.
Points are `radius` units away from the line.
"""
function pointsperpendicular(p1::Point3D, p2::Point3D, radius, angles = [0, π])
    V3 = p2 - p1 # unit vector
    V3 /= distance(Point3D(0.0, 0.0, 0.0), V3)
    e = Point3D(0.0, 0.0, 1.0)
    V1 = crossproduct3D(e, V3) # unit vector perpendicular to both e and V3
    V1 /= distance(Point3D(0.0, 0.0, 0.0), V1)
    V2 = crossproduct3D(V3, V1) # third unit vector perpendicular to both V3 and V1
    # find points P on the circle θ radians relative to vector V1 (x-axis)
    return [p1 + radius * (cos(θ) * V1 + sin(θ) * V2) for θ in angles]
end


module RationalFunctionFit
using ..Polynomials
import ..Polynomials: RationalFunction, indeterminate, constructorof
import ..Polynomials: evalpoly
using LinearAlgebra

"""
    fit(::Type{RationalFunction}, xs::AbstractVector{S}, ys::AbstractVector{T}, m, n; var=:x)

Fit a rational function of the form `pq = (a₀ + a₁x¹ + … + aₘxᵐ) / (1 + b₁x¹ + … + bₙxⁿ)` to the data `(x,y)`.

!!! Note:

    This uses a simple implementation of the Gauss-Newton method
    to solve the non-linear least squares problem:
    `minᵦ Σ(yᵢ - pq(xᵢ,β)²`, where `β=(a₀,a₁,…,aₘ,b₁,…,bₙ)`.

    A more rapidly convergent method is used in the `LsqFit.jl`
    package, and if performance is important, re-expressing the
    problem for use with that package is suggested.

    Further, if an accurate rational function fit of adaptive degrees
    is of interest, the `BaryRational.jl` package provides an
    implementation of the AAA algorithm ("which offers speed,
    flexibility, and robustness we have not seen in other algorithms"
    [Nakatsukasa, Sète,
    Trefethen](https://arxiv.org/pdf/1612.00337.pdf)) and one using
    Floater-Hormann weights [Floater,
    Hormann](https://doi.org/10.1007/s00211-007-0093-y) ("that have no
    real poles and arbitrarily high approximation orders on any real
    interval, regardless of the distribution of the points")

    The [RationalApproximations](https://github.com/billmclean/RationalApproximations) package also has implementations of the AAA algorithm.

    A python libary, [polyrat](https://github.com/jeffrey-hokanson/polyrat), has implementations of other algorithms.

## Example
```
julia> x = variable(Polynomial{Float64})
Polynomial(1.0*x)

julia> pq = (1+x)//(1-x)
(1.0 + 1.0*x) // (1.0 - 1.0*x)

julia> xs = 2.0:.1:3;

julia> ys = pq.(xs);

julia> v = fit(RationalFunction, xs, ys, 2, 2)
(1.0 + 1.0*x - 6.82121e-13*x^2) // (1.0 - 1.0*x + 2.84217e-13*x^2)

julia> maximum(abs, v(x)-pq(x) for x ∈ 2.1:0.1:3.0)
1.06314956838105e-12

julia> using BaryRational

julia> u = aaa(xs,ys)
(::BaryRational.AAAapprox{Vector{Float64}}) (generic function with 1 method)

julia> maximum(abs, u(x)-pq(x) for x ∈ 2.1:0.1:3.0)
4.440892098500626e-16

julia> u(variable(pq)) # to see which polynomial is used
(2.68328 + 0.447214*x - 1.78885*x^2 + 0.447214*x^3) // (2.68328 - 4.91935*x + 2.68328*x^2 - 0.447214*x^3)
```

"""
function Polynomials.fit(::Type{PQ}, xs::AbstractVector{S}, ys::AbstractVector{T}, m, n; var=:x) where {T,S, PQ<:RationalFunction}

    β₁,β₂ = gauss_newton(collect(xs), convert(Vector{float(T)}, ys), m, n)
    P = eltype(PQ)
    T′ = Polynomials._eltype(P) == nothing ? eltype(β₁) : eltype(P)
    X = indeterminate(PQ, var)
    P′ = constructorof(P){T′,X}
    p = P′(β₁)
    q = P′(vcat(1, β₂))

    p // q
end

"""
    fit(::Type{RationalFunction}, r::Polynomial, m, n; var=:x)

Fit a Pade approximant ([`pade_fit`](@ref)) to `r`.

Examples:

```jldoctext
julia> using Polynomials, PolynomialRatios

julia> x = variable()
Polynomial(x)

julia> ex = 1 + x + x^2/2 + x^3/6 + x^4/24 + x^5/120 # Taylor polynomial for e^x
Polynomial(1.0 + 1.0*x + 0.5*x^2 + 0.16666666666666666*x^3 + 0.041666666666666664*x^4 + 0.008333333333333333*x^5)

julia> maximum(abs, exp(x) - fit(RationalFunction, ex, 1,1)(x) for x ∈ 0:.05:0.5)
0.017945395966538547

julia> maximum(abs, exp(x) - fit(RationalFunction, ex, 1,2)(x) for x ∈ 0:.05:0.5)
0.0016624471707165078

julia> maximum(abs, exp(x) - fit(RationalFunction, ex, 2,1)(x) for x ∈ 0:.05:0.5)
0.001278729299871717

julia> maximum(abs, exp(x) - fit(RationalFunction, ex, 2,2)(x) for x ∈ 0:.05:0.5)
7.262205147950951e-5
```
"""
function Polynomials.fit(::Type{RationalFunction},r::Polynomial, m::Integer, n::Integer;var=:x)
    p,q = pade_fit(r, m,n, var=var)
    p // q
end

## ---- Pade
## https://mathworld.wolfram.com/PadeApproximant.html
"""
    pade_fit(r::Polynomial, m,n)

For a polynomial `r` of degree `d ≥ m + n`, find a rational function `p/q` with
`degree(p) ≤ m`, `degree(q) ≤ n` and `q*r - p = x^{m+n+1}*s(x)` for some polynomial `s`.

This implementation sets up a system of equations to identify `p` and `q`.
"""
function pade_fit(p::Polynomial{T}, m::Integer, n::Integer; var=:x) where {T}
    d = degree(p)
    @assert (0 <= m) && (1 <= n) && (m + n <= d)

    # could be much more perfomant
    c = convert(LaurentPolynomial, p) # for better indexing
    cs = [c[m+j-i] for j ∈ 1:n, i ∈ 0:n]

    qs′ = cs[:, 2:end] \ cs[:,1]
    qs = vcat(1, -qs′)

    cs = [c[0 + j - i] for j ∈ 0:m, i∈0:n]
    ps = cs * qs

    Polynomial(ps, var), Polynomial(qs,var)
end

## ---- Least Squares
## avoiding dependency on another package, LsqFit

using LinearAlgebra

# return pair ((a₀,...,aₙ), (b₁,...,bₘ))
function initial_guess(xs::Vector{T}, ys::Vector{S}, n, m) where {T, S}
    # p(x) = a₀ + ... + aₙx^n
    # q(x) = 1 + b₁x + ... + bₘx^m = 1 + r(x)
    # yᵢ + yᵢ * r(xᵢ) = p(xᵢ)
    # yᵢ = p(xᵢ) - r(xᵢ) * yᵢ
    k = n+1+m
    A = zeros(T, k, k)
    A[:,1] .= one(T)
    xs′ = xs[1:k]
    xs′′ = copy(xs′)
    for i ∈ 1:n
        A[:,1+i] = xs′
        xs′ .*= xs′′
    end
    xs′ = -copy(xs′′) .* ys[1:k]
    for i ∈ 1:m
        A[:, 1+n+i] = xs′
        xs′ .*= xs′′
    end
    β = pinv(A) *  ys[1:k]

end

function make_model(n)
    (x, β) -> begin
        β₁, β₂ = β[1:n+1], β[n+2:end]
        evalpoly.(x,(β₁,)) ./ evalpoly.(x, (vcat(1, β₂),))
    end
end

function J!(Jᵣ, xs::Vector{T}, β, n) where {T}
    β₁, β₂ = β[1:n+1],β[n+2:end]
    ps = evalpoly.(xs, (β₁,))
    qs = evalpoly.(xs, (vcat(1, β₂),))
    λ = one(T) ./ qs
    for i ∈ eachindex(β₁)
        Jᵣ[:,1] = λ
        λ .*= xs
    end
    λ = xs .* ps ./ (qs .* qs)
    for i ∈ eachindex(β₂)
        Jᵣ[:, n+1+i]=λ
        λ .*= xs
    end
    nothing
end

function gauss_newton(xs, ys::Vector{T}, n, m, tol=sqrt(eps(T))) where {T}

    β = initial_guess(xs, ys, n, m)
    model = make_model(n)

    Jᵣ = zeros(T, length(xs), 1 + n + m)

    Δ = norm(ys, Inf) * tol

    ϵₘ = norm(ys - model(xs, β), Inf)
    βₘ = copy(β)

    no_steps = 0

    while no_steps < 25
        no_steps += 1

        r = ys - model(xs, β)
        ϵ = norm(r, Inf)
        ϵ < Δ && return (β[1:n+1], β[n+2:end])
        if ϵ < ϵₘ
            ϵₘ = ϵ
            βₘ .= β
        end
        J!(Jᵣ, xs, β, n)
        Δᵦ = pinv(Jᵣ' * Jᵣ) * (Jᵣ' * r)
        β .-= Δᵦ

    end

    @warn "no convergence; returning best fit of many steps"
    return (βₘ[1:n+1], βₘ[n+2:end])
end

end


function fib(n)
    n < 3 ? 1 : fib(n-2) + fib(n-1)
end

"""
   This is code text.
"""
function memoize(f)
    memo = Dict()
    (args...; kwargs...) -> begin
        x = (args, kwargs)
        haskey(memo, x) || (memo[x] = f(args..., kwargs...))
        memo[x]
    end
end

fib! = memoize(fib)

fib!(20)


using Test

"""
    fahrenheit_to_celsius(temp_f::Float)

Converts temperature in Fahrenheit to Celsius.
"""
function fahrenheit_to_celsius(temp_f)
    temp_c = (temp_f - 32.0) * (5.0/9.0)
    return temp_c
end

# This is the test section
@testset "Test fahrenheit_to_celsius" begin
    temp_c = fahrenheit_to_celsius(100.0)
    expected_result = 37.777777
    @test abs(temp_c - expected_result) < 1.0e-6
end


using FractionalDiffEq, Plots

τ=3.1416; h=0.01; α=0.4; tspan = (0, 70)
x0(t) = [sin(t)*cos(t); sin(t)*cos(t); cos(t)^2-sin(t)^2; cos(t)^2-sin(t)^2]
A=[0 0 1 0; 0 0 0 1; 0 -2 0 0; -2 0 0 0]
B=[0 0 0 0; 0 0 0 0; -2 0 0 0; 0 -2 0 0]
f=[0; 0; 0; 0]

prob = FDDEMatrixProblem(α, τ, A, B, f, x0, tspan)
sol=solve(prob, h, MatrixForm())
plot(sol[:, 1], sol[:, 3])


# Transformation layer
data = [
    [1, 5, 2, 4, 3] => 2,
    [1, 2, 3, 2, 1] => 2,
]

# Test transformations without parameters
funcs = Dict(
    CN.tr_identity => [
        data[1].first,
        data[2].first,
    ],
    CN.tr_count_eq => [
        [0, 0, 0, 0, 0],
        [1, 1, 0, 1, 1],
    ],
    CN.tr_count_eq_right => [
        [0, 0, 0, 0, 0],
        [1, 1, 0, 0, 0],
    ],
    CN.tr_count_eq_left => [
        [0, 0, 0, 0, 0],
        [0, 0, 0, 1, 1],
    ],
    CN.tr_count_greater => [
        [4, 0, 3, 1, 2],
        [3, 1, 0, 1, 3],
    ],
    CN.tr_count_lesser => [
        [0, 4, 1, 3, 2],
        [0, 2, 4, 2, 0],
    ],
    CN.tr_count_g_left => [
        [0, 0, 1, 1, 2],
        [0, 0, 0, 1, 3],
    ],
    CN.tr_count_l_left => [
        [0, 1, 1, 2, 2],
        [0, 1, 2, 1, 0],
    ],
    CN.tr_count_g_right => [
        [4, 0, 2, 0, 0],
        [3, 1, 0, 0, 0],
    ],
    CN.tr_count_l_right => [
        [0, 3, 0, 1, 0],
        [0, 1, 2, 1, 0],
    ],
    CN.tr_contiguous_vals_minus => [
        [0, 3, 0, 1, 0],
        [0, 0, 1, 1, 0],
    ],
    CN.tr_contiguous_vals_minus_rev => [
        [4, 0, 2, 0, 0],
        [1, 1, 0, 0, 0],
    ],
)

for (f, results) in funcs
    @info f
    for (key, vals) in enumerate(data)
        @test f(vals.first) == results[key]
        foreach(i -> f(i, vals.first), vals.first)
    end
end

# Test transformations with parameter
funcs_param = Dict(
    CN.tr_count_eq_param => [
        [1, 0, 1, 0, 1],
        [1, 0, 0, 0, 1],
    ],
    CN.tr_count_l_param => [
        [2, 5, 3, 5, 4],
        [4, 5, 5, 5, 4],
    ],
    CN.tr_count_g_param => [
        [2, 0, 1, 0, 0],
        [0, 0, 0, 0, 0],
    ],
    CN.tr_count_bounding_param => [
        [3, 1, 3, 2, 3],
        [5, 3, 1, 3, 5],
    ],
    CN.tr_val_minus_param => [
        [0, 3, 0, 2, 1],
        [0, 0, 1, 0, 0],
    ],
    CN.tr_param_minus_val => [
        [1, 0, 0, 0, 0],
        [1, 0, 0, 0, 1],
    ],
)

for (f, results) in funcs_param
    @info f
    for (key, vals) in enumerate(data)
        @test f(vals.first; param=vals.second) == results[key]
        foreach(i -> f(i, vals.first; param=vals.second), vals.first)
    end
end

# arithmetic layer
@info CN.ar_sum
@test CN.ar_sum(map(p -> p.first, data)) == [2, 7, 5, 6, 4]
@info CN.ar_prod
@test CN.ar_prod(map(p -> p.first, data)) == [1, 10, 6, 8, 3]

# aggregation layer
@info CN.ag_sum
@test CN.ag_sum(data[1].first) == 15
@test CN.ag_sum(data[2].first) == 9

@info CN.ag_count_positive
@test CN.ag_count_positive(data[1].first) == 5
@test CN.ag_count_positive(data[2].first) == 5
@test CN.ag_count_positive([1, 0, 1, 0, 1]) == 3

# Comparison layer
data = [3 => (1, 5), 5 => (10, 5)]

funcs = [
    CN.co_identity => [3, 5],
]

# test no param/vars
for (f, results) in funcs
    @info f
    for (key, vals) in enumerate(data)
        @test f(vals.first) == results[key]
    end
end

funcs_param = [
    CN.co_abs_diff_val_param => [2, 5],
    CN.co_val_minus_param => [2, 0],
    CN.co_param_minus_val => [0, 5],
]

for (f, results) in funcs_param
    @info f
    for (key, vals) in enumerate(data)
        @test f(vals.first; param=vals.second[1]) == results[key]
    end
end

funcs_vars = [
    CN.co_abs_diff_val_vars => [2, 0],
    CN.co_val_minus_vars => [0, 0],
    CN.co_vars_minus_val => [2, 0],
]

for (f, results) in funcs_vars
    @info f
    for (key, vals) in enumerate(data)
        @test f(vals.first, nvars=vals.second[2]) == results[key]
    end
end

funcs_param_dom = [
    CN.co_euclidian_param => [1.4, 2.0],
]

for (f, results) in funcs_param_dom
    @info f
    for (key, vals) in enumerate(data)
        # @info "Updated" f(vals.first, param=vals.second[1], dom_size=vals.second[2]) results key
        @test f(vals.first, param=vals.second[1], dom_size=vals.second[2]) ≈ results[key]
    end
end

funcs_dom = [
    CN.co_euclidian => [1.6, 2.0],
]

for (f, results) in funcs_dom
    @info f
    for (key, vals) in enumerate(data)
        # @info "Updated" f(vals.first, dom_size=vals.second[2]) results key
        @test f(vals.first, dom_size=vals.second[2]) ≈ results[key]
    end
end


using Metatheory

@testset "Reduction Basics" begin
	t = @theory begin
	    a + a => 2a
	    x / x => 1
	    x * 1 => x
	end

    # basic theory to check that everything works
    @test rewrite(:(a + a), t) == :(2a)
    @test rewrite(:(a + a), t) == :(2a)
    @test rewrite(:(a + (x * 1)), t) == :(a + x)
	@test rewrite(:(a + (a * 1)), t; order=:inner) == :(2a)
end

import Base.(+)
@testset "Extending Algebra Operators" begin
    t = @theory begin
        a + a => 2a
    end

    # Let's extend an operator from base, for sake of example
    function +(x::Symbol, y)
        rewrite(:($x + $y), t)
    end

    @test (:x + :x) == :(2x)
end

## Free Monoid

@testset "Free Monoid - Overriding identity" begin
    # support symbol literals
	symbol_monoid = @theory begin
		a ⋅ :ε => a
		:ε ⋅ a => a
		a::Symbol => a
		a::Symbol ⋅ b::Symbol |> Symbol(String(a) * String(b))
		# i |> error("unsupported ", i)
	end;

    @test rewrite(:(ε ⋅ a ⋅ ε ⋅ b ⋅ c ⋅ (ε ⋅ ε ⋅ d) ⋅ e), symbol_monoid; order=:inner) == :abcde
end

## Interpolation should be possible at runtime

@testset "Calculator" begin
	calculator = @theory begin
		x::Number ⊕ y::Number |> x + y
		x::Number ⊗ y::Number |> x * y
		x::Number ⊖ y::Number |> x ÷ y
		x::Symbol => x
		x::Number => x
	end;
	a = 10

	@test rewrite(:(3 ⊕ 1 ⊕ $a), calculator; order=:inner) == 14
end

# FIXME add support for splatting
# @testset "Type assertions and destructuring" begin
#     # let's try type assertions and destructuring
#     t = @theory begin
#         f(a::Number) => a
#         f(a...) => a
#     end

#     @test rewrite(:(f(3)), t) == 3
#     @test rewrite(:(f(2, 3)), t) == [2, 3]

#     # destructuring in right hand
#     n = @theory begin
#         f(a...) => +(a...)
#     end

#     @test rewrite(:(f(2, 3)), n) == :(2 + 3)
#     @test rewrite(:(f(a, b, c, d)), n) == :(((a + b) + c) + d)
# end

## Direct rules
@testset "Direct Rules" begin
    t = @theory begin
        # maps
        a * b |> ((a isa Number && b isa Number) ? a * b : _lhs_expr)
    end
    @test rewrite(:(3 * 1), t) == 3

    t = @theory begin
        # maps
        a::Number * b::Number |> a * b
    end
    @test rewrite(:(3 * 1), t) == 3
end

## Take advantage of subtyping.
# Subtyping in Julia has been formalized in this paper
# [Julia Subtyping: A Rational Reconstruction](https://benchung.github.io/papers/jlsub.pdf)

abstract type Vehicle end
abstract type GroundVehicle <: Vehicle end
abstract type AirVehicle <: Vehicle end
struct Airplane <: AirVehicle end
struct Car <: GroundVehicle end

airpl = Airplane()
car = Car()

@testset "Subtyping" begin
	t = @theory begin
		a::AirVehicle * b => "flies"
		a::GroundVehicle * b => "doesnt_fly"
	end

	sf = rewrite(:($airpl * c), t)
	df = rewrite(:($car * c), t)

    @test sf == "flies"
    @test df == "doesnt_fly"
end

airpl = Airplane()
car = Car()

@testset "Interpolation" begin
	t = @theory begin
		$airpl * b => "flies"
		$car * b => "doesnt_fly"
	end

	sf = rewrite(:($airpl * c), t)
	df = rewrite(:($car * c), t)

    @test sf == "flies"
    @test df == "doesnt_fly"
end

## Multiplicative Monoid

# TODO consider autocompiling from src/Library. This is very ugly.

# FIXME TYPE VARIABLES IN CLASSICAL MATCHCORE ALGORITHM
macro red_commutative_monoid(t, op, id, func)
	zop = Meta.quot(op)
   quote
	   @theory begin
			($op)($id, a) => a
			($op)(a, $id) => a
			# closure
			($op)(a::$t, b::$t) |> ($func)(a,b)
			# associativity reductions implying commutativity
			($op)(a::$t, ($op)(b::$t, c)) |>
				(z = ($func)(a, b); Expr(:call, ($zop), z, c))

			($op)(($op)(a::$t, c), b::$t) |>
				(z = ($func)(a, b); Expr(:call, ($zop), z, c))

			($op)(a::$t, ($op)(c, b::$t)) |>
				(z = ($func)(a, b); Expr(:call, ($zop), z, c))

			($op)(($op)(c, a::$t), b::$t) |>
				(z = ($func)(a, b); Expr(:call, ($zop), z, c))
	   end
   end
end

mult_monoid = @red_commutative_monoid Int (*) 1 (*)

@testset "Multiplication Monoid" begin
	res_macro = rewrite(:(3 * (1 * 2)), mult_monoid)
	res_sym = rewrite(:(3 * (1 * 2)), mult_monoid; order=:inner)
	res_macro_2 = rewrite(:(3 * (1 * 2)), mult_monoid; order=:inner)
	res_macro_3 = rewrite(:(2a * (3 * 2)), mult_monoid; order=:inner)

	@test res_macro == 6
	@test res_macro_2 == 6
	@test res_sym == 6
	@test res_macro_3 == :(12a)
end

# constructs a semantic theory about a an abelian group
# The definition of a group does not require that a ⋅ b = b ⋅ a
# for all elements a and b in G. If this additional condition holds,
# then the operation is said to be commutative, and the group is called an abelian group.
macro red_abelian_group(t, op, id, invop, func)
	zop = Meta.quot(op)
   quote
	   @theory begin
			# identity element
			($op)($id, a) => a
			($op)(a, $id) => a
			# inversibility
			($op)(a, ($invop)(a)) => $id
			($op)(($invop)(a), a) => $id
			# closure
			($op)(a::$t, b::$t) |> ($func)(a,b)
			# inverse
			# associativity reductions
			($op)(a::$t, ($op)(b::$t, c)) |> (z = ($func)(a, b); Expr(:call, ($zop), z, c))
			($op)(($op)(a::$t, c), b::$t) |> (z = ($func)(a, b); Expr(:call, ($zop), z, c))
			($op)(a::$t, ($op)(c, b::$t)) |> (z = ($func)(a, b); Expr(:call, ($zop), z, c))
			($op)(($op)(c, a::$t), b::$t) |> (z = ($func)(a, b); Expr(:call, ($zop), z, c))
	   end
   end
end

addition_group = @red_abelian_group Int (+) 0 (-) (+) ;
@testset "Addition Abelian Group" begin
	zero = rewrite(:((x+y) +  -(x+y)), addition_group)
	@test zero == 0
end

# distributivity of two operations
# example: `@distrib (⋅) (⊕)`
macro red_distrib(outop, inop)
	@assert Base.isbinaryoperator(outop)
	@assert Base.isbinaryoperator(inop)
	quote
		@theory begin
			($outop)(a, ($inop)(b,c)) => ($inop)(($outop)(a,b),($outop)(a,c))
			($outop)(($inop)(a,b), c) => ($inop)(($outop)(a,c),($outop)(b,c))
		end
	end
end
distr = @red_distrib (*) (+)
Z = mult_monoid ∪ addition_group ∪ distr;
@testset "Composing Theories, distributivity" begin
	res_1 = rewrite(:(((2 + (b + -b)) * 3) * (a + b)), Z; order=:inner)
	e = rewrite(:((2 * (3 + b + (4 * 2)))), Z; order=:inner)

	@test_skip res_1 == :(6a+6b)
	@test_skip e == :(22+2b)
end

logids = @theory begin
	x * x => x^2
	x^n * x |> :($x^($(n+1)))
	x * x^n |> :($x^($(n+1)))
	a * (a * b) => a^2 * b
	a * (b * a) => a^2 * b
	(a * b) * a => a^2 * b
	(b * a) * a => a^2 * b
	# catch e as ℯ
	:e => ℯ
	log(a^n) => n * log(a)
	log(x * y) => log(x) * log(y)
	log(1) => 0
	log(:ℯ) => 1
	:ℯ^(log(x)) => x
end;

t = logids ∪ Z;
@testset "Symbolic Logarithmic Identities, Composing Theories" begin
    @test rewrite(:(log(ℯ)), t) == 1
    @test rewrite(:(log(x)), t) == :(log(x))
    @test rewrite(:(log(ℯ^3)), t) == 3
    @test rewrite(:(log(a^3)), t) == :(3 * log(a))
    # Reduce.jl wtf u doing? log(:x^-2 * ℯ^3) = :(log(5021384230796917 / (250000000000000 * x ^ 2)))

    @test rewrite(:(log(x^2 * ℯ^3)), t) == :((6 * log(x)))
    @test_skip rewrite(:(log(x^2 * ℯ^(log(x)))), t; order=:inner) == :(3 * log(x))
end

using Calculus: differentiate
diff = @theory begin
	∂(y)/∂(x) |> differentiate(y, x)
	a * (1/x) => a/x
	a * 0 => 0
	0 * a => 0
end
finalt = t ∪ diff
@testset "Symbolic Differentiation, Accessing Outer Variables" begin
	@test rewrite(:(∂( log(x^2 * ℯ^(log(x))) )/∂(x)), finalt) == :(3/x)
end;

## let's loop this
# @testset "Reduction loop should return." begin
#     t = @theory begin
#         a + b => b + a
#     end

#     @test :(a+b) == rewrite(:(a + b), t)

# 	t = @theory begin
# 	    a + b => b + a
# 	    b + a => a + b
# 	end

# 	@test :(a+b) == rewrite(:(a + b), t)
# end

# safe_var = 0
# ref_var = Ref{Real}(0)

# reft = @theory begin
# 	:safe |> (safe_var = π)
# 	:ref |> (ref_var[] = π)
# end

# @testset "Evil Assignment" begin
# 	rewrite(:(safe), reft; order=:inner)
# 	rewrite(:(ref), reft; order=:inner)

# 	@test safe_var == 0
# 	@test ref_var[] == π

# 	rewrite(ref, reft)
# 	rewrite(safe, reft)

# 	@test safe_var == 0
# 	@test ref_var[] == π
# end;

# dollart = @theory begin
# 	# dollars are interpolated at rule creation time?
# 	$(Expr(:call, :foo, :bar)) => :match
# 	:x + $(3+2) => :weeee
# end

# @testset "Dollar Escapes" begin
# 	r1 = rewrite(:(foo(bar)), dollart; m=@__MODULE__)
# 	r2 = rewrite(:(x + 5), dollart; m=@__MODULE__)

# 	@test r1 == :match
# 	@test r2 == :weeee
# end;


# coding: utf-8

# In[1]:

#bubble
#bubble sort is the brute calculation
#the time complexity can go to n*(n-1), quadratic
#basically we create iterations to compare each element
#for the first round, we try to compare item i with the item right next to it
#we swap the place of two items if the item next to item i is smaller
#we keep applying this rule to the list during the traversal
#in the end, we must have placed the largest item at the end of list
#for the next iteration, we do the same trick
#except the length of traversal is reduced by one
#cuz we have already placed the largest item at the end of list
#we just need to find the second largest and place it right next to the end
#for the third iteration, we try to find the third largest item, etc

# In[2]:

function bubble_sort(arr)

    #for each round n, we find the nth largest number
    #we use reverse order from 1 to end
    #dont forget range starts from 1 in julia
    for i in reverse(1:length(arr))

        #given n, we go through n items in the list
        #compare then swap
        #range always includes the first and the last
        #very different from python
        #thus we use i-1 in the inner loop
        for j in 1:i-1

            #if the item is larger than the one before it
            #we do a simultaneous swap
            if arr[j]>arr[j+1]
                arr[j],arr[j+1]=arr[j+1],arr[j]
            end

        end

    end

    return arr

end

# In[3]:

for _ in 1:100

    test_arr=rand(1000)

    if !(bubble_sort(test_arr)==sort(test_arr))
        printstyled("Erreur",color=:red)
    end

end

# In[4]:

#selection
#selection sort is a simplified version of bubble sort
#selection sort tries to find the nth largest number for each round n
#the next round is n-1, the same as bubble sort
#however, we only do traversal for each round
#we dont make any swap at all
#we simply use extra variable for the largest number and its index throughout traversal
#once we complete the traversal
#we move the largest number to the end of the list with the index we stored
#we repeat this process for each round

# In[5]:

function selection_sort(arr)

    #the first loop is similar bubble sort
    #the difference is that we are not going to use 1
    #we use 2 instead
    #the reason is that we would be left with one element for the last round
    #there is no point of doing so
    for i in reverse(2:length(arr))

        max_value=-Inf
        max_index=NaN

        #the second loop is the same as bubble sort
        for j in 1:i

            #if the value is larger than the extra one
            #we replace the extra value and its index
            if arr[j]>max_value
                max_value=arr[j]
                max_index=j
            end
        end

        #once we find out the ith largest number
        #we swap it with the item at the end of unsorted list
        arr[max_index],arr[i]=arr[i],arr[max_index]

    end
    return arr
end

# In[6]:

for _ in 1:100

    test_arr=rand(1000)

    if !(selection_sort(test_arr)==sort(test_arr))
        printstyled("Erreur",color=:red)
    end

end

# In[7]:

#insertion
#insertion sort is like the reverse of bubble sort
#for bubble sort, you find nth largest number and
#put it at the end throughout iterations
#then we do the same trick on the unsorted part
#insertion is quite the opposite
#we are always dealing with the sorted list
#we start from the first two item
#we do comparison and swap if necessary
#after that we expand the sorted list
#we introduce the next item to the sorted list
#as the sorted list is sorted
#we can reduce the rounds of comparison
#as long as we find where the new item truly belongs to
#we insert it there

# In[8]:

function insertion_sort(arr)

    #finally! we are doing ascending order range function
    #it is straight forward
    #we dont wanna waste memory with one item so we start at 2
    for i in 2:length(arr)

        #we need to create an extra value to store the next item
        target=arr[i]

        #we still use j for second loop
        #we use while instead of if this time
        #we run a descending order loop on j
        #so we initialize j with i
        #which is the length of sorted target
        j=i

        #we do comparison and use j-=1 as iteration
        #j!=1 is to make sure the iteration doesnt go outta the length of the target
        while (j!=1) && (target<arr[j-1])

            #we dont swap
            #we simply shift it backwards
            arr[j]=arr[j-1]
            j-=1
        end

        #when we finish comparison and find out the index where we need to insert
        #we insert it
        arr[j]=target

    end

    return arr

end

# In[9]:

for _ in 1:100

    test_arr=rand(1000)

    if !(insertion_sort(test_arr)==sort(test_arr))
        printstyled("Erreur",color=:red)
    end

end


"""
    BSTNode

Binary node.
"""
mutable struct BSTNode{S,T}
    key::Union{S, Nothing}
    value::Union{T, Nothing}
    left::Union{BSTNode{S,T}, Nothing}
    right::Union{BSTNode{S,T}, Nothing}
    size::Int
    BSTNode{S,T}() where {S,T} = new(nothing, nothing, nothing, nothing, 0)
    BSTNode{S,T}(k::S, v::T) where {S,T} = new(k, v, nothing, nothing, 0)
    BSTNode{S,T}(k::S, v::T, size::Int) where {S,T} = new(k, v, nothing, nothing, size)
    BSTNode{S,T}(k::S, v::T, l::BSTNode{S,T}) where {S,T} = new(k, v, l, nothing, 0)
    BSTNode{S,T}(k::S, v::T, l::BSTNode{S,T}, r::BSTNode{S,T}) where {S,T} = new(k, v, l, r, 0)
end

"""
    BST

Binary search tree. Does not support duplicate keys.
"""
mutable struct BST{S,T}
    root::Union{BSTNode{S,T}, Nothing}
    size::Int
    BST{S,T}() where {S,T} = x = new(BSTNode{S,T}(), 0)
end

"""
    get!(bst, key)
    get!(node, key)

Get item from tree given key. Time complexity O(logn), worst O(n).
"""
get!(bst::BST{S,T}, key::S) where {S,T} = get!(bst.root, key)
get!(bst::BST{S,T}, key::S, parent) where {S,T} = get!(bst.root, key, parent)
function get!(node::BSTNode{S,T}, key::S) where {S,T}
    if node.key === nothing
        return nothing
    elseif key < node.key
        return get!(node.left, key)
    elseif key > node.key
        return get!(node.right, key)
    else
        return node.value
    end
end

"""
    put!(bst, key, value)
    put!(node, key, value)

Put item into tree. Time complexity O(logn), worst O(n).
"""
function put!(bst::BST{S,T}, key::S, value::T) where {S,T}
    bst.root = put!(bst.root, key, value)
    bst.size = size(bst.root)
    return nothing
end

function put!(node::Union{BSTNode{S,T}, Nothing}, key::S, value::T) where {S,T}
    if node === nothing || node.key === nothing
        return BSTNode{S,T}(key, value, 1)
    elseif key < node.key
        node.left = put!(node.left, key, value)
    elseif key > node.key
        node.right = put!(node.right, key, value)
    else
        node.value = value
    end

    node.size = (node.left !== nothing ? size(node.left) : 0) +
                (node.right !== nothing ? size(node.right) : 0) +
                1
    return node
end

"""
    delete!(bst, key)
    delete!(node, parent, key)

Remove node from tree. Time complexity O(logn), worst O(n).

Implements two versions. First is based on lecutre slides (my implementation) with one traversal to find maximum node. The other is the Hibbard deletion which makes two traversals to find and delete maximum node. The Hibbard deletion is much slower. Why?
"""
function delete!(bst::BST{S,T}, key::S) where {S,T}
    delete!(bst.root, bst.root, key)
    #bst.root = delete!(bst.root, key) # Hibbard deletion
    bst.size = bst.root === nothing ? 0 : size(bst.root)
    return nothing
end

function delete!(node::Union{BSTNode{S,T}, Nothing}, parent::Union{BSTNode{S,T}, Nothing}, key::S) where {S,T}
    if node === nothing return nothing end

    # check this first, if last extremely slow due to many recursion calls
    if node.key < key
        delete!(node.right, node, key)
    elseif node.key > key
        delete!(node.left, node, key)
    else
        if node.left === nothing && node.right === nothing
            setparent!(parent, node)
        elseif node.left === nothing
            setparent!(parent, node, node.right)
        elseif node.right === nothing
            setparent!(parent, node, node.left)
        else
            n, p = maximum_node(node.left, node)
            node.key = n.key
            node.value = n.value
            setparent!(p, n)
        end
    end

    node.size = (node.left !== nothing ? size(node.left) : 0) +
                (node.right !== nothing ? size(node.right) : 0) +
                1
    return nothing
end

"""
    setparent!(p, n, v = nothing)

Helper function for delete.
"""
function setparent!(p::BSTNode, n::BSTNode, v::Union{BSTNode, Nothing} = nothing)
    if p.left == n
        p.left = v
    elseif p.right == n
        p.right = v
    else
        p = v
    end
    return nothing
end

"""
    maximum_node(node, parent)

Find maximum node and its parent.
"""
function maximum_node(node::BSTNode, parent::BSTNode)
    if node.right !== nothing
        node.size -= 1
        node, parent = maximum_node(node.right, node)
    end

    return node, parent
end

# Hibbard deletion, from sedgewick and wayne algorithms
function delete!(node::Union{BSTNode{S,T}, Nothing}, key::S) where {S,T}
    if node === nothing return nothing end

    if key < node.key
        node.left = delete!(node.left, key)
    elseif key > node.key
        node.right = delete!(node.right, key)
    else
        if node.right === nothing return node.left end
        if node.left === nothing return node.right end
        t = node

        node = maximum(t.left) # 1st traversal
        node.left = deletemax!(t.left) # 2nd traversal
        node.right = t.right
    end

    node.size = (node.left !== nothing ? size(node.left) : 0) +
                (node.right !== nothing ? size(node.right) : 0) +
                1
    return node
end

"""
    deletemax!(n::BSTNode)

Help function for Hibbard deletion.
"""
function deletemax!(node::BSTNode)
    if node.right === nothing return node.left end
    node.right = deletemax!(node.right)
    node.size = (node.left !== nothing ? size(node.left) : 0) +
                (node.right !== nothing ? size(node.right) : 0) +
                1
    return node
end

"""
    printsort(bst)

Print binary search tree in order.
"""
printsort(bst::BST) = printsort(bst.root)
function printsort(node::BSTNode)
    if node.key === nothing return nothing end

    if node.left !== nothing printsort(node.left) end
    print(" $(node.key) ")
    if node.right !== nothing printsort(node.right) end
end

# """
#     iterate(l, n)

# Iterate over a linked list.
# """
# function iterate(bst::Union{BST{S,T}, BSTNode{S,T}}, r = bst.root) where {S,T}
#     @show r.key
#     if isdefined(r, :left)
#         return r = (r.key, r.left)
#     end
#     if isdefined(r, :right)
#         return r = (r.key, r.right)
#     end
#     #r = nothing
# end

# Could we use iterate to get this?
"""
    collect(bst, l)
    collect(bst, l)

Return a list of BST keys in order.
"""
function collect(bst::BST)
    l = Vector{eltype(bst)}(undef, 0)
    collect(bst.root, l)
    l
end

function collect(node::BSTNode, l)
    if node.key === nothing return nothing end

    if node.left !== nothing collect(node.left, l) end
    push!(l, node.key)
    if node.right !== nothing collect(node.right, l) end
end

"""
    eltype(type)

Type of key elements.
"""
eltype(::BST{S,T}) where {S,T} = S

"""
    length(bst)
    size(bst)
    size(node)

Size of BST or node.
"""
length(bst::BST) = bst.size
size(bst::BST) = length(bst)
size(n::BSTNode) = n.size

"""
    minimum(bst)
    minimum(node)

Find minimum node.
"""
minimum(bst::BST) = minimum(bst.root)
function minimum(node::BSTNode)
    if node.left === nothing
        return node
    end

    return minimum(node.left)
end

"""
    maximum(bst)

Find maximum node.
"""
maximum(bst::BST) = maximum(bst.root)
function maximum(node::BSTNode)
    if node.right === nothing
        return node
    end

    return maximum(node.right)
end

"""
    isbst(node)

Check if tree is a BST, i.e. if the invariant holds.
"""
isbst(bst::BST) = isbst(bst.root)
function isbst(node::Union{BSTNode, Nothing})
    if node === nothing return true end

    if node.left !== nothing
        if node.left.key < node.key
            isbst(node.left)
        else
            return false
        end
    end
    if node.right !== nothing
        if node.right.key > node.key
            isbst(node.right)
        else
            return false
        end
    end

    return true
end


# ## log plot
# ### semilog y

#md # ![FILE_NAME.png](images/FILE_NAME.png)

plot(exp, -5, 5, yscale=:log10, title="semilogy")


using Gen
using Test
import Random

"""
Compute a numerical partial derivative of `f` with respect to the `i`th
argument using finite differences.

If `broadcast` is `false`, then `args[i]` must be a scalar.

If `broadcast` is `true`, then `args[i]` may be an array.  In that case, `f` is
assumed to be the broadcast of a function whose `i`th argument is a scalar.  In
particular, `f` must still operate independently on each element of `args[i]`.
This condition cannot be checked automatically for arbitrary `f::Function`, so
the caller must guarantee it.
"""
function finite_diff(f::Function, args::Tuple, i::Int, dx::Float64;
                     broadcast=false)
    pos_args = Any[args...]
    neg_args = Any[args...]
    if broadcast
        pos_args[i] = copy(args[i]) .+ dx
        neg_args[i] = copy(args[i]) .- dx
        return (f(pos_args...) - f(neg_args...)) ./ (2. * dx)
    else
        pos_args[i] += dx
        neg_args[i] -= dx
        return (f(pos_args...) - f(neg_args...)) / (2. * dx)
    end
end

function finite_diff_vec(f::Function, args::Tuple, i::Int, j::Int, dx::Float64)
    pos_args = Any[deepcopy(args)...]
    pos_args[i][j] += dx
    neg_args = Any[deepcopy(args)...]
    neg_args[i][j] -= dx
    return (f(pos_args...) - f(neg_args...)) / (2. * dx)
end

"""
Compute a numerical partial derivative of `f` with respect to the `i`th argument, a symettric matrix, using finite differences.

Finite differences are applied symettrically to off-diagonal elements of the matrix, ensuring that the modified matrix is still symettric.

When applied to off-diagonal elements, the finite difference is scaled by a factor of 1/2.

The condition that the `i`th argument is a symettric matrix cannot be checked automatically, so the caller must guarantee it.
"""

function finite_diff_mat_sym(f::Function, args::Tuple, i::Int, j::Int, k::Int, dx::Float64)
    pos_args = Any[deepcopy(args)...]
    pos_args[i][j, k] += dx
    neg_args = Any[deepcopy(args)...]
    neg_args[i][j, k] -= dx

    if j!=k
        pos_args[i][k, j] += dx
        neg_args[i][k, j] -= dx
        return (f(pos_args...) - f(neg_args...)) / (4. * dx)
    end

    return (f(pos_args...) - f(neg_args...)) / (2. * dx)
end

const dx = 1e-6

include("autodiff.jl")
include("diff.jl")
include("selection.jl")
include("assignment.jl")
include("dynamic_dsl.jl")
include("optional_args.jl")
include("static_ir/static_ir.jl")
include("static_dsl.jl")
include("tilde_sugar.jl")
include("inference/inference.jl")
include("modeling_library/modeling_library.jl")


using Compose
function sierpinski(n)
  if n == 0
    compose(context(), polygon([(1,1), (0,1), (1/2, 0)]));
  else
    t = sierpinski(n - 1);
    compose(	context(),	(context( 1/4, 0, 1/2, 1/2), t),
			 	(context( 0, 1/2, 1/2, 1/2), t),
				(context( 1/2, 1/2, 1/2, 1/2), t));
  end
end

cx1 = compose(sierpinski(1), linewidth(0.2mm),
			fill(nothing), stroke("black"));
img = SVG("sierp1.svg", 10cm, 8.66cm); draw(img,cx1)

cx3 = compose(sierpinski(3), linewidth(0.2mm),
			fill(nothing), stroke("black"));
img = SVG("sierp3.svg", 10cm, 8.66cm); draw(img,cx3)

cx5 = compose(sierpinski(5), linewidth(0.2mm),
		fill(nothing), stroke("black"));
img = SVG("sierp5.svg", 10cm, 8.66cm); draw(img,cx5)


using Romberg
using Test
using Trapz

@testset "Romberg.jl" begin
    # Test interfaces
    x = range(0, π, length=2^8+1)
    y = sin.(x)

    @test romberg(x, y) == romberg(x, y, 8)

    max_steps = 7
    R = zeros(max_steps+1, max_steps+1)
    @test romberg!(R, x, y)[end] == romberg(x, y, max_steps)

    # Test ispow2(length(x) - 1) == false
    x = range(0, π, length=2^8)
    y = sin.(x)
    @test romberg(x, y) ≈ 2

    x = range(0, π, length=2^8-1)
    y = sin.(x)
    @test romberg(x, y) ≈ 2

    # Test `max_steps` too large
    @test_throws DomainError romberg(x, y, 9)

    # Test length(x) == 1
    x = 0:0
    y = sin.(x)
    R = zeros(1, 1)
    @test romberg(x, y) == 0
    @test romberg!(R, x, y) == zeros(1,1)

    # Test length(x) == 2
    x = range(0, 1, length=2)
    y = x.^2
    @test romberg(x,y) ≈ trapz(x, y)  # integration is inaccurate

    # Test length(x) == 3
    x = range(0, 1, length=3)
    y = x.^2
    @test romberg(x, y) ≈ 1/3

    # Integrate different functions
    x = range(0, π, length=2^8+1)
    y = sin.(x)
    @test romberg(x, y) ≈ 2

    x = range(0, 1, length=2^8+1)
    y = x.^3
    @test romberg(x, y) ≈ 1/4

    x = range(0, π/2, length=2^5+1)
    y = sin.(x).^2
    @test romberg(x, y) ≈ π/4

    m = 3
    n = 4
    x = range(0, π, length=2^8+1)
    y = sin.(m*x).*cos.(n*x)
    @test romberg(x, y) ≈ 2*m/(m^2 - n^2)

    # this one requires lots of samples...
    a = 15
    x = range(0, a, length=2^16+1)
    y = sqrt.(a^2 .- x.^2)
    @test romberg(x, y) ≈ π*a^2/4

    # tricky to integrate
    x = range(1e-15, 1, length=2^16+1)
    y = log.(x)./(1 .+ x)
    @test romberg(x, y) ≈ -π^2/12 atol=1e-4
end


include("cluster.jl")

Nx = Ny = 16
@show steps = 20000

Jc = 0.4406868                              # critical temp reduced coupling const
J = 1.0                                     # coupling constant in Kelvin
Tc = J/Jc                                   # critical temperature in Kelvin
println("Critical temperature Tc = $Tc")

# Try different lattice sizes
Ns = [16, 32, 128]
# 17 distinct temperatures
Ts = [0.5, 1.0, 1.25, 1.5:0.1:2.2, 2.22, 2.24, 2.26, 2.265, 2.27, 2.28, 2.3:0.1:3.0]
# magnetic fields
Bs = [0.00001, 0.0001, 0.001, 0.01, 0.1]

Tvals = [0.001:0.001:2.269]
Mvals = M(Tvals, J)

# maximum correlation distance
mcd = 12

print("Importing Gadfly... ")
import Gadfly
Gadfly.set_default_plot_size(20Gadfly.cm, 12Gadfly.cm)
println("done.")
# for plotting results
function MeansvsT(means, b, desc::String)

    layers = Array(Gadfly.Layer,0);
    for n in 1:length(Ns)
        newlayer = Gadfly.layer(
            x=Ts, y=means[:,n], Gadfly.Geom.point, Gadfly.Geom.line,
            color=["Simulated for N=$(Ns[n])"]);
        layers = cat(1, layers, newlayer)
    end

    return Gadfly.plot(
        layers,
        Gadfly.layer(
            x=Tvals, y=Mvals, Gadfly.Geom.line,
            color=["Predicted magnetization"]),
        Gadfly.layer(
            x=[Tc,Tc], y=[0,1], Gadfly.Geom.line,
            color=["Critical Temperature 2.269K"]),
        Gadfly.Guide.xticks(ticks=[0:0.5:2, 2.269, 2.5, 3]),
        Gadfly.Guide.yticks(ticks=[0:0.1:1]),
        Gadfly.Guide.xlabel("Temperature"),
        Gadfly.Guide.ylabel("Mean normalized absolute magnetization"),
        Gadfly.Guide.title("$desc vs. Temperature for B=$(Bs[b]) at Different Lattice Sizes")
    )
end

# for plotting results
μvsT(b) = Gadfly.plot(
    Gadfly.layer(
        x=Ts, y=μs[:,b], Gadfly.Geom.point, Gadfly.Geom.line,
        color=["Measured for B=$(Bs[b])"]),
    Gadfly.layer(
        x=Tvals, y=Mvals, Gadfly.Geom.line,
        color=["Predicted magnetization"]),
    Gadfly.Guide.xlabel("Temperature"),
    Gadfly.Guide.ylabel("Mean normalized absolute magnetization"),
    Gadfly.Guide.title("Magnetization vs. Temperature for B=$(Bs[b]) with $(Nx*Ny) Particles")
)

# see magnetization distribution
μhist(mags) = Gadfly.plot(
    Gadfly.Guide.title("Magnetization Histogram for B = 0.1"),
    Gadfly.layer(
        x=abs(mags), color=["Absolute Magnetization"],
        Gadfly.Geom.histogram(bincount=100, density=true), order=1),
    Gadfly.layer(
        x=mags, color=["Magnetization"],
        Gadfly.Geom.histogram(bincount=100, density=true), order=2)
)

# Make plots of the data
absmeans = Array(Float64, length(Ts), length(Ns), length(Bs))
absmeans[:,:,1] = readdlm("absmeans-1.0e-5.dat")
absmeans[:,:,2] = readdlm("absmeans-0.0001.dat")
absmeans[:,:,3] = readdlm("absmeans-0.001.dat")
absmeans[:,:,4] = readdlm("absmeans-0.01.dat")
absmeans[:,:,5] = readdlm("absmeans-0.1.dat")

means = Array(Float64, length(Ts), length(Ns), length(Bs))
means[:,:,1] = readdlm("means-1.0e-5.dat")
means[:,:,2] = readdlm("means-0.0001.dat")
means[:,:,3] = readdlm("means-0.001.dat")
means[:,:,4] = readdlm("means-0.01.dat")
means[:,:,5] = readdlm("means-0.1.dat")

plots = Array(Gadfly.Plot, length(Bs))
for i in 1:length(plots)
    plots[i] = MeansvsT(means[:,:,i], i, "Magnetization")
end

absplots = Array(Gadfly.Plot, length(Bs))
for i in 1:length(plots)
    absplots[i] = MeansvsT(absmeans[:,:,i], i, "Absolute Magnetization")
end


# # Polarizability using automatic differentiation
#
# Simple example for computing properties using (forward-mode)
# automatic differentation.
# For a more classical approach and more details about computing polarizabilities,
# see [Polarizability by linear response](@ref).

using DFTK
using LinearAlgebra
using ForwardDiff

## Construct PlaneWaveBasis given a particular electric field strength
## Again we take the example of a Helium atom.
function make_basis(ε::T; a=10., Ecut=30) where T
    lattice=T(a) * I(3)  # lattice is a cube of ``a`` Bohrs
    ## Helium at the center of the box
    atoms     = [ElementPsp(:He, psp=load_psp("hgh/lda/He-q2"))]
    positions = [[1/2, 1/2, 1/2]]

    model = model_DFT(lattice, atoms, positions, [:lda_x, :lda_c_vwn];
                      extra_terms=[ExternalFromReal(r -> -ε * (r[1] - a/2))],
                      symmetries=false)
    PlaneWaveBasis(model; Ecut, kgrid=[1, 1, 1])  # No k-point sampling on isolated system
end

## dipole moment of a given density (assuming the current geometry)
function dipole(basis, ρ)
    @assert isdiag(basis.model.lattice)
    a  = basis.model.lattice[1, 1]
    rr = [a * (r[1] - 1/2) for r in r_vectors(basis)]
    sum(rr .* ρ) * basis.dvol
end

## Function to compute the dipole for a given field strength
function compute_dipole(ε; tol=1e-8, kwargs...)
    scfres = self_consistent_field(make_basis(ε; kwargs...), tol=tol)
    dipole(scfres.basis, scfres.ρ)
end;

# With this in place we can compute the polarizability from finite differences
# (just like in the previous example):
polarizability_fd = let
    ε = 0.01
    (compute_dipole(ε) - compute_dipole(0.0)) / ε
end

# We do the same thing using automatic differentiation. Under the hood this uses
# custom rules to implicitly differentiate through the self-consistent
# field fixed-point problem.
polarizability = ForwardDiff.derivative(compute_dipole, 0.0)
println()
println("Polarizability via ForwardDiff:       $polarizability")
println("Polarizability via finite difference: $polarizability_fd")


immutable Binomial <: DiscreteUnivariateDistribution
    size::Int
    prob::Float64
    function Binomial(n::Real, p::Real)
        n >= zero(n) || error("size must be positive")
        zero(p) <= p <= one(p) || error("prob must be in [0, 1]")
        new(int(n), float64(p))
    end
end

Binomial(size::Integer) = Binomial(size, 0.5)
Binomial() = Binomial(1, 0.5)

minimum(d::Binomial) = 0
maximum(d::Binomial) = d.size

@_jl_dist_2p Binomial binom

function _probs(d::Binomial, f::Int, l::Int)
    n = d.size
    p = d.prob
    b = f - 1
    r = Array(Float64, l - b)
    r[1] = v = pdf(d, f)
    if l > f
        c = p / (1.0 - p)
        for k = f+1:l
            v *= ((n - k + 1) / k * c)
            r[k-b] = v
        end
    end
    return r
end

probs(d::Binomial) = _probs(d, 0, d.size)

function probs(d::Binomial, rgn::UnitRange)
    f, l = rgn[1], rgn[end]
    0 <= f <= l <= d.size || throw(BoundsError())
    _probs(d, f, l)
end

function entropy(d::Binomial; approx::Bool=false)
    n = d.size
    p1 = d.prob

    (p1 == 0.0 || p1 == 1.0 || n == 0) && return 0.0
    p0 = 1.0 - p1

    if approx
        return 0.5 * (log(2.0pi * n * p0 * p1) + 1.0)
    else
        lg = log(p1 / p0)
        lp = n * log(p0)
        s = exp(lp) * lp
        for k = 1:n
           lp += log((n - k + 1) / k) + lg
           s += exp(lp) * lp
        end
        return -s
    end
end

kurtosis(d::Binomial) = (1.0 - 6.0 * d.prob * (1.0 - d.prob)) / var(d)

mean(d::Binomial) = d.size * d.prob

var(d::Binomial) = d.size * d.prob * (1.0 - d.prob)

skewness(d::Binomial) = (1.0 - 2.0 * d.prob) / std(d)

median(d::Binomial) = iround(d.size * d.prob)

# TODO: May need to subtract 1 sometimes
# two modes possible e.g. size odd, p = 0.5
mode(d::Binomial) = d.size > 0 ? iround((d.size + 1.0) * d.prob) : 0

function mgf(d::Binomial, t::Real)
    p = d.prob
    (1.0 - p + p * exp(t))^d.size
end

function cf(d::Binomial, t::Real)
    p = d.prob
    (1.0 - p + p * exp(im * t))^d.size
end

modes(d::Binomial) = iround([d.size * d.prob])

### handling support

isupperbounded(d::Union(Binomial, Type{Binomial})) = true
islowerbounded(d::Union(Binomial, Type{Binomial})) = true
isbounded(d::Union(Binomial, Type{Binomial})) = true

minimum(d::Union(Binomial, Type{Binomial})) = 0
maximum(d::Binomial) = d.size
support(d::Binomial) = 0:d.size

insupport(d::Binomial, x::Real) = isinteger(x) && 0 <= x <= d.size

## Fit model

immutable BinomialStats <: SufficientStats
    ns::Float64   # the total number of successes
    ne::Float64   # the number of experiments
    n::Int        # the number of trials in each experiment

    function BinomialStats(ns::Real, ne::Real, n::Integer)
        new(float64(ns), float64(ne), int(n))
    end
end

function suffstats{T<:Integer}(::Type{Binomial}, n::Integer, x::Array{T})
    ns = zero(T)
    for xi in x
        0 <= xi <= n || throw(DomainError())
        ns += xi
    end
    BinomialStats(ns, length(x), n)
end

function suffstats{T<:Integer}(::Type{Binomial}, n::Integer, x::Array{T}, w::Array{Float64})
    ns = 0.
    ne = 0.
    for i = 1:length(x)
        @inbounds xi = x[i]
        0 <= xi <= n || throw(DomainError())
        @inbounds wi = w[i]
        ns += xi * wi
        ne += wi
    end
    BinomialStats(ns, ne, n)
end

suffstats{T<:Integer}(::Type{Binomial}, data::(Int, Array{T})) = suffstats(Binomial, data...)
suffstats{T<:Integer}(::Type{Binomial}, data::(Int, Array{T}), w::Array{Float64}) = suffstats(Binomial, data..., w)

fit_mle(::Type{Binomial}, ss::BinomialStats) = Binomial(ss.n, ss.ns / (ss.ne * ss.n))

fit_mle{T<:Integer}(::Type{Binomial}, n::Integer, x::Array{T}) = fit_mle(Binomial, suffstats(Binomial, n, x))
fit_mle{T<:Integer}(::Type{Binomial}, n::Integer, x::Array{T}, w::Array{Float64}) = fit_mle(Binomial, suffstats(Binomial, n, x, w))
fit_mle{T<:Integer}(::Type{Binomial}, data::(Int, Array{T})) = fit_mle(Binomial, suffstats(Binomial, data))
fit_mle{T<:Integer}(::Type{Binomial}, data::(Int, Array{T}), w::Array{Float64}) = fit_mle(Binomial, suffstats(Binomial, data, w))


using Flux: onehot, onehotbatch, onecold

@show onehot(1,[1,2,3])
@show onehotbatch([1,2], [1,2,3])
@show onecold([0.3,0.2,0.5],[1,2,3])


function vol_given_values_and_times(vals::Vector, times::Vector, asset::Symbol)
    duration = maximum(times) - minimum(times)
    returns = simple_differencing(vals[2:end], vals[1:(end-1)])
    return sqrt(sum(returns .^ 2)/duration)
end

"""
    default_num_grids(ts::SortedDataFrame)

This gives a default number of intervals to divide a series of ticks over for the purposes of estimating returns and volatility.
### Inputs
* `ts` - The tick data.
### Returns
* An integer for the number of intervals.
"""
function default_num_grids(ts::SortedDataFrame)
    min_ticks = minimum(map( a -> length(ts.groupingrows[a]) , collect(keys(ts.groupingrows)) ))
    return Int(max(floor(min_ticks / 100), 3))
end

"""
    two_scales_volatility(vals::Vector, times::Vector, asset::Symbol, num_grids::Real)

Calculates volatility with the two scales method of Zhang, Mykland, Ait-Sahalia 2005. The amount of time for the grid spacing is by default this is a tenth of the total duration
by default. If this doesn't make sense for your use of it then choose a spacing at which you expect the effect of microstructure noise will be small.
### Inputs
* `vals` - The prices at each instant in time.
* `times` - The times corresponding to each element in `vals`.
* `asset` - The name of the asset.
* `num_grids` - Number of grids used in order in two scales estimation.
### Returns
* A scalar for the estimated volatility of the asset.
* A scalar for the estimated microstructure noise variance.

    two_scales_volatility(ts::SortedDataFrame, assets::Vector{Symbol} = get_assets(ts);
                          num_grids::Real = default_num_grids(ts))

Calculates volatility with the two scales method of Zhang, Mykland, Ait-Sahalia 2005. The amount of time for the grid spacing is by default this is a tenth of the total duration
by default. If this doesn't make sense for your use of it then choose a spacing at which you expect the effect of microstructure noise will be small.
### Inputs
* `ts` - The tick data.
* `assets` - The assets you want to estimate volatilities for.
* `num_grids` - Number of grids used in order in two scales estimation.
### Returns
* A  `Dict` with estimated volatilities for each asset.
* A  `Dict` with estimated microstructure noise variances for each asset.

### References
Zhang L, Mykland PA, Aït-Sahalia Y (2005). "A Tale of Two Time Scales: Determining Integrated Volatility with Noisy High-Frequency Data." Journal of the American Statistical Association, 100(472), 1394–1411. ISSN 01621459. doi:10.1198/016214505000000169.
"""
function two_scales_volatility(vals::Vector, times::Vector, asset::Symbol, num_grids::Real)
    dura  = maximum(times) - minimum(times)
    if (dura < eps()) | (length(vals) < 10) return NaN, NaN end
    num_grids = Int(  min( floor(length(times)/4),   max( 2,  floor(num_grids)  ) ))
    avg_vol   = mean(map(i ->  vol_given_values_and_times(vals[i:num_grids:end], times[i:num_grids:end], asset), 1:num_grids ))
    all_vol   = vol_given_values_and_times(vals, times, asset)

    pure_vol  = ((1- 1/num_grids)^(-1)) * ( avg_vol - (1/num_grids) *all_vol   )
    noise     = (1/(2*length(vals))) * ((all_vol^2)*dura - (pure_vol^2)*dura )
    return pure_vol, noise
end

function two_scales_volatility(ts::SortedDataFrame, assets::Vector{Symbol} = get_assets(ts);
                               num_grids::Real = default_num_grids(ts))
    vols = Dict{Symbol,eltype(ts.df[:,ts.value])}()
    micro_noise_var = Dict{Symbol,eltype(ts.df[:,ts.value])}()
    for a in assets
        vals = ts.df[ts.groupingrows[a],ts.value]
        times = ts.df[ts.groupingrows[a],ts.time]
        if length(vals) < 10
            @warn "There was not enough data to use the two_scales_volatility method. Returning NaN."
            vols[a] = NaN
            micro_noise_var[a] = NaN
        else
            pure_vol, noise = two_scales_volatility(vals, times, a, num_grids)
            vols[a] = pure_vol
            micro_noise_var[a] = max(0.0,noise)
        end
    end
    return vols, micro_noise_var
end


using Combinatorics
import Base: ==

"""
    ConditionalIndependence

Structure representing conditional independency between two variables given conditioned set of variables.
"""
struct ConditionalIndependence
    x::Symbol
    y::Symbol
    cond::Vector{Symbol}
end

==(a::ConditionalIndependence, b::ConditionalIndependence) = a.x == b.x && a.y == b.y && a.cond == b.cond

"""
    implied_conditional_independencies(dag)

From given DAG find all pair-wise conditional independencies of nodes.
Returns vector of `ConditionalIndependence` structures.

# Examples
```jldoctest
julia> using Dagitty

julia> g = DAG(:A => :C, :C => :B)
DAG: {3, 2} directed simple Int64 graph with labels [:A, :B, :C])

julia> implied_conditional_independencies(g)
1-element Vector{ConditionalIndependence}:
 ConditionalIndependence(:A, :B, [:C])
```
"""
function implied_conditional_independencies(dag::DAG)::Vector{ConditionalIndependence}
    max_set = length(dag.labels) - 2
    result = Vector{ConditionalIndependence}()
    nodes = Base.OneTo(length(dag.labels))
    for (n1, n2) ∈ combinations(nodes, 2)
        rest = [n for n ∈ nodes if n != n1 && n != n2]
        for cond ∈ powerset(rest)
            if is_d_separated(dag, [n1], [n2], cond)
                ci = ConditionalIndependence(dag.labels[n1], dag.labels[n2], dag.labels[cond])
                push!(result, ci)
            end
        end
    end
    result
end

"""
    implied_conditional_independencies_min

From given DAG return all pair-wise conditional independencies of nodes with
minimal condition set.

# Examples
```jldoctest
julia> using Dagitty

julia> plant_dag = Dagitty.DAG(:H_0 => :H_1, :F => :H_1, :T => :F)
DAG: {4, 3} directed simple Int64 graph with labels [:F, :H_0, :H_1, :T])

julia> implied_conditional_independencies_min(plant_dag)
3-element Vector{ConditionalIndependence}:
 ConditionalIndependence(:F, :H_0, Symbol[])
 ConditionalIndependence(:H_0, :T, Symbol[])
 ConditionalIndependence(:H_1, :T, [:F])

```
"""
function implied_conditional_independencies_min(dag::DAG)::Vector{ConditionalIndependence}
    s = implied_conditional_independencies(dag)
    [
        a for a ∈ s
        if !any(b -> b.x == a.x && b.y == a.y && b.cond != a.cond && b.cond ⊆ a.cond, s)
    ]
end

export
    ConditionalIndependence,
    implied_conditional_independencies,
    implied_conditional_independencies_min


###############################################################################
#
# Reduce step method to find peak for uni-variable
#
###############################################################################
function find_peak(
            f::F,
            ms::ReduceStepMethod{FT},
            tol::SolutionTolerance{FT}
) where {F<:Function, FT<:AbstractFloat}
    # count iterations
    count_all = 0

    # define the initial step
    @unpack x_ini, x_max, x_min, Δ_ini = ms;

    # initialize the y
    tar_x = x_ini;
    tar_y = f(tar_x);
    new_x = x_min;
    new_y = f(new_x);
    new_y >= tar_y ? (tar_x=new_x; tar_y=new_y;) : nothing;
    new_x = x_max;
    new_y = f(new_x);
    new_y >  tar_y ? (tar_x=new_x; tar_y=new_y;) : nothing;

    # find the solution
    Δx::FT = Δ_ini;
    while true
        # 1. increase the x by Δx till tar_y is bigger
        count_inc = 0
        while true
            new_x = tar_x + Δx;
            new_x > x_max ? break : nothing;
            new_y = f(new_x);
            new_y > tar_y ? (tar_x=new_x; tar_y=new_y;) : break;
            count_inc += 1;
            count_all += 1;
        end

        # 2. decrease the x by Δx till tar_y is bigger
        count_dec = 0
        while count_inc == 0
            new_x = tar_x - Δx;
            new_x < x_min ? break : nothing;
            new_y = f(new_x);
            new_y >= tar_y ? (tar_x=new_x; tar_y=new_y;) : break;
            count_dec += 1;
            count_all += 1;
        end

        # 3. if no update, then 10% the Δx
        if count_inc + count_dec == 0
            Δx /= 10;
        end

        # 4. if break
        Δx <= tol.tol ? break : nothing;
    end

    return tar_x
end


#===============================================================================
  Conjugate Gradient Descent

  We implement a generic conjugate gradient descent method that includes many
  different β-choices, preconditioning, precise line searches, and more.

  A conjugate gradient method here will have to implement a CGUpdate type (CD,
  HZ, HS, ...) that controls the β update via the update_parameter function.

  The update formulae are mostly taken from HZ2006

  Let N be a neighborhood of {x ∈ R^n : f(x) ⩽ f(x₀)}

  [LA; Lipschitz Assumption]: ∃ L < ∞ : ||∇f(x)-∇f(y)|| ⩽ L||x-y|| for x,y ∈ N

  Todos:
  Might consider Daniel 1967 that uses second order
===============================================================================#

abstract type CGUpdate end
struct ConjugateGradient{Tu, TP}
  update::Tu
  P::TP
end
ConjugateGradient(;update=HZ()) = ConjugateGradient(update, nothing)
hasprecon(::ConjugateGradient{<:Any, <:Nothing}) = NoPrecon()
hasprecon(::ConjugateGradient{<:Any, <:Any}) = HasPrecon()

struct CGVars{T1, T2, T3}
    y::T1 # change in successive gradients
    d::T2 # search direction
    α::T3
    β::T3
    ls_success::Bool
end

function prepare_variables(problem, approach::LineSearch{<:ConjugateGradient, <:Any, <:Any}, x0, ∇fz)
    z = x0
    x = copy(z)
    fz, ∇fz = upto_gradient(problem, ∇fz, x)

    fx = copy(fz)
    ∇fx = copy(∇fz)

    Pg = approach.scheme.P isa Nothing ? ∇fz : copy(∇fz)
    return (x=x, fx=fx, ∇fx=∇fx, z=z, fz=fz, ∇fz=∇fz, B=nothing, Pg=Pg)
end

summary(cg::ConjugateGradient) = "Conjugate Gradient Descent ($(summary(cg.update)))"
#===============================================================================
  Conjugate Descent [Fletcher] (CD)

  R. Fletcher. "Practical Methods of Optimization vol. 1: Unconstrained
  Optimization." John Wiley & Sons, New York (1987).
===============================================================================#
struct CD <: CGUpdate end
function update_parameter(mstyle, cg::CD, d, ∇fz, ∇fx, y, P, P∇fz)
  -dot(∇fz, P∇fz)/dot(d, ∇fx)
end

#===============================================================================
  Hager-Zhang (HZ)
===============================================================================#
struct HZ{Tη} <: CGUpdate
    η::Tη # a "forcing term"
end
HZ() = HZ(0.4)
function update_parameter(mstyle, cg::HZ, d, ∇fz, ∇fx, y, P, P∇fz)
    T = eltype(∇fz)
    θ = T(2)
    η = T(cg.η)
    # βHS = update_parameter(mstyle, HS(), d,  ∇fz, ∇fx, y, P, P∇fz)
    # but let's save the dy calculation from begin repeated for
    # efficiency's sake
    dy = dot(d, y)
    βHS = dot(y, P∇fz)/dy

    # Apply preconditioner to y
    Py = apply_preconditioner(mstyle, P, copy(P∇fz), y)
    βN = βHS - θ*dot(y, Py)/dy*dot(d, ∇fz)/dy

    # 2013 version is scale invariant
    Pinv_y = apply_inverse_preconditioner(mstyle, P, copy(P∇fz), y)
    ηk = η*dot(d, ∇fx)/dot(d, Pinv_y)
    # 2006 version
    # ηk = -inv(norm(d)*min(T(cg.η), norm(∇fx)))

    βkp1 = max(βN, ηk)
end

#===============================================================================
  Hestenes-Stiefel (HS)

  Jamming protection from y in the numerator. x-x' small means ∇f(x)-∇f(x')
  should also be small.

  Hestenes, Magnus R., and Eduard Stiefel. "Methods of conjugate gradients for
  solving linear systems." Journal of research of the National Bureau of
  Standards 49, no. 6 (1952): 409-436.
===============================================================================#
struct HS <: CGUpdate end
function update_parameter(mstyle, cg::HS, d, ∇fz, ∇fx, y, P, P∇fz)
  dot(y, P∇fz)/dot(d, y)
end
#===============================================================================
  Fletcher-Reeves (FR)

  Fletcher, Reeves, and Colin M. Reeves. "Function minimization by conjugate
  gradients." The computer journal 7, no. 2 (1964): 149-154.
===============================================================================#
struct FR <: CGUpdate end
function update_parameter(mstyle, cg::FR, d, ∇fz, ∇fx, y, P, P∇fz)
    dot(∇fz, P∇fz)/dot(∇fx, ∇fx)
end
#===============================================================================
  Polak-Ribiére-Polyak (PRP)

  Jamming protection from y in the numerator. x-x' small means ∇f(x)-∇f(x')
  should also be small.

  Polak, Elijah, and Gerard Ribiere. "Note sur la convergence de méthodes de
  directions conjuguées." ESAIM: Mathematical Modelling and Numerical Analysis-
  Modélisation Mathématique et Analyse Numérique 3, no. R1 (1969): 35-43.

  Polyak, Boris T. "The conjugate gradient method in extremal problems." USSR
  Computational Mathematics and Mathematical Physics 9, no. 4 (1969): 94-112.
===============================================================================#
struct PRP{Plus} end
PRP(;plus=true) = PRP{plus}()
function update_parameter(mstyle, ::PRP, d, ∇fz, ∇fx, y, P, P∇fz)
    dot(y, P∇fz)/dot(∇fx, ∇fx)
end
function update_parameter(::PRP{true}, d, ∇fz, ∇fx, y, P, P∇fz)
  βPR = update_parameter(RP{false}(), d, ∇fz, ∇fx, y, P, P∇fz)
  max(0, β)
end

#===============================================================================
  Liu-Storey (LS)

  Liu-Storey is identical to PRP for exact line search.

  Jamming protection from y in the numerator. x-x' small means ∇f(x)-∇f(x')
  should also be small.

  Liu, Y., and C. Storey. "Efficient generalized conjugate gradient algorithms,
  part 1: theory." Journal of optimization theory and applications 69, no. 1
  (1991): 129-137.
===============================================================================#
struct LS <: CGUpdate end
function update_parameter(mstyle, ::LS, d, ∇fz, ∇fx, y, P, P∇fz)
  -dot(y, P∇fz)/dot(d, ∇fx)
end
#===============================================================================
  Dai-Yuan (DY)

  Different from CD and FR. With Wolfe line search it always generates descent
  directions. Global convergence with Lipchitz Assumption [LA].

  Dai, Yu-Hong, and Yaxiang Yuan. "A nonlinear conjugate gradient method with a
  strong global convergence property." SIAM Journal on optimization 10, no. 1
  (1999): 177-182.
===============================================================================#
struct DY <: CGUpdate end
function update_parameter(mstyle, cg::DY, d, ∇fz, ∇fx, y, P, P∇fz)
  dot(∇fz, P∇fz)/dot(d, y)
end

#===============================================================================
  Wei-Yao-Liu (VPRP)

  Yu, Gaohang, Yanlin Zhao, and Zengxin Wei. "A descent nonlinear conjugate
  gradient method for large-scale unconstrained optimization." Applied
  mathematics and computation 187, no. 2 (2007): 636-643.
===============================================================================#
struct VPRP <: CGUpdate end
function update_parameter(mstyle, cg::VPRP, d, ∇fz, ∇fx, y, P, P∇fz)
  a = dot(∇fz, P∇fz)
  b = dot(∇fx, P∇fz)
  c = norm(∇fx, 2)
  num = a - sqrt(a)/c*b
  num/dot(d, y)
end

# using an "initial vectors" function we can initialize s if necessary or nothing if not to save on vectors
function solve(problem::OptimizationProblem, x0, cg::ConjugateGradient, options::OptimizationOptions)
  _solve(problem, x0, LineSearch(cg, HZAW()), options, mstyle(problem))
end
function solve(problem::OptimizationProblem, x0, approach::LineSearch{<:ConjugateGradient, <:LineSearcher}, options::OptimizationOptions)
  _solve(problem, x0, approach, options, mstyle(problem))
end
function _solve(problem::OptimizationProblem, x0, approach::LineSearch{<:ConjugateGradient, <:LineSearcher}, options::OptimizationOptions, mstyle::MutateStyle)
    t0 = time()
    #==============
         Setup
    ==============#
    Tx = eltype(x0)

    objvars = prepare_variables(problem, approach, x0, copy(x0))
    f0, ∇f0 = objvars.fx, norm(objvars.∇fx, Inf) # use user norm

    y, d, α, β = copy(objvars.∇fz), -copy(objvars.∇fx), Tx(0), Tx(0)
    cgvars = CGVars(y, d, α, β, true)

    k = 1
    objvars, P, cgvars = iterate(mstyle, cgvars, objvars, approach, problem, options)
    is_converged = converged(approach, objvars, ∇f0, options)
    while k < options.maxiter && !any(is_converged)
        k += 1
        objvars, P, cgvars = iterate(mstyle, cgvars, objvars, approach, problem, options, P, false)
        is_converged = converged(approach, objvars, ∇f0, options)
    end
    x, fx, ∇fx, z, fz, ∇fz, B = objvars
    return ConvergenceInfo(approach, (beta=β, ρs=norm(x.-z), ρx=norm(x), minimizer=z, fx=fx, minimum=fz, ∇fx=∇fx, ∇fz=∇fz, f0=f0, ∇f0=∇f0, iter=k, time=time()-t0), options)
end
function iterate(mstyle::InPlace, cgvars::CGVars, objvars, approach::LineSearch{<:ConjugateGradient, <:Any, <:Any}, problem::OptimizationProblem, options::OptimizationOptions, P=nothing, is_first=nothing)
    # split up the approach into the hessian approximation scheme and line search
    x, fx, ∇fx, z, fz, ∇fz, B, Pg = objvars
    Tx = eltype(x)

    scheme, linesearch = modelscheme(approach), algorithm(approach)
    y, d, α, β = cgvars.y, cgvars.d, cgvars.α, cgvars.β

    # Move nexts into currs
    fx = fz
    copyto!(x, z)
    copyto!(∇fx, ∇fz)

    # Precondition current gradient and calculate the search direction
    P = update_preconditioner(scheme, x, P)
    P∇fz = apply_preconditioner(mstyle, P, Pg, ∇fz)
    @. d = -P∇fz + β*d

    α_0 = initial(scheme.update, a->value(problem, x.+a.*d), α, x, fx, dot(d, ∇fx), ∇fx, is_first)
    φ = _lineobjective(mstyle, problem, ∇fz, z, x, d, fx, dot(∇fx, d))

    # Perform line search along d
    α, f_α, ls_success = find_steplength(mstyle, linesearch, φ, Tx(1))

    # Calculate final step vector and update the state
    if ls_success
        z = retract(problem, z, x, d, α)
        fz, ∇fz = upto_gradient(problem, ∇fz, z)
        @. y = ∇fz - ∇fx
    else
        # if no succesful search direction is found, reset to gradient
        y .= .-∇fz
    end
    β = update_parameter(mstyle, scheme.update, d, ∇fz, ∇fx, y, P, P∇fz)

    return (x=x, fx=fx, ∇fx=∇fx, z=z, fz=fz, ∇fz=∇fz, B=nothing, Pg=Pg), P, CGVars(y, d, α, β, ls_success)
end

function iterate(mstyle::OutOfPlace, cgvars::CGVars, objvars, approach::LineSearch{<:ConjugateGradient, <:Any, <:Any}, problem::OptimizationProblem, options::OptimizationOptions, P=nothing, is_first=nothing)
    # split up the approach into the hessian approximation scheme and line search
    x, fx, ∇fx, z, fz, ∇fz, B, Pg = objvars
    Tx = eltype(x)
    scheme, linesearch = modelscheme(approach), algorithm(approach)
    d, α, β = cgvars.d, cgvars.α, cgvars.β

    # Move nexts into currs
    fx = fz
    x = copy(z)
    ∇fx = copy(∇fz)

    # Update preconditioner
    P = update_preconditioner(scheme, x, P)
    P∇fz = apply_preconditioner(mstyle, P, Pg, ∇fz)
    # Update current gradient and calculate the search direction
    d = @. -P∇fz + β*d

    α_0 = initial(scheme.update, a-> value(problem, x .+ a.*d), α, x, fx, dot(d, ∇fx), ∇fx, is_first)
    φ = _lineobjective(mstyle, problem, ∇fz, z, x, d, fx, dot(∇fx, d))

    # Perform line search along d
    α, f_α, ls_success = find_steplength(mstyle, linesearch, φ, Tx(1))

    z = retract(problem, z, x, d, α)
    fz, ∇fz = upto_gradient(problem, ∇fz, z)
    y = @. ∇fz - ∇fx
    β = update_parameter(mstyle, scheme.update, d, ∇fz, ∇fx, y, P, P∇fz)
    return (x=x, fx=fx, ∇fx=∇fx, z=z, fz=fz, ∇fz=∇fz, B=nothing, Pg=Pg), P, CGVars(y, d, α, β, ls_success)
end

function initial(cg, φ, α, x, φ₀, dφ₀, ∇fx, is_first)
    T = eltype(x)
    ψ₀ = T(0.01)
    ψ₁ = T(0.1)
    ψ₂ = T(2.0)
    quadstep = true
    if is_first isa Nothing
        if !all(x .≈ T(0)) # should we define "how approx we want?"
            return ψ₀ * norm(x, Inf)/norm(∇fx, Inf)
        elseif !(φ₀ ≈ T(0))
            return ψ₀ * abs(φ₀)/norm(∇fx, 2)^2
        else
            return T(1)
        end
    elseif quadstep
        R = ψ₁*α
        φR = φ(R)
        if φR ≤ φ₀
            c = (φR - φ₀ - dφ₀*R)/R^2
            if c > 0
               return -dφ₀/(T(2)*c) # > 0 by df0 < 0 and c > 0
            end
        end
    end
    return ψ₂*α
end


#   This file is part of Fatou.jl. It is licensed under the MIT license
#   Copyright (C) 2019 Michael Reed

PyPlot.plot(K::FilledSet;c::String="",bare::Bool=false) = PyPlot.imshow(K;cmap=c,bare=bare)

function PyPlot.imshow(K::FilledSet;cmap::String="",bare::Bool=false)
    PyPlot.figure() # plot figure using imshow based in input preferences
    isempty(cmap) && (cmap = K.meta.cmap)
    isempty(cmap) ? PyPlot.imshow(K.meta.iter ? K.iter : K.mix, extent=bounds(K)) :
        PyPlot.imshow(K.meta.iter ? K.iter : K.mix, cmap=cmap, extent=bounds(K))
    PyPlot.tight_layout()
    !bare && PyPlot.title(K)
end

function PyPlot.title(K::FilledSet) # annotate title using LaTeX
    text,t = "f:z\\mapsto $(rdpm(Algebra.latex(K.meta.E))),\\,",LaTeXString(typeplot(K))
    if K.meta.newt # annotate y-axis with Newton's method
        PyPlot.title(latexstring("$text m = $(K.meta.m), ")*t)
        PyPlot.ylabel(L"Fatou\,set:\,"*L"z\,↦\,z-m\,×\,f(z)\,/\,f\,'(z)")
    else
        PyPlot.title(latexstring(text)*t)
    end
    PyPlot.tight_layout()
    PyPlot.colorbar()
end

function orbit(E,f::Function,bi::Matrix{Float64},orb::Int=0,depth::Int=1,incr::Int=384)
    x,N,N2,orbit,bis = real_orb(E,f,bi,orb,depth,incr)
    # prepare for next figure
    PyPlot.figure()
    # plot background lines
    PyPlot.plot(x[:],N[:,1],"k--",x[:],N[:,2])
    # plot orbit cobweb path
    PyPlot.plot(orbit[:,1],orbit[:,2],"r")
    # plot f^2,f^3,f^4,...
    for h ∈ 3:depth+1
        PyPlot.plot(x,N[:,h],lw=1)
    end
    if ~(orb == 0)
        PyPlot.plot(range(bi[1],stop=bi[2],length=length(N2)),N2[:],"gray",marker="x",linestyle=":",lw=1)
        funs = [latexstring("\\phi(x_{0:$orb})")]
        funt = ", IC: \$ x_0 = $(bis[3])\$, \$ n\\in0:$orb\$"
    else
        funs = []
        funt = ""
    end
    # trim graph
    d=1.07
    PyPlot.xlim(bi[1],bi[2])
    PyPlot.ylim(minimum([d*minimum(N[:,2]),0]),maximum([d*maximum(N[:,2]),0]))
    # set title
    fune = rdpm(Algebra.latex(E))
    PyPlot.title(latexstring("\$ x \\mapsto $fune\$$funt"))
    # set legend
    PyPlot.legend(vcat([L"$y=x$",L"$\phi(x)$",L"(x_n,\phi(x_n))"],[latexstring("\\phi^{$x}(x)") for x ∈ 2:depth],funs))
    PyPlot.tight_layout()
    return
end


#
# Jerzy Wroczyński (nr. indeksu 250075)
# Obliczenia Naukowe, Laboratorium
# Lista 5.
#

module blocksys

include("types.jl")
import .types.FL, .types.SquareSparseMatrix

using SparseArrays

# constants
macheps = eps(FL)

"""
### Asserts provided arguments match in size.

### Arguments:
- `A` — the `A` matrix
- `n` — the size of the matrix
- `b` — the `b` vector

### Returns: `Nothing`

"""
function assertMatrixSize(A, n::Int, b::Vector=[0 for _ in 1:n])::Nothing

    # if size(A, 1) == size(A, 2) == n == length(b)
    if n == length(b)
        return
    end

    throw(ErrorException("provided arguments do not match in sizes"))

end

"""
### Asserts `n` is divisible by `l`.

### Arguments:
- `A` — the `A` matrix
- `n` — the size of the matrix
- `b` — the `b` vector

### Returns: `Nothing`

"""
function assertBlockSize(n::Int, l::Int)::Nothing

    if n%l != 0
        throw(ErrorException("n is not divisible by l"))
    end

end

"""
### Calculates the permutation vector of given matrix.
This function implements the *partial pivoting* algorithm that assures that the diagonal of the matrix contains numbers that are not near zero or equal to zero.

### Arguments:
- `A` — the `A` matrix
- `n` — size of the matrix
- `l` — size of each chunk (submatrix)
- `limited` — this flag limits the bounds of the searching loop; it may be used when the function gives a permutation that leads to having a value near zero on the diagonal of the matrix; default value is set to `false`

### Returns
— the pivot vector that simulates all the matrix rows swaps that would have occurred

"""
function partialPivot(A, n::Int, l::Int, limited::Bool = false)::Vector

    assertMatrixSize(A, n)
    assertBlockSize(n, l)

    # save the new order of the rows
    pivot = [x for x in 1 : n]

    for k in 1:Int(n/l)

        for p in (k-1)*l+1 : k*l
            # get current value
            curr = abs(A[pivot[p], p])

            # check if it’s the biggest value
            max_value = curr
            max_index = pivot[p]
            search_end = limited ? k*l : k*l + l
            for i in (p+1) : search_end

                # special case for k=n/l
                # because this is the last submatrix there are no submatrices underneath;
                # thus,the internal loop used for searching max value needs to step earlier
                if i > n
                    break
                end
                # get another value from that column
                potential = A[pivot[i], p]
                if abs(potential) > max_value
                    # its absolute value is bigger — mark as ready for swap
                    max_value = abs(potential)
                    max_index = pivot[i]
                end
            end
            if max_value < macheps
                throw(ErrorException("a value near or equal zero was selected as the main value (diagonal value) — try to run `partialPivot` with `limited` set to `true` or do not use pivoting for this particular case at all"))
            end
            # swap the rows if needed
            if max_value > curr
                pivot[p], pivot[max_index] = pivot[max_index], pivot[p]
            end

        end

    end

    return pivot

end

"""
### Gives LU decomposition of given matrix.

### Arguments:
- `A` — the matrix to decompose
- `b` — the `b` vector (the vector that is on the right side of the `Ax = b` equation)
- `n` — the size of the matrix
- `l` — size of each chunk (submatrix)
- `pivot` — optional pivoting vector that simulates matrix row swap (can be generated using the `partialPivot` function); default value is set to a linear vector [1,…,n]

### Returns a tuple:
- `[1]` — matrices `L` and `U` contained in one matrix for optimization
- `[2]` — altered `b' = b^{(k)}` vector

"""
function decomposeIntoLU(A, b::Vector, n::Int, l::Int, pivot::Vector = [x for x in 1:n])#::Tuple{types.SquareSparseMatrix, Vector}

    assertMatrixSize(A, n)
    assertBlockSize(n, l)

    # this a more optimized way of storing both L and U matrices
    # we copy `A` at the beginning because that’s the start of the sequence:
    # A^{(1)} → A^{(2)} → … → A^{(n)} = U
    # in the end both matrices are triangular, so there is no problem with storing them this way
    # and the `L` matrix does not store any information on the diagonal of the matrix so there is no collision between `U` and `L`
    LU = copy(A)

    bprime = copy(b)

    # these two adjacent `for` loops are just one loop split into `l`-sized chunks
    # — its time complexity is still O(n)
    for k in 1:Int(n/l)
        for p in (k-1)*l+1 : k*l

            # update all rows beneath
            for i in (p+1) : k*l + l
                # special case k = n/l
                if i > n
                    break
                end

                lᵢₚ = LU[pivot[i], p] / LU[pivot[p], p]

                # process the row with new coefficients
                for j in p : k*l + l
                    if j > n
                        break
                    end
                    LU[pivot[i], j] = LU[pivot[i], j] - lᵢₚ * LU[pivot[p], j]
                end

                # adjust the `b` matrix accordingly
                bprime[pivot[i]] = bprime[pivot[i]] - lᵢₚ * bprime[pivot[p]]

                # save the `L` matrix part
                LU[pivot[i], p] = lᵢₚ
            end
        end
    end

    return LU, bprime

end

"""
### Solves the standard equation `Ux = y` where `L` is an upper-triangular matrix.

### Arguments:
- `U` — the upper-triangular matrix
- `y` — the right side of the equation
- `n` — the size of the matrix
- `l` — size of each chunk (submatrix)
- `pivot` — optional pivoting vector that simulates matrix row swap (can be generated using the `partialPivot` function); default value is set to a linear vector [1,…,n]

### Returns
— the `x` output vector (solution to the `Ux = y` equation)

"""
function solveUpperTriangularMatrix(U, y::Vector, n::Int, l::Int, pivot::Vector = [x for x in 1:n])::Array

    # init the output vector
    x = [0.0 for _ in 1:n]

    # calc the values of the output vector
    for k in Int(n/l) : -1 : 1
        for p in k*l : -1 : (k-1)*l+1

            tmp = y[pivot[p]]
            for i in (p+1) : (k*l + l)
                if i > n
                    break
                end
                tmp -= x[pivot[i]] * U[pivot[p], i]
            end
            x[pivot[p]] = tmp/U[pivot[p], p]

        end
    end

    return x

end

"""
### Solves the `LUx = b` equation.

### Arguments:
- `LU` — LU decomposition of a given matrix
- `b` — the `b` vector (the right side of the `LUx = b` equation)
- `n` — the size of the matrix
- `l` — size of each chunk (submatrix)
- `pivot` — optional pivoting vector that simulates matrix row swap (can be generated using the `partialPivot` function); default value is set to a linear vector [1,…,n]

### Returns
— the `x` output vector (solution to the `LUx = b` equation)

"""
function solveFromLU(LU, b::Vector, n::Int, l::Int, pivot::Vector = [x for x in 1:n])::Array

    # Ly = b

    y = zeros(n)

    # calc the values of the intermediary vector
    for k in 1:Int(n/l)
        for p in (k-1)*l+1 : k*l

            tmp = b[pivot[p]]
            for i in (p-1) : -1 : (k-1)*l - l
                if i <= 0
                    break
                end
                tmp -= y[pivot[i]] * LU[pivot[p], i]
            end

            y[pivot[p]] = tmp#/1

        end
    end

    # Ux = y
    return solveUpperTriangularMatrix(LU, y, n, l, pivot)

end

"""
### Solves the `Ax = b` equation.

### Arguments:
- `A` — the input matrix
- `b` — the `b` vector (the right side of the `LUx = b` equation)
- `n` — the size of the matrix
- `l` — size of each chunk (submatrix)
- `pivot` — optional pivoting vector that simulates matrix row swap (can be generated using the `partialPivot` function); default value is set to a linear vector [1,…,n]

### Returns
— the `x` output vector (solution to the `Ax = b` equation)

"""
function gaussianElimination(A, b::Vector, n::Int, l::Int, pivot::Vector = [x for x in 1:n])::Array

    # get the `U` matrix and `b'`
    # here we don’t actually need the `L` matrix, but it is generated either way — but that doesn’t affect performance
    LU, bprime = decomposeIntoLU(A, b, n, l, pivot)

    return solveUpperTriangularMatrix(LU, bprime, n, l, pivot)

end

"""
### Generates the `b` vector with respect to given `A` matrix.
*(Assuming `x` is a vector of ones.)*

### Arguments:
- `A` — the input matrix
- `n` — the size of the matrix
- `l` — size of each chunk (submatrix)

### Returns
— the `b` vector

"""
function generateB(A, n::Int, l::Int)::Vector

    b = zeros(n)

    for k in 1:Int(n/l)
        for p in (k-1)*l+1 : k*l

            for j in k : k*l + l
                if j > n
                    break
                end
                b[p] += A[p, j]
            end

        end
    end

    return b

end

end