text:
module TestDoubleRobustnessATE

include("helper_fns.jl")

using TMLE
using Random
using Test
using Distributions
using MLJ
using StableRNGs
using StatsBase
using HypothesisTests

LogisticClassifier = @load LogisticClassifier pkg=MLJLinearModels verbosity=0
LinearRegressor = @load LinearRegressor pkg=MLJLinearModels verbosity = 0

Ns = [100, 1000, 10000, 100000]

function binary_target_binary_treatment_pb(rng;n=100)
    p_w() = 0.3
    pa_given_w(w) = 1 ./ (1 .+ exp.(-0.5w .+ 1))
    py_given_aw(a, w) = 1 ./ (1 .+ exp.(2w .- 3a .+ 1))
    # Sample from dataset
    Unif = Uniform(0, 1)
    w = rand(rng, Unif, n) .< p_w()
    t = rand(rng, Unif, n) .< pa_given_w(w)
    y = rand(rng, Unif, n) .< py_given_aw(t, w)
    # Convert to dataframe to respect the Tables.jl
    # and convert types
    W = (W=convert(Array{Float64}, w),)
    t = (t=categorical(t),)
    y = categorical(y)
    # Compute the theoretical ATE
    ATE₁ = py_given_aw(1, 1)*p_w() + (1-p_w())*py_given_aw(1, 0)
    ATE₀ = py_given_aw(0, 1)*p_w() + (1-p_w())*py_given_aw(0, 0)
    ATE = ATE₁ - ATE₀

    return t, W, y, ATE
end

"""
From https://www.degruyter.com/document/doi/10.2202/1557-4679.1043/html
The theoretical ATE is 1
"""
function continuous_target_binary_treatment_pb(rng;n=100)
    Unif = Uniform(0, 1)
    W = float(rand(rng, Bernoulli(0.5), n, 3))
    t = rand(rng, Unif, n) .< TMLE.expit(0.5W[:, 1] + 1.5W[:, 2] - W[:,3])
    y = t + 2W[:, 1] + 3W[:, 2] - 4W[:, 3] + rand(rng, Normal(0,1), n)
    # Type coercion
    W = MLJ.table(W)
    t = (t=categorical(t),)
    return t, W, y, 1
end

function continuous_target_categorical_treatment_pb(rng;n=100, control="TT", treatment="AA")
    ft(T) = (T .== "AA") - (T .== "AT") + 2(T .== "TT")
    fw(W) = 2W[:, 1] + 3W[:, 2] - 4W[:, 3]

    W = float(rand(rng, Bernoulli(0.5), n, 3))
    θ = rand(rng, 3, 3)
    softmax = exp.(W*θ) ./ sum(exp.(W*θ), dims=2)
    T = [sample(rng, ["TT", "AA", "AT"], Weights(softmax[i, :])) for i in 1:n]
    y = ft(T) + fw(W) + rand(rng, Normal(0,1), n)

    # Ew[E[Y|t,w]] = ∑ᵤ (ft(T) + fw(w))p(w) = ft(t) + 0.5
    ATE = (ft(treatment) + 0.5) -  (ft(control) + 0.5)
    # Type coercion
    W = MLJ.table(W)
    T = (T=categorical(T),)
    return T, W, y, ATE
end

@testset "Test Double Robustness ATE on continuous_target_categorical_treatment_pb" begin
    # When Q̅ is misspecified but G is well specified
    query = (T=["AA", "TT"],)
    Q̅ = MLJ.DeterministicConstantRegressor()
    G = LogisticClassifier()
    tmle = TMLEstimator(Q̅, G, query)

    abs_mean_rel_errors, abs_vars = asymptotics(
            tmle,
            continuous_target_categorical_treatment_pb,
            StableRNG(123),
            Ns
            )
    @test all(abs_mean_rel_errors .< [26, 9, 3.5, 0.6])
    @test all(abs_vars .< [0.09, 0.02, 0.002, 3.9e-5])

    # When Q̅ is well specified but G is misspecified
    query = (T=["AA", "TT"],)
    Q̅ = LinearRegressor()
    G = ConstantClassifier()
    tmle = TMLEstimator(Q̅, G, query)

    abs_mean_rel_errors, abs_vars = asymptotics(
            tmle,
            continuous_target_categorical_treatment_pb,
            StableRNG(123),
            Ns
            )
    @test all(abs_mean_rel_errors .< [12, 9.1, 3.4, 0.5])
    @test all(abs_vars .< [0.03, 0.02, 0.002, 3.5e-5])
end

@testset "Test Double Robustness ATE on binary_target_binary_treatment_pb" begin
    # When Q̅ is misspecified but G is well specified
    query = (t=[true, false],)
    Q̅ = ConstantClassifier()
    G = LogisticClassifier()
    tmle = TMLEstimator(Q̅, G, query)

    abs_mean_rel_errors, abs_vars = asymptotics(
            tmle,
            binary_target_binary_treatment_pb,
            StableRNG(123),
            Ns
            )
    @test all(abs_mean_rel_errors .< [9.2, 4.1, 1.6, 0.5])
    @test all(abs_vars .< [0.004, 0.001, 4.7e-5, 7.8e-6])

    # When Q̅ is well specified but G is misspecified
    query = (t=[true, false],)
    Q̅ = LogisticClassifier()
    G = ConstantClassifier()
    tmle = TMLEstimator(Q̅, G, query)

    abs_mean_rel_errors, abs_vars = asymptotics(
            tmle,
            binary_target_binary_treatment_pb,
            StableRNG(123),
            Ns
            )
    @test all(abs_mean_rel_errors .< [9.2, 4.1, 1.2, 0.4])
    @test all(abs_vars .< [0.004, 0.001, 4.9e-5, 8.2e-6])
end

@testset "Test Double Robustness ATE on continuous_target_binary_treatment_pb" begin
    # When Q̅ is misspecified but G is well specified
    query = (t=[true, false],)
    Q̅ = MLJ.DeterministicConstantRegressor()
    G = LogisticClassifier()
    tmle = TMLEstimator(Q̅, G, query)

    abs_mean_rel_errors, abs_vars = asymptotics(
            tmle,
            continuous_target_binary_treatment_pb,
            StableRNG(123),
            Ns
            )
    @test all(abs_mean_rel_errors .< [67, 8.1, 2.8, 1.2])
    @test all(abs_vars .< [0.4, 0.01, 0.002, 0.0002])

    # When Q̅ is well specified but G is misspecified
    query = (t=[true, false],)
    Q̅ = LinearRegressor()
    G = ConstantClassifier()
    tmle = TMLEstimator(Q̅, G, query)

    abs_mean_rel_errors, abs_vars = asymptotics(
            tmle,
            continuous_target_binary_treatment_pb,
            StableRNG(123),
            Ns
            )
    @test all(abs_mean_rel_errors .< [14.8, 5.4, 2.3, 0.9])
    @test all(abs_vars .< [0.04, 0.004, 0.0008, 0.0002])
end

@testset "Test multi-queries" begin
    queries = [(T=["AT", "TT"],), (T=["AA", "AT"],)]

    Q̅ = LinearRegressor()
    G = LogisticClassifier()
    tmle = TMLEstimator(Q̅, G, queries...)

    t, W, y, ATE₁ = continuous_target_categorical_treatment_pb(StableRNG(123);n=100, control="TT", treatment="AT")
    _, _, _, ATE₂ = continuous_target_categorical_treatment_pb(StableRNG(123);control="AT", treatment="AA")

    mach = machine(tmle, t, W, y)
    fit!(mach, verbosity=0)

    # Test queryreport accessibility
    @test TMLE.queryreportname(456) == :queryreport_456
    @test TMLE.getqueryreport(mach, 1) isa TMLE.QueryReport
    @test TMLE.getqueryreport(mach, 1) == mach.report.queryreport_1
    @test TMLE.getqueryreport(mach, 2) isa TMLE.QueryReport
    @test TMLE.getqueryreport(mach, 2) == mach.report.queryreport_2
    # Test various ztests
    qr₁ = getqueryreport(mach, 1)
    testresult = ztest(mach, 1)
    @test testresult == ztest(qr₁)
    @test testresult isa HypothesisTests.OneSampleZTest
    @test pvalue(testresult) ≈ 3.474e-25 atol=1e-3
    @test confint(testresult)[1] < ATE₁ < confint(testresult)[2]

    qr₂ = getqueryreport(mach, 2)
    testresult = ztest(mach, 2)
    @test testresult == ztest(qr₂)
    @test testresult isa HypothesisTests.OneSampleZTest
    @test pvalue(testresult) ≈ 3.428e-14 atol=1e-3
    @test confint(testresult)[1] < ATE₂ < confint(testresult)[2]

    testresult = ztest(mach, 1=>2)
    @test pvalue(testresult) ≈ 7.819e-25 atol=1e-3

    bf = briefreport(mach)
    for (i, _) in enumerate(mach.model.queries)
        @test briefreport(getqueryreport(mach, i)) == bf[i]
        @test keys(bf[i]) == (:query, :pvalue, :confint, :estimate, :initial_estimate, :stderror, :mean_inf_curve)
    end

end

end;

true

max_stars_repo_name: olivierlabayle/TMLE

max_issues_repo_name: olivierlabayle/TMLE

max_forks_repo_name: olivierlabayle/TMLE


text:
######
# This file is part of the BivariateCopulas.jl package.
#
#   Definition of bivariate copula type, Joint distribution type and methods
#
#           University of Liverpool, Institute for Risk and Uncertainty
#
#                                           Author: Ander Gray
#                                           Email:  ander.gray@liverpool.ac.uk
######

###
#   Notes: This is a module for performing tests with copulas for plotting, sampling and
#   making bivariate joint distributions
#
#   Would it be possible to represent a copula as an expansion of orthogonal
#   polynomials? Could be used as a sentivitivty analysis for dependancy. Something equivilant
#   to disperive Monte Carlo, where not only the correlatin coefficient is changed but the entire
#   dependancy structure. Could this expansion be used in the arithmetic?
#
#   Or can we do it as sums of the Frechet bounds? (Shuffles of M)
#   See: Nelson page 68.
#
#
###

###
#
#   Known bugs:
#                   ->  Differentiating with the interpolators is really ugly
#
###

#import PyPlot: PyPlot.plot

mutable struct copula <: AbstractCopula
    cdf::Array{Float64,2}
    density::Union{Array{Float64,2},Missing}
    func::Union{Function,Missing}
    param::Union{Float64,Missing}

    function copula(cdf=missing; density=missing, func=missing, param=missing)
        if (ismissing(cdf) && ismissing(density))
            throw(ArgumentError("Cdf and/or density must be provided"))
        end
        if (ismissing(cdf))
            cdf = calcCdf(density)
        end
        C = new(cdf, density, func, param)
        if (ismissing(density))
            density = calcDensity(C)
        end

        return new(cdf, density, func, param)
    end
end

function (obj::copula)(X, Y, useInterp=false)             # Evaluating the copula

    #if (!issorted(X) || !issorted(Y)) throw(ArgumentError("cdf evaluation request must be given in assending order"));end

    if !useInterp
        if !ismissing(obj.func)
            if !ismissing(obj.param)
                return obj.func(X, Y, obj.param)
            end
            return obj.func(X, Y)
        end
    end

    cdf = interpolate(obj.cdf, BSpline(Linear()))

    uu = range(0, 1; length=n)
    sCdf = Interpolations.scale(cdf, uu, uu)
    return sCdf(X, Y)
end

cdf(C::copula, X, Y) = C(X, Y)

