using NeuralPDE
using Quadrature, Cubature, Cuba
using Flux, ModelingToolkit, GalacticOptim, Optim, DiffEqFlux
using Plots
using DelimitedFiles
using QuasiMonteCarlo
import ModelingToolkit: Interval, infimum, supremum

function nernst_planck(strategy, minimizer, maxIters)

    ##  DECLARATIONS
    @parameters t x y z
    @variables c(..)

    Dt = Differential(t)
    Dx = Differential(x)
    Dy = Differential(y)
    Dz = Differential(z)
    Dxx = Differential(x)^2
    Dyy = Differential(y)^2
    Dzz = Differential(z)^2

    ## DOMAINS AND OPERATORS

    # Discretization
    xwidth      = 1.0
    ywidth      = 1.0
    zwidth      = 1.0
    tmax        = 1.0
    xMeshNum    = 10
    yMeshNum    = 10
    zMeshNum    = 10
    tMeshNum    = 10

    dx  = xwidth/xMeshNum
    dy  = ywidth/yMeshNum
    dz  = zwidth/zMeshNum
    dt  = tmax/tMeshNum

    domains = [t ∈ Interval(0.0,tmax),
               x ∈ Interval(0.0,xwidth),
               y ∈ Interval(0.0,ywidth),
               z ∈ Interval(0.0,zwidth)]

    xs = 0.0 : dx : xwidth
    ys = 0.0 : dy : ywidth
    zs = 0.0 : dz : zwidth
    ts = 0.0 : dt : tmax

    # Constants
    D = 1  #dummy
    ux = 10 #dummy
    uy = 10 #dummy
    uz = 10 #dummy

    # Operators
    div = - D*(Dxx(c(t,x,y,z)) + Dyy(c(t,x,y,z)) + Dzz(c(t,x,y,z)))
          + (ux*Dx(c(t,x,y,z)) + uy*Dy(c(t,x,y,z)) + uz*Dz(c(t,x,y,z)))

    # Equation
    eq = Dt(c(t,x,y,z)) + div ~ 0      #NERNST-PLANCK EQUATION

    # Boundary conditions
    bcs = [c(0,x,y,z) ~ 0]

    ## NEURAL NETWORK
    n = 16   #neuron number

    chain = FastChain(FastDense(4,n,Flux.σ),FastDense(n,n,Flux.σ),FastDense(n,1))   #Neural network from Flux library

    indvars = [t,x,y,z]   #independent variables
    depvars = [c(t,x,y,z)]       #dependent (target) variable

    dim = length(domains)

    losses = []
    error = []
    times = []

    dx_err = 0.2

    error_strategy = GridTraining(dx_err)

    initθ = Float64.(DiffEqFlux.initial_params(chain))
    eltypeθ = eltype(initθ)
    parameterless_type_θ = DiffEqBase.parameterless_type(initθ)

    phi = NeuralPDE.get_phi(chain,parameterless_type_θ)
    derivative = NeuralPDE.get_numeric_derivative()

    _pde_loss_function = NeuralPDE.build_loss_function(eq,indvars,depvars,
                                             phi,derivative,nothing,chain,initθ,error_strategy)

    bc_indvars = NeuralPDE.get_variables(bcs,indvars,depvars)
    _bc_loss_functions = [NeuralPDE.build_loss_function(bc,indvars,depvars,
                                              phi,derivative,nothing,chain,initθ,error_strategy,
                                              bc_indvars = bc_indvar) for (bc,bc_indvar) in zip(bcs,bc_indvars)]

    train_sets = NeuralPDE.generate_training_sets(domains,dx_err,[eq],bcs,eltypeθ,indvars,depvars)
    train_domain_set, train_bound_set = train_sets

    pde_loss_functions = [NeuralPDE.get_loss_function(_pde_loss_function,
                                          train_domain_set[1],eltypeθ,
                                          parameterless_type_θ,error_strategy)]

    bc_loss_functions = [NeuralPDE.get_loss_function(bc,
                                          train_set,eltypeθ,
                                          parameterless_type_θ,error_strategy) for (train_set,bc) in zip(train_bound_set, _bc_loss_functions)]
    loss_functions =  [pde_loss_functions;bc_loss_functions]
    function loss_function_(θ,p)
        return sum(map(l->l(θ) ,loss_functions))
    end

    cb_ = function (p,l)
        deltaT_s = time_ns() #Start a clock when the callback begins, this will evaluate questo misurerà anche il calcolo degli uniform error

        ctime = time_ns() - startTime - timeCounter #This variable is the time to use for the time benchmark plot
        append!(times, ctime/10^9) #Conversion nanosec to seconds
        append!(losses, l)
        loss_ = loss_function_(p,nothing)
        append!(error, loss_)

        timeCounter = timeCounter + time_ns() - deltaT_s #timeCounter sums all delays due to the callback functions of the previous iterations

        return false
    end

    @named pde_system = PDESystem(eq, bcs, domains, indvars, depvars)

    discretization = NeuralPDE.PhysicsInformedNN(chain,strategy; init_params =initθ)
    prob = NeuralPDE.discretize(pde_system,discretization)

    timeCounter = 0.0
    startTime = time_ns() #Fix initial time (t=0) before starting the training
    res = GalacticOptim.solve(prob, minimizer, cb = cb_, maxiters=maxIters)

    phi = discretization.phi

    params = res.minimizer

    # Model prediction
    domain = [ts, xs, ys, zs]

    u_predict  = [reshape([phi([t,x,y,z],res.minimizer) for x in xs for y in ys for z in zs],
                 (length(xs),length(ys),length(zs))) for t in ts]

    return [error, params, domain, times]
end

maxIters = [(1,1,1,1000,1000,1000,1000),(1,1,1,300,300,300,300)] #iters for ADAM/LBFGS
# maxIters = [(1,1,1,10,10,10,10),(1,1,1,3,3,3,3)] #iters for ADAM/LBFGS

strategies = [NeuralPDE.QuadratureTraining(quadrature_alg = CubaCuhre(), reltol = 1e-4, abstol = 1e-4, maxiters = 50),
              NeuralPDE.QuadratureTraining(quadrature_alg = HCubatureJL(), reltol = 1e-4, abstol = 1e-4, maxiters = 50, batch = 0),
              NeuralPDE.QuadratureTraining(quadrature_alg = CubatureJLh(), reltol = 1e-4, abstol = 1e-4, maxiters = 50),
              NeuralPDE.QuadratureTraining(quadrature_alg = CubatureJLp(), reltol = 1e-4, abstol = 1e-4, maxiters = 50),
              NeuralPDE.GridTraining(0.2),
              NeuralPDE.StochasticTraining(400 ; bcs_points= 50),
              NeuralPDE.QuasiRandomTraining(400 ; bcs_points= 50)]

strategies_short_name = ["CubaCuhre",
                        "HCubatureJL",
                        "CubatureJLh",
                        "CubatureJLp",
                        "GridTraining",
                        "StochasticTraining",
                        "QuasiRandomTraining"]

minimizers = [ADAM(0.005),
              #BFGS()]
              LBFGS()]

minimizers_short_name = ["ADAM",
                         "LBFGS"]
                        # "BFGS"]

# Run models
error_res =  Dict()
domains = Dict()
params_res = Dict()  #to use same params for the next run
times = Dict()

## Convergence

for strat=1:length(strategies) # strategy
      for min =1:length(minimizers) # minimizer
            # println(string(strategies_short_name[strat], "  ", minimizers_short_name[min]))
            res = nernst_planck(strategies[strat], minimizers[min], maxIters[min][strat])
            push!(error_res, string(strat,min)     => res[1])
            push!(params_res, string(strat,min) => res[2])
            push!(domains, string(strat,min)        => res[3])
            push!(times, string(strat,min)        => res[4])
      end
end

#Plotting the first strategy with the first minimizer out from the loop to initialize the canvas
current_label = string(strategies_short_name[1], " + " , minimizers_short_name[1])
error = Plots.plot(times["11"], error_res["11"], yaxis=:log10, label = current_label)#, xlims = (0,10))#legend = true)#, size=(1200,700))
plot!(error, times["21"], error_res["21"], yaxis=:log10, label = string(strategies_short_name[2], " + " , minimizers_short_name[1]))
plot!(error, times["31"], error_res["31"], yaxis=:log10, label = string(strategies_short_name[3], " + " , minimizers_short_name[1]))
plot!(error, times["41"], error_res["41"], yaxis=:log10, label = string(strategies_short_name[4], " + " , minimizers_short_name[1]))
plot!(error, times["51"], error_res["51"], yaxis=:log10, label = string(strategies_short_name[5], " + " , minimizers_short_name[1]))
plot!(error, times["61"], error_res["61"], yaxis=:log10, label = string(strategies_short_name[6], " + " , minimizers_short_name[1]))
plot!(error, times["71"], error_res["71"], yaxis=:log10, label = string(strategies_short_name[7], " + " , minimizers_short_name[1]))

plot!(error, times["12"], error_res["12"], yaxis=:log10, label = string(strategies_short_name[1], " + " , minimizers_short_name[2]))
plot!(error, times["22"], error_res["22"], yaxis=:log10, label = string(strategies_short_name[2], " + " , minimizers_short_name[2]))
plot!(error, times["32"], error_res["32"], yaxis=:log10, label = string(strategies_short_name[3], " + " , minimizers_short_name[2]))
plot!(error, times["42"], error_res["42"], yaxis=:log10, label = string(strategies_short_name[4], " + " , minimizers_short_name[2]))
plot!(error, times["52"], error_res["52"], yaxis=:log10, label = string(strategies_short_name[5], " + " , minimizers_short_name[2]))
plot!(error, times["62"], error_res["62"], yaxis=:log10, label = string(strategies_short_name[6], " + " , minimizers_short_name[2]))
plot!(error, times["72"], error_res["72"], yaxis=:log10, title = string("Nernst Planck convergence ADAM/LBFGS"), ylabel = "log(error)", xlabel = "t",label = string(strategies_short_name[7], " + " , minimizers_short_name[2]))

using SciMLBenchmarks
SciMLBenchmarks.bench_footer(WEAVE_ARGS[:folder],WEAVE_ARGS[:file])


using StochasticSearch, FactCheck, Base.Test

facts("[Search]") do
    context("optimize and simulated_annealing") do
        function rosenbrock(x::Configuration, parameters::Dict{Symbol, Any} = Dict{Symbol, Any}())
            return (1.0 - x["i0"].value)^2 + 100.0 * (x["i1"].value - x["i0"].value^2)^2
        end
        configuration = Configuration([FloatParameter(-2.0,2.0,0.0,"i0"),
                                       FloatParameter(-2.0,2.0,0.0,"i1")],
                                       "rosenbrock_config")

        tuning_run = Run(cost               = rosenbrock,
                         starting_point     = configuration,
                         methods            = [[:simulated_annealing 2];],
                         stopping_criterion = elapsed_time_criterion,
                         duration           = 2,
                         measurement_method = sequential_measure_mean!)

        search_task = @task optimize(tuning_run)
        result = consume(search_task)
        print(result)
        while result.is_final == false
            result = consume(search_task)
            print(result)
        end
        rr = rosenbrock(result.minimum)
        rc = result.cost_minimum
        @test_approx_eq rc rr
        @fact (rosenbrock(result.minimum) <= rosenbrock(configuration))   --> true
        @fact_throws Exception optimize(rosenbrock, configuration, [:bozo_search])
        println(rosenbrock(result.minimum))
    end
end


"""
```julia
    LogisticSoftMaxLikelihood()
```

The multiclass likelihood with a logistic-softmax mapping: :
```math
p(y=i|{fₖ}₁ᴷ) = σ(fᵢ)/∑ₖ σ(fₖ)
```
where `σ` is the logistic function.
This likelihood has the same properties as [softmax](https://en.wikipedia.org/wiki/Softmax_function).
---

For the analytical version, the likelihood is augmented multiple times. More details can be found in the paper [Multi-Class Gaussian Process Classification Made Conjugate: Efficient Inference via Data Augmentation](https://arxiv.org/abs/1905.09670)
"""
struct LogisticSoftMaxLikelihood{T<:Real} <: MultiClassLikelihood{T}
    Y::AbstractVector{BitVector} #Mapping from instances to classes
    class_mapping::AbstractVector{Any} # Classes labels mapping
    ind_mapping::Dict{Any,Int} # Mapping from label to index
    y_class::AbstractVector{Int64} # GP Index for each sample
    c::AbstractVector{AbstractVector{T}} # Second moment of fₖ
    α::AbstractVector{T} # Variational parameter of Gamma distribution
    β::AbstractVector{T} # Variational parameter of Gamma distribution
    θ::AbstractVector{AbstractVector{T}} # Variational parameter of Polya-Gamma distribution
    γ::AbstractVector{AbstractVector{T}} # Variational parameter of Poisson distribution
    function LogisticSoftMaxLikelihood{T}() where {T<:Real}
        new{T}()
    end
    function LogisticSoftMaxLikelihood{T}(Y::AbstractVector{<:BitVector},
    class_mapping::AbstractVector, ind_mapping::Dict{<:Any,<:Int},y_class::AbstractVector{<:Int}) where {T<:Real}
        new{T}(Y,class_mapping,ind_mapping,y_class)
    end
    function LogisticSoftMaxLikelihood{T}(Y::AbstractVector{<:BitVector},
    class_mapping::AbstractVector, ind_mapping::Dict{<:Any,<:Int},y_class::AbstractVector{<:Int},
    c::AbstractVector{<:AbstractVector{<:Real}}, α::AbstractVector{<:Real},
    β::AbstractVector, θ::AbstractVector{<:AbstractVector},γ::AbstractVector{<:AbstractVector{<:Real}}) where {T<:Real}
        new{T}(Y,class_mapping,ind_mapping,y_class,c,α,β,θ,γ)
    end
end

function LogisticSoftMaxLikelihood()
    LogisticSoftMaxLikelihood{Float64}()
end

function pdf(l::LogisticSoftMaxLikelihood,f::AbstractVector)
    logisticsoftmax(f)
end

function pdf(l::LogisticSoftMaxLikelihood,y::Integer,f::AbstractVector)
    logisticsoftmax(f)[y]
end

function Base.show(io::IO,model::LogisticSoftMaxLikelihood{T}) where T
    print(io,"Logistic-Softmax Likelihood")
end

function init_likelihood(likelihood::LogisticSoftMaxLikelihood{T},inference::Inference{T},nLatent::Integer,nSamplesUsed::Integer,nFeatures::Integer) where T
    if inference isa AnalyticVI || inference isa GibbsSampling
        c = [ones(T,nSamplesUsed) for i in 1:nLatent]
        α = nLatent*ones(T,nSamplesUsed)
        β = nLatent*ones(T,nSamplesUsed)
        θ = [abs.(rand(T,nSamplesUsed))*2 for i in 1:nLatent]
        γ = [abs.(rand(T,nSamplesUsed)) for i in 1:nLatent]
        LogisticSoftMaxLikelihood{T}(likelihood.Y,likelihood.class_mapping,likelihood.ind_mapping,likelihood.y_class,c,α,β,θ,γ)
    else
        return likelihood
    end
end

