text:
using CmdStan
using Plots
using StatsPlots
using AxisArrays
using StatsBase
using Random
using Distributions
using CSV
using LaTeXStrings

##### Some local setups
# If run from command line externally:
#projDir = dirname(@__FILE__)

# If run from Jupyter/Hydrogen, maybe change to suit you:
projDir= "/lhome/johhub/Desktop/ABDA/A7"
tmpDir = projDir*"/tmp1"

#####
#%% Input Data #################################################################
y = [607, 583, 521, 494, 369, 782, 570, 678, 467, 620, 425, 395, 346, 361, 310,
300, 382, 294, 315, 323, 421, 339, 398, 328, 335, 291, 329, 310, 294, 321, 286,
349, 279, 268, 293, 310, 259, 241, 243, 272, 247, 275, 220, 245, 268, 357, 273,
301, 322, 276, 401, 368, 149, 507, 411, 362, 358, 355, 362, 324, 332, 268, 259,
274, 248, 254, 242, 286, 276, 237, 259, 251, 239, 247, 260, 237, 206, 242, 361,
267, 245, 331, 357, 284, 263, 244, 317, 225, 254, 253, 251, 314, 239, 248, 250,
200, 256, 233, 427, 391, 331, 395, 337, 392, 352, 381, 330, 368, 381, 316, 335,
316, 302, 375, 361, 330, 351, 186, 221, 278, 244, 218, 126, 269, 238, 194, 384,
154, 555, 387, 317, 365, 357, 390, 320, 316, 297, 354, 266, 279, 327, 285, 258,
267, 226, 237, 264, 510, 490, 458, 425, 522, 927, 555, 550, 516, 548, 560, 545,
633, 496, 498, 223, 222, 309, 244, 207, 258, 255, 281, 258, 226, 257, 263, 266,
238, 249, 340, 247, 216, 241, 239, 226, 273, 235, 251, 290, 473, 416, 451, 475,
406, 349, 401, 334, 446, 401, 252, 266, 210, 228, 250, 265, 236, 289, 244, 327,
274, 223, 327, 307, 338, 345, 381, 369, 445, 296, 303, 326, 321, 309, 307, 319,
288, 299, 284, 278, 310, 282, 275, 372, 295, 306, 303, 285, 316, 294, 284, 324,
264, 278, 369, 254, 306, 237, 439, 287, 285, 261, 299, 311, 265, 292, 282, 271,
268, 270, 259, 269, 249, 261, 425, 291, 291, 441, 222, 347, 244, 232, 272, 264,
190, 219, 317, 232, 256, 185, 210, 213, 202, 226, 250, 238, 252, 233, 221, 220,
287, 267, 264, 273, 304, 294, 236, 200, 219, 276, 287, 365, 438, 420, 396, 359,
405, 397, 383, 360, 387, 429, 358, 459, 371, 368, 452, 358, 371]

ind = [1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 4, 5, 5, 5, 5,
5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 8,
 8, 8, 8, 8, 9, 9, 9, 9, 9, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10,
 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11, 11,
  11, 11, 11, 11, 11, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12, 12,
   12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13, 13, 13, 14, 14, 14, 14,
   14, 14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 16, 16, 16, 16,
   16, 17, 17, 17, 17, 17, 18, 18, 18, 18, 18, 19, 19, 19, 19, 19, 20, 20, 20,
   20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 20, 21, 21,
   21, 21, 21, 21, 21, 21, 21, 21, 22, 22, 22, 22, 22, 22, 22, 22, 22, 22, 22,
   22, 23, 23, 23, 23, 23, 23, 23, 23, 23, 23, 24, 24, 24, 24, 24, 24, 24, 24,
   24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 24, 25, 25, 25, 25, 25, 25, 25,
   25, 25, 25, 25, 25, 25, 26, 26, 26, 26, 26, 27, 27, 27, 27, 27, 28, 28, 28,
   28, 28, 28, 28, 28, 28, 29, 29, 29, 29, 29, 29, 29, 29, 29, 29, 30, 30, 30,
   30, 30, 30, 31, 31, 31, 31, 31, 32, 32, 32, 32, 32, 33, 34, 34, 34, 34, 34,
   34, 34, 34, 34, 34, 34, 34, 34, 34, 34, 34, 34, 34];

# Indicator for each individual j whether he/she is a child or not:
global child_j = [1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 1,
           0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0];

# Indicator for each observation i whether it comes from a child or not:
child_i = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1,
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1,
1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1,
1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];

# Attempt numbers of each observation i:
x = [1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 1, 2, 3,
4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 1, 2, 3, 4, 5, 6, 7, 8, 9, 1, 2, 3,
4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
15, 16, 17, 18, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
20, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 1, 2,
3, 4, 5, 6, 7, 8, 9, 10, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 1, 2, 3, 4,
5, 6, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5,
6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 1, 2, 3, 4, 5, 6, 7, 8,
9, 10, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 1,
2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 1, 2, 3, 4,
5, 6, 7, 8, 9, 10, 11, 12, 13, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 6, 7,
8, 9, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 1, 2, 3, 4, 5, 6, 1, 2, 3, 4, 5, 1, 2, 3, 4,
5, 1, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18];

# Alternatively, have it all in one array:
x_alt = Array{Int64,1}[
[1, 2, 3, 4, 5],
[1, 2, 3, 4, 5],
[1, 2, 3, 4, 5, 6, 7, 8, 9],
[1], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
[1, 2, 3, 4, 5, 6, 7, 8, 9],
[1, 2, 3, 4, 5],
[1, 2, 3, 4, 5],
[1, 2, 3, 4, 5],
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18],
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20],
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20],
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13],
[1, 2, 3, 4, 5, 6],
[1, 2, 3, 4, 5],
[1, 2, 3, 4, 5],
[1, 2, 3, 4, 5],
[1, 2, 3, 4, 5],
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20],
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20],
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13],
[1, 2, 3, 4, 5],
[1, 2, 3, 4, 5],
[1, 2, 3, 4, 5, 6, 7, 8, 9],
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
[1, 2, 3, 4, 5, 6],
[1, 2, 3, 4, 5],
[1, 2, 3, 4, 5],
[1],
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18]];

#####
#%% Functions ##################################################################
## Jespers HDI:
function hdi(theta_samp,alpha=0.05)
    cred_mass = 1.0-alpha
    ci = zeros(2)
    if length(size(theta_samp))>1   # multidimensional theta
        K,N = size(theta_samp)
        cis = zeros(2,K)
        for k in 1:K    # for each dimension

            ts = theta_samp[k,:]
            sind = sortperm(ts)     # get series of indices that sorts theta values
            sts = ts[sind]  # sorted theta values

            # Shifting a 95% index block from far left to far right and use the
            # samllest interval that can be found during the shift
            N = length(sind)    # number of sample points
            length_ci = Inf     # start with inf HDI interval and shrink
            for i in 1:Int(floor(N*alpha))  # iterate over the 5% lowest values (low to high)
                i2 = Int(floor(N*cred_mass)+i)  # the 5% highest values (low to high), one 95% block away on the other side
                prop_ci = [sts[i],sts[i2]]      # proposed interval
                length_prop_ci = prop_ci[2]-prop_ci[1]
                if length_prop_ci < length_ci
                    ci = prop_ci
                    length_ci = ci[2]-ci[1]
                end
            end
            cis[:,k] = ci

        end
        return cis
    else
        N = length(theta_samp)

        ts = theta_samp
        sind = sortperm(ts)
        sts = ts[sind]

        N = length(sind)
        length_ci = Inf
        for i in 1:Int(floor(N*alpha))
            i2 = Int(floor(N*cred_mass)+i)
            prop_ci = [sts[i],sts[i2]]
            length_prop_ci = prop_ci[2]-prop_ci[1]
            if length_prop_ci < length_ci
                ci = prop_ci
                length_ci = ci[2]-ci[1]
            end
        end
        return ci
    end
end

# A recursive application of the HDI algorithm to find the mode:
function find_mode(sample_pts;runs=2)
    # re-apply the hdi algorithm, until there is only two points or so left
    # so the ci interval becomes smaller and smaller
    # in the end, return an average over the interval
    ci = zeros(2)
    i_left = 0
    i_right = 0

    # Make 50% credibility intervals
    alpha = 0.5;
    cred_mass = 1.0 - alpha;

    # sort the pts:
    sind = sortperm(sample_pts)
    sts = sample_pts[sind]
    N = length(sind)

    if N > runs
        # HDI algorithm, but return indices too:
        length_ci = Inf
        for i in 1:Int(floor(N*alpha))
            i2 = Int(floor(N*cred_mass)+i)
            prop_ci = [sts[i], sts[i2]]
            length_prop_ci = prop_ci[2] - prop_ci[1]
            if length_prop_ci < length_ci
                i_left  = i
                i_right = i2
                ci = prop_ci
                length_ci = ci[2] - ci[1]
            end
        end

        # Recursion:
        ci = find_mode(sts[i_left:i_right],runs=runs)

    else
        ci = [sts[1],sts[end]]
    end

    return ci
end

# A function to make the plots with HDIs
function makeDistributionPlot(X, color="blue"; ann=true, offset=0.0, scale=1.0)
    #%% Mode, Mean, HDIs
    ω = mean(find_mode(X))
    μ_bar = mean(X);
    left,right = hdi(X);

    h = fit(Histogram, X, nbins=200)
    w = h.weights      # weights of each bar
    e = h.edges[1]     # edges of the bars (must be +1 more than bars)

    riemannSum = sum(w.*diff(e))
    w = w ./ riemannSum

    e2 = similar([e[:]; e[:]])  # doubling of the edge vector
    e2[2:2:end] = e             # even indices
    e2[1:2:end] = e             # odd indices
    e2 = e2[2:(end-1)]          # pop the ends

    w2 = similar([w[:]; w[:]])  # doubling of the weight vector
    w2[2:2:end] = w             # even indices
    w2[1:2:end] = w             # odd indices
    w2 = w2 .* scale .+ offset

    w_left = w[sum(e .< left)] * scale + offset  # comparison leaves at max left edge of bar that contains the limit
    w_right = w[sum(e .< right)] * scale + offset
    w_ω = w[sum(e .< ω)] * scale + offset
    w_μ = w[sum(e .< μ_bar)] * scale + offset

    #%% Plotting
    plot(e2, w2,
         linealpha=0.1, linecolor=color,
         fillrange=offset, fillalpha=0.4, fillcolor=color,
         legend=false)

    # or do it by histogram, then we also have vertical lines
    # specify statsplot?
    #histogram(X, bins=100, normalize=:pdf,
    #          color=color, alpha=0.4, linealpha=0.1, legend=false)   # Comes from StatsPlots now

    if ann
    plot!([(left,offset),(left,w_left)],
        linewidth=1.5, color=color, linestyle=:dash,
        annotations = (left, w_left, text("$(round(left,sigdigits=4))",color,rotation=90,:center,:left)))
    plot!([(right,offset),(right,w_right)],
        linewidth=1.5, color=color, linestyle=:dash,
        annotations = (right, w_right, text("$(round(right,sigdigits=4))",color,rotation=90,:center,:left)))
    plot!([(μ_bar,offset),(μ_bar,w_μ)],
        linewidth=1.5, color=color, linestyle=:dash,
        annotations = (μ_bar, 0.5*w_μ, text("mean $(round(μ_bar,sigdigits=4))",color,rotation=90,:bottom,:right)))
    plot!([(ω,offset),(ω,w_ω)],
        linewidth=1.5, color=color, linestyle=:dashdot,
        annotations = (ω, w_ω, text("mode $(round(ω,sigdigits=4))",color,rotation=90,:bottom,:right)))
    else
    plot!([(left,offset),(left,w_left)],
        linewidth=1.5, color=color, linestyle=:dash)
    plot!([(right,offset),(right,w_right)],
        linewidth=1.5, color=color, linestyle=:dash)
    plot!([(μ_bar,offset),(μ_bar,w_μ)],
        linewidth=1.5, color=color, linestyle=:dash)
    plot!([(ω,offset),(ω,w_ω)],
        linewidth=1.5, color=color, linestyle=:dashdot)
    end
end

function makeDistributionPlot!(X, color="blue"; ann=true, offset=0.0, scale=1.0)
    #%% Mode, Mean, HDIs
    ω = mean(find_mode(X))
    μ_bar = mean(X);
    left,right = hdi(X);

    h = fit(Histogram, X, nbins=200)
    w = h.weights      # weights of each bar
    e = h.edges[1]     # edges of the bars (must be +1 more than bars)

    riemannSum = sum(w.*diff(e))
    w = w ./ riemannSum

    e2 = similar([e[:]; e[:]])  # doubling of the edge vector
    e2[2:2:end] = e             # even indices
    e2[1:2:end] = e             # odd indices
    e2 = e2[2:(end-1)]          # pop the ends

    w2 = similar([w[:]; w[:]])  # doubling of the weight vector
    w2[2:2:end] = w             # even indices
    w2[1:2:end] = w             # odd indices
    w2 = w2 .* scale .+ offset

    w_left = w[sum(e .< left)] * scale + offset  # comparison leaves at max left edge of bar that contains the limit
    w_right = w[sum(e .< right)] * scale + offset
    w_ω = w[sum(e .< ω)] * scale + offset
    w_μ = w[sum(e .< μ_bar)] * scale + offset

    #%% Plotting
    plot!(e2, w2,
         linealpha=0.1, linecolor=color,
         fillrange=offset, fillalpha=0.4, fillcolor=color,
         legend=false)

    # or do it by histogram, then we also have vertical lines
    # specify statsplot?
    #histogram!(X, bins=100, normalize=:pdf,
    #          color=color, alpha=0.4, linealpha=0.1, legend=false)   # Comes from StatsPlots now

    if ann
    plot!([(left,offset),(left,w_left)],
        linewidth=1.5, color=color, linestyle=:dash,
        annotations = (left, w_left, text("$(round(left,sigdigits=4))",color,rotation=90,:center,:left)))
    plot!([(right,offset),(right,w_right)],
        linewidth=1.5, color=color, linestyle=:dash,
        annotations = (right, w_right, text("$(round(right,sigdigits=4))",color,rotation=90,:center,:left)))
    plot!([(μ_bar,offset),(μ_bar,w_μ)],
        linewidth=1.5, color=color, linestyle=:dash,
        annotations = (μ_bar, 0.5*w_μ, text("mean $(round(μ_bar,sigdigits=4))",color,rotation=90,:bottom,:right)))
    plot!([(ω,offset),(ω,w_ω)],
        linewidth=1.5, color=color, linestyle=:dashdot,
        annotations = (ω, w_ω, text("mode $(round(ω,sigdigits=4))",color,rotation=90,:bottom,:right)))
    else
    plot!([(left,offset),(left,w_left)],
        linewidth=1.5, color=color, linestyle=:dash)
    plot!([(right,offset),(right,w_right)],
        linewidth=1.5, color=color, linestyle=:dash)
    plot!([(μ_bar,offset),(μ_bar,w_μ)],
        linewidth=1.5, color=color, linestyle=:dash)
    plot!([(ω,offset),(ω,w_ω)],
        linewidth=1.5, color=color, linestyle=:dashdot)
    end
