#
# sure_tools.jl
#
# Provides an implementation of the Generalized Stein
# Umbiased Risk Estimator (GSURE) and of the prevision error.
#
# [Stein, 1981] Stein, C. M. (1981). Estimation of the mean of a
# multivariate normal distribution. The annals of Statistics,
# pages 1135–1151.
#
# [Eldar, 2008] Eldar, Y. C. (2008). Generalized sure for
# exponential families : Applications to regularization.
# IEEE Transactions on Signal Processing, 57(2) :471–481.
#
# [Ramani et al., 2008] Ramani, S., Blu, T., and Unser, M. (2008).
# Monte-Carlo Sure : A Black-Box Optimization of Regularization
# Parameters for General Denoising Algorithms. IEEE Trans.
# on Image Process., 17(9) :1540–1554.
#
#
#------------------------------------------------------
#
# This file is part of Rhapsodie
#
# Copyright (c) 2017-2021 Laurence DENNEULIN (see Licence.md)
#

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

function MSE_data(x_est::Array{T,N}, x_true::Array{T,N}, d::Array{data_table,1}) where {T <: AbstractFloat,N}
    MSE=0.0;
    n=0;
    for  data in d
        res = MSE_data(x_est, x_true, data)
        MSE +=res[1];
        n += res[2];
    end
    return MSE,n
end

function MSE_data(x_est::Array{T,N}, x_true::Array{T,N}, d::data_table) where {T <: AbstractFloat,N}
        minimum(d.weights) >=0 || error("invalid weights");

        res= d.H*(x_est - x_true);
        MSE = vdot( res, d.weights.*res);
        n= count(d.weights .>0);
    return MSE,n
end

function sure_crit(x::Array{T,N},
                   δx::Array{T,N},
                   d::Array{data_table,1},
                   δd::Array{data_table,1};
                   mask::Array{Array{Bool,2},1}=Array{Array{Bool,2},1}()) where {T <: AbstractFloat,N}

    data_fidelity=0.0;
    trace= 0.0;
    n=0;

    for k=1:length(d)
        if !isempty(mask)
            res = sure_crit(x, δx, d[k], δd[k], mask=mask[k])
        else
            res = sure_crit(x, δx, d[k], δd[k])
        end
        data_fidelity += res[1];
        trace += res[2];
        n += res[3];
    end
    return [data_fidelity ; trace; n]

end

function sure_crit(x::Array{T,N},
                   δx::Array{T,N},
                   d::data_table,
                   δd::data_table;
                   mask::Array{Bool,2}=Array{Bool,2}(undef, 0,0)) where {T <: AbstractFloat,N}

        MAP = (d.weights .!=0); #Valid pixels map

        !isempty(mask) && (size(MAP)==size(mask)) && (MAP .*= mask);

        n=count(MAP);
        if n != 0;
            εb=(δd.data -d.data)[MAP]; #Perturbation of variance ε^2
            εd=d.data - d.H*x; #difference between data and model
            εx=(d.H*(δx - x))[MAP]; #difference between model estimated on perturbed data and on data

            data_fidelity = vdot(εd, d.weights.*εd);

            trace = n*vdot(εb, εx)/vdot(εb, εb)
        else
            data_fidelity=0.
            trace=0.
        end
    return [data_fidelity ; trace; n]

end

function SURE(solver::Function,
              par::Array{Float64,1},
              d::Array{data_table,1},
              δd::Array{data_table,1},
              x0::Array{Float64,3};
              mask::Array{Array{Bool,2},1}=Array{Array{Bool,2},1}())

    x  = solver(x0, d, par);
    δx = solver(x0, δd, par);

    return sure_crit(x,δx, d, δd, mask=mask), x
end

function do_data_perturbation(dset::Array{data_table,1},eps::Float64)
    p = data_table[];
    for d in dset
        push!(p,do_data_perturbation(d,eps));
    end
    return p
end

function do_data_perturbation(d::data_table,eps::Float64)
    return data_table(d.data + eps* randn(size(d.data)), d.weights, d.H);
end

function sure_optim(solver, x)#; linear=true)
        res=sure_tools.SURE(solver, x, grad_tools.dataset, data_perturbed, X0nl)[1]
    return res[1] + 2*res[2];
end


using StanSample, DiffEqBayes, OrdinaryDiffEq, ParameterizedFunctions,
      ModelingToolkit, RecursiveArrayTools, Distributions, Random, Test

# Uncomment for local testing only, make sure MCMCChains and StatsPlots are available
using MCMCChains, StatsPlots

#Random.seed!(123)
cd(@__DIR__)
isdir("tmp") && rm("tmp", recursive=true)

include("../../../stan_string.jl")
include("../../../stan_inference.jl")

println("\nOne parameter case 1\n")
f1 = @ode_def begin
  dx = a*x - x*y
  dy = -3y + x*y
end a

u0 = [1.0,1.0]
tspan = (0.0,10.0)
pars = [1.5]

prob1 = ODEProblem(f1,u0,tspan,pars)
sol = solve(prob1,Tsit5())
t = collect(range(1,stop=10,length=10))
randomized = VectorOfArray([(sol(t[i]) + .5randn(2)) for i in 1:length(t)])
data = convert(Array,randomized)
priors = [truncated(Normal(0.7,1),0.1,2)]

bayesian_result = stan_inference(prob1,t,data,priors;num_samples=2000, nchains=4,
  num_warmup=1000,likelihood=Normal)

if success(bayesian_result.return_code)
  sdf = read_summary(bayesian_result.model)
  sdf |> display

  p  = read_samples(bayesian_result.model, output_format=:particles)
  @test mean(p.theta1) ≈ 1.5 atol=5e-1

  # Uncomment for local chain inspection
  p1 = Vector{Plots.Plot{Plots.GRBackend}}(undef, length(pars))
  p2 = Vector{Plots.Plot{Plots.GRBackend}}(undef, length(pars))
  for i in 1:length(pars)
    p1[i] = plot(title="theta.$i")
    plot!(p1[i], p[Symbol("theta.$i")].particles)
    p2[i] = plot(title="theta.$i")
    density!(p2[i], p[Symbol("theta.$i")].particles)
  end
  plot(p1..., p2..., layout=(length(pars),2))
  savefig("$(@__DIR__)/case_1_chains.png")

  v1 = [mean(p[Symbol("u_hat.$i.1")]) for i in 1:10]
  v2 = [mean(p[Symbol("u_hat.$i.2")]) for i in 1:10]
  qa_prey = zeros(10, 3);
  qa_pred = zeros(10, 3);

  df = read_samples(bayesian_result.model, output_format=:dataframe)
  achn = Array(df);
  for i in 1:10
    qa_prey[i, :] = quantile(achn[:, 4+i], [0.055, 0.5, 0.945])
    qa_pred[i, :] = quantile(achn[:, 14+i], [0.055, 0.5, 0.945])
  end
  p1 = plot(1:10, v1; ribbon=(qa_prey[:, 1], qa_prey[:, 3]), color=:lightgrey, leg=false)
  title!("Prey u_hat 89% quantiles")
  plot!(v1, lab="prey", xlab="time", ylab="prey", color=:darkred)

  p2 = plot(1:10, v2; ribbon=(qa_pred[:, 1], qa_pred[:, 3]), color=:lightgrey, leg=false)
  title!("Preditors u_hat 89% quantiles")
  plot!(v2, lab="pred", xlab="time", ylab="preditors", color=:darkblue)

  plot(p1, p2, layout=(2,1))
  savefig("$(@__DIR__)/case_1_pred_prey.png")
end


### A Pluto.jl notebook ###
# v0.17.4

using Markdown
using InteractiveUtils

# ╔═╡ 6a407bd2-802c-11ec-0c58-b3f5f4ac5b8f
md"""
# Functions

Topics:
1. How to declare a function
2. Duck-typing in Julia
3. Mutating vs. non-mutating functions
4. Some higher order functions

"""

# ╔═╡ e75e9339-1d60-4d04-9063-a5885d0113f8
md"""
## How to declare a function
Julia gives us a few different ways to write a function. The first requires the `function` and `end` keywords
"""

# ╔═╡ a637b4f0-bc1a-4071-be73-d00c257d4266
md"### 1. Default variant:"

# ╔═╡ 07ed7ab9-a5a4-418c-b71f-5c2bc681eeb3
function sayhi(name)

    return "Hi $name, it's great to see you!"
end

# ╔═╡ fd562b69-53ba-4554-9897-6c5709e237fe
md"
We can call either of these functions like this:"

# ╔═╡ 809baaa8-339b-4b70-bb21-8c24a9cab24f
sayhi("C-3PO")

# ╔═╡ 200a1cbb-9e71-4996-8d4f-8204a8701344
md"### 2. Shortcut variant:"

# ╔═╡ 53a47fcd-52f7-49bf-a274-97496258ace6
md"Alternatively, we could have declared either of these functions in a single line"

# ╔═╡ 2413ce2b-ee1b-4d3f-ad8d-ca4064b19d7c
sayhi2(name) = "Hi $name, it's great to see you!"

# ╔═╡ 40666547-6209-42fd-a14f-4c944a5e30ab
sayhi2("R2D2")

# ╔═╡ 9e4f43a7-58e7-4b86-b874-e5cb44298e23
md"### 3. Anonymous functions:"

# ╔═╡ 162585b0-936d-4f9f-8667-3e8c5a1b7212
md"""Finally, we could have declared these as "anonymous" functions"""

# ╔═╡ 191be85a-9cce-4af2-bbe8-6c45d7ab30df
name -> "Hi $name, it's great to see you!"

# ╔═╡ 3a74682e-72ff-4f91-9884-52e26c64c427
sayhi3("Chewbacca")

# ╔═╡ bcd284f3-cb2a-482e-ba27-9a5f7a86c046
md"""
## Duck-typing in Julia
> If it quacks like a duck, it's a duck.

Julia functions will just work on whatever inputs make sense.

For example, sayhi works on the name of this minor tv character, written as an integer...

"""

# ╔═╡ bd4182c4-18ed-41f4-a3f1-50aa553b4924
sayhi(55595472)

# ╔═╡ 47d3b1f8-cc8d-4b0e-89b6-f1d7d06337f8
md"
And `f` will work on a matrix."

# ╔═╡ c4b6e75d-0a51-4159-b7f7-292f5e25d1be
A = rand(3, 3)

# ╔═╡ ec0b4d04-44ce-40ae-b807-8ba38a6a8b59
md"""`f` will also work on a string like "hi" because * is defined for string inputs as string concatenation."""

# ╔═╡ f95f8ea9-1478-4ac2-8e22-085fc6516e56
md"""On the other hand, `f` will not work on a vector. Unlike `A^2`, which is well-defined, the meaning of `v^2` for a vector, `v`, is not a well-defined algebraic operation. """

# ╔═╡ 5620baa2-2aed-49eb-84a1-9fcfdee568c9
v = rand(3)

# ╔═╡ 23b408a4-a91a-4828-a7d8-038325e31d64
Base.:^(x::Vector, i::Int64) = x .^ i

