"""
MPC controller

    #Parameters

    -A differential equation x' = f(x,u)
    -States and inputs linearization points
    -States and inputs constraints
    -A sample time for discretization
    -A method, 'l'
    -Weithing matrices with appropriate dimension Q and R and P
    _Horizon length N

    #Example

    regulator = mpc_problem(f, x_cons, u_cons, Ts, method, Q, R, P, N)

"""

function mpc_problem(f, xl::AbstractArray, ul::AbstractArray, x_cons::AbstractArray, u_cons::AbstractArray, Ts::Float64, method::String, Q::AbstractArray, R::Array, N::Integer)
    if method == "l"
        # the dynamical system is linearised
        JacobM = ForwardDiff.jacobian(f, vcat(xl, ul))
        A = JacobM[begin:size(xl,1), begin:size(xl,1)]
        B = JacobM[begin:size(xl,1), size(xl,1)+1:end]

        # the linear system is dicretized
        sysc = ControlSystems.ss(A, B,  Matrix(LinearAlgebra.I, size(A,1), size(A,1)), zeros(size(B,1),size(B,2)))
        sysd = ControlSystems.c2d(sysc, Ts)

        # The tmerinal weighting matrix is computed
        P = ModelPredictiveControl.terminal_weight(sysd.A, sysd.B, Q, R)
        controller = ModelPredictiveControl.linear_mpc(sysd.A, sysd.B, x_cons, u_cons, Q, R, P, N)
    end

    return controller

end

"""
Terminal matrix computation of the terminal cost x'Px. P is the solution to the discrete-time algebraic Riccati equation.

    #Parameters
    -State and input discrete matrices A and B
    _Weighting matrices Q and R

    #Example

    P = terminal_weight(Ad, Bd, Q, R)

"""

function terminal_weight(Ad::AbstractMatrix, Bd::AbstractMatrix, Q::AbstractMatrix, R::AbstractMatrix)
    return ControlSystems.dare(Ad, Bd, Q, R)
end


using Random
using Distributions
using StatsBase
using LinearAlgebra
using SpecialFunctions

using DelimitedFiles

using Printf
using ProgressMeter

include("./utils.jl")
include("./load_data.jl")

###############################################################################
###############################################################################
### disease model

abstract type AbstractDiseaseModel end

"""
    DiseaseModel

define disease model that characeterises per-agent disease progression starting
from infection at time 0, to being infectious after the latent period, to begin
no longer infectious after the infectious period.
"""
struct UnivariateDiseaseModel{T1,T2} <: AbstractDiseaseModel
    P_latent::UnivariateDistribution{T1}
    P_infectious::UnivariateDistribution{T2}
end

struct DeltaDiseaseModel{T<:Real} <: AbstractDiseaseModel
    P_latent::Normal{T}
    P_infectious::Normal{T}

    function DeltaDiseaseModel{T}(mean_latent, mean_infectious) where T
        new{T}(Normal(mean_latent,0), Normal(mean_infectious, 0))
    end
end
DeltaDiseaseModel(mean_latent, mean_infectious) = DeltaDiseaseModel{Float64}(mean_latent, mean_infectious)

# default example is delta
DiseaseDefault() = DeltaDiseaseModel(seconds_from_days(3.0), seconds_from_days(5.0))

# generically distributed periods (cf. [LLoyed-Smith et al, Nature, 2001] for motivation of Gamma)
GammaDiseaseModel(mean_latent, mean_infectious, k) = UnivariateDiseaseModel(Gamma(k, mean_latent/k), Gamma(k, mean_infectious/k))

# literature based guesses for pandemic diseases
Covid19() =  DeltaDiseaseModel(NaN,NaN)
H1N1()    =  DeltaDiseaseModel(NaN,NaN)

###############################################################################
###############################################################################
### Viral load

"""
    viral_load(model, time_range; samples=Int(1e6), seed=1000)

evaluates for a disease `model` the ensemble probability to be infectious along
a `time_range` (for each step).

# Arguments
    * `model`: DiseaseModel
    * `time_range`: range over which the probability is evaluated.

# Keyword argeuments
    * `samples`: the number of disease progressions realized (ensemble size)
    * `seed`: seed of MersenneTwister random number generator
"""
function viral_load(model::AbstractDiseaseModel, time_range; samples::Int=Int(1e6), seed::Int=1000)
    print("... evaluate viral load\n")
    rng = MersenneTwister(seed)
    load = Histogram(time_range)

    # define local function
    time_to_index(time) = StatsBase.binindex(load,time)

    p = Progress(samples, 1, "Viral Load: ", offset=1)
    for i in 1:samples
        next!(p)
        index_start, index_end = time_to_index.(infectious_interval(model, rng))
        if index_end > length(load.weights)
            load.weights[index_start:index_end-1] .+= 1
        else
            load.weights[index_start:index_end] .+= 1
        end
    end

    load = float(load)
    load.weights ./= float(samples)

    return load
end

function infectious_interval(model::AbstractDiseaseModel, rng::AbstractRNG)
    time_infectious_start = rand(rng, model.P_latent)
    time_infectious_end   = time_infectious_start + rand(rng, model.P_infectious)
    return time_infectious_start, time_infectious_end
end

function infectious_interval(model::DeltaDiseaseModel)
    time_infectious_start = mean(model.P_latent)
    time_infectious_end   = time_infectious_start + mean(model.P_infectious)
    return time_infectious_start, time_infectious_end
end

###############################################################################
###############################################################################
### evaluate/sample infectious encounter

"""
generate samples only with start times that are below time_max.

This is a control.
"""
function samples_infectious_encounter_with_infection_in_early_interval(
        model::AbstractDiseaseModel,
        ets::encounter_trains{I,T},
        time_max::Number;
        #optional
        samples::Int=Int(1e6),
        rng::AbstractRNG=MersenneTwister(1000)
    ) where {I,T}
    # manipulate the cumulative sum of valid encounters to only draw from early
    # interval
    times_onset = similar(trains(ets));
    for (i,train) in enumerate(trains(ets))
        times_onset[i] = trains(ets)[i][trains(ets)[i] .< time_max];
    end
    ets_onset = encounter_trains(times_onset, ids(ets), duration(ets), timestep(ets));

    if sum(length.(ets_onset)) > 0
        data = [Int64[] for i in 1:length(ets)]
        p = Progress(samples, 1, "Sample Infectious Encounter: ", offset=1)
        for i in 1:samples
            next!(p)
            index_train, index_time = random_encoutner(rng, ets_onset)
            times = trains(ets)[index_train]
            time = times[index_time]
            interval_count = time .+ infectious_interval(model, rng)
            encounter = count_times_in_interval(times[index_time+1:end], interval_count)
            push!(data[index_train], encounter)
        end
        return data
    else
        error("The chosen interval leaves no valid disease start points")
    end
end

function samples_infectious_encounter_random_onset(
        model::AbstractDiseaseModel,
        ets::encounter_trains{I,T};
        #optional
        weighted_train::Bool=false,
        weighted_time::Bool=false,
        samples::Int=Int(1e6),
        rng::AbstractRNG=MersenneTwister(1000)
    ) where {I,T}
    support_random_time = 0:timestep(ets):duration(ets)
    if weighted_train
        weight_random_train = ProbabilityWeights(length.(ets))
    end
    if weighted_time
        encounter_rate = rate(ets, 0:timestep(ets):seconds_from_days(7))
        weight_random_time = ProbabilityWeights(vcat(encounter_rate.weights*4))
    end

    data = [Int64[] for i in 1:length(ets)]
    p = Progress(samples, 1, "Sample Infectious Encounter: ", offset=1)
    for i in 1:samples
        next!(p)
        # try several times to sample
        count_bad_tries = 0
        # first draw and fix realization of diases to not introduce a bias when
        # long diseasess are not fitting into data set
        interval_infectious = infectious_interval(model, rng)
        while true
            # draw random starting time unitl disases fits into remaining duration
            if weighted_train
                rand_train = sample(rng, 1:length(ets), weight_random_train)
            else
                rand_train = rand(rng, 1:length(ets))
            end
            if weighted_time
                rand_time = sample(rng, support_random_time, weight_random_time)
            else
                rand_time  = rand(rng, support_random_time)
            end
            train = trains(ets)[rand_train]
            interval_count = rand_time .+ interval_infectious
            if last(interval_count) < duration(ets)
                encounter = count_times_in_interval(train[searchsortedlast(train, rand_time)+1:end], interval_count)
                push!(data[rand_train], encounter)
                break
            end
            count_bad_tries += 1
            if count_bad_tries == 1000
                throw(BoundsError(ets, "infectious interval does not fit into contact trains"))
            end
        end
    end
    return data
    return samples_infectious_encounter_random_onset(model, ets, support_random_time = support_random_time, weight_random_time = weight_random_time)
end

"""
    samples_infectious_encounter(model, ets [, samples, rng])

returns a data object (A vector of length trains, each containing a vector with
samples of number of infectious encounter for a particular train). A particular
sample is obtained with `sample_infectious_encounter`.
"""
function samples_infectious_encounter(
        model::AbstractDiseaseModel,
        ets::encounter_trains{I,T};
        #optional
        samples::Int=Int(1e6),
        rng::AbstractRNG=MersenneTwister(1000)
    ) where {I,T}
    data = [Int64[] for i in 1:length(ets)]
    p = Progress(samples, 1, "Sample Infectious Encounter: ", offset=1)
    for i in 1:samples
        next!(p)
        # try several times to sample
        num_encounter, index_train = sample_infectious_encounter(ets, model, rng)
        push!(data[index_train], num_encounter)
    end
    return data
end

"""
   sample_infectious_encounter(source, infectious_interval, rng)
"""
function sample_infectious_encounter(
        ets::encounter_trains{I,T},
        model::AbstractDiseaseModel,
        rng::AbstractRNG
    ) where {I,T}
    count_bad_tries = 0
    # first draw and fix realization of diases to not introduce a bias when
    # long diseasess are not fitting into data set
    interval_infectious = infectious_interval(model, rng)
    while true
        # draw random starting time unitl disases fits into remaining duration
        index_train, index_time = random_encoutner(rng, ets)
        times = trains(ets)[index_train]
        time = times[index_time]
        interval_count = time .+ interval_infectious
        if last(interval_count) < duration(ets)
            encounter = count_times_in_interval(times[index_time+1:end], interval_count)
            return encounter, index_train
        end
        count_bad_tries += 1
        if count_bad_tries == 1000
            throw(BoundsError(ets, "infectious interval does not fit into contact trains"))
        end
    end
end

#special case for deltaDisease
function samples_infectious_encounter(
        model::DeltaDiseaseModel,
        ets::encounter_trains{I,T};
        #optional
    ) where {I,T}
    interval_infectious  = infectious_interval(model)
    train_duration = duration(ets)
    data = [Int64[] for j in 1:length(ets)]
    for (j,train) in enumerate(trains(ets))
        for (i, time_contact) in enumerate(train)
            interval_infectious_contact = time_contact .+ interval_infectious
            # counter infectious counters if remaining train fits into infectious window
            if last(interval_infectious_contact) < train_duration
                encounter = count_times_in_interval(train[i+1:end], interval_infectious_contact)
                push!(data[j], encounter)
            else # break the for loop because all disease will not fit into remaining duration for later times either
                break
            end
        end
    end
    return data
end

"""
    mean_infectious_encounter(model, contact; samples, rng)

estimatee the mean of infectious encounters for a univariate disease
`model` and a source of `contacts` (can be contact_trains object or something
else)
"""
function mean_infectious_encounter(
        model::AbstractDiseaseModel,
        ets::encounter_trains{I,T};
        #optional
        samples::Int=Int(1e6),
        rng::AbstractRNG=MersenneTwister(1000)
    ) where {I,T}
    data = samples_infectious_encounter(model, ets, samples=samples, rng=rng)
    return mean(vcat(data...))
end
function mean_infectious_encounter(
        model::DeltaDiseaseModel,
        ets::encounter_trains{I,T};
        #optional
    ) where {I,T}
    data = samples_infectious_encounter(model, ets)
    return mean_from_samples_infectious_encounter(data)
end
mean_from_samples_infectious_encounter(data) = mean(vcat(data...))

"""
    distribution_infectious_encounter(model, contact; samples, seed)

sample the distribution of infectious encounters for a univariate disease
`model` and a source of `contacts` (can be contact_trains object or something
else)
"""
function distribution_infectious_encounter(
        model::AbstractDiseaseModel,
        ets::encounter_trains{I,T};
        ### optional
        edges=missing,
        samples::Int=Int(1e6),
        rng::AbstractRNG=MersenneTwister(1000)
    ) where {I,T}
    data = samples_infectious_encounter(model, ets)
    return distribution_from_samples_infectious_encounter(data, edges=edges)
end
function distribution_infectious_encounter(
        model::DeltaDiseaseModel,
        ets::encounter_trains{I,T};
        ### optional
        edges=missing,
    ) where {I,T}
    data = samples_infectious_encounter(model, ets)
    return distribution_from_samples_infectious_encounter(data, edges=edges)
end
function distribution_from_samples_infectious_encounter(
        data;
        ### optional
        edges=missing,
    )
    if ismissing(edges)
        list_encounter = vcat(data...)
        dist = fit(Histogram{Float64}, list_encounter, 0:1:maximum(list_encounter)+2)
    else
        dist = fit(Histogram{Float64}, vcat(data...), edges)
    end
    normalize!(dist)
    return dist
end

function count_times_in_interval(times::Vector{T1}, interval_count::Tuple{T2,T2}) where {T1<:Number,T2<:Number}
    encounter = 0
    for time in times
        encounter += ininterval(time, interval_count)
    end
    return encounter
end
function ininterval(time::T1, interval_count::Tuple{T2,T2}) where {T1<:Number,T2<:Number}
    return (getindex(interval_count,1) <= time < getindex(interval_count,2))
end


### A Pluto.jl notebook ###
# v0.11.14

using Markdown
using InteractiveUtils

# ╔═╡ 545bb79c-fbf0-11ea-3455-67441d62f5fe
using Flux

# ╔═╡ 741102a8-fc05-11ea-2352-0bd39f49fcd3
using Statistics

# ╔═╡ 11f5a760-fbf4-11ea-286d-c9f6f30f83ed
using Plots

# ╔═╡ c4878102-fbf0-11ea-298e-eb205469990c
images = Flux.Data.FashionMNIST.images()

