using MonteCarloMarkovKernels
using StaticArrays
using Random
using LinearAlgebra
import Statistics.mean

zero2 = [0.0, 0.0]
μ1 = [-1.0,0.0]
Σ1 = [1.0 1.5 ; 1.5 3.0]
propSigma = [2.0 1.2 ; 1.2 3.0]

Szero2 = SVector{2, Float64}(zero2)
SpropSigma = SMatrix{2, 2, Float64}(propSigma)

## slightly more complicated, but much more efficient approach
## uses StaticArrays and scratch space
function makelogMVN(μ::SVector{d, Float64}, Σ::SMatrix{d, d, Float64}) where d
  invΣ = inv(Σ)
  lognc = - 0.5 * d * log(2 * π) - 0.5 * logdet(Σ)
  function lpi(x::SVector{d, Float64})
    v = x - μ
    return lognc - 0.5*dot(v, invΣ * v)
  end
  return lpi
end

logtarget = makelogMVN(SVector{2, Float64}(μ1), SMatrix{2, 2, Float64}(Σ1))
P_RWM = makeRWMKernel(logtarget, SpropSigma)

seed!(12345)
chain = simulateChain(P_RWM, Szero2, 2^15)

cc = MonteCarloMarkovKernels.cov(chain)

@test maximum(abs.(mean(chain) - μ1)) < 0.1
@test maximum(abs.(cc - Σ1)) < 0.2
@test 0.2 < P_RWM(:acceptanceRate) < 0.5

## test that only custom RNG is used
Random.seed!(1); v1 = rand(); Random.seed!(1)
P_RWM = makeRWMKernel(logtarget, SpropSigma, MersenneTwister(12345))
chain = simulateChain(P_RWM, Szero2, 256)
v2 = rand()
@test v1 == v2


function add_epsilon!(x::ListOfPoints{d}, i, epsilon) where d
  ei = SVector{d,Float64}([(j==i ? epsilon : 0.0) for j=1:d])
  for i=1:length(x)
    x[i] += ei
  end
end

function DiffFun(fun, x0::Vector{ListOfPoints{n_x}}, epsilon=1e-6) where n_x
    xi = deepcopy(x0)
    N = length(x0[1])
    n_m = length(x0)
    r0 = fun(x0)::Vector{ListOfPoints{n_x}}
    JMat = [zeros(n_x, n_x, N) for i=1:n_m]
    for i_x=1:n_x
      for i_m=1:n_m
        add_epsilon!(xi[i_m], i_x, epsilon)
      end
      fi = fun(xi)::Vector{ListOfPoints{n_x}}
      di = (fi-r0)/epsilon
      for i_m=1:n_m
        JMat[i_m][:,i_x,:] = reshape(reinterpret(Float64, vec(di[i_m])), (n_x, N))
        add_epsilon!(xi[i_m], i_x, -epsilon)
      end
    end
    J = [reshape(reinterpret(SMatrix{n_x,n_x,Float64,n_x^2},vec(JMat[i])),(N,)) for i=1:n_m]
    return (r0,J) #::Tuple{Vector{ListOfPoints{n_x}},Vector{Vector{SMatrix{n_x,n_x,Float64,n_x*n_x}}}}
end

struct NewtonResult
    solution
    iterations::Int64
    maxit::Int64
    err_ϵ::Float64
    tol_ϵ::Float64
    err_η::Float64
    tol_η::Float64
    errors::Vector{Float64}
end

function newton(fun::Function, x0::Vector{ListOfPoints{n_x}}, a::Union{Vector{ListOfPoints{n_x}},Nothing}=nothing, b::Union{Vector{ListOfPoints{n_x}},Nothing}=nothing; maxit=10, verbose=false, n_bsteps=5, lam_bsteps=0.5) where n_x

    steps = (lam_bsteps).^collect(0:n_bsteps)

    n_m = length(x0)
    N = length(x0[1])
    x = x0::Vector{ListOfPoints{n_x}}
    err_e_0 = 10.0
    tol_e = 1e-8
    err_η = 10.0
    tol_η = -1.0
    it=0
    errors = Float64[]

    while (it<maxit) && (err_e_0>tol_e)
        it += 1
        R_i, D_i = DiffFun(fun, x) # ::Tuple{Vector{ListOfPoints{n_x}},Vector{Vector{SMatrix{n_x,n_x,Float64,n_x*n_x}}}}
        if !(a isa Nothing)
            PhiPhi!(R_i,x,a,b,D_i)
        end
        err_e_0 = maxabs(R_i)
        push!(errors, err_e_0)
        err_e_1 = err_e_0
        dx = [[D_i[i][n]\R_i[i][n] for n=1:N] for i=1:n_m]::Vector{ListOfPoints{n_x}}
        err_η = maxabs(dx)
        i_bckstps = 0
        new_x = x
        while err_e_1>=err_e_0 && i_bckstps<length(steps)
            i_bckstps += 1
            new_x = x-dx*steps[i_bckstps]
            new_res = fun(new_x)::Vector{ListOfPoints{n_x}} # no diff
            if !(a isa Nothing)
                new_res = [PhiPhi0.(new_res[i],new_x[i],a[i],b[i]) for i=1:n_m]
            end
            err_e_1 = maxabs(new_res)
        end
        x = new_x::Vector{ListOfPoints{n_x}}
        if verbose
            println(i_bckstps, " : ", err_x, " : ", err_e_1, )
        end
    end

    res = NewtonResult(
        x,
        it,
        maxit,
        err_e_0,
        tol_e,
        err_η,
        tol_η,
        errors
    )
    return res
end


using MonteCarloX
using Random
using HypothesisTests
using Distributions
import Distributions.pdf
import Distributions.cdf

""" Testing if sampling single event correctly"""
function test_poisson_single()
    LambdaMaxs = [10.0, 20.0] # testing for two different max-sample-rates
    Lambda(t) = 10.0
    pass = true

    # Regular sampler
    alg = InhomogeneousPoisson()
    for LambdaMax in LambdaMaxs
        rng = MersenneTwister(1000)
        nSamples = 1000
        samples = zeros(nSamples)

        for i in 1:nSamples
            t0 = next_time(alg, rng, Lambda, LambdaMax)
            samples[i] = t0
        end

        test = HypothesisTests.ExactOneSampleKSTest(samples, Exponential(1.0 / 10.0))
        pass &= pvalue(test) > 0.05
    end

    # Decreasing sampler
    alg = InhomogeneousPoissonPiecewiseDecreasing()
    rng = MersenneTwister(1000)
    nSamples = 1000
    samples = zeros(nSamples)

    for i in 1:nSamples
        t0 = next_time(alg, rng, Lambda)
        samples[i] = t0
    end

    test = HypothesisTests.ExactOneSampleKSTest(samples, Exponential(1.0 / 10.0))
    pass &= pvalue(test) > 0.05

    return pass
end

""" Testing if sampling the uniform case correctly"""
function test_poisson_constant()
    LambdaMaxs = [1.0, 2.0] # testing for two different max-sample-rates
    Lambda(t) = 1.0
    pass = true

    alg = InhomogeneousPoisson()
    for LambdaMax in LambdaMaxs
        rng = MersenneTwister(1000)
        nSamples = 1000
        samples = zeros(nSamples)

        t0 = 0
        for i in 1:nSamples
            t0 += next_time(alg, rng, t->Lambda(t + t0), LambdaMax)
            samples[i] = t0 % 1
        end

        test = HypothesisTests.ExactOneSampleKSTest(samples, Uniform(0, 1))
        pass &= pvalue(test) > 0.05
    end

    return pass
end

""" Testing if sampling correctly from a shifted sine-wave distribution."""
function test_poisson_sin_wave()
    LambdaMaxs = [2.0, 3.0] # testing for two different max-sample-rates
    Lambda(t) = sin(t) + 1.0
    pass = true

    alg = InhomogeneousPoisson()
    for LambdaMax in LambdaMaxs
        rng = MersenneTwister(1000)
        nSamples = 1000
        samples = zeros(nSamples)

        t0 = 0
        for i in 1:nSamples
            t0 += next_time(alg, rng, t->Lambda(t + t0), LambdaMax)
            samples[i] = t0 % (2 * pi)
        end

        test = HypothesisTests.ExactOneSampleKSTest(samples, SinDistribution())
        pass &= pvalue(test) > 0.05
    end

    return pass
end

""" Definition of the custom sine-wave distribution for hypothesis testing."""
struct SinDistribution <: ContinuousUnivariateDistribution end

function Distributions.minimum(d::SinDistribution)
    0.0
end

function Distributions.maximum(d::SinDistribution)
    2 * pi
end

function pdf(d::SinDistribution, x::Real)
    insupport(d, x) ? (sin(x) + 1) / (2 * pi) : 0.0
end

function cdf(d::SinDistribution, x::Real)
    insupport(d, x) ? (x - cos(x) + 1.0) / (2 * pi) : Float64(x > 0)
end


# ---
# title: Density plots
# cover: assets/density_plots.png
# description: Visualizing kernel density estimation of data
# author: "[Pietro Vertechi](https://github.com/piever)"
# ---

using AlgebraOfGraphics, CairoMakie
using AlgebraOfGraphics: density
set_aog_theme!() #src

df = (x=randn(1000), c=rand(["a", "b"], 1000))
plt = data(df) * mapping(:x, color=:c) * density(bandwidth=0.5)
fg = draw(plt)

#

df = (x=randn(1000), c=rand(["a", "b"], 1000))
plt = data(df) * mapping(:x, color=:c) * density(bandwidth=0.5) * visual(orientation=:vertical)
"Not yet supported" # hide

# save cover image #src
mkpath("assets") #src
save("assets/density_plots.png", fg) #src


# Run code for previous steps with plotting turned off.
make_plots_orig_5 = isdefined(Main,:make_plots) ? make_plots : true
 make_plots = false
 include("neid_solar_5_calc_rvs_proj.jl")
 make_plots = make_plots_orig_5
 if make_plots
   using Plots
 end
using MultivariateStats

# Set parameters for plotting analysis
plt_order = 42
 plt_order_pix = 4500:5000

fm_perp = RvSpectML.compute_spectra_perp_doppler_shift(spectral_orders_matrix.flux,deriv_orders, ave_good_chunks_rvs )
  idx_good_obs = 1:length(ave_good_chunks_rvs)
  M = fit(PCA, fm_perp[:,idx_good_obs]; maxoutdim=10)
  pca_out = MultivariateStats.transform(M,fm_perp[:,idx_good_obs])
  frac_var_explained = 1.0.-cumsum(principalvars(M))./tvar(M)
  if make_plots
    plt = scatter(frac_var_explained, xlabel="Number of PCs", ylabel="Frac Variance Unexplained")
    display(plt)
  end
  frac_var_explained

if make_plots
  plt_order = 40
  # plt_order_pix = 3301:3800
  RvSpectML.plot_basis_vectors(order_grids, f_mean_orders, deriv_orders, M.proj, idx_plt = spectral_orders_matrix.chunk_map[plt_order][plt_order_pix] )
end

if make_plots
  RvSpectML.plot_basis_scores(plt_times, ave_good_chunks_rvs, pca_out )
end

if make_plots
  RvSpectML.plot_basis_scores_cor( ave_good_chunks_rvs, pca_out)
end


using Test, LinearAlgebra
using WaveProp
using WaveProp.Nystrom
using WaveProp.Geometry
using WaveProp.Integration
using WaveProp.Mesh

@testset "Basic tests" begin
    Geometry.clear!()
    Ω,M   = WaveProp.IO.gmsh_sphere(dim=2,h=0.05)
    Γ     = boundary(Ω)
    mesh  = view(M,boundary(Ω))
    nmesh = NystromMesh(mesh,order=2)
    @test isapprox(sum(qweights(nmesh)),π,atol=0.1)
    Geometry.clear!()
    geo   = Circle(radius=0.5)
    Ω     = Domain(geo)
    M     = meshgen(Ω;h=(0.1))
    mesh  = view(M,boundary(Ω))
    nmesh = NystromMesh(mesh,order=2)
    @test isapprox(sum(qweights(nmesh)),π,atol=0.1)
end

@testset "Mesh quadrature" begin
    @testset "Cube" begin
        # generate a mesh
        Geometry.clear!()
        (lx,ly,lz) = widths = (1.,1.,2.)
        Ω, M  = WaveProp.IO.gmsh_box(;widths=widths)
        ∂Ω = boundary(Ω)
        mesh  = NystromMesh(view(M,∂Ω),order=1)
        A     = 2*(lx*ly + lx*lz + ly*lz)
        @test A ≈ sum(qweights(mesh))
        # generate a Nystrom mesh for volume. Note that since we are creating a
        # volume Nystrom mesh, we pass the kwarg compute_normal=false since
        # normals do not make sense in this case
        mesh  = NystromMesh(view(M,Ω),order=1,compute_normal=false)
        V     = prod(widths)
        # sum only weights corresponding to tetras
        @test V ≈ sum(qweights(mesh))
        # finally generate a Nystrom mesh for the surface AND volume
        Ω₁ = union(Ω,∂Ω)
        mesh  = NystromMesh(view(M,Ω₁),order=1,compute_normal=false)
        @test V+A ≈ sum(qweights(mesh))
    end
    @testset "Sphere" begin
        r = 0.5
        Geometry.clear!()
        Ω, M = WaveProp.IO.gmsh_sphere(;radius=r)
        Γ = boundary(Ω)
        A = 4π*r^2
        mesh = NystromMesh(view(M,Γ),order=2)
        # the coarse tolerance below is because we use flat elements to
        # approximate the surface area and volume of a sphere
        @test isapprox(A,sum(qweights(mesh));atol=0.1)
        # test volume
        mesh = NystromMesh(view(M,Ω),order=2;compute_normal=false)
        A = 4π*r^2
        V = (4/3)*π*r^3
        @test isapprox(V,sum(qweights(mesh));atol=0.1)
    end
    @testset "Disk" begin
        r = rx = ry = 0.5
        Geometry.clear!()
        Ω, M = WaveProp.IO.gmsh_disk(;rx,ry)
        Γ = boundary(Ω)
        mesh = NystromMesh(view(M,Ω),order=2,compute_normal=false)
        A = π*r^2
        # the coarse tolerance below is because we use flat elements to
        # approximate the surface area and volume of a sphere
        @test isapprox(A,sum(qweights(mesh));atol=0.1)
        # test perimeter
        mesh = NystromMesh(view(M,Γ),order=2)
        P = 2π*r
        @test isapprox(P,sum(qweights(mesh));atol=0.1)
    end
    @testset "Append" begin
        Geometry.clear!()
        Ω, M = WaveProp.IO.gmsh_disk()
        Γ = boundary(Ω)
        mesh = NystromMesh(view(M,Ω),order=2,compute_normal=false)
        Σ_mesh = NystromMesh{2,Float64}()
        append!(Σ_mesh,mesh)
        append!(Σ_mesh,mesh)
    end
end


mutable struct SpanningTreeLightGraphsPrimSolver <: SpanningTreeSolver
  graph::SimpleGraph
  weights_matrix::Matrix{Float64}

  function SpanningTreeLightGraphsPrimSolver()
    return new(Graph(0), zeros(0, 0))
  end
end

