text:
using HePPCAT
using ForwardDiff, LinearAlgebra, StableRNGs, Test

# Internal functions
using HePPCAT: homppca
using HePPCAT: ExpectationMaximization, DifferenceOfConcave,
    MinorizeMaximize, ProjectedGradientAscent, RootFinding, StiefelGradientAscent,
    QuadraticSolvableMinorizer, CubicSolvableMinorizer,
    QuadraticMinorizer, OptimalQuadraticMinorizer
using HePPCAT: ArmijoSearch, InverseLipschitz
using HePPCAT: LipBoundU1, LipBoundU2, loglikelihood
using HePPCAT: updateF!, updatev!, updateU!, updateλ!

# Flag to use automatic differentiation
const TEST_WITH_AD = get(ENV,"HEPPCAT_TEST_WITH_AD","false") == "true"

# Load reference implementations
include("ref.jl")

# Convenience functions
factormatrix(M::HePPCATModel) = M.U*sqrt(Diagonal(M.λ))*M.Vt
flatten(M::HePPCATModel,n) = HePPCATModel(M.U,M.λ,M.Vt,vcat(fill.(M.v,n)...))

# Test HePPCATModel type
@testset "HePPCATModel type" begin
    rng = StableRNG(123)
    d, k, v = 10, 3, [4.0,2.0]
    F = randn(rng,d,k)
    U, s, V = svd(F)
    λ = s.^2

    @test HePPCATModel(U,λ,V',v) == HePPCATModel(U,λ,V',Int.(v))
    @test HePPCATModel(U,λ,Matrix{Float64}(I,k,k),v) == HePPCATModel(U,λ,I(k),v)
    @test HePPCATModel(U,λ,V',v) == HePPCATModel(F,v)
    @test HePPCATModel(U,λ,V',v).F ≈ F
end

# Test all updates
T = 5
n, v = (40, 10), (4, 1)
@testset "n=$(n[1:L]), v=$(v[1:L]), d=$d, k=$k" for L = 1:2, d = 5:10:25, k = 1:3
    # Compute test problem
    rng = StableRNG(123)
    F, Z = randn(rng,d,k), [randn(rng,k,n[l]) for l in 1:L]
    Yb = [F*Z[l] + sqrt(v[l])*randn(rng,d,n[l]) for l in 1:L]   # blocked
    Yf = reshape.(collect.(eachcol(hcat(Yb...))),:,1)           # flatten

    @testset "homoscedastic init" begin
        Ur, λr, vr = Ref.homppca(Yb,k)
        Mb = homppca(Yb,k)
        @test Ur*Diagonal(λr)*Ur' ≈ Mb.U*Diagonal(Mb.λ)*Mb.U'
        @test length(Mb.v) == length(Yb)
        @test all(vr .≈ Mb.v)
        Mf = homppca(Yf,k)
        @test Ur*Diagonal(λr)*Ur' ≈ Mf.U*Diagonal(Mf.λ)*Mf.U'
        @test length(Mf.v) == length(Yf)
        @test all(vr .≈ Mf.v)
    end

    @testset "overall function" begin
        Mr = homppca(Yb,k)
        for _ in 1:T
            updatev!(Mr,Yb,ExpectationMaximization())
            updateF!(Mr,Yb,ExpectationMaximization())
        end
        @test Mr !== heppcat(Yb,k,T)
        @test Mr == heppcat(Yb,k,T)
    end

    @testset "block calls" begin
        # Generate sequence of test iterates
        MM = [HePPCATModel(randn(rng,d,k),rand(rng,L))]
        for t in 1:T
            push!(MM, deepcopy(MM[end]))
            updateF!(MM[end],Yb,ExpectationMaximization())
            updatev!(MM[end],Yb,RootFinding())
        end

        # Test v updates
        @testset "updatev! (RootFinding): t=$t" for t in 1:T
            vr = Ref.updatev_roots(MM[t].U,MM[t].λ,MM[t].v,Yb)
            Mb = updatev!(deepcopy(MM[t]),Yb,RootFinding())
            @test vr ≈ Mb.v
        end
        @testset "updatev! (ExpectationMaximization): t=$t" for t in 1:T
            vr = Ref.updatev_em(MM[t].U,MM[t].λ,MM[t].v,Yb)
            Mb = updatev!(deepcopy(MM[t]),Yb,ExpectationMaximization())
            @test vr ≈ Mb.v
        end
        @testset "updatev! (DifferenceOfConcave): t=$t" for t in 1:T
            vr = Ref.updatev_doc(MM[t].U,MM[t].λ,MM[t].v,Yb)
            Mb = updatev!(deepcopy(MM[t]),Yb,DifferenceOfConcave())
            @test vr ≈ Mb.v
        end
        @testset "updatev! (QuadraticSolvableMinorizer): t=$t" for t in 1:T
            vr = Ref.updatev_mm_quad(MM[t].U,MM[t].λ,MM[t].v,Yb)
            Mb = updatev!(deepcopy(MM[t]),Yb,QuadraticSolvableMinorizer())
            @test vr ≈ Mb.v
        end
        @testset "updatev! (CubicSolvableMinorizer): t=$t" for t in 1:T
            vr = Ref.updatev_mm_cubic(MM[t].U,MM[t].λ,MM[t].v,Yb)
            Mb = updatev!(deepcopy(MM[t]),Yb,CubicSolvableMinorizer())
            @test vr ≈ Mb.v
        end

        # Test F updates
        @testset "updateF! (ExpectationMaximization): t=$t" for t in 1:T
            Fr = Ref.updateF_em(factormatrix(MM[t]),MM[t].v,Yb)
            Mb = updateF!(deepcopy(MM[t]),Yb,ExpectationMaximization())
            @test Fr ≈ factormatrix(Mb)
            Mf = updateF!(flatten(deepcopy(MM[t]),n[1:L]),Yf,ExpectationMaximization())
            @test Fr ≈ factormatrix(Mf)
        end

        # Test U updates
        @testset "updateU! (MinorizeMaximize): t=$t" for t in 1:T
            Ur = Ref.updateU_mm(MM[t].U,MM[t].λ,MM[t].v,Yb)
            Mb = updateU!(deepcopy(MM[t]),Yb,MinorizeMaximize())
            @test Ur ≈ Mb.U
            Mf = updateU!(flatten(deepcopy(MM[t]),n[1:L]),Yf,MinorizeMaximize())
            @test Ur ≈ Mf.U
        end
        @testset "updateU! (ProjectedGradientAscent, $stepsize): t=$t" for stepsize in [0.0,0.5,Inf], t in 1:T
            Ur = Ref.updateU_pga(MM[t].U,MM[t].λ,MM[t].v,Yb,stepsize)
            Mb = updateU!(deepcopy(MM[t]),Yb,ProjectedGradientAscent(stepsize))
            @test Ur ≈ Mb.U
            Mf = updateU!(flatten(deepcopy(MM[t]),n[1:L]),Yf,ProjectedGradientAscent(stepsize))
            @test Ur ≈ Mf.U
        end
        @testset "updateU! (ProjectedGradientAscent, Lipschitz): t=$t" for t in 1:T
            Lip = Ref.LipBoundU2(MM[t].λ,MM[t].v,Yb)
            Ur = Ref.updateU_pga(MM[t].U,MM[t].λ,MM[t].v,Yb,1/Lip)
            Mb = updateU!(deepcopy(MM[t]),Yb,ProjectedGradientAscent(InverseLipschitz()))
            @test Ur ≈ Mb.U
            # Mf = updateU!(flatten(deepcopy(MM[t]),n[1:L]),Yf,ProjectedGradientAscent(InverseLipschitz()))  # LipBound2 depends on
            # @test Ur ≈ Mf.U                                                                                # how the blocks are done
        end
        @testset "updateU! (ProjectedGradientAscent, Armijo): t=$t" for t in 1:T
            Ur = Ref.updateU_pga_armijo(MM[t].U,MM[t].λ,MM[t].v,Yb,50,0.8,0.5,1e-4)
            Mb = updateU!(deepcopy(MM[t]),Yb,ProjectedGradientAscent(ArmijoSearch(50,0.8,0.5,1e-4)))
            @test Ur ≈ Mb.U
            Mf = updateU!(flatten(deepcopy(MM[t]),n[1:L]),Yf,ProjectedGradientAscent(ArmijoSearch(50,0.8,0.5,1e-4)))
            @test Ur ≈ Mf.U
        end
        @testset "updateU! (StiefelGradientAscent): t=$t" for t in 1:T
            Ur = Ref.updateU_sga(MM[t].U,MM[t].λ,MM[t].v,Yb,50,0.8,0.5,0.5)
            Mb = updateU!(deepcopy(MM[t]),Yb,StiefelGradientAscent(ArmijoSearch(50,0.8,0.5,0.5)))
            @test Ur ≈ Mb.U
            Mf = updateU!(flatten(deepcopy(MM[t]),n[1:L]),Yf,StiefelGradientAscent(ArmijoSearch(50,0.8,0.5,0.5)))
            @test Ur ≈ Mf.U
        end

        # Test λ updates
        @testset "updateλ! (RootFinding): t=$t" for t in 1:T
            λr = Ref.updateλ_roots(MM[t].U,MM[t].v,Yb)
            Mb = updateλ!(deepcopy(MM[t]),Yb,RootFinding())
            @test λr ≈ Mb.λ
            Mf = updateλ!(flatten(deepcopy(MM[t]),n[1:L]),Yf,RootFinding())
            @test λr ≈ Mf.λ
        end
        @testset "updateλ! (ExpectationMaximization): t=$t" for t in 1:T
            λr = Ref.updateλ_em(MM[t].λ,MM[t].U,MM[t].v,Yb)
            Mb = updateλ!(deepcopy(MM[t]),Yb,ExpectationMaximization())
            @test λr ≈ Mb.λ
            Mf = updateλ!(flatten(deepcopy(MM[t]),n[1:L]),Yf,ExpectationMaximization())
            @test λr ≈ Mf.λ
        end
        @testset "updateλ! (MinorizeMaximize): t=$t" for t in 1:T
            λr = Ref.updateλ_mm(MM[t].λ,MM[t].U,MM[t].v,Yb)
            Mb = updateλ!(deepcopy(MM[t]),Yb,MinorizeMaximize())
            @test λr ≈ Mb.λ
            Mf = updateλ!(flatten(deepcopy(MM[t]),n[1:L]),Yf,MinorizeMaximize())
            @test λr ≈ Mf.λ
        end
        @testset "updateλ! (QuadraticMinorizer): t=$t" for t in 1:T
            λr = Ref.updateλ_quad(MM[t].λ,MM[t].U,MM[t].v,Yb)
            Mb = updateλ!(deepcopy(MM[t]),Yb,QuadraticMinorizer())
            @test λr ≈ Mb.λ
            Mf = updateλ!(flatten(deepcopy(MM[t]),n[1:L]),Yf,QuadraticMinorizer())
            @test λr ≈ Mf.λ
        end
        @testset "updateλ! (DifferenceOfConcave): t=$t" for t in 1:T
            λr = Ref.updateλ_doc(MM[t].λ,MM[t].U,MM[t].v,Yb)
            Mb = updateλ!(deepcopy(MM[t]),Yb,DifferenceOfConcave())
            @test λr ≈ Mb.λ
            Mf = updateλ!(flatten(deepcopy(MM[t]),n[1:L]),Yf,DifferenceOfConcave())
            @test λr ≈ Mf.λ
        end
        @testset "updateλ! (OptimalQuadraticMinorizer): t=$t" for t in 1:T
            λr = Ref.updateλ_opt_quad(MM[t].λ,MM[t].U,MM[t].v,Yb)
            Mb = updateλ!(deepcopy(MM[t]),Yb,OptimalQuadraticMinorizer())
            @test λr ≈ Mb.λ
            # Mf = updateλ!(flatten(deepcopy(MM[t]),n[1:L]),Yf,OptimalQuadraticMinorizer())  # Optimal curvature depends on
            # @test λr ≈ Mf.λ                                                                # how the blocks are done
        end

        # Test F/gradF
        @testset "F/gradF: t=$t" for t in 1:T
            vf = vcat(fill.(MM[t].v,n[1:L])...)
            Fr = Ref.F(MM[t].U,MM[t].λ,vf,Yf)
            Fb = HePPCAT.F(MM[t].U,MM[t].λ,MM[t].v,Yb)
            @test Fr ≈ Fb
            Ff = HePPCAT.F(MM[t].U,MM[t].λ,vf,Yf)
            @test Fr ≈ Ff

            LLr = [Ref.loglikelihood(U,MM[t].λ,vf,Yf) for U in getfield.(MM,:U)]
            Fb = [HePPCAT.F(U,MM[t].λ,MM[t].v,Yb) for U in getfield.(MM,:U)]
            @test LLr ≈ Fb .+ (LLr[1] - Fb[1])    # "values match up to a constant w.r.t U"

            Gr = Ref.gradF(MM[t].U,MM[t].λ,vf,Yf)
            Gb = HePPCAT.gradF(MM[t].U,MM[t].λ,MM[t].v,Yb)
            @test Gr ≈ Gb
            Gf = HePPCAT.gradF(MM[t].U,MM[t].λ,vf,Yf)
            @test Gr ≈ Gf

            if TEST_WITH_AD
                GadLL = ForwardDiff.gradient(U -> Ref.loglikelihood(U,MM[t].λ,vf,Yf),MM[t].U)
                @test GadLL ≈ Gb
                GadF = ForwardDiff.gradient(U -> Ref.F(U,MM[t].λ,vf,Yf),MM[t].U)
                @test GadF ≈ Gb
            end
        end

        # Test Lipschitz bound 1 w.r.t U
        @testset "LipBoundU1: t=$t" for t in 1:T
            vf = vcat(fill.(MM[t].v,n[1:L])...)
            Lipr = Ref.LipBoundU1(MM[t].λ,vf,Yf)
            Lipb = LipBoundU1(MM[t],Yb)
            @test Lipr ≈ Lipb
            Lipf = LipBoundU1(flatten(MM[t],n[1:L]),Yf)
            @test Lipr ≈ Lipf

            if TEST_WITH_AD
                Hess = ForwardDiff.hessian(U -> Ref.F(U,MM[t].λ,vf,Yf),MM[t].U)
                @test (1-1e-12)*opnorm(Hess) <= Lipb
            end
        end

        # Test Lipschitz bound 2 w.r.t U
        @testset "LipBoundU2: t=$t" for t in 1:T
            Lipr = Ref.LipBoundU2(MM[t].λ,MM[t].v,Yb)
            Lipb = LipBoundU2(MM[t],Yb)
            @test Lipr ≈ Lipb

            if TEST_WITH_AD
                Hess = ForwardDiff.hessian(U -> Ref.F(U,MM[t].λ,MM[t].v,Yb),MM[t].U)
                @test (1-1e-12)*opnorm(Hess) <= Lipb
            end
        end

        # Test log-likelihood
        @testset "loglikelihood: t=$t" for t in 1:T
            vf = vcat(fill.(MM[t].v,n[1:L])...)
            LLr = Ref.loglikelihood(MM[t].U,MM[t].λ,vf,Yf)
            LLb = loglikelihood(MM[t],Yb)
            @test LLr ≈ LLb
            LLf = loglikelihood(flatten(MM[t],n[1:L]),Yf)
            @test LLr ≈ LLf
        end
    end

    @testset "flat calls" begin
        # Generate sequence of test iterates
        MM = [HePPCATModel(randn(rng,d,k),rand(rng,sum(n[1:L])))]
        for t in 1:T
            push!(MM, deepcopy(MM[end]))
            updateF!(MM[end],Yf,ExpectationMaximization())
            updatev!(MM[end],Yf,RootFinding())
        end

        # Test v updates
        @testset "updatev! (RootFinding): t=$t" for t in 1:T
            vr = Ref.updatev_roots(MM[t].U,MM[t].λ,MM[t].v,Yf)
            Mf = updatev!(deepcopy(MM[t]),Yf,RootFinding())
            @test vr ≈ Mf.v
        end
        @testset "updatev! (ExpectationMaximization): t=$t" for t in 1:T
            vr = Ref.updatev_em(MM[t].U,MM[t].λ,MM[t].v,Yf)
            Mf = updatev!(deepcopy(MM[t]),Yf,ExpectationMaximization())
            @test vr ≈ Mf.v
        end
        @testset "updatev! (DifferenceOfConcave): t=$t" for t in 1:T
            vr = Ref.updatev_doc(MM[t].U,MM[t].λ,MM[t].v,Yf)
            Mf = updatev!(deepcopy(MM[t]),Yf,DifferenceOfConcave())
            @test vr ≈ Mf.v
        end
        @testset "updatev! (QuadraticSolvableMinorizer): t=$t" for t in 1:T
            vr = Ref.updatev_mm_quad(MM[t].U,MM[t].λ,MM[t].v,Yf)
            Mf = updatev!(deepcopy(MM[t]),Yf,QuadraticSolvableMinorizer())
            @test vr ≈ Mf.v
        end
        @testset "updatev! (CubicSolvableMinorizer): t=$t" for t in 1:T
            vr = Ref.updatev_mm_cubic(MM[t].U,MM[t].λ,MM[t].v,Yf)
            Mf = updatev!(deepcopy(MM[t]),Yf,CubicSolvableMinorizer())
            @test vr ≈ Mf.v
        end

        # Test F updates
        @testset "updateF! (ExpectationMaximization): t=$t" for t in 1:T
            Fr = Ref.updateF_em(factormatrix(MM[t]),MM[t].v,Yf)
            Mf = updateF!(deepcopy(MM[t]),Yf,ExpectationMaximization())
            @test Fr ≈ factormatrix(Mf)
        end

        # Test U updates
        @testset "updateU! (MinorizeMaximize): t=$t" for t in 1:T
            Ur = Ref.updateU_mm(MM[t].U,MM[t].λ,MM[t].v,Yf)
            Mf = updateU!(deepcopy(MM[t]),Yf,MinorizeMaximize())
            @test Ur ≈ Mf.U
        end
        @testset "updateU! (ProjectedGradientAscent, $stepsize): t=$t" for stepsize in [0.0,0.5,Inf], t in 1:T
            Ur = Ref.updateU_pga(MM[t].U,MM[t].λ,MM[t].v,Yf,stepsize)
            Mf = updateU!(deepcopy(MM[t]),Yf,ProjectedGradientAscent(stepsize))
            @test Ur ≈ Mf.U
        end
        @testset "updateU! (ProjectedGradientAscent, Lipschitz): t=$t" for t in 1:T
            Lip = Ref.LipBoundU2(MM[t].λ,MM[t].v,Yf)
            Ur = Ref.updateU_pga(MM[t].U,MM[t].λ,MM[t].v,Yf,1/Lip)
            Mf = updateU!(deepcopy(MM[t]),Yf,ProjectedGradientAscent(InverseLipschitz()))
            @test Ur ≈ Mf.U
        end
        @testset "updateU! (ProjectedGradientAscent, Armijo): t=$t" for t in 1:T
            Ur = Ref.updateU_pga_armijo(MM[t].U,MM[t].λ,MM[t].v,Yf,100,0.8,0.5,1e-4)
            Mf = updateU!(deepcopy(MM[t]),Yf,ProjectedGradientAscent(ArmijoSearch(100,0.8,0.5,1e-4)))
            @test Ur ≈ Mf.U
        end
        @testset "updateU! (StiefelGradientAscent): t=$t" for t in 1:T
            Ur = Ref.updateU_sga(MM[t].U,MM[t].λ,MM[t].v,Yf,50,0.8,0.5,0.5)
            Mf = updateU!(deepcopy(MM[t]),Yf,StiefelGradientAscent(ArmijoSearch(50,0.8,0.5,0.5)))
            @test Ur ≈ Mf.U
        end

        # Test λ updates
        @testset "updateλ! (RootFinding): t=$t" for t in 1:T
            λr = Ref.updateλ_roots(MM[t].U,MM[t].v,Yf)
            Mf = updateλ!(deepcopy(MM[t]),Yf,RootFinding())
            @test λr ≈ Mf.λ
        end
        @testset "updateλ! (ExpectationMaximization): t=$t" for t in 1:T
            λr = Ref.updateλ_em(MM[t].λ,MM[t].U,MM[t].v,Yf)
            Mf = updateλ!(deepcopy(MM[t]),Yf,ExpectationMaximization())
            @test λr ≈ Mf.λ
        end
        @testset "updateλ! (MinorizeMaximize): t=$t" for t in 1:T
            λr = Ref.updateλ_mm(MM[t].λ,MM[t].U,MM[t].v,Yf)
            Mf = updateλ!(deepcopy(MM[t]),Yf,MinorizeMaximize())
            @test λr ≈ Mf.λ
        end
        @testset "updateλ! (QuadraticMinorizer): t=$t" for t in 1:T
            λr = Ref.updateλ_quad(MM[t].λ,MM[t].U,MM[t].v,Yf)
            Mf = updateλ!(deepcopy(MM[t]),Yf,QuadraticMinorizer())
            @test λr ≈ Mf.λ
        end
        @testset "updateλ! (DifferenceOfConcave): t=$t" for t in 1:T
            λr = Ref.updateλ_doc(MM[t].λ,MM[t].U,MM[t].v,Yf)
            Mf = updateλ!(deepcopy(MM[t]),Yf,DifferenceOfConcave())
            @test λr ≈ Mf.λ
        end
        @testset "updateλ! (OptimalQuadraticMinorizer): t=$t" for t in 1:T
            λr = Ref.updateλ_opt_quad(MM[t].λ,MM[t].U,MM[t].v,Yf)
            Mf = updateλ!(deepcopy(MM[t]),Yf,OptimalQuadraticMinorizer())
            @test λr ≈ Mf.λ
        end

        # Test F/gradF
        @testset "F/gradF: t=$t" for t in 1:T
            Fr = Ref.F(MM[t].U,MM[t].λ,MM[t].v,Yf)
            Ff = HePPCAT.F(MM[t].U,MM[t].λ,MM[t].v,Yf)
            @test Fr ≈ Ff

            LLr = [Ref.loglikelihood(U,MM[t].λ,MM[t].v,Yf) for U in getfield.(MM,:U)]
            Ff = [HePPCAT.F(U,MM[t].λ,MM[t].v,Yf) for U in getfield.(MM,:U)]
            @test LLr ≈ Ff .+ (LLr[1] - Ff[1])    # "values match up to a constant w.r.t U"

            Gr = Ref.gradF(MM[t].U,MM[t].λ,MM[t].v,Yf)
            Gf = HePPCAT.gradF(MM[t].U,MM[t].λ,MM[t].v,Yf)
            @test Gr ≈ Gf

            if TEST_WITH_AD
                GadLL = ForwardDiff.gradient(U -> Ref.loglikelihood(U,MM[t].λ,MM[t].v,Yf),MM[t].U)
                @test GadLL ≈ Gf
                GadF = ForwardDiff.gradient(U -> Ref.F(U,MM[t].λ,MM[t].v,Yf),MM[t].U)
                @test GadF ≈ Gf
            end
        end

        # Test Lipschitz bound 1 w.r.t U
        @testset "LipBoundU1: t=$t" for t in 1:T
            Lipr = Ref.LipBoundU1(MM[t].λ,MM[t].v,Yf)
            Lipf = LipBoundU1(MM[t],Yf)
            @test Lipr ≈ Lipf

            if TEST_WITH_AD
                Hess = ForwardDiff.hessian(U -> Ref.F(U,MM[t].λ,MM[t].v,Yf),MM[t].U)
                @test (1-1e-12)*opnorm(Hess) <= Lipf
            end
        end

        # Test Lipschitz bound 2 w.r.t U
        @testset "LipBoundU2: t=$t" for t in 1:T
            Lipr = Ref.LipBoundU2(MM[t].λ,MM[t].v,Yf)
            Lipf = LipBoundU2(MM[t],Yf)
            @test Lipr ≈ Lipf

            if TEST_WITH_AD
                Hess = ForwardDiff.hessian(U -> Ref.F(U,MM[t].λ,MM[t].v,Yf),MM[t].U)
                @test (1-1e-12)*opnorm(Hess) <= Lipf
            end
        end

        # Test log-likelihood
        @testset "loglikelihood: t=$t" for t in 1:T
            LLr = Ref.loglikelihood(MM[t].U,MM[t].λ,MM[t].v,Yf)
            LLf = loglikelihood(MM[t],Yf)
            @test LLr ≈ LLf
        end
    end
end

# Test skew-symmetry on several matrices
@testset "skew" begin
    rng = StableRNG(123)
    @testset "random $n x $n matrix" for n in 1:20
        A = HePPCAT.skew(randn(rng,n,n))
        @test A == HePPCAT.skew(A)
        @test A == -A'
    end
end

# Test orthonormality of StiefelGradientAscent updates
@testset "StiefelGradientAscent orthonormality" begin
    rng = StableRNG(123)
    d, λ, n, v = 100, [4.,2.], [200,200], [0.2,0.4]
    k = length(λ)
    U = svd(randn(rng,d,k)).U
    F = U*sqrt(Diagonal(λ))
    Y = [F*randn(rng,k,nl) + sqrt(vl)*randn(rng,d,nl) for (nl,vl) in zip(n,v)]
    H = HePPCATModel(svd(randn(rng,d,k)).U,λ,Matrix{Float64}(I,k,k),v)
    @testset "iterate $t" for t in 1:200
        updateU!(H,Y,StiefelGradientAscent(ArmijoSearch(10,0.15,0.5,0.005)))
        @test H.U'*H.U ≈ I
    end
end

# Test StiefelGradientAscent line search warning
@testset "StiefelGradientAscent line search warning" begin
    rng = StableRNG(123)
    d, λ, n, v = 100, [4.,2.], [200,200], [0.2,0.4]
    k = length(λ)
    U = svd(randn(rng,d,k)).U
    F = U*sqrt(Diagonal(λ))
    Y = [F*randn(rng,k,nl) + sqrt(vl)*randn(rng,d,nl) for (nl,vl) in zip(n,v)]
    H = HePPCATModel(svd(randn(rng,d,k)).U,λ,Matrix{Float64}(I,k,k),v)
    @test_logs (:warn, "Exceeded maximum line search iterations. Accuracy not guaranteed.") updateU!(deepcopy(H),Y,StiefelGradientAscent(ArmijoSearch(0,0.15,0.5,0.005)))
    @test_logs (:warn, "Exceeded maximum line search iterations. Accuracy not guaranteed.") updateU!(deepcopy(H),Y,StiefelGradientAscent(ArmijoSearch(2,10.0,0.5,0.005)))
end


text:
"""
    stFlux(inputSize::NTuple{N}, m; trainable=false,
                             normalize=true, outputPool=2,
                             poolBy=3 // 2, σ=abs, flatten=false, kwargs...)
    stFlux(inputSize::NTuple{N}, m; trainable = false,
                        normalize = true, outputPool = 2,
                        poolBy= 3//2, σ=abs, scales=(4,4,4),
                        shearLevels=scales) where {N}

Create a scattering transform of depth m (which returns a m+1 depth list of
arrays) that subsamples at a rate of poolBy each layer, using scales[i] and
`shearLevels[i]` at each layer. Normalize means give each layer the same average
weight per path, e.g. since the zeroth layer has one path, give it norm 1, if
the first layer has 16 paths, give it norm 16, etc. This is primarily for
cases where the classification algorithm needs roughly the same order of
magnitude variance.
"""
function stFlux(inputSize::NTuple{N}, m; trainable=false,
    normalize=true, outputPool=2,
    poolBy=3 // 2, σ=abs, flatten=false, kwargs...) where {N}
    # N determines the number of spatial dimensions
    Nd = N - 2
    if length(outputPool) == 1 # replicate for each dimension and layer
        outputPool = ntuple(k -> ntuple(i -> outputPool[1], Nd), m + 1)
    elseif length(outputPool) == m + 1 # replicate for each dimension
        outputPool = ntuple(k -> ntuple(i -> outputPool[k]), m + 1)
    elseif length(outputPool) == Nd
        outputPool = ntuple(k -> outputPool, m + 1)
    end
    if typeof(poolBy) <: Union{<:Rational,<:Integer}
        poolBy = map(ii -> RationPool((ntuple(i -> poolBy, N - 2)...,),
                nExtraDims=nPathDims(ii + 1) + 1), 1:m+1)
    end
    poolBy = makeTuple(m, poolBy)
    argsToEach = processArgs(m + 1, kwargs)

    listOfSizes = [(inputSize..., ntuple(i -> 1, max(i - 1, 0))...) for i = 0:m]
    interstitial = Array{Any,1}(undef, 3 * (m + 1) - 2)
    for i = 1:m
        # first transform
        interstitial[3*i-2] = dispatchLayer(listOfSizes[i], Val(Nd); σ=identity,
            argsToEach[i]...)
        nFilters = length(interstitial[3*i-2].weight)

        pooledSize = poolSize(poolBy[i], listOfSizes[i][1:Nd])
        # then throw away the averaging (we'll pick it up in the actual transform)
        # also, in the first layer, merge the channels into the first shearing,
        # since max pool isn't defined for arbitrary arrays
        if i == 1
            listOfSizes[i+1] = (pooledSize...,
                (nFilters - 1) * listOfSizes[1][Nd+1],
                listOfSizes[i][Nd+2:end]...)
            interstitial[3*i-1] = x -> begin
                ax = axes(x)
                return σ.(reshape(x[ax[1:Nd]..., ax[Nd+1][1:end-1], ax[(Nd+2):end]...],
                    (ax[1:Nd]..., listOfSizes[i+1][Nd+1], ax[end])))
            end
        else
            listOfSizes[i+1] = (pooledSize..., (nFilters - 1),
                listOfSizes[i][Nd+1:end]...)
            interstitial[3*i-1] = x -> begin
                ax = axes(x)
                return σ.(x[ax[1:Nd]..., ax[Nd+1][1:end-1], ax[Nd+2:end]...])
            end
        end
        # then pool
        interstitial[3*i] = poolBy[i]
    end

    # final averaging layer
    interstitial[3*m+1] = dispatchLayer(listOfSizes[end], Val(Nd);
        averagingLayer=true, σ=identity,
        argsToEach[end]...)
    chacha = Chain(interstitial...)
    outputSizes = ([(map(poolSize, outputPool[ii], x[1:Nd])..., x[(Nd+1):end]...) for (ii, x)
                    in
                    enumerate(listOfSizes)]...,)

    # record the settings used pretty kludgy
    settings = Dict(:outputPool => outputPool, :poolBy => poolBy, :σ => σ, :flatten => flatten, (argsToEach...)...)
    return stFlux{Nd,m,typeof(chacha),typeof(outputSizes),typeof(outputPool),typeof(settings)}(chacha, normalize,
        outputSizes, outputPool, settings)
end

function dispatchLayer(listOfSizes, Nd::Val{1}; varargs...)
    waveletLayer(listOfSizes; varargs...)
end

Base.size(a::Tuple{AbstractFFTs.Plan,AbstractFFTs.Plan}) = size(a[1])

# actually apply the transform
function (St::stFlux{Dimension,Depth})(x::T) where {Dimension,Depth,T<:AbstractArray}
    mc = St.mainChain.layers
    if ndims(x) < length(size(mc[1]))
        # the input is missing dimensions, so reshape
        targetSize = size(mc[1])
        x = reshape(x, (size(x)..., targetSize[length(targetSize)-ndims(x):end]...))
    end
    if size(x)[end] != getBatchSize(mc[1])
        res = breakAndAdapt(St, x)
    else
        res = applyScattering(mc, x, ndims(St), St, 0)
    end
    if get(St.settings, :flatten, false) # it may not be defined, in which case we don't do it
        batchSize = size(res[1])[end]
        return cat((reshape(x, (:, batchSize)) for x in res)..., dims=1)
    else
        return ScatteredOut(res, ndims(mc[1]))
    end
end

# adapt changes both the eltype and the container type, while I just want a different container type
function maybeAdapt(contType, x)
    if contType <: CuArray && !(typeof(x) <: CuArray)
        # should be a CuArray but isn't
        return cu(x)
    elseif contType <: Array && !(typeof(x) <: Array)
        # should be an Array but isn't
        return adapt(Array, x)
    else
        return x
    end
end

"""
extract x at adr in the last dimension, and make sure that it has a size of chunkSize
"""
function extractAddPadding(x, adr, chunkSize, N)
    justUsing = x[fill(Colon(), N)..., :, adr]
    if length(adr) < chunkSize
        actualSize = chunkSize - length(adr)
        return cat(justUsing, zeros(eltype(justUsing), size(x)[1:end-1]...,
                actualSize), dims=ndims(justUsing)), length(adr)
    else
        return justUsing, chunkSize
    end
end

trim(x, actualSize) = x[axes(x)[1:end-1]..., 1:actualSize]
function breakAndAdapt(St::stFlux{N,D}, x) where {N,D}
    mc = St.mainChain.layers
    chunkSize = size(mc[1].fftPlan)[end]
    nSteps = ceil(Int, (size(x)[end]) / chunkSize) # the first entry is taken care of already
    containerType = typeof(St.mainChain[1].weight)
    xAxes = axes(x)
    firstAddr = 1+0:min(size(x)[end], chunkSize)
    firstEx, actualSize = extractAddPadding(x, firstAddr, chunkSize, N) # x[xAxes[1:end-1]..., 1:chunkSize]
    # do the first beforehand to get the sizes
    out = applyScattering(mc, maybeAdapt(containerType, firstEx), ndims(St), St, 0)
    # create storage
    outputs = map(out -> zeros(eltype(out), size(out)[1:end-1]..., size(x)[end]...), out)
    # util to write out to outputs at location batchInds
    function writeOut!(out, batchInds, actualSize)
        for jj in 1:length(out)
            mA = maybeAdapt(typeof(x), out[jj])
            oAx = axes(outputs[jj])
            @views outputs[jj][oAx[1:end-1]..., batchInds] = trim(mA, actualSize)
        end
    end
    writeOut!(out, firstAddr, actualSize) # write the first batch out
    for ii = 2:nSteps
        out = 3
        # clear output to force garbage collection, otherwise the gpu may be full
        addr = 1+(ii-1)*chunkSize:min(size(x)[end], ii * chunkSize)
        tmpX, actualSize = extractAddPadding(x, addr, chunkSize, N)
        out = applyScattering(mc, maybeAdapt(containerType, tmpX), ndims(St), St, 0)
        writeOut!(out, addr, actualSize)
    end
    return outputs
end

function applyScattering(c::Tuple, x, Nd, st, M)
    res = first(c)(x)
    if typeof(first(c)) <: ConvFFT
        tmpRes = res[map(x -> Colon(), 1:Nd)..., end, map(x -> Colon(), 1:ndims(res)-Nd-1)...]
        # return a subsampled version of the output at this layer
        poolSizes = (st.outputPool[M+1]..., ntuple(i -> 1, ndims(tmpRes) - Nd - 2)...)
        r = RationPool(st.outputPool[M+1], nExtraDims=nPathDims(M + 1) + 1)
        if st.normalize
            tmpRes = normalize(r(real.(tmpRes)), Nd)
            apld = applyScattering(tail(c), res, Nd, st, M + 1)
            return (tmpRes, apld...)
        else
            tmpRes = r(real.(tmpRes))
            return (tmpRes,
                applyScattering(tail(c), res, Nd, st, M + 1)...)
        end
    else
        # this is either a reshaping layer or a subsampling layer, so no output
        return applyScattering(tail(c), res, Nd, st, M)
    end
end

applyScattering(::Tuple{}, x, Nd, st, M) = tuple() # all of the returns should
# happen along the way, not at the end

"""
    normedX = normalize(x, Nd)
normalize x over the dimensions Nd through ndims(x)-1, e.g. for a Nd=2, and x
that is 4D, then norm(x[:,:,:,j], 2) ≈ size(x,3)
"""
function normalize(x, Nd)
    n = ndims(x)
    totalThisLayer = prod(size(x)[(Nd+1):(n-1)])
    ax = axes(x)
    buf = Zygote.Buffer(x)
    for i = 1:size(x)[end]
        xSlice = x[ax[1:(n-1)]..., i]
        thisExample = totalThisLayer / (sum(abs.(xSlice) .^ 2)) .^ (0.5f0)
        if isnan(thisExample) || thisExample ≈ Inf
            buf[ax[1:(n-1)]..., i] = xSlice
        else
            buf[ax[1:(n-1)]..., i] = xSlice .* thisExample
        end
    end
    return copy(buf)
end

normalize(sct::ScatteredOut, Nd) = ScatteredOut(map(x -> normalize(x, Nd),
        sct.output), sct.k)

normalize(sct::ScatteredFull, Nd) = ScatteredFull(sct.m, sct.k, sct.data,
    map(x -> normalize(x, Nd),
        sct.output))

ext: jl


text:
# Instructions
#
# Script should be run within the package folder.
# To reproduce results, checkout the github package
# https://github.com/mschauer/BridgeSPDE.jl
# with tag arxiv_v2

## From the shell:
# git clone https://github.com/mschauer/BridgeSPDE.jl
# git checkout tags/arxiv_v2

## From julia:
# include("scripts/smoothing.jl")

# install and make local copy
import Pkg

# instantiate reproducible environment
Pkg.activate(@__DIR__)
#Pkg.develop(Pkg.PackageSpec(url="https://github.com/mschauer/BridgeSPDE.jl"))

# move to package directory
#cd(joinpath(Pkg.devdir(), "BridgeSPDE"))
cd(joinpath(@__DIR__, ".."))

using Revise
using Test
using Kalman
using SparseArrays
using BridgeSPDE
using FileIO
using Base.Iterators: flatten
using Statistics
using LinearAlgebra
using GaussianDistributions
using Trajectories
using ImageTransformations
using Random
using Colors
PLOT = true
SAVE = true
if PLOT
    using Makie
end
inner(x::Vector) = dot(x,x)
inner(x::Vector, y::Vector) = dot(x,y)
inner(x, y) = x'*y
inner(x) = x'*x

@assert isfile(joinpath("src", "BridgeSPDE.jl")) # run from package directory

const F0 = Float32

θs = [(F0(0), F0(0))]
for downscale in [8, 4, 3]
    Random.seed!(1)

    global θs
    scale = F0(8 / downscale)

    mkpath(joinpath(@__DIR__, "..", "output$downscale"))

    #ds(img, s) = imresize(img, Int(size(img,1)/downscale), Int(size(img,2)/downscale))
    ds(img, ds) = imresize(img, floor(Int,size(img,1)/ds), floor(Int,size(img,2)/ds))

    meteosat = [ ds(F0.(FileIO.load(joinpath(@__DIR__, "..","data", "meteosatf$frame.png"))), downscale) for frame in [33, 35, 40, 45]]

    if SAVE && PLOT
        for i in eachindex(meteosat)
            FileIO.save(joinpath(@__DIR__, "..", "output$downscale", "meteosat$i.png"), Gray.(meteosat[i]))

            p1 = surface(meteosat[i]; shading=false, show_axis=false, colormap = :deep)
            scale!(p1, 1.0, 1.0, 7.5)
            FileIO.save(joinpath(@__DIR__, "..", "output$downscale", "meteosurf$i.png"), p1)

        end
    end

    y = meteosat[1]

    m, n = size(y)

    Δt = F0(1.0)
    J = gridderiv(F0, m, n)
    Λ0 = gridlaplacian(F0, m, n)
    @test norm(Λ0 + (J[1]+J[2]+J[3]+J[4])) == 0

    J1 = gridderiv(F0, m, n)[1]*F0(scale)
    J2t = gridderiv(F0, m, n)[4]*F0(scale)
    img(x) = image(reshape(x, (m, n)))
    mat(x) = reshape(x, (m, n))

    d = m*n
    Λ = (Λ0 + 0.1f0 * I)*downscale

    q0 = F0(0.5) .+ zeros(F0, d)
    Q0 = zeros(F0, d, d) + F0(2.0)^2*I

    # nominally, the noise should scale with sqrt(pixel in mean)
    # the noise is not really independent, so actually it doesn't scale (so much)
    R = F0(0.05)^2*I*scale
    σ = F0(0.08*scale)
    H = I

    a = σ^2*I
    σᶥ = sparse(cholesky(Λ0+0.2I).L)

    #L = randn(5,5)
    #heatmap([inv(L*L') cov((L'\randn(5, 10000))')])

    l = floor(Int,scale)*7*[2,5,5] # not equidistant
    dt = F0(Δt / l[2])
    T = sum(l)*dt
    droptol = F0(1e-8)
    droptoli = F0(1e-8)

    iters = 1:15
    saveiters = [5, 10, 15]
    for iter in iters
        #iter = 1
        #begin
        println("\nIter $iter\n")
        θ1, θ2 = θs[end]

        t = T
        ν, P = copy(q0), copy(Q0)

        B = -σ^2/2*Λ + θ1*J1 + θ2*J2t
        B0 = -σ^2/2*Λ

        Bt = typeof(B)(B')

        @time begin
            println("backward filter")
            timecor = 0.0
            timecor += @elapsed (ν, P), _ = BridgeSPDE.correct((ν, sparse(P)), (vec(meteosat[end]), sparse(R, d,d)), H, droptol=droptoli)

            droptol!(P, droptol)

            PLOT && image(reshape(ν, (m, n)))
            μ = trajectory((t[1] => Gaussian(ν, P),))

            for segment in 3:-1:1
                for i in 2:l[segment]
                    t = t - dt
                    P = P .- dt*(B*P + P*Bt - a)
                    droptol!(P, droptol)
                    i == l[segment] && @show sparsity(P)
                    ν = ν - B*(ν*dt)
                    push!(μ, t => Gaussian(ν, P))
                end
                t = t - dt
                P = P - dt*(B*P + P*Bt - a)
                ν = ν - B*ν*dt
                #(ν, P), _ = Kalman.correct(Kalman.JosephForm(), Gaussian(ν, P), (Gaussian(vec(meteosat[segment]), R), H))
                timecor += @elapsed (ν, P), _ = BridgeSPDE.correct((ν, P), (vec(meteosat[segment]), sparse(R, d,d)), H, droptol=droptoli)
                droptol!(P, droptol)
                push!(μ, t => Gaussian(ν, P))
            end
            println("correction: $timecor")
        end

        @time begin
            println("\nforward guiding")

            #
            ts = μ.t
            νs = reverse(mean.(μ.x))
            Ps = reverse(cov.(μ.x))

            rs(x) = reshape(x, (m, n))

            x0 = νs[1]
            x = copy(x0);
            x2 = rand(Gaussian(νs[1], Symmetric(Matrix(Ps[1]))))

            global X = trajectory((ts[1]=>reshape(x, (m,n)),))
            global X2 = trajectory((ts[1]=>reshape(x2, (m,n)),))
            for i in 2:length(μ)
        #        global x,t
                x = x + dt*(B*x) + σ^2*(cholesky(Hermitian(Ps[i]) + 100droptol*I)\(νs[i] - x))*dt
                x2 = x2 + dt*(B*x2) + σ^2*(cholesky(Hermitian(Ps[i]) + 100droptol*I)\(νs[i] - x2))*dt + σ*sqrt(dt)*randn(F0, d)
                t += dt
                push!(X, ts[i] => reshape(x, (m, n)))
                push!(X2, ts[i] => reshape(x2, (m, n)))

            end
        end

        @time begin
            println("\nconjugate update")

            μ0 = zeros(F0, 2)
            Γ = zeros(F0, 2,2)
            i = 2
            for i in 2:length(X)
        #        global μ0, Γ
                x = vec(X.x[i-1])
                z = inv(a)*(vec(X.x[i]-X.x[i-1]) - B0*x*dt)
                μ0 += inner([J1*x J2t*x], z)

                Γ += inner(inv(σ)*[J1*x J2t*x])*dt # CHECK
            end
            @show θ1, θ2 = cholesky(Symmetric(Γ))\μ0
            push!(θs, (θ1, θ2))
        end

        @time begin
            println("\ncreate images")

            m3 = quantile(flatten(X.x), 0.17)

            PLOT && image(X.x[1])
            if SAVE && iter in saveiters
                for i in eachindex(X.x)[1:5:end]
                    global out = reshape(clamp.(X.x[i], 0, 1), (m, n))

                    out2 = boundary(F0.(out .> m3))
                    out3 = reshape(clamp.(X2.x[i], 0, 1), (m, n))
                    mi = findmin(map(k->abs(k - i), cumsum(l[2:end]).-l[1]))[2] # which image is closest
                    FileIO.save(joinpath(@__DIR__, "..", "output$downscale", "img$iter-$i.png"), Gray.(out))

                    if PLOT
                        p1 = surface(out; shading=false, show_axis=false, colormap = :deep)
                        scale!(p1, 1.0, 1.0, 7.5)
                        FileIO.save(joinpath(@__DIR__, "..", "output$downscale", "surf$iter-$i.png"), p1)
                        p2 = surface(out3; shading=false, show_axis=false, colormap = :deep)
                        scale!(p2, 1.0, 1.0, 7.5)
                        FileIO.save(joinpath(@__DIR__, "..", "output$downscale", "surfr$iter-$i.png"), p2)

                    end
                end
                run(`ffmpeg -y -r 40 -f image2 -i output$downscale/img$iter-%d.png -vcodec libx264 -crf 25 -vf "pad=ceil(iw/2)*2:ceil(ih/2)*2" -pix_fmt yuv420p img$downscale-$iter.mp4`)
                PLOT && run(`ffmpeg -y -r 40 -f image2 -i output$downscale/surf$iter-%d.png -vcodec libx264 -crf 25  -pix_fmt yuv420p surf$downscale-$iter.mp4`)
                PLOT && run(`ffmpeg -y -r 40 -f image2 -i output$downscale/surfr$iter-%d.png -vcodec libx264 -crf 25  -pix_fmt yuv420p surfr$downscale-$iter.mp4`)

            end
        end
    end

    println("Done downscale= $downscale.")
end

# save thetas in scale of pixels of final resolution
downscale = 3
scale = F0(8 / downscale)
Δt = F0(1.0)

l = floor(Int,scale)*7*[2,5,5] # not equidistant
dt = F0(Δt / l[2])
T = sum(l)*dt

writedlm("thetas.txt", [first.(θs)*scale*T last.(θs)*scale*T])

lines(first.(θs)*scale*T);
lines!(last.(θs)*scale*T)


text:
#-------------------------------------------------------------------
#* EMSO Model Library (EML) Copyright (C) 2004 - 2007 ALSOC.
#*
#* This LIBRARY is free software; you can distribute it and/or modify
#* it under the therms of the ALSOC FREE LICENSE as available at
#* http://www.enq.ufrgs.br/alsoc.
#*
#* EMSO Copyright (C) 2004 - 2007 ALSOC, original code
#* from http://www.rps.eng.br Copyright (C) 2002-2004.
#* All rights reserved.
#*
#* EMSO is distributed under the therms of the ALSOC LICENSE as
#* available at http://www.enq.ufrgs.br/alsoc.
#*
#*--------------------------------------------------------------------
#* Model of basic streams
#*----------------------------------------------------------------------
#* Author: Paula B. Staudt and Rafael de P. Soares
#* $Id$
#*---------------------------------------------------------------------
type work_source
	work_source()=begin
		new(
			power(Dict{Symbol,Any}(
				:Brief=>"Outlet work stream",
				:PosX=>1,
				:PosY=>0.46,
				:Symbol=>"_{out}"
			)),
			[:Work,]
		)
	end
	Work::power
	variables::Array{Symbol,1}
	attributes::Dict{Symbol,Any}
end
export work_source
function atributes(in::work_source,_::Dict{Symbol,Any})
	fields::Dict{Symbol,Any}=Dict{Symbol,Any}()
	fields[:Pallete]=true
	fields[:Icon]="icon/work_source"
	fields[:Brief]="Work stream source"
	drive!(fields,_)
	return fields
end
work_source(_::Dict{Symbol,Any})=begin
	newModel=work_source()
	newModel.attributes=atributes(newModel,_)
	newModel
end


text:
InputFolder = './Images/HeLaMCF10AMovie/Movie1/';
OutputFolder = './Results/Images/HeLaMCF10AMovie/Movie1/';

@iA = '*.tif';

@fxg_mLapThrBinWatTiles [iA] > [L];
params.GRad = 7;
params.Thr = 0.02;
params.GaussianD = 2;
params.MinArea = 20;
/endf

/keep L > tif;

max_stars_repo_name: SebastienTs/LOBSTER

max_issues_repo_name: SebastienTs/LOBSTER

max_forks_repo_name: SebastienTs/LOBSTER


text:
#== # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
#
#   Auxiliary functions to pretty print tables.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # ==#

export @ptconfclean, @ptconf, @pt

# Global configuration object.
const _pt_conf = PrettyTablesConf()

"""
    @ptconfclean()

Clean all global configurations to pretty print tables using the macro `@pt`.

"""
macro ptconfclean()
    return :(clear_pt_conf!(_pt_conf))
end

"""
    @ptconf(expr...)

Add configurations in `expr` to be used with the macro `@pt`.

The expression format must be:

    keyword1 = value1 keyword2 = value2 ...

in which the keywords can be any other possible keyword that can be used in the
function `pretty_table`.

!!! warning

    If a keyword is not supported by the function `pretty_table`, then no error
    message is printed when calling `@ptconf`. However, an error will be thrown
    when `@pt` is called.

"""
macro ptconf(expr...)
    kws = Expr[]

    for ex in expr
        # If the head is :(=), then it must be a configuration. Notice that we
        # will ignore everything that is not an expression like `a = b`.
        if (ex isa Expr) && (ex.head == :(=))
            push!(kws, Expr(:kw, ex.args[1], ex.args[2]))
        end
    end

    return :(set_pt_conf!(_pt_conf, $(kws...)))
end

"""
    @pt(expr...)

Pretty print tables in `expr` to `stdout` using the global configurations
selected with the macro `@ptconf`.

Multiple tables can be printed by passing multiple expressions like:

    @pt table1 table2 table3

The user can select the table header by passing the expression:

    :header = [<Vector with the header>]

Notice that the header is valid only for the next printed table. Hence:

    @pt :header = header1 table1 :header = header2 table2 table3

will print `table1` using `header1`, `table2` using `header2`, and `table3`
using the default header.

# Examples

```julia
julia> @ptconf tf = simple

julia> @pt :header = ["Time","Velocity"] [1:1:10 ones(10)] :header = ["Time","Position"] [1:1:10 1:1:10]
======= ===========
  Time   Velocity
======= ===========
   1.0        1.0
   2.0        1.0
   3.0        1.0
   4.0        1.0
   5.0        1.0
   6.0        1.0
   7.0        1.0
   8.0        1.0
   9.0        1.0
  10.0        1.0
======= ===========
======= ===========
  Time   Position
======= ===========
     1          1
     2          2
     3          3
     4          4
     5          5
     6          6
     7          7
     8          8
     9          9
    10         10
======= ===========

julia> @pt ones(3,3) + I + [1 2 3; 4 5 6; 7 8 9]
========= ======== =========
  Col. 1   Col. 2   Col. 3
========= ======== =========
     3.0      3.0      4.0
     5.0      7.0      7.0
     8.0      9.0     11.0
========= ======== =========
```

# Remarks

When more than one table is passed to this macro, then multiple calls to
`pretty_table` will occur. Hence, the cropping algorithm will behave exactly the
same as printing the tables separately.

"""
macro pt(expr...)
    exprs  = Expr[]
    header = nothing

    for ex in expr
        # If the head is :(=) and the argument is `:header`, then we must set
        # the header.
        if (ex isa Expr) && (ex.head == :(=)) && (ex.args[1] == :(:header))
            header = esc(ex.args[2])
        # If it is not, then we assume that it is a table to be printed.
        else
            # Assemble the arguments.
            args = Expr[]
            header != nothing  && push!(args, header)

            if isempty(args)
                expr = :(pretty_table_with_conf(_pt_conf, $(esc(ex)) ))
            else
                expr = :(pretty_table_with_conf(_pt_conf, $(esc(ex)), $(args...)))
            end

            push!(exprs, expr)
            header = nothing
        end
    end

    return Expr(:block, exprs...)
end


text:
# Autogenerated wrapper script for LibCURL_jll for x86_64-linux-gnu
export curl, libcurl

using LibSSH2_jll
using MbedTLS_jll
using Zlib_jll
using nghttp2_jll
## Global variables
PATH = ""
LIBPATH = ""
LIBPATH_env = "LD_LIBRARY_PATH"

# Relative path to `curl`
const curl_splitpath = ["bin", "curl"]

# This will be filled out by __init__() for all products, as it must be done at runtime
curl_path = ""

# curl-specific global declaration
function curl(f::Function; adjust_PATH::Bool = true, adjust_LIBPATH::Bool = true)
    global PATH, LIBPATH
    env_mapping = Dict{String,String}()
    if adjust_PATH
        if !isempty(get(ENV, "PATH", ""))
            env_mapping["PATH"] = string(PATH, ':', ENV["PATH"])
        else
            env_mapping["PATH"] = PATH
        end
    end
    if adjust_LIBPATH
        if !isempty(get(ENV, LIBPATH_env, ""))
            env_mapping[LIBPATH_env] = string(LIBPATH, ':', ENV[LIBPATH_env])
        else
            env_mapping[LIBPATH_env] = LIBPATH
        end
    end
    withenv(env_mapping...) do
        f(curl_path)
    end
end

# Relative path to `libcurl`
const libcurl_splitpath = ["lib", "libcurl.so"]

# This will be filled out by __init__() for all products, as it must be done at runtime
libcurl_path = ""

# libcurl-specific global declaration
# This will be filled out by __init__()
libcurl_handle = C_NULL

# This must be `const` so that we can use it with `ccall()`
const libcurl = "libcurl.so.4"

"""
Open all libraries
"""
function __init__()
    global artifact_dir = abspath(artifact"LibCURL")

    # Initialize PATH and LIBPATH environment variable listings
    global PATH_list, LIBPATH_list
    # We first need to add to LIBPATH_list the libraries provided by Julia
    append!(LIBPATH_list, [joinpath(Sys.BINDIR, Base.LIBDIR, "julia"), joinpath(Sys.BINDIR, Base.LIBDIR)])
    # From the list of our dependencies, generate a tuple of all the PATH and LIBPATH lists,
    # then append them to our own.
    foreach(p -> append!(PATH_list, p), (LibSSH2_jll.PATH_list, MbedTLS_jll.PATH_list, Zlib_jll.PATH_list, nghttp2_jll.PATH_list,))
    foreach(p -> append!(LIBPATH_list, p), (LibSSH2_jll.LIBPATH_list, MbedTLS_jll.LIBPATH_list, Zlib_jll.LIBPATH_list, nghttp2_jll.LIBPATH_list,))

    global curl_path = normpath(joinpath(artifact_dir, curl_splitpath...))

    push!(PATH_list, dirname(curl_path))
    global libcurl_path = normpath(joinpath(artifact_dir, libcurl_splitpath...))

    # Manually `dlopen()` this right now so that future invocations
    # of `ccall` with its `SONAME` will find this path immediately.
    global libcurl_handle = dlopen(libcurl_path)
    push!(LIBPATH_list, dirname(libcurl_path))

    # Filter out duplicate and empty entries in our PATH and LIBPATH entries
    filter!(!isempty, unique!(PATH_list))
    filter!(!isempty, unique!(LIBPATH_list))
    global PATH = join(PATH_list, ':')
    global LIBPATH = join(LIBPATH_list, ':')

    # Add each element of LIBPATH to our DL_LOAD_PATH (necessary on platforms
    # that don't honor our "already opened" trick)
    #for lp in LIBPATH_list
    #    push!(DL_LOAD_PATH, lp)
    #end
end  # __init__()


text:
##  SCRIPT NAME:   getEquiprobNormalDeviates.jl
##  DATE:          Jnauary 2015
#
#   This script generates N equiprobable line segments in the stationary
#   univariate normal distribution of income Y:
#          Y ~ N [ alpha/1-rho  ;  sigma_epsilon^2/(1-rho)^2 ]
#   where the mean of the above distribution is declared below as 'Ymean'
#   and the standard deviation as 'Ysd'. It delivers two arrays: 'Bounaries'
#   which holds the N+1 bounds of the N equiprobable segments, and
#   'ExpValues' which holds the expected value in each segment.
#
#
##  ALEXANDROS THELOUDIS, UCL
#
#  Needs: A constant called "normBnd"

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

function getEquiprobNormalDeviates(Ymean, Ysd, N)

  #  Initialise output:
  Boundaries = zeros(N+1)
  ExpValues  = zeros(N)

  #  --------------------------------------------------------------------------------------------
  #  Calculate applicable support of univariate normal distribution (first and
  #  last point in the range that I divide equiprobably). Note: in principle
  #  this is +-infinity but I am ruling some extremely unlikely values out:
  Boundaries[1]   = Ymean - normBnd * Ysd
  Boundaries[N+1] = Ymean + normBnd * Ysd

  #  --------------------------------------------------------------------------------------------
  #  Sequentially calculate the remaining bounds of the equiprobable regions:
  for ixi = 2:1:N
    Boundaries[ixi] = quantile(TruncatedNormal(Ymean,Ysd,Boundaries[1],Boundaries[N+1]), (ixi-1)/N)
  end

  #  --------------------------------------------------------------------------------------------
  #  Calculate the expected value of a random variable X ~ N[Ymean,Ysd^2] which
  #  is bounded by the equiprobable line segment above. Note: this is essentially
  #  the mean of a univariate 2-sided truncated normal distribution; read more
  #  on http://en.wikipedia.org/wiki/Truncated_normal_distribution
  for ixi = 1:1:N
    ExpValues[ixi] = Ymean + Ysd * N * (pdf(Normal(),(Boundaries[ixi]-Ymean)/Ysd) - pdf(Normal(),(Boundaries[ixi+1]-Ymean)/Ysd))
  end

  return (Boundaries, ExpValues)

end

max_stars_repo_name: floswald/ucl-econ-julia

max_issues_repo_name: floswald/ucl-econ-julia

max_forks_repo_name: floswald/ucl-econ-julia


text:
## R code 5.3
m5.1 <- quap(
    alist(
        D ~ dnorm( mu , sigma ) ,
        mu <- a + bA * A ,
        a ~ dnorm( 0 , 0.2 ) ,
        bA ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 )
    ) , data = d )

## R code 5.6
d$M <- scale( d$Marriage )
m5.2 <- quap(
    alist(
        D ~ dnorm( mu , sigma ) ,
        mu <- a + bM * M ,
        a ~ dnorm( 0 , 0.2 ) ,
        bM ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 )
    ) , data = d )

## R code 5.10
m5.3 <- quap(
    alist(
        D ~ dnorm( mu , sigma ) ,
        mu <- a + bM*M + bA*A ,
        a ~ dnorm( 0 , 0.2 ) ,
        bM ~ dnorm( 0 , 0.5 ) ,
        bA ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 )
    ) , data = d )

## R code 5.13
m5.4 <- quap(
    alist(
        M ~ dnorm( mu , sigma ) ,
        mu <- a + bAM * A ,
        a ~ dnorm( 0 , 0.2 ) ,
        bAM ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 )
    ) , data = d )

