text:
# The [Neherlab COVID-19](https://neherlab.org/covid19/) forecast model

# Runs with 4 processes (if 4 cores CPU) - Only for optimised libraries
#using Distributed
#TOTAL_PROCS = 4
#Distributed.addprocs(TOTAL_PROCS - Distributed.nprocs())

include("COVID-19-parameters.jl")
include("COVID-19-utils.jl")
include("COVID-19-model.jl")
include("COVID-19-run-model.jl")
include("COVID-19-data.jl")
PRINT_DEBUG = false

using BlackBoxOptim

# Global variable that holds the optimised disease parameters
DiseaseParameters = DISEASE_INIT

# If running for the first time, or no updates for a long time
# updateData()

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

#-- FOR DEBUGGING RUNS
# COUNTRY_LIST = [
#     ("France",                      :north),
#     ("Italy",                       :north),
#     ("Spain",                       :north),
#     ("United Kingdom",              :north)]

# The list is focused on Europe
COUNTRY_LIST = [
    ("Austria",                     :north),
    ("Belgium",                     :north),
    ("Bulgaria",                    :north),
    ("Denmark",                     :north),
    ("France",                      :north),
    ("Germany",                     :north),
    ("Greece",                      :north),
    ("Italy",                       :north),
    ("Netherlands",                 :north),
    ("Norway",                      :north),
    ("Poland",                      :north),
    ("Spain",                       :north),
    ("Sweden",                      :north),
    ("United_Kingdom",              :north)]

countryData = Dict( c => populateCountryData(c, h, useOptimised = false) for (c, h) in COUNTRY_LIST)

plotVignette()
plotCountriestoDisk("__beforeOptim")
saveParameters()

#-------------------------------------------------------------------------------------------------
#--
#-- Global optimisation across all countries
#--

# Print only 3 decimals
using Printf
Base.show(io::IO, f::Float64) = @printf(io, "%1.3f", f)

N_RUNS = 1
for run in 1:N_RUNS

    #-------------------------------------------------------------------------------------------------
    #--
    #-- Optimisition all countries at once

    #-- Build a vector of all the countries' parameters with start date constraints
    fullRange = empty([(0.0, 0.0)])
    for (c1, _) in COUNTRY_LIST
        countryRange = COUNTRY_RANGE
        countryRange[COUNTRY_PARAM_START] = approximateModelStartRange(c1)
        fullRange = vcat(fullRange, countryRange)
    end

    #-- Optimise
    best = best_candidate(bboptimize(sumCountryLossesCountries,
                                    SearchRange = fullRange;
                                    Method = :adaptive_de_rand_1_bin,
                                    MaxTime = 120,
                                    NThreads = Threads.nthreads(),
                                    TraceMode = :compact))

    #-- Store the optimised parameters
    for i in 1:COUNTRY_N
        country, _ = COUNTRY_LIST[i]

        country_start_index = (i - 1) * COUNTRY_N + 1
        country_final_index = (i - 1) * COUNTRY_N + COUNTRY_N

        global countryData[country][:params] = best[country_start_index:country_final_index]
    end

    #-------------------------------------------------------------------------------------------------
    # Update Epidemiology
    updateEpidemiologyOnce(maxtime = 120)

    #-------------------------------------------------------------------------------------------------
    # Outputs
    saveParameters()
    plotCountriestoDisk(repr(now()));

end

plotVignette()

#-------------------------------------------------------------------------------------------------
# Normal runs: Optimise country, then disease after each country
N_RUNS = 10
for run in 1:N_RUNS
    println(); println("OPTIMISING THE WORST HALF OF COUNTRIES BY AVERAGE ERRORS---------------------------")
    scores = allSingleLosses(sorted = true)
    @show scores
    for c1 in first(scores, COUNTRY_N ÷ 2)[:, 2]
        global countryData

        # Make a note of the disease parameters
        p = countryData[c1][:params]
        print(c1); println("-----------------------------")
        print("Before                     "); @show p

        # Determine optimal parameters for each countryw
        countryRange = COUNTRY_RANGE
        countryRange[COUNTRY_PARAM_START] = approximateModelStartRange(c1)

        result = bboptimize(countryData[c1][:lossFunction],
                            SearchRange = countryRange;
                            Method = :adaptive_de_rand_1_bin,
                            MaxTime = 20,
                            TraceMode = :silent)

        print("After "); @show best_candidate(result)

        countryData[c1][:params] = best_candidate(result)
    end

    # After having done all the countries, epidemiology again more seriously
    println("OPTIMISING EPIDEMIOLOGY---------------------------")
    # Optimise across all countries (using "France" just to have disease parameters to optimise)q

    print("Before     "); @show DiseaseParameters
    result = bboptimize(sumCountriesLossDisease,
                        SearchRange = DISEASE_RANGE;
                        Method = :adaptive_de_rand_1_bin,
                        MaxTime = 60,
                        NThreads = Threads.nthreads(),
                        TraceMode = :silent)

    print("After "); @show best_candidate(result)
    println()

    global DiseaseParameters = best_candidate(result)

    saveParameters()
    plotCountriestoDisk(repr(now()))
end

plotVignette()

plotCountriestoDisk(repr(now()))
saveParameters()

allDiseaseParameters = DiseaseParameters
rename!(allDiseaseParameters, DISEASE_NAMES)

@show DiseaseParameters

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

println("OPTIMISING COUNTRIES---------------------------")
for (c1, _) in COUNTRY_LIST

    global countryData

    # Make a note of the disease parameters
    p = countryData[c1][:params]

    # Determine optimal parameters for each country
    result = bboptimize(countryData[c1][:lossFunction],
                        SearchRange = COUNTRY_RANGE;
                        Method = :adaptive_de_rand_1_bin,
                        MaxTime = 20,
                        TraceMode = :compact)

    println(); println(c1)
    print("Before         "); @show p
    print("After  "); @show best_candidate(result)

    countryData[c1][:params] = best_candidate(result)
end

#-------------------------------------------------------------------------------------------------
#--
#-- Global optimisation: all countries at once + disease parameters.
#--
fullRange = empty([(0.0, 0.0)])
for (c1, _) in COUNTRY_LIST
    countryRange = COUNTRY_RANGE
    countryRange[COUNTRY_PARAM_START] = approximateModelStartRange(c1)
    global fullRange = vcat(fullRange, countryRange)
end

fullRange = vcat(DISEASE_RANGE, fullRange)
result = bboptimize(fullEpidemyLoss,
                    SearchRange = fullRange;
                    Method = :adaptive_de_rand_1_bin,
                    MaxTime = 300,
                    NThreads = Threads.nthreads(),
                    TraceMode = :compact)

best = best_candidate(result)
DiseaseParameters = best[1:DISEASE_N]
for i in 1:COUNTRY_LIST_N
    country, _ = COUNTRY_LIST[i]

    country_start_index = DISEASE_N + (i - 1) * COUNTRY_N + 1
    country_final_index = DISEASE_N + (i - 1) * COUNTRY_N + COUNTRY_N

    global countryData[country][:params] = best[country_start_index:country_final_index]
end

plotVignette()
saveParameters()
plotCountriestoDisk(repr(now()));


text:
include("ukbb.jl")
include("estimators.jl")
include("utils.jl")
include("models.jl")
include("genotypes_and_queries.jl")
include("confounders.jl")
include("phenotypes.jl")
include("grm.jl")
include("sieve_plateau.jl")
include("summary.jl")

ext: jl


text:
import OpenCL: Buffer, CmdQueue, Context, Program, Kernel, call, build!

include("potentialCL.jl")
include("diffusionCL.jl")
include("laplacianCL.jl")
include("addCL.jl")
include("alignCL.jl")
include("calcRowCL.jl")
include("deltaCL.jl")
include("smulCL.jl")
include("deltaStep2CL.jl")
include("areaCL.jl")

### Define Module variables
ctx = nothing
NT = nothing
queue = nothing

potentialProgram = nothing
areaProgram = nothing
diffusionProgram = nothing
addProgram = nothing
alignProgram = nothing
rowProgram = nothing
deltaProgram = nothing
smulProgram = nothing
delta2Program = nothing
laplacianProgram = nothing

function createComputeContext{T <: FloatingPoint}(:: Type{T}, c, q)
    global ctx = c
    global NT = T
    global queue = q

    global potentialProgram = Program(ctx, source=getPotentialKernel(T)) |> build!
    global areaProgram = Program(ctx, source=getAreaKernel(T)) |> build!
    global diffusionProgram = Program(ctx, source=getDiffusionKernel(T)) |> build!
    global addProgram = Program(ctx, source=getAddKernel(T)) |> build!
    global alignProgram = Program(ctx, source=getAlignKernel(T)) |> build!
    global rowProgram = Program(ctx, source=getRowKernel(T)) |> build!
    global deltaProgram = Program(ctx, source=getDeltaKernel(T)) |> build!
    global smulProgram = Program(ctx, source=getSMulKernel(T)) |> build!
    global delta2Program = Program(ctx, source=getDelta2Kernel(T)) |> build!
    global laplacianProgram = Program(ctx, source=getLaplacianKernel(T)) |> build!
end

function destroyComputeContext()
    global ctx = nothing
    global queue = nothing
    global fieldResY = nothing
    global fieldResX = nothing

    global potentialProgram = nothing
    global areaProgram = nothing
    global diffusionProgram = nothing
    global addProgram = nothing
    global alignProgram = nothing
    global rowProgram = nothing
    global deltaProgram = nothing
    global smulProgram = nothing
    global delta2Program = nothing
    global laplacianProgram = nothing
end

###
# Define clfunctions
# IMPORTANT: Julia convention is func!(X, Y) overrides X with the operation func(X,Y)
# For the OCL functions the target is always the LAST BUFFER.

diffusion!(buff_Xfield, buff_Xpot, buff_Xlap) = diffusionCL!(buff_Xfield, buff_Xpot, buff_Xlap, fieldResY, fieldResX, ctx, queue, diffusionProgram, NT)
potential!(buff_Xfield, buff_Yfield, buff_Zfield, buff_Xpot, buff_Ypot, repulsion) = potentialCL!(buff_Xfield, buff_Yfield, buff_Zfield, buff_Xpot, buff_Ypot, convert(NT, repulsion), fieldResY, fieldResX, ctx, queue, potentialProgram)
area!(buff_in, buff_out) = areaCL!(buff_in, buff_out, long_direction, fieldResY, fieldResX, ctx, queue, areaProgram, NT)
align!(buff_Xfield, buff_Yfield, buff_OUTfield) = alignCL!(buff_Xfield, buff_Yfield, buff_OUTfield, convert(NT, attractionRate), convert(NT, stepIntegration), fieldResY, fieldResX, ctx, queue, alignProgram)
calcRow!(buff_Xfield, buff_Yfield, buff_Zfield, buff_Wfield, buff_OUT, x, y, z, w) = calcRowCL!(buff_Xfield, buff_Yfield, buff_Zfield, buff_Wfield, buff_OUT, x, y, z, w, fieldResY, fieldResX, ctx, queue, rowProgram, NT)
laplacian!(buff_in, buff_out) = laplacianCL!(buff_in, buff_out, fieldResY, fieldResX, ctx, queue, laplacianProgram)
delta!(buff1, buff2, buff3, buff4, buff5, buff6, buff_out, x1, x2, x3, x4, x5, x6) = deltaCL!(buff1, buff2, buff3, buff4, buff5, buff6, buff_out, x1, x2, x3, x4, x5, x6, fieldResY, fieldResX, ctx, queue, deltaProgram, NT)
delta2!(buff_D, buff_Xfield, buff_Yfield, buff_L,buff_out, decay) = delta2CL!(buff_D, buff_Xfield, buff_Yfield, buff_L, buff_out, decay, fieldResY, fieldResX, ctx, queue, delta2Program, NT)

smul!(X, buff_in, buff_out) =  smulCL!(X, buff_in, buff_out, fieldResY, fieldResX, ctx, queue, smulProgram)
c_add!(buff_in1, buff_in2, buff_out) = addCL!(buff_in1, buff_in2, buff_out, fieldResY, fieldResX, ctx, queue, addProgram)

#Julia style smul and add
smul!(buff, X) =  smulCL!(X, buff, buff, fieldResY, fieldResX, ctx, queue, smulProgram, NT)
c_add!(buff_out, buff_in) = addCL!(buff_out, buff_in, buff_out, fieldResY, fieldResX, ctx, queue, addProgram)

# Standard Julia Convention target first
c_copy!(target, source) = OpenCL.copy!(queue, target, source)
c_read(source) = OpenCL.read(queue, source)

create{T <: FloatingPoint}(:: Type{T}) = Buffer(T, ctx, :rw, fieldResX * fieldResY)
create_n4{T <: FloatingPoint}(:: Type{T}) = Buffer(T, ctx, :rw, fieldResX * fieldResY * 4)
function create_const_n4{T <: FloatingPoint}(x :: T)
    vals = fill(x, fieldResY * fieldResX * 4)
    buff = Buffer(T, ctx, :r, fieldResX * fieldResY * 4)
    OpenCL.copy!(queue, buff, vals)
    return buff
end


text:
Random.seed!(123)

N = 200
X = MLJBase.table(randn(Float32, N, 5));

# TODO: replace Short2 -> Short when
# https://github.com/FluxML/Flux.jl/pull/1618 is resolved:
builder = Short2(σ=identity)
optimiser = Flux.Optimise.ADAM()

losses = []

Random.seed!(123)
y = 1 .+ X.x1 - X.x2 .- 2X.x4 + X.x5
train, test = MLJBase.partition(1:N, 0.7)

@testset_accelerated "NeuralNetworkRegressor" accel begin

    Random.seed!(123)

    basictest(MLJFlux.NeuralNetworkRegressor,
              X,
              y,
              builder,
              optimiser,
              0.7,
              accel)

    # test model is a bit better than constant predictor:
    stable_rng = StableRNGs.StableRNG(123)
    model = MLJFlux.NeuralNetworkRegressor(builder=builder,
                                           acceleration=accel,
                                           rng=stable_rng)
    @time fitresult, _, rpt =
        fit(model, 0, MLJBase.selectrows(X, train), y[train])
    first_last_training_loss = rpt[1][[1, end]]
    push!(losses, first_last_training_loss[2])
#    @show first_last_training_loss
    yhat = predict(model, fitresult, selectrows(X, test))
    truth = y[test]
    goal = 0.9*model.loss(truth .- mean(truth), 0)
    @test model.loss(yhat, truth) < goal

    optimisertest(MLJFlux.NeuralNetworkRegressor,
                  X,
                  y,
                  builder,
                  optimiser,
                  accel)

end

