text:
#' # Modeling data 2

#' ## Building a model
#' Recall that in notebook 3, we saw that we could use a mathematical function
#' to classify an image as an apple or a banana, based on the average amount of
#' green in an image:
#'
#' ![data_flow](data/data_flow.png)
#'
#' ![what_is_model](data/what_is_model.png)

#' A common function for performing this kind of **classification** is the
#' sigmoid that we saw in the last notebook, and that we will now extend by
#' adding two **parameters**, $w$ and $b$:
#'
#' $$\sigma(x; w, b) := \frac{1}{1 + e^{(-wx + b)}}$$
#'
#' $$x = \mathrm{data}$$
#'
#'
#' $$\begin{align}
#' \sigma(x;w,b) &\approx 0 \implies \mathrm{apple} \\
#' \sigma(x;w,b) &\approx 1 \implies \mathrm{banana}
#' \end{align}$$
#'
#' In our mathematical notation above, the `;` in the function differentiates
#' between the **data** and the **parameters**. `x` is the data and is
#' determined from the image. The parameters, `w` and `b`, are numbers which we
#' choose to make our function match the results it should be modeling.

#' Note that in the code below, we don't distinguish between data and
#' parameters - both are just inputs to our function, σ!

using Images, Statistics

# apple = load("C:\\Users\\Yi-Kun.Yang\\work\\org.allnix.julia\\JuliaAcademy\\FundMachLearn/data/10_100.jpg")
# banana = load("C:\\Users\\Yi-Kun.Yang\\work\\org.allnix.julia\\JuliaAcademy\\FundMachLearn/data/104_100.jpg")
apple = load("data/10_100.jpg")
#'
banana = load("data/104_100.jpg")
#'

apple_green_amount = mean(Float64.(green.(apple)))
banana_green_amount = mean(Float64.(green.(banana)))

println("Average green for apple = $apple_green_amount")
println("Average green for banana = $banana_green_amount")

#+ results="hidden"
σ(x, w, b) = 1 / (1 + exp(-w * x + b))
#+

#' What we want is that when we give σ as input the average green for the apple,
#' roughly `x = 0.3385`, it should return as output something close to 0,
#' meaning "apple". And when we give σ the input `x = 0.8808`, it should output
#' something close to 1, meaning "banana".

#' By changing the parameters of the function, we can change the shape of the
#' function, and hence make it represent, or **fit**, the data better!

#' ## Data fitting by varying parameters

# We can understand how our choice of `w` and `b` affects our model by seeing how our values for `w` and `b` change the plot of the $\sigma$ function.

using Plots
gr()   # GR works better for interactive manipulations
#plotly()
#pyplot()

#' Run the code in the next cell. You should see two "sliders" appear, one for
#' `w` and one for `b`.

#' **Game**:
#' Change w and b around until the blue curve, labeled "model", which is the
#' graph of the `\sigma` function, passes through *both* of the data points at
#' the same time.

w = 10.0 # try manipulating w between -10 and 30
b = 10.0 # try manipulating b between 0 and 20

using Interact

# Function for Interact
function myplot(w,b)
      pt = plot()
      plot!(pt, x -> σ(x, w, b), xlim=(-0,1), ylim=(-0.1,1.1), label="model", legend=(0.1,0.9), lw=3)
      scatter!(pt, [apple_green_amount],  [0.0], label="apple", ms=5)   # marker size = 5
      scatter!(pt, [banana_green_amount], [1.0], label="banana", ms=5)

      return pt
end

#' Notice that the two parameters do two very different things. The **weight**,
#' `w`, determines *how fast* the transition between 0 and 1 occurs. It encodes
#' how trustworthy we think our data  actually is, and in what range we should
#' be putting points between 0 and 1 and thus calling them "unsure". The
#' **bias**, `b`, encodes *where* on the $x$-axis the switch should take place.
#' It can be seen as shifting the function left-right. We'll come to understand
#' these *parameters* more in notebook 6.

#' Here are some parameter choices that work well:

w = 25.58 #20.8
b = 15.6 #12.4
pt = myplot(w,b)

#+ results="hidden"
# Run this in Juno or Jupyter to get interative plot
# @manipulate for w=-10:0.2:30,b = 0:0.2:20
#      vbox(hbox(w,b),myplot(w,b))
# end
#+

#' (Note that in this problem there are many combinations of `w` and `b` that
#' fit the data well.)
apple2 = load("data/107_100.jpg")
#'
green_amount = mean(Float64.(green.(apple2)))
@show green_amount
scatter!(pt, [green_amount], [0.0], label="new apple")

#' Our model successfully says that our new image is an apple! Pat yourself on
#' the back: you've actually just trained your first neural network!

#' #### Exercise 1

#' Load the image of a banana in `data/8_100.jpg` as `mybanana`. Edit the code
#' below to calculate the amount of green in `mybanana` and to overlay data for
#' this image with the existing model and data points.
mybanana = load("data/8_100.jpg")
green_amount = mean(Float64.(green.(mybanana)))
@show green_amount
scatter!(pt, [green_amount], [1.0], label="new banana")
#' ## Closing remarks: bigger models, more data, more accuracy

#' That last apple should start making you think: not all apples are red; some
#' are yellow. "Redness" is one attribute of being an apple, but isn't the
#' whole thing. What we need to do is incorporate more ideas into our model by
#' allowing more inputs. However, more inputs would mean more parameters to
#' play with. Also, we would like to have the computer start "learning" on its
#' own, instead of modifying the parameters ourselves until we think it "looks
#' right". How do we take the next step?

#' The first thing to think about is, if you wanted to incorporate more data
#' into the model, how would you change the sigmoid function? Play around with
#' some ideas. But also, start thinking about how you chose parameters. What
#' process did you do to finally end up at good parameters? These two problems
#' (working with models with more data and automatically choosing parameters)
#' are the last remaining step to understanding deep learning.


text:
function make_metadata(ts::IS.TimeSeriesData)
    return IS.time_series_data_to_metadata(typeof(ts))(ts)
end

"""
Helper function that gets all values and then deserializes a full object.
"""
function _deserialize_full(storage, ts)
    ts_metadata = make_metadata(ts)
    return IS.deserialize_time_series(
        IS.SingleTimeSeries,
        storage,
        ts_metadata,
        UnitRange(1, length(ts)),
        UnitRange(1, 1),
    )
end

function test_add_remove(storage::IS.TimeSeriesStorage)
    name = "component1"
    name = "val"
    component = IS.TestComponent(name, 5)
    ts = IS.SingleTimeSeries(data = create_time_array(), name = "test")
    IS.serialize_time_series!(storage, IS.get_uuid(component), name, ts)

    ts2 = _deserialize_full(storage, ts)
    @test TimeSeries.timestamp(IS.get_data(ts2)) == TimeSeries.timestamp(IS.get_data(ts))
    @test TimeSeries.values(IS.get_data(ts2)) == TimeSeries.values(IS.get_data(ts))

    component2 = IS.TestComponent("component2", 6)
    IS.serialize_time_series!(storage, IS.get_uuid(component2), name, ts)

    @test IS.get_num_time_series(storage) == 1

    IS.remove_time_series!(storage, IS.get_uuid(ts), IS.get_uuid(component2), name)

    ## There should still be one reference to the data.
    ts2 = _deserialize_full(storage, ts)
    @test IS.get_data(ts2) isa TimeSeries.TimeArray

    IS.remove_time_series!(storage, IS.get_uuid(ts), IS.get_uuid(component), name)
    @test_throws ArgumentError _deserialize_full(storage, ts)
    return IS.get_num_time_series(storage) == 0
end

function test_add_references(storage::IS.TimeSeriesStorage)
    name = "val"
    component1 = IS.TestComponent("component1", 5)
    component2 = IS.TestComponent("component2", 6)
    ts = IS.SingleTimeSeries(data = create_time_array(), name = "test")
    ts_uuid = IS.get_uuid(ts)
    IS.serialize_time_series!(storage, IS.get_uuid(component1), name, ts)
    IS.add_time_series_reference!(storage, IS.get_uuid(component2), name, ts_uuid)

    # Adding duplicate references is not allowed.
    @test_throws AssertionError IS.add_time_series_reference!(
        storage,
        IS.get_uuid(component2),
        name,
        ts_uuid,
    )

    @test IS.get_num_time_series(storage) == 1

    IS.remove_time_series!(storage, ts_uuid, IS.get_uuid(component1), name)

    # There should still be one reference to the data.
    @test _deserialize_full(storage, ts) isa IS.TimeSeriesData

    IS.remove_time_series!(storage, ts_uuid, IS.get_uuid(component2), name)
    @test_throws ArgumentError _deserialize_full(storage, ts)
    return IS.get_num_time_series(storage) == 0
end

function test_get_subset(storage::IS.TimeSeriesStorage)
    name = "component1"
    name = "val"
    component = IS.TestComponent(name, 1)
    ts = IS.SingleTimeSeries(data = create_time_array(), name = "test")
    IS.serialize_time_series!(storage, IS.get_uuid(component), name, ts)
    ts2 = _deserialize_full(storage, ts)

    @test TimeSeries.timestamp(IS.get_data(ts2)) == TimeSeries.timestamp(IS.get_data(ts))
    rows = UnitRange(3, 8)
    columns = UnitRange(1, 1)
    ts_metadata = make_metadata(ts)
    ts_subset =
        IS.deserialize_time_series(IS.SingleTimeSeries, storage, ts_metadata, rows, columns)
    @test IS.get_data(ts_subset)[1] == IS.get_data(ts2)[rows.start]
    @test length(ts_subset) == length(rows)

    initial_time1 = Dates.DateTime("2020-09-01")
    resolution = Dates.Hour(1)
    initial_time2 = initial_time1 + resolution
    name = "test"
    horizon = 24
    data = SortedDict(initial_time1 => ones(horizon), initial_time2 => ones(horizon))

    ts = IS.Deterministic(data = data, name = name, resolution = resolution)
    IS.serialize_time_series!(storage, IS.get_uuid(component), name, ts)
    ts_metadata = make_metadata(ts)
    rows = UnitRange(1, horizon)
    columns = UnitRange(1, 2)
    ts2 = IS.deserialize_time_series(IS.Deterministic, storage, ts_metadata, rows, columns)
    @test collect(IS.get_initial_times(ts2)) == collect(IS.get_initial_times(ts))
    @test collect(IS.iterate_windows(ts2)) == collect(IS.iterate_windows(ts))

    rows = UnitRange(3, 8)
    columns = UnitRange(1, 2)
    ts_subset =
        IS.deserialize_time_series(IS.Deterministic, storage, ts_metadata, rows, columns)
    @test IS.get_horizon(ts_subset) == length(rows)
    @test IS.get_count(ts_subset) == columns.stop
    @test IS.get_initial_timestamp(ts_subset) ==
          initial_time1 + resolution * (rows.start - 1)

    rows = UnitRange(2, 7)
    columns = UnitRange(1, 1)
    ts_subset =
        IS.deserialize_time_series(IS.Deterministic, storage, ts_metadata, rows, columns)
    @test IS.get_horizon(ts_subset) == length(rows)
    @test IS.get_count(ts_subset) == columns.stop
    @test IS.get_initial_timestamp(ts_subset) ==
          initial_time1 + resolution * (rows.start - 1)
end