function density(C::copula, X, Y)
    den = interpolate(C.density, BSpline(Quadratic(Line(OnCell()))))
    uu = range(0, 1; length=n)
    sDen = Interpolations.scale(den, uu, uu)
    return sDen(X, Y)
end

function (obj::copula)(
    X::ContinuousUnivariateDistribution, Y::ContinuousUnivariateDistribution
)             #Constructing a joint distribution using the copula
    return Joint(X, Y, obj)
end

function sample(C::AbstractCopula, N::Int64=1; plot=false, useInterp=false)

    # Set true to use interpolator
    useInterp = useInterp
    if ismissing(C.func)
        func = 0
        useInterp = true
    else
        func = C.func
    end

    if plot
        return samplePlot(C, N)
    end
    func == Gau && return CholeskyGaussian(N, C.param)   # Use Cholesky decompostition of the Cov matrix for gaussian copula sampling

    x = rand(N)
    y = rand(N)
    ux = x
    uy = zeros(N)
    m = n

    if func == perf
        return hcat(ux, ux)
    end
    if func == opp
        return hcat(ux, 1 .- ux)
    end
    if func == indep
        return hcat(ux, y)
    end

    # if (!ismissing(C.func) && n < 10^5) m = 10^5; end  # If function is provided more acurate sampling, otherwise must use interpolator
    e = sqrt(eps())
    ygrid = range(0; stop=1, length=m)

    #if (C.func == Gau) useInterp = true ;end           # Should you not want to use Cholesky and interpolator has better performance for gaussian

    if func !== 0
        conditional =
            (C(x .+ e / 2, ygrid, useInterp) - C(x .- e / 2, ygrid, useInterp)) ./ e
    else
    end
    conditional = (C(x .+ e / 2, ygrid, useInterp) - C(x .- e / 2, ygrid, useInterp)) ./ e
    for i in 1:N
        #conditional = (C(x[i] + e/2,ygrid, true) - C(x[i] - e/2,ygrid, true))./e
        #conditional[1] = 0; conditional[end]=1;       # Small errors sometimes happen
        #if (conditional[end] < y[i]) uy[i] = 1;      # Should linearly interpolate between cond[end] and 1
        #elseif (conditional[1] > y[i]) uy[i] = 0;      # Should linearly interpolate between cond[end] and 1
        #else uy[i] = findlast(conditional .<= y[i]) / m; end        # Finds Nothing??
        #yitp_linear = interpolate(conditional, BSpline(Linear()))
        #uy[i] = yitp_linear(y[i] * m);

        #uy[i] = findlast(conditional .<= y[i])/m; end
        thisCond = conditional[i, :]
        thisGrid = ygrid

        if !issorted(thisCond)
            thisCond = sort(thisCond)
        end

        idx = unique(z -> thisCond[z], 1:n)    # Some magic: https://discourse.julialang.org/t/unique-indices-method-similar-to-matlab/34446
        # Gets unique indicies from the vector

        thisCond = thisCond[idx]
        thisGrid = ygrid[idx]

        inpl = LinearInterpolation(thisCond, thisGrid; extrapolation_bc=Flat())
        uy[i] = inpl(y[i])
    end
    return hcat(ux, uy)
end

rand(C::AbstractCopula, N::Int64=1; plot=false) = sample(C, N; plot=plot)

function CholeskyGaussian(N=1, correlation=0)
    Cov = ones(2, 2)
    Cov[1, 2] = Cov[2, 1] = correlation
    a = cholesky(Cov).L
    z = transpose(rand(Normal(), N, 2))
    x = a * z
    u = transpose(cdf.(Normal(), x))

    return hcat(u[:, 1], u[:, 2])
end

function conditional(C::AbstractCopula, xVal::Real; plot=false)   # May also work for joint? xVal will be take from invCdf(M1, u1)
    e = 1 / n
    ygrid = range(0; stop=1, length=n)
    if !ismissing(C.func)
        e = (sqrt(eps()))
    end

    conditional = (C(xVal + e / 2, ygrid) - C(xVal - e / 2, ygrid)) ./ e
    if (plot)
        plot2(C.cdf, conditional, xVal)
    end
    return conditional
end

function calcCdf(density::Array{<:Real,2})
    cdf = zeros(size(density))
    for i in 2:n
        for j in 2:n
            cdf[i, j] = (
                density[i, j] / (n^2) + cdf[i, j - 1] + cdf[i - 1, j] - cdf[i - 1, j - 1]
            )
        end
    end
    return cdf
end

function calcDensity(C::AbstractCopula)
    if ismissing(C.func)
        return calcDensity(C.cdf)
    end

    e = sqrt(sqrt(eps()))

    x = range(0; stop=1 - e, length=n)
    y = range(0; stop=1 - e, length=n)

    der1 = (C(x .+ e, y) - C(x, y)) / e      # df/dx
    der2 = (C(x .+ e, y .+ e) - C(x, y .+ e)) / e     # d(df/dx)/dy
    density = (der2 - der1) / e

    return density
end

function calcDensity(
    cdf::Array{<:Real,2}, xRange=range(0, 1; length=n), yRange=range(0, 1; length=n)
)
    xStep = Float64(xRange.step)
    yStep = Float64(yRange.step)

    density = zeros(size(cdf))
    for i in 1:(n - 1)
        for j in 1:(n - 1)
            density[i, j] =
                (cdf[i, j] + cdf[i + 1, j + 1] - cdf[i + 1, j] - cdf[i, j + 1]) /
                (xStep * yStep)

            #   -> Possible alternative:
            #der1 = ( cdf[i + 1, j]      - cdf[i, j]     )/(xStep^2)             # df/dx
            #der2 = ( cdf[i + 1, j + 1 ] - cdf[i, j + 1] )/(yStep^2)             # d(df/dx)/dy
            #density[i,j] = (der2 - der1)
        end
    end

    # Problem: How do you get the end values?
    #   Two solutions:

    #   -> If copula, it's symetric (exact answer exept at two endpoints)
    #density[end,:] = reverse(density[:,1]);
    #density[:,end] = reverse(density[1,:]);

    #   -> If joint, it may not be symetric
    density[end, :] = density[end - 1, :]
    density[:, end] = density[:, end - 1]
    return density
end

function M()
    x = y = range(0; stop=1, length=n)
    return copula(perf(x, y); density=Matrix{Float64}(I, n, n), func=perf)
end

function W()
    x = y = range(0; stop=1, length=n)

    den = zeros(n, n)
    for i in 1:n
        for j in 1:n
            if (i - 1) == n - (j - 1)
                den[i, j] = 1
            end
        end
    end
    return copula(opp(x, y); density=den, func=opp)
end

function Pi()
    x = y = range(0; stop=1, length=n)
    return copula(indep(x, y); density=ones(n, n), func=indep)
end

function Frank(s=1)                                     #   Frank copula
    x = y = range(0; stop=1, length=n)                          #   s>0; 0 for perfect, 1 for indep, inf for oposite
    return copula(F(x, y, s); func=F, param=s)
end

function Gaussian(corr=0)
    x = y = range(0; stop=1, length=n)
    cdf = Gau(x, y, corr)
    #cdf[:,end] = cdf[end,:] = x;    # This removes the NaNs from Liams mvNormCdf. We know the marginals must be uniform
    return copula(cdf; func=Gau, param=corr)
end

function τCopula(τ=0)   # Imprecise copula from kendal tau
    x = y = range(0; stop=1, length=n)

    cdf = kenLB(x, y, τ)

    return copula(cdf; func=kenLB, param=τ)
end

KendalCopula(τ=0) = τCopula(τ)

function ρCopula(ρ=0) # Imprecise copula from Spearman rho
    x = y = range(0; stop=1, length=n)

    cdf = spearLB(x, y, ρ)

    return copula(cdf; func=spearLB, param=ρ)
end

SpearmanCopula(ρ=0) = ρCopula(ρ)

indep(X, Y) = [x * y for x in X, y in Y];
perf(X, Y) = [min(x, y) for x in X, y in Y];
opp(X, Y) = [max(x + y - 1, 0) for x in X, y in Y];
F(X, Y, s=1) = [log(1 + (s^x - 1) * (s^y - 1) / (s - 1)) / log(s) for x in X, y in Y]
function Gau(X, Y, corr=0)
    return [
        bivariate_cdf(quantile.(Normal(), x), quantile.(Normal(), y), corr) for x in X,
        y in Y
    ]
end;

#kenLB(x, y, τ) = max(0, x+y-1, 0.5*( (x+y)-sqrt( (x-y)^2)+1-τ ) )
kenUB(x, y, τ) = min(x, y, 0.5 * ((x + y - 1) + sqrt((x + y - 1)^2 + 1 + τ)))
kenLB(x, y, τ) = [x - kenUB(x, 1 - y, -τ) for x in x, y in y]

#spearϕ(a, b) = 1/6 * ( (max(0, 9*b + 3*sqrt( 9*b^2 - 3*a^6)))^(1/3) + ( max(0,9*b - 3*sqrt(9*b^2 - 3*a^6)))^(1/3) )

function spearϕ(a, b)
    A = 9 * b
    B = max(9 * b^2 - 3 * a^6, 0)
    C = (max(0, A + 3 * sqrt(B)))^(1 / 3)
    D = (max(0, A - 3 * sqrt(B)))^(1 / 3)
    return 1 / 6 * (C + D)
end

#spearLB(x, y, ρ) = max(0, x + y - 1, (x + y)/2 - spearϕ(x - y, 1 - ρ))
spearUB(x, y, ρ) = max(opp(x, y)[1], min(x, y, (x + y - 1) / 2 + spearϕ(x + y - 1, 1 + ρ)))
spearLB(x, y, τ) = [x - spearUB(x, 1 - y, -τ) for x in x, y in y]

