text:
#=
     Find the distribution of a tracer given:
     (a) the pathways described by A or its LU decomposition Alu,
     (b) first-guess boundary conditions and interior sources given by d₀,
     (c) perturbations to the surface boundary condition u
    that best fits observations, y,
    according to the cost function,
    J = (ỹ - Ec)ᵀ W⁻¹ (ỹ - Ec)
    subject to Ac = d₀ + Γ u₀.
    W⁻¹ is a (sparse) weighting matrix.
    See Supplementary Section 2, Gebbie & Huybers 2011.
# Arguments
- `u`: control vector of surface tracer perturbations
- `Alu`: LU decomposition of water-mass matrix A
- `y`: observations on 3D grid
- `d₀`: first guess of boundary conditions and interior sources
- `W⁻`: weighting matrix best chosen as inverse error covariance matrix
- `wet`: BitArray mask of ocean points
=#
using Revise, TMI, Interpolations, Statistics

TMIversion = "modern_90x45x33_GH10_GH12"
A, Alu, γ, TMIfile, L, B = config_from_nc(TMIversion)

# first guess of change to surface boundary conditions
# how many randomly sampled observations?
N = 20

# first guess of change to surface boundary conditions
# ocean values are 0
#u = zerosurfaceboundary(γ)
u = (;surface = zerosurfaceboundary(γ))
uvec = vec(u)

# take synthetic, noisy observations
y, W⁻, ctrue, ytrue, locs, wis = synthetic_observations(TMIversion,"θ",γ,N)

# make a silly first guess for surface
#b = mean(y) * onesurfaceboundary(γ)
b = (;surface = mean(y) * onesurfaceboundary(γ))

# assume temperature known ± 5°C
σb = 5.0
Q⁻ = 1.0/(σb^2)

## gradient checks
# check with forward differences
fg(x) = costfunction_point_obs(x,Alu,b,u,y,W⁻,wis,locs,Q⁻,γ)
f(x) = fg(x)[1]
J0 = f(uvec)
J̃₀,gJ₀ = fg(uvec)
fg!(F,G,x) = costfunction_point_obs!(F,G,x,Alu,b,u,y,W⁻,wis,locs,Q⁻,γ)

ϵ = 1e-3 # size of finite perturbation
ii = rand(1:sum(γ.wet[:,:,1]))
println("gradient check location=",ii)
δu = copy(uvec); δu[ii] += ϵ
∇f_finite = (f(δu) - f(uvec))/ϵ
println("∇f_finite=",∇f_finite)

fg!(J̃₀,gJ₀,(uvec+δu)./2) # J̃₀ is not overwritten
∇f = gJ₀[ii]
println("∇f=",∇f)

# error less than 10 percent?
println("Percent error=",100*abs(∇f - ∇f_finite)/abs(∇f + ∇f_finite))

# optimize the sparse data map with an Optim.jl method
iterations = 10
out = sparsedatamap(uvec,Alu,b,u,y,W⁻,wis,locs,Q⁻,γ,iterations)

# reconstruct by hand to double-check.
ũ = unvec(u,out.minimizer)
b̃ = adjustboundarycondition(b,ũ)

# reconstruct tracer map
c₀ = steadyinversion(Alu,b,γ)
c̃  = steadyinversion(Alu,b̃,γ)

Δc̃ = c̃ - ctrue
Δc₀ = c₀ - ctrue

# plot the difference
level = 15 # your choice 1-33
depth = γ.depth[level]
cntrs = -10:0.5:10
label = "Optimized misfit: Δc̃"
# Help: needs work with continents and labels
planviewplot(Δc̃, depth, cntrs, titlelabel=label)

label = "First guess misfit: Δc₀"
planviewplot(Δc₀, depth, cntrs, titlelabel=label)

# % what is the uncertainty in the surface boundary condition?
# err_d = 1;
# diagonal = false; % or set to false.
# if diagonal
#   S     = 1./err_d.^2;
# else
#   % if diagonal==false
#   % impose spatial smoothing in surface b.c.
#   lengthscale = 4; % horizontal lengthscale in units of gridcells
#   factor = 0.126.*(1/lengthscale)^2 ;
#   load Del2_4deg.mat
#   S = factor.* sparse(1:Nsfc,1:Nsfc,1./err_d.^2);
#   S = S + Del2'*(lengthscale^4.*S*Del2);
# end

# % how much did the data points reduce the error (globally).
# sqrt(sum((c-Tobs).^2)/Nfield)
# sqrt(sum((c0-Tobs).^2)/Nfield)

# % how much did the data points reduce the error (at data points).
# Nobs = length(y);
# sqrt(sum((E*c-y).^2)/Nobs)
# sqrt(sum((E*c0-y).^2)/Nobs)


text:
module BugReporting

# Use README as the docstring of the module:
@doc read(joinpath(dirname(@__DIR__), "README.md"), String) BugReporting

export replay, make_interactive_report

using Base.Filesystem: uperm
using rr_jll
using Zstd_jll
using HTTP, JSON
using AWSCore, AWSS3
using Tar
using Pkg

# https://github.com/JuliaLang/julia/pull/29411
if isdefined(Base, :exit_on_sigint)
    using Base: exit_on_sigint
else
    exit_on_sigint(on::Bool) =
        ccall(:jl_exit_on_sigint, Cvoid, (Cint,), on)
end

const WSS_ENDPOINT = "wss://53ly7yebjg.execute-api.us-east-1.amazonaws.com/test"
const GITHUB_APP_ID = "Iv1.c29a629771fe63c4"
const TRACE_BUCKET = "julialang-dumps"

function check_rr_available()
    if !isdefined(rr_jll, :rr_path)
        error("RR not available on this platform")
    end
end

get_record_flags() =
    split(get(ENV, "JULIA_RR_RECORD_ARGS", ""), ' ', keepempty=false)

struct InvalidPerfEventParanoidError <: Exception
    value
end

function Base.showerror(io::IO, err::InvalidPerfEventParanoidError)
    println(io, "InvalidPerfEventParanoidError")
    print(io, """
    rr needs /proc/sys/kernel/perf_event_paranoid <= 1, but it is $(err.value).
    Change it to 1, or use `JULIA_RR_RECORD_ARGS=-n julia --bug-report=rr` (slow).
    Consider putting 'kernel.perf_event_paranoid = 1' in /etc/sysctl.conf
    or change it temporarily by
        echo 1 | sudo tee /proc/sys/kernel/perf_event_paranoid
    """)
end

# `path` used for testing
function check_perf_event_paranoid(path = "/proc/sys/kernel/perf_event_paranoid")
    isempty(intersect(["-n", "--no-syscall-buffer"], get_record_flags())) || return
    isfile(path) || return  # let `rr` handle this
    value = tryparse(Int, read(path, String))
    value === nothing && return  # let `rr` handle this
    value <= 1 && return
    throw(InvalidPerfEventParanoidError(value))
end

function collect_inner_traces(trace_directory)
    # If this is already a trace directory, return just it
    if isfile(joinpath(trace_directory, "version"))
        return [trace_directory]
    end
    # If this is a directory containing traces, return those
    traces = String[]
    for f in readdir(trace_directory, join=true)
        if isfile(joinpath(f, "version"))
            push!(traces, f)
        end
    end
    return unique(realpath.(traces))
end

function find_latest_trace(trace_directory)
    # What _RR_TRACE_DIR calls a "trace directory" is not what `rr pack` calls a "trace directory"
    # This function allows us to normalize to the inner "latest-trace" directory, if necessary.
    latest_symlink = joinpath(trace_directory, "latest-trace")
    if !isfile(joinpath(trace_directory, "version")) && islink(latest_symlink)
        return realpath(latest_symlink)
    end
    return trace_directory
end

function rr_pack(trace_directory)
    check_rr_available()

    rr() do rr_path
        for dir in collect_inner_traces(trace_directory)
            @debug("rr pack'ing $(dir)")
            run(`$rr_path pack $(dir)`)
        end
    end
end

function rr_record(args...; trace_dir=nothing)
    check_rr_available()
    check_perf_event_paranoid()

    record_flags = get_record_flags()
    rr() do rr_path
        new_env = copy(ENV)
        if trace_dir !== nothing
            new_env["_RR_TRACE_DIR"] = trace_dir
        end
        # Intersperse all given arguments with spaces, then splat:
        rr_cmd = `$(rr_path) record $(record_flags)`
        for arg in args
            rr_cmd = `$(rr_cmd) $(arg)`
        end
        run(ignorestatus(setenv(rr_cmd, new_env)))
    end
end

function download_rr_trace(trace_url; verbose=true)
    Pkg.PlatformEngines.probe_platform_engines!()
    artifact_hash = Pkg.create_artifact() do dir
        mktempdir() do dl_dir
            # Download into temporary directory, unpack into artifact directory
            local_path = joinpath(dl_dir, "trace.tar.zst")
            Pkg.PlatformEngines.download(trace_url, local_path; verbose=verbose)
            proc = zstdmt() do zstdp
                open(`$zstdp --stdout -d $local_path`, "r+")
            end
            Tar.extract(proc, dir)
        end
    end
    return Pkg.artifact_path(artifact_hash)
end

function replay(trace_url)
    if startswith(trace_url, "s3://")
        trace_url = string("https://s3.amazonaws.com/julialang-dumps/", trace_url[6:end])
    end
    if startswith(trace_url, "https://")
        trace_url = download_rr_trace(trace_url)
    end

    if !isdir(trace_url)
        error("Invalid trace location: $(trace_url)")
    end

    rr() do rr_path
        run(`$(rr_path) replay $(find_latest_trace(trace_url))`)
    end
end

function make_interactive_report(report_type, ARGS=[])
    default_julia_args = `--history-file=no`
    if report_type == "rr-local"
        rr_record(`$(Base.julia_cmd()) $default_julia_args`, ARGS)
        return
    elseif report_type == "rr"
        exit_on_sigint(false)  # throw InterruptException on Ctrl-C
        artifact_hash = Pkg.create_artifact() do trace_dir
            rr_record(`$(Base.julia_cmd()) $default_julia_args`, ARGS; trace_dir=trace_dir)
            @info "Preparing trace directory for upload (if your trace is large this may take a few minutes)"
            rr_pack(trace_dir)
        end
        upload_rr_trace(Pkg.artifact_path(artifact_hash))
        return
    elseif report_type == "help"
        show(stdout, "text/plain", @doc(BugReporting))
        println()
        return
    end
    error("Unknown report type: $report_type")
end

const S3_CHUNK_SIZE = 25 * 1024 * 1024 # 25 MB

include("sync_compat.jl")

