using ADCME

function neural_net(args)
    #=
    x:      flattened mesh coordinates x-direction (PyObject: tf.Tensor)
    t:      flattened mesh coordinates time (PyObject: tf.Tensor)
    μ_w:    viscosity of wetting fluid (PyObject: tf.Tensor)
    μ_o:    viscosity of production fluid (PyObject: tf.Tensor)
    config: configuration of neural network i.e. [20, 20, 20, 3] (PyObject: tf.Tensor)
    θ:      already initialised weights and biases
    =#

    sw_nn = squeeze(fc(hcat(args.x, args.t), args.config, args.Theta)) + 1  # return neural network result

    sw_t_nn = tf.gradients(sw_nn, args.t)[1]             # first time derivative of saturation
    sw_x_nn = tf.gradients(sw_nn, args.x)[1]             # first spatial derivative of saturation
    sw_xx_nn = tf.gradients(sw_x_nn, args.x)[1]          # second spatial derivative of saturation

    λ_w = sw_nn^2 / args.μ_w                             # mobility wetting fluid
    λ_o = (1 - sw_nn)^2 / args.μ_o                       # mobility production fluid

    f_w_nn = λ_w / (λ_w + λ_o)                      # flux-function
    f_w_x_nn = tf.gradients(f_w_nn, args.x)[1]           # spatial derivative of flux-function

    return sw_nn, sw_t_nn, sw_xx_nn, f_w_x_nn
end

function get_loss(args, sw_pred, sw_t_pred, sw_xx_pred, f_w_x_pred, prod)
    #=
    sw_train:       training saturation (PyObject: tf.Tensor)
    sw_pred:        predicted saturation (PyObject: tf.Tensor)
    sw_t_pred:      predicted saturation derivative with respect to time (PyObject: tf.Tensor)
    sw_xx_pred:     predicted saturation second derivative with respect to space (PyObject: tf.Tensor)
    f_w_x_pred:     predicted flux-fluncion result derivative with respect to space (PyObject: tf.Tensor)
    ϵ:              1/Pe (PyObject: tf.Tensor)
    calib1:         calibrate saturation term (Bool)
    calib2:         calibrate flux-funcion term (Bool)
    calib3:         calibrate diffusion term (Bool)
    =#

    # define first calibration parameter
    λ_1 = Variable(ones(args.Nx))
    λ_1_multi = repeat(λ_1, args.Nt)

    # define second calibration parameter
    λ_2 = Variable(ones(args.Nx))
    λ_2_multi = repeat(λ_2, args.Nt)

    # compute loss

    # initial loss
    loss_init = sum((sw_pred[1 : args.Nx] - args.sw[1 : args.Nx])^2) / args.Nx
    # injection loss
    loss_inj =  sum((sw_pred[1 : args.Nx : (args.Nx * args.Nt)] -
                                args.sw[1 : args.Nx : (args.Nx * args.Nt)])^2) / args.Nt
    # production loss
    loss_prod = sum((sw_pred[args.Nx : args.Nx : (args.Nx * (args.Nt))] -
                                args.sw[args.Nx : args.Nx : (args.Nx * (args.Nt))])^2)/ args.Nt
    # PDE regularization
    loss_PDE = sum((sw_t_pred * λ_1_multi + f_w_x_pred * λ_2_multi - ϵ * sw_xx_pred)^2) / (args.Nx * args.Nt)

    # full loss
    if prod
        loss = loss_init + loss_inj + loss_prod + loss_PDE
    else
        loss = loss_init + loss_inj + loss_PDE
    end

    return loss, λ_1, λ_2

end

mutable struct CtxNN
    x; t; sw; μ_w; μ_o; ϵ; Theta; config; Nx; Nt
  end

function tfCtxGenNN(x, t, sw, μ_w, μ_o, ϵ, config, Nx, Nt)
    x_tf = constant(x)                          # tf.Tensor of spatial mesh
    t_tf = constant(t)                          # tf.Tensor of time mesh
    μ_w_tf = constant(μ_w)                      # tf.Tensor of wetting viscosity
    μ_o_tf = constant(μ_o)                      # tf.Tensor of production viscosity
    Θ_tf = Variable(fc_init([2; config]))       # xavier initialise weights and biases
    sw_tf = constant(sw)                        # tf.Tensor of test data
    ϵ_tf = constant(ϵ)                          # tf.Tensor of dimensionless diffusion coefficient

    return CtxNN(x_tf, t_tf, sw_tf, μ_w_tf, μ_o_tf, ϵ_tf, Θ_tf, config, Nx, Nt)
end


using Distributions

m = constructdice()
run(m)

mcs = @defsim begin
    # The parameters of the log‐normal distribution fit to Olsen et al.
    # are μ = 1.10704 and σ = 0.264 (Gillingham et al. 2015)
    t2xco2 = LogNormal(1.10704, 0.264)

    save(damages.DAMAGES)
end

res = run(mcs, m, 1000; trials_output_filename="/tmp/dice-2013/trialdata.csv", results_output_dir="/tmp/dice-2013")


using ChaosTools
using Test

"""
    saturation_point(x, y; threshold = 0.01, dxi::Int = 1, tol = 0.2)
Decompose the curve `y(x)` into linear regions using `linear_regions(x, y; dxi, tol)`
and then attempt to find a saturation point where the the first slope
of the linear regions become `< threshold`.

Return the `x` value of the saturation point.
"""
function saturation_point(Ds, E1s; threshold = 0.01, kwargs...)
    lrs, slops = ChaosTools.linear_regions(Ds, E1s; kwargs...)
    i = findfirst(x -> x < threshold, slops)
    return i == 0 ? Ds[end] : Ds[lrs[i]]
end

test_value = (val, vmin, vmax) -> @test vmin <= val <= vmax

@testset "Estimate Delay" begin

    ds = Systems.henon()
    data = trajectory(ds,1000;dt=1)
    x = data[:,1]
    @test estimate_delay(x,"first_zero") <= 2
    @test estimate_delay(x,"first_min")  <= 2
    @test estimate_delay(x,"exp_decay")  <= 2
    # @test 3 <= estimate_delay(x,"mutual_inf"; k=1) <= 10
    # @test 3 <= estimate_delay(x,"mutual_inf"; k=10) <= 10

    ds = Systems.roessler(ones(3))
    dt = 0.01
    data = trajectory(ds,200,dt=dt)
    x = data[:,1]
    @test 1.3 <= estimate_delay(x,"first_zero")*dt <= 1.7
    @test 2.6 <= estimate_delay(x,"first_min")*dt  <= 3.4

    dt = 0.1
    data = trajectory(ds,2000,dt=dt)
    x = data[:,1]
    @test 1.3 <= estimate_delay(x,"first_zero")*dt <= 1.7
    @test 2.6 <= estimate_delay(x,"first_min")*dt  <= 3.4
    @test 1.3 <= estimate_delay(x,"mutual_inf")*dt <= 1.8
    @test 1.3 <= estimate_delay(x,"mutual_inf"; maxtau=150, k = 2)*dt <= 1.7

    ds = Systems.lorenz()
    dt = 0.05
    data = trajectory(ds,1000;dt=dt)
    x = data[500:end,1]
    @test 0.1 <= estimate_delay(x,"mutual_inf")*dt <= 0.4
    @test 0.1 <= estimate_delay(x,"exp_decay")*dt  <= 0.4
    #
    dt = 0.1
    data = trajectory(ds,2000;dt=dt)
    x = data[:,1]
    @test 0.1 <= estimate_delay(x,"exp_decay")*dt  <= 0.4
end

@testset "Estimate Dimension" begin
    s = sin.(0:0.1:1000)
    τ = 15
    Ds = 1:5
    E1s = estimate_dimension(s, τ, Ds)
    E2s = stochastic_indicator(s, τ, Ds)
    @test saturation_point(Ds,E1s; threshold=0.01) == 1

    # @test minimum(E2s) < 0.1 # THIS TEST FAILS

    ds = Systems.roessler();τ=15; dt=0.1
    data = trajectory(ds,1000;dt=dt)
    s = data[:,1]
    Ds = 1:5
    E1s = estimate_dimension(s, τ, Ds)
    E2s = stochastic_indicator(s, τ, Ds)
    @test saturation_point(Ds,E1s; threshold=0.1) ∈ [2, 3]
    @test minimum(E2s) < 0.3

    ds = Systems.lorenz();τ=5; dt=0.01
    data = trajectory(ds,500;dt=dt)
    s = data[:,1]
    Ds = 1:5
    E1s = estimate_dimension(s, τ, Ds)
    E2s = stochastic_indicator(s, τ, Ds)
    @test saturation_point(Ds,E1s; threshold=0.1) ∈ [2, 3]

    # @test minimum(E2s) < 0.1 # THIS TEST FAILS

    #Test against random signal
    E2s = stochastic_indicator(rand(100000), 1, 1:5)
    @test minimum(E2s) > 0.9

end


using ApproxFun, Base.Test

x=Fun(identity);
@test_approx_eq sqrt(cos(π/2*x))(.1) sqrt(cos(.1π/2))

x=Fun(identity,[-2.,2.])
u=sqrt(4-x.^2)/(2π)
@test_approx_eq u(.1) sqrt(4-.1^2)/(2π)
@test_approx_eq sum(u) 1

#this call threw an error, which we check
@test length(values(u))==1

f=Fun(x->x.*cot(π*x/2))
x=Fun(identity)
u=Fun((f./(1-x.^2)).coefficients,JacobiWeight(1.,1.,Interval()))
@test_approx_eq 1./(.1.*cot(π*.1/2)) (1./u)(.1)

@test_approx_eq (x./u)(.1) tan(π*.1/2)

f=Fun(x->exp(-x.^2),Line(0.,0.,-.5,-.5),400)
@test_approx_eq sum(f) sqrt(π)

