function binarySearch(A, t)

  low = 1
  high = length(A)
  while low <= high
    mid = Int(round((low + high) / 2))
    if t == A[mid]
      return mid
    elseif t < A[mid]
      high = mid - 1
    elseif t > A[mid]
      low = mid + 1
    end
  end
  return false

end

function main()

  # list to be sorted
  rng = MersenneTwister(1234);
  list_size = 1e7
  A = Vector(1:list_size) # The list is to be sorted for binary search

  # Search
  search_num = 3000
  index_found = @time binarySearch(A, search_num)

  print("Search for : ")
  println(search_num)

  print("Found index: ")
  println(index_found)

  print("Found value is ... ")
  println(A[index_found])

end

if contains(@__FILE__, PROGRAM_FILE)
    main()
end


# Load Julia packages (libraries)

using Pkg, DrWatson

#@quickactivate "StatisticalRethinkingStan"
using StanSample
using StatisticalRethinking

# Define the Stan language model

stan8_2s = "
data{
    int N;
    vector[N] y;
}
parameters{
    real mu;
    real<lower=0> sigma;
}
model{
    y ~ normal( mu , sigma );
}
";

# Define the Stanmodel.

m8_2s = SampleModel("m8.2s", stan8_2s);

# Input data for cmdstan

m8_2_data = Dict("N" => 2, "y" => [-1, 1]);
m8_2_init = Dict("mu" => 0.0, "sigma" => 1.0);

# Sample using cmdstan

rc8_2s = stan_sample(m8_2s; data=m8_2_data, init=m8_2_init);

# Describe the draws
if success(rc8_2s)
  part8_2s = read_samples(m8_2s, :particles)
  part8_2s |> display
end

# End of m8_2s


@testset "Subgradient Plan" begin
    M = Euclidean(2)
    x = RnPoint([4., 2.])
    x0 = RnPoint([5.,2.])
    o = SubGradientMethodOptions(x0,stopAfterIteration(200), DecreasingStepsize(0.1))
    o.∂ = RnTVector([1., 0.])
    f = y -> distance(M,y,x)
    ∂f = y -> distance(M,x,y) == 0 ? zeroTVector(M,y) : -2*log(M,y,x)/distance(M,x,y)
    p = SubGradientProblem(M,f, ∂f)
    oR = solve(p,o)
    xHat = getSolverResult(p,oR)
    @test getInitialStepsize(p,o) == 0.1
    @test getStepsize!(p,o,1) == 0.1
    @test getLastStepsize(p,o,1) == 0.1
    # Check Fallbacks of Problen
    @test getCost(p,x) == 0.
    @test norm(M,x,getSubGradient(p,x)) == 0
    @test_throws ErrorException getGradient(p,o.x)
    @test_throws ErrorException getProximalMap(p,1.,o.x,1)
end


# ------------------------------------------------------------------------------------------
# # Learning with a single neuron using Flux.jl
#
# In this notebook, we'll use `Flux` to create a single neuron and teach it to learn, as we
# did by hand in notebook 10!
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ### Read in data and process it
#
# Let's start by reading in our data
# ------------------------------------------------------------------------------------------

using CSV
using TextParse
using DataFrames

applecols, applecolnames = TextParse.csvread("data/Apple_Golden_1.dat", '\t')
bananacols, bananacolnames = TextParse.csvread("data/bananas.dat", '\t')

apples = DataFrame(Dict(strip(name)=>col for (name, col) in zip(applecolnames, applecols)))
bananas = DataFrame(Dict(strip(name)=>col for (name, col) in zip(bananacolnames, bananacols)));

# ------------------------------------------------------------------------------------------
# and processing it to extract information about the red and green coloring in our images:
# ------------------------------------------------------------------------------------------

col1 = :red
col2 = :green

x_apples  = [ [apples[i, col1], apples[i, col2]] for i in 1:size(apples)[1] ]
x_bananas = [ [bananas[i, col1], bananas[i, col2]] for i in 1:size(bananas)[1] ]

xs = vcat(x_apples, x_bananas)

ys = vcat( zeros(size(x_apples)[1]), ones(size(x_bananas)[1]) );

# ------------------------------------------------------------------------------------------
# The input data is in `xs` and the labels (true classifications as bananas or apples) in
# `ys`.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ### Using `Flux.jl`
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# Now we can load `Flux` to really get going!
# ------------------------------------------------------------------------------------------

using Flux

# ------------------------------------------------------------------------------------------
# We saw in the last notebook that σ is a built-in function in `Flux`.
#
# Another function that is used a lot in neural networks is called `ReLU`; in Julia, the
# function is called `relu`.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 1
#
# Use the docs to discover what `ReLU` is all about.
#
# `relu.([-3, 3])` returns
#
# A) [-3, 3]

# B) [0, 3]

# C) [0, 0]

# D) [3, 3]

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

# ------------------------------------------------------------------------------------------
# ### Making a single neuron in Flux
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# Let's use `Flux` to build our neuron with 2 inputs and 1 output:
#
#
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# We previously put the two weights in a vector, $\mathbf{w}$. Flux instead puts weights in
# a $1 \times 2$ matrix (i.e. a matrix with 1 *row* and 2 *columns*).
#
# Previously, to compute the dot product of $\mathbf{w}$ and $\mathbf{x}$ we had to use
# either the `dot` function, or we had to transpose the vector $\mathbf{w}$:
#
# ```julia
# # transpose w
# b = w' * x
# # or use dot!
# b = dot(w, x)
# ```
# If the weights are instead stored in a $1 \times 2$ matrix, `W`, then we can simply
# multiply `W` and `x` together instead!
#
# We start off with random values for our parameters and data:
# ------------------------------------------------------------------------------------------

W = rand(1, 2)

x = rand(2)

# ------------------------------------------------------------------------------------------
# Note that the product of `W` and `x` will now be an array (vector) with a single element,
# rather than a single number:
# ------------------------------------------------------------------------------------------

W * x

# ------------------------------------------------------------------------------------------
# This means that our bias `b` is treated as an array when we're using `Flux`:
# ------------------------------------------------------------------------------------------

b = rand(1)

# ------------------------------------------------------------------------------------------
# #### Exercise 2
#
# Write a function `mypredict` that will take a single input, array `x` and use `W`, `b`,
# and built-in `σ` to generate an output prediction (stored as an array). This function
# defines our neural network!
#
# Hint: This function will look very similar to $f_{\mathbf{w},\mathbf{b}}$ from the last
# notebook but has changed since our data structures to store our parameters have changed!
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 3
#
# Define a loss function called `loss`.
#
# `loss` should take two inputs: a vector storing data, `x`, and a vector storing the
# correct "labels" for that data. `loss` should return the sum of the squares of differences
# between the predictions and the correct labels.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## Calculating gradients using Flux: backpropagation
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# For learning, we know that what we need is a way to calculate derivatives of the `loss`
# function with respect to the parameters `W` and `b`. So far, we have been doing that using
# finite differences.
#
# `Flux.jl` instead implements a numerical method called **backpropagation** that calculates
# gradients (essentially) exactly, in an automatic way, by indirectly applying the rules of
# calculus.
# To do so, it provides a new type of object called "tracked" arrays. These are arrays that
# store not only their current value, but also information about gradients, which is used by
# the backpropagation method.
#
# [If you want to understand the maths behind backpropagation, we recommend e.g. [this
# lecture](https://www.youtube.com/watch?v=i94OvYb6noo).]
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# To do so, `Flux` provides a function `param` to define such objects that will contain the
# information for a *param*eter.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# Let's start, as usual, by setting up some random initial values for the parameters:
# ------------------------------------------------------------------------------------------

W_data = rand(1, 2)
b_data = rand(1)

W_data, b_data

# ------------------------------------------------------------------------------------------
# We now set up `Flux.jl` objects that will contain these values *and* their derivatives,
# and allow to propagate
# this information around:
# ------------------------------------------------------------------------------------------

W = param(W_data)
b = param(b_data)

# ------------------------------------------------------------------------------------------
# Here, `param` is a function that `Flux` provides to create an object that represents a
# parameter of a machine learning model, i.e. an object which has both a value and
# derivative information, and such that other objects know how to *keep track* of when it is
# used in an expression.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 4
#
# What type does `W` have?
#
# A) Array (1D)

# B) Array (2D)

# C) TrackedArray (1D)

# D) TrackedArray (2D)

# E) Parameter (1D)

# F) Parameter (2D)

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

# ------------------------------------------------------------------------------------------
# #### Exercise 5
#
# `W` stores not only its current value, but also has space to store gradient information.
# You can access the values and gradient of the weights as follows:
#
# ```julia
# W.data
# W.grad
# ```
#
# At this point, are the values of the weights or the gradient of the weights larger?
#
# A) the values of the weights

# B) the gradient of the weights
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 6
#
# For data `x` and `y` where
#
# ```julia
# x, y = [0.413759, 0.692204], [0.845677]
# ```
# apply the loss function to `x` and `y` to give a new variable `l`. What is the type of
# `l`? (How many dimensions does it have?)
#
# A) Array (0D)

# B) Array (1D)

# C) TrackedArray (0D)

# D) TrackedArray (1D)

# E) Float64

# F) Int64

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

# ------------------------------------------------------------------------------------------
# ### Stochastic gradient descent
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# We can now use these features to reimplement stochastic gradient descent, following the
# method we used in the previous notebook, but now using backpropagation!
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 7
#
# Modify the code from the previous notebook for stochastic gradient descent to use Flux
# instead.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ### Investigating stochastic gradient descent
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# Let's look at the values stored in `b` before we run stochastic gradient descent:
# ------------------------------------------------------------------------------------------

b

# ------------------------------------------------------------------------------------------
# After running `stochastic_gradient_descent`, we find the following:
# ------------------------------------------------------------------------------------------

W_final, b_final = stochastic_gradient_descent(loss, W, b, xs, ys, 100000)

# ------------------------------------------------------------------------------------------
# we can look at the values of `W_final` and `b_final`, which our machine learned to
# generate our desired classification.
# ------------------------------------------------------------------------------------------

W_final

b_final

# ------------------------------------------------------------------------------------------
# #### Exercise 8
#
# Plot the data and the learned function.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 9
#
# Do this plot every so often as the learning process is proceeding in order to have an
# animation of the process.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ### Automation with Flux.jl
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# We will need to repeat the above process for a lot of different systems.
# Fortunately, `Flux.jl` provides us with tools to automate this!
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# Flux allows to create a neuron in a simple way:
# ------------------------------------------------------------------------------------------

using Flux

model = Dense(2, 1, σ)

# ------------------------------------------------------------------------------------------
# The `2` and `1` refer to the number of inputs and outputs, and the neuron is defined using
# the $\sigma$ function.
# ------------------------------------------------------------------------------------------

typeof(model)

# ------------------------------------------------------------------------------------------
# We have made an object of type `Dense`, defined by `Flux`, with the name `model`. This
# represents a "dense neural network layer" (see later for more on neural network layers).
# The parameters that will be modified during the learning process live *inside* the `model`
# object.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 10
#
# Investigate which variables live inside the `model` object and what type they are. How
# does that compare to the call to create the `Dense` object that we started with?
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ### Model object as a function
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# We can apply the `model` object to data just as if it were a standard function:
# ------------------------------------------------------------------------------------------

model(rand(2))

# ------------------------------------------------------------------------------------------
# #### Exercise 11
#
# Prove to yourself that you understand what is going on when we call `model`. Create two
# arrays `W` and `b` with the same elements as `model.W` and `model.b`. Use `W` and `b` to
# generate the same answer that you get when we call `model([.5, .5])`.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ### Using Flux
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# We now need to provide Flux with three pieces of information:
#
# 1. A loss function
# 2. Some training data
# 3. An optimization method
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ### Loss functions
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# Flux has various loss functions built in, for example the mean-squared error (`mse`) that
# we have been using:
# ------------------------------------------------------------------------------------------

loss(x, y) = Flux.mse(model(x), y)

# ------------------------------------------------------------------------------------------
# Another common one is the cross entropy, `Flux.crossentropy`.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ### Data
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# The data can take a couple of different forms.
# One form is a single **iterator**, consisting of pairs $(x, y)$ of data and labels.
# We can achieve this with `zip`.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 12
#
# Use `zip` to "zip together" `xs` and `ys`. Then use the `collect` function to check what
# `zip` actually does.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ### Optimization routine
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# Now we need to tell Flux what kind of optimization routine to use. It has several built
# in; the standard stochastic gradient descent algorithm that we have been using is called
# `SGD`. We must pass it two things: a list of parameter objects which will be modified by
# the optimization routine, and a step size:
# ------------------------------------------------------------------------------------------

opt = SGD([model.W, model.b], 0.01)
# give a list of the parameters that will be modified

# ------------------------------------------------------------------------------------------
# The gradient calculations and parameter updates will be carried out by this optimizer
# function; we do not see those details, but if you are curious, you can, of course, look at
# the `Flux.jl` source code!
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ### Training
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# We now have all the pieces in place to actually **train** our model (a single neuron) on
# the data.
# "Training" refers to using pre-labeled data to learn the function that relates the input
# data to the desired output data given by the labels.
#
# `Flux` provides the function `train!`, which performs a single pass through the data and
# does a single step of optimization using the partial cost function for each data point:
# ------------------------------------------------------------------------------------------

Flux.train!(loss, data, opt)

# ------------------------------------------------------------------------------------------
# We can then just repeat this several times to train the network more and coax it towards
# the minimum of the cost function:
# ------------------------------------------------------------------------------------------

for i in 1:100
    Flux.train!(loss, data, opt)
end

# ------------------------------------------------------------------------------------------
# Now let's look at the parameters after training:
# ------------------------------------------------------------------------------------------

model.W

model.b

# ------------------------------------------------------------------------------------------
# Instead of writing out a list of parameters to modify, `Flux` provides the function
# `params`, which extracts all available parameters from a model:
# ------------------------------------------------------------------------------------------

opt = SGD(params(model), 0.01)

