using Flux, Statistics
using Flux: onehotbatch, onecold, logitcrossentropy
using Printf, BSON

include("model-commons.jl")
include("myplots.jl")

batch_size = 100
img_n = 48
img_m = 48

function make_batch_indices(i, n, batch_size=batch_size)
  istart = (i - 1) * batch_size + 1
  iend = if (istart+batch_size-1 > n) n else istart+batch_size-1 end
  convert(Int,istart):convert(Int,iend)
end

function train_minibatches(X, Y, batch_size=batch_size, img_n=img_n, img_m=img_m)
  n = size(X,2)
  bcount = ceil(n/batch_size)
  batches = []
  for i in 1:bcount
    inds = make_batch_indices(i, n)
    push!(batches, (reshape(X[:,inds], (img_n, img_m, 1, length(inds))), Y[:,inds]))
  end
  batches
end

function test_minibatch(X, Y, n=img_n, m=img_m)
  bsize = size(X,2)
  inds = make_batch_indices(1, bsize, bsize)
  (reshape(X[:,inds], (n, m, 1, bsize)), Y[:,inds])
end

dataset = create_data_splits()
X_train = reshape(dataset.X_train, (img_n, img_m, 1, size(dataset.X_train,2)))
X_test = reshape(dataset.X_test, (img_n, img_m, 1, size(dataset.X_test,2)))
train_set = train_minibatches(dataset.X_train, dataset.Y_train)
test_set = test_minibatch(dataset.X_test, dataset.Y_test)

@info("Constructing model...")
model = Chain(
    # First convolution, operating upon a 48x48 images
    Conv((8, 8), 1=>16, pad=(1,1), stride=2, relu),

    # Second convolution, operating upon a 22x22 image
    Conv((8, 8), 16=>32, pad=(1,1), stride=2, relu),

    # Reshape 3d tensor into a 2d one, at this point it should be (9, 9, 32, N)
    # which is where we get the 2592 in the `Dense` layer below:
    x -> reshape(x, :, size(x, 4)),
    Dense(2592, 10)
) |> gpu

# Make sure our model is nicely precompiled before starting our training loop
model(train_set[1][1])

loss(x, y) = logitcrossentropy(model(x), y)

function accuracy(x, y)
  logits = model(x)
  probs = [logsoftmax(logits[:,i]) for i in 1:size(logits,2)]
  mean(onecold.(probs) .== onecold(y))
end

opt = ADAM(0.001)

@info("Beginning training loop...")
best_acc = 0.0
last_improvement = 0
for epoch_idx in 1:100
    global best_acc, last_improvement
    # Train for a single epoch
    Flux.train!(loss, params(model), train_set, opt)

    # Calculate metrics:
    acc = accuracy(test_set...)
    err = loss(X_train, dataset.Y_train).data
    @info(@sprintf("[%d]: Test accuracy: %.4f", epoch_idx, acc))
    @info(@sprintf("[%d]: Current loss: %.4f", epoch_idx, err))

    # If our accuracy is good enough, quit out.
    if acc >= 0.999
        @info(" -> Early-exiting: We reached our target accuracy of 99.9%")
        break
    end

    # If this is the best accuracy we've seen so far, save the model out
    if acc >= best_acc
        @info(" -> New best accuracy! Saving model out to sound_cnn2.bson")
        BSON.@save "sound_cnn2.bson" model epoch_idx acc
        best_acc = acc
        last_improvement = epoch_idx
    end

    # If we haven't seen improvement in 5 epochs, drop our learning rate:
    if epoch_idx - last_improvement >= 3 && opt.eta > 1e-5
        opt.eta /= 10.0
        @warn(" -> Haven't improved in a while, dropping learning rate to $(opt.eta)!")

        # After dropping learning rate, give it a few epochs to improve
        last_improvement = epoch_idx
    end

    if epoch_idx - last_improvement >= 10
        @warn(" -> We're calling this converged.")
        break
    end
end

# Best Accuracy: 77.4% - File sound_cnn.bson
#   batch_size       = 100
#   spectrogram size = 48x48
#   eta              = 0.001
#   model:
#       Conv((8, 8), 1=>16, pad=(1,1), stride=2, relu)
#       Conv((8, 8), 16=>32, pad=(1,1), stride=2, relu),
#       Dense(2592, 10)
#       softmax


function xlogx(x)
    x*log(x)
end


using ITensors
using ITensorUnicodePlots
using PastaQ
using Zygote

import ITensors: op

function op(
  ::OpName"CRy",
  ::SiteType"S=1/2"; θ)
  c = cos(θ/2)
  s = sin(θ/2)
  return [
  1 0 0 0
  0 1 0 0
  0 0 c -s
  0 0 s c]
end

i1 = Index(2, "S=1/2,i1")
i2 = Index(2, "S=1/2,i2")

Xp1 = state("Xp", i1)
Xm1 = state("Xm", i1)
Xp2 = state("Xp", i2)
Xm2 = state("Xm", i2)
Zp1 = state("Zp", i1)
Zm1 = state("Zm", i1)
Zp2 = state("Zp", i2)
Zm2 = state("Zm", i2)

Id1 = op("Id", i1)
Z1 = op("Z", i1)
X1 = op("X", i1)
Id2 = op("Id", i2)
Z2 = op("Z", i2)
X2 = op("X", i2)

i1 != i2

ZpZm = ITensor(i1, i2)
ZpZm[i1=>1, i2=>2] = 1

Zp1 = state("Zp", i1)
Zp2 = state("Zp", i2)
Zm1 = state("Zm", i1)
Zm2 = state("Zm", i2)

ZpZm = Zp1 * Zm2

ψ = ITensor(i1, i2)
ψ[i1=>1, i2=>2] = 1/√2
ψ[i1=>2, i2=>1] = 1/√2

ψ = (Zp1 * Zm2 + Zm1 * Zp2)/√2

inner(ZpZm, ψ) == 1/√2

_, S, _ = svd(ZpZm, i1);
diag(S) == [1, 0]

_, S, _ = svd(ψ, i1);
diag(S) == [1/√2, 1/√2]

H = ITensor(i1', i2',
i1, i2)
H[i1'=>2, i2'=>1,
i1=>2, i2=>1] = -1
# …

apply(X1, Zp1 * Zm2) ≈ Zm1 * Zm2

@visualize X1 * ZpZm

ZZ = Z1 * Z2
XI = X1 * Id2
IX = Id1 * X2
h = 0.5
H = -ZZ + h * (XI + IX)

inner(Zp1' * Zp2', H, Zp1 * Zp2)

D, _ = eigen(H);
diag(D) ≈ [-√2, -1, 1, √2]

CRy = op("CRy", i1, i2; θ=π/2)

apply(CRy, Zm1 * Zp2) ≈ Zm1 * Xp2

# Optimization
function minimize(f, ∂f, x; nsteps, γ)
  for n in 1:nsteps
    x = x - γ * ∂f(x)
    println("Step = $(n), f(x) = $(f(x))")
  end
  return x
end