function test_clear(storage::IS.TimeSeriesStorage)
    name = "component1"
    name = "val"
    component = IS.TestComponent(name, 5)
    ts = IS.SingleTimeSeries(data = create_time_array(), name = "test")
    IS.serialize_time_series!(storage, IS.get_uuid(component), name, ts)

    ts2 = _deserialize_full(storage, ts)
    @test TimeSeries.timestamp(IS.get_data(ts2)) == TimeSeries.timestamp(IS.get_data(ts))
    @test TimeSeries.values(IS.get_data(ts2)) == TimeSeries.values(IS.get_data(ts))

    IS.clear_time_series!(storage)
    @test_throws ArgumentError _deserialize_full(storage, ts)
end

@testset "Test time series storage implementations" begin
    for in_memory in (true, false)
        test_add_remove(IS.make_time_series_storage(; in_memory = in_memory))
        test_get_subset(IS.make_time_series_storage(; in_memory = in_memory))
        test_clear(IS.make_time_series_storage(; in_memory = in_memory))
    end

    test_add_remove(IS.make_time_series_storage(; in_memory = false, directory = "."))
    test_get_subset(IS.make_time_series_storage(; in_memory = false, directory = "."))
    test_clear(IS.make_time_series_storage(; in_memory = false, directory = "."))
end

@testset "Test copy time series references" begin
    for in_memory in (true, false)
        test_add_remove(IS.make_time_series_storage(; in_memory = in_memory))
        test_add_references(IS.make_time_series_storage(; in_memory = in_memory))
        test_get_subset(IS.make_time_series_storage(; in_memory = in_memory))
        test_clear(IS.make_time_series_storage(; in_memory = in_memory))
    end
end

@testset "Test data format version" begin
    storage = IS.make_time_series_storage(in_memory = false)
    @test IS.read_data_format_version(storage) == IS.TIME_SERIES_DATA_FORMAT_VERSION
end

@testset "Test compression" begin
    in_memory = false
    for type in (IS.CompressionTypes.BLOSC, IS.CompressionTypes.DEFLATE)
        for shuffle in (true, false)
            compression = IS.CompressionSettings(
                enabled = true,
                type = type,
                level = 5,
                shuffle = shuffle,
            )
            test_add_remove(
                IS.make_time_series_storage(;
                    in_memory = in_memory,
                    compression = compression,
                ),
            )
            test_add_references(
                IS.make_time_series_storage(;
                    in_memory = in_memory,
                    compression = compression,
                ),
            )
            test_get_subset(
                IS.make_time_series_storage(;
                    in_memory = in_memory,
                    compression = compression,
                ),
            )
            test_clear(
                IS.make_time_series_storage(;
                    in_memory = in_memory,
                    compression = compression,
                ),
            )
        end
    end
end

@testset "Test isempty" begin
    for in_memory in (true, false)
        storage = IS.make_time_series_storage(in_memory = in_memory)
        @test isempty(storage)
        name = "component1"
        name = "val"
        component = IS.TestComponent(name, 5)
        ts = IS.SingleTimeSeries(data = create_time_array(), name = "test")
        IS.serialize_time_series!(storage, IS.get_uuid(component), name, ts)
        @test !isempty(storage)
    end
end

ext: jl


text:
module Default
    using HTTP
    export default

    function default(req, res, route)
        res[:text] = this -> x -> let
            this.status = 200
            this.body = string(x)
            this.content_type = "text/plain"
        end

        res[:code] = this -> x -> let
            this.status = x
            # TODO: check x and add information?
        end

        res[:html] = this -> x -> let
            this.status = 200
            this.body = x
            this.content_type = "text/html"
        end

        req[:params] = this -> this.params = HTTP.URIs.queryparams(this.uri)

        req.body = read(req._http, String)
        route.next()

        if res.route_status != nothing
            @error "not found" req.method req.uri.uri
            res.code(res.route_status)
        end

        if res.body != nothing
            res.content_type != nothing && push!(res.headers, "Content-Type" => res.content_type)
            res.body isa IO || push!(res.headers, "Content-Length" => string(sizeof(res.body)))
        end
    end
end


text:
module JDQZBench

using JacobiDavidson
using BenchmarkTools

function bench_jdqz_allocs(; n = 10_000, τ = 0.0 + 0.01im)
  srand(4)

  A = 2 * speye(ComplexF64, n) + sprand(ComplexF64, n, n, 1 / n)
  B = 2 * speye(ComplexF64, n) + sprand(ComplexF64, n, n, 1 / n)

  result = jdqz(
    A,
    B,
    BiCGStabl(A, max_mv_products = 10),
    τ = 0.0 + 0.0im,
    pairs = 20,
    max_iter = 1000,
    subspace_dimensions = 10:15,
    tolerance = 1e-8,
    verbosity = 1
  )
end

function bench_jdqz(; n = 1_000, τ = 0.0 + 0.01im)
  srand(4)

  A = 2 * speye(ComplexF64, n) + sprand(ComplexF64, n, n, 1 / n)
  B = 2 * speye(ComplexF64, n) + sprand(ComplexF64, n, n, 1 / n)

  eig_bench = @benchmark eigs($A, $B, nev = 20, sigma = 0.0 + 0.0im, ritzvec = true, tol = 1e-8)

  jdqz_bench = @benchmark $jdqz(
    $A,
    $B,
    $BiCGStabl($A, max_mv_products = 10),
    τ = 0.0 + 0.0im,
    pairs = 20,
    max_iter = 1000,
    subspace_dimensions = 10:15,
    tolerance = 1e-8,
  )

  eig_bench, jdqz_bench
end

end

JDQZBench.bench_jdqz_allocs()
Profile.clear()
@profile JDQZBench.bench_jdqz_allocs()

max_forks_repo_name: haampie/jacobi-davidson-julia


text:
using Yota

include("primitives.jl")
include("coupling.jl")

# function g_func(c::Coupling, z)
    # mask, s, t = c.mask, c.s, c.t
    # zp, zn = mask .* z, (1 .- mask) .* z  # z positive, z negative
    # return zp .+ (zn - t(zp)) .* exp.(-s(zp))
# end

function logprob(c::Coupling, x)
end

# fun(c, x) = sum(g_func(c, f_func(c, x)))

mutable struct RealVNP
    c1::Coupling
    c2::Coupling
end

function RealNVP(xz_len, u_len)
    mlps = [MLP3(xz_len, u_len, u_len, xz_len) for i=1:4]
    mask = vcat([1 for i=1:Int(xz_len / 2)], [0 for i=1:Int(xz_len / 2)])
    RealNVP(Coupling(mlps[1], mlps[2]),
            Coupling(mlps[3], mlps[4]))
end

function main2()
    mask = [0, 0, 0, 0, 0, 1, 1, 1, 1, 1]
    s = MLP3{Float64}(10, 5, 5, 10)
    t = MLP3{Float64}(10, 5, 5, 10)
    c = Coupling(mask, s, t)

    x = rand(10, 20)
    z = f_func(c, x)
    g_func(c, z) .- x
end


text:
# Julia wrapper for header: /usr/include/libavutil/threadmessage.h
# Automatically generated using Clang.jl wrap_c, version 0.0.0

export
    av_thread_message_queue_alloc,
    av_thread_message_queue_free,
    av_thread_message_queue_send,
    av_thread_message_queue_recv,
    av_thread_message_queue_set_err_send,
    av_thread_message_queue_set_err_recv,
    av_thread_message_queue_set_free_func,
    av_thread_message_queue_nb_elems,
    av_thread_message_flush

function av_thread_message_queue_alloc(mq, nelem::Integer, elsize::Integer)
    ccall((:av_thread_message_queue_alloc, libavutil), Cint, (Ptr{Ptr{AVThreadMessageQueue}}, UInt32, UInt32), mq, nelem, elsize)
end

function av_thread_message_queue_free(mq)
    ccall((:av_thread_message_queue_free, libavutil), Cvoid, (Ptr{Ptr{AVThreadMessageQueue}},), mq)
end

function av_thread_message_queue_send(mq, msg, flags::Integer)
    ccall((:av_thread_message_queue_send, libavutil), Cint, (Ptr{AVThreadMessageQueue}, Ptr{Cvoid}, UInt32), mq, msg, flags)
end

function av_thread_message_queue_recv(mq, msg, flags::Integer)
    ccall((:av_thread_message_queue_recv, libavutil), Cint, (Ptr{AVThreadMessageQueue}, Ptr{Cvoid}, UInt32), mq, msg, flags)
end

function av_thread_message_queue_set_err_send(mq, err::Integer)
    ccall((:av_thread_message_queue_set_err_send, libavutil), Cvoid, (Ptr{AVThreadMessageQueue}, Cint), mq, err)
end

function av_thread_message_queue_set_err_recv(mq, err::Integer)
    ccall((:av_thread_message_queue_set_err_recv, libavutil), Cvoid, (Ptr{AVThreadMessageQueue}, Cint), mq, err)
end

function av_thread_message_queue_set_free_func(mq, free_func)
    ccall((:av_thread_message_queue_set_free_func, libavutil), Cvoid, (Ptr{AVThreadMessageQueue}, Ptr{Cvoid}), mq, free_func)
end

function av_thread_message_queue_nb_elems(mq)
    ccall((:av_thread_message_queue_nb_elems, libavutil), Cint, (Ptr{AVThreadMessageQueue},), mq)
end

function av_thread_message_flush(mq)
    ccall((:av_thread_message_flush, libavutil), Cvoid, (Ptr{AVThreadMessageQueue},), mq)
end


text:
using DEBTrait, Test
using Random
Random.seed!(123)

######## benchmark K_D against Tang et al. (2019)
D_S     = 1.4e-9*ones(1)
N_p     = 3000*ones(1,1)
r_c     = 1e-6
V_cell  = (4/3)*pi*r_c^3*ones(1)
ρ_p     = DEBTrait.transporter_density(V_cell, N_p)
K_D_0   = DEBTrait.specific_reference_affinity(ρ_p, V_cell, D_S)
@test K_D_0[1] ≈ 5.7969874109e-05 atol=1e-6
#
pct_sand = 70.0
pct_clay = 30.0
s_sat    = 0.10101010101010102
N_cell   = 10.0*ones(1)
Soil = BaseSoil(pct_sand, pct_clay, s_sat, N_cell)
K_D     = DEBTrait.affinity_constant(ρ_p, V_cell, D_S, Soil.pct_sand, Soil.pct_clay, Soil.s_sat, Soil.N_cell)
@test K_D[1] ≈ 1.447412099595345 atol=1e-6
#
Soil.s_sat = 0.8080808080808082
K_D     = DEBTrait.affinity_constant(ρ_p, V_cell, D_S, Soil.pct_sand, Soil.pct_clay, Soil.s_sat, Soil.N_cell)
@test K_D[1] ≈ 9.390492906747054e-05 atol=1e-6
#
N_SB    = DEBTrait.monomer_uptake_sites(V_cell, ρ_p)