end

#%% Logarithmise and Transform Data ############################################
J = maximum(ind)    # number of individuals
I = length(y)       # number of observations

# Go to log-space:
logy = log.(y)

logMean = mean(logy)
logStd = std(logy)

global trainMean = mean(x)
global trainStd = std(x)

# Mean-centre and scale each element:
zlogy = (logy .- logMean) ./ logStd;

# Standardise training data:
zx = (x .- trainMean) ./ trainStd;

################################################################################
############################# STANS MCMC #######################################

#%% Stan Setup #################################################################
noOfChains = 4
N = 10^6                # total number of samples
N_chain = convert(Int64, N / noOfChains)   # samples per chain
keepchains = false
burnIn = 10^4           # burn-in per chain

modelString = "
data {
    int<lower=1> J;         // number of individuals for which data was observed
    int<lower=1> I;         // observation number in the total vector
    int<lower=1> ID[I];     // individual ID in the total vector
    int<lower=0> K[J];      // input: identifier for whether individual is a kid
    real zx[I];             // (transformed) input: attempt number of the observation
    real zlogy[I];          // (transformed) output: reaction time
}
parameters {
    //real theta0[J];      // intercept for each individual (without reparametrisation)
    //real theta1[J];      // slope for each individual (without reparametrisation)
    real eta0[J];          // helper for reparametrisation
    real eta1[J];          // helper for reparametrisation
    real<lower=0.000001> sigma;    // variation around regression line (same for all)
    real mu0;               // mean for the intercept at group level
    real phi0;              // additive term for the intercept if being a kid
    real mu1;               // mean for the slope at the group level
    real phi1;              // additive term for the slope if being a kid
    real<lower=0> tau0;     // std of the group intercept (same for all individuals)
    real<lower=0> tau1;     // std of the group slope (same for all)
}
transformed parameters {
    real theta0[J];          // intercept for each individual (reparametrised)
    real theta1[J];          // slope for each individual (reparametrised)
    for (j in 1:J) {
        theta0[j] = mu0 + phi0*K[j] + tau0 * eta0[j];
        theta1[j] = mu1 + phi1*K[j] + tau1 * eta1[j];
    }
}
model {
    // tau0 ~ uniform(0,10000);     // if these were used, chain got still stuck
    // tau1 ~ uniform(0,10000);     // if these were used, chain got still stuck
    // using improper prior for taus instead (using lower value in parameter block)

    for (i in 1:I)
        zlogy[i] ~ normal(theta0[ID[i]] + theta1[ID[i]]*zx[i], sigma);
    for (j in 1:J) {
        eta0[j] ~ normal(0, 1);
        eta1[j] ~ normal(0, 1);
    }

    // Without reparametrisation (chain gets stuck):
    //for (j in 1:J) {
    //    theta0[j] ~ normal(mu0 + phi0*K[j], tau0);
    //    theta1[j] ~ normal(mu1 + phi1*K[j], tau1);
    //}

    // no prior is equivalent to a uniform prior
    // however this is an <<improper>> prior and can lead to problems
}
generated quantities {
    // doing this in Julia instead
}
";

#%%### 2) Data for the Stan model, note that variable names must
# correspond to defined model in stan-code
observedData = Dict("I" => I,
                    "J" => J,
                    "zlogy" => zlogy,
                    "ID" => ind,
                    "K" => child_j,
                    "zx" => zx);

### 3) Chain specs
myModel = Stanmodel(
                Sample(save_warmup = false,
                       num_warmup = burnIn,
                       num_samples = N_chain,
                       thin = 1),   # thin: Period between saved samples
                name = "reactionTime-A7",
                model = modelString,
                printsummary = false,
                tmpdir = tmpDir,
                nchains = noOfChains);    # number of chains (4 default)

#%%### 4) Run CmdStan:
# rc:     return code (0 if all is fine)
# chn:    chain results
# cnames: vector of variable names
rc, chn, cnames = stan(myModel,
                       observedData,
                       projDir,
                       diagnostics = false,
                       CmdStanDir = CMDSTAN_HOME);

# Time measures for all together(preprocess, sample, postprocess):
# N=10^4 around 30 seconds.
# N=10^5 around 1.5 minutes
# N=10^6 around 12.5 min, if no chain gets stuck

################################################################################
############################# RESULTS ##########################################

##############
# Acess results and Transform back (undo mean-centering and scaling and go to non-log space)
# See derivation in PDF file instead

# Access the axis array by names:
θ_0 = Array{Float64,2}(undef,N,J)
θ_1 = Array{Float64,2}(undef,N,J)
for j in 1:J
    θ_0[:,j] = 1.0 * chn.value[Axis{:var}("theta0.$j")][:]
    θ_1[:,j] = 1.0 * chn.value[Axis{:var}("theta1.$j")][:]
end
μ_0 = 1.0 * chn.value[Axis{:var}("mu0")][:]
μ_1 = 1.0 * chn.value[Axis{:var}("mu1")][:]
ϕ_0 = 1.0 * chn.value[Axis{:var}("phi0")][:]   # all chains in one sausage
ϕ_1 = 1.0 * chn.value[Axis{:var}("phi1")][:]
σ   = 1.0 * chn.value[Axis{:var}("sigma")][:]
τ_0 = 1.0 * chn.value[Axis{:var}("tau0")][:]
τ_1 = 1.0 * chn.value[Axis{:var}("tau1")][:]

##### RE-READ RESULTS IF CHAIN HAS ALREADY RUN
# only for re-use purposes, please ignore!
using CSV

θ_0 = Array{Float64,2}(undef,N,J)
θ_1 = Array{Float64,2}(undef,N,J)
μ_0 = Array{Float64,1}(undef,N)
μ_1 = Array{Float64,1}(undef,N)
ϕ_0 = Array{Float64,1}(undef,N)
ϕ_1 = Array{Float64,1}(undef,N)
σ   = Array{Float64,1}(undef,N)
τ_0 = Array{Float64,1}(undef,N)
τ_1 = Array{Float64,1}(undef,N)

for i in 1:noOfChains
    # subindices
    from = N_chain*(i-1) + 1
    to = N_chain*i

    # read into a DataFrame:
    chain_csv = CSV.read(tmpDir*"/reactionTime-A7_samples_$i.csv"; comment="#", normalizenames=true)

    θ_0_df = chain_csv[:,r"theta0"]    # regex
    θ_1_df = chain_csv[:,r"theta1"]
    for j in 1:J
        θ_0[from:to,j] = θ_0_df[:,j]  # need to bring on array form, from data frame
        θ_1[from:to,j] = θ_1_df[:,j]  # need to bring on array form, from data frame
    end

    μ_0[from:to,:] = chain_csv[:, r"mu0"][:,1]
    μ_1[from:to,:] = chain_csv[:, r"mu1"][:,1]
    ϕ_0[from:to,:] = chain_csv[:, r"phi0"][:,1]
    ϕ_1[from:to,:] = chain_csv[:, r"phi1"][:,1]
    σ[from:to,:]   = chain_csv[:, r"sigma"][:,1]
    τ_0[from:to,:] = chain_csv[:, r"tau0"][:,1]
    τ_1[from:to,:] = chain_csv[:, r"tau1"][:,1]
end
######

#####
# Un-standardise:
# Individual level:
global θ_0_unscaled = logStd.*θ_0 .- logStd.*θ_1.*trainMean./trainStd .+ logMean
global θ_1_unscaled = logStd.*θ_1./trainStd

global σ_unscaled = σ .* logStd

# Group level (maybe not needed)
global μ_0_unscaled = logStd.*μ_0 .- logStd.*μ_1.*trainMean./trainStd .+ logMean
global ϕ_0_unscaled = logStd.*(ϕ_0 .- ϕ_1.*trainMean./trainStd)

global μ_1_unscaled = logStd.*μ_1./trainStd
global ϕ_1_unscaled = logStd.*ϕ_1./trainStd

global τ_0_unscaled = logStd.* sqrt.(τ_0.^2 .+ τ_1.^2 .* trainMean.^2 ./ trainStd.^2)
global τ_tilde_1_unscaled = sqrt.(2 .* trainMean) .* logStd .* τ_1 ./ trainStd
global τ_1_unscaled = logStd .* τ_1 ./ trainStd

#####
# Get into non-log space (generates function of input):
# Individual level:
function E_y_ind(x,j)
    # transform to zx:
    #zx = (x - trainMean) / trainStd
    unlogExpected = Array{Float64,1}(undef,N)
    unlogExpected = exp.(θ_0_unscaled[:,j] .+
                θ_1_unscaled[:,j] .* x .+
                0.5 .* σ_unscaled.^2);
    return unlogExpected
end

function E_y_group(x,isKid)
    # isKid is either 1 or 0
    # transform to zx:
    #zx = (x - trainMean) / trainStd
    unlogExpected = Array{Float64,1}(undef,N)
    unlogExpected = exp.(μ_0_unscaled + ϕ_0_unscaled .* isKid .+
                    (μ_1_unscaled .+ ϕ_1_unscaled .* isKid) .* x .+
                    0.5 .* (σ_unscaled.^2 .+ τ_0_unscaled.^2 .-
                            τ_tilde_1_unscaled.^2 .* x .+ τ_1_unscaled.^2 .* x^2))
    return unlogExpected
end

function curveSwarm(nrCurves,j,N,α,subpl=0)
    checkpoints = 100
    attempts = range(0,21,length=checkpoints)
    curves = Array{Float64,2}(undef,N,checkpoints)

    for (i, att) in enumerate(attempts)
        curves[:,i] = E_y_ind(att,j)
    end

    if subpl > 0
        plot!(legend=false,ylims=[100,800],subplot=subpl)
    else
        plot(legend=false,ylims=[100,800])
    end
    randIdces = Int.(ceil.(rand(nrCurves).*N))

    for i = 1:nrCurves
        if subpl > 0
            plot!(attempts,curves[randIdces[i],:],
                  color="blue",linealpha=α,subplot=subpl,
                  legend=false,ylims=[100,800])
        else
            plot!(attempts,curves[randIdces[i],:],
                  color="blue",linealpha=α,
                  legend=false,ylims=[100,800])
        end
    end
end

function curveSwarmGroup(nrCurves,isKid,N,α,subpl=0)
    checkpoints = 100
    attempts = range(0,21,length=checkpoints)
    curves = Array{Float64,2}(undef,N,checkpoints)

    for (i, att) in enumerate(attempts)
        curves[:,i] = E_y_group(att,isKid)
    end

    #if subpl > 0
    #    plot(legend=false,ylims=[100,800],subplot=subpl)
    #else
    #    plot(legend=false,ylims=[100,800])
    #end
    randIdces = Int.(ceil.(rand(nrCurves).*N))

    for i = 1:nrCurves
        if subpl > 0
            plot!(attempts,curves[randIdces[i],:],
                  color=RGB(isKid,0,0),linealpha=α,subplot=subpl,
                  legend=false,ylims=[100,800])
        else
            plot!(attempts,curves[randIdces[i],:],
                  color=RGB(isKid,0,0),linealpha=α,
                  legend=false,ylims=[100,800])
        end
    end
end

################################################################################
############################# TASKS ############################################

##### Task 3
for j in [1,3,4]
    ### Diagnostics to compare to Jesper:
    makeDistributionPlot(exp.(θ_1_unscaled[:,j]),"blue")
    plot!(xlabel="exp(theta_1[$j])")
    Plots.savefig(projDir*"/figs/th1$j-Stan.pdf")

    ### Expected reaction times:
    makeDistributionPlot(E_y_ind(1,j),"blue")
    makeDistributionPlot!(E_y_ind(5,j),"red")
    plot!(xlims=[200,800],grid=false,xlabel="E[y]-ind$j")
    Plots.savefig(projDir*"/figs/times-$j-Stan.pdf")

    ### Credible regression lines:
    curveSwarm(200,j,N,.15)
    col = RGB(child_j[j],0,0)
    scatter!(1:1:sum(ind.==j),y[ind.==j],
        markersize=10,markercolor=col,markerstrokecolor=col)
    plot!(grid=false,xlabel="attempt nr",ylabel="reaction time")
    Plots.savefig(projDir*"/figs/swarm-$j-Stan.pdf")
end

### All Regression lines:
plot(layout=(4, 9),size=(1500, 1000),
     legend=false,grid=false,xaxis=false,yaxis=false)
for j in 1:J
    curveSwarm(200,j,N,.15,j)
    if (j-1)%9==0
        myax=true
    else
        myax=false
    end
    col = RGB(child_j[j],0,0)
    scatter!(1:1:sum(ind.==j),y[ind.==j],
        markersize=4,markercolor=col,markerstrokecolor=col,
        subplot=j, ann=(12,700,"j = $j"),xaxis=true,yaxis=myax,
        framestyle=:axes)
end
Plots.savefig(projDir*"/figs/allcurves-Stan.pdf")

### Task 5 (sigma)
# Import sigma from A6
A6_import = CSV.read(projDir*"/export_sigma_A6.csv"; header=false)
σ_A6 = A6_import[:,1]
σ_A6_unscaled = A6_import[:,2]
makeDistributionPlot(σ,"blue")
makeDistributionPlot!(σ_A6,"black")
plot!(grid=false,xlabel=L"\sigma")
Plots.savefig(projDir*"/figs/sigma-comp-Stan.pdf")

##### Diagnostics to compare to Jesper:
makeDistributionPlot(τ_0,"black")
Plots.savefig(projDir*"/figs/tau0-Stan.pdf")
makeDistributionPlot(τ_0_unscaled,"black")
Plots.savefig(projDir*"/figs/tau0-unsc-Stan.pdf")

makeDistributionPlot(σ,"black")
Plots.savefig(projDir*"/figs/sigma-Stan.pdf")
makeDistributionPlot(σ_unscaled,"black")
Plots.savefig(projDir*"/figs/sigma-unsc-Stan.pdf")

##### Swarm of the groups
attempts = range(0,22,length=200)
curve_mean_adult = mean.(E_y_group.(attempts,0))
curve_mean_kid = mean.(E_y_group.(attempts,1))
col = RGB.(child_i,0,0);

plot(attempts,curve_mean_adult,
    color=:black,ylims=[100,800],linewidth=8,linealpha=0.6)