# ╔═╡ 01be0a34-fbf1-11ea-0beb-0521a2d634c4
labels = Flux.Data.FashionMNIST.labels()

# ╔═╡ 459673f4-fbf1-11ea-114a-53b01f61a81b
size(images[1])

# ╔═╡ 78f37170-fbf1-11ea-3dab-07eed2cba687
function get_fashion_label(labels)
	text_label = ["t-shirt", "trouser", "pullover", "dress", "coat", "sandal", "shirt", "sneaker", "bag", "ankle boot"]
	return [text_label[i+1] for i in labels]
end

# ╔═╡ 05baaefc-fbf2-11ea-0a8b-29a1748f44d3
get_fashion_label(labels)

# ╔═╡ a5e15114-fbf3-11ea-1660-d3aaa45ff6c4
images[1]

# ╔═╡ 3e7e073c-fbf4-11ea-03b7-e723e9d36fe6
plot(images[1], title="short")

# ╔═╡ 93b8a5d6-fbf4-11ea-265e-5d6c49cc4242
m = Chain(Dense(784,10), softmax)

# ╔═╡ e007cc00-fbf4-11ea-32d8-f11128553f4f
size(images)

# ╔═╡ ce10f3f2-fc01-11ea-0c30-3945431c0fb2
Float64.(reshape(images[1],(784,1)))

# ╔═╡ 5d11b170-fc04-11ea-1c67-e7df917005f9
train_x = rand(784, 1000)

# ╔═╡ 76a58fd0-fc04-11ea-3ef0-5b61e18dac67
begin
	for i in 1:1000
		train_x[:,i] = Float64.(reshape(images[i],(784,1)))
	end
end

# ╔═╡ 8c989422-fc04-11ea-0a3c-d5ec2e7bc9dd
train_x

# ╔═╡ 9498eb48-fc01-11ea-24f5-c130e859c64f
get_fashion_label(labels)[1]

# ╔═╡ aff167a8-fc01-11ea-3e6c-5323f94be667
begin

	train_y = []
	for i in 1:1000

		push!(train_y, get_fashion_label(labels)[i])
	end
end

# ╔═╡ 1cdac5f0-fc02-11ea-1db6-67ec6ff27617
size(train_y)

# ╔═╡ 480a9cfc-fc03-11ea-0d4e-0d20b008513c
klasses = sort(unique(train_y))

# ╔═╡ c49d6140-fc04-11ea-12ed-612c4eafc3be
onehot_labels = Flux.onehotbatch(train_y, klasses)

# ╔═╡ 997dc9bc-fc04-11ea-2b6c-f1a54753f29a
train_data = Iterators.repeated((train_x, onehot_labels), 150)

# ╔═╡ acf6d226-fc02-11ea-3f18-493597b870d2
loss(x,y) = Flux.crossentropy(m(x), y)

# ╔═╡ 565b20cc-fc05-11ea-2306-53574dd81833
accuracy(x, y, model) = Flux.mean(Flux.onecold(model(x)) .== Flux.onecold(y))

# ╔═╡ c665f7c8-fc02-11ea-3eb8-9b0c019e5d87
opt = Descent(0.3)

# ╔═╡ df3dd4d0-fc04-11ea-0110-b3d5a5301326
Flux.train!(loss, Flux.params(m), train_data, opt)

# ╔═╡ 0bb69ca2-fc05-11ea-1ecf-3116ffb13f62
Flux.onecold(m(train_x[:,8]))

# ╔═╡ 34edea8c-fc05-11ea-1871-a36b2dc53a52
Flux.onecold(onehot_labels[:,8])

# ╔═╡ 5f892cac-fc05-11ea-247e-5d1a02a89a8a
accuracy(train_x, onehot_labels, m)

# ╔═╡ c351ba60-fc05-11ea-32c9-e5e8a65bc239
sum(Flux.onecold(m(train_x)) .== Flux.onecold(onehot_labels))

# ╔═╡ Cell order:
# ╠═545bb79c-fbf0-11ea-3455-67441d62f5fe
# ╠═741102a8-fc05-11ea-2352-0bd39f49fcd3
# ╠═c4878102-fbf0-11ea-298e-eb205469990c
# ╠═01be0a34-fbf1-11ea-0beb-0521a2d634c4
# ╠═459673f4-fbf1-11ea-114a-53b01f61a81b
# ╠═78f37170-fbf1-11ea-3dab-07eed2cba687
# ╠═05baaefc-fbf2-11ea-0a8b-29a1748f44d3
# ╠═a5e15114-fbf3-11ea-1660-d3aaa45ff6c4
# ╠═11f5a760-fbf4-11ea-286d-c9f6f30f83ed
# ╠═3e7e073c-fbf4-11ea-03b7-e723e9d36fe6
# ╠═93b8a5d6-fbf4-11ea-265e-5d6c49cc4242
# ╠═e007cc00-fbf4-11ea-32d8-f11128553f4f
# ╠═ce10f3f2-fc01-11ea-0c30-3945431c0fb2
# ╠═5d11b170-fc04-11ea-1c67-e7df917005f9
# ╠═76a58fd0-fc04-11ea-3ef0-5b61e18dac67
# ╠═8c989422-fc04-11ea-0a3c-d5ec2e7bc9dd
# ╠═9498eb48-fc01-11ea-24f5-c130e859c64f
# ╠═aff167a8-fc01-11ea-3e6c-5323f94be667
# ╠═997dc9bc-fc04-11ea-2b6c-f1a54753f29a
# ╠═1cdac5f0-fc02-11ea-1db6-67ec6ff27617
# ╠═480a9cfc-fc03-11ea-0d4e-0d20b008513c
# ╠═c49d6140-fc04-11ea-12ed-612c4eafc3be
# ╠═acf6d226-fc02-11ea-3f18-493597b870d2
# ╠═565b20cc-fc05-11ea-2306-53574dd81833
# ╠═c665f7c8-fc02-11ea-3eb8-9b0c019e5d87
# ╠═df3dd4d0-fc04-11ea-0110-b3d5a5301326
# ╠═0bb69ca2-fc05-11ea-1ecf-3116ffb13f62
# ╠═34edea8c-fc05-11ea-1871-a36b2dc53a52
# ╠═5f892cac-fc05-11ea-247e-5d1a02a89a8a
# ╠═c351ba60-fc05-11ea-32c9-e5e8a65bc239


@testset "Expectation" begin
    Mod.@polyvar x[1:3]
    p = x[3] - 2x[1]*x[2]^2 + 3x[3]*x[1] - 5x[1]^3
    v = (1,2,3)
    m = dirac(monomials(p), x=>v)
    @test MultivariateMoments.expectation(m, p) == p(x=>v) == MultivariateMoments.expectation(p, m)
    @test_throws ErrorException dot(x[1] * x[2] * x[3], m)
    @test dot(0.5 * x[1] * x[2]^2, m) == 2.0
    @test dot(m, x[1] * x[3]) == 3
end


export smRecovery

###################################################################
# System Matrix recovery using DCT-sparsity and low-rank constraing
###################################################################
function smRecovery(y::Matrix{T}, samplingIdx::Array{Int64}, params::Dict) where T
  shape = params[:shape]

  # sampling operator
  P = SamplingOp(samplingIdx, shape)

  # sparsifying transformation and regularization
  if !haskey(params, :sparseTrafo)
    params[:sparseTrafo] = DCTOp(ComplexF64,params[:shape],2)
  end
  reg1Name = get(params, :reg1, "L1")
  reg1 = Regularization(reg1Name, params[:λ_l1]; params...)

  # setup low rank regularization
  λ_lr = get(params, :λ_lr, 0.0)
  ρ_lr = get(params, :ρ_lr, 0.0)
  if λ_lr==0 || ρ_lr==0
    # no LR regularization
    reg = [reg1]
    ρ = [params[:ρ_l1]]
  elseif length(shape)==2
    # 2d case (use nuclear norm)
    regNN = regularizationNN(λ_lr, shape)
    reg = [reg1, regNN]
    ρ = [params[:ρ_l1], ρ_lr]
  else
    # 3d case (use tensor nuclear norm)
    regTNN1 = regularizationTNN(λ_lr, shape, 1)
    regTNN2 = regularizationTNN(λ_lr, shape, 2)
    regTNN3 = regularizationTNN(λ_lr, shape, 3)
    reg = [reg1, regTNN1, regTNN2, regTNN3]
    ρ = [params[:ρ_l1], ρ_lr, ρ_lr, ρ_lr]
  end

  # precalculate inverse for split Bregman iteration
  precon = invPrecon(samplingIdx, prod(shape); params...)
  params[:precon] = precon
  params[:iterationsCG] = 1

  # normalized measurement
  y2, y_norm = getNormalizedMeasurement(y)

  solver = [SplitBregman(P; reg=deepcopy(reg), ρ=ρ, params...) for i=1:Threads.nthreads()]

  sfMat = zeros(ComplexF64,prod(shape),size(y,2))
  @time Threads.@threads for k=1:size(y,2)
    t = Threads.threadid()
    sfMat[:,k] .= solve(solver[t], y2[k])
    # undo normalization
    sfMat[:,k] *= y_norm[k]
  end

  return sfMat
end

#########################
# low rank regularization
#########################
function regularizationTNN(λ::AbstractFloat, shape::NTuple{3,Int64}, mode::Int64)
  proxMap = (x,y;kargs...)->proxTNN!(x,y,shape,mode)
  reg = Regularization(prox! = proxMap, λ=λ, params=Dict{Symbol,Any}())
end

function proxTNN!(x, λ::Real, shape::NTuple{3,Int64},mode::Int64; kargs...)
  nx,ny,nz = shape
  # flattened tensor
  x = reshape(x,shape)
  if mode==1
    x_flat = reshape(x,nx,ny*nz)
  elseif mode==2
    x_flat = reshape(permutedims(x,[2,3,1]),ny,nx*nz)
  else
    if mode!=3
      @error "mode in proxTNN! should be 1,2 or 3"
    end
    x_flat = reshape(permutedims(x,[3,1,2]),nz,nx*ny)
  end
  # perform singular value thresholding
  usv = svd!(x_flat)
  proxL1!(usv.S,λ)
  x_flat[:] .= vec(usv.U*Diagonal(usv.S)*usv.Vt)
  # undo flattening operation
  if mode==1
    x[:] .= vec(x_flat)
  elseif mode==2
    x[:] .= vec(permutedims(reshape(x_flat,ny,nz,nx),[3,1,2]))
  else
    x[:] .= vec(permutedims(reshape(x_flat,nz,nx,ny),[2,3,1]))
  end
end

# low rank regularization for 2d (using SVD)
function regularizationNN(λ::Real, shape::NTuple{2,Int64} )
  proxMap = (x,y;kargs...)->proxNN!(x,y,shape;kargs...)
  return Regularization(prox! = proxMap, λ=λ)
end