## Local Updates##
function local_updates!(l::LogisticSoftMaxLikelihood,y,μ::NTuple{N,<:AbstractVector},Σ::NTuple{N,<:AbstractVector}) where {T,N}
    l.c .= broadcast((Σ,μ)->sqrt.(Σ+abs2.(μ)),Σ,μ)
    for _ in 1:2
        l.γ .= broadcast(
            (β,c,μ,ψα)->0.5 / β * exp.(ψα) .* safe_expcosh.(-0.5*μ, 0.5*c),                                    [l.β],l.c,μ,[digamma.(l.α)])
        l.α .= 1.0.+(l.γ...)
    end
    l.θ .= broadcast((y,γ,c)->0.5*(y+γ)./c.*tanh.(0.5.*c),
                                    y,l.γ,l.c)
    return nothing
end

function sample_local!(l::LogisticSoftMaxLikelihood{T},y::AbstractVector,f) where {T}
    l.γ .= broadcast(f->rand.(Poisson.(0.5*l.α.*safe_expcosh.(-0.5*f,0.5*f))), f)
    l.α .= rand.(Gamma.(one(T).+(l.γ...),1.0./l.β))
    pg = PolyaGammaDist()
    set_ω!(l,broadcast((y,γ,f)->draw.([pg],y.+γ,μ),y,l.γ,f))
    return nothing
end

## Global Gradient Section ##

@inline ∇E_μ(l::LogisticSoftMaxLikelihood,::AOptimizer,y::AbstractVector) where {T} = 0.5.*(y.-l.γ)
@inline ∇E_Σ(l::LogisticSoftMaxLikelihood,::AOptimizer,y::AbstractVector) where {T} = 0.5.*l.θ

## ELBO Section ##
function expec_log_likelihood(l::LogisticSoftMaxLikelihood{T},i::AnalyticVI,y,μ,Σ) where {T}
    tot = -length(y)*logtwo
    tot += -sum(sum(l.γ.+y))*logtwo
    tot +=  0.5*sum(broadcast((θ,γ,y,μ,Σ)->dot(μ,(y-γ))-dot(θ,abs2.(μ))-dot(θ,Σ),l.θ,l.γ,y,μ,Σ))
    return tot
end

AugmentedKL(l::LogisticSoftMaxLikelihood,y::AbstractVector) = PolyaGammaKL(l,y) + PoissonKL(l) + GammaEntropy(l)

PolyaGammaKL(l::LogisticSoftMaxLikelihood,y) =  sum(broadcast(PolyaGammaKL,y.+l.γ,l.c,l.θ))

PoissonKL(l::LogisticSoftMaxLikelihood) =     sum(broadcast(PoissonKL,l.γ,[l.α./l.β],[digamma.(l.α).-log.(l.β)]))

##  Compute the equivalent of KL divergence between an improper prior p(λ) (``1_{[0,\\infty]}``) and a variational Gamma distribution ##
function GammaEntropy(l::LogisticSoftMaxLikelihood)
    return -sum(l.α)+sum(log,l.β[1])-sum(lgamma,l.α)-dot(1.0.-l.α,digamma.(l.α))
end

## Numerical Gradient Section ##

function grad_samples(model::AbstractGP{T,<:LogisticSoftMaxLikelihood,<:NumericalVI},samples::AbstractMatrix{T},index::Integer) where {T}
    class = model.likelihood.y_class[index]::Int64
    grad_μ = zeros(T,model.nLatent)
    grad_Σ = zeros(T,model.nLatent)
    g_μ = similar(grad_μ)
    nSamples = size(samples,1)
    @views @inbounds for i in 1:nSamples
        σ = logistic.(samples[i,:])
        samples[i,:]  .= logisticsoftmax(samples[i,:])
        s = samples[i,class]
        g_μ .= grad_logisticsoftmax(samples[i,:],σ,class)/s
        grad_μ += g_μ
        grad_Σ += diaghessian_logisticsoftmax(samples[i,:],σ,class)/s - abs2.(g_μ)
    end
    for k in 1:model.nLatent
        model.inference.ν[k][index] = -grad_μ[k]/nSamples
        model.inference.λ[k][index] = grad_Σ[k]/nSamples
    end
end

function log_like_samples(model::AbstractGP{T,<:LogisticSoftMaxLikelihood},samples::AbstractMatrix,index::Integer) where {T}
    class = model.likelihood.y_class[index]
    nSamples = size(samples,1)
    loglike = zero(T)
    for i in 1:nSamples
        σ = logistic.(samples[i,:])
        loglike += log(σ[class])-log(sum(σ))
    end
    return loglike/nSamples
end

function grad_logisticsoftmax(s::AbstractVector{<:Real},σ::AbstractVector{<:Real},i::Integer)
    s[i]*(δ.(i,eachindex(σ)).-s).*(1.0.-σ)
end

function diaghessian_logisticsoftmax(s::AbstractVector{<:Real},σ::AbstractVector{<:Real},i::Integer)
    s[i]*(1.0.-σ).*(
    abs2.(δ.(i,eachindex(σ))-s).*(1.0.-σ)
    -s.*(1.0.-s).*(1.0.-σ)
    -σ.*(δ.(i,eachindex(σ))-s))
end

function hessian_logisticsoftmax(s::AbstractVector{T},σ::AbstractVector{T},i::Integer) where {T<:Real}
    m = length(s)
    hessian = zeros(T,m,m)
    @inbounds for j in 1:m
        for k in 1:m
            hessian[j,k] = (1-σ[j])*s[i]*(
            (δ(i,k)-s[k])*(1.0-σ[k])*(δ(i,j)-s[j])
            -s[j]*(δ(j,k)-s[k])*(1.0-σ[k])
            -δ(k,j)*σ[j]*(δ(i,j)-s[j]))
        end
    end
    return hessian
end


# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
# ==============================================================================
#
#   Compute the nutation and the Equation of the Origins (EO) as in
#   equinox-based IAU-2006 theory.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# References
# ==============================================================================
#
#   [1] Vallado, D. A (2013). Fundamentals of Astrodynamics and Applications.
#       Microcosm Press, Hawthorn, CA, USA.
#
#   [2] ftp://tai.bipm.org/iers/conv2010/chapter5/tab5.2e.txt
#
#   [3] Wallace, P. T., Capitaine, N (2006). Precession-nutation procedures
#       consistent with IAU 2006 resolutions. Astronomy & Astrophysics.
#
#   [4] Capitaine, N., Wallace, P. T (2006). High precision methods for locating
#       the celestial intermediate pole and origin. Astronomy & Astrophysics.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

export mean_obliquity_iau2006, nutation_eo_iau2006

"""
    mean_obliquity_iau2006(jd_tt::Number)

Compute the mean obliquity of the ecliptic [rad] using the equinox-based
IAU-2006 theory in the Julian day `jd_tt` [Terrestiral Time].

The algorithm was obtained in **[1]**.

# Reference

- **[1]**: Wallace, P. T., Capitaine, N (2006). Precession-nutation procedures
    consistent with IAU 2006 resolutions. Astronomy & Astrophysics.
"""
function mean_obliquity_iau2006(jd_tt::Number)
    # Compute the Julian Centuries from `jd_tt`.
    t_tt = (jd_tt - JD_J2000) / 36525

    # Auxiliary variables
    # ===================

    a2d = 1 / 3600
    d2r = π / 180
    a2r = a2d * d2r

    # Mean obliquity of the ecliptic
    # ==============================

    # Compute the mean obliquity of the ecliptic [s].

    # NOTE: This equation is wrong in [1](p. 216, eq. 3-68)!
    # The one used here was obtained in [3].
    mϵ_2000 = @evalpoly(
        t_tt,
        +84381.406,
        -46.836769,
        -0.0001831,
        +0.00200340,
        -0.000000576,
        -0.0000000434
    )

    # Reduce to the interval [0, 2π].
    mϵ_2000 = mod(mϵ_2000 * a2r, 2π)

    return mϵ_2000
end

"""
    nutation_eo_iau2006(jd_tt::Number)

Compute the nutation parameters and the Equation of Origins (EO) at the Julian
Day `jd_tt` [TT] using the equinox-based 2006 IAU Theory of Nutation. Notice
that one can provide corrections for the nutation in obliquity (`δΔϵ_2000`)
[rad] and in longitude (`δΔψ_2000`) [rad] that are usually obtained from IERS
EOP Data (see [`get_iers_eop`](@ref)).

# Returns

- The mean obliquity of the ecliptic [rad].
- The nutation in obliquity of the ecliptic [rad].
- The nutation in longitude [rad].
- The Equation of Origins (EO) [rad].
"""
function nutation_eo_iau2006(
    jd_tt::Number,
    δΔϵ_2000::Number = 0,
    δΔΨ_2000::Number = 0
)
    # Compute the Julian Centuries from `jd_tt`.
    t_tt = (jd_tt - JD_J2000) / 36525

    # Auxiliary variables
    # ===================

    a2d = 1 / 3600
    d2r = π / 180
    a2r = a2d * d2r
    r2d = 180 / π
    d2a = 3600
    r2a = r2d * d2a

    # Fundamental arguments
    # =====================

    # Luni-solar part.
    M_s, M_m, u_Mm, D_s, Ω_m = luni_solar_args_iau2006(jd_tt)

    # Planetary part.
    λ_M☿, λ_M♀, λ_Me, λ_M♂, λ_M♃, λ_M♄, λ_M⛢, λ_M♆, p_λ = planetary_args_iau2006(jd_tt)

    # Mean obliquity of the ecliptic
    # ==============================

    mϵ_2000 = mean_obliquity_iau2006(jd_tt)

    # Nutation in the obliquity
    # ==========================================================================

    Δϵ_2000 = _iau2006_sum(
        (
            COEFS_IAU2006_NUT_OBL0,
            COEFS_IAU2006_NUT_OBL1
        ),
        t_tt,
        M_s,
        M_m,
        u_Mm,
        D_s,
        Ω_m,
        λ_M☿,
        λ_M♀,
        λ_Me,
        λ_M♂,
        λ_M♃,
        λ_M♄,
        λ_M⛢,
        λ_M♆,
        p_λ
    )

    # Apply the correction given EOP data.
    Δϵ_2000 += δΔϵ_2000 * r2a

    # Nutation in longitude
    # ==========================================================================

    ΔΨ_2000 = _iau2006_sum(
        (
            COEFS_IAU2006_NUT_LON0,
            COEFS_IAU2006_NUT_LON1
        ),
        t_tt,
        M_s,
        M_m,
        u_Mm,
        D_s,
        Ω_m,
        λ_M☿,
        λ_M♀,
        λ_Me,
        λ_M♂,
        λ_M♃,
        λ_M♄,
        λ_M⛢,
        λ_M♆,
        p_λ
    )

    # Apply the correction given EOP data.
    ΔΨ_2000 += δΔΨ_2000 * r2a

    # Equation of Origins (EO)
    # ==========================================================================

    ΔEO = _iau2006_sum(
        (
            COEFS_IAU2006_EO_0,
            COEFS_IAU2006_EO_1
        ),
        t_tt,
        M_s,
        M_m,
        u_Mm,
        D_s,
        Ω_m,
        λ_M☿,
        λ_M♀,
        λ_Me,
        λ_M♂,
        λ_M♃,
        λ_M♄,
        λ_M⛢,
        λ_M♆,
        p_λ
    )

    # NOTE: This equation is wrong in [1](p. 216, eq. 3-66)! The coefficients
    # were obtained from [2] and [4](eq. 69).
    EO = @evalpoly(
        t_tt,
        -0.014_506,
        -4612.156_534,
        -1.391_581_7,
        +0.000_000_44,
        +0.000_029_956,
        +3.68e-8
    ) - ΔΨ_2000 * cos(mϵ_2000) - ΔEO

    # Convert to [rad].
    EO *= a2r
    Δϵ_2000 *= a2r
    ΔΨ_2000 *= a2r

    return mϵ_2000, Δϵ_2000, ΔΨ_2000, EO
end


# Archimedian spirals
using Luxor

spiraldata = [
  (-2, "Lituus",      130),
  (-1, "Hyperbolic", 180),
  ( 1, "Archimedes",   1),
  ( 2, "Fermat",       7)]

grid = GridRect(O - (200, 0), 130, 50)
@png begin
    background("grey15")
    for aspiral in spiraldata
        @layer begin
            sethue("antiquewhite")
            setcolor(sethue("antiquewhite"))
            setline(.5)
            Luxor.translate(nextgridpoint(grid))
            spiral(last(aspiral), first(aspiral), period=20π, :stroke)
            strokepath()
            fontsize(16)
            label(aspiral[2], :S, offset=80)
        end
    end
end


module KMeans
  using ..KNeighbors
  using ...Statistics

  function runKMeans(centroids,dataset,k)
    assignment = 1:k
    return KNeighbors.KNearestNeighbors(centroids,dataset,assignment)
  end

  function trainKmeans(training_set,k,seed)
    changed = true;
    #size is k by n
    centroids = initCentroids(training_set,k,seed)
    assignment = 1:k
    y=zeros(Int64,size(training_set)[1],1)
    index = 0;
    while(changed)
      index+=1
      y=KNeighbors.KNearestNeighbors(centroids,training_set,assignment)
      newCentroids = recalculateCentroids(centroids,training_set,y);
      if centroids==newCentroids
        changed=false
      end
      centroids = newCentroids
    end
    return centroids,y;
  end
  #
  function initCentroids(training_set,k,seed)
    indices = []
    srand(seed)
    indices=randperm(size(training_set)[1])[1:k]
    return training_set[indices, :]
  end

  function recalculateCentroids(centroids,training_set, y)
    newCentroids = zeros(Float64,size(centroids))
    for i=1:size(centroids)[1]
      newCentroids[i,:] = Statistics.meanColumn(training_set[find(x->x==i,y),:])
    end
    return newCentroids
  end

end


G = 6.673*10^-11
# Insira os valores da massa e do raio do planeta o qual iremos calcular a velocidade de escape
M = 3.2*10^20
R = 5*10^5
v = sqrt((2*G*M)/R)


### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# ╔═╡ 831a3d08-b4f3-11eb-101b-a19a95d2b62e
begin
	using CSV
	using DataFrames
	using LightGraphs
	using GLM
	using GraphPlot
	using ParserCombinator.Parsers.DOT: nodes, edges, parse_dot, StringID
    using Statistics
end

# ╔═╡ 1d7b6d76-8974-4480-942d-35e895ac0c69

# ╔═╡ cb7406e4-ae75-4c99-b7d5-842055b52aee
df = CSV.read("./data/sodium.csv", DataFrame)
df = select!(df, Not(:Column1))   # drop index column

# ╔═╡ 3c9d639a-dea3-40f5-bec0-be46ecf332e3
begin
    ols = lm(@formula(sbp ~ sodium + age + proteinuria), df)

    xt = copy(df[:,Not(:sbp)])
    xt₁ = copy(xt)
    xt₁.sodium .= 1
    xt₂ = copy(xt)
    xt₂.sodium .= 0

    ate_est = mean(predict(ols, xt₁) - predict(ols, xt₂))
    println("ATE estimate: $ate_est")
end

