# Two-qubit stabilizer states from Table II in
# Garcia et al. Efficient Inner-product Algorithm for Stabilizer States.
# arXiv:1210.6646 (2012)
TWOQUBIT = Ket.(normalize!.(Vector{Complex{Float64}}[
    # Separable states
    [0,0,1,0],[0,0,0,1],[0,0,1,1],[0,0,1,-1],[0,0,1,1im],[0,0,1,-1im],
    [1,0,-1,0],[0,1,0,-1],[1,1,-1,-1],[1,-1,-1,1],[1,1im,-1,-1im],[1,-1im,-1,1im],
    [1,0,-1im,0],[0,1,0,-1im],[1,1,-1im,-1im],[1,-1,-1im,1im],[1,1im,-1im,1],[1,-1im,-1im,-1],
    [1,0,1,0],[0,1,0,1],[1,1,1,1],[1,-1,1,-1],[1,1im,1,1im],[1,-1im,1,-1im],
    [1,0,1im,0],[0,1,0,1im],[1,1,1im,1im],[1,-1,1im,-1im],[1,1im,1im,-1],[1,-1im,1im,1],
    [1,0,0,0],[0,1,0,0],[1,1,0,0],[1,-1,0,0],[1,1im,0,0],[1,-1im,0,0],
    # Entangled states
    [0,1,1,0],[1,0,0,-1],[1,0,0,1],[0,1,-1,0],
    [1,0,0,1im],[0,1,1im,0],[0,1,-1im,0],[1,0,0,-1im],
    [1,1,1,-1],[1,1,-1,1],[1,-1,1,1],[1,-1,-1,-1],
    [1,1im,1,-1im],[1,1im,-1,1im],[1,-1im,1,1im],[1,-1im,-1,-1im],
    [1,1,1im,-1im],[1,1,-1im,1im],[1,-1,1im,1im],[1,-1,-1im,-1im],
    [1,1im,1im,1],[1,1im,-1im,-1],[1,-1im,1im,-1],[1,-1im,-1im,1],
    ]),((2,2),))


"""
Orthogonalize!(V, hv, vv, orth)

Orthogonalize the Krylov vectors using your (my) choice of
methods. Anything other than classical Gram-Schmidt twice (cgs2) is
likely to become an undocumented option. Methods other than cgs2 are
for CI for the linear solver.
"""
function Orthogonalize!(V, hv, vv, orth = "mgs1")
    orthopts = ["mgs1", "mgs2", "cgs1", "cgs2"]
    orth in orthopts || error("Impossible orth spec in Orthogonalize!")
    if orth == "mgs1"
        mgs!(V, hv, vv)
    elseif orth == "mgs2"
        mgs!(V, hv, vv, "twice")
    elseif orth == "cgs1"
        cgs!(V, hv, vv, "once")
    else
        cgs!(V, hv, vv, "twice")
    end
end

"""
cgs!(V, hv, vv, orth="twice")

Classical Gram-Schmidt.
"""
function cgs!(V, hv, vv, orth)
    T = eltype(V)
    k = length(hv)
    @views rk = hv[1:k-1]
    onep = T(1.0)
    zerop = T(0.0)
    pk = zeros(T, k - 1)
    qk = vv
    Qkm = V
    #
    # As the inventors of the BLAS said, "RTFM"!
    #
    # Orthogonalize!
    BLAS.gemv!('T', onep, Qkm, qk, onep, rk)
    BLAS.gemv!('N', -onep, Qkm, rk, onep, qk)
    if orth == "twice"
        #
        # Orthogonalize again!
        BLAS.gemv!('T', onep, Qkm, qk, zerop, pk)
        BLAS.gemv!('N', -onep, Qkm, pk, onep, qk)
        rk .+= pk
    end
    #
    # Keep track of what you did.
    nqk = norm(qk)
    qk ./= nqk
    hv[k] = nqk
end

#
function cgs!(V::SubArray{Float16,2}, hv, vv, orth)
    #
    #   no BLAS
    #
    k = length(hv)
    T = eltype(V)
    onep = T(1.0)
    zerop = T(0.0)
    @views rk = hv[1:k-1]
    pk = zeros(T, size(rk))
    qk = vv
    Qkm = V
    # Orthogonalize
    rk .+= Qkm' * qk
    qk .-= Qkm * rk
    if orth == "twice"
        # Orthogonalize again
        pk .= Qkm' * qk
        qk .-= Qkm * pk
        rk .+= pk
    end
    # Keep track of what you did.
    nqk = norm(qk)
    nqk != 0.0 || println("breakdown")
    nqk == 0.0 || qk ./= nqk
    hv[k] = nqk
end

"""
mgs!(V, hv, vv, orth)
"""
function mgs!(V, hv, vv, orth = "once")
    k = length(hv) - 1
    normin = norm(vv)
    #p=copy(vv)
    @views for j = 1:k
        p = vec(V[:, j])
        hv[j] = p' * vv
        vv .-= hv[j] * p
    end
    hv[k+1] = norm(vv)
    if (normin + 0.001 * hv[k+1] == normin) && (orth == "twice")
        @views for j = 1:k
            p = vec(V[:, j])
            hr = p' * vv
            hv[j] += hr
            vv .-= hr * p
        end
        hv[k+1] = norm(vv)
    end
    nv = hv[k+1]
    #
    # Watch out for happy breakdown
    #
    #if hv[k+1] != 0
    #@views vv .= vv/hv[k+1]
    nv != 0 || println("breakdown")
    if nv != 0
        vv ./= nv
    end
end


using Statistics

mutable struct Acao
    DATAPREGAO::Vector{String}
    CODNEG::String
    QUATOT::Vector{Int}
    PREABE::Vector{Float64}
    PREMAX::Vector{Float64}
    PREMIN::Vector{Float64}
    PREMED::Vector{Float64}
    PREULT::Vector{Float64}
    variancia::Float64
    media::Float64
    desvioPadrao::Float64
    mediaMovelSimplesCurtoPrazo::Vector{Float64}
    mediaMovelSimplesLongoPrazo::Vector{Float64}
    mediaMovelSimplesRetorno5::Vector{Float64}
end

function buildAcao(DATAPREGAO::String, CODNEG::String, QUATOT::Int, PREABE::Float64, PREMAX::Float64, PREMIN::Float64, PREMED::Float64, PREULT::Float64)
    return Acao([DATAPREGAO], CODNEG, [QUATOT], [PREABE], [PREMAX], [PREMIN], [PREMED], [PREULT], 0, 0, 0, Float64[], Float64[], Float64[])
end

function calcularRetorno(serie::Vector{Float64})::Vector{Float64}
    retorno::Vector{Float64} = Float64[]
    anterior = serie[1]

    for valor in serie[2:lastindex(serie)]
        append!(retorno, (valor - anterior) / anterior)
    end

    return retorno
end

function mediaMovelSimples(serieAnterior::Union{Vector{Float64},Nothing}, serie::Vector{Float64}, periodos::Int64)::Vector{Float64}
    _serie = serieAnterior !== nothing ? append!(serieAnterior, serie) : serie
    _mediaMovel::Vector{Float64} = Float64[]
    _ultimoIndice = lastindex(_serie)
    periodoInicial = 1
    periodoAtual = periodos

    while periodoAtual <= _ultimoIndice
        media = 0
        for valor in _serie[periodoInicial:periodoAtual]
            media += valor
        end
        append!(_mediaMovel, media / periodos)
        periodoAtual += 1
        periodoInicial += 1
    end

    return _mediaMovel
end

function mediaMovelExponencial(serieAnterior::Vector{Float64}, serie::Vector{Float64}, periodos::Int16)::Vector{Float64}
    _serie = length(serieAnterior) > 0 ? append!(serieAnterior, serie) : serie
    _mediaMovel::Vector{Float64} = Float64[]
    _ultimoIndice = lastindex(_serie)
    periodoInicial = 1
    periodoAtual = periodos

    while periodoAtual <= _ultimoIndice
        media = 0
        for valor in _serie[periodoInicial:periodoAtual]
            media += valor
        end
        append!(_mediaMovel, media / periodos)
        periodoAtual += 1
        periodoInicial += 1
    end

    return _mediaMovel
end
@time begin
    println("Portifolio de Ações")
    _dados = Float64[]

    arquivo_cotacoes = "/Users/rafaeldias/Downloads/COTAHIST_A2021 3.TXT"
    println("Carregar cotações do arquivo $arquivo_cotacoes")
    _arquivo = open(arquivo_cotacoes, "r")
    _dados = readlines(_arquivo)
    close(_arquivo)

    arquivo_cotacoes = "/Users/rafaeldias/Downloads/COTAHIST_A2022 5.TXT"
    println("Carregar cotações do arquivo $arquivo_cotacoes")
    _arquivo = open(arquivo_cotacoes, "r")
    append!(_dados, readlines(_arquivo))
    close(_arquivo)

    global acoes = Dict{String,Acao}()
    global carteiraAtual = Dict{String,Acao}()

    for _cotacao in _dados
        if length(_cotacao) > 0
            TIPREG = _cotacao[1:2]
            DATAPREGAO = _cotacao[3:10]
            CODNEG = replace(_cotacao[13:24], " " => "")
            TPMERC = _cotacao[25:27]
            QUATOT = _cotacao[153:170]
            PREABE = _cotacao[57:69]
            PREMAX = _cotacao[70:82]
            PREMIN = _cotacao[83:95]
            PREMED = _cotacao[96:108]
            PREULT = _cotacao[109:121]

            # Mercado Fracionario
            if TIPREG == "01" && TPMERC == "020"
                QUATOT = parse(Int, QUATOT)
                PREABE = parse(Int, PREABE) / 100.00 # DIVIDIR POR 100
                PREMAX = parse(Int, PREMAX) / 100.00 # DIVIDIR POR 100
                PREMIN = parse(Int, PREMIN) / 100.00 # DIVIDIR POR 100
                PREMED = parse(Int, PREMED) / 100.00 # DIVIDIR POR 100
                PREULT = parse(Int, PREULT) / 100.00 # DIVIDIR POR 100
                if haskey(acoes, CODNEG)
                    append!(acoes[CODNEG].QUATOT, QUATOT)
                    append!(acoes[CODNEG].PREABE, PREABE)
                    append!(acoes[CODNEG].PREMAX, PREMAX)
                    append!(acoes[CODNEG].PREMIN, PREMIN)
                    append!(acoes[CODNEG].PREMED, PREMED)
                    append!(acoes[CODNEG].PREULT, PREULT)
                else
                    acoes[CODNEG] = buildAcao(DATAPREGAO, CODNEG, QUATOT, PREABE, PREMAX, PREMIN, PREMED, PREULT)
                end
            end

            # Mercado a Vista
            if TIPREG == "01" && TPMERC == "010" && (startswith(CODNEG, "BOVA11"))
                QUATOT = parse(Int, QUATOT)
                PREABE = parse(Int, PREABE) / 100.00 # DIVIDIR POR 100
                PREMAX = parse(Int, PREMAX) / 100.00 # DIVIDIR POR 100
                PREMIN = parse(Int, PREMIN) / 100.00 # DIVIDIR POR 100
                PREMED = parse(Int, PREMED) / 100.00 # DIVIDIR POR 100
                PREULT = parse(Int, PREULT) / 100.00 # DIVIDIR POR 100

                if haskey(acoes, CODNEG)
                    append!(acoes[CODNEG].QUATOT, QUATOT)
                    append!(acoes[CODNEG].PREABE, PREABE)
                    append!(acoes[CODNEG].PREMAX, PREMAX)
                    append!(acoes[CODNEG].PREMIN, PREMIN)
                    append!(acoes[CODNEG].PREMED, PREMED)
                    append!(acoes[CODNEG].PREULT, PREULT)
                else
                    acoes[CODNEG] = buildAcao(DATAPREGAO, CODNEG, QUATOT, PREABE, PREMAX, PREMIN, PREMED, PREULT)
                end
            end
        end
    end
    _dados = nothing
    # Carteira atual
    # ITUB4 6 23.2
    # SBSP3 3 36.17
    # ABEV3 7 15.72
    # BOVA11 1 103.49
    # LUPA3F 1

    somaVariancia = 0.0
    alocacoes = Dict{String,Vector{Float64}}()
    alocacoesAux = Vector{Float64}()
    codigos = Vector{String}()
    for acao::Acao in sort!(collect(values(acoes)), by = x -> var(x.PREULT))
        if acao.CODNEG == "CBAV3F" || acao.CODNEG == "ITUB4F" || acao.CODNEG == "SBSP3F" || acao.CODNEG == "_ABEV3F" || acao.CODNEG == "_BOVA11" || acao.CODNEG == "_LUPA3F" || acao.CODNEG == "ENAT3F"
            global somaVariancia += var(acao.PREULT)
        end
    end
    f = open("cotacoes.json", "w")
    write(f, "[")
    for acao::Acao in sort!(collect(values(acoes)), by = x -> var(x.PREULT))
        acao.variancia = var(acao.PREULT)
        acao.desvioPadrao = std(acao.PREULT)
        acao.media = mean(acao.PREULT)
        acao.mediaMovelSimplesCurtoPrazo = mediaMovelSimples(nothing, acao.PREULT, 5)
        acao.mediaMovelSimplesLongoPrazo = mediaMovelSimples(nothing, acao.PREULT, 30)
        acao.mediaMovelSimplesRetorno5 = mediaMovelSimples(nothing, calcularRetorno(acao.PREULT), 5)

        mediaRetorno = mean(last(acao.mediaMovelSimplesRetorno5, 5))
        mediaVolume = mean(last(acao.QUATOT, 5))

        acimaMedia = 0
        abaixoMedia = 0

        for retorno in acao.mediaMovelSimplesRetorno5
            if retorno > mediaRetorno
                acimaMedia += 1
            else
                abaixoMedia += 1
            end
        end

        if acao.CODNEG == "CBAV3F" || acao.CODNEG == "ITUB4F" || acao.CODNEG == "SBSP3F" || acao.CODNEG == "_BOVA11" || acao.CODNEG == "ENAT3F" || acao.CODNEG == "_LUPA3F" || acao.CODNEG == "_ABEV3F"
            global alocacoes[acao.CODNEG] = [1. / (acao.variancia / somaVariancia)]
            push!(alocacoesAux, 1. / (acao.variancia / somaVariancia))
            push!(codigos, acao.CODNEG)
            write(f, replace("{\"fatorAlocacao\":$(1. / (acao.variancia / somaVariancia)),\"portifolio\":\"SIM\",\"codigo\":\"$(acao.CODNEG)\",\"variancia\":$(acao.variancia),\"desviopadrao\":$(acao.desvioPadrao),\"media\":$(acao.media),\"volume\":$(mean(acao.QUATOT)),\"mmscurto\":$(last(acao.mediaMovelSimplesCurtoPrazo, 5)),\"mmslongo\":$(last(acao.mediaMovelSimplesLongoPrazo, 5)),\"mmsretorno\":$(last(acao.mediaMovelSimplesRetorno5, 5)), \"mediaretorno\":$(mediaRetorno)},\n", "Float64[]" => "[]", "NaN" => "0"))
        end

        # if mediaRetorno > 0 && acimaMedia > abaixoMedia && last(acao.mediaMovelSimplesRetorno5) > mediaRetorno && mediaVolume > 1000
        #     global alocacoes[acao.CODNEG] = [1. / (acao.variancia / somaVariancia)]
        #     push!(alocacoesAux, 1. / (acao.variancia / somaVariancia))
        #     push!(codigos, acao.CODNEG)
        #     write(f, replace("{\"portifolio\":\"NAO\",\"codigo\":\"$(acao.CODNEG)\",\"variancia\":$(acao.variancia),\"desviopadrao\":$(acao.desvioPadrao),\"media\":$(acao.media),\"volume\":$(mean(acao.QUATOT)),\"mmscurto\":$(last(acao.mediaMovelSimplesCurtoPrazo, 5)),\"mmslongo\":$(last(acao.mediaMovelSimplesLongoPrazo, 5)),\"mmsretorno\":$(last(acao.mediaMovelSimplesRetorno5, 5)), \"mediaretorno\":$(mediaRetorno)},\n", "Float64[]" => "[]", "NaN" => "0"))
        # end
    end
    write(f, "{\"codigo\":\"AAAAAA\"}]")
    flush(f)
    close(f)
    acoes = nothing