has_lp_formulation(::SpanningTreeLightGraphsPrimSolver) = false
approximation_ratio(::SpanningTreeLightGraphsPrimSolver) = 1.0
approximation_term(::SpanningTreeLightGraphsPrimSolver) = 0.0
approximation_ratio_budgeted(::SpanningTreeLightGraphsPrimSolver) = 1.0
approximation_term_budgeted(::SpanningTreeLightGraphsPrimSolver) = 0.0
supports_solve_budgeted_linear(::SpanningTreeLightGraphsPrimSolver) = false
supports_solve_all_budgeted_linear(::SpanningTreeLightGraphsPrimSolver) = false

function build!(solver::SpanningTreeLightGraphsPrimSolver, graph::SimpleGraph)
  n = nv(graph)
  solver.graph = graph
  solver.weights_matrix = zeros(n, n)
end

function solve_linear(solver::SpanningTreeLightGraphsPrimSolver, reward::Dict{Tuple{Int, Int}, Float64})
  # Make up a rewards matrix by copying the input dictionary into the right data structure for the shortest path computations.
  fill!(solver.weights_matrix, 0.0)
  for (k, v) in reward
    solver.weights_matrix[k[1], k[2]] = v
  end
  solver.weights_matrix = maximum(solver.weights_matrix) .- solver.weights_matrix

  # Compute the maximum spanning tree.
  mst = prim_mst(solver.graph, solver.weights_matrix)

  # Transform edges into tuples.
  return _mst_solution_normalise(reward, [(src(e), dst(e)) for e in mst])
end


#-------------binary_trees.jl--------------------------------------------------#
#
# Purpose: This implements a basica binary tree, DFS, and a Priority Queue
#
#------------------------------------------------------------------------------#

# Now we need a typealias for the tree
abstract BT

# Type of Binary Tree
type BTree <: BT
    weight::Float64
    bitpath::String
    left::BT
    right::BT
end

# First, we need to create a type for empty nodes
type Empty <: BT end
const et = Empty()

# Depth-First Search
function DFS(node::BTree)
    println(node.bitpath)
    if !isa(node.right, Empty)
        #println(0)
        DFS(node.right)
    end

    if !isa(node.left, Empty)
        #println(1)
        DFS(node.left)
    end

    if isa(node.left, Empty) && isa(node.right, Empty)
        println(node.bitpath)
    end
end

# Function to create binary tree with a certain number of generations
function create_BTree(node::BTree, gen::Int64, max_gen::Int64, bitpath::String)

    node.weight = rand()
    node.bitpath = bitpath

    if gen < max_gen
        node.right = BTree(0, "", et, et)
        node.left = BTree(0, "", et, et)
        create_BTree(node.right, gen+1, max_gen, string(bitpath,1))
        create_BTree(node.left, gen+1, max_gen, string(bitpath,0))
    end

end

function main()

    tree = BTree(0, "", et, et,)
    create_BTree(tree, 0, 3, "")
    DFS(tree)

end

main()


@inline function parent_cell(p::NTuple{N,A}, dxi::NTuple{N,B}) where {N,A,B}
    return ntuple(i -> Int16(p[i] ÷ dxi[i]) + 1, Val(N))
end

# dimension-agnostic fully unrolled euclidean distance
@generated function distance(a::NTuple{N,T}, b::NTuple{N,T}) where {N,T}
    ex = zero(T)
    @inbounds for i in 1:N
        ex = :((a[$i] - b[$i]) * (a[$i] - b[$i]) + $ex)
    end
    return :(√($ex))
end

# check whether particle is inside the grid (includes boundary)
@inline function isinside(px::Real, py::Real, x, y)
    xmin, xmax = extrema(x)
    ymin, ymax = extrema(y)
    @assert (xmin ≤ px ≤ xmax) && (ymin ≤ py ≤ ymax)
end

@inline function isinside(px::Real, py::Real, pz::Real, x, y, z)
    xmin, xmax = extrema(x)
    ymin, ymax = extrema(y)
    zmin, zmax = extrema(z)
    @assert (xmin ≤ px ≤ xmax) && (ymin ≤ py ≤ ymax) && (zmin ≤ pz ≤ zmax)
end

@inline isinside(p::NTuple{2,T1}, x::NTuple{2,T2}) where {T1,T2} = isinside(p[1], p[2], x[1], x[2])

@inline function isinside(p::NTuple{3,T1}, x::NTuple{3,T2}) where {T1,T2}
    return isinside(p[1], p[2], p[3], x[1], x[2], x[3])
end

# normalize coordinates
@inline function normalize_coordinates(
    p::NTuple{N,A}, xi::NTuple{N,B}, dxi::NTuple{N,C}, idx::NTuple{N,D}
) where {N,A,B,C,D}
    return ntuple(i -> (p[i] - xi[i][idx[i]]) * (1 / dxi[i]), Val(N))
end

# compute grid size
function grid_size(x::NTuple{N,T}) where {T,N}
    return ntuple(i -> x[i][2] - x[i][1], Val(N))
end

# Get field F at the corners of a given cell
@inline function field_corners(F::AbstractArray{T,2}, idx::NTuple{2,Integer}) where {T}
    idx_x, idx_y = idx
    return (
        F[idx_x, idx_y], F[idx_x + 1, idx_y], F[idx_x, idx_y + 1], F[idx_x + 1, idx_y + 1]
    )
end

@inline function field_corners(F::AbstractArray{T,3}, idx::NTuple{3,Integer}) where {T}
    idx_x, idx_y, idx_z = idx
    return (
        F[idx_x, idx_y, idx_z],   # v000
        F[idx_x + 1, idx_y, idx_z],   # v100
        F[idx_x, idx_y, idx_z + 1], # v001
        F[idx_x + 1, idx_y, idx_z + 1], # v101
        F[idx_x, idx_y + 1, idx_z],   # v010
        F[idx_x + 1, idx_y + 1, idx_z],   # v110
        F[idx_x, idx_y + 1, idx_z + 1], # v011
        F[idx_x + 1, idx_y + 1, idx_z + 1], # v111
    )
end

@inline particle2tuple(p::NTuple{N,AbstractArray}, ix) where {N} = ntuple(i -> p[i][ix], Val(N))

# 2D random particle generator for regular grids
function random_particles(nxcell, x, y, dx, dy, nx, ny)
    # number of cells
    ncells = (nx - 1) * (ny - 1)
    # allocate particle coordinate arrays
    px, py = zeros(nxcell * ncells), zeros(nxcell * ncells)
    Threads.@threads for i in 1:(nx - 1)
        @inbounds for j in 1:(ny - 1)
            # lowermost-left corner of the cell
            x0, y0 = x[i], y[j]
            # cell index
            cell = i + (nx - 1) * (j - 1)
            for l in 1:nxcell
                px[(cell - 1) * nxcell + l] = rand() * dx + x0
                py[(cell - 1) * nxcell + l] = rand() * dy + y0
            end
        end
    end

    return px, py
end

# 3D random particle generator for regular grids
function random_particles(nxcell, x, y, z, dx, dy, dz, nx, ny, nz)
    # number of cells
    ncells = (nx - 1) * (ny - 1) * (nz - 1)
    # allocate particle coordinate arrays
    px, py, pz = zeros(nxcell * ncells), zeros(nxcell * ncells), zeros(nxcell * ncells)
    Threads.@threads for i in 1:(nx - 1)
        @inbounds for j in 1:(ny - 1), k in 1:(nz - 1)
            # lowermost-left corner of the cell
            x0, y0, z0 = x[i], y[j], z[k]
            # cell index
            cell = i + (nx - 1) * (j - 1) + (nx - 1) * (ny - 1) * (k - 1)
            for l in 1:nxcell
                px[(cell - 1) * nxcell + l] = rand() * dx + x0
                py[(cell - 1) * nxcell + l] = rand() * dy + y0
                pz[(cell - 1) * nxcell + l] = rand() * dz + z0
            end
        end
    end

    return px, py, pz
end


# Paul Rodriguez translation into Julia 0.3.5 of Monica Costa Dias and Corma O'Dea "Dynamic Economics in Practice"
# using Alex code for Getting Equiprobable Shock Points from a Normal Distributions
# March 2015, p.lesmes.11@ucl.ac.uk
# Please improve it!

# ------------------------------------------------------------------------
# DESCRIPTION
# This program solves and simulates a finite period consumption and saving
# problem. There is income that can be uncertain. The income program can be
# hardcoded by the user or can be set to follow a log normal autoregessive
# process

# IMPORTANT: Comment lines 188 and 189 if you want to change the number of
# periods or use shocks generated in Julia

# ------------------------------------------------------------------------
# PREAMBLE
# Ensure that all storage spaces variables, matrices, memory, globals are
# clean from information stored in past runs

using Distributions   # For drwaing distributions
using Optim           # Basic optimization routines
using Grid            # Interpolation tools

# Set here the relevant directory. Please check
cd("C://Users//Paul//Documents//GitHub//ucl-econ-julia//LifecycleCostaDias//v5_julia")

# ----------------------------------------------------------------------

mainDir=pwd()

# Numeric tools
cd("code_numericTools")
require("getEquiprobNormalDeviates.jl")
require("truncate.jl")
require("interp1lin.jl")
require("interp2D.jl")
cd(mainDir)

# Grids and other inputs
cd("code_gridsInputs")
require("checkInputs.jl")
require("getIncomeGrid.jl")
require("getMinandMaxAss.jl")
require("getGrid.jl")
cd(mainDir)

# Programs for the solution
cd("code_solution")
require("solveValueFunction.jl")
require("objectivefunc.jl")
require("utility.jl")
cd(mainDir)

# Programs for the simulation
cd("code_simulation")
require("simNoUncer.jl")
require("simWithUncer.jl")
cd(mainDir)

# Plot stuff
cd("code_plots")
require("plotPaths.jl")
require("plots.jl")
cd(mainDir)

tic()        # start the clock

# ------------------------------------------------------------------------
# DECLARE VARIABLES AND MATRICES THAT WILL BE 'GLOBAL'
# explicitly set variables and matrices to be shared throughout the routine
# as globals

#% ------------------------------------------------------------------------
# NUMERICAL METHODS
# select solution, interpolation and integration methods

#global interpMethod = 'pchip';      # interpolation method - IN JULIA, UP TO THE MOMENT, YOU CAN ONLY USE LINEAR WITH UNEVEN GRIDS
const linearise = 1;               # whether to linearise the slope of EV when using EE - set linearise=1 to do this, else = 0

# ------------------------------------------------------------------------
# NUMERICAL CONSTANTS
# set constants needed in numerical solution and simulation

# precision parameters
#--------------------------------------%
const tol = 1e-10;                 # max allowed error
const minCons = 1e-5;              # min allowed consumption

# where to truncate the normal distributions
#--------------------------------------%
const normBnd = 3;                 #Ignore draws less than -NormalTunc*sigma and greater than normalTrunc*sigma

# information for simulations
#--------------------------------------%
const numSims = 2;                #How many individuals to simulate

# ------------------------------------------------------------------------
# THE ECONOMIC ENVIRONMENT
# Set values of structural economic parameters

const T = 40;                      # Number of time period (If you want to reproduce matlab results, be sure that it has the same number of periods!)
const r = 0.01;                    # Interest rate
const beta = 0.98;                 # Discount factor
const gamma = 1.5;                 # Coefficient of relative risk aversion
const mu = 0;                      # mean of initial log income
const sigma = 0.25;                   # variance of log income
const rho = 0.75;                     # persistency of log income
const Tretire = 41;                # age after which there is no income earned
const borrowingAllowed = 0;        # Is borrowing allowed
const isUncertainty = 1;           # Is income uncertain?
const startA = 0;                  # How much asset do people start life with

# ------------------------------------------------------------------------
# GRIDS
# choose dimensions, set matrices and select methods to construct grids

#The grid for assets
#--------------------------------------%
const numPointsA = 20;             # number of points in the discretised asset grid
gridMethod = "3logsteps";    # method to construct grid. One of equalsteps, logsteps, 3logsteps, 5logsteps or 10logsteps

#The grid for income shocks
#--------------------------------------%
const numPointsY = 5;           #  points in grid for income (should be 2 if hard-coded)
const uncertaintyMethod = 1;    #  =0 if we enter shocks manually, =1 if put shocks on a grid using Tauchen (1986) method
hcIncome = [0.5941 ,   0.8179  ,  1.0000  ,  1.2227  ,  1.6832];  # hard-coded shocks, used if uncertaintyMethod = 1
hcIncPDF = [0.1016 ,   0.2492  ,  0.2983  ,  0.2492  ,  0.1017];  # and respective probabilities

# Check inputs
checkInputs()

## Get income grid
(Ygrid, incTransitionMrx, minInc, maxInc) = getIncomeGrid();

## ------------------------------------------------------------------------
# GET ASSET GRID
# populate grid for assets using 'gridMethod'
# populate matrix borrowingConstraints

(borrowCon, maxAss) = getMinAndMaxAss(borrowingAllowed, minInc, maxInc, startA);

Agrid = Array(Float64,T+1, numPointsA);
for ixt = 1:1:T+1
    Agrid[ixt, :] = getGrid(borrowCon[ixt], maxAss[ixt], numPointsA, gridMethod);
end

#% ------------------------------------------------------------------------
# SOLVE CONSUMER'S PROBLEM
# Get policy function and value function

(policyA1, policyC, val, exVal) = solveValueFunction();

	# Want to compare with Matlab version? In that way you can test how different are
	# simulations for the same policyFunc in Matlab and Julia
	# policyA1=reshape(readdlm("matlablObj\\policyA1matlab.csv", ',' ),40,20,5);

#% ------------------------------------------------------------------------
# SIMULATE CONSUMER'S PATHS
# start from initial level of assets and simulate optimal consumption and
# savings profiles over lifecycle

if isUncertainty == 0
    (ypath, cpath, apath, vpath) = simNoUncer(policyA1, exVal, startA);
else
	# Get the draws, and the feed the simulation with them (sligthly different than the original code)

	# For the discreate shocks of Income case... Julia makes our life way easier!
	srand(1223424) # Seed it!
	randYpath=rand( Categorical(vec(hcIncPDF))  ,(T,numSims));

	# For the normal distributed case
	srand(1223424) # Seed it!
 	e=rand(Normal(0, sigma),(T,numSims)); # normally distributed random draws for the innovation
 	srand(234636)  # Seed it!
 	sig_inc = sigma/ ((1-rho^2)^0.5);
 	logy1   =rand(Normal(mu, sig_inc),numSims);  # a random draw for the initial income

		# IMPORTANT: Want to compare with Matlab version? In that way you can test how different are
		# simulations using the same shocks. NOTICE THAT THIS OVERIDES THE GENERATED SHOCKS FROM
		# THE PREVIOUS LINES
		# Comment this lines if you want to change the number of periods
 		e      = readdlm("matlablObj\\ematlab.csv", ',' );
 		logy1  = readdlm("matlablObj\\logy1matlab.csv", ',' );

    (ypath, cpath, apath, vpath) = simWithUncer(policyA1,exVal, startA,e,logy1,randYpath);
end

toc();     # Stop the clock

#% ------------------------------------------------------------------------
# PLOTS
# Plots some features of the simulation and simulation
# Use a program that supports Gadfly, in these days, either JUNO, IJULIA

if 1==1  # In purspouse, Gadfly is just so slow!! But there are few good alternatives (March 2015)

	# Slow.........
	using Gadfly
	using DataFrames      # For making graphs or exporting data, is better as dataframes

	plotPaths()

	# Now plot value and policy functions
	whichyear = 20;
	plotNode1 = 3;
	plotNodeLast = numPointsA;

	plots();   #Plot Value Functions
end