begin
    ols_causal = lm(@formula(sbp ~ sodium + age), df)

    xt = copy(df[:,Not([:sbp, :proteinuria])])
    xt₁ = copy(xt)
    xt₁.sodium .= 1
    xt₂ = copy(xt)
    xt₂.sodium .= 0

    causal_ate_est = mean(predict(ols_causal, xt₁) - predict(ols_causal, xt₂))
    println("causal ATE estimate: $causal_ate_est")
end

# ╔═╡ 48203bec-8156-489a-b04d-9a613526c85e
function _dot_read_one_graph(pg)
    isdir = pg.directed
    nvg = length(nodes(pg))
    nodedict = Dict(zip(collect(nodes(pg)), 1:nvg))
    if isdir
        g = LightGraphs.DiGraph(nvg)
    else
        g = LightGraphs.Graph(nvg)
    end
    for es in edges(pg)
        s = nodedict[es[1]]
        d = nodedict[es[2]]
        add_edge!(g, s, d)
    end
    return g, collect(nodes(pg))
end

# ╔═╡ b1abedc3-647e-4f64-ad78-c4b0eeff8510
function loaddot(gdot::String, gname::String)
    p = parse_dot(gdot)
    for pg in p
        isdir = pg.directed
        possname = isdir ? StringID("digraph") : StringID("graph")
        name = get(pg.id, possname).id
        name == gname && return _dot_read_one_graph(pg)
    end
    error("Graph $gname not found")
end

# ╔═╡ 7e9da9c2-ac34-4ec9-a6ad-f81d67435c5b
gio = "digraph graphname {
    age -> sodium;
	age -> sbp;
	sodium -> proteinuria;
	sodium -> sbp;
    sbp -> proteinuria;
}"

# ╔═╡ abb14b1e-75fb-4a3f-9973-c6ec4fdb2722

# ╔═╡ 8302c22b-816c-43cf-9d06-60a9daaa56fb
graph, nodelabels = loaddot(gio, "graphname")

# ╔═╡ 4ae575f1-a224-4b98-8142-207a1c9b0094
gplot(graph, layout=circular_layout, nodelabel=nodelabels)

# ╔═╡ ae162157-5768-4935-a1ee-b720fb2d1dc7
begin
	Xt = df[:, ["sodium", "age", "proteinuria"]]
	y = df[!, "sbp"]
end

# ╔═╡ 2aa55dbe-4776-46df-9107-8bd64a545ee4

# ╔═╡ f7a19ddb-391e-49e0-9912-990edc76e92b

# ╔═╡ 41ed3b5e-cdb5-4478-86ce-138f08c09784

# ╔═╡ Cell order:
# ╠═831a3d08-b4f3-11eb-101b-a19a95d2b62e
# ╠═1d7b6d76-8974-4480-942d-35e895ac0c69
# ╠═cb7406e4-ae75-4c99-b7d5-842055b52aee
# ╠═3c9d639a-dea3-40f5-bec0-be46ecf332e3
# ╠═b1abedc3-647e-4f64-ad78-c4b0eeff8510
# ╠═48203bec-8156-489a-b04d-9a613526c85e
# ╠═7e9da9c2-ac34-4ec9-a6ad-f81d67435c5b
# ╠═abb14b1e-75fb-4a3f-9973-c6ec4fdb2722
# ╠═8302c22b-816c-43cf-9d06-60a9daaa56fb
# ╠═4ae575f1-a224-4b98-8142-207a1c9b0094
# ╠═ae162157-5768-4935-a1ee-b720fb2d1dc7
# ╠═2aa55dbe-4776-46df-9107-8bd64a545ee4
# ╠═f7a19ddb-391e-49e0-9912-990edc76e92b
# ╠═41ed3b5e-cdb5-4478-86ce-138f08c09784


struct SampleSize{T<:Real} <: Distributions.DiscreteUnivariateDistribution

  design::Design
  p::T

  function SampleSize{T}(design::Design, p::T) where {T<:Real}

    @checkprob p
    new(design, p)

  end # inner constructor

end # SampleSize

SampleSize(design::Design, p::T) where {T<:Real} = SampleSize{T}(design, p)

design(ss::SampleSize) = ss.design

Base.show(io::IO, d::SampleSize) = print("SampleSize")

Base.size(::SampleSize) = ()

Base.getindex(ss::SampleSize, i) = ss

rand(d::SampleSize) = d.design |> interimsamplesize |> n1 -> Distributions.Binomial(n1, d.p) |> rand |> x1 -> samplesize(d.design, x1)

function pdf(d::SampleSize, n::Int)

  res = 0.0
  for x1 in 0:interimsamplesize(d.design)
    if samplesize(d.design, x1) == n
      res += Distributions.pdf(Distributions.Binomial(interimsamplesize(d.design), d.p), x1)
    end
  end
  return res

end # pdf

minimum(d::SampleSize) = interimsamplesize(d.design)

maximum(d::SampleSize) = maximum(samplesize(d.design))

function quantile(d::SampleSize, q::Real)

    nrange = minimum(d):maximum(d)
    res    = 0
    while Distributions.cdf(d, res) <= q
        res += 1
    end
    return res

end # quantile

function mean(d::SampleSize)

    res = 0.0
    for n in minimum(d):maximum(d)
        res += pdf(d, n)*n
    end
    return(res)

end # mean

function var(d::SampleSize)

    res = 0.0
    m   = mean(d)
    for n in minimum(d):maximum(d)
        res += pdf(d, n)*(n - m)^2
    end
    return(res)

end # var


function getDCResistivitySourcesAndRecAll(n,Mesh;srcSpacing::Vector=[1,1],srcPad::Vector=[4,4],
	                                                               wPad=2,offset=12,offsetRec=1)

	offset    = 12
	offsetRec = 1
	i1  =  srcPad[1]+offset:srcSpacing[1]:n[1]-srcPad[2]-offset
	i2  =  srcPad[3]+offset:srcSpacing[2]:n[2]-srcPad[4]-offset
	nsrc = 2*length(i1)*length(i2)

	println("number of sources =",nsrc)
	nrec = 4
	# generate pFors for x1 dipoles
	Wd   = Array{Array}(nsrc)
	Sources   = spzeros(prod(Mesh.n+1),nsrc)

	# Sources
	cnt = 1
	# x-direction source dipoles
	for i=i1
		for j=i2
        	# source
			S                  = zeros(tuple(n+1...));
			S[i-offset,j,1]  = -1;
        	S[i+offset,j,1]  =  1;
			Sources[:,cnt]     = sparse(vec(S))
			cnt +=1
		end
	end

	# y direction dipoles sources
	for i=i1
		for j=i2
	        	# source
			S              = zeros(tuple(n+1...));
			S[i,j-offset,1]     = -1;
	      S[i,j+offset,1]   =  1;
	      Sources[:,cnt] = sparse(vec(S))
			cnt +=1
		end
	end

	###### Receivers
	# x-direction receiver dipoles
	i1  =  srcPad[1]+offsetRec:1:n[1]-srcPad[2]-offsetRec
	i2  =  srcPad[3]+offsetRec:1:n[2]-srcPad[4]-offsetRec

	Receivers = spzeros(prod(Mesh.n+1),2*length(i1)*length(i2))

	cnt = 1
	for i=i1
		for j=i2
        	# source
			S                  = zeros(tuple(n+1...));
			S[i-offsetRec,j,1]  = -1;
        	S[i+offsetRec,j,1]  =  1;
			Receivers[:,cnt]     = sparse(vec(S))
			cnt +=1

		end
	end

	# y direction dipoles
	for i=i1
		for j=i2
	        	# source
			S              = zeros(tuple(n+1...));
			S[i,j-offsetRec,1]     = -1;
	      S[i,j+offsetRec,1]   =  1;
	      Receivers[:,cnt] = sparse(vec(S))
			cnt +=1

		end
	end

	Wd = ones(nsrc,nsrc)*1e-3

	return Sources,Receivers,Wd
end


"""
Derivatives-related functions
"""

## Derivative types
abstract type AbstractDiffMethod end

struct ForwardAD <: AbstractDiffMethod end
struct ReverseAD <: AbstractDiffMethod end
struct RevZyg <: AbstractDiffMethod end  # only used for gradients (not jacobians)
struct ForwardFD <: AbstractDiffMethod end
struct CentralFD <: AbstractDiffMethod end
struct ComplexStep <: AbstractDiffMethod end
struct UserDeriv <: AbstractDiffMethod end   # user-specified derivatives

FD = Union{ForwardFD, CentralFD, ComplexStep}

"""
convert to type used in FiniteDiff package
"""
function finitediff_type(dtype)
    if isa(dtype, ForwardFD)
        fdtype = Val{:forward}
    elseif isa(dtype, CentralFD)
        fdtype = Val{:central}
    elseif isa(dtype, ComplexStep)
        fdtype = Val{:complex}
    end
    return fdtype
end

# internally used cache for gradients and jacobians (separately)
struct GradOrJacCache{T1,T2,T3,T4}
    f!::T1
    work::T2
    cache::T3
    dtype::T4
end

function gradientcache(dtype::FD, func!, nx, ng)

    df = zeros(nx)
    x = zeros(nx)
    fdtype = finitediff_type(dtype)
    cache = FiniteDiff.GradientCache(df, x, fdtype)

    return GradOrJacCache(func!, nothing, cache, dtype)
end

function gradient!(df, x, cache::GradOrJacCache{T1,T2,T3,T4}
    where {T1,T2,T3,T4<:FD})

    FiniteDiff.finite_difference_gradient!(df, cache.f!, x, cache.cache)

    return nothing
end


using Revise
using ADSeismic
using ADCME
using PyPlot
using DelimitedFiles
# matplotlib.use("Agg")
close("all")
if has_gpu()
  use_gpu()
  gpu = true
else
  gpu = false
end

output_dir = "data/acoustic"
if !ispath(output_dir)
  mkpath(output_dir)
end
figure_dir = "figure/acoustic/marmousi/"
if !ispath(figure_dir)
  mkpath(figure_dir)
end
result_dir = "result/acoustic/marmousi/"
if !ispath(result_dir)
  mkpath(result_dir)
end
# if isfile(joinpath(result_dir, "loss.txt"))
#   rm(joinpath(result_dir, "loss.txt"))
# end

################### Generate synthetic data #####################
reset_default_graph()

ap_sim = load_acoustic_model("models/marmousi2-model-true.mat", IT_DISPLAY=0)
src = load_acoustic_source("models/marmousi2-model-true.mat")
rcv_sim = load_acoustic_receiver("models/marmousi2-model-true.mat")
## For debug, only run the first source
# src = [src[1]]
# rcv_sim = [rcv_sim[1]]

if gpu
  Rs_ = compute_forward_GPU(ap_sim, src, rcv_sim)
else
  [SimulatedObservation!(ap_sim(src[i]), rcv_sim[i]) for i = 1:length(src)]
  Rs_ = [rcv_sim[i].rcvv for i = 1:length(rcv_sim)]
end

sess = Session(); init(sess)

Rs = run(sess, Rs_)

for i = 1:length(src)
    writedlm(joinpath(output_dir, "marmousi-r$i.txt"), Rs[i])
end

## visualize_wavefield if needed
# u = run(sess, ap_sim(src[div(length(src),2)+1]).u)
# visualize_wavefield(u, ap_sim.param)

################### Inversion using Automatic Differentiation #####################
reset_default_graph()

ap_sim = load_acoustic_model("models/marmousi2-model-smooth.mat"; inv_vp=true, IT_DISPLAY=0)
src = load_acoustic_source("models/marmousi2-model-true.mat")
rcv_sim = load_acoustic_receiver("models/marmousi2-model-true.mat")
## For debug, only run the first source
# src = [src[1]]
# rcv_sim = [rcv_sim[1]]

Rs = Array{Array{Float64,2}}(undef, length(src))
for i = 1:length(src)
    Rs[i] = readdlm(joinpath(output_dir, "marmousi-r$i.txt"))
end

if gpu
  losses, gs = compute_loss_and_grads_GPU(ap_sim, src, rcv_sim, Rs, get_collection()[1])
  g = sum(gs)
  loss = sum(losses)
else
  [SimulatedObservation!(ap_sim(src[i]), rcv_sim[i]) for i = 1:length(src)]
  loss = sum([sum((rcv_sim[i].rcvv-Rs[i])^2) for i = 1:length(rcv_sim)])
end

sess = Session(); init(sess)

vp = (run(sess, ap_sim.vp.contents)).^0.5./1e3
scale = mean(vp)
vmax = maximum(vp)
vmin = minimum(vp)

function callback(vs, iter, loss)
  if iter%10==0
    if ndims(vs) == 1
      vp = Array(reshape(vs[1:(ap_sim.param.NY+2)*(ap_sim.param.NX+2)], ap_sim.param.NY+2, ap_sim.param.NX+2))
    else
      vp = vs[1]'
    end
    clf()
    pcolormesh([0:ap_sim.param.NX+1;]*ap_sim.param.DELTAX/1e3,[0:ap_sim.param.NY+1;]*ap_sim.param.DELTAY/1e3,  vp*scale)
    axis("scaled")
    colorbar(shrink=0.4)
    xlabel("x (km)")
    ylabel("z (km)")
    gca().invert_yaxis()
    title("Iteration = $iter")
    savefig(joinpath(figure_dir, "inv_$(lpad(iter,5,"0")).png"), bbox_inches="tight")
    writedlm(joinpath(result_dir, "inv_$(lpad(iter,5,"0")).txt"), vp*scale)

    open(joinpath(result_dir, "loss.txt"), "a") do io
      writedlm(io, loss)
    end
  else
  end
end

if gpu
  LBFGS!(sess, loss, g, get_collection()[1]; callback=callback)
else
  LBFGS!(sess, loss, vars=[get_collection()[1]], callback=callback)
end


using CUDex

SIZE = 100000

x32 = DexArray(rand(Float32,SIZE))
x64 = DexArray(rand(Float64,SIZE))
s32 = rand(Float32)
s64 = rand(Float64)

# sqrt
@time sqrt(x32)
@time sqrt(x64)

# exp
@time exp(x32)
@time exp(x64)

# log
@time log(x32)
@time log(x64)

# tanh
@time tanh(x32)
@time tanh(x64)

# .*
@time x32.*s32
@time x64.*s64


using StanModels, CSV

df = CSV.read(joinpath(@__DIR__, "..", "..", "data", "WaffleDivorce.csv"), delim=';')
mean_ma = mean(df[!, :Marriage])
df[!, :Marriage_s] = convert(Vector{Float64},
  (df[!, :Marriage]) .- mean_ma)/std(df[!, :Marriage]);

mean_mam = mean(df[!, :MedianAgeMarriage])
df[!, :MedianAgeMarriage_s] = convert(Vector{Float64},
  (df[!, :MedianAgeMarriage]) .- mean_mam)/std(df[!, :MedianAgeMarriage]);

df[1:6, [1, 7, 14, 15]]

rethinking_data = "
    Location Divorce  Marriage.s MedianAgeMarriage.s
1    Alabama    12.7  0.02264406          -0.6062895
2     Alaska    12.5  1.54980162          -0.6866993
3    Arizona    10.8  0.04897436          -0.2042408
4   Arkansas    13.5  1.65512283          -1.4103870
5 California     8.0 -0.26698927           0.5998567
"

# Define the Stan language model