plot!(attempts,curve_mean_kid,
      color=:red,ylims=[100,800],linewidth=8,linealpha=0.6)
curveSwarmGroup(200,0,N,.15)
curveSwarmGroup(200,1,N,.15)
plot!(grid=false,xlabel="attempt nr",ylabel="reaction time")
scatter!(x,y,markersize=2,markercolor=col,markerstrokecolor=col,alpha=0.2)
Plots.savefig(projDir*"/figs/swarm-groups-Stan.pdf")

### Chain diagnostics
diagnostics_csv = CSV.read(tmpDir*"/reactionTime-A7_summary.csv"; comment="#", normalizenames=true)
# only keep the names and NEff without underscore:
diagnostics_csv = diagnostics_csv[[occursin(r"^((?!__).)*$", elementname) for elementname in diagnostics_csv[:, :name]], [:name, :N_Eff]]
diagnostics_csv_θ_0 = diagnostics_csv[[occursin(r"theta0.*", elementname) for elementname in diagnostics_csv[:, :name]], [:name, :N_Eff]]
diagnostics_csv_θ_1 = diagnostics_csv[[occursin(r"theta1.*", elementname) for elementname in diagnostics_csv[:, :name]], [:name, :N_Eff]]
diagnostics_csv_rest = diagnostics_csv[[occursin(r"^((?!theta).)*$", elementname) for elementname in diagnostics_csv[:, :name]], [:name, :N_Eff]]

println("Max N_eff θ_0 is ",
    diagnostics_csv_θ_0[diagnostics_csv_θ_0[:,:N_Eff] .>= maximum(diagnostics_csv_θ_0.N_Eff),:][1,1],
    " = ",
    diagnostics_csv_θ_0[diagnostics_csv_θ_0[:,:N_Eff] .>= maximum(diagnostics_csv_θ_0.N_Eff),:][1,2])
println("Min N_eff θ_0 is ",
    diagnostics_csv_θ_0[diagnostics_csv_θ_0[:,:N_Eff] .<= minimum(diagnostics_csv_θ_0.N_Eff),:][1,1],
    " = ",
    diagnostics_csv_θ_0[diagnostics_csv_θ_0[:,:N_Eff] .<= minimum(diagnostics_csv_θ_0.N_Eff),:][1,2])
println("Max N_eff θ_1 is ",
    diagnostics_csv_θ_1[diagnostics_csv_θ_1[:,:N_Eff] .>= maximum(diagnostics_csv_θ_1.N_Eff),:][1,1],
    " = ",
    diagnostics_csv_θ_1[diagnostics_csv_θ_1[:,:N_Eff] .>= maximum(diagnostics_csv_θ_1.N_Eff),:][1,2])
println("Min N_eff θ_1 is ",
    diagnostics_csv_θ_1[diagnostics_csv_θ_1[:,:N_Eff] .<= minimum(diagnostics_csv_θ_1.N_Eff),:][1,1],
    " = ",
    diagnostics_csv_θ_1[diagnostics_csv_θ_1[:,:N_Eff] .<= minimum(diagnostics_csv_θ_1.N_Eff),:][1,2])

for i = 1:size(diagnostics_csv_rest,1)
    println("N_eff of ",diagnostics_csv_rest[i,1]," = ",diagnostics_csv_rest[i,2])
end

chstart = 650000
chend = 655000
plot(chstart:chend,θ_0[chstart:chend,:],legend=false,xlabel="sample nr",ylabel=L"\theta_{0_j}")
Plots.savefig(projDir*"/figs/unstuckChain-theta0-Stan.pdf")

plot(chstart:chend,θ_1[chstart:chend,:],legend=false,xlabel="sample nr",ylabel=L"\theta_{1_j}")
Plots.savefig(projDir*"/figs/unstuckChain-theta1-Stan.pdf")

plot(chstart:chend,σ[chstart:chend],legend=:topleft,xlabel="sample nr",label=L"\sigma")
plot!(chstart:chend,ϕ_0[chstart:chend],xlabel="sample nr",label=L"\varphi_0")
plot!(chstart:chend,ϕ_1[chstart:chend],xlabel="sample nr",label=L"\varphi_1")
plot!(chstart:chend,μ_0[chstart:chend],xlabel="sample nr",label=L"\mu_0")
plot!(chstart:chend,μ_1[chstart:chend],xlabel="sample nr",label=L"\mu_1")
Plots.savefig(projDir*"/figs/unstuckChain-restPars-Stan.pdf")
plot(chstart:chend,τ_0[chstart:chend],xlabel="sample nr",label=L"\tau_0")
plot!(chstart:chend,τ_1[chstart:chend],xlabel="sample nr",label=L"\tau_1")
Plots.savefig(projDir*"/figs/unstuckChain-tau-Stan.pdf")

makeDistributionPlot(τ_1,"red")
plot!(grid=false,xlabel=L"\tau_1")
Plots.savefig(projDir*"/figs/tau1-Stan.pdf")

# Convergence
chstart = 150000
chend = 160000
plot(grid=false)
for c = 0:noOfChains-1
    plot!(chstart:chend,τ_1[chstart+c*N_chain:chend+c*N_chain],
          xlabel="chain sample nr",
          label=latexstring("\\tau_1, chain_$(c+1)"),linealpha=0.4)
end
Plots.savefig(projDir*"/figs/convergence-Stan.pdf")

##### Useful posteriors
plot(grid=false,xlabel=L"\mu_1")
makeDistributionPlot!(μ_1,"black")
Plots.savefig(projDir*"/figs/mu1-Stan.pdf")
mean(exp.(μ_1_unscaled))

plot(grid=false,xlabel=L"\varphi_1")
makeDistributionPlot!(ϕ_1,"blue")
Plots.savefig(projDir*"/figs/phi1-Stan.pdf")
mean(exp.(ϕ_1_unscaled))

plot(grid=false,xlabel="Reaction time",xlims=(200, 600))
makeDistributionPlot!(E_y_group(1,0),"black",ann=false)
makeDistributionPlot!(E_y_group(1,1),"red",ann=false,scale=1,offset=0)
plot!(ann=(230,.005,"x = 1"))
makeDistributionPlot!(E_y_group(5,0),"black",ann=false,scale=1,offset=0.015)
makeDistributionPlot!(E_y_group(5,1),"red",ann=false,scale=1,offset=0.015)
plot!(ann=(230,.020,"x = 5"))
makeDistributionPlot!(E_y_group(10,0),"black",ann=false,scale=1,offset=0.03)
makeDistributionPlot!(E_y_group(10,1),"red",ann=false,scale=1,offset=0.03)
plot!(ann=(230,.035,"x = 10"))
Plots.savefig(projDir*"/figs/groupTimes-Stan.pdf")

##### Just for own reference and playing around with posterior predictions:
zlogy_pred_kid(xin,Ξ) = mean(μ_0) .+ mean(ϕ_1) .+ (mean(μ_1) + mean(ϕ_1))*((xin - trainMean)/trainStd) .+
                    sqrt(mean(σ)^2 + mean(τ_0)^2 + mean(τ_1)^2 * ((xin - trainMean)/trainStd)^2).*Ξ

zlogy_pred_kid2(xin,Ξ,ix) = μ_0[ix] + ϕ_1[ix] + (μ_1[ix] + ϕ_1[ix])*((xin - trainMean)/trainStd) +
                    sqrt(σ[ix]^2 + τ_0[ix]^2 + τ_1[ix]^2 * ((xin - trainMean)/trainStd)^2) * Ξ

zlogy_pred_adult2(xin,Ξ,ix) = μ_0[ix] + μ_1[ix]*((xin - trainMean)/trainStd) +
                    sqrt(σ[ix]^2 + τ_0[ix]^2 + τ_1[ix]^2 * ((xin - trainMean)/trainStd)^2) * Ξ

atts = collect(1:21)
plot()
for i = 1:500
    Ξ = randn()
    randInd = max(Int(round(rand()*N)),1)
    pKid = exp.((zlogy_pred_kid2.(atts,Ξ,randInd)) .* logStd .+ logMean)
    pAdu = exp.((zlogy_pred_adult2.(atts,Ξ,randInd)) .* logStd .+ logMean)
    plot!(atts,pKid,
          color="red",linealpha=0.2,
          legend=false)
    plot!(atts,pAdu,
          color="black",linealpha=0.2,
          legend=false)
end

plot!()


text:
# Car

function car_dynamics!(ẋ::AbstractVector,x::AbstractVector,u::AbstractVector)
    ẋ[1] = u[1]*cos(x[3])
    ẋ[2] = u[1]*sin(x[3])
    ẋ[3] = u[2]
    return nothing
end
n,m = 3,2

car = Model(car_dynamics!,n,m)

ext: jl

max_issues_repo_name: MehrdadZareian/GMARA

max_forks_repo_name: MehrdadZareian/GMARA


text:
export LightSource, DistantLight, PointLight, DiskLight

################################
# Visibility Tester is a ray from point to light source

"""
    VisibilityTester(p1::Point3, p2::Point3, eps1::Float64, eps2::Float64)

Returns a ray to test intersection on the line segment from `p1` to `p2` with the given ray
epsilon values.

"""
function VisibilityTester(p1::Point3, p2::Point3, eps1::Float64, eps2::Float64)
    distance = norm(p2-p1)
    return Ray(p1, (p2-p1)/distance, eps1, distance*(1.0 - eps2), 1)
end

"""
    VisibilityTester(origin::Point3, direction::Vector3, eps::Float64)

Returns a ray to test intersection on the line segment from `origin` infinitely far in the
given direction, with a given ray epsilon.

"""
function VisibilityTester(origin::Point3, direction::Vector3, eps::Float64)
    return Ray(origin, direction, eps, Inf, 1)
end

################################
# Point light

"""
    PointLight{S<:Spectrum} <: LightSource

Point light source. It's located at a given point and radiates isotropically. The intensity
value given is interpreted as the intensity in Watts/m^2 at a distance of 1 meter.

# Fields

- `position::Point3`, the location of the light source.
- `intensity::S`, the intensity at a distance of 1 meter in Watts/m^2.
- `light_to_world::Transformation`, transformation to shift the location of the light
   source.
- `world_to_light::Transformation`, inverse transformation of the above.

"""
struct PointLight{S<:Spectrum} <: LightSource
    position::Point3
    intensity::S
    light_to_world::Transformation
    world_to_light::Transformation
end

"""
    PointLight(L::Spectrum, l2w::Transformation)

Create a new `PointLight` with the given spectrum, located at the origin but shifted by
the given `Transformation`.

"""
function PointLight(L::Spectrum, l2w::Transformation)
    PointLight(l2w(Point3(0)), L, l2w, inv(l2w))
end

"""
    sample_L(light::PointLight, p::Point3)

Get an intensity sample from the given `DistantLight` and a ray from the light towards the
given point. The sample is always the same since the light source is a uniform parallel
light.

Returns:

- Intensity `Spectrum`
- Direction `Vector3` of the light.
- A weight coefficient which is always 1.0
- A ray from the light source to the given point.

"""
function sample_L(light::PointLight, p::Point3)
    distance = norm(light.position - p)
    L = light.intensity / distance^2
    pdf = 1.0
    vis = VisibilityTester(p, light.position, 2e-5, 0.0)
    return L, pdf, vis
end

"""
    generate_ray(light::PointLight, bounds::BoundingSphere, S::Sample)

Returns a ray from the light source at the given `BoundingSphere`, parametrized by the
given random sample.

"""
function generate_ray(light::PointLight, bounds::BoundingSphere, S::Sample)
    ray = ray_from_point(bounds.radius, S.imgX, S.imgY)
    return light.light_to_world(ray)
end

################################
# Distanct light source

"""
    DistantLight{S<:Spectrum} <: LightSource

A point light source infinitely far away, i.e. parallel rays with a given intensity spectrum
interpreted as Watts / m^2.

By default the light source is "located" infinitely far away along the positive Z axis and
the light rays travel in the -Z direction.

# Fields

- `intensity::S`, the intensity of the light in Watts/m^2.
- `light_to_world::Transformation`, transformation to define the direction of the light
   source.
- `world_to_light::Transformation`, inverse transformation of the above.

"""
struct DistantLight{S<:Spectrum} <: LightSource
    intensity::S
    light_to_world::Transformation
    world_to_light::Transformation
end

"""
    function DistantLight(L::Spectrum, l2w::Transformation)

Create a new `DistantLight` given the intensity spectrum (in units of Watts / m^2) and
a `Transformation` used to rotate the direction of the light.

"""
function DistantLight(L::Spectrum, l2w::Transformation)
    DistantLight(L, l2w, inv(l2w))
end

"""
    sample_L(light::DistantLight, p::Point3)

Get an intensity sample from the given `DistantLight` and a ray from the light towards the
given point. The sample is always the same since the light source is a uniform parallel
light.

Returns:

- Intensity `Spectrum`
- A weight coefficient which is always 1.0
- A ray from the light source to the given point.

"""
function sample_L(light::DistantLight, p::Point3)
    return light.intensity, 1.0, VisibilityTester(p, light.light_to_world(NEG_Z_AXIS), 1e-9)
end

"""
    generate_ray(light::DistantLight, bounds::BoundingSphere, S::Sample)

Returns a ray from the light source at the given `BoundingSphere`, parametrized by the
given random sample.

"""
function generate_ray(light::DistantLight, bounds::BoundingSphere, S::Sample)
    ray = ray_parallel(bounds.radius, S.imgX, S.imgY)
    shift = translation(bounds.center)
    return shift(light.light_to_world(ray))
end

###### ########################
# TODO: Area light
struct AreaLight <: LightSource end

################################
# Disk Light (TODO)

"""
    DiskLight{S<:Spectrum} <: LightSource

A light source which is a uniformly illuminated disk far away. Expressed in terms of its
angular radius. By default it's in the positive Z direction, and can be rotated with
a transformation.

# Fields:

- `radius::Float64`, angular radius (in radians) of the source
- `intensity::S`,
- `light_to_world::Transformation`, transformation to adjust the direction of the source.
   Chosen to rotate the +Z axis into the desired direction of the light.
- `world_to_light::Transformation`, inverse transformation of the above

"""
struct DiskLight{S<:Spectrum} <: LightSource
    radius::Float64
    intensity::S
    light_to_world::Transformation
    world_to_light::Transformation
end

"""
    DiskLight(direction::Vector3, radius::Float64, intensity::Spectrum)

Create new `DiskLight` in the desired direction, angular radius and intensity.

"""
function DiskLight(direction::Vector3, radius::Float64, intensity::Spectrum)
    rot_z = rotate_z_to(direction)
    Ω = 2π * (1 - cos(radius))
    return DiskLight(radius, Ω.*intensity, rot_z, inv(rot_z))
end