params(model)

# ------------------------------------------------------------------------------------------
# ## Adding more features
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 13
#
# So far we have just used two features, red and green.
#
# (i) Add a third feature, blue. Plot the new data.
#
# (ii) Train a neuron with 3 inputs and 1 output on the data.
#
# (iii) Can you find a good way to visualize the result?
# ------------------------------------------------------------------------------------------

.nbexports/introductory-tutorials/broader-topics-and-ecosystem/intro-to-ml/14. Tools - Learning with a single neuron using Flux.jl.jl


# internal method essentially the same as Base.replace!(y, (z .=> r)...)
# but more efficient.
# Similar to the behaviour of `Base.replace!` if `z` contain repetions of values in
# `y` then only the transformation corresponding to the first occurence is performed
# i.e `_replace!([1,5,3], [1,4], 4:5)` would return `[4,5,3]` rather than `[5,5,3]`
# (which replaces `1=>4` and then `4=>5`)
function _replace!(y::AbstractVector, z::AbstractVector, r::AbstractVector)
    length(r) == length(z) ||
     throw(DimensionMismatch("`z` and `r` has to be of the same length"))
    @inbounds for i in eachindex(y)
        for j in eachindex(z)
            isequal(z[j], y[i]) && (y[i] = r[j]; break)
        end
    end
    return y
end

"""
 softmax(X::AbstractMatrix{<:Real})
Return the softmax computed in a numerically stable way:
``σ(x) = exp.(x) ./ sum(exp.(x))``
Implementation taken from NNlib.jl.
"""
function softmax(X::AbstractMatrix{<:Real})
    S = copyto!(similar(X, eltype(X)), X)
    return softmax!(S)
end

"""
 softmax!(X::AbstractMatrix{<:Real})
Return the softmax computed in a numerically stable way:
``σ(x) = exp.(x) ./ sum(exp.(x))`` and store the result in the
input matrix
Implementation taken from NNlib.jl.
"""
function softmax!(X::AbstractMatrix{<:Real})
    max_ = maximum(X, dims=2)
    X .= exp.(X .- max_)
    X ./= sum(X, dims=2)
    return X
end


#=
Challenge #259 [Easy] Clarence the Slow Typist

Description

Clarence is a data entry clerk who works at an internet service provider. His
job is to manually enter the IP addresses of all of the ISP's customers into the
database. He does this using a keypad which has the following layout:

1	2	3
4	5	6
7	8	9
.	0

The distance between the centre of horizontally or vertically adjacent keys is
exactly one centimetre. For instance, the distance between the centres of 3 and
9 would be two centimetres. The distance between the centres of 3 and 5 would be
sqrt 2cm. The Pythagoras theorem is sufficient to calculate the distance between
any two keys.
Clarence, as you might expect from one who works in an ISP, uses a very slow and
inefficient system of typing. He uses a single finger and searches for the key,
and then moves his finger to the key, then presses it, and repeats for all of the
digits in the number. You might know of this style as the "eagle search system"
since the finger searches above the keyboard for the correct key before plunging
down for the keypress, like an eagle plunging down for a kill.
For example, here is how Clarence would type out the number 7851:
    * He starts his finger at 7 and pushes the key.
    * He moves his finger to the right 1cm to 8 and pushes the key.
    * He moves his finger upwards 1cm to 5 and pushes the key.
    * He moves his finger diagonally upwards and left sqrt 2cm to 1 and pushes
      the key.
Therefore the total distance that Clarence moved his finger to type in 7851 is 1
+ 1 + sqrt 2 which is about 3.41cm.
Your task is to write a program that calculates the distance Clarence must move
his finger to type in arbitrary IP addresses.

Formal Inputs and Outputs

Input Description

Input is a string that will be in the form
().().().()
where each () is an integer in the range 0 - 999. This represents the IP address
that Clarence must type in. An example input might be:
219.45.143.143
I would also like to point out that inputs such as 0.42.42.42 or 999.999.999.999
are still valid inputs, despite the fact that they are invalid IP addresses.
So you don't need to include any IP address verification code in your program.

Output Description

Output the distance that Clarence must move his finger in order to type in the
specified IP address. Round answers to two decimal places where needed, and use
the cm unit in your output. The output for the example input is 27.38cm (1 + sqrt
8 + sqrt 5 + 2 + 1 + sqrt 5 + 3 + 1 + sqrt 5 + sqrt 13 + 3 + 1 + sqrt 5).

Credit

This challenge was suggested by /u/katyai. If you have any challenge ideas please
share them on /r/dailyprogrammer_ideas and there's a good chance we'll use them!

( source: https://redd.it/4bc3el )

The version of the user Gobbedyret was so simple -- and good -- that inspired me
=#

function keyboard_distance( key_1::Char, key_2::Char )
    sequence::AbstractString = "123456789.0"

    first::Tuple = divrem( search( sequence, key_1 ) - 1, 3 )
    second::Tuple = divrem( search( sequence, key_2 ) - 1, 3 )

    dx::Integer = second[ 2 ] - first[ 2 ]
    dy::Integer = second[ 1 ] - first[ 1 ]

    return ( dx^2 + dy^2 )^0.5
end

function main( )
    input::AbstractString = chomp( readline( STDIN ) )
    @printf( "%.2fcm\n", sum( [ keyboard_distance( i, j ) for ( i, j ) in zip( input[1:end-1], input[2:end] ) ] ) )
end

main( )


using FiniteDiff, LinearAlgebra, SparseArrays, Test, LinearAlgebra,
      BlockBandedMatrices, ArrayInterfaceCore, BandedMatrices,
      ArrayInterfaceBlockBandedMatrices

fcalls = 0
function f(dx,x)
  global fcalls += 1
  for i in 2:length(x)-1
    dx[i] = x[i-1] - 2x[i] + x[i+1]
  end
  dx[1] = -2x[1] + x[2]
  dx[end] = x[end-1] - 2x[end]
  nothing
end
function oopf(x)
  global fcalls += 1
  vcat([-2x[1]+x[2]],[x[i-1]-2x[i]+x[i+1] for i in 2:length(x)-1],[x[end-1]-2x[end]])
end

function second_derivative_stencil(N)
  A = zeros(N,N)
  for i in 1:N, j in 1:N
      (j-i==-1 || j-i==1) && (A[i,j]=1)
      j-i==0 && (A[i,j]=-2)
  end
  A
end

J = FiniteDiff.finite_difference_jacobian(oopf,rand(30))
@test J ≈ second_derivative_stencil(30)
_J = sparse(J)
@test fcalls == 31

_J2 = similar(_J)
fcalls = 0
FiniteDiff.finite_difference_jacobian!(_J2,f,rand(30),colorvec=repeat(1:3,10))
@test fcalls == 4
@test _J2 ≈ _J

_J2 = similar(_J)
fcalls = 0
FiniteDiff.finite_difference_jacobian!(_J2,f,rand(30),Val{:central},colorvec=repeat(1:3,10))
@test fcalls == 6
@test _J2 ≈ _J

_J2 = similar(_J)
fcalls = 0
FiniteDiff.finite_difference_jacobian!(_J2,f,rand(30),Val{:complex},colorvec=repeat(1:3,10))
@test fcalls == 3
@test _J2 ≈ _J

_J2 = similar(_J)
_denseJ2 = collect(_J2)
fcalls = 0
FiniteDiff.finite_difference_jacobian!(_denseJ2,f,rand(30),colorvec=repeat(1:3,10),sparsity=_J2)
@test fcalls == 4
@test sparse(_denseJ2) ≈ _J

_J2 = similar(_J)
_denseJ2 = collect(_J2)
fcalls = 0
FiniteDiff.finite_difference_jacobian!(_denseJ2,f,rand(30),Val{:central},colorvec=repeat(1:3,10),sparsity=_J2)
@test fcalls == 6
@test sparse(_denseJ2) ≈ _J

_J2 = similar(_J)
_denseJ2 = collect(_J2)
fcalls = 0
FiniteDiff.finite_difference_jacobian!(_denseJ2,f,rand(30),Val{:complex},colorvec=repeat(1:3,10),sparsity=_J2)
@test fcalls == 3
@test sparse(_denseJ2) ≈ _J

_J2 = similar(Tridiagonal(_J))
fcalls = 0
FiniteDiff.finite_difference_jacobian!(_J2,f,rand(30),colorvec=repeat(1:3,10))
@test fcalls == 4
@test _J2 ≈ _J

_J2 = similar(_J2)
fcalls = 0
FiniteDiff.finite_difference_jacobian!(_J2,f,rand(30),Val{:central},colorvec=repeat(1:3,10))
@test fcalls == 6
@test _J2 ≈ _J

_J2 = similar(_J2)
fcalls = 0
FiniteDiff.finite_difference_jacobian!(_J2,f,rand(30),Val{:complex},colorvec=repeat(1:3,10))
@test fcalls == 3
@test _J2 ≈ _J

_Jb = BandedMatrices.BandedMatrix(similar(_J2),(1,1))
FiniteDiff.finite_difference_jacobian!(_Jb, f, rand(30), colorvec=repeat(1:3,10))
@test _Jb ≈ _J

_Jtri = Tridiagonal(similar(_J2))
FiniteDiff.finite_difference_jacobian!(_Jtri, f, rand(30), colorvec=repeat(1:3,10))
@test _Jtri ≈ _J

#https://github.com/JuliaDiff/FiniteDiff.jl/issues/67#issuecomment-516871956
function f(out, x)
	x = reshape(x, 100, 100)
	out = reshape(out, 100, 100)
	for i in 1:100
		for j in 1:100
			out[i, j] = x[i, j] + x[max(i -1, 1), j] + x[min(i+1, size(x, 1)), j] +  x[i, max(j-1, 1)]  + x[i, min(j+1, size(x, 2))]
		end
	end
	return vec(out)
end
x = rand(10000)
Jbbb = BandedBlockBandedMatrix(Ones(10000, 10000), fill(100, 100), fill(100, 100), (1, 1), (1, 1))
Jsparse = sparse(Jbbb)
colorsbbb = ArrayInterfaceCore.matrix_colors(Jbbb)
FiniteDiff.finite_difference_jacobian!(Jbbb, f, x, colorvec=colorsbbb)
FiniteDiff.finite_difference_jacobian!(Jsparse, f, x, colorvec=colorsbbb)
@test Jbbb ≈ Jsparse
Jbb = BlockBandedMatrix(similar(Jsparse),fill(100, 100), fill(100, 100),(1,1));
colorsbb = ArrayInterfaceCore.matrix_colors(Jbb)
FiniteDiff.finite_difference_jacobian!(Jbb, f, x, colorvec=colorsbb)
@test Jbb ≈ Jsparse

# Non-square Jacobian test.
# The Jacobian of f_nonsquare! has size (n, 2*n).
function f_nonsquare!(y, x)
    global fcalls += 1
    @assert length(x) == 2*length(y)
    n = length(x) ÷ 2
    x1 = @view x[1:n]
    x2 = @view x[n+1:end]

    @. y = (x1 .- 3).^2 .+ x1.*x2 .+ (x2 .+ 4).^2 .- 3
    return nothing
end

n = 4
x0 = vcat(ones(n).*(1:n) .+ 0.5, ones(n).*(1:n) .+ 1.5)
y0 = zeros(n)
rows = vcat([i for i in 1:n], [i for i in 1:n])
cols = vcat([i for i in 1:n], [i+n for i in 1:n])
sparsity = sparse(rows, cols, ones(length(rows)))
colorvec = vcat(fill(1, n), fill(2, n))

J_nonsquare1 = zeros(size(sparsity))
FiniteDiff.finite_difference_jacobian!(J_nonsquare1, f_nonsquare!, x0)

J_nonsquare2 = similar(sparsity)
for method in [Val(:forward), Val(:central), Val(:complex)]
    cache = FiniteDiff.JacobianCache(copy(x0), copy(y0), copy(y0), method; sparsity, colorvec)
    global fcalls = 0
    FiniteDiff.finite_difference_jacobian!(J_nonsquare2, f_nonsquare!, x0, cache)
    if method == Val(:central)
        @test fcalls == 2*maximum(colorvec)
    elseif method == Val(:complex)
        @test fcalls == maximum(colorvec)
    else
        @test fcalls == maximum(colorvec) + 1
    end
    @test isapprox(J_nonsquare2, J_nonsquare1; rtol=1e-6)
end


function Base._cpow(z::Complex{DoubleFloat{T}}, p::Complex{DoubleFloat{T}}) where {T<:IEEEFloat}
    pᵣ, pᵢ = reim(p)
    r = abs(z)
    θ = angle(z)
    rᵖ = r^pᵣ * exp(-pᵢ*θ)
    ϕ = pᵣ*θ + pᵢ*log(r)
    if isfinite(ϕ)
        return rᵖ * cis(ϕ)
    else
        iszero(rᵖ) && return zero(Complex{DoubleFloat{T}}) # no way to get correct signs of 0.0
        return Complex{DoubleFloat{T}}(NaN) # non-finite phase angle or NaN input
    end
end


@doc raw"""
    GeneralizedGrassmann{n,k,𝔽} <: AbstractEmbeddedManifold{𝔽,DefaultEmbeddingType}

The generalized Grassmann manifold $\operatorname{Gr}(n,k,B)$ consists of all subspaces
spanned by $k$ linear independent vectors $𝔽^n$, where $𝔽  ∈ \{ℝ, ℂ\}$ is either the real- (or complex-) valued vectors.
This yields all $k$-dimensional subspaces of $ℝ^n$ for the real-valued case and all $2k$-dimensional subspaces
of $ℂ^n$ for the second.

The manifold can be represented as

````math
\operatorname{Gr}(n, k, B) := \bigl\{ \operatorname{span}(p)\ \big|\ p ∈ 𝔽^{n × k}, p^\mathrm{H}Bp = I_k\},
````

where $\cdot^{\mathrm{H}}$ denotes the complex conjugate (or Hermitian) transpose and
$I_k$ is the $k × k$ identity matrix. This means, that the columns of $p$
form an unitary basis of the subspace with respect to the scaled inner product, that is a
point on $\operatorname{Gr}(n,k,B)$, and hence the subspace can actually be represented by
a whole equivalence class of representers. For $B=I_n$ this simplifies to the [`Grassmann`](@ref) manifold.

The tangent space at a point (subspace) $p$ is given by

````math
T_x\mathrm{Gr}(n,k,B) = \bigl\{
X ∈ 𝔽^{n × k} :
X^{\mathrm{H}}Bp + p^{\mathrm{H}}BX = 0_{k} \bigr\},
````

where $0_{k}$ denotes the $k × k$ zero matrix.

Note that a point $p ∈ \operatorname{Gr}(n,k,B)$ might be represented by
different matrices (i.e. matrices with $B$-unitary column vectors that span
the same subspace). Different representations of $p$ also lead to different
representation matrices for the tangent space $T_p\mathrm{Gr}(n,k,B)$

The manifold is named after
[Hermann G. Graßmann](https://en.wikipedia.org/wiki/Hermann_Grassmann) (1809-1877).

# Constructor

    GeneralizedGrassmann(n, k, B=I_n, field=ℝ)

Generate the (real-valued) Generalized Grassmann manifold of $n\times k$ dimensional
orthonormal matrices with scalar product `B`.
"""
struct GeneralizedGrassmann{n,k,𝔽,TB<:AbstractMatrix} <:
       AbstractEmbeddedManifold{𝔽,DefaultEmbeddingType}
    B::TB
end

function GeneralizedGrassmann(
    n::Int,
    k::Int,
    B::AbstractMatrix=Matrix{Float64}(I, n, n),
    field::AbstractNumbers=ℝ,
)
    return GeneralizedGrassmann{n,k,field,typeof(B)}(B)
end

@doc raw"""
    check_point(M::GeneralizedGrassmann{n,k,𝔽}, p)

Check whether `p` is representing a point on the [`GeneralizedGrassmann`](@ref) `M`, i.e. its
a `n`-by-`k` matrix of unitary column vectors with respect to the B inner prudct and
of correct `eltype` with respect to `𝔽`.
"""
function check_point(M::GeneralizedGrassmann{n,k,𝔽}, p; kwargs...) where {n,k,𝔽}
    mpv = invoke(
        check_point,
        Tuple{typeof(get_embedding(M)),typeof(p)},
        get_embedding(M),
        p;
        kwargs...,
    )
    mpv === nothing || return mpv
    c = p' * M.B * p
    if !isapprox(c, one(c); kwargs...)
        return DomainError(
            norm(c - one(c)),
            "The point $(p) does not lie on $(M), because x'Bx is not the unit matrix.",
        )
    end
    return nothing
end

@doc raw"""
    check_vector(M::GeneralizedGrassmann{n,k,𝔽}, p, X; kwargs...)

Check whether `X` is a tangent vector in the tangent space of `p` on
the [`GeneralizedGrassmann`](@ref) `M`, i.e. that `X` is of size and type as well as that

````math
    p^{\mathrm{H}}BX + \overline{X^{\mathrm{H}}Bp} = 0_k,
````

where $\cdot^{\mathrm{H}}$ denotes the complex conjugate transpose or Hermitian,
$\overline{\cdot}$ the (elementwise) complex conjugate, and $0_k$ denotes the $k × k$ zero natrix.
"""
function check_vector(M::GeneralizedGrassmann{n,k,𝔽}, p, X; kwargs...) where {n,k,𝔽}
    mpv = invoke(
        check_vector,
        Tuple{typeof(get_embedding(M)),typeof(p),typeof(X)},
        get_embedding(M),
        p,
        X;
        kwargs...,
    )
    mpv === nothing || return mpv
    if !isapprox(p' * M.B * X, -conj(X' * M.B * p); kwargs...)
        return DomainError(
            norm(p' * M.B * X + conj(X' * M.B * p)),
            "The matrix $(X) does not lie in the tangent space of $(p) on $(M), since x'Bv + v'Bx is not the zero matrix.",
        )
    end
    return nothing