function GauCopulaDen(corr=0, X=range(bOt, tOp; length=n), Y=range(bOt, tOp; length=n))
    COR = [1.0 corr; corr 1.0]
    I = [1.0 0; 0 1.0]
    den = zeros(length(X), length(Y))
    for i in 1:length(X)
        for j in 1:length(Y)
            Vec = [invCdf(Normal(), X[i]); invCdf(Normal(), Y[j])]
            den[i, j] = 1 / sqrt(det(COR)) * exp(-1 / 2 * Vec' * (inv(COR) - I) * Vec)
        end
    end
    return den
end

function Arch1(X, Y, theta=0)
    return [
        1 - ((1 - x)^theta + (1 - y)^theta - (1 - x)^theta * (1 - y)^theta)^theta for
        x in X, y in Y
    ]
end;

mvNormCdf(X, Y, cor) = [bivariate_cdf(x, y, cor) for x in X, y in Y];

###
#   Copula rotations
###
function rotate90(C::copula)
    C_func = C.func
    C_param = C.param
    if !ismissing(C_func)
        if !ismissing(C_param)
            function Rotate90(xs, ys, C_param)
                return [y - C_func(1 - x, y, C_param)[1] for x in xs, y in ys]
            end
            u = v = range(0; stop=1, length=n)
            cdf = Rotate90(u, v, C_param)
            cdf[1, :] .= 0
            cdf[:, 1] .= 0
            return copula(cdf; func=Rotate90, param=C_param)
        else
            Rotate90(xs, ys) = [y - C_func(1 - x, y)[1] for x in xs, y in ys]
            u = v = range(0; stop=1, length=n)
            cdf = Rotate90(u, v)
            return copula(cdf; func=Rotate90)
        end
    end

    ###
    #   Here we need to figure out how to rotate the cdf matrix
    ###

    throw(ArgumentError("Currently can only rotate copulas with functions defined"))
end

function rotate180(C::copula)
    C_func = C.func
    C_param = C.param
    if !ismissing(C_func)
        if !ismissing(C_param)
            function Rotate180(xs, ys, C_param)
                return [x + y - 1 + C_func(1 - x, 1 - y, C_param)[1] for x in xs, y in ys]
            end
            u = v = range(0; stop=1, length=n)
            cdf = Rotate180(u, v, C_param)
            cdf[1, :] .= 0
            cdf[:, 1] .= 0
            return copula(cdf; func=Rotate180, param=C_param)
        else
            Rotate180(xs, ys) = [x + y - 1 + C_func(1 - x, 1 - y)[1] for x in xs, y in ys]
            u = v = range(0; stop=1, length=n)
            cdf = Rotate180(u, v)
            return copula(cdf; func=Rotate180)
        end
    end

    ###
    #   Here we need to figure out how to rotate the cdf matrix
    ###

    throw(ArgumentError("Currently can only rotate copulas with functions defined"))
end

function rotate270(C::copula)
    C_func = C.func
    C_param = C.param
    if !ismissing(C_func)
        if !ismissing(C_param)
            function Rotate270(xs, ys, C_param)
                return [x - C_func(x, 1 - y, C_param)[1] for x in xs, y in ys]
            end
            u = v = range(0; stop=1, length=n)
            cdf = Rotate270(u, v, C_param)
            cdf[1, :] .= 0
            cdf[:, 1] .= 0
            return copula(cdf; func=Rotate270, param=C_param)
        else
            Rotate270(xs, ys) = [x - C_func(x, 1 - y)[1] for x in xs, y in ys]
            u = v = range(0; stop=1, length=n)
            cdf = Rotate270(u, v)
            return copula(cdf; func=Rotate270)
        end
    end

    ###
    #   Here we need to figure out how to rotate the cdf matrix
    ###

    throw(ArgumentError("Currently can only rotate copulas with functions defined"))
end

mutable struct Joint <: AbstractJoint
    marginal1::ContinuousUnivariateDistribution
    marginal2::ContinuousUnivariateDistribution
    copula::AbstractCopula

    xRange
    yRange

    cdf::Union{Array{Float64,2},Missing}     # Or should the density/ cdf be calculated on the fly from the copula?
    density::Union{Array{Float64,2},Missing}

    function Joint(marginal1, marginal2, copula)
        if isa(copula, ArchimedeanCopula)
            return new(marginal1, marginal2, copula, missing, missing, missing, missing)
        end

        xRange = range(invCdf(marginal1, bOt); stop=invCdf(marginal1, tOp), length=n)
        yRange = range(invCdf(marginal2, bOt); stop=invCdf(marginal2, tOp), length=n)

        j = new(marginal1, marginal2, copula, xRange, yRange, missing)
        cdf = j(xRange, yRange)
        j = new(marginal1, marginal2, copula, xRange, yRange, cdf)
        density = calcDensity(j)                                    # Looks really ugly when using the interpolator for a marginal

        return new(marginal1, marginal2, copula, xRange, yRange, cdf, density)
    end
end

function (obj::Joint)(X, Y)    # Evaluating the joint cdf
    # BiLinear interpolator if cdf is provided

    if isa(obj.copula, ArchimedeanCopula)
        return cdf.(obj.copula, cdf.(obj.marginal1, X), cdf.(obj.marginal2, Y))
    end

    if !ismissing(obj.copula.func)
        return obj.copula(cdf.(obj.marginal1, X), cdf.(obj.marginal2, Y))  # This is the case where we have the functions
    end
    if ismissing(obj.cdf)
        return obj.copula(cdf.(obj.marginal1, X), cdf.(obj.marginal2, Y))
    end

    cdf1 = interpolate(obj.cdf, BSpline(Quadratic(Line(OnCell()))))
    sCdf1 = Interpolations.scale(cdf1, obj.xRange, obj.yRange)

    return sCdf1(X, Y)

    # Should use new scaling from interpolatorrs
end

cdf(J::Joint, X, Y) = J(X, Y)

function conditional(J::Joint, xVal::Real; plot=false)   # May also work for joint? xVal will be take from invCdf(M1, u1)
    e = 1 / n
    ygrid = J.yRange
    if !ismissing(J.copula.func)
        e = (sqrt(eps()))
    end

    conditional = (J(xVal + e / 2, ygrid) - J(xVal - e / 2, ygrid)) ./ e

    conditional = conditional / pdf(J.marginal1, xVal)

    if (plot)
        plot2(J.cdf, conditional, xVal)
    end
    return conditional
end

function density2(J::Joint, X=J.xRange, Y=J.yRange)
    den = interpolate(J.density, BSpline(Quadratic(Line(OnCell()))))
    sDen = Interpolations.scale(den, X, Y)
    return sDen(X, Y)
end

function density(J::Joint, X=J.xRange, Y=J.yRange)        # An alternative way to construct the Joint. fxy / fx*fy = c[Fx,Fy]. Where c is the copula Density
    # More error found against Multivariate gaussian pdf. Due to interpolation error in the copula density.
    denCop = density(J.copula, cdf.(J.marginal1, X), cdf.(J.marginal2, Y))
    return denCop .* [pdf(J.marginal1, x) * pdf(J.marginal2, y) for x in X, y in Y]
end

function sample(J::AbstractJoint, N::Int64=1)
    copulaSamples = sample(J.copula, N)

    @show size(copulaSamples)
    x = invCdf.(J.marginal1, copulaSamples[:, 1])
    y = invCdf.(J.marginal2, copulaSamples[:, 2])
    return hcat(x, y)
end

rand(J::AbstractJoint, N::Int64=1) = sample(J::AbstractJoint, N)

function invCdf(X::Sampleable{Univariate}, u)
    return quantile.(X, u)                           # We will also need the interpolartor here
end

function calcDensity(J::AbstractJoint)

    #return calcDensity(J.cdf, J.xRange, J.yRange)

    if ismissing(J.copula.func)
        return calcDensity(J.cdf, J.xRange, J.yRange)
    end

    e = sqrt(sqrt(eps()))

    #x = J.xRange - e;f
    #y = J.yRange - e;

    x = range(J.xRange[1]; stop=J.xRange[end] - e, length=length(J.xRange))
    y = range(J.xRange[1]; stop=J.yRange[end] - e, length=length(J.yRange))

    der1 = (J(x .+ e, y) - J(x, y)) / e     # df/dx
    der2 = (J(x .+ e, y .+ e) - J(x, y .+ e)) / e     # d(df/dx)/dy
    density = (der2 - der1) / e                            #d^2f/dxdy

    return density
end

M(M1::ContinuousUnivariateDistribution, M2::ContinuousUnivariateDistribution) = M()(M1, M2)
W(M1::ContinuousUnivariateDistribution, M2::ContinuousUnivariateDistribution) = W()(M1, M2)
function Pi(M1::ContinuousUnivariateDistribution, M2::ContinuousUnivariateDistribution)
    return Pi()(M1, M2)
end
function Frank(
    M1::ContinuousUnivariateDistribution, M2::ContinuousUnivariateDistribution, s=1
)
    return Frank(s)(M1, M2)
end
function Gaussian(
    M1::ContinuousUnivariateDistribution, M2::ContinuousUnivariateDistribution, corr=0
)
    return Gaussian(corr)(M1, M2)
end

struct marginal <: AbstractMarginal
    cdf::Array{Float64,1}
    range
    cdfInp
    invCdfInp

    # Could have density and a calcDensity function for pdf

    function marginal(cdf::Array{Float64,1}, range)
        if (length(cdf) !== length(range))
            throw(ArgumentError("Cdf and range must be the same length"))
        end

        inpl = interpolate(cdf, BSpline(Quadratic(Line(OnCell()))))
        cdfInp = Interpolations.scale(inpl, range)

        # Error here and also in differention
        invCdfInp = LinearInterpolation(cdf, range; extrapolation_bc=Flat())

        return new(cdf, range, cdfInp, invCdfInp)
    end
end

function (obj::marginal)(x)
    return cdf(obj, x)
end

function cdf(D::AbstractMarginal, x::Real)
    if x < D.range[1]
        return 0
    end
    if x > D.range[end]
        return 1
    end

    return D.cdfInp(x)
end

function invCdf(D::AbstractMarginal, u::Real)
    return D.invCdfInp(u)
end

quantile(D::marginal, u) = invCdf(D, u);

function sample(D::marginal, N=1)
    return invCdf(D, rand(N))
end

rand(D::marginal, N=1) = sample(D, N)

function Base.show(io::IO, z::copula)
    statement1 = "Arbitrary"
    statement2 = ""

    if (!ismissing(z.func))
        func = z.func
        if (func == indep)
            func = "π"
        end
        if (func == perf)
            func = "M"
        end
        if (func == opp)
            func = "W"
        end
        statement1 = "$(func)"
    end

    if (!ismissing(z.param))
        func = z.func
        parName = "par"
        if (func == Gau)
            parName = "r"
        end
        if (func == F)
            parName = "s"
        end
        statement2 = "$parName=$(z.param)"
    end

    return print(io, "Copula ~ $statement1($statement2)")
end

function Base.show(io::IO, z::Joint)
    statement1 = "Arbitrary"
    statement2 = ""

    if isa(z.copula, ArchimedeanCopula)
        statement1 = string(typeof(z.copula))
    elseif (!ismissing(z.copula.func))
        func = z.copula.func
        if (func == indep)
            func = "π"
        end
        if (func == perf)
            func = "M"
        end
        if (func == opp)
            func = "W"
        end
        statement1 = "$(func)"
    end

    if isa(z.copula, ArchimedeanCopula)
        statement2 = "ϑ=$(z.copula.ϑ)"
    elseif (!ismissing(z.copula.param))
        func = z.copula.func
        parName = "par"
        if (func == Gau)
            parName = "r"
        end
        if (func == F)
            parName = "s"
        end
        if (func == spearLB)
            parName = "ρ"
        end
        if (func == kenLB)
            parName = "τ"
        end
        statement2 = "$parName=$(z.copula.param)"
    end

    return print(io, "Joint ~ $statement1($statement2, $(z.marginal1), $(z.marginal2))")
end

function Base.show(io::IO, z::marginal)
    return print(io, "Distribution ~ ( range=[$(z.range[1]),$(z.range[end])])")
end

function linInterp(A, x)
    ind0 = findlast(A .<= x)
    x0 = A[ind0]
    x1 = A[ind0 + 1]
    y0 = ind0 / n
    y1 = (ind0 + 1) / n

    return y0 + (x - x0) * (y1 - y0) / (x1 - x0)
end

ext: jl


text:
module Dataset

import ..Flux3D
import SHA

export CustomDataset, ModelNet10, ModelNet40

# Utilities
include("utils.jl")

# Custom Dataset
include("custom.jl")

# ModelNet
include("modelnet/base.jl")
include("modelnet/mn10.jl")
include("modelnet/mn40.jl")

end # module


text:
using Distributions

function fisher_exact_test(contingency_table)
    @assert size(contingency_table) ≡ (2, 2)
    w, y, x, z = contingency_table
    n = w + y + x + z
    p = binomial(w + x, w) * binomial(y + z, y) / binomial(n, w + y)
end

function χ²_test(contingency_table)
    @assert size(contingency_table) ≡ (2, 2)
    w, y, x, z = contingency_table
    n = w + y + x + z
    fA = (w + x) / n
    fB = (w + y) / n

    expected = [
        fA * fB        fA * (1 - fB)
        (1 - fA) * fB  (1 - fA) * (1 - fB)
    ] * n
    χ² = sum(contingency_table .^ 2 ./ expected) - n
    ccdf(Chisq(1), χ²)
end


text:
## Main translation unit
# Need to include Precompile.jl before running this unit
using Distributed
using DistributedArrays

## Some important parameters
dict = Dict{String,Float64}("U" => 4.0, "V" => 1.0)
beta = 200; Niωn = 50
N_it = 3 ## Lowest number is 1: one loop in the process
SubLast = 1 ## Subdivision of last integral (N_it) to be split in #Sublast to be fed to different cores
Full = false ## If you want to compute the results of all the iterations, set to true. Set to false otherwise!
Grid_K = 50
Dims = 1

@everywhere include("Precompile.jl")

if Full
    SuperHF.Hubbard.@assert (N_it*SubLast) <= Base.Sys.CPU_THREADS # Maximum number of cores
    addprocs(N_it*SubLast) ## defines the number of workers (cores)
else
    SuperHF.Hubbard.@assert (SubLast) <= Base.Sys.CPU_THREADS # Maximum number of cores
    addprocs(SubLast)
end

function Gk_conv(vals::Union{SuperHF.Green_Utils.Integral,SuperHF.Green_Utils.Integral1D})
    if isa(vals, SuperHF.Green_Utils.Integral)
        @time @sync for (idx,pid) in enumerate(workers())
            @async c_container[idx] = remotecall_fetch(Dict_funct[N_it][idx], pid, vals) ## To test
        end
        summation_subs = sum(c_container)
        println("Summation: ", summation_subs)

        Gk = 1.0/(vals.iωn - SuperHF.Green_Utils.epsilonk(vals.kx, vals.ky) - summation_subs)

        return Gk
    elseif isa(vals, SuperHF.Green_Utils.Integral1D)
        println("In gk_conv 1D")
        @time @sync for (idx,pid) in enumerate(workers())
            @async c_container[idx] = remotecall_fetch(Dict_funct[N_it][idx], pid, vals) ## To test
        end
        println("Elements: ", c_container)
        summation_subs = sum(c_container)
        println("Summation: ", summation_subs)

        Gk = 1.0/(vals.iωn - SuperHF.Green_Utils.epsilonk1D(vals.kx) - summation_subs)

        return Gk
    end
end

function Lambda(q::Array{Float64,1}, qp::Array{Float64,1}, qpp::Array{Float64,1}, k::Array{Float64,1}, kp::Array{Float64,1}, iωn::Complex{Float64}, HF::SuperHF.Green_Utils.HFVec, Gk1::SuperHF.Green_Utils.Integral, Gk2::SuperHF.Green_Utils.Integral)
    println("IN LAMBDA FUNCTION", "\n\n")

    kernel = inv( inv(SuperHF.Green_Utils.Potential(q[1]-qp[1], q[2]-qp[2], values(HF.dict_v_)...)) +
                inv(SuperHF.Green_Utils.Potential(qpp[1], qpp[2], values(HF.dict_v_)...))*SuperHF.Green_Utils.Potential(q[1]-qp[1]-qpp[1], q[2]-qp[2]-qpp[2], values(HF.dict_v_)...)*Gk_conv(Gk1)*Gk_conv(Gk2) )

    println("kernel value: ", kernel, "\n")
    return kernel
end

function Lambda1D(q::Float64, qp::Float64, k::Float64, kp::Float64, iωn::Complex{Float64}, HF::SuperHF.Green_Utils.HFVec, Gk1::SuperHF.Green_Utils.Integral1D, Gk2::SuperHF.Green_Utils.Integral1D)
    println("IN LAMBDA FUNCTION 1D", "\n\n")

    kernel = inv( inv(SuperHF.Green_Utils.Potential1D(q, values(HF.dict_v_)...)) +
                inv(SuperHF.Green_Utils.Potential1D(qp, values(HF.dict_v_)...))*SuperHF.Green_Utils.Potential1D(q-qp, values(HF.dict_v_)...)*Gk_conv(Gk1)*Gk_conv(Gk2) )

    println("kernel value 1D: ", kernel, "\n")
    return kernel
end

function Susceptibility(q::Array{Float64,1}, qp::Array{Float64,1}, qpp::Array{Float64,1}, k::Array{Float64,1}, kp::Array{Float64,1}, iωn::Complex{Float64}, HF::SuperHF.Green_Utils.HFVec, Gk1::SuperHF.Green_Utils.Integral, Gk2::SuperHF.Green_Utils.Integral, Gks::Array{SuperHF.Green_Utils.Integral,1})
    println("IN SUSCEPTIBILITY FUNCTION", "\n\n")
    sus = Gk_conv(Gks[1])*SuperHF.Green_Utils.Potential(qp[1], qp[2], values(HF.dict_v_)...)*Gk_conv(Gks[2])*Lambda(q,qp,qpp,k,kp,iωn,HF,Gk1,Gk2)*Gk_conv(Gks[3])*Gk_conv(Gks[4])*Gk_conv(Gks[5])*Gk_conv(Gks[6])

    println("Susceptibility: ", sus, "\n\n")
    return sus
end

function Susceptibility1D(q::Float64, qp::Float64, k::Float64, kp::Float64, iωn::Complex{Float64}, HF::SuperHF.Green_Utils.HFVec, Gk1::SuperHF.Green_Utils.Integral1D, Gk2::SuperHF.Green_Utils.Integral1D, Gks::Array{SuperHF.Green_Utils.Integral1D,1})
    println("IN SUSCEPTIBILITY FUNCTION 1D", "\n\n")
    sus = Gk_conv(Gks[1])*SuperHF.Green_Utils.Potential1D(qp, values(HF.dict_v_)...)*Lambda1D(q,qp,k,kp,iωn,HF,Gk1,Gk2)*Gk_conv(Gks[2])*Gk_conv(Gks[3])*Gk_conv(Gks[4])

    println("Susceptibility 1D: ", sus, "\n\n")
    return sus
end

## Initializing the struct members ##
#####################################
model = SuperHF.Green_Utils.HFVec(Niωn,beta,N_it,dict)
Boundaries = Array{Array{Float64,1},1}([[-pi,-pi],[pi,pi]])
qp_array = Array{Array{Float64,1},1}([[a,b] for a in range(-pi,stop=pi,length=Grid_K) for b in range(-pi,stop=pi,length=Grid_K)])
qpp_array = qp_array; k_array = qp_array; kp_array = qp_array
Boundaries1D = Array{Float64,1}([-pi,pi])
qp_array1D = Array{Float64,1}([a for a in range(-pi,stop=pi,length=Grid_K)])
qpp_array1D = qp_array1D; k_array1D = qp_array1D; kp_array1D = qp_array1D

intermediate_integral = SuperHF.Green_Utils.integrate_complex(SuperHF.Green_Utils.Init_Gk, x -> x/2 , 1, model, Boundaries)
println(intermediate_integral(SuperHF.Green_Utils.Integral(pi/4,pi/3,0.01im)))

if Dims == 2
    Dict_funct = SuperHF.Green_Utils.interation_process(model,Boundaries,SubLast)
    if Full
        println("Length of function array: ", length(Dict_funct)*length(Dict_funct[1]), " and number of workers: ", nworkers())
        @assert isa(Dict_funct,Dict{Int64,Array{Function,1}}) && (length(Dict_funct)*length(Dict_funct[1]) == nworkers())
        c_container = Vector{Complex{Float64}}(undef,N_it*SubLast)

        function main()
            ## Part where the convergence (behavior) is checked ##
            ######################################################
            try
                for (ii,iωn) in enumerate(model.Matsubara_G_)
                    if ii == 1
                        println("iwn: ", iωn)
                        val = SuperHF.Green_Utils.Integral(pi, pi/2, iωn)
                        @time @sync for (idx,pid) in enumerate(workers())
                            @async c_container[idx] = remotecall_fetch(Dict_funct[ceil(Int64,idx/N_it)][ceil(Int64,idx/SubLast)], pid, val) ## To test
                            #mapping = pmap((f,x)->f(x), Funct_arr, vals)
                        end
                        summation_subs = sum(c_container[end-SubLast:end])
                        println("Summation last terms: ", summation_subs)
                    end
                end
            catch err
                if typeof(err) == InterruptException
                    println("ALL THE TASKS HAVE BEEN INTERRUPTED","\n")
                    for pid in workers()
                        interrupt(pid)
                    end
                end
            end
            println("Program terminated. Have a nice day!")
            rmprocs(workers())
            return nothing
        end
    else
        c_container = Vector{Complex{Float64}}(undef,SubLast)
        println("Length of function array: ", length(Dict_funct[N_it]), " and number of workers: ", nworkers())
        @assert isa(Dict_funct,Dict{Int64,Array{Function,1}}) && (length(Dict_funct[N_it]) == nworkers())
        function main()
            try
                Matsubara_array_susceptibility = Array{Complex{Float64},1}()
                q = [0.,0.]
                for (ii,iωn) in enumerate(model.Matsubara_G_)
                    k_sum = 0.0+0.0im
                    println("iwn: ", iωn)
                    for qp in qp_array
                        println("In qp: ", qp[1], " ", qp[2])
                        for qpp in qpp_array
                            println("In qpp: ", qpp[1], " ", qpp[2])
                            for k in k_array
                                println("In k: ", k[1], " ", k[2])
                                for kp in kp_array
                                    println("In kp: ", kp[1], " ", kp[2])
                                    Gk1 = SuperHF.Green_Utils.Integral(k[1], k[2], iωn); Gk2 = SuperHF.Green_Utils.Integral(kp[1]+qpp[1]-q[1], kp[2]+qpp[2]-q[2], iωn)
                                    Gks1 = SuperHF.Green_Utils.Integral(k[1]-qpp[1]+q[1], k[2]-qpp[2]+q[2], iωn); Gks2 = SuperHF.Green_Utils.Integral(k[1]-qpp[1]-qp[1]+q[1], k[2]-qpp[2]-qp[2]+q[2], iωn)
                                    Gks3 = SuperHF.Green_Utils.Integral(kp[1]-qp[1], kp[2]-qp[2], iωn); Gks4 = SuperHF.Green_Utils.Integral(kp[1], kp[2], iωn)
                                    Gks5 = SuperHF.Green_Utils.Integral(kp[1]-q[1], kp[2]-q[2], iωn); Gks6 = SuperHF.Green_Utils.Integral(k[1]-qpp[1], k[2]-qpp[2], iωn)
                                    Matsubara_sus = Susceptibility(q, qp, qpp, k, kp, iωn, model, Gk1, Gk2, [Gks1,Gks2,Gks3,Gks4,Gks5,Gks6])
                                    k_sum += Matsubara_sus
                                end
                            end
                        end
                        push!(Matsubara_array_susceptibility,2.0*(1.0/(Grid_K^2))^4*k_sum)
                    end
                end
                tot_susceptibility = (1/model.beta_)^4*sum(Matsubara_array_susceptibility)
                return tot_susceptibility
            catch err
                if typeof(err) == InterruptException
                    println("ALL THE TASKS HAVE BEEN INTERRUPTED","\n")
                    for pid in workers()
                        interrupt(pid)
                    end
                end
            end
            println("Program terminated. Have a nice day!")
            rmprocs(workers())
            return nothing
        end
    end
elseif Dims == 1
    Dict_funct = SuperHF.Green_Utils.interation_process1D(model,Boundaries1D,SubLast,Grid_K)
    c_container = Vector{Complex{Float64}}(undef,SubLast)
    println("Length of function array: ", length(Dict_funct[N_it]), " and number of workers: ", nworkers())
    @assert isa(Dict_funct,Dict{Int64,Array{Function,1}}) && (length(Dict_funct[N_it]) == nworkers())
    function main()
        try
            Matsubara_array_susceptibility = Array{Complex{Float64},1}()
            q = 0.
            for (ii,iωn) in enumerate(model.Matsubara_G_)
                f = open("1D_HF_Susceptibility_calc_first.dat", "a")
                if ii == 1
                    write(f, "#N_it "*"$(N_it)"*" q="*"$(q)"*"\n")
                end
                k_sum = 0.0+0.0im
                println("iwn: ", iωn)
                for qp in qp_array1D
                    println("In qp: ", qp)
                    for k in k_array1D
                        println("In k: ", k)
                        for kp in kp_array1D
                            println("In kp: ", kp)
                            Gk1 = SuperHF.Green_Utils.Integral1D(k, iωn); Gk2 = SuperHF.Green_Utils.Integral1D(kp+qp, iωn)
                            Gks1 = SuperHF.Green_Utils.Integral1D(k-qp+q, iωn); Gks2 = SuperHF.Green_Utils.Integral1D(kp+q, iωn)
                            Gks3 = SuperHF.Green_Utils.Integral1D(kp, iωn); Gks4 = SuperHF.Green_Utils.Integral1D(k-qp, iωn)
                            Matsubara_sus = Susceptibility1D(q, qp, k, kp, iωn, model, Gk1, Gk2, [Gks1,Gks2,Gks3,Gks4])
                            k_sum += Matsubara_sus
                        end
                    end
                end
                k_sum = 2.0*(1.0/(Grid_K))^3*k_sum
                push!(Matsubara_array_susceptibility,k_sum)
                write(f, "$(iωn)"*"\t\t"*"$(k_sum)")
                close(f)
            end
            tot_susceptibility = (1/model.beta_)^3*sum(Matsubara_array_susceptibility)
            return tot_susceptibility
        catch err
            if typeof(err) == InterruptException
                println("ALL THE TASKS HAVE BEEN INTERRUPTED","\n")
                for pid in workers()
                    interrupt(pid)
                end
            end
        end
        println("Program terminated. Have a nice day!")
        rmprocs(workers())
        return nothing
    end
end

main()

max_stars_repo_name: oliviersimard/Hartree-Fock

max_issues_repo_name: oliviersimard/Hartree-Fock

max_forks_repo_name: oliviersimard/Hartree-Fock


text:
# This file is a part of Julia. License is MIT: https://julialang.org/license

## initdefs.jl - initialization and runtime management definitions

"""
    PROGRAM_FILE

A string containing the script name passed to Julia from the command line. Note that the
script name remains unchanged from within included files. Alternatively see
[`@__FILE__`](@ref).
"""
global PROGRAM_FILE = ""

"""
    ARGS

An array of the command line arguments passed to Julia, as strings.
"""
const ARGS = String[]

"""
    exit(code=0)

Quit the program with an exit code. The default exit code is zero, indicating that the
program completed successfully (see also [`quit`](@ref)). In an interactive session,
`exit()` can be called with the keyboard shorcut `^D`.

"""
exit(n) = ccall(:jl_exit, Void, (Int32,), n)
exit() = exit(0)

"""
    quit()

Quit the program indicating successful completion. This function is equivalent to
`exit(0)` (see [`exit`](@ref)). In an interactive session, `quit()` can be called
with the keyboard shorcut `^D`.
"""
quit() = exit()

const roottask = current_task()

is_interactive = false

"""
    isinteractive() -> Bool

Determine whether Julia is running an interactive session.
"""
isinteractive() = (is_interactive::Bool)

"""
    LOAD_PATH

An array of paths as strings or custom loader objects for the `require`
function and `using` and `import` statements to consider when loading
code.
"""
const LOAD_PATH = String[]
const LOAD_CACHE_PATH = String[]

function init_load_path(JULIA_HOME = JULIA_HOME)
    vers = "v$(VERSION.major).$(VERSION.minor)"
    if haskey(ENV, "JULIA_LOAD_PATH")
        prepend!(LOAD_PATH, split(ENV["JULIA_LOAD_PATH"], @static Sys.iswindows() ? ';' : ':'))
    end
    push!(LOAD_PATH, abspath(JULIA_HOME, "..", "local", "share", "julia", "site", vers))
    push!(LOAD_PATH, abspath(JULIA_HOME, "..", "share", "julia", "site", vers))
    #push!(LOAD_CACHE_PATH, abspath(JULIA_HOME, "..", "lib", "julia")) #TODO: add a builtin location?
end

function early_init()
    global const JULIA_HOME = ccall(:jl_get_julia_home, Any, ())
    # make sure OpenBLAS does not set CPU affinity (#1070, #9639)
    ENV["OPENBLAS_MAIN_FREE"] = get(ENV, "OPENBLAS_MAIN_FREE",
                                    get(ENV, "GOTOBLAS_MAIN_FREE", "1"))
    if Sys.CPU_CORES > 8 && !("OPENBLAS_NUM_THREADS" in keys(ENV)) && !("OMP_NUM_THREADS" in keys(ENV))
        # Prevent openblas from starting too many threads, unless/until specifically requested
        ENV["OPENBLAS_NUM_THREADS"] = 8
    end
end

"""
    JULIA_HOME

A string containing the full path to the directory containing the `julia` executable.
"""
:JULIA_HOME

const atexit_hooks = []

"""
    atexit(f)

Register a zero-argument function `f()` to be called at process exit. `atexit()` hooks are
called in last in first out (LIFO) order and run before object finalizers.
"""
atexit(f::Function) = (unshift!(atexit_hooks, f); nothing)

function _atexit()
    for f in atexit_hooks
        try
            f()
        catch err
            show(STDERR, err)
            println(STDERR)
        end
    end
end

max_stars_repo_name: xorJane/julia

max_issues_repo_name: JuliaLangEs/julieta

max_forks_repo_name: JuliaLangEs/julieta


text:
using Dates
using DataStructures
using Parameters
using Printf

using .Networks

###############################################
# Sequence
###############################################

"""
    Sequence

Lists, for each horizon timepoint, the ordered operations to execute at that time.
"""
@with_kw_noshow struct Sequence
    operations::SortedDict{Dates.DateTime, Vector{AbstractRunnable}} = SortedDict{Dates.DateTime, Vector{AbstractRunnable}}()
end

function get_operations(sequence::Sequence)
    return sequence.operations
end

function get_timepoints(sequence::Sequence)
    return collect(keys(get_operations(sequence)))
end

function get_steps(sequence::Sequence, ech::DateTime)
    return get_operations(sequence)[ech]
end

"""
    length of the sequence in terms of time not in terms of number of runnables to execute
"""
function Base.length(sequence::Sequence)
    return length(get_operations(sequence))
end

function Base.show(io::IO, sequence::Sequence)
    pretty_print(io, get_operations(sequence))
end

function get_horizon_timepoints(sequence::Sequence)::Vector{Dates.DateTime}
    return collect(keys(get_operations(sequence)))
end

function get_ech(sequence::Sequence, index::Int)
    if length(sequence) < index
        throw( error("attempt to acess ", length(sequence), "-element Sequence at index ", index, ".") )
    else
        return get_horizon_timepoints(sequence)[index]
    end
end

function add_step!(sequence::Sequence, step_type::Type{T}, ech::Dates.DateTime) where T<:AbstractRunnable
    step_instance = step_type()
    add_step!(sequence, step_instance, ech)
end
function add_step!(sequence::Sequence, step_instance::AbstractRunnable, ech::Dates.DateTime)
    steps_at_ech = get!(sequence.operations, ech, Vector{AbstractRunnable}())
    push!(steps_at_ech, step_instance)
end

###############################################
# Sequence Launch
###############################################

function init!(context_p::AbstractContext, sequence_p::Sequence, check_context::Bool)
    if !isnothing(context_p.out_dir)
        @warn("removing files that start with $(PSCOPFio.OUTPUT_PREFIX) in $(context_p.out_dir)")
        rm_prefixed(context_p.out_dir, PSCOPFio.OUTPUT_PREFIX)
    end
    set_horizon_timepoints(context_p, get_timepoints(sequence_p))

    @info @sprintf("Launching Management Mode : %s", context_p.management_mode.name)
    @info "Launching sequence :"
    @info sequence_p
    @info @sprintf("Interest date (T) : %s", get_target_timepoints(context_p)[1])
    @info @sprintf("target timepoints (TS) : %s", get_target_timepoints(context_p))
    @info @sprintf("Horizon timepoints (ECH=T-t) : %s", get_horizon_timepoints(context_p))
    @info("initial units state: ")
    @info get_generators_initial_state(context_p)

    if check_context && !check(context_p)
        throw( error("Invalid context!") )
    end
end

function run_step!(context_p::AbstractContext, step::AbstractRunnable, ech, next_ech)
    println("-"^20)
    println(typeof(step))
    firmness = compute_firmness(step, ech, next_ech,
                            get_target_timepoints(context_p), context_p)
    trace_firmness(firmness)
    result = run(step, ech, firmness,
                get_target_timepoints(context_p),
                context_p)

    if affects_market_schedule(step)
        @debug "update market schedule based on optimization results"
        # old_market_schedule = deepcopy(get_market_schedule(context_p))
        update_market_schedule!(context_p, ech, result, firmness, step)
        # println("Changes to the market schedule:")
        # trace_delta_schedules(old_market_schedule, get_market_schedule(context_p))
        #TODO : error if !verify
        verify_firmness(firmness, context_p.market_schedule,
                        excluded_ids=get_limitables_ids(context_p))
        PSCOPF.PSCOPFio.write(context_p, get_market_schedule(context_p), "market_")
        PSCOPF.PSCOPFio.write_full(context_p, get_market_schedule(context_p), "market_")
        update_market_flows!(context_p)
        trace_flows(get_market_flows(context_p), get_network(context_p))
    end

    if affects_tso_schedule(step)
        @debug "update TSO schedule based on optimization results"
        # old_tso_schedule = deepcopy(get_tso_schedule(context_p))
        update_tso_schedule!(context_p, ech, result, firmness, step)
        # println("Changes to the TSO schedule:")
        # trace_delta_schedules(old_tso_schedule, get_tso_schedule(context_p))
        #TODO : error if !verify
        verify_firmness(firmness, context_p.tso_schedule,
                        excluded_ids=get_limitables_ids(context_p))
        PSCOPF.PSCOPFio.write(context_p, get_tso_schedule(context_p), "tso_")
        PSCOPF.PSCOPFio.write_full(context_p, get_tso_schedule(context_p), "tso_")
        update_tso_flows!(context_p)
        trace_flows(get_tso_flows(context_p), get_network(context_p))
    end

    if affects_tso_actions(step)
        @debug "update TSO actions based on optimization results"
        update_tso_actions!(context_p,
                            ech, result, firmness, step)
        trace_tso_actions(get_tso_actions(context_p))
    end
    #TODO check coherence between tso schedule and actions

    if ( (affects_market_schedule(step) || affects_tso_schedule(step)) )
        schedules_to_delta = sort([get_market_schedule(context_p), get_tso_schedule(context_p)],
                                    by=x->x.decision_time)
        @printf("changes in %s schedule compared to preceding %s schedule:\n",
                schedules_to_delta[2].decider_type, schedules_to_delta[1].decider_type)
        trace_delta_schedules(schedules_to_delta...)
    end

    return result, firmness
end

function run!(context_p::AbstractContext, sequence_p::Sequence;
                check_context=true)
    init!(context_p, sequence_p, check_context)

    for (steps_index, (ech, steps_at_ech)) in enumerate(get_operations(sequence_p))
        println("-"^50)
        delta = Dates.value(Dates.Minute(get_target_timepoints(context_p)[1]-ech))
        println("ECH : ", ech)
        println("t : M-", delta)
        println("-"^50)
        for step in steps_at_ech
            next_ech = get_next_ech(sequence_p, steps_index, step)
            run_step!(context_p, step, ech, next_ech)
        end
    end
end

"""
returns the ech that comes after ech at which we execute a step of the same DeciderType as the input step
if none exists returns nothing
"""
function get_next_ech(sequence::Sequence, index::Int, decider_step::AbstractRunnable)
    if index >= length(sequence)
        return nothing
    end

    for ech_l in get_horizon_timepoints(sequence)[index+1:end]
        for step_l in get_steps(sequence, ech_l)
            if DeciderType(step_l) == DeciderType(decider_step)
                return ech_l
            end
        end
    end

    return nothing
end

##################################
# Trace
##################################

function trace_firmness(firmness::Firmness)
    for (gen_id,_) in firmness.commitment
        TS_to_decide_for = Vector{DateTime}()
        for (ts, decision_firmness) in firmness.commitment[gen_id]
            if decision_firmness == TO_DECIDE
                push!(TS_to_decide_for, ts)
            end
        end
        if !isempty(TS_to_decide_for)
            msg = @sprintf("TO_DECIDE commitment (Firm) :  %s for %s",
                            gen_id, TS_to_decide_for)
            @info msg
        end
    end

    for (gen_id,_) in firmness.power_level
        TS_to_decide_for = Vector{DateTime}()
        for (ts, decision_firmness) in firmness.power_level[gen_id]
            if decision_firmness == TO_DECIDE
                push!(TS_to_decide_for, ts)
            end
        end
        if !isempty(TS_to_decide_for)
            msg = @sprintf("TO_DECIDE production level (Firm) :  %s for %s",
                            gen_id, TS_to_decide_for)
            @info msg
        end
    end
end

function trace_flows(flows::SortedDict{Tuple{String, DateTime, String}, Float64},
                    network::Networks.Network)
    for ((branch_id, ts, scenario), flow_val) in flows
        branch = Networks.safeget_branch(network, branch_id)
        limit = Networks.get_limit(branch)
        if abs(flow_val) > limit
            @printf("Flow value %f for branch %s, at timestep %s and scenario %s exceeds branch limit (%f)\n",
                    flow_val, branch_id, ts, scenario, limit)
        end
    end
end

function trace_delta_schedules(old_schedule::Schedule, new_schedule::Schedule)
    print_non_firm_changes = false

    println("Commitment updates :")
    trace_delta_schedule_component(old_schedule, new_schedule, get_commitment_sub_schedule, print_non_firm_changes)
    println("Production updates :")
    trace_delta_schedule_component(old_schedule, new_schedule, get_production_sub_schedule, print_non_firm_changes)
end

function trace_delta_schedule_component(old_schedule::Schedule, new_schedule::Schedule,
                                        component_accessor::Function,
                                        print_non_firm_changes=false)
    if !print_non_firm_changes
        println("(only showing firm changes)")
    end
    for (gen_id, _) in new_schedule.generator_schedules
        @assert(haskey(old_schedule.generator_schedules, gen_id))

        #retrieve either commitment or production schedule component
        old_schedule_component = component_accessor(old_schedule, gen_id)
        new_schedule_component = component_accessor(new_schedule, gen_id)

        trace_delta_genschedule_component(gen_id, old_schedule_component, new_schedule_component,
                                        print_non_firm_changes)
    end
end

function trace_delta_genschedule_component(gen_id::String,
                                        old_schedule::SortedDict{Dates.DateTime, UncertainValue{T}},
                                        new_schedule::SortedDict{Dates.DateTime, UncertainValue{T}},
                                        print_non_firm_changes=false) where T
    for (ts, new_uncertain) in new_schedule
        @assert(haskey(old_schedule, ts))
        old_uncertain = old_schedule[ts]
        old_firmness_msg = is_definitive(old_uncertain) ? "firm" : "non-firm"
        new_firmness_msg = is_definitive(new_uncertain) ? "firm" : "non-firm"

        if is_definitive(new_uncertain) && is_definitive(old_uncertain)
            if is_different(get_value(new_uncertain), get_value(old_uncertain))
                @printf("%s\t%s\t%s (%s) --> %s (%s)\n",
                        gen_id, ts, get_value(old_uncertain), old_firmness_msg, get_value(new_uncertain), new_firmness_msg)
            end

        elseif is_definitive(new_uncertain) && !is_definitive(old_uncertain)
            @printf("%s\t%s\t%s (%s)\n",
                    gen_id, ts, get_value(new_uncertain), new_firmness_msg)

        elseif !is_definitive(new_uncertain) && is_definitive(old_uncertain)
            @printf("%s\t%s\t%s (%s) --> (%s)\n",
                    gen_id, ts, get_value(old_uncertain), old_firmness_msg, new_firmness_msg)
            if print_non_firm_changes
                println(new_uncertain.anticipated_value)
            end

        elseif print_non_firm_changes
            for scenario in get_scenarios(new_uncertain)
                old_val = get_value(old_uncertain, scenario)
                new_val = get_value(new_uncertain, scenario)
                if ( (ismissing(new_val) != ismissing(old_val))
                    || ( !ismissing(old_val) && !ismissing(new_val) && is_different(new_val, old_val) )
                    )
                    @printf("%s\t%s\t%s\t%s --> %s\n",
                            gen_id, ts, scenario, old_val, new_val)
                end
            end
        end
    end

end

function trace_tso_actions(tso_actions::TSOActions)
    trace_limitations(tso_actions)
    trace_impositions(tso_actions)
end
function trace_limitations(tso_actions)
    println("TSO limitation actions :")
    for ((gen_id,ts), val_l) in get_limitations(tso_actions)
        @printf("\t%s at %s : %.2f\n", gen_id, ts, val_l)
    end
end
function trace_impositions(tso_actions)
    println("TSO imposition actions :")
    for ((gen_id,ts), imposition) in get_impositions(tso_actions)
        @printf("\t%s at %s : %s\n", gen_id, ts, imposition)
    end
end

max_stars_repo_name: klorel/pscopf

max_issues_repo_name: klorel/pscopf

max_forks_repo_name: klorel/pscopf


text:
module StructWalk

export prewalk, postwalk

"""
Abstract type `WalkStyle`

Subtype `WalkStyle` and overload [`walkstyle`](@ref) to define custom walking behaviors (constructor / children /...).
"""
abstract type WalkStyle end

"""
    walkstyle(::CustomWalkStyle, x::T) where {CumstomWalkStyle <: WalkStyle}

Should return a tuple of length 2-3 with:

    1. A proper constuctor for `T`, can be `identity` if `x` isa leaf.
    2. Children of `x` in a tuple, or empty tuple `()` if `x` is a leaf.
    3. [optional] a bool indicate whether element of 2. is the actual list of children. default to `false`.

For example, since `Array` has 0 `fieldcount`, we doesn't split the value into a tuple as children.
 Instead, we return `(x,)` as children and the extra boolean `true`, so it will `walk`/`map` through `x`
 accordingly.
"""
function walkstyle end

"""
    walkstyle(x)
    walkstyle(::Type{WalkStyle}, x::T) where T

return `T` and a tuple all field values of `x`.
"""
walkstyle(x) = walkstyle(WalkStyle, x)
walkstyle(s::WalkStyle, x) = walkstyle(WalkStyle, x)
function walkstyle(::Type{WalkStyle}, x::T) where T
    n = fieldcount(T)
    isleaf = iszero(n)
    return T.name.wrapper, isleaf ? () : ntuple(i->getfield(x, i), n)
end

# default walkstyle for some types
include("./walkstyle.jl")

"""
    LeafNode(x)

special type for marking non-leaf value as leaf. Use with `prewalk`.

See also: [`prewalk`](@ref)
"""
struct LeafNode{T}
    x::T
end

@nospecialize

walk(_, _, _, x::LeafNode, _) = x.x

walk(f, style, x, inner_walk) = walk(f, f, style, x, inner_walk)
function walk(f, g, style, x, inner_walk)
    S = walkstyle(style, x)
    T, fields = S
    isleaf = isempty(fields)
    isnontuple = length(S) <= 2 ? false : S[3]
    if isleaf
        return f(x)
    else
        h = isnontuple ? v->map(inner_walk, v) : inner_walk
        return g(T(map(h, fields)...))
    end
end

"""
    postwalk(f, [style = WalkStyle], x)

Applies `f` to each node in `x` and return the result.
`f` sees the leaves first and then the transformed node.

# Example

```julia
julia> postwalk(x -> @show(x) isa Integer ? x + 1 : x, (a=2, b=(c=4, d=0)))
x = 2
x = 4
x = 0
x = (c = 5, d = 1)
x = (a = 3, b = (c = 5, d = 1))
(a = 3, b = (c = 5, d = 1))

julia> postwalk(x -> @show(x) isa Integer ? x + 1 : x .+ 1, (3, 5))
x = 3
x = 5
x = (4, 6)
(5, 7)

julia> postwalk(x -> @show(x) isa Integer ? x // 2 : x isa Tuple ? =>(x .+ 1...) : x, (3, 5))
x = 3
x = 5
x = (3//2, 5//2)
5//2 => 7//2

```

See also: [`prewalk`](@ref)
"""
postwalk(f, x) = postwalk(f, WalkStyle, x)
postwalk(f, style, x) = walk(f, style, x, x -> postwalk(f, style, x))

"""
    prewalk(f, [style = WalkStyle], x)

Applies `f` to each node in `x` and return the result.
`f` sees the node first and then the transformed leaves.

*Notice* that it is possible it walk infinitely if you transform a node into non-leaf value.
 Wrapping the non-leaf value with `LeafNode(y)` in `f` to prevent infinite walk.

# Example

```julia
julia> prewalk(x -> @show(x) isa Integer ? x + 1 : x, (a=2, b=(c=4, d=0)))
x = (a = 2, b = (c = 4, d = 0))
x = 2
x = (c = 4, d = 0)
x = 4
x = 0
(a = 3, b = (c = 5, d = 1))

julia> prewalk(x -> @show(x) isa Integer ? x + 1 : x .+ 1, (3, 5))
x = (3, 5)
x = 4
x = 6
(5, 7)

julia> prewalk(x -> @show(x) isa Integer ? StructWalk.LeafNode(x // 2) : x isa Tuple ? =>(x .+ 1...) : x, (3, 5))
x = (3, 5)
x = 4
x = 6
2//1 => 3//1

```

See also: [`postwalk`](@ref), [`LeafNode`](@ref)
"""
prewalk(f, x) = prewalk(f, WalkStyle, x)
prewalk(f, style, x) = walk(identity, style, f(x), x -> prewalk(f, style, x))

@specialize

end


text:
function Base.show(io::IO, mime::MIME"text/plain", dims::DimTuple)
    ctx = IOContext(io, :compact=>true)
    if all(map(d -> !(parent(d) isa AbstractArray) || (parent(d) isa NoLookup), dims))
        for d in dims[1:end-1]
            show(ctx, mime, d)
            print(io, ", ")
        end
        show(ctx, mime, dims[end])
        return 0
    else # Dims get a line each
        haskey(io, :inset) && print(io, "\n")
        inset = get(io, :inset, "")
        lines = 3
        for d in dims[1:end-1]
            print(io, inset)
            show(ctx, mime, d)
            print(io, ",")
            lines += 2 # Often they wrap
            print(io, "\n")
        end
        print(io, inset)
        show(ctx, mime, dims[end])
        return lines
    end
end
function Base.show(io::IO, mime::MIME"text/plain", dim::Dimension)
    get(io, :compact, false) && return show_compact(io::IO, mime, dim)
    print_dimname(io, dim)
    print_dimval(io, mime, val(dim))
end
function Base.show(io::IO, mime::MIME"text/plain", dim::Dimension{Colon})
    print_dimname(io, dim)
end

# compact version for dimensions and lookups
show_compact(io::IO, mime, dim::Dimension{Colon}) = print_dimname(io, dim)
function show_compact(io::IO, mime, dim::Dimension)
    # Print to a buffer and count lengths
    buf = IOBuffer()
    print_dimname(buf, dim)
    nchars = length(String(take!(buf)))
    print_dimval(buf, mime, parent(dim), nchars)
    nvalchars = length(String(take!(buf)))
    # Actually print to IO
    print_dimname(io, dim)
    if nvalchars > 0
        print_dimval(io, mime, val(dim), nchars)
    end
end

dimcolor(io) = get(io, :is_ref_dim, false) ? :magenta : :red

# print dims with description string and inset
function print_dims(io::IO, mime, dims::Tuple)
    if isempty(dims)
        print(io, ": ")
        return 0
    end
    printstyled(io, " with dimensions: "; color=:light_black)
    ctx = IOContext(io, :inset => "  ")
    return show(ctx, mime, dims)
end
# print refdims with description string and inset
function print_refdims(io::IO, mime, refdims::Tuple)
    if isempty(refdims)
        return 0
    end
    printstyled(io, "and reference dimensions: "; color=:light_black)
    ctx = IOContext(io, :inset => "  ", :is_ref_dim=>true, :show_dim_val=>true)
    lines = show(ctx, mime, refdims)
    return lines
end
# print a dimension name
function print_dimname(io, dim::Dim)
    color = dimcolor(io)
    printstyled(io, "Dim{"; color=color)
    printstyled(io, string(":", name(dim)); color=:yellow)
    printstyled(io, "}"; color=color)
end
function print_dimname(io, dim::Dimension)
    printstyled(io, dim2key(dim); color = dimcolor(io))
end

# print the dimension value
function print_dimval(io, mime, val, nchars=0)
    val isa Colon || print(io, " ")
    printstyled(io, val; color=:cyan)
end
function print_dimval(io, mime, lookup::AbstractArray, nchars=0)
    LookupArrays.print_index(io, mime, lookup, nchars)
end
print_dimval(io, mime, lookup::Union{AutoLookup,NoLookup}, nchars=0) = print(io, " ")
function print_dimval(io, mime, lookup::LookupArray, nchars=0)
    print(io, " ")
    ctx = IOContext(io, :nchars=>nchars)
    show(ctx, mime, lookup)
end


text:
using SIMDPirates, VectorizationBase
using Test

for T in (Float16, Float32, Float64)
    @test sizeof(int_type(T)) == sizeof(T)
    @test sizeof(uint_type(T)) == sizeof(T)
    @test significand_bits(T) + exponent_bits(T) + sign_bits(T) == 8*sizeof(T)
    @test significand_mask(T) | exponent_mask(T) | sign_mask(T) ==
        typemax(uint_type(T))
    @test significand_mask(T) ⊻ exponent_mask(T) ⊻ sign_mask(T) ==
        typemax(uint_type(T))

    W = VectorizationBase.pick_vector_width(T)
    x = ntuple(W) do _ Core.VecElement(randn(T)) end
    y = ntuple(W) do _ Core.VecElement(randn(T)) end
    @test SIMDPirates.vall(
        SIMDPirates.vgreater_or_equal(x,y) === ntuple(_ -> Core.VecElement(x[w].value >= y[w].value), W)
    )
    @test SIMDPirates.vecbool_to_mask( SIMDPirates.vless(x,y) ) === SIMDPirates.vgreater_mask(y, x)
end


text:
#####
##### Exploring a player's policy
#####

using DataStructures: Stack

# All state statistics and are expressed with respect to the current player

@kwdef mutable struct StateActionStats
  P     :: Union{Float64, Nothing} = nothing  # Probability given by `think`
  Pmem  :: Union{Float64, Nothing} = nothing  # Recorded π component in memory
  Pmcts :: Union{Float64, Nothing} = nothing  # Percentage of MCTS visits
  Qmcts :: Union{Float64, Nothing} = nothing
  UCT   :: Union{Float64, Nothing} = nothing
  Pnet  :: Union{Float64, Nothing} = nothing
  Qnet  :: Union{Float64, Nothing} = nothing
end

mutable struct StateStats{Action}
  Nmcts :: Union{Int, Nothing}
  Nmem  :: Union{Int, Nothing}
  Vmem  :: Union{Float64, Nothing}
  Vnet  :: Union{Float64, Nothing}
  actions :: Vector{Tuple{Action, StateActionStats}} # Sorted by P
  function StateStats(actions)
    actions_stats = [(a, StateActionStats()) for a in actions]
    new{eltype(actions)}(nothing, nothing, nothing, nothing, actions_stats)
  end
end

player_reward(game, white_reward) =
  GI.white_playing(game) ? white_reward : - white_reward

function evaluate_vnet(oracle, game)
  if GI.game_terminated(game)
    return 0.0
  else
    return oracle(GI.current_state(game))[2]
  end
end

function evaluate_qnet(oracle, game, action, gamma)
  @assert !GI.game_terminated(game)
  next = GI.clone(game)
  GI.play!(next, action)
  wr = GI.white_reward(next)
  r = GI.white_playing(game) ? wr : -wr
  qnext = evaluate_vnet(oracle, next)
  if GI.white_playing(game) != GI.white_playing(next)
    qnext = -qnext
  end
  return r + gamma * qnext
end

player_oracle(p) = nothing
player_oracle(p::MctsPlayer) = p.mcts.oracle
player_oracle(p::NetworkPlayer) = p.network

player_gamma(p) = p.gamma
player_gamma(p::MctsPlayer) = p.mcts.gamma

function state_statistics(game, player, turn, memory=nothing)
  @assert !GI.game_terminated(game)
  state = GI.current_state(game)
  # Make the player think
  actions, π = think(player, game)
  τ = player_temperature(player, game, turn)
  π = apply_temperature(π, τ)
  report = StateStats(actions)
  for i in eachindex(actions)
    report.actions[i][2].P = π[i]
  end
  # Collect MCTS Statistics
  if isa(player, MctsPlayer) && haskey(player.mcts.tree, state)
    mcts = player.mcts
    info = mcts.tree[state]
    ucts = MCTS.uct_scores(info, mcts.cpuct, 0., nothing)
    report.Nmcts = MCTS.Ntot(info)
    for (i, a) in enumerate(actions)
      astats = info.stats[i]
      arep = report.actions[i][2]
      arep.Pmcts = astats.N / max(1, report.Nmcts)
      arep.Qmcts = astats.N > 0 ? astats.W / astats.N : 0.
      arep.UCT = ucts[i]
    end
  end
  # Collect memory statistics
  if !isnothing(memory)
    mem = AlphaZero.merge_by_state(get_experience(memory))
    relevant = findall((ex -> ex.s == state), mem)
    if !isempty(relevant)
      @assert length(relevant) == 1
      e = mem[relevant[1]]
      report.Nmem = e.n
      report.Vmem = e.z
      for i in eachindex(actions)
        report.actions[i][2].Pmem = e.π[i]
      end
    end
  end
  # Collect network statistics
  oracle = player_oracle(player)
  if isa(oracle, AbstractNetwork)
    Pnet, Vnet = oracle(state)
    report.Vnet = Vnet
    for i in eachindex(actions)
      arep = report.actions[i][2]
      arep.Pnet = Pnet[i]
      arep.Qnet = evaluate_qnet(oracle, game, actions[i], player_gamma(player))
    end
  end
  # Sort the actions from best to worst
  sortby = ((a, arep),) -> arep.P
  sort!(report.actions, rev=true, by=sortby)
  return report
end

#####
##### Displaying state statistics
#####

function print_state_statistics(gspec, stats::StateStats)
  prob   = Log.ColType(nothing, x -> fmt(".1f", 100 * x) * "%")
  val    = Log.ColType(nothing, x -> fmt("+.2f", x))
  bigint = Log.ColType(nothing, n -> format(ceil(Int, n), commas=true))
  alabel = Log.ColType(nothing, identity)
  btable = Log.Table([
    ("Nmcts", bigint, r -> r.Nmcts),
    ("Nmem",  bigint, r -> r.Nmem),
    ("Vmem",  val,    r -> r.Vmem),
    ("Vnet",  val,    r -> r.Vnet)],
    header_style=Log.BOLD)
  atable = Log.Table([
    ("",      alabel, r -> GI.action_string(gspec, r[1])),
    ("",      prob,   r -> r[2].P),
    ("Pmcts", prob,   r -> r[2].Pmcts),
    ("Pnet",  prob,   r -> r[2].Pnet),
    ("UCT",   val,    r -> r[2].UCT),
    ("Pmem",  prob,   r -> r[2].Pmem),
    ("Qmcts", val,    r -> r[2].Qmcts),
    ("Qnet",  val,    r -> r[2].Qnet)],
    header_style=Log.BOLD)
  logger = Logger()
  if !all(isnothing, [stats.Nmcts, stats.Nmem, stats.Vmem, stats.Vnet])
    Log.table(logger, btable, [stats])
    Log.sep(logger)
  end
  Log.table(logger, atable, stats.actions)
  Log.sep(logger)
end

# Return the stats
function compute_and_print_state_statistics(exp)
  if !GI.game_terminated(exp.game)
    stats = state_statistics(exp.game, exp.player, exp.turn, exp.memory)
    print_state_statistics(GI.spec(exp.game), stats)
    return stats
  else
    return nothing
  end
end

#####
##### Interactive exploration of the environment
#####

mutable struct Explorer
  game :: AbstractGameEnv
  history :: Stack{AbstractGameEnv}
  player :: AbstractPlayer
  memory :: Union{MemoryBuffer, Nothing}
  turn :: Int
  function Explorer(
      player::AbstractPlayer,
      game::AbstractGameEnv;
      memory=nothing)
    history = Stack{AbstractGameEnv}()
    new(game, history, player, memory, 0)
  end
end

function restart!(exp::Explorer)
  reset_player!(exp.player)
  empty!(exp.history)
  exp.game = GI.init(GI.spec(exp.game))
  exp.turn = 0
end

save_game!(exp::Explorer) = push!(exp.history, GI.clone(exp.game))

# Return true if the command is valid, false otherwise
function interpret!(exp::Explorer, stats, cmd, args=[])
  gspec = GI.spec(exp.game)
  if cmd == "go"
    st = GI.read_state(gspec)
    if isnothing(st)
      println("Invalid state description.")
      return false
    end
    g = GI.init(gspec, st)
    if !isnothing(st)
      save_game!(exp)
      exp.game = g
      exp.turn = 0
      return true
    end
  elseif cmd == "restart"
    restart!(exp)
    return true
  elseif cmd == "undo"
    if !isempty(exp.history)
      exp.game = pop!(exp.history)
      exp.turn -= 1
      return true
    end
  elseif cmd == "do"
    GI.game_terminated(exp.game) && return false
    if length(args) == 0
      a = stats.actions[1][1]
    else
      length(args) == 1 || return false
      a = GI.parse_action(GI.spec(exp.game), args[1])
      isnothing(a) && return false
      a ∈ GI.available_actions(exp.game) || return false
    end
    save_game!(exp)
    GI.play!(exp.game, a)
    exp.turn += 1
    return true
  elseif cmd == "flip"
    save_game!(exp)
    GI.apply_random_symmetry!(exp.game)
    return true
  elseif cmd == "explore"
    isa(exp.player, MctsPlayer) || (return false)
    try
      if isempty(args)
        n = exp.player.niters
      else
        n = parse(Int, args[1])
      end
      MCTS.explore!(exp.player.mcts, exp.game, n)
      return true
    catch e
      isa(e, ArgumentError) && (return false)
      rethrow(e)
    end
  elseif cmd == "dump"
    GI.dump_state(Game)
    return true
  end
  return false
end

function start_explorer(exp::Explorer)
  while true
    # Print the state
    GI.render(exp.game)
    stats = compute_and_print_state_statistics(exp)
    # Interpret command
    while true
      print("> ")
      inp = readline() |> lowercase |> split
      isempty(inp) && return
      cmd = inp[1]
      args = inp[2:end]
      interpret!(exp, stats, cmd, args) && break
    end
    println("")
  end
end

#####
##### The explore methods
#####

"""
    explore(::AbstractPlayer, ::AbstractGameEnv;  [memory])
    explore(::AbstractPlayer, ::AbstractGameSpec; [memory])

Start a command interpreter to explore the internals of a player
through interactive play.

The `memory` argument is an optional reference to a memory buffer.
When provided, additional state statistics are displayed.

# Commands

The following commands are currently implemented:

  - `do [action]`: make the current player perform `action`.
    By default, the action of highest score is played.
  - `explore [num_sims]`: run `num_sims` MCTS simulations from the current
    state (for MCTS players only).
  - `go`: query the user for a state description and go to this state.
  - `flip`: flip the board according to a random symmetry.
  - `undo`: undo the effect of the previous command.
  - `restart`: restart the explorer.
"""
function explore end

function explore(
    player::AbstractPlayer,
    game::AbstractGameEnv;
    memory=nothing)
  return start_explorer(Explorer(player, game, memory=memory))
end

function explore(player::AbstractPlayer, gspec::AbstractGameSpec; args...)
  return explore(player, GI.init(gspec), args...)
end


text:
using OutlierDetectionNetworks
using OutlierDetectionNetworks.Templates
using OutlierDetectionTest

# Test the metadata of all exported detectors
test_meta.(eval.(OutlierDetectionNetworks.MODELS))

data = TestData()
run_test(detector) = test_detector(detector, data)

const encoder, decoder = MLPAutoEncoder(size(data.x_raw, 1), 5, [50,20]; bias=false);

# AE
run_test(AEDetector(encoder=encoder, decoder=decoder))

# DSAD
run_test(DSADDetector(encoder=encoder, decoder=decoder))

# ESAD
run_test(ESADDetector(encoder=encoder, decoder=decoder))

max_stars_repo_name: davnn/OutlierDetectionNetworks

max_issues_repo_name: davnn/OutlierDetectionNetworks

max_forks_repo_name: davnn/OutlierDetectionNetworks


text:
## Data frame test
## not run by default. Too time consuming and relies on external pacakgs
using Mustache, DataFrames

glm_tpl = mt"""
\begin{table}
\begin{tabular}{l @{\quad} rrrr}
{{#colnms}}
&\verb+{{{colnm}}}+
{{/colnms}}\\
{{#mat}}
{{variable}} & {{col1}} & {{col2}} & {{col3}} & {{col4}}\\
{{/mat}}
\end{tabular}
\end{table}
"""

function glm_table(mod)
    tbl = coeftable(mod)

    colnms = DataFrame(colnm=tbl.colnms)
    mat = DataFrame(variable=tbl.rownms)
    for j in 1:size(tbl.mat)[2]
        nm = "col$j"
        mat[Symbol(nm)] = map(x -> @sprintf("%.2f", x), tbl.mat[:,j])
    end

    Mustache.render(glm_tpl,Dict("colnms"=>colnms, "mat"=>mat))
end

using GLM, RDatasets, DataFrames
LifeCycleSavings = dataset("datasets", "LifeCycleSavings")
fm2 = fit(LinearModel, SR ~ Pop15 + Pop75 + DPI + DDPI, LifeCycleSavings)
glm_table(fm2)


text:
using Observables
using Test

@testset "listeners" begin
    r = Observable(0)
    @test r[] == 0
    r[] = 1
    @test r[] == 1

    f = on(r) do x
        r[] == 2
        @test x == 2
    end
    r[] = 2

    off(r, f)
    @test !(f in r.listeners)
    @test_throws KeyError off(r, f)
    r[] = 3 # shouldn't call test
end

@testset "onany" begin
    r = Observable(0)
    tested = Ref(false)
    onany(1, r) do x, y
        @test x == 1
        @test y == 2
        tested[] = true
    end
    r[] = 2
    @test tested[]

    r1 = Observable(0)
    r2 = Observable(0)
    map!(x->x+1, r2, r1)
    @test r2[] == 0
    r1[] = 1
    @test r2[] == 2

    r1 = Observable(2)
    r2 = map(+, r1, 1)
    @test r2[] == 3
    r1[] = 3
    @test r2[] == 4
end

@testset "macros" begin
    a = Observable(2)
    b = Observable(3)
    c = Observables.@map &a + &b
    @test c isa Observable
    @test c[] == 5
    a[] = 100
    sleep(0.1)
    @test c[] == 103

    a = Observable(2)
    b = Observable(3)
    c = Observable(10)
    Observables.@map! c &a + &b
    sleep(0.1)
    @test c[] == 10
    a[] = 100
    sleep(0.1)
    @test c[] == 103

    a = Observable(2)
    b = Observable(3)
    c = Observable(10)
    Observables.@on c[] = &a + &b
    sleep(0.1)
    @test c[] == 10
    a[] = 100
    sleep(0.1)
    @test c[] == 103
end

@testset "async_latest" begin
    o = Observable(0)
    cnt = Ref(0)
    function compute_something(x)
        for i=1:10^7; rand() end
        cnt[] = cnt[] + 1
    end
    o_latest = async_latest(o, 1)

    on(compute_something, o_latest) # compute something on the latest update
    for i=1:5
        o[] = i
    end
    sleep(1)

    @test o_latest[] == 5
    @test cnt[] == 1

    for i=5:-1:1
        @async o[] = i
    end
    sleep(1)

    @test o_latest[] == 1
    @test cnt[] == 2 # only one more

    o = Observable(0)
    cnt[] = 0
    function compute_something(x)
        for i=1:10^8; rand() end
        cnt[] = cnt[] + 1
    end
    sleep(1)
    o_latest = async_latest(o, 3)

    on(compute_something, o_latest) # compute something on the latest update
    for i=1:5
        o[] = i
    end
    sleep(.1)

    @test o_latest[] == 5
    @test cnt[] == 3
end

@testset "throttle" begin
    obs = Observable(1)
    throttled = throttle(2, obs)
    @test throttled[] == 1
    obs[] = 2
    sleep(0.1)
    @test throttled[] == 2
    obs[] = 3
    sleep(1)
    @test throttled[] == 2
    sleep(2)
    @test throttled[] == 3
end

struct A<:Observables.AbstractObservable{Int}; end

struct B{T}<:Observables.AbstractObservable{T}
    output::Observable{T}
end

Observables.observe(b::B) = b.output

@testset "interface" begin
    a = A()
    @test_throws ErrorException a[]
    @test_throws ErrorException a[] = 2
    @test_throws ErrorException Observables.obsid(a)
    @test_throws ErrorException Observables.listeners(a)
    @test eltype(a) == Int

    b = B(Observable("test"))
    @test b[] == "test"
    b[] = "test2"
    sleep(0.1)
    @test b[] == "test2"
    @test isempty(Observables.listeners(b))
    @test eltype(b) == String
end

@testset "pair" begin
    v = Observables.ObservablePair(Observable(1.0), f = exp, g = log)
    @test length(v) == 2
    @test firstindex(v) == 1
    @test lastindex(v) == 2
    firstobs, lastobs = v
    @test firstobs[] == 1.0
    @test lastobs[] ≈ ℯ
    @test v.second[] ≈ ℯ
    @test first(v)[] ≈ 1.0
    @test last(v)[] ≈ ℯ
    @test v[2][] ≈ ℯ

    v.first[] = 0
    @test v.second[] ≈ 1
    v.second[] = 2
    @test v.first[] ≈ log(2)

    obs = Observable(Observable(2))
    o2 = Observables.flatten(obs)

    o2[] = 12
    sleep(0.1)
    @test obs[][] == 12
    obs[][] = 22
    sleep(0.1)
    @test o2[] == 22
    obs[] = Observable(11)
    sleep(0.1)
    @test o2[] == 11

    obs = Observable(Observable(Observable(10)))
    o2 = Observables.flatten(obs)
    @test o2[] == 10

    obs[] = Observable(Observable(13))
    @test o2[] == 13
end

@testset "async" begin
    y = Observable(@async begin
        sleep(0.5)
        return 1 + 1
    end)
    @test !isdefined(y, :val)
    while !isdefined(y, :val)
        sleep(0.001)
    end
    @test y[] == 2

    x = Observable(1)
    y = map(x) do val
        @async begin
            sleep(0.5)
            return val + 1
        end
    end
    @test !isdefined(y, :val)
    while !isdefined(y, :val)
        sleep(0.001)
    end
    @test y[] == 2

    x = Observable(1)
    y = map(x) do val
        Channel() do channel
            for i in 1:10
                put!(channel, i + val)
            end
        end
    end
    @test !isdefined(y, :val)
    vals = Int[]
    on(y) do val
        push!(vals, val)
    end
    while length(vals) != 10
        sleep(0.001)
    end
    @test vals == 2:11
    y = Observable(Channel() do channel
        for i in 1:10
            put!(channel, i + 1)
        end
    end)
    @test !isdefined(y, :val)
    vals = Int[]
    on(y) do val
        push!(vals, val)
    end
    while length(vals) != 10
        sleep(0.001)
    end
    @test vals == 2:11
end

@testset "copy" begin
    obs = Observable{Any}(:hey)
    obs_copy = copy(obs)
    obs[] = 1
    @test obs_copy[] == 1
end


text:
import Base.*

# Group action is x ↦ a*x+b
struct SpaceGroupElement{N,T<:Integer} <: GroupElement
    a::SMatrix{N,N,T}
    b::SVector{N, Rational{T}}
    function SpaceGroupElement{N,T}(a::SMatrix{N,N,T}, b::SVector{N,Rational{T}}) where {N, T<:Integer}
        new(a, b-floor.(b))
    end
end

function SpaceGroupElement(a::SMatrix{N,N,T}, b::SVector{N,Rational{T}}) where {N,T<:Integer}
    SpaceGroupElement{N,T}(a, b)
end

SpaceGroupElement{N,T}() where {N,T<:Integer} =
    SpaceGroupElement(one(SMatrix{N,N,T}), zeros(SVector{N, Rational{T}}))

function *(e1::SpaceGroupElement{N,T}, e2::SpaceGroupElement{N,T}) where {N,T<:Integer}
    SpaceGroupElement(e1.a*e2.a, e1.a*e2.b+e1.b)
end

const SpaceGroup{N,T} = FiniteGroup{SpaceGroupElement{N,T}}