## R code 5.19
data(WaffleDivorce)
d <- list()
d$A <- standardize( WaffleDivorce$MedianAgeMarriage )
d$D <- standardize( WaffleDivorce$Divorce )
d$M <- standardize( WaffleDivorce$Marriage )

m5.3_A <- quap(
    alist(
      ## A -> D <- M
        D ~ dnorm( mu , sigma ) ,
        mu <- a + bM*M + bA*A ,
        a ~ dnorm( 0 , 0.2 ) ,
        bM ~ dnorm( 0 , 0.5 ) ,
        bA ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 ),
      ## A -> M
        M ~ dnorm( mu_M , sigma_M ),
        mu_M <- aM + bAM*A,
        aM ~ dnorm( 0 , 0.2 ),
        bAM ~ dnorm( 0 , 0.5 ),
        sigma_M ~ dexp( 1 )
    ) , data = d )

## R code 5.29
m5.5_draft <- quap(
    alist(
        K ~ dnorm( mu , sigma ) ,
        mu <- a + bN*N ,
        a ~ dnorm( 0 , 1 ) ,
        bN ~ dnorm( 0 , 1 ) ,
        sigma ~ dexp( 1 )
    ) , data=d )

## R code 5.32
m5.5_draft <- quap(
    alist(
        K ~ dnorm( mu , sigma ) ,
        mu <- a + bN*N ,
        a ~ dnorm( 0 , 1 ) ,
        bN ~ dnorm( 0 , 1 ) ,
        sigma ~ dexp( 1 )
    ) , data=dcc )

