using Random, SnpArrays
using LinearAlgebra,  BenchmarkTools
Random.seed!(1234)
import TraitSimulation: snparray_simulation
n = 10
p = 2
d = 2
m = 2

function generateSPDmatrix(n)
	A = rand(n)
	m = 0.5 * (A * A')
	PDmat = m + (n * Diagonal(ones(n)))
end

function generateRandomVCM(n::Int64, p::Int64, d::Int64, m::Int64)
	# n-by-p design matrix
	X = randn(n, p)

	# p-by-d mean component regression coefficient for each trait
	B = hcat(ones(p, 1), rand(p))

	V = ntuple(x -> zeros(n, n), m)
	for i = 1:m-1
	  copy!(V[i], generateSPDmatrix(n))
	end
	copy!(V[end], Diagonal(ones(n))) # last covarianec matrix is identity

	# a tuple of m d-by-d variance component parameters
	Σ = ntuple(x -> zeros(d, d), m)
	for i in 1:m
	  copy!(Σ[i], generateSPDmatrix(d))
	end
	return(X, B, Σ, V)
end

X, B, Σ, V = generateRandomVCM(n, p, d, m)
# specify the variance compoents of the model
varcomp = @vc Σ[1] ⊗ V[1] + Σ[2] ⊗ V[2]
test_vcm1 = VCMTrait(X, B, varcomp)

model_info_vcm = @capture_out show(test_vcm1)
vc_info = @capture_out show(varcomp)

@test model_info_vcm != nothing
@test vc_info != nothing

@test nsamplesize(test_vcm1) == n
@test nvc(test_vcm1) == m
@test ntraits(test_vcm1) == d
@test neffects(test_vcm1) == p

y_vcm = simulate(test_vcm1)

@test size(y_vcm) == size(X*B)
number_independent_simulations = 5
@test length(simulate(test_vcm1, number_independent_simulations)) == number_independent_simulations

@test test_vcm1 isa TraitSimulation.AbstractTraitModel

#testing for types
minor_alle_frequency  = 0.2
nsnps = p
@test snparray_simulation([minor_alle_frequency], nsnps) isa SnpArrays.SnpArray

variance_formula2  = @vc [minor_alle_frequency][:,:] ⊗ V[1] + [minor_alle_frequency][:,:] ⊗ V[1]
trait2 = VCMTrait(["x2 + 3x1", "2 +2x1"], DataFrame(X, :auto), variance_formula2)
sigma, v = vcobjtuple(variance_formula2)
trait2_equivalent = VCMTrait(["x2 + 3x1", "2 +2x1"], DataFrame(X, :auto), [sigma...], [v...])

@test trait2_equivalent.vc[1].V == trait2.vc[1].V

# simulate the SnpArray
G = snparray_simulation([0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5], n)
γ = rand(7, 2)

vcmOBJ =  VCMTrait(X, B, G, γ, varcomp)
vcmOBJ_equivalent =  VCMTrait(X, B, G, γ, [Σ...], [V...])

@test vcmOBJ.vc[1].V == V[1]
@test vcmOBJ_equivalent.vc[1].V == V[1]
# check data was copied correctly
@test vcmOBJ.X == X

vcmOBJ2 = VCMTrait(X, B, varcomp)

@test eltype(simulate(vcmOBJ2)) == Float64

@test length(simulate(vcmOBJ2, number_independent_simulations))  == number_independent_simulations

vcmOBJ2_equivalent = VCMTrait(X, B, [Σ...], [V...])
@test vcmOBJ2.vc[1].V == vcmOBJ2_equivalent.vc[1].V
@test vcmOBJ2.G != nothing
@test vcmOBJ2.μ == X*B

# check if the preallocated memory that will loop over for the matrix normal  (to sum the variance components) before we sum the mean matrix
@test size(vcmOBJ2.Z) == size(vcmOBJ2.μ)

y_alternative = zeros(size(vcmOBJ2.μ))
function testing_allocations(vcmOBJ2, y_alternative)
    TraitSimulation.simulate!(y_alternative, vcmOBJ2)
end
bmrk  = @benchmark testing_allocations(vcmOBJ2, y_alternative)

@test bmrk.memory == 0
@test bmrk.allocs ==  0


# Check conversion and output
for i = 1:10000
    a = {rand(Int8), rand(Int16), rand(Int32), rand(Int64), rand(Int128),
         rand(Uint8), rand(Uint16), rand(Uint32), rand(Uint64), rand(Uint128),
         rand(Float32), rand(Float64)}
    b = map(string, a)
    for j in a
        if j != 0
            @test "[$(BigFloat(j)), $(BigFloat(j))]" == string(Interval(j))
        else
            @test "[$(BigFloat(j)), -$(BigFloat(j))]" == string(Interval(j))
        end
    end
end

# Check precision
a = Interval(1)
@test precision(a) == get_bigfloat_precision()

# Check squaring
a = Interval(-1, 1)
@test (left(square(a)), right(square(a))) == (0, 1)
@test (left(a * a), right(a * a)) == (-1, 1)

# Check -
a = Interval(0, 1)
@test (left(-a), right(-a)) == (-1, 0)
@test signbit(right(-a)) == 1

# Check rand
a = Interval(-1, 1)
@test -1 <= rand(a) <= 1
r = rand(a, 1, 3)
@test size(r) == (1, 3)
@test typeof(r) == Matrix{BigFloat}


using DrWatson
@quickactivate :ExtendedConv

using PyPlot
using Seaborn
using Random
using Flux
using JOLI
using Flux, Flux.Data.MNIST

set_style("whitegrid")
rc("font", family="serif", size=10)
sfmt=matplotlib.ticker.ScalarFormatter(useMathText=true)
sfmt.set_powerlimits((0, 0))
plt.ioff()

sim_name = "schematics"
save_path = plotsdir(sim_name)

nx = 8
ny = 8
nc = 1
batchsize = 1

# Conv Op
w = collect(range(-0.9999f0, 1f0, length=9))
W = reshape(w, 3, 3, 1, 1)
bias = zeros(Float32, 1)

G = Conv(W, bias, identity; stride=1, pad=1)
Gadj = ConvTranspose(W, bias, identity; stride=1, pad=1)

N = nx * ny * nc * batchsize
K = joLinearFunctionFwd_T(
    N, N,
    x -> vec(G(reshape(x, nx, ny, nc, batchsize))),
    y -> vec(Gadj(reshape(y, nx, ny, nc, batchsize))),
    Float32,Float32, name="Conv"
)

# Form explicit convolution matrix
K=K*jo_eye(Float32,nx*ny)

K_vis = deepcopy(K)
K_vis[K_vis.==0f0] .= NaN

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    W[end:-1:1, end:-1:1, 1, 1], cmap="tab20",
    vmin=-1, vmax=1)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "3x3-kernel.png"), fig)
close(fig)

# Toeplitz matrix
fig, ax = subplots(1, 1, figsize=(4, 4))
ax.imshow(
    K_vis, cmap="tab20",
    vmin=-1, vmax=1)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "toeplitz.png"), fig)
close(fig)

# Matricize multiple images
Random.seed!(13)
x1 = 1f0*ones(Float32, nx, ny) + 7f-1*randn(Float32, nx, ny)
x2 = 2f0*ones(Float32, nx, ny) + 7f-1*randn(Float32, nx, ny)
x3 = 3f0*ones(Float32, nx, ny) + 7f-1*randn(Float32, nx, ny)

X = hcat(vec(x1), vec(x2), vec(x3))

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    x1, cmap="Set3",
    vmin=0, vmax=4)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "image-1.png"), fig)
close(fig)

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    x2, cmap="Set3",
    vmin=0, vmax=4)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "image-2.png"), fig)
close(fig)

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    x3, cmap="Set3",
    vmin=0, vmax=4)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "image-3.png"), fig)
close(fig)

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    X, cmap="Set3", aspect=1,
    vmin=0, vmax=4)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "X.png"), fig)
close(fig)

# MNIST images
train_imgs = MNIST.images()
X = Array{Float32}(undef, size(train_imgs[1])..., 1, 3)
Y = Array{Float32}(undef, size(train_imgs[1])..., 1, 3)

x1 = Float32.(train_imgs[1])
x2 = Float32.(train_imgs[2])
x3 = Float32.(train_imgs[3])

X = hcat(vec(x1), vec(x2), vec(x3))

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    x1, interpolation="nearest", cmap="gray",
    vmin=0, vmax=1)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "mnist-1.png"), fig)
close(fig)

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    x2, interpolation="nearest", cmap="gray",
    vmin=0, vmax=1)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "mnist-2.png"), fig)
close(fig)

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    x3, interpolation="nearest", cmap="gray",
    vmin=0, vmax=1)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "mnist-3.png"), fig)
close(fig)

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    X, aspect=.1, interpolation="nearest", cmap="gray",
    vmin=0, vmax=1)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "X-mnist.png"), fig)
close(fig)

# MNIST Conv
nx, ny = size(x1)
nc = 1
batchsize = 1

# Conv Op
Random.seed!(1)
w = randn(Float32, 3, 3)/sqrt(9f0)
W = reshape(w, 3, 3, 1, 1)
bias = zeros(Float32, 1)

G = Conv(W, bias, identity; stride=1, pad=1)
Gadj = ConvTranspose(W, bias, identity; stride=1, pad=1)

N = nx * ny * nc * batchsize
K = joLinearFunctionFwd_T(
    N, N,
    x -> vec(G(reshape(x, nx, ny, nc, batchsize))),
    y -> vec(Gadj(reshape(y, nx, ny, nc, batchsize))),
    Float32,Float32, name="Conv"
)

# Form explicit convolution matrix
K=K*jo_eye(Float32,nx*ny)
Y = K*X

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    reshape(Y[:, 1], nx, ny), interpolation="nearest", cmap="gray",
    vmin=0, vmax=.2)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "y-1.png"), fig)
close(fig)

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    reshape(Y[:, 2], nx, ny), interpolation="nearest", cmap="gray",
    vmin=0, vmax=.2)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "y-2.png"), fig)
close(fig)

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    reshape(Y[:, 3], nx, ny), interpolation="nearest", cmap="gray",
    vmin=0, vmax=.2)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "y-3.png"), fig)
close(fig)

fig, ax = subplots(1, 1, figsize=(4, 4))
im = ax.imshow(
    Y, aspect=.1, interpolation="nearest", cmap="gray",
    vmin=0, vmax=.2)
ax.axes.xaxis.set_visible(false)
ax.axes.yaxis.set_visible(false)
wsave(joinpath(save_path, "Y.png"), fig)
close(fig)


########################################################################
# Level set method
########################################################################
"""
Return an exit flag for the level-set oracle.
"""
function checkExitLvl(gap::Float64, ℓ::Float64, u::Float64, k::Int64, α::Float64, feaTol::Float64, maxIts::Int64, rule::String)
    if rule == "newton"
        # u < α + feaTol && return :feasible              # |Mx-b|^2/2 < α is nearly satisfied
        ℓ > α + feaTol && return :suboptimal            # ℓ is large enough to update τ
    elseif rule == "bisection"
        u < α && return :suboptimal_large               # τ too large
        ℓ > α && return :suboptimal_small               # τ too small
    end
    gap ≤ feaTol * eps(Float64) && return :optimal      # current lasso problem is fully optimal
    k ≥ maxIts && return :iterations                    # out of iterations
    return :noerror                                     # otherwise
end

"""
Dual conditional gradient method for the problem

    minimize 1/2||Mx-b||^2  subj to  gauge(x|A)≤τ

Includes various stopping conditions needed for
the level-set methods.
"""
function solveSubproblem!(dcg::DualCGIterable, τ::Float64, α::Float64, feaTol::Float64, maxIts::Int64, rule::String)
    setRadius!(dcg, τ)
    k = 0; u = 0.0; ℓ = 0.0
    exitFlag = :noerror
    for (gap, r) in dcg
        k = k + 1
        u = norm(r)^2/2
        ℓ = max(α, u - gap)
        exitFlag = checkExitLvl(gap, ℓ, u, k, α, feaTol, maxIts, rule)
        exitFlag == :noerror || break
    end
    s = support(dcg.A, dcg.z)
    return u, ℓ, s, k, exitFlag
end

