#=

  A coin application in Julia ConstraintSolver.jl

  From "The ECLiPSe Book" pages 99f and 234 ff
  The solution in ECLiPSe is at page 236.

  """
  What is the minimum number of coins that allows one to pay _exactly_
  any amount smaller than one Euro? Recall that there are six different
  euro cents, of denomination 1, 2, 5, 10, 20, 50
  """

  There are 4 optimal solutions (8 coins):
    x=[1,2,1,2,1,1]
    x=[1,2,2,1,1,1]
    x=[2,1,1,2,1,1]
    x=[2,1,2,1,1,1]

  This problem was implemented in coins3.jl

  This model implement a related problem:
  Is there a set of coin denominations which make it possible to change
  into 1..99 with less than 8 coins?

  Answer: Yes, there is. For example the following configuration:

  num_coins:7
  variables:[1,2,3,6,12,25,50]
  x:[1,1,1,1,1,1,1]

  Though, it come with a cost of using 7 different coin denominations.

  Also, we can see that there is a set with just 5 different
  denominations for the 8 coin change. However, it requires that
  more than one coin of some of the denominations

  This is done by comment this constraint:
        NumCoins =< 8,

  Result:
    num_coins:8
    variables:[1,2,4,11,33]
    x:[1,1,2,2,2]

Here's the run. As we see it can be done with:
- 8 coins of 5 (and 6) different denominations
- 7 coins of 7 different denominations.

"""
n:5
num_sols:1

solution #1
num_coins:8
coins:[1, 2, 4, 11, 33]
xs:[1, 1, 2, 2, 2]
 15.108249 seconds (24.47 M allocations: 12.097 GiB, 56.45% gc time)

n:6
num_sols:1

solution #1
num_coins:8
coins:[1, 2, 3, 5, 11, 33]
xs:[1, 1, 1, 1, 2, 2]
 21.561126 seconds (34.97 M allocations: 14.951 GiB, 59.19% gc time)

n:7
num_sols:1

solution #1
num_coins:7
coins:[1, 2, 3, 6, 13, 24, 50]
xs:[1, 1, 1, 1, 1, 1, 1]
 15.875412 seconds (16.94 M allocations: 14.988 GiB, 64.24% gc time)
"""

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#
using ConstraintSolver, JuMP
using Cbc, GLPK, Ipopt
const CS = ConstraintSolver
include("constraints_utils.jl")

function coins3_2(n,print_solutions=true,all_solutions=false,timeout=6)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)
    ipopt_optimizer = optimizer_with_attributes(Ipopt.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   "all_solutions"=> all_solutions,
                                                            # "all_optimal_solutions"=>all_solutions,
                                                            "logging"=>[],

                                                            "traverse_strategy" => :BFS,
                                                            # "traverse_strategy" => :DFS,
                                                            # "traverse_strategy" => :DBFS,

                                                            "branch_split" => :Smallest,
                                                            # "branch_split" => :Biggest,
                                                            # "branch_split" => :InHalf,

                                                            # https://wikunia.github.io/ConstraintSolver.jl/stable/options/#branch_strategy-(:Auto)
                                                            "branch_strategy" => :IMPS, # default
                                                            # "branch_strategy" => :ABS, # Activity Based Search
                                                            # "activity.decay" => 0.999, # default 0.999
                                                            # "activity.max_probes" => 10, # default, 10
                                                            # "activity.max_confidence_deviation" => 20, # default 20

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>timeout,

                                                            # "backtrack" => false, # default true
                                                            # "backtrack_sorting" => false, # default true

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))

    # Which coins should we use?
    # we assume that there are no coin of a denomination larger than 50.
    @variable(model, 1 <= coins[1:n] <= 50, Int)

    # Total number of coins used. Can we do it in 8 or less than 8 coins?
    @variable(model, 1 <= num_coins <= 8, Int)

    # Use either 1 or 2 coins of each denomination
    @variable(model, 1 <= xs[1:n] <= 2, Int)

    @constraint(model, coins in CS.AllDifferent())
    increasing_strict(model, coins) # symmetry breaking

    # Checks that all changes from 1 to 99 can be made.
    table = resize_matrix([ [i,j,i * j] for i in 0:99, j in 0:99])
    for j in 1:99
       tmp = @variable(model, [1:n], CS.Integers(0:99*99))

       # scalar_product(model,coins,tmp,j) # non-linear constraint
       # Reformulation:
       t = @variable(model, [1:n], CS.Integers(0:99*99)) # the product of coins[i] * tmp[i]
       for i in 1:n
           @constraint(model, [coins[i], tmp[i], t[i]] in CS.TableSet(table))
           # Ensure that we don't use more coins than in x[i]
           @constraint(model, tmp[i] <= xs[i])
       end
       @constraint(model, j == sum(t))

    end
    @constraint(model,num_coins == sum(xs))

    @objective(model,Min,num_coins)

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                println("solution #$sol")
                coins_val = convert.(Integer,JuMP.value.(coins; result=sol))
                xs_val = convert.(Integer,JuMP.value.(xs; result=sol))
                num_coins_val = convert.(Integer,JuMP.value.(num_coins; result=sol))
                println("num_coins:$num_coins_val")
                println("coins:$coins_val")
                println("xs:$xs_val")

            end
        end
    else
        println("status:$status")
    end

    return status
end

begin
  for n in 1:7
    println("\nn:$n")
    @time status = coins3_2(n ,true,false,125)
    if status == MOI.TIME_LIMIT
      break
    end
  end
end


"""
A kernel used in the bnhls covariance method.
"""
struct HFC_Kernel{R<:Real}
    f::Function
    abs_k_pp_0::R
    k_00::R
    c_star::R
    relative_asymptotic_efficiency::R
    domain::Tuple{R,R}
    function HFC_Kernel(f::Function, abs_k_pp_0::R, k_00::R, c_star::R, relative_asymptotic_efficiency::R, domain::Tuple{R,R}= (-Inf, Inf)) where R<:Real
        return new{R}(f, abs_k_pp_0, k_00, c_star, relative_asymptotic_efficiency, domain)
    end
end

function parzen_func(x::Real)
    if x < 0.5 return 1 -6*(x^2) + 6*(x^3) end
    return x < 1 ? 2*((1-x)^2) : 0.0
end
"""
A parzen kernel used in the bnhls covariance method.
"""
const parzen = HFC_Kernel(parzen_func, 12.0, 0.269, 3.51, 0.97, (-1.0,1.0))
fqs(x)   = abs(x) < 0.001 ? 1.0 : (3/(x^2))*(sin(x)/x - cos(x))
"""
A quadratic_spectral kernel used in the bnhls covariance method.
"""
const quadratic_spectral = HFC_Kernel(fqs, 0.2, 3*pi/5, 0.46, 0.93)
f_fejer(x)   = abs(x) < 0.001 ? 1.0 : (sin(x)/x)^2
"""
A fejer kernel used in the bnhls covariance method.
"""
const fejer  = HFC_Kernel(f_fejer, 2/3, pi/3, 0.84, 0.94)
f_tukey(x)   = (sin(exp(-x)*(pi/2)))^2
"""
A tukey_hanning kernel used in the bnhls covariance method.
"""
const tukey_hanning  = HFC_Kernel(f_tukey, (pi^2)/2, 0.52, 2.16, 1.06)
f_bnhls(x) = (1+x)*exp(-x)
"""
A bnhls_2008 kernel used in the bnhls covariance method.
"""
const bnhls_2008 = HFC_Kernel(f_bnhls, 1.0, 5/4, 0.96, 1.09)

"""
Realised autocovariance.
"""
function realised_autocovariance(returns::Array{R,2}, h::Integer) where R<:Real
    d = size(returns)[2]
    h = abs(h)
    start = h+1
    end_ = size(returns)[1]
    summed = zeros(d,d)
    for j in start:end_
        summed += returns[j,:] * transpose(returns[j-h,:])
    end
    return summed
end

"""
This averages the first few and last few returns. We do this to returns rather than
prices (as suggested in BNHLS 2011).

## References
Barndorff-Nielsen, O., Hansen, P.R., Lunde, A., Shephard, N. 2011. - Section 2.2.
"""
function preaveraging_end_returns(returns::Array{R,2}, m::Integer) where R<:Real
    N = size(returns)[1]
    if m >= (N/2) error("m is too high compared to the dimension of the returns matrix. Not possible to average away this much data.") end
    # Preaveraging of first few returns.
    first_m_rows  = returns[1:m,:]
    for i in 1:m first_m_rows[i,:] = first_m_rows[i,:] .* (-i/m) end
    first_row = sum(first_m_rows, dims = 1)
    # Preaveraging of last few returns.
    last_m_rows   = returns[(N-m+1):N,:]
    for i in 1:m last_m_rows[i,:] = last_m_rows[i,:] .* (i/m) end
    last_row = sum(last_m_rows, dims = 1)
    # returning the completed returns.
    new_returns   = deepcopy(returns[m:(N-m+1),:])
    new_returns[1,:]          = first_row
    new_returns[(N-2*m+2),:]  = last_row
    return new_returns
end

"""
This calculates covariance with the multivariate realised kernel of BNHLS(2011).
"""
function bnhls_covariance_estimate_given_returns(returns::Array{R,2}; kernel::HFC_Kernel{T}, H::Real, m::Integer) where R<:Real where T<:Real
    returns_end_averaged  = preaveraging_end_returns(returns, m)
    N = size(returns)[1]
    termination = Integer(floor(min(N, maximum(abs.(kernel.domain))*H)))
    summed = realised_autocovariance(returns,0) # as kernel.f(0) is always 1.
    for h in 1:termination
        kern = kernel.f(h/H)
        a1 = realised_autocovariance(returns,h)
        summed  += kern .* (a1 + transpose(a1)) # We are doing the negative h values at the same time as positives here.
    end
    return Hermitian(summed)
end

"""
    bnhls_covariance(ts::SortedDataFrame, assets::Vector{Symbol} = get_assets(ts);
                     regularisation::Union{Missing,Symbol} = :covariance_default,
                     regularisation_params::Dict = Dict(), only_regulise_if_not_PSD::Bool = false,
                     kernel::HFC_Kernel{<:Real} = parzen,
                     H::Real = kernel.c_star * ( mean(map(a -> length(ts.groupingrows[a]), assets))   )^0.6,
                     m::Integer = 2)

This calculates covariance with the Multivariate realised kernel oof BNHLS(2011).
### Inputs
* `ts` - The tick data.
* `assets` - The assets you want to estimate volatilities for.
* `regularisation` - A symbol representing what regularisation technique should be used. If missing no regularisation is performed.
* `regularisation_params` - keyword arguments to be consumed in the regularisation algorithm.
* `only_regulise_if_not_PSD` - Should regularisation only be attempted if the matrix is not psd already.
* `kernel` - The kernel used. See the paper for details.
* `H` - The number of lags/leads used in estimation. See the paper for details.
* `m` - The number of end returns to average.
### Returns
* A `CovarianceMatrix`.

### References
Barndorff-Nielsen, O., Hansen, P.R., Lunde, A., Shephard, N. 2011. - The whole paper but particularly 2.2, 2.3 here. Kernels are in table 1. choices of H are discussed in section 3.4 of the paper.
"""
function bnhls_covariance(ts::SortedDataFrame, assets::Vector{Symbol} = get_assets(ts);
                          regularisation::Union{Missing,Symbol} = :covariance_default,
                          regularisation_params::Dict = Dict(), only_regulise_if_not_PSD::Bool = false,
                          kernel::HFC_Kernel{<:Real} = parzen,
                          H::Real = kernel.c_star * ( mean(map(a -> length(ts.groupingrows[a]), assets))   )^0.6,
                          m::Integer = 2)
    at_times = get_all_refresh_times(ts, assets)
    dd_compiled = latest_value(ts, at_times; assets = assets)
    dd = get_returns(dd_compiled)

    if m >= nrow(dd)/2
       @warn string("Cannot estimate the correlation matrix with the bnhls method with only ", nrow(ts.df), " ticks.")
       return make_nan_covariance_matrix(assets, ts.time_period_per_unit)
   end

    returns = Matrix(dd[:, assets])
    cov_mat = bnhls_covariance_estimate_given_returns(returns; kernel = kernel, H = H, m = m)

    # Regularisation
    dont_regulise = ismissing(regularisation) || (only_regulise_if_not_PSD && is_psd_matrix(cov_mat))
    cov_mat = dont_regulise ? cov_mat : regularise(cov_mat, ts, assets, regularisation; regularisation_params... )

    # In some cases we get negative terms on the diagonal with this algorithm.
    negative_diagonals = findall(diag(cov_mat) .< eps())
    covar = make_nan_covariance_matrix(assets, ts.time_period_per_unit)
    if length(negative_diagonals) == 0
        cor, vols = cov2cor_and_vol(cov_mat, duration(ts; in_dates_period = false))
        covar.correlation = cor
        covar.volatility = vols
     end
    return covar
end


using Distributions
using Test
using DataFrames
using CSVFiles
using DataStructures

################################################################################
## JULIA
################################################################################
include("../src/utils.jl")
include("../src/sample_sobol.jl")
include("../src/analyze_sobol.jl")
include("../src/test_functions/ishigami.jl")

# define the (uncertain) parameters of the problem and their distributions
data = SobolData(
    params = OrderedDict(:x1 => Uniform(-3.14159265359, 3.14159265359),
        :x2 => Uniform(-3.14159265359, 3.14159265359),
        :x3 => Uniform(-3.14159265359, 3.14159265359)),
    N = 1000
)

N = data.N
D = length(data.params)

# sampling
julia_samples = sample(data) |> DataFrame
julia_ishigami = ishigami(convert(Matrix, julia_samples)) |> DataFrame

# analysis
julia_A, julia_B, julia_AB, julia_BA = split_output(convert(Matrix, julia_ishigami), N, D, data.calc_second_order)
julia_results = analyze(data, convert(Matrix, julia_ishigami); num_resamples = 1_000)

################################################################################
## Python
################################################################################

# sampling
py_samples = load("data/py_uniform/py_samples.csv", header_exists=false, colnames = ["x1", "x2", "x3"]) |> DataFrame
py_ishigami = load("data/py_uniform/py_ishigami.csv", header_exists=false) |> DataFrame

# analysis
py_A = load("data/py_uniform/py_A.csv", header_exists=false) |> DataFrame
py_B = load("data/py_uniform/py_B.csv", header_exists=false) |> DataFrame
py_AB = load("data/py_uniform/py_AB.csv", header_exists=false) |> DataFrame
py_BA = load("data/py_uniform/py_BA.csv", header_exists=false) |> DataFrame

py_firstorder = load("data/py_uniform/py_firstorder.csv", header_exists=false) |> DataFrame
py_secondorder = load("data/py_uniform/py_secondorder.csv", header_exists=false) |> DataFrame
py_totalorder = load("data/py_uniform/py_totalorder.csv", header_exists=false) |> DataFrame

py_firstorder_conf = load("data/py_uniform/py_firstorder_conf.csv", header_exists=false) |> DataFrame
py_secondorder_conf = load("data/py_uniform/py_secondorder_conf.csv", header_exists=false) |> DataFrame
py_totalorder_conf = load("data/py_uniform/py_totalorder_conf.csv", header_exists=false) |> DataFrame