## R code 5.34
m5.5 <- quap(
    alist(
        K ~ dnorm( mu , sigma ) ,
        mu <- a + bN*N ,
        a ~ dnorm( 0 , 0.2 ) ,
        bN ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 )
    ) , data=dcc )

## R code 5.37
m5.6 <- quap(
    alist(
        K ~ dnorm( mu , sigma ) ,
        mu <- a + bM*M ,
        a ~ dnorm( 0 , 0.2 ) ,
        bM ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 )
    ) , data=dcc )

## R code 5.38
m5.7 <- quap(
    alist(
        K ~ dnorm( mu , sigma ) ,
        mu <- a + bN*N + bM*M ,
        a ~ dnorm( 0 , 0.2 ) ,
        bN ~ dnorm( 0 , 0.5 ) ,
        bM ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 )
    ) , data=dcc )

## R code 5.41
# M -> K <- N
# M -> N
n <- 100
M <- rnorm( n )
N <- rnorm( n , M )
K <- rnorm( n , N - M )
d_sim <- data.frame(K=K,N=N,M=M)

m5.7.A <- quap(
    alist(
        K ~ dnorm( mu , sigma ) ,
        mu <- a + bN*N + bM*M ,
        a ~ dnorm( 0 , 0.2 ) ,
        bN ~ dnorm( 0 , 0.5 ) ,
        bM ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 )
    ) , data=d_sim )

