### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# ╔═╡ 1bcf25d8-71c2-11eb-05ca-6fe7c4e1b24f
begin
	using Pkg
	Pkg.activate("MLJ_env", shared=true)

	# using MLJ
	using CSV
	using DataFrames
	using PlutoUI
	using Random
	using Test
	using Printf
	using Plots
end

# ╔═╡ 96e7042c-73be-11eb-3ac5-9f9998b8c8ff
using LinearAlgebra

# ╔═╡ d17897d0-7540-11eb-1d5e-191caf244165
include("./utils.jl")

# ╔═╡ 9dcffb1c-71c1-11eb-2402-39eaa1f1b5e1
md"""
## Week 4: Ridge Regression (Gradient Descent)
"""

# ╔═╡ c39f9348-71c1-11eb-1691-1d28c6b1d916
md"""
In this notebook we will cover estimating multiple regression weights via gradient descent, performing the following:

  - Write a function to compute the derivative of the regression weights with respect to a single feature
  -  Write gradient descent function to compute the regression weights given an initial weight vector, step size, tolerance, and L2 penalty
"""

# ╔═╡ 3e5d907e-71c2-11eb-2eb3-4375fb21b884
md"""
### Load in house sales data
"""

# ╔═╡ 4c736078-71c2-11eb-225b-816745199225
sales = CSV.File("../../ML_UW_Spec/C02/data/kc_house_test_data.csv";
	header=true) |> DataFrame;

# ╔═╡ 35b5d05a-71cc-11eb-15a9-852c5a534fab
md"""
### Convert to Julia Matrix/Vector
"""

# ╔═╡ a2928dc6-71c2-11eb-3724-29dc1448a0f3
md"""
Now we will write a function that will accept a DataFrame, a list of feature names (e.g. [:sqft_living, :bedrooms') and a target feature e.g. (:price) and will return two things:

 - A matrix whose columns are the desired features plus a constant column (this is how we create an 'intercept')
 - An array containing the values of the output

With this in mind, let's write the `get_data` and `predict_output` functions (moved to `utils.jl`) as follows:

```julia
function get_data(df, features, output)
	df[:, :constant] .= 1.0
	features = [:constant, features...]
	X_matrix = convert(Matrix, select(df, features)) # to get a matrix
	y = df[!, output]                                # => to get a vector
	(X_matrix, y)
end

function predict_output(X::Matrix{T}, weights::Vector{T}) where {T <: Real}
    # assume feature_matrix is a matrix containing the features as columns
	# and weights is a corresponding array
    X * weights
end
```

"""

# ╔═╡ 21635b6e-71c8-11eb-2272-edec98c58339
md"""
### Computing the derivative

We are now going to move to computing the derivative of the regression cost function. Recall that the cost function is the sum over the data points of the squared difference between an observed output and a predicted output, plus the L2 penalty term.

$$Cost(w) = Σ [(prediction - output)^2 + l_2 × Σ_{i=1}^{k} w_i^2]$$

Since the derivative of a sum is the sum of the derivatives, we can take the derivative of the first part (the RSS) as we did in the notebook for the unregularized case in Week 2 and add the derivative of the regularization part. <br />
As we saw, the derivative of the RSS with respect to `w[i]` can be written as:

$$2 × Σ_i [error × [feature_i]]$$

The derivative of the regularization term with respect to `w[i]` is:

$$2 × l_2 \times w_i$$

Summing both, we get

$$2 \times Σ_i [error \times [feature_i] + 2 × l_2 * w_i]$$

That is, the derivative for the weight for feature i is the sum (over data points) of 2 times the product of the error and the feature itself, plus $2 × l_2 * w_i$.

**We will not regularize the constant.**  Thus, in the case of the constant, the derivative is just twice the sum of the errors (without the $2 × l_2 * w_0$ term).

Recall that twice the sum of the product of two vectors is just twice the dot product of the two vectors. Therefore the derivative for the weight for feature_i is just two times the dot product between the values of feature_i and the current errors, plus $2 × l_2 * w_i$.

With this in mind complete the following derivative function which computes the derivative of the weight given the value of the feature (over all data points) and the errors (over all data points).

To decide when to we are dealing with the constant (so we don't regularize it) we added the extra parameter to the call `feature_is_constant` which you should set to `true` when computing the derivative of the constant and `false` otherwise.

"""

# ╔═╡ f377f6d4-71c9-11eb-0249-f3d63e87062b
function feature_derivative_ridge(errors::Vector{T}, feature::Vector{T},
		weights, l2; is_constant=false) where {T <: Real}
	deriv_err = 2 * dot(feature, errors) # ≡ 2 * feature' * errors
	deriv_l2p = l2 * weights
    return is_constant ? deriv_err : deriv_err + deriv_l2p
end

# ╔═╡ 59e3f22a-73bd-11eb-1c0d-150d7d462676
begin
	(ex_features, ex_output) = get_data(sales, [:sqft_living], :price)
	typeof(ex_features), typeof(ex_output)
	ex_weights = Float64[1., 10.]
	test_preds = predict_output(ex_features, ex_weights)
	errors = test_preds - ex_output # prediction errors

	@test feature_derivative_ridge(errors, ex_features[:, 2], ex_weights[2], 1;
		is_constant=false) ≈ sum(errors .* ex_features[:, 2]) * 2. + 20.

	@test feature_derivative_ridge(errors, ex_features[:, 1], ex_weights[1], 1;
		is_constant=true) ≈ sum(errors) * 2.
end

# ╔═╡ 239ff6ca-71cc-11eb-341a-df1453d60a15
md"""
### Gradient Descent [GD]

Now we will write a function that performs a gradient descent. The basic premise is simple. Given a starting point we update the current weights by moving in the negative gradient direction. Recall that the gradient is the direction of increase and therefore the negative gradient is the direction of decrease and we're trying to minimize a cost function.

The amount by which we move in the negative gradient direction is called the *step size* denoted by η. We stop when we are 'sufficiently close' to the optimum. Unlike in Week 2, this time we will set a maximum number of iterations and take gradient steps until we reach this maximum number. If no maximum number is supplied, the maximum should be set 100 by default. (Use default parameter values in Julia.)

With this in mind, write the following gradient descent function below using your derivative function above. For each step in the gradient descent, we update the weight for each feature before computing our stopping criteria.
"""

# ╔═╡ 958f81a6-71cc-11eb-13e8-61dddeb014a7
function regression_gradient_descent(f_matrix::Matrix{T},
		output::Vector{T}, init_weights::Vector{T}, η::T, l2_penalty;
		max_iter=100) where {T <: Real}
    weights = copy(init_weights)
	iter = 0

    while iter < max_iter
		iter += 1
		preds = predict_output(f_matrix, weights) # compute the predictions
        errors = preds .- output                  # compute the errors

		## Update the weights
        for ix ∈ 1:length(weights)                # loop over each weight
            ## Recall that feature_matrix[:, ix] is the feature column
		    ## associated with weights[i]
            deriv_ix = feature_derivative_ridge(errors, f_matrix[:, ix],
				weights[ix],
				l2_penalty;
				is_constant=ix == 1)

			## subtract the step size times the derivative from the CURRENT weight
            weights[ix] -= η * deriv_ix
		end
	end
    weights
end

# ╔═╡ 336f9d62-73c0-11eb-1861-3fb709350f2a
md"""
### Visualizing effect of L2 penalty

The L2 penalty gets its name because it causes weights to have small L2 norms than otherwise. Let's see how large weights get penalized. Let us consider a simple model with 1 feature:
"""

# ╔═╡ e8757854-73c9-11eb-2fb9-c10ac86bf9f0
begin
	sales_train, sales_test = train_test_split(sales; shuffled=false);
	first(sales_train, 3)
end

# ╔═╡ e9114fbe-73cd-11eb-052b-15279aaada23
begin
	s_features = [:sqft_living]
	s_output = :price

	## use: sales_train, sales_test
	(s_feature_mtr, s_output_tr) = get_data(sales_train, s_features, s_output)
	(s_feature_mte, s_output_te) = get_data(sales_test, s_features, s_output)

	s_init_weights = Float64[0., 0.]
	η = 1e-12
	max_iter=1000

	s_l2_penalty = 0.0    # no penalty
	h_l2_penalty = 1e11   # high penalty

	size(s_feature_mtr), size(s_output_tr), size(s_feature_mte), size(s_output_te)
end

# ╔═╡ 47de5248-73c0-11eb-2a0c-85982a16d3a3
s_weights_0_penalty = regression_gradient_descent(s_feature_mtr,
		Vector{Float64}(s_output_tr), s_init_weights, η, s_l2_penalty;
		max_iter)

# ╔═╡ 5cb2aee6-73c3-11eb-0eb0-cfbfcb4c145a
s_weights_h_penalty = regression_gradient_descent(s_feature_mtr,
		Vector{Float64}(s_output_tr), s_init_weights, η, h_l2_penalty;
		max_iter)

# ╔═╡ c20a95dc-73ce-11eb-072e-cfdf28bb8412
(s_weights_0_penalty, s_weights_h_penalty)

# ╔═╡ 87b80b2c-73c3-11eb-3175-0dedf19b1858
begin
	scatter(s_feature_mtr, s_output_tr, color=[:orange], marker=".")
	plot!(s_feature_mtr, predict_output(s_feature_mtr, s_weights_0_penalty), color=[:blue])
	plot!(s_feature_mtr, predict_output(s_feature_mtr, s_weights_h_penalty), color=[:green])
end

# ╔═╡ a2b905e0-73cb-11eb-0aef-79f5e6941f30
md"""
Compute the RSS on the *test* data for the following three sets of weights:
  1. The initial weights (all zeros)
  1. The weights learned with no regularization
  1. The weights learned with high regularization

Which weights perform best?
"""

# ╔═╡ b8f7c40e-73cb-11eb-1e61-8f84439c8cb2
md"""

moved to `utils.jl`

```julia
function (X::Matrix{T}, y, weights::Vector{T})
    preds = predict_output(X, weights)
	sum((preds .- y) .^ 2)
end
```
"""

# ╔═╡ cd05b5fe-73cc-11eb-35a4-952fa79374f4
begin
	## RSS initial weigths
	v0 = calc_rss(s_feature_mte, s_output_te, s_init_weights)

	## RSS weigths no regularization
	v1 = calc_rss(s_feature_mte, s_output_te, s_weights_0_penalty)

	## RSS weigths with high regularization
	v2 = calc_rss(s_feature_mte, s_output_te, s_weights_h_penalty)

	with_terminal() do
		@printf("rss init. weights:     %1.5e\n", v0)
		@printf("rss no reg. weights:   %1.5e\n", v1)
		@printf("rss high reg. weights: %1.5e\n", v2)
	end
end

# ╔═╡ 92029980-73cd-11eb-1de5-9fd4c888d352
md"""

**Quiz Questions**

1. What is the value of the coefficient for `sqft_living` that you learned with no regularization, rounded to 1 decimal place?  What about the one with high regularization?
   - 263.0 (no regularization), 68.7 (with high regularization)

2. Comparing the lines you fit with the with no regularization versus high regularization, which one is steeper?
   - the blue line (with no regularization)

3. What are the RSS on the test data for each of the set of weights above (initial, no regularization, high regularization)?
   - cf. above
"""

# ╔═╡ 27319a50-73d1-11eb-13ee-5b868486e887
md"""
### Running a multiple regression with L2 penalty

Let us now consider a model with 2 features: `[:sqft_living, :sqft_living15]`.

"""

# ╔═╡ 4def015a-73d1-11eb-07be-af419657dd3d
begin
	m_features = [:sqft_living, :sqft_living15]
	m_output = :price

	## use: sales_train, sales_test
	(m_feature_mtr, m_output_tr) = get_data(sales_train, m_features, m_output)
	(m_feature_mte, m_output_te) = get_data(sales_test, m_features, m_output)

	m_init_weights = Float64[0., 0., 0.]

	## As above for the following:
	# η = 1e-12
	# max_iter=1000
	# s_l2_penalty = 0.0    # no penalty
	# h_l2_penalty = 1e11   # high penalty
end

# ╔═╡ 505bc8cc-73d1-11eb-104b-f1fe8928f397
m_weights_0_penalty = regression_gradient_descent(m_feature_mtr,
		Vector{Float64}(m_output_tr), m_init_weights, η, s_l2_penalty;
		max_iter)

# ╔═╡ 502fc1c8-73d1-11eb-3102-4303736bcc5f
m_weights_h_penalty = regression_gradient_descent(m_feature_mtr,
		Vector{Float64}(m_output_tr), m_init_weights, η, h_l2_penalty;
		max_iter)

# ╔═╡ 4ff2dc02-73d1-11eb-333d-07031c47234b
begin
	## RSS initial weigths
	mv0 = calc_rss(m_feature_mte, m_output_te, m_init_weights)

	## RSS weigths no regularization
	mv1 = calc_rss(m_feature_mte, m_output_te, m_weights_0_penalty)

	## RSS weigths with high regularization
	mv2 = calc_rss(m_feature_mte, m_output_te, m_weights_h_penalty)

	with_terminal() do
		@printf("rss init. weights:     %1.5e\n", mv0)
		@printf("rss no reg. weights:   %1.5e\n", mv1)
		@printf("rss high reg. weights: %1.5e\n", mv2)
	end
end

# ╔═╡ 4fdf127a-73d1-11eb-1b41-67a0caaa1eca
begin
	## house price for the 1st house in the test set using the no regularization
	pred_0_no_reg = predict_output(m_feature_mte, m_weights_0_penalty)[1]

	## house price for the 1st house in the test set using high regularization
	pred_0_high_reg = predict_output(m_feature_mte, m_weights_h_penalty)[1]

	with_terminal() do
		@printf("no reg. price: %6.2f / price rounded to nearest dollar: %6.2f\n",
			pred_0_no_reg, round(pred_0_no_reg; digits=0))

		@printf("high reg. price: %6.2f / price rounded to nearest dollar: %6.2f\n",
			pred_0_high_reg, round(pred_0_high_reg; digits=0))

		@printf("actual price: %6.2f\n", m_output_te[1])
	end
end

# ╔═╡ 4fb0a88e-73d1-11eb-3eea-339d461b9e17
abs(pred_0_high_reg - m_output_te[1]), abs(pred_0_no_reg - m_output_te[1])

# ╔═╡ 4f8574ac-73d1-11eb-2b52-dd15c43a1c73
md"""
**Quiz Questions**

 1. What is the value of the coefficient for `sqft_living` that you learned with no regularization, rounded to 1 decimal place?  What about the one with high regularization?
   - 189.3 (no reg), 55.2 (high reg)

 2. What are the RSS on the test data for each of the set of weights above (initial, no regularization, high regularization)?
   -  cf. above

 3. We make prediction for the first house in the test set using two sets of weights (no regularization vs high regularization). Which weights make better prediction *for that particular house*?
   - price from no reg. model

"""