################################################################################
## Testing
################################################################################

@testset "Uniform Sampling" begin
    @test convert(Matrix, julia_samples) ≈ convert(Matrix, py_samples) atol = 1e-9
end

@testset "Uniform Analysis" begin
    @test convert(Matrix, julia_ishigami) ≈ convert(Matrix, py_ishigami) atol = 1e-9
    @test julia_A ≈ convert(Matrix, py_A) atol = 1e-9
    @test julia_B ≈ convert(Matrix, py_B) atol = 1e-9
    @test julia_AB ≈ convert(Matrix, py_AB) atol = 1e-9
    @test julia_BA ≈ convert(Matrix, py_BA) atol = 1e-9
    @test julia_results[:firstorder] ≈ convert(Matrix, py_firstorder) atol = 1e-9
    @test julia_results[:totalorder] ≈ convert(Matrix, py_totalorder) atol = 1e-9

    for i = 1:D
        for j = i+1:D
            @test julia_results[:secondorder][i,j] ≈ convert(Matrix, py_secondorder)[i,j] atol = 1e-9
        end
    end

    # TODO Choose proper tolerance for CI comparison
    @test julia_results[:firstorder_conf] ≈ convert(Matrix, py_firstorder_conf) atol = 1e-1
    @test julia_results[:totalorder_conf] ≈ convert(Matrix, py_totalorder_conf) atol = 1e-1

    for i = 1:D
        for j = i+1:D
            @test julia_results[:secondorder_conf][i,j] ≈ convert(Matrix, py_secondorder_conf)[i,j] atol = 1e-1
        end
    end
end


module DescriptorSystems

using LinearAlgebra
using MatrixEquations
using MatrixPencils
using Polynomials
using Random

import LinearAlgebra: BlasFloat, BlasReal, BlasComplex, copy_oftype, transpose, adjoint, opnorm, normalize, rdiv!
import Base: +, -, *, /, \, (==), (!=), ^, isapprox, iszero, convert, promote_op, size, length, ndims,
             hcat, vcat, hvcat, inv, show, lastindex, require_one_based_indexing, print, show, one, zero, eltype
import MatrixPencils: isregular, rmeval
import Polynomials: AbstractRationalFunction, AbstractPolynomial, poles, isconstant, variable, degree, pqs

export DescriptorStateSpace, AbstractDescriptorStateSpace, dss, dssdata, rdss, rss, iszero, order
export RationalTransferFunction, rtf
export gminreal, gir, gir_lrtran, gbalmr, gsvselect, gss2ss, dss2ss
export gbilin, c2d, dss2rm, dss2pm, timeresp
export confmap, simplify, normalize, poles, gain, zpk, rtfbilin, numpoly, denpoly, isconstant, sampling_time
export blockdiag, eye, rcond
export gdual, ctranspose, inv, ldiv, rdiv, ginv
export append, series, parallel, horzcat, vertcat
export order, evalfr, dcgain, opnorm, freqresp, chess
export gpole, gzero, gpoleinfo, gzeroinfo, gnrank, isregular, isproper, isstable,
       glinfnorm, ghinfnorm, gl2norm, gh2norm, ghanorm, gnugap
export gsdec, grnull, glnull, grange, gcrange, grsol, glsol, grmcover1, grmcover2, glmcover1, glmcover2
export grcf, glcf, grcfid, glcfid, gnrcf, gnlcf, giofac, goifac, grsfg, glsfg
export gnehari, glinfldp, glasol, grasol
# export PencilStateSpace, pss, pssdata

abstract type AbstractDynamicalSystem end
abstract type AbstractLTISystem <: AbstractDynamicalSystem end
abstract type AbstractDescriptorStateSpace <: AbstractLTISystem end

include("types/DescriptorStateSpace.jl")
include("types/RationalTransferFunction.jl")
#include("types/PencilStateSpace.jl")
include("ginv.jl")
include("dss.jl")
include("connections.jl")
include("operations.jl")
include("conversions.jl")
include("order_reduction.jl")
include("analysis.jl")
include("decompositions.jl")
include("factorizations.jl")
include("covers.jl")
include("linsol.jl")
include("nullrange.jl")
include("model_matching.jl")
include("dstools.jl")
include("dsutils.jl")
end


using ITensors
using UnicodePlots

include("ising_model.jl")
include("orthogonalize_eigenvectors.jl")

#
# Compute excitations of the Ising model with DMRG
#

# Define the sites of the Hilbert space
n = 10
i = [Index(2, "S=1/2") for j in 1:n];

# Define the Hamiltonian as an MPO
h = 0.5 # field
H = MPO(ising(n; h), i);

# Define the initial state
psi_init = randomMPS(i; linkdims=2)

# Use DMRG as a reference
println("\nComputing ground state with DMRG")
E0, psi0 = dmrg(H, psi_init; weight=4, nsweeps=10, cutoff=1e-5);

println("\nComputing first excited state with DMRG")
E1, psi1 = dmrg(H, [psi0], psi_init; weight=4, nsweeps=10, cutoff=1e-5, maxdim=10);

println("\nComputing second excited state with DMRG")
E2, psi2 = dmrg(H, [psi0, psi1], psi_init; weight=4, nsweeps=10, cutoff=1e-5, maxdim=10);

# Collect the ground state and excited states in a vector
psis = [psi0, psi1, psi2]
nstates = length(psis)