f=Fun(x->exp(x)/sqrt(1-x.^2),JacobiWeight(-.5,-.5))
@test_approx_eq f(.1) (x->exp(x)/sqrt(1-x.^2))(.1)

S=JacobiWeight(-1.,-1.,Chebyshev([0.,1.]))
D=Derivative(S)

f=Fun(Fun(exp,[0.,1.]).coefficients,S)
x=.1
@test_approx_eq f(x) exp(x)*x^(-1).*(1-x)^(-1)/4
@test_approx_eq (D*f)(x) -exp(x)*(1+(x-3)*x)/(4*(x-1)^2*x^2)

S=JacobiWeight(-1.,0.,Chebyshev([0.,1.]))
D=Derivative(S)

f=Fun(Fun(exp,[0.,1.]).coefficients,S)
x=.1
@test_approx_eq f(x) exp(x)*x^(-1)/2
@test_approx_eq (D*f)(x) exp(x)*(x-1)/(2x^2)

## ODEs

println("    Bessel tests")

for ν in (1.,0.5,2.,3.5)
    println("        ν = $ν")
    S=JacobiWeight(-ν,0.,Chebyshev([0.,1.]))
    D=Derivative(S)
    x=Fun(identity,domain(S))
    L=(x^2)*D^2+x*D+(x^2-ν^2);
    u=linsolve([rdirichlet(S);rneumann(S);L],[bessely(ν,1.),.5*(bessely(ν-1.,1.)-bessely(ν+1.,1.))];
                tolerance=1E-13)
    @test_approx_eq_eps u(.1) bessely(ν,.1) eps(100000.)*max(abs(u(.1)),1)
    u=linsolve([rdirichlet(S),rneumann(S),L],[besselj(ν,1.),.5*(besselj(ν-1.,1.)-besselj(ν+1.,1.))];
                tolerance=1E-13)
    @test_approx_eq_eps u(.1) besselj(ν,.1) eps(100000.)*max(abs(u(.1)),1)

    u=Fun(x->bessely(ν,x),S)
    @test_approx_eq_eps u(.1) bessely(ν,.1) eps(10000.)*max(abs(u(.1)),1)
    u=Fun(x->besselj(ν,x),S)
    @test_approx_eq_eps u(.1) besselj(ν,.1) eps(10000.)*max(abs(u(.1)),1)
end

for ν in (1.,0.5,0.123,3.5)
    println("        ν = $ν")
    S=JacobiWeight(ν,0.,Chebyshev([0.,1.]))
    D=Derivative(S)
    x=Fun(identity,domain(S))
    L=(x^2)*D^2+x*D+(x^2-ν^2);

    u=linsolve([rdirichlet(S),rneumann(S),L],[besselj(ν,1.),.5*(besselj(ν-1.,1.)-besselj(ν+1.,1.))];
                tolerance=1E-10)
    @test_approx_eq_eps u(.1) besselj(ν,.1) eps(1000000.)*max(abs(u(.1)),1)

    u=Fun(x->besselj(ν,x),S)
    @test_approx_eq_eps u(.1) besselj(ν,.1) eps(10000.)*max(abs(u(.1)),1)
end

## f/g bugs

println("    Jacobi singularity tests")

x = Fun(identity)
f = exp(x)./(1-x.^2)

@test_approx_eq f(.1) exp(.1)./(1-.1^2)
f = exp(x)./(1-x.^2).^1
@test_approx_eq f(.1) exp(.1)./(1-.1^2)
f = exp(x)./(1-x.^2).^1.0
@test_approx_eq f(.1) exp(.1)./(1-.1^2)

## 1/f with poles

x=Fun(identity)
f=sin(10x)
g=1/f

@test_approx_eq g(.123) csc(10*.123)

println("    Ray tests")

## Ray

@test Inf in Ray()   # this was a bug

f=Fun(x->exp(-x),[0,Inf])
@test_approx_eq f'(.1) -f(.1)

x=Fun(identity,Ray())
f=exp(-x)
u=integrate(f)
@test_approx_eq (u(1.)-u(0)-1) -f(1)

x=Fun(identity,Ray())
f=x^(-0.123)*exp(-x)
@test_approx_eq integrate(f)'(1.) f(1.)

@test_approx_eq_eps sum(Fun(sech,[0,Inf])) sum(Fun(sech,[0,40.])) 1000000eps()

#Ei (Exp Integral)

y=Fun(Ray())
q=integrate(exp(-y)/y)
@test_approx_eq (q-last(q))(2.) (-0.04890051070806113)

## Line

f=Fun(x->exp(-x^2),Line())

@test_approx_eq f'(0.1) -2*0.1exp(-0.1^2)
@test_approx_eq (Derivative()*f)(0.1) -2*0.1exp(-0.1^2)

## PeriodicLine

d=PeriodicLine()
D=Derivative(d)
f=Fun(x->sech(x-.1),d)

@test_approx_eq_eps (D*f)(.2) -0.0991717226583897  100000eps()
@test_approx_eq_eps (D^2*f)(.2) -0.9752522555114987  1000000eps()

## LogWeight

x=Fun(identity,[-1.,1.])
f=exp(x+1)-1
@test_approx_eq log(f)(0.1) log(f(0.1))

x=Fun(identity,[0.,1.])
f=exp(x)-1
@test_approx_eq log(f)(0.1) log(f(0.1))

## Test divide sing

x=Fun(identity,[0,1])
@test_approx_eq Fun(exp(x)/x-1/x,Chebyshev)(0.1) (exp(0.1)-1)/0.1

x=Fun(identity,[0,1])
f=1/x
p=integrate(f)
@test_approx_eq (p-p(1.))(0.5) log(0.5)

f=1/(1-x)
p=integrate(f)
@test_approx_eq (p-p(0.))(0.5) -log(1-0.5)

y=Fun(Ray())
@test_approx_eq (y^2)(10.) y(10.)^2
@test_approx_eq 1/y(10.) (1/y)(10.)
@test_approx_eq (1/y^2)(10.) 1/y(10.)^2
@test_approx_eq (-1/y^2)'(10.) 2/(y(10.)^3)
@test_approx_eq exp(-1/y^2)(5.) exp(-1/y(5.)^2)

# catch bug from SIE

a=1+10*im;b=2-6*im
d=Curve(Fun(x->1+a*x+b*x^2))

x=Fun(d)
w=sqrt(abs(first(d)-x))*sqrt(abs(last(d)-x))

@test_approx_eq sum(w/(x-2.))/(2π*im) (-4.722196879007759+2.347910413861846im)
@test_approx_eq linesum(w*log(abs(x-2.)))/π (88.5579588360686)

## Dirac Space

a,b=DiracDelta(0.),DiracDelta(1.)
f=Fun(exp)
g=a+0.2b+f
@test_approx_eq pieces(g)[2](0.) 1.
@test_approx_eq g(.1) exp(.1)
@test_approx_eq sum(g) (sum(f)+1.2)

#Checks prevoius bug
δ=DiracDelta()
x=Fun()
w=sqrt(1-x^2)
w+δ

## multiplicities
x=Fun(identity,[-1,1.])
@test_approx_eq (1/x^2)(0.1) 100.
@test_approx_eq (1/x^2)(-0.1) 100.

fc=x*(1+x)^2
@test_approx_eq (1/fc)(0.1) 1/fc(0.1)

fc=x*(1-x)^2
@test_approx_eq (1/fc)(0.1) 1/fc(0.1)

## erf(sqrt(x))

x=Fun([0.,1.])
@test_approx_eq erf(sqrt(x))(0.1) erf(sqrt(0.1))
@test_approx_eq erfc(sqrt(x))(0.1) erfc(sqrt(0.1))

## norm(u-x)

@test_approx_eq norm(Fun(exp,Legendre([0,1]))+sqrt(x)) 2.491141949903508

## Test Jacobi special conversions

S1,S2=JacobiWeight(3.,1.,Jacobi(1.,1.)),JacobiWeight(1.,1.,Jacobi(1.,0.))
f=Fun([1,2,3.],S1)
C=Conversion(S1,S2)
Cf=C*f
@test_approx_eq Cf(0.1) f(0.1)

S1,S2=JacobiWeight(3.,2.,Jacobi(1.,1.)),JacobiWeight(1.,1.,Jacobi(0.,0.))
f=Fun([1,2,3.],S1)
C=Conversion(S1,S2)
Cf=C*f
@test_approx_eq Cf(0.1) f(0.1)


# # How to port NiLang to Zygote
using NiLang, NiLang.AD, Zygote

# Let's start from the Julia native implementation of `norm2` function.
function norm2(x::AbstractArray{T}) where T
    out = zero(T)
    for i=1:length(x)
        @inbounds out += x[i]^2
    end
    return out
end

# Zygote is able to generate correct gradients, but much slower than the original program.
using BenchmarkTools
x = randn(1000);
original_grad = norm2'(x)
@benchmark norm2'($x) seconds=1

# The orignal program is
@benchmark norm2($x) seconds=1

# Then we have the reversible implementation
@i function r_norm2(out::T, x::AbstractArray{T}) where T
    for i=1:length(x)
        @inbounds out += x[i]^2
    end
end

# The gradient generated by NiLang is much faster, which is comparable to the forward program
@benchmark (~r_norm2)(GVar($(norm2(x)), 1.0), $(GVar(x))) seconds=1

# to enjoy the speed of `NiLang` in `Zygote`, just bind the adjoint rule
Zygote.@adjoint function norm2(x::AbstractArray{T}) where T
    out = norm2(x)
    out, δy -> (grad((~r_norm2)(GVar(out, δy), GVar(x))[2]),)
end
@assert norm2'(x) ≈ original_grad

# See, much faster
@benchmark norm2'(x) seconds=1


using TestImages: testimage
using CSV, FileIO, Images, Colors, Augmentor, Distributions