## R code 5.47
m5.8 <- quap(
    alist(
        height ~ dnorm( mu , sigma ) ,
        mu <- a[sex] ,
        a[sex] ~ dnorm( 178 , 20 ) ,
        sigma ~ dunif( 0 , 50 )
    ) , data=d )
precis( m5.8 , depth=2 )

## R code 5.51
d$K <- scale( d$kcal.per.g )
m5.9 <- quap(
    alist(
        K ~ dnorm( mu , sigma ),
        mu <- a[clade_id],
        a[clade_id] ~ dnorm( 0 , 0.5 ),
        sigma ~ dexp( 1 )
    ) , data=d )

## R code 5.53
m5.10 <- quap(
    alist(
        K ~ dnorm( mu , sigma ),
        mu <- a[clade_id] + h[house],
        a[clade_id] ~ dnorm( 0 , 0.5 ),
        h[house] ~ dnorm( 0 , 0.5 ),
        sigma ~ dexp( 1 )
    ) , data=d )


text:
using Test
using MPI

MPI.Init()

#MPI.mpitype_dict[Boundary] = MPI.mpitype_dict[Int]
comm_size = MPI.Comm_size(MPI.COMM_WORLD)
comm_rank = MPI.Comm_rank(MPI.COMM_WORLD) + 1