########
n_microbes      = 1
n_substrates    = 1
Isolate         = BaseGenome(1e6*ones(n_microbes), 10*ones(n_microbes), 1.0*ones(n_microbes), 1.0*ones(n_substrates,n_microbes))
Monomer         = BaseMonomer(1e-10*ones(n_substrates), 6*ones(n_substrates), 3*ones(n_substrates), 0.3*ones(n_substrates))
Soil            = BaseSoil(0.0, 0.0, 1.0, 1*ones(n_microbes))
MetabolismC     = PMetabolismC{BaseGenome, Array{Float64,1}}(Isolate)
MetabolismC.y_EM = ones(n_microbes)
MetabolismC.y_EX = ones(n_microbes)
MetabolismC.α    = zeros(n_microbes)
AssimilationC   = PAssimilation(Isolate, Monomer, Soil, MetabolismC)
@test size(AssimilationC.ρ_p) == (n_substrates, n_microbes)
@test size(AssimilationC.N_SB) == (n_substrates, n_microbes)
@test size(AssimilationC.K_D) == (n_substrates, n_microbes)
@test AssimilationC.K_D_0 ≈ AssimilationC.K_D atol = 1e-4
Soil.s_sat = 0.05
AssimilationC   = PAssimilation(Isolate, Monomer, Soil, MetabolismC)
@test AssimilationC.K_D_0[1] < AssimilationC.K_D[1]

n_microbes      = rand(1:10)
n_substrates    = rand(1:10)
Isolate         = BaseGenome(1e6*ones(n_microbes), 10*ones(n_microbes), 1.0*ones(n_microbes), 1.0*ones(n_substrates,n_microbes))
Monomer         = BaseMonomer(1e-10*ones(n_substrates), 6*ones(n_substrates), 3*ones(n_substrates), 0.3*ones(n_substrates))
Soil            = BaseSoil(0.0, 0.0, 1.0, 10*ones(n_microbes))
MetabolismC     = PMetabolismC{BaseGenome, Array{Float64,1}}(Isolate)
MetabolismC.y_EM = ones(n_microbes)
MetabolismC.y_EX = ones(n_microbes)
MetabolismC.α    = zeros(n_microbes)
AssimilationC   = PAssimilation(Isolate, Monomer, Soil, MetabolismC)
@test size(AssimilationC.ρ_p) == (n_substrates, n_microbes)
@test size(AssimilationC.N_SB) == (n_substrates, n_microbes)
@test size(AssimilationC.K_D) == (n_substrates, n_microbes)


text:
module Callsite

# sketch of a call site macro, which support different stages
#so stuff could look like this:
function remove_slots(slot::Slot)
    sym = symbol("_$(slot.id)")
    if slot.typ == Any
        return sym
    else
        return Expr(:(::), sym, slot.typ)
    end
end
remove_slots(no_slot) = no_slot
function remove_slots(expressions::Vector)
    out = []
    for elem in expressions
        if isa(elem, Expr)
            args = remove_slots(elem.args)
            expr = Expr(elem.head); expr.args = args
            push!(out, expr)
        else
            push!(out, remove_slots(elem))
        end
    end
    out
end
callsite_typed(transformation, f, types) = callsite_lambda(code_typed, transformation, f, types)
callsite_lowered(transformation, f, types) = callsite_lambda(code_lowered, transformation, f, types)
function callsite_lambda(stage, transformation, f, types)
    names = [symbol("_$(i+1)") for i=1:length(types)]
    lam = stage(f, types)[1]
    typenames = [Expr(:(::), name, typ) for (name, typ) in zip(names, types)]

    u_code = Base.uncompressed_ast(lam)
    body = transformation(u_code, typenames)
    funbody = Expr(:block);
    renamed_args = [:($(names[i])::$(types[i]) = types[$i]::$(types[i])) for i=1:length(types)]
    append!(funbody.args, renamed_args)
    append!(funbody.args, remove_slots(body))
    funbody
end
function callsite_llvm(transformation, f, types)
    io = IOBuffer
    ir = takebuf_string(code_llvm(io, f.instance, types))
    expr = transformation(ir, types) # this should rather return IR?
end

callsite_stage_dict = Dict{Any, Function}()

macro callsite(stage, func)
    @assert func.head == :function
    @assert stage == :typed || stage == :lowered || stage == :llvm
    esc(quote
        $func
        Callsite.callsite_stage_dict[typeof($(func.args[1].args[1]))] = Callsite.$(symbol("callsite_$stage"))
    end)
end

@generated function _call{N}(transformation, f, types::NTuple{N}) # this needs to be _1, _2 for inline to actually work
    stage = callsite_stage_dict[transformation]
    stage(transformation.instance, f.instance, tuple(types.parameters...))
end

export _call, @callsite

end # module


text:
abstract type Proposal end

struct Symmetric <: Proposal end
struct Prior <: Proposal end
struct Asymmetric <: Proposal end

const SYMMETRIC = Symmetric()
const PRIOR = Prior()
const ASYMMETRIC = Asymmetric()

struct Metropolis <: InferenceType end
const METROPOLIS = Metropolis()
export METROPOLIS

bare_metropolis_results() = BareResults{Metropolis}(METROPOLIS, DefaultDict{Any, Vector{Any}}([]))
structured_metropolis_results() = StructuredResults{Metropolis}(METROPOLIS, Dict{String, Vector{Any}}(), Float64[])
metropolis_results() = NonparametricSamplingResults{Metropolis}(METROPOLIS, Array{Float64, 1}(), Array{Any, 1}(), Array{Trace, 1}())
metropolis_results(A::DataType, D::DataType) = NonparametricSamplingResults{Metropolis}(METROPOLIS, Array{Float64, 1}(), Array{Any, 1}(), Array{TypedTrace{A, D}, 1}())
symmetric(f :: F) where F <: Function = false

### independent prior metropolis stuff ###

@doc raw"""
    function loglatent(t :: Trace)

Computes the joint log probability of all latent variables in a trace, ``\log p(t) - \ell(t)``.
"""
function loglatent(t :: Trace)
    l = 0.0
    for v in values(t)
        if !v.observed
            l += v.logprob_sum
        end
    end
    l
end

@doc raw"""
    function log_acceptance_ratio(t :: Trace, t_proposed :: Trace, p :: Prior)

Computes the log acceptance ratio of a Metropolis step when using the independent prior proposal
algorithm:

``
\log \alpha = \ell(t_{\text{proposed}}) - \ell(t_{\text{original}})
``
"""
function log_acceptance_ratio(t :: Trace, t_proposed :: Trace, p :: Prior)
    log_proposal_lik = loglikelihood(t_proposed)
    log_orig_lik = loglikelihood(t)
    log_proposal_lik - log_orig_lik
end

@doc raw"""
    function accept(t :: Trace, new_t :: Trace, log_a :: Float64)

Stochastic function that either returns `new_t` if accepted or returns `t` if not accepted.
"""
accept(t :: Trace, new_t :: Trace, log_a :: Float64) = log(rand()) < log_a ? new_t : t

@doc raw"""
    function mh_step(t :: Trace, f; params = (), return_val :: Bool = false)

An independent prior sample Metropolis step.

Given a trace `t` and stochatic function `f` depending on `params...`, generates proposals
from prior draws and accepts based on the likelihood ratio.
"""
function mh_step(t :: Trace, f; params = (), return_val :: Bool = false)
    new_t = trace()
    r = f(new_t, params...)
    log_a = log_acceptance_ratio(t, new_t, PRIOR)
    a = accept(t, new_t, log_a)
    if !return_val
        a
    else
        (r, a)
    end
end

@doc raw"""
    function mh_step(t :: Trace, f, types::Tuple{DataType, DataType}; params = (), return_val :: Bool = false)

An independent prior sample Metropolis step.

Given a trace `t` and stochastic function `f` depending on `params...`, generates proposals
from prior draws and accepts based on the likelihood ratio.
"""
function mh_step(t :: Trace, f, types::Tuple{DataType, DataType}; params = (), return_val :: Bool = false)
    new_t = trace(types[1], types[2])
    r = f(new_t, params...)
    log_a = log_acceptance_ratio(t, new_t, PRIOR)
    a = accept(t, new_t, log_a)
    if !return_val
        a
    else
        (r, a)
    end
end

@doc raw"""
    function mh(f; params = (), burn = 1000, thin = 50, num_iterations = 10000)

Generic Metropolis algorithm using draws from the prior.

Args:

+ `f`: stochastic function. Must have call signature `f(t :: Trace, params...)`
+ `params`: addditional arguments to pass to `f` and each of the proposal kernels.
+ `burn`: number of samples to discard at beginning of markov chain
+ `thin`: keep only every `thin`-th draw. E.g., if `thin = 100`, only every 100-th trace will be kept.
+ `num_iterations`: total number of steps to take in the markov chain
"""
function mh(f; params = (), burn = 1000, thin = 50, num_iterations = 10000)
    results = metropolis_results()
    t = trace()
    r = f(t, params...)
    for n in 1:num_iterations
        r, t = mh_step(t, f; params = params, return_val = true)
        if (n > burn) && (n % thin == 0)
            push!(results.return_values, r)
            push!(results.traces, t)
            push!(results.log_weights, logprob(t))
        end
    end
    results
end

@doc raw"""
    function mh(f, types::Tuple{DataType, DataType}; params = (), burn = 1000, thin = 50, num_iterations = 10000)

Generic Metropolis algorithm using draws from the prior.

Args:

+ `f`: stochastic function. Must have call signature `f(t :: Trace, params...)`
+ `params`: addditional arguments to pass to `f` and each of the proposal kernels.
+ `burn`: number of samples to discard at beginning of markov chain
+ `thin`: keep only every `thin`-th draw. E.g., if `thin = 100`, only every 100-th trace will be kept.
+ `num_iterations`: total number of steps to take in the markov chain
"""
function mh(f, types::Tuple{DataType, DataType}; params = (), burn = 1000, thin = 50, num_iterations = 10000)
    results = metropolis_results(types...)
    t = trace(types...)
    r = f(t, params...)
    for n in 1:num_iterations
        r, t = mh_step(t, f, types; params = params, return_val = true)
        if (n > burn) && (n % thin == 0)
            push!(results.return_values, r)
            push!(results.traces, t)
            push!(results.log_weights, logprob(t))
        end
    end
    results
end

function sample(r :: NonparametricSamplingResults{Metropolis}, k, n :: Int)
    v = r[k]
    if n == 1
        StatsBase.sample(v)
    else
        StatsBase.sample(v, n)
    end
end

function sample(r :: ParametricSamplingResults{Metropolis}, k, n :: Int)
    v = getsampled(r, k)
    if n == 1
        StatsBase.sample(v)
    else
        StatsBase.sample(v, n)
    end
end

### general metropolis stuff ###

@doc raw"""
    function logprob(t0 :: Trace, t1 :: Trace)

Computes the proposal log probability ``q(t_1 | t_0)``.

This expression has two parts: log probability that is generated at the proposed site(s), and
log probability that is generated at the sites that are present in `t1` but not in `t0`.
"""
function logprob(t0 :: Trace, t1 :: Trace)
    log_q = 0.0
    kt0 = keys(t0)
    kt1 = keys(t1)
    # get lp of new addresses in domain
    new_addresses = setdiff(kt1, kt0)  # 0 for static structure
    for a in new_addresses
        log_q += t1[a].logprob_sum
    end
    # get lp of addresses selected by proposal
    common_addresses = intersect(kt0, kt1)
    for a in common_addresses
        if t0[a].value != t1[a].value
            log_q += t1[a].logprob_sum
        end
    end
    log_q
end

@doc raw"""
    function copy_common!(old_t :: Trace, new_t :: Trace)

Copies nodes from `old_t` into `new_t` for all addresses in the intersection
of their address sets.
"""
function copy_common!(old_t :: Trace, new_t :: Trace)
    common_addresses = intersect(keys(old_t), keys(new_t))
    for a in common_addresses
        new_t[a] = old_t[a]
    end
    new_t
end

# proposal signature
# function q(old_trace :: Trace, new_trace :: Trace, params...)