"""
Level-set method for the problem

    min gauge(x | A) subj to 1/2||Mx-b||^2 ≤ α
"""
function level_set(M::AbstractLinearOp, b::Vector{Float64}, A::AbstractAtomicSet;
                   α::Float64 = 0.0,
                   tol::Float64 = 1e-12,
                   gapTol::Float64 = 1+norm(b),
                   maxIts::Int = size(M,2),
                   pr::Bool = true,
                   logger::Bool = true,
                   rule::String = "newton",
                   τmax::Float64 = 1.0,
                   dual_gaps::Vector{Float64} = Vector{Float64}(),
                   infeas::Vector{Float64} = Vector{Float64}()
                   )

    feaTol = (tol * norm(b))^2/2
    # Structure to hold the iterates
    dcg = DualCGIterable(M, b, A)
    # Structure to hold the solution
    sol = Solution(M, b, A)

    # head logging
    m, n = size(M)
    logger && logger_head_lvl(m, n, feaTol, α, maxIts)

    # Check whether b is zero
    if norm(b) < tol
        return sol
    end

    # Initialization at τ = 0
    τ = 0.0
    u = norm(b)^2/2; ℓ = norm(b)^2/2
    if rule == "newton"
        z = getNegGradient(dcg)
        s = support(A, z)
    elseif rule == "bisection"
        τmin = 0.0
    else
        println("invalid updating rule")
        return sol
    end

    # ----------------------------------------------------------------
    # min loop
    # ----------------------------------------------------------------
    k = 0; totIts = 0; exitFlag = :suboptimal_small
    while true
        k = k + 1

        # --------------------------------------------------------------
        # Update level-set radius.
        # --------------------------------------------------------------
        if rule == "newton"
            τ = τ + (ℓ-α)/s
        elseif rule == "bisection"
            if exitFlag == :suboptimal_small
                τmin = τ
            elseif exitFlag == :suboptimal_large
                τmax = τ
            end
            τ = (τmin + τmax)/2
            ρ = τmax - τmin
            @show ρ
        end

        # --------------------------------------------------------------
        # Check for exit.
        # --------------------------------------------------------------
        exitFlag == :iterations && break
        exitFlag == :feasible && break
        exitFlag == :optimal && break

        # --------------------------------------------------------------
        # Solve subproblem to obtain lower minorant defined by (ℓ,s).
        # --------------------------------------------------------------
        u, ℓ, s, minorIts, exitFlag = solveSubproblem!(dcg, τ, α, feaTol, maxIts-totIts, rule)

        # --------------------------------------------------------------
        # Primal recovery.
        # --------------------------------------------------------------
        # if pr && u - α ≤ gapTol
        if pr
            exitFlag = primalrecover!(dcg, sol, α, feaTol, exitFlag)
        end

        # --------------------------------------------------------------
        # record dual gaps and primal recover infeas
        # --------------------------------------------------------------
        push!(dual_gaps, dual_obj_gap(dcg, τmax, α))
        push!(infeas, sol.feasnew)

        # --------------------------------------------------------------
        # Logging and bookkeeping.
        # --------------------------------------------------------------
        logger && logger_level_lvl(α, τ, k, minorIts, ℓ, u, exitFlag, sol.feas)
        totIts += minorIts
    end

    primalrecover!(dcg, sol, α, feaTol, exitFlag)
    sort!(dual_gaps, rev=true)
    sort!(infeas, rev=true)

    # foot logging
    logger&& logger_foot_lvl(b, sol.feas, totIts, u-ℓ)

    return sol
end

function logger_head_lvl(m, n, feaTol, α, maxIts)
    println()
    println("  -------------------------------------------------------------------------")
    println("  Polar Level Set Method")
    println("  -------------------------------------------------------------------------")
    @printf("  %-22s %7d %7s %-20s %7d\n","number of variables",n,"","number of constraints",m)
    @printf("  %-22s %7.2e %7s %-20s %7.2e\n","feasibility tolerance",feaTol,"","α",α)
    @printf("  %-22s %7d \n","max iterations",maxIts)
    println("  -------------------------------------------------------------------------")
    @printf("  %5s   %8s   %8s   %8s   %8s   %8s  %15s  %10s\n",
            "Major","Minor","u-α","ℓ-α","gap","τ","infeas-α","Subproblem")
end

function logger_foot_lvl(b, feas, totIts, gap)
    scale = 1+norm(b)
    normr = sqrt(2*feas)
    println("  -------------------------------------------------------------------------")
    @printf("  %-22s %8.1e %7s %-22s %8.2e\n","residual (abs)",normr,"","optimality gap (abs)",gap)
    @printf("  %-22s %8.1e %7s %-22s %8.2e\n","residual (rel)",normr/scale,"","optimality gap (rel)",gap/scale)
    @printf("  %-22s %8d %7s %-22s %8s\n","total iterations",totIts,"","","")
    println("  -------------------------------------------------------------------------")
end

function logger_level_lvl(α, τ, k, minorItns, ℓ, u, exitFlag, feas)
    gap = u - ℓ
    # optimal = log10(dot(p.b,p.oracle.r)/norm(p.oracle.z,Inf))
    @printf("  %5d   %8d   %8.2e   %8.2e   %8.2e   %8.2e   %12.2e   %s\n",
            k, minorItns, u-α, ℓ-α, gap, τ, feas-α, exitFlag)
end

function dual_obj_gap(p::DualCGIterable, τ::Float64, α::Float64)
    β = τ
    s = support(p.A, p.z)
    y = deepcopy(p.r); y .*= 1 / s
    dobj = norm(y)^2/(2*β) + β*α - dot(y, p.b)
    return abs(dobj + τ)
end


using JackPolynomials
using Test

@testset "Jack" begin
  @testset "Jack at x = 0" begin
    @test isapprox(Jack([0.0; 0.0], Int64[], 2.0), 1.0)
    @test isapprox(Jack([0.0; 0.0; 0.0], [3; 2], 2.0), 0.0)
  end
  @testset "Jack = 0 when l(lambda)>l(x)" begin
    @test isapprox(Jack([1.0; 2.0], [3; 2; 1], 2.0), 0.0)
    @test isapprox(Jack([1.0; 2.0], [3; 2; 1], 1.0), 0.0)
  end
  @testset "Jack = 1 when the partition is empty" begin
    @test isapprox(Jack([1.0; 2.0], Int64[], 2.0), 1.0)
    @test isapprox(Jack([1.0; 2.0], [0; 0], 2.0), 1.0)
  end
  @testset "Jack - some values" begin
    @test isapprox(Jack([2.0; 3.0; 4.0; 5.0], [3; 1], 2.5), 111217.0)
    @test isapprox(Jack([2.0; 3.0; 4.0; 5.0], [4], 2.5), 500338.5)
  end
end

@testset "Jack symbolic" begin
  @testset "Jack at x = 0" begin
    @test isapprox(JackPolynomial(2, Int64[], 2.0)(0.0, 0.0), 1.0)
    @test isapprox(JackPolynomial(3, [3; 2], 2.0)(0.0, 0.0, 0.0), 0.0)
  end
  @testset "Jack = 0 when l(lambda)>l(x)" begin
    @test JackPolynomial(2, [3; 2; 1], 2.0) == 0.0
    @test JackPolynomial(2, [3; 2; 1], 1.0) == 0.0
  end
  @testset "Jack = 1 when the partition is empty" begin
    @test JackPolynomial(2, Int64[], 2.0) == 1.0
    @test JackPolynomial(2, [0; 0], 2.0) == 1.0
  end
  @testset "Jack - some values" begin
    @test isapprox(JackPolynomial(4, [3; 1], 2.5)(2.0, 3.0, 4.0, 5.0), 111217.0)
    @test isapprox(JackPolynomial(4, [4], 2.5)(2.0, 3.0, 4.0, 5.0), 500338.5)
  end
end

@testset "Schur" begin
  @testset "Schur expansion of (sum x_i)^n" begin
    x = [3.0; 4.0; 5.0; 6.0]
    s = Schur(x, [4]) + 3*Schur(x, [3,1]) + 2*Schur(x, [2,2]) +
      3*Schur(x, [2,1,1]) + Schur(x, [1,1,1,1])
    @test isapprox(s, sum(x)^4)
  end
  @testset "Schur = 0 when l(lambda)>l(x)" begin
    @test isapprox(Schur([1.0; 2.0], [3; 2; 1]), 0.0)
  end
  @testset "Schur (3,2)" begin
    x = [0.3, 2.0, 5.0]
    expected = x[1]^3*x[2]^2 + x[1]^3*x[3]^2 + x[1]^3*x[2]*x[3] +
      x[1]^2*x[2]^3 + x[1]^2*x[3]^3 + 2*x[1]^2*x[2]*x[3]^2 +
      2*x[1]^2*x[2]^2*x[3] + x[1]*x[2]*x[3]^3 + 2*x[1]*x[2]^2*x[3]^2 +
      x[1]*x[2]^3*x[3] + x[2]^2*x[3]^3 + x[2]^3*x[3]^2
    @test isapprox(Schur(x, [3; 2]), expected)
  end
end

@testset "Zonal" begin
  @testset "Zonal sum to the trace" begin
    x = [3.0; 4.0; 5.0]
    z = Zonal(x, [3]) + Zonal(x, [2,1]) + Zonal(x, [1,1,1])
    @test isapprox(z, sum(x)^3)
  end
  @testset "Zonal = 0 when l(lambda)>l(x)" begin
    @test isapprox(Zonal([1.0; 2.0], [3; 2; 1]), 0.0)
  end
end

@testset "ZonalQ" begin
  @testset "ZonalQ sum to the trace" begin
    x = [3.0; 4.0; 5.0]
    z = ZonalQ(x, [3]) + ZonalQ(x, [2,1]) + ZonalQ(x, [1,1,1])
    @test isapprox(z, sum(x)^3)
  end
end


#=
This module contains functions that allow operations in extension fields
GF(2^m). It also contains a dictionary with primitive polynomials for
small extension fields.

Author: Ivan A. Moreno Soto
Last updated: 2019/July/02
=#

module ExtensionField

export primitivePolynomials,
    add,
    multiply,
    divide,
    modulo,
    power,
    inverse,
    squareroot,
    polygcd

primitivePolynomials = Dict(
    2 => UInt32(7),
    3 => UInt32(11),
    4 => UInt32(19),
    5 => UInt32(37),
    6 => UInt32(103),
    7 => UInt32(131),
    8 => UInt32(487),
    9 => UInt32(901),
    10 => UInt32(1033),
    11 => UInt32(2053),
    12 => UInt32(4179),
    13 => UInt32(8219),
    14 => UInt32(16707),
    15 => UInt32(32785),
)

"""
    add(a::UInt32, b::UInt32)

XORs two polynomials in GF(2).
"""
function add(a::UInt32, b::UInt32)
    return a ⊻ b
end

"""
    multiply(a::UInt32, b::UInt32)

Multiplies two polynomials in GF(2). It doesn't reduce modulo
an irreducible polynomial.
"""
function multiply(a::UInt32, b::UInt32)
    partial = Array{UInt32, 1}()
    for (i, bit) in enumerate(digits(b, base=2))
        if bit == 1
            push!(partial, a << (i - 1))
        end
    end

    c = UInt32(0)
    for i in partial
        c ⊻= i
    end

    return c
end

"""
    division(a::UInt32, b::UInt32)

Computes a/b. Returns the quotient and the remainder.
"""
function division(a::UInt32, b::UInt32)
    if a == b
        return UInt32(1), UInt32(0)
    elseif length(digits(a, base=2)) < length(digits(b, base=2)) #a < b
        return UInt32(0), a
    end

    quotient = UInt32(0)
    remainder = a
    for i in (length(digits(a, base=2)) - length(digits(b, base=2))):-1:0
        if length(digits(remainder, base=2)) >= length(digits(b << i, base=2))
            remainder ⊻= b << i
            quotient ⊻= UInt32(2^i)
        end
    end

    return quotient, remainder
end

"""
    divide(a::UInt32, b::UInt32)

Computes a/b but only returns the quotient.
"""
function divide(a::UInt32, b::UInt32)
    quotient, _ = division(a, b)
    return quotient
end

"""
    modulo(a::UInt32, b::UInt32)

Computes a/b but only returns the remainder.
"""
function modulo(a::UInt32, b::UInt32)
    _, remainder = division(a, b)
    return remainder
end

"""
    power(x::UInt32, n::Int, g::UInt32)

Computes and returns x^n modulo g in GF(2^m).
"""
function power(x::UInt32, n::Int, g::UInt32)
    if n == 0
        return UInt32(1)
    end

    y = UInt32(1)

    while n > 1
        if n % 2 == 0
            n = n / 2
        else
            y = modulo(multiply(x, y), g)
            n = (n - 1)/2
        end

        x = modulo(multiply(x, x), g)
    end

    return modulo(multiply(x, y), g)
end

"""
    inverse(a::UInt32, g::UInt32)

Computes and returns the inverse of a mod g(x).
"""
function inverse(a::UInt32, g::UInt32)
    # Extended euclidean algorithm.
    t = UInt32(0)
    T = UInt32(1)
    r = g
    R = a

    while R != 0
        q = divide(r, R)
        x = T; T = add(t, multiply(q, T)); t = x
        x = R; R = add(r, multiply(q, R)); r = x
    end

    if r > 1
        error("The element $(a) doesn't have an inverse mod $(g).")
    end

    t = modulo(t, g)

    return t
end

"""
    squareroot(p::UInt32, m)

Computes the square root of p in GF(2^m).
"""
function squareroot(p::UInt32, m)
    if p == 0
        return UInt32(0)
    elseif p == 1
        return UInt32(1)
    end

    return power(p, 2^(m - 1), primitivePolynomials[m])
end

"""
    polygcd(a::UInt32, b::UInt32)

Computes the greatest common divisor of polynomials a and b.
"""
function polygcd(a::UInt32, b::UInt32)
    while b != 0
        t = b
        b = modulo(a, b)
        a = t
    end

    return a
end

end # ExtensionField


#=

  Golomb ruler in Julia ConstraintSolver.jl

  A Golomb ruler is a set of integers (marks) a(1) < ...  < a(n) such
  that all the differences a(i)-a(j) (i > j) are distinct.  Clearly we
  may assume a(1)=0.  Then a(n) is the length of the Golomb ruler.
  For a given number of marks, n, we are interested in finding the
  shortest Golomb rulers.  Such rulers are called optimal.

  See http://www.research.ibm.com/people/s/shearer/grule.html

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#

using ConstraintSolver, JuMP
using Cbc, GLPK, Ipopt
const CS = ConstraintSolver
include("constraints_utils.jl")