# send to next higher process, with wraparound
dest = (comm_rank % comm_size) + 1
if comm_rank > 1
  src = comm_rank - 1
else
  src = comm_size
end

# test simple type

mutable struct NotABits
  a::Any
end

@test_throws ArgumentError MPI.mpitype(NotABits)

struct Boundary
  c::UInt16  # force some padding to be inserted
  a::Int
  b::UInt8
end

MPI.mpitype(Boundary)

arr = [Boundary( (comm_rank + i) % 127, i + comm_rank, i % 64) for i = 1:3]
req_send = MPI.Isend(arr, dest - 1, 1, MPI.COMM_WORLD)

# receive the message
arr_recv = Array{Boundary}(undef, 3)
req_recv = MPI.Irecv!(arr_recv, src - 1, 1, MPI.COMM_WORLD)

MPI.Wait!(req_send)
MPI.Wait!(req_recv)

# check received array
for i=1:3
  bndry_i = arr_recv[i]
  @test bndry_i.a == (src + i)
  @test bndry_i.b == i % 64
  @test bndry_i.c == (src + i) % 127
end

# test nested types
struct Boundary2
  a::UInt32
  b::Tuple{Int, UInt8}
end

MPI.mpitype(Boundary2)

arr = Array{Boundary2}(undef,3)
arr_recv = Array{Boundary2}(undef,3)

for i=1:3
  arr[i] = Boundary2( (comm_rank + i) % 127, ( Int(i + comm_rank), UInt8(i % 64) ) )
end

req_send = MPI.Isend(arr, dest - 1, 1, MPI.COMM_WORLD)
req_recv = MPI.Irecv!(arr_recv, src - 1, 1, MPI.COMM_WORLD)

MPI.Wait!(req_send)
MPI.Wait!(req_recv)

# check received array
for i=1:3
  bndry_i = arr_recv[i]
  @test bndry_i.a == (src + i) % 127
  @test bndry_i.b[1] == (src + i)
  @test bndry_i.b[2] == (i % 64)
end

# test a primitive type
primitive type Primitive16 16 end
primitive type Primitive24 24 end

nfields, blocklengths, displacements, types = MPI.factorPrimitiveType(Primitive16)
@test nfields == 1
@test displacements[1] == 0
@test types[1] == MPI.mpitype(Int16)
@test blocklengths[1] == 1

nfields, blocklengths, displacements, types = MPI.factorPrimitiveType(Primitive24)
@test nfields == 2
@test displacements[1] == 0
@test displacements[2] == 2
@test types[1] == MPI.mpitype(Int16)
@test types[2] == MPI.mpitype(Int8)
@test blocklengths[1] == 1
@test blocklengths[2] == 1

obj = [Ptr{Int}(comm_rank)]
obj_recv = Array{Ptr{Int}}(undef, 1)
recv_req = MPI.Irecv!(obj_recv, src - 1, 2, MPI.COMM_WORLD)
send_req = MPI.Isend(obj, dest - 1, 2, MPI.COMM_WORLD)

MPI.Wait!(recv_req)
MPI.Wait!(send_req)

@test obj_recv[1] == Ptr{Int}(src)

MPI.Barrier(MPI.COMM_WORLD)
MPI.Finalize()
@test MPI.Finalized()


text:
using SurfaceQuasiGeostrophic
using Test
using FFTW
using ProgressMeter
using TimerOutputs

const to = TimerOutput()

@testset "Advection" begin

# Duration of the simulation [in seconds]
advection_duration = 3600*24; # 1 day

# Resolution [even integer]
nx, ny  = 128, 128
lx, ly = 1e6, 1e6
angle_grid = π/4
omega = 2π/(24*60*60)
f0 = 2 * omega * sin( angle_grid )

## Grid [spatial & Fourier]
grid = Grid( lx, nx, ly, ny )

# Gather parameters in the SQG model
sqg = SQG( grid, f0 )

# Initial condition for the buoyancy
init_buoyancy!(sqg)

## Initialisation of the spatial fields

update_velocities!( sqg)

## Hyperviscosity [order & coefficient]
@show sqg.hv_order = 8
@show sqg.hv_val = compute_hypercoef(sqg)

## Choice of time step : CFL
@show dt = compute_dt(sqg)
println("Time step: $dt seconds")

nstep = floor(Int, advection_duration / dt)
@show nstep

# Printing some information
println("Final time: $(nstep*dt÷(3600*24)) days ")

rk4 = TimeSolver(nx, ny)

pbar = Progress(nstep)

for i = 1:nstep

    rk4.b̂ .= sqg.b̂

    @timeit to "advection" update_advection_term!(sqg)

    sqg.b̂ .= rk4.b̂ .+ sqg.â .* dt / 2
    rk4.db̂ .= sqg.â ./ 2

    @timeit to "advection" update_advection_term!(sqg)

    sqg.b̂ .= rk4.b̂ .+ sqg.â .* dt / 2
    rk4.db̂ .+= sqg.â

    @timeit to "advection" update_advection_term!(sqg)

    sqg.b̂ .= rk4.b̂ .+ sqg.â .* dt
    rk4.db̂ .+= sqg.â

    @timeit to "advection" update_advection_term!(sqg)

    rk4.db̂ .+= sqg.â ./ 2

    sqg.b̂ .= rk4.b̂ .+ dt / 3 .* rk4.db̂

    next!(pbar)

end

show(to)

@test true

end

include("init_grid.jl")
include("test_poisson.jl")
include("init_buoyancy.jl")


text:
module GrammaticalEvolution

import Base.*
import Base.sort!
import Base.getindex
import Base.setindex!
import Base.length
import Base.endof
import Base.isless
import Base.push!
import Base.pop!

export Individual, Population
export select_two_individuals, one_point_crossover, mutate!, evaluate!, generate, transform
export length, getindex, endof, setindex!, isless, genome_iterator
export MaxWrapException

include("EBNF.jl")

export Grammar, @grammar, Rule, AndRule, OrRule, parseGrammar
export RangeRule, ReferencedRule, ExprRule, RepeatedRule, Terminal #rcnlee added

type MaxWrapException <: Exception end

abstract Individual
abstract Population

#rcnlee
include("../examples/ExamplePopulation.jl")
export ExampleIndividual, ExamplePopulation
#/rcnlee

# methods that have to be supported by subclasses of population
length{T <: Population}(pop::T) = length(pop.individuals)
getindex{T <: Population}(pop::T, indices...) = pop.individuals[indices...]
push!{T <: Population, S <: Individual}(pop::T, ind::S) = push!(pop.individuals, ind)
pop!{T <: Population}(pop::T) = pop!(pop.individuals)

# methods that have to be supported by subclasses of individuals
length{T <: Individual}(ind::T) = length(ind.genome)
endof{T <: Individual}(ind::T) = endof(ind.genome)
getindex{T <: Individual}(ind::T, indices...) = ind.genome[indices...]
setindex!{T <: Individual}(ind::T, value::Int64, indices) = ind.genome[indices] = value
isless{T <: Individual}(ind1::T, ind2::T) = ind1.fitness < ind2.fitness
getFitness{T <: Individual}(ind::T) = ind.fitness
getCode{T <: Individual}(ind::T) = ind.code
# evaluate(ind::Individual) = nothing
evaluate!{T <: Individual}(grammar::Grammar, ind::T, args...) = error("evaluate! not defined!")

# TODO: this should be distributed
function evaluate!{PopulationType <: Population}(grammar::Grammar, pop::PopulationType, args...)
  for i=1:length(pop)
    if getCode(pop[i]) == nothing #uninitialized
      evaluate!(grammar, pop[i], pop, args...)
    end
  end
end

function sort!{PopulationType <: Population}(pop::PopulationType)
  sort!(pop.individuals)
end

function select_two_individuals{T <: Individual}(individuals::Array{T,1})
  # randomly select two individuals from a population
  while true
    i1 = rand(1:length(individuals))
    i2 = rand(1:length(individuals))

    # make sure the two individuals are not the same
    if i1 != i2
      return (i1, i2)
    end
  end
end

function one_point_crossover{IndividualType <: Individual}(ind1::IndividualType, ind2::IndividualType)
  cross_point = rand(1:length(ind1))
  g1 = vcat(ind1[1:cross_point-1], ind2[cross_point:end])
  g2 = vcat(ind2[1:cross_point-1], ind1[cross_point:end])

  return (IndividualType(g1), IndividualType(g2), cross_point)
end

function mutate!(ind::Individual, mutation_rate::Float64; max_value=1000)
  for i=1:length(ind)
    if rand() < mutation_rate
      ind[i] = rand(1:max_value)
    end
  end
end