# ------------------------------------------------------------------------
# ------------------------------------------------------------------------


module mbDTW
export dtwDist, dtwPath

using mbUtil

function dtwMatrix(A,B,d=(x,y)->norm(x-y))
	m,n = length(A),length(B)
	M = Array{Float64}(m,n)

	M[1,1] = d(A[1],B[1])
	for i in 2:m M[i,1] = M[i-1,1] + d(A[i],B[1]) end
	for j in 2:n M[1,j] = M[1,j-1] + d(A[1],B[j]) end

	for i in 2:m, j in 2:n
		M[i,j] = d(A[i],B[j]) + min(M[i-1, j-1], M[i-1, j], M[i, j-1])
	end

	return M
end

dtwDist = last ∘ dtwMatrix

function dtwPath(A,B)
	M = dtwMatrix(A,B)
	m,n = size(M)

	P = [(m,n)]

	function locmin(M,i,j)
		i≤1 && return (i,j-1)
		j≤1 && return (i-1,j)
		d = [(i-1, j-1), (i, j-1), (i-1, j)]
		rtn = d[indmin(map(i->M[i...], d))]
		return rtn
	end

	while P[1] ≠ (1,1)
		unshift!(P, locmin(M, P[1]...))
	end

	return P
end

end #module


# Tutorial 2.4.2. Closed systems
# ==============================
#
# Physics background
# ------------------
#  Fock-darwin spectrum of a quantum dot (energy spectrum in
#  as a function of a magnetic field)
#
# Kwant features highlighted
# --------------------------
#  - Use of `hamiltonian_submatrix` in order to obtain a Hamiltonian
#    matrix.

# For eigenvalue computation
sla = pyimport("scipy.sparse.linalg")

function make_system(;a=1, t=1.0, r=10)
    # Start with an empty tight-binding system and a single square lattice.
    # `a` is the lattice constant (by default set to 1 for simplicity).

    lat = kwant.lattice.square(a, norbs=1)

    syst = kwant.Builder()

    # Define the quantum dot
    # def circle(pos):
    #     (x, y) = pos
    #     rsq = x ** 2 + y ** 2
    #     return rsq < r ** 2
    function circle(pos)
        (x, y) = pos
        R = hypot(x,y)
        return R < r
    end

    py"""
    def hopx(site1, site2, B):
        # The magnetic field is controlled by the parameter B
        y = site1.pos[1]
        return -$t * $exp(-1j * B * y)
    """
    # function hopx(site1,site2,B)
    #     y = site1.pos[2]
    #     return -t * exp(-1im*B*y)
    # end

    syst[lat.shape(circle, (0, 0))] = 4 * t
    # hoppings in x-direction
    syst[kwant.builder.HoppingKind((1, 0), lat, lat)] = py"hopx"
    # hoppings in y-directions
    syst[kwant.builder.HoppingKind((0, 1), lat, lat)] = -t

    # It's a closed system for a change, so no leads
    return syst
end