m5_3s = "
data {
  int N;
  vector[N] divorce;
  vector[N] marriage_z;
  vector[N] median_age_z;
}
parameters {
  real a;
  real bA;
  real bM;
  real<lower=0> sigma;
}
model {
  vector[N] mu = a + median_age_z * bA + marriage_z * bM;
  sigma ~ uniform( 0 , 10 );
  bA ~ normal( 0 , 1 );
  bM ~ normal( 0 , 1 );
  a ~ normal( 10 , 10 );
  divorce ~ normal( mu , sigma );
}
";

# Define the Stanmodel and set the output format to :mcmcchains.

sm = SampleModel("m5.3s", m5_3s);

# Input data for cmdstan

m5_3_data = Dict("N" => size(df, 1), "divorce" => df[!, :Divorce],
    "marriage_z" => df[!, :Marriage_s], "median_age_z" => df[!, :MedianAgeMarriage_s]);

# Sample using cmdstan

(sample_file, log_file) = stan_sample(sm, data=m5_3_data);

# Rethinking results

rethinking_results = "
       mean   sd  5.5% 94.5% n_eff Rhat
a      9.69 0.22  9.34 10.03  1313    1
bR    -0.12 0.30 -0.60  0.35   932    1
bA    -1.13 0.29 -1.56 -0.67   994    1
sigma  1.53 0.16  1.28  1.80  1121    1
"

if !(sample_file == nothing)
  chn = read_samples(sm)
  describe(chn)
end


import Base.Iterators: flatten
using Catlab.Graphics.Graphviz
import Catlab.Graphics.Graphviz: Graph, Edge

graph_attrs = Attributes(:rankdir=>"LR")
node_attrs  = Attributes(:shape=>"plain", :style=>"filled", :color=>"white")
edge_attrs  = Attributes(:splines=>"splines")

function edgify(δ, transition, reverse::Bool)
    return [Edge(reverse ? ["T_$transition", "S_$k"] :
                           ["S_$k", "T_$transition"],
                 Attributes(:label=>"$(δ[k])", :labelfontsize=>"6"))
             for k in collect(keys(δ)) if δ[k] != 0]
end

"""
    Graph(model::Model)

convert a Model into a GraphViz Graph. Transition are green boxes and states are blue circles. Arrows go from the input states to the output states for each transition.
"""
function Graph(model::Model)
    ks = collect(keys(model.Δ))
    statenodes = [Node(string("S_$s"), Attributes(:shape=>"circle", :color=>"#6C9AC3")) for s in model.S]
    transnodes = [Node(string("T_$k"), Attributes(:shape=>"square", :color=>"#E28F41")) for k in ks]

    stmts = vcat(statenodes, transnodes)
    edges = map(ks) do k
      vcat(edgify(first(model.Δ[k]), k, false),
           edgify(last(model.Δ[k]), k, true))
    end |> flatten |> collect
    stmts = vcat(stmts, edges)
    g = Graphviz.Digraph("G", stmts; graph_attrs=graph_attrs, node_attrs=node_attrs, edge_attrs=edge_attrs)
    return g
end


struct LevelSetToFractions{M, T, J, C}
    mesh::M
    topology::T
    jac::J
    colors::C
end
function LevelSetToFractions(mesh::Mesh)
    topo = topology(mesh)
    sparsity = get_sparsity_pattern(topo)
    jac = float.(sparsity)
    colors = matrix_colors(jac)
    return LevelSetToFractions(mesh, topo, jac, colors)
end
function Base.show(io::IO, ls::LevelSetToFractions)
    nelems = nelements(ls.mesh)
    nnodes = nvertices(ls.mesh)
    println(io, "Level set function: node values -> element volume fractions")
    println(io, "  #nodes = $nnodes \n  #elems = $nelems")
end

function get_sparsity_pattern(topo)
    I = Int[]
    J = Int[]
    V = Bool[]
    for (ei, e) in enumerate(elements(topo))
        for vi in indices(e)
            push!(I, ei)
            push!(J, vi)
            push!(V, true)
        end
    end
    return sparse(I, J, V, nelements(topo), nvertices(topo))
end

function (ls::LevelSetToFractions)(fracs::AbstractVector, C::AbstractVector)
    return volume_fractions!(fracs, ls.topology, ls.mesh, C)
end
function (ls::LevelSetToFractions)(C::AbstractVector)
    fracs = similar(C, nelements(ls.topology))
    return volume_fractions!(fracs, ls.topology, ls.mesh, C)