@doc raw"""
    function mh_step(t :: Trace, f, q; params = (), return_val :: Bool = false)

A generic Metropolis step using an arbitrary proposal kernel.

Given a trace `t`, a stochastic function `f` with signature `f(t :: Trace, params...)` a stochastic function
`q` with signature `q(old_trace :: Trace, new_trace :: Trace, params...)`, generates a proposal from `q` and
accepts based on the log acceptance probability:

``
\log \alpha = \log p(t_{\text{new}}) - \log q(t_{\text{new}}|t_{\text{old}}) - [\log p(t_{\text{old}}) - \log q(t_{\text{old}} | t_{\text{new}})].
``
"""
function mh_step(t :: Trace, f, q; params = (), return_val :: Bool = false)
    new_t = trace()
    f(new_t, params...)
    copy_common!(t, new_t)
    logprob!(t)
    q(t, new_t, params...) # propose into the trace
    new_t, g = replay(f, new_t)
    r = g(params...)  # new joint
    logprob!(new_t)
    log_q_new_given_old = logprob(t, new_t)
    log_q_old_given_new = logprob(new_t, t)
    log_a = new_t.logprob_sum + log_q_old_given_new - t.logprob_sum - log_q_new_given_old
    a = accept(t, new_t, log_a)
    if !return_val
        a
    else
        (r, a)
    end
end

@doc raw"""
    function mh_step(t :: Trace, f, q, types::Tuple{DataType,DataType}; params = (), return_val :: Bool = false)

A generic Metropolis step using an arbitrary proposal kernel.

Given a trace `t`, a stochastic function `f` with signature `f(t :: Trace, params...)` a stochastic function
`q` with signature `q(old_trace :: Trace, new_trace :: Trace, params...)`, generates a proposal from `q` and
accepts based on the log acceptance probability:

``
\log \alpha = \log p(t_{\text{new}}) - \log q(t_{\text{new}}|t_{\text{old}}) - [\log p(t_{\text{old}}) - \log q(t_{\text{old}} | t_{\text{new}})].
``
"""
function mh_step(t :: Trace, f, q, types::Tuple{DataType,DataType}; params = (), return_val :: Bool = false)
    new_t = trace(types[1], types[2])
    f(new_t, params...)
    copy_common!(t, new_t)
    logprob!(t)
    q(t, new_t, params...) # propose into the trace
    new_t, g = replay(f, new_t)
    r = g(params...)  # new joint
    logprob!(new_t)
    log_q_new_given_old = logprob(t, new_t)
    log_q_old_given_new = logprob(new_t, t)
    log_a = new_t.logprob_sum + log_q_old_given_new - t.logprob_sum - log_q_new_given_old
    a = accept(t, new_t, log_a)
    if !return_val
        a
    else
        (r, a)
    end
end

@doc raw"""
    function mh(f, qs :: A; params = (), burn = 100, thin = 10, num_iterations = 10000, inverse_verbosity = 100) where A <: AbstractArray

Generic Metropolis algorithm using user-defined proposal kernels.

Args:

+ `f`: stochastic function. Must have call signature `f(t :: Trace, params...)`
+ `qs`: array-like of proposal kernels. Proposal kernels are applied sequentially in the order that they appear in this array.
    Proposal kernels must have the signature `q(old_t :: Trace, new_t :: Trace, params...)` where it must take in at least the same number of arguments
    in `params` as `f`.
+ `params`: addditional arguments to pass to `f` and each of the proposal kernels.
+ `burn`: number of samples to discard at beginning of markov chain
+ `thin`: keep only every `thin`-th draw. E.g., if `thin = 100`, only every 100-th trace will be kept.
+ `num_iterations`: total number of steps to take in the markov chain
+ `inverse_verbosity`: every `inverse_verbosity` iterations, a stattus report will be logged.
"""
function mh(f, qs :: A; params = (), burn = 100, thin = 10, num_iterations = 10000, inverse_verbosity = 100) where A <: AbstractArray
    results = metropolis_results()
    t = trace()
    f(t, params...)
    local r
    for n in 1:num_iterations
        for q in qs
            (r, t) = mh_step(t, f, q; params = params, return_val = true)
        end
        if (n > burn) && (n % thin == 0)
            push!(results.return_values, r)
            push!(results.traces, t)
            push!(results.log_weights, logprob(t))
        end
        if n % inverse_verbosity == 0
            @info "On iteration $n of $num_iterations"
        end
    end
    results
end

@doc raw"""
    function mh(f, qs :: A, addresses; params = (), burn = 100, thin = 10, num_iterations = 10000, inverse_verbosity = 100) where A <: AbstractArray

Generic Metropolis algorithm using user-defined proposal kernels, returning only a requested subset of
addresses.

Args:

+ `f`: stochastic function. Must have call signature `f(t :: Trace, params...)`
+ `qs`: array-like of proposal kernels. Proposal kernels are applied sequentially in the order that they appear in this array.
    Proposal kernels must have the signature `q(old_t :: Trace, new_t :: Trace, params...)` where it must take in at least the same number of arguments
    in `params` as `f`.
+ `addresses`: *only* values sampled at these addresses will be saved in the `values` field of the
    `BareResults` struct returned.
+ `params`: addditional arguments to pass to `f` and each of the proposal kernels.
+ `burn`: number of samples to discard at beginning of markov chain
+ `thin`: keep only every `thin`-th draw. E.g., if `thin = 100`, only every 100-th trace will be kept.
+ `num_iterations`: total number of steps to take in the markov chain
+ `inverse_verbosity`: every `inverse_verbosity` iterations, a stattus report will be logged.
"""
function mh(f, qs :: A, addresses; params = (), burn = 100, thin = 10, num_iterations = 10000, inverse_verbosity = 100) where A <: AbstractArray
    results = bare_metropolis_results()
    t = trace()
    f(t, params...)
    local r
    for n in 1:num_iterations
        for q in qs
            (r, t) = mh_step(t, f, q; params = params, return_val = true)
        end
        if (n > burn) && (n % thin == 0)
            for a in addresses
                v = a in keys(t) ? t[a].value : nothing
                push!(results.values[a], v)
            end
        end
        if n % inverse_verbosity == 0
            @info "On iteration $n of $num_iterations"
        end
    end
    results
end

@doc raw"""
    function mh(f, qs :: A, types::Tuple{DataType,DataType}; params = (), burn = 100, thin = 10, num_iterations = 10000, inverse_verbosity = 100) where A <: AbstractArray

Generic Metropolis algorithm using user-defined proposal kernels.

Args:

+ `f`: stochastic function. Must have call signature `f(t :: Trace, params...)`
+ `qs`: array-like of proposal kernels. Proposal kernels are applied sequentially in the order that they appear in this array.
    Proposal kernels must have the signature `q(old_t :: Trace, new_t :: Trace, params...)` where it must take in at least the same number of arguments
    in `params` as `f`.
+ `params`: addditional arguments to pass to `f` and each of the proposal kernels.
+ `burn`: number of samples to discard at beginning of markov chain
+ `thin`: keep only every `thin`-th draw. E.g., if `thin = 100`, only every 100-th trace will be kept.
+ `num_iterations`: total number of steps to take in the markov chain
+ `inverse_verbosity`: every `inverse_verbosity` iterations, a stattus report will be logged.
"""
function mh(f, qs :: A, types::Tuple{DataType,DataType}; params = (), burn = 100, thin = 10, num_iterations = 10000, inverse_verbosity = 100) where A <: AbstractArray
    results = metropolis_results(types...)
    t = trace(types...)
    f(t, params...)
    local r
    for n in 1:num_iterations
        for q in qs
            (r, t) = mh_step(t, f, q, types; params = params, return_val = true)
        end
        if (n > burn) && (n % thin == 0)
            push!(results.return_values, r)
            push!(results.traces, t)
            push!(results.log_weights, logprob(t))
        end
        if n % inverse_verbosity == 0
            @info "On iteration $n of $num_iterations"
        end
    end
    results
end

### q is a proposal kernel

@doc raw"""
    function mh_step(rf::RandomField, q, x, log_prob_rf_x)

A Metropolis step to sample from a random field using an arbitrary proposal kernel.

Args:

+ `rf`: a `RandomField` from which to sample
+ `q`: a proposal kernel. This must be a callable that satisfies the following
    requirements:

    + `q(x)` returns a new value `x_prime` that is generated using the input parameters
        `x`, i.e., ``x' \sim q(x'|x)``. For example, `q(x) = rand(Normal(x, 0.25))`
    + `q(x_prime, x)` scores (computes the log probability of) `x_prime` against `x`, i.e.,
        computes ``\log q(x' | x)``

+ `x`: a dict with format `address => value`, the current sampled value.
+ `log_prob_rf_x`: the log probability of `x` under the random field.
"""
function mh_step(rf::RandomField, q, x, log_prob_rf_x)
    x_prime_given_x = q(x)
    log_prob_x_prime_given_x = q(x_prime_given_x, x)
    log_prob_rf_x_prime = rf(x_prime_given_x)
    log_prob_x_given_x_prime = q(x, x_prime_given_x)
    log_a = log_prob_rf_x_prime + log_prob_x_given_x_prime - log_prob_rf_x - log_prob_x_prime_given_x
    accept(x, log_prob_rf_x, x_prime_given_x, log_prob_rf_x_prime, log_a)
end

@doc raw"""
    function mh(rf::RandomField, qs::Vector{T}, val; burn=1000, thin=100, num_iterations=11000) where T

Metropolis Hastings algorithm for sampling from a random field using arbitrary proposal kernels.

Args:

+ `rf`: the random field from which to sample
+ `qs`: vector of proposal kernel callables; see documentation of `mh_step` for specification of proposal kernels
+ `val`: initial guess with which to initialize MH, must be a dict with format `address => value`.
+ `burn`: number of samples to discard at beginning of markov chain
+ `thin`: keep only every `thin`-th draw. E.g., if `thin = 100`, only every 100-th trace will be kept.
+ `num_iterations`: total number of steps to take in the markov chain
"""
function mh(rf::RandomField, qs::Vector{T}, val; burn=1000, thin=100, num_iterations=11000) where T
    log_prob = rf(val)
    results = structured_metropolis_results()
    for name in rf.names
        results.values[name] = []
    end
    for n in 1:num_iterations
        for q in qs
            (val, log_prob) = mh_step(rf, q, val, log_prob)
        end
        (n > burn) && (n % thin == 0) && begin
            for (k, v) in val
                push!(results.values[k], v)
            end
            push!(results.log_weights, log_prob)
        end
    end
    results
end

export Metropolis
export propose, is_symmetric
export mh_step, mh, metropolis_results


text:
# ---
# title: 713. Subarray Product Less Than K
# id: problem713
# author: Indigo
# date: 2021-06-29
# difficulty: Medium
# categories: Array, Two Pointers
# link: <https://leetcode.com/problems/subarray-product-less-than-k/description/>
# hidden: true
# ---
#
# Your are given an array of positive integers `nums`.
#
# Count and print the number of (contiguous) subarrays where the product of all
# the elements in the subarray is less than `k`.
#
# **Example 1:**
#
#
#
#     Input: nums = [10, 5, 2, 6], k = 100
#     Output: 8
#     Explanation: The 8 subarrays that have product less than 100 are: [10], [5], [2], [6], [10, 5], [5, 2], [2, 6], [5, 2, 6].
#     Note that [10, 5, 2] is not included as the product of 100 is not strictly less than k.
#
#
# **Note:**
#
# * `0 < nums.length <= 50000`.
# * `0 < nums[i] < 1000`.
# * `0 <= k < 10^6`.
#
#
## @lc code=start
using LeetCode