end

b = 1600.
a = transpose(alocacoesAux) \ [b]
for i in 1:length(codigos)
    print("$(codigos[i]) $(a[i] * alocacoesAux[i]) -> $(100. * ((a[i] * alocacoesAux[i]) / b))% => $(b * (((a[i] * alocacoesAux[i]) / b)))\n")
end


using StructuralCausalModels, Test

ProjDir = @__DIR__

d_str = "dag{k1 -> k2;k2 -> y;v -> x2;w -> k1;x1 -> v;x1 -> w;x2 -> y;x3 -> w;x3 -> y}"

dag = DAG("test_desc", d_str)

to_dagitty(dag)

#=
fname = joinpath(ProjDir, "test_descendents_03.dot")
to_graphviz(dag, fname)
Sys.isapple() && run(`open -a GraphViz.app $(fname)`)
=#

bs = basis_set(dag)

@testset "basis_set_01" begin

  @test length(bs) == 48
  @test bs[18] == [:x2, :w, :v]
  @test bs[19] == [:x2, :w, :x1]
  @test bs[20] == [:x2, :w, :v, :x1]
  @test bs[44] == [:y, :k1, :w, :k2]
  @test bs[45] == [:y, :k1, :w, :x2, :k2]
  @test bs[46] == [:y, :k1, :w, :x3, :k2]
  @test bs[47] == [:y, :k1, :x2, :x3, :k2]
  @test bs[48] == [:y, :k1, :w, :x2, :x3, :k2]

end


# kspace.jl
# kspace sampling patterns
# 2018-06-13 Jeff Fessler, University of Michigan

using Plots: plot, plot!, scatter, gui
using Test

"""
`kspace = ir_mri_kspace_ga_radial(; Nro=?, Nspoke=?, ...)`

Generate k-space sampling pattern for "golden angle" radial sampling.

option
* `Nro:Int`		number of samples in each readout/spoke, default 256
* `Nspoke::Int`	number of spokes, default 1
* `start::Real`	first angle in series [radians], default pi/2
* `angle::Real`	angular spacing [radians], default GA
* `delta_ro::Real`	readout spacing, default 1/Nro
* `shift::Real`		shift due to gradient delays, default 0
		radial sample locations are `ir * delta_ro`
		where `ir = [-(Nro/2 - 1):1:Nro/2] + shift`
* `show::Bool`		plot k-space locations, default `false`

out
* `kspace`	`[Nro Nspoke 2]` (Float32)
			kx and ky k-space locations for Nspoke*Nro samples
			in interval (-0.5 0.5] for default shift, delta_ro
			so default units are "cycles / sample"

2015-07 Mai Le, original Matlab version
2015-07-04 Jeff Fessler, minor changes to Matlab version
"""
function ir_mri_kspace_ga_radial(;
		Nro::Int = 256,
		Nspoke::Int = 1,
		delta_ro::Real = 1/Nro,
		shift::Real = 0,
		angle::Real = pi*(sqrt(5)-1)/2, # golden angle [radians]
		start::Real = pi/2,
		show::Bool = false,
	)

	isodd(Nro) && throw(ArgumentError("Number along read out not even!"))

	rho = (-(Nro/2 - 1):1:Nro/2) * delta_ro # radial samples
	rho = rho .+ shift * delta_ro # apply gradient delay shift

	phi = (0:Nspoke-1) * angle .+ start # angles

	rho = Float32.(rho)
	phi = Float32.(phi)

	kx = rho * cos.(phi)' # [Nro Nspoke]
	ky = rho * sin.(-phi)'
	kspace = cat(dims=3, kx, ky) # [Nro Nspoke 2]

	if show
		scatter(kx, ky, label="", aspect_ratio=1)
		kmax = delta_ro * Nro / 2
		plot!(xlim = [-1,1]*kmax)
		plot!(ylim = [-1,1]*kmax)
		plot!(xtick = (-1:1) * kmax, xlabel = "k_x")
		plot!(ytick = (-1:1) * kmax, ylabel = "k_y")
		gui()
	end

	return kspace
end

"""
`ir_mri_kspace_ga_radial`
self test
"""
function ir_mri_kspace_ga_radial(test::Symbol)
	test != :test && throw(":test")
	k = ir_mri_kspace_ga_radial(Nro=30, Nspoke=13, start=pi,
		show=true, shift=-0.5, delta_ro=1)
	@test size(k) == (30, 13, 2)
	true
end

# ir_mri_kspace_ga_radial(:test)


const _EPOUT_BETA0_CACHE_ = Dict()