"""
    spaghetti_fun(x, y, a, b, τ, len)

Gives the intensity for a single "spaghetti" at (`x`, `y`) with the params `a`, `b`, `τ` and `len`.

The intensity is equal to:
```math
I(\\mathrm{x,y})=\\left(1-\\frac{f(\\mathrm{x, y})}{τ^{2}}\\right)\\left(1+\\mathrm{e}
^{\\left(\\left(\\mathrm{x}^2 +\\mathrm{y}^2 \\right) -\\mathrm{len}^{2}\\right) /\\left(50 τ^{2}
\\right)}\\right)^{-1} \\mathrm{e}^{-f(\\mathrm{x, y}) /\\left(2 τ^{2}\\right)}
```
where
```math
f(\\mathrm{x, y})=\\left|y-a x^{3}-b x\\right|^{2} /\\left(1+\\left|3 a x^{2}-b\\right|^{2}\\right)
```

See also: [`gen_spaghetti`](@ref), [`gen_one_spaghetti`](@ref), [`gen_spaghetti`](@ref)
"""
function spaghetti_fun(x, y, a, b, τ, len)
    spag(x, y, a, b) = (abs(y - a*x^3 - b*x)^2)/(1 + abs(3a*x^2 - b)^2)
    args = (x, y, a, b)
    return (1 - spag(args...)/τ^2) * (1 + exp(((x^2 + y^2) - len)/(50*τ^2)))^(-1) * exp((-spag(args...))/(2*τ^2))
end

# TODO edit wording
"""
    gen_one_spaghetti(img_size, τ, len, spread)

Generate a single spaghetti image with the params `τ` (thickness), `len` (length) and `spread`.
The coordinates `x` and `y` in the equation are scaled from `(-1, 1)` and the `spread` gives the length of interval
in which the `(-1, 1)` is randomly located. The genereated spaghetti is also randomly rotated.

(Example for spread 10, the interval for `x` and `y` is randomly anywhere between `(-9:1)` and `(-1:9)`)

The intensity is equal to:
```math
I(\\mathrm{x,y})=\\left(1-\\frac{f(\\mathrm{x, y})}{τ^{2}}\\right)\\left(1+\\mathrm{e}
^{\\left(\\left(\\mathrm{x}^2 +\\mathrm{y}^2 \\right) -\\mathrm{len}^{2}\\right) /\\left(50 τ^{2}
\\right)}\\right)^{-1} \\mathrm{e}^{-f(\\mathrm{x, y}) /\\left(2 τ^{2}\\right)}
```
where
```math
f(\\mathrm{x, y})=\\left|y-a x^{3}-b x\\right|^{2} /\\left(1+\\left|3 a x^{2}-b\\right|^{2}\\right)
```
`a` and `b` are chosen randomly from a normal distribution.

See also: [`gen_spaghetti`](@ref), [`spaghetti_fun`](@ref)
"""
function gen_one_spaghetti(img_size, τ, len, spread)
    (a, b) = randn(2)

    c = rand()*(spread-2)+1
    d = rand()*(spread-2)+1
    enlarged_img_size = ceil.(Int, img_size .* sqrt(2)) # so that the rotation doesnt make have to extrapolate

    X = ones(enlarged_img_size[1]) * collect(range(c-spread,c,length=enlarged_img_size[2]))'
    Y = collect(range(d-spread,d,length=enlarged_img_size[1])) * ones(enlarged_img_size[2])'

    rand_rot = Rotate(0:360)
    crop = CropSize(img_size...)
    pl = rand_rot |>
         # ShearX(-5:5) * ShearY(-5:5) |>
         crop

    X_rot, Y_rot = augment((X, Y), pl);

    img = map((x, y) -> spaghetti_fun(x, y, a, b, τ, len), X_rot, Y_rot)
end

"""
    gen_spaghetti(img_size = (200, 200),
                  τ = 0.2,
                  len = 30;
                  count=25,
                  spread=10)

Generate an image with `count` spaghetti images with the params `τ` (thickness), `len` (length) and `spread`.
The coordinates `x` and `y` in the equation are scaled from `(-1, 1)` and the `spread` gives the length of interval
in which the `(-1, 1)` is randomly located. The genereated spaghetti is also randomly rotated.

(Example for spread 10, the interval for `x` and `y` is randomly anywhere between `(-9:1)` and `(-1:9)`)

# Example
```@example
# feature-full spaghetti image
img = gen_spaghetti((256, 256), 0.15, 70, spread=20)
# feature-less spaghetti image
img = gen_spaghetti((256, 256), 1, 50, spread=12)
```

Note: For the exact intensity function see [`gen_one_spaghetti`](@ref)

See also:  [`spaghetti_fun`](@ref)
"""
function gen_spaghetti(img_size = (256, 256),
                       τ = 0.2,
                       len = 30;
                       count=25,
                       spread=10)

    img = zeros(img_size)
    for _ in 1:count
        adding_img = gen_one_spaghetti(img_size, τ, len, spread)
        img = img + adding_img
    end
    img = LAP_julia.normalize_to_zero_one(img)
    return img
end

rand_uniform(a, b) = rand() * (b-a) + a

"""
    rand_H(flow_size)

Generate a semi-random homography transfomration matrix. See the implementation for used constants and distributions.

See also: [`showflow`](@ref), [`gen_tiled_flow`](@ref), [`Flow`](@ref), [`gen_quad_flow`](@ref), [`gen_homo_flow`](@ref)
"""
function rand_H(flow_size)
    a = flow_size[1] * 0.4
    H = [rand(Normal(1, 0.4)) rand(Normal(0, 0.2)) rand_uniform(-a, a);
         rand(Normal(0, 0.2)) rand(Normal(1, 0.4)) rand_uniform(-a, a);
         rand(Normal(0, 0.0002)) rand(Normal(0, 0.0002)) 1.]
    @info H
    return H
end

"""
    gen_homo_flow(flow_size, max_magnitude=10)

Generate a smoothly varying locally constant random flow of size `img_size` and with a maximal
displacement of `max_magnitude` using a randomly generated homography matrix.

See also: [`showflow`](@ref), [`gen_tiled_flow`](@ref), [`Flow`](@ref), [`gen_quad_flow`](@ref), [`rand_H`](@ref)
"""
function gen_homo_flow(flow_size, max_magnitude=10)
    H = rand_H(flow_size)
    flow = LAP_julia.make_flow_from_H(H, flow_size)
    max_len = maximum(LAP_julia.vec_len.(flow))
    flow = flow .* (max_magnitude/max_len)
    @assert median(LAP_julia.vec_len.(flow)) > (0.1*max_magnitude) H
    return flow
end

"""
    gen_anhir(base_path = "/Users/MrTrololord/Documents/anhir/";
              mutate=true,
              diag_pixels=500)

Generate a random ANHIR image pair `target` and `source` images (in this order) from the ANHIR dataset at `base_path`.

See also: [`load_anhir_image_pair`](@ref)
"""
function gen_anhir(base_path = "/Users/MrTrololord/Documents/anhir/";
                   mutate=true,
                   diag_pixels=500)

    loc_table = CSV.read(joinpath(base_path, "location_table.csv"))
    train_rows = loc_table[loc_table[:status] .== "training", :]

    random_row = train_rows[rand(1:size(train_rows, 1)), :]
    target_path = random_row[Symbol("Target image")]
    source_path = random_row[Symbol("Source image")]

    target, source = load_anhir_image_pair(target_path, source_path, base_path=joinpath(base_path, "dataset"))

    if any(size(target) .!= size(source))
        println("sizes: ", size(target), ", ", size(source))
    end
    target, source = pad_images(target, source)

    if mutate
        target = resize_to_diag_size(target, diag_pixels)
        source = resize_to_diag_size(source, diag_pixels)
    end

    return target, source
end

"""
    resize_to_diag_size(img, desired_diag_size)

Resize `img` so that it has `desired_diag_size` pixels on its diagonal.

See also: [`gen_anhir`](@ref), [`load_anhir_image_pair`](@ref)
"""
function resize_to_diag_size(img, desired_diag_size)
    actual_diag_pixels = sqrt(sum((size(img).^2)))
    resize_ratio = desired_diag_size/actual_diag_pixels
    img = imresize(img, ratio = resize_ratio)
    return img, resize_ratio
end

"""
    load_anhir_image_pair(target_path,
                          source_path;
                          base_path = "/Users/MrTrololord/Documents/anhir/dataset")

Load an image pair in the locations `target_path` and `source_path` from the location `base_path`.

See also: [`gen_anhir`](@ref)
"""
function load_anhir_image_pair(target_path,
                               source_path;
                               base_path = "/Users/MrTrololord/Documents/anhir/dataset")

    println(joinpath(base_path, target_path))
    target = Float64.(Gray.(load(joinpath(base_path, target_path))))
    source = Float64.(Gray.(load(joinpath(base_path, source_path))))

    return target, source
end

function load_image_gray(img_path)::Image
    img = Float64.(Gray.(load(img_path)))
end

"""
    gen_lena()

Get a `256x256` grayscale "lena" image.
"""
function gen_lena()
    img = testimage("lena_gray")
    img = Float64.(img)
    return img
end

"""
    gen_quad_flow(img_size, max_magnitude=10)

Generate a smoothly varying locally constant random flow of size `img_size` and with a maximal
displacement of `max_magnitude` using a quadratic function:

```math
f(z) = a + b*z + c*z^2,
```
where:
```math
z = x + y*i, x, y ∈ (0,1)
```

The constants `a`, `b` and `c` are random numbers from the normal distribution with mean 0 and standard deviation 1.

See also: [`showflow`](@ref), [`gen_tiled_flow`](@ref), [`Flow`](@ref)
"""
function gen_quad_flow(img_size, max_magnitude=10)

    a, b, c = randn(), randn(), randn()
    f(z) = a + b*z + c*z^2

    X = ones(img_size[1]) * collect(range(0,1,length=img_size[2]))'
    Y = collect(range(0,1,length=img_size[1])) * ones(img_size[2])'
    A = Y .* im + X

    B = f.(A)
    max_len = maximum(vec_len.(B))
    B = B .* (max_magnitude/max_len)

    return B