# singular-value thresholding
function proxNN!(x, λ::Real, shape::NTuple{2,Int64}; kargs...)
  U,S,V = svd(reshape(x,shape))
  proxL1!(S,λ)
  x[:] .= vec(U*Matrix(Diagonal(S))*V')
end

# fixed rank constraing for 3d (using HOSVD)
function FRTruncation!(x::Array{T,3}, rank::NTuple{3,Int64}) where T
  guvw_r = hosvd(real.(x),rank)
  guvw_i = hosvd(imag.(x),rank)
  x[:] .= compose(guvw_r)+1im*compose(guvw_i)
end

# fixed rank constraing for 2d (using SVD)
function FRTruncation!(x, rank::Int64)
  U,S,V = svd(reshape(x, shape))
  S[rank+1:end] .= 0.0
  x[:] .= vec(U*Matrix(Diagonal(S))*V')
end

##########################
# inverse of normal matrix
##########################
mutable struct diagPrecon{T}
  diag::Vector{T}
end

function ldiv!(y::Vector{T},P::diagPrecon{T},x::Vector{T}) where T
  y[:] .= P.diag .* x
end

function ldiv!(P::diagPrecon{T},x::Vector{T}) where T
  x[:] .= P.diag .* x
end

\(P::diagPrecon{T}, x::Vector{T}) where T = P.diag .* x

function invPrecon(samplingIdx, ncol; ρ::Vector{Float64}=[1.0], kargs...)
  normalP = normalPDiag(samplingIdx, ncol)
  diag = [normalP[i]+sum(ρ) for i=1:length(normalP)]
  diag = diag.^-1
  return diagPrecon(diag)
end

function normalPDiag(samplingIdx, ncol; T=ComplexF64)
  normalP = zeros(T,ncol)
  for i=1:length(samplingIdx)
    normalP[samplingIdx[i]] = 1.0
  end
  return normalP
end

function getNormalizedMeasurement(y::Matrix{T}) where T
  y_norm = [norm(y[:,k]) for k=1:size(y,2)]
  y2 = [ComplexF64.(y[:,k]/y_norm[k]) for k=1:size(y,2)]
  return y2, y_norm
end


function pilot(length, freq; samplerate=1)
    L = round(Int, uconvert(NoUnits, length*samplerate))
    dθ = uconvert(NoUnits, 2π*freq/samplerate)

    cos.((0:L-1) * dθ)
end

function pilotfilt(sig, ω)
    # DSP.jl normalizes to half-cycle/sample, so that 1 is the nyquist,
    # but the rest of this code uses 1 cycle/sample (so 0.5 is nyquist)
    bpf = digitalfilter(Bandpass(2ω*0.97, 2ω/0.97),
                        FIRWindow(;transitionwidth=2ω*0.1))
    filt(bpf, sig)[length(bpf):end]
end

"""
Find a pilot tone in the given signal.
"""
function findpilot(sig, freq, len; samplerate=1)
    # TODO: this isn't robust to a wideband energy increase. It looks for an
    # energy pulse in the target band, but that could also get triggered if
    # everything just gets louder
    ω = freqnorm(freq, samplerate)
    L = durnorm(len, samplerate)
    filtered = pilotfilt(sig, ω)

    env = abs2.(hilbert(filtered))
    _, onset = findmax(xcorr(env, ones(L); padmode=:none))
    onset -= L

    prerollpower = sum(x->x^2, filtered[1:onset-1]) / (onset-1)
    pilotpower = sum(x->x^2, filtered[(0:L-1).+onset]) / L

    snr = pow2db(pilotpower/prerollpower)
    if snr < 6
        @warn @sprintf "Narrowband Pilot SNR only %.2f dB" snr
    end
    onset
end

"""
Compute the average phase difference of `sig` (which should correspond to the
frequency). The frequency is given in cycles/sample
"""
function avgphasediff(sig)
    # we remove the ends to get rid of the ringing from the hilbert transform
    an = hilbert(sig)[2000:end-2000]
    diff = @view(an[2:end]) ./ @view an[1:end-1]
    mean(angle, diff)/2π
end

"""
Takes the given signal and pilot frequency and returns the resample factor
necessary to synchronize the signal.
"""
function pilotsync(sig, freq; samplerate=1)
    ω = freqnorm(freq, samplerate)
    # bpf = gaussian(N, 0.20) .* cos.((0:N-1) * ω)
    # filtered = filt(bpf, sig)[length(bpf):end]
    filtered = pilotfilt(sig, ω)
    ω̂₁ = avgphasediff(filtered)
    @info "frequency estimated from hilbert transform: $(ω̂₁)"

    # estimate amplitude, frequency, and phase
    model(n, p) = @. p[1] * cos(2π*p[2]*n + p[3])

    fit = curve_fit(model, (0:length(filtered)-1), filtered,
                    [1.0, ω̂₁, 0.0], autodiff=:forwarddiff)

    # @info "frequency estimated from lsq fit: $(fit.param[2]/2π*samplerate)"
    @info "fit amplitude: $(fit.param[1])"
    @info "fit frequency: $(fit.param[2] * samplerate)"
    @info "fit phase (rad): $(fit.param[3])"

    fit.param[2] / ω
end

"""
    phaselock(x; alpha=0.05)

Simple PLL (phase-locked loop). Assumes `x` to be a real or complex sinusoidal
signal. Returns a tuple of (freq, phase, err, x2), where each is a vector with
the same length as `x`. `x2` will be real or complex to match `x`.

`alpha` controls how agressive the PLL is in correcting its frequency and phase
estimates. Lower values are more stable but converge more slowly.

## Example

```
using Plots

in_p = 0:0.1:500
x = cos.(in_p) + randn(length(in_p))*0.3

f, p, err, x2 = phaselock(x)
```

Algorithm ported from C code at http://liquidsdr.org/blog/pll-simple-howto/
Released under the MIT License, Copyright (c) 2007 - 2016 Joseph Gaeddert
"""
function phaselock(x::Vector{<:Complex}; alpha=0.05)
    beta = 0.5*alpha^2
    freq = 0.0;
    phase = 0.0;
    freqs = zeros(length(x))
    phases = zeros(length(x))
    errs = zeros(length(x))
    out = zeros(Complex{Float64}, length(x))

    for i in eachindex(x)
        out[i] = exp(im*phase)
        err = angle(x[i]*conj(out[i]))

        freq += beta * err
        phase += freq + alpha * err

        freqs[i] = freq/2π
        phases[i] = phase/2π
        errs[i] = err
    end

    (freqs, phases, errs, out)
end

function phaselock(x::Vector{<:Real}; kwargs...)
    freqs, phases, errs, out = phaselock(hilbert(x); kwargs...)

    (freqs, phases, errs, real.(out))
end


"""
    TfidfTransformer()

The following is taken largely from scikit-learn's documentation:
https://github.com/scikit-learn/scikit-learn/blob/main/sklearn/feature_extraction/text.py

Convert a collection of raw documents to a matrix of TF-IDF features.

"TF" means term-frequency while "TF-IDF" means term-frequency times
inverse document-frequency.  This is a common term weighting scheme in
information retrieval, that has also found good use in document
classification.

The goal of using TF-IDF instead of the raw frequencies of occurrence
of a token in a given document is to scale down the impact of tokens
that occur very frequently in a given corpus and that are hence
empirically less informative than features that occur in a small
fraction of the training corpus.

The formula that is used to compute the TF-IDF for a term `t` of a
document `d` in a document set is `tf_idf(t, d) = tf(t, d) *
idf(t)`. Assuming `smooth_idf=false`, `idf(t) = log [ n / df(t) ] + 1`
where `n` is the total number of documents in the document set and
`df(t)` is the document frequency of `t`. The document frequency is
the number of documents in the document set that contain the term
`t`. The effect of adding “1” to the idf in the equation above is that
terms with zero idf, i.e., terms that occur in all documents in a
training set, will not be entirely ignored. (Note that the idf formula
above differs from that appearing in standard texts, `idf(t) = log [ n
/ (df(t) + 1) ])`.

If `smooth_idf=true` (the default), the constant “1” is added to the
numerator and denominator of the idf as if an extra document was seen
containing every term in the collection exactly once, which prevents
zero divisions: `idf(t) = log [ (1 + n) / (1 + df(t)) ] + 1`.

The parameters `max_doc_freq` and `min_doc_freq` restrict the vocabulary
that the transformer will consider. `max_doc_freq` indicates that terms in only
up to the specified percentage of documents will be considered. For example, if
`max_doc_freq` is set to 0.9, terms that are in more than 90% of documents
will be removed. Similarly, the `min_doc_freq` parameter restricts terms in the
other direction. A value of 0.01 means that only terms that are at least in 1% of
documents will be included.
"""
mutable struct TfidfTransformer <: AbstractTextTransformer
    max_doc_freq::Float64
    min_doc_freq::Float64
    smooth_idf::Bool
end

function TfidfTransformer(; max_doc_freq::Float64 = 1.0, min_doc_freq::Float64 = 0.0, smooth_idf::Bool = true)
    transformer = TfidfTransformer(max_doc_freq, min_doc_freq, smooth_idf)
    message = MMI.clean!(transformer)
    isempty(message) || @warn message
    return transformer
end

struct TfidfTransformerResult
    vocab::Vector{String}
    idf_vector::Vector{Float64}
end

get_result(::TfidfTransformer, idf::Vector{<:AbstractFloat}, vocab::Vector{String}, ::SparseMatrixCSC) =
    TfidfTransformerResult(vocab, idf)

function build_tfidf!(doc_term_mat::SparseMatrixCSC{T},
                      tfidf::SparseMatrixCSC{F},
                      idf_vector::Vector{F}) where {T <: Real, F <: AbstractFloat}
    rows = rowvals(doc_term_mat)
    dtmvals = nonzeros(doc_term_mat)
    tfidfvals = nonzeros(tfidf)
    @assert size(dtmvals) == size(tfidfvals)

    p, n = size(doc_term_mat)

    # TF tells us what proportion of a document is defined by a term
    words_in_documents = F.(sum(doc_term_mat; dims=1))
    oneval = one(F)

    @inbounds for i = 1:n
        for j in nzrange(doc_term_mat, i)
            row = rows[j]
            tfidfvals[j] = dtmvals[j] / max(words_in_documents[i], oneval) * idf_vector[row]
        end
    end

    return tfidf
end

function _transform(::TfidfTransformer,
                    result::TfidfTransformerResult,
                    v::Corpus)
    doc_terms = build_dtm(v, result.vocab)
    tfidf = similar(doc_terms.dtm, eltype(result.idf_vector))
    build_tfidf!(doc_terms.dtm, tfidf, result.idf_vector)

    # here we return the `adjoint` of our sparse matrix to conform to
    # the `n x p` dimensions throughout MLJ
    return adjoint(tfidf)
end

# for returning user-friendly form of the learned parameters:
function MMI.fitted_params(::TfidfTransformer, fitresult)
    vocab = fitresult.vocab
    idf_vector = fitresult.idf_vector
    return (vocab = vocab, idf_vector = idf_vector)
end

## META DATA

MMI.metadata_pkg(TfidfTransformer,
             name="$PKG",
             uuid="7876af07-990d-54b4-ab0e-23690620f79a",
             url="https://github.com/JuliaAI/MLJText.jl",
             is_pure_julia=true,
             license="MIT",
             is_wrapper=false
)

MMI.metadata_model(TfidfTransformer,
               input_scitype = Union{
                   AbstractVector{<:AbstractVector{STB.Textual}},
                   AbstractVector{<:STB.Multiset{<:ScientificNGram}},
                   AbstractVector{<:STB.Multiset{STB.Textual}}
                   },
               output_scitype = AbstractMatrix{STB.Continuous},
               docstring = "Build TF-IDF matrix from raw documents",
               path = "MLJText.TfidfTransformer"
               )


# A quadratic quadrilateral in 3D space.

struct Quadrilateral8_3D{T <: AbstractFloat}
    # The points are assumed to be ordered  in counter clockwise order as follows
    # p₁ = vertex A
    # p₂ = vertex B
    # p₃ = vertex C
    # p₄ = vertex D
    # p₅ = point on the quadratic segment from A to B
    # p₆ = point on the quadratic segment from B to C
    # p₇ = point on the quadratic segment from C to D
    # p₈ = point on the quadratic segment from D to A
    points::NTuple{8, Point_3D{T}}
end

# Constructors
# -------------------------------------------------------------------------------------------------
Quadrilateral8_3D(p₁::Point_3D{T}, p₂::Point_3D{T}, p₃::Point_3D{T}, p₄::Point_3D{T},
                  p₅::Point_3D{T}, p₆::Point_3D{T}, p₇::Point_3D{T}, p₈::Point_3D{T}
        ) where {T <: AbstractFloat} = Quadrilateral8_3D((p₁, p₂, p₃, p₄, p₅, p₆, p₇, p₈))

# Methods
# -------------------------------------------------------------------------------------------------
function (quad8::Quadrilateral8_3D{T})(r::R, s::S) where {T <: AbstractFloat,
                                                          R <: Real,
                                                          S <: Real}
    # See The Visualization Toolkit: An Object-Oriented Approach to 3D Graphics, 4th Edition
    # Chapter 8, Advanced Data Representation, in the interpolation functions section
    ξ = 2T(r) - 1; η = 2T(s) - 1
    return (1 - ξ)*(1 - η)*(-ξ - η - 1)/4*quad8.points[1] +
           (1 + ξ)*(1 - η)*( ξ - η - 1)/4*quad8.points[2] +
           (1 + ξ)*(1 + η)*( ξ + η - 1)/4*quad8.points[3] +
           (1 - ξ)*(1 + η)*(-ξ + η - 1)/4*quad8.points[4] +
                      (1 - ξ^2)*(1 - η)/2*quad8.points[5] +
                      (1 - η^2)*(1 + ξ)/2*quad8.points[6] +
                      (1 - ξ^2)*(1 + η)/2*quad8.points[7] +
                      (1 - η^2)*(1 - ξ)/2*quad8.points[8]
end

function derivatives(quad8::Quadrilateral8_3D{T}, r::R, s::S) where {T <: AbstractFloat,
                                                                     R <: Real,
                                                                     S <: Real}
    # Chain rule
    # ∂Q   ∂Q ∂ξ  ∂Q   ∂Q ∂η
    # -- = -- --, -- = -- --
    # ∂r   ∂ξ ∂r  ∂s   ∂η ∂s
    ξ = 2T(r) - 1; η = 2T(s) - 1
    ∂Q_∂ξ = (1 - η)*(2ξ + η)/4*quad8.points[1] +
            (1 - η)*(2ξ - η)/4*quad8.points[2] +
            (1 + η)*(2ξ + η)/4*quad8.points[3] +
            (1 + η)*(2ξ - η)/4*quad8.points[4] +
                    -ξ*(1 - η)*quad8.points[5] +
                   (1 - η^2)/2*quad8.points[6] +
                    -ξ*(1 + η)*quad8.points[7] +
                  -(1 - η^2)/2*quad8.points[8]

    ∂Q_∂η = (1 - ξ)*( ξ + 2η)/4*quad8.points[1] +
            (1 + ξ)*(-ξ + 2η)/4*quad8.points[2] +
            (1 + ξ)*( ξ + 2η)/4*quad8.points[3] +
            (1 - ξ)*(-ξ + 2η)/4*quad8.points[4] +
                   -(1 - ξ^2)/2*quad8.points[5] +
                     -η*(1 + ξ)*quad8.points[6] +
                    (1 - ξ^2)/2*quad8.points[7] +
                     -η*(1 - ξ)*quad8.points[8]

    return 2*∂Q_∂ξ, 2*∂Q_∂η
end

function area(quad8::Quadrilateral8_3D{T}; N::Int64=15) where {T <: AbstractFloat}
    # Numerical integration required. Gauss-Legendre quadrature over a quadrilateral is used.
    # Let Q(r,s) be the interpolation function for quad8,
    #                             1  1
    # A = ∬ ||∂Q/∂r × ∂Q/∂s||dA = ∫  ∫ ||∂Q/∂r × ∂Q/∂s|| ds dr
    #      D                      0  0
    #
    #       N   N
    #   =   ∑   ∑  wᵢwⱼ||∂Q/∂r(rᵢ,sⱼ) × ∂Q/∂s(rᵢ,sⱼ)||
    #      i=1 j=1
    # N is the square root of the number of points used in the quadrature.
    # See tuning/Quadrilateral8_3D_area.jl for more info on how N = 15 was chosen.
    w, r = gauss_legendre_quadrature(T, N)
    a = T(0)
    for i = 1:N, j = 1:N
        ∂Q_∂r, ∂Q_∂s = derivatives(quad8, r[i], r[j])
        a += w[i]*w[j]*norm(∂Q_∂r × ∂Q_∂s)
    end
    return a
end

function triangulate(quad8::Quadrilateral8_3D{T}, N::Int64) where {T <: AbstractFloat}
    # N is the number of divisions of each edge
    triangles = Vector{Triangle_3D{T}}(undef, 2*(N+1)*(N+1))
    if N === 0
        triangles[1] = Triangle_3D(quad8.points[1], quad8.points[2], quad8.points[3])
        triangles[2] = Triangle_3D(quad8.points[3], quad8.points[4], quad8.points[1])
    else
        for j = 0:N, i = 0:N
            triangles[2*(N+1)*j + 2i + 1] = Triangle_3D(quad8(    i/(N+1),     j/(N+1)),
                                                        quad8((i+1)/(N+1),     j/(N+1)),
                                                        quad8(    i/(N+1), (j+1)/(N+1)))
            triangles[2*(N+1)*j + 2i + 2] = Triangle_3D(quad8(    i/(N+1), (j+1)/(N+1)),
                                                        quad8((i+1)/(N+1),     j/(N+1)),
                                                        quad8((i+1)/(N+1), (j+1)/(N+1)))
        end
    end
    return triangles
end

function intersect(l::LineSegment_3D{T}, quad8::Quadrilateral8_3D{T}; N::Int64 = 20) where {T <: AbstractFloat}
    # Triangulate and intersect each triangle
    triangles = triangulate(quad8, N)
    npoints = 0
    p₁ = Point_3D(T, 0)
    p₂ = Point_3D(T, 0)
    intersections = l .∩ triangles
    bools = map(x->x[1], intersections)
    points = map(x->x[2], intersections)
    npoints = count(bools)
    if npoints === 0
        return false, 0, p₁, p₂
    elseif npoints === 1
        p₁ = points[argmax(bools)]
        return true, 1, p₁, p₂
    elseif npoints === 2
        indices = findall(bools)
        p₁ = points[indices[1]]
        p₂ = points[indices[2]]
        # Check uniqueness
        if p₁ ≈ p₂
            return true, 1, p₁, p₂
        else
            return true, 2, p₁, p₂
        end
    else
        # Account for 3 points and 4 points?
        # If intersection is on edge shared by two triangles on entrance and/or exit 3/4 intersections
        # can be detected
        return true, -1, p₁, p₂
    end
end


using GLMakie

function equidistant_cubed_shell_warp(a, b, c, R = max(abs(a), abs(b), abs(c)))

    r = hypot(a, b, c)

    x1 = R * a / r
    x2 = R * b / r
    x3 = R * c / r

    return x1, x2, x3
end

n = 20

r = range(-1,1,length=n+1)

xa = zeros(n+1,n+1)
xb = zeros(n+1,n+1)
xc = zeros(n+1,n+1)

xa .= r
xb .= r'
xc .= 1

a = [equidistant_cubed_shell_warp(a,b,c)[1] for (a,b,c) in zip(xa,xb,xc)]
b = [equidistant_cubed_shell_warp(a,b,c)[2] for (a,b,c) in zip(xa,xb,xc)]
c = [equidistant_cubed_shell_warp(a,b,c)[3] for (a,b,c) in zip(xa,xb,xc)]

sc = Scene()

wireframe!(sc,
           a,b,c,
           show_axis=false,
           linewidth=1.2)
wireframe!(sc,
           a,b,-c,
           show_axis=false,
           linewidth=1.2)
wireframe!(sc,
           c,a,b,
           show_axis=false,
           linewidth=1.2)
wireframe!(sc,
           -c,a,b,
           show_axis=false,
           linewidth=1.2)
wireframe!(sc,
           b,c,a,
           show_axis=false,
           linewidth=1.2)
wireframe!(sc,
           b,-c,a,
           show_axis=false,
           linewidth=1.2)

sc


lines = readlines(ARGS[1])

function solution1(lines)
    val_matrix = []
    for line in lines
        line_vals = [parse(Int, v) for v in line]
        append!(val_matrix, [line_vals])
    end
    num_rows = length(val_matrix)
    row_width = length(val_matrix[1])
    minima = []
    for (i, row) in enumerate(val_matrix)
        for (j, val) in enumerate(row)
            is_minimal = true
            # check above
            if i > 1
                if val >= val_matrix[i-1][j]
                    is_minimal = false
                end
            end
            # check below
            if i < num_rows
                if val >= val_matrix[i+1][j]
                    is_minimal = false
                end
            end
            # check left
            if j > 1
                if val >= val_matrix[i][j-1]
                    is_minimal = false
                end
            end
            # check right
            if j < row_width
                if val >= val_matrix[i][j+1]
                    is_minimal = false
                end
            end
            if is_minimal
                append!(minima, val + 1)
            end
        end
    end
    return sum(minima)
end

function crawl_basin(basin, val_matrix, num_rows, row_width)
    to_search = Set(basin)
    found = Set(basin)
    while length(to_search) > 0
        check_point = pop!(to_search)
        check_x = check_point[1]
        check_y = check_point[2]
        # check above
        if check_y - 1 >= 1 && val_matrix[check_x][check_y-1] != 9 && !in([check_x, check_y - 1], found)
            push!(to_search, [check_x, check_y - 1])
            push!(found, [check_x, check_y - 1])
        end
        # check below
        if check_y + 1 <= num_rows && val_matrix[check_x][check_y+1] != 9 && !in([check_x, check_y + 1], found)
            push!(to_search, [check_x, check_y + 1])
            push!(found, [check_x, check_y + 1])
        end
        # check left
        if check_x - 1 >= 1 && val_matrix[check_x-1][check_y] != 9 && !in([check_x - 1, check_y], found)
            push!(to_search, [check_x - 1, check_y])
            push!(found, [check_x - 1, check_y])
        end
        # check right
        if check_x + 1 <= row_width && val_matrix[check_x+1][check_y] != 9 && !in([check_x + 1, check_y], found)
            push!(to_search, [check_x + 1, check_y])
            push!(found, [check_x + 1, check_y])
        end
    end
    return [x for x in found]
end

function solution2(lines)
    val_matrix = []
    for line in lines
        line_vals = [parse(Int, v) for v in line]
        append!(val_matrix, [line_vals])
    end
    num_rows = length(val_matrix)
    row_width = length(val_matrix[1])
    minima = []
    basins = []
    for (i, row) in enumerate(val_matrix)
        for (j, val) in enumerate(row)
            is_minimal = true
            basin = [[i, j]]
            # check above
            if i > 1
                if val >= val_matrix[i-1][j]
                    is_minimal = false
                else
                    if val_matrix[i-1][j] != 9
                        append!(basin, [[i - 1, j]])
                    end

                end
            end
            # check below
            if i < num_rows
                if val >= val_matrix[i+1][j]
                    is_minimal = false
                else
                    if val_matrix[i+1][j] != 9
                        append!(basin, [[i + 1, j]])
                    end
                end
            end
            # check left
            if j > 1
                if val >= val_matrix[i][j-1]
                    is_minimal = false
                else
                    if val_matrix[i][j-1] != 9
                        append!(basin, [[i, j - 1]])
                    end
                end
            end
            # check right
            if j < row_width
                if val >= val_matrix[i][j+1]
                    is_minimal = false
                else
                    if val_matrix[i][j+1] != 9
                        append!(basin, [[i, j + 1]])
                    end
                end
            end
            if is_minimal
                append!(minima, val + 1)
                append!(basins, [basin])
            end
        end
    end

    crawled_basins = []
    for basin in basins
        append!(crawled_basins, length(crawl_basin(basin, val_matrix, num_rows, row_width)))
    end

    rev_lengths = reverse(sort(crawled_basins))

    return rev_lengths[1] * rev_lengths[2] * rev_lengths[3]
end

println(solution1(lines))

println(solution2(lines))


export omega, omega_unscaled

# Tools to compute the sparsity pattern

function omega(h, M, Y, X, Ny, Nx, ϵY, L::LinearTransform)
    Ne = size(X, 2)
    Omega = zeros(Nx+Ny, Nx+Ny)

    precisionY = (1/ϵY^2)
    sqprecisionY = precisionY^2

    # Fill block
    @showprogress for i=1:Ne
        Xi  = X[:, i]

        # Compute Jacobian of the forward model
        J = Jacobian(h, Xi, Ny, Nx)

        # Compute Hessian of the forward model
        H = Hessian(h, Xi, Ny, Nx)

        # Compute Hessian of the log pdf of the transport map
        Hlog = hess_x_log_pdf(M, X[:,i:i])

        # Compute Hessian of the log pdf of the transport map
        # Fill Ωxy block
        @inbounds for k=1:Ny
            for j=1:Nx
                Omega[Ny+j, k] += (J[k,j]*L.L.diag[k]*L.L.diag[Ny+j])^2
            end
        end

        # Fill Ωxx
        @inbounds for k=1:Nx
            for j=1:Nx
                Jj = view(J,:,j)
                Jk = view(J,:,k)
                Hjk = view(H,:,j,k)
                # @show L.L.diag[Ny+k]*L.L.diag[Ny+j]*precisionY*dot(Jj, Jk)
                # @show Hlog[1,j,k]
                # @show (L.L.diag[Ny+k]*L.L.diag[Ny+j])^2*dot(Hjk, Hjk)*precisionY
                Omega[Ny+j, Ny+k] += (-L.L.diag[Ny+k]*L.L.diag[Ny+j]*precisionY*dot(Jj, Jk) + Hlog[1,j,k])^2 +
                                  (L.L.diag[Ny+k]*L.L.diag[Ny+j])^2*dot(Hjk, Hjk)*precisionY
            end
        end
    end

    view(Omega, Ny+1:Ny+Nx, 1:Ny) .*= sqprecisionY
    view(Omega, 1:Ny, Ny+1:Ny+Nx) .= view(Omega, Ny+1:Ny+Nx, 1:Ny)'

    rmul!(Omega, 1.0/Ne)

    # Fill Ωyy with the identity
    @inbounds for i=1:Ny
        Omega[i,i] = 1.0
    end

    return Omega
end

function omega_unscaled(h, M, Y, X, Ny, Nx, ϵY, L::LinearTransform)
    Ne = size(X, 2)
    Ω = zeros(Nx+Ny, Nx+Ny)

    precisionY = (1/ϵY^2)
    sqprecisionY = precisionY^2

    # Fill block
    @showprogress for i=1:Ne
        Xi  = X[:, i]

        # Compute Jacobian of the forward model
        J = Jacobian(h, Xi, Ny, Nx)

        # Compute Hessian of the forward model
        H = Hessian(h, Xi, Ny, Nx)

        # Compute Hessian of the log pdf of the transport map
        Hlog = hess_x_log_pdf(M, X[:,i:i])

        # Compute Hessian of the log pdf of the transport map
        # Fill Ωxy block
        @inbounds for k=1:Ny
            for j=1:Nx
                Ω[Ny+j, k] += J[k,j]^2
            end
        end

        # Fill Ωxx
        @inbounds for k=1:Nx
            for j=1:Nx
                Jj = view(J,:,j)
                Jk = view(J,:,k)
                Hjk = view(H,j,k)
                Ω[Ny+j, Ny+k] += (-precisionY*dot(Jj, Jk) + 1/(L.L.diag[Ny+k]*L.L.diag[Ny+j])*Hlog[1,j,k])^2 +
                                 dot(Hjk, Hjk)*precisionY
            end
        end
    end

    view(Ω, Ny+1:Ny+Nx, 1:Ny) .*= sqprecisionY
    view(Ω, 1:Ny, Ny+1:Ny+Nx) .= view(Ω, Ny+1:Ny+Nx, 1:Ny)'

    rmul!(Ω, 1/Ne)

    # Fill Ωyy with the identity
    @inbounds for i=1:Ny
        Ω[i,i] = 1.0
    end

    return Ω
end


using DynamicalSystemsBase, Test
using LinearAlgebra: norm
@testset "Projected integrator" begin

@testset "Lorenz" begin
    # Set β to 10, we have a stable fixed point
    ds = Systems.lorenz([0.0, 10.0, 1.0]; σ = 10.0, ρ = 28.0, β = 10)
    psys = projected_integrator(ds, 1:2, [0.0])
    @test get_state(psys) == get_state(ds)[1:2]

    reinit!(psys, [1.0, 1.0])
    for j in 1:5
      step!(psys, 1.0)
    end
    y = get_state(psys)
    @test length(y) == 2
    @test abs(sum(y - [16.43, 16.43])) < 0.01

    # Test complete state function
    complete = (y) -> [y[1], y[2], y[2] + 1]
    psys = projected_integrator(ds, 1:2, complete)
    reinit!(psys, [2.0, 1.0])
    y = get_state(psys)
    @test y[1] == 2
    @test length(y) == 2

    for j in 1:5
      step!(psys, 1.0)
    end
    u = get_state(psys)
    @test abs(sum(u - [16.43, 16.43])) < 0.01

    # Test projection function on a sphere of unit radius
    projection(u) = u/norm(u)
    complete = y -> 10y
    psys = projected_integrator(ds, projection, complete)
    @test norm(get_state(psys)) == 1
    reinit!(psys, ones(3))
    @test psys.integ.u[1] == 10

    for j in 1:5
      step!(psys, 1.0)
    end
    u = get_state(psys)
    @test abs(sum(u[1:2] - [0.461, 0.461])) < 0.01

    # Trajectory, continuous
    psys = projected_integrator(ds, projection, complete)
    tr = trajectory(psys, 10)
    @test dimension(tr) == 3
    @test norm(tr[end]) == 1

    # Trajectory, discrete
    ds = Systems.towel()
    psys = projected_integrator(ds, [1,2], [0.0])
    @test get_state(psys) == [0.085, -0.121]
    @test dimension(psys) == 2
    tr = trajectory(psys, 10)
    @test dimension(tr) == 2
end

end


b=3;a=@equ a=$b #:a ≖ 3
@test a==Equation(:a,3)
P=@equ P=Ten($(map(x->pi^x,1:3)),i)
@test floor(P.rhs&@equ(i=1))==3
arr=@equs(e=$ℯ, pi=$pi, M=$(rand(3,2)))
@test arr[1].rhs == ℯ
@test arr[2].rhs == pi

@test Ten(:I,[:i,:i])&@equ(I=[1 0;0 1])==2
@test Ten(:A,[:i,:i])&@equ(A=[:a 0;0 :b])==:a+:b
@test Ten(:A,:i)*Ten(:B,:j)&@equs(A=[1,2,3],B=[3,2,1], j=i)==10
@test Ten(:A,[:j,:i,:i])*Ten(:B,:j)&@equs(A=ones(3,3,3), B=[1,2,3])==18
@test (Alt([:i,:j,:k])*Ten([:a1,:a2,:a3],:j)*Ten([:b1,:b2,:b3],:k)&@equ i=1 )==simplify(:a2*:b3+-1.0*:a3*:b2)

x=@equ x=a*b^sqrt(y)+c/d
@test (x&@equs(a=3, b=2, y=9, c=8, d=4)).rhs==3*2^sqrt(9)+8/4

tri=@equ c^2=a^2+b^2
@test sqrt(tri)==simplify(:c ≖ Equations.Sqrt(:a*:a+:b*:b))

energy=@equ E=m*c^2
c=@equ c=299792458
m=@equ m=3*n
n=@equ n=9
@test (energy&c&m&n).rhs == 3*9*299792458^2

r=(Der(:x^:n,:x)-Der(-0.1*:x^:m,:x)+1/:a*Der(:a*sqrt(:x),:x))&relations["Der"]
@test r==simplify(:n*Equations.Pow(:x,:n+-1)+0.1*:m*Equations.Pow(:x,:m+-1)+0.5*Equations.Pow(:x,-0.5))

relation=@equ Log(:a,:a)=1
@test Log(:e)&relation ==1
@test Log(9,9)&relation ==1
@test Log(:x+:y,:x+:y)&relation ==1

eq=3*:x^2-5*:x+1.5 ≖ 0
meq=eq&quadratic
@test evaluate(eq.lhs,Dict(:x=>meq[1].rhs))==0

rel=@equ Oneable(a)*x*z=y
@test :q*:r&rel ==:y

eq=3*:x^2-5*:x+1.5 ≖ 0
meq=eq&quadratic
@assert evaluate(eq.lhs,Dict(:x=>meq[1].rhs))==0
f1(ex::Expression)=ex[1][1]
f2(fac::Factor)=3*fac
@test (:a+:b)&[f1,f2]==3*:a
f3(eq::Equation)=eq'
f4(eq::Equation)=sqrt(eq)
@test (@equ(a=b^2)&[f3,f4]).lhs==:b

l=U(:l,:meter);t=U(:t,:second);v=l/t
@test v==U(:l*Equations.╱(:t),:meter*Equations.╱(:second))

rule=Der(:a*:x,:x)≖:a #equivalent to Equation(Der(:a*:x,:x),:a)
ex=Der(3*:x,:x)
m=matches(ex,rule)[1] #equivalent to ex&rule
@test m==ex&rule==3

eq=Equation(:x*:z+:y)
@test eq.rhs==0
@test !isempty(matches(eq))
eq=Equation(:x^2,9)
@test matches(eq,Sqrt)[1].rhs==3

meq=matches(:x^2+:a*:x≖0,Div)[1]
try
	evaluate(meq,Dict(:x=>0))
	@test false
catch er
	@test true
end


using ModelingToolkit, LinearAlgebra, DiffEqJump, Test, Distributions
MT = ModelingToolkit

@parameters t k[1:20]
@variables A(t) B(t) C(t) D(t)
rxs = [Reaction(k[1], nothing, [A]),            # 0 -> A
       Reaction(k[2], [B], nothing),            # B -> 0
       Reaction(k[3],[A],[C]),                  # A -> C
       Reaction(k[4], [C], [A,B]),              # C -> A + B
       Reaction(k[5], [C], [A], [1], [2]),      # C -> A + A
       Reaction(k[6], [A,B], [C]),              # A + B -> C
       Reaction(k[7], [B], [A], [2], [1]),      # 2B -> A
       Reaction(k[8], [A,B], [A,C]),            # A + B -> A + C
       Reaction(k[9], [A,B], [C,D]),            # A + B -> C + D
       Reaction(k[10], [A], [C,D], [2], [1,1]), # 2A -> C + D
       Reaction(k[11], [A], [A,B], [2], [1,1]), # 2A -> A + B
       Reaction(k[12], [A,B], nothing, [3,1], nothing),           # 3A+B -> 0
       Reaction(k[13], [A,B,C], [C,D], [1,3,4], [Poisson(2), 3]),          # A+3B+4C -> k_pois * C + 3D
       Reaction(k[14], nothing, [A], nothing, [Geometric(1 // 3)]),               # 0 -> k_geom * A
       Reaction(k[15]*A/(2+A), [A], nothing; only_use_rate=true), # A -> 0 with custom rate
       Reaction(k[16], [A], [B]; only_use_rate=true),             # A -> B with custom rate.
       Reaction(k[17]*A*exp(B), [C], [D], [2], [Bernoulli(1)]),   # 2C -> D with non constant rate.
       Reaction(k[18]*B, nothing, [B], nothing, [2]),             # 0 -> 2B with non constant rate.
       Reaction(k[19]*t, [A], [B]),                                # A -> B with non constant rate.
       Reaction(k[20]*t*A, [B,C], [D],[2,1],[2])                  # 2A +B -> 2C with non constant rate.
  ]
rs = ReactionSystem(rxs,t,[A,B,C,D],k)
odesys = convert(ODESystem,rs)
sdesys = convert(SDESystem,rs)

# test show
io = IOBuffer()
show(io, rs)
str = String(take!(io))
@test count(isequal('\n'), str) < 30

# hard coded ODE rhs
function oderhs(u,k,t)
       A = u[1]; B = u[2]; C = u[3]; D = u[4];
       du = zeros(eltype(u),4)
       du[1] = k[1] - k[3]*A + k[4]*C + 2*k[5]*C - k[6]*A*B + k[7]*B^2/2 - k[9]*A*B - k[10]*A^2 - k[11]*A^2/2 - k[13]*A*B^3*C^4/144 - 3*k[12]*A^3*B/6 + 2*k[14] - k[15]*A/(2+A) - k[16] - k[19]*t*A
       du[2] = -k[2]*B + k[4]*C - k[6]*A*B - k[7]*B^2 - k[8]*A*B - k[9]*A*B + k[11]*A^2/2 - 3*k[13]*A*B^3*C^4/144 - k[12]*A^3*B/6 + k[16] + 2*k[18]*B + k[19]*t*A - 2*k[20]*t*A*B^2*C
       du[3] = k[3]*A - k[4]*C - k[5]*C  + k[6]*A*B + k[8]*A*B + k[9]*A*B + k[10]*A^2/2 - 2*k[13]*A*B^3*C^4/144 - 2*k[17]*A*exp(B)*C^2/2 - k[20]*t*A*B^2*C
       du[4] = k[9]*A*B + k[10]*A^2/2 + 3*k[13]*A*B^3*C^4/144 + k[17]*A*exp(B)*C^2/2 + 2*k[20]*t*A*B^2*C
       du
end

# sde noise coefs
function sdenoise(u,k,t)
       A = u[1]; B = u[2]; C = u[3]; D = u[4];
       G = zeros(eltype(u),length(k),length(u))
       z = zero(eltype(u))

       G = [sqrt(k[1]) z  z z;
            z -sqrt(k[2]*B) z z;
            -sqrt(k[3]*A) z sqrt(k[3]*A) z;
            sqrt(k[4]*C) sqrt(k[4]*C) -sqrt(k[4]*C) z;
            2*sqrt(k[5]*C) z -sqrt(k[5]*C) z;
            -sqrt(k[6]*A*B) -sqrt(k[6]*A*B) sqrt(k[6]*A*B) z;
            sqrt(k[7]*B^2/2) -2*sqrt(k[7]*B^2/2) z z;
            z -sqrt(k[8]*A*B) sqrt(k[8]*A*B) z;
            -sqrt(k[9]*A*B) -sqrt(k[9]*A*B) sqrt(k[9]*A*B) sqrt(k[9]*A*B);
            -2*sqrt(k[10]*A^2/2) z sqrt(k[10]*A^2/2) sqrt(k[10]*A^2/2);
            -sqrt(k[11]*A^2/2) sqrt(k[11]*A^2/2) z z;
            -3*sqrt(k[12]*A^3*B/6) -sqrt(k[12]*A^3*B/6) z z;
            -sqrt(k[13]*A*B^3*C^4/144) -3*sqrt(k[13]*A*B^3*C^4/144) -2*sqrt(k[13]*A*B^3*C^4/144) 3*sqrt(k[13]*A*B^3*C^4/144);
            2*sqrt(k[14]) z z z;
            -sqrt(k[15]*A/(2+A)) z z z;
            -sqrt(k[16]) sqrt(k[16]) z z;
            z z -2*sqrt(k[17]*A*exp(B)*C^2/2) sqrt(k[17]*A*exp(B)*C^2/2);
            z 2*sqrt(k[18]*B) z z;
            -sqrt(k[19]*t*A) sqrt(k[19]*t*A) z z;
            z -2*sqrt(k[20]*t*A*B^2*C) -sqrt(k[20]*t*A*B^2*C) +2*sqrt(k[20]*t*A*B^2*C);
            # burst variances
            z z sqrt(k[13]*A*B^3*C^4/144) * sqrt(2.) z;
            sqrt(k[14])*sqrt(6.) z z z;
            z z z z]';

       return G
end

# test by evaluating drift terms
p  = rand(length(k))
u  = rand(length(k))
t  = 0.
du = oderhs(u,p,t)
odesys = convert(ODESystem,rs)
of = ODEFunction{false}(odesys, states(rs), parameters(rs))
du2 = of.f(u,p,t)
@test norm(du-du2) < 100*eps()

# test by evaluating drift and diffusion terms
p  = rand(length(k))
u  = rand(length(k))
t  = 0.
du = oderhs(u,p,t)
G  = sdenoise(u,p,t)
sdesys = convert(SDESystem,rs)
sf = SDEFunction{false}(sdesys, states(rs), parameters(rs))
du2 = sf.f(u,p,t)
@test norm(du-du2) < 100*eps()
G2 = sf.g(u,p,t)
@test norm(G-G2) < 100*eps()

# test conversion to NonlinearSystem
ns = convert(NonlinearSystem,rs)
fnl = eval(generate_function(ns)[2])
dunl = similar(du)
fnl(dunl,u,p)
@test norm(du-dunl) < 100*eps()

# tests the noise_scaling argument.
p  = rand(length(k)+1)
u  = rand(length(k))
t  = 0.
G  = p[21]*sdenoise(u,p,t)
@variables η
sdesys_noise_scaling = convert(SDESystem,rs;noise_scaling=η)
sf = SDEFunction{false}(sdesys_noise_scaling, states(rs), parameters(sdesys_noise_scaling))
G2 = sf.g(u,p,t)
@test norm(G-G2) < 100*eps()

# tests the noise_scaling vector argument.
p  = rand(length(k)+3)
u  = rand(length(k))
t  = 0.
G  = vcat(fill(p[21],8),fill(p[22],5),fill(p[23],7),[p[22],p[23],p[23]])' .* sdenoise(u,p,t)
@variables η[1:3]
sdesys_noise_scaling = convert(SDESystem,rs;noise_scaling=vcat(fill(η[1],8),fill(η[2],5),fill(η[3],7)))
sf = SDEFunction{false}(sdesys_noise_scaling, states(rs), parameters(sdesys_noise_scaling))
G2 = sf.g(u,p,t)
@test norm(G-G2) < 100*eps()

# tests using previous parameter for noise scaling
p  = rand(length(k)+3)
p[21], p[22], p[23] = p[13], p[14], p[17]
u  = rand(length(k))
t  = 0.
G  = [p p p p]' .* sdenoise(u,p,t)
sdesys_noise_scaling = convert(SDESystem,rs;noise_scaling=vcat(k))
sf = SDEFunction{false}(sdesys_noise_scaling, states(rs), parameters(sdesys_noise_scaling))
G2 = sf.g(u,p,t)
@test norm(G-G2) < 100*eps()

# test with JumpSystem
js = convert(JumpSystem, rs)

midxs = 1:12
cidxs = 13:18
vidxs = 19:20
@test all(map(i -> typeof(equations(js)[i]) <: DiffEqJump.MassActionJump, midxs))
@test all(map(i -> typeof(equations(js)[i]) <: DiffEqJump.ConstantRateJump, cidxs))
@test all(map(i -> typeof(equations(js)[i]) <: DiffEqJump.VariableRateJump, vidxs))

pars = rand(length(k)); u0 = rand(1:10,4); ttt = rand();
jumps = Vector{Union{ConstantRateJump, MassActionJump, VariableRateJump}}(undef,length(rxs))

jumps[1] = MassActionJump(pars[1], Vector{Pair{Int,Int}}(), [1 => 1]);
jumps[2] = MassActionJump(pars[2], [2 => 1], [2 => -1]);
jumps[3] = MassActionJump(pars[3], [1 => 1], [1 => -1, 3 => 1]);
jumps[4] = MassActionJump(pars[4], [3 => 1], [1 => 1, 2 => 1, 3 => -1]);
jumps[5] = MassActionJump(pars[5], [3 => 1], [1 => 2, 3 => -1]);
jumps[6] = MassActionJump(pars[6], [1 => 1, 2 => 1], [1 => -1, 2 => -1, 3 => 1]);
jumps[7] = MassActionJump(pars[7], [2 => 2], [1 => 1, 2 => -2]);
jumps[8] = MassActionJump(pars[8], [1 => 1, 2 => 1], [2 => -1, 3 => 1]);
jumps[9] = MassActionJump(pars[9], [1 => 1, 2 => 1], [1 => -1, 2 => -1, 3 => 1, 4 => 1]);
jumps[10] = MassActionJump(pars[10], [1 => 2], [1 => -2, 3 => 1, 4 => 1]);
jumps[11] = MassActionJump(pars[11], [1 => 2], [1 => -1, 2 => 1]);
jumps[12] = MassActionJump(pars[12], [1 => 3, 2 => 1], [1 => -3, 2 => -1]);

jumps[13] = ConstantRateJump((u,p,t) -> p[13]*u[1]*binomial(u[2],3)*binomial(u[3],4),
                              integrator -> (integrator.u[1] -= 1; integrator.u[2] -= 3; integrator.u[3] += 2 - 4; integrator.u[4]+=3));
jumps[14] = ConstantRateJump((u,p,t) -> p[14], integrator -> (integrator.u[1] += 1));

jumps[15] = ConstantRateJump((u,p,t) -> p[15]*u[1]/(2+u[1]), integrator -> (integrator.u[1] -= 1))
jumps[16] = ConstantRateJump((u,p,t) -> p[16], integrator -> (integrator.u[1] -= 1; integrator.u[2] += 1;))
jumps[17] = ConstantRateJump((u,p,t) -> p[17]*u[1]*exp(u[2])*binomial(u[3],2), integrator -> (integrator.u[3] -= 2; integrator.u[4] += 1))
jumps[18] = ConstantRateJump((u,p,t) -> p[18]*u[2], integrator -> (integrator.u[2] += 2))

jumps[19] = VariableRateJump((u,p,t) -> p[19]*u[1]*t, integrator -> (integrator.u[1] -= 1; integrator.u[2] += 1))
jumps[20] = VariableRateJump((u,p,t) -> p[20]*t*u[1]*binomial(u[2],2)*u[3], integrator -> (integrator.u[2] -= 2; integrator.u[3] -= 1; integrator.u[4] += 2))

statetoid = Dict(state => i for (i,state) in enumerate(states(js)))
parammap = map((x,y)->Pair(x,y),parameters(js),pars)
for i in midxs
  maj = MT.assemble_maj(equations(js)[i], statetoid, ModelingToolkit.substituter(parammap),eltype(pars))
  @test abs(jumps[i].scaled_rates - maj.scaled_rates) < 100*eps()
  @test jumps[i].reactant_stoch == maj.reactant_stoch
  @test jumps[i].net_stoch == maj.net_stoch
end
for i in cidxs
  crj = MT.assemble_crj(js, equations(js)[i], statetoid)
  @test isapprox(crj.rate(u0,p,ttt), jumps[i].rate(u0,p,ttt))
  fake_integrator1 = (u=zeros(4),p=p,t=0); fake_integrator2 = deepcopy(fake_integrator1);
  crj.affect!(fake_integrator1);
  jumps[i].affect!(fake_integrator2);
  # Won't have equality in general due to randomness in bursts
  #@test fake_integrator1 == fake_integrator2
end
for i in vidxs
  crj = MT.assemble_vrj(js, equations(js)[i], statetoid)
  @test isapprox(crj.rate(u0,p,ttt), jumps[i].rate(u0,p,ttt))
  fake_integrator1 = (u=zeros(4),p=p,t=0.); fake_integrator2 = deepcopy(fake_integrator1);
  crj.affect!(fake_integrator1); jumps[i].affect!(fake_integrator2);
  @test fake_integrator1 == fake_integrator2
end


function hartree_pot_bc(rho::Array{Float64,2}, Ls::Float64, YukawaK::Float64, epsil0::Float64)
# Calculate the Hartree potential and the energy in periodic boundary potential
# this function is called from the one in the Ham. This is anotated to give the jit
# compiler more information

    Ns_glb = length(rho);
    kx = 2*pi* vcat(collect(0:Ns_glb/2),
                    collect(-Ns_glb/2+1:-1))./Ls ;
    # Use Yukawa rather than the bare Coulomb potential.  The method does
    # not depend critically on the fact that the Coulomb kernel is to be
    # used.
    invkx2 = [0;1 ./(kx[2:end].^2 .+ YukawaK^2)];
    # invkx2 = [1/YukawaK^2;1./(kx(2:end).^2+YukawaK^2)];
    # The 1/YukawaK^2 factor is not important for neutral defect calculation, but
    # might be important for the charged defect calculation.
    Vtemp = 4*pi*invkx2.*fft(rho)  # we can decorate this one
    V  = real(ifft(Vtemp));
 return V / epsil0;
end

############################################################################
##  Trying different optimizations

function hartree_pot_bc_opt(rho::Array{Float64,2}, Ls::Float64,
                            YukawaK::Float64, epsil0::Float64)
# Calculate the Hartree potential and the energy in periodic boundary potential
# this function is called from the one in the Ham. This is anotated to give the jit
# compiler more information
    rhoFourier = rfft(rho)
    inv_yukawa_fourier_mult!(rhoFourier,Ls,YukawaK)
    V  = irfft(rhoFourier, size(rho,1));
    return V / epsil0;
end

@inline function inv_yukawa_fourier_mult!(R::Array{ComplexF64,2}, Ls::Float64,
                                          YukawaK::Float64 )
    c = (2 * pi / Ls)^2
    c1 = 4*pi
    R[1] = 0
    @inbounds @simd for ii = 2:size(R)[1]
        R[ii] = c1*R[ii]/((ii-1)^2*c + YukawaK^2);
    end
end

# using 1D vector in general should produce faster resutls

function hartree_pot_bc_opt_vec(rho::Vector{Float64}, Ls::Float64,
                                YukawaK::Float64, epsil0::Float64)
# Calculate the Hartree potential and the energy in periodic boundary potential
# this function is called from the one in the Ham. This is anotated to give the jit
# compiler more information
    rhoFourier = rfft(rho)
    inv_yukawa_fourier_mult_vec!(rhoFourier,Ls,YukawaK)
    V  = irfft(rhoFourier, size(rho,1))/ epsil0;
    return V;
end

@inline function inv_yukawa_fourier_mult_vec!(R::Vector{ComplexF64}, Ls::Float64,
                                              YukawaK::Float64 )
    c = (2 * pi / Ls)^2
    c1 = 4*pi
    R[1] = 0
    @inbounds @simd for ii = 2:size(R)[1]
        R[ii] = c1*R[ii]/((ii-1)^2*c + YukawaK^2);
    end
end


# 1:2*length(layerconfig) -> encoder lstm (2k-1 for weight and 2k for bias in k'th encoder layer)
# 2*length(layerconfig)+1:4*length(layerconfig) -> decoder lstm (2k-1+2length for weight and 2k+2length for bias in kth decoder layer)
function initparams(atype, layerconfig, encembed, decembed, invocab, outvocab, winit)
    len = length(layerconfig)
    parameters = Array(Any, 4*length(layerconfig) + 4)

    # embedding and prediction layers
    parameters[end] = winit * randn(outvocab, decembed) # decoder embedding
    parameters[end-1] = winit * randn(invocab, encembed) # encoder embedding
    parameters[end-2] = winit * randn(layerconfig[end], outvocab) # final layer weight
    parameters[end-3] = zeros(1, outvocab) # final layer bias

    # encoder lstm
    input = encembed
    for k=1:len
        parameters[2k-1] = winit * randn(input+layerconfig[k], 4layerconfig[k])
        parameters[2k] = zeros(1, 4layerconfig[k])
        parameters[2k][1:layerconfig[k]] = 1 # forget gate bias
        input = layerconfig[k]
    end

    # decoder lstm
    input2 = decembed
    for k=1:len
        parameters[2k-1+2len] = winit * randn(input2+layerconfig[k], 4*layerconfig[k])
        parameters[2k+2len] = zeros(1, 4layerconfig[k])
        parameters[2k+2len][1:layerconfig[k]] = 1 # forget gate bias
        input2 = layerconfig[k]
    end
    return map(p->convert(atype, p), parameters)
end

# 1:2length(layerconfig) -> encoder states (k is cell and k-1 is hidden for kth layer)
# 2length(layerconfig)+1:4length(layerconfig) -> decoder states
# first hidden layer of the decoder left nothing intentionally
function initstates(atype, layerconfig, batchsize)
    len = length(layerconfig)
    states = Array(Any, 4len)

    # initialize encoder states
    for k=1:len
        states[2k-1] = zeros(batchsize, layerconfig[k])
        states[2k] = zeros(batchsize, layerconfig[k])
    end

    # initialize decoder states
    for k=1:len
        states[2k-1+2len] = zeros(batchsize, layerconfig[k])
        states[2k + 2len] = zeros(batchsize, layerconfig[k])
    end
    states = map(s->convert(atype, s), states)
    states[2len+1] = nothing
    return states
end

function lstm(weight, bias, hidden, cell, input)
    gates   = hcat(input,hidden) * weight .+ bias
    hsize   = size(hidden,2)
    forget  = sigm(gates[:,1:hsize])
    ingate  = sigm(gates[:,1+hsize:2hsize])
    outgate = sigm(gates[:,1+2hsize:3hsize])
    change  = tanh(gates[:,1+3hsize:end])
    cell    = cell .* forget + ingate .* change
    hidden  = outgate .* tanh(cell)
    return (hidden,cell)
end

# multi layer lstm forward function
function forward(parameters, states, input)
    x = input
    for i=1:2:length(states)
        (states[i], states[i+1]) = lstm(parameters[i], parameters[i+1], states[i], states[i+1], x)
        x = states[i]
    end
    return x
end

function loss(parameters, states, inputseq, outseq)
    total = 0.0
    count = 0.0
    atype = typeof(AutoGrad.getval(parameters[1]))
    hlayers = length(states)/4
    hlayers = convert(Int, hlayers)

    # go through the encoder
    enstates = states[1:2hlayers]
    final_hidden = nothing
    for i=1:length(inputseq)
        input = convert(atype, inputseq[i])
        x = input * parameters[end-1]
        final_hidden = forward(parameters[1:2hlayers], enstates, x)
    end

    # initialize the hidden state of the decoder, decstates[1] is the initial hidden of decoder lstm
    decstates = states[2hlayers+1:4hlayers]
    decstates[1] = final_hidden

    # go through the decoder
    for i=1:length(outseq)-1
        input = convert(atype, outseq[i])
        x = input * parameters[end]
        hidden = forward(parameters[2hlayers+1:4hlayers], decstates, x)
        ypred = hidden * parameters[end-2] .+ parameters[end-3]
        ynorm = logp(ypred, 2)
        ygold = convert(atype, outseq[i+1])
        count += size(ygold, 1)
        total += sum(ygold .* ynorm)
    end
    return - total / count
end

lossgradient = grad(loss)


using HyperbolicPlane, Plots

function circle_example()
    P = HPoint(1,pi/6)
    C = HCircle(P,1.5)
    X,Y,Z = points_on_circle(C)
    T = X+Y+Z
    set_color(T,:red)
    set_radius(P,2)
    L1 = bisector(X,Y)
    L2 = bisector(X,Z)
    L3 = bisector(Y,Z)
    set_color(L1,:green)
    set_color(L2,:green)
    set_color(L3,:green)
    plot()
    draw(C,T,L1,L2,L3,P,HPlane())
    finish()
end

function rand_circle_example()
    X,Y,Z = [ RandomHPoint() for j=1:3 ]
    C = HCircle()
    try
        C = HCircle(X,Y,Z)
    catch
        println("Sorry: The three points we generated are not contained in a common circle")
        println("Please try again")
        return
    end
    set_thickness(C,2)
    P = get_center(C)

    T = X+Y+Z
    set_color(T,:red)
    set_radius(P,2)
    L1 = bisector(X,Y)
    L2 = bisector(X,Z)
    L3 = bisector(Y,Z)
    set_color(L1,:green)
    set_color(L2,:green)
    set_color(L3,:green)

    plot()
    draw(C,T,L1,L2,L3,P,HPlane())
    finish()
end


reload("gencond.jl")

function genba( kappa, m, n, kind )
  kappa = log10(kappa)
  sqrt_kappa = sqrt(kappa)
  V = qr(randn(m,m))[1]
  D = logspace(0, -kappa, m)
  B = scale(V, D)*V'
  W = qr(randn(n,n))[1]
  if kind == 1
    A = scale(V[:,m-n+1:m], logspace(0,-sqrt_kappa,n))*W'
  elseif kind == 2
    A = scale(V[:,1:n], logspace(0,-sqrt_kappa,n))*W'
  elseif kind == 3
    A = scale([V[:,m-cld(n,2)+1:m] V[:,1:fld(n,2)]],logspace(0,-sqrt_kappa,n))*W'
  elseif kind == 4
    A = gencond(10^sqrt_kappa, m, n)
  else
    U = [V[:,m-cld(n,2)+1:m] V[:,1:fld(n,2)]]
    S = [D[m-cld(n,2)+1:m]; D[1:fld(n,2)]]
    A = scale(U,1./sqrt(S))*W'
  end
  return sparse(triu(B)),A
end


function kendallsTau{T<:AbstractFloat}(Y::Matrix{T})

  n, p = size(Y)
  covM = zeros(T, (p, p))
  denom = 0

  for a=1:p, b=a:p
    if a == b
      covM[a,b] = 1.
      continue
    end
    t = zero(T)
    for i=1:n, j=i+1:n
      @inbounds t += sign((Y[i, a] - Y[j, a])*(Y[i, b] - Y[j, b]))
    end
    t = 2. * t / ((n-1)*n)
    t = sin(π/2*T)
    covM[a,b] = t
    covM[b,a] = t
  end
  covM
end

function nonparanormalCorrelation{T<:AbstractFloat}(Y::Matrix{T})
  n, p = size(Y)
  stdNormal = Normal()
  δn = 1. / (4 * n^(1./4.) * sqrt(π*log(n)))
  IX = zeros(Y)
  for j=1:p
    sortperm!(view(IX, :, j), view(Y, :, j))
  end
  scale!(IX, 1./n)
  for j=1:p, i=1:n
    if IX[i,j] < δn
      IX[i,j] < δn
    elseif IX[i,j] >= 1. - δn
      IX[i,j] = 1. - δn
    end
    IX[i,j] = quantile(stdNormal, IX[i,j])
  end
  return cor(IX)
end


@diffusion_process SIR begin
    :parameters
    (α, β, σ1, σ2) --> Float64

    :conjugate
    phi(t, u) --> (
        (0.0, 0.0),
        ((k - u[1] - u[2])*u[1], 0.0),
        (-u[1], u[1]),
        (0.0, 0.0),
        (0.0, 0.0),
        (0.0, 0.0),
    )
    nonhypo(x) --> x
    num_non_hypo --> 2

    :additional
    statespace -->BoundedStateSpace(((1,2), (0.0,0.0)), ((1,2), (1.0,1.0)))
end

function b(t, x, P::SIR)
    ℝ{2}(
        P.α*(1 - u[1] - u[2])*u[1] - P.β*u[1],
        P.β*u[1]
    )
end

function σ(t, x, P::SIR)
    @SMatrix Float64[
    -P.σ1*sq((1 - u[1] - u[2])*u[1])  -P.σ2*sq(u[1]);
    0.0   P.σ2*sq.(u[1])
]


using Base.Test
using DataFrames, NLreg

const sd1 = readtable(Pkg.dir("NLreg","data","sd1.csv.gz"))
sd1[:ID] = pool(sd1[:ID])

# nl = gnfit(NonlinearLS(Logsd1(:(CONC ~ TIME), sd1)),true)

pl = fit([LogTr] * BolusSD1(CONC ~ TIME, sd1),true)

@test_approx_eq coef(pl) [1.1429558452268844,-1.4036989045425874]

pl = fit(AsympReg([2.:7],[18.6,22.6,25.1,27.2,29.1,30.1]),true)

const bod = DataFrame(Time=[1.:5,7],Demand=[8.3,10.3,19.0,16.0,15.6,19.8])

pl = fit(AsympOrig(Demand ~ Time,bod),true)


using ArbitrageFreeSABR, Test

@testset "HaganSet" begin
    forward = 1.0; expiry = 1.0;
    alpha = 0.35; beta = 0.25; nu = 1.0; rho = -0.1;
    N = 500; timesteps = 5; nd= 4;
    maturity = SABRMaturity(alpha,beta,rho,nu,forward,expiry,ArbitrageFreeSABRModel())
    density = makeTransformedDensityLawsonSwayne(maturity, N, timesteps, nd)
    @test isapprox(priceTransformedDensity(density, true, forward, ArbitrageFreeSABR.midpoint), 0.149701955629, atol=1e-10)
    @test isapprox(priceTransformedDensity(density, true, forward, ArbitrageFreeSABR.linear), 0.14970195574320477, atol=1e-10)
end

@testset "ImpliedVolatilitySet" begin
    price = 271.43234885190117
    strike = 275.0
    f = 356.73063159822254
    tte = 1.5917808219178082
    isCall = false
    tolerance = 1e-8
    vol = impliedVolatilityLiSORTS(price, isCall, strike, f, tte, 1.0, 0.0, tolerance,64)
    @test isapprox(blackScholesFormula(isCall, strike, f, vol*vol*tte, 1.0, 1.0), price, atol=tolerance)
end

@testset "BachelierVolatilitySet" begin
    price = 271.43234885190117
    strike = 275.0
    f = 356.73063159822254
    tte = 1.5917808219178082
    isCall = false
    tolerance = 1e-13
    vol = bachelierVolatility(price, isCall, strike, f, tte,1.0)
    call = f - strike + price
    tv = call - max(f-strike,0.0)
    y = tv/abs(f-strike)*4*sqrt(pi)
    @test isapprox(bachelierFormula(isCall, strike, f, vol*vol*tte, 1.0), price, atol=tolerance)
end


using DataFrames, CSV, Tables
using StanSample
using Test

df = CSV.read(joinpath(@__DIR__, "..", "..", "data", "WaffleDivorce.csv"), DataFrame);

stan5_1_t = "
data {
 int < lower = 1 > N; // Sample size
 vector[N] D; // Outcome
 vector[N] A; // Predictor
}

parameters {
 real a; // Intercept
 real bA; // Slope (regression coefficients)
 real < lower = 0 > sigma;    // Error SD
}

transformed parameters {
    vector[N] mu;
    mu = a + + bA * A;
}

model {
  a ~ normal( 0 , 0.2 );
  bA ~ normal( 0 , 0.5 );
  sigma ~ exponential( 1 );
  D ~ student_t( 2, mu , sigma );
}
generated quantities{
    vector[N] log_lik;
    for (i in 1:N)
        log_lik[i] = student_t_lpdf(D[i] | 2, mu[i], sigma);
}
";

begin
    data = (N=size(df, 1), D=df.Divorce, A=df.MedianAgeMarriage,
        M=df.Marriage)
    m5_1s_t = SampleModel("m5.1s_t", stan5_1_t)
    rc5_1s_t = stan_sample(m5_1s_t; data)

    if success(rc5_1s_t)
        post5_1s_t_df = read_samples(m5_1s_t, :dataframe)
    end
end

if success(rc5_1s_t)
    nt5_1s_t = read_samples(m5_1s_t, :namedtuple)
    df5_1s_t = read_samples(m5_1s_t, :dataframe)
    log_lik_1_t = nt5_1s_t.log_lik'
    a5_1s_t, cnames = read_samples(m5_1s_t, :array; return_parameters=true);
end

st5_1_t = convert_a3d(a5_1s_t, cnames, Val(:table));
@test names(st5_1_t)[end] == Symbol("log_lik.50")
@test size(DataFrame(st5_1_t)) == (4000, 103)

mu = matrix(st5_1_t, "mu")
@test size(mu) == (4000, 50)


# subtypes should implement:
# transform_op(x, reg) -> LinearOperator
# projection_op(supp, reg) -> LinearOperator
# proximal!(x, thresh, reg) -> vector
# lipschitz_scale(reg) -> value
# support(x, reg) -> vector
# TODO: document above properly
abstract type Regularizer end

support(x, thresh, reg::Regularizer) = abs.(transform_op(reg) * x) .> thresh
proximal(x, thresh, reg::Regularizer) = proximal!(copy(x), thresh, reg) # fallback for out-of-place proximal

## L1 regularizer
struct L1 <: Regularizer
    size::Int
end

support(x, reg::L1) = support(x, 1e-3, reg)
lipschitz_scale(::L1) = 2.

# proximal

softthresh(val, thresh) = max(abs(val) - thresh, 0) .* sign(val)
function proximal!(x, thresh, reg::L1)
    x .= softthresh.(x, thresh)
    x
end

# transform

transform_op(reg::L1) = Identity(reg.size)

# projection

struct L1Project <: LinearMap{Float64}
    p::Int
    suppinv::AbstractVector{<:Int}
end

Base.size(P::L1Project) = (sum(P.suppinv .> 0), P.p)
L1ProjectTranspose = LinearMaps.TransposeMap{<:Any, <:L1Project}

function L1Project(p::Int, supp::AbstractVector{Bool})
    suppinv = zeros(Int, p)
    index = 1
    for i in 1:p
        if supp[i]
            suppinv[i] = index
            index += 1
        end
    end
    L1Project(p, suppinv)
end

Base.:(*)(P::L1Project, x::AbstractVector) = x[P.suppinv .> 0]
Base.:(*)(Pt::L1ProjectTranspose, y::AbstractVector) = [Pt.lmap.suppinv[i] > 0 ? y[Pt.lmap.suppinv[i]] : 0.0 for i in 1:Pt.lmap.p] # TODO: use zero element instead?
mul!(uflat::AbstractVecOrMat, P::L1Project, yflat::AbstractVector) = (uflat[:] = P * yflat)
mul!(uflat::AbstractVecOrMat, Pt::L1ProjectTranspose, yflat::AbstractVector) = (uflat[:] = Pt * yflat)

projection_op(supp, reg::L1) = L1Project(reg.size, supp)

## TV regularizer

# TODO: type-genericness?
struct TVProximalWork
    Δ::AbstractArray{Float64}
    diff::AbstractArray{Float64}
end

struct TV <: Regularizer
    size::Tuple
    wrap::Bool
    work::TVProximalWork
end

TV(size::Tuple) = TV(size, true)
function TV(size::Tuple, wrap::Bool)
    Δ = Array{Float64}(undef, size)
    diff = Array{Float64}(undef, size)
    TV(size, wrap, TVProximalWork(Δ, diff))
end

support(x, reg::TV) = support(x, 0.005, reg)
lipschitz_scale(::TV) = 50

# proximal

function proximal!(x, thresh, reg::TV)
    x = reshape(x, reg.size)
    D = length(reg.size)
    Δ = reg.work.Δ
    diff = reg.work.diff
    Δ .= 0
    for i in 1:D
        xroll = roll(x, -1, i)
        diff .= (d -> min(2 * D * thresh, abs(d) / 2) * sign(d)).(xroll .- x)
        Δ .+= diff ./ (2 * D)
        Δ .-= roll(diff, 1, i) ./  (2 * D)
    end
    x .+= Δ
    reshape(x, prod(reg.size))
end

function proximal2(x, thresh, reg::TV)
    y = x[:]
    proxTV!(y, thresh, shape=reg.size, iterations=20)
    y
end

# transform

roll(x, shift, dim) = ShiftedArrays.circshift(x, (zeros(dim-1)..., shift, zeros(ndims(x)-dim)...))

struct Gradient <: LinearMap{Float64}
    size::Tuple
    dim::Int
    wrap::Bool
end

Base.size(D::Gradient) = (prod(D.size), prod(D.size))
GradientTranspose = LinearMaps.TransposeMap{<:Any, <:Gradient}
function Base.:(*)(D::Gradient, xflat::AbstractVector)
    x = reshape(xflat, D.size)
    y = x .- roll(x, -1, D.dim)
    !D.wrap && (y[end] = 0)
    y[:]
end
function Base.:(*)(Dt::GradientTranspose, xflat::AbstractVector)
    D = Dt.lmap
    x = reshape(xflat, D.size)
    y = x .- roll(x, 1, D.dim)
    !D.wrap && (y[1] = 0)
    y[:]
end

mul!(uflat::AbstractVecOrMat, D::Gradient, yflat::AbstractVector) = (uflat[:] = D * yflat)
mul!(uflat::AbstractVecOrMat, Dt::GradientTranspose, yflat::AbstractVector) = (uflat[:] = Dt * yflat)

function transform_op(reg::TV)
    ndims = length(reg.size)
    vcat([Gradient(reg.size, dim, reg.wrap) for dim in 1:ndims]...) # TODO: (style) should I avoid these sort of splatted arrays?
end

# projection

# TODO (relevant everywhere): proper typing for things like 2D arrays
# TODO: type generic code (no Float64's)
struct TVProject{N} <: LinearMap{Float64}
    size::NTuple{N, Int}
    regions::AbstractVector{AbstractVector{CartesianIndex{N}}}
end

Base.size(P::TVProject) = (length(P.regions), prod(P.size))
TVProjectTranspose = LinearMaps.TransposeMap{<:Any, <:TVProject}

function TVProject(size::Tuple, regions::AbstractVector{AbstractVector{CartesianIndex}})
    TVProject{length(size)}(size, regions) # TODO: check if this what I'm supposed to do
end

# TODO: type support properly
function TVProject(size::Tuple, supp::AbstractVector{<:Bool}, wrap::Bool)
    ndims = length(size)
    explored = fill(false, size)
    regions = AbstractVector{CartesianIndex{ndims}}[]
    supp = reshape(supp, (size..., ndims))
    for start_pt in CartesianIndices(size)
        explored[start_pt] && continue
        region = [start_pt]
        explored[start_pt] = true;
        s = Stack{CartesianIndex{ndims}}();
        push!(s, start_pt)
        while !isempty(s)
            pt = pop!(s)
            pttup = Tuple(pt)
            for i in 1:ndims
                if wrap || pt[i] < size[i]
                    pt_right = CartesianIndex(pttup[1:i-1]..., 1 + mod(pt[i], size[i]), pttup[i+1:ndims]...)
                    if !supp[pt, i] && !explored[pt_right]
                        explored[pt_right] = true
                        push!(s, pt_right)
                        push!(region, pt_right)
                    end
                end
                if wrap || pt[i] > 1
                    pt_left = CartesianIndex(pttup[1:i-1]..., 1 + mod(pt[i]-2, size[i]), pttup[i+1:ndims]...)
                    if !supp[pt_left, i] && !explored[pt_left]
                        explored[pt_left] = true
                        push!(s, pt_left)
                        push!(region, pt_left)
                    end
                end
            end
        end
        push!(regions, region)
    end
    TVProject(size, regions)
end

function Base.:(*)(P::TVProject, x::AbstractVector)
    x = reshape(x, P.size)
    region_norm(region) = sum(x[pt] for pt in region) / sqrt(length(region))
    map(region_norm, P.regions)
end

function region_vals(P::TVProject, x::AbstractVector)
    x = reshape(x, P.size)
    function region_vals(region)
        vals = [x[pt] for pt in region]
        sum(vals) / length(region), maximum(vals) - minimum(vals), length(vals)
    end
    map(region_vals, P.regions)
end

function Base.:(*)(Pt::TVProjectTranspose, y::AbstractVector)
    P = Pt.lmap
    x = zeros(P.size)
    for (val, region) in zip(y, P.regions)
        x[region] .= val / sqrt(length(region))
    end
    x[:]
end

projection_op(supp, reg::TV) = TVProject(reg.size, supp, reg.wrap)
mul!(uflat::AbstractVecOrMat, P::TVProject, yflat::AbstractVector) = (uflat[:] = P * yflat)
mul!(uflat::AbstractVecOrMat, Pt::TVProjectTranspose, yflat::AbstractVector) = (uflat[:] = Pt * yflat)


mutable struct MyCallable
    f
end
(a::MyCallable)(x) = a.f(x)
@testset "callables" begin
    @testset "univariate" begin
    # FIXME
    end
    @testset "multivariate" begin
        function rosenbrock(x::Vector)
           return (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2
        end

        a = MyCallable(rosenbrock)

        optimize(a, rand(2), Optim.Options())
    end
end


function test_l2_constraint(backend::Backend, T, eps)
  println("-- Testing L2 constraint on $(typeof(backend)){$T}...")
  # this simulates a convolutional filter and applies
  # the l2 constraint to it
  n_filters = 5
  coef = 0.2
  param = rand(T, 2,3,4,n_filters) - 0.5
  param_after = zeros(T, size(param))
  param_blob = make_blob(backend, param)

  cons = L2Cons(coef)
  constrain!(backend, cons, param_blob)
  copy!(param_after, param_blob)
  param_after = reshape(param_after, size(param))
  for f=1:n_filters
    norm2 = vecnorm(param_after[:, :, :, f])
    @test norm2 <= coef + eps
  end

  # this is the same as above but for fully connected weights
  n_input = 10
  n_out   = 12
  param = rand(T, n_input,n_out) - 0.5
  param_after = zeros(T, size(param))
  param_blob = make_blob(backend, param)

  cons = L2Cons(coef)
  constrain!(backend, cons, param_blob)
  copy!(param_after, param_blob)
  param_after = reshape(param_after, size(param))
  for f=1:n_out
    norm2 = vecnorm(param_after[:, f])
    @test norm2 <= coef + eps
  end

  # The case for bias
  len = 10
  param = rand(T, len) - 0.5
  param_after = similar(param)
  param_blob = make_blob(backend, param)
  cons = L2Cons(coef)
  constrain!(backend, cons, param_blob)
  copy!(param_after, param_blob)
  @test vecnorm(param_after) <= coef+eps
end

function test_l2_constraint(backend::Backend)
  test_l2_constraint(backend, Float32, 1e-5)
  test_l2_constraint(backend, Float64, 1e-5)
end

if test_cpu
  test_l2_constraint(backend_cpu)
end
if test_gpu
  test_l2_constraint(backend_gpu)
end


### Solve linear system by applying the implicit ULV factorization
#
# as seen in
# Chandrasekaran, S., Gu, M., & Pals, T. (2006). A Fast $ULV$ Decomposition Solver for Hierarchically Semiseparable Representations.
# SIAM Journal on Matrix Analysis and Applications, 28(3), 603–622. https://doi.org/10.1137/S0895479803436652
#
# Written by Boris Bonev, Nov. 2020

## function for direct solution using the implicit ULV factorization
function ulvfactsolve(hssA::HssMatrix{T}, b::Matrix{T}) where T
  if isleaf(hssA)
    return hssA.D\b
  else
    z = zeros(size(hssA,2), size(b,2))
    _, _, _, _, _, _, _, QV = _ulvfactsolve!(hssA, b, z, 0; rootnode=true)
    z = _ulvfactsolve_topdown!(QV, z)
    return z
  end
end

# core routine to reduce the rows and triangularize the diagonal block via QR/LQ decompositions and immediately apply them to b
function _ulvreduce!(D::Matrix{T}, U::Matrix{T}, V::Matrix{T}, b::Matrix{T}) where T
  T <: Complex ? adj = 'C' : adj = 'T'
  m, n = size(D)
  m == size(U,1) || throw(DimensionMismatch("D and U have different first dimensions"))
  n == size(V,1) || throw(DimensionMismatch("D and V have different second dimensions"))
  k = min(size(U, 2), m)
  nk = min(m-k,n)
  ind = 1:m-k
  cind = m-k+1:m
  # can't be compressed, exit early
  if k >= m
    # @warn "Encountered a full-rank block with k=$(k). Clustering might not yield best performance!"
    u = zeros(size(V,2), size(b,2))
    zloc = Matrix{T}(undef, 0, size(b,2))
    lqf = (Matrix{T}(undef, 0, n), Vector{T}(undef, 0))
  else
    # form QL decomposition of the row generators and apply it
    qlf = geqlf!(U);
    U = tril(U[end-k+1:end,:]) # k x k block
    ormql!('L', adj, qlf..., D) # transform the diagonal block
    ormql!('L', adj, qlf..., b) # transform the right-hand side
    # Form the LQ decomposition of the first m-k rows of D
    lqf = gelqf!(D[1:end-k,:])
    L1 = tril(lqf[1])
    L1 = L1[:,1:nk]
    L2 = ormlq!('R', adj, lqf..., D[end-k+1:end,:]) # update the bottom block of the diagonal block
    zloc = trsm('L', 'L', 'N', 'N', 1., L1, b[ind,:])
    b = b[cind, :] .- L2[:,1:nk] * zloc # remove contribution in the uncompressed parts
    V = ormlq!('L', 'N', lqf..., V) # compute the updated off-diagonal generators on the right
    u = V[ind,:]' * zloc # compute update vector to be passed on
    # pass on uncompressed parts of the problem
    D = L2[:, nk+1:end]
    V = V[nk+1:end, :]
  end
  return D, U, V, b, zloc, u, m-k, nk, lqf
end

# tree traversal that handles the hierarchical ULV factorization
function _ulvfactsolve!(hssA::HssMatrix{T}, b::Matrix{T}, z::Matrix{T}, co::Int; rootnode=false) where T
  if isleaf(hssA)
    cols = collect(co .+ (1:size(hssA,2)))
    D = copy(hssA.D); U = copy(hssA.U); V = copy(hssA.V)
    D, U, V, b, zloc, u, mk, nk, lqf =_ulvreduce!(D, U, V, b)
    z[cols[1:mk], :] = zloc
    QV = BinaryNode((cols, lqf...))
  else
    m1, n1 = hssA.sz1; m2, n2 = hssA.sz2
    b1 = b[1:m1, :]; b2 = b[m1+1:end, :] # performance could be further improved by getting rid of those allocations
    b1, u1, D1, U1, V1, cols1, nk1, QV1 = _ulvfactsolve!(hssA.A11, b1, z, co)
    b2, u2, D2, U2, V2, cols2, nk2, QV2 = _ulvfactsolve!(hssA.A22, b2, z, co+n1)

    # merge nodes to form new diagonal block
    b = [b1; b2] .- [U1*hssA.B12*u2; U2*hssA.B21*u1]
    D = [D1 U1*hssA.B12*V2'; U2*hssA.B21*V1' D2]
    cols = [cols1[nk1+1:end]; cols2[nk2+1:end]] # to re-adjust local numbering

    U = [U1*hssA.R1; U2*hssA.R2]
    V = [V1*hssA.W1; V2*hssA.W2]

    # now if this is the rootnode we should directly solve?
    if rootnode
      z[cols, :] = D\b
      u = Matrix{T}(undef,0,size(b,2))
      mk, nk = size(D)
      QV = BinaryNode{Tuple{Vector{Int}, Matrix{T}, Vector{T}}}()
    else
      D, U, V, b, zloc, u, mk, nk, lqf = _ulvreduce!(D, U, V, b)
      u .= u .+ hssA.W1'*u1 .+ hssA.W2'*u2
      z[cols[1:mk], :] = zloc
      QV = BinaryNode((cols, lqf...))
    end
    QV.left = QV1; QV.right = QV2
  end
  return b, u, D, U, V, cols, nk, QV
end

function _ulvfactsolve_topdown!(QV::BinaryNode{Tuple{Vector{Int}, Matrix{T}, Vector{T}}}, z::Matrix{T}) where T
  T <: Complex ? adj = 'C' : adj = 'T'
  if !isnothing(QV.data)
    (cols, Q, tau) = QV.data
    z[cols, :] = ormlq!('L', adj, Q, tau, z[cols, :])
  end
  if !isnothing(QV.left) z = _ulvfactsolve_topdown!(QV.left, z) end
  if !isnothing(QV.right) z = _ulvfactsolve_topdown!(QV.right, z) end
  return z
end


##### Weighted var & std

## var

function Base.varzm(v::RealArray, wv::WeightVec)
    n = length(v)
    length(wv) == n || throw(DimensionMismatch("Inconsistent array lengths."))
    w = values(wv)
    s = 0.0
    for i = 1:n
        @inbounds s += w[i] * abs2(v[i])
    end
    return s / sum(wv)
end

Base.varm(v::RealArray, m::Real, wv::WeightVec) = _moment2(v, m, wv)

function Base.var(v::RealArray, wv::WeightVec; mean=nothing)
    mean == 0 ? Base.varzm(v, wv) :
    mean == nothing ? varm(v, Base.mean(v, wv), wv) :
    varm(v, mean, wv)
end

## var along dim

Base.varzm!(R::AbstractArray, A::RealArray, wv::WeightVec, dim::Int) =
    scale!(_wsum_general!(R, Abs2Fun(), A, values(wv), dim, true), inv(sum(wv)))

Base.varm!(R::AbstractArray, A::RealArray, M::RealArray, wv::WeightVec, dim::Int) =
    scale!(_wsum_centralize!(R, Abs2Fun(), A, M, values(wv), dim, true), inv(sum(wv)))

function var!(R::AbstractArray, A::RealArray, wv::WeightVec, dim::Int; mean=nothing)
    if mean == 0
        Base.varzm!(R, A, wv, dim)
    elseif mean == nothing
        Base.varm!(R, A, Base.mean(A, wv, dim), wv, dim)
    else
        # check size of mean
        for i = 1:ndims(A)
            dA = size(A,i)
            dM = size(mean,i)
            if i == dim
                dM == 1 || throw(DimensionMismatch("Incorrect size of mean."))
            else
                dM == dA || throw(DimensionMismatch("Incorrect size of mean."))
            end
        end
        Base.varm!(R, A, mean, wv, dim)
    end
end

Base.varm(A::RealArray, M::RealArray, wv::WeightVec, dim::Int) =
    Base.varm!(Array(Float64, Base.reduced_dims(A, dim)), A, M, wv, dim)

Base.var(A::RealArray, wv::WeightVec, dim::Int; mean=nothing) =
    var!(Array(Float64, Base.reduced_dims(A, dim)), A, wv, dim; mean=mean)

## std

Base.stdm(v::RealArray, m::Real, wv::WeightVec) = sqrt(varm(v, m, wv))
Base.std(v::RealArray, wv::WeightVec; mean=nothing) = sqrt(var(v, wv; mean=mean))

Base.stdm(v::RealArray, m::RealArray, dim::Int) = Base.sqrt!(varm(v, m, dim))
Base.stdm(v::RealArray, m::RealArray, wv::WeightVec, dim::Int) = sqrt(varm(v, m, wv, dim))
Base.std(v::RealArray, wv::WeightVec, dim::Int; mean=nothing) = sqrt(var(v, wv, dim; mean=mean))

##### Fused statistics

mean_and_var(A::RealArray) = (m = mean(A); (m, varm(A, m)))
mean_and_std(A::RealArray) = (m = mean(A); (m, stdm(A, m)))

mean_and_var(A::RealArray, wv::WeightVec) = (m = mean(A, wv); (m, varm(A, m, wv)))
mean_and_std(A::RealArray, wv::WeightVec) = (m = mean(A, wv); (m, stdm(A, m, wv)))

mean_and_var(A::RealArray, dim::Int) = (m = mean(A, dim); (m, varm(A, m, dim)))
mean_and_std(A::RealArray, dim::Int) = (m = mean(A, dim); (m, stdm(A, m, dim)))

mean_and_var(A::RealArray, wv::WeightVec, dim::Int) = (m = mean(A, wv, dim); (m, varm(A, m, wv, dim)))
mean_and_std(A::RealArray, wv::WeightVec, dim::Int) = (m = mean(A, wv, dim); (m, stdm(A, m, wv, dim)))

##### General central moment

function _moment2(v::RealArray, m::Real)
    n = length(v)
    s = 0.0
    for i = 1:n
        @inbounds z = v[i] - m
        s += z * z
    end
    s / n
end

function _moment2(v::RealArray, m::Real, wv::WeightVec)
    n = length(v)
    s = 0.0
    w = values(wv)
    for i = 1:n
        @inbounds z = v[i] - m
        @inbounds s += (z * z) * w[i]
    end
    s / sum(wv)
end

function _moment3(v::RealArray, m::Real)
    n = length(v)
    s = 0.0
    for i = 1:n
        @inbounds z = v[i] - m
        s += z * z * z
    end
    s / n
end

function _moment3(v::RealArray, m::Real, wv::WeightVec)
    n = length(v)
    s = 0.0
    w = values(wv)
    for i = 1:n
        @inbounds z = v[i] - m
        @inbounds s += (z * z * z) * w[i]
    end
    s / sum(wv)
end

function _moment4(v::RealArray, m::Real)
    n = length(v)
    s = 0.0
    for i = 1:n
        @inbounds z = v[i] - m
        s += abs2(z * z)
    end
    s / n
end

function _moment4(v::RealArray, m::Real, wv::WeightVec)
    n = length(v)
    s = 0.0
    w = values(wv)
    for i = 1:n
        @inbounds z = v[i] - m
        @inbounds s += abs2(z * z) * w[i]
    end
    s / sum(wv)
end

function _momentk(v::RealArray, k::Int, m::Real)
    n = length(v)
    s = 0.0
    for i = 1:n
        @inbounds z = v[i] - m
        s += (z ^ k)
    end
    s / n
end

function _momentk(v::RealArray, k::Int, m::Real, wv::WeightVec)
    n = length(v)
    s = 0.0
    w = values(wv)
    for i = 1:n
        @inbounds z = v[i] - m
        @inbounds s += (z ^ k) * w[i]
    end
    s / sum(wv)
end

function moment(v::RealArray, k::Int, m::Real)
    k == 2 ? _moment2(v, m) :
    k == 3 ? _moment3(v, m) :
    k == 4 ? _moment4(v, m) :
    _momentk(v, k, m)
end

function moment(v::RealArray, k::Int, m::Real, wv::WeightVec)
    k == 2 ? _moment2(v, m, wv) :
    k == 3 ? _moment3(v, m, wv) :
    k == 4 ? _moment4(v, m, wv) :
    _momentk(v, k, m, wv)
end

moment(v::RealArray, k::Int) = moment(v, k, mean(v))
moment(v::RealArray, k::Int, wv::WeightVec) = moment(v, k, mean(v, wv), wv)

##### Skewness and Kurtosis

# Skewness
# This is Type 1 definition according to Joanes and Gill (1998)
function skewness(v::RealArray, m::Real)
    n = length(v)
    cm2 = 0.0   # empirical 2nd centered moment (variance)
    cm3 = 0.0   # empirical 3rd centered moment
    for i = 1:n
        @inbounds z = v[i] - m
        z2 = z * z

        cm2 += z2
        cm3 += z2 * z
    end
    cm3 /= n
    cm2 /= n
    return cm3 / sqrt(cm2 * cm2 * cm2)  # this is much faster than cm2^1.5
end

function skewness(v::RealArray, wv::WeightVec, m::Real)
    n = length(v)
    length(wv) == n || throw(DimensionMismatch("Inconsistent array lengths."))
    cm2 = 0.0   # empirical 2nd centered moment (variance)
    cm3 = 0.0   # empirical 3rd centered moment
    w = values(wv)

    @inbounds for i = 1:n
        x_i = v[i]
        w_i = w[i]
        z = x_i - m
        z2w = z * z * w_i
        cm2 += z2w
        cm3 += z2w * z
    end
    sw = sum(wv)
    cm3 /= sw
    cm2 /= sw
    return cm3 / sqrt(cm2 * cm2 * cm2)  # this is much faster than cm2^1.5
end

skewness(v::RealArray) = skewness(v, mean(v))
skewness(v::RealArray, wv::WeightVec) = skewness(v, wv, mean(v, wv))

# (excessive) Kurtosis
# This is Type 1 definition according to Joanes and Gill (1998)
function kurtosis(v::RealArray, m::Real)
    n = length(v)
    cm2 = 0.0  # empirical 2nd centered moment (variance)
    cm4 = 0.0  # empirical 4th centered moment
    for i = 1:n
        @inbounds z = v[i] - m
        z2 = z * z
        cm2 += z2
        cm4 += z2 * z2
    end
    cm4 /= n
    cm2 /= n
    return (cm4 / (cm2 * cm2)) - 3.0
end

function kurtosis(v::RealArray, wv::WeightVec, m::Real)
    n = length(v)
    length(wv) == n || throw(DimensionMismatch("Inconsistent array lengths."))
    cm2 = 0.0  # empirical 2nd centered moment (variance)
    cm4 = 0.0  # empirical 4th centered moment
    w = values(wv)

    @inbounds for i = 1 : n
        x_i = v[i]
        w_i = w[i]
        z = x_i - m
        z2 = z * z
        z2w = z2 * w_i
        cm2 += z2w
        cm4 += z2w * z2
    end
    sw = sum(wv)
    cm4 /= sw
    cm2 /= sw
    return (cm4 / (cm2 * cm2)) - 3.0
end

kurtosis(v::RealArray) = kurtosis(v, mean(v))
kurtosis(v::RealArray, wv::WeightVec) = kurtosis(v, wv, mean(v, wv))