function num_subarray_productLess_than_k(nums::Vector{Int}, k::Int)
    k <= 1 && return 0
    res, left = 0, 1
    p = 1
    for (right, val) in enumerate(nums)
        p *= val
        while p >= k
            p ÷= nums[left]
            left += 1
        end
        res += right - left + 1
    end
    return res
end
## @lc code=end


text:
# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.10.3
#   kernelspec:
#     display_name: Julia 1.7.0
#     language: julia
#     name: julia-1.7
# ---

# %% [markdown]
# https://twitter.com/iitenki_moruten/status/1467683477474930688

# %%
# Original: https://github.com/moruten/julia-code/blob/1dafca3e1a4e3b36445c2263e440f6e4056b90aa/2021-12-6-test-no1.ipynb

using Plots
using Random
using Distributions
using QuadGK

#============関数定義============================================#
function action(x)
    0.5*x*x
end

function deriv_action(x)
    x
end

function hamiltonian(x,p)
    action(x) + 0.5*p*p
end

function HMC_update(x,Nt,dt)
    #backup
    x_old = x
    p     = rand(Normal(0,1))

    #check
    H_ini = hamiltonian(x,p)
    x     = molecular_dynamics!(x,p,Nt,dt)
    H_fin = hamiltonian(x,p)

    r  = rand()
    ΔH = H_fin-H_ini
    if r < exp(-ΔH)
        return x,1 #accept
    else
        return x_old,0
    end
end

function molecular_dynamics!(x,p,Nt,dt)
    force = 0.0
    #1/2step
        x = x + p * 0.5*dt
    #1~Nt-1 step
        for j=1:(Nt-1)
            force = deriv_action(x)
            p = p - dt * force
            x = x + dt * p
        end
    #Nt step
        p = p - dt * deriv_action(x)
        x = x + 0.5 * dt * p
        return x
end
#============関数終わり============================================#

#============計算=========================================================================#
  #セットアップ
    Ntest        = 300000
    Nt           = 20
    dt           = 1.0/Nt
    conf_vec     = zeros(Ntest)
    accept_count = 0
    ret          = 0
    x            = 0.0

    sumxx        = 0.0
    sumx         = 0.0

  #計算
    for i=1:Ntest
        x,ret         = HMC_update(x,Nt,dt)
        accept_count += ret
        conf_vec[i]   = x

        sumx         += x
        sumxx        += x*x
    end

    println("P(accept) = $(accept_count/Ntest)")
    println("<x>       = $(sumx/Ntest)")
    println("<xx>      = $(sumxx/Ntest)")

#=======確認=============================================================================#
    xr = range(-5,5,length=1000)
    f1(x) = exp(-0.5*x^2)
    f2(x) = exp(-x^2)
    Z1,error1 = quadgk(f1,-Inf,Inf)
    Z2,error2 = quadgk(f2,-Inf,Inf)
    g1(x) = f1(x)/Z1
    g2(x) = f2(x)/Z2
    histogram(conf_vec,norm=:true,label="data")
    plot!(xr,[g1.(xr) g2.(xr)],lw=3,label=["exp(-0.5x^2)/Z1" "exp(-x^2)/Z2"])
#========================================================================================#

# %%
# Revised 1

using Plots
using Random
using Distributions
using QuadGK

#============関数定義============================================#
function action(x)
    0.5*x*x
end

function deriv_action(x)
    x
end

function hamiltonian(x,p)
    action(x) + 0.5*p*p
end

function HMC_update(x,Nt,dt)
    #backup
    x_old = x
    p     = rand(Normal(0,1))

    #check
    H_ini = hamiltonian(x,p)
    x, p = molecular_dynamics!(x,p,Nt,dt)    # <========== ここ
    H_fin = hamiltonian(x,p)

    r  = rand()
    ΔH = H_fin-H_ini
    if r < exp(-ΔH)
        return x,1 #accept
    else
        return x_old,0
    end
end

function molecular_dynamics!(x,p,Nt,dt)
    force = 0.0
    #1/2step
        x = x + p * 0.5*dt
    #1~Nt-1 step
        for j=1:(Nt-1)
            force = deriv_action(x)
            p = p - dt * force
            x = x + dt * p
        end
    #Nt step
        p = p - dt * deriv_action(x)
        x = x + 0.5 * dt * p
        return x, p                          # <========== ここ
end
#============関数終わり============================================#

#============計算=========================================================================#
  #セットアップ
    Ntest        = 300000
    Nt           = 20
    dt           = 1.0/Nt
    conf_vec     = zeros(Ntest)
    accept_count = 0
    ret          = 0
    x            = 0.0

    sumxx        = 0.0
    sumx         = 0.0

  #計算
    for i=1:Ntest
        x,ret         = HMC_update(x,Nt,dt)
        accept_count += ret
        conf_vec[i]   = x

        sumx         += x
        sumxx        += x*x
    end

    println("P(accept) = $(accept_count/Ntest)")
    println("<x>       = $(sumx/Ntest)")
    println("<xx>      = $(sumxx/Ntest)")

#=======確認=============================================================================#
    xr = range(-5,5,length=1000)
    f1(x) = exp(-0.5*x^2)
    f2(x) = exp(-x^2)
    Z1,error1 = quadgk(f1,-Inf,Inf)
    Z2,error2 = quadgk(f2,-Inf,Inf)
    g1(x) = f1(x)/Z1
    g2(x) = f2(x)/Z2
    histogram(conf_vec,norm=:true,label="data")
    plot!(xr,[g1.(xr) g2.(xr)],lw=3,label=["exp(-0.5x^2)/Z1" "exp(-x^2)/Z2"])
#========================================================================================#

# %%
# Revised 1 - second test

using Plots
using Random
using Distributions
using QuadGK

#============関数定義============================================#
function action(x)
    3(x^2 - 1)^2
end

function deriv_action(x)
    6x*(x^2 - 1)
end

function hamiltonian(x,p)
    action(x) + 0.5*p*p
end

function HMC_update(x,Nt,dt)
    #backup
    x_old = x
    p     = rand(Normal(0,1))

    #check
    H_ini = hamiltonian(x,p)
    x, p = molecular_dynamics!(x,p,Nt,dt)     # <========== ここ
    H_fin = hamiltonian(x,p)

    r  = rand()
    ΔH = H_fin-H_ini
    if r < exp(-ΔH)
        return x,1 #accept
    else
        return x_old,0
    end
end

function molecular_dynamics!(x,p,Nt,dt)
    force = 0.0
    #1/2step
        x = x + p * 0.5*dt
    #1~Nt-1 step
        for j=1:(Nt-1)
            force = deriv_action(x)
            p = p - dt * force
            x = x + dt * p
        end
    #Nt step
        p = p - dt * deriv_action(x)
        x = x + 0.5 * dt * p
        return x, p                    # <========== ここ
end
#============関数終わり============================================#

#============計算=========================================================================#
  #セットアップ
    Ntest        = 300000
    Nt           = 20
    dt           = 1.0/Nt
    conf_vec     = zeros(Ntest)
    accept_count = 0
    ret          = 0
    x            = 0.0

    sumxx        = 0.0
    sumx         = 0.0

  #計算
    for i=1:Ntest
        x,ret         = HMC_update(x,Nt,dt)
        accept_count += ret
        conf_vec[i]   = x

        sumx         += x
        sumxx        += x*x
    end

    println("P(accept) = $(accept_count/Ntest)")
    println("<x>       = $(sumx/Ntest)")
    println("<xx>      = $(sumxx/Ntest)")

#=======確認=============================================================================#
    xr = range(-2,2,length=1000)
    f(x) = exp(-action(x))
    Z,error1 = quadgk(f,-Inf,Inf)
    g(x) = f(x)/Z
    histogram(conf_vec,norm=:true,label="data")
    plot!(xr,g.(xr),lw=3,label="exp(-action(x))/Z", legend=:outertop)
#========================================================================================#

# %%
# Revised 2

using Plots
using Random
using Distributions
using QuadGK

#============関数定義============================================#
function action(x)
    x^2/2
end

function deriv_action(x)
    x
end

function hamiltonian(x,p)
    action(x) + 0.5*p*p
end

function HMC_update(x,Nt,dt)
    #backup
    x_old = x
    p     = rand(Normal(0,1))

    #check
    H_ini = hamiltonian(x,p)
    x, p = molecular_dynamics!(x,p,Nt,dt)     # <========== ここ
    H_fin = hamiltonian(x,p)

    r  = rand()
    ΔH = H_fin-H_ini
    if r < exp(-ΔH)
        return x,1 #accept
    else
        return x_old,0
    end
end

function molecular_dynamics!(x,p,Nt,dt)
    p -= deriv_action(x) * dt/2
    x += p * dt
    for j in 2:Nt
        p -= deriv_action(x) * dt
        x += p * dt
    end
    p -= deriv_action(x) * dt/2
    return x, p                    # <========== ここ
end
#============関数終わり============================================#

#============計算=========================================================================#
  #セットアップ
    Ntest        = 300000
    Nt           = 20
    dt           = 1.0/Nt
    conf_vec     = zeros(Ntest)
    accept_count = 0
    ret          = 0
    x            = 0.0

    sumxx        = 0.0
    sumx         = 0.0

  #計算
    for i=1:Ntest
        x,ret         = HMC_update(x,Nt,dt)
        accept_count += ret
        conf_vec[i]   = x

        sumx         += x
        sumxx        += x*x
    end

    println("P(accept) = $(accept_count/Ntest)")
    println("<x>       = $(sumx/Ntest)")
    println("<xx>      = $(sumxx/Ntest)")

#=======確認=============================================================================#
    xr = range(-5,5,length=1000)
    f(x) = exp(-action(x))
    Z,error1 = quadgk(f,-Inf,Inf)
    g(x) = f(x)/Z
    histogram(conf_vec,norm=:true,label="data",alpha=0.3,bin=100)
    plot!(xr,g.(xr),lw=3,label="exp(-action(x))/Z", legend=:outertop)
#========================================================================================#

# %%
# Revised 2 - second test

using Plots
using Random
using Distributions
using QuadGK

#============関数定義============================================#
function action(x)
    3(x^2 - 1)^2
end

function deriv_action(x)
    6x*(x^2 - 1)
end

function hamiltonian(x,p)
    action(x) + 0.5*p*p
end

function HMC_update(x,Nt,dt)
    #backup
    x_old = x
    p     = rand(Normal(0,1))

    #check
    H_ini = hamiltonian(x,p)
    x, p = molecular_dynamics!(x,p,Nt,dt)     # <========== ここ
    H_fin = hamiltonian(x,p)

    r  = rand()
    ΔH = H_fin-H_ini
    if r < exp(-ΔH)
        return x,1 #accept
    else
        return x_old,0
    end
end

function molecular_dynamics!(x,p,Nt,dt)
    p -= deriv_action(x) * dt/2
    x += p * dt
    for j in 2:Nt
        p -= deriv_action(x) * dt
        x += p * dt
    end
    p -= deriv_action(x) * dt/2
    return x, p                    # <========== ここ
end
#============関数終わり============================================#