end

"""
    gen_uniform_flow(flow_size=(200, 200), vector=1 + 1im, max_magnitude=vec_len(vector))

Generate a uniform flow of size `flow_size`, where every displacement vector is `vector`, scaled by `max_magnitude`.

Note: `max_magnitude` = `vec_len(vector)` by default,
so if nothing is enetered as `max_magnitude` the flow is made up of `vector` and is not scaled.

See also: [`showflow`](@ref), [`gen_tiled_flow`](@ref), [`Flow`](@ref), [`gen_quad_flow`](@ref)
"""
function gen_uniform_flow(flow_size=(200, 200),
                          vector=1 + 1im,
                          max_magnitude=vec_len(vector))

    flow = ones(flow_size) .* vector

    max_len = vec_len(vector)
    flow = flow .* (max_magnitude/max_len)
    return flow
end

"""
    gen_tiled_flow(flow_size::Tuple{T, T}=(200, 200),
    max_magnitude::Real=20, tile_size=nothing; filter_amp=nothing)::Flow where {T <: Integer}

Generate a smoothly varying random flow. The flow parameters are set by the function arguments.

It works by generating a tiled flow where each tile has a random uniform flow. Then it smooths these tiles by filtering with a gaussian.

# Arguments
- `flow_size::Tuple{T, T}=(200, 200)`: dimensions of the flow.
- `max_magnitude::Real=20`: maximum allowed amplitude of the displacement.
- `tile_size=nothing`: size of random uniform flow tiles that make up the start matrix. (The larger the slower the flow.) Note: If set to `flow_size` or larger it will generate a uniform pixel shift in a random direction.
- `filter_amp=nothing`: size of the gaussian filter which is used to smooth the random start matrix.

See also: [`showflow`](@ref), [`Flow`](@ref), [`gen_chess`](@ref), [`gen_quad_flow`](@ref)
"""
function gen_tiled_flow(flow_size::Tuple{T, T}=(200, 200),
                        max_magnitude::Real=10,
                        tile_size=nothing,
                        filter_amp=nothing)::Flow where {T <: Integer}

    # set default values
    if tile_size == nothing
        tile_size = ceil(Int64, flow_size[1]/6)
    end
    if filter_amp == nothing
        filter_amp = ceil(Int64, tile_size/2)
        if isodd(filter_amp)
            filter_amp += 1
        end
    end

    tile_count = ceil.(Int64, flow_size ./ tile_size)

    uv_base = Array{Float64}(undef, tile_count..., 2)
    uv_base[:,:,1] .= -max_magnitude .+ (2 .* max_magnitude .* rand(tile_count...))
    uv_base[:,:,2] .= -max_magnitude .+ (2 .* max_magnitude .* rand(tile_count...))

    uv_flow = Array{Float64}(undef, tile_size*tile_count[1], tile_size*tile_count[2], 2)
    uv_flow[:,:,1] .= repeat(uv_base[:,:,1], inner=(tile_size, tile_size))
    uv_flow[:,:,2] .= repeat(uv_base[:,:,2], inner=(tile_size, tile_size))

    # cut off what we dont want
    #rand_flow = zeros(Complex{Float64}, flow_size)
    rand_flow = uv_flow[1:flow_size[1], 1:flow_size[2], 1] .+ (im * uv_flow[1:flow_size[1], 1:flow_size[2], 2]);

    # blur to make it continuous
    rand_flow = smooth_with_gaussian!(rand_flow, [filter_amp, filter_amp])

    # set maximum magnitude to actually be the max_magnitude
    max_len = maximum(vec_len.(rand_flow))
    rand_flow = rand_flow .* (max_magnitude/max_len)

    return rand_flow
end

"""
    gen_chess(tile_size::Integer=50, board_size::Integer=4)

Create a chessboard image with `board_size` tiles in each dimension, where each tile is `tile_size` pixels in each dimension.
Note that `board_size` has to be even.

See also: [`imgshow`](@ref), [`gen_tiled_flow`](@ref)
"""
function gen_chess(tile_size::Integer=50, board_size::Integer=4)
    mini_board = [zeros(tile_size, tile_size) ones(tile_size, tile_size);
                  ones(tile_size, tile_size) zeros(tile_size, tile_size)]

    chessboard = repeat(mini_board, outer=(convert(Integer, (board_size/2)), convert(Integer, (board_size/2))))
    return chessboard
end

"""
    gen_init(type::Symbol=:lena; flow_args=[])

Create the usual testing data; `img`, `imgw`, `flow` using the given parameters.

# Arguments
- `img_type::Symbol=:lena`: what base image is used. [Options: `:lena`, `:chess`]
- `flow_type::Symbol=:quad`: what flow generation function is used. [Options: `:tiled`, `:quad`, `:uniform`, `:homo`]

# Keyword Arguments
- `flow_args=[]`: arguments passed to the flow generation function besides the flow size.
- `img_args=[]`: arguments passed to the img generation function if `:chess` is chosen.
- `flipped::Bool=false`: the image will be flipped along the horizontal axis, so that the origin is at the bottom.

# Example
```@example
# chess image, warped chess image, flow with maximal displacement 20 generated by the `gen_quad_flow` function
img, imgw, flow = gen_init(:chess, :quad, flow_args=[20])

See also: [`gen_anhir`](@ref), [`gen_lena`](@ref), [`gen_spaghetti`](@ref), [`gen_tiled_flow`](@ref), [`gen_quad_flow`](@ref), [`gen_uniform_flow`](@ref)
```
"""
function gen_init(img_type::Symbol=:lena, flow_type::Symbol=:quad; flow_args=[], img_args=[], flipped::Bool=false)
    if img_type == :lena
        img = gen_lena()
    elseif img_type == :chess
        img = gen_chess((Int64.(img_args))...)
    elseif img_type == :spaghetti
        img = gen_spaghetti(img_args...)
    end

    if flipped
        img = reverse(img, dims=1)
    end

    if flow_type == :quad
        flow = gen_quad_flow(size(img), flow_args...)
    elseif flow_type == :tiled
        flow = gen_tiled_flow(size(img), flow_args...)
    elseif flow_type == :uniform
        flow = gen_uniform_flow(size(img), flow_args...)
    elseif flow_type == :homo
        flow = gen_homo_flow(size(img), flow_args...)
    end

    imgw = warp_img(img, -real(flow), -imag(flow), border_strat=:zeros)

    return img, imgw, flow
end


seed!(12345)

function testCategorical1(n::Int64, p::Vector{Float64})
  vs::Vector{Int64} = sampleCategorical(n, p)
  counts = Vector{Int64}(undef, length(p))
  fill!(counts, 0)
  for i in 1:n
    v::Int64 = vs[i]
    @inbounds counts[vs[i]] += 1
  end
  return testGOFMultinomial(p, counts)
end

function testCategorical2(n::Int64, p::Vector{Float64})
  counts = Vector{Int64}(undef, length(p))
  fill!(counts, 0)
  for i in 1:n
    v::Int64 = sampleCategorical(p)
    @inbounds counts[v] += 1
  end
  return testGOFMultinomial(p, counts)
end

@test testCategorical1(100, [0.5; 0.1; 0.25; 0.15]) > 0.01
@test testCategorical1(10000, [0.5; 0.1; 0.25; 0.15]) > 0.01

@test testCategorical2(100, [0.5; 0.1; 0.25; 0.15]) > 0.01
@test testCategorical2(10000, [0.5; 0.1; 0.25; 0.15]) > 0.01


export scatter1d

# depth-first scattering tree traversal
"""
    scatter1d_layer(S, i, j, nDone, x_hat, Transform, FilterBanks)

Apply one layer of the wavelet scattering transform, calling future layers recursively.

# Arguments
 - `S::Scattered`: Scattered object for holding scattering coefficients.
 - `i::Int64`: Index of the current scattering layer.
 - `j::Int64`: Max subsampling accepted for the current node
 - `x_hat::Vector{Complex{Float64}}`: Fourier transformed input data vector.
 - `Transform::ScatteringTransform1d`: ScatteringTransform1d object containing information about the total transform.
 - `FilterBanks::Vector{FilterBank1d}`: Filter banks for each layer of the scattering transform.
"""
function scatter1d_layer(S::Scattered, i::Int64, j::Int64, nDone::Vector{Int64}, x_hat::Vector{Complex{Float64}}, Transform::ScatteringTransform1d, FilterBanks::Vector{FilterBank1d}; subsample::Bool=true)
    for ψ in FilterBanks[i].Λ
        if ψ.j > j || subsample==false
            # convolve signal with morlet wavelet, transform back to the time domain, and take the modulus
            U_hat = fft(abs.(ifft(x_hat .* ψ.ψ)))

            # stop scattering when we reach the last layer
            if i < Transform.D
                scatter1d_layer(S, i+1, ψ.j, nDone, U_hat, Transform, FilterBanks, subsample=subsample)
            end

            # apply lowpass filter
            # result should be real-valued, so cast to real
            S.Coeff[i+1][:, nDone[i]] .= real.(ifft(U_hat .* FilterBanks[i].ϕ.ϕ))
            nDone[i]+=1
        end
    end
end