"""
top_keep is fraction of population (sorted by fitness) that is directly kept
rand_frac is fraction of new population that is randomly generated
top_seed is fraction of population (sorted by fitness) that is used to seed the remaining part of population
prob_mutation is the probability of mutating an individual
mutation_rate is the mutation rate on a mutating individual
"""
function generate{PopulationType <: Population}(grammar::Grammar, population::PopulationType,
    top_keep::Float64, top_seed::Float64, rand_frac::Float64, prob_mutation::Float64,
    mutation_rate::Float64, args...)
  # sort population
  sort!(population)

  #rcnlee: changed#####
  # take top performers for cross-over
  n_keep = floor(Int64, length(population)*top_keep) #top n_keep is copied from top_performers
  n_seed = floor(Int64, length(population)*top_seed) #top n_seed is used to seed remaining generation
  n_rand = floor(Int64, length(population)*rand_frac) #rand_frac fraction of population is random
  top_performers = population[1:n_seed] #candidates for cross-over and mutation

  # create a new population
  genome_size = length(population[1])

  #randomly sample n_rand individuals
  new_population = PopulationType(n_rand, genome_size,
                                  population.best_fitness, #maintain info from previous generation
                                  population.best_ind,
                                  population.best_at_eval,
                                  population.totalevals) #random population of size n_rand

  for i = 1:n_keep #n_keep of old population is directly copied
    push!(new_population, population[i])
  end
  ##########

  # fill in the rest by mating top performers
  while length(new_population) < length(population)
    (i1, i2) = select_two_individuals(top_performers)
    (ind1, ind2) = one_point_crossover(top_performers[i1], top_performers[i2])
    push!(new_population, ind1)
    push!(new_population, ind2)
  end

  # mutate the crossed-over portion of population
  for j=(n_rand+n_keep+1):length(population)
    if rand() < prob_mutation
      mutate!(new_population[j], mutation_rate)
    end
  end

  # evaluate new population and re-sort
  evaluate!(grammar, new_population, args...)
  sort!(new_population)

  # it's possible that we might have added too many individuals, so trim down if necessary
  while length(new_population) > length(population)
    pop!(new_population)
  end

  return new_population
end

# stateful iterator that keeps track of its current position, wraps the position
# when the maximum length is reached, and emits an exception when the maximum
# number of wraps occurs
type GenomeIterator
  #consts
  size::Int64
  maxwraps::Int64

  #states
  i::Int64
  wraps::Int64
end

function GenomeIterator(size::Int64, maxwraps::Int64;
                        i::Int64=0, wraps::Int64=0)
  return GenomeIterator(size, maxwraps, i, wraps)
end

function Base.consume(pos::GenomeIterator)
  pos.i += 1
  if pos.i > pos.size
    pos.wraps += 1
    pos.i = 1
  end
  if pos.wraps > pos.maxwraps
    throw(MaxWrapException())
  end
  return pos.i
end

function transform(grammar::Grammar, ind::Individual; maxwraps=2)
  pos = GenomeIterator(length(ind), maxwraps)
  value = transform(grammar, grammar.rules[:start], ind, pos)
  return value
end

function transform(grammar::Grammar, rule::OrRule, ind::Individual, pos::GenomeIterator)
  idx = (ind[consume(pos)] % length(rule.values))+1
  value = transform(grammar, rule.values[idx], ind, pos)

  if rule.action !== nothing
    value = rule.action(value)
  end

  return value
end

function transform(grammar::Grammar, rule::RangeRule, ind::Individual, pos::GenomeIterator)
  value = (ind[consume(pos)] % length(rule.range))+rule.range.start

  if rule.action !== nothing
    value = rule.action(value)
  end

  return value
end

function transform(grammar::Grammar, rule::ReferencedRule, ind::Individual, pos::GenomeIterator)
  return transform(grammar, grammar.rules[rule.symbol], ind, pos)
end

function transform(grammar::Grammar, rule::Terminal, ind::Individual, pos::GenomeIterator)
  return rule.value
end

function transform(grammar::Grammar, rule::AndRule, ind::Individual, pos::GenomeIterator)
  values = [transform(grammar, subrule, ind, pos) for subrule in rule.values]

  if rule.action !== nothing
    values = rule.action(values)
  end

  return values
end

function transform(grammar::Grammar, sym::Symbol, ind::Individual, pos::GenomeIterator)
  return sym
end

function transform(grammar::Grammar, q::QuoteNode, ind::Individual, pos::GenomeIterator)
  return q.value
end

function transform(grammar::Grammar, rule::ExprRule, ind::Individual, pos::GenomeIterator)
  args = [transform(grammar, arg, ind, pos) for arg in rule.args]
  return Expr(args...)
end

# It's very unlikely these two methods will be useful -- the maximum size of the genome is arbritrarily high, so
# you'll likely end up with mostly large numbers
# function transform(grammar::Grammar, rule::ZeroOrMoreRule, ind::Individual, pos::GenomeIterator)
#   # genome value gives number of time to repeat
#   reps = ind[consume(pos)]

#   # invoke given rule reps times
#   values = [transform(grammar, rule.rule, ind, pos) for i=1:reps]

#   if rule.action !== nothing
#     values = rule.action(values)
#   end

#   return values
# end

# function transform(grammar::Grammar, rule::OneOrMoreRule, ind::Individual, pos::GenomeIterator)
#   # genome value gives number of time to repeat
#   reps = ind[consume(pos)]

#   # enforce that it's at least one
#   if reps == 0
#     reps = 1
#   end

#   # invoke given rule reps times
#   values = [transform(grammar, rule.value, ind, pos) for i=1:reps]

#   if rule.action !== nothing
#     values = rule.action(values)
#   end

#   return values
# end

function transform(grammar::Grammar, rule::RepeatedRule, ind::Individual, pos::GenomeIterator)
  # genome value gives number of time to repeat
  reps = ind[consume(pos)]
  range = (rule.range.stop - rule.range.start)
  reps = (reps % range) + rule.range.start

  # invoke given rule reps times
  values = [transform(grammar, rule.value, ind, pos) for i=1:reps]

  if rule.action !== nothing
    values = rule.action(values)
  end

  return values
end

end # module

max_stars_repo_name: rcnlee/GrammaticalEvolution

max_issues_repo_name: rcnlee/GrammaticalEvolution

max_forks_repo_name: rcnlee/GrammaticalEvolution


text:
#################### Heidelberger and Welch Diagnostic ####################

function heideldiag(x::Vector{T}; alpha::Real=0.05, eps::Real=0.1,
                             etype=:imse, start::Integer=1, args...) where {T<:Real}
  n = length(x)
  delta = trunc(Int, 0.10 * n)
  y = x[trunc(Int, n / 2):end]
  S0 = length(y) * mcse(y, etype; args...)^2
  i, pvalue, converged, ybar = 1, 1.0, false, NaN
  while i < n / 2
    y = x[i:end]
    m = length(y)
    ybar = mean(y)
    B = cumsum(y) - ybar * collect(1:m)
    Bsq = (B .* B) ./ (m * S0)
    I = sum(Bsq) / m
    pvalue = 1.0 - pcramer(I)
    converged = pvalue > alpha
    if converged
      break
    end
    i += delta
  end
  halfwidth = sqrt(2.0) * erfinv(1.0 - alpha) * mcse(y, etype; args...)
  passed = halfwidth / abs(ybar) <= eps
  [i + start - 2, converged, round(pvalue, digits = 4), ybar, halfwidth, passed]
end

function heideldiag(c::AbstractChains; alpha::Real=0.05, eps::Real=0.1,
                    etype=:imse, args...)
  _, p, m = size(c.value)
  vals = Array{Float64}(undef, p, 6, m)
  for j in 1:p, k in 1:m
    vals[j, :, k] = heideldiag(c.value[:, j, k], alpha=alpha, eps=eps,
                               etype=etype, start=first(c.range); args...)
  end
  hdr = header(c) * "\nHeidelberger and Welch Diagnostic:\n" *
        "Target Halfwidth Ratio = $eps\nAlpha = $alpha\n"
  ChainSummary(vals, c.names, ["Burn-in", "Stationarity", "p-value", "Mean",
                               "Halfwidth", "Test"], hdr)
end


text:
module DataProvider

@static in("BasePlus", keys(Pkg.installed())) ? nothing : Pkg.clone("https://github.com/zhmz90/BasePlus.jl.git")

using BasePlus

include("data_provider.jl")

end # module


text:
ENV["GKSwstype"] = 100 # gr segfault workaround
using Robotlib, LinearAlgebra, Random, Statistics, Plots
using Test
import Robotlib: ad, adi, I4

function simulateCalibration1(N)
    n = 6
    q = 2π * rand(N, n)

    wn = zeros(3, n)
    pn = zeros(3, n)
    xin = zeros(6, n + 1)
    wn[:, 1] = [0, 0, 1]
    pn[:, 1] = [0, 0, 0]
    wn[:, 2] = [0, 1, 0]
    pn[:, 2] = [0, 0, 0.3]
    wn[:, 3] = [0, 1, 0]
    pn[:, 3] = [0, 0, 1]
    wn[:, 4] = [1, 0, 0]
    pn[:, 4] = [1, 0, 1]
    wn[:, 5] = [0, 1, 0]
    pn[:, 5] = [1, 0, 1]
    wn[:, 6] = [1, 0, 0]
    pn[:, 6] = [1, 0, 1]
    T0 = [
    0 -1 0 1.2
    0 0 -1 0
    1 0 0 1
    0 0 0 1
    ]
    for i = 1:n
        xin[:, i] = [-skew(pn[:, i]) * wn[:, i]; wn[:, i]]
    end
    xin[:, n+1] = twistcoords(logT(T0))
    xinmod = deepcopy(xin)
    for i = 1:n+1
        xinmod[:, i] += [0.001randn(3); 0.01π / 180 * randn(3)]
        #conformize(xinmod[:,i])
    end
    Ta = zeros(4, 4, N)
    for i = 1:N
        Ta[:, :, i] = fkinePOE(xin, q[i, :])
    end
    return q, xin, T0, xinmod, Ta
end

@testset "Robotlib tests" begin
    Random.seed!(1)

    @testset "calibNAXP" begin
        @info "Testing calibNAXP"
        include("testCalibNAXP.jl")

    end
    @testset "Force" begin
        @info "Testing Force"
        include("testForce.jl")

    end

    @testset "DH" begin
        @info "Testing DH"
        DH2400()
        DH7600()
        DHYuMi()
        dh = DHtest()
        DH2twistsPOE(dh)
        DH2twistsLPOE(dh)

    end

    @testset "Plotting" begin
        @info "Testing Plotting"
        q, xin, T0, xinmod, Ta = simulateCalibration1(10)
        trajplot(Ta)
        trajplot3(Ta)
    end

    @testset "Utils" begin
        @info "Testing Utils"
        R = orthonormal(randn(3,3))
        @test Rangle(R,R) < 1e-7
        R2 = rpy2R(1*pi/180,0,0)*R
        @test Rangle(R,R2,true) <= 1.0000001

        q = Quaternion(R)
        @test Robotlib.Quaternions.rotationmatrix(q) ≈ R

        r,p,y = R2rpy(R, conv="xyz")
        @test rpy2R(r,p,y, "xyz") ≈ R

        r1 = [-1.0 0.0 0.0
              0.0 0.0 1.0
              0.0 1.0 0.0]
        log_r1 = logR(r1)
        log_r1_true = [0.0 -π/sqrt(2) π/sqrt(2)
                       π/sqrt(2) 0.0 0.0
                       -π/sqrt(2) 0.0 0.0]
        @test isapprox(log_r1, log_r1_true)

        r2 = [-1.0 0.0 0.0
              0.0 -1.0 0.0
              0.0 0.0 1.0]
        log_r2 = logR(r2)
        log_r2_true = [0.0 -π 0.0
                       π 0.0 0.0
                       0.0 0.0 0.0]
        @test isapprox(log_r2, log_r2_true)

        w1 = [1e-13; 0; 0]
        r1 = expω(w1)
        w2 = skewcoords(r1)
        @test isapprox(w1, w2)

        w3 = [1e-14; -1e-15; 2e-15]
        r2 = expω(w3)
        w4 = skewcoords(r2)
        @test isapprox(w3, w4)
        q3 = 123.0
        r3 = expω(w3, q3)
        w5 = skewcoords(r3)
        @test isapprox(w3*q3, w5)
    end

    @testset "Calibration and kinematics" begin

        @testset "Inverse kinematics" begin
            @info "Testing Inverse kinematics"

            dh = DH7600();
            q0 = ones(6)
            xi = DH2twistsPOE(dh)
            T0 = fkinePOE(xi,q0)
            Tt = Matrix(T0)
            Tt[1:3,4] += 0.1*[1,1,1]
            Tt[1:3,1:3] *= rpy2R(pi/180,0,0)

            @time q,err = ikinePOE(xi,Tt,q0,verbose=true)
            @test err[end] < sqrt(eps())

            # pp = semilogy(err)
            # display(pp)

        end

        @testset "Jacobian" begin
            @info "Testing Jacobian"

            dh = DH7600();
            q = [1,1,0,0,0,0];
            q = randn(6)
            AAA,BBB,T0,Ti0,Tn0 = jacobian(zeros(6),dh, I4);
            Jn,J0,Tjac,Tijac,Tnjac = jacobian(q,dh, I4);

            xiL = DH2twistsLPOE(Tn0)
            xi = DH2twistsPOE(Tn0)
            Jsb, Jbb = jacobianPOE(q,xi)

            @test Jsb ≈ J0
            @test_broken Jbb ≈ Jn

            # display(round.(J0, digits=3))
            # display(round.(Jsb, digits=3))
        end

        N = 100

        # println("===== Testing calibPOE =====")
        # q, xin,T0, xinmod, Ta= Robotlib.simulateCalibration_POE(N)
        # xic,et,er = calibPOE(xinmod,Ta,q,maxiter=150, λ = 100.0)
        # println("Initial error: ", norm(xin[:,1:7]-xinmod[:,1:7]))
        # println("Final error: ", norm(xin[:,1:7]-xic[:,1:7]))
        # @test et[et .!= 0][end] < 1e-12
        # @test et[et .!= 0][end] < 1e-12

        println("===== Testing calibLPOE =====")
        q, xin, Tn0, Tn0mod, Ta =
            Robotlib.simulateCalibration_LPOE(N)
        @time Tn0c, xic, et, er = calibLPOE(
            xin,
            Tn0mod,
            Ta,
            q,
            maxiter = 10,
            λ = 100.0,
        )
        @test et[et.!=0][end] < 1e-12
        @test er[er.!=0][end] < 1e-12
        println("Initial error: ", sqrt(sum((Tn0 - Tn0mod) .^ 2)))
        println("Final error: ", sqrt(sum((Tn0 - Tn0c) .^ 2)))
        @test sqrt(sum((Tn0 - Tn0c) .^ 2)) < 0.8sqrt(sum((Tn0 - Tn0mod) .^ 2))

        println("===== Testing calibLPOEdual =====")
        q, xin, Tn0, Tn0mod, Ta =
            Robotlib.simulateCalibration_LPOE_dual(N)
        @time Tn0c, xic, et, er = calibLPOEdual(
            xin,
            Tn0mod,
            q,
            maxiter = 6,
            λ = 0.01,
        )
        println(
            "Error between Tn0 and Tn0mod: ",
            sqrt(
                sum(
                    (Tn0[1:3, 4, :] - Tn0mod[1:3, 4, :, 1]) .^ 2 +
                    (Tn0[1:3, 4, :] - Tn0mod[1:3, 4, :, 2]) .^ 2,
                ) / N,
            ),
        )
        println(
            "Error between Tn0 and Tn0c  : ",
            sqrt(
                sum(
                    (Tn0[1:3, 4, :] - Tn0c[1:3, 4, :, 1]) .^ 2 +
                    (Tn0[1:3, 4, :] - Tn0c[1:3, 4, :, 2]) .^ 2,
                ) / N,
            ),
        )
        global ei = 0.0
        global ec = 0.0
        for i = 1:N
            global ei, ec
            T1 = fkineLPOE(Tn0mod[:, :, :, 1], xin[:, :, 1], q[i, :, 1])
            T2 = fkineLPOE(Tn0mod[:, :, :, 2], xin[:, :, 2], q[i, :, 2])
            ei += norm(twistcoords(logT(Ta[:, :, i] * trinv(T1))))
            ei += norm(twistcoords(logT(Ta[:, :, i] * trinv(T2))))
            T1 = fkineLPOE(Tn0c[:, :, :, 1], xin[:, :, 1], q[i, :, 1])
            T2 = fkineLPOE(Tn0c[:, :, :, 2], xin[:, :, 2], q[i, :, 2])
            ec += norm(twistcoords(logT(Ta[:, :, i] * trinv(T1))))
            ec += norm(twistcoords(logT(Ta[:, :, i] * trinv(T2))))
        end
        println(
            "Initial error: ",
            round(ei / N, digits = 5),
            " Calibrated error: ",
            round(ec / N, digits = 5),
        )
        @test ec < ei

    end

    @testset "calibAXYB" begin
        @info "Testing calibAXYB"
        include("test_calibAXYB.jl")
    end

    @testset "Frames" begin
        @info "Testing Frames"
        include("testFrames.jl")
    end