end
function ChainRulesCore.rrule(ls::LevelSetToFractions, C::AbstractVector)
    forwarddiff_color_jacobian!(ls.jac, ls, C, colorvec = ls.colors)
    return ls(C), Δ -> begin
        (NoTangent(), ls.jac' * Δ)
    end
end


"""
    ZN{N} <: DAS
ZN symmetry singleton type where N is an Int.
"""
struct ZN{N} <: DAS end
Base.show(io::IO, ::ZN{N}) where N = print(io,"Z",N)

"""
    ZNCharges{N} <: DASCharges
singleton-type for collections of `ZNCharge{N}` objects. Only the type parameter
is provided - `ZNCharges{N}` doesn't have a field.
Upon iteration yields `ZNCharge{N}(i)` for `i = 0,...,N-1`

# Example
```julia-repl
julia>a = ZNCharges{2}()
julia>foreach(println, a)
ZNCharge{2}(0)
ZNCharge{2}(1)
```
"""
struct ZNCharges{N} <: DASCharges end
Base.show(io::IO, s::ZNCharges{N}) where N = print(io,"Z",N,"Charges(",0:(N-1),")")

⊕(::ZNCharges{N}, ::ZNCharges{N}) where N = ZNCharges{N}()
Base.inv(::ZNCharges{N}) where N = ZNCharges{N}()
Base.length(::ZNCharges{N}) where N = N
Base.getindex(s::ZNCharges{N}, i::Int) where N = ZNCharge{N}(i-1)
chargestype(::ZN{N}) where N = ZNCharges{N}
Base.eltype(::Union{ZNCharges{N},Type{ZNCharges{N}}}) where N = ZNCharge{N}
Base.intersect(::ZNCharges{N}, ::ZNCharges{N}) where N = ZNCharges{N}()
Base.intersect(::ZNCharges, ::ZNCharges) = ArgumentError("cannot intersect different Z charges")

"""
    ZNCharge{N} <: DASCharge
holds the charge of a ZN symmetry as an integer.
The integer is taken `mod N` s.t. the charge is always between `0` and `N-1`

# Example
```julia-repl
julia>a = ZNCharge{2}(1);
julia>a ⊕ a
ZNCharge{2}(0)

julia>a = ZNCharge{2}(3) == ZNCharge{2}(1)
true
```
"""
struct ZNCharge{N} <: DASCharge
    ch::Int
    ZNCharge{N}(a) where {N} = new(mod(a,N))
end
Base.show(io::IO, s::ZNCharge{N}) where N = print(io,"Z",N,"Charge(",s.ch,")")

⊕(a::ZNCharge{N}, b::ZNCharge{N}) where N = ZNCharge{N}(a.ch + b.ch)
function chargeindex(ch::ZNCharge{N}, chs::ZNCharges{N}) where N
    ch in chs || throw(ArgumentError(string(ch, " not in ", chs)))
    (ch.ch + 1)
end
chargetype(::ZN{N}) where N = ZNCharge{N}
Base.zero(::Type{ZNCharge{N}}) where N = ZNCharge{N}(zero(Int))
Base.:+(chs::ZNCharges{N}, ch::ZNCharge{N}) where N = ZNCharges{N}()

const Z2 = ZN{2}
const Z2Charges = ZNCharges{2}
const Z2Charge = ZNCharge{2}


"""
This module contains:

  - Data types:
    + `SparseMatrixCSR`:Compressed Sparse Row (CSR) sparse matrix implementation with Bi-based indexing.
    + `SymSparseMatrixCSR`: Symmetric Compressed Sparse Row sparse matrix implementation with Bi-based indexing.

  - Procedures:
    + `push_coo!`: Helper function to build COO arrays for further building a SparseMatrix
    + `finalize_coo!`: Finalization of COO arrays building.
    + `sparsecsr`: Analogous method to [`sparse`](@ref) function to build a SparseMatrixCSR.
    + `symsparsecsr`: Analogous method to [`sparse`](@ref) function to build a SymSparseMatrixCSR.
    + `colvals`: Analogous method to [`rowvals`](@ref) to return `colvals` array.
    + `hasrowmajororder`: Return `true` if matrix values are ordered by row.
    + `hascolmajororder`: Return `true` if matrix values are ordered by column.
    + `getptr`: Return the pointer array of a SparseMatrix (`rowptr` or `colptr` depending on the SparseMatrix type)
    + `getindices`: Return the indices array of a SparseMatrix (`rowval` or `colval` depending on the SparseMatrix type)

  - Overloaded procedures:
    + `*`: SparseMatrix-Vector product.
    + `mul!`: SparseMatrix-Vector product.
    + `nnz`: Return the number of stored (filled) elements in a sparse array.
    + `nonzeros`: Return `nzval` array.
    + `nzrange`: Return the range of indices for a particular row or column (Depending on the SparseMatrix type)
    + `findnz`: Return a tuple (I, J, V) where I and J are the row and column indices.
    + `rowvals`: Return row indices or raises an error (Depending on the SparseMatrix type)
    + `convert`: Type conversion

"""
module SparseMatricesCSR

using SparseArrays
using LinearAlgebra

import Base: convert, size, getindex, show, count, *
import LinearAlgebra: mul!
import SparseArrays: nnz, getnzval, nonzeros, nzrange, findnz, rowvals

export SparseMatrixCSR
export SymSparseMatrixCSR
export push_coo!, finalize_coo!, sparsecsr, symsparsecsr, colvals
export hasrowmajororder, hascolmajororder, getptr, getindices

include("SparseMatrixCSC.jl")
include("SparseMatrixCSR.jl")
include("SymSparseMatrixCSR.jl")

end # module


function zad39(A,el1)
    a,b = size(A)
    s = 0
    V = []
    for i in 1:a
        for j in 1:b
         if i == j && A[i,j] > el1
            s +=A[i,j]
         end
         if rem(i,2)==1 && rem(A[i,j],el1) == 0
            push!(V,A[i,j])
         end
        end
    end

    return V, s
end


import Catalyst
import Catalyst: @reaction_network
import ModelingToolkit
import ModelingToolkit: FnType

"""
    gene_expression_system(; mean_s=50, mean_x=mean_s, corr_time_s=1.0, corr_time_x=0.1, u0=SA[mean_s, mean_x], dtimes=0:0.1:2.0)

Creates a system for a very simple model of gene expression.

# Model Description

This model consists of four reactions:
```
                κ
reaction 1:  ∅ ---> S

                λ
reaction 2:  S ---> ∅

                ρ
reaction 3:  S ---> S + X

                μ
reaction 4:  X ---> ∅
```

The first two reactions specify the evolution of the input signal `S`,
and last two reactions specify the evolution of the output `X`. Thus,
both the input and output signal are modeled as a simple birth-death
process, however the birth rate of `X` increases with higher copy numbers
of  `S`.

# Examples

The values of the reaction rates can be specified directly as follows:
```jldoctest
using PathWeightSampling
system = PathWeightSampling.gene_expression_system(kappa = 10.0, lambda = 0.1, rho = 1.0, mu = 1.0)

# output

SimpleSystem with 4 reactions
Input variables: S(t)
Output variables: X(t)
Initial condition:
    S(t) = 50
    X(t) = 50
Parameters:
    κ = 10.0
    λ = 0.1
    ρ = 1.0
    μ = 1.0
```

Alternatively, the reaction rates can be specified indirectly through the following arguments:
- `mean_s`: The average copy number of S when the system has relaxed to steady state.
- `mean_x`: The average copy number of X when the system has relaxed to steady state.
- `corr_time_s`: The input signal correlation time. The shorter this time is, the faster the fluctuations in the input signal.
- `corr_time_x`: The output signal correlation time. This sets the timescale of output fluctuations.

```jldoctest
using PathWeightSampling
system = PathWeightSampling.gene_expression_system(mean_s=25, mean_x=50, corr_time_s=1.0, corr_time_x=0.75)

# output

SimpleSystem with 4 reactions
Input variables: S(t)
Output variables: X(t)
Initial condition:
    S(t) = 25
    X(t) = 50
Parameters:
    κ = 25.0
    λ = 1.0
    ρ = 2.666666666666666
    μ = 1.3333333333333333
```
"""
function gene_expression_system(;
    mean_s=50,
    mean_x=mean_s,
    corr_time_s=1.0,
    corr_time_x=0.1,
    kappa=mean_s / corr_time_s,
    lambda=1 / corr_time_s,
    mu=1 / corr_time_x,
    rho=mu * mean_x / mean_s,
    u0=SA[mean_s, mean_x],
    dtimes=0:0.1:2.0
)
    sn = @reaction_network begin
        κ, ∅ --> S
        λ, S --> ∅
    end κ λ

    xn = @reaction_network begin
        ρ, S --> X + S
        μ, X --> ∅
    end ρ μ

    ps = [kappa, lambda]
    px = [rho, mu]

    SimpleSystem(sn, xn, u0, ps, px, dtimes)
end

function chemotaxis_system(;
    mean_L=20,
    num_receptors=100,
    Y_tot=500,
    L_timescale=1.0,
    LR_timescale=0.01,
    LR_ratio=0.5,
    Y_timescale=0.1,
    Y_ratio=1 / 2,
    q=0,
    dtimes=0:0.1:2.0
)
    mean_LR = num_receptors * LR_ratio
    mean_R = num_receptors - mean_LR

    mean_Yp = Y_tot * Y_ratio
    mean_Y = Y_tot - mean_Yp

    eq_L = mean_L * exp(-q)

    sn = @reaction_network begin
        κ, ∅ --> L
        λ, L --> ∅
    end κ λ

    rn = @reaction_network begin
        ρ, L + R --> L + LR
        μ, LR --> R
    end ρ μ

    xn = @reaction_network begin
        δ, LR + Y --> Yp + LR
        χ, Yp --> Y
    end δ χ

    u0 = round.([mean_L, mean_R, mean_LR, mean_Y, mean_Yp])
    ps = [mean_L, 1 / L_timescale]

    ρ = inv(eq_L * LR_timescale * (1 + mean_R / mean_LR))
    μ = eq_L * ρ * mean_R / mean_LR

    pr = [ρ, μ]

    χ = inv(Y_timescale * (1 + Y_ratio))
    δ = χ * Y_ratio / mean_LR

    px = [δ, χ]

    ComplexSystem(sn, rn, xn, u0, ps, pr, px, dtimes)
end

function chemotaxis_parameters(;
    E₀=3.0,
    Kₐ=500,
    Kᵢ=25,
    δg=log(Kₐ / Kᵢ),
    δf=-1.5,
    k⁺=0.05,
    k⁺ₐ=k⁺,
    k⁺ᵢ=k⁺,
    k⁻ₐ=Kₐ * k⁺ₐ,
    k⁻ᵢ=Kᵢ * k⁺ᵢ,
    a_star=0.5,
    γ=1 / 10,
    k_B=(1 - a_star) * γ / abs(δf),
    k_R=a_star * γ / abs(δf)
)
    [
        E₀,
        δg,
        δf,
        k⁺ₐ,
        k⁺ᵢ,
        k⁻ₐ,
        k⁻ᵢ,
        k_B,
        k_R
    ]
end

"""
    cooperative_chemotaxis_system(;
        lmax = 3,
        mmax = 9,
        n_clusters = 100,
        n_chey = 10_000,

        mean_l = 50,
        tau_l = 1.0,

        phosphorylate = 2000.0 / (n_chey * n_clusters),
        dephosphorylate = 2000.0 / (n_chey),
        dtimes = 0:0.1:20.0,
        varargs...
    )

Create a system for a complex chemotaxis model.

# Model Description

This model describes the bacterial chemotaxis signaling network.

```
Rml -> Rm(l+1), with rate kon L(t) (lmax - l): ligand binding to active state
Rm(l+1) -> Rml, with rate koff_A l: ligand unbinding
Rml -> R(m+1)l with rate km: methylation rate
R(m+1)l -> Rml with rate kdm: demethylation rate
```

# Examples

Create a chemotaxis system with default parameters.
```jldoctest
using PathWeightSampling
PathWeightSampling.cooperative_chemotaxis_system()

# output

ComplexSystem with 175 reactions
Input variables: L(t)
Latent variables: R_0_0(t), R_0_1(t), R_0_2(t), R_0_3(t), R_0_4(t), R_0_5(t), R_0_6(t), R_0_7(t), R_0_8(t), R_0_9(t), R_1_0(t), R_1_1(t), R_1_2(t), R_1_3(t), R_1_4(t), R_1_5(t), R_1_6(t), R_1_7(t), R_1_8(t), R_1_9(t), R_2_0(t), R_2_1(t), R_2_2(t), R_2_3(t), R_2_4(t), R_2_5(t), R_2_6(t), R_2_7(t), R_2_8(t), R_2_9(t), R_3_0(t), R_3_1(t), R_3_2(t), R_3_3(t), R_3_4(t), R_3_5(t), R_3_6(t), R_3_7(t), R_3_8(t), R_3_9(t)
Output variables: Yp(t), Y(t)
Initial condition:
    L(t) = 50
    R_0_0(t) = 100
    R_0_1(t) = 0
    R_0_2(t) = 0
    R_0_3(t) = 0
    R_0_4(t) = 0
    R_0_5(t) = 0
    R_0_6(t) = 0
    R_0_7(t) = 0
    R_0_8(t) = 0
    R_0_9(t) = 0
    R_1_0(t) = 0
    R_1_1(t) = 0
    R_1_2(t) = 0
    R_1_3(t) = 0
    R_1_4(t) = 0
    R_1_5(t) = 0
    R_1_6(t) = 0
    R_1_7(t) = 0
    R_1_8(t) = 0
    R_1_9(t) = 0
    R_2_0(t) = 0
    R_2_1(t) = 0
    R_2_2(t) = 0
    R_2_3(t) = 0
    R_2_4(t) = 0
    R_2_5(t) = 0
    R_2_6(t) = 0
    R_2_7(t) = 0
    R_2_8(t) = 0
    R_2_9(t) = 0
    R_3_0(t) = 0
    R_3_1(t) = 0
    R_3_2(t) = 0
    R_3_3(t) = 0
    R_3_4(t) = 0
    R_3_5(t) = 0
    R_3_6(t) = 0
    R_3_7(t) = 0
    R_3_8(t) = 0
    R_3_9(t) = 0
    Yp(t) = 0
    Y(t) = 10000
Parameters:
    κ = 50.0
    λ = 1.0
    E0 = 3.0
    δg = 2.995732273553991
    δf = -1.5
    lba = 0.05
    lbi = 0.05
    lda = 25.0
    ldi = 1.25
    mda = 0.03333333333333333
    mbi = 0.03333333333333333
    μ = 8.57142857142857
    ρ = 0.028571428571428564
```

"""
function cooperative_chemotaxis_system(;
    lmax=3,
    mmax=9,
    n_clusters=100,
    n_chey=10_000,
    mean_l=50,
    tau_l=1.0,
    tau_y=0.1,
    phi_y=1.0 / 6.0,
    dephosphorylate=inv(tau_y * (1 + phi_y)),
    phosphorylate=dephosphorylate * phi_y / (n_clusters / 2),
    dtimes=0:0.1:20.0, aggregator=DiffEqJump.RSSACR(),
    varargs...
)
    sn = @reaction_network begin
        κ, ∅ --> L
        λ, L --> ∅
    end κ λ

    rn = Catalyst.make_empty_network()

    Catalyst.@parameters t E0 δg δf lba lbi lda ldi mda mbi ρ
    Catalyst.@variables L(t) Y(t) Yp(t)

    Catalyst.addspecies!(rn, L)

    Catalyst.addparam!(rn, E0)
    Catalyst.addparam!(rn, δg)
    Catalyst.addparam!(rn, δf)
    Catalyst.addparam!(rn, lba)
    Catalyst.addparam!(rn, lbi)
    Catalyst.addparam!(rn, lda)
    Catalyst.addparam!(rn, ldi)
    Catalyst.addparam!(rn, mda)
    Catalyst.addparam!(rn, mbi)

    xn = @reaction_network begin
        μ, Yp --> Y
    end μ

    Catalyst.addparam!(xn, ρ)

    spmap = Dict()
    for l = 0:lmax, m = 0:mmax
        receptor_species = ModelingToolkit.Num(ModelingToolkit.variable(Symbol("R_", l, "_", m), T=FnType{Tuple{Any},Real}))(t)

        spmap[(l, m)] = receptor_species

        Catalyst.addspecies!(xn, receptor_species)
        Catalyst.addspecies!(rn, receptor_species)
    end

    p_active(l, m) = 1 / (1 + exp(E0 + l * δg + m * δf))

    for l = 0:lmax, m = 0:mmax

        if l > 0
            ligand_bind = Catalyst.Reaction((lba * p_active(l - 1, m) + lbi * (1 - p_active(l - 1, m))) * (lmax + 1 - l), [spmap[(l - 1, m)], L], [spmap[(l, m)], L])
            ligand_unbind = Catalyst.Reaction((lda * p_active(l, m) + ldi * (1 - p_active(l, m))) * l, [spmap[(l, m)]], [spmap[(l - 1, m)]])
            Catalyst.addreaction!(rn, ligand_bind)
            Catalyst.addreaction!(rn, ligand_unbind)
        end

        if m > 0
            demethylate_active = Catalyst.Reaction(mda * p_active(l, m), [spmap[(l, m)]], [spmap[(l, m - 1)]])
            methylate_inactive = Catalyst.Reaction(mbi * (1 - p_active(l, m - 1)), [spmap[(l, m - 1)]], [spmap[(l, m)]])

            Catalyst.addreaction!(rn, demethylate_active)
            Catalyst.addreaction!(rn, methylate_inactive)
        end

        # every receptor phosphorylates Y with rate ρ if the receptor is active
        # if the receptor is inactive, no phosphorylation can happen
        receptor = spmap[(l, m)]
        phosphorylation = Catalyst.Reaction(ρ * p_active(l, m), [Y, receptor], [Yp, receptor])
        Catalyst.addreaction!(xn, phosphorylation)
    end

    joint = ModelingToolkit.extend(xn, ModelingToolkit.extend(rn, sn))

    u0 = zeros(Int, length(spmap) + 3)
    u0[1] = round(Int, mean_l)
    u0[Catalyst.speciesmap(joint)[spmap[(0, 0)]]] = n_clusters
    u0[Catalyst.speciesmap(joint)[Y]] = n_chey

    ps = [mean_l / tau_l, 1.0 / tau_l]
    pr = chemotaxis_parameters(; varargs...)
    px = [dephosphorylate, phosphorylate]

    ComplexSystem(sn, rn, xn, u0, ps, pr, px, dtimes; aggregator=aggregator)
end

function sde_chemotaxis_system(;
    velocity_decay=0.862,
    velocity_noise=sqrt(2 * velocity_decay * 157.1),
    gradient_steepness=0.2e-3,
    harmonic_rate=nothing,
    kwargs...
)
    ModelingToolkit.@variables t V(t) L(t)
    ModelingToolkit.@parameters λ σ g ω₀

    D = ModelingToolkit.Differential(t)

    if harmonic_rate === nothing
        eqs = [
            D(V) ~ -λ * V,
            D(L) ~ g * L * V
        ]
        noiseeqs = [σ, 0]
        sparams = [λ, σ, g]
        ps = [velocity_decay, velocity_noise, gradient_steepness]
    else
        c₀ = 100.0 # μM
        eqs = [
            D(V) ~ -ω₀^2 * (L - c₀) - λ * V,
            D(L) ~ g * c₀ * V
        ]
        noiseeqs = [σ, 0]
        sparams = [λ, σ, g, ω₀]
        ps = [velocity_decay, velocity_noise, gradient_steepness, harmonic_rate]
    end

    ModelingToolkit.@named sn = SDESystem(eqs, noiseeqs, t, [V, L], sparams)

    system = PathWeightSampling.cooperative_chemotaxis_system(; kwargs...)

    Catalyst.addspecies!(system.rn, V)
    st = ModelingToolkit.states(system.rn)
    perm = circshift(collect(1:length(st)), 1)
    Catalyst.reorder_states!(system.rn, perm)

    sds = PathWeightSampling.SDEDrivenSystem(
        sn, system.rn, system.xn,
        vcat(0.0, system.u0[1], system.u0[2:end]), # u0
        ps, system.pr, system.px,
        system.dtimes
    )
end

getname(sym) = String(ModelingToolkit.operation(sym).name)

function parse_receptor(species)
    sname = getname(species)
    mtch = match(r"([A-Z])_(\d+)_(\d+)", sname)
    if mtch !== nothing
        l = parse(Int, mtch.captures[2])
        m = parse(Int, mtch.captures[3])
        (l, m)
    else
        nothing
    end
end

function receptor_states(rs::ReactionSystem)
    smap = Catalyst.speciesmap(rs)
    xf = KeepSomething() do (species, index)
        result = parse_receptor(species)
        if result === nothing
            nothing
        else
            result => index
        end
    end
    smap |> xf
end

receptor_states(system::ComplexSystem) = receptor_states(reaction_network(system))

function active_receptors(conf::SRXconfiguration, system::ComplexSystem)
    rstates = receptor_states(system)

    E0 = system.pr[1]
    δg = system.pr[2]
    δf = system.pr[3]

    p_active(l, m) = 1 / (1 + exp(E0 + l * δg + m * δf))

    p_a = rstates |> Map(((l, m), i)::Pair -> i => p_active(l, m)) |> collect

    f = function (u)
        sum(p_a) do (i, p)
            p * u[i]
        end
    end

    merge_trajectories(conf.s_traj, conf.r_traj, conf.x_traj) |> Map((u, t, i)::Tuple -> (ensurevec(f(u)), t, i)) |> collect_trajectory
end

precompile(gene_expression_system, ())
precompile(chemotaxis_system, ())
precompile(cooperative_chemotaxis_system, ())


export HDPHyperparam

"Hierarchical Dirichlet Process Mixture Model Hyperparameters"
struct HDPHyperparam <: AbstractHyperparam

  γ_a::Float64
  γ_b::Float64

  α_a::Float64
  α_b::Float64

  "default values"
  HDPHyperparam(;γ_a = 1.0, γ_b = 1.0, α_a = 1.0, α_b = 1.0) = new(γ_a, γ_b, α_a, α_b)

end

"Hierarchical Dirichlet Process Mixture Model Data Object"
mutable struct HDPData <: AbstractModelData

  # Energy
  energy::Float64

  # Distributions
  distributions::Vector{ConjugatePostDistribution}

  # Assignments
  assignments::Vector{Vector{Int}}

	# Weights
  weights::Vector{Vector{Float64}}

end

mutable struct HDPBuffer{T <: Real} <: AbstractModelBuffer

  # samples
  X::Vector{Vector{T}}

  # number of groups
  N0::Int

  # indecies of groups
  N0idx::Vector{Int}

  # number of samples per group
  Nj::Vector{Int}

  # indecies of samples per group
  Njidx::Array{Vector}

  # assignments
  Z::Vector{Vector{Int}}

  # number of active cluster
  K::Int

  # number of samples per group per cluster
  C::Array{Int, 2}

  # total number of tables
  totalnt::Array{Int}

  # number of clusters per table
  classnt::Array{Int}

  # distributions
  G::Vector{ConjugatePostDistribution}

  # base distribution
  G0::ConjugatePostDistribution

  # concentration parameter
  α::Float64

  # concentration parameters of clusters
  β::Array{Float64}

  # gamma
  γ::Float64
end

"HDP initialization using random assignments."
function init(X::Vector{Vector{T}}, model::HDP, init::RandomInitialisation) where T <: Real

    Z = Vector{Vector{Int}}(length(X))
    G = ConjugatePostDistribution[]

    K = init.k

    for i in 1:length(X)
      N = length(X[i])
      Z[i] = rand(1:K, N)

      for c in 1:K
        idx = findall(Z[i] .== c)
        if length(idx) > 0
          if size(G, 1) < c
            push!(G, add(model.H, X[i][idx]))
          else
            add!(G[c], X[i][idx])
          end
        end
      end
    end

    N0 = length(X)
    N0idx = collect(1:N0)

    Nj = [length(x) for x in X]
    Njidx = [collect(1:N) for N in Nj]

    C = reduce(hcat, [counts(z, 1:K) for z in Z])

    # init step
    β = ones(K) ./ K
    classnt = randnumtable(model.α * β[:,ones(Int, N0)]', C')
    totalnt = sum(classnt, 1)

    # set alpha vector of Dirichlet Distribution to sample β
    a = zeros(K + 1)
    a[1:end-1] = totalnt
    a[end] = model.γ

    # update beta
    β = rand(Dirichlet(a));

    return HDPBuffer{T}(
      X,
      N0,
      N0idx,
      Nj,
      Njidx,
      Z,
      K,
      C,
      totalnt,
      classnt,
      G,
      model.H,
      model.α,
      β,
      model.γ)
end

"Single Gibbs sweep of HDP training using beta variables."
function gibbs!(B::HDPBuffer)

    shuffle!(B.N0idx)
    prob = zeros(B.K + 1) * -Inf

    for j in B.N0idx

        # get samples of group
        data = B.X[j]
        z = B.Z[j]

        shuffle!(B.Njidx[j])

        for i in B.Njidx[j]

            cluster = z[i]

            # remove data item from model
            remove!(B.G[cluster], data[i])
            B.C[cluster, j] -= 1

            if isdistempty(B.G[cluster])

                # remove cluster
                B.G = B.G[[1:cluster-1; cluster+1:end]]

                for jj in 1:B.N0
                    B.Z[jj][B.Z[jj] .> cluster] -= 1
                end

                B.β = B.β[[1:cluster-1; cluster+1:end]]

                B.classnt = B.classnt[:,[1:cluster-1; cluster+1:end]]
                B.totalnt = B.totalnt[[1:cluster-1; cluster+1:end]]'

                B.K -= 1

                B.C = B.C[[1:cluster-1; cluster+1:end],:]
                prob = zeros(B.K + 1) * -Inf
            end

            # compute log likelihood
            for k in 1:B.K
              llh = logpostpred(B.G[k], data[i])[1]
              prior = B.C[k, j] + B.β[k] * B.α
              prob[k] = llh + log( prior )
            end

            prob[B.K + 1] = logpostpred(B.G0, data[i])[1] + log( B.β[B.K+1] * B.α )
            prob = exp(prob - maximum(prob))

            c = randomindex(prob)

            # add data to model
            B.Z[j][i] = c

            if c > B.K
                # create new cluster
                B.K += 1
                B.G = vcat(B.G, deepcopy(B.G0))
                b = rand(Dirichlet([1, B.γ]))
                b = b * B.β[end]
                B.β = vcat(B.β, 1)
                B.β[end-1:end] = b
                B.C = vcat(B.C, zeros(Int, 1, B.N0))
                prob = zeros(B.K + 1) * -Inf
            end

            B.C[c, j] += 1
            add!(B.G[c], [data[i]])

        end

    end

    # sample number of tables
    kk = maximum([0, B.K - length(B.totalnt)])
    B.totalnt = hcat(B.totalnt - sum(B.classnt, 1), zeros(Int, 1, kk))
    B.classnt = randnumtable(B.α .* B.β[:,ones(Int, B.N0)]', B.C')
    B.totalnt = B.totalnt + sum(B.classnt, 1)

    # update beta weights
    a = zeros(B.K + 1)
    a[1:end-1] = B.totalnt
    a[end] = B.γ

    B.β = rand(Dirichlet(a))

    B
end

function updateenergy!(B::HDPData, X::Vector{Vector{T}}) where T <: Real

  E = 0
  numl = 0

  for j in length(X)

      # get samples of group
      data = X[j]

      pp = 0.0
      for i in length(data)
        for k = 1:length(B.distributions)
          pp += logpostpred( B.distributions[k], data[i] )[1] + log(B.weights[j][k])
        end
      end
      E += pp
      numl += 1

  end

  B.energy = E - log(numl)

end

function sampleparameters!(B::HDPBuffer, P::HDPHyperparam)
  totalnt = sum(B.classnt, 1)
  B.γ = resampleα(B.γ, sum(B.totalnt), B.K, k = P.γ_a, Θ = P.γ_b, maxiter = 10)
  B.α = resampleα(B.α, B.Nj, B.totalnt, k = P.α_a, Θ = P.α_b, maxiter = 10)
end

function extractpointestimate(B::HDPBuffer)
  W = map(j -> map(k -> B.C[k, j] + B.β[k] * B.α, 1:B.K), B.N0idx)
  model = HDPData(0.0, deepcopy(B.G), deepcopy(B.Z), deepcopy(W) )
  updateenergy!(model, B.X)
  model
end


#
# Computes the total potential energy of the system
#

function utotal(p,data)
  utotal = 0.
  for ip in 1:(data.N)-1
    for jp in ip+1:(data.N)
      rij = pbcseparation(p[ip],p[jp],data.side)
      if rij < data.cutoff
        utotal += upair(rij,data)
      end
    end
  end
  return utotal
end


#=

Author: Shunsuke Hori

=#

using PyPlot

"""
If we expand the implicit function for h_j(n_k) then we find that
it is a quadratic. We know that h_j(n_k) > 0 so we can get its
value by using the quadratic form
"""

function h_j(j::Integer, nk::Real, s1::Real, s2::Real,
             θ::Real, δ::Real, ρ::Real)
    # Find out who's h we are evaluating
    if j == 1
        sj = s1
        sk = s2
    else
        sj = s2
        sk = s1
    end

    # Coefficients on the quadratic a x^2 + b x + c = 0
    a = 1.0
    b = ((ρ + 1 / ρ) * nk - sj - sk)
    c = (nk * nk - (sj * nk) / ρ - sk * ρ * nk)

    # Positive solution of quadratic form
    root = (-b + sqrt(b * b - 4 * a * c)) / (2 * a)

    return root
end

"""
Determine whether (n1, n2) is in the set DLL
"""
DLL(n1::Real, n2::Real,
    s1_ρ::Real, s2_ρ::Real,
    s1::Real, s2::Real,
    θ::Real, δ::Real, ρ::Real) =
        (n1 <= s1_ρ) && (n2 <= s2_ρ)

"""
Determine whether (n1, n2) is in the set DHH
"""
DHH(n1::Real, n2::Real,
    s1_ρ::Real, s2_ρ::Real,
    s1::Real, s2::Real,
    θ::Real, δ::Real, ρ::Real) =
        (n1 >= h_j(1, n2, s1, s2, θ, δ, ρ)) && (n2 >= h_j(2, n1, s1, s2, θ, δ, ρ))

"""
Determine whether (n1, n2) is in the set DHL
"""
DHL(n1::Real, n2::Real,
    s1_ρ::Real, s2_ρ::Real,
    s1::Real, s2::Real,
    θ::Real, δ::Real, ρ::Real) =
        (n1 >= s1_ρ) && (n2 <= h_j(2, n1, s1, s2, θ, δ, ρ))

"""
Determine whether (n1, n2) is in the set DLH
"""
DLH(n1::Real, n2::Real,
    s1_ρ::Real, s2_ρ::Real,
    s1::Real, s2::Real,
    θ::Real, δ::Real, ρ::Real) =
        (n1 <= h_j(1, n2, s1, s2, θ, δ, ρ)) && (n2 >= s2_ρ)

"""
Takes a current value for (n_{1, t}, n_{2, t}) and returns the
values (n_{1, t+1}, n_{2, t+1}) according to the law of motion.
"""
function one_step(n1::Real, n2::Real,
                  s1_ρ::Real, s2_ρ::Real,
                  s1::Real, s2::Real,
                  θ::Real, δ::Real, ρ::Real)
    # Depending on where we are, evaluate the right branch
    if DLL(n1, n2, s1_ρ, s2_ρ, s1, s2, θ, δ, ρ)
        n1_tp1 = δ * (θ * s1_ρ + (1 - θ) * n1)
        n2_tp1 = δ * (θ * s2_ρ + (1 - θ) * n2)
    elseif DHH(n1, n2, s1_ρ, s2_ρ, s1, s2, θ, δ, ρ)
        n1_tp1 = δ * n1
        n2_tp1 = δ * n2
    elseif DHL(n1, n2, s1_ρ, s2_ρ, s1, s2, θ, δ, ρ)
        n1_tp1 = δ * n1
        n2_tp1 = δ * (θ * h_j(2, n1, s1, s2, θ, δ, ρ) + (1 - θ) * n2)
    elseif DLH(n1, n2, s1_ρ, s2_ρ, s1, s2, θ, δ, ρ)
        n1_tp1 = δ * (θ * h_j(1, n2, s1, s2, θ, δ, ρ) + (1 - θ) * n1)
        n2_tp1 = δ * n2
    end

    return n1_tp1, n2_tp1
end

"""
Given an initial condition, continues to yield new values of `n1` and `n2`
"""
new_n1n2(n1_0::Real, n2_0::Real,
         s1_ρ::Real, s2_ρ::Real,
         s1::Real, s2::Real,
         θ::Real, δ::Real, ρ::Real) =
            one_step(n1_0, n2_0, s1_ρ, s2_ρ, s1, s2, θ, δ, ρ)

"""
Takes initial values and iterates forward to see whether
the histories eventually end up in sync.

If countries are symmetric then as soon as the two countries have the
same measure of firms then they will by synchronized -- However, if
they are not symmetric then it is possible they have the same measure
of firms but are not yet synchronized. To address this, we check whether
firms stay synchronized for `npers` periods with Euclidean norm

##### Parameters
----------
- `n1_0` : `Real`,
Initial normalized measure of firms in country one
- `n2_0` : `Real`,
Initial normalized measure of firms in country two
- `maxiter` : `Integer`,
Maximum number of periods to simulate
- `npers` : `Integer`,
Number of periods we would like the countries to have the same measure for

##### Returns
-------
- `synchronized` : `Bool`,
Did they two economies end up synchronized
- `pers_2_sync` : `Integer`,
The number of periods required until they synchronized
"""
function pers_till_sync(n1_0::Real, n2_0::Real,
                        s1_ρ::Real, s2_ρ::Real,
                        s1::Real, s2::Real,
                        θ::Real, δ::Real, ρ::Real,
                        maxiter::Integer, npers::Integer)

    # Initialize the status of synchronization
    synchronized = false
    pers_2_sync = maxiter
    iters = 0

    nsync = 0

    while (~synchronized) && (iters < maxiter)
        # Increment the number of iterations and get next values
        iters += 1

        n1_t, n2_t = new_n1n2(n1_0, n2_0, s1_ρ, s2_ρ, s1, s2, θ, δ, ρ)

        # Check whether same in this period
        if abs(n1_t - n2_t) < 1e-8
            nsync += 1
        # If not, then reset the nsync counter
        else
            nsync = 0
        end

        # If we have been in sync for npers then stop and countries
        # became synchronized nsync periods ago
        if nsync > npers
            synchronized = true
            pers_2_sync = iters - nsync
        end
        n1_0, n2_0 = n1_t, n2_t
    end
    return synchronized, pers_2_sync
end

function create_attraction_basis{TR <: Real}(s1_ρ::TR, s2_ρ::TR,
                                 s1::TR, s2::TR, θ::TR, δ::TR, ρ::TR,
                                 maxiter::Integer, npers::Integer, npts::Integer)
    # Create unit range with npts
    synchronized, pers_2_sync = false, 0
    unit_range = linspace(0.0, 1.0, npts)

    # Allocate space to store time to sync
    time_2_sync = Matrix{TR}(npts, npts)
    # Iterate over initial conditions
    for (i, n1_0) in enumerate(unit_range)
        for (j, n2_0) in enumerate(unit_range)
            synchronized, pers_2_sync = pers_till_sync(n1_0, n2_0, s1_ρ, s2_ρ,
                                                        s1, s2, θ, δ, ρ,
                                                        maxiter, npers)
            time_2_sync[i, j] = pers_2_sync
        end
    end

    return time_2_sync
end

# == Now we define a type for the model == #

"""
The paper "Globalization and Synchronization of Innovation Cycles" presents
a two country model with endogenous innovation cycles. Combines elements
from Deneckere Judd (1985) and Helpman Krugman (1985) to allow for a
model with trade that has firms who can introduce new varieties into
the economy.

We focus on being able to determine whether two countries eventually
synchronize their innovation cycles. To do this, we only need a few
of the many parameters. In particular, we need the parameters listed
below

##### Parameters
----------
- `s1` : `Real`,
Amount of total labor in country 1 relative to total worldwide labor
- `θ` : `Real`,
A measure of how mcuh more of the competitive variety is used in
production of final goods
- `δ` : `Real`,
Percentage of firms that are not exogenously destroyed every period
- `ρ` : `Real`,
Measure of how expensive it is to trade between countries
"""
struct MSGSync{TR <: Real}
    s1::TR
    s2::TR
    s1_ρ::TR
    s2_ρ::TR
    θ::TR
    δ::TR
    ρ::TR
end

function MSGSync(s1::Real=0.5, θ::Real=2.5,
                 δ::Real=0.7, ρ::Real=0.2)
    # Store other cutoffs and parameters we use
    s2 = 1 - s1
    s1_ρ = min((s1 - ρ * s2) / (1 - ρ), 1)
    s2_ρ = 1 - s1_ρ

    model=MSGSync(s1, s2, s1_ρ, s2_ρ, θ, δ, ρ)

    return model
end

unpack_params(model::MSGSync) =
    model.s1, model.s2, model.θ, model.δ, model.ρ, model.s1_ρ, model.s2_ρ

"""
Simulates the values of `n1` and `n2` for `T` periods

##### Parameters
----------
- `n1_0` : `Real`, Initial normalized measure of firms in country one
- `n2_0` : `Real`, Initial normalized measure of firms in country two
- `T` : `Integer`, Number of periods to simulate

##### Returns
-------
- `n1` : `Vector{TR}(ndim=1) where TR <: Real`,
A history of normalized measures of firms in country one
- `n2` : `Vector{TR}(ndim=1) where TR <: Real`,
A history of normalized measures of firms in country two
"""
function simulate_n{TR <: Real}(model::MSGSync, n1_0::TR, n2_0::TR, T::Integer)
    # Unpack parameters
    s1, s2, θ, δ, ρ, s1_ρ, s2_ρ = unpack_params(model)

    # Allocate space
    n1 = Vector{TR}(T)
    n2 = Vector{TR}(T)

    # Simulate for T periods
    for t in 1:T
        # Get next values
        n1[t], n2[t] = n1_0, n2_0
        n1_0, n2_0 = new_n1n2(n1_0, n2_0, s1_ρ, s2_ρ, s1, s2, θ, δ, ρ)
    end

    return n1, n2
end

"""
Takes initial values and iterates forward to see whether
the histories eventually end up in sync.

If countries are symmetric then as soon as the two countries have the
same measure of firms then they will by synchronized -- However, if
they are not symmetric then it is possible they have the same measure
of firms but are not yet synchronized. To address this, we check whether
firms stay synchronized for `npers` periods with Euclidean norm

##### Parameters
----------
- `n1_0` : `Real`,
Initial normalized measure of firms in country one
- `n2_0` : `Real`,
Initial normalized measure of firms in country two
- `maxiter` : `Integer`,
Maximum number of periods to simulate
- `npers` : `Integer`,
Number of periods we would like the countries to have the same measure for

##### Returns
-------
- `synchronized` : `Bool`,
Did they two economies end up synchronized
- `pers_2_sync` : `Integer`,
The number of periods required until they synchronized
"""
function pers_till_sync(model::MSGSync, n1_0::Real, n2_0::Real,
                        maxiter::Integer=500, npers::Integer=3)
    # Unpack parameters
    s1, s2, θ, δ, ρ, s1_ρ, s2_ρ = unpack_params(model)

    return pers_till_sync(n1_0, n2_0, s1_ρ, s2_ρ, s1, s2,
                          θ, δ, ρ, maxiter, npers)
end

"""
Creates an attraction basis for values of n on [0, 1] X [0, 1] with npts in each dimension
"""
function create_attraction_basis(model::MSGSync;
                                 maxiter::Integer=250,
                                 npers::Integer=3,
                                 npts::Integer=50)
    # Unpack parameters
    s1, s2, θ, δ, ρ, s1_ρ, s2_ρ = unpack_params(model)

    ab = create_attraction_basis(s1_ρ, s2_ρ, s1, s2, θ, δ,
                                  ρ, maxiter, npers, npts)
    return ab
end


eye(n) = LinearAlgebra.I + zeros(n, n)

mutable struct PropCov
  value::Matrix{Float64}
  buffer::Matrix{Float64}
  iter::Int
  batchsize::Int
  scale::Float64
  momentum::Float64
  eps::Float64
end

function PropCov(value::Matrix{Float64}; iter::Int=0,
                 batchsize::Int=200, momentum::Float64=0.1,
                 eps::Float64=1e-6)
  # Get dimensions of prop cov
  K, K2 = size(value)
  @assert K == K2

  # Momentum needs to be between 0 and 1
  @assert 0 <= momentum <= 1

  # Create buffer
  buffer = zeros(Float64, batchsize, K)

  # Scale for proposal covariance
  scale = 2.38 * 2.38 / K

  return PropCov(value, buffer, iter, batchsize, scale, momentum, eps)
end

function PropCov(K::Int; sigma::Float64=1.0, iter::Int=0, batchsize::Int=200,
                 momentum::Float64=0.1, eps::Float64=1e-6)
  return PropCov(eye(K) * sigma, iter=iter, batchsize=batchsize,
                 momentum=momentum, eps=eps)
end

nparam(pc::PropCov) = size(pc.buffer, 2)

function update!(pc::PropCov, x::Vector{Float64})
  # Increment iteration
  pc.iter += 1

  # Compute current buffer index
  buffer_index = let
    idx = mod(pc.iter, pc.batchsize)
    idx == 0 ? pc.batchsize : idx
  end

  # Cache current value in buffer
  pc.buffer[buffer_index, :] .= deepcopy(x)

  if buffer_index == pc.batchsize
    # Update prop cov
    pc.value .= let
      S = cov(pc.buffer) * pc. scale * pc.momentum + pc.value * (1 - pc.momentum)
      S + eye(nparam(pc)) * pc.eps
    end
  end
end

# K = 30
# pc = Metropolis.PropCov(K, sigma=1e-2)
# for i in 1:10000 Metropolis.update!(pc, randn(K)) end
# pc.value


using LazySequences

# Cons
c = cons(1, [42])
# Test first(s::Cons)
@assert first(c) == 1
# Test rest(s::Cons)
@assert first(rest(c)) == 42

ct = cat([1], [42])
# Test first(s::Cat)
@assert first(ct) == 1
# Test rest(s::Cat)

# Test getindex implementation
fibs = cat([0, 1], @lazyseq map(+, rest(fibs), fibs))
@assert fibs[1] == 0
@assert fibs[2] == 1
@assert fibs[3] == 1
@assert fibs[4] == 2


#=

=#
function sortededgesbyweight(G::DictGraph)
    @assert !G.directed "the graph must be undirected"
    𝑄 = PriorityQueue{Tuple{Int,Int}, Int}()
    for k in keys(G.relation)
        for link in G.relation[k]
            local t = (k, link.u)
            (!haskey(𝑄, t) && !haskey(𝑄, reverse(t))) && push!(𝑄, t => link.w)
        end
    end
    return 𝑄

end

function mstkruskal(G::DictGraph)
    T = DictGraph()
    Q = sortededgesbyweight(G)
    DS = IntDisjointSets(length(G.relation))

    while !isempty(Q)
        e = dequeue_pair!(Q)
        if !in_same_set(DS, e.first[1] + 1, e.first[2] + 1)
            union!(DS, e.first[1] + 1, e.first[2] + 1)
            addedge!(T, e.first[1], e.first[2]; weight = e.second)
        end
    end

    return T
end


@testset "Non_negative" begin

    p, p_inv = Non_negative()

    x = abs.(randn((10, 10)))
    x2 = 100 .*randn((10, 10))

    @test x ≈ p(p_inv(x))
    @test x ≈ p_inv(p(x))

    @test all(p(x2) .>= 0)
end

@testset "Map_0_1" begin

    p, p_inv = Map_0_1()

    x = abs.(randn((10, 10)))
    x2 = 100 .*randn((10, 10))

    @test x ≈ p_inv(p(x))

    @test all(p(x2) .>= 0)
    @test all(p(x2) .<= 1)
end

@testset "Pow4_positive" begin
    p, p_inv = Pow4_positive()

    x = abs.(randn((10, 10)))
    x2 = 100 .*randn((10, 10))

    @test x ≈ p(p_inv(x))
    @test x ≈ p_inv(p(x))

    @test all(p(x2) .>= 0)
end

@testset "Piecewise_positive" begin
    p, p_inv = Piecewise_positive()

    x = abs.(randn((10, 10)))
    x2 = 100 .*randn((10, 10))

    @test x ≈ p(p_inv(x))
    @test x ≈ p_inv(p(x))

    @test all(p(x2) .>= 0)

    function f(x)
        @test all(.≈(Zygote.gradient(x -> sum(p(x)), x)[1], (p(x .+ 1e-8) .- p(x))./1e-8, rtol=1e-4))
        @test Zygote.gradient(x -> sum(p(x)), x)[1] ≈ DeconvOptim.f_pw_pos_grad(x)
    end

    f([1.1, 12312.2, -10.123, 22.2, -123.23, 0])
end

@testset "Abs_positive" begin
    p, p_inv = Abs_positive()

    x = abs.(randn((10, 10)))
    x2 = 100 .*randn((10, 10))

    @test x ≈ p(p_inv(x))
    @test x ≈ p_inv(p(x))

    @test all(p(x2) .>= 0)
end


@testset "ResidueRingMultGrp" begin

  function verify_order(g::Hecke.NfOrdQuoRingElem,o)
    g == 0 && return false
    g^o == 1 || return false
    for l in keys(factor(o).fac)
      g^div(o,l) == 1 && return false
    end
    return true
  end

  function verify_order(g,i,o)
    iszero(g) && return false
    powermod(g,o,i.gen_one) - 1 in i || return false
    for l in keys(factor(o).fac)
      powermod(g,div(o,l),i.gen_one) - 1  in i && return false
    end
    return true
  end

  @testset "multiplicative_group" begin
     Qx,  x = PolynomialRing(FlintQQ, "x");

    @testset "K = Q" begin
      K,  a = NumberField(x,"a");
      O = maximal_order(K)

      @testset "m0 = <$n>" for n in 1:50
        m0 = O(n)*O
        Q = NfOrdQuoRing(O, m0)
        G, M = multiplicative_group(Q)
        @test issnf(G)
        SNF = G.snf
        gens = M.generators
        @test length(SNF) == length(gens)
        for i in 1:length(gens)
          @test verify_order(gens[i],SNF[i])
        end
        H = Hecke.multgrp_of_cyclic_grp(n)
        @test Hecke.isisomorphic(G, H)
        @test M(G(G.snf)) == Q(1)
        for g in gens
          for exp in -5:10
            el = g^exp
            @test M(M\el) == el
          end
        end
      end

      @testset "m0 = <1361>^3000" begin
        m0 = ideal(O,O(1361))^100
        Q = NfOrdQuoRing(O,m0)
        G, M = multiplicative_group(Q)
        @test issnf(G)
        SNF = G.snf
        gens = M.generators
        @test length(SNF) == length(gens)
        for i in 1:length(gens)
          @test verify_order(gens[i],SNF[i])
        end
        @test M(G(G.snf)) == Q(O(1))
        for g in gens
          for exp in -5:10
            el = g^exp
            @test M(M\el) == el
          end
        end
      end

    end

    @testset "K = Q[√2]" begin
       K,  a = NumberField(x^2 - 2,"a")
      O = maximal_order(K)

      @testset "m0 = <$n>" for n in 1:100
        m0 = ideal(O,O(n))
        Q = NfOrdQuoRing(O,m0)
        G, M = multiplicative_group(Q)
        @test issnf(G)
        SNF = G.snf
        gens = M.generators
        @test length(SNF) == length(gens)
        for i in 1:length(gens)
          @test verify_order(gens[i], SNF[i])
        end
        @test M(G(G.snf)) == Q(O(1))
        for g in gens
          for exp in -5:10
            el = g^exp
            @test M(M\el) == el
          end
        end
      end
    end

    @testset "f = x^2-x+15" begin
      f = x^2-x+15
       K,  a = NumberField(f,"a");
      O = maximal_order(K)

      @testset "m0 = <2>" begin
        m0 = ideal(O,O(2))
        Q = NfOrdQuoRing(O,m0)
        G, M = multiplicative_group(Q)
        @test issnf(G)
        SNF = G.snf
        gens = M.generators
        @test length(SNF) == length(gens)
        for i in 1:length(gens)
          @test verify_order(gens[i],SNF[i])
        end
        @test order(G) == 3
        H = abelian_group([ 3 ])
        @test Hecke.isisomorphic(G,H)
        @test M(G(G.snf)) == Q(O(1))
        for g in gens
          for exp in -5:10
            el = g^exp
            @test M(M\el) == el
          end
        end
      end

      @testset "m0 = <4>" begin
        m0 = ideal(O,O(4))
        Q = NfOrdQuoRing(O,m0)
        G, M = multiplicative_group(Q)
        @test issnf(G)
        SNF = G.snf
        gens = M.generators
        @test length(SNF) == length(gens)
        for i in 1:length(gens)
          @test verify_order(gens[i],SNF[i])
        end
        @test order(G) == 12
        H = abelian_group([ 3, 2, 2 ])
        @test Hecke.isisomorphic(G,H)
        @test M(G(G.snf)) == Q(O(1))
        for g in gens
          for exp in -5:10
            el = g^exp
            @test M(M\el) == el
          end
          el = g^(-8)
          @test M(M\el) == el
        end
      end

      @testset "m0 = <5>" begin
        m0 = ideal(O,O(5))
        Q = NfOrdQuoRing(O,m0)
        G, M = multiplicative_group(Q)
        @test issnf(G)
        SNF = G.snf
        gens = M.generators
        @test length(SNF) == length(gens)
        for i in 1:length(gens)
          @test verify_order(gens[i], SNF[i])
        end
        @test M(G(G.snf)) == Q(O(1))
        @test order(G) == 16
        H = abelian_group([4, 4])
        @test Hecke.isisomorphic(G,H)
        for g in gens
          for exp in [ -3, 2, 4 ]
            el = g^exp
            @test M(M\el) == el
          end
        end
      end

      @testset "m0 = <20>" begin
        m0 = ideal(O,O(20))
        Q = NfOrdQuoRing(O,m0)
        G, M = multiplicative_group(Q)
        @test issnf(G)
        SNF = G.snf
        gens = M.generators
        @test length(SNF) == length(gens)
        for i in 1:length(gens)
          @test verify_order(gens[i],SNF[i])
        end
        @test M(G(G.snf)) == Q(O(1))
        @test order(G) == 192
        H = abelian_group([3, 2, 2, 4, 4])
        @test Hecke.isisomorphic(G,H)
        for g in gens
          for exp in [ -4, 1, 7 ]
            el = g^exp
            @test M(M\el) == el
          end
        end
      end
    end

    @testset "x^12+..." begin
      #= f = x^12+7*x^11-97*x^10-859*x^9+2558*x^8+38839*x^7+30012*x^6-710649*x^5-2189082*x^4+2534629*x^3+25314673*x^2+43623088*x+28168151 =#
      #=  K,  a = NumberField(f,"a"); =#
      #= O = maximal_order(K) =#

      #= @testset "m0 = <3>" begin =#
      #=   println("m0 = <3>") =#
      #=   m0 = ideal(O,O(3)) =#
      #=   Q = NfOrdQuoRing(O,m0) =#
      #=   G = domain(Hecke.multiplicative_group(Q)) =#
      #=   println(snf(G)[1].snf) =#
      #=   @test order(G) == 512000 =#
      #=   structt = [80,80,80] =#
      #=   H = abelian_group(structt) =#
      #=   @test Hecke.is_isomorphic(G,H) =#
      #=   @test M(G(G.snf)) == Q(O(1)) =#
      #=   for g in gens =#
      #=     for exp in -5:10 =#
      #=       el = g^exp =#
      #=       @assert M(M\el) == el =#
      #=     end =#
      #=   end =#
      #= end =#

      #= @testset "m0 = <4>" begin =#
      #=   println("m0 = <4>") =#
      #=   m0 = ideal(O,O(4)) =#
      #=   Q = NfOrdQuoRing(O,m0) =#
      #=   G = domain(Hecke.multiplicative_group(Q)) =#
      #=   println(snf(G)[1].snf) =#
      #=   @test order(G) == 14745600 =#
      #=   structt = [15,2,2,2,2,15,4,4,4,4,2,2,2,2] =#
      #=   H = abelian_group(structt) =#
      #=   @test Hecke.is_isomorphic(G,H) =#
      #=   @test M(G(G.snf)) == Q(O(1)) =#
      #=   for g in gens =#
      #=     for exp in -5:10 =#
      #=       el = g^exp =#
      #=       @assert M(M\el) == el =#
      #=     end =#
      #=   end =#
      #= end =#

      #= @testset "m0 = <5>" begin =#
      #=   println("m0 = <5>") =#
      #=   m0 = ideal(O,O(5)) =#
      #=   Q = NfOrdQuoRing(O,m0) =#
      #=   G = domain(Hecke.multiplicative_group(Q)) =#
      #=   println(snf(G)[1].snf) =#
      #=   @test order(G) == 187500000 =#
      #=   structt = [24,5,5,5,5,5,5,4,5,5,5] =#
      #=   H = abelian_group(structt) =#
      #=   @test Hecke.is_isomorphic(G,H) =#
      #=   @test M(G(G.snf)) == Q(O(1)) =#
      #=   for g in gens =#
      #=     for exp in -5:10 =#
      #=       el = g^exp =#
      #=       @assert M(M\el) == el =#
      #=     end =#
      #=   end =#
      #= end =#

      #= @testset "m0 = <60>" begin =#
      #=   println("m0 = <60>") =#
      #=   m0 = ideal(O,O(60)) =#
      #=   Q = NfOrdQuoRing(O,m0) =#
      #=   G = domain(Hecke.multiplicative_group(Q)) =#
      #=   println(snf(G)[1].snf) =#
      #=   @test order(G) == 14155776*fmpz(10)^14 =#
      #=   structt = [15,2,2,2,2,15,4,4,4,4,2,2,2,2,80,80,80,24,5,5,5,5,5,5,4,5,5,5] =#
      #=   H = abelian_group(structt) =#
      #=   @test Hecke.is_isomorphic(G,H) =#
      #=   @test M(G(G.snf)) == Q(O(1)) =#
      #=   for g in gens =#
      #=     for exp in -5:10 =#
      #=       el = g^exp =#
      #=       @assert M(M\el) == el =#
      #=     end =#
      #=   end =#
      #= end =#
    end
  end

  @testset "_multgrp_mod_pv" begin
     Qx,  x = PolynomialRing(FlintQQ, "x");

    @testset "K = Q" begin
       K,  a = NumberField(x,"a");
      O = maximal_order(K)

      @testset "i = <$pnum>^$v" for (pnum,v) in [(pnum,v) for pnum in [ x for x in 1:50 if isprime(fmpz(x))], v in [1,2,4,17]]
        #p = ideal(O,O(pnum))
        p = prime_decomposition(O, pnum)[1][1]
        #p = collect(keys(factor(p)))[1]
        pv = p^v
        G, mG = Hecke._multgrp_mod_pv(p, v, pv)
        S, StoG, mS = snf(G, mG)
        gens = mS.generators
        d = S.snf
        @test length(gens) == length(d)
        for i in 1:length(gens)
          @test verify_order(gens[i].elem, pv, d[i])
        end
        H = Hecke.multgrp_of_cyclic_grp(fmpz(pnum)^v)
        @test Hecke.isisomorphic(S, H)
        # Test discrete logarithm
        Q = NfOrdQuoRing(O, pv)
        for g in mG.generators
          for exp in [ -1, 1, 6 ]
            el = g^exp
            @test mG(mG\el) == el
          end
        end
      end
    end

    @testset "f = x^2-x+15" begin
      f = x^2-x+15
       K,  a = NumberField(f,"a");
      O = maximal_order(K)
      I = O(20)*O

      @testset "p = <$(p.gen_one),$(p.gen_two)>, v = $v" for (p,v) in factor(I)
        pv = p^v
        G, mG = Hecke._multgrp_mod_pv(p, v, pv)
        S, StoG, mS = snf(G, mG)
        gens = mS.generators
        structure = S.snf
        @test length(gens) == length(structure)
        for i in 1:length(gens)
          @test verify_order(gens[i].elem, pv, structure[i])
        end
        # Test discrete logarithm
        Q = NfOrdQuoRing(O, pv)
        for g in mG.generators
          for exp in [ -1, 0, 1, 2 ]
            el = g^exp
            @test mG(mG\el) == el
          end
        end
      end
    end

    @testset "f = x^3+8*x^2+6*x-17" begin
      f = x^3+8*x^2+6*x-17
       K,  a = NumberField(f,"a");
      O = maximal_order(K)
      p = O(3)*O
      p = collect(keys(factor(p)))[1]
      structures = Vector{fmpz}[[26],[26,3,3,3],[26,9,9,9],[26,27,27,27]]
      @testset "v = $v" for v in 1:length(structures)
        pv = p^v
        G, mG = Hecke._multgrp_mod_pv(p, v, pv)
        S, StoG, mS = snf(G, mG)
        gens = mS.generators
        d = S.snf
        @test length(gens) == length(d)
        for i in 1:length(gens)
          @test verify_order(gens[i].elem, pv, d[i])
        end
        @test order(G) == prod(structures[v])
        H = abelian_group(structures[v])
        @test Hecke.isisomorphic(S, H)
        # Test discrete logarithm
        Q = NfOrdQuoRing(O, pv)
        for g in mG.generators
          for exp in [ -2, 1, 3 ]
            el = g^exp
            @test mG(mG\el) == el
          end
        end
      end
    end

    @testset "f = x^4+11*x^3-19*x^2-8*x+7" begin
      f = x^4+11*x^3-19*x^2-8*x+7
       K,  a = NumberField(f,"a");
      O = maximal_order(K)
      p = ideal(O, fmpz(3), O(2+a+a^2))
      p = collect(keys(factor(p)))[1]
      structures = Vector{fmpz}[[8],[8,3,3],[8,3,3,3,3],[8,3,3,3,3,9],[8,3,3,9,9,9],[8,3,9,9,9,27]]
      @testset "v = $v" for v in 1:length(structures)
        pv = p^v
        G, mG = Hecke._multgrp_mod_pv(p, v, pv)
        S, StoG, mS = snf(G, mG)
        gens = mS.generators
        d = S.snf
        @test length(gens) == length(d)
        for i in 1:length(gens)
          @test verify_order(gens[i].elem, pv, d[i])
        end
        @test order(G) == prod(structures[v])
        H = abelian_group(structures[v])
        @test Hecke.isisomorphic(S, H)
        # Test discrete logarithm
        Q = NfOrdQuoRing(O,pv)
        for g in mG.generators
          for exp in [ -4, 1, 6 ]
            el = g^exp
            @test mG(mG\el) == el
          end
        end
      end
    end

    @testset "f = x^6+6x^5-12*x^4-x^3-6*x^2+9*x+20" begin
      f = x^6+6x^5-12*x^4-x^3-6*x^2+9*x+20
       K,  a = NumberField(f,"a");
      O = maximal_order(K)
      p = ideal(O, fmpz(3), O(2+2*a+a^2))
      p = collect(keys(factor(p)))[1]
      structures = Vector{fmpz}[[8],[8,3,3],[8,3,3,3,3],[8,9,9,3,3],[8,9,9,3,3,3,3],[8,9,9,9,9,3,3]]
      @testset "v = $v" for v in 1:length(structures)
        pv = p^v
        G, mG = Hecke._multgrp_mod_pv(p, v, pv)
        S, StoG, mS = snf(G, mG)
        gens = mS.generators
        d = S.snf
        @test length(gens) == length(d)
        for i in 1:length(gens)
          @test verify_order(gens[i].elem, pv, d[i])
        end
        @test order(G) == prod(structures[v])
        H = abelian_group(structures[v])
        @test Hecke.isisomorphic(S, H)
        # Test discrete logarithm
        Q = NfOrdQuoRing(O, pv)
        for g in mG.generators
          for exp in [ -1, 0, 1, 6 ]
            el = g^exp
            @test mG(mG\el) == el
          end
        end
      end
    end
  end

  @testset "_multgrp_mod_p" begin
     Qx,  x = PolynomialRing(FlintQQ, "x");

    @testset "K = Q" begin
       K,  a = NumberField(x, "a");
      O = maximal_order(K)

      @testset "p = <$(pnum)>" for pnum in [ x for x in 1:50 if isprime(fmpz(x))]
        p = O(pnum)*O
        @test isprime(p)
        #p = collect(keys(factor(p)))[1]
        G, mG = Hecke._multgrp_mod_p(p)
        @test length(mG.generators) == 1
        g = mG.generators[1]
        n = exponent(G)
        @test n == pnum - 1
        @test verify_order(g, p, n)
        for exp in [ 0, 1, 2, n - 2, n - 1, n, 3234239 ]
          @test mG.discrete_logarithm(g^exp)[1] == mod(fmpz(exp), n)
        end
      end
    end

    @testset "K = Q[√2]" begin
       K,  a = NumberField(x^2-2,"a");
      O = maximal_order(K)

      primeideals = Vector{Hecke.NfOrdIdl}()
      for pnum in [ x for x in 1:40 if isprime(fmpz(x))]
        append!(primeideals, collect(keys(factor(O(pnum)*O))))
      end

      @testset "p = <$(p.gen_one), $(p.gen_two)>" for p in primeideals
        @test isprime(p)
        G, mG = Hecke._multgrp_mod_p(p)
        @test length(mG.generators) == 1
        g = mG.generators[1]
        n = exponent(G)
        @test !iszero(g)
        @test !(g in p)
        @test verify_order(g, p, n)
        Q = NfOrdQuoRing(O, p)
        for exp in [ 0, 1, 2, n - 2, n - 1, n ]
          @test mG.discrete_logarithm((Q(g)^exp).elem)[1] == mod(fmpz(exp), n)
        end
      end
    end

    @testset "K = Q[x]/<x^6+6*x^5-12*x^4-x^3-6*x^2+9*x+20>" begin
       K,  a = NumberField(x^6+6*x^5-12*x^4-x^3-6*x^2+9*x+20,"a");
      O = maximal_order(K)

      primeideals = Vector{Hecke.NfOrdIdl}()
      for pnum in [ x for x in 1:20 if isprime(fmpz(x)) ]
        append!(primeideals, collect(keys(factor(O(pnum)*O))))
      end

      @testset "p = <$(p.gen_one), $(p.gen_two)>" for p in primeideals
        @test isprime(p)
        G, mG = Hecke._multgrp_mod_p(p)
        @test length(mG.generators) == 1
        g = mG.generators[1]
        n = exponent(G)
        @test !iszero(g)
        @test !(g in p)
        @test verify_order(g, p, n)
        Q = NfOrdQuoRing(O, p)
        for exp in [ 24, n - 345 ]
          @test mG.discrete_logarithm((Q(g)^exp).elem)[1] == mod(fmpz(exp), n)
        end
      end
    end

    @testset "K = Q[x]/<x^10-x^9+x^8-x^7+x^6-x^5+x^4-x^3+x^2-x+1>" begin
       K,  a = NumberField(x^10-x^9+x^8-x^7+x^6-x^5+x^4-x^3+x^2-x+1,"a");
      O = maximal_order(K)

      primeideals = Vector{Hecke.NfOrdIdl}()
      for pnum in [ 2, 3, 5, 11 ]
        append!(primeideals, collect(keys(factor(O(pnum)*O))))
      end

      @testset "p = <$(p.gen_one), $(p.gen_two)>" for p in primeideals
        @test isprime(p)
        G, mG = Hecke._multgrp_mod_p(p)
        @test length(mG.generators) == 1
        g = mG.generators[1]
        n = exponent(G)
        @test !iszero(g)
        @test !(g in p)
        @test verify_order(g, p, n)
        Q = NfOrdQuoRing(O, p)
        for exp in [ 50, n - 30 ]
          @test mG.discrete_logarithm((Q(g)^exp).elem)[1] == mod(fmpz(exp), n)
        end
      end
    end
  end

  @testset "_1_plus_p_mod_1_plus_pv" begin
     Qx,  x = PolynomialRing(FlintQQ, "x");

    @testset "Method: $method" for method in [:quadratic,:artin_hasse,:p_adic]

      @testset "K = Q" begin
         K,  a = NumberField(x,"a");
        O = maximal_order(K)

        @testset "p = <$(pnum)>, v = $(v)" for pnum in [ x for x in 1:30 if isprime(fmpz(x))], v in [1,2,3,4,11,30]
          p = O(pnum)*O
          @test isprime(p)
          pv = p^v
          #p = collect(keys(factor(i)))[1]
          G, mG = Hecke._1_plus_p_mod_1_plus_pv(p, v, pv, fmpz(pnum)^v; method = method)
          @test issnf(G)
          @test length(mG.generators) == ngens(G)
          # Test generators
          for i in 1:length(mG.generators)
            if G.snf[i] != 1
              @test parent(mG.generators[i]) == O
              @test mG.generators[i] - 1 in p
              @test verify_order(mG.generators[i], pv, G.snf[i])
            end
          end
          # Test structure/relations
          H = Hecke.multgrp_of_cyclic_grp(fmpz(pnum)^v)
          I = Hecke._multgrp_mod_p(p)[1]
          @test order(G) == div(order(H), order(I))
          check, J = Hecke._1_plus_pa_mod_1_plus_pb_structure(p, 1, v)
          if check
            @test Hecke.isisomorphic(G, J)
          end
          # Test discrete logarithm
          Q = NfOrdQuoRing(O, pv)
          for g in mG.generators
            for exp in -1:2
              el = Q(g)^exp
              @test Q(mG(mG\el.elem)) == el
            end
          end
        end
      end

      @testset "K = Q[√2]" begin
         K,  a = NumberField(x^2 - 2,"a")
        O = maximal_order(K)

        primeideals = Vector{Hecke.NfOrdIdl}()
        for pnum in [2,3,5,7,19]
          fac = factor(ideal(O,O(pnum)))
          ks = collect(keys(fac))
          vs = collect(values(fac))
          all(isone,vs) && append!(primeideals,ks)
        end

        @testset "p = <$(p.gen_one), $(p.gen_two)>, v = $(v)" for p in primeideals, v in [1,4,11]
          pv = p^v
          G, mG = Hecke._1_plus_p_mod_1_plus_pv(p, v, pv, minimum(p, copy = false)^v; method = method)
          @test issnf(G)
          @test length(mG.generators) == ngens(G)
          # Test generators
          for i in 1:length(mG.generators)
            if G.snf[i] != 1
              @test parent(mG.generators[i]) == O
              @test mG.generators[i] - 1 in p
              @test verify_order(mG.generators[i], pv, G.snf[i])
            end
          end
          # Test structure
          check, J = Hecke._1_plus_pa_mod_1_plus_pb_structure(p, 1, v)
          if check
            @test Hecke.isisomorphic(G, J)
          end
          # Test discrete logarithm
          Q = NfOrdQuoRing(O, pv)
          for g in mG.generators
            for exp in [ -1, 1, 2 ]
              el = Q(g)^exp
              @test Q(mG(mG\el.elem)) == el
            end
          end
        end
      end

      @testset "K = Q[x]/<f = x^6+...>" begin
        f = x^6 + 6*x^5 - 12*x^4 - x^3 - 6*x^2 + 9*x + 20
         K,  a = NumberField(f,"a")
        O = maximal_order(K)

        primeideals = Vector{Hecke.NfOrdIdl}()
        for pnum in [2,11,13]
          fac = factor(ideal(O,O(pnum)))
          ks = collect(keys(fac))
          vs = collect(values(fac))
          all(isone,vs) && append!(primeideals,ks)
        end

        @testset "p = <$(p.gen_one), $(p.gen_two)>, v = $(v)" for p in primeideals, v in [2,9]
          pv = p^v
          G, mG = Hecke._1_plus_p_mod_1_plus_pv(p, v, pv, minimum(p, copy = false)^v; method = method)
          @test issnf(G)
          @test length(mG.generators) == ngens(G)
          # Test generators
          for i in 1:length(mG.generators)
            if G.snf[i] != 1
              @test parent(mG.generators[i]) == O
              @test mG.generators[i] - 1 in p
              @test verify_order(mG.generators[i], pv, G.snf[i])
            end
          end
          # Test structure
          check, J = Hecke._1_plus_pa_mod_1_plus_pb_structure(p, 1, v)
          if check
            @test Hecke.isisomorphic(G, J)
          end
          # Test discrete logarithm
          Q = NfOrdQuoRing(O, pv)
          for g in mG.generators
            for exp in [-1, 2]
              el = Q(g)^exp
              @test Q(mG(mG\el.elem)) == el
            end
          end
        end
      end
    end
  end

  @testset "Non-maximal orders" begin
    Qx, x = FlintQQ["x"]

    K1, a1 = number_field(x^3 - 2, "a1")
    OK1 = maximal_order(K1)
    O1 = Order(K1, [K1(1), 10*a1, 100*a1^2])
    F1 = conductor(O1, OK1)
    Q1, mQ1 = quo(O1, F1)
    G1, mG1 = multiplicative_group(Q1)
    @test G1.snf == [ fmpz(2), fmpz(10), fmpz(20) ]
    for i = 1:10
      g = G1(map(fmpz, rand(-100:100, ngens(G1))))
      h = G1(map(fmpz, rand(-100:100, ngens(G1))))
      @test mG1(g)*mG1(h) == mG1(g + h)
    end

    K2, a2 = number_field(x^3 - 12*x^2 - 6324*x + 459510, "a2")
    O2 = EquationOrder(K2)
    OK2 = maximal_order(O2)
    F2 = conductor(O2, OK2)
    Q2, mQ2 = quo(O2, F2)
    G2, mG2 = multiplicative_group(Q2)
    @test G2.snf == [ fmpz(2), fmpz(4), fmpz(4), fmpz(4368) ]
    for i = 1:10
      g = G2(map(fmpz, rand(-100:100, ngens(G2))))
      h = G2(map(fmpz, rand(-100:100, ngens(G2))))
      @test mG2(g)*mG2(h) == mG2(g + h)
    end

    K3, a3 = number_field(x^5 - 1389*x^4 + 512066*x^3 - 11859166*x^2 + 83453925*x - 211865821, "a3")
    OK3 = maximal_order(K3)
    O3 = EquationOrder(K3)
    F3 = conductor(O3, OK3)
    Q3, mQ3 = quo(O3, F3)
    G3, mG3 = multiplicative_group(Q3)
    @test G3.snf == [ fmpz(2), fmpz(2), fmpz(2), fmpz(2), fmpz(4), fmpz(36100148955782880) ]
    for i = 1:10
      g = G3(map(fmpz, rand(-100:100, ngens(G3))))
      h = G3(map(fmpz, rand(-100:100, ngens(G3))))
      @test mG3(g)*mG3(h) == mG3(g + h)
    end
  end
end