@inline solid_angle(D::DiskLight) = 2π * (1 - cos(D.radius))

sample_L(light::DiskLight, p::Point3) = sample_L(light, p, rand(), rand())

function sample_L(light::DiskLight, p::Point3, u1::Real, u2::Real)
    direction = light.light_to_world(disk_light_sample(light.radius, u1, u2))
    return light.intensity, 1.0, VisibilityTester(p, direction, 2e-9)
end

"""
    disk_light_sample(radius::Real, u1::Real, u2::Real)

Create a vector which points from origin to a point on the surface of a disk with given
angular radius centered on the +Z axis. It is parametrized by two numbers, `u1` and `u2`,
which range from zero to one. Using uniform random numbers will result in the surface of
the disk to be uniformly sampled.

"""
function disk_light_sample(radius::Real, u1::Real, u2::Real)
    cos_th = cos(radius)
    z = cos_th + (1.0 - cos_th) * u1
    phi = 2π * u2
    r = sqrt(1.0 - z^2)
    return Vector3(r*cos(phi), r*sin(phi), z)
end

"""
    generate_ray(light::DiskLight, bounds::BoundingSphere, S::Sample)

Returns a ray from the light source at the given `BoundingSphere`, parametrized by the
given random sample.

"""
function generate_ray(light::DiskLight, bounds::BoundingSphere, S::Sample)
    # Pick a direction, i.e. a point on the (infinitely distant) disk
    # and make the transformation to rotate Z into that direction
    V = disk_light_sample(light.radius, S.lensU, S.lensV)
    rot = light.light_to_world * rotate_z_to(V)

    # Make a ray from infinitely far along the Z axis and rotate it into the direction
    # given by T. Also shift it relative to the origin of the bounding sphere.

    ray = ray_parallel(bounds.radius, S.imgX, S.imgY)
    T = translation(bounds.center) * rot
    return T(ray)
end

max_stars_repo_name: dronir/GRAYTR

max_issues_repo_name: dronir/GRAYTR

max_forks_repo_name: dronir/GRAYTR


text:
fast_flags(fast::Bool) = fast ? "nsz arcp contract afn reassoc" : "nsz contract"
# fast_flags(fast::Bool) = fast ? "fast" : "nsz contract"

const LLVM_TYPES = IdDict{Type{<:NativeTypes},String}(
  Float16 => "half",
  Float32 => "float",
  Float64 => "double",
  Bit => "i1",
  Bool => "i8",
  Int8 => "i8",
  UInt8 => "i8",
  Int16 => "i16",
  UInt16 => "i16",
  Int32 => "i32",
  UInt32 => "i32",
  Int64 => "i64",
  UInt64 => "i64",
  # Int128 => "i128",
  # UInt128 => "i128",
  # UInt256 => "i256",
  # UInt512 => "i512",
  # UInt1024 => "i1024",
)
const JULIA_TYPES = IdDict{Type{<:NativeTypes},Symbol}(
  Float16 => :Float16,
  Float32 => :Float32,
  Float64 => :Float64,
  Int8 => :Int8,
  Int16 => :Int16,
  Int32 => :Int32,
  Int64 => :Int64,
  # Int128 => :Int128,
  UInt8 => :UInt8,
  UInt16 => :UInt16,
  UInt32 => :UInt32,
  UInt64 => :UInt64,
  # UInt128 => :UInt128,
  Bool => :Bool,
  Bit => :Bit,
  # UInt256 => :UInt256,
  # UInt512 => :UInt512,
  # UInt1024 => :UInt1024,
)
const LLVM_TYPES_SYM = IdDict{Symbol,String}(
  :Float16 => "half",
  :Float32 => "float",
  :Float64 => "double",
  :Int8 => "i8",
  :Int16 => "i16",
  :Int32 => "i32",
  :Int64 => "i64",
  # :Int128 => "i128",
  :UInt8 => "i8",
  :UInt16 => "i16",
  :UInt32 => "i32",
  :UInt64 => "i64",
  # :UInt128 => "i128",
  :Bool => "i8",
  :Bit => "i1",
  :Nothing => "void",
  # :UInt256 => "i256",
  # :UInt512 => "i512",
  # :UInt1024 => "i1024",
)
const TYPE_LOOKUP = IdDict{Symbol,Type{<:NativeTypes}}(
  :Float16 => Float16,
  :Float32 => Float32,
  :Float64 => Float64,
  :Int8 => Int8,
  :Int16 => Int16,
  :Int32 => Int32,
  :Int64 => Int64,
  # :Int128 => Int128,
  :UInt8 => UInt8,
  :UInt16 => UInt16,
  :UInt32 => UInt32,
  :UInt64 => UInt64,
  # :UInt128 => UInt128,
  :Bool => Bool,
  :Bit => Bit,
  # :UInt256 => UInt256,
  # :UInt512 => UInt512,
  # :UInt1024 => UInt1024
)
const JULIA_TYPE_SIZE = IdDict{Symbol,Int}(
  :Float16 => 2,
  :Float32 => 4,
  :Float64 => 8,
  :Int8 => 1,
  :Int16 => 2,
  :Int32 => 4,
  :Int64 => 8,
  # :Int128 => 16,
  :UInt8 => 1,
  :UInt16 => 2,
  :UInt32 => 4,
  :UInt64 => 8,
  # :UInt128 => 16,
  :Bool => 1,
  :Bit => 1,
  # :UInt256 => 32,
  # :UInt512 => 64,
  # :UInt1024 => 128,
)

function _get_alignment(W::Int, sym::Symbol)::Int
    sym === :Bit && return 1
    T = TYPE_LOOKUP[sym]
    if W > 1
        Base.datatype_alignment(_Vec{W,T})
    else
        Base.datatype_alignment(T)
    end
end

const JULIAPOINTERTYPE = 'i' * string(8sizeof(Int))

vtype(W, typ::String) = isone(abs(W)) ? typ : "<$W x $typ>"
vtype(W, T::DataType) = vtype(W, LLVM_TYPES[T])
vtype(W, T::Symbol) = vtype(W, get(LLVM_TYPES_SYM, T, T))
function push_julia_type!(x, W, T)
    if W ≤ 1
        push!(x, T)
        nothing
    else
        push!(x, Expr(:curly, :_Vec, W, T))
        nothing
    end
end
function append_julia_type!(x, Ws, Ts)
    for i ∈ eachindex(Ws)
        push_julia_type!(x, Ws[i], Ts[i])
    end
end

ptr_suffix(T) = "p0" * suffix(T)
ptr_suffix(W, T) = suffix(W, ptr_suffix(T))
suffix(W::Int, s::String) = W == -1 ? s : 'v' * string(W) * s
suffix(W::Int, T) = suffix(W, suffix(T))
suffix(::Type{Ptr{T}}) where {T} = "p0" * suffix(T)
suffix_jlsym(W::Int, s::Symbol) = suffix(W, suffix(T))
function suffix(T::Symbol)::String
    if T === :Float64
        "f64"
    elseif T === :Float32
        "f32"
    else
        string('i', 8JULIA_TYPE_SIZE[T])
    end
end
suffix(@nospecialize(T))::String = suffix(JULIA_TYPES[T])

# Type-dependent LLVM constants
function llvmconst(T, val)::String
    iszero(val) ? "zeroinitializer" : "$(LLVM_TYPES[T]) $val"
end
function llvmconst(::Type{Bool}, val)::String
    Bool(val) ? "i1 1" : "zeroinitializer"
end
function llvmconst(W::Int, @nospecialize(T), val)::String
    isa(val, Number) && iszero(val) && return "zeroinitializer"
    typ = (LLVM_TYPES[T])::String
    '<' * join(("$typ $(val)" for _ in Base.OneTo(W)), ", ") * '>'
end
function llvmconst(W::Int, ::Type{Bool}, val)::String
    Bool(val) ? '<' * join(("i1 $(Int(val))" for _ in Base.OneTo(W)), ", ") * '>' : "zeroinitializer"
end
function llvmconst(W::Int, v::String)::String
    '<' * join((v for _ in Base.OneTo(W)), ", ") * '>'
end
# function llvmtypedconst(T, val)
#     typ = LLVM_TYPES[T]
#     iszero(val) && return "$typ zeroinitializer"
#     "$typ $val"
# end
# function llvmtypedconst(::Type{Bool}, val)
#     Bool(val) ? "i1 1" : "i1 zeroinitializer"
# end
function _llvmcall_expr(ff, WR, R, argt)
    if WR ≤ 1
        Expr(:call, :ccall, ff, :llvmcall, R, argt)
    else
        Expr(:call, :ccall, ff, :llvmcall, Expr(:curly, :_Vec, WR, R), argt)
    end
end

function llvmname(op::String, WR::Int, WA, T::Symbol, TA::Symbol)
    lret = LLVM_TYPES_SYM[T]
    ln = WR ≤ 1 ? "llvm.$op" : "llvm.$op.$(suffix(WR,T))"
    (isone(abs(WR)) || T !== TA) ? ln * '.' * suffix(maximum(WA),TA) : ln
end

function build_llvmcall_expr(op, WR, R::Symbol, WA, TA, ::Nothing = nothing)
    ff = llvmname(op, WR, WA, R, first(TA))
    argt = Expr(:tuple)
    append_julia_type!(argt.args, WA, TA)
    call = _llvmcall_expr(ff, WR, R, argt)
    for n ∈ eachindex(TA)
        push!(call.args, Expr(:call, :data, Symbol(:v, n)))
    end
    Expr(:block, Expr(:meta, :inline), isone(abs(WR)) ? call : Expr(:call, :Vec, call))
end
function build_llvmcall_expr(op, WR, R::Symbol, WA, TA, flags::String)
    lret = LLVM_TYPES_SYM[R]
    lvret = vtype(WR, lret)
    lop = llvmname(op, WR, WA, R, first(TA))
    instr = "$lvret $flags @$lop"
    larg_types = map(vtype, WA, TA)
    decl = "declare $lvret @$(lop)(" * join(larg_types, ", ") * ')'
    args_for_call = ("$T %$(n-1)" for (n,T) ∈ enumerate(larg_types))
    instrs = """%res = call $flags $lvret @$(lop)($(join(args_for_call, ", ")))
        ret $lvret %res"""
    args = Expr(:curly, :Tuple);
    append_julia_type!(args.args, WA, TA)
    arg_syms = Vector{Expr}(undef, length(TA))
    for n ∈ eachindex(TA)
        arg_syms[n] = Expr(:call, :data, Symbol(:v,n))
    end
    # println(instrs)
    if WR ≤ 1
        llvmcall_expr(decl, instrs, R, args, lvret, larg_types, arg_syms)
    else
        llvmcall_expr(decl, instrs, Expr(:curly, :_Vec, WR, R), args, lvret, larg_types, arg_syms)
    end
end

@static if VERSION ≥ v"1.6.0-DEV.674"
    function llvmcall_expr(
        decl::String, instr::String, ret::Union{Symbol,Expr}, args::Expr, lret::String, largs::Vector{String}, arg_syms::Vector, callonly::Bool = false
    )
        mod = """
            $decl

            define $lret @entry($(join(largs, ", "))) alwaysinline {
            top:
                $instr
            }
        """
        # attributes #0 = { alwaysinline }
        call = Expr(:call, LLVMCALL, (mod::String, "entry")::Tuple{String,String}, ret, args)
        for arg ∈ arg_syms
            push!(call.args, arg)
        end
        call = Expr(:(::), call, ret)
        if first(lret) === '<'
            call = Expr(:call, :Vec, call)
        end
        callonly && return call
        Expr(:block, Expr(:meta, :inline), call)
        # Expr(:block, Expr(:meta, :inline), )
    end
else
    function llvmcall_expr(
        decl::String, instr::String, ret::Union{Symbol,Expr}, args::Expr, lret::String, largs::Vector{String}, arg_syms::Vector, callonly::Bool = false
    )
        call = Expr(:call, LLVMCALL, (decl, instr), ret, args)
        foreach(arg -> push!(call.args, arg), arg_syms)
        if first(lret) === '<'
            call = Expr(:call, :Vec, call)
        end
        callonly && return call
        Expr(:block, Expr(:meta, :inline), call)
    end
end


text:
#=
# function utils
=#
export GenFunction, stringToFunction, add_genfunction, makeFunction, makeFunctions
# export @stringToFunction, @makeFunction, @makeFunctions

# Stores everything we could need about a generated function.
struct GenFunction      # example:
    name::String        # "genfunction_7"
    args::String        # "x"
    str::String         # "sin(x)"
    expr                # expression: sin(x) NOTE: could be an Expr, Symbol, Number,
    func                # handle for: genfunction_7(x) = sin(x)
end

# Generates a function from a string
function stringToFunction(name, args, fun)
    return eval(Meta.parse(name*"("*args*")="*fun));
end

# Adds the GenFunction to a global array of generated functions
function add_genfunction(genfun)
    global genfunc_count += 1;
    global genfunctions = [genfunctions ; genfun];
    log_entry("Generated function: "*genfun.name);
end

# Makes a GenFunction and adds it to the
# args is a string like "x,y,z"
# fun is a string like "sin(x)*y + 3*z"
function makeFunction(args, fun)
    name = "genfunction_"*string(Femshop.genfunc_count);
    if typeof(fun) == Expr
        ex = fun;
        strfun = string(fun);
        nf = GenFunction(name, args, strfun, ex, stringToFunction(name, args, strfun));
    else
        ex = Meta.parse(fun);
        nf = GenFunction(name, args, fun, ex, stringToFunction(name, args, fun));
    end
    add_genfunction(nf);
end

# Makes either: a constant number, a genfunction, or an array of genfunctions
function makeFunctions(ex; args="x=0,y=0,z=0,t=0,node_index=1")
    nfuns = 0;
    if typeof(ex) <: Array
        for i=1:length(ex)
            if typeof(ex[i]) == String
                tmp = makeFunction(args, ex[i]);
                nfuns += 1;
            # else # It could be a number or function handle
            end
        end
    else
        if typeof(ex) == String
            makeFunction(args, ex);
            nfuns = 1;
        # else # It could be a number or function handle
        end
    end
    return nfuns;
end

### Old macro versions below. keep for now.

# # Generates a function
# macro stringToFunction(name, args, fun)
#     return esc(quote
#         eval(Meta.parse($name*"("*$args*")="*$fun));
#     end)
# end

# # Makes a GenFunction and adds it to the
# # args is a string like "x,y,z"
# # fun is a string like "sin(x)*y + 3*z"
# macro makeFunction(args, fun)
#     return esc(quote
#         let
#             name = "genfunction_"*string(Femshop.genfunc_count);
#             if typeof($fun) == Expr
#                 ex = $fun;
#                 strfun = string($fun);
#                 nf = GenFunction(name, $args, strfun, ex, @stringToFunction(name, $args, strfun));
#             else
#                 ex = Meta.parse($fun);
#                 nf = GenFunction(name, $args, $fun, ex, @stringToFunction(name, $args, $fun));
#             end
#             add_genfunction(nf);
#         end
#     end)
# end

