nums = parse.(Int, readlines("day10/input.txt"))

push!(nums, 0)
push!(nums, 3+maximum(nums))
sort!(nums)

println(nums)

dp = Dict()
function dp_proxy(i)
    if i in keys(dp)
        return dp[i]
    else
        result = connect_suffix(i)
        dp[i] = result
        return result
    end
end

function connect_suffix(i)
    global nums
    num_ways = 0
    if i == length(nums)
        return 1
    end
    if i + 1 <= length(nums) && (nums[i]+1 <= nums[i + 1] <= nums[i]+3)
        num_ways += dp_proxy(i+1)
    end
    if i + 2 <= length(nums) && (nums[i]+1 <= nums[i + 2] <= nums[i]+3)
        num_ways += dp_proxy(i+2)
    end
    if i + 3 <= length(nums) && (nums[i]+1 <= nums[i + 3] <= nums[i]+3)
        num_ways += dp_proxy(i+3)
    end
    return num_ways
end

println(connect_suffix(1))

# holy hell DP! that was exciting!


#=
AirFoilCoordinatesArraysActivity.jl
Jacob Child
April 23rd, 2022

Equations Used
Zu(Zupper) = zbar + t/2
Zl(Zlower) = zbar -t/2
t = 10m(.2969sqrt(x) - .1620x - .3537x^2 + .2843x^3 - .1015x^4)
if X<=p
        zbar = c*(2*p*x-x^2)/(p^2)
    elseif x > p
        zbar = c*(1-2*p+2*p*x-x^2)/((1-p)^2)
    end

Pseudocode- Take in NACA 4- Series Airfoil numbers- ie max camber (c),
    max camber position (p), thickness (t), and an array (x) to plot at.
    For example NACA 2412 = c: .02, p: .04, t: .12 (all in meters I believe)
    Do that by asking the user for c,p,t, then creating x (hardcode it for simplicity)
    Use the functions previously written to calculate zu and zl.

    Goal = clean code, lots of well commented functions
=#

"""
user_input Function to take in the 4 NACA numbers
"""
function user_input()
    println("Input the 4 NACA numbers ie 2412 \nNACA ")
    input1 = readline()
    println("\nDo you want Cosine Spacing? 1 = Yes, 0 = No\n")
    input2 = readline()
    if input2 == "1"
        println("\nHow many points do you want?\n")
        stepsize = pi / parse(Int64,readline() )
        #println("\n", stepsize)
        theta = collect(0.0:stepsize:pi)    #creates numbers from 0 to pi
        x = (1/2) * (1 .- cos.(theta))  #should give proper Spacing, assuming c (chord length) = 1

    else
        println("\nNot using Cosine Spacing")
        x = collect(0:0.1:1)   #10 points

    end

    c = parse(Int64, input1[1])
    c = c/100  #converts the chord to the proper units   #can do c//100 and that leaves it in fraction form
    p = parse(Int64, input1[2])
    p = p/10    #converts the max camber position to the proper units
    t = parse(Int64, input1[3:4])
    t = t/100
    #println("\n c:", c, "\n p:", p, "\n t:", t)    #prints out the actual NACA values

    return c, p, t, input1, x
end

"""
thickness_calculations(x,p)
Takes in the x coordinates, and max thickness (p), calculates thickness, and ouputs the array

"""
function thickness_calculations(x,t)
    m = t
    TArray = 10*m*(.2969*sqrt(x) - .1260*x - .3537*x^2 + .2843*x^3 - .1015*x^4)
    return TArray
end

"""
zbar_calculations(c,p,x)
Takes in c, p, and x values and outputs zbar
"""
function zbar_calculations(c,p,x)
    zbar = Float64[] #makes zbar an array of zeros the length of x, hopefully?
    for i in x
        if i <= p
             zbarInput = c*(2*p*i - i^2)/(p^2)
             push!(zbar,zbarInput)

        elseif i > p
             zbarInput = c*(1 - 2*p + 2*p*i - i^2)/((1-p)^2)
             push!(zbar,zbarInput)

        else
             zbar = "you wenti"
             break

        end

    end
    return zbar
end

"""
zupper_lower(zbar, TArray)
calculates zu and zl and returns both
"""
function zupper_lower(zbar, TArray)

    zu = zbar .+ TArray/2
    zl = zbar .- TArray/2

    return zu, zl
end

#Actual code- See how it goes

c, p, t, fromuser, x = user_input()  #function to get user input
#why is collect better than x = range(0:0.010:1)?
TArray = thickness_calculations.(x, t)   #calculates the thicknesses
zbarn = zbar_calculations(c, p, x) #calculate and returns zbar
zu, zl = zupper_lower(zbarn, TArray) #calculates and returns zupper and z lower

#outputs

#println("\n \n", zu, "\n \n", zl) #for vs code

using DataFrames #only seems to work in the Julia REPL

DataFrame()

Data = DataFrame("X" => x, "Zupper"=> zu, "Zlower" => zl, "Thickness" => TArray, "Zbar" => zbarn)

display(Data)

#look into aspect ratio for plots
 using Plots
 plt = plot(x,zu)   #saving the plot to an element
 plot!(x,zl)    #adding/updating to that plot
 #Plot settings
 plot!(xlims = (-.1, 1.1), title = ("NACA $fromuser"), xlabel = "Position", ylabel = "Position")
 plot!(aspectratio = :equal)

 display(plt)   #displaying the plot by calling the saved element


using Revise
using LinearAlgebra, Plots, Setfield, Parameters

using PseudoArcLengthContinuation
const PALC = PseudoArcLengthContinuation

TY = Float64
AF = Array{TY}
#################################################################
if 1==1
	ENV["GKSwstype"] = "nul"    # needed for the GR backend on headless servers
	using Plots
	using SixelTerm
else
	using GR
	GR.inline("iterm")
end

using CuArrays
CuArrays.allowscalar(false)
# import LinearAlgebra: mul!, axpby!
# mul!(x::CuArray, y::CuArray, α::T) where {T <: Number} = (x .= α .* y)
# mul!(x::CuArray, y::T, α::CuArray) where {T <: Number} = (x .= α .* y)
# axpby!(a::T, X::CuArray, b::T, Y::CuArray) where {T <: Number} = (Y .= a .* X .+ b .* Y)