"""
    Make MaxEnt-EP but applaying MaxEnt starting from a EP solution at beta zero.
    It is important to pass a epout solution computed with beta zero.
    It is faster than the fully integrated version 'maxent_ep' but less accurate.
    It is used for secundary computations as estimating a beta which lead to a
    given objective value and so use this beta for the fully integrated version of
    the algorithm
"""
function fast_maxent_ep(model::MetNet, epout_β0::EPOut,
        alpha::Real,  beta_vec::AbstractVector)

    #= Anna's code scales fluxes to [-1,1], but it
    doesn't rescale a,d back to the original range =#
    maxflux = max(maximum(abs.(model.lb)), maximum(abs.(model.ub)))
    a = epout_β0.sol.a * maxflux;
    d = epout_β0.sol.b * maxflux^2;
    #= a, d are the mean and variance of the univariate Gaussians
    used as site approximations in the EP =#

    #= v, Σ are the mean vector and covariance matrix of Q,
    the full multivariate Gaussian (in Braunstein et al paper) =#
    if haskey(_EPOUT_BETA0_CACHE_, epout_β0)
        Σ = _EPOUT_BETA0_CACHE_[epout_β0]["Σ"]
        v = _EPOUT_BETA0_CACHE_[epout_β0]["v"]
    else
        S = model.S
        b = model.b
        D = Diagonal(1.0 ./ d)
        invΣ = Matrix(alpha*S'*S + D)
        Σ = inv(invΣ)
        v = Σ*(alpha*S'*b + D*a)

        empty!(_EPOUT_BETA0_CACHE_)
        _EPOUT_BETA0_CACHE_[epout_β0] = Dict()
        _EPOUT_BETA0_CACHE_[epout_β0]["Σ"] = Σ
        _EPOUT_BETA0_CACHE_[epout_β0]["v"] = v
    end

    #= From Cossios papar (see README) =#
    w = v + Σ * beta_vec
    wn = (d .* w - diag(Σ) .* a) ./ (d .- diag(Σ)) # mean of the non-truncated marginals
    Σnn = d .* diag(Σ) ./ (d .- diag(Σ)) # variances of the non-truncated marginals

    # Fix case of blocked reactoins
    tns = Truncated.(Normal.(wn, sqrt.(Σnn)), model.lb, model.ub)

    return EPOut(wn, Σnn, mean.(tns), var.(tns), epout_β0.sol, epout_β0.status, epout_β0.iter)
end


"""
$(TYPEDEF)
$(TYPEDFIELDS)

"""
struct QuantLogUniform <: LogUniformQuantDist
    low::NestedFloat
    high::NestedFloat
    q::NestedFloat
end

function quantloguniform(low::Float64, high::Float64, q::Float64)::Float64
    return round(exp(uniform(low, high)) / q) * q
end

loguniformquant(x::QuantLogUniform, low, high, q) = quantloguniform(low, high, q)


# Defazio, Bach, Lacoste-Julien: SAGA: A fast incremental gradient method
# with support for non-strongly convex composite objectives.
# In: Advances in neural information processing systems, pp. 1646–1654 (2014).
#
# Reddi, Sra, Poczos, and Smola, "Proximal stochastic methods for nonsmooth
# nonconvex finite-sum optimization." In Advances in Neural Information
# Processing Systems, pp. 1145–1153 (2016).
#
# Schmidt, Le Roux, Bach, "Minimizing finite sums with the stochastic average gradient"
# Mathematical Programming, 162(1-2), 83-112 (2017).
#

using LinearAlgebra
using ProximalOperators
using ProximalAlgorithms.IterationTools
using Printf
using Base.Iterators
using Random

export solution

include("SAGA_basic.jl")

struct SAGA{R<:Real}
    γ::Maybe{R}
    maxit::Int
    verbose::Bool
    freq::Int
    SAG_flag::Bool
    function SAGA{R}(;
        γ::Maybe{R} = nothing,
        maxit::Int = 10000,
        verbose::Bool = false,
        freq::Int = 1000,
        SAG_flag::Bool = false,
    ) where {R}
        @assert γ === nothing || γ > 0
        @assert maxit > 0
        @assert freq > 0
        new(γ, maxit, verbose, freq, SAG_flag)
    end
end

function (solver::SAGA{R})(
    x0::AbstractArray{C};
    F = nothing,
    g = ProximalOperators.Zero(),
    L = nothing,
    N = N,
) where {R,C<:RealOrComplex{R}}

    stop(state::SAGA_basic_state) = false
    disp(it, state) = @printf "%5d | %.3e  \n" it state.γ

    F === nothing && (F = fill(ProximalOperators.Zero(), (N,)))

    # dispatching the structure
    iter = SAGA_basic_iterable(F, g, x0, N, L, solver.γ, solver.SAG_flag)
    iter = take(halt(iter, stop), solver.maxit)
    iter = enumerate(iter)
    num_iters, state_final = nothing, nothing
    for (it_, state_) in iter  # unrolling the iterator
        # see https://docs.julialang.org/en/v1/manual/interfaces/index.html
        if solver.verbose && mod(it_, solver.freq) == 0
            disp(it_, state_)
        end
        num_iters, state_final = it_, state_
    end
    if solver.verbose && mod(num_iters, solver.freq) !== 0
        disp(num_iters, state_final)
    end # for the final iteration
    return solution(state_final), num_iters
end

"""
    SAGA([γ, maxit, verbose, freq])

Instantiate the SAGA algorithm  for solving convex optimization problems of the form

    minimize 1/N sum_{i=1}^N f_i(x) + g(x)

If `solver = SAGA(args...)`, then the above problem is solved with

	solver(x0, [F, g, N, L])

where F is an array containing f_i's, x0 is the initial point, and L is an array of
smoothness moduli of f_i's; it is optional when γ is provided.

Optional keyword arguments are:
* `γ`: stepsize
* `L`: an array of smoothness moduli of f_i's
* `maxit::Integer` (default: `10000`), maximum number of iterations to perform.
* `verbose::Bool` (default: `true`), whether or not to print information during the iterations.
* `freq::Integer` (default: `100`), frequency of verbosity.

References:

[1] Defazio, Bach, Lacoste-Julien, "SAGA: A fast incremental gradient method
with support for non-strongly convex composite objectives"
In: Advances in neural information processing systems, pp. 1646–1654 (2014).

[2] Reddi, Sra, Poczos, and Smola, "Proximal stochastic methods for nonsmooth
nonconvex finite-sum optimization" In Advances in Neural Information
Processing Systems, pp. 1145–1153 (2016).
"""

SAGA(::Type{R}; kwargs...) where {R} = SAGA{R}(; kwargs...)
SAGA(; kwargs...) = SAGA(Float64; kwargs...)

"""
If `solver = SAGA(args...)`, then

    itr = iterator(solver, x0, [F, g, N, L])

is an iterable object. Note that [maxit, verbose, freq] fields of the solver are ignored here.

The solution at any given state can be obtained using solution(state), e.g.,
for state in Iterators.take(itr, maxit)
    # do something using solution(state)
end

See https://docs.julialang.org/en/v1/manual/interfaces/index.html
and https://docs.julialang.org/en/v1/base/iterators/ for a list of iteration utilities
"""

function iterator(
    solver::SAGA{R},
    x0::AbstractArray{C};
    F = nothing,
    g = ProximalOperators.Zero(),
    L = nothing,
    N = N,
) where {R,C<:RealOrComplex{R}}

    F === nothing && (F = fill(ProximalOperators.Zero(), (N,)))
    # dispatching the iterator
    iter = SAGA_basic_iterable(F, g, x0, N, L, solver.γ, solver.SAG_flag)

    return iter
end

"""
    SAG([γ, maxit, verbose, freq])

Instantiate the SAG algorithm  for solving convex optimization problems of the form

    minimize 1/N sum_{i=1}^N f_i(x) + g(x)

If `solver = SAG(args...)`, then the above problem is solved with

    solver(x0, [F, g, N, L])

where F is an array containing f_i's, x0 is the initial point, and L is an array of
smoothness moduli of f_i's; it is optional when γ is provided.

Optional keyword arguments are:
* `γ`: stepsize
* `L`: an array of smoothness moduli of f_i's
* `maxit::Integer` (default: `10000`), maximum number of iterations to perform.
* `verbose::Bool` (default: `true`), whether or not to print information during the iterations.
* `freq::Integer` (default: `100`), frequency of verbosity.

References:

[1] Schmidt, Le Roux, Bach, "Minimizing finite sums with the stochastic average gradient"
Mathematical Programming, 162(1-2), 83-112 (2017).

"""

"""
If `solver = SAG(args...)`, then

    itr = iterator(solver, x0, [F, g, N, L])

is an iterable object. Note that [maxit, verbose, freq] fields of the solver are ignored here.

The solution at any given state can be obtained using solution(state), e.g.,
for state in Iterators.take(itr, maxit)
    # do something using solution(state)
end

See https://docs.julialang.org/en/v1/manual/interfaces/index.html
and https://docs.julialang.org/en/v1/base/iterators/ for a list of iteration utilities
"""

SAG(::Type{R}; kwargs...) where {R} = SAGA{R}(; kwargs..., SAG_flag = true)
SAG(; kwargs...) = SAG(Float64; kwargs...)


#############################################################################
# setdiff.jl
# Handles set differences of combinatorial sets.
#############################################################################

import Base: setdiff

export setdiff, get_elements

mutable struct SetDiffAtom <: CombiSet
  head::Symbol
  id_hash::UInt64
  children::Tuple
  elements::ValOrNothing
  baseset::AbstractArray
  cardinality::Int

  function SetDiffAtom(elements::ValOrNothing, children::Tuple, baseset::AbstractArray)
    this = new(:setdiff, 0, children, elements, baseset, length(baseset))
    this.id_hash = object_id(this)
    return this
  end
end

function setdiff(set1, set2::AbstractArray)
  set01 = Set(set1)
  set02 = Set(set2)
  return collect(setdiff(set01, set02))
end

function setdiff(set1::CombiSet, set2::AbstractArray)
  newset = SetDiffAtom(nothing, (set1, Constant(set2)), set1.baseset)
  return newset
end

function setdiff(set1::AbstractArray, set2::CombiSet)
  newset = SetDiffAtom(nothing, (Constant(set1), set2), set2.baseset)
  return newset
end

function setdiff(set1::CombiSet, set2::CombiSet)
  baseset = collect(union(Set(set1.baseset), Set(set2.baseset)))
  newset = SetDiffAtom(nothing, (set1, set2), baseset)
  return newset
end

function get_elements(x::SetDiffAtom)
  if isa(x.children[1], Constant)
    set1 = Set(x.children[1].value)
    set2 = x.children[2].elements
  elseif isa(x.children[2], Constant)
    set1 = x.children[1].elements
    set2 = Set(x.children[2].value)
  else
    set1 = x.children[1].elements
    set2 = x.children[2].elements
  end
  x.elements = setdiff(set1, set2)
  return x.elements
end


# Also see this:
# https://github.com/FluxML/Flux.jl/pull/524

# TODO:
# Possibly remove this file when Zygote.jl is officially the AD for Flux?

# cumsum
# https://discourse.julialang.org/t/how-to-create-tracked-cumsum-using-flux-jl/17772/2
Base.cumsum(x::TrackedArray; dims=1) = track(cumsum, x, dims)
@grad function cumsum(x::TrackedArray, dims)
  return cumsum(x.data, dims=dims), function(Δ)
    return (reverse(cumsum(reverse(Δ, dims=dims), dims=dims), dims=dims), nothing)
  end
end

# cumprod
# https://stackoverflow.com/questions/40916955/how-to-compute-gradient-of-cumprod-safely
Base.cumprod(x::TrackedArray; dims=1) = track(cumprod, x, dims)
@grad function cumprod(x::TrackedArray, dims)
  return cumprod(x.data, dims=dims), function(Δ)
    reverse(cumsum(reverse(cumprod(x, dims=dims) .* Δ, dims=dims), dims=dims), dims=dims) ./ x, nothing
  end
end

# reverse
# https://github.com/FluxML/Tracker.jl/blob/master/src/lib/array.jl
# Base.reverse(x::TrackedArray; dims=1) = track(reverse, x, dims)
# @grad function reverse(x::TrackedArray, dims)
#   return reverse(x.data, dims=dims), function(Δ)
#     return reverse(Δ, dims=dims), nothing
#   end
# end


using Zygote
using ForwardDiff
using BenchmarkTools

function besselj(ν, z; atol=1e-8)
    k = 0
    s = (z/2)^ν / factorial(ν)
    out = s
    while abs(s) > atol
        k += 1
        s *= (-1) / k / (k+ν) * (z/2)^2
        out += s
    end
    out
end

function grad_besselj_manual(ν, z; atol=1e-8)
    (besselj(ν-1, z; atol=atol) - besselj(ν+1, z); atol=atol)/2
end

println("Benchmarking Julia")
display(@benchmark besselj(2, 1.0))
println("Benchmarking Manual")
display(@benchmark grad_besselj_manual(2, 1.0))
println("Benchmarking Zygote")
display(@benchmark Zygote.gradient(besselj, 2, 1.0))
println("Benchmarking ForwardDiff")
display(@benchmark ForwardDiff.derivative(x->besselj(2, x), 1.0))


"""
    incrementalProblem(c, α, x, X, pd)

Solves incremental problem with specified costs `c`, parameter ``α \\in [0, 1]``,
first stage solutions `x` and a list of constraints `X` defining a set of feasible solutions.
It is subproblem of of [`evaluationProblem`](@ref) and [`adversarialProblem`](@ref).

Check section 4 _Solving the problems by MIP formulations_ of [publication](https://arxiv.org/abs/1811.06719)
for more information about this algorithm.

# Arguments
- `c`: is a vector of a nonnegative nominal second stage costs.
- `α`: fixed number belonging to ``[0, 1]``
- `x`: first stage solution.
- `X`: is a set of feasible solutions represented as a list functions, each of which
    accepts a list of JuMP variables as an argument and returns a JuMP linear constraint.
- `pd`: an instance of [`ProblemDescriptor`](@ref)
"""
function incrementalProblem(c, α, x, X, pd)
    n = size(c, 1)

    model = Model(solver=CplexSolver(CPXPARAM_ScreenOutput = getProperty("incrementalProblem.cplexLogging")))
    if ndims(c) == 1
        @variable(model, y[1:n], Bin)
        if hasEqualCardinalityProperty(pd)
            @variable(model, z[1:n], Bin)
        end
    elseif ndims(c) == 2
        m = size(c, 2)
        @variable(model, y[1:n, 1:m], Bin)
        if hasEqualCardinalityProperty(pd)
            @variable(model, z[1:n, 1:m], Bin)
        end
    else
        throw(ArgumentError("Unsupported number of dimensions ($(ndims(c)))"))
    end

    @objective(model, Min, vecdot(c, y))

    if hasEqualCardinalityProperty(pd)
        k = ceil(n * (1 - α))
        @constraint(model, sum(z) >= k)
        @constraint(model, z .<= y)
        @constraint(model, z .<= x)
    else
        @constraint(model, vecdot(x, 1 - y) <= α * sum(x))
    end

    # Constraints defining set of feasible solutions
    for constraint in X
        JuMP.addconstraint(model, constraint(y))
    end

    status = solve(model)
    return (getvalue(y), getobjectivevalue(model)) # InexactError is thrown here while solving computeAdversarialLowerBound -> evaluationProblem
end


using JLD

function sample_from_prior()
    # this just does rejection sampling to
    # stay in the identified region of the
    # MA(2) model. I'm too lazy to reflect
    # the bad draws
    ok = false
    theta1 = 0.
    theta2 = 0.
    while !ok
        theta1 = 4. * rand() - 2.
        theta2 = 2. * rand() - 1.
        ok = (theta2 > -1. + theta1) & (theta2 > -1. - theta1)
    end
    return [theta1 theta2]

end

function ma2(theta, n)
    e = randn(n+2)
    y = e[3:end] + theta[1]*e[2:end-1] + theta[2]*e[1:end-2]
end

function lags(x,p)
    cols = size(x,2)
    result = zeros(size(x,1), cols*p)
    for i = 1:p
        result[:,i*cols-cols+1:i*cols] = [ones(i,cols); x[1:end-i,:]]
    end
    return result
end

# fits an AR(P) model to MA(2) data
function aux_stat(theta, n, P)
    y = ma2(theta,n)
    x = lags(y,P)
    y = y[P+1:end]
    x = [ones(n-P,1) x[P+1:end,:]]
    Z = x\y
    return Z
end

n = 100 # sample size
R = 1000000 # replications
P = 10 # order of AR fit to the data
thetas = zeros(R,2)
Zs = zeros(R,P+1)
for i = 1:R
    if mod(i,10000)== 0.
        println(i)
    end
    theta = sample_from_prior()
    thetas[i,:] = theta'
    Zs[i,:] = aux_stat(theta, n, P)
end

save("MA2data.jld","thetas", thetas, "Zs", Zs)


import CLBLAS
using OpenCL: cl

const clblas = CLBLAS
clblas.setup()

device, ctx, queue = clblas.get_next_compute_context()

N = Unsigned(5)
data = Vector{Float32}(5)
data[1] = 1.1
data[2] = 2.22
data[3] = 3.333
data[4] = 4.4444
data[5] = 5.55555

bufX = cl.Buffer(Float32, ctx, (:r, :copy),hostbuf=data)

bufNRM2 = cl.Buffer(Float32, ctx, :w, 1)
scratchBuff = cl.Buffer(Float32, ctx, :rw, 2 * length(data))

offNRM2 = offx = Unsigned(0)
incx = cl.cl_int(1)
ncq = cl.cl_uint(1)
clq = queue.id
newl = cl.cl_uint(0)

event = cl.UserEvent(ctx, retain=true)
ptrEvent = [event.id]

clblas.clblasSnrm2( N, bufNRM2.id, offNRM2, bufX.id, offx, incx, scratchBuff.id, ncq, [clq], newl, C_NULL, ptrEvent);

cl.api.clWaitForEvents(cl.cl_uint(1), ptrEvent)

result = Vector{Float32}(1)
cl.enqueue_read_buffer(queue, bufNRM2, result, Unsigned(0), nothing, true)

expected = BLAS.nrm2(5, data, 1)
println(result)
println("------------")
println(expected)

try
    @assert(isapprox(norm(expected - result[1]), zero(Float32)))
    info("success!")
finally
    clblas.teardown()
end


"""
    InflationTrimmedMeanEq <: InflationFunction

Función de inflación para computar la media truncada equiponderada

## Utilización
    function (inflfn::InflationTrimmedMeanEq)(base::VarCPIBase{T}) where T

Define cómo opera InflationTrimmedMeanEq sobre un objeto de tipo VarCPIBase.

### Ejemplo
```julia-repl
julia> mtfn = InflationTrimmedMeanEq(25, 75.5)
(::InflationTrimmedMeanEq) (generic function with 5 methods)
julia> mtfn(gt00) # gt00 es de tipo VarCPIBase
120-element Vector{Float32}:
 1.3350569
 0.81459785
 0.5427902
 0.44377768
 0.3310551
 0.6161327
 ⋮
 0.284297
 0.20947386
 0.298732
 0.25540668
 0.2260508
 0.3456037
```
"""
Base.@kwdef struct InflationTrimmedMeanEq <: InflationFunction
    l1::Float32
    l2::Float32
end

# Métodos para crear funciones de inflación a partir de enteros
"""
    InflationTrimmedMeanEq(l1::Real,l2::Real)
Nos permite utilizar recortes que no necesariamente son de tipos de punto flotante, como por ejemplo: enteros, fracciones, Float32 y Float64.

# Ejemplo:
```julia-repl
julia> mtfn = InflationTrimmedMeanEq(25, 75.5)
(::InflationTrimmedMeanEq) (generic function with 5 methods)
```
"""
function InflationTrimmedMeanEq(l1::Real,l2::Real)
    # Obtener los recortes adecuados en tiempo de construcción
    InflationTrimmedMeanEq(Float32(l1), Float32(l2))
end

"""
    measure_name(inflfn::InflationTrimmedMeanEq)

Nos indica qué medida se utiliza para una instancia de una función de inflación.

# Ejemplo:

```julia-repl
julia> mtfn = InflationTrimmedMeanEq(15.5,75.5)
julia> measure_name(mtfn)
"Media Truncada Equiponderada (15.5 , 75.5)"
```
"""
function measure_name(inflfn::InflationTrimmedMeanEq)
    l1 = string(round(inflfn.l1, digits=2))
    l2 = string(round(inflfn.l2, digits=2))
    "Media Truncada Equiponderada (" * l1 * ", " * l2 * ")"
end

# Extendemos `params`, que devuelve los parámetros de la medida de inflación
CPIDataBase.params(inflfn::InflationTrimmedMeanEq) = (inflfn.l1, inflfn.l2)

# Operación de InflationTrimmedMeanEq sobre VarCPIBase para obtener el resumen
# intermensual de esta metodología
function (inflfn::InflationTrimmedMeanEq)(base::VarCPIBase{T}) where T
    # Obtener los percentiles de recorte
    l1 = inflfn.l1
    l2 = inflfn.l2
    # l1 = min(inflfn.l1, inflfn.l2)
    # l2 = max(inflfn.l1, inflfn.l2)

    # Determinamos en dónde truncar
    q1      = Int(ceil(length(base.w) * l1 / 100))
    q2      = Int(floor(length(base.w) * l2 / 100))
    outVec  = Vector{T}(undef, periods(base))

    if q1 == 0
        q1 = 1
    end

    # para cada t: ordenamos, truncamos y obtenemos la media.
    Threads.@threads for i in 1:periods(base)

        # Creamos una vista de cada fila: ahora temporal almacena una referencia
        # a la fila de base.v, sin crear nueva memoria
        temporal = @view base.v[i,:]
        # Ordenamos el vector y almacenamos en uno nuevo: `sorted_data`
        sorted  = sort(temporal)

        # No es necesario generar esta asignación porque la función mean puede
        # acceder a una vista del arreglo entre las posiciones q1 y q2, sin
        # alojar nueva memoria
        # temporal    = temporal[q1:q2]

        # Por lo que la siguiente operación es la de obtener el promedio entre
        # dichas posiciones, sin reservar nueva memoria
        @inbounds outVec[i]  = mean(@view sorted[q1:q2])
    end
    return outVec
end

# Método para recibir argumentos en forma de tupla
InflationTrimmedMeanEq(factors::Tuple{Real, Real}) = InflationTrimmedMeanEq(
    convert(Float32, factors[1]),
    convert(Float32, factors[2])
)

# Método para recibir argumentos en forma de vector
function InflationTrimmedMeanEq(factor_vec::Vector{<:Real})
    length(factor_vec) != 2 && return @error "Dimensión incorrecta del vector"
    InflationTrimmedMeanEq(
        convert(Float32, factor_vec[1]),
        convert(Float32, factor_vec[2])
    )
end


"""
    LegendreDerivativeOperator{T<:Real, Grid}

A derivative operator on a nonperiodic Lobatto-Legendre grid with scalar type
`T` computing the first derivative using a Legendre expansion.
"""
struct LegendreDerivativeOperator{T<:Real} <: AbstractDerivativeOperator{T}
    jac::T
    Δx::T
    grid::Vector{T} # N nodes, including the left and the right boundary
    basis::LobattoLegendre{T}

    function LegendreDerivativeOperator(xmin::T, xmax::T, basis::LobattoLegendre{T}) where {T<:Real}
        grid = map_from_canonical.(basis.nodes, xmin, xmax, basis)
        jac = 2 / (xmax - xmin)
        Δx = inv(jac)

        new{T}(jac, Δx, grid, basis)
    end
end

"""
    LegendreDerivativeOperator(xmin::T, xmax::T, N::Int) where {T<:Real}

Construct the `LegendreDerivativeOperator` on a uniform grid between `xmin` and
`xmax` using `N` nodes and `N-1` Legendre modes.
"""
function LegendreDerivativeOperator(xmin::T, xmax::T, N::Int) where {T<:Real}
    @argcheck N >= 2

    basis = LobattoLegendre(N-1, T)

    LegendreDerivativeOperator(xmin, xmax, basis)
end

function legendre_derivative_operator(xmin::T, xmax::T, N::Int) where {T<:Real}
    LegendreDerivativeOperator(xmin, xmax, N)
end

derivative_order(D::LegendreDerivativeOperator) = 1
LinearAlgebra.issymmetric(D::LegendreDerivativeOperator) = false
integrate(func, u, D::LegendreDerivativeOperator) = D.Δx*integrate(func, u, D.basis)
function evaluate_coefficients(u, D::LegendreDerivativeOperator,
                               npoints=2*size(D,2)+1)
    evaluate_coefficients(u, D.basis, npoints)
end

function Base.show(io::IO, D::LegendreDerivativeOperator{T}) where {T}
    x = grid(D)
    print(io, "First derivative operator {T=", T, "} \n")
    print(io, "on the Lobatto Legendre nodes in [", first(x), ", ", last(x),
                "] using ", length(x), " nodes. \n")
end

function mul!(dest::AbstractVector{T}, D::LegendreDerivativeOperator, u::AbstractVector{T}) where {T}
    @unpack jac, basis = D
    N, _ = size(D)
    @boundscheck begin
        @argcheck N == length(u)
        @argcheck N == length(dest)
    end

    mul!(dest, basis.D, u)
    dest .*= jac

    nothing
end

function left_boundary_weight(D::LegendreDerivativeOperator)
    @inbounds retval = D.Δx * D.basis.weights[1]
    retval
end

function right_boundary_weight(D::LegendreDerivativeOperator)
    @inbounds retval = D.Δx * D.basis.weights[end]
    retval
end

function ConstantFilter(basis::LobattoLegendre{T}, filter, tmp::Array) where {T}
    Np1 = length(basis.nodes)
    coefficients = Array{T}(undef, Np1)
    set_filter_coefficients!(coefficients, filter)
    modal2nodal = legendre_vandermonde(basis)
    nodal2modal = FactorisationWrapper(factorize(modal2nodal))

    ConstantFilter(coefficients, nodal2modal, modal2nodal, tmp, filter)
end

function ConstantFilter(basis::LobattoLegendre{T}, filter, TmpEltype::DataType=T) where {T}
    Np1 = length(basis.nodes)
    tmp = Array{TmpEltype}(undef, Np1)

    ConstantFilter(basis, filter, tmp)
end

"""
    ConstantFilter(D::LegendreDerivativeOperator, filter, TmpEltype=T)

Create a modal filter with constant parameters adapted to the Legendre
derivative operator `D` with parameters given by the filter function `filter`.
"""
function ConstantFilter(D::LegendreDerivativeOperator{T}, filter, TmpEltype=T) where {T}
    ConstantFilter(D.basis, filter, TmpEltype)
end

# TODO: LegendreSuperSpectralViscosity


using BayesianDataFusion
using Base.Test

X = sprand(50, 100, 0.1)
Xr = sparse_csr(X)

@test size(X) == size(Xr)
@test size(X,1) == size(Xr,1)
@test size(X,2) == size(Xr,2)
@test eltype(X) == eltype(Xr)

y1 = rand(100)
y2 = rand(50)

u1 = X * y1
u2 = X' * y2

v1 = Xr * y1
v2a = Ac_mul_B(Xr, y2)
v2b = At_mul_B(Xr, y2)

@test_approx_eq u1 v1
@test_approx_eq u2 v2a
@test_approx_eq u2 v2b

@test_approx_eq full(At_mul_B(X,X)) full(At_mul_B(Xr,Xr))

rows = [1, 2, 2, 4]
cols = [2, 1, 3, 3]
vals = [0.1, 0.2, 0.15, 0.3]
z    = rand(3)
w1   = sparse(rows, cols, vals) * z
w2   = sparse_csr(rows, cols, vals) * z
@test_approx_eq w1 w2


# f(x)=hcat(x...)
Fhcat(x...)=(hcat(x...),nothing)
Fhcat_inplace(value,auxvalue,x...)=copy!(value,hcat(x...))

function Dhcat(derivativeIDX,f_c,faux_c,grad_c,grad_n,x...)
    startind=1
    for i=1:length(x)
        endind=startind+length(x[i])-1
        if pointer(x[i])==pointer(x[derivativeIDX]) # use pointers so we can deal with hcat([A B A])
            axpy!(1.0,grad_c[startind:endind],grad_n)
        end
        startind=endind+1
    end
end

if PROC=="GPU"

    function Fhcat_inplace(value::CudaArray,auxvalue,x::CudaArray...) # inplace
        totallength=1
        for i in 1:length(x)
            copyinto!(value,x[i],totallength)
            totallength+=length(x[i])
        end
    end

    function Dhcat(derivativeIDX,f_c,faux_c,grad_c,grad_n::CudaArray,x...)
        startind=1
        for i=1:length(x)
            endind=startind+length(x[i])-1
            if pointer(x[i])==pointer(x[derivativeIDX]) # use pointers so we can deal with vcat([A B A])
                copyfrom_update!(grad_n,grad_c,startind)
            end
            startind=endind+1
        end
    end

end

Derivative[Fhcat]=Dhcat # Define dictionary lookup
Inplace[Fhcat]=Fhcat_inplace

import Base.hcat
hcat(n::ADnode)=ADnode(Fhcat,n)

hcat(n::Array{ADnode,1})=n
hcat(n::ArrayADnode)=ADnode(Fhcat,n)

export hcat


using LinearAlgebra

import ClimaCore:
    Domains, Fields, Geometry, Meshes, Operators, Spaces, Topologies
import ClimaCore.Geometry: ⊗

using OrdinaryDiffEq: ODEProblem, solve, SSPRK33

import Logging
import TerminalLoggers
Logging.global_logger(TerminalLoggers.TerminalLogger())

const parameters = (
    ϵ = 0.1,  # perturbation size for initial condition
    l = 0.5, # Gaussian width
    k = 0.5, # Sinusoidal wavenumber
    ρ₀ = 1.0, # reference density
    c = 2,
    g = 10,
)

domain = Domains.RectangleDomain(
    Domains.IntervalDomain(
        Geometry.XPoint(-2π),
        Geometry.XPoint(2π),
        periodic = true,
    ),
    Domains.IntervalDomain(
        Geometry.YPoint(-2π),
        Geometry.YPoint(2π),
        periodic = true,
    ),
)

n1, n2 = 16, 16
Nq = 4
Nqh = 7
mesh = Meshes.RectilinearMesh(domain, n1, n2)
grid_topology = Topologies.Topology2D(mesh)
quad = Spaces.Quadratures.GLL{Nq}()
space = Spaces.SpectralElementSpace2D(grid_topology, quad)

Iquad = Spaces.Quadratures.GLL{Nqh}()
Ispace = Spaces.SpectralElementSpace2D(grid_topology, Iquad)

function init_state(coord, p)
    x, y = coord.x, coord.y
    # set initial state
    ρ = p.ρ₀

    # set initial velocity
    U₁ = cosh(y)^(-2)

    # Ψ′ = exp(-(y + p.l / 10)^2 / 2p.l^2) * cos(p.k * x) * cos(p.k * y)
    # Vortical velocity fields (u₁′, u₂′) = (-∂²Ψ′, ∂¹Ψ′)
    gaussian = exp(-(y + p.l / 10)^2 / 2p.l^2)
    u₁′ = gaussian * (y + p.l / 10) / p.l^2 * cos(p.k * x) * cos(p.k * y)
    u₁′ += p.k * gaussian * cos(p.k * x) * sin(p.k * y)
    u₂′ = -p.k * gaussian * sin(p.k * x) * cos(p.k * y)

    u = Geometry.UVVector(U₁ + p.ϵ * u₁′, p.ϵ * u₂′)
    # set initial tracer
    θ = sin(p.k * y)

    return (ρ = ρ, ρu = ρ * u, ρθ = ρ * θ)
end

y0 = init_state.(Fields.coordinate_field(space), Ref(parameters))

function flux(state, param)
    ρ, ρu, ρθ = state.ρ, state.ρu, state.ρθ
    u = ρu / ρ
    return (
        ρ = ρu,
        ρu = ((ρu ⊗ u) + (param.g * ρ^2 / 2) * LinearAlgebra.I),
        ρθ = ρθ * u,
    )
end

function energy(state, param)
    ρ, ρu = state.ρ, state.ρu
    u = ρu / ρ
    return ρ * (u.u^2 + u.v^2) / 2 + param.g * ρ^2 / 2
end

function total_energy(y, parameters)
    sum(energy.(y, Ref(parameters)))
end

function rhs!(dydt, y, _, t)

    I = Operators.Interpolate(Ispace)
    div = Operators.WeakDivergence()
    R = Operators.Restrict(space)

    rparameters = Ref(parameters)

    @. dydt = -R(div(flux(I(y), rparameters)))

    Spaces.weighted_dss!(dydt)
    return dydt
end

# Next steps:
# 1. add the above to the design docs (divergence + over-integration + DSS)
# 2. add boundary conditions

dydt = similar(y0)
rhs!(dydt, y0, nothing, 0.0)

# Solve the ODE operator
prob = ODEProblem(rhs!, y0, (0.0, 80.0))
sol = solve(
    prob,
    SSPRK33(),
    dt = 0.02,
    saveat = 1.0,
    progress = true,
    progress_message = (dt, u, p, t) -> t,
)

ENV["GKSwstype"] = "nul"
using ClimaCorePlots, Plots
Plots.GRBackend()

dirname = "cg"
path = joinpath(@__DIR__, "output", dirname)
mkpath(path)

anim = Plots.@animate for u in sol.u
    Plots.plot(u.ρθ, clim = (-1, 1))
end
Plots.mp4(anim, joinpath(path, "tracer.mp4"), fps = 10)

Es = [total_energy(u, parameters) for u in sol.u]
Plots.png(Plots.plot(Es), joinpath(path, "energy.png"))

function linkfig(figpath, alt = "")
    # buildkite-agent upload figpath
    # link figure in logs if we are running on CI
    if get(ENV, "BUILDKITE", "") == "true"
        artifact_url = "artifact://$figpath"
        print("\033]1338;url='$(artifact_url)';alt='$(alt)'\a\n")
    end
end

linkfig(
    relpath(joinpath(path, "energy.png"), joinpath(@__DIR__, "../..")),
    "Total Energy",
)


#!/usr/bin/env julia

#struct session
#    ifvl::Bool
#    ifad::Bool
#    ifpr::Bool
#    ifps::Bool
#
#    nx1::Int
#    ny1::Int
#end

using SpectralElements

using FastGaussQuadrature
using Plots,LinearAlgebra,SparseArrays

using Zygote

#using NLopt,Optim
using DiffEqFlux
using Flux
using Statistics

#----------------------------------------------------------------------#
# size
#----------------------------------------------------------------------#
nx1 = 8; Ex = 5;
ny1 = 8; Ey = 5;

nx2 = nx1-2; nxd = Int(ceil(1.5*nx1));
ny2 = nx1-2; nyd = Int(ceil(1.5*ny1));

nxp = 5*nx1;
nyp = 5*ny1;

#----------------------------------------------------------------------#
# nodal operators
#----------------------------------------------------------------------#
zr1,wr1 = gausslobatto(nx1);  zs1,ws1 = gausslobatto(ny1);
zr2,wr2 = gausslobatto(nx2);  zs2,ws2 = gausslobatto(ny2);
zrd,wrd = gausslobatto(nxd);  zsd,wsd = gausslobatto(nyd);
zrp     = linspace(-1,1,nxp); zsp     = linspace(-1,1,nyp);

Jr1d = interpMat(zrd,zr1); Js1d = interpMat(zsd,zs1);
Jr2d = interpMat(zrd,zr2); Js2d = interpMat(zsd,zs2);
Jr21 = interpMat(zr1,zr2); Js21 = interpMat(zs1,zs2);
Jr1p = interpMat(zrp,zr1); Js1p = interpMat(zsp,zs1);

Dr1 = derivMat(zr1); Ds1 = derivMat(zs1);
Dr2 = derivMat(zr2); Ds2 = derivMat(zs2);
Drd = derivMat(zrd); Dsd = derivMat(zsd);

#----------------------------------------------------------------------#
# boundary conditions
#----------------------------------------------------------------------#
ifperiodicX = false
ifperiodicY = false

Ix1 = Matrix(I,Ex*nx1,Ex*nx1);
Iy1 = Matrix(I,Ey*ny1,Ey*ny1);

Rx1 = Ix1[2:end-1,:];
Ry1 = Iy1[2:end-1,:];

if(ifperiodicX) Rx1 = Ix1; end
if(ifperiodicY) Ry1 = Iy1; end

M1 = diag(Rx1'*Rx1) * diag(Ry1'*Ry1)'; # hom. dir BC

Ix1 = nothing
Iy1 = nothing
Rx1 = nothing
Ry1 = nothing
#----------------------------------------------------------------------#
# mapping
#----------------------------------------------------------------------#

# Q: global -> local op, Q': local -> global
Qx1 = semq(Ex,nx1,ifperiodicX); Qx2 = semq(Ex,nx2,ifperiodicX);
Qy1 = semq(Ey,ny1,ifperiodicY); Qy2 = semq(Ey,ny2,ifperiodicY);

# gather scatter op
QQtx1 = Qx1*Qx1'; QQtx2 = Qx2*Qx2';
QQty1 = Qy1*Qy1'; QQty2 = Qy2*Qy2';

# weight for inner products
mult1 = ones(nx1*Ex,ny1*Ey);
mult1 = gatherScatter(mult1,QQtx1,QQty1);
mult1 = @. 1 / mult1;

mult2 = ones(nx2*Ex,ny2*Ey);
mult2 = gatherScatter(mult2,QQtx2,QQty2);
mult2 = @. 1 / mult2;
#----------------------------------------------------------------------#
# geometry
#----------------------------------------------------------------------#
x1e,_ = semmesh(Ex,nx1); y1e,_ = semmesh(Ey,ny1);
x2e,_ = semmesh(Ex,nx2); y2e,_ = semmesh(Ey,ny2);
xde,_ = semmesh(Ex,nxd); yde,_ = semmesh(Ey,nyd);

x1,y1 = ndgrid(x1e,y1e);
x2,y2 = ndgrid(x2e,y2e);
xd,yd = ndgrid(xde,yde);

# deform grid with gordonhall
function usrgeom(r,s)
    x = @. 0.5*(r+1)
    y = @. 0.5*(s+1)
return x,y
end

x1,y1 = usrgeom(x1,y1)
x2,y2 = usrgeom(x2,y2)
xd,yd = usrgeom(xd,yd)

Jac1,Jaci1,rx1,ry1,sx1,sy1 = jac(x1,y1,Dr1,Ds1);
Jac2,Jaci2,rx2,ry2,sx2,sy2 = jac(x2,y2,Dr2,Ds2);
Jacd,Jacid,rxd,ryd,sxd,syd = jac(xd,yd,Drd,Dsd);

wx1 = kron(ones(Ex,1),wr1); wy1 = kron(ones(Ey,1),ws1);
wx2 = kron(ones(Ex,1),wr2); wy2 = kron(ones(Ey,1),ws2);
wxd = kron(ones(Ex,1),wrd); wyd = kron(ones(Ey,1),wsd);

# diagonal mass matrix
B1  = Jac1 .* (wx1*wy1');
B2  = Jac2 .* (wx2*wy2');
Bd  = Jacd .* (wxd*wyd');
Bi1 = 1 ./ B1;

#----------------------------------------------------------------------#
# debugging
#----------------------------------------------------------------------#

## don't use lumped mass matrices for now
#Jr1d = []
#Js1d = []
#rxd = rx1
#ryd = ry1
#sxd = sx1
#syd = sy1
#Bd = B1

#----------------------------------------------------------------------#
# case setup
#----------------------------------------------------------------------#
V  = 0.4
pa = 5
ε  = 1e-3
α  = 1e-8
f1 = @. 1e-2 + 0*x1
#M1[:,end] .= 1;
#M1[1,:]   .= 1;
af(p)   = @. 0.5*(tanh(p)+1)
kond(a) = @. ε + (1-ε)*a^pa

function problem(a) # topology parameter, forcing

    Ba = Bd .* ABu(Js1d,Jr1d,a);

    visc  = kond(a)
    viscd = ABu(Js1d,Jr1d,visc);
    G11 = @. viscd * Ba * (rxd * rxd + ryd * ryd);
    G12 = @. viscd * Ba * (rxd * sxd + ryd * syd);
    G22 = @. viscd * Ba * (sxd * sxd + syd * syd);

    #Ba = Zygote.hook(d->hmp(d),Ba);
    lhs(v) = lapl(v,M1,Jr1d,Js1d,QQtx1,QQty1,Dr1,Ds1,G11,G12,G22,mult1);

    f   = f1;

    rhs = mass(f,M1,Ba,Jr1d,Js1d,QQtx1,QQty1,mult1);

    return lhs, rhs
end

function opM(v) # preconditioner
    return v
end

function solver(opA,rhs,adj::Bool)
    if adj
        # zero out sensitivities at bounary points
        rhs = mask(rhs,M1);
        # ensure continuity and elem boundaries
        rhs = mult1 .* gatherScatter(rhs,QQty1,QQtx1);

        return pcg(rhs,opA,opM,mult1,false);
    else
        return pcg(rhs,opA,opM,mult1,false);
    end
end

function model(a)
    u = linsolve(a,problem,solver) # adjoint support thru Zygote
end

function loss(p)
    a = af(p)
    u = model(a) # du = 1 .* B1 .* a .* f
    f = f1
    adx,ady = grad(a,Dr1,Ds1,rx1,ry1,sx1,sy1);

    ad   = ABu(Js1d,Jr1d,a  );
    ud   = ABu(Js1d,Jr1d,u  );
    fd   = ABu(Js1d,Jr1d,f  );
    adxd = ABu(Js1d,Jr1d,adx);
    adyd = ABu(Js1d,Jr1d,ady);

    ll = @. fd*ud + α*(adxd^2+adyd^2);
    l  = sum(Bd.*ad.*ll);
    return l, u
end

#----------------------------------------------------------------------#
# DiffEqFlux.sciml_train
#----------------------------------------------------------------------#
p0 = @. 0.6 + 0.0*x1

global param = []
callback = function (p, l, pred; doplot = true)
  global param = p
  return false
end

result = DiffEqFlux.sciml_train(loss,p0
                                ,ADAM(1e-1),cb=Flux.throttle(callback,1)
                                ,maxiters=10)

p=param;dp=Zygote.gradient((p)->loss(p)[1], p)[1];
u= model(af(p));

Jp  = ABu(Js1p,Jr1p,p );
Jdp = ABu(Js1p,Jr1p,dp);
Ju  = ABu(Js1p,Jr1p,u );
pl=plot(heatmap(af(Jp)),heatmap(Jdp),heatmap(Ju));
display(pl)

display(loss(p0)[1])
display(loss(p )[1])
#----------------------------------------------------------------------#
# NLopt.Optimize
#----------------------------------------------------------------------#
#function fmin(a::Vector, grad::Vector)
#    a = reshape(a,nx1*Ex,ny1*Ey)
#    if length(grad)>0
#        grad[:] = Zygote.gradient((a)->loss(a)[1], a)[1][:]
#    end
#    return loss(a)[1]
#end

#function fc(a::Vector, grad::Vector)
#    a = reshape(a,nx1*Ex,ny1*Ey)
#    ineq(a) = sum(B1.*a) .- V
#    if length(grad)>0
#        grad[:] = Zygote.gradient((a)->ineq(a), a)[1][:]
#    end
#    return ineq(a)
#end

## set up and run optimizer
#a0 = rand(nx1*Ex,ny1*Ey)
#a0 = 0.4*ones(nx1*Ex,ny1*Ey)

#fig = mesh(x1,y1,model(a0)); display(fig);

#opt = NLopt.Opt(:LD_MMA, length(a0))
#opt.min_objective = fmin
#NLopt.inequality_constraint!(opt, fc)
#opt.lower_bounds = 0.
#opt.upper_bounds = 1.
#opt.maxeval = 20
#opt.xtol_abs = 1e-8
#opt.xtol_rel = 1e-8

#(optf,optx,ret) = NLopt.optimize(opt, a0[:])

#opta = reshape(optx,nx1*Ex,ny1*Ey)
#numevals = opt.numevals
#println("started with loss ",loss(a0)[1])
#println("got $optf after $numevals iterations (returned $ret)")
#fig = mesh(x1,y1,model(opta)); display(fig);
#----------------------------------------------------------------------#
nothing


mutable struct EmpiricalMean{T<:Real,V<:AbstractVector{<:Real}} <: PriorMean{T}
    C::V
    opt::Optimizer
end

"""
**EmpiricalMean**
```julia`
function EmpiricalMean(c::V=1.0;opt::Optimizer=Adam(α=0.01)) where {V<:AbstractVector{<:Real}}
```
Construct a constant mean with values `c`
Optionally give an optimizer `opt` (`Adam(α=0.01)` by default)
"""
function EmpiricalMean(c::V=1.0;opt::Optimizer=Adam(α=0.01)) where {V<:AbstractVector{<:Real}}
    EmpiricalMean{eltype(c),V}(c,opt)
end

function update!(μ::EmpiricalMean{T},grad::AbstractVector{T}) where {T<:Real}
    μ.C .+= update!(μ.opt,grad)
end

function (μ::EmpiricalMean{T})(x::AbstractMatrix) where {T<:Real}
    @assert size(x,1)==length(μ)
    return μ.C
end


module CP_EM

include("../Misc.jl")

using .Misc
using Distributions, SpecialFunctions

export standard_EM, dual_EM

function standard_EM(X::Array{ℜ,3}, R::Ƶ ; μ::ℜ=3.5, γ::ℜ=0.1, EPOCHS::Ƶ=1, ϵ::ℜ=1e-16) where {ℜ<:Real, Ƶ<:Int}
    I::Ƶ, J::Ƶ, K::Ƶ = size(X)

    a::ℜ, b::ℜ = I*J*K*γ, γ/μ
    α_R, α_IR, α_JR, α_KR = fill(a/R,R), fill(a/(I*R),I,R), fill(a/(J*R),J,R), fill(a/(K*R),K,R)

    log_λ = log(rand(Gamma(a,1.0/b)))
    log_θ_R = rand(Dirichlet(α_R .+ 1.0/R))
    log_θ_IR = reshape([log(θ_ir) for r=1:R for θ_ir=rand(Dirichlet(α_IR[:,r] .+ 1.0/I))],I,R)
    log_θ_JR = reshape([log(θ_jr) for r=1:R for θ_jr=rand(Dirichlet(α_JR[:,r] .+ 1.0/J))],J,R)
    log_θ_KR = reshape([log(θ_kr) for r=1:R for θ_kr=rand(Dirichlet(α_KR[:,r] .+ 1.0/K))],K,R)

    s, S_R, S_IR, S_JR, S_KR, S⁺ = zeros(ℜ,R), zeros(ℜ,R), zeros(ℜ,I,R), zeros(ℜ,J,R), zeros(ℜ,K,R), sum(X)
    q, log_ρ, log_ρ_ijk::ℜ = Array{ℜ}(undef,R), Array{ℜ}(undef,R), 0.0

    X_full = similar(X)

    for eph=1:EPOCHS
        S_R .= 0.0
        S_IR .= 0.0
        S_JR .= 0.0
        S_KR .= 0.0
        S⁺ = 0.0

        for k=1:K, j=1:J, i=1:I #order of traversal is important
            log_ρ .= log_λ .+ log_θ_R .+ log_θ_IR[i,:] .+ log_θ_JR[j,:] .+ log_θ_KR[k,:]
            log_ρ_ijk = logsumexp(log_ρ)
            q .= exp.(log_ρ .- log_ρ_ijk)

            X_full[i,j,k] = isnan(X[i,j,k]) ? exp(log_ρ_ijk) : X[i,j,k]
            s .= X_full[i,j,k] .* q

            S_R .+= s
            S_IR[i,:] .+= s
            S_JR[j,:] .+= s
            S_KR[k,:] .+= s
            S⁺ += X_full[i,j,k]
        end

        log_λ = log(max(S⁺+a-1.0, ϵ)) - log(b+1.0)
        log_θ_R .= log.(max.(S_R.+α_R.-1.0, ϵ))
        log_θ_IR .= log.(max.(S_IR.+α_IR.-1.0, ϵ))
        log_θ_JR .= log.(max.(S_JR.+α_JR.-1.0, ϵ))
        log_θ_KR .= log.(max.(S_KR.+α_KR.-1.0, ϵ))

        log_θ_R .-= logsumexp(log_θ_R)
        log_θ_IR .-= logsumexp(log_θ_IR;dims=1)
        log_θ_JR .-= logsumexp(log_θ_JR;dims=1)
        log_θ_KR .-= logsumexp(log_θ_KR;dims=1)
    end
    return X_full, log_λ, (log_θ_R, log_θ_IR, log_θ_JR, log_θ_KR)
end

function dual_EM(X::Array{ℜ,3}, R::Ƶ; μ::ℜ=3.5, γ::ℜ=0.1, EPOCHS::Ƶ=1, ϵ::ℜ=1e-16) where {ℜ<:Real, Ƶ<:Int}
    I::Ƶ, J::Ƶ, K::Ƶ = size(X)

    a::ℜ, b::ℜ = I*J*K*γ, γ/μ
    α_R, α_IR, α_JR, α_KR = fill(a/R,R), fill(a/(I*R),I,R), fill(a/(J*R),J,R), fill(a/(K*R),K,R)

    log_λ = log(rand(Gamma(a,1.0/b)))
    log_θ_R = rand(Dirichlet(α_R .+ 1.0/R))
    log_θ_IR = reshape([log(θ_ir) for r=1:R for θ_ir=rand(Dirichlet(α_IR[:,r] .+ 1.0/I))],I,R)
    log_θ_JR = reshape([log(θ_jr) for r=1:R for θ_jr=rand(Dirichlet(α_JR[:,r] .+ 1.0/J))],J,R)
    log_θ_KR = reshape([log(θ_kr) for r=1:R for θ_kr=rand(Dirichlet(α_KR[:,r] .+ 1.0/K))],K,R)

    s, S_R, S_IR, S_JR, S_KR, S⁺ = zeros(ℜ,R), zeros(ℜ,R), zeros(ℜ,I,R), zeros(ℜ,J,R), zeros(ℜ,K,R), sum(X)
    q, log_ρ, log_ρ_ijk::ℜ = Array{ℜ}(undef,R), Array{ℜ}(undef,R), 0.0

    X_full = similar(X)

    for eph=1:EPOCHS
        S_R .= 0.0
        S_IR .= 0.0
        S_JR .= 0.0
        S_KR .= 0.0
        S⁺ = 0.0

        for k=1:K, j=1:J, i=1:I #order of traversal is important
            log_ρ .= log_λ .+ log_θ_R .+ log_θ_IR[i,:] .+ log_θ_JR[j,:] .+ log_θ_KR[k,:]
            log_ρ_ijk = logsumexp(log_ρ)
            q .= exp.(log_ρ .- log_ρ_ijk)

            X_full[i,j,k] = isnan(X[i,j,k]) ? exp(log_ρ_ijk) : X[i,j,k]
            s .= finucan(X_full[i,j,k],q)

            S_R .+= s
            S_IR[i,:] .+= s
            S_JR[j,:] .+= s
            S_KR[k,:] .+= s
            S⁺ += X_full[i,j,k]
        end
        log_λ = digamma(S⁺+a) - log(b+1.0)
        log_θ_R .= digamma.(S_R.+α_R) .- digamma(S⁺+a)
        log_θ_IR .= digamma.(S_IR.+α_IR) .- digamma.(S_R.+α_R)'
        log_θ_JR .= digamma.(S_JR.+α_JR) .- digamma.(S_R.+α_R)'
        log_θ_KR .= digamma.(S_KR.+α_KR) .- digamma.(S_R.+α_R)'
    end
    return X_full, S⁺, (S_R, S_IR, S_JR, S_KR)
end

end


# Check ChangepointSampler works for different distributions

using Distributions

num_samples = 1000
λ = 50                 # Frequency of changes

println("Running ChangepointSampler tests...")

# Normal mean change sampler
mu = Normal(1.0, 1.0)
σ = 1.0
Y = ChangepointSampler(()->Normal(rand(mu), σ), λ)
rand(Y, num_samples)
@changepoint_sampler num_samples λ Normal(mu, σ)

# Normal variance change sampler
μ = 1.0
σ = Uniform(2.0, 15.0)
Y = ChangepointSampler(()->Normal(μ, rand(σ)), λ)
rand(Y, num_samples)
@changepoint_sampler num_samples λ Normal(μ, σ)

# Exponential change sampler
μ = Uniform(0.0, 10.0)
Y = ChangepointSampler(()->Exponential(rand(μ)), λ)
rand(Y, num_samples)
@changepoint_sampler num_samples λ Exponential(μ)

# Poisson change sampler
μ = Uniform(0.0, 10.0)
Y = ChangepointSampler(()->Poisson(rand(μ)), λ)
rand(Y, num_samples)
@changepoint_sampler num_samples λ Poisson(μ)


using InteractiveChaos, Makie

i = 1
p_index = 1

systems = [(Systems.logistic(), 2.5, 4.0, "r"),
           (Systems.henon(), 0.8, 1.4, "a"),
           (Systems.standardmap(), 0.6, 1.2, "k")]

ds, p_min, p_max, parname = systems[1]

oddata = interactive_orbitdiagram(
                             ds, i, p_index, p_min, p_max;
                             parname = parname, n = 100)

ps, us = scaleod(oddata...)


#=

  Assignment 1:
  https://edu.swi-prolog.org/mod/assign/view.php?id=242
  """
  http://cplint.eu/p/urns_and_balls.swinb

  Urns and balls

  Suppose you have two urns: urn1 contains 40 blue balls and 20 red balls and urn2 contains 25
  blue balls and 30 red balls.

  You throw an unbiased coin and, if it turns out head, you draw a ball from the first urn,
  it it turns out tails you draw a ball from the second urn.

  Write a program modeling this process and a query for answering the question
  "What is the probability of drawing a blue ball?"

  Write the program and the query in the cells below. Then save the notebook and submit the url.

  """
  The exact answer of drawing a blue call is
     0.5*40/60 + 0.5*25/55 = 0.56060606060606060606

  Cf ~/cplint/course_urns_and_balls.pl
     ~/blog/urns_and_balls.blog
     ~/psi/urns_and_balls.blog
     ~/webppl/urns_and_balls.wppl

=#

using Turing, StatsPlots, DataFrames
using ReverseDiff, Zygote, Tracker
# Turing.setadbackend(:reversediff)
# Turing.setadbackend(:zygote)
# Turing.setadbackend(:tracker)
include("jl_utils.jl")

#=
Summary Statistics
  parameters      mean       std   naive_se      mcse          ess      rhat
      Symbol   Float64   Float64    Float64   Float64      Float64   Float64

        coin    1.5008    0.5000     0.0025    0.0025   39377.7877    1.0000
        draw    1.4408    0.4965     0.0025    0.0025   39019.4117    1.0000

Quantiles
  parameters      2.5%     25.0%     50.0%     75.0%     97.5%
      Symbol   Float64   Float64   Float64   Float64   Float64

        coin    1.0000    1.0000    2.0000    2.0000    2.0000
        draw    1.0000    1.0000    1.0000    2.0000    2.0000

length(data): 40001
Dict{Any,Any} with 4 entries:
  ["tail", "blue"] => 8991
  ["head", "red"]  => 6654
  ["tail", "red"]  => 10979
  ["head", "blue"] => 13377
p(tail):0.49923751906202346
p(blue):0.5591860203494913
  2.282226 seconds (8.63 M allocations: 532.511 MiB, 7.70% gc time)

=#

data = []
@model function urns_and_balls()
    tail = 1
    head = 2
    coin ~ Categorical(simplex([0.5,0.5]))

    blue = 1
    red = 2
    draw ~ coin == head ? Categorical(simplex([40,20])) : Categorical(simplex([25,30]))
    push!(data,[
                coin == tail ? "tail" : "head"
                draw == blue ? "blue" : "red"
    ])
end

model = urns_and_balls()

num_chains = 4
# chains = sample(model, Prior(), MCMCThreads(), 1000, num_chains)

# chains = sample(model, MH(), 40_000)
# chains = sample(model, MH(), MCMCThreads(), 40_000, num_chains)

# chains = sample(model, PG(20), MCMCThreads(), 10_000, num_chains)
# chains = sample(model, PG(20), 40_000)
# chains = sample(model, IS(), MCMCThreads(), 10_000, num_chains)
chains = sample(model, IS(), 40_000)

# chains = sample(model, NUTS(1000,0.65), MCMCThreads(), 40_000, num_chains)
# chains = sample(model, Gibbs(MH(:zlabels),NUTS(1000,0.65,:m,:b,:sigma)), MCMCThreads(), 40_000, num_chains)

display(chains)
println("length(data): ", length(data))

h = make_hash(data)
display(h)
# println("data:", make_hash(data))
# println("p(blue): ", (df[!:,2] .== 2)/length(data))

println("p(tail):", sum([d[1] == "tail" for d in data])/length(data))
println("p(blue):", sum([d[2] == "blue" for d in data])/length(data))

# df = DataFrame(chains)


using Random
using Test

function cubic_hermite_test()
    # Need at least two samples:
    xs = ones(1)
    ys = ones(1)
    dydxs = ones(1)
    @test_throws DomainError CubicHermite(xs, ys, dydxs)

    # All vectors must be the same size:
    xs = ones(2)
    xs[2] = 2
    ys = ones(2)
    dydxs = ones(3)
    @test_throws DomainError CubicHermite(xs, ys, dydxs)

    xs = ones(5)
    for i = 2:length(xs)
        xs[i] = i
    end
    ys = ones(5)
    dydxs = zeros(5)
    ch = CubicHermite(xs, ys, dydxs)
    for i = 1:length(xs)-1
        x = Float64(i)
        @test ch(x) == 1.0
        @test gradient(ch, x) == 0.0
        @test hessian(ch, x) == 0.0
        x = x + 0.25
        @test ch(x) == 1.0
        @test gradient(ch, x) == 0.0
        @test hessian(ch, x) == 0.0
        x = x + 0.25
        @test ch(x) == 1.0
        @test gradient(ch, x) == 0.0
        @test hessian(ch, x) == 0.0
        x = x + 0.25
        @test ch(x) == 1.0
        @test gradient(ch, x) == 0.0
        @test hessian(ch, x) == 0.0
    end
    # Ensure that the ight endpoint query doesn't read past the end of the array:
    @test ch(5.0) == 1.0
    @test gradient(ch, 5.0) == 0.0
    @test hessian(ch, 5.0) == 0.0

    # Now linear functions:
    a = 7.2
    b = 9.6
    for i = 1:length(xs)
        ys[i] = a * xs[i] + b
        dydxs[i] = a
    end
    ch = CubicHermite(xs, ys, dydxs)
    for i = 1:length(xs)-1
        x = Float64(i)
        @test ch(x) ≈ a * x + b
        @test gradient(ch, x) ≈ a
        @test abs(hessian(ch, x)) < 3e-14
        x = x + 0.25
        @test ch(x) ≈ a * x + b
        @test gradient(ch, x) ≈ a
        @test abs(hessian(ch, x)) < 3e-14
        x = x + 0.25
        @test ch(x) ≈ a * x + b
        @test gradient(ch, x) ≈ a
        @test abs(hessian(ch, x)) < 3e-14
        x = x + 0.25
        @test ch(x) ≈ a * x + b
        @test gradient(ch, x) ≈ a
        @test abs(hessian(ch, x)) < 3e-14
    end
    @test ch(last(xs)) ≈ a * last(xs) + b
    @test gradient(ch, last(xs)) ≈ a
    @test abs(hessian(ch, last(xs))) < 3e-14

    # Now the interpolation condition:
    xs = zeros(50)
    ys = zeros(50)
    dydxs = zeros(50)
    xs[1] = 0.0
    ys[1] = randn()
    dydxs[1] = randn()
    for i = 2:50
        xs[i] = xs[i-1] + rand() + 0.1
        ys[i] = randn()
        dydxs[i] = randn()
    end

    ch = CubicHermite(xs, ys, dydxs)

    for i = 1:50
        @test ch(xs[i]) ≈ ys[i]
        @test gradient(ch, xs[i]) ≈ dydxs[i]
    end

    # Now quadratics:
    a = 1.0 / 8
    b = -3
    c = -2
    for i = 1:50
        ys[i] = a * xs[i] * xs[i] + b * xs[i] + c
        dydxs[i] = 2 * a * xs[i] + b
    end
    ch = CubicHermite(xs, ys, dydxs)
    for i = 1:200
        x = rand() * last(xs)
        @test ch(x) ≈ a * x * x + b * x + c
        @test gradient(ch, x) ≈ 2 * a * x + b
        @test hessian(ch, x) ≈ 2 * a
    end
    x = last(xs)
    @test ch(x) ≈ a * x * x + b * x + c
    @test gradient(ch, x) ≈ 2 * x * a + b
    @test hessian(ch, x) ≈ 2 * a
    # Cannot extrapolate:
    @test_throws DomainError ch(x + 0.1)
    @test_throws DomainError ch(xs[1] - 0.1)
end

@testset "Hermite" begin
    cubic_hermite_test()
end


function nestlist( f::Function, x0, n::Integer)
	(n ≤ 0) ? [x0] : begin
		list = Vector{typeof(x0)}(undef, n+1)
		list[1] = x0
		for i in 1:n
			list[i+1] = f(list[i])
		end
		list
	end
end

function bin2dec(bvect::BitVector)
    v = 1
	s = 0
    for i in view(bvect,length(bvect):-1:1)
        s += v*i
        v <<= 1
    end
    s
end


using IntervalWavelets
using Test

@testset "Filters" begin
    filter_coefficients = rand(3)
    f = Filter(filter_coefficients, -1:1)

    @test support(f) == -1:1
    @test support_boundaries(f) == (-1, 1)
    @test length(f) == 3

    @test f[-1] == filter_coefficients[1]
    @test f[0] == filter_coefficients[2]
    @test f[1] == filter_coefficients[3]

    @test f[2] == 0.0
    @test f[-2] == 0.0
end

@testset "Interior filters" begin
    p = rand(2:8)

    @testset "Construct interior filters $phase" for phase in [:min, :symmlet]
        h = interior_filter(p, phase)
        @test vanishing_moments(h) == p
        @test length(h) == 2*p
    end

    @testset "Daubechies filters are normalized in l2 $phase" for phase in [:min, :symmlet]
        h = interior_filter(p, phase)
        @test sum(coefficients(h).^2) ≈ 1 atol = sqrt(eps())
    end

    @testset "Failure constructing interior filters" begin
        @test_throws DomainError interior_filter(0, :symmlet)
        @test_throws DomainError interior_filter(9, :symmlet)

        @test_throws MethodError interior_filter(p, :foo)
    end
end

@testset "Boundary filters" begin
    @testset "Construct boundary filter" for s in [IntervalWavelets.LEFT, IntervalWavelets.RIGHT]
        p = 2

        b = boundary_filters(p, s)
        @test side(b) == s
        @test vanishing_moments(b) == p
        @test length(filters(b)) == p
    end

    @testset "Failure constructing boundary filter" begin
        @test_throws DomainError boundary_filters(0, IntervalWavelets.LEFT)
        @test_throws DomainError boundary_filters(9, IntervalWavelets.LEFT)
    end
end


# ---
# title: 1584. Min Cost to Connect All Points
# id: problem1584
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Union Find
# link: <https://leetcode.com/problems/min-cost-to-connect-all-points/description/>
# hidden: true
# ---
#
# You are given an array `points` representing integer coordinates of some
# points on a 2D-plane, where `points[i] = [xi, yi]`.
#
# The cost of connecting two points `[xi, yi]` and `[xj, yj]` is the **manhattan
# distance** between them: `|xi - xj| + |yi - yj|`, where `|val|` denotes the
# absolute value of `val`.
#
# Return  _the minimum cost to make all points connected._ All points are
# connected if there is **exactly one** simple path between any two points.
#
#
#
# **Example 1:**
#
# ![](https://assets.leetcode.com/uploads/2020/08/26/d.png)
#
#
#
#     Input: points = [[0,0],[2,2],[3,10],[5,2],[7,0]]
#     Output: 20
#     Explanation:![](https://assets.leetcode.com/uploads/2020/08/26/c.png)
#     We can connect the points as shown above to get the minimum cost of 20.
#     Notice that there is a unique path between every pair of points.
#
#
# **Example 2:**
#
#
#
#     Input: points = [[3,12],[-2,5],[-4,1]]
#     Output: 18
#
#
# **Example 3:**
#
#
#
#     Input: points = [[0,0],[1,1],[1,0],[-1,1]]
#     Output: 4
#
#
# **Example 4:**
#
#
#
#     Input: points = [[-1000000,-1000000],[1000000,1000000]]
#     Output: 4000000
#
#
# **Example 5:**
#
#
#
#     Input: points = [[0,0]]
#     Output: 0
#
#
#
#
# **Constraints:**
#
#   * `1 <= points.length <= 1000`
#   * `-106 <= xi, yi <= 106`
#   * All pairs `(xi, yi)` are distinct.
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


using Dates

this_year = year(today(Dates.today())
this_month = month(Dates.today())
this_day = day(Dates.today())

function read_ohlc(stock::String, fm::Int, fd::Int, fy::Int, tm::Int, td::Int, ty::Int, period::String)

# take care of yahoo's 0 indexing for month
    fm -= 1
    tm -= 1

    ydata     = readlines(`curl -s "http://ichart.finance.yahoo.com/table.csv?s=$stock&a=$fm&b=$fd&c=$fy&d=$tm&e=$td&f=$ty&g=$period"`)
    numstring = ydata[2:end]
    sa  = split(numstring[1], ',')'

    for i in 2:length(numstring)
        sa  = [sa ; split(numstring[i], ',')']
    end

    time_conversion = map(x -> parse("yyyy-MM-dd", x), convert(Array{UTF16String}, sa[:,1]))

    ohlc = DataFrame(quote
        Date  = $time_conversion
        Open  = float($sa[:,2])
        High  = float($sa[:,3])
        Low   = float($sa[:,4])
        Close = float($sa[:,5])
        Vol   =   int($sa[:,6])
        Adj   = float($sa[:,7])
    end)

    flipud(ohlc)

end

# monthly data

window = 3 # years

monthly_ohlc_over_window(stock::String) = read_ohlc(stock::String, this_month, this_day, this_year - window, this_month,  this_day, this_year, "m")

# TODO make the portfolio list unlimited in length
function returns{T <: Union(ASCIIString, UTF8String)}(list::Array{T})

    str1 = list[1]
    str2 = list[2]
    str3 = list[3]
    str4 = list[4]
    str5 = list[5]

    s1 = monthly_ohlc_over_window(list[1])
    s2 = monthly_ohlc_over_window(list[2])
    s3 = monthly_ohlc_over_window(list[3])
    s4 = monthly_ohlc_over_window(list[4])
    s5 = monthly_ohlc_over_window(list[5])

    r1 = log_return(s1["Adj"])
    r2 = log_return(s2["Adj"])
    r3 = log_return(s3["Adj"])
    r4 = log_return(s4["Adj"])
    r5 = log_return(s5["Adj"])

    ohlc = DataFrame(quote
        $str1 = $r1
        $str2 = $r2
        $str3 = $r3
        $str4 = $r4
        $str5 = $r5
    end)
    ohlc[2:end, :]
end

# function read_monthlies(filename::String)

#     ohlc  = read_table(filename)
#     time_conversion =  Calendar.parse("MMM-yy", ohlc[:,1])

#     within!(ohlc, quote
#         Date = $(time_conversion)
#     end)

# # flipud(df)
# end


"""
```
binarize(Intermodes(), img)
```

Under the assumption that the image histogram is bimodal the histogram is
smoothed using a length-3 mean filter until two modes remain. The binarization
threshold is then set to the average value of the two modes.

# Output

Returns the binarized image as an `Array{Gray{Bool},2}`.

# Arguments

The function argument is described in more detail below.

##  `img`

An `AbstractArray` representing an image. The image is automatically converted
to `Gray` in order to construct the requisite graylevel histogram.

# Example

Binarize the "cameraman" image in the `TestImages` package.

```julia
using TestImages, ImageBinarization

img = testimage("cameraman")
img_binary = binarize(Intermodes(), img)
```

## Reference

C. A. Glasbey, “An Analysis of Histogram-Based Thresholding Algorithms,” CVGIP: Graphical Models and Image Processing, vol. 55, no. 6, pp. 532–537, Nov. 1993. [doi:10.1006/cgip.1993.1040](https://doi.org/10.1006/cgip.1993.1040)
"""
function binarize(algorithm::Intermodes,  img::AbstractArray{T,2}) where T <: Colorant
  img₀₁ = zeros(Gray{Bool}, axes(img))
  edges, counts = build_histogram(img,  256)
  t = find_threshold(HistogramThresholding.Intermodes(), counts[1:end], edges)
  for i in CartesianIndices(img)
    img₀₁[i] = img[i] < t ? 0 : 1
  end
  img₀₁
end


function test_cic( nstep, nbpart)

    @show nstep, nbpart

    nx     = 128
    ny     = 128

    xp  = zeros(Float64, nbpart)
    yp  = zeros(Float64, nbpart)
    vx  = zeros(Float64, (nx+1,ny+1))
    vy  = zeros(Float64, (nx+1,ny+1))

    xmin, xmax = -5.0, 5.0
    ymin, ymax = -5.0, 5.0
    dx = (xmax-xmin)/nx
    dy = (ymax-ymin)/ny
    dt = 0.1

    xp .= rand(nbpart)
    yp .= rand(nbpart)

    for j = 1:ny+1, i = 1:nx+1
        vx[i,j] =   (j-1)*dy + xmin
        vy[i,j] = - (i-1)*dx + ymin
    end

    for istep = 1:nstep

        for k=1:nbpart

           xp[k] = xmin + mod(xp[k]-xmin,xmax - xmin)
           yp[k] = ymin + mod(yp[k]-ymin,ymax - ymin)

           xc = (xp[k] - xmin)/dx
           yc = (yp[k] - ymin)/dy

           i = trunc(Int64,xc)
           j = trunc(Int64,yc)

           dpx = xc - i
           dpy = yc - j

           i += 1
           j += 1

           a1 = (1.0-dpx) * (1.0-dpy)
           a2 = (    dpx) * (1.0-dpy)
           a3 = (    dpx) * (    dpy)
           a4 = (1.0-dpx) * (    dpy)

           vpx = ( a1 * vx[i  ,j  ] + a2 * vx[i+1,j  ]
                 + a3 * vx[i+1,j+1] + a4 * vx[i  ,j+1] )
           vpy = ( a1 * vy[i  ,j  ] + a2 * vy[i+1,j  ]
                 + a3 * vy[i+1,j+1] + a4 * vy[i  ,j+1] )

           xp[k] += vpx * dt
           yp[k] += vpy * dt

        end

        #println(sum(xp)/nbpart, " - ", sum(yp)/nbpart)

    end

end

test_cic( 1, 1) #trigger comoilation
@time test_cic( 1000, 409600)


"""
  Returns the a_j,k and c_k,p Coefficients considering the approximation of the
  exponential term exp(-t_j * z_p) = sum( k : a_j,k * c_k,p )
"""
function get_AC_coefficients(K::Int64
                            , times::Vector
                            , z_p::Vector{Complex{T}}
                            , alpha::Float64
                            , m::Float64
                            ; method = "nfft") where T

  if method == "leastsquare"
    A,C = get_AC_coefficients_lsqr(K,times,z_p)
  elseif method == "nfft"
    A,C = get_AC_coefficients_nfft(K,times,z_p,alpha,m)
  elseif method == "hist"
    A,C = get_AC_coefficients_hist_lsqr(K,times,z_p)
  else
    A = zeros(Complex{T},length(times),K)
    C = zeros(Complex{T},K,length(z_p))
    error("Not implemented yet")
  end
  return A,C
end

function _precompute_coeffs_nfft(K::Int64
                                , times::Vector     # readout times
                                , z_p::Vector{Complex{T}}      # Correction map
                                , alpha::Float64   # oversampling factor of NFFT
                                , m::Float64       # kernel size of NFFT
                                ; adaptK::Bool=true
                                , centering::Bool=true) where T

  # Centering time and correctionterm to lower computational effort
  t_hat = (times[1] + times[end])/2
  z_hat = minimum(real(z_p)) + maximum(real(z_p)) + 1im*(minimum(imag(z_p)) + maximum(imag(z_p)))
  z_hat *= 0.5
  if !centering
    t_hat=0.0
    z_hat=0.0
  end
  # Centering readouttimes and correctionmap
  times = times .- t_hat
  z_p = z_p .- z_hat

  # Calculating timescaling factor so it satisfies: t_j / T in [0.5,0.5)
  t_scale = maximum(abs.(times)) / (0.5-1e-7)

  # determining N and K,
  # K can be choosen freely with tradeoff comp. complexity vs precision
  N = Int64(4*ceil( maximum(abs.(z_p)) * maximum(abs.(times)) / (2*pi) ))
  if adaptK
    K = ceil(Int64, alpha*N + 2*m)
    # println("Calculated K for NFFT: ",K)
  end

  win, win_hat = NFFT.getWindow(:kaiser_bessel)

  return win, win_hat, K, N, t_scale, times, z_p
end

function get_A_coefficients_nfft(K::Int64
                                , times::Vector     # readout times
                                , z_p::Vector{Complex{T}}       # Correction map
                                , alpha::Float64    # oversampling factor(NFFT)
                                , m::Float64        # kernel size(NFFT)
                                ; adaptK::Bool=true
                                , centering::Bool=true) where T

  win, win_hat, K, N, t_scale, times, z_p = _precompute_coeffs_nfft(K, times, z_p, alpha, m; adaptK=adaptK,centering=centering)

  A = zeros(Complex{T}, length(times), K)

  for kappa=1:K
    for j=1:length(times)
        argument = times[j]/t_scale - (kappa - K/2)/(K - 2*m)
        A[j,kappa] = win(argument,N*alpha,m,alpha)
    end
  end
  return A
end

function get_C_coefficients_nfft(K::Int64
                                , times::Vector     # readout times
                                , z_p::Vector{Complex{T}}       # Correction map
                                , alpha::Float64    # oversampling factor(NFFT)
                                , m::Float64       # kernel size(NFFT)
                                ; adaptK::Bool=true
                                , centering::Bool=true) where T

  win, win_hat, K, N, t_scale, times, z_p = _precompute_coeffs_nfft(K, times, z_p, alpha, m; adaptK=adaptK,centering=centering)

  C = zeros(Complex{T}, K, length(z_p))

  for p=1:length(z_p)
    for kappa=1:K
        fourierarg = 1im * z_p[p]*t_scale/(2pi)
        exparg = -z_p[p]*t_scale * ( (kappa - K/2)/(K-2*m) )
        C[kappa,p] = (1/ ( (K-2*m)* 1/(N*alpha)*win_hat(fourierarg,N*alpha,m,alpha))) * exp(exparg)
    end
  end
  return C
end

"""
  Returns the a_j,k and c_k,p Coefficients considering the approximation of the
  exponential term exp(-t_j * z_p) = sum( k : a_j,k * c_k,p )

  Using the nfft method
"""
function get_AC_coefficients_nfft(K::Int64
                                , times::Vector     # readout times
                                , z_p::Vector{Complex{T}}       # Correction map
                                , alpha::Float64   # oversampling factor(NFFT)
                                , m::Float64        # kernel size(NFFT)
                                ; adaptK::Bool=true
                                , centering::Bool=true) where T

  A = get_A_coefficients_nfft(K, times, z_p, alpha, m; adaptK=adaptK, centering=centering)
  C = get_C_coefficients_nfft(K, times, z_p, alpha, m; adaptK=adaptK, centering=centering)
  return A, C
end

"""
  Returns the a_j,k and c_k,p Coefficients considering the approximation of the
  exponential term exp(-t_j * z_p) = sum( k : a_j,k * c_k,p )

  Using the least squares method and time-segmentation
"""
function get_AC_coefficients_lsqr(K::Int64, times::Vector{T}, z_p::Vector{Complex{T}}) where T
  A = zeros(Complex{T},length(times),K)
  C = zeros(Complex{T},K,length(z_p))
  G = zeros(Complex{T},length(z_p),K)

  t_hat = zeros(K)
  t_hat = collect(range(times[1], stop=times[length(times)], length=K))

  z_p = vec(z_p)
  for kappa=1:K
    for p=1:length(z_p)
      G[p,kappa] = exp(-t_hat[kappa]*z_p[p])
      C[kappa,p] = exp(-t_hat[kappa]*z_p[p])
    end
  end

  systemMat = G'*G
  b = [exp(-times[j]*z_p[p]) for p=1:length(z_p), j=1:length(times)]
  At = systemMat \ (G'*b)
  A .= transpose(At)

  return A,C
end

"""
  Returns the a_j,k Coefficients considering the approximation of the
  exponential term exp(-t_j * z_p) = sum( k : a_j,k * c_k,p )

  Using the least squares method with histogramms and time-segmentation
"""
function get_A_coefficients_hist_lsqr(K::Int64, times::Vector, z_p::Vector{Complex{T}}) where T

  numBins = 10*K
  if real.(z_p)== zeros(length(z_p))
    fmax_im = maximum(imag.(z_p))
    fmin_im = minimum(imag.(z_p))
    δf = (fmax_im-fmin_im)/(numBins-1)
    edges_im = [fmin_im-δf/2 + i*δf for i = 0:numBins]
    h = fit(Histogram, imag.(z_p),edges_im,closed=:left)
    z_center = 0.5*1im*vec( [edges_im[i]+edges_im[i+1] for i=1:numBins] )
    z_weights = h.weights
  elseif imag.(z_p)== zeros(length(z_p))
    fmax_re = maximum(real.(z_p))
    fmin_re = minimum(real.(z_p))
    δf = (fmax_re-fmin_re)/(numBins-1)
    edges_re = collect(fmin_re-δf/2.:δf:fmax_re+δf/2)
    h = fit(Histogram, real.(z_p),edges_re,closed=:left)
    z_center = 0.5*vec( [edges_re[i]+edges_re[i+1] for i=1:numBins] )
    z_weights = h.weights
  else
    nbins = ceil(Int, sqrt(numBins))
    fmax_re = maximum(real.(z_p))
    fmin_re = minimum(real.(z_p))
    δf_re = (fmax_re-fmin_re)/(nbins-1)
    edges_re = collect(fmin_re-δf/2.:δf:fmax_re+δf/2)
    fmax_im = maximum(imag.(z_p))
    fmin_im = minimum(imag.(z_p))
    δf = (fmax_im-fmin_im)/(nbins-1)
    edges_im = collect(fmin_im-δf/2.:δf:fmax_im+δf/2)
    h = fit(Histogram, (real.(z_p),imag.(z_p)),(edges_re,edges_im),closed=:left)
    z_center = 0.5*vec( [edges_re[i]+edges_re[i+1]+1im*(edges_im[j]+edges_im[j+1]) for i=1:numBins[1], j=1:numBins[2] ] )
    z_weights = vec(h.weights)
    numBins = nbins^2
  end

  A = zeros(Complex{T},length(times),K)
  G = zeros(Complex{T},numBins,K)

  t_hat = zeros(K)
  t_hat = collect(range(times[1], stop=times[end], length=K))

  z_p = vec(z_p)
  for p=1:numBins
    for kappa=1:K
      G[p,kappa] = exp(-t_hat[kappa]*z_center[p])
    end
  end

  systemMat = G'*Diagonal(z_weights)*G
  b = [z_weights[p]*exp(-times[j]*z_center[p]) for p=1:numBins, j=1:length(times)]
  At = systemMat \ (G'*b)
  A .= transpose(At)

  return A
end

"""
  Returns the c_k,p Coefficients considering the approximation of the
  exponential term exp(-t_j * z_p) = sum( k : a_j,k * c_k,p )

  Using the least squares method with histogramms
"""
function get_C_coefficients_hist_lsqr(K::Int64, times::Vector, z_p::Vector{Complex{T}}) where T

  C = zeros(Complex{T},K,length(z_p))
  t_hat = collect(range(times[1],stop=times[end],length=K))
  # z_p = vec(z_p)
  for p=1:length(z_p)
    for kappa=1:K
      C[kappa,p] = exp(-t_hat[kappa]*z_p[p])
    end
  end

  return C
end

"""
  Returns the a_j,k and c_k,p Coefficients considering the approximation of the
  exponential term exp(-t_j * z_p) = sum( k : a_j,k * c_k,p )

  Using the least squares method and frequency-segmentation
"""
function get_AC_coefficients_lsqr_fs(K::Int64, times::Vector, z_p::Vector{Complex{T}}) where T

  ω = imag.(z_p) # field map
  # ω_hat = collect(range(minimum(ω), stop=maximum(ω),length=K)) # translate frequencies
  ω_hat = loydMax(vec(ω),K, iterations=1000)

  A = zeros(Complex{T},length(times),K)
  G = zeros(Complex{T},length(times),K)

  for kappa=1:K
    for j=1:length(times)
      G[j,kappa] = exp(-1im*times[j]*ω_hat[kappa])
      A[j,kappa] = exp(-1im*times[j]*ω_hat[kappa])
    end
  end

  systemMat = G'*G
  b = [exp(-times[j]*z_p[p]) for j=1:length(times), p=1:length(z_p)]
  C = systemMat \ (G'*b)

  return A,C
end

"""
  Returns the a_j,k Coefficients considering the approximation of the
  exponential term exp(-t_j * z_p) = sum( k : a_j,k * c_k,p )

  Using the least squares method with histogramms and frequency-segmentation
"""
function get_A_coefficients_hist_lsqr_fs(K::Int64, times::Vector, z_p::Vector{Complex{T}}) where T
  ω = imag.(z_p) # field map
  # ω_hat = collect(range(minimum(ω), stop=maximum(ω),length=K)) # translate frequencies
  ω_hat = loydMax(vec(ω),K, iterations=1000)

  A = zeros(Complex{T},length(times),K)
  for kappa=1:K
    for j=1:length(times)
      A[j,kappa] = exp(-1im*times[j]*ω_hat[kappa])
    end
  end

  return A
end

function get_A_coefficients_svd(K::Int64, times::Vector, z_p::Vector{Complex{T}}; numBins=10*K, K_tol::T=1.e-4) where T
  ω = imag.(z_p)

  # create histrogram of field map
  fmax = maximum(ω)
  fmin = minimum(ω)
  δf = (fmax-fmin)/(numBins-1)
  edges = [fmin-δf/2 + i*δf for i = 0:numBins]
  h = fit( Histogram, ω, edges, closed=:left )
  z_center = 0.5*1im * ( edges[1:numBins].+edges[2:numBins+1] )
  z_weights = h.weights

  # from reduced exponential matrix
  A = zeros(Complex{T}, length(times), numBins)
  for j=1:numBins, i=1:length(times)
    A[i,j] = sqrt(z_weights[j])*exp(-z_center[j]*times[i])
  end

  # compute SVD
  U,S,V = psvd(A, rank=2*K)
  # estimate accuracy of the approximation and determine rank threshold
  acc = sqrt.(cumsum(S.^2)./sum(S.^2))
  K0 = findfirst(x->x>=(1-K_tol-1e-14), acc)
  @debug "accuracy: $(acc)"
  @info " K_tol=$(K_tol) => chose K= $(K0)"

  return U[:,1:K0]
end

"""
  Returns the c_k,p Coefficients considering the approximation of the
  exponential term exp(-t_j * z_p) = sum( k : a_j,k * c_k,p )

  Uses the LSQR-method for a given set of a_j,k-coefficients
"""
function get_C_coefficients_lsqr(times::Vector, z_p::Vector{Complex{T}}, A::Matrix{Complex{T}}, numSamp::Int64, step=10) where T

  systemMat = adjoint(A[1:step:numSamp,:])* A[1:step:numSamp,:]
  b = [exp(-times[j]*z_p[i]) for j=1:step:numSamp, i=1:length(z_p)]
  C = systemMat \ ( adjoint(A[1:step:numSamp,:])*b )

  return C
end


@testset "Wigner3j" begin
    w3j_prealloc = zeros(Cdouble, 5)
    for j2 in 0:2, j3 = 0:2, m2 in -j2:j2, m3 in -j3:j3
        jmax = j2 + j3
        for m1 in -jmax:jmax
            w3j, jmin, jmax = Wigner3j(j2, j3, m1, m2, m3)
            Wigner3j!(w3j_prealloc, j2, j3, m1, m2, m3)
            @test w3j_prealloc[1:(jmax - jmin + 1)] == w3j
        end
    end

    @testset "known values" begin
        for j in 0:20, m in -j:j
            w3j, jmin, jmax = Wigner3j(j, 0, m, -m, 0)
            @test length(w3j) == 1
            @test w3j[1] ≈ (-1)^(j-m)/√(2j+1)
        end
    end
end


"""
    Gamma(α,θ)

The *Gamma distribution* with shape parameter `α` and scale `θ` has probability density
function

```math
f(x; \\alpha, \\theta) = \\frac{x^{\\alpha-1} e^{-x/\\theta}}{\\Gamma(\\alpha) \\theta^\\alpha},
\\quad x > 0
```

```julia
Gamma()          # Gamma distribution with unit shape and unit scale, i.e. Gamma(1, 1)
Gamma(α)         # Gamma distribution with shape α and unit scale, i.e. Gamma(α, 1)
Gamma(α, θ)      # Gamma distribution with shape α and scale θ

params(d)        # Get the parameters, i.e. (α, θ)
shape(d)         # Get the shape parameter, i.e. α
scale(d)         # Get the scale parameter, i.e. θ
```

External links

* [Gamma distribution on Wikipedia](http://en.wikipedia.org/wiki/Gamma_distribution)

"""
struct Gamma{T<:Real} <: ContinuousUnivariateDistribution
    α::T
    θ::T
    Gamma{T}(α, θ) where {T} = new{T}(α, θ)
end

function Gamma(α::T, θ::T; check_args::Bool=true) where {T <: Real}
    @check_args Gamma (α, α > zero(α)) (θ, θ > zero(θ))
    return Gamma{T}(α, θ)
end

Gamma(α::Real, θ::Real; check_args::Bool=true) = Gamma(promote(α, θ)...; check_args=check_args)
Gamma(α::Integer, θ::Integer; check_args::Bool=true) = Gamma(float(α), float(θ); check_args=check_args)
Gamma(α::Real; check_args::Bool=true) = Gamma(α, one(α); check_args=check_args)
Gamma() = Gamma{Float64}(1.0, 1.0)

@distr_support Gamma 0.0 Inf

#### Conversions
convert(::Type{Gamma{T}}, α::S, θ::S) where {T <: Real, S <: Real} = Gamma(T(α), T(θ))
Base.convert(::Type{Gamma{T}}, d::Gamma) where {T<:Real} = Gamma{T}(T(d.α), T(d.θ))
Base.convert(::Type{Gamma{T}}, d::Gamma{T}) where {T<:Real} = d

#### Parameters

shape(d::Gamma) = d.α
scale(d::Gamma) = d.θ
rate(d::Gamma) = 1 / d.θ

params(d::Gamma) = (d.α, d.θ)
partype(::Gamma{T}) where {T} = T

#### Statistics

mean(d::Gamma) = d.α * d.θ

var(d::Gamma) = d.α * d.θ^2

skewness(d::Gamma) = 2 / sqrt(d.α)

kurtosis(d::Gamma) = 6 / d.α

function mode(d::Gamma)
    (α, θ) = params(d)
    α >= 1 ? θ * (α - 1) : error("Gamma has no mode when shape < 1")
end

function entropy(d::Gamma)
    (α, θ) = params(d)
    α + loggamma(α) + (1 - α) * digamma(α) + log(θ)
end

mgf(d::Gamma, t::Real) = (1 - t * d.θ)^(-d.α)

cf(d::Gamma, t::Real) = (1 - im * t * d.θ)^(-d.α)

function kldivergence(p::Gamma, q::Gamma)
    # We use the parametrization with the scale θ
    αp, θp = params(p)
    αq, θq = params(q)
    θp_over_θq = θp / θq
    return (αp - αq) * digamma(αp) - loggamma(αp) + loggamma(αq) -
        αq * log(θp_over_θq) + αp * (θp_over_θq - 1)
end

#### Evaluation & Sampling

@_delegate_statsfuns Gamma gamma α θ

gradlogpdf(d::Gamma{T}, x::Real) where {T<:Real} =
    insupport(Gamma, x) ? (d.α - 1) / x - 1 / d.θ : zero(T)

function rand(rng::AbstractRNG, d::Gamma)
    if shape(d) < 1.0
        # TODO: shape(d) = 0.5 : use scaled chisq
        return rand(rng, GammaIPSampler(d))
    elseif shape(d) == 1.0
        return rand(rng, Exponential(d.θ))
    else
        return rand(rng, GammaGDSampler(d))
    end
end

function sampler(d::Gamma)
    if shape(d) < 1.0
        # TODO: shape(d) = 0.5 : use scaled chisq
        return GammaIPSampler(d)
    elseif shape(d) == 1.0
        return sampler(Exponential(d.θ))
    else
        return GammaGDSampler(d)
    end
end

#### Fit model

struct GammaStats <: SufficientStats
    sx::Float64      # (weighted) sum of x
    slogx::Float64   # (weighted) sum of log(x)
    tw::Float64      # total sample weight

    GammaStats(sx::Real, slogx::Real, tw::Real) = new(sx, slogx, tw)
end

function suffstats(::Type{<:Gamma}, x::AbstractArray{T}) where T<:Real
    sx = zero(T)
    slogx = zero(T)
    for xi = x
        sx += xi
        slogx += log(xi)
    end
    GammaStats(sx, slogx, length(x))
end

function suffstats(::Type{<:Gamma}, x::AbstractArray{T}, w::AbstractArray{Float64}) where T<:Real
    n = length(x)
    if length(w) != n
        throw(DimensionMismatch("Inconsistent argument dimensions."))
    end

    sx = zero(T)
    slogx = zero(T)
    tw = zero(T)
    for i = 1:n
        @inbounds xi = x[i]
        @inbounds wi = w[i]
        sx += wi * xi
        slogx += wi * log(xi)
        tw += wi
    end
    GammaStats(sx, slogx, tw)
end

function gamma_mle_update(logmx::Float64, mlogx::Float64, a::Float64)
    ia = 1 / a
    z = ia + (mlogx - logmx + log(a) - digamma(a)) / (abs2(a) * (ia - trigamma(a)))
    1 / z
end

function fit_mle(::Type{<:Gamma}, ss::GammaStats;
    alpha0::Float64=NaN, maxiter::Int=1000, tol::Float64=1e-16)

    mx = ss.sx / ss.tw
    logmx = log(mx)
    mlogx = ss.slogx / ss.tw

    a::Float64 = isnan(alpha0) ? (logmx - mlogx)/2 : alpha0
    converged = false

    t = 0
    while !converged && t < maxiter
        t += 1
        a_old = a
        a = gamma_mle_update(logmx, mlogx, a)
        converged = abs(a - a_old) <= tol
    end

    Gamma(a, mx / a)
end


include("molzahn.jl")

function minimize_sliwak(cliques::AbstractVector, cliquetree::AbstractVector, alpha, beta, m, nb_lc)
    costs = [sliwak_heuristic(cliques, edge, alpha, beta, m, nb_lc) for edge in cliquetree]
    index = argmin(costs)
    return cliquetree[index]
end

function sliwak_heuristic(cliques::AbstractVector, edge, alpha, beta, m, nb_lc)
    clique1 = cliques[edge[1]]
    clique2 = cliques[edge[2]]
    size_intersect = length(intersect(clique1, clique2))
    size_clique1 = length(clique1)
    size_clique2 = length(clique2)
    size_fusion_clique = size_clique1 + size_clique2 - size_intersect
    nb_lc_new = nb_lc - size_intersect * (2 * size_intersect + 1)
    q2 = (m + nb_lc_new)^3
    cost = alpha * (- size_clique1^3 - size_clique2^3 + size_fusion_clique^3)
           + beta * q2
    return cost
end

function merge_sliwak(cliques::AbstractVector, cliquetree::AbstractVector;
                      alpha, beta, m, nb_lc,
                      treshold::Union{Int,Function}=treshold_percent(cliques, 0.1),
                      kmax=3, kwargs...)
    edges = Set()
    if typeof(treshold) <: Int
        treshold_value = treshold
        treshold = x -> length(x) >= treshold_value
    end
    k = 0
    while treshold(cliques)
        edge = k <= kmax ?
                    minimize_sliwak(cliques, cliquetree, alpha, beta, nb_lc, m) :
                    minimize_molzahn(cliques, cliquetree)
        size_intersect = length(intersect(cliques[edge[1]], cliques[edge[2]]))
        nb_lc -= size_intersect * (2 * size_intersect + 1)
        edges = vcat(edges..., get_added_edges(cliques[edge[1]], cliques[edge[2]])...)
        cliques, cliquetree = merge_cliques(cliques, cliquetree, edge)
        k += 1
    end
    return cliques, cliquetree, edges
end


using Distances
using NearestNeighbors
using StaticArrays

export NeighborList, neighborlist

struct NeighborList
    j::Vector{Vector{Int}}
    R::Vector{Vector{Float64}}
    r::Vector{Vector{SVector{3,Float64}}}
end
Base.length(nn::NeighborList) = length(nn.j)

function get_distance(L::SVector{3,<:AbstractFloat}, x::SVector{3,<:AbstractFloat}, y::SVector{3,<:AbstractFloat})
    broadcast(L, x, y) do Li, xi, yi
        if Li == Inf
            xi - yi
        else
            d = mod(xi - yi, Li)
            2d < Li ? d : d - Li
        end
    end
end

function neighborlist(A::AbstractSystem{3}, rcutoff::Float64)
    # Convert Positions to Matrix for Ball tree
    X = [SVector{3}(ustrip.(p)) for p ∈ position(A)]

    # Create Metric for Periodic Boundary Conditions
    periodic = periodicity(A)
    L = @SVector [periodic[i] ? ustrip(bounding_box(A)[i][i]) : Inf for i = 1:3]
    d = Distances.PeriodicEuclidean(L)

    # Build Ball tree
    tree = BallTree(X, d)

    # Intialize empty vectors
    j = Vector{Vector{Int64}}(undef, length(X))              # Neighbors
    R = Vector{Vector{Float64}}(undef, length(X))            # Distances
    r = Vector{Vector{SVector{3,Float64}}}(undef, length(X)) # Positions

    # Fill vectors
    for n in 1:length(X)
        neighbors = filter(m -> m > n, inrange(tree, X[n], rcutoff, true))
        j[n] = zeros(Int64, length(neighbors))
        R[n] = zeros(Float64, length(neighbors))
        r[n] = zeros(SVector{3,Float64}, length(neighbors))
        for (i, m) in enumerate(neighbors)
            rr = get_distance(L, X[n], X[m])
            j[n][i] = m
            R[n][i] = norm(rr)
            r[n][i] = rr
        end
    end
    NeighborList(j, R, r)
end