"""
    BaB(optimizer, ϵ::Float64)

BaB uses branch and bound to estimate the range of the output node.

# Problem requirement
1. Network: any depth, ReLU activation, single output
2. Input: hyperrectangle
3. Output: hyperrectangle (1d interval)

# Return
`CounterExampleResult` or `ReachabilityResult`

# Method
Branch and bound.
For branch, it uses iterative interval refinement.
For bound, it computes concrete bounds by sampling, approximated bound by optimization.
- `optimizer` default `GLPKSolverMIP()`
- `ϵ` is the desired accurancy for termination, default `0.1`.

# Property
Sound and complete.

# Reference
[R. Bunel, I. Turkaslan, P. H. Torr, P. Kohli, and M. P. Kumar,
"A Unified View of Piecewise Linear Neural Network Verification,"
*ArXiv Preprint ArXiv:1711.00455*, 2017.](https://arxiv.org/abs/1711.00455)
"""
@with_kw struct BaB <: Solver
    optimizer = GLPK.Optimizer
    ϵ::Float64 = 0.1
end

function solve(solver::BaB, problem::Problem)
    (u_approx, u, x_u) = output_bound(solver, problem, :max)
    (l_approx, l, x_l) = output_bound(solver, problem, :min)
    bound = Hyperrectangle(low = [l], high = [u])
    reach = Hyperrectangle(low = [l_approx], high = [u_approx])

    output = problem.output
    if reach ⊆ output
        return ReachabilityResult(:holds, [reach])
    end
    high(bound) > high(output) && return CounterExampleResult(:violated, x_u)
    low(bound)  < low(output)  && return CounterExampleResult(:violated, x_l)
    return ReachabilityResult(:unknown, [reach])
end

function output_bound(solver::BaB, problem::Problem, type::Symbol)
    nnet = problem.network
    global_concrete, x_star = concrete_bound(nnet, problem.input, type)
    global_approx = approx_bound(nnet, problem.input, solver.optimizer, type)
    doms = Tuple{Float64, Hyperrectangle}[(global_approx, problem.input)]
    index = ifelse(type == :max, 1, -1)
    while index * (global_approx - global_concrete) > solver.ϵ
        dom, doms = pick_out(doms) # pick_out implements the search strategy
        subdoms = split_dom(dom[2]) # split implements the branching rule
        for i in 1:length(subdoms)
            dom_concrete, x = concrete_bound(nnet, subdoms[i], type) # Sample
            dom_approx = approx_bound(nnet, subdoms[i], solver.optimizer, type)
            if index * (dom_concrete - global_concrete) > 0
                (global_concrete, x_star) = (dom_concrete, x)
            end
            if index * (dom_approx - global_concrete) > 0
                add_domain!(doms, (dom_approx, subdoms[i]), type)
            end
        end
        length(doms) == 0 && return (global_approx, global_concrete, x_star)
        global_approx = doms[1][1]
    end
    return (global_approx, global_concrete, x_star)
end

# Always pick the first dom
function pick_out(doms)
    return (doms[1], doms[2:end])
end

function add_domain!(doms::Vector{Tuple{Float64, Hyperrectangle}}, new::Tuple{Float64, Hyperrectangle}, type::Symbol)
    rank = length(doms) + 1
    index = ifelse(type == :max, 1, -1)
    for i in 1:length(doms)
        if index * (new[1] - doms[i][1]) >= 0
            rank = i
            break
        end
    end
    insert!(doms, rank, new)
end

# Always split the longest input dimension
function split_dom(dom::Hyperrectangle)
    max_value, index_to_split = findmax(dom.radius)
    return split_interval(dom, index_to_split)
end

# For simplicity
function concrete_bound(nnet::Network, subdom::Hyperrectangle, type::Symbol)
    points = [center(subdom), low(subdom), high(subdom)]
    vals = [sum(compute_output(nnet, p)) for p in points]
    i = (type == :min) ? argmin(vals) : argmax(vals)
    return vals[i], points[i]
end

function approx_bound(nnet::Network, dom::Hyperrectangle, optimizer, type::Symbol)
    model = Model(optimizer)
    model[:bounds] = get_bounds(nnet, dom, false)
    model[:before_act] = true
    z = init_vars(model, nnet, :z, with_input=true)

    add_set_constraint!(model, dom, first(z))
    encode_network!(model, nnet, TriangularRelaxedLP())
    index = ifelse(type == :max, 1, -1)
    o = sum(last(z))
    @objective(model, Max, index * o)
    optimize!(model)
    termination_status(model) == OPTIMAL && return value(o)
    error("Could not find bound for dom: ", dom)
end


function rec(G, dp1, dp2, v, p = -1)
    MOD = 10^9 + 7

    for v′ ∈ G[v]
        v′ != p && rec(G, dp1, dp2, v′, v)
    end

    for v′ ∈ G[v]
        if v′ != p
            dp1[v] = dp1[v] * (dp1[v′] + dp2[v′]) % MOD
            dp2[v] = dp2[v] * dp1[v′] % MOD
        end
    end

end

function main()
    MOD = 10^9 + 7
    N = parse(Int, readline())
    G = [Vector{Int}() for _ in 1:N]

    for _ = 1:N-1
        x, y = parse.(Int, split(readline()))
        push!(G[x], y)
        push!(G[y], x)
    end

    dp1 = ones(Int, N)
    dp2 = ones(Int, N)
    rec(G, dp1, dp2, 1)

    println((dp1[1] + dp2[1]) % MOD)
end

main()


function part1(input)
    reindeers = parse_line.(eachline(input))
    return foldl(max, flying_distance.(reindeers, 2503))
end

function part2(input)
    reindeers = parse_line.(eachline(input))
    return maximum(scores(reindeers, 2503))
end

parse_line(line) = parse.(Int, split(line)[[4, 7, 14]])

function flying_distance(reindeer, t)
    v, t1, t2 = reindeer
    return v * (t1 * (t ÷ (t1 + t2)) + min(t1, t % (t1 + t2)))
end

function scores(reindeers, t)
    n = length(reindeers)
    positions = fill(0, n)
    scores = fill(0, n)
    for t = 1:t
        for i = 1:n
            v, t1, t2 = reindeers[i]
            if mod1(t, t1 + t2) <= t1
                positions[i] += v
            end
        end
        scores += positions .== maximum(positions)
    end
    return scores
end


using JuMP

using LinearAlgebra

function formulate_nearest_good_box_qp(
    Q::Matrix,
    c::Vector,
    y_max::Float64,
    x_target::Vector{Float64},
)
    n = length(c)
    @assert n == size(Q, 1) == size(Q, 2)
    model = JuMP.Model()
    JuMP.@variable(model, 0 ≤ x[1:n] ≤ 1)
    JuMP.@constraint(model, -0.5dot(x, Q * x) - dot(c, x) ≤ y_max)

    JuMP.@variable(model, t[1:n] >= 0)
    for i = 1:n
        JuMP.@constraint(model, t[i] >= x[i] - x_target[i])
        JuMP.@constraint(model, t[i] >= -x[i] + x_target[i])
    end
    JuMP.@objective(model, Min, sum(t))
    return JuMP.backend(model)
end

using Justitia
using MathOptInterface
const MOI = MathOptInterface
using QuadraticRelaxations

const MULTIPLICATIVE_FACTOR_TO_OPTIMAL = 0.95

mutable struct NearestGoodBoxQPInstance <: Justitia.AbstractInstance
    box_qp_model::MOI.ModelLike
    Q::Matrix{Float64}
    c::Vector{Float64}
    primal_factory::Any
    nearest_good_model::MOI.ModelLike
    y_max::Float64
    cutoff::Float64

    # ASSUMPTION: BoxQP instance is minimization
    NearestGoodBoxQPInstance(
        box_qp_model::MOI.ModelLike,
        Q::Matrix{Float64},
        c::Vector{Float64},
        primal_factory,
    ) = new(box_qp_model, Q, c, primal_factory)
end