# # Makes either: a constant number, a genfunction, or an array of genfunctions
# macro makeFunctions(ex)
#     return esc(quote
#         nfuns = 0;
#         args = "x=0,y=0,z=0,t=0";
#         if typeof($ex) <: Array
#             ### If it's an array, it can't be mixed type! At least for now.
#             # There's a problem with nfuns disappearing in the loop.
#             if typeof($ex[1]) <: Number
#                 nfuns = 0;
#             else
#                 for i=1:length($ex)
#                     if typeof($ex[i]) == String
#                         tmp = @makeFunction(args, $ex[i]);
#                     end
#                 end
#                 nfuns = length($ex)
#             end
#         else
#             if typeof($ex) == String
#                 @makeFunction(args, $ex);
#                 nfuns = 1;
#             elseif typeof($ex) <: Number
#                 nfuns = 0;
#             end
#         end
#         nfuns
#     end)
# end


text:
using Pkg
Pkg.activate(".")
Pkg.add(["Markdown",
        "InteractiveUtils",
        "Plots",
        "Pluto",
        "PlutoUI",
        "LaTeXStrings",
        "LiteQTL",
        "DelimitedFiles",
        "CSV",
        "DataFrames",
        "ColorSchemes",
        "Colors"])

max_stars_repo_name: senresearch/GNDemo

max_issues_repo_name: senresearch/GNDemo

max_forks_repo_name: senresearch/GNDemo


text:
# load and plot the imbibition data
include("rel_perms.jl")
include("imbibition_IMPES.jl")

plot(size=(500,400), xtickfont = font(10, "Courier"), ytickfont=font(10, "Courier"),
  xlabel="t [day]", ylabel="Recovery Factor", legendfont=font(10, "Courier"),
  guidefont=font(12, "Courier"))
for SF in 0.1:0.2:0.9
  (t,R)=imb_impes(SF)
  save("imb_$SF-ww.jld", "t", t, "R", R)
  plot!(t/(24*3600), R, xlims=(0.0,50.0), ylims=(0, 1), linewidth=2,
  label="$SF Sal. Factor")
  savefig("imb_recovery_$SF-WW.png")
end
plot!()

max_stars_repo_name: simulkade/peteng

max_issues_repo_name: simulkade/peteng

max_forks_repo_name: simulkade/peteng


