using Distributions
using Base.Test

# Core testing procedure

function test_mixture(g::UnivariateMixture, n::Int, ns::Int)
    X = zeros(n)
    for i = 1:n
        X[i] = rand(g)
    end

    K = ncomponents(g)
    pr = probs(g)
    @assert length(pr) == K

    # mean
    mu = 0.0
    for k = 1:K
        mu += pr[k] * mean(component(g, k))
    end
    @test mean(g) ≈ mu

    # evaluation of cdf
    cf = zeros(n)
    for k = 1:K
        c_k = component(g, k)
        for i = 1:n
            cf[i] += pr[k] * cdf(c_k, X[i])
        end
    end

    for i = 1:n
        @test cdf(g, X[i]) ≈ cf[i]
    end
    @test cdf.(g, X) ≈ cf

    # evaluation
    P0 = zeros(n, K)
    LP0 = zeros(n, K)
    for k = 1:K
        c_k = component(g, k)
        for i = 1:n
            x_i = X[i]
            P0[i,k] = pdf(c_k, x_i)
            LP0[i,k] = logpdf(c_k, x_i)
        end
    end

    mix_p0 = P0 * pr
    mix_lp0 = log.(mix_p0)

    for i = 1:n
        @test pdf(g, X[i])                  ≈ mix_p0[i]
        @test logpdf(g, X[i])               ≈ mix_lp0[i]
        @test componentwise_pdf(g, X[i])    ≈ vec(P0[i,:])
        @test componentwise_logpdf(g, X[i]) ≈ vec(LP0[i,:])
    end

    @test pdf.(g, X)                  ≈ mix_p0
    @test logpdf.(g, X)               ≈ mix_lp0
    @test componentwise_pdf(g, X)    ≈ P0
    @test componentwise_logpdf(g, X) ≈ LP0

    # sampling
    Xs = rand(g, ns)
    @test isa(Xs, Vector{Float64})
    @test length(Xs) == ns
    @test isapprox(mean(Xs), mean(g), atol=0.01)
end

function test_mixture(g::MultivariateMixture, n::Int, ns::Int)
    X = zeros(length(g), n)
    for i = 1:n
        X[:,i] = rand(g)
    end

    K = ncomponents(g)
    pr = probs(g)
    @assert length(pr) == K

    # mean
    mu = 0.0
    for k = 1:K
        mu += pr[k] * mean(component(g, k))
    end
    @test mean(g) ≈ mu

    # evaluation
    P0 = zeros(n, K)
    LP0 = zeros(n, K)
    for k = 1:K
        c_k = component(g, k)
        for i = 1:n
            x_i = X[:,i]
            P0[i,k] = pdf(c_k, x_i)
            LP0[i,k] = logpdf(c_k, x_i)
        end
    end

    mix_p0 = P0 * pr
    mix_lp0 = log.(mix_p0)

    for i = 1:n
        x_i = X[:,i]
        @test pdf(g, x_i)                  ≈ mix_p0[i]
        @test logpdf(g, x_i)               ≈ mix_lp0[i]
        @test componentwise_pdf(g, x_i)    ≈ vec(P0[i,:])
        @test componentwise_logpdf(g, x_i) ≈ vec(LP0[i,:])
    end

    @show g
    @show size(X)
    @show size(mix_p0)
    @test pdf(g, X)                  ≈ mix_p0
    @test logpdf(g, X)               ≈ mix_lp0
    @test componentwise_pdf(g, X)    ≈ P0
    @test componentwise_logpdf(g, X) ≈ LP0

    # sampling
    Xs = rand(g, ns)
    @test isa(Xs, Matrix{Float64})
    @test size(Xs) == (length(g), ns)
    @test isapprox(vec(mean(Xs, 2)), mean(g), atol=0.1)
    @test isapprox(cov(Xs, 2)      , cov(g) , atol=0.1)
end

function test_params(g::AbstractMixtureModel)
    C = eltype(g.components)
    pars = params(g)
    mm = MixtureModel(C, pars...)
    @test g.prior == mm.prior
    @test g.components == mm.components
end

function test_params(g::UnivariateGMM)
    pars = params(g)
    mm = UnivariateGMM(pars...)
    @test g == mm
end

# Tests

println("    testing UnivariateMixture")

g_u = MixtureModel(Normal, [(0.0, 1.0), (2.0, 1.0), (-4.0, 1.5)], [0.2, 0.5, 0.3])
@test isa(g_u, MixtureModel{Univariate, Continuous, Normal})
@test ncomponents(g_u) == 3
test_mixture(g_u, 1000, 10^6)
test_params(g_u)
@test minimum(g_u) == -Inf
@test maximum(g_u) == Inf

g_u = MixtureModel([TriangularDist(-1,2,0),TriangularDist(-.5,3,1),TriangularDist(-2,0,-1)])
@test minimum(g_u) == -2.0
@test maximum(g_u) == 3.0
@test insupport(g_u, 2.5) == true
@test insupport(g_u, 3.5) == false

g_u = UnivariateGMM([0.0, 2.0, -4.0], [1.0, 1.2, 1.5], Categorical([0.2, 0.5, 0.3]))
@test isa(g_u, UnivariateGMM)
@test ncomponents(g_u) == 3
test_mixture(g_u, 1000, 10^6)
test_params(g_u)
@test minimum(g_u) == -Inf
@test maximum(g_u) == Inf

println("    testing MultivariateMixture")
g_m = MixtureModel(
    IsoNormal[ MvNormal([0.0, 0.0], 1.0),
               MvNormal([0.2, 1.0], 1.0),
               MvNormal([-0.5, -3.0], 1.6) ],
    [0.2, 0.5, 0.3])
@test isa(g_m, MixtureModel{Multivariate, Continuous, IsoNormal})
@test length(components(g_m)) == 3
@test length(g_m) == 2
@test insupport(g_m, [0.0, 0.0]) == true
test_mixture(g_m, 1000, 10^6)
test_params(g_m)


module TestVariables

using TinyModia
using ModiaPlot
using Measurements
using StaticArrays
using Test

@testset "... Test variable declarations" begin

v1 = input | output | flow | Var(min=0, nominal=10, init=5)
@test v1 == (class = :Var, input = true, output = true, flow = true, min = 0, nominal = 10, init = 5)

v2 = parameter | 5 | Var(info="v2 info")
@test v2 == (class = :Var, parameter = true, value = 5, info = "v2 info")

v3 = "string value" | parameter | info"v3 info"
@test v3 == (class = :Var, parameter = true, value = "string value", info = "v3 info")

v4 = parameter | 5.5u"m/s" | Var(min=0)
@test v4 == (class = :Var, parameter = true, value = 5.5u"m*s^-1", min = 0)

v5 = parameter | Var(5, min=0)
@test v5 == (class = :Var, parameter = true, value = 5, min = 0)

Pin = Model( v = potential | Var(nominal=10), i = flow )
@test Pin == (class = :Model, v = (class = :Var, potential = true, nominal = 10), i = (class = :Var, flow = true,))

v6 = Var(2.0)
@test v6 == (class = :Var, value = 2.0,)

v7 = Var(2.0 ± 0.1)
@test v7 == (class = :Var, value = 2.0 ± 0.1,)

v8 = Var(SVector(1.0, 2.0, 3.0))
@test v8 == (class = :Var, value = [1.0, 2.0, 3.0],)

v9 = v4 | Var(parameter=false)
@test v9 == (class = :Var, parameter = false, value = 5.5u"m*s^-1", min = 0)

v10 = v4 | Var(parameter=nothing)
@test v10 == (class = :Var, value = 5.5u"m*s^-1", min = 0)

v11 = v4 | interval(1,10)
@test v11 == (class = :Var, parameter = true, value = 5.5u"m*s^-1", min = 1, max = 10)

TestVar1 = Model(
    p = parameter | -1 | info"A parameter",
    s = parameter | "string value" | info"A string parameter",
    x = Var(init=0.1),
    y = Var(start=0.1), # Not used
    equations = :[
        der(x) = p*x+1
    ]
)
#@showModel TestVar1
model = @instantiateModel(TestVar1, log=false, logCode=false)
simulate!(model, merge=Map(p=-2, x=0.2), requiredFinalStates = [0.4593994150057028]) # x.init is not changed
plot(model, "x")

end

end


@testset "Weighting" begin
  @testset "BlockWeighter" begin
    sdomain = RegularGrid(100,100)
    sdata = georef(DataFrame(z=rand(10000)), RegularGrid(100,100))
    for d in [sdomain, sdata]
      w = weight(d, BlockWeighter(10,10))
      @test length(unique(w)) == 1
      @test w[1] == 1 / 100
    end
  end

  @testset "DensityRatioWeighter" begin
    n = 1000; Random.seed!(123)

    r1 = Normal(0, 2)
    r2 = MixtureModel([Normal(-2,1), Normal(2,2)], [0.2, 0.8])

    z1 = sort(rand(r1, n))
    z2 = sort(rand(r2, n))

    d1 = georef(DataFrame(z=z1), PointSet(reshape(1:n,1,:)))
    d2 = georef(DataFrame(z=z2), PointSet(reshape(1:n,1,:)))

    w = weight(d1, DensityRatioWeighter(d2))

    if visualtests
      @plottest begin
        plot( z1, pdf.(r1, z1), size=(800,400), label="source")
        plot!(z1, pdf.(r2, z1), label="target")
        plot!(z1, w .* pdf.(r1, z1), label="approx")
      end joinpath(datadir,"density-ratio.png") !istravis
    end
  end
end


using LowRankModels

# test losses in losses.jl
srand(1);

losses = [
QuadLoss(),
QuadLoss(10),
L1Loss(),
L1Loss(5.2),
HuberLoss(),
HuberLoss(4),
HuberLoss(3.1, crossover=3.2),
PeriodicLoss(2*pi),
PeriodicLoss(2*pi, 4),
PoissonLoss(20),
PoissonLoss(22,4.1),
OrdinalHingeLoss(1,10),
OrdinalHingeLoss(2,7,5),
LogisticLoss(),
LogisticLoss(0.2),
WeightedHingeLoss(),
WeightedHingeLoss(11),
WeightedHingeLoss(1.5, case_weight_ratio=4.3),
MultinomialLoss(4),
MultinomialLoss(6, .5),
# OrdisticLoss(5),
MultinomialOrdinalLoss(3)
] #tests what should be successful constructions

# TODO: do some bad constructions and test that they fail with catches
bad_losses = [
:(QuadLoss(10,RealDomain)),
:(HuberLoss(3.1, 3.2)),
:(PeriodicLoss(scale=2*pi)),
:(PeriodicLoss(2*pi, scale=4)),
:(PeriodicLoss())
]
for expression in bad_losses
	try
		eval(expression);
		println("test FAILED for $expression")
	catch
		println("test PASSED for $expression (failed to construct)")
	end
end

m,n,k = 1000, length(losses), 5;
d = embedding_dim(losses)
X_real, Y_real = 2*randn(m,k), 2*randn(k,d);
XY_real = X_real*Y_real;

# tests default imputations and implicit domains
# we can visually inspect the differences between A and A_real to make sure imputation is right
A = impute(losses, XY_real)

regscale = 1
yregs = Array(Regularizer, length(losses))
for i=1:length(losses)
	if typeof(losses[i]) == MultinomialOrdinalLoss ||
		typeof(losses[i]) == OrdisticLoss
		yregs[i] = OrdinalReg(QuadReg(regscale))
	else
		yregs[i] = QuadReg(regscale)
	end
end

# tests all the M-estimators with scale=false, offset=false
glrm = GLRM(A, losses, QuadReg(regscale), yregs, 5, scale=false, offset=false);

# interestingly adding an offset to a model with multidimensional ordinal data causes a segfault
# but let's test the offset for everything but ordinals
# oops we still get a segfault...
# tamecols = [typeof(losses[i]) !== MultinomialOrdinalLoss &&
# 			typeof(losses[i]) !== OrdisticLoss
# 			for i=1:length(losses)]
# glrm = GLRM(A[:, tamecols],
# 	losses[tamecols],
# 	QuadReg(regscale),
# 	yregs[tamecols],
# 	5, scale=false, offset=true)

# tests eval and grad
@time X,Y,ch = fit!(glrm);

# tests initialization
init_svd!(glrm)


using KernelStat
using ARIMA
using DependentBootstrap

#-------------------------------
# SIMULATION FUNCTIONS
#-------------------------------
#Simulate kernel functions
const numObsBandwidthSim = 500::Int
const maxCorLag = 6::Int
const kernelList = [KernelUniform(), KernelEpanechnikov(), KernelGaussian(), KernelPR1993FlatTop(), KernelP2003FlatTop(), KernelPP2002Trap(), KernelPP2002Smooth(), KernelPR1994SB(numObsBandwidthSim)]::Vector{KernelFunction}
const bandwidthList = [BandwidthWhiteNoise(), BandwidthBartlett(), BandwidthP2003(numObsBandwidthSim)]
const hacList = [HACVarianceBasic(KernelUniform(0, maxCorLag), BandwidthP2003(numObsBandwidthSim)), HACVarianceBasic(KernelGaussian(1), BandwidthP2003(numObsBandwidthSim)), HACVarianceBasic(KernelEpanechnikov(1), BandwidthP2003(numObsBandwidthSim)), HACVarianceBasic(KernelBartlett(1), BandwidthP2003(numObsBandwidthSim)), HACVarianceBasic(KernelPR1994SB(numObsBandwidthSim, 0.1), BandwidthMax())]

#------------------------------
# TEST FUNCTIONS
#------------------------------
function testevaluate()
	for k = 1:length(kernelList)
		kF = kernelList[k]
		aD = activedomain(kF)
		if aD.start == -Inf || aD.stop == Inf
			aD.start == -Inf ? (tempStart = -10) : (tempStart = aD.start)
			aD.stop == Inf ? (tempStop = 10) : (tempStop = aD.stop)
			aD = UnitRange(tempStart, tempStop)
		end
		xIn = [aD.start:(aD.stop - aD.start)/10:aD.stop]
		xOut = [ evaluate(xIn[n], kF) for n = 1:length(xIn) ]
		println("Kernel func = " * string(kF) * ". evaluate of active domain = " * string(xOut))
	end
	println("")
	println("")
	return(true)
end

function testbandwidth()
	for k = 0:maxCorLag
		k == 0 ? maCoef = Array(Float64, 0) : maCoef = [ 1 / 2^p for p = 1:k ]
		for j = 1:length(bandwidthList)
			bC = bandwidthList[j]
			mVec = Array(Int, 10)
			for i = 1:10
				x = simulate(numObsBandwidthSim, ARIMAModel(maCoef=maCoef))
				(mVec[i], xVar, xCovVec) = bandwidth(x, bC)
			end
			println("Bandwidth method = " * string(bC))
			println("    MA Order = " * string(k))
			println("    Bandwidth estimates = " * string(mVec))
			println("")
		end
	end
	println("")
	println("")
	return(true)
end