function golomb_ruler(n=6,print_solutions=true,all_solutions=true,timeout=6)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)
    ipopt_optimizer = optimizer_with_attributes(Ipopt.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   # "all_solutions"=> all_solutions,
                                                            "all_optimal_solutions"=>all_solutions,
                                                            "logging"=>[],

                                                            # "traverse_strategy"=>:BFS,
                                                            "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            # "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            "branch_split"=>:InHalf,

                                                            # https://wikunia.github.io/ConstraintSolver.jl/stable/options/#branch_strategy-(:Auto)
                                                            "branch_strategy" => :IMPS, # default
                                                            # "branch_strategy" => :ABS, # Activity Based Search
                                                            # "activity.decay" => 0.999, # default 0.999
                                                            # "activity.max_probes" => 6, # default, 10
                                                            # "activity.max_confidence_deviation" => 10, # default 20

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=> timeout,

                                                            # "backtrack" => false, # default true
                                                            # "backtrack_sorting" => false, # default true

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))

    # m = 2^(n-1) -1
    m = n^2
    @variable(model, 0 <= x[1:n] <= m, Int)

    @constraint(model, x[1] == 0)
    @constraint(model, x in CS.AllDifferent())
    increasing(model,x)

    # This don't work:
    # """Each variable must be an integer and bounded."""
    # diffs = [x[i]-x[j] for i in 1:n, j in 1:n if i < j]
    # println("diffs: $diffs")
    # This works:
    diffs = []
    for i in 1:n, j in 1:n
        if i != j
            d = @variable(model,integer=true, lower_bound = -m, upper_bound = m )
            @constraint(model, d == x[i]-x[j])
            push!(diffs,d)
        end
    end

    @constraint(model, diffs in CS.AllDifferent())

    # Symmetry breaking
    if n > 2
        @constraint(model, x[2] - x[1] < x[n] - x[n-1] )
    end
    @constraint(model, diffs[1] < diffs[end])

    @objective(model,Min,x[n])

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        # println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                # println("solution #$sol")
                x_val = convert.(Integer,JuMP.value.(x; result=sol))
                diffs_val = convert.(Integer,JuMP.value.(diffs; result=sol))
                println("x:$x_val")
                println("diffs:$diffs_val")
                println()

            end
        end
    else
        println("status:$status")
    end

    return status
end # end golomb_ruler

# @time golomb_ruler(8,true,false) # 1.5

# n  time
# -------
#  8   1.3s
#  9  13.3s
# 10  > 60s
for n in 2:10
    println("\nn:$n")
    @time status = golomb_ruler(n,true,false,60)
    if status == MOI.TIME_LIMIT
        break
    end
end


settings = class_settings();

context("nlp") do
    # - x - y

    a = sparse([-2.0, -2.0, 0.0])
    b = 0.0
    P = spdiagm([0.0, 0.0, 0.0])
    obj = nl_quad(a,P,b)

    # x^2 + y^2 - s - 4.0
    # x^2 + y^2 >= 4.0
    a1 = sparse([0.0, 0.0, -1.0])
    P1 = spdiagm([2.0, 2.0, 0.0])
    b1 = -4.0
    constr1 = nl_quad(a1,P1,b1)

    a1 = sparse([0.0, 0.0, -1.0])
    P1 = spdiagm([2.0, 2.5, 0.0])
    b1 = -4.0
    constr1_b = nl_quad(a1,P1,b1)

    # x + y + s - 1.0 = 0.0
    # x + y = 1.0
    a2 = sparse([1.0, 1.0, 0.0])
    P2 = spdiagm([0.0, 0.0, 0.0])
    b2 = -1.0 #-4.0
    constr2 = nl_quad(a2,P2,b2)

    # x + x^2 - y = 0
    # x + x^2 = y
    a3 = sparse([1.0, -1.0, 0.0])
    P3 = spdiagm([1.0, 0.0, 0.0])
    b3 = 0.0 #-4.0
    constr3 = nl_quad(a3,P3,b3)

    # x^2 + y^2 - 6x - 6y - s + 17 = 0
    # (x-3)^2 + (y-3)^2 - s - 1 = 0
    # (x-3)^2 + (y-3)^2 <= 1
    # ball function?
    a3 = sparse([-6.0, -6.0, -1.0])
    P3 = spdiagm([1.0, 1.0, 0.0])
    b3 = -17.0 #-4.0
    constr4 = nl_quad(a3,P3,b3)

    # convex problem
    println("CONVEX INFEASIBLE PROBLEM")
    nlp = class_nlp(3, obj, [constr2, constr4])
    vars = class_variables(3,2);
    status, vars = homogeneous_algorithm(nlp, vars, settings)
    @show x_scaled(vars)
    @show y_scaled(vars)

    println("NONCONVEX INFEASIBLE PROBLEM 1")
    nlp = class_nlp(3, obj, [constr1, constr2])
    vars = class_variables(3,2);
    status, vars = homogeneous_algorithm(nlp, vars, settings)
    @show x_scaled(vars)
    @show y_scaled(vars)

    println("NONCONVEX INFEASIBLE PROBLEM 2")
    nlp = class_nlp(3, obj, [constr1_b, constr2])
    vars = class_variables(3,2);
    status, vars = homogeneous_algorithm(nlp, vars, settings)
    @show x_scaled(vars)
    @show y_scaled(vars)

    nlp = class_nlp(3, obj, [constr1, constr3])
    vars = class_variables(3,2);
    status, vars = homogeneous_algorithm(nlp, vars, settings)
    @show x_scaled(vars)
    @show y_scaled(vars)
    #@show internal_eval_hesslag_prod(nlp,x_scaled(vars),y_scaled(vars))

    nlp = class_nlp(3, obj, [constr2, constr3])
    vars = class_variables(3,2);
    status, vars = homogeneous_algorithm(nlp, vars, settings)
    @show x_scaled(vars)
    @show y_scaled(vars)

    #@fact status --> :locally_optimal
    #@fact norm(x_scaled(vars) - [2.0,2.0,0.0]) --> less_than(1e-6)
    #@fact tau(vars) --> greater_than(0.1)
    #@fact 0.0 --> less_than(kappa(vars))
    #@fact kappa(vars) --> less_than(1e-6)
end


using DynamicalSystems
using Plots

# Simulate an discrete dynamical system (Henon map) using DynamicalSystems.jl

@inline @inbounds function hoop(x, p, n)
    dx1 = 1.0 - p[1]*x[1]^2 + x[2]
    dx2 = p[2]*x[1]
    return SVector{2}(dx1, dx2)
end
# Jacobian:
@inline @inbounds function hoop_jac(x, p, n)
    J = @SMatrix [-2*p[1]*x[1]      1.0;
                p[2]            0.0]
    return J
end

ds = DiscreteDynamicalSystem(hoop, [0.0; 0.0], [1.4, 0.3], hoop_jac)

data = trajectory(ds, 10000; dt=1)

plot(data[:, 1], data[:, 2], seriestype=:scatter, markersize=0.5)


function VelocityScale(z,h,L,ustar,θv=nothing,wθv=nothing)
  xi=z/L
  psis=holtslag_psis(xi)
  psim=holtslag_psim(xi)
  if xi>0
      wt=ustar/psis
      wm=wt
  else
      if z/h<0.1
          wt=ustar/psis
          wm=ustar/psim
      else
          a=7.2
          xi1=0.1*h/L
          psis1=holtslag_psis(xi1)
          psim1=holtslag_psim(xi1)
          #wθv<0 ? println("wθv",wθv) : nothing
          #println("θv",θv)
          wstar=((gravity_mks/θv)*wθv*h)^(1/3)
          # this part is different than Holtslag et. al to ensure continueos
          # velositiy scale at the surface and the mixed layer above.
          ct=(ustar/wstar)^3*(1/psis1^3-1)
          cm=(ustar/wstar)^3*(1/psim1^3-1)
          wm=(ustar^3+cm*wstar^3)^(1/3)
          wt=(ustar^3+ct*wstar^3)^(1/3)
          #Pr=max(min(psis1/psim1*abs(xi1)+a*karman*(0.1*(wstar/wm)),4),1)
          #wt=wm/Pr
      end
  end
  return wt,wm
end


"""
    gpfitpwm(...)

Estimate the GP parameters with the probability weighted moments.

# Implementation

Estimation with the probability weighted moments, as described by [Hosking and Wallis (1987)](https://www.jstor.org/stable/1269343?seq=1),
is only possible in the stationary case.

See also [`gpfitpwm`](@ref) for the other methods, [`gpfit`](@ref), [`gpfitbayes`](@ref) and [`ThresholdExceedance`](@ref).

# Reference

Hosking, J. R. M. and Wallis, J. R. (1987). Parameter and quantile estimation for the Generalized Pareto distribution,
    *Technometrics*, 29:339-349.
"""
function gpfitpwm end

"""
    gpfitpwm(y::Vector{<:Real})

Estimate the GP parameters with the probability weighted moments.

"""
function gpfitpwm(y::Vector{<:Real})::pwmEVA

    model = ThresholdExceedance(Variable("y", y))

    fittedmodel = gpfitpwm(model)

    return fittedmodel

end

"""
    gpfitpwm(df::DataFrame, datacol::Symbol)

Estimate the GP parameters with the probability weighted moments.

Block maxima data are in the column `datacol` of the dataframe `df`.

"""
function gpfitpwm(df::DataFrame, datacol::Symbol)::pwmEVA

    model = ThresholdExceedance(Variable(string(datacol), df[:, datacol]))

    fittedmodel = gpfitpwm(model)

    return fittedmodel

end

"""
    gpfitpwm(model::ThresholdExceedance)

Estimate the GP parameters with the probability weighted moments.

"""
function gpfitpwm(model::ThresholdExceedance)::pwmEVA

    model = validatestationarity(model)

    y = model.data.value

    a₀ = pwm(y,1,0,0)
    a₁ = pwm(y,1,0,1)

    σ̂ = 2*a₀*a₁/(a₀ - 2*a₁)

    k = a₀ / (a₀ - 2*a₁) - 2

    ξ̂ = - k
    ϕ̂ = log(σ̂)

    θ̂ = [ϕ̂; ξ̂]

    fm = pwmEVA{ThresholdExceedance, GeneralizedPareto}(model, θ̂)

    return fm

end


using Distributed
using SharedArrays

@everywhere import FIRRemez
@everywhere import PyPlot
@everywhere import Random
@everywhere import Printf

import FileIO
import Printf

using FFTW
#import JLD
import XLSX
import BSON

fig_num = 1
PyPlot.close("all")

Random.seed!(25)

fig_num = 1

L = 25
#L = 1000
if length(ARGS) > 1
    L = tryparse(Int, ARGS[1])
end

## low-pass.
passband_ω_Float64 = 0.025*π # change this.
if length(ARGS) > 2
    passband_ω_Float64 = tryparse(Float64, ARGS[2])
end

stopband_ω_Float64 = 0.025*2*π # change this.
if length(ARGS) > 3
    stopband_ω_Float64 = tryparse(Float64, ARGS[3])
end

save_name_tag = "lowpass"
if length(ARGS) > 4
    save_name_tag = ARGS[4]
end

passband_ω = convert(BigFloat, passband_ω_Float64)
stopband_ω = convert(BigFloat, stopband_ω_Float64)

f = xx->FIRRemez.cosinetransitionlowpassfunction(xx,passband_ω,stopband_ω)

println("passband is ", passband_ω)
println("stopband is ", stopband_ω)

filter_type_num = 1 # symmetric, odd length.

savefig_flag = false
plot_output_flag = false

output_folder = "./output"
opt_params_file_name = "../default/config_opt.txt" # location of the configuration file.

# omit opt_params_file_name is using default configuration file provided by the package.
h_BigFloat, X_BigFloat, fig_num = FIRRemez.frontendfilterdesign(filter_type_num, f,
                            L, fig_num, savefig_flag, plot_output_flag,
                            opt_params_file_name)

h = convert(Vector{Float64},h_BigFloat)
X = convert(Vector{Float64}, X_BigFloat)

Printf.@printf("Order of Chebyshev is %d, length of filter is %d, FIR type %d\n", L, length(h), filter_type_num)

band_info_string = Printf.@sprintf("passband_%f_stopband_%f", passband_ω_Float64, stopband_ω_Float64)
save_name = Printf.@sprintf("type%d_%s_L%d_%s.bson",
    filter_type_num, save_name_tag, L, band_info_string)

passband = passband_ω_Float64
stopband = stopband_ω_Float64

BSON.bson(joinpath(output_folder, save_name),
    h = h, X = X, passband = passband, stopband = stopband)
#JLD.save(joinpath(output_folder, save_name), "h" = h, "X" = X, "passband" = passband, "stopband" = stopband)
#
save_name = Printf.@sprintf("type%d_%s_L%d_%s.xlsx",
    filter_type_num, save_name_tag, L, band_info_string)
XLSX.openxlsx(joinpath(output_folder, save_name), mode="w") do xf
    sheet = xf[1]
    XLSX.rename!(sheet, "filter")

    sheet["A1", dim = 1] = ["passband"; passband]
    sheet["B1", dim = 1] = ["passband"; passband]
    sheet["C1", dim = 1] = ["Chebyshev node positions"; X]
    sheet["D1", dim = 1] = ["coefficients"; h]
end

#fig_num = FIRRemez.plotmagnitudersp(h, fig_num, "filter's magnitude response")

import Plots
Plots.plotly()

ω_set_fft, DFT_evals, ω_set, DTFT_evals = FIRRemez.getfreqrsp(h)

mag_rsp_G = abs.(DTFT_evals)
mag_rsp_fft = abs.(DFT_evals)

plot_obj = Plots.plot( ω_set_fft,
mag_rsp_fft,
title = "Magnitude spectrum",
label = "DFT",
seriestype = :line,
ticks = :native,
hover = ω_set_fft,
linewidth = 4,
size = (1200, 900))

