import Base: +, -, ==

immutable BoundingBox{T}
  i::T
  j::T
  h::T    # height
  w::T    # width
end

eps = 1e-6

BoundingBox() = BoundingBox(0,0,0,0)
BoundingBox(a, b, c, d) = BoundingBox(promote(a,b,c,d)...)

"""
Add bounding boxes to find BB of their union
"""
function +(bbA::BoundingBox, bbB::BoundingBox)
  i = min(bbA.i, bbB.i)
  j = min(bbA.j, bbB.j)
  h = max(bbA.h+bbA.i, bbB.h+bbB.i) - i
  w = max(bbA.w+bbA.j, bbB.w+bbB.j) - j
  return BoundingBox(i,j,h,w)
end

"""
Subtract bounding boxes to find BB of their intersection
"""
function -(bbA::BoundingBox, bbB::BoundingBox)
  i = max(bbA.i, bbB.i)
  j = max(bbA.j, bbB.j)
  h = min(bbA.h+bbA.i, bbB.h+bbB.i)-i
  w = min(bbA.w+bbA.j, bbB.w+bbB.j)-j
  if h <= 0 || w <= 0
    bb = BoundingBox(NaN, NaN, NaN, NaN)
  else
    bb = BoundingBox(i,j,h,w)
  end
  return bb
end

"""
Boolean if bounding boxes intersect
"""
function intersects(bbA::BoundingBox, bbB::BoundingBox)
  bb = bbA - bbB
  return !isnan(bb.i)
end

"""
Test if BoundingBoxes have the same origin and dimensions
"""
function ==(bbA::BoundingBox, bbB::BoundingBox)
  return bbA.i == bbB.i && bbA.j == bbB.j && bbA.w == bbB.w && bbA.h == bbB.h
end

"""
Convert bounding box object to a polygon point list (counter-clockwise)
"""
function bb_to_pts(r)
  return [r.i r.j;
          r.i+r.h r.j;
          r.i+r.h r.j+r.w;
          r.i r.j+r.w;
          r.i r.j];
end

"""
Return xmin, ymin, xmax, ymax of a bounding box (opposing corner definition)
"""
function minsandmax(bb)
  return bb.i, bb.j, bb.i+bb.h, bb.j+bb.w
end

"""
`FIND_MESH_BB` - Find bounding box around mesh

    BoundingBox(ilow, jlow, height, width) = find_mesh_bb(nodes)

* `nodes`: 2xN matrix of mesh nodes
* `BoundingBox`: smallest integer-valued rectangle containing all mesh nodes

"""
function find_mesh_bb(nodes)
    ilow = floor(Int64,minimum(nodes[:,1]))
    jlow = floor(Int64,minimum(nodes[:,2]))
    ihigh = ceil(Int64,maximum(nodes[:,1]))
    jhigh = ceil(Int64,maximum(nodes[:,2]))
    return BoundingBox(ilow, jlow, ihigh-ilow, jhigh-jlow)
end

"""
Snap bounding box to integer values (smallest rectangle containing original)
Returns BoundingBox of integers
"""
function snap_bb(bb)
  r = bb_to_pts(bb)
  i = floor(Int,bb.i)
  j = floor(Int,bb.j)
  h = ceil(Int,r[2,1]) - i
  w = ceil(Int,r[3,2]) - j
  return BoundingBox{Int}(i, j, h, w)
end

"""
Apply affine transform to points in a bounding box & find new bounding box
"""
function tform_bb(bb, tform)
  tform_pts = [bb_to_pts(bb) ones(size(bb_to_pts(bb),1),1)] * tform
  i = minimum(tform_pts[:,1])
  j = minimum(tform_pts[:,2])
  h = maximum(tform_pts[:,1])-i
  w = maximum(tform_pts[:,2])-j
  return BoundingBox(i, j, h, w)
end

"""
Convert Tuple for image size into a BoundingBox at (1,1)
"""
function sz_to_bb(sz)
  return BoundingBox(0, 0, sz[1], sz[2])
end

"""
Convert a BoundingBox to its snapped sizes
"""
function bb_to_sz(bb)
  bb = snap_bb(bb);
  return bb.h, bb.w
end

"""
Get bounding box offset
"""
function get_offset(bb::BoundingBox)
  return [bb.i, bb.j]
end

"""
Get height & width of bounding box
"""
function get_size(bb::BoundingBox)
  return [bb.h, bb.w]
end

"""
Get four-tuple of bounding box upper-left & lower-right in i,j coordinates
"""
function get_bounds(bb::BoundingBox)
  return (bb.i, bb.j, bb.i+bb.h, bb.j+bb.w)
end

"""
Get four-tuple of bounding box upper-left & dimensions in x,y coordinates

Convention used by Cairo package
"""
function get_rect(bb::BoundingBox)
  return (bb.j, bb.i, bb.w, bb.h)
end

"""
Return the product of the width & height of bounding box
"""
function get_area(bb::BoundingBox)
  return bb.w*bb.h
end

"""
Convert bounding box to tuple of ranges for easy array slicing
"""
function bb_to_slice(bb::BoundingBox{Int64})
  return (bb.i):(bb.i+bb.h-1), (bb.j):(bb.j+bb.w-1)
end

function bb_to_slice(bb::BoundingBox{Float64})
  return round(Int64, bb.i):round(Int64, bb.i+bb.h-1),
              round(Int64, bb.j):round(Int64, bb.j+bb.w-1)
end

"""
Convert tuple of ranges to bounding box
"""
function slice_to_bb(slice)
  return BoundingBox(slice[1][1], slice[2][1],
                      slice[1][end]-slice[1][1]+1, slice[2][end]-slice[2][1]+1)
end

"""
Shift bounding box by 2-element array
"""
function translate_bb(bb::BoundingBox, offset)
  return BoundingBox(bb.i + offset[1], bb.j + offset[2], bb.h, bb.w)
end

"""
Transform bounding box by scaling matrix and snap for nearest integer
"""
function scale_bb(bb::BoundingBox{Int64}, scale)
  tform = make_scale_matrix(scale)
  tbb = tform_bb(bb, tform)
  return snap_bb(tbb)
end

"""
Transform bounding box by scaling matrix
"""
function scale_bb(bb::BoundingBox{Float64}, scale)
  tform = make_scale_matrix(scale)
  return tform_bb(bb, tform)
end


lines = readlines(ARGS[1])

function solution1(lines)
    horizontal = 0
    depth = 0
    for line in lines
        command, distance = split(line)
        int_distance = parse(Int, distance)
        if command == "up"
            depth -= int_distance
        elseif command == "down"
            depth += int_distance
        elseif command == "forward"
            horizontal += int_distance
        else
            throw(Exception(command, "unknown input"))
        end
    end
    return depth * horizontal
end

function solution2(lines)
    horizontal = 0
    depth = 0
    aim = 0
    for line in lines
        command, distance = split(line)
        int_distance = parse(Int, distance)
        if command == "up"
            aim -= int_distance
        elseif command == "down"
            aim += int_distance
        elseif command == "forward"
            horizontal += int_distance
            depth += aim * int_distance
        else
            throw(Exception(command, "unknown input"))
        end
    end
    return depth * horizontal
end

println(solution1(lines))
println(solution2(lines))


### A Pluto.jl notebook ###
# v0.12.18

using Markdown
using InteractiveUtils

# ╔═╡ c2e4a0d0-5cae-11eb-199f-3b8bc7ca4867
begin
	using ReinforcementLearning
	using Flux
	using Statistics
	using Plots
	using SparseArrays
end

# ╔═╡ 65614b3c-5cae-11eb-0613-93393f90cfc0
md"""
# Chapter 10 Mountain Car
"""

# ╔═╡ 054805aa-5caf-11eb-1f4a-cd04a590a023
md"""
The `MountainCarEnv` is already provided in `ReinforcementLearning.jl`. So we can use it directly here. Note that by default this environment will terminate at the maximum step of `200`. While in the example on the book there's no such restriction.
"""

# ╔═╡ db653a64-5cae-11eb-0a6d-cd071416fc50
env = MountainCarEnv()

# ╔═╡ dc826048-5cae-11eb-3141-1b589c25852a
S = state_space(env)

# ╔═╡ 02252420-5caf-11eb-2352-251cf8b60517
md"""
First let's define a `Tiling` structure to encode the state.
"""

# ╔═╡ e1e987aa-5cae-11eb-1f24-69fefb87d4e0
begin
	struct Tiling{N,Tr<:AbstractRange}
		ranges::NTuple{N,Tr}
		inds::LinearIndices{N,NTuple{N,Base.OneTo{Int}}}
	end

	Tiling(ranges::AbstractRange...) =Tiling(
		ranges,
		LinearIndices(Tuple(length(r) - 1 for r in ranges))
	)

	Base.length(t::Tiling) = reduce(*, (length(r) - 1 for r in t.ranges))

	function Base.:-(t::Tiling, xs)
		Tiling((r .- x for (r, x) in zip(t.ranges, xs))...)
	end

	encode(range::AbstractRange, x) = floor(Int, div(x - range[1], step(range)) + 1)

	encode(t::Tiling, xs) = t.inds[CartesianIndex(Tuple(map(encode,  t.ranges, xs)))]
end

# ╔═╡ a0f8a23c-5caf-11eb-01d2-1706e563dac1
begin
	ntilings = 8
	ntiles = 8
	tiling = Tiling(
		(
			range(r.left, step=(r.right-r.left)/ntiles, length=ntiles+2)
			for r in S
		)...
	)
	offset = map(x-> x.right - x.left, S) ./ (ntiles * ntilings)
	tilings = [tiling - offset .* (i-1) for i in 1:ntilings]
end

# ╔═╡ d62c2d40-5caf-11eb-36e4-5bb8ee95c24e
md"""
The rest parts are simple, we initialize `agent` and `env`, then roll out experiments:
"""

# ╔═╡ 214a577a-5cb0-11eb-356b-ab06ad454eef
function create_env_agent(α=2e-4, n=0)
    env = StateOverriddenEnv(
        MountainCarEnv(;max_steps=10000),
        s -> sparse(map(t -> encode(t, s), tilings), 1:8, ones(8), 81, 8) |> vec
    )

    agent = Agent(
        policy=QBasedPolicy(
            learner=TDLearner(
                approximator=LinearQApproximator(
                    n_state=81*8,
                    n_action=3,
                    opt = Descent(α)
                    ),
                method=:SARSA,
                n=n
                ),
            explorer=GreedyExplorer()
            ),
        trajectory=VectorSARTTrajectory(;state=Vector{Int})
    )

    env, agent
end

# ╔═╡ 3901d7b4-5cb0-11eb-2cc3-9ff34896f465
X = range(S[1].left, stop=S[1].right, length=40)

# ╔═╡ 3c0e7bea-5cb0-11eb-3190-598e56c887ff
Y = range(S[2].left, stop=S[2].right, length=40)

# ╔═╡ 230ee9ae-5cb0-11eb-0fe5-6fd8f5482af4
function show_approximation(n)
    env, agent = create_env_agent()
    run(agent, env, StopAfterEpisode(n))
    [
		agent.policy.learner.approximator(env.f([p, v])) |> maximum
        for p in X, v in Y
	]
end

# ╔═╡ 4ba11496-5cb0-11eb-202c-7f2baba059bf
n = 10

# ╔═╡ 4905a328-5cb0-11eb-03f1-0f6f27350ade
plot(X, Y, -show_approximation(n), linetype=:wireframe)

# ╔═╡ 68779ca2-5cb0-11eb-11e4-9fdb4ea0989f
begin
	fig_10_2 = plot(legend=:topright)
	n_runs = 5  # quite slow here, need revisit
	for α in [0.1/8, 0.2/8, 0.5/8]
		avg_steps_per_episode = zeros(500)
		for _ in 1:n_runs
			env, agent = create_env_agent(α)
			hook = StepsPerEpisode()
			run(agent, env, StopAfterEpisode(500; is_show_progress=false),hook)
			avg_steps_per_episode .+= hook.steps
		end
		plot!(fig_10_2, avg_steps_per_episode ./ n_runs, yscale=:log10, label="α=$α")
	end
	fig_10_2
end

# ╔═╡ 79e4a002-5cb0-11eb-049a-db8707d6c400
begin
	function run_once(α, n;is_reduce=true, n_episode=50)
		env, agent = create_env_agent(α, n)
		hook = StepsPerEpisode()
		run(agent, env, StopAfterEpisode(n_episode; is_show_progress=false),hook)
		is_reduce ? mean(hook.steps) : hook.steps
	end
	fig_10_3 = plot()
	plot!(fig_10_3, mean(run_once(0.5/8, 1; is_reduce=false, n_episode=500) for _ in 1:10), yscale=:log10)
	plot!(fig_10_3, mean(run_once(0.3/8, 8; is_reduce=false, n_episode=500) for _ in 1:10), yscale=:log10)
	fig_10_3
end

# ╔═╡ 832c5ce0-5cb0-11eb-06f6-4f9ad5b707a2
begin
	fig_10_4 = plot(legend=:topright)
	for (A, n) in [(0.4:0.1:1.7, 1), (0.3:0.1:1.6, 2), (0.2:0.1:1.4, 4), (0.2:0.1:0.9, 8), (0.2:0.1:0.7, 16)]
		plot!(fig_10_4, A, [mean(run_once(α/8, n) for _ in 1:5) for α in A], label="n = $n")
	end
	fig_10_4
end

# ╔═╡ Cell order:
# ╟─65614b3c-5cae-11eb-0613-93393f90cfc0
# ╠═c2e4a0d0-5cae-11eb-199f-3b8bc7ca4867
# ╟─054805aa-5caf-11eb-1f4a-cd04a590a023
# ╠═db653a64-5cae-11eb-0a6d-cd071416fc50
# ╠═dc826048-5cae-11eb-3141-1b589c25852a
# ╟─02252420-5caf-11eb-2352-251cf8b60517
# ╠═e1e987aa-5cae-11eb-1f24-69fefb87d4e0
# ╠═a0f8a23c-5caf-11eb-01d2-1706e563dac1
# ╟─d62c2d40-5caf-11eb-36e4-5bb8ee95c24e
# ╠═214a577a-5cb0-11eb-356b-ab06ad454eef
# ╠═3901d7b4-5cb0-11eb-2cc3-9ff34896f465
# ╠═3c0e7bea-5cb0-11eb-3190-598e56c887ff
# ╠═230ee9ae-5cb0-11eb-0fe5-6fd8f5482af4
# ╠═4ba11496-5cb0-11eb-202c-7f2baba059bf
# ╠═4905a328-5cb0-11eb-03f1-0f6f27350ade
# ╠═68779ca2-5cb0-11eb-11e4-9fdb4ea0989f
# ╠═79e4a002-5cb0-11eb-049a-db8707d6c400
# ╠═832c5ce0-5cb0-11eb-06f6-4f9ad5b707a2


# ---
# title: 931. Minimum Falling Path Sum
# id: problem931
# author: zhwang
# date: 2022-02-23
# difficulty: Medium
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/minimum-falling-path-sum/description/>
# hidden: true
# ---
#
# Given a **square** array of integers `A`, we want the **minimum** sum of a
# _falling path_ through `A`.
#
# A falling path starts at any element in the first row, and chooses one element
# from each row.  The next row's choice must be in a column that is different
# from the previous row's column by at most one.
#
#
#
# **Example 1:**
#
#
#
#     Input: [[1,2,3],[4,5,6],[7,8,9]]
#     Output: 12
#     Explanation:
#     The possible falling paths are:
#
#
#   * `[1,4,7], [1,4,8], [1,5,7], [1,5,8], [1,5,9]`
#   * `[2,4,7], [2,4,8], [2,5,7], [2,5,8], [2,5,9], [2,6,8], [2,6,9]`
#   * `[3,5,7], [3,5,8], [3,5,9], [3,6,8], [3,6,9]`
#
# The falling path with the smallest sum is `[1,4,7]`, so the answer is `12`.
#
#
#
# **Constraints:**
#
#   * `1 <= A.length == A[0].length <= 100`
#   * `-100 <= A[i][j] <= 100`
#
#
## @lc code=start
using LeetCode