function Justitia.prep_instance!(instance::NearestGoodBoxQPInstance)
    let primal_model = MOI.instantiate(instance.primal_factory)
        MOI.copy_to(primal_model, instance.box_qp_model)
        MOI.optimize!(primal_model)
        @assert MOI.get(primal_model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
        instance.y_max = MOI.get(primal_model, MOI.ObjectiveValue())
        @assert instance.y_max < 0
    end
    instance.nearest_good_model = formulate_nearest_good_box_qp(
        instance.Q,
        instance.c,
        MULTIPLICATIVE_FACTOR_TO_OPTIMAL * instance.y_max,
        0.5 * ones(length(instance.c)),
    )
    let primal_model = MOI.instantiate(instance.primal_factory)
        MOI.copy_to(primal_model, instance.nearest_good_model)
        MOI.optimize!(primal_model)
        if MOI.get(primal_model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
            instance.cutoff = MOI.get(primal_model, MOI.ObjectiveValue())
        else
            instance.cutoff = 1e30
        end
    end
    return nothing
end

function _create_nearest_good_instance(family::String, instance::String, primal_factory)
    Q, c = parse_box_qp(joinpath(INSTANCE_DIR, family, instance))
    return NearestGoodBoxQPInstance(formulate_box_qp(Q, c), Q, c, primal_factory)
end

function Justitia.build_model(
    approach::BoxQPApproach,
    instance::NearestGoodBoxQPInstance,
    config::Dict{String},
)
    @assert collect(keys(config)) == ["factory_config"]
    src_model = approach.transformation(instance.nearest_good_model)
    dest_model =
        MOI.instantiate(approach.factory(config["factory_config"], instance.cutoff))
    MOI.copy_to(dest_model, src_model, copy_names = false)
    return Justitia.MOIModel(dest_model)
end

Base.@kwdef mutable struct NearestGoodBoxQPResult <: Justitia.AbstractResult
    termination_status::MOI.TerminationStatusCode
    primal_status::MOI.ResultStatusCode
    solve_time_sec::Float64
    x::Union{Nothing,Vector{Float64}} = nothing
    primal_bound::Float64 = NaN
    dual_bound::Float64 = NaN
    cutoff::Float64
    node_count::Int = -1
end

const _NEAREST_GOOD_BOXQP_RESULT_FIELDS = [
    :termination_status,
    :primal_status,
    :primal_bound,
    :dual_bound,
    :cutoff,
    :solve_time_sec,
    :node_count,
]

function Justitia.CSVRecord(filename::String, ::Type{NearestGoodBoxQPResult})
    # Don't write feasible solution to CSV
    return Justitia.CSVRecord(
        filename,
        vcat(
            ["instance", "approach"],
            [string(field) for field in _NEAREST_GOOD_BOXQP_RESULT_FIELDS],
        ),
    )
end

function Justitia.record_result!(
    table::Justitia.CSVRecord,
    result::NearestGoodBoxQPResult,
    instance_name::String,
    approach_name::String,
)
    println(
        table.fp,
        join(
            vcat(
                [instance_name, approach_name],
                [getfield(result, field) for field in _NEAREST_GOOD_BOXQP_RESULT_FIELDS],
            ),
            ",",
        ),
    )
    return flush(table.fp)
end

function Justitia.tear_down(
    model::Justitia.MOIModel,
    instance::NearestGoodBoxQPInstance,
    ::Type{NearestGoodBoxQPResult},
)
    opt = model.opt
    result = NearestGoodBoxQPResult(
        termination_status = MOI.get(opt, MOI.TerminationStatus()),
        primal_status = MOI.get(opt, MOI.PrimalStatus()),
        solve_time_sec = MOI.get(opt, MOI.SolveTime()),
        cutoff = instance.cutoff,
    )
    if result.primal_status == MOI.FEASIBLE_POINT
        num_vars = MOI.get(opt, MOI.NumberOfVariables())
        all_var_vals =
            [MOI.get(opt, MOI.VariablePrimal(), MOI.VariableIndex(i)) for i = 1:num_vars]
        result.x = all_var_vals
        n = length(instance.c)
        @assert n == size(instance.Q, 1) == size(instance.Q, 2)
        x_val = all_var_vals[1:n]
    end
    try
        primal_bound = MOI.get(opt, MOI.ObjectiveValue())
        result.primal_bound = primal_bound
    catch ArgumentError
        @warn "Primal bound not available; proceeding anyway."
    end
    try
        dual_bound = MOI.get(opt, MOI.ObjectiveBound())
        result.dual_bound = dual_bound
    catch ArgumentError
        @warn "Dual bound not available; proceeding anyway."
    end
    try
        node_count = MOI.get(opt, MOI.NodeCount())
        result.node_count = node_count
    catch ArgumentError
        @warn "Node count not available; proceeding anyway."
    end
    return result
end


function GD_ellispe(∇f::Function, b, x_k, y_k, ϵ)
    k = 1
    x_vals = [x_k]
    y_vals = [y_k]
    while norm(∇f(x_k,y_k)) > ϵ
        x_k = b*((b-1)/(b+1))^k
        y_k = ((1-b)/(b+1))^k
        k+=1

        push!(x_vals, x_k)
        push!(y_vals, y_k)
    end
    return x_vals, y_vals
end

function Quad_Backtracking(f::Function, ∇f::Function, P, x_k, y_k, α, β, ϵ)

    x_vals = []
    y_vals = []

    ∇fx = ∇f(x_k, y_k)

    for i in 1:200
        if norm(∇fx) < ϵ
            break
        end
        ∇fx = ∇f(x_k, y_k)
        fx = f(x_k, y_k)
        Δxk = P\∇fx
        ∇fxT_Δxk = ∇fx'Δxk
        t = 1
        while f(x_k + t*Δxk[1], y_k + t*Δxk[2]) > fx + α*t*∇fxT_Δxk
            t = β * t
        end
        x_k = x_k + t*Δxk[1]
        y_k = y_k + t*Δxk[2]

        push!(x_vals, x_k)
        push!(y_vals, y_k)
    end
    return x_vals, y_vals
end

function updating_Hessien_Descent(f::Function, ∇f::Function, H::Function, x_k, y_k, l, α, β, ϵ)

    x_vals = []
    y_vals = []

    P = -1 * H(x_k, y_k)
    for i in 1:200
        ∇fx = ∇f(x_k, y_k)
        if norm(∇fx) < ϵ
            break
        end

        if i % l == 0
            P = -1 * H(x_k, y_k)
        end

        fx = f(x_k, y_k)
        Δxk = P\∇fx
        ∇fxT_Δxk = ∇fx'Δxk
        t=1
        while f(x_k + t*Δxk[1], y_k + t*Δxk[2]) > fx + α*t*∇fxT_Δxk
            t = β * t
        end

        x_k = x_k + t*Δxk[1]
        y_k = y_k + t*Δxk[2]

        push!(x_vals, x_k)
        push!(y_vals, y_k)
    end
    return x_vals, y_vals
end

;


@doc raw"""
    PSFModels.Moffat([T=Float64]; fwhm, x=0, y=0, amp=1, alpha=1, maxsize=3, extent=maxsize .* fwhm)
    PSFModels.Moffat([T=Float64]; fwhm, pos, amp=1, alpha=1, maxsize=3, extent=maxsize .* fwhm)
    PSFModels.Moffat([T=Float64]; fwhm, r, theta, amp=1, alpha=1, maxsize=3, extent=maxsize .* fwhm)

Two dimensional Moffat model. The position can be specified in `(x, y)` coordinates as a `Tuple`, `AbstractVector`, or as separate arguments. By default the model is placed at the origin. The position can also be given as a polar coordinate using `r`/`ρ` and `theta`/`θ`, optionally centered around `origin`.

The `fwhm` can be a scalar (isotropic) or a vector/tuple (diagonal). For efficient calculations, we recommend using [StaticArrays](https://github.com/JuliaArrays/StaticArrays.jl). Here, `maxsize` is a multiple of the fwhm, and can be given as a scalar or as a tuple for each axis. The `extent` defines the bounding box for the model and is used for the default rendering size.

# Functional form
```
f(x | x̂, FWHM, α) = A / (1 + ||x - x̂|| / (FWHM / 2)^2)^α
```
where `x̂` and `x` are position vectors (indices) `||⋅||` represents the square-distance, and `FWHM` is the full width at half-maximum. If `FWHM` is a vector or tuple, the weighting is applied along each axis.
"""
struct Moffat{T,FT,VT<:AbstractVector,IT<:Tuple} <: PSFModel{T}
    fwhm::FT
    pos::VT
    amp::T
    alpha::Int
    indices::IT

    function Moffat(::Type{T}, fwhm::FT, pos::VT, amp::T, alpha, indices::IT) where {T,FT,VT<:AbstractVector,IT<:Tuple}
        new{T,FT,VT,IT}(fwhm, pos, amp, alpha, indices)
    end
end

## constructors
# default type is Float64
Moffat(; kwargs...) = Moffat(Float64; kwargs...)
function Moffat(T; fwhm, amp=one(T), alpha=1, maxsize=3, extent = maxsize .* fwhm, position...)
    # get the position from keyword distpatch
    pos = _position(; position...)
    return Moffat(T, fwhm, pos, convert(T, amp), alpha, indices_from_extent(pos, extent))
end

Base.size(m::Moffat) = map(length, m.indices)
Base.axes(m::Moffat) = m.indices

# scalar case
function (m::Moffat{T})(point::AbstractVector) where T
    hwhm = m.fwhm / 2
    Δ = sqeuclidean(point, m.pos)
    val = m.amp / (1 + Δ / hwhm^2)^m.alpha
    return convert(T, val)
end

# vector case
function (m::Moffat{T,<:Union{AbstractVector,Tuple}})(point::AbstractVector) where T
    weights = SA[(2 / first(m.fwhm))^2, (2 / last(m.fwhm))^2]
    Δ = wsqeuclidean(point, m.pos, weights)
    val = m.amp  / (1 + Δ)^m.alpha
    return convert(T, val)
end


struct ClusterInfo
    data::Matrix{Float64}
    label::Int
    index::Vector{Int}
    size::Int
    df::Float64
    center::Vector{Float64}
    cov::Matrix{Float64}
    function ClusterInfo(x, index, kmeans, label::Int)
        c = findall(kmeans.labels.==label)
        data = x[c, :]
        index = index[c]
        cluster_size = size(data, 1)
        features = size(data, 2)
        df = features*(features+3)/2
        centers = kmeans.μ[label, :]
        data_cov = cov(data)
        return new(data, label, index, cluster_size, df, centers, data_cov)
    end
end

function build_cluster_info(x, k_means::Kmeans; index=nothing)
    if index === nothing
        index = collect(1:size(x, 1))
    end
    labels = 1:n_clusters(k_means)

    return [ClusterInfo(x, index, k_means, label) for label in labels]
end

function log_likehood(c::ClusterInfo)
    return sum([logpdf(c.data[i, :], c.center, c.cov) for i in 1 : size(c.data, 1)])
end
function cluster_bic(c::ClusterInfo)
    return -2*log_likehood(c)+c.df*log(c.size)
end

"""
    Xmeans(; kinit=2, max=300, th=1e-4)
Xmeans method. This algorithm uses BIC to determine if a cluster should be split in two.
In other words, you don't need to give the number of clusters, you can cluster only with the data.

# Parameters:
- `kinit`: initial number of class
- `max`: maximum number of repitition
- `th`: converge threshold

# Example
```jldoctest
julia> model = Xmenas()
Kmeans{3}(Matrix{Float64}(undef, 0, 0), 100000000, 0.0001)

julia> using HorseML.Clustering: fit!

julia> fit!(model, x)

julia> model.labels |> size
(100, 3)
```
"""
mutable struct Xmeans
    kinit::Int
    max::Int64
    th::Float64
    labels::Vector{Int}
    centers::Vector{Vector{Float64}}
    log_likehoods::Vector{Float64}
    sizes::Vector{Int}
    Xmeans(; kinit=2, max=300, th=1e-4) = new(kinit, max, th, [], [], [], [])
end

function fit!(model::Xmeans, x)
    xmeans_clusters = ClusterInfo[]

    kmeans = Kmeans(model.kinit, max=model.max, th=model.th)
    fit!(kmeans, x)
    clusters = build_cluster_info(x, kmeans)
    recursively_split!(model, clusters, xmeans_clusters)

    labels = Array{Int}(undef, size(x, 1))
    for (i, c) in enumerate(xmeans_clusters)
        labels[c.index] .= i
    end
    model.labels = labels
    model.centers = [cluster.center for cluster in xmeans_clusters]
    model.log_likehoods = [log_likehood(cluster) for cluster in xmeans_clusters]
    model.sizes = [cluster.size for cluster in xmeans_clusters]
end

function recursively_split!(
        model::Xmeans,
        clusters::Vector{ClusterInfo},
        xmeans_clusters::Vector{ClusterInfo}
    )
    for cluster in clusters
        if cluster.size <= 3
            push!(xmeans_clusters)
            continue
        end
        kmeans = Kmeans(2, max=model.max, th=model.th)
        fit!(kmeans, cluster.data)
        c1, c2 = build_cluster_info(cluster.data, kmeans, index=cluster.index)

        β = norm(c1.center-c2.center) / sqrt(det(c1.cov) + det(c2.cov))
        α = 0.5 / norm_cdf(β)
        bic = -2*(cluster.size * log(α) + log_likehood(c1) + log_likehood(c2)) + 2*cluster.df*log(cluster.size)

        if bic < cluster_bic(cluster)
            recursively_split!(model, [c1, c2], xmeans_clusters)
        else
            push!(xmeans_clusters, cluster)
        end
    end
end


# TODO Move this to GPLikelihoods.jl
@doc raw"""
    LaplaceLikelihood(β::Real)

Likelihood with a [Laplace distribution](https://en.wikipedia.org/wiki/Laplace_distribution):
```math
    p(y|f,\beta) = \frac{1}{2\beta}\exp\left(-\frac{|y-f|}{\beta}\right)
```

## Arguments
- `β::Real`, scale parameter
"""
struct LaplaceLikelihood{Tβ<:Real} <: AbstractLikelihood
    β::Tβ
end

LaplaceLikelihood() = LaplaceLikelihood(1.0)

(lik::LaplaceLikelihood)(f::Real) = Laplace(f, lik.β)

function (lik::LaplaceLikelihood)(f::AbstractVector{<:Real})
    return Product(lik.(f))
end

@inline laplace_λ(lik::LaplaceLikelihood) = inv(abs2(2 * lik.β))

aux_field(::LaplaceLikelihood, Ω) = getproperty(Ω, :ω)

function init_aux_variables(rng::AbstractRNG, ::LaplaceLikelihood, n::Int)
    return TupleVector((; ω=rand(rng, InverseGamma(), n)))
end

function init_aux_posterior(T::DataType, lik::LaplaceLikelihood, n::Int)
    λ = laplace_λ(lik)
    return For(TupleVector(; μ=zeros(T, n))) do φ
        NTDist(InverseGaussian(φ.μ, λ))
    end
end

function aux_full_conditional(lik::LaplaceLikelihood, y::Real, f::Real)
    return NTDist(InverseGaussian(inv(2 * lik.β * abs(y - f)), 2 * laplace_λ(lik)))
end

function aux_posterior!(
    qΩ, lik::LaplaceLikelihood, y::AbstractVector, qf::AbstractVector{<:Normal}
)
    φ = only(qΩ.inds)
    map!(φ.μ, y, qf) do yᵢ, qfᵢ
        inv(2 * lik.β * sqrt(second_moment(qfᵢ, yᵢ)))
    end
    return qΩ
end

function auglik_potential(::LaplaceLikelihood, Ω, y::AbstractVector)
    return (2 * Ω.ω .* y,)
end

function auglik_precision(::LaplaceLikelihood, Ω, ::AbstractVector)
    return (2 * Ω.ω,)
end

function expected_auglik_potential(::LaplaceLikelihood, qΩ, y::AbstractVector)
    return (2 * tvmean(qΩ).ω .* y,)
end

function expected_auglik_precision(::LaplaceLikelihood, qΩ, ::AbstractVector)
    return (2 * tvmean(qΩ).ω,)
end

function logtilt(
    lik::LaplaceLikelihood, Ω::TupleVector, y::AbstractVector, f::AbstractVector{<:Real}
)
    return length(y) * (loggamma(1//2) - log(sqrtπ) - log(2 * lik.β)) +
           mapreduce(+, aux_field(lik, Ω), y, f) do ωᵢ, yᵢ, fᵢ
        -abs2(yᵢ - fᵢ) * ωᵢ
    end
end

function logtilt(lik::LaplaceLikelihood, ω::Real, y::Real, f::Real)
    return loggamma(1//2) - log(sqrtπ) - log(2 * lik.β) - abs2(y - f) * ω
end

function expected_logtilt(lik::LaplaceLikelihood, qΩ, y, qf::AbstractVector{<:Normal})
    return length(y) * (loggamma(1//2) - log(sqrtπ) - log(2 * lik.β)) +
           mapreduce(+, y, qf, @ignore_derivatives marginals(qΩ)) do yᵢ, qfᵢ, qωᵢ
        -second_moment(qfᵢ, yᵢ) * ntmean(qωᵢ).ω
    end
end

function aux_prior(lik::LaplaceLikelihood, y::AbstractVector{<:Real})
    return For(length(y)) do _
        NTDist(aux_prior(lik))
    end
end

aux_prior(lik::LaplaceLikelihood) = InverseGamma(1//2, laplace_λ(lik))

function aux_kldivergence(::LaplaceLikelihood, qΩ::For, pΩ::For)
    return mapreduce(+, marginals(qΩ), marginals(pΩ)) do qωᵢ, pωᵢ
        μ = mean(dist(qωᵢ))
        λ = scale(dist(pωᵢ))
        log(2λ) / 2 - log(2π) / 2 - log(λ) / 2 + loggamma(1//2) + λ / μ
    end
end


# Generate the data sets used in the comparisons of mix-SQP and the
# first-order methods (EM and projected gradient descent).

# SCRIPT PARAMETERS
# -----------------
n        = 20000;
m        = 20;
filename = "simdata-n=20000-m=20.csv";

# SET UP ENVIRONMENT
# ------------------
using Random
using Distributions
using DelimitedFiles
include("../code/datasim.jl");
include("../code/likelihood.jl");

# Initialize the sequence of pseudorandom numbers.
Random.seed!(1);

# SIMULATE DATA
# -------------
# Generate the n x m conditional likelihood matrix.
z  = normtmixdatasim(n);
sd = autoselectmixsd(z,nv = m);
L  = normlikmatrix(z,sd = sd);

# Save the matrix to a CSV file.
writedlm(filename,L,',');


function DerivativeMatrixSingle(OrdPoly)
(wmat,xw)=GaussLobattoQuad(OrdPoly);
DS=zeros(OrdPoly+1,OrdPoly+1);
for i=1:OrdPoly+1
  for j=1:OrdPoly+1
    DS[i,j]=DLagrange(xw[i],xw,j);
  end
end
w = vec(wmat)
DW=-inv(diagm(w))*DS'*diagm(w);
DV=2*DS;
DV[1,1]=2*DS[1,1]+1/w[1];
DV[OrdPoly+1,OrdPoly+1]=2*DS[OrdPoly+1,OrdPoly+1]-1/w[OrdPoly+1];
DVT=DV';
B=zeros(OrdPoly+1,OrdPoly+1);
B[1,1]=-1/w[1];
B[OrdPoly+1,OrdPoly+1]=1/w[OrdPoly+1];
return (DW,DS,DV,DVT,B)
end


module ScenTreesMakie

using Plots
import Plots.plot
import Plots.plot!

using StatsBase
using Distributions, Random
rng = MersenneTwister(01012019);
using Statistics, LinearAlgebra
using ProgressMeter
using JuMP, Clp
gr()

include("TreeStructure.jl")
include("LatticeStructure.jl")
include("Approx.jl")
include("StochPaths.jl")

include("KernelDensityEstimation.jl")
include("plots.jl")
include("helper.jl")
include("Wasserstein.jl")

export Tree, Lattice,
        tree_approximation!,lattice_approximation,

        stage,height,leaves,nodes,root, children, part_tree,build_probabilities!, checkTree,
        structure, part_lattice, combine_tree,
        gaussian_path, running_maximum, kernel_scenarios,
        sample_path,  plot_path!, plot, plot!,

        path_ident,

        ontoSimplex!, summary,

        distFunction, Wasserstein, Sinkhorn, nestedWasserstein

end


function generate_sphere_system(n, m)
	x0 = zeros(n)

	Rs = rand(m) .+ 1
	centers = zeros(n, m)

	for i = 1:m
	    dir = randn(n)
	    normalize!(dir)
	    centers[:,i] = x0 + Rs[i]*dir
	end

	sphere(x, c, R) = dot(x, x) + dot(c, c) - 2*dot(x, c) - R^2

	function c!(cval, x)
	    for i = 1:m
	        cval[i] = sphere(x, view(centers, :, i), Rs[i])
	    end
	end

	function jac!(J, cval, x)
		fill!(J, 0.0)

	    for i = 1:m
	    	cval[i] = sphere(x, view(centers, :, i), Rs[i])
	    	J[i,:] .= 2.0.*(x .- view(centers, :, i))
	    end
	end

	return x0, c!, jac!
end

function generate_sin_system(n, m)
	x0 = zeros(n)

	function c!(cval, x)
	    for i = 1:m
	        cval[i] = x[2*i] - sin(x[2*i-1])
	    end
	end

	function jac!(J, cval, x)
		fill!(J, 0.0)

	    for i = 1:m
	    	cval[i] = x[2*i] - sin(x[2*i-1])
	    	J[i, 2*i] = 1.0
	    	J[i, 2*i-1] = -cos(x[2*i-1])
	    end
	end

	return x0, c!, jac!
end

@testset verbose=true "Retractions                 " begin

	n = 1000
	m = 100

	x0, c!, jac! = generate_sin_system(n, m)

	# do SVD
	cval = zeros(m)
	cval2 = zeros(m)
	J = zeros(m,n)
	jac!(J, cval, x0)

	U, S, Vt = svd(Array(J'))
	Vt = Array(Vt);

	nr = LFPSQP.NR(U, S, Vt, 1.0, 1000, LFPSQP.NRWork(m), false, LFPSQP.InequalityData())
	pp = LFPSQP.ProjPenalty(jac!, U, S, Vt, m, 0.01, 1.0, 100, 200, LFPSQP.ProjPenaltyWork(m, n, m, n), false, LFPSQP.InequalityDecomp(), LFPSQP.InequalityData())

	# generate random step and project onto tangent plane
	step = randn(n)
	step -= U*(U'*step)
	step .*= 5.0 / norm(step)

	xtilde = x0 + step
	xtilde_copy = copy(xtilde)
	xnew = zeros(n)

	c!(cval, xtilde)
	# @show norm(cval)

	@testset "Newton retraction" begin
		for tol in [1e-6, 1e-8]
			nr.tol = tol

			flag, i, _ = LFPSQP.retract!(cval, xnew, c!, xtilde, x0, nr)
			c!(cval2, xnew)

			# @show flag, i
			@test flag == 0 && norm(cval, Inf) < tol
			@test all(cval .== cval2)			# cval should be c! evaluated at xnew
			@test all(xtilde .== xtilde_copy)
			@test abs(dot(step, xnew - xtilde)) < 1e-6
		end
	end

	@testset "CG iteration" begin
		# solve (μI + J'J)x = b
		J = randn(m, n)
		p = zeros(n)
		z = zeros(n)
		tmp_m = zeros(m)
		tol = 1e-6
		maxiter = 100

		for μ in [1e-1, 1e-2, 1e-4]
			x = zeros(n) # initial guess
			b = randn(n)
			r = copy(b)

			M! = LFPSQP.no_precondition

			flag, i = LFPSQP.pcg!(μ, J, M!, x, r, p, z, tmp_m, tol, maxiter)

			@test flag == 0
			@test norm(r) < tol
			@test norm(μ.*x + J'*(J*x) - b) < tol

			# test exact preconditioning
			x = zeros(n) # initial guess
			b = randn(n)
			r = copy(b)

			exact_inverse = inv(μ*Matrix(I, n, n) + J'*J)
			M!(z, r) = z .= (μ*Matrix(I, n, n) + J'*J) \ r

			flag, i = LFPSQP.pcg!(μ, J, M!, x, r, p, z, tmp_m, tol, maxiter)

			@test flag == 0 && i == 1
			@test norm(r) < tol
			@test norm(μ.*x + J'*(J*x) - b) < tol
		end
	end

	@testset "Project penalty retraction" begin
		for tol in [1e-6, 1e-8, 1e-10]
			pp.tol = tol

			flag, i, pcg_i = LFPSQP.retract!(cval, xnew, c!, xtilde, x0, pp)
			c!(cval2, xnew)

			# @show flag, i, pcg_i
			@test flag == 0 && norm(cval, Inf) < tol
			@test all(cval .== cval2)			# cval should be c! evaluated at xnew
			@test all(xtilde .== xtilde_copy)
			@test norm(step) >= norm(xnew - x0) - tol
		end
	end
end


module day17

data = Array{Bool}(hcat(map(line -> map(c -> c == '#', collect(line)), readlines("day17_input.txt"))...)')
steps = 6

# Part 1
function simulate(current)
    next = copy(current)
    region = CartesianIndices(current)
    ifirst, ilast = first(region), last(region)
    iunit = oneunit(ifirst)
    for i in region
        neighbours = [current[j] for j in max(ifirst, i - iunit):min(ilast, i + iunit) if j != i]
        if (current[i] && count(neighbours) ∉ [2, 3]) next[i] = false end
        if (!current[i] && count(neighbours) == 3) next[i] = true end
    end
    next
end

function runall(dim)
    grid = zeros(Bool, map(+, (size(data)..., fill(1, dim-2)...), tuple(fill(steps*2, dim)...)))
    s = steps + 1
    e = s + first(size(data)) - 1
    grid[s:e, s:e, fill(s:s, dim-2)...] = data
    for i in 1:steps
        grid = simulate(grid)
    end
    grid
end
println("Part 1: active count = $(sum(runall(3)))")

# Part 2
println("Part 2: active count = $(sum(runall(4)))")

end


import .DayCounts: yearfraction, Actual365, DateRoll, adjust
import .TermStructure: YieldTermStructure, ConstantContinuousYieldCurve, discount, startdate, frequency, compounding, daycount

struct SingleStock <: Contract
end

abstract type AbstractCoreModel <: AbstractModel

end

"""
    CoreModel(startdate, startprice, interestrate, carryrate)

Contains the core parameters for the model of `SingleStock`:
* `startprice`: initial price at `startdate`
* `yieldcurve`: risk free rate of return for the numeraire.
* `carrycurve`: the carry rate, i.e. the net income for holding the asset:
  - for stocks this is typically positive (i.e. dividends)
  - for commodities this is typically negative (i.e. cost-of-carry)

The `yieldcurve` and `carryrate` are all specified on a continously compounded, Actual/365 basis.

The exact dynamics of the model are unspecified, but has the property that

```math
E[S_t] = S_0 \\times e^{(r-q)t}
```

"""
struct CoreModel{T, S<:YieldTermStructure, R<:YieldTermStructure} <: AbstractCoreModel
    startprice::T
    yieldcurve::S
    carrycurve::R
end

CoreModel(startprice::T, yieldcurve::S, carryrate::Float64=0.0) where {T, S} =
    CoreModel(startprice, yieldcurve,
              ConstantContinuousYieldCurve(daycount(yieldcurve), carryrate, startdate(yieldcurve)))

CoreModel(startdate::Date, startprice::T, yieldrate::V, carryrate::Float64=0.0) where {T,V} =
    CoreModel(startprice,
              ConstantContinuousYieldCurve(Actual365(), yieldrate, startdate),
              ConstantContinuousYieldCurve(Actual365(), carryrate, startdate))

numeraire(m::CoreModel) = unit(m.startprice)
startdate(m::CoreModel) = startdate(m.yieldcurve)

yearfractionto(m::CoreModel, dt::Date) =
    yearfraction(daycount(m.yieldcurve),  startdate(m.yieldcurve),  dt)

function value(m::CoreModel, c::WhenAt{Receive{T}}) where T
    value(m, c.c) * discount(m.yieldcurve, maturitydate(c))
end

value(m::CoreModel, c::SingleStock) = m.startprice
function value(m::CoreModel, c::WhenAt{SingleStock})
    m.startprice * discount(m.carrycurve, maturitydate(c))
end

function forwardprice(m::CoreModel, ::SingleStock, dt::Date)
    m.startprice * discount(m.carrycurve, dt) / discount(m.yieldcurve, dt)
end

"""A model for the time value of money
"""
struct YieldModel{S<:YieldTermStructure, T<:DateRoll, U<:HolidayCalendar} <: AbstractModel
    yieldcurve::S
    dateroll::T
    holidaycalendar::U
end

value(m::YieldModel, c::When{At, Receive{T}}) where {T} = value(m, c.c) * discount(m.yieldcurve, adjust(m.dateroll, m.holidaycalendar, maturitydate(c)))


const g = @SMatrix [1  0  0  0;
                    0 -1  0  0;
                    0  0 -1  0;
                    0  0  0 -1] # Minkowski space metric

function dynEq(v,type,τ)
    x = v[1:4]
    u = v[5:8]

    q = type.constants.q
    mₑ = type.constants.mₑ

    dx = u
    du = q/mₑ * Fμν(x, type) * g * u

    return vcat(dx,du)
end

function Mμν(v)
    x⁰, x¹, x², x³, u⁰, u¹, u², u³ = v
    Lx = x²*u³ - x³*u²
    Ly = x³*u¹ - x¹*u³
    Lz = x¹*u² - x²*u¹
    # Nx = x¹*u⁰ - x⁰*u¹
    # Ny = x²*u⁰ - x⁰*u²
    # Nz = x³*u⁰ - x⁰*u³
    L² = Lx^2 + Ly^2 + Lz^2
    # N² = Nx^2 + Ny^2 + Nz^2
    #the physical quantities are m* [1,2,3] , m^2*[4], m/c* [5,6,7] , m^2/c^2*[8],
    return (Lx = Lx, Ly = Ly, Lz = Lz, L² = L², #Nx = Nx, Ny = Ny, Nz = Nz, N² = N²,
            x⁰ = x⁰, x¹ = x¹, x² = x², x³ = x³, u⁰ = u⁰, u¹ = u¹, u² = u², u³ = u³)
end

function σA(A,A²)
    .√abs.((A²-A.^2))
end

function Mμντ(X)
    [Mμν(xτ) for xτ in X.u]
end

function UniformAnnulus(r₁,r₂)
    θ = rand(Uniform(0,2*π))
    u = rand(Uniform(r₁,r₂))
    v = rand(Uniform(0,r₁+r₂))
    ρ = v<u ? u : r₁+r₂ - u
    [SVector{3}(ρ*cos(θ), ρ*sin(θ),0), ρ]
end

R = [UniformAnnulus(0,maxR) for i in 1:N]
RVec = VectorOfArray(R)
idx = sortperm(R, lt=(x,y)->x[2]<y[2])
R₀ = RVec[1,:][idx]
ρ₀ = RVec[2,:][idx]

function xu_cat(X,U,T)
    x0 = c*T
    x = vcat([x0],X)
    u0 = √(c*c + U⋅U)
    u = vcat([u0],U)
    SVector{8}(vcat(x,u))
end

function problem_set(R₀, p, U₀ = [0,0,0])
    V₀ = [xu_cat(r₀,U₀,τi) for r₀ in R₀]
    prob = ODEProblem(dynEq, V₀[1], (τi, τf), p)

    function prob_func(prob,i,repeat)
        ODEProblem(dynEq, V₀[i], (τi, τf), p) #return
    end

    function output_func(sol, i)
        (Mμντ(sol), false) # put sol as argument, if you want the actual trajectories
    end

    EnsembleProblem(prob; prob_func, output_func)
end

# new
function Moment_Avg(M, J; moment_func = identity, f = 1, l = N)
    N_samples = 1 + l - f
    func(i,t) = moment_func(getproperty.((M[i][t], ), J)...)
    [ThreadsX.sum(func(i,t) for i in f:l)/N_samples for t in 1:time_samples+1]
end

# function Radial_Avg(M, args_funcs, bins = 100) # J = (:Lz, (:x¹, :x²))
# data_points = zeros(bins, 1 + sum(length.(args_funcs)) - length(args_funcs))

# particle_idx = 1
# current_M = M[particle_idx][end] # current_particle data at the final time
# current_ρ = ρ₀[particle_idx] # ⟺ hypot(m[begin].x¹, m[begin].x²), but faster since in memory already

#     for i in 1:bins
#         bin_averages = zeros(1 + 2length(args_funcs))
#         tally = 0
#         while current_ρ < i * maxR/bins
#             bin_averages[1]     += current_ρ
#             bin_averages[2:end] += [func(getproperty.((current_M,), args[begin])...)
#                                                         for args in args_funcs
#                                                             for func in args[2:end]]
#             tally += 1
#             try
#                 Midx += 1
#                 current_M = M[particle_idx][end]
#                 current_ρ = ρ₀[particle_idx]
#             catch err
#                 # isa(err, BoundsError)
#                 break
#             end
#         end
#         data_points[i,:] = bin_averages/tally
#     end
#     return data_points
# end

# args_funcs = [((:Lz,), identity, x -> x^2), ((:x¹, :x²), hypot, (x,y) -> x^2 + y^2)]

function Radial_Avg(M, bins = 100)
    LzStats = zeros(bins, 3)
    ρStats  = zeros(bins, 3)
    Midx = 1
    m = M[Midx][end]
    ρ = ρ₀[Midx]
    for i in 1:bins
        ρ_avg, Lz_, Lz²_, ρ_, ρ²_ = 0., 0., 0., 0., 0.
        tally = 0
        while ρ < i * maxR/bins
            ρ_avg += ρ
            Lz_ += m.Lz
            Lz²_ += m.Lz^2
            ρ_ += hypot(m.x¹, m.x²)
            ρ²_ += m.x¹^2 + m.x²^2
            tally += 1
            try Midx < N
                Midx += 1
                m = M[Midx][end]
                ρ = ρ₀[Midx]
            catch err
                break
            end
        end
        rt = ρ_avg/tally
        LzStats[i,:] .= rt, Lz_/tally, Lz²_/tally
        ρStats[i,:] .= rt, ρ_/tally,  ρ²_/tally
    end
    return LzStats, ρStats
end

common_settings = (fontfamily = "computer modern",
                   dpi = 360,
                   label = false,
                   fillalpha=.3,
                   color_palette = :Set1_4,
                   linewidth = 2)

for type in lasers
    pol = polarization(type)

    folder_setup("$(pol)_p$(type.p)_m$(type.m)") # change folder for each simulation
    folder_setup("particles") # separate particles from fields

    file_ext = "$(pol)_p$(type.p)_m$(type.m).$fmt"

    write(io, "$(pol), p = $(type.p), m = $(type.m)\n")

    plot_data = [plot([],[], label = false) for i in 1:8]

    eprob = problem_set(R₀, type)
    M = solve(eprob, Vern9(), EnsembleThreads(), abstol = 1e-15, reltol = 1e-15,
            saveat = (τf-τi)/time_samples, trajectories = N)

    roots = rootsLUT[(p = type.p, m = abs(type.m))]
    idxs = zeros(Int, length(roots))                   # initialize index list

    root_idx = 2                                       # first root is always zero, start from root 2
    for (i,v) in enumerate(ρ₀)                         # ρ₀ is ordered!
        try
            if v > w₀*roots[root_idx]
                idxs[root_idx] = i-1                   # take the last one which respects the condition
                root_idx += 1
            end
        catch err
            # isa(err, BoundsError)
            break
        end
    end

    for j in 2:(length(roots))
        z,  Lz,  L² = Moment_Avg.((M,), (:x³, :Lz, :L²); f = 1+idxs[j-1], l = idxs[j])
        z², Lz², L⁴ = Moment_Avg.((M,), (:x³, :Lz, :L²); moment_func = x -> x^2, f = 1+idxs[j-1], l = idxs[j])
        ρ  = Moment_Avg(M, (:x¹, :x²); moment_func = hypot, f = 1+idxs[j-1], l = idxs[j])
        # ρ² = Moment_Avg(M, (:x¹, :x²); moment_func = (x, y) -> x^2 + y^2, f = 1+idxs[j-1], l = idxs[j])

        σz  = σA(z, z²)/w₀
        # σρ  = σA(ρ, ρ²)/w₀ # not worth computing because it's too wide and breaks plots
        σLz = σA(Lz, Lz²)/w₀/π₀
        σL² = σA(L², L⁴)/w₀/π₀/w₀/π₀

        write(io,
            "from $(roots[j-1]/w₀) to $(roots[j]/w₀)\n"*
            "z ± σz $(z[end]/w₀) ± $(σz[end])\n"*
            "Lz ± σLz $(Lz[end]/w₀/π₀) ± $(σLz[end])\n"*
            "L² ± σL² $(L²[end]/w₀/π₀/w₀/π₀) ± $(σL²[end])\n"
             )

        plot_data[1] = plot!(plot_data[1], Δτ, z/λ;
                                common_settings...,
                                ribbon = σz,
                                seriescolor = j - 1, # these cannot be moved outside due to variable scope
                                xlabel = L"\tau/T_0",
                                ylabel = L"\overline{z}/\lambda")

        # the average radius for particles homogeneously distributed in the annulus between root[i] and root[i+1]
        mean_ρ = 2/3*(roots[j]^3 - roots[j-1]^3)/(roots[j]^2 - roots[j-1]^2)

        plot_data[2] = plot!(plot_data[2], Δτ, ρ/w₀ .- mean_ρ;
                                common_settings...,
                                # ribbon = σρ, too big relative to ρ changes, makes plot unreadable
                                seriescolor = j - 1,
                                xlabel = L"\tau/T_0",
                                ylabel = L"\Delta\overline{\rho}/w_0")

        plot_data[3] = plot!(plot_data[3], Δτ, Lz/w₀/π₀;
                                common_settings...,
                                ribbon = σLz,
                                seriescolor = j - 1,
                                xlabel = L"\tau/T_0",
                                ylabel = L"\overline{L_z}/(w_0 \pi_0)")

        plot_data[4] = plot!(plot_data[4], Δτ, L²/w₀/π₀/w₀/π₀;
                                common_settings...,
                                ribbon = (L²/w₀/π₀/w₀/π₀, σL²),
                                seriescolor = j - 1,
                                xlabel = L"\tau/T_0",
                                ylabel = L"\overline{L^2}/(w_0^2 \pi_0^2)")

    end
    write(io,"\n")

    add_info!(plot_data[1], type)
    savefig(plot_data[1], "z_t_$file_ext")
    add_info!(plot_data[2], type)
    savefig(plot_data[2], "rho_t_$file_ext")
    add_info!(plot_data[3], type)
    savefig(plot_data[3], "Lz_t_$file_ext")
    add_info!(plot_data[4], type)
    savefig(plot_data[4], "L2_t_$file_ext")

    X = [M[i][begin].x¹ for i in 1:N]/w₀ # pulling out the initial x coordinates
    Y = [M[i][begin].x² for i in 1:N]/w₀ # pulling out the initial y coordinates
    Lz = [M[i][end].Lz for i in 1:N] # pulling out the final Lz for colouring
    Lzmax = max(abs.(extrema(Lz))...) # finding maximum
    # colours = Lz/Lzmax # doing a normalization

    write(io, "Maximum Lz transfered = $Lzmax\n\n")

    plot_data[7] = scatter(X,Y;
        marker_z = Lz,
        markershape = circ,
        markerstrokealpha = 0,
        markersize = 1.2,
        seriescolor = cgrad(:PuOr_9, rev = true),
        clims = (-Lzmax, Lzmax),
        fontfamily = "computer modern",
        dpi = 360,
        label = false,
        xlabel = L"x/w_0", ylabel = L"y/w_0",
        aspect_ratio = 1,
        xticks = -4:1:4,
        right_margin = 0.9*Plots.cm,)
    add_info!(plot_data[7], type, true)
    savefig(plot_data[7], "XYLz_$file_ext")

    @time begin
    particle_range = 1:50:N
    Xt  = [[M[i][t].x¹/w₀ for t in 1:time_samples] for i in particle_range]
    Yt  = [[M[i][t].x²/w₀ for t in 1:time_samples] for i in particle_range]
    Zt  = [[M[i][t].x³/λ  for t in 1:time_samples] for i in particle_range]
    Lzt = [[M[i][t].Lz    for t in 1:time_samples] for i in particle_range]
    plot_data[8] = plot(Xt, Yt, Zt;
        line_z = VectorOfArray(Lzt),
        right_margin = 0.9*Plots.cm,
        seriescolor = cgrad(:PuOr_9, rev = true),
        xlabel = L"x/w_0",
        ylabel = L"y/w_0",
        zlabel = L"z/\lambda",
        common_settings...)
    # add_info!(plot_data[8], type) broken for 3D plots
    savefig(plot_data[8], "XYZ_$file_ext")
    end

    LzS, ρS = Radial_Avg(M)

    plot_data[5] = plot(LzS[:,1]/w₀, LzS[:,2]/π₀/w₀,
                        label = false,
                        ribbon = σA(LzS[:,2], LzS[:,3])/π₀/w₀)
    plot_data[5] = add_rect!(plot_data[5], LzS[:,1]/w₀, LzS[:,2]/π₀/w₀, roots;
                        common_settings...,
                        ribbon = σA(LzS[:,2], LzS[:,3])/π₀/w₀,
                        seriescolor = 4,
                        xlabel = L"\rho_0/w_0",
                        ylabel = L"\overline{L_z}(\rho_0)/(w_0 \pi_0)",
                        )
    add_info!(plot_data[5], type)
    savefig(plot_data[5], "Lz_rho_$file_ext")

    plot_data[6] = plot(ρS[:,1]/w₀, ρS[:,2]/w₀ .- ρS[:,1]/w₀;
                        label = false,
                        ribbon = σA(ρS[:,2], ρS[:,2])/w₀)
    plot_data[6] = add_rect!(plot_data[6],
                        ρS[:,1]/w₀, ρS[:,2]/w₀ .- ρS[:,1]/w₀, roots;
                        common_settings...,
                        ribbon = σA(ρS[:,2], ρS[:,3])/w₀,
                        seriescolor = 4,
                        xlabel = L"\rho_0/w_0",
                        ylabel = L"\Delta\overline{\rho}_{\mathrm{final}}(\rho_0)/w_0")
    add_info!(plot_data[6], type)
    savefig(plot_data[6], "rho_rho_$file_ext")

    serialize("plots.jls", plot_data)
    serialize("radial_ρ_data.jls", LzS)
    serialize("radial_Lz_data.jls", LzS)
    # serialize("plot_data.jls", M) save at your own peril, too large to be read efficiently
    cd("../..") # return to output_* base folder
end


mutable struct LinearGlobalFunctionApproximator{W <: AbstractArray} <: GlobalFunctionApproximator
    weights::W
end

function fit!(lgfa::LinearGlobalFunctionApproximator, dataset_input::AbstractMatrix{T},
              dataset_output::AbstractArray{T}) where T
    # TODO: Since we are ASSIGNING to weights here, does templating even matter? Does the struct even matter?
    lgfa.weights = llsq(dataset_input, dataset_output, bias=false)
end

function compute_value(lgfa::LinearGlobalFunctionApproximator, v::AbstractArray{T}) where T
    return dot(lgfa.weights, v)
end


using Test
using LinearAlgebra
using BEAST
using CompScienceMeshes
using StaticArrays

y = point(2,0,0)
j = ẑ
κ = 1.0

gj(x) = exp(-im*κ*norm(x-y))/(4*π*norm(x-y)) * j
ccg = BEAST.CurlCurlGreen(κ, j, y)

function curlh(f,x,h)
    e1 = point(1,0,0)
    e2 = point(0,1,0)
    e3 = point(0,0,1)

    d1f = (f(x+h*e1) - f(x-h*e1))/(2*h)
    d2f = (f(x+h*e2) - f(x-h*e2))/(2*h)
    d3f = (f(x+h*e3) - f(x-h*e3))/(2*h)

    return @SVector[
        d2f[3] - d3f[2],
        d3f[1] - d1f[3],
        d1f[2] - d2f[1]]
end

cgh(x) = curlh(gj,x,h)
ccgh(x) = curlh(cgh,x,h)

h = 0.01
x = point(1,1,1)
a = ccg(x)
b = ccgh(x)
@show norm(x-y)
@test norm(a-b) < 1e-5

cccg = curl(ccg)
cccgh(x) = curlh(ccg,x,h)

# @show a = cccg(x)
# @show b = cccgh(x)
@test norm(a-b) < 1e-5


function esf_gen!(S::AbstractArray{T,1}, x::Array{Array{T,1},1}, a::Array{Array{D,1},1}, m::Array{D,1}, ni::Array{D,1}) where {D <: Integer, T <: Real}
  n = length(x)
  fill!(S,zero(T))
  S[1] = one(T)
  @inbounds for col in 1:n
    for row in m[col]:-1:1
      for k in 1:ni[col]
        if row - a[col][k] >= 0
          S[row+1] += x[col][k] * S[row - a[col][k] + 1]
        end
      end
    end
  end
end

function esf_gen(x::AbstractArray{T,1}, a::AbstractArray{D,1}, lengths::AbstractArray{D,1}) where {D <: Integer, T <: Real}
  n = cumsum(lengths)
  m = cumsum(map(maximum,a))

  S = Vector{T}(undef,m[end]+1)
  esf_gen!(S,x,a,n,m)
  return S
end


using Random
using StatsBase
using Distributions
using StatsPlots
using DataFrames
using ShiftedArrays
using Measures

Random.seed!(123)

begin
    arrival_rates = (0.1, 0.1)
    job_size = 10
    private_service_rate(job_size) = 1 / (1.5 * job_size)
    public_service_rate(job_size) = 1 / (0.7 * job_size)
    p1_prob_public = 0.576
    p2_prob_public = 0.576
    segment_time = 10000
    total_time = 10000
end

function run_simulation(arrival_rates, job_size, private_service_rate, public_service_rate, p1_prob_public, p2_prob_public; total_time = 1000, return_ft = false)
    function generate_interarrival_times(d::Distribution, time)
        interarrival_times = [rand(d)]
            # initial_n = time/mean(d)
            # append!(interarrival_times, rand(d, initial_n))
        while sum(interarrival_times) <= time
            append!(interarrival_times, rand(d))
        end
        return interarrival_times
    end

    job_interarrival_times = generate_interarrival_times(Exponential(1 / sum(arrival_rates)), total_time)
    total_jobs = length(job_interarrival_times)
    job_which_player = 1 .+ rand(Bernoulli(arrival_rates[1] / sum(arrival_rates)), total_jobs)
    df = DataFrame(order = 1:total_jobs, interarrival_times = job_interarrival_times, player = job_which_player)
    df.public = @. ifelse(df.player == 1, rand(Bernoulli(p1_prob_public)), rand(Bernoulli(p2_prob_public)))
    df.service_duration = @. ifelse(
        df.public,
        rand(Exponential(1 / public_service_rate(job_size))),
        rand(Exponential(1 / private_service_rate(job_size)))
    )
    df.arrival_time = cumsum(df.interarrival_times)
    df.finished_time = df.arrival_time .+ df.service_duration

    # df_public = filter(a -> a.public, df)
    df_public = subset(df, :public, view = true)
    df_private = subset(df, :public => ByRow(!=(true)), view = true)
    df_private1 = subset(df_private,  :player => ByRow(==(1)), view = true)
    df_private2 = subset(df_private,  :player => ByRow(==(2)), view = true)

    function correct_finished_time!(df)
        df_return = view(df, :, :)
        @inbounds for i = 2:size(df_return, 1)
            if lag(df_return.finished_time)[i] > df_return.arrival_time[i]
                df_return[i, :finished_time] = lag(df_return.finished_time)[i] + df_return.service_duration[i]
            end
        end
        return df_return
    end

    df_public = correct_finished_time!(df_public)
    df_private1 = correct_finished_time!(df_private1)
    df_private2 = correct_finished_time!(df_private2)

    function make_events(arrivals_df)
        df = arrivals_df
        function f1(df, column, event_result, length_df = size(df)[1])
            times_df = select(df, [:player, :public, column])
            rename!(times_df, column => :time)
            times_df[!, :event] = repeat([event_result], length_df)
            return times_df
        end

        arrivaltimes_df = f1(arrivals_df, :arrival_time, 1)
        servicetimes_df = f1(arrivals_df, :finished_time, -1)

        return vcat(arrivaltimes_df, servicetimes_df)
    end

    if return_ft
        return select(df, [:player, :public, :arrival_time, :service_duration, :finished_time])
    end

    df_events = sort(make_events(df), :time)
    df_events_simple = select(df_events, [:time, :event])
    transform!(df_events_simple,
        :event => cumsum => :cumevent,
        :time => Base.Fix2(lag, -1) => :endtime,
    )
    df_events_simple[:, :group] .= -1
    transform!(df_events_simple,
        :cumevent => Base.Fix2(lag, -1) => :nextevent
    )
    transform!(df_events, [:player, :public] => ((x, y) -> ifelse.(y, 0, x)) => :group)
    gby = groupby(df_events, :group)
    transform!(gby,
        :event => cumsum => :cumevent,
        :time => Base.Fix2(lag, -1) => :endtime,
    )
    transform!(gby, :cumevent => Base.Fix2(lag, -1) => :nextevent)
    return(df_events_simple, gby)
end

df_events_simple, gby = run_simulation(arrival_rates, job_size, private_service_rate, public_service_rate, p1_prob_public, p2_prob_public)
gby2 = gby
gby = DataFrame(gby)
dfs = vcat(gby, df_events_simple, cols=:union)

colour_dict = Dict(-1 => :black, 0 => :red, 1 => :blue, 2 => :green)
label_dict = Dict(-1 => "Total", 0 => "Public", 1 => "Player 1 Private", 2 => "Player 2 Private")
translate(x) = colour_dict[x]
dfs[:, :groupcolour] = translate.(dfs[!, :group])

function plot_path(df, end_time)
    p = plot(title = "Simulation of a queue", xlabel = "Time", ylabel = "Queue length")
    @inbounds for i in 1:(size(df)[1] - 1)
        plot!(p,
            [df[i, :time], df[i, :endtime]],
            [df[i, :cumevent], df[i, :cumevent]],
            colour = df[i, :groupcolour],
            label = false
        )
        plot!(p,
            [df[i, :endtime], df[i, :endtime]],
            [df[i, :cumevent], df[i, :nextevent]],
            colour = df[i, :groupcolour],
            label = false)
    end
    tempdf = combine(groupby(dfs, :group)) do sdf; sdf[1, :] end
    tempdf[:, :endtime] = tempdf[:, :time]
    tempdf[:, :time] .= 0
    @inbounds for i in 1:(size(tempdf)[1])
        plot!(p,
            [tempdf[i, :time], tempdf[i, :endtime]],
            [0, 0],
            colour = tempdf[i, :groupcolour],
            label = label_dict[tempdf[i, :group]]
        )
        plot!(p,
            [tempdf[i, :endtime], tempdf[i, :endtime]],
            [0, 1],
            colour = tempdf[i, :groupcolour],
            label = false)
    end

    tempdf2 = combine(groupby(dfs, :group)) do sdf; sdf[end, :] end
    final_time = dfs[argmax(skipmissing(dfs[:, :endtime])), :endtime]
    tempdf2[:, :endtime] .= final_time
    @inbounds for i in 1:(size(tempdf2)[1])
        plot!(p,
            [tempdf2[i, :time], tempdf2[i, :endtime]],
            [0, 0],
            colour = tempdf2[i, :groupcolour],
            label = false
        )
    end
    xlims!(p, (0, end_time))
        return p
end

# p2 = plot_path(dfs, total_time)

@df dfs plot(
    :time,
    :cumevent,
    group = :groupcolour
)

begin
JSs = 0.1:0.1:1
AR_JS_ratio = 1
ARs = AR_JS_ratio ./ JSs

# res = DataFrame(player = Int64[], public = Bool[], JS = Float64[], mean_response_time = Float64[])
resval_mean = Float64[]
resval_var = Float64[]
for (AR, JS) in zip(ARs, JSs)
    temp_mean = Float64[]
    temp_var = Float64[]
    print("AR: ", round(AR, digits = 4), ", JS: ", round(JS, digits = 4))
    num_repetitions = JS * 10
    for i in 1:num_repetitions
        if mod(i, floor(num_repetitions / 5)) == 0 print(".") end
        _, gby = run_simulation(
            AR,
            JS,
            private_service_rate,
            public_service_rate,
            p1_prob_public,
            p2_prob_public,
            total_time = 10000
        )
        gby = DataFrame(gby)
        sort!(gby, [:player, :public, :time])

        arrivals = subset(gby, :event => ByRow(isequal(1)))
        services = subset(gby, :event => ByRow(isequal(-1))) |>
            x -> select(x, [:time]) |>
            x -> rename(x, :time => :dep_time)

        gby = hcat(arrivals, services) |>
            x -> select(x, [:player, :public, :time, :dep_time]) |>
            x -> groupby(x, [:player, :public]) |>
            x -> transform(x, [:time, :dep_time] => ((x, y) -> (y - x)) => :response_time) |>
            x -> combine(x, :response_time => (x -> (mean_RT = mean(skipmissing(x)), sd_RT = var(skipmissing(x)))) => AsTable)
        append!(temp_mean, gby[1, 1])
        append!(temp_var, gby[1, 2])
    end
    println()
    append!(resval_mean, mean(temp_mean))
    append!(resval_var, mean(temp_var))
end
end
### Mean and Variance of Response Time Plots ###
begin
include("../helper_functions.jl")
function variance_of_cost(p, q, λ1, λ2, α, μ, ν)
    (p * public_cost(p, q, λ1, λ2, α, μ, ν))^2 + ((1 - p) * private_cost(p, λ1, α, μ))^2
end

Hyperexponential(parameters, weights) = MixtureModel(Exponential, parameters, weights)
RTDist(p, q, λ1, λ2, α, μ, ν) = Hyperexponential([(1/(ν(α) - λ1*p - λ2*q)), (1/(μ(α) - λ1*(1-p)))], [p, 1 - p])
RT = RTDist(0.57, 0.57, 1, 1, 1, private_service_rate, public_service_rate)
RTs = RTDist.(0.57, 0.57, ARs, ARs, JSs, private_service_rate, public_service_rate)
plot_mean = plot(JSs, resval_mean, xlabel = "Job Size", ylabel = "Mean Response Time", label = "Empirical Mean")
plot!(JSs,
    mean.(RTs),
    title = "Mean Response Time by Job Size",
    xlabel = "Job Size",
    ylabel = "Mean Response Time",
    label = "Theoretical Mean",
    size = [400, 300],
    legend = :topleft)

plot_variance = plot(JSs, resval_var, xlabel = "Job Size", ylabel = "Variance of Response Time", label = "Empirical Variance")
plot!(JSs,
    var.(RTs),
    title = "Variance by Job Size",
    xlabel = "Job Size",
    ylabel = "Variance of Response Time",
    label = "Theoretical Variance",
    size = [400, 300],
    legend = :topleft)

mean_var_plot = plot(plot_mean, plot_variance, size = [800, 300], margin = 12pt)
end

savefig("outputs/simulated_variance.pdf")

plot(JSs, resval_mean ./ cost.(0.57, 0.57, ARs, ARs, JSs, private_service_rate, public_service_rate))
plot(JSs, resval_var ./ variance_of_cost.(0.57, 0.57, ARs, ARs, JSs, private_service_rate, public_service_rate))

cost(0.57, 0.57, 1, 1, 1, private_service_rate, public_service_rate)
cost(0.57, 0.57, 2, 2, 0.5, private_service_rate, public_service_rate)
# sort!(gby, [:player, :public, :time])

# arrivals = subset(gby, :event => ByRow(isequal(1)))
# services = subset(gby, :event => ByRow(isequal(-1))) |>
#     x -> select(x, [:time]) |>
#     x -> rename(x, :time => :dep_time)

# hcat(arrivals, services) |>
#     x -> select(x, [:player, :public, :time, :dep_time]) |>
#     x -> groupby(x, [:player, :public]) |>
#     x -> combine(x, [:time, :dep_time] => ((x, y) -> (mean(y - x))) => :mean_response_time) |>
#     x -> combine(x, :mean_response_time => mean)[1, 1]

C = 1
αs = 0.1:0.01:5
plot(αs, ForwardDiff.derivative.(α -> cost(0.57, 0.57, C/α, C/α, α, private_service_rate, public_service_rate), αs))

sim_df = run_simulation(
            1,
            1,
            private_service_rate,
            public_service_rate,
            p1_prob_public,
            p2_prob_public;
            total_time = 100000,
            return_ft = true
) |>
    x -> transform(x, [:finished_time, :arrival_time] => (-) => "response_time")

# function generate_response_times(N, p, q, λ1, λ2, α, μ, ν)
#     pub = rand(Exponential(1/(ν(α) - λ1*p - λ2*q)), N)
#     pr = rand(Exponential(1/(μ(α) - λ1*(1-p))), N)
#     @. (p * pub) + ((1-p) * pr)
# end

# histogram(generate_response_times(100000, 0.57, 0.57, 1, 1, 1, private_service_rate, public_service_rate))
# cost(0.57, 0.57, 1, 1, 1, private_service_rate, public_service_rate)

density_plot = density(log.(rand(RT, 1000000)), lw = 4, label = "Theoretical Response Distribution")
density!(log.(sim_df.response_time), lw = 4, label = "Empirical Response Distribution")
plot!(xlims = [-6, 4], legend = :topleft, size = (500, 300),
    xlabel = "log(x)",
    ylabel = "f(x)",
    title = "Comparison of Densities")

# savefig("outputs/densities.pdf")

mean(RT)
mean(sim_df.response_time)
var(RT)
var(sim_df.response_time)


# Test complete Lanczos factorization
@testset "Complete Lanczos factorization" begin
    @testset for T in (Float32, Float64, ComplexF32, ComplexF64)
        @testset for orth in (cgs2, mgs2, cgsr, mgsr) # tests fail miserably for cgs and mgs
            A = rand(T,(n,n))
            v = rand(T,(n,))
            A = (A+A')
            iter = LanczosIterator(A, v, orth)
            fact = initialize(iter)
            while length(fact) < n
                expand!(iter, fact)
            end

            V = hcat(basis(fact)...)
            H = rayleighquotient(fact)
            @test normres(fact) < 10*n*eps(real(T))
            @test V'*V ≈ I
            @test A*V ≈ V*H
        end
    end
end

# Test complete Arnoldi factorization
@testset "Complete Arnoldi factorization" begin
    @testset for T in (Float32, Float64, ComplexF32, ComplexF64)
        @testset for orth in (cgs, mgs, cgs2, mgs2, cgsr, mgsr)
            A = rand(T,(n,n))
            v = rand(T,(n,))
            iter = ArnoldiIterator(A, v, orth)
            fact = initialize(iter)
            while length(fact) < n
                expand!(iter, fact)
            end

            V = hcat(basis(fact)...)
            H = rayleighquotient(fact)
            factor = (orth == cgs || orth == mgs ? 100 : 10)
            @test normres(fact) < factor*n*eps(real(T))
            @test V'*V ≈ I
            @test A*V ≈ V*H
        end
    end
end

# Test complete Arnoldi factorization
@testset "Complete Golub-Kahan-Lanczos factorization" begin
    @testset for T in (Float32, Float64, ComplexF32, ComplexF64)
        @testset for orth in (cgs2, mgs2)
            A = rand(T,(n,n))
            v = A*rand(T,(n,)) # ensure v is in column space of A
            iter = GKLIterator(A, v, orth)
            fact = initialize(iter)
            while length(fact) < n
                expand!(iter, fact)
            end

            U = hcat(basis(fact, :U)...)
            V = hcat(basis(fact, :V)...)
            B = rayleighquotient(fact)
            @test normres(fact) < 10*n*eps(real(T))
            @test U'*U ≈ I
            @test V'*V ≈ I
            @test A*V ≈ U*B
            @test A'*U ≈ V*B'
        end
    end
end

# Test incomplete Lanczos factorization
@testset "Incomplete Lanczos factorization" begin
    @testset for T in (Float32, Float64, ComplexF32, ComplexF64)
        @testset for orth in (cgs2, mgs2, cgsr, mgsr) # tests fail miserably for cgs and mgs
            A = rand(T,(N,N))
            v = rand(T,(N,))
            A = (A+A')
            iter = @inferred LanczosIterator(A, v, orth)
            krylovdim = n
            fact = @inferred initialize(iter)
            while normres(fact) > eps(real(T)) && length(fact) < krylovdim
                @inferred expand!(iter, fact)

                V = hcat(basis(fact)...)
                H = rayleighquotient(fact)
                r = residual(fact)
                β = normres(fact)
                e = rayleighextension(fact)
                @test V'*V ≈ I
                @test norm(r) ≈ β
                @test A*V ≈ V*H + r*e'
            end

            fact = @inferred shrink!(fact, div(n,2))
            V = hcat((@inferred basis(fact))...)
            H = @inferred rayleighquotient(fact)
            r = @inferred residual(fact)
            β = @inferred normres(fact)
            e = @inferred rayleighextension(fact)
            @test V'*V ≈ I
            @test norm(r) ≈ β
            @test A*V ≈ V*H + r*e'
        end
    end
end

# Test incomplete Arnoldi factorization
@testset "Incomplete Arnoldi factorization" begin
    @testset for T in (Float32, Float64, ComplexF32, ComplexF64)
        @testset for orth in (cgs, mgs, cgs2, mgs2, cgsr, mgsr)
            A = rand(T,(N,N))
            v = rand(T,(N,))
            iter = @inferred ArnoldiIterator(A, v, orth)
            krylovdim = 3*n
            fact = @inferred initialize(iter)
            while normres(fact) > eps(real(T)) && length(fact) < krylovdim
                @inferred expand!(iter, fact)

                V = hcat(basis(fact)...)
                H = rayleighquotient(fact)
                r = residual(fact)
                β = normres(fact)
                e = rayleighextension(fact)
                @test V'*V ≈ I
                @test norm(r) ≈ β
                @test A*V ≈ V*H + r*e'
            end

            fact = @inferred shrink!(fact, div(n,2))
            V = hcat((@inferred basis(fact))...)
            H = @inferred rayleighquotient(fact)
            r = @inferred residual(fact)
            β = @inferred normres(fact)
            e = @inferred rayleighextension(fact)
            @test V'*V ≈ I
            @test norm(r) ≈ β
            @test A*V ≈ V*H + r*e'
        end
    end
end

# Test incomplete Arnoldi factorization
@testset "Incomplete GKL factorization" begin
    @testset for T in (Float32, Float64, ComplexF32, ComplexF64)
        @testset for orth in (cgs2, mgs2)
            A = rand(T,(N,N))
            v = rand(T,(N,))
            iter = @inferred GKLIterator(A, v, orth)
            krylovdim = 3*n
            fact = @inferred initialize(iter)
            while normres(fact) > eps(real(T)) && length(fact) < krylovdim
                @inferred expand!(iter, fact)

                U = hcat(basis(fact, :U)...)
                V = hcat(basis(fact, :V)...)
                B = rayleighquotient(fact)
                r = residual(fact)
                β = normres(fact)
                e = rayleighextension(fact)
                @test U'*U ≈ I
                @test V'*V ≈ I
                @test norm(r) ≈ β
                @test A*V ≈ U*B + r*e'
                @test A'*U ≈ V*B'
            end

            fact = @inferred shrink!(fact, div(n,2))
            U = hcat((@inferred basis(fact, :U))...)
            V = hcat((@inferred basis(fact, :V))...)
            B = @inferred rayleighquotient(fact)
            r = @inferred residual(fact)
            β = @inferred normres(fact)
            e = @inferred rayleighextension(fact)
            @test U'*U ≈ I
            @test V'*V ≈ I
            @test norm(r) ≈ β
            @test A*V ≈ U*B + r*e'
            @test A'*U ≈ V*B'
        end
    end
end


using TrafficAssignment
using Test

include("test_functions.jl")

test_tntp()

summarize_ta_data()

data_time = time()
ta_data = load_ta_network("SiouxFalls")

println("Data Loading Completed, time:", time() - data_time, " seconds")

link_flow, link_travel_time, objective = ta_frank_wolfe(ta_data, method=:bfw, step=:exact, log=:on, tol=1e-3, max_iter_no=5)

link_flow, link_travel_time, objective = ta_frank_wolfe(ta_data, method=:cfw, step=:newton, log=:on, tol=1e-3, max_iter_no=5)

link_flow, link_travel_time, objective = ta_frank_wolfe(ta_data, method=:fw, step=:exact, log=:on, tol=1e-3, max_iter_no=5)
@assert abs( objective - 4.963799502172654e6 ) < 1e6


using BandedMatrices

function sband(n::Integer, a::Vector{T}) where T
    k = length(a)
    SH = T<:Complex ? Hermitian : Symmetric
    SH(BandedMatrix([-i => fill(a[i+1], n-i) for i = 0:min(k,n)-1]...), :L)
end

function example_1(n, T=Float64)
    tt = sband(n, T[4, 1])
    ss = sband(n, T[2, 1])
    tt, ss
end

det0(A, B, s) = det(Matrix(A - s * B))
function det1(A, B, s)
    h = sqrt(eps(max(abs(s), 1.0)))
    ( det0(A, B, s + h) - det0(A, B, s - h) ) / (2*h)
end
function det2(A, B, s)
    h = eps((max(abs(s), 1)))^(1/3)
    ( det0(A, B, s + h) + det0(A, B, s - h) -2*det0(A, B, s) ) / h^2
end

@testset "basics n=$n s=$s" for n = [1, 2, 3, 10], s = [1.0, 0.5]
    A, B = sband(n, [0.1]), sband(n, [1.0])
    κ, η, ζ = detderivates(A, B, s)
    @test κ == n
    @test η ≈ -n / (0.1 - s)
    @test ζ ≈ n / (0.1 - s)^2
end

@testset "small real $n" for n = 1:4
    A, B = sband(n, [6.0, -4, 1]), sband(n, [2.0, 1])
    s = 0.001
    d0 = det0(A, B, s)
    κ, η, ζ = detderivates(A, B, s)
    @test κ == 0
    @test isapprox(η, det1(A, B, s) / d0; rtol=1e-3)
    @test isapprox(η^2 - ζ, det2(A, B, s) / d0; atol=1e-5, rtol=1e-3)
end

@testset "small complex $n" for n = 1:4
    A, B = sband(n, [6.0, -4, 1+0im]), sband(n, [2.0, 1im])
    s = 0.001
    d0 = det0(A, B, s)
    κ, η, ζ = detderivates(A, B, s)
    @test η isa Real
    @test ζ isa Real
    @test κ == 0
    @test isapprox(η, det1(A, B, s) / d0; rtol=1e-3)
    @test isapprox(η^2 - ζ, det2(A, B, s) / d0; atol=1e-5, rtol=1e-3)
end

kappa(x) = x < 0.0 ? 0 : 2
kappa1(x) = x < 1.0 ? 0 : 3
kappa2(x) = x < 0.2 ? 0 : x < 1 ? 1 : x < 2 ? 3 : typemax(Int)

MatrixAlgebra.detderivates!(f::Function, x) = f(x), 0.0, 0.0

@testset "bisection bounds 0-2" begin
    lb, ub = eigbounds(1, 2, kappa, 1.0)
    @test  all(isapprox.(lb, 0.0; atol=eps()))
    @test  all(isapprox.(ub, 0.0; atol=eps()))
end
@testset "bisection bounds 1-3" begin
    lb, ub = eigbounds(1, 3, kappa1, 1.0)
    @test  all(isapprox.(lb, 1.0; atol=eps()))
    @test  all(isapprox.(ub, 1.0; atol=eps()))
end
@testset "bisection bounds 2-4" begin
    lb, ub = eigbounds(1, 4, kappa2, 1.0)
    @test all(isapprox.(lb, [0, 1, 1, 1]))
    @test all(isapprox.(ub, [0.5, 1, 1, 3]))
end
@testset "bisection bounds 2-5" begin
    lb, ub = eigbounds(1, 5, kappa2, 1.0)
    @test all(isapprox.(lb, [0, 1, 1, 2, 2]))
    @test all(isapprox.(ub, [0.5, 1, 1, 2, 2]))
end


using BlockArrays
using LinearAlgebra
using JuMP
using Ipopt
using Plots
using PyPlot
using Plots.PlotMeasures
using Formatting
using Printf

pyplot()
PyPlot.matplotlib.rc("text", usetex=true)
PyPlot.matplotlib.rc("font", family="serif")
Plots.scalefontsizes(1.5)

φ(x, n, σ) = exp(-x^2 / (2*σ^2)) * x^n

function dφ(x, n, σ)
    if n == 0
        exp(-x^2 / (2*σ^2)) * (-x/σ^2)
    else
        exp(-x^2 / (2*σ^2)) * (n * x^(n-1) - x^(n+1)/σ^2)
    end
end

K(x, y, σ, N) = sum(φ(x, n, σ) * φ(y, n, σ) for n=0:N)
∂K(x, y, σ, N) = sum(φ(x, n, σ) * dφ(y, n, σ) for n=0:N)
∂∂K(x, y, σ, N) = sum(dφ(x, n, σ) * dφ(y, n, σ) for n=0:N)

function D(x, σ, N)
    s = length(x)
    D0 = [K(x[i], x[j], σ, N) for i=1:s, j=1:s]
    D1 = [∂K(x[i], x[j], σ, N) for i=1:s, j=1:s]
    D2 = [∂∂K(x[i], x[j], σ, N) for i=1:s, j=1:s]
    [D0 D1; D1 D2]
end

function compute_ηV(x0, σ, N)
    s = length(x0)

    model = Model(with_optimizer(Ipopt.Optimizer, print_level=0))
    @variable(model, p[1:N+1])

    for i=1:s
        @constraint(model, sum(p[n+1] * φ(x0[i], n, σ) for n=0:N) == 1)
        @constraint(model, sum(p[n+1] * dφ(x0[i], n, σ) for n=0:N) == 0)
    end

    @objective(model, Min, sum(p[n]^2 for n=1:N+1))
    optimize!(model)

    pV = [value(p[n]) for n=1:N+1]
    ηV(x) = sum(φ(x, n, σ) * pV[n+1] for n=0:N)
end

fetitle = FormatExpr(L"$\sigma={}$")

function plot_ηV(x0, ηV, σ)
    s = length(x0)
    numpoints = 500

    x_tab = range(-1, stop=1, length=numpoints)
    y_tab = [ηV(x) for x in x_tab]

    Plots.plot(x_tab, y_tab, gridcolor="lightgrey", title=format(fetitle, σ))
    hline!([1], linestyle=:dash, linecolor=:black)

    if min(y_tab...) < -0.05
        hline!([-1], linestyle=:dash, linecolor=:black)
    end

    plot!(x0, ones(s), line=:stem, linestyle=:dash, marker=:circle, markersize=5,
          color="red")
    plot!(xlabel=L"$x$", legend=false, yformatter = yi -> @sprintf("%0.2f", yi))
    plot!(bottom_margin=5mm, right_margin=10mm, left_margin=5mm, top_margin=5mm)
end

N = 9

l = @layout[a; [b c d]; [e f g]]

p0 = Plots.plot((1:1)', label=L"$\eta_V$")
plot!(p0, (1:1)', label=L"$x_0$", color="red", linestyle=:dash,
      legend=:bottom, framestyle=:none)

x0 = [-0.8, 0.1, 0.3, 0.9]

σ = 0.285
ηV = compute_ηV(x0, σ, N)
p1 = plot_ηV(x0, ηV, σ)

σ = 0.315
ηV = compute_ηV(x0, σ, N)
p2 = plot_ηV(x0, ηV, σ)

σ = 0.375
ηV = compute_ηV(x0, σ, N)
p3 = plot_ηV(x0, ηV, σ)

x0 = [-0.1, 0.0, 0.05, 0.1]

σ = 0.05
ηV = compute_ηV(x0, σ, N)
p4 = plot_ηV(x0, ηV, σ)

σ = 0.3
ηV = compute_ηV(x0, σ, N)
p5 = plot_ηV(x0, ηV, σ)

σ = 0.5
ηV = compute_ηV(x0, σ, N)
p6 = plot_ηV(x0, ηV, σ)

Plots.plot(p0, p1, p2, p3, p4, p5, p6,
           layout=l, link=:x, size=(1000, 750),
           left_margin=5mm, right_margin=10mm)


@add_benchmark function sdp_nuclear_norm_atom(handle_problem)
    y = Semidefinite(3)
    p = minimize(nuclearnorm(y), y[2,1]<=4, y[2,2]>=3, y[3,3]<=2)
    handle_problem(p)
end

@add_benchmark function sdp_operator_norm_atom(handle_problem)
    y = Variable((3,3))
    p = minimize(opnorm(y), y[2,1]<=4, y[2,2]>=3, sum(y)>=12)
    handle_problem(p)
end

@add_benchmark function sdp_sigma_max_atom(handle_problem)
    y = Variable((3,3))
    p = minimize(sigmamax(y), y[2,1]<=4, y[2,2]>=3, sum(y)>=12)
    handle_problem(p)
end

@add_benchmark function sdp_lambda_max_atom(handle_problem)
    y = Semidefinite(3)
    p = minimize(lambdamax(y), y[1,1]>=4)
    handle_problem(p)
end

@add_benchmark function sdp_lambda_min_atom(handle_problem)
    y = Semidefinite(3)
    p = maximize(lambdamin(y), tr(y)<=6)
    handle_problem(p)
end

@add_benchmark function sdp_matrix_frac_atom(handle_problem)
    probs = [
        let
            x = [1, 2, 3]
            P = Variable(3, 3)
            p = minimize(matrixfrac(x, P), P <= 2*eye(3), P >= 0.5 * eye(3))
        end,
        let
            x = Variable(3)
            P = Variable(3, 3)
            p = minimize(matrixfrac(x, P), lambdamax(P) <= 2, x[1] >= 1)
        end
    ]

    handle_problem.(probs)
end

@add_benchmark function sdp_sum_squares_atom(handle_problem)
    n = 10
    A = rand(n,n) + im*rand(n,n)
    A = A + A' # now A is hermitian
    x = ComplexVariable(n,n)
    objective = sumsquares(A - x)
    c1 = x in :SDP
    p = minimize(objective, c1)
    handle_problem(p)
end


using GLVisualize, GeometryTypes, GLAbstraction
using Colors, Reactive, FileIO

if !isdefined(:runtests)

	window = glscreen()
	timesignal = bounce(linspace(0f0,1f0,360))
end
t 		   = const_lift(*, timesignal, 10f0)
radius     = 200f0
w,h 	   = widths(window)
middle     = Point2f0(w/2, h/2)
circle_pos = Point2f0[(Point2f0(sin(i), cos(i))*radius)+middle for i=linspace(0, 2pi, 20)]
rotation   = Vec2f0[normalize(Vec2f0(middle)-Vec2f0(p)) for p in circle_pos]
scales     = map(t) do t
    Vec2f0[Vec2f0(30, ((sin(i+t)+1)/2)*60) for i=linspace(0, 2pi, 20)]
end

circles = visualize(
	(CIRCLE, circle_pos),
	rotation=rotation, scale=scales,
)

_view(circles, window)

if !isdefined(:runtests)
	renderloop(window)
end


@testset "Links" begin

  #Logit link
  @test isequal(Unmarked.invlink(0.5, Unmarked.LogitLink()), logistic(0.5))
  @test isequal(Unmarked.invlink([0.3,0.5], Unmarked.LogitLink()),
                logistic.([0.3,0.5]))

  #Log link
  @test isequal(Unmarked.invlink(0.5, Unmarked.LogLink()), exp(0.5))
  @test isequal(Unmarked.invlink([0.3,0.5], Unmarked.LogLink()),
                exp.([0.3,0.5]))

end


module GalerkinSparseGrids

using HCubature
using JLD2
if haskey(ENV, "MKLROOT")
    using MKLSparse
end
using ODE
using Printf
using SparseArrays
using StaticArrays
using Plots

# The following script files are used:

include("schemes.jl")                    # CartesianIndex Manipulation Methods
include("additional_tools.jl")            # Helper functions not specific to any part of the package
include("dg_basis.jl")                    # Gram-Schmidt procedure for DG basis functions in 1-D
include("1d_dg_functions.jl")             # Explicitly building the 1-D Galerkin Basis
include("dg_methods.jl")                # Multidimensional hierarchical & sparse coefficients
include("dg_vmethods.jl")                # Going between a dictionary & a vector of coeffs
include("derivative_matrix_elements.jl")# 1-D symbolic piecewise derivative
include("1d_derivative.jl")                # Constructing ideal 1D derivative matrix using boundary terms
include("multidim_derivative.jl")        # Multidimensional DG Derivatives in full & sparse bases
include("pdes.jl")                        # Solving the n-D wave equation with periodic boundary
include("1d_nodal_basis.jl")            # Construct 1D nodal basis for multiplication
include("multidim_nodal_basis.jl")        # Construct Multidimensional nodal basis for multiplication

include("basic_function_exact_coeffs.jl") # Exact coefficients for certain easy functions
include("error_measure.jl")                # Monte Carlo Methods to measure error
include("tensor_construct.jl")            # Quickly calculates coeffs of simple tensors of functions
include("traveling_wave.jl")              # Constructs initial data for traveling wave

export

coeffs_DG,
reconstruct_DG,

get_size,
V2D,
D2V,
V2Dref,
D2Vref,
vcoeffs_DG,

transform,
make_modal2point_matrices,
make_point2modal_matrices,

D_matrix,
grad_matrix,
laplacian_matrix,

nodal2points_1D,
points2nodal_1D,
nodal2heir_1D,
hier2points_1D,
points2hier_1D,

get_one_modal,
get_xi_modal,
get_r2_modal,
get_xi_point,
get_r2_point,

mcerr,
mcerr2,

tensor_construct,

wave_evolve_1D,
wave_evolve,
vlasov_evolve,
traveling_wave,

periodic_pos_DLF_matrix,
periodic_hier_DLF_matrix,
periodic_DLF_matrix,
hier2pos,
spzeros,
get_size,
cutoff

end # module


@doc raw"""
    MultinomialSymmetric{n} <: AbstractMultinomialDoublyStochastic{N}

The multinomial symmetric matrices manifold consists of all symmetric $n×n$ matrices with
positive entries such that each column sums to one, i.e.

````math
\begin{aligned}
\mathcal{SP}(n) \coloneqq \bigl\{p ∈ ℝ^{n×n}\ \big|\ &p_{i,j} > 0 \text{ for all } i=1,…,n, j=1,…,m,\\
& p^\mathrm{T} = p,\\
& p\mathbf{1}_n = \mathbf{1}_n
\bigr\},
\end{aligned}
````

where $\mathbf{1}_n$ is the vector of length $n$ containing ones.

It is modeled as an [`DefaultIsometricEmbeddingType`](@ref), [`AbstractEmbeddedManifold`](@ref)
via the [`AbstractMultinomialDoublyStochastic`](@ref) type, since it shares a few functions
also with [`AbstractMultinomialDoublyStochastic`](@ref), most and foremost projection of
a point from the embedding onto the manifold.

The tangent space can be written as

````math
T_p\mathcal{SP}(n) \coloneqq \bigl\{
X ∈ ℝ^{n×n}\ \big|\ X = X^{\mathrm{T}} \text{ and }
X\mathbf{1}_n = \mathbf{0}_n
\bigr\},
````

where $\mathbf{0}_n$ is the vector of length $n$ containing zeros.

More details can be found in Section IV[^DouikHassibi2019].

# Constructor

    MultinomialSymmetric(n)

Generate the manifold of matrices $\mathbb R^{n×n}$ that are doubly stochastic and symmetric.

[^DouikHassibi2019]:
    > A. Douik, B. Hassibi:
    > AbstractManifold Optimization Over the Set of Doubly Stochastic Matrices: A Second-Order Geometry,
    > IEEE Transactions on Signal Processing 67(22), pp. 5761–5774, 2019.
    > doi: [10.1109/tsp.2019.2946024](http://doi.org/10.1109/tsp.2019.2946024),
    > arXiv: [1802.02628](https://arxiv.org/abs/1802.02628).
"""
struct MultinomialSymmetric{N} <: AbstractMultinomialDoublyStochastic{N} end

function MultinomialSymmetric(n::Int)
    return MultinomialSymmetric{n}()
end

@doc raw"""
    check_point(M::MultinomialSymmetric, p)

Checks whether `p` is a valid point on the [`MultinomialSymmetric`](@ref)`(m,n)` `M`,
i.e. is a symmetric matrix with positive entries whose rows sum to one.
"""
function check_point(M::MultinomialSymmetric{n}, p; kwargs...) where {n}
    mpv = invoke(check_point, Tuple{supertype(typeof(M)),typeof(p)}, M, p; kwargs...)
    mpv === nothing || return mpv
    # the embedding checks for positivity and unit sum columns, by symmetry we would get
    # the same for the rows, so checking symmetry is the only thing left, we can just use
    # the corresponding manifold for that
    return check_point(SymmetricMatrices(n, ℝ), p)
end
@doc raw"""
    check_vector(M::MultinomialSymmetric p, X; kwargs...)

Checks whether `X` is a valid tangent vector to `p` on the [`MultinomialSymmetric`](@ref) `M`.
This means, that `p` is valid, that `X` is of correct dimension, symmetric, and sums to zero
along any row.
"""
function check_vector(M::MultinomialSymmetric{n}, p, X; kwargs...) where {n}
    mpv = invoke(
        check_vector,
        Tuple{supertype(typeof(M)),typeof(p),typeof(X)},
        M,
        p,
        X;
        kwargs...,
    )
    mpv === nothing || return mpv
    # from the embedding we know that columns sum to zero, only symmety is left, i.e.
    return check_vector(SymmetricMatrices(n, ℝ), p, X; kwargs...)
end

function decorated_manifold(::MultinomialSymmetric{N}) where {N}
    return MultinomialMatrices(N, N)
end

embed!(::MultinomialSymmetric, q, p) = copyto!(q, p)
embed!(::MultinomialSymmetric, Y, ::Any, X) = copyto!(Y, X)

@doc raw"""
    manifold_dimension(M::MultinomialSymmetric{n}) where {n}

returns the dimension of the [`MultinomialSymmetric`](@ref) manifold
namely
````math
\operatorname{dim}_{\mathcal{SP}(n)} = \frac{n(n-1)}{2}.
````
"""
@generated function manifold_dimension(::MultinomialSymmetric{n}) where {n}
    return div(n * (n - 1), 2)
end

@doc raw"""
    project(M::MultinomialSymmetric{n}, p, Y) where {n}

Project `Y` onto the tangent space at `p` on the [`MultinomialSymmetric`](@ref) `M`, return the result in `X`.
The formula reads
````math
    \operatorname{proj}_p(Y) = Y - (α\mathbf{1}_n^{\mathrm{T}} + \mathbf{1}_n α^{\mathrm{T}}) ⊙ p,
````
where $⊙$ denotes the Hadamard or elementwise product and $\mathbb{1}_n$ is the vector of length $n$ containing ones.
The two vector $α ∈ ℝ^{n×n}$ is given by solving
````math
    (I_n+p)α =  Y\mathbf{1},
````
where $I_n$ is teh $n×n$ unit matrix and $\mathbf{1}_n$ is the vector of length $n$ containing ones.

"""
project(::MultinomialSymmetric, ::Any, ::Any)

function project!(::MultinomialSymmetric{n}, X, p, Y) where {n}
    α = (I + p) \ sum(Y, dims=2) # Formula (49) from 1802.02628
    return X .= Y .- (repeat(α, 1, 3) .+ repeat(α', 3, 1)) .* p
end

@generated function representation_size(::MultinomialSymmetric{n}) where {n}
    return (n, n)
end

@doc raw"""
    retract(M::MultinomialSymmetric, p, X, ::ProjectionRetraction)

compute a projection based retraction by projecting $p\odot\exp(X⨸p)$ back onto the manifold,
where $⊙,⨸$ are elementwise multiplication and division, respectively. Similarly, $\exp$
refers to the elementwise exponentiation.
"""
retract(::MultinomialSymmetric, ::Any, ::Any, ::ProjectionRetraction)

function retract!(M::MultinomialSymmetric, q, p, X, ::ProjectionRetraction)
    return project!(M, q, p .* exp.(X ./ p))
end

"""
    vector_transport_to(M::MultinomialSymmetric, p, X, q)

transport the tangent vector `X` at `p` to `q` by projecting it onto the tangent space
at `q`.
"""
vector_transport_to(::MultinomialSymmetric, ::Any, ::Any, ::Any, ::ProjectionTransport)

function vector_transport_to!(M::MultinomialSymmetric, Y, p, X, q, ::ProjectionTransport)
    project!(M, Y, q, X)
    return Y
end

function Base.show(io::IO, ::MultinomialSymmetric{n}) where {n}
    return print(io, "MultinomialSymmetric($(n))")
end


unsafe_mulby_1000(x) = (x<<10) - (x << 5) + (x << 3)

safe_mulby_1000(x::Int64) = x <= 9_223_372_036_854_775 ? mulby_100(x) :
                                                         ArgumentError("$(x) is too large")
safe_mulby_1000(x::Int32) = x <= 2_147_483 ? mulby_100(x) : ArgumentError("$(x) is too large")

# nanoseconds per microsecond
# microseconds per millisecond
# milliseconds per second
# safe for 0 < x < 268435456  (2^28)
unsafe_fld_1000(x::T) where {T} =
    ((x >> 3) * 34_359_739) >> 32

fld_1000(x::T) where {T<:Union{Int64,UInt64}} =
    if (x < 434_934_000) # 2^28 = 268_435_456
    # if x & ~0xffffffff == 0
        unsafe_fld_1000(x)
    else
        fld(x, 1_000)
    end

function fldmod_1000(x::T) where {T<:Union{Int64,UInt64}}
    if abs(x) <= 9_780_955_816
        quotient = if (x < 434_934_000)          # in 2^28 .. 2^29
            unsafe_fld_1000(x)
        else
            fld(x, 1_000)
        end

        remainder = x - unsafe_mulby_1000(quotient)

        if signbit(quotient) && signbit(remainder)
            quotient -= 1
            remainder += 1000
        end
    else
        quotient, remainder = fldmod(x, 1000)
    end

    quotient, remainder
end

# nanoseconds per millisecond
# microseconds per second

function mulby1000000(x)
    y = Int128(x)
    z = (y << 6) - y
    z = (z << 5) - z
    z = (z << 3) + y
    z << 6
end

# safe for x:2^2:Int64 <= 2^27 - 1
unsafe_fld_1_000_000(x::T) where {T<:Union{Int64,UInt64}} =
    ((x >> 6) * 274_878) >> 32

fld_1_000_000(x::T) where {T<:Union{Int64,UInt64}} =
    if (x < 134_217_728)          # 2^27 -1
        unsafe_fld_1_000_000(x)
    else
        fld(x, 1_000_000)
    end

function fldmod_1_000_000(x::T) where {T<:Union{Int64,UInt64}}
    quotient = if (x < 134_217_728)          # 2^27 -1
        unsafe_fld_1_000_000(x)
    else
        fld(x, 1_000_000)
    end

    remainder = x - quotient * 1_000_000
    quotient, remainder
end

# 86_400 (SecondsPerDay)
function mulby86400(x)
    y = x
    z = (y << 4) - y
    z = (z << 4) - z
    z = (z << 2) - z
    z << 7
end

# 86_400_000 (MillisecondsPerDay)
function mulby86400000(x)
    y = x
    z = (y << 5) - y
    z = (z << 2) + y
    z = (z << 4) - z
    z = (z << 4) - z
    z = (z << 2) - z
    z << 10
end

# 60 (SecondsPerMinute, MinutesPerHour)
function mulby60(x)
    y = x
    z = (y << 4) - y
    z << 2
end

# 3600 = 60*60 (SecondsPerHour)
function mulby3600(x)
    y = x
    z = (y << 4) - y
    z = (z << 4) - z
    z << 4
end

# 1440 = 24*60 (MinutesPerDay)
function mulby1440(x)
    y = x
    z = (y << 4) - y
    z = (z << 2) - z
    z << 5
end

# 1_000_000_000 (NanosecondsPerSecond)
function mulby_1000000000(x)
    y = x
    z = (y << 6) - y
    z = (z << 5) - z
    z = (z << 3) + y
    z = (z << 2) + z
    z = (z << 2) + z
    z = (z << 2) + z
    z << 9
end

#=

reference

Euclidean Affine Functions and Applications to Calendar Algorithms
by Cassio Neri, Lorenz Schneider
    arXiv:2102.06959v1 [cs.DS] 13 Feb2021
    https://arxiv.org/pdf/2102.06959.pdf
    pg 11
=#

mod2p32(x) = (x & 0xffff)

unsafe_fld_3600(x) =
    (1193047 * x) >> 32

function fld_3600(x)
    if x < 2_257_200
        unsafe_fld_3600(x)
    else
        fld(x, 3600)
    end
end

unsafe_mod_3600(x) =
    3600 * mod2p32(1193047 * x) >> 32

function mod_3600(x)
    if x < 2_255_762
        unsafe_mod_3600(x)
    else
        mod(x, 3600)
    end
end

function fldmod_3600(x)
    if x < 2_255_762
        (unsafe_fld_3600(x), unsafe_mod_3600(x))
    else
        fldmod(x, 3600)
    end
end

unsafe_fld_60(x) =
    mod2p32(71_582_789 * x) >> 32

function fld_60(x)
    if x < 97_612_920
        unsafe_fld_60(x)
    else
        fld(x, 60)
    end
end

unsafe_mod_60(x) =
    60 * mod2p32(71_582_789 * x) >> 32

function mod_60(x)
    if x < 97_612_920
        unsafe_mod_60(x)
    else
        mod(x, 60)
    end
end

function fldmod_60(x)
    if x < 2_255_762
        (unsafe_fld_60(x), unsafe_mod_60(x))
    else
        fldmod(x, 60)
    end
end

unsafe_fld_10(x) =
    mod2p32(429_496_730 * x) >> 32

function fld_10(x)
    if x < 1_073_741_829
        unsafe_fld_10(x)
    else
        fld(x, 10)
    end
end

unsafe_mod_10(x) =
    10 * mod2p32(429_496_730 * x) >> 32

function mod_10(x)
    if x < 1_073_741_829
        unsafe_mod_10(x)
    else
        mod(x, 10)
    end
end

function fldmod_10(x)
    if x < 1_073_741_829
        (unsafe_fld_10(x), unsafe_mod_10(x))
    else
        fldmod(x, 10)
    end
end

#=

const twopow32 = typemax(UInt32)
mod2p32(x) = (x & 0xffff)

const k3600 =   1_193_047     #  n_div_3600 = mod2p32(k3600 * mod32(n))
const   k60 =  71_582_789     #    n_div_60 = mod2p32(  k60 * mod32(n))
const   k10 = 429_496_730     #    n_div_10 = mod2p32(  k10 * mod32(n))

const k3600 =   1_193_047     #  n_mod_3600 = (k3600 * mod32(n)) ÷ k3600
const   k60 =  71_582_789     #    n_mod_60 = (  k60 * mod32(n)) ÷ k60
const   k10 = 429_496_730     #    n_mod_10 = (  k10 * mod32(n)) ÷ k10

const k3600 =   1_193_047     #  n_div_3600 = mod2p32(k3600 * mod32(n))
const   k60 =  71_582_789     #    n_div_60 = mod2p32(  k60 * mod32(n))
const   k10 = 429_496_730     #    n_div_10 = mod2p32(  k10 * mod32(n))

const k3600 =   1_193_047     #  n_mod_3600 = (k3600 * mod32(n)) ÷ k3600
const   k60 =  71_582_789     #    n_mod_60 = (  k60 * mod32(n)) ÷ k60
const   k10 = 429_496_730     #    n_mod_10 = (  k10 * mod32(n)) ÷ k10

const k3600 =   1_193_047     #  n_div_3600 = mod2p32(k3600 * mod32(n))
const   k60 =  71_582_789     #    n_div_60 = mod2p32(  k60 * mod32(n))
const   k10 = 429_496_730     #    n_div_10 = mod2p32(  k10 * mod32(n))

const k3600 =   1_193_047     #  n_mod_3600 = (k3600 * mod32(n)) ÷ k3600
const   k60 =  71_582_789     #    n_mod_60 = (  k60 * mod32(n)) ÷ k60
const   k10 = 429_496_730     #    n_mod_10 = (  k10 * mod32(n)) ÷ k10

=#

#=
    divmod10(k)   -> (k ÷   10, k %   10)     ∀n ∈ [0,     2_257_200]
    divmod60(k)   -> (k ÷   60, k %   60)     ∀n ∈ [0,    97_612_920]
    divmod3600(k) -> (k ÷ 3600, k % 3600)     ∀n ∈ [0, 1_073_741_830]

    div10(k)    -> k ÷   10   ↤  ( 429_496_730 * n ) & twopow32
    div60(k)    -> k ÷   60   ↤  (  71_582_789 * n ) & twopow32
    div3600(k)  -> k ÷ 3600   ↤  (   1_193_047 * n ) & twopow32

    mod10(k)    -> k %   10   ↤  ( 429_496_730 * n ) % twopow32
    mod60(k)    -> k %   60
    mod3600(k)  -> k % 3600
=#


board = [[" ", " ", " "],
         [" ", " ", " "],
         [" ", " ", " "]]

player1 = "X"
player2 = "O"

player1_win = -1
tie = 0
player2_win = 1

function print_board(board)
    println("    1   2   3")
    println(" 1  ", join(board[1], " | "))
    println("   ———┼———┼———")
    println(" 2  ", join(board[2], " | "))
    println("   ———┼———┼———")
    println(" 3  ", join(board[3], " | "))
end

function check_available(board)
    available_cells = []
    for i in 1:3, j in 1:3
        if board[i][j] == " "
            push!(available_cells, (i, j))
        end
    end
    return available_cells
end

function check_win(board)

    # Rows
    for i in 1:3
        if (board[i][1] == board[i][2] == board[i][3]) && (board[i][1] in [player1, player2])
            if (board[i][1] == player1)
                return player1_win
            end
            return player2_win
        end
    end

    # Column
    for i in 1:3
        if (board[1][i] == board[2][i] == board[3][i]) && (board[1][i] in [player1, player2])
            if (board[1][i] == player1)
                return player1_win
            end
            return player2_win
        end
    end

    # Principal diagonal
    if (board[1][1] == board[2][2] == board[3][3]) && (board[1][1] in [player1, player2])
        if (board[1][1] == player1)
            return player1_win
        end
        return player2_win
    end

    # Other diag
    if (board[1][3] == board[2][2] == board[3][1]) && (board[1][3] in [player1, player2])
        if (board[1][3] == player1)
            return player1_win
        end
        return player2_win
    end

    # Tie
    if length(check_available(board)) == 0
        return tie
    end

    # No one wins, game still goes on
    return nothing
end

function make_move!(board, player, (x, y))
    board[x][y] = player
end

function worst_move(board)
    worst_score = Inf

    worst_pos = (0, 0)

    for i in 1:3, j in 1:3
        if board[i][j] == " "

            board[i][j] = player2

            score = minimax(board, 0, true)

            board[i][j] = " "

            if score < worst_score
                worst_score = score
                worst_pos = (i, j)
            end
        end
    end
    return worst_pos
end

function ai_move!(board, difficulty)
    if difficulty == "loser"
        make_move!(board, player2, worst_move(board))
    elseif difficulty == "easy"
        pos = rand(check_available(board))
        make_move!(board, player2, pos)
    elseif difficulty == "medium"
        if rand() < 0.3
            make_move!(board, player2, best_move(board))
        else
            pos = rand(check_available(board))
            make_move!(board, player2, pos)
        end
    elseif difficulty == "hard"
        if rand() < 0.7
            make_move!(board, player2, best_move(board))
        else
            pos = rand(check_available(board))
            make_move!(board, player2, pos)
        end
    elseif difficulty == "impossible"
        make_move!(board, player2, best_move(board))
    end

end

function minimax(board, depth, is_maximising)
    result = check_win(board)

    if result !== nothing
        return result
    elseif is_maximising

        best_score = -Inf

        for i in 1:3, j in 1:3
            if board[i][j] == " "

                board[i][j] = player2

                score = minimax(board, depth + 1, false)

                board[i][j] = " "

                best_score = max(score, best_score)
            end
        end

        return best_score
    else
        best_score = Inf

        for i in 1:3, j in 1:3
            if board[i][j] == " "

                board[i][j] = player1

                score = minimax(board, depth + 1, true)

                board[i][j] = " "

                best_score = min(score, best_score)
            end
        end

        return best_score
    end

end

function best_move(board)
    best_score = -Inf

    best_pos = (0, 0)

    for i in 1:3, j in 1:3
        if board[i][j] == " "

            board[i][j] = player2

            score = minimax(board, 0, false)

            board[i][j] = " "

            if score > best_score
                best_score = score
                best_pos = (i, j)
            end
        end
    end
    return best_pos
end

function move_input(board, prompt, params)

    print(prompt)
    txt = split(readline())

    if txt == ["exit"]
        exit()
    end

    if length(txt) == params && txt[1] in string.(collect(1:3)) && txt[2] in string.(collect(1:3)) && (board[parse(Int, txt[1])][parse(Int, txt[2])] == " ")
        return txt
    else
        println("Please enter a valid command (See README)")
        move_input(board, prompt, params)
    end
end

function generalised_input(prompt, same_line, choices, params)
    same_line ? print(prompt) : println(prompt)

    txt = split(lowercase(readline()))

    if txt == ["exit"]
        exit()
    end

    if length(txt) == params && txt[1] in choices
        return txt
    else
        println("Please enter a valid command (See README)")
        generalised_input(prompt, same_line, choices, params)
    end
end

function singleplayer(board, (player1, player2), (player1_win, player2_win, tie))
    diff_levels = """
    Choose a difficulty level:
    - `loser` - AI always makes worst move
    - `easy` - AI plays randomly
    - `medium` - AI makes some mistakes
    - `hard` - AI makes very few mistakes
    - `impossible` - AI always makes best move. Impossible to win.
    """

    println(diff_levels)
    difficulty = generalised_input("Difficulty: ", true, ["loser", "easy", "medium", "hard", "impossible"], 1)[1]

    println("Human is X, AI is O. Human goes first.")
    print_board(board)

    while true

        txt = move_input(board, "Enter position of cell (row and column seperated by space): ", 2)

        make_move!(board, player1, (parse(Int, txt[1]), parse(Int, txt[2])))

        if check_win(board) == player1_win
            println("Human wins! The humans are still worthy!")
            break
        elseif check_win(board) == player2_win
            println("AI wins. Humanity has no hope.")
            break
        elseif check_win(board) == tie
            println("Tie! Have we reached the singularity?")
            break
        end

        ai_move!(board, difficulty)
        print_board(board)

        if check_win(board) == player1_win
            println("Human wins! The human are still worthy!")
            break
        elseif check_win(board) == player2_win
            println("AI wins. Humanity has no hope.")
            break
        elseif check_win(board) == tie
            println("Tie! Have we reached the singularity?")
            break
        end
    end
end

function multiplayer(board, (player1, player2), (player1_win, player2_win, tie))
    println("Player 1 is X, Player 2 is O. Player 1 goes first.")
    print_board(board)

    while true

        txt = move_input(board, "Enter position of cell (row and column seperated by space): ", 2)

        make_move!(board, player1, (parse(Int, txt[1]), parse(Int, txt[2])))
        print_board(board)

        if check_win(board) == player1_win
            println("Player 1 Wins!")
            break
        elseif check_win(board) == player2_win
            println("Player 2 Wins!")
            break
        elseif check_win(board) == tie
            println("Tie!")
            break
        end

        txt = move_input(board, "Enter position of cell (row and column seperated by space): ", 2)

        make_move!(board, player2, (parse(Int, txt[1]), parse(Int, txt[2])))

        print_board(board)

        if check_win(board) == player1_win
            println("Player 1 Wins!")
            break
        elseif check_win(board) == player2_win
            println("Player 2 Wins!")
            break
        elseif check_win(board) == tie
            println("Tie!")
            break
        end
    end
end

function play_again()
    replay = generalised_input("Play again? (y/n): ", true, ["y", "n"], 1)[1]
    if replay == "y"
        run(`julia $PROGRAM_FILE`)
        exit()
    else
        exit()
    end
end

mode = generalised_input("Welcome to TicTacToe! Would you like to play singleplayer (1) or multiplayer (2)?", false, ["1", "2"], 1)[1]

if mode == "1"
    singleplayer(board, (player1, player2), (player1_win, player2_win, tie))
else
    multiplayer(board, (player1, player2), (player1_win, player2_win, tie))
end

play_again()


using Plots, Random, Interact, Distributions, LinearAlgebra, DelimitedFiles

# N = 30 #number of osc
# plot([rand(Distributions.Binomial(30,0.1)) for i in 1:100], st=:histogram)
# rand([0,-rand(Distributions.Binomial(30,1/150))-0.0001*rand()])
function randomHessian(N; dims=2, p=0.0)
    p = p == 0.0 ? 1/(5*N) : p
    m = zeros(dims,N,N)
    for k in 1:dims
        # m[k,:,:] = [rand([0,-rand(1:N)-0.1*rand()]) for j in 1:N, i in 1:N]
        m[k,:,:] = [rand([0,-rand(Distributions.Binomial(N,p))-0.0001*rand()]) for j in 1:N, i in 1:N]
        m[k,:,:] = (m[k,:,:] + m[k,:,:]')/2
        for i in 1:N
            m[k,i,i] = -sum(m[k,i,:])+m[k,i,i]
        end
    end
    m ./ maximum(m, dims=3)
end

function expHessian(N; dims=2)
    m = zeros(dims,N,N)
    for k in 1:dims
        m[k,:,:] = [-exp(-(i-j)^2) for i in 1:N, j in 1:N]
        m[k,:,:] = (m[k,:,:] + m[k,:,:]')/2
        for i in 1:N
            m[k,i,i] = -sum(m[k,i,:])+m[k,i,i]
        end
    end
    m
end

function nNeighHessian(N; dims=2, n=2)
    m = zeros(dims,N,N)
    for k in 1:dims
        m[k,:,:] = [ abs(i-j) < n ? -1 : 0 for i in 1:N, j in 1:N]
        m[k,:,:] = (m[k,:,:] + m[k,:,:]')/2
        for i in 1:N
            m[k,i,i] = -sum(m[k,i,:])+m[k,i,i]
        end
    end
    m
end

function genHessian(N, type; dims=2, n=2)
    if type == "random"
        return randomHessian(N; dims)
    elseif type == "exp"
        return expHessian(N; dims)
    elseif type == "nNeigh"
        return nNeighHessian(N; dims,n)
    end
end

function gen_Cryst(N; L=[5,5,5])
    init_pos = zeros(3,N)
    for i in 0:N-1
        init_pos[1,i+1] = i % L[1]
        init_pos[2,i+1] = (i ÷ L[1]) % L[2]
        init_pos[3,i+1] = i ÷ (L[1]*L[2])
    end
    init_pos
end

function generateData(N; dims=2, steps=500, Δt=0.01, arrange = "Crystal",
            typeHessian="random", amp = 0.4, nModes = 0, idxStart = 1)
    if arrange == "Crystal"
        L = [Int(floor(N^(1/dims))) for i in 1:dims]
        if dims==1
            append!(L,1)
        end
        init_pos = gen_Cryst(N; L)
    end

    M = genHessian(N, typeHessian; dims)
    r = zeros(dims, N, steps)
    vel = zeros(dims, N, steps)
    nModes = nModes == 0 ? N : nModes
    for k in 1:dims
        d,v = eigen(M[k,:,:])
        # r[k,:,:] = hcat([sum([amp * d[i]^2 .* v[:,i] .* sin(2π*t*d[i]) for i in idxStart:nModes]) .+ init_pos[k,:] for t in 0:Δt:(steps-1)*Δt]...)
        r[k,:,:] = hcat([sum([amp .* v[:,i] .* sin(2π*t*d[i]) for i in idxStart:nModes]) .+ init_pos[k,:] for t in 0:Δt:(steps-1)*Δt]...)
        vel[k,:,:] = hcat([sum([amp .* v[:,i] .* cos(2π*t*d[i]) * 2π * d[i] for i in idxStart:nModes]) for t in 0:Δt:(steps-1)*Δt]...)
    end
    permutedims(r,(1,3,2)), permutedims(vel,(1,3,2)), M
end


module SparseGrids

import Compat.view

import FastGaussQuadrature: gausshermite
import Iterators: collect, product

include("grids.jl")
include("nested.jl")

# package code goes here
export
	sparsegrid,
	tensorgrid,
	listNdq,
	combvec,
	symmetrize!,
	kpn

end # module


module SimpleDiscretes

import Base: minimum, maximum
import Statistics

import Distributions
const Dst = Distributions

export SimpleDiscrete

struct SimpleDiscrete{T1<:Real, T2<:Real} <: Dst.DiscreteUnivariateDistribution
    xs::Vector{T1}
    ps::Vector{T2}
end

maximum(d::SimpleDiscrete) = isempty(d.xs) ? -Inf64 : maximum(d.xs)
minimum(d::SimpleDiscrete) = isempty(d.xs) ? +Inf64 : minimum(d.xs)

Statistics.mean(d::SimpleDiscrete) = sum(d.xs .* d.ps)
Statistics.var(d::SimpleDiscrete)  = Statistics.mean(d.xs.^2) - Statistics.mean(d.xs)^2

Dst.insupport(d::SimpleDiscrete, x::Real) = minimum(d) <= x && x <= maximum(d)

function Dst.cdf(d::SimpleDiscrete, x::Real)
    return sum(pᵢ for (xᵢ, pᵢ) in zip(md.mass_points, md.mass_probs) if xᵢ <= x)
end

function Dst.pdf(d::SimpleDiscrete{T1,T2}, x::Real) where {T1,T2}
    for (xᵢ , pᵢ) in zip(d.xs,d.ps)
        xᵢ ≈ x && return pᵢ
    end
    return zero(T2)
end

Dst.ncomponents(d::SimpleDiscrete) = length(d.xs)

Dst.params(d::SimpleDiscrete) = (d.xs, d.ps)

end


tanh(x::Float64, ::Type{RoundingMode{:Nearest}}) =
   sinh(x, RoundingMode{:Nearest}) / cosh(x, RoundingMode{:Nearest})

function tanh(x::Float64, ::Type{RoundingMode{:Down}})
   t = sinh(x, RoundingMode{:Nearest})
   c = cosh(x, RoundingMode{:Nearest})
   with_rounding(Float64, RoundDown) do
       t = t/c
   end
   t
end
function tanh(x::Float64, ::Type{RoundingMode{:Up}})
   t = sinh(x, RoundingMode{:Nearest})
   c = cosh(x, RoundingMode{:Nearest})
   with_rounding(Float64, RoundUp) do
       t = t/c
   end
   t
end

for (fn, invfn) in ((:csc,:sin),(:sec,:cos),(:cot,:tan))
    @eval begin
        function ($fn)(x::Float64, ::Type{RoundingMode{:Up}})
            y= $(invfn)(x, RoundingMode{:Nearest})
            with_rounding(Float64, RoundDown) do
               y = 1.0 / y
            end
            y
        end
        function ($fn)(x::Float64, ::Type{RoundingMode{:Down}})
            y= $(invfn)(x, RoundingMode{:Nearest})
            with_rounding(Float64, RoundUp) do
               y = 1.0 / y
            end
            y
        end
        function ($fn)(x::Float64, ::Type{RoundingMode{:Nearest}})
            y = $(invfn)(x, RoundingMode{:Nearest})
            z = $(fn)(y)
            s = sign(y-z)
            if s < 0.0
               with_rounding(Float64, RoundDown) do
                   y = 1.0 / y
               end
            elseif x > 1.0
               with_rounding(Float64, RoundUp) do
                   y = 1.0 / y
               end
            else
               y = 1.0 / y
            end
            y
        end

    end
end


using Gadfly, Distributions

set_default_plot_size(6inch, 3inch)

plot(x=rand(Poisson(20), 1000), Scale.x_discrete, Geom.histogram)


import Statistics

lines = readlines("input.txt")
matrix = parse.(Bool, hcat(collect.(lines)...))

gamma = Statistics.mean(eachcol(matrix)) .>= 0.5
epsilon = @. !gamma

function bitarr_to_int(arr)
    return sum(arr .* (2 .^ collect(length(arr)-1:-1:0)))
end

println(bitarr_to_int(gamma) * bitarr_to_int(epsilon))