end

max_forks_repo_name: baggepinnen/Robotlib


text:
module JavaCall
export JavaObject, JavaMetaClass,
       jint, jlong, jbyte, jboolean, jchar, jshort, jfloat, jdouble,
       JObject, JClass, JMethod, JString,
       @jimport, jcall, jfield, isnull,
       getname, getclass, listmethods, getreturntype, getparametertypes, classforname,
       narrow

# using Compat, Compat.Dates

# using Sys: iswindows, islinux, isunix, isapple

import DataStructures: OrderedSet
import Libdl
using Dates

@static if Sys.iswindows()
    using WinReg
end

import Base: convert, unsafe_convert, unsafe_string

include("jnienv.jl")
include("jvm.jl")
include("core.jl")
include("convert.jl")
include("reflect.jl")

function __init__()
	findjvm()
	global create = Libdl.dlsym(libjvm, :JNI_CreateJavaVM)
end

end # module


text:
# Based off https://github.com/openpmix/openpmix/blob/4d07260d9f79bb7f328b1fc9107b45e683cf2c4e/examples/client.c
using PMIx

import PMIx: nspace

function main()
    pid = Libc.getpid()
    @info "Client running" pid

    # init us - note that the call to "init" includes the return of
    # any job-related info provided by the RM. This includes any
    # debugger flag instructing us to stop-in-init. If such a directive
    # is included, then the process will be stopped in this call until
    # the "debugger release" notification arrives

    myproc = PMIx.init()
    @info "Client initialized" ns=nspace(myproc) rank=myproc.rank pid

    # TODO: Register event handler

    # job-related info is found in our nspace, assigned to the
    # wildcard rank as it doesn't relate to a specific rank. Setup
    # a name to retrieve such values
    proc = PMIx.Proc(myproc.nspace, PMIx.API.PMIX_RANK_WILDCARD)

    # TODO: Check debugger
    # TODO: Check local topo

    # get our universe size
    value = PMIx.get(proc, PMIx.API.PMIX_UNIV_SIZE)
    @info "Universe size" ns=nspace(myproc) size=PMIx.get_number(value)

    PMIx.finalize()
    exit()
end

if !isinteractive()
    main()
end

max_issues_repo_name: JuliaParallel/PMIx

max_forks_repo_name: JuliaParallel/PMIx


text:
module AdjacencyGraphs

using ArrayInterface
using Base: @propagate_inbounds, OneTo
using LinearAlgebra
using Metadata
import Metadata: NoMetadata, no_metadata, metadata_keys, metadata_type, metadata
using Static
using SparseArrays
using SparseArrays: rowvals, getcolptr, AbstractSparseMatrixCSC, fkeep!

using LightGraphs
import LightGraphs: is_directed, add_edge!, rem_edge!, has_edge, vertices, edges,
    induced_subgraph, add_vertex!, add_vertices!, rem_vertex!, rem_vertices!, inneighbors,
    outneighbors, neighbors, indegree, outdegree, degree, all_neighbors, common_neighbors,
    density, num_self_loops, degree_histogram, weights

export AdjMap, AdjDiMap, WAdjMap, WAdjDiMap

const SUnitRange{F,L} = ArrayInterface.OptionallyStaticUnitRange{StaticInt{F},StaticInt{L}}

SUnitRange{F,L}() where {F,L} = static(F):static(L)

include("metadata.jl")
include("adjacency_map.jl")
include("interface.jl")
include("adjedge.jl")
include("core.jl")
include("utils.jl")

LightGraphs.nv(g::AdjacencyMap) = nvertices(g)

@inline LightGraphs.ne(g::AdjacencyMap) = nedges(g)

end


text:
#=============================================
Generate 𝐹 and ∇ₓ𝐹 from user input
=============================================#

"""
    F, ∇ₓF = state_function_and_Jacobian(Ts, Gs, nb)

Returns the state function `F` and its jacobian, `∇ₓF`.
"""
function state_function_and_Jacobian(Ts::Tuple, Gs::Tuple, nb)
    nt = length(Ts)
    tracers(x) = state_to_tracers(x, nb, nt)
    tracer(x,i) = state_to_tracer(x, nb, nt, i)
    T(p) = blockdiag([Tⱼ(p) for Tⱼ in Ts]...) # Big T (linear part)
    function G(x,p)
        xs = tracers(x)
        return reduce(vcat, [Gⱼ(xs..., p) for Gⱼ in Gs]) # nonlinear part
    end
    F(x,p) = G(x,p) - reduce(vcat, [Tⱼ(p) * tracer(x,j) for (j,Tⱼ) in enumerate(Ts)])
    ∇ₓG(x,p) = local_jacobian(Gs, x, p, nt, nb)     # Jacobian of nonlinear part
    ∇ₓF(x,p) = ∇ₓG(x,p) - T(p)       # full Jacobian ∇ₓ𝐹(𝑥) = -T + ∇ₓ𝐺(𝑥)
    return F, ∇ₓF
end
function state_function_and_Jacobian(T, G)
    F(x,p) = G(x,p) - T(p) * x                     # full 𝐹(𝑥) = -T 𝑥 + 𝐺(𝑥)
    ∇ₓG(x,p) = sparse(Diagonal(𝔇(G(x .+ ε, p))))     # Jacobian of nonlinear part
    ∇ₓF(x,p) = ∇ₓG(x,p) - T(p)       # full Jacobian ∇ₓ𝐹(𝑥) = -T + ∇ₓ𝐺(𝑥)
    return F, ∇ₓF
end

function inplace_state_function_and_Jacobian(Ts::Tuple, Gs::Tuple, nb)
    nt = length(Ts)
    tracers(x) = state_to_tracers(x, nb, nt)
    tracer(x,i) = state_to_tracer(x, nb, nt, i)
    T(p) = blockdiag([Tⱼ(p) for Tⱼ in Ts]...) # Big T (linear part)
    #F(x,p) = G(x,p) - T(p) * x                     # full 𝐹(𝑥) = -T 𝑥 + 𝐺(𝑥)
    function F!(dx,x,p)
        xs = tracers(x)
        for (j, (Tⱼ, Gⱼ!)) in enumerate(zip(Ts, Gs))
            ij = tracer_indices(nb,nt,j)
            @views dx[ij] .= Gⱼ!(dx[ij], xs..., p)
            @views dx[ij] .-= Tⱼ(p) * x[ij]
        end
        return dx
    end
    ∇ₓG(x,p) = inplace_local_jacobian(Gs, x, p, nt, nb)     # Jacobian of nonlinear part
    ∇ₓF(x,p) = ∇ₓG(x,p) - T(p)       # full Jacobian ∇ₓ𝐹(𝑥) = -T + ∇ₓ𝐺(𝑥)
    return F!, ∇ₓF
end

export state_function_and_Jacobian, inplace_state_function_and_Jacobian

"""
    F, ∇ₓF = state_function_and_Jacobian(Ts, Gs, nb)

Returns the state function `F` and its jacobian, `∇ₓF`.
"""
function split_state_function_and_Jacobian(Ts::Tuple, Ls::Tuple, NLs::Tuple, nb)
    nt = length(Ts)
    tracers(x) = state_to_tracers(x, nb, nt)
    T(p) = blockdiag([Tⱼ(p) for Tⱼ in Ts]...) # Big T (linear part)
    NL(x,p) = reduce(vcat, [NLⱼ(tracers(x)..., p) for NLⱼ in NLs]) # nonlinear part
    L(x,p) = reduce(vcat, [Lⱼ(tracers(x)..., p) for Lⱼ in Ls]) # nonlinear part
    F(x,p) = NL(x,p) + L(x,p) - T(p) * x                     # full 𝐹(𝑥) = -T 𝑥 + 𝐺(𝑥)
    ∇ₓNL(x,p) = local_jacobian(NLs, x, p, nt, nb)     # Jacobian of nonlinear part
    ∇ₓL(p) = local_jacobian(Ls, zeros(nt*nb), p, nt, nb)     # Jacobian of nonlinear part
    ∇ₓF(x,p) = ∇ₓNL(x,p) + ∇ₓL(p) - T(p)       # full Jacobian ∇ₓ𝐹(𝑥) = -T + ∇ₓ𝐺(𝑥)
    return F, L, NL, ∇ₓF, ∇ₓL, ∇ₓNL, T
end
function split_state_function_and_Jacobian(T, L, NL, nb)
    F(x,p) = NL(x,p) + L(x,p) - T(p) * x                     # full 𝐹(𝑥)
    ∇ₓNL(x,p) = sparse(Diagonal(𝔇(NL(x .+ ε,p))))     # Jacobian of nonlinear part
    ∇ₓL(p) = sparse(Diagonal(𝔇(L(zeros(nb) .+ ε,p))))     # Jacobian of nonlinear part
    ∇ₓF(x,p) = ∇ₓNL(x,p) + ∇ₓL(p) - T(p)       # full Jacobian ∇ₓ𝐹(𝑥) = -T + ∇ₓ𝐺(𝑥)
    return F, L, NL, ∇ₓF, ∇ₓL, ∇ₓNL, T
end
export split_state_function_and_Jacobian

function local_jacobian(Gs, x, p, nt, nb)
    return reduce(vcat, [local_jacobian_row(Gⱼ, x, p, nt, nb) for Gⱼ in Gs])
end
function inplace_local_jacobian(Gs, x, p, nt, nb)
    return reduce(vcat, [inplace_local_jacobian_row(Gⱼ, x, p, nt, nb) for Gⱼ in Gs])
end

𝔇(x) = DualNumbers.dualpart.(x)      # dual part

function local_jacobian_row(Gⱼ, x, p, nt, nb)
    e(j) = kron([j == k for k in 1:nt], trues(nb))
    tracers(x) = state_to_tracers(x, nb, nt)
    return reduce(hcat, [sparse(Diagonal(𝔇(Gⱼ(tracers(x + ε * e(j))..., p)))) for j in 1:nt])
end
function inplace_local_jacobian_row(Gⱼ, x, p, nt, nb)
    e(j) = kron([j == k for k in 1:nt], trues(nb))
    tracers(x) = state_to_tracers(x, nb, nt)
    return reduce(hcat, [sparse(Diagonal(𝔇(Gⱼ(Vector{Dual{Float64}}(undef,nb), tracers(x + ε * e(j))..., p)))) for j in 1:nt])
end

#=============================================
Generate 𝑓 and derivatives from user input
=============================================#

function generate_objective(ωs, μx, σ²x, v, ωp, μp, σ²p)
    nt, nb = length(ωs), length(v)
    tracers(x) = state_to_tracers(x, nb, nt)
    f(x, p) = ωp * mismatch(p, μp, σ²p) +
        sum([ωⱼ * mismatch(xⱼ, μⱼ, σⱼ², v) for (ωⱼ, xⱼ, μⱼ, σⱼ²) in zip(ωs, tracers(x), μx, σ²x)])
    return f
end

function generate_∇ₓobjective(ωs, μx, σ²x, v, ωp, μp, σ²p)
    nt, nb = length(ωs), length(v)
    tracers(x) = state_to_tracers(x, nb, nt)
    ∇ₓf(x, p) = reduce(hcat, [ωⱼ * ∇mismatch(xⱼ, μⱼ, σⱼ², v) for (ωⱼ, xⱼ, μⱼ, σⱼ²) in zip(ωs, tracers(x), μx, σ²x)])
    return ∇ₓf
end

function generate_∇ₚobjective(ωs, μx, σ²x, v, ωp, μp, σ²p)
    nt, nb = length(ωs), length(v)
    tracers(x) = state_to_tracers(x, nb, nt)
    ∇ₚf(x, p) = ωp * ∇mismatch(p, μp, σ²p)
    return ∇ₚf
end

generate_objective_and_derivatives(ωs, μx, σ²x, v, ωp, μp, σ²p) =
      generate_objective(ωs, μx, σ²x, v, ωp, μp, σ²p),
    generate_∇ₓobjective(ωs, μx, σ²x, v, ωp, μp, σ²p),
    generate_∇ₚobjective(ωs, μx, σ²x, v, ωp, μp, σ²p)

export generate_objective, generate_∇ₓobjective, generate_∇ₚobjective
export generate_objective_and_derivatives

"""
    mismatch(x, xobs, σ²xobs, v)

Volume-weighted mismatch of modelled tracer `x` against observed mean, `xobs`, given observed variance, `σ²xobs`, and volumes `v`.
"""
function mismatch(x, xobs, σ²xobs, v)
    δx = x - xobs
    W = Diagonal(v ./ σ²xobs)
    return 0.5 * transpose(δx) * W * δx / (transpose(xobs) * W * xobs)
end

mismatch(x, ::Missing, args...) = 0

"""
    ∇mismatch(x, xobs, σ²xobs, v)

Adjoint of the gradient of `mismatch(x, xobs, σ²xobs, v)`.
"""
function ∇mismatch(x, xobs, σ²xobs, v)
    δx = x - xobs
    W = Diagonal(v ./ σ²xobs)
    return transpose(W * δx) / (transpose(xobs) * W * xobs)