function upload_rr_trace(trace_directory)
    # Auto-pack this trace directory if it hasn't already been:
    sample_directory = joinpath(trace_directory, "latest-trace")
    if isdir(sample_directory) && uperm(sample_directory) & 0x2 == 0
        @info "`$sample_directory` not writable. Skipping `rr pack`."
    else
        rr_pack(trace_directory)
    end

    c = Channel()
    t = @async HTTP.WebSockets.open(WSS_ENDPOINT) do ws
        write(ws, "Hello Server, if it's not too much trouble, please send me S3 credentials")
        x = readavailable(ws)
        put!(c, JSON.parse(String(x))["connectionId"])
        # This will block until the user has completed the authentication flow
        x = readavailable(ws)
        push!(c, JSON.parse(String(x)))
    end
    bind(c, t)
    connectionId = take!(c)
    println()
    println("""
    ### IMPORTANT =============================================================
    You are about to upload a trace directory to a publicly accessible location.
    Such traces contain any information that was accessed by the traced
    executable during its execution. This includes any code loaded, any
    secrets entered, the contents of any configuration files, etc.
    DO NOT proceed, if you do not wish to make this information publicly available.
    By proceeding you explicitly agree to waive any privacy interest in the
    uploaded information.
    ### =======================================================================
    """)
    println("To upload a trace, please authenticate, by visiting:\n")
    println("\thttps://github.com/login/oauth/authorize?client_id=$GITHUB_APP_ID&state=$(HTTP.escapeuri(connectionId))")
    println("\n", "You can cancel upload by `Ctrl-C`.")
    flush(stdout)
    s3creds = try
        take!(c)
    catch err
        if err isa InterruptException
            println()
            println("Cancel uploading the trace.")
            return
        end
        rethrow()
    end

    println()
    @info "Uploading Trace directory"

    creds = AWSCore.AWSCredentials(
        s3creds["AWS_ACCESS_KEY_ID"],
        s3creds["AWS_SECRET_ACCESS_KEY"],
        s3creds["AWS_SESSION_TOKEN"])
    aws = AWSCore.aws_config(creds = creds, region="us-east-1")

    # Tar it up
    proc = zstdmt() do zstdp
        open(`$zstdp -`, "r+")
    end

    t = @async begin try
        upload = s3_begin_multipart_upload(aws, TRACE_BUCKET, s3creds["UPLOAD_PATH"])
        tags = Vector{String}()
        i = 1
        @Base.Experimental.sync begin
            while isopen(proc)
                buf = Vector{UInt8}(undef, S3_CHUNK_SIZE)
                n = readbytes!(proc, buf)
                n < S3_CHUNK_SIZE && resize!(buf, n)
                resize!(tags, i)
                let partno = i, buf=buf
                    @async begin
                        try
                            tags[partno] = s3_upload_part(aws, upload, partno, buf)
                        catch e
                            close(proc)
                            rethrow(e)
                        end
                    end
                end
                i += 1
            end
        end
        s3_complete_multipart_upload(aws, upload, tags)
    catch e
        Base.showerror(stderr, e)
    end
    end

    # Start the Tar creation process, the file will be uploaded as it's created
    Tar.create(trace_directory, proc)
    close(proc.in)

    wait(t)
    println("Uploaded to https://s3.amazonaws.com/$TRACE_BUCKET/$(s3creds["UPLOAD_PATH"])")
end

end # module

ext: jl


text:
# module Petkovsek

using SymPy

# export algpoly

macro log(var)
    println("$var: ", eval(var))
    return eval(var)
end

function fallingfactorial(x, j)
    result = 1
    for i in 0:j - 1
        result *= (x - i)
    end
    return result
end

function algpoly(polylist, f, n)
    # Construct polynomials qj that arise from treating the shift operator
	# as the difference operator plus the identity operator.
    qlist = []
    for j in 0:length(polylist)-1
        qj = 0 * n
        for i in j:length(polylist)-1
            qj += binomial(i, j) * polylist[i+1]
        end
        push!(qlist, qj)
    end

    # Find all candidates for the degree bound on output polynomial.
    b = maximum([(degree(poly, n) - (j - 1)) for (j, poly) in enumerate(qlist)])
    first = degree(f, n) - b
    second = -1 * b - 1
    lc = [LC(Poly(p,n)) for p in qlist]

    alpha = 0 * n
    for j in 0:length(qlist)-1
        if (degree(qlist[j+1], n) - j) == b
            alpha += lc[j+1] * fallingfactorial(n, j)
        end
    end

    deg = polyroots(alpha)
    third = maximum(keys(deg))
    d = max(first, second, third, 0) |> Int64
    # Use method of undetermined coefficients to find the output polynomial.
    varlist = symset("a", d+1)
    pol = 0 * n
    for i in 0:d
        pol += n^i * varlist[i+1]
    end
    solution = -1 * f
    for (i, poly) in enumerate(polylist)
        solution += subs(pol, (n, n + (i-1))) * poly
    end
    coef = coeffs(Poly(solution, n))
    filter!(e->e!=n*0, coef)
    if isempty(coef)
        return Dict([v => v for v in varlist])
    end
    return solve(coef, varlist)
end

function factors(expr)
    c, list = factor_list(expr)
    result = [x^y for (x,y) in list]
    if c != 1
        push!(result, c)
    end
    return result
end

function alghyper(polylist, n)
    p = polylist[1]
    alist = factors(p)
    for (i,p) in enumerate(alist)
        alist[i] /= LC(Poly(p, n))
    end
    if !(1*n^0 in alist)
        push!(alist, 1*n^0)
    end

    d = length(polylist)
    p = subs(polylist[end], (n,n-d+2))
    blist = factors(p)
    for (i,p) in enumerate(blist)
        blist[i] /= LC(Poly(p,n))
    end
    if !(1*n^0 in blist)
        push!(blist, 1*n^0)
    end

    solutions = []
    for aelem in alist
        for belem in blist
            plist = []
            for i in 0:d-1
                pi = polylist[i+1]
                for j in 0:i-1
                    pi *= subs(aelem, (n, n+j))
                end
                for j in i:d-1
                    pi *= subs(belem, (n, n+j))
                end
                push!(plist, pi)
            end

            m = maximum([degree(Poly(p, n)) for p in plist])
            alpha = [coeff(expand(p), n^m) for p in plist]
            @syms z
            zpol = 0*z
            for i in 0:length(alpha) - 1
                zpol += alpha[i+1]*z^i
            end

            vals = [key for (key,val) in polyroots(zpol) if key != 0]

            for x in vals
                polylist2 = [x^(i-1)*p for (i,p) in enumerate(plist)]
                polysols = algpoly(polylist2, 0*n, n)
                if isempty(polysols)
                    continue
                end
                polysols = collect(values(polysols))
                filter!(e->e!=n*0, polysols)
                if length(polysols) > 0
                    c = 0*n
                    for (i,p) in enumerate(polysols)
                        c += n^(i-1) * p
                    end
                    s = x * aelem/belem * subs(c, (n, n+1))/c
                    push!(solutions, simplify(s))
                end
            end
        end
    end
    return solutions
end

function tohg(sol, n)
    facts = factors(sol)
    result = []
    for f in facts
        if has(f, n)
            f = f |> subs(n, n-1)
            c = coeff(f, n)
            f = f / c
            push!(result, factorial(f))
            if c != 1
                push!(result, c^n)
            end
        else
            push!(result, f^n)
        end
    end
    return prod(result)
end

function hyper()

# end # module

# using SymPy
# using Petkovsek

@syms n y

alghyper([2*n*(n+1), -(n^2 +3*n-2), n-1], n)

# algpoly([3, -n, n-1], 0*n, n)
# algpoly([n-1, -n, 3], 0*n, n)
println("alghyper: ", [tohg(sol, n) for sol in alghyper([2*n*(n+1), -(n^2 +3*n-2), n-1], n)])
println("algpoly: ", algpoly([n*(n + 1), -1*n^2 - 3*n + 2, 2*n - 2], 0*n,n))


text:
slices(n) = 1./tan.(linspace(0,π/2,n+2)[2:end-1])

function RCNP(::Type{Min}, m, α, Ƶ1, Ƶ2)

	@show α
	mSUP = copy(m); mINF = copy(m)
	@constraint(mSUP, copy(Ƶ2, mSUP).aff >= α*copy(Ƶ1, mSUP).aff)
	@constraint(mINF, copy(Ƶ2, mINF).aff <= α*copy(Ƶ1, mINF).aff)
	s1 = @suppress solve(mSUP, method=:lexico)
	s2 = @suppress solve(mINF, method=:lexico)

	zA = last(getY_N(mSUP))
	xA = [getvalue(JuMP.Variable(m, i), 2) for i = 1:vm.numCols]

	zB = first(getY_N(mINF))
	xB = [getvalue(JuMP.Variable(m, i), 1) for i = 1:vm.numCols]

	@show zA, zB

	if zA[1] <= zB[1] && zA[2] <= zB[2]
		@constraint(mINF, copy(Ƶ2,mINF).aff <= zA[2] - 1e-4)
		s2 = @suppress solve(mINF, method=:lexico)
		@show s2
		if s2 == :Infeasible
			return Nullable((xA, zA)), Nullable{typeof((xA,zA))}()
		end
		zB = first(getY_N(mINF))
		xB = [getvalue(JuMP.Variable(m, i), 1) for i = 1:vm.numCols]
	elseif zB[1] <= zA[1] && zB[2] <= zA[2]
		@constraint(mSUP, copy(Ƶ1,mSUP).aff <= zB[1] - 1e-4)
		s1 = @suppress solve(mSUP, method=:lexico)
		@show s1
		if s1 == :Infeasible
			return Nullable{typeof((xB,zB))}(), Nullable((xB, zB))
		end
		zA = last(getY_N(mSUP))
		xA = [getvalue(JuMP.Variable(m, i), 2) for i = 1:vm.numCols]
	end
	println("###")
	return Nullable((xA,zA)), Nullable((xB,zB))
end

function RCNP(::Type{Max}, m, α, Ƶ1, Ƶ2)

	@show α
	mSUP = copy(m); mINF = copy(m)
	@constraint(mSUP, copy(Ƶ2, mSUP).aff >= α*copy(Ƶ1, mSUP).aff)
	@constraint(mINF, copy(Ƶ2, mINF).aff <= α*copy(Ƶ1, mINF).aff)
	s1 = @suppress solve(mSUP, method=:lexico)
	s2 = @suppress solve(mINF, method=:lexico)

	zA = first(getY_N(mSUP))
	xA = [getvalue(JuMP.Variable(m, i), 1) for i = 1:m.numCols]

	zB = last(getY_N(mINF))
	xB = [getvalue(JuMP.Variable(m, i), 1) for i = 1:m.numCols]

	@show zA, zB

	if zA[1] >= zB[1] && zA[2] >= zB[2]
		@constraint(mINF, copy(Ƶ2,mINF).aff >= zA[2] + 1e-4)
		s2 = @suppress solve(mINF, method=:lexico)
		@show s2
		if s2 == :Infeasible
			return Nullable((xA, zA)), Nullable{typeof((xA,zA))}()
		end
		zB = last(getY_N(mINF))
		xB = [getvalue(JuMP.Variable(m, i), 1) for i = 1:m.numCols]
	elseif zB[1] >= zA[1] && zB[2] >= zA[2]
		@constraint(mSUP, copy(Ƶ1,mSUP).aff >= zB[1] + 1e-4)
		s1 = @suppress solve(mSUP, method=:lexico)
		@show s1
		if s1 == :Infeasible
			return Nullable{typeof((xB,zB))}(), Nullable((xB, zB))
		end
		zA = first(getY_N(mSUP))
		xA = [getvalue(JuMP.Variable(m, i), 1) for i = 1:m.numCols]
	end
	@show zA, zB
	println("###")
	return Nullable((xA,zA)), Nullable((xB,zB))
end

function select_slices(YN, k)
	K = length(YN)
	@assert k <= K÷2
	G = OffsetArray(Vector{Int}, 0:k)
	G[0] = [0]
	G[k] = [K-1]
	for i = 1:k-1
		G[i] = collect(2i-1:K-2k+2i-1)
	end
	@show G

	d = Dict{Tuple{Int,Int}, Int}()
	for col = 0:k-1
		for i in G[col] .+ 1
			for j in G[col+1] .+ 1
				if col == 0
					d[i,j] = abs((YN[i][1]-YN[j][1]) * (YN[i][2]-YN[j][2]))
				else
					if j >= i+2
						d[i,j] = abs((YN[i+1][1]-YN[j][1]) * (YN[i+1][2]-YN[j][2]))
					end
				end
			end
		end
	end

	m = Model(solver = CplexSolver(CPX_PARAM_SCRIND = 0))
	@variable(m, x[keys(d)], Bin)
	@variable(m, wmax, Int)

	@objective(m, Min, wmax)

	edges_1_x = Iterators.filter(x->first(x)==1, keys(d))
	edges_x_K = Iterators.filter(x->last(x)==K, keys(d))

	@constraint(m, sum(x[e] for e in edges_1_x) == 1)
	@constraint(m, sum(x[e] for e in edges_x_K) == 1)
	@constraint(m, [i = keys(d)], wmax >= x[i]*d[i])
	@constraint(m, sum(x) == k)
	for i = 2:K-1
		in_edges = Iterators.filter(x->last(x)==i, keys(d))
		out_edges = Iterators.filter(x->first(x)==i, keys(d))
		@constraint(m, sum(x[j] for j in in_edges) == sum(x[l] for l in out_edges))
	end

	@show solve(m)
	@show getobjectivevalue(m)

	x_res = getvalue(x)

	res=Int[]
	i = 1
	j = 2
	while j != K
		if i==1 && j==2 && x_res[(i,j)] == 1.
			push!(res, 2)
			i=2
		elseif length(res) == k-1
			@assert x_res[(i,K)] == 1 res
			push!(res, K)
			j = K
		else
			j = findfirst(x->x_res[(i,x)]==1, (i+2):K) + (i+1)
			push!(res, j)
			i = j
		end
	end
	res