"""
    scatter1d(x, Transform, FilterBanks)

Calculate the 1d wavelet scattering transform.

# Arguments
 - `x::Vector{Float64}`: Time series.
 - `Transform::ScatteringTransform1d`: Scattering transform object.
 - `FilterBanks::Vector{FilterBank1d}`: Filter banks for each layer of the scattering transform.
"""
function scatter1d(x::Vector{Float64}, Transform::ScatteringTransform1d, FilterBanks::Vector{FilterBank1d}; subsample::Bool=true)
    @assert length(FilterBanks)==Transform.D "The number of filter banks should equal the scattering transform depth.
                                              Got $(length(FilterBanks)) filter banks and $(Transform.D) layers."

    # Transform the signal to the frequency domain
    x_hat = FFTW.fft(x)
    # create Scattered object to hold the scattering coefficients
    S = Scattered1d(FilterBanks, subsample=subsample)

    # zeroth order scattering coefficients are a simple averaging of the
    S.Coeff[1] .= real.(ifft(x_hat .* FilterBanks[1].ϕ.ϕ))
    nDone=ones(Int, Transform.D)
    # higher order scattering coefficients are computed recursively
    # seed j=0 to allow all filters for the first layer
    scatter1d_layer(S, 1, -1, nDone, x_hat, Transform, FilterBanks, subsample=subsample)

    return S
end

# breadth-first scattering tree traversal
# TODO: allow any number of layers with this approach
"""
    scatter1d(x, Transform, FilterBanks)

Calculate the 1d wavelet scattering transform.

# Arguments
 - `x::Vector{Float64}`: Time series.
 - `Transform::ScatteringTransform1d`: Scattering transform object.
 - `FilterBanks::Vector{FilterBank1d}`: Filter banks for each layer of the scattering transform.
"""
function scatter1d(x::Vector{Float64}, Transform::ScatteringTransform1d, FilterBanks::Vector{FilterBank1dBlock})
    @assert length(FilterBanks)==Transform.D "The number of filter banks should equal the scattering transform depth.
                                              Got $(length(FilterBanks)) filter banks and $(Transform.D) layers."
    @assert Transform.D <=2 "The number of layers should be 2 or fewer. Got $(Transform.D)"

    # Transform the signal to the frequency domain
    x_hat = FFTW.fft(x)
    # create Scattered object to hold the scattering coefficients
    S = Scattered1d(FilterBanks)

    # zeroth order scattering coefficients are a simple averaging of the
    S.Coeff[1] .= real.(ifft(x_hat .* FilterBanks[1].ϕ))

    # first order scattering coefficients
    U1 = Array{Float64,2}(undef, size(FilterBanks[1].Λ))
    U1 .= abs.(cwt(x, FilterBanks[1]))
    S.Coeff[2] .= real.(ifft(fft(U1, 1) .* FilterBanks[1].ϕ, 1))

    # second order scattering coefficients
    U2i = Array{Float64,2}(undef, size(FilterBanks[2].Λ))
    nDone=0
    for i=1:size(U1,2)
        U2i .= abs.(cwt(U1[:,i], FilterBanks[2]))
        S.Coeff[3][:, nDone+1:nDone+size(U2i,2)] .= real.(ifft(fft(U2i,1) .* FilterBanks[2].ϕ, 1))
        nDone+=size(U2i,2)
    end

    return S
end


using quantumoptics

N = 200
xmin = -10.
xmax = 10.
omega = 2.

x0 = 1.
p0 = 0.
sigma0 = 1.

T = [0.:0.2:2.;]

basis_position = quantumoptics.particle.PositionBasis(xmin, xmax, N)
basis_position2d = compose(basis_position, basis_position)
basis_momentum = quantumoptics.particle.MomentumBasis(basis_position)
basis_momentum2d = compose(basis_momentum, basis_momentum)

xsamplepoints = quantumoptics.particle.samplepoints(basis_position)
psamplepoints = quantumoptics.particle.samplepoints(basis_momentum)

# Tpx = quantumoptics.particle.FFTOperator(basis_momentum, basis_position)
# Txp = quantumoptics.particle.FFTOperator(basis_position, basis_momentum)

# Position Basis
@time p = SparseOperator(quantumoptics.particle.momentumoperator(basis_position))
@time x = SparseOperator(quantumoptics.particle.positionoperator(basis_position))
@time V = SparseOperator(basis_position, diagm(complex(omega)^2*xsamplepoints.^2))
@time I = sparse_identity(basis_position)

@time p2 = p^2/2
@time Hkin2d = (tensor(p2, I) + tensor(I, p2))
@time V2d = tensor(V, I) + tensor(I, V)
@time H2d = Hkin2d + V2d

@time psi0 = quantumoptics.particle.gaussianstate(basis_position, x0, p0, sigma0)
@time psi0_2d = tensor(psi0, psi0)
@time tout, psi_t = timeevolution.schroedinger(T, psi0_2d, H2d)
@time tout, psi_t = timeevolution.schroedinger(T, psi0_2d, H2d)

# for i=1:length(T)
#     println(expect(p_mom, psi_mom_t[i]), ", ", expect(p_pos, psi_pos_t[i]))
#     println(expect(x_mom, psi_mom_t[i]), ", ", expect(x_pos, psi_pos_t[i]))
#     println(norm(Txp*psi_mom_t[i] - psi_pos_t[i]))
#     println(norm(psi_mix_t[i] - psi_pos_t[i]))
#     println("================================================")
# end


#=
network.jl
~~~~~~~~~~

A module to implement the stochastic gradient descent learning
algorithm for a feedforward neural network. Gradients are calculated
using backpropagation. Note that I have focused on making the code
simple, easily readable, and easily modifiable. It is not optimized,
and omits many desirable features.

Original Python code: Copyright © 2012-2022 Michael Nielsen
Julia reimplementation: Copyright © 2018-2022 Derik Kauffman
Licensed under MIT license
=#

using Random

struct Network
	sizes::Vector{Int}
	biases::Vector{Vector{Float32}}
	weights::Vector{Matrix{Float32}}
end

"""
	Network(sizes::Vector{<:Integer})

The list `sizes` contains the number of neurons in the
respective layers of the network. For example, if the list
was [2, 3, 1] then it would be a three-layer network, with the
first layer containing 2 neurons, the second layer 3 neurons,
and the third layer 1 neuron. The biases and weights for the
network are initialized randomly, using a Gaussian
distribution with mean 0 and variance 1. Note that the first
layer is assumed to be an input layer, and by convention we
won't set any biases for those neurons, since biases are only
ever used in computing the outputs from later layers.
"""
function Network(sizes::Vector{<:Integer})
	biases = randn.(sizes[2:end])
	weights = randn.(zip(sizes[2:end], sizes[1:end-1]))
	Network(sizes, biases, weights)
end

"""
	feedforward(net::Network, a::AbstractVector{<:Real})

Return the output of the network if `a` is input.
"""
function feedforward(net::Network, a::AbstractVector{<:Real})
	for (b, w) = zip(net.biases, net.weights)
		a = σ.(w * a .+ b)
	end
	return a
end

"""
The neural network's output is assumed to be the index of whichever
neuron in the final layer has the highest activation.
"""
classify(net::Network, a::AbstractVector{<:Real}) = argmax(feedforward(net, a)) - 1

"""
Train the neural network using mini-batch stochastic gradient
descent. The `training_data` is a list of tuples `(x, y)`
representing the training inputs and the desired outputs.
The other non-optional parameters are self-explanatory.
If `test_data` is provided then the network will be evaluated
against the test data after each epoch, and partial progress
printed out. This is useful for tracking progress, but slows things
down substantially.
"""
function SGD!(net::Network, training_data, epochs::Integer,
	  mini_batch_size::Integer, η::Real; test_data=nothing)
	if test_data !== nothing n_test = length(test_data) end
	for j in 1:epochs
		shuffle!(training_data)
		mini_batches = reshape(training_data, :, mini_batch_size)
		for mini_batch in eachrow(mini_batches)
			update_mini_batch!(net, mini_batch, η)
		end

		if test_data !== nothing
			println("Epoch $j: $(evaluate(net, test_data)) / $n_test")
		else
			println("Epoch $j complete")
		end
	end
end

"""
	update_mini_batch!(net::Network, mini_batch, η::Real)

Update the network's weights and biases by applying
gradient descent using backpropagation to a single mini batch.
The `mini_batch` is a list of tuples `(x, y)`, and `η`
is the learning rate.
"""
function update_mini_batch!(net::Network, mini_batch, η::Real)
	∇b = zero.(net.biases)
	∇w = zero.(net.weights)
	for (x, y) in mini_batch
		Δ∇b, Δ∇w = backprop(net, x, y)
		∇b .+= Δ∇b
		∇w .+= Δ∇w
	end
	net.biases  .-= Float32(η / length(mini_batch)) .* ∇b # Eq. 21
	net.weights .-= Float32(η / length(mini_batch)) .* ∇w # Eq. 20
end