# ╔═╡ Cell order:
# ╟─9dcffb1c-71c1-11eb-2402-39eaa1f1b5e1
# ╟─c39f9348-71c1-11eb-1691-1d28c6b1d916
# ╠═1bcf25d8-71c2-11eb-05ca-6fe7c4e1b24f
# ╠═3e5d907e-71c2-11eb-2eb3-4375fb21b884
# ╠═4c736078-71c2-11eb-225b-816745199225
# ╠═d17897d0-7540-11eb-1d5e-191caf244165
# ╟─35b5d05a-71cc-11eb-15a9-852c5a534fab
# ╟─a2928dc6-71c2-11eb-3724-29dc1448a0f3
# ╟─21635b6e-71c8-11eb-2272-edec98c58339
# ╠═96e7042c-73be-11eb-3ac5-9f9998b8c8ff
# ╠═f377f6d4-71c9-11eb-0249-f3d63e87062b
# ╠═59e3f22a-73bd-11eb-1c0d-150d7d462676
# ╟─239ff6ca-71cc-11eb-341a-df1453d60a15
# ╠═958f81a6-71cc-11eb-13e8-61dddeb014a7
# ╟─336f9d62-73c0-11eb-1861-3fb709350f2a
# ╠═e8757854-73c9-11eb-2fb9-c10ac86bf9f0
# ╠═e9114fbe-73cd-11eb-052b-15279aaada23
# ╠═47de5248-73c0-11eb-2a0c-85982a16d3a3
# ╠═5cb2aee6-73c3-11eb-0eb0-cfbfcb4c145a
# ╠═c20a95dc-73ce-11eb-072e-cfdf28bb8412
# ╠═87b80b2c-73c3-11eb-3175-0dedf19b1858
# ╟─a2b905e0-73cb-11eb-0aef-79f5e6941f30
# ╟─b8f7c40e-73cb-11eb-1e61-8f84439c8cb2
# ╠═cd05b5fe-73cc-11eb-35a4-952fa79374f4
# ╟─92029980-73cd-11eb-1de5-9fd4c888d352
# ╟─27319a50-73d1-11eb-13ee-5b868486e887
# ╠═4def015a-73d1-11eb-07be-af419657dd3d
# ╠═505bc8cc-73d1-11eb-104b-f1fe8928f397
# ╠═502fc1c8-73d1-11eb-3102-4303736bcc5f
# ╠═4ff2dc02-73d1-11eb-333d-07031c47234b
# ╠═4fdf127a-73d1-11eb-1b41-67a0caaa1eca
# ╠═4fb0a88e-73d1-11eb-3eea-339d461b9e17
# ╟─4f8574ac-73d1-11eb-2b52-dd15c43a1c73


# This files stores constraints or concepts that are not yet available within usual_constraints
# The list is based on https://arxiv.org/abs/2009.00514

# SECTION - Generic constraints
# intension: predicate over variables (DONE)
# extension: supports or conflicts with a set of configurations

# SECTION - Constraints defined from Languages
# regular: recognize a regular automaton
# mdd: follow a multi-values decision diagram

# SECTION - Comparison-based Constraints
# all_different: done in Constraints.jl
# candidate: all_different_except (DONE)
# all_equal: done in Constraints.jl
# all_equal_param: done (not in XCSP3)
# ordered: increasing (current ordered in Constraints.jl), strictly_increasing (DONE), decreasing, strictly_decreasing

# SECTION - Counting and Summing Constraints
# sum or linear or linear_sum: sum_equal_param (DONE)
# candidate: more generic linear equation (DONE)
# count: count, among, at_least, at_most
# n_values: n_values, at_least_n_values, at_most_n_values
# candidate: n_values_except (ignores the values in except)
# cardinality: (global_)cardinality

# SECTION - Connection Constraints
# maximum: all values lower than param (DONE)

# minimum
concept_minimum(x; param) = minimum(x) < param

const _minimum = Constraint(
    concept = concept_minimum,
    error = make_error(:minimum),
)

push!(BENCHED_CONSTRAINTS, :minimum => _minimum)

# element: param ∈ values(variables)
# channel: if xᵢ = paramⱼ then xⱼ = paramᵢ

# SECTION - Packing and Scheduling Constraints
# no_overlap: disjunctive (no_overlap_1D), diffn (kD)
# cummulative: check pdf :p

# SECTION - Constraints on Graphs
# circuit: verify that the sequence of vertices forms a circuit (where xᵢ = j => (i,j)∈E)

# SECTION - Elementary Constraints
# instantiation: x == param

## no overlap 1D and 2D (commented)
function concept_no_overlap(x; param)
    for i in 1:(length(x)-1), j in i:length(x)
        x[i] + param > x[j] && x[j] + param > x[i] && return false
    end
    return true
end

# function concept_no_overlap(x; param::AbstractVector)
#     for i in 1:(length(x)-1), j in i:length(x)
#         x[i] + param[i] > x[j] && x[j] + param[j] > x[i] && return false
#     end
#     return true
# end

const no_overlap = Constraint(
    concept = concept_no_overlap,
    error = make_error(:no_overlap),
)

push!(BENCHED_CONSTRAINTS, :no_overlap => no_overlap)

## maximum
concept_maximum(x; param) = maximum(x) < param

const _maximum = Constraint(
    concept = concept_maximum,
    error = make_error(:maximum),
)

push!(BENCHED_CONSTRAINTS, :maximum => _maximum)

## intension
concept_intension(x; param::Base.Function) = param(x)

const _intension = Constraint(
    concept = concept_intension,
    error = make_error(:intension),
)

push!(BENCHED_CONSTRAINTS, :intension => _intension)

## linear
concept_linear(x1, op1::Base.Function, x2, op2::Base.Function, param) = op2(op1(x1,x2),param)

const _linear = Constraint(
    concept = concept_linear,
    error = make_error(:intension),
)

push!(BENCHED_CONSTRAINTS, :linear => _linear)

## all_different_except

concept_all_different_except(x, s) = allunique(filter(element -> element ∉ s , x))

const _all_different_except = Constraint(
    concept = concept_all_different_except,
    error = make_error(:all_different_except),
)

push!(BENCHED_CONSTRAINTS, :all_different_except => _all_different_except)

## strictly_increasing

concept_strictly_increasing(x; param=nothing) = issorted(x) && length(x) == length(Set(x))

const _strictly_increasing = Constraint(
    concept = concept_strictly_increasing,
    error = make_error(:strictly_increasing),
)

push!(BENCHED_CONSTRAINTS, :strictly_increasing => _strictly_increasing)


#= root3
function root3_2(x::T) where {T<:AbstractFloat}
    hi = cbrt(x)
    chi, clo = powr3_2(hi)
    d = (-chi + x) - clo
    lo = d / (T(3)*hi*hi)
    return hi, lo
end

@inline root3_(x::T) where {T<:AbstractFloat} = root3_2(x)
=#


##  SCRIPT NAME:   solveValueFunction.jl
##  DATE:          January 2015
#
#   This script solves the household problem by value function iteration.
#   The function obtains the value function for each time period and the policy
#   function (i.e. optimal next-period asset choice). From there I can work out
#   the optimal consumption level. The solution approach taken is by backwards
#   recursion. The optimization each period is carried out using NLopt.
#
#
##  ALEXANDROS THELOUDIS, UCL

##  -------------------------------------------------------------------------------------------
cd("C:\\Dropbox\\JuliaMeeting3")

#   Load packages:
using Grid
using NLopt

#   Load functions:
include("objectiveFunction.jl")
include("utilityH.jl")
include("utilityW.jl")
include("equalityConstraints1.jl")
include("mUtilH_cH.jl")
include("mUtilW_cW.jl")
include("mUtilH_k.jl")
include("mUtilW_k.jl")

##  ---------------------------------------------------------------------------------------------
#   Economy constants:
const r                  = 0.01            # Deterministic interest rate
const beta               = 0.98            # Discount factor
const cminH              = .1              # Min consumption husband (normalised /1000)
const cminW              = .1              # Min consumption wife (normalised /1000)
const kmin               = .1              # Min (inputs to) public consumption (normalised /1000)
const maxAnnWorkingHrsH  = 2.0             # 40 hrs/week times 50 weeks/year (normalised /1000)
const maxAnnWorkingHrsW  = 2.0             # 40 hrs/week times 50 weeks/year (normalised /1000)

#   Pareto weights related constants:
const paretoH            = 0.50            # Pareto weight husband
const paretoW            = 1 - paretoH     # Pareto weight wife

#   Preference-related constants (husband):
const gammaCH            = 1.5             # CRRA on private consumption
const gammaKH            = 1.5             # CRRA on public consumption
const alphaCH            = 0.5             # Weight of private consumption

#   Preference-related constants (wife):
const gammaCW            = 1.5             # CRRA on private consumption
const gammaKW            = 1.5             # CRRA on public consumption
const alphaCW            = 0.5             # Weight of private consumption

#   Solution constant:
const lengthA1           = 50              # Length of asset grid tomorrow

##  ---------------------------------------------------------------------------------------------
#   Current and future state:
A         = 10          # current assets
wH        = 5           # husband's wage today
wW        = 5           # wife's wage today
priceC    = 1.5         # price of private consumption

#   Calculate min and max assets tmr:
# ..lowest assets: cannot die with debts
lbA1      = 0
# ..highest assets: everything saved
ubA1      = (A + wH*maxAnnWorkingHrsH + wW*maxAnnWorkingHrsW - priceC*(cminH + cminW) -kmin)*(1+r)
# ..generate an asset grid for tmr assets (to be used in the EVal interpolation)
xcoord   = range(lbA1, (ubA1 - lbA1)/(lengthA1-1), lengthA1)

#   Expected Value function tomorrow (tmr is terminal period):
EVal1     = zeros(lengthA1)

#  ------------------------------------------------------------------------------------
#  Declare the relationship (linear) between next period's assets and Eval1:
VA!i     = CoordInterpGrid(xcoord, EVal1, BCnil, InterpLinear)

#  ------------------------------------------------------------------------------------
#  Define function to be optimized:

# Set properties of NL optimization problem:
valueFunctionOpt = Opt(:LD_SLSQP, 4)          # algorithm and dimensionality (:LN_COBYLA  :LD_SLSQP)
xtol_rel!(valueFunctionOpt,1e-6)   # tolerance

# Set bounds of optimization:
low_bound = [lbA1 ; cminH ; cminW ; kmin]     # lower bounds
up_bound  = [ubA1 ; Inf   ; Inf   ; Inf ]     # upper bounds
lower_bounds!(valueFunctionOpt,low_bound)     # lower bounds
upper_bounds!(valueFunctionOpt,up_bound)      # upper bounds

# Set objective function, min/max mode, and constraints:
max_objective!(valueFunctionOpt, (x,g) -> objectiveFunction(x, g, A, wH, wW, VA!i))
equality_constraint!(valueFunctionOpt, (x,g) -> equalityConstraints1(x, g, A, wH, wW, priceC), 1e-6)

# Set starting value and initialize optimization:
start_point = [ (ubA1 + lbA1)/2 ;
                (ubA1 - lbA1)/3 ;
                (ubA1 - lbA1)/3 ;
                (ubA1 - lbA1)/3 ]
startpt_valueFunctionOpt     = start_point             # starting point
(fval,x_optimal,return_flag) = optimize(valueFunctionOpt,[startpt_valueFunctionOpt])

# Store results of optimization:
policyA1 = maximum(x_optimal[1])    # policy for next period's assets
policyCH = maximum(x_optimal[2])    # policy for husband's consumption
policyCW = maximum(x_optimal[3])    # policy for wife's consumption
policyK  = maximum(x_optimal[4])    # policy for public consumption
Val      = fval                     # value function today

# Print results:
println("NLopt finished solving. It found an fval=",Val," with CH=",policyCH,", CW=",policyCW,", K=",policyK," and next period's assets are ",policyA1)


using Distributions
import StatsBase.ecdf
#using Plots

function s_weibull( x::Float64, k::Float64=1.0, λ::Float64=1.0)
    return 1.0-cdf( Weibull(k, λ), x)
end

#using PyPlot
# figure()
# plot( collect(0.:1.0/100.:5.),[s_weibull(i, 2.2, 1.75) for i in collect(0.:1.0/100.:5.)], linewidth=1.1, label = "\$ Y\\^{(0)} survival with \$l=1\$"  )
# plot( collect(0.:1.0/100.:5.),[s_weibull(i, 1.2, 2.2) for i in collect(0.:1.0/100.:5.)], linewidth=1.1, label = "\$ Y^{(1)} survival with \$l=0\$"   )
# plot( collect(0.:1.0/100.:5.),[s_weibull(i, 5.3, 1.5) for i in collect(0.:1.0/100.:5.)], linewidth=1.1, label = "\$ Y^{(2)} survival with \$l=0\$" )
# ylabel("\$ S(t) \$")
# xlabel("\$ t \$")
# title("Simulated data true survival functions")
# legend(loc=1)
# savefig("weibulls.png")

c_sw = [5.0,2.2]
c_ws = [1.5,3.4]
c_m = [2.1,2.0]

function s_comprisk_weibull( x::Float64, k₁::Float64=1.0, k₂::Float64=1.0,
                        λ₁::Float64=1.0,  λ₂::Float64=1.0)
    return exp( log( ccdf( Weibull(k₁, λ₁), x) ) + log( ccdf( Weibull(k₂, λ₂), x) ) )
end

# plot( collect(0.:1.0/100.:5.),[s_comprisk_weibull(i, 5.0, 1.2, 1.8, 2.2) for i in collect(0.:1.0/100.:5.)] )
# plot!( collect(0.:1.0/100.:5.),[s_comprisk_weibull(i, 5.0, 4.3, 1.8, 1.5) for i in collect(0.:1.0/100.:5.)] )

function sim_study_sampler_comprisk( n::Int64, δ::Float64, thr::Float64,
        k₀::Float64=1.0, λ₀::Float64=1.0, k₁::Float64=1.0, λ₁::Float64=1.0,
        k₂::Float64=1.0, λ₂::Float64=1.0,
        μ_co::Float64=1.0, σ_co::Float64=0.75, d_co::Float64=0.0, u_co::Float64=2.0)
    samp_Z = zeros( Float64, n)
    samp_δ = zeros( Float64, n)
    samp_Y = zeros( Float64, n)
    for i = 1:n
        samp_Z[i] = rand( TruncatedNormal( μ_co, σ_co, d_co, u_co) )
        if samp_Z[i] < thr
            samp_Y[i] = min( rand( Weibull( k₁, λ₁/(1.0-δ)^(1.0/k₁) ) ),
                                 rand( Weibull( k₀, λ₀/(δ)^(1.0/k₀) ) ) )
        else
            samp_Y[i] = min( rand( Weibull( k₂, λ₂/(1.0-δ)^(1.0/k₂) ) ),
                                 rand( Weibull( k₀, λ₀/(δ)^(1.0/k₀) ) ) )
        end
    end
    return [ samp_Y, samp_Z ]
end

function true_surv_comprisk( x::Float64, z::Float64, δ::Float64, thr::Float64,
    k₀::Float64=1.0, λ₀::Float64=1.0, k₁::Float64=1.0, λ₁::Float64=1.0,
    k₂::Float64=1.0, λ₂::Float64=1.0)
    if z <= thr
        return s_comprisk_weibull( x, k₁, k₀,
        λ₁/(1.0-δ)^(1.0/k₁), λ₀/(δ)^(1.0/k₀) )
    else
        return s_comprisk_weibull( x, k₂, k₀,
        λ₂/(1.0-δ)^(1.0/k₂), λ₀/(δ)^(1.0/k₀) )
    end