end


text:
using DiffEqBayes, OrdinaryDiffEq, ParameterizedFunctions, Distances, StatsBase, RecursiveArrayTools
using Test

# One parameter case
f1 = @ode_def begin
  dx = a*x - x*y
  dy = -3y + x*y
end a
u0 = [1.0,1.0]
tspan = (0.0,10.0)
prob1 = ODEProblem(f1,u0,tspan,[1.5])
sol = solve(prob1,Tsit5())
t = collect(range(1,stop=10,length=10))
randomized = VectorOfArray([(sol(t[i]) + .01randn(2)) for i in 1:length(t)])
data = convert(Array,randomized)
priors = [Normal(1.5,0.01)]

bayesian_result = abc_inference(prob1,Tsit5(),t,data,priors;
                                num_samples=500,ϵ = 0.001)

@test mean(bayesian_result.parameters, weights(bayesian_result.weights)) ≈ 1.5 atol=0.1

sol = solve(prob1,Tsit5(),save_idxs=[1])
randomized = VectorOfArray([(sol(t[i]) + .01randn(1)) for i in 1:length(t)])
data = convert(Array,randomized)

bayesian_result = abc_inference(prob1,Tsit5(),t,data,priors;
                                num_samples=500,ϵ = 0.001,save_idxs=[1])

@test mean(bayesian_result.parameters, weights(bayesian_result.weights)) ≈ 1.5 atol=0.1

# custom distance-function
weights_ = ones(size(data)) # weighted data
for i = 1:3:length(data)
    weights_[i] = 0
    data[i] = 1e20 # to test that those points are indeed not used
end
distfn = function (d1, d2)
    d = 0.0
    for i=1:length(d1)
        d += (d1[i] - d2[i])^2 * weights_[i]
    end
    return sqrt(d)
end
bayesian_result = abc_inference(prob1,Tsit5(),t,data,priors;
                                num_samples=500, ϵ = 0.001,
                                distancefunction = distfn)

@test mean(bayesian_result.parameters, weights(bayesian_result.weights)) ≈ 1.5 atol=0.1

# Four parameter case
f1 = @ode_def begin
  dx = a*x - b*x*y
  dy = -c*y + d*x*y
end a b c d
u0 = [1.0,1.0]
tspan = (0.0,10.0)
p = [1.5,1.0,3.0,1.0]
prob1 = ODEProblem(f1,u0,tspan,p)
sol = solve(prob1,Tsit5())
t = collect(range(1,stop=10,length=10))
randomized = VectorOfArray([(sol(t[i]) + .01randn(2)) for i in 1:length(t)])
data = convert(Array,randomized)
priors = [Truncated(Normal(1.5,1),0,2),Truncated(Normal(1.0,1),0,1.5),
          Truncated(Normal(3.0,1),0,4),Truncated(Normal(1.0,1),0,2)]

bayesian_result = abc_inference(prob1,Tsit5(),t,data,priors; num_samples=500)

meanvals = mean(bayesian_result.parameters, weights(bayesian_result.weights), 1)

@test meanvals[1] ≈ 1.5 atol=3e-1
@test meanvals[2] ≈ 1.0 atol=3e-1
@test meanvals[3] ≈ 3.0 atol=3e-1
@test meanvals[4] ≈ 1.0 atol=3e-1


text:
# Creating a photodynamical model of TRAPPIST-1 based on the Spitzer data:
# Create a function to read in the Spitzer datasets:

using DelimitedFiles
using Statistics
using PyPlot
using Optim
if ~@isdefined(CGS)
  include("../../../src/CGS.jl")
  using Main.CGS
end

# Now, call ttv to get a list of transit times and pair these with the data:

include("../../../src/NbodyGradient/src/ttv.jl")

# Read in initial orbital elements:
elements = readdlm("../../../data/elements_noprior_students.txt",',')

# Set integration parameters:
n=8
nplanet = n-1
t0 = 7257.93115525
h  = 0.06
tmax = 1600.0

# Set up array for holding times of transit:
ntt = zeros(Int64,n);
for i=2:n
  ntt[i] = ceil(Int64,tmax/elements[i,2])+1
end
tt1 = zeros(3,n,maximum(ntt));
# Save a counter for the actual number of transit times of each planet:
count1 = zeros(Int64,n);
# Set the "size" of the star (only checks that close to transit):
rstar = 1e11
# Call the ttv function:
@time dq=ttvbv_elements!(n,t0,h,tmax,elements,tt1,count1,0.0,0,0,rstar)

# Now, call Limbdark to create a model for each planet:

using Limbdark

# Take the radius ratios and limb-darkening parameters from
# Delrez et al. (2018):

depth_ch2  = [0.7277,0.6940,0.3566,0.4802,0.634,0.764,0.346] .* 0.01

rp = sqrt.(depth_ch2) # Ratio of planet radius to stellar radius

# Fractional radius errors from Delrez et al.:
srp = [0.028,0.028,0.020,0.025,0.026,0.029,0.025]./[1.127,1.100,0.788,0.915,1.052,1.154,0.777].*rp

# Set limb-darkening parameters:

u_ch1 = [0.168,0.244]
u_ch2 = [0.161,0.208]

# Set transit durations:
T = [36.19,42.31,49.33,55.92,63.14,68.53,76.92] ./(60*24)

# Set impact parameters:
b = [0.157,0.148,0.08,0.240,0.337,0.406,0.392]

# Set the stellar density:

include("../../../src/regress.jl")