#============計算=========================================================================#
  #セットアップ
    Ntest        = 300000
    Nt           = 20
    dt           = 1.0/Nt
    conf_vec     = zeros(Ntest)
    accept_count = 0
    ret          = 0
    x            = 0.0

    sumxx        = 0.0
    sumx         = 0.0

  #計算
    for i=1:Ntest
        x,ret         = HMC_update(x,Nt,dt)
        accept_count += ret
        conf_vec[i]   = x

        sumx         += x
        sumxx        += x*x
    end

    println("P(accept) = $(accept_count/Ntest)")
    println("<x>       = $(sumx/Ntest)")
    println("<xx>      = $(sumxx/Ntest)")

#=======確認=============================================================================#
    xr = range(-2,2,length=1000)
    f(x) = exp(-action(x))
    Z,error1 = quadgk(f,-Inf,Inf)
    g(x) = f(x)/Z
    histogram(conf_vec,norm=:true,label="data",alpha=0.3,bin=100)
    plot!(xr,g.(xr),lw=3,label="exp(-action(x))/Z", legend=:outertop)
#========================================================================================#

# %%

max_stars_repo_path: 0025/HMC revised.jl

max_stars_repo_name: genkuroki/public

max_issues_repo_path: 0025/HMC revised.jl

max_issues_repo_name: genkuroki/public

max_forks_repo_path: 0025/HMC revised.jl

max_forks_repo_name: genkuroki/public


text:
using LinearAlgebra

function build_projection_matrix(vs::Lar.Points; y0 = false)::Matrix{Float64}

#    u1 = LinearAlgebra.normalize(vs[:, 2] - vs[:, 1])
    u1 = vs[:, 2] - vs[:, 1]
    i = 2 + argmax(map(i -> norm(Lar.cross(u1, vs[:, i] - vs[:, 1])), 3 : size(vs, 2)))
#    u2 = LinearAlgebra.normalize(vs[:, i] - vs[:, 1])
#    u3 = LinearAlgebra.normalize(Lar.cross(u1, u2))
#    if y0  u2 = LinearAlgebra.normalize(Lar.cross(u1, u3))  end
    u2 = vs[:, i] - vs[:, 1]
    u3 = Lar.cross(u1, u2)
    if y0  u2 = Lar.cross(u1, u3)  end
    T = Matrix{Float64}(LinearAlgebra.I, 4, 4)
    T[1:3, 4] = - vs[:,1]
    M = Matrix{Float64}(LinearAlgebra.I, 4, 4)
    M[1:3, 1:3] = [u1 u2 u3]
    return M' * T
end

"""
    build_rototranslation_matrix(vs::Lar.Points)::Matrix{Float64}

Generates the rototranslation matrix that maps `vs[1,2]` -> ([0 0 0; x 0 0]')
"""
function build_rototranslation_matrix(vs::Lar.Points)::Matrix{Float64}

    Ivs = [vs; ones(1, size(vs, 2))]
    T   = Matrix{Float64}(LinearAlgebra.I, 4, 4)
    Rx  = Matrix{Float64}(LinearAlgebra.I, 4, 4)
    Ry  = Matrix{Float64}(LinearAlgebra.I, 4, 4)

    # Translate first point to Origin
    T[1:3, 4] = - vs[:,1]
    Tvs = T * Ivs

    # Rotate along x-axis to xz-plan the 1->2 edge (y = 0)
    #   |  1   0   0  |       sin(θ) = Py (= θ[1])
    #   |  0   c  -s  |   =>
    #   |  0   s   c  |       cos(θ) = Pz (= θ[2])
    if !isapprox(norm(Tvs[2:3, 2]), 0.0, atol = 1e-10)
        θ = LinearAlgebra.normalize(Tvs[2:3, 2])
        Rx[2, 2] =  θ[2]
        Rx[3, 3] =  θ[2]
        Rx[2, 3] = -θ[1]
        Rx[3, 2] =  θ[1]
    end
    RTvs = (Rx * T) * Ivs

    # Rotate along y-axis to x-axis the 1->2 edge (y = 0 && z = 0)
    #   |  c   0   s  |       sin(σ) = Pz (= σ[2])
    #   |  0   1   0  |   =>
    #   | -s   0   c  |       cos(σ) = Px (= σ[1])
    if !isapprox(norm(RTvs[[1, 3], 2]), 0.0, atol = 1e-10)
        σ = LinearAlgebra.normalize(RTvs[[1, 3], 2])
        Ry[1, 1] =  σ[1]
        Ry[3, 3] =  σ[1]
        Ry[1, 3] =  σ[2]
        Ry[3, 1] = -σ[2]
    end

    return Ry * Rx * T
end

"""
    eval_ord_angle(model::CAGD.Model, dim::Int, lo_cell::Int)::Array{Int,1}

Evaluates the `dim`-cell circular ordering w.r.t the `dim-1` cell `lo_cell`.
"""
function eval_ord_angle(model, dim, lo_cell; atol = 1e-7)
    if dim == 1
        return CAGD.eval_ord_edges(model, lo_cell, atol = atol)
    elseif dim == 2
        return CAGD.eval_ord_faces(model, lo_cell, atol = atol)
    end
    throw(ArgumentError("Dimension $dim not coded for ordering"))
end