end

function decorated_manifold(M::GeneralizedGrassmann{N,K,𝔽}) where {N,K,𝔽}
    return Euclidean(N, K; field=𝔽)
end

@doc raw"""
    distance(M::GeneralizedGrassmann, p, q)

Compute the Riemannian distance on [`GeneralizedGrassmann`](@ref)
manifold `M`$= \mathrm{Gr}(n,k,B)$.

Let $USV = p^\mathrm{H}Bq$ denote the SVD decomposition of
$p^\mathrm{H}Bq$, where $\cdot^{\mathrm{H}}$ denotes the complex
conjugate transposed or Hermitian. Then the distance is given by
````math
d_{\mathrm{Gr}(n,k,B)}(p,q) = \operatorname{norm}(\operatorname{Re}(b)).
````
where

````math
b_{i}=\begin{cases}
0 & \text{if} \; S_i ≥ 1\\
\arccos(S_i) & \, \text{if} \; S_i<1.
\end{cases}
````
"""
function distance(M::GeneralizedGrassmann, p, q)
    p ≈ q && return zero(real(eltype(p)))
    a = svd(p' * M.B * q).S
    return sqrt(sum(x -> abs2(acos(clamp(x, -1, 1))), a))
end

@doc raw"""
    exp(M::GeneralizedGrassmann, p, X)

Compute the exponential map on the [`GeneralizedGrassmann`](@ref) `M`$= \mathrm{Gr}(n,k,B)$
starting in `p` with tangent vector (direction) `X`. Let $X^{\mathrm{H}}BX = USV$ denote the
SVD decomposition of $X^{\mathrm{H}}BX$. Then the exponential map is written using

````math
\exp_p X = p V\cos(S)V^\mathrm{H} + U\sin(S)V^\mathrm{H},
````

where $\cdot^{\mathrm{H}}$ denotes the complex conjugate transposed or Hermitian and the
cosine and sine are applied element wise to the diagonal entries of $S$.
"""
exp(::GeneralizedGrassmann, ::Any...)

function exp!(M::GeneralizedGrassmann, q, p, X)
    norm(M, p, X) ≈ 0 && return copyto!(q, p)
    d = svd(X' * M.B * X)
    V = d.Vt
    S = abs.(sqrt.(d.S))
    U = X * (V / Diagonal(S))
    z = p * V * Diagonal(cos.(S)) * V + U * Diagonal(sin.(S)) * V
    return copyto!(q, project(M, z))
end

@doc raw"""
    injectivity_radius(M::GeneralizedGrassmann)
    injectivity_radius(M::GeneralizedGrassmann, p)

Return the injectivity radius on the [`GeneralizedGrassmann`](@ref) `M`,
which is $\frac{π}{2}$.
"""
injectivity_radius(::GeneralizedGrassmann) = π / 2
injectivity_radius(::GeneralizedGrassmann, ::ExponentialRetraction) = π / 2
injectivity_radius(::GeneralizedGrassmann, ::Any) = π / 2
injectivity_radius(::GeneralizedGrassmann, ::Any, ::ExponentialRetraction) = π / 2
eval(
    quote
        @invoke_maker 1 AbstractManifold injectivity_radius(
            M::GeneralizedGrassmann,
            rm::AbstractRetractionMethod,
        )
    end,
)

@doc raw"""
    inner(M::GeneralizedGrassmann, p, X, Y)

Compute the inner product for two tangent vectors `X`, `Y` from the tangent space
of `p` on the [`GeneralizedGrassmann`](@ref) manifold `M`. The formula reads

````math
g_p(X,Y) = \operatorname{tr}(X^{\mathrm{H}}BY),
````

where $\cdot^{\mathrm{H}}$ denotes the complex conjugate transposed or Hermitian.
"""
inner(M::GeneralizedGrassmann{n,k}, p, X, Y) where {n,k} = dot(X, M.B * Y)

function Base.isapprox(M::GeneralizedGrassmann, p, X, Y; kwargs...)
    return isapprox(sqrt(inner(M, p, zero_vector(M, p), X - Y)), 0; kwargs...)
end
function Base.isapprox(M::GeneralizedGrassmann, p, q; kwargs...)
    return isapprox(distance(M, p, q), 0.0; kwargs...)
end

@doc raw"""
    log(M::GeneralizedGrassmann, p, q)

Compute the logarithmic map on the [`GeneralizedGrassmann`](@ref) `M`$ = \mathcal M=\mathrm{Gr}(n,k,B)$,
i.e. the tangent vector `X` whose corresponding [`geodesic`](@ref) starting from `p`
reaches `q` after time 1 on `M`. The formula reads

````math
\log_p q = V\cdot \operatorname{atan}(S) \cdot U^\mathrm{H},
````

where $\cdot^{\mathrm{H}}$ denotes the complex conjugate transposed or Hermitian.
The matrices $U$ and $V$ are the unitary matrices, and $S$ is the diagonal matrix
containing the singular values of the SVD-decomposition

````math
USV = (q^\mathrm{H}Bp)^{-1} ( q^\mathrm{H} - q^\mathrm{H}Bpp^\mathrm{H}).
````

In this formula the $\operatorname{atan}$ is meant elementwise.
"""
log(::GeneralizedGrassmann, ::Any...)

function log!(M::GeneralizedGrassmann{n,k}, X, p, q) where {n,k}
    z = q' * M.B * p
    At = q' - z * p'
    Bt = z \ At
    d = svd(Bt')
    return X .= view(d.U, :, 1:k) * Diagonal(atan.(view(d.S, 1:k))) * view(d.Vt, 1:k, :)
end

@doc raw"""
    manifold_dimension(M::GeneralizedGrassmann)

Return the dimension of the [`GeneralizedGrassmann(n,k,𝔽)`](@ref) manifold `M`, i.e.

````math
\dim \operatorname{Gr}(n,k,B) = k(n-k) \dim_ℝ 𝔽,
````

where $\dim_ℝ 𝔽$ is the [`real_dimension`](@ref) of `𝔽`.
"""
function manifold_dimension(::GeneralizedGrassmann{n,k,𝔽}) where {n,k,𝔽}
    return k * (n - k) * real_dimension(𝔽)
end

"""
    mean(
        M::GeneralizedGrassmann,
        x::AbstractVector,
        [w::AbstractWeights,]
        method = GeodesicInterpolationWithinRadius(π/4);
        kwargs...,
    )

Compute the Riemannian [`mean`](@ref mean(M::AbstractManifold, args...)) of `x` using
[`GeodesicInterpolationWithinRadius`](@ref).
"""
mean(::GeneralizedGrassmann{n,k} where {n,k}, ::Any...)

function Statistics.mean!(
    M::GeneralizedGrassmann{n,k},
    p,
    x::AbstractVector,
    w::AbstractVector;
    kwargs...,
) where {n,k}
    return mean!(M, p, x, w, GeodesicInterpolationWithinRadius(π / 4); kwargs...)
end

@doc raw"""
    project(M::GeneralizedGrassmann, p)

Project `p` from the embedding onto the [`GeneralizedGrassmann`](@ref) `M`, i.e. compute `q`
as the polar decomposition of $p$ such that $q^{\mathrm{H}}Bq$ is the identity,
where $\cdot^{\mathrm{H}}$ denotes the Hermitian, i.e. complex conjugate transpose.
"""
project(::GeneralizedGrassmann, ::Any)

function project!(M::GeneralizedGrassmann, q, p)
    s = svd(p)
    e = eigen(s.U' * M.B * s.U)
    qsinv = e.vectors * Diagonal(1 ./ sqrt.(e.values))
    q .= s.U * qsinv * e.vectors' * s.V'
    return q
end

@doc raw"""
    project(M::GeneralizedGrassmann, p, X)

Project the `n`-by-`k` `X` onto the tangent space of `p` on the
[`GeneralizedGrassmann`](@ref) `M`, which is computed by

````math
\operatorname{proj_p}(X) = X - pp^{\mathrm{H}}B^\mathrm{T}X,
````

where $\cdot^{\mathrm{H}}$ denotes the complex conjugate transposed or Hermitian
and $\cdot^{\mathrm{T}}$ the transpose.
"""
project(::GeneralizedGrassmann, ::Any, ::Any)

function project!(M::GeneralizedGrassmann, Y, p, X)
    A = p' * M.B' * X
    copyto!(Y, X - p * Hermitian((A + A') / 2))
    return Y
end

@doc raw"""
    representation_size(M::GeneralizedGrassmann{n,k})

Return the represenation size or matrix dimension of a point on the [`GeneralizedGrassmann`](@ref)
`M`, i.e. $(n,k)$ for both the real-valued and the complex value case.
"""
@generated representation_size(::GeneralizedGrassmann{n,k}) where {n,k} = (n, k)

@doc raw"""
    retract(M::GeneralizedGrassmann, p, X, ::PolarRetraction)

Compute the SVD-based retraction [`PolarRetraction`](@ref) on the
[`GeneralizedGrassmann`](@ref) `M`, by
[`project`](@ref project(M::GeneralizedGrassmann, p))ing $p + X$ onto `M`.
"""
retract(::GeneralizedGrassmann, ::Any, ::Any, ::PolarRetraction)

function retract!(M::GeneralizedGrassmann, q, p, X, ::PolarRetraction)
    project!(M, q, p + X)
    return q
end
function retract!(M::GeneralizedGrassmann, q, p, X, ::ProjectionRetraction)
    project!(M, q, p + X)
    return q
end

function Base.show(io::IO, M::GeneralizedGrassmann{n,k,𝔽}) where {n,k,𝔽}
    return print(io, "GeneralizedGrassmann($(n), $(k), $(M.B), $(𝔽))")
end

@doc raw"""
    vector_transport_to(M::GeneralizedGrassmann, p, X, q, ::ProjectionTransport)

Compute the vector transport of the tangent vector `X` at `p` to `q`,
using the [`project`](@ref project(::GeneralizedGrassmann, ::Any...))
of `X` to `q`.
"""
vector_transport_to(::GeneralizedGrassmann, ::Any, ::Any, ::Any, ::ProjectionTransport)

function vector_transport_to!(M::GeneralizedGrassmann, Y, p, X, q, ::ProjectionTransport)
    project!(M, Y, q, X)
    return Y
end

@doc raw"""
    zero_vector(M::GeneralizedGrassmann, p)

Return the zero tangent vector from the tangent space at `p` on the
[`GeneralizedGrassmann`](@ref) `M`, which is given by a zero matrix the same size as `p`.
"""
zero_vector(::GeneralizedGrassmann, ::Any...)

zero_vector!(::GeneralizedGrassmann, X, p) = fill!(X, 0)


using FinancialMonteCarlo

S0 = 100.0;
K = 100.0;
r = 0.02;
T = 1.0;
d = 0.01;
D = 90.0;

Nsim = 10000;
Nstep = 300;
sigma = 0.2;
variate_ = FinancialMonteCarlo.ControlVariates(Forward(T), MonteCarloConfiguration(1000, 100))
variate_2 = FinancialMonteCarlo.ControlVariates(AsianFloatingStrikeOption(T), MonteCarloConfiguration(1000, 100))
mc = MonteCarloConfiguration(Nsim, Nstep, variate_);
mc2 = MonteCarloConfiguration(Nsim, Nstep, variate_2);
mc1 = MonteCarloConfiguration(Nsim, Nstep);
toll = 1e-3;

spotData1 = ZeroRate(r);

FwdData = Forward(T)
EUData = EuropeanOption(T, K)
AMData = AmericanOption(T, K)
BarrierData = BarrierOptionDownOut(T, K, D)
AsianFloatingStrikeData = AsianFloatingStrikeOption(T)
AsianFixedStrikeData = AsianFixedStrikeOption(T, K)
Model = BlackScholesProcess(sigma, Underlying(S0, d));

AsianPrice1 = pricer(Model, spotData1, mc, AsianFloatingStrikeData);
AsianPrice2 = pricer(Model, spotData1, mc, AsianFixedStrikeData);
AsianPrice2 = pricer(Model, spotData1, mc2, AsianFixedStrikeData);


"""
Model for ZIP Load model given by:

P_zip = P_power + P_current * (V / V0) + P_impedance * (V / V0)^2
Q_zip = Q_power + Q_current * (V / V0) + Q_impedance * (V / V0)^2

with V = sqrt(V_r^2 + V_i^2) and V0 the voltage magnitude from the power flow solution

The current taken for the load is computed as:
I_zip = (P_zip + j Q_zip)^* / (V_r + j V_i)^*
I_zip = (P_zip - j Q_zip) / (V_r - j V_i)

For constant impedance it is obtained:
Iz_re = (1 / V0)^2 * (V_r * P_impedance + V_i * Q_impedance)
Iz_im = (1 / V0)^2 * (V_i * P_impedance - V_r * Q_impedance)

For constant current it is obtained:
Ii_re = (1 / V0) * ( (V_r * P_current + V_i * Q_current) / V )
Ii_im = (1 / V0) * ( (V_i * P_current - V_r * Q_current) / V )

For constant power it is obtained:
Ip_re =  (V_r * P_power + V_i * Q_power) / V^2
Ip_im =  (V_i * P_power - V_r * Q_power) / V^2

Model for Exponential Load model given by:

P_exp = P0 * (V / V0)^α
Q_exp = Q0 * (V / V0)^β

The current taken for the load is computed as:
I_exp = (P_exp + j Q_exp)^* / (V_r + j V_i)^*
I_exp = (P_exp - j Q_exp) / (V_r - j V_i)

It results:
Ir_exp = V_r * P0 * (V^(α - 2) / V0^α) + V_i * Q0 * (V^(β - 2)/ V0^β)
Ii_im  = V_i * P0 * (V^(α - 2) / V0^α) - V_r * Q0 * (V^(β - 2)/ V0^β)

"""
function mdl_zip_load!(
    voltage_r::T,
    voltage_i::T,
    current_r::AbstractArray{T},
    current_i::AbstractArray{T},
    wrapper::StaticLoadWrapper,
) where {T <: ACCEPTED_REAL_TYPES}
    # Read power flow voltages
    #V0_mag_inv = 1.0 / get_V_ref(wrapper)
    V0_mag_inv = 1.0 / PSY.get_magnitude(PSY.get_bus(wrapper))
    V0_mag_sq_inv = V0_mag_inv^2

    V_mag = sqrt(voltage_r^2 + voltage_i^2)
    V_mag_inv = 1.0 / V_mag
    V_mag_sq_inv = V_mag_inv^2

    # Load device parameters
    P_power = get_P_power(wrapper)
    P_current = get_P_current(wrapper)
    P_impedance = get_P_impedance(wrapper)
    Q_power = get_Q_power(wrapper)
    Q_current = get_Q_current(wrapper)
    Q_impedance = get_Q_impedance(wrapper)
    exp_params = get_exp_params(wrapper)
    Ir_exp = zero(T)
    Ii_exp = zero(T)

    # Compute ZIP currents
    Iz_re = V0_mag_sq_inv * (voltage_r * P_impedance + voltage_i * Q_impedance)
    Iz_im = V0_mag_sq_inv * (voltage_i * P_impedance - voltage_r * Q_impedance)

    Ii_re = V0_mag_inv * V_mag_inv * (voltage_r * P_current + voltage_i * Q_current)
    Ii_im = V0_mag_inv * V_mag_inv * (voltage_i * P_current - voltage_r * Q_current)

    Ip_re = V_mag_sq_inv * (voltage_r * P_power + voltage_i * Q_power)
    Ip_im = V_mag_sq_inv * (voltage_i * P_power - voltage_r * Q_power)

    # Compute Exponential Currents
    if !isempty(exp_params)
        for tuple in exp_params
            P0 = tuple.P_exp
            α = tuple.P_coeff
            Q0 = tuple.Q_exp
            β = tuple.Q_coeff

            # Compute Auxiliary terms
            V_coeff_active = V_mag^(α - 2.0) * V0_mag_inv^α
            V_coeff_reactive = V_mag^(β - 2.0) * V0_mag_inv^β

            # Compute currents
            Ir_exp += voltage_r * P0 * V_coeff_active + voltage_i * Q0 * V_coeff_reactive
            Ii_exp += voltage_i * P0 * V_coeff_active - voltage_r * Q0 * V_coeff_reactive
        end
    end

    # Update current
    current_r[1] += -(Iz_re + Ii_re + Ip_re + Ir_exp) #in system pu flowing out
    current_i[1] += -(Iz_im + Ii_im + Ip_im + Ii_exp) #in system pu flowing out

    return
end


include("../src/modeling/util.jl")

function model(t :: Trace, data)
    loc = sample(t, "loc", Normal(0.0, 10.0))
    scale = sample(t, "scale", Gamma(3.0, 3.0))
    for (i, d) in enumerate(data)
        observe(t, "data $i", Normal(loc, scale), d)
    end
end

@testset "nested sampling 1" begin
    true_loc = 8.0
    true_scale = 2.0
    data = rand(4) .* true_scale .+ true_loc
    @info "True loc = $true_loc"
    @info "True scale = $true_scale"

    for n in [3, 4, 8, 15, 25, 50]
        @info "Nested sampling inference with $n points"
        @time results = nested(model; params = (data,), num_points = n)
        mean_loc = mean(results, "loc")
        mean_std = mean(results, "scale")
        @info "Posterior mean loc = $mean_loc"
        @info "Posterior mean scale = $mean_std"
        log_Z = logsumexp(results.log_weights)
        @info "log p(x) ≈ $log_Z"
    end
end


using PyPlot, Distributions

using TransitionalMCMC

# Prior Bounds
lb, ub  = -5, 5

# Prior Density and sampler
logprior(x) = logpdf(Uniform(lb, ub), x[1]) + logpdf(Uniform(lb, ub), x[2])
priorRnd(Nsamples) = rand(Uniform(lb, ub), Nsamples, 2)

# Log Likelihood
function logLik(x)
    return -1 * ((x[1]^2 + x[2] - 11)^2 + (x[1] + x[2]^2 - 7)^2)
end

Nsamples = 200

samps, acc = tmcmc(logLik, logprior, priorRnd, Nsamples, 5, 2)

plt.scatter(samps[:,1], samps[:,2])


using ModelingToolkit
using SteadyStateDiffEq
using Test

@parameters t r
@variables x(t)
@derivatives D'~t
eqs = [D(x) ~ x^2-r]
de = ODESystem(eqs)

for factor in [1e-1, 1e0, 1e10], u0_p in [(2.34,2.676),(22.34,1.632),(.3,15.676),(0.3,0.006)]
    u0 = [x => factor*u0_p[1]]
    p = [r => factor*u0_p[2]]
    ss_prob = SteadyStateProblem(de,u0,p)
    sol = solve(ss_prob,SSRootfind()).u[1]
    @test abs(sol^2 - factor*u0_p[2]) < 1e-8
    ss_prob = SteadyStateProblemExpr(de,u0,p)
    sol_expr = solve(eval(ss_prob),SSRootfind()).u[1]
    @test all(x->x==0,sol-sol_expr)
end


function bessely0(x::BigFloat)
    if x <= zero(x)
        if iszero(x)
            return -Inf
        else
            return throw(DomainError(x, "NaN result for non-NaN input."))
        end
    elseif isinf(x)
        return zero(x)
    end

    xx = abs(x)

    if xx <= 2.0
        z = xx * xx
        p = evalpoly(z, Y0_2N) / evalpoly(z, Y0_2D)
        p = TWOOPI(BigFloat) * log(x) * besselj0(x) + p
        return p
    end

    xinv = inv(xx)
    z = xinv * xinv
    if xinv <= 0.25
        if xinv <= 0.125
            if xinv <= 0.0625
                p = evalpoly(z, P16_IN) / evalpoly(z, P16_ID)
                q = evalpoly(z, Q16_IN) / evalpoly(z, Q16_ID)
            else
                p = evalpoly(z, P8_16N) / evalpoly(z, P8_16D)
                q = evalpoly(z, Q8_16N) / evalpoly(z, Q8_16D)
            end
        elseif xinv <= 0.1875
            p = evalpoly(z, P5_8N) / evalpoly(z, P5_8D)
	        q = evalpoly(z, Q5_8N) / evalpoly(z, Q5_8D)
        else
            p = evalpoly(z, P4_5N) / evalpoly(z, P4_5D)
	        q = evalpoly(z, Q4_5N) / evalpoly(z, Q4_5D)
        end
    else
        if xinv <= 0.375
            if xinv <= 0.3125
                p = evalpoly(z, P3r2_4N) / evalpoly(z, P3r2_4D)
	            q = evalpoly(z, Q3r2_4N) / evalpoly(z, Q3r2_4D)
            else
                p = evalpoly(z, P2r7_3r2N) / evalpoly(z, P2r7_3r2D)
	            q = evalpoly(z, Q2r7_3r2N) / evalpoly(z, Q2r7_3r2D)
            end
        elseif xinv <= 0.4375
            p = evalpoly(z, P2r3_2r7N) / evalpoly(z, P2r3_2r7D)
	        q = evalpoly(z, Q2r3_2r7N) / evalpoly(z, Q2r3_2r7D)
        else
            p = evalpoly(z, P2_2r3N) / evalpoly(z, P2_2r3D)
	        q = evalpoly(z, Q2_2r3N) / evalpoly(z, Q2_2r3D)
        end
    end

    p = 1.0 + z * p
    q = z * xinv * q
    q = q - 0.125 * xinv
    c = cos(x)
    s = sin(x)
    ss = s - c
    cc = s + c
    z = -cos(x + x)
    if s * c < zero(x)
        cc = z / ss
    else
        ss = z / cc
    end
    z = ONEOSQPI(BigFloat) * (p * ss + q * cc) / sqrt(x)
    return z
end


# Old linalg.jl functions, for comparison with the new ones

#### SVD, i.e. UDV decomposition
function decompose_udv!(A::Matrix{T}) where T<:Number
    Base.LinAlg.LAPACK.gesvd!('A','A',A)
end

function decompose_udv(A::Matrix{T}) where T<:Number
    X = copy(A)
    decompose_udv!(X)
    return X
end

#### QR, i.e. UDT decomposition
function decompose_udt(A::AbstractMatrix{C}) where C<:Number
    F = qr(A, Val(true))
    p = F.p
    @views p[p] = collect(1:length(p))
    # D = abs.(real(diag(triu(R))))
    D = abs.(real(diag(F.R)))
    T = (Diagonal(1. ./ D) * F.R)[:, p]
    return Matrix(F.Q), D, T
end

#### Other
function expm_diag!(A::Matrix{T}) where T<:Number
    F = eigfact!(A)
    return F[:vectors] * spdiagm(0 => exp(F[:values])) * ctranspose(F[:vectors])
end

function lu_det(M)
    L, U, p = lu(M)
    return prod(diag(L)) * prod(diag(U))
end


module FF

export ff

import Base: /, *, +, -, promote_rule
import BitIntegers: UBI, @uint256_str, UInt256

const _P_ = uint256"21888242871839275222246405745257275088548364400416034343698204186575808495617"

struct ff <: Integer
    v::UInt256
    ff(v::T) where T <: Integer = begin
        t = UInt256(mod(v, _P_))
        return (t >= 0) ? new(t) : new(t + _P_)
    end
    ff(v::ff) = v
end

@generated promote_rule(::Type{ff}, ::Type{<:UBI}) = ff
(::Type{T})(x::ff) where T <: UBI = T(x.v)

(+)(a::ff, b::ff) = ff(a.v + b.v)
(-)(a::ff, b::ff)  = ff(_P_ + a.v - b.v)
(-)(a::ff)  = ff(_P_ - a.v)
(*)(a::ff, b::ff) = ff(widemul(a.v, b.v) % _P_)
(/)(a::ff, b::ff) = ff(widemul(a.v, invmod(b.v, _P_)) % _P_)

iszero(x::ff) = iszero(x.v)
end #module


using ApproxFunOrthogonalPolynomials, ApproxFunRational,
 ApproxFunFourier, ApproxFunBase, ApproxFun, AbstractIterativeSolvers, Plots

### Scalar case
L = 1.
g = x -> .2exp(-x^2)
G1 = Fun(zai(g),OscLaurent(1.,L),100) + Fun(zai(g),OscLaurent(0.,L),100)
h = x -> .2/(1 + x^2)
H2 = Fun(zai(h),OscLaurent(-1.,L),100)
𝓒 = Cauchy(-1)
𝓢 = Cauchy(1)

G=zeros(Fun,2,2)
G[1,1] = G1
G[2,2] = G1
G[1,2] = H2
G[2,1] = H2
G = copy(Fun(G))

One = pad(Fun(1.,OscLaurent(0.0,L)),40)
b = G*Fun([One;0*One])

op = x -> condense(x - G*(𝓒*x))
out = GMRES(op,b,⋅,1e-3,100, condense)
u = sum([out[2][i]*out[1][i] for i=1:length(out[2]) ])

id = Array([1. 0.; 0. 1.])
𝓢 = Cauchy(1)
𝓒 = Cauchy(-1)
c = id*[1.0; 0.]

id + G(.1)

jump = x -> (𝓢*u)(x) + c - (id + G(x))*((𝓒*u)(x)+c)
jump(.11)


"""
Uses `hunting-circular-orbits.jl` and the Page & Thorne (1974) eq. (12)
to simulate a flux model.
Works for a < 0.97, then the orbit finder starts to panic close to the ISCO.
"""

# using Revise

using GeodesicTracer
using ComputedGeodesicEquations
using GeodesicBase
using AccretionFormulae

import CarterBoyerLindquist # for rms function
import LinearAlgebra

using StaticArrays
using Interpolations

# include("/Users/lucyackland-snow/Desktop/DevEnv1/dev/AccretionFormulae/src/johannsen.jl")
include("hunting-circular-orbits.jl")

# for Kerr, this is just `u[2]`, but to illustrate
# how you could do this for other metrics
function sqrt_g_tilde(m, u)
    metric = GeodesicBase.metric(m, u)
    metric_no_θ = metric[setdiff(1:4, 3), setdiff(1:4, 3)]

    sqrt(-LinearAlgebra.det(metric_no_θ))
end
"""
Johannsen 2014 eq 28
"""
function sqrt_g_tilde(r; a=0.97, M=0, ϵ_3=0, α_13=0, α_22=0, α_52=0)
    sqrt_g_tilde =   (r/abs(1 + (α_13 * M^3)/(r^3) -
                        (α_22 * a^2 * M^2)/(r^4) +
                        (α_13 * a^2 * M^3)/(r^5))) *
                √(((1 + (ϵ_3 * M^3)/(r^3))^3)/
                (1 + (α_52 * M^2)/(r^2)))
end

# gradient function
#
# we just compute discrete derivatives, since this is
# sufficient when we have small `r`.
# using interpolations can make the resolution better
# and increase the accuracy, but this is a good quick-and-dirty way
# to see if our results are going in the right direction
function ∇r(array, r)
    deriv = diff(array) ./ diff(r)
    # ensure same dimension using a lazy extension
    push!(deriv, deriv[end])
    deriv
end

"""
    f_approx(Lz, ∇rwt, ∇rLz_array, wt_array)
A discretized version of `f` from Page & Thorne (1974) eq. (12):
```math
f(r_n) \\approx - \\left. \\left[
   \\frac{1}{L_z} \\left(\\frac{\\partial \\dot{u}^t}{\\partial r} \\right)
\\right] \\right\\rvert_{r = r_n}
\\sum_{i=0}^{n} \\left\\{ \\left. \\left[
   \\frac{1}{\\dot{u}^t} \\left( \\frac{\\partial L_z}{\\partial r}\\right)
\\right] \\right\\rvert_{r = r_i}
\\right\\}
```
"""
function f_approx(Lz, ∇rwt, ∇rLz_array, wt_array, r; a=0.97, M=1.0, ϵ_3=0, α_13=0, α_22=0, α_52=0)
    -(∇rwt / Lz) * sum(∇rLz_array ./ wt_array) * r/sqrt_g_tilde(r; a, M, ϵ_3, α_13, α_22, α_52)
end

# use this to see where things break down / if they break down
function plot_derivative_test(r_range, res)
    lzs = res[:, 2]
    # naive derivatives
    derivs = ∇r(lzs, r_range)

    plot(r_range, lzs, label = "Lz", c = :black, legend = false)
    plot!(r_range, derivs, label = "derivative", c = :black, style = :dot, lw = 1.2)
end

# test plot
function f_approx_test(r_range, res; a=0.97, M=1.0, ϵ_3=0, α_13=0, α_22=0, α_52=0)
    wts = res[:, 4]
    Lzs = res[:, 2]

    ∇rwt_array = ∇r(wts, r_range)
    ∇rLz_array = ∇r(Lzs, r_range)

    fs = map(2:length(r_range)) do i
        @views f_approx(Lzs[i], ∇rwt_array[i], ∇rLz_array[1:i], wts[1:i], r_range[i]; a=a, M=M, ϵ_3=ϵ_3, α_13=α_13, α_22=α_22, α_52=α_52)
    end

    # we rescale here for a better comparison
    fs = fs ./ maximum(fs)

    # we lose a point, so select all but last `r_range`
    # since our derivative function needs at least 2 points  work
    # plot(r_range[2:end], fs, c = :black, xlabel = "r", ylabel = "f", label = "simulated f")
end

# this works up until about 0.97 then, the orbit search
# starts to break down

function init(;m, a=0.97, M=1.0, disc_radius=50.0, ϵ_3=0, α_13=0, α_22=0, α_52=0)
    # m = BoyerLindquist(M = M, a = a)
    # m = JohannsenAD(M=1.0, a=a)
    u = @SVector [0.0, 13.0, deg2rad(90), 0.0]

    # making the steps too big over-estimates the gradient
    r_range = CarterBoyerLindquist.rms(m.M, m.a):0.01:disc_radius
    res = find_orbit_range(;m, r_range = r_range, a = m.a)

    # creating an interpolated function of f

    fs = f_approx_test(r_range, res; a=a, M=M, ϵ_3=ϵ_3, α_13=α_13, α_22=α_22, α_52=α_52)
    return fs, r_range
end

function f_approx_func(r, fs, r_range; a=0.97, M=1.0)

    r_range_short = collect(r_range)
    pop!(r_range_short)
    interp_linear_f = LinearInterpolation(r_range_short, fs,extrapolation_bc=Line())
    interp_linear_f(r)
end

G = 6.67e-11
c = 3e8
L_☼ = 3.8e26
M_☼ = 1.99e30
σ_SB = 5.67e-8
η = 0.1

function diss_approx(mdot, r, a_star, M, fs, r_range)
    diss_approx = ((c^6) / (G^2)) * mdot * f_approx_func(r, fs, r_range; a=a_star, M=M) / (4 * π * r)
end

function temperature_approx(r, a_star, M, fs, r_range)
    m_dot = AccretionFormulae.mdot(M)
    temperature_approx = (diss_approx(m_dot, r, a_star, M, fs, r_range) / σ_SB)^(1 / 4)
end

function observed_temperature_approx(r, a_star, M, g, fs, r_range)
    T = temperature_approx(r, a_star, M, fs, r_range)
    observed_temperature_approx = g * T
end


using ContinuumArrays, FillArrays, InfiniteArrays, Plots

T = Chebyshev()
C = Ultraspherical(2)
D = Derivative(axes(T,1))
A = (C\(D^2*T))+100(C\T)

n = 100
c = [T[[-1,1],1:n]; A[1:n-2,1:n]] \ [1;2;zeros(n-2)]
u = T*Vcat(c,Zeros(∞))

xx = range(-1,1;length=1000)
plot(xx,u[xx])


# Advent of Code 2020 - day 18

list = readlines("day18_input.txt")

function reduce_left_first(line)
    while occursin("(", line)
        part = match(r".*(\([^\)]*\)).*", line)[1]
        line = replace(line, part => reduce_left_first(part[2:end-1]))
    end
    while occursin("+", line) || occursin("*", line)
        parts = match(r"(\d+\s[\+|\*]\s\d+)(.*)", line)
        line = string(eval(Meta.parse(parts[1]))) * parts[2]
    end
    return parse(Int, line)
end

function reduce_plus_first(line)
    while occursin("(", line)
        part = match(r".*(\([^\)]*\)).*", line)[1]
        line = replace(line, part => reduce_plus_first(part[2:end-1]))
    end
    while occursin("+", line)
        parts = match(r"([^\+]*\s|)(\d+\s\+\s\d+)(.*)", line)
        line = parts[1] * string(eval(Meta.parse(parts[2]))) * parts[3]
    end
    return eval(Meta.parse(line))
end

function evaluate_list_left_first(list)
    total = 0
    for i in 1:length(list)
        total += reduce_left_first(list[i])
    end
    return total
end

function evaluate_list_plus_first(list)
    total = 0
    for i in 1:length(list)
        total += reduce_plus_first(list[i])
    end
    return total
end

@show reduce_left_first("1 + 2 * 3 + 4 * 5 + 6") == 71
@show reduce_left_first("1 + (2 * 3) + (4 * (5 + 6))") == 51
@show reduce_left_first("2 * 3 + (4 * 5)") == 26
@show reduce_left_first("5 + (8 * 3 + 9 + 3 * 4 * 3)") == 437
@show reduce_left_first("5 * 9 * (7 * 3 * 3 + 9 * 3 + (8 + 6 * 4))") == 12240
@show reduce_left_first("((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2") == 13632

@show evaluate_list_left_first(list) == 1402255785165

@show reduce_plus_first("1 + 2 * 3 + 4 * 5 + 6") == 231
@show reduce_plus_first("1 + (2 * 3) + (4 * (5 + 6))") == 51
@show reduce_plus_first("2 * 3 + (4 * 5)") == 46
@show reduce_plus_first("5 + (8 * 3 + 9 + 3 * 4 * 3)") == 1445
@show reduce_plus_first("5 * 9 * (7 * 3 * 3 + 9 * 3 + (8 + 6 * 4))") == 669060
@show reduce_plus_first("((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2") == 23340

@show evaluate_list_plus_first(list) == 119224703255966

nothing


immutable SquaredError <: Loss
end

@inline function loss_value(loss::SquaredError, p::Real, y::Real)
    return (p - y) * (p - y) # == (p - y)^2
end

Base.show(io::IO, loss::SquaredError) = print(io, "Squared Error")

@inline function loss_derivative(loss::SquaredError, p::Real, y::Real)
    return 2 * (p - y) * (p * (1 - p))
end


module effective_interaction_rotated

# Optimized functions for calculation of Omega and Gamma

function Omega(a, θ)
    ξ = 2*pi*a
    cosθpow2 = cos(θ)^2
    cosξdivξ = cos(ξ)/ξ
    return 3.0/4*((1.0-3.0*cosθpow2) * (sin(ξ)+cosξdivξ)/ξ^2 - (1.0-cosθpow2)*cosξdivξ)
end

function Gamma(a, θ)
    ξ = 2*pi*a
    cosθpow2 = cos(θ)^2
    if ξ<1e-3
        return 3/2*(2/3 + 1/15*ξ^2 *(-2 + cosθpow2) + ξ^4*(1/140-cosθpow2/210))
    else
        sinξ = sin(ξ)/ξ
        return 3/2*((1-cosθpow2)*sinξ + (1-3*cosθpow2) * (cos(ξ)-sinξ)/ξ^2)
    end
end

function OmegaGamma(a, cosθpow2)
    cosξ = cospi(2*a)
    sinξ = sinpi(2*a)
    omega = 0.75*((1.0-3.0*cosθpow2) * (sinξ+cosξ/(2*pi*a))/(2*pi*a)^2 - (1.0-cosθpow2)*cosξ/(2*pi*a))
    gamma = 1.5*((1-cosθpow2)*sinξ/(2*pi*a) + (1-3*cosθpow2) * (cosξ-sinξ/(2*pi*a))/(2*pi*a)^2)
    return omega, gamma
end

function OmegaGamma_orthogonal(a)
    cosξ = cospi(2*a)
    sinξ = sinpi(2*a)
    omega = 0.75*((sinξ+cosξ/(2*pi*a))/(2*pi*a)^2 - cosξ/(2*pi*a))
    gamma = 1.5*(sinξ/(2*pi*a) + (cosξ-sinξ/(2*pi*a))/(2*pi*a)^2)
    return omega, gamma
end

function Omega_orthogonal(a)
    ξ = 2*pi*a
    cosξdξ = cos(ξ)/ξ
    return 3.0/4.0*(-cosξdξ + (sin(ξ)+cosξdξ)/ξ^2)
end

function Gamma_orthogonal(a)
    ξ = 2*a
    sincξ = sinc(ξ)
    return 3.0/2.0*(sincξ + (cospi(ξ)-sincξ)/(pi*ξ)^2)
end

function map_effectiveinteraction(a::Vector{T}, f::Function) where T<:Real
    omega_list = float(T)[]
    gamma_list = float(T)[]
    for ai=a
        omega, gamma = f(ai, e_dipole)
        push!(omega_list, omega)
        push!(gamma_list, gamma)
    end
    return omega_list, gamma_list
end

# Finite symmetric systems

# function triangle_orthogonal(a::Float64)
#     omega_eff = 2*Omega_orthogonal(a)
#     gamma_eff = 2*Gamma_orthogonal(a)
#     return omega_eff, gamma_eff
# end

"""
    effective_interaction_rotated.square_orthogonal(a, Nδ)

Effective Omega and Gamma for a square.

The polarization axis is orthogonal to the square plane.

# Arguments
* `a`: Edge length.
* `Nδ`: Phase shift (Number of atoms in 2π).
"""
function square_orthogonal(a::Real, Nδ::Int)
    @assert 0 <= Nδ < 4
    δ = 2.0*pi*Nδ/4
    Ωcos = 2*Omega_orthogonal(a)*cos(δ) + Omega_orthogonal(sqrt(2.)*a)*cos(2*δ)
    Γcos = 2*Gamma_orthogonal(a)*cos(δ) + Gamma_orthogonal(sqrt(2.)*a)*cos(2*δ)
    return Ωcos, Γcos
end

"""
    effective_interaction_rotated.cube_orthogonal(a, dϕ)

Effective Omega and Gamma for a cube.

The polarization axis is orthogonal to the xy faces.

# Arguments
* `a`: edge length.
* `dϕ`: Phase shift.
"""
function cube_orthogonal(a::Real, dϕ)
    omega_eff, gamma_eff = square_orthogonal(a, 0)
    sqrt2 = sqrt(2.)
    sqrt3 = sqrt(3.)
    Θdiag = atan(sqrt2, 1.)
    omega_eff += (Omega(a, 0.) + 2*Omega(sqrt(2.)*a, pi/4.) + Omega(sqrt3*a, Θdiag))*cos(dϕ)
    gamma_eff += (Gamma(a, 0.) + 2*Gamma(sqrt(2.)*a, pi/4.) + Gamma(sqrt3*a, Θdiag))*cos(dϕ)
    return omega_eff, gamma_eff
end

"""
    effective_interaction_rotated.chain_orthogonal(a, N, dϕ)

Effective Omega and Gamma for an infinite chain.

The polarization axis is orthogonal to the chain and the calculation is
done by adding N spins left and N spins right of a central spin.

# Arguments
* `a`: Spin-spin distance.
* `N`: Number of included spins.
* `dϕ`: Phase shift between neighboring spins.
"""
function chain_orthogonal(a::T, N::Int, dϕ::Real) where T<:Real
    omega_eff::float(T) = 0.
    gamma_eff::float(T) = 0.
    d::float(T) = 0.
    ϕ::float(T) = 0.
    for j=1:N
        d += a
        ϕ += dϕ
        omega_eff += Omega_orthogonal(d)*cos(ϕ)
        gamma_eff += Gamma_orthogonal(d)*cos(ϕ)
    end
    return 2*omega_eff, 2*gamma_eff
end

end # module


@info "Starting..."
using LinearAlgebra, StaticArrays, DifferentialEquations, Makie, ODEInterfaceDiffEq
using BSON: @save, @load
include("system.jl")
include("util.jl")
include("plotting.jl")

@info "Defining ODE system..."
N = M = 3;
p = (k_m=0.2, l₀ = 1.0)
t = (0.0, 60.0)

u₀ = uniform_grid(M, N, p[:l₀])
for i in 1:N
    u₀[3,i,1] += 0.3*p[:l₀]
    u₀[3,i,M] -= 0.3*p[:l₀]
end

@info "Displaying initial positions."
display(plot_lindo(u₀))

@info "Integrating ODE system..."
prob = ODEProblem(rigid!, u₀, t, p)
@time sol = solve(prob, progress=true, force_dtmin=true, DifferentialEquations.TRBDF2())

@save "$(pwd())/data/$(N)x$(M).bson" sol
@info "Saved solution to $(pwd())/data/$(N)x$(M).bson"

@info "Rendering..."
render(sol)
@info "Done."


#!/usr/bin/julia
#
# Helper functions for constructing the Hamiltonian.
#
export sec_dv, energy, ke, pe;

"""
    sec_dv(grid::UniformGrid{S,1}) where S

Construct matrix representation of discretised second order partial derivative
acting on position space.
"""
function sec_dv(grid::UniformGrid{S,1}) where S
  rows = length(grid);
  # Second order method.
  #return LinearAlgebra.SymTridiagonal(repeat([-2.0/grid.Δ^2], rows),
  #                      repeat([1.0/grid.Δ^2], rows-1));

  # Second order method explicitly constructed.
  M = zeros(Float64, rows, rows);
  for i ∈ 1:rows
    # Set incomplete end values to zero to prevent introduction of spurious
    # boundary terms
    if i > 1 && i <= rows-1
      M[i,i-1] = -1;
    #end
    M[i,i] = 2;
    #if i <= rows-1
      M[i,i+1] = -1;
    end
  end
  #M *= -Rational(1/grid.Δ^2);
  M *= -(1/grid.Δ^2);
  return M;

  #=
  # Infinite order expression for 2nd derivative as per
  #J. Chem. Phys. 96, 1982 (1992) Daniel T. Colbert and William H. Miller
  # https://aip-scitation-org.ezphost.dur.ac.uk/doi/10.1063/1.462100
  M = Matrix{Float64}(undef, rows, rows);
  for i ∈ 1:rows
    for j ∈ 1:rows
      if (i == j)
        M[i,i] = π^2/3;
      else
        M[i,j] = (-1)^(i-j) * 2/(i-j)^2;
      end
    end
  end
  M *= -1/grid.Δ^2;
  return M;=#

  # 5 point function
  #M = Matrix{Rational}(zero(Rational), rows, rows);
  #M = zeros(big(Rational), rows, rows);
  #=
  M = zeros(Float64, rows, rows);
  for i ∈ 3:rows-3
    M[i,i-2] = 1/12;
    M[i,i+2] = 1/12;
    M[i,i-1] = -4/3;
    M[i,i+1] = -4/3;
    M[i,i] = 5/2;
    #=
    if i > 1
      M[i,i-1] = -4//3;
      if (i > 2) M[i,i-2] = 1//12; end
    end
    M[i,i] = 5//2;
    if i <= rows-1
      M[i,i+1] = -4//3;
      if (i <= rows-2) M[i,i+2] = 1//12; end
    end=#
  end
  #M *= -Rational(1/grid.Δ^2);
  M *= -(1/grid.Δ^2);
  return M;=#
end

"""
    energy(ψ::WaveFunction{T,1,S,G}, H::SeparableHamiltonian) where
            {T,S,G<:UniformGrid}

Calculate energy of the given wave function with the kinetic energy calculated
in momentum space (uses a fast Fourier transform if wave function is in
configuration space).

See also: [`ke`](@ref), [`pe`](@ref)
"""
function energy(ψ::WaveFunction{T,1,S,G}, H::SeparableHamiltonian) where
            {T,S,G<:UniformGrid}
  return ke(ψ, H) + pe(ψ,H);
end
#In the future this function may implement a different underlying energy calculation.
function energy(ψ::WaveFunction{T,1,S,G}, H::SeparableHamiltonian) where
            {T,S<:ConfigurationSpace,G<:UniformGrid}
  return ke(ψ, H) + pe(ψ, H);
end

"""
    energy(ψ::WaveFunction{T,1,S,G}, H::SeparableHamiltonian,
           g::UniformGrid{S,1}) where {T,S<:ConfigurationSpace,G<:UniformGrid}

Calculate energy of the given wave function with the kinetic energy determined
using matrix multiplication in configuration space.

See also: [`ke`](@ref), [`pe`](@ref)
"""
function energy(ψ::WaveFunction{T,1,S,G}, H::SeparableHamiltonian,
                g::UniformGrid{S,1}) where {T,S<:ConfigurationSpace,G<:UniformGrid}
  return ke(ψ, g) + pe(ψ, H);
end
# TODO: Implement method of finding energy for non-separable Hamiltonians.

"""
    ke(ψ::WaveFunction{T,1,S,G}, H::SeparableHamiltonian) where
            {T,S<:MomentumSpace,G<:UniformGrid}

Compute kinetic energy fully in momentum space.

Uses a one-dimensional wave function defined on a `UniformGrid`.
"""
function ke(ψ::WaveFunction{T,1,S,G}, H::SeparableHamiltonian) where
            {T,S<:MomentumSpace,G<:UniformGrid}
  return real(sum(conj(ψ) .* H.T .* ψ)) * ψ.Δ;
end

"""
    ke(ψ::WaveFunction{T,1,S,G}, g::UniformGrid{S,1}) where
      {T,S<:ConfigurationSpace,G<:UniformGrid}

Compute kinetic energy in configuration space using E_k = \frac{1}{2} p^2
where the wave function is defined on a one-dimensional `UniformGrid`.
"""
function ke(ψ::WaveFunction{T,1,S,G}, g::UniformGrid{S,1}) where
            {T,S<:ConfigurationSpace,G<:UniformGrid}
  return real(ψ' * -0.5sec_dv(g) * ψ) * ψ.Δ;
end

"""
    ke(ψ::WaveFunction{T,1,S,G}, H::SeparableHamiltonian) where
      {T,S<:ConfigurationSpace,G<:UniformGrid}

Compute kinetic energy using a Fourier transform to momentum space and then
using the kinetic component of Hamiltonian.

The wave function will be defined on a one-dimensional `UniformGrid`.
"""
function ke(ψ::WaveFunction{T,1,S,G}, H::SeparableHamiltonian) where
            {T,S<:ConfigurationSpace,G<:UniformGrid}
  return real(sum(conj(ψ) .* FFTW.ifft(H.T .* FFTW.fft(ψ)))) * ψ.Δ;
end
function ke(ψ::WaveFunction{T,1,S,G}, H::SeparableHamiltonian{Q}) where
            {T<:Unitful.Quantity,S<:ConfigurationSpace,G<:UniformGrid,Q<:Unitful.Quantity}
  return ustrip(real(sum(conj(ψ) .* FFTW.ifft(ustrip(H.T) .* FFTW.fft(ψ)))) * ψ.Δ) * unit(Q);
end

"""
    pe(ψ::WaveFunction{T,1,S,G}, H::SeparableHamiltonian) where
            {T,S<:ConfigurationSpace,G<:UniformGrid}

Calculate potential energy of wave function.

Wave function is in configuration space and defined on one-dimensional uniform
grid.
"""
function pe(ψ::WaveFunction{T,1,S,G}, H::SeparableHamiltonian) where
            {T,S<:ConfigurationSpace,G<:UniformGrid}
  # Calculate energy due to spatial potential.
  Ex = real(sum(conj(ψ) .* H.V .* ψ)) * ψ.Δ;
end

# TODO: Implement determination of P.E. for momentum space wave function.


# the julia iteration function

function julia(z, maxiter::Int64)
    c=-0.8+0.156im
    for n = 1:maxiter
        if abs(z) > 2
            return n-1
        end
        z = z^2 + c
    end
    return maxiter
end


# A Julia translation of the log function of Yeppp:
#  https://bitbucket.org/MDukhan/yeppp/src/b8db687e912bbd7e2a26dd93c348ef2d7a5febdc/library/headers/yepBuiltin.h?at=default#cl-1128

# Available under following terms:
#
#                          Yeppp! library header
#
# This file is part of Yeppp! library and licensed under the New BSD license.
#
# Copyright (C) 2010-2012 Marat Dukhan
# Copyright (C) 2012-2013 Georgia Institute of Technology
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of the Georgia Institute of Technology nor the
#       names of its contributors may be used to endorse or promote products
#       derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

function log_yeppp(x::Float64)
    xu = reinterpret(Uint64,x)
    if xu >= 0x0010_0000_0000_0000
        k = int(xu >> 52) & 0x07ff - 1023
        s = reinterpret(Float64,(xu & 0x000f_ffff_ffff_ffff) | 0x3ff0_0000_0000_0000)
    else
        offset = leading_zeros(xu)-11
        k = -1022 - offset
        s = reinterpret(Float64,((xu << offset) & 0x000f_ffff_ffff_ffff)
                        | 0x3ff0_0000_0000_0000)
    end
    if s >= 0x1.6A09E667F3BCDp+0 # sqrt(2)
        k += 1
        s *= 0.5
    end
    fk = Float64(k)

    t = s-1.0
    pt = @horner(t,
            -0x1.FFFFFFFFFFFF2p-2,
            0x1.5555555555103p-2,
            -0x1.00000000013C7p-2,
            0x1.9999999A43E4Fp-3,
            -0x1.55555554A6A2Bp-3,
            0x1.249248DAE4B2Ap-3,
            -0x1.FFFFFFBD8606Dp-4,
            0x1.C71C90DB06248p-4,
            -0x1.9999C5BE751E3p-4,
            0x1.745980F3FB889p-4,
            -0x1.554D5ACD502ABp-4,
            0x1.3B4ED39194B87p-4,
            -0x1.25480A82633AFp-4,
            0x1.0F23916A44515p-4,
            -0x1.EED2E2BB64B2Ep-5,
            0x1.EA17E14773369p-5,
            -0x1.1654764F478ECp-4,
            0x1.0266CD08DB2F2p-4,
            -0x1.CC4EC078138E3p-6)

    rf = t*(t*pt)+t
    f = fk*0x1.62E42FEFA3800p-1 + (fk*0x1.EF35793C76730p-45 + rf)
    if xu < 0
        NaN
    elseif xu == 0
        -Inf
    elseif !(x < Inf)
        x
    else
        f
    end
end

@vectorize_1arg Real log_yeppp


"""
    BernoulliLikelihood(l=logistic)

Bernoulli likelihood is to be used if we assume that the
uncertainity associated with the data follows a Bernoulli distribution.
The link `l` needs to transform the input `f` to the domain [0, 1]

```math
    p(y|f) = \\operatorname{Bernoulli}(y | l(f))
```
On calling, this would return a Bernoulli distribution with `l(f)` probability of `true`.
"""
struct BernoulliLikelihood{Tl<:AbstractLink} <: AbstractLikelihood
    invlink::Tl
end

BernoulliLikelihood(l=logistic) = BernoulliLikelihood(link(l))

(l::BernoulliLikelihood)(f::Real) = Bernoulli(l.invlink(f))

(l::BernoulliLikelihood)(fs::AbstractVector{<:Real}) = Product(map(l, fs))


push!(LOAD_PATH, "Modules")
using Colors
using ColorVectorSpace
using IntervalSets
using Revise
using JuliRay
Base.@irrational ° 0.0174532925199432957692369076848861271344 (big(pi)/big(180))

function Smooth(x::Real)
    f(x) = ifelse(x>0,exp(-1/x),0)
    return f(2x/√3)/(f(2x/√3)+f(2/√3-2x/√3))
end
function Smooth(a::Real, b::Real, x::Real)
    return Smooth((x-a)/(b-a))
end
function Smooth(dict::Dict{A,B} where A <: Real where B <: Any, t::Real)
    T=sort(collect(keys(dict)))
    L=length(T)
    V=[dict[T[i]] for i in 1:L]

    if t ≤ T[1]
        return V[1]
    end
    for i in 2:L
        if t ≤ T[i]
            return V[i-1]+(V[i]-V[i-1])*Smooth(T[i-1],T[i],t)
        end
    end
    return V[end]
end

N=120
for i in 0:N-1
    t04 = 4i/N
    t = Smooth(Dict(0=>0, 1=>π/2, 2=>π, 3=>3π/2, 4=>2π),t04)
    catenoid(u) = [cos(u[2])*cosh(u[1]),sin(u[2])*cosh(u[1]),u[1]]
    helicoid(u) = [cos(u[2])*sinh(u[1]),sin(u[2])*sinh(u[1]),u[2]]

    𝒑(u) = cos(t)*catenoid(u)+sin(t)*helicoid(u)
    D1 = -π/2..π/2
    D2 = -π..π
    mesh = (36,36)

    nn = 18

    col(u) = Float64(xor(sin(nn*u[1])<0,sin(nn*u[2])>0))*RGB(0,1,0.5) + Float64(xor(sin(nn*u[1])>0,sin(nn*u[2])>0))*RGB(0,0.5,1)
    obj = ColoredParametricSurface(𝒑, D1, D2, mesh=mesh, color=col)
    lng = 30°+180°*i/N
    lat = 20°
    lookat = [0,0,1]
    zoom = 0.1
    render(obj,camera=LngLatCamera(lng=lng,lat=lat,pers=0.2,zoom=zoom,lookat=lookat),name="CatenoidHelicoid",index=i+1)
end


f_qbftest(t) = 1/(2+cos(2*pi*t))^2

function test_qbf()
    N = 32
    for degree in 1:4
        basis = BSplineTranslatesBasis(N, degree, 0, 1)
        coef_direct = [DomainIntegrals.integral(t -> f_qbftest(t)*BasisFunctions.unsafe_eval_element(basis, i, t), FrameFun.support(basis, i)) for i in 1:length(basis)]
        for oversamplingfactor in 1:5
            A = IntegralEquations2D.qbf_operator(basis, degree, oversamplingfactor)
            coef = A * f_qbftest
            @test norm(coef-coef_direct) < 1e-4
        end
    end

    # Do a BigFloat test
    N = 500
    degree = 11
    oversamplingfactor = 6
    basis = BSplineTranslatesBasis(N, degree, big(0.), big(1.))
    A = IntegralEquations2D.qbf_operator(basis, degree, oversamplingfactor)
    coef = A * f_qbftest
    # c10 = integral(t -> f_qbftest(t)*BasisFunctions.unsafe_eval_element(basis, 10, t), FrameFun.support(basis, 10))
    c10 = BigFloat("0.004990536452129881961287340961088526093027349985606")
    @test abs(coef[10]-c10) < 1e-45
end


# this is a little Monte Carlo exercise that illustrates that
# the OLS estimator is efficient. It compares the OLS estimator
# to an estimator that is the average of the OLS estimators using
# 3 subsamples
using Plots
function main()
pyplot()
reps = 10000 # number of Monte Carlo reps.
n = 21 # sample size
sig = 2.0  # st. dev. of errors

x = [ones(n,1) randn(n,1)]  # x is fixed over repeated samples
beta = [1.0, 2.0] # true beta
PopRegLine = x*beta

betas = zeros(reps,2) # holder for results

for i = 1:reps
	# generate dep var
	e = sig*randn(n)
	y = PopRegLine + e
	# the OLS estimator
	betas[i,1] = (y\x)[2]
	# the average of split sample estimator
	y1 = y[1:7,:]
	y2 = y[8:14,:]
	y3 = y[15:21,:]
	x1 = x[1:7,:]
	x2 = x[8:14,:]
	x3 = x[15:21,:]
	beta_ss1 = y1\x1
	beta_ss2 = y2\x2
	beta_ss3 = y3\x3
	beta_ss = (beta_ss1 + beta_ss3 + beta_ss3)/3.0 # average the 3 estimators
	betas[i,2] = beta_ss[2]
end

histogram(betas[:,1], label="", nbins=30, show=true)
#savefig("efficiency1.svg")
histogram(betas[:,2], label="", nbins=30, reuse=false, show=true)
#savefig("efficiency2.svg")
return
end
main()


using DiffEqSensitivity, Test

A = reshape([1,0,2,3],2,2)
function f_morris(p)
    A*p
end

function linear_batch(X)
    A= 7
    B= 0.1
    @. A*X[1,:]+B*X[2,:]
end

function ishi(X)
    A= 7
    B= 0.1
    sin(X[1]) + A*sin(X[2])^2+ B*X[3]^4 *sin(X[1])
end

function ishi_linear(X)
    A= 7
    B= 0.1
    [sin(X[1]) + A*sin(X[2])^2+ B*X[3]^4 *sin(X[1]),A*X[1]+B*X[2]]
end

function ishi_linear_batch(X)
    A= 7
    B= 0.1
    X1 = @. sin(X[1,:]) + A*sin(X[2,:])^2+ B*X[3,:]^4 *sin(X[1,:])
    X2 = @. A*X[1,:]+B*X[2,:]
    vcat(X1',X2')
end

lb = -ones(4)*π
ub = ones(4)*π

m = gsa(f_morris,Morris(p_steps=[10,10],total_num_trajectory=1000,num_trajectory=150),[[1,5],[1,5]])
@test m.means[:,1] ≈ A[:,1] atol=1e-12
@test m.means[:,2] ≈ A[:,2] atol=1e-12
@test m.variances ≈ reshape([0,0,0,0],2,2) atol=1e-12

m = gsa(f_morris,Morris(p_steps=[10,10],relative_scale=true,total_num_trajectory=1000,num_trajectory=150),[[1,5],[1,5]])
@test m.means[2,1] ≈ 0 atol=1e-12
@test m.variances[2,1] ≈ 0 atol=1e-12
@test m.means[1,2] < m.means[2,2]

m = gsa(linear_batch,Morris(p_steps=[10,10],num_trajectory=10000),[[1,5],[1,5]],batch=true)
@test m.means ≈ [7.0  0.1] atol = 1e-2

m = gsa(ishi, Morris(num_trajectory=100000), [[lb[i],ub[i]] for i in 1:4])
@test m.means[1,:] ≈ [2.25341,4.40246,2.5049,0.0] atol = 5e-2

m = gsa(ishi_linear,Morris(num_trajectory=100000),[[lb[i],ub[i]] for i in 1:4])
@test m.means[1,:] ≈ [2.25341,4.40246,2.5049,0.0] atol = 5e-2
@test m.means[2,:] ≈ [7.0,0.1,0.0,0.0] rtol = 1e-12

m = gsa(ishi_linear_batch,Morris(num_trajectory=100000),[[lb[i],ub[i]] for i in 1:4],batch=true)
@test m.means[1,:] ≈ [2.25341,4.40246,2.5049,0.0] atol = 5e-2
@test m.means[2,:] ≈ [7.0,0.1,0.0,0.0] rtol = 1e-12


using SymPy.SpecialFuncs
using SymPy: Sym, sqrt, conjugate, symbols, PI, simplify,
             expand_func, rewrite, N, gamma
if VERSION >= v"0.5.0-dev+7720"
    using Base.Test
else
    using BaseTestNext
    const Test = BaseTestNext
end

@testset "SpecialFuns" begin
    a, b, x, y = symbols("a, b, x, y")
    n, m, θ, ϕ = symbols("n, m, theta, phi")
    ν = symbols("nu", integer=true)

    @test fresnels(Sym(0)) == 0
    @test fresnels(Sym(oo)) == Sym(1)/2
    @test diff(fresnels(x), x) == sin(PI*x^2/2)
    #@test evalf(fresnels(Sym(2)), 30) == Sym(parse(BigFloat, "0.343415678363698242195300815958"))

    @test  diff(fresnelc(x), x) == cos(PI*x^2/2)

    @test Ei(Sym(-1)) == Ei(exp(1im*PI))
    @test diff(Ei(x), x) == exp(x)/x
    @test diff(Si(x), x) == sin(x)/x
    @test diff(Ci(x), x) == cos(x)/x

    @test airyai(Sym(0)) == 3^(Sym(1)/3)/(3*gamma(Sym(2)/3))
    @test airybi(Sym(0)) == 3^(Sym(5)/6)/(3*gamma(Sym(2)/3))

    @test jacobi(0, a, b, x) == 1
    @test jacobi(1, a, b, x) == a/2 - b/2 + x*(a/2 + b/2 + 1)
    @test jacobi(n, 0, 0, x) == legendre(n, x)
    @test jacobi(n, a, b, -x) == (-1)^n*jacobi(n, b, a, x)
    @test conjugate(jacobi(n, a, b, x)) == jacobi(n, conjugate(a), conjugate(b), conjugate(x))

    @test gegenbauer(0, a, x) == 1
    @test gegenbauer(1, a, x) == 2a*x
    @test gegenbauer(2, a, x) == -a + x^2*(2a^2 + 2a)
    @test conjugate(gegenbauer(n, a, x)) == gegenbauer(n, conjugate(a), conjugate(x))
    @test diff(gegenbauer(n, a, x), x) == 2a*gegenbauer(n-1, a+1, x)

    @test chebyshevt(0, x) == 1
    @test chebyshevt(1, x) == x
    @test chebyshevt(2, x) == 2x^2-1

    @test chebyshevu(0, x) == 1
    @test chebyshevu(1, x) == 2x
    @test chebyshevu(2, x) == 4x^2 - 1
    @test chebyshevu(n, 0) == cos(π*n/2)
    @test chebyshevu(n, 1) == n + 1

    @test legendre(0, x) == 1
    @test legendre(1, x) == x
    @test legendre(2, x) == 3x^2/2 - 1/2
    @test legendre(n, x) == legendre(n, x)
    @test diff(legendre(n,x), x) == n*(x*legendre(n, x) - legendre(n - 1, x))/(x^2 - 1)

    @test assoc_legendre(0,0, x) == 1
    @test assoc_legendre(1,0, x) == x
    @test assoc_legendre(1,1, x) == -sqrt(-x^2 + 1)
    #@test assoc_legendre(n,m,x) == assoc_legendre(n, m, x)

    @test hermite(0, x) == 1
    @test hermite(1, x) == 2x
    @test hermite(2, x) == 4x^2 - 2
    @test hermite(n, x) == hermite(n, x)
    @test diff(hermite(n,x), x) == 2n*hermite(n - 1, x)
    @test hermite(n, -x) == (-1)^n*hermite(n, x)

    @test laguerre(0, x) == 1
    @test laguerre(1, x) == -x + 1
    @test laguerre(2, x) == x^2/2 - 2x + 1
    @test laguerre(3, x) == -x^3/6 + 3x^2/2 - 3x + 1
    @test diff(laguerre(n, x), x) == -assoc_laguerre(n - 1, 1, x)

    @test assoc_laguerre(0, a, x) == 1
    @test assoc_laguerre(1, a, x) == a - x + 1
    @test assoc_laguerre(2, a, x) == a^2/2 + 3a/2 + x^2/2 + x*(-a - 2) + 1
    @test assoc_laguerre(n, 0, x) == laguerre(n, x)
    @test diff(assoc_laguerre(n, a, x), x) == -assoc_laguerre(n - 1, a + 1, x)

    @test Ynm(n, m, θ, ϕ) == Ynm(n, m, θ, ϕ)
    @test Ynm(n, -m, θ, ϕ) == (-1)^m*exp(-2im*m*ϕ)*Ynm(n, m, θ, ϕ)
    @test Ynm(n, m, -θ, ϕ) == Ynm(n, m, θ, ϕ)
    @test Ynm(n, m, θ, -ϕ) == exp(-2im*m*ϕ)*Ynm(n, m, θ, ϕ)
    @test expand(simplify(Ynm(0, 0, θ, ϕ)), func=true) == 1/(2*sqrt(PI))

    @test Ynm_c(n, m, θ, ϕ) == conjugate(Ynm(n, m, θ, ϕ))

    @test diff(hankel1(n, x), x) == hankel1(n - 1, x)/2 - hankel1(n + 1, x)/2
    @test diff(hankel2(n, x), x) == hankel2(n - 1, x)/2 - hankel2(n + 1, x)/2

    @test expand_func(jn(0, x)) == sin(x)/x
    @test expand_func(jn(1, x)) == sin(x)/x^2 - cos(x)/x
    @test rewrite(jn(ν, x), "besselj") == sqrt(2PI/x)*besselj(ν + Sym(1)/2, x)/2
    VERSION < v"0.6.0-dev" && @test rewrite(jn(ν, x), "bessely") == (-1)^ν*sqrt(2PI/x)*bessely(-ν - Sym(1)/2, x)/2
    u = N(jn(2, 5.2+0.3im), 20)
    @test norm(real(u) - 0.099419756723640344491) <= 1e-15 && norm(imag(u) + 0.054525080242173562897) <= 1e-15

    @test expand_func(yn(0, x)) == -cos(x)/x
    @test expand_func(yn(1, x)) == -cos(x)/x^2-sin(x)/x
    VERSION < v"0.6.0-dev" && @test rewrite(yn(ν, x), "besselj") == (-1)^(ν + 1)*sqrt(2PI/x)*besselj(-ν - Sym(1)/2, x)/2
    @test rewrite(yn(ν, x), "bessely") == sqrt(2PI/x)*bessely(ν + Sym(1)/2, x)/2
    @test N(yn(2, 5.2+0.3im)) ≈ 0.18525034196069722536 + 0.014895573969924817587im

    # gamma, beta and related functions
    @test gamma(Sym(1)) == 1
    @test gamma(Sym(3)/2) == sqrt(PI)/2

    @test diff(polygamma(Sym(0), x), x) == polygamma(Sym(1), x)
    @test diff(polygamma(Sym(0), x), x, 2) == polygamma(Sym(2), x)

    @test diff(beta(x, y), x) == (polygamma(Sym(0), x) - polygamma(Sym(0), x + y)) * beta(x, y)
    @test diff(beta(x, y), y) == (polygamma(Sym(0), y) - polygamma(Sym(0), x + y)) * beta(x, y)

    # test numerical consistency with Julia functions
    @test N(gamma(Sym(4.1))) ≈ gamma(4.1)
    @test N(polygamma(Sym(2), Sym(3.2))) ≈ polygamma(2, 3.2)
    VERSION >= v"0.5.0" && @test N(beta(Sym(1)+1im, Sym(1)+1im)) ≈ beta(1.0+1im, 1.0+1im)

    # Elliptic-type functions

    @test elliptic_k(Sym(0)) == PI/2
    @test N(elliptic_k(Sym(1.0 + im))) ≈ 1.50923695405127 + 0.625146415202697im

    @test N(elliptic_f(Sym(3.0 + im/2), Sym(1.0 + im))) ≈ 2.909449841483 + 1.74720545502474im

    @test elliptic_e(Sym(0)) == PI/2
    @test N(elliptic_e(Sym(2.0 - im))) ≈ 0.991052601328069 + 0.81879421395609im

    @test elliptic_pi(Sym(0), Sym(0)) == PI/2
    @test N(elliptic_pi(Sym(1.0 - im/3), Sym(2.0 + im))) ≈ 3.29136443417283 + 0.32555634906645im

    # Bessel-type functions
    @test diff(besselj(n, x), x) == (besselj(n - 1, x) - besselj(n + 1, x))/2
    @test rewrite(besselj(n, x), "jn") == sqrt(2x/PI)*jn(n - 1/2, x)

    @test diff(bessely(n, x), x) == (bessely(n - 1, x) - bessely(n + 1, x))/2
    @test rewrite(bessely(n, x), "yn") == sqrt(2x/PI)*yn(n - 1/2, x)

    @test diff(besseli(n, x), x) == (besseli(n - 1, x) + besseli(n + 1, x))/2

    @test diff(besselk(n, x), x) == -(besselk(n - 1, x) + besselk(n + 1, x))/2

    # test numerical consistency with Julia functions
    if VERSION < v"0.6.0-dev"
      @test N(besselj(3.2, Sym(1.5))) ≈ besselj(3.2, 1.5)
      @test N(bessely(3.2, Sym(1.5))) ≈ bessely(3.2, 1.5)
      @test N(besseli(3.2, Sym(1.5))) ≈ besseli(3.2, 1.5)
      @test N(besselk(3.2, Sym(1.5))) ≈ besselk(3.2, 1.5)
    end

    @test expand_func(x*hyper([1, 1], [2], -x)) == log(x + 1)
end


using Statistics
N = 10_000         # number of nodes
m = 30             # number of active links per active node
mu = 0.1            # probability I -> S
tmax = 100         # maximum duration of the simulation
Ni = Int(0.1 * N)     # initial number of infects
ActMin = 0.001      # minimum activation probability
ActMax = 1          # maximum activation probability
gamma = -2.22       # exponent of the power law

AgentAct = rand(N, 1);
AgentAct = ((ActMax.^(gamma + 1) - ActMin.^(gamma + 1)) .* AgentAct .+
            ActMin.^(gamma + 1)).^(1.0 / (gamma + 1))
AvgDegree = 2 * m * mean(AgentAct)

NTrials = 5         # number of trials for each value of lambda

tolerance = 1e-15   # termination tolerance
TWindow = 1000      # time window for evaluating steady state
TMinConvergence = 500

# health state of agents S=0 I=1
global AgentState = zeros(N)
AgentState[1:Ni] .= 1.0

# diffferent lambda
# lambda=0.1:0.1:1;
# IAvgLambda=zeros(length(lambda))
# beta=zeros(length(lambda))

lambda = 0.2
termination = 1.0
t = 1;
ITrend = zeros(0)

 for t = 1:tmax
    AdjMatrix = zeros(N, N)
    ActiveIndicesBool = [rand() < agentAct for agentAct in AgentAct]
    indx = [k[1] for k in findall(ActiveIndicesBool)]
    global AgentState=AgentState
    for i in indx
        pick_from = setdiff(1:N, i)
        ContactedAgents = rand(pick_from, m)
        AdjMatrix[i, ContactedAgents] .= 1
        AdjMatrix[ContactedAgents, i] .= 1 # connections are bilateral!
    end

    global AgentStateNext = AgentState

    InfectedIndicesBool = [agentState == 1 for agentState in AgentState]
    InfectedIndices = [k[1] for k in findall(InfectedIndicesBool)]

    for j in InfectedIndices
        ConnectedSusceptibles = Bool.(AdjMatrix[j, :]) .& (.~Bool.(AgentState)) # find individuals connected to an infected that are susceptible
        ConnSuscIndices = findall(ConnectedSusceptibles) # find their indices
        AgentStateNext[ConnSuscIndices] = [rand() < lambda for v in ConnSuscIndices] # infection process S->I on the susceptible and connected;
    end

    InfectedIndices = [k[1] for k in findall(Bool.(AgentStateNext)) ]# find Infected again and
    AgentStateNext[InfectedIndices] = [~(rand() < mu) for v in InfectedIndices] # process I->S (it is independent from the network connections)
    AgentState = AgentStateNext

    push!(ITrend, Int(sum(AgentState)))
    if t > (TMinConvergence + TWindow)
        termination = (abs(mean(ITrend[t - TWindow:t - 1]) - mean(ITrend[t - TWindow + 1:t])) < tolerance);
        if termination
            break
        end
    end

end


using eFEMpart
include("MeshTransform.jl")

point = [rand(),rand()]

p1 = 2*[-1.0,-1.0]
p2 = 2*[1.0,-1.0]
p3 = 2*[1.0,1.0]
p4 = 2*[-1.0,1.0]
quad = [p1,p2,p3,p4]

s,t = reverseQuadMap(quad,point)


using ChangesOfVariables, InverseFunctions, ArraysOfArrays, Statistics
using Optimisers
using Plots

using EuclidianNormalizingFlows: CenterStretch, JohnsonTrafo, HouseholderTrafo
using EuclidianNormalizingFlows: optimize_whitening

f_fwd_true =
    CenterStretch([4.0], [1.0], [0.0]) ∘
    JohnsonTrafo([10.0], [3.5], [10.0], [1.0])

XW = randn(1, 10^5)
stephist(XW[1,:], nbins = 100; normalize = true)

X = f_fwd_true(XW)
stephist!(X[1,:], nbins = 100; normalize = true)

initial_trafo =
    JohnsonTrafo([0.0], [5.0], [0.0], [5.0]) ∘
    inverse(CenterStretch([0.0], [1.0], [0.0])) ∘
    JohnsonTrafo([0.0], [5.0], [0.0], [5.0]) ∘
    inverse(CenterStretch([0.0], [1.0], [0.0]))

optimizer = ADAGrad()

smpls = nestedview(X)
nbatches = 100
nepochs = 10

r = optimize_whitening(smpls, initial_trafo, optimizer, nbatches = nbatches, nepochs = nepochs)

XW2 = r.result(X)
stephist(XW[1,:], nbins = 100; normalize = true)
stephist!(XW2[1,:], nbins = 100; normalize = true)

using EuclidianNormalizingFlows: mvnormal_negll_trafo, mvnormal_negll_trafograd
batches = collect(Iterators.partition(smpls, round(Int, length(smpls)/nbatches)))
plot(r.negll_history, label = "batch negative log-likelihood")
ref_neg_ll = mean([mvnormal_negll_trafo(inverse(f_fwd_true), flatview(b)) for b in batches])
hline!([ref_neg_ll], label = "initial")
ref_neg_ll = mean([mvnormal_negll_trafo(initial_trafo, flatview(b)) for b in batches])
hline!([ref_neg_ll], label = "target")


function CX(ALPHA,EL)
      A=[-0.0990   -0.0480   -0.0220   -0.0400   -0.0830
         -0.0810   -0.0380   -0.0200   -0.0380   -0.0730
         -0.0810   -0.0400   -0.0210   -0.0390   -0.0760
         -0.0630   -0.0210   -0.0040   -0.0250   -0.0720
         -0.0250    0.0160    0.0320    0.0060   -0.0460
          0.0440    0.0830    0.0940    0.0620    0.0120
          0.0970    0.1270    0.1280    0.0870    0.0240
          0.1130    0.1370    0.1300    0.0850    0.0250
          0.1450    0.1620    0.1540    0.1000    0.0430
          0.1670    0.1770    0.1610    0.1100    0.0530
          0.1740    0.1790    0.1550    0.1040    0.0470
          0.1660    0.1670    0.1380    0.0910    0.0400]
      #
      S= 0.2*ALPHA
      K= floor(S)+3
      if (K <=1)
            K = 2
      end
      if (K >=12)
            K = 11
      end
      DA= S - K+3
      L = K + fix(1.1*sign(DA))
      #
      S= EL/12.0
      M= floor(S)+3
      if(M <= 1)
            M = 2
      end
      if(M >=  5)
            M = 4
      end
      DE= S - M+3
      N= M + fix(1.1*sign(DE))
      K=convert(Int64, K)
      L=convert(Int64, L)
      M=convert(Int64, M)
      N=convert(Int64, N)
      #
      T= A[K,M]
      U= A[K,N]
      V= T + abs(DA) * (A[L,M] - T)
      W= U + abs(DA) * (A[L,N] - U)
      return cxv = V + (W-V)  * abs(DE)
end


export
    CoherentState,
    SqueezedState,
    ThermalState,
    SqueezedThermalState

##############
# pure state #
##############

"""
    CoherentState(α::ComplexVec{<:Real}; dim=DIM, rep=StateVector)

Coherent state is defined as the eigenstate of annihilation operator.

* `α`: Eigenvalue of annihilation operator.
* `dim`: Maximum photon number for truncate, default is $DIM.
* `rep`: In which representation, default is `StateVector`.

``\\hat{a} | \\alpha \\rangle = \\alpha | \\alpha \\rangle``

This constructor will construct ``| \\alpha \\rangle = \\hat{D}(\\alpha) | 0 \\rangle``
"""
function CoherentState(T::Type{<:Number}, α::ComplexVec; dim=DIM, rep=StateVector)
    return displace!(VacuumState(T, dim=dim, rep=rep), α)
end

CoherentState(α::ComplexVec; dim=DIM, rep=StateVector) = CoherentState(ComplexF64, α, dim=dim, rep=rep)

"""
    SqueezedState(ξ::ComplexVec{<:Real}; dim=DIM, rep=StateVector)

Squeezed state is defined if its electric field strength for some phases
has a quantum uncertainty smaller than that of a coherent state.

* `ξ`: Squeezing factor
* `dim`: Maximum photon number for truncate, default is $DIM.
* `rep`: In which representation, default is `StateVector`.

This constructor will construct ``| \\xi \\rangle = \\hat{S}(\\xi) | 0 \\rangle``
"""
function SqueezedState(T::Type{<:Number}, ξ::ComplexVec; dim=DIM, rep=StateVector)
    return squeeze!(VacuumState(T, dim=dim, rep=rep), ξ)
end

SqueezedState(ξ::ComplexVec; dim=DIM, rep=StateVector) = SqueezedState(ComplexF64, ξ, dim=dim, rep=rep)

###############
# mixed state #
###############

bose_einstein(n::Number, n̄::Real) = n̄^n / (1 + n̄)^(n+1)

bose_einstein(n̄::Real) = n -> bose_einstein(n, n̄)

"""
    ThermalState(n̄::Real; dim=DIM)

Thermal state is a mixed state with photon number distribution described by Bose-Einstein distribution.

* `n̄`: Average photon number at temperature T.
* `dim`: Maximum photon number for truncate, default is $DIM.
"""
ThermalState(T::Type{<:Number}, n̄::Real; dim=DIM) = StateMatrix(diagm(bose_einstein(n̄).(T.(0:dim-1))), dim)
ThermalState(n̄::Real; dim=DIM) = ThermalState(ComplexF64, n̄, dim=dim)

"""
    SqueezedThermalState(ξ::ComplexVec{<:Real}, n̄::Real; dim=DIM)

Squeezed state is defined if its electric field strength for some phases
has a quantum uncertainty smaller than that of a coherent state.

* `ξ`: Squeezing factor
* `n̄`: Average photon number at temperature T.
* `dim`: Maximum photon number for truncate, default is $DIM.

This constructor will construct ``\\rho = \\hat{S}(\\xi) \\rho_{th} \\hat{S}(\\xi)^{T}``
"""
function SqueezedThermalState(T::Type{<:Number}, ξ::ComplexVec, n̄::Real; dim=DIM)
    return squeeze!(ThermalState(T, n̄, dim=dim), ξ)
end

SqueezedThermalState(ξ::ComplexVec, n̄::Real; dim=DIM) = SqueezedThermalState(ComplexF64, ξ, n̄; dim=dim)


#
using DiffMesh, LinearAlgebra

n = 100

dl = rand(n-1)
dd = rand(n)
du = rand(n-1)

Td = Tridiagonal(dl,dd,du)
Tt = Tdg(dl,dd,du)

u = rand(n)

@test Td *u ≈ Tt *u atol=1e-12
@test Td'*u ≈ Tt'*u atol=1e-12