# check different resources (CPU1, CUDALibs, etc)) give about the same loss:
reference = losses[1]
@test all(x->abs(x - reference)/reference < 1e-6, losses[2:end])

Random.seed!(123)
ymatrix = hcat(1 .+ X.x1 - X.x2, 1 .- 2X.x4 + X.x5);
y = MLJBase.table(ymatrix);

losses = []

@testset_accelerated "MultitargetNeuralNetworkRegressor" accel begin

    Random.seed!(123)

    basictest(MLJFlux.MultitargetNeuralNetworkRegressor,
              X,
              y,
              builder,
              optimiser,
              1.0,
              accel)

    # test model is a bit better than constant predictor
    model = MLJFlux.MultitargetNeuralNetworkRegressor(acceleration=accel,
                                                      builder=builder)
    @time fitresult, _, rpt =
        fit(model, 0, MLJBase.selectrows(X, train), selectrows(y, train))
    first_last_training_loss = rpt[1][[1, end]]
    push!(losses, first_last_training_loss[2])
#   @show first_last_training_loss
    yhat = predict(model, fitresult, selectrows(X, test))
    truth = ymatrix[test,:]
    goal = 0.8*model.loss(truth .- mean(truth), 0)
    @test model.loss(Tables.matrix(yhat), truth) < goal

    optimisertest(MLJFlux.MultitargetNeuralNetworkRegressor,
              X,
              y,
              builder,
              optimiser,
              accel)

end

# check different resources (CPU1, CUDALibs, etc)) give about the same loss:
reference = losses[1]
@test all(x->abs(x - reference)/reference < 1e-6, losses[2:end])

true


text:
using Marble
reload("Marble")

# Streams
path = Marble.mrbldir("streams/$(string(Base.Random.uuid4()))")
settings = Marble.get_settings(path)
Marble.create_paths(settings)
cache = Marble.Cache("$(settings["paths"]["cache"])/hashes.json")
doc = Marble.MarbleDoc("stream", readstring("$(Pkg.dir("Marble"))//test/docs/test.md"), cache, settings)
Marble.parse(doc)
Marble.process(doc)
Marble.render(doc)
Marble.template(doc)

# File
basename = Marble.get_basename(doc)
texfile = "$(doc.settings["paths"]["base"])/$(Marble.get_basename(doc)).tex"
builddir = abspath(doc.settings["paths"]["build"])
command = `latexmk -$(doc.settings["texcmd"]) -shell-escape -halt-on-error $texfile`
Marble.runindir(builddir) do
    run(command)
end

# Build File
using Marble
reload("Marble")
file_path = Pkg.dir("Marble") * "/test/docs/test.md"
fpath = "/Users/dustinrb/Desktop/t.md"
cd("/Users/dustinrb/Desktop/")
build_path = Marble.get_build_dir(file_path)
run(`open $build_path`)
Marble.clean_tex(file_path)
reload("Marble")
Marble.build_file(fpath, out="/Users/dustinrb/Desktop/test.pdf")
Marble.build_file(file_path, out="/Users/dustinrb/Desktop/test2.pdf")

# Build folder
fpath = "/Users/dustinrb/Desktop/test2"
build_dir(fpath)

# Init Document
using Marble
reload("Marble")
Marble.init_dir("/Users/dustinrb/Desktop/testdoc_$(randstring(3))")
Marble.init_dir("/Users/dustinrb/Desktop/testdoc_$(randstring(3))"; template="test")
Marble.init_dir("/Users/dustinrb/Desktop/testdoc_$(randstring(3))"; template="fail")

run(`open $(pwd())`)
run(`open $path`)


text:
## Compile rand vars into callable functions
## =========================================

dims(X::JuliaRandVar) = X.dims

"Evaluate a random variable on an element of scenario ω"
(X::JuliaRandVar)(ω::Omega) = Base.invokelatest(X.func, ω)

"Compile SymbolicRandVar into JuliaRandVar"
convert{T}(::Type{JuliaRandVar{T}}, X::SymbolicRandVar{T}) =
  JuliaRandVar{T}(compile(X), dims(X))

all_functional_randvars = union(real_real_real, real_real, real_floating,
                                real_real_bool, ((:IfElseRandVar, :ifelse),),
                                bool_bool_bool, ((:NotRandVar, :!),))
for (name,op) in all_functional_randvars
  eval(
  quote
  function convert(::Type{Expr}, X::$name)
    Expr(:call, $op, [convert(Expr, arg) for arg in args(X)]...)
  end
  end)
end

convert(::Type{Expr}, X::ConstantRandVar) = X.val
convert(::Type{Expr}, X::OmegaRandVar) = :(ω[$(X.dim)])

"lambda `Expr` of `X` e.g., `ω -> w * b + c`"
lambda_expr(X::SymbolicRandVar) = Expr(:(->), :ω,convert(Expr,X))

"Compile `X` into Julia lambda"
compile(X::SymbolicRandVar)::Function = eval(lambda_expr(X))

"Convert Elementary `X` to `ω -> quantile_X(ω)`"
function convert{T <: ElementaryRandVar}(::Type{Expr}, X::T)
  params = [convert(Expr, arg) for arg in args(X)]
  Expr(:call, :quantile, distribution_type(T), :(ω[$X.dim]), params...)
end

"Convert a random variable into an executable random variable"
executionalize{T}(X::RandVar{T}) = convert(JuliaRandVar{T}, X)


text:
using Test, CMDimData

function printsep(title)
	println("\n", title, "\n", repeat("-", 80))
end

function show_testset_description()
	@info Test.get_testset().description
end

@testset "CMDimData tests" begin
	testfiles = ["EasyPlot.jl", "EasyData.jl"]

	for testfile in testfiles
		include(testfile)
	end

end #testset

:Test_Complete


text:
export QBasedPolicy

"""
    QBasedPolicy(;learner::Q, explorer::S)

Use a Q-`learner` to generate the estimations of actions and use `explorer` to get the action.
"""
Base.@kwdef struct QBasedPolicy{Q<:AbstractLearner,E<:AbstractExplorer} <: AbstractPolicy
    learner::Q
    explorer::E
end

(π::QBasedPolicy)(obs, ::MinimalActionSet) = obs |> π.learner |> π.explorer
(π::QBasedPolicy)(obs, ::FullActionSet) =
    π.explorer(π.learner(obs), get_legal_actions_mask(obs))

function RLBase.update!(p::QBasedPolicy, t::AbstractTrajectory)
    experience = extract_experience(t, p)
    isnothing(experience) || update!(p.learner, experience)
end

RLBase.update!(
    p::QBasedPolicy,
    m::AbstractEnvironmentModel,
    t::AbstractTrajectory,
    n::Int,
) = update!(p.learner, m, t, n)

RLBase.extract_experience(trajectory::AbstractTrajectory, p::QBasedPolicy) =
    extract_experience(trajectory, p.learner)
RLBase.get_prob(p::QBasedPolicy, obs, ::MinimalActionSet) =
    get_prob(p.explorer, p.learner(obs))
RLBase.get_prob(p::QBasedPolicy, obs, ::FullActionSet) =
    get_prob(p.explorer, p.learner(obs), get_legal_actions_mask(obs))

RLBase.get_priority(p::QBasedPolicy, experience) = get_priority(p.learner, experience)


text:
module Manifests
include("DBFs.jl"); using .DBFs;
include("PointArrays.jl"); using .PointArrays;

# using JLD2
using BigArrays
using BigArrays.GSDicts

using ..RealNeuralNetworks.NodeNets
#import ..RealNeuralNetworks.NodeNets.DBFs
#import ..RealNeuralNetworks.NodeNets.PointArrays
using OffsetArrays
using JSON
#import Distributed: pmap

const MIP_LEVEL = 4

export Manifest

struct Manifest
    # the bigarray for cutout of segmentation
    ba          ::AbstractBigArray
    # the id of object
    obj_id      ::Integer
    # the unit range list for cutout
    rangeList   ::Vector
end

"""
Parameters:
    dir
"""
function Manifest(manifestDirPath::AbstractString, manifestKey::AbstractString,
                  bigArrayPath::AbstractString, mip::Integer )
    println("big array path: ", bigArrayPath)

    ba = BigArray( GSDict( bigArrayPath ); mip=mip )
    h = GSDict( manifestDirPath )
    d = h[manifestKey]
    ranges = d[:fragments]
    Manifest( ranges, ba )
end

"""
example: ["770048087:0:2968-3480_1776-2288_16912-17424"]
"""
function Manifest( ranges::Vector, ba::BigArray{D,T} ) where {D,T}
    obj_id = Meta.parse( split(ranges[1], ":")[1] )
    obj_id = convert(T, obj_id)
    ranges = map(x-> split(x,":")[end], ranges)
    rangesList = map( BigArrays.Indexes.string2unit_range, ranges )
    # convert from z,y,x to x,y,z
    # this is due to a bug in chunkflow, we should not need this in the future
    rangesList = map( reverse, rangesList)
    @show rangesList
    Manifest( ba, obj_id, rangesList )
end

function Base.length(self::Manifest) length(get_range_list(self)) end

function get_range_list(self::Manifest) self.rangeList end

"""
the voxel offset in the neuroglancer precomputed info file
"""
function get_voxel_offset(self::Manifest)
    voxel_offset =  self.ba.kvStore.configDict[:offset]
    #real voxel offset should based on the highest resolution 5x5x45
    # we are using mip level 4 with (80x80x45 nm) now
    # voxel_offset = map((x,y)->UInt32(x*y), voxel_offset, (2^MIP_LEVEL, 2^MIP_LEVEL,1))
    voxel_offset = Vector{UInt32}(voxel_offset)
    @show voxel_offset
    return (voxel_offset...,)
end

"""
iterate the chunks containing the neuron with specified neuronId
build point cloud and dbf when iterating the chunks
"""
function trace(self::Manifest)
    println("extract point clouds and distance from boundary fields ...")

    pointCloudDBFList = map(x->_get_point_cloud_dbf(self, x), self.rangeList )

    pointClouds = map( x->x[1], pointCloudDBFList )
    pointCloud = vcat(pointClouds ...)
    dbfs = map(x->x[2], pointCloudDBFList)
    dbf = vcat(dbfs ...)
    # save temporal variables for debug
    # @save "/tmp/$(neuronId).jld" pointClouds, pointCloud, dbf
    println("skeletonization from global point cloud and dbf using RealNeuralNetworks algorithm...")
    @time nodeNet = NodeNet(pointCloud; dbf=dbf)
    return nodeNet
end

function _get_point_cloud_dbf(self::Manifest, ranges::Vector)
    # example: [2456:2968, 1776:2288, 16400:16912]
    offset = (map(x-> UInt32(x.start-1), ranges)...,)
    seg = self.ba[ranges...] |> parent
    bin_im = DBFs.create_binary_image( seg; obj_id = self.obj_id )
    @assert any(bin_im)
    point_cloud = PointArrays.from_binary_image( bin_im )
    # distance from boundary field
    dbf = DBFs.compute_DBF(point_cloud, bin_im)
    PointArrays.add_offset!(point_cloud, offset)
    # no need to use voxel_offset since the file name encoded the global coordinate
    # PointArrays.add_offset!(point_cloud, get_voxel_offset(self))
    return point_cloud, dbf
end

end # module


text:
"""
    LJForceMagnitude(r,s,d,ϵ)

Returns magnitude of force resulting from Lennard-Jones-type potential.
r is the distance between particles, s is a scale parameter for the
strength of interaction, d is the equilibrium distance between particles,
and ϵ is the interaction extension, typically a small positive number.
"""
function LJForceMagnitude(r::T,s::T,d::T,ϵ::T) where T<:Real
    return ( r > (1+ϵ)*d ? zero(T) : 2*s*d/r^2*(1 - d/r) )
end

"""
    ForceCalculation(r,s,d,ϵ,Δx,Δy)

Returns 2D cohesion force between two particles. s is a scale parameter for
the strength of interaction, d is the equilibrium distance between particles,
and ϵ is the interaction extension, typically a small positive number.
"""
function ForceCalculation(s::T,d::T,ϵ::T,Δx::T,Δy::T) where T<:Real
    r = sqrt(Δx^2 + Δy^2)
    tx = Δx/r
    ty = Δy/r

    LJmag = LJForceMagnitude(r,s,d,ϵ)

    Fx = tx*LJmag
    Fy = ty*LJmag

    return Fx,Fy
end

"""
    computeCohesion_backup!(cfX,cfY,particleList,rc,ϵ)

Backup method to directly compute cohesion force on all particles. cfX and cfY are arrays storing
forces on each particle. particleList is an arrays of the particles. r is the distance between
particles, s is a scale parameter for the strength of interaction, and ϵ is the interaction
extension, typically a small positive number.

Note that this method is *slow* for large number of particles. It is O(N^2) where N is the number
of particles (the length of nodeList).
"""
function compute_cohesion_force_backup!(cfX::Vector{T},cfY::Vector{T},particleList::Vector{P},s::T,ϵ::T) where {T<:Real,P<:AbstractParticle}
    Nparticles = length(particleList)

    fill!(cfX,zero(T))
    fill!(cfY,zero(T))

    for ti=1:Nparticles, tj=(ti+1):Nparticles       # takes advantage of force anti-symmetry
        Δx = particleList[tj].pos.x - particleList[ti].pos.x
        Δy = particleList[tj].pos.y - particleList[ti].pos.y

        d = particleList[ti].radius + particleList[tj].radius

        fx,fy = ForceCalculation(s,d,ϵ,Δx,Δy)
        cfX[ti] += fx; cfY[ti] += fy
        cfX[tj] -= fx; cfY[tj] -= fy
    end
    nothing
end

function compute_adhesion_force!(afX::Vector{T},afY::Vector{T},particleList::Vector{P},wallList::Vector{W},pointOnWall::Point2D{T},s::T,ϵ::T) where {T<:Real,P<:AbstractParticle,W<:AbstractWall}
    fill!(afX,zero(T))
    fill!(afY,zero(T))

    for tp=1:length(particleList), tw=1:length(wallList)
        get_nearest_point!(pointOnWall,particleList[tp],wallList[tw])
        if is_in_line(wallList[tw],pointOnWall)
            Δx = pointOnWall.x - particleList[tp].pos.x
            Δy = pointOnWall.y - particleList[tp].pos.y

            d = particleList[tp].radius + wallList[tw].thickness/2

            fx,fy = ForceCalculation(s,d,ϵ,Δx,Δy)
            afX[tp] += fx
            afY[tp] += fy
        end
    end
    nothing
end

function compute_gravity_force!(gfX,gfY,n_particles,G)
    for ti=1:n_particles
        gfX[ti] = 0.0
        gfY[ti] = -G
    end
    nothing