function compute_model_all(tt1,count,rho_T1,texp,b,rp,u_ch2,neach)
      # Create a photometric model for all data.
      # Number of points per transit:
      # Create a matrix to hold the simulations:
      phot_mod = zeros(sum(count),neach)
      # Create a time array for each transit:
      t_mod = zeros(sum(count),neach)
      # Create array to tell me which transit belongs to which planet:
      ip_trans = zeros(Int64,sum(count))

      # Density of T1 star in units of solar density:
      fac = rho_T1^(1//3)*AU/RSUN

      # Exposure time in days:
      texp_day = texp/(60*24)

      # Two limb-darkening coefficients:
      nu = 2

      itrans = 0
      # Initialize transit structure (no derivatives for now):
      trans =  Limbdark.transit_init(0.1,0.5,u_ch2,false)
      # Loop over each planet:

      for ip=1:nplanet
      # Carry out a fit to the simulated transits:
        tt_ref1 = zeros(count1[ip+1])
        fn = zeros(Float64,2,count1[ip+1])
        sig = ones(count1[ip+1])
        tti1 = tt1[1,ip+1,1:count1[ip+1]]
        for j=1:count1[ip+1]
          fn[1,j] = 1.0
          fn[2,j] = round(Int64,(tti1[j]-elements[ip+1,3])/elements[ip+1,2])
          epoch = round((tti1[j]-elements[ip+1,3])/elements[ip+1,2])
        end
        coeff,cov = regress(fn,tti1,sig)
        println(ip," ",coeff)
        tt_ref1 = coeff[1] .+coeff[2] .*fn[2,:]
      # Now, Loop over each transit:
      for ifile=1:count[ip+1]
        # Create an time array for the current transit:
        tcurr = tt_ref1[ifile] .+ (collect(-div(neach,2):div(neach,2)-1) .+0.5) .*texp_day
        # First set the limb-darkening parameters
        # for the current light curve

        # Number of exposure times for this case:
        nt = neach
        # Set up array to hold the time-integrated flux for each planet:
        favg1 = zeros(5+nu,nt)
        # Set up arrays to hold statistics on integrations:
        neval1 = zeros(Int64,nt)
        depthmax1 = zeros(Int64,nt)
        # Loop over each transit:
        itrans += 1
        ip_trans[itrans] = ip
        ttbv = tt1[:,ip+1,ifile]
        # Now, create a model.  Convert impact parameter and
        # sky velocity to units of radius of the star:
        param = [ttbv[1],ttbv[2]*fac,ttbv[3]*fac]
        # Set the radius ratio and impact parameter to the current planet:
        param[3] = b[ip]
        trans.r = rp[ip]
        Limbdark.integrate_lightcurve!(trans,param,tcurr,texp_day,favg1,nt,1e-11,40,neval1,depthmax1)
        phot_mod[itrans,:] .= favg1[1,:]
        t_mod[itrans,:] .= tcurr
      end
    end
    return t_mod,ip_trans,phot_mod
end

rho_T1 = 52.5363620668614

neach = 400

t_mod_all,ip_trans, phot_mod_all = compute_model_all(tt1,count1,rho_T1,0.5,b,rp,u_ch2,neach)

pygui(false)
# Now, emplace into an image:
pname = ["b","c","d","e","f","g","h"]
maxdepth = minimum(phot_mod_all)
fig,axes = subplots(1,7,figsize=[12,5])
# Range of the transit x-axis:
texp_plot = 0.5
t_xaxis = (collect(-div(neach,2):div(neach,2)-1) .+0.5) .*texp_plot
tmin = (t_xaxis[1]-0.5*texp_plot)/60
tmax = (t_xaxis[neach]+0.5*texp_plot)/60
stretch = 1
for i=1:nplanet
  ax = axes[i]
  if i > 3
    img = zeros(count1[i+1]*stretch,neach)
    i0 = sum(count1[1:i])
    for j=1:count1[i+1]
      for k=1:stretch
        img[(j-1)*stretch+k,:] = phot_mod_all[i0+j,:]
      end
    end
  else
    img = zeros(count1[i+1]*stretch,div(neach,2))
    i0 = sum(count1[1:i])
    for j=1:count1[i+1]
      for k=1:stretch
        img[(j-1)*stretch+k,:] = phot_mod_all[i0+j,div(neach,4)+1:neach-div(neach,4)]
      end
    end
  end

  # Set one pixel equal to maximum depth:
  img[1,175]=maxdepth
  print(i," ",size(img))
  if i > 3
    #ax.imshow(img,aspect="auto",cmap="winter",interpolation="nearest",origin="lower",extent = [tmin,tmax,0,1600])
    ax.imshow(img,aspect="auto",cmap="winter",interpolation="nearest",origin="lower",extent = [tmin,tmax,0,1600])
  else
    #ax.imshow(img,aspect="auto",cmap="winter",interpolation="nearest",origin="lower",extent = [tmin/2,tmax/2,0,1600])
    ax.imshow(img,aspect="auto",cmap="winter",interpolation = "nearest",origin="lower",extent = [tmin/2,tmax/2,0,1600])
  end
  ax.title.set_text(pname[i])
  #ax.set_xticks(fontsize=8)
  #ax.title(pname[i])
  #ax.tick_params(axis="x",labelsize=6)
  if i == 4
    ax.set_xlabel(L"Time [hr] - $\mathdefault{t_n}$")
  end
  if i == 1
    ax.set_ylabel(L"Time [days] - $\mathdefault{t_0}$")
  else
        ax.set_yticks([])
  end
end
subplots_adjust(hspace=0,wspace=0.05)
savefig("../T1_riverplot.pdf",bbox_inches="tight")


text:
using Configurations: @option

export FormatConfig

"""
    FormatConfig(delimiter, newline, indent, float, int, bool)

Specify the formatting options of an `Input` or part of an `Input`.

# Arguments
- `delimiter::String`: the delimiter between objects. We suggest `" "`.
- `newline::String`: the line terminator. Unix systems and macOS consider `'\\n'` as a line terminator, while Windows supports `'\\r\\n'`. We suggest `"\\n"`.
- `indent::String`: the empty spaces at the beginning of a line. We suggest `' '^4`.
- `float::String`: the format specification for `AbstractFloat`.
- `int::String`: the format specification for `Integer`s.
- `bool::String`: the format specification for `Bool`s. For Fortran software, it could be `".%s."`.
"""
@option mutable struct FormatConfig
    delimiter::String
    newline::String
    indent::String
    float::String
    int::String
    bool::String
end


text:
module Transforms

import Base:   ==
import ACE: read_dict, write_dict, transform, transform_d, inv_transform
abstract type DistanceTransform end

export PolyTransform, IdTransform, MorseTransform, AgnesiTransform

poly_trans(p, r0, r) = @fastmath(((1+r0)/(1+r))^p)

poly_trans_d(p, r0, r) = @fastmath((-p/(1+r0)) * ((1+r0)/(1+r))^(p+1))

poly_trans_inv(p, r0, x) = ( (1+r0)/(x^(1/p)) - 1 )

# TODO: generalise the distance transform to allow
#       ((a + r0) / (a + r) )^p

@doc raw"""
Implements the distance transform
```math
   x(r) = \Big(\frac{1 + r_0}{1 + r}\Big)^p
```

Constructor:
```
PolyTransform(p, r0)
```
"""
struct PolyTransform{TP, T} <: DistanceTransform
   p::TP
   r0::T
end

write_dict(T::PolyTransform) =
   Dict("__id__" => "ACE_PolyTransform", "p" => T.p, "r0" => T.r0)

PolyTransform(D::Dict) = PolyTransform(D["p"], D["r0"])

read_dict(::Val{:ACE_PolyTransform}, D::Dict) = PolyTransform(D)

transform(t::PolyTransform, r::Number) = poly_trans(t.p, t.r0, r)

transform_d(t::PolyTransform, r::Number) = poly_trans_d(t.p, t.r0, r)

inv_transform(t::PolyTransform, x::Number) = poly_trans_inv(t.p, t.r0, x)

(t::PolyTransform)(x) = transform(t, x)

"""
`IdTransform`: Implements the distance transform `z -> z`;
Primarily used for the z-coordinate for the EnvPairPots

Constructor: `IdTransform()`
"""
struct IdTransform <: DistanceTransform
end

write_dict(T::IdTransform) =  Dict("__id__" => "ACE_IdTransform")
IdTransform(D::Dict) = IdTransform()
read_dict(::Val{:ACE_IdTransform}, D::Dict) = IdTransform(D)
transform(t::IdTransform, z::Number) = z
transform_d(t::IdTransform, r::Number) = one(r)
inv_transform(t::IdTransform, x::Number) = x

read_dict(::Val{:SHIPs_IdTransform}, D::Dict) =
   read_dict(Val{:ACE_IdTransform}(), D)

@doc raw"""
Implements the distance transform
```math
   x(r) = \exp( - \lambda (r/r_0) )
```

Constructor:
```
MorseTransform(lambda, r0)
```
"""
struct MorseTransform{T} <: DistanceTransform
   lambda::T
   r0::T
end

write_dict(T::MorseTransform) =
   Dict("__id__" => "ACE_MorseTransform", "lambda" => T.p, "r0" => T.r0)
MorseTransform(D::Dict) = MorseTransform(D["lambda"], D["r0"])
read_dict(::Val{:ACE_MorseTransform}, D::Dict) = MorseTransform(D)
transform(t::MorseTransform, r::Number) = exp(- t.lambda * (r/t.r0))
transform_d(t::MorseTransform, r::Number) = (-t.lambda/t.r0) * exp(- t.lambda * (r/t.r0))
inv_transform(t::MorseTransform, x::Number) = - t.r0/t.lambda * log(x)
(t::MorseTransform)(x) = transform(t, x)

@doc raw"""
Implements the distance transform
```math
   x(r) = \frac{1}{1 + a (r/r_0)^p}
```
with default $a = (p-1)/(p+1)$. That default is chosen such that
$|x'(r)|$ is maximised at $r = r_0$. Default for $p$ is $p = 2$. Any value
$p > 1$ is permitted.

Constructor:
```
AgnesiTransform(r0, [p, [, a]])
```
"""
struct AgnesiTransform{T, TP} <: DistanceTransform
   r0::T
   p::TP
   a::T
   c::T
end

AgnesiTransform(r0, p=2, a=(p-1)/(p+1)) = (@assert p > 1;
                             AgnesiTransform(r0, p, a, - a * p / r0))

write_dict(T::AgnesiTransform) =
Dict("__id__" => "ACE_AgnesiTransform", "r0" => T.r0, "p" => T.p, "a" => T.a)
AgnesiTransform(D::Dict) = AgnesiTransform(D["r0"], D["p"], D["a"])
read_dict(::Val{:ACE_AgnesiTransform}, D::Dict) = AgnesiTransform(D)
transform(t::AgnesiTransform, r::Number) = @fastmath 1 / (1 + t.a * (r/t.r0)^t.p)
transform_d(t::AgnesiTransform, r::Number) = (s1 = (r/t.r0); s2 = s1^(t.p-1);
                                         @fastmath t.c * s2 / (1+t.a * s2*s1)^2)
inv_transform(t::AgnesiTransform, x::Number) = t.r0 * ( (1/x-1)/t.a )^(1/t.p)
(t::AgnesiTransform)(x) = transform(t, x)

end


text:
# Calculate the tidal prism
function prism(Q,ts)
    Qc = Q[:]
    Qc[Q.<0] = 0.0
    TidalFluxStats.integrate(Qc,ts)
end

function prism(Q::Discharge,M::Mask)
    st,sq = segment(Q,M)
    [prism(sq[i],st[i]) for i in eachindex(st)]
end

function adjustbias_prism(Q::Discharge,M::Mask)
    t,q = adjustbias_discharge(Q,M)
    [prism(q[i],t[i]) for i in eachindex(q)]
end


text:
module RVM

using MLBase
using MLKernels

#__precompile__()

@inline logsig(x) = 1.0/(1.0+exp(-x))

export RVMspec, RVMfit, fit, predict

type RVMSpec
    kernel::MLKernels.Kernel
    n_iter::Integer
    normalizer::Function
    threshold::Real
    tol::Real
end

type RVMFit{R <: Real}
    kernel::MLKernels.Kernel
    normal::Function
    labelmap::MLBase.LabelMap
    converged::Bool
    steps::Integer
    w::Vector{R}
    RV::Matrix{R}
end

function _predicted_probability{R <: Real}(Model::RVMFit, Features::AbstractMatrix{R})
    X = Model.normal(Features)
    _ = MLKernels.kernelmatrix(Model.kernel, X, Model.RV)
    ϕ = [_ ones(size(_, 1))] # augment with bias

    P = _classify(Model.w, ϕ)
    P
end

function predict{R <: Real}(Model::RVMFit, Values::AbstractMatrix{R})
    P = _predicted_probability(Model, Values)
    _ = convert(Vector{Int64}, P .>= 0.5)
    Y = _ .+ 1 # booleans map to {0, 1} , {en,de}coders map to {1, 2}
    labeldecode(Model.labelmap, Y), P
end

function fit{R <: Real}(S::RVMSpec, Obs::AbstractMatrix{R}, ObsT::AbstractVector)
    # sample count, features
    n, p = size(Obs)

    # standardize normalizing process
    normalize = S.normalizer(Obs)
    X = normalize(Obs)

    # Relevance Vectors
    RV = X

    # standardize encoding process
    lm = labelmap(ObsT)
    t  = labelencode(lm, ObsT) - 1

    # design matrix / basis functions
    _ = MLKernels.kernelmatrix(S.kernel, X, X)
    ϕ = [_ ones(size(_, 1))] # augment with bias
    basis_count = size(ϕ, 2)

    # initialize priors N(0, 1)
    μ  = zeros(basis_count)
    α₀ = ones(basis_count) .* 1e0
    α₁ = α₀[1:end]
    #β  = 1e0

    # reporting information
    steps = S.n_iter
    succ  = false

    for i in 1:S.n_iter

        # update values
        μ, Σ = _solve_post_prob(μ, α₁, ϕ, t)
        γ  = 1 - (α₁ .* diag(Σ))
        α₁ = γ ./ (μ .^ 2)

        # identify informative vectors
        keep::Vector{Bool} = α₁ .< S.threshold
        if !any(keep)
            keep[1] = true # require at least one vecor
        end
        keep[end] = true # save bias

        # downselect uninformative vectors and weights
        α₀ = α₀[keep]
        α₁ = α₁[keep]
        γ  = γ[keep]
        ϕ  = ϕ[:, keep]
        Σ  = Σ[keep, keep]
        μ  = μ[keep]
        RV = RV[keep[1:end-1], :]

        # check for brreak
        Δα = maximum(abs(α₁ .- α₀))
        if (Δα < S.tol) & (steps > 2)
            steps = i
            succ  = true
            break
        end
        α₀ = α₁[1:end]
    end

    # Trained model as object
    modelfit::RVMFit{R} =
        RVMFit(S.kernel, normalize, lm, succ, steps, μ, RV)
end

function _classify(μ::AbstractVector, ϕ::AbstractMatrix)
    #@show "_classify"
    _ = ϕ * μ
    y = [logsig(x) for x in _]
    y
end

function _log_solve_post_prob(w::AbstractVector, α::AbstractVector,
                              ϕ::AbstractMatrix, t::AbstractVector)
    #@show "_log_post_probability"
    A = diagm(α)
    y = _classify(w, ϕ)
    pos = y[t .== 1]
    neg = y[t .!= 1]

    log_p = (0.5 * w' * A * w)[1] - sum(log(pos)) - sum(log(1 - neg))
end

function _hessian(w::AbstractVector, α::AbstractVector,
                  ϕ::AbstractMatrix, t::AbstractVector)
    #@show "_hessian"
    A = diagm(α)
    y = _classify(w, ϕ)
    B = diagm(y .* (1 - y))

    H = (ϕ' * B * ϕ + A)
    H
end

function _gradient(w::AbstractVector, α::AbstractVector,
                   ϕ::AbstractMatrix, t::AbstractVector)
    #@show "_gradient"
    A = diagm(α)
    y = _classify(w, ϕ)
    G = ϕ' * (t .- y) - (A * w)
    -G
end

function _newton_method(X₀::AbstractVector, F::Function, # Initial, Function
                        ∇::Function, ∇∇::Function)       # Gradient, Hessian
    #@show "_newton_method"
    for _ in 1:50
        #@show F(X₀) # uncomment to watch converge
        X₁ = X₀ - inv(∇∇(X₀)) * ∇(X₀)

        # convergence
        ΔX = maximum(abs(X₁ .- X₀))
        X₀ = X₁[1:end]
        if ΔX < 1e-6
            break
        end
    end
    X₀
end

function _solve_post_prob(w::AbstractVector, α::AbstractVector,
                          ϕ::AbstractMatrix, t::AbstractVector)
    w = _newton_method(w,
                       (w) -> _log_solve_post_prob(w, α, ϕ, t),
                       (w) -> _gradient(w, α, ϕ, t),
                       (w) -> _hessian(w, α, ϕ, t),
                       )

    Σ = inv(_hessian(w, α, ϕ, t))
    w, Σ
end

end


text:
export get_entities_by_color

"""
    get_entities_by_color(dim::Int64=-1)

Return a dictionary of NTuple{4, Int32} RGBA color keys and Vector{NTuple{2,Int32}}
values, representing the dim tags of the entities with that color.
If `dim` == -1, return all entities, otherwise return only entities of dimension `dim`.

Note, entities in gmsh have color (0, 0, 255, 0) by default, so these entities are excluded.
"""
function get_entities_by_color(dim::Int64=-1)
    color_dict = Dict{NTuple{4,Int32}, Vector{NTuple{2,Int32}}}()
    for ent in gmsh.model.get_entities(dim)
        edim, tag = ent
        color = gmsh.model.get_color(edim, tag)
        if color != (0, 0, 255, 0) # Default color when unassigned
            if color ∉ keys(color_dict)
                color_dict[color] = NTuple{2,Int32}[]
            end
            push!(color_dict[color], ent)
        end
    end
    return color_dict
end

max_stars_repo_name: khurrumsaleem/MOCNeutronTransport

max_issues_repo_name: khurrumsaleem/MOCNeutronTransport

max_forks_repo_name: khurrumsaleem/MOCNeutronTransport


text:
# sparse linear algebra functions that perform operations between dense and sparse vectors

export axpyi!, axpyi, sctr!, sctr, gthr!, gthr, gthrz!, roti!, roti

"""
    axpyi!(alpha::BlasFloat, X::CuSparseVector, Y::CuVector, index::SparseChar)

Computes `alpha * X + Y` for sparse `X` and dense `Y`.
"""
axpyi!(alpha::BlasFloat, X::CuSparseVector, Y::CuVector, index::SparseChar)

for (fname,elty) in ((:cusparseSaxpyi, :Float32),
                     (:cusparseDaxpyi, :Float64),
                     (:cusparseCaxpyi, :ComplexF32),
                     (:cusparseZaxpyi, :ComplexF64))
    @eval begin
        function axpyi!(alpha::Number,
                        X::CuSparseVector{$elty},
                        Y::CuVector{$elty},
                        index::SparseChar)
            $fname(handle(), nnz(X), alpha, nonzeros(X), nonzeroinds(X), Y, index)
            Y
        end
        function axpyi(alpha::Number,
                       X::CuSparseVector{$elty},
                       Y::CuVector{$elty},
                       index::SparseChar)
            axpyi!(alpha,X,copy(Y),index)
        end
        function axpyi(X::CuSparseVector{$elty},
                       Y::CuVector{$elty},
                       index::SparseChar)
            axpyi!(one($elty),X,copy(Y),index)
        end
    end
end

"""
    gthr!(X::CuSparseVector, Y::CuVector, index::SparseChar)

Sets the nonzero elements of `X` equal to the nonzero elements of `Y` at the same indices.
"""
gthr!(X::CuSparseVector, Y::CuVector, index::SparseChar)
for (fname,elty) in ((:cusparseSgthr, :Float32),
                     (:cusparseDgthr, :Float64),
                     (:cusparseCgthr, :ComplexF32),
                     (:cusparseZgthr, :ComplexF64))
    @eval begin
        function gthr!(X::CuSparseVector{$elty},
                       Y::CuVector{$elty},
                       index::SparseChar)
            $fname(handle(), nnz(X), Y, nonzeros(X), nonzeroinds(X), index)
            X
        end
        function gthr(X::CuSparseVector{$elty},
                      Y::CuVector{$elty},
                      index::SparseChar)
            gthr!(copy(X),Y,index)
        end
    end
end

"""
    gthrz!(X::CuSparseVector, Y::CuVector, index::SparseChar)

Sets the nonzero elements of `X` equal to the nonzero elements of `Y` at the same indices, and zeros out those elements of `Y`.
"""
gthrz!(X::CuSparseVector, Y::CuVector, index::SparseChar)
for (fname,elty) in ((:cusparseSgthrz, :Float32),
                     (:cusparseDgthrz, :Float64),
                     (:cusparseCgthrz, :ComplexF32),
                     (:cusparseZgthrz, :ComplexF64))
    @eval begin
        function gthrz!(X::CuSparseVector{$elty},
                        Y::CuVector{$elty},
                        index::SparseChar)
            $fname(handle(), nnz(X), Y, nonzeros(X), nonzeroinds(X), index)
            X,Y
        end
        function gthrz(X::CuSparseVector{$elty},
                       Y::CuVector{$elty},
                       index::SparseChar)
            gthrz!(copy(X),copy(Y),index)
        end
    end
end

"""
    roti!(X::CuSparseVector, Y::CuVector, c::BlasFloat, s::BlasFloat, index::SparseChar)

Performs the Givens rotation specified by `c` and `s` to sparse `X` and dense `Y`.
"""
roti!(X::CuSparseVector, Y::CuVector, c::BlasFloat, s::BlasFloat, index::SparseChar)
for (fname,elty) in ((:cusparseSroti, :Float32),
                     (:cusparseDroti, :Float64))
    @eval begin
        function roti!(X::CuSparseVector{$elty},
                       Y::CuVector{$elty},
                       c::Number,
                       s::Number,
                       index::SparseChar)
            $fname(handle(), nnz(X), nonzeros(X), nonzeroinds(X), Y, c, s, index)
            X,Y
        end
        function roti(X::CuSparseVector{$elty},
                      Y::CuVector{$elty},
                      c::Number,
                      s::Number,
                      index::SparseChar)
            roti!(copy(X),copy(Y),c,s,index)
        end
    end
end

"""
    sctr!(X::CuSparseVector, Y::CuVector, index::SparseChar)

Set `Y[:] = X[:]` for dense `Y` and sparse `X`.
"""
sctr!(X::CuSparseVector, Y::CuVector, index::SparseChar)
for (fname,elty) in ((:cusparseSsctr, :Float32),
                     (:cusparseDsctr, :Float64),
                     (:cusparseCsctr, :ComplexF32),
                     (:cusparseZsctr, :ComplexF64))
    @eval begin
        function sctr!(X::CuSparseVector{$elty},
                       Y::CuVector{$elty},
                       index::SparseChar)
            $fname(handle(), nnz(X), nonzeros(X), nonzeroinds(X), Y, index)
            Y
        end
        function sctr(X::CuSparseVector{$elty},
                      index::SparseChar)
            sctr!(X, CUDA.zeros($elty, X.dims[1]),index)
        end
    end
end


text:
# Run mixture c-posterior on real-data examples
module RealData
using Distributions
using PyPlot
draw() = (pause(0.001); get_current_fig_manager()[:window][:raise_]())

# Settings
alphas = [20,100,500,1000,Inf]  # robustification params to use
Ns = [10^5,10^5,10^5,10^5,10^6]  # number of MCMC iterations for each alpha
nburns = [10^4,10^4,10^4,10^4,2.5*10^5]  # number of iterations to discard as burn-in for each alpha
data_ID = "Shapley"  # which dataset to use
m = 20 # maximum number of components

# Load data
if data_ID=="birth"
    dataset = readdlm("datasets/gestday.dat",' ')
    xmin,xmax = 180,380
elseif data_ID=="Shapley"
    dataset = readdlm("datasets/Shapley_galaxy.dat",' ')
    dataset = convert(Array{Float64,1}, dataset[:,4])/1000
    xmin,xmax = -5,50
elseif data_ID=="height"
    #D = readdlm("datasets/height_DEFG.dat",' ')
    D = readdlm("datasets/height_FG.dat",' ')
    subset = (18 .<= D[:,1])  # ages 18+
    dataset = vec(D[subset, 2])
    dataset = 0.393701*dataset  # convert from centimeters to inches
    xmin,xmax = minimum(dataset),maximum(dataset)
end
x = convert(Array{Float64,1}, dataset[:])
n = length(x)

# Stuff for mixture weights
a = 1/m # Gamma parameters
G = Gamma(a,1/1) # prior on weights (before conditioning on s>0)
lambda = 1.0  # Poisson parameter of the limiting prior distribution on k as m -> Inf
c = quantile(G,1-lambda/m)  # cutoff (Note: This choice makes p(k) \propto Binomial(k|m,lambda/m)I(k>0).)
sigma = 0.25 # scale for weight proposals
log_v_prior(v_i) = logpdf(G,v_i) # prior on latent weights v_i

# Stuff for mixture components
m0m,s0m = mean(x),std(x)  # data-dependent choices here
m0l = log(4/var(x))  # expect stddev of components to be around 1/2 the stddev of the overall data
s0l = 2.0

# Inference stuff
normpdf(x,m,l) = sqrt(l/(2*pi))*exp(-0.5*l*(x-m).*(x-m))
normlogpdf(x,m,s) = -0.5*log(2*pi) - log(s) - 0.5*(x-m)*(x-m)/(s*s)
likelihood(x,t) = normpdf(x,t[1],t[2])  # t=[mean,precision]
log_theta_prior(t) = normlogpdf(t[1],m0m,s0m) + normlogpdf(log(t[2]),m0l,s0l)
sm,sl = 0.2*s0m, 0.2
theta_prop(t) = Float64[t[1]+sm*randn(), t[2]*exp(randn()*sl)]
log_theta_prop(t,tp) = normlogpdf(tp[1],t[1],sm) + normlogpdf(log(tp[2]),log(t[2]),sl)
new_thetas(m) = [Float64[m0m,exp(m0l)] for i=1:m]

# Sampler code
include("core.jl")

# Run sampler
n_alphas = length(alphas)
k_posteriors = zeros(m,n_alphas)
for (i_a,alpha) in enumerate(alphas)

    # Run sampler
    n_samples = Ns[i_a]
    nburn = nburns[i_a]
    zeta = (1/n)/(1/n + 1/alpha)
    p,theta,k_r,v_r,art,arv,m_r,s_r = sampler(x,n_samples,m,c,sigma,zeta)

    # Compute posterior on k
    edges,counts = hist(k_r[nburn+1:end],[0:m])
    k_posteriors[:,i_a] = counts/(n_samples-nburn)

    # Plot traceplot of k
    figure(10,figsize=(8,3.2)); clf(); hold(true)
    plot(k_r,"k.",markersize=2)
    savefig("real-traceplot-alpha=$alpha.png",dpi=150)
    #draw()

    # Plot typical sample density
    figure(i_a,figsize=(8,3.2)); clf(); hold(true)
    subplots_adjust(bottom=0.2)
    xs = linspace(xmin,xmax,1000)
    density = zeros(length(xs))
    for i = 1:m
        if p[i]>0
            pixs = p[i]*normpdf(xs,theta[i][1],theta[i][2])
            density += pixs
            plot(xs,pixs)
        end
    end
    plot(xs,density,"k--",linewidth=3)
    xlim(xmin,xmax)
    if alpha<Inf; title("\$\\alpha = $(repr(int(alpha)))\$",fontsize=17)
    else; title("\$\\alpha = \\infty\$",fontsize=17)
    end
    xlabel("\$x\$",fontsize=16)
    ylabel("\$\\mathrm{density}\$",fontsize=15)
    #draw()
    savefig("real-density-alpha=$alpha.png",dpi=150)
end

# Plot posterior on k
figure(11,figsize=(8,3.2)); clf(); hold(true)
subplots_adjust(bottom=0.2)
colors = "bgyrm"
shapes = "ds^v*"
for (i_a,alpha) in enumerate(alphas)
    lab = "\$\\alpha = " * (alpha<Inf ? "$(repr(int(alpha)))\$" : "\\infty\$")
    plot(1:m,vec(k_posteriors[:,i_a]),"$(colors[i_a])$(shapes[i_a])-",label=lab,markersize=8,linewidth=2)
end
xlim(0,m)
xticks(1:m)
xlabel("\$k \\,\\,\\mathrm{(\\#\\,of\\,components)}\$",fontsize=17)
ylabel("\$\\Pi(k|\\mathrm{data})\$",fontsize=15)
legend(loc="upper right",numpoints=1)
draw()
savefig("real-k-posteriors.png",dpi=150)

# Plot data distribution
figure(12,figsize=(8,3.2)); clf(); hold(true)
subplots_adjust(bottom=0.2)
edges,counts = hist(x,linspace(xmin,xmax,81))
dx = edges[2]-edges[1]
bar(edges[1:end-1],(counts/n)/dx,dx,color="w")
xlim(xmin,xmax)
title("\$\\mathrm{Data\\,distribution}\$",fontsize=17)
xlabel("\$x\$",fontsize=16)
ylabel("\$\\mathrm{density}\$",fontsize=15)
draw()
savefig("real-data.png",dpi=150)

end # module

max_stars_repo_name: jwmi/CoarsenedPosterior

max_issues_repo_name: jwmi/CoarsenedPosterior

max_forks_repo_name: jwmi/CoarsenedPosterior


text:
module TestData
    importall Base # so that we get warnings for conflicts
    using Base.Test
    using DataTables
    using Compat

    #test_group("NullableArray creation")
    nvint = NullableArray(Nullable{Int}[1, 2, Nullable(), 4])
    nvint2 = NullableArray(5:8)
    nvint3 = NullableArray(5:8)
    nvflt = NullableArray(Nullable{Float64}[1.0, 2.0, Nullable(), 4.0])
    nvstr = NullableArray(Nullable{Compat.ASCIIString}["one", "two", Nullable(), "four"])
    dvdict = NullableArray(Dict, 4)    # for issue #199

    #test_group("constructors")
    dt1 = DataTable(Any[nvint, nvstr], [:Ints, :Strs])
    dt2 = DataTable(Any[nvint, nvstr])
    dt3 = DataTable(Any[nvint])
    dt4 = DataTable(Any[NullableArray(1:4), NullableArray(1:4)])
    dt5 = DataTable(Any[NullableArray([1,2,3,4]), nvstr])
    dt6 = DataTable(Any[nvint, nvint, nvstr], [:A, :B, :C])
    dt7 = DataTable(x = nvint, y = nvstr)
    @test size(dt7) == (4, 2)
    @test isequal(dt7[:x], nvint)

    #test_group("description functions")
    @test size(dt6, 1) == 4
    @test size(dt6, 2) == 3
    @test names(dt6) == [:A, :B, :C]
    @test names(dt2) == [:x1, :x2]
    @test names(dt7) == [:x, :y]

    #test_group("ref")
    @test isequal(dt6[2, 3], Nullable("two"))
    @test isnull(dt6[3, 3])
    @test isequal(dt6[2, :C], Nullable("two"))
    @test isequal(dt6[:B], nvint)
    @test size(dt6[[2,3]], 2) == 2
    @test size(dt6[2,:], 1) == 1
    @test size(dt6[[1, 3], [1, 3]]) == (2, 2)
    @test size(dt6[1:2, 1:2]) == (2, 2)
    @test size(head(dt6,2)) == (2, 3)
    # lots more to do

    #test_group("assign")
    dt6[3] = NullableArray(["un", "deux", "troix", "quatre"])
    @test isequal(dt6[1, 3], Nullable("un"))
    dt6[:B] = NullableArray([4, 3, 2, 1])
    @test isequal(dt6[1,2], Nullable(4))
    dt6[:D] = NullableArray([true, false, true, false])
    @test isequal(dt6[1,4], Nullable(true))
    delete!(dt6, :D)
    @test names(dt6) == [:A, :B, :C]
    @test size(dt6, 2) == 3

    #test_group("null handling")
    @test nrow(dt5[completecases(dt5), :]) == 3
    @test nrow(dropnull(dt5)) == 3
    returned = dropnull(dt4)
    @test dt4 == returned && dt4 !== returned
    @test nrow(dropnull!(dt5)) == 3
    returned = dropnull!(dt4)
    @test dt4 == returned && dt4 === returned

    #test_context("SubDataTables")

    #test_group("constructors")
    # single index is rows
    sdt6a = view(dt6, 1)
    sdt6b = view(dt6, 2:3)
    sdt6c = view(dt6, [true, false, true, false])
    @test size(sdt6a) == (1,3)
    sdt6d = view(dt6, [1,3], :B)
    @test size(sdt6d) == (2,1)

    #test_group("ref")
    @test isequal(sdt6a[1,2], Nullable(4))

    #test_context("Within")
    #test_group("Associative")

    #test_group("DataTable")
    srand(1)
    N = 20
    #Cast to Int64 as rand() behavior differs between Int32/64
    d1 = NullableArray(rand(map(Int64, 1:2), N))
    d2 = NullableCategoricalArray(Nullable{String}["A", "B", Nullable()])[rand(map(Int64, 1:3), N)]
    d3 = NullableArray(randn(N))
    d4 = NullableArray(randn(N))
    dt7 = DataTable(Any[d1, d2, d3], [:d1, :d2, :d3])

    #test_group("groupby")
    gd = groupby(dt7, :d1)
    @test length(gd) == 2
    # @test isequal(gd[2]["d2"], CategoricalVector["A", "B", Nullable(), "A", Nullable(), Nullable(), Nullable(), Nullable()])
    @test isequal(sum(gd[2][:d3]), sum(dt7[:d3][Vector(dt7[:d1]) .== 2]))

    g1 = groupby(dt7, [:d1, :d2])
    g2 = groupby(dt7, [:d2, :d1])
    @test isequal(sum(g1[1][:d3]), sum(g2[1][:d3]))

    res = Nullable(0.0)
    for x in g1
        res += sum(x[:d1])
    end
    @test isequal(res, sum(dt7[:d1]))

    @test aggregate(DataTable(a=1), identity) == DataTable(a_identity=1)

    dt8 = aggregate(dt7[[1, 3]], sum)
    @test isequal(dt8[1, :d1_sum], sum(dt7[:d1]))

    dt8 = aggregate(dt7, :d2, [sum, length], sort=true)
    @test isequal(dt8[1:2, :d2], NullableCategoricalArray(["A", "B"]))
    @test size(dt8, 1) == 3
    @test size(dt8, 2) == 5
    @test sum(dt8[:d1_length]) == N
    @test all(dt8[:d1_length] .> 0)
    @test dt8[:d1_length] == [4, 5, 11]
    @test isequal(dt8, aggregate(groupby(dt7, :d2, sort=true), [sum, length]))
    @test isequal(dt8[1, :d1_length], 4)
    @test isequal(dt8[2, :d1_length], 5)
    @test isequal(dt8[3, :d1_length], 11)
    @test isequal(dt8, aggregate(groupby(dt7, :d2), [sum, length], sort=true))

    dt9 = dt7 |> groupby([:d2], sort=true) |> [sum, length]
    @test isequal(dt9, dt8)
    dt9 = aggregate(dt7, :d2, [sum, length], sort=true)
    @test isequal(dt9, dt8)

    dt10 = DataTable(
        Any[[1:4;], [2:5;], ["a", "a", "a", "b" ], ["c", "d", "c", "d"]],
        [:d1, :d2, :d3, :d4]
    )

    gd = groupby(dt10, [:d3], sort=true)
    ggd = groupby(gd[1], [:d3, :d4], sort=true) # make sure we can groupby subdatatables
    @test ggd[1][1, :d3] == "a"
    @test ggd[1][1, :d4] == "c"
    @test ggd[1][2, :d3] == "a"
    @test ggd[1][2, :d4] == "c"
    @test ggd[2][1, :d3] == "a"
    @test ggd[2][1, :d4] == "d"

    #test_group("reshape")
    d1 = DataTable(a = NullableArray(repeat([1:3;], inner = [4])),
                   b = NullableArray(repeat([1:4;], inner = [3])),
                   c = NullableArray(randn(12)),
                   d = NullableArray(randn(12)),
                   e = NullableArray(map(string, 'a':'l')))

    stack(d1, :a)
    d1s = stack(d1, [:a, :b])
    d1s2 = stack(d1, [:c, :d])
    d1s3 = stack(d1)
    d1m = melt(d1, [:c, :d, :e])
    @test isequal(d1s[1:12, :c], d1[:c])
    @test isequal(d1s[13:24, :c], d1[:c])
    @test isequal(d1s2, d1s3)
    @test names(d1s) == [:variable, :value, :c, :d, :e]
    @test isequal(d1s, d1m)
    d1m = melt(d1[[1,3,4]], :a)
    @test names(d1m) == [:variable, :value, :a]

    # Test naming of measure/value columns
    d1s_named = stack(d1, [:a, :b], variable_name=:letter, value_name=:someval)
    @test names(d1s_named) == [:letter, :someval, :c, :d, :e]
    d1m_named = melt(d1[[1,3,4]], :a, variable_name=:letter, value_name=:someval)
    @test names(d1m_named) == [:letter, :someval, :a]

    stackdt(d1, :a)
    d1s = stackdt(d1, [:a, :b])
    d1s2 = stackdt(d1, [:c, :d])
    d1s3 = stackdt(d1)
    d1m = meltdt(d1, [:c, :d, :e])
    @test isequal(d1s[1:12, :c], d1[:c])
    @test isequal(d1s[13:24, :c], d1[:c])
    @test isequal(d1s2, d1s3)
    @test names(d1s) == [:variable, :value, :c, :d, :e]
    @test isequal(d1s, d1m)
    d1m = meltdt(d1[[1,3,4]], :a)
    @test names(d1m) == [:variable, :value, :a]

    d1s_named = stackdt(d1, [:a, :b], variable_name=:letter, value_name=:someval)
    @test names(d1s_named) == [:letter, :someval, :c, :d, :e]
    d1m_named = meltdt(d1, [:c, :d, :e], variable_name=:letter, value_name=:someval)
    @test names(d1m_named) == [:letter, :someval, :c, :d, :e]

    d1s[:id] = NullableArray([1:12; 1:12])
    d1s2[:id] =  NullableArray([1:12; 1:12])
    d1us = unstack(d1s, :id, :variable, :value)
    d1us2 = unstack(d1s2)
    d1us3 = unstack(d1s2, :variable, :value)
    @test isequal(d1us[:a], d1[:a])
    @test isequal(d1us2[:d], d1[:d])
    @test isequal(d1us2[:3], d1[:d])

    d2 = DataTable(id1 = [:a, :a, :a, :b],
                   id2 = [:A, :B, :B, :B],
                   id3 = [:t, :f, :t, :f],
                   val = [.1, .2, .3, .4])

    #test_group("merge")

    srand(1)
    dt1 = DataTable(a = shuffle!(NullableArray(1:10)),
                    b = NullableArray(rand([:A,:B], 10)),
                    v1 = NullableArray(randn(10)))

    dt2 = DataTable(a = shuffle!(NullableArray(1:5)),
                    b2 = NullableArray(rand([:A,:B,:C], 5)),
                    v2 = NullableArray(randn(5)))

    m1 = join(dt1, dt2, on = :a, kind=:inner)
    @test isequal(m1[:a], dt1[:a][dt1[:a].values .<= 5]) # preserves dt1 order
    m2 = join(dt1, dt2, on = :a, kind = :outer)
    @test isequal(m2[:a], dt1[:a]) # preserves dt1 order
    @test isequal(m2[:b], dt1[:b]) # preserves dt1 order
    # TODO: Re-enable
    m2 = join(dt1, dt2, on = :a, kind = :outer)
    # @test isequal(m2[:b2],
    #               NullableArray(Nullable{String}["A", "B", "B", "B", "B",
    #                                              Nullable(), Nullable(),
    #                                              Nullable(), Nullable(), Nullable()]))
    # @test isequal(m2[:b2],
    #               NullableArray(Nullable{String}["B", "B", "B", "C", "B",
    #                                              Nullable(), Nullable(),
    #                                              Nullable(), Nullable(), Nullable()]))

    dt1 = DataTable(a = NullableArray([1, 2, 3]),
                    b = NullableArray(["America", "Europe", "Africa"]))
    dt2 = DataTable(a = NullableArray([1, 2, 4]),
                    c = NullableArray(["New World", "Old World", "New World"]))

    m1 = join(dt1, dt2, on = :a, kind = :inner)
    @test isequal(m1[:a], NullableArray([1, 2]))

    m2 = join(dt1, dt2, on = :a, kind = :left)
    @test isequal(m2[:a], NullableArray([1, 2, 3]))

    m3 = join(dt1, dt2, on = :a, kind = :right)
    @test isequal(m3[:a], NullableArray([1, 2, 4]))

    m4 = join(dt1, dt2, on = :a, kind = :outer)
    @test isequal(m4[:a], NullableArray([1, 2, 3, 4]))

    # test with nulls (issue #185)
    dt1 = DataTable()
    dt1[:A] = NullableArray(["a", "b", "a", Nullable()])
    dt1[:B] = NullableArray([1, 2, 1, 3])

    dt2 = DataTable()
    dt2[:A] = NullableArray(["a", Nullable(), "c"])
    dt2[:C] = NullableArray([1, 2, 4])

    m1 = join(dt1, dt2, on = :A)
    @test size(m1) == (3,3)
    @test isequal(m1[:A], NullableArray(["a","a", Nullable()]))

    m2 = join(dt1, dt2, on = :A, kind = :outer)
    @test size(m2) == (5,3)
    @test isequal(m2[:A], NullableArray(["a", "b", "a", Nullable(), "c"]))

    srand(1)
    dt1 = DataTable(
        a = NullableArray(rand([:x,:y], 10)),
        b = NullableArray(rand([:A,:B], 10)),
        v1 = NullableArray(randn(10))
    )

    dt2 = DataTable(
        a = NullableArray([:x,:y][[1,2,1,1,2]]),
        b = NullableArray([:A,:B,:C][[1,1,1,2,3]]),
        v2 = NullableArray(randn(5))
    )
    dt2[1,:a] = Nullable()

    # # TODO: Restore this functionality
    # m1 = join(dt1, dt2, on = [:a,:b])
    # @test isequal(m1[:a], NullableArray(["x", "x", "y", "y", fill("x", 5)]))
    # m2 = join(dt1, dt2, on = ["a","b"], kind = :outer)
    # @test isequal(m2[10,:v2], Nullable())
    # @test isequal(m2[:a],
    #               NullableArray(Nullable{String}["x", "x", "y", "y",
    #                                              "x", "x", "x", "x", "x", "y",
    #                                              Nullable(), "y"])

    srand(1)
    function spltdt(d)
        d[:x1] = map(x -> x[1], d[:a])
        d[:x2] = map(x -> x[2], d[:a])
        d[:x3] = map(x -> x[3], d[:a])
        d
    end
    dt1 = DataTable(
        a = ["abc", "abx", "axz", "def", "dfr"],
        v1 = randn(5)
    )
    dt1 = spltdt(dt1)
    dt2 = DataTable(
        a = ["def", "abc","abx", "axz", "xyz"],
        v2 = randn(5)
    )
    dt2 = spltdt(dt2)

    # m1 = join(dt1, dt2, on = :a)
    # m2 = join(dt1, dt2, on = [:x1, :x2, :x3])
    # @test isequal(sort(m1[:a]), sort(m2[:a]))

    # test nonunique() with extra argument
    dt1 = DataTable(a = NullableArray(["a", "b", "a", "b", "a", "b"]),
                    b = NullableArray(1:6),
                    c = NullableArray([1:3;1:3]))
    dt = vcat(dt1, dt1)
    @test find(nonunique(dt)) == collect(7:12)
    @test find(nonunique(dt, :)) == collect(7:12)
    @test find(nonunique(dt, Colon())) == collect(7:12)
    @test find(nonunique(dt, :a)) == collect(3:12)
    @test find(nonunique(dt, [:a, :c])) == collect(7:12)
    @test find(nonunique(dt, [1, 3])) == collect(7:12)
    @test find(nonunique(dt, 1)) == collect(3:12)

    # Test unique() with extra argument
    @test isequal(unique(dt), dt1)
    @test isequal(unique(dt, :), dt1)
    @test isequal(unique(dt, Colon()), dt1)
    @test isequal(unique(dt, 2:3), dt1)
    @test isequal(unique(dt, 3), dt1[1:3,:])
    @test isequal(unique(dt, [1, 3]), dt1)
    @test isequal(unique(dt, [:a, :c]), dt1)
    @test isequal(unique(dt, :a), dt1[1:2,:])

    #test unique!() with extra argument
    unique!(dt, [1, 3])
    @test isequal(dt, dt1)
end


text:
using ForwardDiff

export AutodiffOptiModel

mutable struct AutodiffOptiModel <: AbstractOptiModel
  meta::OptiModelMeta
  f
  c
end

#  An AutodiffOptiModel instance is an AbstractOptiModel using ForwardDiff
#  to compute the derivatives. The represented problem is defined as
#
#    optimize    f(x) + g(x)
#    subject to  c(x) in S
#
#  Syntax:
#    AutodiffOptiModel(f, x0)
#    AutodiffOptiModel(f, x0, c, y0)

#=========================================================================#
function AutodiffOptiModel end

"
      AutodiffOptiModel(f, x0)
"
function AutodiffOptiModel(f, x0::AbstractVector{T};
                     name::String = "Generic", minimize::Bool=true) where T
  nvar = length(x0)
  @lencheck nvar x0
  meta = OptiModelMeta(nvar, 0, x0=x0, minimize=minimize, name=name)
  return AutodiffOptiModel(meta, f, x->T[])
end

"
      AutodiffOptiModel(f, x0, c, y0)
"
function AutodiffOptiModel(f, x0::AbstractVector{T}, c, y0::AbstractVector{T};
                     name::String = "Generic", minimize::Bool=true) where T
  nvar = length(x0)
  ncon = length(y0)
  @lencheck nvar x0
  @lencheck ncon y0
  meta = OptiModelMeta(nvar, ncon, x0=x0, y0=y0, minimize=minimize, name=name)
  return AutodiffOptiModel(meta, f, c)
end

#=========================================================================#
function obj(prob::AutodiffOptiModel, x::AbstractVector)
  @lencheck prob.meta.nvar x
  return prob.f(x)
end

function grad!(prob::AutodiffOptiModel, x::AbstractVector, dfx::AbstractVector)
  @lencheck prob.meta.nvar x dfx
  ForwardDiff.gradient!(dfx, prob.f, x)
  return dfx
end

#function prox!(prob::AutodiffOptiModel, x::AbstractVector, a::Real, z::AbstractVector)
#  @lencheck prob.meta.nvar x z
#  prox!(z, prob.g, x, a)
#  return z
#end

function cons!(prob::AutodiffOptiModel, x::AbstractVector, c::AbstractVector)
  @lencheck prob.meta.nvar x
  @lencheck prob.meta.ncon c
  c .= prob.c(x)
  return c
end

#function jac(prob::AutodiffOptiModel, x::AbstractVector)
#  @lencheck prob.meta.nvar x
#  return ForwardDiff.jacobian(prob.c, x)
#end

function jprod!(prob::AutodiffOptiModel, x::AbstractVector, v::AbstractVector, Jv::AbstractVector)
  @lencheck prob.meta.nvar x v
  @lencheck prob.meta.ncon Jv
  Jv .= ForwardDiff.derivative(t -> prob.c(x + t * v), 0)
  return Jv
end

function jtprod!(prob::AutodiffOptiModel, x::AbstractVector, v::AbstractVector, Jtv::AbstractVector)
  @lencheck prob.meta.nvar x Jtv
  @lencheck prob.meta.ncon v
  Jtv .= ForwardDiff.gradient(x -> dot(prob.c(x), v), x)
  return Jtv
end

#function objprox!(prob::AutodiffOptiModel, x::AbstractVector, a, z::AbstractVector)
#  @lencheck prob.meta.nvar x z
#  gz = prox!(z, prob.g, x, a)
#  return gz, z
#end
#
#function proj!(prob::AutodiffOptiModel, v::AbstractVector, p::AbstractVector)
#  @lencheck prob.meta.ncon v p
#  proj!(p, prob.S, v)
#  return p
#end


text:
# This fits Unteraargletscher (in its ITMIX incarnation) and compares the
# model on some radar lines (done with the f3.jl and f4.jl scripts)

# copied and adapted from ../itmix2-runs.jl

##########################################
## Fitting
using Parameters
include("../itmix-setup.jl")

######
# Choose glacier & options
######
plotyes = false
itmixversion = 2
r_tracks = [5,34,42]

glacier = :Unteraar
gid = BM.ITMIXGlacier(glacier, 2)
gl,gb,pp,pm,pn,pl = BM.init_forward(gid)

pm = BM.MPara(pm,
              sigma_h_model=20,
              error_model=[BM.ErrorModel.sqrerr, BM.ErrorModel.abserr][1],
              nthin_h=1)

####################
# Do it
####################
runtyp = [:test,:testmid,:prodlow,:prodmid,:prod][3]
fit_target = [BM.FitTarget.h, BM.FitTarget.iv, BM.FitTarget.h_iv, BM.FitTarget.length][3]
n_1d_btilde = 3
n_1d_fsl = 3
n_1d_temp = 1
rm_iv = false

date = "$(lowercase(Dates.monthabbr(Dates.today())))$(VAWTools.int2str2(Dates.day( Dates.today())))"
fl = "figs/f2-production-length995-btildes-0.3-ball-$date-$(join(r_tracks, "-"))-$fit_target-$(pm.error_model).jls"
println("\n $fl")

(gln, gbn, th0d,
 theta0, logposterior, logposterior1d,
 logprior, logprior_debug, loglikelihood, loglikelihood1d,
 fwdm_fn, fwdm1d_fn, pmcmc_defaults, fit_target) =
     setup_one_exp(r_tracks, gl, gb, pp, pm, pn,
                   n_1d_btilde, n_1d_fsl, n_1d_temp,
                   runtyp, fit_target, rm_iv,
                   th0_with_sigma=true)

# some tests:
print("Time to run forward model:")
@time fwdsol = fwdm_fn(theta0.th0);
# display(BM.plot_run(fwdsol, reuse=false, plot2d=[:h,:iv,:ivlog][1]))
print("Prior value:")
@show logprior(theta0.th0)[1]
#logprior_debug(theta0.th0)[1]
print("Posterior value:")
@show logposterior(theta0.th0)[1]
print("Time to calculate posterior:")
@time logposterior(theta0.th0)[1]

########
# max posterior estimation
#######
if false
    using Optim
    println("Optimizing for maximum posterior...")
    theta_max2 = optimize(theta_vec -> -logposterior(theta_vec)[1], float(theta0.th0), iterations=200, f_tol=1.0)
    # these don't work so well:
    #   theta_max2 = optimize(theta_vec -> -logposterior(theta_vec)[1], float(theta0.th0), GradientDescent())
    #theta_max2 = optimize(theta_vec -> -logposterior(theta_vec)[1], float(theta0.th0),  LBFGS())

    @show theta_max2
    error("asdf")
    theta0.th0[:] = theta_max2.minimizer
end

###############
## MCMC fitting
###############
println("Starting MCMC...")
pmcmc = BM.MCMCNum(;pmcmc_defaults...)

# sample posterior
sol = BM.mcmc(logposterior, theta0, pmcmc; verbose=true,
              misc=Dict(:logposterior=>logposterior,
                        :logposterior1d=>logposterior1d,
                        :logprior=>logprior,
                        :logprior_debug=>logprior_debug,
                        :loglikelihood=>loglikelihood,
                        :loglikelihood1d=>loglikelihood1d,
                        :fwdm_fn=>fwdm_fn,
                        :fit_target=>fit_target))
varnames = BM.get_varnames(theta0)
using KissMCMC
print_results(sol.thetas, sol.solverstats[:accept_ratio], sol.solverstats[:eff_sample_sizes], names=varnames)

# sample the prior to get that distribution too:
println("\n\nStarting MCMC of the prior to get that distribution too...")

pmcmc_prior = BM.MCMCNum(pmcmc, niter=10^5,
                         nthin = 50,
                         nburnin = 5*10^4,
                         nchains=100)

sol_prior = BM.mcmc(logprior, theta0, pmcmc_prior; verbose=true)
print_results(sol_prior.thetas, sol_prior.solverstats[:accept_ratio], sol_prior.solverstats[:eff_sample_sizes], names=varnames)

open(fl, "w") do io
    serialize(io, (sol, sol_prior))
end


text:
"""
structs.jl déclare les nouveaux types utilisés dans ce projet.
"""

import Base.show

# objet Station modélisant une station de navette
struct Station
	ind::Int
    name::String
    latitude::Float64
    longitude::Float64
end

Base.show(io::IO, x::Station) = print(io, x.name)

# objet Subdivision modélisant un lotissement urbain
struct Townsite
    name::String
    latitude::Float64
    longitude::Float64
    density::Int
    nearStations::Array{Station,1}
end

Base.show(io::IO, x::Townsite) = print(io, x.name)

getDensity(townsite::Townsite) = townsite.density

# objet Node modélisant un état du graphe Graphe
struct Node
    name::String
    shuttle::Int
    indStation::Int
    state::State
    sense::Sense
end

struct Graph
    name::String
    nbNodes::Int
    nbArcs::Int
end

max_stars_repo_name: LucasBaussay/ProjetODC

max_issues_repo_name: LucasBaussay/ProjetODC

max_forks_repo_name: LucasBaussay/ProjetODC


text:
# module stringrep

# import ..dlgo: Player, black, white, Move, Board

const BOARDCOLS = "ABCDEFGHJKLMNOPQRST"
const BOARDCOLSIDX = Dict(letter => i for (i, letter) in enumerate(BOARDCOLS))

function print_move(player::Player, move::Move)::Nothing
    move_str = if move.is_pass
        "passes"
    elseif move.is_resign
        "resigns"
    else
        string(BOARDCOLS[move.point.col]) * string(move.point.row)
    end

    println(tosymbol(player), " ", move_str)
end

function print_board(board::Board)
    for r in reverse(1:board.num_rows)
        bump = (r <= 9) ? " " : ""
        line = " "
        for c in 1:board.num_cols
            line *= tosymbol(board[r, c].color)
        end
        println(bump, r, line)
    end

    println("\n   " * BOARDCOLS[1:board.num_cols])
end

print_board(game_state::GameState) = print_board(game_state.board)

function point_from_coords(coords::AbstractString)::Point
    @assert length(coords) >= 2
    col = BOARDCOLSIDX[coords[1]]
    row = parse(Int, coords[2:end])
    return Point(row, col)
end
# end


text:
using LinearAlgebra

a = [1, 2, 3, 5, 5, 5, 6]
b = [1, 1, 2, 2, 3, 5]

function DTW(a, b)
    table = fill(Inf, size(a)[1]+1, size(b)[1]+1)
    table[1, 1] = 0
    for i in 2:size(table, 1)
        for j in 2:size(table, 2)
            table[i, j] = abs(a[i-1] - b[j-1])  # Calculate distance of current signal indices
            table[i, j] += min(Inf, table[i-1, j], table[i-1, j-1], table[i, j-1])
        end
    end
    return table[end, end], minimum_warp_path(table)
end

function minimum_warp_path(table)
    current_pos = size(table) |> collect
    at_start_index = false
    shortest_path = []
    push!(shortest_path, copy(current_pos))  # We start at the end and work back

    while at_start_index == false
        current_value = table[current_pos...]
        println(current_pos)
        # We search the direct left, left and up, and upward directions
        # for the smallest distance.
        indexes = ((current_pos[1] - 1, current_pos[2]),      # Left
                   (current_pos[1] - 1, current_pos[2] - 1),  # Up and Left
                   (current_pos[1], current_pos[2] - 1))      # Up

        values = [table[idx...] for idx in indexes]
        min_arg = argmin(values)  # Get direction had the smallest distance

        if min_arg == 1         # Left
            current_pos[1] -= 1
        elseif min_arg == 2     # Up and Left
            current_pos[1] -= 1
            current_pos[2] -= 1
        else                    # Up
            current_pos[2] -= 1
        end

        # Stop once you're at the beginning
        if current_pos == [1, 1]
            at_start_index = true
        end
        pushfirst!(shortest_path, copy(current_pos))
    end
    shortest_path = vcat(shortest_path'...)  # reshape
    return shortest_path
end


text:
export CUTENSORNETError

struct CUTENSORNETError <: Exception
    code::cutensornetStatus_t
end

Base.convert(::Type{cutensornetStatus_t}, err::CUTENSORNETError) = err.code

Base.showerror(io::IO, err::CUTENSORNETError) =
    print(io, "CUTENSORNETError: ", description(err), " (code $(reinterpret(Int32, err.code)), $(name(err)))")

name(err) = string(err.code)

function description(err::CUTENSORNETError)
    if err.code == CUTENSORNET_STATUS_SUCCESS
        "the operation completed successfully"
    elseif err.code == CUTENSORNET_STATUS_NOT_INITIALIZED
        "the library was not initialized"
    elseif err.code == CUTENSORNET_STATUS_ALLOC_FAILED
        "the resource allocation failed"
    elseif err.code == CUTENSORNET_STATUS_INVALID_VALUE
        "an invalid value was used as an argument"
    elseif err.code == CUTENSORNET_STATUS_ARCH_MISMATCH
        "an absent device architectural feature is required"
    elseif err.code == CUTENSORNET_STATUS_EXECUTION_FAILED
        "the GPU program failed to execute"
    elseif err.code == CUTENSORNET_STATUS_INTERNAL_ERROR
        "an internal operation failed"
    elseif err.code == CUTENSORNET_STATUS_NOT_SUPPORTED
        "the API is not supported by the backend."
    elseif err.code == CUTENSORNET_STATUS_LICENSE_ERROR
        "error checking current licensing."
    elseif err.code == CUTENSORNET_STATUS_CUBLAS_ERROR
        "a call to CUBLAS did not succeed."
    elseif err.code == CUTENSORNET_STATUS_CUDA_ERROR
        "some unknown CUDA error has occurred."
    elseif err.code == CUTENSORNET_STATUS_INSUFFICIENT_WORKSPACE
        "the workspace on the device is too small to execute."
    elseif err.code == CUTENSORNET_STATUS_INSUFFICIENT_DRIVER
        "the driver version is insufficient."
    elseif err.code == CUTENSORNET_STATUS_IO_ERROR
        "an error occurred related to file IO."
    elseif err.code == CUTENSORNET_STATUS_VERSION_MISMATCH
        "the dynamically linked cuTENSOR library is incompatible."
    else
        "no description for this error"
    end
end

## API call wrapper

# outlined functionality to avoid GC frame allocation
@noinline function throw_api_error(res)
    if res == CUTENSORNET_STATUS_ALLOC_FAILED
        throw(OutOfGPUMemoryError())
    else
        throw(CUTENSORNETError(res))
    end
end

macro check(ex, errs...)
    check = :(isequal(err, CUTENSORNET_STATUS_ALLOC_FAILED))
    for err in errs
        check = :($check || isequal(err, $(esc(err))))
    end

    quote
        res = @retry_reclaim err->$check $(esc(ex))
        if res != CUTENSORNET_STATUS_SUCCESS
            throw_api_error(res)
        end

        nothing
    end
end


text:
#===============================================================================

        Defines the struct `MCMC` which is used to determine an MCMC
        sampling procedure.

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

"""
    mutable struct MCMC
        updates::Vector
        updates_and_decorators::Vector
        backend::MCMCBackend
        schedule::MCMCSchedule
        workspace::GlobalWorkspace
    end

A struct for defining a list of updates and specifying the backend algorithm
doing the MCMC sampling. Any other functionality is done solely by the internal
mechanizms of function the `run!`.

    MCMC(
        updt_and_decor::Vector{<:Union{MCMCUpdate,MCMCUpdateDecorator}};
        backend=GenericMCMCBackend()
    )

The main constructor of an `MCMC` struct. It accepts an array of updates and
update decorators `updates_and_decorators` which together constitute a single
MCMC step. Additionally, `backend` specialized to a particular, overarching MCMC
algorithm (if it exists). An example of a backend would be
`DiffusionMCMCBackend` (from the package `DiffusionMCMC.jl`).
"""
mutable struct MCMC
    updates::Vector
    updates_and_decorators::Vector
    backend::MCMCBackend
    schedule::MCMCSchedule
    workspace::GlobalWorkspace

    function MCMC(
            updt_and_decor::Vector{<:Union{MCMCUpdate,MCMCUpdateDecorator}};
            backend=GenericMCMCBackend()
        )
        new(
            strip_decorators(updt_and_decor),
            updt_and_decor,
            backend
        )
    end
end

"""
    strip_decorators(updt_and_decor)

Remove all decorators from a list
"""
strip_decorators(updt_and_decor) = filter(u->!isdecorator(u), updt_and_decor)

"""
    get_decorators(updt_and_decor)

Retrieve all decorators from a list
"""
get_decorators(updt_and_decor) = filter(u->isdecorator(u), updt_and_decor)

"""
    init!(
        mcmc::MCMC,
        num_mcmc_steps,
        data,
        θinit,
        exclude_updates=[];
        kwargs...
    )

Initialize the schedule and the global workspace of the MCMC sampler for a given
set of updates (already saved in `mcmc.updates`), a total number of MCMC
iterations given by `num_mcmc_steps`, observed dataset `data`. `θinit` is the
initial value of the main parameter that the MCMC sampling is done for,
`exclude_updates` lists the update indices and the repspective ranges of mcmc
iterations from which these updates are supposed to be omitted from and `kwargs`
lists all additional named arguments passed for creating a global workspace.
"""
function init!(
        mcmc::MCMC,
        num_mcmc_steps,
        data,
        θinit,
        exclude_updates=[];
        kwargs...
    )
    mcmc.workspace = init_global_workspace(
        mcmc.backend,
        num_mcmc_steps,
        mcmc.updates,
        data,
        θinit;
        kwargs...
    )
    mcmc.schedule = MCMCSchedule(
        num_mcmc_steps,
        length(mcmc.updates),
        exclude_updates;
        extra_schedule_params(
            mcmc.workspace,
            mcmc.updates_and_decorators;
            kwargs...
        )...
    )
end

function extra_schedule_params(
        workspace::GlobalWorkspace,
        updates_and_decorators;
        kwargs...
    )
    tuple()
end


text:
function _next_state! end

abstract type ExpressionIterator end

"""
    ExpressionIterator(grammar::Grammar, max_depth::Int, sym::Symbol)

An iterator over all possible expressions of a grammar up to max_depth with start symbol sym.

Types of search depends on the order of production rules in the given grammar: BFS - terminals come first; DFS: nonterminals come first
"""
mutable struct ContextFreeIterator <: ExpressionIterator
    grammar::Grammar
    max_depth::Int
    sym::Symbol
end

Base.IteratorSize(::ExpressionIterator) = Base.SizeUnknown()

Base.eltype(::ExpressionIterator) = RuleNode

function Base.iterate(iter::ExpressionIterator)
    grammar, sym, max_depth = iter.grammar, iter.sym, iter.max_depth
    node = RuleNode(grammar[sym][1])
    if isterminal(grammar, node)
        return (deepcopy(node), node)
    else
        node, worked =  _next_state!(node, grammar, max_depth)
        while !worked
            # increment root's rule
            rules = grammar[sym]
            i = something(findfirst(isequal(node.ind), rules), 0)
            if i < length(rules)
                node, worked = RuleNode(rules[i+1]), true
                if !isterminal(grammar, node)
                    node, worked = _next_state!(node, grammar, max_depth)
                end
            else
                break
            end
        end
        return worked ? (deepcopy(node), node) : nothing
    end
end

function Base.iterate(iter::ExpressionIterator, state::RuleNode)
    grammar, max_depth = iter.grammar, iter.max_depth
    node, worked = _next_state!(state, grammar, max_depth)

    while !worked
        # increment root's rule
        rules = grammar[iter.sym]
        i = something(findfirst(isequal(node.ind), rules), 0)
        if i < length(rules)
            node, worked = RuleNode(rules[i+1]), true
            if !isterminal(grammar, node)
                node, worked = _next_state!(node, grammar, max_depth)
            end
        else
            break
        end
    end
    return worked ? (deepcopy(node), node) : nothing
end

"""
    count_expressions(grammar::Grammar, max_depth::Int, sym::Symbol)

Count the number of possible expressions of a grammar up to max_depth with start symbol sym.
"""
function count_expressions(grammar::Grammar, max_depth::Int, sym::Symbol)
    retval = 0
    for root_rule in grammar[sym]
        node = RuleNode(root_rule)
        if isterminal(grammar, node)
            retval += 1
        else
            node, worked = _next_state!(node, grammar, max_depth)
            while worked
                retval += 1
                node, worked = _next_state!(node, grammar, max_depth)
            end
        end
    end
    return retval
end

"""
    count_expressions(iter::ExpressionIterator)

Count the number of possible expressions in the expression iterator.
"""
count_expressions(iter::ExpressionIterator) = count_expressions(iter.grammar, iter.max_depth, iter.sym)