function min_falling_path_sum(matrix::Vector{Vector{Int}})::Int
    n = length(matrix)
    function f(l1, l2)
        return [
            num + minimum(l1[max(i - 1, 1):min(i + 1, n)]) for (i, num) in enumerate(l2)
        ]
    end
    return minimum(foldl(f, matrix))
end

## @lc code=end


# flexible representations of priors

# Collection of priors from -- False Discovery Rates: A new deal

ash_spiky =  MixtureModel( [Normal(0,0.25); Normal(0,0.5); Normal(0,1); Normal(0,2)] ,
                           [0.4; 0.2; 0.2; 0.2 ])

ash_nearnormal = MixtureModel([Normal(0,1); Normal(0,2)], [1/3; 2/3])

ash_flattop = MixtureModel([Normal(x,0.5) for x=-1.5:0.5:1.5])

ash_skew = MixtureModel([Normal(-2,2), Normal(-1,1.5), Normal(0,1), Normal(1,1)],
                        [1/4; 1/4; 1/3; 1/6])

ash_bignormal = Normal(0,4)

ash_bimodal = MixtureModel( [Normal(-2,1); Normal(2,1)],
                            [0.5;0.5] )

# Distributions used for simulations in this paper

hard_dist = MixtureModel([ Normal(-0.2,.2), Normal(0,0.9)],[0.7, 0.3])
easy_dist = MixtureModel([Normal(-1.5,.2), Normal(1.5, .2)])


using Sundials
function exotherm(t, x, dx; n=1, a=1, b=1)
  p = x[2]^n * exp(x[1])
  dx[1] = p - a*x[1]
  dx[2] = -b*p
  return(dx)
end;

t = linspace(0.0,5.0,1001);
fexo(t,x,dx) = exotherm(t, x, dx, a=0.6, b=0.1);

x1 = Sundials.cvode(fexo, [0.9, 1.0], t);
using Winston
plot(t,x1[:,1]))


# SPDX-License-Identifier: X11
# 2020-11-22
# Sum and Product (300pt)

function main()
  S,P = parse.(Int, (split ∘ readline)())

  det::BigInt = big(S)^2 - 4P
  if det < 0
    println("No")
    return
  end

  rtdet = √det
  if !isinteger(rtdet)
    println("No")
    return
  end

  x,y = S + Int(rtdet), S - Int(rtdet)
  println(x & 1 == 0 && x > 0 &&
          y & 1 == 0 && y > 0
          ? "Yes" : "No")
end

main()


module _Inplace

using Brainstorm._Math._GMP: set!, add!, sub!, mul!, div!, neg!, isqrt!, rsh!, lsh!

const guard_digits = 5
const eps_digits = 5

pi(digits::Integer) = begin
    o = big(10)^(digits + guard_digits + eps_digits)
    osqr = o^2
    t = big(0)

    # a = 1
    a = big(0)
    set!(a, o)

    # b = 1 / sqrt(2)
    set!(t, osqr)
    lsh!(t, 1)
    isqrt!(t)
    b = osqr ÷ t

    s = big(0)
    d = a - b

    n = 0
    ϵ = 10^eps_digits

    while d >= ϵ
        # s += 2^n * d^2
        mul!(d, d)
        lsh!(d, n)
        add!(s, d)

        # a, b = (a + b) / 2, sqrt(a * b)
        set!(t, a)
        add!(t, b)
        rsh!(t, 1)

        mul!(b, a)
        isqrt!(b)
        set!(a, t)

        # d = a - b
        set!(d, a)
        sub!(d, b)

        n += 1
    end

    # 4*a^2 / (1 - s)
    neg!(s)
    add!(s, osqr)
    mul!(s, big(10)^(guard_digits + eps_digits))
    mul!(a, a)
    mul!(a, o)
    lsh!(a, 2)
    div!(a, s)
end

end


using Surrogates
using PolyChaos

#1D
n = 20
lb = 0.0
ub = 4.0
f = x-> 2*x
x = sample(n,lb,ub,SobolSample())
y = f.(x)

my_pce = PolynomialChaosSurrogate(x,y,lb,ub)
val = my_pce(2.0)
add_point!(my_pce,3.0,6.0)
my_pce_changed = PolynomialChaosSurrogate(x,y,lb,ub,op=Uniform01OrthoPoly(1))

#ND
n = 60
lb = [0.0,0.0]
ub = [5.0,5.0]
f = x-> x[1]*x[2]
x = sample(n,lb,ub,SobolSample())
y = f.(x)

my_pce = PolynomialChaosSurrogate(x,y,lb,ub)
val = my_pce((2.0,2.0))
add_point!(my_pce,(2.0,3.0),6.0)

op1 = Uniform01OrthoPoly(1)
op2 = Beta01OrthoPoly(2,2,1.2)
ops = [op1,op2]
multi_poly = MultiOrthoPoly(ops,min(1,2))
my_pce_changed = PolynomialChaosSurrogate(x,y,lb,ub,op = multi_poly)


using LazySets, MathematicalPredicates

function crane_specification()
    # initial set
    X0 = Hyperrectangle([2.5, 0.0, 0.0, 0.0, 0.0, 0.0],
                        [2.5, 0.0, 0.2, 0.1, 0.0, 0.0])

    # safety property: x1 ≥ -1.8
    property = is_contained_in(HalfSpace([-1., 0.], 1.8))

    # time horizon: 15 time units
    time_horizon = 15.0

    # specification
    O = Dict(:T => time_horizon, :property => property)

    return X0, O
end


# ==================================
# Motor
#
# system type: continuous LTI system
# state dimension: 8
# input dimension: 2
# ==================================
using ReachabilityBenchmarks, MathematicalSystems, LazySets, SparseArrays