Plots.plot!(plot_obj, ω_set, mag_rsp_G, label = "DTFT",
seriestype = :line,
linestyle = :dash,
linewidth = 4)

save_name = Printf.@sprintf("type%d_%s_L%d_%s.html",
    filter_type_num, save_name_tag, L, band_info_string)
Plots.savefig(plot_obj, joinpath(output_folder, save_name))


export Regularizer
export NoRegu, L2Regu, L1Regu
export forward, backward

abstract Regularizer

immutable NoRegu <: Regularizer
  coefficient :: AbstractFloat # not used, just for consistent API
end
NoRegu() = NoRegu(0.0)

immutable L2Regu <: Regularizer
  coefficient :: AbstractFloat
end

immutable L1Regu <: Regularizer
  coefficient :: AbstractFloat
end

############################################################
# No regularization
############################################################
# This function should return a number (the regularization) to be added to the objective function value
function forward(backend::Backend, regu :: NoRegu, global_regu::AbstractFloat, param :: Blob)
  # 0, since no regularization
  return convert(eltype(param), 0)
end

# This function should compute the gradient of the regularizer and add it to the gradient blob. Note
# the gradient blob already contains computed gradient, make sure to ADD to instead of to overwrite it.
function backward(backend::Backend, regu :: NoRegu, global_regu::AbstractFloat, param :: Blob, gradient :: Blob)
  # do nothing, since no regularization
end

############################################################
# L2 regularization
############################################################
function forward(backend::CPUBackend, regu :: L2Regu, global_regu::AbstractFloat, param :: Blob)
  return regu.coefficient * global_regu * vecnorm(param.data)^2
end
function backward(backend::CPUBackend, regu :: L2Regu, global_regu::AbstractFloat, param :: Blob, gradient :: Blob)
  BLAS.axpy!(length(param), convert(eltype(param), 2 * regu.coefficient * global_regu), pointer(param.data), 1, pointer(gradient.data), 1)
end

############################################################
# L1 regularization
############################################################
function forward(backend::CPUBackend, regu :: L1Regu, global_regu::AbstractFloat, param :: Blob)
  return regu.coefficient * global_regu * sum(abs(param.data))
end
function backward(backend::CPUBackend, regu :: L1Regu, global_regu::AbstractFloat, param :: Blob, gradient :: Blob)
  coef = convert(eltype(param), regu.coefficient * global_regu)
  len = length(param)
  @simd for i = 1:len
    @inbounds gradient.data[i] += coef * sign(param.data[i])
  end
end


export doubledescription
# Naive implementation of the double description Algorithm
# See JuliaPolyhedra/ConvexHull.jl for an efficient implementation

polytypefor(::Type{T}) where {T <: Integer} = Rational{BigInt}
polytypefor(::Type{Float32}) = Float64
polytypefor(::Type{T}) where {T} = T

polyvectortype(a) = a
# TODO sparse halfspaces does not mean sparse points
polyvectortype(::Type{<:SparseVector{T}}) where T = Vector{T}

dualtype(RepT::Type{<:Representation}) = dualtype(RepT, polyvectortype(vectortype(RepT)))
function dualfullspace(rep::Representation, d::FullDim, ::Type{T}) where T
    dualfullspace(rep, d, T, polyvectortype(similar_type(vectortype(rep), d, T)))
end
function dualfullspace(rep::Representation{T}) where T
    dualfullspace(rep, FullDim(rep), polytypefor(T))
end

"""
    doubledescription(h::HRepresentation)

Computes the V-representation of the polyhedron represented by `h` using the Double-Description algorithm [1, 2].

    doubledescription(V::VRepresentation)

Computes the H-representation of the polyhedron represented by `v` using the Double-Description algorithm [1, 2].

[1] Motzkin, T. S., Raiffa, H., Thompson, G. L. and Thrall, R. M.
The double description method
*Contribution to the Theory of Games*, *Princeton University Press*, **1953**

[2] Fukuda, K. and Prodon, A.
Double description method revisited
*Combinatorics and computer science*, *Springer*, **1996**, 91-111
"""
function doubledescription end

function print_v_summary(v::VRep)
    print("$(npoints(v)) points, $(nrays(v)) rays and $(nlines(v)) lines")
end

function intersect_and_remove_redundancy(v, hs, h; verbose=0)
    if eltype(hs) <: HalfSpace
        str = "halfspace"
    else
        str = "hyperplane"
    end
    for (i, hel) in enumerate(hs)
        if verbose >= 1
            println("Intersecting $str $i/$(length(hs))")
        end
        v_int = v ∩ hel
        if verbose >= 3
            print("Removing duplicates: ")
            print_v_summary(v_int)
            println(".")
        end
        # removeduplicates is cheaper than removevredundancy since the latter
        # needs to go through all the hrep element
        v_dup = removeduplicates(v_int)
        if verbose >= 3
            print("Removing redundancy: ")
            print_v_summary(v_dup)
            println(".")
        end
        v = removevredundancy(v_dup, h)
        if verbose >= 2
            print("After intersection:  ")
            print_v_summary(v)
            println(".")
        end
    end
    return v
end

function doubledescription(h::HRepresentation; kws...)
    # The redundancy of V-elements are measured using
    # the number of hyperplane they are in. If there are
    # redundant hyperplanes, it does not matter since no point
    # will be inside them but if there are duplicates it is a problem

    # FIXME Note that removevredundancy, uses `h` which contains all hyperplanes and halfspaces
    # already taken into account but also all the other ones. We should check that this
    # is the right idea.
    h = removeduplicates(h)
    v = dualfullspace(h)
    checkvconsistency(v)
    v = intersect_and_remove_redundancy(v, hyperplanes(h), h; kws...)
    v = intersect_and_remove_redundancy(v, halfspaces(h), h; kws...)
    return v
end

function doubledescription(v::VRepresentation{T}; kws...) where {T}
    checkvconsistency(v)
    lv = convert(LiftedVRepresentation{T, Matrix{T}}, v)
    R = -lv.R
    vl = doubledescription(MixedMatHRep{T}(R, zeros(T, size(R, 1)), lv.linset); kws...)
    LiftedHRepresentation{T}(vl.R, vl.Rlinset)
end


using ZigZagBoomerang
using LinearAlgebra
using Random
using SparseArrays
using Distributions
using Statistics
using ZigZagBoomerang: idot, idot_moving!
using ReverseDiff
using Test
using FileIO
using GLM

println("Sparse logistic regression")

# Design matrix
Random.seed!(2)
sparsity(A, d = 3) = round(nnz(A)/length(A), digits=d)
include("sparsedesign.jl")
# create mock design matrix with 2 categorical explanatory variables
# and their interaction effects and 2 continuous explanatory variable
A = sparse_design([20,20], 2, 4*20)
n, p = size(A)
@show n, p
Γ = A'*A
@show sparsity(A), sparsity(Γ)

const m = 1 # m Bernoulli samples for each set of covariates
println("Av. number of regressors per column ", mean(sum(A .!= 0, dims=1)), ", row ", mean(sum(A .!= 0, dims=2)))