"""
	backprop(net::Network, x, y)

Return a tuple `(∇b, ∇w)` representing the gradient for the
cost function Cₓ. `∇b` and `∇w` are layer-by-layer lists of arrays,
similar to `net.biases` and `net.weights`.
"""
function backprop(net::Network, x, y)
	∇b = zero.(net.biases)
	∇w = zero.(net.weights)

	# feedforward
	activation = x
	activations = [x] # list to store all the activations, layer by layer
	zs = typeof(x)[]  # list to store all the z vectors, layer by layer
	for (b, w) = zip(net.biases, net.weights)
		z = w * activation .+ b
		push!(zs, z)
		activation = σ.(z)
		push!(activations, activation)
	end

	# backward pass
	δ = cost_derivative(activations[end], y) .* σ′.(zs[end]) # Eq. BP1
	∇b[end] = δ
	∇w[end] = δ * activations[end-1]'

	for l in length(net.sizes)-1:2
		z = zs[l]
		δ = (net.weights[l+1]' * δ) .* σ′.(z) # Eq. BP2
		∇b[l] = δ # Eq. BP3
		∇w[l] = δ * activations[l-1]' # Eq. BP4
	end
	return ∇b, ∇w
end

"""
	evaluate(net::Network, test_data)

Return the number of test inputs for which the neural
network outputs the correct result.
"""
function evaluate(net::Network, test_data)
	return count((classify(net, x) == y) for (x, y) in test_data)
end

"""
	cost_derivative(output_activations, y)

Return the vector of partial derivatives ∂Cₓ/∂a for the output activations.
"""
cost_derivative(output_activations, y) = output_activations .- onehot(y)

#### Miscellaneous functions
"""The sigmoid function."""
σ(z) = 1/(1 + exp(-z))

"""Derivative of the sigmoid function."""
σ′(z) = σ(z) * (1 - σ(z))

"""
	onehot(j::Integer)

Return a 10-dimensional unit vector with a 1 in the jth
position and zeros elsewhere. This is used to convert a digit (0-9)
into a corresponding desired output from the neural network.
"""
function onehot(j::Integer)
	e = falses(10)
	e[j+1] = true
	return e
end

# MNIST loading code
using MLDatasets

function load_data()
	training_inputs = reshape(MNIST.traintensor(Float32), 784, :)
	training_labels = MNIST.trainlabels()
	training_data = zip(eachcol(training_inputs), training_labels)

	test_inputs = reshape(MNIST.testtensor(Float32), 784, :)
	test_labels = MNIST.testlabels()
	test_data = zip(eachcol(test_inputs), test_labels)
	println("Loaded data")
	return training_data, test_data
end

function test_mnist()
	training_data, test_data = load_data()
	net = Network([784, 30, 10])
	SGD!(net, training_data, 30, 10, 0.1, test_data=test_data)
	return net
end

function test_mnist_quick()
	training_data, test_data = load_data()
	net = Network([784, 10])
	SGD!(net, training_data, 5, 10, 3.0, test_data=test_data)
	return net
end


"""
TT version of the restarted GMRES algorithm
"""

#takes a hessenberg matrix and returns q,H resp. orthogonal and upper triangular such that q^T H_out = H_in
function qr_hessenberg(H)
    m = size(H,2)
    q = Matrix{Float64}(I,m+1,m+1)
    T = H
    for i in 1:m
        R = Matrix{Float64}(I,m+1,m+1)
        s = T[i+1,i]/sqrt(T[i,i]^2+T[i+1,i]^2)
        c = T[i,i]/sqrt(T[i,i]^2+T[i+1,i]^2)
        R[i:(i+1),i:(i+1)] = [c s; -s c]
        q = R*q
        T = R*T
    end
    return q,T
end

function tt_gmres(A::ttoperator,b::ttvector,x0::ttvector;Imax=500,tol=1e-8,m=30,hist=false,γ_list=Float64[])
    V = Array{ttvector}(undef,m)
    W = Array{ttvector}(undef,m)
    H = zeros(m+1,m)
    r0 = tt_compression_par(tt_add(b,mult_a_tt(-1.0,tt_compression_par(mult(A,x0)))))
    β = sqrt(tt_dot(r0,r0))
    V[1] = mult_a_tt(1/β,r0)
    W[1] = tt_compression_par(mult(A,V[1]))
    H[1,1] = tt_dot(W[1],V[1])
    W[1] = tt_compression_par(tt_add(W[1],mult_a_tt(-H[1,1],V[1])))
    H[2,1] = sqrt(tt_dot(W[1],W[1]))
    q,r = qr_hessenberg(H[1:2,1])
    γ = abs(β*q[2,1]) #γ = \| Ax_j-b \|_2
    if hist
        γ_list = vcat(γ_list,γ)
    end
    j = 1
    if Imax <=0 || isapprox(H[j+1,j],0.,atol=tol) || isapprox(γ,0,atol=tol)
        if hist
            return x0,γ_list,H[2,1]
        else
            return x0,[γ],H[2,1]
        end
    else
        while j <= min(m-1,Imax) && !isapprox(H[j+1,j],0.,atol=tol) && !isapprox(γ,0,atol=tol)
            V[j+1] = mult_a_tt(1/H[j+1,j],W[j])
            j+=1
            W[j] = mult(A,V[j])
            W[j] = tt_compression_par(W[j])
            for i in 1:j
                H[i,j] = tt_dot(W[j],V[i])
                W[j] = tt_add(W[j],mult_a_tt(-H[i,j],V[i]))
                W[j] = tt_compression_par(W[j])
            end
            H[j+1,j] = sqrt(tt_dot(W[j],W[j]))
            q,r = qr_hessenberg(@view H[1:(j+1),1:j])
            γ = abs(β*q[j+1,1])
            if hist
                γ_list = vcat(γ_list,γ)
            end
        end
        z = r[1:j,1:j]\q[1:j,1]
        for i in 1:j
            x0 = tt_add(x0,mult_a_tt(β*z[i],V[i]))
        end
        x0 = tt_compression_par(x0)
        return tt_gmres(A,b,x0,tol=tol,Imax=Imax-j,m=m,hist=hist,γ_list=γ_list)
    end
end

function tt_cg(A::ttoperator,b::ttvector,x0::ttvector;Imax=500,tol=1e-8)
    p = tt_compression_par(tt_add(b,mult_a_tt(-1.0,tt_compression_par(mult(A,x0)))))
    r = p
    j=1
    res= zeros(Imax)
    res[1] = sqrt(tt_dot(p,p))
    while j < Imax && res[j]>tol
        Ap = mult(A,p)
        Ap = tt_compression_par(Ap)
        a = res[j]^2/tt_dot(p,Ap)
        x0 = tt_add(x0,mult_a_tt(a,p))
        x0 = tt_compression_par(x0)
        r = tt_add(r,mult_a_tt(-a,Ap))
        r = tt_compression_par(r)
        res[j+1] = sqrt(tt_dot(r,r))
        p = tt_add(r,mult_a_tt(res[j+1]^2/res[j]^2,p))
        p = tt_compression_par(p)
        j+=1
    end
    return x0, res[1:j]
end

"""
A_k : n_k x n_k x R_{k-1} x R_k
X_k : n_k x r^X_{k-1} x r^X_k
B_k : n_k x r^B_{k-1} x r^B_k
Ql : R_{k-1} r^X_{k-1} x r^B_{k-1}
Qr : R_k r^X_k x r^B_k
"""
function init_core(A_k,dim_X,B_k,Ql,Qr)
    X_k = zeros(dim_X...)
    #B right shape
    B = reshape(B_k,:,size(B_k,3)) # n_k r^B_{k-1} x r^B_k
    B = reshape(B*(Qr'),size(B_k,1),size(B_k,2),size(A_k,4),size(X_k,3)) #n_k x r^B_{k-1} x R_k x r^X_k
    B = permutedims(B,[2,1,3,4]) #r^B_{k-1} x n_k x R_k x r^X_k
    B = Ql*reshape(B,size(B,1),:) #R_{k-1} r^X_{k-1} x n_k R_k r^X_k
    B = permutedims(reshape(B,size(A_k,3),size(X_k,2),size(A_k,1),size(A_k,4),size(X_k,3)),[3,1,4,2,5]) #n_k x R_{k-1} x R_k x r^X_{k-1} x r^X_k
    B = reshape(B,size(A_k,1)*size(A_k,3)*size(A_k,4),dim_X[2]*dim_X[3]) #n_k R_{k-1} R_k x r^X_{k-1} r^X_k
    A = reshape(permutedims(A_k,[1,3,4,2]),size(A_k,1)*size(A_k,3)*size(A_k,4),size(A_k,2))
    ua,sa,va = svd(A)
    X = va*inv(Diagonal(sa))*ua'*B
    println(norm(A*X-B))
    return reshape(X,dim_X[1],dim_X[2],dim_X[3])
end

function rand_norm(n,m) #m>=n
    A = randn(n,m)
    for i in 1:m
        A[:,i] = A[:,i]./norm(A[:,i])
    end
    return A
end

function rand_struct_orth(r_A,r_X,r_b)
    A = zeros(r_X,r_A,r_b)
    q1 = rand_norm(r_A,r_b)
    q2 = rand_orthogonal(r_b,r_X)
    for ia in 1:r_A
        for ix in 1:r_X
            for ib in 1:r_b
                A[ix,ia,ib] = q1[ia,ib]*q2[ib,ix]
            end
        end
    end
    return reshape(A,r_A*r_X,r_b)
end

function init(A::ttoperator,b::ttvector,opt_rks)
    @assert(A.tto_dims == b.ttv_dims,DimensionMismatch)
    d = length(A.tto_dims)
    opt_rks = vcat([1],opt_rks)
    Q_list = Array{Array{Float64},1}(undef,d+1)
    ttvec = Array{Array{Float64},1}(undef,d)
    Q_list[1] = [1]
    Q_list[d+1] = [1]
    for k in 1:(d-1)
        Q_list[k+1] = rand_struct_orth(A.tto_rks[k],opt_rks[k+1],b.ttv_rks[k])
    end
    for k in 1:d
        ttvec[k] = init_core(A.tto_vec[k],[A.tto_dims[k],opt_rks[k],opt_rks[k+1]],b.ttv_vec[k],Q_list[k],Q_list[k+1])
    end
    return ttvector(ttvec,A.tto_dims,opt_rks[2:(d+1)],ones(Int64,d))
end

#automatically determines the initial tt ranks
function init_adapt(A::ttoperator,b::ttvector)
    d = length(A.tto_dims)
    opt_rks = ones(Int64,d)
    for k in 1:(d-1)
        opt_rks[k] = lcm(A.tto_rks[k],b.ttv_rks[k])/A.tto_rks[k]
    end
    println(opt_rks)
    return init(A,b,opt_rks)
end


import DSGE: n_model_states, n_backward_looking_states

path = dirname(@__FILE__)

### Model
m = KrusellSmith()

### Steady State
@load "$path/reference/steady_state.jld2" saved_l saved_c saved_mu saved_K saved_KF

@testset "Steady State" begin
    @test saved_l ≈ m[:lstar].value
    @test saved_c ≈ m[:cstar].value
    @test saved_mu ≈ m[:μstar].value
    @test saved_K ≈ m[:Kstar].value
    @test @test_matrix_approx_eq saved_KF m[:KFstar].value
end

### Jacobian
m.testing = true
JJ       = DSGE.jacobian(m)
@load "$path/reference/jacobian.jld2" saved_JJ nw

# we will always order things XP YP X Y
# convention is that capital letters generally refer to indices
# XP
LMP   = 1     :nw
LELLP = nw+1  :2*nw
KKP   = 2*nw+1

# YP
ZP    = 2*nw+2
MP    = 2*nw+3:3*nw+2
ELLP  = 3*nw+3:4*nw+2

# X
LM    = 4*nw+3:5*nw+2
LELL  = 5*nw+3:6*nw+2
KK    = 6*nw+3

# Y
Z     = 6*nw+4
M     = 6*nw+5:7*nw+4
ELL   = 7*nw+5:8*nw+4

# create objects needed for solve.jl
F1 = 1     :nw
F2 = nw+1  :2*nw
F3 = 2*nw+1:3*nw
F4 = 3*nw+1:4*nw
F5 = 4*nw+1:4*nw+1
F6 = 4*nw+2:4*nw+2

@testset "Jacobian" begin
    @testset "Euler Equation" begin
        @test @test_matrix_approx_eq saved_JJ[F1, KKP]  JJ[F1, KKP]
        @test @test_matrix_approx_eq saved_JJ[F1, ZP]   JJ[F1, ZP]
        @test @test_matrix_approx_eq saved_JJ[F1, ELLP] JJ[F1, ELLP]
        @test @test_matrix_approx_eq saved_JJ[F1, ELL]  JJ[F1, ELL]
    end

    @testset "Kolmogorov Forward Equation" begin
        @test @test_matrix_approx_eq saved_JJ[F2, LM]   JJ[F2, LM]
        @test @test_matrix_approx_eq saved_JJ[F2, LELL] JJ[F2, LELL]
        @test @test_matrix_approx_eq saved_JJ[F2, KK]   JJ[F2, KK]
        @test @test_matrix_approx_eq saved_JJ[F2, Z]    JJ[F2, Z]
        @test @test_matrix_approx_eq saved_JJ[F2, M]    JJ[F2, M]
    end

    @testset "LM(t+1) = M(t)" begin
        @test @test_matrix_approx_eq saved_JJ[F3, LMP]   JJ[F3, LMP]
        @test @test_matrix_approx_eq saved_JJ[F3, M]     JJ[F3, M]
    end

    @testset "LELL(t+1) = ELL(t)" begin
        @test @test_matrix_approx_eq saved_JJ[F4, LELLP]   JJ[F4, LELLP]
        @test @test_matrix_approx_eq saved_JJ[F4, ELL]     JJ[F4, ELL]
    end

    @testset "LOM K" begin
        @test @test_matrix_approx_eq saved_JJ[F5, ELL]   JJ[F5, ELL]
        @test @test_matrix_approx_eq saved_JJ[F5, M]     JJ[F5, M]
        @test @test_matrix_approx_eq saved_JJ[F5, KKP]   JJ[F5, KKP]
    end

    @testset "TFP" begin
        @test @test_matrix_approx_eq saved_JJ[F6, ZP]    JJ[F6, ZP]
        @test @test_matrix_approx_eq saved_JJ[F6, Z]     JJ[F6, Z]
    end
end

### Klein
m.testing = false
TTT_jump, TTT_state = klein(m)
@load "$path/reference/klein.jld2" gx hx

@testset "Solve" begin
    @test @test_matrix_approx_eq gx TTT_jump
    @test @test_matrix_approx_eq hx TTT_state
end

### Filtering

# Setup
N    = 200
m    = KrusellSmith()
endo = m.endogenous_states

# Model Solution/Transition Equation
TTT, RRR, CCC = solve(m)

# Measurement Equation
meas = measurement(m, TTT, RRR, CCC)
ZZ  = meas[:ZZ]
DD  = meas[:DD]
EE  = fill(0.1, (1,1))
QQ  = meas[:QQ]

# Generate measurement errors and shocks
srand(42)
u_t     = EE*randn(n_observables(m), N)
ε_t     = QQ*randn(n_shocks_exogenous(m), N)

####################################
# Simulate states forward N periods
####################################
s_t = zeros(n_model_states(m), N)
for t in 1:N
    # Initialize with shock around steady state
    if t == 1
        s_t[:, t] = RRR*ε_t[:, t]
    else
        s_t[:, t] = TTT*s_t[:, t-1] + RRR*ε_t[:, t]
    end
end

# Simulation in terms of grid values
simulated_log_gdp = vec(s_t[endo[:z′_t], :] + (m[:α]/m[:Kstar])*s_t[endo[:K′_t], :])
meas_log_gdp      = simulated_log_gdp + u_t[1, :]

##########
# Filter!
##########

# Measurement equation
data = ZZ*s_t + u_t   # Should be the same as meas_log_gdp

# Initialize with 0 initial condition
s_0     = zeros(n_model_states(m))
P_0     = RRR*QQ*RRR'

# Testing the simulation and filter setup
file = jldopen("$path/reference/simulate_and_filter.jld2", "r")
saved_u_t  = read(file, "u_t")
saved_ε_t  = read(file, "eps_t")
saved_s_t  = read(file, "s_t")
saved_s_jump_t = read(file, "s_jump_t")
saved_simulated_log_gdp = read(file, "simulated_log_gdp")
saved_meas_log_gdp = read(file, "meas_log_gdp")
saved_data = read(file, "data")
saved_s_0  = read(file, "s_0")
saved_P_0  = read(file, "P_0")
close(file)

@testset "Simulate and Filter" begin
    @test saved_u_t ≈ u_t
    @test saved_ε_t ≈ ε_t
    @test saved_s_t ≈ s_t[1:n_backward_looking_states(m), :]
    @test saved_s_jump_t ≈ s_t[n_backward_looking_states(m)+1:end, :]
    @test saved_simulated_log_gdp ≈ simulated_log_gdp
    @test saved_meas_log_gdp ≈ meas_log_gdp
    @test saved_data ≈ data
    @test saved_s_0 ≈ s_0[1:n_backward_looking_states(m)]
    @test saved_P_0 ≈ P_0[1:n_backward_looking_states(m), 1:n_backward_looking_states(m)]
end

# Testing the filter inputs against the outputs
EE_zero = zeros(1, 1)
# reference_output = load("$path/reference/filter_inputs_output.jld2")
@load "$path/reference/filter_inputs_output.jld2" y loglik
test = kalman_filter(y, TTT, RRR, CCC, QQ, ZZ,
                     DD, EE_zero, s_0, P_0)

@testset "Check Filter Output" begin
    @test test[1] ≈ loglik
end


@doc raw"""
Abstract type for indicator functions $\mathbb{R}^{2n} \rightarrow
\left\{0, 1\right\}$ where $n = 1, 2 \text{ or } 3$.
"""
abstract type AbstractIndicator end

@doc raw"""
    SeparableIndicator(χ₁, χ₂)

Type for separable indicator function, that is for such an indicator
function which can be written as $\chi(x,y) = \chi_1(x)\chi_2(y)$.

`χ1` and `χ2` must be functions of one argument which return a value
of `Bool` type.

**NB**: This indicator function is not symmetric (i.e. $\chi(x,y) \ne
\chi(y,x)$). This behaviour is intentional. For example you can write
such an indicator, so the corresponding correlation function is
sensitive to the spatial orientation of a system.

*"That one, too fat! This one, too tall! This one… too symmetrical!"*
"""
struct SeparableIndicator <: AbstractIndicator
    χ1 :: Function
    χ2 :: Function
end

"""
    SeparableIndicator(χ)

This is equvalent of writing `SeparableIndicator(χ, χ)`.
"""
SeparableIndicator(χ :: Function) = SeparableIndicator(χ, χ)

@doc raw"""
    InseparableIndicator(χ)

Type for inseparable indicator function, that is for such an indicator
function which cannot be written as $\chi(x,y) = \max(\Chi(x, y), \Chi(y,
x))$, where $\Chi(x, y) = \chi_1(x)\chi_2(y)$.

`χ` must be a function of two arguments which returns a value of `Bool`
type.
"""
struct InseparableIndicator <: AbstractIndicator
    χ :: Function
end

indicator_function(x :: InseparableIndicator) = x.χ
indicator_function(x :: SeparableIndicator)   = x.χ1, x.χ2


using Flux.Tracker, Statistics, DelimitedFiles
using Flux.Tracker: Params, gradient, update!
using DelimitedFiles, Statistics
using Flux: gpu

# This replicates the housing data example from the Knet.jl readme. Although we
# could have reused more of Flux (see the mnist example), the library's
# abstractions are very lightweight and don't force you into any particular
# strategy.

cd(@__DIR__)

isfile("housing.data") ||
  download("https://raw.githubusercontent.com/MikeInnes/notebooks/master/housing.data",
           "housing.data")

rawdata = readdlm("housing.data")'

# The last feature is our target -- the price of the house.
split_ratio = 0.1 # For the train test split

x = rawdata[1:13,:] |> gpu
y = rawdata[14:14,:] |> gpu

# Normalise the data
x = (x .- mean(x, dims = 2)) ./ std(x, dims = 2)

# Split into train and test sets
split_index = floor(Int,size(x,2)*split_ratio)
x_train = x[:,1:split_index]
y_train = y[:,1:split_index]
x_test = x[:,split_index+1:size(x,2)]
y_test = y[:,split_index+1:size(x,2)]

# The model
W = param(randn(1,13)/10) |> gpu
b = param([0.]) |> gpu

predict(x) = W*x .+ b
meansquarederror(ŷ, y) = sum((ŷ .- y).^2)/size(y, 2)
loss(x, y) = meansquarederror(predict(x), y)

η = 0.1
θ = Params([W, b])

for i = 1:10
  g = gradient(() -> loss(x_train, y_train), θ)
  for x in θ
    update!(x, -g[x]*η)
  end
  @show loss(x_train, y_train)
end

# Predict the RMSE on the test set
err = meansquarederror(predict(x_test),y_test)
println(err)


@testset "592.fraction-addition-and-subtraction.jl" begin
    @test fraction_addition("-1/2+1/2") == "0/1"
    @test fraction_addition("-1/2+1/2+1/3") == "1/3"
    @test fraction_addition("1/3-1/2") == "-1/6"
    @test fraction_addition("5/3+1/3") == "2/1"
end


const EARTH_RADIUS = 6371.0  # km

struct Layer
    name::AbstractString
    r_min
    r_max
    density::Polynomial
end

struct EarthModel
    layers::Vector{Layer}
end

PREM = EarthModel(
    [
        Layer("Inner Core", 0.0, 1221.5, Polynomial([13.0885, 0.0, -8.8381])),
        Layer("Outer Core", 1221.5, 3480.0, Polynomial([12.5815, -1.2638, -3.6426, -5.5281])),
        Layer("Lower Mantle", 3480.0, 5701.0, Polynomial([7.9565, -6.4761, 5.5283, -3.0807])),
        Layer("Transition Zone 1", 5701.0, 5771.0, Polynomial([5.3197, -1.4836])),
        Layer("Transition Zone 2", 5771.0, 5971.0, Polynomial([11.2494, -8.0298])),
        Layer("Transition Zone 3", 5971.0, 6151.0, Polynomial([7.1089, -3.8045])),
        Layer("LVZ", 6151.0, 6291.0, Polynomial([2.6910, 0.6924])),
        Layer("LID", 6291.0, 6346.6, Polynomial([2.6910, 0.6924])),
        Layer("Crust 1", 6346.6, 6356.0, Polynomial([2.9])),
        Layer("Crust 2", 6356.0, 6368.0, Polynomial([2.6])),
        Layer("Ocean", 6368.0, 6371.0, Polynomial([1.020]))
    ]
)

function (m::EarthModel)(r)
    for layer ∈ m.layers
        if r < layer.r_max && r >= layer.r_min
            return layer.density(r / m.layers[end].r_max)
        end
    end
end

function tracklength(zenith, zposition)
"""
$(SIGNATURES)

Total path length through earth from detector position

# Arguments
- `zenith`: Zenith angle of the path with respect to the detector frame [rad]
- `zposition` Distance below the surface of the Earth (positive value) [km]
"""
    if isapprox(zenith, pi)
        return 2*EARTH_RADIUS - zposition
    elseif isapprox(zenith, 0)
        return zposition
    end

    theta = pi - zenith
    zprime = EARTH_RADIUS - zposition
    sintheta = sin(theta)
    costheta = cos(theta)
    halfcoord = sqrt(EARTH_RADIUS^2 - (zprime * sintheta)^2)
    chordsection = zprime * costheta
    if isapprox(zenith, pi/2)
        return halfcoord
    else
        return halfcoord + chordsection
    end
end

function prempath(zenith::Vector{Float64}, zposition; samples=100, discrete_densities=nothing)
"""
$(SIGNATURES)

# Arguments
- `zenith::Vector{Float64}`: Zenith angles of the paths with respect to the detector frame [rad]
- `zposition::Float64` Distance below the surface of the Earth (positive value) [km]
- `samples` The number of steps with equal distance
- `discrete_densities` List of density values to be used for discretization
"""
    map(z->prempath(z, zposition; samples=samples, discrete_densities=discrete_densities), zenith)
end

function prempath(zenith::T, zposition; samples=100, discrete_densities=nothing) where {T <: Number}
"""
$(SIGNATURES)

# Arguments
- `zenith::Float64`: Zenith angle of the path with respect to the detector frame [rad]
- `zposition::Float64` Distance below the surface of the Earth (positive value) [km]
- `samples` The number of steps with equal distance
"""
    trklen = tracklength(zenith, zposition)
    x = Array(range(0.0; stop=trklen, length=samples))
    sections = (x[2:end] - x[1:end-1])
    total_pathlen = 0.5 * (x[2:end] + x[1:end-1])
    zprime = EARTH_RADIUS - zposition
    radii = map(x -> sqrt(zprime^2 + x^2 + 2*zprime*x*cos(zenith)), total_pathlen)
    densities = PREM.(radii)
    if !isnothing(discrete_densities)
        idx = map(d->searchsortedfirst(discrete_densities,d), densities)
        densities = map(i->discrete_densities[i], idx)
    end
    Path(densities, sections)
end


using Base: BitSignedSmall
# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    MinDistanceSampling(α, ρ=0.65, δ=100, metric=Euclidean())

Generate sample from geometric object such that all pairs of
points are at least `α` units of distance away from each other.
Optionally specify the relative radius `ρ` for the packing
pattern, the oversampling factor `δ` and the `metric`.

This method is sometimes referred to as Poisson disk sampling
or blue noise sampling in the computer graphics community.

## References

* Lagae, A. & Dutré, P. 2007. [A Comparison of Methods for
  Generating Poisson Disk Distributions]
  (https://onlinelibrary.wiley.com/doi/abs/10.1111/j.1467-8659.2007.01100.x)
* Bowers et al. 2010. [Parallel Poisson disk sampling with
  spectrum analysis on surfaces](https://dl.acm.org/doi/10.1145/1882261.1866188)
* Medeiros et al. 2014. [Fast adaptive blue noise on polygonal surfaces]
  (https://www.sciencedirect.com/science/article/abs/pii/S1524070313000313)
"""
struct MinDistanceSampling{T,M} <: ContinuousSamplingMethod
  α::T
  ρ::T
  δ::Int
  metric::M
end

MinDistanceSampling(α::T; ρ=T(0.65), δ=100, metric=Euclidean()) where {T} =
  MinDistanceSampling(α, ρ, δ, metric)

function sample(rng::AbstractRNG, Ω::DomainOrData, method::MinDistanceSampling)
  # retrive parameters
  α = method.α
  ρ = method.ρ
  δ = method.δ
  m = method.metric

  # total volume/area of the object
  V = sum(measure, Ω)

  # expected number of Poisson samples
  # for relative radius (Lagae & Dutré 2007)
  N = 2V/√3 * (ρ/α)^2

  # number of oversamples (Medeiros et al. 2014)
  O = ceil(Int, δ * N)

  # oversample the object
  points = sample(rng, Ω, HomogeneousSampling(O))

  # collect points into point set
  𝒫 = PointSet(collect(points))

  # discard points that do not satisfy distance criterion
  sample(rng, 𝒫, BallSampling(α, metric=m))
end


using Manifolds, Manopt, Test

@testset "Nelder Mead Options" begin
    o = NelderMeadOptions(Euclidean(2))
    o2 = NelderMeadOptions(Euclidean(2), o.population)
    @test o.x == o2.x
    @test o.population == o2.population
    @test get_options(o) == o
end


# m6.8s.jl

using Pkg, DrWatson

using MonteCarloMeasurements
using StanSample
using StatisticalRethinking

N = 100
df = DataFrame(
    :h0 => rand(Normal(10,2 ), N),
    :treatment => vcat(zeros(Int, Int(N/2)), ones(Int, Int(N/2)))
);
df[!, :fungus] = [rand(Binomial(1, 0.5 - 0.4 * df[i, :treatment]), 1)[1] for i in 1:N]
df[!, :h1] = [df[i, :h0] + rand(Normal(5 - 3 * df[i, :fungus]), 1)[1] for i in 1:N]

stan6_8 = "
data {
    int <lower=1> N;
    vector[N] h0;
    vector[N] h1;
    vector[N] treatment;
}
parameters{
    real a;
    real bt;
    real<lower=0> sigma;
}
model {
  vector[N] mu;
    vector[N] p;
    a ~ lognormal(0, 0.2);
    bt ~ normal(0, 0.5);
    sigma ~ exponential(1);
    for ( i in 1:N ) {
    p[i] = a + bt*treatment[i];
    mu[i] = h0[i] * p[i];
    }
    h1 ~ normal(mu, sigma);
}
"

m6_8_data = Dict(
    :N => nrow(df),
    :h0 => df[:, :h0],
    :h1 => df[:, :h1],
    :treatment => df[:, :treatment]
)

m6_8s = SampleModel("m6.8s", stan6_8)

rc6_8s = stan_sample(m6_8s; data=m6_8_data)

if success(rc6_8s)
    part6_8s = read_samples(m6_8s, :particles);
    part6_8s |> display
end

# End of m6.8s.jl


"""
    uniform(dims...)

Return an `Array` of size `dims` containing random variables taken from a uniform
distribution in the interval ``[-1, 1]``.

# Examples
```jldoctest; setup = :(using Random; Random.seed!(0))
julia> uniform(2, 3)
2×3 Array{Float32,2}:
 0.548721  -0.524116  -0.743921
 0.822376   0.735768   0.0525029
```
"""
uniform(dims...) = (rand(dims...) .- 0.5f0) * 2

"""
    SIREN_init(SIREN_init(omega_0::Real, is_first::Bool, dims::Integer...)

Return an `Array` of size `dims` containing random variables.
The distribution is supposed to be an implementation of the SIREN initialization
from the original paper (Sec 3.2, Appdx 1, Appdx 1.5). My explanation will probably
be wrong so check out their explanation.

# Examples
```jldoctest; setup = :(using Random; Random.seed!(0))
julia> SIREN_init(30,false,3,2)
3×2 Array{Float64,2}:
  0.0305138  -0.0304217
  0.0386888  -0.0208474
 -0.031625   -0.0279565
```
"""
function SIREN_init(omega_0::Real, is_first::Bool, dims::Integer...)
    w = uniform(dims...)

    fan_in = 1
    s = size(w)
    for i in 1:(length(s)-1)
      fan_in *= s[i]
    end

    w .*= sqrt(6/fan_in)

    if !is_first
      w ./= omega_0
    end

    return w
end