#  A delay eigenvalue problem with one eigenvalue equal to one.
function dep1()
    A0=([1 2 3 ; 4 5 6; 1 -1 3]);
    A1=((-A0+[1 0 3;0 0 -1;0 0 10])*exp(1));
    Q=[1 0 3; 1 1 -4; 2 3 1];
    A0=Q\(A0*Q);
    A1=Q\(A1*Q);
    nep=DEP([A0,A1],[0,1])
    return nep;
end

# A symmetric delay eigenvalue problem with double eigenvalues
# Examle from H. Voss and M. M. Betcke, Restarting iterative projection methods for Hermitian nonlinear eigenvalue problems with minmax property, Numer. Math., 2017
function dep_symm_double(n::Integer=100)
    LL=-sparse(1:n,1:n,2*ones(n))+sparse(2:n,1:n-1,ones(n-1),n,n)+sparse(1:n-1,2:n,ones(n-1),n,n)

    x = range(0, stop = pi, length = n)
    h=x[2]-x[1];
    LL=LL/(h^2)
    LL=kron(LL,LL)

    b=broadcast((x,y)->-100*abs(sin(x+y)),x,transpose(x))
    a=broadcast((x,y)->8*sin(x)*sin(y),x,transpose(x))
    B=sparse(1:n^2,1:n^2,b[:])
    A=LL+sparse(1:n^2,1:n^2,a[:])

    nep=DEP([A,B],[0,2.0])
    return nep
end

# A delay eigenvalue problem with a double non-semisimple eigenvalue in λ=3πi
# Examle from E. Jarlebring, Convergence factors of Newton methods for nonlinear eigenvalue problems, LAA, 2012
function dep_double()
    denom = 8+5*pi;
    a1 = 2/5 *(65*pi + 32)/(denom);
    a2 = 9*pi^2*(13+5*pi)/(denom);
    a3 = 324/5 *pi^2*(5*pi+4)/(denom);
    b1 = (260*pi + 128 + 225*pi^2)/(10*denom);
    b2 = 45*pi^2/denom;
    b3 = 81*pi^2*(40*pi + 32 + 25*pi^2)/(10*denom);
    A0 = [ 0    1    0;  0    0    1;  -a3  -a2  -a1];
    A1 = [ 0    0    0;  0    0    0;  -b3  -b2  -b1];

    tau=1.0;
    nep=DEP([A0,A1],[0,tau])
    return nep
end

#  A quadratic problem with real eigenvalues
# Four smallest eigenvalues of the problem:
# 1.0e+03 * [ -2.051741417993845  -0.182101627437811  -0.039344930222838  -0.004039879577113]
function real_quadratic()
    A0 = [ 4     0     1     1;
           0     2     1     1;
           1     1     6    -2;
           1     1    -2     3];
    A1 = [ 167  -140    95  -131;
          -140   327    54    85;
            95    54   235   -81;
          -131    85   -81   181];
    A2 = [ 2     1    -1    -1;
           1     5    -3     2;
          -1    -3     3     0;
          -1     2     0     3];
    A=[A0,A1,A2]
    nep=PEP(A)
    return nep
end

# A quadratic delay eigenvalue problem in "The infinite Bi-Lanczos method for nonlinear eigenvalue problems",  Sarah W. Gaaf and Elias Jarlebring
function qdep0()
    qdepbase=joinpath(dirname(@__FILE__()),
                      "converted_misc" ,"qdep_infbilanczos_")
    A0=read_sparse_matrix(qdepbase * "A0.txt", Int64)
    A1=read_sparse_matrix(qdepbase * "A1.txt", Int64)
    tau::Float64 = 1
    quadfun = S -> S^2
    constfun = S -> one(S)
    expfun = S -> exp(-tau*S)

    AA = [-one(A0), A0, A1]
    fi = [quadfun, constfun, expfun]
    return SPMF_NEP(AA, fi)
end

# A quadratic delay eigenvalue problem in "A linear eigenvalue algorithm for the  nonlinear eigenvalue problem",  Elias Jarlebring, Wim Michiels, Karl Meerbergen \\
function qdep1()
    A0=[ 0.3000   -0.6000         0    0.4000
        -0.3000    0.4000   -0.8000    1.9000
         0.1000   -1.6000   -1.3000         0
        -1.4000   -0.9000    0.2000    0.9000];
    A1=[ 0.8000    0.2000   -1.3000   -0.3000
        -1.1000    0.9000    1.2000    0.5000
         0.5000    0.2000   -1.6000   -1.3000
         0.7000    0.4000   -0.4000         0];
    return SPMF_NEP([one(A0), A0, A1], [λ -> -λ^2, λ -> one(λ), λ -> exp(-λ)])
end

# A quadratic eigenvalue problem with chosen eigenvalues
function qep_fixed_eig(n::Integer=5, E::AbstractVecOrMat=NaN*ones(2*n))
    if any(isnan.(E))
        msws_rng = MSWS_RNG()
        E=vec(gen_rng_mat(msws_rng,2*n,1))
    end
    A1 = diagm(0 => E[1:n])
    A2 = diagm(0 => E[n+1:2*n])
    K = one(A1)
    nep=PEP([A1*A2,-A1-A2,K])
    return nep
end

# This problem stems from
# L. P. Shayer and S. A. Campbell.  Stability, bifurcation and multistability in a system of two coupled neurons with multiple time delays. SIAM J. Applied Mathematics , 61(2):673–700, 2000
# It is also a benchmark example in DDE-BIFTOOL
function neuron0()
    pars = [1/2; -1; 1; 2.34; 0.2; 0.2 ; 1.5] .+ 0im
    kappa = pars[1]
    beta = pars[2]
    A = [0 pars[3]; pars[4] 0]

    x = [0; 0]  # The zero (trivial) stationary solution
    # A non-trivial stationary solution
    #x=[3.201081590416643561697725111745656884148241428177442574927999582405266342752249e-01
    #   5.096324796647208606096018689631125587762848405395086474417800152349531876959548e-01]

    tauv = [0;0.2;0.2;1.5]
    A0 = -kappa * Matrix(1.0I, 2, 2)
    A1 = A[2,1] * [0 0; (1-tanh(x[2])^2) 0]
    A2 = A[1,2] * [0 (1-tanh(x[1])^2); 0 0]
    A3 = beta * diagm(0 => [(1-tanh(x[1])^2), (1-tanh(x[2])^2)])
    return DEP([A0, A1, A2, A3], tauv)
end

# DEP modelling a beam
function beam(n::Integer=100)
    h=1/n;
    ee = ones(n);
    A0 = spdiagm(-1 => ee[1:n-1], 0 => -2*ee, 1 => ee[1:n-1]);
    A0[end,end]=1/h;
    A0[end,end-1]=-1/h;
    A1=sparse([n],[n],[1.0]); # A1=en*en'
    tau=1.0;
    return DEP([A0,A1],[0,tau]);
end