At = SparseMatrixCSC(A')
γ0 = 0.01
# Data from the model
xtrue = 5*randn(p).*(rand(p) .< 0.1)
sigmoid(x) = inv(one(x) + exp(-x))
lsigmoid(x) = -log(one(x) + exp(-x))
y_ = Matrix(rand(n, m) .< sigmoid.(A*xtrue))
y = sum(y_, dims=2)[:]
ny = m .- y
# prior and potential
# prior precision
ϕ(x, A, y::Vector) = γ0*dot(x,x)/2 - sum(y .* lsigmoid.(A*x) + (m .- y) .* lsigmoid.(-A*x))

# Sparse gradient
# helper functions for sparse gradient
function fdot(A::SparseMatrixCSC, At::SparseMatrixCSC, f, j, x, y)
   rows = rowvals(A)
   vals = nonzeros(A)
   s = zero(eltype(A))
   @inbounds for i in nzrange(A, j)
       s += vals[i]*y[rows[i]]*f(idot(At, rows[i], x))
   end
   s
end

sigmoidn(x) = sigmoid(-x)
nsigmoid(x) = -sigmoid(x)
# helper function for sparse gradient estimate through subsampling
function fdotr(A::SparseMatrixCSC, At::SparseMatrixCSC, j, x, μ, y, ny, k)
   rows = rowvals(A)
   vals = nonzeros(A)
   s = zero(eltype(A))
   r = nzrange(A, j)
   sampler = Random.SamplerRangeNDL(r)
   l = length(r)
   @inbounds for _ in 1:k
       i = rand(sampler)
       u = idot(At, rows[i], x)
       s += l/k*vals[i]*y[rows[i]]*sigmoidn(u)
       s += l/k*vals[i]*ny[rows[i]]*nsigmoid(u)
       u0 = idot(At, rows[i], μ)
       s -= l/k*vals[i]*y[rows[i]]*sigmoidn(u0)
       s -= l/k*vals[i]*ny[rows[i]]*nsigmoid(u0)   end
   s
end

# the same, but advancing only the coordinates necessary with `idot_moving!`
function fdot_moving(A::SparseMatrixCSC, At::SparseMatrixCSC, j, t, x, θ, t′, F, μ, y, ny, k)
   rows = rowvals(A)
   vals = nonzeros(A)
   s = zero(eltype(A))
   r = nzrange(A, j)
   sampler = Random.SamplerRangeNDL(r)
   l = length(r)
   @inbounds for _ in 1:k
       i = rand(sampler)
       u = idot_moving!(At, rows[i], t, x, θ, t′, F)
       s += l/k*vals[i]*y[rows[i]]*sigmoidn(u)
       s += l/k*vals[i]*ny[rows[i]]*nsigmoid(u)
       u0 = idot(At, rows[i], μ)
       s -= l/k*vals[i]*y[rows[i]]*sigmoidn(u0)
       s -= l/k*vals[i]*ny[rows[i]]*nsigmoid(u0)
   end
   s
end

# Gradient (found using http://www.matrixcalculus.org/ S. Laue, M. Mitterreiter, and J. Giesen. A Simple and Efficient Tensor Calculus, AAAI 2020.)
∇ϕ(x, A, At, y) = γ0*x - A'*(y .* sigmoidn.(A*x)) - A'*((m .- y).*nsigmoid.(A*x))

# Element i of the gradient exploiting sparsity
∇ϕ(x, i, A, At, y, ny = m .- y) = γ0*x[i] - fdot(A, At, sigmoidn, i, x, y) - fdot(A, At, nsigmoid, i, x, ny)
# Element i of the gradient exploiting sparsity and random subsampling
∇ϕr(x, i, A, At, μ, y, ny = m .- y, k = 5) = γ0*x[i] - fdotr(A, At, i, x, μ, y, ny, k)
# The same, but exploiting that only dependencies in subsamples need to be evaluated
∇ϕmoving(t, x, θ, i, t′, F, A, At, μ, y, ny = m .- y, k = 5) = γ0*x[i] - fdot_moving(A, At, i, t, x, θ, t′, F, μ, y, ny, k)
# Tests, to be sure
#@test norm(ReverseDiff.gradient(x->ϕ(x, A, y), xtrue) - ∇ϕ(xtrue, A, At, y)) < 1e-7
#@test norm(ReverseDiff.gradient(x->ϕ(x, A, y), xtrue) - [∇ϕ(xtrue, i, A, At, y) for i in 1:p]) < 1e-7

#res = glm(A, y/m, Binomial(m), LogitLink())
#x0 = coef(res)
#norm(xtrue - x0)

# Some Newton steps towards the mode as starting point
println("Newton steps")
x0 = 0.1rand(p)
@time for i in 1:30
    global x0
    H = hcat((sparse(ReverseDiff.gradient(x -> ∇ϕ(x, i, A, At, y, ny), x0)) for i in 1:p)...)
    x0 = x0 - (Symmetric(0.00I + H))\[∇ϕ(x0, i, A, At, y, ny) for i in 1:p]
end
norm(xtrue - x0)

# Note that the A'A has the same sparsity structure as the Hessian of ϕ
# We can use it to tune the doubly local ZigZig with spdmp sparsifying the
# precision matrix estimate
Γ0 = A'*A
nonzeros(Γ0) .= 1

Γ = hcat((sparse(ReverseDiff.gradient(x -> ∇ϕ(x, i, A, At, y), x0)) for i in 1:p)...)
μ = copy(x0)

Γdrop = droptol!(copy(Γ), 1e-2)
@assert nnz(diag(Γdrop)) == p
println("Precision Γ is ", 100*nnz(Γ)/length(Γ), "% filled")
println("Dropped precision is ", 100*nnz(Γdrop)/length(Γdrop), "% filled")

# Random initial values
t0 = 0.0

# Define ZigZag
c = 0.01*ones(p) # Rejection bounds
c0 = copy(c)
σ = sqrt.(diag(inv(Matrix(Γ))))
#σ = (Vector(diag(Γ))).^(-0.5) # cheaper

Z = ZigZag(Γ, μ, σ, ρ=0.5, λref=0.00)
Zdiag = ZigZag(sparse(Diagonal(Γ)), μ, σ, ρ=0.5, λref=0.00)
Zdrop = ZigZag(Γdrop, μ, σ, ρ=0.5, λref=0.00)
T = 2000.0

θ0 = rand([-1.0,1.0], p) .* σ

# Run sparse ZigZag for T time units and collect trajectory
println("Distance starting point")
@show norm(x0 - xtrue)

println("Run spdmp")
trace0, u0, (acc0,num0), c0 = @time spdmp(∇ϕmoving, t0, x0, θ0, T, c, Zdrop, SelfMoving(), A, At, μ, y, m .- y, 10; structured = true, adapt=true, factor=5)

c = 0.01*ones(p)
κ = (γ0/sqrt(2π))/(1/0.9 -1)*ones(p)
trace, u, (acc,num), c = @time sspdmp(∇ϕmoving, t0, x0, θ0, T, c, Zdrop, κ, SelfMoving(), A, At, μ, y, m .- y, 10; structured = true, adapt=true, factor=5)

using Pkg
Pkg.activate(@__DIR__)
cd(@__DIR__)
using Makie, AbstractPlotting
using ColorSchemes, GeometryBasics

@time traj0 = collect(discretize(trace0, 0.1))
@time traj = collect(discretize(trace, 0.1))

xhat0 = [median(getindex.(last.(traj0), i)) for i in eachindex(μ)]
xhat = [median(getindex.(last.(traj), i)) for i in eachindex(μ)]
xcol = xtrue.*(xtrue .!= 0.0) .- (xtrue .== 0.0)*16.0
fi0 = linesegments(repeat(1:p, inner = 2), vec([xhat0 xtrue]'), linewidth = 0.5, resolution = (1200,900))
scatter!(1:p, xhat0, color = xcol, colormap=:berlin, markersize=7, strokewidth = 0.5,
            marker = map(x -> x != 0 ? :circle : :x, xtrue))
fi0
fi = linesegments(repeat(1:p, inner = 2), vec([xhat xtrue]'), linewidth = 0.5, resolution = (1200,900))
scatter!(1:p, xhat, color = xcol, colormap=:berlin, markersize=7,
            strokewidth = 0.5,marker = map(x -> x != 0 ? :circle : :x, xtrue))
fi

save("logistic.png", fi0)
save("logisticsticky.png", fi)

@show acc/num
@show extrema(c ./ c0)

function confusion(xhat)
    Dict(
    "false negative" => sum((xhat .== 0) .& (xtrue .!= 0)),
    "true positive" => sum((xhat .!= 0) .& (xtrue .!= 0)),
    "true negative" => sum((xhat .== 0) .& (xtrue .== 0)),
    "false positive" => sum((xhat .!= 0) .& (xtrue .== 0)),
    )
end

confusion(xhat)

confusion(abs.(xhat0) .> 1.9)

dt = T/4000
x̂ = mean(x for (t,x) in discretize(trace, dt))
X = Float64[]
ts = Float64[]
for (t,x) in discretize(trace, dt)
    append!(ts, t)
    append!(X, x)
end
X = reshape(X, p, length(X)÷p)

[xtrue x̂ abs.(xtrue - x̂)]

@show norm(xtrue - μ)
@show norm(xtrue - x̂)

println("Plot")
using Makie
using Colors
using GoldenSequences

ps = [4, 22, 50, 378, p-2, p-1]
p0 = length(ps)
cs = map(x->RGB(x...), (Iterators.take(GoldenSequence(3), p0)))
pis = []
fig = Figure(resolution = (1200, 1200))
for i in 1:p0
    global p_1
    p_i = fig[i, 1] = Axis(fig)
    if i == 1
        p_1 = p_i
    end
    lines!(p_i, ts, X[ps[i], :], color=cs[i])
    lines!(p_i, [0, T], [xtrue[ps[i]], xtrue[ps[i]]], color=cs[i], linewidth=2.0)
    p_i.ylabel = "var$(ps[i])"
    p_i.xlabel = "t"
    i > 0 && linkaxes!(p_1, p_i)
end
#fig[0, :] = Label(fig, "Sparse design logistic regression n=$(m*n), p=$p", textsize = 30)
fig
save("logistic$(typeof(Z).name).png", fig)

#p2 = title(vbox(hbox([text("$p", show_axis=false, textsize=10) for p in ps]...),
#    [hbox([i==j ? lines(ts, X[i,:]) : lines(X[i, :], X[j, :]) for i in ps]...) for j in ps]...), "$ps")

p1


let

        # A �~F~R C �~F~P B
        bn = BayesNet()
        push!(bn, StaticCPD(:a, Categorical([1.0,0.0])))
        push!(bn, StaticCPD(:b, Categorical([0.0,1.0])))
        push!(bn, CategoricalCPD{Bernoulli}(:c, [:a, :b], [2,2], [Bernoulli(0.1), Bernoulli(0.2), Bernoulli(1.0), Bernoulli(0.4)]))

        t5 = gibbs_sample(bn, 5, 100; thinning=5, consistent_with=Assignment(),
             variable_order=nothing, time_limit=nothing,
             error_if_time_out=true, initial_sample=nothing)
        @test size(t5) == (5, 3)
        @test t5[!,:a] == [1,1,1,1,1]
        @test t5[!,:b] == [2,2,2,2,2]
        @test t5[!,:c] == [1,1,1,1,1]

        config = GibbsSampler(burn_in=100, thinning=5)
        t5 = rand(bn, config, 5)
        @test size(t5) == (5, 3)
        @test t5[!,:a] == [1,1,1,1,1]
        @test t5[!,:b] == [2,2,2,2,2]
        @test t5[!,:c] == [1,1,1,1,1]

        bn2 = BayesNet()
        push!(bn2, StaticCPD(:a, Categorical([0.5,0.5])))
        push!(bn2, StaticCPD(:b, Categorical([0.5,0.5])))
        push!(bn2, CategoricalCPD{Categorical}(:c, [:a, :b], [2,2], [
               Categorical([1.0, 0, 0, 0]), Categorical([0, 1.0, 0, 0]),
               Categorical([0, 0, 1.0, 0]), Categorical([0, 0, 0, 1.0])]))

        t6 = gibbs_sample(bn2, 5, 100; thinning=5, consistent_with=Assignment(:c=>1),
             variable_order=nothing, time_limit=nothing,
             error_if_time_out=true, initial_sample=nothing)
        @test t6[!,:a] == [1,1,1,1,1]
        @test t6[!,:b] == [1,1,1,1,1]
        @test t6[!,:c] == [1,1,1,1,1]

	config.evidence = Assignment(:c=>1)
	t6 = rand(bn2, config, 5)
	@test t6[!,:a] == [1,1,1,1,1]
        @test t6[!,:b] == [1,1,1,1,1]
        @test t6[!,:c] == [1,1,1,1,1]

        t7 = gibbs_sample(bn2, 5, 100; thinning=5, consistent_with=Assignment(:c=>2),
             variable_order=nothing, time_limit=nothing,
             error_if_time_out=true, initial_sample=nothing)
        @test t7[!,:a] == [2,2,2,2,2]
        @test t7[!,:b] == [1,1,1,1,1]
        @test t7[!,:c] == [2,2,2,2,2]

        rand()
	config.evidence = Assignment(:c=>2)
        t7 = rand(bn2, config, 5)
        @test t7[!,:a] == [2,2,2,2,2]
        @test t7[!,:b] == [1,1,1,1,1]
        @test t7[!,:c] == [2,2,2,2,2]

        bn3 = BayesNet()
        push!(bn3, StaticCPD(:a, Normal(2.5, 1.5)))
        push!(bn3, StaticCPD(:b, Categorical([0.3,0.3,0.4])))
        push!(bn3, LinearGaussianCPD(:c, [:a, :b], [0.5, 0.5], 1.0, 0.5))

        t8 = gibbs_sample(bn3, 5, 100; thinning=5, consistent_with=Assignment(),
             variable_order=nothing, time_limit=nothing,
             error_if_time_out=true, initial_sample=nothing)
        @test size(t8) == (5, 3)

        # unlikely c
        t9 = gibbs_sample(bn3, 5, 100; thinning=5, consistent_with=Assignment(:c=>1.0),
             variable_order=nothing, time_limit=nothing,
             error_if_time_out=true, initial_sample=nothing)
        @test size(t9) == (5, 3)

        # use optional parameters and border cases for other parameters
        t10 = gibbs_sample(bn3, 5, 0; thinning=0, consistent_with=Assignment(:c=>2.0, :b=>1),
             variable_order=(Vector{Symbol}([:c, :a, :b])), time_limit=1000000,
             error_if_time_out=false, initial_sample=Assignment(:a=>1, :b=>1, :c=>2.0))
        @test size(t10)[2] == 3 && size(t10)[1] <= 5

	config.burn_in = 0
	config.evidence = Assignment(:c=>2.0, :b=>1)
	config.variable_order = Vector{Symbol}([:c, :a, :b])
	config.time_limit = 1000000
	config.error_if_time_out = false
	config.initial_sample = Assignment(:a=>1, :b=>1, :c=>2.0)
	t10 = rand(bn3, config, 5)
        @test size(t10)[2] == 3 && size(t10)[1] <= 5

        # test early return from time limit
        t11 = gibbs_sample(bn3, 25000, 100; thinning=4, consistent_with=Assignment(),
             variable_order=nothing, time_limit= 1000,
             error_if_time_out=false, initial_sample=nothing)
        @test size(t10)[2] == 3 && size(t10)[1] <= 25000

	config = GibbsSampler(burn_in=100, thinning=4, time_limit = 1000, error_if_time_out=false)
	t11 = rand(bn3, config, 25000)
        @test size(t10)[2] == 3 && size(t10)[1] <= 25000

        d_bn = DiscreteBayesNet()
        push!(d_bn, DiscreteCPD(:a, [1.0,0.0]))
        push!(d_bn, DiscreteCPD(:b, [0.0,1.0]))
        push!(d_bn, CategoricalCPD{Categorical{Float64,Vector{Float64}}}(:c, [:a, :b], [2,2],
                        [Categorical([0.1, 0.9]), Categorical([0.2, 0.8]), Categorical([1.0, 0.0]), Categorical([0.4, 0.6])]))

	t12 = gibbs_sample(d_bn, 5, 100; thinning=5, consistent_with=Assignment(),
             variable_order=nothing, time_limit=nothing,
             error_if_time_out=true, initial_sample=nothing)
        @test size(t12) == (5, 3)
        @test t12[!,:a] == [1,1,1,1,1]
        @test t12[!,:b] == [2,2,2,2,2]
        @test t12[!,:c] == [1,1,1,1,1]

        t13 = gibbs_sample(d_bn, 5, 100; thinning=5, consistent_with=Assignment(:c=>1),
             variable_order=nothing, time_limit=nothing,
             error_if_time_out=true, initial_sample=nothing)
        @test size(t13) == (5, 3)
        @test t13[!,:a] == [1,1,1,1,1]
        @test t13[!,:b] == [2,2,2,2,2]
        @test t13[!,:c] == [1,1,1,1,1]

        # test bad parameters given to gibbs_sample
        pass = false
        try
            error_test = gibbs_sample(bn, 0, 100; thinning=5, consistent_with=Assignment(),
                 variable_order=nothing, time_limit=nothing,
                 error_if_time_out=true, initial_sample=nothing)
        catch e
            pass = true
        end
        @test pass

        pass = false
        try
            error_test = gibbs_sample(bn, 5, -1; thinning=5, consistent_with=Assignment(),
                 variable_order=nothing, time_limit=nothing,
                 error_if_time_out=true, initial_sample=nothing)
        catch e
            pass = true
        end
        @test pass

        pass = false
        try
            error_test = gibbs_sample(bn, 5, 100; thinning=-1, consistent_with=Assignment(),
                 variable_order=nothing, time_limit=nothing,
                 error_if_time_out=true, initial_sample=nothing)
        catch e
            pass = true
        end
        @test pass

        pass = false
        try
            v_order = Vector{Symbol}([:a, :c])
            error_test = gibbs_sample(bn, 5, 100; thinning=5, consistent_with=Assignment(),
                 variable_order=v_order, time_limit=nothing,
                 error_if_time_out=true, initial_sample=nothing)
        catch e
            pass = true
        end
        @test pass

        pass = false
        try
            v_order = Vector{Symbol}([:a, :c])
            config = GibbsSampler(burn_in=100, variable_order=v_order, thinning=5)
            error_test = rand(bn, config, 5)
        catch e
            pass = true
        end
        @test pass

        pass = false
        try
            v_order = Vector{Symbol}([:a, :c, :b, :d])
            error_test = gibbs_sample(bn, 5, 100; thinning=5, consistent_with=Assignment(),
                 variable_order=v_order, time_limit=nothing,
                 error_if_time_out=true, initial_sample=nothing)
        catch e
            pass = true
        end
        @test pass

        pass = false
        try
            error_test = gibbs_sample(bn, 5, 100; thinning=5, consistent_with=Assignment(),
                 variable_order=nothing, time_limit=0,
                 error_if_time_out=true, initial_sample=nothing)
        catch e
            pass = true
        end
        @test pass

        pass = false
        try
            error_test = gibbs_sample(bn, 5, 100; thinning=5, consistent_with=Assignment(),
                 variable_order=nothing, time_limit=nothing,
                 error_if_time_out=true, initial_sample=Assignment(:a => 1, :b => 1))
        catch e
            pass = true
        end
        @test pass

        pass = false
        try
            config = GibbsSampler(burn_in=100, thinning=5, initial_sample=Assignment(:a => 1, :b => 1))
            error_test = rand(bn, config, 5)
        catch e
            pass = true
        end
        @test pass

        pass = false
        try
            error_test = gibbs_sample(bn, 5, 100; thinning=5, consistent_with=Assignment(:c => 2),
                 variable_order=nothing, time_limit=nothing,
                 error_if_time_out=true, initial_sample=Assignment(:a => 1, :b => 1, :c => 1))
        catch e
            pass = true
        end
        @test pass

        pass = false
        try
            error_test = gibbs_sample(bn2, 5, 100; initial_sample=Assignment(:a => 1, :b => 1, :c => 3))
        catch e
            pass = true
        end
        @test pass
end


using Aqua

Aqua.test_all(GeometricAlgebra, ambiguities=false, unbound_args=false)


#=

  Number lock problem in Turing.jl

  From Presh Talwalkar (MindYourDecisions)
  """
  Puzzles like this have been shared with the dubious claim that "only a
  genius can solve" them. But they are still fun problems so let's work one
  out.

  A number lock requires a 3 digit code. Based on these hints, can you crack
  the code?

   682 - one number is correct and in the correct position
   645 - one number is correct but in the wrong position
   206 - two numbers are correct but in the wrong positions
   738 - nothing is correct
   780 - one number is correct but in the wrong position

  Video:  https://youtu.be/-etLb-8sHBc
  """

  This model use the same approach as
  - MiniZinc model: http://hakank.org/minizinc/number_lock.mzn
  - Picat model: http://hakank.org/picat/number_lock.pi

  Summary Statistics
  parameters      mean       std   naive_se      mcse       ess      rhat   ess_per_sec
      Symbol   Float64   Float64    Float64   Float64   Float64   Float64       Float64

        x[1]    0.0000    0.0000     0.0000    0.0000       NaN       NaN           NaN
        x[2]    5.0000    0.0000     0.0000    0.0000       NaN       NaN           NaN
        x[3]    2.0000    0.0000     0.0000    0.0000       NaN       NaN           NaN

   Distributions of variable (num:0)
   [0, 5, 2]	=>	10000 (1.0)

  Cf ~/webppl/number_lock2.wppl

=#

using Turing, StatsPlots
include("jl_utils.jl")

@model function number_lock2(y, hints)
    # m = 5 # number of rows (i.e. number of hints)
    # n = 3 # number of columns (i.e. length of the numbers)
    m,n = size(y)

    # The unknown number
    x ~ filldist(DiscreteUniform(0,9),n)

    #=
    # Number of digits that has correct position and value.
    num_correct_positions = tzeros(m)
    for r in 1:m
        num_correct_positions[r] ~ Dirac(sum([y[r,i] == x[i] ? 1 : 0 for i in 1:n]))
    end

    # Number of digits that has correct value (wether in correct position or not)
    num_correct_values = tzeros(m)
    for r in 1:m
        num_correct_values[r] ~ Dirac(sum([y[r,j] == x[i] ? 1 : 0  for i in 1:n for j in 1:n]))
    end

    for r in 1:m
        hints[r,1] ~ Dirac(num_correct_positions[r])
        hints[r,2] ~ Dirac(num_correct_values[r])
    end
    =#

    # It's faster if we skip the two intermediate arrays num_correct_position and num_correct_values
    for r in 1:m
        hints[r,1] ~ Dirac(sum([y[r,i] == x[i] ? 1 : 0 for i in 1:n]))
        hints[r,2] ~ Dirac(sum([y[r,j] == x[i] ? 1 : 0  for i in 1:n for j in 1:n]))
    end

    return x
end

function run_number_lock2(y, hints)

    model = number_lock2(y,hints)

    # chns = sample(model, Prior(), 10_000)
    # chns = sample(model, MH(), 10_000)
    # chns = sample(model, PG(5), 10_000)
    chns = sample(model, SMC(), 10_000)
    # chns = sample(model, IS(), 10_000)

    display(chns)
    # display(plot(chns))

    chains_params = Turing.MCMCChains.get_sections(chns, :parameters)
    genq = generated_quantities(model, chains_params)
    show_var_dist_pct(genq)
    println()

end

# This instance: [0,5,2]
# The number combinations
y = [6 8 2
     6 4 5
     2 0 6
     7 3 8
     7 8 0]

# Hints
hints = [1 1
         0 1
         0 2
         0 0
         0 1]

run_number_lock2(y, hints)

# Another instance: https://twitter.com/natebrix/status/1418256269052321793
# The only changed is that the second row is
#   6 1 4
# instead of
#   6 4 5

# Solution:
# [0, 4, 2]	=>	10000 (1.0)#
#
y = [6 8 2
     6 1 4 # <- changed from 6 4 5
     2 0 6
     7 3 8
     7 8 0]

# Hints
hints = [1 1
         0 1
         0 2
         0 0
         0 1]
run_number_lock2(y, hints)


# ----------------------------------------------------------------------------------------------------------
# l-∞ norm
# ----------------------------------------------------------------------------------------------------------
"""
    Linf

An empty structure representing l-infinity norm.
"""
struct Linf <: Norm end

"""
    name(d::Linf)

Returns the full name of the l-infinity norm.
"""
name(d::Linf) = "l-infinity norm"

"""
    normtype(d::Linf)

Returns the type of l-infinity norm.
"""
normtype(d::Linf) = Inf

"""
    optimal(s::Our, d::Linf, m::DRO; kwargs...)

Returns the optimal solution of the DRO model 'm' with l-infinity norm.
"""
function optimal(s::Our, d::Linf, m::DRO)
    c      = .- m.c
    perm   = sortperm(c)
    c      = c[perm]
    p      = m.q[perm]
    i      = 1
    j      = length(c)
    δ_down = 0

    while i < j
        δ1    = min(1 - p[i], m.ε)
        δ2    = 0
        p[i] += δ1

        while δ2 < δ1
            if min(p[j], m.ε - δ_down) >= δ1 - δ2
                p[j]   -= δ1 - δ2
                δ_down += δ1 - δ2
                break
            else
                δ2     += min(p[j], m.ε - δ_down)
                p[j]   -= min(p[j], m.ε - δ_down)
                δ_down  = 0
                j      -= 1

                if i == j
                    p[i] += - δ1 + δ2
                    break
                end
            end
        end
        i += 1
    end
    return p[invperm(perm)]
end

# ----------------------------------------------------------------------------------------------------------
# l-1 norm
# ----------------------------------------------------------------------------------------------------------
"""
    Lone

An empty structure representing l-1 norm.
"""
struct Lone <: Norm end

"""
    name(d::Lone)

Returns the full name of the l-1 norm.
"""
name(d::Lone) = "l-1 norm"

"""
    normtype(d::Lone)

Returns the type of l-1 norm.
"""
normtype(d::Lone) = 1

"""
    optimal(s::Our, d::Lone, m::DRO)

Returns the optimal solution of the DRO model 'm' with l-1 norm.
"""
function optimal(s::Our, d::Lone, m::DRO; kwargs...)
    c    = .- m.c
    perm = sortperm(c)
    c    = c[perm]
    p    = m.q[perm]
    i    = 1
    j    = length(c)
    δ_up = 0

    while i < j && δ_up < m.ε/2
        δ1   = min(1 - p[i], m.ε/2-δ_up)
        δ_up += δ1
        δ2   = 0
        p[i] += δ1

        while δ2 < δ1
            if p[j] >= δ1 - δ2
                p[j] += - δ1 + δ2
                break
            else
                δ2 = δ2 + p[j]
                p[j] = 0
                j -= 1
                if i == j
                    p[i] += - δ1 + δ2
                    break
                end
            end
        end
        i += 1
    end
    return p[invperm(perm)]
end

# ----------------------------------------------------------------------------------------------------------
# l-2 norm
# ----------------------------------------------------------------------------------------------------------
"""
    Ltwo

An empty structure representing l-2 norm.
"""
struct Ltwo <: Norm end

"""
    name(d::Ltwo)

Returns the full name of the l-2 norm.
"""
name(d::Ltwo) = "l-2 norm"

"""
    normtype(d::Ltwo)

Returns the type of l-2 norm.
"""
normtype(d::Ltwo) = 2

"""
    initial(d::Ltwo, m::DRO)

Returns the initial point for finding the root of the function `h` using the newton method.
"""
function initial(d::Ltwo, m::DRO)
    f(μ)  = h(d, m, μ)
    μ     = 10
    f_val = f(μ)
    add_eval()

    while f_val >= 0
        f == 0 && return μ

        μ    *= 10
        f_val = f(μ)
        add_eval()
    end
    return μ
end

function g(d::Ltwo, m::DRO, μ::Real)
    n     = length(m.q)
    s     = sort(μ.*m.q .+ m.c)
    λ     = 0
    g     = μ
    g_old = g

    for i in n-1:-1:1
        g_old = g
        g    -= (n - i)*(s[i + 1] - s[i])
        if g <= 0
            λ = s[i] - (s[i + 1] - s[i])/(g_old - g)*g
            break
        end
    end
    if g > 0
        λ = Statistics.mean(m.c)
    end
    return λ
end

function h(d::Ltwo, m::DRO, μ::Real; λ::Real = g(d, m, μ))
    return sum((min.(λ .- m.c, μ .* m.q)).^2) - m.ε^2 * μ^2
end

function ∇h(d::Ltwo, m::DRO, μ::Real; λ::Real = g(d, m, μ))
    q_i = @view m.q[λ .- m.c .>  μ .* m.q]
    c_i = @view m.c[λ .- m.c .<= μ .* m.q]

    return 2*μ*(sum(abs2, q_i) - m.ε^2) - 2*sum(q_i)*sum(λ .- c_i)/length(c_i)
end

"""
    optimal(s::Our, d::Ltwo, m::DRO; kwargs...)

Returns the optimal solution of the DRO model 'm' with l-2 norm.
"""
function optimal(s::Our, d::Ltwo, m::DRO; kwargs...)
    Ilen = length(m.Imax)
    Isum = sum(m.q[m.Imax])
    p    = zero(m.q)
    p[m.Imax] .= m.q[m.Imax] .+ 1/Ilen .- Isum/Ilen

    if sum(abs2, p - m.q) <= m.ε^2
        return p
    else
        μ = newton(d, m; kwargs...)
        λ = g(d, m, μ)
        return max.(m.q .- (λ .- m.c)./μ, 0)
    end
end

"""
    optimal(s::Philpott, d::Philpott, m::DRO; atol::Real = 1e-10)

Returns the optimal solution given by Philpott's algorithm of the DRO model 'm' with l-2 norm.
"""
function optimal(s::Philpott, d::Ltwo, m::DRO; atol::Real = 1e-10, kwargs...)
    n      = length(m.q)
    c_mean = Statistics.mean(m.c)
    c_std  = Statistics.stdm(m.c, c_mean; corrected = false)
    p      = m.q .+ m.ε .* (m.c .- c_mean) ./ (sqrt(n) * c_std)

    isapprox(c_std, 0, atol = atol)         && return m.q
    m.ε <= sqrt(n / (n - 1)) * minimum(m.q) && return p

    I  = collect(1:n)
    j  = 0
    a1 = 0
    a2 = 0
    a3 = sqrt(n)*m.ε

    @views for I_len in n-1:-1:1

        minimum(p[I]) >= 0 && return p

        ## find index j
        J   = I[p[I] .< 0]
        ε_p = @. (a1^2 + (c_std * (I_len * m.q[J] + a1) / (m.c[J] - c_mean))^2) / I_len + a2
        j   = J[findmin(ε_p)[2]]

        deleteat!(I, searchsortedfirst(I, j))
        p[j] = 0

        ## update p
        a1 += m.q[j]
        a2 += m.q[j]^2
        a3 = sqrt(I_len*(m.ε^2 - a2) - a1^2)

        c_mean = ((I_len + 1) * c_mean - m.c[j]) / I_len
        c_std  = Statistics.stdm(m.c[I], c_mean; corrected = false)

        @. p[I] = m.q[I] + (a1 + a3 * (m.c[I] - c_mean) / c_std) / I_len

        if isapprox(c_std, 0, atol = atol)
            @. p[I] = m.q[I] + a1 / I_len
            return p
        end
    end
    p[I] = 1
    return p
end


module GradientTest

import Calculus

using Base.Test
using ForwardDiff

include(joinpath(dirname(@__FILE__), "utils.jl"))

#############################
# rosenbrock hardcoded test #
#############################

x = [0.1, 0.2, 0.3]
v = rosenbrock(x)
g = [-9.4, 15.6, 52.0]

for c in (Chunk{1}(), Chunk{2}(), Chunk{3}())

    # single-threaded #
    #-----------------#
    @test_approx_eq g ForwardDiff.gradient(rosenbrock, x, c)

    out = similar(x)
    ForwardDiff.gradient!(out, rosenbrock, x, c)
    @test_approx_eq out g

    out = GradientResult(x)
    ForwardDiff.gradient!(out, rosenbrock, x, c)
    @test_approx_eq ForwardDiff.value(out) v
    @test_approx_eq ForwardDiff.gradient(out) g

    # multithreaded #
    #---------------#
    if ForwardDiff.IS_MULTITHREADED_JULIA
        @test_approx_eq g ForwardDiff.gradient(rosenbrock, x, c; multithread = true)

        out = similar(x)
        ForwardDiff.gradient!(out, rosenbrock, x, c; multithread = true)
        @test_approx_eq out g

        out = GradientResult(x)
        ForwardDiff.gradient!(out, rosenbrock, x, c; multithread = true)
        @test_approx_eq ForwardDiff.value(out) v
        @test_approx_eq ForwardDiff.gradient(out) g
    end
end

########################
# test vs. Calculus.jl #
########################

for f in VECTOR_TO_NUMBER_FUNCS
    v = f(X)
    g = ForwardDiff.gradient(f, X)
    @test_approx_eq_eps g Calculus.gradient(f, X) FINITEDIFF_ERROR
    for c in CHUNK_SIZES
        for usecache in (true, false)
            println("  ...testing $f with (chunk size = $c) and (usecache = $usecache)")
            chunk = Chunk{c}()

            # single-threaded #
            #-----------------#
            out = ForwardDiff.gradient(f, X, chunk; usecache = usecache)
            @test_approx_eq out g

            out = similar(X)
            ForwardDiff.gradient!(out, f, X, chunk; usecache = usecache)
            @test_approx_eq out g

            out = GradientResult(X)
            ForwardDiff.gradient!(out, f, X, chunk; usecache = usecache)
            @test_approx_eq ForwardDiff.value(out) v
            @test_approx_eq ForwardDiff.gradient(out) g

            # multithreaded #
            #---------------#
            if ForwardDiff.IS_MULTITHREADED_JULIA
                out = ForwardDiff.gradient(f, X, chunk; multithread = true, usecache = usecache)
                @test_approx_eq out g

                out = similar(X)
                ForwardDiff.gradient!(out, f, X, chunk; multithread = true, usecache = usecache)
                @test_approx_eq out g

                out = GradientResult(X)
                ForwardDiff.gradient!(out, f, X, chunk; multithread = true, usecache = usecache)
                @test_approx_eq ForwardDiff.value(out) v
                @test_approx_eq ForwardDiff.gradient(out) g
            end
        end
    end
end

end # module


API_token = "" #input your Etherscan API token
using HTTP, JSON, StatsBase, Statistics, Pkg, Plots, CSV
pyplot()

fb = 6290000 #first block to start from

log_array = [] #for stroing event logs from API call
for i in 1:400 #the logs come in batches of 1000 — get as many as possible
    print(i,", ")
    link = "https://api.etherscan.io/api?module=logs&action=getLogs&fromBlock=$fb&toBlock=latest&address=0xa52e014b3f5cc48287c2d483a3e026c32cc76e6d&topic0=0x8dd0b145385d04711e29558ceab40b456976a2b9a7d648cc1bcd416161bf97b9&apikey=$API_token"

    r = HTTP.request("GET", link)
    r = String(r.body)

    x = JSON.parse(r)["result"]
    push!(log_array, x)
    fb = parse(Int64, x[end]["blockNumber"]) + 1 #for tracking progress
    println(fb)

end

result_array = Int64[] #for storing Oracle results
for i in 1:length(log_array)
    for j in 1:length(log_array[i])
        result = parse(Int64, log_array[i][j]["data"][129:130], base = 16) #results stored in hex at index 129:130
        push!(result_array, result)
    end
end
filter!(x->x !== 0, result_array) #results of 0 == failed tx

#CSV.write() — once you've pulled from API once, save results as CSV to save time in future
result_array = CSV.read("/Users/tuckercahillchambers/Desktop/Crypto projects/Etheroll_results.csv")
result_array = Matrix{Int64}(result_array)
result_array = result_array[:,1]

random = rand(1:100, 269754)
x = countmap([r for r in random])
plot(x, seriestype = :bar, color = :purple,
    xticks = (0:5:100),
    xlab = "Result", ylab = "Occurrences",
    title = "Histogram of results from rand()",
    legend = false)

std(random)

Provable = countmap([r for r in result_array])
plot(Provable, seriestype = :bar,
    xticks = (0:5:100),
    xlab = "Result", ylab = "Occurrences",
    title = "Histogram of results from Etheroll",
    legend = false)

median(collect(values(Provable)))

#if result below is true, results > 90 are less likely to repeat
#

cnt_array = []
for i in 1:length(result_array)-1
    cnt = 0
    while result_array[i] > 90
        cnt += 1
        i += 1
    end
    push!(cnt_array, cnt)
end
cnt_array

c = countmap([run for run in cnt_array])
sort(collect(c), by=x->x[1])[3:end]

#if result below is true, then there would be advantage
#in betting a low number (i.e. risking a lower-prob bet, at e.g 30)
#if 3 oracle results come back < 11, BUT WHY?
#because there seems to be ~1/5 probability of seeing another result < 11

cnt_array = []
for i in 1:length(result_array)-1
    cnt = 0
    while result_array[i] < 11
        cnt += 1
        i += 1
    end
    push!(cnt_array, cnt)
end
cnt_array

c = countmap([run for run in cnt_array])
sort(collect(c), by=x->x[1])[3:end]

t = zeros(Float64, 100,3)

for i in 1:length(result_array) - 1
    for j in 1:length(t)
        if j == result_array[i]
            if result_array[i+1] > j
                t[j, 1] += 1
            elseif result_array[i+1] < j
                t[j, 2] += 1
            end
        end
    end
end

r = zeros(Float64, 100,3)

for i in 1:length(random) - 1
    for j in 1:length(r)
        if j == random[i]
            if random[i+1] > j
                r[j, 1] += 1
            elseif random[i+1] < j
                r[j, 2] += 1
            end
        end
    end
end

for i in 1:size(t)[1]
    t[i,3] = round(minimum(t[i,1:2])/maximum(t[i,1:2]), digits = 2)
end

for i in 1:size(r)[1]
    r[i,3] = round(minimum(r[i,1:2])/maximum(r[i,1:2]), digits = 2)
end

x = collect(t[:,3])
y =  collect(r[:,3])
plot(x, xticks = collect(0:5:100), yticks = collect(0:0.1:1), ylim = (0,1.0), label = "Etheroll", lw = 3,
    ylab = "Ratio", xlab = "Result", title = "[n+1]>n / [n+1]<n")
plot!(y, lc = :purple, label = "rand()", lw = 3)

for i in 1:size(t)[1]
    t[i,3] = round(t[i,1]/t[i,2], digits = 2)
end

for i in 1:size(r)[1]
    r[i,3] = round(r[i,1]/r[i,2], digits = 2)
end

plot(t[:,3], seriestype = :scatter, xticks = collect(1:9:100),
    legend = false,
    title = "[n+1] > n  /  [n+1] < n",
    xlab = "n (Oraclize result)", ms = 6)
plot!(r[:,3], seriestype = :scatter,
    legend = false, color = :purple, ms = 4)

#What is the MEAN distance between appearances of same oracle result?
cnt = countmap([c for c in result_array])
cnt = sort(collect(cnt), by=x->x[1])

test = []
y = 0
for i in 1:100
    x = findall(x->x == i, result_array)[2:end] - findall(x->x == i, result_array)[1:end-1]
    push!(test, mean(x))
end
test_mean = round.(test, digits = 1)

#What is the MEDIAN distance between appearances of same oracle result?
cnt = countmap([c for c in result_array])
cnt = sort(collect(cnt), by=x->x[1])

test = []
test_geo = []
y = 0
for i in 1:100
    x = findall(x->x == i, result_array)[2:end] - findall(x->x == i, result_array)[1:end-1]
    push!(test, median(x))
    push!(test_geo, geomean(x))
end
test = round.(test, digits = 1)
test_geo = round.(test_geo, digits = 1)

#Median gap between consecutive appearances of
plot(test, seriestype = :scatter, reg = true, mc = :blue, ms = 6, label = "Median",
    ylab = "Value", xlab = "Result", title = "Median and Geometric mean gaps")

#MEAN gap between consecutive appearances of
plot(test_mean, seriestype = :scatter, reg = true,
    xticks = (0:10:100), label = "Mean",
    title = "Mean and Median gaps btwn consecutive occurrences", ms = 6,
    xlab = "Result", ylab = "Mean gap", lc = :red, lw = 2)
plot!(test, seriestype = :scatter, reg = true, mc = :green, ms = 6, lw = 2, lc = :red, label = "Median",
    ylab = "Value", xlab = "Result")

#What is the MEAN distance between appearances of same oracle result?
cnt = countmap([c for c in random])
cnt = sort(collect(cnt), by=x->x[1])

test = []
y = 0
for i in 1:100
    x = findall(x->x == i, random)[2:end] - findall(x->x == i, random)[1:end-1]
    push!(test, mean(x))
end
test_mean = round.(test, digits = 1)

#What is the MEDIAN distance between appearances of same oracle result?
cnt = countmap([c for c in random])
cnt = sort(collect(cnt), by=x->x[1])

test = []
test_geo = []
y = 0
for i in 1:100
    x = findall(x->x == i, random)[2:end] - findall(x->x == i, random)[1:end-1]
    push!(test, median(x))
    push!(test_geo, geomean(x))
end
test = round.(test, digits = 1)
test_geo = round.(test_geo, digits = 1)

#Median gap between consecutive appearances of
plot(test, seriestype = :scatter, reg = true, mc = :blue, ms = 6, label = "Median",
    ylab = "Value", xlab = "Result", title = "Median and Geometric mean gaps")

#MEAN gap between consecutive appearances of
plot(test_mean, seriestype = :scatter, reg = true,
    xticks = (0:10:100), label = "Mean",
    title = "Mean and Median gaps btwn consecutive occurrences", ms = 6, mc = :pink,
    xlab = "Result", ylab = "Mean gap", lc = :red, lw = 2)
plot!(test, seriestype = :scatter, reg = true, mc = :purple, ms = 6, lw = 2, lc = :red, label = "Median",
    ylab = "Value", xlab = "Result")

#What is the MEDIAN distance between appearances of same oracle result?
cnt = countmap([c for c in result_array])
cnt = sort(collect(cnt), by=x->x[1])

test = []
test_geo = []
y = 0
for i in 1:100
    x = findall(x->x == i, result_array)[2:end] - findall(x->x == i, result_array)[1:end-1]
    push!(test, median(x))
    push!(test_geo, geomean(x))
end
test = round.(test, digits = 1)
test_geo = round.(test_geo, digits = 1)

#Median gap between consecutive appearances of
plot(test, seriestype = :scatter, reg = true, mc = :blue, ms = 6, label = "Median",
    ylab = "Value", xlab = "Result", title = "Median and Geometric mean gaps")
plot!(test_geo, seriestype = :scatter, reg = true, mc = :green, ms = 6, label = "Geometric mean")

x = findall(x->x == 1, result_array)[2:end] - findall(x->x == 1, result_array)[1:end-1]
c = countmap([c for c in x])
sort(collect(c), by=x->x[1])

x = collect(keys(sort(collect(c), by=x->x[1])))
y = collect(values(sort(collect(c), by=x->x[1])))

x = getindex.(x, 1)
y = getindex.(y, 2)

bar(x,y)

#Simulate using Provable results (result_array)

super_array = []
for r in 1:9
    stack_array = []
    print(r, ": ")
        bet_amt = 0.2
        tx_cost = 0.000025
        stack = 0 #initial value of ETH stack
        X = 1 #set 'bet under' amount (setting X as 1 == 'bet under' 2)
        for i in collect(1:3:Int64(length(result_array)))
            cnt = 0
            bet = 0
            if result_array[i] == X #prompt the iterative betting
                bet += bet_amt + tx_cost #make bet
                stack -= bet #subtract bet from stack
                i += 1 #move ahead one result in result_array

                #iterate over next results, for run of r, betting each time
                while result_array[i] !== X && cnt <= r && i !== length(result_array)
                    #print(bet)
                    if bet > 0.2
                        bet = bet*1.5 - tx_cost #update bet — TESTING FOR INCREMENTAL STRATEGY
                        if bet < 0.2 #this loop is used if you want t have bet amount decrease iteratively
                            bet = 0.2
                        end
                    end
                    stack -= bet
                    cnt += 1
                    i += 1
                    push!(stack_array, stack)
                end
                cnt = 0
                if result_array[i] == X
                    stack += bet * (99 - X)/X + bet_amt
                end
                push!(stack_array, stack)
            end
        end
    push!(super_array, stack_array)
    println(round(stack, digits=0))
end

nums = ["1" "2" "3" "4" "5" "6" "7" "8" "9"]
plot(super_array, lw = 2, label = nums,
    ylab = "ETH won", xlab = "Number of bets", title = "Strategy performance")

nums = ["1" "2" "3" "4" "5" "6" "7" "8" "9"]
plot(super_array, lw = 2, label = nums,
    ylab = "ETH won", xlab = "Number of bets", title = "Strategy performance")


##### Generic methods #####

"""
    size(d::MatrixDistribution)

Return the size of each sample from distribution `d`.
"""
size(d::MatrixDistribution)

"""
    length(d::MatrixDistribution)

The length (*i.e* number of elements) of each sample from the distribution `d`.
"""
Base.length(d::MatrixDistribution)

"""
    mean(d::MatrixDistribution)

Return the mean matrix of `d`.
"""
mean(d::MatrixDistribution)

# sampling

rand!(d::MatrixDistribution, A::AbstractArray{M}) where {M<:Matrix} = _rand!(sampler(d), A)

"""
    rand(d::MatrixDistribution, n)

Draw a sample matrix from the distribution `d`.
"""
rand(d::MatrixDistribution, n::Int) = _rand!(sampler(d), Vector{Matrix{eltype(d)}}(n))

# pdf & logpdf

_pdf(d::MatrixDistribution, x::AbstractMatrix{T}) where {T<:Real} = exp(_logpdf(d, x))

"""
    logpdf(d::MatrixDistribution, AbstractMatrix)

Compute the logarithm of the probability density at the input matrix `x`.
"""
function logpdf(d::MatrixDistribution, x::AbstractMatrix{T}) where T<:Real
    size(x) == size(d) ||
        throw(DimensionMismatch("Inconsistent array dimensions."))
    _logpdf(d, x)
end

"""
    pdf(d::MatrixDistribution, x::AbstractArray)

Compute the probability density at the input matrix `x`.
"""
function pdf(d::MatrixDistribution, x::AbstractMatrix{T}) where T<:Real
    size(x) == size(d) ||
        throw(DimensionMismatch("Inconsistent array dimensions."))
    _pdf(d, x)
end

function _logpdf!(r::AbstractArray, d::MatrixDistribution, X::AbstractArray{M}) where M<:Matrix
    for i = 1:length(X)
        r[i] = logpdf(d, X[i])
    end
    return r
end

function _pdf!(r::AbstractArray, d::MatrixDistribution, X::AbstractArray{M}) where M<:Matrix
    for i = 1:length(X)
        r[i] = pdf(d, X[i])
    end
    return r
end

function logpdf!(r::AbstractArray, d::MatrixDistribution, X::AbstractArray{M}) where M<:Matrix
    length(X) == length(r) ||
        throw(DimensionMismatch("Inconsistent array dimensions."))
    _logpdf!(r, d, X)
end

function pdf!(r::AbstractArray, d::MatrixDistribution, X::AbstractArray{M}) where M<:Matrix
    length(X) == length(r) ||
        throw(DimensionMismatch("Inconsistent array dimensions."))
    _pdf!(r, d, X)
end

function logpdf(d::MatrixDistribution, X::AbstractArray{M}) where M<:Matrix
    T = promote_type(partype(d), eltype(M))
    _logpdf!(Array{T}(size(X)), d, X)
end

function pdf(d::MatrixDistribution, X::AbstractArray{M}) where M<:Matrix
    T = promote_type(partype(d), eltype(M))
    _pdf!(Array{T}(size(X)), d, X)
end

"""
    _logpdf(d::MatrixDistribution, x::AbstractArray)

Evaluate logarithm of pdf value for a given sample `x`. This function need not perform dimension checking.
"""
_logpdf(d::MatrixDistribution, x::AbstractArray)

##### Specific distributions #####

for fname in ["wishart.jl", "inversewishart.jl"]
    include(joinpath("matrix", fname))
end


module SudokuSolver

export solvesudoku

include("heap.jl")
include("sudoku.jl")
include("puzzleexamples.jl")
end


######### Jags line program example  ###########

using Mamba, Jags

ProjDir = joinpath(dirname(@__FILE__))
cd(ProjDir) do

  line = "
  model {
    for (i in 1:n) {
          mu[i] <- alpha + beta*(x[i] - x.bar);
          y[i]   ~ dnorm(mu[i],tau);
    }
    x.bar   <- mean(x[]);
    alpha    ~ dnorm(0.0,1.0E-4);
    beta     ~ dnorm(0.0,1.0E-4);
    tau      ~ dgamma(1.0E-3,1.0E-3);
    sigma   <- 1.0/sqrt(tau);
  }
  "

  monitors = Dict(
    "alpha" => true,
    "beta" => true,
    "tau" => true,
    "sigma" => true
    )

  jagsmodel = Jagsmodel(
    name="line2",
    model=line,
    monitor=monitors,
    ncommands=4, nchains=1,
    #deviance=true, dic=true, popt=true,
    pdir=ProjDir
    );

  println("\nJagsmodel that will be used:")
  jagsmodel |> display

  data = Dict{String, Any}(
    "x" => [1, 2, 3, 4, 5],
    "y" => [1, 3, 3, 3, 5],
    "n" => 5
  )

  inits = [
    Dict("alpha" => 0,"beta" => 0,"tau" => 1,".RNG.name" => "base::Wichmann-Hill"),
    Dict("alpha" => 1,"beta" => 2,"tau" => 1,".RNG.name" => "base::Marsaglia-Multicarry"),
    Dict("alpha" => 3,"beta" => 3,"tau" => 2,".RNG.name" => "base::Super-Duper"),
    Dict("alpha" => 5,"beta" => 2,"tau" => 5,".RNG.name" => "base::Mersenne-Twister")
  ]

  println("Input observed data dictionary:")
  data |> display
  println("\nInput initial values dictionary:")
  inits |> display
  println()

  sim = jags(jagsmodel, data, inits, ProjDir)
  describe(sim)
  println()

  ## Brooks, Gelman and Rubin Convergence Diagnostic
  try
    gelmandiag(sim1, mpsrf=true, transform=true) |> display
  catch e
    #println(e)
    gelmandiag(sim, mpsrf=false, transform=true) |> display
  end

  ## Geweke Convergence Diagnostic
  gewekediag(sim) |> display

  ## Highest Posterior Density Intervals
  hpd(sim) |> display

  ## Cross-Correlations
  cor(sim) |> display

  ## Lag-Autocorrelations
  autocor(sim) |> display

  ## Deviance Information Criterion
  #dic(sim) |> display

  ## Plotting
  p = plot(sim, [:trace, :mean, :density, :autocor], legend=true);
  draw(p, nrow=4, ncol=4, filename="$(jagsmodel.name)-summaryplot", fmt=:svg)
  draw(p, nrow=4, ncol=4, filename="$(jagsmodel.name)-summaryplot", fmt=:pdf)

  # Below will only work on OSX, please adjust for your environment.
  # JULIA_SVG_BROWSER is set from environment variable JULIA_SVG_BROWSER
  @static is_apple() ? if isdefined(Main, :JULIA_SVG_BROWSER) && length(JULIA_SVG_BROWSER) > 0
          for i in 1:4
            isfile("$(jagsmodel.name)-summaryplot-$(i).svg") &&
              run(`open -a $(JULIA_SVG_BROWSER) "$(jagsmodel.name)-summaryplot-$(i).svg"`)
          end
        end : println()

end #cd


# Utility functions for field tracing.

export select_seeds

"""
	 select_seeds(x, y; nsegment=(5, 5))

Generate uniform seeding points in the grid range `x` and `y` in `nsegment`. If `nsegment`
specified, use the keyword input, otherwise it will be overloaded by the 3D version seed
generation function!
"""
function select_seeds(x, y; nsegment=(5, 5))
   xmin, xmax = extrema(x)
   ymin, ymax = extrema(y)

   dx = (xmax - xmin) / nsegment[1]
   dy = (ymax - ymin) / nsegment[2]

   xrange = xmin .+ dx*((1:nsegment[1]) .- 0.5)
   yrange = ymin .+ dy*((1:nsegment[2]) .- 0.5)

   seeds = zeros(eltype(x[1]), 2, prod(nsegment))
   for i in 0:prod(nsegment)-1
      seeds[1,i+1] = xrange[i % nsegment[1] + 1]
      seeds[2,i+1] = yrange[i ÷ nsegment[2] + 1]
   end
   seeds
end

"""
	 select_seeds(x, y, z; nsegment=(5, 5, 5))

Generate uniform seeding points in the grid range `x`, `y` and `z` in `nsegment`.
"""
function select_seeds(x, y, z; nsegment=(5, 5, 5))
   xmin, xmax = extrema(x)
   ymin, ymax = extrema(y)
   zmin, zmax = extrema(z)

   dx = (xmax - xmin) / nsegment[1]
   dy = (ymax - ymin) / nsegment[2]
   dz = (zmax - zmin) / nsegment[3]

   xrange = xmin .+ dx*((1:nsegment[1]) .- 0.5)
   yrange = ymin .+ dy*((1:nsegment[2]) .- 0.5)
   zrange = zmin .+ dz*((1:nsegment[3]) .- 0.5)

   seeds = zeros(eltype(x[1]), 3, prod(nsegment))
   for k in 1:nsegment[3], j in 1:nsegment[2], i in 1:nsegment[1]
      seeds[1,i+(j-1)*nsegment[2]+(k-1)*nsegment[2]*nsegment[3]] = xrange[i]
      seeds[2,i+(j-1)*nsegment[2]+(k-1)*nsegment[2]*nsegment[3]] = yrange[j]
      seeds[3,i+(j-1)*nsegment[2]+(k-1)*nsegment[2]*nsegment[3]] = zrange[k]
   end
   seeds
end


doc"""
`radpoisson(r,rho,l)`

solves the poisson equation for a radial density `rho` in
the `l`-angular momentum channel

computes the hartree potential

$v_{h,lm}(r) = \int_0^r dr_<\, r_<^2 g_l(r_<,r) \rho_{lm}(r_<) +
	\int_r^\infty dr_>\, r_>^2 g_l(r,r_>) \rho_{lm}(r_>)$,

where $g_l(r_<,r_>) = r_<^l/r_>^{l+1}$

uses a 3rd order adams-moulton multistep integrator to compute
each integral

`r` (LogarithmicGrid) - logarithmic grid (`loggrid`)

`rho` (AbstractVector) - radial density $\rho_{lm}(r)$ evaluated on `r`

`l` (Int) - l-quantum number

e.g.
```julia
l = 0
r = loggrid(-13,5,100)
rho = exp(-2r)
vh = radpoisson(r,rho,l)
q = multipolemoment(r,rho,l)
plot(r,vh)
plot!(r,q./r.^(l+1))  # asymptotic form of hartree potential
xlims!(0,10)
ylims!(0,0.3)
```

[`V. Blum et al., CPC 180 (2009)` p. 2185-2186]
"""
function radpoisson(r::LogarithmicGrid, rho::AbstractVector, l::Int)
    a = similar(r)
    b = similar(r)
    poisson_multistep_a(a,r,rho,r.dx,l)
    poisson_multistep_b(b,reverse(r),reverse(rho),-r.dx,l)
    a + reverse(b)
end

@adamsmoulton3_multistep 	poisson_multistep_a 	-(l+1)*y[j] + r[j]^2*rho[j]
@adamsmoulton3_multistep 	poisson_multistep_b 	     l*y[j] - r[j]^2*rho[j]

doc"""
`multipolemoment(r,rho,l)`

computes the multipole moment of radial density `rho` in
the `l`-angular momentum channel

$q = \int dr\, r^{l+2} \rho_{lm}(r)$.

`r` (LogarithmicGrid) - logarithmic grid (`loggrid`)

`rho` (AbstractVector) - radial density $\rho_{lm}(r)$ evaluated on `r`

`l` (Int) - l-quantum number

"""
multipolemoment(r,rho,l) = integrate(r,r.^(l+2).*rho)

# @fastmath function poisson_multistep_rule(f_jm1::Float64, f_j::Float64, f_jp1::Float64, y_jm1::Float64, y_j::Float64, h::Float64, c::Float64)
#     (h*h*(f_jp1 + 10f_j + f_jm1 + c*(10y_j + y_jm1)) + 24y_j - 12y_jm1)/(12 - c*h*h)
# end

# @fastmath @inbounds function poisson_multistep_integrator!(y::Vector{Float64}, f::Vector{Float64}, h::Float64,
#                                                      y1::Float64, y2::Float64, c::Float64)
#     y[1] = y1
#     y[2] = y2
#     for j in 2:length(f)-1
#         y[j+1] = poisson_multistep_rule(f[j-1], f[j], f[j+1], y[j-1], y[j], h, c)
#     end
# end

# @fastmath @inbounds function poisson_multistep_integrator(f::Vector{Float64}, h::Float64, y1::Float64, y2::Float64, c::Float64)
#     y = similar(f)
#     poisson_multistep_integrator!(y,f,h,y1,y2,c)
#     y
# end

# @multistep poisson_multistep_integrator y[j+1] = (h*h*(f[j+1] + 10f[j] + f[j-1] + c*(10y[j] + y[j-1])) + 24y[j] - 12y[j-1])/(12 - c*h*h)

# @multistep numerov_multistep y[j+1] = ( (2-5h*h*g[j]/6)*y[j] - (1+h*h*g[j-1]/12)*y[j-1] + h*h/12*(s[j+1]+10s[j]+s[j-1]) ) / (1+h*h*g[j+1]/12)

# function radpoisson(r::LogarithmicGrid, rho::AbstractVector)
#     const c = 0.25
#     N = integrate(r, 4pi*r.^2.*rho)
#     f = -4pi.*r.^2.5.*rho
#     U = similar(r)
#     poisson_multistep_integrator!(U, 0., 1e-16, f, c, r.dx)
#     U .*= sqrt(r)
#     R = r[end]
#     U += (N - U[end])/R.*r
#     vh = U./r
#     vh/4pi
# end

# @fastmath function multistep_rule1(f_jm1::Float64, f_j::Float64, f_jp1::Float64, y_jm1::Float64, y_j::Float64, h::Float64)
#     2y_j - y_jm1 + h*h*(f_jp1 + 10f_j + f_jm1)/12
# end

# @fastmath @inbounds function multistep_integrator1!(y::Vector{Float64}, f::Vector{Float64},
#                                                       h::Float64, y1::Float64, y2::Float64)
#     y[1] = y1
#     y[2] = y2
#     for j in 2:length(f)-1
#         y[j+1] = multistep_rule1(f[j-1], f[j], f[j+1], y[j-1], y[j], h)
#     end
# end

# @fastmath @inbounds function multistep_integrator1(f::Vector{Float64}, h::Float64, y1::Float64, y2::Float64)
#     y = similar(f)
#     multistep_integrator1!(y,f,h,y1,y2)
#     y
# end

# uniform radial grid
# function poissonrad(rho,r)
#     N = integrate(r, 4pi*r.^2.*rho)
#     dr = r[1] - r[0]
#     f = -4pi.*r.*rho
#     U = similar(r)
#     multistep_integrator1!(U, f, dr, 0., 1e-16)
#     R = r[end]
#     U += (N - U[end])/R.*r
#     vh = U./r
#     vh
# end

# function hartreeatomicrad(r::Grid, rho::AbstractArray, l::Int)
#     cumsum(r.w.*r.^(2+l).*rho)./r.^(l+1) + r.^l.*flipdim(cumsum(flipdim((r.w.*r.^(2-l-1).*rho),1)),1)
# end