end

# Kaplan-Meier estiamtor function
function KMest(tiempo, muerte)
    t_KMest = sort( unique( tiempo ) )
    if t_KMest[1] != 0
        t_KMest = [0; t_KMest]
    end
    n_KMest = [count( i->(i >= j), tiempo) for j in t_KMest]
    d_KMest = [sum( muerte[ findall(j->j==i,tiempo) ] ) for i in t_KMest]
    s_KMest = cumprod( 1 .- d_KMest./n_KMest )
    return t_KMest, s_KMest
end

# plot( collect(0.:1.0/100.:5.), [true_surv_comprisk( i, 0.5, 0.1, 1.0, 5.0, 1.8, 1.2, 2.2, 4.3, 1.5) for i in collect(0.:1.0/100.:5.)])
# plot!( collect(0.:1.0/100.:5.), [true_surv_comprisk( i, 1.5, 0.1, 1.0, 5.0, 1.8, 1.2, 2.2, 4.3, 1.5) for i in collect(0.:1.0/100.:5.)])


using Cubature
# using Cuba
# using MonteCarloIntegration

include("vegas.jl")

function koopman(g,prob,u0,p,args...;kwargs...)
    g(solve(remake(prob,u0=u0,p=p),args...;kwargs...))
end

function koopman_cost(u0s,ps,g,prob,args...;maxevals=0,
                      ireltol = 1e-2, iabstol=1e-2, use_vegas=false, kwargs...)
    n = length(u0s)
    function _f(x)
    u0 = x[1:n]
    p  = x[n+1:end]
    k = koopman(g,prob,u0,p,args...;kwargs...)
    w = prod(pdf(a,b) for (a,b) in zip(u0s,u0))*
        prod(pdf(a,b) for (a,b) in zip(ps,p))
    k*w
    end
    xs = [u0s;ps]
    st = minimum.(xs)
    en = maximum.(xs)
    @show st
    @show en

    if use_vegas
        # vegas((x,f) -> f[1] = _f(st .+ x ./ (en .- st)) * sum(en .- st), rtol = ireltol)
        vegas(_f, st, en)
    else
        hcubature(_f, minimum.(xs), maximum.(xs);
            reltol=ireltol, abstol=iabstol, maxevals = maxevals)
    end
end

function montecarlo_cost(u0s,ps,g,prob,args...;num_monte,kwargs...)
  prob_func = function (prob,i,repeat)
    remake(prob,u0=rand.(u0s),p=rand.(ps))
  end
  output_func = (sol,i) -> (g(sol),false)
  monte_prob = MonteCarloProblem(prob;
                                 output_func = output_func,
                                 prob_func = prob_func)
  mean(solve(monte_prob,args...;num_monte=num_monte,kwargs...).u)
end


using Atum
using Atum.Euler

using PGFPlotsX
using StaticArrays: SVector

function sod(law, x⃗)
  FT = eltype(law)
  ρ = x⃗[1] < 1 // 2 ? 1 : 1 // 8
  ρu⃗ = SVector(FT(0))
  p = x⃗[1] < 1 // 2 ? 1 : 1 // 10
  ρe = Euler.energy(law, ρ, ρu⃗, p)
  SVector(ρ, ρu⃗..., ρe)
end

import Atum: boundarystate
boundarystate(law::EulerLaw, n⃗, x⃗, q⁻, _) = sod(law, x⃗)