function sine_nep()
    data_dir=joinpath(dirname(@__FILE__()),  "converted_sine")
    A0=read_sparse_matrix(joinpath(data_dir,"sine_A0.txt"), Int64);
    A1=read_sparse_matrix(joinpath(data_dir,"sine_A1.txt"), Int64);
    A2=read_sparse_matrix(joinpath(data_dir,"sine_A2.txt"), Int64);
    V=Matrix(read_sparse_matrix(joinpath(data_dir,"sine_V.txt"), Int64));
    Q=Matrix(read_sparse_matrix(joinpath(data_dir,"sine_Q.txt"), Int64));

    n=size(A0,1);
    Z=spzeros(n,n);
    pep=PEP([A0,A1,Z,Z,A2]);
    # Matrix  sine function. Note that the Term is rank two which is not exploited here
    sin_nep=SPMF_NEP([V*Q'], [S-> sin(S)]);
    nep=SPMFSumNEP(pep,sin_nep) # Note: nep has a low-rank term
    return nep;
end

# From the tutorial on NEP-PACK online user's manual (except we use LowRankNEP)
function schrodinger_movebc(n=1000,L0=1,L1=8,α=25*pi/2,V0=10.0)

    # Discreatization matrices
    xv=Vector(range(0,stop=L0,length=n))
    h=xv[2]-xv[1];
    n=size(xv,1);
    V=x->1+sin(α*x);
    Dn=spdiagm(-1 => [ones(n-2);0]/h^2, 0 => -2*ones(n-1)/h^2, 1 => ones(n-1)/h^2)
    Vn=spdiagm(0 => [V.(xv[1:end-1]);0]);
    In=spdiagm(0 => [ones(n-1);0])
    # Construct functions
    f1=S->one(S);
    f2=S->-S;
    hh=S-> sqrt(S+V0*one(S))
    g=S-> cosh((L1-L0)*hh(S))
    f=S-> inv(hh(S))*sinh((L1-L0)*hh(S))

    # First NEP
    nep1=SPMF_NEP([Dn-Vn,In],[f1,f2]);

    # Create the Low-Rank NEP
    Lv1=sparse([zeros(n-1,1);1.0]);
    Lv2=sparse([zeros(n-1,1);1.0]);
    Uv1=sparse([zeros(n-1,1);1.0]);
    Uv2=sparse([zeros(n-3,1); [1/(2*h), -2/h, 3/(2*h)]])
    nep2=LowRankFactorizedNEP([Lv1,Lv2],[Uv1,Uv2],[g,f]);

    # Create a sum of the two
    nep=SumNEP(nep1,nep2);
    return nep
end


#################################################
#
#   Type definitions and properties
#
#################################################

"""
Joint finite distribution represented by probability mass function.
"""
mutable struct JointPMF{T<:AbstractFloat}
    varlist::SVarList
    p::Array{T}

    function JointPMF{T}(vl::SVarList, p::Array{T}) where T<:AbstractFloat
        nv = length(vl)
        ndims(p) == nv ||
            throw(DimensionMismatch(
                "ndims(p) does not match the number of variables."))
        for i = 1:nv
            size(p,i) == cardinality(vl[i]) ||
                throw(DimensionMismatch(
                    "size(p,i) does not match the corresponding space size."))
        end
        new(vl, p)
    end
end

JointPMF{T<:AbstractFloat}(vl::VarList{SVar}, p::Array{T}) = JointPMF{T}(vl, p)

function JointPMF{T<:AbstractFloat}(vnames::Vector{String}, p::Array{T})
    vlst = SVar[dvar(v, size(p,i)) for (i, v) in enumerate(vnames)]
    JointPMF{T}(VarList(vlst), p)
end

function JointPMF{T<:AbstractFloat,N}(vnames::NTuple{N,String}, p::Array{T,N})
    vlst = SVar[dvar(v, size(p,i)) for (i, v) in enumerate(vnames)]
    JointPMF{T}(VarList(vlst), p)
end

nvars(jp::JointPMF) = length(jp.varlist)
vars(jp::JointPMF) = jp.varlist
probs(jp::JointPMF) = jp.p

function show{T}(io::IO, jp::JointPMF{T})
    println(io, "JointPMF{$T} with $(nvars(jp)) variables:")
    for v in jp.varlist
        println(io, "  $(v.id): $(values(v.space))")
    end
end

#################################################
#
#   Compute marginals and conditionals
#
#################################################

function marginal_i{T}(jp::JointPMF{T}, vidx::Int)
    if nvars(jp) == 1
        vidx == 1 || throw(BoundsError("Variable index out of range."))
        return copy(probs(jp))::Vector{T}
    else
        reduc_region = Utils.range_rm(nvars(jp), vidx)
        _mp = sum(probs(jp), reduc_region)
        return reshape(_mp, size(_mp, vidx))::Vector{T}
    end
end

function marginal_i{T}(jp::JointPMF{T}, vinds)
    @assert !isa(vinds, Int)
    !isempty(vinds) || throw(ArgumentError("Input indices are empty."))

    # process the case for single variable
    if length(vinds) == 1
        return marginal_i(jp, first(vinds))
    end

    # marginalize out unspecified variables
    vinds_ = collect(vinds)::Vector{Int}
    reduc_region = Utils.range_rm(nvars(jp), vinds_)
    if isempty(reduc_region)
        # no remaining variables, just keep p
        mp = copy(probs(jp))
    else
        s_vinds = issorted(vinds_) ? vinds_ : sort(vinds_)
        # reduce along unspecified dimensions
        _mp = sum(probs(jp), reduc_region)
        # squeeze out reduced dimensions
        rsiz = ntuple(i->size(_mp, s_vinds[i]), length(s_vinds))
        mp = reshape(_mp, rsiz)
    end
    # make output, permute dimensions according to specified order
    issorted(vinds)    ? mp :
    length(vinds) == 2 ? transpose(mp) :
                         permutedims(mp, Utils.isortperm(vinds_))
end

marginal(jp::JointPMF, vids) = marginal_i(jp, indexof(vids, vars(jp)))

function conditional_i(jp::JointPMF, vinds, cinds)
    # no conditions, degenerate to marginal distribution
    if isempty(cinds)
        return marginal_i(jp, vinds)
    end

    # compute the marginal of both targets and conditions
    mp = marginal_i(jp, Int[vinds..., cinds...])

    # obtain the conditional probabilities
    mp ./= sum(mp, 1:length(vinds))
end

conditional(jp::JointPMF, vs, cs) =
    conditional_i(jp, indexof(vs, vars(jp)), indexof(cs, vars(jp)))


using OrdinaryDiffEq, DiffEqSensitivity, ForwardDiff, Zygote, Test
A = [0. 1.; 1. 0.; 0 0; 0 0];
B = [1. 0.; 0. 1.; 0 0; 0 0];

utarget = A;
const T = 10.0;

function f(u, p, t)
    return -p[1]*u # just a silly example to demonstrate the issue
end

u0 = [1.0 0.0; 0.0 1.0; 0.0 0.0; 0.0 0.0];

tspan = (0.0, T)
tsteps = 0.0:T/100.0:T

p = [1.7, 1.0, 3.0, 1.0]

prob_ode = ODEProblem(f, u0, tspan, p);

fd_ode = ForwardDiff.gradient(p) do p
    sum(last(solve(prob_ode, Tsit5(),p=p,abstol=1e-12,reltol=1e-12)))
end

grad_ode = Zygote.gradient(p) do p
    sum(last(solve(prob_ode, Tsit5(),p=p,abstol=1e-12,reltol=1e-12)))
end[1]

@test fd_ode ≈ grad_ode rtol=1e-6

grad_ode = Zygote.gradient(p) do p
    sum(last(solve(prob_ode, Tsit5(),p=p,abstol=1e-12,reltol=1e-12, sensealg = InterpolatingAdjoint())))
end[1]

@test fd_ode ≈ grad_ode rtol=1e-6


# https://adventofcode.com/2020/day/10
using DataStructures
using BenchmarkTools

input = read("2020/data/day_10.txt", String)

test_input = """
28
33
18
42
31
14
46
20
48
47
24
23
49
45
19
38
39
11
1
32
25
35
8
17
7
9
4
2
34
10
3
"""

function parse_input(input)
    numbers = [0]
    inp = sort(parse.(Int, split(strip(input), "\n")))
    append!(numbers, inp)
    append!(numbers, numbers[end] + 3)
    numbers
end

function part_1(input)
    numbers = parse_input(input)
    counts = counter(numbers[2:end] - numbers[1:end - 1])
    counts[3] * counts[1]
end
# @info part_1(test_input)
@info part_1(input)

function part_2(input)
    numbers = parse_input(input)
    routes = zeros(Int, length(numbers))
    for (idx, n) in enumerate(numbers)
        if idx == 1
            routes[idx] = 1
        else
            is_connector = numbers[idx] - 3 .<= numbers .< numbers[idx]
            routes[idx] = sum(routes[is_connector])
        end
    end
    routes[end]
end
# @info part_2(test_input)
@info part_2(input)
# @btime part_2(input)


using Random
using HDF5 # ProgressMeter
# τ >= 2  period and n sites
periodicSignal(n, τ) = [i % τ ==0 ? 1.0 : 0.0 for i in 0:n-1]

function signal_plus_noise(signal, σ)
    noise = σ .* abs.(randn(length(signal)))
    signalx =  signal .+ noise
    signalx./maximum(signalx)
end
function signal_plus_noise(signal, σ, doble)
    noise = σ .* randn(length(signal))
    signalx =  signal .+ noise
    minsig = minimum(signalx)
    signalx = signalx .- minsig
    signalx ./maximum(signalx)
end

function gen_samples_signal(signal, σ, num_samples)
    set_samples = zeros(length(signal), num_samples)
    for i in 1:num_samples
        set_samples[:, i] = signal_plus_noise(signal, σ)
    end
    set_samples
end

function gen_samples_signal(signal, σ, num_samples, doble)
    set_samples = zeros(length(signal), num_samples)
    for i in 1:num_samples
        set_samples[:, i] = signal_plus_noise(signal, σ, doble)
    end
    set_samples
end

function gen_data_set(all_signals, σ, num_samples)
    set_samples = gen_samples_signal(all_signals[1, :], σ, num_samples)
    labels = fill(1, num_samples)
    for i in 2:size(all_signals)[1]
        set_samplesx = gen_samples_signal(all_signals[i, :], σ, num_samples)
        set_samples = hcat(set_samples, set_samplesx)
        labelsx = fill(i, num_samples)
        labels = vcat(labels,labelsx)
    end
    set_samples, labels
end
function gen_data_set(all_signals, σ, num_samples, doble)
    set_samples = gen_samples_signal(all_signals[1, :], σ, num_samples, doble)
    labels = fill(1, num_samples)
    for i in 2:size(all_signals)[1]
        set_samplesx = gen_samples_signal(all_signals[i, :], σ, num_samples, doble)
        set_samples = hcat(set_samples, set_samplesx)
        labelsx = fill(i, num_samples)
        labels = vcat(labels,labelsx)
    end
    set_samples, labels
end

function gen_signals(n_sites, periods)
    all_signals = zeros(length(periods), n_sites)
    i = 1
    for τ in periods
        all_signals[i, :] = periodicSignal(n_sites, τ)
        i += 1
    end
    all_signals
end


#Demo 3: Monochrome 2x2 array of subplots
#-------------------------------------------------------------------------------
using InspectDR
using Colors

#==Input
===============================================================================#
#Constants
#-------------------------------------------------------------------------------
black = RGB24(0, 0, 0)
gray50 = RGB24(.5, .5, .5)

#Input data
#-------------------------------------------------------------------------------
x=collect(-10:0.1:10)
titles = ["Linear", "Quadratic", "Cubic", "Quartic"]

linetypes = []
for lstyle in [:solid, :dash, :dot, :dashdot]
	push!(linetypes, line(color=black, width=3, style=lstyle))
	push!(linetypes, line(color=gray50, width=3, style=lstyle))
end

#==Generate plot
===============================================================================#
mplot = InspectDR.Multiplot(title="Powers of X")
mplot.ncolumns = 2
xlabel = "X-Axis (X-Unit)"
ylabel = "Y-Axis (Y-Unit)"

plotlist = InspectDR.Plot
for i in 1:4
	plot = add(mplot, InspectDR.Plot2D) #Generate empty plot
	plot.layout.legend.enabled = true
	plot.layout.legend.width = 80
	a = plot.annotation
		a.xlabel = xlabel
		a.ylabel = ylabel
		a.title = titles[i]

	for scalei in 1:8
#		id = "$scalei(x/10)^$i"
		id = "ref×$scalei"
		wfrm = add(plot, x, scalei*((x/10).^i), id=id)
		wfrm.line = linetypes[scalei]
	end
end

#Display
#-------------------------------------------------------------------------------
gplot = display(InspectDR.GtkDisplay(), mplot)

#==Save multi-plot to file
===============================================================================#
InspectDR.write_png("export_monochrome.png", mplot)
InspectDR.write_svg("export_monochrome.svg", mplot)

:DONE


@testset "GeodesicComplex" begin

    @testset "Construction & data structure properties" begin
        for (pts, metric, r) in datasets
            lmks, cover = GeodesicComplexes.getcover(nothing, pts, r, KDTree(pts, metric))
            @test all(1 .≤ lmks .≤ length(pts))
            @test reduce(union, cover) == Set(1:length(pts))
            @test issorted(lmks)
            @test all(eachindex(lmks)) do i
                all(evaluate(metric, pts[lmks[i]], pts[j]) ≤ r for j in cover[i])
            end
            gc = GeodesicComplex(pts, r, metric = metric)

            # TODO: indexing
            lmrks = landmarks(gc)
            othrs = nonlandmarks(gc)
            @test allunique(lmrks)
            @test allunique(othrs)
            @test length(intersect(lmrks, othrs)) == 0
            @test length(lmrks) + length(othrs) == length(points(gc))
            @test Set(union(lmrks, othrs)) == Set(points(gc))
            nlandmarks(gc) == length(lmrks)

            @test all(weight(e) ≈ evaluate(metric, lmrks[src(e)], lmrks[dst(e)])
                      for e in edges(gc))
            @test all(r ≤ weight(e) ≤ 2r
                      for e in edges(gc))
        end
    end

    @testset "LightGraphs interface" begin
        for (pts, metric, r) in datasets
            gc = GeodesicComplex(pts, r, metric = metric)
            # Run some functions to see that the interface is defined properly
            @test betweenness_centrality(gc) isa Array
            @test dijkstra_shortest_paths(gc, 1) isa LightGraphs.DijkstraState
        end
    end

    @testset "Plotting" begin
        # Count number of series.
        d = Dict{Symbol, Any}()
        for (pts, metric, r) in datasets
            gc = GeodesicComplex(pts, r, metric = metric)
            # ¯\_(ツ)_/¯
            @test length(RecipesBase.apply_recipe(d, gc)) == 2
        end
    end
end


# computation using analytical expression

# define cache
struct PBM_cache{cnumtype, cdentype, cprobtype, cpredxtype, crestype}
  cnum::cnumtype
  cden::cdentype
  cprob::cprobtype
  cpredx::cpredxtype
  cres::crestype
end

# get optimal predictions for specific sample
function get_pred_opt_x_PBM!(x, distribution, p_range,PBMc)
  @unpack cnum, cden, cprob, cpredx, cres = PBMc

  cnum[1] = zero(eltype(cnum[1]))
  cden[1] = zero(eltype(cden[1]))

  # compute numerator and denominator
  for p in p_range
    cprob[1] = distribution(x, p)
    cnum[1] += cprob[1]*p
    cden[1] += cprob[1]
  end

  if cden[1] == zero(eltype(cden[1]))
    return zero(eltype(cden[1]))
  else
    return cnum[1]/cden[1]
  end
end

# get mean optimal predictions for fixed value of tuning parameter
function get_pred_opt_p_PBM!(samples, distribution, p_range, p_tar,PBMc)
  @unpack cnum, cden, cprob, cpredx, cres = PBMc

  cres[1] = zero(eltype(cres[1]))

  for x in samples
    cpredx[1] = get_pred_opt_x_PBM!(x, distribution, p_range,PBMc)
    cprob[1] = distribution(x, p_tar)

    # contribution to optimal prediction
    cres[1] += cprob[1]*cpredx[1]

    # contribution to optimal loss
    cres[2] += cprob[1]*(p_tar-cpredx[1])^2
  end
  return nothing
end

# compute optimal predictions and indicators, as well as optimal loss of PBM
function get_indicators_PBM_analytical(samples, distribution, p_range, dp)

  caches=[PBM_cache([zero(eltype(dp))], [zero(eltype(dp))], [zero(eltype(dp))],[zero(eltype(dp))], [zero(eltype(dp)), zero(eltype(dp))]) for i in 1:Threads.nthreads()]
  pred_PBM_opt = zeros(eltype(dp), length(p_range))

  # start parallel computation for sampled values of tuning parameter
  Threads.@threads for indxp in collect(1:length(p_range))
    get_pred_opt_p_PBM!(samples, distribution, p_range, p_range[indxp], caches[Threads.threadid()])
    pred_PBM_opt[indxp] = caches[Threads.threadid()].cres[1]
  end

  # sum up all individual loss contributions
  cres2 = zero(eltype(dp))
  for i in 1:Threads.nthreads()
    cres2 += caches[i].cres[2]
  end

  # compute indicator using symmetric difference quotient
  return pred_PBM_opt, map(x -> x-1,(circshift(pred_PBM_opt, -1).-circshift(pred_PBM_opt, 1))./(2*dp))[2:end-1], cres2/length(p_range)
end

# computation using neural networks

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

  pred = NN(pnn)(input)
  labels = reshape([p_range[Int(i)] for i in data[3, :]], 1, length(data[3, :]))
  freq = reshape(data[2, :], 1, length(data[2, :]))

  return sum(freq.*((pred.-labels).^2))
end

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

  pred = NN(pnn)(input)
  labels = reshape([p_range[Int(i)] for i in dataset[3, :]], 1, length(dataset[3, :]))

  return sum((pred.-labels).^2)/length(dataset[3, :])
end

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

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

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

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

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

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

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

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

# train NN in PBM based on ``stochastic'' loss function
function train_PBM_stochastic(NN,pnn, dataset, epochs, p_range, dp, opt, batchsize, batchsize_stochastic, n_batches, n_batches_train_stochastic, inputs; verbose=false, saveat=epochs)
  NN_logger = [zeros(eltype(p_range[1]), length(pnn))]
  pred_logger = [(zeros(eltype(p_range[1]), length(p_range)), zeros(eltype(p_range[1]), length(p_range)-2), zeros(eltype(p_range[1]), 1))]

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

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

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

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

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

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

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

# compute predictions, indicators, and loss value based on current NN
function predict_PBM(dataset, p_range, pnn, dp, NN, batchsize, n_batches, inputs; calc_loss=false, loss=zero(eltype(p_range[1])))

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

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

    pred = NN(pnn)(input)
    for indxp in 1:length(rand_int)
      predictions[Int(data[3, indxp])] += data[2, indxp]*pred[1, indxp]
    end

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

  if calc_loss
    loss = loss/length(p_range)
  end

  # compute indicator using symmetric difference quotient
  return predictions,map(x -> x-1,(circshift(predictions, -1).-circshift(predictions, 1))./(2*dp))[2:end-1], [loss]
end

# apply PBM using neural networks
function get_indicators_PBM_numerical(pnn, NN, dataset, epochs, p_range, dp, opt, inputs; verbose=false, trained=false, saveat=epochs, batchsize=length(dataset[1,:]), stochastic=false, n_batches_train_stochastic=10, batchsize_stochastic=length(dataset[1,:]), dataset_train=dataset)

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

  if !trained
    if stochastic
      losses, NN_logger, pred_logger = train_PBM_stochastic(NN, pnn, dataset, epochs, p_range, dp, opt, batchsize, batchsize_stochastic, n_batches, n_batches_train_stochastic, inputs, verbose=verbose, saveat=saveat)
    else
      losses, NN_logger, pred_logger = train_PBM_weighted(NN, pnn, dataset, epochs, p_range, dp, opt, batchsize, n_batches, inputs, verbose=verbose, saveat=saveat)
    end

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

  else
    predictions, indicator, loss = predict_PBM(dataset, p_range, pnn, dp, NN, batchsize, n_batches, inputs, calc_loss=true)
    losses = zeros(eltype(dp), epochs)
    NN_logger = [pnn]
    pred_logger = [(predictions, indicator, loss)]
  end

  return pred_logger, losses, NN_logger
end


export Normal

doc"""
    Normal(mean, var)

Generator of ndarray with a normal distribution.

# Arguments
* mean: Mean of the random values.
* var: Variance of the random values.
"""
struct Normal
    mean
    var

    function Normal(mean, var)
        @assert var >= 0
        new(mean, var)
    end
end

function (n::Normal)(::Type{T}, dims::Int...) where T
    r = randn(T, dims)
    r .*= sqrt(n.var)
    r .+= n.mean
    r
end


using DataFrames, CSV
using TimeSeries, Dates

function pct_change(input::AbstractVector{<:Number})
    [i == 1 ? missing : (input[i]-input[i-1])/input[i-1] for i in eachindex(input)]
end

# This is the path to the Stocks directory in your extracted folder for the data
cd("/Path/to/data/price-volume-data-for-all-us-stocks-etfs/Stocks")

master_data = DataFrame(date = Date(2020, 5, 8), symbol = "NEWSTAT", price = 69.69)
for i in 1:730
    master_data[string(i) * "_pc"] = 0.00001
end

global index = 0
for (root, dirs, files) in walkdir(".")
    for file in files
        symbol_index = findfirst(".", file)[1]
        symbol = uppercase(SubString(file, 1, (symbol_index - 1)))
        println(symbol)

        data = convert(Matrix, CSV.read(joinpath(root, file)))

        data = data[:, 1:end .!= 2]
        data = data[:, 1:end .!= 2]
        data = data[:, 1:end .!= 2]
        data = data[:, 1:end .!= 3]
        data = data[:, 1:end .!= 3]

        data = convert(DataFrame, data)
        colnames = ["date", "price"]
        names!(data, Symbol.(colnames))
        symbols = [symbol]
        for i in 2:size(data, 1)
            append!(symbols, [symbol])
        end
        data["symbol"] = symbols
        permutecols!(data, [:date, :symbol, :price])

        data.date = Date.(data.date)

        for i in 1:730
            data_add = [0.00001]
            counter = 1
            for row in eachrow(data)
                if counter == 1
                    counter += 1
                    continue
                end

                pc = 0.00001
                if (counter - i) >= 1
                    oldPrice = data[(counter - i), 3]
                    pc = (row.price - oldPrice) / oldPrice
                end
                append!(data_add, [pc])

                counter += 1
            end
            data[string(i) * "_pc"] = data_add
        end

        append!(master_data, data)

        global index += 1
        if index % 500 == 0
            println("*** Writing data ***")
            # This path should be where you want the output files to go
            CSV.write("/Path to output files/master_data_" * string(Int(floor(index / 200))) * ".csv", master_data, writeheader=true)

            global master_data = DataFrame(date = Date(2020, 5, 8), symbol = "NEWSTAT", price = 69.69)
            for i in 1:730
                master_data[string(i) * "_pc"] = 0.00001
            end
        end
    end
end


using Bravais
using Test

import StaticArrays
using LinearAlgebra
using FastClosures

debug = false

let lattice = BravaisLattice{3}([4,6,8])
    @test length(lattice) == 4*6*8
    @test size(lattice) == (4*6*8,)

    # NOTE: assumes hypercubic by default, but BravaisLattice does not
    # define any concept of nearest neighbors
    @test realspace(lattice, [1,4,3]) == [1,4,3]
end

let lattice = TriangularLattice([4,6])
    @test length(lattice) == 4*6
    @test realspace(lattice, [1,0]) == [1.0, 0]
    @test realspace(lattice, [0,1]) == [0.5, sqrt(3)/2]
    @test realspace(lattice, 7) == [1.0, 0]
end

@test length(BravaisLattice(StaticArrays.@SVector[4,6])) == 4*6
@test length(HypercubicLattice{4}([4,6,8,10])) == 4*6*8*10
@test length(TriangularLattice([4,6])) == 4*6
@test length(HoneycombLattice([3,5])) == 3*5*2
@test length(KagomeLattice([8,6])) == 8*6*3

@test_throws BoundsError checkbounds(SquareLattice([4,6]), 0)
@test_throws BoundsError checkbounds(SquareLattice([4,6]), 4*6+1)

lattices = Any[
    "Basic BravaisLattice",
    @inferred(BravaisLattice{2}([4,6])),
    @inferred(BravaisLattice(StaticArrays.@SVector[4,6])),

    "1D and effectively 1D",
    @inferred(ChainLattice([8])),
    @inferred(CubicLattice([1, 1, 8])),
    @inferred(CubicLattice([1, 8, 1])),

    "1D Twisted",
    @inferred(ChainLattice([6], Diagonal([6]), [0//1])),
    @inferred(ChainLattice([6], Diagonal([6]), [1//2])),
    @inferred(ChainLattice([6], Diagonal([6]), [1//5])),

    "2D",
    @inferred(SquareLattice([4,6])),
    @inferred(SquareLattice([6,4])),

    "2D twisted",
    @inferred(SquareLattice([4,6], Diagonal([4,6]), [1//7, 3//13])),
    @inferred(SquareLattice([6,4], Diagonal([6,4]), [1//7, 3//13])),

    "2D twisted cylinder",
    @inferred(SquareLattice([6,4], Diagonal([6,0]), [1//7, 0//1])),

    "2D twisted helix",
    @inferred(SquareLattice([6,4], [6 0; 1 4], [1//7, 3//13])),

    "3D like",
    @inferred(CubicLattice([4,2,6])),
    @inferred(HypercubicLattice{4}([4,2,6,1])),
    @inferred(HypercubicLattice(StaticArrays.@MVector[4,2,6,1])),

    "3D open twisted helix",
    @inferred(CubicLattice([4,6,8], [4 0 0; 0 6 0; 1 0 8], [1//7, 0//1, 3//13])),

    "Other lattices",
    @inferred(TriangularLattice([4,6])),
    @inferred(HoneycombLattice([3,5])),
    @inferred(KagomeLattice([8,6])),
]

for lattice in lattices
    if isa(lattice, AbstractString)
        debug && println(lattice)
        continue
    end

    len = @inferred length(lattice)
    @test size(lattice) == (len,)
    @test eachindex(lattice) == 1:len

    @test_throws BoundsError lattice[0]
    @test_throws BoundsError lattice[len+1]

    @static if VERSION >= v"1.2.0-DEV"
        # Relies on https://github.com/JuliaLang/julia/pull/27516
        @test @inferred(Nothing, findfirst(isequal(maxcoords(lattice)), lattice)) === nothing
    else
        @test findfirst(isequal(maxcoords(lattice)), lattice) === nothing
    end
    @inferred(in(@inferred(getindex(lattice, 1)),lattice))

    last = false
    for (i, site) in enumerate(lattice)
        checkbounds(lattice, i)
        @test lattice[i] == site
        @test site in lattice
        @test findfirst(isequal(site), lattice) == i
        @test last == false
        if i == length(lattice)
            last = true
        end

        @inferred wraparound(lattice, site) # FIXME: actually test something here

        @inferred realspace(lattice, site) # FIXME: actually test something here
        realspace(lattice, site, ones(Int, ndimensions(lattice))) # FIXME: actually test something here
    end
    @test last == true

    @inferred dimensions(lattice)
    d = @inferred ndimensions(lattice)
    η = @inferred twist(lattice)
    M = @inferred repeater(lattice)

    @test ishelical(lattice) == !isdiag(repeater(lattice))

    @test isbravais(lattice) == (lattice === bravais(lattice))

    # Check that translating multiple times, in total by an $\mathbf(A)_i$
    # vector, brings us back to where we started (and picking up a phase due
    # to $η_i$).
    sz = length(lattice)
    for i in 1:d
        if M[i,i] == 0
            @test_throws ArgumentError LatticeTranslationCache(lattice, i)
            continue
        end

        # Initialize our array of each index, along with a beginning
        # phase of zero
        indices = Tuple{Int, Rational{Int}}[(z, 0//1) for z in 1:sz]

        # Translate as many times as we need to in each dimension
        for j in 1:d
            if M[i,j] != 0
                ltrc = LatticeTranslationCache(lattice, j)
                for z in 1:M[i,j]
                    new_indices = Tuple{Int, Rational{Int}}[]
                    sizehint!(new_indices, sz)
                    for (s, old_η) in indices
                        newidx, wrap_η = @inferred translateη(lattice, s, j)
                        @test (newidx, wrap_η) == @inferred translateη(ltrc, s)
                        push!(new_indices, (newidx, wrap_η + old_η))
                    end
                    indices = new_indices
                end
            end
        end

        # assert that everything is as expected (i.e., overall we got the
        # identity element while picking up some phase)
        for z in 1:sz
            @test indices[z] == (z, η[i])
        end
    end

    if !isbravais(lattice)
        continue
    end

    @test transpose(primvecs(lattice)) * recivecs(lattice) ≈ Matrix(2π * I, ndimensions(lattice), ndimensions(lattice))

    # Check the momenta across the boundary conditions
    n_k_idx = @inferred nmomenta(lattice)
    @test 1:n_k_idx == @inferred eachmomentumindex(lattice)
    for site in lattice
        for k_idx in 1:n_k_idx
            for i in 1:d
                if M[i,i] != 0
                    site2 = site + vec(M[i,:])
                    exp1 = exp(im * (kdotr(lattice, k_idx, site) + 2π * η[i]))
                    exp2 = exp(im * @inferred(kdotr(lattice, k_idx, site2)))
                    idx, η_wrap = @inferred wraparoundη(lattice, site2)
                    exp3 = exp(im * (kdotr(lattice, k_idx, idx) + 2π * η_wrap))
                    @test exp1 ≈ exp2
                    @test exp1 ≈ exp3
                    for charge in (1, 3)
                        k_total = @inferred momentum(lattice, k_idx, charge)
                        exp1 = exp(im * (kdotr(lattice, k_total, site) + 2π * η[i] * charge))
                        exp2 = exp(im * kdotr(lattice, k_total, site2))
                        exp3 = exp(im * (kdotr(lattice, k_total, idx) + 2π * η_wrap * charge))
                        @test exp1 ≈ exp2
                        @test exp1 ≈ exp3
                    end
                    k = @inferred momentumspace(lattice, k_idx)
                    exp1 = exp(im * (dot(k, realspace(lattice, site)) + 2π * η[i]))
                    exp2 = exp(im * dot(k, realspace(lattice, site2)))
                    site3, wrap = @inferred wraparound_site(lattice, site2)
                    exp3 = exp(im * (dot(k, realspace(lattice, site3)) + 2π * dot(wrap, twist(lattice))))
                    @test exp1 ≈ exp2
                    @test exp1 ≈ exp3
                end
            end
        end
    end
end

let lattice = KagomeLattice([2,3])
    # try a valid wraparound
    @inferred wraparound(lattice, [4,0,0])
    # invalid basis index
    @test_throws ArgumentError wraparound(lattice, [4,0,3])
    # attempt to "translate" in direction of the basis index
    @test_throws ArgumentError translateη(lattice, 1, 3)
end

# Invalid wraparound for OBC
let lattice = ChainLattice([8], Diagonal([0]))
    @inferred wraparound(lattice, [0])
    @inferred wraparound(lattice, [7])
    @test_throws ArgumentError wraparound(lattice, [19])
    @test_throws ArgumentError wraparound(lattice, [-1])
end

function test_neighbors(lattice, pairs, neigh=Val{1})
    mypairs = Set(pairs)
    @closure neighborsη(lattice, neigh) do i, j, η
        pop!(mypairs, (i, j))
        # FIXME: check that η is correct
    end
    @test isempty(mypairs)
end

pairs = [
    (1, 2),
    (2, 3),
    (3, 4),
    (4, 5),
    (5, 6),
    (6, 1),
]
test_neighbors(ChainLattice([6], Diagonal([6]), [1//5]), pairs)
test_neighbors(SquareLattice([6, 1]), pairs)
test_neighbors(SquareLattice([1, 6]), pairs)
test_neighbors(CubicLattice([6, 1, 1]), pairs)
test_neighbors(HypercubicLattice{4}([1, 1, 6, 1]), pairs)

pairs = [
    (1, 3),
    (2, 4),
    (3, 5),
    (4, 6),
    (5, 7),
    (6, 8),
    (7, 1),
    (8, 2),
    (1, 2),
    (3, 4),
    (5, 6),
    (7, 8),
    ]
# XXX:
#test_neighbors(SquareLattice([4, 2]), pairs)
#test_neighbors(HypercubicLattice{4}([4, 1, 2, 1]), pairs)
# make sure open boundary conditions in the 2 direction is equivalent
# (unless i decide to remove this feature that doesn't double count these
# links)
test_neighbors(SquareLattice([4, 2], Diagonal([4, 0])), pairs)

# 4x3 Open
pairs = [
    (1, 4),
    (2, 5),
    (3, 6),
    (4, 7),
    (5, 8),
    (6, 9),
    (7, 10),
    (8, 11),
    (9, 12),
    (10, 1),
    (11, 2),
    (12, 3),
    (1, 2),
    (2, 3),
    (4, 5),
    (5, 6),
    (7, 8),
    (8, 9),
    (10, 11),
    (11, 12),
]
test_neighbors(SquareLattice([4, 3], Diagonal([4, 0])), pairs)

# 4x4
pairs = [
    (1, 2),
    (2, 3),
    (3, 4),
    (4, 1),
    (5, 6),
    (6, 7),
    (7, 8),
    (8, 5),
    (9, 10),
    (10, 11),
    (11, 12),
    (12, 9),
    (13, 14),
    (14, 15),
    (15, 16),
    (16, 13),
    (1, 5),
    (2, 6),
    (3, 7),
    (4, 8),
    (5, 9),
    (6, 10),
    (7, 11),
    (8, 12),
    (9, 13),
    (10, 14),
    (11, 15),
    (12, 16),
    (13, 1),
    (14, 2),
    (15, 3),
    (16, 4),
]
test_neighbors(SquareLattice([4, 4]), pairs)
test_neighbors(HypercubicLattice{4}([1, 4, 1, 4]), pairs)
test_neighbors(HypercubicLattice{4}([4, 1, 1, 4]), pairs)

#=
TEST(HypercubicLattice, NearestNeighbors_2x4) (
    std::set<std::pair<size_t, size_t> > pairs = (
        (0, 1),
        (1, 2),
        (2, 3),
        (3, 0),
        (4, 5),
        (5, 6),
        (6, 7),
        (7, 4),
        (0, 4),
        (1, 5),
        (2, 6),
        (3, 7)
    )
    test_neighbors(HypercubicLattice((2, 4)), pairs)
    test_neighbors(HypercubicLattice((1, 2, 4, 1)), pairs)
)

template <class LatticeType>
static void test_next_nearest_neighbors(const LatticeType &lattice, const std::set<std::pair<size_t, size_t> > &pairs)
(
    auto mypairs = pairs
    lattice.next_nearest_neighbors([&mypairs](size_t i, size_t j, boost::rational<int> eta) (
        const size_t erased = mypairs.erase((i, j))
        if (!erased)
            std::cerr << i << ' ' << j << '\n'
        ASSERT_EQ(erased, 1)

        // fixme: check that eta is correct
        (void) eta
#if 0
        std::cerr << eta << std::endl
#endif
    ))
    ASSERT_EQ(mypairs.size(), 0)
)

TEST(HypercubicLattice, NextNearestNeighbors_OBC) (
    std::set<std::pair<size_t, size_t> > pairs = (
        (0, 2),
        (1, 3),
        (2, 4),
        (3, 5)
    )
    Eigen::Matrix<int, 1, 1> repeater
    repeater(0, 0) = 0
    test_next_nearest_neighbors(HypercubicLattice((6), (0), repeater), pairs)
)

TEST(HypercubicLattice, NextNearestNeighbors_PBC) (
    std::set<std::pair<size_t, size_t> > pairs = (
        (0, 2),
        (1, 3),
        (2, 4),
        (3, 5),
        (4, 0),
        (5, 1)
    )
    test_next_nearest_neighbors(HypercubicLattice((6)), pairs)
)
=#

function test_neighbor_sublattices(lattice, allowed, neigh=Val{1})
    @inferred isbipartite(lattice)
    @inferred istripartite(lattice)

    @closure neighbors(lattice, neigh) do i, j, wrap
        ind1 = @inferred sublattice_index(lattice, i)
        ind2 = sublattice_index(lattice, j)
        @test ind1 in allowed
        @test ind2 in allowed
        @test ind1 != ind2
    end
end

@test isbipartite(SquareLattice([4,4]))
@test !isbipartite(SquareLattice([3,3]))
@test isbipartite(SquareLattice([3,3], Diagonal([0,0])))
@test !isbipartite(SquareLattice([4,3]))
@test !isbipartite(SquareLattice([4,1])) # if two dimensions are specified, we don't consider this bipartite.
@test !istripartite(SquareLattice([4,4]))

let lattice = SquareLattice([4, 3], Diagonal([4,0]))
    test_neighbor_sublattices(lattice, [0,1])
    @test sublattice_index(lattice, 1) == 0
    @test sublattice_index(lattice, 2) == 1
    @test sublattice_index(lattice, 3) == 0
    @test sublattice_index(lattice, 4) == 1
    @test sublattice_index(lattice, 5) == 0
    @test sublattice_index(lattice, 6) == 1
    @test sublattice_index(lattice, 7) == 0
    @test sublattice_index(lattice, 8) == 1
    @test sublattice_index(lattice, 9) == 0
    @test sublattice_index(lattice, 10) == 1
    @test sublattice_index(lattice, 11) == 0
    @test sublattice_index(lattice, 12) == 1
end

let lattice = SquareLattice([4, 3])
    @test_throws ArgumentError sublattice_index(lattice, 1)
end

let lattice = SquareLattice([4, 2])
    test_neighbor_sublattices(lattice, [0,1])
    @test sublattice_index(lattice, 1) == 0
    @test sublattice_index(lattice, 2) == 1
    @test sublattice_index(lattice, 3) == 1
    @test sublattice_index(lattice, 4) == 0
    @test sublattice_index(lattice, 5) == 0
    @test sublattice_index(lattice, 6) == 1
    @test sublattice_index(lattice, 7) == 1
    @test sublattice_index(lattice, 8) == 0
end

let lattice = ChainLattice([6])
    test_neighbor_sublattices(lattice, [0,1])
    @test sublattice_index(lattice, 1) == 0
    @test sublattice_index(lattice, 2) == 1
    @test sublattice_index(lattice, 3) == 0
    @test sublattice_index(lattice, 4) == 1
    @test sublattice_index(lattice, 5) == 0
    @test sublattice_index(lattice, 6) == 1
end

let lattice = SquareLattice([6,1])
    @test_throws ArgumentError sublattice_index(lattice, 1)
end

let lattice = TriangularLattice([3,3])
    @test !isbipartite(lattice)
    @test istripartite(lattice)
end
let lattice = TriangularLattice([4,3])
    @test !isbipartite(lattice)
    @test !istripartite(lattice)
    @test_throws ArgumentError sublattice_index(lattice, 1)
end
let lattice = TriangularLattice([4,3], Diagonal([0,3]))
    @test !isbipartite(lattice)
    @test istripartite(lattice)
end
let lattice = TriangularLattice([4,3], Diagonal([0,0]))
    @test !isbipartite(lattice)
    @test istripartite(lattice)
end

let lattice = TriangularLattice([3,3])
    test_neighbor_sublattices(lattice, [0,1,2])
    @test sublattice_index(lattice, 1) == 0
    @test sublattice_index(lattice, 2) == 1
    @test sublattice_index(lattice, 3) == 2
    @test sublattice_index(lattice, 4) == 2
    @test sublattice_index(lattice, 5) == 0
    @test sublattice_index(lattice, 6) == 1
    @test sublattice_index(lattice, 7) == 1
    @test sublattice_index(lattice, 8) == 2
    @test sublattice_index(lattice, 9) == 0
end

let lattice = HoneycombLattice([5,5])
    @test isbipartite(lattice)
    @test !istripartite(lattice)
    test_neighbor_sublattices(lattice, [0,1])
    for (i, site) in enumerate(lattice)
        @test sublattice_index(lattice, i) == sublattice_index(lattice, site) == site[end] == ((i ⊻ 1) & 1)
    end
end

let lattice = KagomeLattice([4,3])
    @test !isbipartite(lattice)
    @test istripartite(lattice)
    test_neighbor_sublattices(lattice, [0,1,2])
    for (i, site) in enumerate(lattice)
        @test sublattice_index(lattice, i) == sublattice_index(lattice, site) == site[end]
    end
end

function test_neighbor_distances(lattice, neigh=Val{1}, expected_squared=1)
    @closure neighbors(lattice, neigh) do i, j, wrap
        diff = realspace(lattice, i) - realspace(lattice, j) - realspace(bravais(lattice), wrap .* dimensions(lattice))
        dist_squared = dot(diff, diff)
        @test expected_squared ≈ dist_squared atol=1e-8
    end
end

test_neighbor_distances(ChainLattice([8]))
test_neighbor_distances(ChainLattice([8], Diagonal([0])))
test_neighbor_distances(ChainLattice([8]), Val{1}, 1)
test_neighbor_distances(ChainLattice([8], Diagonal([0])), Val{1}, 1)
test_neighbor_distances(ChainLattice([8]), Val{2}, 4)
test_neighbor_distances(ChainLattice([8], Diagonal([0])), Val{2}, 4)
test_neighbor_distances(ChainLattice([8]), Val{5}, 25)
test_neighbor_distances(ChainLattice([8], Diagonal([0])), Val{5}, 25)
test_neighbor_distances(SquareLattice([4,6]))
test_neighbor_distances(SquareLattice([4,6], Diagonal([0,0])))
test_neighbor_distances(SquareLattice([4,6]), Val{2}, 2)
test_neighbor_distances(SquareLattice([4,6], Diagonal([0,0])), Val{2}, 2)
test_neighbor_distances(TriangularLattice([4,6]))
test_neighbor_distances(TriangularLattice([4,6], Diagonal([0,0])))
test_neighbor_distances(TriangularLattice([4,6]), Val{2}, 3)
test_neighbor_distances(TriangularLattice([4,6], Diagonal([0,0])), Val{2}, 3)
test_neighbor_distances(TriangularLattice([4,6]), Val{3}, 4)
test_neighbor_distances(TriangularLattice([4,6], Diagonal([0,0])), Val{3}, 4)
test_neighbor_distances(HoneycombLattice([4,6]))
test_neighbor_distances(HoneycombLattice([4,6], Diagonal([0,0])))
test_neighbor_distances(KagomeLattice([4,6]))
test_neighbor_distances(KagomeLattice([4,6], Diagonal([0,0])))


# # Linear shell
#
# ![](linear_shell.png)
#-
# ## Introduction
#
# In this example we show how shell elements can be analyzed in Ferrite.jl. The shell implemented here comes from the book
# "The finite elment method - Linear static and dynamic finite element analysis" by Hughes (1987), and a brief description of it is
# given at the end of this tutorial.  The first part of the tutorial explains how to set up the problem.

# ## Setting up the problem
using Ferrite
using ForwardDiff

function main() #wrap everything in a function...

# First we generate a flat rectangular mesh. There is currently no built-in function for generating
# shell meshes in Ferrite, so we have to create our own simple mesh generator (see the
# function `generate_shell_grid` further down in this file).
#+
nels = (10,10)
size = (10.0, 10.0)
grid = generate_shell_grid(nels, size)

# Here we define the bi-linear interpolation used for the geometrical description of the shell.
# We also create two quadrature rules for the in-plane and out-of-plane directions. Note that we use
# under integration for the inplane integration, to avoid shear locking.
#+
ip = Lagrange{2,RefCube,1}()
qr_inplane = QuadratureRule{2,RefCube}(1)
qr_ooplane = QuadratureRule{1,RefCube}(2)
cv = CellScalarValues(qr_inplane, ip)

# Next we distribute displacement dofs,`:u = (x,y,z)` and rotational dofs, `:θ = (θ₁,  θ₂)`.
#+
dh = DofHandler(grid)
push!(dh, :u, 3, ip)
push!(dh, :θ, 2, ip)
close!(dh)

# In order to apply our boundary conditions, we first need to create some edge- and vertex-sets. This
# is done with `addedgeset!` and `addvertexset!` (similar to `addfaceset!`)
#+
addedgeset!(grid, "left",  (x) -> x[1] ≈ 0.0)
addedgeset!(grid, "right", (x) -> x[1] ≈ size[1])
addvertexset!(grid, "corner", (x) -> x[1] ≈ 0.0 && x[2] ≈ 0.0 && x[3] ≈ 0.0)

# Here we define the boundary conditions. On the left edge, we lock the displacements in the x- and z- directions, and all the rotations.
#+
ch = ConstraintHandler(dh)
add!(ch,  Dirichlet(:u, getedgeset(grid, "left"), (x, t) -> (0.0, 0.0), [1,3])  )
add!(ch,  Dirichlet(:θ, getedgeset(grid, "left"), (x, t) -> (0.0, 0.0), [1,2])  )

# On the right edge, we also lock the displacements in the x- and z- directions, but apply a precribed roation.
#+
add!(ch,  Dirichlet(:u, getedgeset(grid, "right"), (x, t) -> (0.0, 0.0), [1,3])  )
add!(ch,  Dirichlet(:θ, getedgeset(grid, "right"), (x, t) -> (0.0, pi/10), [1,2])  )

# In order to not get rigid body motion, we lock the y-displacement in one fo the corners.
#+
add!(ch,  Dirichlet(:θ, getvertexset(grid, "corner"), (x, t) -> (0.0), [2])  )

close!(ch)
update!(ch, 0.0)

# Next we define relevant data for the shell, such as shear correction factor and stiffness matrix for the material.
# In this linear shell, plane stress is assumed, ie $\\sigma_{zz} = 0 $. Therefor, the stiffness matrix is 5x5 (opposed to the normal 6x6).
#+
κ = 5/6 # Shear correction factor
E = 210.0
ν = 0.3
a = (1-ν)/2
C = E/(1-ν^2) * [1 ν 0   0   0;
                ν 1 0   0   0;
                0 0 a*κ 0   0;
                0 0 0   a*κ 0;
                0 0 0   0   a*κ]

data = (thickness = 1.0, C = C); #Named tuple

# We now assemble the problem in standard finite element fashion
#+
nnodes = getnbasefunctions(ip)
ndofs_shell = ndofs_per_cell(dh)

K = create_sparsity_pattern(dh)
f = zeros(Float64, ndofs(dh))

ke = zeros(ndofs_shell, ndofs_shell)
fe = zeros(ndofs_shell)

celldofs = zeros(Int, ndofs_shell)
cellcoords = zeros(Vec{3,Float64}, nnodes)

assembler = start_assemble(K, f)
for cellid in 1:getncells(grid)
    fill!(ke, 0.0)

    celldofs!(celldofs, dh, cellid)
    getcoordinates!(cellcoords, grid, cellid)

    #Call the element routine
    integrate_shell!(ke, cv, qr_ooplane, cellcoords, data)

    assemble!(assembler, celldofs, fe, ke)
end

# Apply BC and solve.
#+
apply!(K, f, ch)
a = K\f

# Output results.
#+
vtk_grid("linear_shell", dh) do vtk
    vtk_point_data(vtk, dh, a)
end

end; #end main functions

# Below is the function that creates the shell mesh. It simply generates a 2d-quadrature mesh, and appends
# a third coordinate (z-direction) to the node-positions.
function generate_shell_grid(nels, size)
    _grid = generate_grid(Quadrilateral, nels, Vec((0.0,0.0)), Vec(size))
    nodes = [(n.x[1], n.x[2], 0.0) |> Vec{3} |> Node  for n in _grid.nodes]
    cells = [Quadrilateral3D(cell.nodes) for cell in _grid.cells]

    grid = Grid(cells, nodes)

    return grid
end;

# ## The shell element
#
# The shell presented here comes from the book "The finite elment method - Linear static and dynamic finite element analysis" by Hughes (1987).
# The shell is a so called degenerate shell element, meaning it is based on a continuum element.
# A brief describtion of the shell is given here.

#md # !!! note
#md #     This element might experience various locking phenomenas, and should only be seen as a proof of concept.

# ##### Fiber coordinate system
# The element uses two coordinate systems. The first coordianate system, called the fiber system, is created for each
# element node, and is used as a reference frame for the rotations. The function below implements an algorthim that return the
# fiber directions, $\boldsymbol{e}^{f}_{a1}$, $\boldsymbol{e}^{f}_{a2}$ and $\boldsymbol{e}^{f}_{a3}$, at each node $a$.
function fiber_coordsys(Ps::Vector{Vec{3,Float64}})

    ef1 = Vec{3,Float64}[]
    ef2 = Vec{3,Float64}[]
    ef3 = Vec{3,Float64}[]
    for P in Ps
        a = abs.(P)
        j = 1
        if a[1] > a[3]; a[3] = a[1]; j = 2; end
        if a[2] > a[3]; j = 3; end

        e3 = P
        e2 = Tensors.cross(P, basevec(Vec{3}, j))
        e2 /= norm(e2)
        e1 = Tensors.cross(e2, P)

        push!(ef1, e1)
        push!(ef2, e2)
        push!(ef3, e3)
    end
    return ef1, ef2, ef3

end;

# ##### Lamina coordinate system
# The second coordinate system is the so called Lamina Coordinate system. It is
# created for each integration point, and is defined to be tangent to the
# mid-surface. It is in this system that we enforce that plane stress assumption,
# i.e. $\sigma_{zz} = 0$. The function below returns the rotation matrix, $\boldsymbol{q}$, for this coordinate system.
function lamina_coordsys(dNdξ, ζ, x, p, h)

    e1 = zero(Vec{3})
    e2 = zero(Vec{3})

    for i in 1:length(dNdξ)
        e1 += dNdξ[i][1] * x[i] + 0.5*h*ζ * dNdξ[i][1] * p[i]
        e2 += dNdξ[i][2] * x[i] + 0.5*h*ζ * dNdξ[i][1] * p[i]
    end

    e1 /= norm(e1)
    e2 /= norm(e2)

    ez = Tensors.cross(e1,e2)
    ez /= norm(ez)

    a = 0.5*(e1 + e2)
    a /= norm(a)

    b = Tensors.cross(ez,a)
    b /= norm(b)

    ex = sqrt(2)/2 * (a - b)
    ey = sqrt(2)/2 * (a + b)

    return Tensor{2,3}(hcat(ex,ey,ez))
end;

# ##### Geometrical description
# A material point in the shell is defined as
# ```math
# \boldsymbol x(\xi, \eta, \zeta) = \sum_{a=1}^{N_{\text{nodes}}} N_a(\xi, \eta) \boldsymbol{\bar{x}}_{a} + ζ \frac{h}{2} \boldsymbol{\bar{p}_a}
# ```
# where $\boldsymbol{\bar{x}}_{a}$ are nodal positions on the mid-surface, and $\boldsymbol{\bar{p}_a}$ is an vector that defines the fiber direction
# on the reference surface. $N_a$ arethe shape functions.
#
# Based on the defintion of the position vector, we create an function for obtaining the Jacobian-matrix,
# ```math
# J_{ij} = \frac{\partial x_i}{\partial \xi_j},
# ```
function getjacobian(q, N, dNdξ, ζ, X, p, h)

    J = zeros(3,3)
    for a in 1:length(N)
        for i in 1:3, j in 1:3
            _dNdξ = (j==3) ? 0.0 : dNdξ[a][j]
            _dζdξ = (j==3) ? 1.0 : 0.0
            _N = N[a]

            J[i,j] += _dNdξ * X[a][i]  +  (_dNdξ*ζ + _N*_dζdξ) * h/2 * p[a][i]
        end
    end

    return (q' * J) |> Tensor{2,3,Float64}
end;

# ##### Strains
# Small deformation is assumed,
# ```math
# \varepsilon_{ij}= \frac{1}{2}(\frac{\partial u_{i}}{\partial x_j} + \frac{\partial u_{j}}{\partial x_i})
# ```
# The displacement field is calculated as:
# ```math
# \boldsymbol u = \sum_{a=1}^{N_{\text{nodes}}} N_a \bar{\boldsymbol u}_{a} +
#  N_a ζ\frac{h}{2}(\theta_{a2} \boldsymbol e^{f}_{a1} - \theta_{a1} \boldsymbol e^{f}_{a2})
#
# ```
# The gradient of the displacement (in the lamina coordinate system), then becomes:
# ```math
# \frac{\partial u_{i}}{\partial x_j} = \sum_{m=1}^3 q_{im} \sum_{a=1}^{N_{\text{nodes}}} \frac{\partial N_a}{\partial x_j} \bar{u}_{am} +
#  \frac{\partial(N_a ζ)}{\partial x_j} \frac{h}{2} (\theta_{a2} e^{f}_{am1} - \theta_{a1} e^{f}_{am2})
# ```
function strain(dofvec::Vector{T}, N, dNdx, ζ, dζdx, q, ef1, ef2, h) where T

    u = reinterpret(Vec{3,T}, dofvec[1:12])
    θ = reinterpret(Vec{2,T}, dofvec[13:20])

    dudx = zeros(T, 3, 3)
    for m in 1:3, j in 1:3
        for a in 1:length(N)
            dudx[m,j] += dNdx[a][j] * u[a][m] + h/2 * (dNdx[a][j]*ζ + N[a]*dζdx[j]) * (θ[a][2]*ef1[a][m] - θ[a][1]*ef2[a][m])
        end
    end

    dudx = q*dudx
    ε = [dudx[1,1], dudx[2,2], dudx[1,2]+dudx[2,1], dudx[2,3]+dudx[3,2], dudx[1,3]+dudx[3,1]]
    return ε

end;

# ##### Main element routine
# Below is the main routine that calculates the stiffness matrix of the shell element.
# Since it is a so called degenerate shell element, the code is similar to that for an standard continuum element.
function integrate_shell!(ke, cv, qr_ooplane, X, data)
    nnodes = getnbasefunctions(cv)
    ndofs = nnodes*5
    h = data.thickness

    #Create the directors in each node.
    #Note: For a more general case, the directors should
    #be input parameters for the element routine.
    p = zeros(Vec{3}, nnodes)
    for i in 1:nnodes
        a = Vec{3}((0.0, 0.0, 1.0))
        p[i] = a/norm(a)
    end

    ef1, ef2, ef3 = fiber_coordsys(p)

    for iqp in 1:getnquadpoints(cv)

        dNdξ = cv.dNdξ[:,iqp]
        N = cv.N[:,iqp]

        for oqp in 1:length(qr_ooplane.weights)

            ζ = qr_ooplane.points[oqp][1]

            q = lamina_coordsys(dNdξ, ζ, X, p, h)
            J = getjacobian(q, N, dNdξ, ζ, X, p, h)
            Jinv = inv(J)

            dζdx = Vec{3}((0.0, 0.0, 1.0)) ⋅ Jinv
            dNdx = [Vec{3}((dNdξ[i][1], dNdξ[i][2], 0.0)) ⋅ Jinv for i in 1:nnodes]

            #For simplicity, use automatic differentiation to construct the B-matrix from the strain.
            B = ForwardDiff.jacobian(
                (a) -> strain(a, N, dNdx, ζ, dζdx, q, ef1, ef2, h), zeros(Float64, ndofs) )

            dV = det(J) * cv.qr_weights[iqp] * qr_ooplane.weights[oqp]
            ke .+= B'*data.C*B * dV
        end
    end
end;

# Run everything:
main()


"""
    BHMM([a::AbstractVector{T}, ]A::AbstractMatrix{T}, B::AbstractVector{<:Categorical}) where F where T
Build an BHMM with transition matrix `A` and observations distributions `B`.
If the initial state distribution `a` is not specified, a uniform distribution is assumed.
Observations distributions can be of different types (for example `Normal` and `Exponential`).
However they must be of the same dimension (all scalars or all multivariates).
# Example
```julia
hmm = BHMM([0.9 0.1; 0.1 0.9], [Normal(0,1), Normal(10,1)])
```
"""
struct BHMM{T} <: AbstractHMM{Univariate}
    a::AbstractVector{T}
    A::AbstractMatrix{T}
    B::AbstractVector{<:Categorical}
    partition::String
    BHMM{T}(a, A, B, partition="") where {T} = assert_BHMM(a, A, B) && new(a, A, B, partition)
end

BHMM(a::AbstractVector{T}, A::AbstractMatrix{T}, B::AbstractVector{<:Categorical}, partition ...) where {T} = BHMM{T}(a, A, B, partition ...)
BHMM(A::AbstractMatrix{T}, B::AbstractVector{<:Categorical}, partition...) where {T} = BHMM{T}(ones(size(A)[1])/size(A)[1], A, B, partition...)

copy(hmm::BHMM) = BHMM(copy(hmm.a), copy(hmm.A), copy(hmm.B), hmm.partition)
size(hmm::BHMM) = (length(hmm.B), length(hmm.B[1].p))

"""
    assert_BHMM(a, A, B)
Throw an `ArgumentError` if the initial state distribution and the transition matrix rows does not sum to 1,
and if the observations distributions does not have the same dimensions.
"""
function assert_BHMM(a::AbstractVector,
                    A::AbstractMatrix,
                    B::AbstractVector{<:Categorical})
    !isprobvec(a) && throw(ArgumentError("Initial state vector a is not a valid probability vector!"))
    !istransmat(A) && throw(ArgumentError("Transition matrix A is not valid!"))

    !all([length(d.p) for d in B] .== length(B[1].p)) && throw(ArgumentError("All distributions must have the same dimensions"))
    !(length(a) == size(A,1) == length(B)) && throw(ArgumentError("Length of initial state vector a, dimension of transition matrix A, and number of distributions B are not the same!"))
    return true
end

function Base.show(io::IO, hmm::BHMM)
    println(io, "Background BHMM")
    println(io, "State Initial and Transition Probabilities")
    print(io, "a: ")
    show(io, hmm.a)
    println(io)
    print(io, "A: ")
    display(hmm.A)
    println(io)
    println(io, "INFORMATIVE SYMBOLS BY STATE")
    for (n,d) in enumerate(hmm.B)
        print(io, "K$n ")
        print_emitters(d)
    end
end


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

# %% [markdown]
# https://twitter.com/iitenki_moruten/status/1467683477474930688

# %%
# Original: https://github.com/moruten/julia-code/blob/1dafca3e1a4e3b36445c2263e440f6e4056b90aa/2021-12-6-test-no1.ipynb

using Plots
using Random
using Distributions
using QuadGK

#============関数定義============================================#
function action(x)
    0.5*x*x
end

function deriv_action(x)
    x
end

function hamiltonian(x,p)
    action(x) + 0.5*p*p
end

function HMC_update(x,Nt,dt)
    #backup
    x_old = x
    p     = rand(Normal(0,1))

    #check
    H_ini = hamiltonian(x,p)
    x     = molecular_dynamics!(x,p,Nt,dt)
    H_fin = hamiltonian(x,p)

    r  = rand()
    ΔH = H_fin-H_ini
    if r < exp(-ΔH)
        return x,1 #accept
    else
        return x_old,0
    end
end

function molecular_dynamics!(x,p,Nt,dt)
    force = 0.0
    #1/2step
        x = x + p * 0.5*dt
    #1~Nt-1 step
        for j=1:(Nt-1)
            force = deriv_action(x)
            p = p - dt * force
            x = x + dt * p
        end
    #Nt step
        p = p - dt * deriv_action(x)
        x = x + 0.5 * dt * p
        return x
end
#============関数終わり============================================#

#============計算=========================================================================#
  #セットアップ
    Ntest        = 300000
    Nt           = 20
    dt           = 1.0/Nt
    conf_vec     = zeros(Ntest)
    accept_count = 0
    ret          = 0
    x            = 0.0

    sumxx        = 0.0
    sumx         = 0.0

  #計算
    for i=1:Ntest
        x,ret         = HMC_update(x,Nt,dt)
        accept_count += ret
        conf_vec[i]   = x

        sumx         += x
        sumxx        += x*x
    end

    println("P(accept) = $(accept_count/Ntest)")
    println("<x>       = $(sumx/Ntest)")
    println("<xx>      = $(sumxx/Ntest)")

#=======確認=============================================================================#
    xr = range(-5,5,length=1000)
    f1(x) = exp(-0.5*x^2)
    f2(x) = exp(-x^2)
    Z1,error1 = quadgk(f1,-Inf,Inf)
    Z2,error2 = quadgk(f2,-Inf,Inf)
    g1(x) = f1(x)/Z1
    g2(x) = f2(x)/Z2
    histogram(conf_vec,norm=:true,label="data")
    plot!(xr,[g1.(xr) g2.(xr)],lw=3,label=["exp(-0.5x^2)/Z1" "exp(-x^2)/Z2"])
#========================================================================================#

# %%
# Revised 1

using Plots
using Random
using Distributions
using QuadGK

#============関数定義============================================#
function action(x)
    0.5*x*x
end

function deriv_action(x)
    x
end

function hamiltonian(x,p)
    action(x) + 0.5*p*p
end

function HMC_update(x,Nt,dt)
    #backup
    x_old = x
    p     = rand(Normal(0,1))

    #check
    H_ini = hamiltonian(x,p)
    x, p = molecular_dynamics!(x,p,Nt,dt)    # <========== ここ
    H_fin = hamiltonian(x,p)

    r  = rand()
    ΔH = H_fin-H_ini
    if r < exp(-ΔH)
        return x,1 #accept
    else
        return x_old,0
    end
end

function molecular_dynamics!(x,p,Nt,dt)
    force = 0.0
    #1/2step
        x = x + p * 0.5*dt
    #1~Nt-1 step
        for j=1:(Nt-1)
            force = deriv_action(x)
            p = p - dt * force
            x = x + dt * p
        end
    #Nt step
        p = p - dt * deriv_action(x)
        x = x + 0.5 * dt * p
        return x, p                          # <========== ここ
end
#============関数終わり============================================#

#============計算=========================================================================#
  #セットアップ
    Ntest        = 300000
    Nt           = 20
    dt           = 1.0/Nt
    conf_vec     = zeros(Ntest)
    accept_count = 0
    ret          = 0
    x            = 0.0

    sumxx        = 0.0
    sumx         = 0.0

  #計算
    for i=1:Ntest
        x,ret         = HMC_update(x,Nt,dt)
        accept_count += ret
        conf_vec[i]   = x

        sumx         += x
        sumxx        += x*x
    end

    println("P(accept) = $(accept_count/Ntest)")
    println("<x>       = $(sumx/Ntest)")
    println("<xx>      = $(sumxx/Ntest)")

#=======確認=============================================================================#
    xr = range(-5,5,length=1000)
    f1(x) = exp(-0.5*x^2)
    f2(x) = exp(-x^2)
    Z1,error1 = quadgk(f1,-Inf,Inf)
    Z2,error2 = quadgk(f2,-Inf,Inf)
    g1(x) = f1(x)/Z1
    g2(x) = f2(x)/Z2
    histogram(conf_vec,norm=:true,label="data")
    plot!(xr,[g1.(xr) g2.(xr)],lw=3,label=["exp(-0.5x^2)/Z1" "exp(-x^2)/Z2"])
#========================================================================================#

# %%
# Revised 1 - second test

using Plots
using Random
using Distributions
using QuadGK

#============関数定義============================================#
function action(x)
    3(x^2 - 1)^2
end

function deriv_action(x)
    6x*(x^2 - 1)
end

function hamiltonian(x,p)
    action(x) + 0.5*p*p
end

function HMC_update(x,Nt,dt)
    #backup
    x_old = x
    p     = rand(Normal(0,1))

    #check
    H_ini = hamiltonian(x,p)
    x, p = molecular_dynamics!(x,p,Nt,dt)     # <========== ここ
    H_fin = hamiltonian(x,p)

    r  = rand()
    ΔH = H_fin-H_ini
    if r < exp(-ΔH)
        return x,1 #accept
    else
        return x_old,0
    end
end

function molecular_dynamics!(x,p,Nt,dt)
    force = 0.0
    #1/2step
        x = x + p * 0.5*dt
    #1~Nt-1 step
        for j=1:(Nt-1)
            force = deriv_action(x)
            p = p - dt * force
            x = x + dt * p
        end
    #Nt step
        p = p - dt * deriv_action(x)
        x = x + 0.5 * dt * p
        return x, p                    # <========== ここ
end
#============関数終わり============================================#

#============計算=========================================================================#
  #セットアップ
    Ntest        = 300000
    Nt           = 20
    dt           = 1.0/Nt
    conf_vec     = zeros(Ntest)
    accept_count = 0
    ret          = 0
    x            = 0.0

    sumxx        = 0.0
    sumx         = 0.0

  #計算
    for i=1:Ntest
        x,ret         = HMC_update(x,Nt,dt)
        accept_count += ret
        conf_vec[i]   = x

        sumx         += x
        sumxx        += x*x
    end

    println("P(accept) = $(accept_count/Ntest)")
    println("<x>       = $(sumx/Ntest)")
    println("<xx>      = $(sumxx/Ntest)")

#=======確認=============================================================================#
    xr = range(-2,2,length=1000)
    f(x) = exp(-action(x))
    Z,error1 = quadgk(f,-Inf,Inf)
    g(x) = f(x)/Z
    histogram(conf_vec,norm=:true,label="data")
    plot!(xr,g.(xr),lw=3,label="exp(-action(x))/Z", legend=:outertop)
#========================================================================================#

# %%
# Revised 2

using Plots
using Random
using Distributions
using QuadGK

#============関数定義============================================#
function action(x)
    x^2/2
end

function deriv_action(x)
    x
end

function hamiltonian(x,p)
    action(x) + 0.5*p*p
end

function HMC_update(x,Nt,dt)
    #backup
    x_old = x
    p     = rand(Normal(0,1))

    #check
    H_ini = hamiltonian(x,p)
    x, p = molecular_dynamics!(x,p,Nt,dt)     # <========== ここ
    H_fin = hamiltonian(x,p)

    r  = rand()
    ΔH = H_fin-H_ini
    if r < exp(-ΔH)
        return x,1 #accept
    else
        return x_old,0
    end
end

function molecular_dynamics!(x,p,Nt,dt)
    p -= deriv_action(x) * dt/2
    x += p * dt
    for j in 2:Nt
        p -= deriv_action(x) * dt
        x += p * dt
    end
    p -= deriv_action(x) * dt/2
    return x, p                    # <========== ここ
end
#============関数終わり============================================#

#============計算=========================================================================#
  #セットアップ
    Ntest        = 300000
    Nt           = 20
    dt           = 1.0/Nt
    conf_vec     = zeros(Ntest)
    accept_count = 0
    ret          = 0
    x            = 0.0

    sumxx        = 0.0
    sumx         = 0.0

  #計算
    for i=1:Ntest
        x,ret         = HMC_update(x,Nt,dt)
        accept_count += ret
        conf_vec[i]   = x

        sumx         += x
        sumxx        += x*x
    end

    println("P(accept) = $(accept_count/Ntest)")
    println("<x>       = $(sumx/Ntest)")
    println("<xx>      = $(sumxx/Ntest)")

#=======確認=============================================================================#
    xr = range(-5,5,length=1000)
    f(x) = exp(-action(x))
    Z,error1 = quadgk(f,-Inf,Inf)
    g(x) = f(x)/Z
    histogram(conf_vec,norm=:true,label="data",alpha=0.3,bin=100)
    plot!(xr,g.(xr),lw=3,label="exp(-action(x))/Z", legend=:outertop)
#========================================================================================#

# %%
# Revised 2 - second test

using Plots
using Random
using Distributions
using QuadGK

#============関数定義============================================#
function action(x)
    3(x^2 - 1)^2
end

function deriv_action(x)
    6x*(x^2 - 1)
end

function hamiltonian(x,p)
    action(x) + 0.5*p*p
end

function HMC_update(x,Nt,dt)
    #backup
    x_old = x
    p     = rand(Normal(0,1))

    #check
    H_ini = hamiltonian(x,p)
    x, p = molecular_dynamics!(x,p,Nt,dt)     # <========== ここ
    H_fin = hamiltonian(x,p)

    r  = rand()
    ΔH = H_fin-H_ini
    if r < exp(-ΔH)
        return x,1 #accept
    else
        return x_old,0
    end
end

function molecular_dynamics!(x,p,Nt,dt)
    p -= deriv_action(x) * dt/2
    x += p * dt
    for j in 2:Nt
        p -= deriv_action(x) * dt
        x += p * dt
    end
    p -= deriv_action(x) * dt/2
    return x, p                    # <========== ここ
end
#============関数終わり============================================#

#============計算=========================================================================#
  #セットアップ
    Ntest        = 300000
    Nt           = 20
    dt           = 1.0/Nt
    conf_vec     = zeros(Ntest)
    accept_count = 0
    ret          = 0
    x            = 0.0

    sumxx        = 0.0
    sumx         = 0.0

  #計算
    for i=1:Ntest
        x,ret         = HMC_update(x,Nt,dt)
        accept_count += ret
        conf_vec[i]   = x

        sumx         += x
        sumxx        += x*x
    end

    println("P(accept) = $(accept_count/Ntest)")
    println("<x>       = $(sumx/Ntest)")
    println("<xx>      = $(sumxx/Ntest)")

#=======確認=============================================================================#
    xr = range(-2,2,length=1000)
    f(x) = exp(-action(x))
    Z,error1 = quadgk(f,-Inf,Inf)
    g(x) = f(x)/Z
    histogram(conf_vec,norm=:true,label="data",alpha=0.3,bin=100)
    plot!(xr,g.(xr),lw=3,label="exp(-action(x))/Z", legend=:outertop)
#========================================================================================#

# %%


# coding: utf-8

# In[1]:

#explanation can be found in dynamic programming version
# https://github.com/je-suis-tm/recursion/blob/master/edit%20distance%20dynamic%20programming.jl
#the only problem with recursion is that it is so freaking slow
#recursion is so inefficient in any programming language
#although it looks much more elegant than dynamic programming

# In[2]:

function edit_distance(text1,text2)

    if isempty(text1) || isempty(text2)

        return max(length(text1),length(text2))

    end

    #we are comparing characters here
    #to get string, we should do end:end
    if text1[end]==text2[end]

        replacement=0

    else

        replacement=1

    end

    steps=min(edit_distance(text1[1:end-1],text2)+1,
        edit_distance(text1,text2[1:end-1])+1,
        edit_distance(text1[1:end-1],text2[1:end-1])+replacement)

    return steps

end

# In[3]:

println(edit_distance("arsehole","asshoe"))


abstract type Transform end

Base.broadcastable(tr::Transform) = Ref(tr)

struct ComposedTransform{TR1<:Transform,TR2<:Transform} <: Transform
  """
  Outer transform.
  """
  t1::TR1
  """
  Inner transform.
  """
  t2::TR2
end
(t::ComposedTransform)(p::AbstractVector) = t.t1(t.t2(p))

AbstractTrees.children(tr::ComposedTransform) = (tr.t1, tr.t2)

transforms(tr::ComposedTransform) = Leaves(tr)

Base.:∘(t1::Transform, t2::Transform) = ComposedTransform(t1, t2)

inv(t::ComposedTransform) = inv(t.t2) ∘ inv(t.t1)

struct Transformed{O,TR<:Transform}
  obj::O
  transf::TR
  Transformed{O,TR}(obj::O, args...) where {O,TR<:Transform} = Transformed{O,TR}(obj, TR(args...))
  Transformed{O,TR}(obj::O, transf::TR) where {O,TR<:Transform} = new{O,TR}(obj, transf)
  Transformed{TR1,TR2}(obj::TR1, transf::TR2) where {TR1<:Transformed,TR2<:Transform} = Transformed(obj.obj, transf ∘ obj.transf)
  Transformed(obj::O, transf::TR) where {O,TR<:Transform} = Transformed{O,TR}(obj, transf)
end
(tr::Transformed)(p) = tr.obj(inv(tr.transf)(p))

(≈)(x::Transformed, y::Transformed) = typeof(x) == typeof(y) && x.obj ≈ y.obj && x.transf ≈ y.transf

include("transforms/scaling.jl")
include("transforms/translation.jl")
include("transforms/rotation.jl")


import MPI

include("types.jl")
include("init.jl")
include("get_sim_params.jl")
include("update.jl")
include("io.jl")
include("get_sim_params.jl")

function main()

#temp1_init: temperature init on borders
temp1_init = 10.0
#temp2_init: temperature init inside domain
temp2_init = -10.0
#diffusion coefficient
k0 = Float64(1)

#define rank of root process
root = 0

#MPI Initialization
MPI.Init()
comm = MPI.COMM_WORLD
my_id = MPI.Comm_rank(comm)
nproc = MPI.Comm_size(comm)

#simulation parameters
params_int = Array{Int64}(undef, 5)
params_double = Array{Float64}(undef, 2)
if (my_id == root)
    params_int, params_double, output_path = get_sim_params(ARGS)
end

MPI.Barrier(comm)
MPI.Bcast!(params_int, 5, 0, comm)
MPI.Bcast!(params_double, 2, 0, comm)

size_x    = params_int[1]
size_y    = params_int[2]
nx_domains = params_int[3]
ny_domains = params_int[4]
maxStep   = params_int[5]
dt1       = params_double[1]
epsilon   = params_double[2]

#Warning message if dimensions and number of processes don't match
if ((my_id == root) && (nproc != (nx_domains * ny_domains)))
    println("ERROR - Number of processes not equal to number of subdomains")
end

#Various other variables
size_global_x = size_x + 2
size_global_y = size_y + 2
hx = Float64(1.0 / size_global_x)
hy = Float64(1.0 / size_global_y)
dt2 = 0.25 * min(hx, hy)^2 / k0 #fraction of CFL condition
size_total_x = size_x + 2 * nx_domains + 2 #including ghost cells
size_total_y = size_y + 2 * ny_domains + 2 #including ghost cells

#Take a right time step for convergence
if (dt1 >= dt2)
    if (my_id == 0)
        println()
        println("Time step too large ==> Taking convergence criterion")
    end
    dt = dt2
else
    dt = dt1
end

#2D solution including ghost cells
u0 = twoDArray(Float64, size_total_x, size_total_y)
u = twoDArray(Float64, size_total_x, size_total_y)

#Allocate coordinates of processes (start cell, end cell)
xs = oneDArray(Int, nproc)
xe = oneDArray(Int, nproc)
ys = oneDArray(Int, nproc)
ye = oneDArray(Int, nproc)

#Size of each physical domain
xcell = Int(size_x / nx_domains)
ycell = Int(size_y / ny_domains)

#allocate flattened (1D) local physical solution (i.e., relative to sub-domain)
u_local = oneDArray(Float64, xcell * ycell)
#allocate flattened (1D) global physical solution
u_global = oneDArray(Float64, size_x * size_y)

#find processes surrounding mine (if they exist)
my_neighbors = neighbors(my_id, nproc, nx_domains, ny_domains)

#compute coordinates of processes for each sub-domain
process_coordinates!(xs, ys, xe, ye, xcell, ycell, nx_domains, ny_domains, nproc)

#initialize domain
init_values(u0, size_total_x, size_total_y, temp1_init, temp2_init)

#update ghost cells
updateBound!(u0, size_total_x, size_total_y, my_neighbors, comm,
            my_id, xs, ys, xe, ye, xcell, ycell, nproc)

#Initialize step, time and convergence boolean
step = 0
t = 0.0
converged = false

#Starting time
if (my_id==0)
    time_init = time()
end

#Main loop : until convergence
while (!converged)
    #increment step and time
    step += 1
    t += dt
    #perform one step of the explicit scheme
    local_diff = computeNext!(u0, u, size_total_x, size_total_y, dt, hx, hy,
        my_id, xs, ys, xe, ye, nproc, k0)

    #update ghost cells
    updateBound!(u0, size_total_x, size_total_y, my_neighbors, comm,
        my_id, xs, ys, xe, ye, xcell, ycell, nproc)

    MPI.Barrier(comm)

    #sum local_diff to get global difference
    global_diff = MPI.Allreduce(local_diff, MPI.SUM, comm)
    global_diff = sqrt(global_diff)
    #break if convergence reached or step greater than maxStep
    if ((global_diff <= epsilon) || (step >= maxStep))
        converged = true
    end
end

#get ending time
if (my_id == 0)
    time_final = time()
    #elapsed time
    elapsed_time = time_final - time_init
    println("Elapsed time = ", elapsed_time)
    println("Steps = ", step)
end

#find solution on my sub-domain (as a 1-dimensional array)
i = 1
for j = ys[my_id+1]:ye[my_id+1]
    u_local[(i-1)*xcell+1:i*xcell] = u0[xs[my_id+1]:xe[my_id+1],j]
    i = i+1
end

#gather local solution to global solution (as a 1-dimensional array)
u_global = MPI.Gather(u_local, root, comm)

#write to disk
if (my_id == 0)
    write_to_disk(u_global, nx_domains, ny_domains, xcell, ycell, temp1_init, output_path)
end

MPI.Finalize()

end

main()


# see: http://www.novaroma.org/via_romana/numbers.html

function digits(number)
    digits_list = Int[]
    m = 10
    while (m / 10) <= number
        digit = number % m
        push!(digits_list, digit / (m / 10))
        number -= digit
        m *= 10
    end
    return digits_list
end

function to_roman(number)
    if number <= 0 || number > 3000
        error("number cannot be converted")
    end

    roman = [
        ['I', 'V'], # ones
        ['X', 'L'], # test
        ['C', 'D'], # hundreads
        ['M'],      # thousands
    ]

    digits_list = digits(number)
    r = ""

    for i in 1:length(digits_list)
        d = digits_list[i]
        if 1 <= d <= 3
            for _ in 1:d
                r *= roman[i][1]
            end
        elseif d == 4
            r *= roman[i][2] * roman[i][1]
        elseif 5 <= d <= 8
            for _ in 1:(d - 5)
                r *= roman[i][1]
            end
            r *= roman[i][2]
        elseif d == 9
            r *= roman[i+1][1] * roman[i][1]
        end
    end

    return reverse(r)
end


module ShaleDrillingLikelihood_ThreadUtils_Test

DOBTIME = false

using ShaleDrillingLikelihood

using Test
using StatsFuns
using Random
using BenchmarkTools
using Base.Threads

using Base.Threads: nthreads, @threads, Atomic, atomic_add!

using ShaleDrillingLikelihood: getrange,
    Mapper,
    Mapper2,
    AbstractThreadMapper,
    getrange,
    nextrange!,
    stop,
    step,
    add_iter!

function h2(x)
    n = length(x)
    mapper = Mapper2(n)
    ld = stop(mapper)
    batch_size = step(mapper)

    s = Atomic{Float64}(0.0)
    @threads for j in 1:nthreads()
        while true
            k = add_iter!(mapper)
            batch_start = 1 + (k-1) * batch_size
            batch_end = min(k * batch_size, ld)
            batch_start > ld && break
            y = 0.0
            @inbounds @simd for i in batch_start:batch_end
                y += exp(x[i])
            end
            atomic_add!(s, y)
        end
    end
    return s[]
end

function h3(x)
    mapper = Mapper(length(x))
    s = Atomic{Float64}(0.0)

    @threads for j in 1:nthreads()
        while true
            y = 0.0
            irng = nextrange!(mapper)
            isnothing(irng) && break
            @inbounds @simd for i in irng
                y += exp(x[i])
            end
            atomic_add!(s, y)
        end
    end

    return s[]
end

function f(x)
    s = Atomic{Float64}(0.0)
    n = length(x)
    @threads for k in 1:nthreads()
        y = 0.0
        @inbounds @simd for i in getrange(n)
            y += exp(x[i])
        end
        atomic_add!(s, y)
    end
    s[]
end

function g(x)
    n = length(x)
    s = 0.0
    @inbounds @simd for i in 1:n
        s += exp(x[i])
    end
    s
end

gg(x) = mapreduce(exp, +, x)

function ggg(x)
    n = length(x)
    s = 0.0
    for i in 1:n
        s += exp(x[i])
    end
    s
end

z = rand(10^3)
@testset "summation: getrange" begin
    println("--------")

    println("Mapper")
    _h3 = h3(z)
    DOBTIME && @btime h3($z)

    println("Mapper2")
    _h2 = h2(z)
    DOBTIME && @btime h2($z)

    println("getrange")
    _f = f(z)
    DOBTIME && @btime f($z)

    println("unthreaded")
    DOBTIME && @btime g($z)
    _g = g(z)

    println("mapreduce")
    DOBTIME && @btime gg($z)

    println("vanilla")
    DOBTIME && @btime ggg($z)

    @test _f ≈ _g
    @test _g ≈ _h2
    @test _g ≈ _h3
end

end


println(zero(Float32))
println(zero(Float64))
println(zero(Int128))
println(one(Float32))
println(one(Float64))

println(zero(BigInt))
println(zero(BigFloat))

#: 0.0
#: 0.0
#: 0
#: 1.0
#: 1.0

#: 0
#: 0.0


# ---
# title: 927. Three Equal Parts
# id: problem927
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Math, Binary Search, Greedy
# link: <https://leetcode.com/problems/three-equal-parts/description/>
# hidden: true
# ---
#
# Given an array `A` of `0`s and `1`s, divide the array into 3 non-empty parts
# such that all of these parts represent the same binary value.
#
# If it is possible, return **any** `[i, j]` with `i+1 < j`, such that:
#
#   * `A[0], A[1], ..., A[i]` is the first part;
#   * `A[i+1], A[i+2], ..., A[j-1]` is the second part, and
#   * `A[j], A[j+1], ..., A[A.length - 1]` is the third part.
#   * All three parts have equal binary value.
#
# If it is not possible, return `[-1, -1]`.
#
# Note that the entire part is used when considering what binary value it
# represents.  For example, `[1,1,0]` represents `6` in decimal, not `3`.  Also,
# leading zeros are allowed, so `[0,1,1]` and `[1,1]` represent the same value.
#
#
#
# **Example 1:**
#
#
#
#     Input: [1,0,1,0,1]
#     Output: [0,3]
#
#
# **Example 2:**
#
#
#
#     Input: [1,1,0,1,1]
#     Output: [-1,-1]
#
#
#
# **Note:**
#
#   1. `3 <= A.length <= 30000`
#   2. `A[i] == 0` or `A[i] == 1`
#
#
#
#
## @lc code=start
using LeetCode

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


## Logical operators for (Sym,Sym)

## XXX Experimental! Not sure these are such a good idea ...
## but used with piecewise
@compat Base.:&(x::Sym, y::Sym) = PyCall.pycall(x.x["__and__"], Sym, y)
@compat Base.:|(x::Sym, y::Sym) =  PyCall.pycall(x.x["__or__"], Sym, y)
!(x::Sym)         =       PyCall.pycall(x.x["__invert__"], Sym)::Sym

## use ∨, ∧, ¬ for |,&,! (\vee<tab>, \wedge<tab>, \neg<tab>)
∨(x::Sym, y::Sym) = x | y
∧(x::Sym, y::Sym) = x & y
¬(x::Sym) = !x

#=

In SymPy, symbolic equations are not represented by `=` or `==`,
rather ther function `Eq` is used. Here we use the unicode
`\Equal<tab>` for an infix operator. There are also unicode values to represent `<`, `<=`, `>=`. `>`. These are

* `<` is `\ll<tab>`
* `<=` is `\le<tab>
* `==` is `\Equal<tab>`
* `>=` is `\ge<tab>`
* `>` is `\gg<tab>`

=#

## These are useful with pairwise and plot_implicit, but they cause  issues elsewhere with linear algebra functions
#Base.isless(a::Sym, b::Sym) = Lt(a,b)
#Base.isequal(a::Sym, b::Sym) = Eq(a,b)
# ⩵(a::Sym, b::Sym) = isequal(a,b)
#export ⩵
## Instead we have:
## We use unicode for visual appeal of infix operators, but the Lt, Le, Eq, Ge, Gt are the proper way:

"This is `\ll<tab>` mapped as an infix operator to `Lt`"
(≪)(a::Sym, b::Sym) = Lt(a,b)  # \ll<tab>
(≪)(a::Sym, b::Number) = Lt(a,Sym(b))  # \ll<tab>
(≪)(a::Number, b::Sym) = Lt(Sym(a),b)  # \ll<tab>

## compat for v0.4
@compat Base.:≤(a::Sym, b::Sym) = Le(a,b)  # \le<tab>
@compat Base.:≤(a::Sym, b::Number) = Le(a,Sym(b))  # \le<tab>
@compat Base.:≤(a::Number, b::Sym) = Le(Sym(a),b)  # \le<tab>

"For hashing, we use equality at the python level."
## Base.isequal(x::Sym, y::Sym) = x.x == y.x
==(x::Sym, y::Sym) = x.x == y.x

"For infix `Eq` one can use \Equal<tab> unicode operator"
(⩵)(a::Sym, b::Sym) = Eq(a,b)  # \Equal<tab>
(⩵)(a::Sym, b::Number) = Eq(a,Sym(b))  # \Equal<tab>
(⩵)(a::Number, b::Sym) = Eq(Sym(a),b)  # \Equal<tab>

# compat for v0.4
@compat Base.:>=(a::Sym, b::Sym) = Ge(a,b)
@compat Base.:>=(a::Sym, b::Number) = Ge(a,Sym(b))
@compat Base.:>=(a::Number, b::Sym) = Ge(Sym(a),b)

(≫)(a::Sym, b::Sym) = Gt(a,b)
(≫)(a::Sym, b::Number) = Gt(a,Sym(b))
(≫)(a::Number, b::Sym) = Gt(Sym(a),b)

export ≪,⩵,≫

# ## XXX need a decision.
# ## These are useful for `piecewise` and `plot_implicit`, but mess up generic code
# ## we likely only need to define isequal, ==, isless and < for generic code, but there is
# ## no way for < to return a boolean in general, save for (x < y -> float(x) < float(y)???)

# <(x::Sym,  y::Sym) = Lt(x,y)
# <(x::Sym, y::Number) = x < convert(Sym, y)
# <=(x::Sym, y::Sym) = Le(x,y)
# <=(x::Sym, y::Number) = x <= convert(Sym, y)
# >=(x::Sym, y::Sym) = Ge(x,y)
# >=(x::Sym, y::Number) =  x >= convert(Sym, y)
# >(x::Sym, y::Sym)  = Gt(x,y)
# >(x::Sym, y::Number) = x > convert(Sym, y)
# ## hacky, but == is something else to SymPy. (This is \Equal<tab>
# ==(x::Sym, y::Number) = Eq(x,y) #(x <= y) & (x >= y)

# <(x::Number, y::Sym)  = y > x
# <=(x::Number, y::Sym) = y >= x
# >=(x::Number, y::Sym) = y <= x
# >(x::Number, y::Sym)  = y < x

function !={T <: Real}(x::Sym, y::T)
    try
        x = convert(Float64, x)
        x != y
    catch
        true
    end
end
function !={T <: Complex}(x::Sym, y::T)
    try
        x = complex(x)
        x != y
    catch
        true
    end
end

function init_logical()
    ## Try this. It works with `sort`
    global const SympyTRUE = sympy_meth(:Lt, 0,1)
    Base.isless(x::Sym, y::Sym) = sympy_meth(:Lt,x, y) == SympyTRUE
end