function testhacvariance()
	for k = 0:maxCorLag
		k == 0 ? maCoef = Array(Float64, 0) : maCoef = [ 1 / 2^p for p = 1:k ]
		x = simulate(numObsBandwidthSim, ARIMAModel(maCoef=maCoef))
		meanBoot = dbootstrapstatistic(x, BootstrapParam(x, statistic=mean))
		bootTrueVar = var(sqrt(numObsBandwidthSim) * meanBoot)
		xMean = [ mean(simulate(numObsBandwidthSim, ARIMAModel(maCoef=maCoef))) for i = 1:2000 ]
		simTrueVar = var(sqrt(numObsBandwidthSim) * xMean)
		for j = 1:length(hacList)
			hC = hacList[j]
			vVec = Array(Float64, 10)
			bwVec = Array(Int, 10)
			for i = 1:10
				x = simulate(numObsBandwidthSim, ARIMAModel(maCoef=maCoef))
				(vVec[i], bwVec[i]) = hacvariance(x, hC)
			end
			show(hC)
			println("True MA Order = " * string(k))
			println("Bandwidth estimates = " * string(bwVec))
			println("Simulated true variance estimate = " * string(simTrueVar))
			println("Bootstrapped true variance estimate = " * string(bootTrueVar))
			println("HAC estimates = " * string(vVec))
			println("")
		end
	end
	println("")
	println("")
	return(true)
end

testhacvariance()

function testall(KK::Int)
	testevaluate()
	testbandwidth()
	testhacvariance()
	println("Tests complete.")
end


todisplay = VERSION <= VersionNumber("1.2") ? true : false

using Random
using SparseArrays
using LinearAlgebra
using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation
using CAGD
if todisplay
	using ViewerGL
	GL = ViewerGL
	include("../examples/3d/views.jl")
end

function build_copFE(FV::Lar.Cells, EV::Lar.Cells)
	copFE = Lar.u_coboundary_1(FV, EV) # unsigned
	faceedges = [findnz(copFE[f,:])[1] for f=1:size(copFE,1)]

	f_edgepairs = Array{Array{Int64,1}}[]
	for f=1:size(copFE,1)
		edgepairs = Array{Int64,1}[]
		for v in FV[f]
			push!(edgepairs, [e for e in faceedges[f] if v in EV[e]])
		end
		push!(f_edgepairs, edgepairs)
	end
	copFE.=0
	for f=1:size(copFE,1)
		facepairs = sort(sort.(f_edgepairs[f]))
		# setting a reference
		copFE[f, facepairs[1][1]] = 1
		for (e1,e2) in facepairs
			@assert copFE[f, e1] != 0
			if copFE[f, e2] == 0
				copFE[f, e2] =
					-((EV[e1][1]==EV[e2][1])||(EV[e1][2]==EV[e2][2])) * copFE[f,e1]+
					((EV[e1][1]==EV[e2][2])||(EV[e1][2]==EV[e2][1])) * copFE[f,e1]
			end
		end
	end
	return copFE
end

Random.seed!(1234);

scaling = 1.2
no_sphere = 10
expl = 2.0
pt1, pt2 = 5, 10
#pt1, pt2 = 10, 10

carry = Array{Tuple{Array{Float64,2},Array{Array{Int64,1},1}}}(undef, no_sphere)
for i = 1 : no_sphere
	carry[i] = Lar.struct2lar(Lar.Struct([
		Lar.t(rand(0.5:1e-2:5, 3)...),
		Lar.sphere(rand(1:1e-2:3) * scaling)([pt1, pt2])
	]))
end
V, FV = Lar.struct2lar(Lar.Struct(carry))
FV = sort(sort.(FV))

model = CAGD.Model(V);
function getFaceEdges(fV::Array{Int,1})
    return [fV[1], fV[2]], [fV[1], fV[3]], [fV[2], fV[3]]
end
EV = unique([(map(f -> getFaceEdges(f), FV)...)...])
FE = [[3*i+1, 3*i+2, 3*i+3] for i = 0 : length(FV)-1]
CAGD.addModelCells!(model, 1, EV, signed = true)
CAGD.addModelCells!(model, 2, build_copFE(FV, EV))
model = CAGD.mergeModelVertices(model, signed_merge=true)

if todisplay  displayModel(model)  end

split_model = CAGD.facesplitting(model)

if todisplay  displayModel(split_model)  end

congr_model = CAGD.mergeModelVertices(split_model, signed_merge=true)

if todisplay  displayModel(congr_model)  end

