"""
*UtilityModel*

`UtilityModel` is an abstract utility-based model object
````
"""
abstract type UtilityModel end

"""
*mean*

`mean` generic method for computing mean

- `model`: a model M <: UtilityModel
- `gamble`: a gamble object

Function Signature
````julia
mean(model::UtilityModel, gamble::Gamble)
````
"""
function mean(model::UtilityModel, gamble::Gamble)
    sort!(model, gamble)
    weights = compute_weights(model, gamble)
    utility = compute_utility(model, gamble)
    return weights'*utility
end

"""
*var*

`var` a generic method for computing the variance of the gamble

- `model`: a model M <: UtilityModel
- `gamble`: a gamble object

Function Signature
````julia
var(model::UtilityModel, gamble::Gamble)
````
"""
function var(model::UtilityModel, gamble::Gamble)
    @unpack p = gamble
    utility = compute_utility(model, gamble)
    eu = mean(model, gamble)
    return  sum(p .* (utility .- eu).^2)
end

"""
*std*

`std` a generic method for computing the standard deviation of the gamble

- `model`: a model M <: UtilityModel
- `gamble`: a gamble object

Function Signature
````julia
std(model::UtilityModel, gamble::Gamble)
````
"""
std(model::UtilityModel, gamble::Gamble) = sqrt(var(model, gamble))

"""
*compute_weights*

`compute_weights` a generic method for computing decision weights

- `model`: a model M <: UtilityModel
- `gamble`: a gamble object

Function Signature
````julia
compute_weights(model::UtilityModel, gamble::Gamble)
````
"""
compute_weights(model::UtilityModel, gamble::Gamble) = gamble.p

"""
*sort!*

`sort!` a generic method for sorting gamble probabilities and values. The generic method
does not sort the gambles

- `model`: a model M <: UtilityModel
- `gamble`: a gamble object

Function Signature
````julia
sort!(model::UtilityModel, gamble)
````
"""
function sort!(model::UtilityModel, gamble)
    return nothing
end


using GalacticOptim, Optim, Test, Random