# State optimization
function E(ψ)
  ψHψ = inner(ψ', H, ψ)
  ψψ = inner(ψ, ψ)
  return ψHψ / ψψ
end
∂E(ψ) = gradient(E, ψ)[1]

ψ0 = (Zp1 * Zp2 + Zm1 * Zm2)/√2
E(ψ0) == -1
norm(∂E(ψ0)) == 2

ψ = minimize(E, ∂E, ψ0; nsteps=10, γ=0.1)
E(ψ0) == -1
E(ψ) # ≈ -1.4142131 ≈ -√2

# Circuit optimization
U(θ,i1,i2) = [
  op("Ry",i1;θ=θ[1]),
  op("Ry",i2;θ=θ[2]),
  op("CX",i1,i2),
  op("Ry",i1;θ=θ[3]),
  op("Ry",i2;θ=θ[4])]

u(θ) = [
  ("Ry",1,(;θ=θ[1])),
  ("Ry",2,(;θ=θ[2])),
  ("CNOT",1,2),
  ("Ry",1,(;θ=θ[3])),
  ("Ry",2,(;θ=θ[4]))]

U(θ,i1,i2) = buildcircuit([i1,i2],u(θ))

ψ0 = Zp1 * Zp2
function E(θ)
  ψθ = apply(U(θ,i1,i2),ψ0)
  return inner(ψθ',H,ψθ)
end
∂E(ψ) = gradient(E, ψ)[1]

θ0 = [0, 0, 0, 0]
θ = minimize(E, ∂E, θ0; nsteps=40, γ=0.5)
E(θ0) == -1
E(θ) # ≈ -1.4142077 ≈ -√2

# Fidelity optimization
ψ0 = Zp1 * Zp2
ψ = (Zp1 * Zp2 + Zm1 * Zm2) / √2

function F(θ)
  ψθ = apply(U(θ,i1,i2),ψ0)
  return -abs2(inner(ψ,ψθ))
end
∂F(θ) = gradient(F, θ)[1]

θ0 = [0, 0, 0, 0]
θ = minimize(F, ∂F, θ0; nsteps=50, γ=0.1)
F(θ0) == -0.5
F(θ) # ≈ -0.9938992 ≈ -1


using ProofOfConcept
using Test

@testset "Barycentric coordinates" begin
    t = Triangle([0, 0], [1, 0], [0, 1])
    p1 = Point([0.1, 0.1])
    bc = get_barycentric_coordinates(p1, t)
    @test p1 == bc * t

    p2 = Point([0.2, 0.1])
    bc = get_barycentric_coordinates(p2, t)
    @test p2 == bc * t

    p3 = Point([0.8, 0.1])
    bc = get_barycentric_coordinates(p3, t)
    @test p3 == bc * t

    p4 = Point([1.8, -2.1])
    bc = get_barycentric_coordinates(p4, t)
    @test p4 == bc * t
end

@testset "is_inside" begin
    t = Triangle([0, 0], [1, 0], [0, 1])

    p1 = Point([0.1, 0.1])
    @test is_inside(p1, t)
    bc = get_barycentric_coordinates(p1, t)
    @test is_inside(bc)

    p2 = Point([0.2, 0.1])
    @test is_inside(p2, t)
    bc = get_barycentric_coordinates(p2, t)
    @test is_inside(bc)

    p3 = Point([0.8, 0.1])
    @test is_inside(p3, t)
    bc = get_barycentric_coordinates(p3, t)
    @test is_inside(bc)

    p4 = Point([1.8, -2.1])
    @test !is_inside(p4, t)
    bc = get_barycentric_coordinates(p4, t)
    @test !is_inside(bc)
end

@testset "Refinement criteria" begin
    t = Triangle([[0, 0] [1, 0] [0, 1]], [1, 2, 3])
    tol = 0.1

    p1 = Point([0.1, 0.1], 5)
    @test refinement_criteria(p1, t, tol)

    p1 = Point([0, 0], 1.09)
    @test !refinement_criteria(p1, t, tol)

end

@testset "Distribute points" begin

    h = reshape(collect(1:6*5), 6, 5)
    m = TerrainMap(0, 0, 10, 10, h)
    idx = CartesianIndices(h)

    t1 = Triangle([[15, 15] [45, 0] [20, 40]], [1, 2, 3])
    map_t1 = view(idx, findall(look_for_points_inside_triangle(m, t1), idx))

    @test all(map_t1 == [CartesianIndex(4, 2), CartesianIndex(3, 3), CartesianIndex(4, 3), CartesianIndex(3, 4), CartesianIndex(3, 5)])

    # Now we are splitting the triangle
    t2 = Triangle([[15, 15] [35, 15] [20, 40]], [1, 2, 3])
    map_t2 = view(map_t1, findall(look_for_points_inside_triangle(m, t2), map_t1))

    @test all(map_t2 == [CartesianIndex(3, 3), CartesianIndex(4, 3), CartesianIndex(3, 4), CartesianIndex(3, 5)])

    t3 = Triangle([[15, 15] [45, 0] [35, 15]], [1, 2, 3])
    map_t3 = view(map_t1, findall(look_for_points_inside_triangle(m, t3), map_t1))

    @test all(map_t3 == [CartesianIndex(4, 2)])

end

@testset "Distribute and mark for refinement" begin

    tol = 0.1

    h = reshape(collect(1:6*5), 6, 5)
    m = TerrainMap(0, 0, 10, 10, h)
    idx = CartesianIndices(h)

    t1 = Triangle([[15, 15] [45, 0] [20, 40]], [1, 2, 3])
    (refine_t1, map_t1) = distribute_and_mark_for_refinement(m, idx, t1, tol)
    @test refine_t1
    @test all(map_t1 == [CartesianIndex(4, 2), CartesianIndex(3, 3), CartesianIndex(4, 3), CartesianIndex(3, 4), CartesianIndex(3, 5)])

    # Now we are splitting t1 into t2 and t3
    t2 = Triangle([[15, 15] [35, 15] [20, 40]], [1, 2, 3])
    (refine_t2, map_t2) = distribute_and_mark_for_refinement(m, map_t1, t2, tol)
    @test refine_t2
    @test all(map_t2 == [CartesianIndex(3, 3), CartesianIndex(4, 3), CartesianIndex(3, 4), CartesianIndex(3, 5)])

    t3 = Triangle([[15, 15] [45, 0] [35, 15]], [1, 2, 3])
    (refine_t3, map_t3) = distribute_and_mark_for_refinement(m, map_t1, t3, tol)
    @test refine_t3
    @test all(map_t3 == [CartesianIndex(4, 2)])

    # Now we are splitting t2 into t4 and t5
    t4 = Triangle([[15, 15] [25, 15] [20, 40]], [1, 2, 3])
    (refine_t4, map_t4) = distribute_and_mark_for_refinement(m, map_t2, t4, tol)
    @test refine_t4
    @test all(map_t4 == [CartesianIndex(3, 3), CartesianIndex(3, 4), CartesianIndex(3, 5)])

    t5 = Triangle([[25, 15] [35, 15] [20, 40]], [1, 2, 3])
    (refine_t5, map_t5) = distribute_and_mark_for_refinement(m, map_t2, t5, tol)
    @test refine_t5
    @test all(map_t5 == [CartesianIndex(4, 3), CartesianIndex(3, 5)])
end


function is_it(u0,H,tol)
    prob = ODEProblem(Eq_of_M,u0,(0.,10.))
    sol=solve(prob, RK4(),reltol=1e-5,abstol=1e-7)
    A=sol[:,:]

    Z=A[1,:]+im *A[2,:]
    W=A[3,:]+im *A[4,:]
    A=zeros(0)
    h=1/(2*H)
    a=(2+h-2*sqrt(h+1))/h;
    N=length(sol.t)
    lin_impulse=im*(1+a )*ones(N)

    z_1pos=.5*(Z+lin_impulse+W)
    z_1neg=.5*(Z-lin_impulse-W)
    z_2pos=.5*(-Z+lin_impulse-W)
    z_2neg=.5*(-Z-lin_impulse+W)

    #1pos and 1neg bound
    d_11=abs.(z_1pos-z_1neg)
    #1pos and 2neg bound
    d_12=abs.(z_1pos-z_2neg)

    delta_11=abs(d_11[1]-d_11[end])
    delta_12=abs(d_12[1]-d_12[end])

    Z=zeros(0)
    W=zeros(0)
    lin_impulse=zeros(0)
    z_1pos=zeros(0)
    z_1neg=zeros(0)
    z_2pos=zeros(0)
    z_2neg=zeros(0)
    d_11=zeros(0)
    d_12=zeros(0)

    return delta_11<tol || delta_12<tol
    # return delta_11, delta_12

end


#-------------------------------------------------------------------------
function calculatekernelmatrix(X, Y, k, hyper)
#-------------------------------------------------------------------------

    K = zeros(length(X), length(Y))

    calculatekernelmatrix!(K, X, Y, k, hyper)

    return K

end

#-------------------------------------------------------------------------
function calculatekernelmatrix!(K, X, Y, k, hyper)
#-------------------------------------------------------------------------
    @assert(size(K, 1) == length(X))
    @assert(size(K, 2) == length(Y))

    # calculate kernel matrix
    for (i, xᵢ) in enumerate(X), (j, xⱼ) in enumerate(Y)
        @inbounds K[i,j] = k(xᵢ, xⱼ, hyper) + ((i==j) ? 1e-12 : 0.0)
    end
end

#-------------------------------------------------------------------------
function calculatekernelmatrix!(K, dist2, k, hyper)
#-------------------------------------------------------------------------
    @assert(size(K, 1) == size(dist2, 1))
    @assert(size(K, 2) == size(dist2, 2))

    copyto!(K, k(dist2, hyper))

    nothing
end

#-------------------------------------------------------------------------
function rbf(xᵢ, xⱼ, hyper)
#-------------------------------------------------------------------------
    makepos(hyper[1]) * exp(-0.5 * makepos(hyper[2]) * (xᵢ-xⱼ)^2)
end

#-------------------------------------------------------------------------
function rbf(dist2, hyper)
#-------------------------------------------------------------------------
    makepos(hyper[1]) .* exp.(-0.5 .* makepos(hyper[2]) .* dist2)
end

#-------------------------------------------------------------------------
function OU(xᵢ, xⱼ, hyper)
#-------------------------------------------------------------------------
    makepos(hyper[1]) * exp(-0.5 * makepos(hyper[2]) * abs(xᵢ-xⱼ))
end

#-------------------------------------------------------------------------
function OU(dist2, hyper)
#-------------------------------------------------------------------------
    makepos(hyper[1]) .* exp.(-0.5 .* makepos(hyper[2]) .* sqrt.(dist2))
end

#-------------------------------------------------------------------------
function makepos(x)
#-------------------------------------------------------------------------
    exp(x)
end


struct DGSM <: GSAMethod
    crossed::Bool
end

DGSM(;crossed::Bool = false) = DGSM(crossed)

mutable struct DGSMResult{T}
    a::Array{T,1}
    absa::Array{T,1}
    asq::Array{T,1}
    sigma::Array{T,1}
    tao::Array{T,1}
    crossed::Union{Nothing,Array{T,2}}
    abscrossed::Union{Nothing,Array{T,2}}
    crossedsq::Union{Nothing,Array{T,2}}
end

"""
The inputs for DGSM are as follows:
1.f:
    This is the input function based on which the values of DGSM are to be evaluated
    Eg- f(x) = x[1]+x[2]^2
        This is function in 2 variables
2.samples:
    Depicts the number of sampling set of points to be used for evaluation of E(a), E(|a|) and E(a^2)
    a = partial derivative of f wrt x_i
3.distri:
    Array of distribution of respective variables
    Eg- dist = [Normal(5,6),Uniform(2,3)]
    for two variables
4.crossed:
    A string(True/False) which act as indicator for computation of DGSM crossed indices
    Eg- a True value over there will lead to evauation of crossed indices
"""
function gsa(f, method::DGSM, distr::AbstractArray; samples::Int, kwargs...)

    k = length(distr)

    #XX is the matrix consisting of 'samples' number of sampling based on respective
    #distributions of variables

    XX = [rand.(distr) for x = 1:samples]

    #function to evaluate gradient of f wrt x
    grad(x)= ForwardDiff.gradient(f,x)

    #function to evaluate hessian of f wrt x
    hess(x) = ForwardDiff.hessian(f,x)

    #Evaluating the derivatives with AD

    dfdx = [grad(XX[i]) for i = 1:samples]
    dfdx = reduce(hcat,dfdx)
    dfdx = dfdx'

    #Evaluating E(a) E(|a|) and E(a^2)

    a = [mean(dfdx[:,i]) for i in 1:k]
    asq = [mean(dfdx[:,i].^2) for i in 1:k]
    absa = [mean(abs.(dfdx[:,i])) for i in 1:k]

    sigma = zeros(Float64,k)
    tao = zeros(Float64,k)

    #Evaluating tao_i for all input parameters

    for i in 1:k
        for j in 1:samples
            tao[i] += (dfdx[j + (i-1)*samples]^2)*(1 - 3*XX[j][i] + XX[j][i]^2)/6
        end
        tao[i] = tao[i]/samples
    end

    #Evaluating sigma_i for all input parameters

    for i in 1:k
        for j in 1:samples
            sigma[i] += 0.5*(XX[j][i])*(1-XX[j][i])*dfdx[j + (i-1)*samples]^2
        end
        sigma[i] = sigma[i]/samples
    end

    if method.crossed == true

        #Evaluating the derivatives with AD
        dfdxdy = [hess(XX[i]) for i in 1 : samples]

        crossed = zeros(Float64,k,k)
        crossedsq = zeros(Float64,k,k)
        abscrossed = zeros(Float64,k,k)

        #computing the elements of crossed, crossedsq, abscrossed

        for a in 1:k
            for b in a+1:k
                crossed[b + (a-1)*k] = mean(dfdxdy[i][b + (a-1)*k] for i in 1:samples)
                crossed[a + (b-1)*k] = crossed[b + (a-1)*k]
                crossedsq[b + (a-1)*k] = mean(dfdxdy[i][b + (a-1)*k]^2 for i in 1:samples)
                crossedsq[a + (b-1)*k] = crossedsq[b + (a-1)*k]
                abscrossed[b + (a-1)*k] = mean(abs(dfdxdy[i][b + (a-1)*k]) for i in 1:samples)
                abscrossed[a + (b-1)*k] = abscrossed[b + (a-1)*k]
            end
        end

    else
    	return DGSMResult(a, absa, asq, sigma, tao, nothing, nothing, nothing)
    end

    return DGSMResult(a, absa, asq, sigma, tao, crossed, abscrossed, crossedsq)
    #returns a struct of 7 elements i.e. a, absa, asq, sigma(all 4 arrays) and crossed, abscrossed, crossedsq (all 3 matrices)
end


"""Linear interpolation of a variable u in the x-direction.
m,n = size(ux) must be m+1,n = size(u)."""
function Ix!(ux::AbstractMatrix,u::AbstractMatrix)
    m, n = size(ux)
    @boundscheck (m+1,n) == size(u) || throw(BoundsError())

    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            ux[i,j] = one_half*(u[i+1,j] + u[i,j])
        end
    end
end

""" Linear interpolation a variable u in the y-direction.
    m,n = size(uy) must be m,n+1 = size(u)."""
function Iy!(uy::AbstractMatrix,u::AbstractMatrix)
    m,n = size(uy)
    @boundscheck (m,n+1) == size(u) || throw(BoundsError())

    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            uy[i,j] = one_half*(u[i,j+1] + u[i,j])
        end
    end
end

""" Bilinear interpolation a variable u in x and y-direction.
m,n = size(uxy) must be m+1,n+1 = size(u). """
function Ixy!(uxy::AbstractMatrix,u::AbstractMatrix)
    m,n = size(uxy)
    @boundscheck (m+1,n+1) == size(u) || throw(BoundsError())

    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            uxy[i,j] = one_quart*(u[i,j] + u[i+1,j] + u[i,j+1] + u[i+1,j+1])
        end
    end
end


#=
  Euler #15 in Julia.

  Problem 15
  """
  Starting in the top left corner of a 2×2 grid, there are 6 routes
  (without backtracking) to the bottom right corner.

  How many routes are there through a 20×20 grid?
  """

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

=#

include("Euler.jl")

function prodlist(from,to)
    return prod(from:to)
end

# 0.01229874s
function euler15a()
    return prodlist(BigInt(21),40) / prodlist(2,20);
end

run_euler(euler15a)


cd(@__DIR__) # changes the directory to the current directory, the default I guess is the HOME
using Pkg; Pkg.activate("."); Pkg.instantiate()

#= Pkg.add("NeuralPDE")
Pkg.add("Flux")
Pkg.add("ModelingToolkit")
Pkg.add("GalacticOptim")
Pkg.add("Optim")
Pkg.add("DiffEqFlux")
Pkg.add("Plots") =#

using NeuralPDE, Flux, ModelingToolkit, GalacticOptim, Optim, DiffEqFlux
import ModelingToolkit: Interval, infimum, supremum

@parameters t, x
@variables u(..)
Dxx = Differential(x)^2
Dtt = Differential(t)^2
Dt = Differential(t)

# PDE
C=1
eq = Dtt(u(t,x)) ~ C^2*Dxx(u(t,x))

# Initial & boundary conditions
bcs = [
    u(t,0) ~ 0., # for all t > 0
    u(t,1) ~ 0., # for all t > 0
    u(0,x) ~ x*(1. - x), # for all 0 < x < 1
    Dt(u(0,x)) ~ 0. # for all 0 < x < 1
]

# Space-time domains
domains = [
    t ∈ Interval(0.0, 1.0),
    x ∈ Interval(0.0, 1.0)
]

# Discretization
dx = 0.1

println("antes de chain")

# Neural Network
chain = FastChain(FastDense(2,16,Flux.σ),FastDense(16,16,Flux.σ),FastDense(16,1))
initθ = Float64.(DiffEqFlux.initial_params(chain))
discretization = PhysicsInformedNN(chain, GridTraining(dx); init_params = initθ)

@named pde_system = PDESystem(eq,bcs,domains,[t,x],[u(t,x)])
prob = discretize(pde_system,discretization)

cb = function (p,l)
    println("Current loss is: $l")
    return false
end

println("antes de optimizer")

# optimizer
opt = Optim.BFGS()
res = GalacticOptim.solve(prob,opt; cb = cb, maxiters=1200)
phi = discretization.phi

using Plots

ts,xs = [infimum(d.domain):dx:supremum(d.domain) for d in domains]
analytic_sol_func(t,x) =  sum([(8/(k^3*pi^3)) * sin(k*pi*x)*cos(C*k*pi*t) for k in 1:2:50000])

u_predict = reshape([first(phi([t,x],res.minimizer)) for t in ts for x in xs],(length(ts),length(xs)))
u_real = reshape([analytic_sol_func(t,x) for t in ts for x in xs], (length(ts),length(xs)))

diff_u = abs.(u_predict .- u_real)
p1 = plot(ts, xs, u_real, linetype=:contourf,title = "analytic");
p2 =plot(ts, xs, u_predict, linetype=:contourf,title = "predict");
p3 = plot(ts, xs, diff_u,linetype=:contourf,title = "error");
plot(p1,p2,p3)


using Distributions
using Random, Test

function test_location_scale_normal(μ::Float64,σ::Float64,μD::Float64,σD::Float64)
    ρ = Normal(μD,σD)
    d = LocationScale(μ,σ,ρ)
    dref = Normal(μ+σ*μD,σ*σD)

    @test minimum(d) == minimum(dref)
    @test maximum(d) == maximum(dref)
    @test extrema(d) == (minimum(d), maximum(d))

    #### Conversions

    @test convert(LocationScale{Float16}, μ, σ, ρ) == LocationScale{Float16,typeof(ρ)}(Float16(μ), Float16(σ), ρ)
    @test convert(LocationScale{Float16}, d) == LocationScale{Float16,typeof(ρ)}(Float16(μ), Float16(σ), ρ)

    @test location(d) == μ
    @test scale(d) == σ
    @test params(d) == (μ,σ,ρ)
    @test partype(d) == Float64

    #### Statistics

    @test mean(d) ≈ mean(dref)
    @test median(d) ≈ median(dref)
    @test mode(d) ≈ mode(dref)
    @test modes(d) ≈ modes(dref)

    @test var(d) ≈ var(dref)
    @test std(d) ≈ std(dref)
    @test skewness(d) ≈ skewness(dref)
    @test kurtosis(d) ≈ kurtosis(dref)

    @test isplatykurtic(d) == isplatykurtic(dref)
    @test isleptokurtic(d) == isleptokurtic(dref)
    @test ismesokurtic(d) == ismesokurtic(dref)

    @test entropy(d) ≈ entropy(dref)
    @test mgf(d,-0.1) ≈ mgf(dref,-0.1)

    #### Evaluation & Sampling

    insupport(d,0.4) == insupport(dref,0.4)
    @test pdf(d,0.1) ≈ pdf(dref,0.1)
    @test pdf.(d,2:4) ≈ pdf.(dref,2:4)
    @test logpdf(d,0.4) ≈ logpdf(dref,0.4)
    @test loglikelihood(d,[0.1,0.2,0.3]) ≈ loglikelihood(dref,[0.1,0.2,0.3])
    @test cdf(d,μ-0.4) ≈ cdf(dref,μ-0.4)
    @test logcdf(d,μ-0.4) ≈ logcdf(dref,μ-0.4)
    @test ccdf(d,μ-0.4) ≈ ccdf(dref,μ-0.4) atol=1e-100
    @test logccdf(d,μ-0.4) ≈ logccdf(dref,μ-0.4) atol=1e-16
    @test quantile(d,0.1) ≈ quantile(dref,0.1)
    @test quantile(d,0.5) ≈ quantile(dref,0.5)
    @test quantile(d,0.9) ≈ quantile(dref,0.9)
    @test cquantile(d,0.1) ≈ cquantile(dref,0.1)
    @test cquantile(d,0.5) ≈ cquantile(dref,0.5)
    @test cquantile(d,0.9) ≈ cquantile(dref,0.9)
    @test invlogcdf(d,log(0.2)) ≈ invlogcdf(dref,log(0.2))
    @test invlogcdf(d,log(0.5)) ≈ invlogcdf(dref,log(0.5))
    @test invlogcdf(d,log(0.8)) ≈ invlogcdf(dref,log(0.8))
    @test invlogccdf(d,log(0.2)) ≈ invlogccdf(dref,log(0.2))
    @test invlogccdf(d,log(0.5)) ≈ invlogccdf(dref,log(0.5))
    @test invlogccdf(d,log(0.8)) ≈ invlogccdf(dref,log(0.8))

    r = Array{Float64}(undef, 100000)
    rand!(d,r)
    @test mean(r) ≈ mean(dref) atol=0.01
    @test std(r) ≈ std(dref) atol=0.01
    @test cf(d, -0.1) ≈ cf(dref,-0.1)
    @test gradlogpdf(d, 0.1) ≈ gradlogpdf(dref, 0.1)

end

println("\ttesting LocationScale")
test_location_scale_normal(0.3,0.2,0.1,0.2)
test_location_scale_normal(-0.3,0.1,-0.1,0.3)
test_location_scale_normal(1.3,0.4,-0.1,0.5)


"A `HyperRectangle` is the cartesian product of intervals."
abstract type HyperRectangle{T} <: ProductDomain{T} end

boundingbox(d::HyperRectangle) = d

"Compute all corners of the hyperrectangle."
function corners(d::HyperRectangle)
	left = infimum(d)
	right = supremum(d)
    N = length(left)
    corners = [zeros(numtype(d), N) for i in 1:2^N]
    # All possible permutations of the corners
    for i=1:2^length(left)
        for j=1:N
            corners[i][j] = ((i>>(j-1))%2==0) ? left[j] : right[j]
        end
    end
    corners
end

"A `Cube` is a hyperrectangle with equal side lengths in each dimension."
abstract type Cube{T} <: HyperRectangle{T} end

"A cube in a fixed N-dimensional Euclidean space."
const EuclideanCube{N,T} = Cube{SVector{N,T}}

"A cube with vector elements of variable length."
const VectorCube{T} = Cube{Vector{T}}

# This will cause a warning if vectors of different length are used with cubes
iscompatiblepair(x::AbstractVector, d::VectorCube) = length(x) == dimension(d)

"The unit cube is the domain `[0,1]^d`."
abstract type UnitCube{T} <: Cube{T} end

component(d::UnitCube, i::Int) =
    (1 <= i <= dimension(d) || throw(BoundsError); UnitInterval{numtype(d)}())

volume(d::UnitCube) = 1

"A unit cube that is specified by the element type `T`."
struct StaticUnitCube{T} <: UnitCube{T}
end

StaticUnitCube() = StaticUnitCube{SVector{3,Float64}}()
StaticUnitCube(::Val{N}) where {N} = StaticUnitCube{SVector{N,Float64}}()

StaticUnitCube{T}(n::Int) where {T} =
    (@assert n == euclideandimension(T); StaticUnitCube{T}())
StaticUnitCube{T}(::Val{N}) where {N,T} =
    (@assert N == euclideandimension(T); StaticUnitCube{T}())

similardomain(d::StaticUnitCube, ::Type{T}) where {T<:StaticTypes} =
    StaticUnitCube{T}()
similardomain(d::StaticUnitCube, ::Type{T}) where {T} =
    DynamicUnitCube{T}(dimension(d))

components(d::StaticUnitCube{SVector{N,T}}) where {N,T} =
    ntuple(x->UnitInterval{T}(), Val(N))

"The unit cube in a fixed N-dimensional space."
const EuclideanUnitCube{N,T} = StaticUnitCube{SVector{N,T}}

EuclideanUnitCube{N}() where {N} = EuclideanUnitCube{N,Float64}()

const UnitSquare{T} = EuclideanUnitCube{2,T}

"A unit cube whose dimension is specified by a field."
struct DynamicUnitCube{T} <: UnitCube{T}
    dimension   ::  Int

    DynamicUnitCube{T}(n::Int) where {T} = new(n)
    DynamicUnitCube{T}(n::Int) where {T<:StaticTypes} =
        (@assert n == euclideandimension(T); new(n))
end

DynamicUnitCube(n::Int) = DynamicUnitCube{Vector{Float64}}(n)

dimension(d::DynamicUnitCube) = d.dimension

components(d::DynamicUnitCube) = map(x->UnitInterval{numtype(d)}(), 1:dimension(d))

similardomain(d::DynamicUnitCube, ::Type{T}) where {T} =
    DynamicUnitCube{T}(d.dimension)
similardomain(d::DynamicUnitCube, ::Type{T}) where {T <: StaticTypes} =
    StaticUnitCube{T}()

"The unit cube with vector elements of a given dimension."
const VectorUnitCube{T} = DynamicUnitCube{Vector{T}}

VectorUnitCube(n::Int = 3) = VectorUnitCube{Float64}(n)
VectorUnitSquare() = VectorUnitCube(2)

UnitCube(n::Int) = DynamicUnitCube(n)
UnitCube(::Val{N} = Val(3)) where {N} = EuclideanUnitCube{N}()

UnitCube{T}(n::Int) where {T <: StaticTypes} = StaticUnitCube{T}(n)
UnitCube{T}(::Val{N}) where {N,T} = StaticUnitCube{T}(Val(N))
UnitCube{T}() where {T <: StaticTypes} = StaticUnitCube{T}()
UnitCube{T}(n::Int) where {T} = DynamicUnitCube{T}(n)

UnitCube(domains::UnitInterval...) = UnitCube(domains)
UnitCube(domains::NTuple{N,UnitInterval{T}}) where {N,T} =
    UnitCube{SVector{N,T}}(domains)
UnitCube(domains::SVector{N,UnitInterval{T}}) where {N,T} =
    UnitCube{SVector{N,T}}(domains)

UnitCube{T}(domains::UnitInterval...) where {T} = UnitCube{T}(domains)
UnitCube{T}(domain::NTuple{N,<:UnitInterval}) where {N,T<:SVector{N}} =
    StaticUnitCube{T}()
UnitCube{T}(domain::SVector{N,<:UnitInterval}) where {N,T<:SVector{N}} =
    StaticUnitCube{T}()

# Constructor: careful about ambiguities with FixedInterval arguments below
ProductDomain(domains::UnitInterval...) = UnitCube(domains...)
ProductDomain(domains::NTuple{N,<:UnitInterval}) where {N} = UnitCube(domains)
ProductDomain(domains::SVector{N,<:UnitInterval}) where {N} = UnitCube(domains)
ProductDomain(domains::AbstractVector{<:UnitInterval{T}}) where {T} =
    VectorUnitCube{T}(length(domains))
ProductDomain{T}(domains::UnitInterval...) where {N,S,T<:SVector{N,S}} =
	UnitCube{T}(domains...)
ProductDomain{T}(domains::NTuple{N,<:UnitInterval}) where {N,S,T<:SVector{N,S}} =
	UnitCube{T}(domains)
ProductDomain{T}(domains::SVector{N,<:UnitInterval}) where {N,S,T<:SVector{N,S}} =
	UnitCube{T}(domains)
ProductDomain{T}(domains::AbstractVector{<:UnitInterval}) where {S,T<:Vector{S}} =
	VectorUnitCube{S}(length(domains))

## Display:
show(io::IO, d::EuclideanUnitCube{3,Float64}) = print(io, "UnitCube()")
show(io::IO, d::EuclideanUnitCube{N,Float64}) where {N} = print(io, "UnitCube(Val($(N)))")
show(io::IO, d::UnitSquare{Float64}) = print(io, "UnitSquare()")
show(io::IO, d::UnitSquare{T}) where {T} = print(io, "UnitSquare{$(T)}()")
show(io::IO, d::VectorUnitCube{Float64}) = print(io, "UnitCube($(dimension(d)))")

# set the display stencils to [] to opt-out of composite display for the types above
Display.displaystencil(d::EuclideanUnitCube{N,Float64}) where {N} = []
Display.displaystencil(d::UnitSquare) = []
Display.displaystencil(d::VectorUnitCube{Float64}) = []
Display.object_parentheses(d::EuclideanUnitCube{N,Float64}) where {N} = false
Display.object_parentheses(d::UnitSquare) = false
Display.object_parentheses(d::VectorUnitCube{Float64}) = false

"An N-dimensional rectangle is the cartesian product of closed intervals."
struct Rectangle{T} <: HyperRectangle{T}
    a   ::  T
    b   ::  T

    function Rectangle{T}(a::S,b::S) where {S,T}
        @assert length(a)==length(b)
        new(a,b)
    end
end

dimension(d::Rectangle) = length(d.a)
component(d::Rectangle, i::Int) = ClosedInterval(d.a[i],d.b[i])
components(d::Rectangle) = map(ClosedInterval, d.a, d.b)

Rectangle(a, b) = Rectangle(promote(a,b)...)
Rectangle(a::T, b::T) where {T} = Rectangle{T}(a, b)
Rectangle(a::NTuple{N,T}, b::NTuple{N,T}) where {N,T} =
    Rectangle(SVector{N,T}(a), SVector{N,T}(b))

Rectangle(domains::Tuple) = Rectangle(domains...)
Rectangle(domains::ClosedInterval...) = Rectangle(promote_domains(domains)...)
Rectangle(domains::ClosedInterval{T}...) where {T} =
    Rectangle(map(infimum, domains), map(supremum, domains))
Rectangle(domains::AbstractVector{<:ClosedInterval}) =
    Rectangle(map(infimum, domains), map(supremum, domains))
Rectangle(domains::Domain...) =
    error("The Rectangle constructor expects two points or a list of intervals (closed).")

Rectangle{T}(domains::Tuple) where {T} = Rectangle{T}(domains...)
Rectangle{T}(domains::ClosedInterval...) where {T} =
	Rectangle{T}(map(infimum, domains), map(supremum, domains))
Rectangle{T}(domains::AbstractVector{<:ClosedInterval}) where {T} =
    Rectangle{T}(map(infimum, domains), map(supremum, domains))
Rectangle{T}(domains::Domain...) where {T} =
    error("The Rectangle constructor expects two points or a list of intervals (closed).")

ProductDomain(domains::ClosedInterval...) = Rectangle(domains...)
ProductDomain(domains::AbstractVector{<:ClosedInterval}) =
    Rectangle(map(infimum, domains), map(supremum, domains))
ProductDomain{T}(domains::ClosedInterval...) where {T} = Rectangle{T}(domains...)
ProductDomain{T}(domains::AbstractVector{<:ClosedInterval}) where {T} =
    Rectangle{T}(map(infimum, domains), map(supremum, domains))

"The N-fold cartesian product of a fixed interval."
struct FixedIntervalProduct{N,T,D} <: Cube{SVector{N,T}}
end

component(d::FixedIntervalProduct{N,T,D}, i::Int) where {N,T,D} =
    (1 <= i <= dimension(d) || throw(BoundsError); D())
components(d::FixedIntervalProduct{N,T,D}) where {N,T,D} =
    ntuple(x->D(), Val(N))

volume(d::FixedIntervalProduct{N,T,D}) where {N,T,D} = volume(D())^N

FixedIntervalProduct(domains::NTuple{N,D}) where {N,D <: FixedInterval} =
	FixedIntervalProduct{N,eltype(D),D}()
FixedIntervalProduct(domains::SVector{N,D}) where {N,D <: FixedInterval} =
	FixedIntervalProduct{N,eltype(D),D}()

const ChebyshevProductDomain{N,T} = FixedIntervalProduct{N,T,ChebyshevInterval{T}}
ChebyshevProductDomain(::Val{N}) where {N} = ChebyshevProductDomain{N}()
ChebyshevProductDomain{N}() where {N} = ChebyshevProductDomain{N,Float64}()

ProductDomain(domains::D...) where {D <: FixedInterval} =
	FixedIntervalProduct(domains)
ProductDomain(domains::NTuple{N,D}) where {N,D <: FixedInterval} =
	FixedIntervalProduct(domains)
ProductDomain(domains::SVector{N,<:FixedInterval}) where {N} =
	FixedIntervalProduct(domains)
ProductDomain{T}(domains::D...) where {N,S,T<:SVector{N,S},D <: FixedInterval} =
	FixedIntervalProduct(convert.(Ref(Domain{S}), domains))
ProductDomain{T}(domains::NTuple{N,D}) where {N,S,T<:SVector{N,S},D <: FixedInterval} =
	FixedIntervalProduct(convert.(Ref(Domain{S}), domains))
ProductDomain{T}(domains::SVector{N,<:FixedInterval}) where {N,S,T<:SVector{N}} =
	FixedIntervalProduct(domains)


### A Pluto.jl notebook ###
# v0.17.1

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local iv = try Base.loaded_modules[Base.PkgId(Base.UUID("6e696c72-6542-2067-7265-42206c756150"), "AbstractPlutoDingetjes")].Bonds.initial_value catch; b -> missing; end
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : iv(el)
        el
    end
end

# ╔═╡ a6261f66-40ca-11ec-3021-ed2173550635
begin
	using Pkg
	Pkg.activate("../../")
	using Revise
end

# ╔═╡ 77e6389e-3491-45b6-9f76-459b15a8922d
using TestImages, ImgProcMic, ImageShow, Random, PlutoTest, Plots, FourierTools, Statistics

# ╔═╡ ee4b6b2d-119e-4ac3-a1f5-500ff045bfc5
using FFTW, PlutoUI

# ╔═╡ c7150c2c-afeb-4158-b381-3341bb678fd5
using Noise

# ╔═╡ 7148e705-8b90-46ea-8763-778897daf9a7
md"# Homework 03

In this homework we cover several topics like image correction, Fourier transforms, Fourier Shifting and Sampling.
"

# ╔═╡ 83d9c252-2a8d-47a3-b247-a62b6a175201
md"## 1. Flat Fielding
As discussed in the lecture, raw images might be affected by dirty optics or sensor errors.
In this exercise we want to apply such a correction to simulated images.
"

# ╔═╡ 39c623c7-6063-4360-a2b7-95ea42c5262d
begin
	img = Float32.(testimage("mandril_gray.tif"))
	gray_show(img)
end

# ╔═╡ eacd8fdf-971c-46ca-914f-f73f65e28ca8
md"### Simulate Deteriorated Image
In the first step, we degrade our image with some artifacts.
We add an offset to our image which is caused by sensor artifacts.
Second, we add some dirty spots to the image which are for example caused by dust on the sensor.
"

# ╔═╡ 25dde5cb-9ed6-498e-b511-50ed0640d0c0
"""
	simulate_bad_microscope(img)

The ideal `img` is degraded by artifacts.
"""
function simulate_bad_microscope(img::AbstractArray{T}) where T
	# generate same random numbers!
	rng = MersenneTwister(2021);
	# some weird dark field slope offset
	dark = zero.(img)
	dark .+= range(0,0.3, length=size(img, 1))

	xs = 1:size(img, 1)

	ff = one.(img)
	for i = 1:20 # emulates some dust spots as modulations to the flatfield ff
		x, y = rand(rng, xs), rand(rng, xs)
		sca = abs(randn(rng) .* T(0.2))
		r = sqrt.((xs .- x).^2 .+ (xs' .- y).^2)
		ff .-= T(0.2) .* sinc.(r .* sca)
	end

	return abs.(dark .+ ff .* img)
end

# ╔═╡ a6c024a4-7cb2-4f8d-be24-8727d75679da
img_dirty = simulate_bad_microscope(img);

# ╔═╡ d5c127ee-1ecf-4ed0-9cd9-6fcdefa95847
gray_show(img_dirty)

# ╔═╡ 6e95cbcb-0d89-4eca-8599-b266076708a7
md"## 1.1 Task
Extract both a flat field and a dark field image with `simulate_bad_microscope`.
Use the function `simulate_bad_microscope` to obtain a measured flat field and measured dark field.
"

# ╔═╡ 910cb7fa-b1fc-4106-acc4-3bf56edfb83d
flat_field = simulate_bad_microscope(ones(eltype(img), size(img)));

# ╔═╡ cbebb311-c521-4fa6-a909-1a4df8da4d2d
gray_show(flat_field)

# ╔═╡ 8bb737cb-2d59-49ba-9b89-afef6ca7bbf8
dark_field = simulate_bad_microscope(zeros(eltype(img), size(img)));

# ╔═╡ 52bc8b3d-542e-4a35-b38c-35840f4075d4
gray_show(dark_field)

# ╔═╡ e3448ffe-5fbf-4bed-bf03-6eb9031e6c12
"""
	correct_image(img_dirty, dark_field, flat_field)

Correct a dirty image `img_dirty` by providing the `dark_field` and `flat_field`.

Conceptually we divide by the `flat_field` since the measured `flat_field`
describes the deviations from an ideal white image.
The `dark_field` has to be subtracted before since that only contains sensor artifacts.
"""
function correct_image(img_dirty, dark_field, flat_field)
	return (img_dirty .- dark_field) ./ (flat_field .- dark_field)
end

# ╔═╡ d9308666-efaa-4fcc-a083-f3d12dbab1e2
img_fixed = correct_image(img_dirty, dark_field, flat_field);

# ╔═╡ 493b6bb4-1dfc-4a82-9360-40944b29826c
gray_show(img_fixed)

# ╔═╡ c80a2ce8-6620-4de1-bb78-5af97caf466d
md"### 1.1 Test"

# ╔═╡ e2e470e4-ff73-4dd5-973c-14e1535b01ea
PlutoTest.@test img_fixed ≈ img

# ╔═╡ 0a957f4a-a975-4087-b458-3236940eda43
PlutoTest.@test img ≈ correct_image(img, zero.(img), one.(img))

# ╔═╡ af9866f4-2b03-4dc4-9f8a-e02ae46e2f9c
md"# 2 Sampling
If a signal is not properly sampled, we can see a lot of artifacts which are called aliasing.
Try to determine experimentally (via the slider) which is the correct value for Nyquist sampling.

Afterwards, calculate the correct sampling in the variable $N$!

Don't state the pure number but rather a short equation which Julia calculates (something like `1 * 2 * 3`).
Do you see a difference between your expected and your calculated value?
"

# ╔═╡ 7a481af3-4ec8-4a58-bb84-e9d3d5381570
md"N=$(@bind N Slider(2:300, show_value=true))"

# ╔═╡ 4006ccf7-df8e-440c-8b3c-af525bb50ed4
begin
	xs = range(0, 8, length=N+1)[begin:end-1]
	f(x) =  sin(x * 2π * 5)
end

# ╔═╡ c871cdbd-b560-46f2-918c-45d24d2f9b75
plot(xs, f.(xs))

# ╔═╡ f6b8cd92-c1ab-445b-8b4f-04f521b72cad
#           for Nyquist    because of the        because of the interval length
N_correct = 1 +    2 *         5                    *   8

# ╔═╡ 51b2d7be-fe94-44e9-8e67-73aafb636ef1
begin
	xs_correct = range(0, 8, length=round(Int, N_correct + 1))[begin:end-1]
	plot(xs_correct, f.(xs_correct))
end

# ╔═╡ a19c8850-9a7d-4e26-908b-a067a9006e84
md"No tests, since they would reveal the answer."

# ╔═╡ b606e4d0-8cac-4028-8e38-113050531720
plot(resample(f.(xs_correct), (400 ,)))

# ╔═╡ b0c67bd6-b5ef-4359-8727-1a3bfefcf759
md"# 3 FFTs

Another very important topics are FFTs!
The standard library in many languages is based on FFTW.
In Julia FFTW.jl is the recommended library.

The discrete Fourier transform is given by

$$X_k = \sum_{n=0}^{N-1} x_n \exp\left(-i2 \pi \frac{kn}{N} \right)$$
and the inverse by

$$x_n = \frac1{N}\sum_{k=0}^{N-1} X_k \exp\left(i2 \pi \frac{kn}{N} \right)$$.

The FFT algorithm evaluates those equations very efficiently ($\mathcal O(N \cdot \log N)$) by avoiding the explicit sum ($\mathcal O(N^2)$).

Further, the FFT calculates the frequencies in an (maybe) unexpected format
* Data: $[x_1, x_2, ..., x_N]$
* FFT of Data: $[f_0, f_1, ..., f_{N/2}, f_{-N/2}, f_{-N/2+1}, ...,f_{-1}]$
* (for real data negative and positive frequencies are the complex conjugate)

Sometimes we want to have the frequencies centered, which means a `fftshift` operations is needed
* FFT of Data followed by fftshift: $[f_{-N/2}, f_{-N/2+1},..., f_{-1}, f_0, f_1,..., f_{N/2-1}, f_{N/2}]$
"

# ╔═╡ 7276f9e0-80b8-4794-9125-fca43246f818
"""
    my_dft(x::Vector{T})::Vector{Complex{T}}

Calculate the DFT of `data` according to the equation given above.
Note, the output should be a complex Vector!
"""
function my_dft(x::Vector{T})::Vector{Complex{T}} where T
	X = zeros(Complex{T}, size(x))

	# length of x
	N = size(x, 1)

	# loop over output elements
	for k ∈ 0:N-1
		# loop over inputs
		for n ∈ 0:N-1
			X[k + 1] += x[n + 1] * exp(-im * T(2π) * k * n / N)
		end
	end

	return X
end

# ╔═╡ f054e266-e16c-4311-9e5e-1dbb9d024a63
typeof(π)

# ╔═╡ e895cff6-4018-495f-9a0e-1e56ddcfd17d
arr = [1f0,2,3,4]

# ╔═╡ 14d68904-4de5-4f1d-81f6-dae8bf50b749
my_dft(arr)

# ╔═╡ 7307dd33-cb79-48c0-9e73-78ed70e9c628
fft(arr)

# ╔═╡ 52061097-5be4-4b1a-a4f6-fd030a5a52ba
arr123 = randn((1_000))

# ╔═╡ 66cdd72e-ec6f-49ec-8338-c32e744cc359
@time fft(arr123)

# ╔═╡ 3c77082d-20e9-4541-8426-f8457cc728ec
@time my_dft(arr123)

# ╔═╡ 31b60200-4d70-4a3c-995f-949b19800549
# From worker 4:       0.000315 seconds (32 allocations: 33.844 KiB)
# From worker 4:       0.052112 seconds (1 allocation: 15.750 KiB)

# ╔═╡ 393d03e2-461d-480a-9719-f33a4eef3a6f
md"### 3.1 Test"

# ╔═╡ df12fd86-f29d-4926-b360-daf4cb40af86
begin
	arr_r = randn((13,))
	PlutoTest.@test fft(arr_r) ≈ my_dft(arr_r)
end

# ╔═╡ 37eedc2e-27bd-4d1c-a83b-6ec14aba55ab
begin
	arr_r2 = randn((14,))
	PlutoTest.@test fft(arr_r2) ≈ my_dft(arr_r2)
end

# ╔═╡ 5313dbed-bfd7-4bb8-b39c-89c55e50d1bb
md"## 3.2 Task
Calculate the FFT of the following dataset, such that center frequency is in the center!

Use `fft` and `fftshift`.
"

# ╔═╡ e1eac3cd-e6ce-4743-9119-5ee0c69b8e00
begin
	data_odd = [1.7142857142857142, 1.643544850921574, 1.0471849694228637, 1.383330948397344, 1.8300917955219322, 1.2951185346328968, 1.086443186817675]
	data_even = [-2.5, 1.5, -2.5, 3.5]
end

# ╔═╡ 481961ac-ad7c-4768-a7aa-eed34143226b
ffts(x) = fftshift(fft(x))

# ╔═╡ 25a2de2a-612b-4ed3-8aad-d5caf5d59c30
[1,0,0,0,0]

# ╔═╡ 2a74167b-7056-47e1-861b-bf7c1810a1f0
fftshift([1,0,0,0,0])

# ╔═╡ c95f9463-de4d-47d4-abfd-f707641d2021
[1,0,0,0]

# ╔═╡ 6fc3a300-7b96-4b21-8487-fdf4b5eb7521
fftshift([1,0,0,0])

# ╔═╡ c24b2457-4396-48c6-a21f-6d982d80776e
length([1,2,3,4,5]) ÷ 2 + 1

# ╔═╡ 27aa0d27-1cb1-4827-80c7-48741cc0312e
length([1,2,3,4]) ÷ 2 + 1  # \div or div(N, 2)

# ╔═╡ fc3bb52b-f9a2-42c0-b34b-347d0a137e2d
ffts(data_even)

# ╔═╡ a55851b6-3f8d-4af3-8f09-a3777c56df1b
ffts(data_odd)

# ╔═╡ e46d67be-d974-4f07-8b3e-02fff3e749d0
md"## 3.3 Task

Now do the inverse operation such that `data ≈ iffts(ffts(x))`

Check out `ifftshift`.
"

# ╔═╡ deb9eaf0-ee89-49c5-90b8-570d1a88f8fc
iffts(x) = ifft(ifftshift(x))

# ╔═╡ 90dd14ae-fb42-400e-86b4-9ae9e63716fc

# ╔═╡ 34e1c198-abe0-470e-a4bc-348ec09a2032
ifftshift([0,0,1,0])

# ╔═╡ ffb5665d-ecac-4474-b2a4-9685cc40ad7e
ifftshift([0,0,1,0,0])

# ╔═╡ 6ca02109-6114-411f-bfb3-bd638ab93218

# ╔═╡ 2875999a-a61d-491a-8a42-0552a87d3c6a
md"""
### 3.2 and 3.3 Test
"""

# ╔═╡ 333a253e-ce8b-477a-9db2-f99c034ae048
PlutoTest.@test real(iffts(ffts(data_even))) ≈ data_even

# ╔═╡ e81990c8-2194-472c-a102-b0d010f3a266
 PlutoTest.@test real(ffts(iffts(data_odd))) ≈ data_odd

# ╔═╡ f3fc2ea5-033c-4789-8ce6-0c6b15b607b7
 PlutoTest.@test real(iffts(ffts(data_odd))) ≈ data_odd

# ╔═╡ 90596b8a-53d4-457f-9e93-522e6c5dc35d
 PlutoTest.@test real(iffts(ffts(data_even))) ≈ data_even

# ╔═╡ 4e54ddc6-b9cd-472f-b496-18055ee5b667
md"## 3.4 Task
Try to reproduce this $\cos$ curve.
The reason it looks so _sharp_ is that we only sample slightly above the Nyquist limit.

However, you are only allowed to use `fft`, `fftshift`, `ifftshift`, `ifft` and array manipulation (like `x[34] = 1723.12`).
And you can also take the `real` part at the end.

Please explain with some #comments what you did.
"

# ╔═╡ 9844ec03-b169-40a0-ad56-1c07b11bb886
begin
	x = range(0, 4, length=21)[begin:end-1]
	y = cos.(x .* 2π)
	plot(x, y, mark="-*")
end

# ╔═╡ 49c39cea-895e-429a-909c-d1438d0dfd00
function reproduce_cos(y)
	out = zeros(size(y))
	# we go to Fourier space with center frequncy in the center
	out_ft = ffts(out)

	out_ft[15] = 10
	out_ft[7] = 10

	# we go real space again
	out = real(iffts(out_ft))
	return out
end

# ╔═╡ fcdb96af-30bd-42c2-bf45-ecd13f27b57d
collect(fftshift(FFTW.fftfreq(length(y))))

# ╔═╡ 5c4f40df-cde3-4140-ace2-6c0e1b60ebcf
begin
	plot(x, y, mark="-*")
	plot!(x, reproduce_cos(y) .+  0.1, mark="-*")
end

# ╔═╡ 79578e4d-7527-459b-8915-ee5cb9ccad24
md"### 3.4 Test"

# ╔═╡ a74c4034-fa19-4360-9f5e-191d5ba4001a
PlutoTest.@test real(reproduce_cos(copy(y))) ≈ y

# ╔═╡ 8ce249a3-0881-49e5-8ffd-5df21bda50f4
md"## 3.5 Fourier Shift Theorem

The discrete Fourier transform is given by

$$X_k = \sum_{n=0}^{N-1} x_n \exp\left(-i2 \pi \frac{kn}{N} \right)$$
and the inverse by

$$x_n = \frac1{N}\sum_{k=0}^{N-1} X_k \exp\left(i2 \pi \frac{kn}{N} \right)$$.
"

# ╔═╡ 13131381-11b3-4db4-9351-fb7264c94ea4
md"## 3.4 Task
Proof the Fourier shift theorem (with LaTeX) which states

$x_{n+\Delta n} = \mathcal{F}^{-1}\left[ X_{k} \cdot f(k, \Delta n)\right]$

Derive how the function $f(k, \Delta n)$ looks like!

hint: by clicking on the eye symbol next to this task, you can see how to embed LaTeX code in markdown comments.
"

# ╔═╡ bc3ae07f-3877-480f-bd0b-633ce21172e3
md"### Derivation
$x_{n+\Delta n} =  \frac1{N}\sum_{k=0}^{N-1} X_k \exp\left(i2 \pi \frac{k(n+ \Delta n)}{N} \right)$

$=\frac1{N}\sum_{k=0}^{N-1} X_k \exp\left(i2 \pi \frac{kn}{N} \right)\exp\left(i2 \pi \frac{k\Delta n}{N} \right)$

$=\frac1{N}\sum_{k=0}^{N-1} \left( X_k \exp\left(i2 \pi \frac{k\Delta n}{N} \right) \right) \exp\left(i2 \pi \frac{kn}{N} \right)$

$=\mathcal{F}^{-1} \left[ X_k \exp\left(i2 \pi \frac{k\Delta n}{N} \right) \right]$

$f(k, \Delta n) = \exp\left(i2 \pi \frac{k\Delta n}{N} \right)$

"

# ╔═╡ 8bf6ba5b-1d06-41de-b22a-b58c1eaf2c83
"""
	shift(x::AbstractArray{T, N}, Δn::Number)

Shifts a (real) array via the Fourier shift theorem by the amount of
`Δn`.

If `Δn ∈ ℕ` then `circshift(x, -Δn) ≈ shift(x, Δn).`
Otherwise a sub-pixel shift is obtained which is equivalent
to a sinc interpolation.

## Examples
```julia
julia> shift([1.0, 2.0, 3.0, 4.0], 1)
4-element Vector{Float64}:
 2.0
 3.0
 4.0
 1.0

julia> shift([1.0, 2.0, 3.0, 4.0], -1)
4-element Vector{Float64}:
 4.0
 1.0
 2.0
 3.0

julia> shift([1.0, 2.0, 3.0, 4.0], -0.5)
4-element Vector{Float64}:
 2.5
 1.085786437626905
 2.5
 3.914213562373095
```
"""
function shift(x::AbstractArray{T, N3}, Δn::Number) where {T, N3}
	x_ft = ffts(x)

	N = size(x, 1)
	k = (-N÷2) : 1 : (- N÷2 + N - 1)

	x_ft.*= exp.(1im .* 2π .* k .* Δn ./ N)

	return real(iffts(x_ft))

	# alternative solution
	#x_ft = fft(x)
	#f = fftfreq(size(x, 1))
	#phase = exp.(1im * T(2π) .* f .* Δn)
	#return real(ifft(x_ft .* phase))
end

# ╔═╡ 64507581-5751-496e-87b6-88eb320ae44c
N2 = 5

# ╔═╡ 47c0516a-4a37-4823-a447-1207d5d4b8b4
k = (-N2÷2) : 1 : (- N2÷2 + N2 - 1)

# ╔═╡ 250559fb-1098-47ab-a55c-55a89497873d
length(k)

# ╔═╡ a21d8b13-20db-4df1-a2c7-451f8942466b
begin
	x2 = 0:0.1:3
	y2 = sin.(x2.*3) .+ cos.(x2 .* 4).^2
end;

# ╔═╡ 420d1b15-8dc9-4c9a-b91f-9cc4b8101b92
md"
 $\Delta n$=$(@bind Δn Slider(0:0.01:100, show_value=true))"

# ╔═╡ 54ebf4c4-6692-4a40-a83e-c548440b3b85
begin
	plot(shift(y2, Δn), label="FFT based shift", mark="-*")
	plot!(circshift(y2, round(Int, .-Δn)), label="rounded to integer with circshift", mark="-*")
end

# ╔═╡ bbe52a98-9811-4c3a-9b4e-b90a1f5c163f
md"### 3.5 Test"

# ╔═╡ 73fd7360-ffdf-4870-a301-e39a982e0517
arr4 = randn((37,));

# ╔═╡ 37a815a0-c6cd-4ee8-9f46-ee4aa5259690
PlutoTest.@test circshift(arr4, -12) ≈ shift(arr4, 12)

# ╔═╡ 3c09615a-e9c4-462b-9f23-5d5676aa29b7
PlutoTest.@test FourierTools.shift(arr4, -13.1) ≈ shift(arr4, 13.1)

# ╔═╡ d3f2648a-a736-4812-b721-70298151a1a7
md"# 4 Noise Beyond Bandlimit

As seen in the lecture, noise can move information even beyond
the bandlimit of the optical system.
In this task, we want to demonstrate this behaviour with a simulation.
For that we need a few helper functions, try to fill the missing lines.

## 4.1 rr2 Function
In the first step, we create a rr2 function.
"

# ╔═╡ 9ef7592b-92f9-48bb-aa3e-9710bf251379
begin
	"""
		rr2(T, s, center=s .÷ 2 + 1)

	Returns a 2D array which stores the squared distance to the center.

	## Examples
	```julia
	julia> rr2(Float32, (5,5))
	5×5 Matrix{Float32}:
	 8.0  5.0  4.0  5.0  8.0
	 5.0  2.0  1.0  2.0  5.0
	 4.0  1.0  0.0  1.0  4.0
	 5.0  2.0  1.0  2.0  5.0
	 8.0  5.0  4.0  5.0  8.0

	julia> rr2((4,4))
	4×4 Matrix{Float64}:
	 8.0  5.0  4.0  5.0
	 5.0  2.0  1.0  2.0
	 4.0  1.0  0.0  1.0
	 5.0  2.0  1.0  2.0

	julia> rr2((3,3), (3,3))
	3×3 Matrix{Float64}:
	 8.0  5.0  4.0
	 5.0  2.0  1.0
	 4.0  1.0  0.0
	```
	"""
	function rr2(T::DataType, s, center=s .÷ 2 .+ 1)
		is = T.((1:s[1]) .- center[1]).^2
		js = T.(((1:s[2]).- center[2]).^2)'
		return is .+ js
	end

	function rr2(s, center=s .÷ 2 .+ 1)
		rr2(Float64, s, center)
	end
end

# ╔═╡ 9e63206c-82a8-4744-a5a4-a42281077c4f
(1:5).^2 .+ ((1:5).^2)'

# ╔═╡ 6899ad5c-3d31-40d7-94f4-042286b6c1c4
((1:5).^2)'

# ╔═╡ b62519f5-26dd-49ca-afa4-2c72de7ef8a6
rr2((4,4))

# ╔═╡ fc25378b-11b1-4383-b144-75c7dc6e9a1f
rr2((5,5))

# ╔═╡ bd45654d-15e0-4717-962b-d5a302de4a9b
PlutoTest.@test rr2(Float32, (4, 3), (2.5, 3)) == Float32[6.25 3.25 2.25; 4.25 1.25 0.25; 4.25 1.25 0.25; 6.25 3.25 2.25]

# ╔═╡ c6ed3889-0d44-4180-9f81-3095055c6e54
PlutoTest.@test rr2((3, 3), (3, 3)) == [8.0 5.0 4.0; 5.0 2.0 1.0; 4.0 1.0 0.0]

# ╔═╡ 5a16e94e-b01b-4897-9631-41d345229be2
md"## 4.2 circ Function
Having a function which creates the radial distance, we now create a function which returns an array which is 1 inside an circle of radius `r` and 0 outside of it.
"

# ╔═╡ 06a666f9-c99e-41e5-b703-a5e826721d04
begin
	"""
		circ(size, radius)

	`size` the size of the resulting array
	and `radius` the radius of the circle
	"""
	circ(size, radius) = rr2(size) .< radius^2
end

# ╔═╡ 255a1f5a-fe84-4ae2-85c3-a5a512a040ec
md"
if $|a| < |b|$ then also $a^2 < b^2$
for $a,b \in \mathrm{R}^+$
"

# ╔═╡ 9f24aec3-b93c-4e13-a227-0e4deefda611

# ╔═╡ 592e0ed4-bd97-4bb6-a826-b1e2f5ddbd59
circ((5,5), 2)

# ╔═╡ 205a744b-e8a1-4acc-ba6b-7a7307f9a028
md"## 4.2 Test"

# ╔═╡ daf4e474-0db4-4865-859d-96c31e7947f0
PlutoTest.@test circ((10, 2), 2.76) ≈ Bool[0 0; 0 0; 0 0; 1 1; 1 1; 1 1; 1 1; 1 1; 0 0; 0 0]

# ╔═╡ 7fb8fedf-5684-4874-8f17-d01debaff477
PlutoTest.@test circ((4, 5), 2.1) ≈ Bool[0 0 1 0 0; 0 1 1 1 0; 1 1 1 1 1; 0 1 1 1 0]

# ╔═╡ 9b9fb85b-13d0-43a2-ac5f-713593117b96
md" ## 4.3 Frequency Filter
Now we combine all methods together such that we frequency filter
our image.

Procedure:
* Transform image to Fourier space
* take only the inner part of a circle a `radius` of the Fourier space
* go to real space and take real part
"

# ╔═╡ 2a48072b-fc5a-4044-9aa9-5661483ffbfc
function frequency_filter(img, radius)
	img_ft = ffts(img)

	img_ft .*= circ(size(img), radius)

	return real(iffts(img_ft))
end

# ╔═╡ 643bf329-cce9-4c61-931d-65455b59357d
md"r=$(@bind r Slider(1:512, show_value=true))"

# ╔═╡ 2cafea2d-b4b6-4ede-b899-720a71034fda
gray_show(frequency_filter(img, r))

# ╔═╡ 7467608a-ce6d-4f29-8cd8-d766395f7129
md"## 4.4 Noise Breaks Resolution Limit?
We demonstrate now, that noise beyond the frequency limit still contains information

### Procedure
* First frequency filter our image such that it definetely bandlimited
* Apply Poissonian and Gaussian noise
* Now do the exact opposite of the above low-pass filter and take only the high frequency content!
"

# ╔═╡ 6d7fd369-1f6c-487d-b399-c8cbbfede916
function simulate(img, noise_function, radius)
	img_bandlimited = frequency_filter(img, radius)
	img_noisy = noise_function(img_bandlimited)

	img_noisy_ft = ffts(img_noisy)

	# a high pass filter now!
	img_noisy_ft .*= (1 .- circ(size(img), radius))

	return real(iffts(img_noisy_ft))
end

# ╔═╡ e8bd9cb9-bd20-4619-93da-20e7829a61b4
# anonymous noise function we are going to pass to `simulate`
begin
	p(x) = poisson(x, 100000)
	g(x) = add_gauss(x, 0.1)
end

# ╔═╡ 8a479712-7290-4c00-af7c-44e11ec81d0d
md"r=$(@bind r2 Slider(1:300, show_value=true))"

# ╔═╡ 372e190b-0344-4001-9c72-171610b41e9c
img_noisy = abs.(simulate(img, p, r2));

# ╔═╡ b17cd903-bce6-4291-9d71-ea8a8422295a
gray_show(img_noisy, set_one=true)

# ╔═╡ fc91074a-99fc-4a13-802c-740872bf7502
md"
Mean is $(round(mean(img_noisy), sigdigits=3))

Maximum is $(round(maximum(img_noisy), sigdigits=3))

Minimum is $(round(minimum(img_noisy), sigdigits=3))
"

# ╔═╡ fac6c740-93c1-4e12-9610-f3c621f84043
md"## 5.1 Undersampling
Demonstration of aliasing.

Write a function `undersample` which only takes every `n-th` point of an array, keeping the first pixel of the array.
"

# ╔═╡ 15aa846c-7b81-4094-8b5b-dd6575e82afe
img3 = Float32.(testimage("resolution_test_512"));

# ╔═╡ be3b0811-c05a-4e78-9d34-8c52195dd7d0
"""
	undersample(x, factor)

```julia-repl
julia> undersample([1,2,3,4,5,6,7,8,9,10], 2)
5-element Vector{Int64}:
 1
 3
 5
 7
 9

julia> undersample([1,2,3,4,5,6,7,8,9,10], 3)
4-element Vector{Int64}:
  1
  4
  7
 10

julia> undersample([1,2,3,4,5,6,7,8,9,10], 4)
3-element Vector{Int64}:
 1
 5
 9
```
"""
function undersample(x, factor)
	# more Julia idiomatic
	return x[begin:factor:end]
	#return x[1:factor:size(x, 1)]
end

# ╔═╡ 282acd5a-12e8-43f3-8de9-1b7b913a476c
collect(1:5:100)

# ╔═╡ e5188537-b955-4637-bab5-a76b6d74c189
PlutoTest.@test undersample([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 2) == [1, 3, 5, 7, 9]

# ╔═╡ c2b3cdc7-60b5-4287-8268-f95184d452b7
PlutoTest.@test undersample([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 3) == [1, 4, 7, 10]

# ╔═╡ 4e3a272e-0280-4e27-b014-3f9da5ac8039
PlutoTest.@test undersample([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 4) == [1, 5, 9]

# ╔═╡ a76e4081-7c52-4548-b513-ba03954378db

# ╔═╡ f5fbd0da-fe4c-4162-be5e-f9d9d5b27d87
md"## 5.1 Undersampling Visualization

In this plot you can see a nicely oversampled $\sin$ curve.
"

# ╔═╡ a07cb23e-6fca-4c42-a6b6-a525b2b5e232
xs2 = range(0, 10π, length=500)

# ╔═╡ c373046e-04b3-4159-80c0-8fbec141833f
f2(x) = sin(5 * x)

# ╔═╡ 9c1ed9aa-9d52-425d-892d-af3b8a981663
ys2 = f2.(xs2);

# ╔═╡ 75006409-0338-470d-a41d-fa1aba5c4fca
plot(xs2, ys2, xlabel="x pos", ylabel="amplitude")

# ╔═╡ f2b3d8a0-4bea-4443-9bc3-37258106c3d3
md"
Now we only take each `undersample_factor` data point of the series.

`undersample_factor` = $(@bind undersample_factor Slider(1:40, show_value=true))

Analyze what happens once you're below the Nyquist frequency.
See further what happens if you undersample then even more.
"

# ╔═╡ f6b06d6e-b268-45d6-a7e4-d8614cd0884d
y2_undersampled = undersample(ys2, undersample_factor);

# ╔═╡ eab1fe28-f336-4b33-89a3-4c65ed2a818f
plot(fftfreq(length(y2_undersampled), 1 / 10π * 500 / undersample_factor)[1:length(y2_undersampled)÷2],abs.(rfft(y2_undersampled))[1:end-1], xlabel="Positive Frequencies", ylabel="Absolute Value of Frequency amplitude", mark="-*")

# ╔═╡ Cell order:
# ╠═a6261f66-40ca-11ec-3021-ed2173550635
# ╠═77e6389e-3491-45b6-9f76-459b15a8922d
# ╟─7148e705-8b90-46ea-8763-778897daf9a7
# ╟─83d9c252-2a8d-47a3-b247-a62b6a175201
# ╠═39c623c7-6063-4360-a2b7-95ea42c5262d
# ╟─eacd8fdf-971c-46ca-914f-f73f65e28ca8
# ╠═25dde5cb-9ed6-498e-b511-50ed0640d0c0
# ╠═a6c024a4-7cb2-4f8d-be24-8727d75679da
# ╠═d5c127ee-1ecf-4ed0-9cd9-6fcdefa95847
# ╟─6e95cbcb-0d89-4eca-8599-b266076708a7
# ╠═910cb7fa-b1fc-4106-acc4-3bf56edfb83d
# ╠═cbebb311-c521-4fa6-a909-1a4df8da4d2d
# ╠═8bb737cb-2d59-49ba-9b89-afef6ca7bbf8
# ╠═52bc8b3d-542e-4a35-b38c-35840f4075d4
# ╠═e3448ffe-5fbf-4bed-bf03-6eb9031e6c12
# ╠═d9308666-efaa-4fcc-a083-f3d12dbab1e2
# ╠═493b6bb4-1dfc-4a82-9360-40944b29826c
# ╟─c80a2ce8-6620-4de1-bb78-5af97caf466d
# ╠═e2e470e4-ff73-4dd5-973c-14e1535b01ea
# ╠═0a957f4a-a975-4087-b458-3236940eda43
# ╟─af9866f4-2b03-4dc4-9f8a-e02ae46e2f9c
# ╠═ee4b6b2d-119e-4ac3-a1f5-500ff045bfc5
# ╠═7a481af3-4ec8-4a58-bb84-e9d3d5381570
# ╠═4006ccf7-df8e-440c-8b3c-af525bb50ed4
# ╠═c871cdbd-b560-46f2-918c-45d24d2f9b75
# ╠═f6b8cd92-c1ab-445b-8b4f-04f521b72cad
# ╠═51b2d7be-fe94-44e9-8e67-73aafb636ef1
# ╟─a19c8850-9a7d-4e26-908b-a067a9006e84
# ╠═b606e4d0-8cac-4028-8e38-113050531720
# ╟─b0c67bd6-b5ef-4359-8727-1a3bfefcf759
# ╠═7276f9e0-80b8-4794-9125-fca43246f818
# ╠═f054e266-e16c-4311-9e5e-1dbb9d024a63
# ╠═e895cff6-4018-495f-9a0e-1e56ddcfd17d
# ╠═14d68904-4de5-4f1d-81f6-dae8bf50b749
# ╠═7307dd33-cb79-48c0-9e73-78ed70e9c628
# ╠═52061097-5be4-4b1a-a4f6-fd030a5a52ba
# ╠═66cdd72e-ec6f-49ec-8338-c32e744cc359
# ╠═3c77082d-20e9-4541-8426-f8457cc728ec
# ╠═31b60200-4d70-4a3c-995f-949b19800549
# ╟─393d03e2-461d-480a-9719-f33a4eef3a6f
# ╠═df12fd86-f29d-4926-b360-daf4cb40af86
# ╠═37eedc2e-27bd-4d1c-a83b-6ec14aba55ab
# ╟─5313dbed-bfd7-4bb8-b39c-89c55e50d1bb
# ╠═e1eac3cd-e6ce-4743-9119-5ee0c69b8e00
# ╠═481961ac-ad7c-4768-a7aa-eed34143226b
# ╠═25a2de2a-612b-4ed3-8aad-d5caf5d59c30
# ╠═2a74167b-7056-47e1-861b-bf7c1810a1f0
# ╠═c95f9463-de4d-47d4-abfd-f707641d2021
# ╠═6fc3a300-7b96-4b21-8487-fdf4b5eb7521
# ╠═c24b2457-4396-48c6-a21f-6d982d80776e
# ╠═27aa0d27-1cb1-4827-80c7-48741cc0312e
# ╠═fc3bb52b-f9a2-42c0-b34b-347d0a137e2d
# ╠═a55851b6-3f8d-4af3-8f09-a3777c56df1b
# ╟─e46d67be-d974-4f07-8b3e-02fff3e749d0
# ╠═deb9eaf0-ee89-49c5-90b8-570d1a88f8fc
# ╠═90dd14ae-fb42-400e-86b4-9ae9e63716fc
# ╠═34e1c198-abe0-470e-a4bc-348ec09a2032
# ╠═ffb5665d-ecac-4474-b2a4-9685cc40ad7e
# ╠═6ca02109-6114-411f-bfb3-bd638ab93218
# ╟─2875999a-a61d-491a-8a42-0552a87d3c6a
# ╠═333a253e-ce8b-477a-9db2-f99c034ae048
# ╠═e81990c8-2194-472c-a102-b0d010f3a266
# ╠═f3fc2ea5-033c-4789-8ce6-0c6b15b607b7
# ╠═90596b8a-53d4-457f-9e93-522e6c5dc35d
# ╟─4e54ddc6-b9cd-472f-b496-18055ee5b667
# ╠═9844ec03-b169-40a0-ad56-1c07b11bb886
# ╠═49c39cea-895e-429a-909c-d1438d0dfd00
# ╠═fcdb96af-30bd-42c2-bf45-ecd13f27b57d
# ╠═5c4f40df-cde3-4140-ace2-6c0e1b60ebcf
# ╟─79578e4d-7527-459b-8915-ee5cb9ccad24
# ╠═a74c4034-fa19-4360-9f5e-191d5ba4001a
# ╟─8ce249a3-0881-49e5-8ffd-5df21bda50f4
# ╟─13131381-11b3-4db4-9351-fb7264c94ea4
# ╟─bc3ae07f-3877-480f-bd0b-633ce21172e3
# ╠═8bf6ba5b-1d06-41de-b22a-b58c1eaf2c83
# ╠═64507581-5751-496e-87b6-88eb320ae44c
# ╠═47c0516a-4a37-4823-a447-1207d5d4b8b4
# ╠═250559fb-1098-47ab-a55c-55a89497873d
# ╠═a21d8b13-20db-4df1-a2c7-451f8942466b
# ╟─420d1b15-8dc9-4c9a-b91f-9cc4b8101b92
# ╠═54ebf4c4-6692-4a40-a83e-c548440b3b85
# ╟─bbe52a98-9811-4c3a-9b4e-b90a1f5c163f
# ╠═73fd7360-ffdf-4870-a301-e39a982e0517
# ╠═37a815a0-c6cd-4ee8-9f46-ee4aa5259690
# ╠═3c09615a-e9c4-462b-9f23-5d5676aa29b7
# ╟─d3f2648a-a736-4812-b721-70298151a1a7
# ╠═9ef7592b-92f9-48bb-aa3e-9710bf251379
# ╠═9e63206c-82a8-4744-a5a4-a42281077c4f
# ╠═6899ad5c-3d31-40d7-94f4-042286b6c1c4
# ╠═b62519f5-26dd-49ca-afa4-2c72de7ef8a6
# ╠═fc25378b-11b1-4383-b144-75c7dc6e9a1f
# ╠═bd45654d-15e0-4717-962b-d5a302de4a9b
# ╠═c6ed3889-0d44-4180-9f81-3095055c6e54
# ╟─5a16e94e-b01b-4897-9631-41d345229be2
# ╠═06a666f9-c99e-41e5-b703-a5e826721d04
# ╟─255a1f5a-fe84-4ae2-85c3-a5a512a040ec
# ╠═9f24aec3-b93c-4e13-a227-0e4deefda611
# ╠═592e0ed4-bd97-4bb6-a826-b1e2f5ddbd59
# ╟─205a744b-e8a1-4acc-ba6b-7a7307f9a028
# ╠═daf4e474-0db4-4865-859d-96c31e7947f0
# ╠═7fb8fedf-5684-4874-8f17-d01debaff477
# ╟─9b9fb85b-13d0-43a2-ac5f-713593117b96
# ╠═2a48072b-fc5a-4044-9aa9-5661483ffbfc
# ╠═2cafea2d-b4b6-4ede-b899-720a71034fda
# ╠═643bf329-cce9-4c61-931d-65455b59357d
# ╟─7467608a-ce6d-4f29-8cd8-d766395f7129
# ╠═c7150c2c-afeb-4158-b381-3341bb678fd5
# ╠═6d7fd369-1f6c-487d-b399-c8cbbfede916
# ╠═e8bd9cb9-bd20-4619-93da-20e7829a61b4
# ╠═8a479712-7290-4c00-af7c-44e11ec81d0d
# ╠═372e190b-0344-4001-9c72-171610b41e9c
# ╠═b17cd903-bce6-4291-9d71-ea8a8422295a
# ╟─fc91074a-99fc-4a13-802c-740872bf7502
# ╟─fac6c740-93c1-4e12-9610-f3c621f84043
# ╠═15aa846c-7b81-4094-8b5b-dd6575e82afe
# ╠═be3b0811-c05a-4e78-9d34-8c52195dd7d0
# ╠═282acd5a-12e8-43f3-8de9-1b7b913a476c
# ╠═e5188537-b955-4637-bab5-a76b6d74c189
# ╠═c2b3cdc7-60b5-4287-8268-f95184d452b7
# ╠═4e3a272e-0280-4e27-b014-3f9da5ac8039
# ╠═a76e4081-7c52-4548-b513-ba03954378db
# ╠═f5fbd0da-fe4c-4162-be5e-f9d9d5b27d87
# ╠═a07cb23e-6fca-4c42-a6b6-a525b2b5e232
# ╠═c373046e-04b3-4159-80c0-8fbec141833f
# ╠═9c1ed9aa-9d52-425d-892d-af3b8a981663
# ╠═75006409-0338-470d-a41d-fa1aba5c4fca
# ╠═f2b3d8a0-4bea-4443-9bc3-37258106c3d3
# ╠═f6b06d6e-b268-45d6-a7e4-d8614cd0884d
# ╠═eab1fe28-f336-4b33-89a3-4c65ed2a818f


module Test_ode_pendulum_with_Measurement

#=
This test model demonstrates how to simulate a simple pendulum
with DifferentialEquations.jl using different types:

- model1: states and parameters are Float64
- model2: states and some parameters are Measurement{Float64}
          and the remaining parameters are Float64.

DifferentialEquations.jl stores by default t, x, der_x in its result data structure
at communication points (this can be switched off)
In this test model, additionally the complete model (all variables)
are stored at communication points at every time instant as well (= Vector{Model})
=#

using DifferentialEquations, Measurements, StaticArrays, PyPlot

mutable struct Model{R<:Real, F<:Real}
    m::R  # mass
    L::R  # length
    d::R  # damping
    g::F  # gravity (not used as Measurment)
    ϕ::R  # angle
    ω::R  # der(angle)
    α::R  # der(der(angle))
    r::SVector{2,R}   # absolute position of end point

    function Model{R,F}(m::R, L::R, d::R, g::F) where {R<:Real, F<:Real}
        new(m, L, d, g, 0, 0, 0, SVector{2,R}(0,0))
    end
end

Model(;m::R=0, L::R=0, d::R=0, g::F=0) where {R<:Real, F<:Real} = Model{R,F}(m,L,d,g)

#Model(m::R, L::R, d::R, g::F) where {R<:Real, F<:Real} = Model{R,F}(m,L,d,g)
#Model(;m::Real=0, L::Real=0, d::Real=0, g::Real=0) = Model(promote(m,L,d),...)

"simplePendulum: Compute m::Model at actual time instant"
function simplePendulum!(m::Model, t, x)::Nothing
    m.ϕ = x[1]
    m.ω = x[2]
    m.α = (-m.m*m.g*m.L*sin(m.ϕ) - m.d*m.ω) / (m.m*m.L^2)
    m.r = @SVector [m.L*sin(m.ϕ), -m.L*cos(m.ϕ)]
    return nothing
end

"SimplePendulum1: Called by integrator"
function simplePendulum1!(der_x, x, m::Model, t)
    simplePendulum!(m, t, x)
    der_x[1] = m.ω
    der_x[2] = m.α
end

"SimplePendulum2: Called at every communication point"
function simplePendulum2(x, t, integrator)
    m::Model = integrator.p
    simplePendulum!(m, t, x)
    return deepcopy(m)
end

t_start = 0.0
t_inc   = 0.01
t_end   = 7.0
tspan   = (t_start, t_end)
tspan2  = t_start:t_inc:t_end

# Problem 1 with Float64
model1 = Model(m=1.0, L=1.0, d=0.2, g=9.81)
# println("model1 = $model1")
x1₀ = [π/2, 0.0]

saved_values1=SavedValues(typeof(x1₀[1]), Model)
cb1 = SavingCallback(simplePendulum2, saved_values1, saveat=tspan2)

prob1 = ODEProblem(simplePendulum1!, x1₀, tspan, model1)
sol1 = solve(prob1, Tsit5(), reltol = 1e-6, saveat=t_inc, callback=cb1)
#prob1b = deepcopy(prob1)
#sol3 = solve(prob1, Tsit5(), reltol = 1e-6, saveat=t_inc)

figure(1)
clf()
plot(sol1.t, getindex.(sol1.u, 1), label="\$\\varphi\$")
grid(true)
legend()
title("SimplePendulum with Float64")

figure(2)
clf()
t = saved_values1.t
m = saved_values1.saveval
r = getfield.(m, :r)
r1 = getindex.(r,1)
r2 = getindex.(r,2)
plot(t, r1, label="r1")
plot(t, r2, label="r2")
grid(true)
legend()
title("SimplePendulum with Float64")

# Problem 2 with Measurement
model2 = Model(m=1.0 ± 0.1, L=1.00 ± 0.01, d=0.2± 0.05, g=9.81)
# println("model2 = $model2")
# println("typeof(model2.r) = ", typeof(model2.r))
x2₀ = [π/2 ± 0.1, 0.0 ± 0]

saved_values2=SavedValues(typeof(x2₀[1]), Model)
cb2 = SavingCallback(simplePendulum2, saved_values2, saveat=tspan2)

prob2 = ODEProblem(simplePendulum1!, x2₀, tspan, model2)
sol2  = solve(prob2, Tsit5(), reltol = 1e-6, saveat=t_inc, callback=cb2)

ϕ = getindex.(sol2.u, 1)
ϕm = Measurements.value.(ϕ)
ϕu = Measurements.uncertainty.(ϕ)
ϕmax = ϕm + ϕu
ϕmin = ϕm - ϕu

figure(3)
clf()
plot(sol2.t, ϕmax, label="\$\\varphi_{max}\$")
plot(sol2.t, ϕmin, label="\$\\varphi_{min}\$" )
plot(sol2.t, ϕm  , label="\$\\varphi_{mean}\$" )
grid(true)
legend()
title("SimplePendulum with Measurement{Float64}")

end


@model function form_model(y :: Vector{Float64})
    N = length(y)
    s = Vector(undef, N)

    start ~ Uniform(-1, 1)

    s[1] ~ Normal(start, 0.05)
    y[1] ~ Normal(s[1], 0.5)

    for i ∈ 2:N
        s[i] ~ Normal(s[i-1], 0.05)
        y[i] ~ Normal(s[i], 0.5)
    end
end

function run_form_analysis(games :: Dict{String, Vector{Game}})
    results = Dict{String, Pair}()
    teams = [team for team ∈ keys(games)]
    pbar = Progress(length(teams))
    Threads.@threads for team ∈ teams
        N = 100
        g = NUTS()
        samples = sample(form_model(discounted_score.(games[team])), g, N);

        s = @pipe describe(group(samples, :s))[1][:, :mean] |> vec
        s_σ = @pipe describe(group(samples, :s))[1][:, :std] |> vec

        results[team] = Pair(s, s_σ)
        next!(pbar)
    end
    return results
end


function close_eqs(sys::MomentEquations, closure_exp::OrderedDict,
                   closure::OrderedDict, polynorm::Bool)

    # TODO: improve performance

    closed_eqs = Equation[]
    for eq in sys.odes.eqs

        closed_rhs = substitute(eq.rhs, closure_exp)
        # apply binomial expansion on the expressions
        if polynorm # depending on the functional form
            closed_rhs = expand(closed_rhs)
        else
            closed_rhs = expand_expr(closed_rhs)
        end
        push!(closed_eqs, Equation(eq.lhs, closed_rhs))
    end

    iv = get_iv(sys.odes)
    ps = get_ps(sys.odes)

    vars = sys.odes.states[1:(length(sys.iter_1)+length(sys.iter_m))]

    odename = Symbol(nameof(sys), "_closed")
    odes = ODESystem(closed_eqs, iv, vars, ps; name=odename)

    ClosedMomentEquations(odes, closure, sys)

end

"""
    moment_closure(sys::MomentEquations, closure::String, binary_vars::Array{Int,1}=Int[])

Given `MomentEquations`, apply the specified moment closure approximation and return
the [`ClosedMomentEquations`](@ref).

The supported `closure` options are:
- [`"zero"`](@ref zero_closure)
- [`"normal"`](@ref normal_closure)
- [`"log-normal"`](@ref log-normal_closure)
- [`"poisson"`](@ref poisson_closure)
- [`"gamma"`](@ref gamma_closure)
- [`"derivative matching"`](@ref derivative_matching)
- [`"conditional gaussian"`](@ref conditional_gaussian_closure)
- [`"conditional derivative matching"`](@ref conditional_derivative_matching)

# Notes
- `binary_vars` *must* be specified for conditional closures as an array of indices of all species
  (as in [`speciesmap`](@ref)) which molecule number is a Bernoulli variable. Although not necessary
  for other closures, specifying `binary_vars` is recommended as the properties of Bernoulli variables
  will be used to remove the redundant moment equations and simplify the expressions, which can
  significantly improve numerical stability.
"""
function moment_closure(sys::MomentEquations, closure::String, binary_vars::Array{Int,1}=Int[])

    if closure == "zero"
        zero_closure(sys, binary_vars)
    elseif closure == "normal"
        normal_closure(sys, binary_vars)
    elseif closure == "log-normal"
        log_normal_closure(sys, binary_vars)
    elseif closure == "poisson"
        poisson_closure(sys, binary_vars)
    elseif closure == "gamma"
        gamma_closure(sys, binary_vars)
    elseif closure == "derivative matching"
        derivative_matching(sys, binary_vars)
    elseif closure == "conditional gaussian"
        conditional_gaussian_closure(sys, binary_vars)
    elseif closure == "conditional derivative matching"
        conditional_derivative_matching(sys, binary_vars)
    else
        error(closure*" closure does not exist")
    end

end


"""
Provides the basic functionality to interact with the Collatz conjecture. The
parameterisation uses the same (P,a,b) notation as Conway's generalisations.
Besides the function and reverse function, there is also functionality to retrieve
the hailstone sequence, the "stopping time"/"total stopping time", or tree-graph.

# Examples
```
julia> import Pkg; Pkg.add("Collatz"); using Collatz
```
"""
module Collatz
#TODO: ^ add jldoctest back to the above example when this is in general

export _ErrMsg, _CC, _KNOWN_CYCLES
export collatz_function, reverse_collatz_function, hailstone_sequence, stopping_time, tree_graph

"The four known cycles (besides 0->0), for the default parameterisation."
const _KNOWN_CYCLES = [[1, 4, 2], [-1, -2], [-5, -14, -7, -20, -10], [-17, -50, -25, -74, -37, -110, -55, -164, -82, -41, -122, -61, -182, -91, -272, -136, -68, -34]]
"The value up to which has been proven numerically, for the default parameterisation."
const __VERIFIED_MAXIMUM = 295147905179352825856
"The value down to which has been proven numerically, for the default parameterisation."
const __VERIFIED_MINIMUM = -272  #&TODO: Check the actual lowest bound.

"""
Error message constant.
"""
module _ErrMsg
const SANE_PARAMS_P = "'P' should not be 0 ~ violates modulo being non-zero."
const SANE_PARAMS_A = "'a' should not be 0 ~ violates the reversability."
end # module _ErrMsg
import ._ErrMsg

"""
Cycle Control: Descriptive flags to indicate when some event occurs in the
hailstone sequences, when set to verbose, or stopping time check. A module used
to wrap an enum to reference the values as `_CC.ABC` rather than `ABC::CC`

# Examples
```jldoctest
julia> string(_CC.STOPPING_TIME)
"STOPPING_TIME"
julia> string(_CC.TOTAL_STOPPING_TIME)
"TOTAL_STOPPING_TIME"
julia> string(_CC.CYCLE_INIT)
"CYCLE_INIT"
julia> string(_CC.CYCLE_LENGTH)
"CYCLE_LENGTH"
julia> string(_CC.MAX_STOP_OOB)
"MAX_STOP_OOB"
julia> string(_CC.ZERO_STOP)
"ZERO_STOP"
```
"""
module _CC
# The elements of an enum are accessed as string(ABC::CC) or just string(ABC) so
# wrap the enum in a module so they can be referenced as string(_CC.ABC)
@enum CC begin
    STOPPING_TIME
    TOTAL_STOPPING_TIME
    CYCLE_INIT
    CYCLE_LENGTH
    MAX_STOP_OOB  # ~ "out of bounds"
    ZERO_STOP
end # @enum CC
end # module _CC
import ._CC

"""
    __assert_sane_parameterisation(P, a, b)

Handles the sanity check for the parameterisation (P,a,b) required by both
the function and reverse function.

# Args
- `P::Integer`: Modulus used to devide n, iff n is equivalent to (0 mod P).
- `a::Integer`: Factor by which to multiply n.
- `b::Integer`: Value to add to the scaled value of n.
"""
function __assert_sane_parameterisation(P::Integer, a::Integer, b::Integer)
    # Sanity check (P,a,b) ~ P absolutely can't be 0. a "could" be zero
    # theoretically, although would violate the reversability (if ~a is 0 then a
    # value of "b" as the input to the reverse function would have a pre-emptive
    # value of every number not divisible by P). The function doesn't _have_ to
    # be reversable, but we are only interested in dealing with the class of
    # functions that exhibit behaviour consistant with the collatz function. If
    # _every_ input not divisable by P went straight to "b", it would simply
    # cause a cycle consisting of "b" and every b/P^z that is an integer. While
    # P in [-1, 1] could also be a reasonable check, as it makes every value
    # either a 1 or 2 length cycle, it's not strictly an illegal operation.
    # "b" being zero would cause behaviour not consistant with the collatz
    # function, but would not violate the reversability, so no check either.
    # " != 0" is redundant for python assertions.
    @assert P != zero(P) _ErrMsg.SANE_PARAMS_P
    @assert a != zero(a) _ErrMsg.SANE_PARAMS_A
end

"""
    collatz_function(n; P=2, a=3, b=1)

Returns the output of a single application of a Collatz-esque function.

# Args
- `n::Integer`: The value on which to perform the Collatz-esque function

# Kwargs
- `P::Integer=2`: Modulus used to devide n, iff n is equivalent to (0 mod P).
- `a::Integer=3`: Factor by which to multiply n.
- `b::Integer=1`: Value to add to the scaled value of n.

# Examples
```jldoctest
julia> collatz_function(5)
16
```
```jldoctest
julia> collatz_function(14, P=7)
2
```
```jldoctest
julia> collatz_function(15, P=7, a=5, b=3)
78
```
"""
function collatz_function(n::Integer; P::Integer=2, a::Integer=3, b::Integer=1)
    __assert_sane_parameterisation(P,a,b)
    if n%P == 0
        # In Julia, "Integer//Integer" creates a "Rational{Int64}" type, rather
        # than strict integer division (which is done instead with "÷"). See;
        # https://docs.julialang.org/en/v1/manual/mathematical-operations/
        # https://docs.julialang.org/en/v1/manual/complex-and-rational-numbers/
        # So we can either "n÷P", or "numerator(n//P)"
        return n÷P
    else
        return (a*n+b)
    end
end

"""
    reverse_collatz_function(n; P=2, a=3, b=1)

Returns the output of a single application of a Collatz-esque reverse function.

# Args
- `n::Integer`: The value on which to perform the reverse Collatz function

# Kwargs
- `P::Integer=2`: Modulus used to devide n, iff n is equivalent to (0 mod P).
- `a::Integer=3`: Factor by which to multiply n.
- `b::Integer=1`: Value to add to the scaled value of n.

# Examples
```jldoctest
julia> print(reverse_collatz_function(1))
[2]
```
```jldoctest
julia> print(reverse_collatz_function(4))
[1, 8]
```
```jldoctest
julia> print(reverse_collatz_function(3, P=-3, a=-2, b=-5))
[-9, -4]
```
"""
function reverse_collatz_function(n::Integer; P::Integer=2, a::Integer=3, b::Integer=1)
    __assert_sane_parameterisation(P,a,b)
    # Every input can be reversed as the result of "n/P" division, which yields
    # "Pn"... {f(n) = an + b}≡{(f(n) - b)/a = n} ~ if n was such that the
    # muliplication step was taken instead of the division by the modulus, then
    # (f(n) - b)/a) must be an integer that is not in (0 mod P). Because we're
    # not placing restrictions on the parameters yet, although there is a better
    # way of shortcutting this for the default variables, we need to always
    # attempt (f(n) - b)/a)
    pre_values = [P*n]
    if (n-b)%a == 0 && (n-b)%(P*a) != 0
        push!(pre_values, ((n-b)÷a))
        sort!(pre_values)
    end
    return pre_values
end

"""
    __initial_value_outside_verified_range(x)

Checks if the initial value is greater than __VERIFIED_MAXIMUM or less than
__VERIFIED_MINIMUM. Only intended for the default parameterisation.

# Args
- `x::Integer`: The initial value to check if it is within range or not.
"""
function __initial_value_outside_verified_range(x::Integer)
    return (__VERIFIED_MAXIMUM < x) || (x < __VERIFIED_MINIMUM)
end

"""
    __stopping_time_terminus(n, total_stop)

Provides the appropriate lambda to use to check if iterations on an initial
value have reached either the stopping time, or total stopping time.

# Args
- `n::Integer`: The initial value to confirm against a stopping time check.
- `total_stop::Bool`: If false, the lambda will confirm that iterations
    of n have reached the oriented stopping time to reach a value closer
    to 0. If true, the lambda will simply check equality to 1.
"""
function __stopping_time_terminus(n::Integer, total_stop::Bool)
    # https://docs.julialang.org/en/v1/manual/functions/#man-anonymous-functions
    if total_stop
        return (function (x) x == 1 end)
    elseif n >= 0
        return (function (x) x < n && x > 0 end)
    else
        return (function (x) x > n && x < 0 end)
    end
end

#TODO: Make the hailstone calc arbitrary integer safe!
"""
    hailstone_sequence(initial_value; P=2, a=3, b=1, max_total_stopping_time=1000, total_stopping_time=true, verbose=true)

Returns a list of successive values obtained by iterating a Collatz-esque function.

Until either 1 is reached, or the total amount of iterations
exceeds max_total_stopping_time. Unless total_stopping_time is False,
which will terminate the hailstone at the "stopping time" value, i.e. the
first value less than the initial value. While the sequence has the
capability to determine that it has encountered a cycle, the cycle from "1"
wont be attempted or reported as part of a cycle, regardless of default or
custom parameterisation, as "1" is considered a "total stop".

# Args
- `initial_value::Integer`: The value to begin the hailstone sequence from.

# Kwargs
- `P::Integer=2`: Modulus used to devide n, iff n is equivalent to (0 mod P).
- `a::Integer=3`: Factor by which to multiply n.
- `b::Integer=1`: Value to add to the scaled value of n.
- `max_total_stopping_time::Integer=1000`: Maximum amount of times to iterate the
    function, if 1 is not reached.
- `total_stopping_time::Bool=true`: Whether or not to execute until the "total"
    stopping time (number of iterations to obtain 1) rather than the regular stopping
    time (number of iterations to reach a value less than the initial value).
- `verbose::Bool=true`: If set to verbose, the hailstone sequence will include
    control string sequences to provide information about how the sequence
    terminated, whether by reaching a stopping time or entering a cycle.

# Examples
```jldoctest
julia> print(hailstone_sequence(16, verbose=false))
[16, 8, 4, 2, 1]
```
```jldoctest
julia> print(hailstone_sequence(16))
Any[16, 8, 4, 2, 1, Any[Collatz._CC.TOTAL_STOPPING_TIME, 4]]
```
```jldoctest
julia> print(hailstone_sequence(761, P=5, a=2, b=3, verbose=false))
[761, 1525, 305, 61, 125, 25, 5, 1]
```
# Example cycle!
```jldoctest
julia> print(hailstone_sequence(-56))
Any[-56, -28, Collatz._CC.CYCLE_INIT, Any[-14, -7, -20, -10, -5], Any[Collatz._CC.CYCLE_LENGTH, 5]]
```
```jldoctest
julia> print(hailstone_sequence(-56, verbose=false))
[-56, -28, -14, -7, -20, -10, -5, -14]
```
"""
function hailstone_sequence(initial_value::Integer; P::Integer=2, a::Integer=3, b::Integer=1, max_total_stopping_time::Integer=1000, total_stopping_time::Bool=true, verbose::Bool=true)
    # Call out the collatz_function before any magic returns to trap bad values.
    _ = collatz_function(initial_value,P=P,a=a,b=b)
    # 0 is always an immediate stop.
    if initial_value == 0; if verbose; return [[_CC.ZERO_STOP, 0]]; else; return [0]; end; end
    # 1 is always an immediate stop, with 0 stopping time.
    if initial_value == 1; if verbose; return [[_CC.TOTAL_STOPPING_TIME, 0]]; else; return [1]; end; end
    terminate = __stopping_time_terminus(initial_value, total_stopping_time)
    # Start the hailstone sequence.
    _max_total_stopping_time = max(max_total_stopping_time, 1)
    hailstone = [initial_value] # Initialises as Vector{Integer}
    if verbose; hailstone = Vector{Any}(hailstone); end
    cyclic = (function (x) x in hailstone end)
    for k in 1:_max_total_stopping_time
        _next = collatz_function(last(hailstone),P=P,a=a,b=b)
        # Check if the next hailstone is either the stopping time, total
        # stopping time, the same as the initial value, or stuck at zero.
        if terminate(_next)
            push!(hailstone, _next)
            if verbose
                if _next == 1
                    m = _CC.TOTAL_STOPPING_TIME
                else
                    m = _CC.STOPPING_TIME
                end
                push!(hailstone, [m, length(hailstone)-1])
            end
            break
        end
        if cyclic(_next)
            cycle_init = 0
            for j in 0:(length(hailstone)-1)
                if hailstone[length(hailstone)-j] == _next
                    # Because Julia is 1-based indexed, add 1 to the J
                    cycle_init = j+1
                    break
                end
            end
            if verbose
                hailstone = append!(hailstone[1:(length(hailstone)-cycle_init)], [_CC.CYCLE_INIT, hailstone[(length(hailstone)-cycle_init+1):length(hailstone)], [_CC.CYCLE_LENGTH, cycle_init]])
            else
                push!(hailstone, _next)
            end
            break
        end
        if _next == 0
            push!(hailstone, 0)
            if verbose
                push!(hailstone, [_CC.ZERO_STOP, -(length(hailstone)-1)])
            end
            break
        end
        push!(hailstone, _next)
    end
    # Julia has no for ~ else yet https://github.com/JuliaLang/julia/issues/1289
    if length(hailstone) >= max_total_stopping_time
        if verbose
            push!(hailstone, [_CC.MAX_STOP_OOB, _max_total_stopping_time])
        end
    end
    return hailstone
end

"""
    stopping_time(initial_value; P=2, a=3, b=1, max_stopping_time=1000, total_stopping_time=false)

Returns the stopping time, the amount of iterations required to reach a
value less than the initial value, or nothing if max_stopping_time is exceeded.

Alternatively, if total_stopping_time is True, then it will instead count
the amount of iterations to reach 1. If the sequence does not stop, but
instead ends in a cycle, the result will be infinity. If (P,a,b) are such
that it is possible to get stuck on zero, the result will be the negative of
what would otherwise be the "total stopping time" to reach 1, where 0 is
considered a "total stop" that should not occur as it does form a cycle of
length 1.

# Args
- `initial_value::Integer`: The value for which to find the stopping time.

# Kwargs
- `P::Integer=2`: Modulus used to devide n, iff n is equivalent to (0 mod P).
- `a::Integer=3`: Factor by which to multiply n.
- `b::Integer=1`: Value to add to the scaled value of n.
- `max_stopping_time::Integer=1000`: Maximum amount of times to iterate the function, if the stopping
    time is not reached. IF the max_stopping_time is reached, the function will return nothing.
- `total_stopping_time::Bool=false`: Whether or not to execute until the "total"
    stopping time (number of iterations to obtain 1) rather than the regular
    stopping time (number of iterations to reach a value less than the initial value).

# Examples
```jldoctest
julia> stopping_time(5)
3
julia> stopping_time(27)
96
julia> stopping_time(21, P=5, a=2, b=3, total_stopping_time=true)
Inf
```
"""
function stopping_time(initial_value::Integer; P::Integer=2, a::Integer=3, b::Integer=1, max_stopping_time::Integer=1000, total_stopping_time::Bool=false)
    # The information is contained in the verbose form of a hailstone sequence.
    # Although the "max_~_time" for hailstones is name for "total stopping" time
    # and the "max_~_time" for this "stopping time" function is _not_ "total",
    # they are handled the same way, as the default for "total_stopping_time"
    # for hailstones is true, but for this, is false. Thus the naming difference
    end_msg = last(hailstone_sequence(initial_value, P=P, a=a, b=b, verbose=true, max_total_stopping_time=max_stopping_time, total_stopping_time=total_stopping_time))
    # For total/regular/zero stopping time, the value is already the same as
    # that present, for cycles we report infinity instead of the cycle length,
    # and for max stop out of bounds, we report nothing instead of the max stop cap
    return get(Dict(_CC.TOTAL_STOPPING_TIME => end_msg[2],
                    _CC.STOPPING_TIME => end_msg[2],
                    _CC.CYCLE_LENGTH => Inf, # infinity
                    _CC.ZERO_STOP => end_msg[2],
                    _CC.MAX_STOP_OOB => nothing,
                    ), end_msg[1], nothing)
end

# TODO: Determine why the output of print(tree_graph(1, 3)) in the example in the below
# yields a "Dict{Any, Any}(Collatz._CC.CYCLE_INIT => 1, 8 => Dict{Any, Any}())" in jldoctest
# but a "Dict{Any, Any}(8 => Dict{Any, Any}(), Collatz._CC.CYCLE_INIT => 1)" from REPL?
# Both examples are covered in tests, which doesn't appear to care about order, but jldoctest
# is just a stringy parser?? so can't jldoctest these two.
"""
    tree_graph(initial_value, max_orbit_distance; P=2, a=3, b=1, __cycle_prevention=nothing)

Returns nested dictionaries that model the directed tree graph up to a
maximum nesting of max_orbit_distance, with the initial_value as the root.

# Args
- `initial_value::Integer`: The root value of the directed tree graph.
- `max_orbit_distance::Integer`: Maximum amount of times to iterate the reverse
    function. There is no natural termination to populating the tree graph, equivalent
    to the termination of hailstone sequences or stopping time attempts, so this is not
    an optional argument like max_stopping_time / max_total_stopping_time, as it is the
    intended target of orbits to obtain, rather than a limit to avoid uncapped computation.

# Kwargs
- `P::Integer=2`: Modulus used to devide n, iff n is equivalent to (0 mod P).
- `a::Integer=3`: Factor by which to multiply n.
- `b::Integer=1`: Value to add to the scaled value of n.

# Internal Kwargs
- `__cycle_prevention::Union{Set{Integer},Nothing}=nothing`: Used to prevent cycles
    from precipitatingby keeping track of all values added across previous nest depths.

# Examples
```
julia> print(tree_graph(1, 3))
Dict{Int64, Dict{Any, Any}}(1 => Dict(2 => Dict{Any, Any}(4 => Dict{Any, Any}(Collatz._CC.CYCLE_INIT => 1, 8 => Dict{Any, Any}()))))
```
```
julia> print(tree_graph(4, 3))
Dict{Int64, Dict{Any, Any}}(4 => Dict(8 => Dict{Any, Any}(16 => Dict{Any, Any}(5 => Dict{Any, Any}(), 32 => Dict{Any, Any}())), 1 => Dict{Any, Any}(2 => Dict{Any, Any}(Collatz._CC.CYCLE_INIT => 4))))
```
"""
function tree_graph(initial_value::Integer, max_orbit_distance::Integer; P::Integer=2, a::Integer=3, b::Integer=1, __cycle_prevention::Union{Set{Integer},Nothing}=nothing)
    # Call out the reverse_collatz_function before any magic returns to trap bad values.
    _ = reverse_collatz_function(initial_value,P=P,a=a,b=b)
    tgraph = Dict(initial_value=>Dict())
    if max(0, max_orbit_distance) == 0; return tgraph; end
    # Handle cycle prevention for recursive calls ~
    # Shouldn't use a mutable object initialiser for a default.
    if __cycle_prevention == nothing; __cycle_prevention = Set{Integer}(); end
    push!(__cycle_prevention, initial_value)
    for branch_value in reverse_collatz_function(initial_value, P=P, a=a, b=b)
        if branch_value in __cycle_prevention
            tgraph[initial_value][_CC.CYCLE_INIT] = branch_value
        else
            tgraph[initial_value][branch_value] = tree_graph(branch_value, max_orbit_distance-1, P=P, a=a, b=b, __cycle_prevention=__cycle_prevention)[branch_value]
        end
    end
    return tgraph
end

end # module


# FFT test

Fs = 1000;            # Sampling frequency
T = 1/Fs;             # Sampling period

FFTW.set_num_threads(Sys.CPU_CORES)

function prepare_input(input_size)
	t = (0:input_size-1)*T;        # Time vector

	S = 0.7*sin.(2pi*50t) + sin.(2pi*120t);
	return S + 2*randn(size(t));
end

function test_function(input_data)
	return fft(input_data)
end


using ImplicitEquations
using Test

f(x,y) = y-x
g(x,y) = y+x

ImplicitEquations.GRAPH(f ⩵ 0, -5, 5, -5, 5, 2^4, 2^4)
ImplicitEquations.GRAPH(f ≦ 0, -5, 5, -5, 5, 2^4, 2^4)
ImplicitEquations.GRAPH(f ≪  g, -5, 5, -5, 5, 2^4, 2^4)

ImplicitEquations.GRAPH(f ≶  g, -5, 5, -5, 5, 2^4, 2^4)
ImplicitEquations.GRAPH(f ≷  g, -5, 5, -5, 5, 2^4, 2^4)
ImplicitEquations.GRAPH(Neq(f, g), -5, 5, -5, 5, 2^4, 2^4)

ImplicitEquations.GRAPH(Eq(f, g), -5, 5, -5, 5, 2^4, 2^4)
ImplicitEquations.GRAPH(f ⩵ g, -5, 5, -5, 5, 2^4, 2^4)


#!/usr/bin/env julia
# -*- coding: utf-8 -*-
# File   : sphere_approx.jl
# License: MIT
# Author : Andrei Leonard Nicusan <a.l.nicusan@bham.ac.uk>
# Date   : 21.05.2021

using GLMakie
using Polyhedra
using ColorSchemes
using LinearAlgebra

import QHull
lib = QHull.Library()

# Return points on a sphere, generated with the method of Markus Deserno from
# https://www.cmu.edu/biolphys/deserno/pdf/sphere_equi.pdf
function sphere(;centre = [0., 0., 0.], radius = 1., n = 16)
    nc = 0
    a = 4π * radius^2 / n
    d = √a

    mθ = floor(π / d)

    dθ = π / mθ
    dϕ = a / dθ

    points = zeros(n, 3)

    for m in 0:mθ - 1
        θ = π * (m + 0.5) / mθ
        mϕ = floor(2π * sin(θ) / dϕ)
        for p in 0:mϕ - 1
            ϕ = 2π * p / mϕ
            nc += 1

            points[nc, 1] = centre[1] + radius * sin(θ) * cos(ϕ)
            points[nc, 2] = centre[2] + radius * sin(θ) * sin(ϕ)
            points[nc, 3] = centre[3] + radius * cos(θ)
        end
    end

    return points[1:nc, :]
end

# Very dense circle points
sphere_points = sphere(n = 512)
sphere_poly = polyhedron(vrep(sphere_points), lib)

# Plot polygonal approximations and area ratio
cm = ColorSchemes.magma
# GLMakie.mesh(Polyhedra.Mesh{3}(sphere_poly), color = get(cm, 0.))

nverts = collect(10:25:500)   # [16, 32, 64, 128, 256]
ratios = zeros(length(nverts))
colors = [get(cm, i) for i in LinRange(0.1, 0.9, length(nverts))]

for (i, n) in enumerate(nverts)
    local sphere_points = sphere(centre = [0., 0. - i / 10, 0.], n = n)

    local sphere_poly = polyhedron(vrep(sphere_points), lib)
    # GLMakie.mesh!(Polyhedra.Mesh{3}(sphere_poly), color = colors[i])

    ratios[i] = Polyhedra.volume(sphere_poly) / (4 / 3 * pi)
end

current_figure()

# Ploting
GLMakie.lines(nverts, ratios)


module TestParseJuliaPrograms

using Test

using Catlab, Catlab.Theories, Catlab.WiringDiagrams
using Catlab.Programs
using Catlab.Programs.ParseJuliaPrograms: normalize_arguments

@present C(FreeBiproductCategory) begin
  W::Ob
  X::Ob
  Y::Ob
  Z::Ob
  f::Hom(X,Y)
  g::Hom(Y,Z)
  h::Hom(Z,W)
  l::Hom(otimes(Z,W),otimes(Y,X))
  m::Hom(otimes(X,Y),otimes(W,Z))
  n::Hom(otimes(W,Z),otimes(X,Y))
end

# Parsing
#########

W, X, Y, Z = (generator(C, name) for name in [:W, :X, :Y, :Z])
f, g, h, m, n = (generator(C, name) for name in [:f, :g, :h, :m, :n])

# Generator.

parsed = @program(C, (x::X) -> f(x))
@test parsed == to_wiring_diagram(f)
@test parsed isa WiringDiagram{BiproductCategory}

# Composition: one-dimensional.

parsed = @program(C, (x::X) -> g(f(x)))
@test parsed == to_wiring_diagram(compose(f,g))

parsed = @program C (x::X) begin
  y = f(x)
  z = g(y)
  return z
end
@test parsed == to_wiring_diagram(compose(f,g))

parsed = @program C (v::X) begin
  v = f(v)
  v = g(v)
  v = h(v)
  v
end
@test parsed == to_wiring_diagram(compose(f,g,h))

# Composition: multidimensional.

parsed = @program C (x::X, y::Y) begin
  w, z = m(x, y)
  x2, y2 = n(w, z)
  return (x2, y2)
end
@test parsed == to_wiring_diagram(compose(m,n))

parsed = @program(C, (x::X, y::Y) -> n(m(x,y)))
@test parsed == to_wiring_diagram(compose(m,n))

# Constants and co-constants.

I = munit(FreeBiproductCategory.Ob)
c = add_generator!(C, Hom(:c, I, X))
d = add_generator!(C, Hom(:d, X, I))

parsed = @program(C, () -> c())
@test parsed == to_wiring_diagram(c)

parsed = @program(C, (x::X) -> c(d(x)))
@test parsed == to_wiring_diagram(compose(d,c))

parsed = @program C (x::X) begin
  d(x)
  c()
end
@test parsed == to_wiring_diagram(compose(d,c))

# Diagonals: implicit syntax.

parsed = @program(C, (x::X) -> (f(x),f(x)))
@test parsed == to_wiring_diagram(compose(mcopy(X),otimes(f,f)))

parsed = @program(C, (x::X) -> nothing)
@test parsed == to_wiring_diagram(delete(X))

# Codiagonals: special syntax.

parsed = @program(C, (x1::X, x2::X) -> f([x1,x2]))
@test parsed == to_wiring_diagram(compose(mmerge(X),f))

parsed = @program(C, (x1::X, x2::X) -> [f(x1),f(x2)])
@test parsed == to_wiring_diagram(compose(otimes(f,f),mmerge(Y)))

parsed = @program(C, (x::X, y::Y) -> n([m(x,y), m(x,y)]))
@test parsed == to_wiring_diagram(compose(
  mcopy(otimes(X,Y)), otimes(m,m), mmerge(otimes(W,Z)), n))

parsed = @program(C, () -> f([]))
@test parsed == to_wiring_diagram(compose(create(X),f))

# Explicit syntax for special objects and morphisms.

parsed = @program(C, (x::X) -> id{X}(x))
@test parsed == to_wiring_diagram(id(X))

parsed = @program(C, (x::X) -> mcopy{X}(x))
@test parsed == to_wiring_diagram(mcopy(X))

parsed = @program(C, (x::X) -> delete{X}(x))
@test parsed == to_wiring_diagram(delete(X))

parsed = @program(C, (x1::X, x2::X) -> mmerge{X}(x1,x2))
@test parsed == to_wiring_diagram(mmerge(X))

parsed = @program(C, () -> create{X}())
@test parsed == to_wiring_diagram(create(X))

parsed = @program(C, (x::X, y::Y) -> mcopy{otimes{X,Y}}(x,y))
@test parsed == to_wiring_diagram(mcopy(otimes(X,Y)))

parsed = @program(C, (x::X, y::Y) -> delete{otimes{X,Y}}(x,y))
@test parsed == to_wiring_diagram(delete(otimes(X,Y)))

parsed = @program(C, (x::X) -> Δ{X}(x))
@test parsed == to_wiring_diagram(mcopy(X))

parsed = @program(C, (x1::X, x2::X) -> ∇{X}(x1,x2))
@test parsed == to_wiring_diagram(mmerge(X))

parsed = @program(C, (xy::otimes{X,Y}) -> m(xy))
@test parsed == to_wiring_diagram(m)

parsed = @program C (xy::otimes{X,Y}) begin
  x, y = xy
  (f(x), g(y))
end
@test parsed == to_wiring_diagram(otimes(f,g))

parsed = @program(C, (xy::otimes{X,Y}, wz::otimes{W,Z}) -> (m(xy), n(wz)))
@test parsed == to_wiring_diagram(otimes(m,n))

# Helper function: normalization of arguments.

normalize(args...) = normalize_arguments(Tuple(args))

@test normalize(nothing) == ()
@test normalize(1) == ([1],)
@test normalize(1,nothing,2) == ([1],[2])
@test normalize(1,2,3) == ([1],[2],[3])
@test normalize((1,2,3)) == ([1],[2],[3])
@test normalize((1,2),3) == ([1],[2],[3])
@test normalize(((1,2),3)) == ([1],[2],[3])

@test normalize([]) == ([],)
@test normalize([1]) == ([1],)
@test normalize([1,2]) == ([1,2],)
@test normalize([1,2],[3,4]) == ([1,2],[3,4])
@test normalize([(1,2),(3,4)]) == ([1,3],[2,4])
@test normalize([[(1,2),(3,4)],(5,6)]) == ([1,3,5],[2,4,6])

# Roundtrip
###########

""" Convert morphism expression to Julia code and then back to wiring diagram.
"""
function roundtrip(f::HomExpr)::WiringDiagram
  expr = compile_expr(f, arg_types=first.(collect(dom(f))))
  parse_wiring_diagram(C, expr)
end

function test_roundtrip(f::HomExpr)
  @test roundtrip(f) == to_wiring_diagram(f)
end

f, g, h, l, m, n = (generator(C, name) for name in [:f, :g, :h, :l, :m, :n])

test_roundtrip(f)
test_roundtrip(id(X))
test_roundtrip(compose(f,g))
test_roundtrip(otimes(f,h))
test_roundtrip(compose(m,n))
test_roundtrip(compose(l, braid(Y,X), m))
test_roundtrip(compose(mcopy(X), otimes(f,f)))
test_roundtrip(compose(f, mcopy(Y), otimes(g,g)))

end


module FESpaceWithConstantFixedTests
using Gridap
using Gridap.Geometry
using Gridap.FESpaces
using Test

domain = (0,1,0,1)
partition = (4,4)
model = CartesianDiscreteModel(domain,partition)

trian = get_triangulation(model)

order = 2

V = TestFESpace(
  model=model,
  reffe=:Lagrangian,
  valuetype=Float64,
  order=order,
  conformity=:L2)

V0 = FESpaceWithConstantFixed(V,true,rand(1:num_free_dofs(V)))
test_single_field_fe_space(V0)

@test Gridap.FESpaces.get_constant_approach(V0) == Gridap.FESpaces.FixConstant()

uh0 = interpolate(V0) do x
    sin(4*pi*(x[1]+x[2]^2)) + 3
end
#using Gridap.Visualization
#
#writevtk(trian,"trian",nsubcells=20,cellfields=["uh0"=>uh0])

end # module


# ASYBC.M
# This function does bias-correction based on the formula by Pope (1990)
# Luca Brugnolini 2014
# Based on Kilian

function  asybc(A,SIGMA,t,K,p)

# Input: Slope parameter estimates in compaion form,
#        variance of estimated residuals in companion form,
#        sample length, number of variables, lag order
# Output: Bias-corrected slope parameter estimates
T = t-p;
# vecSIGMAY=inv(eye((K*p)^2)-kron(A,A))*vec(SIGMA);  % Lutkepohl p.29 (2.1.39)
# SIGMAs = sparse(SIGMA[:]); As = sparse(A); eyes = sparse(eye((K*p)^2));
# vecSIGMAY = (eyes - kron(As, As))\SIGMAs; #***

vecSIGMAY = (eye((K*p)^2)-kron(A,A))\vec(SIGMA);
SIGMAY = reshape(vecSIGMAY, K*p, K*p);
I = eye(K*p, K*p);
B = A';

# There are K*p eigenvalues by construction
peigen = eigvals(A);
sumeig = zeros(K*p, K*p);

    for h = 1:K*p
        sumeig = sumeig + (peigen[h].\(I - peigen[h]*B)); #***

    end;

bias = SIGMA*(inv(I - B) + B/(I-B^2) + sumeig)/(SIGMAY); #***

Abias = - bias/T;

#############################################################

bcstab = 9;         # Arbitrary default value > 1
delta = 1;            # Adjustment factor

while bcstab >= 1

# Adjust bias-correction proportionately
global bcA = A-delta*Abias;
bcmod = abs(eigvals(bcA));

if  any(bcmod.>= 1)
    bcstab = 1;

else
    bcstab = 0;

end

delta += - 0.01;

if delta <= 0
   bcstab = 0;

end

end

return bcA

end


using ApproximateBayesianComputation

data = rand(Distributions.Normal(2.0, 0.5), 100)
sample_prior = make_sample_prior(Distributions.Normal(0.0, 0.5))
forward_model(μ::Float64) =  rand(Distributions.Normal(μ, 0.5), 100)

summaryStatistics1 = mean(data)
threshold1 = 0.05
compute_distance1(x::Array{Float64, 1}, y::Float64) =  abs(mean(x) - y)

summaryStatistics2 = [mean(data), median(data)]
threshold2 = [0.05, 0.05]
compute_distance2(x::Array{Float64, 1}, y::Array{Float64, 1}) = abs([mean(x), median(x)] - y)

N = 200

@time out00 = abc_standard(summaryStatistics1, N, 0.5, sample_prior, forward_model, compute_distance1)

@time out0 = abc_standard(summaryStatistics1, N, 0.25, sample_prior, forward_model, compute_distance1)

@time out1 = abc_standard(summaryStatistics1, N, 0.05, sample_prior, forward_model, compute_distance1)

@time out11 = abc_standard(summaryStatistics1, N, 0.01, sample_prior, forward_model, compute_distance1)

@time out12 = abc_standard(summaryStatistics1, N, 0.005, sample_prior, forward_model, compute_distance1)

@time out2 = abc_standard(summaryStatistics2, N, threshold2, sample_prior, forward_model, compute_distance2)


module PoissonDGTests

using Test
using Gridap
import Gridap: ∇
#using LinearAlgebra

#domain = (0,1,0,1)
#partition = (4,4)
#model = CartesianDiscreteModel(domain,partition)
#const h = (domain[2]-domain[1]) / partition[1]

using Gridap.Geometry: DiscreteModelMock
model = DiscreteModelMock()
const h = 1

order = 2
const γ = 10

Ω = Triangulation(model)
Γ = BoundaryTriangulation(model)
Λ = SkeletonTriangulation(model)

degree = order
dΩ = Measure(Ω,degree)
dΓ = Measure(Γ,degree)
dΛ = Measure(Λ,degree)

n_Γ = get_normal_vector(Γ)
n_Λ = get_normal_vector(Λ)

u(x) = x[1]^2 + x[2]
∇u(x) = VectorValue( 2*x[1], one(x[2]) )
Δu(x) = 2
f(x) = - Δu(x)
∇(::typeof(u)) = ∇u

V = TestFESpace(model,ReferenceFE(lagrangian,Float64,order),conformity=:L2)
U = TrialFESpace(V,u)

a(u,v) =
  ∫( ∇(v)⋅∇(u) )*dΩ +
  ∫( (γ/h)*v*u  - v*(n_Γ⋅∇(u)) - (n_Γ⋅∇(v))*u )*dΓ +
  ∫( (γ/h)*jump(v*n_Λ)⋅jump(u*n_Λ) - jump(v*n_Λ)⋅mean(∇(u)) -  mean(∇(v))⋅jump(u*n_Λ) )*dΛ

l(v) =
  ∫( v*f )*dΩ +
  ∫( (γ/h)*v*u - (n_Γ⋅∇(v))*u )*dΓ

op = AffineFEOperator(a,l,U,V)

uh = solve(op)

e = u - uh

l2(u) = sqrt(sum( ∫( u⊙u )*dΩ ))
h1(u) = sqrt(sum( ∫( u⊙u + ∇(u)⊙∇(u) )*dΩ ))

el2 = l2(e)
eh1 = h1(e)
ul2 = l2(uh)
uh1 = h1(uh)

@test el2/ul2 < 1.e-8
@test eh1/uh1 < 1.e-7

end # module


nbits(U) = round(Int, log2(size(U,1)))

@generated function bits(n, ::Type{Val{N}}) where N
  :($([:(Bool((n >> $(i-1)) & 0x01)) for i = N:-1:1]...),)
end

@generated function int(bits::NTuple{N}) where N
  reduce((x, b) -> :(($x<<1)+$b), [:(bits[$i]) for i = 1:N]; init = 0)
end

int(n::Bool) = int((n,))

function classical(f, ::Type{Val{N}}) where N
  U = zeros(Int, 2^N, 2^N)
  for i = 0:2^N-1
    U[int(f(bits(i, Val{N})...))+1, i+1] = 1
  end
  return U
end

classical(f, N) = classical(f, Val{N})

permutation(is::NTuple{N}) where N =
  classical((x...) -> map(i -> x[i], is), Val{N})

pad(U, l, r) = foldr(⊗, [repeated(I, l)..., U, repeated(I, r)...])

headfirst(n, is...) = permutation((is..., setdiff(1:n, is)...))

function measure(n, m)
  M = zeros(Int, 2^n, 2^(n+m))
  for i = 0:2^n-1, j = 0:2^m-1
    mn = i<<m+j
    M[i+1, mn+1] = 1
  end
  return M
end


constants=[:me≖9.1094e-31,:qe≖1.6022e-19,:kB≖1.3807e-23,:ε0≖8.8542e-12]

# Solar wind
ne=:ne≖10.0^7
Te=:Te≖10.0*:qe
variables=[ne,Te]
# Debye length lD
lD=:λD≖sqrt(:ε0*:Te/(:ne*:qe^2)) #valid when Te>>Ti
# plasma frequency wp
wp=:ωp≖sqrt(:ne*:qe^2/(:ε0*:me))
# particles in a debye cube ND
ND=:ND≖:ne*:λD^3

@test (lD&variables&constants).rhs ≈ 7.433892903446526
@test (wp&lD&variables&constants).rhs ≈ 178400.95597166813
@test (ND&lD&variables&constants).rhs ≈ 4.108174668936227e9

# Drifts
ve=:ve≖Cross(:E⊥,:B)/norm(:B)^2
ex=[1,0,0];ey=[0,1,0];ez=[0,0,1]
B=:B≖:Bt*Vec(ez)
Bt=:Bt≖1.5
r=:r≖1
dBt=:dBt≖0.3
Eperp=:E⊥≖-:r/2*:dBt*Vec(ey)
vars=[B,Bt,Eperp,r,dBt]

@test (ve&vars).rhs[1][1] ≈ -0.1
@test (ve&vars).rhs[1][2]==Vec([1,0,0])

# Temperatures
equ=Equation
Ekin=equ(:Ekin,0.5*:Tpa+:Tpe)
parfac=1
Tpa1=equ(:Tpa1,parfac*:T0) #is parallel temperature affected?
Tpe1=equ(:Tpe1,2*:T0)
Ek1=equ(:Ek1,0.5*:Tpa1+:Tpe1)
Ek2=equ(:Ek1,0.5*:T2+:T2)
Ek1s=Ek1&Tpa1&Tpe1 #s=solved
T2=(Ek2&Ek1s)/1.5
T2=equ(T2.rhs,T2.lhs)
Tpa3=equ(:Tpa3,1/parfac*:T2) #is it?
Tpe3=equ(:Tpe3,:T2/2)
Ek3=equ(:Ek3,0.5*:Tpa3+:Tpe3)
Ekf=[equ(:Ekf,0.5*:Tf+:Tf),equ(:Ekf,:Ek3)]
Tf=(Ekf[1]&Ekf[2]&Ek3&Tpa3&Tpe3&T2)/1.5

@test T2.rhs[1][1] ≈ 1.6666666666666665
@test Tf.lhs[1][1] ≈ 1.111111111111111


module Day01

using AdventOfCode2019

function day01(input::String = readInput(joinpath(@__DIR__, "..", "data", "day01.txt")))
    masses = parse.(Int, split(input))
    fuel = masses .÷ 3 .- 2

    solution2 = 0
    for elem in fuel
        while elem > 0
            solution2 += elem
            elem = elem ÷ 3 - 2
        end
    end

    return [fuel |> sum, solution2]
end

end # module


using SkewDist, Distributions, Base.Test

println("Testing MvSkewNormal...")

function cov2cor(X::Matrix{Float64})
    A = similar(X)
    d = size(X,1)
    ωinv = 1.0./sqrt(diag(X))
    for i in 1:d, j in 1:d
        A[i,j] = X[i,j]*ωinv[i]*ωinv[j]
    end
    A
end

srand(1)
N = 100000 # Sample size

ξ = [0.0, -0.1]
α = [1.0, 0.5]
Ω = rand_pd(2)

dist = MvSkewNormal(ξ, Ω, α)
rand(dist)

Z = rand(dist, N)

@test_approx_eq_eps mean(Z, 2) mean(dist) 1e-2
@test_approx_eq_eps var(Z, 2) var(dist) 1e-2
@test_approx_eq_eps Base.covm(Z, mean(Z, 2), 2) cov(dist) 1e-2


using FactorizationMachines
using Base.Test

include("test_read_libsvm.jl")
include("test_common.jl")
include("test_reco.jl")

include("matrix_factorization/test_identity.jl")
include("matrix_factorization/test_rank1.jl")