function run(A, FT, N, K)
  Nq = N + 1

  law = EulerLaw{FT, 1}()

  cell = LobattoCell{FT, A}(Nq)
  v1d = range(FT(0), stop=FT(1), length=K+1)
  grid = brickgrid(cell, (v1d,); periodic=(false,))

  dg = ESDGSEM(; law, cell, grid,
               volume_numericalflux = EntropyConservativeFlux(),
               surface_numericalflux = RusanovFlux())

  cfl = FT(1 // 4)
  dt = cfl * step(v1d) / N / Euler.soundspeed(law, FT(1), FT(1))
  timeend = FT(2 // 10)

  q = sod.(Ref(law), points(grid))

  odesolver = LSRK54(dg, q, dt)
  solve!(q, timeend, odesolver)

  @pgf begin
    ρ, ρu, ρe = components(q)
    p = Euler.pressure.(Ref(law), ρ, ρu, ρe)
    u = ρu ./ ρ
    x = vec(first(components(points(grid))))

    fig = @pgf GroupPlot({group_style= {group_size="2 by 2"}})
    ρ_plot = Plot({no_marks}, Coordinates(x, vec(ρ)))
    u_plot = Plot({no_marks}, Coordinates(x, vec(u)))
    E_plot = Plot({no_marks}, Coordinates(x, vec(ρe)))
    p_plot = Plot({no_marks}, Coordinates(x, vec(p)))

    push!(fig, {}, ρ_plot)
    push!(fig, {}, u_plot)
    push!(fig, {}, E_plot)
    push!(fig, {}, p_plot)

    path = mkpath(joinpath("output", "euler", "sod"))
    pgfsave(joinpath(path, "sod.pdf"), fig)
  end
end

let
  A = Array
  FT = Float64
  N = 4
  K = 32

  errf = run(A, FT, N, K)
end


# ==================
# Dependencies
# ==================
using ReachabilityAnalysis, CarlemanLinearization
using Plots, LaTeXStrings, LinearAlgebra, SparseArrays
using Plots.PlotMeasures
using DifferentialEquations

using CarlemanLinearization: _error_bound_specabs_R

include("../utils.jl")

# ==================
# Model definition
# ==================

# the following are constants of the model
# the domain of x is [-L0/2, L0/2]
L0 = 1.0     # domain length
ν = 0.05     # kinematic viscosity
U0 = 1.0
Re = 20
ν = U0 * L0 / Re
Tnl = L0 / U0
Tmax = Tnl/2

# ================================
# Spatial flowpipe visualization
# ================================

# no-op
ReachabilityAnalysis.convexify(R::ReachabilityAnalysis.AbstractReachSet) = R

function flowpipe_spatial(sol, t, n)
    solt = sol(t) |> convexify

    xdom = range(-L0/2, L0/2, length=n)
    xdomred = xdom[2:end-1]

    X = [set(overapproximate(Projection(solt, i:i), Interval)) for i in 1:dim(sol)]
    R = ReachSet(CartesianProductArray(X), tspan(solt))
    arr = [Interval(xdomred[i]) × CartesianProductArray(X).array[i] for i in 1:length(xdomred)]
    arr = vcat(Interval(xdom[1]) × Interval(0), arr, Interval(xdom[end]) × Interval(0))
    U = UnionSetArray(arr)
    Uch = [ConvexHull(U.array[i], U.array[i+1]) for i in 1:length(U.array)-1] |> UnionSetArray

    return Uch
end

# model for solution with TM
@taylorize function burgers!(du, u, p, t)
    local nf = length(u)
    n = nf + 2
    local Δx = L0/(n-1)
    local c1 = ν/Δx^2
    local c2 = 1/(4Δx)

    du[1] = c1*(-2*u[1]+u[2]) - c2*(u[2]^2)
    du[nf] = c1*(u[nf-1]-2*u[nf]) - c2*(-u[nf-1]^2)

    for i in 2:nf-1
        du[i] = c1*(u[i-1]-2*u[i]+u[i+1]) - c2*(u[i+1]^2 - u[i-1]^2)
    end
end

# ====================
# Carleman method
# ====================

using DynamicPolynomials

# we remove the endpoints
function burgers_carlin(nf)
    n = nf + 2
    Δx = L0/(n-1)
    c1 = ν/Δx^2
    c2 = 1/(4Δx)

    # fix the endpoints
    n = nf # rename to simplify the rest of the code
    dv = c1 * fill(-2.0, n)
    ev = c1 * fill(1.0, n-1)

    F1 = SymTridiagonal(dv, ev)

    F2 = zeros(n, n^2)
    x, = @polyvar x[1:n]
    y = kron(x, x)

    idx = findfirst(y, x[2]^2)
    F2[1, idx] = 1.0

    idx = findfirst(y, x[n]^2)
    F2[n, idx] = -1.0

    for i in 2:n-1
        idx = findfirst(y, x[i-1]^2)
        F2[i, idx] = -1.0

        idx = findfirst(y, x[i+1]^2)
        F2[i, idx] = 1.0
    end
    F2 .*= -c2
    F2 = sparse(F2)

    return F1, F2
end

# ===================
# Solution method
# ===================

function _solve_burgers_carlin(; n, N, width0=0, δ=0.02)

    nf = n-2
    F1, F2 = burgers_carlin(nf)

    A = build_matrix(F1, F2, N)

    xdom = range(-L0/2, L0/2, length=n)
    c0 = Singleton(-U0*sin.(2*π*xdom[2:end-1]/L0))
    if width0 == 0
        X0 = overapproximate(c0, Hyperrectangle);
    else
        H0 = Hyperrectangle(zeros(nf), fill(width0, nf))
        X0 = overapproximate(c0 ⊕ H0, Hyperrectangle)
    end

    dirs = _template(n=nf, N=N)
    # amodel = Forward(inv=true, setops=:lazy, sih=:concrete) # FIXME set options
    alg = LGG09(δ=δ, template=dirs, approx_model=Forward())

    sol_carlin = _solve_CARLIN(X0, F1, F2; alg=alg, N=N, T=Tmax, bloat=false, kron_pow_algorithm="symbolic");
    sol_carlin_spatial = flowpipe_spatial(sol_carlin, Tnl/2, n);
end

# ===================
# TM solution method
# ===================

function _solve_burgers_TM(; n, width0=0)

    nf = n - 2

    xdom = range(-L0/2, L0/2, length=n)
    c0 = Singleton(-U0*sin.(2*π*xdom[2:end-1]/L0))
    if width0 == 0
        X0 = overapproximate(c0, Hyperrectangle)
    else
        H0 = Hyperrectangle(zeros(nf), fill(width0, nf))
        X0 = overapproximate(c0 ⊕ H0, Hyperrectangle)
    end

    prob = IVP(BlackBoxContinuousSystem(burgers!, nf), X0)
    solTM = solve(prob, tspan=(0.0, Tmax), alg=TMJets20(), ensemble=true, trajectories=1);
    solzTM = overapproximate(solTM, Zonotope);

    sol_spatial = flowpipe_spatial(solzTM, Tnl/2, n);
end

function _get_R(; n)
    nf = n-2
    F1, F2 = burgers_carlin(nf);
    xdom = range(-L0/2, L0/2, length=n)
    x0 = -U0*sin.(2*π*xdom[2:end-1]/L0);
    R, Re_λ1 = _error_bound_specabs_R(x0, Matrix(F1), Matrix(F2); check=true)
end

# ================================
# Results
# ================================

n = 10

# Initial point
_solve_burgers_TM(n=n); tt = @elapsed _solve_burgers_TM(n=n)
print(io, "Burgers, TMJets, initial point, -, $(tt)\n")

_solve_burgers_carlin(n=n, N=2, δ=0.01); tt = @elapsed _solve_burgers_carlin(n=n, N=2, δ=0.01)
print(io, "Burgers, Carleman, initial point, N=2, $(tt)\n")

_solve_burgers_carlin(n=n, N=3, δ=0.01); tt = @elapsed _solve_burgers_carlin(n=n, N=3, δ=0.01)
print(io, "Burgers, Carleman, initial point, N=3, $(tt)\n")

# Initial set
width0 = 0.03
_solve_burgers_TM(n=n, width0=width0); tt = @elapsed _solve_burgers_TM(n=n, width0=width0)
print(io, "Burgers, TMJets, initial set, -, $(tt)\n")

_solve_burgers_carlin(n=n, N=2, δ=0.01, width0=width0); tt = @elapsed _solve_burgers_carlin(n=n, N=2, δ=0.01, width0=width0)
print(io, "Burgers, Carleman, initial set, N=2, $(tt)\n")

_solve_burgers_carlin(n=n, N=3, δ=0.01, width0=width0); tt = @elapsed _solve_burgers_carlin(n=n, N=3, δ=0.01, width0=width0)
print(io, "Burgers, Carleman, initial set, N=3, $(tt)\n")

# ==================
# Figures
# ==================

# initial point
function figure_burger_point()

    fig = plot(legend=:topright, xlab=L"x", ylab=L"u(x, 0.5)",
               legendfontsize=25,
               tickfont=font(25, "Times"),
               guidefontsize=25,
               xguidefont=font(35, "Times"),
               yguidefont=font(35, "Times"),
               xtick = xticklatex([-0.50, -0.25, 0.0, 0.25, 0.50], 2),
               ytick = xticklatex([-1.0, -0.5, 0.0, 0.5, 1.0], 2),
               xlims=(-0.5, 0.5),
               ylims=(-1.05, 1.0),
               bottom_margin=10mm, left_margin=6mm, right_margin=6mm, top_margin=3mm, size=(900, 600))

    # model values
    n = 10

    # initial condition (smooth)
    xdom = range(-L0/2, L0/2, length=200)
    c0 = Singleton(-U0*sin.(2*π*xdom/L0))
    plot!(fig, xdom, element(c0), seriestype=:path, c=:black, marker=:none, ls=:dash, lab=L"u(x, 0)", lw=2.0, alpha=1.)

    # initial cond, samples
    xdom = range(-L0/2, L0/2, length=n)
    c0 = Singleton(-U0*sin.(2*π*xdom/L0))
    plot!(fig, xdom, element(c0), seriestype=:path, c=:grey, marker=:square, ls=:dash, lab="", lw=2.0, alpha=1.)

    solN2_carlin = _solve_burgers_carlin(n=n, N=2, δ=0.01)
    solN3_carlin = _solve_burgers_carlin(n=n, N=3, δ=0.01);

    sol_TM = _solve_burgers_TM(n=n);

    plot!(fig, solN3_carlin, c=:orange, lab=L"N = 3", alpha=1., lc=:orange)
    plot!(fig, solN2_carlin, c=:red, lab=L"N = 2", alpha=1., lc=:red)
    plot!(fig, sol_TM, c=:blue, lab=L"\textrm{TM}", alpha=1., lc=:blue)

    return fig
end

fig = figure_burger_point()
savefig(fig, joinpath(TARGET_FOLDER, "figure_burger_point.pdf"))

# ==================
# initial set
# ==================

function figure_burger_set()

    fig = plot(legend=:topright, xlab=L"x", ylab=L"u(x, 0.5)",
               legendfontsize=25,
               tickfont=font(25, "Times"),
               guidefontsize=25,
               xguidefont=font(35, "Times"),
               yguidefont=font(35, "Times"),
               xtick = xticklatex([-0.50, -0.25, 0.0, 0.25, 0.50], 2),
               ytick = xticklatex([-1.0, -0.5, 0.0, 0.5, 1.0], 2),
               xlims=(-0.5, 0.5),
               ylims=(-1.05, 1.0),
               bottom_margin=10mm, left_margin=6mm, right_margin=6mm, top_margin=3mm, size=(900, 600))

    # model values
    n = 10
    width0 = 0.03

    # initial condition (smooth)
    xdom = range(-L0/2, L0/2, length=200)
    c0 = Singleton(-U0*sin.(2*π*xdom/L0))
    plot!(fig, xdom, element(c0), seriestype=:path, c=:black, marker=:none, ls=:dash, lab=L"u(x, 0)", lw=2.0, alpha=1.)

    # initial cond, samples
    xdom = range(-L0/2, L0/2, length=n)
    c0 = Singleton(-U0*sin.(2*π*xdom/L0))
    plot!(fig, xdom, element(c0), seriestype=:path, c=:grey, marker=:square, ls=:dash, lab="", lw=2.0, alpha=1.)

    solN2_carlin = _solve_burgers_carlin(n=n, N=2, width0=width0, δ=0.01)
    solN3_carlin = _solve_burgers_carlin(n=n, N=3, width0=width0, δ=0.01);

    sol_TM = _solve_burgers_TM(n=n, width0=width0);

    plot!(fig, solN3_carlin, c=:orange, lab=L"N = 3", alpha=1., lc=:orange)
    plot!(fig, solN2_carlin, c=:red, lab=L"N = 2", alpha=1., lc=:red)
    plot!(fig, sol_TM, c=:blue, lab=L"\textrm{TM}", alpha=1., lc=:blue)

    return fig
end

fig = figure_burger_set()
savefig(fig, joinpath(TARGET_FOLDER, "figure_burger_set.pdf"))


using QuantEcon

srand(42)
lm = LakeModel(d=0.0, b=0.0)
T = 5000                        # Simulation length

α, λ = lm.α, lm.λ
P = [(1 - λ)     λ;
     α       (1 - α)]

mc = MarkovChain(P, [0; 1])     # 0=unemployed, 1=employed
xbar = rate_steady_state(lm)

s_path = simulate(mc, T; init=2)
s_bar_e = cumsum(s_path) ./ (1:T)
s_bar_u = 1 - s_bar_e
s_bars = [s_bar_u s_bar_e]

titles = ["Percent of time unemployed" "Percent of time employed"]

fig, axes = subplots(2, 1, figsize=(10, 8))

for (i, ax) in enumerate(axes)
    ax[:plot](1:T, s_bars[:, i], c="blue", lw=2, alpha=0.5)
    ax[:hlines](xbar[i], 0, T, "r", "--")
    ax[:set](title=titles[i])
    ax[:grid]("on")
end


# Mixture models
distset = [
    MixtureModel([Uniform(), Normal(), LogNormal()]),
    MixtureModel([Uniform(1, 10), Gamma()]),
    MixtureModel([Uniform(1, 100), Normal(1, 1000)]),
    MixtureModel([DiscreteUniform(1, 10), DiscreteUniform(1, 10)])
]

for dist in distset
    println(dist)
    μ = mean(dist)
    σ = std(dist)
    # No convenience call.
    E = expectation(dist)
    @test E(x -> x) ≈ μ
    @test E(x -> x^2) - μ^2 ≈ σ^2
    # Convenience call.
    @test expectation(x -> x, dist) ≈ μ
    @test expectation(x -> x^2, dist) - μ^2 ≈ σ^2
    # Stress tests.
    # Many nodes.
    E2 = expectation(dist, n = 100)
    @test E2(x -> x) ≈ μ
    @test E2(x -> x^2) - μ^2 ≈ σ^2
end

# Addition of expectations
e1 = expectation(Normal())
e2 = expectation(Uniform())
e3 = expectation(Gamma())

@test ((e1 + e2) + e3)(identity) == ((e1 + e2) + e3)(identity) == (e1 + e2 + e3)(identity)
@test (e1 + e2 + e3)(x -> x^2) ≈ 3.333333333333342

# Linear operator behavior.

# Error handling.


using TetGen
using LinearAlgebra

"""
     random_delaunay(;npoints=20)

Create a Delaunay tetrahedralization  from a random set of points.
This is  a tetrahedralization of the convex hull of the points.
If the points are in general position (which we assume here),
it is unique.
"""
function random_delaunay(;npoints=20)
    input=TetGen.RawTetGenIO{Cdouble}(pointlist=rand(3,npoints))
    tetrahedralize(input, "Q")
end

"""
   cube(;vol=1)

Tetrahedralization of cube witn maximum tetrahedron volume.
"""
function cube(;vol=1)
    input=TetGen.RawTetGenIO{Cdouble}()
    input.pointlist=[0 0 0;
                     1 0 0;
                     1 1 0;
                     0 1 0;
                     0 0 1;
                     1 0 1;
                     1 1 1;
                     0 1 1]'

    TetGen.facetlist!(input,[1 2 3 4;
                             5 6 7 8;
                             1 2 6 5;
                             2 3 7 6;
                             3 4 8 7;
                             4 1 5 8]')
    tetrahedralize(input, "pQa$(vol)")
end

"""
   cubewithhole(;vol=1)

Tetrahedralization of cube witn maximum tetrahedron volume.
"""
function cubewithhole(;vol=1)
    input=TetGen.RawTetGenIO{Cdouble}()
    outerpoints=[-1 -1 -1;
                 1 -1 -1;
                 1 1 -1;
                 -1 1 -1;
                 -1 -1 1;
                 1 -1 1;
                 1 1 1;
                 -1 1 1;]'

    input.pointlist=hcat(outerpoints,outerpoints/2)

    outerfacets=[1 2 3 4;
                 5 6 7 8;
                 1 2 6 5;
                 2 3 7 6;
                 3 4 8 7;
                 4 1 5 8;
                 ]'

    input.holelist=[0 0 0;]'

    TetGen.facetlist!(input,hcat(outerfacets,outerfacets.+8))
    tetrahedralize(input, "pQa$(vol)")
end

"""
   cube_localref()

   Tetrahedralization of cube with local refinement callback
"""

function cube_localref()

    tetunsuitable() do pa,pb,pc,pd
        vol=det(hcat(pb-pa,pc-pa,pd-pa))/6
        center=0.25*(pa+pb+pc+pd)-[0.5,0.5,0.5]
        vol> 0.05*norm(center)^2.5
    end

    input=TetGen.RawTetGenIO{Cdouble}()
    input.pointlist=[0 0 0;
                     1 0 0;
                     1 1 0;
                     0 1 0;
                     0 0 1;
                     1 0 1;
                     1 1 1;
                     0 1 1]'

    TetGen.facetlist!(input,[1 2 3 4;
                             5 6 7 8;
                             1 2 6 5;
                             2 3 7 6;
                             3 4 8 7;
                             4 1 5 8]')
    tetrahedralize(input, "pQa")
end

"""
   prism(;vol=1)

Tetrahedralization of a prism with maximum tetrahedron volume.
"""
function prism(vol=2)
    input=TetGen.RawTetGenIO{Cdouble}()
    input.pointlist=[0 0 0;
                     1 0 0;
                     0 1 0;
                     0 0 1;
                     1 0 1;
                     0 1 1]'

    TetGen.facetlist!(input,[[1,2,3],
                             [4,5,6],
                             [1,2,5,4],
                             [2,3,6,5],
                             [3,1,4,6]])

    tetrahedralize(input, "pQa$(vol)")
end

"""
   prism(;vol1=0.01, vol2=0.2)

Tetrahedralization of a prism with two regions with different tet volumes
"""
function material_prism(;vol1=0.01, vol2=0.1)
    input=TetGen.RawTetGenIO{Cdouble}()

    input.pointlist=[0 0 0;
                     1 0 0;
                     0 1 0;
                     0 0 1;
                     1 0 1;
                     0 1 1;
                     0 0 2;
                     1 0 2;
                     0 1 2]'

    TetGen.facetlist!(input,[[1,2,3],
                             [7,8,9],
                             [1,2,5,4],
                             [2,3,6,5],
                             [3,1,4,6],
                             [1,2,5,4].+3,
                             [2,3,6,5].+3,
                             [3,1,4,6].+3])

    input.facetmarkerlist=[1,2,3,3,3,3,3,3]
    input.regionlist=[0.1 0.1 0.5 1 vol1;
                      0.1 0.1 1.5 2 vol2]'

    tetrahedralize(input, "paAqQ")
end

"""
   cutprism(;vol=0.05)

Tetrahedralization of a prism with  a prism cut out in the middle.
This tests the use of facet holes.
"""
function cutprism(;vol=0.05)
    input=TetGen.RawTetGenIO{Cdouble}()
    input.pointlist=[0 0 0;
                     1 0 0;
                     0 1 0;

                     0 0 10;
                     1 0 10;
                     0 1 10;

                     -1 -1 0;  # 7
                     2 -1 0;
                     2 2 0;
                     -1 2 0;

                     -1 -1 10;  # 11
                     2 -1 10;
                     2 2 10;
                     -1 2 10]'

    push!(input.facetlist,RawFacet([Cint[1, 2, 3],Cint[7, 8, 9, 10]], [0.1 0.1 0.0;]))
    push!(input.facetlist,RawFacet([Cint[4, 5, 6],Cint[11,12,13,14]], [0.1 0.1 1.0;]))
    push!(input.facetlist,RawFacet([Cint[1, 2, 5, 4 ]], Array{Cdouble,2}(undef,0,0)))
    push!(input.facetlist,RawFacet([Cint[2, 3, 6, 5 ]], Array{Cdouble,2}(undef,0,0)))
    push!(input.facetlist,RawFacet([Cint[3, 1, 4, 6 ]], Array{Cdouble,2}(undef,0,0)))
    push!(input.facetlist,RawFacet([Cint[7, 8, 12,11]], Array{Cdouble,2}(undef,0,0)))
    push!(input.facetlist,RawFacet([Cint[8, 9, 13,12]], Array{Cdouble,2}(undef,0,0)))
    push!(input.facetlist,RawFacet([Cint[9, 10,14,13]], Array{Cdouble,2}(undef,0,0)))
    push!(input.facetlist,RawFacet([Cint[10,7, 11,14]], Array{Cdouble,2}(undef,0,0)))

    input.facetmarkerlist=[1,2,3,3,3,4,5,6,7]
    input.regionlist=[ -0.1 -0.1 0.1 1 vol;]'
    tetrahedralize(input, "paqAQ")
end


# # Simple diffusion example
#
# This script provides our simplest example of Oceananigans.jl functionality:
# the diffusion of a one-dimensional Gaussian. This example demonstrates
#
#   * how to load `Oceananigans.jl`;
#   * how to instantiate an `Oceananigans.jl` `Model`;
#   * how to set an initial condition with a function;
#   * how to time-step a model forward, and finally
#   * how to look at results.
#
# ## Using `Oceananigans.jl`
#
# To use `Oceananigans.jl` after it has been installed, we bring
# `Oceananigans.jl` functions and names into our 'namespace' by writing

using Oceananigans

# We also use `PyPlot.jl` for plotting and `Printf` to format plot legends:

using PyPlot, Printf

# ## Instantiating and configuring a `Model`
#
# To begin using Oceananigans, we instantiate a `Model` by calling the
# `Model` constructor:

model = Model(
    grid = RegularCartesianGrid(size = (1, 1, 128), length = (1, 1, 1)),
    closure = ConstantIsotropicDiffusivity(κ = 1.0)
)

# The keyword arguments `grid` and `closure` indicate that
# our model grid is Cartesian with uniform grid spacing, that our diffusive
# stress and tracer fluxes are determined by diffusion with a constant
# diffusivity `κ` (note that we do not use viscosity in this example).

# Note that by default, a `Model` has no-flux boundary condition on all
# variables. Next, we set an initial condition on our "passive tracer",
# temperature. Our objective is to observe the diffusion of a Gaussian.

## Build a Gaussian initial condition function with width `δ`:
δ = 0.1
Tᵢ(x, y, z) = exp( -(z + 0.5)^2 / (2δ^2) )

## Set `model.tracers.T` to the function `Tᵢ`:
set!(model, T=Tᵢ)

# ## Running your first `Model`
#
# Finally, we time-step the model forward using the function
# `time_step!`, with a time-step size that ensures numerical stability.

## Time-scale for diffusion across a grid cell
cell_diffusion_time_scale = model.grid.Δz^2 / model.closure.κ.T

## The function `time_step!` executes `Nt` time steps with step size `Δt`
## using a second-order Adams-Bashforth method
time_step!(model, Nt = 1000, Δt = 0.1 * cell_diffusion_time_scale)

# ## Visualizing the results
#
# We use `PyPlot.jl` to look at the results.

## A convenient function for generating a label with the Current model time
tracer_label(model) = @sprintf("\$ t=%.3f \$", model.clock.time)

## Create a figure with `PyPlot.jl`
close("all")
fig, ax = subplots()
title("Diffusion of a Gaussian")
xlabel("Tracer concentration")
ylabel(L"z")

## Plot initial condition
plot(Tᵢ.(0, 0, model.grid.zC), model.grid.zC, "--", label=L"t=0")

## Plot current solution
plot(interior(model.tracers.T)[1, 1, :], model.grid.zC, label=tracer_label(model))
legend()
gcf()

# Interesting! Running the model even longer makes even more interesting results.

for i = 1:3
    time_step!(model, Nt = 1000, Δt = 0.1 * cell_diffusion_time_scale)
    plot(interior(model.tracers.T)[1, 1, :], model.grid.zC, label=tracer_label(model))
end

legend()
gcf()


#=
    The file contains the implementation of the FGLM algorithm,
    an algorithm for ordering conversion
=#

# does not insert duplicates
function insertsorted!(array, item, rev=false)
    pos = searchsortedfirst(array, item, rev=rev)
    if !isempty(array) && pos <= length(array) && array[pos] == item
        return array
    end
    insert!(array, pos, item)
    return array
end

# TODO: revise this criterion
function insert_nexts!(nextmonomials, monom)
    xs = gens(parent(monom))
    for new_monom in monom .* xs
        insertsorted!(nextmonomials, new_monom)
    end
end

"""
    Represents the given polynomial as an element from Kⁿ
    vector space with respect to monomials occuring in monoms
"""
function decompose_poly(f, monoms)
    v = zeros(base_ring(f), length(monoms))
    for (j, c, m) in zip(1:length(f), coefficients(f), monomials(f))
        ind = findfirst(c -> c == m, monoms)
        v[ind] = c
    end
    v
end

"""
    Represents the given vector from Kⁿ as a polynomial
    in the given monomials
"""
function compose_poly(v, monoms)
    f = zero(parent(first(monoms)))
    for (i, m) in enumerate(monoms)
        f += m * v[i]
    end
    f
end

struct RrefBasis{C}
    rref_vectors::Vector{Vector{C}} # vectors stored in rref form
    monoms::Vector{MPoly{C}}   # monomials corresponding to columns of rref
    new_monoms::Vector{MPoly{C}} # monomials from the new basis
end

RrefBasis{C}() where {C} = RrefBasis(Vector{Vector{C}}(), Vector{MPoly{C}}(), Vector{MPoly{C}}())

function update_basis!(basis::RrefBasis, newpoly)
    ground = base_ring(newpoly)

    n = length(basis.monoms)
    for new_monom in monomials(newpoly)
        (new_monom in basis.monoms) && continue
        push!(basis.monoms, new_monom)
    end

    for vec in basis.rref_vectors
        append!(vec, zeros(ground, length(basis.monoms) - n))
    end
end

function solve_in_rref(rref_vectors, vector)
    coeffs = zeros(parent(first(vector)), length(rref_vectors))
    for (i, basis_vector) in enumerate(rref_vectors)
        pivot = findfirst(!iszero, basis_vector)
        iszero(vector[pivot]) && continue
        coeffs[i] = vector[pivot]
        vector .-= coeffs[i] .* basis_vector
    end
    exists = iszero(vector)
    coeffs, exists, vector
end

function linear_relation!(basis::RrefBasis, poly)
    vector = decompose_poly(poly, basis.monoms)
    coeffs, exists, vector = solve_in_rref(basis.rref_vectors, vector)
    coeffs, exists, vector
end

function update_rref!(basis::RrefBasis, new_monom, new_vector, coeffs)
    push!(basis.new_monoms, new_monom)
    pivot = findfirst(!iszero, new_vector)

    n_coeff = new_vector[pivot]
    new_vector .//= n_coeff

    for (i, c) in enumerate(coeffs)
        basis.new_monoms[end] -= c * basis.new_monoms[i]
    end
    basis.new_monoms[end] *= inv(n_coeff)

    es = basis.rref_vectors
    for i in 1:length(es)
        iszero(es[i][pivot]) && continue
        basis.new_monoms[i] -= es[i][pivot] * basis.new_monoms[end]
        es[i] .-= es[i][pivot] .* new_vector
    end
    push!(basis.rref_vectors, new_vector)

end

function linear_combination(basis::RrefBasis, coeffs)
    sum(c*x  for (c, x) in zip(coeffs, basis.new_monoms))
end

"""
    Converts the basis `G` into the :lex ordering
"""
function fglm(ring, gbexps, gbcoeffs, rng)

    # TODO #
    RR, _ = PolynomialRing(GF(Int(ring.ch)), ["x$i" for i in 1:length(gbcoeffs)], ordering=:degrevlex)
    println(typeof(RR))
    G = export_basis(RR, gbexps, gbcoeffs)
    ###

    origR  = parent(first(G))
    origxs = gens(origR)
    ground = base_ring(origR)
    T      = elem_type(origR)

    @assert ordering(origR) == :degrevlex
    @assert characteristic(ground) > 0

    # @info "converting from $(ordering(origR)) to :lex with fglm"

    newR, _ = PolynomialRing(ground, string.(origxs), ordering=:lex)

    newG      = T[]            #  lex
    staircase = T[]            #  orderless
    nextmonomials = T[newR(1)] # lex

    i = 0
    basis = RrefBasis{elem_type(ground)}()

    while !isempty(nextmonomials)
        i += 1
        if i > 10000
            @error "Something probably went wrong in FGLM"
            break
        end
        if i % 10 == 0
            # @debug "fglm iter $i" size(newG)
        end

        m_lex = popfirst!(nextmonomials)
        m_deg = change_base_ring(ground, m_lex, parent=origR)

        # do we even need this ?
        if any(x -> is_term_divisible(m_lex, x), staircase)
            continue
        end

        # we convert to origR to ensure the output lives there
        nf = normal_form(m_deg, G)

        # pad vectors with zeros and add new monoms
        update_basis!(basis, nf)

        # try to find linear relation between vectors of basis
        # resulting to nf
        λ, exists, reduced = linear_relation!(basis, nf)

        if exists
            poly = m_lex - linear_combination(basis, λ)
            push!(newG, poly)
            push!(staircase, m_lex)
        else
            update_rref!(basis, m_lex, reduced, λ)
            insert_nexts!(nextmonomials, m_lex)
        end
    end

    _, newexps, newcoeffs = convert_to_internal(newG)

    newexps, newcoeffs
end


function gate(::GateName"pauli_channel", dims::Tuple = (2,);
              pauli_ops = ["Id","X","Y","Z"],
              error_probabilities = prepend!(zeros(length(pauli_ops)^length(dims)-1),1))
  @assert sum(error_probabilities) ≈ 1
  @assert all(dims .== 2)
  N = length(dims)
  length(error_probabilities) > (1 << 10) && error("Hilbert space too large")
  error_probabilities ./= sum(error_probabilities)
  kraus = zeros(Complex{Float64},1<<N,1<<N,length(pauli_ops)^N)
  basis = vec(reverse.(Iterators.product(fill(pauli_ops,N)...)|>collect))
  for (k,ops) in enumerate(basis)
    kraus[:,:,k] = √error_probabilities[k] * reduce(kron,gate.(ops))
  end
  return kraus
end

gate(::GateName"bit_flip", dims::Tuple = (2,); p::Number) =
  gate("pauli_channel", dims; error_probabilities = prepend!(p/(2^length(dims)-1) * ones(2^length(dims)-1), 1-p), pauli_ops = ["Id","X"])

gate(::GateName"phase_flip", dims::Tuple = (2,); p::Number) =
  gate("pauli_channel", dims; error_probabilities = prepend!(p/(2^length(dims)-1) * ones(2^length(dims)-1), 1-p), pauli_ops = ["Id","Z"])

function gate(::GateName"AD", dims::Tuple = (2,); γ::Real)
  N = length(dims)
  @assert all(dims .== 2)
  kraus = zeros(2,2,2)
  kraus[:,:,1] = [1 0
                  0 sqrt(1-γ)]
  kraus[:,:,2] = [0 sqrt(γ)
                  0 0]
  N == 1 && return kraus
  k1 = kraus[:,:,1]
  k2 = kraus[:,:,2]
  T = vec(Iterators.product(fill([k1,k2],N)...)|>collect)
  K = zeros(Float64,1<<N,1<<N,1<<N)
  for x in 1:1<<N
    K[:,:,x] = reduce(kron,T[x])
  end
  return K
end
#
# To accept the gate name "amplitude_damping"
gate(::GateName"amplitude_damping", dims::Tuple = (2,); kwargs...) = gate("AD", dims; kwargs...)

function gate(::GateName"PD", dims::Tuple = (2,); γ::Real)
  N = length(dims)
  @assert all(dims .== 2)
  kraus = zeros(2,2,2)
  kraus[:,:,1] = [1 0
                  0 sqrt(1-γ)]
  kraus[:,:,2] = [0 0
                  0 sqrt(γ)]

  N == 1 && return kraus
  k1 = kraus[:,:,1]
  k2 = kraus[:,:,2]
  T = vec(Iterators.product(fill([k1,k2],N)...)|>collect)
  K = zeros(Float64,1<<N,1<<N,1<<N)
  for x in 1:1<<N
    K[:,:,x] = reduce(kron,T[x])
  end
  return K
end

# To accept the gate name "phase_damping"
gate(::GateName"phase_damping", dims::Tuple = (2,); kwargs...) = gate("PD", dims; kwargs...)
#
# To accept the gate name "dephasing"
gate(::GateName"dephasing", dims::Tuple = (2,); kwargs...) = gate("PD", dims; kwargs...)

# make general n-qubit
gate(::GateName"DEP", dims::Tuple = (2,); p::Number) =
  gate("pauli_channel", dims; error_probabilities = prepend!(p/(4^length(dims)-1) * ones(4^length(dims)-1), 1-p), pauli_ops = ["Id","X","Y","Z"])

# To accept the gate name "depolarizing"
gate(::GateName"depolarizing", dims::Tuple = (2,); kwargs...) =
  gate("DEP", dims; kwargs...)

function insertnoise(circuit::Vector{<:Vector{<:Any}}, noisemodel::Tuple; gate = nothing)#idlenoise::Bool = false)
  max_g_size = maxgatesize(circuit)

  # single noise model for all
  if noisemodel[1] isa String
    tmp = []
    for k in 1:max_g_size
      push!(tmp, k => noisemodel)
    end
    #if idlenoise
    #  push!(tmp, "idle" => noisemodel)
    #end
    noisemodel = Tuple(tmp)
  end

  noisycircuit = []
  for layer in circuit
    noisylayer = []
    for g in layer
      push!(noisylayer, g)
      applynoise = (isnothing(gate) ? true :
                    gate isa String ? g[1] == gate : g[1] in gate)
      #@applytogate = gate isa String ? g[1] == gate : g[1] in gate
      #if isnothing(gate) || applytogate
      if applynoise
        nq = g[2]
        # n-qubit gate
        if nq isa Tuple
          gatenoiseindex = findfirst(x -> x == length(nq), first.(noisemodel))
          isnothing(gatenoiseindex) && error("Noise model not defined for $(length(nq))-qubit gates!")
          gatenoise = last(noisemodel[gatenoiseindex])
          # Check whether the n-qubit Kraus channel has been defined
          # n -qubit Kraus operator
          noisecheck = PastaQ.gate(gatenoise[1], Tuple(repeat([2], length(nq))); gatenoise[2]...)
          # if the single-qubit copy is return, use productnoise, and throw a warning
          if size(noisecheck,1) < 1<<length(nq)
            @warn "$(length(nq))-qubit Kraus operators for the $(gatenoise[1]) noise not defined. Applying tensor-product noise instead.\n"
            # tensor-product noise
            for q in nq
              push!(noisylayer,(gatenoise[1], q, gatenoise[2]))
            end
          else
            # correlated n-qubit noise
            push!(noisylayer,(gatenoise[1], nq, gatenoise[2]))
          end
        # 1-qubit gate
        else
          gatenoiseindex = findfirst(x -> x == 1, first.(noisemodel))
          isnothing(gatenoiseindex) && error("Noise model not defined for 1-qubit gates!")
          gatenoise = last(noisemodel[gatenoiseindex])
          push!(noisylayer, (gatenoise[1], nq, gatenoise[2]))
        end
      end
    end
    #if idlenoise
    #  gatenoiseindex = findfirst(x -> x == "idle", first.(noisemodel))
    #  isnothing(gatenoiseindex) && error("Noise model not defined for idling qubits!")
    #  gatenoise = last(noisemodel[gatenoiseindex])
    #  busy_qubits = vcat([collect(g[2]) for g in layer]...)
    #  idle_qubits = filter(y -> y ∉ busy_qubits, 1:nqubits(circuit))
    #  for n in idle_qubits
    #    push!(noisylayer, (gatenoise[1], n, gatenoise[2]))
    #  end
    #end
    push!(noisycircuit, noisylayer)
  end
  return noisycircuit
end

insertnoise(circuit::Vector{<:Any}, noisemodel::Tuple; kwargs...) =
  insertnoise([circuit], noisemodel; kwargs...)[1]

function maxgatesize(circuit::Vector{<:Vector{<:Any}})
  maxsize = 0
  for layer in circuit
    for g in layer
      maxsize = length(g[2]) > maxsize ? length(g[2]) : maxsize
    end
  end
  return maxsize
end

maxgatesize(circuit::Vector{<:Any}) =
  maxgatesize([circuit])


# ---------------------- #
# Piecewise linear Basis #
# ---------------------- #

function Basis(::Lin, breaks::Vector, evennum::Int=0)
    n = length(breaks)  # 28
    !(issorted(breaks)) && error("breaks should be increasing")

    if evennum != 0
        if length(breaks) == 2
            breaks = linspace(breaks[1], breaks[2], evennum)
        else
            if length(breaks) < 2
                error("breaks must have at least 2 elements")
            end

            if any(abs(diff(diff(breaks))) .> 5e-15*mean(abs(breaks)))
                error("Breaks not evenly spaced")
            end
            evennum = length(breaks)
        end
    end
    n = length(breaks)
    a = breaks[1]
    b = breaks[end]
    Basis(Lin(), n, a, b, LinParams(breaks, evennum))
end

# define methods for LinParams type
Basis(p::LinParams) = Basis(Lin(), p.breaks, p.evennum)

nodes(p::LinParams) = p.breaks

function derivative_op(p::LinParams, order::Int=1)
    breaks, evennum = p.breaks, p.evennum

    newbreaks = breaks
    n = length(breaks)
    D = Array(SparseMatrixCSC{Float64, Int}, abs(order))

    for i in 1:order
        d = 1./diff(newbreaks)
        d = sparse([1:n-1; 1:n-1], [1:n-1; 2:n], [-d; d], n-1, n)
        if i > 1
            D[i] = d*D[i-1]
        else
            D[1] = d
        end
        newbreaks = (newbreaks[1:end-1]+newbreaks[2:end])/2
        n = n-1
    end

    for i in -1:-1:order
        newbreaks=[dot([3, -1], newbreaks[1:2]);
                   (newbreaks[1:end-1]+newbreaks[2:end]);
                   dot([-1, 3], newbreaks[end-1:end])]/2
        d = diff(newbreaks)'
        n = n+1
        d = tril(repmat(d, n, 1), -1)
        if i<-1
            D[-i] = d*D[-i-1]
        else
            D[1] = d
        end
        #adjustment to make value at original left endpoint equal 0
        if evennum > 0
            temp = evalbase(LinParams(newbreaks, length(newbreaks)),
                            breaks[1], 0)*D[-i]
        else
            temp = evalbase(LinParams(newbreaks, 0), breaks[1], 0)*D[-i]
        end
        D[-i] = D[-i]-repmat(temp, length(newbreaks), 1)
    end

    params = LinParams(newbreaks, evennum > 0 ? evennum : 0)

    return D, params

end

function _prep_evalbase(p::LinParams, x::Real)
    n = length(p.breaks)
    if p.evennum != 0
        _ind = fix((x-p.breaks[1])*((n-1)/(p.breaks[end]-p.breaks[1])))
        ind = clamp(_ind+1, 1, n-1)
    else
        ind = lookup(p.breaks, x, 3)
    end
    1, n, [ind]
end

function _prep_evalbase(p::LinParams, x::AbstractArray)
    m = size(x, 1)
    n = length(p.breaks)

    # Determine the maximum index of
    #   the breakpoints that are less than or equal to x,
    #   (if x=b use the index of the next to last breakpoint).
    if p.evennum != 0
        ind = fix((x-p.breaks[1]).*((n-1)./(p.breaks[end]-p.breaks[1]))) + 1
        clamp!(ind, 1, n-1)
    else
        ind = lookup(p.breaks, x, 3)
    end

    return m, n, ind
end

function evalbase(::Type{SparseMatrixCSC}, p::LinParams,
                  x::Union{Real,AbstractArray}=nodes(p), order::Int=0)
    # 46-49
    if order != 0
        D, params = derivative_op(p, order)
        B = evalbase(params, x, 0) * D[end]
        return B
    end

    m, n, ind = _prep_evalbase(p, x)
    z = similar(x)
    for i in 1:length(x)
        z[i] = (x[i]-p.breaks[ind[i]])/(p.breaks[ind[i]+1]-p.breaks[ind[i]])
    end

    return sparse(vcat(1:m, 1:m), vcat(ind, ind+1), vcat(1-z, z), m, n)
end

function evalbase(::Type{SplineSparse}, p::LinParams,
                  x::Union{Real,AbstractArray}=nodes(p), order::Int=0)
    # 46-49
    if order != 0
        error("derivatives un-supported right now")
    end

    m, n, ind = _prep_evalbase(p, x)

    z = Array(eltype(x), 2*length(x))
    for i in 1:length(x)
        ix = 2i
        z[ix] = (x[i]-p.breaks[ind[i]])/(p.breaks[ind[i]+1]-p.breaks[ind[i]])
        z[ix-1] = 1 - z[ix]
    end

    return SplineSparse(1, 2, n, z, ind)
end

evalbase(p::LinParams, x::Union{Real,AbstractArray}=nodes(p), order::Int=0) =
    evalbase(SparseMatrixCSC, p, x, order)

function evalbase(p::LinParams, x, order::AbstractArray{Int})
    out = Array(SparseMatrixCSC{Float64,Int}, size(order))

    for I in eachindex(order)
        out[I] = evalbase(p, x, order[I])
    end
    return out
end


import Base.+
Base.:+(s::AbstractString, n::Real) = string((x = tryparse(Int, s)) isa Int ? x + 1 : parse(Float64, s) + 1)
@show "125" + 1
@show "125.15" + 1
@show "1234567890987654321" + 1


function terms(orb::RelativisticOrbital, w::Int=one(Int))
    j = kappa_to_j(orb.κ)
    w <= 2j+1 || throw(ArgumentError("w=$w too large for $orb orbital"))

    2w ≥ 2j+1 && (w = convert(Int, 2j) + 1 - w)
    w == 0 && return [zero(HalfInteger)]
    w == 1 && return [j]

    # Forms full Cartesian product of all mⱼ, not necessarily the most
    # performant method.
    mⱼs = filter(allunique, collect(allchoices([-j:j for i = 1:w])))
    MJs = map(x -> reduce(+, x), mⱼs) # TODO: make map(sum, mⱼs) work

    Js = HalfInteger[]

    while !isempty(MJs)
        # Identify the maximum MJ and associate it with J.
        MJmax = maximum(MJs)
        N = count(isequal(MJmax), MJs)
        append!(Js, repeat([MJmax], N))
        # Remove the -MJ:MJ series, N times.
        for MJ = -MJmax:MJmax
            deleteat!(MJs, findall(isequal(MJ), MJs)[1:N])
        end
    end

    # Do we really want unique here?
    sort(unique(Js))
end

# This is a workaround until seniority number are implemented for
# jj-coupled subshells.
intermediate_terms(orb::RelativisticOrbital, w::Int=one(Int)) =
    terms(orb, w)

export terms, intermediate_terms


struct NNKolmogorov{C,O,S} <: NeuralNetDiffEqAlgorithm
    chain::C
    opt::O
    sdealg::S
end
NNKolmogorov(chain  ; opt=Flux.ADAM(0.1) , sdealg = EM()) = NNKolmogorov(chain , opt , sdealg)

function DiffEqBase.solve(
    prob::KolmogorovPDEProblem,
    alg::NeuralNetDiffEqAlgorithm;
    abstol = 1f-6,
    verbose = false,
    maxiters = 300,
    save_everystep = false,
    dt,
    kwargs...
    )

    tspan = prob.tspan
    phi = prob.phi
    xspan = prob.xspan
    sigma = prob.sigma
    μ = prob.μ
    d = prob.d
    ts = tspan[1]:dt:tspan[2]
    xs = xspan[1]:0.001:xspan[2]
    N = size(ts)
    T = tspan[2]
    #hidden layer
    chain  = alg.chain
    opt    = alg.opt
    sdealg = alg.sdealg
    ps     = Flux.params(chain)
    xi     = rand(xs , d , N[1])
    #Finding Solution to the SDE having initial condition xi. Y = Phi(S(X , T))
    sdeproblem = SDEProblem(μ,sigma,xi,tspan)
    sol = solve(sdeproblem, sdealg ,dt=0.01 , save_everystep=false , kwargs...)
    x_sde = sol[end]
    y = phi(x_sde)
    data   = Iterators.repeated((xi , y), maxiters)

    #MSE Loss Function
    loss(x , y) =Flux.mse(chain(x), y)

    cb = function ()
        l = loss(xi, y)
        verbose && println("Current loss is: $l")
        l < abstol && Flux.stop()
    end

    Flux.train!(loss, ps, data, opt; cb = cb)
    xi , chain(xi)
 end #solve


# # Spectrally accurate Nyström method for the Helmholtz equation in 2D
# ## Carlos Pérez Arancibia (cperez@mat.uc.cl)

using WaveProp, WaveProp.Geometry, WaveProp.Integration, WaveProp.Nystrom, WaveProp.Mesh
using LinearAlgebra
using Plots

pyplot()

#=

## Validation

=#
k   = 20
pde = Helmholtz(dim=2,k=k)
# geo = Kite()
δ = 1e-4
ref = ReferenceLine() # the [0,1] segment
bnd = ParametricEntity(ref) do u
     x = cos(2π*u[1])
     y = (δ+x^2)*sin(2π*u[1])
     SVector(x,y)
end
geo = ParametricBody(boundary=bnd)
G   = SingleLayerKernel(pde)
xₛ  = (0.5,0)
u   = (x) -> G(xₛ,x)
xₒ  = (3,-1)
subfig = plot(geo,label="Γ")
plot!(subfig,[xₛ],m=:o,ms=10,label="source location")
plot!(subfig,[xₒ],m=:o,ms=10,label="test location")
η     = k
Ω,M   = meshgen(geo)
Γ     = boundary(Ω)
niter = 7
er    = zeros(niter)
nn    = 40*collect(1:niter)

for i in 1:niter
     quad_rule = TrapezoidalP(nn[i])
     mesh      = NystromMesh(M,Γ;quad_rule)
     S = SingleLayerOperator(pde,mesh)
     D = DoubleLayerOperator(pde,mesh)
     Smk = Nystrom.assemble_mk(S)
     Dmk = Nystrom.assemble_mk(D)
     L     = I/2 + Dmk - im*η*Smk # create a dense matrix
     rhs = γ₀(u,mesh)
     σ   = L\rhs
     𝒮 = SingleLayerPotential(pde,mesh)
     𝒟 = DoubleLayerPotential(pde,mesh)
     uₙ = (x) -> 𝒟[σ](x) - im*η*𝒮[σ](x)
     xx = (5,5)
     er[i] = abs(u(xx) - uₙ(xx))/abs(u(xx))
end
fig = plot(nn,er,m=:x,yscale=:log10,ylabel="log₁₀(error)",xlabel="n",label="(uₙ - u)/u")

niter = 9
er_dim    = []
er_ldim    = []
nn    = []
quad_rule = GaussLegendre(6)
for i in 1:niter
     Ω,M   = meshgen(geo;gridsize=(1/2)^i)
     Γ     = boundary(Ω)
     mesh      = NystromMesh(M,Γ;quad_rule)
     S = SingleLayerOperator(pde,mesh)
     D = DoubleLayerOperator(pde,mesh)
     Sdim = Nystrom.assemble_dim(S)
     Ddim = Nystrom.assemble_dim(D)
     Sldim = Nystrom.assemble_ldim(S)
     Dldim = Nystrom.assemble_ldim(D)
     Ldim     = I/2 + Ddim - im*η*Sdim # create a dense matrix
     Lldim     = I/2 + Dldim - im*η*Sldim # create a dense matrix
     rhs = γ₀(u,mesh)
     σdim    = Ldim\rhs
     σldim   = Lldim\rhs
     𝒮 = SingleLayerPotential(pde,mesh)
     𝒟 = DoubleLayerPotential(pde,mesh)
     xx = (5,5)
     uₙ = (x) -> 𝒟[σdim](x) - im*η*𝒮[σdim](x)
     push!(er_dim,abs(u(xx) - uₙ(xx))/abs(u(xx)))
     uₙ = (x) -> 𝒟[σldim](x) - im*η*𝒮[σldim](x)
     push!(er_ldim,abs(u(xx) - uₙ(xx))/abs(u(xx)))
     push!(nn,length(mesh))
end
fig = plot(nn,er_dim,m=:x,xscale=:log2,yscale=:log10,ylabel="log₁₀(error)",xlabel="n",label="gdim")
plot!(fig,nn,er_ldim,m=:x,xscale=:log2,yscale=:log10,ylabel="log₁₀(error)",xlabel="n",label="ldim")

plot(fig,subfig,layout=(2,1))

#===============================================
We consider a kite-shaped curve given by:
```math
     \Gamma = \{(\cos(t)+0.65\cos(2t)-0.65, 1.5\sin(t))\in\mathbb R^2,t\in (0,2\pi)\}
```
and we let $\Omega\subset\mathbb{R}^2$ be the (open) domain enclosed by
$\Gamma$, i.e., $\Gamma=\partial\Omega$.

We can construct this **parametric curve** and plot it through the `Geometry` module:
===============================================#

geo = Kite()
plot(geo)

#===============================================
The boundary value problem that we consider corresponds to an incident planewave
$u^{\rm inc}(x_1,x_2) = \exp(ik(x_1\cos\theta+x_2\sin\theta))$ where
$\theta$ is the angle of incidence and $k>0$ is the wavenumber, that impinges on
$\Gamma$ giving rise to a scattered field $u$ which satisfies:
```math
\begin{aligned}
\Delta u+k^2 u &=0\quad\text{in}\quad\mathbb R^2\setminus\overline\Omega,\\
u &=-u^{\rm inc}\quad\text{on}\quad\Gamma,\\
\lim_{|x|\to\infty}|x|^{1/2}\left\{\frac{\partial u}{\partial |x|}-iku\right\} &= 0
\end{aligned}
```

We proceed to define the partial differential equation (**PDE**) that we wish so solve,
as well as the incident wavefield
===============================================#
θ   = 0.
d   = (cos(θ),sin(θ))
uᵢ  = (x) -> exp.(im*k*(x⋅d))
∂uᵢ(x,ν) = im*k*(ν⋅d)*uᵢ(x)

#=
In order to solve the exterior boundary value problem, we resort to the combined field integral equation formulation, whereby we seek the scattered field in the form:

$u(x) = (\mathcal D-i\eta S)[\varphi](x),\quad x\in\mathbb R^2\setminus\Gamma,$

where $\mathcal D$ and $\mathcal S$ are the double- and single-layer potentials defined as:

$(\mathcal D\varphi)(x) = \int_{\Gamma}\frac{\partial G(x,y)}{\partial n(y)}\varphi(y)ds(y)\quad\text{and}\quad (\mathcal S\varphi)(x) = \int_{\Gamma}G(x,y)\varphi(y)ds(y).$

Here $G(x,y) = \frac{i}4 H_0^{(1)}(k|x-y|)$ is the free-space Green's function for Helmholtz equation and $\varphi:\Gamma\to\mathbb C$ is an unknown density function that will be determined by solving an integral equation on curve $\Gamma$.

The constant $\eta>0$ is intruduced so as to ensure uniqueness of solutions of the boundary integral equation. We now define $\eta$ and the tuple _coeff_ that will be used below to construct the combined field potential and operator:

Imposing the Dirichlet boundary condition on $\Gamma$, we arrive at the following second-kind integral equation:

$\left(\frac{I}{2}+K-i\eta S\right)\varphi = -u^{\rm inc}\quad\text{on}\quad\Gamma,$

where operators above are the so-called double- and single-layer integral operators defined as:

$(K\varphi)(x) = \int_{\Gamma}\frac{\partial G(x,y)}{\partial n(y)}\varphi(y)ds(y)\quad\text{and}\quad(S\varphi)(x) = \int_{\Gamma}G(x,y)\varphi(y)ds(y),\qquad x\in\Gamma.$

where we have used the jump relations:

$$\lim_{\epsilon\to 0\pm}(\mathcal D\varphi)(x+\epsilon\nu(x))=\pm\frac{\varphi(x)}{2}+(K\varphi)(x)$$

and

$$\lim_{\epsilon\to 0\pm}(\mathcal S\psi)(x+\epsilon\nu(x))=(S\psi)(x)$$

for $x\in\Gamma.$

We will approximate the integral operators $S$ and $K$ using a Nyström method, for which we must specify (i) a
mesh of $\Gamma$ and (ii) a quadrature rule to be used in each element of the mesh. This is done as follows
=#

n     = 10*k
quad_rule = TrapezoidalP(n)
mesh  = NystromMesh(M,Γ;quad_rule)

# Now that an appropriate mesh is available, we can create the integral operators:
S = SingleLayerOperator(pde,mesh)
D = DoubleLayerOperator(pde,mesh)

#=====
The `S` and `D` variables are `IntegralOperator`s, which are in fact a light wrapper around the
information needed to defined these operators mathematically. They are subtypes of `AbstractMatrix`,
and behave as such. For instance, we can now *solve* for the unknown density $\sigma$ as follows:
=====#

η   = k
L   = I/2 + D - im*η*S # create a dense matrix
rhs = -γ₀(uᵢ,mesh)
σ   = L\rhs

#=

There are two possible issues with the solution method we just employed: **efficiency** and **accuracy**. We will return to those in a bit, but to get to the end of things, let us now visualize the solution obtained:

=#

𝒮 = SingleLayerPotential(pde,mesh)
𝒟 = DoubleLayerPotential(pde,mesh)
uₛ = (x) -> 𝒟[σ](x) - im*η*𝒮[σ](x)
u  = x -> uᵢ(x) + uₛ(x) # the total field
xgrid = ygrid = -5:(2π/k/10):5
U  = [(x,y) ∈ mesh ? NaN+im*NaN : u((x,y)) for y in ygrid, x in xgrid]
fig1 = heatmap(xgrid,ygrid,real.(U),clims=(-2,2))
plot!(fig1,mesh,lw=6,lc=:white)

#=

Because the integral operators we just used are (weakly) singular, special care is needed to compute the
entries of `S` and `D` if high-order convergence is to be obtained. There are several techniques for doing
so, and deciding which to use is highly application dependent. For two-dimensional smooth obstacles admiting a global parametrization, the most accurate method is probably due to Martensen and Kussmaul. To *assemble* the operators `S` and `K` usign this method, we can simply write

=#

Smk = Nystrom.assemble_mk(S)
Dmk = Nystrom.assemble_mk(D)

# The rest follows exactly as before:

L   = I/2 + Dmk - im*η*Smk # create a dense matrix
σ   = L\rhs
poly = copy(qnodes(mesh))
push!(poly,qnodes(mesh)[1])
U    = [ (x,y) in mesh ? NaN+im*NaN : u((x,y)) for y in ygrid, x in xgrid]
fig2 = heatmap(xgrid,ygrid,real.(U),clims=(-2,2))
plot!(fig2,mesh,lw=6,lc=:white)

# The difference between the naive method and the Martensen-Kusmall method is
# evident. In particular the shadow:
plot(fig1,fig2,layout=(1,2),colorbar=:false,ylims=(-5,5))

# Adding more scatterers is simple:
geo   = [Kite(radius=0.1,center=(xc,yc)) for yc in -1:1:1, xc in -1:0.5:1] |> vec
Ω,M   = meshgen(geo)
Γ     = boundary(Ω)
mesh  = NystromMesh(M,Γ;quad_rule)
S     = SingleLayerOperator(pde,mesh)
D     = DoubleLayerOperator(pde,mesh)
Smk   = Nystrom.assemble_mk(S)
Dmk   = Nystrom.assemble_mk(D)
L     = I/2 + Dmk - im*η*Smk # create a dense matrix
rhs   = -γ₀(uᵢ,mesh)
σ     = L\rhs
## now output the solution.
𝒮     = SingleLayerPotential(pde,mesh)
𝒟     = DoubleLayerPotential(pde,mesh)
uₛ = (x) -> 𝒟[σ](x) - im*η*𝒮[σ](x)
u  = x -> uᵢ(x) .+ uₛ(x) # the total field
xgrid = ygrid = -2:(2π/k/10):2
U  = [(x,y) in mesh ? NaN+im*NaN : u((x,y)) for y in ygrid, x in xgrid]
# U  = [u(SVector(x,y)) for y in ygrid, x in xgrid]
heatmap(xgrid,ygrid,real.(U),clims=(-2,2))
plot!(mesh,lw=2,lc=:black)
plot!(size=(2000,2000))

#=

## Adding your own shapes

Not everybody wants to solve scattering by an arrays of kites, so lets create our own shape and solve scattering problem using that.
=#
## define the shape throught its boundary
ref = ReferenceLine() # the [0,1] segment
bnd = ParametricEntity(ref) do u
     r = 1 + 0.25*cos(10π*u[1])
     SVector(r*cos(2π*u[1]),r*sin(2π*u[1]))
end
geo = ParametricBody(boundary=bnd)
## solve the scattering problem
Ω,M   = meshgen(geo)
Γ     = boundary(Ω)
mesh  = NystromMesh(M,Γ;quad_rule=TrapezoidalP(1000))
S     = SingleLayerOperator(pde,mesh)
D     = DoubleLayerOperator(pde,mesh)
Smk   = Nystrom.assemble_mk(S)
Dmk   = Nystrom.assemble_mk(D)
L     = I/2 + Dmk - im*η*Smk # create a dense matrix
rhs   = -γ₀(uᵢ,mesh)
σ     = L\rhs
## now output the solution.
𝒮     = SingleLayerPotential(pde,mesh)
𝒟     = DoubleLayerPotential(pde,mesh)
uₛ = (x) -> 𝒟[σ](x) - im*η*𝒮[σ](x)
u  = x -> uᵢ(x) .+ uₛ(x) # the total field
xgrid = ygrid = -3:(2π/k/10):3
U  = [(x,y) ∈ mesh ? NaN+im*NaN : u(SVector(x,y)) for y in ygrid, x in xgrid]
heatmap(xgrid,ygrid,real.(U),clims=(-2,2))
plot!(mesh,lw=4,lc=:white)
plot!(size=(2000,2000))

# # Composite quadrature rules
k = 40
pde = Helmholtz(k=k,dim=2)
Ω,M   = meshgen(geo;gridsize=0.01)
Γ     = boundary(Ω)
mesh  = NystromMesh(M,Γ;quad_rule=TrapezoidalP(20))
S     = SingleLayerOperator(pde,mesh)
D     = DoubleLayerOperator(pde,mesh)
Smk   = Nystrom.assemble_dim(S)
Dmk   = Nystrom.assemble_dim(D)
L     = I/2 + Dmk - im*k*Smk # create a dense matrix
rhs   = -γ₀(uᵢ,mesh)
σ     = L\rhs
## now output the solution.
𝒮     = SingleLayerPotential(pde,mesh)
𝒟     = DoubleLayerPotential(pde,mesh)
uₛ = (x) -> 𝒟[σ](x) - im*k*𝒮[σ](x)
u  = x -> uᵢ(x) .+ uₛ(x) # the total field
xgrid = ygrid = -3:(2π/k/10):3
U  = [(x,y) ∈ mesh ? NaN+im*NaN : u(SVector(x,y)) for y in ygrid, x in xgrid]
heatmap(xgrid,ygrid,real.(U),clims=(-2,2))
plot!(mesh,lw=1,lc=:white)


module Evolve

import Laplacian

# Return the x values associated with a given choice of grid divisions and grid
# width
function X(L :: Real, N :: Integer, xmin :: Real)
    return (L/N)*[0:N]+xmin
end
X(L :: Real, N :: Integer) = X(L, N, 0)

# Return a time evolution matrix A for a given potential energy vector and box width.
function A(V :: Vector, L :: Real)

    ϵ = L / (length(V) - 1)

    AA  = SymTridiagonal(0.25(1+ϵ^2*V)im+0.5,-0.125im*ones(length(V)-1))
    AA.dv[1] = AA.dv[end] = 1
    AA.ev[1] = AA.ev[end] = zero(typeof(AA[1,1]))
    return AA

end
A(V :: Function, L :: Real, N :: Integer, xmin :: Real) = A(V(X(L,N,xmin)),L)
A(V :: Function, L :: Real, N :: Integer) = A(V, L, N, 0)

# Return a reverse time evolution matrix B for a given potential energy vector
# and box width.
function B(V :: Vector, L :: Real)

    ϵ = L / (length(V) - 1)

    AA  = SymTridiagonal(0.25(1+ϵ^2*V)im-0.5,-0.125im*ones(length(V)-1))
    AA.dv[1] = AA.dv[end] = 1
    AA.ev[1] = AA.ev[end] = zero(typeof(AA[1,1]))
    return AA

end
B(V :: Function, L :: Real, N :: Integer, xmin :: Real) = B(V(X(L,N,xmin)),L)
B(V :: Function, L :: Real, N :: Integer) = B(V, L, N, 0)

# Return a constructor for an update matrix A with the specified potential energy
# function
function A(V :: Function)
    return (L :: Real, N :: Integer, xmin :: Real) -> A(V, L, N, xmin)
end

# Return the Hamiltonian for a given potential energy vector and box width.
function H(V :: Vector, L :: Real)

    ϵ = L / (length(V) - 1)

    AA = SymTridiagonal((ϵ^-2)-V, -0.5(ϵ^-2)ones(length(V)-1))
    AA.dv[1] = AA.dv[end] = 1
    AA.ev[1] = AA.ev[end] = zero(typeof(AA[1,1]))
    return AA

end
H(V :: Function, L :: Real, N :: Integer, xmin :: Real) = H(V(X(L,N,xmin)),L)
H(V :: Function, L :: Real, N :: Integer) = H(V, L, N, 0)

# Method for setting ϕ₀
function Φ(k₀ :: Real, σ :: Real, x₀ :: Real)

    return (L :: Real, N :: Integer, xmin :: Real) -> begin
        xs = (L/N) .* [0:N] + xmin
        swp = map((x) -> exp(im*k₀*x - (x-x₀)^2 / (2σ^2)), xs)
        swp *= (1 / (π*σ^2))^(1/4)          # normalize
        swp[1] = swp[end] = 0               # enforce ∞ square well boundary
        return swp
    end

end
Φ(k₀ :: Real, σ :: Real) = Φ(k₀, σ, 0)          # start in the center

# Function for evolving a given state vector through n additional steps
# (first step not in result). Keep running until done or until test fails.
function evolve(A :: AbstractArray, ϕ :: Vector, n :: Integer, test :: Function)

    ϕs = zeros(Complex, length(ϕ), n)

    # Solve for χ using Thomas algorithm
    χ  = tommy(A, ϕ)
    ϕs[:,1] = χ - ϕ

    for j in 1:(n-1)
        if test(ϕs[:,j])
            ϕs = ϕs[:,1:j]
            break
        end
        χ = tommy(A, ϕs[:,j])
        ϕs[:,j+1] = χ - ϕs[:,j]
    end

    return ϕs

end
evolve(A :: AbstractArray, ϕ :: Vector, n :: Integer) = evolve(A, ϕ, n, x->true)

# Method for evolving a state vector by passing functions
evolve(A :: Function, ϕ :: Function, n :: Integer, L :: Real, N :: Integer, xmin :: Real) = evolve(A(L,N,xmin), ϕ(L,N,xmin), n)

# Solve Ax = b where A is tridiagonal using the Thomas algorithm
# based on MATLAB code: http://www.mathworks.com/matlabcentral/fileexchange/1359-fast-tridiagonal-system-solver/content/thomas.m
function tommy(a :: Vector, b :: Vector, c :: Vector, d :: Vector)

    n = length(a)
    m = zeros(typeof(a[1]), n)
    l = zeros(typeof(a[1]), n-1)
    y = zeros(typeof(a[1]), n)
    x = zeros(typeof(a[1]), n)

    m[1] = a[1]
    y[1] = d[1]

    # Perform LU factorization and forward substitution to solve Ly = d for y
    for i in 2:n
        l[i-1] = c[i-1] / m[i-1]
        m[i] = a[i] - l[i-1] * b[i-1]

        y[i] = d[i] - l[i-1] * y[i-1]
    end

    x[n] = y[n] / m[n]

    # Begin backward substitution to solve Ux = y for x
    for i in (n-1):-1:1
        x[i] = (y[i] - b[i] * x[i+1]) / m[i]
    end

    return x

end

tommy(A :: Matrix, d :: Vector) = tommy(diag(A), diag(A, -1), diag(A, +1), d)
tommy(A :: Tridiagonal, d :: Vector) = tommy(A.d, A.du, A.dl, d)
tommy(A :: SymTridiagonal, d :: Vector) = tommy(A.dv, A.ev, A.ev, d)

# Create a pocket function centered at zero
function pocket(L :: Real, h :: Real , b :: Real, s :: Real, w :: Real, xmin :: Real, x :: Real)
    if x < (xmin + 0.5L - 0.5w - s) || x > (xmin + 0.5L + 0.5w +s)
        V = 0
    elseif x < (xmin + 0.5L - 0.5w) || x > (xmin + 0.5L + 0.5w)
        V = h
    else
        V = b
    end
    return V
end

end


import Base.Broadcast: broadcasted

"""
    .+(x1::Var, x2::Var)
"""
function broadcasted(::typeof(+), x1::Var, x2::Var)
    ydata = x1.data .+ x2.data
    Var(ydata, ∇broadcasted!, (+,x1,x2))
end
broadcasted(::typeof(+), x1::Node, x2) = Node(broadcasted, (+,x1,x2))
broadcasted(::typeof(+), x1, x2::Node) = Node(broadcasted, (+,x1,x2))

function ∇broadcasted!(y::Var, ::typeof(+), x1::Var, x2::Var)
    isnothing(x1.grad) || ∇broadcast_plus!(y.grad, x1.grad)
    isnothing(x2.grad) || ∇broadcast_plus!(y.grad, x2.grad)
end

function ∇broadcast_plus!(gy::UniArray{T}, gx::UniArray{T}) where T
    if length(gy) == length(gx)
        addto!(gx, gy)
    else
        dims = ()
        for i = 1:ndims(gy)
            if size(gx,i) == 1 && size(gy,i) > 1
                dims = tuple(dims..., i)
            end
        end
        addto!(gx, sum(gy,dims=dims))
    end
end

"""
    .-(x1::Var, x2::Var)
"""
function broadcasted(::typeof(-), x1::Var, x2::Var)
    ydata = x1.data .- x2.data
    Var(ydata, ∇broadcasted!, (-,x1,x2))
end
broadcasted(::typeof(-), x1::Node, x2) = Node(broadcasted, (-,x1,x2))
broadcasted(::typeof(-), x1, x2::Node) = Node(broadcasted, (-,x1,x2))

function ∇broadcasted!(y::Var, ::typeof(-), x1::Var, x2::Var)
    isnothing(x1.grad) || ∇broadcast_plus!(y.grad, x1.grad)
    isnothing(x2.grad) || ∇broadcast_minus!(y.grad, x2.grad)
end

function ∇broadcast_minus!(gy::UniArray{T}, gx::UniArray{T}) where T
    if length(gy) == length(gx)
        axpy!(T(-1), gy, gx)
    else
        dims = ()
        for i = 1:ndims(gy)
            if size(gx,i) == 1 && size(gy,i) > 1
                dims = tuple(dims..., i)
            end
        end
        axpy!(T(-1), sum(gy,dims=dims), gx)
    end
end

"""
    .*(x1::Var, x2::Var)
"""
function broadcasted(::typeof(*), x1::Var, x2::Var)
    ydata = x1.data .* x2.data
    Var(ydata, ∇broadcasted!, (*,x1,x2))
end
broadcasted(::typeof(*), x1::Node, x2) = Node(broadcasted, (*,x1,x2))
broadcasted(::typeof(*), x1, x2::Node) = Node(broadcasted, (*,x1,x2))

function ∇broadcasted!(y::Var, ::typeof(*), x1::Var, x2::Var)
    isnothing(x1.grad) || ∇dottimes!(y.grad, x2.data, x1.grad)
    isnothing(x2.grad) || ∇dottimes!(y.grad, x1.data, x2.grad)
end

function ∇dottimes!(gy::UniArray{T}, x2::UniArray{T}, gx1::UniArray{T}) where T
    g = gy .* x2
    if length(gx1) == length(g)
        addto!(gx1, g)
    else
        dims = ()
        for i = 1:ndims(gy)
            if size(gx1,i) == 1 && size(g,i) > 1
                dims = tuple(dims..., i)
            end
        end
        addto!(gx1, sum(g,dims=dims))
    end
end


@doc raw"""
Axisymmetric tokamak equilibrium in (r,θ,ϕ) coordinates with covariant
components of the vector potential given by
```math
A (r, \theta, \phi) = B_0 \, \bigg( 0 , \, \frac{r R_0}{\cos (\theta)} - \bigg( \frac{R_0}{\cos (\theta)} \bigg)^2 \, \ln \bigg( \frac{R}{R_0} \bigg) , \, - \frac{r^2}{2 q_0} \bigg)^T ,
```
resulting in the magnetic field with covariant components
```math
B (r, \theta, \phi) = \frac{B_0}{q_0} \, \bigg( 0 , \, \frac{r^2}{R}, \, q_0 R_0 \bigg)^T ,
```
where $R = R_0 + r \cos \theta$.

Parameters:
 * `R₀`: position of magnetic axis
 * `B₀`: B-field at magnetic axis
 * `q₀`: safety factor at magnetic axis
"""
struct AxisymmetricTokamakToroidal{T <: Number} <: AnalyticEquilibrium
    name::String
    R₀::T
    B₀::T
    q₀::T

    function AxisymmetricTokamakToroidal{T}(R₀::T, B₀::T, q₀::T) where T <: Number
        new("AxisymmetricTokamakEquilibriumToroidal", R₀, B₀, q₀)
    end
end

AxisymmetricTokamakToroidal(R₀::T=1.0, B₀::T=1.0, q₀::T=2.0) where T <: Number = AxisymmetricTokamakToroidal{T}(R₀, B₀, q₀)

function Base.show(io::IO, equ::AxisymmetricTokamakToroidal)
    print(io, "Axisymmetric Tokamak Equilibrium in Toroidal Coordinates with\n")
    print(io, "  R₀ = ", equ.R₀, "\n")
    print(io, "  B₀ = ", equ.B₀, "\n")
    print(io, "  q₀ = ", equ.q₀)
end

function X(x::AbstractArray{T,1}, equ::AxisymmetricTokamakToroidal) where {T <: Number}
    R(x,equ) * cos(ϕ(x,equ))
end

function Y(x::AbstractArray{T,1}, equ::AxisymmetricTokamakToroidal) where {T <: Number}
    R(x,equ) * sin(ϕ(x,equ))
end

function Z(x::AbstractArray{T,1}, equ::AxisymmetricTokamakToroidal) where {T <: Number}
    r(x,equ) * sin(θ(x,equ))
end

function R(x::AbstractArray{T,1}, equ::AxisymmetricTokamakToroidal) where {T <: Number}
    equ.R₀ + r(x,equ) * cos(θ(x,equ))
end

function R²(x::AbstractArray{T,1}, equ::AxisymmetricTokamakToroidal) where {T <: Number}
    R(x,equ)^2
end

function r(x::AbstractArray{T,1}, equ::AxisymmetricTokamakToroidal) where {T <: Number}
    x[1]
end

function r²(x::AbstractArray{T,1}, equ::AxisymmetricTokamakToroidal) where {T <: Number}
    r(x,equ)^2
end

function θ(x::AbstractArray{T,1}, equ::AxisymmetricTokamakToroidal) where {T <: Number}
    x[2]
end

function ϕ(x::AbstractArray{T,1}, equ::AxisymmetricTokamakToroidal) where {T <: Number}
    x[3]
end

function J(x::AbstractArray{T,1}, equ::AxisymmetricTokamakToroidal) where {T <: Number}
    r(x,equ) * R(x,equ)
end

function periodicity(x::AbstractArray{T,1}, equ::AxisymmetricTokamakToroidal) where {T <: Number}
    p = zero(x)
    p[2] = 2π
    p[3] = 2π
    return p
end

function A₁(x::AbstractArray{T,1}, equ::AxisymmetricTokamakToroidal) where {T <: Number}
    zero(T)
end

function A₂(x::AbstractArray{T,1}, equ::AxisymmetricTokamakToroidal) where {T <: Number}
    equ.B₀ * equ.R₀ / cos(θ(x,equ))^2 * ( r(x,equ) * cos(θ(x,equ)) - equ.R₀ * log(R(x,equ) / equ.R₀) )
end

function A₃(x::AbstractArray{T,1}, equ::AxisymmetricTokamakToroidal) where {T <: Number}
    - equ.B₀ * r²(x,equ) / equ.q₀ / 2
end

function g₁₁(x::AbstractArray{T,1}, equ::AxisymmetricTokamakToroidal) where {T <: Number}
    one(T)
end

function g₂₂(x::AbstractArray{T,1}, equ::AxisymmetricTokamakToroidal) where {T <: Number}
    r²(x, equ)
end

function g₃₃(x::AbstractArray{T,1}, equ::AxisymmetricTokamakToroidal) where {T <: Number}
    R²(x, equ)
end

macro axisymmetric_tokamak_equilibrium_toroidal(R₀, B₀, q₀)
    generate_equilibrium_code(AxisymmetricTokamakToroidal(R₀, B₀, q₀); output=false)
end


using Test
using ForwardDiff: Dual, gradient, partials

using DiffEqBase: ODE_DEFAULT_NORM
const internalnorm = ODE_DEFAULT_NORM

val = rand(10)
par = rand(10)
u = Dual.(val, par)
reference(val, par) = sqrt((sum(abs2,val) + sum(abs2,par)) / (length(val) + length(par)))
dual_real = internalnorm(u, 1)
dual_dual = internalnorm(u, u[1])
@test reference(val, par) ≈ dual_real
@test reference(val, par) ≈ dual_dual
@test partials(dual_dual, 1) ≈ gradient(x->internalnorm(x, x[1]), val)'par


import ContinuumArrays: apply, MulQuasiMatrix

@testset "$(rpad("Bernstein Basis Tests",80))" begin

    l = Bernstein(2)
    d = Derivative(axes(l,1))

    @test apply(*,d,l) isa BernsteinDerivative
    @test d*l isa BernsteinDerivative
    @test l' isa BernsteinDerivative

    @test basis(l) == l.b
    @test nbasis(l) == 2
    @test eachindex(l) == 0:1
    @test order(l) == 2
    @test degree(l) == 1

    l1 = Bernstein(Float64, 2)
    l2 = Bernstein(Integer, 2)
    l3 = Bernstein(Float64, 3)

    @test hash(l) == hash(l1)
    @test hash(l) == hash(l2)
    @test hash(l) != hash(l3)

    @test l == l1
    @test l == l2
    @test l != l3

    @test  isequal(l, l1)
    @test !isequal(l, l2)
    @test !isequal(l, l3)

    y = rand(5)

    z1 = parent([ l[y[i], j] for i in eachindex(y), j in eachindex(l)])
    z2 =   hcat([ l[y[i], :] for i in eachindex(y)]...)'
    z3 =   hcat([ l[y,    j] for j in eachindex(l)]...)
    z4 = parent(  l[y,    :] )

    @test z1 == z2 == z3 == z4

    z1 = parent([ (d*l)[y[i], j] for i in eachindex(y), j in eachindex(l)])
    z2 =   hcat([ (d*l)[y[i], :] for i in eachindex(y)]...)'
    z3 =   hcat([ (d*l)[y,    j] for j in eachindex(l)]...)
    z4 = parent(  (d*l)[y,    :] )

    @test z1 == z2 == z3 == z4

    @test l(0.0, 0) == 1.0
    @test l(0.5, 0) == 0.5
    @test l(1.0, 0) == 0.0

    @test l(0.0, 1) == 0.0
    @test l(0.5, 1) == 0.5
    @test l(1.0, 1) == 1.0

    @test l[0.0, 0] == 1.0
    @test l[0.5, 0] == 0.5
    @test l[1.0, 0] == 0.0

    @test l[0.0, 1] == 0.0
    @test l[0.5, 1] == 0.5
    @test l[1.0, 1] == 1.0

    @test l[0, 0] == +1.0
    @test l[1, 0] ==  0.0
    @test l[2, 0] == -1.0

    @test l[0, 1] == 0.0
    @test l[1, 1] == 1.0
    @test l[2, 1] == 2.0

    @test (d*l)[0.0, 0] == -1.0
    @test (d*l)[0.5, 0] == -1.0
    @test (d*l)[1.0, 0] == -1.0

    @test (d*l)[0.0, 1] == +1.0
    @test (d*l)[0.5, 1] == +1.0
    @test (d*l)[1.0, 1] == +1.0

    @test (d*l)[0, 0] == -1.0
    @test (d*l)[1, 0] == -1.0
    @test (d*l)[2, 0] == -1.0

    @test (d*l)[0, 1] == +1.0
    @test (d*l)[1, 1] == +1.0
    @test (d*l)[2, 1] == +1.0

end


using Test, OrdinaryDiffEq

using SparseArrays, LinearSolve
using LinearAlgebra, Random
N = 30
AA = sprand(MersenneTwister(12),  N, N, 0.5)
mm = sprand(MersenneTwister(123), N, N, 0.5)
A = DiffEqArrayOperator(AA)
M = DiffEqArrayOperator(mm'mm)
u0 = ones(N)
prob = ODEProblem(ODEFunction(A; mass_matrix=M), u0, (0.0, 1.0))

for alg in [Rosenbrock23(), Rosenbrock23(linsolve=KLUFactorization())]
  sol = solve(prob, alg)
  @test sol.destats.njacs == 0
  @test sol.destats.nw == 1
end

## OOP

foop(u, p, t) = jac(u, p, t) * u
jac(u, p, t) = spdiagm(0=>p)
paramjac(u, p, t) = SparseArrays.spdiagm(0=>u)

n = 2
p = collect(1.0:n)
u0 = ones(n)
tspan = [0.0, 1]
odef = ODEFunction(foop; jac=jac, jac_prototype=jac(u0, p, 0.0), paramjac=paramjac)

function g_helper(p; alg=Rosenbrock23(linsolve=LUFactorization()))
    prob = ODEProblem(odef, u0, tspan, p)
    soln = Array(solve(prob, alg; u0=prob.u0, p=prob.p, abstol=1e-4, reltol=1e-4))[:, end]
    return soln
end
function g(p; kwargs...)
    soln = g_helper(p; kwargs...)
    return sum(soln)
end

@test isapprox(exp.(p), g_helper(p; alg=Rosenbrock23(linsolve=KLUFactorization())); atol=1e-3, rtol=1e-3)
@test isapprox(exp.(p), g_helper(p; alg=Rosenbrock23(linsolve=UMFPACKFactorization())); atol=1e-1, rtol=1e-1)
@test isapprox(exp.(p), g_helper(p; alg=Rosenbrock23(linsolve=KrylovJL_GMRES())); atol=1e-1, rtol=1e-1)

@test isapprox(exp.(p), g_helper(p; alg=Rodas4(linsolve=KLUFactorization())); atol=1e-3, rtol=1e-3)
@test isapprox(exp.(p), g_helper(p; alg=Rodas4(linsolve=UMFPACKFactorization())); atol=1e-1, rtol=1e-1)
@test isapprox(exp.(p), g_helper(p; alg=Rodas4(linsolve=KrylovJL_GMRES())); atol=1e-1, rtol=1e-1)

@test isapprox(exp.(p), g_helper(p; alg=Rodas5(linsolve=KLUFactorization())); atol=1e-3, rtol=1e-3)
@test isapprox(exp.(p), g_helper(p; alg=Rodas5(linsolve=UMFPACKFactorization())); atol=1e-1, rtol=1e-1)
@test isapprox(exp.(p), g_helper(p; alg=Rodas5(linsolve=KrylovJL_GMRES())); atol=1e-1, rtol=1e-1)

@test isapprox(exp.(p), g_helper(p; alg=ImplicitEuler(linsolve=KLUFactorization())); atol=1e-1, rtol=1e-1)
@test isapprox(exp.(p), g_helper(p; alg=ImplicitEuler(linsolve=UMFPACKFactorization())); atol=1e-1, rtol=1e-1)
@test isapprox(exp.(p), g_helper(p; alg=ImplicitEuler(linsolve=KrylovJL_GMRES())); atol=1e-1, rtol=1e-1)

@test isapprox(exp.(p), g_helper(p; alg=TRBDF2(linsolve=KLUFactorization())); atol=1e-1, rtol=1e-1)
@test isapprox(exp.(p), g_helper(p; alg=TRBDF2(linsolve=UMFPACKFactorization())); atol=1e-1, rtol=1e-1)
@test isapprox(exp.(p), g_helper(p; alg=TRBDF2(linsolve=KrylovJL_GMRES())); atol=1e-1, rtol=1e-1)

@test isapprox(exp.(p), g_helper(p; alg=KenCarp47(linsolve=KLUFactorization())); atol=1e-1, rtol=1e-1)
@test isapprox(exp.(p), g_helper(p; alg=KenCarp47(linsolve=UMFPACKFactorization())); atol=1e-1, rtol=1e-1)
@test isapprox(exp.(p), g_helper(p; alg=KenCarp47(linsolve=KrylovJL_GMRES())); atol=1e-1, rtol=1e-1)


using DelimitedFiles
using StatsBase

function answer1()
    input = "input"
    initial_fish = parse.(Int,split(readline(input),","))
    fishes_now = copy(initial_fish)
    new_to_add = 0
    for day in 1:80
        # @show fishes_now
        new_to_add = 0
        for i in  eachindex(fishes_now)
            if fishes_now[i] == 0
                new_to_add +=1
                fishes_now[i] = 6
            else
                fishes_now[i] -=1
            end
        end
        append!(fishes_now, repeat([8],new_to_add))
    end
    fishes_now
end

@show answer1()

"Do not try to bruteforce part 2..."
function answer2()
    input = "input"
    initial_fish = parse.(Int,split(readline(input),","))
    fishes_now = zeros(Int,9)
    days = 256
    for fish in initial_fish
        fishes_now[fish+1] += 1
    end
    for day in 1:1:days
        fishes_to_add = fishes_now[1]
        for i in 2:1:9
            fishes_now[i-1] = copy(fishes_now[i])
        end
        fishes_now[7] += fishes_to_add
        fishes_now[9] = fishes_to_add
    end
    sum(fishes_now)
end

@show answer2()


# A simple 2D transmission full waveform inversion experiment
# to showcase how this framework operates. Much of the
#

using JOLI   # Linear operators
using PyPlot # For plotting
using WAVEFORM

# Turn off annoying JOLI type checks
JOLI.jo_type_mismatch_error_set(false)
JOLI.jo_type_mismatch_warn_set(false)

# number of model points in each direction (z,x)
n = 101 .* [1;1];

# grid spacing in metres
d = 10.0 .*[1;1];

# location of the coordinate axis origin
o = 0.0 .* [1;1];

# Frequencies
freqs = [5.0;10.0;15.0];

# Source wavelet time shift
t0 = 0.0;

# Peak frequency of Ricker Wavelet
f0 = 10.0;

# Units of velocity parameter vector
unit = "m/s";

# Background velocity (m/s)
vel_background = 2000;

# Model domain range
L = n.*d;

# Source grid definition
xsrc = 0.0:100.0:L[2];
ysrc = [0.0];
zsrc = [10.0];

# Receiver grid definition
xrec = 0:10.0:L[2];
yrec = [0.0];
zrec = [950.0];

# Model type containing the domain geometry
model = Model{Int64,Float64}(n,d,o,t0,f0,unit,freqs,xsrc,ysrc,zsrc,xrec,yrec,zrec);
nsrc = length(model.xsrc); nfreq = length(model.freq);

comp_n = n;
comp_d = d;
comp_o = o;
npml = 20*ones(Int64,2,2);

# Stencil used to discretize the PDE
pde_scheme = helm2d_chen9p;

# If true,
cut_pml = true;

implicit_matrix = false;

# Objective function
misfit = least_squares;

# Binary mask to indicate which sources/frequencies to use (default: all of them)
srcfreqmask = trues(nsrc,nfreq);

# Use a direct solver for 2D problems
lsopts = LinSolveOpts(solver=:lufact);

# Options for discretizing + solving the Helmholtz PDEs
opts = PDEopts{Int64,Float64}(pde_scheme,comp_n,comp_d,comp_o,cut_pml,implicit_matrix,npml,misfit,srcfreqmask,lsopts);

v0 = vel_background*ones(n...);
v = copy(v0);
v[div(n[1],3):2*div(n[1],3),div(n[2],3):2*div(n[2],3)] = 1.25.*vel_background;
v = vec(v);
v0 = vec(v0);

# True velocity
imshow(v);

Q = diagm(ones(nsrc));
D = forw_model(v,Q,model,opts);

# Plot a frequency slice
imshow(real(D[:,1:nsrc]),aspect="auto");

obj! = construct_pde_misfit(v,Q,D,model,opts,batch_mode=false)

using OptimPackNextGen
proj! = (xproj,x)->project_bounds!(x,minimum(v),maximum(v),xproj);
vest = spg(obj!,proj!,v0,3,maxfc=50,verb=true)
imshow(reshape(vest,model.n...),vmin=minimum(v),vmax=maximum(v))


# Discrete.
hist = [0=>2, 1=>3, 2=>5]
erp = Discrete(hist)
@test all([x in support(erp) for x in 0:2])
@test length(support(erp)) == 3
for x in 0:2
    @test_approx_eq(log(hist[x]/10),score(erp,x))
end

@test all([rand(erp) in 0:2 for _ in 1:5])

# Hellinger distance.
p = Categorical([0.25,0.25,0.5])
q = Discrete([1=>2, 3=>3])
@test_approx_eq_eps 0.36885 hellingerdistance(p,q) 1e-5
q = Discrete([4=>1]) # Different support to p.
@test_throws ErrorException hellingerdistance(p,q)