gift_model = deepcopy(congr_model)
FC, bicon_comps = CAGD.tgw(congr_model, 3)
CAGD.addModelCells!(gift_model, 3, convert(Lar.ChainOp, FC'))

if todisplay  viewExplode(gift_model)  end


using ExponentialAction, LinearAlgebra, SparseArrays, Test

Tsingle = (Float32, ComplexF32)
Tdouble = (Float64, ComplexF64)

@testset "expv" begin
    n = 10
    @testset "expv(t::$Tt, A::$MT{$TA}, B::Array{$TB,$(length(Bdims2)+1)}), tscale=$tscale, shift=$shift" for
        MT in (Matrix, Diagonal, Bidiagonal, Tridiagonal),
        Tset in (Tsingle, Tdouble),
        Bdims2 in ((), (4,)),
        Tt in Tset,
        TA in Tset,
        TB in Tset,
        tscale in Tt.((0.1, 1, 10)),
        shift in (true, false)

        t = tscale * randn(Tt)
        if MT <: Bidiagonal
            A = MT(randn(TA, n, n), :U)
        else
            A = MT(randn(TA, n, n))
        end
        B = randn(TB, n, Bdims2...)
        T = Base.promote_eltype(t, A, B)
        rT = real(T)
        @inferred expv(t, A, B; shift=shift)
        @test expv(t, A, B; shift=shift) ≈ exp(Matrix(t * A)) * B
        @test eltype(expv(t, A, B; shift=shift)) === T
    end

    @testset "expv(t::$Tt, A::SparseMatrixCSC{$TA}, B::Array{$TB,$(length(Bdims2)+1)}), tscale=$tscale, shift=$shift" for
        Tset in (Tsingle, Tdouble),
        Bdims2 in ((), (4,)),
        Tt in Tset,
        TA in Tset,
        TB in Tset,
        tscale in Tt.((0.1, 1, 10)),
        shift in (true, false)

        t = tscale * randn(Tt)
        A = VERSION ≥ v"1.1" ? sprandn(TA, n, n, 1 / n) : TA.(sprandn(n, n, 1 / n))
        B = randn(TB, n, Bdims2...)
        T = Base.promote_eltype(t, A, B)
        rT = real(T)
        @inferred expv(t, A, B; shift=shift)
        @test expv(t, A, B; shift=shift) ≈ exp(Matrix(t * A)) * B
        @test eltype(expv(t, A, B; shift=shift)) === T
    end

    @testset "errors if tolerance too low" begin
        t = randn()
        A = randn(10, 10)
        B = randn(10, 10)

        @test_throws DomainError expv(t, A, B; tol=eps(Float64) / 100)
        @test_throws DomainError expv(big(t), big.(A), big.(B))
    end
end


#=
    Following the description in
        _Exploring corner transfer matrices and corner tensors for the classical
        simulation of quantum lattice systems_
    by Orus

    this implementation follows the description in III.B, specifically
    the simplified one-directional 1D-method as in e.g. Fig 12, 13
    DOI: 10.1103/PhysRevB.85.205117
=#
struct transconjCTMIterable{T,
        TA<:AbstractTensor{T,4},
        TC<:AbstractTensor{T,2},
        TT<:AbstractTensor{T,3}} <: AbstractCTMIterable
    A::TA
    χ::Int
    Cinit::Union{TC,Nothing}
    Tsinit::Union{Tuple{TT,TT},Nothing}
end

function transconjctmiterable(A::DTensor{T,4}, χ::Int,
        Cinit = nothing,
        Tsinit = nothing) where T
    transconjCTMIterable{T,DTensor{T,4},DTensor{T,2},DTensor{T,3}}(A, χ, Cinit, Tsinit)
end

function transconjctmiterable(A::DASTensor{T,4,SYM,CHS,SS,CH}, χ::Int,
        Cinit = nothing,
        Tsinit = nothing) where {T,N,SYM,CHS,SS,CH}
    TC = DASTensor{T,2,SYM,CHS,SS,CH}
    TT = DASTensor{T,3,SYM,CHS,SS,CH}
    TA = DASTensor{T,4,SYM,CHS,SS,CH}
    transconjCTMIterable{T,TA,TC,TT}(A, χ, Cinit, Tsinit)
end

struct transconjCTMState{S,
        TA <: AbstractTensor,
        TC <: AbstractTensor,
        TT <: AbstractTensor} <: AbstractCTMState
    C::TC
    Ts::Tuple{TT,TT}
    oldsvdvals::Vector{S}
    diffs::Vector{S}
    n_it::Ref{Int}
end

function iterate(iter::transconjCTMIterable{S,TA,TC,TT}) where {S,TA,TC,TT}
    @unpack A, χ, Cinit, Tsinit = iter
    C  = isnothing(Cinit)  ? initializeC(A,χ) : Cinit
    Ts = isnothing(Tsinit) ? (initializeT(A,χ),initializeT(A,χ)) : Tsinit

    l = ifelse(C isa DASTensor, 2χ, χ)
    oldsvdvals = zeros(real(S),l)
    state = transconjCTMState{real(S),TA,TC,TT}(C, Ts, oldsvdvals, [], Ref(0))
    return state, state
end

function iterate(iter::transconjCTMIterable, state::transconjCTMState)
    @unpack A, χ = iter
    @unpack C, Ts, oldsvdvals, diffs, n_it = state
    T1, T2 = Ts

    xmove!((C,T2),(T1,A), χ)
    ymove!((C,T1),(T2,A), χ)

    vals = diag(tensorsvd(C)[2])
    maxval = maximum(vals)
    apply!(C, x -> x .= x ./ maxval)
    apply!(T1, x -> x .= x ./ norm(T1))
    apply!(T2, x -> x .= x ./ norm(T2))
    normalize!(vals,1)

    #compare
    push!(diffs, sum(abs, oldsvdvals - vals))
    oldsvdvals[:] = vals
    n_it[] += 1

    return state, state
end

function xmove!((C,T2),(T1,A), χ)
    # grow
    @tensor begin
        Cp[1,2,3]      := C[1,-1]    * T1[-1,2,3]
        T2p[1,2,3,4,5] := T2[1,-1,5] * A[2,3,4,-1]
    end

    #renormalize
    U, = tensorsvd(Cp, ((1,2),(3,)), svdtrunc = svdtrunc_maxχ(χ))
    @tensor begin
        C[1,2]    = U'[-1,-2,1] * Cp[-1,-2,2]
        T2[1,2,3] = U'[-1,-2,1] * T2p[-1,-2,2,-3,-4] * U[-4,-3,3]
    end
end

function ymove!((C,T1),(T2,A), χ)
    # grow
    @tensor begin
        Cp[1,2,3]      := C[-1,3]    * T2[1,2,-1]
        T1p[1,2,3,4,5] := T1[1,-1,5] * A[3,4,-1,2]
    end

    #renormalize
    u,s,Vd = tensorsvd(Cp, ((1,),(2,3)), svdtrunc = svdtrunc_maxχ(χ))
    @tensor begin
        C[1,2]    = Cp[1,-1,-2] * Vd'[2,-1,-2]
        T1[1,2,3] = Vd[1,-1,-2] * T1p[-2,-1,2,-3,-4] * Vd'[3,-3,-4]
    end
end


using JuMP, EAGO, Gurobi

m = Model(with_optimizer(EAGO.Optimizer, relaxed_optimizer = Gurobi.Optimizer(OutputFlag=0)))

# ----- Variables ----- #
x_Idx = Any[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
@variable(m, x[x_Idx])
JuMP.set_lower_bound(x[5], 0.0)
JuMP.set_lower_bound(x[4], 0.0)
JuMP.set_lower_bound(x[2], 0.0)
JuMP.set_lower_bound(x[3], 0.0)
JuMP.set_lower_bound(x[1], 10.0)
JuMP.set_upper_bound(x[1], 2000.0)
JuMP.set_upper_bound(x[2], 16000.0)
JuMP.set_upper_bound(x[3], 120.0)
JuMP.set_upper_bound(x[4], 5000.0)
JuMP.set_upper_bound(x[5], 2000.0)
JuMP.set_lower_bound(x[6], 85.0)
JuMP.set_upper_bound(x[6], 93.0)
JuMP.set_lower_bound(x[7], 90.0)
JuMP.set_upper_bound(x[7], 95.0)
JuMP.set_lower_bound(x[8], 3.0)
JuMP.set_upper_bound(x[8], 12.0)
JuMP.set_lower_bound(x[9], 1.2)
JuMP.set_upper_bound(x[9], 4.0)
JuMP.set_lower_bound(x[10], 145.0)
JuMP.set_upper_bound(x[10], 162.0)

# ----- Constraints ----- #
@NLconstraint(m, e1, -x[1]*(1.12+0.13167*x[8]-0.00667* (x[8])^2)+x[4] == 0.0)
@constraint(m, e2, -x[1]+1.22*x[4]-x[5] == 0.0)
@NLconstraint(m, e3, -0.001*x[4]*x[9]*x[6]/(98-x[6])+x[3] == 0.0)
@NLconstraint(m, e4, -(1.098*x[8]-0.038* (x[8])^2)-0.325*x[6]+x[7] == 57.425)
@NLconstraint(m, e5, -(x[2]+x[5])/x[1]+x[8] == 0.0)
@constraint(m, e6, x[9]+0.222*x[10] == 35.82)
@constraint(m, e7, -3*x[7]+x[10] == -133.0)

# ----- Objective ----- #
@NLobjective(m, Max, 0.063*x[4]*x[7] - 5.04*x[1] - 0.035*x[2] - 10*x[3] - 3.36*x[5]) # CAN BE QUADRATIC

JuMP.optimize!(m)

run_time = backend(m).optimizer.model.optimizer._run_time
#objective_value =
#objective_bound =
#println("objective between: ")
println("run time: $run_time")


using Distributions
using SpecialFunctions
using JLD2
using FileIO
using Distributed

function minalpha(β)
    α =   (exp(1)* sqrt( π))/β
    return α

end

function takememory(β,v,α)
    prob =Probability_of_memory_usage(β,v,α)
    wahrscheinlikeit = rand()
    result=false
    if wahrscheinlikeit < prob
        result=true
    end
    result
end

function Probability_of_memory_usage(β,v,α)
    prob = 1- exp(-β *energymemory(v,α))
    prob
end

function energymemory(v,α)
    energy =α*exp(-(α*v)^2)*sqrt(1/ π)
    energy
end

function freeenergy(β,v,α)
    energy = log( exp(β *energymemory(v,α)) -1)/β
    energy
end

function dissipation(β,v,α)
    diss = α*α/β * (v)^2
    diss
end

function newv_drift(δ,n)
    vnew=sqrt(n)δ+randn()*δ*0.05*sqrt(n)
    vnew
end

function onestep(β,δ,oldv,α,n)

    if (takememory(β,oldv,α))
        free=freeenergy(β,oldv,α)
        diss=dissipation(β,oldv,α)
        unet=free -diss
        nextv=newv_drift(δ,n)
        n+=1
        naiveuse=0
    else
        unet = 0
        free=0
        diss=0
        nextv=newv_drift(δ,1)
        naiveuse=1
        n=2
    end

    return (unet,nextv,naiveuse,free,diss,n)
end

function livesycle(β,δ,α,meanencounters,reps)
    meanunet=0 #initializing for all the rounds for fist cost
    meanfree=0
    meandiss=0
    naiveuses = 0
    for round in 1:reps
        localnetu= 0
        netu = 0
        localfree= 0
        free = 0
        localdiss= 0
        diss = 0

        v=newv_drift(δ,1) #start the cirus at the memory wich is allways at 0
        naiveuses += 1/reps
        n=2
        for encounter in 2:meanencounters  #start at 1 because initialization is allways at 0 in the first round without any prior memory
             naiveuse =0    # Used naive for the first virus
            (netu,v,naiveuse,free,diss,n)=onestep(β,δ,v,α,n)
            localnetu+=netu
            localfree+=free
            localdiss+=diss
            naiveuses += naiveuse/reps
        end

        meanunet+= localnetu/reps
        meanfree+= localfree/reps
        meandiss+= localdiss/reps
    end

    return (meanunet,naiveuses,meanfree,meandiss)
end

function findoptalphaoflive_drift_real(β,δ,meanencounters,reps,αsteps,αmax)

    #if (minalpha(β) >αmax)
    #    println("amax is to small for the chosen beta")
    #    return 0,0,0,0,0,0
    #end
    start= 0.01
    α = range(start,stop=αmax,length=αsteps)

    data=livesycle.(β,δ,α,meanencounters,reps)
    usesmaive=zeros(αsteps)
    memuses=zeros(αsteps)
    utility=zeros(αsteps)
    free=zeros(αsteps)
    diss=zeros(αsteps)
    for ii in 1:αsteps
        utility[ii] = data[ii][1]
        usesmaive[ii] = data[ii][2]
        free[ii] = data[ii][3]
        diss[ii] = data[ii][4]
    end
    memuses= meanencounters./usesmaive

    amax= argmax(utility)
    optalpha = α[amax]
    Nusesatmax=memuses[amax]
    argmaxn = argmax(memuses)
    maxNpossible=memuses[argmaxn]
    utiliymax = utility[amax]
    freemax=free[amax]
    dissmax=diss[amax]

    return Nusesatmax,utiliymax,optalpha,maxNpossible,freemax,dissmax
end

function find_optimal_alpha_for_range_of_delta(β,meanencounters,reps,αsteps,αmax,deltasteps)

    δ = range(0.0001,stop=1.5/αmax,length= deltasteps)
    #δ = range(0.325,stop=.4,length= deltasteps)./sqrt(αmax)
    memstepsatmax =zeros(deltasteps)
    utilityopt =zeros(deltasteps)
    alphaopt =zeros(deltasteps)
    posnmax =zeros(deltasteps)
    freemax =zeros(deltasteps)
    dissmax =zeros(deltasteps)
    for i =1:deltasteps
        memstepsatmax[i],utilityopt[i],alphaopt[i],posnmax[i],freemax[i],dissmax[i]=findoptalphaoflive_drift_real(β,δ[i],meanencounters,reps,αsteps,αmax)
    end
    return memstepsatmax,utilityopt,alphaopt,posnmax,δ,freemax,dissmax

end

function find_optimal_alpha_for_range_delta_and_beta(meanencounters,reps,αsteps,αmax,deltasteps,betasteps)

    β = range(0.1,length=betasteps,stop=10.0)

    memstepsatmax= zeros(betasteps,deltasteps)
    utilityopt= zeros(betasteps,deltasteps)
    alphaopt= zeros(betasteps,deltasteps)
    posnmax= zeros(betasteps,deltasteps)
    freemax =zeros(betasteps,deltasteps)
    dissmax =zeros(betasteps,deltasteps)
    δ=zeros(deltasteps)
    Threads.@threads for ii in 1:betasteps      # for parallel execution
    # for ii in 1:betasteps                     # non parallel
    memstepsatmax[ii,:],utilityopt[ii,:],alphaopt[ii,:],posnmax[ii,:],δ,freemax[ii,:],dissmax[ii,:] = find_optimal_alpha_for_range_of_delta(β[ii],meanencounters,reps,αsteps,αmax,deltasteps)
    end
    introduction="This gives the data meory_used_at_the_max  utility_atmax  otimal_alpha posible_maximum_memoryuses delta and beta and free_at_the_max and diss_at_the_max the data is Gamma first and delta second index"
 return memstepsatmax,utilityopt,alphaopt,posnmax,δ,β,freemax,dissmax,introduction

end


"""
    Ke_VL_model.jl

Written by Carl Whitfield, University of Manchester, 2021

Only to be included by definitions.jl
"""

#===================== VL trajectory params  =====================#

#====> Ke et al. 2021 (fitted) <====#
#log means of (1) peak viral load (copies/ml), (2) peak time (days from infection)
#(3) 1/r -- inverse growth rate (days) (4) 1/d -- inverse decay rate (days)
const Kmvμ = [17.546426, 1.392914, -1.197782, -0.671588]
#log covariance matrix
const KmvΣ = [0.909403 0.043808 -0.010543 0.080628;
              0.043808 0.032527 0.028866 0.004503;
             -0.010543 0.028866 0.028640 0.000383;
              0.080628 0.004503 0.000383 0.026590]
const Log_VLparams_dist = MvNormal(Kmvμ, KmvΣ)  #Distribution used to generate params

#=================================================================#

#=================== Infectivity vs. VL params ===================#

#====> Ke et al. 2021 (fitted) <====#
#log means of parameters J and h (Inf = J VL^h / (VL^h + Km^h))
const Kmvinfμ = [1.782, -0.108]
#log covars
const KmvinfΣ = [0.3387 0.0265; 0.0265 0.1270]
const Jscale = 4.188/exp(Kmvinfμ[1] + KmvinfΣ[1,1]/2)  #normalisation of mean peak inf (a.u.)
const Kinf_Km = 4.0e6
const Log_infparams_dist = MvNormal(Kmvinfμ, KmvinfΣ)  #Distribution used to generate params

# KmvΣ = reshape(KmvΣ,(4,4))
function generate_ke_VL_params()
    Ps = exp.(rand(MvNormal(Kmvμ, KmvΣ)))
    return log10(Ps[1]), Ps[2], 1/Ps[3], 1/Ps[4]
end

function generate_ke_inf_params()
    return exp.(rand(MvNormal(Kmvinfμ, KmvinfΣ)))
end

function infectivity_Ke(PVL::Float64, r::Float64, d::Float64, tp::Float64, J::Float64, h::Float64, T::Int64)
    #if viral load is piecewise linear, and infectivity is hill function, what is cumulative infectivity?
    m_up = r/log(10)
    m_down = d/log(10)

    cum_inf = zeros(T+1)
    t_inf = collect(-1:(T-1)) .+ 0.5
    t_inf[1] = 0
    t_inf[T+1] = T - 1.0

    V0 = (PVL - m_up*tp)
    K_m = 4.0e6^h
    a = h*r
    PInf = Jscale*J
    c1 = K_m * 10^(-V0*h)
    cond1 = (t_inf .> 0) .* (t_inf .<=  tp)
    cum_inf[cond1] = PInf .* (log.(exp.(a .* t_inf[cond1]) .+ c1)
                              .- log(c1 + 1.0)) ./ a
    ci_peak = PInf .* (log.(exp.(a * (tp)) .+ c1) .- log(c1 + 1)) ./ a
    cond2 = (t_inf .>  tp)
    b = -h*d
    c2 = K_m * 10^(-PVL*h)
    cum_inf[cond2] = ci_peak .+ PInf .* (log.(exp.(b .* (t_inf[cond2] .- tp)) .+ c2)
                              .- log(c2 + 1.0)) ./ b

    inf = (cum_inf[2:(T+1)] .- cum_inf[1:T]) ./ (t_inf[2:(T+1)] .- t_inf[1:T])

    return inf
end


using TrajectoryOptimization
using Altro
using RobotDynamics
using Attitude
using SatelliteDynamics
using StaticArrays
using LinearAlgebra
using MATLAB
const TO = TrajectoryOptimization
const RD = RobotDynamics

include(joinpath(dirname(dirname(@__FILE__)),"ks_functions.jl"))

struct KSopt <: TO.AbstractModel
    μ::Float64
    tscale::Float64
    dscale::Float64
    uscale::Float64
end
function RobotDynamics.dynamics(model::KSopt, x, u_dot)

    # unpack state
    p = SVector(x[1],x[2],x[3],x[4])
    p_prime = SVector(x[5],x[6],x[7],x[8])
    h = x[9]
    #r = x[10]
    #z = x[11]
    u = SVector(x[12],x[13],x[14])

    P_j2 = SVector(u[1],u[2],u[3])/model.uscale
    # Levi-Civita transformation matrix
    L = L_fx(p)

    # append 0 to end of P_j2 acceleration vector
    P_j2 = SVector{4}([P_j2;0]*model.tscale^2/model.dscale)

    p_dp = -(h/2)*p + ((dot(p,p))/(2))*L'*P_j2

    return SVector(p_prime[1],p_prime[2],p_prime[3],p_prime[4],
                   p_dp[1],p_dp[2],p_dp[3],p_dp[4],
                   -2*dot(p_prime,L'*P_j2),
                   2*p'*p_prime,
                   2*(p_prime[1]*p[3] + p[1]*p_prime[3] + p_prime[2]*p[4] + p[2]*p_prime[4]),
                   u_dot[1],u_dot[2],u_dot[3])

end

# scaling
dscale = 1e7 # m
tscale = 20000 # s
uscale = 10000.0
μ = 3.986004418e14 *(tscale^2)/(dscale^3)
model = KSopt(μ,tscale, dscale, uscale)
Base.size(::KSopt) = 14,3
n,m = size(model)

# initial conditions
r_eci = @SVector [6.578137000000001e6,0,0]
v_eci = @SVector [0,9111.210131355681,4642.393437617197]
R_0 = r_eci/dscale
V_0 = v_eci*tscale/dscale
p_0 = u_from_x(R_0)                          # u
p_prime_0 = uprime_from_xdot(V_0,p_0)     # du/ds
h_0 = μ/norm(R_0) - .5*norm(V_0)^2         # m^2/s^2
t_0 = 0                                    # seconds
x0 = SVector{n}([p_0; p_prime_0; h_0; norm(R_0); R_0[3];zeros(3)])

# time stuff
N = 1501
dt = 3e-2
tf = dt*(N-1)

# cost function
scaling_num = 10
r_Desired = 42e6/dscale#4.2#6.77814
R_scale = .01
Q = Diagonal(SVector{n}([0*ones(9);scaling_num;100;R_scale*ones(3)]))
R = 0.001*Diagonal(SVector(R_scale,R_scale,R_scale))
q = SVector{n}([zeros(9);-r_Desired*scaling_num;0;zeros(3)])
costfun = QuadraticCost(Q,R,q=q)
costfun_term = QuadraticCost(Q,R,q=q)
obj = Objective(costfun, costfun_term, N)

# constraints
cons = ConstraintList(n,m,N)
u_bnd = 10.0
maxThrust = NormConstraint(n, m, u_bnd, TO.SecondOrderCone(), 12:14)
TO.add_constraint!(cons, maxThrust, 1:N)

xf = randn(n)
xf[11] = 4.2
xf = SVector{n}(xf)
prob = Problem(model, obj,xf, tf, x0=x0, constraints=cons,integration=RD.RK4)
opts = SolverOptions(verbose=2,
                     show_summary = true,
                     iterations = 20000,
                     projected_newton = false,
                     constraint_tolerance = 1e-4, penalty_scaling = 2.0, penalty_initial = 1e-2)
solver = ALTROSolver(prob,opts)

solve!(solver)

# get time stuff
function get_time_transfer(solver,dt,tscale)
    """get the time in days for a run. This allows us to remove t from state"""
    X = states(solver)
    t = 0.0
    t_hist = zeros(length(X))
    for i = 1:(length(X)-1)

        p0 = X[i][1:4]
        R0 = dot(p0,p0)

        p1 = X[i+1][1:4]
        R1 = dot(p1,p1)

        t += (R0 + R1)*dt/2
        t_hist[i+1] = copy(t)
    end
    t_days = t*tscale/(24*3600)
    t_hist = t_hist*tscale/(24*3600)
    return t_days, t_hist
end

t_days,t_hist = get_time_transfer(solver,dt,tscale)

X = states(solver)
U = controls(solver)

r_eci_hist = dscale*mat_from_vec([x_from_u(X[i][1:4]) for i = 1:length(X)])
v_eci_hist = (dscale/tscale)*mat_from_vec([xdot_from_u(X[i][1:4],X[i][5:8]) for i = 1:length(X)])

mat"
figure
hold on
plot3($r_eci_hist(1,:),$r_eci_hist(2,:),$r_eci_hist(3,:))
%axis equal
hold off
"

Xm = mat_from_vec(X)
U_real = Xm[12:14,:]

mat"
figure
hold on
title('Thrust Vector')
plot($t_hist,$U_real')
ylabel('Thrust Acceleration (0.0001m/s^2)')
xlabel('Time (days)')
hold off
xlim([0,$t_hist(end)])
legend('u_x','u_y','u_z')
%saveas(gcf,'long_traj.png')
"

U_real_vec = [X[i][12:14] for i = 1:length(X)]
Unorm = [norm(U_real_vec[i]) for i = 1:length(U)]

mat"
figure
hold on
plot($Unorm)
hold off
"

println("Time in days:  $t_days")# using TrajectoryOptimization

# now let's get orbital elements
const SD = SatelliteDynamics

oe_hist = [SD.sCARTtoOSC([r_eci_hist[:,i];v_eci_hist[:,i]]) for i = 1:size(r_eci_hist,2)]

a_hist = [oe_hist[i][1] for i = 1:length(oe_hist)]
e_hist = [oe_hist[i][2] for i = 1:length(oe_hist)]
i_hist = [rad2deg(oe_hist[i][3]) for i = 1:length(oe_hist)]

mat"
figure
hold on
subplot(3,1,1)
hold on
plot($t_hist,$a_hist)
subplot(3,1,2)
hold on
plot($t_hist,$e_hist)
subplot(3,1,3)
hold on
plot($t_hist,$i_hist)
hold off
"

# using JLD2
# X = states(solver)
# U = controls(solver)
# @save "rate_control_transfers/30_v3_day_transfer.jld2" X U t_hist


# Non-Linear Regression Tests

f(X_, beta_) = exp.(beta_[1]*X_[:,1] + beta_[2]*log.(X_[:,2]))
Ynlin = f(X, beta) + 0.1*randn(TEST_SIZE_1, 1)

# Optimize test
OptimizeRegression = Optimize(f, [1.0, 1.0])
res = nlinfit(X, Ynlin, OptimizeRegression)

true


# Exponential Arm Model

type Exponential <: BanditArmBase
    armDist::Distributions.Exponential

    function Exponential( λ::Real ) # λ is the mean parameter
        new( Distributions.Exponential(λ) )
    end
end

function pull!( arm::Exponential )
    return Distributions.rand( arm.armDist )
end

function tick!( arm::Exponential )
    # Do nothing
    nothing
end

function reset!( arm::Exponential )
    # Do nothing
    nothing
end


using SMLMData
using StatsBase

"""
    smld_combined = combinelocalizations(smld::SMLMData.SMLD2D)

Combine clustered localizations in `smld` into higher precision localizations.

# Description
This function combines localizations in `smld` that share the same value of
smld.connectID.  Localizations are combined assuming they arose from
independent measurements of the same position with Gaussian errors.
"""
function combinelocalizations(smld::SMLMData.SMLD2D)
    # Isolate and sort some arrays from `smld`.
    connectID = smld.connectID
    sortindices = sortperm(connectID)
    connectID = connectID[sortindices]
    x = smld.x[sortindices]
    y = smld.y[sortindices]
    σ_x = smld.σ_x[sortindices]
    σ_y = smld.σ_y[sortindices]
    photons = smld.photons[sortindices]
    σ_photons = smld.σ_photons[sortindices]
    bg = smld.bg[sortindices]
    σ_bg = smld.σ_bg[sortindices]
    framenum = smld.framenum[sortindices]
    datasetnum = smld.datasetnum[sortindices]

    # Loop over clusters and combine their localization data as appropriate.
    nperID = counts(connectID)
    nperID = nperID[nperID .!= 0]
    ncumulative = [0; cumsum(nperID)]
    nclusters = length(nperID)
    smld_combined = deepcopy(smld)
    for nn in 1:nclusters
        indices = (1:nperID[nn]) .+ ncumulative[nn]
        smld_combined.x[nn] = StatsBase.mean(x[indices], weights(1.0 / σ_x[indices] .^ 2))
        smld_combined.y[nn] = StatsBase.mean(y[indices], weights(1.0 / σ_y[indices] .^ 2))
        smld_combined.σ_x[nn] = sqrt(1.0 / sum(1.0 ./ σ_x[indices] .^ 2))
        smld_combined.σ_y[nn] = sqrt(1.0 / sum(1.0 ./ σ_y[indices] .^ 2))
        smld_combined.photons[nn] = sum(photons[indices])
        smld_combined.σ_photons[nn] = sqrt(sum(σ_photons[indices] .^ 2))
        smld_combined.bg[nn] = sum(bg[indices])
        smld_combined.σ_bg[nn] = sqrt(sum(σ_bg[indices] .^ 2))
        smld_combined.connectID[nn] = connectID[indices[1]]
        smld_combined.framenum[nn] = framenum[indices[1]]
        smld_combined.datasetnum[nn] = datasetnum[indices[1]]
    end
    smld_combined = SMLMData.isolatesmld(smld_combined, 1:nclusters)

    return smld_combined
end


#!/usr/bin/julia

function fib(n)
    f = BigInt[1 1; 1 0]
    for i = 1:n
        fn = f ^ i
        println(i, ": ", fn[2, 1])
    end
end

fib(parse(Int, ARGS[1]))


using LsqFit

include("utils.jl")
function dfa(z; scales=3:0.5:8, offset=-1)
    # cummulative sum of (z - <z>)
    zcum = cumsum(z .- mean(z))

    # define the rulers for each x-window
    scales = collect(scales)
    rulers = convert.(Int64,round.(2 .^ scales))

    # each rulers will have rms averaged for each window in the data
    fluct = zeros(length(rulers))

    for (k, ℓ) in enumerate(rulers)
        # x data is the same for all windows in this scale
        xdata = collect(1:ℓ)

        #each element contail the data in one window
        yvec = strided(zcum,ℓ,offset)

        #for each window one rms
        rms = zeros(length(yvec))

        for (i,ydata) in enumerate(yvec)
            fit = curve_fit((x, p) -> p[1] .+ p[2].*x, xdata, ydata, [1.0, 1.0])
            σ = stderror(fit)
            rms[i] = sqrt(mean(fit.resid.^2))
            # yfit = model(xdata, fit.param)
            # rms2 = sqrt(mean((ydata .- yfit).^2))
        end
        fluct[k] = mean(rms)
    end

    # now fit (fluct vs. scale)
    x = log10.(rulers)
    y = log10.(fluct)
    fit = curve_fit((x, p) -> p[1] .+ p[2].*x ,x, y, [.0, .0])
    h = fit.param[2]-1.0
    return  h,x,y
end

function dfa2d(z; scales=3:0.5:8, offset=-1)
    L = size(z)[1]
    hurst = 0
    n = 0
    for i=10:10:L-10
        h, x, y = dfa(z[:,i]; scales=scales, offset=offset)
        hurst += h
        n+=1
    end
    for i=10:10:L-10
        h, x, y = dfa(z[i,:]; scales=scales, offset=offset)
        hurst += h
        n+=1
    end
    hurst = hurst/n
end


using MAT

useproc("GPU") # GPU is much faster than CPU
#useproc("CPU")

Ntrain=60000
Ntest=10000
BatchSize=100
TrainingIts=5000 # number of training iterations
include("loadmnist.jl")
images,label=loadmnist()
label=convert(Array{Int,2},label)
r=randperm(size(images,2))
data=images[:,r]; label=label[r]
class=zeros(10,70000)+0.00001
for i=1:70000
    class[label[i]+1,i]=0.9999
end
traindata=data[:,1:Ntrain]
testdata=data[:,Ntrain+1:Ntrain+Ntest]
trainclass=class[:,1:Ntrain]
testclass=class[:,Ntrain+1:Ntrain+Ntest]

# Construct the DAG function:
H=[784 3000 200 100 50 10] # number of units in each layer

L=length(H) # number of hidden layers
# node indices:
bias=Array(ADnode,L) # weight index
w=Array(ADnode,L) # weight index
h=Array(ADnode,L) # hidden layer index (note that I call the input layer h[1])
#hh=Array(ADnode,L) # hidden layer index (note that I call the input layer h[1])
dropout=Array(ADnode,L) # hidden layer index (note that I call the input layer h[1])
StartCode()
class=ADnode() # node has no parents
x=ADnode()
dropout[1]=ADnode()
h[1]=diagm(dropout[1])*x
for i=2:L-1
    dropout[i]=ADnode()
    w[i]=ADvariable()
    bias[i]=ADvariable()
    h[i]=diagm(dropout[i])*rectlinAXplusBias(w[i],h[i-1],bias[i])
end
w[L]=ADvariable()
h[L]= w[L]*h[L-1]
loss=KLsoftmax(class,h[L]) # the loss we are minimising must be the final node in the graph.
net=EndCode()

# instantiate parameter nodes and inputs:
net.value[x]=traindata[:,1:BatchSize]
net.value[class]=trainclass[:,1:BatchSize]
for i=2:L
    net.value[w[i]]=sign(randn(H[i],H[i-1]))/sqrt(H[i-1])
end
net.value[dropout[1]]=ones(H[1])
for i=2:L-1
    net.value[bias[i]]=zeros(H[i])
    net.value[dropout[i]]=ones(H[i],1)
end

net=compile(net) # compile the DAG and preallocate memory
#gradcheck(net)

@gpu CUDArt.init([0])
@gpu net=convert(net,"GPU")

tmp=net.value[1]

ADforward!(net)

dropoutprob=zeros(L)
dropoutprob[1]=0
#for i=2:L-1
#    dropoutprob[i]=0.2
#end
dropoutprob[2]=0.5
println("Training: using $(net.gpu==true? "GPU" : "CPU") ")
error=Array(Float64,0)
ParsToUpdate=Parameters(net)
meangrad=zeros(TrainingIts)
meanv=zeros(TrainingIts)
velo=NesterovInit(net)
minibatchstart=1 # starting datapoint for the minibatch
for iter=1:TrainingIts
    LearningRate=0.1/(1+iter/1500)
    minibatchstart,minibatch=GetBatch(minibatchstart,BatchSize,Ntrain)
    net.value[x]=cArray(PROC,traindata[:,minibatch]) # select batch of data
    net.value[class]=cArray(PROC,trainclass[:,minibatch]) # select batch of data
    for i=2:L-1
        net.value[dropout[i]]=cArray(PROC,1.0*(rand(H[i],1).>dropoutprob[i]))
    end
    ADforward!(net)
    ADbackward!(net)
    push!(error,extract(net.value[net.FunctionNode])[1])
    printover("iteration $iter: training loss = $(error[iter])")
    for par in ParsToUpdate
        #GradientDescentUpdate!(net.value[par],net.gradient[par],LearningRate)
        NesterovGradientDescentUpdate!(net.value[par],net.gradient[par],velo[par],LearningRate,iter)
    end
end
#For prediction scaling the weights is equivalent to using a fixed dropout value:
for i=1:L-1
    net.value[dropout[i]]=cArray((1-dropoutprob[i])*ones(H[i],1))
end
ADforward!(net)

net=convert(net,"CPU") # convert to CPU for easy results analysis

net.value[h[L]]
classpredtrain=argcolmax(softmax(net.value[h[L]]))
classtrain=argcolmax(net.value[class])
println("\nlast batch train accuracy = $(mean(classpredtrain.==classtrain))")

println("Testing: using $(net.gpu==true? "GPU" : "CPU") ")
# evaluate the test predictions:
net.value[x]=testdata
ForwardPassList!(net,ExcludeNodes=[class])
# ADforward! by default computes values inplace and does no additional memory allocation on the graph. If we change the input value dimensions (which happens here since there are a different number of test points to train points), we need to reallocatte memory based on the shapes of the new values on the graph. This only needs to be done once.
ADforward!(net,AllocateMemory=true) # reallocate since shape of net.value[x] changed

classpredtest=argcolmax(softmax(net.value[h[L]]))
classtest=argcolmax(testclass)
println("test accuracy = $(mean(classpredtest.==classtest))")

#end


@testset "AABB" begin
    @testset "Construction" begin
        for T in [Float16, Float32, Float64]
            e = one(T)
            z = zero(T)
            a = AABB(SVector(-2e,-2e,-2e),SVector(-e,-e,-e))
            b = AABB(SVector(-2e,-2e,-2e),SVector(-3e,-3e,-3e))
            c = AABB(a, b)
            @test c.pmin == SVector(-3e,-3e,-3e)
            @test c.pmax == SVector(-e,-e,-e)
        end
    end
    @testset "Intersections" begin
        for T in [Float16, Float32, Float64]
            e = one(T)
            for ϵ in [eps(T), e/10]
                z = zero(T)

                finite_aabb = AABB(SVector(-e,-e,-e), SVector(e,e,e))

                r₁ = Ray(SVector(z, z, 2e), SVector(z, z, -e))
                @test intersect(r₁, finite_aabb)

                r₂ = Ray(SVector(z, z, 2e), SVector(z, z, e))
                @test !intersect(r₂, finite_aabb)

                r₃ = Ray(SVector(z, z, 2e), SVector(e,e,-(e+ϵ)))
                @test intersect(r₃, finite_aabb)

                r₄ = Ray(SVector(z, z, 2e), SVector(e,e,-(e-ϵ)))
                @test !intersect(r₄, finite_aabb)

                r₅ = Ray(SVector(z, e-ϵ, 2e), SVector(z,z,-e))
                @test intersect(r₅, finite_aabb)

                r₆ = Ray(SVector(z, e+ϵ, 2e), SVector(z,z,-e))
                @test !intersect(r₆, finite_aabb)
            end
        end
    end
end


using VortexDistributions

# simple test field
Nx = 400; Ny = 400
Lx = 200; Ly = Lx
x = LinRange(-Lx/2,Lx/2, Nx+1)[1:end-1]; y = LinRange(-Ly/2,Ly/2, Ny+1)[1:end-1];
psi0 = one.(x*y') |> complex

psi = Torus(copy(psi0), x, y);

xp = 0.0
yp = 0.0
vp = PointVortex(xp,yp,1)
sp = ScalarVortex(vp)

vortex!(psi, sp)

vfound = findvortices(psi)

vortm = keep_vortices(vfound)

@test length(vfound) == length(vortm)


 @testset "Direct Ellipse Fit" begin
    @info "Test: Direct Ellipse Fit"

    @testset "Numerical" begin
        truth = SVector(4.053658735425687e-6,
                        -1.3487952962436518e-7,
                        1.8534714216409155e-5,
                        -0.001965777847001004,
                        -0.009242570343232327,
                         0.9999553541288334)

        𝛉 = fit_ellipse(data, DirectEllipseFit())
        @test all(𝛉 .≈ truth)
    end

end


include("euler/euler.jl")
using .BigIntegers: BigInteger
using .Numbers: digits_sum
using BenchmarkTools
BenchmarkTools.DEFAULT_PARAMETERS.samples = 100

function compute(n::Int)::Int
    current, previous = BigInteger(1), BigInteger(1)
    for i ∈ 2:n + 1
        current, previous = (i % 3 == 1 ? 2i ÷ 3 : 1) * current + previous, current
    end
    return digits_sum(current)
end

compute(10)

compute(100)

@benchmark compute(100)


### CONCRETE TYPE: DIRECT PROX EVALUATION
# prox! is computed using a Cholesky factorization of Q + I/gamma.
# The factorization is cached and recomputed whenever gamma changes

using LinearAlgebra
using SparseArrays
using SuiteSparse

mutable struct QuadraticDirect{R <: Real, M <: AbstractMatrix{R}, V <: AbstractVector{R}, F <: Factorization} <: Quadratic
    Q::M
    q::V
    gamma::R
    temp::V
    fact::F
    function QuadraticDirect{R, M, V, F}(Q::M, q::V) where {R <: Real, M <: AbstractMatrix{R}, V <: AbstractVector{R}, F <: Factorization}
        if size(Q, 1) != size(Q, 2) || length(q) != size(Q, 2)
            error("Q must be squared and q must be compatible with Q")
        end
        new(Q, q, -1, similar(q))
    end
end

function QuadraticDirect(Q::M, q::V) where {R <: Real, I <: Integer, M <: SparseMatrixCSC{R, I}, V <: AbstractVector{R}}
    QuadraticDirect{R, M, V, SuiteSparse.CHOLMOD.Factor{R}}(Q, q)
end

function QuadraticDirect(Q::M, q::V) where {R <: Real, M <: DenseMatrix{R}, V <: AbstractVector{R}}
    QuadraticDirect{R, M, V, Cholesky{R, M}}(Q, q)
end

function (f::QuadraticDirect{R, M, V, F})(x::AbstractArray{R}) where {R, M, V, F}
    mul!(f.temp, f.Q, x)
    return 0.5*dot(x, f.temp) + dot(x, f.q)
end

function prox!(y::AbstractArray{R}, f::QuadraticDirect{R, M, V, F}, x::AbstractArray{R}, gamma::R=R(1)) where {R, M, V, F <: Cholesky}
    if gamma != f.gamma
        factor_step!(f, gamma)
    end
    y .= x./gamma
    y .-= f.q
    # Qy = U'Uy = b, therefore y = U\(U'\b)
    LAPACK.trtrs!('U', 'C', 'N', f.fact.factors, y)
    LAPACK.trtrs!('U', 'N', 'N', f.fact.factors, y)
    mul!(f.temp, f.Q, y)
    fy = 0.5*dot(y, f.temp) + dot(y, f.q)
    return fy
end

function prox!(y::AbstractArray{R}, f::QuadraticDirect{R, M, V, F}, x::AbstractArray{R}, gamma::R=R(1)) where {R, M, V, F <: SuiteSparse.CHOLMOD.Factor}
    if gamma != f.gamma
        factor_step!(f, gamma)
    end
    f.temp .= x./gamma
    f.temp .-= f.q
    y .= f.fact\f.temp
    mul!(f.temp, f.Q, y)
    fy = 0.5*dot(y, f.temp) + dot(y, f.q)
    return fy
end

function factor_step!(f::QuadraticDirect{R, M, V, F}, gamma::R) where {R, M <: DenseMatrix{R}, V, F}
    f.gamma = gamma
    f.fact = cholesky(f.Q + I/gamma)
end

function factor_step!(f::QuadraticDirect{R, M, V, F}, gamma::R) where {R, I, M <: SparseMatrixCSC{R, I}, V, F}
    f.gamma = gamma
    f.fact = ldlt(f.Q; shift = 1/gamma)
end

function gradient!(y::AbstractArray{R}, f::QuadraticDirect{R, M, V, F}, x::AbstractArray{R}) where {R, M, V, F}
    mul!(y, f.Q, x)
    y .+= f.q
    return 0.5*(dot(x, y) + dot(x, f.q))
end

function prox_naive(f::QuadraticDirect, x::AbstractArray{R}, gamma::R=one(R)) where R
    y = (gamma*f.Q + I)\(x - gamma*f.q)
    fy = 0.5*dot(y, f.Q*y) + dot(y, f.q)
    return y, fy
end


##########################################################################
# Channels
abstract Channel

immutable SinglecarrierChannel <: Channel
    H::Array{Matrix{Complex128},2}
    Ns::Vector{Int}
    Ms::Vector{Int}
    K::Int
    I::Int

    large_scale_fading_factor::Matrix{Float64}
end
SinglecarrierChannel(H, Ns, Ms, K, I) =
    SinglecarrierChannel(H, Ns, Ms, K, I, ones(K, I))

immutable MulticarrierChannel <: Channel
    H::Array{Array{Complex128,3},2}
    Ns::Vector{Int}
    Ms::Vector{Int}
    K::Int
    I::Int
    Lc::Int

    large_scale_fading_factor::Matrix{Float64}
end
MulticarrierChannel(H, Ns, Ms, K, I, Lc) =
    MulticarrierChannel(H, Ns, Ms, K, I, Lc, ones(K, I))

get_channel_gains_dB(channel::SinglecarrierChannel) =
    Float64[ 20*log10(vecnorm(channel.H[k,i])) for k = 1:channel.K, i = 1:channel.I ]

function get_average_channel_gains_dB(channel::MulticarrierChannel)
    average_channel_gains = zeros(Float64, channel.K, channel.I)

    for k = 1:channel.K
        for i = 1:channel.I
            for l = 1:channel.Lc
                average_channel_gains[k,i] += (1/channel.Lc)*vecnorm(channel.H[k,i,l])^2
            end
        end
    end

    return 10*log10(average_channel_gains)
end


@inline xorshift_rotl(x::UInt64, k::Int) = (x >>> (0x3f & -k)) | (x << (0x3f & k))
@inline xorshift_rotl(x::UInt32, k::Int) = (x >>> (0x1f & -k)) | (x << (0x1f & k))

"""
SplitMix64: only for initializing a random seed.
"""
@inline function splitmix64(x::UInt64)
    x += 0x9e3779b97f4a7c15
    x = (x ⊻ (x >> 30)) * 0xbf58476d1ce4e5b9
    x = (x ⊻ (x >> 27)) * 0x94d049bb133111eb
    x ⊻ (x >> 31)
end

# for xoshiro256
function init_seed(seed::UInt64)
    x2 = splitmix64(seed)
    x3 = splitmix64(x2)
    x4 = splitmix64(x3)
    (seed, x2, x3, x4)
end


module DiscreteChoiceModels

using OnlineStats, Tables, Printf, Statistics, LinearAlgebra, DataFrames, Dagger, ForwardDiff, FLoops, SplittablesBase
import FunctionWrappers: FunctionWrapper

include("LogitModel.jl")
include("util.jl")
include("utilitymacro.jl")
include("mnl.jl")
include("compute_utility.jl")

export multinomial_logit, @utility, @β, @dummy_code

end


#=
Copyright (c) 2018-2022 Chris Coey, Lea Kapelevich, and contributors

This Julia package Hypatia.jl is released under the MIT license; see LICENSE
file in the root directory or at https://github.com/chriscoey/Hypatia.jl

portfolio rebalancing problem
=#

struct PortfolioJuMP{T <: Real} <: ExampleInstanceJuMP{T}
    num_stocks::Int
    epinormeucl_constr::Bool # add an L2 ball constraint
    epinorminf_constrs::Bool # add L1 and Linfty ball constraints
end

function build(inst::PortfolioJuMP{T}) where {T <: Float64}
    @assert xor(inst.epinormeucl_constr, inst.epinorminf_constrs)
    num_stocks = inst.num_stocks

    returns = rand(num_stocks)
    sigma_half = randn(num_stocks, num_stocks)
    x = randn(num_stocks)
    gamma = sum(abs, sigma_half * x) / norm(x)
    A = randn(div(num_stocks, 2), num_stocks)

    model = JuMP.Model()
    JuMP.@variable(model, invest[1:num_stocks])
    JuMP.@objective(model, Max, dot(returns, invest))
    JuMP.@constraint(model, sum(invest) == 0)
    JuMP.@constraint(model, A * invest .== 0)

    if inst.epinormeucl_constr
        JuMP.@constraint(model, vcat(gamma / sqrt(num_stocks),
            sigma_half * invest) in JuMP.SecondOrderCone())
    end
    if inst.epinorminf_constrs
        JuMP.@constraint(model, vcat(1, invest) in
            MOI.NormInfinityCone(1 + num_stocks))
        JuMP.@constraint(model, vcat(gamma, sigma_half * invest) in
            MOI.NormOneCone(1 + num_stocks))
    end

    return model
end


using GLVisualize, GLAbstraction, Reactive, GeometryTypes, Colors, GLWindow

if !isdefined(:runtests)
    window = glscreen()
    addprocs(1)
end

description = """
Example showing off how to run computations in a different process and then
visualizing them.
"""

const workerid = workers()[]

@everywhere begin
    using GeometryTypes
    function spin(dir, α)
        x, y = dir
        Point(x*cos(α) - y*sin(α), x*sin(α) + y*cos(α))
    end
    function fractal_step!(
            a, b, depth, angles,
            result = Point2f0[a], levels = Float32[depth] # tmp too not allocate
        )
        depth == 0 && return result, levels, b
        N = length(angles)
        diff = (b - a)
        len = norm(diff) / (N - 1)
        nth_segment = normalize(diff)
        for n = 1:N
            b = a + spin(nth_segment * len, angles[n]) # rotate with current angle
            _, _, b = fractal_step!(a, b, depth-1, angles, result, levels) # recursion step
            if n < N
                push!(result, b)
                push!(levels, depth)
            end
            nth_segment = normalize(b - a)
            a = b
        end
        result, levels, b
    end
    function generate_fractal(angles, depth = 5)
        tmp = zeros(Point2f0, length(angles))
        angles = map(deg2rad, angles)
        result, levels, b = fractal_step!(Point2f0(0,0), Point2f0(300,0), round(Int, depth), angles)
        push!(result, b)
        push!(levels, depth)
        result, depth .- levels
    end
end

editarea, viewarea = x_partition(window.area, 30.0)
editscreen = Screen(
    window, area = editarea,
    color = RGBA{Float32}(0.1f0, 0.1f0, 0.1f0, 1f0),
    stroke = (1f0, RGBA{Float32}(0.13f0, 0.13f0, 0.13f0, 13f0))
)
viewscreen = Screen(window, area=viewarea, color = RGBA(0.1f0, 0.1f0, 0.1f0, 1f0))

controls = [
    :angles => Vec4f0(0.0, 80.0, -140.0, 80.0),
    :colormap => map(RGBA{Float32}, colormap("Blues", 10)),
    :thickness => 0.5f0,
    :iterations => 6,
]

menu, controls = GLVisualize.extract_edit_menu(controls, editscreen, true)
_view(menu, editscreen, camera = :fixed_pixel)

const channel2 = Channel{Tuple{Vector{Point2f0}, Vector{Float32}}}(1)
v0 = (Point2f0[0,0,0,0], Float32[0,0,0,0])
put!(channel2, v0)
line_level = foldp(v0, controls[:angles], controls[:iterations]) do v0, angles, iter
    if isready(channel2)
        points = take!(channel2)
        task = @async put!(
            channel2,
            remotecall_fetch(generate_fractal, workerid, angles, iter)
        )
        return points
    end
    v0
end

_view(visualize(
    map(first, line_level), :lines,
    intensity = map(last, line_level),
    thickness = controls[:thickness],
    color_map = controls[:colormap],
    color_norm = Vec2f0(0,10),
), viewscreen, camera=:orthographic_pixel)

if !isdefined(:runtests)
    renderloop(window)
end


# This file was generated, do not modify it.

using Pkg # hideall
Pkg.activate("_literate/EX-breastcancer/Project.toml")
Pkg.update()
macro OUTPUT()
    return isdefined(Main, :Franklin) ? Franklin.OUT_PATH[] : "/tmp/"
end;

using UrlDownload
using DataFrames
using PrettyPrinting
using PyPlot
using MLJ
ioff() # hide
MLJ.color_off(); # hide

RANDOM_SEED = 42;

url = "https://archive.ics.uci.edu/ml/machine-learning-databases/breast-cancer-wisconsin/wdbc.data";
feature_names = ["ID", "Class", "mean radius", "mean texture", "mean perimeter", "mean area", "mean smoothness", "mean compactness", "mean concavity", "mean concave points", "mean symmetry", "mean fractal dimension", "radius error", "texture error", "perimeter error", "area error", "smoothness error", "compactness error", "concavity error", "concave points error", "symmetry error", "fractal dimension error", "worst radius", "worst texture", "worst perimeter", "worst area", "worst smoothness", "worst compactness", "worst concavity", "worst concave points", "worst symmetry", "worst fractal dimension"]
data = urldownload(url, true, format = :CSV, header = feature_names);

figure(figsize=(8, 6))
hist(data.Class)
xlabel("Classes")
ylabel("Number of samples")
plt.savefig(joinpath(@OUTPUT, "Target_class.svg")); # hide

df = DataFrame(data)[:, 2:end];

pprint(first(df,10))

pprint(describe(df))

pprint(schema(df))

coerce!(df, :Class => OrderedFactor{2});

y, X = unpack(df, ==(:Class),name->true, rng = RANDOM_SEED);

transformer_instance = Standardizer()
transformer_model = machine(transformer_instance, X)
fit!(transformer_model)
X = MLJ.transform(transformer_model, X);

train, test = partition(eachindex(y), 0.8, shuffle=true, rng=RANDOM_SEED);

for m in models(matching(X, y))
    println("Model name = ",m.name,", ","Prediction type = ",m.prediction_type,", ","Package name = ",m.package_name);
end

model_names=Vector{String}();
loss_acc=[];
loss_ce=[];
loss_f1=[];

figure(figsize=(8, 6))
for m in models(matching(X, y))
    if m.prediction_type==Symbol("probabilistic") && m.package_name=="ScikitLearn" && m.name!="LogisticCVClassifier"
        #Excluding LogisticCVClassfiier as we can infer similar baseline results from the LogisticClassifier

        #Capturing the model and loading it using the @load utility
        model_name=m.name
        package_name=m.package_name
        eval(:(clf = @load $model_name pkg=$package_name verbosity=1))

        #Fitting the captured model onto the training set
        clf_machine = machine(clf(), X, y)
        fit!(clf_machine, rows=train)

        #Getting the predictions onto the test set
        y_pred = MLJ.predict(clf_machine, rows=test);

        #Plotting the ROC-AUC curve for each model being iterated
        fprs, tprs, thresholds = roc(y_pred, y[test])
        plot(fprs, tprs,label=model_name);

        #Obtaining different evaluation metrics
        ce_loss = mean(cross_entropy(y_pred,y[test]))
        acc = accuracy(mode.(y_pred), y[test])
        f1_score = f1score(mode.(y_pred), y[test])

        #Adding the different obtained values of the evaluation metrics to the respective vectors
        push!(model_names, m.name)
        append!(loss_acc, acc)
        append!(loss_ce, ce_loss)
        append!(loss_f1, f1_score)
    end
end

#Adding labels and legend to the ROC-AUC curve
xlabel("False Positive Rate")
ylabel("True Positive Rate")
legend(loc="best", fontsize="xx-small")
title("ROC curve")
plt.savefig(joinpath(@OUTPUT, "breastcancer_auc_curve.svg")) # hide

model_info=DataFrame(ModelName=model_names,Accuracy=loss_acc,CrossEntropyLoss=loss_ce,F1Score=loss_f1);

pprint(sort!(model_info,[:CrossEntropyLoss]));


@testset "Test DIRECTIONAL.S2 on short predefined sequence" begin
    a = [1, 1, 0, 1, 0, 0, 0, 1, 1, 1, 1, 0]

    corr = mean(Directional.s2(a, 1; len = 12))
    @test corr ≈ [7, 4, 3, 2, 1, 1, 2, 3, 2, 2, 1, 0] ./ [12 - x for x in 0:11]

    corr = mean(Directional.s2(a, 0; len = 12))
    @test corr ≈ [5, 2, 2, 1, 1, 1, 1, 1, 0, 1, 0, 0] ./ [12 - x for x in 0:11]

    corr = mean(Directional.s2(a, 1; len = 12, periodic = true))
    @test corr ≈ [7, 4, 4, 4, 3, 4, 4, 4, 3, 4, 4, 4] ./ 12

    corr = mean(Directional.s2(a, 0; len = 12, periodic = true))
    @test corr ≈ [5, 2, 2, 2, 1, 2, 2, 2, 1, 2, 2, 2] ./ 12
end

@testset "Test DIRECTIONAL.L2 on short predefined sequences" begin
    a = [1, 0, 1, 0, 1, 1, 1]

    corr = mean(Directional.l2(a, 1; len = 7))
    @test corr ≈ [5, 2, 1, 0, 0, 0, 0] ./ [7 - x for x in 0:6]

    corr = mean(Directional.l2(a, 0; len = 7))
    @test corr ≈ [2, 0, 0, 0, 0, 0, 0] ./ [7 - x for x in 0:6]

    corr = mean(Directional.l2(a, 1; len = 7, periodic = true))
    @test corr ≈ [5, 3, 2, 1, 0, 0, 0] ./ 7

    a = [1, 1, 0, 1, 0]
    @test Directional.l2(a, 1; len = 5).success ==
        Directional.l2(a, 1; len = 5, periodic = true).success

    a = [0, 1, 1, 0, 1]
    @test Directional.l2(a, 1; len = 5).success ==
        Directional.l2(a, 1; len = 5, periodic = true).success
end

@testset "Test on array with all elements equal to 1" begin
    a = fill(1, (50, 50, 50))
    for p in (false, true)
        for func in (Directional.s2, Directional.l2, Directional.c2)
            @test all(x -> x ≈ 1,
                      mean(func(a, 1; periodic = p, directions = known_directions)))
        end
    end

    @test length(pore_size(a, 0).weights) == 0
end


using FinEtools
using FinEtools.AlgoDeforLinearModule
using FinEtools.MeshExportModule

mu = 756000
nu = 0.49546;
E = 2*(1.0 + nu) * mu;
width = 48.0; height = 44.0; thickness  = 10.0;
free_height  = 16.0;
Mid_edge  = [48.0 52.0 0.0];# Location of tracked  deflection
magn = 1.0/(free_height*thickness);# Density of applied load
nt = 1
n = 16*nt;# number of elements per side
tolerance = minimum([width, height])/n/1000.;#Geometrical tolerance

fens,fes = T10block(width, height, thickness, n, n, nt)

# Reshape into a trapezoidal panel
for i = 1:count(fens)
    fens.xyz[i,2] = fens.xyz[i,2]+(fens.xyz[i,1]/width)*(height -fens.xyz[i,2]/height*(height-free_height));
end

# Clamped edge of the membrane
l1 = selectnode(fens; box=[0.,0.,-Inf, Inf,-Inf, Inf], inflate = tolerance)
ess1 = FDataDict("displacement"=>  0.0, "component"=> 1, "node_list"=>l1)
ess2 = FDataDict("displacement"=>  0.0, "component"=> 2, "node_list"=>l1)
ess3 = FDataDict("displacement"=>  0.0, "component"=> 3, "node_list"=>l1)

# Traction on the opposite edge
boundaryfes =  meshboundary(fes);
Toplist  = selectelem(fens, boundaryfes, box= [width, width, -Inf, Inf, -Inf, Inf], inflate=  tolerance);
el1femm = FEMMBase(IntegDomain(subset(boundaryfes, Toplist), SimplexRule(2, 3)))
flux1 = FDataDict("traction_vector"=>[0.0, 0.0, +magn],
    "femm"=>el1femm
    )

# Make the region
MR = DeforModelRed3D
material = MatDeforElastIso(MR,  0.0, E, nu, 0.0)
region1 = FDataDict("femm"=>FEMMDeforLinearMST10(MR,
    IntegDomain(fes, SimplexRule(3, 2)), material))

modeldata = FDataDict("fens"=>fens,
 "regions"=>[region1],
 "essential_bcs"=>[ess1, ess2, ess3],
 "traction_bcs"=>[flux1]
 )

# Call the solver
modeldata = AlgoDeforLinearModule.linearstatics(modeldata)

u = modeldata["u"]
geom = modeldata["geom"]

# Extract the solution
nl = selectnode(fens, box=[Mid_edge[1],Mid_edge[1],Mid_edge[2],Mid_edge[2], -Inf, Inf],
          inflate=tolerance);
theutip = mean(u.values[nl,:])
println("displacement =$(theutip[2]) as compared to converged $convutip")

modeldata["postprocessing"] = FDataDict("file"=>"cook_like-ew",
   "quantity"=>:Cauchy, "component"=>:xy)
modeldata = AlgoDeforLinearModule.exportstresselementwise(modeldata)
fld = modeldata["postprocessing"]["exported"][1]["field"]
println("range of Cauchy_xy = $((minimum(fld.values), maximum(fld.values)))")
File = modeldata["postprocessing"]["exported"][1]["file"]
@async run(`"paraview.exe" $File`)

modeldata["postprocessing"] = FDataDict("file"=>"cook_like-ew-vm",
   "quantity"=>:vm, "component"=>1)
modeldata = AlgoDeforLinearModule.exportstresselementwise(modeldata)
fld = modeldata["postprocessing"]["exported"][1]["field"]
println("range of vm = $((minimum(fld.values), maximum(fld.values)))")
File = modeldata["postprocessing"]["exported"][1]["file"]
@async run(`"paraview.exe" $File`)

modeldata["postprocessing"] = FDataDict("file"=>"cook_like-ew-pressure",
   "quantity"=>:pressure, "component"=>1)
modeldata = AlgoDeforLinearModule.exportstresselementwise(modeldata)
fld = modeldata["postprocessing"]["exported"][1]["field"]
println("range of pressure = $((minimum(fld.values), maximum(fld.values)))")
File = modeldata["postprocessing"]["exported"][1]["file"]
@async run(`"paraview.exe" $File`)

modeldata["postprocessing"] = FDataDict("file"=>"cook_like-ew-princ1",
   "quantity"=>:princCauchy, "component"=>1)
modeldata = AlgoDeforLinearModule.exportstresselementwise(modeldata)
fld = modeldata["postprocessing"]["exported"][1]["field"]
println("range of princCauchy Max = $((minimum(fld.values), maximum(fld.values)))")
File = modeldata["postprocessing"]["exported"][1]["file"]
@async run(`"paraview.exe" $File`)

modeldata["postprocessing"] = FDataDict("file"=>"cook_like-ew-princ3",
   "quantity"=>:princCauchy, "component"=>3)
modeldata = AlgoDeforLinearModule.exportstresselementwise(modeldata)
fld = modeldata["postprocessing"]["exported"][1]["field"]
println("range of princCauchy Min = $((minimum(fld.values), maximum(fld.values)))")
File = modeldata["postprocessing"]["exported"][1]["file"]
@async run(`"paraview.exe" $File`)

true


abstract AbstractAffineMap <: Transformation

"""
    Translation(v) <: AbstractAffineMap
    Translation(dx, dy)       (2D)
    Translation(dx, dy, dz)   (3D)

Construct the `Translation` transformation for translating Cartesian points by
an offset `v = (dx, dy, ...)`
"""
immutable Translation{V <: AbstractVector} <: AbstractAffineMap
    v::V
end
Translation(x::Tuple) = Translation(SVector(x))
Translation(x,y) = Translation(SVector(x,y))
Translation(x,y,z) = Translation(SVector(x,y,z))
Base.show(io::IO, trans::Translation) = print(io, "Translation$((trans.v...))")

function (trans::Translation{V}){V}(x)
    x + trans.v
end

Base.inv(trans::Translation) = Translation(-trans.v)

function compose(trans1::Translation, trans2::Translation)
    Translation(trans1.v + trans2.v)
end

transform_deriv(trans::Translation, x) = I
transform_deriv_params(trans::Translation, x) = I

function Base.isapprox(t1::Translation, t2::Translation; kwargs...)
    isapprox(t1.v, t2.v; kwargs...)
end

"""
    LinearMap <: AbstractAffineMap
    LinearMap(M)

A general linear transformation, constructed using `LinearMap(M)`
for any `AbstractMatrix` `M`.
"""
immutable LinearMap{M <: AbstractMatrix} <: AbstractAffineMap
    m::M
end
Base.show(io::IO, trans::LinearMap)   = print(io, "LinearMap($(trans.m))") # TODO make this output more petite

function (trans::LinearMap{M}){M}(x)
    trans.m * x
end

Base.inv(trans::LinearMap) = LinearMap(inv(trans.m))

compose(t1::LinearMap, t2::LinearMap) = LinearMap(t1.m * t2.m)

function Base.isapprox(t1::LinearMap, t2::LinearMap; kwargs...)
    isapprox(t1.m, t2.m; kwargs...)
end

function Base.isapprox(t1::LinearMap, t2::Translation; kwargs...)
    isapprox(vecnorm(t1.m), 0; kwargs...) &&
        isapprox(vecnorm(t2.v),0; kwargs...)
end

function Base.isapprox(t1::Translation, t2::LinearMap; kwargs...)
    isapprox(vecnorm(t1.v), 0; kwargs...) &&
        isapprox(vecnorm(t2.m),0; kwargs...)
end

function Base.:(==)(t1::LinearMap, t2::Translation)
    vecnorm(t1.m) == 0 &&
        0 == vecnorm(t2.v)
end

function Base.:(==)(t1::Translation, t2::LinearMap)
    vecnorm(t1.v) == 0 &&
        vecnorm(t2.m) == 0
end

transform_deriv(trans::LinearMap, x) = trans.m
# TODO transform_deriv_params

"""
    AffineMap <: AbstractAffineMap

A concrete affine transformation.  To construct the mapping `x -> M*x + v`, use

    AffineMap(M, v)

where `M` is a matrix and `v` a vector.  An arbitrary `Transformation` may be
converted into an affine approximation by linearizing about a point `x` using

    AffineMap(trans, [x])

For transformations which are already affine, `x` may be omitted.
"""
immutable AffineMap{M <: AbstractMatrix, V <: AbstractVector} <: AbstractAffineMap
    m::M
    v::V
end

function (trans::AffineMap{M,V}){M,V}(x)
    trans.m * x + trans.v
end

# Note: the expression `Tx - dT*Tx` will have large cancellation error for
# large Tx!  However, changing the order of applying the matrix and
# translation won't fix things, because then we'd have `Tx*(x-x0)` which
# also can incur large cancellation error in `x-x0`.
"""
    AffineMap(trans::Transformation, x0)

Create an Affine transformation corresponding to the differential transformation
of `x0 + dx` according to `trans`, i.e. the Affine transformation that is
locally most accurate in the vicinity of `x0`.
"""
function AffineMap(trans::Transformation, x0)
    dT = transform_deriv(trans, x0)
    Tx = trans(x0)
    AffineMap(dT, Tx - dT*x0)
end

Base.show(io::IO, trans::AffineMap) = print(io, "AffineMap($(trans.m), $(trans.v))") # TODO make this output more petite

function compose(t1::Translation, t2::LinearMap)
    AffineMap(t2.m, t1.v)
end

function compose(t1::LinearMap, t2::Translation)
    AffineMap(t1.m, t1.m * t2.v)
end

function compose(t1::AffineMap, t2::AffineMap)
    AffineMap(t1.m * t2.m, t1.v + t1.m * t2.v)
end

function compose(t1::AffineMap, t2::LinearMap)
    AffineMap(t1.m * t2.m, t1.v)
end

function compose(t1::LinearMap, t2::AffineMap)
    AffineMap(t1.m * t2.m, t1.m * t2.v)
end

function compose(t1::AffineMap, t2::Translation)
    AffineMap(t1.m, t1.v + t1.m * t2.v)
end

function compose(t1::Translation, t2::AffineMap)
    AffineMap(t2.m, t1.v + t2.v)
end

function Base.inv(trans::AffineMap)
    m_inv = inv(trans.m)
    AffineMap(m_inv, m_inv * (-trans.v))
end

function Base.isapprox(t1::AffineMap, t2::AffineMap; kwargs...)
    isapprox(t1.m, t2.m; kwargs...) &&
        isapprox(t1.v, t2.v; kwargs...)
end

function Base.isapprox(t1::AffineMap, t2::Translation; kwargs...)
    isapprox(vecnorm(t1.m), 0; kwargs...) &&
        isapprox(t1.v, t2.v; kwargs...)
end

function Base.isapprox(t1::Translation, t2::AffineMap; kwargs...)
    isapprox(vecnorm(t2.m), 0; kwargs...) &&
        isapprox(t1.v, t2.v; kwargs...)
end

function Base.isapprox(t1::AffineMap, t2::LinearMap; kwargs...)
    isapprox(t1.m, t2.m; kwargs...) &&
        isapprox(vecnorm(t1.v), 0; kwargs...)
end

function Base.isapprox(t1::LinearMap, t2::AffineMap; kwargs...)
    isapprox(t1.m, t2.m; kwargs...) &&
        isapprox(0, vecnorm(t2.v); kwargs...)
end

function Base.:(==)(t1::AffineMap, t2::Translation)
    vecnorm(t1.m) == 0 &&
        t1.v == t2.v
end

function Base.:(==)(t1::Translation, t2::AffineMap)
    vecnorm(t2.m) == 0 &&
        t1.v == t2.v
end

function Base.:(==)(t1::AffineMap, t2::LinearMap)
    t1.m == t2.m &&
        vecnorm(t1.v) == 0
end

function Base.:(==)(t1::LinearMap, t2::AffineMap)
    t1.m == t2.m &&
        0 == vecnorm(t2.v)
end

transform_deriv(trans::AffineMap, x) = trans.m
# TODO transform_deriv_params


"""
    LDA_X(ρ; α = 2/3)

Gives Slater's approximation of the exchange energy density
``\\epsilon_x``, and potential ``V_x``. Slater’s Xα.

## References

1. P. A. M. Dirac, Math. Proc. Cambridge Philos. Soc. 26, 376 (1930) (doi: 10.1017/S0305004100016108)
2. F. Bloch, Z. Phys. 57, 545 (1929) (doi: 10.1007/BF01340281)
3. J.C. Slater, Phys. Rev. 81, 385 (1951) (doi: 10.1103/PhysRev.81.385)
4. J.C. Slater, T.M. Wilson, and J.H. Wood, Phys. Rev. 179, 28 (1969)
"""
function LDA_X(ρ; α = 2/3)
  vx = -3α * (3ρ/8π)^(1/3)
  ϵx = 3/4 * vx
  return vx, ϵx
end

"""
    LDA_C_CHACHIYO(ρ; a  = -0.01554535, b₁ = 20.4562557, b  = 20.4562557)

Gives Chachiyo's approximation of the correlation energy density
``\\epsilon_x``, and potential ``V_x``.

## References

1. T. Chachiyo, J. Chem. Phys. 145, 021101 (2016) (doi: 10.1063/1.4958669)
"""
function LDA_C_CHACHIYO(ρ; a  = -0.01554535, b₁ = 20.4562557, b  = 20.4562557)
  rs = (3/4π / ρ)^(1/3)
  ϵc = a * log(1 + b₁ / rs + b / rs^2)
  vc = ϵc + a * (b₁ * rs + 2b) / (3 * (rs^2 + b₁ * rs + b))
  return vc, ϵc
end

"""
    LDA_C_VWN(ρ; a = 0.0310907, b = 3.72744, c = 12.9352, x0 = -0.10498)

Gives VWN5 approximation of the correlation energy density
``\\epsilon_x``, and potential ``V_x``.

## References

1. S.H. Vosko, L. Wilk, and M. Nusair, Can. J. Phys. 58, 1200 (1980) (doi: 10.1139/p80-159)
"""
function LDA_C_VWN(ρ; a = 0.0310907, b = 3.72744, c = 12.9352, x0 = -0.10498)
  rs = (3/4π/ρ)^(1/3)
  q = sqrt(4.0 * c - b * b)
  f1 = 2.0 * b / q
  f2 = b * x0 / (x0 * x0 + b * x0 + c)
  f3 = 2.0 * (2.0 * x0 + b) / q
  rs12 = sqrt(rs)
  fx = rs + b * rs12 + c
  qx = atan(q / (2.0 * rs12 + b) )
  ϵc = a * (log(rs/fx) + f1 * qx - f2 * (log((rs12 - x0)^2 / fx) + f3 * qx))
  tx = 2.0 * rs12 + b
  tt = tx * tx + q * q
  vc = ϵc - rs12 * a / 6.0 * (2.0 / rs12 - tx / fx - 4.0 * b / tt - f2 * (2.0 / (rs12 - x0) - tx / fx - 4.0 * (2.0 * x0 + b) / tt))
  return vc, ϵc
end

# TODO:
# V. V. Karasiev, J. Chem. Phys. 145, 157101 (2016) (doi: 10.1063/1.4964758)
# LDA_C_KARASIEV_MOD(ρ) = LDA_C_CHACHIYO(ρ, b₁ = 21.7392245)


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

"""
    BisectPointPartitioner(normal, point)

A method for partitioning spatial data into two half spaces
defined by a `normal` direction and a reference `point`.
"""
struct BisectPointPartitioner{T,N} <: AbstractPartitioner
  normal::SVector{N,T}
  point::SVector{N,T}
end

BisectPointPartitioner(normal::NTuple{N,T},
                  point::NTuple{N,T}=ntuple(i->zero(T), N)) where {T,N} =
  BisectPointPartitioner{T,N}(normalize(SVector(normal)), SVector(point))

function partition(object::AbstractSpatialObject{T,N},
                   partitioner::BisectPointPartitioner{T,N}) where {T,N}
  n = partitioner.normal
  p = partitioner.point

  x = MVector{N,T}(undef)

  left  = Vector{Int}()
  right = Vector{Int}()
  for location in 1:npoints(object)
    coordinates!(x, object, location)
    if (x - p) ⋅ n < zero(T)
      push!(left, location)
    else
      push!(right, location)
    end
  end

  SpatialPartition(object, [left,right])
end


module Day9
export solvePart1, solvePart2

numberOfPlayers = 418
lastMarblePoints = 71339

mutable struct Marble
    next
    prev
    number::Int64
end

function indexCounterClockwise(ind, numPlaces, len)
    i = ind - numPlaces
    while i < 0
        i += len
    end
    return i
end

function indexClockwise(ind, numPlaces, len)
    i = ind + numPlaces
    while i > len
        i -= len
    end
    return i
end

function removeMarbleSevenPlacesCounterClockwise(current::Marble)
    toRemove = current
    for i in 1:7
        toRemove = toRemove.prev
    end
    before = toRemove.prev
    after = toRemove.next
    before.next = after
    after.prev = before
    return toRemove
end

function winningElfScore(numberOfPlayers, lastNumberScore)
    # Start with a marble whose previous and next is itself
    firstMarble = Marble(nothing, nothing, 0)
    firstMarble.next = firstMarble
    firstMarble.prev = firstMarble

    whoseTurn = 1
    points = Dict()
    current = firstMarble
    for nextMarbleNumber in 1:lastNumberScore
        if nextMarbleNumber % 23 == 0
            removed = removeMarbleSevenPlacesCounterClockwise(current)
            toGain = nextMarbleNumber + removed.number
            points[whoseTurn] = get(points, whoseTurn, 0) + toGain
            current = removed.next
        else
            before = current.next
            after = before.next
            toInsert = Marble(before, after, nextMarbleNumber)
            before.next = toInsert
            after.prev = toInsert
            toInsert.prev = before
            toInsert.next = after
            current = toInsert
        end

        whoseTurn += 1
        if whoseTurn > numberOfPlayers
            whoseTurn = 1
        end
    end
    return max(values(points)...)
end

function solvePart1()
    winningElfScore(numberOfPlayers, lastMarblePoints)
end

function solvePart2()
    winningElfScore(numberOfPlayers, lastMarblePoints*100)
end
end


# SPDX-License-Identifier: X11
# 2020-11-09 (Timed)
# Almost GCD (200pt)

function main()
  readline()
  A = parse.(Int, split(readline()))

  max,gcd = -1,-1
  for i ∈ 2:maximum(A)
    if (g = gcdness(i, A)) > gcd
      max = i
      gcd = g
    end
  end

  println(max)
end

function gcdness(k::Integer, A::Array{<:Integer})::Integer
  cnt = 0
  for a ∈ A
    a % k == 0 && (cnt += 1)
  end
  return cnt
end

main()


using QuantEcon
using Optim

# utility and marginal utility functions
u(x) = log(x)
du(x) = 1 / x

"""
Income fluctuation problem

##### Fields

- `r::Real` : Strictly positive interest rate
- `R::Real` : The interest rate plus 1 (strictly greater than 1)
- `bet::Real` : Discount rate in (0, 1)
- `b::Real` :  The borrowing constraint
- `Pi::Matrix{T} where T<:Real` : Transition matrix for `z`
- `z_vals::Vector{T} where T<:Real` : Levels of productivity
- `asset_grid::AbstractArray` : Grid of asset values

"""
struct ConsumerProblem{T <: Real}
    r::T
    R::T
    bet::T
    b::T
    Pi::Matrix{T}
    z_vals::Vector{T}
    asset_grid::AbstractArray
end

function ConsumerProblem(;r=0.01,
                         bet=0.96,
                         Pi=[0.6 0.4; 0.05 0.95],
                         z_vals=[0.5, 1.0],
                         b=0.0,
                         grid_max=16,
                         grid_size=50)
    R = 1 + r
    asset_grid = linspace(-b, grid_max, grid_size)

    ConsumerProblem(r, R, bet, b, Pi, z_vals, asset_grid)
end

"""
Apply the Bellman operator for a given model and initial value.

##### Arguments

- `cp::ConsumerProblem` : Instance of `ConsumerProblem`
- `v::Matrix`: Current guess for the value function
- `out::Matrix` : Storage for output
- `;ret_policy::Bool(false)`: Toggles return of value or policy functions

##### Returns

None, `out` is updated in place. If `ret_policy == true` out is filled with the
policy function, otherwise the value function is stored in `out`.

"""
function bellman_operator!(cp::ConsumerProblem,
                          V::Matrix,
                          out::Matrix;
                          ret_policy::Bool=false)

    # simplify names, set up arrays
    R, Pi, bet, b = cp.R, cp.Pi, cp.bet, cp.b
    asset_grid, z_vals = cp.asset_grid, cp.z_vals
    z_idx = 1:length(z_vals)

    # value function when the shock index is z_i
    vf = interp(asset_grid, V)

    opt_lb = 1e-8

    # solve for RHS of Bellman equation
    for (i_z, z) in enumerate(z_vals)
        for (i_a, a) in enumerate(asset_grid)

            function obj(c)
                EV = dot(vf.(R*a+z-c, z_idx), Pi[i_z, :]) # compute expectation
                return -u(c)  - bet * EV
            end
            res = optimize(obj, opt_lb, R.*a.+z.+b)
            c_star = Optim.minimizer(res)

            if ret_policy
                out[i_a, i_z] = c_star
            else
               out[i_a, i_z] = - Optim.minimum(res)
            end

        end
    end
    out
end

bellman_operator(cp::ConsumerProblem, V::Matrix; ret_policy=false) =
    bellman_operator!(cp, V, similar(V); ret_policy=ret_policy)

"""
Extract the greedy policy (policy function) of the model.

##### Arguments

- `cp::CareerWorkerProblem` : Instance of `CareerWorkerProblem`
- `v::Matrix`: Current guess for the value function
- `out::Matrix` : Storage for output

##### Returns

None, `out` is updated in place to hold the policy function

"""
get_greedy!(cp::ConsumerProblem, V::Matrix, out::Matrix) =
    update_bellman!(cp, V, out, ret_policy=true)

get_greedy(cp::ConsumerProblem, V::Matrix) =
    update_bellman(cp, V, ret_policy=true)

"""
The approximate Coleman operator.

Iteration with this operator corresponds to policy function
iteration. Computes and returns the updated consumption policy
c.  The array c is replaced with a function cf that implements
univariate linear interpolation over the asset grid for each
possible value of z.

##### Arguments

- `cp::CareerWorkerProblem` : Instance of `CareerWorkerProblem`
- `c::Matrix`: Current guess for the policy function
- `out::Matrix` : Storage for output

##### Returns

None, `out` is updated in place to hold the policy function

"""
function coleman_operator!(cp::ConsumerProblem, c::Matrix, out::Matrix)
    # simplify names, set up arrays
    R, Pi, bet, b = cp.R, cp.Pi, cp.bet, cp.b
    asset_grid, z_vals = cp.asset_grid, cp.z_vals
    z_idx = 1:length(z_vals)
    gam = R * bet

    # policy function when the shock index is z_i
    cf = interp(asset_grid, c)

    # compute lower_bound for optimization
    opt_lb = 1e-8

    for (i_z, z) in enumerate(z_vals)
        for (i_a, a) in enumerate(asset_grid)
            function h(t)
                cps = cf.(R*a+z-t, z_idx) # c' for each z'
                expectation = dot(du.(cps), Pi[i_z, :])
                return abs(du(t) - max(gam * expectation, du(R*a+z+b)))
            end
            opt_ub = R*a + z + b  # addresses issue #8 on github
            res = optimize(h, min(opt_lb, opt_ub - 1e-2), opt_ub,
                           method=Optim.Brent())
            out[i_a, i_z] = Optim.minimizer(res)
        end
    end
    out
end

"""
Apply the Coleman operator for a given model and initial value

See the specific methods of the mutating version of this function for more
details on arguments
"""
coleman_operator(cp::ConsumerProblem, c::Matrix) =
    coleman_operator!(cp, c, similar(c))

function initialize(cp::ConsumerProblem)
    # simplify names, set up arrays
    R, bet, b = cp.R, cp.bet, cp.b
    asset_grid, z_vals = cp.asset_grid, cp.z_vals
    shape = length(asset_grid), length(z_vals)
    V, c = Array{Float64}(shape...), Array{Float64}(shape...)

    # Populate V and c
    for (i_z, z) in enumerate(z_vals)
        for (i_a, a) in enumerate(asset_grid)
            c_max = R*a + z + b
            c[i_a, i_z] = c_max
            V[i_a, i_z] = u(c_max) ./ (1 - bet)
        end
    end

    return V, c
end


# This file is a part of BAT.jl, licensed under the MIT License (MIT).

using BAT
using Test

using LinearAlgebra
using ValueShapes, Distributions, ArraysOfArrays
using ForwardDiff, Zygote, DistributionsAD
using InverseFunctions, ChangesOfVariables

@testset "test_distribution_transform" begin
    function test_back_and_forth(trg_d, src_d)
        @testset "transform $(typeof(trg_d).name) <-> $(typeof(src_d).name)" begin
            src_v = rand(src_d)
            trg_v = BAT.apply_dist_trafo(trg_d, src_d, src_v)
            src_v_reco = BAT.apply_dist_trafo(src_d, trg_d, trg_v)

            @test src_v ≈ src_v_reco

            let vs_trg = varshape(trg_d), vs_src = varshape(src_d)
                f = unshaped_x -> inverse(vs_trg)(BAT.apply_dist_trafo(trg_d, src_d, vs_src(unshaped_x)))
                ref_ladj = logpdf(src_d, src_v) - logpdf(trg_d, trg_v)
                @test ref_ladj ≈ logabsdet(ForwardDiff.jacobian(f, inverse(vs_src)(src_v)))[1]
            end
        end
    end

    function get_trgxs(trg_d, src_d, X)
        return (x -> BAT.apply_dist_trafo(trg_d, src_d, x)).(nestedview(X))
    end

    function get_trgxs(trg_d, src_d::Distribution{Univariate}, X)
        return (x -> BAT.apply_dist_trafo(trg_d, src_d, x)).(X)
    end

    function test_dist_trafo_moments(trg_d, src_d)
        @testset "check moments of trafo $(typeof(trg_d).name) <- $(typeof(src_d).name)" begin
            X = flatview(rand(src_d, 10^5))
            trgxs = get_trgxs(trg_d, src_d, X)
            unshaped_trgxs = broadcast(unshaped, trgxs, Ref(varshape(trg_d)))
            @test isapprox(mean(unshaped_trgxs), mean(unshaped(trg_d)), atol = 0.1)
            @test isapprox(cov(unshaped_trgxs), cov(unshaped(trg_d)), rtol = 0.1)
        end
    end

    stduvuni = BAT.StandardUvUniform()
    stduvnorm = BAT.StandardUvNormal()

    uniform1 = Uniform(-5.0, -0.01)
    uniform2 = Uniform(0.01, 5.0)

    normal1 = Normal(-10, 1)
    normal2 = Normal(10, 5)

    stdmvnorm1 = BAT.StandardMvNormal(1)
    stdmvnorm2 = BAT.StandardMvNormal(2)

    stdmvuni2 = BAT.StandardMvUniform(2)

    standnorm2_reshaped = ReshapedDist(stdmvnorm2, varshape(stdmvnorm2))

    mvnorm = MvNormal([0.3, -2.9], [1.7 0.5; 0.5 2.3])
    beta = Beta(3,1)
    gamma = Gamma(0.1,0.7)
    dirich = Dirichlet([0.1,4])

    ntdist = NamedTupleDist(
        a = uniform1,
        b = mvnorm,
        c = [4.2, 3.7],
        x = beta,
        y = gamma
    )

    test_back_and_forth(stduvuni, stduvuni)
    test_back_and_forth(stduvnorm, stduvnorm)
    test_back_and_forth(stduvuni, stduvnorm)
    test_back_and_forth(stduvnorm, stduvuni)

    test_back_and_forth(stdmvuni2, stdmvuni2)
    test_back_and_forth(stdmvnorm2, stdmvnorm2)
    test_back_and_forth(stdmvuni2, stdmvnorm2)
    test_back_and_forth(stdmvnorm2, stdmvuni2)

    test_back_and_forth(beta, stduvnorm)
    test_back_and_forth(gamma, stduvnorm)

    test_dist_trafo_moments(normal2, normal1)
    test_dist_trafo_moments(uniform2, uniform1)

    test_dist_trafo_moments(beta, gamma)

    test_dist_trafo_moments(beta, stduvnorm)
    test_dist_trafo_moments(gamma, stduvnorm)

    test_dist_trafo_moments(mvnorm, stdmvnorm2)
    test_dist_trafo_moments(dirich, stdmvnorm1)

    test_dist_trafo_moments(mvnorm, stdmvuni2)
    test_dist_trafo_moments(stdmvuni2, mvnorm)

    test_dist_trafo_moments(stdmvnorm2, stdmvuni2)

    test_dist_trafo_moments(mvnorm, standnorm2_reshaped)
    test_dist_trafo_moments(standnorm2_reshaped, mvnorm)
    test_dist_trafo_moments(stdmvnorm2, standnorm2_reshaped)
    test_dist_trafo_moments(standnorm2_reshaped, standnorm2_reshaped)

    test_back_and_forth(ntdist, BAT.StandardMvNormal(5))
    test_back_and_forth(ntdist, BAT.StandardMvUniform(5))

    let
        mvuni = product_distribution([Uniform(), Uniform()])

        x = rand()
        @test_throws ArgumentError BAT.apply_dist_trafo(stduvnorm, mvnorm, x)
        @test_throws ArgumentError BAT.apply_dist_trafo(stduvnorm, stdmvnorm1, x)
        @test_throws ArgumentError BAT.apply_dist_trafo(stduvnorm, stdmvnorm2, x)

        x = rand(2)
        @test_throws ArgumentError BAT.apply_dist_trafo(mvuni, mvnorm, x)
        @test_throws ArgumentError BAT.apply_dist_trafo(mvnorm, mvuni, x)
        @test_throws ArgumentError BAT.apply_dist_trafo(stduvnorm, mvnorm, x)
        @test_throws ArgumentError BAT.apply_dist_trafo(stduvnorm, stdmvnorm1, x)
        @test_throws ArgumentError BAT.apply_dist_trafo(stduvnorm, stdmvnorm2, x)
    end

    let
        primary_dist = NamedTupleDist(x = Normal(2), c = 5)
        f = x -> NamedTupleDist(y = Normal(x.x, 3), z = MvNormal([1.3 0.5; 0.5 2.2]))
        trg_d = @inferred(HierarchicalDistribution(f, primary_dist))
        src_d = BAT.StandardMvNormal(totalndof(varshape(trg_d)))
        test_back_and_forth(trg_d, src_d)
        test_dist_trafo_moments(trg_d, src_d)
    end

    #=
    using Cuba
    function integrate_over_unit(density::AbstractDensity)
        vs = varshape(density)
        f_cuba(source_x, y) = y[1] = exp(logdensityof(density)(vs(source_x)))
        Cuba.vegas(f_cuba, 1, 1).integral[1]
    end
    =#

    @testset "Custom cdf and quantile for dual numbers" begin
        Dual = ForwardDiff.Dual

        @test BAT._trafo_cdf(Normal(Dual(0, 1, 0, 0), Dual(1, 0, 1, 0)), Dual(0.5, 0, 0, 1)) == cdf(Normal(Dual(0, 1, 0, 0), Dual(1, 0, 1, 0)), Dual(0.5, 0, 0, 1))
        @test BAT._trafo_cdf(Normal(0, 1), Dual(0.5, 1)) == cdf(Normal(0, 1), Dual(0.5, 1))

        @test BAT._trafo_quantile(Normal(0, 1), Dual(0.5, 1)) == quantile(Normal(0, 1), Dual(0.5, 1))
        @test BAT._trafo_quantile(Normal(Dual(0, 1, 0, 0), Dual(1, 0, 1, 0)), Dual(0.5, 0, 0, 1)) == quantile(Normal(Dual(0, 1, 0, 0), Dual(1, 0, 1, 0)), Dual(0.5, 0, 0, 1))
    end

    for VT in (NamedTuple, ShapedAsNT)
        src_dist = unshaped(NamedTupleDist(VT, a = Weibull(), b = MvNormal([1.3 0.6; 0.6 2.4])))
        f = BAT.DistributionTransform(Normal, src_dist)
        x = rand(src_dist)
        InverseFunctions.test_inverse(f, x)
        ChangesOfVariables.test_with_logabsdet_jacobian(f, x, ForwardDiff.jacobian)
    end

    @testset "trafo broadcasting" begin
        dist = NamedTupleDist(a = Weibull(), b = Exponential())
        smpls = bat_sample(dist, IIDSampling(nsamples = 100)).result
        trafo = BAT.DistributionTransform(Normal, dist)
        @inferred(broadcast(trafo, smpls)) isa DensitySampleVector
        smpls_tr = trafo.(smpls)
        smpls_tr_cmp = [trafo(s) for s in smpls]
        @test smpls_tr == smpls_tr_cmp
	    @test @inferred(varshape(trafo)) == varshape(trafo.source_dist)
	    @test @inferred(trafo(varshape(trafo))) == varshape(trafo.target_dist)
    end

    @testset "trafo composition" begin
        dist1 = @inferred(NamedTupleDist(a = Normal(), b = Uniform(), c = Cauchy()))
        dist2 = @inferred(NamedTupleDist(a = Exponential(), b = Weibull(), c = Beta()))
        normal1 = Normal()
        normal2 = Normal(2)

        trafo = @inferred(BAT.DistributionTransform(dist1, dist2))
        inv_trafo = @inferred(inverse(trafo))

        composed_trafo = @inferred(∘(trafo, inv_trafo))
        @test composed_trafo.source_dist == composed_trafo.target_dist == dist1
        @test composed_trafo ∘ trafo == trafo
        @test_throws ArgumentError  trafo ∘ composed_trafo

        trafo = @inferred(BAT.DistributionTransform(normal1, normal2))
        @test_throws ArgumentError trafo ∘ trafo
    end

    @testset "full density transform" begin
        likelihood = NamedTupleDist(a = Normal(), b = Exponential())
        prior = NamedTupleDist(a = Normal(), b = Gamma())
        posterior_density = PosteriorDensity(likelihood, prior)

        posterior_density_trafod = @inferred(bat_transform(PriorToUniform(), posterior_density, FullDensityTransform()))

        @test posterior_density_trafod.result.orig.likelihood.dist == likelihood
        @test posterior_density_trafod.result.orig.prior.dist == prior

        @test posterior_density_trafod.result.trafo.target_dist isa BAT.StandardMvUniform

        lower_bounds = Float32.([-10, -10, -10])
        upper_bounds = Float32.([10, 10, 10])
        rect_bounds = @inferred(BAT.HyperRectBounds(lower_bounds, upper_bounds))
            mvn = @inferred(product_distribution(Normal.(randn(3))))
        dist_density = @inferred(BAT.DistributionDensity(mvn, rect_bounds))

        dist_density_trafod = @inferred(bat_transform(PriorToUniform(), dist_density, FullDensityTransform()))

        @test dist_density_trafod.trafo.target_dist isa BAT.StandardMvUniform
        @test dist_density_trafod.result.orig == dist_density
        @test dist_density_trafod.trafo.source_dist == dist_density_trafod.result.trafo.source_dist == mvn

        @test dist_density_trafod.trafo(varshape(dist_density_trafod.trafo)) == @inferred(varshape(dist_density))

        dist_density_trafod = @inferred(bat_transform(PriorToGaussian(), dist_density, FullDensityTransform()))

        @test dist_density_trafod.trafo.target_dist isa BAT.StandardMvNormal
        @test dist_density_trafod.result.orig == dist_density
        @test dist_density_trafod.trafo.source_dist == dist_density_trafod.result.trafo.source_dist == mvn

        @test dist_density_trafod.trafo(varshape(dist_density_trafod.trafo)) == @inferred(varshape(dist_density))
    end

    @testset "trafo autodiff pullbacks" begin
        # ToDo: Test for type stability and fix where necessary.

        xs = rand(5)
        @test Zygote.jacobian(BAT._pushfront, xs, 42)[1] ≈ ForwardDiff.jacobian(xs -> BAT._pushfront(xs, 1), xs)
        @test Zygote.jacobian(BAT._pushfront, xs, 42)[2] ≈ vec(ForwardDiff.jacobian(x -> BAT._pushfront(xs, x[1]), [42]))
        @test Zygote.jacobian(BAT._pushback, xs, 42)[1] ≈ ForwardDiff.jacobian(xs -> BAT._pushback(xs, 1), xs)
        @test Zygote.jacobian(BAT._pushback, xs, 42)[2] ≈ vec(ForwardDiff.jacobian(x -> BAT._pushback(xs, x[1]), [42]))
        @test Zygote.jacobian(BAT._rev_cumsum, xs)[1] ≈ ForwardDiff.jacobian(BAT._rev_cumsum, xs)
        @test Zygote.jacobian(BAT._exp_cumsum_log, xs)[1] ≈ ForwardDiff.jacobian(BAT._exp_cumsum_log, xs) ≈ ForwardDiff.jacobian(cumprod, xs)

        src_v = [0.6, 0.7, 0.8, 0.9]
        f = inverse(BAT.DistributionTransform(Uniform, DistributionsAD.TuringDirichlet([3.0, 4.0, 5.0, 6.0, 7.0])))
        @test isapprox(ForwardDiff.jacobian(f, src_v), Zygote.jacobian(f, src_v)[1], rtol = 10^-4)
        f = inverse(BAT.DistributionTransform(Uniform, Dirichlet([3.0, 4.0, 5.0, 6.0, 7.0])))
        @test isapprox(ForwardDiff.jacobian(f, src_v), Zygote.jacobian(f, src_v)[1], rtol = 10^-4)
        f = inverse(BAT.DistributionTransform(Normal, Dirichlet([3.0, 4.0, 5.0, 6.0, 7.0])))
        @test isapprox(ForwardDiff.jacobian(f, src_v), Zygote.jacobian(f, src_v)[1], rtol = 10^-4)
    end
end

@testset "bat_transform_defaults" begin
    mvn = @inferred(product_distribution([Normal(-1), Normal(), Normal(1)]))
    uniform_prior = @inferred(product_distribution([Uniform(-3, 1), Uniform(-2, 2), Uniform(-1, 3)]))

    posterior_uniform_prior = @inferred(PosteriorDensity(mvn, uniform_prior))
    posterior_gaussian_prior = @inferred(PosteriorDensity(mvn, mvn))

    @test @inferred(bat_transform(PriorToGaussian(), posterior_uniform_prior)).result.prior.dist == @inferred(BAT.StandardMvNormal(3))
    @test @inferred(bat_transform(PriorToUniform(), posterior_gaussian_prior)).result.prior.dist == @inferred(BAT.StandardMvUniform(3))
    @test @inferred(bat_transform(NoDensityTransform(), posterior_uniform_prior)).result.prior.dist == uniform_prior
    pd = @inferred(product_distribution([Uniform() for i in 1:3]))
    density = @inferred(BAT.DistributionDensity(pd))
    @test @inferred(bat_transform(NoDensityTransform(), density)).result.dist == density.dist

    # ToDo: Improve comparison for bounds so `.dist` is not required here:
    @inferred(bat_transform(PriorToUniform(), convert(AbstractDensity, BAT.StandardUvUniform()))).result.dist == convert(AbstractDensity, BAT.StandardUvUniform()).dist
    @inferred(bat_transform(PriorToUniform(), convert(AbstractDensity, BAT.StandardMvUniform(4)))).result.dist == convert(AbstractDensity, BAT.StandardMvUniform(4)).dist
    @inferred(bat_transform(PriorToGaussian(), convert(AbstractDensity, BAT.StandardUvNormal()))).result.dist == convert(AbstractDensity, BAT.StandardUvNormal()).dist
    @inferred(bat_transform(PriorToGaussian(), convert(AbstractDensity, BAT.StandardMvNormal(4)))).result.dist == convert(AbstractDensity, BAT.StandardMvNormal(4)).dist
end


using DynamicBounds, Plots
pyplot()

x0(p) = [34.0; 20.0; 0.0; 0.0; 16.0; 0.0]
function f!(du, u, p, t)
    du[1] = -p[1]*u[1]*u[2] + p[2]*u[3] + p[6]*u[6]
    du[2] = -p[1]*u[1]*u[2] + p[2]*u[3] + p[3]*u[3]
    du[3] =  p[1]*u[1]*u[2] - p[2]*u[3] - p[3]*u[3]
    du[4] =  p[3]*u[3] - p[4]*u[4]*u[5] + p[5]*u[6]
    du[5] = -p[4]*u[4]*u[5] + p[5]*u[6] + p[6]*u[6]
    du[6] =  p[4]*u[4]*u[5] - p[5]*u[6] - p[6]*u[6]
    return
end

function add_trajectory(plt, integrator, pL, pU)
    ratio = rand(6)
    pstar = pL.*ratio .+ pU.*(1.0 .- ratio)
    setall!(integrator, ParameterValue(), pstar)
    integrate!(integrator)
    plot!(plt, integrator.local_problem_storage.integrator_t,
    integrator.local_problem_storage.pode_x[1,:], label="", linecolor = :green,
    markershape = :+, markercolor = :green, linestyle = :dash, markersize = 2, lw=0.75)
    return plt
end

# Example from Improved relaxations for the parametric solutionsof ODEs using differential inequalities
function generate_plot()
    tspan = (0.0,18.0e-5*250)
    pL = [0.1; 0.033; 16.0; 5.0; 0.5; 0.3]
    pU = 10.0*pL

    prob = ODERelaxProb(f!, tspan, x0, pL, pU)

    A = [0.0 -1.0 -1.0  0.0  0.0  0.0;
         0.0  0.0  0.0  0.0 -1.0 -1.0;
         1.0 -1.0 0.0 1.0 -1.0 0.0]
    b = [-20.0; -16.0; -2.0]
    set!(prob, PolyhedralConstraint(A, b))

    xL = zeros(6)
    xU = [34.0; 20.0; 20.0; 34.0; 16.0; 16.0]
    set!(prob, ConstantStateBounds(xL,xU))

    Scott2013(prob, calculate_relax = false, calculate_subgradient = false)
    integrator = Scott2013(prob, calculate_relax = false, calculate_subgradient = false)

    relax!(integrator)

    plt = plot(integrator.relax_t, integrator.relax_lo[1,:], label="Interval Bounds", linecolor = :darkblue, linestyle = :dash, lw=1.5, legend=:bottomleft)
    plot!(plt, integrator.relax_t, integrator.relax_hi[1,:], label="", linecolor = :darkblue, linestyle = :dash, lw=1.5)
    ylabel!("x[1] (M)")
    xlabel!("Time (seconds)")

    ratio = rand(6)
    pstar = pL.*ratio .+ pU.*(1.0 .- ratio)
    setall!(integrator, ParameterValue(), pstar)
    integrate!(integrator)
    plot!(plt, integrator.local_problem_storage.integrator_t,
    integrator.local_problem_storage.pode_x[1,:], label="Trajectories", linecolor = :green,
    markershape = :+, markercolor = :green, linestyle = :dash, markersize = 2, lw=0.75)

    for i in 1:39
        add_trajectory(plt, integrator, pL, pU)
    end
    plt
end

plt = generate_plot()
Plots.eps(plt, "C:\\Users\\wilhe\\Desktop\\Package Development Work\\sample_plot")

#=
function add_trajectory(plt, integrator, pL, pU)
    ratio = rand(6)
    pstar = pL.*ratio .+ pU.*(1.0 .- ratio)
    setall!(integrator, ParameterValue(), pstar)
    integrate!(integrator)
    plot!(plt, integrator.local_problem_storage.integrator_t,
    integrator.local_problem_storage.pode_x[1,:], label="", lw=1)
    return plt
end

plt = add_trajectory(plt, integrator, pL, pU)
plt = add_trajectory(plt, integrator, pL, pU)
plot(plt)
show(plt)
=#

#invariant = DenseLinearInvariant(A, b)
#xnew = Interval.(xL,xU)
#xprior = Interval.(xL,xU)

#=
using Plots
plot([integrator.local_problem_storage.pode_x[1,:],
      integrator.local_problem_storage.pode_x[2,:],
      integrator.local_problem_storage.pode_x[3,:],
      integrator.local_problem_storage.pode_x[4,:],
      integrator.local_problem_storage.pode_x[5,:],
      integrator.local_problem_storage.pode_x[6,:]])
=#

#using Profile, ProfileView
#test_func(x::Int) = relax!(integrator)
#@profview test_func(1)
#@profview test_func(2)
#=
t = integrator.relax_t
println("length(t): $(length(t))")

#cv1 = integrator.relax_cv[1,:]
#cc1 = integrator.relax_cc[1,:]
lo1 = integrator.relax_lo[1,:]
hi1 = integrator.relax_hi[1,:]
val1 = integrator.local_problem_storage.pode_x[1,:]

#cv2 = integrator.relax_cv[2,:]
#cc2 = integrator.relax_cc[2,:]
lo2 = integrator.relax_lo[2,:]
hi2 = integrator.relax_hi[2,:]
val2 = integrator.local_problem_storage.pode_x[2,:]

p1 = Plots.plot(t, lo1, label="lo[1,:]", title="Dim1")
#Plots.plot!(t, cc1, label="cc[1,:]")
#Plots.plot!(t, cv1, label="cv[1,:]")
Plots.plot!(t, hi1, label="hi[1,:]")
Plots.plot!(integrator.local_problem_storage.integrator_t, val1, label="val[1,:]")

prob = integrator.relax_ode_prob
#@code_warntype prob.f.f(ScottA2013a_pL, ScottA2013a_pL, ScottA2013a_pL, ScottA2013a_tf)
=#