text:
##
# Ultraspherical
##

struct UltrasphericalWeight{T,Λ} <: AbstractJacobiWeight{T}
    λ::Λ
end

UltrasphericalWeight{T}(λ) where T = UltrasphericalWeight{T,typeof(λ)}(λ)
UltrasphericalWeight(λ) = UltrasphericalWeight{typeof(λ),typeof(λ)}(λ)

==(a::UltrasphericalWeight, b::UltrasphericalWeight) = a.λ == b.λ

function getindex(w::UltrasphericalWeight, x::Number)
    x ∈ axes(w,1) || throw(BoundsError())
    (1-x^2)^(w.λ-one(w.λ)/2)
end

struct Ultraspherical{T,Λ} <: AbstractJacobi{T}
    λ::Λ
end
Ultraspherical{T}(λ::Λ) where {T,Λ} = Ultraspherical{T,Λ}(λ)
Ultraspherical(λ::Λ) where Λ = Ultraspherical{Float64,Λ}(λ)
Ultraspherical(P::Legendre{T}) where T = Ultraspherical(one(T)/2)
function Ultraspherical(P::Jacobi{T}) where T
    P.a == P.b || throw(ArgumentError("$P is not ultraspherical"))
    Ultraspherical(P.a+one(T)/2)
end

Ultraspherical(::ChebyshevU{T}) where T = Ultraspherical{T}(1)

const WeightedUltraspherical{T} = WeightedBasis{T,<:UltrasphericalWeight,<:Ultraspherical}

WeightedUltraspherical(λ) = UltrasphericalWeight(λ) .* Ultraspherical(λ)
WeightedUltraspherical{T}(λ) where T = UltrasphericalWeight{T}(λ) .* Ultraspherical{T}(λ)

==(a::Ultraspherical, b::Ultraspherical) = a.λ == b.λ
==(::Ultraspherical, ::ChebyshevT) = false
==(::ChebyshevT, ::Ultraspherical) = false
==(C::Ultraspherical, ::ChebyshevU) = isone(C.λ)
==(::ChebyshevU, C::Ultraspherical) = isone(C.λ)

###
# interrelationships
###

Jacobi(C::Ultraspherical{T}) where T = Jacobi(C.λ-one(T)/2,C.λ-one(T)/2)

########
# Jacobi Matrix
########