# ╔═╡ 6a407bd2-802c-11ec-3ed6-a7acf451b25a
begin
	bigbreak = HTML("
"^4);
	using PlutoUI
	using Test
end

# ╔═╡ ed5addf3-f534-4ab7-944d-ca948c177ce4
bigbreak

# ╔═╡ e9cf1e4e-e5fa-46de-ab86-f7a35e5b1302
bigbreak

# ╔═╡ dbd09e73-afa6-4065-baee-342951e10031
bigbreak

# ╔═╡ aa60f167-fa1f-4a7e-9d87-eeba0146ed0f
bigbreak

# ╔═╡ f8ce588a-cf7f-4b61-ac9f-afb520003571
function f(x)
    return x^2
end

# ╔═╡ 928f37b1-9059-4b6e-b733-89997203c138
f2(x) = x^2

# ╔═╡ b7a066ac-7d19-4449-bb1e-bf7fcab4eadc
f2(42)

# ╔═╡ 51955b6a-9508-49ca-b435-ba54fb8deb1a
f4(x) = x^2

# ╔═╡ 6b2b763f-6223-4c50-93a0-54680c4313ff
f3 = x -> x^2

# ╔═╡ 353c23a2-4e6d-401f-8a0c-8038f4ef0998
f3(42)

# ╔═╡ 652b18db-731f-4283-b5dc-26dc95bf25a5
bigbreak

# ╔═╡ 5bea284c-9eed-467a-b916-1ba6bc99495f
md"""
## Mutating vs. non-mutating functions

By convention, functions followed by `!` alter their contents and functions lacking `!` do not.

For example, let's look at the difference between `sort` and `sort!`.
"""

# ╔═╡ ee1468af-b2d0-4505-ad59-4b04437d0751
let
	w = [3, 1, 2]
	w = sort(w)
	w
end

# ╔═╡ 1589bbe3-15c4-495b-aae2-eac2f421b2f8
md"""
`sort(w)` returns a sorted array that contains the same elements as `w`, but `w` is left unchanged.

On the other hand, when we run `sort!(w)`, the contents of `w` are sorted within the array `w`.
"""

# ╔═╡ efb84537-41b6-41af-9205-9df193d49dcc
let
	w = [3, 1, 2]
	sort!(w)
	w
end

# ╔═╡ d52cb505-1492-4e64-869a-d598fd99d361
bigbreak

# ╔═╡ 2bc704f5-82e8-4a43-b5ae-1cdbfc18747b
md"""
## Optional arguments and keyword arguments

There are three kind of arguments:

- **Required arguments**, first in the list
- **Optional arguments**, with `= value` to indicate default choice (order matters!)
- **Keword arguments**, after a semicolon `;`. (order does not matter, only keywords)
"""

# ╔═╡ d451cbb3-fd29-46c8-a4ff-a36db76c97ca
md"""
```julia
function fnc( req1, req2,
		      op1 = dv1, opt2 = dv2;
   			  kw1 = kval1, kw2 = kval2)

end
```

"""

# ╔═╡ 3eb97bc4-0ea7-466c-a237-f40fc8b43a71
function g(x, y = 1.0; a = 2, b = 3)
	return (;x,y,a,b)
end

# ╔═╡ 3b57b7eb-d4ea-4668-a702-48d1950e8d63
md"You can call this function in this way:"

# ╔═╡ ce4831c8-e098-4837-901f-383a445a2998
g("🐱"; b = "🐶")

# ╔═╡ 216540a3-0c24-4854-b17e-cd1623e93c43
g("🐱", "🐶")

# ╔═╡ 2de2c2bb-1b17-4085-a0d7-de322fb21370
g("🐱"; a = "🐶")

# ╔═╡ 8eee9f7d-7645-4ff1-b4ef-04dbb4da386e
g("🐱")

# ╔═╡ 3c315931-12fb-4b86-b745-9053cb8ccf68
md"_It is recommended (but optional) to separate optional arguments and keyword arguments with ';'_"

# ╔═╡ e8a86b02-ecc9-479b-97b9-68e3f6ba410a
bigbreak

# ╔═╡ 20f9c597-b1a9-4f32-bca4-ad174f3b78ed
md"""
## Some higher order functions

### map

`map` is a "higher-order" function in Julia that *takes a function* as one of its input arguments.
`map` then applies that function to every element of the data structure you pass it. For example, executing

```julia
map(f, [1, 2, 3])
```
will give you an output array where the function `f` has been applied to all elements of `[1, 2, 3]`
```julia
[f(1), f(2), f(3)]
```
"""

# ╔═╡ 2b69252d-cf04-40b6-95d9-eacf369c4dc5
md"""
Here we've squared all the elements of the vector `[1, 2, 3]`, rather than squaring the vector `[1, 2, 3]`.

To do this, we could have passed to `map` an anonymous function rather than a named function, such as
"""

# ╔═╡ 88dbda57-a845-404e-b680-31c2a473794e
x -> x^3

# ╔═╡ b53c7e60-559c-4d92-99bb-f970f44e48cb
md"via"

# ╔═╡ d0c94953-b7a5-4749-aa4b-b5286233767c
map(x -> x^3, [1, 2, 3])

# ╔═╡ ebb4b925-9c00-471e-b248-d7f5c7794ca0
md"and now we've cubed all the elements of `[1, 2, 3]`!"

# ╔═╡ 76d3ebff-1e69-4f7a-bb55-a10f5f83c806
with_terminal() do
	println("test")
end

# ╔═╡ 783648ba-d6b5-4cdf-91ac-13323255428a
bigbreak

# ╔═╡ e452fd6d-fb31-43b6-b399-0863784115fc
md"""
## broadcast

`broadcast` is another higher-order function like `map`. `broadcast` is a generalization of `map`, so it can do every thing `map` can do and more. The syntax for calling `broadcast` is the same as for calling `map`
"""

# ╔═╡ 69e8fbb4-38e2-402b-89a0-018b4083d7a1
md"""
and again, we've applied `f` (squared) to all the elements of `[1, 2, 3]` - this time by "broadcasting" `f`!

Some syntactic sugar for calling `broadcast` is to place a `.` between the name of the function you want to `broadcast` and its input arguments. For example,

```julia
broadcast(f, [1, 2, 3])
```
is the same as
```julia
f.([1, 2, 3])
```
"""

# ╔═╡ a20f8214-0126-43f4-8a56-bcfed741463f
md"""
Notice again how different this is from calling
```julia
f([1, 2, 3])
```
We can square every element of a vector, but we can't square a vector!
"""

# ╔═╡ 0bdca8db-d6f6-4cc3-8661-98c5cb72fdae
md"""
To drive home the point, let's look at the difference between

```julia
f(A)
```
and
```julia
f.(A)
```
for a matrix `A`:
"""

# ╔═╡ 963169eb-9244-4bbd-8337-65fec67c725c
A

# ╔═╡ 7b411fc5-7ab1-4052-a199-3195a931860a
B = [i + 3*j for j in 0:2, i in 1:3]

# ╔═╡ 99edddcb-73ac-4848-9540-d8ac22f80e5f

# ╔═╡ 3dafca9a-56ee-4a83-9853-917efaf6a93d
md"""
As before we see that for a matrix, `A`,
```
f(A) = A^2 = A * A
```

On the other hand,
"""

# ╔═╡ e56b633e-8b74-4629-b564-f9a4a656626f
f(x,y) = 10*x + y

# ╔═╡ 87ceb82b-94a7-4e5a-84da-f56e24279755
f(42)

# ╔═╡ 9b8a0c73-27f4-48af-b13b-0af1711440b3
f(A)

# ╔═╡ 04a0dbd8-d8f2-434b-8968-63ee989b3131
f("hi")

# ╔═╡ 770cafe7-11e8-4b1c-86c9-b51e4c20ec42
# This won't work:
f(v)

# ╔═╡ 01590c72-ec3e-4db5-957a-02308a6a5e6c
map(f, [1, 2, 3])

# ╔═╡ 928bf981-ba2d-4eb7-9b31-8dd88be33976
broadcast(f, [1, 2, 3])

# ╔═╡ 413e16a6-c8af-43b5-85ab-e6bd16690006
f.( [1, 2, 3] )

# ╔═╡ 33b1ce0b-371d-4af8-af52-7328e4e5b70b
f(B) == B*B == B^2

# ╔═╡ f3efe46a-f1bc-4792-884d-b29ef98595df
f.(B) == B .^ 2

# ╔═╡ 1537e7c8-4e74-4eb5-8a26-7dbb3af600ab
f.( [1,2], 2)

# ╔═╡ a18cdba8-8969-4715-992e-2f3ea339c6e1
md"""
contains the squares of all the entries of `A`.

This dot syntax for broadcasting allows us to write relatively complex compound elementwise expressions in a way that looks natural/closer to mathematical notation. For example, we can write
"""

# ╔═╡ 59fea39e-dea8-4412-8cd3-64b4cede04db
A .+ 2 .* f.(A) ./ A

# ╔═╡ e29c6b03-1128-4928-9878-289886dbef38
A .+ 2 .* f.(A) ./ A

# ╔═╡ 4477326e-bef1-4b74-a24b-c421fcb60c18
md"or"

# ╔═╡ f0dead06-58f6-4c97-9f14-5000b79be8cb
@. A + 2 * f(A) / A

# ╔═╡ 429ec6d9-f80c-4ca7-85b8-db8096f826ea
md"instead of"

# ╔═╡ c6f86e8a-05e8-4ad3-8e7c-3598b22a4779
broadcast(x -> x + 2 * f(x) / x, A)

# ╔═╡ 60eea80b-a510-455b-a2a4-d16366303b87
md"and the two will do exactly the same!"

# ╔═╡ 63032f05-4635-438b-9c56-5571620a180d
bigbreak

# ╔═╡ 62d01357-fdc5-4bd2-9b32-d7515c0c44f6
md"""### Fusing the dot operator

Composition of `broadcast` is suboptimal, as in each step an intermediate resulat is stored, i.e.
```julia
A .+ B .+ C == broadcast( (r+c) -> r + c, broadcast( (a,b) -> a+b, A, B), C)
```

But with a hand _macro_ we can fuse those broadcast calls, i.e.
```julia
@. A + B + C  ==   broadcast( (a,b,c) -> a+b+c, A, B, C)
```
"""

# ╔═╡ e4f1a223-ff88-4588-94b1-1bd6aeb0e856
A .* A .* A

# ╔═╡ 94070b96-502c-4685-bdb9-f6ee6fe55a06
@. A * A * A

# ╔═╡ 281ac3be-aaea-4078-9fb0-391997924306

# ╔═╡ 96e07c5b-6a14-4aa0-8d96-99aaf19a235f
@allocated @. A * A * A

# ╔═╡ 7617e358-c1bf-42b8-b0e3-6e50f993ec2e
@allocated A .* A .* A

# ╔═╡ 6b3f422f-e8cc-480c-aede-296808d063fc
bigbreak

# ╔═╡ fc405e7f-d062-41ed-bcc3-5599e80e3467
md"""
## Exercises

6.1

Write a function `add_one` that adds 1 to its input.
"""

# ╔═╡ 1e1b4cca-0753-440b-9eb2-6f8382964e71

# ╔═╡ 68c25eb9-c682-4f7d-8c82-1e786557de01

# ╔═╡ 86edb75d-a484-4d7f-bf9a-bc6b97cce3ec
( @isdefined(add_one) &&
hasmethod(add_one, Tuple{Union{Any,Int,Float64}}) &&
add_one(11) == 12)

# ╔═╡ 185ffedc-8462-4ef3-add4-9ba371e236ee
bigbreak

# ╔═╡ 56ead0a8-4c9b-4e07-9a8c-0c56b4b30c21
md"""
6.2

Use map or broadcast to increment every element of matrix `B` by 1 and assign it to a variable `A1`.
"""

# ╔═╡ 8e725749-3c7a-413d-8838-49e471c8b47e
B

# ╔═╡ bd5e73ca-a683-476b-8239-ddf57cba89ba
begin
	A1 = B
end

# ╔═╡ e11ca969-a16a-451d-9f6f-5f826d18a55f

# ╔═╡ b9d3ae00-42a2-4007-a0f9-3327d656a425
A1 == [2 3 4; 5 6 7; 8 9 10]

# ╔═╡ e83429b7-77e9-4efe-a365-827b7234c360
bigbreak

# ╔═╡ 93609093-4353-44a3-be9f-5ac41ec69c81
md"""
6.3

Use the broadcast dot syntax to increment every element of matrix `A1` by `1` and store it in variable `A2`
"""

# ╔═╡ 43917a91-235f-49ab-9a60-dc0dfd9bcb4d
begin
	A2 = A1
end

# ╔═╡ da127939-3adc-40c9-9031-db9dd3de847b

# ╔═╡ 15459bb0-bd18-4d48-b39a-193f7a01b1f8
A2 == [3 4 5; 6 7 8; 9 10 11]

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[compat]
PlutoUI = "~0.7.32"
"""

# ╔═╡ 00000000-0000-0000-0000-000000000002
PLUTO_MANIFEST_TOML_CONTENTS = """
# This file is machine-generated - editing it directly is not advised

julia_version = "1.7.1"
manifest_format = "2.0"

[[deps.AbstractPlutoDingetjes]]
deps = ["Pkg"]
git-tree-sha1 = "8eaf9f1b4921132a4cff3f36a1d9ba923b14a481"
uuid = "6e696c72-6542-2067-7265-42206c756150"
version = "1.1.4"

[[deps.ArgTools]]
uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"

[[deps.Artifacts]]
uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"

[[deps.Base64]]
uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

[[deps.ColorTypes]]
deps = ["FixedPointNumbers", "Random"]
git-tree-sha1 = "024fe24d83e4a5bf5fc80501a314ce0d1aa35597"
uuid = "3da002f7-5984-5a60-b8a6-cbb66c0b333f"
version = "0.11.0"

[[deps.CompilerSupportLibraries_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"

[[deps.Dates]]
deps = ["Printf"]
uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"

[[deps.Downloads]]
deps = ["ArgTools", "LibCURL", "NetworkOptions"]
uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"

[[deps.FixedPointNumbers]]
deps = ["Statistics"]
git-tree-sha1 = "335bfdceacc84c5cdf16aadc768aa5ddfc5383cc"
uuid = "53c48c17-4a7d-5ca2-90c5-79b7896eea93"
version = "0.8.4"

[[deps.Hyperscript]]
deps = ["Test"]
git-tree-sha1 = "8d511d5b81240fc8e6802386302675bdf47737b9"
uuid = "47d2ed2b-36de-50cf-bf87-49c2cf4b8b91"
version = "0.0.4"

[[deps.HypertextLiteral]]
git-tree-sha1 = "2b078b5a615c6c0396c77810d92ee8c6f470d238"
uuid = "ac1192a8-f4b3-4bfe-ba22-af5b92cd3ab2"
version = "0.9.3"

[[deps.IOCapture]]
deps = ["Logging", "Random"]
git-tree-sha1 = "f7be53659ab06ddc986428d3a9dcc95f6fa6705a"
uuid = "b5f81e59-6552-4d32-b1f0-c071b021bf89"
version = "0.2.2"

[[deps.InteractiveUtils]]
deps = ["Markdown"]
uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"

[[deps.JSON]]
deps = ["Dates", "Mmap", "Parsers", "Unicode"]
git-tree-sha1 = "8076680b162ada2a031f707ac7b4953e30667a37"
uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
version = "0.21.2"

[[deps.LibCURL]]
deps = ["LibCURL_jll", "MozillaCACerts_jll"]
uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"

[[deps.LibCURL_jll]]
deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"

[[deps.LibGit2]]
deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"

[[deps.LibSSH2_jll]]
deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"

[[deps.Libdl]]
uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"

[[deps.LinearAlgebra]]
deps = ["Libdl", "libblastrampoline_jll"]
uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

[[deps.Logging]]
uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"

[[deps.Markdown]]
deps = ["Base64"]
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"

[[deps.MbedTLS_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"

[[deps.Mmap]]
uuid = "a63ad114-7e13-5084-954f-fe012c677804"

[[deps.MozillaCACerts_jll]]
uuid = "14a3606d-f60d-562e-9121-12d972cd8159"

[[deps.NetworkOptions]]
uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"

[[deps.OpenBLAS_jll]]
deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"]
uuid = "4536629a-c528-5b80-bd46-f80d51c5b363"

[[deps.Parsers]]
deps = ["Dates"]
git-tree-sha1 = "0b5cfbb704034b5b4c1869e36634438a047df065"
uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
version = "2.2.1"

[[deps.Pkg]]
deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

[[deps.PlutoUI]]
deps = ["AbstractPlutoDingetjes", "Base64", "ColorTypes", "Dates", "Hyperscript", "HypertextLiteral", "IOCapture", "InteractiveUtils", "JSON", "Logging", "Markdown", "Random", "Reexport", "UUIDs"]
git-tree-sha1 = "ae6145ca68947569058866e443df69587acc1806"
uuid = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
version = "0.7.32"

[[deps.Printf]]
deps = ["Unicode"]
uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"

[[deps.REPL]]
deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"

[[deps.Random]]
deps = ["SHA", "Serialization"]
uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

[[deps.Reexport]]
git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b"
uuid = "189a3867-3050-52da-a836-e630ba90ab69"
version = "1.2.2"

[[deps.SHA]]
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"

[[deps.Serialization]]
uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"

[[deps.Sockets]]
uuid = "6462fe0b-24de-5631-8697-dd941f90decc"

[[deps.SparseArrays]]
deps = ["LinearAlgebra", "Random"]
uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

[[deps.Statistics]]
deps = ["LinearAlgebra", "SparseArrays"]
uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

[[deps.TOML]]
deps = ["Dates"]
uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"

[[deps.Tar]]
deps = ["ArgTools", "SHA"]
uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"

[[deps.Test]]
deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[[deps.UUIDs]]
deps = ["Random", "SHA"]
uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"

[[deps.Unicode]]
uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"

[[deps.Zlib_jll]]
deps = ["Libdl"]
uuid = "83775a58-1f1d-513f-b197-d71354ab007a"

[[deps.libblastrampoline_jll]]
deps = ["Artifacts", "Libdl", "OpenBLAS_jll"]
uuid = "8e850b90-86db-534c-a0d3-1478176c7d93"

[[deps.nghttp2_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"

[[deps.p7zip_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
"""

# ╔═╡ Cell order:
# ╟─6a407bd2-802c-11ec-3ed6-a7acf451b25a
# ╟─6a407bd2-802c-11ec-0c58-b3f5f4ac5b8f
# ╟─e75e9339-1d60-4d04-9063-a5885d0113f8
# ╟─a637b4f0-bc1a-4071-be73-d00c257d4266
# ╠═07ed7ab9-a5a4-418c-b71f-5c2bc681eeb3
# ╠═f8ce588a-cf7f-4b61-ac9f-afb520003571
# ╟─fd562b69-53ba-4554-9897-6c5709e237fe
# ╠═809baaa8-339b-4b70-bb21-8c24a9cab24f
# ╠═87ceb82b-94a7-4e5a-84da-f56e24279755
# ╟─ed5addf3-f534-4ab7-944d-ca948c177ce4
# ╟─200a1cbb-9e71-4996-8d4f-8204a8701344
# ╟─53a47fcd-52f7-49bf-a274-97496258ace6
# ╠═2413ce2b-ee1b-4d3f-ad8d-ca4064b19d7c
# ╠═928f37b1-9059-4b6e-b733-89997203c138
# ╠═51955b6a-9508-49ca-b435-ba54fb8deb1a
# ╠═40666547-6209-42fd-a14f-4c944a5e30ab
# ╠═b7a066ac-7d19-4449-bb1e-bf7fcab4eadc
# ╟─e9cf1e4e-e5fa-46de-ab86-f7a35e5b1302
# ╟─9e4f43a7-58e7-4b86-b874-e5cb44298e23
# ╟─162585b0-936d-4f9f-8667-3e8c5a1b7212
# ╠═191be85a-9cce-4af2-bbe8-6c45d7ab30df
# ╠═6b2b763f-6223-4c50-93a0-54680c4313ff
# ╠═3a74682e-72ff-4f91-9884-52e26c64c427
# ╠═353c23a2-4e6d-401f-8a0c-8038f4ef0998
# ╟─dbd09e73-afa6-4065-baee-342951e10031
# ╟─bcd284f3-cb2a-482e-ba27-9a5f7a86c046
# ╠═bd4182c4-18ed-41f4-a3f1-50aa553b4924
# ╟─47d3b1f8-cc8d-4b0e-89b6-f1d7d06337f8
# ╠═c4b6e75d-0a51-4159-b7f7-292f5e25d1be
# ╠═9b8a0c73-27f4-48af-b13b-0af1711440b3
# ╟─ec0b4d04-44ce-40ae-b807-8ba38a6a8b59
# ╠═04a0dbd8-d8f2-434b-8968-63ee989b3131
# ╟─f95f8ea9-1478-4ac2-8e22-085fc6516e56
# ╠═5620baa2-2aed-49eb-84a1-9fcfdee568c9
# ╟─aa60f167-fa1f-4a7e-9d87-eeba0146ed0f
# ╠═770cafe7-11e8-4b1c-86c9-b51e4c20ec42
# ╠═23b408a4-a91a-4828-a7d8-038325e31d64
# ╟─652b18db-731f-4283-b5dc-26dc95bf25a5
# ╟─5bea284c-9eed-467a-b916-1ba6bc99495f
# ╠═ee1468af-b2d0-4505-ad59-4b04437d0751
# ╟─1589bbe3-15c4-495b-aae2-eac2f421b2f8
# ╠═efb84537-41b6-41af-9205-9df193d49dcc
# ╟─d52cb505-1492-4e64-869a-d598fd99d361
# ╟─2bc704f5-82e8-4a43-b5ae-1cdbfc18747b
# ╟─d451cbb3-fd29-46c8-a4ff-a36db76c97ca
# ╠═3eb97bc4-0ea7-466c-a237-f40fc8b43a71
# ╟─3b57b7eb-d4ea-4668-a702-48d1950e8d63
# ╠═ce4831c8-e098-4837-901f-383a445a2998
# ╠═216540a3-0c24-4854-b17e-cd1623e93c43
# ╠═2de2c2bb-1b17-4085-a0d7-de322fb21370
# ╠═8eee9f7d-7645-4ff1-b4ef-04dbb4da386e
# ╟─3c315931-12fb-4b86-b745-9053cb8ccf68
# ╟─e8a86b02-ecc9-479b-97b9-68e3f6ba410a
# ╟─20f9c597-b1a9-4f32-bca4-ad174f3b78ed
# ╠═01590c72-ec3e-4db5-957a-02308a6a5e6c
# ╟─2b69252d-cf04-40b6-95d9-eacf369c4dc5
# ╠═88dbda57-a845-404e-b680-31c2a473794e
# ╟─b53c7e60-559c-4d92-99bb-f970f44e48cb
# ╠═d0c94953-b7a5-4749-aa4b-b5286233767c
# ╟─ebb4b925-9c00-471e-b248-d7f5c7794ca0
# ╠═76d3ebff-1e69-4f7a-bb55-a10f5f83c806
# ╟─783648ba-d6b5-4cdf-91ac-13323255428a
# ╟─e452fd6d-fb31-43b6-b399-0863784115fc
# ╠═928bf981-ba2d-4eb7-9b31-8dd88be33976
# ╟─69e8fbb4-38e2-402b-89a0-018b4083d7a1
# ╠═413e16a6-c8af-43b5-85ab-e6bd16690006
# ╟─a20f8214-0126-43f4-8a56-bcfed741463f
# ╟─0bdca8db-d6f6-4cc3-8661-98c5cb72fdae
# ╠═963169eb-9244-4bbd-8337-65fec67c725c
# ╠═7b411fc5-7ab1-4052-a199-3195a931860a
# ╠═99edddcb-73ac-4848-9540-d8ac22f80e5f
# ╠═33b1ce0b-371d-4af8-af52-7328e4e5b70b
# ╟─3dafca9a-56ee-4a83-9853-917efaf6a93d
# ╠═f3efe46a-f1bc-4792-884d-b29ef98595df
# ╠═e56b633e-8b74-4629-b564-f9a4a656626f
# ╠═1537e7c8-4e74-4eb5-8a26-7dbb3af600ab
# ╟─a18cdba8-8969-4715-992e-2f3ea339c6e1
# ╠═59fea39e-dea8-4412-8cd3-64b4cede04db
# ╠═e29c6b03-1128-4928-9878-289886dbef38
# ╟─4477326e-bef1-4b74-a24b-c421fcb60c18
# ╠═f0dead06-58f6-4c97-9f14-5000b79be8cb
# ╟─429ec6d9-f80c-4ca7-85b8-db8096f826ea
# ╠═c6f86e8a-05e8-4ad3-8e7c-3598b22a4779
# ╟─60eea80b-a510-455b-a2a4-d16366303b87
# ╟─63032f05-4635-438b-9c56-5571620a180d
# ╠═62d01357-fdc5-4bd2-9b32-d7515c0c44f6
# ╠═e4f1a223-ff88-4588-94b1-1bd6aeb0e856
# ╠═94070b96-502c-4685-bdb9-f6ee6fe55a06
# ╠═281ac3be-aaea-4078-9fb0-391997924306
# ╠═96e07c5b-6a14-4aa0-8d96-99aaf19a235f
# ╠═7617e358-c1bf-42b8-b0e3-6e50f993ec2e
# ╟─6b3f422f-e8cc-480c-aede-296808d063fc
# ╠═fc405e7f-d062-41ed-bcc3-5599e80e3467
# ╠═1e1b4cca-0753-440b-9eb2-6f8382964e71
# ╟─68c25eb9-c682-4f7d-8c82-1e786557de01
# ╟─86edb75d-a484-4d7f-bf9a-bc6b97cce3ec
# ╟─185ffedc-8462-4ef3-add4-9ba371e236ee
# ╟─56ead0a8-4c9b-4e07-9a8c-0c56b4b30c21
# ╠═8e725749-3c7a-413d-8838-49e471c8b47e
# ╠═bd5e73ca-a683-476b-8239-ddf57cba89ba
# ╟─e11ca969-a16a-451d-9f6f-5f826d18a55f
# ╟─b9d3ae00-42a2-4007-a0f9-3327d656a425
# ╟─e83429b7-77e9-4efe-a365-827b7234c360
# ╠═93609093-4353-44a3-be9f-5ac41ec69c81
# ╠═43917a91-235f-49ab-9a60-dc0dfd9bcb4d
# ╟─da127939-3adc-40c9-9031-db9dd3de847b
# ╟─15459bb0-bd18-4d48-b39a-193f7a01b1f8
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


using Test
using Quiqbox

@testset "OneBody.jl tests" begin
    errT1 = 1e-10
    errT2 = 1e-10
    nucs = fill("H", 2)
    cens = [[-0.7, 0.0, 0.0], [ 0.7, 0.0, 0.0]]
    bf1 = genBasisFunc(cens[1], ("STO-3G", "H"))[]
    bf2 = genBasisFunc(cens[2], ("STO-3G", "H"))[]
    bs = [bf1, bf2]

    S = [1.0 0.6593182058508896; 0.6593182058508896 1.0]
    @test isapprox(overlap(bf1, bf1), ones(1,1), atol=errT1)
    @test isapprox(overlap(bf1, bf2), S[2]*ones(1,1), atol=errT1)
    @test isapprox(overlaps(bs), S)

    V = [-1.8804408905227634 -1.1948346220535715; -1.1948346220535715 -1.8804408905227632]
    @test isapprox(nucAttraction(bf1, bf2, nucs, cens), V[2]*ones(1,1), atol=errT1)
    @test isapprox(nucAttractions(bs, nucs, cens), V, atol=errT1)

    T = [0.7600318799755844 0.23645465829079276; 0.23645465829079276 0.7600318799755844]
    @test isapprox(elecKinetic(bf1, bf2), T[2]*ones(1,1,1), atol=errT1)
    @test isapprox(elecKinetics(bs), T)

    Hc = coreH(bs, nucs, cens)
    @test isapprox(Hc, T+V, atol=errT1)
    @test isapprox(coreHij(bf1, bf1, nucs, cens)[], Hc[1], atol=errT2)
    @test isapprox(coreHij(bf2, bf1, nucs, cens)[],
                   coreHij(bf1, bf2, nucs, cens)[], atol=errT2)
    @test isapprox(coreHij(bf2, bf1, nucs, cens)[], Hc[3], atol=errT2)
    @test isapprox(coreHij(bf2, bf2, nucs, cens)[], Hc[4], atol=errT2)
end


@testset "462.minimum-moves-to-equal-array-elements-ii.jl" begin
    @test min_moves2([-9, -17, -18, 12, -22, 3, -23, -15, 27, 23, 14, -16, 7, -7, -24, 22, -6, -10, -6, 2, -19, 18, 7, 16, 12, 9, -28, -25, 2, 4, 27, -19, -3, 21, 10, 14, -8, -11, 14, -17, -7, -22, -3, -25, 18, -29, -24, -13, 3, -4]) == 695
    @test min_moves2([1, 2, 3]) == 2
end


"""
	Modify a cluster by adding or removing vertices by picking at each step
	the vertex that has the maximum value of (Deg_external - Deg_Internal).
	Each vertex can be added in/removed only once.
"""
function refineCut(G::SparseMatrixCSC{Tv,Ti}, s::Array{Int64,1}) where {Tv,Ti}

	n = max(G.n, G.m)

	# inset - keep track of vertices in the low conductance set
	inset = zeros(Bool, n)
	for u in s
		inset[u] = true
	end

	# pq_in - costs are maintained to add elements to the set
	# pq_out - costs are maintained to remove elements from the set
	pq_in = DataStructures.PriorityQueue(Base.Order.Reverse)
	pq_out = DataStructures.PriorityQueue(Base.Order.Forward)

	for u in 1:n
		cost = 0

		for i in 1:deg(G,u)
			v = nbri(G,u,i)
			if inset[v]
				cost = cost + weighti(G,u,i)
			else
				cost = cost - weighti(G,u,i)
			end
		end

		if inset[u] == false
			pq_in[u] = cost
		end
		if inset[u] == true
			pq_out[u] = cost
		end
	end

	improve = true
	tries = 0
	while improve && tries < n
		improve = false

		# try to add a new vertex into the set
		if !isempty(pq_in)
			u,contrib = DataStructures.peek(pq_in)

			if contrib >= 0
				improve = true
				DataStructures.dequeue!(pq_in)
				inset[u] = true

				for i in 1:deg(G,u)
					v = nbri(G,u,i)

					if v in keys(pq_in)
						pq_in[v] = pq_in[v] + weighti(G,u,i)
					end
					if v in keys(pq_out)
						pq_out[v] = pq_out[v] + weighti(G,u,i)
					end
				end
			end
		end

		# try to remove a vertex from the set
		if !isempty(pq_out)
			u,contrib = DataStructures.peek(pq_out)
			if contrib < 0
				improve = true
				DataStructures.dequeue!(pq_out)
				inset[u] = false

				for i in 1:deg(G,u)
					v = nbri(G,u,i)
					if v in keys(pq_in)
						pq_in[v] = pq_in[v] - weighti(G,u,i)
					end
					if v in keys(pq_out)
						pq_out[v] = pq_out[v] - weighti(G,u,i)
					end
				end
			end
		end
	end

	news = (Int64)[]
	for i in 1:n
		if inset[i]
			push!(news, i)
		end
	end

	return news
end

"""
	Modify a cluster by passing through all the vertices exactly once and
	adding/removing them based on the value of (Deg_external - Deg_Internal).
"""
function dumbRefineCut(G::SparseMatrixCSC{Tv,Ti}, s::Array{Int64,1}) where {Tv,Ti}

	n = max(G.n, G.m)

	news = BitSet(s)

	for v in 1:n
		nrbsA = 0
		nrbsnotA = 0
		for i in 1:deg(G, v)
			u = nbri(G, v, i)
			if u in news
				nrbsA = nrbsA + weighti(G, v, i)
			else
				nrbsnotA = nrbsnotA + weighti(G, v, i)
			end
		end

		if nrbsA >= nrbsnotA
			push!(news, v)
		else
			if v in news
				pop!(news, v)
			end
		end
	end

	return collect(news)

end


using Makie
using MakieLayout

include("sirmodel.jl")

R0=2.0
N=1000
γ=0.1
β=γ*R0/N

sol = solveSIR(N, β, γ)

rep = calcEffectiveReproduction(sol)

scene, layout = layoutscene(resolution=(1200,900))

xticks=0:10:160

ax = layout[1,1] = LAxis(scene, title="SIR Model,N=$N,β=$β,γ=$γ", xticks=xticks)

ls = [lines!(ax, sol.t, sol[1,:], color=:blue),
      lines!(ax, sol.t, sol[2,:], color=:red),
      lines!(ax, sol.t, sol[3,:], color=:green)]

layout[1,1] = LLegend(scene, ls, ["Susceptible", "Infectious", "Removed"],
                      width=Auto(), height=Auto(), halign=:left, valign=:center,
                      tellwidth=false, tellheight=false, margin=(10,10,10,10))

bx = layout[2,1] = LAxis(scene, title="Effective Reproduce. \0=$R0",xticks=xticks)
lines!(bx, rep[2], rep[1], color=:blue)

#display(scene)
#readline()

save("plot.png",scene)


# The entry format is (cudaname, julianame, kernelcode)
# With single name entries cudaname=julianame and kernelcode=name(xi,yi).
# I commented out functions if I don't know the Julia equivalent.

binary_ops = [
    ("add","+","xi+yi"),
    ("sub","-","xi-yi"),
    ("mul","*","xi*yi"),
    ("div","/","xi/yi"),
    ("pow","^","pow(xi,yi)"),
    ("max","max","(xi>yi?xi:yi)"),
    ("min","min","(xi<yi?xi:yi)"),
    ("eq","==","xi==yi"),
    ("ne","!=","xi!=yi"),
    ("gt",">","xi>yi"),
    ("ge",">=","xi>=yi"),
    ("lt","<","xi<yi"),
    ("le","<=","xi<=yi"),
    # "hypot",
    # "rhypot",
    # "atan2",
    # "frexp",
    # "ldexp",
    # "scalbn",
    # "scalbln",
    # "fmod",
    # "remainder",
    # "mod",
    # "fdim",
    ("invxback","invxback","(-xi*yi*yi)"),
    ("reluback","reluback","(yi>0?xi:0)"),
    ("eluback","eluback","(yi>0?xi:xi*(1+yi))"),
    ("seluback","seluback","(yi>0?1.0507009873554805*xi:xi*(1.7580993408473773+yi))"),
    ("sigmback","sigmback","(xi*yi*(1-yi))"),
    ("tanhback","tanhback","(xi*(1-yi*yi))"),
    # ("rpow","rpow","pow(yi,xi)"),   # need this for Array.^Scalar -> cuda bug #108 switching to CUDA.jl for pow
]

unary_ops_with_int_degree = [
    # cuda does not define negative degrees for bessel, fix it here:
    ("jn", "besselj", "(d>=0 ? jn(d,xi) : d%2==0 ? jn(-d,xi) : -jn(-d,xi))"),
    ("yn", "bessely", "(d>=0 ? yn(d,xi) : d%2==0 ? yn(-d,xi) : -yn(-d,xi))"), # bessely not defined for negative x!
]

# The following list comes from the NVIDIA math docs with some extras.
# http://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#mathematical-functions-appendix
# The entry format is (cudaname, julianame, kernelcode)
# With single name entries cudaname=julianame and kernelcode=name(xi).
# I commented out functions if I don't know the Julia equivalent.
unary_ops = [
("abs2", "abs2", "(xi*xi)"),
("abs", "abs", "(xi<0?-xi:xi)"),
"acos",
"acosh",
"asin",
"asinh",
"atan",
"atanh",
"cbrt",
"ceil",
"cos",
"cosh",
"cospi",
# ("cyl_bessel_i0", "besseli0"),  # besseli0,i1 is not defined in SpecialFunctions
# ("cyl_bessel_i1", "besseli1"),
"erf",
"erfc",
"erfcinv",
"erfcx",
"erfinv",
"exp",
"exp10",
"exp2",
"expm1",
"floor",
# "ilogb",
("invx", "invx", "1/xi"),
("j0", "besselj0"),
("j1", "besselj1"),
("gamma_impl", "gamma"),
("lgamma", "lgamma"), # TODO: SpecialFunctions 0.8: lgamma(x::Real)` is deprecated, use `(logabsgamma(x))[1]` instead. Other alternative is loggamma, throws a DomainError if gamma(x) is negative.
("digamma_impl", "digamma"),
("trigamma_impl", "trigamma"),
#
# "llrint",
# "llround",
"log",
"log10",
"log1p",
"log2",
# "logb",
# "lrint",
# "lround",
# "nearbyint",
("neg", "-", "-xi"),
# "normcdf",
# "normcdfinv",
("one", "one", "1"),
# "rcbrt",
("relu", "relu", "(xi>0?xi:0)"),
("elu", "elu", "(xi>0?xi:exp(xi)-1)"),
("selu", "selu", "1.0507009873554805*(xi>0?xi:1.6732632423543778*(exp(xi)-1))"),
# "rint",
"round",
# "rsqrt",
("sigm", "sigm", "(xi>=0?1/(1+exp(-xi)):(exp(xi)/(1+exp(xi))))"),
("sign", "sign", "(xi>0?1:xi<0?-1:0)"),
"sin",
"sinh",
"sinpi",
"sqrt",
"tan",
"tanh",
# "tgamma",
"trunc",
("y0","bessely0"),
("y1","bessely1"),
("zero", "zero", "0"),
]

reduction_ops = [
# The entry format is (cudaname, julianame, merge, item, init)
# ai is the accumulator, xi is the array element
("sum","sum","ai+xi","xi","0"),
("prod","prod","ai*xi","xi","1"),
("maximum","maximum","(ai>xi ? ai : xi)","xi","(-INFINITY)"),
("minimum","minimum","(ai<xi ? ai : xi)","xi","INFINITY"),
("sumabs","sumabs","ai+xi","(xi<0 ? -xi : xi)","0"),
("sumabs2","sumabs2","ai+xi","(xi*xi)","0"),
("maxabs","maxabs","(ai>xi ? ai : xi)","(xi<0 ? -xi : xi)","0"),
("minabs","minabs","(ai<xi ? ai : xi)","(xi<0 ? -xi : xi)","INFINITY"),
("countnz","countnz","ai+xi","(xi!=0)","0"),
]


module BPPrivInfo

export integrated_payoff
export expected_payoff
export V
export λ
export upper_integral_λ
export Λ
export μ
export optimise
export UnivariateDensity
export discretise

using LinearAlgebra
using Optim
using QuadGK
using Distributions
using NLsolve
using Quadrature
using Expectations
using ForwardDiff
using Zygote

# Calls to include need to come after `using`
include("UnivariateDensity.jl")
include("discrete.jl")

const derivative = ForwardDiff.derivative
const CUD = ContinuousUnivariateDistribution
const UD = UnivariateDistribution
const E(d::UD, n::Integer=500) = expectation(d; n=n)
const Density = UnivariateDensity

const likelihoodratio(p) = p / (1 - p)
const LR = likelihoodratio
likelihood_ratio_multiplier(x, y) = (likelihoodratio(y) - 1) / (likelihoodratio(y) - likelihoodratio(x))
const LRM = likelihood_ratio_multiplier

# Assume sender chooses a cutoff type in [0,1/2]
const lb = 0.0
const ub = 0.5
# Assume Uniform Distribution for default arguements
const unidist = Uniform(lb, ub)
const f(t) = pdf(unidist, t)
const unidens = Density(f, lb, ub)

# Sender Ex-Post Payoff
function payoff(type::T, cutoff::T) where {T <: Real}
    pay = type + (1 - type) * likelihoodratio(cutoff)
    if cutoff ≤ 1/2 && cutoff ≤ type
        return pay
    elseif cutoff ≤ type
        return one(pay)
    else
        return zero(pay)
    end
end
payoff(type::Real, cutoff::Real) = payoff(promote(type, cutoff)...)

# Methods for two-cutoff case
function payoff(type::T, pL::T, pH::T) where {T <: Real}
    pay = type + LR(pL) * (1 - type)
    if pL ≤ type ≤ pH
        return pay
    else
        return zero(pay)
    end
end
payoff(type::Real, pL::Real, pH::Real) = payoff(promote(type, pL, pH)...)

# Sender Expected Payoffs
payoff_integrand(t, c, f::Function=f) = payoff(t, c) * f(t)
integrated_payoff(cutoff, density::Function=f, ub::Real=ub) =
        quadgk(t -> payoff_integrand(t, cutoff, density), cutoff, ub)
integrated_payoff(cutoff::Real, density::Density=unidens) =
        integrate(t -> payoff(t, cutoff), density, cutoff)

"""
Sender expected payoff as a function of the cutoff and the type distribution
"""
expected_payoff(cutoff::Real, dist::UD=unidist; n::Integer) = E(dist, n)(t -> payoff(t, cutoff))
expected_payoff(cutoff, f::Function=f, ub::Real=ub) = integrated_payoff(cutoff, f, ub)[1]
expected_payoff(cutoff, density::Density) = expected_payoff(cutoff, density.f, density.ub)

# Methods for the two-cutoff case
expected_payoff(pL::T, pH::T, dist::UD; n::Integer) where {T <: Real} = LRM(pL, pH) * E(dist, n)(t -> payoff(t, pL, pH)) + 1 - cdf(dist, pH)
expected_payoff(pL::T, pH::T, f::Function, ub::Real) where {T <: Real} = LRM(pL, pH) * expected_payoff(pL, f, pH) + quadgk(f, pH, ub)[1]
expected_payoff(pL::T, pH::T, density::Density) where {T<: Real} = expected_payoff(pL, pH, density.f, density.ub)

expected_payoff(pL::Real, pH::Real, dist::UD; n::Integer) = expectedpayoff(promote(pL, pH)..., dist; n = n)
expected_payoff(pL::Real, pH::Real, f::Function, ub::Real) = expectedpayoff(promote(pL, pH)..., f, ub)
expected_payoff(pL::Real, pH::Real, density::Density) = expectedpayoff(promote(pL, pH)..., density)

"""
Alias for expected_payoff
"""
V(cutoff::Real, dist::UD=unidist; n::Integer=500) = expected_payoff(cutoff, dist; n = n)
V(cutoff, f::Function, ub::Real) = expected_payoff(cutoff, f, ub)
V(cutoff, density::Density=unidens) = V(cutoff, density.f, density.ub)

V(pL, pH, dist::UD; n=Integer=500) = expected_payoff(pL, pH, dist; n = n)
V(pL, pH, f::Function, ub::Real) = expected_payoff(pL, pH, f, ub)
V(pL, pH, density::Density) = expected_payoff(pL, pH, density)

# Lagrange Multipliers
series_integrand(p::T, t::T) where {T <: Real} = t ≥ p ? 1 - t : zero(t)
series_integrand(p::Real, t::Real) = series_integrand(promote(p, t)...)
series_integral(p, dist::UD, n::Integer=500) = E(dist, n)(t -> series_integrand(p, t))
integral_term(p, dist::UD, n::Integer=500) = series_integral(p, dist, n) / (1 - p)^2

series_integrand(t, f::Function) = (1 - t) * f(t)
series_integral(p, f::Function, ub::Real) = quadgk(t -> series_integrand(t, f), p, ub)
integral_term(p, f::Function, ub::Real) = series_integral(p, f, ub)[1] / (1 - p)^2

# Multiplier for IC constraint
"""
Lagrange multiplier for the IC (reporting) constraint
"""
λ(p::Real, dist::UD=unidist, n::Integer=500) = pdf(dist, p) - integral_term(p, dist, n)
λ(p, f::Function, ub::Real) = f(p) - integral_term(p, f, ub)
λ(p, density::Density=unidens) = λ(p, density.f, density.ub)

upper_integral_λ(p, f::Function, ub::Real) = quadgk(t -> λ(t, f, ub), p, ub)
upper_integral_λ(p, density::Density=unidens) = upper_integral_λ(p, density.f, density.ub)

"""
``\\Lambda (p, dist) = \\int_p^{1/2} \\lambda(t, dist) \\, \\mathrm{d}t``
"""
Λ(p::Real, dist::CUD=unidist, n::Integer=500) =
                            E(Uniform(), n)(t -> t ≥ p ? λ(t, dist, n) : zero(t))
Λ(p, f::Function, ub::Real) = upper_integral_λ(p, f, ub)[1]
Λ(p, density::Density=unidens) = Λ(p, density.f, density.ub)

# Multiplier for bound constraint on π_G
"""
``\\mu`` is the Lagrange multiplier for the constraint that ``\\pi_G`` must be a probability.
"""
μ(p::Real, dist::CUD=unidist, n::Integer=500) = 2p * pdf(dist, p) - integral_term(p, dist, n)
μ(p, f::Function, ub::Real) = 2p * f(p) - integral_term(p, f, ub)
μ(p, density::Density=unidens) = μ(p, density.f, density.ub)

# ∂μ/∂p
# Use 5000 nodes for dμdp or else plot looks terrible
"""
``\\partial \\mu / \\partial p``
"""
dμdp(p::Real, dist::CUD=unidist, n::Integer=500) =
                (2 + 1 / (1 - p)) * pdf(dist, p) +
                2p * derivative(t -> pdf(dist, t), p) +
                2integral_term(p, dist, n) / (1 - p)

dμdp(p, f::Function=f, ub::Real=ub) =
                (2 + 1 / (1 - p)) * f(p) +
                2p * derivative(f, p) +
                2integral_term(p, f, ub) / (1 - p)

dμdp(p, d::Density) = dμdp(p, d.f, d.ub)

LR′(x) = 1 / (1 - x)^2
function dVdpH(pL, pH, dist::CUD; n::Integer=500)
    fpH = pdf(dist, pH)
    int_mul = LR′(pH) * (1 - LR(pL)) / (LR(pH) - LR(pL))^2
    term1 = fpH * (-1 + LRM(pL, pH) * (pH + LR(pL) * (1 - pH)))
    term2 = E(dist, n)(t -> payoff(t, pL, pH))
    return term1 + int_mul * term2
end

function dVdpH(pL, pH, f::Function)
    fpH = f(pH)
    int_mul = LR′(pH) * (1 - LR(pL)) / (LR(pH) - LR(pL))^2
    term1 = fpH * (-1 + LRM(pL, pH) * (pH + LR(pL) * (1 - pH)))
    term2 = quadgk(t -> payoff(t, pL, pH), pL, pH)[1]
    return term1 + int_mul * term2
end

dVdpH(pL, pH, dens::Density) = dVdpH(pL, pH, dens.f)

function dVdpL(pL, pH, dist::CUD; n::Integer=500)
    lrmul = LRM(pL, pH)
    fpL = pdf(dist, pL)
    integrand(x) = pL ≤ x ≤ pH ? 1 - x : zero(x)
    integral1 = E(dist, n)(integrand)
    term1 = -2pL * fpL + LR′(pL) * integral1
    term2 = LR′(pL) * E(dist, n)(t -> payoff(t, pL, pH)) / (LR(pH) - LR(pL))
    return lrmul *(term1 + term2)
end

function dVdpL(pL, pH, f::Function)
    lrmul = LRM(pL, pH)
    fpL = f(pL)
    integrand(x) = pL ≤ x ≤ pH ? 1 - x : zero(x)
    integral1 = quadgk(integrand, pL, pH)[1]
    term1 = -2pL * fpL + LR′(pL) * integral1
    term2 = LR′(pL) * quadgk(t -> payoff(t, pL, pH), pL, pH)[1] / (LR(pH) - LR(pL))
    return lrmul *(term1 + term2)
end

dVdpL(pL, pH, dens::Density) = dVdpL(pL, pH, dens.f)

function ∇expected_payoff!(G, x, dist::CUD; n::Integer=500)
    G[1] = dVdpL(x[1], x[2], dist; n = n)
    G[2] = dVdpH(x[1], x[2], dist; n = n)
end

function ∇expected_payoff!(G, x, f::Function)
    G[1] = dVdpL(x[1], x[2], f)
    G[2] = dVdpH(x[1], x[2], f)
end

# This next function is not necessary, but keeping it helps standardise the
# structure of the code.
∇expected_payoff!(G, x, dens::Density) = ∇expected_payoff!(G, x, dens.f)

function optimise(dist::UD=unidist; n::Integer=5000, initx=[eps(0.0), 1 - eps(1.0)], alg=Brent(), inner_optimizer=BFGS())
    ub = maximum(dist)
    lb = minimum(dist)
    if ub > 1/2 && cdf(dist,ub) - cdf(dist,1/2) > 0
        lower = [lb, 1/2]
        upper = [1/2, ub]
        initx = [max(initx[1], lb + eps(lb)), min(initx[2], ub - eps(ub))]
        gradient!(G, x) = -∇expected_payoff!(G, x, dist; n = n)
        return optimize(x -> -expected_payoff(x[1], x[2], dist; n = n), gradient!, lower, upper, initx, Fminbox(inner_optimizer))
    else
        return optimize(x -> -expected_payoff(x, dist; n = n), lb, ub, alg)
    end
end

function optimise(f::Function, lb::Real, ub::Real; initx=[eps(0.0), 1 - eps(1.0)], alg=Brent(), inner_optimizer=BFGS())
    if ub > 1/2 && quadgk(f, 1/2, ub)[1] > 0
        lower = [lb, 1/2]
        upper = [1/2, ub]
        initx = [max(initx[1], lb + eps(lb)), min(initx[2], ub - eps(ub))]
        gradient!(G, x) = -∇expected_payoff!(G, x, f)
        return optimize(x -> -expected_payoff(x[1], x[2], f, ub), gradient!, lower, upper, initx, Fminbox(inner_optimizer))
    else
        objective(x) = -expected_payoff(x, f, ub)
        return optimize(x -> -expected_payoff(x, f, ub), lb, ub, alg)
    end
end

optimise(density::Density; initx=[eps(0.0), 1 - eps(1.0)], alg=Brent(), inner_optimizer=BFGS()) = optimise(density.f, density.ub; alg=alg)

end # module


using ChebyExp, MatrixDepot

# random 10x10
n = 1000
R = Matrix{Float64}(n,2)
A = Matrix{Float64}(10,10)
for i in 1:n
    rand!(A)
    R[i,1] = norm(expm(A)*expm(-A)-I,2)/norm(expm(A),2)
    R[i,2] = norm(chebyexp(A)*chebyexp(-A)-I,2)/norm(chebyexp(A),2)
end
a,b = mean(R,1), std(R,1)
@printf("""
Random 10x10 matrices:
expm error:     %1.2G ± %1.2G
chebyexp error: %1.2G ± %1.2G
""",a[1],b[1],a[2],b[2])

# Chebyshev and Forsythe
R = Matrix{Float64}(2,2)
for (i,A) in enumerate([matrixdepot("chebspec",10), matrixdepot("forsythe",10)])
    R[i,1] = norm(expm(A)*expm(-A)-I,2)/norm(expm(A),2)
    R[i,2] = norm(chebyexp(A)*chebyexp(-A)-I,2)/norm(chebyexp(A),2)
end
a,b = mean(R,1), std(R,1)
@printf("""
Chebyshev and Forsythe 10x10 matrices:
expm error:     %1.2G ± %1.2G
chebyexp error: %1.2G ± %1.2G
""",a[1],b[1],a[2],b[2])

# Ill-conditioned
R = Matrix{Float64}(length(matrixdepot("ill-cond")),2)
for (i,nameA) in enumerate(matrixdepot("ill-cond"))
    A = matrixdepot(nameA,Float64,16)
    norm(full(A))>5 && (A /= norm(full(A))/4)
    R[i,1] = norm(expm(full(A))*expm(-full(A))-I,2)/norm(expm(full(A)),2)
    R[i,2] = norm(chebyexp(A)*chebyexp(-A)-I,2)/norm(chebyexp(A),2)
end
a,b = mean(R,1), std(R,1)
@printf("""
Ill-conditioned 16x16 matrices:
expm error:     %1.2G ± %1.2G
chebyexp error: %1.2G ± %1.2G
""",a[1],b[1],a[2],b[2])


using Test
using JuMP
using Cbc
using ConditionalJuMP
using Polyhedra
using LinearAlgebra

# In this example, we search for the point closest to the origin in the union
# of a list of convex polyhedra

# This implementation lets us add disjunctions and implications of the form
#   x in P
# where x is a vector and P is a Polyhedron
function ConditionalJuMP.Conditional(op::typeof(in), x::AbstractVector, P::HRepresentation)
    (&)([@?(h.a' * x <= h.β) for h in allhalfspaces(P)]...)
end

# A simple L1-norm objective we can minimize (since the Cbc
function l1norm_objective(x::AbstractVector{Variable})
    model = first(x).m
    y = @variable(model, [1:length(x)], basename="y", lowerbound=0)
    for i in 1:length(x)
        @constraint(model, y[i] >= x[i])
        @constraint(model, y[i] >= -x[i])
    end
    sum(y)
end

# P1 is a box from [-1.5, -1] to [-0.5, 1]
P1 = hrep(vcat(Matrix(I, 2, 2), -Matrix(I, 2, 2)), [-0.5, 1, 1.5, 1])

# P2 is a box from [1, -1] to [2, 1]
P2 = hrep(vcat(Matrix(I, 2, 2), -Matrix(I, 2, 2)), [2., 1, -1, 1])

m = Model(solver=CbcSolver(logLevel=0))
@variable(m, -5 <= x[1:2] <= 5)
@disjunction(m, x in P1, x in P2)
@objective(m, Min, l1norm_objective(x))

@test sum(m.colCat .== :Bin) == 1
solve(m)
x = getvalue(x)
@test x ≈ [-0.5, 0]


"""
    greedy(g)

Create a tour which starts at node 1 and goes to the nearest unvisited neighbor in each step.
Return the tour as a list of nodes and the cost of the tour.
If no tour is possible return `nothing, Inf`. If no tour was found but there isn't a proof that no tour exists
return `nothing, NaN`.
"""
function greedy(g)
    N = nv(g)
    sorted_neighbors = Vector{Vector{Int}}(undef, N)
    for i in 1:N
        vec_neighbors = collect(neighbors(g,i))
        NN = length(vec_neighbors)
        weights = zeros(NN)
        for j in 1:NN
            weights[j] = get_prop(g, i, vec_neighbors[j], :weight)
        end
        order = sortperm(weights)
        sorted_neighbors[i] = vec_neighbors[order]
        p = findfirst(>(0),weights[order])
        # if three edges are fixed
        if p > 3
            return nothing, Inf
        end
        # if both end points are fixed we can fix them here
        if p == 3
            sorted_neighbors[i] = sorted_neighbors[i][1:2]
        end
    end

    visited = zeros(Bool, N)
    node_next = ones(Int, N)
    tour = zeros(Int, N)
    tour[1] = 1
    visited[1] = true
    cp = 1
    backtrack_c = 0
    while 1 <= cp < N
        # simply give up searching :D
        if backtrack_c == 1000000
            return nothing, NaN
        end
        neighbor = next_neighbor!(g, tour, cp, node_next, visited, sorted_neighbors[tour[cp]])
        if neighbor === nothing
            backtrack_c += 1
            node_next[tour[cp]] = 1
            visited[tour[cp]] = false
            cp -= 1
            if cp == 0
                break
            end
        else
            visited[neighbor] = true
            cp += 1
            tour[cp] = neighbor
        end
    end
    return tour, get_tour_cost(g, tour)
end

function next_neighbor!(g, tour, cp, node_next, visited, sorted_neighbors)
    for neighbor_id in node_next[tour[cp]]:length(sorted_neighbors)
        node_next[tour[cp]] += 1
        neighbor = sorted_neighbors[neighbor_id]
        visited[neighbor] && continue
        if (cp == length(tour)-1 && !has_edge(g, neighbor, 1))
            continue
        end

        return sorted_neighbors[neighbor_id]
    end
    return nothing
end

function two_opt(tour, tour_cost, cost_matrix)
    swapped = true
    while swapped
        swapped = false
        for i in 1:length(tour)
            for j in i+1:length(tour)
                swap_cost = get_swap_cost(tour, i, j, cost_matrix)
                if swap_cost < -1e-6
                    tour_cost += swap_cost
                    swap!(tour, i, j)
                    swapped = true
                    break
                end
            end
            swapped && break
        end
    end
    return tour, tour_cost
end

function get_swap_cost(tour, i, j, cost_matrix)
    cost = -cost_matrix[tour[i], tour[mod1(i+1,end)]]
    cost -= cost_matrix[tour[j], tour[mod1(j+1,end)]]

    cost += cost_matrix[tour[i], tour[j]]
    cost += cost_matrix[tour[mod1(i+1,end)], tour[mod1(j+1,end)]]
    return cost
end

function swap!(tour, i, j)
    reverse!(tour, i+1, j)
end


using Makie
using MultipleViewGeometry, Test, Random
using MultipleViewGeometry.ModuleCostFunction
using MultipleViewGeometry.ModuleTypes
using MultipleViewGeometry.ModuleConstraints
using MultipleViewGeometry.ModuleConstruct
using MultipleViewGeometry.ModuleDraw
using MultipleViewGeometry.ModuleMove
using LinearAlgebra
using StaticArrays
using GeometryTypes

# Fix random seed.
Random.seed!(1234)
# Construct two camera matrices and parametrise two planar surfaces.
f = 50
image_width = 640 / 10
image_height = 480 / 10
𝐊₁ = @SMatrix [f 0 0 ;
               0 f 0 ;
               0 0 1 ]
𝐑₁ = SMatrix{3,3,Float64,9}(rotxyz(0, 25*(pi/180), 0))
𝐭₁ = [-30.0, 0.0, -5.0]

𝐊₂ = @SMatrix [f 0 0 ;
               0 f 0 ;
               0 0 1 ]

𝐑₂ = SMatrix{3,3,Float64,9}(rotxyz(0, -25*(pi/180), 0))
𝐭₂ = [30.0, 0.0, 5.0]

# Normals and distance from origin
𝐧₁ = [0.0, 0.0, 1.0]
d₁ = 55.0

# Normals and distance from origin
𝐧₁ = [0.0, 0.0, 1.0]
d₁ = 55.0

𝐧₂ = [0.5, -0.2, 2.0]
d₂ = 145.0

# Generates a random point on the plane centered around a point on the plane
# that is closest to the origin.
function random_points_on_plane(𝐧::AbstractArray, d::Real, extent::Real, N::Int)
    # Generate vector 𝐰 on a plane through the origin with normal vector 𝐧.
    first(𝐧) == 0 ? 𝐰 = cross(𝐧,SVector(1.0,0.0,0.0)) : 𝐰 = cross(𝐧,SVector(0.0,0.0,1.0))
    points = Array{SVector{3,Float64},1}(undef,N)
    for n = 1:N
        # Rotate 𝐰 randomly around the axis 𝐧.
        θ = rand() * 2*pi
        𝐤 = 𝐧 / norm(𝐧)
        𝐯 = 𝐰*cos(θ) + cross(𝐤,𝐰)*sin(θ) + 𝐤*dot(𝐤,𝐰)*(1-cos(θ))
        # Scale the vector so that it lies in the interval [0, extent)
        𝐯 = (rand() * extent) * 𝐯
        # Translate the vector so that it lies on the plane parametrised by 𝐧 and d.
        𝐯 = 𝐯 + d*(𝐧/norm(𝐧)^2)
        points[n] = 𝐯
    end
    points
end

𝒳₁ = random_points_on_plane(𝐧₁,d₁, 20, 50)
𝒳₂ = random_points_on_plane(𝐧₂,d₂, 20, 50)
X₁ = reshape(reinterpret(Float64,𝒳₁),(3,length(𝒳₁)))
X₂ = reshape(reinterpret(Float64,𝒳₂),(3,length(𝒳₂)))

scale = 20.0f0
x = Vec3f0(0); baselen = 0.2f0 * scale ; dirlen = 1f0 * scale
# create an array of differently colored boxes in the direction of the 3 axes
rectangles = [
    (HyperRectangle(Vec3f0(x), Vec3f0(dirlen, baselen, baselen)), RGBAf0(1,0,0,1)),
    (HyperRectangle(Vec3f0(x), Vec3f0(baselen, dirlen, baselen)), RGBAf0(0,1,0,1)),
    (HyperRectangle(Vec3f0(x), Vec3f0(baselen, baselen, dirlen)), RGBAf0(0,0,1,1))
]
meshes = map(GLNormalMesh, rectangles)

scene = mesh(merge(meshes))
scatter!(scene, X₁[1,:],X₁[2,:], X₁[3,:], markersize = 3, color = :red)
scatter!(scene, X₂[1,:],X₂[2,:], X₂[3,:], markersize = 3, color = :blue)

world_basis = (Vec(1.0, 0.0, 0.0), Vec(0.0, 1.0, 0.0), Vec(0.0, 0.0, 1.0))
camera_basis = (Point(0.0, 0.0, 0.0), Vec(-1.0, 0.0, 0.0), Vec(0.0, -1.0, 0.0), Vec(0.0, 0.0, 1.0))
picture_basis = (Point(0.0, 0.0), Vec(-1.0, 0.0), Vec(0.0, -1.0))

camera₁ = Pinhole(image_width, image_height, f, camera_basis..., picture_basis...)
camera₂ = Pinhole(image_width, image_height, f, camera_basis..., picture_basis...)
relocate!(camera₁, 𝐑₁, 𝐭₁)
relocate!(camera₂, 𝐑₂, 𝐭₂)

draw!(camera₁, scene)
draw!(camera₂, scene)

𝐑₁′, 𝐭₁′ = ascertain_pose(camera₁, world_basis... )
𝐊₁′ = obtain_intrinsics(camera₁, CartesianSystem())
𝐑₂′, 𝐭₂′ = ascertain_pose(camera₂, world_basis... )
𝐊₂′ = obtain_intrinsics(camera₂, CartesianSystem())

𝐏₁ = construct(ProjectionMatrix(),𝐊₁′,𝐑₁′,𝐭₁′)
𝐏₂ = construct(ProjectionMatrix(),𝐊₂′,𝐑₂′,𝐭₂′)

# Set of corresponding points.
ℳ₁ = project(camera₁,𝐏₁,𝒳₁)
ℳ₁ʹ= project(camera₂,𝐏₂,𝒳₁)
ℳ₂ = project(camera₁,𝐏₁,𝒳₂)
ℳ₂ʹ= project(camera₂,𝐏₂,𝒳₂)

M₁ = reshape(reinterpret(Float64,ℳ₁),(2,length(ℳ₁)))
M₂ = reshape(reinterpret(Float64,ℳ₂),(2,length(ℳ₂)))
M₁ʹ = reshape(reinterpret(Float64,ℳ₁ʹ),(2,length(ℳ₁ʹ)))
M₂ʹ = reshape(reinterpret(Float64,ℳ₂ʹ),(2,length(ℳ₂ʹ)))

# Visualize the set of corresponding points
scene = Scene()
scatter!(scene,M₁[1,:], M₁[2,:], markersize = 1, color = :red, limits = FRect(0, 0, 64, 48.0))
scatter!(scene,M₂[1,:], M₂[2,:], markersize = 1, color = :blue, limits = FRect(0, 0, 64, 48.0))

scene = Scene()
scatter!(scene,M₁ʹ[1,:], M₁ʹ[2,:], markersize = 1, color = :red, limits = FRect(0, 0, 64, 48.0))
scatter!(scene,M₂ʹ[1,:], M₂ʹ[2,:], markersize = 1, color = :blue, limits = FRect(0, 0, 64, 48.0))

# We will construct a pair of homography matrices and then construct a pair of
# projection matrices from the homographies.

𝐇₁ = construct(HomographyMatrix(),𝐊₁′,𝐑₁′,-𝐭₁′,𝐊₂′,𝐑₂′,-𝐭₂′,𝐧₁,d₁)
𝐇₂ = construct(HomographyMatrix(),𝐊₁′,𝐑₁′,-𝐭₁′,𝐊₂′,𝐑₂′,-𝐭₂′,𝐧₂,d₂)

context = ProjectionMatrices(HomographyMatrices(), Chojnacki(), TwoViews())
𝐎₁, 𝐎₂ = construct(context, (𝐇₁,𝐇₂))
𝐅₁ = construct(FundamentalMatrix(), 𝐏₁, 𝐏₂)
𝐅₂ = construct(FundamentalMatrix(), 𝐎₁, 𝐎₂)

𝐅₁ = 𝐅₁ / 𝐅₁[3,3]
𝐅₂ = 𝐅₂ / 𝐅₂[3,3]

# Use homographies to determine corresponding coordinates in the second image.
𝒪₁ = similar(ℳ₁ʹ)
𝒪₂ = similar(ℳ₂ʹ)
for n = 1:length(ℳ₁)
    𝐦 = hom(ℳ₁[n])
    𝐦ʹ = 𝐇₁*𝐦
    𝒪₁[n] = hom⁻¹(𝐦ʹ)
end
for n = 1:length(ℳ₂)
    𝐦 = hom(ℳ₂[n])
    𝐦ʹ = 𝐇₂*𝐦
    𝒪₂[n] = hom⁻¹(𝐦ʹ)
end

#TODO Test that 𝒪₁ .== ℳ₁ʹ and 𝒪₂ .== M₂ʹ


import SpecialFunctions: trigamma

"""
    OhmicBath

Ohmic bath object to hold a particular parameter set.

**Fields**
- `η` -- strength.
- `ωc` -- cutoff frequence.
- `β` -- inverse temperature.
"""
struct OhmicBath <: AbstractBath
    η::Float64
    ωc::Float64
    β::Float64
end

"""
$(SIGNATURES)

Construct OhmicBath from parameters with physical unit: `η`--unitless interaction strength; `fc`--cutoff frequency in GHz; `T`--temperature in mK.
"""
function Ohmic(η, fc, T)
    ωc = 2 * π * fc
    β = temperature_2_β(T)
    OhmicBath(η, ωc, β)
end

"""
$(SIGNATURES)

Calculate spectral density ``γ(ω)`` of `bath`.
"""
function spectrum(ω, bath::OhmicBath)
    if isapprox(ω, 0.0, atol=1e-9)
        return 2 * pi * bath.η / bath.β
    else
        return 2 * pi * bath.η * ω * exp(-abs(ω) / bath.ωc) / (1 - exp(-bath.β * ω))
    end
end

"""
$(SIGNATURES)

Calculate the two point correlation function ``C(τ)`` of `bath`.
"""
function correlation(τ, bath::OhmicBath)
    x2 = 1 / bath.β / bath.ωc
    x1 = 1.0im * τ / bath.β
    bath.η * (trigamma(-x1 + 1 + x2) + trigamma(x1 + x2)) / bath.β^2
end

"""
$(SIGNATURES)

Calculate the polaron transformed correlation function of `bath`.
"""
function polaron_correlation(τ, bath::OhmicBath)
    res = (1 + 1.0im * bath.ωc * τ)^(-4 * bath.η)
    if !isapprox(τ, 0, atol=1e-9)
        x = π * τ / bath.β
        res *= (x / sinh(x))^(4 * bath.η)
    end
    res
end

@inline polaron_correlation(t1, t2, bath::OhmicBath) = polaron_correlation(t1 - t2, bath)

function info_freq(bath::OhmicBath)
    println("ωc (GHz): ", bath.ωc / pi / 2)
    println("T (GHz): ", temperature_2_freq(β_2_temperature(bath.β)))
end

function Base.show(io::IO, ::MIME"text/plain", m::OhmicBath)
    print(
        io,
        "Ohmic bath instance:\n",
        "η (unitless): ",
        m.η,
        "\n",
        "ωc (GHz): ",
        m.ωc / pi / 2,
        "\n",
        "T (mK): ",
        β_2_temperature(m.β),
    )
end


"""
$(DocStringExtensions.README)
"""
module DataDrivenDiffEq

using DocStringExtensions
using LinearAlgebra
using DiffEqBase
using ModelingToolkit
using QuadGK
using Statistics
using DSP
using FiniteDifferences, DataInterpolations
using Compat
using DocStringExtensions

abstract type AbstractKoopmanOperator <: Function end;

include("./optimizers/Optimize.jl")
using .Optimize

export set_threshold!, set_threshold
export STRRidge, ADMM, SR3

export ADM

include("./basis.jl")
export Basis
export dynamics
export free_parameters

include("./koopman/algorithms.jl")
export DMDPINV, DMDSVD, TOTALDMD

include("./koopman/koopman.jl")
export eigen, eigvals, eigvecs
export modes, frequencies
export is_discrete, is_continuous
export operator, generator
export inputmap, outputmap, updatable, isstable

include("./koopman/linearkoopman.jl")
export LinearKoopman, update!

include("./koopman/nonlinearkoopman.jl")
export NonlinearKoopman, reduce_basis

include("./koopman/exact_dmd.jl")
export DMD, gDMD

include("./koopman/dmdc.jl")
export DMDc, gDMDc

include("./koopman/extended_dmd.jl")
export EDMD, gEDMD

include("./sindy/results.jl")
export SparseIdentificationResult
export print_equations
export get_coefficients, get_error, get_sparsity, get_aicc

include("./sindy/sindy.jl")
export SINDy
export sparse_regression, sparse_regression!

function SInDy(Y, X, basis; opt = STRRidge(), kwargs...)
    @warn("SInDy has been deprecated. Use SINDy to recover the same functionality.")
    SINDy(Y, X, basis, opt; kwargs...)
end

function ISInDy(Y, X, basis; opt = ADM(), kwargs...)
    @warn("ISInDy has been deprecated. Use ISINDy to recover the same functionality.")
    ISINDy(Y, X, basis, opt; kwargs...)
end

export SInDy, ISInDy

include("./sindy/isindy.jl")
export ISINDy

include("./system_conversions.jl")

include("./utils.jl")
export AIC, AICC, BIC
export optimal_shrinkage, optimal_shrinkage!
export savitzky_golay
export burst_sampling, subsample

include("./basis_generators.jl")
export chebyshev_basis, monomial_basis, polynomial_basis
export sin_basis, cos_basis, fourier_basis

end # module


module RandomMatrices

using Random

export unitvector, orthogonal, normal, general, symmetric, hermitian

using Random: GLOBAL_RNG
using LinearAlgebra:  reflector!, QR

import Base.randn, Base.rand

"""
unitvector(m, T::Type [, rng::AbstractRNG]) -> Vector{T}

  create random unit vector, uniformly distributed on unit sphere of `m` dimensions.
  `T` may be any real or complex float type.
"""
function unitvector(m::Integer, S::Union{Type{T},Type{Complex{T}}}=Float64, rng::AbstractRNG = GLOBAL_RNG) where T<:AbstractFloat

  unitvector!(zeros(S, m), rng)
end

"""
  unitvector!(v::AbstractVector{T} [, rng::AbstractRNG]) -> v

  Same as unitvector, populating the given vector or 1-dim view `v`.
"""
function unitvector!(v::Union{AbstractVector{T},AbstractVector{Complex{T}}}, rng::AbstractRNG = GLOBAL_RNG) where T<:AbstractFloat

  n = length(v)
  r = zero(T)
  sqnh = sqrt(n) / 2
  while r < sqnh
    randn!(rng, v)
    r = norm(v)
  end
  normalize!(v)
end

"""
  orthogonal(m, n, ::Type{T}=Float64 [, rng::AbstractRNG]) -> LinearAlgebra.QRCompactWYQ

Generate random orthogonal real or unitary complex matrix.
"""
function orthogonal(m::Integer, n::Integer, S::Type=Float64, rng::AbstractRNG=GLOBAL_RNG)
  S = S == Complex ? Complex{Float64} : S
  S <: Union{T,Complex{T}} where T <: AbstractFloat || error("illegal type $S")
  m >= n || error("more independent vectors($n) than dimension($m) requested")
  A = zeros(S, m, n)
  τ = zeros(S, min(m,n))

  n2 = n == m && (S <: Real) && rand(Bool) ? n - 1 : n
  for k = 1:n2
    x = view(A, k:m, k)
    unitvector!(x)
    τk = reflector!(x)
    τ[k] = τk
  end
  QR(A, τ)[:Q]
end

# modify rand and randn to produce BigFloat and complex values
function rand(rng::AbstractRNG, ::Type{BigFloat})
  bigFloat_randfill(rng, rand(rng, Float64))
end
function randn(rng::AbstractRNG, ::Type{BigFloat})
  bigFloat_randfill(rng, randn(rng, Float64))
end
function rand(rng::AbstractRNG, ::Type{Complex{T}}) where T<:AbstractFloat
  a, b = 2rand(T, 2) - 1
  while hypot(a,b) > 1
    a, b = 2rand(T, 2) - 1
  end
  complex(a, b)
end

"""
Append random bits in BigFloat representation of input small float number.
"""
function bigFloat_randfill(rng::AbstractRNG, a::T) where {T<:Union{Float64,Float32,Float16}}
  missing = precision(BigFloat) - precision(a)
  bigi = BigInt(1) << missing
  add = ( rand(rng, BigInt(0):bigi-1) / bigi )
  add *= big"2.0"^(exponent(a) - precision(a))
  add = copysign(add, a)
  a + add
end

"""
  diagonal(dia::AbstractVector{T}) -> Array{T,2}

Create quasi diagonal real matrix with diagonal blocks given by `dia`.
`dia` must contain only real and pairs of conjugate complex numbers.
If complex output matrix is required, use `diagm(dia)`.
"""
function diagonal(dia::Union{AbstractVector{T}, AbstractVector{Complex{T}}}) where T<:AbstractFloat

  m = length(dia)
  mr = m
  j = 1
  while j <= m
    a = dia[j]
    if !isreal(a)
      if j >= m || imag(a) != -imag(dia[j+1])
        throw(InexactError())
      else
        j += 1
      end
    end
    j += 1
  end

  A = zeros(T, mr, mr)

  k = 1
  while k <= m
    dr = real(dia[k])
    di = imag(dia[k])
    if di == 0
      A[k,k] = dr
      k += 1
    else
      A[k,k] = A[k+1,k+1] = dr
      A[k+1,k] = di
      A[k,k+1] = -di
      k += 2
    end
  end
  A
end

"""
  normal(dia::AbstractVector, [,rng]) -> Array{T,2}

Create real or complex normal matrix with eigenvalues given by dia.
"""
function normal(dia::Union{AbstractVector{T}, AbstractVector{Complex{T}}}, S::Union{Type{T},Type{Complex{T}}}=Float64, rng::AbstractRNG = GLOBAL_RNG) where T<:AbstractFloat

  D = S<:Real ? diagonal(dia) : diagm(dia)
  m = size(D, 1)
  Q = orthogonal(m, m, S, rng)
  Q * D * Q'
end

"""
  general(dia::Vector{<:AbstractFloat}, m, n, S::Type=Float64, [,rng]) -> Array{S,2}

Create general real or complex (m,n)-matrix with singular values given.

"""
function general(singular_values::AbstractVector{T}, m::Int, n::Int, S::Union{Type{T},Type{Complex{T}}}=Float64, rng::AbstractRNG=GLOBAL_RNG) where T<:AbstractFloat

  j = length(singular_values)
  j <= min(m, n) || error("too many singular values")
  all(singular_values .> 0) || error("singular values must be positive")

  D = zeros(S, m, n)
  for k = 1:j
    D[k,k] = singular_values[k]
  end
  orthogonal(m, m, S, rng) * D * orthogonal(n, n, S, rng)'
end

end # module


# function poiseuille(y::Real, canal::CanalConfig, μ, dpdx)
#     h = canal.H
#     vmax = -1 / μ * dpdx * h^2 / 2
#     vmax * (1 - (y / h)^2)
# end

# poiseuille(y::AbstractVector, c::CanalConfig, μ, dpdx) = poiseuille.(y, Ref(c), Ref(μ), Ref(dpdx))

function poiseuille(y::Real, gb::GridBox, Q)
    uc = 1.5 * Q / gb.D
    uc * (1 - (y / (0.5*gb.D))^2)
end

poiseuille(y::AbstractVector, c::GridBox, Q) = poiseuille.(y, Ref(c), Ref(Q))

function poiseuille!(grid::AbstractGrid, c::GridBox, Q::Real)
    ys = getys(grid)
    for i in eachindex(grid)
        grid[i] = poiseuille(ys[i[2]], c, Q)
    end
end

# poiseuille(y::Real, canal::CanalConfig, params::SimParam) = poiseuille(y, canal, params.Q)
poiseuille!(y::AbstractGrid, gb::GridBox, params::SimParam) = poiseuille!(y, gb, params.Q)

poiseuille!(sim::Simulation) = poiseuille!(sim.u, sim.u.grid, sim.params)
poiseuille(sim::Simulation, y) = poiseuille(y, sim.u.grid, sim.params.Q)


"""
```
segments                = meanshift(img, spatial_radius, range_radius; iter=50, eps=0.01))
```
Segments the image using meanshift clustering. Returns a `SegmentedImage`.

Parameters:
-    img                            = input grayscale image
-    spatial_radius/range_radius    = bandwidth parameters of truncated normal kernel. Controlling the size of the kernel determines the resolution of the mode detection.
-    iter/eps                       = stopping criterion for meanshift procedure. The algorithm stops after iter iterations or if kernel center moves less than eps distance in an update step, whichever comes first.
"""
function meanshift(img::Array{CT, 2}, spatial_radius::Real, range_radius::Real; iter::Int = 50, eps::Real = 0.01) where CT

    rows, cols = size(img)
    rowbins = Int(floor(rows/spatial_radius))
    colbins = Int(floor(cols/spatial_radius))
    colorbins = Int(floor(1/range_radius))

    colorbin(x)::Int = min(floor(x/range_radius) + 1, colorbins)
    rowbin(x)::Int = min(floor(x/spatial_radius) +1, rowbins)
    colbin(x)::Int = min(floor(x/spatial_radius) +1, colbins)

    buckets = Array{Vector{CartesianIndex{2}}}(undef, rowbins, colbins, colorbins)
    for i in CartesianIndices(size(buckets))
        buckets[i]=Array{CartesianIndex{2}}(undef, 0)
    end

    for i in CartesianIndices(size(img))
        push!( buckets[rowbin(i[1]), colbin(i[2]), colorbin(img[i])], i)
    end

    function dist2(a::SVector, b::SVector)::Float64
        return ((a[1]-b[1])/spatial_radius)^2 + ((a[2]-b[2])/spatial_radius)^2 + ((a[3]-b[3])/range_radius)^2
    end

    function neighborhood_mean(pt::SVector{3, Float64})::SVector{3, Float64}
        den = 0.0
        num = SVector(0.0, 0.0, 0.0)

        R = CartesianIndices(size(buckets))
        I1, Iend = first(R), last(R)
        I = CartesianIndex(rowbin(pt[1]), colbin(pt[2]), colorbin(pt[3]))

        for J in _colon(max(I1, I-_oneunit(I)), min(Iend, I+_oneunit(I)))
            for point in buckets[J]
                if dist2(pt, SVector(point[1], point[2], img[point])) <= 1
                    num += SVector(point[1], point[2], img[point])
                    den += 1
                end
            end
        end

        return den<=0 ? pt : num/den
    end

    modes = Array{Float64}(undef, 3, rows*cols)

    for i in CartesianIndices(size(img))
        pt::SVector{3, Float64} = SVector(i[1], i[2], img[i])
        for j in 1:iter
            nextpt = neighborhood_mean(pt)
            if dist2(pt, nextpt) < eps^2
                break
            end
            pt = nextpt
        end
        modes[:, i[1] + (i[2]-1)*rows] = [pt[1]/spatial_radius, pt[2]/spatial_radius, pt[3]/range_radius]
    end

    clusters = dbscan(modes, 1.414)
    num_segments = length(clusters)
    TM = meantype(CT)
    result              = similar(img, Int)
    labels              = Array(1:num_segments)
    region_means        = Dict{Int, TM}()
    region_pix_count    = Dict{Int, Int}()

    cluster_idx = 0
    for cluster in clusters
        cluster_idx += 1
        region_pix_count[cluster_idx] = cluster.size
        for index in [cluster.core_indices; cluster.boundary_indices]
            i, j = (index-1)%rows + 1, floor(Int, (index-1)/rows) + 1
            result[i, j] = cluster_idx
            region_means[cluster_idx] = get(region_means, cluster_idx, zero(TM)) + img[i, j]
        end
        region_means[cluster_idx] /= region_pix_count[cluster_idx]
    end

    return SegmentedImage(result, labels, region_means, region_pix_count)
end


using SymPy
if VERSION >= v"0.5.0-dev+7720"
    using Base.Test
else
    using BaseTestNext
    const Test = BaseTestNext
end

@testset "Logical" begin
    ## More logical expressions
    @vars x y
    (x ≪ 0) ∧ Le(x*y,1) ∨ (x ⩵ y) ∧ (¬(x ≫ 3))
end


__precompile__()
module LESFilter

using InplaceRealFFT

export rfftfreq, fftfreq, rfftfreqn, fftfreqn, lesfilter!, lesfilter

function rfftfreq(n::Integer,s::Real)
  d = 2π*s
  Float64[(n/2 - i)/d for i = n/2:-1:0]
end

function fftfreq(n::Integer,s::Real)
  d = 2π*s
  if iseven(n)
    return vcat(Float64[(n/2 - i)/(d) for i = n/2:-1:1],Float64[-i/(d) for i = n/2:-1:1])
  else return vcat(Float64[(n/2 - i)/(d) for i = n/2:-1:0],Float64[-i/(d) for i = (n-1)/2:-1:1])
  end
end

function rfftfreqn(n::Integer,d::Real)
  Float64[(n/2 - i)/d for i = n/2:-1:0]
end

function fftfreqn(n::Integer,d::Real)
  if iseven(n)
    return vcat(Float64[(n/2 - i)/(d) for i = n/2:-1:1],Float64[-i/(d) for i = n/2:-1:1])
  else return vcat(Float64[(n/2 - i)/(d) for i = n/2:-1:0],Float64[-i/(d) for i = (n-1)/2:-1:1])
  end
end

function lesfilter(field::AbstractArray{<:Real,3} , fil::String, boxdim::Real,lengths::Tuple{Real,Real,Real})
  nx,ny,nz = size(field)
  xs,ys,zs = lengths

  fieldhat = rfft(field)

  kx2 = rfftfreq(nx,xs).^2
  ky2 = fftfreq(ny,ys).^2
  kz2 = fftfreq(nz,zs).^2

  loopgaussian!(fieldhat,kx2,ky2,kz2,boxdim)

  return irfft(fieldhat,nx)
end

function loopgaussian!(fieldhat::AbstractArray{<:Complex,3},kx2::AbstractVector,ky2::AbstractVector,kz2::AbstractVector,boxdim::Real)
  aux = -((π*boxdim)^2)/6
  Threads.@threads for k = 1:length(kz2)
    for j = 1:length(ky2)
      @simd for i = 1:length(kx2)
        @inbounds fieldhat[i,j,k] = fieldhat[i,j,k]*exp((kx2[i]+ky2[j]+kz2[k])*aux)
      end
    end
  end
  nothing
end

function loopcutoff!(fieldhat::AbstractArray{<:Complex,3},kx2::AbstractVector,ky2::AbstractVector,kz2::AbstractVector,boxdim::Real)
  aux = ((1/2boxdim)^2)
  Threads.@threads for k = 1:length(kz2)
    for j = 1:length(ky2)
      @simd for i = 1:length(kx2)
        @inbounds fieldhat[i,j,k] = fieldhat[i,j,k]*((kx2[i]+ky2[j]+kz2[k]) < aux)
      end
    end
  end
  nothing
end

function loopbox!(fieldhat::AbstractArray{<:Complex,3},kx2::AbstractVector,ky2::AbstractVector,kz2::AbstractVector,boxdim::Real)
  aux = π*boxdim
  a = fieldhat[1]
  Threads.@threads for k = 1:length(kz2)
    for j = 1:length(ky2)
      @fastmath @simd for i = 1:length(kx2)
        mk = sqrt(kx2[i]+ky2[j]+kz2[k])
        @inbounds fieldhat[i,j,k] = fieldhat[i,j,k]*sinpi(mk*boxdim)/(mk*aux)
      end
    end
  end
  fieldhat[1] = a
  nothing
end

function loopanigaussian!(fieldhat::AbstractArray{<:Complex,3},kx2::AbstractVector,ky2::AbstractVector,kz2::AbstractVector,boxdimxy::Real,boxdimz::Real)
  auxxy = -((π*boxdimxy)^2)/6
  auxz = -((π*boxdimz)^2)/6
  Threads.@threads for k = 1:length(kz2)
    for j = 1:length(ky2)
      @simd for i = 1:length(kx2)
        @inbounds fieldhat[i,j,k] = fieldhat[i,j,k]*exp((kx2[i]+ky2[j])*auxxy + kz2[k]*auxz)
      end
    end
  end
  nothing
end

function lesfilter(field::AbstractPaddedArray{<:Real,3} , fil::String, boxdim::Real,lengths::NTuple{3,Real})
  newfield = copy(field)

  lesfilter!(newfield,fil=fil,boxdim=boxdim,lengths=lengths)

  return newfield
end

function lesfilter!(field::AbstractPaddedArray{<:Real,3}, fil::String, boxdim::Real, lengths::NTuple{3,Real})
  nx,ny,nz = size(real(field))
  xs,ys,zs = lengths

  fieldhat = complex(rfft!(field))

  kx2 = rfftfreq(nx,xs).^2
  ky2 = fftfreq(ny,ys).^2
  kz2 = fftfreq(nz,zs).^2

  if fil == "G"
    loopgaussian!(fieldhat,kx2,ky2,kz2,boxdim)
  elseif fil == "C"
    loopcutoff!(fieldhat,kx2,ky2,kz2,boxdim)
  elseif fil == "B"
    loopbox!(fieldhat,kx2,ky2,kz2,boxdim)
  end

  return irfft!(field)
end

function lesfilter!(field::AbstractPaddedArray{<:Real,3}, fil::String, boxdims::NTuple{2,Real}, lengths::NTuple{3,Real})
  nx,ny,nz = size(real(field))
  xs,ys,zs = lengths

  fieldhat = complex(rfft!(field))

  kx2 = rfftfreq(nx,xs).^2
  ky2 = fftfreq(ny,ys).^2
  kz2 = fftfreq(nz,zs).^2

  if fil == "G"
    loopanigaussian!(fieldhat,kx2,ky2,kz2,boxdims[1],boxdims[2])
  elseif fil == "C"
    #loopanicutoff!(fieldhat,kx2,ky2,kz2,boxdim,boxdimz)
  elseif fil == "B"
    #loopanibox!(fieldhat,kx2,ky2,kz2,boxdim,boxdimz)
  end
  return irfft!(field)
end

end # module


using Plexifications
using Base.Test: @test, @test_throws

@test_throws ErrorException begin
    inv(BiPerplex(0, 0.0))
end

@test begin
    !iszerodivisor(BiPerplex(0,1,2,0))
end

@test begin
    iszerodivisor(BiPerplex(1,2,3,4))
end

@test begin
    a = Plexifications.randomBigFloat()
    isreal(BiPerplex(a))
end

@test begin
    a = 1
    !isreal(BiPerplex(a, a))
end

@test begin
    z = random(BiPerplex{Int})
    z == +(z)
end

@test begin
    a = Plexifications.randomBigFloat()
    real(BiPerplex(a)) == a
end

@test begin
    io = IOBuffer()
    show(io, BiPerplex(1, 2, 3, 4))
    l = String(take!(io))
    r = "[1: 1, R: 2, Q: 3, QR: 4]"
    l == r
end

@test begin
    x = random(BiPerplex{BigFloat})
    y = random(BiPerplex{BigFloat})
    x != y
end

@test begin
    a = Plexifications.randomBigInt()
    z = random(BiPerplex{BigInt})
    l = a + z
    r = z + a
    l == r
end

@test begin
    a = Plexifications.randomBigInt()
    z = random(BiPerplex{BigInt})
    l = a - z
    r = -(z - a)
    l == r
end

@test begin
    x = random(BiPerplex{BigInt})
    y = random(BiPerplex{BigInt})
    l = x + y
    r = y + x
    l == r
end

@test begin
    x = random(Prplex{BigInt})
    y = random(BiPerplex{BigInt})
    l = x + y
    r = y + x
    l == r
end

@test begin
    x = random(BiPerplex{BigInt})
    y = random(BiPerplex{BigInt})
    l = x * y
    r = y * x
    l == r
end

@test begin
    x = random(Prplex{BigInt})
    y = random(BiPerplex{BigInt})
    l = x * y
    r = y * x
    l == r
end

@test begin
    x = random(BiPerplex{BigInt})
    y = random(BiPerplex{BigInt})
    l = x - y
    r = -(y - x)
    l == r
end

@test begin
    x = random(Prplex{BigInt})
    y = random(BiPerplex{BigInt})
    l = x - y
    r = -(y - x)
    l == r
end

@test begin
    x = random(BiPerplex{BigInt})
    y = random(BiPerplex{BigInt})
    z = random(BiPerplex{BigInt})
    l = (x + y) + z
    r = x + (y + z)
    l == r
end

@test begin
    x = random(BiPerplex{BigInt})
    y = random(BiPerplex{BigInt})
    z = random(BiPerplex{BigInt})
    l = (x * y) * z
    r = x * (y * z)
    l == r
end

@test begin
    x = random(BiPerplex{BigInt})
    barR(barR(x)) == x
end

@test begin
    x = random(BiPerplex{BigInt})
    barQ(barQ(x)) == x
end

@test begin
    x = random(BiPerplex{BigInt})
    barQ(barR(x)) == barR(barQ(x))
end

@test begin
    x = random(BiPerplex{Rational{BigInt}})
    inv(inv(x)) == x
end

@test begin
    x = random(BiPerplex{BigInt})
    y = random(BiPerplex{BigInt})
    l = barR(x * y)
    r = barR(x) * barR(y)
    l == r
end

@test begin
    x = random(BiPerplex{BigInt})
    y = random(BiPerplex{BigInt})
    l = barR(x + y)
    r = barR(x) + barR(y)
    l == r
end

@test begin
    x = random(BiPerplex{BigInt})
    y = random(BiPerplex{BigInt})
    l = barR(x - y)
    r = barR(x) - barR(y)
    l == r
end

@test begin
    x = random(BiPerplex{BigInt})
    y = random(BiPerplex{BigInt})
    l = barQ(x * y)
    r = barQ(x) * barQ(y)
    l == r
end

@test begin
    x = random(BiPerplex{BigInt})
    y = random(BiPerplex{BigInt})
    l = barQ(x + y)
    r = barQ(x) + barQ(y)
    l == r
end

@test begin
    x = random(BiPerplex{BigInt})
    y = random(BiPerplex{BigInt})
    l = barQ(x - y)
    r = barQ(x) - barQ(y)
    l == r
end

@test begin
    x = random(BiPerplex{Rational{BigInt}})
    y = random(BiPerplex{Rational{BigInt}})
    l = inv(x * y)
    r = inv(y) * inv(x)
    l == r
end

@test begin
    x = random(BiPerplex{Rational{BigInt}})
    y = random(BiPerplex{Rational{BigInt}})
    l = inv(x / y)
    r = y / x
    l == r
end

@test begin
    x = random(BiPerplex{Rational{BigInt}})
    y = random(BiPerplex{Rational{BigInt}})
    l = inv(x \ y)
    r = y \ x
    l == r
end

@test begin
    x = random(BiPerplex{BigInt})
    y = random(BiPerplex{BigInt})
    z = random(BiPerplex{BigInt})
    l = x * (y + z)
    r = (x * y) + (x * z)
    l == r
end

@test begin
    a = Plexifications.randomBigInt()
    x = random(BiPerplex{BigInt})
    y = random(BiPerplex{BigInt})
    l = a * (x + y)
    r = (a * x) + (a * y)
    l == r
end

@test begin
    a = Plexifications.randomBigInt()
    x = random(BiPerplex{BigInt})
    l = a * x
    r = x * a
    l == r
end

@test begin
    x = random(BiPerplex{BigInt})
    y = random(BiPerplex{BigInt})
    z = random(BiPerplex{BigInt})
    l = (x + y) * z
    r = (x * z) + (y * z)
    l == r
end

@test begin
    x = random(BiPerplex{BigInt})
    y = random(BiPerplex{BigInt})
    z = random(BiPerplex{BigInt})
    l = x * (y - z)
    r = (x * y) - (x * z)
    l == r
end

@test begin
    x = random(BiPerplex{BigInt})
    y = random(BiPerplex{BigInt})
    z = random(BiPerplex{BigInt})
    l = (x - y) * z
    r = (x * z) - (y * z)
    l == r
end

@test begin
    x = random(BiPerplex{BigInt})
    n = norm(x)
    n > big(0) || n < big(0) || n == big(0)
end

@test begin
    x = random(BiPerplex{BigInt})
    y = random(BiPerplex{BigInt})
    l = norm(x * y)
    r = norm(x) * norm(y)
    l == r
end

@test_throws ErrorException begin
    inv(BiPerplex(0, 0.0))
end

@test_throws ErrorException begin
    random(BiPerplex{Int}) / 0
end

@test_throws ErrorException begin
    0 \ random(BiPerplex{Int})
end

@test begin
    a = Plexifications.randomBigRational()
    z = random(BiPerplex{Rational{BigInt}})
    l = z / a
    r = a \ z
    l == r
end

@test begin
    a = Plexifications.randomBigRational()
    z = random(BiPerplex{Rational{BigInt}})
    l = a / z
    r = z \ a
    l == r
end

@test begin
    z = random(BiPerplex{BigInt})
    z == z + zero(BiPerplex{BigInt})
end

@test begin
    z = random(BiPerplex{BigInt})
    z == z * one(BiPerplex{BigInt})
end

@test begin
    z = random(BiPerplex{BigInt})
    z == z + zero(z)
end

@test begin
    z = random(BiPerplex{BigInt})
    z == z * one(z)
end

@test begin
    x = random(BiPerplex{BigInt})
    y = random(BiPerplex{BigInt})
    l = commutator(x, y)
    r = zero(BiPerplex{BigInt})
    l == r
end

@test begin
    x = random(BiPerplex{BigInt})
    y = random(BiPerplex{BigInt})
    l = alternatorL(x, y)
    r = zero(BiPerplex{BigInt})
    l == r
end

@test begin
    x = random(BiPerplex{BigInt})
    y = random(BiPerplex{BigInt})
    l = alternatorR(x, y)
    r = zero(BiPerplex{BigInt})
    l == r
end

@test begin
    x = random(BiPerplex{BigInt})
    y = random(BiPerplex{BigInt})
    l = flexator(x, y)
    r = zero(BiPerplex{BigInt})
    l == r
end

@test begin
    x = random(BiPerplex{BigInt})
    y = random(BiPerplex{BigInt})
    z = random(BiPerplex{BigInt})
    l = associator(x, y, z)
    r = zero(BiPerplex{BigInt})
    l == r
end

@test begin
    x = random(BiPerplex{BigInt})
    y = random(BiPerplex{BigInt})
    z = random(BiPerplex{BigInt})
    l = jacobiator(x, y, z)
    r = 3 * x * y * z
    l == r
end

@test begin
    l = BiPerplex(Prplex(1,2))
    r = BiPerplex(1,2,0)
    l == r
end

@test begin
    l = BiPerplex(1, 2, 3.0)
    r = BiPerplex(1.0, 2.0, 3.0)
    l == r
end

@test begin
    l = BiPerplex(1, 2, 3, 4.0)
    r = BiPerplex(1.0, 2.0, 3.0, 4.0)
    l == r
end


using RenyiEntropyOfMagic

N = parse(Int, ARGS[1])
prepare_states(N, 0.1:0.1:2.5)


using ADCME
using PyCall
using LinearAlgebra
using PyPlot
using Random
using DelimitedFiles
using SparseArrays
Random.seed!(233)

mpi_init()
sp = sprand(10,10,0.3)
SP = mpi_SparseTensor(sp)
SPt = SP'
sess = Session(); init(sess)
Array(run(sess, SP))-Array(run(sess, SPt))'


function flux(u, q, w, v, x, t, mu, eta)
    z = x[1];
    r = x[2];
    f = mu[1]*r*q;
    return f;
end
function source(u, q, w, v, x, t, mu, eta)
    z = x[1];
    r = x[2];
    s = sin(r)/exp(z);
    return s;
end
function ubou(u, q, w, v, x, t, mu, eta, uhat, n, tau)
    z = x[1];
    r = x[2];
    uexact = exp(-z)*cos(r);
    ub = [u uexact uexact uexact];
    return ub;
end
function fbou(u, q, w, v, x, t, mu, eta, uhat, n, tau)
    f = flux(u, q, w, v, x, t, mu, eta);
    tm = f[1]*n[1] + f[2]*n[2] + tau[1]*(u[1]-uhat[1]);
    fb = [0.0 tm tm tm];
    return fb;
end
function initu(x, mu, eta)
    u0 = 0.0;
    return u0;
end


using StanSample, Test

ProjDir = @__DIR__
cd(ProjDir)

gq = "
  data {
    int<lower=0> N;
    int<lower=0> y[N];
  }
  parameters {
    real<lower=0.00001> Theta;
  }
  model {
    Theta ~ gamma(4, 1000);
    for (n in 1:N)
      y[n] ~ exponential(Theta);
  }
  generated quantities{
    real y_pred;
    y_pred = exponential_rng(Theta);
  }
";

gq_data = Dict(
  "N" => 3,
  "y" => [100, 950, 450]
);

stanmodel = SampleModel("Generate_quantities", gq);

(sample_file, log_file) = stan_sample(stanmodel; data=gq_data)

if !(sample_file == Nothing)
  stan_generate_quantities(stanmodel, 1)
end


@testset "Second order" begin
    @testset "Bits $i" for i in 3:5
        x = 0:1 << i
        y_approx = fpcos.(Int16.(x), Val(i))
        y = cos.(x / 1 << i * π / 2) .* 1 << i
        @test all(abs.(y_approx .- y) .< 2.0)
    end
end


# active information functions load correctly
@test Inform._active_info_found
@test Inform._active_info != C_NULL
@test isa(Inform._active_info, Ptr{Void})

@test Inform._active_info_ensemble_found
@test Inform._active_info_ensemble != C_NULL
@test isa(Inform._active_info_ensemble, Ptr{Void})

## Individual time series

# series too short
@test isnan(activeinfo([],2))
@test isnan(activeinfo([1],2))

# history too long
@test isnan(activeinfo([1,1],2))
@test !isnan(activeinfo([1,1,0],2))

# encoding error
let series = [2,1,0,0,1,0,0,1]
    @test !isnan(activeinfo(series,2,b=3))
    @test isnan(activeinfo(series,2,b=2))
end

@test_approx_eq_eps activeinfo([1,1,0,0,1,0,0,1],2) 0.918296 1e-6

@test_approx_eq_eps activeinfo([1,0,0,0,0,0,0,0,0], 2) 0.000000 1e-6
@test_approx_eq_eps activeinfo([0,0,1,1,1,1,0,0,0], 2) 0.305958 1e-6
@test_approx_eq_eps activeinfo([1,0,0,0,0,0,0,1,1], 2) 0.347458 1e-6
@test_approx_eq_eps activeinfo([0,0,0,0,0,1,1,0,0], 2) 0.399533 1e-6
@test_approx_eq_eps activeinfo([1,1,1,0,0,0,0,1,1], 2) 0.305958 1e-6
@test_approx_eq_eps activeinfo([0,0,0,1,1,1,1,0,0], 2) 0.305958 1e-6
@test_approx_eq_eps activeinfo([0,0,0,0,0,0,1,1,0], 2) 0.347458 1e-6

@test_approx_eq_eps activeinfo([3,3,3,2,1,0,0,0,1], 2, b=4) 0.635471 1e-6
@test_approx_eq_eps activeinfo([2,2,3,3,3,3,2,1,0], 2, b=4) 0.635471 1e-6
@test_approx_eq_eps activeinfo([2,2,2,2,2,2,1,1,1], 2, b=4) 0.234783 1e-6

## Ensemble time series
let series = [ 1 0; 1 0; 0 0; 0 1; 1 0; 0 0; 0 0; 1 1]
    @test_approx_eq_eps activeinfo(series, 2) 0.459148 1e-6
end
let series = [
    1 0 1 1 0 0 1 0 0;
    0 0 0 0 0 0 1 0 0;
    0 1 0 0 0 0 1 0 0;
    0 1 0 0 0 0 0 1 0;
    0 1 0 0 0 1 0 1 0;
    0 1 0 0 1 1 0 1 0;
    0 0 0 0 1 0 0 1 1;
    0 0 1 1 0 0 1 0 1;
    0 0 1 1 0 0 1 0 0]
    @test_approx_eq_eps activeinfo(series, 2) 0.308047 1e-6
end

let series = [
    3 3 3 2 1 0 0 0 1;
    2 2 3 3 3 3 2 1 0;
    0 0 0 0 1 1 0 0 0;
    1 1 0 0 0 1 1 2 2]'
    @test_approx_eq_eps activeinfo(series, 2) 0.6621455 1e-6
end


using LinearAlgebra
using Statistics
using Random

function coherence(X, ord=1, model=5, iter=200, ax=1)
	X = makebinary(X)

	if ord == 1
		X = allzeros(X)
		Y = ordinate(X, ax)
	end

	Abs = embedabs(Y)

	EAbs = zeros(Int, iter)
	for i = 1:iter
		null = randomnull(X, model)
		if ord == 1
			EX = allzeros(null)
			EY = ordinate(EX, ax)
		end
		EAbs[i] = embedabs(EY)
	end

	MA = mean(EAbs)
	SA = std(EAbs)

	# z-test
	return Abs, MA, SA
end

function turnover(X, ord=1, Comm=0, iter=200, ax=1)
	X = makebinary(X)

	if ord == 1
		X = allzeros(X)
		Y = ordinate(X, ax)
	end

	Y = boundfill(Y, Comm)
	Re = checkstat(Y)

	ERe = zeros(iter)
	for i = 1:iter
		null = replacenull(Y, Comm)
		ERe[i] = checkstat(null)
	end
	MR = mean(ERe)
	SR = std(ERe)

	return Re, MR, SR
end

function clumping(X, ord=1, Comm=0, ax=1)
	if Comm == 0
		X = transpose(X)
	end

	X = makebinary(X)

	if ord == 1
		X = allzeros(X)
		Y = ordinate(X)
	end

	r,c = size(Y)
	ComBnd = zeros(c,1)
	M = 0
	ComBndChi = 0

	for i = 1:r
		ind = findall(Y[i,:] .== 1)
		First = ind[1];
		Last = ind[length(ind)]
		for j = 1:c
			if First == j
				ComBnd[j] += 1
			end
			if Last == j
				ComBnd[j] += 1
			end
		end
	end

	TotComBnd = r * 2 - ComBnd[1] - ComBnd[c]
	ExpComBnd = TotComBnd / (c - 2)
	df = -1

	for i = 2:(c-1)
		M += ((ComBnd[i]/TotComBnd)*((ComBnd[i]-1)/(TotComBnd-1)))
		ComBndChi += ((ComBnd[i] - ExpComBnd)^2/ExpComBnd)
		df += 1
	end

	M *= c-2

	return M, ComBndChi, df
end

function replacenull(X, Comm)
	sites, species = size(X)
	null = zeros(Int, sites, species)
	if Comm == 0
		for i = 1:species
			ind = findall(X[:,i] .== 1)
			len = length(ind)
			rd = Int(ceil(rand()*(sites-len+1)))
			null[rd:(rd+len-1),i] .= 1
		end
	elseif Comm == 1
		for i = 1:sites
			ind = findall(X[i,:] .== 1)
			len = length(ind)
			rd = Int(ceil(rand()*(species-len+1)))
			null[i,rd:(rd+len-1)] .= 1
		end
	end
	return null
end

function checkstat(X)
	r,c = size(X)
	C = 0
	for i = 1:(c-1)
		Xi = X[:,i]
		R1 = sum(Xi)
		for j = (i+1):c
			Xj = X[:,j]
			S = sum(Xi.&Xj)
			R2 = sum(Xj)
			C += (R1-S) * (R2-S)
		end
	end
	return C
end

function boundfill(X, Comm)
	sites, species = size(X)
	Y = copy(X)
	if Comm == 0
		for i = 1:species
			ind = findall(X[:,i] .== 1)
			First = ind[1]
			Last = ind[length(ind)]
			Y[First:Last,i] .= 1
		end
	elseif Comm == 1
		for i = 1:sites
			ind = findall(X[i,:] .== 1)
			First = ind[1]
			Last = ind[length(ind)]
			Y[i,First:Last] .= 1
		end
	end
	return Y
end

function makebinary(X)
	Y = Int.(X .> 0)
	return Y
end

function allzeros(X, rc=3)
	Y = X
	if rc == 3 || rc == 1
		Rsum = sum(X, dims=2)
		Rind = findall(Rsum[:] .> 0)
		Y = X[Rind,:]
	end
	if rc == 3 || rc == 2
		Csum = sum(X, dims=1)
		Cind = findall(Csum[:] .> 0)
		Y = Y[:,Cind]
	end
	return Y
end

function ordinate(X, ax=1)
	V, W = ca(X)
	ir = sortperm(V[ax+1,:])
	ic = sortperm(W[ax+1,:])
	Y = X[ir,ic]
	Y = Y[size(Y,1):-1:1,:]
	return Y
end

function ca(X)
	r = sum(X, dims=2)
	c = sum(X, dims=1)
	N = sum(r)
	R = diagm(0 => r[:])
	C = diagm(0 => c[:])

	Risq = sqrt(inv(R))
	Cisq = sqrt(inv(C))

	M = Risq * X * Cisq
	Mt = transpose(M)
	P = M * Mt
	Q = Mt * M

	vals, vecs = eig(P)
	UP = transpose(vecs)
	vals, vecs = eig(Q)
	UQ = transpose(vecs)

	V = sqrt(N) * UP * Risq
	W = sqrt(N) * UQ * Cisq

	return V, W
end

function eig(X, n=size(X,1))
	vals, vecs = eigen(X)

	if !isreal(vals)
		mask = isreal.(vals)
		image_ = abs.(.!mask .* vals)
		real_ = mask .* vals
		vals = Float64.(real_ + image_)
	end
	if !isreal(vecs)
		mask = isreal.(vecs)
		image_ = abs.(.!mask .* vecs)
		real_ = mask .* vecs
		vecs = Float64.(real_ + image_)
	end

	ind = sortperm(-vals)
	vals = vals[ind]
	vecs = vecs[:,ind]

	for j = 1:size(vecs, 2)
		if sum(vecs[:,j]) < 0
			vecs[:,j] = -vecs[:,j];
		end
	end

	return vals[1:n], vecs[:,1:n]
end

function embedabs(X)
	sites, species = size(X)
	FirstSpecies, LastSpecies, FirstSite, LastSite = boundfind(X)
	Abs = 0
	for i = 1:sites
		for j = 1:species
			if X[i,j] == 0
				flag = 0
				if j > FirstSpecies[i] && j < LastSpecies[i]
					Abs += 1
					flag = 1
				end
				if flag == 0 && i > FirstSite[j] && i < LastSite[j]
					Abs += 1
				end
			end
		end
	end
	return Abs
end

function boundfind(X)
	r,c = size(X)
	FirstSpecies = zeros(r,1)
    LastSpecies = zeros(r,1)
    FirstSite = zeros(1,c)
	LastSite = zeros(1,c)
	for i = 1:r
        ind = findall(X[i,:] .== 1)
        FirstSpecies[i] = ind[1];
        LastSpecies[i] = ind[length(ind)];
	end
	for i = 1:c
        ind = findall(X[:,i] .== 1);
        FirstSite[i] = ind[1];
        LastSite[i] = ind[length(ind)];
	end
	return FirstSpecies, LastSpecies, FirstSite, LastSite
end

# only binary mode is implemented
function randomnull(X, model=1)
	r,c = size(X)

	Richness = sum(X, dims=2)
	Incidence = sum(X, dims=1)
	N = sum(Richness)

    FixRichness = ones(r,1)
    EqRichness = ones(r,1) * (1/r)
    PropRichness = Richness / N
    FixIncidence = ones(1,c)
    EqIncidence = ones(1,c) * (1/c)
	PropIncidence = Incidence / N

	null = zeros(r,c)
	if model == 1
		P = EqRichness * EqIncidence
		null = assignsp(P, N)
	elseif model == 2
		P = EqRichness * FixIncidence
		for i = 1:c
			Pi = reshape(P[:,i], (:,1))
			null[:,i] = assignsp(Pi, Incidence[i])
		end
	elseif model == 3
		P = FixRichness * EqIncidence
		for i = 1:r
			Pi = reshape(P[i,:], (1,:))
			null[i,:] = assignsp(Pi, Richness[i])
		end;
	elseif model == 4
		P = PropRichness * FixIncidence
		for i = 1:c
			Pi = reshape(P[:,i], (:,1))
			null[:,i] = assignsp(Pi, Incidence[i])
		end;
	elseif model == 5
		P = FixRichness * PropIncidence
		for i = 1:r
			Pi = reshape(P[i,:], (1,:))
			null[i,:] = assignsp(Pi, Richness[i])
		end
	elseif model == 6
		P = PropRichness * EqIncidence
        null = assignsp(P, N)
	elseif model == 7
		P = EqRichness * PropIncidence
		null = assignsp(P, N)
	elseif model == 8
		P = PropRichness * PropIncidence
		null = assignsp(P, N)
	elseif model == 9
		null = swap(X)
	elseif model == 10
		null = randnest(X)
	end

	return null
end

# not functionally idential to Swap.m, but similar and quicker
function swap(X)
	r,c = size(X)
	ir = randperm(r)
	return X[ir,:]
end

function randnest(X)
	r,c = size(X)
	Incidence = sum(X, dims=1)
	SpecFreq = Incidence / r
	null = ones(r, c)
	mask = rand(r, c) .< repeat(SpecFreq, 5)
	return mask .* null
end

function assignsp(P, N, con=1)
	r,c = size(P)
	null = zeros(r,c)
	X = P[:]

	X = X / sum(X)
	C = cumsum(X)

	while true
		for i = 1:N
			while true
				ind = findall(C .>= rand())
				if null[ind[1]] != 1
					null[ind[1]] = 1
					break
				end
			end
		end
		if con == 1
			break
		end
	end

	return null
end


using Test

using QuantumCircuits
using QuantumCircuits.QCircuits.Math

using QuantumCircuits.Simulation.Gates

gates = [ZZ, XX, YY]
θs = [0, π/4, π/2, 3π/4, π, 5π/4, 3π/2, 7π/4, 2π]

for g in gates
    for θ in θs
        simqc = QCircuit(2)
        g(simqc, 0, 1, θ, true)
        m1 = tomatrix(simqc)

        simqc = QCircuit(2)
        g(simqc, 0, 1, θ, false)
        m2 = tomatrix(simqc)

        @test unitary_error(m1, m2) < 1e-8
    end
end


module QGA
export minimize

using Printf
#using StatsBase
using Random
using Statistics
using LinearAlgebra

function ent(t, f, fmin, w)
    w .= exp2.(t.*(fmin.-f))
    z = sum(w)
    w .*= t.*(fmin.-f)
    log2(z)-sum(w)/z
end

function selection_entropy!(w, k, t, f, fmin)
    if ent(t, f, fmin, w) > k
        while true
            t = t*1.001
            if ent(t, f, fmin, w) < k
                return t
            end
#             if t > 1e15
#                 @warn "probably infinite loop" f
#             end
        end
    else
        smax = log2(length(f))
        while true
            t = t/1.001
            s = ent(t, f, fmin, w)
            if s > k
                return t
            elseif s >= smax - 1e-2
                @warn "reached maximum entropy $smax"
                return t
            end
        end
    end
end

function minimize(fitness, x, f, w, r, xm, xnew, s, n,
        FminTarget, sdmin, trace, MaxEvals, mumin; evals = 0)
    fmin, fmini = findmin(f)
    beta = 0.1/std(f)
    sd = 0.0
    dim = size(x, 1)
    local halt
    if trace < Inf
        println("Iteration     Fmin     Beta")
    end
    while true
        beta = selection_entropy!(w, s, beta, f, fmin)
        w .= exp2.(beta.*(fmin.-f))
        z = sum(w)
        z2 = sum(abs2, w)
        fm = mapreduce(*, +, f, w)/z
        if sdmin > 0.0
            sd = mapreduce((f,w)-> (f-fm)*w, +, f, w)/(z - z2/z)
            if sd < sdmin
                halt = :sdmin
                break
            end
        end
        if evals % trace == 0
            @printf("%9i %1.2e %1.2e\n",
                    evals, fmin, beta)
        end
        if mumin
            xm .= x[:, fmini]
        else
            for i = 1:dim
                xm[i] = 0.0
                for j = 1:n
                    xm[i] += x[i, j]*w[j]
                end
                xm[i] /= z
            end
        end

        randn!(r)
        r .= r.*sqrt.(w)/sqrt(z-z2/z)
        mul!(xnew, x, r)
        xnew .+= xm.*(1-sum(r))

        j = argmin(w)
        f[j] = fitness(xnew)
        x[:, j] .= xnew
        evals += 1
        if f[j] < fmin
            fmin = f[j]
            fmini = j
        end
        if fmin < FminTarget
            halt = :FminTarget
            break
        elseif evals >= MaxEvals
            halt = :MaxEvals
            break
        elseif sum(isequal(f[j]), f) > 1
            halt = :DupFit
            break
        end
    end
    Dict(:fmin => fmin, :halt => halt, :xmin => x[:,fmini],
         :n => n, :mumin => mumin, :sd => sd,
         :evals => evals, :s=> s)
end

function minimize(fitness, xm0, xsd0, s;
        FminTarget = -Inf, sdmin = 1e-16,
        trace = Inf, MaxEvals = 10000*length(xm0),
        mumin = true,
        n = floor(Int64, exp2(s+1)))
    dim = length(xm0)
    x = randn(dim, n).*xsd0 .+ xm0
    f = vec(mapslices(fitness, x; dims=1))
    w = zeros(n)
    r = zeros(n)
    xm = zeros(dim)
    xnew = zeros(dim)
    minimize(fitness, x, f, w, r, xm, xnew, s, n,
        FminTarget, sdmin, trace, MaxEvals, mumin, evals = n)
end

end


bitsof(::Type{T}) where {T} = sizeof(T) * 8
bitsof(x::T) where {T} = sizeof(T) * 8

#=
    maxmag(q) whichever is the larger, abs(numerator) or abs(denominator)
    leading_zeros_maxmag(q) count of leading 0bits in maxmag(q)

    msbitidx(q) 1-based bit index of the most significant bit in maxmag(q)
                zero iff iszero(maxmag(q))
=#
maxmag(q::Rational{T}) where {T<:Integer} = max(abs(q.num), abs(q.den))
leading_zeros_maxmag(q::Rational{T}) where {T<:Integer} = leading_zeros(maxmag(q))

msbitidx(q::Rational{T}) where {T<:Integer} = bitsof(T) - leading_zeros_maxmag(q)
msbitidx(q1::Rational{T}, q2::Rational{T}) where {T<:Integer} = msbitidx(q1) + msbitidx(q2)

maxmag(q::FastRational{T}) where {T<:FastInt} = max(abs(q.num), abs(q.den))  # q.den != typemin(T)
leading_zeros_maxmag(q::FastRational{T}) where {T<:FastInt} = leading_zeros(maxmag(q))

msbitidx(q::FastRational{T}) where {T<:FastInt} = bitsof(T) - leading_zeros(maxmag(q))
msbitidx(q1::FastRational{T}, q2::FastRational{T}) where {T<:FastInt} = msbitidx(q1) + msbitidx(q2)

#=
msbitidx(q1::T) + msbitidx(q2::T)
    =  (bitsof(T) - leading_zeros(maxmag(q1))) + (bitsof(T) - leading_zeros(maxmag(q2)))
    =  (bitsof(T) + bitsof(T)) - (leading_zeros(maxmag(q1)) + leading_zeros(maxmag(q2)))
    =  2*bitsof(T) - (leading_zeros(maxmag(q1)) + leading_zeros(maxmag(q2)))
=#

mayoverflow(q1::T, q2::T) where {T} = bitsof(T) <= msbitidx(q1, q2)
#=
   = bitsof(T) <= 2*bitsof(T) - (leading_zeros(maxmag(q1)) + leading_zeros(maxmag(q2)))
   = (leading_zeros(maxmag(q1)) + leading_zeros(maxmag(q2))) <= 2*bitsof(T) - bitsof(T)
   = (leading_zeros(maxmag(q1)) + leading_zeros(maxmag(q2))) <= bitsof(T)
=#

mayoverflow(i1::T, i2::T) where {T<:Integer} =
    (leading_zeros(i1) + leading_zeros(i2)) <= bitsof(T)

#=
   this was used in deriving the faster version immediately below

   mayoverflow(q1::Rational{T}, q2::Rational{T}) where {T<:Integer} =
        bitsof(T) >= leading_zeros_maxmag(q1) + leading_zeros_maxmag(q2)
=#

mayoverflow(q1::Rational{T}, q2::Rational{T}) where {T<:Integer} =
    (bitsof(T)<<1) >= leading_zeros(q1.num) + leading_zeros(q1.den) +
                      leading_zeros(q2.num) + leading_zeros(q2.den)

mayoverflow(q1::FastRational{T}, q2::FastRational{T}) where {T<:FastInt} =
    (bitsof(T)<<1) >= leading_zeros(q1.num) + leading_zeros(q1.den) +
                      leading_zeros(q2.num) + leading_zeros(q2.den)