end

function compute_cohesion_force_cl!(cfX::Vector{T},cfY::Vector{T},cl::CellList,pList::Vector{P},s::T,ϵ::T) where {T<:Real,P<:AbstractParticle}
    # zero-out cohesion
    fill!(cfX,zero(T))
    fill!(cfY,zero(T))

    # go over cell lists
    for cellInd = 1:length(cl.cells)
        # compute cell-cell (same cell) forces
        numNodes = length(cl.cells[cellInd].particleIDList)
        for ti=1:numNodes, tj=(ti+1):numNodes
            nodeI = cl.cells[cellInd].particleIDList[ti]
            nodeJ = cl.cells[cellInd].particleIDList[tj]

            Δx = pList[nodeJ].pos.x - pList[nodeI].pos.x
            Δy = pList[nodeJ].pos.y - pList[nodeI].pos.y

            d = pList[nodeI].radius + pList[nodeJ].radius

            fx,fy = ForceCalculation(s,d,ϵ,Δx,Δy)
            cfX[nodeI] += fx; cfY[nodeI] += fy
            cfX[nodeJ] -= fx; cfY[nodeJ] -= fy
        end

        # compute cell-neighbor forces
        for ncInd in cl.cells[cellInd].neighborList
            numNeighborNodes = length(cl.cells[ncInd].particleIDList)
            for ti=1:numNodes, tj=1:numNeighborNodes
                nodeI = cl.cells[cellInd].particleIDList[ti]
                nodeJ = cl.cells[ncInd].particleIDList[tj]

                Δx = pList[nodeJ].pos.x - pList[nodeI].pos.x
                Δy = pList[nodeJ].pos.y - pList[nodeI].pos.y

                d = pList[nodeI].radius + pList[nodeJ].radius

                fx,fy = ForceCalculation(s,d,ϵ,Δx,Δy)
                cfX[nodeI] += fx
                cfY[nodeI] += fy
            end
        end
    end

    nothing
end


text:
include("hyperelasticity.jl")

ts = ExplicitSolverTime(Δt, NT)
ubd, abd = compute_boundary_info(domain, globaldata, ts)
Fext = compute_external_force(domain, globaldata, ts)

d0 = zeros(2domain.nnodes)
v0 = zeros(2domain.nnodes)
a0 = zeros(2domain.nnodes)

# general neural network
function nn_law(ε)
  ae(ε, [20,20,20,3])
end

# # free energy
# function nn_law(ε)
#   φ = squeeze(ae(ε, [20,20,20,1]))
#   tf.gradients(φ, ε)
# end

# # parameter
# c1 = Variable(0.5) * 1e-1
# c2 = Variable(0.5) * 1e-1
# function nn_law(ε)
#   compute_stress_rivlin_saunders(ε, c1, c2)
# end

d, v, a= ExplicitSolver(globaldata, domain, d0, v0, a0, Δt, NT, nn_law, Fext, ubd, abd; strain_type="finite")

idx = @. (0:n)*m + 1
idx = [idx; (@. idx + domain.nnodes)]
loss = sum((d[:, idx] - d_[:, idx])^2)
sess = Session(); init(sess)
@info run(sess, loss)
# error()

for i = 1:100
  BFGS!(sess, loss, 1000)
  ADCME.save(sess, "data1_$i.mat")
end


text:
using ValueHistoriesLogger
using Test

@testset "ValueHistoriesLogger.jl" begin
    # Write your own tests here.
end


text:
# symmetric, skew-symmetric and symmetric checks
"""
```julia
symmetric(::SecondOrderTensor)
symmetric(::FourthOrderTensor)
```
Computes the symmetric part of a second or fourth order tensor.
For a fourth order tensor, the symmetric part is the same as the minor symmetric part.
Returns a `SymmetricTensor`.

**Example:**

```jldoctest
julia> A = rand(Tensor{2,2})
2×2 ContMechTensors.Tensor{2,2,Float64,4}:
 0.590845  0.566237
 0.766797  0.460085

julia> symmetric(A)
2×2 ContMechTensors.SymmetricTensor{2,2,Float64,3}:
 0.590845  0.666517
 0.666517  0.460085
```
"""
@inline symmetric(S1::SymmetricTensors) = S1

@generated function symmetric{dim, T}(t::Tensor{2, dim, T})
    exps = Expr[]
    for j in 1:dim, i in j:dim
        if i == j
            push!(exps, :(get_data(t)[$(compute_index(Tensor{2, dim}, i, j))]))
        else
            I = compute_index(Tensor{2, dim}, i, j)
            J = compute_index(Tensor{2, dim}, j, i)
            push!(exps, :((get_data(t)[$I] + get_data(t)[$J])/2))
        end
    end
    exp = Expr(:tuple, exps...)
    return quote
        $(Expr(:meta, :inline))
        SymmetricTensor{2, dim}($exp)
    end
end

"""
```julia
minorsymmetric(::FourthOrderTensor)
```
Computes the minor symmetric part of a fourth order tensor, returns a `SymmetricTensor{4}`.
"""
@generated function minorsymmetric{dim}(t::Tensor{4, dim})
    exps = Expr[]
    for l in 1:dim, k in l:dim, j in 1:dim, i in j:dim
        if i == j && k == l
            push!(exps, :(data[$(compute_index(Tensor{4, dim}, i, j, k, l))]))
        else
            I = compute_index(Tensor{4, dim}, i, j, k, l)
            J = compute_index(Tensor{4, dim}, j, i, k, l)
            K = compute_index(Tensor{4, dim}, i, j, k, l)
            L = compute_index(Tensor{4, dim}, i, j, l, k)
            push!(exps, :((data[$I] + data[$J] + data[$K] + data[$L]) / 4))
        end
    end
    exp = Expr(:tuple, exps...)
    return quote
        $(Expr(:meta, :inline))
        data = get_data(t)
        SymmetricTensor{4, dim}($exp)
    end
end

@inline minorsymmetric(t::SymmetricTensors) = t

@inline symmetric(t::Tensor{4}) = minorsymmetric(t)

"""
```julia
majorsymmetric(::FourthOrderTensor)
```
Computes the major symmetric part of a fourth order tensor, returns a `Tensor{4}`.
"""
@generated function majorsymmetric{dim}(t::FourthOrderTensor{dim})
    exps = Expr[]
    for l in 1:dim, k in 1:dim, j in 1:dim, i in 1:dim
        if i == j == k == l || i == k && j == l
            push!(exps, :(data[$(compute_index(get_base(t), i, j, k, l))]))
        else
            I = compute_index(get_base(t), i, j, k, l)
            J = compute_index(get_base(t), k, l, i, j)
            push!(exps, :((data[$I] + data[$J]) / 2))
        end
    end
    exp = Expr(:tuple, exps...)
    return quote
        $(Expr(:meta, :inline))
        data = get_data(t)
        Tensor{4, dim}($exp)
    end
end