function eval_ord_faces(model, τ; atol = 1e-7)
    proj_matrix = false
    # Get EV array of arrays representation for Lar.pointInPolygonClassification
    EV = Lar.cop2lar(model.T[1])
    # Get faces of τ's petals
    faces = model.T[2][:, τ].nzind
    # Get vertices idx of τ
    edge = model.T[1][τ, :].nzind
    # Get other vertices idx of petals (but τ's)
    FVs = [setdiff((abs.(model.T[1])'*abs.(model.T[2][f, :])).nzind, edge) for f in faces]
    # Get geometry of the first petal (with τ's vertices first)
    Gface = model.G[:, [edge..., FVs[1]...]]
    # Get transformation matrix that maps τ to (0,0,0)->(x,0,0)
    if proj_matrix
        M = CAGD.build_projection_matrix(Gface, y0 = true)
    else
        M = CAGD.build_rototranslation_matrix(Gface)
    end
    # Project the other points such that each face is normal to x axis
    PG = M[1:3, :] * [model.G; ones(1, size(model, 0, 2))]
    # Each face describe in yz a straight line.
    # Faces are sorted by the `angles` their xy-ray form w.r.t. the y = 0 positive ray
    angles = zeros(length(faces))
    for fidx = 1 : length(faces)
        # If the face lies on xz plan, then the angle of the projection on yz is either -π/2 or +π/2
        if sum(1 .- isapprox.(PG[2, FVs[fidx]], 0.0, atol = atol)) == 0
            # Look for a non trivial point (z coord non null)
            for v in FVs[fidx]  if !isapprox(PG[3, v], 0.0, atol = atol)
                # Check whether face is on the z-halfplane described by `v`'s ray or not
                # `side_pt` is located "near" the middle of `τ` on the halfplane of `v`
                side_pt = [PG[1, edge[2]]/2 sign(PG[3, v])*1e-10]
                # Take the edges of the face as array of arrays for pointInPolygonClassification
                EVs = EV[model.T[2][faces[fidx], :].nzind]
                # Check wheter v is on the other side or the same halfplane
                other_side=Lar.pointInPolygonClassification(PG[[1,3],:], EVs)(side_pt)=="p_out"
                # the 90° angle is either on the side of `v` or not
                angles[fidx] = (-1)^other_side * sign(PG[3, v]) * π / 2
                break
            end  end
        else
            # If not so, the face is a non vertical segment if projected on yz plan
            # Search for a consisten `v`, that is y coord not null
            for v in FVs[fidx]  if !isapprox(PG[2, v], 0.0, atol = atol)
                # Check whether the face is on the halfplane described by `v`'s ray or not
                #  to do so, consider the xy projection of the face and check wether
                #  `side_pt` (located "near" the middle of `τ` on the halfplane of `v`)
                #  is in the face projection
                side_pt = [PG[1, edge[2]]/2 PG[2, edge[2]]+sign(PG[2, v])*1e-10]
                # Take the edges of the face as array of arrays for pointInPolygonClassification
                EVs = EV[model.T[2][faces[fidx], :].nzind]
                # Check wheter v is on the other side or the same halfplane
                other_side=Lar.pointInPolygonClassification(PG[[1,2],:], EVs)(side_pt)=="p_out"
                # Copute the angle of `v`'s ray, reoriented according the face halfplane
                angles[fidx] = atan((-1)^other_side * PG[3, v], (-1)^other_side * PG[2, v])
                break
            end  end
        end
    end

    # faces are sorted according to right-hand rule
    if proj_matrix angles = -angles end
    ord_faces = sortperm(angles)

    return faces[ord_faces]
end

function eval_ord_edges(model, τ; G = model.G, atol = 1e-7)
    edges = model.T[1][:, τ].nzind
    EVs = [setdiff((model.T[1][e, :]).nzind, τ)[1] for e in edges]
    angles = [
        [edges[idx] atan(G[2, EVs[idx]], G[1, EVs[idx]])]
        for idx = 1 : length(edges)
        if !isapprox(G[:, EVs[idx]], [0.,0.], atol = atol)
    ]
    sort!(angles, lt = (a, b) -> a[2] < b[2])
    return map(a->Int(a[1]), angles)
end

"""
    pointInSolid(
        model::CAGD.Model,
        point::Array{Float64,1},
        cycle::Int,
        signed = false
    )::Int8

Computes if `point` is internal (`1`), external (`-1`) or onto (`0`) `cycle` solid

Checks the incidences of the `x`-ray from `point` with `model.T[3][cycle,:]`
boundary faces to determine wether it is internal, external or onto it.

By default it consider the cycle to be the internal boundary of a cell.
If `signed` is set `true` then it also consider the orientation of the cycle
(_i.e_ if the cycle is the external boundary). The model needs being oriented.
"""
function pointInSolid(model, point, cycle, signed = false; atol = 1e-8)

    function doBBoxIntersect(bbox) #dimension free
        flag = point[1] <= bbox[1, 2]
        for d = 2 : length(point)
            flag &= bbox[d, 1] <= point[d] <= bbox[d, 2]
        end
        return flag
    end

    if signed
        throw(ArgumentError("signed not coded yet"))
    end
    # Vertices of `cycle`
    vs = CAGD.getModelCellVertices(model, 3, cycle)
    # Check if `point` is a vertex
    for vert in vs  if isapprox(vs, vert)  return 0  end  end
    # Reshaping vertices
    vs = hcat(vs...)
    # Bounding box of vs
    minimum = mapslices(x->min(x...), vs, dims=2)
    maximum = mapslices(x->max(x...), vs, dims=2)
    # x-line vertices (outside of bbox)
    vmin = point
    # adding a little delta to avoid coplanarity
    vmax = [maximum[1]+1; point[2] + 2*atol; point[3] + 2*atol]

    # initialization of counter
    counter = 0
    face_cycles = CAGD.getModelLoCell(model, 3, cycle)
    face_bboxes = CAGD.getModelBoundingBoxes(model, 2, face_cycles, atol = atol*1e-2)
    for face_idx = 1 : length(face_cycles)  if doBBoxIntersect(face_bboxes[face_idx])
        # To check incidence on each face, project it to z=0
        Gface, Gfaceidx = CAGD.getModelCellGeometry(model, 2, face_idx, true)
        # Build the projection matrix
        M = CAGD.build_projection_matrix(Gface)
        # Project the points
        PG = (M[1:3, :] * [model.G; ones(1, size(model, 0, 2))])
        Pvmin = M[1:3, :] * [vmin; 1]
        Pvmax = M[1:3, :] * [vmax; 1]

        if isapprox(Pvmin[3], 0.0, atol = atol) && isapprox(Pvmax[3], 0.0, atol = atol)
            # if point is collinear with face, then it is either on face or not.
            # If not, then the face has even contribute and can be avoided
            if CAGD.point_in_face(Pvmin[1:2], PG[1:2, :], model.T[1])
                return 0
            end
        else
            # If point is not collinear, retrieve the x-ray point on face plan
            Pvplan = Pvmin[1:2] + (Pvmax[1:2] - Pvmin[1:2]).*((Pvmin[3])/(Pvmax[3]-Pvmin[3]))
            # If it is in, then it adds one to the counter
            if CAGD.point_in_face(Pvplan, PG[1:2, :], model.T[1])
                # check it is really in and not on the border
                @assert Lar.point_in_face(Pvplan, PG[1:2, :], model.T[1])
                counter += 1
            end
        end
    end  end
    return counter % 2 == 1
end

"""
    getExternalBoundary(model::CAGD.Model, deg::Int, cycles::Lar.ChainOp, comp::Array{Int,1})::Int

Evaluates the external cycle of the component described by `deg`-`cycles` `comp`

Consider a component made of the `comp` cycles among `deg`-cells and
evaluates wich of those is the external one (in `comp` ordering).

#TODO improve the function checking if the longest cycle is the outer
"""
function getExternalBoundary(model, deg, cycles, comp)
    return findmax([length(cycles[:, comp[i]].nzind) for i = 1 : length(comp)])[2]
end


text:
module splines

export filter_coeffs, interpolant_cspline, filter_coeffs
export eval_UC_spline, eval_UC_spline!
export prefilter!

include("csplines.jl")
include("splines_filter.jl")
include("cubic_prefilter.jl")

function interpolant_cspline(a, b, orders, V)

    coefs = filter_coeffs(a, b, orders, V)

    function fun(s::Array{Float64,2})
        return eval_UC_spline(a, b, orders, coefs, s)
    end

    function fun(p::Float64...)
        return fun([p...]')
    end

    return fun

end

end


text:
function EvaluateCℓMCMCStep(Cosmology::AbstractCosmology, Density::AbstractConvolvedDensity,
    Bias::AbstractBias, IA::AbstractIntrinsicAlignment, CosmoGrid::AbstractCosmologicalGrid)
    backgroundquantities = BackgroundQuantities(HZArray= zeros(length(CosmoGrid.ZArray)),
    rZArray=zeros(length(CosmoGrid.ZArray)))
    ComputeBackgroundQuantitiesGrid!(CosmoGrid, backgroundquantities, Cosmology)
    ClassyParams = Initializeclassy(Cosmology)
    Pmm = PowerSpectrum(PowerSpectrumLinArray =
    zeros(length(CosmoGrid.KArray), length(CosmoGrid.ZArray)),
    PowerSpectrumNonlinArray = zeros(length(CosmoGrid.KArray), length(CosmoGrid.ZArray)),
    InterpolatedPowerSpectrum = zeros(length(CosmoGrid.ℓBinCenters),
    length(CosmoGrid.ZArray)))
    EvaluatePowerSpectrum!(ClassyParams, CosmoGrid, Pmm)
    ComputeLimberArray!(CosmoGrid, backgroundquantities)
    InterpolatePowerSpectrumLimberGrid!(CosmoGrid, backgroundquantities, Pmm,
    BSplineCubic())
    ExtractGrowthFactor!(backgroundquantities, Pmm)
    GCW = GCWeightFunction(WeightFunctionArray =
    zeros(length(Density.DensityNormalizationArray), length(CosmoGrid.ZArray)),
    BiasKind = Bias)
    CosmoCentral.ComputeBiasGrid!(CosmoGrid, GCW, Density)
    CosmoCentral.ComputeWeightFunctionGrid!(GCW, Density, CosmoGrid,
    backgroundquantities, Cosmology)
    WLW = WLWeightFunction(WeightFunctionArray =
    zeros(length(Density.DensityNormalizationArray), length(CosmoGrid.ZArray)),
    LensingEfficiencyArray = zeros(length(Density.DensityNormalizationArray),
    length(CosmoGrid.ZArray)), IntrinsicAlignmentModel = IA)
    ComputeLensingEfficiencyGrid!(WLW, Density, CosmoGrid,
    backgroundquantities, Cosmology, CustomLensingEfficiency())
    ComputeIntrinsicAlignmentGrid!(CosmoGrid, WLW, Density,
    backgroundquantities, Cosmology)
    ComputeWeightFunctionGrid!(WLW, Density, CosmoGrid, backgroundquantities,
    Cosmology)
    CℓLL = Cℓ(CℓArray = zeros(length(CosmoGrid.ℓBinCenters),
    length(WLW.WeightFunctionArray[:, 1]),
    length(WLW.WeightFunctionArray[:, 1])))
    ComputeCℓ!(CℓLL, WLW, WLW, backgroundquantities, Cosmology,
    CosmoGrid, Pmm, CustomSimpson())
    FudgeFactor = 0.9919859719438879
    CℓLL.CℓArray ./= FudgeFactor
    return CℓLL
end


text:
module QuasiMonteCarlo

using Sobol, LatinHypercubeSampling, LatticeRules, Distributions

abstract type SamplingAlgorithm end

"""
```julia
GridSample{T}
```

The grid is given by `lb:dx[i]:ub` in the ith direction.
"""
struct GridSample{T} <: SamplingAlgorithm
    dx::T
end

"""
```julia
struct UniformSample <: SamplingAlgorithm end
```

Uniformly distributed random numbers.
"""
struct UniformSample <: SamplingAlgorithm end

"""
```julia
struct SobolSample <: SamplingAlgorithm end
```

Samples using the Sobol sequence using Sobol.jl.
"""
struct SobolSample <: SamplingAlgorithm end

"""
```julia
struct LatinHypercubeSample <: SamplingAlgorithm end
```

Samples using a Latin Hypercube using LatinHypercubeSampling.jl

```julia
LatinHypercubeSample(threading=false)
```

Keyword arguments:

- `threading`: whether to use threading. Default is false, i.e. serial.
"""
Base.@kwdef struct LatinHypercubeSample{T} <: SamplingAlgorithm
    threading::T = false
end

"""
```julia
struct LatticeRuleSample <: SamplingAlgorithm end
```

Samples using a randomly-shifted rank-1 lattice rule
using LatticeRules.jl
"""
struct LatticeRuleSample <: SamplingAlgorithm end

"""
```julia
struct LowDiscrepancySample{T} <: SamplingAlgorithm
```

`base[i]` is the base in the ith direction.
"""
struct LowDiscrepancySample{T} <: SamplingAlgorithm
    base::T
end

"""
```julia
struct RandomSample <: SamplingAlgorithm end
```
"""
struct RandomSample <: SamplingAlgorithm end

"""
```julia
struct KroneckerSample{A,B} <: SamplingAlgorithm
```

`KroneckerSample(alpha, s0)` for a Kronecker sequence, where alpha is an length-d vector of irrational numbers (often sqrt(d)) and s0 is a length-d seed vector (often 0).
"""
struct KroneckerSample{A,B} <: SamplingAlgorithm
    alpha::A
    s0::B
end

"""
```julia
struct GoldenSample <: SamplingAlgorithm end
```
"""
struct GoldenSample <: SamplingAlgorithm end

"""
```julia
struct SectionSample{T} <: SamplingAlgorithm
```

`SectionSample(x0, sampler)` where `sampler` is any sampler above and `x0` is a vector of either `NaN` for a free dimension or some scalar for a constrained dimension.
"""
struct SectionSample{T} <: SamplingAlgorithm
    x0::Vector{T}
    sa::SamplingAlgorithm
end

"""
```julia
A = QuasiMonteCarlo.sample(n,lb,ub,sample_method)
```

where:

- `n` is the number of points to sample.
- `lb` is the lower bound for each variable. The length determines the dimensionality.
- `ub` is the upper bound.
- `sample_method` is the quasi-Monte Carlo sampling strategy. Note that any Distributions.jl
  distribution can be used in addition to any `SamplingAlgorithm`.
"""
function sample end

"""
sample(n,lb,ub,S::GridSample)
Returns a tuple containing numbers in a grid.
"""
function sample(n,lb,ub,S::GridSample)
    dx = S.dx
    if lb isa Number
        return vec(rand(lb:S.dx:ub,n))
    else
        d = length(lb)
        x = [[rand(lb[j]:dx[j]:ub[j]) for j = 1:d] for i in 1:n]
        return reduce(hcat,x)
    end
end

"""
sample(n,lb,ub,::UniformRandom)
Returns a tuple containing uniform random numbers.
"""
function sample(n,lb,ub,::UniformSample)
    if lb isa Number
        return rand(Uniform(lb,ub),n)
    else
        d = length(lb)
        x = [[rand(Uniform(lb[j],ub[j])) for j in 1:d] for i in 1:n]
        return reduce(hcat,x)
    end
end

"""
sample(n,lb,ub,::SobolSampling)
Returns a tuple containing Sobol sequences.
"""
function sample(n,lb,ub,::SobolSample)
    s = SobolSeq(lb,ub)
    skip(s,n)
    if lb isa Number
        return [next!(s)[1] for i = 1:n]
    else
        return reduce(hcat,[next!(s) for i = 1:n])
    end
end

"""
sample(n,lb,ub,T::LatinHypercubeSample)
Returns a tuple containing LatinHypercube sequences.
"""
function sample(n,lb,ub,T::LatinHypercubeSample)
    threading = T.threading
    d = length(lb)
    if lb isa Number
        x = vec(LHCoptim(n,d,1; threading=threading)[1])
        # x∈[0,n], so affine transform
        return @. (ub-lb) * x/(n) + lb
    else
        lib_out = collect(float(LHCoptim(n,d,1; threading=threading)[1])')
        # x∈[0,n], so affine transform column-wise
        @inbounds for c = 1:d
            lib_out[c, :] = (ub[c]-lb[c])*lib_out[c, :]/n .+ lb[c]
        end
        return lib_out
    end
end

"""
sample(n,lb,ub,::LatticeRuleSample)
Returns a matrix with the `n` rank-1 lattice points in each column if `lb` is a vector, or a vector with the `n` rank-1 lattice points if `lb` is a number.
"""
function sample(n,lb,ub,::LatticeRuleSample)
    if lb isa Number
        lat = ShiftedLatticeRule(1)
        pts = reduce(vcat,Iterators.take(lat, n))
        # transform from [0, 1] to [lb, ub]
        @inbounds for i in 1:n
            pts[i] = (ub-lb)*pts[i] + lb
        end
        return pts
    else
        d = length(lb)
        lat = ShiftedLatticeRule(d)
        pts = reduce(hcat,Iterators.take(lat, n))
        # transform from [0, 1]^d to [lb, ub]
        @inbounds for j in 1:n
            for i in 1:d
                pts[i, j] = (ub[i]-lb[i])*pts[i, j] + lb[i]
            end
        end
        return pts
    end
end

"""
sample(n,lb,ub,S::LowDiscrepancySample)
Low-discrepancy sample:
- Dimension 1: Van der Corput sequence
- Dimension > 1: Halton sequence
If dimension d > 1, all bases must be coprime with one other.
"""
function sample(n,lb,ub,S::LowDiscrepancySample)
    @assert length(lb) == length(ub)

    d = length(lb)
    t = float(eltype(lb)) # infer data type for return values
    if d == 1
        #Van der Corput
        b = S.base
        x = zeros(t, n)
        for i = 1:n
            expansion = digits(i,base = b)
            L = length(expansion)
            val = zero(t)
            for k = 1:L
                val += expansion[k]*float(b)^(-(k-1)-1)
            end
            x[i] = val
        end
        # It is always defined on the unit interval, resizing:
        return @. (ub-lb) * x + lb
    else
        #Halton sequence
        x = zeros(t,d,n)
        for j = 1:d
            b = S.base[j]
            for i = 1:n
                val = zero(t)
                expansion = digits(i, base = b)
                L = length(expansion)
                for k = 1:L
                    val += expansion[k]*float(b)^(-(k-1)-1)
                end
                x[j,i] = val
            end
        end
        #Resizing
        # x∈[0,1], so affine transform column-wise
        @inbounds for c = 1:d
            x[c, :] = (ub[c]-lb[c])*x[c, :] .+ lb[c]
        end
        return x
    end
end

"""
sample(n,d,K::KroneckerSample)

Returns a Tuple containing numbers following the Kronecker sample
"""
function sample(n,lb,ub,K::KroneckerSample)
    d = length(lb)
    alpha = K.alpha
    s0 = K.s0
    if d == 1
        x = zeros(n)
        @inbounds for i = 1:n
            x[i] = (s0+i*alpha)%1
        end
        return @. (ub-lb) * x + lb
    else
        x = zeros(n,d)
        @inbounds for j = 1:d
            for i = 1:n
                x[i,j] = (s0[j] + i*alpha[j])%i
            end
        end
        #Resizing
        # x∈[0,1], so affine transform column-wise
        @inbounds for c = 1:d
            x[:,c] = (ub[c]-lb[c])*x[:,c] .+ lb[c]
        end

        y = collect(x')
        return y
    end
end

function sample(n,lb,ub, G::GoldenSample)
    d = length(lb)
    if d == 1
        x = zeros(n)
        g = (sqrt(5)+1)/2
        a = 1.0/g
        for i = 1:n
            x[i] = (0.5+a*i)%1
        end
        return @. (ub-lb) * x + lb
    else
        x = zeros(n,d)
        for j = 1:d
            #Approximate solution of x^(d+1) = x + 1, a simple newton is good enough
            y = 2.0
            for s = 1:10
                g = (1+y)^(1/(j+1))
            end
            a = 1.0/g
            for i = 1:n
                x[i,j] = (0.5+a*i)%1
            end
        end
        @inbounds for c = 1:d
            x[:,c] = (ub[c]-lb[c])*x[:,c] .+ lb[c]
        end
        y = collect(x')
        return y
    end
end

fixed_dimensions(
        section_sampler::SectionSample)::Vector{Int64} = findall(
    x->x == false, isnan.(section_sampler.x0))

free_dimensions(
        section_sampler::SectionSample)::Vector{Int64} = findall(
    x->x == true, isnan.(section_sampler.x0))

"""
sample(n,lb,ub,K::SectionSample)

Returns Tuples constrained to a section.

In surrogate-based identification and control, optimization can alternate between unconstrained sampling in the full-dimensional parameter space, and sampling constrained on specific sections (e.g. a planes in a 3D volume),

A SectionSampler allows sampling and optimizing on a subset of 'free' dimensions while keeping 'fixed' ones constrained.
The sampler is defined as in e.g.

`section_sampler_y_is_10 = SectionSample([NaN64, NaN64, 10.0, 10.0], UniformSample())`

where the first argument is a Vector{T} in which numbers are fixed coordinates and `NaN`s correspond to free dimensions, and the second argument is a SamplingAlgorithm which is used to sample in the free dimensions.
"""
function sample(n,lb,ub,section_sampler::SectionSample)
    if lb isa Number
        if isnan(section_sampler.x0[1])
            return sample(n, lb, ub, section_sampler.sa)
        else
            return fill(section_sampler.x0[1], n)
        end
    else
        d_free = free_dimensions(section_sampler)
        new_samples = sample(n, lb[d_free], ub[d_free], section_sampler.sa)
        out_as_vec = collect(repeat(section_sampler.x0', n, 1)')
        for y in 1:size(out_as_vec,2)
            for (xi, d) in enumerate(d_free)
                out_as_vec[d,y] = new_samples[xi, y]
            end
        end
        return out_as_vec
    end
end

"""
sample(n,d,D::Distribution)
Returns a tuple containing numbers distributed as D.
"""
function sample(n,d,D::Distribution)
    if d == 1
        return rand(D,n)
    else
        x = [[rand(D) for j in 1:d] for i in 1:n]
        return reduce(hcat, x)
    end
end

"""
```julia
k=2
As = QuasiMonteCarlo.generate_design_matrices(n,lb,ub,sample_method,k)
```

which returns `As` which is an array of `k` design matrices `A[i]` that are
all sampled from the same low-discrepancy sequence.
"""
function generate_design_matrices(n,lb,ub,sampler,num_mats = 2)
    @assert length(lb) == length(ub)
    d = length(lb)
    out = sample(n, repeat(lb,num_mats), repeat(ub,num_mats),sampler)
    [out[(j*d+1):((j+1)*d),:] for j in 0:num_mats-1]
end

export GridSample, UniformSample, SobolSample, LatinHypercubeSample, LatticeRuleSample,
       RandomSample, LowDiscrepancySample, GoldenSample, KroneckerSample, SectionSample

end # module


text:
function attach_dirichlet(cellmatvec,cellvals,cellmask=Fill(true,length(cellvals)))
  k = AttachDirichletMap()
  lazy_map(k,cellmatvec,cellvals,cellmask)
end

struct ZeroVectorMap <: Map end

function Arrays.return_cache(k::ZeroVectorMap,mat::AbstractMatrix)
  v = similar(mat,eltype(mat),size(mat,1))
  fill!(v,zero(eltype(v)))
  CachedArray(v)
end

function Arrays.evaluate!(cache,k::ZeroVectorMap,mat::AbstractMatrix)
  setsize!(cache,(size(mat,1),))
  v = cache.array
  fill!(v,zero(eltype(v)))
  v
end

function Arrays.return_cache(k::ZeroVectorMap,mat::MatrixBlock)
  ni,nj = size(mat)
  ai = testitem(mat)
  ci = return_cache(k,ai)
  vi = evaluate!(ci,k,ai)
  c = Vector{typeof(ci)}(undef,ni)
  array = Vector{typeof(vi)}(undef,ni)
  touched = fill(false,ni)
  for i in 1:ni
    for j in 1:nj
      if !touched[i] && mat.touched[i,j]
        c[i] = return_cache(k,mat.array[i,j])
        touched[i] = true
      end
    end
  end
  ArrayBlock(array,touched), c
end

function Arrays.evaluate!(cache,k::ZeroVectorMap,mat::MatrixBlock)
  r,c = cache
  ni,nj = size(mat)
  fill!(r.touched,false)
  for j in 1:nj
    for i in 1:ni
      if !r.touched[i] && mat.touched[i,j]
        r.array[i] = evaluate!(c[i],k,mat.array[i,j])
        r.touched[i] = true
      end
    end
  end
  r
end

struct AttachDirichletMap <: Map
  muladd::MulAddMap{Int}
  AttachDirichletMap() = new(MulAddMap(-1,1))
end

function Arrays.return_value(k::AttachDirichletMap,matvec::Tuple,vals,mask)
  mat, vec = matvec
  vec_with_bcs = return_value(k.muladd,mat,vals,vec)
  (mat, vec_with_bcs)
end

function Arrays.return_cache(k::AttachDirichletMap,matvec::Tuple,vals,mask)
  mat, vec = matvec
  return_cache(k.muladd,mat,vals,vec)
end

function Arrays.evaluate!(cache,k::AttachDirichletMap,matvec::Tuple,vals,mask)
  if mask
    mat, vec = matvec
    vec_with_bcs = evaluate!(cache,k.muladd,mat,vals,vec)
    (mat, vec_with_bcs)
  else
    matvec
  end
end

function Arrays.return_value(k::AttachDirichletMap,mat,vals,mask)
  cv = return_value(ZeroVectorMap(),mat)
  (mat,cv)
end

function Arrays.return_cache(k::AttachDirichletMap,mat,vals,mask)
  cm = return_cache(*,mat,vals)
  cv = return_cache(ZeroVectorMap(),mat)
  (cm,cv)
end

function Arrays.evaluate!(cache,k::AttachDirichletMap,mat,vals,mask)
  cm, cv = cache
  if mask
    vec_with_bcs = evaluate!(cm,*,mat,vals)
    rmul!(vec_with_bcs,-1)
  else
    vec_with_bcs = evaluate!(cv,ZeroVectorMap(),mat)
  end
  (mat, vec_with_bcs)
end


text:
using OrbitalTrajectories

using Test
using Unitful

function check_libration_points(model)
    L = libration_points(model)
    @test 0.5 < L[1][1] < 1.0        # L1 (X)
    @test L[1][2] == L[1][3] == 0.0  # L1 (Y & Z)
    @test 1.0 < L[2][1] < 1.5        # L2 (X)
    @test L[2][2] == L[2][3] == 0.0  # L2 (Y & Z)
    @test -1.5 < L[3][1] < -1.0      # L3 (X)
    @test L[3][2] == L[3][3] == 0.0  # L3 (Y & Z)
    @test 0.0 < L[4][1] < 1.0        # L4 (X)
    @test 0.0 < L[4][2] < 1.0        # L4 (Y)
    @test L[4][1] == L[5][1]         # L4 & L5 (X)
    @test L[4][2] == -L[5][2]        # L4 & L5 (Y)
    @test L[4][3] == L[5][3] == 0.0  # L4 & L5 (Z)
    return L
end

@testset "3BP Properties" begin
    # Cross-check the Mars-Phobos system (values from Wikipedia)
    MarsPhobos = OrbitalTrajectories.Dynamics.R3BPSystemProperties(:mars, :phobos)
    @test MarsPhobos.L ≈ 9376.0u"km" rtol=1e-5
    @test MarsPhobos.V ≈ 2.138u"km/s" rtol=1e-5
    @test MarsPhobos.T ≈ 0.31891023u"d"/2π rtol=1e-5
    @test MarsPhobos.R1 ≈ [3396.2, 3396.2, 3376.2]u"km"  rtol=1e-1
    @test MarsPhobos.R2 ≈ [13.5, 11., 9.]u"km"  rtol=1e-1

    # Check the relationship between libration points
    check_libration_points(MarsPhobos)

    SunMars = OrbitalTrajectories.Dynamics.R3BPSystemProperties(:sun, :mars)
    L = check_libration_points(SunMars)

    # Validate against constants in [DeiTos2018]
    @test SunMars.L ≈ 2.279497905330276e8u"km"                          rtol=1e-4
    @test SunMars.V ≈ 24.128831378998047u"km/s"                         rtol=1e-4
    @test SunMars.R2[1] ≈ 3396.19u"km"                                  rtol=1e-4
    @test SunMars.T ≈ 109.3425420965616u"d"                             rtol=1e-4
    @test L[1][1] ≈ (1 - SunMars.μ) + (-1.082385474e6u"km" / SunMars.L) rtol=1e-4  # (with respect to Mars)
    @test SunMars.e ≈ 0.0935643512                                      rtol=1e-4

    # XXX: These values as listed in [DeiTos2018] are possibly wrong?
    #@test SunMars.μ ≈ 3.227154876045166e-6      rtol=1e-4
end