# How close are our eigenvectors to eigenstates?
# We can use the energy variance:
#
# <psi|H^2|psi> - <psi|H|psi>^2
#
# to find out.
#
# Hints:
#
# 1. You can use `inner(psi', H, psi)` to compute `<psi|H|psi>`.
# 2. You can use `apply(H, psi)` to compute `H|psi>`.
for psi in psis
  Hpsi = apply(H, psi)
  @show inner(Hpsi, Hpsi) - inner(psi', H, psi)^2
end

# Which states did we find? The `expect` function
# computes the local expectation value of an operator
# on each site

# Use the UnicodePlots function `lineplot` to plot
# <Sz> and <Sx> for our eigenvectors
for i in 1:nstates
  plt = lineplot(expect(psis[i], "Sz"); title="Ising model state $i from DMRG", name="Sz", xlim=[1,10], ylim=[-0.5,0.5])
  lineplot!(plt, expect(psis[i], "Sx"); name="Sx")
  display(plt)
end

# Construct an effective Hamiltonian `Heff` for the states we found
Heff = [inner(psis[i]', H, psis[j]) for i in 1:nstates, j in 1:nstates]

# Construct the overlap matrix `N` for the states we found
N = [inner(psis[i], psis[j]) for i in 1:nstates, j in 1:nstates]

println("\nHeff = <psi_i|H|psi_j> from DMRG")
display(Heff)

println("\nN = <psi_i|psi_j> from DMRG")
display(N)

#
# Orthogonalize the eigenvectors by solving the generalized eigenvector
# equations:
#
# <psi_i|H|psi_j><psi_j|v> = lambda_i * <psi_j|v>
#
# See `orthogonalize_eigenvectors.jl` for details.
#
psis = orthogonalize_eigenvectors(H, psis)

# Replot <Sz> and <Sx> for our eigenvectors
for i in 1:nstates
  plt = lineplot(expect(psis[i], "Sz"); title="Ising model state $i from DMRG after orthogonalization", name="Sz", xlim=[1,10], ylim=[-0.5,0.5])
  lineplot!(plt, expect(psis[i], "Sx"); name="Sx")
  display(plt)
end

# Construct an effective Hamiltonian `Heff` for the states we found
Heff = [inner(psis[i]', H, psis[j]) for i in 1:nstates, j in 1:nstates]

# Construct the overlap matrix `N` for the states we found
N = [inner(psis[i], psis[j]) for i in 1:nstates, j in 1:nstates]

println("\nHeff = <psi_i|H|psi_j> from DMRG")
display(Heff)

println("\nN = <psi_i|psi_j> from DMRG")
display(N)


@doc doc"""
    Torus{N} <: AbstractPowerManifold

The n-dimensionsl torus is the $n$-dimensional product of the [´Circle`](@ref).

The Circle is stored internally within `manifold`, such that all functions of
[`AbstractPowerManifold`](@ref) can be used directly.
"""
struct Torus{N} <: AbstractPowerManifold{Circle{ℝ},MultidimentionalArrayPowerRepresentation}
    manifold::Circle{ℝ}
end

Torus(n::Int) = Torus{n}(Circle())

@doc doc"""
    check_manifold_point(M::Torus{n},x)

check whether `x` is a valid point on the [`GraphManifold`](@ref)
"""
check_manifold_point(::Torus, ::Any)
function check_manifold_point(M::Torus{N}, x; kwargs...) where {N}
    if length(x) != N
        return DomainError(
            length(x),
            "The number of elements in `x` ($(length(x))) does not match the dimension of the torus ($(N)).",
        )
    end
    return check_manifold_point(PowerManifold(M.manifold, N), x; kwargs...)
end

function check_tangent_vector(M::Torus{N}, x, v; kwargs...) where {N}
    if length(x) != N
        return DomainError(
            length(x),
            "The number of elements in `x` ($(length(x))) does not match the dimension of the torus ($(N)).",
        )
    end
    if length(v) != N
        return DomainError(
            length(v),
            "The number of elements in `v` ($(length(v))) does not match the dimension of the torus ($(N)).",
        )
    end
    return check_tangent_vector(PowerManifold(M.manifold, N), x, v; kwargs...)
end

get_iterator(M::Torus{N}) where {N} = 1:N

@generated manifold_dimension(::Torus{N}) where {N} = N

@generated representation_size(::Torus{N}) where {N} = (N,)


using GlobalSensitivity, SafeTestsets
using Test

const GROUP = get(ENV, "GROUP", "All")
const is_APPVEYOR = Sys.iswindows() && haskey(ENV,"APPVEYOR")
const is_TRAVIS = haskey(ENV,"TRAVIS")

@time begin

if GROUP == "All" || GROUP == "GSA"
    @time @safetestset "Morris Method" begin include("morris_method.jl") end
    @time @safetestset "Sobol Method" begin include("sobol_method.jl") end
    @time @safetestset "DGSM Method" begin include("DGSM.jl") end
    @time @safetestset "eFAST Method" begin include("eFAST_method.jl") end
    @time @safetestset "RegressionGSA Method" begin include("regression_sensitivity.jl") end
    @time @safetestset "DeltaMoment Method" begin include("delta_method.jl") end
end

end


begin
    using LinearAlgebra, DataFrames,Optim, ForwardDiff, BenchmarkTools,Distributions,
        Expectations, QuantEcon, Statistics, GLM
    using Distributions: invsqrt2π, log2π, sqrt2, invsqrt2
    using DCDC
    using Test
    using Distributed, Suppressor
    using Plots
end

begin "Dynamic Decision Process"
    σ₀ = 1;
    ρ = 0.15;
    β = 1 / ( 1 + ρ);
    r = 0.2;
    α = 0.9;
    A = 1 + r;
    ddc = DynamicDecisionProcess(σ₀,β,α,A);
    # The log function has a solution
    value_function(s) = (log(1 - β * α) + (β * α) * log(β * α) /(1-β * α) ) / (1 - β) + α * log(s) /(1-β * α);
    policy_function(s) = (1 - β * α) * A * s^α;
    ddc.ValueFn.y = value_function.(ddc.ValueFn.xdata);
    ddc.PolicyFn.y = policy_function.(ddc.PolicyFn.xdata[:,1]);
    # computeEquilibrium(ddc);
    # ee = EulerEquation(σ₀,β,α);
end

# "Check Approximation of Policy Function"
begin
    scatter(ddc.PolicyFn.xdata[:,1],ddc.PolicyFn.y);
    scatter!(ddc.PolicyFn.xdata[:,1],ddc.PolicyFn(ddc.PolicyFn.xdata));
    scatter!(ddc.PolicyFn.xdata[:,1],policy_function.(ddc.PolicyFn.xdata[:,1]));
    # scatter!(ddc.PolicyFn.xdata[:,1],A * ddc.PolicyFn.xdata[:,1].^α)
end

# "Check Approximation of Value Function"
begin
    scatter(ddc.ValueFn.xdata[:,1] ,ddc.ValueFn.y);
    scatter!(ddc.ValueFn.xdata[:,1],ddc.ValueFn(ddc.ValueFn.xdata[:,1]));
end

begin "Check whether the optimal choice is correct"
    c_opt = deepcopy(ddc.PolicyFn.y);
    s = ddc.ValueFn.xdata;
    η = ddc.PolicyFn.xdata[:,2];
    ϵ = minimum(s)*0.9;
    s = ddc.PolicyFn.xdata[:,1];
    lb = ϵ;
    for n = 1:ddc.nSolve
        ub =  ddc.trans(0,s[n]) - ϵ;
        @show c_opt[n] = find_optim(s[n],η[n],lb,ub,ddc.β,ddc.trans,ddc.util,ddc.ValueFn);
    end
    s̃ = ddc.trans(c_opt,s);
    y=ddc.util(c_opt) + ddc.β * ddc.ValueFn(s̃);
    scatter(ddc.PolicyFn.xdata[:,1],c_opt);
    scatter!(ddc.PolicyFn.xdata[:,1],ddc.PolicyFn.y);
end

begin
    scatter(ddc.ValueFn.xdata[:,1],ddc.ValueFn.y);
    scatter!(ddc.ValueFn.xdata[:,1],value_function.(ddc.ValueFn.xdata[:,1]));
    scatter!(ddc.ValueFn.xdata[1:end-1,1],y);
end

begin
    c_opt = deepcopy(ddc.PolicyFn.y);
    s = ddc.PolicyFn.xdata[:,1];
    s̃ = ddc.trans(c_opt,s);
    y=ddc.util(c_opt) + ddc.β * ddc.ValueFn(s̃);
    scatter(ddc.ValueFn.xdata[:,1],ddc.ValueFn.y);
    scatter!(ddc.ValueFn.xdata[:,1],value_function.(ddc.ValueFn.xdata[:,1]));
    scatter!(ddc.ValueFn.xdata[1:end-1,1],y)
end

begin
    ee_scatter=[check_ee(ddc) for i = 1:100];
    @show mean(ee_scatter)
    scatter(ee_scatter);
end

#_____________________________________________________________________________
# linearize estimation method
#_____________________________________________________________________________

begin
    β_MC=zeros(10);
    for i = 1:10

        nM = 100;
        nT = 2;
        data = simulate_ddc(nM,nT,ddc);
        # data =simulate_ee(nM,nT,ee);
        # The moment condition is
        # u'(c_t) = β d_Trans(c_t,x_t) u'(c_t+1)
        a_t = []; a_t1 = [];
        s_t = []; s_t1 = [];
        for t = 1:(nT-1)
             a_t  = vcat(a_t,data.action[:,t]);
             s_t  = vcat(s_t, [s.s[1] for s in data.state[:,t]] );
             a_t1 = vcat(a_t1,data.action[:,t+1])
            s_t1  = vcat(s_t1, [s.s[1] for s in data.state[:,t+1]] );
        end
        a_t=convert(Array{Float64,1},a_t);
        s_t=convert(Array{Float64,1},s_t);
        a_t1=convert(Array{Float64,1},a_t1);
        s_t1=convert(Array{Float64,1},s_t1);
        dtrans=ddc.dtrans;
        dtrans_s = (a,s) -> dtrans(a,s)[2].data;
        R = dtrans_s.(a_t1,s_t1);
        y=log.(a_t1) - log.(a_t)
        X=hcat(ones(length(y)),log.(R));
        b=inv(X'*X)*X'*y;
        @show β_MC[i]=exp(b[1]/b[2]);
    end

    @show b=inv(X'*X)*X'*y;
    @show 1/b[2];
    @show "True beta"
    @show β
    @show exp(b[1]/b[2]);
    "End";
end

begin "The estimation"
    policy_approx = ApproxFn(s_t1,a_t1,:epan,2);
    scatter(s_t1,a_t1);
    scatter!(s_t1,policy_approx(s_t1));

end
# plot!(ee.PolicyFn.xdata[:,1],ee.PolicyFn.y);
# scatter(ddc.PolicyFn.xdata[:,2]);
# checked = [check_ee(ddc) for i = 1:100];
# scatter(checked[checked .< 100])
#

#_____________________________________________________________________________
# Nonlinearize estimation method
#_____________________________________________________________________________

using Flux #Get derivatives
using Flux.Tracker
using Flux.Tracker: update!
# Solve the Euler equations

lb = [0,0];
ub = [1,100];

σ̂=2;β̂=0.8;

begin
    nM = 100;
    nT = 3;
    data = simulate_ddc(nM,nT,ddc);
    # data =simulate_ee(nM,nT,ee);
    # The moment condition is
    # u'(c_t) = β d_Trans(c_t,x_t) u'(c_t+1)
    a_t = [];
    s_t = [];
    for t = 1:(nT)
         global a_t  = vcat(a_t,data.action[:,t]);
         global s_t  = vcat(s_t, [s.s[1] for s in data.state[:,t]] );
    end
    a_t=convert(Array{Float64,1},a_t);
    s_t=convert(Array{Float64,1},s_t);
end

trans=ddc.trans;
function _moment(θ,trans,a_t,s_t,policy_approx)
    σ̂,β̂=θ; s_t1_simul=trans(a_t,s_t);
    a_t1_simul=policy_approx(s_t1_simul);
    util=Utility(σ̂); dutil=(c,η) -> Tracker.gradient(util,c,η)[1].data;
    # η=randn(100)./100;
    dc_t=dutil.(a_t,zeros(length(a_t)));
    dc_t1=dutil.(a_t1_simul,zeros(length(a_t)));
    dtrans_s = (a,s) -> dtrans(a,s)[2].data;
    R = dtrans_s.(a_t1_simul,s_t1_simul);
    ϵ= (β̂ * dc_t1 .* R ) ./ dc_t .-1;
    ϵ.* s_t;
end

function f_obj(θ)
    CF =_moment(θ,ddc.trans,a_t,s_t,policy_approx);
    sum(CF)^2;
end

result = optimize(f_obj,[0.1,0.01])
result.minimizer
f_obj(result.minimizer)
f_obj([β,1]+randn(2))

# Results of Optimization Algorithm
#  * Algorithm: Nelder-Mead
#  * Starting Point: [0.1,0.01]
#  * Minimizer: [0.47126793573588627,1.0128364461302402]
#  * Minimum: 2.226862e-11
#  * Iterations: 49
#  * Convergence: true
#    *  √(Σ(yᵢ-ȳ)²)/n < 1.0e-08: true
#    * Reached Maximum Number of Iterations: false
#  * Objective Calls: 98


module ConvergenceMHDTest

using Gridap
using Test
using Polynomials: fit

u(x) = VectorValue(x[1]^4+x[2]^4,-x[1]^2-x[2]^3,one(x[1])+sin(x[2]^4))
p(x) = sin(x[1]+x[2]+x[3])
j(x) = VectorValue(x[1]^4+x[2],-x[1]-cos(x[2]),one(x[1])+x[3]^4)
φ(x) = sin(x[1]+x[2]+x[3])
B = VectorValue(1.0,1.0,1.0)

∇u(x) = ∇(u)(x)
Δu(x) = Δ(u)(x)
∇p(x) = ∇(p)(x)
∇φ(x) = ∇(φ)(x)

f_u(x) = (∇u(x)')⋅u(x) - Δu(x) + ∇p(x) - j(x)×B
f_p(x) = (∇⋅u)(x)
f_j(x) = j(x) + ∇φ(x) - u(x)×B
f_φ(x) = (∇⋅j)(x)

g_u(x) = u(x)
g_j(x) = j(x)

eu_l2 = Vector{Float64}()
eu_h1 = Vector{Float64}()
ep_l2 = Vector{Float64}()
ej_l2 = Vector{Float64}()
ej_hdiv = Vector{Float64}()
eφ_l2 = Vector{Float64}()

order = 2
nxs = 2:5
domain = (-0.5,0.5,-0.5,0.5,-0.5,0.5)

for n=nxs
  partition = (n,n,n)
  model = CartesianDiscreteModel(domain,partition)

  labels = get_face_labeling(model)
  add_tag_from_tags!(labels,"dirichlet_u",collect(1:24))
  add_tag_from_tags!(labels,"dirichlet_j",collect(1:24))
  add_tag_from_tags!(labels,"neumann_u",collect(25:26))
  add_tag_from_tags!(labels,"neumann_j",collect(25:26))

  Vu = FESpace(model, ReferenceFE(lagrangian,VectorValue{3,Float64},order);
      conformity=:H1, dirichlet_tags="dirichlet_u")

  Vp = FESpace(model, ReferenceFE(lagrangian,Float64,order-1,space=:P);
      conformity=:L2)

  Vj = FESpace(model, ReferenceFE(raviart_thomas,Float64,order-1);
      conformity=:Hdiv, dirichlet_tags="dirichlet_j")

  Vφ = FESpace(model, ReferenceFE(lagrangian,Float64,order-1,space=:Q);
      conformity=:L2)

  U = TrialFESpace(Vu,g_u)
  P = TrialFESpace(Vp)
  J = TrialFESpace(Vj,g_j)
  Φ = TrialFESpace(Vφ)

  Y = MultiFieldFESpace([Vu, Vp, Vj, Vφ])
  X = MultiFieldFESpace([U, P, J, Φ])

  trian = Triangulation(model)
  degree = 2*(order)
  dΩ = Measure(trian,degree)

  uk = interpolate(u,U)
  function a(X,Y)
    u  , p  , j  , φ   = X
    v_u, v_p, v_j, v_φ = Y

    ∫((∇(u)'⋅uk)⋅v_u + inner(∇(u),∇(v_u)) - p*(∇⋅v_u) - (j×B)⋅v_u +
      (∇⋅u)*v_p +
      j⋅v_j - φ*(∇⋅v_j) - (u×B)⋅v_j +
      (∇⋅j)*v_φ )*dΩ
  end

  function l(Y)
    v_u, v_p, v_j, v_φ = Y

    ∫( v_u⋅f_u + v_p*f_p + v_j⋅f_j + v_φ*f_φ )*dΩ
  end

  btrian_u = BoundaryTriangulation(model,tags=["neumann_u"])
  dΓ_u = Measure(btrian_u,degree)
  nb_u = get_normal_vector(btrian_u)

  function l_Γ_u(Y)
    v_u, v_p, v_j, v_φ = Y

    ∫( v_u⋅(nb_u⋅∇u) - (nb_u⋅v_u)*p)*dΓ_u
  end

  btrian_j = BoundaryTriangulation(model,tags=["neumann_j"])
  dΓ_j = Measure(btrian_j,degree)
  nb_j = get_normal_vector(btrian_j)

  function l_Γ_j(Y)
    v_u, v_p, v_j, v_φ = Y

    ∫( -(v_j⋅nb_j)*φ)*dΓ_j
  end

  b(Y) = l(Y) + l_Γ_u(Y) + l_Γ_j(Y)

  op  = AffineFEOperator(a,b,X,Y)

  xh = solve(op)
  uh, ph, jh, φh = xh

  eu = uh - u
  ep = ph - p
  ej = jh - j
  eφ = φh - φ

  l2(v) = sqrt(sum(∫(v⋅v)*dΩ ))
  h1(v) = sqrt(sum(∫(v⋅v + inner(∇(v),∇(v)) )*dΩ ))
  hdiv(v) = sqrt(sum(∫(v⋅v + inner((∇⋅v),(∇⋅v)) )*dΩ ))

  append!(eu_l2, l2(eu) )
  append!(eu_h1, h1(eu) )
  append!(ep_l2, l2(ep) )
  append!(ej_l2, l2(ej) )
  append!(ej_hdiv, hdiv(ej) )
  append!(eφ_l2, l2(eφ) )

end

slope_eu_l2 = fit([log(x) for x in nxs], [log(y) for y in eu_l2], 1)[end]
slope_eu_h1 = fit([log(x) for x in nxs], [log(y) for y in eu_h1], 1)[end]
slope_ep_l2 = fit([log(x) for x in nxs], [log(y) for y in ep_l2], 1)[end]
slope_ej_l2 = fit([log(x) for x in nxs], [log(y) for y in ej_l2], 1)[end]
slope_ej_hdiv = fit([log(x) for x in nxs], [log(y) for y in ej_hdiv], 1)[end]
slope_eφ_l2 = fit([log(x) for x in nxs], [log(y) for y in eφ_l2], 1)[end]

@test abs(slope_eu_l2 + order + 1.0) < 2e-1
@test abs(slope_eu_h1 + order) < 2e-1
@test abs(slope_ep_l2 + order) < 2e-1
@test abs(slope_ej_l2 + order + 1.0) < 2e-1
@test abs(slope_ej_hdiv + order) < 2e-1
@test abs(slope_eφ_l2 + order) < 2e-1

end #module


const ϵ = 1e-8

"""
    TruncatedADAGrad(η=0.1, τ=1.0, n=100)

Implements a truncated version of AdaGrad in the sense that only the `n` previous gradient norms are used to compute the scaling rather than *all* previous. It has parameter specific learning rates based on how frequently it is updated.

## Parameters
  - η: learning rate
  - τ: constant scale factor
  - n: number of previous gradient norms to use in the scaling.
```
## References
[ADAGrad](http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf) optimiser.
Parameters don't need tuning.

[TruncatedADAGrad](https://arxiv.org/abs/1506.03431v2) (Appendix E).
"""
mutable struct TruncatedADAGrad
    eta::Float64
    tau::Float64
    n::Int

    iters::IdDict
    acc::IdDict
end

function TruncatedADAGrad(η = 0.1, τ = 1.0, n = 100)
    TruncatedADAGrad(η, τ, n, IdDict(), IdDict())
end

function apply!(o::TruncatedADAGrad, x, Δ)
    T = eltype(Tracker.data(Δ))

    η = o.eta
    τ = o.tau

    g² = get!(
        o.acc,
        x,
        [zeros(T, size(x)) for j = 1:o.n]
    )::Array{typeof(Tracker.data(Δ)), 1}
    i = get!(o.iters, x, 1)::Int

    # Example: suppose i = 12 and o.n = 10
    idx = mod(i - 1, o.n) + 1 # => idx = 2

    # set the current
    @inbounds @. g²[idx] = Δ^2 # => g²[2] = Δ^2 where Δ is the (o.n + 2)-th Δ

    # TODO: make more efficient and stable
    s = sum(g²)

    # increment
    o.iters[x] += 1

    # TODO: increment (but "truncate")
    # o.iters[x] = i > o.n ? o.n + mod(i, o.n) : i + 1

    @. Δ *= η / (τ + sqrt(s) + ϵ)
end

"""
    DecayedADAGrad(η=0.1, pre=1.0, post=0.9)

Implements a decayed version of AdaGrad. It has parameter specific learning rates based on how frequently it is updated.

## Parameters
  - η: learning rate
  - pre: weight of new gradient norm
  - post: weight of histroy of gradient norms
```
## References
[ADAGrad](http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf) optimiser.
Parameters don't need tuning.
"""
mutable struct DecayedADAGrad
    eta::Float64
    pre::Float64
    post::Float64

    acc::IdDict
end

DecayedADAGrad(η = 0.1, pre = 1.0, post = 0.9) = DecayedADAGrad(η, pre, post, IdDict())

function apply!(o::DecayedADAGrad, x, Δ)
    T = eltype(Tracker.data(Δ))

    η = o.eta
    acc = get!(o.acc, x, fill(T(ϵ), size(x)))::typeof(Tracker.data(x))
    @. acc = o.post * acc + o.pre * Δ^2
    @. Δ *= η / (√acc + ϵ)
end


module test_julia_ntuple

using Test

# https://discourse.juliacn.com/t/topic/1770/3
using Base.Cartesian: @ntuple

@generated add_tuples(t1::NTuple{N,Any}, t2::NTuple{N,Any}) where {N} = :(@ntuple $N i -> t1[i] + t2[i])

@test add_tuples((1,2,3), (4,5,6)) == (5, 7, 9)

@test ntuple(x->2x, 5) == (2, 4, 6, 8, 10)

end


function f(g, w, v)
    n = size(w, 1)
    q = HermitianSemidefinite(n)
    problem = maximize(real(tr(w * q')), [ Ψ(g, q) ⪯ v ])
    solve!(problem, () -> SCS.Optimizer(verbose=0, eps=solver_eps))
    return problem.optval
end

function g(g, w)
    D = g.D

    #z = HermitianSemidefinite(g.n)
    #problem = minimize(real(tr(D * z)), [ z ⪰ w, z in g.S1 ])
    #problem = minimize(real(tr(z)), [ z ⪰ √D * w * √D, z in g.S1 ])

    z = variable_in_space(g.S1)
    problem = minimize(real(tr(D * z)), [ z ⪰ w ])
    #problem = minimize(real(tr(z)), [ z ⪰ √D * w * √D ])

    solve!(problem, () -> SCS.Optimizer(verbose=0, eps=solver_eps))
    return problem.optval
end

#function h(g, w, y)
#    D = g.D
#
#    z = variable_in_space(g.S1)
#    problem = minimize(real(tr(y * √D * z * √D)), [ z ⪰ w ])
#    #problem = minimize(real(tr(D * √y * z * √y)), [ z ⪰ w ])
#    #problem = minimize(real(tr(z)), [ z ⪰ √D * √y * w * √y * √D ])
#
#    solve!(problem, () -> SCS.Optimizer(verbose=0, eps=solver_eps))
#    return problem.optval
#end

Random.seed!(0)

#sig = [1 2; 2 2]
#sig = [1 1; 2 3]
#sig = [2 3]
#sig = [2 2]
sig = [3 2; 2 3]

S = random_S0Graph(sig)
T = complement(S)

w = random_bounded(S.n)

@time opt0 = dsw(S, w, eps=solver_eps)[1]
@time opt1, x, y = dsw_via_complement(T, w, eps=solver_eps)
@test opt1 ≈ opt0  atol=tol
@time opt2 = dsw(T, y, eps=solver_eps)[1]
@test opt2 ≈ 1  atol=tol
@time opt3 = dsw(vertex_graph(S), √y * w * √y, eps=solver_eps)[1]
@test opt3 ≈ opt0  atol=tol
# ϑ(S, w) = max{ ϑ(S0, √y * w * √y) / ϑ(T, y) : y ∈ S1 }
@test opt0 ≈ opt3 / opt2  atol=tol

@test f(S, w, y) ≈ opt3  atol=tol
@test g(S, √y * w * √y) ≈ opt3  atol=tol


using Test
using Clustering

@testset "fuzzy_cmeans()" begin

Random.seed!(34568)

m = 3
n = 1000
k = 5

x = rand(m,n)

@testset "fuzziness = 2.0" begin
    fuzziness = 2.0
    Random.seed!(34568)
    r = fuzzy_cmeans(x, k, fuzziness)
    @test isa(r, FuzzyCMeansResult{Float64})
    @test size(r.centers) == (m,k)
    @test size(r.weights) == (n,k)
    @test sum(r.weights, dims=2) ≈ fill(1.0, n)
    @test all(0 .<= r.weights .<= 1)
end

@testset "fuzziness = 3.0" begin
    fuzziness = 3.0
    Random.seed!(34568)
    r = fuzzy_cmeans(x, k, fuzziness)
    @test isa(r, FuzzyCMeansResult{Float64})
    @test size(r.centers) == (m,k)
    @test size(r.weights) == (n,k)
    @test sum(r.weights, dims=2) ≈ fill(1.0, n)
    @test all(0 .<= r.weights .<= 1)
end

@testset "Abstract data matrix" begin
    fuzziness = 2.0
    Random.seed!(34568)
    r = fuzzy_cmeans(view(x, :, :), k, fuzziness)
    @test isa(r, FuzzyCMeansResult{Float64})
    @test size(r.centers) == (m,k)
    @test size(r.weights) == (n,k)
    @test sum(r.weights, dims=2) ≈ fill(1.0, n)
    @test all(0 .<= r.weights .<= 1)
end

end


# traceradvdiff.jl's test function

import FourierFlows.TracerAdvDiff

"""
    test_constvel(; kwargs...)

Advects a gaussian concentration c0(x, y, t) with a constant velocity flow
u(x, y, t) = uvel and v(x, y, t) = vvel and compares the final state with
cfinal = c0(x-uvel*tfinal, y-vvel*tfinal)
"""
function test_constvel(stepper, dt, nsteps)

  nx, Lx = 128, 2π
  uvel, vvel = 0.2, 0.1
  u(x, y) = uvel
  v(x, y) = vvel

  prob = TracerAdvDiff.Problem(; nx=nx, Lx=Lx, kap=0.0, u=u, v=v, dt=dt, stepper=stepper)
  s, v, p, g = prob.state, prob.vars, prob.params, prob.grid

  σ = 0.1
  c0func(x, y) = 0.1*exp.(-(x.^2+y.^2)/(2σ^2))

  c0 = c0func.(g.X, g.Y)
  tfinal = nsteps*dt
  cfinal = c0func.(g.X .- uvel*tfinal, g.Y .- vvel*tfinal)

  TracerAdvDiff.set_c!(prob, c0)

  stepforward!(prob, nsteps)
  TracerAdvDiff.updatevars!(prob)

  isapprox(cfinal, v.c, rtol=g.nx*g.ny*nsteps*1e-12)
end

"""
    test_timedependenttvel(; kwargs...)

Advects a gaussian concentration c0(x, y, t) with a time-varying velocity flow
u(x, y, t) = uvel and v(x, y, t) = vvel*sign(-t+tfinal/2) and compares the final
state with cfinal = c0(x-uvel*tfinal, y)
"""
function test_timedependentvel(stepper, dt, tfinal)

  nx, Lx = 128, 2π
  nsteps = round(Int, tfinal/dt)

  if !isapprox(tfinal, nsteps*dt, rtol=1e-12)
      error("tfinal is not multiple of dt")
  end

  uvel, vvel = 0.2, 0.1
  u(x, y, t) = uvel
  v(x, y, t) = t <= tfinal/2 ? vvel : -vvel

  prob = TracerAdvDiff.ConstDiffProblem(; nx=nx, Lx=Lx, kap=0.0, u=u, v=v, dt=dt, stepper=stepper)
  s, v, p, g = prob.state, prob.vars, prob.params, prob.grid

  σ = 0.1
  c0func(x, y) = 0.1*exp.(-(x.^2+y.^2)/(2σ^2))

  c0 = c0func.(g.X, g.Y)
  tfinal = nsteps*dt
  cfinal = c0func.(g.X .- uvel*tfinal, g.Y)

  TracerAdvDiff.set_c!(prob, c0func)

  stepforward!(prob, nsteps)
  TracerAdvDiff.updatevars!(prob)

  isapprox(cfinal, v.c, rtol=g.nx*g.ny*nsteps*1e-12)
end

"""
    test_diffusion(; kwargs...)

Diffuses a gaussian concentration c0(x, y, t) and compares the final state with
the analytic solution of the heat equation, cfinal
"""
function test_diffusion(stepper, dt, tfinal; steadyflow = true)

  nx  = 128
  Lx  = 2π
  kap = 0.01
  nsteps = round(Int, tfinal/dt)

  if !isapprox(tfinal, nsteps*dt, rtol=1e-12)
      error("tfinal is not multiple of dt")
  end

  grid = TwoDGrid(nx, Lx)
  prob = TracerAdvDiff.ConstDiffProblem(; steadyflow=steadyflow, grid=grid, nx=nx,
    Lx=Lx, kap=kap, dt=dt, stepper=stepper)
  s, v, p, g = prob.state, prob.vars, prob.params, prob.grid

  c0ampl, σ = 0.1, 0.1
  c0func(x, y) = c0ampl*exp.(-(x.^2+y.^2)/(2σ^2))

  c0 = c0func.(g.X, g.Y)
  tfinal = nsteps*dt
  σt = sqrt(2*kap*tfinal + σ^2)
  cfinal = c0ampl*σ^2/σt^2 * exp.(-(g.X.^2 + g.Y.^2)/(2*σt^2))

  TracerAdvDiff.set_c!(prob, c0)

  stepforward!(prob, nsteps)
  TracerAdvDiff.updatevars!(prob)

  isapprox(cfinal, v.c, rtol=g.nx*g.ny*nsteps*1e-12)
end

"""
    test_hyperdiffusion(; kwargs...)

Diffuses a gaussian concentration c0(x, y, t) using hyperdiffusivity and
compares the final state with the analytic solution of the heat equation, cfinal
"""
function test_hyperdiffusion(stepper, dt, tfinal; steadyflow = true)

    nx  = 128
    Lx  = 2π
    kap = 0.0   # no diffusivity
    eta = kap   # no diffusivity
   kaph = 0.01  # hyperdiffusivity coeff
  nkaph = 1     # nkaph=1 converts hyperdiffusivity to plain diffusivity
                # so we can compare with the analytic solution of heat equation

  nsteps = round(Int, tfinal/dt)

  if !isapprox(tfinal, nsteps*dt, rtol=1e-12)
      error("tfinal is not multiple of dt")
  end

   g = TwoDGrid(nx, Lx)

  u, v = 0*g.X, 0*g.X

  vs = TracerAdvDiff.Vars(g)
  pr = TracerAdvDiff.ConstDiffSteadyFlowParams(eta, kap, kaph, nkaph, u, v, g)
  eq = TracerAdvDiff.Equation(pr, g)
  ts = FourierFlows.autoconstructtimestepper(stepper, dt, eq.LC, g)
  prob = FourierFlows.Problem(g, vs, pr, eq, ts)

  c0ampl, σ = 0.1, 0.1
  c0func(x, y) = c0ampl*exp.(-(x.^2+y.^2)/(2σ^2))

  c0 = c0func.(g.X, g.Y)
  tfinal = nsteps*dt
  σt = sqrt(2*kaph*tfinal + σ^2)
  cfinal = c0ampl*σ^2/σt^2 * exp.(-(g.X.^2 + g.Y.^2)/(2*σt^2))

  TracerAdvDiff.set_c!(prob, c0)

  stepforward!(prob, nsteps)
  TracerAdvDiff.updatevars!(prob)

  isapprox(cfinal, vs.c, rtol=g.nx*g.ny*nsteps*1e-12)
end

# --
# Run tests
# --

stepper = "RK4"

dt, nsteps  = 1e-2, 40
@test test_constvel(stepper, dt, nsteps)

dt, tfinal  = 0.002, 0.1
@test test_timedependentvel(stepper, dt, tfinal)

dt, tfinal  = 0.005, 0.1
@test test_diffusion(stepper, dt, tfinal; steadyflow=true)

dt, tfinal  = 0.005, 0.1
@test test_diffusion(stepper, dt, tfinal; steadyflow=false)

dt, tfinal  = 0.005, 0.1
@test test_hyperdiffusion(stepper, dt, tfinal)


#=
# 2D heat equation, Dirichlet bc, CG
=#

### If the Finch package has already been added, use this line #########
using Finch # Note: to add the package, first do: ]add "https://github.com/paralab/Finch.git"

### If not, use these four lines (working from the examples directory) ###
# if !@isdefined(Finch)
#     include("../Finch.jl");
#     using .Finch
# end
##########################################################################

init_finch("heat2d");

useLog("heat2dlog")

# Set up the configuration
domain(2)                   # dimension
functionSpace(order=2)      # basis polynomial order
timeStepper(RK4)            # time stepper (optional second arg is CFL#)

# Specify the problem
mesh(QUADMESH, elsperdim=10) # 10x10 elements

u = variable("u")            # make a scalar variable with symbol u
testSymbol("v")              # sets the symbol for a test function

T = 1;
timeInterval(T)              # The time interval is 0 to T
initial(u, "abs(x-0.5)+abs(y-0.5) < 0.2 ? 1 : 0") # initial condition

boundary(u, 1, DIRICHLET, 0)  # boundary condition for BID 1 is Dirichlet with value 0

# Write the weak form
coefficient("f", "0.5*sin(6*pi*x)*sin(6*pi*y)")
weakForm(u, "Dt(u*v) + 0.01 * dot(grad(u),grad(v)) - f*v")

solve(u);

## Uncomment below to plot ##
# # solution is stored in the variable's "values"
# using Plots
# pyplot();
# display(plot(Finch.grid_data.allnodes[1,:], Finch.grid_data.allnodes[2,:], u.values[:], st = :surface))

finalize_finch() # Finish writing and close any files


using JITrench

f(x) = 2x^2 + 5x + 10

x = Variable(2.0, name="x")
#=
name: x
values: 2.0
creator: User-Defined (nothing)
=#

y = f(x)
#=
name: nothing
values: 28.0
creator: JITrench.Add
=#

backward!(y)

@show x.grad
# x.grad = Variable(13.0)


# # Barotropic QG beta-plane turbulence over topography
#
#md # This example can be run online via [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/barotropicqgtopography.ipynb).
#md # Also, it can be viewed as a Jupyter notebook via [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/barotropicqgtopography.ipynb).
#
# An idealized version of the Southern Ocean. We solve the barotropic
# quasi-geostrophic eddy dynamics in a flud with variable depth $H-h(x,y)$. We
# also include an "Antarctic Circumpolar Current," i.e.,  a domain-average zonal
# velocity $U(t)$ which is forced by constant wind  stress $F$ and influenced by
# bottom drag and topographic form stress. The equations solved are:
# $\partial_t \nabla^2\psi + \mathsf{J}(\psi-U y, \nabla^2\psi + \beta y + \eta) = -\mu\nabla^2\psi$
# and
# $\partial_t U = F - \mu U -  \langle\psi\partial_x\eta\rangle$.

using FourierFlows, Plots, Printf

using FFTW: irfft

import GeophysicalFlows.BarotropicQG
import GeophysicalFlows.BarotropicQG: energy, meanenergy, enstrophy, meanenstrophy

# ## Choosing a device: CPU or GPU

dev = CPU()     # Device (CPU/GPU)
nothing # hide

# ## Numerical parameters and time-stepping parameters
     nx = 128            # 2D resolution = nx^2
stepper = "FilteredRK4"  # timestepper
    dt  = 0.1            # timestep
 nsteps = 10000          # total number of time-steps
 nsubs  = 25             # number of time-steps for intermediate logging/plotting (nsteps must be multiple of nsubs)
nothing # hide

# ## Physical parameters

Lx = 2π        # domain size
 ν = 4e-15     # viscosity
nν = 4         # viscosity order
 β = 1.4015    # the y-gradient of planetary PV
 μ = 1e-2      # linear drag
 F = 0.0012    # normalized wind stress forcing on domain-averaged zonal flow U(t) flow
nothing # hide

# Define the topographic potential vorticity, $f_0 h(x, y)/H$
topoPV(x, y) = 2cos(4x)*cos(4y)
nothing # hide

# and the forcing function $F$ (here forcing is constant in time) that acts on the domain-averaged $U$ equation.
calcFU(t) = F
nothing # hide

# ## Problem setup
# We initialize a `Problem` by providing a set of keyword arguments,
prob = BarotropicQG.Problem(dev; nx=nx, Lx=Lx, β=β, eta=topoPV,
                  calcFU=calcFU, ν=ν, nν=nν, μ=μ, dt=dt, stepper=stepper)
nothing # hide

# and define some shortcuts.
sol, cl, vs, pr, gr = prob.sol, prob.clock, prob.vars, prob.params, prob.grid
x, y = gr.x, gr.y
nothing # hide

# ## Setting initial conditions

# Our initial condition is simply fluid at rest.
BarotropicQG.set_zeta!(prob, zeros(gr.nx, gr.ny))

# ## Diagnostics

# Create Diagnostics -- `energy`, `meanenergy`, `enstrophy`, and `meanenstrophy` functions are imported at the top.
E = Diagnostic(energy, prob; nsteps=nsteps)
Q = Diagnostic(enstrophy, prob; nsteps=nsteps)
Emean = Diagnostic(meanenergy, prob; nsteps=nsteps)
Qmean = Diagnostic(meanenstrophy, prob; nsteps=nsteps)
diags = [E, Emean, Q, Qmean]
nothing # hide

# ## Output

# We choose folder for outputing `.jld2` files and snapshots (`.png` files).
filepath = "."
plotpath = "./plots_barotropicqgtopography"
plotname = "snapshots"
filename = joinpath(filepath, "barotropicqgtopography.jld2")
nothing # hide

# Do some basic file management,
if isfile(filename); rm(filename); end
if !isdir(plotpath); mkdir(plotpath); end
nothing # hide

# and then create Output.
get_sol(prob) = sol # extracts the Fourier-transformed solution
get_u(prob) = irfft(im*gr.lr.*gr.invKrsq.*sol, gr.nx)
out = Output(prob, filename, (:sol, get_sol), (:u, get_u))
nothing # hide

# ## Visualizing the simulation

# We define a function that plots the potential vorticity field and the evolution
# of energy and enstrophy.

function plot_output(prob)

  pq = heatmap(x, y, vs.q,
                 c = :balance,
              clim = (-2, 2),
       aspectratio = 1,
             xlims = (-gr.Lx/2, gr.Lx/2),
             ylims = (-gr.Ly/2, gr.Ly/2),
            xticks = -3:3,
            yticks = -3:3,
            xlabel = "x",
            ylabel = "y",
             title = "∇²ψ + η",
        framestyle = :box)

  pE = plot(2,
             label = ["eddy energy" "mean energy"],
         linewidth = 2,
             alpha = 0.7,
             xlims = (-0.1, 10.1),
             ylims = (0, 0.0008),
            xlabel = "μt")

  pQ = plot(2,
             label = ["eddy enstrophy" "mean enstrophy"],
         linewidth = 2,
             alpha = 0.7,
             xlims = (-0.1, 10.1),
             ylims = (-0.02, 0.12),
            xlabel = "μt")

  l = @layout [ a{0.5w} grid(2, 1) ]
  p = plot(pq, pE, pQ, layout=l, size = (900, 600))

  return p
end
nothing # hide

# ## Time-stepping the `Problem` forward

# We time-step the `Problem` forward in time.

p = plot_output(prob)

startwalltime = time()

anim = @animate for j=0:Int(nsteps/nsubs)

  cfl = cl.dt*maximum([maximum(vs.U.+vs.u)/gr.dx, maximum(vs.v)/gr.dy])

  log = @sprintf("step: %04d, t: %d, cfl: %.2f, E: %.4f, Q: %.4f, walltime: %.2f min",
    cl.step, cl.t, cfl, E.data[E.i], Q.data[Q.i], (time()-startwalltime)/60)

  if j%(2000/nsubs)==0; println(log) end

  p[1][1][:z] = Array(vs.q)
  p[1][:title] = "∇²ψ + η, μt="*@sprintf("%.2f", μ*cl.t)
  push!(p[2][1], μ*E.t[E.i], E.data[E.i])
  push!(p[2][2], μ*Emean.t[Emean.i], Emean.data[Emean.i])
  push!(p[3][1], μ*Q.t[Q.i], Q.data[Q.i])
  push!(p[3][2], μ*Qmean.t[Qmean.i], Qmean.data[Qmean.i])

  stepforward!(prob, diags, nsubs)
  BarotropicQG.updatevars!(prob)

end

mp4(anim, "barotropicqg_acc.mp4", fps=18)

# Note that since mean flow enstrophy is $Q_U = \beta U$ it can attain negative values.

# ## Save

# Finally save the last snapshot.
savename = @sprintf("%s_%09d.png", joinpath(plotpath, plotname), cl.step)
savefig(savename)


export Subsampling_1D
###### define the subsampling function
using Random
# regular subsample
function Regular_subsample(ns::Int64,p::Int64)
    # ns: the number of receivers
    #   p=4 --> 75 % subsampling
    #   p=2 --> 50 % subsampling
    selected_locs = collect(1:p:ns);
    return selected_locs
end

# uniform subsampling % white noise subsample
function Uniform_subsample(ns::Int64,p::Int64)
    # ns: the number of receivers
    # p
    ratio = 1/p;
    Num = Int(round(ns*ratio));
	rand_locs = randperm(ns);
    selected_locs = rand_locs[1:Num];
    # return the subsampling locations for these data
    return sort(selected_locs)
end

#Optimal jittered subsampling %kind of blue noise subsample
function Optimal_Jittered_subsample(ns::Int64, p::Int64)
    # ns: the number of receivers
    # p
    b = collect(1:p)
    ind = []
    for i = 1:Int(round(ns/p))
        push!(ind,rand(b)+(i-1)*p)
    end
    return ind
end

#Blue noise subsampling w/ Farthest Point sampling
#Eldar, Yuval, et al. "The farthest point strategy for progressive image sampling." IEEE Transactions on Image Processing 6.9 (1997): 1305-1315.
function Blue_Noise_subsample(ns::Int64, p::Int64)
    # ns: the number of receivers
    # p
    ind = []
    #step 1: randomly choose one point
    push!(ind, randperm(ns)[1])
    ratio = 1/p;
    Num = Int(round(ns*ratio));
    for i = 1:Num-1
        #step 2: randomly choose a set of points
        cand_ind = rand(filter(x->!(x in ind), collect(1:ns)), p);
        Min_ind = []
        #step 3: calculate the smallest distance between the random set and the existed subpoints
        for j = 1:length(cand_ind)
            push!(Min_ind,minimum(abs.(cand_ind[j] .- ind)));
        end
        #step 4: find the largest distance among the random set
        Idx = findall(x-> x == maximum(Min_ind), Min_ind);
        #step 5: add one point
        push!(ind, cand_ind[Idx][1])
    end
    return sort(unique(ind))
end

function Subsampling_1D(ns::Int64, p::Int64; SubMethod = "Optimal_Jittered")
    #Regular: Period subsampling
    #Uniform: Uniform random subsampling (White noise subsampling)
    #Optimal_Jittered: Optimal Jittered subsampling
    #Blue_Noise: Blue noise subsampling w/ Farthest Point sampling
    SubMethod == "Regular" && return Regular_subsample(ns,p)
    SubMethod == "Uniform" && return Uniform_subsample(ns,p)
    SubMethod == "Blue_Noise" && return Blue_Noise_subsample(ns,p)
    return Optimal_Jittered_subsample(ns, p)
end


using OrdinaryDiffEq, DiffEqDevTools, Test

function nonauto1(u, p, t)
  x, _ = u
  [
    t*x
    0
  ]
end

function nonauto2(u, p, t)
  _, y = u
  [
   y,
   t*y
  ]
end

function analytic(u0,p,t)
  x0, y0 = u0
  et = exp(t^2/2)
  [
   et*(x0 + t*y0),
   et*y0
  ]
end

u0 = [1.1, 2.2]
tspan = (0.0,1.0)
prob1 = ODEProblem(ODEFunction{true}((du,u,p,t)->du .= nonauto1(u,p,t) .+ nonauto2(u,p,t),
                              analytic=analytic),
                              u0,tspan)
prob2 = ODEProblem(ODEFunction{false}((u,p,t)-> nonauto1(u,p,t) .+ nonauto2(u,p,t),
                              analytic=analytic),
                              u0,tspan)
prob3 = SplitODEProblem(SplitFunction{true}((du,u,p,t)->du .= nonauto1(u,p,t),
                                       (du,u,p,t)->du .= nonauto2(u,p,t),
                              analytic=analytic),
                              u0,tspan)
prob4 = SplitODEProblem(SplitFunction{false}(nonauto1,
                                        nonauto2,
                              analytic=analytic),
                              u0,tspan)

testTol = 0.2

kwargs = (reltol=1e-16, abstol=1e-16)
for prob in [prob1, prob2, prob3, prob4]
  dts = 1 .//2 .^(12:-1:8)
  sim = test_convergence(dts,prob,KenCarp3(); kwargs...)
  @test sim.𝒪est[:l∞] ≈ 3 atol=testTol
  dts = 1 .//2 .^(9:-1:6)
  sim = test_convergence(dts,prob,CFNLIRK3(); kwargs...)
  @test sim.𝒪est[:l∞] ≈ 3 atol=testTol
  sim = test_convergence(dts,prob,KenCarp4(); kwargs...)
  @test sim.𝒪est[:l∞] ≈ 4 atol=testTol
  sim = test_convergence(dts,prob,KenCarp47(); kwargs...)
  @test sim.𝒪est[:l∞] ≈ 4 atol=testTol
  dts = 1 .//2 .^(7:-1:4)
  sim = test_convergence(dts,prob,KenCarp5(); kwargs...)
  @test sim.𝒪est[:l∞] ≈ 5 atol=testTol
  sim = test_convergence(dts,prob,KenCarp58(); kwargs...)
  @test sim.𝒪est[:l∞] ≈ 5 atol=testTol
end
for prob in [prob1, prob2]
  dts = 1 .//2 .^(7:-1:4)
  sim = test_convergence(dts,prob,ESDIRK54I8L2SA(); kwargs...)
  @test sim.𝒪est[:l∞] ≈ 5 atol=testTol
end

#=
prob = prob2
dts = 1 .//2 .^(9:-1:6)
sim = test_convergence(dts,prob,KenCarp4())
sim.𝒪est[:l∞]

using Plots
plot(sim)
=#


#=

  Murder mystery

  From
  Andy Gordpn:
  "Reverend Bayes, meet Countess Lovelace: Probabilistic Programming for Machine Learning"
  https://channel9.msdn.com/Events/Lang-NEXT/Lang-NEXT-2012/Reverend-Bayes-meet-Countess-Lovelace-Probabilistic-Programming-for-Machine-Learning
  Around @14:00

  """
  Miss Scarlett dunnit 30%. Col Mustard dunnit 70%.
  Scarlett uses gun 3%, uses pipe 97%.
  Mustard uses gun 80%, uses pipe 20%.
  ...
  We found a gun at the Scene.
  What is the probability that Scarlett dunnit?
  """

  * Original question: we found a gun, what is the
    probability that Scarlett did it?

   It's not very likely that Scarlett did it!

   Summary Statistics
   parameters      mean       std   naive_se      mcse         ess      rhat   ess_per_sec
      Symbol   Float64   Float64    Float64   Float64     Float64   Float64       Float64

    scarlett    0.0142    0.1183     0.0012    0.0013   7957.8849    1.0001      605.3925
     mustard    0.7024    0.4572     0.0046    0.0057   7373.2272    1.0002      560.9150
     withGun    1.0000    0.0000     0.0000    0.0000         NaN       NaN           NaN
    withPipe    0.2031    0.4023     0.0040    0.0038   8323.5860    0.9999      633.2131

  * If we had found a pipe instead:

    It's a little more probable that Mustard did it. But Scarlett is still under
    investigation for it.

   Summary Statistics
   parameters      mean       std   naive_se      mcse         ess      rhat   ess_per_sec
      Symbol   Float64   Float64    Float64   Float64     Float64   Float64       Float64

    scarlett    0.6620    0.4731     0.0047    0.0059   6800.0219    1.0003      441.2161
     mustard    0.6979    0.4592     0.0046    0.0054   6727.5710    0.9999      436.5151
     withGun    0.2908    0.4542     0.0045    0.0060   6720.3716    1.0002      436.0480
    withPipe    1.0000    0.0000     0.0000    0.0000         NaN       NaN           NaN

  * Here's the priors (after we just heard about the murder)
    Mustard probably did it and it's more likely that it was with a Gun than a knife.

   Summary Statistics
   parameters      mean       std   naive_se      mcse         ess      rhat   ess_per_sec
      Symbol   Float64   Float64    Float64   Float64     Float64   Float64       Float64

    scarlett    0.2923    0.4548     0.0045    0.0049   9712.4171    1.0002    33035.4322
     mustard    0.7020    0.4574     0.0046    0.0052   8739.3445    1.0002    29725.6616
     withGun    0.5762    0.4942     0.0049    0.0052   9747.4247    1.0000    33154.5059
    withPipe    0.4247    0.4943     0.0049    0.0046   9917.5952    0.9999    33733.3169

  cf ~/blog/murder_mystery.blog
     ~/psi/murder_mystery.psi
     ~/webppl/murder_mystery.wppl

=#

using Turing, StatsPlots
include("jl_utils.jl")

@model function murder_mystery()
    scarlett ~ flip(0.30)
    mustard ~ flip(0.70)

    withGun ~ scarlett ? flip(0.03) : flip(0.80)
    withPipe ~ scarlett ? flip(0.97) : flip(0.20)

    true ~ Dirac(withGun == true)
    # true ~ Dirac(withPipe == true)
    # true ~ Dirac(withGun < withPipe)
    # true ~ Dirac(withGun > withPipe)

end

model = murder_mystery()

# chns = sample(model, Prior(), 10_000)
chns = sample(model, MH(), 10_000)
# chns = sample(model, PG(15), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 10_000)

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


using DiffEqBase, StochasticDiffEq, DiffEqDevTools, Test, LinearAlgebra, Random

u0 = ones(2)
A = [-3/2 1/20
      1/20 -3/2]
B = [1/5 1/100
    1/100 1/5]

function f(du,u,p,t)
  mul!(du,A,u)
end
function σ(du,u,p,t)
  mul!(@view(du[:,1]),B,u)
  mul!(@view(du[:,2]),B,u)
end

function (::typeof(f))(::Type{Val{:analytic}},u0,p,t,W)
  tmp = (A-(B^2))*t + B*W[1] + B*W[2]
  exp(tmp)*u0
end

prob2 = SDEProblem(f,σ,u0,(0.0,1.0),noise_rate_prototype=rand(2,2))

sol2 = solve(prob2,EM(),dt=1/100)

dts = 1 ./ 2 .^ (14:-1:7)

println("First Test")
Random.seed!(100)
sim  = test_convergence(dts,prob2,EM(),numMonte=100)
@test abs(sim.𝒪est[:l2]-0.5) < 0.1

# using Plots; plot(sim)

u0 = rand(2)
A = [2.0 -1.0; -1.0 2.0]
B = [1/5 1/100
    1/100 1/5]

function f(du,u,p,t)
  mul!(du,A,u)
  du .+= 1.01u
end
function σ(du,u,p,t)
  mul!(@view(du[:,1]),B,u)
  mul!(@view(du[:,2]),B,u)
end

function (::typeof(f))(::Type{Val{:analytic}},u0,p,t,W)
 tmp = (A+1.01I-(B^2))*t + B*sum(W)
 exp(tmp)*u0
end

prob2 = SDEProblem(f,σ,u0,(0.0,1.0),noise_rate_prototype=rand(2,2))

sol2 = solve(prob2,EM(),dt=1/100)

dts = 1 ./ 2 .^ (14:-1:7)

println("Second Test")
Random.seed!(100)
sim  = test_convergence(dts,prob2,EM(),numMonte=50)
# Superconvergence
@test abs(sim.𝒪est[:l2]-1.0) < 0.1

# using Plots; plot(sim)


function RTAU(DP)
      if (DP<=25.0)
          return rtauv=1.0
      elseif (DP>=50.0)
          return rtauv=0.1
      else
          return rtauv=1.9-.036*DP
      end
end


using Random
using Plots

function simulate_group(number_of_people::Int, group::Array, number_of_simulations::Int)
    local numerator = 0.0
    for i in 1:number_of_simulations::Int
        for k in 1:number_of_people
            group[k] = rand(1:365)
        end
        #local number_of_turns = rand(1:365, number_of_people)
        if length(unique(group)) < number_of_people
            numerator += 1.0
        end
    end
    return(numerator)
end

number_of_simulations = 4000

probability_shared_birthday_results = Array{Float32}(undef, 499)
time_per_simulation = Array{Float32}(undef, 499)

for num_people in 1:499
    local group = Array{Int16}(undef, num_people+1)
    stats = @timed numer = simulate_group(num_people+1, group, number_of_simulations)
    probability_shared_birthday_results[num_people] = numer/number_of_simulations
    time_per_simulation[num_people] = stats.time
end

plot([2:1:500], probability_shared_birthday_results, xlim=[2, 500], dpi=600, xlabel="Number of People", ylabel="Probability of Shared Birthdays", legend = false)
png("birthdayproblem.png")

plot([2:1:500], time_per_simulation, xlim=[2, 500], dpi=600, xlabel="Number of People", ylabel="Simulation Time (s)", legend = false)
png("birthdayproblemtime.png")


using BenchmarkTools
using MPI
using Trilinos
using Test

MPI.Init()

#Kokkos.initialize(Kokkos.OpenMP(), 4)

comm = Teuchos.MpiComm(MPI.CComm(MPI.COMM_WORLD))

my_rank = Teuchos.getRank(comm)
if my_rank != 0
  redirect_stdout(open("/dev/null", "w"))
  redirect_stderr(open("/dev/null", "w"))
end

function indefinite(n)
  # Generate an indefinite "hard" matrix
  srand(1)
  A = 20 * sparse(1.0I, n, n) + sprand(n, n, 20.0 / n)
  A = (A + A') / 2
  x = ones(n)
  b = A * x

  return A, b, ones(n)
end

# The following Laplace 2D implementation is based on:
# https://github.com/lruthotto/KrylovMethods.jl/blob/master/benchmarks/benchmark2DLaplacian.ipynb

"""
ibc,iin = getBoundaryIndices(n::Int)

returns indices of boundary and interior nodes of regular mesh with n^2 cells.
"""
function getBoundaryIndices(n::Int)
    ids = reshape(collect(1:(n+1)^2),n+1,n+1)
    iin = vec(ids[2:end-1,2:end-1])
    ibc = setdiff(vec(ids),vec(iin))
    return ibc,iin
end

"""
L,Lin,Lib = getLaplacian(n::Int)

returns discrete Laplacian, interior part, and boundary part, on regular mesh with n^2 cells.
"""
function getLaplacian(n::Int)
    h  = 1/n
    d2x = spdiagm(-1 => fill(-1.0/h,n), 0 => fill(2.0/h,n+1), 1 => fill(-1.0/h,n))
    d2x[1] /= 2
    d2x[end] /= 2
    L  = kron(sparse(1.0I, n+1, n+1),d2x) + kron(d2x,sparse(1.0I, n+1, n+1))

    # split into boundary and interior part
    ibc,iin = getBoundaryIndices(n)
    Lin = L[iin,iin]
    Lib = L[iin,ibc]
    return L,Lin,Lib
end
"""
g = getBoundaryConditions(n::Int,fctn=(x,y)->5.*y*sin(pi*x))

discretizes function on regular grid with n^2 cells and returns boundary and interior values as well as the source term value
"""
function getBoundaryConditions(n::Int,dirichlet = (x,y) -> (1-x^2)*(1-y^2), sourceterm = (x, y) -> 2*((1-x^2)+(1-y^2)))
    x  = linspace(0,1,n+1)
    g = map(Float64,[dirichlet(x1,x2) for x1 in x, x2 in x]')
    s = map(Float64,[sourceterm(x1,x2) for x1 in x, x2 in x]')
    ibc,iin = getBoundaryIndices(n)
    return vec(g)[ibc], vec(g)[iin], vec(s)[iin]
end

function laplace2d(n=1000)
  L,Lin,Lib = getLaplacian(n)
  ubc,refsol,f      = getBoundaryConditions(n)
  return Lin,vec(-Lib*ubc)+f,refsol
end

function tpetra_system(S)
  A,b,x = S
  n = length(b)
  rowmap = Tpetra.Map(n, 0, comm) #, Kokkos.KokkosOpenMPWrapperNode)
  At = Tpetra.CrsMatrix(A,rowmap)

  rmap = Tpetra.getRangeMap(At)
  bt = Tpetra.Vector(rmap)
  xt = Tpetra.Vector(rmap)

  num_my_elements = Tpetra.getNodeNumElements(rowmap)

  bv = Tpetra.device_view(bt)
  xv = Tpetra.device_view(xt)
  for local_row = 0:num_my_elements-1
    global_row = Tpetra.getGlobalElement(rmap, local_row)
    bv[local_row] = b[global_row+1]
    xv[local_row] = x[global_row+1]
  end

  return At, bt, xt
end

"""
Builds a Belos / Tpetra linear system, solving it once and checking the result against the reference solution
"""
function build_checked_system(f, solver_name, prec_type, n, tol, restart::Int, maxiter::Int)
  A, b, refsol = tpetra_system(f(n))

  println("Tolerance = ", tol, "; restart = ", restart, "; max #iterations = ", maxiter)

  params = Trilinos.default_parameters(solver_name)
  solver_params = params["Linear Solver Types"]["Belos"]["Solver Types"][solver_name]
  solver_params["Convergence Tolerance"] = tol
  solver_params["Verbosity"] = Belos.StatusTestDetails + Belos.FinalSummary + Belos.TimingDetails
  solver_params["Maximum Restarts"] = Int32(restart)
  solver_params["Block Size"] = Int32(maxiter)
  solver_params["Maximum Iterations"] = Int32(maxiter)
  params["Preconditioner Type"] = prec_type

  solver = TpetraSolver(A, params)

  sol = solver \ b
  v = Tpetra.device_view(sol)
  refv = Tpetra.device_view(refsol)
  maxerr = 0
  for i in LinearIndices(v)
    maxerr = max(abs(v[i] - refv[i]), maxerr)
  end
  @test maxerr < 100*tol

  return solver, b
end

sol_ndef(solver_name; n = 500000, tol = 1e-5, restart::Int = 15, maxiter::Int = 1500) = build_checked_system(indefinite, solver_name, "None", n, tol, restart, maxiter)

"""
Note: n is not the size of the matrix, but n in an n × n grid!
"""
sol_posdef(solver_name; n = 1000, tol = 1e-10, maxiter::Int = 1500) = build_checked_system(laplace2d, solver_name, "Ifpack2", n, tol, 0, maxiter)

suite = BenchmarkGroup()

suite["GMRES"] = @benchmarkable x=A\b setup=((A,b) = sol_ndef("GMRES"))
suite["BICGSTAB"] = @benchmarkable x=A\b setup=((A,b) = sol_ndef("BICGSTAB"))
suite["CG"] = @benchmarkable A\b setup=((A,b) = sol_posdef("CG",n=300))

result = run(suite)
if(my_rank == 0)
  display(result)
end

MPI.Finalize()


# # Fractional Kelvin-Voigt
#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/notebooks/model_kv.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/notebooks/model_kv.ipynb)

using RHEOS
## include a helper function for plotting
include("assets/plothelper.jl");
#-

Fract_KelvinVoigt

# #### Constitutive Equation
# ```math
# \sigma(t) = c_\alpha \frac{d^{\alpha} \epsilon(t)}{dt^{\alpha}}+ c_{\beta} \frac{d^\beta \epsilon(t)}{dt^\beta}
# ```

# ```math
# \text{for}\; \ 0 \leq \beta \leq \alpha \leq 1
# ```

# #### Relaxation Modulus
# ```math
# G(t) =\frac{c_{\alpha} }{\Gamma(1-\alpha)} t^{-\alpha}+\frac{c_{\beta} }{\Gamma(1-\beta)} t^{-\beta}
# ```

# #### Creep Modulus
# ```math
# J(t) = \frac{t^{\alpha}}{c_\alpha} E_{\alpha-\beta,1+\alpha}\left(-\frac{c_\beta}{c_\alpha} t^{\alpha-\beta}\right)
# ```

# #### Storage Modulus
# ```math
# G^{\prime}(\omega) = c_\alpha \omega^\alpha \cos\left(\alpha \frac{\pi}{2}\right) + c_\beta \omega^\beta \cos \left( \beta \frac{\pi}{2}\right)
# ```

# #### Loss Modulus
# ```math
# G^{\prime\prime}(\omega) = c_\alpha \omega^\alpha \sin\left(\alpha \frac{\pi}{2}\right) + c_\beta \omega^\beta \sin \left( \beta \frac{\pi}{2}\right)
# ```

# # Fractional (Spring) Kelvin-Voigt

FractS_KelvinVoigt
#-

models = Vector{RheoModel}()

## plot moduli for varying α
for alpha in [0.1, 0.3, 0.5, 0.7, 0.9]

    push!(models, RheoModel(FractS_KelvinVoigt, (cₐ = 1.0, a = alpha, k = 1.0)))

end

plotmodel(models, ymaxG = 5.0)

# # Fractional (Dashpot) Kelvin-Voigt

FractD_KelvinVoigt
#-

models = Vector{RheoModel}()

## plot moduli for varying β
for beta in [0.1, 0.3, 0.5, 0.7, 0.9]

    push!(models, RheoModel(FractD_KelvinVoigt, (η = 10, cᵦ= 1.0, β = beta)))

end

plotmodel(models, ymaxG = 3.0)

# # Kelvin-Voigt

KelvinVoigt
#-

models = Vector{RheoModel}()

## plot moduli for varying k
for η in [0.1, 1.0, 5.0]

    push!(models, RheoModel(KelvinVoigt, (η = η, k = 1)))

end

plotmodel(models)


### A Pluto.jl notebook ###
# v0.16.0

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ f4653ca6-5cba-11eb-0935-4d963ea430c8
begin
	using Pkg
	Pkg.add(["Plots", "PlutoUI","StatsBase", "DataStructures", "Unitful", "PhysicalConstants", "UnitfulEquivalences", "DataFrames", "CSV", "Dates", "TimeZones", "StatsPlots", "ImageFiltering", "OffsetArrays"])

	using Plots
	plotly()
	using PlutoUI
	using StatsBase
	using DataStructures
	using Unitful
	using PhysicalConstants.CODATA2018: c_0, h, ħ
	using UnitfulEquivalences
	using DataFrames
	using CSV
	using TimeZones
	using Dates
	using StatsPlots
	using ImageFiltering
	using OffsetArrays
end

# ╔═╡ cd4fe1e2-5cbb-11eb-392d-37cd89200e04
md"""
__Author__: Kristjan Laht, B59990
"""

# ╔═╡ e0f2ceba-5cba-11eb-1f14-ad4ef1619c20
md"""
# Functions in physics
1. In the figure below you see a graph of some function. It is a square function.
"""

# ╔═╡ b9290aee-5cba-11eb-16e7-c92402c6f6b8
md"""
### 2. Physical example

This function could be the function of velocity with regards to time with a constant _jerk_.

$$v(t) = \frac{jt^2}{2} + a_0t + v_0$$
$$1 \frac{m}{s} = \frac{1 \frac{m}{s^3} \cdot (1 s)^2}{2} + 1 \frac{m}{s^2} \cdot 1s + 1 \frac{m}{s}$$

* _$v$_ velocity - it's the first deriative of position, showing the change of change of postion.
* __$v_0$__ starting velocity
* __$a_0$__ starting acceleration. acceleration is the first deriative of velocity and the second deriative of position.
* __$j$__ jerk - the third deriative of position. constantly changing velocity.
* __$t$__ time
"""

# ╔═╡ cc76810a-5cba-11eb-3784-e9ccb5bc8ce2
md"""
### 3. Integral

$$v(t) = \frac{jt^2}{2} + a_0t + v_0$$
$$ds = v\ dt =  \left( \frac{jt^2}{2} + a_0t + v_0 \right) dt$$
$$\int_{s_0}^s ds = \int_0^t \left( \frac{jt^2}{2} + a_0t + v_0 \right) dt$$
$$\int_{s_0}^s ds = \int_0^t \left( \frac{jt^2}{2} + a_0t + v_0 \right) dt$$
$$s-s_0 = \frac{jt^3}{6} + \frac{a_0t^2}{2} + v_0t $$
$$s = \frac{jt^3}{6} + \frac{a_0t^2}{2} + v_0t + s_0$$

I chose to derive this integral with the limits of $(s_0, s)$ to get the positional displacement in familiar terms. (where $s_0$ is the starting positional displacement and $s$ is the final displacement)

The integral of the constant jerk velocity function with regards to time is the positional displacement physical quantity.
"""

# ╔═╡ ccc56cde-5cba-11eb-0601-bf2f29d0639a
md"""

### 4. Property of Nature

The square function describes the concept of compound growth. How even a linearly changing variable, when summed up over time, changes polynominally. This can happen when measuring speed and the jerk increases linearly.

"""

# ╔═╡ ccf312c4-5cba-11eb-24dc-c53f2185596b
square(a,x,c) = a*x^2 + x;

# ╔═╡ 282bfa7a-5cbb-11eb-3e20-818cbc5cadbb
@bind j Slider(2:50, show_value=true)

# ╔═╡ 72d21bfc-5cd1-11eb-2a88-5b7e0e16a951
@bind v_0 Slider(0:50, show_value=true)

# ╔═╡ 327c1456-5cbb-11eb-2537-29ee89edf2be
@bind s_0 Slider(-500:500, show_value=true);

# ╔═╡ 2aa87dd4-5cbb-11eb-3465-fbfff6973a43
x = 0:1:50;

# ╔═╡ e42deb9e-5cba-11eb-05bc-f99846852e79
begin
	plot(x,square.(j, x, v_0), xlabel="t (second)", ylabel="v (m/s)", label="v(t)")
	plot!([0], [v_0], seriestype = :scatter, label="v₀")
end

# ╔═╡ 687a0c52-5cbb-11eb-3b69-d3f0ba61397d
md"""
# Second function
#### 1.
In the figure below you see a graph of some function. It is an exponential decay function.
"""

# ╔═╡ 47a12e5c-5cc3-11eb-38c0-13546d4e6b2a
md"""
### 3. Derivative

The derivative of the discharge function is:

$$\frac{dV}{dt} = -\frac{V_0}{RC} e^{ \frac{-t}{RC}}$$

Which, if we move the $$C$$ to the other side, gets us the __current__ in the capacitor:

$$I(t) = C \frac{dV}{dt} = -\frac{V_0}{R} e^{\frac{-t}{RC}}$$

This only makes sense with regards to positive time t ∈ (0;∞) and positive voltage V ∈ (0;∞).

"""

# ╔═╡ d9b2a2f6-5da5-11eb-137a-1b28fc1e8851
md"""
### 4. Property of Nature

The energy stored inside a charged capactitor is potential energy. That is, it only exists as a relation to some other state (a discharged capacitor, perhaps). An ideal capacitor can be modeled with a constant *capacitance* $$C = \frac{Q}{V}$$. Therefore the voltage at any single time point is $$V(t) = \frac{Q(t)}{C}$$. That is, the voltage is linearly related to the charge stored inside the capacitor. Since the energy stored is potential energy, by way of positive charge on one plate and negative charge on the other plate, that energy/charge "wants" to settle back down onto a balanced equilibrium. The size of that "want" is proportional to how much imbalance there is.

That means that the discharge starts out fast, but slows down as the amount of imbalance drops off.
"""

# ╔═╡ 83c2566e-5cbd-11eb-073f-1d4a3d5518b8
exponential_decay(A, b, x) = A / exp(b*x)

# ╔═╡ 444988e0-5cbf-11eb-2406-f169032118cf
@bind V_0 Slider(0.2:0.2:4, show_value=true, default=3.4)

# ╔═╡ 5925a4d6-5cbf-11eb-3eeb-53518a7955eb
@bind RC Slider(0.2:0.2:10, show_value=true, default=5)

# ╔═╡ 68c1c22c-5cbb-11eb-3eb0-330e89f3f39d
begin
	plot(x, exponential_decay.(V_0, 1/RC, x), ylabel="V (volt)", xlabel="t (second)", label="V₀ e^(-t/RC)")
	plot!([0], [V_0], seriestype = :scatter, label="V₀")
end

# ╔═╡ a3077b76-5cc4-11eb-040f-251e086b20e3
RC_circuit_url = "https://upload.wikimedia.org/wikipedia/commons/a/a4/Discharging_capacitor.svg"

# ╔═╡ 1a449718-5cc1-11eb-1946-f114a25dc73d
md"""
### 2. Physical example
An example of exponential decay is the voltage of a simple __resistor-capacitor__ circuit, where the capacitor starts out charged and once the circuit is closed, the capacitor discharges its energy through the resistor.

$(Resource(RC_circuit_url, :width => 150))

$$V(t) = \frac{V_0}{ e^{\frac{t}{RC}}}$$

* __V(t)__ - function of voltage (V) over time (t)
* __V₀__ - voltage of capacitor at t = 0
* __t__ - time (in seconds)
* __R__ - resistance of resistor (in Ω)
* __C__ - capacitance of capaciotr (in farads)
* __e__ - the base of the natural logarithm

RC is the __exponential time constant__ τ.

"""

# ╔═╡ Cell order:
# ╟─f4653ca6-5cba-11eb-0935-4d963ea430c8
# ╟─cd4fe1e2-5cbb-11eb-392d-37cd89200e04
# ╟─e0f2ceba-5cba-11eb-1f14-ad4ef1619c20
# ╠═e42deb9e-5cba-11eb-05bc-f99846852e79
# ╟─b9290aee-5cba-11eb-16e7-c92402c6f6b8
# ╟─cc76810a-5cba-11eb-3784-e9ccb5bc8ce2
# ╟─ccc56cde-5cba-11eb-0601-bf2f29d0639a
# ╠═ccf312c4-5cba-11eb-24dc-c53f2185596b
# ╠═282bfa7a-5cbb-11eb-3e20-818cbc5cadbb
# ╠═72d21bfc-5cd1-11eb-2a88-5b7e0e16a951
# ╟─327c1456-5cbb-11eb-2537-29ee89edf2be
# ╠═2aa87dd4-5cbb-11eb-3465-fbfff6973a43
# ╟─687a0c52-5cbb-11eb-3b69-d3f0ba61397d
# ╠═68c1c22c-5cbb-11eb-3eb0-330e89f3f39d
# ╟─1a449718-5cc1-11eb-1946-f114a25dc73d
# ╟─47a12e5c-5cc3-11eb-38c0-13546d4e6b2a
# ╟─d9b2a2f6-5da5-11eb-137a-1b28fc1e8851
# ╠═83c2566e-5cbd-11eb-073f-1d4a3d5518b8
# ╠═444988e0-5cbf-11eb-2406-f169032118cf
# ╠═5925a4d6-5cbf-11eb-3eeb-53518a7955eb
# ╟─a3077b76-5cc4-11eb-040f-251e086b20e3


@testset "detrend()" begin
    # use detrend on polynomials and see if results are close to 0
    # (add 1. and see if results ≈ 1. as a workaround)
    for n ∈ 2:4
        @test begin
            x = collect(Float64, 1:20)
            y = x.^n
            all(yᵢ + 1.0 ≈ 1.0 for yᵢ ∈ nsi.detrend(x, y, n)) == true
        end
    end
    @test begin
        x = collect(Float64, 1:20)
        y = @. 0.1 + 2x - 0.01x^2
        all(yᵢ + 1.0 ≈ 1.0 for yᵢ ∈ nsi.detrend(x, y, 2)) == true
    end
end

@testset "NoiseSample" begin

    # test construction

    let
        ns1 = NoiseSample{Int}([1, 2, 3], [4, 5, 6])
        ns2 = NoiseSample([1, 2, 3], [4, 5, 6])
        ns3 = NoiseSample(Curve([1, 2, 3], [4, 5, 6]))
        ns4 = NoiseSample([4, 5, 6])
        for ns in (ns1, ns2, ns3, ns4)
            @test length(ns) == 3
            @test eltype(ns) == Int
            @test ns.x == [1, 2, 3]
            @test ns.y == [-1, 0, 1]
        end
    end

    # test basic operation

    noise = NoiseSample([1.0, 2.0], [4.0, 4.0])
    @test length(noise) == 2
    @test noise[end][1] == 2.0
    @test noise[end][2] + 1 ≈ 1.0
    @test [noise; noise] == NoiseSample([1.0, 2.0, 1.0, 2.0], [4.0, 4.0, 4.0, 4.0])

    # x and y need to have same length
    @test_throws ArgumentError NoiseSample([1.0], [4.0, 4.0])
    @test_throws ArgumentError NoiseSample([1.0, 2.0], [4.0])
end;

@testset "GaussianNoiseModel" begin
    @test eltype(GaussianNoiseModel(1)) == Int
    @test eltype(GaussianNoiseModel(1.)) == Float64
    @test_throws MethodError GaussianNoiseModel(1 + 1im)
end

@testset "MvGaussianNoiseModel" begin
    @test eltype(MvGaussianNoiseModel(1, 1, 1)) == Int
    @test eltype(MvGaussianNoiseModel(1., 1., 1.)) == Float64
    @test_throws MethodError MvGaussianNoiseModel(1 + 1im, 1, 1)
end

@testset "gauss_kernel()" begin
    @test nsi.gauss_kernel(0, [2, 1]) == 4.0
    @test nsi.gauss_kernel(2.0, [2, 2]) == 4.0 * exp(-0.5)
    @test nsi.gauss_kernel(4.0, [2, 1]) == 4.0 * exp(-8)
    @test nsi.gauss_kernel(4.0, [0, 1]) == 0.0
end

@testset "get_cov()" begin
    let f = 1.0000001
        @test_throws ArgumentError nsi.get_cov(1.0, -1, 2.0, 1.0)
        @test nsi.get_cov(1.0, 3, 2.0, 1.0) == 4 .* [f         exp(-0.5) exp(-2)  ;
                                                     exp(-0.5) f         exp(-0.5);
                                                     exp(-2)   exp(-0.5) f         ]
    end
end

@testset "add_noise()" begin
    let # use MvGaussianNoiseModel
        c = Curve([1., 2., 3.], [4., 5., 6.])
        nm = MvGaussianNoiseModel(1., 3., 5.)
        noise = add_noise(c, nm)
        @test length(noise) == 3
        @test noise.y |> first |> x -> x.particles |> length == 10_000
        noise = add_noise(c, nm, 1000)
        @test noise.y |> first |> x -> x.particles |> length == 1000
    end
    let # use GaussianNoiseModel
        c = Curve([1., 2., 3.], [4., 5., 6.])
        nm = GaussianNoiseModel(1.)
        noise = add_noise(c, nm)
        @test length(noise) == 3
        @test noise.y |> first |> x -> x.particles |> length == 10_000
        noise = add_noise(c, nm, 1000)
        @test noise.y |> first |> x -> x.particles |> length == 1000
    end
end

@testset "generate_noise() (correlated)" begin
    let
        nm = MvGaussianNoiseModel(1.0, 3.0, 5.0)
        noise = nsi.generate_noise(nm, 100, 2000)
        @test length(noise) == 100
        @test noise |> first |> x -> x.particles |> length == 2000
        ns = NoiseSample(nsi.get_draw.(1, noise))
        noise2 = nsi.generate_noise(ns, 100)
        @test noise2 |> first |> x -> x.particles |> length == 100
    end
end

@testset "generate_noise() (uncorrelated)" begin
    let
        nm = GaussianNoiseModel(1.0)
        noise = nsi.generate_noise(nm, 100, 2000)
        @test length(noise) == 100
        @test noise |> first |> x -> x.particles |> length == 2000
    end
end

@testset "fit_noise()" begin
    # this also tests estimate_autocov()
    let
        seed!(2)
        samples = 100
        len = 2_000
        for (δxᵢ, αᵢ, λᵢ) ∈ [(1.0, 3.0, 5.0),
                             (0.5, 2.0, 0.5),
                             (2.0, 0.5, 3.0)]
            nm = MvGaussianNoiseModel(δxᵢ, αᵢ, λᵢ)
            noise = nsi.generate_noise(nm, len, samples)
            A = Float64[]
            Λ = Float64[]
            for i in 1:samples
                x = δxᵢ:δxᵢ:(len * δxᵢ) |> collect
                ns = NoiseSample(x, nsi.get_draw.(i, noise))
                noise_param = fit_noise(ns)
                push!(A, noise_param.α)
                push!(Λ, noise_param.λ)
            end
            # To approach the true values, we need to have very long
            # noise samples, which gets computationally expensive
            # quickly. As a workaround, shorter noise samples
            # are used here, and the fitted parameters are averaged
            # over several trials.
            # When using a noise sample length of 2000 and averaging
            # over 100 trials, the fitted parameters should reproduce
            # the input parameters with 5% accuray:
            @test mean(A) ≈ αᵢ atol=5e-2
            @test mean(Λ) ≈ λᵢ atol=5e-2
        end
    end
end


## im2col

a = [1 4
     2 5
     3 6]
b = [7 10
     8 11
     9 12]
x = [a b]
cols = im2col(x,size(a),"distinct")
@test cols == [vec(a) vec(b)]
@test cols == im2col!(similar(cols),x,size(a),size(a))
@test_throws DimensionMismatch im2col!(Array(Int,1,2),x,size(a),size(a))

x = randn(4,1,7)
sz = (2,1,3)
f = rand(sz)
@test_approx_eq  vec(f)' * im2col(x,sz,"sliding")  imfilter(x,f,"inner")

x = reshape(1:(3*7),3,7)
@test x[:] == im2col(x,(1,1))[:]
@test im2col(x,(3,1),(1,3)) == x[:,1:3:end]
@test_throws DimensionMismatch im2col(x,(3,1),(1,4))
@test_throws DimensionMismatch im2col(x,(4,1))
@test_throws ArgumentError im2col(x,(1,1),"foo")

## col2im

szI = (4,3,2)
szW  = (2,1,2)
N = length(szI)
onetuple = ntuple(d->1,N)
I = randn(szI)

# identity
@test I == col2im(reshape(I,1,length(I)),onetuple,szI,"distinct")
@test I == col2im(reshape(I,1,length(I)),onetuple,szI,"sliding")
@test I == col2im(reshape(I,1,length(I)),onetuple,szI,onetuple)

# distinct
cols = im2col(I,szW,"distinct")
@test I == col2im(cols,szW,szI,"distinct")
@test I == col2im(cols,szW,szI,szW)
@test I == col2im!(zeros(I),cols,szW,szW)

# sliding with averaging
cols = im2col(I,szW,"sliding")
@test I == col2im(cols,szW,szI)
@test I == col2im(cols,szW,szI,"sliding")
@test I == col2im(cols,szW,szI,"sliding","average")
@test I == col2im(cols,szW,szI,onetuple)
@test I == col2im(cols,szW,szI,ntuple(d->1,N),+,0,/)

# sliding
I = rand(3,3)
cols = im2col(I,(2,2),"sliding")
w = [1 2 1; 2 4 2; 1 2 1]
ref = I .* w
@test_approx_eq  ref  col2im(cols,(2,2),(3,3),"sliding","sum")
@test_approx_eq  ref  col2im(cols,(2,2),(3,3),(1,1),+,0,nothing)
@test_approx_eq  ref  col2im!(zeros(I),cols,(2,2),(1,1),+)
ref = I .^ w
@test_approx_eq  ref  col2im(cols,(2,2),(3,3),(1,1),*,1,nothing)
@test_approx_eq  ref  col2im!(ones(I),cols,(2,2),(1,1),*)
ref = I
@test_approx_eq  ref  col2im(cols,(2,2),(3,3),"sliding","average")
@test_approx_eq  ref  col2im(cols,(2,2),(3,3),(1,1),+,0,/)
ref = I .* w - w
@test_approx_eq  ref  col2im(cols,(2,2),(3,3),(1,1),+,0,-)

# errors
I = rand(Int,4,4); cols = im2col(I,(2,2),"sliding")
@test_throws ArgumentError col2im(cols,(2,2),(4,4),"sliding","average")
I = rand(4,4); cols = im2col(I,(2,2),"sliding")
@test_throws DimensionMismatch col2im(cols,(2,2),(3,4))
@test_throws DimensionMismatch col2im(cols,(2,2),(4,4),"distinct")
@test_throws ArgumentError col2im(cols,(2,2),(4,4),"sliding","foo")
@test_throws ArgumentError col2im(cols,(2,2),(4,4),"foo")


"""
    transpose_to_axes(coord, b, n)

Convert from Hilbert index to a coordinatefor `b` bits
in `n` dimensions.
"""
function transpose_to_axes!(X::Vector{T}, b, n) where {T <: Integer}
    N = T(2) << (b - 1)
    # Gray decode by H^(H/2)
    t = X[n] >> 1
    for io = n:-1:2
        X[io] ⊻= X[io - 1]
    end
    X[1] ⊻= t
    # Undo excess work
    Q = T(2)
    while Q != N
        P = Q - one(T)
        for jo = n:-1:1
            if (X[jo] & Q) != zero(T)
                X[1] ⊻= P  # invert
            else  # exchange
                t = (X[1] ⊻ X[jo]) & P
                X[1] ⊻= t
                X[jo] ⊻= t
            end
        end
        Q <<= 1
    end
end

function axes_to_transpose!(X::Vector{T}, b, n) where {T <: Integer}
    M = one(T) << (b - 1)
    # Inverse undo
    Q = M
    while Q > one(T)
        P = Q - one(T)
        for io = 1:n
            if (X[io] & Q) != zero(T)
                X[1] ⊻= P
            else
                t = (X[1] ⊻ X[io]) & P
                X[1] ⊻= t
                X[io] ⊻= t
            end
        end
        Q >>= 1
    end
    # Gray encode
    for jo = 2:n
        X[jo] ⊻= X[jo - 1]
    end
    t2 = zero(T)
    Q = M
    while Q > one(T)
        if (X[n] & Q) != 0
            t2 ⊻= (Q - one(T))
        end
        Q >>= one(T)
    end
    for ko = 1:n
        X[ko] ⊻= t2
    end
end

"""
Takes a vector of length `n` and places the bits of all `n` integers
into a single integer. The vector's 1st component is the most significant bit.
"""
function interleave_transpose(::Type{T}, X::Vector, b, n) where {T}
    h = zero(T)
    for i in 0:(b - 1)
        for d in 1:n
            ith_bit = (X[d] & (1<<i)) >> i
            h |= ith_bit << (i * n + n - d)
        end
    end
    h
end

"""
Takes a single integer and places its values into components of a vector,
bit-by-bit.
"""
function outerleave_transpose!(X::Vector{T}, h, b, n) where {T <: Integer}
    X .= zero(T)
    for i in 0:(b-1)
        for d in 1:n
            ith_bit = (h & (one(h) << (i * n + n - d))) >> (i * n + n - d)
            X[d] |= ith_bit << i
        end
    end
end

"""
Takes a vector of length `n` and places the bits of all `n` integers
into a single integer. The vector's 1st component is the least significant bit.
"""
function interleave_transpose_low(::Type{T}, X::Vector{T}, b, n) where {T}
    h = zero(T)
    for i in 0:(b - 1)
        for d in 1:n
            h |= ((X[d] & (one(T)<<i))) << (i*(n - 1) + d - 1)
        end
    end
    h
end

function outerleave_transpose_low!(X::Vector{T}, h, b, n) where {T <: Integer}
    X .= zero(T)
    for i in 0:(b-1)
        for d in 1:n
            X[d] |= (h & (one(T) << (i * n + d - 1))) >> (i * (n - 1) + d - 1)
        end
    end
end

"""
    GlobalGray(b, n)
    GlobalGray(T, b, n)

`T` is a data type for the Hilbert index. It can be signed or unsigned,
as long as it has at least `n * b` bits. `n` is the number of dimensions,
and `b` is the bits per dimension, so each axis value should be between
0 and ``2^b - 1``, inclusive, for the zero-based interface. They should be
between 1 and ``2^b``, inclusive, for the one-based interface.

The GlobalGray algorithm is an n-dimensional Hilbert curve with a simplified
implementation. It follows an article, "Programming the Hilbert Curve," by
John Skilling, 707 (2004), http://dx.doi.org/10.1063/1.1751381.
I call it "Global Gray" because the insight of the article
is that a single, global Gray code can be applied to all
np bits of a Hilbert length.
"""
struct GlobalGray{T} <: HilbertAlgorithm{T}
    b::Int
    n::Int
end

axis_type(gg::GlobalGray) = large_enough_unsigned(gg.b)

function GlobalGray(b, n)
    ttype = large_enough_unsigned(b * n)
    GlobalGray{ttype}(b, n)
end

function GlobalGray(::Type{T}, b, n) where {T}
    GlobalGray{T}(b, n)
end

function encode_hilbert_zero!(g::GlobalGray{T}, X::Vector)::T where {T}
    axes_to_transpose!(X, g.b, g.n)
    interleave_transpose(T, X, g.b, g.n)
end

"""
    encode_hilbert_zero(ha::HilbertAlgorithm{T}, X::Vector{A})

Takes an n-dimensional vector `X` and returns a single integer of type
`T` which orders `X` to improve spatial locality. The input `X` has multiple
axes and the output is called a Hilbert index. This version is zero-based,
so each axis counts from 0, and the smallest Hilbert index is 0.
"""
function encode_hilbert_zero(g::GlobalGray{T}, X::Vector)::T where {T}
    Y = copy(X)
    encode_hilbert_zero!(g, Y)
end

"""
    decode_hilbert_zero!(ha::HilbertAlgorithm{T}}, X::Vector{A}, h::T)

Given a Hilbert index, `h`, computes an n-dimensional coordinate `X`. The type of
the Hilbert index is large enought to contain the bits of all dimensions of the
axis vector, `X`.
"""
function decode_hilbert_zero!(g::GlobalGray{T}, X::Vector, h::T) where {T}
    outerleave_transpose!(X, h, g.b, g.n)
    transpose_to_axes!(X, g.b, g.n)
end


"""
This type is used to store the solution to the problem presented
in the "Wald Friedman" notebook presented on the QuantEcon website.

Solution
----------
J : vector(Float64)
    Discretized value function that solves the Bellman equation
lb : scalar(Real)
    Lower cutoff for continuation decision
ub : scalar(Real)
    Upper cutoff for continuation decision
"""
type WFSolution
    J::Vector{Float64}
    lb::Real
    ub::Real
end

"""
This type is used to solve the problem presented in the "Wald Friedman"
notebook presented on the QuantEcon website.

Parameters
----------
c : scalar(Real)
    Cost of postponing decision
L0 : scalar(Real)
    Cost of choosing model 0 when the truth is model 1
L1 : scalar(Real)
    Cost of choosing model 1 when the truth is model 0
f0 : vector(Float64)
    A finite state probability distribution
f1 : vector(Float64)
    A finite state probability distribution
m : scalar(Int64)
    Number of points to use in function approximation
"""
immutable WaldFriedman
    c::Real
    L0::Real
    L1::Real
    f0::Vector{Float64}
    f1::Vector{Float64}
    m::Int64
    pgrid::LinSpace{Float64}
    sol::WFSolution
end

function WaldFriedman(c, L0, L1, f0, f1; m=25)
    pgrid = linspace(0.0, 1.0, m)

    # Renormalize distributions so nothing is "too" small
    f0 = clamp(f0, 1e-8, 1-1e-8)
    f1 = clamp(f1, 1e-8, 1-1e-8)
    f0 = f0 / sum(f0)
    f1 = f1 / sum(f1)
    J = zeros(m)
    lb = 0.
    ub = 0.

    WaldFriedman(c, L0, L1, f0, f1, m, pgrid, WFSolution(J, lb, ub))
end

"""
This function takes a value for the probability with which
the correct model is model 0 and returns the mixed
distribution that corresponds with that belief.
"""
function current_distribution(wf::WaldFriedman, p)
    return p*wf.f0 + (1-p)*wf.f1
end

"""
This function takes a value for p, and a realization of the
random variable and calculates the value for p tomorrow.
"""
function bayes_update_k(wf::WaldFriedman, p, k)
    f0_k = wf.f0[k]
    f1_k = wf.f1[k]

    p_tp1 = p*f0_k / (p*f0_k + (1-p)*f1_k)

    return clamp(p_tp1, 0, 1)
end

"""
This is similar to `bayes_update_k` except it returns a
new value for p for each realization of the random variable
"""
function bayes_update_all(wf::WaldFriedman, p)
    return clamp(p*wf.f0 ./ (p*wf.f0 + (1-p)*wf.f1), 0, 1)
end

"""
For a given probability specify the cost of accepting model 0
"""
function payoff_choose_f0(wf::WaldFriedman, p)
    return (1-p)*wf.L0
end

"""
For a given probability specify the cost of accepting model 1
"""
function payoff_choose_f1(wf::WaldFriedman, p)
    return p*wf.L1
end

"""
This function evaluates the expectation of the value function
at period t+1. It does so by taking the current probability
distribution over outcomes:

    p(z_{k+1}) = p_k f_0(z_{k+1}) + (1-p_k) f_1(z_{k+1})

and evaluating the value function at the possible states
tomorrow J(p_{t+1}) where

    p_{t+1} = p f0 / ( p f0 + (1-p) f1)

Parameters
----------
p : scalar
    The current believed probability that model 0 is the true
    model.
J : function (interpolant)
    The current value function for a decision to continue

Returns
-------
EJ : scalar
    The expected value of the value function tomorrow
"""
function EJ(wf::WaldFriedman, p, J)
    # Pull out information
    f0, f1 = wf.f0, wf.f1

    # Get the current believed distribution and tomorrows possible dists
    # Need to clip to make sure things don't blow up (go to infinity)
    curr_dist = current_distribution(wf, p)
    tp1_dist = bayes_update_all(wf, p)

    # Evaluate the expectation
    EJ = dot(curr_dist, J.(tp1_dist))

    return EJ
end

"""
For a given probability distribution and value function give
cost of continuing the search for correct model
"""
function payoff_continue(wf::WaldFriedman, p, J)
    return wf.c + EJ(wf, p, J)
end

"""
Evaluates the value function for a given continuation value
function; that is, evaluates

    J(p) = min( (1-p)L0, pL1, c + E[J(p')])

Uses linear interpolation between points
"""
function bellman_operator(wf::WaldFriedman, J)
    c, L0, L1, f0, f1 = wf.c, wf.L0, wf.L1, wf.f0, wf.f1
    m, pgrid = wf.m, wf.pgrid

    J_out = zeros(m)
    J_interp = LinInterp(pgrid, J)

    for (p_ind, p) in enumerate(pgrid)
        # Payoff of choosing model 0
        p_c_0 = payoff_choose_f0(wf, p)
        p_c_1 = payoff_choose_f1(wf, p)
        p_con = payoff_continue(wf, p, J_interp)

        J_out[p_ind] = min(p_c_0, p_c_1, p_con)
    end

    return J_out
end

"""
This function takes a value function and returns the corresponding
cutoffs of where you transition between continue and choosing a
specific model
"""
function find_cutoff_rule(wf::WaldFriedman, J)
    m, pgrid = wf.m, wf.pgrid

    # Evaluate cost at all points on grid for choosing a model
    p_c_0 = payoff_choose_f0(wf, pgrid)
    p_c_1 = payoff_choose_f1(wf, pgrid)

    # The cutoff points can be found by differencing these costs with
    # the Bellman equation (J is always less than or equal to p_c_i)
    lb = pgrid[searchsortedlast(p_c_1 - J, 1e-10)]
    ub = pgrid[searchsortedlast(J - p_c_0, -1e-10)]

    return lb, ub
end

function solve_model(wf; tol=1e-7)
    bell_op(x) = bellman_operator(wf, x)
    J =  compute_fixed_point(bell_op, zeros(wf.m), err_tol=tol, print_skip=5)

    wf.sol.J = J
    wf.sol.lb, wf.sol.ub = find_cutoff_rule(wf, J)
    return J
end

"""
This function takes an initial condition and simulates until it
stops (when a decision is made).
"""
function simulate(wf::WaldFriedman, f; p0=0.5)
    # Check whether vf is computed
    if sumabs(wf.sol.J) < 1e-8
        solve_model(wf)
    end

    # Unpack useful info
    lb, ub = wf.sol.lb, wf.sol.ub
    drv = DiscreteRV(f)

    # Initialize a couple useful variables
    decision_made = false
    decision = 0
    p = p0
    t = 0

    while !decision_made
        # Maybe should specify which distribution is correct one so that
        # the draws come from the "right" distribution
        k = rand(drv)
        t = t+1
        p = bayes_update_k(wf, p, k)
        if p < lb
            decision_made = true
            decision = 1
        elseif p > ub
            decision_made = true
            decision = 0
        end
    end

    return decision, p, t
end

"""
Uses the distribution f0 as the true data generating
process
"""
function simulate_tdgp_f0(wf::WaldFriedman; p0=0.5)
    decision, p, t = simulate(wf, wf.f0; p0=p0)

    if decision == 0
        correct = true
    else
        correct = false
    end

    return correct, p, t
end

"""
Uses the distribution f1 as the true data generating
process
"""
function simulate_tdgp_f1(wf::WaldFriedman; p0=0.5)
    decision, p, t = simulate(wf, wf.f1; p0=p0)

    if decision == 1
        correct = true
    else
        correct = false
    end

    return correct, p, t
end

"""
Simulates repeatedly to get distributions of time needed to make a
decision and how often they are correct.
"""
function stopping_dist(wf::WaldFriedman; ndraws=250, tdgp="f0")
    if tdgp=="f0"
        simfunc = simulate_tdgp_f0
    else
        simfunc = simulate_tdgp_f1
    end

    # Allocate space
    tdist = Array(Int64, ndraws)
    cdist = Array(Bool, ndraws)

    for i in 1:ndraws
        correct, p, t = simfunc(wf)
        tdist[i] = t
        cdist[i] = correct
    end

    return cdist, tdist
end


### A Pluto.jl notebook ###
# v0.12.6

using Markdown
using InteractiveUtils

# ╔═╡ a80cb638-1dde-11eb-08dd-593b9630a62e
using Pkg, DrWatson

# ╔═╡ a80cedd8-1dde-11eb-2ac5-ff913973a868
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Plots, Images
end

# ╔═╡ 3bc81896-1dde-11eb-22ff-3372c06964a9
md"## Listing 1.13"

# ╔═╡ 924b85f8-1ddf-11eb-025f-6f6f1ba174c0
Pkg.add("ImageIO")

# ╔═╡ a80d6ee8-1dde-11eb-3a36-c1d068d268f0
begin
	img = load(datadir("stars.png"))
	gImg = red.(img)*0.299 + green.(img)*0.587 + blue.(img)*0.114
	rows, cols = size(img)
end

# ╔═╡ a8187068-1dde-11eb-290c-919a177c2f04
"Highest intensity pixel: ", findmax(gImg)

# ╔═╡ a818e0ca-1dde-11eb-2d10-6be0b50f7358
function boxBlur(image,x,y,d)
    if x<=d || y<=d || x>=cols-d || y>=rows-d
        return image[x,y]
    else
        total = 0.0
        for xi = x-d:x+d
            for yi = y-d:y+d
                total += image[xi,yi]
            end
        end
        return total/((2d+1)^2)
    end
end

# ╔═╡ a821fab6-1dde-11eb-36fd-db97e9de0689
begin
	blurImg = [boxBlur(gImg,x,y,5) for x in 1:cols, y in 1:rows]

	yOriginal, xOriginal = argmax(gImg).I
	yBoxBlur, xBoxBlur   = argmax(blurImg).I

	heatmap(gImg, c=:Greys, yflip=true)
end

# ╔═╡ a8223dee-1dde-11eb-3345-630187ae9ff6
p1 = scatter!((xOriginal, yOriginal), ms=60, ma=0, msw=4, msc=:red)

# ╔═╡ a82b3c52-1dde-11eb-1e09-f3ad27bbd9ee
heatmap(blurImg, c=:Greys, yflip=true)

# ╔═╡ a82bc988-1dde-11eb-17a5-67a907d95a5e
p2 = scatter!((xBoxBlur, yBoxBlur), ms=60, ma=0, msw=4, msc=:red)

# ╔═╡ a831c28e-1dde-11eb-2f74-53e1521437e3
plot(p1, p2, size=(800, 400), ratio=:equal, xlims=(0,cols), ylims=(0,rows),
	colorbar_entry=false, border=:none, legend=:none)

# ╔═╡ 0349ba5c-1de2-11eb-359e-cfe15670884f
md"## End of listing 1.13"

# ╔═╡ Cell order:
# ╟─3bc81896-1dde-11eb-22ff-3372c06964a9
# ╠═a80cb638-1dde-11eb-08dd-593b9630a62e
# ╠═a80cedd8-1dde-11eb-2ac5-ff913973a868
# ╠═924b85f8-1ddf-11eb-025f-6f6f1ba174c0
# ╠═a80d6ee8-1dde-11eb-3a36-c1d068d268f0
# ╠═a8187068-1dde-11eb-290c-919a177c2f04
# ╠═a818e0ca-1dde-11eb-2d10-6be0b50f7358
# ╠═a821fab6-1dde-11eb-36fd-db97e9de0689
# ╠═a8223dee-1dde-11eb-3345-630187ae9ff6
# ╠═a82b3c52-1dde-11eb-1e09-f3ad27bbd9ee
# ╠═a82bc988-1dde-11eb-17a5-67a907d95a5e
# ╠═a831c28e-1dde-11eb-2f74-53e1521437e3
# ╟─0349ba5c-1de2-11eb-359e-cfe15670884f