"""
```julia
skew(::SecondOrderTensor)
```
Computes the skew-symmetric (anti-symmetric) part of a second order tensor, returns a `Tensor{2}`.
"""
@inline skew(S1::Tensor{2}) = (S1 - S1.') / 2
@inline skew{dim,T}(S1::SymmetricTensor{2,dim,T}) = zero(Tensor{2,dim,T})

# Symmetry checks
@inline Base.issymmetric(t::Tensor{2, 1}) = true
@inline function Base.issymmetric(t::Tensor{2, 2})
    data = get_data(t)
    @inbounds return data[2] == data[3]
end
@inline function Base.issymmetric(t::Tensor{2, 3})
    data = get_data(t)
    @inbounds begin
        return (data[2] == data[4] &&
                data[3] == data[7] &&
                data[6] == data[8])
    end
end

function isminorsymmetric{dim}(t::Tensor{4, dim})
    @inbounds for l in 1:dim, k in l:dim, j in 1:dim, i in j:dim
        if t[i,j,k,l] != t[j,i,k,l] || t[i,j,k,l] != t[i,j,l,k]
            return false
        end
    end
    return true
end

isminorsymmetric(::SymmetricTensor{4}) = true

function ismajorsymmetric{dim}(t::FourthOrderTensor{dim})
    @inbounds for l in 1:dim, k in l:dim, j in 1:dim, i in j:dim
        if t[i,j,k,l] != t[k,l,i,j]
            return false
        end
    end
    return true
end

Base.issymmetric(t::Tensor{4}) = isminorsymmetric(t)

Base.issymmetric(::SymmetricTensors) = true


text:
mutable struct FixedType <: NamedType
    type::String
    namespace::String
    name::String
    aliases::Union{Vector{String}, Nothing}
    size::Int
end

==(x::FixedType, y::FixedType) = x.name == y.name && x.size == y.size
FixedType(name, size, namespace="", aliases=nothing) =
    FixedType("fixed", namespace, name, aliases, size)
FixedType() = FixedType("fixed", "", "", nothing, 0)
StructTypes.StructType(::Type{FixedType}) = StructTypes.Mutable()

juliatype(x::FixedType) = NTuple{x.size, UInt8}
schematype(::StructTypes.ArrayType, ::Type{NTuple{N, UInt8}}, name="Fixed_$N") where {N} =
    FixedType(name, N)

function writevalue(::Binary, ::FixedType, x::NTuple{N, UInt8}, buf, pos, len, opts) where {N}
    @check N
    for i = 1:N
        @inbounds buf[pos] = x[i]
        pos += 1
    end
    return pos
end

nbytes(::FixedType, x::NTuple{N, UInt8}) where {N} = N

mutable struct FixedClosure
    pos::Int
end

@inline function (f::FixedClosure)(i::Int, buf)
    @inbounds b = buf[f.pos]
    f.pos += 1
    return b
end

function readvalue(::Binary, ::FixedType, ::Type{NTuple{N, UInt8}}, buf, pos, len, opts) where {N}
    f = FixedClosure(pos)
    x = ntuple(i -> f(i, buf), Val(N))
    return x, f.pos
end

skipvalue(::Binary, ::FixedType, ::Type{NTuple{N, UInt8}}, buf, pos, len, opts) where {N} = pos + N


text:
# The type hierarchy is as follows:
# abstract Ball
# |-> abstract UnitBall: radius is 1
#     |-> StaticUnitBall: dimension is part of type
#     |-> DynamicUnitBall: dimension is specified by int field
# |-> GenericBall: stores center and radius. Here, the dimension can be
#     obtained from the center, so there is only one type.
# There are aliases for SVector{N,T} and Vector{T}:
#   EuclideanBall, VectorBall,
#   EuclideanUnitBall (StaticUnitBall), VectorUnitBall (DynamicUnitBall).

"""
Supertype of balls for which elements satisfy `norm(x) < radius(ball)` (open ball)
or `norm(x) <= radius(ball)` (closed ball).
"""
abstract type Ball{T,C} <: Domain{T} end

"A ball in a fixed N-dimensional Euclidean space."
const EuclideanBall{N,T,C} = Ball{SVector{N,T},C}

"A ball with vector elements of variable length."
const VectorBall{T,C} = Ball{Vector{T},C}

iscompatiblepair(x::AbstractVector, d::VectorBall) = length(x) == dimension(d)

const ClosedBall{T} = Ball{T,:closed}
const OpenBall{T} = Ball{T,:open}

isclosedset(::ClosedBall) = true
isclosedset(::OpenBall) = false
isopenset(d::Ball) = !isclosedset(d)

# Convenience constructors for the abstract type
Ball() = UnitBall()
Ball{T}() where {T} = UnitBall{T}()
Ball{T,C}() where {T,C} = UnitBall{T,C}()
Ball(radius::Number) = GenericBall(radius)
Ball{T}(radius::Number) where {T} = GenericBall{T}(radius)
Ball{T,C}(radius::Number) where {T,C} = GenericBall{T,C}(radius)
Ball(radius::Number, center) = GenericBall(radius, center)
Ball{T}(radius::Number, center) where {T} = GenericBall{T}(radius, center)
Ball{T,C}(radius::Number, center) where {T,C} = GenericBall{T,C}(radius, center)

indomain(x, d::OpenBall) = norm(x-center(d)) < radius(d)
indomain(x, d::ClosedBall) = norm(x-center(d)) <= radius(d)
approx_indomain(x, d::OpenBall, tolerance) = norm(x-center(d)) < radius(d)+tolerance
approx_indomain(x, d::ClosedBall, tolerance) = norm(x-center(d)) <= radius(d)+tolerance

# A closed ball always contains at least its center point.
isempty(d::ClosedBall) = false
isempty(d::OpenBall) = radius(d) == 0

==(d1::Ball, d2::Ball) = isclosedset(d1)==isclosedset(d2) &&
    radius(d1)==radius(d2) && center(d1)==center(d2)

function issubset(d1::Ball, d2::Ball)
    if dimension(d1) == dimension(d2)
        if center(d1) == center(d2)
            if radius(d1) < radius(d2)
                true
            elseif radius(d1) == radius(d2)
                isclosedset(d2) || isopenset(d1)
            else
                false
            end
        else
            dist = norm(center(d1)-center(d2))
            if dist + radius(d1) < radius(d2)
                true
            else
                false
            end
        end
    else
        false
    end
end

# We choose the center of the ball here. Concrete types should implement 'center'
point_in_domain(d::Ball) = center(d)

"The unit ball."
abstract type UnitBall{T,C} <: Ball{T,C} end

radius(::UnitBall) = 1
center(d::UnitBall{T}) where {T<:StaticTypes} = zero(T)

UnitBall(n::Int) = DynamicUnitBall(n)
UnitBall(::Val{N} = Val(3)) where {N} = EuclideanUnitBall{N}()

UnitBall{T}(n::Int) where {T <: StaticTypes} = StaticUnitBall{T}(n)
UnitBall{T}(::Val{N}) where {N,T} = StaticUnitBall{T}(Val(N))
UnitBall{T}() where {T <: StaticTypes} = StaticUnitBall{T}()
UnitBall{T}(n::Int) where {T} = DynamicUnitBall{T}(n)

UnitBall{T,C}(n::Int) where {T <: StaticTypes,C} = StaticUnitBall{T,C}(n)
UnitBall{T,C}(::Val{N}) where {N,T,C} = StaticUnitBall{T,C}(Val(N))
UnitBall{T,C}() where {T <: StaticTypes,C} = StaticUnitBall{T,C}()
UnitBall{T,C}(n::Int) where {T,C} = DynamicUnitBall{T,C}(n)

const ClosedUnitBall{T} = UnitBall{T,:closed}
const OpenUnitBall{T} = UnitBall{T,:open}

indomain(x, d::OpenUnitBall) = norm(x) < 1
indomain(x, d::ClosedUnitBall) = norm(x) <= 1
approx_indomain(x, d::OpenUnitBall, tolerance) = norm(x) < 1+tolerance
approx_indomain(x, d::ClosedUnitBall, tolerance) = norm(x) <= 1+tolerance

==(d1::UnitBall, d2::UnitBall) = isclosedset(d1)==isclosedset(d2) &&
    dimension(d1) == dimension(d2)

issubset(d1::UnitBall, d2::UnitBall) =
    dimension(d1) == dimension(d2) && (isclosedset(d2) || isopenset(d1))

convert(::Type{SublevelSet}, d::UnitBall{T,C}) where {T,C} =
    SublevelSet{T,C}(norm, radius(d))
convert(::Type{SublevelSet{T}}, d::UnitBall{S,C}) where {T,S,C} =
    SublevelSet{T,C}(norm, radius(d))

"The unit ball with fixed dimension(s) specified by the element type."
struct StaticUnitBall{T,C} <: UnitBall{T,C}
end

StaticUnitBall() = StaticUnitBall{Float64}()
StaticUnitBall(::Val{N}) where {N} = StaticUnitBall{SVector{N,Float64}}()

StaticUnitBall{T}() where {T} = StaticUnitBall{T,:closed}()

StaticUnitBall{T}(n::Int) where {T} =
    (@assert n == euclideandimension(T); StaticUnitBall{T}())
StaticUnitBall{T}(::Val{N}) where {N,T} =
    (@assert N == euclideandimension(T); StaticUnitBall{T}())

StaticUnitBall{T,C}(n::Int) where {T,C} =
    (@assert n == euclideandimension(T); StaticUnitBall{T,C}())
StaticUnitBall{T,C}(::Val{N}) where {N,T,C} =
    (@assert N == euclideandimension(T); StaticUnitBall{T,C}())

"The unit ball in a fixed N-dimensional space."
const EuclideanUnitBall{N,T,C} = StaticUnitBall{SVector{N,T},C}

EuclideanUnitBall{N}() where {N} = EuclideanUnitBall{N,Float64}()

similardomain(d::StaticUnitBall{S,C}, ::Type{T}) where {S,C,T<:StaticTypes} =
    StaticUnitBall{T,C}()
similardomain(d::StaticUnitBall{S,C}, ::Type{T}) where {S,C,T} =
    DynamicUnitBall{T,C}(dimension(d))

const UnitDisk{T} = UnitBall{SVector{2,T},:closed}
UnitDisk() = UnitDisk{Float64}()

## A StaticUnitBall{<:Number} equals the interval [-1,1]  (open or closed)
convert(::Type{Interval}, d::StaticUnitBall{T,:closed}) where {T <: Number} =
    ChebyshevInterval{T}()
convert(::Type{Interval}, d::StaticUnitBall{T,:open}) where {T <: Number} =
    OpenInterval{T}(-1, 1)

canonicaldomain(::Equal, d::StaticUnitBall{T}) where {T<:Number} = convert(Interval, d)

"The unit ball with variable dimension stored in a data field."
struct DynamicUnitBall{T,C} <: UnitBall{T,C}
    dimension   ::  Int

    DynamicUnitBall{T,C}(n::Int) where {T,C} = new(n)
    DynamicUnitBall{T,C}(n::Int) where {T<:StaticTypes,C} =
        (@assert n == euclideandimension(T); new(n))
end

DynamicUnitBall(n::Int) = DynamicUnitBall{Vector{Float64}}(n)
DynamicUnitBall{T}(n::Int) where {T} = DynamicUnitBall{T,:closed}(n)

dimension(d::DynamicUnitBall) = d.dimension

center(d::DynamicUnitBall{T}) where {T} = zeros(eltype(T), dimension(d))

"The unit ball with vector elements of a given dimension."
const VectorUnitBall{T,C} = DynamicUnitBall{Vector{T},C}

VectorUnitBall(n::Int = 3) = VectorUnitBall{Float64}(n)
VectorUnitDisk() = VectorUnitBall(2)

similardomain(d::DynamicUnitBall{S,C}, ::Type{T}) where {S,C,T} =
    DynamicUnitBall{T,C}(dimension(d))
similardomain(d::DynamicUnitBall{S,C}, ::Type{T}) where {S,C,T<:StaticTypes} =
    StaticUnitBall{T,C}()

"A `GenericBall` is a ball with a given radius and center."
struct GenericBall{T,C,S} <: Ball{T,C}
    radius  ::  S
    center  ::  T
end

GenericBall() = GenericBall(1.0)
GenericBall(radius::S) where {S<:Number} = GenericBall(radius, zero(SVector{3,float(S)}))
GenericBall(radius::S, center::T) where {S<:Number,U,T<:Vector{U}} =
    GenericBall{Vector{promote_type(S,U)}}(radius, center)
GenericBall(radius::S, center::T) where {S<:Number,N,U,T<:SVector{N,U}} =
    GenericBall{SVector{N,promote_type(S,U)}}(radius, center)
GenericBall(radius::S, center::T) where {S<:Number,T<:AbstractVector} =
    GenericBall{Vector{promote_type(S,eltype(T))}}(radius, center)
GenericBall(radius::S, center::T) where {S<:Number,T<:Number} =
    GenericBall{promote_type(S,T)}(radius, center)
GenericBall{T}(radius) where {T} = GenericBall{T,:closed}(radius)
GenericBall{T}(radius, center) where {T} = GenericBall{T,:closed}(radius, center)
GenericBall{T,C}(radius) where {T,C} = GenericBall{T,C}(radius, zero(T))
GenericBall{T,C}(radius, center) where {T,C} = GenericBall{T,C,eltype(T)}(radius, center)

radius(d::GenericBall) = d.radius
center(d::GenericBall) = d.center

dimension(d::GenericBall) = length(d.center)

canonicaldomain(d::GenericBall{T,C}) where {T,C} = UnitBall{T,C}(dimension(d))
mapfrom_canonical(d::GenericBall) = AffineMap(radius(d), center(d))

boundingbox(d::GenericBall) =
    map_boundingbox(boundingbox(canonicaldomain(d)), mapfrom_canonical(d))

# Preserve the `Ball` type under affine maps which preserve shape
map_domain(m::GenericAffineMap{T,S}, d::Ball{U,C}) where {T<:AbstractVector,S<:Number,U<:AbstractVector,C} =
    Ball{U,C}(m.A * radius(d), m.A * center(d) + m.b)
map_domain(m::AffineMap{T}, d::Ball{U,C}) where {T<:Number,U<:Number,C} =
    Ball{promote_type(U,T),C}(m.A * radius(d), m.A * center(d) + m.b)
map_domain(m::GenericLinearMap{T,S}, d::Ball{U,C}) where {S<:Number,T<:AbstractVector,U<:AbstractVector,C} =
    Ball{promote_type(T,U),C}(m.A * radius(d), m.A * center(d))
map_domain(m::LinearMap{T}, d::Ball{U,C}) where {T<:Number,U<:Number,C} =
    Ball{promote_type(T,U),C}(m.A * radius(d), m.A * center(d))
map_domain(m::LinearMap{T}, d::Ball{U,C}) where {T<:Number,U,C} =
    Ball{U,C}(m.A * radius(d), m.A * center(d))
map_domain(m::Translation{T}, d::Ball{S,C}) where {T<:AbstractVector,S<:AbstractVector,C} =
    Ball{S,C}(radius(d), center(d) + m.b)
map_domain(m::Translation{T}, d::Ball{S,C}) where {T<:Number,S<:Number,C} =
    Ball{S,C}(radius(d), center(d) + m.b)

show(io::IO, d::EuclideanUnitBall{3,Float64,:closed}) = print(io, "UnitBall()")
show(io::IO, d::EuclideanUnitBall{N,Float64,:closed}) where {N} = print(io, "UnitBall(Val($(N)))")
show(io::IO, d::EuclideanUnitBall{3,Float64,:open}) = print(io, "UnitBall()  (open)")
show(io::IO, d::EuclideanUnitBall{N,Float64,:open}) where {N} = print(io, "UnitBall(Val($(N)))  (open)")
Display.object_parentheses(d::EuclideanUnitBall{N,Float64,:open}) where {N} = true
show(io::IO, d::EuclideanUnitBall{2,Float64,:closed}) = print(io, "UnitDisk()")
show(io::IO, d::UnitDisk{T}) where {T} = print(io, "UnitDisk{$(T)}()")
show(io::IO, d::UnitBall{T}) where {T<:Number} = print(io, "UnitBall{$(T)}()")
show(io::IO, d::UnitBall{T,:open}) where {T<:Number} = print(io, "UnitBall{$(T)}()  (open)")
Display.object_parentheses(d::UnitBall{T,:open}) where {T<:Number} = true
show(io::IO, d::VectorUnitBall{Float64,:closed}) = print(io, "UnitBall($(dimension(d)))")
show(io::IO, d::VectorUnitBall{Float64,:open}) = print(io, "UnitBall($(dimension(d)))  (open)")
show(io::IO, d::Ball) = print(io, "Ball($(radius(d)), $(center(d)))")

# The type hierarchy of spheres parallels that of Ball above:
# abstract Sphere
# |-> abstract UnitSphere: radius is 1
#     |-> StaticUnitSphere: dimension is part of type
#     |-> DynamicUnitSphere: dimension is specified by int field
# |-> GenericSphere: stores center and radius. Here, the dimension can be
#     obtained from the center, so there is only one type.
# There are aliases for SVector{N,T} and Vector{T}.

"Supertype of spherical domains for which elements satisfy `norm(x) == radius(sphere)`."
abstract type Sphere{T} <: Domain{T} end

indomain(x, d::Sphere) = norm(x-center(d)) == radius(d)
approx_indomain(x, d::Sphere, tolerance) =
    radius(d)-tolerance <= norm(x-center(d)) <= radius(d)+tolerance

# Convenience constructors for the abstract type
Sphere() = UnitSphere()
Sphere{T}() where {T} = UnitSphere{T}()
Sphere(radius::Number) = GenericSphere(radius)
Sphere{T}(radius::Number) where {T} = GenericSphere{T}(radius)
Sphere(radius::Number, center) = GenericSphere(radius, center)
Sphere{T}(radius::Number, center) where {T} = GenericSphere{T}(radius, center)

isempty(::Sphere) = false

isclosedset(::Sphere) = true
isopenset(::Sphere) = false

point_in_domain(d::Sphere) = center(d) + unitvector(d)

==(d1::Sphere, d2::Sphere) = radius(d1)==radius(d2) && center(d1)==center(d2)

"A hypersphere in a fixed N-dimensional Euclidean space."
const EuclideanSphere{N,T} = Sphere{SVector{N,T}}

"A sphere with vector elements of variable length."
const VectorSphere{T} = Sphere{Vector{T}}

"The unit sphere."
abstract type UnitSphere{T} <: Sphere{T} end

radius(::UnitSphere) = 1
center(d::UnitSphere{T}) where {T<:StaticTypes} = zero(T)

indomain(x, d::UnitSphere) = norm(x) == 1
approx_indomain(x, d::UnitSphere, tolerance) =
    1-tolerance <= norm(x) <= 1+tolerance

UnitSphere(n::Int) = DynamicUnitSphere(n)
UnitSphere(::Val{N} = Val(3)) where {N} = EuclideanUnitSphere{N}()

UnitSphere{T}(n::Int) where {T <: StaticTypes} = StaticUnitSphere{T}(n)
UnitSphere{T}(::Val{N}) where {N,T} = StaticUnitSphere{T}(Val(N))
UnitSphere{T}() where {T <: StaticTypes} = StaticUnitSphere{T}()
UnitSphere{T}(n::Int) where {T} = DynamicUnitSphere{T}(n)

==(d1::UnitSphere, d2::UnitSphere) = dimension(d1)==dimension(d2)

issubset1(d1::UnitSphere, d2::UnitBall) =
    dimension(d1) == dimension(d2) && isclosedset(d2)

convert(::Type{LevelSet}, d::UnitSphere{T}) where {T} = LevelSet{T}(norm, radius(d))
convert(::Type{LevelSet{T}}, d::UnitSphere) where {T} = LevelSet{T}(norm, radius(d))

"The unit sphere with fixed dimension(s) specified by the element type."
struct StaticUnitSphere{T} <: UnitSphere{T}
end

StaticUnitSphere() = StaticUnitSphere{Float64}()
StaticUnitSphere(::Val{N}) where {N} = StaticUnitSphere{SVector{N,Float64}}()

StaticUnitSphere{T}(n::Int) where {T} =
    (@assert n == euclideandimension(T); StaticUnitSphere{T}())
StaticUnitSphere{T}(::Val{N}) where {N,T} =
    (@assert N == euclideandimension(T); StaticUnitSphere{T}())

similardomain(d::StaticUnitSphere, ::Type{T}) where {T<:StaticTypes} =
    StaticUnitSphere{T}()
similardomain(d::StaticUnitSphere, ::Type{T}) where {T} =
    DynamicUnitSphere{T}(dimension(d))

"The unit sphere in a fixed N-dimensional Euclidean space."
const EuclideanUnitSphere{N,T} = StaticUnitSphere{SVector{N,T}}

EuclideanUnitSphere{N}() where {N} = EuclideanUnitSphere{N,Float64}()

"The unit circle in 2D."
const UnitCircle{T} = UnitSphere{SVector{2,T}}
UnitCircle() = UnitCircle{Float64}()

"The unit sphere with variable dimension."
struct DynamicUnitSphere{T} <: UnitSphere{T}
    dimension   ::  Int

    DynamicUnitSphere{T}(n::Int) where {T} = new(n)
    DynamicUnitSphere{T}(n::Int) where {T<:StaticTypes} =
        (@assert n == euclideandimension(T); new(n))
end

DynamicUnitSphere(n::Int) = DynamicUnitSphere{Vector{Float64}}(n)

dimension(d::DynamicUnitSphere) = d.dimension

center(d::DynamicUnitSphere{T}) where {T} = zeros(eltype(T), dimension(d))

"The unit sphere with vector elements of a given dimension."
const VectorUnitSphere{T} = DynamicUnitSphere{Vector{T}}

VectorUnitSphere(n::Int = 3) = VectorUnitSphere{Float64}(n)
VectorUnitCircle() = VectorUnitSphere(2)

similardomain(d::DynamicUnitSphere, ::Type{T}) where {T} =
    DynamicUnitSphere{T}(d.dimension)
similardomain(d::DynamicUnitSphere, ::Type{T}) where {T <: StaticTypes} =
    StaticUnitSphere{T}()

"A `GenericSphere` is a sphere with a given radius and center."
struct GenericSphere{T,S} <: Sphere{T}
    radius  ::  S
    center  ::  T
end

GenericSphere() = GenericSphere(1.0)
GenericSphere(radius::S) where {S<:Number} =
    GenericSphere(radius, zero(SVector{3,float(S)}))
GenericSphere(radius::S, center::T) where {S<:Number,U,T<:Vector{U}} =
    GenericSphere{Vector{promote_type(S,U)}}(radius, center)
GenericSphere(radius::S, center::T) where {S<:Number,N,U,T<:SVector{N,U}} =
    GenericSphere{SVector{N,promote_type(S,U)}}(radius, center)
GenericSphere(radius::S, center::T) where {S<:Number,T<:AbstractVector} =
    GenericSphere{Vector{promote_type(S,eltype(T))}}(radius, center)
GenericSphere{T}(radius) where {T} = GenericSphere{T}(radius, zero(T))
GenericSphere{T}(radius, center) where {T} =
    GenericSphere{T,eltype(T)}(radius, center)

radius(d::GenericSphere) = d.radius
center(d::GenericSphere) = d.center

dimension(d::GenericSphere) = length(d.center)

canonicaldomain(d::GenericSphere{T}) where {T} = UnitSphere{T}(dimension(d))
mapfrom_canonical(d::GenericSphere) = AffineMap(radius(d), center(d))

boundingbox(d::GenericSphere) =
    map_boundingbox(boundingbox(canonicaldomain(d)), mapfrom_canonical(d))

# Preserve the `Sphere` type under affine maps which preserve shape
map_domain(m::GenericAffineMap{T,S}, d::Sphere{U}) where {T<:AbstractVector,S<:Number,U<:AbstractVector} =
    Sphere{U}(m.A * radius(d), m.A * center(d) + m.b)
map_domain(m::AffineMap{T}, d::Sphere{U}) where {T<:Number,U<:Number} =
    Sphere{promote_type(U,T)}(m.A * radius(d), m.A * center(d) + m.b)
map_domain(m::GenericLinearMap{T,S}, d::Sphere{U}) where {S<:Number,T<:AbstractVector,U<:AbstractVector} =
    Sphere{promote_type(T,U)}(m.A * radius(d), m.A * center(d))
map_domain(m::LinearMap{T}, d::Sphere{U}) where {T<:Number,U<:Number} =
    Sphere{promote_type(T,U)}(m.A * radius(d), m.A * center(d))
map_domain(m::LinearMap{T}, d::Sphere{U}) where {T<:Number,U} =
    Sphere{U}(m.A * radius(d), m.A * center(d))
map_domain(m::Translation{T}, d::Sphere{S}) where {T<:AbstractVector,S<:AbstractVector} =
    Sphere{S}(radius(d), center(d) + m.b)
map_domain(m::Translation{T}, d::Sphere{S}) where {T<:Number,S<:Number} =
    Sphere{S}(radius(d), center(d) + m.b)

show(io::IO, d::EuclideanUnitSphere{3,Float64}) = print(io, "UnitSphere()")
show(io::IO, d::EuclideanUnitSphere{N,Float64}) where {N} = print(io, "UnitSphere(Val($(N)))")
show(io::IO, d::EuclideanUnitSphere{2,Float64}) = print(io, "UnitCircle()")
show(io::IO, d::UnitCircle{T}) where {T} = print(io, "UnitCircle{$(T)}()")
show(io::IO, d::VectorUnitSphere{Float64}) = print(io, "UnitSphere($(dimension(d)))")
show(io::IO, d::UnitSphere{T}) where {T<:Number} = print(io, "UnitSphere{$(T)}()")
show(io::IO, d::Sphere) = print(io, "Sphere($(radius(d)), $(center(d)))")

boundary(d::UnitBall{T}) where {T} = UnitSphere{T}(dimension(d))
boundary(d::Ball{T}) where {T} = Sphere{T}(radius(d), center(d))

interior(d::StaticUnitBall{T}) where {T} = StaticUnitBall{T,:open}()
interior(d::DynamicUnitBall{T}) where {T} = DynamicUnitBall{T,:open}(dimension(d))
interior(d::Ball{T}) where {T} = Ball{T,:open}(radius(d), center(d))

closure(d::StaticUnitBall{T}) where {T} = StaticUnitBall{T,:closed}()
closure(d::DynamicUnitBall{T}) where {T} = DynamicUnitBall{T,:closed}(dimension(d))
closure(d::Ball{T}) where {T} = Ball{T,:closed}(radius(d), center(d))

boundingbox(d::UnitBall{T}) where {T<:Number} = ChebyshevInterval{T}()
boundingbox(d::UnitBall{SVector{N,T}}) where {N,T} =
    ChebyshevProductDomain{N,T}()
boundingbox(d::UnitBall{T}) where {T} =
    Rectangle{T}(-ones(eltype(T), dimension(d)), ones(eltype(T), dimension(d)))

boundingbox(d::UnitSphere{T}) where {T<:Number} = ChebyshevInterval{T}()
boundingbox(d::UnitSphere{SVector{N,T}}) where {N,T} =
    ChebyshevProductDomain{N,T}()
boundingbox(d::UnitSphere{T}) where {T} =
    Rectangle{T}(-ones(eltype(T), dimension(d)), ones(eltype(T), dimension(d)))

################
# Applications
################

"Create a cylinder with given radius and length."
cylinder(::Type{T} = Float64) where {T} = UnitDisk{T}() × UnitInterval{T}()
cylinder(radius, length) = cylinder(promote(radius, length)...)
cylinder(radius::T, length::T) where {T} = (radius .* UnitDisk{T}()) × (0..length)

"Create an ellipse curve with semi-axes lengths `a` and `b` respectively."
ellipse(a::Number, b::Number) = ellipse(promote(a,b)...)
ellipse(a::T, b::T) where {T <: Number} = LinearMap(a, b).(UnitCircle{T}())

"Create an ellipse-shaped domain with semi-axes lengths `a` and `b` respectively."
ellipse_shape(a::Number, b::Number) = ellipse_shape(promote(a,b)...)
ellipse_shape(a::T, b::T) where {T <: Number} = LinearMap(a, b).(UnitDisk{T}())

"""
The map `[cos(2πt), sin(2πt)]` from `[0,1]` to the unit circle in `ℝ^2`.
"""
struct UnitCircleMap{T} <: Map{T} end

UnitCircleMap() = UnitCircleMap{Float64}()

mapsize(m::UnitCircleMap) = (2,)

applymap(m::UnitCircleMap{T}, t) where {T} = SVector(cos(2*T(pi)*t), sin(2*T(pi)*t))

function jacobian(m::UnitCircleMap{T}, t) where {T}
    a = 2*T(pi)
    SVector(-a*sin(a*t), a*cos(a*t))
end

# we know that the differential volume is 2*pi
diffvolume(m::UnitCircleMap{T}) where {T} = ConstantMap{T}(2*T(pi))

"""
`AngleMap` is a left inverse of `UnitCircleMap`. A 2D vector `x` is projected onto
the intersection point with the unit circle of the line connecting `x` to the
origin. The angle of this point, scaled to the interval `[0,1)`, is the result.
"""
struct AngleMap{T} <: Map{SVector{2,T}}
end

AngleMap() = AngleMap{Float64}()

function applymap(m::AngleMap{T}, x) where {T}
    twopi = 2*convert(T, pi)
    θ = atan(x[2],x[1])
    if θ < 0
        # atan2 returns an angle in (-π,π], convert to [0,2π) using periodicity.
        θ += twopi
    end
    # And divide by 2π to scale to [0,1)
    θ / twopi
end

mapsize(m::AngleMap) = (1,2)

function jacobian(m::AngleMap{T}, t) where {T}
    x = t[1]; y = t[2]
    twopi = 2*convert(T, pi)
    v = SVector(one(T) / (1+(y/x)^2) * (-y/x^2) * 1/twopi,
            one(T) / (1+(y/x)^2) * one(T)/x * 1/twopi)
    transpose(v)
end

leftinverse(m::UnitCircleMap{T}) where {T} = AngleMap{T}()
leftinverse(m::UnitCircleMap, x) = leftinverse(m)(x)
rightinverse(m::AngleMap{T}) where {T} = UnitCircleMap{T}()
rightinverse(m::AngleMap, x) = rightinverse(m)(x)

canonicaldomain(::Parameterization, d::UnitCircle{T}) where {T} = UnitInterval{T}()
mapfrom_canonical(::Parameterization, d::UnitCircle{T}) where {T} = UnitCircleMap{T}()

## The complex plane

const ComplexUnitCircle{T} = StaticUnitSphere{Complex{T}}
const ComplexUnitDisk{T,C} = StaticUnitBall{Complex{T},C}

ComplexUnitCircle() = ComplexUnitCircle{Float64}()
ComplexUnitDisk() = ComplexUnitDisk{Float64}()
ComplexUnitDisk{T}() where {T} = ComplexUnitDisk{T,:closed}()

show(io::IO, d::ComplexUnitCircle{Float64}) = print(io, "ComplexUnitCircle()")
show(io::IO, d::ComplexUnitDisk{Float64,:closed}) = print(io, "ComplexUnitDisk()")
show(io::IO, d::ComplexUnitDisk{Float64,:open}) = print(io, "ComplexUnitDisk()  (open)")
show(io::IO, d::ComplexUnitDisk{T,:closed}) where {T} = print(io, "ComplexUnitDisk{$(T)}()")
show(io::IO, d::ComplexUnitDisk{T,:open}) where {T} = print(io, "ComplexUnitDisk{$(T)}()  (open)")

canonicaldomain(::Parameterization, d::ComplexUnitCircle{T}) where {T} =
    UnitInterval{T}()
mapfrom_canonical(::Parameterization, d::ComplexUnitCircle{T}) where {T} =
    VectorToComplex{T}() ∘ UnitCircleMap{T}()


text:
# This file was generated, do not modify it. # hide
t_inference = print(round(@elapsed begin #hide
postsample = dynamicHMC(post; ℓ=ℓ)
end; sigdigits=4)) #hide

max_stars_repo_name: cscherrer/InformativePrior

max_issues_repo_name: cscherrer/InformativePrior

max_forks_repo_name: cscherrer/InformativePrior


text:
# Given a state and it's environments, we can act on it

"""
    One-site derivative
"""
function ac_prime(x::MPSTensor,pos::Int,mps::Union{FiniteMPS,InfiniteMPS,MPSComoving},cache)
    ham=cache.opp

    toret=zero(x)
    for (i,j) in opkeys(ham,pos)
        @tensor toret[-1,-2,-3]+=leftenv(cache,pos,mps)[i][-1,5,4]*x[4,2,1]*ham[pos,i,j][5,-2,3,2]*rightenv(cache,pos,mps)[j][1,3,-3]
    end
    for (i,j) in scalkeys(ham,pos)
        scal = ham.Os[pos,i,j];
        @tensor toret[-1,-2,-3]+=leftenv(cache,pos,mps)[i][-1,5,4]*(scal*x)[4,-2,1]*rightenv(cache,pos,mps)[j][1,5,-3]
    end

    return toret
end

function ac_prime(x::GenericMPSTensor{S,3},pos::Int,mpo,cache) where S
    ham=cache.opp

    toret=zero(x)
    for (i,j) in keys(ham,pos)
        opp = ham[pos,i,j]

        if isbelow(ham,i)
            @tensor toret[-1,-2,-3,-4] +=   leftenv(cache,pos,mpo)[i][-1,8,7]*
                                            x[7,2,-3,1]*
                                            opp[8,-2,3,2]*
                                            rightenv(cache,pos,mpo)[j][1,3,-4]
        else
            @tensor toret[-1,-2,-3,-4] +=   leftenv(cache,pos,mpo)[i][-1,6,7]*
                                            x[7,-2,4,2]*
                                            opp[6,4,5,-3]*
                                            rightenv(cache,pos,mpo)[j][2,5,-4]
        end
    end

    return toret
end
function ac_prime(x::MPSTensor, row::Int,col::Int,mps::Union{InfiniteMPS,MPSMultiline}, pars::PerMPOInfEnv)
    @tensor toret[-1 -2;-3]:=leftenv(pars,row,col,mps)[-1,2,1]*x[1,3,4]*(pars.opp[row,col])[2,-2,5,3]*rightenv(pars,row,col,mps)[4,5,-3]
end

"""
    Two-site derivative
"""
function ac2_prime(x::MPOTensor,pos::Int,mps::Union{FiniteMPS,InfiniteMPS,MPSComoving},cache)
    ham=cache.opp

    toret=zero(x)

    for (i,j) in keys(ham,pos)
        for k in 1:ham.odim
            if contains(ham,pos+1,j,k)
                #can be sped up for scalar fields
                @tensor toret[-1,-2,-3,-4]+=leftenv(cache,pos,mps)[i][-1,7,6]*x[6,5,3,1]*ham[pos,i,j][7,-2,4,5]*ham[pos+1,j,k][4,-3,2,3]*rightenv(cache,pos+1,mps)[k][1,2,-4]
            end
        end

    end

    return toret
end
function ac2_prime(x::AbstractTensorMap,pos::Int,mpo,cache::FinEnv{<:ComAct})
    ham=cache.opp

    toret=zero(x)
    for (i,j) in keys(ham,pos)
        for (k,l) in keys(ham,pos+1)
            if j!=k
                continue
            end
            opp1 = ham[pos,i,j]
            opp2 = ham[pos+1,k,l]

            if isbelow(ham,i)
                @tensor toret[-1,-2,-3,-4,-5,-6] += leftenv(cache,pos,mpo)[i][-1,2,1]*
                                                    x[1,3,-3,5,-5,7]*
                                                    opp1[2,-2,4,3]*
                                                    opp2[4,-4,6,5]*
                                                    rightenv(cache,pos+1,mpo)[l][7,6,-6]
            else
                @tensor toret[-1,-2,-3,-4,-5,-6] += leftenv(cache,pos,mpo)[i][-1,2,1]*
                                                    x[1,-2,3,-4,5,7]*
                                                    opp1[2,3,4,-3]*
                                                    opp2[4,5,6,-5]*
                                                    rightenv(cache,pos+1,mpo)[l][7,6,-6]
            end
        end
    end

    return toret
end
function ac2_prime(x::MPOTensor, row::Int,col::Int,mps::Union{InfiniteMPS,MPSMultiline}, pars::PerMPOInfEnv)
    @tensor toret[-1 -2;-3 -4]:=leftenv(pars,row,col,mps)[-1,2,1]*
                                x[1,3,4,5]*
                                pars.opp[row,col][2,-2,6,3]*
                                pars.opp[row,col+1][6,-3,7,4]*
                                rightenv(pars,row,col+1,mps)[5,7,-4]
end

"""
    Zero-site derivative (the C matrix to the right of pos)
"""
function c_prime(x::MPSBondTensor,pos::Int,mps::Union{FiniteMPS,InfiniteMPS,MPSComoving},cache)
    toret=zero(x)
    ham=cache.opp

    for i in 1:ham.odim
        @tensor toret[-1,-2]+=leftenv(cache,pos+1,mps)[i][-1,2,1]*x[1,3]*rightenv(cache,pos,mps)[i][3,2,-2]
    end

    return toret
end
function c_prime(x::TensorMap, row::Int,col::Int, mps::Union{InfiniteMPS,MPSMultiline}, pars::PerMPOInfEnv)
    @tensor toret[-1;-2] := leftenv(pars,row,col+1,mps)[-1,3,1]*x[1,2]*rightenv(pars,row,col,mps)[2,3,-2]
end


text:
struct Equation{DT <: Number, FT <: Function, FPRE <: Union{Function,Nothing}, FPOST <: Union{Function,Nothing}}
    f::FT
    x₀::Array{DT,2}
    f_pre::FPRE
    f_post::FPOST

    function Equation(f::FT, x₀::AbstractArray{DT,2}; f_preproc::FPRE=nothing, f_postproc::FPOST=nothing) where {DT, FT, FPRE, FPOST}
        new{DT, FT, FPRE, FPOST}(f, Array(x₀), f_preproc, f_postproc)
    end
end

ndims(equ::Equation) = length(axes(equ.x₀,1))
nsamples(equ::Equation) = length(axes(equ.x₀,2))

preprocessing(::Equation{<:Number, <:Function, Nothing}, ::AbstractArray) = nothing
postprocessing(::Equation{<:Number, <:Function, <:Union{Function,Nothing}, Nothing}, ::AbstractArray) = nothing

preprocessing(equ::Equation{<:Number, <:Function, <:Function}, x₀::AbstractArray) = equ.f_pre(x₀)
postprocessing(equ::Equation{<:Number, <:Function, <:Union{Function,Nothing}, <:Function}, x₁::AbstractArray) = equ.f_post(x₁)

max_stars_repo_name: michakraus/JuliaAtHomeAssignments

max_issues_repo_name: michakraus/JuliaAtHomeAssignments

max_forks_repo_name: michakraus/JuliaAtHomeAssignments


text:
# add JSExpr  # in package mode
using JSExpr

# using the @js macro:
@js function cube(arg)
  return arg * arg * arg
end
# JSString("function cube(arg){return (arg*arg*arg)}")

# showing interpolation:
var1 = 108
callback = @js n -> n + $var1
# JSString("(function (n){return (n+108)})")

# using the js" " macro:
js"var 1 is: $var1"
# JSString("var 1 is: 108")

message = "hi"
fun1 = js"
       function () {
           alert($message) // you can interpolate Julia variables!
       }
       "
# JSString("function () {\n    alert(\"hi\") // you can interpolate Julia variables!\n}\n")


text:
using Distributions
using Base.Test

@time @testset "linear reg chow test" begin
    p = 7
    n = 500
    n_env = 3
    env = repeat(1:n_env, inner=n)
    X = randn(n * n_env, p)
    beta = randn(p)
    y = X * beta - 1. + randn(n * n_env)
    rej, p_value, conf_intervals = InvariantCausal.conditional_inv_test_chow(X, y, env, n_env)
    @test rej == false
    y[1:n] = X[1:n, :] * (beta / 2) - 1. + randn(n)  # change environment 1
    rej, p_value, conf_intervals = InvariantCausal.conditional_inv_test_chow(X, y, env, n_env)
    @test rej == true
end

@time @testset "chow test" begin
    m = 100
    p_values = map(x -> (
        n = 300;
        p = 7;
        X1 = randn(2 * n, p);
        X2 = randn(n, p);
        beta = randn(p);
        y1 = X1 * beta + 1 + randn(2 * n);
        y2 = X2 * beta + 1 + randn(n);
        two_sample_chow(X1, X2, y1, y2)), 1:m)
    @test abs(mean(p_values .< 0.05) - 0.05) < 2 * sqrt(0.05 * 0.95 / m)
end


text:
const ChecksumAlgorithm = String

@dict_readable struct ExceptionBreakpointsFilter <: Outbound
    filter::String
    label::String
    default::Union{Missing,Bool}
end

@dict_readable struct ColumnDescriptor <: Outbound
    attributeName::String
    label::String
    format::Union{Missing,String}
    type::Union{Missing,String} # default: "string"
    width::Union{Missing,Int}
end

@dict_readable struct Capabilities <: Outbound
    supportsConfigurationDoneRequest::Union{Missing,Bool}
    supportsFunctionBreakpoints::Union{Missing,Bool}
    supportsConditionalBreakpoints::Union{Missing,Bool}
    supportsHitConditionalBreakpoints::Union{Missing,Bool}
    supportsEvaluateForHovers::Union{Missing,Bool}
    exceptionBreakpointFilters::Vector{ExceptionBreakpointsFilter}
    supportsStepBack::Union{Missing,Bool}
    supportsSetVariable::Union{Missing,Bool}
    supportsRestartFrame::Union{Missing,Bool}
    supportsGotoTargetsRequest::Union{Missing,Bool}
    supportsStepInTargetsRequest::Union{Missing,Bool}
    supportsCompletionsRequest::Union{Missing,Bool}
    supportsModulesRequest::Union{Missing,Bool}
    additionalModuleColumns::Vector{ColumnDescriptor}
    supportedChecksumAlgorithms::Vector{ChecksumAlgorithm}
    supportsRestartRequest::Union{Missing,Bool}
    supportsExceptionOptions::Union{Missing,Bool}
    supportsValueFormattingOptions::Union{Missing,Bool}
    supportsExceptionInfoRequest::Union{Missing,Bool}
    supportTerminateDebuggee::Union{Missing,Bool}
    supportsDelayedStackTraceLoading::Union{Missing,Bool}
    supportsLoadedSourcesRequest::Union{Missing,Bool}
    supportsLogPoints::Union{Missing,Bool}
    supportsTerminateThreadsRequest::Union{Missing,Bool}
    supportsSetExpression::Union{Missing,Bool}
    supportsTerminateRequest::Union{Missing,Bool}
    supportsDataBreakpoints::Union{Missing,Bool}
    supportsReadMemoryRequest::Union{Missing,Bool}
    supportsDisassembleRequest::Union{Missing,Bool}
end

@dict_readable struct Message <: Outbound
    id::Int
    format::String
    variables::Union{Missing,Dict{String,String}}
    sendTelemetry::Union{Missing,Bool}
    showUser::Union{Missing,Bool}
    url::Union{Missing,String}
    urlLabel::Union{Missing,String}
end

@dict_readable struct DAModule <: Outbound
    id::Union{Int,String}
    name::String
    path::Union{Missing,String}
    isOptimized::Union{Missing,Bool}
    isUserCode::Union{Missing,Bool}
    version::Union{Missing,String}
    symbolStatus::Union{Missing,String}
    dateTimeStamp::Union{Missing,String}
    addressRange::Union{Missing,String}
end

@dict_readable struct ModulesViewDescriptor <: Outbound
    columns::Vector{ColumnDescriptor}
end

@dict_readable struct Thread <: Outbound
    id::Int
    name::String
end

@dict_readable struct Checksum <: Outbound
    algorithm::ChecksumAlgorithm
    checksum::String
end

@dict_readable struct Source <: Outbound
    name::Union{Missing,String}
    path::Union{Missing,String}
    sourceReference::Union{Missing,Int}
    presentationHint::Union{Missing,String}
    origin::Union{Missing,String}
    sources::Union{Missing,Vector{Source}}
    adapterData::Union{Missing,Any}
    checksums::Union{Missing,Vector{Checksum}}
end

@dict_readable struct StackFrame <: Outbound
    id::Int
    name::String
    source::Union{Missing,Source}
    line::Int
    column::Int
    endLine::Union{Missing,Int}
    endColum::Union{Missing,Int}
    instructionPointerReference::Union{Missing,String}
    moduleId::Union{Missing,Int,String}
    presentationHint::Union{Missing,String}
end

@dict_readable struct Scope <: Outbound
    name::String
    presentationHint::Union{Missing,String}
    variablesReference::Int
    namedVariables::Union{Missing,Int}
    indexedVariables::Union{Missing,Int}
    expensive::Bool
    source::Union{Missing,Source}
    line::Union{Missing,Int}
    column::Union{Missing,Int}
    endLine::Union{Missing,Int}
    endColum::Union{Missing,Int}
end

@dict_readable struct VariablePresentationHint <: Outbound
    kind::Union{Missing,String}
    attributes::Union{Missing,Vector{String}}
    visibility::Union{Missing,String}
end

@dict_readable struct Variable <: Outbound
    name::String
    value::String
    type::Union{Missing,String}
    presentatonHint::Union{Missing,VariablePresentationHint}
    evaluateName::Union{Missing,String}
    variablesReference::Int
    namedVariables::Union{Missing,Int}
    indexedVariables::Union{Missing,Int}
    memoryReference::Union{Missing,String}
end

@dict_readable struct SourceBreakpoint <: Outbound
    line::Int
    column::Union{Missing,Int}
    condition::Union{Missing,String}
    hitCondition::Union{Missing,String}
    logMessage::Union{Missing,String}
end

@dict_readable struct FunctionBreakpoint <: Outbound
    name::String
    condition::Union{Missing,String}
    hitCondition::Union{Missing,String}
end

const DataBreakpointAccessType = String

@dict_readable struct DataBreakpoint <: Outbound
    dataId::String
    accessType::Union{Missing,DataBreakpointAccessType}
    condition::Union{Missing,String}
    hitCondition::Union{Missing,String}
end

@dict_readable struct Breakpoint <: Outbound
    id::Union{Missing,Int}
    verified::Bool
    message::Union{Missing,String}
    source::Union{Missing,Source}
    line::Union{Missing,Int}
    column::Union{Missing,Int}
    endLine::Union{Missing,Int}
    endColumn::Union{Missing,Int}
end

function Breakpoint(verified::Bool)
    return Breakpoint(
        missing,
        verified,
        missing,
        missing,
        missing,
        missing,
        missing,
        missing
    )
end

@dict_readable struct StepInTarget <: Outbound
    id::Int
    label::String
end

@dict_readable struct GotoTarget <: Outbound
    id::Int
    label::String
    line::Int
    column::Union{Missing,Int}
    endLine::Union{Missing,Int}
    endColumn::Union{Missing,Int}
    instructionPointerReference::Union{Missing,String}
end

const CompletionItemType = String

@dict_readable struct CompletionItem <: Outbound
    label::String
    text::Union{Missing,String}
    type::CompletionItemType
    start::Union{Missing,Int}
    length::Union{Missing,Int}
end

@dict_readable struct ValueFormat <: Outbound
    hex::Union{Missing,Bool}
end

@dict_readable struct StackFrameFormat <: Outbound
    parameters::Union{Missing,Bool}
    parameterTypes::Union{Missing,Bool}
    parameterNames::Union{Missing,Bool}
    parameterValues::Union{Missing,Bool}
    line::Union{Missing,Bool}
    # module::Union{Missing,Bool}
    includeAll::Union{Missing,Bool}
end

const ExceptionBreakMode = String
ExceptionBreakModes = ["never", "always", "unhandled", "userUnhandled"]

@dict_readable struct ExceptionPathSegment <: Outbound
    negate::Union{Missing,Bool}
    names::Vector{String}
end

@dict_readable struct ExceptionOptions <: Outbound
    path::Union{Missing,Vector{ExceptionPathSegment}}
    breakMode::ExceptionBreakMode
end

@dict_readable struct ExceptionDetails <: Outbound
    message::Union{Missing,String}
    typeName::Union{Missing,String}
    fullTypeName::Union{Missing,String}
    evaluateName::Union{Missing,String}
    stackTrace::Union{Missing,String}
    innerException::Union{Missing,Vector{ExceptionDetails}}
end

@dict_readable struct DisassembledInstruction <: Outbound
    address::String
    instructionBytes::Union{Missing,String}
    instruction::String
    symbol::Union{Missing,String}
    location::Union{Missing,Source}
    line::Union{Missing,Int}
    column::Union{Missing,Int}
    endLine::Union{Missing,Int}
    endColumn::Union{Missing,Int}
end

@dict_readable struct BreakpointLocation <: Outbound
    line::Int
    column::Union{Missing,Int}
    endLine::Union{Missing,Int}
    endColumn::Union{Missing,Int}
end


text:
using Flux, Tracker, TextAnalysis
using BSON: load
using Flux: onehot, onehotbatch, reset!

mutable struct BiLSTM_CNN_CRF_Model{C, W, L, D, O, A}
    labels::Array{String, 1} # List of Labels
    chars_idx::Dict{Char, Int64} # Dict that maps chars to indices in W_Char_Embed
    words_idx::Dict{String, Int64} # Dict that maps words to indices in W_word_Embed
    conv1::C # Convolution Layer over W_Char_Embed to give character representation
    W_Char_Embed::W # Weights for character embeddings
    W_word_Embed::W # Further trained GloVe Embeddings
    forward_lstm::L # Forward LSTM
    backward::L # Backward LSTM
    d_out::D # Dense_out
    c::O # CRF
    init_α::A
    UNK_Word_idx::Integer
    UNK_char_idx::Integer
    # device::Symbol # Cpu or Gpu
end

# BiLSTM_CNN_CRF_Model(labels, chars_idx, words_idx) =
            # BiLSTM_CNN_CRF_Model(labels, chars_idx, words_idx, :cpu)

function BiLSTM_CNN_CRF_Model(labels, chars_idx, words_idx, UNK_char_idx, UNK_Word_idx; CHAR_EMBED_DIMS=25, WORD_EMBED_DIMS=100,
                              CNN_OUTPUT_SIZE=30, CONV_PAD= (0,2), CONV_WINDOW_LENGTH = 3, LSTM_STATE_SIZE = 200)
    n = length(labels)
    init_α = fill(-10000, (n + 2, 1))
    init_α[n + 1] = 0

    BiLSTM_CNN_CRF_Model(labels, chars_idx, words_idx, Conv((CHAR_EMBED_DIMS, CONV_WINDOW_LENGTH), 1=>CNN_OUTPUT_SIZE, pad=CONV_PAD),
                rand(CHAR_EMBED_DIMS, length(chars_idx)), rand(WORD_EMBED_DIMS, length(words_idx)),
                LSTM(CNN_OUTPUT_SIZE + WORD_EMBED_DIMS, LSTM_STATE_SIZE), LSTM(CNN_OUTPUT_SIZE + WORD_EMBED_DIMS, LSTM_STATE_SIZE),
                Dense(LSTM_STATE_SIZE * 2, length(labels) + 2), CRF(n), init_α, UNK_Word_idx, UNK_char_idx)
end

function BiLSTM_CNN_CRF_Model(labels, chars_idx, words_idx, UNK_char_idx,UNK_Word_idx, weights_path)
    n = length(labels)
    init_α = fill(-10000, (n + 2, 1))
    init_α[n + 1] = 0

    conv1 = load(joinpath(weights_path, "conv_cpu.bson"))[:conv_cpu]
    W_word_Embed = load(joinpath(weights_path, "W_word_cpu.bson"))[:W_word_cpu].data
    W_Char_Embed = load(joinpath(weights_path, "W_char_cpu.bson"))[:W_char_cpu].data
    forward_lstm = load(joinpath(weights_path, "forward_lstm.bson"))[:forward_lstm_cpu]
    backward = load(joinpath(weights_path, "backward_lstm.bson"))[:backward_lstm_cpu]
    d_out = load(joinpath(weights_path, "d_cpu.bson"))[:d_cpu]
    c = load(joinpath(weights_path, "crf.bson"))[:crf_cpu]

    BiLSTM_CNN_CRF_Model(labels, chars_idx, words_idx, conv1, W_Char_Embed, W_word_Embed,
                forward_lstm, backward, d_out, c, init_α, UNK_Word_idx, UNK_char_idx)
end

function (a::BiLSTM_CNN_CRF_Model)(x)
    char_features = Chain(x -> reshape(x, size(x)..., 1,1),
                          a.conv1,
                          x -> maximum(x, dims=2),
                          x -> reshape(x, length(x),1))
    input_embeddings((w, cs)) = vcat(a.W_word_Embed * w, char_features(a.W_Char_Embed * cs))
    backward_lstm(x) = reverse((a.backward).(reverse(x)))
    bilstm_layer(x) = vcat.((a.forward_lstm).(x), backward_lstm(x))
    m = Chain(x -> input_embeddings.(x),
              bilstm_layer,
              x -> (a.d_out).(x))

    oh_outs = viterbi_decode(a.c, m(x), a.init_α)
    reset!(a.backward)
    reset!(a.forward_lstm)
    [a.labels[oh.ix] for oh in oh_outs]
end

# function load(m::BiLSTM_CNN_CRF_Model, weights_path)
#     m.conv1 = BSON.load(joinpath(weights_path, "conv_cpu.bson"))[:conv_cpu]
#     println("ConvLoaded")
#     m.W_word_Embed = BSON.load(joinpath(weights_path, "W_word_cpu.bson"))[:W_word_cpu].data
#     println("W Word loaded")
#     m.W_Char_Embed = BSON.load(joinpath(weights_path, "W_char_cpu.bson"))[:W_char_cpu].data
#     println("W char loaded")
#     m.forward_lstm = BSON.load(joinpath(weights_path, "forward_lstm.bson"))[:forward_lstm_cpu]
#     m.backward = BSON.load(joinpath(weights_path, "backward_lstm.bson"))[:backward_lstm_cpu]
#     m.d_out = BSON.load(joinpath(weights_path, "d_cpu.bson"))[:d_cpu]
#     m.c = BSON.load(joinpath(weights_path, "crf.bson"))[:crf_cpu]
# end

onehotinput(m::BiLSTM_CNN_CRF_Model, word) = (onehot(get(m.words_idx, lowercase(word), m.UNK_Word_idx), 1:length(m.words_idx)),
                onehotbatch([get(m.chars_idx, c, m.UNK_char_idx) for c in word], 1:length(m.chars_idx)))


text:
export insolation, solar_flux_and_cos_sza

"""
    insolation(θ::FT, d::FT, param_set::APS) where {FT <: Real}

Returns the insolation given the zenith angle and earth-sun distance
param_set is an AbstractParameterSet from CLIMAParameters.jl.
"""
function insolation(θ::FT, d::FT, param_set::APS) where {FT <: Real}
    S0::FT = tot_solar_irrad(param_set)
    AU::FT = astro_unit()
    # set max. zenith angle to π/2, insolation should not be negative
    if θ > FT(π)/FT(2)
        θ = FT(π)/FT(2)
    end
    # weighted irradiance (3.12)
    S = S0 * (AU / d)^2
    # TOA insolation (3.15)
    F = S * cos(θ)
    return F
end

"""
    solar_flux_and_cos_sza(date::DateTime,
                      longitude::FT,
                      latitude::FT,
                      param_set::APS) where {FT <: Real}

Returns the top-of-atmosphere (TOA) solar flux, i.e.
the total solar irradiance (TSI) weighted by the earth-sun distance
and cos(solar zenith angle) for input to RRTMGP.jl
param_set is an AbstractParameterSet from CLIMAParameters.jl.
"""
function solar_flux_and_cos_sza(date::DateTime,
                           longitude::FT,
                           latitude::FT,
                           param_set::APS) where {FT <: Real}
    S0::FT = tot_solar_irrad(param_set)
    AU::FT = astro_unit()
    # θ = solar zenith angle, ζ = solar azimuth angle, d = earth-sun distance
    θ, ζ, d = instantaneous_zenith_angle(date, longitude, latitude, param_set)
    # set max. zenith angle to π/2, insolation should not be negative
    if θ > FT(π)/FT(2)
        θ = FT(π)/FT(2)
    end
    μ = cos(θ)
    # TOA solar flux (3.12)
    S = S0 * (AU / d)^2
    # return inputs needed for RRTMGP.jl
    return S, μ
end


text:
type NCE <: Model; rnn; tst; trn1; trn2; rng; psample; prob; alias; noise; kqvec;
    function NCE(rnn::Function; rng=Base.GLOBAL_RNG, o...)
        rnn = Net(rnn; o...)                            # x[X,B]->h[H,B]
        trn1 = Net(nce_trn; o..., ninputs=2)            # h[H,B]->p[B,B]
        trn2 = Net(nce_trn; o..., ninputs=2)            # h[H,B]->p[K,B]
        tst  = Net(nce_tst; o...)                       # h[H,B]->p[V,B]
        new(rnn, tst, trn1, trn2, rng)
    end
end

params(r::NCE)=vcat(params(r.rnn), params(r.tst))       # trn1,2 and tst share parameters
reset!(r::NCE)=map(reset!, (r.rnn, r.tst, r.trn1, r.trn2))

@knet function nce_trn(h,r; vocab=0, psample=0, nsample=0, kqvec=psample*nsample, nce_winit=Gaussian(0,.01), nce_binit=Constant(log(1/vocab)), o...)
    w  = par(; init=nce_winit, dims=(vocab,0))
    b  = par(; init=nce_binit, dims=(vocab,1))
    q  = arr(; init=kqvec)
    rw = dot(r,w)
    rb = dot(r,b)
    rq = dot(r,q)
    y  = dot(rw,h)
    s  = add(rb,y)
    p  = nce(rq,s)
end

@knet function nce_tst(h; o...)
    w  = par()
    b  = par()
    y1 = dot(w,h)
    y2 = add(b,y1)
    y3 = soft(y2)
end

# m.tst: input, par(w), par(b), dot, add, soft
# m.trn1,2: input(h), input(r), par(w), par(b), par(q), dot(rw), dot(rb), dot(rq), dot(y), add(s), nce(q)
function nce_init(m::NCE, data)
    m.tst.out0[1] == nothing || return
    (x,ygold) = item2xy(first(data))
    @assert x!=nothing && ygold!=nothing
    # This should allocate all params:
    h = forw(m.rnn,  x...; trn=true, seq=true)
    r = csc2csr(ygold)
    y = forw(m.trn1, h, r; trn=true, seq=true)
    kqvec = CudaArray(m.kqvec[ygold.rowval])
    g = back(m.trn1, kqvec, nce_loss_real; seq=true, getdx=1)
    back(m.rnn, g; seq=true)
    for i=3:4, n=(:op,:out,:out0,:dif,:dif0,:tmp)
        m.trn2.(n)[i]  = m.trn1.(n)[i]
        m.tst.(n)[i-1] = m.trn1.(n)[i]
    end
    m.trn1.params = filter(x->isa(x,Par), m.trn1.op)
    m.trn2.params = filter(x->isa(x,Par), m.trn2.op)
    m.tst.params = filter(x->isa(x,Par), m.tst.op)
end

function test(m::NCE, data, loss; o...)
    sumloss = numloss = 0
    reset!(m)
    for item in data
        if item != nothing
            (x,ygold) = item2xy(item)
            h = forw(m.rnn,  x...; trn=false, seq=true, o...)
            ypred = forw(m.tst, h; trn=false, seq=true, o...)
            sumloss += loss(ypred, ygold); numloss += 1
        else
            reset!(m)
        end
    end
    return sumloss/numloss
end

function train(m::NCE, data, loss=nothing; nsample=0, psample=nothing, gclip=0, gcheck=false, maxnorm=nothing, losscnt=nothing, o...) # t:7351
    loss==nothing || Base.warn_once("NCE uses its own loss function for training, the user provided one is ignored.")
    nsample==0 && (nsample=100; Base.warn_once("Using nsample=100"))
    psample==nothing && (v=m.trn1.op[3].dims[1];psample=fill(Float32(1/v),v); Base.warn_once("Using uniform noise dist"))
    nce_alias_init(m, psample, nsample)
    nce_init(m, data)
    reset!(m)
    for item in data
        if item != nothing
            (x,ygold) = item2xy(item)
            isa(ygold, SparseMatrixCSC) || error("NCE expects ygold to be a SparseMatrixCSC")
            h = forw(m.rnn, x...; trn=true, seq=true, o...)                                     # t:1668

            r = csc2csr(ygold)                                                                  # t:159  copies the transpose to gpu
            y = forw(m.trn1, h, r; trn=true, seq=true, o...)                                    # t:460
            losscnt != nothing && (losscnt[1] += nce_loss_real(y); losscnt[2] += 1)             # t:94

            nce_alias_sample(m)                                                                 # t:56
            r = csc2csr(m.noise)                                                                # t:109
            y = forw(m.trn2, h, r; trn=true, seq=true, o...)                                    # t:523
            losscnt != nothing && (losscnt[1] += nce_loss_noise(y); losscnt[2] += nsample)      # t:134

        else # end of sequence
            while m.rnn.sp > 0
                hgrad1 = back(m.trn2, Any[], nce_loss_noise; seq=true, getdx=1, o...)           # t:864
                hgrad = copy(hgrad1)                                                            # t:28
                hgrad2 = back(m.trn1, Any[], nce_loss_real;  seq=true, getdx=1, o...)           # t:660
                axpy!(1, hgrad2, hgrad)                                                         # t:9
                back(m.rnn, hgrad; seq=true, o...)                                              # t:2446
            end
            gcheck && break
            g = (gclip > 0 || maxnorm!=nothing ? gnorm(m) : 0)
            update!(m; gclip=(g > gclip > 0 ? gclip/g : 0), o...)
            if maxnorm != nothing
                w=wnorm(m)
                w > maxnorm[1] && (maxnorm[1]=w)
                g > maxnorm[2] && (maxnorm[2]=g)
            end
            reset!(m)
        end
    end
end

# See Vose's alias method for sampling from
# http://keithschwarz.com/darts-dice-coins
function nce_alias_init(m::NCE, psample, nsample)
    vocab = length(psample)
    isdefined(m, :noise) && size(m.noise) == (vocab, nsample) ||
    (m.noise = sponehot(eltype(psample), length(psample), nsample))
    isdefined(m, :psample) && isequal(m.psample, psample) && return
    m.psample = copy(psample)
    m.kqvec = nsample * psample
    m.prob = similar(psample)
    m.alias = similar(psample, Int)
    small = Int[]
    large = Int[]
    p = vocab * psample
    for i in 1:vocab; push!((p[i] < 1 ? small : large), i); end
    while !isempty(small) && !isempty(large)
        l = pop!(small)
        g = pop!(large)
        m.prob[l] = p[l]
        m.alias[l] = g
        p[g] = (p[g] + p[l]) - 1
        push!((p[g] < 1 ? small : large), g)
    end
    for i in large; m.prob[i]=1; end
    for i in small; m.prob[i]=1; end
end

function nce_alias_sample(m::NCE)
    (vocab, nsample) = size(m.noise)
    for n=1:nsample
        i = ceil(Int, rand(m.rng) * vocab)
        w = (rand(m.rng) <= m.prob[i] ? i : m.alias[i])
        m.noise.rowval[n] = w
    end
end

# TODO: will have to implement masks like softloss?

# J_real = -logp(d=1|y)
# ypred is Matrix(B,B), ypred[i,j] is the p(d=1) given the i'th gold word on the j'th sentence (only ypred[i,i] is used)
nce_loss_real(ypred, ygold=nothing; o...)=vecnorm(log!(diag(ypred)),1)/size2(ypred,2)

# J_noise = -logp(d=0|y) = -log(1-p(d=1|y))
# ypred is Matrix(K,B), ypred[i,j] is the p(d=1) given the i'th gold word on the j'th sentence
nce_loss_noise(ypred, ygold=nothing; o...)=vecnorm(log!(axpb!(ypred,similar(ypred);a=-1,b=1)),1)/size2(ypred,2)

# dJ_real/dp = -1/p
nce_loss_real(ypred, ygold, ygrad; o...)=diagm!(axpb!(diag(ypred); a=-1/size2(ypred,2),p=-1), ygrad)

# dJ_noise/dp = 1/(1-p)
nce_loss_noise(ypred, ygold, ygrad; o...)=axpb!(axpb!(ypred,ygrad; a=-1, b=1); p=-1, a=1/size2(ypred,2))

csc2csr{T}(x::SparseMatrixCSC{T})=CudaSparseMatrixCSR{T}(CudaArray(convert(Vector{Cint},x.colptr)), CudaArray(convert(Vector{Cint},x.rowval)), CudaArray(x.nzval), (x.n,x.m), convert(Cint,length(x.nzval)), device())

# Define main, test, noise, gold as four separate networks.
# If we leave w,b outside of networks they will get unnecessarily copied.
# More troubling, they will not get automatically updated.
# A lot of net forw and back code will be duplicated.
# If we construct a full nce_trn network how do we share its parameters with nce_tst?
# I think we'll just have to do that and figure nce_tst later.
# In nce_trn it would be nice to split the noise and the real batch steps.
# so we have three possible phases: test, noise, gold.  noise and gold have their own forw/back/loss.
# they share parameters and structure.
# we want to compute h only once and pass it to both noise and gold.
# so noise and gold should be partial nets with input h only (gold also gets ygold).
# how do we share their weights?  figure that out later.
# run test first so w,b initialized, then operate on the others to point to w,b.

# Call nce_trn with r = random one-hot-rows matrix to get noise sample scores
# Call nce_trn with r = ygold' to get real sample scores
# DONE: figure out w/b sharing, can it be a par argument?
# problem is it cannot be initialized at construction, we only get allocation after first sample
# run the first sample through the test network to get real w and b Par's and then point to them.

# DONE: r = generate r for noise samples: 1hr
# keep nsample and psample as training parameters, they should not be stored in the model
# figure out the genious sampling algorithm and implement again
# try to reuse the same sparse csr array to avoid allocation

# Note that we had to create trn1 and trn2 because the r dimensions are different.  Keep that in mind when designing conditional
# compiler and/or minibatch resize.

# DONE: back needs to give us hgrad but not rgrad: 30 mins
# DONE: make sure going back our updates are sparse and efficient. 1hr

# TODO: a lot of things only work on gpu

# DONE: init bias = -log vocab
# DONE: check losscnt
# DONE: check loss functions
# TODO: cannot infer size of lstm if batchsize=1

### DEAD CODE

# function nce_loss_real(ypred::Matrix, kqvec::Vector, ygrad::Matrix; o...)
#     fill!(ygrad,0)
#     n=size(ygrad,1)
#     for i=1:n
#         s = ypred[i,i]
#         kq = kqvec[i]
#         ygrad[i,i] = -(kq / (exp(s) + kq))/n
#     end
#     return ygrad
# end

# @gpu function nce_loss_real{T}(ypred::CudaMatrix{T}, kqvec::CudaVector{T}, ygrad::CudaMatrix{T}; o...)
#     fill!(ygrad,0)
#     # kq = CudaArray(size(m.noise,2) * m.psample[ygold.rowval])
#     T <: Float32 ? ccall((:nce_grad_real_32,libknet),Void,(Cint,Ptr{Cfloat},Ptr{Cfloat},Ptr{Cfloat}), size(ypred,1), ypred, kqvec, ygrad) :
#     T <: Float64 ? ccall((:nce_grad_real_64,libknet),Void,(Cint,Ptr{Cdouble},Ptr{Cdouble},Ptr{Cdouble}), size(ypred,1), ypred, kqvec, ygrad) :
#     error("$T not supported")
#     return ygrad
# end

# @gpu nce_loss_real{T}(ypred::CudaMatrix{T}, kqvec::Vector{T}, ygrad::CudaMatrix{T}; o...)=nce_loss_real(ypred, CudaArray(kqvec), ygrad; o...)

# # J = -log(kq(y))+log(exp s(y)+kq(y))
# # ypred is (K,B); ypred[k,b] is the score of k'th noise sample in b'th sentence
# # kqvec is (K,); kqvec[k] is K * prob(k'th noise sample)
# function nce_loss_noise(ypred::Matrix, kqvec::Vector; o...)
#     (K,B) = size(ypred)
#     J = 0.0
#     for k=1:K
#         kq = kqvec[k]
#         logkq = log(kq)
#         for b=1:B
#             s = ypred[k,b]
#             J += log(exp(s) + kq) - logkq
#         end
#     end
#     return J/B
# end

# @gpu function nce_loss_noise{T}(ypred::CudaMatrix{T}, kqvec::CudaVector{T}; o...)
#     (K,B) = size(ypred)
#     ytemp = similar(ypred)
#     T <: Float32 ? ccall((:nce_loss_noise_32,libknet),Void,(Cint,Cint,Ptr{Cfloat},Ptr{Cfloat},Ptr{Cfloat}), K, B, ypred, kqvec, ytemp) :
#     T <: Float64 ? ccall((:nce_loss_noise_64,libknet),Void,(Cint,Cint,Ptr{Cdouble},Ptr{Cdouble},Ptr{Cdouble}), K, B, ypred, kqvec, ytemp) :
#     error("$T not supported")
#     vecnorm(ytemp,1)/B
# end

# @gpu nce_loss_noise{T}(ypred::CudaMatrix{T}, kqvec::Vector{T}; o...)=nce_loss_noise(ypred,CudaArray(kqvec); o...)

# # dJ/ds = (exp s(y))/(exp s(y) + kq(y))
# function nce_loss_noise(ypred::Matrix, kqvec::Vector, ygrad::Matrix; o...)
#     (K,B) = size(ypred)
#     for k=1:K
#         for b=1:B
#             exps = exp(ypred[k,b])
#             ygrad[k,b] = (exps/(exps+kqvec[k]))/B
#         end
#     end
#     return ygrad
# end

# @gpu function nce_loss_noise{T}(ypred::CudaMatrix{T}, kqvec::CudaVector{T}, ygrad::CudaMatrix{T}; o...)
#     (K,B) = size(ypred)
#     T <: Float32 ? ccall((:nce_grad_noise_32,libknet),Void,(Cint,Cint,Ptr{Cfloat},Ptr{Cfloat},Ptr{Cfloat}), K, B, ypred, kqvec, ygrad) :
#     T <: Float64 ? ccall((:nce_grad_noise_64,libknet),Void,(Cint,Cint,Ptr{Cdouble},Ptr{Cdouble},Ptr{Cdouble}), K, B, ypred, kqvec, ygrad) :
#     error("$T not supported")
#     return ygrad
# end

# @gpu nce_loss_noise{T}(ypred::CudaMatrix{T}, kqvec::Vector{T}, ygrad::CudaMatrix{T}; o...)=nce_loss_noise(ypred,CudaArray(kqvec),ygrad; o...)


text:
using Base.Threads
using ThreadPools
using LazyArrays

import Base: last, iterate
export last, iterate

export ThreadedRootSearch

struct ThreadBuffer{T, V <: AbstractVector{T}}
    root_regions::V
    indeterminate_regions::V
end
ThreadBuffer(region::T) where T = ThreadBuffer(T[], T[])

"""
    explore_region!(buffer, region, contractor, f, tol)

Explore `region` for roots of `f` using `contractor` and store findings in
`buffer`, up to search tolerance `tol`.
Returns parts of `region` which are left unexplored.
"""
function explore_region!(buffer, region, contractor, f, tol)
    image = f(region)
    if !contains_root(image) || _isempty(image)
        return empty.(region)
    end

    contraction = contractor(region)
    if strict_isinterior(contraction, region)
        push!(buffer.root_regions, contraction)
        return empty.(region)
    end

    if diam(region) < tol
        push!(buffer.indeterminate_regions, region)
        return empty.(region)
    end

    return region .∩ contraction
end

struct ThreadedRootSearch{B}
    buffers::B
    contractor
    f
    tol::Float64
    target_task_number::Int
    function ThreadedRootSearch(
        f,
        region,
        contractor,
        tol = default_tolerance,
        target_task_number = nothing,
        n = nothing
    )
        n = isnothing(n) ? nthreads() : n
        buffers = map(i -> ThreadBuffer(region), Tuple(1:n))
        push!(buffers[1].indeterminate_regions, region)
        return new{typeof(buffers)}(
            buffers,
            contractor,
            f,
            tol,
            isnothing(target_task_number) ? 1000n : target_task_number
        )
    end
end

Base.IteratorSize(::Type{ThreadedRootSearch}) = Base.SizeUnknown()
Base.eltype(::Type{ThreadedRootSearch{B}}) where {
    N, T, V, B <: NTuple{N, ThreadBuffer{T, V}}
} = NTuple{2, ApplyVector{T, typeof(vcat), NTuple{N, V}}}

"""
    (search::ThreadedRootSearch)(region, generation, maxgeneration)

Advances the `search` by exploring and, if necessary, bisecting `region`.
Returns `true` when a region is explored up to search tolerance.
Returns `false` when the next step (`generation`) would exceed `maxgeneration`.
If exploration requires further bisection, returns an iterator of `Task`s.
Use `guarded_reduce(&, tasks)` to recursively fetch the boolean result.
"""
function (search::ThreadedRootSearch)(region, generation, maxgeneration)
    buffer = search.buffers[threadid()]

    _region = explore_region!(buffer, region, search.contractor, search.f, search.tol)
    if _isempty(_region)
        return true
    end

    bisection = bisect(_region)

    if generation == maxgeneration
        append!(buffer.indeterminate_regions, bisection)
        return false
    end

    tasks = map(bisection) do i
        Threads.@spawn search(i, generation + 1, maxgeneration)
    end
    # we cannot wait for results here because we need to free up the thread
    return tasks
end

guarded_reduce(op, b) = b
guarded_reduce(op, t::Task) = guarded_reduce(op, fetch(t))
guarded_reduce(op, ts::NTuple{N}) where N =
    mapreduce(t -> guarded_reduce(op, t), op, ts)
guarded_reduce(op, ts::AbstractVector) =
    mapreduce(t -> guarded_reduce(op, t), op, ts)

function pop_unfinished!(buffers)
    return mapreduce(vcat, buffers) do buffer
        unfinished_regions = copy(buffer.indeterminate_regions)
        # remove unfinished regions from buffers without freeing memory
        empty!(buffer.indeterminate_regions)
        return unfinished_regions
    end
end

function chunks(regions, chunksize)
    n = length(regions)
    return map(1:chunksize:n) do i
        return @view regions[i:min(n, i + chunksize - 1)]
    end
end

estimate_generations(target_task_number, nregions) =
    max(1, floor(Int, log2(target_task_number / nregions)))

function (search::ThreadedRootSearch{B})(unfinished_regions) where {N, B <: NTuple{N}}
    return mapreduce(&, chunks(unfinished_regions, N)) do regions
        max_generations = estimate_generations(search.target_task_number, length(regions))
        tasks = qmap(r -> search(r, 1, max_generations), regions)
        return guarded_reduce(&, tasks)
    end
end

function collect_regions(buffers::B)  where {
    N, T, V, B <: NTuple{N, ThreadBuffer{T, V}}
}
    root_regions = ApplyVector{T, typeof(vcat), NTuple{N, V}}(
        vcat,
        map(b -> b.root_regions, buffers)
    )
    indeterminate_regions =  ApplyVector{T, typeof(vcat), NTuple{N, V}}(
        vcat,
        map(b -> b.indeterminate_regions, buffers)
    )
    return root_regions, indeterminate_regions
end

function iterate(search::ThreadedRootSearch, completed = false)
    completed && return nothing
    unfinished_regions = pop_unfinished!(search.buffers)
    completed = unfinished_regions |> search
    return collect_regions(search.buffers), completed
end

function last(search::T) where T <: ThreadedRootSearch
    local res
    for s in search
        res = s
    end
    return res
end


text:
using Plots
gr()
include("PRM.jl")
#include("RRT.jl")

plot([1,2],[2,3])

#cost, isPathFound, nlist = rrtPathPlanner(30)
#@show isPathFound
#@show cost
#@show length(nlist)
#@show nlist

iterList = [40 80 100 120 150 180 200 300]
nTrials = 30
cost = 0
listCosts = Vector{Float32}()
listpSucc= Vector{Float32}()

#
# print("\n\n")
# for testIter in iterList
    # totalcost =0
    # idx = 0
    # @show testIter
    # nSuccess = 0
    # cost = 0
    # while idx < nTrials
        # cost, isPathFound, nlist = rrtPathPlanner(testIter)
        # if isPathFound
            # nSuccess += 1
        # end
        # totalcost += cost
        # idx += 1
    # end
    # avgCost = totalcost /  nSuccess
    # @show nSuccess
    # pSucc = nSuccess / nTrials
    # @show avgCost
    # @show listCosts
    # @show listpSucc
    # print("\n\n")
    # #@printf("For the iter of %d the avg cost was %d across %d trials", maxIter, avgCost, nTrials)
    # push!(listCosts, avgCost)
    # push!(listpSucc, pSucc)
# end
#
# @show iterList
# @show listCosts
# @show listpSucc
#
# sizeplot = (300,300)
#
# pRRTcost = plot(iterList, listCosts', show=true, size=sizeplot, seriestype=:scatter,
     # legend=false, yaxis=((0,100), 0:20:100), xaxis=((0,320), 0:50:300), color=:black)
#
# title!("RRT nIter vs pathcost")
# xlabel!("iterations requested of RRT alg")
# ylabel!("euclidean path cost")
#
#
# pRRTsuccess = scatter(iterList, listpSucc', legend=false, yaxis=((0,1), 0:0.1:1), xaxis=((0,320), 0:50:300), color=:orange)
# title!("RRT nIter vs pSuccess")
# ylabel!("# P(success) = numSucc / 30 trials")
#
# plot(pRRTcost, pRRTsuccess, layout=(1,2),legend=false)
#

#####################################

print("\n\n")
connectDist = 15
iterList = [40 80 100 120 150 180 200 300]

for testIter in iterList
    totalcost =0
    idx = 0
    @show testIter
    nSuccess = 0
    cost = 0
    while idx < nTrials
        nodeslist, edgeslist = preprocessPRM(testIter, connectDist)
        cost, isPathFound, winPath = queryPRM(start, goal, nodeslist, edgeslist)
        if isPathFound
            nSuccess += 1
        end
        totalcost += cost
        idx += 1
    end
    avgCost = totalcost /  nSuccess
    @show nSuccess
    pSucc = nSuccess / nTrials
    @show avgCost
    @show listCosts
    @show listpSucc
    print("\n\n")
    #@printf("For the iter of %d the avg cost was %d across %d trials", maxIter, avgCost, nTrials)
    push!(listCosts, avgCost)
    push!(listpSucc, pSucc)
end

@show iterList
@show listCosts
@show listpSucc

sizeplot = (300,300)

pPRMcost = plot(iterList, listCosts', show=true, size=sizeplot, seriestype=:scatter,
     legend=false, yaxis=((0,100), 0:20:100), xaxis=((0,320), 0:50:300), color=:black)

title!("RRT nIter vs pathcost")
xlabel!("iterations requested of RRT alg")
ylabel!("euclidean path cost")

pPRMsuccess = scatter(iterList, listpSucc', legend=false, yaxis=((0,1), 0:0.1:1), xaxis=((0,320), 0:50:300), color=:orange)
title!("RRT nIter vs pSuccess")
ylabel!("# P(success) = numSucc / 30 trials")

plot(pRRTcost, pRRTsuccess, layout=(1,2),legend=false)

########################################
#pRRTsuccess[:axis]

#y = rand(100)
#lot(0:10:100,rand(11,4),lab="lines",w=3,palette=:grays,fill=0,α=0.6)
#scatter!(y,zcolor=abs(y - 0.5),m=(:heat,0.8,stroke(1,:green)),ms=10 * abs(y - 0.5) + 4,lab="grad")

#cost, isPathFound, nlist = rrtPathPlanner(40) #maxIter
#plotPath(isPathFound,nlist)