function plot_spectrum(syst, Bfields)

    energies = Array{Float64}(undef,15,0)
    for B in Bfields
        # Obtain the Hamiltonian as a sparse matrix
        ham_mat = syst.hamiltonian_submatrix(params=Dict(:B=>B), sparse=true)

        # we only calculate the 15 lowest eigenvalues
        ev = sla.eigsh(ham_mat.tocsc(), k=15, sigma=0,
                       return_eigenvectors=false)

        energies = hcat(energies,ev)
    end

    pyplot.figure()
    pyplot.plot(Bfields, energies')
    pyplot.xlabel("magnetic field [arbitrary units]")
    pyplot.ylabel("energy [t]")
    pyplot.gcf()
end

function sorted_eigs(ev)
    evals, evecs = ev
    perm = sortperm(evals)
    return evals[perm],evecs[:,perm]
    # evals, evecs = map(np.array, zip(*sorted(zip(evals, evecs.transpose()))))
    # return evals, evecs.transpose()
end

function plot_wave_function(syst, B=0.001)
    # Calculate the wave functions in the system.
    ham_mat = syst.hamiltonian_submatrix(sparse=true, params=Dict(:B=>B))
    evals, evecs = sorted_eigs(sla.eigsh(ham_mat.tocsc(), k=20, sigma=0))

    # Plot the probability density of the 10th eigenmode.
    kwant.plotter.map(syst, abs2.(evecs[:, 9]),
                      colorbar=false, oversampling=1)
end

function plot_current(syst, B=0.001)
    # Calculate the wave functions in the system.
    ham_mat = syst.hamiltonian_submatrix(sparse=true, params=Dict(:B=>B))
    evals, evecs = sorted_eigs(sla.eigsh(ham_mat.tocsc(), k=20, sigma=0))

    # Calculate and plot the local current of the 10th eigenmode.
    J = kwant.operator.Current(syst)
    current = J(evecs[:, 9], params=Dict(:B=>B))
    kwant.plotter.current(syst, current, colorbar=false)
end


@testset "7.3 Inverse hyperbolic tangent" begin
    include("7.3.2 (d x)^m (a+b arctanh(c x^n))^p.jl")
    include("7.3.3 (d+e x)^m (a+b arctanh(c x^n))^p.jl")
    include("7.3.4 u (a+b arctanh(c x))^p.jl")
    include("7.3.5 u (a+b arctanh(c+d x))^p.jl")
    include("7.3.6 Exponentials of inverse hyperbolic tangent functions.jl")
    include("7.3.7 Inverse hyperbolic tangent functions.jl")
end

test/7 Inverse hyperbolic functions/7.3 Inverse hyperbolic tangent/.jl


module Blobs

export Blob

using ..Elements
import ..Motions: induce_velocity, mutually_induce_velocity!, self_induce_velocity!, advect

#== Type definition ==#

"""
    Blob <: Elements.Element

An immutable structure representing a regularized point source/vortex

## Fields
- `z`: position
- `S`: strength/circulation
- `δ`: regularization radius

## Constructors

"""
struct Blob{T <: Number} <: Element
    z::ComplexF64
    S::T
    δ::Float64
    Blob{T}(z, s::Real, δ) where T <: Complex = new(z, im*s, δ)
    Blob{T}(z, s, δ)       where T <: Complex = new(z, s, δ)
    Blob{T}(z, s::Real, δ) where T <: Real    = new(z, s, δ)
end
Elements.kind(::Blob) = Singleton
Elements.kind(::Type{Blob{T}}) where T = Singleton

(b::Blob{T})(; z = b.z, S = b.S, δ = b.δ) where T = Blob{T}(z, S, δ)

#== Methods to be extended ==#

Elements.position(b::Blob) = b.z
Elements.streamfunction(z::ComplexF64, b::Blob) = real(-0.5b.S*log(z - b.z)/π)

Elements.complexpotential(z::ComplexF64, b::Blob) = -0.5im*b.S*log(z - b.z)/π

blob_kernel(z, δ) = 0.5im*z/(π*(abs2(z) + δ^2))

function induce_velocity(z::ComplexF64, b::Blob, t)
    b.S*blob_kernel(z - b.z, b.δ)
end

function induce_velocity(target::Blob, source::Blob, t)
    δ = √(0.5(target.δ^2 + source.δ^2))
    source.S'*blob_kernel(target.z - source.z, δ)
end

function mutually_induce_velocity!(ws₁, ws₂,
                                   blobs₁::Vector{Blob{T₁}},
                                   blobs₂::Vector{Blob{T₂}}, t) where {T₁, T₂}
    for (s, source) in enumerate(blobs₁)
        for (t, target) in enumerate(blobs₂)
            δ = √(0.5(target.δ^2 + source.δ^2))
            K = blob_kernel(target.z - source.z, δ)
            ws₂[t] += source.S'*K
            ws₁[s] -= target.S'*K
        end
    end
    nothing
end

function self_induce_velocity!(ws, blobs::Vector{Blob{T}}, t) where T
    N = length(blobs)

    for s in 1:N, t in s+1:N
        δ = √(0.5(blobs[t].δ^2 + blobs[s].δ^2))
        K = blob_kernel(blobs[t].z -blobs[s].z, δ)
        ws[t] += blobs[s].S'*K
        ws[s] -= blobs[t].S'*K
    end
    ws
end

function advect(b::Blob{T}, w::ComplexF64, Δt::Float64) where T
    Blob{T}(b.z + w*Δt, b.S, b.δ)
end
#
#function Base.show(io::IO, b::Blob)
#    print(io, "Vortex Blob: z = $(round(b.z, digits=3)), Γ = $(round(b.Γ, digits=3)), δ = $(round(b.δ, digits=3))")
#end

end


@testset "shortest path" begin
    using STMO
    using STMO.ShortestPath

    graph = Dict('A' => [(0.5, 'B'), (1.5, 'H')],
        'B' => [(1.0, 'C')],
        'C' => [(1.0, 'D'), (1.0, 'E')],
        'D' => [(2.0, 'F')],
        'E' => [(1.0, 'G')],
        'F' => [(0.1, 'E')],
        'G' => [(1.0, 'F')],
        'H' => [(0.9, 'E')])

    distances, previous = dijkstra(graph, 'A')

    @test distances['D'] ≈ 2.5
    @test previous['E'] == 'H'

    path = reconstruct_path(previous, 'A', 'H')
    @test first(path) == 'A'
    @test last(path) == 'H'

    path_d, cost_AH = dijkstra(graph, 'A', 'H')
    @test path_d == path
    @test cost_AH ≈ distances['H']

    path_as, cost_as = a_star(graph, 'A', 'H', x -> 0.0)
    @test path_as == path
    @test cost_as ≈ distances['H']

end


using Test
using TrigPolys

@testset "Arithmetic" begin
    p1 = random_trig_poly(102)
    p2 = random_trig_poly(201)
    x = range(0; stop=2*pi, length=2000)

    @test isapprox((-p1).(x), -(p1.(x)))
    @test isapprox((p1 + p2).(x), p1.(x) .+ p2.(x))
    @test isapprox((p1 - p2).(x), p1.(x) .- p2.(x))
    @test isapprox((p1 * p2).(x), p1.(x) .* p2.(x))

    @test TrigPolys.a0(p1 + 3) == TrigPolys.a0(p1) + 3
    @test TrigPoly(pi) == TrigPoly(pi)
    @test p1 + p2 == p2 + p1
    @test p1 * p2 == p2 * p1
    @test p1 + pi == pi + p1
    @test p1 * pi == pi * p1
    @test p1 / pi == p1 * (1/pi)
end

@testset "Type Conversion and Promotion" begin
    F1 = Float32
    F2 = Float64
    VF1 = Vector{F1}
    VF2 = Vector{F2}
    @test typeof(TrigPoly(F2(4))+F1(3)) == TrigPoly{F2, VF2}
    @test typeof(convert(TrigPoly{F2, VF2}, F1(.3))) == TrigPoly{F2, VF2}
    @test typeof(convert(TrigPoly{F1, VF1}, F2(.3))) == TrigPoly{F1, VF1}
    @test typeof(TrigPoly(F1(2))+TrigPoly(F2(3))) == TrigPoly{F2, VF2}
    @test typeof(TrigPoly(F1(2))+F1(3)) == TrigPoly{F1, VF1}
end

@testset "Evaluate and interpolate" begin
    n = 1000
    p = random_trig_poly(n)
    s = p.n*2+1
    samples = [2*pi*(i-1)/s for i in 1:s]
    B = hcat([TrigPolys.basis(p.n, xi) for xi in samples]...)

    u = randn(s)
    @test isapprox(evaluate(u), B'*u)
    @test isapprox(evaluate(p), p.(samples))
    @test isapprox(evaluateT(vec(u)), B*u)
    @test isapprox(interpolatev(evaluate(u)), u)
end


# brutal force caculations for accuracy benchmark
# partition function
function partition_function_exact(
    Ns::Int64, N::Int64, expβϵ::Array{T,1}
) where {T<:FloatType}

    config_list = config_list_generator(Ns::Int64, N::Int64)
    Z = 0

    for i in config_list
        Z += prod(expβϵ[findall(!iszero, expβϵ .* i)])
    end

    Z

end

# mean-occupation number
function occupancy_exact(
    Ns::Int64, N::Int64, expβϵ::Array{T,1}
) where {T<:FloatType}

    config_list = config_list_generator(Ns::Int64, N::Int64)
    Z = 0

    for i in config_list
        Z += prod(expβϵ[findall(!iszero, expβϵ .* i)])
    end

    Z = real(Z)
    n̄ = complex(zeros(Ns))

    for i = 1 : Ns
        for j in config_list
            if j[i] > 0
                n̄[i] += prod(expβϵ[findall(!iszero, expβϵ .* j)]) / Z
            end
        end
    end

    Z, n̄

end

# use deep-first search to generate all fermion configurations
function config_list_generator(Ns::Int64, N::Int64)

    config_Array = zeros(Int64, Ns)
    config_list = Array{Int64,1}[]
    DFS(1, Ns, N, config_Array, config_list)

    config_list
end

function DFS(site_index, Ns, N, config_Array, config_list)

    if N == 0 || site_index == Ns
        if N <= 1
            config_Array[site_index : Ns] .= N
            push!(config_list, copy(config_Array))
            return
        else
            return
        end
    end

    for i = 0 : 1
        config_Array[site_index] = i
        DFS(site_index + 1, Ns, N - i, config_Array, config_list)
    end

end

# Projection Method
function particle_projection(
    Ns::Int64, N::Int64, expβϵ::Array{T,1}, FLAG_OccNum::Bool
) where {T<:FloatType}

    φ = [2 * π * m / Ns for m = 1 : Ns]
    expβμ = (abs(expβϵ[Ns - N + 1]) +
            abs(expβϵ[Ns - N])) / 2     # chemical potential

    η = [prod(1 .+ exp(im * φ[m]) / expβμ * expβϵ) for m = 1 : Ns]
    Z = sum(exp.(-im * φ) .^ N .* η) / Ns * expβμ ^ N

    # skip the occupation number projection
    if FLAG_OccNum == false
        return real(Z)
    end

    n̄ = [sum(exp.(im * φ) / expβμ * expβϵ[i] ./ (1 .+ exp.(im * φ) / expβμ * expβϵ[i]) .*
    exp.(-im * φ) .^ N .* η) * expβμ ^ N / Ns / Z for i = 1 : Ns]

    return real(Z), n̄

end

function test_matrix_product(σlist::Array{Int64,2}, lattice::Lattice, qmc::QMC)

    MatProdUp = Matrix{Float64}(I, lattice.Ns, lattice.Ns)
    MatProdDn = Matrix{Float64}(I, lattice.Ns, lattice.Ns)

    for i = 1 : size(σlist)[2]
        σ = σlist[:, i]
        Bj_Up, Bj_Dn = singlestep_matrix(σ, lattice, qmc)
        MatProdUp = Bj_Up * MatProdUp
        MatProdDn = Bj_Dn * MatProdDn
    end

    MatProdUp, MatProdDn

end

function test_matrix_product_svd(
    A::Array{Float64, 2}, Ns::Int64, L::Int64
    )

    num_cluster_size = 10

    eigvalsA = eigvals(A)

    MatProd_U = Matrix{Float64}(I, Ns, Ns)
    MatProd_D = Matrix{Float64}(I, Ns, Ns)
    MatProd_V = Matrix{Float64}(I, Ns, Ns)

    for i = 1 : div(L, num_cluster_size)

        MatProd = Matrix{Float64}(I, Ns, Ns)

        for j = 1 : num_cluster_size
            MatProd = A * MatProd
        end

        BUD = MatProd * MatProd_U * MatProd_D
        SVD_BUD = LinearAlgebra.LAPACK.gesvd!('A', 'A', BUD)
        MatProd_U = SVD_BUD[1]
        MatProd_D = diagm(0 => SVD_BUD[2])
        MatProd_V = SVD_BUD[3] * MatProd_V

    end

    eigs_prop = eigvals(MatProd_D * MatProd_V * MatProd_U; permute = true, scale = true, sortby = abs)
    #eigs_prop = LinearAlgebra.LAPACK.geevx!('B', 'V', 'V', 'N', MatProd_D * MatProd_V * MatProd_U)

    eigvalsA .^ L, eigs_prop
    #MatProd_D

end

function test_matrix_product_qr(
    A::Array{Float64, 2}, Ns::Int64, L::Int64
    )

    num_cluster_size = 10

    eigvalsA = eigvals(A)

    MatProd_Q = Matrix{Float64}(I, Ns, Ns)
    MatProd_D = Matrix{Float64}(I, Ns, Ns)
    MatProd_T = Matrix{Float64}(I, Ns, Ns)

    for i = 1 : div(L, num_cluster_size)

        MatProd = Matrix{Float64}(I, Ns, Ns)

        for j = 1 : num_cluster_size
            MatProd = A * MatProd
        end

        QRCP_update!(MatProd_Q, MatProd_D, MatProd_T, MatProd, 'R')

    end

    eigs_prop = eigvals(MatProd_D * MatProd_T * MatProd_Q; permute = true, scale = true, sortby = abs)

    eigvalsA .^ L, eigs_prop

end

function measure_particlenumber_GC(
    lattice::Lattice, UUp_Array::Array{T1,3}, UDn_Array::Array{T2,3},
    expβμ::Float64, expβϵ_up::Array{T3,1}, expβϵ_dn::Array{T4,1}
) where {T1<:FloatType, T2<:FloatType, T3<:FloatType, T4<:FloatType}

    nup, ndn = 0, 0
    n̄_up_gc = expβμ * expβϵ_up ./ (1 .+ expβμ * expβϵ_up)
    n̄_dn_gc = expβμ * expβϵ_dn ./ (1 .+ expβμ * expβϵ_dn)

    for i = 1 : lattice.Ns
        nup += sum(UUp_Array[:, i, i] .* n̄_up_gc)
        ndn += sum(UDn_Array[:, i, i] .* n̄_dn_gc)
    end

    real(nup), real(ndn)

end

function measure_energy_GC(
    lattice::Lattice, UUp_Array::Array{T1,3}, UDn_Array::Array{T2,3},
    expβμ::Float64, expβϵ_up::Array{T3,1}, expβϵ_dn::Array{T4,1}
) where {T1<:FloatType, T2<:FloatType, T3<:FloatType, T4<:FloatType}

    EOneBody, ETwoBody = 0, 0
    n̄_up_gc = expβμ * expβϵ_up ./ (1 .+ expβμ * expβϵ_up)
    n̄_dn_gc = expβμ * expβϵ_dn ./ (1 .+ expβμ * expβϵ_dn)

    for i in eachindex(view(lattice.KMatrix, 1 : lattice.Ns, 1 : lattice.Ns))
        if lattice.KMatrix[i] != 0
            Dij_up = sum(UUp_Array[:, i[1], i[2]] .* n̄_up_gc)
            Dij_dn = sum(UDn_Array[:, i[1], i[2]] .*n̄_dn_gc)
            EOneBody += -lattice.t * (Dij_up + Dij_dn)
        end
    end

    for i = 1 : lattice.Ns
        Umatrix = UUp_Array[:, i, i] * UDn_Array[:, i, i]'
        nmatrix = n̄_up_gc * n̄_dn_gc'
        Diiii = sum(Umatrix .* nmatrix)
        ETwoBody += lattice.U * Diiii
    end

    real(EOneBody), real(ETwoBody)
end

function mc_integration_test(lattice::Lattice, qmc::QMC)

    io = open("GCETest", "a+")
    for n = 1 : qmc.sample_size
        # initialize a random configuration
        σlist = 2 * (rand(Float64, (lattice.Ns, qmc.L)) .< 0.5) .- 1
        # propagate and diagonalize
        EigvalsUp, EigvecsUp, EigvalsDn, EigvecsDn = matrix_product_stablized(σlist, lattice, qmc)
        # GCE partition function
        expβμ = exp(qmc.Δτ * qmc.L * qmc.μ)
        ZGCUp = prod(1 .+ expβμ * EigvalsUp)
        ZGCDn = prod(1 .+ expβμ * EigvalsDn)
        Sgn_GC = sign(ZGCUp) * sign(ZGCDn)
        # CE partition function
        ZCUp, n̄up = ProjectionPF(lattice.Ns, lattice.nup, EigvalsUp)
        ZCDn, n̄dn = ProjectionPF(lattice.Ns, lattice.ndn, EigvalsDn)
        Sgn_C = sign(ZCUp) * sign(ZCDn)

        if Sgn_GC * Sgn_C == -1
            println(io, "GCE Sign: ", Sgn_GC)
            println(io, σlist)
        end

    end

    close(io)

end

function mc_integration_GC(lattice::Lattice, qmc::QMC)
    # initialization
    nup_Array = Vector{Float64}()   # particle number (spin up)
    ndn_Array = Vector{Float64}()   # particle number (spin down)
    #Ek_Array = Vector{Float64}()    # kinetic energy
    #Ep_Array = Vector{Float64}()    # potential energy
    #E_Array = Vector{Float64}()     # total energy
    Sgn_Array = Vector{Int64}()     # sign
    # initialize a random configuration
    σlist = 2 * (rand(Float64, (lattice.Ns, qmc.L)) .< 0.5) .- 1
    # propagate and diagonalize
    EigvalsUp, EigvecsUp, EigvalsDn, EigvecsDn = matrix_product_stablized(σlist, lattice, qmc)
    # GC partition function
    expβμ = exp(qmc.Δτ * qmc.L * qmc.μ)
    ZGCUp = prod(1 .+ expβμ * EigvalsUp)
    ZGCDn = prod(1 .+ expβμ * EigvalsDn)
    Sgn = sign(ZGCUp) * sign(ZGCDn)

    for n = 1 : qmc.sample_size

        σlist_New = copy(σlist)

        # random flip a spin / new configuration
        flip = rand(1 : qmc.L * lattice.Ns)
        σlist_New[flip] *= -1
        EigvalsUp_New, EigvecsUp_New, EigvalsDn_New, EigvecsDn_New =
                    matrix_product_stablized(σlist_New, lattice, qmc)

        ZGCUp_New = prod(1 .+ expβμ * EigvalsUp_New)
        ZGCDn_New = prod(1 .+ expβμ * EigvalsDn_New)
        Sgn_New = sign(ZGCUp_New) * sign(ZGCDn_New)
        # caculate the acceptence ratio
        R = abs((ZGCUp_New / ZGCUp) * (ZGCDn_New / ZGCDn))
        AcceptRatio = R / (1 + R)

        if rand() <= AcceptRatio   # accept
            # adopt the new configuration
            σlist = σlist_New
            ZGCUp, ZGCDn = ZGCUp_New, ZGCDn_New
            EigvalsUp, EigvecsUp = EigvalsUp_New, EigvecsUp_New
            EigvalsDn, EigvecsDn = EigvalsDn_New, EigvecsDn_New
            Sgn = Sgn_New

            # measurement
            UUp_Array, UDn_Array = overlap_coefficient(lattice.Ns, EigvecsUp, EigvecsDn)
            nup, ndn =  measure_particlenumber_GC(lattice, UUp_Array, UDn_Array,
                                                #expβμ, EigvalsUp, EigvalsDn)
            #Ek, Ep = measure_energy_GC(lattice, UUp_Array, UDn_Array,
                                    expβμ, EigvalsUp, EigvalsDn)

            # store data
            push!(Sgn_Array, Sgn)
            push!(nup_Array, nup * Sgn)
            push!(ndn_Array, ndn * Sgn)
            #push!(Ek_Array, Ek * Sgn)
            #push!(Ep_Array, Ep * Sgn)
            #push!(E_Array, Ek * Sgn + Ep * Sgn)
        else                        # reject
            # measurement
            UUp_Array, UDn_Array = overlap_coefficient(lattice.Ns, EigvecsUp, EigvecsDn)
            nup, ndn =  measure_particlenumber_GC(lattice, UUp_Array, UDn_Array,
                                            #expβμ, EigvalsUp, EigvalsDn)
            #Ek, Ep = measure_energy_GC(lattice, UUp_Array, UDn_Array,
                                    expβμ, EigvalsUp, EigvalsDn)

            # store data
            push!(Sgn_Array, Sgn)
            push!(nup_Array, nup * Sgn)
            push!(ndn_Array, ndn * Sgn)
            #push!(Ek_Array, Ek * Sgn)
            #push!(Ep_Array, Ep * Sgn)
            #push!(E_Array, Ek * Sgn + Ep * Sgn)
        end
    end

    nup_Mean = mean(nup_Array[qmc.BurnIn : qmc.sample_size])
    nup_Var = varm(nup_Array[qmc.BurnIn : qmc.sample_size], nup_Mean) / qmc.sample_size
    ndn_Mean = mean(ndn_Array[qmc.BurnIn : qmc.sample_size])
    ndn_Var = varm(ndn_Array[qmc.BurnIn : qmc.sample_size], ndn_Mean) / qmc.sample_size

    #EMean = mean(E_Array[qmc.BurnIn : qmc.sample_size])
    #EVar = varm(E_Array[qmc.BurnIn : qmc.sample_size], EMean) / qmc.sample_size
    #EkMean = mean(Ek_Array[qmc.BurnIn : qmc.sample_size])
    #EkVar = varm(Ek_Array[qmc.BurnIn : qmc.sample_size], EkMean) / qmc.sample_size
    #EpMean = mean(Ep_Array[qmc.BurnIn : qmc.sample_size])
    #EpVar = varm(Ep_Array[qmc.BurnIn : qmc.sample_size], EpMean) / qmc.sample_size

    avgSgn = mean(Sgn_Array[qmc.BurnIn : qmc.sample_size])

    nup_Mean, nup_Var, ndn_Mean, ndn_Var, avgSgn
    #return EMean, EVar, EkMean, EkVar, EpMean, EpVar, avgSgn
end

function main_gc()
    lattice, qmc = GetInputs()

    BlockNUp_Array = Vector{Float64}()
    BlockNDn_Array = Vector{Float64}()
    #BlockEk_Array = Vector{Float64}()
    #BlockEp_Array = Vector{Float64}()
    #BlockE_Array = Vector{Float64}()
    #BlockEkVar_Array = Vector{Float64}()
    #BlockEpVar_Array = Vector{Float64}()
    #BlockEVar_Array = Vector{Float64}()
    BlockSgn_Array = Vector{Float64}()

    filename = string("GCE_output_mu", qmc.μ, "_U", lattice.U, "_beta", qmc.Δτ * qmc.L)
    outputfile = open(filename, "a+")

    @printf(outputfile, "       Grand Canonical Ensemble AFQMC for 2D Hubbard Model     \n\n")
    @printf(outputfile, "mu: %.4f\n", qmc.μ)
    @printf(outputfile, "Number of Blocks: %d\n", qmc.blocks)
    @printf(outputfile, "Length of Img Timestep: %.4f\n", qmc.Δτ)
    @printf(outputfile, "Stablization Interval: %d\n\n", qmc.cluster_size)

    for i = 1 : qmc.blocks
        BlockNUp, BlockNUpVar, BlockNDn, BlockNDnVar, avgSgn = mc_integration_GC(lattice, qmc)
        BlockNUp = BlockNUp / avgSgn
        BlockNDn = BlockNDn / avgSgn

        @printf(outputfile, "Block # %d out of %d\n", i, qmc.blocks)
        @printf(outputfile, "--------------------\n")
        @printf(outputfile, "Average Sign: %.3f\n", avgSgn)
        @printf(outputfile, "NUp            NDn            NTot\n")
        @printf(outputfile, "%.5f     %.5f       %.5f\n", BlockNUp, BlockNDn, BlockNUp + BlockNDn)
        @printf(outputfile, "Error:\n")
        @printf(outputfile, "%.5f     %.5f       %.5f\n\n", sqrt(BlockNUpVar), sqrt(BlockNDnVar), sqrt(BlockNUpVar) + sqrt(BlockNDnVar))

        push!(BlockNUp_Array, BlockNUp)
        push!(BlockNDn_Array, BlockNDn)
        push!(BlockSgn_Array, avgSgn)

        #BlockE, BlockEVar, BlockEk, BlockEkVar, BlockEp, BlockEpVar, avgSgn = mc_integration_GC(lattice, qmc)
        #BlockE = BlockE / avgSgn
        #BlockEk = BlockEk / avgSgn
        #BlockEp = BlockEp / avgSgn

        #@printf(outputfile, "Block # %d out of %d\n", i, qmc.blocks)
        #@printf(outputfile, "--------------------\n")
        #@printf(outputfile, "Average Sign: %.3f\n", avgSgn)
        #@printf(outputfile, "Ek            Ep            Etot\n")
        #@printf(outputfile, "%.5f     %.5f       %.5f\n", BlockEk, BlockEp, BlockE)
        #@printf(outputfile, "Error:\n")
        #@printf(outputfile, "%.5f     %.5f       %.5f\n\n", sqrt(BlockEkVar), sqrt(BlockEpVar), sqrt(BlockEVar))

        #push!(BlockEk_Array, BlockEk)
        #push!(BlockEkVar_Array, BlockEkVar)
        #push!(BlockEp_Array, BlockEp)
        #push!(BlockEpVar_Array, BlockEpVar)
        #push!(BlockE_Array, BlockE)
        #push!(BlockEVar_Array, BlockEVar)
        #push!(BlockSgn_Array, avgSgn)
    end

    NUp_Mean = mean(BlockNUp_Array)
    NDn_Mean = mean(BlockNDn_Array)

    @printf(outputfile, "--------------------\n")
    @printf(outputfile, "Results for this run:\n")
    @printf(outputfile, "Average Sign: %.3f\n", mean(BlockSgn_Array))
    @printf(outputfile, "NUp            NDn            NTot\n")
    @printf(outputfile, "%.5f     %.5f       %.5f\n", NUp_Mean, NDn_Mean, NUp_Mean + NDn_Mean)

    #EkMean = mean(BlockEk_Array)
    #EkVar = sum(BlockEkVar_Array * (qmc.sample_size - 1) .+
    #qmc.sample_size * (BlockEk_Array .- EkMean) .^ 2) / (qmc.blocks * qmc.sample_size - 1)
    #EpMean = mean(BlockEp_Array)
    #EpVar = sum(BlockEpVar_Array * (qmc.sample_size - 1) .+
    #qmc.sample_size * (BlockEp_Array .- EpMean) .^ 2) / (qmc.blocks * qmc.sample_size - 1)
    #EMean = mean(BlockE_Array)
    #EVar = sum(BlockEVar_Array * (qmc.sample_size - 1) .+
    #qmc.sample_size * (BlockE_Array .- EMean) .^ 2) / (qmc.blocks * qmc.sample_size - 1)

    #@printf(outputfile, "--------------------\n")
    #@printf(outputfile, "Results for this run:\n")
    #@printf(outputfile, "Average Sign: %.3f\n", mean(BlockSgn_Array))
    #@printf(outputfile, "Ek            Ep            Etot\n")
    #@printf(outputfile, "%.5f     %.5f       %.5f\n", EkMean, EpMean, EMean)
    #@printf(outputfile, "Error:\n")
    #@printf(outputfile, "%.5f     %.5f       %.5f\n\n", sqrt(EkVar / qmc.blocks), sqrt(EpVar / qmc.blocks), sqrt(EVar / qmc.blocks))

    close(outputfile)

end


# A coupled convex quadratic programming problem with 14 variables and 7 Nodes

# Add processors
addprocs(2)

# Define functions for the algorithm on all processors
include("PDIP_Functions.jl")

# Define the Nodes
N7=Node(7,3,[7,14],nothing,nothing,[],[],[],[2 -1;-1 2],[1 1]',[1 10],[1]',[1 2],[1]',[],[],[],[],[],[],[],[],[],[],[])
N6=Node(6,3,[6,13],nothing,nothing,[],[],[],[2 -1;-1 2],[1 1]',[1 9],[1]',[1 2],[1]',[],[],[],[],[],[],[],[],[],[],[])
N4=Node(4,2,[4,11],nothing,nothing,[],[],[],[2 -1;-1 2],[1 1]',[1 3],[1]',[1 2],[1]',[],[],[],[],[],[],[],[],[],[],[])
N3=Node(3,1,[3,10],nothing,nothing,[],[],[],[2 -1;-1 2],[1 1]',[1 6],[1]',[1 2],[1]',[],[],[],[],[],[],[],[],[],[],[])
N5=Node(5,3,[5,12,13,14],nothing,[N6,N7],[],[],[],[5.8 .8 .5 .9;.8 4.1 .6 .5;.5 .6 4.2 .9;.9 .5 .9 4.1],[1 1 1 1]',[1 4 2 6;2 6 4 12],[1 2]',[1 2 3 4],[1]',[],[],[],[],[],[],[],[],[],[],[])
N2=Node(2,2,[2,9,11,12],nothing,[N4,N5],[],[],[],[5.8 .8 .5 .9;.8 4.1 .6 .5;.5 .6 4.2 .9;.9 .5 .9 4.1],[1 1 1 1]',[1 3 5 2;1 4 3 7],[1 2]',[1 2 3 4],[1]',[],[],[],[],[],[],[],[],[],[],[])
N1=Node(1,1,[1,8,9,10],nothing,[N2,N3],[],[],[],[5.8 .8 .5 .9;.8 4.1 .6 .5;.5 .6 4.2 .9;.9 .5 .9 4.1],[1 1 1 1]',[1 4 5 6;1 2 3 7;3 2 5 9],[1 2 3]',[1 2 3 4],[1]',[],[],[],[],[],[],[],[],[],[],[])

# Initialization
(n,q) = (14,7)           # Number of variables and Nodes
x_0 = .01 * ones(14,1)   # Initial starting point (should be feasible w.r.t inequality constraints)
Beta = .8                # Beta for backtracking line search
alpha = 0.1              # alpha for backtracking line search
mu = 15                  # mu for updating the barrier coefficient
epsilon_feas = 1e-6      # Primal and dual feasibility threshold
epsilon = 1e-6           # Surrogate duality gap threshold

x = x_0

setparent(N1)             # Set parent Nodes

setchildrenprocsid(N1)    # Set the processor id of child Nodes in which they are defined

SepRes(N1)                # Define separators and residuals

Preprocessing_EquCons(N1) # Perform preprocessing on equality constraints (QR factorization)

(m,p) = find_m_p(N1)      # Find the number of equality and inequality constraints

init_lambda_v_x(N1,x)     # initialize x, lambda and v

############################################
# Send N2 to processor 2
subtree_to_send = Base.deepcopy(N2)
subtree_to_send.parent = nothing  # No need to provide N2 with information regarding its parent Node
sendto(2, subtree=subtree_to_send)

# Send N5 to processor 3
subtree_to_send = nothing        # Clean the previously defined subtree_to_send
subtree_to_send = Base.deepcopy(N5)
subtree_to_send.parent = nothing  # No need to provide N2 with information regarding its parent Node
sendto(3, subtree=subtree_to_send)
############################################
# initialize the algorithm parameters
iter = 1
step = [0.];
back_no = [0.];

sur_duality = find_surduality(N1)
sur_duality=[sur_duality]

(res_primal_tot,res_dual_tot,res_cent_tot) =  TotalResiduals(N1,mu*m/sur_duality[1],n)
res_primal_tot_norm = (1/p) * norm(res_primal_tot)
res_dual_tot_norm = (1/m) * norm(res_dual_tot)

###############################################
# Start iterating !

while (sur_duality[iter] > epsilon || res_primal_tot_norm > epsilon_feas || res_dual_tot_norm > epsilon_feas)

  # Update the parameter t
  t1 = mu*m/sur_duality[iter];

  # Find the search direction
  MessagePassingDirection(N1,t1,n)

  # Find the corresponding Delta_lambda
  find_Delta_lambda(N1)

  # Find a proper step size
  StepInfo = MessagePassingStepSize(N1,t1,alpha,Beta,n)
  push!(step,StepInfo[1])
  push!(back_no,StepInfo[2])

  if back_no[iter+1]==1
    println("Algorithm could not find a proper step size!")
    break
  end

  # Update x, lambda and v
  Update_lambda_v_x(N1,step[iter+1])

  # Increment the iterator
  iter += 1

  # Find the surrogate duality gap for checking the termination criteria
  push!(sur_duality,0)
  sur_duality[iter] = find_surduality(N1)

  # Find the primal and dual residuals for checking the termination criteria
  (res_primal_tot,res_dual_tot,res_cent_tot) =  TotalResiduals(N1,t1,n)
  res_primal_tot_norm = (1/p) * norm(res_primal_tot)
  res_dual_tot_norm = (1/m) * norm(res_dual_tot)
end

# Display x, lambda and v locally on each Node
display(N1)
display(N3)
display(@fetchfrom 2 subtree)
display(@fetchfrom 2 subtree.children[1])
display(@fetchfrom 3 subtree)
display(@fetchfrom 3 subtree.children[1])
display(@fetchfrom 3 subtree.children[2])

# Display the algorithm parameters
println("Iteration number is $(iter-1)")                                # number of primal-dual interior-point iterations
println("Step size at each iteration is $(step)")                       # Step sizes at each iteration
println("Surrogate duality gap at each iteration is $(sur_duality)")    # Surrogate duality gap at each iteration


using PPInterpolation, LaverySpline, Test

@testset "HaganCurveLavery" begin
    #for quintic: first pass compute cubic first derivative, second pass, use cubic b, and optimize second derivative.
    t = [0.0, 0.1, 1.0, 2.0, 3.0, 4.0, 9.0, 20.0, 30.0]
    y = [0.06, 0.06, 0.06, 0.05, 0.05, 0.05, 0.065, 0.06, 0.06]
    t = [0.0, 0.1, 1.0, 4.0, 9.0, 20.0, 30.0]
    y = [0.081, 0.081, 0.07, 0.044, 0.07, 0.04, 0.03]
    t = [0, 0.1, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 2.0, 2.5, 3.0, 3.5]
    y = [0.012, 0.008, 0.0104, 0.00995, 0.010633333333333333, 0.010849999999999999, 0.010179999999999998, 0.009816666666666665, 0.008612499999999999, 0.007789999999999998, 0.007741666666666665, 0.007635714285714285]

    # t = [0.0, 0.1, 1, 2, 3, 4, 5, 6, 7, 30.0]
    # f = [0.01, 0.01, 0.01, 0.01, 0.01, 0.03, 0.03, 0.03, 0.03, 0.02]
    # y = zeros(length(t))
    # y[1] = f[1]
    # y[2:end] = (y[1:end-1] .* t[1:end-1] + (t[2:end]-t[1:end-1]) .* f[2:end]) ./ t[2:end]
    # y[3:end-2] = f[3:end-2] ./ t[3:end-2]
    # println(y)
    yt = zeros(length(t))
    yt[1:end] = y[1:end] .* t[1:end]
    spline = makeQuinticLaverySpline(t, yt, 1e-3, 1e-2, 3)
    nSample = 1000
    for i = 1:nSample
        ti = (t[end]) * (i) / (nSample)
        val = evaluateDerivative(spline, ti)
        zval = evaluate(spline, ti) / ti
        println(ti, " QuinticLavery3 ", val, " ", zval)
    end
end


export adjnodes, adjnzval, degree, dependency, expci

""" degree vector of an adjacency graph """
degree(G::SparseMatrixCSC) = vec(full(sum(G,1)))

"""
 Given an adj. matrix of a graph, return index of nodes
 adjacent to node i
"""
adjnodes(A::SparseMatrixCSC,i) = A.rowval[nzrange(A,i)]

"""
 Given an adj. matrix of a graph, return values of nodes
 adjacent to node i
"""
adjnzval(A::SparseMatrixCSC,i) = A.nzval[nzrange(A,i)]

"""
probability that any of i's neighbors will be conserved (netal)
"""
function dependency(G::SparseMatrixCSC, deg=degree(G))
    dep = 1./deg
    map!(x -> ifelse(isinf(x), 0, x), dep, dep)
end

"""
 approx expected number of conserved interactions (netal)
"""
function expci(G::SparseMatrixCSC, dep=dependency(G))
    map(i -> sum(dep[adjnodes(G,i)]), 1:size(G,1))
end

"""
expci(..)[i,j] = I[i,j] = expected # of conserved edges if i is aligned to j
"""
function expci(G1::SparseMatrixCSC,G2::SparseMatrixCSC,
               deg1=degree(G1),
               deg2=degree(G2),
               expci1=expci(G1,dependency(G1,deg1)),
               expci2=expci(G2,dependency(G2,deg2));
               normalize=true)
    z = normalize ? max(maximum(deg1),maximum(deg2)) : 1
    min.(expci1./z, (expci2./z)')
end


# Equations of motion and simulation functions for the double mass-spring-damper system

using Plots
using DelimitedFiles

frequency = 1/3.0*[1.0, 1.0]  # Hz
stiffness = (frequency * (2 * pi)).^2
damping = 0.05
dt = 1/(200*maximum(frequency))

K0 = [sum(stiffness) -stiffness[2];
        -stiffness[2] sum(stiffness)]
K(x) = K0 .+ sqrt(sum(x.^2))*K0
C(x) = damping * K0

function xddot(x, xdot, f)
    return f - K(x)*x - C(x)*xdot
end

function staticSolve(f)
    return K\f
end

function sim(f)
    N = size(f)[2]
    x = [0; 0]
    xdot = [0; 0]
    xOut = zeros(2, N)
    for ii in 1:N
        println("Iteration $ii")
        println("x: $x, xdot: $xdot")
        println("f: $(f[:,ii])")
        # x = staticSolve(f[:,ii])
        xOut[:,ii] = x
        xdot = xdot + dt*xddot(x, xdot, f[:,ii])
        x = x + dt*xdot
    end
    return xOut
end

tt = 0:dt:3/minimum(frequency)
fApplied = zeros(2, length(tt))
noise = sinpi.(2*2*maximum(frequency).*tt)
fApplied[2,:] = noise*0.05 .+ (sinpi.(2*1/2.0.*minimum(frequency)*tt)).*(1/2.0*(1 .- cospi.((2*3/tt[length(tt)]).*tt)))
# fApplied[2,:] .= 1.0
xOut = sim(fApplied * 10*maximum(stiffness))

# plot(tt, xOut[1,:]/maximum(xOut[1,:]), label="Output")
# plot!(tt, fApplied[2,:]/maximum(fApplied[2,:]), label="Input", alpha=0.2)
# # scatter(xOut[1,:], fApplied[2,:])
# xlabel!("Time, s")
# # ylabel!("Output")
# # ylabel!("Input")
# savefig("spring.pdf")

f = open("spring.csv", "w")
for ii in 1:length(tt)
    write(f, "$(tt[ii]),$(fApplied[2,ii]),$(xOut[1,ii])\n")
end
close(f)


using DoubleFloats

using SemiLagrangian:
    UniformMesh,
    AbstractInterpolation,
    points,
    dotprod,
    compute_ee,
    compute_ke,
    PoissonConst,
    PoissonVar,
    getpoissonvar,
    initcoef!,
    compute_charge!,
    compute_elfield,
    totuple,
    Lagrange,
    Advection,
    AdvectionData,
    advection!,
    NoTimeOpt,
    SimpleThreadsOpt,
    SplitThreadsOpt,
    PrepareFftBig,
    standardsplit,
    strangsplit,
    triplejumpsplit,
    order6split,
    hamsplit_3_11,
    getenergyall

function initmesh(t_deb, t_end, t_size)
    t_step = (t_end - t_deb) ./ t_size
    return totuple(UniformMesh.(t_deb, t_end, t_size)), t_step
end

function test_poisson(T::DataType, isfft = true)

    t_debsp = T.([-1 // 1, -10 // 1, -3 // 1])
    t_endsp = T.([3 // 1, 6 // 1, 5 // 1])
    t_szsp = (8, 4, 16)
    Nsp = 3
    t_debv = T.([-3 // 1, -9 // 1, 1 // 1])
    t_endv = T.([1 // 1, 7 // 1, 5 // 1])
    t_szv = (4, 8, 4)
    base_dt = one(T) / 80
    Nv = 3

    N = Nsp + Nv

    t_meshsp, t_stepsp = initmesh(t_debsp, t_endsp, t_szsp)
    t_meshv, t_stepv = initmesh(t_debv, t_endv, t_szv)

    interp = Lagrange(3, T)
    adv = Advection(
        (t_meshsp..., t_meshv...),
        map(x -> interp, 1:6),
        base_dt,
        [([1, 2, 3, 4, 5, 6], 3, 1, false), ([4, 5, 6, 1, 2, 3], 3, 2, false)],
    )

    tab = rand(T, t_szsp..., t_szv...)

    pc = PoissonConst(adv; isfftbig = isfft)

    @test pc.v_square == dotprod(points.(adv.t_mesh[N-Nv+1:N])) .^ 2

    # #    println("t_perms=$(pc.t_perms)")
    # @test pc.t_perms == (
    #     [1, 2, 3, 6, 5, 4],
    #     [2, 1, 3, 4, 6, 5],
    #     [3, 2, 1, 4, 5, 6], # space states
    #     [4, 5, 6, 1, 2, 3],
    #     [5, 4, 6, 1, 2, 3],
    #     [6, 5, 4, 1, 2, 3],  # velocity states
    # )

    pvar = PoissonVar(pc)

    advdata = AdvectionData(adv, tab, pvar)

    @test compute_ke(advdata) == compute_ke(totuple(t_meshsp), totuple(t_meshv), tab)

    advdata.state_gen = 2
    initcoef!(pvar, advdata)
    rhoref = zeros(T, t_szsp)
    compute_charge!(rhoref, t_meshv, tab)

    @test rhoref == pvar.rho

    pfft = if isfft
        PrepareFftBig(t_szsp, T; numdims = Nsp, dims = ntuple(x -> x, Nsp))
    else
        missing
    end

    refelfield = compute_elfield(t_meshsp, rhoref, pfft)

    #    prec = (T==BigFloat) ? 1e-70 : 1e-15

    for i = 1:Nsp
        #        @test isapprox(refelfield[i], pvar.t_elfield[i], atol=prec, rtol=prec)
        @test isapprox(refelfield[i], pvar.t_elfield[i])
    end

    #    @test isapprox(compute_ee(t_meshsp, refelfield), compute_ee(advdata), atol=prec, rtol=prec)
    @test isapprox(compute_ee(t_meshsp, refelfield), compute_ee(advdata))

end

@testset "test compute_ee" begin
    t_deb = [big"-1" // 1, -10 // 1, -3 // 1, -1 // 1]
    t_end = [big"3" // 1, 6 // 1, 5 // 1, 1 // 1]
    t_sz = [4, 4, 8, 16]
    t_step = (t_end - t_deb) ./ t_sz
    tt_mesh = UniformMesh.(t_deb, t_end, t_sz)
    t_mesh = totuple(tt_mesh)
    N = 4
    t = ntuple(x -> rationalize.(BigInt, rand(Float64, totuple(t_sz))), N)

    resref = sum(map(x -> prod(step, t_mesh) * sum(x .^ 2), t))

    @test resref == compute_ee(t_mesh, t)

end

function test_poisson_real(T::DataType, timeopt)
    dt = T(1 // 10)
    nbdt = 10
    epsilon = 0.5
    spmin, spmax, nsp = T(0), T(4big(pi)), 128
    vmin, vmax, nv = -T(10), T(10), 128

    mesh_sp = UniformMesh(spmin, spmax, nsp)
    mesh_v = UniformMesh(vmin, vmax, nv)

    tabst = [([1, 2], 1, 1, true), ([2, 1], 1, 2, true)]

    interp = Lagrange(9, T)

    adv = Advection((mesh_sp, mesh_v), [interp, interp], dt, tabst, timeopt = timeopt)

    fct_sp(x) = epsilon * cos(x / 2) + 1
    fct_v(v) = exp(-v^2 / 2) / sqrt(2T(pi))

    lgn_sp = fct_sp.(mesh_sp.points)
    lgn_v = fct_v.(mesh_v.points)

    data = dotprod((lgn_sp, lgn_v))

    pvar = getpoissonvar(adv)

    advd = AdvectionData(adv, data, pvar)

    for i = 1:nbdt
        while advection!(advd)
        end
    end
    elenergy = compute_ee(advd)
    kinenergy = compute_ke(advd)

    return elenergy, kinenergy
end
function test_poisson_split(
    sz::NTuple{2,Int},
    interp::Vector{I},
    dt::T,
    nbdt::Int,
    tcoef::Vector{T},
) where {T,I<:AbstractInterpolation{T}}
    t_max = dt * nbdt
    epsilon = 0.5
    spmin, spmax, nsp = T(0), T(4big(pi)), sz[1]
    vmin, vmax, nv = -T(10), T(10), sz[2]

    mesh_sp = UniformMesh(spmin, spmax, nsp)
    mesh_v = UniformMesh(vmin, vmax, nv)

    tabst = [([2, 1], 1, 1, true), ([1, 2], 1, 2, true)]

    adv = Advection((mesh_sp, mesh_v), interp, dt, tabst; tab_coef = tcoef)

    fct_sp(x) = epsilon * cos(x / 2) + 1
    fct_v(v) = exp(-v^2 / 2) / sqrt(2T(pi))

    lgn_sp = fct_sp.(mesh_sp.points)
    lgn_v = fct_v.(mesh_v.points)

    data = dotprod((lgn_sp, lgn_v))

    pvar = getpoissonvar(adv)

    advd = AdvectionData(adv, data, pvar)

    enall = getenergyall(advd)
    minen = maxen = enall
    for i = 1:nbdt
        while advection!(advd)
        end
        enall = getenergyall(advd)
        minen = min(enall, minen)
        maxen = max(enall, maxen)
        diff = maxen - minen
        @show advd.time_cur, diff
    end

    return maxen - minen
end
@testset "Poisson Float64" begin
    test_poisson(Float64, true)
    test_poisson(Float64, false)
end

@testset "Poisson BigFloat" begin
    test_poisson(BigFloat)
end

@testset "Poisson Double64" begin
    test_poisson(Double64)
end

@testset "Poisson Threads" begin
    res = test_poisson_real(Float64, NoTimeOpt)
    @test res == test_poisson_real(Float64, SimpleThreadsOpt)
    @test res == test_poisson_real(Float64, SplitThreadsOpt)
end

function test_split(T, nbdt, split, order)
    lag = Lagrange(19, T)
    dt = one(T) / nbdt
    sz = (64, 50)
    res1 = test_poisson_split(sz, [lag, lag], dt, nbdt, split(dt))
    nbdt *= 2
    dt = one(T) / nbdt
    res2 = test_poisson_split(sz, [lag, lag], dt, nbdt, split(dt))

    @show order, res1, res2, res1 / res2
    @test (1.05 * res1) / res2 > 2^order
end

@testset "Poisson split" begin
    @time test_split(Float64, 5, standardsplit, 1)
    @time test_split(Float64, 5, strangsplit, 2)
    @time test_split(Double64, 10, triplejumpsplit, 4)
    @time test_split(Double64, 5, order6split, 6)
    @time test_split(Double64, 5, hamsplit_3_11, 6)
end


# functions related to F distribution

# R implementations
# For pdf and logpdf we use the Julia implementation
using .RFunctions:
    fdistcdf,
    fdistccdf,
    fdistlogcdf,
    fdistlogccdf,
    fdistinvcdf,
    fdistinvccdf,
    fdistinvlogcdf,
    fdistinvlogccdf

# Julia implementations
fdistpdf(ν1::Real, ν2::Real, x::Real) = exp(fdistlogpdf(ν1, ν2, x))

fdistlogpdf(ν1::Real, ν2::Real, x::Real) = fdistlogpdf(promote(ν1, ν2, x)...)
function fdistlogpdf(ν1::T, ν2::T, x::T) where {T<:Real}
    # we ensure that `log(x)` does not error if `x < 0`
    ν1ν2 = ν1 / ν2
    y = max(x, 0)
    val = (xlogy(ν1, ν1ν2) + xlogy(ν1 - 2, y) - xlogy(ν1 + ν2, 1 + ν1ν2 * y)) / 2 - logbeta(ν1 / 2, ν2 / 2)
    return x < 0 ? oftype(val, -Inf) : val
end


"""
    GSDist <: ContinuousUnivariateDistribution

Generalized S-Distribution for approximating univariate distributions.

# Arguments
- `F₀`: the reference quantile.
- `x₀`: the value at the reference quantile.
- `α, g, k, γ`: parameters of the Generalized S-Distribution
- `dist`: the underlying probability distribution.
- `F`: the cdf of the distribution.
"""
struct GSDist <: ContinuousUnivariateDistribution
    F₀::Real
    x₀::Real
    α::Real
    g::Real
    k::Real
    γ::Real
    dist::Union{UnivariateDistribution, Nothing} # underlying distribution
    F::Union{Function, Nothing}                  # CDF
    function GSDist(F₀, x₀, α, g, k, γ, dist, F)
        F₀ < 0 || F₀ > 1 && throw(DomainError(F₀, "F₀ must be between 0 and 1"))
        α ≤ 0 && throw(DomainError(α, "α must be a positive real number"))
        g ≤ 0 && throw(DomainError(g, "g must be a positive real number"))
        k ≤ 0 && throw(DomainError(k, "k must be a positive real number"))
        γ ≤ 0 && throw(DomainError(γ, "γ must be a positive real number"))

        new(F₀, x₀, α, g, k, γ, dist, F)
    end
end
GSDist(F₀, x₀, α, g, k, γ) = GSDist(F₀, x₀, α, g, k, γ, nothing, nothing)
GSDist(F₀, x₀, α, g, k, γ, dist) = GSDist(F₀, x₀, α, g, k, γ, dist, _piecewise_cdf(α, g, k, γ, dist))

"""
    GSDist(dist::UnivariateDistribution, F₀::Real=0.5;
    n::Int=21, diff::Function=three_point_midpoint, h::Real=1.0)

# Arguments
- `dist::UnivariateDistribution`: a univariate distribution.
- `F₀::Real`: the reference quantile to center the distribution around.
- `n::Integer`: the number of points on the distribution to use for least squares fit.
- `diff::Function`: one of the finite difference methods for gradient approximation.
- `h::Real`: the denominator used in the diff function.
"""
function GSDist(dist::UnivariateDistribution, F₀::Real=0.5;
    n::Int=21, diff::Function=three_point_midpoint, h::Real=1.0)
    F₀ < 0 || F₀ > 1 && throw(DomainError(F₀, "F₀ must be between 0 and 1"))
    F₀ = Float64(F₀)

    l, u = extrema(dist)
    l = isinf(l) ? quantile(dist, 1e-6) : l
    u = isinf(u) ? quantile(dist, 1 - 1e-6) : u

    ql, qu = cdf.(dist, (l,u))

    q = range(ql, qu, length=n)
    X = typeof(dist) <: DiscreteUnivariateDistribution ? Set(quantile.(dist, q)) : quantile.(dist, q)

    FX = cdf.(dist, X)
    dF = y -> diff(x -> cdf(dist, x), y, h)
    fX = typeof(dist) <: DiscreteUnivariateDistribution ? dF.(X) : pdf.(dist, X)

    f = (t, p) -> p[1] * t.^p[2] .* (1.0 .- t.^p[3]).^p[4]

    p0 = [1, 0.5, 1, 0.5]
    fit = curve_fit(f, FX, fX, p0)
    p = coef(fit)

    F = _piecewise_cdf(p..., dist)

    GSDist(F₀, quantile(dist, F₀), p..., dist, F)
end

params(G::GSDist) = (G.α, G.g, G.k, G.γ)

function quantile(D::GSDist, p::Real)
    q = _gsd_quantile(p, D.F₀, D.x₀, D.α, D.g, D.k, D.γ)
    if isnan(q)
        @warn "Unable to calculate the quantile of the GSDist. Falling back to the quantile of the underlying distribution."
        quantile(D.dist, p)
    else
        q
    end
end

function mean(G::GSDist)
    m = _gsd_mean(G.F₀, G.x₀, G.α, G.g, G.k, G.γ)
    if isnan(m)
        @warn "Unable to calculate the mean of the GSDist. Falling back to the mean of the underlying distribution."
        mean(G.dist)
    else
        m
    end
end

function var(G::GSDist)
    m1 = mean(G)
    m2 = try
        _gsd_moment(G.F₀, G.x₀, G.α, G.g, G.k, G.γ; moment=2)
    catch y
        if isa(y, DomainError)
            @warn "Unable to calculate the variance of the GSDist. Falling back to the variance of the underlying distribution."
            return var(G.dist)
        end
    end
    m2 - m1^2
end

std(G::GSDist) = sqrt(var(G))

cdf(G::GSDist, x::Real) = G.F(x)

function pdf(G::GSDist, x::Real)
    F = cdf(G, x)
    α, g, k, γ = params(G)
    α * F^g * (1 - F^k)^γ
end

function logpdf(G::GSDist, x::Real)
    F = cdf(G, x)
    α, g, k, γ = params(G)
    log(α) + g*log(F) + γ * log(1 - F^k)
end

minimum(G::GSDist) = minimum(G.dist)
maximum(G::GSDist) = maximum(G.dist)

insupport(G::GSDist, x::Real) = minimum(G) ≤ x ≤ maximum(G)
insupport(G::GSDist, x::VecOrMat{<:Real}) = insupport.(G, x)

function mode(G::GSDist)
    α, g, k, γ = params(G)
    x = (g / (k*γ + g))^(1/k)
    quantile(G, x)
end

# skewness(G::GSDist) = _gsd_moment(G.F₀, G.x₀, G.α, G.g, G.k, G.γ; moment=3)
# kurtosis(G::GSDist) = _gsd_moment(G.F₀, G.x₀, G.α, G.g, G.k, G.γ; moment=4)
# entropy(G::GSDist, b::Real)
# mgf(G::GSDist, t::Any)
# cf(G::GSDist, t::Any)


using Random

include("get_div_grad.jl")

function test_cr()
  cr_tol = 1.0e-6

  # Cubic spline matrix _ case: ‖x*‖ > Δ
  n = 10
  A = spdiagm(-1 => ones(n-1), 0 => 4 * ones(n), 1 => ones(n-1))
  b = A * [1:n;]

  (x, stats) = cr(A, b)
  r = b - A * x
  resid = norm(r) / norm(b)
  @printf("CR: Relative residual: %8.1e\n", resid)
  @test(resid <= cr_tol)
  @test(stats.solved)

  # Code coverage
  (x, stats) = cr(Matrix(A), b)
  show(stats)

  radius = 0.75 * norm(x)
  (x, stats) = cr(A, b, radius=radius)
  show(stats)
  @test(stats.solved)
  @test abs(norm(x) - radius) ≤ cr_tol * radius

  # Sparse Laplacian
  A = get_div_grad(16, 16, 16)
  b = randn(size(A, 1))
  itmax = 0
  # case: ‖x*‖ > Δ
  radius = 10.
  (x, stats) = cr(A, b, radius=radius)
  xNorm = norm(x)
  r = b - A * x
  resid = norm(r) / norm(b)
  @printf("CR: Relative residual: %8.1e\n", resid)
  @test abs(xNorm - radius) ≤ cr_tol * radius
  @test(stats.solved)
  # case: ‖x*‖ < Δ
  radius = 30.
  (x, stats) = cr(A, b, radius=radius)
  xNorm = norm(x)
  r = b - A * x
  resid = norm(r) / norm(b)
  @printf("CR: Relative residual: %8.1e\n", resid)
  @test(resid <= cr_tol)
  @test(stats.solved)

  radius = 0.75 * xNorm
  (x, stats) = cr(A, b, radius=radius)
  show(stats)
  @test(stats.solved)
  @test(abs(radius - norm(x)) <= cr_tol * radius)

  opA = LinearOperator(A)
  (xop, statsop) = cr(opA, b, radius=radius)
  @test(abs(radius - norm(xop)) <= cr_tol * radius)

  n = 100
  itmax = 2 * n
  B = LBFGSOperator(n)
  Random.seed!(0)
  for i = 1:5
    push!(B, rand(n), rand(n))
  end
  b = B * ones(n)
  (x, stats) = cr(B, b)
  @test norm(x - ones(n)) ≤ cr_tol * norm(x)
  @test stats.solved

  # Test b == 0
  (x, stats) = cr(A, zeros(size(A,1)))
  @test x == zeros(size(A,1))
  @test stats.status == "x = 0 is a zero-residual solution"

  # Test integer values
  A = [4 -1 0; -1 4 -1; 0 -1 4]
  b = [7; 2; -1]
  (x, stats) = cr(A, b)
  @test stats.solved

  # Test with Jacobi (or diagonal) preconditioner
  A = ones(10,10) + 9 * I
  b = 10 * [1:10;]
  M = 1/10 * opEye(10)
  (x, stats) = cr(A, b, M=M, itmax=10);
  show(stats)
  r = b - A * x;
  resid = norm(r) / norm(b);
  @printf("CR: Relative residual: %8.1e\n", resid);
  @test(resid <= cr_tol);
  # @test(stats.solved);
end

test_cr()


module DiffPrivacy
using Random, Distributions

function get_one_parallel_database(A, remove_index)

    # Returns an array A, without the element at position remove_index, of size (1, length(A) - 1)

    return vcat(A[1:remove_index-1], A[remove_index+1:length(A)])'
end

function get_multiple_parallel_databases(A)

    # Returns Array of length(A) with all the parallel databases generated and stored in parallel_databases

    parallel_databases = []

    for i in 1:length(A)
        x = get_one_parallel_database(A, i)
        push!(parallel_databases, x)
    end

    return parallel_databases
end

function sum_sensitivity(A)

    db = A
    pdbs = get_multiple_parallel_databases(A)

    full_db_result = sum(db)

    max_distance = 0
    for pdb in pdbs
        pdb_result = sum(pdb)

        db_distance = abs(pdb_result - full_db_result)

        if(db_distance > max_distance)
            max_distance = db_distance
        end
    end

    return max_distance
end

function mean_sensitivity(A)

    db = A
    pdbs = get_multiple_parallel_databases(A)

    full_db_result = sum(db) / length(db)

    max_distance = 0
    for pdb in pdbs
        pdb_result = sum(pdb) / length(pdb)

        db_distance = abs(pdb_result - full_db_result)

        if(db_distance > max_distance)
            max_distance = db_distance
        end
    end

    return max_distance
end

function sum_sensitive_query(A, threshold)
    if sum(A) > threshold
        return 1
    else
        return 0
    end
end

function sum_sensitivity_threshold(A, threshold)

    db = A
    pdbs = get_multiple_parallel_databases(A)

    full_db_result = sum_sensitive_query(db, threshold)

    max_distance = 0
    for pdb in pdbs
        pdb_result = sum_sensitive_query(pdb, threshold)

        db_distance = abs(pdb_result - full_db_result)

        if(db_distance > max_distance)
            max_distance = db_distance
        end
    end

    return max_distance
end

function add_gaussian_noise(A, mu, sigma)
    d = rand(Gaussian(mu, sigma), size(A))
    A = A + d
    return A
end

function add_laplacian_noise(A, loc, scale)
    d = rand(Laplacian(loc, scale), size(A))
    A = A + d
    push!(A, d)
    return A
end

function add_normal_noise(A, mu, sigma)
    d = rand(Normal(mu, sigma), size(A))
    A = A + d
    return A
end

export DiffPrivacy

end # module


function ccc(x::AbstractVector, y::AbstractVector; lim=(0, Inf))
    x = vcat(x...)
    y = vcat(y...)
    @assert length(x) == length(y)
    mask = trues(length(x))
    for i in eachindex(x)
        mask[i] = !isnan(x[i]) & !isnan(y[i]) & (x[i]>lim[1]) & (x[i]<lim[2]) & (y[i]>lim[1]) & (y[i]<lim[2])
    end
    x = x[mask]
    y = y[mask]

    mx = mean(x)
    my = mean(y)
    sx = var(x) * (length(x) - 1) / length(x)
    sy = var(y) * (length(y) - 1) / length(y)
    sxy = sum((x .- mx) .* (y .- my)) / length(x)
    ccc_value = 2 * sxy / (sx + sy + (mx - mx)^2)
    return ccc_value
end

function ccc(fits::Dict; lim=(0, Inf))
    models = keys(fits)
    ccc_values = Dict()
    for model_a in models
        inner_dict = Dict()
        for model_b in models
            inner_dict[model_b] = ccc(fits[model_a], fits[model_b], lim=lim)
        end
        ccc_values[model_a] = inner_dict
    end
    return ccc_values
end

function apply_mask(; data, mask)
    @assert length(size(data)) >= length(size(mask))
    mask_indices = findall(mask)
    return data[mask_indices, :]
end

function negatives_to_zero!(x::AbstractArray)
    for i in eachindex(x)
        x[i] = x[i] < 0 ? 0 : x[i]
    end
    return x
end

function crop(data::AbstractArray; mask=nothing)
    if isnothing(mask)
        mask = @. !isnan(data) & (data > 0)
    end
    xlim = findall(vec(sum(mask, dims = [2,3])) .> 0)
    ylim = findall(vec(sum(mask, dims = [1,3])) .> 0)
    return data[xlim, ylim, :]
end
function crop(data::NamedTuple; mask=nothing)
    for key in keys(data)
        data[key] = crop(data[key]; mask = mask)
    end
    return data
end


"""
```
koopman_smoother(y, T, R, C, Q, Z, D, E, s_0, P_0, s_pred, P_pred;
    Nt0 = 0)

koopman_smoother(regime_indices, y, Ts, Rs, Cs, Qs, Zs, Ds, Es,
    s_0, P_0, s_pred, P_pred; Nt0 = 0)
```

This is a Kalman smoothing program based on S.J. Koopman's \"Disturbance
Smoother for State Space Models\" (Biometrika, 1993), as specified in
Durbin and Koopman's \"A Simple and Efficient Simulation Smoother for
State Space Time Series Analysis\" (Biometrika, 2002).

Unlike other Kalman smoothing programs, there is no need to invert
singular matrices using the Moore-Penrose pseudoinverse (`pinv`), which
should lead to efficiency gains and fewer inversion problems. Also, the
states vector and the corresponding matrices do not need to be augmented
to include the shock innovations. Instead they are saved automatically
in the `ϵ_smth` matrix.

The state space is given by:

```
s_{t+1} = C + T*s_t + R*ϵ_t    (transition equation)
y_t     = D + Z*s_t + u_t      (measurement equation)

ϵ_t ∼ N(0, Q)
u_t ∼ N(0, E)
Cov(ϵ_t, u_t) = 0
```

### Inputs

- `y`: `Ny` x `Nt` matrix containing data `y_1, ... , y_T`
- `s_0`: `Ns` x 1 initial state vector
- `P_0`: `Ns` x `Ns` initial state covariance matrix
- `s_pred`: `Ns` x `Nt` matrix of one-step predicted state vectors `s_{t|t-1}`
  (from the Kalman filter)
- `P_pred`: `Ns` x `Ns` x `Nt` array of mean squared errors `P_{t|t-1}` of
  predicted state vectors

**Method 1 only:** state-space system matrices `T`, `R`, `C`, `Q`, `Z`,
`D`, `E`. See `?kalman_filter`

**Method 2 only:** `regime_indices` and system matrices for each regime `Ts`,
`Rs`, `Cs`, `Qs`, `Zs`, `Ds`, `Es`. See `?kalman_filter`

where:

- `Nt`: number of periods for which we have data
- `Ns`: number of states
- `Ne`: number of shocks
- `Ny`: number of observables

### Keyword Arguments

- `Nt0`: if greater than 0, the returned smoothed states and
  shocks matrices will be shorter by that number of columns (taken from the
  beginning)

### Outputs

- `s_smth`: `Ns` x `Nt` matrix of smoothed states `s_{t|T}`
- `ϵ_smth`: `Ne` x `Nt` matrix of smoothed shocks `ϵ_{t|T}`
"""
function koopman_smoother(y::AbstractMatrix,
    T::Matrix{S}, R::Matrix{S}, C::Vector{S}, Q::Matrix{S},
    Z::Matrix{S}, D::Vector{S}, E::Matrix{S},
    s_0::Vector{S}, P_0::Matrix{S}, s_pred::Matrix{S}, P_pred::Array{S, 3};
    Nt0::Int = 0) where {S<:AbstractFloat}

    Nt = size(y, 2)
    koopman_smoother(UnitRange{Int}[1:Nt], y, Matrix{S}[T], Matrix{S}[R], Vector{S}[C],
        Matrix{S}[Q], Matrix{S}[Z], Vector{S}[D], Matrix{S}[E],
        s_0, P_0, s_pred, P_pred; Nt0 = Nt0)
end

function koopman_smoother(regime_indices::Vector{UnitRange{Int}}, y::AbstractMatrix,
    Ts::Vector{Matrix{S}}, Rs::Vector{Matrix{S}}, Cs::Vector{Vector{S}}, Qs::Vector{Matrix{S}},
    Zs::Vector{Matrix{S}}, Ds::Vector{Vector{S}}, Es::Vector{Matrix{S}},
    s_0::Vector{S}, P_0::Matrix{S}, s_pred::Matrix{S}, P_pred::Array{S, 3};
    Nt0::Int = 0) where {S<:AbstractFloat}

    # Dimensions
    n_regimes = length(regime_indices)
    Nt = size(y, 2)     # number of periods of data
    Ns = size(Ts[1], 1) # number of states
    Ne = size(Rs[1], 2) # number of shocks

    # Call disturbance smoother
    s_dist, _ = koopman_disturbance_smoother(regime_indices, y, Ts, Rs, Qs,
                                             Zs, Ds, Es, s_pred, P_pred; Nt0 = 0)

    # Initialize outputs
    s_smth = zeros(S, Ns, Nt)
    ϵ_smth = zeros(S, Ne, Nt)

    s_t = zeros(S, Ns) # initialize dummy value s.t. s_t is in scope

    for i = 1:n_regimes
        # Get state-space system matrices for this regime
        T, R, Q = Ts[i], Rs[i], Qs[i]

        for t in regime_indices[i]
            r_t = s_dist[:, t]
            s_t = if t == 1
                s_0 + P_0*r_t
            else
                T*s_t + R*Q*R'*r_t
            end

            s_smth[:, t] = s_t
            ϵ_smth[:, t] = Q*R'*r_t
        end
    end

    # Trim the presample if needed
    if Nt0 > 0
        insample = Nt0+1:Nt

        s_smth = s_smth[:, insample]
        ϵ_smth = ϵ_smth[:, insample]
    end

    return s_smth, ϵ_smth
end

"""
```
koopman_disturbance_smoother(y, T, R, Q, Z, D, E, s_pred, P_pred;
    Nt0 = 0)

koopman_smoother(regime_indices, y, Ts, Rs, Qs, Zs, Ds, Es,
    s_pred, P_pred; Nt0 = 0)
```

This disturbance smoother is intended for use with the state smoother
`koopman_smoother` from S.J. Koopman's \"Disturbance Smoother for State Space
Models\" (Biometrika, 1993), as specified in Durbin and Koopman's \"A Simple and
Efficient Simulation Smoother for State Space Time Series Analysis\"
(Biometrika, 2002).

The state space is given by:

```
s_{t+1} = C + T*s_t + R*ϵ_t    (transition equation)
y_t     = D + Z*s_t + u_t      (measurement equation)

ϵ_t ∼ N(0, Q)
u_t ∼ N(0, E)
Cov(ϵ_t, u_t) = 0
```

### Inputs

- `y`: `Ny` x `Nt` matrix containing data `y_1, ... , y_T`
- `s_pred`: `Ns` x `Nt` matrix of one-step predicted state vectors `s_{t|t-1}`
  (from the Kalman filter)
- `P_pred`: `Ns` x `Ns` x `Nt` array of mean squared errors `P_{t|t-1}` of
  predicted state vectors

**Method 1 only:** state-space system matrices `T`, `R`, `Q`, `Z`, `D`, `E`. See
`?kalman_filter`

**Method 2 only:** `regime_indices` and system matrices for each regime `Ts`,
`Rs`, `Qs`, `Zs`, `Ds`, `Es`. See `?kalman_filter`

where:

- `Nt`: number of periods for which we have data
- `Ns`: number of states
- `Ne`: number of shocks
- `Ny`: number of observables

### Keyword Arguments

- `Nt0`: if greater than 0, the returned smoothed disturbances and shocks
  matrices will be shorter by that number of columns (taken from the beginning)

### Outputs

- `s_dist`: `Ns` x `Nt` matrix of transition equation disturbances `r_t`
- `y_dist`: `Ny` x `Nt` matrix of measurement equation disturbances `e_t`
"""
function koopman_disturbance_smoother(y::AbstractArray,
    T::Matrix{S}, R::Matrix{S}, Q::Matrix{S},
    Z::Matrix{S}, D::Vector{S}, E::Matrix{S},
    s_pred::Matrix{S}, P_pred::Array{S, 3}; Nt0::Int = 0) where {S<:AbstractFloat}

    Nt = size(y, 2)
    koopman_disturbance_smoother(UnitRange{Int}[1:Nt], y, Matrix{S}[T], Matrix{S}[R],
        Matrix{S}[Q], Matrix{S}[Z], Vector{S}[D], Matrix{S}[E],
        s_pred, P_pred; Nt0 = Nt0)
end

function koopman_disturbance_smoother(regime_indices::Vector{UnitRange{Int}}, y::AbstractArray,
    Ts::Vector{Matrix{S}}, Rs::Vector{Matrix{S}}, Qs::Vector{Matrix{S}},
    Zs::Vector{Matrix{S}}, Ds::Vector{Vector{S}}, Es::Vector{Matrix{S}},
    s_pred::Matrix{S}, P_pred::Array{S, 3}; Nt0::Int = 0) where {S<:AbstractFloat}

    # Dimensions
    Nt = size(y ,    2) # number of periods of data
    Ns = size(Ts[1], 1) # number of states
    Ne = size(Rs[1], 2) # number of shocks
    Ny = size(Zs[1], 1) # number of observables

    # Initialize outputs
    s_dist = zeros(S, Ns, Nt)
    y_dist = zeros(S, Ny, Nt)

    r_t = zeros(S, Ns) # r_0 = 0

    for i = length(regime_indices):-1:1
        # Get state-space system matrices for this regime
        T, R, Q = Ts[i], Rs[i], Qs[i]
        Z, D, E = Zs[i], Ds[i], Es[i]

        for t in reverse(regime_indices[i])
            # Keep rows of measurement equation corresponding to nonmissing observables
            nonmissing = .!map(x -> ismissing(x) ? true : isnan(x), y[:, t]) #.!ismissing.(y[:, t])
            y_t = y[nonmissing, t]
            Z_t = Z[nonmissing, :]
            D_t = D[nonmissing]
            E_t = E[nonmissing, nonmissing]

            s_pred_t = s_pred[:, t]            # s_{t|t-1}
            P_pred_t = P_pred[:, :, t]         # P_{t|t-1} = Var s_{t|t-1}

            y_pred_t = Z_t*s_pred_t + D_t      # y_{t|t-1} = Z*s_{t|t-1} + D
            V_pred_t = Z_t*P_pred_t*Z_t' + E_t # V_{t|t-1} = Var y_{t|t-1} = Z*P_{t|t-1}*Z' + E
            dy = y_t - y_pred_t                # dy = y_t - y_{t|t-1} = prediction error
            K = T*P_pred_t*Z_t'/V_pred_t       # K_t = T*P_{t|t-1}'Z'/V_{t|t-1} = Kalman gain

            e_t = V_pred_t\dy - K'*r_t         # e_t     = (1/V_{t|t-1})dy - K_t'*r_t
            r_t = Z_t'*e_t + T'*r_t            # r_{t-1} = Z'*e_t + T'*r_t

            s_dist[:,      t] = r_t
            y_dist[nonmissing, t] = e_t
        end # of loop backward through this regime's periods
    end # of loop backward through regimes

    # Trim the presample if needed
    if Nt0 > 0
        insample = Nt0+1:Nt
        s_dist = s_dist[:, insample]
        y_dist = y_dist[:, insample]
    end

    return s_dist, y_dist
end


using DataFrames, Gadfly

set_default_plot_size(3.3inch, 3.3inch)

srand(123)
D = convert(DataFrame, 99*rand(4, 4)+0.5)

xsc  = Scale.x_continuous(minvalue=0.0, maxvalue=100)
ysc  = Scale.y_continuous(minvalue=0.0, maxvalue=100)

plot(D, x=:x1, y=:x2, xend=:x3, yend=:x4, Geom.segment(arrow=true), xsc, ysc)
plot(D, x=:x1, y=:x2, xend=:x3, yend=:x4, Geom.vector, xsc, ysc)

plot(z=(x,y)->x*exp(-(x^2+y^2)),
        xmin=[-2], xmax=[2], ymin=[-2], ymax=[2],
        Geom.vectorfield(scale=0.4, samples=17),
        Scale.x_continuous(minvalue=-2.0, maxvalue=2.0),
        Scale.y_continuous(minvalue=-2.0, maxvalue=2.0),
        Guide.xlabel("x"), Guide.ylabel("y"), Guide.colorkey(title="z")
    )

Z = fill(1.0, 5, 5)
Z[2:4, 2:4] = 2.0
Z[3,3] = 3.0

plot(z=Z,
        Geom.vectorfield,
        Scale.x_continuous(minvalue=1.0, maxvalue=5.0),
        Scale.y_continuous(minvalue=1.0, maxvalue=5.0),
        Guide.xlabel("x"), Guide.ylabel("y"),
        Theme(key_position=:none)
    )


##
using Plots
using FFTW
using DSP
sr = 48000
baseFreq = 20
overSamp = 2

maxHarms = floor(Int, sr / (3.0+baseFreq) + 0.5)
v = maxHarms
v -= 1
v |= v >> 1
v |= v >> 2
v |= v >> 4
v |= v >> 8
v |= v >> 16
v += 1
tableLen = v * 2 * overSamp

topFreq = baseFreq * 2 / sr

## sawtooth
nHarms = 32
a = zeros(Float64,tableLen)
for i=1:nHarms
    temp = -1.0/i
    a[i+1]=-temp
    a[tableLen-i+1]=temp
end
# plot(a)

A = fft(a)
plot([real.(A),imag.(A)])

##
shift = 1280
x = imag.(A)
xs = circshift(x,shift)
p1 = plot([x,xs])
X = ifft(xs)
p2 = plot([a[1:32],abs.(X[1:32])])
p3 = plot(unwrap(angle.(X[1:32])))
#p2 = plot([a[1:32],real.(X[1:32]),imag.(X[1:32])],ylims=(-1,1))
plot(p1,p2,p3,layout=(3,1))

##
anim = @animate for shift=range(1,tableLen,step=64)
    x = imag.(A)
    xs = circshift(x,shift)
    p1 = plot([x,xs])
    X = ifft(xs)
    p2 = plot([a[1:nHarms],abs.(X[1:nHarms])])
    p3 = plot(unwrap(angle.(X[1:nHarms])),ylim=(-nHarms*pi,nHarms*pi))
    #p2 = plot([a[1:32],real.(X[1:32]),imag.(X[1:32])],ylims=(-1,1))
    plot(p1,p2,p3,layout=(3,1),size=(600,800))
end

gif(anim,"analysis/test.gif",fps=30)


using LoopFieldCalc

# Define current loop with radius of 50 cm, current of 1 A, centered at the origin
loop = LoopFieldCalc.CurrentLoop(
  radius = 0.5,         # radius in meters
  current = 1.0,        # current in Amperes
  center = LoopFieldCalc.CartesianPoint(0.0, 0.0, 0.0)    # cartesian coordinates of the loop center
)

# Define grid points along x, y, and z axes. We have a 2D grid of points in the y-z plane.
xs = 0.0:0.0
ys = -1.0:0.01:1.0
zs = -2.0:0.01:2.0

# Define current loops. Here, we only have the single loop, defined above
loops = [loop]

# Compute the magnetic field components
Bx, By, Bz = LoopFieldCalc.field_on_grid(loops, xs, ys, zs)

LoopFieldCalc.write_field("examples/example_1.dat", ys, zs, By[1, :, :], Bz[1, :, :])


using LinearAlgebra

function E(r :: Vector{Float64}, t :: Float64) :: Vector{Float64}
    [0., 0., .01]
end

function B(r :: Vector{Float64}, t :: Float64) :: Vector{Float64}
    [0., 0., 1.]
end

function boris_step!(r :: Vector{Float64}, v :: Vector{Float64},
                     qm :: Float64, dt :: Float64, t :: Float64) :: Tuple{Vector, Vector}
    dtqm = dt * qm
    r += .5 * dt .* v
    B_rt = B(r, t)
    E_rt = E(r, t)
    p = .5 * dtqm .* B_rt
    a_sq = .25 * dtqm^2 * dot(B_rt, B_rt)
    v += .5 * dtqm .* E_rt
    v_prime = v .+ cross(v, p)
    v += 2. .* cross(v_prime, p) / (1 + a_sq)
    v += .5 * dtqm .* E_rt
    r += .5 * dt .* v
    return r, v
end

qm = 1.
dt = .25
steps = 100
r = zeros(steps, 3)
v = zeros(steps, 3)
r[1,:] = [0. 0. 0.]
v[1,:] = [.5 0. 0.]

function test!()
    t = 0.
    for i=2:steps
        t += .5 * dt
        next = boris_step!(r[i-1,:], v[i-1,:], qm, dt, t)
        r[i,:] = next[1]
        v[i,:] = next[2]
        t += .5 * dt
    end
    return r, v
end


using InteractiveDynamics, DynamicalSystems, GLMakie

ds = Systems.henonheiles()

# Grid of initial conditions at given energy:
energy(x,y,px,py) = 0.5(px^2 + py^2) + potential(x,y)
potential(x, y) = 0.5(x^2 + y^2) + (x^2*y - (y^3)/3)
function generate_ics(E, n)
    ys = range(-0.4, stop = 1.0, length = n)
    pys = range(-0.5, stop = 0.5, length = n)
    ics = Vector{Vector{Float64}}()
    for y in ys
        V = potential(0.0, y)
        V ≥ E && continue
        for py in pys
            Ky = 0.5*(py^2)
            Ky + V ≥ E && continue
            px = sqrt(2(E - V - Ky))
            ic = [0.0, y, px, py]
            push!(ics, [0.0, y, px, py])
        end
    end
    return ics
end

density = 15
tfinal = 2000.0
tgali = 1000.0
E = 0.13
ics = generate_ics(E, density)

tinteg = tangent_integrator(ds, 4)

regularity = Float64[]; psos = Dataset{2, Float64}[]
trs = Dataset{3, Float64}[]
@time for u in ics
    # compute gali (using advanced usage)
    reinit!(tinteg, u, orthonormal(4,4))
    push!(regularity, gali(tinteg, tgali, 1, 1e-12)[2][end]/tgali)
    push!(psos, poincaresos(ds, (1, 0.0), 2000.0; u0 = u, idxs = [2, 4]))
    tr = trajectory(ds, 200.0, u)[:, [1, 2, 4]]
    push!(trs, tr)
end

# %%

trajectory_highlighter(psos, regularity; α = 0.05, hname = "regularity")

trajectory_highlighter(trs[1:10:end], regularity[1:10:end];
nbins = 10, α = 0.05, linewidth = 4.0, hname = "regularity")


scenario(nx = 60, ny = 60, nz = 60; case = :torus2) = begin
    F = Float32
    vol = zeros(F, nx, ny, nz)

    sx, sy, sz = size(vol) ./ F(16)
    tx = F(nx) / 2sx
    ty = F(ny) / 2sy + F(1.5)
    tz = F(nz) / 2sz

    r = F(1.85)
    R = F(4)

    cushin(x, y, z) =
        z^2 * x^2 - z^4 - 2z * x^2 + 2z^3 + x^2 - z^2 - (x^2 - z)^2 - y^4 - 2x^2 * y^2 -
        y^2 * z^2 +
        2y^2 * z +
        y^2

    torus2(x, y, z) =
        ((x^2 + y^2 + z^2 + R^2 - r^2)^2 - 4R^2 * (x^2 + y^2)) *
        ((x^2 + (y + R)^2 + z^2 + R^2 - r^2)^2 - 4R^2 * ((y + R)^2 + z^2))

    sphere(x, y, z) = (
        ((x - 2)^2 + (y - 2)^2 + (z - 2)^2 - 1) *
        ((x + 2)^2 + (y - 2)^2 + (z - 2)^2 - 1) *
        ((x - 2)^2 + (y + 2)^2 + (z - 2)^2 - 1)
    )

    plane(x, y, z) = x + y + z - 3

    cassini(x, y, z) =
        (x^2 + y^2 + z^2 + F(0.45)^2)^2 - 16 * F(0.45)^2 * (x^2 + z^2) - F(0.5)^2

    blooby(x, y, z) = x^4 - 5x^2 + y^4 - 5y^2 + z^4 - 5z^2 + F(11.8)

    hyperboloid(x, y, z) = x^2 + y^2 - z^2 - 1

    callback = (;
        cushin = cushin,
        torus2 = torus2,
        sphere = sphere,
        plane = plane,
        cassini = cassini,
        blooby = blooby,
        hyperboloid = hyperboloid,
    )[case]

    for k ∈ 1:nz, j ∈ 1:ny, i ∈ 1:nx
        vol[i, j, k] = callback((i - 1) / sx - tx, (j - 1) / sy - ty, (k - 1) / sz - tz)
    end

    MC(vol, Int32)
end

output(PlyIO::Module, m::MC, fn::AbstractString = "test.ply") = begin
    ply = PlyIO.Ply()
    push!(
        ply,
        PlyIO.PlyElement(
            "vertex",
            PlyIO.ArrayProperty("x", Float32[v[1] for v ∈ m.vertices]),
            PlyIO.ArrayProperty("y", Float32[v[2] for v ∈ m.vertices]),
            PlyIO.ArrayProperty("z", Float32[v[3] for v ∈ m.vertices]),
            PlyIO.ArrayProperty("nx", Float32[n[1] for n ∈ m.normals]),
            PlyIO.ArrayProperty("ny", Float32[n[2] for n ∈ m.normals]),
            PlyIO.ArrayProperty("nz", Float32[n[3] for n ∈ m.normals]),
        ),
    )
    vertex_indices = PlyIO.ListProperty("vertex_indices", UInt8, Int32)
    for i ∈ eachindex(m.triangles)
        push!(vertex_indices, m.triangles[i] .- 1)
    end
    push!(ply, PlyIO.PlyElement("face", vertex_indices))

    PlyIO.save_ply(ply, fn, ascii = true)
    return
end


starting_numbers = open("starting_numbers.txt") do f
    parse.(Int, split(readline(f), ","))
end

function play_game(starting::Vector{Int64}, upto::Int64)::Int64
    last_seen = Dict{Int64, Tuple{Int64, Int64}}()
    prev = starting[end]

    for (i, v) = enumerate(starting)
        last_seen[v] = (i, 0)
        # println("Turn $(i): $(v)")
    end

    for i in (length(starting)+1):upto
        if last_seen[prev][2] == 0
            x = 0
        else
            x = i - 1 - last_seen[prev][2]
            # println("$(prev): $(i-1) - $(recent)")
        end
        # println("Turn $(i): $(x)")
        if haskey(last_seen, x)
            last_seen[x] = (i, last_seen[x][1])
        else
            last_seen[x] = (i, 0)
        end
        prev = x
    end

    return prev
end

# println(play_game([0,3,6], 10) == 0)
# println(play_game([0,3,6], 2020) == 436)
# println(play_game([1,3,2], 2020) == 1)
# println(play_game([2,1,3], 2020) == 10)
# println(play_game([1,2,3], 2020) == 27)
# println(play_game([2,3,1], 2020) == 78)
# println(play_game([3,2,1], 2020) == 438)
# println(play_game([3,1,2], 2020) == 1836)

@time println(play_game(starting_numbers, 2020))
@time println(play_game(starting_numbers, 30000000))


#------------------------------
# Non-linear ODE system with two unkown parameters
#------------------------------

#------------------------------
# Setup
#------------------------------

using Pkg;

# all paths are relative to the repository main folder

Pkg.activate(".")
Pkg.instantiate()
Pkg.status()

using Distributions
using Random
using Flux
using DiffEqFlux
using OrdinaryDiffEq
using Plots
using LaTeXStrings
gr()

include("../src/simulation.jl")
include("../src/model.jl")
include("../src/plotting.jl")

#------------------------------
# Define and obtain ground-truth developments
#------------------------------

true_u0 = Float32[2, 2]
tspan = (0.0f0, 10.0f0)
true_odeparams_group1 = Float32[1.0, 1.0, 1.0, 0.5]
true_odeparams_group2 = Float32[0.5, 1.0, 1.0, 2.0]

lvprob1 = ODEProblem(lotka_volterra, true_u0, tspan, true_odeparams_group1)
lvprob2 = ODEProblem(lotka_volterra, true_u0, tspan, true_odeparams_group2)

# true trajectory of each group
dt=0.1
sol_group1 = solve(lvprob1, Tsit5(), saveat = dt);
sol_group2 = solve(lvprob2, Tsit5(), saveat = dt);
#plot(sol_group1)
#plot(sol_group2)

#------------------------------
# Simulate data based on true trajectories
#------------------------------

n = 200
p = 10
q, q_info = 50, 20
baseline_simulation = "groupsonly"

Random.seed!(1234)
xs, tvals, group1, group2 = generate_xs(n, p, true_u0, sol_group1, sol_group2, t_start=0.0, maxntps=1, σ_var=0.1, σ_ind=0.1); # default vals: t_start=1.5, t_end=10, maxntps=10, dt=0.1, σ_var=0.1, σ_ind=0.5
if baseline_simulation == "trueparams"
    x_baseline = generate_baseline(n, q, q_info, group1, true_odeparams_group1, true_odeparams_group2); # defaul vals: σ_info=0.1, σ_noise=0.1
elseif baseline_simulation == "groupsonly"
    x_baseline = generate_baseline(n, q, q_info, group1,  σ_info=0.5, σ_noise=0.5); # default vals: σ_info=1, σ_noise=1
else
    error("Please select either 'trueparams' or 'groupsonly' as mode for simulation of baseline variables")
end

# look at data:
data = simdata(xs, x_baseline, tvals, group1, group2);

plot(plot_truesolution(2, data, sol_group1, sol_group2, showdata=false),
    plot_truesolution(1,  data, sol_group1, sol_group2, showdata=false),
    plot_truesolution(2, data, sol_group1, sol_group2, showdata=true),
    plot_truesolution(1,  data, sol_group1, sol_group2, showdata=true),
    layout = (2,2),
    legend = false
)

#------------------------------
# Define and train model
#------------------------------

zdim = nODEparams = 2
m = init_vae(p, q, zdim, nODEparams, lvprob2, seed=84)
L = loss_wrapper(m)
ps = getparams(m)
opt = ADAM(0.0001)
trainingdata = zip(xs, x_baseline, tvals);
evalcb() = @show(mean(L.(xs, x_baseline, tvals)))
evalcb_zs() = eval_z_trajectories(xs, x_baseline, tvals, group1, sol_group1, sol_group2, m, dt)

for epoch in 1:40
    Flux.train!(L, ps, trainingdata, opt)
    evalcb()
    evalcb_zs()
end

#------------------------------
# Look at results
#------------------------------

plot(allindsplot(2, data, m, sol_group1, sol_group2),
    allindsplot(1, data, m, sol_group1, sol_group2, showlegend=false),
    layout = (1,2)
)

# individual plots
p1 = plot_individual_solutions(46, xs, x_baseline, tvals, group1, sol_group1, sol_group2, m, dt)
p2 = plot_individual_solutions(182, xs, x_baseline, tvals, group1, sol_group1, sol_group2, m, dt, showlegend=false)

# solutions across the entire dataset
p3 = allindsplot(2, data, m, sol_group1, sol_group2)
p4 = allindsplot(1, data, m, sol_group1, sol_group2, showlegend=false)
plot(p3,p4, layout = (1,2))