function jacobimatrix(P::Ultraspherical{T}) where T
    λ = P.λ
    _BandedMatrix(Vcat((((2λ-1):∞) ./ (2 .*((zero(T):∞) .+ λ)))',
                        Zeros{T}(1,∞),
                        ((one(T):∞) ./ (2 .*((zero(T):∞) .+ λ)))'), ∞, 1, 1)
end

# These return vectors A[k], B[k], C[k] are from DLMF. Cause of MikaelSlevinsky we need an extra entry in C ... for now.
function recurrencecoefficients(C::Ultraspherical)
    λ = C.λ
    n = 0:∞
    (2(n .+ λ) ./ (n .+ 1), Zeros{typeof(λ)}(∞), (n .+ (2λ-1)) ./ (n .+ 1))
end

##########
# Derivatives
##########

# Ultraspherical(1)\(D*Chebyshev())
@simplify *(D::Derivative{<:Any,<:ChebyshevInterval}, S::ChebyshevU) = D * Ultraspherical(S)

# Ultraspherical(1/2)\(D*Legendre())
@simplify function *(D::Derivative{<:Any,<:ChebyshevInterval}, S::Legendre)
    T = promote_type(eltype(D),eltype(S))
    A = _BandedMatrix(Ones{T}(1,∞), ∞, -1,1)
    ApplyQuasiMatrix(*, Ultraspherical{T}(3/2), A)
end

# Ultraspherical(λ+1)\(D*Ultraspherical(λ))
@simplify function *(D::Derivative{<:Any,<:ChebyshevInterval}, S::Ultraspherical)
    A = _BandedMatrix(Fill(2S.λ,1,∞), ∞, -1,1)
    ApplyQuasiMatrix(*, Ultraspherical{eltype(S)}(S.λ+1), A)
end

# Ultraspherical(λ-1)\ (D*w*Ultraspherical(λ))
@simplify function *(D::Derivative{<:Any,<:AbstractInterval}, WS::WeightedUltraspherical)
    w,S = WS.args
    λ = S.λ
    T = eltype(WS)
    if iszero(w.λ)
        D*S
    elseif w.λ == λ == 1
        A = _BandedMatrix((-(1:∞))', ∞, 1,-1)
        ApplyQuasiMatrix(*, ChebyshevTWeight{T}() .* ChebyshevT{T}(), A)
    elseif w.λ == λ
        n = (0:∞)
        A = _BandedMatrix((-one(T)/(2*(λ-1)) * ((n.+1) .* (n .+ (2λ-1))))', ∞, 1,-1)
        ApplyQuasiMatrix(*, WeightedUltraspherical{T}(λ-1), A)
    else
        error("Not implemented")
    end
end

##########
# Conversion
##########

\(A::Ultraspherical, B::Legendre) = A\Ultraspherical(B)
\(A::Legendre, B::Ultraspherical) = Ultraspherical(A)\B

function \(A::Ultraspherical, B::Jacobi)
    Ã = Jacobi(A)
    Diagonal(Ã[1,:]./A[1,:]) * (Ã\B)
end
function \(A::Jacobi, B::Ultraspherical)
    B̃ = Jacobi(B)
    (A\B̃)*Diagonal(B[1,:]./B̃[1,:])
end

function \(U::Ultraspherical{<:Any,<:Integer}, C::ChebyshevT)
    T = promote_type(eltype(U), eltype(C))
    (U\Ultraspherical{T}(1)) * (ChebyshevU{T}()\C)
end

function \(U::Ultraspherical{<:Any,<:Integer}, C::ChebyshevU)
    T = promote_type(eltype(U), eltype(C))
    U\Ultraspherical(C)
end

function \(C2::Ultraspherical{<:Any,<:Integer}, C1::Ultraspherical{<:Any,<:Integer})
    λ = C1.λ
    T = promote_type(eltype(C2), eltype(C1))
    if C2.λ == λ+1
        _BandedMatrix( Vcat(-(λ ./ ((0:∞) .+ λ))', Zeros(1,∞), (λ ./ ((0:∞) .+ λ))'), ∞, 0, 2)
    elseif C2.λ == λ
        Eye{T}(∞)
    elseif C2.λ > λ
        (C2 \ Ultraspherical(λ+1)) * (Ultraspherical(λ+1)\C1)
    else
        error("Not implemented")
    end
end

function \(C2::Ultraspherical, C1::Ultraspherical)
    λ = C1.λ
    T = promote_type(eltype(C2), eltype(C1))
    if C2.λ == λ+1
        _BandedMatrix( Vcat(-(λ ./ ((0:∞) .+ λ))', Zeros(1,∞), (λ ./ ((0:∞) .+ λ))'), ∞, 0, 2)
    elseif C2.λ == λ
        Eye{T}(∞)
    else
        error("Not implemented")
    end
end

function \(w_A::WeightedUltraspherical, w_B::WeightedUltraspherical)
    wA,A = w_A.args
    wB,B = w_B.args

    if wA == wB
        A \ B
    elseif B.λ == A.λ+1 && wB.λ == wA.λ+1
        λ = A.λ
        _BandedMatrix(Vcat(((2λ:∞) .* ((2λ+1):∞) ./ (4λ .* (λ+1:∞)))',
                            Zeros(1,∞),
                            (-(1:∞) .* (2:∞) ./ (4λ .* (λ+1:∞)))'), ∞, 2,0)
    else
        error("not implemented for $A and $wB")
    end
end

\(A::Legendre, wB::WeightedUltraspherical) = Ultraspherical(A) \ wB

function \(A::Ultraspherical, w_B::WeightedUltraspherical)
    (UltrasphericalWeight(zero(A.λ)) .* A) \ w_B
end


text:
using ConceptDrift
using Test
using Statistics
using Distributions

@testset "ConceptDrift.jl" begin
    # Write your tests here.
end

@testset "gaussian_estimator.jl" begin
    g = ConceptDrift.GaussianEstimator()
    @test g.weight_sum == 0.0
    @test g.mean == 0.0
    @test g.variance_sum == 0.0
    test_stream_A = [1.0, 1.0, 2.0]
    for (i, value) in enumerate(test_stream_A)
        g = ConceptDrift.add_weight(g, value, 1.0)
        @test g.weight_sum == i
        @test g.mean == mean(test_stream_A[1:i])
        if i > 1
            @test isapprox(ConceptDrift.variance(g), var(test_stream_A[1:i]))
            @test isapprox(ConceptDrift.stdev(g), std(test_stream_A[1:i]))
            @test isapprox(ConceptDrift.pdf(g, 0.1), pdf(Normal(g.mean, ConceptDrift.stdev(g)), 0.1))
        end
    end
    g = ConceptDrift.GaussianEstimator()
    @test g.weight_sum == 0.0
    @test g.mean == 0.0
    @test g.variance_sum == 0.0
    test_stream_unweighted = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 3.0, 3.0, 4.0, 4.0, 1.0, 1.0]
    test_stream_weighted = [1.0, 1.0, 1.0, 3.0, 4.0, 1.0]
    for (i, value) in enumerate(test_stream_weighted)
        g = ConceptDrift.add_weight(g, value, 2.0)
        @test g.weight_sum == i*2
        @test g.mean == mean(test_stream_unweighted[1:i*2])
        if i > 1
            @test isapprox(ConceptDrift.variance(g), var(test_stream_unweighted[1:i*2]))
            @test isapprox(ConceptDrift.stdev(g), std(test_stream_unweighted[1:i*2]))
            @test isapprox(ConceptDrift.pdf(g, 0.1), pdf(Normal(g.mean, ConceptDrift.stdev(g)), 0.1))
        end
    end
    g = ConceptDrift.GaussianEstimator()
    @test g.weight_sum == 0.0
    @test g.mean == 0.0
    @test g.variance_sum == 0.0
    test_stream_neg = [-1.0, -1.0, -1.0, 1.0, 1.0, -1.0, 3.0, 3.0, -4.0, 4.0, -1.0, 1.0]
    for (i, value) in enumerate(test_stream_neg)
        g = ConceptDrift.add_weight(g, value, 1.0)
        @test g.weight_sum == i
        @test isapprox(g.mean, mean(test_stream_neg[1:i]))
        if i > 1
            @test isapprox(ConceptDrift.variance(g), var(test_stream_neg[1:i]))
            @test isapprox(ConceptDrift.stdev(g), std(test_stream_neg[1:i]))
            @test isapprox(ConceptDrift.pdf(g, 0.1), pdf(Normal(g.mean, ConceptDrift.stdev(g)), 0.1))
        end
    end

end

@testset "naive_bayes.jl" begin
    @test ConceptDrift.argmax_votes(Dict(1=>0.9, 0=>0.1)) == 1
    @test ConceptDrift.argmax_votes(Dict(1=>0.5, 0=>0.1, 2=>0.4)) == 1
    @test ConceptDrift.argmax_votes(Dict(0=>0.5, 1=>0.1, 2=>0.4)) == 0
    @test ConceptDrift.argmax_votes(Dict(0=>-0.5, 1=>0.1, 2=>0.4)) == 2

    c = ConceptDrift.NaiveBayes()
    ConceptDrift.partial_fit!(c, [0.5, 0.1], 0)
    @test c.observed_class_distribution[0] == 1
    @test collect(keys(c.observed_class_distribution)) == [0]
    attr1_estimator = c.attribute_observers[1].distribution_per_class[0]
    @test typeof(attr1_estimator) == ConceptDrift.GaussianEstimator
    @test isapprox(attr1_estimator.weight_sum, 1.0)
    @test isapprox(attr1_estimator.mean, 0.5)
    @test isapprox(attr1_estimator.variance_sum, 0.0)
    attr2_estimator = c.attribute_observers[2].distribution_per_class[0]
    @test typeof(attr2_estimator) == ConceptDrift.GaussianEstimator
    @test isapprox(attr2_estimator.weight_sum, 1.0)
    @test isapprox(attr2_estimator.mean, 0.1)
    @test isapprox(attr2_estimator.variance_sum, 0.0)

    ConceptDrift.partial_fit!(c, [1.0, 0.0], 0)
    @test c.observed_class_distribution[0] == 2
    @test collect(keys(c.observed_class_distribution)) == [0]
    attr1_estimator = c.attribute_observers[1].distribution_per_class[0]
    @test typeof(attr1_estimator) == ConceptDrift.GaussianEstimator
    @test isapprox(attr1_estimator.weight_sum, 2.0)
    @test isapprox(attr1_estimator.mean, 0.75)
    @test isapprox(attr1_estimator.variance_sum, 0.125)
    attr2_estimator = c.attribute_observers[2].distribution_per_class[0]
    @test typeof(attr2_estimator) == ConceptDrift.GaussianEstimator
    @test isapprox(attr2_estimator.weight_sum, 2.0)
    @test isapprox(attr2_estimator.mean, 0.05)
    @test isapprox(attr2_estimator.variance_sum, 0.005)

    ConceptDrift.partial_fit!(c, [1.0, 0.0], 1)
    @test c.observed_class_distribution[0] == 2
    @test c.observed_class_distribution[1] == 1
    @test collect(keys(c.observed_class_distribution)) == [0, 1]
    attr1_estimator = c.attribute_observers[1].distribution_per_class[0]
    @test typeof(attr1_estimator) == ConceptDrift.GaussianEstimator
    @test isapprox(attr1_estimator.weight_sum, 2.0)
    @test isapprox(attr1_estimator.mean, 0.75)
    @test isapprox(attr1_estimator.variance_sum, 0.125)
    attr2_estimator = c.attribute_observers[2].distribution_per_class[0]
    @test typeof(attr2_estimator) == ConceptDrift.GaussianEstimator
    @test isapprox(attr2_estimator.weight_sum, 2.0)
    @test isapprox(attr2_estimator.mean, 0.05)
    @test isapprox(attr2_estimator.variance_sum, 0.005)
    attr1_estimator = c.attribute_observers[1].distribution_per_class[1]
    @test typeof(attr1_estimator) == ConceptDrift.GaussianEstimator
    @test isapprox(attr1_estimator.weight_sum, 1.0)
    @test isapprox(attr1_estimator.mean, 1.0)
    @test isapprox(attr1_estimator.variance_sum, 0.0)
    attr2_estimator = c.attribute_observers[2].distribution_per_class[1]
    @test typeof(attr2_estimator) == ConceptDrift.GaussianEstimator
    @test isapprox(attr2_estimator.weight_sum, 1.0)
    @test isapprox(attr2_estimator.mean, 0.0)
    @test isapprox(attr2_estimator.variance_sum, 0.0)

    c = ConceptDrift.NaiveBayes()
    ConceptDrift.partial_fit!(c, [0.5, 0.1], 0)
    @test ConceptDrift.predict_proba(c, [0.5, 0.1]) == Dict(0=>1)
    @test ConceptDrift.predict(c, [0.5, 0.1]) == 0
    ConceptDrift.partial_fit!(c, [1.0, 0.0], 0)
    @test ConceptDrift.predict_proba(c, [0.5, 0.1]) == Dict(0=>3.8612941052021554)
    @test ConceptDrift.predict(c, [0.5, 0.1]) == 0
    ConceptDrift.partial_fit!(c, [0.1, 1.0], 1)
    @test ConceptDrift.predict_proba(c, [0.5, 0.1]) == Dict(0=>2.57419607013477,1=>0.0)
    @test ConceptDrift.predict(c, [0.5, 0.1]) == 0
    @test ConceptDrift.predict_proba(c, [0.1, 1.0]) == Dict(0=>4.9975160518252554e-40,1=>0.3333333333333333)
    @test ConceptDrift.predict(c, [0.1, 1.0]) == 1
end

@testset "data\\random_tree_generator.kl" begin
    @test ConceptDrift.get_cat_value_from_onehot([0, 0, 1], 1, 0, 3) == 3
    @test ConceptDrift.get_cat_value_from_onehot([0, 0, 1], 1, 0, 3) == 3
    @test ConceptDrift.get_cat_value_from_onehot([0, 1, 0], 1, 0, 3) == 2
    @test ConceptDrift.get_cat_value_from_onehot([1, 0, 0], 1, 0, 3) == 1
    @test ConceptDrift.get_cat_value_from_onehot([1, 0, 0, 0, 0, 1], 2, 0, 3) == 3
    @test ConceptDrift.get_cat_value_from_onehot([1, 0, 0, 0, 1, 0], 2, 0, 3) == 2
    @test ConceptDrift.get_cat_value_from_onehot([1, 0, 0, 1, 0, 0], 2, 0, 3) == 1
    @test ConceptDrift.get_cat_value_from_onehot([3, 1, 0, 0, 0, 0, 1], 3, 1, 3) == 3
    @test ConceptDrift.get_cat_value_from_onehot([3, 1, 0, 0, 0, 1, 0], 3, 1, 3) == 2
    @test ConceptDrift.get_cat_value_from_onehot([3, 1, 0, 0, 1, 0, 0], 3, 1, 3) == 1
end

ext: jl


text:
# backward
export backward!

function backward!(l::lossLayer)
  if l.lock==false
    l.∂rᵢ_L=l.∂Cost(l.rₒ,l.rᵢ)
  end
end

function backward!(p::predictLayer)
  if p.lock==false
    p.∂rₒ_L=p.top[1].∂rᵢ_L[p.topRange[1],:]
    for i=2:length(p.top)
      p.∂rₒ_L+=p.top[i].∂rᵢ_L[p.topRange[i],:]
    end
    p.∂rᵢ_L=reshape(p.∂Predict(p.rᵢ,p.rₒ) ⊗ p.∂rₒ_L,size(p.∂rᵢ_L)...)
  end
end
function backward!(h::hiddenLayer)
  if h.lock==false
    h.∂rₒ_L=h.top[1].∂rᵢ_L[h.topRange[1],:]
    for i=2:length(h.top)
      h.∂rₒ_L+=h.top[i].∂rᵢ_L[h.topRange[i],:]
    end
    h.∂i_L=h.∂rₒ_L.*h.∂A(h.I)
    h.∂rᵢ_L=h.W'*h.∂i_L
    h.∂w_L=reshape(h.∂i_L⊗h.rᵢ,size(h.∂w_L)...)
    h.∂b_L=h.∂i_L
  end
end
function backward!(d::dataLayer)
end
function backward!(v::voidLayer)
end
function backward!(net::Array{Layer})
  for idx=order[end:-1:1]
    backward!(net[idx])
  end
end


text:
tempSum = 0
numOfCells = 0

function countAvgTemp(r::Robot)
    dir = Ost
    move_while_can!(r, dir)
    dir = inverse(dir)
    while !isborder(r,Nord)
      move!(r, Nord)
      move_while_can!(r, dir)
      dir = inverse(dir)
    end
    if numOfCells > 0
        return tempSum/numOfCells
    else
        return 0
    end
end

function move_while_can!(r::Robot, side)
    checkCell(r)
    while !isborder(r, side)
        move!(r, side)
        checkCell(r)
    end
end

function checkCell(r::Robot)
   global tempSum, numOfCells
   if ismarker(r)
      tempSum += temperature(r)
      numOfCells += 1
   end
end

function inverse(side)
    return HorizonSide(mod(Int(side) + 2, 4))
end


text:
### 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

# ╔═╡ 9112b09c-5a21-11ec-0847-1971b3c9a6cd
begin
	using Pkg
	Pkg.activate("../../")
	using Revise
end

# ╔═╡ 4b66d9ba-eb36-4f53-8e87-28b6eadbd828
using FFTW, Colors, ImageShow, TestImages, Tullio, Zygote, Noise, IndexFunArrays, FourierTools, ImgProcMic, PlutoUI, DeconvOptim, PlutoTest, Optim, Statistics, Plots

# ╔═╡ 9bc5b039-5c13-46bc-8d3b-9edae93d164a
using ComponentArrays, Images, Random, MicroscopyTools

# ╔═╡ 606ad4de-213f-4aac-a0a9-9a6827c05c29
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

# ╔═╡ e7a503cf-c59c-4737-95f5-b8f38548833a
"""
	calc_psf(arr_size, radius)

Calculate a simple PSF. `arr_size` is the output size and `radius` is the radius
in Fourier space.
The output array is normalized by its sum to 1.
"""
function calc_psf(arr_size, radius)
	# TODO
	rr2 = circ(arr_size, radius)
	h = abs2.(ift(rr2))
	h ./= sum(h)
	return h
end

# ╔═╡ 3194824d-ed15-479a-af10-a1803f9f20bc
md"# 1 Lucy Richardson Deconvolution
In the last exercise we implemented the Wiener filter which is suited for Gaussian noise.
However, in microscopy the most dominant noise is usually Poisson noise.
The LR deconvolution algorithm is derived on a Poisson noise assumption.

Based on the lecture notes, implement the LR algorithm.
"

# ╔═╡ 7ef1aaf2-081c-4885-9ab3-4606cf0656b4
"""
	LR(measured, psf, N)

`measured` is the measured image. `psf` is the point spread function.
`N` is the number of iterations.
"""
function LR(measured, psf, N)
	otf = rfft(psf)
	otf_c = conj.(otf)

	rec = one.(measured)

	conv(object, o) = irfft(rfft(object) .* o, size(measured, 1))

	for i = 1:N
		rec .*= conv(measured ./ conv(rec, otf_c), otf)
	end
	return rec
end

# ╔═╡ 400aeebf-21bd-4b58-9b64-ccec3071df63
# testimage
img = Float32.(testimage("resolution_test_512"));

# ╔═╡ d7741c07-0982-4e57-bf95-a3ee0bc6f0b1
# calculate the PSF with a radius of 30
psf = ifftshift(Float32.(calc_psf(size(img), 30)));

# ╔═╡ aa15a3ed-0f97-4ef9-9704-0ede98426b6c
# our measured image, degraded with poisson
measured = 0.0001f0 .+ poisson(FourierTools.conv(img, psf), 100);

# ╔═╡ f4e2cc11-e86e-4afc-a021-aed27283a5f1
gray_show(measured)

# ╔═╡ dc9fa103-6a00-4129-b5d1-c1b36216d13d
md"
After

iterations=$(@bind iterations Slider(1:300, show_value=true))

the image should look _better_ than the original.
Take that as test!
"

# ╔═╡ 6ba1619b-036a-480d-b3bb-88ea6f5b259e
rec = LR(measured, psf, iterations);

# ╔═╡ f9bb5a74-067d-4afd-ae61-7f41f2e1dbca
gray_show(rec)

# ╔═╡ 0d7e91be-e03e-454c-9ebd-c0b68ca93848
md"# 2 (Smoothed) Total Variation Regularizer
In this part we want to implement a Total Variation Regularizer.
It's purpose is to impose a constraint of _piecewise smoothness_ to the image.
Formally, it is defined as

$\text{TV}(\text{img}) = \sum_{i,j} \sqrt{\epsilon + \left|\text{img}[i, j] - \text{img}[i+1, j]\right|^2 + \left|\text{img}[i, j] - \text{img}[i, j+1]\right|^2}$

where $\epsilon$ is a small value prevent a singularity in the derivative.

The output of this equation is a scalar value and can be understood as a loss function.
If we add this term to our loss function of the LR algorithm, the optimizer has to minimize not only the LR loss term but also this additional regularizer.
Often this can introduce nicer results.

As you learnt in the last homework, we also need the gradient of the `TV` function.
Fortunately, for that purpose we can use the automatic differentiation package [Zygote.jl](https://github.com/FluxML/Zygote.jl).
"

# ╔═╡ 72cdab83-0162-4951-8751-b91cb6fd0e49
md"For the implementation of the regularizer itself, we use the package [Tullio.jl](https://github.com/mcabbott/Tullio.jl). _Tullio.jl_ speeds up the calculation and is conventient to use.
For simple usage, see below:"

# ╔═╡ f8c665e4-e07f-4995-ad84-5a31658be5f4
begin
	# example how @tullio works
	arr = [1.0, 2.0, 3.0]
	@tullio arr_ex = abs2(arr[i] - arr[i-1])
end

# ╔═╡ bfb88c5e-0e45-428b-a66c-b9c37f633c6b
"""
	TV(img::AbstractArray{T, N}, ϵ=T(1e-6)) where {T, N}

Calculates the TV regularizer. `ϵ` is a small positive value to make it differentiable.
"""
function TV(img::AbstractArray{T, N}, ϵ=T(1e-6)) where {T, N}
	# insert tullio expression here!
	@tullio r = sqrt(ϵ + abs2(img[i, j] - img[i, j+1]) + abs2(img[i, j] - img[i + 1, j]))
end

# ╔═╡ 49143f59-2f8f-46e9-ac32-c4681164ef82
TV(img)

# ╔═╡ e33d9b1f-dc9b-49c0-a6e3-b1e2266df999
# the gradient can be calculated likes this!
# it describes how the value of the TV change for each pixel value
gradient(TV, img)[1]

# ╔═╡ 8a9da6e2-528a-41e5-9bc5-750885d94ca0
md"## 2 Test"

# ╔═╡ 53364e6f-b672-4c60-8fe1-845ba85399ba
PlutoTest.@test typeof(TV([1.0 2; 3 4])) <: Number

# ╔═╡ 30e313ee-aca2-4469-97cf-08acd489aaef
PlutoTest.@test TV([1.0 2; 3 4]) ≈ 2.2360682011065762

# ╔═╡ 4c91d0da-201a-4281-91a2-2de626ab4a30
PlutoTest.@test Zygote.gradient(TV, [1.0 2; 3 4])[1] ≈ [-1.34164   0.447214; 0.894427  0.0] rtol=1e-5

# ╔═╡ 340851e8-a66d-4643-83d1-48fb775dde8e
md"# 3 Lucy Richardson with TV
Additionally to the plain LR we want to incorporate the TV regularizer in this task.

The gradient of `TV` can be calculated with `gradient(TV, arr)[1]`.

In the lectures it is sketched how to add the regularizer loss to the normal loss function.
Copy your `LR` from above and try to incorporate the gradient of the `TV` loss value.
The `TV` gradient value is weighted with `λ`.
"

# ╔═╡ cfc01fbb-5ee4-4909-ba67-2034506e55e8
function LR_TV(measured, psf, λ, N)
	otf = rfft(psf)
	otf_c = conj.(otf)

	buffer = rfft(img)
	rec = one.(measured)

	conv(object, o) = irfft(rfft(object) .* o, size(measured, 1))

	for i = 1:N
		rec .*= (conv(measured ./ conv(rec, otf_c), otf) .- λ .* gradient(TV, rec)[1])
	end
	return rec
end

# ╔═╡ ae7c45b6-c8a7-4fe0-b922-9ab40fcce372
# that can easily take a minute. Zygote is very sloooooooow on the first call
rec2 = LR_TV(measured, psf, 0.02, 400)

# ╔═╡ 8c91142a-5b2d-4dfc-9764-87a66f90f0cb
gray_show(rec2)

# ╔═╡ dcba043b-db09-4adc-8239-a6439a9c074e
md" ## 4 Bonus
Anyone who can beat [DeconvOptim.jl](https://github.com/roflmaostc/DeconvOptim.jl) in terms of speed (with equal reconstruction quality), immediatly gets the bonus for the exam.
"

# ╔═╡ bbb7dda6-84c2-4172-a888-85a3658c1556
reg = DeconvOptim.TV(num_dims=2);

# ╔═╡ 231698fd-bb16-4895-9278-0d41dfafdb7e
rec_DO,r = deconvolution(measured, psf, regularizer=reg, iterations=20, λ=0.02);

# ╔═╡ 45f089ec-a7cb-4941-ad04-7a6c6c55dbef
r

# ╔═╡ d725c173-0123-40ed-aaa0-feec1dcd7a25
gray_show(rec_DO)

# ╔═╡ dfc9d8f3-6b16-476a-bba9-1c7f55229f57
md"# 5 Single Molecule Localization

In this part we fit the position, intensity and the standard deviation of a point source emitter to a measured image.
"

# ╔═╡ 954f06a6-ac27-4596-a978-e1274b2532d8
md"### Hint
Use the following pattern:
"

# ╔═╡ 46185370-aaf3-48c3-b297-b3a2843a0435
1:100

# ╔═╡ 7c51a261-a188-469d-8020-876b1b1dcbf0
(1:100)'

# ╔═╡ fae6889b-65f5-4b90-972f-56187b69fee5
"""
	gauss(I, σ, x_offset, y_offset)

Returns an image of size `(256, 256)` with a Gauss peak
centered at `(x_offset, y_offset)` with standard deviation `σ` and intensity
`I` (which is multiplied to the ideal Gauss)
"""
function gauss(I, σ::T, x_offset::T, y_offset::T) where T
	# TODO
	s = (256, 256)
	x = 1:s[1]
	y = (1:s[2])'
	@. I * exp(-T(0.5) * (abs2(x - x_offset) + abs2(y-y_offset)) / σ^2)
end

# ╔═╡ 23b9ea69-99d5-4a49-8e36-cd25323c8858
md"## 5.1 Test Gauss"

# ╔═╡ e9c17987-029e-4bd9-824e-47175f4d6924
z = zeros((256, 256));

# ╔═╡ 046ec177-9394-4b7d-bead-68590e169f1b
z[10, 10] = 42.0;

# ╔═╡ eb9b2246-c9d9-442e-bde4-48ed14bb3fa6
PlutoTest.@test gauss(42.0, 0.000001, 10.0, 10.0) ≈ z

# ╔═╡ 9cdf2e65-af36-4421-a15b-fdc4c87569b4
PlutoTest.@test gauss(11.0, 100000000000000000000.0, 0.0, 0.0) ≈ 11.0 .* ones((256, 256))

# ╔═╡ 7ecb1d8d-334c-4043-9378-adcbe0b0bf35
PlutoTest.@test (gauss(10.0, 2.0, 3.0, 3.0))[1:5, 1:5] ≈ [3.6787944117144233 5.352614285189903 6.065306597126334 5.352614285189903 3.6787944117144233; 5.352614285189903 7.788007830714049 8.824969025845954 7.788007830714049 5.352614285189903; 6.065306597126334 8.824969025845954 10.0 8.824969025845954 6.065306597126334; 5.352614285189903 7.788007830714049 8.824969025845954 7.788007830714049 5.352614285189903; 3.6787944117144233 5.352614285189903 6.065306597126334 5.352614285189903 3.6787944117144233]

# ╔═╡ 9b5b3495-0d3b-4f7c-9f51-04245ebd59c2
md"## 5.2 Forward Model
In the next part, we assemble a forward model mapping a set of parameters to an image.
"

# ╔═╡ 26950064-0127-41fd-9973-6c170bd6b8fe
"""
	draw_parameters(N)

Randomly draws `N` parameters for the Gauss peak.
"""
function draw_parameters(N)
	# make the random numbers reproducible
	Random.seed!(40)
	σ = [rand(1:0.01:5.0) for i = 1:N]
	x_offset = [rand(10:0.1:240) for i = 1:N]
	y_offset = [rand(10:0.1:240) for i = 1:N]
	I = [rand(0.5:0.1:1.0) for i = 1:N]
	return ComponentVector(σ=σ, x_offset=x_offset, y_offset=y_offset, I=I)
end

# ╔═╡ ef9c0737-ffbc-498b-80d2-2e93558c508f
# draws randomly a few parameters
# the output is a ComponentVector which is needed for Optim
params = draw_parameters(5);

# ╔═╡ 70a9e03f-d44f-4cce-a82d-898e63f789fd
# all fields can be access with params.FIELD, see below

# ╔═╡ 70453553-7dbc-46f1-a141-f4d19f2ae3d4
params.σ

# ╔═╡ 6ac55320-b94f-44c2-9669-1d4e7eb3ddee
params.x_offset

# ╔═╡ 73416e0d-ed2b-4b95-aa43-ce58de7f29aa
"""
	forward(params)

Maps the `params` to an output image.
"""
function forward(params)
	sum(gauss.(params.I, params.σ, params.x_offset, params.y_offset))
end

# ╔═╡ 56d566e3-632c-4852-839b-adf505122ad2
measurement = randn() * 0.3 .+ poisson(forward(params), 50);

# ╔═╡ 26f1a0c2-3787-4335-958b-12719e2901bf
# noisy measurement of the point sources
gray_show(measurement)

# ╔═╡ 2b3c151a-d7bd-40e4-97d7-f121c9da23a6
md"## 5.3 Loss Function"

# ╔═╡ 32bd26f1-0d04-4589-a2d5-1e01abce4a7d
"""
	loss(measurement, pred)

Compares the prediction `pred` and the `measurement` under the L2 norm.
"""
function loss(measurement, pred)
	# TODO
	sum(abs2.(measurement .- pred))
end

# ╔═╡ dc66401d-2581-4d6f-8b36-d15a2d21391a
md"## 5.4 Reconstruction"

# ╔═╡ daa77619-e7a8-4c6a-9854-d8ef21042277
"""
	find_peaks(arr, bin_factor)

Finds the peaks in the measurement `arr`.
To get a better SND we bin the measurement by a binning factor `bin_factor`.
"""
function find_peaks(arr, bin_factor)
	arr_s = bin(arr, (bin_factor,bin_factor))
	map(x -> x .* bin_factor, Tuple.(findlocalmaxima(arr_s)))
end

# ╔═╡ 1a505ff5-4745-42c5-bc40-bc1075fcd4d6
"""
	init_parameters(N, measurement)

Initializes a good set of parameters for the reconstruction.
`N` is the number of PSFs and `measurement` the measured image.

The critical part is to find good initial positions for the peaks.
For that we use `find_peaks`.

"""
function init_params(N, measurement, bin_factor, σ_guess)
	# find a σ_guess which results in a good estimate
	σ = [σ_guess for i = 1:N]

	peaks = find_peaks(measurement, bin_factor)
	x_offset = [peaks[i][1] for i =1:N]
	y_offset = [peaks[i][2] for i =1:N]

	# you can change that as well! Is it critical?
	I = [1.0 for i = 1:N]

	return ComponentVector(σ=σ, x_offset=x_offset, y_offset=y_offset, I=I)
end

# ╔═╡ 74f8a126-7e5c-4e1d-90bd-fe66065d5faf
begin
	# initial guess
	# try to find an binning factor which results in good peaks
	bin_factor = 5 # TODO
	σ_guess = 5.30 # TODO
	params_guess = init_params(5, measurement, bin_factor, σ_guess);
end

# ╔═╡ 3c38cdec-99d8-4c11-bedb-87137ea119ba
md"
The left image is the measured image.
The center image the initial guess.
The right image is the difference map suited to judge the position of the peaks.
"

# ╔═╡ 0a3cfe9b-888c-420a-9465-bef61d058c17
gray_show([measurement forward(params_guess) abs.(measurement .- forward(params_guess))])

# ╔═╡ 22d4842a-908f-485d-9bf4-96c053b106c6
md"### 5.5 Optim Specific Parts for the Reconstruction"

# ╔═╡ c3996433-e27d-4227-8e18-9531a30ea389
begin
	# the function maps the params to a single scalar value
	function f(params)
		loss(measurement, forward(params))
	end
	# this function calculates the
	g!(G, params) = isnothing(G) ? G : G .= gradient(f, params)[1]
end

# ╔═╡ d30c5e75-8a58-4231-8164-c7d783b9bc02
res = optimize(f, g!, params_guess, LBFGS(), Optim.Options(iterations=50))

# ╔═╡ 90b98a76-1bc7-43e0-927c-1cf94d35ffd0
output_optim = forward(res.minimizer);

# ╔═╡ b5880735-c615-440c-a0a0-56a6a024494a
gray_show([measurement output_optim measurement .- output_optim])

# ╔═╡ efaf6cfb-bf37-4cc2-be30-abbc6706e905
heatmap(measurement .- output_optim, title="Difference map")

# ╔═╡ 19a1ff11-1276-46ab-b84b-a66898522068
md"## Check if the values look plausible"

# ╔═╡ 060ef38d-312b-4487-b59d-0ba288b16daf
md"The operator `|>` in an expression like `f(x) |> g` is the same as `g(f(x))`. It's only a convenvient wrapper"

# ╔═╡ e60c3150-7658-4a3d-819f-b1e6293a27fb
params.σ |> sort

# ╔═╡ 17ed7c8d-4d12-4008-8745-e96f36ce85b1
res.minimizer.σ |> sort

# ╔═╡ e677b9f8-38ee-41cd-b319-ef1cc48b17e2
params.I |> sort

# ╔═╡ 854f732f-441b-45b1-9a00-2b7ae854dc47
res.minimizer.I |> sort

# ╔═╡ 45abec82-b841-4d8a-b2ae-bba738352da3
res.minimizer.y_offset |> sort

# ╔═╡ ff1c6ed2-e4b1-4f6a-87df-44f350bb50f7
params.y_offset |> sort

# ╔═╡ 8ffb693f-f5ca-45a5-8d1b-5c96885118d0
PlutoTest.@test (res.minimizer.x_offset |> sort) ≈ (params.x_offset |> sort) rtol=0.01

# ╔═╡ 67db2159-c2f0-4130-8e82-6b75907db82f
PlutoTest.@test (res.minimizer.y_offset |> sort) ≈ (params.y_offset |> sort) rtol=0.01

# ╔═╡ Cell order:
# ╠═9112b09c-5a21-11ec-0847-1971b3c9a6cd
# ╠═4b66d9ba-eb36-4f53-8e87-28b6eadbd828
# ╟─606ad4de-213f-4aac-a0a9-9a6827c05c29
# ╟─e7a503cf-c59c-4737-95f5-b8f38548833a
# ╟─3194824d-ed15-479a-af10-a1803f9f20bc
# ╠═7ef1aaf2-081c-4885-9ab3-4606cf0656b4
# ╠═400aeebf-21bd-4b58-9b64-ccec3071df63
# ╠═d7741c07-0982-4e57-bf95-a3ee0bc6f0b1
# ╠═aa15a3ed-0f97-4ef9-9704-0ede98426b6c
# ╠═f4e2cc11-e86e-4afc-a021-aed27283a5f1
# ╟─dc9fa103-6a00-4129-b5d1-c1b36216d13d
# ╠═6ba1619b-036a-480d-b3bb-88ea6f5b259e
# ╠═f9bb5a74-067d-4afd-ae61-7f41f2e1dbca
# ╟─0d7e91be-e03e-454c-9ebd-c0b68ca93848
# ╟─72cdab83-0162-4951-8751-b91cb6fd0e49
# ╠═f8c665e4-e07f-4995-ad84-5a31658be5f4
# ╠═bfb88c5e-0e45-428b-a66c-b9c37f633c6b
# ╠═49143f59-2f8f-46e9-ac32-c4681164ef82
# ╠═e33d9b1f-dc9b-49c0-a6e3-b1e2266df999
# ╟─8a9da6e2-528a-41e5-9bc5-750885d94ca0
# ╠═53364e6f-b672-4c60-8fe1-845ba85399ba
# ╠═30e313ee-aca2-4469-97cf-08acd489aaef
# ╠═4c91d0da-201a-4281-91a2-2de626ab4a30
# ╟─340851e8-a66d-4643-83d1-48fb775dde8e
# ╠═cfc01fbb-5ee4-4909-ba67-2034506e55e8
# ╠═ae7c45b6-c8a7-4fe0-b922-9ab40fcce372
# ╠═8c91142a-5b2d-4dfc-9764-87a66f90f0cb
# ╟─dcba043b-db09-4adc-8239-a6439a9c074e
# ╠═bbb7dda6-84c2-4172-a888-85a3658c1556
# ╠═231698fd-bb16-4895-9278-0d41dfafdb7e
# ╠═45f089ec-a7cb-4941-ad04-7a6c6c55dbef
# ╠═d725c173-0123-40ed-aaa0-feec1dcd7a25
# ╟─dfc9d8f3-6b16-476a-bba9-1c7f55229f57
# ╠═9bc5b039-5c13-46bc-8d3b-9edae93d164a
# ╟─954f06a6-ac27-4596-a978-e1274b2532d8
# ╠═46185370-aaf3-48c3-b297-b3a2843a0435
# ╠═7c51a261-a188-469d-8020-876b1b1dcbf0
# ╠═fae6889b-65f5-4b90-972f-56187b69fee5
# ╟─23b9ea69-99d5-4a49-8e36-cd25323c8858
# ╠═e9c17987-029e-4bd9-824e-47175f4d6924
# ╠═046ec177-9394-4b7d-bead-68590e169f1b
# ╠═eb9b2246-c9d9-442e-bde4-48ed14bb3fa6
# ╠═9cdf2e65-af36-4421-a15b-fdc4c87569b4
# ╠═7ecb1d8d-334c-4043-9378-adcbe0b0bf35
# ╟─9b5b3495-0d3b-4f7c-9f51-04245ebd59c2
# ╠═26950064-0127-41fd-9973-6c170bd6b8fe
# ╠═ef9c0737-ffbc-498b-80d2-2e93558c508f
# ╠═70a9e03f-d44f-4cce-a82d-898e63f789fd
# ╠═70453553-7dbc-46f1-a141-f4d19f2ae3d4
# ╠═6ac55320-b94f-44c2-9669-1d4e7eb3ddee
# ╠═73416e0d-ed2b-4b95-aa43-ce58de7f29aa
# ╠═56d566e3-632c-4852-839b-adf505122ad2
# ╠═26f1a0c2-3787-4335-958b-12719e2901bf
# ╟─2b3c151a-d7bd-40e4-97d7-f121c9da23a6
# ╠═32bd26f1-0d04-4589-a2d5-1e01abce4a7d
# ╟─dc66401d-2581-4d6f-8b36-d15a2d21391a
# ╠═daa77619-e7a8-4c6a-9854-d8ef21042277
# ╠═1a505ff5-4745-42c5-bc40-bc1075fcd4d6
# ╠═74f8a126-7e5c-4e1d-90bd-fe66065d5faf
# ╟─3c38cdec-99d8-4c11-bedb-87137ea119ba
# ╠═0a3cfe9b-888c-420a-9465-bef61d058c17
# ╟─22d4842a-908f-485d-9bf4-96c053b106c6
# ╠═c3996433-e27d-4227-8e18-9531a30ea389
# ╠═d30c5e75-8a58-4231-8164-c7d783b9bc02
# ╠═90b98a76-1bc7-43e0-927c-1cf94d35ffd0
# ╠═b5880735-c615-440c-a0a0-56a6a024494a
# ╠═efaf6cfb-bf37-4cc2-be30-abbc6706e905
# ╟─19a1ff11-1276-46ab-b84b-a66898522068
# ╟─060ef38d-312b-4487-b59d-0ba288b16daf
# ╠═e60c3150-7658-4a3d-819f-b1e6293a27fb
# ╠═17ed7c8d-4d12-4008-8745-e96f36ce85b1
# ╠═e677b9f8-38ee-41cd-b319-ef1cc48b17e2
# ╠═854f732f-441b-45b1-9a00-2b7ae854dc47
# ╠═45abec82-b841-4d8a-b2ae-bba738352da3
# ╠═ff1c6ed2-e4b1-4f6a-87df-44f350bb50f7
# ╠═8ffb693f-f5ca-45a5-8d1b-5c96885118d0
# ╠═67db2159-c2f0-4130-8e82-6b75907db82f


text:
# ------------------------------------------------------------------
# Licensed under the ISC License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    LeaveBallOut(ball; loss=Dict())

Leave-`ball`-out (a.k.a. spatial leave-one-out) validation.
Optionally, specify `loss` function from `LossFunctions.jl`
for some of the variables.

    LeaveBallOut(radius; ndims=3, loss=Dict())

By default, use Euclidean ball of given `radius` in space
with `ndims` dimensions.

## References

* Le Rest et al. 2014. [Spatial leave-one-out cross-validation
  for variable selection in the presence of spatial autocorrelation]
  (https://onlinelibrary.wiley.com/doi/full/10.1111/geb.12161)
"""
struct LeaveBallOut{B<:BallNeighborhood} <: AbstractErrorEstimator
  ball::B
  loss::Dict{Symbol,SupervisedLoss}
end

LeaveBallOut(ball::BallNeighborhood; loss=Dict()) =
  LeaveBallOut{typeof(ball)}(ball, loss)

LeaveBallOut(radius::Real; ndims=3, loss=Dict()) =
  LeaveBallOut(BallNeighborhood{ndims}(radius), loss=loss)

function Base.error(solver::AbstractLearningSolver,
                    problem::LearningProblem,
                    eestimator::LeaveBallOut)
  sdata = sourcedata(problem)
  ovars = outputvars(task(problem))
  ball  = eestimator.ball
  loss  = eestimator.loss
  for var in ovars
    if var ∉ keys(loss)
      loss[var] = defaultloss(sdata[1,var])
    end
  end

  # efficient neighborhood search
  searcher = NeighborhoodSearcher(sdata, ball)

  # pre-allocate memory for coordinates
  coords = MVector{ndims(sdata),coordtype(sdata)}(undef)

  solutions = pmap(1:npoints(sdata)) do i
    coordinates!(coords, sdata, i)

    # points inside and outside ball
    inside  = search(coords, searcher)
    outside = [j for j in 1:npoints(sdata) if j ∉ inside]

    # setup and solve learning sub-problem
    subproblem = LearningProblem(view(sdata, outside),
                                 view(sdata, [i]),
                                 task(problem))
    solve(subproblem, solver)
  end

  result = pmap(ovars) do var
    y = [sdata[i,var] for i in 1:npoints(sdata)]
    ŷ = [solutions[i][1,var] for i in 1:npoints(sdata)]
    var => value(loss[var], y, ŷ, AggMode.Mean())
  end

  Dict(result)
end


text:
function alt_return(mdl::Model)

    Ψₙₑᵤ = copy(mdl.Λ' * ζ * mdl.λ)
    Λₙₑᵤ = copy(vec( sum( ((mdl.DR .== 0) .* mdl.Ψ * mdl.λ ), dims = 1 ) + mdl.Λ' * ( 1 - ζ ) * mdl.λ) )
    mdl.Ψ = copy(zeros(size(mdl.Ψ)))
    mdl.Ψ[1,:] = copy(Ψₙₑᵤ)
    mdl.Λ = copy(Λₙₑᵤ)
    return( mdl.Ψ, mdl.Λ)

end


text:
function db(t,n,k)
    out = binomial(n,k).*(k.*t.^(k.-1).*(1 .-t).^(n-k) .- t.^k.*(n-k).*(1 .-t).^(n-k-1));
    return out
end


text:
module XlsxWriter

#=
http://xlsxwriter.readthedocs.io/
=#

import Base.close

using PyCall
using Dates

const xlsxwriter = PyNULL()

function __init__()
	if haskey(ENV, "XLSXWRITER_PATH")
		if xlsxwriter == PyNULL()
			pushfirst!(PyVector(pyimport("sys")["path"]), ENV["XLSXWRITER_PATH"])
			copy!(xlsxwriter, pyimport("xlsxwriter"))
		end
	else
		throw("Need to set ENV[\"XLSXWRITER_PATH\"] to the path where the xlsxwriter folder can be found. For me that is: raw\"C:\\Users\\matt\\repos\\XlsxWriter.py\\\" and reload or call XlsxWriter.__init__() ")
	end
end

export add_worksheet!, add_chartsheet!, add_format!, set_properties!, set_custom_property!, set_calc_mode!, set_column!, set_row!, add_chart!, close, rc2cell, colNtocolA

export add_series!, set_x_axis!, set_y_axis!, set_x2_axis!, set_y2_axis!, combine!, set_size!, set_title!, set_legend!, set_chartarea!, set_plotarea!, set_style!, set_table!, set_up_down_bars!, set_drop_lines!, set_high_low_lines!, set_blanks_as!, show_hidden_data!, set_rotation!, set_hole_size!

export set_font_name!, set_font_size!, set_font_color!, set_bold!, set_italic!, set_underline!, set_font_strikeout!, set_font_script!, set_num_format!, set_locked!, set_hidden!, set_align!, set_center_across!, set_text_wrap!, set_rotation!, set_indent!, set_shrink!, set_text_justlast!, set_pattern!, set_bg_color!, set_fg_color!, set_border!, set_bottom!, set_top!, set_left!, set_right!, set_border_color!, set_bottom_color!, set_top_color!, set_left_color!, set_right_color!, set_diag_border!, set_diag_type!, set_diag_color!

export set_chart!, write!, write_string!, write_blank!, write_formula!, write_datetime!, write_bool!, write_url!, write_number!, write_array_formula!, write_row!, write_column!, write_matrix!, define_name!, worksheets, get_worksheet_by_name, set_first_sheet!, merge_range!, freeze_panes!, split_panes!, Xls, write_matrix!, data_validation!, conditional_format!, add_table!, add_sparkline!, activate!, select!, hide!, set_first_sheet!, t!, protect!, set_zoom!, set_tab_color!, set_landscape!, set_portrait!, set_paper!, set_margins!, set_header!, set_header!, set_footer!, set_footer!, get_name, write_rich_string!, insert_image!, insert_chart!, insert_textbox!, insert_button!, write_comment!, show_comments!, set_comments_author!, autofilter!, filter_column!, filter_column_list!, set_selection!, set_default_row!, outline_settings!, set_vba_name!, add_vba_project!

# type definitions

struct Url
	url::AbstractString
end

struct Workbook
	py::PyObject
	io::Union{AbstractString, IOBuffer}
	Workbook() = Workbook(IOBuffer())
	Workbook(fn::AbstractString, opts::Dict=Dict()) = fn == "" ? Workbook(IOBuffer(), opts) : new(xlsxwriter.Workbook(fn, opts), fn)
	function Workbook(io::IOBuffer, opts::Dict=Dict())
		if !haskey(opts, "in_memory")
			opts["in_memory"] = true
		end
		new(xlsxwriter.Workbook(io, opts), io)
	end
end

struct Worksheet
	py::PyObject
end

struct Format
	py::PyObject
end

struct Chart
	py::PyObject
end

struct Chartsheet
	py::PyObject
end

const Data = Union{Real, AbstractString, DateTime, Date, Bool, Url}

@enum CALC_MODES manual  auto_except_tables auto
const calc_modes = Dict(manual=>"manual", auto_except_tables=>"auto_except_tables", auto=>"auto")

# end of type definitions so export them

export Url, Workbook, Worksheet, Format, Chart, Chartsheet, Data, manual, auto_except_tables, auto

#####

pyfmt(fmt) = fmt === nothing ? nothing : fmt.py

function colNtocolA(n::Int64)
	a = string(Char(mod(n, 26) + 65))
	n = div(n, 26)
	while n > 0
		a = string(Char(mod(n, 26) + 64)) * a
		n = div(n, 26)
	end
	a
end

rc2cell(row::Int64, col::Int64) = "$(colNtocolA(col))$(row+1)"

function cell2rc(cell::AbstractString)
	col = 0
	while length(cell) > 0 && cell[1]>='A'
		col *= 26
		col += 1 + Int(cell[1]-'A')
		cell = cell[2:end]
	end
	parse(Int, cell)-1, col-1
end

add_worksheet!(wb::Workbook) = Worksheet(wb.py[:add_worksheet]())
add_worksheet!(wb::Workbook, name::AbstractString) = Worksheet(wb.py[:add_worksheet](name))
define_name!(wb::Workbook, name::AbstractString, target::AbstractString) = wb.py[:define_name](name, target)
worksheets(wb::Workbook) = wb.py[:worksheets]()

close(wb::Workbook) = wb.py[:close]()

get_worksheet_by_name(wb::Workbook) = wb.py[:get_worksheet_by_name]()
set_properties!(wb::Workbook, options::Dict) = wb.py[:set_properties](options)
add_format!(wb::Workbook, options::Dict=Dict()) = Format(wb.py[:add_format](options))
set_custom_property!(wb::Workbook, name::AbstractString, value::Data) = wb.py[:set_custom_property](name, value)
set_calc_mode!(wb::Workbook, mode::CALC_MODES) = wb.py[:set_calc_mode](calc_modes[mode])
add_chart!(wb::Workbook, options::Dict) = Chart(wb.py[:add_chart](options))
add_chartsheet!(wb::Workbook) = Chartsheet(wb.py[:add_chartsheet]())
set_size!(wb::Workbook, width::Int64, height::Int64) = wb.py[:set_size](width, height)
add_vba_project!(wb::Workbook, filename::AbstractString, is_stream::Bool=false) = wb.py[:add_vba_project](filename, is_stream)
set_vba_name!(wb::Workbook, name::AbstractString) = wb.py[:set_vba_name](name)
use_zip64!(wb::Workbook) = wb.py[:use_zip64]()

set_column!(ws::Worksheet, first_col::Int64, last_col::Int64, width::Real; fmt=nothing, options::Dict=Dict()) = ws.py[:set_column](first_col, last_col, width, pyfmt(fmt), options)
set_column!(ws::Worksheet, cols::AbstractString, width::Real; fmt=nothing, options::Dict=Dict()) = ws.py[:set_column](cols, width, pyfmt(fmt), options)

set_row!(ws::Worksheet, row::Int64; height=nothing, fmt=nothing, options::Dict=Dict()) = ws.py[:set_row](row, height, pyfmt(fmt), options)

# write_string / write_formula

write!(ws::Worksheet, cell::AbstractString, data; fmt=nothing) = write!(ws, cell2rc(cell)..., data, fmt=fmt)
write_string! = write!

function write!(ws::Worksheet, row::Int64, col::Int64, data::String; fmt=nothing)
	if length(data) > 0
		if data[1] == '=' || (length(data) > 1 && (data[1:2] == "{=" && data[end] == '}'))
			write!(ws, row, col, :write_formula, data, fmt=fmt)
		else
			write!(ws, row, col, :write_string, data, fmt=fmt)
		end
	else
		write!(ws, row, col, :write_blank, data, fmt=fmt)
	end
end
#write_formula! = write!

function write!(ws::Worksheet, row::Int64, col::Int64, fn::Symbol, data::Data; fmt=nothing)
	ws.py[fn](row, col, data, pyfmt(fmt))
	1
end

# convert r,c into cell format

write!(ws::Worksheet, row::Int64, col::Int64, num::Real; fmt=nothing) = write!(ws, row, col, :write_number, num, fmt=fmt)
write_number! = write!

write!(ws::Worksheet, row::Int64, col::Int64; fmt=nothing) = write!(ws, row, col, :write_blank, fmt=fmt)
write_blank! = write!

write!(ws::Worksheet, row::Int64, col::Int64, dt::DateTime; fmt=nothing) = write!(ws, row, col, :write_datetime, dt, fmt=fmt)
write!(ws::Worksheet, row::Int64, col::Int64, dt::Date; fmt=nothing) = write!(ws, row, col, :write_datetime, dt, fmt=fmt)
write_datetime! = write!

write!(ws::Worksheet, row::Int64, col::Int64, bool::Bool; fmt=nothing) = write!(ws, row, col, :write_boolean, bool, fmt=fmt)
write_bool! = write!

write!(ws::Worksheet, row::Int64, col::Int64, u::Url; fmt=nothing) = write!(ws, row, col, :write_url, u.url, fmt=fmt)
write_url! = write!

function write_matrix!(ws::Worksheet, row::Int64, col::Int64, data::Matrix; fmt=nothing)
	re = size(data, 1)-1
	ce = size(data, 2)-1
	if re > ce
		for c in 0:ce
			write_column!(ws, row, col+c, vec(data[:, c+1]), fmt=fmt)
		end
	else
		for r in 0:re
			write_row!(ws, row+r, col, vec(data[r+1, :]), fmt=fmt)
		end
	end
end

write_formula!(ws::Worksheet, row::Int64, col::Int64, formula::AbstractString; fmt=nothing, result::Data=0) = ws.py[:write_formula](row, col, formula, pyfmt(fmt), result)

write_formula!(ws::Worksheet, cell::AbstractString, formula::AbstractString; fmt=nothing, result::Data=0) = ws.py[:write_formula](cell2rc(cell)..., formula, pyfmt(fmt), result)

function write_array_formula!(ws::Worksheet, first_row::Int64, first_col::Int64, last_row::Int64, last_col::Int64, formula::AbstractString; fmt=nothing)
	ws.py[:write_array_formula](first_row, first_col, last_row, last_col, formula, pyfmt(fmt))
end

function write_array_formula!(ws::Worksheet, first_row::Int64, first_col::Int64, formula::AbstractString; fmt=nothing)
	ws.py[:write_array_formula](first_row, first_col, first_row, first_col, formula, pyfmt(fmt))
end

function write_array_formula!(ws::Worksheet, first_cell::AbstractString, last_cell::AbstractString, formula::AbstractString; fmt=nothing)
	ws.py[:write_array_formula](cell2rc(first_cell)..., cell2rc(last_cell)..., formula, pyfmt(fmt))
end

function write_array_formula!(ws::Worksheet, cell::AbstractString, formula::AbstractString; fmt=nothing)
	if occursin(':', cell) > 0
		first, last = split(cell, ':')
	else
		first = last = cell
	end

	ws.py[:write_array_formula](cell2rc(first)..., cell2rc(last)..., formula, pyfmt(fmt))
end

function write_row!(ws::Worksheet, row::Int64, col::Int64, data::Array; fmt=nothing)
	ws.py[:write_row](row, col, vec(data), pyfmt(fmt))
	length(vec(data))
end

function write_row!(ws::Worksheet, cell::AbstractString, data::Array; fmt=nothing)
	ws.py[:write_row](cell2rc(cell)..., vec(data), pyfmt(fmt))
	length(vec(data))
end

function write_column!(ws::Worksheet, row::Int64, col::Int64, data::Array; fmt=nothing)
	ws.py[:write_column](row, col, vec(data), pyfmt(fmt))
	length(vec(data))
end

function write_column!(ws::Worksheet, cell::AbstractString, data::Array; fmt=nothing)
	ws.py[:write_column](cell2rc(cell)..., vec(data), pyfmt(fmt))
	length(vec(data))
end

function write_rich_string!(ws::Worksheet, row::Int64, col::Int64, parts...)
	unwrap(f::Format) = f.py
	unwrap(a) = a
	ws.py[:write_rich_string](row, col, [unwrap(p) for p in parts]...)
end

function write_rich_string!(ws::Worksheet, cell::AbstractString, parts...)
	unwrap(f::Format) = f.py
	unwrap(a) = a
	ws.py[:write_rich_string](cell2rc(cell)..., [unwrap(p) for p in parts]...)
end

set_first_sheet!(ws::Worksheet) = ws.py[:set_first_sheet]()

merge_range!(ws::Worksheet, first_row::Int64, first_col::Int64, last_row::Int64, last_col::Int64, contents; fmt=nothing) = ws.py[:merge_range](first_row, first_col, last_row, last_col, contents, pyfmt(fmt))

freeze_panes!(ws::Worksheet, row::Int64, col::Int64) = ws.py[:freeze_panes](row, col)
freeze_panes!(ws::Worksheet, cell::AbstractString) = ws.py[:freeze_panes](cell2rc(cell)...)
freeze_panes!(ws::Worksheet, row::Int64, col::Int64, top_row::Int64) = ws.py[:freeze_panes](row, col, top_row)
freeze_panes!(ws::Worksheet, row::Int64, col::Int64, top_row::Int64, left_col::Int64) = ws.py[:freeze_panes](row, col, top_row, left_col)

split_panes!(ws::Worksheet,x::Int64, y::Int64) = ws.py[:split_panes](x, y)
split_panes!(ws::Worksheet,x::Int64, y::Int64, top_row::Int64) = ws.py[:split_panes](x, y, top_row)
split_panes!(ws::Worksheet,x::Int64, y::Int64, top_row::Int64, left_col::Int64) = ws.py[:split_panes](x, y, top_row, left_col)

data_validation!(ws::Worksheet, first_row::Int64, first_col::Int64, last_row::Int64, last_col::Int64, options::Dict) = ws.py[:data_validation](first_row, first_col, last_row, last_col, options)
data_validation!(ws::Worksheet, first_cell::AbstractString, last_cell::AbstractString, options::Dict) = ws.py[:data_validation](cell2rc(first_cell)..., cell2rc(last_cell)..., options)

function conditional_format!(ws::Worksheet, first_row::Int64, first_col::Int64, last_row::Int64, last_col::Int64, options::Dict)
	for k in collect(keys(options))
		if typeof(options[k]) == XlsxWriter.Format
			options[k] = options[k].py
		end
	end
	ws.py[:conditional_format](first_row, first_col, last_row, last_col, options)
end

conditional_format!(ws::Worksheet, first_cell::AbstractString, last_cell::AbstractString, options::Dict) = conditional_format!(ws, cell2rc(first_cell)..., cell2rc(last_cell)..., options)
function conditional_format!(ws::Worksheet, cell::AbstractString, options::Dict)
	if occursin(':', cell) > 0
		f, s = split(cell, ':')
		conditional_format!(ws, cell2rc(f)..., cell2rc(s)..., options)
	else
		conditional_format!(ws, cell2rc(cell)..., cell2rc(cell)..., options)
	end
end

add_table!(ws::Worksheet, first_row::Int64, first_col::Int64, last_row::Int64, last_col::Int64, options::Dict) = ws.py[:add_table](first_row, first_col, last_row, last_col, options)
add_table!(ws::Worksheet, first_cell::AbstractString, last_cell::AbstractString, options::Dict) = ws.py[:add_table](cell2rc(first_cell), cell2rc(last_cell), options)

add_sparkline!(ws::Worksheet, row::Int64, col::Int64, options::Dict) = ws.py[:add_sparkline](row, col, options)
add_sparkline!(ws::Worksheet, cell::AbstractString, options::Dict) = ws.py[:add_sparkline](cell2rc(cell)..., options)

insert_image!(ws::Worksheet, row::Int64, col::Int64, image::AbstractString, options::Dict=Dict()) = ws.py[:insert_image](row, col, image, options)
insert_image!(ws::Worksheet, cell::AbstractString, image::AbstractString, options::Dict=Dict()) = ws.py[:insert_image](cell2rc(cell)..., image, options)

insert_chart!(ws::Worksheet, row::Int64, col::Int64, ch::Chart, options::Dict=Dict()) = ws.py[:insert_chart](row, col, ch.py, options)
insert_chart!(ws::Worksheet, cell::AbstractString, ch::Chart, options::Dict=Dict()) = ws.py[:insert_chart](cell2rc(cell)..., ch.py, options)

insert_textbox!(ws::Worksheet, row::Int64, col::Int64, text::AbstractString, options::Dict=Dict()) = ws.py[:insert_textbox](row, col, text, options)
insert_textbox!(ws::Worksheet, cell::AbstractString, text::AbstractString, options::Dict=Dict()) = ws.py[:insert_textbox](cell2rc(cell)..., text, options)

insert_button!(ws::Worksheet, row::Int64, col::Int64, options::Dict) = wb.py[:insert_button](row, col, options)
insert_button!(ws::Worksheet, cell::AbstractString, options::Dict) = wb.py[:insert_button](cell2rc(cell)..., options)

write_comment!(ws::Worksheet, row::Int64, col::Int64, comment::AbstractString, options::Dict=Dict()) = ws.py[:write_comment](row, col, comment, options)
write_comment!(ws::Worksheet, cell::AbstractString, comment::AbstractString, options::Dict=Dict()) = ws.py[:write_comment](cell2rc(cell)..., comment, options)

show_comments!(ws::Worksheet) = ws.py[:show_comments]()
set_comments_author!(ws::Worksheet, author::AbstractString) = ws.py[:set_comments_author](author)

autofilter!(ws::Worksheet, first_row::Int64, first_col::Int64, last_row::Int64, last_col::Int64) = ws.py[:autofilter](first_row, first_col, last_row, last_col)
autofilter!(ws::Worksheet, cells::AbstractString) = ws.py[:autofilter](cells)

filter_column!(ws::Worksheet, col::Int64, criteria::AbstractString) = ws.py[:filter_column](col, criteria)
filter_column!(ws::Worksheet, col::AbstractString, criteria::AbstractString) = ws.py[:filter_column](col, criteria)

filter_column_list!(ws::Worksheet, col::Int64, criteria::Array{AbstractString}) = ws.py[:filter_column](col, vec(criteria))
filter_column!(ws::Worksheet, col::AbstractString, criteria::Array{AbstractString}) = ws.py[:filter_column](col, vec(criteria))

set_selection!(ws::Worksheet, first_row::Int64, first_col::Int64, last_row::Int64, last_col::Int64) = ws.py[:set_selection](first_row, first_col, last_row, last_col)
set_selection!(ws::Worksheet, cells::AbstractString) = ws.py[:set_selection](cells)

set_default_row!(ws::Worksheet, height::Float64=15.; hide_unused_rows::Bool=false) = ws.py[:set_default_row](height, hide_unused_rows)

outline_settings!(ws::Worksheet, visible::Bool=true, symbols_below::Bool=true, symbols_right::Bool=true, auto_style::Bool=false) = ws.py[:outline_settings](visible, symbols_below, symbols_right, auto_style)

# worksheet

hide_zero!(sh::Worksheet) = sh.py[:hide_zero]()

# worksheet or chartsheet
const Sheet = Union{Worksheet, Chartsheet}
activate!(sh::Sheet) = sh.py[:activate]()
select!(sh::Sheet) = sh.py[:select]()
hide!(sh::Sheet) = sh.py[:hide]()
set_first_sheet!(sh::Sheet) = sh.py[:set_first_sheet]()
right_to_left!(sh::Sheet) = sh.py[:right_to_left]()
protect!(sh::Sheet, password::AbstractString, options::Dict) = sh.py[:protect](password, options)
set_zoom!(sh::Sheet, zoom::Int64) = sh.py[:set_zoom](zoom)
set_tab_color!(sh::Sheet, color::AbstractString) = sh.py[:set_tab_color](color)
set_landscape!(sh::Sheet) = sh.py[:set_landscape]()
set_portrait!(sh::Sheet) = sh.py[:set_portrait]()
set_paper!(sh::Sheet, index::Int64) = sh.py[:set_paper](index)
set_margins!(sh::Sheet, left::Float64=0.7, right::Float64=0.7, top::Float64=0.75, bottom::Float64=0.75) = sh.py[:set_margins](left, right, top, bottom)
set_header!(sh::Sheet, header::AbstractString, options::Dict=Dict()) = sh.py[:set_header](header, options)
set_header!(sh::Sheet, options::Dict) = sh.py[:set_header]("", options)
set_footer!(sh::Sheet, footer::AbstractString, options::Dict=Dict()) = sh.py[:set_footer](footer, options)
set_footer!(sh::Sheet, options::Dict) = sh.py[:set_footer]("", options)
get_name(sh::Sheet) = sh.py[:get_name]()

set_font_name!(fmt::Format, opt::AbstractString) = fmt.py[:set_font_name](opt)
set_font_size!(fmt::Format, opt::Int64) = fmt.py[:set_font_size](opt)
set_font_color!(fmt::Format, opt::AbstractString) = fmt.py[:set_font_color](opt)
set_bold!(fmt::Format, opt::Bool=true) = fmt.py[:set_bold](opt)
set_italic!(fmt::Format, opt::Bool=true) = fmt.py[:set_italic](opt)
set_underline!(fmt::Format, opt::Bool=true) = fmt.py[:set_underline](opt)
set_font_strikeout!(fmt::Format, opt::Bool=true) = fmt.py[:set_font_strikeout](opt)
set_font_script!(fmt::Format, opt::Int64) = fmt.py[:set_font_script](opt)
set_num_format!(fmt::Format, opt::AbstractString) = fmt.py[:set_font_num_format](opt)
set_locked!(fmt::Format, opt::Bool=true) = fmt.py[:set_locked](opt)
set_hidden!(fmt::Format, opt::Bool=true) = fmt.py[:set_hidden](opt)
set_align!(fmt::Format, opt::AbstractString) = fmt.py[:set_align](opt)
set_center_across!(fmt::Format, opt::AbstractString) = fmt.py[:set_center_across](opt)
set_text_wrap!(fmt::Format, opt::Bool=true) = fmt.py[:set_text_wrap](opt)
set_rotation!(fmt::Format, opt::Int64) = fmt.py[:set_rotation](opt)
set_indent!(fmt::Format, opt::Int64) = fmt.py[:set_indent](opt)
set_shrink!(fmt::Format, opt::Bool=true) = fmt.py[:set_shrink](opt)
set_text_justlast!(fmt::Format, opt::Bool=true) = fmt.py[:set_text_justlast](opt)
set_pattern!(fmt::Format, opt::Int64) = fmt.py[:set_pattern](opt)
set_bg_color!(fmt::Format, opt::AbstractString) = fmt.py[:set_bg_color](opt)
set_fg_color!(fmt::Format, opt::AbstractString) = fmt.py[:set_fg_color](opt)
set_border!(fmt::Format, opt::Int64) = fmt.py[:set_border](opt)
set_bottom!(fmt::Format, opt::Int64) = fmt.py[:set_bottom](opt)
set_top!(fmt::Format, opt::Int64) = fmt.py[:set_top](opt)
set_left!(fmt::Format, opt::Int64) = fmt.py[:set_left](opt)
set_right!(fmt::Format, opt::Int64) = fmt.py[:set_right](opt)
set_border_color!(fmt::Format, opt::AbstractString) = fmt.py[:set_border_color](opt)
set_bottom_color!(fmt::Format, opt::AbstractString) = fmt.py[:set_bottom_color](opt)
set_top_color!(fmt::Format, opt::AbstractString) = fmt.py[:set_top_color](opt)
set_left_color!(fmt::Format, opt::AbstractString) = fmt.py[:set_left_color](opt)
set_right_color!(fmt::Format, opt::AbstractString) = fmt.py[:set_right_color](opt)
set_diag_border!(fmt::Format, opt::Int64) = fmt.py[:set_diag_border](opt)
set_diag_type!(fmt::Format, opt::Int64) = fmt.py[:set_diag_type](opt)
set_diag_color!(fmt::Format, opt::AbstractString) = fmt.py[:set_diag_color](opt)

# Chart

add_series!(ch::Chart, options::Dict=Dict()) = ch.py[:add_series](options)
set_x_axis!(ch::Chart, options::Dict=Dict()) = ch.py[:set_x_axis](options)
set_y_axis!(ch::Chart, options::Dict=Dict()) = ch.py[:set_y_axis](options)
set_x2_axis!(ch::Chart, options::Dict=Dict()) = ch.py[:set_x2_axis](options)
set_y2_axis!(ch::Chart, options::Dict=Dict()) = ch.py[:set_y2_axis](options)
combine!(ch1::Chart, ch2::Chart) = ch1.py[:combine](ch2.py)
set_size!(ch::Chart, options::Dict=Dict()) = ch.py[:set_size](options)
set_title!(ch::Chart, options::Dict=Dict()) = ch.py[:set_title](options)
set_legend!(ch::Chart, options::Dict=Dict()) = ch.py[:set_legend](options)
set_chartarea!(ch::Chart, options::Dict=Dict()) = ch.py[:set_chartarea](options)
set_plotarea!(ch::Chart, options::Dict=Dict()) = ch.py[:set_plotarea](options)
set_style!(ch::Chart, style_id::Int64) = ch.py[:set_style](style_id)
set_table!(ch::Chart, options::Dict=Dict()) = ch.py[:set_table](options)
set_up_down_bars!(ch::Chart, options::Dict=Dict()) = ch.py[:set_up_down_bars](options)
set_drop_lines!(ch::Chart, options::Dict=Dict()) = ch.py[:set_drop_lines](options)
set_high_low_lines!(ch::Chart, options::Dict=Dict()) = ch.py[:set_high_low_lines](options)
show_blanks_as!(ch::Chart, option::AbstractString) = ch.py[:show_blanks_as](option)
show_hidden_data!(ch::Chart) = ch.py[:show_hidden_data]()
set_rotation!(ch::Chart, angle::Int64) = ch.py[:set_rotation](angle)
set_hole_size!(ch::Chart, hole::Int64) = ch.py[:set_hole_size](hole)

# Chartsheet

set_chart!(cs::Chartsheet, ch::Chart) = cs.py[:set_chart](ch.py)

end


text:
function get_adaption_of_snapshot(snapshot, reference_snapshot, rel_tolerance, min_samples, max_samples)
    all_organisms = get_organisms(snapshot)
    all_reference_organisms = get_organisms(reference_snapshot)

    adaption = estimate(rel_tolerance, min_samples, max_samples) do
        while true
            organism = rand(all_organisms)
            reference_organism = rand(all_reference_organisms)
            reference_id = get_id(reference_snapshot, reference_organism)

            genotype = get_genotype(snapshot, organism)
            reference_genotype = get_genotype(reference_snapshot, reference_organism)

            reference_fitness = get_fitness(reference_snapshot, reference_id, reference_genotype)
            fitness = get_fitness(reference_snapshot, reference_id, genotype)

            if reference_fitness == 0 continue end

            return fitness / reference_fitness
        end
    end

    return adaption
end

max_stars_repo_name: tochsner/ALifeBenchmark

max_issues_repo_name: tochsner/ALifeBenchmark

max_forks_repo_name: tochsner/ALifeBenchmark


text:
# token
struct Failed end
failed = Failed()

# core
cache(f) = (target, body) ->
    gensym("cache_$target") |> TARGET ->
    Expr(
      :block,
      :($TARGET = $target),
      :($(f(TARGET, body)))
    )

wildcard(target, body) = body

capture(sym::Symbol) = (target, body) ->
  Expr(:let, Expr(:block, :($sym = $target)), body)

literal(v) = (target, body) ->
  :($v == $target ? $body : $failed)

and(p1, p2) = (target, body) ->
	p1(target, p2(target, body))

or(p1, p2) = (target, body) ->
  let e1 = p1(target, body),
      e2 = p2(target, body),
      RET = gensym("or_$target")
      Expr(
        :block,
        :($RET = $e1),
        :($RET === $failed ? $e2 : $RET)
      )
  end

# C(p1, p2, .., pn)
# pattern = (target: code, remainder: code) -> code
recog(recogniser, get_field) = elts -> recogniser(
  function (target, body)
    n = elts |> length
    foldr(1:n, init=body) do i, last
      sub_target = get_field(target, i)
      elts[i](sub_target, last)
    end
  end)

typed_as(T) = f -> (target, body) ->
   let FN = gensym("f_$target"),
       TARGET = gensym("$target"),
       f1 = Expr(:function, :($FN($TARGET::$T)), :($(f(TARGET, body)))),
       f2 = Expr(:function, :($FN($TARGET)), failed)
       Expr(:block, f1, f2, :($FN($target)))
   end

match(target) = pairs ->
  let RET = gensym("or_$target"),
      TAG = gensym("var_$target")
      foldr(pairs, init=:(error("non-exhaustive"))) do (case, body), last
          expr = case(TAG, body)
          Expr(
            :block,
            :($RET = $expr),
            :($RET === $failed ? $last : $RET)
        )
      end |> x -> Expr(:block, :($TAG = $target), x)
  end

# pervasives

tuple_pat(elts) =
  let recogniser = typed_as(NTuple{elts |> length, Any}),
      get_field(tag, i) = :($tag[$i])
      recog(recogniser, get_field)(elts)
  end

pred_pat(pred) = (target, body) ->
  let cond = :($pred($target))
    :($cond ? $body : $failed)
  end

guard(cond) = (target, body) -> :($cond ? $body : $failed)

array_pat(elts) =
  let recogniser = typed_as(Vector),
      get_field(tag, i) = :($tag[$i])
      recog(recogniser, get_field)(elts)
  end

array_pat2(elts1, star, elts2) = # with ...
  let (n1, n2) = map(length, [elts1, elts2]),
      recogniser = f -> typed_as(Vector)(
          let n = n1 + n2
              check = n > 0 ?
                pred_pat(Expr(:(->), :x, :(x >= $(n1 + n2)))) :
                wildcard
              and(f, check)
          end
      ),
      get_field(tag, i) =
          i <= n1 ?
            :($tag[$i]) :
          i > n1 + 1 ?
            :($tag[end - $(n2 + n1 + 1 - i)]) :
          :($tag[$(n1 + 1):end - $n2])
      recog(recogniser, get_field)([elts1..., star, elts2...])
  end

# test

@match target begin
  (a, _, 3) => expr
end

tuple_pat(
  (
    capture(:a),
    wildcard,
    literal(3)
  )
)(:var, :expr1) |> println

# array_pat2(
#   [literal(1)],
#   capture(:a),
#   [literal(2)]
# )(:a, :res) |> println

# p1 = tuple_pat(
#   (
#     capture(:a),
#     wildcard,
#     literal(3)
#   )
# )

# p1(:var, :expr1) |> println

# match(target) = pairs ->
#   let RET = gensym("or_$target"),
#       TAG = gensym("var_$target")
#       foldr(pairs, init=:(error("non-exhaustive"))) do (case, body), last
#           expr = case(TAG, body)
#           Expr(
#             :block,
#             :($RET = $expr),
#             :($RET === $failed ? $last : $RET)
#         )
#       end |> x -> Expr(:block, :($TAG = $target), x)
#   end

# macro test_ast(x)
#   match(x)([
#     p1 => :(a + 1),
#     wildcard => 5
#   ]) |> esc
# end

# match(:x)([
#     p1 => :(a + 1),
#     wildcard => 5
# ]) |> println

# function f1(x)
#   @test_ast x
# end

# function f2(x)
#   if x isa Tuple && length(x) == 3
#     x[1] + 1
#   else
#     5
#   end

# end
# using InteractiveUtils, BenchmarkTools
# @code_warntype f1((1, 2, 3))
# println()
# @code_warntype f2((1, 2, 3))
# @btime f1((1, 2, 3))
# @btime f2((1, 2, 3))

max_stars_repo_name: thautwarm/Slides

max_issues_repo_name: thautwarm/Slides

max_forks_repo_name: thautwarm/Slides


text:
# pull out DataFrames w/ file names and obs parameters
function sort_bisector_data(;dir::String=soldir, write::Bool=false)
    # glob the files
    @assert isdir(dir)
    files = glob("*.bisect.fits", dir * "bisectors")
    @assert !isempty(files)

    # extract file parametersl put in data frame
    df = DataFrame(fpath=String[], fname=String[], datetime=DateTime[],
                   wave=String[], mu=String[], axis=String[], ion=String[])
    for i in 1:length(files)
        push!(df, extract_bisector_params(files[i]))
    end

    # sort dataframe in place
    sort!(df, [:axis, :mu, :datetime], rev=(false, true, false))

    # write to CSV
    if write
        CSV.write(df.wave[1] * "_bisectors.csv", df)
    end
    return df
end

# extract observation parameters from string name
function extract_bisector_params(s::String)
    # parse out filename
    paths = split(s, "/")
    fpath = join(paths[1:end-1], "/") * "/"
    fname = paths[end]
    fcomp = split(fname, "_")

    # parse out parameters
    datetimes = DateTime(fcomp[3], "YYYYmmdd-HHMMSS")
    wavelength = split(fcomp[4],"clv")[2]
    muposition = split(fcomp[5], ".")[1]
    if length(fcomp) == 7
        helio_axis = split(fcomp[6], ".")[1]
        ionspecies = split(fcomp[6], ".")[4]
    else
        helio_axis = "c"
        ionspecies = split(fcomp[5], ".")[4]
    end
    return [fpath fname datetimes wavelength muposition helio_axis ionspecies]
end

# simple function to read out wavelength/bisector vectors
function read_bisector(filename::String; masknans::Bool=false)
    wav = Array{Float64,1}[]
    bis = Array{Float64,1}[]
    FITS(filename) do f
        # λs in 1st extension, bis value in 2nd
        wav = FITSIO.read(f[1], :, :, 1)
        bis = FITSIO.read(f[1], :, :, 2)
    end

    # mask NaNs
    if masknans
        mask = .!isnan.(wav)
        wav = wav[mask]
        bis = bis[mask]
    end
    return wav, bis
end

function get_extension_dims(filename::String)
    out = 0
    FITS(filename) do f
        out = size(f[1])
    end
    return out
end

function get_number_waves(filename::String)
    dims = get_extension_dims(filename)
    return dims[1]
end

function get_number_times(filename::String)
    dims = get_extension_dims(filename)
    return dims[2]
end

# make one large time series for a given solar position
function stitch_time_series(df::DataFrame; adjust_mean::Bool=false, contiguous_only::Bool=false)
    # find out size of data
    if contiguous_only
        Nfil=1
    else
        Nfil = size(df, 1)
    end
    Ntim = map(get_number_times, df.fpath[i] * df.fname[i] for i in 1:Nfil)
    Nwav = map(get_number_waves, df.fpath[i] * df.fname[i] for i in 1:Nfil)

    # allocate array
    wavall = zeros(maximum(Nwav), sum(Ntim))
    bisall = zeros(maximum(Nwav), sum(Ntim))

    # loop over the files, filling arrays
    for i in 1:Nfil
        wav, bis = read_bisector(df.fpath[i] * df.fname[i])
        wavall[:, sum(Ntim[1:i-1])+1:sum(Ntim[1:i])] .= wav
        bisall[:, sum(Ntim[1:i-1])+1:sum(Ntim[1:i])] .= bis
    end

    # adjust the mean if adjust_mean == true
    if adjust_mean
        # find the mean for the first chunk of bisectors
        wavgroup1 = wavall[:, 1:Ntim[1]]
        meanwav1 = mean(wavgroup1, dims=2)[:,1]
        for i in 2:Nfil
            # find the mean for the nth group of bisectors
            wavgroupn = wavall[:, sum(Ntim[1:i-1])+1:sum(Ntim[1:i])]
            meanwavn = mean(wavgroupn, dims=2)[:,1]

            # find the distance between the means and correct by it
            meandist = meanwavn - meanwav1
            wavall[:, sum(Ntim[1:i-1])+1:sum(Ntim[1:i])] .-= meandist
        end
    end
    return wavall, bisall
end

function sort_spectrum_data(;dir::String=soldir, write::Bool=false)
    # glob the files
    @assert isdir(dir)
    files = glob("*.chvtt.fits", dir * "spectra/")
    @assert !isempty(files)

    # extract file parametersl put in data frame
    df = DataFrame(fpath=String[], fname=String[], datetime=DateTime[],
                   wave=String[], mu=String[], axis=String[])
    for i in 1:length(files)
        push!(df, extract_line_params(files[i]))
    end

    # sort dataframe in place
    sort!(df, [:axis, :mu, :datetime], rev=(false, true, false))

    # write to CSV
    if write
        CSV.write(df.wave[1] * "_lines.csv", df)
    end
    return df
end

function sort_width_data(;dir::String=soldir, write::Bool=false)
    # glob the files
    @assert isdir(dir)
    files = glob("*widths.fits", dir * "widths/")
    @assert !isempty(files)

    # extract file parametersl put in data frame
    df = DataFrame(fpath=String[], fname=String[], datetime=DateTime[],
                   wave=String[], mu=String[], axis=String[])
    for i in 1:length(files)
        push!(df, extract_width_params(files[i]))
    end

    # sort dataframe in place
    sort!(df, [:axis, :mu, :datetime], rev=(false, true, false))

    # write to CSV
    if write
        CSV.write(df.wave[1] * "_widths.csv", df)
    end
    return df
end

function extract_line_params(s::String)
    # parse out filename
    paths = split(s, "/")
    fpath = join(paths[1:end-1], "/") * "/"
    fname = paths[end]
    fcomp = split(fname, "_")

    # parse out parameters
    datetimes = DateTime(fcomp[3], "YYYYmmdd-HHMMSS")
    wavelength = split(fcomp[4],"clv")[2]
    muposition = split(fcomp[5], ".")[1]
    if length(fcomp) == 6
        helio_axis = split(fcomp[6], ".")[1]
    else
        helio_axis = "c"
    end
    return [fpath fname datetimes wavelength muposition helio_axis]
end

function extract_width_params(s::String)
    # parse out filename
    paths = split(s, "/")
    fpath = join(paths[1:end-1], "/") * "/"
    fname = paths[end]
    fcomp = split(fname, "_")

    # parse out parameters
    datetimes = DateTime(fcomp[3], "YYYYmmdd-HHMMSS")
    wavelength = split(fcomp[4],"clv")[2]
    muposition = split(fcomp[5], ".")[1]
    if length(fcomp) == 7
        helio_axis = split(fcomp[6], ".")[1]
        # ionspecies = split(fcomp[6], ".")[4]
    else
        helio_axis = "c"
        # ionspecies = split(fcomp[5], ".")[4]
    end
    return [fpath fname datetimes wavelength muposition helio_axis]
end

function bin_spectrum(wavs::AbstractArray{T,2}, spec::AbstractArray{T,2};
                      binsize::Integer=10) where T<:Real
    @assert size(wavs) == size(spec)
    @assert binsize >= 1
    @assert binsize < size(wavs,2)

    # get sizes, allocate memory
    nwave = size(wavs,1)
    ntime = size(wavs,2)
    nbins = floor(Int, ntime/binsize)
    wavsb = zeros(nwave, nbins)
    specb = zeros(nwave, nbins)

    # do the binning
    for i in 0:(nbins-1)
        wavsb[:,i+1] = mean(wavs[:, i*binsize+1:(i+1)*binsize], dims=2)
        specb[:,i+1] = mean(spec[:, i*binsize+1:(i+1)*binsize], dims=2)
    end
    return wavsb, specb
end

function read_spectrum(filename::String)
    # Primary HDU has spectra + noise in frame 1/2
    spec = []
    nois = []
    head = []
    wavs = []
    FITS(filename) do f
        spec = FITSIO.read(f[1])[:,:,1]
        nois = FITSIO.read(f[1])[:,:,2]
        head = read_header(f[1])
        wavs = FITSIO.read(f[2]) .* 1.0e10  # convert to angstroms
    end
    return wavs, convert.(Float64, spec)
end

function write_the_fits(fname::String, depth::AbstractArray{T,2}, width::AbstractArray{T,2}) where T<:Real
    FITS(fname, "w") do io
        write(io, cat(depth, width, dims=3))
    end
    return nothing
end

max_stars_repo_name: palumbom/GRASS

max_issues_repo_name: palumbom/GRASS

max_forks_repo_name: palumbom/GRASS


text:
using Test, ReDiverse
cd("..")

########################################
ReDiverse.print_sst("Compile the C++ codes")
ReDiverse.make()

########################################
ReDiverse.print_sst("Test tools and data")

@testset "Tools" begin
    @test isfile(ReDiverse.beagle) || "beagle.jar not ready in bin"
    @test isfile(ReDiverse.plink) || "Plink not ready in bin"
end

@testset "Genotypes" begin
    @test length(readdir("data/genotypes/german/ld/")) == 11
    @test length(readdir("data/genotypes/german/md/")) == 4
    @test length(readdir("data/genotypes/german/v2/")) == 69
    @test length(readdir("data/genotypes/german/v3/")) == 732
    @test length(readdir("data/genotypes/dutch/10690/")) == 1994
    @test length(readdir("data/genotypes/dutch/10993/")) == 102
    @test length(readdir("data/genotypes/dutch/11483/")) == 109
    @test length(readdir("data/genotypes/dutch/50kv3/")) == 44
    @test length(readdir("data/genotypes/dutch/54609/")) == 343
    @test length(readdir("data/genotypes/dutch/777k/")) == 10
    @test length(readdir("data/genotypes/norge/")) == 9
end

@testset "Maps" begin
    for i in ["BovineSNP50_v3_A2.csv",
              "cow.50k.v1.tsv",
              "cow.50k.v2.tsv",
              "dutch.snp",
              "hd-mnft.csv"]
        @test isfile("data/maps/illumina/" * i)
    end
end


text:
## This file autogenerated by ext/HistoricalStdlibGenerator/generate_historical_stdlibs.jl

using Base: UUID

# Julia standard libraries with duplicate entries removed so as to store only the
# first release in a set of releases that all contain the same set of stdlibs.
const STDLIBS_BY_VERSION = [
    v"1.0.0" => Dict(
        UUID("2a0f44e3-6c83-55bd-87e4-b1978d98bd5f") => "Base64",
        UUID("8bf52ea8-c179-5cab-976a-9e18b702a9bc") => "CRC32c",
        UUID("ade2ca70-3891-5945-98fb-dc099432e06a") => "Dates",
        UUID("8bb1440f-4735-579b-a4ab-409b98df4dab") => "DelimitedFiles",
        UUID("8ba89e20-285c-5b6f-9357-94700520ee1b") => "Distributed",
        UUID("7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee") => "FileWatching",
        UUID("9fa8497b-333b-5362-9e8d-4d0656e87820") => "Future",
        UUID("b77e0a4c-d291-57a0-90e8-8db25a27a240") => "InteractiveUtils",
        UUID("76f85450-5226-5b5a-8eaa-529ad045b433") => "LibGit2",
        UUID("8f399da3-3557-5675-b5ff-fb832c97cbdb") => "Libdl",
        UUID("37e2e46d-f89d-539d-b4ee-838fcccc9c8e") => "LinearAlgebra",
        UUID("56ddb016-857b-54e1-b83d-db4d58db5568") => "Logging",
        UUID("d6f4376e-aef5-505a-96c1-9c027394607a") => "Markdown",
        UUID("a63ad114-7e13-5084-954f-fe012c677804") => "Mmap",
        UUID("44cfe95a-1eb2-52ea-b672-e2afdf69b78f") => "Pkg",
        UUID("de0858da-6303-5e67-8744-51eddeeeb8d7") => "Printf",
        UUID("9abbd945-dff8-562f-b5e8-e1ebf5ef1b79") => "Profile",
        UUID("3fa0cd96-eef1-5676-8a61-b3b8758bbffb") => "REPL",
        UUID("9a3f8284-a2c9-5f02-9a11-845980a1fd5c") => "Random",
        UUID("ea8e919c-243c-51af-8825-aaa63cd721ce") => "SHA",
        UUID("9e88b42a-f829-5b0c-bbe9-9e923198166b") => "Serialization",
        UUID("1a1011a3-84de-559e-8e89-a11a2f7dc383") => "SharedArrays",
        UUID("6462fe0b-24de-5631-8697-dd941f90decc") => "Sockets",
        UUID("2f01184e-e22b-5df5-ae63-d93ebab69eaf") => "SparseArrays",
        UUID("10745b16-79ce-11e8-11f9-7d13ad32a3b2") => "Statistics",
        UUID("4607b0f0-06f3-5cda-b6b1-a6196a1729e9") => "SuiteSparse",
        UUID("8dfed614-e22c-5e08-85e1-65c5234f0b40") => "Test",
        UUID("cf7118a7-6976-5b1a-9a39-7adc72f591a4") => "UUIDs",
        UUID("4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5") => "Unicode",
    ),
    v"1.6.0" => Dict(
        UUID("0dad84c5-d112-42e6-8d28-ef12dabb789f") => "ArgTools",
        UUID("56f22d72-fd6d-98f1-02f0-08ddc0907c33") => "Artifacts",
        UUID("2a0f44e3-6c83-55bd-87e4-b1978d98bd5f") => "Base64",
        UUID("8bf52ea8-c179-5cab-976a-9e18b702a9bc") => "CRC32c",
        UUID("e66e0078-7015-5450-92f7-15fbd957f2ae") => "CompilerSupportLibraries_jll",
        UUID("ade2ca70-3891-5945-98fb-dc099432e06a") => "Dates",
        UUID("8bb1440f-4735-579b-a4ab-409b98df4dab") => "DelimitedFiles",
        UUID("8ba89e20-285c-5b6f-9357-94700520ee1b") => "Distributed",
        UUID("f43a241f-c20a-4ad4-852c-f6b1247861c6") => "Downloads",
        UUID("7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee") => "FileWatching",
        UUID("9fa8497b-333b-5362-9e8d-4d0656e87820") => "Future",
        UUID("781609d7-10c4-51f6-84f2-b8444358ff6d") => "GMP_jll",
        UUID("b77e0a4c-d291-57a0-90e8-8db25a27a240") => "InteractiveUtils",
        UUID("4af54fe1-eca0-43a8-85a7-787d91b784e3") => "LazyArtifacts",
        UUID("b27032c2-a3e7-50c8-80cd-2d36dbcbfd21") => "LibCURL",
        UUID("deac9b47-8bc7-5906-a0fe-35ac56dc84c0") => "LibCURL_jll",
        UUID("76f85450-5226-5b5a-8eaa-529ad045b433") => "LibGit2",
        UUID("e37daf67-58a4-590a-8e99-b0245dd2ffc5") => "LibGit2_jll",
        UUID("a83860b7-747b-57cf-bf1f-3e79990d037f") => "LibOSXUnwind_jll",
        UUID("29816b5a-b9ab-546f-933c-edad1886dfa8") => "LibSSH2_jll",
        UUID("183b4373-6708-53ba-ad28-60e28bb38547") => "LibUV_jll",
        UUID("745a5e78-f969-53e9-954f-d19f2f74f4e3") => "LibUnwind_jll",
        UUID("8f399da3-3557-5675-b5ff-fb832c97cbdb") => "Libdl",
        UUID("37e2e46d-f89d-539d-b4ee-838fcccc9c8e") => "LinearAlgebra",
        UUID("56ddb016-857b-54e1-b83d-db4d58db5568") => "Logging",
        UUID("3a97d323-0669-5f0c-9066-3539efd106a3") => "MPFR_jll",
        UUID("d6f4376e-aef5-505a-96c1-9c027394607a") => "Markdown",
        UUID("c8ffd9c3-330d-5841-b78e-0817d7145fa1") => "MbedTLS_jll",
        UUID("a63ad114-7e13-5084-954f-fe012c677804") => "Mmap",
        UUID("14a3606d-f60d-562e-9121-12d972cd8159") => "MozillaCACerts_jll",
        UUID("ca575930-c2e3-43a9-ace4-1e988b2c1908") => "NetworkOptions",
        UUID("4536629a-c528-5b80-bd46-f80d51c5b363") => "OpenBLAS_jll",
        UUID("05823500-19ac-5b8b-9628-191a04bc5112") => "OpenLibm_jll",
        UUID("efcefdf7-47ab-520b-bdef-62a2eaa19f15") => "PCRE2_jll",
        UUID("44cfe95a-1eb2-52ea-b672-e2afdf69b78f") => "Pkg",
        UUID("de0858da-6303-5e67-8744-51eddeeeb8d7") => "Printf",
        UUID("9abbd945-dff8-562f-b5e8-e1ebf5ef1b79") => "Profile",
        UUID("3fa0cd96-eef1-5676-8a61-b3b8758bbffb") => "REPL",
        UUID("9a3f8284-a2c9-5f02-9a11-845980a1fd5c") => "Random",
        UUID("ea8e919c-243c-51af-8825-aaa63cd721ce") => "SHA",
        UUID("9e88b42a-f829-5b0c-bbe9-9e923198166b") => "Serialization",
        UUID("1a1011a3-84de-559e-8e89-a11a2f7dc383") => "SharedArrays",
        UUID("6462fe0b-24de-5631-8697-dd941f90decc") => "Sockets",
        UUID("2f01184e-e22b-5df5-ae63-d93ebab69eaf") => "SparseArrays",
        UUID("10745b16-79ce-11e8-11f9-7d13ad32a3b2") => "Statistics",
        UUID("4607b0f0-06f3-5cda-b6b1-a6196a1729e9") => "SuiteSparse",
        UUID("bea87d4a-7f5b-5778-9afe-8cc45184846c") => "SuiteSparse_jll",
        UUID("fa267f1f-6049-4f14-aa54-33bafae1ed76") => "TOML",
        UUID("a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e") => "Tar",
        UUID("8dfed614-e22c-5e08-85e1-65c5234f0b40") => "Test",
        UUID("cf7118a7-6976-5b1a-9a39-7adc72f591a4") => "UUIDs",
        UUID("4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5") => "Unicode",
        UUID("83775a58-1f1d-513f-b197-d71354ab007a") => "Zlib_jll",
        UUID("05ff407c-b0c1-5878-9df8-858cc2e60c36") => "dSFMT_jll",
        UUID("8f36deef-c2a5-5394-99ed-8e07531fb29a") => "libLLVM_jll",
        UUID("8e850ede-7688-5339-a07c-302acd2aaf8d") => "nghttp2_jll",
        UUID("3f19e933-33d8-53b3-aaab-bd5110c3b7a0") => "p7zip_jll",
    ),
    v"1.7.0" => Dict(
        UUID("0dad84c5-d112-42e6-8d28-ef12dabb789f") => "ArgTools",
        UUID("56f22d72-fd6d-98f1-02f0-08ddc0907c33") => "Artifacts",
        UUID("2a0f44e3-6c83-55bd-87e4-b1978d98bd5f") => "Base64",
        UUID("8bf52ea8-c179-5cab-976a-9e18b702a9bc") => "CRC32c",
        UUID("e66e0078-7015-5450-92f7-15fbd957f2ae") => "CompilerSupportLibraries_jll",
        UUID("ade2ca70-3891-5945-98fb-dc099432e06a") => "Dates",
        UUID("8bb1440f-4735-579b-a4ab-409b98df4dab") => "DelimitedFiles",
        UUID("8ba89e20-285c-5b6f-9357-94700520ee1b") => "Distributed",
        UUID("f43a241f-c20a-4ad4-852c-f6b1247861c6") => "Downloads",
        UUID("7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee") => "FileWatching",
        UUID("9fa8497b-333b-5362-9e8d-4d0656e87820") => "Future",
        UUID("781609d7-10c4-51f6-84f2-b8444358ff6d") => "GMP_jll",
        UUID("b77e0a4c-d291-57a0-90e8-8db25a27a240") => "InteractiveUtils",
        UUID("47c5dbc3-30ba-59ef-96a6-123e260183d9") => "LLVMLibUnwind_jll",
        UUID("4af54fe1-eca0-43a8-85a7-787d91b784e3") => "LazyArtifacts",
        UUID("b27032c2-a3e7-50c8-80cd-2d36dbcbfd21") => "LibCURL",
        UUID("deac9b47-8bc7-5906-a0fe-35ac56dc84c0") => "LibCURL_jll",
        UUID("76f85450-5226-5b5a-8eaa-529ad045b433") => "LibGit2",
        UUID("e37daf67-58a4-590a-8e99-b0245dd2ffc5") => "LibGit2_jll",
        UUID("29816b5a-b9ab-546f-933c-edad1886dfa8") => "LibSSH2_jll",
        UUID("183b4373-6708-53ba-ad28-60e28bb38547") => "LibUV_jll",
        UUID("745a5e78-f969-53e9-954f-d19f2f74f4e3") => "LibUnwind_jll",
        UUID("8f399da3-3557-5675-b5ff-fb832c97cbdb") => "Libdl",
        UUID("37e2e46d-f89d-539d-b4ee-838fcccc9c8e") => "LinearAlgebra",
        UUID("56ddb016-857b-54e1-b83d-db4d58db5568") => "Logging",
        UUID("3a97d323-0669-5f0c-9066-3539efd106a3") => "MPFR_jll",
        UUID("d6f4376e-aef5-505a-96c1-9c027394607a") => "Markdown",
        UUID("c8ffd9c3-330d-5841-b78e-0817d7145fa1") => "MbedTLS_jll",
        UUID("a63ad114-7e13-5084-954f-fe012c677804") => "Mmap",
        UUID("14a3606d-f60d-562e-9121-12d972cd8159") => "MozillaCACerts_jll",
        UUID("ca575930-c2e3-43a9-ace4-1e988b2c1908") => "NetworkOptions",
        UUID("4536629a-c528-5b80-bd46-f80d51c5b363") => "OpenBLAS_jll",
        UUID("05823500-19ac-5b8b-9628-191a04bc5112") => "OpenLibm_jll",
        UUID("efcefdf7-47ab-520b-bdef-62a2eaa19f15") => "PCRE2_jll",
        UUID("44cfe95a-1eb2-52ea-b672-e2afdf69b78f") => "Pkg",
        UUID("de0858da-6303-5e67-8744-51eddeeeb8d7") => "Printf",
        UUID("9abbd945-dff8-562f-b5e8-e1ebf5ef1b79") => "Profile",
        UUID("3fa0cd96-eef1-5676-8a61-b3b8758bbffb") => "REPL",
        UUID("9a3f8284-a2c9-5f02-9a11-845980a1fd5c") => "Random",
        UUID("ea8e919c-243c-51af-8825-aaa63cd721ce") => "SHA",
        UUID("9e88b42a-f829-5b0c-bbe9-9e923198166b") => "Serialization",
        UUID("1a1011a3-84de-559e-8e89-a11a2f7dc383") => "SharedArrays",
        UUID("6462fe0b-24de-5631-8697-dd941f90decc") => "Sockets",
        UUID("2f01184e-e22b-5df5-ae63-d93ebab69eaf") => "SparseArrays",
        UUID("10745b16-79ce-11e8-11f9-7d13ad32a3b2") => "Statistics",
        UUID("4607b0f0-06f3-5cda-b6b1-a6196a1729e9") => "SuiteSparse",
        UUID("bea87d4a-7f5b-5778-9afe-8cc45184846c") => "SuiteSparse_jll",
        UUID("fa267f1f-6049-4f14-aa54-33bafae1ed76") => "TOML",
        UUID("a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e") => "Tar",
        UUID("8dfed614-e22c-5e08-85e1-65c5234f0b40") => "Test",
        UUID("cf7118a7-6976-5b1a-9a39-7adc72f591a4") => "UUIDs",
        UUID("4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5") => "Unicode",
        UUID("83775a58-1f1d-513f-b197-d71354ab007a") => "Zlib_jll",
        UUID("05ff407c-b0c1-5878-9df8-858cc2e60c36") => "dSFMT_jll",
        UUID("8f36deef-c2a5-5394-99ed-8e07531fb29a") => "libLLVM_jll",
        UUID("8e850b90-86db-534c-a0d3-1478176c7d93") => "libblastrampoline_jll",
        UUID("8e850ede-7688-5339-a07c-302acd2aaf8d") => "nghttp2_jll",
        UUID("3f19e933-33d8-53b3-aaab-bd5110c3b7a0") => "p7zip_jll",
    ),
]
# Next, we also embed a list of stdlibs that must _always_ be treated as stdlibs,
# because they cannot be resolved in the registry; they have only ever existed within
# the Julia stdlib source tree, and because of that, trying to resolve them will fail.
const UNREGISTERED_STDLIBS = Dict(
    UUID("2a0f44e3-6c83-55bd-87e4-b1978d98bd5f") => "Base64",
    UUID("8bf52ea8-c179-5cab-976a-9e18b702a9bc") => "CRC32c",
    UUID("ade2ca70-3891-5945-98fb-dc099432e06a") => "Dates",
    UUID("8bb1440f-4735-579b-a4ab-409b98df4dab") => "DelimitedFiles",
    UUID("8ba89e20-285c-5b6f-9357-94700520ee1b") => "Distributed",
    UUID("7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee") => "FileWatching",
    UUID("9fa8497b-333b-5362-9e8d-4d0656e87820") => "Future",
    UUID("b77e0a4c-d291-57a0-90e8-8db25a27a240") => "InteractiveUtils",
    UUID("76f85450-5226-5b5a-8eaa-529ad045b433") => "LibGit2",
    UUID("8f399da3-3557-5675-b5ff-fb832c97cbdb") => "Libdl",
    UUID("37e2e46d-f89d-539d-b4ee-838fcccc9c8e") => "LinearAlgebra",
    UUID("56ddb016-857b-54e1-b83d-db4d58db5568") => "Logging",
    UUID("d6f4376e-aef5-505a-96c1-9c027394607a") => "Markdown",
    UUID("a63ad114-7e13-5084-954f-fe012c677804") => "Mmap",
    UUID("44cfe95a-1eb2-52ea-b672-e2afdf69b78f") => "Pkg",
    UUID("de0858da-6303-5e67-8744-51eddeeeb8d7") => "Printf",
    UUID("9abbd945-dff8-562f-b5e8-e1ebf5ef1b79") => "Profile",
    UUID("3fa0cd96-eef1-5676-8a61-b3b8758bbffb") => "REPL",
    UUID("9a3f8284-a2c9-5f02-9a11-845980a1fd5c") => "Random",
    UUID("9e88b42a-f829-5b0c-bbe9-9e923198166b") => "Serialization",
    UUID("1a1011a3-84de-559e-8e89-a11a2f7dc383") => "SharedArrays",
    UUID("6462fe0b-24de-5631-8697-dd941f90decc") => "Sockets",
    UUID("2f01184e-e22b-5df5-ae63-d93ebab69eaf") => "SparseArrays",
    UUID("10745b16-79ce-11e8-11f9-7d13ad32a3b2") => "Statistics",
    UUID("8dfed614-e22c-5e08-85e1-65c5234f0b40") => "Test",
    UUID("cf7118a7-6976-5b1a-9a39-7adc72f591a4") => "UUIDs",
    UUID("4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5") => "Unicode",
)


text:
using Yao
using QuAlgorithmZoo
using KrylovKit

function ed_groundstate(h::MatrixBlock)
    E, V = eigsolve(h |> mat, 1, :SR, ishermitian=true)
    println("Ground State Energy is $(E[1])")
    ArrayReg(V[1])
end

N = 5
c = random_diff_circuit(N, N, [i=>mod(i,N)+1 for i=1:N], mode=:Merged) |> autodiff(:QC)
dispatch!(c, :random)
hami = heisenberg(N)

# vqe ground state
vqe_solve!(c, hami)


text:
module JacProp

export Trajectory, default_activations, ModelTrainer, ADModelTrainer, sample_jacprop, push!, train!, display_modeltrainer, getter

using LTVModelsBase, Parameters, Reexport, Lazy, Juno, FunctionEnsembles
@reexport using LTVModels, Flux, ValueHistories, IterTools, MLDataUtils
using Flux: back!, truncate!, treelike, train!, mse, testmode!, params, jacobian, throttle
using Flux.Optimise: Param, optimiser, RMSProp, expdecay
using DSP, Plots, DataStructures, ReverseDiff#, InteractNext, Observables
export plot, plot!, scatter, scatter!
const RDiff = ReverseDiff
const default_activations = [swish, Flux.sigmoid, tanh, elu]
const IT = IterTools
using Statistics, LinearAlgebra, Random

abstract type AbstractModelTrainer end

"""
ModelTrainer

- `models` A vector of functions
- `opts` A vector of optimizers
- `losses` A vector of loss functions
- `R1` Parameter drift covariance for Kalman smoother
- `R2` Measurement noise covariance
- `P0` Initial covariance matrix for Kalman smoother
- `P` Covariance magnification of NN prior
- `cb` can be either `()->()` or a a Function that creates a closure around `loss` and `data` (loss,data)->(()->())

See also [`LTVModels`](@ref),  [`LTVModels.KalmanModel`](@ref)
"""
@with_kw mutable struct ModelTrainer{cbT,tT} <: AbstractModelTrainer
    models
    opts
    losses
    R1                = I
    R2                = 10_000I
    P0                = 10_000I
    σdivider::Float64 = 10.
    P::Float64        = 1.0
    trajs::Vector{Trajectory} = Trajectory[]
    cb::cbT
    modelhistory = []
    trace::tT = History(Float64)
end

@with_kw mutable struct ADModelTrainer{mT,tT} <: AbstractModelTrainer
    model::mT
    opt
    λ::Float64 = 1.
    testdata::Trajectory
    trajs::Vector{Trajectory} = Trajectory[]
    modelhistory = []
    trace::tT = History(Float64)
    tracev::tT = History(Float64)
    normalizer::Matrix{Float64} = ones(1,1)
end

include("utilities.jl")

to_callback(cb,args...) = methods(cb).mt.max_args > 1 ? cb(args...) : cb

"""
train!(mt::ModelTrainer; epochs=1, jacprop=0, useprior=true, trace = mt.trace)

See also [`ModelTrainer`](@ref)
"""
function Flux.train!(mt::ModelTrainer; epochs=1, jacprop=0, useprior=true, trace = mt.trace)
    @assert !isempty(mt.trajs) "No data in ModelTrainer"
    @unpack models,opts,losses = mt
    data1 = todata(mt)
    jacprop > 0 && (ltvmodels = fit_models(mt, useprior))
    startepoch = length(trace)+1
    @progress for epoch = startepoch:startepoch+epochs-1
        data = if jacprop > 0
            data2 = [todata(sample_jacprop(mt, ltvmodels)) for i = 1:jacprop]
            chain(data1, data2...)
        else
            data1
        end
        for (model,loss, opt) in zip(models,losses,opts)
            train!(loss, data, opt, cb=to_callback(mt.cb, epoch,loss,data1, trace, model, mt))
        end
    end
    push!(mt.modelhistory, deepcopy(models))
    trace
end

function Flux.train!(mt::ADModelTrainer; epochs=1, jacprop=0, trace = mt.trace, cb=()->())
    @assert !isempty(mt.trajs) "No data in ModelTrainer"
    @unpack model,opt,testdata,λ,trace,tracev = mt
    # mt.λ = 0.
    if epochs <= 0
        push!(mt.modelhistory, deepcopy(model))
        return
    end
    w          = model.w
    data       = todata(mt)
    datat      = todata([testdata])
    mt.normalizer = ones(size(data[1][2], 1),size(data[1][1], 1))
    predfun = model isa ADSystem ? pred : predd
    f(x)   = predfun(w,x,model.nx)
    losses = map(data) do d
        x,y        = d
        lossfun    = loss2(w,x,y,mt)
        # lossfun    = (w...)->cost(w,mt.model.nx, x,y)
        # println("Compiling tape ")
        tape = RDiff.GradientTape(lossfun, w) |> RDiff.compile
        # println(" Done")
        increment!(trace, 1, lossfun(w...))
        increment!(tracev, 1, cost(f,datat))
        lossfun, tape
    end
    g          = similar.(w)
    # epochs > 0 && plot(reuse=false)
    startepoch = last(trace)[1]+1
    @progress for epoch = startepoch:startepoch+epochs-1
        # epoch == 500 && (mt.λ = λ)
        # lossfun = loss(w,x,y,mt)
        for (i,(lossfun,tape)) in enumerate(losses)
            RDiff.gradient!(g,tape,w)
            for (opti,gi) in zip(opt,g)
                opti(gi, epoch)
            end
            increment!(tracev, epoch, cost(f,datat))
            increment!(trace, epoch, lossfun(w...))
            if epoch % 5 == 0
                # mt.normalizer = find_normalizer(w,data[i],mt)
                # if i == length(losses) && Distributed.myid() == 1
                #     cb(model)
                #     # plot(trace, reuse=true)
                #     # plot!(tracev)
                #     # gui()
                      # println("Losses: ", last(trace), last(tracev)[2])
                # end
            end
        end
    end
    push!(mt.modelhistory, deepcopy(model))
    trace, tracev
end

function fit_models(mt::ModelTrainer, useprior)
    map(mt.trajs) do traj
        ltvmodel = KalmanModel(mt, traj, useprior=useprior)
    end
end

function sample_jacprop(mt::ModelTrainer, ltvmodels)
    perturbed_trajs = map(mt.trajs, ltvmodels) do traj, ltvmodel
        @unpack x,u,nx,nu = traj
        xa = x .+ std(x,2)/mt.σdivider .* randn(size(x))
        ua = u .+ std(u,2)/mt.σdivider .* randn(size(u))
        ya = LTVModels.predict(ltvmodel, xa, ua)
        Trajectory(xa,ua,ya)
    end
end

function LTVModels.KalmanModel(mt::ModelTrainer, t::Trajectory, ms=mt.models;
    P=mt.P, useprior=!isempty(mt.modelhistory))
    useprior = useprior && !isempty(mt.modelhistory)
    @unpack x,u,nx,nu = t
    @unpack R1,R2,P0  = mt
    T                 = length(t)
    if useprior
        model    = KalmanModel(zeros(nx,nx,T),zeros(nx,nu,T),zeros(1,1,T),false)
        Jm, Js   = jacobians(ms, t)
        Pt       = cat(3,[diagm(Js[:,i].^2 .+ P) for i=1:T]...) # TODO: magic number
        fx       = cat(3,[reshape(Jm[1:nx^2,i], nx,nx) for i=1:T]...)
        fu       = cat(3,[reshape(Jm[nx^2+1:end,i], nx,nu) for i=1:T]...)
        prior    = KalmanModel(fx,fu,Pt,false)
        ltvmodel = KalmanModel(model, prior, x,u,R1,R2,P0, extend = true, printfit=false)
    else # Don't use prior if it's the first time training
        ltvmodel = KalmanModel(x,u,R1,R2,P0, extend = true, printfit=false)
    end
    ltvmodel
end

todata(trajs::Vector{Trajectory}) = [(traj.xu,traj.y) for traj in trajs]
todata(mt::AbstractModelTrainer) = todata(mt.trajs)
function todata2(trajs::Vector{Trajectory})
    xu = hcat(getfield.(trajs,:xu)...)
    y = hcat(getfield.(trajs,:y)...)
    xu, y
end
todata2(mt::AbstractModelTrainer) = todata2(mt.trajs)

Base.push!(mt::AbstractModelTrainer, t::Trajectory) = push!(mt.trajs, t)
Base.push!(mt::AbstractModelTrainer, data::Matrix...) = push!(mt.trajs, Trajectory(data...))

for T in (ModelTrainer, ADModelTrainer)
    @eval (mt::($T))(; kwargs...) = train!(mt; kwargs...)
    @eval function (mt::($T))(t::Trajectory; kwargs...)
        push!(mt,t)
        train!(mt; kwargs...)
    end
    @eval (mt::($T))(data::Matrix...; kwargs...) = mt(Trajectory(data...); kwargs...)
    @eval Lazy.@forward ($T).models predict, simulate, Flux.jacobian, eigvalplot, jacplot
end

include("plots.jl")

end # module


text:
using Distributed
# using ClusterManagers
using Logging

# logger = SimpleLogger(stdout, Logging.Debug)
# old_logger = global_logger(logger)

try
    num_tasks = parse(Int, ENV["SLURM_NTASKS"])
    @info "Number of slurm tasks: $num_tasks"
#    addprocs(SlurmManager(num_tasks), unbuffered="")#, topology=:master_worker)
catch KeyError
    @warn "Not inside a slurm job"
end

@info "Number of processes $(nprocs())"

@everywhere begin
    using Pkg
    Pkg.activate(".")
end

@everywhere using QContinuousMeasurement
@everywhere using TimerOutputs

using HDF5
using DelimitedFiles
using Logging

using ArgParse
using ProgressMeter

# Parse the arguments passed to the script
s = ArgParseSettings()

@add_arg_table! s begin
    "--Nj"
        help = "Number of spins"
        default = 1
        arg_type = Int64
    "--Ntraj"
        help = "Number of trajectories"
        default = 1
        arg_type = Int64
    "--Tfinal", "-T"
        help = "Final time"
        arg_type = Float64
        default = 1.0
    "--dt"
        help = "dt"
        arg_type = Float64
        default = 0.0001
    "--kind"
        help = "Independent noise rate k_ind"
        arg_type = Float64
        default = 1.0
    "--kcoll"
        help = "Collective noise rate k_coll"
        arg_type = Float64
        default = 0.0
    "--Gamma"
        help = "Monitoring noise rate Gamma"
        arg_type = Float64
        default = 1.0
    "--omega"
        help = "Hamiltonian frequency omega"
        arg_type = Float64
        default = 1.0
    "--eta"
        help = "Measurement efficiency"
        arg_type = Float64
        default = 1.0
    "--mu"
        help = "Initial squeezing parameter"
        arg_type = Float64
        default = 0.0
    "--outpoints"
        help = "Number of output points"
        arg_type = Int64
        default = 200
    "--liouvillianfile"
        help = "npz file with the liovuillian data"
        arg_type = String
        default = nothing
    "--outdir"
        help = "Output directory"
        arg_type = String
        default = "."
end

args = parse_args(s)

jobid = "SLURM_JOB_ID" in keys(ENV) ? jobid = ENV["SLURM_JOB_ID"] : ""
filename = string("$(args["outdir"])/sim_Nj_$(args["Nj"])_Ntraj_$(args["Ntraj"])_Tf_$(args["Tfinal"])_" *
            "dt_$(args["dt"])_kind_$(args["kind"])_kcoll_$(args["kcoll"])_Gamma_$(args["Gamma"])_" *
            "omega_$(args["omega"])_eta_$(args["eta"])_mu_$(args["mu"])_$(jobid)")

let pkgstr = ""
    for (uuid, pkg) in Pkg.dependencies()
        if pkg.is_direct_dep
            pkgstr *= "\t$(pkg.name) $(pkg.version)"
            if pkg.is_tracking_repo
                pkgstr *= " #$(pkg.git_revision) ($(pkg.git_source))"
            elseif pkg.is_tracking_path
                pkgstr *= " $(pkg.source)"
            end
            pkgstr *= "\n"
        end
    end
    @info "Installed packages\n$pkgstr"
end

@info "Output filename: $filename"

Ntraj = args["Ntraj"]

@everywhere to = TimerOutput()

modelparams = ModelParameters(Nj=args["Nj"],
                              kind=args["kind"],
                              Gamma=args["Gamma"],
                              omega=args["omega"],
                              eta=args["eta"],
                              dt=args["dt"],
                              Tfinal=args["Tfinal"],
                              outpoints=args["outpoints"],
                              mu=args["mu"])

@info "Output every $(modelparams._outsteps) steps"

@info "Initializing model..."

init_time = @elapsed begin
    if args["kind"] == 0.0
        @everywhere model = CollectiveDephasingModel($modelparams)
        @everywhere initial_state = fixedj_sss($modelparams.Nj, $modelparams.mu)
    else
        @assert args["kcoll"] == 0.0
        @everywhere model = LocalDephasingModel($modelparams, $args["liouvillianfile"])
        @everywhere initial_state = blockdiag_sss($modelparams.Nj, $modelparams.mu)
    end
end

xi = squeezing_param(initial_state)
@info "Initial spin squeezing $(xi)"
# println(model)
# println(initial_state)

@info "Done in $init_time seconds..."

# Opens the FileWriter
writer = FileWriter("$filename.h5", modelparams, args["Ntraj"], ["FI", "QFI", "FIstrong", "xi2y", "Jx"])

progress_channel = RemoteChannel(() -> Channel{Bool}(1000))

p = Progress(Ntraj * length(get_time(model)), barglyphs=BarGlyphs("[=> ]"), dt=60, barlen=10)

@async while take!(progress_channel)
    next!(p)
end

if "JULIA_NUM_THREADS" in keys(ENV)
    @everywhere numthreads = $(parse(Int, ENV["JULIA_NUM_THREADS"]))
    @info("Number of threads: $numthreads")
else
    @info("JULIA_NUM_THREADS not set, using 1 thread")
    @everywhere numthreads = 1
end

# @everywhere begin
#     numthreads = 1
#     try
#         if "JULIA_NUM_THREADS" in keys(ENV)
#             numthreads = parse(Int, ENV["JULIA_NUM_THREADS"])
#             @info("Number of threads: $numthreads")
#         else
#             @info("JULIA_NUM_THREADS not set, using 1 thread")
#             numthreads = 1
#         end
#     catch
#         @info "Uffa!"
#     end
# end

@everywhere function thread_simulate_trajectory(model, initial_state, ntraj, file_channel, progress_channel)
    pid = Distributed.myid()
    nth = Threads.nthreads()
    Threads.@threads for i in 1:ntraj
        tid = Threads.threadid()

        simulate_trajectory(model, initial_state, file_channel, progress_channel)
        @debug "Trajectory done from thread $tid of $nth on worker $pid."
    end

    @debug "Batch of $ntraj done!"
end

file_channel = writer.channel

function prepare_batches(n, batch_size)
    if n % batch_size == 0
        return [batch_size for i = 1:(n ÷ batch_size)]
    else
        m = Int(floor(n / batch_size))
        vcat([[batch_size for i = 1:m], n - (m * batch_size)]...)
    end
end

trajectory_time = @elapsed begin
   pmap(x -> thread_simulate_trajectory(model, initial_state, x, file_channel, progress_channel), prepare_batches(Ntraj, numthreads), on_error=identity)
end

@info "pmap done"
put!(progress_channel, false)

@info "Trajectory simulation time: $trajectory_time"

for i in eachindex(workers())
    fetch(@spawnat i begin
        try
            result = read(`grep VmHWM /proc/$(getpid())/status`, String)
            peakmem = tryparse(Int, String(match(r"(\d+)", result)[1]))

            @info "Peak memory in GB: $(peakmem / 1024^2)"
        catch er
            @info er
        end
    end)
end

close(writer)