function motor_model()
    # system matrix
    I = [1, 2, 2, 3, 3, 3, 3, 4, 5, 6, 6, 7, 7, 7, 7, 8]
    J = [2, 3, 2, 1, 2, 3, 4, 1, 6, 7, 6, 5, 6, 7, 8, 5]
    V = [1, 8487.2, -1.0865, -2592.1, -21.119, -698.91, -141399.0, 1.0, 1.0,
         8487.2, -1.0865, -2592.1, -21.119, -698.91, -141399.0, 1.0]
    A = sparse(I, J, V)

    # input matrix
    B = sparse([4, 8], [1, 2], [-1.0, -1.0])

    # state domain
    X = Universe(8)

    # input domain
    U = Hyperrectangle([0.23, 0.3], [0.07, 0.1])

    # continuous LTI system
    S = @system(x' = Ax + Bu, x ∈ X, u ∈ U)

    return S
end


# Set up return distribution
using Random; Random.seed!(15)
n_stocks = 10
μ = rand(MvNormal(zeros(n_stocks), 0.1I))
Σ = rand(InverseWishart(n_stocks, diagm(ones(n_stocks))))

# Imports
using JuMP
using Ipopt

function opt_portfolio(R)
    # Preallocate model
    model = Model(Ipopt.Optimizer)

    # Set up variables
    @variable(model, -1 <= q[n_stocks] <= 1)
    @objective(model, Min, q'Σ*q)
    @constraint(model, sum(q) == 1)
    @constraint(model, q'μ == R)
    optimize!(model)

    # Extract optimal values
    q_star = value.(q)
    return q_star'Σ*q_star, q_star'μ
end


"""
    BacktrackingLineSearch{C, T}

`BacktrackingLineSearch` starts from an initial step length guess and gradually backtracks
until the Armijo sufficient decrease condition is met.

# Fields
- `c::C`: the constant for the Armijo condition in the open interval (0, 1)
- `scale::T`: the scale factor to shrink the step length
"""
struct BacktrackingLineSearch{C, T}
    c::C
    scale::T

    function BacktrackingLineSearch(c, scale)
        0 < c < 1 || error("`c` must be in the open interval (0, 1)")
        0 < scale < 1 || error("Scale factor must be between 0 and 1")
        return new{typeof(c), typeof(scale)}(c, scale)
    end
end

"""
    BacktrackingLineSearch(; kwargs...)

Initiate `BacktrackingLineSearch`.

# Keywords
- `c=1e-4`: the constant for the Armijo condition in the open interval (0, 1)
- `scale=0.5`: the scale factor to shrink the step length
"""
BacktrackingLineSearch(; c=1e-4, scale=0.5) = BacktrackingLineSearch(c, scale)

function (bls::BacktrackingLineSearch)(f, state, p, α₀::T) where {T}
    x = state.x
    y = state.f
    ∇y = state.∇f
    dϕ₀ = ∇y⋅p
    c = bls.c
    scale = T(bls.scale)
    while f(x + α₀*p) > y + c*α₀*dϕ₀
        α₀ *= scale
    end
    return α₀
end


using .Plots

"""
    animate_draws(uc::UncertainCurve, bnds::Vector{UncertainBound}; n=20, fps=5, filepath=nothing, kw...)

Create a gif animation showing the first `n` random draws.
Requires the Plots.jl package.

Create an animation of the Monte-Carlo iterations when integrating the
`UncertainCurve` `uc` using the `UncertainBound`s `bnds`.

# Keyword arguments

`n`: number of draws that are included in animation

`fps`: frames-per-second

`filepath`: if provided, the gif-animation will be stored using this filepath,
otherwise a temporary file will be created.

Further keyword arguments are passed on the `Plots.plot()` function.
"""
function animate_draws(
    uc::T, bnds::Vector{S}; n=20, fps=5, filepath=nothing, label="", kw...
) where {T<:UncertainCurve, S<:UncertainBound}
    # get minimum and maximum y value of all curve draws
    # in the selected `n` samples, used to determine the
    # y-span when plotting
    miny, maxy = [
        mm(curve.y for curve ∈ [NoisySignalIntegration.get_draw(i, uc) for i ∈ 1:n]) |> mm
        for mm in (minimum, maximum)
    ]
    anim = @animate for i in 1:n
        plot(NoisySignalIntegration.get_draw(i, uc), bnds, i; label=label, ylim=(miny*0.9, maxy*1.1), kw...)
    end
    fp = filepath === nothing ? tempname()*".gif" : filepath
    return gif(anim, fp, fps=fps)
end

export animate_draws


#=
1D Burger's equation using 1st order Godunov flux.
=#

### If the Finch package has already been added, use this line #########
using Finch # Note: to add the package, first do: ]add "https://github.com/paralab/Finch.git"

### If not, use these four lines (working from the examples directory) ###
# if !@isdefined(Finch)
#     include("../Finch.jl");
#     using .Finch
# end
##########################################################################

init_finch("FVburger1d");

useLog("FVburger1dlog", level=3)

# Configuration setup
domain(1)
solverType(FV)
timeStepper(RK4, cfl=1)

# Mesh
n = 70 # number of elements
mesh(LINEMESH, elsperdim=n)

# Variables and BCs
u = variable("u", location=CELL)
boundary(u, 1, NO_BC)

v = variable("v", location=CELL)
boundary(v, 1, NO_BC)

# Time interval and initial condition
T = 1;
timeInterval(T)
initial(v, "sin(pi*2*x)")
initial(u, "x<0.5 ? -1 : 1")

# The flux term of the conservation equation
# See the file fv_ops.jl for the definition of burgerGodunov(u, f)
flux([u,v], ["0.1*burgerGodunov(u,u*u)", "0.1*burgerGodunov(v,v*v)"])

# printLatex(u)

# @exportCode("fvburgercode") # uncomment to export generated code to a file
# @importCode("fvburgercode") # uncomment to import code from a file

solve([u,v])

# output_values([u,v], "burger1d", format="vtk");

finalize_finch()

##### Uncomment below to plot #####

x = Finch.fv_info.cellCenters[:]
n = length(x);
u_ic = zeros(n);
v_ic = zeros(n);
for i=1:n
    u_ic[i] = x[i]<0.5 ? -1 : 1
    v_ic[i] = sin(2*pi*x[i])
end

using Plots
pyplot();
pu = plot([x x], [u.values[:] u_ic], markershape=:circle, label=["Godunov" "initial"])
pv = plot([x x], [v.values[:] v_ic], markershape=:circle, label=["Godunov" "initial"])
display(plot(pu, pv, layout=2))


struct DiscontinuousGalerkin1DGrid{N, C}
                node_coordinates :: N
    element_to_node_connectivity :: C
end

"""
Generate simple equidistant grid with N elements.
"""
function DiscontinuousGalerkin1DGrid(domain, N)
    xₗ, xᵣ = domain

    node_coordinates = [xₗ + (xᵣ - xₗ) * (i - 1) / N for i in 1:N+1]

    # read element to node connectivity
    element_to_node_connectivity = zeros(N, 2)
    for n in 1:N
        element_to_node_connectivity[n, 1] = n
        element_to_node_connectivity[n, 2] = n+1
    end

    return DiscontinuousGalerkin1DGrid{typeof(node_coordinates), typeof(element_to_node_connectivity)}(
        node_coordinates, element_to_node_connectivity)
end


using GridArrays

using Test, LinearAlgebra, DomainSets, Plots, StaticArrays

using GridArrays: MaskedGrid, IndexSubGrid, randompoint,
    component, components, productgrid

using GridArrays.Test: test_generic_grid, test_interval_grid,
    grid_iterator1, grid_iterator2

function delimit(s::AbstractString)
    println()
    println("## ",s)
end

types = (Float64,BigFloat)

include("test_grids.jl")
include("test_gauss.jl")
include("test_periodiccartesianindices.jl")
include("test_boundingbox.jl")
include("test_gauss.jl")
include("test_subgrids.jl")
include("test_randomgrid.jl")
include("test_plots.jl")


struct CompleteGraph{T<:Integer} <: LG.AbstractGraph{T}
    nv::T
end

LG.eltype(::CompleteGraph{T}) where {T} = T
LG.edgetype(::CompleteGraph{T}) where {T} = LG.Edge{T}
LG.is_directed(::Type{<:CompleteGraph}) = false
LG.nv(g::CompleteGraph) = g.nv
LG.ne(g::CompleteGraph) = div(nv(g) * (nv(g)-1), 2)
LG.vertices(g::CompleteGraph) = 1:LG.nv(g)

LG.edges(g::CompleteGraph) = [Edge(i, j) for i in 1:nv(g)-1 for j in i+1:nv(g)]

function LG.outneighbors(g::CompleteGraph, v)
    return [v1 for v1 in vertices(g) if v1 != v]
end

LG.inneighbors(g::CompleteGraph, v) = outneighbors(g, v)

LG.has_vertex(g::CompleteGraph, v) = 1 <= v <= LG.nv(g)

function LG.has_edge(g::CompleteGraph, v1, v2)
    if !has_vertex(g, v1) || !has_vertex(g, v2)
        return false
    end
    return v1 != v2
end

# =======================================================
#         overrides
# =======================================================

# we use this check so that we have the same convention as in LightGraphs
LG.is_connected(g::CompleteGraph) = nv(g) > 0

function LG.connected_components(g::CompleteGraph)

    nv(g) == 0 && return typeof(vertices(g))[]
    return [vertices(g)]
end

# has_self_loops is defined in terms of this
LG.num_self_loops(::CompleteGraph) = 0

LG.is_bipartite(g::CompleteGraph) = (nv(g) <= 2)

function LG.squash(g::CompleteGraph)
    nvg = nv(g)
    for T ∈ (UInt8, UInt16, UInt32, UInt64)
        nv(g) < typemax(T) && return CompleteGraph(T(nvg))
    end
end

# ---- degree -----------------------------------

function LG.Δ(g::CompleteGraph)
    nvg = nv(g)
    return ifelse(nvg == 0, typemin(Int), Int(nvg) - 1)
end

LG.Δout(g::CompleteGraph) = Δ(g)
LG.Δin(g::CompleteGraph) = Δ(g)

function LG.δ(g::CompleteGraph)
    nvg = nv(g)
    return ifelse(nvg == 0, typemax(Int), Int(nvg) - 1)
end

LG.δout(g::CompleteGraph) = δ(g)
LG.δin(g::CompleteGraph) = δ(g)

LG.degree(g::CompleteGraph, vs::AbstractVector=vertices(g)) = Fill(Δ(g), length(vs))

LG.indegree(g::CompleteGraph, vs::AbstractVector=vertices(g)) = degree(g, vs)
LG.outdegree(g::CompleteGraph, vs::AbstractVector=vertices(g)) = degree(g, vs)


### A Pluto.jl notebook ###
# v0.11.14

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 2069c2c0-01b4-11eb-3cc9-a7c60467a92a
# Load packages
begin
	using Pkg
	Pkg.activate(joinpath(Pkg.envdir(), "dev-env"))
	using FractalTools
	using Plots
	plotly()
	default(:label, "")
	default(:lw, 2);
end

# ╔═╡ 24b8a340-01b5-11eb-0551-318932fe4a10
@bind d html"<input type=range min=0 max=1 step=1e-1>"

# ╔═╡ 2caf47f8-047f-11eb-1f15-db8e661dc9ce
@bind niter html"<input type=range min=1 max=10 step=1>"

# ╔═╡ 69dbf05e-047f-11eb-3d6e-892350bec24f
niter, d

# ╔═╡ 402f7b62-01b5-11eb-2789-7bdb04c0d694
begin
	f(xi) = sin(2π * xi)

	# Consrucct data
	xi, dx, xf = 0, 0.1, 1
	x = collect(xi : dx : xf)
	y = sin.(2π * x)
	dvar = d * ones(length(x) - 1)

	# Interplate
	interp = interpolate(x, y, dvar, niter=niter)

	# Plots
	xd = collect(xi : 0.1dx : xf)
	plot(xd, f.(xd))
	yd = interp.(xd)
	plot!(xd, yd)
	scatter!(x, y, marker=(:circle, 3))
end

# ╔═╡ Cell order:
# ╠═2069c2c0-01b4-11eb-3cc9-a7c60467a92a
# ╠═24b8a340-01b5-11eb-0551-318932fe4a10
# ╠═2caf47f8-047f-11eb-1f15-db8e661dc9ce
# ╠═69dbf05e-047f-11eb-3d6e-892350bec24f
# ╠═402f7b62-01b5-11eb-2789-7bdb04c0d694


# phase space plots
using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using DynamicalSystems, PyPlot, Random

cs = COLORS[[1,2,3]]
tscolor = cs[1]

fig = figure(figsize = (figx, figx/2))

ax = subplot(231)

ds = Systems.standardmap(k=1)
ic = [[0.1, 0.1], [2.5, 0.4], [1.88, 3.25]]
xs = range(0, stop = 2π, length = 7);
ys = range(0, stop = 2π, length = 3);
ys = copy(xs)
iters = 2000
for (i, c) in enumerate(ic)
    tr = trajectory(ds, iters, c)
    ax.scatter(columns(tr)..., s = 5, c = cs[i])
end
# ax.set_xticks([0, 2π])
ax.set_xticklabels([])
ax.set_yticklabels([])
# ax.set_yticks([0, 2π])
# ax.set_yticks([0, 2π])
ax.set_yticklabels([])
ax.set_xlabel("\$\\theta\$", labelpad = -5)
ax.set_xlim(0, 2π)
ax.set_ylim(0, 2π)
ax.set_ylabel("\$v\$",labelpad=0)
ax.set_title("Standard map", pad = 25)

axts = subplot(234)
n = 100
tr = trajectory(ds, n, ic[1])
axts.plot(0:n, tr[:, 2], c = tscolor, marker = "o", lw = 1.0)
axts.set_yticks([0, 2π])
axts.set_yticklabels(["0", "2\$\\pi\$"])
axts.set_ylabel("\$v\$",labelpad=-15)
axts.set_xticks([0, n])
axts.set_xlabel("\$n\$", labelpad = -20)

ds = Systems.lorenz()
using3D()
axlo3d = subplot(232, projection = "3d")
axlo3d.set_title("Lorenz-63", pad = 24)

llw = 2.0
tr = trajectory(ds, 100; Ttr = 100, Δt = 0.01)
axlo3d.plot3D(columns(tr)..., color = cs[1], lw = llw/2)
tr = trajectory(ds, 10; Ttr = 1000, Δt = 0.01)
axlo3d.plot3D(columns(tr)..., color = cs[2], lw = llw)
tr = trajectory(ds, 10, [10,20,40.0]; Ttr = 10, Δt = 0.01)
axlo3d.plot3D(columns(tr)..., color = cs[3], lw = llw)
axlo3d.set_xlabel("\$x\$", labelpad = -10)
axlo3d.set_ylabel("\$y\$", labelpad = -10)
axlo3d.set_zlabel("\$z\$", labelpad = -10)
axlo3d.set_xticklabels([])
axlo3d.set_yticklabels([])
axlo3d.set_zticklabels([])
axlo3d.dist = 8

# ax.set_xticks([-20,20])
# ax.set_yticks([-20,20])
# ax.set_zticks([10,40])

tn = 30
axlo2d = subplot(235)
tr = trajectory(ds, tn; Ttr = 40, Δt = 0.01)
axlo2d.plot(0:0.01:tn, tr[:, 1], c = tscolor)

axlo2d.set_yticks([-15,15])
axlo2d.set_ylabel("\$x\$", labelpad = -40)
axlo2d.set_xticks([0, tn])
axlo2d.set_xlabel("\$t\$", labelpad = -20)

ds = Systems.henonheiles()
u0s = (
    [0.0, -0.25, 0.42, 0.0], # chaotic
    [0.0, 0.1, 0.5, 0.0], # quasiperiodic
    [0.0, 0.30266571044921875, 0.4205654433900762, 0.0], # periodic
)

axhe3d = subplot(233, projection = "3d")
axhe3d.clear()
axhe3d.set_title("Hénon–Heiles", pad = 24)
for (i, u0) in enumerate(u0s)
    tr = trajectory(ds, 50, u0; Ttr = 50)
    axhe3d.plot3D(tr[:, 1], tr[:, 2], tr[:, 3], c = COLORS[i], alpha = 1,
    lw = i == 3 ? 4 : 2)
end

axhe3d.set_xticklabels([])
axhe3d.set_yticklabels([])
axhe3d.set_zticklabels([])
axhe3d.set_xlabel("\$x\$"; labelpad = -10)
axhe3d.set_ylabel("\$y\$"; labelpad = -10)
axhe3d.set_zlabel("\$v_x\$",labelpad=-10)
axhe3d.dist = 8

axhe2d = subplot(236)
n = 150
tr = trajectory(ds, n, [0.0, -0.25, 0.48, 0.0]; Ttr = 100)
axhe2d.plot(0:0.01:n, tr[:, 2], c = tscolor)
axhe2d.set_xticks([0, n])
axhe2d.set_yticks([-0.5, 0.5])
axhe2d.set_ylabel("\$y\$", labelpad = -40)
axhe2d.set_xlabel("\$t\$", labelpad = -20)

fig.tight_layout(pad = 0.3)
# WARNING! Run the above command once more before saving!
wsave(plotsdir("1", "trajectories"), fig)


# Script to replicate table 2 (regression from equation 11)

using DataFrames, FixedEffectModels, RegressionTables, StatsBase, Statistics, LinearAlgebra, XLSX, PrettyTables

# -------------- Previous scripts -----------------------------------------------------------------------------------------------------------------------

include("src/data_wrangling.jl") # formats the raw input data obtained from WIOD

# -------------- Country-sector I/O-Table ---------------------------------------------------------------------------------------------------------------

function country_sector_level_GVC_indicators(df::DataFrame, year::Int64)

    ID, FD, GO, VA, IV, other = obtain_matrices(df, year) # N*S×N*S, N*S×N, N*S×1, N*S×1, N*S×N, N*S×1

    # net inventory correction
    iv_correction = [GO[i]/(GO[i]-IV[i, j]) for i in 1:n_ctrys*n_sectors, j in 1:n_ctrys] # N*S×N, to match country wise inventory changes
    # notice difference in column size, need spread inventory correction of country column to country-sector columns
    # i.e. ID[1,1:35]/iv_correction[1,1], however problem with indexing => ceil(Int, 34/35)=1, ceil(Int, 36/35)=2
    ID_corrected = [ID[i, j]*iv_correction[i, ceil(Int, j/n_sectors)] for i in 1:n_ctrys*n_sectors, j in 1:n_ctrys*n_sectors] # N*S×N*S
    FD_corrected = [FD[i, j]*iv_correction[i, j] for i in 1:n_ctrys*n_sectors, j in 1:n_ctrys]
    FD_corrected = [sum(FD_corrected[i, :]) for i in 1:n_ctrys*n_sectors]

    replace!(ID_corrected, NaN => 0.0) # sometimes zeros are interpreted as NaN (particular for sector 35, private households with employed persons)
    replace!(FD_corrected, NaN => 0.0)

    # input/output coefficient matrices
    A = [ID_corrected[i, j] / GO[j] for i in 1:n_ctrys*n_sectors, j in 1:n_ctrys*n_sectors] # N*S×N*S, iterate output over destination country
    B = [ID_corrected[i, j] / GO[i] for i in 1:n_ctrys*n_sectors, j in 1:n_ctrys*n_sectors] # N*S×N*S, iterate output over reporting country

    replace!(A, NaN => 0.0) # sometimes zeros are interpreted as NaN (particular for sector 35, private households with employed persons)
    replace!(B, NaN => 0.0)

    # country-sector level statistics
    FD_GO = FD_corrected ./ GO # N*S×1
    VA_GO = VA ./ GO # N*S×1
    U_num = inv(I - A) * GO # N*S×1, numerator of equation (5)
    D_num = transpose(GO) * inv(I - B) # N*S×1, numerator of equation (7), notice difference in matrix algebra

    U = U_num ./ GO # N*S×1
    D = transpose(D_num) ./ GO # N*S×1, transpose to have column vector

    # collect data in a DataFrame
    summary_stats = hcat(fill(year, n_ctrys*n_sectors), repeat(ctrys, inner=n_sectors), repeat(1:35, outer=n_ctrys), FD_GO, VA_GO, U, D)
    t_country_sector_wide = DataFrame(summary_stats, :auto)
    rename!(t_country_sector_wide, [:year, :country, :sector, :FD_GO, :VA_GO, :U, :D])
    transform!(t_country_sector_wide, [:FD_GO, :VA_GO, :U, :D] .=> ByRow(x -> round(x, digits=3)) .=> [:FD_GO, :VA_GO, :U, :D])

    return t_country_sector_wide
end

# gather data for all years in a single data frame
t_tab2 = DataFrame(year=Int[], country=String[], sector=Int[], FD_GO=Float64[], VA_GO=Float64[], U=Float64[], D=Float64[])

for i in years
    append!(t_tab2, country_sector_level_GVC_indicators(df, i))
end

XLSX.writetable("clean/t_tab2.xlsx", t_tab2, overwrite=true) # export table

# -------------- Table A1 --------------------------------------------------------------------------------------------------------------------------------

# import data from local drive
t_tab2 = DataFrame(XLSX.readtable("clean/t_tab2.xlsx", "Sheet1")...) # still a lot of NaN
transform!(t_tab2, :country .=> ByRow(string) .=> :country, renamecols=false) # change type to String
transform!(t_tab2, [:year, :sector] .=> ByRow(Int) .=> [:year, :sector], renamecols=false) # change type to Int64
transform!(t_tab2, [:FD_GO, :VA_GO, :U, :D] .=> ByRow(Float64) .=> [:FD_GO, :VA_GO, :U, :D], renamecols=false) # change type to Float64

# missing: create table A.1 with summary statistics of country-industry GVC measures (how to treat NaN?)
# treat NaN = 0
t_tab2_long = stack(t_tab2, Not([:year, :country, :sector]))
transform!(t_tab2_long, :value => ByRow(x -> ifelse(isnan(x), 0.0, x)) => :value, renamecols=false) # set NaN = 0.0
transform!(t_tab2_long, :sector => ByRow(x -> ifelse(x in goods_sectors, "goods", "services")) => :sector_dummy, renamecols=false) # sector dummy for grouping
gdf = groupby(t_tab2_long, :variable)

t_tabA1_1 = combine(gdf, nrow, :value .=> [x -> percentile(x, 10), median, x -> percentile(x, 90), mean, std] .=> [:perc_10, :median, :perc_90, :mean, :std])
t_tabA1_1[:, :sector_dummy] .= "all industries"

gdf = groupby(t_tab2_long, [:variable, :sector_dummy])
t_tabA1_2 = combine(gdf, nrow, :value .=> [x -> percentile(x, 10), median, x -> percentile(x, 90), mean, std] .=> [:perc_10, :median, :perc_90, :mean, :std])

t_tabA1 = vcat(t_tabA1_2, t_tabA1_1) # combine tables
sort!(t_tabA1, [:variable, :sector_dummy])

open("images/summary_stats_country_industry_GVC.txt", "w") do f
    pretty_table(f,t_tabA1; backend = Val(:html), nosubheader=true, standalone=false) # only prints html code for table itself
end

# -------------- Table 2 --------------------------------------------------------------------------------------------------------------------------------

# use package: FixedEffectModels

# different ways of specifying the same regression, i.e. using either @formula, Symbols, Strings
# FixedEffectModels.reg(data_reg, @formula(FD_GO ~ year + fe(country_sector)), Vcov.cluster(:year, :country_sector))
# FixedEffectModels.reg(data_reg, term(:VA_GO) ~ term(:year) + fe(:country_sector), Vcov.cluster(:year, :country_sector))
# FixedEffectModels.reg(data_reg, term("VA_GO") ~ term("year") + fe(term("country_sector")), Vcov.cluster(:year, :country_sector))

#data_reg = subset(t_tab2, :VA_GO => ByRow(x -> !isnan(x))) # filter out rows with NaN
data_reg = transform(t_tab2, [:FD_GO, :VA_GO, :U, :D] .=> ByRow(x -> ifelse(isnan(x), 0.0, x)) .=> [:FD_GO, :VA_GO, :U, :D], renamecols=false) # set NaN = 0.0
# setting NaN = 0 considerably improves match with table 2
data_reg[:, :country_sector] .= data_reg[:, :country] .* "_" .* string.(data_reg[:, :sector]) # considerably improves R² compared to use columns individually

rr_FD_GO_1 = FixedEffectModels.reg(data_reg, @formula(FD_GO ~ year + fe(country_sector)), Vcov.cluster(:country_sector), save=true)
rr_FD_GO_2 = FixedEffectModels.reg(data_reg, @formula(FD_GO ~ year + fe(country_sector)), contrasts = Dict(:year => DummyCoding(base=1995)),
 Vcov.cluster(:country_sector), save=true)
rr_VA_GO_1 = FixedEffectModels.reg(data_reg, @formula(VA_GO ~ year + fe(country_sector)), Vcov.cluster(:country_sector), save=true)
rr_VA_GO_2 = FixedEffectModels.reg(data_reg, @formula(VA_GO ~ year + fe(country_sector)), contrasts = Dict(:year => DummyCoding(base=1995)),
 Vcov.cluster(:country_sector), save=true)
rr_U_1 = FixedEffectModels.reg(data_reg, @formula(U ~ year + fe(country_sector)), Vcov.cluster(:country_sector), save=true)
rr_U_2 = FixedEffectModels.reg(data_reg, @formula(U ~ year + fe(country_sector)), contrasts = Dict(:year => DummyCoding(base=1995)),
 Vcov.cluster(:country_sector), save=true)
rr_D_1 = FixedEffectModels.reg(data_reg, @formula(D ~ year + fe(country_sector)), Vcov.cluster(:country_sector), save=true)
rr_D_2 = FixedEffectModels.reg(data_reg, @formula(D ~ year + fe(country_sector)), contrasts = Dict(:year => DummyCoding(base=1995)),
 Vcov.cluster(:country_sector), save=true)

# use package: RegressionTables

# create .png out of ascii output to include in README.md as unformatted html output does not look good
RegressionTables.regtable(rr_FD_GO_1, rr_FD_GO_2, rr_VA_GO_1, rr_VA_GO_2, rr_U_1, rr_U_2, rr_D_1, rr_D_2;
 renderSettings = asciiOutput(), regression_statistics=[:nobs, :r2], print_fe_section=true, estimformat="%0.4f") # output on the console

# html output
RegressionTables.regtable(rr_FD_GO_1, rr_FD_GO_2, rr_VA_GO_1, rr_VA_GO_2, rr_U_1, rr_U_2, rr_D_1, rr_D_2;
 renderSettings = htmlOutput("images/table2_html.txt"), regression_statistics=[:nobs, :r2], print_fe_section=true, estimformat="%0.4f")

# regressions coincide very well, possible source of difference is the additional 300 observations of Antras and Chor (2018)
# have for VA, U, D (i.e. must have fewer NaN, but how?) => set NaN = 0


using DistStat, LinearAlgebra,Random,Test, Pkg

type=[Float64,Float32]

if get(ENV,"JULIA_MPI_TEST_ARRAYTYPE","") == "CuArray"
    using CUDA
    A= CuArray
else
    A = Array
end

for T in type
    a=A{T}(reshape(collect(1:81),9,9))
    a_dist=distribute(a)

    @test isapprox(opnorm(a),opnorm(a_dist))

end


#=
An optimization oracle is a linear optimization oracle for the set S, i.e., it solves
    z^\ast(c) :=    min c^Tw
                    s.t. w \in S
The solution of the above problem returned by the oracle is referred to as w^\ast(\cdot) in
the paper

A specific implementation of an optimization oracle should take a single input argument c,
and should return the pair (z^\ast(c), w^\ast(c))

This file contains the code for the portfolio problem optimization oracle, which
relies on the JuMP and Gurobi packages.
=#

using JuMP, Gurobi, LinearAlgebra

"""
    portfolio_simplex_oracle_jump_basic(c, Sigma, gamma)

Optimization oracle in JuMP for the portfolio problem:
    min  c^T w
    s.t. w^T Sigma w <= gamma
         e^T w <= 1
         w >= 0
Here c is equal to minus the vector of returns, Sigma is the covariance matirx of the returns,
and gamma is the desired risk level in terms of the overall variance of the portfolio. Short
selling is not allowed.
"""
function portfolio_simplex_oracle_jump_basic(c::Vector{Float64}, Sigma::Matrix{Float64}, gamma::Float64)
    d = length(c)

    mod = Model(with_optimizer(Gurobi.Optimizer, OutputFlag = 0))
    @variable(mod, w[1:d] >= 0)
    @constraint(mod, sum(w[i] for i = 1:d) <= 1)
    @constraint(mod, w'*Sigma*w <= gamma)

    @objective(mod, Min, dot(c, w))

    optimize!(mod)
    z_ast = objective_value(mod)
    w_ast = value.(w)

    return (z_ast, w_ast)
end

"""
Smarter way of implementing the oracle in JuMP. Constructs the feasible region from
(Sigma, gamma) ahead of time and then returns an oracle that is just a function of c.
"""
function portfolio_simplex_jump_setup(Sigma::Matrix{Float64}, gamma::Float64; gurobiEnv = Gurobi.Env())
    (d, d2) = size(Sigma)

    if d != d2
        error("Sigma dimensions don't match")
    end

    mod = Model(with_optimizer(Gurobi.Optimizer, gurobiEnv))
    @variable(mod, w[1:d] >= 0)
    @constraint(mod, sum(w[i] for i = 1:d) <= 1)
    @constraint(mod, w'*Sigma*w <= gamma)

    function local_portfolio_oracle(c::Vector{Float64})
        @objective(mod, Min, dot(c, w))

        optimize!(mod)
        z_ast = objective_value(mod)
        w_ast = value.(w)

        return (z_ast, w_ast)
    end

    return c -> local_portfolio_oracle(c)
end


using Test

using LinearAlgebra

using CLIMA.TurbulenceConvection.TriDiagSolvers
TDMA = TriDiagSolvers

@testset "TriDiagSolvers" begin
  N = 2:10
  for n in N
    dl = rand(n-1)
    du = rand(n-1)
    d = rand(n);

    A = Array(Tridiagonal(dl, d, du))
    b = rand(length(d))

    x_correct = inv(A)*b

    xtemp = zeros(n)
    gamma = zeros(n-1)
    beta = zeros(n)
    x_TDMA = zeros(n)

    solve_tridiag!(x_TDMA, b, dl, d, du, n, xtemp, gamma, beta)
    tol = eps(Float32)

    err = [abs(x-y) for (x, y) in zip(x_correct, x_TDMA)]
    # @show eps(Float64)
    # @show "Full", err
    if !all([x<tol for x in err])
        @show tol
        @show "Full", err
    end
    @assert all([x<tol for x in err])

    init_β_γ!(beta, gamma, dl, d, du, n)
    solve_tridiag_stored!(x_TDMA, b, dl, beta, gamma, n, xtemp)

    err = [abs(x-y) for (x, y) in zip(x_correct, x_TDMA)]
    # @show "Stored", err
    @assert all([x<tol for x in err])

    dl_mod = zeros(n)
    dl_mod[2:end] = dl
    du_mod = zeros(n)
    du_mod[1:end-1] = du
    TDMA.solve_tridiag_old(n, b, dl_mod, d, du_mod)
    x_TDMA = b
    err = [abs(x-y) for (x, y) in zip(x_correct, x_TDMA)]
    # @show "old", err
    @assert all([x<tol for x in err])

  end
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    BlockWeighting(sides)
    BlockWeighting(side₁, side₂, ...)

A weighting method that assigns weights to points in spatial object
based on blocks of given `sides`. The number `n` of points inside a
block determines the weights `1/n` of these points.
"""
struct BlockWeighting{T,N} <: WeightingMethod
  sides::SVector{N,T}
end

BlockWeighting(sides::NTuple{N,T}) where {N,T} =
  BlockWeighting{T,N}(sides)

BlockWeighting(sides::Vararg) = BlockWeighting(sides)

weight(object, method::BlockWeighting) =
  weight(domain(object), method)

function weight(object::AbstractDomain, method::BlockWeighting)
  p = partition(object, BlockPartition(method.sides))

  weights = Vector{Float64}(undef, nelms(object))
  for s in subsets(p)
    n = length(s)
    weights[s] .= 1/n
  end

  SpatialWeights(object, weights)
end


#=
Try optim in high dimension and matrix
=#
using Optim
using BenchmarkTools

function f(x)
    return (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2
end

function g!(G, x)
    G[1] = -2.0 * (1.0 - x[1]) - 400.0 * (x[2] - x[1]^2) * x[1]
    G[2] = 200.0 * (x[2] - x[1]^2)
end

function h!(H, x)
    H[1, 1] = 2.0 - 400.0 * x[2] + 1200.0 * x[1]^2
    H[1, 2] = -400.0 * x[1]
    H[2, 1] = -400.0 * x[1]
    H[2, 2] = 200.0
end

initial_x = zeros(2)

# 27μs
@btime Optim.minimizer(optimize(f, g!, h!, initial_x, BFGS()))
# 39μs
@btime Optim.minimizer(optimize(f,initial_x, BFGS()))
# 31μs
@btime Optim.minimizer(optimize(f, initial_x, BFGS(); autodiff = :forward))


using SparseArrays

"""
Sparse N-mode array.

`vals` contains all (`L`) non-zero elements.
`pos` is `N`x`L` matrix of non-zero elements indices.
`dims` is `NTuple` of array dimensions.
"""
struct SparseArray{T, N} <: AbstractArray{T, N}
    vals::Vector{T}
    pos::Vector{CartesianIndex{N}}
    dims::NTuple{N,Int}
end

Base.size(a::SparseArray) = a.dims
Base.size(a::SparseArray, i::Int) = a.dims[i]
Base.length(a::SparseArray) = prod(a.dims)
SparseArrays.nnz(a::SparseArray) = length(a.vals)

"""
Convert dense N-dimensional array into N-dimensional a `SparseArray`.
"""
function SparseArray{T,N}(arr::DenseArray{T,N}) where {T,N}
    pos = Vector{CartesianIndex{N}}()
    vals = Vector{T}()

    @inbounds for (i, cart) in enumerate(CartesianIndices(arr))
        (arr[i] == 0) && continue
        push!(vals, arr[i])
        push!(pos, cart)
    end

    return SparseArray{T,N}(vals, pos, size(arr))
end

SparseArray(arr::DenseArray{T, N}) where {T, N} = SparseArray{T,N}(arr)

"""
SS-HOPM (Shifted Symmetric Higher-order Power Method)
for computing tensor eigenpars.
"""
function sshopm(tnsr::AbstractArray{T,N},
                alpha::Real;
                tol::Float64=1e-5,
                maxiter::Int=1000,
                verbose::Bool=false) where {T,N}

    r = size(tnsr, 1)
    all(isequal(r), size(tnsr)) || throw(DimensionMismatch("Input tensor should have the same dimension in all modes, got $(size(tnsr))."))
    x = randn(r)
    x .*= 1/norm(x)
    x_old = similar(x)
    converged = false
    niters = 0
    pb = Progress(maxiter, "High-order power iterations ")
    while !converged && niters < maxiter
        update!(pb, niters)

        copyto!(x_old, x)
        x .= flipsign.(nmul(tnsr, x) .+ alpha .* x, alpha)
        x .*= 1/norm(x)
        converged = norm(x .- x_old) < tol
        niters += 1
    end

    verbose && _iter_status(converged, niters, maxiter)
    return (dot(x, nmul(tnsr, x)), flipsign.(x, alpha))
end

# (N-1)-way tensor × vector contraction
function nmul(tnsr::StridedArray{T,N}, x::Vector{T}) where {T, N}
    v = copy(tnsr)
    for i in 2:N
        v = tensorcontract(v, collect(i-1:N), x, N)
    end
    return v
end

# (N-1)-way sparse tensor × vector contraction
function nmul(tnsr::SparseArray{T,N}, x::Vector{T}) where {T, N}
    v = zeros(T, size(tnsr, 1))
    @inbounds for i in 1:nnz(tnsr)
        posi = tnsr.pos[i]
        v[posi[1]] += tnsr.vals[i] * prod(ntuple(k -> x[posi[k+1]], N-1))
    end
    return v
end


# Speed test for decomposing matrices in Julia.

function speed_test(func, orderN = 1000)
  sizes = Any[]
  for i in 0:int(log10(orderN))
    push!(sizes, 1 * 10^i)
    push!(sizes, 2 * 10^i)
    push!(sizes, 5 * 10^i)
  end

  y = zeros(length(sizes))
  for i in 1:length(sizes)
    n = sizes[i]
    A = randn(n, n)
    A = A'*A
    print("Testing with n = $(n)")
    tic()
    res = func(A)
    y[i] = toq()
    println(", time = $(y[i]), ")
  end

  return sizes, y
end

n_eig, t_eig = speed_test(eig, 1000)
n_chol, t_chol = speed_test(chol, 1000)

# Linear model
lin_eig = linreg(t_eig, n_eig)


module PrinzSingleCompartment

using Test, OrdinaryDiffEq

include(normpath(@__DIR__, "..", "demo", "prinz_kinetics.jl"))

@testset "Prinz Single-Compartment Neuron" begin

channels = [NaV(100mS/cm^2),
            CaT(0mS/cm^2),
            CaS(4mS/cm^2),
            KA(0mS/cm^2),
            KCa(15mS/cm^2),
            Kdr(50mS/cm^2),
            H(.02mS/cm^2),
            leak(.03mS/cm^2)]

@named neuron = Soma(channels, gradients, area = area, V0 = -50mV, aux = [calcium_conversion]);

@test length.([equations(neuron.sys),
               states(neuron.sys),
               parameters(neuron.sys)]) == [39,40,19]

t = 2000
simul_sys = Simulation(neuron, time = t*ms, system = true)

@test length.([equations(simul_sys),
               states(simul_sys),
               parameters(simul_sys)]) == [15,15,17]

# Prinz STG neuron hand-written reference implementation

# Steady-state functions
# Input vectors
V0 = -50.0
u0 = zeros(Float64, 13)
p  = zeros(Float64,17)

# Initial state
u0[1]  = 0.05 # Caᵢ
u0[2]  = V0 # Vₘ
u0[3]  = (1.0 + exp(-4.82042 - (0.189036V0)))^-1 # NaV₊m
u0[4]  = (1.0 + exp(9.44015 + 0.19305V0))^-1 # NaV₊h
u0[5]  = (1.0 + exp(-3.76389 - (0.138889V0)))^-1 # CaT₊m
u0[6]  = (1.0 + exp(5.83636 + 0.181818V0))^-1 # CaT₊h
u0[7]  = (1.0 + exp(-4.07407 - (0.123457V0)))^-1 # CaS₊m
u0[8]  = (1.0 + exp(9.67742 + 0.16129V0))^-1 # CaS₊h
u0[9]  = (1.0 + exp(-3.12644 - (0.114943V0)))^-1 # KA₊m
u0[10] = (1.0 + exp(11.6122 + 0.204082V0))^-1 # KA₊h
u0[11] = 0.05*((3.0 + 0.05)^-1)*((1.0 + exp(-2.24603 - (0.0793651V0)))^-1) # KCa₊m
u0[12] = (1.0 + exp(-1.04237 - (0.0847458V0)))^-1 # Kdr₊m
u0[13] = (1.0 + exp(13.6364 + 0.181818V0))^-1 # H₊m

# Parameters
p[1]  = 0.001 # cₘ
p[2]  = 0.000628 # aₘ
p[3]  = 200 # τCa
p[4]  = 0.05 # Ca∞
p[5]  = 14.96 # f
p[6]  = 50.0 # ENa
p[7]  = -80.0 # EK
p[8]  = -20.0 # EH
p[9]  = -50.0 # El
p[10] = 100.0 # NaV₊gbar
p[11] = 0.0 # CaT₊gbar
p[12] = 4.0 # CaS₊gbar
p[13] = 0.0 # KA₊gbar
p[14] = 15.0 # KCa₊gbar
p[15] = 50.0 # Kdr₊gbar
p[16] = 0.02 # H₊gbar
p[17] = 0.03 # leak₊g

function prinz_neuron!(du, u, p, t)
    @inbounds let Caᵢ         = u[1],
                  Vₘ          = u[2],
                  NaV₊m       = u[3],
                  NaV₊h       = u[4],
                  CaT₊m       = u[5],
                  CaT₊h       = u[6],
                  CaS₊m       = u[7],
                  CaS₊h       = u[8],
                  KA₊m        = u[9],
                  KA₊h        = u[10],
                  KCa₊m       = u[11],
                  Kdr₊m       = u[12],
                  H₊m         = u[13],

                  cₘ          = p[1],
                  aₘ          = p[2],
                  τCa         = p[3],
                  Ca∞         = p[4],
                  f           = p[5],
                  ENa         = p[6],
                  EK          = p[7],
                  EH          = p[8],
                  El          = p[9],
                  NaV₊gbar    = p[10],
                  CaT₊gbar    = p[11],
                  CaS₊gbar    = p[12],
                  KA₊gbar     = p[13],
                  KCa₊gbar    = p[14],
                  Kdr₊gbar    = p[15],
                  H₊gbar      = p[16],
                  leak₊g      = p[17],
                  ECa = (Vₘ + 12.199306599503455 * log(0.0003333333333333333 * max(Caᵢ, 0.001)))
        let
            du[1] =
                (
                    Ca∞ - Caᵢ - 1000 * f *
                    (CaS₊gbar * aₘ * CaS₊h * CaS₊m^3 *
                    ECa +
                    CaT₊gbar * aₘ * CaT₊h * CaT₊m^3 *
                    ECa)
                   ) * inv(τCa)

            du[2] =
                (
                    -aₘ * leak₊g * (Vₘ - El) +
                    -H₊gbar * aₘ * H₊m * (Vₘ - EH) +
                    -KCa₊gbar * aₘ * KCa₊m^4 * (Vₘ - EK) +
                    -Kdr₊gbar * aₘ * Kdr₊m^4 * (Vₘ - EK) +
                    -CaS₊gbar * aₘ * CaS₊h * CaS₊m^3 * ECa +
                    -CaT₊gbar * aₘ * CaT₊h * CaT₊m^3 * ECa +
                    -KA₊gbar * aₘ * KA₊h * KA₊m^3 * (Vₘ - EK) +
                    -NaV₊gbar * aₘ * NaV₊h * NaV₊m^3 * (Vₘ - ENa)
                ) * inv(aₘ) * inv(cₘ)

            du[3] =
                (-NaV₊m + inv(1.0 + exp(-4.8204158790170135 - 0.1890359168241966Vₘ))) *
                inv(2.64 - 2.52 * inv(1 + exp(-4.8 - 0.04Vₘ)))

            du[4] =
                0.7462686567164178 * (1.0 + exp(-6.29 - 0.1Vₘ)) *
                (-NaV₊h + inv(1.0 + exp(9.44015444015444 + 0.19305019305019305Vₘ))) *
                inv(1.5 + inv(1.0 + exp(9.694444444444445 + 0.2777777777777778Vₘ)))

            du[5] =
                (-CaT₊m + inv(1.0 + exp(-3.7638888888888893 + -0.1388888888888889Vₘ))) *
                inv(43.4 - 42.6 * inv(1.0 + exp(-3.321951219512195 - 0.04878048780487805Vₘ)))

            du[6] =
                (-CaT₊h + inv(1.0 + exp(5.836363636363637 + 0.18181818181818182Vₘ))) *
                inv(210.0 - 179.6 * inv(1.0 + exp(-3.2544378698224854 - 0.0591715976331361Vₘ)))

            du[7] =
                (-CaS₊m + inv(1.0 + exp(-4.074074074074074 - 0.1234567901234568Vₘ))) *
                inv(2.8 + 14.0 * inv(exp(-5.384615384615385 - 0.07692307692307693Vₘ) + exp(2.7 + 0.1Vₘ)))

            du[8] =
                (-CaS₊h + inv(1.0 + exp(9.67741935483871 + 0.16129032258064516Vₘ))) *
                inv(120.0 + 300.0 * inv(exp(-4.0625 - 0.0625Vₘ) +
                                        exp(6.111111111111111 + 0.1111111111111111Vₘ)))

            du[9] =
                (-KA₊m + inv(1.0 + exp(-3.1264367816091956 - 0.1149425287356322Vₘ))) *
                inv(23.2 - 20.8 * inv(1.0 + exp(-2.164473684210526 - 0.06578947368421052Vₘ)))

            du[10] =
                (-KA₊h + inv(1.0 + exp(11.612244897959181 + 0.2040816326530612Vₘ))) *
                inv(77.2 - 58.4 * inv(1.0 + exp(-1.4679245283018867 - 0.03773584905660377Vₘ)))
            du[11] =
                (-KCa₊m + Caᵢ * inv(3.0 + Caᵢ) *
                inv(1.0 + exp(-2.246031746031746 - 0.07936507936507936Vₘ))) *
                inv(180.6 + -150.2 * inv(1.0 + exp(-2.026431718061674 - 0.04405286343612335Vₘ)))

            du[12] =
                (-Kdr₊m + inv(1.0 + exp(-1.0423728813559323 - 0.0847457627118644Vₘ))) *
                inv(14.4 - 12.8 * inv(1.0 + exp(-1.4739583333333335 - 0.052083333333333336Vₘ)))

            du[13] =
                1 // 2 *
                (exp(-1.8671328671328669 + 0.06993006993006992Vₘ) +
                 exp(-14.629310344827585 - 0.08620689655172414Vₘ)) *
                (-H₊m + inv(1.0 + exp(13.636363636363637 + 0.18181818181818182Vₘ)))
            nothing
        end
    end
end

byhand_prob = ODEProblem{true}(prinz_neuron!, u0, (0.,2000.), p)
mtk_prob = ODAEProblem(simul_sys, [], (0., t), [])
byhand_sol = solve(byhand_prob, Rosenbrock23(), reltol=1e-8, abstol=1e-8, saveat=0.025);
current_mtk_sol = solve(mtk_prob, Rosenbrock23(), reltol=1e-8, abstol=1e-8, saveat=0.025);

tsteps = 0.0:0.025:2000.0
byhand_out = Array(byhand_sol(tsteps, idxs=2))
current_mtk_out = current_mtk_sol(tsteps)[Vₘ]

@test isapprox(byhand_out, current_mtk_out, rtol=0.001)

end # testset
end # module


using Random

const Torus = Int32

# Float64の値を[-0.5,0.5)に射影
FLoat2Torus05(x::Float64) = x - floor(x-0.5) - 1.0

# Float64の値を[0.0, 1.0)に射影
Float2Torus01(x::Float64) = x - floor(x)

# [-0.5,0.5)をTorus32に変換
dtot32(d::Float64) = trunc(Int32, d*2^32)

function gaussian_float(rng::AbstractRNG, sigma::Float64, dims...)
    randn(rng, dims...) .* sigma
end

# モジュラー正規分布に従う乱数生成
function modular_gaussian(rng::AbstractRNG, sigma::Float64, dims...)
    dtot32(Float2Torus05(gaussian_float(rng, sigma, dims...)))
end

# Torus32上の一様分布に従う乱数生成
function rand_uniform_Torus32(rng::AbstractRNG, dims...)
    rand(rng, Torus32, dims...)
end


function solution_error(num_pts, u_extended)
    x = range(0,1,length=num_pts)
    y = range(0,1,length=num_pts)'

    u_true = zeros((num_pts, num_pts));

    for ix = 1:num_pts
        for iy = 1:num_pts
            u_true[ix,iy]=u(x[ix],y[iy])
        end
    end

    error = u_true-u_extended;
    error = error[:]
    return x,y, error
end


"""
    mle(model, θ, vc=" ")

Maximum likelihood estimation.

* model should be a function that gives the vector of log likelihoods of
the observations
* θ is the parameter vector
* vc is the method for computing the variance, default is sandwich
* execute mleresults() for an example, or edit(mleresults,()) to see the code.

"""
function mle(model, θ, vc=1)
    avg_obj = θ -> -mean(vec(model(θ))) # average log likelihood
    thetahat, objvalue, converged = fminunc(avg_obj, θ) # do the minimization of -logL
    objvalue = -objvalue
    obj = θ -> vec(model(θ)) # unaveraged log likelihood
    n = size(obj(θ),1) # how many observations?
    scorecontrib = Calculus.jacobian(obj, vec(thetahat), :central)
    I = cov(scorecontrib)
    J = Calculus.hessian(avg_obj, vec(thetahat), :central)
    Jinv = inv(J)
    if vc==2
        V = Jinv/n      # other possibilities
    elseif vc==3
        V = inv(I)/n
    else
        V= Jinv*I*Jinv/n # sandwich form is preferred
    end
    return thetahat, objvalue, V, converged
end

function mle()
    println("mle(), with no arguments, runs mleresults(), which show usage")
    mleresults()
end


"""
  log_meixner_ortho_poly(a, b, x)

Compute the logarithm of the Meixner polynomial
``M_{b}(a; -(a + b - 1), exp(x))`` where `a` and `b` are positive integers, with
`a >= b`.

Define
``u_{n} = 1 + (1 + a (1 - exp(x)) - n) / ((c + 1 - n) exp(x))``
and
``v_{n} = (n - 1) / ((c + 1 - n) exp(x))``

If ``y_{n}`` represents the Meixner polynomial of degree `n`, then
``y_{n} = u_{n} y_{n - 1} + v_{n} y_{n - 2}``
where ``y_{0} = 1`` and ``y_{1} = 1 + a * (1 - exp(x)) / (c * exp(x))``
"""
function log_meixner_ortho_poly(
  a::F,
  b::F,
  x::F
)::F where {
  F <: AbstractFloat
}
  lom = log1mexp(-x)
  lem = logexpm1(x)

  # define H_{b} = y_{b} / y_{b - 1}, i. e. y_{b} = H_{b} * y_{b - 1}
  # y_{0} = 1, y_{1} = 1 + a * (1 - x) / ((a + b - 1) * x)
  tmp = log1mexp(log(a / (a + b - 1)) + lom)

  # v = [log(H), log(y), compensation_error]
  v = [tmp, tmp, F(0)]

  n = F(1)
  while n < b
    n += 1
    kbn_meixner_ortho_poly_iteration!(v, a, n, a + b - 1, lom, lem)
  end

  # finally add the compensation error
  v[2] + v[3]
end

"""
  log_meixner_ortho_poly_joint(a, b, x)

Compute the logarithm of the Meixner polynomials
``M_{b}(a; -(a + b - 1), exp(x))``,
``M_{b - 1}(a - 1; -(a + b - 2), exp(x))``, and
``M_{b - 2}(a - 2; -(a + b - 3), exp(x))`` at the same time. `a` and `b` are
positive integers, `a >= b`. We assume that `a` and `b` are both greater than 1.
"""
function log_meixner_ortho_poly_joint(
  a::F,
  b::F,
  x::F
)::Tuple{F, F, F} where {
  F <: AbstractFloat
}
  lom = log1mexp(-x)
  lem = logexpm1(x)

  # define H_{b} = y_{b} / y_{b - 1}, i. e. y_{b} = H_{b} * y_{b - 1}
  # y_{0} = 1, y_{1} = 1 + a * (1 - x) / (c * x)

  # M_{b}(a; -(a + b - 1), x)
  c1 = a + b - 1
  tmp = log1mexp(log(a / c1) + lom)
  v1 = [tmp, tmp, F(0)]

  # M_{b - 1}(a - 1; -(a + b - 2), x)
  c2 = a + b - 2
  tmp = log1mexp(log((a - 1) / c2) + lom)
  v2 = [tmp, tmp, F(0)]

  if b > 2
    # M_{b - 2}(a - 2; -(a + b - 3), x)
    c3 = a + b - 3
    tmp = log1mexp(log((a - 2) / c3) + lom)
    v3 = [tmp, tmp, F(0)]

    n = F(1)
    while n < b - 2
      n += 1
      kbn_meixner_ortho_poly_iteration!(v1, a, n, c1, lom, lem)
      kbn_meixner_ortho_poly_iteration!(v2, a - 1, n, c2, lom, lem)
      kbn_meixner_ortho_poly_iteration!(v3, a - 2, n, c3, lom, lem)
    end

    # We still need few iterations
    kbn_meixner_ortho_poly_iteration!(v1, a, b - 1, c1, lom, lem)
    kbn_meixner_ortho_poly_iteration!(v2, a - 1, b - 1, c2, lom, lem)

    kbn_meixner_ortho_poly_iteration!(v1, a, b, c1, lom, lem)

    # add compensation errors and return the final values
    v1[2] + v1[3], v2[2] + v2[3], v3[2] + v3[3]
  else
    # M_{2}(a; -(a + 1), x)
    # M_{1}(a - 1; -(a + 2), x)
    # M_{0}(a - 2; -(a + 3), x)
    kbn_meixner_ortho_poly_iteration!(v1, a, F(2), c1, lom, lem)
    v1[2], v2[2], 0
  end
end

"""
    kbn_meixner_ortho_poly_iteration!(v, a, n, c, x)

Basic Kahan iteration operations for computing Meixner orthogonal polynomial.
"""
@inline function kbn_meixner_ortho_poly_iteration!(
  v::Vector{F},
  a::F,
  n::F,
  c::F,
  # lom = log(1 - exp(-x))
  lom::F,
  # lem = log(exp(x) - 1)
  lem::F
) where {
  F <: AbstractFloat
}
  # new value for log(y_{n} / y_{n - 1})
  v[1] = log1mexp(log(a / (c + 1 - n)) + lom +
                  log1mexp(log((n - 1) / a) + logexpm1(-v[1]) - lem))

  # running sum including error
  s = v[2] + v[1]

  # Neumaier modification to Kahan summation
  v[3] += if abs(v[2]) >= abs(v[1])
    (v[1] - s) + v[2]
  else
    (v[2] - s) + v[1]
  end

  v[2] = s

  nothing
end


using QuantumCumulants
using ModelingToolkit, OrdinaryDiffEq
using Test

@testset "two-level-laser" begin

N = 10

Δ = cnumbers((Symbol(:Δ_, i) for i=1:N)...)
g = cnumbers((Symbol(:g_, i) for i=1:N)...)
γ = cnumbers((Symbol(:γ_, i) for i=1:N)...)
ν = cnumbers((Symbol(:ν_, i) for i=1:N)...)
@cnumbers κ

h_cavity = FockSpace(:cavity)
h_atoms = [NLevelSpace(Symbol(:atom, i),(:g,:e)) for i=1:N]
h = tensor(h_cavity, h_atoms...)

a = Destroy(h, :a)
σ(i,j,k) = Transition(h,Symbol(:σ_,k),i,j,k+1)

ops = begin
    ops_ = [a'*a;[σ(:e,:e,k) for k=1:N];[a'*σ(:g,:e,k) for k=1:N]]
    for i=1:N
        for j=i+1:N
            push!(ops_, σ(:e,:g,i)*σ(:g,:e,j))
        end
    end
    ops_
end

n_eqs = div(N*(N-1),2) + 2N + 1

H = sum(Δ[i]*σ(:e,:e,i) for i=1:N) + sum(g[i]*(a'*σ(:g,:e,i) + a*σ(:e,:g,i)) for i=1:N)

J = [a;[σ(:g,:e,k) for k=1:N];[σ(:e,:g,k) for k=1:N]]
Jdagger = adjoint.(J)
rates = [κ,γ...,ν...]
he = meanfield(ops, H, J; Jdagger=Jdagger, rates=rates, simplify=true, order=2)

missed = find_missing(he)

subs = Dict(missed .=> 0)

he_nophase = substitute(he, subs)

@test isempty(find_missing(he_nophase))

eqs_mtk = equations(he_nophase)

sys = ODESystem(he_nophase)

u0 = zeros(ComplexF64, length(ops))
p0 = [κ => 1,
    (γ .=> 0.25 .* ones(N))...,
    (ν .=> 4 .* ones(N))...,
    (g .=> 1.5 .* ones(N))...,
    (Δ .=> ones(N))...,]

prob = ODEProblem(sys,u0,(0.0,10.0),p0,jac=true,sparse=true)

sol = solve(prob,RK4())

n = real.(sol[average(a'*a)])
pe = sol[average(σ(:e,:e,1))]

# Test with complete and custom filter
ϕ(::Destroy) = -1
ϕ(::Create) = 1
function ϕ(t::Transition)
    if t.i != t.j
        t.i == :e && return 1
        return -1
    else
        return 0
    end
end
function ϕ(q::QuantumCumulants.QMul)
    p = 0
    for arg ∈ q.args_nc
        p += ϕ(arg)
    end
    return p
end
ϕ(avg::Average) = ϕ(avg.arguments[1])
phase_invariant(x) = iszero(ϕ(x))

he_n = meanfield(a'*a, H, J; rates=rates)
complete!(he_n;filter_func=phase_invariant)

@test length(he_n.equations) == length(ops)
@test isempty(find_missing(he_n))

sys_comp = ODESystem(he_n)
prob_comp = ODEProblem(sys_comp,u0,(0.0,10.0),p0)

sol_comp = solve(prob_comp,RK4())

@test getindex.(sol.u, 1) ≈ getindex.(sol_comp.u, 1)

end # testset


module DataAndPlotting

using HDF5
using Plots

#Wave parameters
const ki = 3
const kj = 3

#grid size
const ni = 3000
const nj = 2000

#i ranges from 1 to ni
#i - 1 ranges from 0 to ni - 1
#x ranges from 0 to 1
function index2coord(i, ni)
    @assert ni > 1
    @assert 1 <= i <= ni
    (i - 1) / (ni - 1)
end

#set up standing wave
export initialize
function initialize()
    arr = zeros(Float64, ni, nj)
    for j in 1:nj, i in 1:ni
        x = index2coord(i, ni)
        y = index2coord(j, nj)
        arr[i,j] = sin(2π*ki*x) * sin(2π*kj*y)
    end
    arr

end

#write data to file
export output
function output(arr::Array{Float64,2}, filename::String)
    h5write(filename, "data", arr)
end

#Read data from file
export input
function input(filename::String)
    arr = h5read(filename, "data") ## (1:ni, 1:nj)
end

#Create a plot
export makeplot
function makeplot(arr::Array{Float64,2}, filename::String)
    heatmap(arr, clim=(-1.0,+1.0), color=:viridis)
    savefig(filename)
end

end # module


# --------------
# Initialization
# --------------
Pkg.add("ArrayViews")
Pkg.add("Docile")
include("../src/LLLplus.jl")

using LLLplus
using Base.Test

# --------------
# tests with small matrices
# --------------
println("tests with small matrices...")

# Matrix from http://home.ie.cuhk.edu.hk/~wkshum/wordpress/?p=442
H =[1   9   1   2;
    1   8   8   3;
    7   4   5   1;
    2   6   7   1];
B1=[2   3  -2  -4
    3  -1   2   1
    1   1   6  -4
    1   3  -1   3];
T1=[0   0   1  -1
    0   1  -1   0
    0   0   0   1
    1  -3   3  -2];
(B,T) = lll(H);
@test sum(abs(B-B1))==0
@test sum(abs(T-T1))==0
println("...done")

# --------------
# Run tests with random matrices
# --------------

println("\nIn all the following tests, the first time includes the "*
        "JIT compilation; \nfor the second execution the "*
        "compilation is already done and the time\n"*
        "should be faster.\n")

# Time LLL decomposition of a 1000x1000 real matrix with randn entries
N = 1000;
H = randn(N,N);
println("Testing LLL on $(N)x$(N) real matrix...")
@time (B,T,Q,R) = lll(H);
@time (B,T,Q,R) = lll(H);

# Time LLL, Seysen, VBLAST decompositions of a 10x10 complex matrix with
# randn entries
N = 10;
println("\nTesting LLL on $(N)x$(N) complex matrix...")
H = randn(N,N) + im*randn(N,N);
@time (B,T,Q,R) = lll(H);
@time (B,T,Q,R) = lll(H);
# println("Testing Seysen on same $(N)x$(N) complex matrix...")
# @time (B,T) = seysen(H);
# @time (B,T) = seysen(H);
println("Testing VBLAST on same $(N)x$(N) complex matrix...")
@time (W,P,B) = vblast(H);
@time (W,P,B) = vblast(H);

# Test sphere decoder
Ns = 100000;
N = 2;
println("\nTesting sphere decoder on $(Ns) samples of $(N)x$(N) BPSK system...")
NN = randn(N,Ns)/sqrt(100);
H = [.6 -2.6; -.5 .95];
C = [-1,1];
Z = rand(1:2,N,Ns);
X = C[Z];
Y = H*X+NN;
@time Zt = hard_sphere(Y,H,2);
@time Zt = hard_sphere(Y,H,2);
Zh = Zt+1;
errRate = sum(abs(Z-Zh))/Ns;
println("Error Rate is $(errRate). It should be zero or very small.\n")

# --------------
# test norm for matrix from http://www.latticechallenge.org/
# --------------

println("Testing now with 200x200 matrix from latticechallenge.org.")
println("The col norm of the input should be 233, col norm of the "*
        "reduced bases should be 30.")
mat = readdlm("challenge-200.mod",Int64) #run from parent directory

nrms = zeros(200,1)
for ix=1:200
    nrms[ix] = norm(mat[:,ix])
end
println("max col-norm of input is $(maximum(nrms))")

@time (B,T) = lll(mat);
nrms = zeros(200,1)
for ix=1:200
    nrms[ix] = norm(B[:,ix])
end
maxlll = maximum(nrms)
println("max col-norm of lll-reduced basis is $(maxlll)")
@test maxlll<=30+1e-6

@time (B,T) = seysen(mat);
nrms = zeros(200,1)
for ix=1:200
    nrms[ix] = norm(B[:,ix])
end
maxSeysen = maximum(nrms)
println("max column norm of seysen-reduced basis is $(maxSeysen)")
@test maxSeysen<=30+1e-6


using BenchmarkTools
using Distances

##
function nu_L2!(nu, xa, xb, l)
  nu .= 0.0
  for ka = axes(xa, 2)
    @simd for kd = axes(xa, 1)
      @inbounds nu[ka] = nu[ka] + 0.5*(((xa[kd, ka] - xb[kd])/l[kd])^2)
    end
  end
end

function kern_sqexp!(out, nu)
  out[:] = exp.(-nu)
end

function build_K!(K, xa, xb, l, kern!, nufunc!)
  # Build a kernel matrix using a function `kern` to construct columns/rows
  nu = Array{Float64}(undef, size(xa, 2))
  for kb = axes(xb, 2)
    @views nufunc!(nu, xa, xb[:, kb], l)
    @views kern!(K[:, kb], nu)
  end
end

function kern_sqexp_elem(nu)
  exp(-nu)
end

@inline function nu_L2_elem(xa, xb, l)
  # mapreduce(k -> ((xa[k] - xb[k])/l[k])^2, +, axes(xa, 1)) / 2
  out = 0.0
  @inbounds @simd for k in axes(xa, 1)
    out += ((xa[k] - xb[k])/l[k])^2
  end
  out = 0.5*out
end

@inline function nu_L2_index(xa, xb, ka, kb, l)
  out = 0.0
  @inbounds @simd for k in axes(xa, 1)
    out += ((xa[k, ka] - xb[k, kb])/l[k])^2
  end
  out = 0.5*out
end

function build_K_elem!(K, xa, xb, l, kern::Function, nufunc::Function)
  # Build a kernel matrix using a function `kern` to construct elements
  @inbounds for kb = axes(xb, 2)
    for ka = axes(xa, 2)
      #@views nu = sqeuclidean(xa[:, ka], xb[:, kb])
      #@views nu = nu_L2_index(xa, xb, ka, kb, l)
      @views nu = nufunc(xa[:, ka], xb[:, kb], l)
      K[ka, kb] = kern(nu)
      #@views nu = nu_L2_elem(xa[:, ka], xb[:, kb], l)
      #K[ka, kb] = kern_sqexp_elem(nu)
    end
  end
end

nd = 4
na = 4096
nb = na

l = ones(nd)
xa = rand(nd, na)
xb = xa
K = zeros((na, nb))
##
@time build_K_elem!(K, xa, xb, l, kern_sqexp_elem, nu_L2_elem)
#@code_native build_K_elem!(K, xa, xb, l, kern_sqexp_elem, nu_L2_elem)
##
@time build_K!(K, xa, xb, l, kern_sqexp!, nu_L2!)
##
# @btime @views K = [kern_sqexp_elem(nu_L2_index(xa, xb, ka, kb, l)) for ka in axes(xa, 2), kb in axes(xb, 2)]
# print()
##
# using KernelFunctions

# k₁ = SqExponentialKernel()
# ##
# @time kernelmatrix!(K₁,k₁,xa,obsdim=2)
# print()


using StatsBase
using PyPlot

function meshgrid(x, y)
    X = [i for i in x, j in 1:length(y)]
    Y = [j for i in 1:length(x), j in y]
    return X, Y
end

x = -5 : 0.13 : 5
y = -5 : 0.13 : 5
(X,Y) = meshgrid(x, y)
Z = X + Y*im

Mu1 = log.(Z)./(2*pi)
Mu2 = log.(-sqrt.(Z))./(2*pi)
Mu3 = log.(sqrt.(sqrt.(Z)))./(2*pi)

mean( abs.(Mu2) ./ abs.(Mu1))

theta = -pi : 0.01 : pi
r = 0.5
z = r .* exp.(im*theta)
z2 = sqrt.(z)
z3 = -z2

figure()
plot(theta,  real.(z), imag.(z) )
plot(theta,  real.(z2), imag.(z2), "--" )
axis("equal")

figure()
plot(theta, abs.(log.(z)) )
plot(theta,  abs.(log.(z2)), "--"  )
plot(theta,  abs.(log.(z3)), ":"  )


module TrustRegionMethods

export trust_region_solver, TrustRegionResult, ForwardDiff_wrapper, SolverStoppingCriterion,
    Dogleg, GeneralizedEigenSolver

using ArgCheck: @argcheck
import DiffResults
using DocStringExtensions: FIELDS, FUNCTIONNAME, SIGNATURES, TYPEDEF
import ForwardDiff
using KrylovKit: eigsolve
using LinearAlgebra: diag, Diagonal, dot, I, issuccess, lu, norm, Symmetric, UniformScaling
using UnPack: @unpack

####
#### building blocks
####

"""
$(TYPEDEF)

A *linear model* for system of nonlinear equations, relative to the origin.

The L2 norm induces the objective function

```math
m(p) = p' J' r + 1/2 p' J' J p
```
approximating some `f(x) = \\| r(x) \\|^2_2` as

```math
1/2 \\| f(x + p) \\|_2^2 \approx 1/2 \\| r + J p \\|_2^2 = \\| r \\|_2^2 + p'J'r + 1/2 p' J' J p
```
"""
struct ResidualModel{TR,TJ}
    r::TR
    J::TJ
    function ResidualModel(r::TR, J::TJ) where {TR,TJ}
        @argcheck all(isfinite, r) "Non-finite residuals $(r)."
        @argcheck all(isfinite, J) "Non-finite Jacobian $(J)."
        n = length(r)
        @argcheck size(J) ≡ (n, n) "Non-conformable residual and Jacobian."
        new{TR,TJ}(r, J)
    end
end

"""
$(TYPEDEF)

A *quadratic model* for a minimization problem, relative to the origin, with the objective
function

```math
m(p) = p' g + 1/2 g' A g
```
"""
struct MinimizationModel{TG,TA}
    g::TG
    A::TA
    # FIXME constructor enforces symmetry of `A`, document, but do we actually rely on it?
end

"""
$(SIGNATURES)

Convert a model for a residual to a minimization model using the L2 norm.
"""
function MinimizationModel(model::ResidualModel)
    @unpack r, J = model
    MinimizationModel(J' * r, Symmetric(J' * J))
end

"""
$(SIGNATURES)

Find the Cauchy point (the minimum of the linear part, subject to ``\\| p \\| ≤ Δ``) of the problem.

Return three values:

1. the Cauchy point vector `pC`,

2. the (Euclidean) *norm* of `pC`

3. a boolean indicating whether the constraint was binding.

Caller guarantees non-zero gradient.
"""
function cauchy_point(Δ::Real, model::ResidualModel)
    @unpack r, J = model
    g = J' * r
    q = g' * (J' * J) * g
    g_norm = norm(g, 2)
    @argcheck g_norm > 0
    τ = if q ≤ 0              # practically 0 (semi-definite form) but allow for float error
        one(q)
    else
        min(one(q), g_norm^3 / (Δ * q))
    end
    (-τ * Δ / g_norm) .* g, τ * Δ, τ ≥ 1
end

"""
$(SIGNATURES)

Finds `\tau` such that

```math
\\| p_C + τ D \\|_2 = Δ
```

`pC_norm` is `\\| p_C \\|_2^2`, and can be provided when available.

Caller guarantees that

1. `pC_norm ≤ Δ`,

2. `norm(pC + D, 2) ≥ Δ`.

3. `dot(pC, D) ≥ 0`.

These ensure a meaningful solution `0 ≤ τ ≤ 1`. None of them are checked explicitly.
"""
function dogleg_boundary(Δ, D, pC, pC_norm = norm(pC, 2))
    a = sum(abs2, D)
    b = 2 * dot(D, pC)
    c = abs2(pC_norm) - abs2(Δ)
    (-b + √(abs2(b) - 4 * a * c)) / (2 * a)
end

"""
$(SIGNATURES)

Return `(F, is_valid_F)` where

1. `F` is a form (usually a factorization) of the argument that supports `F \\ r` for
   vectors `r`,

2. `is_valid_F` is a boolean indicating whether `F` can be used in this form (eg `false` for
   a singular factorization).
"""
function _factorize(J::AbstractMatrix)
    LU = lu(J; check = false)
    LU, issuccess(LU)
end

function _factorize(J::Diagonal)
    J, all(x -> x > 0, diag(J))
end

"""
$(SIGNATURES)

Calculate the unconstrained optimum of the model, return its norm as the second value.

When the second value is *infinite*, the unconstrained optimum should not be used as this
indicates a singular problem.
"""
function unconstrained_optimum(model::ResidualModel)
    @unpack r, J = model
    F, is_valid_F = _factorize(J)
    if is_valid_F
        pU = -(F \ r)
        pU, norm(pU, 2)
    else
        ∞ = oftype(one(eltype(r)) / one(eltype(F)), Inf)
        fill(∞, length(r)), ∞
    end
end

####
#### solvers
####

"""
`$(FUNCTIONNAME)(method, Δ, model)`

Optimize the `model` with the given constraint `Δ` using `method`. Returns the following
values:

- the optimum,
- the norm of the optimum,
- a boolean indicating whether the constraint binds.
"""
function solve_model end

"""
The dogleg method.

A simple and efficient heuristic for solving local trust region problems. Does not
necessarily find the optimum, but improves on the Cauchy point.
"""
struct Dogleg end

"""
$(SIGNATURES)

Implementats the branch of the dogleg method where it is already determined that the
unconstrained optimum is lies outside the `Δ` ball. Returns the same kind of results as
[`solve_model`](@ref).
"""
function dogleg_implementation(Δ, model, pU, pU_norm)
    pC, pC_norm, on_boundary = cauchy_point(Δ, model)
    if on_boundary || isinf(pU_norm)
        pC, pC_norm, on_boundary
    else
        D = pU .- pC
        τ = dogleg_boundary(Δ, D, pC, pC_norm)
        pC .+ D .* τ, Δ, true
    end
end

function solve_model(::Dogleg, Δ, model::ResidualModel)
    pU, pU_norm = unconstrained_optimum(model)
    if pU_norm ≤ Δ
        pU, pU_norm, false
    else
        dogleg_implementation(Δ, model, pU, pU_norm)
    end
end

####
#### Generalized eigenvalue solver (Adachi et al 2017).
####
#### NOTE:
####
#### - notation mostly follows the paper
####
#### - implementation has ellipsoidal norm in some parts, for future generalizations,
####   currently unused
####
#### - “hard case” is not implemented yet, falls back to dogleg

"""
Generalized eigenvalue solver (Adachi et al 2017).
"""
struct GeneralizedEigenSolver end

"""
$(SIGNATURES)

Kernel for the generalized eigenvalue solver (methods specialize on the ellipsoidal norm
matrix `B`). Solves the `M̃(λ)` pencil in Adachi et al (2017), eg Algorithm 5.1.

Returns

- `λ`, the largest eigenvalue,

- the gap to the next eigenvalue (see note below)

- `y1` and `y2`, the two parts of the corresponding generalized eigenvalue,

All values are real, methods are type stable.

When theorerical assumptions are violated, `gap` will be non-finite and a debug statement is
emitted. No other values should be used in this case.
"""
function ges_kernel(Δ, model::MinimizationModel, B::UniformScaling)
    @unpack g, A = model
    n = length(g)
    G = ((g * g') ./ abs2(Δ))
    M = [-A G; B -A]
    λs, vs, info = eigsolve(M, 2, :LR)
    @argcheck info.converged ≥ 2 "Eigensolver did not converge."
    T = promote_type(Float64, eltype(M))
    ϵ = √eps(T)
    is_practically_real(z) = imag(z) ≤ n*(ϵ*abs(z) + ϵ)
    λ1, λ2 = λs
    v1 = first(vs)
    if is_practically_real(λ1) && all(is_practically_real, v1)
        λ = real(λ1)::T
        gap = abs(λ - λ2)::T
        y = real.(v1)::Vector{T}
        y1 = y[1:n]
        y2 = y[(n + 1):end]
        λ, gap, y1, y2
    else
        @debug "rightmost eigenvalue not real" M λs vs
        ∞ = T(Inf)
        y∞ = fill(∞, n)
        ∞, ∞, y∞, y∞
    end
end

"""
$(SIGNATURES)

Ellipsoidal norm ``\\| x \\|_B = x'Bx``.
"""
ellipsoidal_norm(x, ::UniformScaling) = norm(x, 2)

function solve_model(::GeneralizedEigenSolver, Δ, model::ResidualModel)
    B = I                       # FIXME: we hardcode Euclidean norm for now
    pU, pU_norm = unconstrained_optimum(model)
    if pU_norm < Δ
        pU, pU_norm, false
    else
        model′ = MinimizationModel(model)
        λ, gap, y1, y2 = ges_kernel(Δ, model′, B)
        τ = √(eps(typeof(λ)) / gap)
        if isfinite(gap) && norm(y1, 2) > τ
            # “easy” case, we can generate a candidate
            @unpack g = model′
            p = (-sign(dot(g, y2)) * Δ / ellipsoidal_norm(y1, B)) .* y1
            p, Δ, true
        else
            # FIXME hard case, this needs to be implemented
            # here we bail and fall back to dogleg
            # NOTE this is also the fallback for eigensolver gone wrong
            @debug "hard case not implemented, falling back to dogleg"
            dogleg_implementation(Δ, model, pU, pU_norm)
        end
    end
end

####
#### trust region steps
####

struct TrustRegionParameters{T}
    η::T
    Δ̄::T
    function TrustRegionParameters(η::T, Δ̄::T) where {T <: Real}
        @argcheck 0 < η < 0.25
        @argcheck Δ̄ > 0
        new{T}(η, Δ̄)
    end
end

TrustRegionParameters(η, Δ̄) = TrustRegionParameterS(promote(η, Δ̄)...)

TrustRegionParameters(; η = 0.125, Δ̄ = Inf) = TrustRegionParameters(η, Δ̄)

"""
$(SIGNATURES)

Ratio between predicted (using `model`, at `p`) and actual reduction (taken from the
residual value `r′`). Will return an arbitrary negative number for infeasible coordinates.
"""
function reduction_ratio(model::ResidualModel, p, r′)
    @unpack r, J = model
    r2 = sum(abs2, r)
    r′2 = sum(abs2, r′)
    ρ = (r2 - r′2) / (r2 - sum(abs2, r .+ J * p))
    isfinite(ρ) ? ρ : -one(ρ)
end

"""
$(SIGNATURES)

Checks residual and Jacobian and throws an appropriate error if they are not both finite.

Internal, each evaluation of `f` should call this.
"""
function _check_residual_Jacobian(x, fx)
    finite_value = all(isfinite, fx.residual)
    finite_Jacobian = all(isfinite, fx.Jacobian)
    if !finite_value || !finite_Jacobian
        if finite_value == finite_Jacobian == false
            msg = "residual and Jacobian"
        else
            msg = !finite_value ? "residual" : "Jacobian"
        end
        DomainError(x, "Non-finite values in $(msg)")
    end
end

"""
$(SIGNATURES)

Take a trust region step using `local_method`.

`f` is the function that returns the residual and the Jacobian (see
[`trust_region_solver`](@ref)).

`Δ` is the trust region radius, `x` is the position, `fx = f(x)`. Caller ensures that the
latter is feasible.
"""
function trust_region_step(parameters::TrustRegionParameters, local_method, f, Δ, x, fx)
    @unpack η, Δ̄ = parameters
    model = ResidualModel(fx.residual, fx.Jacobian)
    p, p_norm, on_boundary = solve_model(local_method, Δ, model)
    x′ = x .+ p
    fx′ = f(x′)
    _check_residual_Jacobian(x′, fx′)
    ρ = reduction_ratio(model, p, fx′.residual)
    Δ′ =
        if ρ < 0.25
            p_norm / 4
        elseif ρ > 0.75 && on_boundary
            min(2 * Δ, Δ̄)
        else
            Δ
        end
    if ρ ≥ η
        Δ′, x′, fx′            # reasonable reduction, use new position
    else
        Δ′, x, fx              # very small reduction, try again from original
    end
end

struct SolverStoppingCriterion{T <: Real}
    residual_norm_tolerance::T
    function SolverStoppingCriterion(residual_norm_tolerance::T) where {T}
        @argcheck residual_norm_tolerance > 0
        new{T}(residual_norm_tolerance)
    end
end

function SolverStoppingCriterion(; residual_norm_tolerance = √eps())
    SolverStoppingCriterion(residual_norm_tolerance)
end

function check_stopping_criterion(nsc::SolverStoppingCriterion, fx)
    r_norm = norm(fx.residual, 2)
    converged = r_norm ≤ nsc.residual_norm_tolerance
    converged, (converged = converged , residual_norm = r_norm)
end

"""
$(TYPEDEF)

# Fields

$(FIELDS)
"""
struct TrustRegionResult{T<:Real,TX<:AbstractVector{T},TFX}
    "The final trust region radius."
    Δ::T
    "The last value (the root only when converged)."
    x::TX
    "`f(x)` at the last `x`."
    fx::TFX
    "The Euclidean norm of the residual at `x`."
    residual_norm::T
    "A boolean indicating convergence."
    converged::Bool
    "Number of iterations (≈ number of function evaluations)."
    iterations::Int
end

function TrustRegionResult(Δ::T1, x::AbstractVector{T2}, fx, residual_norm::T3, converged,
                           iterations) where {T1 <: Real, T2 <: Real, T3 <: Real}
    T = promote_type(T1, T2, T3)
    TrustRegionResult(T(Δ), T.(x), fx, T(residual_norm), converged, iterations)
end

function Base.show(io::IO, ::MIME"text/plain", trr::TrustRegionResult)
    @unpack Δ, x, fx, residual_norm, converged, iterations = trr
    _sig(x) = round(x; sigdigits = 3)
    print(io, "Nonlinear solver using trust region method ")
    if converged
        printstyled(io, "converged"; color = :green)
    else
        printstyled(io, "didn't converge"; color = :red)
    end
    println(io, " after $(iterations) steps")
    print(io, "  with ")
    printstyled(io, "‖x‖₂ = $(_sig(residual_norm)), Δ = $(_sig(Δ))\n"; color = :blue)
    println(io, "  x = ", _sig.(x))
    println(io, "  r = ", _sig.(fx.residual))
end

"""
$(SIGNATURES)

Solve `f ≈ 0` using trust region methods, starting from `x`.

`f` is a callable (function) that

1. takes vectors real numbers of the same length as `x`,

2. returns a value with properties `residual` and `Jacobian`, containing a vector and a
   conformable matrix, with finite values for both or a non-finite residual (for infeasible
   points; Jacobian is ignored). A `NamedTuple` can be used for this, but any structure with
   these properties will be accepted. Importantly, this is treated as a single object and
   can be used to return extra information (a “payload”), which will be ignored by this
   function.

Returns a [`TrustRegionResult`](@ref) object.
"""
function trust_region_solver(f, x;
                             parameters = TrustRegionParameters(),
                             local_method = Dogleg(),
                             stopping_criterion = SolverStoppingCriterion(),
                             maximum_iterations = 500,
                             Δ = 1.0)
    @argcheck Δ > 0
    fx = f(x)
    _check_residual_Jacobian(x, fx)
    iterations = 1
    while true
        Δ, x, fx = trust_region_step(parameters, local_method, f, Δ, x, fx)
        iterations += 1
        do_stop, convergence_statistics = check_stopping_criterion(stopping_criterion, fx)
        reached_max_iter = iterations ≥ maximum_iterations
        if do_stop || reached_max_iter
            @unpack converged, residual_norm = convergence_statistics
            return TrustRegionResult(Δ, x, fx, residual_norm, converged, iterations)
        end
    end
end

####
#### AD shims
####

###
### ForwardDiff
###

"""
$(TYPEDEF)

A buffer for wrapping a (nonlinear) mapping `f` to calculate the Jacobian with
`ForwardDiff.jacobian`.
"""
struct ForwardDiffBuffer{TX,TF,TR,TC}
    "Buffer for inputs."
    x::TX
    "A function that maps vectors to residual vectors."
    f::TF
    "`DiffResults` buffer for Jacobians."
    result::TR
    "Gradient configuration."
    cfg::TC
end

"""
$(SIGNATURES)

Wrap an ``ℝⁿ`` function `f` in a callable that can be used in [`trust_region_solver`](@ref)
directly. Remaining parameters are passed on to `ForwardDiff.JacobianConfig`, and can be
used to set eg chunk size.

Non-finite residuals will be treated as infeasible.

```jldoctest
julia> f(x) = [1.0 2; 3 4] * x - ones(2)
f (generic function with 1 method)

julia> ff = ForwardDiff_wrapper(f, 2)
AD wrapper via ForwardDiff for ℝⁿ→ℝⁿ function, n = 2

julia> ff(ones(2))
(r = [2.0, 6.0], J = [1.0 2.0; 3.0 4.0])

julia> trust_region_solver(ff, [100.0, 100.0])
Nonlinear solver using trust region method converged after 9 steps
  with ‖x‖₂ = 3.97e-15, Δ = 128.0
  x = [-1.0, 1.0]
  r = [-1.78e-15, 3.55e-15]
```
"""
function ForwardDiff_wrapper(f, n, jacobian_config...)
    x = zeros(n)
    result = DiffResults.JacobianResult(x)
    cfg = ForwardDiff.JacobianConfig(f, x, jacobian_config...)
    ForwardDiffBuffer(x, f, result, cfg)
end

function Base.show(io::IO, ::MIME"text/plain", fdb::ForwardDiffBuffer)
    n = length(DiffResults.value(fdb.result))
    print(io, "AD wrapper via ForwardDiff for ℝⁿ→ℝⁿ function, n = $(n)")
end

function (fdb::ForwardDiffBuffer)(y)
    @unpack x, f, result, cfg = fdb
    # copy to our own buffer to work with types other than Float64
    ForwardDiff.jacobian!(result, f, copy!(x, y), cfg)
    residual = copy(DiffResults.value(result))
    (residual = residual, Jacobian = copy(DiffResults.jacobian(result)))
end

end # module


using DrWatson
@quickactivate "ExercisesRepo"
using DelimitedFiles
cd(@__DIR__)

# %% 1 = Mackie Glass timeseries
using DelayDiffEq
function mackey_glass(du,u,h,p,t)
    beta,n,gamma,tau = p
    hist = h(p, t-tau)[1]
    du[1] = (beta*hist)/(1+hist^n) - gamma * u[1]
end
h(p,t) = 0
tau_d = 44
n = 10
β = 0.2
γ = 0.1
δt = 0.5
p = (β,n,γ,tau_d)

tspan = (0.0, 12000.0)
u0 = [1.0]

prob = DDEProblem(mackey_glass,u0,h,tspan,p; constant_lags=tau_d)
alg = MethodOfSteps(Tsit5())
sol = solve(prob,alg; adaptive=false, dt=δt)

s = [u[1] for u in sol.u]
s = s[4001:end]
writedlm("1.csv", s)

# %% 2 = Roessler trajectory (1st parameter set)
using DynamicalSystemsBase
roe = Systems.roessler([1.0, 0, 0]; a=0.2, b=0.2, c=5.7)
sroe = trajectory(roe, 500; dt = 0.05, Ttr = 100.0)
writedlm("2.csv", sroe)

# %% 3 = Lorenz trajectory (1st parameter set, dense time sampling)
lo = Systems.lorenz([0, 10.0, 0.0]; σ=10, ρ=28, β=8/3)
tr = trajectory(lo, 100; dt = 0.01, Ttr = 100)
writedlm("3.csv", tr)

# %% 4 = Roessler trajectory, low sampling, 2nd parameter set
ro = Systems.roessler([1.0, 1.0, 1.0], a=0.15, b = 0.2, c=10)
tr = trajectory(ro, 1250; dt = 0.2, Ttr = 100)
writedlm("4.csv", tr)

# %% 5 = collection of periodic and quasiperiodic orbits
# some from Henon Heiles
using DynamicalSystems, OrdinaryDiffEq, PyPlot, DelimitedFiles
diffeq = (alg = Vern9(), reltol = 1e-9, abstol = 1e-9)
hh = Systems.henonheiles()

xs = Vector{Float64}[]
for u in (
        [0.0, 0.1, 0.5, 0.0], # quasiperiodic
        [0.0, -0.0910355, 0.459522, -0.173339], # higher order quasiperiodic
        [0.0, 0.30266568064689636, 0.42056546332015626, 0.0] # periodic
    )

    tr = trajectory(hh, 500.0, u; Ttr = 10, dt = 0.1, diffeq...)
    x = regularize(tr[:, 1])
    x .+= 0.05randn(length(x))
    push!(xs, x)
end

# Also add a quasi-periodic orbit from Hose Noover for good measure
nh = Systems.nosehoover([0, 1.0, 0])
tr = trajectory(nh, 500.0; Ttr = 10, dt = 0.1, diffeq...)
x = regularize(tr[:, 1])
x .+= 0.05randn(length(x))
push!(xs, x)

# chaotic & period-4 roessler
ro = Systems.roessler(ones(3))
for c in (4.0, 5.7)
    set_parameter!(ro, 3, c)
    tr = trajectory(ro, 500.0; Ttr = 1000, dt = 0.1, diffeq...)
    x = regularize(tr[:, 1])
    x .+= 0.05randn(length(x))
    push!(xs, x)
end

# from standard map (quasiperiod around period 3)
sm = Systems.standardmap()

for (i, u) in enumerate((
        SVector(0.8121, 1.6243),
        SVector(0.877, 1.565),
    ))
    tr = trajectory(sm, 5000, u; Ttr = 10)
    x = regularize(tr[:, 1])
    i == 1 && (x .+= 0.02randn(length(x)))
    push!(xs, x)
end

# Fourth part is just oscillatory ARMA, a-la chapter 7
using Random, Statistics
Random.seed!(77163)
η = randn(15000)
s = ones(15000)
for n in 5:15000
    s[n] = 1.625s[n-1] - 0.284s[n-2] - 0.355s[n-3] + η[n] - 0.26η[n-1]
end
x = s[end-5000:end]
x ./= std(x)
push!(xs, x)

# Last part 3 cosines, to confuse for periodic → quasiperiodic
t = range(0, 100π; length = 5001)
x = @. cos(t) + 0.5cos(2t + π/3) + 0.2cos(13t + π/5)
x ./= std(x)
x .+= 0.05randn(length(x))
push!(xs, x)

writedlm("5.csv", hcat(xs...))

# %% 6 = towel map trajectory
using DynamicalSystems, DelimitedFiles

ds = Systems.towel()

tr = trajectory(ds, 10000; Ttr = 100)

writedlm("6.csv", tr)

# %% 7 = Non-stationary dynamics of logistic map
ds = Systems.logistic()
rs = 3.60:0.0001:3.9
integ = integrator(ds, π/4)
x = [integ.u]; sizehint!(x, length(rs)*3)

for r in rs
    set_parameter!(ds, 1, r)
    for i in 1:3
        DynamicalSystems.step!(integ)
        push!(x, integ.u)
    end
end

x .+= 0.05randn(length(x))

writedlm("7.csv", x)

# %% 8 = oscillatory ARMA, embedded in 4d
using ARFIMA, Random

rng = MersenneTwister(5132786517)
x = arma(rng, 10000, 0.4, SVector(1.624, -0.284, -0.355), SVector(-0.96))
τ = estimate_delay(x, "mi_min")
e = embed(x, 4, τ)

# This looks awesome!
# using3D()
# cols = columns(e)
# scatter3D(cols[1], cols[2], cols[3]; c = cols[4])
writedlm("8.csv", x)

# %% 9 = 6dimensional Lorenz96 with 2% AR noise - Periodic or Quasiperiodic (don't know!)
using DynamicalSystems
D = 6
u = fill(0.1, D)
u[1] = 0.2
u[3] = 0.3
lo = Systems.lorenz96(D, u; F = 5.5)

tr = regularize(trajectory(lo, 100.0; Ttr = 10.0))
co = columns(tr)
for x in co
    noise = 0.02randn(length(x))
    noise = arma(length(x), 0.4, SVector(0.624, -0.284, -0.355, +0.355))
    noise ./= std(noise) / (0.02)
    x .+= noise
end

writedlm("9.csv", hcat(co...))

# scatter3D(co[1], co[2], co[3]; c = co[4])

# %% 10 = 5dimensional Lorenz96 with 2% AR noise - Chaotic trajectory
using DynamicalSystems
D = 5
u = fill(0.1, D)
u[1] = 0.2
u[3] = 0.3
lo = Systems.lorenz96(D, u; F = 8.0)

tr = regularize(trajectory(lo, 100.0; Ttr = 10.0))
co = columns(tr)
for x in co
    noise = 0.02randn(length(x))
    noise = arma(length(x), 0.4, SVector(0.624, -0.284, -0.355, +0.355))
    noise ./= std(noise) / (0.02)
    x .+= noise
end

# scatter3D(co[1], co[2], co[3]; c = co[4])
writedlm("10.csv", hcat(co...))


println("\n\n\nPerformance of FinanceLib.FixedIncomes.Bonds functions")

import FinanceLib as Fl
import FinanceLib.FixedIncomes.Bonds as Bd
import FinanceLib.FixedIncomes.Bonds.Rates as Rt

# @code_warntype Bd.priceCouponBond(Fl.PeriodSeries([(1, 0.09), (2, 0.10)]), 0.06) * 1000

# println("Coupon bond")
# @code_warntype Bd.priceCouponBond(Fl.PeriodSeries([(1, 0.09), (2, 0.10), (3, 0.11)]), 0.05) * 100

# rS = Fl.PeriodSeries([(1, 0.05), (2, 0.06), (3, 0.07), (4, 0.08), (5, 0.09)])
# @code_warntype Bd.ytmCouponBond(rS, 0.1)

# @code_warntype Rt.parToSpotRates([(1,0.05), (2,0.0597), (3, 0.0691), (4, 0.0781)])

# @code_warntype Rt.parToSpotRates(Fl.RateCurve{Fl.NomRate}([0.05, 0.0597, 0.0691, 0.0781], 1))

# @code_warntype Rt.parToSpotRates(Rt.spotToParRates(Fl.RateCurve{Fl.NomRate}([0.05, 0.06, 0.07, 0.08], 2))).rate[3]

# @code_warntype Bd.priceCouponBond(0.07, Bd.CouponBond(0.05, 2, 3))

# @code_warntype Bd.priceCouponBond(Fl.RateCurve{Fl.NomRate}([0.05, 0.055, 0.06, 0.07, 0.075, 0.085],2), Bd.CouponBond(0.05, 2, 3))

# @code_warntype Bd.ytmCouponBond(Bd.CouponBond(0.05, 2, 3), 0.916183660871172)

# println("Benchmarks\n")

# @btime Bd.priceCouponBond(Fl.PeriodSeries([(1, 0.09), (2, 0.10), (3, 0.11)]), 0.05) * 100

# @btime Bd.ytmCouponBond(Bd.CouponBond(0.05, 2, 3), 0.916183660871172)

# # @btime Bd.ytmCouponBond(rS, 0.1)

# # @btime Rt.parToSpotRates([(1,0.05), (2,0.0597), (3, 0.0691), (4, 0.0781)])

# rx = Fl.RateCurve([0.05, 0.0597, 0.0691, 0.0781], 1)

# @btime Rt.parToSpotRates(rx)


doc"""
Compute a rigorous bound for a Taylor1Model.
"""
bound(f::Taylor1Model) = bound(f, f.I)

doc"""
Compute a rigorous bound for a `Taylor1Model` `f` evaluated over the interval `I`.
This interval must be a subset of `f.I`.
"""
function bound(f::Taylor1Model, I)
    if !(I ⊆ f.I)
        error("Cannot evaluate `Taylor1Model` on interval $I that is not included in interval of definition, $(f.I)")
    end

    return bound(f.p, f.x0, I) + f.Δ
end

function bound(p::Taylor1, x0, I)

    B = zero(I)
    n = degree(p)

    for i = n:-1:0
        B = B * (I - x0) + bound(p[i])
    end

    return B
end

bound(f::TaylorN) = evaluate(f)  # can replace by better polynomial bounder

# bound(f::Taylor1Model) = bound(f.p, f.x0, f.I)

bound(x::Interval) = x


"""
heq_aa()
Duplicate the part of Table 4 in Toth-Kelley for l^2 optimization
with c=.99.

Compare solution with one from nsoli.
"""
function heq_aa()
    c = 0.99
    tol = 1.e-8
    fcount = [11, 10, 10, 11, 12, 12]
    anormref = [4.0, 5.4, 5.4, 5.4, 5.4, 5.4]
    condref = [1.0, 2.e2, 1.9e5, 1.9e7, 5.5e9, 6.5e10]
    n = 500
    u0 = ones(n)
    hdata = heqinit(u0, c)
    FS = zeros(n)
    FPS = zeros(n, 20)
    rtol = tol
    atol = tol
    maxit = 100
    mmax = 6
    Vstore = zeros(n, 3 * mmax + 2)
    itrecords = zeros(6, 4)
    houtn = nsoli(
        heqf!,
        u0,
        FS,
        FPS;
        pdata = hdata,
        rtol = rtol,
        atol = atol,
        lmaxit = 10,
        eta = 0.01,
    )
    #
    # Solve H-equation with Anderson(m) for several values of m
    #
    for m = 1:6
        houta = aasol(
            HeqFix!,
            u0,
            m,
            Vstore;
            maxit = maxit,
            pdata = hdata,
            rtol = rtol,
            atol = atol,
        )
        #
        # Keep the books.
        #
        itrecords[m, 1] = hresults(houta.solution, houtn.solution)
        itrecords[m, 2] = length(houta.history)
        itrecords[m, 3] = norm(houta.stats.condhist, Inf)
        itrecords[m, 4] = norm(houta.stats.alphanorm, Inf)
    end
    #
    # Grade the results. I only use the coefficent norm and the condition
    # numbers for my own research.
    #
    solok = (norm(itrecords[:, 1], Inf) < 1.e-7)
    solok || println("Solution error in Anderson H solve")
    histok = (norm(itrecords[:, 2] - fcount, Inf) < 1.e-5)
    histok || println("History error in Anderson H solve")
    condok = (hresults(itrecords[:, 3], condref) < 1.e-1)
    condok || println("Condition error in Anderson H solve")
    normok = (hresults(itrecords[:, 4], anormref) < 1.e-1)
    normok || println("Coefficient norm error in Anderson H solve")
    #return (itrecords = itrecords)
    return solok && histok && condok && normok
end

function hresults(x, y)
    vdiff = (x - y) ./ abs.(x)
    return norm(vdiff, Inf)
end


#########################################################################
#    Testing script for gradients of distributions
#########################################################################

include("helper_diff.jl")

## variables of different dimension for testing
v0ref = 2.
v1ref = [2., 3, 0.1, 0, -5]
v2ref = [-1. 3 0 ; 0 5 -2]

## continuous distributions
@test_combin    logpdf(Normal(mu, sigma), x)      sigma->sigma<=0?0.1:sigma size(mu)==size(sigma) &&
													(ndims(mu)==0 || size(x)==size(mu))

@test_combin    logpdf(Uniform(a, b), x)    a->a-10 b->b+10 size(a)==size(b) && (ndims(a)==0 || size(x)==size(a))

@test_combin    logpdf(Weibull(sh, sc), x)    sh->sh<=0?0.1:sh  sc->sc<=0?0.1:sc  x->x<=0?0.1:x size(sh)==size(sc) &&
													(ndims(sh)==0 || size(x)==size(sh))

@test_combin    logpdf(Gamma(sh, sc), x)    sh->sh<=0?0.1:sh  sc->sc<=0?0.1:sc x->x<=0?0.1:x  size(sh)==size(sc) &&
												(ndims(sc)==0 || size(sc)==size(x))

@test_combin    logpdf(Beta(a,b),x)         x->clamp(x, 0.01, 0.99) a->a<=0?0.1:a b->b<=0?0.1:b (size(a)==size(b)) &&
												(ndims(a)==0 || size(x)==size(a))

@test_combin    logpdf(TDist(df),x)         df->df<=0?0.1:df    (size(df)==size(x)) || ndims(df)==0  # fail

@test_combin    logpdf(Exponential(sc),x)   sc->sc<=0?0.1:sc  x->x<=0?0.1:x   (size(sc)==size(x)) || ndims(sc)==0  # fail

@test_combin    logpdf(Cauchy(mu,sc),x)   sc->sc<=0?0.1:sc  (size(mu)==size(sc)) &&
								(ndims(mu)==0 || size(x)==size(mu))

@test_combin    logpdf(LogNormal(lmu,lsc),x)   lsc->lsc<=0?0.1:lsc x->x<=0?0.1:x (size(lmu)==size(lsc)) &&
								(ndims(lmu)==0 || size(x)==size(lmu))

# FIXME, Laplace broken
# @test_combin    logpdf(Laplace(loc,sc),x)   sc->sc<=0?0.1:sc (size(loc)==size(sc)) &&
#  								(ndims(loc)==0 || size(x)==size(loc))

## discrete distributions
#  the variable x being an integer should not be derived against

# note for Bernoulli : having prob=1 or 0 is ok but will make the numeric differentiator fail => not tested
@test_combin logpdf(Bernoulli(prob),x) prob prob->clamp(prob, 0.01, 0.99) x->float64(x>0) size(prob)==size(x)||ndims(prob)==0  #fail

@test_combin logpdf(Poisson(l),x) l l->l<=0?0.1:l x->iround(abs(x)) size(l)==size(x)||ndims(l)==0  #fail

@test_combin logpdf(Binomial(n,prob),x) prob prob->clamp(prob, 0.01, 0.99) x->iround(abs(x)) n->iround(abs(n)+10) (size(n)==size(prob)) &&
								(ndims(n)==0 || size(x)==size(n))

#########################################################################
#   misc. tests
#########################################################################

@test_throws ErrorException deriv1(:(logpdf(Bernoulli(1), x)), [1.])
@test_throws ErrorException deriv1(:(logpdf(Poisson(1), x)), [1.])
@test_throws ErrorException deriv1(:(logpdf(Binomial(3, 0.5), x)), [0.])
#@test_throws ErrorException deriv1(:(logpdf(Binomial(x, 0.5), 2)), [0.])

##  ref  testing
deriv1(:(x[2]),              v1ref)
deriv1(:(x[2:3]),            v1ref)
deriv1(:(x[2:end]),          v1ref)

deriv1(:(x[2:end]),          v2ref)
deriv1(:(x[2]),              v2ref)
deriv1(:(x[2:4]),            v2ref)
deriv1(:(x[:,2]),            v2ref)
deriv1(:(x[1,:]),            v2ref)
deriv1(:(x[2:end,:]),        v2ref)
deriv1(:(x[:,2:end]),        v2ref)

deriv1(:(x[2]+x[1]),          v2ref)
deriv1(:(log(x[2]^2+x[1]^2)), v2ref)

# fail case when individual elements of an array are set several times
# model = :(x::real(3); y=x; y[2] = x[1] ; y ~ TestDiff())


"""
    restrict(refspace, element1, element2)

Computes the restriction of a set of local shape functions on `element1` as linear
combinations of the set of local shape functions on `element2`. More precisely `restrict`
returns an `NxM` matrix `P` such that the `i`-th local shape ``g_i`` function on element2
can be written as:

``g_i = sum_{j=1}^{M} P_{ij} f_j``
"""
function restrict end


import Statistics: mean, std, stdm, var, varm

mean(a::KnetArray; o...) = (b = sum(a; o...); b .* convert(eltype(b),(length(b)/length(a))))
mean(f::Base.Callable, a::KnetArray) = sum(f, a) / length(a)
std(x::KnetArray, args...; kws...) = sqrt.(var(x, args...; kws...))
stdm(x::KnetArray, args...; kws...) = sqrt.(varm(x, args...; kws...))
var(x::KnetArray; corrected::Bool=true, mean=nothing, dims=:)=_varm(x, mean; corrected=corrected, dims=dims)
varm(x::KnetArray, m; corrected::Bool=true, dims=:)=_varm(x, m; corrected=corrected, dims=dims)

function _varm(x, m; corrected::Bool=true, dims=:)
    if m === nothing; m = mean(x,dims=dims); end
    s = sum(abs2, x .- m; dims=dims)
    r = length(x) ÷ length(s) - Int(corrected)
    s ./ r
end

if first(methods(varm, (AutoGrad.Value,Any))).module !== AutoGrad # TODO: delete after AutoGrad 1.1.4
    stdm(x::AutoGrad.Value, args...; kws...) = sqrt.(varm(x, args...; kws...))
    varm(x::AutoGrad.Value, m; corrected::Bool=true, dims=:)=_varm(x, m; corrected=corrected, dims=dims)
end

# TODO: stdm, varm to be deprecated after Julia 1.x


function calculate_sum(upper_limit)
    s = 0
    for i in 1 : upper_limit - 1
        if i % 3 == 0 || i % 5 == 0
            s += i
        end
    end
    println(s)
end

calculate_sum(1000)