end
∇mismatch(x, ::Missing, args...) = transpose(zeros(length(x)))

# TODO
# Talk about it with FP
# Assumptions:
# 1. The prior distributions of the parameters, p, are log-normal
# 2. The values `mean_obs` and `variance_obs` are the non-log mean and variance,
# Then the mean and variance of the prior of log(p) are
# logμ = log(μ / √(1 + σ² / μ²))
# logσ² = log(1 + σ² / μ²)
# These are the values we use for the mismatch
"""
    mismatch(p, m, v)

Returns the mismatch of the model parameters `p` against observations.
Assumes priors have a log-normal distributions.
`m` and `v` are the non-log mean and variances,
and are converted to their log counterparts in the mismatch formula.
"""
function mismatch(p, m, v)
    μ = LNμ(m, v)
    σ² = LNσ²(m, v)
    δλ = log.(optvec(p)) .- μ
    W = Diagonal(1 ./ σ²)
    return 0.5 * transpose(δλ) * W * δλ
end
function ∇mismatch(p, m, v)
    μ = LNμ(m, v)
    σ² = LNσ²(m, v)
    δλ = log.(optvec(p)) .- μ
    W = Diagonal(1 ./ σ²)
    return transpose(W * δλ ./ optvec(p))
end

# Lognormal mean and variance and variance conversions (from Wikipedia)
LNμ(m, v) = log.(m ./ sqrt.(1 .+ v ./ m.^2))
LNσ²(m, v) = log.(1 .+ v ./ m.^2)
LNm(μ, σ²) = exp.(μ + σ² / 2)
LNv(μ, σ²) = (exp.(σ²) .- 1) .* exp.(2μ + σ²)

#=============================================
Generate multi-tracer norm
=============================================#

function volumeweighted_norm(nt, v)
    w = repeat(v, nt)
    return nrm(x) = transpose(x) * Diagonal(w) * x
end


text:
using Revise
using Images
using CUDA, CUDAKernels, KernelAbstractions
using Tullio
CUDA.allowscalar(false)
using ShirleyRayTracer

function add_random_scene!(scene)

	push!(scene, Sphere(Point3(0,-1000,0), 1000.0, Lambertian(0.5, 0.5, 0.5)))

	rand_material(p) = if p < 0.8
				Lambertian()
			elseif p < 0.95
				rf = randf(0.5, 1)
				Metal(rf, rf, rf, 0.5rand())
			else
				Dielectric(1.5)
			end

	for a in -11:10, b in -11:10
		center = Point3(a + 0.9rand(), 0.2, b + 0.9rand())
		if ShirleyRayTracer.magnitude(center - Point3(4, 0.2, 0)) > 0.9
			push!(scene, Sphere(center, 0.2, rand_material(rand())))
		end
	end

	push!(scene, Sphere(Point3(0, 1, 0), 1.0, Dielectric(1.5)))
	push!(scene, Sphere(Point3(-4, 1, 0), 1.0, Lambertian(0.4,0.2,0.1)))
	push!(scene, Sphere(Point3(4, 1, 0), 1.0, Metal(0.7,0.6,0.5, 0.0)))
end

function add_random_scene_lambertian!(S, scene)

	# S = Sphere
	# S = Sphere{ShirleyRayTracer.Material}
	# S = Sphere{ShirleyRayTracer.Lambertian}

	push!(scene, S(Point3(0,-1000,0), 1000.0, Lambertian(0.5, 0.5, 0.5)))

	rand_material(p) = if p < 0.8
				Lambertian(0.5,0.3,0.5)
			elseif p < 0.95
				rf = randf(0.5, 1)
				Lambertian(rf, rf, 0.5rand())
			else
				Lambertian(0.1,0.1,0.1)
			end

	function rand_hitable(center, p)
		# if p < 0.5
			S(center, 0.2, rand_material(rand()))
		# else
		# 	ShirleyRayTracer.Sphere2(center, 0.2, rand_material(rand()))
		# end
	end

	for a in -11:10, b in -11:10
		center = Point3(a + 0.9rand(), 0.2, b + 0.9rand())
		if ShirleyRayTracer.magnitude(center - Point3(4, 0.2, 0)) > 0.9
			push!(scene, rand_hitable(center, rand()))
		end
	end

	push!(scene, S(Point3(0, 1, 0), 1.0, Lambertian(0.5,0.3,0.1)))
	push!(scene, S(Point3(-4, 1, 0), 1.0, Lambertian(0.4,0.2,0.1)))
	push!(scene, S(Point3(4, 1, 0), 1.0, Lambertian(0.7,0.6,0.5)))
end

struct DummyHitable{T} <: ShirleyRayTracer.Hitable; end

filename="render.png"
image_width=640
aspect_ratio=16/9
samples_per_pixel=5
max_depth=10
image_height = round(Int, image_width / aspect_ratio)
# function main(;filename="render.png", image_width=640, aspect_ratio=16/9, samples_per_pixel=5, max_depth=5)

image_height = round(Int, image_width / aspect_ratio)

# S = Sphere
# S = Sphere{ShirleyRayTracer.Material}
# S = Sphere{ShirleyRayTracer.Lambertian}

# H = ShirleyRayTracer.Hitable
# H = Sphere
# H = Sphere{ShirleyRayTracer.Material}
H = Union{Sphere{Lambertian}, Sphere{Metal}, Sphere{Dielectric}}
# H = Sphere{Lambertian}

# H = Union{H, Sphere{ShirleyRayTracer.Material}}
# H = Union{H, DummyHitable}
# H = Union{H, DummyHitable{Int}}

hitables = Vector{H}()

add_random_scene!(hitables)
# add_random_scene_lambertian!(S, hitables);
# push!(hitables, S(Point3(0,-1000,0), 1000.0, Lambertian(0.5, 0.5, 0.5)))

cam = Camera(Point3(13.,2.,3.), zero(Point3), Vec3(0,1,0), 20, aspect_ratio, 0.1, 10.0)

scene = Scene(cam, hitables)

image = Array{RGB{Float32}, 2}(undef, image_height, image_width);

scene = scene |> cu;
image = image |> cu;

nothing

##

@time begin
	render!(image, scene, image_width, image_height, samples_per_pixel, max_depth)
	isa(image, CuArray) && CUDA.synchronize()
end

@time image_out = image |> Array;
save(filename, image_out)


text:
##########################
# Homogenous wind stress #
# Constant temperature   #
##########################

"""
    HomogeneousBox <: AbstractSimpleBoxProblem

Container structure for a simple box problem with wind-stress.
Lˣ = zonal (east-west) length
Lʸ = meridional (north-south) length
H  = height of the ocean
τₒ = maximum value of wind-stress (amplitude)
"""
struct HomogeneousBox{T, BC} <: AbstractSimpleBoxProblem
    Lˣ::T
    Lʸ::T
    H::T
    τₒ::T
    boundary_conditions::BC
    function HomogeneousBox{FT}(
        Lˣ,             # m
        Lʸ,             # m
        H;              # m
        τₒ = FT(1e-1),  # N/m²
        BC = (
            OceanBC(Impenetrable(NoSlip()), Insulating()),
            OceanBC(Impenetrable(NoSlip()), Insulating()),
            OceanBC(Penetrable(KinematicStress()), Insulating()),
        ),
    ) where {FT <: AbstractFloat}
        return new{FT, typeof(BC)}(Lˣ, Lʸ, H, τₒ, BC)
    end
end

"""
    ocean_init_state!(::HomogeneousBox)

initialize u,v with random values, η with 0, and θ with a constant (20)

# Arguments
- `p`: HomogeneousBox problem object, used to dispatch on
- `Q`: state vector
- `A`: auxiliary state vector, not used
- `localgeo`: the local geometry, not used
- `t`: time to evaluate at, not used
"""
function ocean_init_state!(m::HBModel, p::HomogeneousBox, Q, A, localgeo, t)
    Q.u = @SVector [0, 0]
    Q.η = 0
    Q.θ = 20

    return nothing
end

"""
    kinematic_stress(::HomogeneousBox)

jet stream like windstress

# Arguments
- `p`: problem object to dispatch on and get additional parameters
- `y`: y-coordinate in the box
"""
@inline kinematic_stress(p::HomogeneousBox, y, ρ) =
    @SVector [(p.τₒ / ρ) * cos(y * π / p.Lʸ), -0]

@inline kinematic_stress(p::HomogeneousBox, y) =
    @SVector [-p.τₒ * cos(π * y / p.Lʸ), -0]


text:
@testset "categorical" begin
    x = 1:5
    @test MLJModelInterface.categorical(x) == categorical(x)
end
@testset "int, classes, decoder" begin
    N = 10
    mix = shuffle(0:N - 1)

    Xraw = broadcast(x->mod(x,N), rand(Int, 2N, 3N))
    Yraw = string.(Xraw)

    # to turn a categ matrix into a ordinary array with categorical
    # elements. Needed because broacasting the identity gives a
    # categorical array in CategoricalArrays >0.5.2
    function matrix_(X)
        ret = Array{Any}(undef, size(X))
        for i in eachindex(X)
            ret[i] = X[i]
        end
        return ret
    end

    X = categorical(Xraw)
    x = X[1]
    Y = categorical(Yraw)
    y = Y[1]
    V = matrix_(X)
    W = matrix_(Y)
    Xo = levels!(deepcopy(X), mix)
    xo = Xo[1]
    Yo = levels!(deepcopy(Y), string.(mix))
    yo = Yo[1]
    Vo = matrix_(Xo)
    Wo = matrix_(Yo)

    raw(x::MLJBase.CategoricalElement) = x.pool.index[x.level]

    @test raw.(classes(xo)) == xo.pool.levels
    @test raw.(classes(yo)) == yo.pool.levels

    # getting all possible elements from one:
    @test raw.(X) == Xraw
    @test raw.(Y) == Yraw
    @test raw.(classes(xo)) == levels(Xo)
    @test raw.(classes(yo)) == levels(Yo)

    # broadcasted encoding:
    @test int(X) == int(V)
    @test int(Y) == int(W)

    @test int(X; type=Int8) isa AbstractArray{Int8}

    # encoding is right-inverse to decoding:
    d = decoder(xo)
    @test d(int(Vo)) == Vo # ie have the same elements
    e = decoder(yo)
    @test e(int(Wo)) == Wo

    # classes return value is in correct order:
    @test int(classes(xo)) == 1:length(classes(xo))

    # int is based on ordering not index
    v = categorical(['a', 'b', 'c'], ordered=true)
    @test int(v) == 1:3
    levels!(v, ['c', 'a', 'b'])
    @test int(v) == [2, 3, 1]

    # Errors
    @test_throws DomainError int("g")
end

@testset "matrix, table" begin
    B = rand(UInt8, (4, 5))
    @test matrix(DataFrame(B)) == B
    @test matrix(table(B)) == B
    @test matrix(table(B), transpose=true) == B'

    X  = (x1=rand(5), x2=rand(5))

    @test table(X, prototype=DataFrame()) == DataFrame(X)
    T = table((x1=(1,2,3), x2=(:x, :y, :z)))
    @test selectcols(T, :x1) == [1, 2, 3]

    v = categorical(11:20)
    A = hcat(v, v)
    tab = table(A)
    @test selectcols(tab, 1) == v

    @test matrix(B) == B
    @test matrix(B, transpose=true) == permutedims(B)
end

@testset "select etc" begin
    N = 10
    A = broadcast(x->Char(65+mod(x,5)), rand(Int, N, 5))
    X = CategoricalArrays.categorical(A)
    df = DataFrame(A)
    df.z  = 1:N
    tt = TypedTables.Table(df)
    rt = Tables.rowtable(tt)
    ct = Tables.columntable(tt)

    @test selectcols(nothing, 4:6)   == nothing
    @test selectcols(df, 4:6)        == df[:,4:6]
    @test selectcols(df, [:x1, :z])  == df[:,[:x1, :z]]
    @test selectcols(df, :x2)        == df.x2
    @test selectcols(nothing, 4:6)   == nothing
    @test selectcols(df, 2)          == df.x2
    @test selectrows(df, 4:6)        == selectrows(df[4:6, :], :)
    @test selectrows(df, 1)          == selectrows(df[1:1, :], :)
    @test selectrows(nothing, 4:6)   == nothing
    @test MLJBase.select(df, 2, :x2) == df[2, :x2]
    @test MLJBase.select(nothing, 2, :x) == nothing
    s = schema(df)
    @test nrows(df) == N

    @test selectcols(tt, 4:6) ==
        selectcols(TypedTables.Table(x4=tt.x4, x5=tt.x5, z=tt.z), :)
    @test selectcols(tt, [:x1, :z]) ==
        selectcols(TypedTables.Table(x1=tt.x1, z=tt.z), :)
    @test selectcols(tt, :x2) == tt.x2
    @test selectcols(tt, 2)   == tt.x2
    @test selectrows(tt, 4:6) == selectrows(tt[4:6], :)
    @test nrows(tt) == N
    @test MLJBase.select(tt, 2, :x2) == tt.x2[2]

    @test selectrows(rt, 4:6) == rt[4:6]
    @test selectrows(rt, :)   == rt

    @test selectrows(rt, 5)   == rt[5,:]
    @test nrows(rt) == N

    @test Tables.rowtable(selectrows(ct, 4:6)) == rt[4:6]
    @test selectrows(ct, :) == ct
    @test Tables.rowtable(selectrows(ct, 5))[1] == rt[5,1]

    # vector accessors
    v = rand(Int, 4)
    @test selectrows(v, 2:3) == v[2:3]
    @test selectrows(v, 2)   == [v[2]]
    @test nrows(v)           == 4

    v = categorical(collect("asdfasdf"))
    @test selectrows(v, 2:3) == v[2:3]
    @test selectrows(v, 2)   == [v[2]]
    @test nrows(v) == 8

    # matrix accessors
    A = rand(5, 10)
    @test selectrows(A, 2:4) == A[2:4,:]
    @test selectrows(A, 2:4) == A[2:4,:]
    @test selectrows(A, 2)   == A[2:2,:]

    A = rand(5, 10) |> categorical
    @test selectrows(A, 2:4) == A[2:4,:]
    @test selectrows(A, 2:4) == A[2:4,:]
    @test selectrows(A, 2)   == A[2:2,:]

    @test nrows(A) == 5

    # TypedTables
    v = categorical(collect("asdfasdf"))
    df = DataFrame(v=v, w=v)
    @test selectcols(df, :w) == v
    tt = TypedTables.Table(df)
    @test selectcols(tt, :w) == v
end