export FreeLSProcessTomo,
       LSProcessTomo

function build_liou_process_predictor(prep::Vector, obs::Vector)
    exps = [ vec(p)*vec(o)' for o in obs, p in prep ]
    return reduce(vcat, map(m->vec(m)', vec(exps)) )
end

function build_choi_process_predictor(prep::Vector, obs::Vector)
    exps = [ QuantumInfo.choi_liou_involution(vec(o)*vec(p)') for o in obs, p in prep ]
    return reduce(vcat, map(m->vec(m)', vec(exps)) )
end

function isstate(m::Matrix)
    LinearAlgebra.ishermitian(m) && isapprox(LinearAlgebra.tr(m),1.0)
end

"""
`FreeLSProcessTomo(states::Vector,obs::Vector)`

Free (unconstrained) least-squares process tomography algorithm. It is
constructed from a vector of states that are used to probe a process, and
a vector of observables that are used to measure the probe states.
"""
struct FreeLSProcessTomo
    statecount::Int
    obscount::Int
    inputdim::Int
    outputdim::Int
    pred::Matrix{ComplexF64}
    function FreeLSProcessTomo(states::Vector,obs::Vector)
        @assert all([LinearAlgebra.ishermitian(o) for o in obs]) "Observables must be Hermitian"
        @assert all([isstate(o) for o in states])  "States must be valid density matrices"

        pred = build_liou_process_predictor(states, obs)
        outputdim = size(pred,1)
        inputdim = size(pred,2)
        new(length(states),length(obs),inputdim,outputdim,pred)
    end
end

function predict(method::FreeLSProcessTomo, process::Matrix)
    reshape(real(method.pred*vec(process)), (method.statecount,method.obscount))
end

function fit(method::FreeLSProcessTomo,
             means::Vector)
    if length(means) != method.outputdim
        error("The number of sample means does not match the required number of experiments")
    end
    d = round(Int, size(method.pred, 2) |> sqrt)
    reg = method.pred \ means
    return (reshape(reg,d,d),
            LinearAlgebra.norm(method.pred*reg-means,2)^2,
            :Optimal)
end

function fit(method::FreeLSProcessTomo,
             means::Vector,
             vars::Vector{Float64})
    if length(means) != method.outputdim
        error("The number of sample means does not match the required number of experiments")
    end
    if length(vars) != method.outputdim
        error("The number of sample variances does not match the required number of experiments")
    end
    if any(vars.<0)
        error("Samples variances must be positive for generalized least squares process tomography.")
    end
    d = round(Int, size(method.pred,2) |> sqrt)
    # weighted least squares
    reg = (LinearAlgebra.Diagonal(1 ./ sqrt.(vars)) * method.pred) \ (LinearAlgebra.Diagonal(1 ./ sqrt.(vars)) * means)
    return (reshape(reg,d,d),
            LinearAlgebra.norm((method.pred*reg-means) ./ sqrt.(vars),2)^2,
            :Optimal)
end

function fit(method::FreeLSProcessTomo,
             means::Matrix)
    if size(means) != (method.statecount, method.obscount)
        error("Mean matrix must be $((method.statecount, method.obscount)), but is $(size(means))")
    end
    fit(method, vec(means))
end

function fit(method::FreeLSProcessTomo,
             means::Matrix,
             vars::Matrix{Float64})
    if size(means) != (method.statecount, method.obscount)
        error("Mean matrix must be $((method.statecount, method.obscount)), but is $(size(means))")
    elseif size(vars) != (method.statecount, method.obscount)
        error("Pointwise variance matrix must be $((method.statecount, method.obscount)), but is $(size(means))")
    end
    fit(method, vec(means), vec(vars))
end

"""
`LSProcessTomo(states::Vector,obs::Vector)`

Least-squares process tomography algorithm where the result is
contrained to be completely positive (CP) and trace perserving
(TP). It is constructed from a vector of states that are used to probe
a process, and a vector of observables that are used to measure the
probe states.

"""
struct LSProcessTomo
    statecount::Int
    obscount::Int
    inputdim::Int
    outputdim::Int
    pred::Matrix{ComplexF64}
    function LSProcessTomo(states::Vector,obs::Vector)
        @assert all([LinearAlgebra.ishermitian(o) for o in obs]) "Observables must be Hermitian"
        @assert all([isstate(o) for o in states]) "States must be valid density matrices"

        pred = build_choi_process_predictor(states, obs)
        outputdim = size(pred,1)
        inputdim = size(pred,2)
        new(length(states),length(obs),inputdim,outputdim,pred)
    end
end

function predict(method::LSProcessTomo, process::Matrix)
    reshape(real(method.pred*vec(QuantumInfo.choi_liou_involution(process))), (method.statecount,method.obscount))
end

begin
    global fit
    tomo_dim = 0
    ptrb = SparseArrays.spzeros(0,0)

    function trb_sop(da,db)
        sop = zeros(da^2,(da*db)^2)
        for i=0:da-1
            for j=0:da-1
                for k=0:db-1
                    sop += vec(QuantumInfo.ketbra(i,j,da))*vec(kron(QuantumInfo.ketbra(i,j,da),QuantumInfo.ketbra(k,k,db)))'
                end
            end
        end
        return SparseArrays.sparse(sop)
    end

    function fit(method::LSProcessTomo,
                 means::Vector{Float64},
                 vars = ones(length(means));
                 #solver = MosekSolver(LOG=0))
                 solver = SCS.SCSSolver(verbose=0, max_iters=10_000, eps = 1e-8))

        if length(means) != length(vars) || size(method.pred,1) != length(means)
            error("Size of observations and predictons do not match.")
        end
        d4 = size(method.pred,2)
        d2 = Int(sqrt(d4))
        d  = Int(sqrt(d2))

        # We assume that the predictions are always real-valued
        # and we need to do the complex->real translation manually since
        # Convex.jl does not support complex numbers yet
        rpred = real(method.pred)
        ipred = -imag(method.pred)
        ivars = 1 ./ sqrt.(vars)

        if d != tomo_dim
            ptrb = trb_sop(d,d)
        end

        ρr = Convex.Variable(d2,d2)
        ρi = Convex.Variable(d2,d2)

        # least squares problem
#       NOTE: Notation "^2" is not working here.
#        problem = Convex.minimize( Convex.norm( (means - (rpred*vec(ρr) + ipred*vec(ρi))) .* ivars, 2 )^2 )
        problem = Convex.minimize( Convex.square(Convex.norm( (means - (rpred*vec(ρr) + ipred*vec(ρi))) .* ivars, 2 ) ))

        # CP constraint
        problem.constraints += Convex.isposdef([ρr ρi; -ρi ρr])
        # TP constraint
        problem.constraints += reshape(ptrb*vec(ρr),d,d) == QuantumInfo.eye(d)
        problem.constraints += reshape(ptrb*vec(ρi),d,d) == zeros(d,d)

        Convex.solve!(problem, solver)

        return QuantumInfo.choi_liou_involution((ρr.value + 1im*ρi.value)), problem.optval, problem.status
    end
end

function fit(method::LSProcessTomo,
             means::Matrix,
             vars::Matrix{Float64}=ones(size(means)))
    if size(means) != (method.statecount, method.obscount)
        error("Matrix of sample means must be $((method.statecount, method.obscount)), but is $(size(means))")
    end
    if size(vars) != (method.statecount, method.obscount)
        error("Matrix of sample variances must be $((method.statecount, method.obscount)), but is $(size(means))")
    end
    fit(method, vec(means), vec(vars))
end


# ------------------------------------------------------------------------------
# h-multigrid example
# ------------------------------------------------------------------------------

using LFAToolkit
using LinearAlgebra

# setup
mesh = Mesh2D(1.0, 1.0)
p = 1
numbercomponents = 1
numberfineelements1d = 2
dimension = 2
ctofbasis = TensorH1LagrangeHProlongationBasis(
    p + 1,
    numbercomponents,
    dimension,
    numberfineelements1d,
);

# operators
finediffusion =
    GalleryMacroElementOperator("diffusion", p + 1, p + 2, numberfineelements1d, mesh);
coarsediffusion = GalleryOperator("diffusion", p + 1, p + 2, mesh);

# Chebyshev smoother
chebyshev = Chebyshev(finediffusion)

# h-multigrid preconditioner
multigrid = HMultigrid(finediffusion, coarsediffusion, chebyshev, [ctofbasis])

# compute operator symbols
A = computesymbols(multigrid, [3], [1, 1], [π, π])
eigenvalues = real(eigvals(A))

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


using Plots

f(x, y) = (1/2pi)*exp(-(x^2+y^2)/2)

function main()
    plotlyjs()
    xs = -5.:0.01:5.
    ys = -1.:0.01:1.
    zs = [f(x, y) for y in ys, x in xs]
    plot(xs, ys, zs, linetype=:surface)
end


"""
    particle dynamics
    - 3D particle subject to contact forces

    - configuration: q = (x, y, z) ∈ R³
    - impact force (magnitude): n ∈ R₊
    - friction force: b ∈ R²
        - contact force: λ = (b, n) ∈ R² × R₊
        - friction coefficient: μ ∈ R₊

    Discrete Mechanics and Variational Integrators
        pg. 363
"""
struct Particle
    n::Int
    m::Int
    d::Int

    mass # mass
    gravity # gravity
    friction_coeff # friction coefficient

    nq # configuration dimension
    nu # control dimension
end

# mass matrix
function mass_matrix(model)
    m = model.mass

    Diagonal(@SVector [m, m, m])
end

# gravity
function gravity(model, q)
    m = model.mass
    g = model.gravity

    @SVector [0.0, 0.0, m * g]
end

# signed distance function
function signed_distance(model, q)
    q[3]
end

# contact force Jacobian
function jacobian(model, q)
    Diagonal(@SVector ones(model.nq))
end

# dynamics
function dynamics(model, q1, q2, q3, u, λ, h)
    nq = model.nq
    nu = model.nu
    SVector{nq}(mass_matrix(model) * (2.0 * q2 - q1 - q3) / h
        - h * gravity(model, q2)
        + h * jacobian(model, q3)' * (u + λ))
end


# Code derived from metabolicEP (https://github.com/anna-pa-m/Metabolic-EP)

function idxlicols(X; tol::Float64=1e-10)

    sum(abs2,X) == 0 && (return(Array{Int,1}(), Array{Int,2}() ))
    Q,R,E = qr(X, ColumnNorm())
    diagr = abs.(diag(R))
    r = findall(diagr .>= tol*diagr[1])[end]
    idx = sort(E[1:r])
    return idx
end

function echelonize(X::T,v; eps::Real=1e-10) where T <:DenseArray
    M,N = size(X)

    idxrow = idxlicols(Matrix(X'))
    Mred = length(idxrow)

    idxind = idxlicols(X)
    idxdep = setdiff(1:N,idxind)
    newidx = vcat(idxind,idxdep)
    Tv = @view X[idxrow,newidx]
    iTv = inv(Tv[1:Mred, 1:Mred])
    res = iTv * Tv
    # trimming zeros
    for i in eachindex(res)
        abs(res[i]) < eps && (res[i] = zero(res[i]))
    end
    bnew = iTv * v[idxrow]
    # trimming ones
    for i in 1:Mred
        abs(1.0 - res[i,i]) < eps && (res[i,i] = one(res[i,i]))
    end

    return idxdep, newidx, res, bnew

end


"""
    AbstractLink

Abstract type defining maps from R^n -> X.
They can be applied by calling `link(x)`.

A series of definitions are given in http://web.pdx.edu/~newsomj/mvclass/ho_link.pdf
"""
abstract type AbstractLink end

"""
	ChainLink(links)

Create a composed chain of different links.
"""
struct ChainLink{Tls} <: AbstractLink
    links::Tls
end

(l::ChainLink)(x) = foldl((x, l) -> l(x), l.ls; init=x)

"""
    Link(f)

General construction for a link with a function `f`.
"""
struct Link{F} <: AbstractLink
    f::F
end

link(f) = Link(f)
link(l::AbstractLink) = l

(l::Link)(x) = l.f(x)

Base.inv(l::Link) = Link(InverseFunctions.inverse(l.f))

"""
    BijectiveSimplexLink(link)

Wrapper to preprocess the inputs by adding a `0` at the end before passing it to
the link `link`.
This is a necessary step to work with simplices.
For example with the [`SoftMaxLink`](@ref), to obtain a `n`-simplex leading to
`n+1` categories for the [`CategoricalLikelihood`](@ref),
one needs to pass `n+1` latent GP.
However, by wrapping the link into a `BijectiveSimplexLink`, only `n` latent are needed.
"""
struct BijectiveSimplexLink{L} <: AbstractLink
    link::L
end

(l::BijectiveSimplexLink)(f::AbstractVector{<:Real}) = l.link(vcat(f, 0))

# alias
const LogLink = Link{typeof(log)}
const ExpLink = Link{typeof(exp)}

const InvLink = Link{typeof(inv)}

const SqrtLink = Link{typeof(sqrt)}
const SquareLink = Link{typeof(InverseFunctions.square)}

const LogitLink = Link{typeof(logit)}
const LogisticLink = Link{typeof(logistic)}

const ProbitLink = Link{typeof(norminvcdf)}
const NormalCDFLink = Link{typeof(normcdf)}

const SoftMaxLink = Link{typeof(softmax)}

"""
    LogLink()

`log` link, f:ℝ⁺->ℝ . Its inverse is the [`ExpLink`](@ref).
"""
LogLink() = Link(log)

"""
    ExpLink()

`exp` link, f:ℝ->ℝ⁺. Its inverse is the [`LogLink`](@ref).
"""
ExpLink() = Link(exp)

"""
    InvLink()

`inv` link, f:ℝ/{0}->ℝ/{0}. It is its own inverse.
"""
InvLink() = Link(inv)

"""
    SqrtLink()

`sqrt` link, f:ℝ⁺∪{0}->ℝ⁺∪{0}. Its inverse is the [`SquareLink`](@ref).
"""
SqrtLink() = Link(sqrt)

"""
    SquareLink()

`^2` link, f:ℝ->ℝ⁺∪{0}. Its inverse is the [`SqrtLink`](@ref).
"""
SquareLink() = Link(InverseFunctions.square)

"""
    LogitLink()

`log(x/(1-x))` link, f:[0,1]->ℝ. Its inverse is the [`LogisticLink`](@ref).
"""
LogitLink() = Link(logit)

"""
    LogisticLink()

`1/(1+exp(-x))` link. f:ℝ->[0,1]. Its inverse is the [`LogitLink`](@ref).
"""
LogisticLink() = Link(logistic)

"""
    ProbitLink()

`ϕ⁻¹(y)` link, where `ϕ⁻¹` is the `invcdf` of a `Normal` distribution, f:[0,1]->ℝ.
Its inverse is the [`NormalCDFLink`](@ref).
"""
ProbitLink() = Link(norminvcdf)

"""
    NormalCDFLink()

`ϕ(y)` link, where `ϕ` is the `cdf` of a `Normal` distribution, f:ℝ->[0,1].
Its inverse is the [`ProbitLink`](@ref).
"""
NormalCDFLink() = Link(normcdf)

(::ChainLink{<:Tuple{LogLink,NormalCDFLink}})(x) = normlogcdf(x) # Specialisation for log + normal cdf

"""
    SoftMaxLink()

`softmax` link, i.e `f(x)ᵢ = exp(xᵢ)/∑ⱼexp(xⱼ)`.
f:ℝⁿ->Sⁿ⁻¹, where Sⁿ⁻¹ is an [(n-1)-simplex](https://en.wikipedia.org/wiki/Simplex)
It has no defined inverse
"""
SoftMaxLink() = Link(softmax)


function golomb(n, L, ::Val{:raw})
    m = model(; kind=:golomb)

    # Add variables
    d = domain(0:L)
    foreach(_ -> variable!(m, d), 1:n)

    # Extract error function from usual_constraint
    e1 = (x; param=nothing, dom_size=n) -> error_f(
        usual_constraints[:all_different])(x; param, dom_size
    )
    e2 = (x; param=nothing, dom_size=n) -> error_f(
        usual_constraints[:all_equal_param])(x; param, dom_size
    )
    e3 = (x; param=nothing, dom_size=n) -> error_f(
        usual_constraints[:dist_different])(x; param, dom_size
    )

    # # Add constraints
    constraint!(m, e1, 1:n)
    constraint!(m, x -> e2(x; param=0), 1:1)
    for i in 1:(n - 1), j in (i + 1):n, k in i:(n - 1), l in (k + 1):n
        (i, j) < (k, l) || continue
        constraint!(m, e3, [i, j, k, l])
    end

    # Add objective
    objective!(m, o_dist_extrema)

    return m
end

function golomb(n, L, ::Val{:JuMP})
    m = JuMP.Model(CBLS.Optimizer)

    @variable(m, 0 ≤ X[1:n] ≤ L, Int)

    @constraint(m, X in AllDifferent()) # different marks
    @constraint(m, X in Ordered()) # for output convenience, keep them ordered

    # No two pairs have the same length
    for i in 1:(n - 1), j in (i + 1):n, k in i:(n - 1), l in (k + 1):n
        (i, j) < (k, l) || continue
        @constraint(m, [X[i], X[j], X[k], X[l]] in DistDifferent())
    end

    # Add objective
    @objective(m, Min, ScalarFunction(maximum))

    return m, X
end

"""
    golomb(n, L=n²)

Model the Golomb problem of `n` marks on the ruler `0:L`. The `modeler` argument accepts :raw, and :JuMP (default), which refer respectively to the solver internal model, the MathOptInterface model, and the JuMP model.
"""
golomb(n, L=n^2; modeler = :JuMP) = golomb(n, L, Val(modeler))


export AR1process,
    TimeSeriesProcess,
    RandomWalk,
    AbstractAR1Process

abstract type TimeSeriesProcess end
abstract type AbstractAR1Process{T<:Real} <: TimeSeriesProcess end

# ----------------------------
# AR1
# ----------------------------

struct AR1process{T<:Real} <: AbstractAR1Process{T}
    mu::T
    ar::T
    sigsq::T
    function AR1process(mu::T,ar::T,sigsq::T) where {T}
        -1 < ar < 1 || @warn("nonstationary: |ar| parameter >= 1")
        sigsq > 0 || throw(DomainError(sigsq))
        return new{T}(mu,ar,sigsq)
    end
end

mu(x::AR1process) = x.mu
ar(x::AR1process) = x.ar
sigsq(x::AbstractAR1Process) = x.sigsq
condvar(x::AbstractAR1Process) = sigsq(x)
condstd(x::AbstractAR1Process) = sqrt(condvar(x))
lrvar(x::AR1process) = sigsq(x) / (1 - ar(x)^2)
lrstd(x::AR1process) = sqrt(lrvar(x))
condmean(x::AbstractAR1Process, yt) = mu(x) + ar(x)*yt
lrmean(x::AR1process) = mu(x)/(1-ar(x))

# ----------------------------
# Random walk
# ----------------------------

struct RandomWalk{T} <: AbstractAR1Process{T}
    sigsq::T
    function RandomWalk(sigsq::T) where {T}
        sigsq > 0 || throw(DomainError(sigsq))
        return new{T}(sigsq)
    end
end

ar(x::RandomWalk) = 1
mu(x::RandomWalk) = 0

# ----------------------------
# Bivariate Brownian
# ----------------------------

function tauchen_1d(x::AbstractAR1Process, xgrid::AbstractRange)
    return tauchen_1d(xgrid, xt -> condmean(x,xt), condvar(x))
end

function tauchen_1d(x::AR1process, n::Integer; kwargs...)
    xgrid = MakeGrid(x, n; kwargs...)
    return tauchen_1d(xgrid, xt -> condmean(x,xt), condvar(x))
end

# ----------------------------
# Make Grid only
# ----------------------------

function MakeGrid(x::AR1process, n::Real; m=3)
    lrstdev = sqrt(lrvar(x))
    centered = range(-m*lrstdev; stop=m*lrstdev, length=n)
    return lrmean(x) .+ centered
end

function MakeGrid(x::AR1process, n::Real, z::AbstractVector; m=3)
    rng = MakeGrid(x,n)
    xz = (x,z)
    len = length(rng)
    a = minimum(minimum.(xz))
    b = maximum(maximum.(xz))
    return range(a; stop=b, length=len)
end

function MakeGrid(p::AR1process, zvec, length=15; process=TestPriceProcess())
    zmin = min(minimum(zvec), lrmean(process)-3*lrstd(process))
    zmax = max(maximum(zvec), lrmean(process)+3*lrstd(process))
    zrng = range(zmin, zmax; length=length)
    return tuple(zrng)
end

@deprecate approxgrid(args...) MakeGrid(args...)

# ----------------------------
# simulations
# ----------------------------

starting_dist(ar1::AR1process) = Normal(lrmean(ar1), lrstd(ar1))
starting_dist(ar1::RandomWalk) = Normal(0, condstd(ar1))

function simulate(ar1::AbstractAR1Process{T}, y0::Real, n::Integer) where {T}
    dist_cond   = Normal(zero(T), condstd(ar1))
    y = Vector{T}(undef,n)
    y[1] = y0
    Distributions.rand!(dist_cond, view(y, 2:n))
    for t in 2:n
        y[t] += condmean(ar1, y[t-1])
    end
    return y
end

function simulate(ar1::AbstractAR1Process, n::Integer)
    dist_uncond = starting_dist(ar1)
    y0 = rand(dist_uncond)
    simulate(ar1, y0, n)
end


Base.@kwdef mutable struct Benchmark
    b::Vector{Float64} # benchmark weights
    c::Vector{Float64} # Climate costs per $ mv
    σ::Vector{Float64} # Volatility
    ρ::Matrix # Correlation matrix
    Σ::Union{Matrix,Nothing} = nothing # Covariance Matrix
end

ρ = [
    1 0.8 0.7 0.6 0.7 0.5 0.7 0.6
    0.8 1 0.75 0.65 0.5 0.6 0.5 0.65
    0.7 0.75 1 0.8 0.7 0.7 0.7 0.7
    0.6 0.65 0.8 1 0.85 0.8 0.75 0.75
    0.7 0.5 0.7 0.85 1 0.6 0.8 0.65
    0.5 0.6 0.7 0.8 0.6 1 0.5 0.7
    0.7 0.5 0.7 0.75 0.8 0.5 1 0.8
    0.6 0.65 0.7 0.75 0.65 0.7 0.8 1
]

our_benchmark = Benchmark(b = [0.23, 0.19, 0.17, 0.13, 0.09, 0.08, 0.06, 0.05],
                        c= [0.06, 0.8, 0.39, 0.01, 0.05, 0.13, 0.5, 0.45],
                        σ = [0.22, 0.20, 0.25, 0.18, 0.35, 0.23, 0.13, 0.29],
                        ρ = ρ)

function get_benchmark_climate_costs(bench::Benchmark, mv::Int)::Float64
    return - round(bench.b' * bench.c * mv; digits= 2)
end

get_benchmark_climate_costs(our_benchmark, 1000)

Base.@kwdef mutable struct TiltedPortfolio
    x::Vector{Float64} # Optimal weights
end

other_portfolio = TiltedPortfolio(x = [0.30, 0.10, 0.08, 0.02, 0.48, 0.01, 0.005, 0.0055])

function get_climate_benefits(portfolio::TiltedPortfolio, bench::Benchmark, mv::Int)::Float64
    return -1 * round((portfolio.x - bench.b)' * bench.c * mv;digits = 5)
end

get_climate_benefits(other_portfolio, our_benchmark,1000)

check = other_portfolio.x' * our_benchmark.c * 1000

using LinearAlgebra
function get_cov_from_corr!(bench::Benchmark)::Benchmark
    D = diagm(bench.σ) # diagonal matrix of volatiltiy
    bench.Σ = D * bench.ρ * D # multiplying columns and rows of correlation matrix by volatiltiy
    return bench
end

get_cov_from_corr!(our_benchmark)

function tracking_error_volatility(portfolio::TiltedPortfolio, bench::Benchmark)::Float64
    return sqrt((portfolio.x - bench.b)' * bench.Σ * (portfolio.x - bench.b))
end

# get_excess_score(portfolio, our_benchmark)

# tracking_error_volatility(portfolio, our_benchmark)

using JuMP, COSMO # for the optimization

# tilt the benchmark according to gamma
function portfolio_tilting(bench::Benchmark, γ::Float64)::TiltedPortfolio
    n = length(bench.c) # number of assets
    model = JuMP.Model(COSMO.Optimizer)
    # the optimal weights we want to find
    @variable(model, x[1:n])
    @objective(model, Min, 1/2 * x' * bench.Σ * x - x' * (γ * -1 * bench.c + bench.Σ * bench.b))
    @constraint(model, zeros(n) .<= x .<= ones(n))
    @constraint(model, ones(n)' * x == 1)
    JuMP.optimize!(model)
    portfolio = TiltedPortfolio(x = JuMP.value.(x))
    return portfolio
end

# simulate for gammas between 0 and 10
using Plots
gammas = [i for i in 0.0:0.001:0.012]
climate_benefits = zeros(length(gammas))
excess_te = zeros(length(gammas))

for i in 1:length(gammas)
    new_portfolio = portfolio_tilting(our_benchmark, gammas[i])
    climate_benefits[i] = get_climate_benefits(new_portfolio, our_benchmark, 1000) # suppose 1000 billion of MV for the benchmark
    excess_te[i] = tracking_error_volatility(new_portfolio, our_benchmark) * 100
end

plot(excess_te, climate_benefits, xlabel = "Tracking error volatility (in %)",
                                ylabel = "Climate benefits for society (in billion USD)",
                                label = "")


@testset "Potential" begin
    # CarrierCapture.harmonic potential generation and finding the crossing
    Qmin = -25
    Qmax = 35
    NQ = 3000
    Q = range(Qmin, stop=Qmax, length=NQ)
    nev = 30
    hw = 0.02

    pot1 = potential(); pot1.name = "pot1"
    pot1.Q0 = 0; pot1.E0 = 1
    pot1.nev = nev
    pot1.func = x -> CarrierCapture.harmonic(x, hw; E₀ = pot1.E0, Q₀ = pot1.Q0)
    pot1.Q = Q
    pot1.E = pot1.func.(Q)

    pot2 = potential(); pot2.name = "pot2"
    pot2.Q0 = 2; pot2.E0 = 0
    pot2.nev = nev
    pot2.func = x -> CarrierCapture.harmonic(x, hw; E₀ = pot2.E0, Q₀ = pot2.Q0)
    pot2.Q = Q
    pot2.E = pot2.func.(Q)

    solve_pot!(pot1)
    solve_pot!(pot2)

    # eigenvalue test
    @test isapprox(pot1.ϵ, 1 .+ pot2.ϵ ; rtol=1E-5)

    Qx, Ex = find_crossing(pot1, pot2)
    @test Qx ≈ -4.22519859246400
    @test Ex ≈ 1.8541447195665985

    # Spline fit test
    nev = 60
    name = "D"

    Q1 = [30.66721918 29.860133 29.05306268 28.24612992 27.43911044 26.63197583 25.82507425 25.01797262 24.21096115 23.40387798 22.59690043 21.78991117 20.98281579 20.17570172 19.36884219 18.56169249 17.75463179 16.94772679 16.14061031 15.33347439 14.52663309 13.71945696 12.91240658 12.10544441 11.29841968 10.4913338 9.684370388 8.877289725 8.070184138]
    E1 = [8.0902 7.5970 7.0749 6.5242 5.9461 5.3451 4.7300 4.1147 3.5182 2.9637 2.4769 2.0819 1.7972 1.6315 1.5800 1.6237 1.7301 1.8586 1.9687 2.0283 2.0190 2.0673 1.9910 2.0528 1.9730 2.0857 2.4550 3.1653 4.3448]

    pot = potential(); pot.name = name
    pot.nev = nev
    pot.Q0 = Q1[findmin(E1)[2]]; pot.E0 = 1.69834
    pot.QE_data = DataFrame(Q = Q1[:], E = E1[:])
    pot.QE_data.E .+= - minimum(pot.QE_data.E) + pot.E0
    pot.Q = Q

    pot.func_type = "spline"
    ## spline fitting parameters
    params = Dict()
    params["order"] = 4
    params["smoothness"] = 0.001
    params["weight"] = [1 1 1 1 1 0.5 0.4 0.4 0.5 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1]

    fit_pot!(pot, Q; params=params)
    @test pot.E == pot.func(Q)

    fit_pot!(pot, Q)
    @test pot.E == pot.func(Q)

    pot.func_type = "bspline"
    fit_pot!(pot, Q)
    @test pot.E == pot.func(Q)

    pot.func_type = "polyfunc"
    params["poly_order"] = 4
    params["p0"] = [1, 2, 3, 4, 6]
    fit_pot!(pot, Q; params=params)
    @test pot.E ≈ pot.func(Q) atol = 1e-8

    # potential types
    # CarrierCapture.harmonic
    # double well
    # polyfunc
    # morse
    # get_spline
    # get_bspline
end


using Distributions
using StatUtils
using Test
using StableRNGs

function test_bootstrapper()
    RNG = StableRNG(1234)
    x = rand(RNG, 1000)
    B = StatUtils.run_bootstrap(1000,sum,x;RNG=RNG)
    @test B.μ ≈ 508.1411134730461
    @test B.σ ≈ 9.045400206380133
    println("Bootstrapper test passed")

end

function test_fdr()
	pvalues = [0.0001, 0.0004, 0.0019, 0.0095, 0.0201, 0.0278, 0.0298, 0.0344, 0.0459, 0.3240, 0.4262, 0.5719, 0.6528, 0.7590, 1.00]
	pidx = StatUtils.fdr(pvalues)
	@test pidx == [1,2,3,4]
	println("FDR test passed")
end

function test_groupby()
	groups = StatUtils.groupby([1,2,3], [1,2,3])
	@test Dict{Int64, Array{Int64,1}}(1=>[1], 2=>[2], 3=>[3]) == groups
	println("groupby test passed")
end

@testset "Bootstrapper" begin
    test_bootstrapper()
end
@testset "False Discovery Rate" begin
    test_fdr()
end
@testset "Group by" begin
    test_groupby()
end

@testset "Bootstrap median" begin
    RNG = StableRNG(1234)
    x1 = rand(RNG, Float64, (1,100))
    μ,σ = StatUtils.bootstrap_median(x1, 10_000, RNG)
    @test μ[1] ≈ 0.4977072659335834
    @test σ[1] ≈ 0.04304170028498684
end

@testset "Spearman pv" begin
    #no correlation
    RNG = StableRNG(1234)
    x1 = rand(RNG, 100)
    x2 = rand(RNG, 100)
    cc,pv = StatUtils.spearmanr(x1,x2, :both, RNG)
    @test cc ≈ -0.023438343834383438
    @test pv ≈ 0.999

    #small non-significant correlation
    RNG = StableRNG(1234)
    x1 = rand(RNG, 100)
    x2 = 0.1*x1 +  0.9*rand(RNG, 100)
    cc,pv = StatUtils.spearmanr(x1,x2, :right, RNG)
    @test cc ≈ 0.09610561056105611
    @test pv ≈ 0.189
    #strong significant correlation
    RNG = StableRNG(1234)
    x1 = rand(RNG, 100)
    x2 = 0.5*x1 +  0.5*rand(RNG, 100)
    cc,pv = StatUtils.spearmanr(x1,x2, :right, RNG)
    @test cc ≈ 0.6623702370237023
    @test pv ≈ 0.0
end

@testset "Percentile" begin
    @test_throws ArgumentError Percentile(-0.1)
    @test_throws ArgumentError Percentile(101.0)
    @test typemax(Percentile) == 100.0
    @test typemin(Percentile) == 0.0
end

@testset "Generalized Poisson" begin
    G = GeneralizedPoisson(1.5370664608871476,-0.2658194383776511)
    @test pdf(G,0) ≈ 0.21501092012028916
    @test log(pdf(G,0)) ≈ logpdf(G,0)
    @test log(pdf(G,1)) ≈ logpdf(G,1)
    @test mean(G) ≈ 1.2142857142857142
    @test var(G) ≈ 0.7578397212543557
    @test sum(pdf.(G,0:7)) ≈ 1.0000005196278676
end

@testset "Bootstrap regression" begin
    #generate data
    RNG = StableRNG(1234)
    x = range(0.0, stop=1.0, length=50)
    y = 0.5 .+ 0.5*x .+ 0.05*randn(RNG, length(x))
    μ, σ, xx = StatUtils.bootstrap_regression(x,y;RNG=RNG)
    sse = sum(abs2, μ-y)
    sq = sum(abs2, y .- 0.5)
    @test sse/sq < 0.05
end

@testset "Robust regression" begin
    RNG = StableRNG(1234)
    x = range(0.0, stop=1.0, length=20)
    y = 0.5 .+ 0.1*x .+ 0.3*randn(RNG, length(x))
    β0 = rand(RNG, Float64, 2)
    β = StatUtils.robust_regression(x,y, β0)
    @test β[1] ≈ 0.23571370032430217
    @test β[2] ≈ 0.38013900282542895
end

@testset "Angular histogram" begin
    θ = [0.0, π/4, π/2, -π, -π/4, -π/2, -3*π/4]
    counts, binidx, bins = StatUtils.angular_histogram(θ, π/4)
    @test bins ≈ [-π, -3*π/4, -π/2, -π/4, 0.0, π/4, π/2, 3π/4]
    @test counts == [1, 1, 1, 1, 1, 1, 1, 0]
    @test binidx == [5, 6, 7, 1, 4, 3, 2]
end

@testset "Grouping" begin
    x = [1,3,3,1,5,7,7,1,5]
    y = StatUtils.compress(x)
    @test y == [1,2,2,1,3,4,4,1,3]
    X = fill(0.0, 3,4,10,6)
    grouping = fill(0, 10)
    grouping[1:5] .= 1
    grouping[6:10] .= 2
    X[:,:,findall(grouping.==1),:] .= reshape([1.0,2.0,3.0,4.0,5.0], 1,1,5,1)
    X[:,:,findall(grouping.==2),:] .= reshape([6.0, 7.0, 8.0, 9.0, 10.0], 1, 1, 5,1)
    sidx = [7, 10, 2, 5, 8, 6, 1, 3, 9, 4]
    X = X[:,:,sidx,:]
    μ = mean(X, grouping[sidx])
    @test size(μ) == (3,4,2,6)
    @test μ[:,:,1,:] ≈ fill(3.0, 3,4,6)
    @test μ[:,:,2,:] ≈ fill(8.0, 3,4,6)
end


using Measures
include("helper_functions.jl")

begin
    λ1 = 1
    λ2 = 1
    α = 1
    μ(α) = 1 / (1.5 * (α))
    ν(α) = 1 / (0.7 * (α))
    parameters = (λ1, λ2, α, μ, ν)
    opponent_parameters = (λ2, λ1, α, μ, ν)
    p, q = equilibrium_probability(parameters...)
    iter = 0:0.001:1
end

cost(p, q, parameters...) ≈ cost(q, p, opponent_parameters...)
costs = cost.(iter, q, parameters...)

function find_asymptotes(vec, iter)
    sv = sign.(vec)
    iter[sv .!= circshift(sv, 1)]
end

asymptotes = find_asymptotes(costs, iter)

nash_plot = plot(iter[costs .> 0], costs[costs .> 0], label = "Response Time")
plot!(ylims = (0, 25), xlims = (0, 1), size = (600, 250), palette = :tab10)
plot!([p, p], [0, cost(p, q, parameters...)], ls = :dash, lc = :black, label = "Minimum")
plot!([0, p], [cost(p, q, parameters...), cost(p, q, parameters...)], ls = :dash, lc = :black, label = "")
vline!([asymptotes], ls = :dash, label = "Asymptotes")
plot!(xlabel = "Proportion Sent to Public Server", ylabel = "Expected Response Time", legend = :topleft, margin = 2.5mm)
savefig("outputs/nash_plot_poster.pdf")


abstract type ActivationFunction end
abstract type NaiveActivationFunction <: ActivationFunction end
abstract type TrainedActivationFunction <: ActivationFunction end

struct Cube <: NaiveActivationFunction end
(::Cube)(x::T) where {T <: Number} = x^3.0

struct Linear <: NaiveActivationFunction end
(::Linear)(x::T) where {T <: Number} = x

struct ReLU <: NaiveActivationFunction end
(::ReLU)(x::T) where {T <: Number} = max(zero(x), x)

struct Sigmoid <: NaiveActivationFunction end
(::Sigmoid)(x::T) where {T <: Number} = 1.0 / (1.0 + exp(-x))

struct Square <: NaiveActivationFunction end
(::Square)(x::T) where {T <: Number} = x^2.0

struct Tanh <: NaiveActivationFunction end
(::Tanh)(x::T) where {T <: Number} = tanh(x)

struct KDEComparator{T₁ <: KernelDensity.UnivariateKDE,
                     T₂ <: KernelDensity.UnivariateKDE} <: TrainedActivationFunction
    pdf₁::KernelDensity.InterpKDE{T₁}
    pdf₂::KernelDensity.InterpKDE{T₂}

    function KDEComparator(samples::T₁,
                           targets::T₂,
                           weights::T₃;
                           boundary_offset::Int = 1) where {T₁ <: AbstractMatrix,
                                                            T₂ <: AbstractVector,
                                                            T₃ <: AbstractVector}
        X = samples
        t = targets
        w = weights

        h = vec(w' * X)
        h₀ = h[t .!= 1]
        h₁ = h[t .== 1]

        boundary = (min(minimum(h₀), minimum(h₁)) - boundary_offset,
                    max(maximum(h₀), maximum(h₁)) + boundary_offset)

        kde₀ = KernelDensity.kde_lscv(h₀, npoints = 100, boundary = boundary)
        kde₁ = KernelDensity.kde_lscv(h₁, npoints = 100, boundary = boundary)

        pdf₀ = KernelDensity.InterpKDE(kde₀)
        pdf₁ = KernelDensity.InterpKDE(kde₁)

        param(::KernelDensity.InterpKDE{T}) where {T} = T
        new{param(pdf₀),param(pdf₁)}(pdf₀, pdf₁)
    end
end

function (c::KDEComparator)(x::T) where {T <: Number}
    Distributions.pdf(c.pdf₁, x) - Distributions.pdf(c.pdf₂, x)
end

function train_kde_comparators(::Type{T₁},
                               samples::T₂,
                               targets::T₃,
                               n_neurons::Int;
                               boundary_offset::Int = 1) where {T₁ <: Number,
                                                                T₂ <: AbstractMatrix,
                                                                T₃ <: AbstractVector}
    X = samples
    t = targets
    L = n_neurons
    D = first(size(X))
    W = gaussian_projection_matrix(T₁, L, D)
    fs = [KDEComparator(X, t, W[l,:], boundary_offset = boundary_offset) for l in 1:L]
    W, fs
end


"""
	similarity_order(wv::WordVectors, alpha=-0.65)

Post-processes the word embeddings `wv` so that the embeddings capture more
information than directly apparent through a linear transformation that
adjusts the similarity order of the model. The function returns a new
`WordVectors` object containing the processed embeddings.

# Arguments
  * `wv::WordVectors` the word embeddings
  # `alpha::AbstractFloat` the `α` parameter of the algorithm (default -0.65)

# References:
  * [Artetxe et al. \"Uncovering divergent linguistic information in
     word embeddings with lessons for intrinsic and extrinsic evaluation\",
     2018](https://arxiv.org/pdf/1809.02094.pdf)
"""
function similarity_order(wv::WordVectors{S,T,H},
                          alpha::T=T(-0.65)
                         ) where {S<:AbstractString, T<:Real, H<:Integer}
    n = length(wv.vocab)
    X = wv.vectors
    M = X*X'
    L, Q = eigen(M)
    Wₐ= Q .* (L.^alpha)
    return WordVectors{S,T,H}(wv.vocab, Wₐ*X, wv.vocab_hash)
end


function orthonormalize!( LF::LF3dGrid, v )
    ΔV = LF.LFx.h * LF.LFy.h * LF.LFz.h
    ortho_gram_schmidt!( v )
    v[:,:] = (1/sqrt(ΔV))*v  # the [:,:] is significant
end

function orthonormalize( LF::LF3dGrid, v )
    ΔV = LF.LFx.h * LF.LFy.h * LF.LFz.h
    return ortho_gram_schmidt( v )/sqrt(ΔV)
end


@testset "Combined Mean and Variance" begin

    @testset "Two samples" begin
        test_two = (x1,x2) -> begin
            xc = vcat(x1,x2)
            meanc, varc = combined_mean_and_var(x1,x2)
            @test abs(mean(xc) - meanc) < 1e-12
            @test abs(var(xc) - varc) < 1e-12
        end

        x1 = rand(30_000)
        x2 = rand(20_000)
        test_two(x1,x2)

        x1 = rand(ComplexF64, 30_000)
        x2 = rand(ComplexF64, 20_000)
        test_two(x1,x2)
    end

    @testset "Three samples" begin
        test_three = (x1,x2,x3) -> begin
            xc = vcat(x1,x2,x3)
            meanc, varc = combined_mean_and_var(x1,x2,x3)
            @test abs(mean(xc) - meanc) < 1e-12
            @test abs(var(xc) - varc) < 1e-12

            meanc_three, varc_three = combined_mean_and_var_three(x1,x2,x3)
            @test isapprox(meanc, meanc_three)
            @test isapprox(varc, varc_three)
        end

        x1 = rand(30_000)
        x2 = rand(20_000)
        x3 = rand(40_000)
        test_three(x1, x2, x3)

        x1 = rand(ComplexF64, 30_000)
        x2 = rand(ComplexF64, 20_000)
        x3 = rand(ComplexF64, 40_000)
        test_three(x1, x2, x3)
    end

    @testset "N samples" begin
        test_N = (xs...) -> begin
            xc = vcat(xs...)
            meanc, varc = combined_mean_and_var(xs...)
            @test abs(mean(xc) - meanc) < 1e-12
            @test abs(var(xc) - varc) < 1e-12
        end

        lengths = [30_000, 20_000, 40_000]
        N = 5

        xs = [rand(Float64, rand(lengths)) for _ in 1:N]
        test_N(xs...)

        xs = [rand(ComplexF64, rand(lengths)) for _ in 1:N]
        test_N(xs...)
    end

    @testset "N samples (ns, μs, vs)" begin
        test_N_moments = (ns, μs, vs, mean_exact, var_exact) -> begin
            meanc, varc = combined_mean_and_var(ns, μs, vs)
            @test abs(mean_exact - meanc) < 1e-12
            @test abs(var_exact - varc) < 1e-12
        end

        lengths = [30_000, 20_000, 40_000]
        N = 5

        xs = [rand(Float64, rand(lengths)) for _ in 1:N]
        xc = vcat(xs...)
        test_N_moments(length.(xs), mean.(xs), var.(xs), mean(xc), var(xc))

        xs = [rand(ComplexF64, rand(lengths)) for _ in 1:N]
        xc = vcat(xs...)
        test_N_moments(length.(xs), mean.(xs), var.(xs), mean(xc), var(xc))
    end
end


###################################################
## rect.jl
## Rectangular virtual city
###################################################

"""
  `rectNetwork` : n*m points rectangular network
 - all road same type (5)
 - `distX` and `distY` are the width and height of a block (in meters)
"""
function rectNetwork(width::Int = 5, height::Int = 10; distX::Float64=200., distY::Float64 = 200.)
        #nodes are indexed as follow :
        #  12
        #  34 (for width=2,2)

        function coordToLoc(i,j) # matrix indexation, i= row j=column
            return j + (i-1)*width
        end
        function locToCoord(n)
            return (div((n-1),width) + 1, ((n-1) % width) + 1)
        end

        #Create nodes
        nodes = Array{Node}(width*height)
        for i in 1:height, j in 1:width
            nodes[coordToLoc(i,j)] = Node((j-(width+1)/2)*distX, ((height+1)/2 - i)*distY)
        end
        roads = Dict{Tuple{Int,Int},Road}()

        #Vertical roads
        for i in 1:(height-1), j in 1:width
            n1 = coordToLoc(i,j); n2 = coordToLoc(i+1,j)
            dist = distanceCoord(nodes[n1],nodes[n2])
            roads[(n1,n2)] = Road(n1,n2,dist, 5)
            roads[(n2,n1)] = Road(n2,n1,dist, 5)
        end

        #Horizontal roads
        for i in 1:height, j in 1:width - 1
            n1 = coordToLoc(i,j); n2 = coordToLoc(i,j+1)
            dist = distanceCoord(nodes[n1],nodes[n2])
            roads[(n1,n2)] = Road(n1,n2,dist, 5)
            roads[(n2,n1)] = Road(n2,n1,dist, 5)
        end
        return Network(nodes, roads)
end

squareNetwork(width::Int = 5; distance::Float64=200.) = rectNetwork(width, width, distX=distance, distY=distance)


module Lattice1D
    using SparseArrays
    export H0, Hnoise, Hamiltonian
    """
        H0(N; ϵ=0., ν=1.)

    Constructs the noiseless Hamiltonian for the 1D lattice with `N` nodes,
    self-energy ``\\epsilon`` and tunneling constant ``\\nu``.
    """
    function H0(n; ϵ=0., ν=1.)
        H = spdiagm(0 => ϵ*ones(n), 1=>ν*ones(n-1), -1=> ν*ones(n-1))
        if n > 2
            H[1,end] = H[end,1] = ν
        end
        return H
    end

    """
        Hnoise(N, noise_id)

    Returns the noisy part of the Hamiltonian for the 1D lattice with `N` nodes
    for a specific state of the noise fluctuators specified by `noise_id`.
    """
    function Hnoise(n, noise_id)
        if noise_id < 1 || noise_id > 2^n
            throw(ArgumentError("noise_id must be between 1 and 2^n"))
        end
        noise = 2 * digits(noise_id .- 1, base=2, pad=n) .- 1
        H = spdiagm(-1 => noise[2:end], 1 => noise[2:end])
        H[1, end] = H[end, 1] = noise[1]
        return H
    end

    """
        Hamiltonian(n, ν, ν0=1.)

    Returns an array of Hamiltonians of size `n` for all possible configurations
    of noise.

    # Arguments
        * `n::Integer`: the dimension of the Hilbert space
        * `ν::Float`: the strength of the noise
        * `ν0::Float`: The rate of the noiseless Hamiltonian (ν0 = 1.)
    """
    function Hamiltonian(n::Integer, ν::Float64, ν0=1.)
        return [H0(n;ν=ν0) + ν*Hnoise(n, id) for id = 1:2^n]
    end

end


using PolyChaos, Test, LinearAlgebra
import QuadGK: quadgk

function integration_test(opq::OrthoPolyQ,dim::Int64,name::String;tol::Float64=1e-6,
                          dom::Tuple{Float64,Float64}=opq.op.meas.dom)
    op::OrthoPoly = opq.op
    N = op.deg
    w::Function = op.meas.w
    ind = zeros(Int64,dim)
    @testset "$name" begin
        for ind_ in Iterators.product([collect(1:N) for i=1:dim]...)
            ind[:] .= ind_[:]
            s1 = computeSP(ind,opq)
            f(t) = prod(evaluate(ind[i],t,op) for i=1:dim)*w(t)
            # f(t) = evaluate(ind[1],t,op)*evaluate(ind[2],t,op)*evaluate(ind[3],t,op)*w(t)
            s2 = quadgk(f,dom[1],dom[2])[1]
            if abs(s1)<=1e-3
                @test isapprox(s1,s2;atol=tol)
            else
                @test isapprox(abs(s1/s2),1;atol=tol)
            end
        end
    end
end
##

N, Nrec, dim, tol = 3, 1000, 3, 1e-4

op = OrthoPoly("gaussian",N;Nrec=Nrec);
opq = OrthoPolyQ(op);
integration_test(opq,dim,"gaussian";tol=tol,dom=(-10.,10.))

α, β = 1.32, 3.23
op = OrthoPoly("beta01",N,Dict(:shape_a=>α,:shape_b=>β);Nrec=Nrec);
opq = OrthoPolyQ(op);
set_beta = integration_test(opq,dim,"beta01";tol=tol)

op = OrthoPoly("gamma",N,Dict(:shape=>1.0,:rate=>1.);Nrec=Nrec);
opq = OrthoPolyQ(op);
set_gamma = integration_test(opq,dim,"gamma";tol=tol,dom=(0.,Inf))


using IterTools
using ProgressMeter

function find_first_invalid_number(input_numbers, window_length)

    @showprogress for (j, i) in enumerate(window_length + 1:length(input_numbers))
        input_subset = input_numbers[j:j + window_length - 1]
        valid_sums = Set([sum(subset) for subset in subsets(input_subset, 2)])
        if !(input_numbers[i] in valid_sums)
            return (input_numbers[i], i)
        end
    end
end

# Invalid number is 507622668

function find_consecutives_that_sum_to_total(input_numbers, bad_number, bad_number_idx)
    @showprogress for i in 1:bad_number_idx
        j = i + 1
        running_sum = 0
        sequence = []
        while running_sum < bad_number
            running_sum += input_numbers[j]
            push!(sequence, input_numbers[j])
            j += 1

            if running_sum == bad_number
                println("Sequence: $sequence")
                min_seq = minimum(sequence)
                max_seq = maximum(sequence)
                println("sum of min and max: $(min_seq + max_seq)")
            end
        end
    end
end

function parse_file_to_input(input_file)
    input_numbers = open(input_file) do file
        file_string = read(file, String)
        [parse(Int64, line) for line in split(file_string, "\r\n")]
    end
    input_numbers
end

function main(input_file, window_length)
    input_numbers = parse_file_to_input(input_file)
    bad_numbers_and_idx = find_first_invalid_number(input_numbers, window_length)
    println("$bad_numbers_and_idx")
end

function part2(input_file)
    input_numbers = parse_file_to_input(input_file)
    find_consecutives_that_sum_to_total(input_numbers, 507622668, 634)
end

# main("day9_input.txt", 25)
part2("day9_input.txt")


######################################################################
# corner_initializer.jl
######################################################################

"""
`ci = CornerInitializer()`

Takes a trajectory to the farthest corner.
"""
mutable struct CornerInitializer <: Initializer end

# Assumes we our domain is 2d
# corner 1 = 0,0
# corner 2 = 1,0
# corner 3 = 0,1
# corner 4 = 1,1
function initialize(ci::CornerInitializer, em::ErgodicManager, tm::TrajectoryManager)

	# dimensionality of state space
	x0 = tm.x0
	n = length(x0)

	# first corner: 0,0
	dx1 = 0.0 - x0[1]
	dy1 = 0.0 - x0[2]
	dc1 = sqrt(dx1*dx1 + dy1*dy1)

	# second corner: 1,0
	dx2 = 1.0 - x0[1]
	dy2 = 0.0 - x0[2]
	dc2 = sqrt(dx2*dx2 + dy2*dy2)

	# third corner: 0,1
	dx3 = 0.0 - x0[1]
	dy3 = 1.0 - x0[2]
	dc3 = sqrt(dx3*dx3 + dy3*dy3)

	# fourth corner: 1,1
	dx4 = 1.0 - x0[1]
	dy4 = 1.0 - x0[2]
	dc4 = sqrt(dx4*dx4 + dy4*dy4)

	bi = indmax([dc1,dc2,dc3,dc4])

	bc = [0.,0.]
	if bi == 2
		bc = [1.,0.]
	elseif bi == 3
		bc = [0.,1.]
	elseif bi == 4
		bc = [1.,1.]
	end

	# could now do point initializer to the selected corner
	return initialize(PointInitializer(bc), em, tm)
end


"""
    x=cvp(y,H,infinite=Val{true},Umax=1)

Solve the problem `argmin_x ||y-Hx||` for integer x using the technique from
the paper below.  The input vector `y` is of length `n`, with `H` of
dimension `n` by `n`, and the returned vector `x` of length `n`. If
`infinite==Val{true}` then we search the (infinite) lattice, otherwise we
search integers in `[-Umax,Umax]`.  At present cvp does not handle complex
numbers.

Uses alg from "Faster Recursions in Sphere Decoding" Arash Ghasemmehdi, Erik
Agrell, IEEE Transactions on Information Theory, vol 57, issue 6 , June 2011.

# Examples
```jldoctest
julia> H=[1 2; 3 4]; Q,R=qr(H); uhat = cvp(Q'*[0,2],R)
2-element Array{Float64,1}:
  2.0
 -1.0

julia> n=100;H=randn(n,n);Q,R=qr(H);

julia> u=Int.(rand(0:1e10,n));y=H*u+rand(n)/100;

julia> uhat=cvp(Q'*y,R); sum(abs.(u-uhat))
0.0

julia> n=500;H=randn(n,n);Q,R=qr(H);

julia> u=Int.(rand(-1:1,n));y=H*u+rand(n)/10;

julia> uhat=cvp(Q'*y,R,Val{false},1); sum(abs.(u-uhat))
0.0
```
"""
function cvp(r::AbstractArray{Td,1},G::AbstractArray{Td,2},
               infinite=Val{true},Umax=1) where {Td<:Number}

    roundGA(x) = roundFinite(x,Td(Umax))
# nx=1
    n = length(r)
# mxNx=Int(ceil(log2(n)*100000))
    C = Inf
    i = n+1
    d = ones(Int,n)*n
    λ = zeros(n+1)
    F = zeros(Td,n,n)
    F[:,n] = r
    p = zeros(n)
    u = NaN*zeros(Int,n)
    û = NaN*zeros(Int,n)
    Δ = zeros(Int,n)

    begin
        @label LOOP

#        nx+=1
        while λ[i]<C
            if i != 1
                i-=1
                for j=d[i]:-1:i+1
                    F[i,j-1] = F[i,j] - G[i,j]*u[j]
                end
                p[i] = F[i,i]/G[i,i]
                if infinite==Val{true}
                    u[i] = round(p[i])
                else
                    u[i] = roundGA(p[i])
                end
                y = (p[i]-u[i])*G[i,i]
                Δ[i] = signGA(y)
                λ[i] = λ[i+1]+y*y
            else
                û .= u
                C = λ[1]
            end
        end
        m = i

        while λ[i]≥C
#            if i==n || nx>mxNx
            if i==n
                # if nx>mxNx
                #     @warn("terminated search after $(nx) iterations.")
                # end
                return û
            else
                i+=1
                if infinite ≠ Val{true}
                    y = typemax(Td)
                end
                u[i]+=Δ[i]
                Δ[i]=-Δ[i]-signGA(Δ[i])
                if infinite ≠ Val{true}
                    if -Umax ≤ u[i] ≤ Umax
                        y = (p[i]-u[i])*G[i,i]
                    else
                        u[i]+=Δ[i]
                        Δ[i]=-Δ[i]-signGA(Δ[i])
                        if -Umax ≤ u[i] ≤ Umax
                            y = (p[i]-u[i])*G[i,i]
                        end
                    end
                else
                    y = (p[i]-u[i])*G[i,i]
                end
                λ[i] = λ[i+1]+y*y
            end
        end
        d[m:i-1] .= i
        for j=m-1:-1:1
            if d[j]<i
                d[j]=i
            else
                @goto LOOP
            end
        end
        @goto LOOP
    end
end

signGA(x) = x<=0 ? -one(x) : one(x)

function roundFinite(x,Umax)
    y=round(x)
    return  abs(y)>Umax ?  Umax*signGA(y) : y
end

"""
    b = svp(B)

Find the shortest basis vector `b` for the lattice formed by the matrix
`B`. This solves the 'shortest vector problem' (SVP).

We use the [`cvp`](@ref) function in the library, which is not necessarily
the fastest SVP solver.  Roughly follows the CVP-to-SVP reduction in
http://web.eecs.umich.edu/~cpeikert/lic15/lec06.pdf

# Examples
```jldoctest
julia> H=[1 2; 3 4]; svp(H)
2-element Array{Int64,1}:
 -1
  1

julia> H= BigFloat.([2.5 2; 3 4]); svp(H)
2-element Array{BigFloat,1}:
  0.50
 -1.0

```
"""
function svp(B::AbstractArray{Td,2}) where Td
    m,n= size(B)
    V = zeros(Td,m,n)
    c = zeros(Td,n)

    for i = 1:n
        Bi = copy(B)
        Bi[:,i] *=2
        Q,R = qr(Bi)
        vc = cvp(Q'*B[:,i],R)
        V[:,i] = Bi*vc-B[:,i]
        c[i] = V[:,i]'*V[:,i]
    end
    idx = sortperm(c)
    return V[:,idx[1]]
end


using InfiniteLinearAlgebra, BandedMatrices, LazyArrays, Test
import InfiniteLinearAlgebra: UpperHessenbergQ, LowerHessenbergQ, tail_de, _BandedMatrix, QL, InfToeplitz

@testset "HessenbergQ" begin
    @testset "finite UpperHessenbergQ" begin
        c,s = cos(0.1), sin(0.1); q = [c s; s -c];
        Q = UpperHessenbergQ(Fill(q,1))
        @test size(Q) == (size(Q,1),size(Q,2)) == (2,2)
        @test Q ≈ q
        Q = UpperHessenbergQ(Fill(q,2))
        @test bandwidths(Q) == (1,2)
        @test Q ≈ [q zeros(2); zeros(1,2) 1] * [1 zeros(1,2); zeros(2) q]
        @test Q' isa LowerHessenbergQ
        @test Q' ≈ [1 zeros(1,2); zeros(2) q'] * [q' zeros(2); zeros(1,2) 1]

        q = qr(randn(2,2) + im*randn(2,2)).Q
        Q = UpperHessenbergQ(Fill(q,1))
        @test Q ≈ q
        Q = UpperHessenbergQ(Fill(q,2))
        @test bandwidths(Q) == (1,2)
        @test Q ≈ [q zeros(2); zeros(1,2) 1] * [1 zeros(1,2); zeros(2) q]
        @test Q' isa LowerHessenbergQ
        @test Q' ≈ [1 zeros(1,2); zeros(2) q'] * [q' zeros(2); zeros(1,2) 1]
    end

    @testset "finite LowerHessenbergQ" begin
        c,s = cos(0.1), sin(0.1); q = [c s; s -c];
        Q = LowerHessenbergQ(Fill(q,1))
        @test size(Q) == (size(Q,1),size(Q,2)) == (2,2)
        @test Q ≈ q
        Q = LowerHessenbergQ(Fill(q,2))
        @test bandwidths(Q) == (2,1)
        @test Q ≈ [1 zeros(1,2); zeros(2) q] * [q zeros(2); zeros(1,2) 1]
        @test Q' isa UpperHessenbergQ
        @test Q' ≈ [q' zeros(2); zeros(1,2) 1]  * [1 zeros(1,2); zeros(2) q']

        q = qr(randn(2,2) + im*randn(2,2)).Q
        Q = LowerHessenbergQ(Fill(q,1))
        @test Q ≈ q
        Q = LowerHessenbergQ(Fill(q,2))
        @test bandwidths(Q) == (2,1)
        @test Q ≈ [1 zeros(1,2); zeros(2) q] * [q zeros(2); zeros(1,2) 1]
        @test Q' isa UpperHessenbergQ
        @test Q' ≈ [q' zeros(2); zeros(1,2) 1]  * [1 zeros(1,2); zeros(2) q']
    end

    @testset "infinite-Q" begin
        c,s = cos(0.1), sin(0.1); q = [c s; s -c];
        Q = LowerHessenbergQ(Fill(q,∞))
        @test Q[1,1] ≈ 0.9950041652780258
        Q = UpperHessenbergQ(Fill(q,∞))
        @test Q[1,1] ≈ 0.9950041652780258
    end

    @testset "QRPackedQ,QLPackedQ -> LowerHessenbergQ,UpperHessenbergQ" begin
        for T in (Float64, ComplexF64)
            A = brand(T,10,10,1,1)
            Q,R = qr(A)
            @test UpperHessenbergQ(Q) ≈ Q
            Q,L = ql(A)
            @test LowerHessenbergQ(Q) ≈ Q
        end
    end
end


n = 7  # rand(1:10)
m = 17 # rand(1:20)
A = matrixdepot("pei", n, m)
B = ones(n,n)
@test issymmetric(A)
@test matrixdepot("pei", n) - triu(B) - tril(B) ≈ zeros(n,n)
println("'pei' passed test...")


"""
    x = spqr_rank_deflation(...)

SPQR_RANK_DEFLATION constructs pseudoinverse or basic solution using deflation.

 Called from spqr_basic and spqr_cod after these routines call spqr,
 spqr_rank_ssi and spqr_rank_form_basis.  The input parameters are as used
 in spqr_basic and spqr_cod.
 The parameter call_from indicates the type of call:
      call_from = 1 indicates a call from spqr_basic
      call_from = 2 indicates a call from spqr_cod
 Output:
   x -- for call_from = 1 x is a basic solution to
                 min || b - A * x ||.               (1)
        The basic solution has n - (rank returned by spqr) zero components.
        For call_from = 2 x is an approximate pseudoinverse solution to (1).
 Not user-callable.

 Algorithm:   R * wh = ch or R' * wh = ch is solved where ch is described
              in the code and R comes from the QR factorizations in
              spqr_basic or spqr_cod. R is triangular and potentially
              numerically singular with left and right singular vectors for
              small singular values stored in U and V.  When R is numerically
              singular deflation (see SIAM SISC, 11:519-530, 1990) to
              calculate an approximate truncated singular value solution to
              the triangular system.  Orthogonal transformations
              are applied to wh to obtain the solutions x to (1).

"""
function spqr_rank_deflation(call_from, R::AbstractMatrix{T}, U, V, C, m, n, rank_spqr,
                    numerical_rank, prow, pcol, Q::AbstractMatrix{T}; opts...) where T

# disable nearly-singular matrix warnings, and save the current state
## TODO: how? warning_state = warning ('off', 'MATLAB:nearlySingularMatrix') ;

    opts, nsvals_large, start_with_A_transpose =
        get_opts(opts, :nsvals_large, :start_with_A_transpose)

    if isempty(C)
        return zeros(T, m)
    end
    # restrict to columns of U, V belonging to small singular Values of R (< tol)
    if numerical_rank < rank_spqr
        k = size(U, 2)
        U = view(U, 1:m, nsvals_large+1:k)
        V = view(V, 1:m, nsvals_large+1:k)
    end

   if start_with_A_transpose || call_from == 1

       ch = view(C, 1:rank_spqr, 1:size(C,2))
       if numerical_rank == rank_spqr
           wh = R \ ch
       else
           # use deflation (see SIAM SISC, 11:519-530, 1990) to calculate an
           # approximate truncated singular value solution to R * wh = ch
           wh = ch - U * (U' * ch)
           wh = R \ wh
           wh = wh - V * (V' * wh)
        end
        if call_from == 2
            wh[prow,:] = wh;
        end
        wh = [wh ; zeros(T, n - rank_spqr,size(C,2)) ];
        if call_from == 1
            x = wh
            x[pcol,:] = x
        else
            x = spqr_qmult(Q, wh, 1)
        end

    else

        ch = view(C, prow, 1:size(C,2))
        if numerical_rank == rank_spqr
            wh = R' \ ch
        else
            # use deflation (see SIAM SISC, 11:519-530, 1990) to calculate an
            # approximate truncated singular value solution to R' * wh = ch
            wh = ch - V * (V' * ch)
            wh = R' \ wh
            wh = wh - U * (U' * wh)
        end
        wh = [wh ; zeros(T, n - rank_spqr, size(C,2)) ];
        wh = spqr_qmult(Q, wh, 1)
        x[pcol,:] = wh

    end

    # restore the warning back to what it was
    ## TODO: warning (warning_state) ;

    x
end


# Create a default kernel from two parameters k[1] and k[2]
make_kernel(k) = softplus(k[1]) * (SqExponentialKernel() ∘ ScaleTransform(softplus(k[2])))

# Computes the optimal closed form solution for the variational posterior
# q(u) (e.g. # https://krasserm.github.io/2020/12/12/gaussian-processes-sparse/
# equations (11) & (12)). Assumes a ZeroMean function.
function optimal_variational_posterior(fu, fx, y)
    fu.f.mean isa AbstractGPs.ZeroMean ||
        error("The exact posterior requires a GP with ZeroMean.")
    σ² = fx.Σy[1]
    Kuf = cov(fu, fx)
    Kuu = Symmetric(cov(fu))
    Σ = (Symmetric(cov(fu) + (1 / σ²) * Kuf * Kuf'))
    m = ((1 / σ²) * Kuu * (Σ \ Kuf)) * y
    S = Symmetric(Kuu * (Σ \ Kuu))
    return MvNormal(m, S)
end


using JuMP
using Cbc
model = Model(Cbc.Optimizer)
# set_optimizer_attribute(model, "cuts","off")

month_num = 8
month = 1:month_num

# paramerters
q = [5000] # product fix cost
p = [100] # cost per product
d = [400, 400, 800, 800, 1200, 1200, 1200, 1200] # demand of product by month
s_init = [200] # initial stock
h = [5] # stock cost per product month

@variable(model, x[j in month] >= 0)
@variable(model, s[j in month] >= 0)
@variable(model, y[j in month], Bin)

# 库存平衡
@constraint(model ,dem_sat1, s_init[1] + x[1] == d[1] + s[1])
@constraint(model, dem_sat[j in month[2:end]], s[j - 1] + x[j] == d[j] + s[j] )

# 最大产量约束
@constraint(model, vub[j in month], x[j] <= sum(d[k]  for k in j:month_num) * y[j])

# 目标函数
@expression(model ,cost, sum(p[1] * x[j] + q[1] * y[j] for j in month) + sum(h[1] * s[j] for j in month[1:end - 1]))
@objective(model, Min, cost)
optimize!(model)

# 显示结果
cost_val = value.(cost)
x_val = value.(x)
s_val = value.(s)
@show cost_val;
@show x_val;
@show s_val;


"""
    loglike_beta(nij, A, cij)

Profile likelihood for the frequency response vectors.

*Reference*: Based on Equation (9) of Stompor et al., MNRAS, 392, 216 (2009)

# Arguments
- `nij::Array{<:AbstractFloat,2}`: symmetric noise covariance matrix with the dimention of `(nν, nν)` where `nν` is the number of frequency bands.
- `A::Array{<:AbstractFloat,2}`: `nν`-by-`nc` matrix of the frequency response, for `nc` components in sky.
    - E.g., ``A = [a B]`` where `a = ones(nν)` for CMB and `B` is a `nν`-by-`nc-1` matrix for the frequency response of foreground components.
- `cij::Array{<:AbstractFloat,2}`: symmetric covariance matrix with the dimention of `(nν, nν)`.
"""
function loglike_beta(
    nij::Array{T,2},
    A::Array{T,2},
    cij::Array{T,2},
) where {T<:AbstractFloat}
    if size(nij)[1] ≠ size(nij)[2] || size(cij)[1] ≠ size(cij)[2]
        throw(DimensionMismatch("covariance matrix must be a square matrix"))
    elseif size(nij)[1] ≠ size(A)[1] || size(nij)[1] ≠ size(cij)[1]
        throw(DimensionMismatch("dimensions of input matrices do not match"))
    else
        N = Symmetric(nij)
        M = Symmetric(cij)
    end
    Ninv = inv(N)
    lnlike = 0.5 * tr(Ninv * A * inv(A' * Ninv * A) * A' * Ninv * M)
end

"""
    loglike_beta(nij, A, d)

Profile likelihood for the frequency response vectors.

*Reference*: Equation (9) of Stompor et al., MNRAS, 392, 216 (2009)

# Arguments
- `nij::Array{<:AbstractFloat,2}`: symmetric noise covariance matrix with the dimention of `(nν, nν)` where `nν` is the number of frequency bands.
- `A::Array{<:AbstractFloat,2}`: `nν`-by-`nc` matrix of the frequency response, for `nc` components in sky.
    - E.g., ``A = [a B]`` where `a = ones(nν)` for CMB and `B` is a `nν`-by-`nc-1` matrix for the frequency response of foreground components.
- `d::Array{<:AbstractFloat,1}`: data vector for a given pixel (or any other appropriate domain). The number of elements is `nν`.
"""
function loglike_beta(
    nij::Array{T,2},
    A::Array{T,2},
    d::Array{T,1},
) where {T<:AbstractFloat}
    if size(nij)[1] ≠ size(nij)[2]
        throw(DimensionMismatch("covariance matrix must be a square matrix"))
    elseif size(nij)[1] ≠ size(A)[1] || size(nij)[1] ≠ length(d)
        throw(DimensionMismatch("dimensions of matrices/vector do not match"))
    else
        M = Symmetric(nij)
    end
    x = M \ d
    s = inv(A' * inv(M) * A) * A'x
    lnlike = 0.5 * x'A * s
end

"""
    loglike_beta_grad(nij, A, dAdβ, d)

Derivative of the profile likelihood for the frequency response vectors with respect to a foreground parameter

*Reference*: Equation (A1) of Errard et al., PRD, 84, 063005 (2011)

# Arguments
- `nij::Array{<:AbstractFloat,2}`: symmetric noise covariance matrix with the dimention of `(nν, nν)` where `nν` is the number of frequency bands.
- `A::Array{<:AbstractFloat,2}`: `nν`-by-`nc` matrix of the frequency response, for `nc` components in sky.
    - E.g., ``A = [a B]`` where `a = ones(nν)` for CMB and `B` is a `nν`-by-`nc-1` matrix for the frequency response of foreground components.
- `dAdβ::Array{<:AbstractFloat,2}`: `nν`-by-`nc` matrix of the derivative of the frequency response with respect to a foreground parameter.
    - E.g., ``dAdβ = [zeros(nν) dsynch/dβs zeros(nν)]`` where `zeros(nν)` for CMB and dust because they do not depend on the synchrotron index `βs`.
- `d::Array{<:AbstractFloat,1}`: data vector for a given pixel (or any other appropriate domain). The number of elements is `nν`.
"""
function loglike_beta_grad(
    nij::Array{T,2},
    A::Array{T,2},
    dAdβ::Array{T,2},
    d::Array{T,1},
) where {T<:AbstractFloat}
    if size(nij)[1] ≠ size(nij)[2]
        throw(DimensionMismatch("covariance matrix must be a square matrix"))
    elseif size(nij)[1] ≠ size(A)[1] ||
           size(nij)[1] ≠ length(d) ||
           size(A)[1] ≠ size(dAdβ)[1] ||
           size(A)[2] ≠ size(dAdβ)[2]
        throw(DimensionMismatch("dimensions of matrices/vector do not match"))
    else
        M = Symmetric(nij)
    end
    x = M \ d
    s = inv(A' * inv(M) * A) * A'x
    lnlike_grad = s' * dAdβ' * (x - M \ (A * s))
end

"""
    loglike_beta_hess(nij, A, dAdβI, dAdβJ, d)

Hessian of the profile likelihood for the frequency response vectors with respect to foreground parameters

*Reference*: Equation (5) of Errard et al., PRD, 84, 063005 (2011)

# Arguments
- `nij::Array{<:AbstractFloat,2}`: symmetric noise covariance matrix with the dimention of `(nν, nν)` where `nν` is the number of frequency bands.
- `A::Array{<:AbstractFloat,2}`: `nν`-by-`nc` matrix of the frequency response, for `nc` components in sky.
    - E.g., ``A = [a B]`` where `a = ones(nν)` for CMB and `B` is a `nν`-by-`nc-1` matrix for the frequency response of foreground components.
- `dAdβI::Array{<:AbstractFloat,2}` and `dAdβJ::Array{<:AbstractFloat,2}`: `nν`-by-`nc` matrix of the derivative of the frequency response with respect to a foreground parameter.
    - E.g., ``dAdβI = [zeros(nν) dsynch/dβs zeros(nν)]`` where `zeros(nν)` for CMB and dust because they do not depend on the synchrotron index `βs`.
- `d::Array{<:AbstractFloat,1}`: data vector for a given pixel (or any other appropriate domain). The number of elements is `nν`.
"""
function loglike_beta_hess(
    nij::Array{T,2},
    A::Array{T,2},
    dAdβI::Array{T,2},
    dAdβJ::Array{T,2},
    d::Array{T,1},
) where {T<:AbstractFloat}
    if size(nij)[1] ≠ size(nij)[2]
        throw(DimensionMismatch("covariance matrix must be a square matrix"))
    elseif size(nij)[1] ≠ size(A)[1] ||
           size(nij)[1] ≠ length(d) ||
           size(A)[1] ≠ size(dAdβI)[1] ||
           size(A)[2] ≠ size(dAdβI)[2] ||
           size(A)[1] ≠ size(dAdβJ)[1] ||
           size(A)[2] ≠ size(dAdβJ)[2]
        throw(DimensionMismatch("dimensions of matrices/vector do not match"))
    else
        M = Symmetric(nij)
    end
    x = M \ d
    AtMinvAinv = inv(A' * inv(M) * A)
    s = AtMinvAinv * A'x
    u = M \ (dAdβJ * s)
    v = M \ (A * AtMinvAinv * A'u)
    lnlike_hess = tr(dAdβI' * (v - u) * s')
end


using Distributions
import Distributions: cdf, pdf

@doc raw"""
Weibull count model.

**Weibull count model** is a discrete distribution obtained as the *count model with inter-arrival times assumed to be i.i.d. Weibull distributed* [1, Section 3.1].

In words, a Weibull count model models the *number of events* which occur within some time `t`, where the time *between* each of the events are (i.i.d.) Weibull distributed.

The PDF is given by the following formula
```math
P \big( I(t) = i \big) = \sum_{j = 1}^{\infty} \frac{(-1)^{j + i} (\lambda t^c)^j a_j^i}{\Gamma(cj + 1)}, \quad i = 0, 1, 2, \dots
```
where
```math
\alpha_j^0 = \frac{\Gamma(cj + 1)}{\Gamma(j + 1)}, \quad j = 0, 1, 2, \dots
```
and
```math
\alpha_j^{i + 1} = \sum_{m = i}^{j - 1} \alpha_m^i \frac{\Gamma(cj - cm + 1)}{\Gamma(j - m + 1)}, \quad i = 0, 1, 2, \dots, \quad j = i + 1, i + 2, \dots
```
The derivation can be found in [1, Section 3.1].

## Notes
The computation of the PDF, CDF, etc. for a Weibull counting model requires approximation of an infinite series, which is highly unstable. We therefore make use of `BigInt` and `BigFloat` for all computations under the hood.
Most of the methods related to `WeibullCountModel` are dependent on the series approximation, and will therefore throw a `WeibullCountModel.DivergedError` upon failing to converge.

## References
[1] Adrian, M., Bradlow, E., Fader, P., & McShane, B., Count models based on weibull interarrival times, CoRR, (),  (2013).

"""
struct WeibullCountModel <: Distribution{Univariate, Discrete}
    λ::BigFloat
    c::BigFloat
    t::BigFloat
end

WeibullCountModel(λ::Number, c::Number, t::Number) = WeibullCountModel(BigFloat(λ), BigFloat(c), BigFloat(1.0))
WeibullCountModel(λ::Number, c::Number) = WeibullCountModel(λ, c, 1.0)

# PDF
pdf(d::WeibullCountModel, x::Array{BigInt}) = begin
    # extract the actual result of the computation
    weibull_count_pdf.(x, d.λ, d.c, d.t)[1]
end
pdf(d::WeibullCountModel, x::AbstractArray) = pdf.(d, BigInt.(x))

pdf(d::WeibullCountModel, x::BigInt) = weibull_count_pdf(x, d.λ, d.c)

# make explicit to prevent ambiguity from `Distributions.jl`
for T ∈ [Int64, Integer, Real, Number]
    eval(quote pdf(d::WeibullCountModel, x::$T) = pdf(d, BigInt(x)) end)
end

# CDF
cdf(d::WeibullCountModel, x::Array{BigInt}) = begin
    # extract the actual result of the computation
    weibull_count_cdf.(x, d.λ, d.c)[1]
end
cdf(d::WeibullCountModel, x::AbstractArray) = cdf(d, BigInt.(x))

cdf(d::WeibullCountModel, x::BigInt) = weibull_count_cdf(x, d.λ, d.c)
# make explicit to prevent ambiguity from `Distributions.jl`
for T ∈ [Int64, Integer, Real, Number]
    eval(quote cdf(d::WeibullCountModel, x::$T) = cdf(d, BigInt(x)) end)
end


using Distributions
#using Distributions: check_args

export NegativeBinomialConvolution, pmf, cdf

immutable NegativeBinomialConvolution <: DiscreteUnivariateDistribution
    rs::Array{Float64,1}
    ps::Array{Float64,1}
    logrs::Array{Float64,1}
    logps::Array{Float64,1}
    logqs::Array{Float64,1}
    mean::Float64
    stddev::Float64

    """
    This is based on the paper:
    "On the sums of compound negative binomial and gamma random variables" by Vallaisamy, 2009
    http://projecteuclid.org/euclid.jap/1238592129
    """
    function NegativeBinomialConvolution(rs::Array{Float64,1}, ps::Array{Float64,1})
        # @check_args(NegativeBinomialConvolution, minimum(rs) > zero(rs[1]))
        # @Distributions.check_args(NegativeBinomialConvolution, minimum(ps) > zero(ps[1]))
        # @Distributions.check_args(NegativeBinomialConvolution, Base.maximum(ps) < one(ps[1]))

        # just add the means and variances to get the convolution mean and variance
        meanVal = sum([(ps[i]*rs[i])/(1-ps[i]) for i in 1:length(ps)])
        varVal = sum([(ps[i]*rs[i])/(1-ps[i])^2 for i in 1:length(ps)])

        new(rs, ps, log(rs), log(ps), log(1 - ps), meanVal, sqrt(varVal))
    end
    NegativeBinomialConvolution(rs::Array{Real,1}, ps::Array{Real,1}) = NegativeBinomialConvolution(convert(Array{Float64,1}, rs), convert(Array{Float64,1}, ps))
    NegativeBinomialConvolution(rs::Array{Int64,1}, ps::Array{Float64,1}) = NegativeBinomialConvolution(convert(Array{Float64,1}, rs), convert(Array{Float64,1}, ps))
end

@Distributions.distr_support NegativeBinomialConvolution 0 Inf

params(d::NegativeBinomialConvolution) = (d.rs, d.ps)

"Call f on all possible ways n non-negative integers can sum to equal s."
function allsums_apply(f::Function, n::Int, s::Int)
    if n == 1
        f([s])
    elseif n == 2
        for i in 0:s
            f(Int64[i,s-i])
        end
    elseif n == 3
        for i in 0:s
            for j in 0:s
                if i + j > s break end
                f(Int64[i,j,s-i-j])
            end
        end
    elseif n == 4
        for i in 0:s
            for j in 0:s
                if i + j > s break end
                for k in 0:s
                    if i + j + k > s break end
                    f(Int64[i,j,k,s-i-j-k])
                end
            end
        end
    elseif n == 5
        for i in 0:s
            for j in 0:s
                if i + j > s break end
                for k in 0:s
                    if i + j + k > s break end
                    for l in 0:s
                        if i + j + k + l > s break end
                        f(Int64[i,j,k,l,s-i-j-k-l])
                    end
                end
            end
        end
    else
        error("n = $n is not yet implemented.")
    end
end

"Compute the exact PMF at the given point with a sum."
function Distributions.pdf(d::NegativeBinomialConvolution, s::Int)
    exp(Distributions.logpdf(d::NegativeBinomialConvolution, s::Int))
end

function Distributions.logpdf(d::NegativeBinomialConvolution, s::Int)
    total = 0.0

    vals = Float64[]
    function inner_term(m)
        val = 0.0
        for j in 1:length(d.ps)
            binVal = lgamma(d.rs[j] + m[j]) - lgamma(m[j]+1) - lgamma(d.rs[j])
            val += binVal + d.rs[j]*d.logps[j] + m[j]*d.logqs[j]
        end
        push!(vals, val)
    end
    allsums_apply(inner_term, length(d.ps), s)

    maxVal = Base.maximum(vals)
    maxVal + log(sum(map(x->exp(x-maxVal), vals)))
end

function Distributions.cdf(d::NegativeBinomialConvolution, s::Int)
    total = 0.0
    for i in 0:s
        total += pdf(d, i)
    end
    total
end


module KroneckerTools

using LinearAlgebra
using LinearAlgebra.BLAS
using QuasiTriangular
import QuasiTriangular: A_mul_B!, At_mul_B!, A_mul_Bt!
import Base.convert

include("ExtendedMul.jl")
export a_mul_kron_b!, a_mul_b_kron_c!, kron_at_kron_b_mul_c!, a_mul_b_kron_c_d!, at_mul_b_kron_c!, a_mul_b_kron_ct!

"""
Content:
c = a*(b ⊗ b ⊗ ... ⊗ b)
a_mul_kron_b!(c::AbstractMatrix, a::AbstractMatrix, b::AbstractMatrix, order::Int64, work1::AbstractVector, work2::AbstractVector)
a_mul_kron_b!(c::AbstractMatrix, a::AbstractMatrix, b::AbstractMatrix, order::Int64)
a_mul_kron_b!(c::AbstractMatrix, a::AbstractMatrix, b::AbstractVector,work::AbstractVector)
d = a*b*(c ⊗ c ⊗ ... ⊗ c)
a_mul_b_kron_c!(d::AbstractMatrix, a::AbstractMatrix, b::AbstractMatrix, c::AbstractMatrix, order::Int64,work1::AbstractVector,work2::AbstractVector)
d = a^T*b*(c ⊗ c ⊗ ... ⊗ c)
at_mul_b_kron_c!(d::AbstractMatrix, a::AbstractMatrix, b::AbstractMatrix, c::AbstractMatrix, order::Int64, work1::AbstractVector, work2::AbstractVector)
d = a*b*(c^T ⊗ c^T ⊗ ... ⊗ c^T)
a_mul_b_kron_ct!(d::AbstractMatrix, a::AbstractMatrix, b::AbstractMatrix, c::AbstractMatrix, order::Int64, work1::AbstractVector, work2::AbstractVector)
e = a*b*(c ⊗ d ⊗ ... ⊗ d)
a_mul_b_kron_c_d!(e::AbstractMatrix, a::AbstractMatrix, b::AbstractMatrix, c::AbstractMatrix, d::AbstractMatrix, order::Int64, work1::AbstractVector, work2::AbstractVector)
d = (a^T ⊗ a^T ⊗ ... ⊗ a^T ⊗ b)*c
kron_at_kron_b_mul_c!(d::AbstractVector, offset_d::Int64, a::AbstractMatrix, order::Int64, b::AbstractMatrix, c::AbstractVector, offset_c::Int64, work1::AbstractVector, work2::AbstractVector, offset_w::Int64)
kron_at_kron_b_mul_c!(d::AbstractVector, a::AbstractMatrix, order::Int64, b::AbstractMatrix, c::AbstractVector, work1::AbstractVector, work2::AbstractVector)
kron_at_kron_b_mul_c!(a::AbstractMatrix, order::Int64, b::AbstractMatrix, c::AbstractVector, offset_c::Int64, work::AbstractVector)
kron_at_kron_b_mul_c!(a::AbstractMatrix, order::Int64, b::AbstractMatrix, c::AbstractVector, work::AbstractVector)
c = (a ⊗ a ⊗ ... ⊗ a)*b
kron_a_mul_b!(c::AbstractVector, a::AbstractMatrix, order::Int64, b::AbstractVector, q::Int64, work1::AbstractVector, work2::AbstractVector)
c = (I_p ⊗ a ⊗ I_q)*b
kron_mul_elem!(c::AbstractVector, offset_c::Int64, a::AbstractMatrix, b::AbstractVector, offset_b::Int64, p::Int64, q::Int64)
kron_mul_elem!(c::AbstractVector, a::AbstractMatrix, b::AbstractVector, p::Int64, q::Int64)
c = (I_p ⊗ a^T ⊗ I_q)*b
kron_mul_elem_t!(c::AbstractVector, offset_c::Int64, a::AbstractMatrix, b::AbstractVector, offset_b::Int64, p::Int64, q::Int64)
kron_mul_elem_t!(c::AbstractVector, a::AbstractMatrix, b::AbstractVector, p::Int64, q::Int64)
"""

"""
    kron_mul_elem!(p::Int64, q::Int64, m::Int64, a::AbstractMatrix, b::AbstractVector, c::AbstractVector)

Performs (I_p ⊗ a ⊗ I_q) b, where m,n = size(a). The result is stored in c.
"""
function kron_mul_elem!(c::AbstractVector, offset_c::Int64, a::AbstractMatrix, b::AbstractVector, offset_b::Int64, p::Int64, q::Int64)
    m, n = size(a)
    length(b) >= n*p*q || throw(DimensionMismatch("The dimension of vector b, $(length(b)) doesn't correspond to order, ($p, $q)  and the dimensions of the matrix, $(size(a))"))
    length(c) >= m*p*q || throw(DimensionMismatch("The dimension of the vector c, $(length(c)) doesn't correspond to order, ($p, $q)  and the dimensions of the matrix, $(size(a))"))

    begin
        if p == 1 && q == 1
            # a*b
            A_mul_B!(c, offset_c, a, b, offset_b, 1)
        elseif q == 1
            #  (I_p ⊗ a)*b = vec(a*[b_1 b_2 ... b_p])
            A_mul_B!(c, offset_c, a, b, offset_b, p)
        elseif p == 1
            # (a ⊗ I_q)*b = (b'*(a' ⊗ I_q))' = vec(reshape(b,q,m)*a')
            A_mul_Bt!(c, offset_c, b, offset_b, q, n, a)
        else
            # (I_p ⊗ a ⊗ I_q)*b = vec([(a ⊗ I_q)*b_1 (a ⊗ I_q)*b_2 ... (a ⊗ I_q)*b_p])
            mq = m*q
            nq = n*q
            for i=1:p
                A_mul_Bt!(c, offset_c, b, offset_b, q, n, a)
                offset_b += nq
                offset_c += mq
            end
        end
    end
end

function kron_mul_elem!(c::AbstractVector, a::AbstractMatrix, b::AbstractVector, p::Int64, q::Int64)
    kron_mul_elem!(c, 1, a, b, 1, p, q)
end

"""
    kron_mul_elem_t!(p::Int64, q::Int64, m::Int64, a::AbstractMatrix, b::AbstractVector, c::AbstractVector)

Performs (I_p ⊗ a' ⊗ I_q) b, where m,n = size(a). The result is stored in c.
"""
function kron_mul_elem_t!(c::AbstractVector, offset_c::Int64, a::AbstractMatrix, b::AbstractVector, offset_b::Int64, p::Int64, q::Int64)
    m, n = size(a)
    length(b) >= m*p*q || throw(DimensionMismatch("The dimension of vector b, $(length(b)) doesn't correspond to order, ($p, $q)  and the dimensions of the matrix, $(size(a))"))
    length(c) >= n*p*q || throw(DimensionMismatch("The dimension of the vector c, $(length(c)) doesn't correspond to order, ($p, $q)  and the dimensions of the matrix, $(size(a))"))

    begin
        if p == 1 && q == 1
            # a'*b
            At_mul_B!(c, offset_c, a, b, offset_b, 1)
        elseif q == 1
            #  (I_p ⊗ a')*b = vec(a'*[b_1 b_2 ... b_p])
            At_mul_B!(c, offset_c, a, b, offset_b, p)
        elseif p == 1
            # (a' ⊗ I_q)*b = (b'*(a ⊗ I_q))' = vec(reshape(b,q,m)*a)
            A_mul_B!(c, offset_c, b, offset_b, q, m, a)
        else
            # (I_p ⊗ a' ⊗ I_q)*b = vec([(a' ⊗ I_q)*b_1 (a' ⊗ I_q)*b_2 ... (a' ⊗ I_q)*b_p])
            mq = m*q
            nq = n*q
            for i=1:p
                A_mul_B!(c, offset_c, b, offset_b, q, m, a)
                offset_b += mq
                offset_c += nq
            end
        end
    end
end

function kron_mul_elem_t!(c::AbstractVector, a::AbstractMatrix, b::AbstractVector, p::Int64, q::Int64)
    kron_mul_elem_t!(c, 1, a, b, 1, p, q)
end

"""
    a_mul_kron_b!(c::AbstractVector, a::AbstractVecOrMat, b::AbstractMatrix, order::Int64)

Performs a*(b ⊗ b ⊗ ... ⊗ b). The solution is returned in matrix c. order indicates the number of occurences of b

We use vec(a*(b ⊗ b ⊗ ... ⊗ b)) = (b' ⊗ b' ⊗ ... ⊗ b' ⊗ I)vec(a)

"""
function a_mul_kron_b!(c::AbstractMatrix, a::AbstractMatrix, b::AbstractMatrix, order::Int64, work1::AbstractVector, work2::AbstractVector)
    ma, na = size(a)
    mb, nb = size(b)
    mc, nc = size(c)
    mborder = mb^order
    nborder = nb^order
    na == mborder || throw(DimensionMismatch("The number of columns of a, $na, doesn't match the number of rows of b, $mb, times order = $order"))
    mc == ma || throw(DimensionMismatch("The number of rows of c, $mc, doesn't match the number of rows of a, $ma"))
    nc == nborder || throw(DimensionMismatch("The number of columns of c, $nc, doesn't match the number of columns of b, $nb, times order = $order"))

    # copy input
    copyto!(work1,a)
    n1 = ma*na
    n2 = Int(n1*nb/mb)
    v1 = view(work1,1:n1)
    v2 = view(work2,1:n2)
    for q=0:order-1
        kron_mul_elem_t!(v2,b,v1,mb^(order-q-1),nb^q*ma)
        if q < order - 1
            # update dimension
            n2 = Int(n2*nb/mb)
            # swap and resize work vectors
            v1parent = v1.parent
            v1 = v2
            v2 = view(v1parent,1:n2)
        end
    end
    copy!(c,v2)
end

function a_mul_kron_b!(c::AbstractMatrix, a::AbstractMatrix, b::AbstractMatrix, order::Int64)
    ma, na = size(a)
    mb, nb = size(b)
    mc, nc = size(c)
    mborder = mb^order
    nborder = nb^order
    na == mborder || throw(DimensionMismatch("The number of columns of a, $na, doesn't match the number of rows of b, $mb, times order = $order"))
    mc == ma || throw(DimensionMismatch("The number of rows of c, $mc, doesn't match the number of rows of a, $ma"))
    nc == nborder || throw(DimensionMismatch("The number of columns of c, $nc, doesn't match the number of columns of b, $nb, times order = $order"))
    mb == nb || throw(DimensionMismatch("B must be a square matrix"))

    avec = vec(a)
    cvec = vec(c)
    for q=0:order-1
#        vavec = view(avec,1:ma*mb^(order-q)*nb^q)
#        vcvec = view(cvec,1:ma*mb^(order-q-1)*nb^(q+1))
        kron_mul_elem_t!(cvec,b,avec,mb^(order-q-1),nb^q*ma)
        if q < order - 1
            copy!(avec,cvec)
        end
    end
end

"""
    a_mul_b_kron_c!(d::AbstractVecOrMat, a::AbstractVecOrMat, b::AbstractMatrix, c::AbstractMatrix, order::Int64)

Performs a*B*(c ⊗ c ⊗ ... ⊗ c). The solution is returned in matrix or vector d. order indicates the number of occurences of c. c must be a square matrix

We use vec(a*b*(c ⊗ c ⊗ ... ⊗ c)) = (c' ⊗ c' ⊗ ... ⊗ c' ⊗ a)vec(b)

"""
function a_mul_b_kron_c!(d::AbstractMatrix, a::AbstractMatrix, b::AbstractMatrix, c::AbstractMatrix, order::Int64,work1::AbstractVector,work2::AbstractVector)
    ma, na = size(a)
    mb, nb = size(b)
    mc, nc = size(c)
    md, nd = size(d)
    na == mb || throw(DimensionMismatch("The number of columns of a, $(size(a,2)), doesn't match the number of rows of b, $(size(b,1))"))
    nb == mc^order || throw(DimensionMismatch("The number of columns of b, $(size(b,2)), doesn't match the number of rows of c, $(size(c,1)), times order, $order"))
    (ma == md && nc^order == nd) || throw(DimensionMismatch("Dimension mismatch for D: $(size(d)) while ($ma, $(nc^order)) was expected"))
    mul!(reshape(view(work1,1:ma*nb),ma,nb),a,b)
    for q=0:order-1
        kron_mul_elem_t!(work2,c,work1,mc^(order-q-1),nc^q*ma)
        if q < order - 1
            copy!(work1,work2)
        end
    end
    copyto!(d,work2)
end

"""
    function kron_at_kron_b_mul_c!(d::AbstractVector, offset_d::Int64, a::AbstractMatrix, order::Int64, b::AbstractMatrix, c::AbstractVector, offset_c::Int64, work1::AbstractVector, work2::AbstractVector, offset_w::Int64)
computes d = (a^T ⊗ a^T ⊗ ... ⊗ a^T ⊗ b)c using data from c at offset_c and work vectors work1 and work2 at offset_w
"""
function kron_at_kron_b_mul_c!(d::AbstractVector, offset_d::Int64, a::AbstractMatrix, order::Int64, b::AbstractMatrix, c::AbstractVector, offset_c::Int64, work1::AbstractVector, work2::AbstractVector, offset_w::Int64)
    mb,nb = size(b)
    if order == 0
        A_mul_B!(d,offset_d,b,1,mb,nb,c,offset_c,1)
    else
        ma, na = size(a)
        #    length(work) == naorder*mb  || throw(DimensionMismatch("The dimension of vector , $(length(c)) doesn't correspond to order, ($order)  and the dimension of the matrices a, $(size(a)), and b, $(size(b))"))
        p = ma^order
        kron_mul_elem!(work1, 1, b, c, offset_c, p, 1)
        p = Int(p/ma)
        q = mb
        for i = 1:order
            kron_mul_elem_t!(work2, offset_w, a, work1, 1, p, q)
            if i < order
                copyto!(work1, 1, work2, offset_w, p*na*q)
                p = Int(p/ma)
                q *= na
            end
        end
        copyto!(d, offset_d, work2, offset_w, p*na*q)
    end
end

"""
    function kron_at_kron_b_mul_c!(d::AbstractVector, a::AbstractMatrix, order::Int64, b::AbstractMatrix, c::AbstractVector, work1::AbstractVector, work2::AbstractVector)
computes d = (a^T ⊗ a^T ⊗ ... ⊗ a^T ⊗ b)c using work vectors work1 and work2
"""
function kron_at_kron_b_mul_c!(d::AbstractVector, a::AbstractMatrix, order::Int64, b::AbstractMatrix, c::AbstractVector, work1::AbstractVector, work2::AbstractVector)
    mb,nb = size(b)
    if order == 0
        A_mul_B!(d,1,b,1,mb,nb,c,1,1)
    else
        ma, na = size(a)
        #    length(work) == naorder*mb  || throw(DimensionMismatch("The dimension of vector , $(length(c)) doesn't correspond to order, ($order)  and the dimension of the matrices a, $(size(a)), and b, $(size(b))"))
        p = ma^order
        kron_mul_elem!(work1, b, c, p, 1)
        p = Int(p/ma)
        q = mb
        for i = 1:order
            kron_mul_elem_t!(work2, a, work1, p, q)
            if i < order
                copyto!(work1, 1, work2, 1, p*na*q)
                p = Int(p/ma)
                q *= na
            end
        end
        copyto!(d,1,work2,1,p*na*q)
    end
end

"""
    function kron_at_kron_b_mul_c!(a::AbstractMatrix, order::Int64, b::AbstractMatrix, c::AbstractVector, offset_c::Int64, work::AbstractVector)
updates c at offset_c with (a^T ⊗ a^T ⊗ ... ⊗ a^T ⊗ b)c using c and work as work vectors
"""
function kron_at_kron_b_mul_c!(a::AbstractMatrix, order::Int64, b::AbstractMatrix, c::AbstractVector, offset_c::Int64, work::AbstractVector)
    kron_at_kron_b_mul_c!(work, 1, a, order, b, c, offset_c, work, c, offset_c)
    copy!(c, work)
end

"""
    function kron_at_kron_b_mul_c!(a::AbstractMatrix, order::Int64, b::AbstractMatrix, c::AbstractVector, work::AbstractVector)
updates c  with (a^T ⊗ a^T ⊗ ... ⊗ a^T ⊗ b)c using c and work as work vectors
"""
function kron_at_kron_b_mul_c!(a::AbstractMatrix, order::Int64, b::AbstractMatrix, c::AbstractVector, work::AbstractVector)
    kron_at_kron_b_mul_c!(work, a, order, b, c, work, c)
    copy!(c, work)
end

function kron_at_mul_b!(c::AbstractVector, a::AbstractMatrix, order::Int64, b::AbstractVector, q::Int64, work1::AbstractVector, work2::AbstractVector)
    ma,na = size(a)
#    length(work) == naorder*mb  || throw(DimensionMismatch("The dimension of vector , $(length(c)) doesn't correspond to order, ($order)  and the dimension of the matrices a, $(size(a)), and b, $(size(b))"))
    p = ma^(order-1)
    s = p*q*na
    copyto!(work1,b)
    for i = 1:order
        kron_mul_elem_t!(work2, a, work1, p, q)
        if i < order
            copyto!(work1, 1, work2, 1, s)
            p = Int(p/ma)
            q *= na
            s = p*q*na
        end
    end
    copyto!(c,1,work2,1,s)
end

function kron_a_mul_b!(c::AbstractVector, a::AbstractMatrix, order::Int64, b::AbstractVector, q::Int64, work1::AbstractVector, work2::AbstractVector)
    ma,na = size(a)
#    length(work) == naorder*mb  || throw(DimensionMismatch("The dimension of vector , $(length(c)) doesn't correspond to order, ($order)  and the dimension of the matrices a, $(size(a)), and b, $(size(b))"))
    p = na^(order-1)
    s = p*q*ma
    copyto!(work1,b)
    for i = 1:order
        kron_mul_elem!(work2,a,work1,p,q)
        if i < order
            copyto!(work1, 1, work2, 1, s)
            p = Int(p/na)
            q *= ma
            s = p*q*ma
        end
    end
    copyto!(c,1,work2,1,s)
end

function at_mul_b_kron_c!(d::AbstractMatrix, a::AbstractMatrix, b::AbstractMatrix, c::AbstractMatrix, order::Int64, work1::AbstractVector, work2::AbstractVector)
    ma, na = size(a)
    mb, nb = size(b)
    mc, nc = size(c)
    if mc <= nc
        At_mul_B!(work1, 1, a, 1, na, ma, b, 1, nb)
        kron_at_mul_b!(vec(d), c, order, work1, na, vec(d), work2)
    else
        kron_at_mul_b!(work1, c, order, b, na, work1, work2)
        At_mul_B!(vec(d), 1, a, 1, na, ma, work1, 1, nc^order)
    end
end

function a_mul_b_kron_ct!(d::AbstractMatrix, a::AbstractMatrix, b::AbstractMatrix, c::AbstractMatrix, order::Int64, work1::AbstractVector, work2::AbstractVector)
    ma, na = size(a)
    mb, nb = size(b)
    mc, nc = size(c)
    if mc <= nc
        A_mul_B!(work1, 1, a, 1, ma, na, b, 1, nb)
        kron_a_mul_b!(vec(d), c, order, work1, ma, work1, work2)
    else
        kron_a_mul_b!(work1, b, order, c, mb, work1, work2)
        At_mul_B!(vec(d), 1, a, 1, na, ma, work1, 1, nc^order)
    end
end

function a_mul_kron_b!(c::AbstractMatrix, a::AbstractMatrix, b::AbstractVector,work::AbstractVector)
    ma, na = size(a)
    order = length(b)
    mc, nc = size(c)
    mborder = 1
    nborder = 1
    for i = 1:order
        mb, nb = size(b[i])
        mborder *= mb
        nborder *= nb
    end
    na == mborder || throw(DimensionMismatch("The number of columns of a, $na, doesn't match the number of rows of matrices in b, $mborder"))
    mc == ma || throw(DimensionMismatch("The number of rows of c, $mc, doesn't match the number of rows of a, $ma"))
    nc == nborder || throw(DimensionMismatch("The number of columns of c, $nc, doesn't match the number of columns of matrices in b, $nborder"))
    mborder <= nborder || throw(DimensionMismatch("the product of the number of rows of the b matrices needs to be smaller or equal to the product of the number of columns. Otherwise, you need to call a_mul_kron_b!(c::AbstractMatrix, a::AbstractMatrix, b::Vector{AbstractMatrix},work::AbstractVector)"))

    copyto!(work,a)
    mb, nb = size(b[1])
    vwork = view(work,1:ma*na)
    cvec = vec(c)
    mc = ma*na
    p = Int(mborder/size(b[order],1))
    q = ma
    for i = order:-1:1
        mb, nb = size(b[i])
        mc = Int(mc*nb/mb)
        vcvec = view(cvec,1:mc)
        kron_mul_elem_t!(cvec,b[i],work,p,q)
        if i > 1
            p = Int(p/mb)
            q = q*nb
            vwork = view(work,1:mc)
            copy!(vwork,vcvec)
        end
    end
end

"""
    a_mul_b_kron_c_d!(d::AbstractVecOrMat, a::AbstractVecOrMat, b::AbstractMatrix, c::AbstractMatrix, order::Int64)

Performs e = a*b*(c ⊗ d ⊗ ... ⊗ d). The solution is returned in matrix or vector e order indicates the number of occurences of d. c and d must be square matrices

We use vec(a*b*(c ⊗ d ⊗ ... ⊗ d)) = (c' ⊗ d' ⊗ ... ⊗ d' ⊗ a)vec(b)

"""
function a_mul_b_kron_c_d!(e::AbstractMatrix, a::AbstractMatrix, b::AbstractMatrix, c::AbstractMatrix, d::AbstractMatrix, order::Int64, work1::AbstractVector, work2::AbstractVector)

     ma, na = size(a)
     mb, nb = size(b)
     mc, nc = size(c)
     md, nd = size(d)
     me, ne = size(e)
     na == mb || throw(DimensionMismatch("The number of columns of a, $(size(a,2)), doesn't match the number of rows of b, $(size(b,1))"))
     nb == mc*md^(order-1) || throw(DimensionMismatch("The number of columns of b, $(size(b,2)), doesn't match the number of rows of c, $(size(c,1)), and d, $(size(d,1)) for order, $order"))
     (ma == me && nc*nd^(order-1) == ne) || throw(DimensionMismatch("Dimension mismatch for e: $(size(e)) while ($ma, $(nc*nd^(order-1))) was expected"))
     A_mul_B!(work1, 1, a, 1, ma, mb, vec(b), 1, nb)
     p = mc*md^(order - 2)
     q = ma
     for i = 0:order - 2
        kron_mul_elem_t!(work2, 1, d, work1, 1, p, q)
        copy!(work1,work2)
        p = Int(p/md)
        q *= nd
    end
    kron_mul_elem_t!(work2, c, work1, 1,q)
    copyto!(e, 1, work2, 1, ma*nc*nd^(order-1))
end

convert(::Type{Array{Float64, 2}}, x::Base.ReshapedArray{Float64,2,SubArray{Float64,1,Array{Float64,1},Tuple{UnitRange{Int64}},true},Tuple{}}) = unsafe_wrap(Array,pointer(x.parent.parent,x.parent.indexes[1][1]),x.dims)

end


"""
	@author: Alois Pichler
	Created September 2018
		provides ontoSimplex: function to project onto the simplex
"""

function ontoSimplex!(probabilities::Vector{Float64})   # projects the probabilities onto the simplex
    for i= 1:length(probabilities)
        if probabilities[i] < 0.0   # make them nonnegative
            probabilities[i]= 0.0
        end
    end
	summ= sum(probabilities)
	if summ <= 0.0
		probabilities= fill!(probabilities, 1.0/ length(probabilities))
		summ= sum(probabilities)
		@info "ontoSimplex: created probabilities." maxlog=1
	end
	if summ != 1
		probabilities./= summ
		tmpi= rand(1:length(probabilities)) 	# modify a random index
		summ= -sum(probabilities[1:tmpi-1])+ 1.0- sum(probabilities[tmpi+1:end])
		probabilities[tmpi]= max(0, summ)
		@info "ontoSimplex: modified probabilities." maxlog=1
	end
    nothing     # modifies probabilities
end


## sets module
## TODO: SymSets class so that we can better manage dispatch?

"""
`SymPy` has the [sets](http://docs.sympy.org/latest/modules/sets.html) module that provides:

* `FiniteSet` to specify a finite set of values
* `Interval` to specify an interval
* `Union` to specify a union of finite sets and/or intervals.

The `FiniteSet` and `Interval` constructors are provided, unions are formed from these.

The `S` module provides some built-in sets: `S.Reals`, `S.Integers`, `S.Naturals`, `S.Naturals0`, and `S.UniversalSet`, `S.EmptySet`.

Various methods for sets and intervals are provided, though not all are defined for unions:

* `powerset`
* `union`, `intersection`, `contains`, `boundary`
* `complement`, `is_disjoint`, `is_subset`, `is_superset`, `is_proper_subset`
* `sup`, `inf`, `measure`
* `is_Union`, `is_Interval`, `is_FiniteSet`
* `as_relational`, `left_open`, `right_open`, `is_left_unbounded`, `is_right_unbounded`

Examples:
```
a,b = Interval(0,1), Interval(1//2,2, false, true)  # [0,1], [1/2, 2)
union(a, b)     # [0,2), also  a ∪ b
intersect(a,b)  # [1/2,1],  also a ∩ b
contains(a, 3/2), contains(b, 3/2) # false, true
inf(a), sup(b)  # 1, 2
right_open(b)   # true
complement(b, S.Reals)  # (-∞, 1/2) ∪ [2, ∞)
complement(b, a)  # [0, 1/2)  (This is a \ b, but that notation is not provided)

x = symbols("x")
as_relational(a, x)  #  0 ≤ x ∧ x ≤ 1

s = FiniteSet(1,2)
powerset(s) # {∅, {1}, {2}, {1, 2}}
elements(s)   # collect(args(s))
```

SymPy `FiniteSet`s are not iterable, but can be collected with `collect(args(u))`.

SymPy `Union`s are not iterable, but can be collected after calling `args` via `collect(args(s))`.

"""
sypmy_sets = nothing

module S
using SymPy
function init_set()
    S = sympy[:S]
    global Reals = S[:Reals]
    global UniversalSet = S[:UniversalSet]
    global Naturals = S[:Naturals]
    global Naturals0 = S[:Naturals0]
    global Integers = S[:Integers]
    global EmptySet = S[:EmptySet]
end
end

export S, FiniteSet, Interval, ProductSet, ConditionSet
export powerset, imageset
export contains,  boundary, sup, measure
export is_disjoint
export is_proper_subset, is_subset, is_superset
export Interval
export as_relational
export is_left_unbounded, is_right_unbounded, left_open, right_open
export is_FiniteSet, is_Union, is_Interval

"Convert a SymPy finite set to a Set of Sym objects"
function Base.convert(::Type{Set}, s::Sym)
    is_FiniteSet(s) || throw(ArgumentError("`s` must be a finite set"))
    s1 = Set(Any[u for u in s.x])
    for el in copy(s1)
        if !isa(el, Set) && !(el.x[:is_Symbol]) && is_FiniteSet(s) && length(el.x) > 1
            ## replace python FiniteSet with julian Set
            setdiff!(s1, Set([el]))
            push!(s1, Set(convert(Set, el)))
        end
    end
    s1
end

## to convert a Union to an array of intervals
function _union_to_intervals(s::Sym)
    is_Union(s) || throw(ArgumentError("`s` must be a Union of Intervals"))
    collect(args(s))
end

"Find power set of set"
powerset(s::Sym) = s[:powerset]()

"Does I contain x?"
Base.contains(I::Sym, x) = (I[:contains](x) == Sym(true))
Base.in(x::Number, I::Sym) = contains(I, x)

"Elements of finite set"
elements(x::Sym) = PyCall.pyeval("[i for i in s]", s=project(x))
export elements

VERSION < v"0.5.0-" && eval(Expr(:import, :Base, :complement))
"Complement of set within the universe"
complement(I::Sym, U::Sym=S.Reals) = I[:complement](U)
export complement

"boundary, returnsa set"
boundary(I::Sym) = I.x[:boundary]

VERSION < v"0.5.0-" && eval(Expr(:import, :Base, :inf))
"Infinum of I"
inf(I::Sym) = I.x[:inf]
export inf

"Intersection of two intervals"
Base.intersect(I::Sym, J::Sym) = I[:intersect](J)

"Are `I` and `J` disjoint?"
is_disjoint(I::Sym, J::Sym) = I[:is_disjoint](J) == Sym(true)

"Is `J` a proper subset of `I`?"
is_proper_subset(I::Sym, J::Sym) = I[:is_proper_subset](J) == Sym(true)

"Is `J` a  subset of `I`?"
is_subset(I::Sym, J::Sym) = I[:is_subset](J) == Sym(true)

"Is `I` a  superset of `J`?"
is_superset(I::Sym, J::Sym) = I[:is_superset](J) == Sym(true)

"Lebesgue mesuare of an interval"
measure(I::Sym) = I.x[:measure]

"Supremum of a set"
sup(I::Sym) = I.x[:sup]

"Union of two intervals"
Base.union(I::Sym, J::Sym) = I[:union](J)

"Symmetric difference of two intervals, in one or other, but not both: `(I ∪ J) \ (I  ∩ J)"
Base.symdiff(I::Sym, J::Sym) = complement(intersection(I,J), union(I,J))

"rexpress I in terms of relations involving variable `x`"
as_relational(I::Sym, x::Sym) = I[:as_relational](x)

# not implemented: `end` => `sup`, `left` =>

"Looks like (-oo, a)?"
is_left_unbounded(I::Sym) = I.x[:is_left_unbounded] == Sym(true)

"Looks like (a, oo)?"
is_right_unbounded(I::Sym) = I.x[:is_right_unbounded] == Sym(true)

"Looks like (a, b...?"
left_open(I::Sym) = I.x[:left_open] == Sym(true)

"Looks like ..., b)?"
right_open(I::Sym) = I.x[:right_open] == Sym(true)

"Is `I` a finite set?"
is_FiniteSet(I::Sym) = I.x[:is_FiniteSet]

"Is `I` an interval?"
is_Interval(I::Sym) = I.x[:is_Interval]

"Is `I` a union?"
is_Union(I::Sym) = I.x[:is_Union]

function init_sets()
    S.init_set()

"""
    `FiniteSet(1,2,3)`, `FiniteSet(1:10...)`: create a finite set

See also `Interval` to create subsets of the real line.
"""
    "FiniteSet: http://docs.sympy.org/latest/modules/sets.html"
    global FiniteSet(args...) = sympy_meth(:FiniteSet, args...)

    "ProductSet: http://docs.sympy.org/latest/modules/sets.html"
    global ProductSet(args...) = sympy_meth(:ProductSet, args...)

    """
    Means to filter a set to pull out elements by a condition.

    ConditionSet:  A set of elements which satisfies a given condition.

    `ConditionSet(x, condition, S) = {x | condition(x) == true for x in S}`

    (Introduced in SymPy 1.0)
    """
    global ConditionSet(var::Sym, pred::Sym, S) = sympy_meth(:ConditionSet, var, pred, S)

    """
    Represent a ComplexRegion.
    working?
```
    I, J = Interval(0,1), Interval(0,2)
    R = ComplexRegion(I * J)
    1 in R        # true
    2 + 1im in R  # false
```
"""
    ComplexRegion(IJ::Sym; kwargs...) = sympy_meth(:ComplexRegion, IJ; kwargs...)

    "imageset: http://docs.sympy.org/latest/modules/sets.html"
    global imageset(fn::Function, args...) = begin
        x = Sym("x")
        imageset(x, fn(x), args...)
    end
    global imageset(args...) = sympy_meth(:imageset, args...)

    ## Interval

"""
Create an interval object

```
Interval(0,1) # [0,1]
Interval(0,1,true, false) # (0,1]
```

"""
    global Interval(l,r,left_open=false, right_open=false) = sympy_meth(:Interval, Sym(l), Sym(r), left_open ,right_open)

end