const TY = Float64
const AF = CuArray{TY}
#################################################################
# to simplify plotting of the solution
heatmapsol(x) = heatmap(reshape(Array(x), Nx, Ny)', color=:viridis)
norminf(x) = maximum(abs.(x))

Nx = 2^8
Ny = 2^8
lx = 8pi * 2
ly = 2*2pi/sqrt(3) * 2

X = -lx .+ 2lx/Nx * collect(0:Nx-1)
Y = -ly .+ 2ly/Ny * collect(0:Ny-1)

sol0 = [(cos(x) .+ cos(x/2) * cos(sqrt(3) * y/2) ) for x in X, y in Y]
		sol0 .= sol0 .- minimum(vec(sol0))
		sol0 ./= maximum(vec(sol0))
		sol0 = sol0 .- 0.25
		sol0 .*= 1.7
		# heatmap(sol0, color=:viridis)

using AbstractFFTs, FFTW, KrylovKit
import Base: *, \

# Making the linear operator a subtype of PALC.LinearSolver is handy as we will use it
# in the Newton iterations.
struct SHLinearOp{Treal, Tcomp, Tl1, Tplan, Tiplan} <: PALC.AbstractLinearSolver
	tmp_real::Treal         # temporary
	tmp_complex::Tcomp      # temporary
	l1::Tl1
	fftplan::Tplan
	ifftplan::Tiplan
end

struct SHEigOp{Tsh <: SHLinearOp} <: PALC.AbstractEigenSolver
	sh::Tsh
end

function SHLinearOp(Nx, lx, Ny, ly; AF = Array{TY})
	# AF is a type, it could be CuArray{TY} to run the following on GPU
	k1 = vcat(collect(0:Nx/2), collect(Nx/2+1:Nx-1) .- Nx)
	k2 = vcat(collect(0:Ny/2), collect(Ny/2+1:Ny-1) .- Ny)
	d2 = [(1-(pi/lx * kx)^2 - (pi/ly * ky)^2)^2 + 1. for kx in k1, ky in k2]
	tmpc = Complex.(AF(zeros(Nx, Ny)))
	return SHLinearOp(AF(zeros(Nx, Ny)), tmpc, AF(d2), plan_fft!(tmpc), plan_ifft!(tmpc))
end

function *(c::SHLinearOp, u)
	c.tmp_complex .= Complex.(u)
	c.fftplan * c.tmp_complex
	c.tmp_complex .= c.l1 .* c.tmp_complex
	c.ifftplan * c.tmp_complex
	c.tmp_real .= real.(c.tmp_complex)
	return copy(c.tmp_real)
end

function \(c::SHLinearOp, u)
	c.tmp_complex .= Complex.(u)
	c.fftplan * c.tmp_complex
	c.tmp_complex .=  c.tmp_complex ./ c.l1
	c.ifftplan * c.tmp_complex
	c.tmp_real .= real.(c.tmp_complex)
	return copy(c.tmp_real)
end

function (sh::SHLinearOp)(J, rhs; shift = 0., tol =  1e-9)
	u, l, ν = J
	udiag = l .+ 1 .+ 2ν .* u .- 3 .* u.^2 .- shift
	res, info = KrylovKit.linsolve( du -> -du .+ sh \ (udiag .* du), sh \ rhs, tol = tol, maxiter = 6)
	return res, true, info.numops
end

function (sheig::SHEigOp)(J, nev::Int; σ = 0.3)
	u, l, ν = J
	sh = sheig.sh
	udiag = l .+ 1 .+ 2ν .* u .- 3 .* u.^2
	N = size(u)

	A = du -> sh(J, du; shift = σ, tol = 1e-5)[1]

	# we adapt the krylov dimension as function of the requested eigenvalue number
	vals, vec, info = KrylovKit.eigsolve(A, AF(rand(eltype(u), N)), nev, :LM, tol = 1e-10, maxiter = 40, verbosity = 2, issymmetric = true, krylovdim = max(40, nev + 10))
	@show 1 ./vals .+ σ
	return 1 ./vals .+ σ, vec, true, info.numops
end

function F_shfft(u, p)
	@unpack l, ν, L = p
	return -(L * u) .+ ((l+1) .* u .+ ν .* u.^2 .- u.^3)
end

J_shfft(u, p) = (u, p.l, p.ν)

L = SHLinearOp(Nx, lx, Ny, ly, AF = AF)
Leig = SHEigOp(L) # for eigenvalues computation
# Leig((sol_hexa, -0.1, 1.3), 20; σ = 0.5)

par = (l = -0.1, ν = 1.3, L = L)

@time F_shfft(sol0, par); # 0.008022 seconds (12 allocations: 1.500 MiB)

opt_new = PALC.NewtonPar(verbose = true, tol = 1e-6, linsolver = L, eigsolver = Leig)
	sol_hexa, hist, flag = @time PALC.newton(
				F_shfft, J_shfft,
				AF(sol0), par,
				opt_new, normN = norminf)
	println("--> norm(sol) = ", norminf(sol_hexa))

heatmapsol(sol_hexa)
####################################################################################################
# trial using IterativeSolvers

# function (sh::SHLinearOp)(J, rhs)
# 	u, l, ν = J
# 	udiag = l .+ 1 .+ 2ν .* u .- 3 .* u.^2
# 	res, info = res, info = KrylovKit.linsolve( u -> -u .+ sh \ (udiag .* u), sh \ rhs, tol = 1e-9, maxiter = 6)
# 	return res, true, info.numops
# end
####################################################################################################
deflationOp = DeflationOperator(2.0, (x, y)->dot(x, y), 1.0, [sol_hexa])

opt_new = @set opt_new.maxIter = 250
outdef, _, flag, _ = @time PALC.newton(
				F_shfft, J_shfft,
				0.4 .* sol_hexa .* AF([exp(-1(x+0lx)^2/25) for x in X, y in Y]),
				par,
				opt_new, deflationOp, normN = x-> maximum(abs.(x)))
		println("--> norm(sol) = ", norm(outdef))
		heatmapsol(outdef) |> display
		flag && push!(deflationOp, outdef)

####################################################################################################
opts_cont = ContinuationPar(dsmin = 0.001, dsmax = 0.007, ds= -0.005, pMax = 0.2, pMin = -1.0, theta = 0.5, plotEveryNsteps = 5, newtonOptions = setproperties(opt_new; tol = 1e-6, maxIter = 15), maxSteps = 88,
	computeEigenValues = true,
	detectBifurcation = 0,
	precisionStability = 1e-5,
	saveEigenvectors = false,
	nev = 11 )

	br, u1 = @time PALC.continuation(
		F_shfft, J_shfft,
		deflationOp.roots[2], par, (@lens _.l),
		opts_cont, plot = true, verbosity = 2,
		plotSolution = (x, p;kwargs...)->heatmap!(reshape(Array(x), Nx, Ny)'; color=:viridis, kwargs...),
		printSolution = (x, p) -> norm(x), normC = norminf
		)

# branches = [br]
# push!(branches, br)
# plot(branches)
#
# plot(br)


@testset "minimal_intersection_test" begin
    @test intersect(interval(1.0,3.0), interval(2.1,4.0)) == interval(2.1,3.0)
    @test intersect(interval(1.0,3.0), interval(3.0,4.0)) == interval(3.0,3.0)
    @test intersect(interval(1.0,3.0), emptyinterval()) == emptyinterval()
    @test intersect(entireinterval(), emptyinterval()) == emptyinterval()
    @test intersect(interval(1.0,3.0), entireinterval()) == interval(1.0,3.0)
end

@testset "minimal_convex_hull_test" begin
    @test hull(interval(1.0,3.0), interval(2.1,4.0)) == interval(1.0,4.0)
    @test hull(interval(1.0,1.0), interval(2.1,4.0)) == interval(1.0,4.0)
    @test hull(interval(1.0,3.0), emptyinterval()) == interval(1.0,3.0)
    @test hull(emptyinterval(), emptyinterval()) == emptyinterval()
    @test hull(interval(1.0,3.0), entireinterval()) == entireinterval()
end


export PowerMeasure
using Base: eltype

"""
    struct PowerMeasure{M,N}
        μ::M
        size::NTuple{N,Int}
    end

A power measure is a product of a measure with itself. The number of elements in
the product determines the dimensionality of the resulting support.

Note that power measures are only well-defined for integer powers.

The nth power of a measure μ can be written μ^x.
"""
struct PowerMeasure{M,N}
    μ::M
    size::NTuple{N,Int}
end

import Base

function Base.:^(μ::AbstractMeasure, n::Integer)
    PowerMeasure(μ, (n,))
end

function Base.:^(μ::AbstractMeasure, size::NTuple{N,I}) where {N, I <: Integer}
    PowerMeasure(μ, size)
end

sampletype(d::PowerMeasure{M,N}) where {M,N} = AbstractArray{sampletype(d.μ), N}

export rand!

function rand!(result::AbstractArray, d::PowerMeasure)
    @inbounds for j in eachindex(result)
        result[j] = rand(d.μ)
    end
    return result
end

function Base.rand(d::PowerMeasure)
    result = Array{sampletype(d.μ), length(d.size)}(undef, d.size...)
    return rand!(result, d)
end


# we write a function that returns the row index of the (absolutely) largest element in column k of the matrix A
function findLargestInCol(A, k)

	# we get column k
	col = A[:,k]

	# we take elementwise absolute value
	absCol = broadcast(abs, col)

	# we get the index of the largest element
	index = argmax(absCol)

	# we return index
	return index
end


# ### subplot (3)

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

a = plot(sin, 0, 2pi, xlabel="x1")
b = plot(cos, 0, 2pi, xlabel="x2")
c = histogram(randn(1000), xlabel="x3")
d = plot(x->exp(-x^2), -3, 3, xlabel="x4")

plot(a, b, c, d, layout=(4,1), size=(300, 800))


@testset "Testing normalizing functions" begin
    @testset "normtime" begin
        nt1 = rand()

        normed1 = normtime(nt1,nt1+1)
        normed2 = normtime(nt1, nt1+1, 65)

        @test length(normed1) === 100
        @test length(normed2) === 65
        @test first(normed1) === nt1
        @test first(normed2) === nt1
        @test last(normed1) < (nt1 + 1)
        @test last(normed1) < (nt1 + 1)
    end

    @testset "normalizeevents" begin
        @test prod(normalizeevents(1, 2, 1.5) .=== [ 51 ])
        @test prod(normalizeevents(1, 2, [1.25, 1.75]) .=== [26, 76])
        @test prod(normalizeevents(1, 2, [1.2, 1.8], 10) .=== [3, 9])
    end

    @testset "timenormalize" begin
        t200 = [ normtime(0, 2*pi, 200); normtime(2*pi, 2*2*pi, 200) ]
        t100 = [ normtime(0, 2*pi, 100); normtime(2*pi, 2*2*pi, 100) ]

        x200 = sin.(t200)
        x100 = sin.(t100)

        normx200 = timenormalize(x200, [1, 201, 401], 100)

        @test prod(isapprox.(x100, normx200; atol=1e-10))
    end
end


function rms{T<:Real}(s::Array{T,1})
    # TODO: add argument valdiation
    m = zero(T)
    for si in s
        m += si * si
    end
    sqrt(m/length(s))
end

function rms{T<:Real}(s::Array{Complex{T},1})
    # TODO: add argument valdiation
    m = zero(T)
    for si in s
        m += abs2(si)
    end
    sqrt(m/length(s))
end


# Determine the mesh sensitivity for the default unexposed HTC and the recovered value of
# the exposed HTC. This is ran post-factum of the recovery calculation to justify it.
# In case the correct mesh was not picked the adjoint calculation muts be repeated with
# the converged value of the mesh size.

using NeuralThermalODE

const testfile = "Fig_67.csv_Result.csv"
const test_data = read_data(testfile)

# This is either the raw HTC or a spline fit.
const htc_file = "checkpoint_htc_360.csv"
const htc_recovered = read_htc(htc_file)

const sinkT = test_data[end, 2]
const ambT = test_data[1, 2]

tspan = (Float64(test_data[1, 1]), Float64(test_data[end, 1]))

const p = Vector(htc_recovered[2])

# Pick correct value.
const t = 0.00635

# Study the mesh sensitivity.
const N1 = 10
const N2 = 20
const N3 = 40
const N4 = 80

dx = t / N1
u = [ambT for i in 1:N1]
prob = ODEProblem(heat_transfer, u, tspan, p, saveat = 1.0)
sol1 = solve(prob, Tsit5(), saveat = 1.0)

dx = t / N2
u = [ambT for i in 1:N2]
prob = ODEProblem(heat_transfer, u, tspan, p, saveat = 1.0)
sol2 = solve(prob, Tsit5(), saveat = 1.0)

dx = t / N3
u = [ambT for i in 1:N3]
prob = ODEProblem(heat_transfer, u, tspan, p, saveat = 1.0)
sol3 = solve(prob, Tsit5(), saveat = 1.0)

dx = t / N4
u = [ambT for i in 1:N4]
prob = ODEProblem(heat_transfer, u, tspan, p, saveat = 1.0)
sol4 = solve(prob, Tsit5(), saveat = 1.0)

plot(sol1.t[:], sol1[end, :], label = "N = 10", color = :black, linestyle = :dash)
plot!(sol2.t[:], sol2[end, :], label = "N = 20", color = :black, linestyle = :dot)
plot!(sol3.t[:], sol3[end, :], label = "N = 40", color = :black, linestyle = :dashdot)
plot!(sol4.t[:], sol4[end, :], label = "N = 80", color = :black)
plot!(fmt = :svg,
      grid = false,
      xlabel = "Time (min)",
      ylabel = "Temperature (°C)",
      xlims = (0, 200),
      ylims = (-200, 50),
      xticks = 0:25:200,
      yticks = -200:25:50,
      minorticks = true)

savefig("mesh_sensitivity.svg")


@testset "Matmul" begin

    # Test multiplications with Double64 and Float64
    A1 = [1.0 2 3; 4 5 6; 7 8 9]
    B1 = [7.0 8; 9 10; 11 12]
    A1B1 = A1*B1

    DA1DB1 = Double64.(A1)*Double64.(B1)
    DA1FB1 = Double64.(A1)*B1
    FA1DB1 = A1*Double64.(B1)

    @test DA1DB1 == Double64.(A1B1)
    @test DA1FB1 == Double64.(A1B1)
    @test FA1DB1 == Double64.(A1B1)
    @test DA1DB1 isa Array{Double64,2}
    @test DA1FB1 isa Array{Double64,2}
    @test FA1DB1 isa Array{Double64,2}

    # Also test transpose
    @test Double64.(B1)'Double64.(A1)' == Double64.(A1B1')
    # Incompatible Dimension
    @test_throws DimensionMismatch Double64.(B1)*Double64.(A1)

    #Test with zero inner dimension (zero result)
    A2 = Array{Float64,2}(undef, 2, 0)
    B2 = Array{Float64,2}(undef, 0, 3)
    A2B2 = A2*B2

    DA2DB2 = Double64.(A2)*Double64.(B2)
    DA2FB2 = Double64.(A2)*B2
    FA2DB2 = A2*Double64.(B2)

    @test DA2DB2 == Double64.(A2B2)
    @test DA2FB2 == Double64.(A2B2)
    @test FA2DB2 == Double64.(A2B2)
    @test DA2DB2 isa Array{Double64,2}
    @test DA2FB2 isa Array{Double64,2}
    @test FA2DB2 isa Array{Double64,2}

    #Test with zero outer dimensions (empty result)
    A3 = Array{Float64,2}(undef, 0, 3)
    B3 = Array{Float64,2}(undef, 3, 0)
    A3B3 = A3*B3

    DA3DB3 = Double64.(A3)*Double64.(B3)
    DA3FB3 = Double64.(A3)*B3
    FA3DB3 = A3*Double64.(B3)

    @test DA3DB3 == Double64.(A3B3)
    @test DA3FB3 == Double64.(A3B3)
    @test FA3DB3 == Double64.(A3B3)
    @test DA3DB3 isa Array{Double64,2}
    @test DA3FB3 isa Array{Double64,2}
    @test FA3DB3 isa Array{Double64,2}

    #Test multiplications with ComplexD64
    C1  = [1.0+2*im 3 4*im; 5 6 7-im]
    C1T = transpose(C1)
    C2  = [im; 0; 5-7*im]
    C1C2 = C1*C2
    C1C1T = C1*C1T
    DC1DC2  = ComplexDF64.(C1)*ComplexDF64.(C2)
    DC1DC1T = ComplexDF64.(C1)*ComplexDF64.(C1T)
    DC1FC2  = ComplexDF64.(C1)*C2
    FC1DC2  = C1*ComplexDF64.(C2)

    @test DC1DC2  == ComplexDF64.(C1C2)
    @test DC1DC1T == ComplexDF64.(C1C1T)
    @test DC1FC2  == ComplexDF64.(C1C2)
    @test FC1DC2  == ComplexDF64.(C1C2)
    # ComplexD64 is not a type, it is a constructor function
    @test DC1DC2  isa Array{Complex{DoubleFloat{Float64}},1}
    @test DC1DC1T isa Array{Complex{DoubleFloat{Float64}},2}
    @test DC1FC2  isa Array{Complex{DoubleFloat{Float64}},1}
    @test FC1DC2  isa Array{Complex{DoubleFloat{Float64}},1}

end


# export tutte

"""
`tutte(G::SimpleGraph, outside)` gives `G` a Tutte embedding in which
the list of vertices in `outside` form the outer boundary. The graph
should be connected (ideally, 3-connected) and, if planar and `outside`
defines a face, the embedding will be crossing free.

`tutte(G::SimpleGraph)` assumes `G` has a rotation system in which case a
largest face will be selected to be `outside`.
"""
function _tutte(G::SimpleGraph{T}, outside::Vector{T}) where {T}
    if !issubset(outside, G.V)
        error("Some of the proposed outer vertices are not in this graph")
    end
    n = NV(G)
    VV = vlist(G)

    outside = unique(outside)  # remove dups

    lookup(v) = findfirst([v == w for w in VV])

    A = zeros(n, n)
    r = sqrt(n)

    rhs_x = zeros(n)
    rhs_y = zeros(n)

    no = length(outside)
    for j = 1:no
        theta = 2 * pi * (j - 1) / no + pi / 2
        v = outside[j]
        idx = lookup(v)
        rhs_x[idx] = r * cos(theta)
        rhs_y[idx] = r * sin(theta)
    end

    for v in G.V
        idx = lookup(v)
        A[idx, idx] = 1

        if !in(v, outside)
            d = deg(G, v)
            for w in G[v]
                j = lookup(w)
                A[idx, j] = -1 / d
            end
        end
    end

    x = A \ rhs_x
    y = A \ rhs_y

    d = Dict{T,Vector{Float64}}()
    for v in VV
        idx = lookup(v)
        d[v] = [x[idx]; y[idx]]
    end
    embed(G, d)
end

function _tutte(G::SimpleGraph)
    FF = faces(G)
    Fmax = first(FF)
    for F in FF
        if length(F) > length(Fmax)
            Fmax = F
        end
    end
    _tutte(G, first.(Fmax))
end


struct AbstractSimplex{T} end

export Simplex
struct Simplex{T}  <: AbstractVector{T}
    weights::Vector{T}
end

function Simplex(w)
    N = length(w)
    T = eltype(w)
    Simplex{T}(w)
end

Base.length(s::Simplex{T}) where {T} = length(s.weights)
Base.size(s::Simplex{T}) where {T} = (length(s.weights),)
Base.getindex(s::Simplex{T}, n) where {T} = s.weights[n]

struct Log{T}
    val :: T
end

function Base.:*(x::Log{Tx}, y::Log{Ty}) where {Tx <: Real, Ty <: Real}
    Log(x.val + y.val)
end

Base.log(x::Log{T}) where {T <: Real} = x.val

struct AbstractMix end

# struct MixFor
#     dists :: For
#     weights :: Simplex
# end

struct MixVec{D,W}
    dists :: Vector{D}
    weights :: Simplex{W}
end

export Mix

function Mix(dists :: Vector{D}, weights :: Vector{W}) where {D,W}
    MixVec{D,W}(dists, Simplex(weights))
end

# Mix(w::Vector) = dists -> Mix(dists, log.(w))

function Base.rand(mix::MixVec)
    # This is the "Gumbel max trick"  for categorical sampling
    (j_max,lw_max) = (0,-Inf)

    for j in eachindex(mix.weights)
        lw_gumbel = log(mix.weights[j]) + rand(Gumbel())
        if lw_gumbel > lw_max
            (j_max,lw_max) = (j, lw_gumbel)
        end
    end

    @inbounds mix.dists[j_max] |> rand
end

# function Base.rand(mix::Mix, N::Int)
#     x1 = rand(mix)
#     x = Vector{typeof(x1)}(undef, N)
#     @inbounds x[1] = x1
#     @inbounds for n in 2:N
#         x[n] = rand(mix)
#     end
#     x
# end

xform(mix::MixVec) = xform(mix.dists[1])

function Dists.logpdf(mix::MixVec, x)
    ℓ = 0.0
    @simd for j in eachindex(mix.weights)
        @inbounds ℓ += log(mix.weights[j]) + logdensity(mix.dists[j], x)
    end
    ℓ
end


module Tests
using StandardizedMatrices, StatsBase, Statistics, LinearAlgebra, Test

n, p, k = 100, 5, 7
x = randn(n, p)
x2 = zscore(x, 1)
z = StandardizedMatrix(x)
b = randn(p)
b2 = randn(n)
b3 = randn(p, k)
b4 = randn(n, k)

@testset "Multiplication with Vector" begin
	@test isapprox(z * b, x2 * b)
	storage = zeros(n)
	mul!(storage, z, b)
	@test isapprox(storage, x2 * b)
	storage = zeros(p)
	mul!(storage, z', b2)
	@test isapprox(storage, x2' * b2)
end

@testset "Multiplication with Matrix" begin
	@test isapprox(z * b3, x2 * b3)
	storage = zeros(n, k)
	mul!(storage, z, b3)
	@test isapprox(storage, x2 * b3)
	storage = zeros(p, k)
	mul!(storage, z', b4)
	@test isapprox(storage, x2' * b4)
end

@testset "Indexing" begin
	for i in eachindex(z)
		@test z[i] == x2[i]
	end
	for j in 1:size(z, 2), i in 1:size(z, 1)
		@test z[i, j] == x2[i, j]
	end
end

end #module


# # Unitary matrix operations without allocation
# A unitary matrix features uniform eigenvalues and reversibility. It is widely used as an approach to ease the gradient exploding and vanishing problem and the memory wall problem.
# One of the simplest ways to parametrize a unitary matrix is representing a unitary matrix as a product of two-level unitary operations. A real unitary matrix of size $N$ can be parametrized compactly by $N(N-1)/2$ rotation operations
#
# ```math
#    {\rm ROT}(a!, b!, \theta)  = \left(\begin{matrix}
#        \cos(\theta) & - \sin(\theta)\\
#        \sin(\theta)  & \cos(\theta)
#    \end{matrix}\right)
#    \left(\begin{matrix}
#        a!\\
#        b!
#    \end{matrix}\right),
# ```
#
# where $\theta$ is the rotation angle, `a!` and `b!` are target registers.

using NiLang, NiLang.AD

@i function umm!(x!, θ)
    @safe @assert length(θ) ==
            length(x!)*(length(x!)-1)/2
    k ← 0
    for j=1:length(x!)
        for i=length(x!)-1:-1:j
            k += 1
            ROT(x![i], x![i+1], θ[k])
        end
    end

    k → length(θ)
end

# Here, the ancilla `k` is deallocated manually by specifying its value, because we know the loop size is $N(N-1)/2$.
# We define the test functions in order to check gradients.

@i function isum(out!, x::AbstractArray)
    for i=1:length(x)
        out! += x[i]
    end
end

@i function test!(out!, x!::Vector, θ::Vector)
   umm!(x!, θ)
   isum(out!, x!)
end

# Let's print the program output

out, x, θ = 0.0, randn(4), randn(6);
@instr Grad(test!)(Val(1), out, x, θ)
x

# We can erease the gradient field by uncomputing the gradient function.
# If you want, you can differentiate it twice to obtain Hessians.
# However, we suggest using ForwardDifferentiation over our NiLang program, this is more efficient.

@instr (~Grad(test!))(Val(1), out, x, θ)
x

# In the above testing code, `Grad(test)` attaches a gradient field to each element of `x`. `~Grad(test)` is the inverse program that erase the gradient fields.
# Notably, this reversible implementation costs zero memory allocation, although it changes the target variables inplace.


module Test_pslifting

using PeriodicSystems
using DescriptorSystems
using Symbolics
using Test
using LinearAlgebra
using ApproxFun

println("Test_liftings")

# # selected sequence to illustrate failure of BlockMatrices.jl v0.16.16
# using LinearAlgebra
# using ApproxFun
# Af = Fun(t -> [0 1; -10*cos(t) -24-10*sin(t)],Fourier(0..2π));
# D = Derivative(domain(Af));
# ND = [D 0I; 0I D];
# Aop = Af - ND;
# NA = 322
# RW = Aop[1:NA,1:NA]

@testset "test_liftings" begin

# using Floquet based approach
At = PeriodicFunctionMatrix(t -> [0 1; -10*cos(t) -24-10*sin(t)],2pi);
@time ev = psceig(At; reltol = 1.e-10)
@time ev = psceig(At,500; reltol = 1.e-10)
@test sort(ev) ≈ sort([0;-24])

# using Fourier series
Afun = FourierFunctionMatrix(Fun(t -> [0 1; -10*cos(t) -24-10*sin(t)],Fourier(0..2π)))
ev1 = psceig(Afun,50)
@test isapprox(sort(real(ev)),sort(real(ev1)),rtol=1.e-6) && norm(imag(ev1)) < 1.e-10

ev3 = psceig(Afun,60)
@test isapprox(sort(real(ev)),sort(real(ev3)),rtol=1.e-6) && norm(imag(ev3)) < 1.e-10

# using Toeplitz operator truncation
Ahr = convert(HarmonicArray,PeriodicFunctionMatrix(t -> [0 1; -10*cos(t) -24-10*sin(t)],2pi));
ev2 = psceig(Ahr,50)
@test sort(ev) ≈ sort(real(ev2[sortperm(imag(ev2),by=abs)][1:2]))

# example Zhou, Hagiwara SCL 2002 period pi/2 and pi
# using Floquet based approach
At = PeriodicFunctionMatrix(t -> [-1-sin(2*t)^2 2-0.5*sin(4*t); -2-0.5*sin(4*t) -1-cos(2*t)^2],pi/2);
At2 = PeriodicFunctionMatrix(t -> [-1-sin(2*t)^2 2-0.5*sin(4*t); -2-0.5*sin(4*t) -1-cos(2*t)^2],pi);
@time ev = psceig(At; reltol = 1.e-10)
@time ev = psceig(At,500; reltol = 1.e-10) # complex characteristic exponents
@test sort(real(ev)) ≈ sort([-1;-2]) && sort(imag(ev)) ≈ [2;2]
@time ev2 = psceig(At2; reltol = 1.e-10)
@time ev2 = psceig(At2,500; reltol = 1.e-10) # real characteristic exponents
@test sort(ev2) ≈ sort([-1;-2])

# using Fourier series
Afun = FourierFunctionMatrix(Fun(t -> [-1-sin(2*t)^2 2-0.5*sin(4*t); -2-0.5*sin(4*t) -1-cos(2*t)^2],Fourier(0..π/2)))
ev3 = psceig(Afun)
@test sort(real(ev3)) ≈ sort([-1;-2]) && sort(imag(ev3)) ≈ [2;2]

# using Toeplitz operator truncation period pi/2
Ahr = convert(HarmonicArray,PeriodicFunctionMatrix(t -> [-1-sin(2*t)^2 2-0.5*sin(4*t); -2-0.5*sin(4*t) -1-cos(2*t)^2],pi/2));
Asym = convert(PeriodicSymbolicMatrix,Ahr); Asym.F
ev3 = psceig(Ahr,5)
@test sort(real(ev3)) ≈ sort([-1;-2]) && sort(imag(ev3)) ≈ [2;2]

# simple period
ev4 = psceig(Ahr,50)
@test sort(real(ev4)) ≈ sort([-1;-2]) && sort(imag(ev4)) ≈ [2;2]

# double period
ev5 = psceig(Ahr,50; P = 2)
@test sort(real(ev5)) ≈ sort([-1;-2]) && norm(imag(ev5)) < 1.e-10

# using Fourier series truncation period pi
ev4 = psceig(Afun,P = 2)
@test sort(real(ev4)) ≈ sort([-1;-2]) && norm(imag(ev4)) < 1.e-10

# using Toeplitz operator truncation period pi
Ahr2 = convert(HarmonicArray,PeriodicFunctionMatrix(t -> [-1-sin(2*t)^2 2-0.5*sin(4*t); -2-0.5*sin(4*t) -1-cos(2*t)^2],pi));
Asym2 = convert(PeriodicSymbolicMatrix,Ahr2); Asym2.F
ev3 = psceig(Ahr2,5)
@test sort(real(ev3)) ≈ sort([-1;-2]) && norm(imag(ev3)) < 1.e-10

ev4 = psceig(Ahr,50; P = 2)
@test sort(real(ev4)) ≈ sort([-1;-2]) && norm(imag(ev4)) < 1.e-10

# constant matrix case
Q = convert(HarmonicArray,PeriodicFunctionMatrix([-1 0; 0 -2],pi/2));
N = 5; ev1 = psceig(Q,N)
@test sort(real(ev1)) ≈ sort([-1;-2]) && norm(imag(ev1)) < 1.e-10

QF = convert(FourierFunctionMatrix,Q)
N = 5; ev2 = psceig(QF,N)
@test sort(real(ev2)) ≈ sort([-1;-2]) && norm(imag(ev2)) < 1.e-10

# Floquet approach
A = PeriodicFunctionMatrix(t -> [-1-sin(2*t)^2 2-0.5*sin(4*t); -2-0.5*sin(4*t) -1-cos(2*t)^2],pi/2);
@time ev = psceig(A; reltol = 1.e-10)
@time ev = psceig(A,50; reltol = 1.e-10)

# Fourier approach
Afun = FourierFunctionMatrix(Fun(t -> [-1-sin(2*t)^2 2-0.5*sin(4*t); -2-0.5*sin(4*t) -1-cos(2*t)^2],Fourier(0..π)))
ev1 = psceig(Afun)
@test sort(real(ev)) ≈ sort(real(ev1)) && norm(imag(ev1)) < 1.e-10

ev2 = psceig(Afun,5)
@test sort(real(ev)) ≈ sort(real(ev2)) && norm(imag(ev2)) < 1.e-10

ev3 = psceig(Afun,20)
@test sort(real(ev)) ≈ sort(real(ev3)) && norm(imag(ev3)) < 1.e-10

# # lossy Mathieu differential equation
# k = -1.; ξ = 0.05; β =0.2; ωh = 2; # unstable feedback
# k = -.5; ξ = 0.05; β =0.2; ωh = 2; # stable feedback
# Ahr = convert(HarmonicArray,PeriodicFunctionMatrix(t -> [0 1; k*(1-2β*cos(ωh*t)) -2*ξ],2pi));
# Asym = convert(PeriodicSymbolicMatrix,Ahr); Asym.F
# N = 5; ev = eigvals(hr2bt(Ahr,N)-phasemat(Ahr,N))
# ev[sortperm(imag(ev),by=abs)][1:2]

# Bitanti-Colaneri's book p.26
a = PeriodicFunctionMatrix(t -> [-1+sin(t) 0; 1-cos(t) -3],2*pi);
b = PeriodicFunctionMatrix(t ->  [-1-cos(t);2-sin(t)], 2*pi);
c = [0 1]; d = [0];

# using Harmonic Array based lifting
psyschr = ps(HarmonicArray,a,b,c,d)
ev = psceig(psyschr.A,20)

syshr = ps2fls(psyschr,20)
p = gpole(syshr,atol=1.e-7);  p = p[sortperm(imag(p),by=abs)][1:2];
@test sort(real(p)) ≈ sort(real(ev))  && norm(imag(p)) < 1.e-10

z = gzero(syshr,atol=1.e-7); z = z[isfinite.(z)]; # Question: How to handle infinite zeros?
ζ = z[sortperm(imag(z),by=abs)][1]

Z = convert(HarmonicArray,PeriodicFunctionMatrix(t -> [(-2+3*sin(t))/(2-sin(t))],2*pi))
ρ = real(Z.values[:,:,1])[1,1]
@test real(ζ) ≈ ρ

# using Fourier Function Matrix based lifting
psysc = ps(FourierFunctionMatrix,a,b,c,d);
ev1 = psceig(psysc.A,30)

sys = ps2frls(psysc,20);
p = gpole(sys,atol=1.e-7);  p = p[sortperm(imag(p),by=abs)][1:4];
@test minimum(abs.(p.-ev1[1])) < 1.e-10 && minimum(abs.(p.-ev1[2])) < 1.e-10

z = gzero(sys,atol=1.e-7); z = z[isfinite.(z)]; # Question: How to handle infinite zeros?
z = z[sortperm(imag(z),by=abs)][1:4];
@test minimum(abs.(z.-ρ)) < 1.e-10

# Zhou-Hagiwara Automatica 2002
β = 0.5
a1 = PeriodicFunctionMatrix(t -> [-1-sin(2*t)^2 2-0.5*sin(4*t); -2-0.5*sin(4*t) -1-cos(2*t)^2],pi);
#γ(t) = mod(t,pi) < pi/2 ? sin(2*t) : 0
γ = chop(Fun(t -> mod(t,pi) < pi/2 ? sin(2*t) : 0, Fourier(0..pi)),1.e-7);
b1 = PeriodicFunctionMatrix(t ->  [0; 1-2*β*(mod(t,float(pi)) < pi/2 ? sin(2*t) : 0 )], pi);
b1 = PeriodicFunctionMatrix(t ->  [0; 1-2*β*γ(t)], pi);
c = [1 1]; d = [0];

# using Harmonic Array based lifting
psyschr = ps(HarmonicArray,a1,b1,c,d);
ev = psceig(psyschr.A)

syshr = ps2fls(psyschr,60);
p = gpole(syshr,atol=1.e-7);  p = p[sortperm(imag(p),by=abs)][1:2]
@test sort(real(p)) ≈ sort(real(ev))  && norm(imag(p)) < 1.e-10

z = gzero(syshr,atol=1.e-7); z = z[isfinite.(z)]; #  Question: How to handle infinite zeros?
ζ = z[sortperm(imag(z),by=abs)][1]
@test ζ ≈ -3.5

# using Fourier Function Matrix based lifting
psysc = ps(FourierFunctionMatrix,a1,b1,c,d);
ev = psceig(psysc.A)

sys = ps2frls(psysc,60);
p = gpole(sys,atol=1.e-7);  p = p[sortperm(imag(p),by=abs)][1:2]
@test sort(real(p)) ≈ sort(real(ev))  && norm(imag(p)) < 1.e-10

z = gzero(sys,atol=1.e-7); z = z[isfinite.(z)]; #  Question: How to handle infinite zeros?
z = z[sortperm(imag(z),by=abs)][1]
@test z[1] ≈ -3.5

# # Ziegler's column

# β = 2pi; η =0.5; λ = 0.5;
# M = [1 3/8; 3/2 1];
# K0 = [3/8 3/16; -3/4 3/4]; K1 = λ*[-1 η; 0 4η-4];
# A = PeriodicFunctionMatrix(τ -> [eye(2) zeros(2,2); zeros(2,2) -M\(K0+cos(β*τ).*K1)],2*pi/β)

# # Mathieu equation: determination of minimum eigenvalue
# q = 2; a = 2; a = -1.513956885056448; a = -1
# Ahr = convert(HarmonicArray,PeriodicFunctionMatrix(t -> [0 1; a-2*q*cos(2*t) 0],pi));
# Asym = convert(PeriodicSymbolicMatrix,Ahr); Asym.F
# ev = psceig(Ahr)
# N = 20; ev = eigvals(hr2bt(Ahr,N)-phasemat(Ahr,N))
# ev[sortperm(imag(ev),by=abs)][1:2]

# N = 5; ev1 = eigvals(hr2bt1(Ahr,N,2)-phasemat1(Ahr,N,2))
# ev1[sortperm(imag(ev1),by=abs)][1:2]

# discrete liftings
Ad = PeriodicMatrix([[1. 0; 0 0], [1 1;1 1], [0 1; 1 0]], 6, nperiod = 2);
Bd = PeriodicMatrix( [[ 1; 0 ], [ 1; 1]] ,2);
Cd = PeriodicMatrix( [[ 1 1 ], [ 1 0]] ,2);
Dd = PeriodicMatrix( [[ 1 ]], 1);
psys = PeriodicStateSpace(Ad,Bd,Cd,Dd);
sys = ps2ls(psys)
sys1 = ps2ls(psys, 7, ss = true)
@test iszero(sys-sys1)

Ad = PeriodicMatrix([[1. 0], [1;1]],2);
Bd = PeriodicMatrix( [[ 1 ], [ 1; 1]] ,2);
Cd = PeriodicMatrix( [[ 1 1 ], [ 1 ]] ,2);
Dd = PeriodicMatrix( [[ 1 ]], 1);
psys = PeriodicStateSpace(Ad,Bd,Cd,Dd);
sys = ps2ls(psys)
sys1 = ps2ls(psys, 7, ss = true)
@test iszero(sys-sys1)

sys = ps2ls(psys, 2)
sys1 = ps2ls(psys, 8, ss = true)
@test iszero(sys-sys1)

sys = ps2ls(psys, cyclic = true)
sys1 = ps2ls(psys, 3, cyclic = true)
@test iszero(sys-sys1)

Ad = PeriodicArray(rand(Float32,2,2,10),10);
Bd = PeriodicArray(rand(2,1,2),2);
Cd = PeriodicArray(rand(1,2,3),3);
Dd = PeriodicArray(rand(1,1,1), 1);
psys = PeriodicStateSpace(Ad,Bd,Cd,Dd);
sys = ps2ls(psys)
sys1 = ps2ls(psys, ss = true)
@test iszero(sys-sys1)

sys = ps2ls(psys, 2)
sys1 = ps2ls(psys, 2, ss = true)
@test iszero(sys-sys1)

sys = ps2ls(psys, cyclic = true)
sys1 = ps2ls(psys, 31, cyclic = true)
@test iszero(sys-sys1)

end # test

end # module


# We shall say that an n-digit number is pandigital if it makes use of all the
# digits 1 to n exactly once; for example, the 5-digit number, 15234, is 1
# through 5 pandigital.
#
# The product 7254 is unusual, as the identity, 39 × 186 = 7254, containing
# multiplicand, multiplier, and product is 1 through 9 pandigital.
#
# Find the sum of all products whose multiplicand/multiplier/product identity
# can be written as a 1 through 9 pandigital.
#
# HINT: Some products can be obtained in more than one way so be sure to only
# include it once in your sum.

using ProjectEulerSolutions

# Uses factorization to find sets of digits from the products.  Shortcuts full
# comparisons when the total number of digits is incorrect.
function p032solution_factors(n::Integer=3)::Integer

    startval = 10^(fld(n, 2) - 1)
    endval = 10^fld(n, 2)

    desired_digits = Set{Integer}(1:n)
    products = Integer[]
    for k in startval:endval
        facs = factors(k)
        factorlen = length(facs)
        numdigits = map(ndigits, facs[2:end])
        numdigits = numdigits .+ reverse(numdigits) .+ ndigits(k)
        # Only look at half the factors
        for i in 2:cld(factorlen, 2)
            if numdigits[i-1] == n
                alldigits = Set(vcat(digits(facs[i]),
                                     digits(facs[factorlen - i + 2]),
                                     digits(k)))
                if (desired_digits == alldigits)
                    push!(products, k)
                end
            end
        end
    end
    return sum(Set(products))
end

# Uses a digits from multipliers of the correct lengths.  If a, b, c are the
# number of digits in the multiplicands and product, then a + b + c = n.  We
# know that a + b < c < a + b + 1, so 2(a + b) < n < 2(a + b) + 1.  So we pick
# a: 1 < a < 0.5*n and b: 0.5*n - 0.5 - a < b < 0.5*n - a
function p032solution_multipliers(n::Integer=5)::Integer
    products = Integer[]
    desired_digits = Set{Integer}(1:n)
    r = n / 2
    as = 2:floor(Integer, 10^r)
    for a in as
        na = ndigits(a)
        minb = max(10, floor(Integer, 10^(r - 0.5 - na)))
        maxb = max(minb, floor(Integer, 10^(r - na)))

        for b in minb:maxb
            c = a * b
            if na + ndigits(b) + ndigits(c) == n
                alldigits = Set(vcat(digits(a), digits(b), digits(c)))
                if (desired_digits == alldigits)
                    push!(products, c)
                end
            end
        end
    end
    return sum(Set(products))
end

p032 = Problems.Problem(Dict("Factors" => p032solution_factors,
                             "Multipliers" => p032solution_multipliers))

Problems.benchmark(p032, 9)


type NewtonSolver <: Solver1D
  solverInfo::SolverInfo
end
NewtonSolver(maxEvals::Int = 100, lowerBoundEnforced::Bool = false, upperBoundEnforced::Bool = false, lowerBound::Float64 = 0.0, upperBound::Float64 = 0.0) =
  NewtonSolver(SolverInfo(maxEvals, lowerBoundEnforced, upperBoundEnforced, lowerBound, upperBound))

# Newton Solver (safe)
function _solve(solver::NewtonSolver, f::Function, accuracy::Float64, xMin::Float64, xMax::Float64, fxMin::Float64, fxMax::Float64, root::Float64, eval_num::Int)
  # Orient the search so that f(xl) < 0
  if fxMin < 0.0
    xl = xMin
    xh = xMax
  else
    xh = xMin
    xl = xMax
  end

  # the stepsize before last
  dxold = xMax - xMin

  # and the last step
  dx = dxold

  froot = f(root)
  dfroot = f(root, Derivative()) # derivative, hurray multiple dispatch!
  eval_num += 1

  while eval_num < solver.solverInfo.maxEvals
    # bisect if out of range or not decreasing fast enough
    if ((root - xh) * dfroot - froot) * ((root - xl) * dfroot - froot) > 0.0 || abs(2.0 * froot) > abs(dxold * dfroot)
      dxold = dx
      dx = (xh - xl) / 2.0
      root = xl + dx
    else
      dxold = dx
      dx = froot / dfroot
      root -= dx
    end

    # Convergence criterion
    if abs(dx) < accuracy
      # root = f(root) # check this
      eval_num += 1
      return root
    end

    froot = f(root)
    dfroot = f(root, Derivative())
    eval_num += 1
    if froot < 0.0
      xl = root
    else
      xh = root
    end
  end

  error("Maximum number of function evals exceeded!")
end


mutable struct MSetLPSolver <: MSetSolver
  solver
  m # Int
  n_arms # Int

  # Optimisation model (reused by solve_linear).
  model # ::Model
  x # ::Vector{Variable}

  function MSetLPSolver(solver)
    return new(solver, nothing, nothing, nothing, nothing)
  end
end

has_lp_formulation(::MSetLPSolver) = true
supports_solve_budgeted_linear(::MSetLPSolver) = true
supports_solve_all_budgeted_linear(::MSetLPSolver) = false
approximation_ratio(::MSetLPSolver) = 1.0
approximation_term(::MSetLPSolver) = 0.0
approximation_ratio_budgeted(::MSetLPSolver) = 1.0
approximation_term_budgeted(::MSetLPSolver) = 0.0

copy(solver::MSetLPSolver) = MSetLPSolver(solver.solver)

function build!(solver::MSetLPSolver, m::Int, n_arms::Int)
  solver.m = m
  solver.n_arms = n_arms
  solver.model = Model(solver.solver)
  solver.x = @variable(solver.model, [1:n_arms], binary=true, lower_bound=0.0)

  for i in 1:n_arms
    set_name(solver.x[i], "x_$i")
  end

  @constraint(solver.model, sum(solver.x) <= m)
end

function get_lp_formulation(solver::MSetLPSolver, rewards::Dict{T, Float64}) where T
  return solver.model,
    sum(rewards[i] * solver.x[i] for i in keys(rewards)),
    Dict{Int, JuMP.VariableRef}(i => solver.x[i] for i in keys(rewards))
end


@testset "Intersection of two lines" begin

    line_a = Line([0, 0], [1, 0])
    line_b = Line([0, 0], [0, 1])
    @test intersect(line_a, line_b) == [0, 0]

    line_b = Line([5, -3], [0, 1])
    @test intersect(line_a, line_b) == [5, 0]

    @test intersect(Line([0, 0], [1, 1]), Line([1, 0], [1, -1])) == [0.5, 0.5]

    @test intersect(Line([0, 0, 0], [1, 1, 1]), Line([1, 1, 0], [0, 0, 1])) ≈ ones(3)

    line = Line([0, 0], [1, 1])
    line_almost_parallel = Line([1, 0], [1, 1.01])
    intersect(line, line_almost_parallel)

    message = "The lines are parallel."
    @test_throws ArgumentError(message) intersect(Line([0, 0], [1, 1]), Line([1, 0], [1, 1]))

    message = "The lines are not coplanar."
    line_a = Line([0, 0, 0], [1, 0, 0])
    line_b = Line([0, 1, 0], [0, 0, 1])
    @test_throws ArgumentError(message) intersect(line_a, line_b)
end

@testset "Intersection of plane and line" begin

    origin = [0, 0, 0]

    line = Line(origin, [0, 0, 1])
    plane = Plane(origin, [0, 0, 1])

    @test intersect(plane, line) == [0, 0, 0]
    @test intersect(plane, Line([3, 5, 2], [0, 0, 1])) == [3, 5, 0]

    message = "The line and plane are parallel."
    @test_throws ArgumentError(message) intersect(plane, Line(origin, [1, 0, 0]))

    # The line is almost parallel to the plane.
    # This throws an error if the tolerance is large enough.
    line_almost_parallel = Line([0, 0, 1], [1, 0, 1e-2])
    intersect(plane, line_almost_parallel)
    @test_throws ArgumentError(message) intersect(plane, line_almost_parallel; atol=1e-2)
end

@testset "Intersection of circle/sphere and line" begin

    for (circle_or_sphere, line, point_expected_a, point_expected_b) in [
        (Circle([0, 0], 1), Line([0, 0], [1, 0]), [-1, 0], [1, 0]),
        (Circle([0, 0], 1), Line([0, 0], [0, 1]), [0, -1], [0, 1]),
        (Circle([0, 0], 1), Line([0, 1], [1, 0]), [0, 1], [0, 1]),
        (
            Circle([0, 0], 1),
            Line([0, 0.5], [1, 0]),
            [-√3 / 2, 0.5],
            [√3 / 2, 0.5],
        ),
        (Circle([1, 0], 1), Line([0, 0], [1, 0]), [0, 0], [2, 0]),
        (Circle([1.5, 0], 1), Line([0, 0], [1, 0]), [0.5, 0], [2.5, 0]),
        (Sphere([0, 0, 0], 1), Line([0, 0, 0], [1, 0, 0]), [-1, 0, 0], [1, 0, 0]),
    ]

        point_a, point_b = intersect(circle_or_sphere, line)

        @test point_a == point_expected_a && point_b == point_expected_b
    end

    template = "The {} and line do not intersect."

    for (circle_or_sphere, line) in [
        (Circle([0, 0], 1), Line([0, 2], [1, 0])),
        (Sphere([0, 0, 0], 1), Line([0, 2, 0], [1, 0, 0])),
        (Circle([-5, 2], 1), Line([0, 0], [1, 0])),
        (Sphere([0, 0, 3], 1), Line([0, 0, 0], [1, 0, 0])),
        (Circle([0, 0], 0.5), Line([0, 1], [1, 0])),
        (Sphere([0, 0, 0], 0.5), Line([0, 1, 0], [1, 0, 0])),
        (Sphere([0, 0, 0], 1), Line([0, 1, 1], [1, 0, 0])),
    ]

        if circle_or_sphere isa Circle
            name = "circle"
        else
            name = "sphere"
        end

        message = format(template, name)

        @test_throws ArgumentError(message) intersect(circle_or_sphere, line)
    end
end

@testset "Intersection of two planes" begin

    for (plane_a, plane_b, line_expected) in [
        (Plane(zeros(3), [0, 0, 1]), Plane(zeros(3), [0, 1, 0]), Line(zeros(3), [1, 0, 0])),
        (Plane(zeros(3), [0, 0, 1]), Plane([0, 0, 1], [1, 0, 1]), Line([1, 0, 0], [0, 1, 0])),
        (Plane(zeros(3), [-1, 1, 0]), Plane([8, 0, 0], [1, 1, 0]), Line([4, 4, 0], [0, 0, -1])),
        (Plane([-1, 0, 0], [-1, 0, 1]), Plane([1, 0, 0], [1, 0, 1]), Line([0, 0, 1], [0, 1, 0])),
    ]
        line_intersection = intersect(plane_a, plane_b)
        @test line_intersection ≈ line_expected
    end

    message = "The planes are parallel."
    for (plane_a, plane_b, atol, throws_error) in [
        (Plane(zeros(3), [0, 0, 1]), Plane(zeros(3), [0, 0, 1]), 0, true),
        (Plane(zeros(3), [0, 0, 1]), Plane(zeros(3), [0, 0, -5]), 0, true),
        (Plane(zeros(3), [0, 0, 1]), Plane(zeros(3), [0, 1e-2, 1]), 0, false),
        (Plane(zeros(3), [0, 0, 1]), Plane(zeros(3), [0, 1e-2, 1]), 1e-3, true),
    ]
        if throws_error
            @test_throws ArgumentError(message) intersect(plane_a, plane_b, atol=atol)
        else
            intersect(plane_a, plane_b, atol=atol)
        end
    end
end


"""
    Pendulum{R}

The most simple canonical nonlinear system, consisting of a simple pendulum controlled by
torque at the base. Has 2 state and 1 control.

# Constructor
    Pendulum(; kwargs...)

with keyword arguments
* `mass` - mass of the pendulum, in kg (default = 1.0)
* `length` - length of the pendulum, in m (default = 0.5)
* `b` - friction coefficient (default=0.1)
* `I` - inertia, in kg⋅m² (default=0.25)
* `g` - gravity, in kg/m² (default=9.81)
"""
@with_kw mutable struct Pendulum{T} <: AbstractModel
    mass::T = 1.
    length::T = 0.5
    b::T = 0.1
    lc::T = 0.5
    I::T = 0.25
    g::T = 9.81
end

RobotDynamics.state_dim(::Pendulum) = 2
RobotDynamics.control_dim(::Pendulum) = 1

function dynamics(p::Pendulum, x, u)
    m = p.mass * p.lc * p.lc
    @SVector [x[2],
              u[1]/m - p.g*sin(x[1])/p.lc - p.b*x[2]/m]
end
@inline Base.position(::Pendulum, x) = @SVector zeros(3)
orientation(::Pendulum, x) = expm((pi-x[1])*@SVector [1,0,0.])


# Tests for validated_integ

using TaylorModels
# using LinearAlgebra: norm
using Test
# using Random

const _num_tests = 1_000

setformat(:full)

# NOTE: IntervalArithmetic v0.16.0 includes this function; but
# IntervalRootFinding is bounded to use v0.15.x
interval_rand(X::Interval{T}) where {T} = X.lo + rand(T) * (X.hi - X.lo)
interval_rand(X::IntervalBox) = interval_rand.(X)

function test_integ(fexact, t0, qTM, q0, δq0)
    normalized_box = symmetric_box(length(q0), Float64)
    # Time domain
    domt = domain(qTM[1])
    # Random time (within time domain) and random initial condition
    δt = rand(domt)
    δtI = (δt .. δt) ∩ domt
    q0ξ = interval_rand(δq0)
    q0ξB = IntervalBox([(q0ξ[i] .. q0ξ[i]) ∩ δq0[i] for i in eachindex(q0ξ)])
    # Box computed to overapproximate the solution at time δt
    q = evaluate.(evaluate.(qTM, δtI), (normalized_box,))
    # Box computed from the exact solution must be within q
    bb = all(fexact(t0+δtI, q0 .+ q0ξB) .⊆ q)
    # Display details if bb is false
    bb || @show(t0, domt, remainder.(qTM),
            δt, δtI, q0ξ, q0ξB, q,
            fexact(t0+δtI, q0 .+ q0ξB))
    return bb
end

@testset "Tests for `validated_integ`" begin
    @testset "falling_ball!" begin
        @taylorize function falling_ball!(dx, x, p, t)
            dx[1] = x[2]
            dx[2] = -one(x[1])
            nothing
        end
        exactsol(t, t0, x0) = (x0[1] + x0[2]*(t-t0) - 0.5*(t-t0)^2, x0[2] - (t-t0))

        # Initial conditions
        tini, tend = 0.0, 10.0
        normalized_box = symmetric_box(2, Float64)
        q0 = [10.0, 0.0]
        δq0 = 0.25 * normalized_box
        X0 = IntervalBox(q0 .+ δq0)

        # Parameters
        abstol = 1e-20
        orderQ = 2
        orderT = 4
        ξ = set_variables("ξₓ ξᵥ", order=2*orderQ, numvars=length(q0))

        @testset "Forward integration 1" begin
            tTM, qv, qTM = validated_integ(falling_ball!, X0, tini, tend, orderQ, orderT, abstol)

            @test length(qv) == length(qTM[1, :]) == length(tTM)

            end_idx = lastindex(tTM)
            # Random.seed!(1)
            for it = 1:_num_tests
                n = rand(2:end_idx)
                @test test_integ((t,x)->exactsol(t,tini,x), tTM[n], qTM[:,n], q0, δq0)
            end

            tTMf, qvf, qTMf = validated_integ(falling_ball!, X0, tini, tend, orderQ, orderT, abstol,
                adaptive=false)
            @test length(qvf) == length(qv)
            @test qTM == qTMf

            # initializaton with a Taylor model
            X0tm = qTM[:, 1]
            tTM2, qv2, qTM2 = validated_integ(falling_ball!, X0tm, tini, tend, orderQ, orderT, abstol)
            @test qTM == qTM2
        end

        @testset "Forward integration 2" begin
            tTM, qv, qTM = validated_integ2(falling_ball!, X0,
                tini, tend, orderQ, orderT, abstol)

            @test length(qv) == length(qTM[1, :]) == length(tTM)

            # Random.seed!(1)
            end_idx = lastindex(tTM)
            for it = 1:_num_tests
                n = rand(2:end_idx)
                @test test_integ((t,x)->exactsol(t,tini,x), tTM[n], qTM[:,n], q0, δq0)
            end

            # initializaton with a Taylor model
            X0tm = qTM[:, 1]
            tTM2, qv2, qTM2 = validated_integ2(falling_ball!, X0tm, tini, tend, orderQ, orderT, abstol)
            @test qTM == qTM2
        end

        # Initial conditions
        tini, tend = 10.0, 0.0
        q0 = [10.0, 0.0]
        δq0 = IntervalBox(-0.25 .. 0.25, 2)
        X0 = IntervalBox(q0 .+ δq0)

        @testset "Backward integration 1" begin
            tTM, qv, qTM = validated_integ(falling_ball!, X0, tini, tend, orderQ, orderT, abstol)

            @test length(qv) == length(qTM[1, :]) == length(tTM)

            # Random.seed!(1)
            end_idx = lastindex(tTM)
            for it = 1:_num_tests
                n = rand(2:end_idx)
                @test test_integ((t,x)->exactsol(t,tini,x), tTM[n], qTM[:,n], q0, δq0)
            end

            tTMf, qvf, qTMf = validated_integ(falling_ball!, X0, tini, tend, orderQ, orderT, abstol,
                adaptive=false)
            @test length(qvf) == length(qv)
            @test all(qTM .== qTMf)

            # initializaton with a Taylor model
            X0tm = qTM[:, 1]
            tTM2, qv2, qTM2 = validated_integ(falling_ball!, X0tm, tini, tend, orderQ, orderT, abstol)
            @test qTM == qTM2

            tTM2f, qv2f, qTM2f = validated_integ(falling_ball!, X0tm, tini, tend, orderQ, orderT, abstol,
                adaptive=false)
            @test length(qv2f) == length(qv2)
            @test all(qTM .== qTM2f)
        end

        @testset "Backward integration 2" begin
            tTM, qv, qTM = validated_integ2(falling_ball!, X0,
            tini, tend, orderQ, orderT, abstol)

            @test length(qv) == length(qTM[1, :]) == length(tTM)

            # Random.seed!(1)
            end_idx = lastindex(tTM)
            for it = 1:_num_tests
                n = rand(2:end_idx)
                @test test_integ((t,x)->exactsol(t,tini,x), tTM[n], qTM[:,n], q0, δq0)
            end
        end
    end

    @testset "x_square!" begin
        @taylorize function x_square!(dx, x, p, t)
            dx[1] = x[1]^2
            nothing
        end

        exactsol(t, x0) = 1 / (1/x0[1] - t)

        tini, tend = 0., 0.45
        normalized_box = symmetric_box(1, Float64)
        abstol = 1e-15
        orderQ = 5
        orderT = 20
        q0 = [2.]
        δq0 = 0.0625 * normalized_box
        X0 = IntervalBox(q0 .+ δq0)
        ξ = set_variables("ξₓ", numvars=1, order=2*orderQ)

        @testset "Forward integration 1" begin
            tTM, qv, qTM = validated_integ(x_square!, X0, tini, tend, orderQ, orderT, abstol)

            @test length(qv) == length(qTM[1, :]) == length(tTM)

            # Random.seed!(1)
            end_idx = lastindex(tTM)
            for it = 1:_num_tests
                n = rand(1:end_idx)
                @test test_integ((t,x)->exactsol(t,x), tTM[n], qTM[:,n], q0, δq0)
            end

            tTMf, qvf, qTMf = validated_integ(x_square!, X0, tini, tend, orderQ, orderT, abstol,
                adaptive=false)
            @test length(qvf) == length(qv)
            @test all(qTMf .== qTM)

            # initializaton with a Taylor model
            X0tm = copy(qTM[:, 1])
            tTM2, qv2, qTM2 = validated_integ(x_square!, X0tm, tini, tend, orderQ, orderT, abstol)
            @test qTM == qTM2
        end

        @testset "Forward integration 2" begin
            tTM, qv, qTM = validated_integ2(x_square!, X0, tini, tend, orderQ, orderT, abstol)

            @test length(qv) == length(qTM[1, :]) == length(tTM)

            # Random.seed!(1)
            end_idx = lastindex(tTM)
            for it = 1:_num_tests
                n = rand(1:end_idx)
                @test test_integ((t,x)->exactsol(t,x), tTM[n], qTM[:,n], q0, δq0)
            end
        end
    end

    @testset "Pendulum with constant torque" begin
        @taylorize function pendulum!(dx, x, p, t)
            si = sin(x[1])
            aux = 2 *  si
            dx[1] = x[2]
            dx[2] = aux + 8*x[3]
            dx[3] = zero(x[1])
            nothing
        end
        # Conserved quantity
        ene_pendulum(x) = x[2]^2/2 + 2 * cos(x[1]) - 8 * x[3]

        # Initial conditions
        tini, tend = 0.0, 12.0
        q0 = [1.1, 0.1, 0.0]
        δq0 = IntervalBox(-0.1 .. 0.1, -0.1 .. 0.1, 0..0)
        X0 = IntervalBox(q0 .+ δq0)
        ene0 = ene_pendulum(X0)

        # Parameters
        abstol = 1e-10
        orderQ = 3
        orderT = 10
        ξ = set_variables("ξ", order=2*orderQ, numvars=length(q0))

        tTM, qv, qTM = validated_integ(pendulum!, X0, tini, tend, orderQ, orderT, abstol);
        @test all(ene0 .⊆ ene_pendulum.(qv))

        tTM, qv, qTM = validated_integ2(pendulum!, X0, tini, tend, orderQ, orderT, abstol,
            validatesteps=32);
        @test all(ene0 .⊆ ene_pendulum.(qv))

        # Initial conditions 2
        q0 = [1.1, 0.1, 0.0]
        δq0 = IntervalBox(-0.1 .. 0.1, -0.1 .. 0.1, -0.01 .. 0.01)
        X0 = IntervalBox(q0 .+ δq0)
        ene0 = ene_pendulum(X0)

        tTM, qv, qTM = validated_integ(pendulum!, X0, tini, tend, orderQ, orderT, abstol);
        @test all(ene0 .⊆ ene_pendulum.(qv))

        tTM, qv, qTM = validated_integ2(pendulum!, X0, tini, tend, orderQ, orderT, abstol,
            validatesteps=32);
        @test all(ene0 .⊆ ene_pendulum.(qv))
    end
end


module Day13

export get_inputs, get_solution1, get_solution2

## Input getting
function get_inputs()
    test_input1 = test_input2 = read_input(IOBuffer(
        """
        939
        7,13,x,x,59,x,31,19
        """
    ))
    test_output1 = 295
    test_output2 = 1068781
    data = read_input(joinpath(@__DIR__, "input.txt"))
    return (; test_input1, test_input2, test_output1, test_output2, data)
end

function read_input(io)
    timestamp, ids = readlines(io)
    return parse(Int, timestamp), split(ids, ',')
end

## Solution functions
# Part 1
parse_input1(ids) = parse.(Int, filter(id -> id != "x", ids))

function get_solution1(data)
    ts, ids = data[1], parse_input1(data[2])
    times = ids .- rem.(ts, ids)
    i = argmin(times)
    return times[i] * ids[i]
end

# Part 2
function parse_input2(input)
    ids = Int[]
    offsets = Int[]
    for (i, val) in enumerate(input)
        if val!="x"
            push!(ids, parse(Int, val))
            push!(offsets, i-1)
        end
    end
    return ids .=> offsets
end

function get_start_time(t0, step, id, offset)
    offset -= (offset ÷ id) * id
    t = t0
    counter = (t0 ÷ id) + 1
    delta = counter * id - t0
    delta_step = step % id
    while delta != offset
        delta = mod(delta - delta_step, id)
        t += step
    end
    return t
end

function get_solution2(data)
    input = data[2]
    outs = parse_input2(input)
    t = 0
    id1, _ = outs[1]
    for i in 2:length(outs)
        id, offset = outs[i]
        t = get_start_time(t, id1, id, offset)
        id1 = lcm(id1, id)
    end
    return t
end

end


Base.@kwdef struct SmoothAnalysis
    npoints::Int=200
    span::Float64=0.75
    degree::Int=2
end

function (l::SmoothAnalysis)(input::ProcessedLayer)
    output = map(input) do p, _
        x, y = p
        model = Loess.loess(x, y; l.span, l.degree)
        x̂ = range(extrema(x)..., length=l.npoints)
        ŷ = Loess.predict(model, x̂)
        return (x̂, ŷ), (;)
    end
    plottype = Makie.plottype(output.plottype, Lines)
    return ProcessedLayer(output; plottype)
end

"""
    smooth(; span=0.75, degree=2, npoints=200)

Fit a loess model. `span` is the degree of smoothing, typically in `[0,1]`.
Smaller values result in smaller local context in fitting.
`degree` is the polynomial degree used in the loess model.
`npoints` is the number of points used by Makie to draw the line
"""
smooth(; options...) = transformation(SmoothAnalysis(; options...))


using myJuliaUtils
using Test
using LinearAlgebra
using Statistics

@testset "testing genLPM function - Monte Carlo test" begin
    T = 10000
    K = 400000
    X = [ones(T) randn(T)];
    lammbda = [rand(),randn() / 10.]

    lambbdaMean = 0. .* lammbda
    y = BitArray(undef, T)
    @time for i = 1:K
        y .= genLPM(lammbda, T, X)
        lambbdaMean .+= X \ y

        if mod(i, 100000) == 0
            println("keep")
        end

    end
    lambbdaMean ./= K

    @test_broken isapprox(lambbdaMean, lammbda, atol = 0.01)

end

@testset "testing genMC function" begin
    noStates = round(Int, 2 + rand() * 50)
    S0 = floor(Int, 1 + rand() * noStates)
    PI = eye(noStates)
    T = 100

    @test all(simMC(PI, T, S0 = S0) .== S0)

    noStates = round(Int, 2 + rand() * 20)
    S0 = floor(Int, 1 + rand() * noStates)

    A = randn(noStates, noStates)
    A = abs.(A)
    A = A ./ sum(A;dims = 2)
    B = copy(A)
    A[end,:] = [zeros(1, noStates - 1) 1]
    B[1,:] = [1 zeros(1, noStates - 1)]

    F = simMC(A, 1000, S0 = S0)
    @test F[end] == noStates
    @test F[1] == S0

    F = simMC(B, 1000, S0 = S0)
    @test F[end] == 1
    @test F[1] == S0
end


export
    DDproblem,
    value_iteration_step,
    transition_matrix,
    policy_iteration_step,
    DDsolution,
    solve_iteratively,
    simulate_transitions

# discrete state, discrete time dynamic programming

"""
Discrete state, discrete time dynamic programing problem with `N`
states and `M` actions.

`β` is the discount factor (infinite horizon).

`U` is an (N,M) matrix that describes the period payoffs for state `N`
and action `M`. For actions that are now allowed, use `-Inf`.

`P` is an `N`-element vector of `(M,N)` transition matrices, ie
`P[n][m,n′]` is the probability of state `n′` given state `n` and
action `m`.
"""
immutable DDproblem{Tβ <: Real, TU <: AbstractMatrix, TP <: AbstractMatrix}
    β::Tβ
    U::TU
    P::Vector{TP}
    function DDproblem(β, U, P)
        (N, M) = size(U)
        @assert length(P) == N "wrong number of elements in P"
        @assert all(size(p) == (M,N) for p in P) "non-conformable matrix in P"
        new(β, U, P)
    end
end

DDproblem{Tβ, TU, TP}(β::Tβ, U::TU, P::Vector{TP}) = DDproblem{Tβ, TU, TP}(β, U, P)

function show(io::IO, dd::DDproblem)
    (N,M) = size(dd.U)
    println(io, "Discrete state, discrete time dynamic programming problem")
    println(io, "  with $(N) states and $(M) actions, discount factor $(dd.β)")
end

"""
Value iteration step. Given value function `V`, return the new value
function and the indexes of optimal policies.
"""
function value_iteration_step(dd::DDproblem, V)
    @unpack β, U, P = dd
    optimal_Vms = [findmax(@view(U[n, :]) + β*P[n]*V) for n in indices(U, 1)]
    (first.(optimal_Vms), last.(optimal_Vms))
end

"""
Return a transition matrix between states given the optimal policies.

It will be the same type as the elements of `dd.P`, tus
"""
function transition_matrix(dd::DDproblem, optimal_ms)
    @unpack U, P = dd
    N = size(U, 1)
    M = zeros(N, N)
    for (n,m) in enumerate(optimal_ms)
        M[n, :] = P[n][m, :]
    end
    convert(eltype(P), M)
end

"""
Policy iteration step. Given a value function `V`, return the new
value function and the indexes of optimal policies.
"""
function policy_iteration_step(dd::DDproblem, V)
    optimal_ms = value_iteration_step(dd, V)[2]
    u = [dd.U[n,m] for (n,m) in enumerate(optimal_ms)]
    M = transition_matrix(dd, optimal_ms)
    V = (I - dd.β * M) \ u
    (V, optimal_ms)
end

"""
Structure for storing the result of iterative solution techniques.

`V` is the value function.
`ms` is the index of actions (only optimal when converged).
`iterations` is the number of iterations performed.
`converged` is a boolean which is `true` when convergence is achieved.
"""
immutable DDsolution
    V
    ms
    iterations
    converged
end

function show(io::IO, dd::DDsolution)
    print(io, "DDsolution ")
    if dd.converged
        print_with_color(:green, io, "converged")
    else
        print_with_color(:red, io, "did not converge")
    end
    println(io, " after $(dd.iterations) iterations")
end

"""
Solve a problem using value and policy iteration, using initial values
`V`. Successive value functions are compared with the L¹ norm,
convergence is achieved the difference between value functions is
below `ϵ`.

For the first `value_iter` steps, use value iteration, and after that,
policy iteration.
"""
function solve_iteratively(dd::DDproblem;
                           V = zeros(size(dd.U, 1)),
                           value_iter = 20,
                           max_iter = 100,
                           ϵ = √eps(Float64))
    for iteration in 1:max_iter
        (V′, ms) = iteration ≤ value_iter ? value_iteration_step(dd, V) :
            policy_iteration_step(dd, V)
        if norm(V′-V, 1) ≤ ϵ
            return DDsolution(V′, ms, iteration, true)
        elseif iteration == max_iter
            return DDsolution(V′, ms, iteration, false)
        else
            V .= V′
        end
    end
end

"""
Given a transition matrix `M` between states, an initial state `n0`,
return a sequence of `T` simulated transitions.
"""
function simulate_transitions(M, n0, T)
    ps = [Categorical(collect(M[n, :])) for n in indices(M, 1)]
    n = n0
    [(n_prev = n; n = rand(ps[n]); n_prev) for i in 1:T]
end


using StructOpt
using DataStructures, LinearAlgebra, PGFPlotsX

function main()
    ## build Lasso pball reg problem, with solution on partly smooth manifold ||x||=1, dimension 2.
    problems = []

    ## Lasso 2d
    sf = 1.0
    A = sf * Matrix{Float64}([1. 1.; 0. 1.])
    y = sf * Vector([1., 2.])
    λ = 1.0
    pb = LassoPb{regularizer_l1}(A, y, λ, size(A, 2), Vector([0, 1]))
    objopt, xopt = get_optimsol(pb)
    @show objopt, xopt
    objopt = 1.5

    xstart = Vector([4, 2]) + 50*Vector([1, 1])

    push!(problems, (
        # name = "Lasso-2d-IstaFista",
        name = "Lasso-2d",
        pb = pb,
        x0 = xstart,
        xopt = xopt,
        objopt = objopt,
    ))

    ## Build algorithms
    algorithms = []

    αuser = 0.1
    push!(algorithms, (
        name="ISTA",
        updatefunc=extra_ISTA,
        params=Dict(
            :αuser => αuser,
            :printstep => 1,
        ),
        # params=Dict(),
    ))

    push!(algorithms, (
        name="FISTA",
        updatefunc=extra_FISTA,
        # params=Dict(:αuser=>αuser),
        params=Dict(
            :αuser => αuser,
            :printstep => 1,
        ),
    ))

    push!(algorithms, (
        name="T1",
        updatefunc=extra_CondInertia,
        params=Dict(
            :id_testkind=>:default,
            :αuser => αuser,
            :printstep => 1,
        ),
    ))

    push!(algorithms, (
        name="T2",
        updatefunc=extra_CondPredInertia,
        params=Dict(
            :printstep => 1,
            :αuser => αuser,
            :printstep => 1,
            :proxgrad_per_it => 2,
        ),
    ))

    push!(algorithms, (
        name="MFISTA",
        updatefunc=extra_MFISTA,
        params=Dict(
            :αuser => αuser,
            :printstep => 1,
        ),
    ))

    ## Start from interesting point and optimize with ISTA, FISTA.
    problem_to_algstats = OrderedDict()
    for (pb_id, pb_data) in enumerate(problems)
        printstyled("\n----- Solving problem $pb_id, $(pb_data.name)\n", color=:light_blue)

        algo_to_stats = OrderedDict()
        for (algo_id, algo) in enumerate(algorithms)
            printstyled("\n---- Algorithm $algo_id, $(algo.name)\n", color=:light_blue)

            xopt, hist = solve_proxgrad(pb_data.pb, pb_data.x0, algo.updatefunc; algo.params...)
            algo_to_stats[algo] = (hist=hist, xopt=xopt)
        end

        problem_to_algstats[pb_data] = algo_to_stats
    end

    ####################################################
    ## Display suboptimality, iterates position.
    FIGS_FOLDER = "./figs"
    basename(pwd()) == "src" && (FIGS_FOLDER = joinpath("..", FIGS_FOLDER))
    !ispath(FIGS_FOLDER) && mkpath(FIGS_FOLDER)

    for (pb, algo_to_stats) in problem_to_algstats
        println("Generating graphs for ", pb.name)

        ## Suboptimality plot
        build_save_suboptimality(pb, algo_to_stats, FIGS_FOLDER)

        ## Steplength plot
        # build_save_steplength(pb, algo_to_stats, FIGS_FOLDER)

        if pb.pb.n == 2
            reg = typeof(pb.pb).parameters[1]

            ## Iterates position
            ps = []

            xmin, xmax = -1, 4
            ymin, ymax = -1, 2

            ## Plot contour
            x = xmin:(xmax - xmin) / 100:xmax
            y = ymin:(ymax - ymin) / 100:ymax
            f = (x,y) -> 0.5 * norm(pb.pb.A * [x, y] - pb.pb.y)^2
            add_contour!(ps, f, x, y)

            for (algo, stats) in algo_to_stats
                coords = [ (point[1], point[2]) for (iter, point) in stats.hist[:iter_x]]

                add_traj!(ps, coords, algo.name, params=get_iterates_algoparams(algo.name))
            end

            ## Plot optimal point
            add_point!(ps, xopt)

            ## Plot manifolds
            coords = [(xmin, 0), (xmax, 0)]; add_manifold!(ps, coords)
            coords = [(0, ymin), (0, ymax)]; add_manifold!(ps, coords)

            fig = pgf_build_iteratesfig(ps, xmin, xmax, ymin, ymax)

            pgfsave(joinpath(FIGS_FOLDER, "$(pb.name)_iterates.pdf"), fig)
            pgfsave(joinpath(FIGS_FOLDER, "$(pb.name)_iterates.tikz"), fig)
        end
    end

    println("\n")
    for (pb, algo_to_stats) in problem_to_algstats
        @show pb.name
        for (algo, stats) in algo_to_stats
            @show algo.name
            display(stats[:xopt])
        end
    end

    return
end

main()


module FCG

include("../Misc.jl")
using .Misc

using Distributions, SpecialFunctions, Einsum, Combinatorics
import Base.Iterators: product

export generate, log_marginal

function generate(T::Ƶ, dims::Vararg{Ƶ}; γ::ℜ=0.1) where {ℜ<:Real, Ƶ<:Int}
    len = prod(dims)
    θ_flat = rand(Dirichlet(γ,len))

    X::Array{ℜ} = reshape(rand(Multinomial(T,θ_flat)),dims)

    return X, reshape(θ_flat,dims)
end

function log_marginal(X::Array{ℜ,N}; μ::ℜ=nanmean(X), γ::ℜ=0.1) where {ℜ<:Real, Ƶ<:Int, N}
    dims, len = size(X), length(X)
    X₊ = sum(X)

    a::ℜ, b::ℜ = len*γ, γ/μ

    log_PX = a*log(b) - (a + X₊)*log(b + 1.0) + lgamma(a + X₊) - lgamma(a) - sum(lgamma.(X .+ 1.0))
    log_PX += sum(lgamma.(γ .+ X)) - lgamma(a + X₊) + lgamma(a) - len*lgamma(γ)

    return log_PX
end

end