text:
int_rules_2_1 = @theory begin

    #= ::Subsection::Closed:: =#
    #= 2.1*(c+d*x)^m*(a+b*(F^(g*(e+f*x)))^n)^p =#
    $UseGamma = False;
    @apply_utils Antiderivative((~(c') + ~(d') * ~x) ^ ~(m') * (~(b') * (~F) ^ (~(g') * (~(e') + ~(f') * ~x))) ^ ~(n'), ~x) => ((~c + ~d * ~x) ^ ~m * (~b * (~F) ^ (~g * (~e + ~f * ~x))) ^ ~n) / (~f * ~g * ~n * log(~F)) - ((~d * ~m) / (~f * ~g * ~n * log(~F))) * Antiderivative((~c + ~d * ~x) ^ (~m - 1) * (~b * (~F) ^ (~g * (~e + ~f * ~x))) ^ ~n, ~x) <-- FreeQ([~F, ~b, ~c, ~d, ~e, ~f, ~g, ~n], ~x) && (GtQ(~m, 0) && (IntegerQ(2 * ~m) && Not(TrueQ($UseGamma))))
    @apply_utils Antiderivative((~(c') + ~(d') * ~x) ^ ~m * (~(b') * (~F) ^ (~(g') * (~(e') + ~(f') * ~x))) ^ ~(n'), ~x) => ((~c + ~d * ~x) ^ (~m + 1) * (~b * (~F) ^ (~g * (~e + ~f * ~x))) ^ ~n) / (~d * (~m + 1)) - ((~f * ~g * ~n * log(~F)) / (~d * (~m + 1))) * Antiderivative((~c + ~d * ~x) ^ (~m + 1) * (~b * (~F) ^ (~g * (~e + ~f * ~x))) ^ ~n, ~x) <-- FreeQ([~F, ~b, ~c, ~d, ~e, ~f, ~g, ~n], ~x) && (LtQ(~m, -1) && (IntegerQ(2 * ~m) && Not(TrueQ($UseGamma))))
    @apply_utils Antiderivative((~F) ^ (~(g') * (~(e') + ~(f') * ~x)) / (~(c') + ~(d') * ~x), ~x) => ((~F) ^ (~g * (~e - (~c * ~f) / ~d)) / ~d) * ExpIntegralEi((~f * ~g * (~c + ~d * ~x) * log(~F)) / ~d) <-- FreeQ([~F, ~c, ~d, ~e, ~f, ~g], ~x) && Not(TrueQ($UseGamma))
    @apply_utils Antiderivative((~(c') + ~(d') * ~x) ^ ~(m') * (~F) ^ (~(g') * (~(e') + ~(f') * ~x)), ~x) => (((-(~d)) ^ ~m * (~F) ^ (~g * (~e - (~c * ~f) / ~d))) / ((~f) ^ (~m + 1) * (~g) ^ (~m + 1) * log(~F) ^ (~m + 1))) * Gamma(~m + 1, ((-(~f) * ~g * log(~F)) / ~d) * (~c + ~d * ~x)) <-- FreeQ([~F, ~c, ~d, ~e, ~f, ~g], ~x) && IntegerQ(~m)
    @apply_utils Antiderivative((~F) ^ (~(g') * (~(e') + ~(f') * ~x)) / sqrt(~(c') + ~(d') * ~x), ~x) => (2 / ~d) * Subst(Antiderivative((~F) ^ (~g * (~e - (~c * ~f) / ~d) + (~f * ~g * (~x) ^ 2) / ~d), ~x), ~x, sqrt(~c + ~d * ~x)) <-- FreeQ([~F, ~c, ~d, ~e, ~f, ~g], ~x) && Not(TrueQ($UseGamma))
    @apply_utils Antiderivative((~(c') + ~(d') * ~x) ^ ~m * (~F) ^ (~(g') * (~(e') + ~(f') * ~x)), ~x) => ((-((~F) ^ (~g * (~e - (~c * ~f) / ~d))) * (~c + ~d * ~x) ^ FracPart(~m)) / (~d * ((-(~f) * ~g * log(~F)) / ~d) ^ (IntPart(~m) + 1) * ((-(~f) * ~g * log(~F) * (~c + ~d * ~x)) / ~d) ^ FracPart(~m))) * Gamma(~m + 1, ((-(~f) * ~g * log(~F)) / ~d) * (~c + ~d * ~x)) <-- FreeQ([~F, ~c, ~d, ~e, ~f, ~g, ~m], ~x) && Not(IntegerQ(~m))
    @apply_utils Antiderivative((~(c') + ~(d') * ~x) ^ ~(m') * (~(b') * (~F) ^ (~(g') * (~(e') + ~(f') * ~x))) ^ ~n, ~x) => ((~b * (~F) ^ (~g * (~e + ~f * ~x))) ^ ~n / (~F) ^ (~g * ~n * (~e + ~f * ~x))) * Antiderivative((~c + ~d * ~x) ^ ~m * (~F) ^ (~g * ~n * (~e + ~f * ~x)), ~x) <-- FreeQ([~F, ~b, ~c, ~d, ~e, ~f, ~g, ~m, ~n], ~x)
    @apply_utils Antiderivative((~(c') + ~(d') * ~x) ^ ~(m') * (~a + ~(b') * ((~F) ^ (~(g') * (~(e') + ~(f') * ~x))) ^ ~(n')) ^ ~(p'), ~x) => Antiderivative(ExpandIntegrand((~c + ~d * ~x) ^ ~m, (~a + ~b * ((~F) ^ (~g * (~e + ~f * ~x))) ^ ~n) ^ ~p, ~x), ~x) <-- FreeQ([~F, ~a, ~b, ~c, ~d, ~e, ~f, ~g, ~m, ~n], ~x) && IGtQ(~p, 0)
    @apply_utils Antiderivative((~(c') + ~(d') * ~x) ^ ~(m') / (~a + ~(b') * ((~F) ^ (~(g') * (~(e') + ~(f') * ~x))) ^ ~(n')), ~x) => (~c + ~d * ~x) ^ (~m + 1) / (~a * ~d * (~m + 1)) - (~b / ~a) * Antiderivative(((~c + ~d * ~x) ^ ~m * ((~F) ^ (~g * (~e + ~f * ~x))) ^ ~n) / (~a + ~b * ((~F) ^ (~g * (~e + ~f * ~x))) ^ ~n), ~x) <-- FreeQ([~F, ~a, ~b, ~c, ~d, ~e, ~f, ~g, ~n], ~x) && IGtQ(~m, 0)
    #= Antiderivative(((~c')+(~d')*(~x))^(~m')/((~a)+(~b')*((~F)^((~g')*((~e')+(~f')*(~x))))^(~n')),~x)  := -(c+d*x)^m/(a*f*g*n*log(F))*log(1+a/(b*(F^(g*(e+f*x)))^n)) + d*m/(a*f*g*n*log(F))*Antiderivative((c+d*x)^(m-1)*log(1+a/(b*(F^(g*(e+f*x)))^n) ),x) <-- FreeQ([F,a,b,c,d,e,f,g,n],x) && IGtQ(m,0) =#
    @apply_utils Antiderivative((~(c') + ~(d') * ~x) ^ ~(m') * (~a + ~(b') * ((~F) ^ (~(g') * (~(e') + ~(f') * ~x))) ^ ~(n')) ^ ~p, ~x) => (1 / ~a) * Antiderivative((~c + ~d * ~x) ^ ~m * (~a + ~b * ((~F) ^ (~g * (~e + ~f * ~x))) ^ ~n) ^ (~p + 1), ~x) - (~b / ~a) * Antiderivative((~c + ~d * ~x) ^ ~m * ((~F) ^ (~g * (~e + ~f * ~x))) ^ ~n * (~a + ~b * ((~F) ^ (~g * (~e + ~f * ~x))) ^ ~n) ^ ~p, ~x) <-- FreeQ([~F, ~a, ~b, ~c, ~d, ~e, ~f, ~g, ~n], ~x) && (ILtQ(~p, 0) && IGtQ(~m, 0))
    @apply_utils Antiderivative((~(c') + ~(d') * ~x) ^ ~(m') * (~a + ~(b') * ((~F) ^ (~(g') * (~(e') + ~(f') * ~x))) ^ ~(n')) ^ ~p, ~x) => With([u = IntHide((~a + ~b * ((~F) ^ (~g * (~e + ~f * ~x))) ^ ~n) ^ ~p, ~x)], Dist((~c + ~d * ~x) ^ ~m, u, ~x) - ~d * ~m * Antiderivative((~c + ~d * ~x) ^ (~m - 1) * u, ~x)) <-- FreeQ([~F, ~a, ~b, ~c, ~d, ~e, ~f, ~g, ~n], ~x) && (IGtQ(~m, 0) && LtQ(~p, -1))
    @apply_utils Antiderivative((~u) ^ ~(m') * (~(a') + ~(b') * ((~F) ^ (~(g') * ~v)) ^ ~(n')) ^ ~(p'), ~x) => Antiderivative(NormalizePowerOfLinear(~u, ~x) ^ ~m * (~a + ~b * ((~F) ^ (~g * ExpandToSum(~v, ~x))) ^ ~n) ^ ~p, ~x) <-- FreeQ([~F, ~a, ~b, ~g, ~n, ~p], ~x) && (LinearQ(~v, ~x) && (PowerOfLinearQ(~u, ~x) && (Not(LinearMatchQ(~v, ~x) && PowerOfLinearMatchQ(~u, ~x)) && IntegerQ(~m))))
    @apply_utils Antiderivative((~u) ^ ~(m') * (~(a') + ~(b') * ((~F) ^ (~(g') * ~v)) ^ ~(n')) ^ ~(p'), ~x) => Module([uu = NormalizePowerOfLinear(~u, ~x), z], z = If(PowerQ(uu) && FreeQ(uu(2), ~x), uu(1) ^ (~m * uu(2)), uu ^ ~m); (uu ^ ~m / z) * Antiderivative(z * (~a + ~b * ((~F) ^ (~g * ExpandToSum(~v, ~x))) ^ ~n) ^ ~p, ~x)) <-- FreeQ([~F, ~a, ~b, ~g, ~m, ~n, ~p], ~x) && (LinearQ(~v, ~x) && (PowerOfLinearQ(~u, ~x) && (Not(LinearMatchQ(~v, ~x) && PowerOfLinearMatchQ(~u, ~x)) && Not(IntegerQ(~m)))))
    @apply_utils Antiderivative((~(c') + ~(d') * ~x) ^ ~(m') * (~a + ~(b') * ((~F) ^ (~(g') * (~(e') + ~(f') * ~x))) ^ ~(n')) ^ ~(p'), ~x) => Unintegrable((~c + ~d * ~x) ^ ~m * (~a + ~b * ((~F) ^ (~g * (~e + ~f * ~x))) ^ ~n) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~e, ~f, ~g, ~m, ~n, ~p], ~x)
end

max_stars_repo_path: src/rules/2 Exponentials/2.1 (c+d x)^m (a+b (F^(g (e+f x)))^n)^p.jl


text:
module BNP_WMReg_Joint

using Dates
using LinearAlgebra
using SpecialFunctions

using Distributions
using PDMats
using StatsBase
using Roots

using BayesInference # personal package

include("general.jl")
include("mcmc.jl")
include("sqfChol.jl")
include("update_alloc.jl")
include("update_weights.jl")
include("update_eta_Met.jl")
include("update_G0.jl")
include("update_variable_selection_local.jl")
include("update_variable_selection_global.jl")
include("densityEstimation.jl")

end


text:
# Author: Gabrio Rizzuti, rizzuti.gabrio@gatech.edu
# Date: August, 2020
# Copyright: Georgia Institute of Technology, 2020
#
# Run FWI on Sleipner data with added random noise

## Module loading

using DrWatson
@quickactivate "UncertaintyQuantificationAVP"
using UncertaintyQuantificationAVP
using LinearAlgebra, PyPlot, Statistics, Optim
using Random; seed = 1; Random.seed!(seed)

## Loading synthetic data and inversion parameters

# Data
data_dict = wload(datadir("Sleipner", "Sleipner_data.bson"))
dat_n = data_dict[:dat_n]
src_pars = data_dict[:src_pars]
abc_geom = data_dict[:abc_geom]
acc = data_dict[:acc_mod]

# Model
model_dict = wload(projectdir()*"/scripts/Sleipner/Sleipner_model.bson")
Mtrue = model_dict[:Mtrue]
Mbg = model_dict[:Mbg]

# Inversion parameters
invpars_dict = wload(datadir("Sleipner", "Sleipner_invpars.bson"))
Mpr = invpars_dict[:Mpr]
M0  = invpars_dict[:M0]
σ2_n = invpars_dict[:σ2_n]
ϵ2 = invpars_dict[:ϵ2]
σ2_pr = invpars_dict[:σ2_pr]
mmin = invpars_dict[:mmin]
mmax = invpars_dict[:mmax]

# MAP result
Mmap = wload(datadir("Sleipner", "Sleipner_MAP.bson"))[:Mmap]

## Loss function

# Pre- and post-processing for log-posterior
nz = Mbg.geom.nz
dz = Mbg.geom.dz
function proj_bounds(m::Array{Float32}, mmin, mmax)
    m_ = copy(m)
    m_[m_ .< mmin] .= mmin
    m_[m_ .> mmax] .= mmax
    return m_
end
prec(m::Array{Float32, 3}) = Model(Mbg.geom, proj_bounds(reshape(m, nz, :), mmin, mmax))
post(g::Model) = reshape(g.m, (nz, 1, :))

# Log-likelihood
F = seismod_fun(dat_n.pars, src_pars; abc_geom = abc_geom, acc = acc)
negloglike_fun(m::Model, dat::Data; compute_grad::Bool = false) = negLogLikelihood_gauss(dat, m, σ2_n, F; compute_grad = compute_grad)

# Log-prior
D = derivative_linop(Mbg.geom)
# neglogpr_fun(m::Model, Mpr::Model; compute_grad::Bool = false) = negLogPrior_L1smooth(m, σ2_pr, ϵ2; m0 = Mpr, D = D, compute_grad = compute_grad)
neglogpr_fun(m::Model, Mnoise::Model; compute_grad::Bool = false) = negLogPrior_gauss(m, σ2_pr; b = Mnoise, m0 = Mpr, D = D, compute_grad = compute_grad)

# Log-posterior
neglogpost_fun(m::Model, dat::Data, Mnoise::Model; compute_grad::Bool = false) = negLogPost(m, (m; compute_grad=compute_grad)->negloglike_fun(m, dat; compute_grad=compute_grad), (m; compute_grad=compute_grad)->neglogpr_fun(m, Mnoise; compute_grad=compute_grad); compute_grad = compute_grad)

# Overloading log-posterior for tensor-type input+pre-/post-conditioning
function neglogpost_fun(m::Array{Float32, 3}, dat::Data, Mnoise::Model; compute_grad::Bool = false)
    nb = size(m, 3)
    if compute_grad
        f, g_ = neglogpost_fun(prec(m), dat, Mnoise; compute_grad = compute_grad)
        return sum(f)/nb, post(g_)
    else
        f = neglogpost_fun(prec(m), dat, Mnoise; compute_grad = compute_grad)
        return sum(f)/nb
    end
end
function neglogpost_fun!(F, G::Union{Nothing, Array{Float32, 3}}, m::Array{Float32, 3}, dat::Data, Mnoise::Model)
    if G == nothing
        f = neglogpost_fun(m, dat, Mnoise; compute_grad = false)
    else
        f, g = neglogpost_fun(m, dat, Mnoise; compute_grad = true)
        G .= g
    end
    return f
end

## Optimization

M0 = deepcopy(Mbg)
method = LBFGS()
niter = 100
optimopt = Optim.Options(iterations = niter, store_trace = true, show_trace = true)#, show_every = 1)

n_noise = 2^7
noise_d = sqrt(σ2_n)*randn(Float32, size(dat_n.d)..., n_noise)
noise_m = sqrt(σ2_pr)*randn(Float32, size(Mpr.m)..., n_noise)
Mrto = Array{Float32, 2}(undef, nz, n_noise)
fval = Array{Array{Float32, 1}, 1}(undef, n_noise)
time = @elapsed Threads.@threads for i = 1:n_noise
    println(i, "/", n_noise)
    neglogpost_fun_n!(F, G, m) = neglogpost_fun!(F, G, m, dat_n+noise_d[:, :, :, i], Model(Mbg.geom, noise_m[:, :, i]))
    result = optimize(Optim.only_fg!(neglogpost_fun_n!), reshape(M0.m, :, 1, 1), method, optimopt)
    Mrto[:, i] = prec(Optim.minimizer(result)).m
    fval[i] = Optim.f_trace(result)
end
println(time)

## Saving results

results_dict = Dict(:noise_d => noise_d, :noise_m => noise_m, :Mrto => Mrto, :fval => fval)
wsave(datadir("Sleipner", "Sleipner_RTO.bson"), results_dict)


text:
# This file establishes the default initial conditions, boundary conditions and sources
# for the baroclinicwave_problem experiment, following [Ullrich2014](@cite)

# Override default CLIMAParameters for consistency with literature on this case
CLIMAParameters.Planet.press_triple(::EarthParameterSet) = 610.78

struct BaroclinicWaveProblem{BCS, ISP, ISA, WP, BS, MP} <: AbstractAtmosProblem
    boundaryconditions::BCS
    init_state_prognostic::ISP
    init_state_auxiliary::ISA
    perturbation::WP
    base_state::BS
    moisture_profile::MP
end
function BaroclinicWaveProblem(
    physics::AtmosPhysics;
    boundaryconditions = (AtmosBC(physics;), AtmosBC(physics;)),
    perturbation = nothing,
    base_state = nothing,
    moisture_profile = nothing,
)
    # Set up defaults
    if isnothing(perturbation)
        perturbation = DeterministicPerturbation()
    end
    if isnothing(base_state)
        base_state = BCWaveBaseState()
    end
    if isnothing(moisture_profile)
        moisture_profile = MoistLowTropicsMoistureProfile()
    end

    problem = (
        boundaryconditions,
        init_gcm_experiment!,
        (_...) -> nothing,
        perturbation,
        base_state,
        moisture_profile,
    )
    return BaroclinicWaveProblem{typeof.(problem)...}(problem...)
end

problem_name(::BaroclinicWaveProblem) = "BaroclinicWave"

setup_source(::BaroclinicWaveProblem) = (Gravity(), Coriolis())


text:
analog_input_configs = Dict{AbstractString,Number}(
    "referenced single-ended" => Val_RSE,
    "non-referenced single-ended" => Val_NRSE,
    "differential" => Val_Diff )

"""
`analog_input(channel, config, range) -> task`

`analog_input(task, channel, config, range)`

create an analog input channel, and a new task if one is not specified
"""
function analog_input(channel::ASCIIString; terminal_config="differential", range=nothing)
    t = AITask()
    analog_input(t, channel, terminal_config=terminal_config, range=range)
    t
end

function analog_input(t::AITask, channel::ASCIIString;
        terminal_config="differential", range=nothing)
    if range==nothing
        device::ASCIIString = split(channel,'/')[1]
        range=float(analog_input_ranges(device)[end,:])
    end
    catch_error( CreateAIVoltageChan(t.th,
            pointer(channel),
            pointer(""),
            analog_input_configs[terminal_config],
            range[1], range[2],
            Val_Volts,
            convert(Ptr{UInt8},C_NULL)) )
    nothing
end

"""
`analog_output(channel, range) -> task`

`analog_output(task, channel, range)`

create an analog output channel, and a new NIDAQ task if one is not specified
"""
function analog_output(channel::ASCIIString; range=nothing)
    t = AOTask()
    analog_output(t, channel, range=range)
    t
end

function analog_output(t::AOTask, channel::ASCIIString; range=nothing)
    if range==nothing
        device::ASCIIString = split(channel,'/')[1]
        range=float(analog_output_ranges(device)[end,:])
    end
    catch_error( CreateAOVoltageChan(t.th,
            pointer(channel),
            pointer(""),
            range[1], range[2],
            Val_Volts,
            convert(Ptr{UInt8},C_NULL)) )
    nothing
end

read_analog_cfunctions = Dict{Type,Function}(
    Float64 => ReadAnalogF64,
    Int16 => ReadBinaryI16,
    Int32 => ReadBinaryI32,
    UInt16 => ReadBinaryU16,
    UInt32 => ReadBinaryU32 )

function read(t::AITask, precision::DataType, num_samples_per_chan::Integer = -1)
    num_channels = getproperties(t)["NumChans"][1]
    num_samples_per_chan_read = Int32[0]
    if num_samples_per_chan==-1;  num_samples_per_chan=1024;  end
    data = Array(precision, num_samples_per_chan*num_channels)
    catch_error( read_analog_cfunctions[precision](t.th,
        convert(Int32,num_samples_per_chan),
        1.0,
        reinterpret(Bool32,Val_GroupByChannel),
        pointer(data),
        convert(UInt32,num_samples_per_chan*num_channels),
        pointer(num_samples_per_chan_read),
        reinterpret(Ptr{Bool32},C_NULL)) )
    data = data[1:num_samples_per_chan_read[1]*num_channels]
    num_channels==1 ? data : reshape(data, (num_samples_per_chan, convert(Int64,num_channels)))
end

for (cfunction, types) in (
        (WriteAnalogF64, Float64),
        (WriteBinaryI16, Int16),
        (WriteBinaryI32, Int32),
        (WriteBinaryU16, UInt16),
        (WriteBinaryU32, UInt32))
    @eval function write(t::AOTask, data::Matrix{$types})
        num_samples_per_chan::Int32 = size(data, 1)
        data = reshape(data, length(data))
        num_samples_per_chan_written = Int32[0]
        catch_error( $cfunction(t.th,
            num_samples_per_chan,
            reinterpret(Bool32, UInt32(false)),
            1.0,
            reinterpret(Bool32,Val_GroupByChannel),
            pointer(data),
            pointer(num_samples_per_chan_written),
            reinterpret(Ptr{Bool32},C_NULL)) )
        num_samples_per_chan_written[1]
    end
    @eval write(t::AOTask, data::Vector{$types}) =
        write(t, reshape(data,(length(data),1)))
end


text:
module Missings

export allowmissing, disallowmissing, ismissing, missing, missings,
       Missing, MissingException, levels, coalesce, passmissing, nonmissingtype,
       skipmissings

using Base: ismissing, missing, Missing, MissingException

@static if VERSION < v"1.3.0-alpha.121"
    nonmissingtype(::Type{S}) where {S} = Core.Compiler.typesubtract(S, Missing)
else
    using Base: nonmissingtype
end

import DataAPI

# vector constructors
missings(dims::Dims) = fill(missing, dims)
missings(::Type{T}, dims::Dims) where {T >: Missing} = fill!(Array{T}(undef, dims), missing)
missings(::Type{T}, dims::Dims) where {T} = fill!(Array{Union{T, Missing}}(undef, dims), missing)
missings(dims::Integer...) = missings(dims)
missings(::Type{T}, dims::Integer...) where {T} = missings(T, dims)

"""
    allowmissing(x::AbstractArray)

Return an array equal to `x` allowing for [`missing`](@ref) values,
i.e. with an element type equal to `Union{eltype(x), Missing}`.

When possible, the result will share memory with `x` (as with [`convert`](@ref)).

See also: [`disallowmissing`](@ref)
"""
allowmissing(x::AbstractArray{T}) where {T} = convert(AbstractArray{Union{T, Missing}}, x)

"""
    disallowmissing(x::AbstractArray)

Return an array equal to `x` not allowing for [`missing`](@ref) values,
i.e. with an element type equal to `nonmissingtype(eltype(x))`.

When possible, the result will share memory with `x` (as with [`convert`](@ref)).
If `x` contains missing values, a `MethodError` is thrown.

See also: [`allowmissing`](@ref)
"""
disallowmissing(x::AbstractArray{T}) where {T} = convert(AbstractArray{nonmissingtype(T)}, x)

# Iterators
"""
    Missings.replace(itr, replacement)

Return an iterator over the elements in `itr` which replaces [`missing`](@ref) values with
`replacement`.

When applicable, the size of `itr` is preserved.
The returned object can be indexed using indices of `itr` if the latter is indexable,
and `eachindex` and `keys` return the indices of `itr`.

If the type of `replacement` differs from the element type of `itr`,
it will be converted to it.

See also: [`skipmissing`](@ref), [`Missings.fail`](@ref)

# Examples
```jldoctest
julia> collect(Missings.replace([1, missing, 2], 0))
3-element Array{Int64,1}:
 1
 0
 2

julia> collect(Missings.replace([1 missing; 2 missing], 0))
2×2 Array{Int64,2}:
 1  0
 2  0

```
"""
replace(itr, replacement) = EachReplaceMissing(itr, convert(eltype(itr), replacement))
struct EachReplaceMissing{T, U}
    x::T
    replacement::U
end
Base.IteratorSize(::Type{<:EachReplaceMissing{T}}) where {T} =
    Base.IteratorSize(T)
Base.IteratorEltype(::Type{<:EachReplaceMissing{T}}) where {T} =
    Base.IteratorEltype(T)
Base.length(itr::EachReplaceMissing) = length(itr.x)
Base.size(itr::EachReplaceMissing) = size(itr.x)
Base.axes(itr::EachReplaceMissing) = axes(itr.x)
Base.eltype(itr::EachReplaceMissing) = nonmissingtype(eltype(itr.x))
Base.eachindex(itr::EachReplaceMissing) = eachindex(itr.x)
Base.keys(itr::EachReplaceMissing) = keys(itr.x)

@inline function Base.iterate(itr::EachReplaceMissing)
    st = iterate(itr.x)
    st === nothing && return nothing
    v, s = st
    return (v isa Missing ? itr.replacement : v, s)
end

@inline function Base.iterate(itr::EachReplaceMissing, state)
    st = iterate(itr.x, state)
    st === nothing && return nothing
    v, s = st
    return (v isa Missing ? itr.replacement : v, s)
end

Base.@propagate_inbounds function Base.getindex(itr::EachReplaceMissing, I...)
    v = itr.x[I...]
    return v === missing ? itr.replacement : v
end

"""
    Missings.fail(itr)

Return an iterator over the elements in `itr` which will throw a [`MissingException`](@ref)
if a [`missing`](@ref) value is found.
The returned object can be indexed using indices of `itr` if the latter is indexable.
Indices corresponding to missing values are not valid: even though they are not skipped
by `keys` and eachindex, a `MissingException` is thrown when trying to use them.

Use [`collect`](@ref) to obtain an `Array` containing the resulting values.
If `itr` is an array, the resulting array will have the same dimensions.

See also: [`skipmissing`](@ref), [`Missings.replace`](@ref)

# Examples
```jldoctest
julia> collect(Missings.fail([1 2; 3 4]))
2×2 Array{Int64,2}:
 1  2
 3  4

julia> collect(Missings.fail([1, missing, 2]))
ERROR: MissingException: missing value encountered by Missings.fail
[...]
```
"""
fail(itr) = EachFailMissing(itr)
struct EachFailMissing{T}
    x::T
end
Base.IteratorSize(::Type{EachFailMissing{T}}) where {T} =
    Base.IteratorSize(T)
Base.IteratorEltype(::Type{EachFailMissing{T}}) where {T} =
    Base.IteratorEltype(T)
Base.length(itr::EachFailMissing) = length(itr.x)
Base.size(itr::EachFailMissing) = size(itr.x)
Base.axes(itr::EachFailMissing) = axes(itr.x)
Base.eltype(itr::EachFailMissing) = nonmissingtype(eltype(itr.x))
Base.eachindex(itr::EachFailMissing) = eachindex(itr.x)
Base.keys(itr::EachFailMissing) = keys(itr.x)

@inline function Base.iterate(itr::EachFailMissing)
    st = iterate(itr.x)
    st === nothing && return nothing
    v, s = st
    ismissing(v) && throw(MissingException("missing value encountered by Missings.fail"))
    return (v::eltype(itr), s)
end

@inline function Base.iterate(itr::EachFailMissing, state)
    st = iterate(itr.x, state)
    st === nothing && return nothing
    v, s = st
    ismissing(v) && throw(MissingException("missing value encountered by Missings.fail"))
    return (v::eltype(itr), s)
end

Base.@propagate_inbounds function Base.getindex(itr::EachFailMissing, I...)
    v = itr.x[I...]
    v === missing && throw(MissingException("the value at index $I is missing"))
    return v
end

const levels = DataAPI.levels

struct PassMissing{F} <: Function
    f::F
end

function (f::PassMissing{F})(x) where {F}
    if @generated
        return x === Missing ? missing : :(f.f(x))
    else
        return x === missing ? missing : f.f(x)
    end
end

function (f::PassMissing{F})(xs...) where {F}
    if @generated
        for T in xs
            T === Missing && return missing
        end
        return :(f.f(xs...))
    else
        return any(ismissing, xs) ? missing : f.f(xs...)
    end
end

"""
    passmissing(f)

Return a function that returns `missing` if any of its positional arguments
are `missing` (even if their number or type is not consistent with any of the
methods defined for `f`) and otherwise applies `f` to these arguments.

`passmissing` does not support passing keyword arguments to the `f` function.

# Examples
```jldoctest
julia> passmissing(sqrt)(4)
2.0

julia> passmissing(sqrt)(missing)
missing

julia> passmissing(sqrt).([missing, 4])
2-element Array{Union{Missing, Float64},1}:
  missing
   2.0

julia> passmissing((x,y)->"\$x \$y")(1, 2)
"1 2"

julia> passmissing((x,y)->"\$x \$y")(missing)
missing
```
"""
passmissing(f) = PassMissing{Core.Typeof(f)}(f)

"""
   skipmissings(args...)

Return a tuple of iterators wrapping each of the iterators in `args`, but
skipping elements at positions where at least one of the iterators returns `missing`
(listwise deletion of missing values).

# Examples
```
julia> x = [1, 2, missing, 4]; y = [1, 2, 3, missing];

julia> tx, ty = skipmissings(x, y)
(Missings.SkipMissings{Array{Union{Missing, Int64},1},Tuple{Array{Union{Missing, Int64},1}}}
(Union{Missing, Int64}[1, 2, missing, 4], (Union{Missing, Int64}[1, 2, 3, missing],)), Missi
ngs.SkipMissings{Array{Union{Missing, Int64},1},Tuple{Array{Union{Missing, Int64},1}}}(Union
{Missing, Int64}[1, 2, 3, missing], (Union{Missing, Int64}[1, 2, missing, 4],)))

julia> collect(tx)
2-element Array{Int64,1}:
 1
 2

```
"""
function skipmissings(args...)
    Base.depwarn("Current design of skipmissings is deprecated. In future releases " *
                 "this function may be redesigned or removed", :skipmissings)
    if isempty(args)
        throw(ArgumentError("Must input one or more arguments"))
    end

    if args isa Tuple{Vararg{AbstractArray}}
        if !all(x -> length(x) == length(args[1]), args)
            throw(ArgumentError("All arguments must have the same length"))
        end

        if !all(x -> eachindex(x) == eachindex(args[1]), args)
            throw(ArgumentError("All arguments must have the same indices"))
        end
    end

    ntuple(length(args)) do i
        s = setdiff(1:length(args), i)
        SkipMissings(args[i], args[s])
    end
end

struct SkipMissings{V, T}
    x::V
    others::T
end

Base.@propagate_inbounds function _anymissingindex(others::Tuple{Vararg{AbstractArray}}, i)
   for oth in others
        oth[i] === missing && return true
    end

    return false
end

@inline function _anymissingiterate(others::Tuple, state)
    for oth in others
        y = iterate(oth, state)
        y !== nothing && first(y) === missing && return true
    end

    return false
end

const SkipMissingsofArrays = SkipMissings{V, T} where
    {V <: AbstractArray, T <: Tuple{Vararg{AbstractArray}}}

function Base.show(io::IO, mime::MIME"text/plain", itr::SkipMissings{V}) where V
    print(io, SkipMissings, '{', V, '}', '(', itr.x, ')', " comprised of " *
          "$(length(itr.others) + 1) iterators")
end

Base.IteratorSize(::Type{<:SkipMissings}) = Base.SizeUnknown()
Base.IteratorEltype(::Type{<:SkipMissings{V}}) where {V} = Base.IteratorEltype(V)
Base.eltype(::Type{<:SkipMissings{V}}) where {V} = nonmissingtype(eltype(V))
Base.IndexStyle(itr::SkipMissings) = Base.IndexStyle(itr.x)

function Base.iterate(itr::SkipMissings, state=1)
    x_itr = iterate(itr.x, state)
    x_itr === nothing && return nothing
    x_item, x_state = x_itr
    while true
        x_item === missing || _anymissingiterate(itr.others, state) || break
        x_itr = iterate(itr.x, x_state)
        x_itr === nothing && return nothing
        state = x_state
        x_item, x_state = x_itr
    end
    return x_item, x_state
end

function Base.iterate(itr::SkipMissingsofArrays, state=0)
    eix = eachindex(itr.x)
    ind_itr = iterate(eix, state)
    ind_itr === nothing && return nothing
    ind_item, ind_state = ind_itr
    @inbounds x_item = itr.x[ind_item]
    @inbounds while true
        x_item === missing || _anymissingindex(itr.others, ind_item) || break
        ind_itr = iterate(eix, ind_state)
        ind_itr === nothing && return nothing
        ind_item, ind_state = ind_itr
        x_item = itr.x[ind_item]
    end
    return x_item, ind_state
end

Base.IndexStyle(::Type{<:SkipMissings{V}}) where {V} = Base.IndexStyle(V)

function Base.eachindex(itr::SkipMissingsofArrays)
    @inbounds Iterators.filter(eachindex(itr.x)) do i
        itr.x[i] !== missing && !_anymissingindex(itr.others, i)
    end
end

function Base.keys(itr::SkipMissingsofArrays)
    @inbounds Iterators.filter(keys(itr.x)) do i
        itr.x[i] !== missing && !_anymissingindex(itr.others, i)
    end
end

@inline function Base.getindex(itr::SkipMissingsofArrays, i)
    @boundscheck checkbounds(itr.x, i)
    @inbounds xi = itr.x[i]
    if xi === missing || @inbounds _anymissingindex(itr.others, i)
        throw(MissingException("the value at index $i is missing for some element"))
    end
    return xi
end

Base.mapreduce(f, op, itr::SkipMissingsofArrays) =
    Base._mapreduce(f, op, Base.IndexStyle(itr), itr)

function Base._mapreduce(f, op, ::IndexLinear, itr::SkipMissingsofArrays)
    A = itr.x
    local ai
    inds = LinearIndices(A)
    i = first(inds)
    ilast = last(inds)
    @inbounds while i <= ilast
        ai = A[i]
        ai === missing || _anymissingindex(itr.others, i) || break
        i += 1
    end
    i > ilast && return Base.mapreduce_empty(f, op, Base.eltype(itr))
    a1::eltype(itr.x) = ai
    i += 1
    @inbounds while i <= ilast
        ai = A[i]
        ai === missing || _anymissingindex(itr.others, i) || break
        i += 1
    end
    i > ilast && return Base.mapreduce_first(f, op, a1)
    # We know A contains at least two non-missing entries: the result cannot be nothing
    something(Base.mapreduce_impl(f, op, itr, first(inds), last(inds)))
end

Base._mapreduce(f, op, ::IndexCartesian, itr::SkipMissingsofArrays) = mapfoldl(f, op, itr)

Base.mapreduce_impl(f, op, A::SkipMissingsofArrays, ifirst::Integer, ilast::Integer) =
    Base.mapreduce_impl(f, op, A, ifirst, ilast, Base.pairwise_blocksize(f, op))

# Returns nothing when the input contains only missing values, and Some(x) otherwise
@noinline function Base.mapreduce_impl(f, op, itr::SkipMissingsofArrays,
                                       ifirst::Integer, ilast::Integer, blksize::Int)
    A = itr.x
    if ifirst == ilast
        @inbounds a1 = A[ifirst]
        if a1 === missing
            return nothing
        elseif _anymissingindex(itr.others, ifirst)
            return nothing
        else
            return Some(Base.mapreduce_first(f, op, a1))
        end
    elseif ifirst + blksize > ilast
        # sequential portion
        local ai
        i = ifirst
        @inbounds while i <= ilast
            ai = A[i]
            ai === missing || _anymissingindex(itr.others, i) || break
            i += 1
        end
        i > ilast && return nothing
        a1 = ai::eltype(itr)
        i += 1
        @inbounds while i <= ilast
            ai = A[i]
            ai === missing || _anymissingindex(itr.others, i) || break
            i += 1
        end
        i > ilast && return Some(Base.mapreduce_first(f, op, a1))
        a2 = ai::eltype(itr)
        i += 1
        v = op(f(a1), f(a2))
        @simd for i = i:ilast
            @inbounds ai = A[i]
            ai === missing || @inbounds _anymissingindex(itr.others, i) || (v = op(v, f(ai)))
        end
        return Some(v)
    else
        # pairwise portion
        imid = (ifirst + ilast) >> 1
        v1 = Base.mapreduce_impl(f, op, itr, ifirst, imid, blksize)
        v2 = Base.mapreduce_impl(f, op, itr, imid+1, ilast, blksize)
        if v1 === nothing && v2 === nothing
            return nothing
        elseif v1 === nothing
            return v2
        elseif v2 === nothing
            return v1
        else
            return Some(op(something(v1), something(v2)))
        end
    end
end

"""
    filter(f, itr::SkipMissings)

Return a vector similar to the array wrapped by the given `SkipMissings` iterator
but skipping all elements with a `missing` value in one of the iterators passed
to `skipmissing` and elements for which `f` returns `false`. This method
only applies when all iterators passed to `skipmissings` are arrays.

# Examples
```
julia> x = [missing; 2:9]; y = [1:9; missing];

julia> mx, my = skipmissings(x, y);

julia> filter(isodd, mx)
4-element Array{Int64,1}:
 3
 5
 7
 9

```
"""
function filter(f, itr::SkipMissingsofArrays)
    x = itr.x
    y = similar(x, eltype(itr), 0)
    for i in eachindex(x)
        @inbounds xi = x[i]
        if xi !== missing && @inbounds !_anymissingindex(itr.others, i) && f(xi)
            push!(y, xi)
        end
    end
    y
end

end # module


text:
# Portions translated from SLICOT-Reference distribution
# Copyright (c) 2002-2020 NICONET e.V.
function run_mb4dpz(datfile, io=stdout)
    NIN = 5
    NOUT = 6
    NMAX = 10
    LDA  = NMAX
    LDC = NMAX
    LDDE = NMAX
    LDVW = NMAX
    A = Array{ComplexF64,2}(undef, LDA,NMAX)
    C = Array{ComplexF64,2}(undef, LDC,NMAX)
    DE = Array{ComplexF64,2}(undef, LDDE,NMAX)
    VW = Array{ComplexF64,2}(undef, LDVW,NMAX)
    DWORK = Array{Float64,1}(undef, 8*NMAX)
    LSCALE = Array{Float64,1}(undef, NMAX)
    RSCALE = Array{Float64,1}(undef, NMAX)
    f = open(datfile,"r")
    readline(f)
    vs = split(readline(f))
    N = parse(BlasInt, vs[1])
    JOB =  vs[2][1]
    THRESH = parse(Float64, replace(vs[3],'D'=>'E'))
    if ( N<=0 || N>NMAX )
        @error "Illegal N=$N"
    end

    vs = String[]
    _isz,_jsz = (N,N)
    while length(vs) < _isz*_jsz
        append!(vs, replace.(split(readline(f)),'D'=>'E'))
    end
    for i in 1:_isz
       _i0 = (i-1)*_jsz
       A[i,1:_jsz] .= parsex.(ComplexF64, vs[_i0+1:_i0+_jsz])
    end

    vs = String[]
    _isz,_jsz = (N,N+1)
    while length(vs) < _isz*_jsz
        append!(vs, replace.(split(readline(f)),'D'=>'E'))
    end
    for i in 1:_isz
       _i0 = (i-1)*_jsz
       DE[i,1:_jsz] .= parsex.(ComplexF64, vs[_i0+1:_i0+_jsz])
    end

    vs = String[]
    _isz,_jsz = (N,N)
    while length(vs) < _isz*_jsz
        append!(vs, replace.(split(readline(f)),'D'=>'E'))
    end
    for i in 1:_isz
       _i0 = (i-1)*_jsz
       C[i,1:_jsz] .= parsex.(ComplexF64, vs[_i0+1:_i0+_jsz])
    end

    vs = String[]
    _isz,_jsz = (N,N+1)
    while length(vs) < _isz*_jsz
        append!(vs, replace.(split(readline(f)),'D'=>'E'))
    end
    for i in 1:_isz
       _i0 = (i-1)*_jsz
       VW[i,1:_jsz] .= parsex.(ComplexF64, vs[_i0+1:_i0+_jsz])
    end
    close(f)
    # interp call 1

    ILO, INFO, IWARN = SLICOT.mb4dpz!( JOB, N, THRESH, A, DE, C, VW, LSCALE, RSCALE, DWORK)
    @test INFO == 0
    INFO == 0 || return

    # interp output 1
    println(io, "A:")
    _nc = N
    _nr = N
    show(io, "text/plain", A[1:_nr,1:_nc])
    println(io)

# interp output 2
    println(io, "DE:")
    _nc = N+1
    _nr = N
    show(io, "text/plain", DE[1:_nr,1:_nc])
    println(io)

# interp output 3
    println(io, "C:")
    _nc = N
    _nr = N
    show(io, "text/plain", C[1:_nr,1:_nc])
    println(io)

# interp output 4
    println(io, "VW:")
    _nc = N+1
    _nr = N
    show(io, "text/plain", VW[1:_nr,1:_nc])
    println(io)

    println(io, "ILO = $ILO")
# interp output 5
    println(io, "LSCALE:")
    _nr = N
    show(io, "text/plain", LSCALE[1:_nr])
    println(io)

# interp output 6
    println(io, "RSCALE:")
    _nr = N
    show(io, "text/plain", RSCALE[1:_nr])
    println(io)

    if ( LSAME( JOB, 'S' ) || LSAME( JOB, 'B' ) )
        if ( !( THRESH==-2 || THRESH==-4 ) )
            # interp output 5
            println(io, "initial 1-norms:")
            _nr = 2
            show(io, "text/plain", DWORK[1:_nr])
            println(io)

            println(io, "final 1-norms:")
            show(io, "text/plain", DWORK[3:4])
            println(io)

# interp output 7
            println(io, "final threshold:")
            show(io, "text/plain", DWORK[5])
            println(io)
        else
            @warn "mb4dlz returned iwarn=$IWARN"
        end
    end
end # run_mb4dpz()


text:
#' ---
#' title: Network Flows
#' ---

#' **Originally Contributed by**: Arpit Bhatia

#' In graph theory, a flow network (also known as a transportation network) is a directed graph where
#' each edge has a capacity and each edge receives a flow. The amount of flow on an edge cannot exceed the capacity of the edge.
#' Often in operations research, a directed graph is called a network, the vertices are called nodes and the edges are called arcs.
#' A flow must satisfy the restriction that the amount of flow into a node equals the amount of flow out of it,
#' unless it is a source, which has only outgoing flow, or sink, which has only incoming flow.
#' A network can be used to model traffic in a computer network, circulation with demands, fluids in pipes,
#' currents in an electrical circuit, or anything similar in which something travels through a network of nodes.

using JuMP
using GLPK

#' ## The Shortest Path Problem
#' Suppose that each arc $(i, j)$ of a graph is assigned a scalar cost $a_{i,j}$, and
#' suppose that we define the cost of a forward path to be the sum of the costs of its arcs.
#' Given a pair of nodes, the shortest path problem is to find a forward path that connects these nodes and has minimum cost.

#' $$
#' \begin{align*}
#' \min && \sum_{\forall e(i,j) \in E} a_{i,j} \times x_{i,j} \\
#' s.t. && b(i) = \sum_j x_{ij} - \sum_k x_{ki} = \begin{cases}
#' 1 &\mbox{if $i$ is the starting node,} \\
#' -1 &\mbox{if $i$ is the ending node,} \\
#' 0 &\mbox{otherwise.} \end{cases} \\
#' && x_{e} \in \{0,1\} && \forall e \in E
#' \end{align*}
#' $$

#'

G = [
0 100 30  0  0;
0   0 20  0  0;
0   0  0 10 60;
0  15  0  0 50;
0   0  0  0  0
]

n = size(G)[1]

shortest_path = Model(with_optimizer(GLPK.Optimizer))

@variable(shortest_path, x[1:n,1:n], Bin)
@constraint(shortest_path, [i = 1:n, j = 1:n; G[i,j] == 0], x[i,j] == 0) # Arcs with zero cost are not a part of the path as they do no exist
@constraint(shortest_path, [i = 1:n; i != 1 && i != 2], sum(x[i,:]) == sum(x[:,i])) # Flow conservation constraint
@constraint(shortest_path, sum(x[1,:]) - sum(x[:,1]) == 1) # Flow coming out of source = 1
@constraint(shortest_path, sum(x[2,:]) - sum(x[:,2]) == -1) # Flowing coming out of destination = -1 i.e. Flow entering destination = 1
@objective(shortest_path, Min, sum(G .* x))

optimize!(shortest_path)
@show objective_value(shortest_path);
@show value.(x);

#' ## The Assignment Problem
#' Suppose that there are $n$ persons and $n$ objects that we have to match on a one-to-one basis.
#' There is a benefit or value $a_{i,j}$ for matching person $i$ with object $j$, and
#' we want to assign persons to objects so as to maximize the total benefit.
#' There is also a restriction that person $i$ can be assigned to object $j$ only if $(i, j)$ belongs to a given set of pairs $A$.
#' Mathematically, we want to find a set of person-object pairs $(1, j_{1}),..., (n, j_{n})$ from $A$ such that
#' the objects $j_{1},...,j_{n}$ are all distinct, and the total benefit $\sum_{i=1}^{y} a_{ij_{i}}$ is maximized.

#' $$
#' \begin{align*}
#' \max && \sum_{(i,j) \in A} a_{i,j} \times y_{i,j} \\
#' s.t. && \sum_{\{j|(i,j) \in A\}} y_{i,j} = 1 && \forall i = \{1,2....n\} \\
#' && \sum_{\{i|(i,j) \in A\}} y_{i,j} = 1 && \forall j = \{1,2....n\} \\
#' && y_{i,j} \in \{0,1\} && \forall (i,j) \in \{1,2...k\}
#' \end{align*}
#' $$

#'

G = [
6 4 5 0;
0 3 6 0;
5 0 4 3;
7 5 5 5;
]

n = size(G)[1]

assignment = Model(with_optimizer(GLPK.Optimizer))
@variable(assignment, y[1:n,1:n], Bin)
@constraint(assignment, [i = 1:n], sum(y[:,i]) == 1) # One person can only be assigned to one object
@constraint(assignment, [j = 1:n], sum(y[j,:]) == 1) # One object can only be assigned to one person
@objective(assignment, Max, sum(G .* y))

optimize!(assignment)
@show objective_value(assignment);
@show value.(y);

#' ## The Max-Flow Problem
#' In the max-flow problem, we have a graph with two special nodes: the $source$, denoted by $s$, and the $sink$, denoted by $t$.
#' The objective is to move as much flow as possible from $s$ into $t$ while observing the capacity constraints.

#' $$
#' \begin{align*}
#' \max && \sum_{v:(s,v) \in E} f(s,v) \\
#' s.t. && \sum_{u:(u,v) \in E} f(u,v)  = \sum_{w:(v,w) \in E} f(v,w) && \forall v \in V - \{s,t\} \\
#' && f(u,v) \leq c(u,v) && \forall (u,v) \in E \\
#' && f(u,v) \geq 0 && \forall (u,v) \in E
#' \end{align*}
#' $$

#'

G = [
0 3 2 2 0 0 0 0
0 0 0 0 5 1 0 0
0 0 0 0 1 3 1 0
0 0 0 0 0 1 0 0
0 0 0 0 0 0 0 4
0 0 0 0 0 0 0 2
0 0 0 0 0 0 0 4
0 0 0 0 0 0 0 0
]

n = size(G)[1]

max_flow = Model(with_optimizer(GLPK.Optimizer))

@variable(max_flow, f[1:n,1:n] >= 0)
@constraint(max_flow, [i = 1:n, j = 1:n], f[i,j] <= G[i,j]) # Capacity constraints
@constraint(max_flow, [i = 1:n; i != 1 && i != 8], sum(f[i,:]) == sum(f[:,i])) # Flow conservation contraints
@objective(max_flow, Max, sum(f[1, :]))

optimize!(max_flow)
@show objective_value(max_flow);
@show value.(f);


text:
using ComplexMixtures
using PDBTools

# Load system PDB file
system = readPDB("./simulation/equilibrated.pdb")

# The complete trajectory file can be downloaded from:
# https://drive.google.com/file/d/1BuXJ8AjBeduMSD2CkDJLDNxAAD2QNNg6/view?usp=sharing
trajectory_file = "../trajectories/traj_Glyc.dcd"

# Select the atoms corresponding to glycerol and water (using PDBTools)
glyc = select(system,"resname GLLM")
water = select(system,"water")

# Compute Glycerol-Glycerol auto correlation mddf
solute = Selection(glyc,natomspermol=14)
traj = Trajectory(trajectory_file,solute) # solute and solvent are the same

# We define a large solute domain (large dbulk) to obtain a good convergence
# for the KB integral. The mddf converges at much shorter distances.
opt = ComplexMixtures.Options(dbulk=20.)
mddf_glyc = mddf(traj,opt)

# Save results for later analysis
save(mddf_glyc,"./results/mddf_glyc.json")

# Compute water-glycerol mddf
solvent = Selection(water,natomspermol=3)
traj = Trajectory(trajectory_file,solute,solvent)
opt = ComplexMixtures.Options(dbulk=20.)
mddf_water_glyc = mddf(traj,opt)

# Save results for later analysis
save(mddf_glyc,"./results/mddf_glyc.json")


text:
using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:3] <= 1)

                     @variable(m, -5.992552298143351 <= q <= 4.750344613921084)

                     add_NL_constraint(m, :(softplus(0.47930083070023954 + -0.2811056889538768*softplus(-0.30287094621097754 + 0.8508304098923269*$(x[1]) + 0.544182872577263*$(x[2]) + -0.2945971031198069*$(x[3])) + -0.8918376604733185*softplus(0.6507439169566887 + 0.9833163152014732*$(x[1]) + -0.36891084266703134*$(x[2]) + -0.5752518327899794*$(x[3])) + -0.8083246377220759*softplus(-0.1240101222843979 + -0.304115262711357*$(x[1]) + -0.19393989104697695*$(x[2]) + -0.50508829642924*$(x[3]))) + softplus(-0.740550943643274 + -0.007926005977438866*softplus(-0.30287094621097754 + 0.8508304098923269*$(x[1]) + 0.544182872577263*$(x[2]) + -0.2945971031198069*$(x[3])) + 0.328386397895855*softplus(0.6507439169566887 + 0.9833163152014732*$(x[1]) + -0.36891084266703134*$(x[2]) + -0.5752518327899794*$(x[3])) + 0.6266373732109729*softplus(-0.1240101222843979 + -0.304115262711357*$(x[1]) + -0.19393989104697695*$(x[2]) + -0.50508829642924*$(x[3]))) + softplus(-0.3322291763851517 + -0.3879865017411772*softplus(-0.30287094621097754 + 0.8508304098923269*$(x[1]) + 0.544182872577263*$(x[2]) + -0.2945971031198069*$(x[3])) + 0.13951112377587105*softplus(0.6507439169566887 + 0.9833163152014732*$(x[1]) + -0.36891084266703134*$(x[2]) + -0.5752518327899794*$(x[3])) + -0.16669374723112762*softplus(-0.1240101222843979 + -0.304115262711357*$(x[1]) + -0.19393989104697695*$(x[2]) + -0.50508829642924*$(x[3]))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m

max_stars_repo_name: PSORLab/RSActivationFunctions

max_issues_repo_name: PSORLab/RSActivationFunctions

max_forks_repo_name: PSORLab/RSActivationFunctions


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct IoK8sApiCoreV1Taint <: SwaggerModel
    effect::Any # spec type: Union{ Nothing, String } # spec name: effect
    key::Any # spec type: Union{ Nothing, String } # spec name: key
    timeAdded::Any # spec type: Union{ Nothing, IoK8sApimachineryPkgApisMetaV1Time } # spec name: timeAdded
    value::Any # spec type: Union{ Nothing, String } # spec name: value

    function IoK8sApiCoreV1Taint(;effect=nothing, key=nothing, timeAdded=nothing, value=nothing)
        o = new()
        validate_property(IoK8sApiCoreV1Taint, Symbol("effect"), effect)
        setfield!(o, Symbol("effect"), effect)
        validate_property(IoK8sApiCoreV1Taint, Symbol("key"), key)
        setfield!(o, Symbol("key"), key)
        validate_property(IoK8sApiCoreV1Taint, Symbol("timeAdded"), timeAdded)
        setfield!(o, Symbol("timeAdded"), timeAdded)
        validate_property(IoK8sApiCoreV1Taint, Symbol("value"), value)
        setfield!(o, Symbol("value"), value)
        o
    end
end # type IoK8sApiCoreV1Taint

const _property_map_IoK8sApiCoreV1Taint = Dict{Symbol,Symbol}(Symbol("effect")=>Symbol("effect"), Symbol("key")=>Symbol("key"), Symbol("timeAdded")=>Symbol("timeAdded"), Symbol("value")=>Symbol("value"))
const _property_types_IoK8sApiCoreV1Taint = Dict{Symbol,String}(Symbol("effect")=>"String", Symbol("key")=>"String", Symbol("timeAdded")=>"IoK8sApimachineryPkgApisMetaV1Time", Symbol("value")=>"String")
Base.propertynames(::Type{ IoK8sApiCoreV1Taint }) = collect(keys(_property_map_IoK8sApiCoreV1Taint))
Swagger.property_type(::Type{ IoK8sApiCoreV1Taint }, name::Symbol) = Union{Nothing,eval(Meta.parse(_property_types_IoK8sApiCoreV1Taint[name]))}
Swagger.field_name(::Type{ IoK8sApiCoreV1Taint }, property_name::Symbol) =  _property_map_IoK8sApiCoreV1Taint[property_name]

function check_required(o::IoK8sApiCoreV1Taint)
    (getproperty(o, Symbol("effect")) === nothing) && (return false)
    (getproperty(o, Symbol("key")) === nothing) && (return false)
    true
end

function validate_property(::Type{ IoK8sApiCoreV1Taint }, name::Symbol, val)
end


text:
using Test
using PotentialCalculation.IdenticalTools
using PotentialCalculation.Atoms
using PotentialCalculation.Molecules

@testset "Molecules" begin
    m = Molecule{AtomOnlySymbol}(["A", "B", "Ca"])
    mm = Molecule{AtomOnlySymbol}(AtomOnlySymbol.(["A", "B", "Ca"]))
    mi = MoleculeIdenticalInformation{AtomOnlySymbol}(["A", "B", "Ca"])
    mmi = MoleculeIdenticalInformation{AtomOnlySymbol}(AtomOnlySymbol.(["A", "B", "Ca"]))

    mt = convert(typeof(mi), m)
    @test mt.atoms == m.atoms

    show(devnull, mm)

    @test length(m) == 3
    @test length(m) == length(mi)

    makeidentical!(mi, (1,2))

    @test ! areidentical(mi, (3,2))
    @test areidentical(mi, (1,2))

    @test try
        makeidentical!(mi, (3,4))
    catch
        true
    end

    @test try
        makeidentical!(mi, (0,2))
    catch
        true
    end

end