rosenbrock(x, p) =  (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2
x0 = zeros(2)
_p  = [1.0, 100.0]

f = OptimizationFunction(rosenbrock, GalacticOptim.AutoForwardDiff())
l1 = rosenbrock(x0, _p)
prob = OptimizationProblem(f, x0, _p)
sol = solve(prob, SimulatedAnnealing())
@test 10*sol.minimum < l1

Random.seed!(1234)
prob = OptimizationProblem(f, x0, _p, lb=[-1.0, -1.0], ub=[0.8, 0.8])
sol = solve(prob, SAMIN())
@test 10*sol.minimum < l1

using CMAEvolutionStrategy
sol = solve(prob, CMAEvolutionStrategyOpt())
@test 10*sol.minimum < l1

rosenbrock(x, p=nothing) =  (1 - x[1])^2 + 100 * (x[2] - x[1]^2)^2

l1 = rosenbrock(x0)
prob = OptimizationProblem(rosenbrock, x0)
sol = solve(prob, NelderMead())
@test 10*sol.minimum < l1

cons= (x,p) -> [x[1]^2 + x[2]^2]
optprob = OptimizationFunction(rosenbrock, GalacticOptim.AutoForwardDiff();cons= cons)

prob = OptimizationProblem(optprob, x0)

sol = solve(prob, ADAM(0.1), maxiters = 1000)
@test 10*sol.minimum < l1

sol = solve(prob, BFGS())
@test 10*sol.minimum < l1

sol = solve(prob, Newton())
@test 10*sol.minimum < l1

sol = solve(prob, Optim.KrylovTrustRegion())
@test 10*sol.minimum < l1

prob = OptimizationProblem(optprob, x0, lcons = [-Inf], ucons = [Inf])
sol = solve(prob, IPNewton())
@test 10*sol.minimum < l1

prob = OptimizationProblem(optprob, x0, lcons = [-5.0], ucons = [10.0])
sol = solve(prob, IPNewton())
@test 10*sol.minimum < l1

prob = OptimizationProblem(optprob, x0, lcons = [-Inf], ucons = [Inf], lb = [-500.0,-500.0], ub=[50.0,50.0])
sol = solve(prob, IPNewton())
@test sol.minimum < l1

function con2_c(x,p)
    [x[1]^2 + x[2]^2, x[2]*sin(x[1])-x[1]]
end

optprob = OptimizationFunction(rosenbrock, GalacticOptim.AutoForwardDiff();cons= con2_c)
prob = OptimizationProblem(optprob, x0, lcons = [-Inf,-Inf], ucons = [Inf,Inf])
sol = solve(prob, IPNewton())
@test 10*sol.minimum < l1

cons_circ = (x,p) -> [x[1]^2 + x[2]^2]
optprob = OptimizationFunction(rosenbrock, GalacticOptim.AutoForwardDiff();cons= cons_circ)
prob = OptimizationProblem(optprob, x0, lcons = [-Inf], ucons = [0.25^2])
sol = solve(prob, IPNewton())
@test sqrt(cons(sol.minimizer,nothing)[1]) ≈ 0.25 rtol = 1e-6

optprob = OptimizationFunction(rosenbrock, GalacticOptim.AutoZygote())
prob = OptimizationProblem(optprob, x0)
sol = solve(prob, ADAM(), maxiters = 1000, progress = false)
@test 10*sol.minimum < l1

prob = OptimizationProblem(optprob, x0, lb=[-1.0, -1.0], ub=[0.8, 0.8])
sol = solve(prob, Fminbox())
@test 10*sol.minimum < l1

prob = OptimizationProblem(optprob, x0, lb=[-1.0, -1.0], ub=[0.8, 0.8])
@test_broken @test_nowarn sol = solve(prob, SAMIN())
@test 10*sol.minimum < l1

using NLopt
prob = OptimizationProblem(optprob, x0)
sol = solve(prob, Opt(:LN_BOBYQA, 2))
@test 10*sol.minimum < l1

sol = solve(prob, Opt(:LD_LBFGS, 2))
@test 10*sol.minimum < l1

prob = OptimizationProblem(optprob, x0, lb=[-1.0, -1.0], ub=[0.8, 0.8])
sol = solve(prob, Opt(:LD_LBFGS, 2))
@test 10*sol.minimum < l1

sol = solve(prob, Opt(:G_MLSL_LDS, 2), nstart=2, local_method = Opt(:LD_LBFGS, 2), maxiters=10000)
@test 10*sol.minimum < l1

# using MultistartOptimization
# sol = solve(prob, MultistartOptimization.TikTak(100), local_method = NLopt.LD_LBFGS)
# @test 10*sol.minimum < l1

# using QuadDIRECT
# sol = solve(prob, QuadDirect(); splits = ([-0.5, 0.0, 0.5],[-0.5, 0.0, 0.5]))
# @test 10*sol.minimum < l1

using Evolutionary
sol = solve(prob, CMAES(μ =40 , λ = 100),abstol=1e-15)
@test 10*sol.minimum < l1

#=
using BlackBoxOptim
prob = GalacticOptim.OptimizationProblem(optprob, x0, lb=[-1.0, -1.0], ub=[0.8, 0.8])
sol = solve(prob, BBO())
@test 10*sol.minimum < l1

prob_multi = EnsembleOptimizationProblem([prob, prob])
sol_multi = solve(prob_multi, BBO(:borg_moea), FitnessScheme=ParetoFitnessScheme{2}(is_minimizing=true), ϵ=0.05, MaxSteps=500000)
@test 10*sum(sol_multi.minimum) < l1
=#

using ModelingToolkit
rosenbrock(x, p) = (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2
x0 = zeros(2)
_p = [1.0, 100.0]

f = OptimizationFunction(rosenbrock,ModelingToolkit.AutoModelingToolkit(),x0,_p,grad=true,hess=true)
prob = OptimizationProblem(f,x0,_p)
sol = solve(prob,Optim.Newton())


# computation using analytical expression

# define cache
struct LBC_cache{cP1type, cP2type, cprobtype, cpredxtype, clabeltype, crestype}
  cP1::cP1type
  cP2::cP2type
  cprob::cprobtype
  cpredx::cpredxtype
  clabel::clabeltype
  cres::crestype
end

# get optimal predictions for specific sample
function get_pred_opt_x_LBC!(x, distribution, p_range, p_tar, LBCc)
  @unpack cP1, cP2, cprob, cpredx, clabel, cres = LBCc

  cP1[1] = zero(eltype(cP1[1]))
  cP2[1] = zero(eltype(cP2[1]))
  for p in p_range
    cprob[1] = distribution(x, p)
    if p <= p_tar
      cP1[1] += cprob[1]
    else
      cP2[1] += cprob[1]
    end
  end

  if cP1[1] + cP2[1] == zero(eltype(cP1[1]))
    return zero(eltype(cP1[1]))
  else

    # adjust depending on the choice of data labelling
    # alternative: cP2[1]/(cP1[1]+cP2[1])
    return cP1[1]/(cP1[1]+cP2[1])
  end
end

function get_e_opt_p_LBC!(samples, distribution, p_range, p_tar, LBCc)
  @unpack cP1, cP2, cprob, cpredx, clabel, cres = LBCc

  cres[1] = zero(eltype(cres[1]))
  cres[2] = zero(eltype(cres[2]))
  for p in p_range

    # set label depending on current bipartition point p_tar
    if p <= p_tar
      clabel[1] = one(eltype(clabel[1]))
    else
      clabel[1] = zero(eltype(clabel[1]))
    end

    for x in samples
      cpredx[1] = get_pred_opt_x_LBC!(x, distribution, p_range, p_tar, LBCc)
      cprob[1] = distribution(x, p)

      # contribution to optimal loss
      cres[2] += cprob[1]*crossentropy(cpredx[1], clabel[1])

      # contribution to optimal classification error
      cres[1] += cprob[1]*min(cpredx[1], one(eltype(cpredx[1]))-cpredx[1])

    end
  end
  return nothing
end

# compute optimal indicator and loss of LBC across entire range of tuning parameter
function get_indicators_LBC_analytical(samples, distribution, p_range, p_range_LBC)

  caches=[LBC_cache([zero(eltype(p_range[1]))], [zero(eltype(p_range[1]))], [zero(eltype(p_range[1]))], [zero(eltype(p_range[1]))], [zero(eltype(p_range[1]))], [zero(eltype(p_range[1])),  zero(eltype(p_range[1]))]) for i in 1:length(p_range_LBC)]

  # start parallel computation for sampled values of tuning parameter
  @sync for indxp in collect(1:length(p_range_LBC))
    Threads.@spawn get_e_opt_p_LBC!(samples, distribution, p_range, p_range_LBC[indxp], caches[indxp])
  end

  # construct optimal indicator and loss
  indicator_LBC_opt = zeros(eltype(p_range_LBC[1]), length(p_range_LBC))
  loss_LBC_opt = zeros(eltype(p_range_LBC[1]), length(p_range_LBC))
  for indxp in collect(1:length(p_range_LBC))
    indicator_LBC_opt[indxp] = 1-caches[indxp].cres[1]/length(p_range)
    loss_LBC_opt[indxp] = caches[indxp].cres[2]/length(p_range)
  end

  return indicator_LBC_opt, loss_LBC_opt
end

# compute optimal indicator and loss of LBC at fixed value of tuning parameter
function get_indicators_LBC_analytical_fixed_p(samples, distribution, p_range, p_range_LBC, indxp)

  cache=LBC_cache([zero(eltype(p_range[1]))], [zero(eltype(p_range[1]))], [zero(eltype(p_range[1]))], [zero(eltype(p_range[1]))], [zero(eltype(p_range[1]))], [zero(eltype(p_range[1])), zero(eltype(p_range[1]))])

  get_e_opt_p_LBC!(samples, distribution, p_range, p_range_LBC[indxp], cache)
  acc_LBC_opt = 1-cache.cres[1]/length(p_range)
  loss_LBC_opt = cache.cres[2]/length(p_range)

  return acc_LBC_opt, loss_LBC_opt
end

# computation using neural networks

# loss function where each prediction is weighted by probability associated with input
function main_loss_LBC_weighted(NN, pnn, dataset, p_tar, p_range, inputs)
  indices = convert.(Int, dataset[1,:])
  input = inputs[:, indices]

  pred = NN(pnn)(input)
  pred = Flux.softmax(pred)
  loss = zero(eltype(p_tar))
  for i in 1:size(pred)[2]
    if p_range[Int(dataset[3, i])] <= p_tar
      loss += dataset[2, i]*MLP.crossentropy(pred[1, i], one(eltype(pred[1, i])))
    else
      loss += dataset[2, i]*MLP.crossentropy(pred[1, i], zero(eltype(pred[1, i])))
    end
  end
  return loss
end

# loss function without weighting
function main_loss_LBC_stochastic(NN, pnn, dataset, p_tar, p_range, inputs)
  indices = convert.(Int, dataset[1, :])
  input = inputs[:, indices]

  pred = NN(pnn)(input)
  pred = Flux.softmax(pred)
  loss = zero(eltype(p_tar))
  for i in 1:size(pred)[2]
    if p_range[Int(dataset[3, i])] <= p_tar
      loss += MLP.crossentropy(pred[1, i], one(eltype(pred[1, i])))
    else
      loss += MLP.crossentropy(pred[1, i], zero(eltype(pred[1, i])))
    end
  end
  return loss/length(dataset[1, :])
end

# train NN in LBC based on weighted loss function
function train_LBC_weighted(NN, pnn, dataset, epochs, p_range, opt, p_tar, batchsize, n_batches, inputs; verbose=false, saveat=epochs)
  NN_logger = [zeros(eltype(p_range[1]), length(pnn))]
  pred_logger = [(zeros(eltype(p_range[1]), 1), zeros(eltype(p_range[1]), 1))]

  pnn_best = zeros(eltype(pnn[1]), length(pnn))
  best_loss = zero(eltype(pnn[1]))

  losses = zeros(eltype(p_range[1]), epochs)
  indices = collect(1:size(dataset)[2])
  for epoch in 1:epochs
    Random.shuffle!(indices)
    grad = zeros(eltype(p_range[1]), length(pnn))

    # compute loss and gradient batchwise
    for batch in 1:n_batches
      randint = get_batches(indices, batchsize, batch)
      data = reshape(dataset[:, randint], 3, length(randint))
      val, back = Flux.Zygote.pullback(p -> main_loss_LBC_weighted(NN, p,data, p_tar, p_range, inputs), pnn)
      grad .+= back(one(val))[1]
      losses[epoch] += val
    end

    # update NN parameters based on overall gradient
    Flux.Optimise.update!(opt, pnn, grad./length(p_range))
    losses[epoch] = losses[epoch]/length(p_range)

    # keep track of best performing NN
    if epoch == 1
      best_loss = losses[epoch]
    elseif losses[epoch] < best_loss
      println("better")
      println(best_loss)
      println(losses[epoch])
      pnn_best = deepcopy(pnn)
      best_loss = losses[epoch]
    end

    # save at regular intervals
    if epoch % saveat == 0
      push!(NN_logger, pnn_best)
      push!(pred_logger, predict_LBC(dataset, p_range, pnn_best, NN, batchsize, n_batches, p_tar, inputs, calc_loss=true))
    end

    verbose && println("epoch: $epoch")
    verbose && println("loss: "*string(losses[epoch]))
  end
  return losses, NN_logger, pred_logger
end

# train NN in LBC based on ``stochastic'' loss function
function train_LBC_stochastic(NN, pnn, dataset, epochs, p_range, opt, p_tar, batchsize, n_batches_stochastic, n_batches, inputs; verbose=false, saveat=epochs)
  NN_logger = [zeros(eltype(p_range[1]), length(pnn))]
  pred_logger = [(zeros(eltype(p_range[1]), 1), zeros(eltype(p_range[1]), 1))]

  pnn_best = deepcopy(pnn)
  best_loss = zero(eltype(pnn[1]))

  losses = zeros(eltype(p_range[1]), epochs)
  for epoch in 1:epochs

    # compute loss and gradient batchwise
    for batch in 1:n_batches_stochastic
      randint = sample(1:length(dataset[1, :]), Weights(dataset[2, :]), batchsize)
      data = reshape(dataset[:, randint], 3, length(randint))
      val, back = Flux.Zygote.pullback(p -> main_loss_LBC_stochastic(NN, p, data, p_tar, p_range, inputs), pnn)
      grad = back(one(val))[1]

      # update NN parameters based on batch gradient
      Flux.Optimise.update!(opt, pnn, grad)
      losses[epoch] += val
    end
    losses[epoch] = losses[epoch]/n_batches_stochastic

    # keep track of best performing NN
    if epoch == 1
      best_loss = losses[epoch]
    elseif losses[epoch] < best_loss
      println("better")
      println(best_loss)
      println(losses[epoch])
      pnn_best = deepcopy(pnn)
      best_loss = losses[epoch]
    end

    # save at regular intervals
    if epoch % saveat == 0
      push!(NN_logger, pnn_best)
      push!(pred_logger, predict_LBC(dataset, p_range, pnn_best, NN, batchsize, n_batches, p_tar, inputs,calc_loss=true))
    end

    verbose && println("epoch: $epoch")
    verbose && println("loss: "*string(losses[epoch]))
  end
  return losses, NN_logger, pred_logger
end

# compute indicator and loss value based on current NN
function predict_LBC(dataset, p_range, pnn, NN, batchsize, n_batches, p_tar, inputs; calc_loss=false, loss=zero(eltype(p_range[1])))
  error = zero(eltype(p_range[1]))

    # compute indicator and loss batch-wise
  indices = collect(1:size(dataset)[2])
  for batch in 1:n_batches
    randint = get_batches(indices, batchsize, batch)
    data = reshape(dataset[:, randint], 3, length(randint))

    indicess = convert.(Int, data[1, :])
    input = inputs[:, indicess]

    pred = NN(pnn)(input)
    pred = Flux.softmax(pred)
    for indxp in 1:length(randint)
      if p_range[Int(data[3, indxp])]<=p_tar
        label = one(eltype(pred[1, indxp]))
      else
        label = zero(eltype(pred[1, indxp]))
      end

      if pred[1, indxp] <= 0.5*one(eltype(pred[1, indxp]))
        binary_pred = zero(eltype(pred[1, indxp]))
      else
        binary_pred = one(eltype(pred[1, indxp]))
      end

      error += data[2, indxp]*abs(label-binary_pred)
    end

    if calc_loss
      loss += main_loss_LBC_weighted(NN, pnn, data, p_tar, p_range, inputs)
    end

  end

  if calc_loss
    loss = loss/length(p_range)
  end

  return [1-error/length(p_range)], [loss]
end

# perform LBC for fixed bipartition point using neural networks
function get_indicators_LBC_numerical_fixed_p(pnn, NN, dataset, epochs, p_range, dp, opt, p_range_LBC, indx_ptar, inputs; verbose=false, trained=false, saveat=epochs, batchsize=length(dataset[1, :]), stochastic=false, n_batches_train_stochastic=10, batchsize_stochastic=length(dataset[1, :]), train_opt=false)

  n_batches = ceil(eltype(batchsize), size(dataset)[2]/batchsize)
  p_tar = p_range_LBC[indx_ptar]

  if !trained
    if stochastic
      losses, NN_logger, pred_logger = train_LBC_stochastic(NN, pnn, dataset, epochs, p_range, deepcopy(opt), p_tar, batchsize_stochastic, n_batches_train_stochastic, n_batches, inputs, verbose=verbose, saveat=saveat)

    else
      losses, NN_logger, pred_logger = train_LBC_weighted(NN, pnn, dataset, epochs, p_range, deepcopy(opt), p_tar, batchsize, n_batches, inputs, verbose=verbose, saveat=saveat)
    end

    pred_logger = pred_logger[2:end]
    NN_logger = NN_logger[2:end]

  else
    accuracy, loss = predict_LBC(dataset, p_range, pnn, NN, batchsize, n_batches, p_range_LBC[indx_ptar], inputs, calc_loss=true)

    losses = zeros(eltype(p_range[1]), epochs)
    NN_logger = [pnn]
    pred_logger = [(accuracy, loss)]

  end

  return pred_logger, losses, NN_logger
end


using ScikitLearn
using RDatasets
using CSV
using DataFrames
using MLLabelUtils
using StatsBase: standardize, mean
@sk_import model_selection:train_test_split
@sk_import preprocessing:OneHotEncoder
iris = dataset("datasets", "iris")

# If you want to use ScikitLearn, you need to convert DataFrames to Arrays
# ScikitLearn only accepts data as arrays.
X = convert(Array, iris[!, 1:4])
y = convert(Array, iris[!, 5])

# train_test_split is a function from ScikitLearn, and splits the data
# according to the value passed to test_size. Here we are saying
# that we need to allocate 33% of the raw data as test data
X_train, X_test, y_train, y_test = train_test_split(X, y; test_size=0.33, random_state=42)

data = dataset("datasets", "airquality")
describe(data)
describe(dropmissing(data))

dropmissing!(data) # <- this will permananty change the data by dropping missing values

data = dataset("datasets", "airquality") # re-loading because we need a fresh copy of the data for next e.g.
# In some other cases, you just want to replace the missing values with the mean of the data or with 0.
# In those cases you can use coalesce function from julia's base library

# To replace the missing values with 0
data.Ozone = coalesce.(data.Ozone, 0)
# Let's check the data once again
describe(data)

# For Solar.R let's replace it with the mean
data."Solar.R" = coalesce.(data."Solar.R", mean(skipmissing(data."Solar.R")))
describe(data)

standardize(UnitRangeTransform, X; dims=1)
standardize(ZScoreTransform, X; dims=1, center=false, scale=true)

onehotencoded = convertlabel(LabelEnc.OneOfK{Float32}, y; obsdim=1)

using Discretizers
bin_edges = [0.1, 0.5, 0.7]
lindisc = LinearDiscretizer(bin_edges);

cat = [:yes, :no, :maybe]
catdisc = CategoricalDiscretizer(cat)
decode(catdisc, z)
z = encode(lindisc, a)


using DelayEmbeddings, StatsBase, LinearAlgebra
export PseudoPeriodic, noiseradius

"""
    PseudoPeriodic(d, τ, ρ, shift = true)

Create surrogates suitable for pseudo-periodic signals. They retain the periodic structure
of the signal, while inter-cycle dynamics that are either deterministic or correlated
noise are destroyed (for appropriate `ρ` choice).
Therefore these surrogates are suitable to test the null hypothesis
that the signal is a periodic orbit with uncorrelated noise[^Small2001].

Arguments `d, τ, ρ` are as in the paper, the embedding dimension, delay time and
noise radius. The method works by performing a delay coordinates ambedding via the
library [DynamicalSystems.jl](https://juliadynamics.github.io/DynamicalSystems.jl/dev/embedding/reconstruction/).
See its documentation for choosing appropriate values for `d, τ`. For `ρ`, we have implemented
the method proposed in the paper in the function [`noiseradius`](@ref).

The argument `shift` is not discussed in the paper. If `shift=false`
we adjust the algorithm so that there is little phase shift between the
periodic component of the original and surrogate data.

See also [`CycleShuffle`](@ref).

[^Small2001]: Small et al., Surrogate test for pseudoperiodic time series data, [Physical Review Letters, 87(18)](https://doi.org/10.1103/PhysRevLett.87.188101)
"""
struct PseudoPeriodic{T<:Real} <: Surrogate
    d::Int
    τ::Int
    ρ::T
    shift::Bool
end
PseudoPeriodic(d, t, r) = PseudoPeriodic(d, t, r, true)

function surrogenerator(x::AbstractVector, pp::PseudoPeriodic, rng = Random.default_rng())
    # in the following symbol `y` stands for `s` of the paper
    d, τ = getfield.(Ref(pp), (:d, :τ))
    N = length(x)
    z = embed(x, d, τ)
    Ñ = length(z)
    w = zeros(eltype(z), Ñ-1) # weights vector
    y = Dataset([z[1] for i in 1:N])
    s = similar(x)
    init = (y = y, w = w, z = z)
    return SurrogateGenerator(pp, x, s, init, rng)
end

function (sg::SurrogateGenerator{<:PseudoPeriodic})()
    y, z, w = getfield.(Ref(sg.init), (:y, :z, :w))
    ρ, shift = getfield.(Ref(sg.method), (:ρ, :shift))
    pseudoperiodic!(sg.s, y, sg.x, z, w, ρ, shift, sg.rng)
end
# Low-level method, also used in `noiseradius`
function pseudoperiodic!(s, y, x, z, w, ρ, shift, rng)
    N, Ñ = length.((x, z))
    y[1] = shift ? rand(rng, z.data) : z[1]
    @inbounds for i in 1:N-1
        w .= (exp(-norm(z[t] - y[i])/ρ) for t in 1:Ñ-1)
        j = sample(rng, 1:Ñ-1, pweights(w))
        y[i+1] = z[j+1]
    end
    for (i, p) in enumerate(y); s[i] = p[1]; end
    return s
end

"""
    noiseradius(x::AbstractVector, d::Int, τ, ρs, n = 1) → ρ

Use the proposed* algorithm of[^Small2001] to estimate optimal `ρ` value for
[`PseudoPeriodic`](@ref) surrogates, where `ρs` is a vector of possible `ρ` values.
*The paper is ambiguous about exactly what to calculate. Here we count how many times
we have pairs of length-2 that are identical in `x` and its surrogate, but **are not**
also part of pairs of length-3.

This function directly returns the arg-maximum of the evaluated distribution of these counts
versus `ρ`, use `TimeseriesSurrogates._noiseradius` with same arguments to get the actual
distribution. `n` means to repeat τhe evaluation `n` times, which increases accuracy.

[^Small2001]: Small et al., Surrogate test for pseudoperiodic time series data, [Physical Review Letters, 87(18)](https://doi.org/10.1103/PhysRevLett.87.188101)
"""
function noiseradius(x::AbstractVector, d::Int, τ, ρs, n = 1)
    l2n = _noiseradius(x, d, τ, ρs, n)
    return ρs[argmax(l2n)]
end

function _noiseradius(x::AbstractVector, d::Int, τ, ρs, n = 1)
    l2n = noiseradius(surrogenerator(x, PseudoPeriodic(d, τ, ρs[1])), ρs, n)
end

function noiseradius(sg::SurrogateGenerator{<:PseudoPeriodic}, ρs, n = 1)
    l2n = zero(ρs) # length-2 number of points existing in both timeseries
    y, z, w = getfield.(Ref(sg.init), (:y, :z, :w))
    x, s = sg.x, sg.s
    N = length(sg.x)
    @inbounds for _ in 1:n
        for (ℓ, ρ) in enumerate(ρs)
            s̄ = pseudoperiodic!(s, y, x, z, w, ρ, sg.method.shift, sg.rng)
            for i in 1:N-1
                # TODO: This can be optimized heavily: checking x[i+2] already tells us
                # that we shouldn't check x[i+1] on the next iteration.
                l2n[ℓ] += count(j -> s̄[j]==x[i] && s̄[j+1]==x[i+1] && s̄[j+2]!=x[i+2], 1:N-2)
            end
        end
    end
    # TODO: here we can directly find maximum and return it as a single number
    return l2n
end


using LatticeBasisReduction
using Base.Test

# Custom rounding
@test LatticeBasisReduction.round_ties_down(1.3) == 1
@test LatticeBasisReduction.round_ties_down(1.5) == 1
@test LatticeBasisReduction.round_ties_down(1.6) == 2
@test LatticeBasisReduction.round_ties_down(2) == 2
@test LatticeBasisReduction.round_ties_down(-0.5) == -1

# Verify gram schmidt process with example 3.3 from Bremner
eps = 1e-3
@test_approx_eq_eps gram_schmidt([3 -1 5; -5 2 -1; -3 9 2]) [3 -1 5; -109/35 48/35 15/7; 1521/566 1859/283 -169/566] eps

# Verify LLL algorithm
x = [[-2 7 7 -5]
     [3 -2 6 -1]
     [2 -8 -9 -7]
     [8 -9 6 -4]]

@test_approx_eq_eps LLL(x, 1) [2 3 1 1; 2 0 -2 -4; -2 2 3 -3; 3 -2 6 -1] eps
@test_approx_eq_eps LLL(x, 1.) [2 3 1 1; 2 0 -2 -4; -2 2 3 -3; 3 -2 6 -1] eps
@test_approx_eq_eps LLL(Array{Float64}(x), 1) [2 3 1 1; 2 0 -2 -4; -2 2 3 -3; 3 -2 6 -1] eps

# Verify diophantine approximation for a single number
@test diophantine(float(π)) == 355//113

# Verify gauss for minimal basis of a lattice in R^2
x = [-56, 43]; y = [95, -73]
@test gauss(x,y) == ([1,1],[-1,2])
@test gauss(y,x) == ([1,1],[-1,2])

# Verify Simulataneous Diophantine Approximation
nums, dem = sda(sqrt(2), sqrt(3), sqrt(5), sqrt(7), ε = eps)
@test_approx_eq_eps nums[1]/dem sqrt(2) eps
@test_approx_eq_eps nums[2]/dem sqrt(3) eps
@test_approx_eq_eps nums[3]/dem sqrt(5) eps
@test_approx_eq_eps nums[4]/dem sqrt(7) eps


function loglikelihoods!(LL::AbstractMatrix, hmm::AbstractHMM{Univariate}, observations)
    T, K = size(observations, 1), size(hmm, 1)
    @argcheck size(LL) == (T, K)
    @inbounds for i in OneTo(K), t in OneTo(T)
        LL[t, i] = logpdf(hmm.B[i], observations[t])
    end
end

function loglikelihoods!(LL::AbstractMatrix, hmm::AbstractHMM{Multivariate}, observations)
    T, K = size(observations, 1), size(hmm, 1)
    @argcheck size(LL) == (T, K)
    @inbounds for i in OneTo(K), t in OneTo(T)
        LL[t, i] = logpdf(hmm.B[i], view(observations, t, :))
    end
end

"""
    loglikelihoods(hmm, observations; robust) -> Matrix

Return the log-likelihood per-state and per-observation.

**Output**
- `Matrix{Float64}`: log-likelihoods matrix (`T x K`).

**Example**
```julia
using Distributions, HMMBase
hmm = HMM([0.9 0.1; 0.1 0.9], [Normal(0,1), Normal(10,1)])
y = rand(hmm, 1000)
LL = likelihoods(hmm, y)
```
"""
function loglikelihoods(hmm::AbstractHMM, observations; logl = nothing, robust = false)
    (logl !== nothing) && deprecate_kwargs("logl")
    T, K = size(observations, 1), size(hmm, 1)
    LL = Matrix{Float64}(undef, T, K)

    loglikelihoods!(LL, hmm, observations)
    if robust
        replace!(LL, -Inf => nextfloat(-Inf), Inf => log(prevfloat(Inf)))
    end
    LL
end


# For this PS3, I primary used Prof.Ransom's code as a referecne point to write down my code
# Also, I tried to write explantation with each line of code to better understand what is going on
# Also Prof. Ransom mentioned that its extremely imp to visualize what is going on before writng a loop - This has helped me quite a bit

cd("D:\\Semester 3_Fall 2020\\ECON 6453 - Econometrics III\\ProblemSets\\PS3-gev")
using Random, Distributions, JLD2, DataFrames, CSV, LinearAlgebra, Statistics, FreqTables, TexTables, Optim, HTTP, GLM, ForwardDiff

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 4: All Wrapping Function
#:::::::::::::::::::::::::::::::::::::::::::::::::::
function problemset_3()

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 1: Multinomial logit with Alternate Specific Covariates
#:::::::::::::::::::::::::::::::::::::::::::::::::::
url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2020/master/ProblemSets/PS3-gev/nlsw88w.csv"

df = CSV.read(HTTP.get(url).body)

X = [df.age df.white df.collgrad]
Z = hcat(df.elnwage1, df.elnwage2, df.elnwage3, df.elnwage4,df.elnwage5, df.elnwage6, df.elnwage7, df.elnwage8)
y = df.occupation

function mlogit(alpha, X, Z, y)
    beta = alpha[1:end-1] # Indexing the parameter so coefficient on X is the first element
    gamma = alpha[end] # Indexing the parameter so coefficient on Z is the last element
    K = size(X,2) # K is the no. of covariates in X, which in this case are 3
    J = length(unique(y)) # J is the no. of coefficients, which are 8 in this case (So J-1 = 7)
    N = length(y) # N is a vector with the same length as y (not sure why we are defining N)
    bigY = zeros(N,J) # Simply defining the matrix of dependent variables Y, which are going to be dummies in case of Multinomial Logit

    for j=1:J # Simply defining j to be something, next step, j represents column
        bigY[:,j] = y.==j # Making a dummies for each of the corresponding occupation in y
    end
    bigBeta = [reshape(beta,K,J-1) zeros(K)] # Reshaping into a matrix with K rows and J-1 columns
    # Adding a column of zeros in the end to make it conform with bigY and also to normalize the parameters of the last occupation to be zero

    T = promote_type(eltype(X),eltype(alpha)) # To get the standard errors we promote the type
    num   = zeros(T,N,J) # Define numerator T-type , N-rows, J-columns
    dem   = zeros(T,N) # Define denominator T-type , N-rows
    for j=1:J
        num[:,j] = exp.(X*bigBeta[:,j] .+ (Z[:,j] .- Z[:,J])*gamma) # Numerator Likelihood function
        dem .+= num[:,j] # Denominator Likelihood function: Simply the Sum of all the numerators
    end

    P = num./repeat(dem,1,J) # Putting together numerator and denominator

    loglike = -sum( bigY.*log.(P) ) # Since Optim minimizes the function, we take its negative to maximize

    return loglike # This is used for storing the function parameters
end

startvals = [2*rand(7*size(X,2)).-1; .1]
td = TwiceDifferentiable(alpha -> mlogit(alpha, X, Z, y), startvals; autodiff = :forward)

# Use Optim #
α_hat_multinomial = optimize(td, startvals, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
α_hat_multinomial_ad = α_hat_multinomial.minimizer

# Find Hessian #
H  = Optim.hessian!(td, α_hat_multinomial_ad) # Finding the Hessian Matrix
α_hat_multinomial_ad_se = sqrt.(diag(inv(H))) # Standard Errors of alpha_hat

println([α_hat_multinomial_ad α_hat_multinomial_ad_se])

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 2: Interpretation of gamma_hat
#:::::::::::::::::::::::::::::::::::::::::::::::::::
#  Since it is a linear-log model, The coefficient gamma/100 represents the change in utility with a 1% increase in wage.

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 3: Nested Logit Model
#:::::::::::::::::::::::::::::::::::::::::::::::::::
function nested_logit(alpha, X, Z, y, nesting_structure)

    beta = alpha[1:end-3]
    lambda = alpha[end-2:end-1]
    gamma = alpha[end]
    K = size(X,2)
    J = length(unique(y))
    N = length(y)
    bigY = zeros(N,J)
    for j=1:J
        bigY[:,j] = y.==j
    end
    bigBeta = [repeat(beta[1:K],1,length(nesting_structure[1])) repeat(beta[K+1:2K],1,length(nesting_structure[2])) zeros(K)]
    # Betas here vary with the nest structure

    T = promote_type(eltype(X),eltype(alpha))
    num   = zeros(T,N,J)
    lidx  = zeros(T,N,J) # Linear Index: Numerator is divided by lambda in choice probabilities
    # lidx is a part of the numerator without the exponent
    dem   = zeros(T,N)

    for j=1:J
        if j in nesting_structure[1]
            lidx[:,j] = exp.( (X*bigBeta[:,j] .+ (Z[:,j] .- Z[:,J])*gamma)./lambda[1] )
        elseif j in nesting_structure[2]
            lidx[:,j] = exp.( (X*bigBeta[:,j] .+ (Z[:,j] .- Z[:,J])*gamma)./lambda[2] )
        else
            lidx[:,j] = exp.(zeros(N))
        end
    end

    for j=1:J
        if j in nesting_structure[1]
            num[:,j] = lidx[:,j].*sum(lidx[:,nesting_structure[1][:]];dims=2).^(lambda[1]-1)
        elseif j in nesting_structure[2]
            num[:,j] = lidx[:,j].*sum(lidx[:,nesting_structure[2][:]];dims=2).^(lambda[2]-1)
            # only sum up the colums of the nesting structure
        else
            num[:,j] = lidx[:,j]
        end

        dem .+= num[:,j]
    end

    P = num./repeat(dem,1,J)

    loglike = -sum( bigY.*log.(P) )

    return loglike
end

nesting_structure = [[1 2 3], [4 5 6 7]]
startvals = [2*rand(2*size(X,2)).-1; 1; 1; .1]

td2 = TwiceDifferentiable(alpha -> nested_logit(alpha, X, Z, y, nesting_structure), startvals; autodiff = :forward)

# Using Optim
α_hat_nlogit = optimize(td2, startvals, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
α_hat_nlogit_ad = α_hat_nlogit.minimizer

# Estimating Hessian
H  = Optim.hessian!(td2, α_hat_nlogit_ad)
α_hat_nlogit_ad_se = sqrt.(diag(inv(H)))

println([α_hat_nlogit_ad α_hat_nlogit_ad_se])

return nothing
end

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 4: All Wrapping Function
#:::::::::::::::::::::::::::::::::::::::::::::::::::
problemset_3()


# # Layouts versus drawings of wiring diagrams
#
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/graphics/layouts_vs_drawings.ipynb)
#
# In Catlab, layout and drawing (rendering) of wiring diagrams are mostly
# decoupled. This notebook shows how to lay out diagrams using Graphviz's
# rank-based layout or Catlab's series-parallel layout and then render them
# using Compose.jl or TikZ.

# The morphism we will visualize is:

using Catlab.Theories

X = Ob(FreeSymmetricMonoidalCategory, :X)
f, g, h = (Hom(sym, X, X) for sym in (:f, :g, :h))

expr = otimes(f, compose(f,g), compose(f,g,h))

# Let's convert this expression into a wiring diagram. This yields a purely
# combinatorial object, as evidenced by its underlying graph.

using Catlab.WiringDiagrams, Catlab.Graphics

diagram = to_wiring_diagram(expr)
WiringDiagrams.graph(diagram)

# ## Graphviz layout

# Calling `to_graphviz` both lays out and draws the diagram, entirely within
# Graphviz.

to_graphviz(diagram, orientation=LeftToRight)

# To get just the layout from Graphviz, we call `graphviz_layout` instead. We
# can then render this layout using Compose.jl. Note that the Graphviz layout
# has units in points.

import Compose

layout = graphviz_layout(diagram, orientation=LeftToRight)
layout_to_composejl(layout, base_unit=Compose.pt)

# The same layout can be rendered in TikZ:

import TikzPictures

layout_to_tikz(layout, base_unit="1pt")

# ## Series-parallel layout

# Catlab has its own layout system based on series-parallel decomposition. In
# this case, the layout exactly recovers the structure of the morphism
# expression created at the beginning.

layout = layout_diagram(FreeSymmetricMonoidalCategory, diagram,
                        orientation=LeftToRight)
layout_to_composejl(layout)
#-
layout_to_tikz(layout)


const use_return  = haskey(ENV, "USE_RETURN" ) ? parse(Bool, ENV["USE_RETURN"] ) : false
const USE_GPU     = haskey(ENV, "USE_GPU"    ) ? parse(Bool, ENV["USE_GPU"]    ) : false
const do_viz      = haskey(ENV, "DO_VIZ"     ) ? parse(Bool, ENV["DO_VIZ"]     ) : true
const do_save     = haskey(ENV, "DO_SAVE"    ) ? parse(Bool, ENV["DO_SAVE"]    ) : false
const do_save_viz = haskey(ENV, "DO_SAVE_VIZ") ? parse(Bool, ENV["DO_SAVE_VIZ"]) : false
const nx          = haskey(ENV, "NX"         ) ? parse(Int , ENV["NX"]         ) : 512
###
using ParallelStencil
using ParallelStencil.FiniteDifferences1D
@static if USE_GPU
    @init_parallel_stencil(CUDA, Float64, 1)
else
    @init_parallel_stencil(Threads, Float64, 1)
end
using Plots, Printf, LinearAlgebra, MAT

@parallel function compute_flux!(qHx, qHx2, H, D, θr_dτ, dx)
    @all(qHx)  = (@all(qHx) * θr_dτ - D * @d(H) / dx) / (1.0 + θr_dτ)
    @all(qHx2) = -D * @d(H) / dx
    return
end

@parallel function compute_update!(H, Hold, qHx, dτ_ρ, dt, dx)
    @inn(H) = (@inn(H) +  dτ_ρ * (@inn(Hold) / dt - @d(qHx) / dx)) / (1.0 + dτ_ρ / dt)
    return
end

@parallel function check_res!(ResH, H, Hold, qHx2, dt, dx)
    @all(ResH) = -(@inn(H) - @inn(Hold)) / dt - @d(qHx2) / dx
    return
end

@views function diffusion_1D_()
    # Physics
    lx     = 10.0       # domain size
    D      = 1          # diffusion coefficient
    ttot   = 1.0        # total simulation time
    dt     = 0.2        # physical time step
    # Numerics
    # nx     = 2*256      # numerical grid resolution
    tol    = 1e-8       # tolerance
    itMax  = 1e5        # max number of iterations
    nout   = 10         # tol check
    CFL    = 0.99       # CFL number
    # Derived numerics
    dx     = lx / nx      # grid size
    Vpdτ   = CFL * dx
    Re     = π + sqrt(π^2 + (lx^2 / D / dt)) # Numerical Reynolds number
    θr_dτ  = lx / Vpdτ / Re
    dτ_ρ   = Vpdτ * lx / D / Re
    xc     = LinRange(dx/2, lx - dx/2, nx)
    # Array allocation
    qHx    = @zeros(nx-1)
    qHx2   = @zeros(nx-1)
    ResH   = @zeros(nx-2)
    dH     = @zeros(nx-2)
    # Initial condition
    H0     = Data.Array(exp.(-(xc .- lx/2).^2 / D))
    Hold   = @ones(nx) .* H0
    H      = @ones(nx) .* H0
    t = 0.0; it = 0; ittot = 0; nt = Int(ceil(ttot/dt))
    # Physical time loop
    while it < nt
        iter = 0; err = 2 * tol
        # Pseudo-transient iteration
        while err > tol && iter < itMax
            @parallel compute_flux!(qHx, qHx2, H, D, θr_dτ, dx)
            @parallel compute_update!(H, Hold, qHx, dτ_ρ, dt, dx)
            iter += 1
            if iter % nout == 0
                @parallel check_res!(ResH, H, Hold, qHx2, dt, dx)
                err = norm(ResH) / sqrt(length(ResH))
            end
        end
        ittot += iter; it += 1; t += dt
        Hold .= H
        if isnan(err) error("NaN") end
    end
    # Analytic solution
    Hana = 1 / sqrt(4 * (ttot + 1 / 4)) * exp.(-(xc .- lx/2).^2 / (4 * D * (ttot + 1 / 4)))
    @printf("Total time = %1.2f, time steps = %d, nx = %d, iterations tot = %d, error vs analytic = %1.2e \n", round(ttot, sigdigits=2), it, nx, ittot, norm(Array(H) - Hana) / sqrt(nx))
    # Visualise
    if do_viz plot(xc, Array(H0), linewidth=3); display(plot!(xc, [Array(H) Array(Hana)], legend=false, framestyle=:box, linewidth=3, xlabel="lx", ylabel="H", title="linear diffusion (nt=$it, iters=$ittot)")) end
    if do_save
        !ispath("../output") && mkdir("../output")
        open("../output/out_diff_1D_lin.txt","a") do io
            println(io, "$(nx) $(ittot) $(nt)")
        end
    end
    if do_save_viz
        !ispath("../out_visu") && mkdir("../out_visu")
        matwrite("../out_visu/diff_1D_lin.mat", Dict("H_1D"=> Array(H), "xc_1D"=> Array(xc)); compress = true)
    end
    return xc, H
end

if use_return
    xc, H = diffusion_1D_();
else
    diffusion_1D = begin diffusion_1D_(); return; end
end


using PyPlot
using Distributions

function plot_histogram(distribution, n)
    epsilon_values = rand(distribution, n)  # n draws from distribution
    plt[:hist](epsilon_values)
end

lp = Laplace()
lg = Logistic()
cc = Cauchy()
prt = Pareto()
#plot_histogram(lp, 500)
#plot_histogram(lg, 500)
#plot_histogram(cc, 500)
plot_histogram(prt, 100)
legend()


using TensorFlow
include("mnist_loader.jl")

loader = DataLoader()

sess = Session(Graph())

x = placeholder(Float32)
y_ = placeholder(Float32)

W = Variable(zeros(Float32, 784, 10))
b = Variable(zeros(Float32, 10))

run(sess, global_variables_initializer())

y = nn.softmax(x*W + b)

cross_entropy = reduce_mean(-reduce_sum(y_ .* log(y), axis=[2]))
train_step = train.minimize(train.GradientDescentOptimizer(.00001), cross_entropy)

correct_prediction = indmax(y, 2) .== indmax(y_, 2)
accuracy=reduce_mean(cast(correct_prediction, Float32))

for i in 1:1000
    batch = next_batch(loader, 100)
    run(sess, train_step, Dict(x=>batch[1], y_=>batch[2]))
end

testx, testy = load_test_set()

println(run(sess, accuracy, Dict(x=>testx, y_=>testy)))


using Plasmo
using JuMP
using Gurobi

graph = ModelGraph()
setsolver(graph, GurobiSolver())

m1 = Model()
@variable(m1,x[1:2], Bin)
@constraint(m1, x[1] + x[2] <= 1)
@objective(m1, Max, 16x[1] + 10x[2])
n1 = add_node!(graph)
setmodel(n1,m1)

m2 = Model()
@variable(m2, x[1:2], Bin)
@variable(m2, y[3:4], Bin)
@constraint(m2, sum(y[i] for i in 3:4) <= 1)
@constraint(m2, 8x[1] + 2x[2] + y[3] + 4y[4] <= 10)
@objective(m2, Max, 4y[4])
n2 = add_node!(graph)
setmodel(n2,m2)

@linkconstraint(graph, [i in 1:2], n1[:x][i] == n2[:x][i])

#Get all of the link constraints in a graph
links = Plasmo.getlinkconstraints(graph)
for link in links
    println(link)
end

solve(graph)

println("n1[:x]= ",JuMP.getvalue(n1[:x]))
println("n2[:x]= ",JuMP.getvalue(n2[:x]))
println("n2[:y]= ",JuMP.getvalue(n2[:y]))


@testset "CSAMaxOrd" begin

  function sum_of_two_squares(a::fmpz)
    for i=1:Int(root(a,2))
      for j=1:Int(root(a,2))
        if a==i^2+j^2
          return true
        end
      end
    end
    return false
  end

  Qx,x=PolynomialRing(FlintQQ, "x")
  for b in Hecke.squarefree_up_to(100)[2:end]
    K,a=NumberField(x^2-b)
    O=maximal_order(K);
    cocval=Array{nf_elem,2}(2, 2)
    G=[Hecke.NfToNfMor(K,K,a),Hecke.NfToNfMor(K,K,-a)]
    cocval[1,1]=K(1)
    cocval[1,2]=K(1)
    cocval[2,1]=K(1)
    cocval[2,2]=K(-1)
    A=Hecke.CrossedProductAlgebra(K,G,cocval)
    if Hecke.issplit(A)
      @test sum_of_two_squares(fmpz(b))
    else
      @test !sum_of_two_squares(fmpz(b))
    end
  end

  A=Hecke.quaternion_algebra(4,36)
  @test Hecke.issplit(A)
  A=Hecke.quaternion_algebra(-1,-1)
  O=Hecke.AlgAssOrd(A, [A[i] for i=1:4])
  @test Hecke.schur_index_at_real_plc(O)==2

end


struct MixRule{M<:ThermoModel,X} <: ThermoModel
    model::M
    xn::X
end

function mol_fraction(st::MixRule{M,X}) where {M,X}
    return st.xn
end

function thermomodel(st::MixRule{M,X}) where {M,X}
    return st.model
end

abstract type AbstractMixtureRule end

"""
    LinarMixing(p)

corresponds to the mix rule Σxᵢpᵢ.

"""
struct LinearMixing{M} <: AbstractMixtureRule
    prop::M
end

function mixrule(model::T,x) where T<: LinearMixing
    return dot(model.prop,x)
end
"""
    CuadraticMixing(p,f=(x,y)/2,Aij= nothing)

corresponds to the mix rule ΣAᵢⱼxᵢxⱼf(pᵢ,pⱼ)
where f is an averaging function. That means:
- f(pᵢ,pⱼ) = f(pⱼ,pᵢ)
- f(pₖ,pₖ) = pₖ

`Aij` is the weighed interaction matrix. if not provided, A[i,j] == Aᵢⱼ == 1
"""
struct QuadraticMixing{P,F,M} <: AbstractMixtureRule
    prop::P
    func::F
    Aij::M
end

function QuadraticMixing(prop)
    return QuadraticMixing(prop,(x,y)->0.5*(x+y),nothing)
end

function QuadraticMixing(prop,f)
    return QuadraticMixing(prop,f,nothing)
end

function mixrule(model::QuadraticMixing,x)
    return mixing_rule(model.func, x, model.prop,model.Aij)
end


## import Helmholtz-solver library
using WavesAndEigenvalues.Helmholtz
## load mesh from file
mesh=Mesh("../NTNU/NTNU.msh",scale=1.0)
#create unit and full mesh from half-cell
# Important: full merges all the domains that are copied from the unit cell
# unit: keeps the domains separate and appends numbers!
doms=[("Interior",:full),("Inlet",:full), ("Outlet_high",:full), ("Outlet_low",:full), ("Flame",:unit),]#
collect_lines!(mesh)
unit_mesh=extend_mesh(mesh,doms,unit=true)
full_mesh=extend_mesh(mesh,doms,unit=false)

## Helmholtz solver
#define speed of sound field
function speedofsound(x,y,z)
    if z<0.415
        return 347.0#m/s
    else
        return 850.0#m/s
    end
end

c=generate_field(unit_mesh,speedofsound,0)
C=generate_field(full_mesh,speedofsound,0)
#D=Dict()

#vtk_write("NTNU_c",full_mesh,D)
##describe model
D=Dict()
D["Interior"]=(:interior,())
D["Outlet_high"]=(:admittance, (:Y_in,1E15))
D["Outlet_low"]=(:admittance, (:Y_out,1E15))
##discretize models
L=discretize(full_mesh, D, C, order=:1,)
l=discretize(unit_mesh, D, c, order=:1, b=:b)
## solve model using beyn
#(type Γ by typing \Gamma and Enter)
Γ=[150.0+5.0im, 150.0-5.0im, 1000.0-5.0im, 1000.0+5.0im].*2*pi #corner points for the contour (in this case a rectangle)
Ω, P = beyn(L,Γ,l=10,N=64, output=true)
ω, p = beyn(l,Γ,l=10,N=128, output=true)
##verify beyn's solution by local householder iteration
tol=1E-10
D=Dict()
for idx=1:length(Ω)
    sol,nn,flag=householder(L,Ω[idx],maxiter=6,order=5,tol=tol,n_eig_val=3,output=false);
    mode="[$(string(round((Ω[idx]/2/pi),digits=2)))]Hz"
    if abs(sol.params[:ω]-Ω[idx])<5*tol
        println("Kept:$mode")
        D["abs:"*mode]=abs.(sol.v)./maximum(abs.(sol.v)) #note that in Julia any function can be performed elementwise by putting a . before the paranthesis
        D["phase:"*mode]=angle.(sol.v)
    else
        convmode="[$(string(round((sol.params[:ω]/2/pi),digits=2)))]Hz"
        println("Excluded:$mode because converged to $convmode")
    end
end
##
D=Dict()
D["speedofsound"]=C
vtk_write("NTNU",full_mesh,D)
##

sol,nn,flag=householder(l,1000*2*pi,maxiter=6,order=1,tol=1E-9,n_eig_val=1,output=true);
mode="[$(string(round((sol.params[:ω]/2/pi),digits=2)))]Hz"
v=bloch_expand(unit_mesh,sol)
D["abs Bloch:"*mode]=abs.(v)./maximum(abs.(v)) #note that in Julia any function can be performed elementwise by putting a . before the paranthesis
D["phase Bloch:"*mode]=angle.(v)
#ol,nn,flag=householder(L,sol.params[:ω],maxiter=6,order=5,tol=1E-9,n_eig_val=3,output=true);
Sol,nn,flag=householder(L,sol.params[:ω],maxiter=6,order=5,tol=1E-5,n_eig_val=3,output=true);

mode="[$(string(round((Sol.params[:ω]/2/pi),digits=2)))]Hz"
D["abs:"*mode]=abs.(Sol.v)./maximum(abs.(Sol.v)) #note that in Julia any function can be performed elementwise by putting a . before the paranthesis
D["phase:"*mode]=angle.(Sol.v)
vtk_write("NTNU",full_mesh,D)

##

##

D=Dict()
D["speedofsound"]=C
Sol,nn,flag=householder(L,500*2*pi,maxiter=6,order=5,tol=1E-5,n_eig_val=3,output=true);
mode="[$(string(round((Sol.params[:ω]/2/pi),digits=2)))]Hz"
D["abs:"*mode]=abs.(Sol.v)./maximum(abs.(Sol.v)) #note that in Julia any function can be performed elementwise by putting a . before the paranthesis
D["phase:"*mode]=angle.(Sol.v)
vtk_write("NTNU",full_mesh,D)

##


import Base: sin, cos, tan, cot, sec, csc, atan

import Base: sinh, cosh, tanh, coth, sech, csch

import Base: abs, exp, log

import Base: +, -, *, /, ^

struct pentacomplex
    i1
    i2
end

function mat(penta::pentacomplex)
    i1, i2 = penta.i1, penta.i2
    return [i1 -i2; i2 i1]
end

function mat2pentacomplex(matrix)
    out1, out2 = matrix[1, 1], matrix[2, 1]
    return pentacomplex(out1, out2)
end

## Basic operators: +, -, *, /, ^

function +(penta1::pentacomplex, penta2::pentacomplex)
    result=mat(penta1)+mat(penta2)
    return mat2pentacomplex(result)
end

+(n, penta::pentacomplex) = mat2pentacomplex(n .+mat(penta))

+(penta::pentacomplex, n) = +(n, penta)

-(penta1::pentacomplex, penta2::pentacomplex) = mat2pentacomplex(mat(penta1)-mat(penta2))
-(n, penta::pentacomplex) = mat2pentacomplex(n .-mat(penta))
-(penta::pentacomplex, n) = -(n, penta)
-(penta::pentacomplex) = mat2pentacomplex(-mat(penta))

*(penta1::pentacomplex, penta2::pentacomplex) = mat2pentacomplex(mat(penta1)*mat(penta2))
*(n, penta::pentacomplex) = mat2pentacomplex(n*mat(penta))
*(penta::pentacomplex, n) = *(n, penta)

/(penta1::pentacomplex, penta2::pentacomplex) = mat2pentacomplex(mat(penta1)/mat(penta2))
/(n, penta::pentacomplex) = mat2pentacomplex(n*inv(mat(penta)))
/(penta::pentacomplex, n) = mat2pentacomplex(mat(penta)/n)

^(penta::pentacomplex, power) = mat2pentacomplex(mat(penta)^power)

## Basic functions: abs, exp, log

abs(penta::pentacomplex) = sqrt(penta.i1^2+penta.i2^2)

function exp(penta::pentacomplex)
    i1, i2 = penta.i1, penta.i2
    out1 = exp(i1)*cos(i2)
    out2 = exp(i1)*sin(i2)
    return pentacomplex(out1, out2)
end

## Basic math functions: sin, cos, tan, cot, sec, csc

function sin(penta::pentacomplex)
    i1, i2 = penta.i1, penta.i2
    out1 = cosh(i2)*sin(i1)
    out2 = sinh(i2)*cos(i1)
    return pentacomplex(out1, out2)
end

function cos(penta::pentacomplex)
    i1, i2 = penta.i1, penta.i2
    out1 = cosh(i2)*cos(i1)
    out2 = -sinh(i2)*sin(i1)
    return pentacomplex(out1, out2)
end

tan(penta::pentacomplex) = sin(penta)/cos(penta)
cot(penta::pentacomplex) = cos(penta)/sin(penta)
sec(penta::pentacomplex) = 1/cos(penta)
csc(penta::pentacomplex) = 1/sin(penta)

## Basic hyperbolic functions: sinh, cosh, tanh, coth, sech, csch

function sinh(penta::pentacomplex)
    i1, i2 = penta.i1, penta.i2
    out1 = cosh(i1)*cos(i2)
    out2 = sinh(i1)*sin(i2)
    return pentacomplex(out1, out2)
end

function cosh(penta::pentacomplex)
    i1, i2 = penta.i1, penta.i2
    out1 = sinh(i1)*cos(i2)
    out2 = cosh(i1)*sin(i2)
    return pentacomplex(out1, out2)
end

tanh(penta::pentacomplex) = sinh(penta)/cosh(penta)
coth(penta::pentacomplex) = cosh(penta)/sinh(penta)
sech(penta::pentacomplex) = 1/cosh(penta)
csch(penta::pentacomplex) = 1/sinh(penta)

function image5(penta::pentacomplex)
    i1, i2 = penta.i1, penta.i2
    return image4(i2)
end

function pentaderivative(f, point, h)
    result = image5(f(pentacomplex(quadcomplex(tricomplex(bicomplex(point+im*h, h), h), h), h)))/h^5
    return result
end


module GrassmannAverages

function grassmann_median(X, K)
    K >= 0 || throw(ArgumentError("Number of components must be positive"));

    N = size(X, 1);
    D = size(X, 2);
    eps = 1e-4;  "Too small of error is problematic"

    basis = Array{Real}(D, K);
    for k=1:K
        "q represents the k-th principal component"
        q = randn(D);
        q /= norm(q);

        while true
            dotsign = sign(X * q);

            "Compute the weighted average of the elements. Robustness is achieved by element-wise median."
            newQ = Array{Real}(D);
            for j=1:D
                newQ[j] = median(X[: , j] .* dotsign);
            end
            newQ /= norm(newQ);

            "If the L_1 difference between q and newQ is small, then terminate."
            max_diff = 0;
            for j=1:D
                max_diff = max(max_diff, abs(q[j] - newQ[j]));
            end
            if max_diff < eps
                break
            end

            q = newQ;
        end

        "Store the basis vector and subtract from X"
        basis[: , k] = q;
        X -= (X * q) * q';
    end

    return basis;
end

end # module


# ellipse_im.jl

using MIRT: ellipse_im, ellipse_im_fast, ellipse_im_params
using MIRT: image_geom, disk_phantom_params, shepp_logan_parameters
using MIRTjim: jim

using Plots: plot
#using Printf: @sprintf
using Test: @test, @test_throws, @inferred

"""
ellipse_im_show()
"""
function ellipse_im_show( ; over::Int = 2^2)
	ig = image_geom(nx=2^8, ny=2^8+2, fov=250)
#	ig.offset_y = 75.6 / ig.dy

	x0 = ellipse_im(ig, oversample=over)
	p0 = jim(ig.x, ig.y, x0, title="Shepp Logan", clim=(0.9,1.1))

	x1 = ellipse_im(ig, :shepplogan_emis, oversample=over)
	p1 = jim(ig.x, ig.y, x1, title="Shepp Logan Emission")

	x2 = ellipse_im(ig, :southpark)
	p2 = jim(ig.x, ig.y, x2, title="South Park")

	x3 = ellipse_im(ig, :shepplogan_brainweb)
	p3 = jim(ig.x, ig.y, x3, title="Shepp Logan Brainweb")

	plot(p0,p1,p2,p3)
end

#=

# compare to aspire
function ellipse_im_aspire()
	nx = 2^6
	ig = image_geom(nx=nx, ny=nx+2, fov=2^7)
	ell = [10, 20, 30, 40, 50, 1]

	file = tempname() * "t.fld"
	file = "tt.fld"
	over = 2^2
#	com = @sprintf "echo y | op -chat 99 ellipse %s %d %d  %g %g %g %g %g %g %d" file ig.nx ig.ny ell ./ [ig.dx, ig.dx, ig.dx, ig.dx, 1, 1] log2(over)+1
	tmp = ""
	for i=1:4
		tmp = tmp * " $(ell[i] / ig.dx)"
	end
	for i=5:6
		tmp = tmp * " $(ell[i])"
	end
	op = "/Users/fessler/bin/mi64-gcc/op"
	com = "$op -chat 99 ellipse $file $(ig.nx) $(ig.ny)$tmp $(log2(over)+1)"
	com = "/bin/ls /Users/fessler/bin/mi64-gcc" # todo: why fails?
	@show com
#	cm = @cmd "$com"
	cm = `"$com"`
	@show cm

	run(cm)
return true
	asp = fld_read(file)
	p1 = jim(asp, title="aspire")

	area_asp = sum(asp) * abs(ig.dx * ig.dy)

	jul = ellipse_im(ig, ell, oversample=over) # 'type', types{ii})
	area_jul = sum(jul) * abs(ig.dx * ig.dy)

	p2 = jim(jul, title="jul")
	p3 = jim((jul-asp)*over^2, "difference: (jul-asp)*over^2")

	@show maximum(abs.(jul - mat)) / ell[6] * over^2
	!isapprox(jul, mat) && throw("approximation error")
#	max_percent_diff(mat, asp)

	area_real = pi * ell[3] * ell[4] * ell[6]
	@show area_real, area_asp, area_jul

	true
end
=#

#=
% ellipse_im_profile()
function ellipse_im_profile
ig = image_geom('nx', 2^9, 'ny', 2^9+2', 'fov', 250)
profile on
x0 = ellipse_im(ig, [], 'oversample', 3, 'type', 'fast')
profile off
profile report
=#

#@test ellipse_im_aspire() # todo

fov = 100
@inferred shepp_logan_parameters(fov, fov)
shepp_logan_parameters(fov, fov, case=:kak)
shepp_logan_parameters(fov, fov, case=:emis)
shepp_logan_parameters(fov, fov, case=:brainweb)

# test various ways of calling
ellipse_im(20)
#@inferred ellipse_im(20,22) # todo-i: why fails?
ellipse_im(20,22)
ellipse_im(30, :shepplogan_emis, oversample=2)

ig = image_geom(nx=80, dx=1)
ellipse_im_params(ig, :kak)

params = shepp_logan_parameters(ig.fovs..., case=:brainweb)
@inferred ellipse_im_fast(ig.nx, ig.ny, params, ig.dx, ig.dy,
		ig.offset_x, ig.offset_y, 0, 2, false)
ellipse_im(ig, :disks)
ellipse_im(ig, params, oversample=2)
ellipse_im(ig, params, how=:fast, replace=true, rot=30)

@test_throws String ellipse_im(ig, :bad)
@test_throws String ellipse_im(ig, params, how=:bad)
@test_throws String ellipse_im(ig, params')
@test_throws String shepp_logan_parameters(ig.fovs..., case=:bad)

ellipse_im(ig, params)

#=
x1 = ellipse_im(100, :shepplogan_emis)
x2 = ellipse_im(100, :shepplogan_emis, oversample=2)
plot(jim(x1), jim(x2), jim(x2-x1))
=#

ellipse_im_show()


using DataStructures
using LightGraphs

# For plotting
using GraphRecipes
import Plots
Plots.pyplot()

function find_element_num(n::T) where {T}
    s = Stack{T}()
    push!(s, n)

    num = 1
    while(length(s) > 0)
        num += 1
        temp = pop!(s)
        for neighbor in temp.neighbors
            push!(s, neighbor)
        end
    end
    return num
end

# We'll talk about this later if we have time!
function GraphRecipes.graphplot(n::T) where {T}

    # Constructing a LightGraphs graph
    element_num = find_element_num(n)

    g = SimpleGraph(element_num)

    s = Stack{T}()
    push!(s, n)

    while(length(s) > 0)
        temp = pop!(s)

        for neighbor in temp.neighbors
            add_edge!(g,neighbor.ID, temp.ID)
            push!(s, neighbor)
        end
    end
    graphplot(g;curves=false, markersize=0.2, names=1:element_num)

end

function create_tree(num_row::Int64, num_neighbor::Int64; mode="Recursive")

    if mode == "recursive"
        return create_tree(num_row, num_neighbor, [1])
    elseif mode == "stack"
        return create_tree_stack(num_row, num_neighbor, [1])
    elseif mode == "queue"
        return create_tree_queue(num_row, num_neighbor)
    else
        println("error mode ", mode, " not found for tree construction!")
    end
end

function create_tree_queue(num_row::Int64, num_neighbor::Int64)

    # Creating tree
    tree = create_tree(num_row, num_neighbor, [1])

    # relabelling elements
    q = Queue{Node}()
    enqueue!(q, tree)

    ID = 1
    while(length(q) > 0)
        temp = dequeue!(q)
        temp.ID = ID
        ID += 1
        for neighbor in temp.neighbors
            enqueue!(q, neighbor)
        end
    end

    return tree

end

function create_tree_stack(num_row::Int64, num_neighbor::Int64, ID::Vector)
    final_tree = Node(ID[1])
    if (num_row == 0)
        return final_tree
    end

    #preallocating for stack ordering
    for i = 1:num_neighbor
        push!(final_tree.neighbors, Node(ID[1]))
    end

    for i = num_neighbor:-1:1
        ID[1] += 1
        neighbor = create_tree(num_row - 1, num_neighbor, ID)
        final_tree.neighbors[i] = neighbor
    end

    return final_tree
end

function create_tree(num_row::Int64, num_neighbor::Int64, ID::Vector)
    final_tree = Node(ID[1])
    if (num_row == 0)
        return final_tree
    end

    for i = 1:num_neighbor
        ID[1] += 1
        neighbor = create_tree(num_row - 1, num_neighbor, ID)
        push!(final_tree.neighbors, neighbor)
    end

    return final_tree
end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.10.3
#   kernelspec:
#     display_name: Julia 1.6.4
#     language: julia
#     name: julia-1.6
# ---

# %% [markdown]
# StaticArrays version: https://github.com/genkuroki/public/blob/main/0018/HMC%20leapfrog%20StaticArrays.ipynb

# %%
module My

using ConcreteStructs: @concrete
using ForwardDiff: derivative
using Parameters: @unpack
using Random: default_rng

@concrete struct LFProblem ϕ; H; F; dt; nsteps end

"""ϕ should be a potential function."""
function LFProblem(ϕ; dt = 0.05, nsteps = 100)
    H(x, v, param) = v^2/2 + ϕ(x, param)
    F(x, param) = -derivative(x -> ϕ(x, param), x)
    LFProblem(ϕ, H, F, dt, nsteps)
end

"""Numerically solve Hamilton's equation of motion with leapfrog method"""
function solve(lf::LFProblem, x, v, param)
    @unpack F, dt, nsteps = lf
    v = v + F(x, param)*dt/2
    x = x + v*dt
    for _ in 2:nsteps
        v = v + F(x, param)*dt
        x = x + v*dt
    end
    v = v + F(x, param)*dt/2
    x, v
end

@inline function _update!(lf::LFProblem{dim}, x, T, param, rng) where dim
    @unpack H = lf
    v = randn(rng, T)
    xnew, vnew = solve(lf, x, v, param)
    dH = H(xnew, vnew, param) - H(x, v, param)
    rand(rng) ≤ exp(-dH) ? xnew : x
end

"""Hamiltonian Monte Carlo"""
function HMC(lf::LFProblem, param = nothing;
        niters = 10^5, nskips = 1, thin = 1, nwarmups = 0, rng = default_rng(),
        x0 = randn(rng))
    T = typeof(x0)
    x = x0
    for _ in 1:nwarmups
        x = _update!(lf, x, T, param, rng)
    end
    sample = Vector{T}(undef, niters)
    for i in 1:niters
        #or _ in 1:thin
            x = _update!(lf, x, T, param, rng)
        #end
        @inbounds sample[i] = x
    end
    sample
end

end

# %%
using Plots
using QuadGK
using BenchmarkTools

# %%
ϕ(x, param=nothing) = x^2/2
lf = My.LFProblem(ϕ)
@time X = My.HMC(lf)
histogram(X; norm=true, alpha=0.3, label="HMC LF sample")
plot!(x -> exp(-x^2/2)/√(2π), -4, 4, label="std Gaussian", lw=2)

# %%
@btime My.HMC($lf);

# %%
ϕ2(x, param=nothing) = x^4 - 4x^2
Z = quadgk(x -> exp(-ϕ2(x)), -Inf, Inf)[1]
lf2 = My.LFProblem(ϕ2)
@time X2 = My.HMC(lf2)
histogram(X2; norm=true, alpha=0.3, label="HMC LF sample", bin=100)
plot!(x -> exp(-ϕ2(x))/Z, -2.5, 2.5; label="exp(-ϕ2(x))/Z", lw=2)
plot!(; legend=:outertop)

# %%
@btime My.HMC($lf2);

# %% [markdown]
# https://github.com/genkuroki/public/blob/main/0019/HMC%20leapfrog.ipynb

# %%
ϕ3(x, param=nothing) = x^2*(x - 3)^2/2
Z = quadgk(x -> exp(-ϕ3(x)), -Inf, Inf)[1]
lf3 = My.LFProblem(ϕ3)
@time X3 = My.HMC(lf3; niters=10^6)
histogram(X3; norm=true, alpha=0.3, label="HMC LF sample", bin=100)
plot!(x -> exp(-ϕ3(x))/Z, -1.5, 4.5; label="exp(-ϕ3(x))/Z", lw=2)
plot!(; legend=:outertop)

# %%
@btime My.HMC($lf3; niters=10^6);

# %%
plot(X[1:10000]; label="", size=(800, 200), lw=0.5)

# %%
plot(X2[1:10000]; label="", size=(800, 200), lw=0.5)

# %%
plot(X3[1:10000]; label="", size=(800, 200), lw=0.5)

# %%
ϕ4(x, a) = a * (x^2 - 1)^2
a = [3, 4, 5, 6, 7, 8]
XX = Vector{Float64}[]
ZZ = Float64[]
PP = []
for i in eachindex(a)
    Z = quadgk(x -> exp(-ϕ4(x, a[i])), -Inf, Inf)[1]
    push!(ZZ, Z)
    lf = My.LFProblem(ϕ4)
    @time X = My.HMC(lf, a[i])
    flush(stdout)
    push!(XX, X)
    P = plot()
    histogram!(X; norm=true, alpha=0.3, label="HMC LF sample", bin=100, c=i)
    plot!(x -> exp(-ϕ4(x, a[i]))/Z, -2, 2; label="exp(-ϕ2(x))/Z", lw=2, c=i)
    plot!(; legend=false, xtick=-2:0.5:2)
    title!("ϕ(x) = a(x² - 1)²  for  a = $(a[i])", titlefontsize=9)
    push!(PP, P)
end
plot(PP...; size=(800, 450), layout=(3, 2))

# %%
QQ = []
for i in eachindex(a)
    Q = plot(XX[i][1:10000]; ylim=(-1.5, 1.5), label="", c=i, lw=0.5)
    title!("ϕ(x) = a(x² - 1)²  for  a = $(a[i])", titlefontsize=9)
    push!(QQ, Q)
end
plot(QQ...; size=(800, 900), layout=(length(a), 1))

# %%
# https://github.com/moruten/julia-code/blob/c0dfc2443d6b74256364e698b3edb37f98214ce7/Test-2021-9-3-no4.ipynb

#HMC method
using Plots
using StatsBase
using Random
using Printf
using Distributions
using BenchmarkTools
using QuadGK
using SpecialFunctions
using Zygote

# Generating Gaussian distribution (1-dim) by using HMC method

# # of configurations
n_config = 1_00_000

# size of a configuration
n_size = 1

# # of step of Leapfrog method
n_tau = 100

# time period of Hamilton equation
τ = 5.0

function S(x)
    sum(0.5*x[i]^2 for i=1:n_size)
end

function dSdx(x,k)
    return x
end

function H(x,p)
    K = sum(p[i]^2 for i=1:n_size)
    return S(x) + 0.5*K
end

#１変数更新のHMC法
function HMC_1dim(x)
        #初期化
        p_ini = rand(Normal(0,1),n_size)
        p_fin = zeros(n_size)
        x_ini = zeros(n_size)
        x_fin = zeros(n_size)

        #初期配位を保存
        for i=1:n_size x_ini[i] = x[i] end

        #初期配位のハミルトニアンを計算
        H_ini = H(x_ini,p_ini)

        #リープフロッグで時間発展
        leapfrog!(x,p_ini)

        #時間発展後の配位を保存
        for i=1:n_size x_fin[i] = x[i]; p_fin[i] = p_ini[i] end

        #時間発展後のハミルトニアンを計算
        H_fin = H(x_fin,p_fin)

        #Metropolis check
        r=rand()
        if r<exp(H_ini - H_fin)
            for i=1:n_size x[i] = x_fin[i] end
        else
            for i=1:n_size x[i] = x_ini[i] end
        end
end

#x,pをリープフロッグ法で更新する．
function leapfrog!(x,p)
    Δτ  = τ/n_tau
    p0  = 0.0
    p1  = 0.0
    x05 = 0.0
    x15 = 0.0
    for k=1:n_size
        #initial setting : τ0
        p0  = p[k]
        x05 = x[k] + p0 * 0.5*Δτ

        #(n-1)step : τ0 -> τn-1
        for i=1:n_tau-1
            p1  = p0 - dSdx(x05,k)*Δτ
            x15 = x05 + p1*Δτ

            p0  = p1
            x05 = x15
        end

        #final step : τn-1 -> τn
        p_end = p1 - dSdx(x15,k)*Δτ
        x_end = x15 + p_end*0.5*Δτ

        x[k]  = x_end
        p[k]  = p_end
    end
end

""" main """
#set of configurations
Ω = zeros(n_config)

#configuration
x = zeros(n_size)

@time for j=1:n_config
    HMC_1dim(x)
    Ω[j] = x[1]
end

Ω

gauss(x) = exp(-0.5*x^2)/sqrt(2*pi)
xr = range(-4,4,length=10001)

plot(Ω,st=:histogram,nbins=50,norm=:pdf,label="HMC's samples")
plot!(xr,gauss.(xr),label="Gaussian") |> display

# %%
function HMC_orig!()
    for j=1:n_config
        HMC_1dim(x)
        Ω[j] = x[1]
    end
end

""" main """
#set of configurations
Ω = zeros(n_config)

#configuration
x = zeros(n_size)

@btime HMC_orig!()

# %%
length(Ω), length(X)

# %%
1.3/0.057

# %%


# simple kernel_function

function kappa(x1,x2,y1,y2,rho)
    d = sqrt((x1-y1)^2+(x2-y2)^2);
    return (1+d/rho+d^2/(3*rho^2))*exp(-d/rho)
end

function Delta_x_kappa(x1,x2,y1,y2,rho)
    d = sqrt((x1-y1)^2+(x2-y2)^2);
    return (d^2-2*rho*(rho+d))/(3*rho^4)*exp(-d/rho)
end

function Delta_x_y_kappa(x1,x2,y1,y2,rho)
    d = sqrt((x1-y1)^2+(x2-y2)^2);
    return (8*rho^2+d^2-7*rho*d)/(3*rho^6)*exp(-d/rho)
end


using ArnoldiMethod
using ArnoldiMethodTransformations
using LinearAlgebra
using SparseArrays
using Statistics
using Test

N = 20
# construct fixed eval matrix in random basis

@testset "generalized eigenvalue problem: " begin
    for i ∈ 1:500

        A = rand(N,N)
        B = N*I + rand(ComplexF64,N,N)
        # find eigenpairs closest to 5.001 (cannot be 5 as algorithm is unstable if σ is exactly an eval)
        decomp, hist = partialschur(A,B,5.001)
        # get evecs
        λ, v = partialeigen(decomp,5.001)
        @test norm(A*v-B*v*diagm(0=>λ))/prod(size(v)) ≤ sqrt(eps(float(N)))N

        As = sparse(A)
        Bs = sparse(B)
        # find eigenpairs closest to 5.001 (cannot be 5 as algorithm is unstable if σ is exactly an eval)
        decomp, hist = partialschur(As,Bs,5.001)
        # get evecs
        λ, v = partialeigen(decomp,5.001)
        @test norm(A*v-Bs*v*diagm(0=>λ))/prod(size(v)) ≤ sqrt(eps(float(N)))N
    end
end


using Stan, Turing

@model gdemo(x) = begin
  s ~ InverseGamma(2,3)
  m ~ Normal(0, sqrt(s))
  x[1] ~ Normal(m, sqrt(s))
  x[2] ~ Normal(m, sqrt(s))
  return s, m
end

model_f = gdemo([1.5, 2.0])

chn = sample(model_f, 2000, 1000, false, 1, Stan.Adapt(), Stan.Hmc())


# Factorize two different matrices (one real, one complex) and solve for multiple rhs and free memory
using Distributed
if nworkers()<3
    addprocs(3)
end
@everywhere using MUMPSjInv
using Test

include("getDivGrad.jl");

A  = getDivGrad(24,20,22);
A2 = getDivGrad(30,31,33);

n  = size(A,1);
n2 = size(A2,1);
A  = A + im*spdiagm(0=>rand(n));

nrhs = 10;
rhs = randn(n,nrhs) + im*randn(n,nrhs);
rhs2 = randn(n2,nrhs);

@test nworkers()>=3

println("Factorize complex matrix on worker 1")
@time begin
F1 = remotecall_fetch(factorMUMPS,workers()[1], A,1)
end
dump(F1)

println("Factorize real matrix on worker 2")
@time begin
	F2 = remotecall_fetch(factorMUMPS,workers()[2],A2,1)
end
dump(F2)

println("Solve complex system on worker 1")
@time begin
	x = remotecall_fetch(applyMUMPS,workers()[3],F1,rhs)
end
err = zeros(nrhs)
for i=1:nrhs
        err[i] =  norm(A*x[:,i]-rhs[:,i]) / norm(rhs[:,i])
end
@test maximum(err) < 1e-14

println("Solve real system on worker 2")
@time begin
	x2 = remotecall_fetch(applyMUMPS,workers()[2],F2,rhs2)
end
err = zeros(nrhs)
for i=1:nrhs
        err[i] =  norm(A2*x2[:,i]-rhs2[:,i]) / norm(rhs2[:,i]);
end
@test maximum(err) < 1e-14

println("Free memory on worker 1")
destroyMUMPS(F1)
# remotecall_fetch(destroyMUMPS,workers()[1],F1);

println("Free memory on worker 2")
destroyMUMPS(F2)
# remotecall_fetch(destroyMUMPS,workers()[2],F2);

println("DONE!")


### Degree sequence functions

function find_startindex(degseq_rotated::Vector{<:Integer})
    min_value = Inf64
    min_index = 0
    current_sum = 0
    size = length(degseq_rotated)
    for i = 1:length(degseq_rotated)
        current_sum += degseq_rotated[i]-1
        if current_sum < min_value
            min_value = current_sum
            min_index = i
        end
    end
    return (min_index+1) % size
end

function fix_degseq(degseq_rotated::Vector{<:Integer})
    start_index = find_startindex(degseq_rotated)
    return circshift(degseq_rotated, -(start_index-1))
end

function degrees_count(size::Int, dist::DiscreteUnivariateDistribution, maxdegree::Int)
    buf = 16 # This is enough for almost all real applications
    degcount = Vector{Int}(undef, buf)

    while true
        cur = 0
        total_node = 0
        total_deg = 0
        p_cum = 0.0
        p_cur = 0.0
        p_para::Float64 = 0.0
        current_degcount = 0
        deg = 0

        while total_deg < size - 1 && total_node < size && deg <= maxdegree
            p_cur = pdf(dist, deg)
            p_para = p_cur/(1-p_cum)
            # This happens at the boundary
            if p_para > 1.0 || p_para < 0.0
                p_para = 1.0
            end
            current_degcount = rand(Binomial(size-total_node, p_para))
            total_node += current_degcount
            total_deg += deg*current_degcount
            if deg < length(degcount)
                degcount[deg+1] = current_degcount
            else
                push!(degcount, current_degcount)
            end
            p_cum += p_cur
            deg += 1
        end

        if total_deg == size - 1 && total_node == size
            return degcount[1:deg]
        end
    end
end

function degrees(size::Int, degcount::AbstractArray{T, 1}) where T <: Integer
    degseq = Array{Int, 1}(undef, size)
    pos = 1
    for deg in 0:length(degcount)-1
        n = degcount[deg+1]
        degseq[pos:pos+n-1] .= deg
        pos = pos + n
    end
    degseq
end

### Cayley Tree

"A conditional Galton-Watson tree with Poission(1) offspring distrbiution."
struct CayleyTree <: CondGWTree
    spec::TreeSpec
end
"""
    CayleyTree(size::Int)

Construct a `CayleyTree` of size `size`.
"""
CayleyTree(size::Int) = CayleyTree(TreeSpec(size, "Cayley Tree"))

function degrees_rotate(tree::CayleyTree, nbatch=2048)
    # See Svante Example 12.1 and Luc's simulation paper
    treesize = size(tree)
    degseq = zeros(Int32, treesize)
    batch = Vector{Int}(undef, nbatch)
    i = 0
    while i + nbatch < treesize
        rand!(batch, 1:treesize)
        for j = 1:nbatch
            @inbounds degseq[batch[j]] += 1
        end
        i += nbatch
    end
    for j = i+1:treesize-1
        box = rand(1:treesize)
        degseq[box] += 1
    end
    degseq
end

function degrees(tree::CondGWTree)
    degseq = degrees_rotate(tree)
    fix_degseq(degseq)
end

### Binary Tree

"A conditional Galton-Watson tree with `2 Bernoulli(1/2)` offspring distrbiution."
struct BinaryTree <: CondGWTree
    spec::TreeSpec
end

"""
    BinaryTree(size::Int)

Construct a `BinaryTree` of size `size`.
"""
function BinaryTree(size::Int)
    if iseven(size)
        size += 1
    end
    BinaryTree(TreeSpec(size, "Binary Tree"))
end

function degrees_rotate(tree::BinaryTree)
    treesize = size(tree)
    degseq = Array{Int32}(undef, treesize)
    lefanum = convert(Int, (treesize + 1)/2)

    # initialize the array
    degseq[1:lefanum] .= 0
    degseq[lefanum+1:end] .= 2

    shuffle!(degseq)
end

### General conditioned Galton-Watson tree

"A conditional Galton-Watson tree."
struct GeneralCondGWTree <: CondGWTree
    "Specification"
    spec::TreeSpec
    "Offspring distribution"
    dist::DiscreteUnivariateDistribution
end

distribution(tree::GeneralCondGWTree) = tree.dist

"""
    GeneralCondGWTree(size::Int, maxdegree::Int, dist::DiscreteUnivariateDistribution)

Construct a `GeneralCondGWTree` of size `size` and offspring distribution `dist`.
"""
function GeneralCondGWTree(size::Int, maxdegree::Int, dist::DiscreteUnivariateDistribution)
    GeneralCondGWTree(TreeSpec(size, "Conditional Galton-Watson Tree with offspring distribution $dist", maxdegree), dist)
end

"""
    GeneralCondGWTree(size::Int, maxdegree::Int, dist::DiscreteUnivariateDistribution)

Construct a `GeneralCondGWTree` of size `size`, maximum degree `maxdegree`, and offspring
distribution `dist`.
"""
function GeneralCondGWTree(size::Int, dist::DiscreteUnivariateDistribution)
    GeneralCondGWTree(size, size, dist)
end

"""
    CatalanTree(size::Int)

Construct a `GeneralCondGWTree` with `Geometric(1/2)` offspring distribution of size `size`.
"""
CatalanTree(size::Int) = GeneralCondGWTree(TreeSpec(size, "Catalan Tree"), Geometric(0.5))

"""
    DAryTree(size::Int, d::Int)

Construct a `GeneralCondGWTree` with `Binomial(d, 1/d)` offspring distribution of size `size`.
"""
DAryTree(size::Int, d::Int) = GeneralCondGWTree(TreeSpec(size, "$d-ary Tree", d), Binomial(d, 1/d))

"""
    MotzkinTree(size::Int)

Construct a `GeneralCondGWTree` with `DiscreteUniform(0, 2)` offspring distribution of size `size`.
"""
MotzkinTree(size::Int) = GeneralCondGWTree(TreeSpec(size, "Motzkin Tree", 2), DiscreteUniform(0, 2))

function degrees_rotate(tree::GeneralCondGWTree)
    degcount = degrees_count(size(tree), distribution(tree), maxdegree(tree))
    degseq = degrees(size(tree), degcount)
    shuffle!(degseq)
end

"""
    treegraph(tree)::FixedTreeGraph

Either convert a fixed `tree` to a `FixedTreeGraph`, or generate a random `FixedTreeGraph` according
to `tree`.
"""
function treegraph end

function treegraph(tree::FiniteTree)::FixedTreeGraph
    degseq = degrees(tree)
    walker = GraphWalker(size(tree))
    walk(degseq, walker, true)
    walker.tree_digraph
end

function treegraph(degseq::Vector{Int})::FixedTreeGraph
    walker = GraphWalker(length(degseq))
    walk(degseq, walker, true)
    walker.tree_digraph
end

### RandomRecursiveTree

"""
A random recursive tree.
"""
struct RandomRecursiveTree <: FiniteRandomTree
    spec::TreeSpec
end

"""
    RandomRecursiveTree(size::Int)

Construct a `RandomRecursiveTree` of size `size`.
"""
RandomRecursiveTree(size::Int) = RandomRecursiveTree(TreeSpec(size, "Random Recursive Tree"))

"""
    degrees(tree)

Return an iterator that traverses the degree sequence of `tree` in DFS order.
"""
function degrees end

function degrees(tree::RandomRecursiveTree)
    return treegraph(tree)
end

function treegraph(tree::RandomRecursiveTree)::FixedTreeGraph
    treesize = size(tree)
    graph = FixedTreeGraph(treesize)

    for current_node in 2:treesize
        parent_node = rand(1:current_node-1)
        addedge!(graph, parent_node, current_node)
    end

    graph
end


export NormalWeightedMeanPrecision

import StatsFuns: log2π, invsqrt2π

struct NormalWeightedMeanPrecision{T <: Real} <: ContinuousUnivariateDistribution
    xi :: T # Weighted mean: xi = w * μ
    w  :: T
end

NormalWeightedMeanPrecision(xi::Real, w::Real)       = NormalWeightedMeanPrecision(promote(xi, w)...)
NormalWeightedMeanPrecision(xi::Integer, w::Integer) = NormalWeightedMeanPrecision(float(xi), float(w))
NormalWeightedMeanPrecision(xi::Real)                = NormalWeightedMeanPrecision(xi, one(xi))
NormalWeightedMeanPrecision()                        = NormalWeightedMeanPrecision(0.0, 1.0)

Distributions.@distr_support NormalWeightedMeanPrecision -Inf Inf

Distributions.support(dist::NormalWeightedMeanPrecision) = Distributions.RealInterval(minimum(dist), maximum(dist))

weightedmean(dist::NormalWeightedMeanPrecision) = dist.xi

function mean_cov(dist::NormalWeightedMeanPrecision)
    v = cov(dist)
    μ = v * weightedmean(dist)
    return (μ, v)
end

Distributions.mean(dist::NormalWeightedMeanPrecision)    = var(dist) * weightedmean(dist)
Distributions.median(dist::NormalWeightedMeanPrecision)  = mean(dist)
Distributions.mode(dist::NormalWeightedMeanPrecision)    = mean(dist)
Distributions.var(dist::NormalWeightedMeanPrecision)     = inv(dist.w)
Distributions.std(dist::NormalWeightedMeanPrecision)     = sqrt(var(dist))
Distributions.cov(dist::NormalWeightedMeanPrecision)     = var(dist)
Distributions.invcov(dist::NormalWeightedMeanPrecision)  = dist.w
Distributions.entropy(dist::NormalWeightedMeanPrecision) = (1 + log2π - log(precision(dist))) / 2

Distributions.pdf(dist::NormalWeightedMeanPrecision, x::Real)    = (invsqrt2π * exp(-abs2(x - mean(dist)) * precision(dist) / 2)) * sqrt(precision(dist))
Distributions.logpdf(dist::NormalWeightedMeanPrecision, x::Real) = -(log2π - log(precision(dist)) + abs2(x - mean(dist)) * precision(dist)) / 2

Base.precision(dist::NormalWeightedMeanPrecision)     = invcov(dist)
Base.eltype(::NormalWeightedMeanPrecision{T}) where T = T

Base.convert(::Type{ NormalWeightedMeanPrecision }, xi::Real, w::Real) = NormalWeightedMeanPrecision(xi, w)
Base.convert(::Type{ NormalWeightedMeanPrecision{T} }, xi::Real, w::Real) where { T <: Real } = NormalWeightedMeanPrecision(convert(T, xi), convert(T, w))

vague(::Type{ <: NormalWeightedMeanPrecision }) = NormalWeightedMeanPrecision(0.0, tiny)

prod_analytical_rule(::Type{ <: NormalWeightedMeanPrecision }, ::Type{ <: NormalWeightedMeanPrecision }) = ProdAnalyticalRuleAvailable()

function Base.prod(::ProdPreserveType, left::NormalWeightedMeanPrecision, right::NormalWeightedMeanPrecision)
    xi = weightedmean(left) + weightedmean(right)
    w  = precision(left) + precision(right)
    return NormalWeightedMeanPrecision(xi, w)
end

function Base.prod(::ProdAnalytical, left::NormalWeightedMeanPrecision, right::NormalWeightedMeanPrecision)
    xi = weightedmean(left) + weightedmean(right)
    w  = precision(left) + precision(right)
    return NormalWeightedMeanPrecision(xi, w)
end


module AlignSpectra

using HDF5, PyPlot, DynamicTimeWarp
using Dierckx, Optim

include("MonotoneSpline.jl")

export MonotoneSpline, MonotoneSplineLogLog, compose_spline,
    compose_splinelog, matchspectra

"""Find <npeaks> distinct peaks in a histogram with contents <h>.
The algorithm is to label the fullest bin as the 1st peak, and penalize
all other bins by P=d/(d+<distscale>), where d is the bin distance to the
nearest peak. The bin with the 2nd-highest value of P*h is the 2nd peak,
the penalties P are recomputed, and the selection step is repeated until
<npeaks> peaks are found and return as a sorted list of bin numbers.

Returns: peak_bin_numbers::Vector{eltype(h)}, a vector of length <npeaks>.
"""
function distinctpeaks(h::Vector, npeaks::Integer, distscale::Number)
    if npeaks > length(h)
        error("Histogram has only $(length(h)) bins: cannot find $(npeaks).")
    end
    bin = collect(1:length(h))
    peaks = Int[indmax(h)]
    dist = abs(peaks[1] .- bin)
    penalty = dist ./ (distscale+dist)
    for pknum in 2:npeaks
        for p in peaks[2:end]
            dist = abs(p .- bin)
            penalty = min(penalty, dist./(distscale+dist))
        end
        push!(peaks, indmax(h .* penalty))
    end
    sort(peaks)
end

"""Replace a DTW path (p1,p2) with (x,y) where the x values are unique and
the y values are an average over all those p2 where p1==x."""
function uniquepath(p1::Vector, p2::Vector)
    @assert length(p1)==length(p2)

    minx = p1[1]
    x = unique(p1)
    sort!(x)
    y = zeros(Float64, length(x))
    n = zeros(Int, length(x))

    bin = 1
    for i=1:length(p2)
        while x[bin] < p1[i]
            bin += 1
        end
        n[bin] += 1
        y[bin] += p2[i]
    end
    float(x), y./n
end

"""Find the best cubic spline with knots at {k}, such that the Euclidean distance to all pairs
of points {x_i, y_i} is minimized and such that the spline has "natural boundary conditions"
(i.e., it has zero second derivative at the extreme knots, k[1] and k[end]).
"""
function findbestspline(x::Vector, y::Vector, k::Vector)
    k = sort(k)
    const m = length(k)
    const n = length(x)
    @assert length(x) == length(y)

    Mdiag = ones(Float64, m)
    for i=2:m-1
        Mdiag[i] = 2.*(k[i+1]-k[i-1])
    end
    Mdu = zeros(Float64, m-1)
    Mdl = zeros(Float64, m-1)
    for i=1:m-2
        Mdu[i+1] = k[i+2]-k[i+1]
        Mdl[i] = k[i+1]-k[i]
    end
    M = Tridiagonal(Mdl, Mdiag, Mdu)

    Ndiag = zeros(Float64, m)
    for i=2:m-1
        Ndiag[i] = 6.*(1./(k[i]-k[i-1]) - 1./(k[i+1]-k[i]))
    end
    Ndl = zeros(Float64, m-1)
    Ndu = zeros(Float64, m-1)
    for i=2:m-1
        Ndl[i-1] = -6./(k[i]-k[i-1])
        Ndu[i] = +6./(k[i]-k[i-1])
    end
    N = Tridiagonal(Ndl, Ndiag, Ndu)

    # s'' = R s where s are the f(x) values at x={k} and s'' are 2nd derivatives there
    R = inv(M) * full(N)

    V = zeros(Float64, n, m)
    for i=1:n
        xi = x[i]

        # Below lowest knot
        if xi<=k[1]
            thisk = 1
            dr = k[2]-k[1]
            B=(xi-k[thisk])/dr
            A=1-B
            V[i,thisk] = A
            V[i,thisk+1] = B

        # Above highest knot
        elseif xi>=k[m]
            thisk = m-1
            dr = k[end]-k[end-1]
            B=(xi-k[thisk])/dr
            A=1-B
            V[i,thisk] = A
            V[i,thisk+1] = B

        # In the proper spline range
        else
            thisk = findn(xi.>=k)[1][1]
            if thisk>=m || thisk<1
                @show thisk, k, xi
            end
            dr = k[thisk+1]-k[thisk]
            B=(xi-k[thisk])/dr
            A=1-B
            C = (A^3-A)*(dr^2)/6.
            D = (B^3-B)*(dr^2)/6.
            V[i,thisk] = A
            V[i,thisk+1] = B
            V[i,:] += C*R[thisk,:]
            V[i,:] += D*R[thisk+1,:]
        end
    end
    knotvals = pinv(V)*y
end

"""Find the best cubic spline according to the Dierckx.Spline1D function with
weight vector w."""
function findbestspline_numeric(x::Vector, y::Vector, w::Vector)
    Dierckx.Spline1D(x, y, w=w, bc="nearest", s=1.0)
end

"""Find the best cubic spline with knots at {k}, according to Dierckx.Spline1D with
fixed knots.
"""
function findbestspline_knots(x::Vector, y::Vector, knots::Vector)
    @show knots
    Dierckx.Spline1D(x, y, knots, bc="nearest", k=3)
end

function basicdtw(data1::Vector{Float32}, data2::Vector{Float32}) #c1::Vector{Int}, c2::Vector{Int})
    clf()
    subplot(321)
    c1,b,_=plt.hist(data1, 8000, [0,8000], histtype="step", color="r", normed=true)
    db = 0.5*(b[2]-b[1])
    binctrs = b[1:end-1]+db
    pk1 = distinctpeaks(c1, 8, 100.)
    for p in pk1
        plt.plot(binctrs[p], c1[p], "ro")
    end

    c2,_,_=plt.hist(data2, 8000, [0,8000], histtype="step", color="k", normed=true)

    d1d2_distance = DynamicTimeWarp.Distance.poissonpenalty(length(data1), length(data2))
    dist,p1,p2 = DynamicTimeWarp.fastdtw(c1, c2, 10, d1d2_distance)
    subplot(322)
    u1,u2 = uniquepath(binctrs[p1],binctrs[p2])
    plot(u1,u2,"g")
    for p in pk1
        plt.plot(binctrs[p], u2[p], "ro")
    end

    f = splinelogspace(binctrs[pk1], float(u2[pk1]), bc="extrapolate")
    #   f = MonotoneSpline(binctrs[pk1], float(u2[pk1]), bc="extrapolate")
    plot(u1, f(u1), "k")

    data1e = f(float(data1))
    subplot(323)
    c1,b,_=plt.hist(data1e, 8000, [0,8000], histtype="step", color="r", normed=true)
    c2,_,_=plt.hist(data2, 8000, [0,8000], histtype="step", color="k", normed=true)
    dist,p1,p2 = DynamicTimeWarp.fastdtw(c1, c2, 20, d1d2_distance)
    subplot(324)
    u1,u2 = uniquepath(binctrs[p1],binctrs[p2])
    plot(u1,u2,"g")
    pk1 = distinctpeaks(c1, 8, 100.)
    for p in pk1
        plt.plot(binctrs[p], u2[p], "ro")
    end

    #   @show pk1, binctrs[pk1], u2[pk1]
    g = splinelogspace(binctrs[pk1], float(u2[pk1]), bc="extrapolate")
    plot(u1, g(u1), "k")

    data1f = g(float(data1e))
    subplot(313)
    c1,b,_=plt.hist(data1f, 8000, [0,8000], histtype="step", color="r", normed=true)
    c2,_,_=plt.hist(data2, 8000, [0,8000], histtype="step", color="k", normed=true)
    return data1f
end

"""Take a single DTW step to align data1 and data2.
Finds "distinct peaks" in data1 histogram (of a minimum size and separation).
Finds the FastDTW path to align the spectra (using a Poisson-inspired distance
function) via fastdtw and then "uniquifies" the path using uniquepath so that
no x value (data1 bin number) is hit twice. Finally, it returns only the
peak locations in the two spaces, but doing lookups in the unique path.

Return: (p1, p2) where p1 is a list of peak locations in the data1
histogram, and p2 is the location of these same peaks in data2, as
determined by FastDTW.

Uses a fixed histogram binning of 0:2:8000 at this point."""
function one_dtw_step(hist1::Vector, hist2::Vector, binctrs::Vector)

    pk1 = distinctpeaks(hist1, 8, 100.)

    d1d2_distance = DynamicTimeWarp.Distance.poissonpenalty(sum(hist1), sum(hist2))
    dist,p1,p2 = DynamicTimeWarp.fastdtw(hist1, hist2, 20, d1d2_distance)
    @printf("The Poisson DTW path distance: %f\n", dist)
    u1,u2 = uniquepath(binctrs[p1],binctrs[p2])
    return binctrs[pk1], float(u2[pk1])
end

function matchspectra{T<:Number}(histograms::Matrix{T}; histrange::Range=0:0.5:8000)
    N = size(histograms)[2]
    if N == 1
        dummyspline = AlignSpectra.MonotoneSplineLogLog([1,10], [1,10]; bc="extrapolate")
        return [dummyspline], histograms[:,1]
    end

    if N == 2
        binedges = collect(histrange)
        db = binedges[2]-binedges[1]
        binctrs = binedges[1:end-1]+db*0.5
        Nbins = length(binctrs)

        n1, n2 = sum(histograms[:,1]), sum(histograms[:,2])
        r,s = one_dtw_step(histograms[:,1], histograms[:,2], binctrs)
        middle = (r*n1 .+ s*n2) / (n1+n2)
        f1 = AlignSpectra.MonotoneSplineLogLog(r, middle; bc="extrapolate")
        f2 = AlignSpectra.MonotoneSplineLogLog(s, middle; bc="extrapolate")

        # Now convert histograms to equivalent values, transform them, and
        # re-histogram the result. Approximate this by random sampling for
        # allowed values in the bin (for bins with <10 samples), or by evenly
        # spreading out the values to
        combined = zeros(histograms[:,1])
        for i = 1:length(combined)
            nh = histograms[i,1]
            if nh > 0
                Nval = div(nh, 8) + zeros(Int, 8)
                for j=1:nh-sum(Nval)
                    Nval[rand(1:8)] += 1
                end
                @assert sum(Nval) == nh
                values = linspace(binedges[i]+db/16., binedges[i+1]-db/16, 8)
                fvalues = f1(values)
                for (n, fv) in zip(Nval, fvalues)
                    bnum = round(Int, div(fv-binedges[1], db))+1
                    if bnum < Nbins && bnum > 0
                        combined[bnum] += n
                    end
                end
            end

            nh = histograms[i,2]
            if nh > 0
                Nval = div(nh, 8) + zeros(Int, 8)
                for j=1:nh-sum(Nval)
                    Nval[rand(1:8)] += 1
                end
                values = linspace(binedges[i]+db/16., binedges[i+1]-db/16, 8)
                fvalues = f2(values)
                for (n, fv) in zip(Nval, fvalues)
                    bnum = round(Int, div(fv-binedges[1], db))+1
                    if bnum < Nbins && bnum > 0
                        combined[bnum] += n
                    end
                end
            end
        end
        return [f1,f2], combined
    end

    if N == 3
        curvesA, combinedA = matchspectra(histograms[:,1:2]; histrange=histrange)
        input = hcat(combinedA, histograms[:,3])
        curvesB, combined = matchspectra(input; histrange=histrange)
        f1 = compose_splinelog(curvesA[1], curvesB[1])
        f2 = compose_splinelog(curvesA[2], curvesB[1])
        f3 = curvesB[2]
        return [f1,f2,f3], combined
    end

    # Break the values vector up into equal-size groups (extra in N2 if N is odd)
    N1 = div(N,2)
    N2 = N - N1

    @assert N > 3
    histograms1 = histograms[:,1:N1]
    histograms2 = histograms[:,N1+1:end]

    curves1, combined1 = matchspectra(histograms1; histrange=histrange)
    curves2, combined2 = matchspectra(histograms2; histrange=histrange)

    @show N, typeof(combined1), eltype(combined1)
    @show N2, typeof(combined2), eltype(combined2)
    inputs = hcat(combined1, combined2)
    curves3, combined3 = matchspectra(inputs; histrange=histrange)
    finalcurves = []
    for c in curves1
        push!(finalcurves, compose_splinelog(c, curves3[1]))
    end
    for c in curves2
        push!(finalcurves, compose_splinelog(c, curves3[2]))
    end

    finalcurves, combined3
end

end # module


using PrivateMultiplicativeWeights
using CSV
import Plots

weights = CSV.read("our_synthetic_weights.csv");
d = Vector(weights[:weight]);
# Magic number is number of samples (from python notebook)
h = Histogram(d, 38462)

epsilons = reverse([10.0^-n for n in range(1, stop=6)])
mses = []
max_errors = []

for ε in epsilons
    println("Running MWEM with ε=$ε")
    mw = mwem(SeriesRangeQueries(length(d)), h, MWParameters(epsilon=ε, iterations=500));
    mse = mean_squared_error(mw)
    println("Mean squared error = ", mse)
    append!(mses, mse)
end

Plots.plot(epsilons, mses, xlabel = "Epsilon", ylabel="MSE", title="Utility vs Privacy", xaxis=:log)

Scripts/Round 3/Baysically Measure Zero/mwem.jl


#=
    We use two sets of infixable comparatives.

    The conventional symbols {==,!=,<,<=,>,>=} are defined using
       the Arb C library's implementation of eq,ne,lt,le,gt,ge.

    The predecessor/successor symbols {≃,≄,≺,≼,≻,≽} are defined from
       Hend Dawood's non-strict total ordering for interval values.
       (q.v. Hend's Master's thesis:
        Interval Mathematics Foundations, Algebraic Structures, and Applications)

=#

for (op,cop) in ((:(==), :(arb_eq)), (:(!=), :(arb_ne)),
                 (:(<=), :(arb_le)), (:(>=), :(arb_ge)),
                 (:(<), :(arb_lt)),  (:(>), :(arb_gt))  )
  @eval begin
    function ($op)(a::T, b::T) where {T <: ArbFloat}
        return Bool(ccall(@libarb($cop), Cint, (Ref{ArbFloat}, Ref{ArbFloat}), a, b) )
    end
    ($op)(a::ArbFloat{P}, b::ArbFloat{Q}) where {P,Q} = ($op)(promote(a,b)...)
    ($op)(a::T, b::R) where {T <: ArbFloat,R <: Real} = ($op)(promote(a,b)...)
    ($op)(a::R, b::T) where {T <: ArbFloat,R <: Real} = ($op)(promote(a,b)...)
  end
end

function (≃)(a::T, b::T) where {T <: ArbFloat}
    return Bool(ccall(@libarb(arb_eq), Cint, (Ref{ArbFloat}, Ref{ArbFloat}), a, b))
end
simeq(a::T, b::T) where {T <: ArbFloat} = (≃)(a,b)

function (≄)(a::T, b::T) where {T <: ArbFloat}
    return !Bool(ccall(@libarb(arb_eq), Cint, (Ref{ArbFloat}, Ref{ArbFloat}), a, b))
end
nsime(a::T, b::T) where {T <: ArbFloat} = (≄)(a,b)

function (⪰)(a::T, b::T) where {T <: ArbFloat}
    alo, ahi = bounds(a)
    blo, bhi = bounds(b)
    return (alo < blo) || ((alo == blo) & (ahi <= bhi))
end
succeq(a::T, b::T) where {T <: ArbFloat} = (⪰)(a,b)

function (≻)(a::T, b::T) where {T <: ArbFloat} # (a ≼ b) & (a ≄ b)
    alo, ahi = bounds(a)
    blo, bhi = bounds(b)
    return (alo < blo) || ((alo == blo) & (ahi < bhi))
end
succ(a::T, b::T) where {T <: ArbFloat} = (≻)(a,b)

function (⪯)(a::T, b::T) where {T <: ArbFloat}
    alo, ahi = bounds(a)
    blo, bhi = bounds(b)
    return (alo > blo) || ((alo == blo) & (ahi >= bhi))
end
preceq(a::T, b::T) where {T <: ArbFloat} = (⪯)(a,b)

function (≺)(a::T, b::T) where {T <: ArbFloat}
    alo, ahi = bounds(a)
    blo, bhi = bounds(b)
    return (alo > blo) || ((alo == blo) & (ahi > bhi))
end
prec(a::T, b::T) where {T <: ArbFloat} = (≺)(a,b)

# for sorted ordering
isequal(a::T, b::T) where {T <: ArbFloat} = !(a != b)
isless(a::T, b::T) where {T <: ArbFloat} = succ(a,b)
isequal(a::Nothing, b::T) where {T <: ArbFloat} = false
isequal(a::T, b::Nothing) where {T <: ArbFloat} = false
isless(a::Nothing, b::T) where {T <: ArbFloat} = true
isless(a::T, b::Nothing) where {T <: ArbFloat} = true

function max(x::T, y::T) where {T <: ArbFloat}
    (isnan(x) || isnan(y)) && return x
    if isinf(x)
       isposinf(x) && return x
       isneginf(x) && return y
    elseif isinf(y)
       isposinf(y) && return y
       isneginf(y) && return x
    end
    return (x + y + abs(x - y))/2
end

function min(x::T, y::T) where {T <: ArbFloat}
    (isnan(x) || isnan(y)) && return x
    if isinf(x)
       isposinf(x) && return y
       isneginf(x) && return x
    elseif isinf(y)
       isposinf(y) && return x
       isneginf(y) && return y
    end
    return (x + y - abs(x - y))/2
end
#=
min{T<:ArbFloat}(a::T, b::T) = smartvalue(a) < smartvalue(b) ? a : b
max{T<:ArbFloat}(a::T, b::T) = smartvalue(b) < smartvalue(a) ? a : b

function min2{T<:ArbFloat}(x::T, y::T)
    return
        if donotoverlap(x,y)
            x < y ? x : y
        else
            xlo, xhi = bounds(x)
            ylo, yhi = bounds(y)
            lo,hi = min(xlo, ylo), min(xhi, yhi)
            md = (hi+lo)/2
            rd = (hi-lo)/2
            midpoint_radius(md, rd)
        end
end

function max2{T<:ArbFloat}(x::T, y::T)
    return
        if donotoverlap(x,y)
            return x > y ? x : y
        else
            xlo, xhi = bounds(x)
            ylo, yhi = bounds(y)
            lo,hi = max(xlo, ylo), max(xhi, yhi)
            md = (hi+lo)/2
            rd = (hi-lo)/2
            return midpoint_radius(md, rd)
        end
end
=#
#=
function max{T<:ArbFloat}(x::T, y::T)
    return ((x>=y) | !(y<x)) ? x : y
end
=#

function minmax(x::T, y::T) where {T <: ArbFloat}
   return min(x,y), max(x,y) # ((x<=y) | !(y>x)) ? (x,y) : (y,x)
end

# experimental ≖ ≗
(eq)(x::T, y::T) where {T <: ArbFloat} = !(x != y)
(eq)(x::ArbFloat{P}, y::ArbFloat{Q}) where {P,Q} = (eq)(promote(x,y)...)
(eq)(x::T1, y::T2) where {T1 <: ArbFloat,T2 <: Real} = (eq)(promote(x,y)...)
(eq)(x::T2, y::T1) where {T1 <: ArbFloat,T2 <: Real} = (eq)(promote(x,y)...)
(≗)(x::T, y::T) where {T <: ArbFloat} = !(x != y)
(≗)(x::ArbFloat{P}, y::ArbFloat{Q}) where {P,Q} = (≗)(promote(x,y)...)
(≗)(x::T1, y::T2) where {T1 <: ArbFloat,T2 <: Real} = (≗)(promote(x,y)...)
(≗)(x::T2, y::T1) where {T1 <: ArbFloat,T2 <: Real} = (≗)(promote(x,y)...)
(neq)(x::T, y::T) where {T <: ArbFloat} = donotoverlap(x, y)
(neq)(x::ArbFloat{P}, y::ArbFloat{Q}) where {P,Q} = (donotoverlap)(promote(x,y)...)
(neq)(x::T1, y::T2) where {T1 <: ArbFloat,T2 <: Real} = (donotoverlap)(promote(x,y)...)
(neq)(x::T2, y::T1) where {T1 <: ArbFloat,T2 <: Real} = (donotoverlap)(promote(x,y)...)


@doc raw"""
    BrownianMotionTorus(n::Int)

Returns a hidden state model corresponding to a Brownian motion on an `n`-dimensional torus, with initial condition drawn uniformly at random.
"""
struct BrownianMotionTorus <: HiddenStateModel{Vector{Float64}, ContinuousTime}
    n::Int
end

#####################
### BASIC METHODS ###
#####################

initial_condition(::BrownianMotionTorus) = Uniform(0,2π)

state_dim(model::BrownianMotionTorus) = model.n

"""
    noise_dim(model)

Returns the dimension of the Brownian motion ``W_t`` in the diffusion model ``dX_t = f(X_t)dt + g(X_t)dW_t``.
"""
noise_dim(model::BrownianMotionTorus) = model.n

function initialize(model::BrownianMotionTorus)
    x = rand(initial_condition(model), state_dim(model))
    return x
end

function Base.show(io::IO, ::MIME"text/plain", model::BrownianMotionTorus)
    print(io, "Brownian motion on the ", state_dim(model), "-torus for the hidden state
    type of hidden state:                   ", state_dim(model),"-dimensional vector
    number of independent Brownian motions: ", noise_dim(model),"
    initial condition:                      uniform")
end

function Base.show(io::IO, model::BrownianMotionTorus)
    print(io, "Brownian motion on the ", state_dim(model), "-torus")
end

######################
### TIME EVOLUTION ###
######################

function (model::BrownianMotionTorus)(x::AbstractVector{T}, dt) where T
    return mod2pi.(x .+ sqrt(dt) .* randn(T, noise_dim(model)))
end

function (model::BrownianMotionTorus)(x::AbstractMatrix, dt)
    N   = size(x, 2)
    out = zeros(eltype(x), state_dim(model), N)
    sqr = sqrt(dt)
    @inbounds for b in 1:N
        for a in 1:state_dim(model)
            out[a,b] += mod2pi(x[a,b] + sqr * randn(eltype(x)))
        end
    end
    return out
end

function propagate!(x::AbstractVector, model::BrownianMotionTorus, dt)
    N   = size(x, 2)
    sqr = sqrt(dt)
    @inbounds for a in 1:state_dim(model)
        x[a] +=  sqr * randn(eltype(x))
        x[a]  =  mod2pi(x[a])
    end
    return x
end

function propagate!(x::AbstractMatrix, model::BrownianMotionTorus, dt)
    N   = size(x, 2)
    sqr = sqrt(dt)
    @inbounds for b in 1:N
        for a in 1:state_dim(model)
            x[a,b] += sqr * randn(eltype(x))
            x[a,b]  = mod2pi(x[a,b])
        end
    end
    return x
end

################################
### CONVENIENCE CONSTRUCTORS ###
################################

const BrownianMotionCircle() = BrownianMotionTorus(1)


using StochasticDiffEq, DiffEqProblemLibrary, Test, Random
using DiffEqProblemLibrary.SDEProblemLibrary: importsdeproblems; importsdeproblems()
import DiffEqProblemLibrary.SDEProblemLibrary: prob_sde_linear_stratonovich, prob_sde_2Dlinear_stratonovich
Random.seed!(100)
dts = 1 ./2 .^(10:-1:2) #14->7 good plot

prob = prob_sde_linear_stratonovich
sim  = test_convergence(dts,prob,EulerHeun(),numMonte=Int(5e2))
@test abs(sim.𝒪est[:l2]-1) < 0.1

sim  = test_convergence(dts,prob,LambaEulerHeun(),numMonte=Int(5e2))
@test abs(sim.𝒪est[:l2]-1) < 0.1

sim  = test_convergence(dts,prob,ISSEulerHeun(),numMonte=Int(5e2))
@test abs(sim.𝒪est[:l2]-1) < 0.1

sim  = test_convergence(dts,prob,ImplicitEulerHeun(),numMonte=Int(5e2))
@test abs(sim.𝒪est[:l2]-1) < 0.1

sim  = test_convergence(dts,prob,ImplicitEulerHeun(theta=1),numMonte=Int(5e2))
@test abs(sim.𝒪est[:l2]-1) < 0.1

sim  = test_convergence(dts,prob,ImplicitEulerHeun(symplectic=true),numMonte=Int(5e2))
@test abs(sim.𝒪est[:l2]-1) < 0.1

sim = test_convergence(dts,prob,RKMil(interpretation=:Stratonovich),numMonte=Int(5e2))
@test abs(sim.𝒪est[:l2]-1) < 0.2

sim  = test_convergence(dts,prob,ImplicitRKMil(interpretation=:Stratonovich),numMonte=Int(5e2))
@test abs(sim.𝒪est[:l2]-1) < 0.1

println("Now 2D")

prob = prob_sde_2Dlinear_stratonovich

sim  = test_convergence(dts,prob,EulerHeun(),numMonte=Int(5e1))
@test abs(sim.𝒪est[:l2]-1) < 0.1

sim  = test_convergence(dts,prob,LambaEulerHeun(),numMonte=Int(5e1))
@test abs(sim.𝒪est[:l2]-1) < 0.1

sim  = test_convergence(dts,prob,ISSEulerHeun(),numMonte=Int(5e1))
@test abs(sim.𝒪est[:l2]-1) < 0.1

sim  = test_convergence(dts,prob,ImplicitEulerHeun(),numMonte=Int(5e1))
@test abs(sim.𝒪est[:l2]-1) < 0.1

sim  = test_convergence(dts,prob,ImplicitEulerHeun(theta=1),numMonte=Int(5e1))
@test abs(sim.𝒪est[:l2]-1) < 0.1

sim  = test_convergence(dts,prob,ImplicitEulerHeun(symplectic=true),numMonte=Int(5e1))
@test abs(sim.𝒪est[:l2]-1) < 0.1

println("RKMils")

sim = test_convergence(dts,prob,RKMil(interpretation=:Stratonovich),numMonte=Int(1e2))
@test abs(sim.𝒪est[:l2]-1) < 0.2

sim  = test_convergence(dts,prob,ImplicitRKMil(interpretation=:Stratonovich),
                        numMonte=Int(1e2))
@test abs(sim.𝒪est[:l2]-1) < 0.1

sim  = test_convergence(dts,prob,ImplicitRKMil(theta=1,interpretation=:Stratonovich),
                        numMonte=Int(1e2))
@test abs(sim.𝒪est[:l2]-1) < 0.1

sim  = test_convergence(dts,prob,ImplicitRKMil(symplectic=true,interpretation=:Stratonovich),
                        numMonte=Int(1e2))
@test abs(sim.𝒪est[:l2]-1) < 0.1

Random.seed!(200)
sol = solve(prob,EulerHeun(),dt=1/4)


#=
  From BLOG example/aircraft-static.blog

  Distributions of variable numAircraft (num:0)
  3.00000 =>    3191  (0.319100)
  4.00000 =>    2608  (0.260800)
  2.00000 =>    1726  (0.172600)
  5.00000 =>    1385  (0.138500)
  6.00000 =>     657  (0.065700)
  1.00000 =>     221  (0.022100)
  7.00000 =>     192  (0.019200)
  8.00000 =>      13  (0.001300)
  9.00000 =>       7  (0.000700)

  location[1]
  Summary Stats:
  Length:         10000
  Missing Count:  0
  Mean:           149.209455
  Minimum:        100.031478
  1st Quartile:   123.426126
  Median:         149.309147
  3rd Quartile:   174.100740
  Maximum:        199.992805

  Also see ~/webppl/aircraft_static.wppl

=#

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

include("jl_utils.jl")

@model function aircraft_static()
    epsilon = 0.05

    # Poisson might get 0 which is not good for this model
    pp ~ Poisson(3.0)
    numAircraft ~ Dirac(pp == 0 ? 1 : pp)

    blipPerAircraft = tzeros(numAircraft)
    location = tzeros(numAircraft)
    blipLocation = tzeros(numAircraft)
    for b in 1:numAircraft
        blipPerAircraft[b] ~ Poisson(1)
        location[b] ~ Uniform(100.0,200.0)
        if blipPerAircraft[b] > 0
            blipLocation[b] ~ Normal(location[b],1.0)
        else
            blipLocation[b] ~ Uniform(90,210)
        end
    end

    sumBlips ~ Dirac(sum(blipPerAircraft))

    true ~ Dirac(sumBlips == 4)
    locationB0 ~ Dirac(location[1])
end

model = aircraft_static()

num_chns = 4
# chns = sample(model, Prior(), 1000)
# chns = sample(model, MH(), 100_000)
chns = sample(model, PG(5), 10_000)
# chns = sample(model, IS(), 10_000)
# chns = sample(model, SMC(), 40_000)
# chns = sample(model, SMC(), MCMCThreads(), 10_000, num_chns)

display(chns)
# display(plot(chns))

show_var_dist_pct(chns, :numAircraft)

println("\nlocation[1]")
display(summarystats(chns[Symbol("location[1]")]))