using FrankWolfe
using LinearAlgebra
using Test

n = Int(1e3)
k = n

f(x) = dot(x, x)
function grad!(storage, x)
    @. storage = 2 * x
end

# pick feasible region
lmo = FrankWolfe.ProbabilitySimplexOracle{Rational{BigInt}}(1); # radius needs to be integer or rational

# compute some initial vertex
x0 = FrankWolfe.compute_extreme_point(lmo, zeros(n));

x, v, primal, dual_gap0, trajectory = FrankWolfe.frank_wolfe(
    f,
    grad!,
    lmo,
    x0,
    max_iteration=k,
    line_search=FrankWolfe.Agnostic(),
    print_iter=k / 10,
    verbose=false,
    emphasis=FrankWolfe.blas,
)

xmem, vmem, primal, dual_gap, trajectory = FrankWolfe.frank_wolfe(
    f,
    grad!,
    lmo,
    copy(x0),
    max_iteration=k,
    line_search=FrankWolfe.Agnostic(),
    print_iter=k / 10,
    verbose=true,
    emphasis=FrankWolfe.memory,
)

@time xstep, _ = FrankWolfe.frank_wolfe(
    f,
    grad!,
    lmo,
    x0,
    max_iteration=k,
    line_search=FrankWolfe.RationalShortstep(),
    L=2,
    print_iter=k / 10,
    verbose=true,
)

@test eltype(x) == eltype(xmem) == eltype(xstep) == Rational{BigInt}
@test xmem == x
@test abs(f(xstep) - f(x)) <= 1e-3
@test vmem == v


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

using Markdown
using InteractiveUtils

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

# ╔═╡ fafae38e-e852-11ea-1208-732b4744e4c2
md"_homework 0, version 2_"

# ╔═╡ 7308bc54-e6cd-11ea-0eab-83f7535edf25
# edit the code below to set your name and kerberos ID (i.e. email without @mit.edu)

student = (name = "Andreas Theodosiou", kerberos_id = "at694")

# press the ▶ button in the bottom right of this cell to run your edits
# or use Shift+Enter

# you might need to wait until all other cells in this notebook have completed running.
# scroll down the page to see what's up

# ╔═╡ cdff6730-e785-11ea-2546-4969521b33a7
md"""

Submission by: **_$(student.name)_** ($(student.kerberos_id)@mit.edu)
"""

# ╔═╡ a2181260-e6cd-11ea-2a69-8d9d31d1ef0e
md"""
# Homework 0: Getting up and running

First of all, **_welcome to the course!_** We are excited to teach you about real world applications of scientific computing, using the same tools that we work with ourselves.

Before we start next week, we'd like everyone to **submit this zeroth homework assignment**. It will not affect your grade, but it will help us get everything running smoothly when the course starts. If you're stuck or don't have much time, just fill in your name and ID and submit 🙂
"""

# ╔═╡ 094e39c8-e6ce-11ea-131b-07c4a1199edf

# ╔═╡ 31a8fbf8-e6ce-11ea-2c66-4b4d02b41995

# ╔═╡ 339c2d5c-e6ce-11ea-32f9-714b3628909c
md"## Exercise 1 - _Square root by Newton's method_

Computing the square of a number is easy -- you just multiply it with itself.

But how does one compute the square root of a number?

##### Algorithm:

Given: $x$

Output: $\sqrt{x}$

1. Take a guess `a`
1. Divide `x` by `a`
1. Set a = the average of `x/a` and `a`. (The square root must be between these two numbers. Why?)
1. Repeat until `x/a` is roughly equal to `a`. Return `a` as the square root.

In general, you will never get to the point where `x/a` is _exactly_ equal to `a`. So if our algorithm keeps going until `x/a == a`, then it will get stuck.

So instead, the algorithm takes a parameter `error_margin`, which is used to decide when `x/a` and `a` are close enough to halt.
"

# ╔═╡ 56866718-e6ce-11ea-0804-d108af4e5653
md"### Exercise 1.1

Step 3 in the algorithm sets the new guess to be the average of `x/a` and the old guess `a`.

This is because the square root must be between the numbers `x/a` and `a`. Why?
"

# ╔═╡ bccf0e88-e754-11ea-3ab8-0170c2d44628
ex_1_1 = md"""
The square root of a number $x$ is in the range $(1, x)$. Thus, for any value of $a$, where $a>0$, if $a>\sqrt{x}$ then the interval $(\frac{x}{a}, a)$ includes the square root and if $a<\sqrt{x}$ then the interval $(a, \frac{x}{a})$ includes the square root.
"""

# you might need to wait until all other cells in this notebook have completed running.
# scroll down the page to see what's up

# ╔═╡ e7abd366-e7a6-11ea-30d7-1b6194614d0a
if !(@isdefined ex_1_1)
	md"""Do not change the name of the variable - write you answer as `ex_1_1 = "..."`"""
end

# ╔═╡ d62f223c-e754-11ea-2470-e72a605a9d7e
md"### Exercise 1.2

Write a function newton_sqrt(x) which implements the above algorithm."

# ╔═╡ 4896bf0c-e754-11ea-19dc-1380bb356ab6
function newton_sqrt(x, error_margin=0.01, a=x / 2) # a=x/2 is the default value of `a`
	while (abs(x/a - a) >= error_margin)
		println("x is $(x)")
		println("a is $(a)")
		a = (a + x/a) / 2
	end
	return a # this is wrong, write your code here!
end

# ╔═╡ 7a01a508-e78a-11ea-11da-999d38785348
newton_sqrt(2)

# ╔═╡ 682db9f8-e7b1-11ea-3949-6b683ca8b47b
let
	result = newton_sqrt(2, 0.01)
	if !(result isa Number)
		md"""
!!! warning "Not a number"
    `newton_sqrt` did not return a number. Did you forget to write `return`?
		"""
	elseif abs(result - sqrt(2)) < 0.01
		md"""
!!! correct
    Well done!
		"""
	else
		md"""
!!! warning "Incorrect"
    Keep working on it!
		"""
	end
end

# ╔═╡ 088cc652-e7a8-11ea-0ca7-f744f6f3afdd
md"""
!!! hint
    `abs(r - s)` is the distance between `r` and `s`
"""

# ╔═╡ c18dce7a-e7a7-11ea-0a1a-f944d46754e5
md"""
!!! hint
    If you're stuck, feel free to cheat, this is homework 0 after all 🙃

    Julia has a function called `sqrt`
"""

# ╔═╡ 5e24d95c-e6ce-11ea-24be-bb19e1e14657
md"## Exercise 2 - _Sierpinksi's triangle_

Sierpinski's triangle is defined _recursively_:

- Sierpinski's triangle of complexity N is a figure in the form of a triangle which is made of 3 triangular figures which are themselves Sierpinski's triangles of complexity N-1.

- A Sierpinski's triangle of complexity 0 is a simple solid equilateral triangle
"

# ╔═╡ 6b8883f6-e7b3-11ea-155e-6f62117e123b
md"To draw Sierpinski's triangle, we are going to use an external package, [_Compose.jl_](https://giovineitalia.github.io/Compose.jl/latest/tutorial). Let's set up a package environment and add the package.

A package contains a coherent set of functionality that you can often use a black box according to its specification. There are [lots of Julia packages](https://juliahub.com/ui/Home).
"

# ╔═╡ 851c03a4-e7a4-11ea-1652-d59b7a6599f0
# setting up an empty package environment
begin
	import Pkg
	Pkg.activate(mktempdir())
	Pkg.Registry.update()
end

# ╔═╡ d6ee91ea-e750-11ea-1260-31ebf3ec6a9b
# add (ie install) a package to our environment
begin
	Pkg.add("Compose")
	# call `using` so that we can use it in our code
	using Compose
end

# ╔═╡ 5acd58e0-e856-11ea-2d3d-8329889fe16f
begin
	Pkg.add("PlutoUI")
	using PlutoUI
end

# ╔═╡ dbc4da6a-e7b4-11ea-3b70-6f2abfcab992
md"Just like the definition above, our `sierpinksi` function is _recursive_: it calls itself."

# ╔═╡ 02b9c9d6-e752-11ea-0f32-91b7b6481684
complexity = 3

# ╔═╡ 1eb79812-e7b5-11ea-1c10-63b24803dd8a
if complexity == 3
	md"""
Try changing the value of **`complexity` to `5`** in the cell above.

Hit `Shift+Enter` to affect the change.
	"""
else
	md"""
**Great!** As you can see, all the cells in this notebook are linked together by the variables they define and use. Just like a spreadsheet!
	"""
end

# ╔═╡ d7e8202c-e7b5-11ea-30d3-adcd6867d5f5
md"### Exercise 2.1

As you can see, the total area covered by triangles is lower when the complexity is higher."

# ╔═╡ f22222b4-e7b5-11ea-0ea0-8fa368d2a014
md"""
Can you write a function that computes the _area of `sierpinski(n)`_, as a fraction of the area of `sierpinski(0)`?

So:
```
area_sierpinski(0) = 1.0
area_sierpinski(1) = 0.??
...
```
"""

# ╔═╡ ca8d2f72-e7b6-11ea-1893-f1e6d0a20dc7
function area_sierpinski(n)
	return (3/4)^n
end

# ╔═╡ 71c78614-e7bc-11ea-0959-c7a91a10d481
if area_sierpinski(0) == 1.0 && area_sierpinski(1) == 3 / 4
	md"""
!!! correct
    Well done!
	"""
else
	md"""
!!! warning "Incorrect"
    Keep working on it!
	"""
end

# ╔═╡ c21096c0-e856-11ea-3dc5-a5b0cbf29335
md"**Let's try it out below:**"

# ╔═╡ 52533e00-e856-11ea-08a7-25e556fb1127
md"Complexity = $(@bind n Slider(0:6, show_value=true))"

# ╔═╡ c1ecad86-e7bc-11ea-1201-23ee380181a1
md"""
!!! hint
    Can you write `area_sierpinksi(n)` as a function of `area_sierpinski(n-1)`?
"""

# ╔═╡ c9bf4288-e6ce-11ea-0e13-a36b5e685998

# ╔═╡ a60a492a-e7bc-11ea-0f0b-75d81ce46a01
md"That's it for now, see you next week!"

# ╔═╡ b3c7a050-e855-11ea-3a22-3f514da746a4
if student.kerberos_id === "jazz"
	md"""
!!! danger "Oops!"
    **Before you submit**, remember to fill in your name and kerberos ID at the top of this notebook!
	"""
end

# ╔═╡ d3625d20-e6ce-11ea-394a-53208540d626

# ╔═╡ dfdeab34-e751-11ea-0f90-2fa9bbdccb1e
triangle() = compose(context(), polygon([(1, 1), (0, 1), (1 / 2, 0)]))

# ╔═╡ b923d394-e750-11ea-1971-595e09ab35b5
# It does not matter which order you define the building blocks (functions) of the
# program in. The best way to organize code is the one that promotes understanding.

function place_in_3_corners(t)
	# Uses the Compose library to place 3 copies of t
	# in the 3 corners of a triangle.
	# treat this function as a black box,
	# or learn how it works from the Compose documentation here https://giovineitalia.github.io/Compose.jl/latest/tutorial/#Compose-is-declarative-1
	compose(context(),
			(context(1 / 4,   0, 1 / 2, 1 / 2), t),
			(context(0, 1 / 2, 1 / 2, 1 / 2), t),
			(context(1 / 2, 1 / 2, 1 / 2, 1 / 2), t))
end

# ╔═╡ e2848b9a-e703-11ea-24f9-b9131434a84b
function sierpinski(n)
	if n == 0
		triangle()
	else
		t = sierpinski(n - 1) # recursively construct a smaller sierpinski's triangle
		place_in_3_corners(t) # place it in the 3 corners of a triangle
	end
end

# ╔═╡ 9664ac52-e750-11ea-171c-e7d57741a68c
sierpinski(complexity)

# ╔═╡ df0a4068-e7b2-11ea-2475-81b237d492b3
sierpinski.(0:6)

# ╔═╡ 147ed7b0-e856-11ea-0d0e-7ff0d527e352
md"""

Sierpinski's triangle of complexity $(n)

 $(sierpinski(n))

has area **$(area_sierpinski(n))**

"""

# ╔═╡ Cell order:
# ╠═fafae38e-e852-11ea-1208-732b4744e4c2
# ╟─cdff6730-e785-11ea-2546-4969521b33a7
# ╠═7308bc54-e6cd-11ea-0eab-83f7535edf25
# ╟─a2181260-e6cd-11ea-2a69-8d9d31d1ef0e
# ╟─094e39c8-e6ce-11ea-131b-07c4a1199edf
# ╟─31a8fbf8-e6ce-11ea-2c66-4b4d02b41995
# ╟─339c2d5c-e6ce-11ea-32f9-714b3628909c
# ╟─56866718-e6ce-11ea-0804-d108af4e5653
# ╠═bccf0e88-e754-11ea-3ab8-0170c2d44628
# ╟─e7abd366-e7a6-11ea-30d7-1b6194614d0a
# ╟─d62f223c-e754-11ea-2470-e72a605a9d7e
# ╠═4896bf0c-e754-11ea-19dc-1380bb356ab6
# ╠═7a01a508-e78a-11ea-11da-999d38785348
# ╟─682db9f8-e7b1-11ea-3949-6b683ca8b47b
# ╟─088cc652-e7a8-11ea-0ca7-f744f6f3afdd
# ╟─c18dce7a-e7a7-11ea-0a1a-f944d46754e5
# ╟─5e24d95c-e6ce-11ea-24be-bb19e1e14657
# ╟─6b8883f6-e7b3-11ea-155e-6f62117e123b
# ╠═851c03a4-e7a4-11ea-1652-d59b7a6599f0
# ╠═d6ee91ea-e750-11ea-1260-31ebf3ec6a9b
# ╠═5acd58e0-e856-11ea-2d3d-8329889fe16f
# ╟─dbc4da6a-e7b4-11ea-3b70-6f2abfcab992
# ╠═e2848b9a-e703-11ea-24f9-b9131434a84b
# ╠═9664ac52-e750-11ea-171c-e7d57741a68c
# ╠═02b9c9d6-e752-11ea-0f32-91b7b6481684
# ╟─1eb79812-e7b5-11ea-1c10-63b24803dd8a
# ╟─d7e8202c-e7b5-11ea-30d3-adcd6867d5f5
# ╠═df0a4068-e7b2-11ea-2475-81b237d492b3
# ╟─f22222b4-e7b5-11ea-0ea0-8fa368d2a014
# ╠═ca8d2f72-e7b6-11ea-1893-f1e6d0a20dc7
# ╟─71c78614-e7bc-11ea-0959-c7a91a10d481
# ╟─c21096c0-e856-11ea-3dc5-a5b0cbf29335
# ╟─52533e00-e856-11ea-08a7-25e556fb1127
# ╟─147ed7b0-e856-11ea-0d0e-7ff0d527e352
# ╟─c1ecad86-e7bc-11ea-1201-23ee380181a1
# ╟─c9bf4288-e6ce-11ea-0e13-a36b5e685998
# ╟─a60a492a-e7bc-11ea-0f0b-75d81ce46a01
# ╟─b3c7a050-e855-11ea-3a22-3f514da746a4
# ╟─d3625d20-e6ce-11ea-394a-53208540d626
# ╠═dfdeab34-e751-11ea-0f90-2fa9bbdccb1e
# ╠═b923d394-e750-11ea-1971-595e09ab35b5


module NedelecFESpacesTests
##
using Test
using Gridap
using Gridap.MapApply
using Gridap.CellValuesGallery
using LinearAlgebra

D = 3
order = 2

dom = fill(1,D)
part = Tuple(fill(5,D))
model = CartesianDiscreteModel(domain=(0.0,1.0,0.0,1.0,0.0,1.0), partition=part)
# model = CartesianDiscreteModel(domain=(0.0,1.0,0.0,1.0), partition=part)

p = Polytope(fill(HEX_AXIS,D)...)
reffe = NedelecRefFE(p,Float64,order)

grid = Grid(model,D)
trian = Triangulation(grid)
graph = GridGraph(model)

labels = FaceLabels(model)
# tags = [5,6,7,8]
tags = Int[]
fun(x) = VectorValue(x[2],x[1],x[2]*x[1])
# fun(x) = VectorValue(1.0,1.0)
# tags = Int[]

fesp = ConformingFESpace(reffe,trian,graph,labels,tags)
dfesp = TrialFESpace(fesp,fun)

trian = Triangulation(model)
quad = CellQuadrature(trian,degree=2)

uh = interpolate(dfesp,fun)
uh.free_dofs
uh.diri_dofs

e = fun - uh
uh.free_dofs
uh.diri_dofs

el2 = sqrt(sum(integrate(inner(e,e),trian,quad)))
@test el2 < 1.0e-10

##

end # module


function bipartite_matching_test()
    W = sprand(10,8,0.5)
    bipartite_matching(W)
    bipartite_matching([10;12;13],[1;2;3],[3;2;4])

    A = [5 1 1 1 1;
     1 5 1 1 1;
     1 1 5 1 1;
     1 1 1 5 1;
     1 1 1 1 5;
     1 1 1 1 1]
    A = sparse(A)
    M1 = bipartite_matching(A)
    (m1,m2) = edge_list(M1)

    if M1.weight != 25 || !isequal(m1,m2)
        error("bipartite matching failed")
    end

    A = ones(Int64,6,5) - eye(Int64,6,5)
    A = sparse(A')
    (ei,ej,ev) = csr_to_sparse(A.colptr,A.rowval,A.nzval)
    ai = [1;2;3;4;5;ei]
    aj = [1;2;3;4;5;ej]
    av = [5;5;5;5;5;ev]

    M2 = bipartite_matching(av,ai,aj)
    mi = MatrixNetworks.edge_indicator(M2,ai,aj)
    mitrue = zeros(Int64,length(av))
    mitrue[1:5] = 1

    if M2.weight != 25 || !isequal(m1,m2) || sum(mi) !=5 || !isequal(mi,mitrue)
        error("bipartite matching failed")
    end

    return true
end


import Distributions.partype
import PDMats.dim
import Base.Random.rand
doc"""
    LKJ(η, p)

The *LKJ correlation matrix* with shape parameter \eta.

```julia
LKJ(eta, p)        # LKJ distribution with shape parameter \eta and dimension p

dim(d)            # Get the dimensions of LKJ correlation matrix
size(d)           # Get the size of correlation matrix as a tuple
params(d)         # Get the parameters, i.e. (mu, sig)

rand(d, n)          # Draw random samples from LKJ distribution, by default n=1
meanlogdet(d)       # Expected log determinant
meanlog(d, h)
entropy(d)
```

"""

immutable LKJ{T<:Real} <: ContinuousMatrixDistribution
    η::T
    p::Int

    function LKJ(η, p)
        η > zero(η) || error("η should be greater than zero.")
        new(η, p)
    end
end

LKJ(η::Real, p::Integer) = LKJ{typeof(η)}(η, p)
LKJ(η::Integer, p::Integer) = LKJ{Float64}(Float64(η), p)

@inline partype{T<:Real}(d::LKJ{T}) = T

"""
Draw n random samples with replacement from the LKJ distribution (n = 1 by default).
"""
function rand(d::LKJ, n::Integer)
    η = d.η
    p = d.p
    idxvec = [η + (p - i)/2. for i in 2:p] # generate Beta distribution indices
    L = zeros(p, p)
    for j in 1:length(idxvec)
        L[j+1:end, j] = 2 * rand(Distributions.Beta(idxvec[j], idxvec[j]), p - j) - 1
    end
    L = _cpc_to_corr(L, p)
    S = PDMats.PDMat(Base.LinAlg.Cholesky(full(L), :L)) # consistent with Stan Lib
    rand(S.mat, n)
end

rand(d::LKJ) = rand(d::LKJ, 1)

function _cpc_to_corr(L::AbstractMatrix, p::Real)  # consistent with Stan Lib
    #p = dim(d)
    LL = eye(p)
    for j in 1:p
        for i in j+1:p
            sumsqs = 0 # faster than dot(vec(L[i, 1:j]), vec(L[i, 1:j]))
            for ij in 1:j
                sumsqs += LL[i, ij]^2
            end
            LL[i, j] = L[i, j] * sqrt(1 - sumsqs)
        end
        sumsqs = 0
        for ij in 1:j-1
            sumsqs += LL[j, ij]^2
        end
        LL[j, j] = sqrt(1 - sumsqs)
    end
    LowerTriangular(LL)
end

dim(d::LKJ) = d.p
size(d::LKJ) = (d.p, d.p)
params(d::LKJ) = (d.η, d.p)

"""
E[log |A|] where
A ~ LKJ(eta) of dimension p
"""
function meanlogdet(d::LKJ)
    η = d.η
    p = d.p
    idxvec = [η + (p - i)/2. for i in 2:p]
    ldsum = 0
    for i in 1:length(idxvec)
        ldsum += (digamma(idxvec[i]) - digamma(2 * idxvec[i])) * (p - i)
    end
    ldsum
end

"""
E_q[log p(x)] where
p(x) = LKJ(h, p)  # p = matrix dimension
q(x) = LKJ(eta, p)
"""
function meanlog(d::LKJ, h::Real)
    η = d.η
    p = d.p
    betaq = [η + (p - i)/2. for i in 2:p]
    betap = [h + (p - i)/2. for i in 2:p]
    Elogx = 0
    for i in 1:length(betaq) # compute from Dirichlet mean log
        Elogx += (p - i) * (2(betap[i] - 1) .* (digamma(betaq[i]) - digamma(2betaq[i])) - (2lgamma(betap[i]) - lgamma(2betap[i])))
    end
    Elogx
end

"""
Compute entropy for LKJ distribution by Dirichelet entropy
"""
function entropy(d::LKJ)
    η = d.η
    p = d.p
    idxvec = [η + (p - i)/2. for i in 2:p]
    H = 0
    for i in 1:length(idxvec) # compute from Dirichlet entropy
        H += (p - i) * (2*lgamma(idxvec[i]) - lgamma(2idxvec[i]) + 2(idxvec[i] - 1) .* digamma(2idxvec[i]) - 2 * (idxvec[i] - 1) .* digamma(idxvec[i]))
    end
    H
end


struct Boat_timing
    T_in::Float64
    T_out::Float64
    T_1::Float64
    T_2::Float64
    T_t::Float64
end

# T_1 = 2 * (0.7 + T_in + T_out)
Boat_timing(T_active, T_passive) =
    Boat_timing(0, 0, 2T_active, 2T_passive, T_active + T_passive)

Boat_timing() = Boat_timing(0.7, 0.5)

struct Boat_angle_ranges
    γ_OA::Tuple{Float64,Float64}
    θ_t::Tuple{Float64,Float64}
    θ_k::Tuple{Float64,Float64}
end

Boat_angle_ranges() = Boat_angle_ranges((30.0, 120.0), (22.5, 22.5), (90.0, 90.0))

function Base.deg2rad(angles::Boat_angle_ranges)
    Boat_angle_ranges(deg2rad.(angles.γ_OA), deg2rad.(angles.θ_t), deg2rad.(angles.θ_k))
end

struct Boat_angles
    γ_OA::Float64
    θ_t::Float64
    θ_k::Float64
    γ′_OA::Float64
    θ′_t::Float64
    θ′_k::Float64
    γ′′_OA::Float64
    θ′′_t::Float64
    θ′′_k::Float64
end

function Boat_angles(angle_ranges::Boat_angle_ranges, bt::Boat_timing, t)
    aranges = angle_ranges.γ_OA, angle_ranges.θ_t, angle_ranges.θ_k
    γoa, θt, θk = map(aranges) do r
        _angle(t, bt, r[1], r[2])
    end
    γ′oa, θ′t, θ′k = map(aranges) do r
        _angle′(t, bt, r[1], r[2])
    end
    γ′′oa, θ′′t, θ′′k = map(aranges) do r
        _angle′′(t, bt, r[1], r[2])
    end
    Boat_angles(γoa, θt, θk, γ′oa, θ′t, θ′k, γ′′oa, θ′′t, θ′′k)
end

Boat_angles() = Boat_angles(deg2rad(Boat_angle_ranges()), Boat_timing(), 0.0)

t_periodic(t, bt::Boat_timing) = mod(t, bt.T_t)

function η(t, bt::Boat_timing)
    # to - czas w danym okresie
    to = t_periodic(t, bt)

    if to < bt.T_in
        to / bt.T_in
    elseif to < 0.5bt.T_1 - bt.T_out
        one(t)
    elseif to < 0.5bt.T_1
        1 - (to - 0.5bt.T_1 + bt.T_out) / bt.T_out
    else
        zero(t)
    end
end

function kpars(to, bt::Boat_timing)
    #KPARS Helper to get phase parameters.
    k1 = to <= 0.5bt.T_1
    (k1, !k1) .* 1
end

function _angle(t, bt::Boat_timing, vmin, vmax)
    to = t_periodic(t, bt)
    k1, k2 = kpars(to, bt)
    T_1, T_2 = bt.T_1, bt.T_2
    0.5(vmax + vmin) -
    0.5(vmax - vmin) * (k1 * cospi(2to / T_1) - k2 * cospi(2(to - 0.5T_1) / T_2))
end

function _angle′(t, bt::Boat_timing, vmin, vmax)
    to = t_periodic(t, bt)
    k1, k2 = kpars(to, bt)
    T_1, T_2 = bt.T_1, bt.T_2
    -π *
    ((k2 * sinpi(2(T_1 / 2 - to) / T_2)) / T_2 + (k1 * sinpi(2to / T_1)) / T_1) *
    (vmin - vmax)
end

function _angle′′(t, bt::Boat_timing, vmin, vmax)
    to = t_periodic(t, bt)
    k1, k2 = kpars(to, bt)
    T_1, T_2 = bt.T_1, bt.T_2
    2π^2 *
    (k2 * cospi(2(T_1 / 2 - to) / T_2) / T_2^2 - k1 * cospi(2to / T_1) / T_1^2) *
    (vmin - vmax)
end


@doc raw"""
    postproc(model::Potts, param::Parameter, obs::MCObservableSet)

# Observables to be calculated
In the following, $m$ is total magnetization per site and $\epsilon$ is total energy per site.

- `"Binder Ratio"`
    - $R := \frac{\left \langle m^4 \right \rangle}{\left \langle m^2 \right\rangle^2}$
- `"Susceptibility"`
    - $\chi := \frac{N}{T}\left(\left\langle m^2\right\rangle\right)$
- `"Connected Susceptibility"`
    - $\frac{N}{T}\left(\left\langle m^2\right\rangle - \left\langle |m| \right\rangle^2\right)$
- `"Specific Heat"`
    - $\frac{N}{T^2}\left(\left\langle \epsilon^2\right\rangle - \left\langle \epsilon \right\rangle^2\right)$
"""
function postproc(model::Potts, param::Parameter, obs::MCObservableSet)
    nsites = numsites(model)
    T = convert(Float64, param["T"])
    beta = 1.0/T

    jk = jackknife(obs)
    jk["Binder Ratio"] = jk["Magnetization^4"] / (jk["Magnetization^2"]^2)
    jk["Susceptibility"] = (nsites*beta)*jk["Magnetization^2"]
    jk["Connected Susceptibility"] = (nsites*beta)*(jk["Magnetization^2"] - jk["|Magnetization|"]^2)
    jk["Specific Heat"] = (nsites*beta*beta)*(jk["Energy^2"] - jk["Energy"]^2)
    return jk
end


#=

  Ski assignment in Julia ConstraintSolver.jl

  From
  Jeffrey Lee Hellrung, Jr.: PIC 60, Fall 2008, Final Review, December 12, 2008
  http://www.math.ucla.edu/~jhellrun/course_files/Fall%25202008/PIC%252060%2520-%2520Data%2520Structures%2520and%2520Algorithms/final_review.pdf
  """
  5. Ski Optimization! Your job at Snapple is pleasant but in the winter you've
  decided to become a ski bum. You've hooked up with the Mount Baldy Ski Resort.
  They'll let you ski all winter for free in exchange for helping their ski rental
  shop with an algorithm to assign skis to skiers. Ideally, each skier should
  obtain a pair of skis whose height matches his or her own height exactly.
  Unfortunately, this is generally not possible. We define the disparity between
  a skier and his or her skis to be the absolute value of the difference between
  the height of the skier and the pair of skis. Our objective is to find an
  assignment of skis to skiers that minimizes the sum of the disparities.
  ...
  Illustrate your algorithm by explicitly filling out the A[i, j] table for the
  following sample data:
    * Ski heights: 1, 2, 5, 7, 13, 21.
    * Skier heights: 3, 4, 7, 11, 18.
  """

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#
using ConstraintSolver, JuMP
using Cbc, GLPK, Ipopt
const CS = ConstraintSolver
include("constraints_utils.jl")

function ski_assignment(print_solutions=true,all_solutions=false, all_optimal_solutions=false)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)
    ipopt_optimizer = optimizer_with_attributes(Ipopt.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   "all_solutions"=> all_solutions,
                                                            "all_optimal_solutions"=>all_optimal_solutions,
                                                            "logging"=>[],

                                                            "traverse_strategy"=>:BFS,
                                                            # "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            # "branch_split"=>:InHalf,

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>6,

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))

    ski_heights   =  [1, 2, 5, 7, 13, 21]
    skier_heights =  [3, 4, 7, 11, 18]

    num_skis = length(ski_heights)
    num_skiers = length(skier_heights)

    @variable(model, 1 <= x[1:num_skiers] <= num_skis, Int)
    @constraint(model, x in CS.AllDifferent())

    # minimize the differences of ski height and skier's height (z)
    t = 100
    @variable(model, 0 <= z <= t, Int)
    @variable(model, sxi[1:num_skiers], CS.Integers(ski_heights))
    @variable(model, 0 <= abs_diff[1:num_skiers] <= t, Int)

    # We must use element here...
    for i in 1:num_skiers
        # find the ski height for skier i
        my_element(model,x[i],ski_heights,sxi[i])
        # the difference beteen ski height and skier's height
        # my_abs(model,sxi[i],skier_heights[i],abs_diff[i])
        my_abs(model,sxi[i]-skier_heights[i],abs_diff[i])
    end
    @constraint(model,z==sum(abs_diff))

    @objective(model,Min, z)

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                # println("solution #$sol")
                xx = convert.(Integer,JuMP.value.(x; result=sol))
                zz = convert.(Integer,JuMP.value.(z; result=sol))
                sxix = convert.(Integer,JuMP.value.(sxi; result=sol))
                abs_diffx = convert.(Integer,JuMP.value.(abs_diff; result=sol))
                println("x:$xx  z:$zz")
                println("sxi:$sxix")
                println("abs_diff:$abs_diffx")
                println()
            end
        end
    else
        println("status:$status")
    end

    return status
end

@time ski_assignment()


##########################
# FunctionCell Interface #
##########################
"""
	ppca([;kwargs])

Constructs an untrained cell that when piped data inside, calculates the
probabilistic PCA projection matrix of the input data. We assume that `d`
is the dimesionality of the input dataset, `n` the number of samples

# Keyword arguments (same as in `MultivariateStats`)
  * `method` is the core algorithm used: `:ml`, `:em` and `:bayes` (default `:ml`)
  * `maxoutdim` is the number of output dimensions (default `<number of variables>-1`)
  * `mean` is the mean vector; can be `nothing`, `0` or precomputed mean (default `nothing`)
  * `tol` convergence tolerance (default `1.0e-6`)
  * `tot` maximum number of iterations (default `1000`)

Read the `MultivariateStats.jl` documentation for more information.
"""
ppca(;kwargs...) = FunctionCell(ppca, (), ModelProperties(), kwtitle("PPCA", kwargs);kwargs...)

############################
# DataCell/Array Interface #
############################
"""
	ppca(x, [;kwargs])

Trains a function cell that when piped data into will compute the principal components of
the observations based on the projection matrix extracted from `x`.
"""
# Training
ppca(x::T where T<:CellData; kwargs...) = ppca(getx!(x);kwargs...)
ppca(x::T where T<:AbstractVector; kwargs...) = ppca(mat(x, LearnBase.ObsDim.Constant{2}());kwargs...)
ppca(x::T where T<:AbstractMatrix; kwargs...)  = begin

	ppcadata = fit(MultivariateStats.PPCA, getobs(x); kwargs...)

	# Build model properties
	modelprops = ModelProperties(size(ppcadata.W,1), size(ppcadata.W,2))

	# Returned trained cell
	FunctionCell(ppca, Model(ppcadata, modelprops), kwtitle("PPCA",kwargs))
end

# Execution
ppca(x::T where T<:CellData, model::Model{<:MultivariateStats.PPCA}) =
	datacell(ppca(getx!(x), model), gety(x))
ppca(x::T where T<:AbstractVector, model::Model{<:MultivariateStats.PPCA}) =
	ppca(mat(x, LearnBase.ObsDim.Constant{2}()), model)
ppca(x::T where T<:AbstractMatrix, model::Model{<:MultivariateStats.PPCA}) =
	MultivariateStats.transform(model.data, getobs(x))

"""
	ppcar(M::Cell), where `M=ppca(x)`

Trains a cell based on a previously trained PPCA function cell. When piped data into,
the returned output will be the reconstruction of the observations in the original
space. Obviously, the input data (e.g. data to be reconstructed) is expected to be
principal components extracted with `M`.

Read the `MultivariateStats.jl` documentation for more information.
"""
# Training for the reconstruction tranform (fixed cell)
ppcar(x::T where T<:CellFunT{<:Model{<:MultivariateStats.PPCA}}) = begin

	# Build model properties

	modelprops = ModelProperties(getx(x).properties.odim, getx(x).properties.idim)

	FunctionCell(ppcar, Model(getx(x).data, modelprops), "PPCA: reconstruct")
end

# Execution
ppcar(x::T where T<:CellData, model::Model{<:MultivariateStats.PPCA}) =
	datacell(ppcar(getx!(x), model), gety(x))
ppcar(x::T where T<:AbstractVector, model::Model{<:MultivariateStats.PPCA}) =
	ppcar(mat(x, LearnBase.ObsDim.Constant{2}()), model)
ppcar(x::T where T<:AbstractMatrix, model::Model{<:MultivariateStats.PPCA})::Matrix{Float64} =
	MultivariateStats.reconstruct(model.data, getobs(x))


using SyntheticNetworks
using EmbeddedGraphs
using LightGraphs
using SimpleTraits
using Statistics

# defining parameters for graph initialization

n = 20
n0 = 10
p = 0.2
q = 0.3
r = 1. / 3.
s = 0.1
u = 0.0 # currently not supported

function synthetic_network(n, n0, p, q, r, s, u)
    RPG = RandomPowerGrid(n, n0, p, q, r, s, u)
    network = generate_graph(RPG) # network is an EmbeddedGraph object
    return network.graph
end

g = synthetic_network(n, n0, p, q, r, s, u)

## statistics

# average node degree
function average_node_degree(g::SimpleGraph)
    sum(degree(g))/nv(g)
end

average_node_degree(g)

# global clustering coefficient
global_clustering_coefficient(g)

# average shortest path length
mean(floyd_warshall_shortest_paths(g).dists)

## assortativity

function assortativity(g::AbstractGraph{T}) where T
    P = promote_type(Int64, T) # at least Int64 to reduce risk of overflow
    nue  = ne(g)
    sjk = sj = sk = sjs = sks = zero(P)
    for d in edges(g)
        j = P(outdegree(g, src(d)) - 1)
        k = P(indegree(g, dst(d)) - 1)
        sjk += j*k
        sj  += j
        sk  += k
        sjs += j^2
        sks += k^2
    end
    return assortativity_coefficient(g, sjk, sj, sk, sjs, sks, nue)
end

#=
assortativity coefficient for directed graphs:
see equation (21) in M. E. J. Newman: Mixing patterns in networks, Phys. Rev. E 67, 026126 (2003),
http://arxiv.org/abs/cond-mat/0209450
=#
@traitfn function assortativity_coefficient(g::::IsDirected, sjk, sj, sk, sjs, sks, nue)
    return (sjk - sj*sk/nue) / sqrt((sjs - sj^2/nue)*(sks - sk^2/nue))
end

#=
assortativity coefficient for undirected graphs:
see equation (4) in M. E. J. Newman: Assortative mixing in networks, Phys. Rev. Lett. 89, 208701 (2002),
http://arxiv.org/abs/cond-mat/0205405/
=#
@traitfn function assortativity_coefficient(g::::(!IsDirected), sjk, sj, sk, sjs, sks, nue)
    return (sjk/nue - ((sj + sk)/(2*nue))^2) / ((sjs + sks)/(2*nue) - ((sj + sk)/(2*nue))^2)
end

assortativity(g)

# end
##

# for plotting

open("julia_networks.txt", "a") do io
    for i in 1:1000
        graph = synthetic_network(n, n0, p, q, r, s, u)
        node_degree = average_node_degree(graph)
        assortativity_value = assortativity(graph)
        average_shortes_path_length = mean(floyd_warshall_shortest_paths(graph).dists)
        clust_coefficient = global_clustering_coefficient(graph)
        write(io, string(node_degree), "\t", string(assortativity_value), "\t", string(average_shortes_path_length), "\t", string(clust_coefficient), "\n")
        #write(io, string(i))
    end
end


using FFTW
using Plots
using Roots
using Optim
using Distributions
using Random
using LinearAlgebra
using ProgressMeter
using Primes
using LaTeXStrings

"""
Nazarov's ``magic'' Phi function x*arctan(x)-(1/2)log(1+x*x)
"""
function nazarov_phi(x::Real)
    return x*atan(x) - 0.5*log(1+x*x)
end

"""
Derivative Phi'(x)
"""
function nazarov_phiprime(x::Real)
    return atan(x)
end

"""
Nazarov's I(f)
"""
function nazarov_potential(fv::Array)
    s = 0
    for x in fv
        s += nazarov_phi(x)
    end
    return s
end

"""
Real orthonormal basis for the DFT. We only return non-DC terms.
"""
function dft_mat(n::Integer)
    mat = zeros(n,n)
    q = 2*pi/n
    h_pt = ceil(Int, (n+1)/2)
    # ceil((n+1)/2) cosines, floor((n-1)/2) sines
    for i=1:h_pt
        for j=1:n
            mat[i,j] = sqrt(2)*cos(q*(i-1)*(j-1))
        end
    end
    for i=h_pt+1:n
        for j=1:n
            mat[i,j] = sqrt(2)*sin(q*(i-h_pt)*(j-1))
        end
    end
    mat[1,:] /= sqrt(2)
    if iseven(n)
        mat[h_pt,:] /= sqrt(2)
    end
    return mat
end

function nazarov_psi(n::Integer)
    return dft_mat(n)[2:end,:]
end

"""
Compute the Nazarov constant for the DFT. In the paper's notation, this is M(n).
The paper gives the asymptotics, here one can numerically study it.
"""
function nazarov_const(n::Integer)
    mat = dft_mat(n)
    nc = 3*pi/2
    max_rat = 0
    for i=1:n
        nrat = 1/norm(mat[i,:],1)
        if (nrat > max_rat)
            max_rat = nrat
        end
    end
    max_rat *= n
    return nc*max_rat*max_rat
end

function plot_nazarov_const(n1::Integer, n2::Integer)
    x = primes(n1,n2)
    plot(x, [nazarov_const(i) for i in x])
end

"""
Some functions for checking the l1 norm beta(n) for the DFT. Asymptotics are given
in the paper. Finite n behavior can be studied with these routines, basically
the above nazarov_const is governed by the minimal value of the l1 norm over the
divisors of n.
"""
function l1_cos(n::Integer)
    sum = 0
    omega = 2*pi/n
    for i=0:n-1
        sum += abs(cos(omega*i))
    end
    sum /= n
    return sum-2/pi
end

function plot_l1_cos(n1::Integer, n2::Integer)
    x = n1:n2
    y = [l1_cos(n) for n in x]
    p = plot(x,y)
    savefig(p, "/tmp/l1_dft.pdf")
end

"""
Seed the greedy local optimum search randomly, works well in practice.
"""
function nazarov_seed_cortege(n::Integer)
    return rand([1,-1], n)
end

"""
Perform the local optimum search, given a vector that describes the power control.
"""
function nazarov_search!(a::Array, psi::Array, cortege::Array{Int})
    n = length(cortege)
    vn = length(psi[1,:])
    iter_order = [i for i in 1:n]
    fv = sum([cortege[j]*a[j]*psi[j,:] for j=1:n])
    cur_pot = nazarov_potential(fv)
    iter = 0
    eval_ct = 1
    is_opt = false
    println("-----------------------------------------------------------------")
    println("Performing local optimum search...")
    println("(iter, pot)")
    while !is_opt
        iter += 1
        print("($iter, $(cur_pot/vn))")
        randperm!(iter_order)
        is_opt = true
        for j in iter_order
            fv -= 2*cortege[j]*a[j]*psi[j,:]
            new_pot = nazarov_potential(fv)
            eval_ct += 1
            if new_pot > cur_pot
                is_opt = false
                cur_pot = new_pot
                cortege[j] = -cortege[j]
                break
            end
            fv += 2*cortege[j]*a[j]*psi[j,:]
        end
    end
    println("\n---------------------------------------------------------------")
    println("Function evaluations (scaled by n^2): $(eval_ct/(n*n))")
    println("\n---------------------------------------------------------------")
end

"""
Perform the shift/scale to get a [0,1] sequence from the locally optimal cortege.
We do not do the scaling of the paper, as it is conservative and meant for theorems.
The scaling here works better in practice and is natural.
"""
function _nazarov_aperture(a::Array, psi::Array, cortege::Array{Int})
    n = length(cortege)
    vn = length(psi[1,:])
    fv = sum([cortege[j]*a[j]*psi[j,:] for j=1:n])
    aper = zeros(vn)
    for i=1:vn
        aper[i] = nazarov_phiprime(fv[i])
    end
    # flip signs if positive
    if (sum(aper) > 0)
        aper = -aper
    end

    min_val = minimum(aper)
    max_val = maximum(aper)
    # scaling is trivial if nonnegative (very unlikely!)
    if (min_val > 0)
        warn("Something really weird, everything is nonnegative a priori!")
        aper /= max_val
        return aper
    end
    # the standard case. Note that this scaling is not the same used in the proof,
    # as the one used in the proof is a conservative one for theoretical soundness,
    # and we can improve upon it in code. Note that this scaling does not guarantee
    # \rho <= 0.5, though it is very close to it.
    for i=1:vn
        aper[i] = (aper[i]-min_val)/(max_val-min_val)
    end
    return aper
end

"""
General function that hooks things together: just takes in an n, a power control
vector (not necessarily normalized), and calls the above routines to get an aperture.
"""
function nazarov_aperture(n::Integer, powers::Array)
    psi = nazarov_psi(n)
    a = powers/norm(powers)
    cortege = nazarov_seed_cortege(n-1)
    nazarov_search!(a, psi, cortege)
    #redirect_stdout(() -> nazarov_search!(a, psi, cortege), open("/dev/null", "w"))
    return _nazarov_aperture(a, psi, cortege)
end

"""
Sanity checks on the generated aperture, and computation of its spectrum
"""
function nazarov_ahat(n::Integer, powers::Array, ctx::FFTW.cFFTWPlan)
    aperture = nazarov_aperture(n, powers)
    @assert(minimum(aperture) >= 0)
    @assert(maximum(aperture) <= 1)
    ahat = ctx * aperture
    ahat = [abs(x) for x in ahat]
    rho = ahat[1]/n
    println("rho (Nazarov): $rho")
    return ahat
end

"""
Scene correlation functions, generated by sampling at n equally spaced points
in [0,1] of a nonnegative, bounded, continuous function with symmetry about 1/2.
"""

"""
We first define the underlying d(x), for 0 <= x <= 0.5. All are underscored to
signify their "internal" nature.
"""
function _exp_corr(x::Real, theta::Real, beta::Real)
    return theta*(beta^x)
end

function _iid_corr(x::Real, theta::Real)
    return theta
end

function _char_corr(x::Real, theta::Real, beta::Real, gamma::Real)
    if (x < beta - gamma)
        return theta
    elseif (x > beta + gamma)
        return 0
    else
        return theta*(beta + gamma - x)/(2*gamma)
    end
end

function _power_corr(x::Real, theta::Real, beta::Real, gamma::Real)
    return theta/(1+beta*(x^gamma))
end

"""
General d vector generator based on sampling and symmetry about 1/2.
"""
function gen_corr(corr_fn, n::Integer)
    f = x -> (x < 0.5) ? corr_fn(x) : corr_fn(1-x)
    d = (1/n)*ones(n)
    for i=1:n
        j = i-1
        d[i] *= f(j/n)
    end
    return d
end

"""
Specializations to various cases of interest.
"""
function exp_corr(theta::Real, beta::Real, n::Integer)
    return gen_corr(x -> _exp_corr(x, theta, beta), n)
end

function iid_corr(theta::Real, n::Integer)
    return gen_corr(x -> _iid_corr(x, theta), n)
end

function char_corr(theta::Real, beta::Real, gamma::Real, n::Integer)
    return gen_corr(x -> _char_corr(x, theta, beta, gamma), n)
end

function power_corr(theta::Real, beta::Real, gamma::Real, n::Integer)
    return gen_corr(x -> _power_corr(x, theta, beta, gamma), n)
end

"""
Computation of power loss factor, i.e using a suboptimal rho. Here, a
(corresponding to W/J) is fixed. rhos describes an array of wherever we have
or consider spectrally flat constructions.
"""
function _power_loss_factor!(a::Real, rhos::Array)
    num = 2*a - 2*sqrt(a*(a+1))+1
    max_denom = 0
    for x in rhos
        denom = x*(1-x)/(a+x)
        if denom > max_denom
            max_denom = denom
        end
    end
    return num/max_denom
end

"""
Here we optimize (i.e maximize) _power_loss_factor over a, with an array of rhos
given (e.g rhos=[0.5, 0.25]). One may somewhat painfully do this rigorously;
this is just a sanity check.
"""
function power_loss_factor(rhos)
    # The 10 is completely ad-hoc, but works fine (see the plot)
    x = range(0, stop=10, length=10^8)
    y = zeros(length(x))
    i = 1
    @showprogress for i=1:length(x)
        y[i] = _power_loss_factor(x[i], rhos)
    end
    println("Loss factor: $(maximum(y))")
    x = range(0, stop=10, length=10^3)
    y = [_power_loss_factor(a, rhos) for a in x]
    p = plot(x, y)
    savefig(p, "/tmp/power_loss.pdf")
end

"""
LMMSE expressions.
"""
function eval_mmse(d::Array, n::Integer, t::Real, w::Real, j::Real, ahat::Array)
    mmse = 0
    rho = ahat[1]/n
    gamma = t/(n*(w+rho*j))
    for i=1:n
        mmse += 1/((1/d[i])+gamma*ahat[i]*ahat[i])
    end
    return mmse
end

function eval_mmse_flat(d::Array, rho::Real, n::Integer, t::Real, w::Real, j::Real)
    mmse = 0
    r = n*rho
    gamma = t/(n*(w+rho*j))
    mmse += 1/((1/d[1])+gamma*r*r)
    r = sqrt((n*n*rho*(1-rho))/(n-1))
    for i=2:n
        mmse += 1/((1/d[i])+gamma*r*r)
    end
    return mmse
end

function eval_mmse_rand(d::Array, rho::Real, n::Integer, t::Real, w::Real, j::Real, ctx::FFTW.cFFTWPlan)
    a = zeros(n)
    pd = Bernoulli(rho)
    ntr = round(Int, 400) # 400 is arbitrary, larger will make plot smoother
    mmse_arr = zeros(ntr)
    gamma = t/(n*(w+rho*j))
    for i=1:ntr
        rand!(pd, a)
        ahat = ctx * a
        mmse = 0
        for i=1:n
            r = abs(ahat[i])
            mmse += 1/((1/d[i])+gamma*r*r)
        end
        mmse_arr[i] = mmse
    end
    mmse = sum(mmse_arr)/ntr
    return mmse
end

function eval_mmse_lens(d::Array, n::Integer, t::Real, w::Real, j::Real)
    mmse = 0
    for i=1:n
        mmse += 1/((1/d[i])+((t*n)/(w+j)))
    end
    return mmse
end

function _mmse_power_sum(c::Real, d::Array, gamma::Real)
    n = length(d)
    p = 0
    for i=2:n
        p += max(0, (c-(1/d[i]))/gamma)
    end
    return p
end

function frac(x::Real)
    return x-floor(x)
end

"""
Waterfilling computation for power control.
"""
function waterfill_powers!(powers::Array, d::Array, rho::Real, n::Integer, t::Real, w::Real, j::Real)
    gamma = t/(n*(w+rho*j))
    # A hack for t=0
    if (gamma == 0)
        for i=1:n-1
            j = i+1
            powers[i] = 1/sqrt(n-1)
        end
        return
    end
    r = n*rho
    p = n*(floor(r)+frac(r)*frac(r))-r*r
    ub = (1/d[1])+gamma*p
    c = find_zero(x -> _mmse_power_sum(x,d,gamma)-p, (0,1.5*ub))
    for i=1:n-1
        j = i+1
        powers[i] = max(0, (c-(1/d[j]))/gamma)
    end
end

"""
Lower bound on LMMSE based on waterfilling.
"""
function mmse_lb(d::Array, rho::Real, n::Integer, t::Real, w::Real, j::Real)
    gamma = t/(n*(w+rho*j))
    # A hack for t=0
    if (gamma == 0)
        return sum(d)
    end
    r = n*rho
    p = n*(floor(r)+frac(r)*frac(r))-r*r
    ub = (1/d[1])+gamma*p
    c = find_zero(x -> _mmse_power_sum(x,d,gamma)-p, (0,1.5*ub))
    mmse = 0
    p = n*n*rho*rho
    mmse += 1/((1/d[1])+gamma*p)
    for i=2:n
        p = max(0, (c-(1/d[i])))
        mmse += 1/((1/d[i])+p)
    end
    return mmse
end

"""
Read the values of n for which we have spectrally flat constructions from
data/ into a dictionary of n->rho_arr[n].
"""
function get_spec_flat()
    fdict = Dict("data/1_2_vals.txt" => 0.5, "data/1_4_vals.txt" => 0.25, "data/1_8_vals.txt" => 0.125)
    spec_flat_dict = Dict{Int64,Array{Float64,1}}()
    for (file, rho) in fdict
        for line in eachline(file)
            n = parse(Int64, line)
            if !haskey(spec_flat_dict, n)
                spec_flat_dict[n] = [rho]
            else
                push!(spec_flat_dict[n], rho)
            end
        end
    end
    return spec_flat_dict
end

"""
Finally, the function that generates the plots of the paper, saved in "/tmp/mmse.pdf".
Specifically, set globally:
julia> n=677; theta=1; w=0.001; j=0.001; start=-50; stop=50; res=256;

For plot (a), do:
julia> beta = 1; gamma=0; plot_lmmse(n,theta,beta,gamma,w,j,start,stop,res);
For plot (b), do:
julia> beta = 0.02, gamma=0.005; plot_lmmse(n,theta,beta,gamma,w,j,start,stop,res);
"""
function plot_lmmse(n::Integer, theta::Real, beta::Real, gamma::Real, w::Real, j::Real, start::Real=0, len::Real=10, res=256)
    d = char_corr(theta, beta, gamma, n)
    fft_ctx = plan_fft(zeros(n))
    waterfill = zeros(n-1)
    x = range(start, length=res, stop=len)
    y_lb = zeros(length(x))
    popt = zeros(length(x))
    y_prng = zeros(length(x))
    y_spec = zeros(length(x))
    y_nazarov = zeros(length(x))
    spec_flat_dict = get_spec_flat()
    if !haskey(spec_flat_dict, n)
        error("No spectrally flat construction available at n=$n")
    else
        rho_arr = spec_flat_dict[n]
    end
    i = 1
    @showprogress for t in x
        t /= 10
        t = 10^t
        t *= n
        # lower bound
        res = Optim.optimize(rho->mmse_lb(d,rho,n,t,w,j), 0, 1)
        y_lb[i] = 10*log10(Optim.minimum(res))
        popt[i] = Optim.minimizer(res)
        waterfill_powers!(waterfill,d,popt[i],n,t,w,j)
        # random sequence
        res = Optim.optimize(rho->eval_mmse_rand(d,rho,n,t,w,j,fft_ctx),0,1)
        y_prng[i] = 10*log10(Optim.minimum(res))
        # spectrally flat, rho in rho_arr
        y_spec[i] = 10*log10(minimum([eval_mmse_flat(d,rho,n,t,w,j) for rho in rho_arr]))
        # Nazarov construction
        ahat_nazarov = nazarov_ahat(n, waterfill, fft_ctx)
        y_nazarov[i] = 10*log10(eval_mmse(d,n,t,w,j,ahat_nazarov))
        i += 1
    end
    # force pyplot, GR has issues with LaTeX labels
    pyplot()
    xlabel_str = L"10\log_{10}(t/n)"
    p = plot(x, [y_lb y_prng y_spec y_nazarov],
             label=["lower bound" "optimal random on-off" "spectrally flat" "Nazarov"],
             xlabel=xlabel_str,
             ylabel="LMMSE (dB)")
    plot_str = "/tmp/mmse.pdf"
    savefig(p, plot_str)
    q = plot(x, popt,
             title="optimal p vs time for iid scene at\n (n,theta,w,j)=($n,$theta,$w,$j)",
             label="optimal p",
             xlabel=xlabel_str,
             ylabel="p")
    plot_str = "/tmp/optp_iid.pdf"
    savefig(q, plot_str)
end


function besselj_zero_approx1(ν::Integer, m::Integer)
    μ = 4.0*ν^2
    a = (m+ν/2-1/4)*π
    a8 = 1/(8a)
    a82 = a8^2
    j = (a -(μ-1)*a8*(1 + a82*(4*(7μ-31)/3 +a82*( 32*(83*μ^2-982*μ+3779)/15 +
         a82*(64*(6949*μ^3-153855*μ^2+1585743*μ-6277237)/105))))
        )
end

function besseljprime_zero_approx1(ν::Integer, m::Integer)
    #if (ν==0 && m==1)
    #    return 0.0
    #end
    if ν==0
        m += 1
    end
    μ = 4*ν^2
    b = (m+ν/2-3/4)*π
    b8 = 1/(8*b)
    b82 = b8^2
    j = (b - b8*(μ+3 + b82*(4*(7*μ^2+82*μ-9)/3 +b82*(32*(83*μ^3+2075*μ^2-3039*μ+3537)/15 +
             b82*(64*(6949*μ^4+296492*μ^3-1248002*μ^2+7414380*μ-5853627)/105)))))
end

function airy_zero_approx(m::Integer)
    if m<=10
        z =[-2.338107410459767,
            -4.087949444130970,
            -5.520559828095552,
            -6.786708090071759,
            -7.944133587120853,
            -9.022650853340980,
           -10.04017434155809,
           -11.00852430373326,
           -11.93601556323626,
           -12.82877675286576][m]
    else
        t = 3/8*π*(4*m-1)
        t2 = t^(-2)
        z = -t^(2/3)*(1+t2*(5/48+t2*(-5/36+t2*(77125/82944+t2*(
            -108056875/6967296+ t2*162375596875/334430208)))))
    end
    z
end

function airyprime_zero_approx(m::Integer)
    if m<=10
        z = [-1.0187929716474710890,
             -3.2481975821798365379,
             -4.8200992111787356394,
             -6.1633073556394865476,
             -7.3721772550477701771,
             -8.4884867340197221329,
             -9.5354490524335474707,
            -10.527660396957407282,
            -11.475056633480245295,
            -12.384788371845747325][m]
    else
        t = 3/8*π*(4*m-3)
        t2 = t^(-2)
        z = -t^(2/3)*(1+t2*(-7/48+t2*(35/288+t2*(-181223/207360+t2*(
            18683371/1244160 - t2*9114588436/191102976)))))
    end
    z
end

function besselj_zero_approx2(ν::Integer, m::Integer; prime::Bool=false)
    if prime
        ζ = ν^(-2/3)*airyprime_zero_approx(m)
    else
         ζ = ν^(-2/3)*airy_zero_approx(m)
    end
    y = 2/3*(-ζ)^(3/2)
    if y>100000
        x = π/2
    elseif y<1
        p = (3*y)^(1/3)
        p2 = p^2
        x = p*(1+p2*(-2/15+p2*(3/175+p2*(-2/1575))))
    else
        p = 1/(y+π/2)
        p2 = p^2
        x = π/2 - p*(1+p2*(2/3+p2*(13/15+p2*(146/105+p2*(781/315+p2*16328/3465)))))
    end
    x2 = (y+x)^2
    r = (x-atan(x+y))/x2
    x = x - (1+x2)*r*(1+r/(x+y))

    z = 1/cos(x)
    h = sqrt(ζ*(1-z^2))
    if prime
        g1 = z/ζ*h*(7/(48*ζ) + h*(7/(z^2-1)+9
        )/24)
        j = ν*z + g1/ν
    else
        f1 = -z/ζ*h*(5/(48*ζ) + h*(5/(z^2-1)+3)/24)
        j = ν*z + f1/ν
    end
    j
end

function besselj_zero_approx(ν::Integer, m::Integer)
  if m>=ν
     j = besselj_zero_approx1(ν, m)
  else
     j = besselj_zero_approx2(ν, m)
  end
  j
end

function besseljprime_zero_approx(ν::Integer, m::Integer)
  if m>=ν
        j = besseljprime_zero_approx1(ν, m)
  else
        j = besselj_zero_approx2(ν, m, prime=true)
  end
  j
end

function besselj_zero_iter(ν::Integer, z::AbstractFloat)
    T = typeof(z)
    ɛ = eps(T)
    for i = 1:200
        J = besselj(ν, z)
        if abs(J) < 1000ɛ
            break
        end
        Jprime = besselj(ν-1, z) - ν*J/z
        z -= J/Jprime
        if i==200
            println("200 iterations, res=",abs(J))
        end
    end
    return z
end

function besseljprime_zero_iter(ν::Integer, z::AbstractFloat)
    T = typeof(z)
    ɛ = eps(T)
    for i = 1:200
        J = besselj(ν, z)
        J1 = besselj(ν-1, z) - ν*J/z
        if abs(J1) < 1000ɛ
            break
        end
        J2 = -J1/z + ((ν/z)^2-1)*J
        z -= J1/J2
        if i==200
            println("200 iterations, res=",abs(J1))
        end
    end
    return z
end

besselj_zero(ν, m; T=Float64) = besselj_zero_iter(ν, convert(T, besselj_zero_approx(ν, m)))

besseljprime_zero(ν::Integer, m::Integer; T=Float64) =
    besseljprime_zero_iter(ν, convert(T, besseljprime_zero_approx(ν, m)))


using LinearAlgebra
I(3)


using Distributions: truncated

## Extra distributions ##

@dist shifted_poisson(rate, offset::Int) = poisson(rate) + offset
@dist delta(val) = uniform_discrete(val, val)

struct TruncatedNormal <: Gen.Distribution{Float64} end

"""
    trunc_normal(mu::Real, std::Real, lb::Real, ub::Real)
Samples a `Float64` value from a normal distribution.
"""
const trunc_normal = TruncatedNormal()

(d::TruncatedNormal)(mu, std, lb, ub) = Gen.random(d, mu, std, lb, ub)

Gen.random(::TruncatedNormal, mu::Real, std::Real, lb::Real, ub::Real) =
    rand(truncated(Distributions.Normal(mu, std), lb, ub))

function Gen.logpdf(::TruncatedNormal, x::Real, mu::Real, std::Real, lb::Real, ub::Real)
    d = truncated(Distributions.Normal(mu, std), lb, ub)
    untrunc_lpdf = Distributions.logpdf(d.untruncated, x)
    if d.tp > 0
        untrunc_lpdf - d.logtp
    elseif cdf(d.untruncated, lb) ≈ 1.0
        untrunc_lpdf - Distributions.logccdf(d.untruncated, lb)
    elseif cdf(d.untruncated, ub) ≈ 0.0
        untrunc_lpdf - Distributions.logcdf(d.untruncated, ub)
    end
end

Gen.is_discrete(::TruncatedNormal) = false

## Other utility functions ##

function segment_find(idx, seglens)
    count = 0
    for (i, l) in enumerate(seglens)
        count += l
        if count >= idx return (i, idx-count+l) end
    end
    return nothing
end


## Logistic regression

module logistic_regression

using Distributions

struct data
    y::Vector
    x::Array{Float64}
    m::Vector
end

## Logistic regression nobs
function nobs(data::data)
    nx = size(data.x)[1]
    ny = length(data.y)
    nm = length(data.m)
    if (nx != ny)
        error("number of rows in of x is not equal to the length of y")
    elseif (nx != nm)
        error("number of rows in of x is not equal to the length of m")
    elseif (ny != nm)
        error("length of y is not equal to the length of m")
    end
    nx
end

function loglik(theta::Vector,
                data::data,
                i::Int64)
    eta = sum(data.x[i, :] .* theta)
    mu = exp.(eta)./(1 .+ exp.(eta))
    ll = data.y[i] .* log.(mu) + (data.m[i] - data.y[i]) .* log.(1 .- mu)
    ll
end

function ef(theta::Vector,
            data::data,
            i::Int64)
    eta = sum(data.x[i, :] .* theta)
    mu = exp.(eta)./(1 .+ exp.(eta))
    data.x[i, :] * (data.y[i] - data.m[i] * mu)
end

function simulate(theta::Vector,
                  x::Matrix{Float64},
                  m::Vector{Int64})
    n = size(x)[2]
    y = rand.(Binomial.(m, cdf.(Logistic(), x * theta)))
    data(y, x, m);
end

function loglik_alt(theta::Vector,
                    data::data,
                    i::Int64)
    eta = sum(data.x[i, :] .* theta)
    mu = exp.(eta)./(1 .+ exp.(eta))
    ll = data.y[i] .* log.(mu) + (data.m[i] - data.y[i]) .* log.(1 .- mu)
    ll - 10 * sum(theta.^2) / length(data.y)
end

function l2_penalty_obj(theta::Vector,
                        data::data)
    - 10 * sum(theta.^2)
end

function l2_penalty_ef(theta::Vector,
                       data::data)
    - 20 * theta
end

end


## Some Polynomial utilities

## conveniences for AbstractAlgebra.jl
# get value of GF(p) in (-p/2, p/2)
#function value(x::AbstractAlgebra.gfelem)
function value(x)
    # reach in to get a, q
    if hasfield(typeof(parent(x)), :p)
        q = parent(x).p
        a = x.d
        a <= q ÷ 2 ? a : a - q
    else
        x
    end
end

#value(x) = x

function poly_coeffs(f::AbstractAlgebra.Generic.Poly)
#    println("poly coef f=$f d=$(degree(f))")
    d = degree(f)
    d < 0 && return [value(coeff(f, 0))]
    [value(coeff(f, i)) for i in 0:degree(f)]
end

poly_coeffs(as) = as

variable(f::AbstractAlgebra.Generic.Poly) = gen(parent(f))

function maxnorm(f::AbstractAlgebra.Generic.Poly)
    norm(poly_coeffs(f), Inf)
end

function onenorm(f::AbstractAlgebra.Generic.Poly)
    norm(poly_coeffs(f), 1)
end

## Conversion to Generic.Poly
# If f = a0 + a1*y + ... an*y^n then compute
# a0*x + ... an*x^n for x a poly. (Could be different ring, ...)
function as_poly(f::AbstractAlgebra.Generic.Poly, x::AbstractAlgebra.Generic.Poly)
    degree(f) < 0 && return zero(x)
    sum(value(coeff(f, i)) * x^i for i in 0:degree(f))
end

function as_poly(as, x::AbstractAlgebra.Generic.Poly)
    sum( ai * x^(i-1) for  (i,ai) in enumerate(as))
end

# make poly over Z?
function as_poly(f::AbstractAlgebra.Generic.Poly{T}, x::String="x") where {T}
    R,y = ZZ[x]
    sum(value(coeff(f, i)) * y^i for i in 0:degree(f))
end

function as_poly(as, x0::String="x")
    R, x = ZZ[x0]
    as_poly(as, x)
end

# Z[x] with coefficients in (-p/2, p/2)
function _modp(x, p)
    a = mod(x,p)
    a > p ÷ 2 ? a - p : a
end

# make poly over x with coefficients in (-p/2...p/2)
function as_poly_modp(f, p, x::AbstractAlgebra.Generic.Poly)
    as = poly_coeffs(f)
    as_poly(_modp.(as, p), x)
end

# make poly in GP from coefficiens of x
# !!! NOTE
#     ?GF When check == false, no check is made, but the behaviour of the
#     resulting object is undefined if p is composite.
function as_poly_modp(f, p, x0::String="x")
    R, x = GF(p, check=false)[x0]
    as = poly_coeffs(f)
    as_poly(_modp.(as, p), x)
end

# Create a polynomial over Zp[x] from the coefficients
# as is (a0,a1, ..., an) as an interable
function as_poly_Zp(f::AbstractAlgebra.Generic.Poly, p, x)
    as_poly_Zp(poly_coeffs(f), p, x)
end

function as_poly_Zp(as, p::S, var="x") where {S}
    T = eltype(as)
    R = promote_type(T, S)
    q = convert(R, p)

    F = GF(q)
    R,x = PolynomialRing(F, var)
    sum(ai*x^(i-1) for (i,ai) in enumerate(as))
end


using ModelingToolkit
using DataFrames
using CSV
using Tables

# State variable and parameters
@variables t
@variables Rtg13A_c(t) Rtg13I_c(t) Rtg13A_n(t) Rtg13I_n(t) Rtg3I_c(t) Rtg3A_c(t) Rtg3A_n(t) Rtg3I_n(t) Rtg1_c(t) Rtg1_n(t)
@variables s(t) Rtg2I_c(t) Rtg2A_c(t)
@variables Mks(t) Rtg2Mks_c(t) Bmh(t) BmhMks(t)
@parameters ΣRtg1 ΣRtg2 ΣRtg3 ΣMks ΣBmh
@parameters ksV = 1.0 ksD = 1.0 n_S = 1.0 k2I = 1.0 mul_S = 1.0
@parameters k2M = 1.0 kn2M = 1.0 kBM = 1.0 knBM = 1.0
@parameters k13I = 1.0 k13IV = 1.0 k13ID = 1.0 k3A_c = 1.0 k3I_c = 1.0 k3I_n = 1.0
@parameters k13_c = 1.0 kn13_c = 1.0 k13_n = 1.0 kn13_n = 1.0
@parameters k1in = 1.0 k1out = 1.0 k3inA = 1.0 k3outA = 1.0 k3inI = 1.0 k3outI = 1.0

# Distributions of RTG1 and RTG3 proteins
rtg13_nucleus(u) = u[Rtg13I_n] + u[Rtg13A_n]
rtg13_cytosol(u) = u[Rtg13I_c] + u[Rtg13A_c]
rtg3_nucleus(u) = rtg13_nucleus(u) + u[Rtg3A_n] + u[Rtg3I_n]
rtg3_cytosol(u) = rtg13_cytosol(u) + u[Rtg3A_c] + u[Rtg3I_c]
rtg1_nucleus(u) = rtg13_nucleus(u) + u[Rtg1_n]
rtg1_cytosol(u) = rtg13_cytosol(u) + u[Rtg1_c]

# Mapping from symbols to MTK parameters
const symbol2mtk = Dict(Symbol(k) => k for k in (ΣRtg1, ΣRtg2, ΣRtg3, ΣMks, ΣBmh, ksV, ksD, n_S, k2I, mul_S, k2M, kn2M, kBM, knBM, k13I, k13IV, k13ID, k3A_c, k3I_c, k3I_n, k13_c, kn13_c, k13_n, kn13_n, k1in, k1out, k3inA, k3outA, k3inI, k3outI))

# Utility functions

"""Loads a CSV file and returns a RowTable (An tuple of NamedTuple for each row)"""
load_data(filename) = Tables.rowtable(dropmissing(CSV.read(filename, DataFrame), disallowmissing=true))

"""Load boolean conditions for model validation."""
load_conditions(filename="boolean_table_RTG13.csv") = load_data(joinpath(@__DIR__, "data", filename))

"""Load model parameters"""
function load_parameters(filename="solution_rtgMTK_optim.csv")
    paramsets = load_data(joinpath(@__DIR__, "data", filename))
    map(paramsets) do p
        Dict(symbol2mtk[k] => i for (k, i) in pairs(p))
    end
end


using Test, ApproxFun, SingularIntegralEquations

@testset "Nonlocal operators" begin
    α = 2.5
    δ = 0.1
    for sp in (Fourier(), SinSpace(), CosSpace(), Laurent(), Hardy{true}(PeriodicSegment()), Hardy{false}(PeriodicSegment()))
        L = NonlocalLaplacian(sp, α, δ)
        @test domain(L) == domain(sp)
        @test domainspace(L) == rangespace(L) == sp
        @test bandwidths(L) == (0,0)
    end

    NonlocalLaplacian(PeriodicSegment(), α, δ) === NonlocalLaplacian(Space(PeriodicSegment()), α, δ)

    L = NonlocalLaplacian(SinSpace(), α, δ)
    @test norm(L[10,10] + 9.836381919022900199967236210084047918127677209354580364785826568010653044652657e+01) < 100*norm(L[10,10])*eps()
    @test norm(L[1000,1000] + 1.670424124232011223460118736778899327539655416751949620343795973043019808627733e+05) < 100*norm(L[1000,1000])*eps()
end


using RobustAndOptimalControl, ControlSystems, MonteCarloMeasurements, Plots
unsafe_comparisons(true)

using RobustAndOptimalControl: δr, δc

# example in section 3.7.1, spinning satellite
a = 10
P = ss([0 a; -a 0], I(2), [1 a; -a 1], 0)
K = ss(I(2))

w = 2π .* exp10.(LinRange(-2, 2, 500))
S, PS, CS, T = RobustAndOptimalControl.gangoffour2(P, K)
sigmaplot(S, w, lab="S")
sigmaplot!(T, w, c=2, lab="T", ylims=(0.01, 45))
# Both sensitivity functions are very large, expect a non-robust system!

## Parametric uncertainty
a = 10*(1 + δr(100))
P = ss([0 a; -a 0], I(2), [1 a; -a 1], 0)

Sp, PSp, CSp, Tp = RobustAndOptimalControl.gangoffour2(P, K)
sigmaplot(Sp, w, lab="S")
sigmaplot!(Tp, w, c=2, lab="T", ylims=(0.01, 45))
# Not only are sensitivity functions large, they vary a lot under the considered uncertainty

## Add complex diagonal multiplicative input uncertainty
# With input uncertainty of magnitude ϵ < 1 / σ̄(T) we are guaranteed robust stability (even for “full-block complex perturbations")
a = 10
P = ss([0 a; -a 0], I(2), [1 a; -a 1], 0)

W0 = makeweight(0.2, 2)
W = I(2) + W0 * diagm([δc(100), δc(100)])
Ps = P*W
Ss, PSs, CSs, Ts = RobustAndOptimalControl.gangoffour2(Ps, K)
sigmaplot(Ss, w, lab="S")
sigmaplot!(Ts, w, c=2, lab="T", ylims=(0.01, 45))
# Under this uncertainty, the sensitivity could potentially be sky high., note how some of the 100 realizations peak much higher than the others

##

# 3.7.2 Motivating robustness example no. 2: Distillation Process
M = [87.8 -86.4; 108.2 -109.6]
G = ss(tf(1, [75, 1])) * M
RGA = relative_gain_array(G, 0)
sum(abs, RGA) # A good estimate of the true condition number, which is 141.7
# large elments in the RGA indicate a process that is difficult to control

# We consider the following inverse-based controller, which may also be looked
# upon as a steady-state decoupler with a PI controlle
k1 = 0.7
Kinv = ss(tf(k1*[75, 1], [1, 0])) * inv(M)

# reference filter
F = tf(1, [5, 1])

sigmaplot(input_sensitivity(G, Kinv), w)
sigmaplot!(output_sensitivity(G, Kinv), w, c=2)
# Sensitivity looks nice, how about step response
plot(step(feedback(G*Kinv)*F, 20))
# Looks excellent..

# We consider again the input gain uncertainty as in the previous example,
# and we select ϵ1 = 0.2 and ϵ2 = 0.2. We then have
G′ = G * diagm([1 + 0.2, 1 - 0.2])
plot!(step(feedback(G′*Kinv)*F, 20), l=:dash)
# Looks very poor! The system was not robust to simultaneous input uncertainty!

# We can also do this with a complex, diagonal input uncertainty that grows with frequency
W0 = makeweight(0.2, 1, 2) # uncertainty goes from 20% at low frequencies to 200% at high frequencies
W = I(2) + W0 * diagm([δc(100), δc(100)])
Gs = G*W
res = step(c2d(feedback(Gs*Kinv)*F, 0.01), 20)
mcplot!(res.t, abs.(res.y[:, :, 1]'), alpha=0.3)
mcplot!(res.t, abs.(res.y[:, :, 2]'), alpha=0.3)
# The system is very sensitive to complex input uncertainty

# How about real input uncertainty?
W = I(2) + W0 * diagm([δr(100), δr(100)])
Gs = G*W

plot(step(feedback(G*Kinv)*F, 20))
plot!(step(feedback(G′*Kinv)*F, 20), l=:dash)
res = step(c2d(feedback(Gs*Kinv)*F, 0.01), 20)
mcplot!(res.t, abs.(res.y[:, :, 1]'), alpha=0.3)
mcplot!(res.t, abs.(res.y[:, :, 2]'), alpha=0.3)
# Also bad

Si = input_sensitivity(Gs, Kinv)
sigmaplot(Si, w, c=1, lab="Si")
So = output_sensitivity(Gs, Kinv)
sigmaplot!(So, w, c=2, lab="So")
# The sensitivity at the plant output is enormous. A low sensitivity with the nominal system does not guarantee robustness!


module Tutorial5
using FEniCS

mesh = UnitSquareMesh(8, 8)
V = FunctionSpace(mesh, "P", 1)
#in the original tutorial, these are calculated via SymPy
#although it has no actual FEniCS code in calculating those, so I have simply
#stated the result
u_code = "x[0]+2*x[1]+1"
f_code = "-10*x[0]-20*x[1]-10"

calc(var) = 1+var*var

u_D = Expression(u_code, degree=2)
bc = DirichletBC(V, u_D, "on_boundary")
u=FeFunction(V)
v = TestFunction(V)
f = Expression(f_code, degree=2)
F = calc(u)*dot(grad(u), grad(v))*dx - f*v*dx
nlvsolve(F,u,bc)
end#module


# # Fitting and Predicting - Frequency Data
#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/notebooks/fitpredictFreq.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/notebooks/fitpredictFreq.ipynb)

# RHEOS can also fit models to dynamic mechanical analysis data from oscillatory tests.

using RHEOS

# ## Fitting

# #### Step 1: Loading Experimental Data

# RHEOS has a convenient function to import data from CSV files (more information in [File I/O](@ref) section).

data = importcsv("assets/data_freq.csv", ω_col = 1, Gp_col = 2, Gpp_col = 3)
rheofreqdatatype(data)

# The frequency, storage modulus and loss modulus are stored into a [`RheoFreqData`](@ref) struct.

# #### Step 2: Model Fit

# Let's fit a Fractional Kelvin-Voigt model. The first argument is our data, the second argument tells RHEOS which model to fit ([`RheoModelClass`](@ref)). The function will return a [`RheoModel`](@ref) object, i.e. a [`RheoModelClass`](@ref) with fixed values for the parameters. The function [`dynamicmodelfit`](@ref) possesses all the additional arguments (lower and upper bounds, initial parameters, tolerance) as for the fitting of time data. Note that initial parameters, upper bounds and lower bounds can be partially specified and the rest will be filled in automatically.

## Lower bounds
Lo = (cₐ = 0.0, a = 0.01, cᵦ = 0.0, β = 0.01)
## Upper bounds
Hi = (cₐ = 1e2, a = 0.99, cᵦ = 1e2, β = 0.99)
## Initial parameters
P0 = (cₐ = 1.0, a = 0.36, cᵦ = 1.0, β = 0.03)

FractKV_model = dynamicmodelfit(data, Fract_KelvinVoigt, weights = "log", lo = Lo, hi = Hi, p0 = P0)

# For the fitting process RHEOS relies on the optimistion package [NLopt.jl](https://nlopt.readthedocs.io/en/latest/). RHEOS makes use of a local derivative free algorithm, specifically the Tom Rowan's "Subplex" algorithm originally introduced in his [PhD thesis at The University of Texas at Austin, 1990](https://dl.acm.org/doi/book/10.5555/100816).

# The storage and loss moduli can sometimes occupy different orders of magnitude. This can cause problems during fitting as the optimisation routine will weight errors at the higher orders of magnitude more strongly than those at the lower orders of magnitude. This is a more general problem often faced during multi-objective optimisation problems. RHEOS offers four rescaling options (weight):
# - The `local` option rescales the cost at each point by the point itself.
# - The `log` approach simply re-scales all the storage and loss moduli and their predicted values logarithmically before finding the error between them.
# - A third option in RHEOS divides the cost at each point by the mean value of the storage and loss modulus respectively depending on which cost is being calculated. This can work well but performance is hindered if the storage or loss moduli individually vary over many orders of magnitude (`mean`).
# - The fourth option offered by RHEOS is simply manual weightings provided by the user for each modulus (`manual`).
# In the first two cases, the benefit arises from the fact that the optimisation weighting is rebalanced in favour of smaller values. The logarithmic rescaling method seems to work particularly well but runs into problems if the storage or loss moduli are exactly 0 at any frequency, this is due to the negative singularity of the logarithmic function for the 0 argument. 0 arguments can also cause problems using the `local` option due to division by 0.

# ## Predicting

# RHEOS allows the user to simulate the frequency response of a model (with defined parameters, [`RheoModel`](@ref) struct) to an imposed loading.

# Given an incomplete data set (frequency only) and model with values substituted into parameters ([`RheoModel`](@ref)), return a new "complete" dataset based on the model with the simulated missing variable (storage and loss).

# #### Assessing the Quality of Fit

# The ability of predicting model's response is first exploited to assess the quality of the fits above. The incomplete [`RheoFreqData`](@ref) variable is defined by extracting the frequency from the original data or by defining a new frequency vector.

data_ext = onlyfreq(data)
## Alternatively, a RheoFreqData with only frequency data can be generated as
## data_ext = frequencyspec(ω_start = 1.0e-2, ω_end = 1.0e2, logstep = 0.1)

rheofreqdatatype(data_ext)

# For the prediction, RHEOS' function [`dynamicmodelpredict`](@ref) requires the incomplete data set and a model with fixed parameters ([`RheoModel`](@ref)). For the assessment of the fitting quality the [`RheoModel`](@ref) is the output of the fitting function.

fractKV_predict = dynamicmodelpredict(data_ext, FractKV_model)

## Now we can plot data and model together for comparison
using PyPlot
fig, ax = subplots(1, 1, figsize = (7, 5))
ax.loglog(data.ω, data.Gp, "o", markersize = 5, color = "blue")
ax.loglog(data.ω, data.Gpp, "o", markersize = 5, color = "red")
ax.loglog(fractKV_predict.ω, fractKV_predict.Gp, "--", color = "blue")
ax.loglog(fractKV_predict.ω, fractKV_predict.Gpp, "--", color = "red")
ax.set_xlabel("Frequency")
ax.set_ylabel("Storage and Loss moduli")
#!nb fig #hide

# #### Simulate Different Model Behaviours

# The ability of predicting model's response can be exploited to simulate the behaviour of the model to other external loading conditions. For example, we will explore the response of the fitted model to a creep experiment.

# First we need to define an incomplete [`RheoTimeData`](@ref) struct, which can be achieved via the data generation functions provided in RHEOS (see section [Generating Data](@ref)).

## Define timeline
dσ = timeline(t_end = 10)
## and a step in stress
dσ = stressfunction(dσ, hstep())

## we can now predict the creep response of the Maxwell model
FractKV_creepPredict = modelsteppredict(dσ, FractKV_model)
## Visualisation of the simulated response
fig, ax = subplots(1, 1, figsize = (7, 5))
ax.plot(FractKV_creepPredict.t, FractKV_creepPredict.ϵ)
ax.set_xlabel("Time")
ax.set_ylabel("Strain")
#!nb fig #hide

# **Reference frequency data**: Deng, Linhong, et al. "Fast and slow dynamics of the cytoskeleton." Nature materials 5.8 (2006): 636.


# ---
# title: 34. Find First and Last Position of Element in Sorted Array
# id: problem34
# author: zhwang
# date: 2022-02-13
# difficulty: Medium
# categories: Array, Binary Search
# link: <https://leetcode.com/problems/find-first-and-last-position-of-element-in-sorted-array/description/>
# hidden: true
# ---
#
# Given an array of integers `nums` sorted in ascending order, find the starting
# and ending position of a given `target` value.
#
# If `target` is not found in the array, return `[-1, -1]`.
#
# **Follow up:**  Could you write an algorithm with `O(log n)` runtime
# complexity?
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [5,7,7,8,8,10], target = 8
#     Output: [3,4]
#
#
# **Example 2:**
#
#
#
#     Input: nums = [5,7,7,8,8,10], target = 6
#     Output: [-1,-1]
#
#
# **Example 3:**
#
#
#
#     Input: nums = [], target = 0
#     Output: [-1,-1]
#
#
#
#
# **Constraints:**
#
#   * `0 <= nums.length <= 105`
#   * `-109 <= nums[i] <= 109`
#   * `nums` is a non-decreasing array.
#   * `-109 <= target <= 109`
#
#
## @lc code=start
using LeetCode

function search_range(nums::Vector{Int}, target::Int)::Vector{Int}
    function get_left_border(target)
        left, right = 1, n
        while left <= right
            mid = (left + right) >> 1
            if nums[mid] >= target
                right = mid - 1
            else
                left = mid + 1
            end
        end
        return left
    end
    n = length(nums)
    left_border = get_left_border(target)
    (left_border > n || nums[left_border] != target) && return [-1, -1]
    return [left_border - 1, get_left_border(target + 1) - 2]
end
## @lc code=end


# Starting random numbers example for Learning Julia
using Random

# generate a random number between 0 and 1
println(rand(), "\n")

# pick a random element from a given collection
for i in 1:3
    println(rand([1, 2, 3, 4, 5, 6])) # choosing 3 random items from the specified array
end
println()

# populate an array with random values
arr = rand(UInt8, 5) # populate an array with 5 unsigned Int8
println(arr, "\n")

# generate a random string
println(randstring("ABCD")) # default: generate a size 8 random string
                            # selected from the specified string
println(randstring('a':'z', 6))
println()

# shuffle a list of elements
vowels = ["A", "E", "I", "O", "U"]
Random.shuffle!(vowels)
println(vowels, "\n")

# use the seed function to control the random number generator
Random.seed!(1000)
println(rand(), ", ", rand())
println("--------")
Random.seed!(1000)
println(rand(), ", ", rand())


# example script for running ChhabraJensen (direct) method

using DataFrames, CSV, GLM, Plots, Plots.PlotMeasures

include("ChhabraJensen.jl")

dataFile = DataFrame(CSV.File("exampleTS.csv", header = false))

dataSeries = dataFile[:,1] |> vec

# check for negative points
minimum(dataSeries)

#= in this example, we have negative values. ChhabraJensen requires positive values so we need to perform a transform. Per Kelty-Stephen et al.
simply adding by some constant is not sufficient. Per [[citation needed]], we can perform a sigmoid transoform:
=#

dataSeries = 1 ./(1 .+ ℯ .^(1 .* dataSeries))

# parameters
qValues = [-10:10;]
scales = [3:10;]

mf_results = ChhabraJensen(dataSeries, qValues, scales)

# Plots

# mf-spectrum plot
mf_spectrum_plot = plot(mf_results["alpha"], mf_results["falpha"],seriestype=:scatter, legend=false, xlabel = "α(q)", ylabel = "f(q)", title=string("mfw: ", round(mf_results["mfw"];digits=3)))

# q-order plot
Dq_q_plot = plot(mf_results["Dq"], mf_results["q"], seriestype=:scatter, legend=false, xlabel = "q", ylabel = "Dq")

# combined plots
plot(Dq_q_plot, mf_spectrum_plot, size = (1000,500), margin = 25px)


@testset "Truth Tables" begin
    a, b, c = LogicalSymbol.([:a, :b, :c])

    # a ∧ b  (a and b)
    st = a ∧ b
    and_tt = truthtable(st)
    @test all(and_tt[!, :a] .== [false, true, false, true])
    @test all(and_tt[!, :b] .== [false, false, true, true])
    @test all(and_tt[!, string(st)] .== [false, false, false, true])

    # b ∨ a  (b or a)
    st = b ∨ a
    or_tt = truthtable(st)
    @test all(or_tt[!, :b] .== [false, false, true, true])
    @test all(or_tt[!, :a] .== [false, true, false, true])
    @test all(or_tt[!, string(st)] .== [false, true, true, true])

    # material implication
    st = a → b
    mi_tt = truthtable(st)
    @test all(mi_tt[!, string(st)] .== [true, false, true, true])

    # slightly more complex statement
    st = (a → b) ∧ c
    mi_tt = truthtable(st)
    @test all(mi_tt[!, string(st)] .== [false, false, false, false, true, false, true, true])
end


@defcomp sealeveldamages begin
    regions = Index()

    SLRDAMAGES = Variable(index=[time, regions])

    slrmultiplier = Parameter(index=[regions])
    slrelasticity = Parameter(index=[regions])
    slrdamlinear = Parameter(index=[regions])
    slrdamquadratic = Parameter(index=[regions])
    TOTALSLR = Parameter(index=[time])
    YGROSS = Parameter(index=[time, regions])

    function run_timestep(p, v, d, t)
        #Define function for SLRDAMAGES
        for r in d.regions
            if is_first(t)
                v.SLRDAMAGES[t,r] = 0.
            else
                v.SLRDAMAGES[t,r] = 100. * p.slrmultiplier[r] * (p.TOTALSLR[t-1] * p.slrdamlinear[r] + p.TOTALSLR[t-1]^2 * p.slrdamquadratic[r]) * (p.YGROSS[t-1,r] / p.YGROSS[TimestepIndex(1),r])^(1/p.slrelasticity[r])
            end
        end
    end
end


using AngularMomentumAlgebra.Dipoles

@testset "Wigner–Eckart" begin
    @testset "Dipoles" begin
        x,y,z = 𝐫̂

        os = os"k[s-h]"

        @testset "z" begin
            for o in os
                ℓ = o.ℓ
                mℓs = -ℓ:ℓ
                for mℓ in mℓs
                    @test matrix_element((ℓ,mℓ),z,(ℓ+1,mℓ)) ≈ √((ℓ+mℓ+1)*(ℓ-mℓ+1)/((2ℓ+3)*(2ℓ+1)))
                end
            end
        end

        @testset "x" begin
            for o in os
                ℓ = o.ℓ
                mℓs = -ℓ:ℓ
                for mℓ in mℓs
                    a = SpinOrbital(o, mℓ, half(1))
                    b = SpinOrbital(Orbital(o.n, ℓ+1), mℓ+1, half(1))
                    # TODO Figure out where the minus sign comes from
                    @test dot(a,x,b) ≈ -0.5*√((ℓ+mℓ+2)*(ℓ+mℓ+1)/((2ℓ+3)*(2ℓ+1)))
                end
            end
        end
    end
end


@testset "Polynomial equality constraint with domain with $(factory.constructor)" for factory in sdp_factories
    @polyvar(x, y)

    m = SOSModel(factory)

    @constraint(m, x^3 + x*y^2 == x, domain=(@set x^2 + y^2 >= 1 && x^2 + y^2 <= 1))

    JuMP.optimize!(m)

    @test JuMP.primal_status(m) == MOI.FeasiblePoint
end


module FDTD1D
export make_advance_fields, YeeMesh1D

struct YeeMesh1D{T<:AbstractFloat, S<:Integer}
    ey::Vector{T}
    ez::Vector{T}
    hy::Vector{T}
    hz::Vector{T}
    N::S
    function YeeMesh1D{T}(N::S) where {T, S}
        new{T,S}(zeros(T, N+1), zeros(T, N+1), zeros(T, N), zeros(T, N), N)
    end
end

struct PML{T<:Integer, S<:AbstractFloat}
    N::T
    σ_max::S
    r₁::Vector{S}
    r₂::Vector{S}
    function PML(N::T, σ_max::S, Δx, Δt) where {T, S}
        σ = [σ_max*(i/2N)^3 for i = 1:2N]
        r₁ = exp.(-Δt.*σ)
        r₂ = (1.0 .- r₁)./(Δx.*σ)
        new{T,S}(N, σ_max, r₁, r₂)
    end
end

function make_advance_fields(f::YeeMesh1D{T,S}, cfl, pulse_shape, Δt, Δx, x_min, pml::PML = PML(10,1e3,Δt,Δx)) where {T,S}
    Nx = f.N

    function generate_fields_x_min!(t)
        f.ey[pml.N+2] -= cfl*pulse_shape.y(t, x_min + Δx)
        f.ez[pml.N+2] -= cfl*pulse_shape.z(t, x_min + Δx)

        f.hz[pml.N+2] -= cfl*pulse_shape.y(t + 0.5*Δt, x_min + 1.5*Δx)
        f.hy[pml.N+2] -= cfl*pulse_shape.z(t + 0.5*Δt, x_min + 1.5*Δx)
    end

    function update_ey!(jy)
        for i = 2:pml.N+1
            f.ey[i] = pml.r₁[1+2*(pml.N-i+1)]*f.ey[i] - pml.r₂[1+2*(pml.N-i+1)]*(f.hz[i] - f.hz[i-1])
        end
        for i = pml.N+2:Nx-pml.N
            f.ey[i] -= cfl*(f.hz[i] - f.hz[i-1])
        end
        for i = Nx-pml.N+1:Nx
            f.ey[i] = pml.r₁[1+2*(i-Nx+pml.N-1)]*f.ey[i] - pml.r₂[1+2*(i-Nx+pml.N-1)]*(f.hz[i] - f.hz[i-1])
        end
        for i = 1:Nx
            f.ey[i] += jy[i]
        end
    end

    function update_ez!(jz)
        for i = 2:pml.N+1
            f.ez[i] = pml.r₁[1+2*(pml.N-i+1)]*f.ez[i] + pml.r₂[1+2*(pml.N-i+1)]*(f.hy[i] - f.hy[i-1])
        end
        for i = pml.N+2:Nx-pml.N
            f.ez[i] += cfl*(f.hy[i] - f.hz[i-1])
        end
        for i = Nx-pml.N+1:Nx
            f.ez[i] = pml.r₁[1+2*(i-Nx+pml.N-1)]*f.ez[i] + pml.r₂[1+2*(i-Nx+pml.N-1)]*(f.hy[i] - f.hy[i-1])
        end
        for i = 1:Nx
            f.ez[i] += jz[i]
        end
    end

    function update_hy!()
        for i = 1:pml.N
            f.hy[i] = pml.r₁[2*(pml.N-i+1)]*f.hy[i] + pml.r₂[2*(pml.N-i+1)]*(f.ez[i+1] - f.ez[i])
        end
        for i = pml.N+1:Nx-pml.N
            f.hy[i] += cfl*(f.ez[i+1] - f.ez[i])
        end
        for i = Nx-pml.N+1:Nx
            f.hy[i] = pml.r₁[2*(i-Nx+pml.N)]*f.hy[i] + pml.r₂[2*(i-Nx+pml.N)]*(f.ez[i+1] - f.ez[i])
        end
    end

    function update_hz!()
        for i = 1:pml.N
            f.hz[i] = pml.r₁[2*(pml.N-i+1)]*f.hz[i] - pml.r₂[2*(pml.N-i+1)]*(f.ey[i+1] - f.ey[i])
        end
        for i = pml.N+1:Nx-pml.N
            f.hz[i] -= cfl*(f.ey[i+1] - f.ey[i])
        end
        for i = Nx-pml.N+1:Nx
            f.hz[i] = pml.r₁[2*(i-Nx+pml.N)]*f.hz[i] - pml.r₂[2*(i-Nx+pml.N)]*(f.ey[i+1] - f.ey[i])
        end
    end

    function make_step!(j)
        update_ey!(j.y)
        update_ez!(j.z)

        update_hz!()
        update_hy!()
    end

    function advance_fields!(t, j)
        generate_fields_x_min!(t)
        make_step!(j)
    end

    return advance_fields!
end

end


function majority(s::Array)
    # Create a counter dictionary Dict(element => number of times element appears in input array)
    counter = Dict()
    length_s = 0
    # Count the number of times each element in s appears
    for i in s
       length_s += 1
       if i in keys(counter)
           counter[i] += 1
       else
           counter[i] = 1
       end
    end
    # Loop through the items in the counter dictionary, return majority element
    for (key, value) in counter
       if value > length_s/2
           return key
       end
    end
    # If no majority element, return a message:
    return "No majority element"
end

# The example case from the prompt
println(majority([2,2,3,7,5,7,7,7,4,7,2,7,4,5,6,7,7,8,6,7,7,8,10,12,29,30,19,10,7,7,7,7,7,7,7,7,7]))
# And another case
println(majority([1,1,1,2,3,4,1,1,1]))
# And a case with no majority element
println(majority([9,9,9,9,8,8,8,8]))


using StatPlots, StatsBase, DataFrames, RDatasets, Distributions, ProgressMeter#, Iterators

doc_path(args...) = joinpath(@__DIR__, args...)

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

gr()
# define the Lorenz attractor
mutable struct Lorenz
    dt; σ; ρ; β; x; y; z
end

function step!(l::Lorenz)
    dx = l.σ*(l.y - l.x)       ; l.x += l.dt * dx
    dy = l.x*(l.ρ - l.z) - l.y ; l.y += l.dt * dy
    dz = l.x*l.y - l.β*l.z     ; l.z += l.dt * dz
end

attractor = Lorenz((dt = 0.02, σ = 10., ρ = 28., β = 8//3, x = 1., y = 1., z = 1.)...)

# initialize a 3D plot with 1 empty series
plt = plot3d(1, xlim=(-25,25), ylim=(-25,25), zlim=(0,50),
            title = "Lorenz Attractor", marker = 2)

# build an animated gif by pushing new points to the plot, saving every 10th frame
anim = @gif for i=1:1500
    step!(attractor)
    push!(plt, attractor.x, attractor.y, attractor.z)
end every 10

mv(anim.filename, doc_path("index", "lorenz_attractor.gif"), force = true)

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

pyplot(leg=false, ticks=nothing) #change to the pyplot backend and define some defaults
x = y = range(-5, stop = 5, length = 40)
zs = zeros(0,40)
n = 100

# create a progress bar for tracking the animation generation
prog = Progress(n,1)

anim = @gif for i in range(0, stop = 2π, length = n)
    f(x,y) = sin(x + 10sin(i)) + cos(y)

    # create a plot with 3 subplots and a custom layout
    l = Plots.@layout [a{0.7w} b; c{0.2h}]
    p = plot(x, y, f, st = [:surface, :contourf], layout=l)

    # induce a slight oscillating camera angle sweep, in degrees (azimuth, altitude)
    plot!(p[1], camera=(15*cos(i), 40))

    # add a tracking line
    fixed_x = zeros(40)
    z = map(f, fixed_x, y)
    plot!(p[1], fixed_x, y, z, line = (:black, 5, 0.2))
    vline!(p[2], [0], line = (:black, 5))

    # add to and show the tracked values over time
    global zs = vcat(zs, z')
    plot!(p[3], zs, alpha = 0.2, palette = cgrad(:blues).colors)

    # increment the progress bar
    next!(prog)
end

mv(anim.filename, doc_path("index", "waves.gif"), force = true)

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

gr(leg=false, bg=:lightgrey)

# Create a filled contour and boxplot side by side.
plot(contourf(randn(10,20)), boxplot(rand(1:4,1000),randn(1000)))

# Add a histogram inset on the heatmap.
# We set the (optional) position relative to bottom-right of the 1st subplot.
# The call is `bbox(x, y, width, height, origin...)`, where numbers are treated as "percent of parent"
histogram!(randn(1000), inset = (1, bbox(0.05,0.05,0.5,0.25,:bottom,:right)), ticks=nothing, subplot=3, bg_inside=nothing)

# Add sticks floating in the window (inset relative to the window, as opposed to being relative to a subplot)
sticks!(randn(100), inset = bbox(0,-0.2,200Plots.px,100Plots.px,:center), ticks=nothing, subplot=4)
png(doc_path("layout", "layout_5"))
theme(:default)

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

gr()
# load a dataset
iris = dataset("datasets", "iris");

# Scatter plot with some custom settings
@df iris scatter(:SepalLength, :SepalWidth, group=:Species,
    title = "My awesome plot",
    xlabel = "Length", ylabel = "Width",
    m=(0.5, [:cross :hex :star7], 12),
    bg=RGB(.2,.2,.2))

png(doc_path("index", "iris"))

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

# iris = dataset("datasets","iris")
@df iris marginalhist(:PetalLength, :PetalWidth)
png(doc_path("tutorial", "plot_recipes"))

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

@df iris scatter(:SepalLength, :SepalWidth, group=:Species,
    m=(0.5, [:+ :h :star7], 12), bg=RGB(.2,.2,.2))
png(doc_path("input", "dataframes"))

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

dfr = DataFrame(a = 1:10, b = 10*rand(10), c = 10 * rand(10))
# Plot the DataFrame by declaring the points by the column names
@df dfr plot(:a, [:b :c]) # x = :a, y = [:b :c]. Notice this is two columns!
png(doc_path("tutorial", "user_recipes_1"))

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

@df dfr scatter(:a, :b, title="My DataFrame Scatter Plot!") # x = :a, y = :b
png(doc_path("tutorial", "user_recipes_2"))

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

gr()
x = 1:10; y = rand(10); # These are the plotting data
plot(x,y)
png(doc_path("tutorial", "lines_1"))

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

x = 1:10; y = rand(10,2) # 2 columns means two lines
plot(x,y)
png(doc_path("tutorial", "lines_2"))

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

z = rand(10)
plot!(x,z)
png(doc_path("tutorial", "lines_3"))

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

x = 1:10; y = rand(10,2) # 2 columns means two lines
p = plot(x,y)
z = rand(10)
plot!(p,x,z)
png(doc_path("tutorial", "lines_4"))

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

x = 1:10; y = rand(10,2) # 2 columns means two lines
plot(x,y,title="Two Lines",label=["Line 1" "Line 2"],lw=3)
png(doc_path("tutorial", "attr_1"))

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

xlabel!("My x label")
png(doc_path("tutorial", "attr_2"))

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

x = 1:10; y = rand(10,2) # 2 columns means two lines
plotlyjs() # Set the backend to Plotly
plot(x,y,title="This is Plotted using Plotly")
png(doc_path("tutorial", "backends_1"))

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

gr() # Set the backend to GR
plot(x,y,title="This is Plotted using GR") # This plots using GR
png(doc_path("tutorial", "backends_2"))

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

gr() # We will continue onward using the GR backend
plot(x,y,seriestype=:scatter,title="My Scatter Plot")
png(doc_path("tutorial", "scatter_1"))

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

gr()
scatter(x,y,title="My Scatter Plot")
png(doc_path("tutorial", "scatter_2"))

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

y = rand(10,4)
plot(x,y,layout=(4,1))
png(doc_path("tutorial", "subplots_1"))

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

p1 = plot(x,y) # Make a line plot
p2 = scatter(x,y) # Make a scatter plot
p3 = plot(x,y,xlabel="This one is labelled",lw=3,title="Subtitle")
p4 = histogram(x,y) # Four histograms each with 10 points? Why not!
plot(p1,p2,p3,p4,layout=(2,2),legend=false)
png(doc_path("tutorial", "subplots_2"))

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

plot(Normal(3,5),lw=3)
png(doc_path("tutorial", "type_recipes"))

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

y = rand(100,4) # Four series of 100 points each
violin(["Series 1" "Series 2" "Series 3" "Series 4"],y,leg=false)
png(doc_path("tutorial", "series_recipes_1"))

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

boxplot!(["Series 1" "Series 2" "Series 3" "Series 4"],y,leg=false)
png(doc_path("tutorial", "series_recipes_2"))

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

gr()

# 10 data points in 4 series
xs = 0 : 2π/10 : 2π
data = [sin.(xs) cos.(xs) 2sin.(xs) 2cos.(xs)]

# We put labels in a row vector: applies to each series
labels = ["Apples" "Oranges" "Hats" "Shoes"]

# Marker shapes in a column vector: applies to data points
markershapes = [:circle, :star5]

# Marker colors in a matrix: applies to series and data points
markercolors = [:green :orange :black :purple
            :red   :yellow :brown :white]

plot(xs, data, label = labels, shape = markershapes, color = markercolors,
 markersize = 10)

png(doc_path("input", "columns_are_series"))

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

plotlyjs()

function rectangle_from_coords(xb,yb,xt,yt)
    [
        xb yb
        xt yb
        xt yt
        xb yt
        xb yb
        NaN NaN
    ]
end

some_rects=[
    rectangle_from_coords(1 ,1 ,5 ,5 )
    rectangle_from_coords(10,10,15,15)
    ]
other_rects=[
    rectangle_from_coords(1 ,10,5 ,15)
    rectangle_from_coords(10,1 ,15,5 )
    ]

plot(some_rects[:,1], some_rects[:,2],label="some group")
plot!(other_rects[:,1], other_rects[:,2],label="other group")

png(doc_path("input", "unconnected"))

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

# itr = repeatedly(()->rand(10), 20)
# anim = animate(itr, ylims=(0,1), c=:red, fps=5)
# gif(anim, doc_path("animations", "rand_anim.gif"), fps = 15)

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

gr(size=(300,300), leg=false)

@userplot MyCount
Plots.@recipe function f(mc::MyCount)
    # get the array from the args field
    arr = mc.args[1]

    T = typeof(arr)
    if T.parameters[1] == Float64
        seriestype := :histogram
        arr
    else
        seriestype := :bar
        cm = countmap(arr)
        x = sort!(collect(keys(cm)))
        y = [cm[xi] for xi=x]
        x, y
    end
end
mycount(rand(500))

png(doc_path("contributing", "my_count_1"))

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

mycount(rand(["A","B","C"],100))
png(doc_path("contributing", "my_count_2"))

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

mycount(rand(1:500, 500))
png(doc_path("contributing", "my_count_3"))

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

pyplot()
plot(x, y, line = (0.5, [4 1 0], [:path :scatter :density]),
    marker=(10, 0.5, [:none :+ :none]), fill=0.5,
    orientation = :h, title = "My title")
png(doc_path("pipeline", "pipeline_1"))

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

gr()
mutable struct MyVecWrapper
    v::Vector{Float64}
end
mvw = MyVecWrapper(rand(100))

@recipe function f(mvw::MyVecWrapper)
    markershape --> :circle
    markersize  --> 30
    mvw.v
end

plot(
    plot(mvw.v),
    plot(mvw)
    )

png(doc_path("pipeline", "pipeline_2"))

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

scatter(rand(100), group = rand(1:3, 100), marker = (10,0.3,[:s :o :x]))
png(doc_path("pipeline", "pipeline_3"))

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

# create a 2x2 grid, and map each of the 4 series to one of the subplots
plot(rand(100,4), layout = 4)
png(doc_path("layout", "layout_1"))

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

# create a 4x1 grid, and map each of the 4 series to one of the subplots
plot(rand(100,4), layout = (4,1))
png(doc_path("layout", "layout_2"))

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

plot(rand(100,4), layout = grid(4,1,heights=[0.1,0.4,0.4,0.1]))
png(doc_path("layout", "layout_3"))

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

l = Plots.@layout [  a{0.3w} [grid(3,3)
                                 b{0.2h} ]]
plot(
    rand(10,11),
    layout = l, label = "", seriestype = [:bar :scatter :path],
    title = ["($i)" for j = 1:1, i=1:11], titleloc = :right, titlefont = font(8)
    )
png(doc_path("layout", "layout_4"))


using Pkg
Pkg.activate(".")

using Plots, LaTeXStrings, PrettyTables, DataFrames, LinearAlgebra,
		  BenchmarkTools, ForwardDiff, Printf, Interact, Markdown,
          Logging

import SIAMFANLEquations

function ata_table(data::AbstractMatrix, labels::AbstractVector;
							format = :md, kwargs...)
	if format == :md
		return pretty_table(String, data, labels;
					 backend=:text,
					 tf = tf_markdown,
					 kwargs...) |> Markdown.parse
	elseif format == :html
		return pretty_table(String, data, labels;
					 backend=:html, tf=tf_html_minimalist,
					 nosubheader=true,
					 kwargs...) |> HTML
	else
		error("unknown table format")
	end
end

function ata_table(args...; T = Float64, format = :md, kwargs...)

	labels = []
	data = Matrix{T}(undef, length(args[1][1]), 0)
	formatters = []
	for (iarg, arg) in enumerate(args)
		data = [data arg[1]]
		if length(arg) > 1
			push!(labels, arg[2])
		else
			push!(labels, "_")
		end
		if length(arg) > 2
			# push a formatter
			push!(formatters, ft_printf(arg[3], [iarg]))
		end
	end
	if format == :md
		return pretty_table(String, data, labels;
					 backend=:text,
					 tf = tf_markdown,
				     formatters = tuple(formatters...),
					 kwargs...) |> Markdown.parse
	elseif format == :html
		return pretty_table(String, data, labels;
					 backend=:html, tf=tf_html_minimalist,
					 nosubheader=true,
    				 formatters = formatters,
					 kwargs...) |> HTML
	else
		error("unknown table format")
	end
end

function SIAMFANLEquations.nsoli(f, x0; kwargs...)
	f!(FV, x) = (FV[:] .= f(x); return nothing)
	FS = similar(x0)
	FPS = similar(x0, (length(x0), length(x0)))
	result = nsoli(f!, x0, FS, FPS; kwargs...)
	return result.solution
end

function chebbasis(x, N, a=-1, b=1)
	x = (x - a) / (b - a) - (x - b) / (a - b)
	T = zeros(N)
	T[1] = 1.0
	T[2] = x
	for n = 3:N
		T[n] = 2 * x * T[n-1] - T[n-2]
	end
	return T
end

"""
Barycentric interpolation with a Chebyshev grid with N grid points.
The interpolant is evaluated at points `x`.
"""
function chebbary(x, F::Vector)
    N = length(F)-1
	 X = [ cos(j*π/N) for j = N:-1:0 ]
    p = 0.5 * ( F[1] ./ (x .- X[1]) + (-1)^N * F[N+1] ./(x .- X[N+1]) )
    q = 0.5 * (1.0 ./ (x .- X[1]) + (-1)^N ./ (x .- X[N+1]))
    for n = 1:N-1
        p += (-1)^n * F[n+1] ./ (x .- X[n+1])
        q += (-1)^n ./ (x .- X[n+1])
    end
    return p ./ q
end

"""
generate a grid on which to plot errors; this is chosen to avoid
any grid points since barycentric interpolation is not defined
on those.
"""
errgrid(Np) = range(-1+0.000123, stop=1-0.000321, length=Np)

"""
Fast Chebyshev Transform: Nodal values -> Coefficients
"""
function fct(A::AbstractVector)
	N = length(A)
	F = ifft([A[1:N]; A[N-1:-1:2, :]])
   return [[F[1]]; 2*F[2:(N-1)]; [F(N)]]
end

"""
Inverse Fast Chebyshev Transform: Coefficients -> Nodal values
"""
function ifct(A)
	N = length(A)
	F = fft([ [A[1]]; A[2:N]/2; A[N-1:-1:2]/2 ])
	return F[1:N]
end

module IRLSQ
	using LinearAlgebra, Printf

	function wlsq(A, y, w)
		W = Diagonal(sqrt.(w))
		return qr(W * A) \ (W * y)
	end

	function irlsq(A, y; tol=1e-5, maxnit = 100, γ = 1.0, γmin = 1e-6, verbose=true)
		M, N = size(A)
		@assert M == length(y)
		wold = w = ones(M) / M
		res = 1e300
		x = zeros(N)
		verbose  && @printf("  n   | ||f-p||_inf |  extrema(w) \n")
		verbose  && @printf("------|-------------|---------------------\n")
		for nit = 1:maxnit
			x = wlsq(A, y, w)

			resnew = norm(y - A * x, Inf)
			verbose  && @printf(" %4d |   %.2e  |  %.2e  %.2e \n", nit, resnew, extrema(w)...)

			# update
			wold = w
			res = resnew
			wnew = w .* (abs.(y - A * x).^γ .+ 1e-15)
			wnew /= sum(wnew)
			w = wnew
		end
		return x, w, res
	end

	function cheb_basis(x::T, N) where {T}
		B = zeros(T, N+1)
		B[1] = one(T)
		B[2] = x
		for k = 2:N
			B[k+1] = 2 * x * B[k] - B[k-1]
		end
		return B
	end

	eval_chebpoly(F̃, x) = dot(F̃, cheb_basis(x, length(F̃)-1))

	function bestcheb(f::Function, X, N)
		y = f.(X)
		A = zeros(length(X), N)
		for (irow, x) in enumerate(X)
			A[irow, :] = cheb_basis(x, N-1)
		end
		F̃, w, res = irlsq(A, y)
		return F̃
	end
end


#=
Generate MarkovChain and DiscreteDP instances randomly.

@author : Daisuke Oyama

=#
import StatsBase: sample
import QuantEcon: MarkovChain, DiscreteDP

# random_markov_chain

"""
Return a randomly sampled MarkovChain instance with n states.

##### Arguments

- `n::Integer` : Number of states.

##### Returns

- `mc::MarkovChain` : MarkovChain instance.

##### Examples

```julia
julia> using QuantEcon

julia> mc = random_markov_chain(3)
Discrete Markov Chain
stochastic matrix:
3x3 Array{Float64,2}:
 0.281188  0.61799   0.100822
 0.144461  0.848179  0.0073594
 0.360115  0.323973  0.315912

```

"""
function random_markov_chain(n::Integer)
    p = random_stochastic_matrix(n)
    mc = MarkovChain(p)
    return mc
end

"""
Return a randomly sampled MarkovChain instance with n states, where each state
has k states with positive transition probability.

##### Arguments

- `n::Integer` : Number of states.

##### Returns

- `mc::MarkovChain` : MarkovChain instance.

##### Examples

```julia
julia> using QuantEcon

julia> mc = random_markov_chain(3, 2)
Discrete Markov Chain
stochastic matrix:
3x3 Array{Float64,2}:
 0.369124  0.0       0.630876
 0.519035  0.480965  0.0
 0.0       0.744614  0.255386

```

"""
function random_markov_chain(n::Integer, k::Integer)
    p = random_stochastic_matrix(n, k)
    mc = MarkovChain(p)
    return mc
end

# random_stochastic_matrix

"""
Return a randomly sampled n x n stochastic matrix with k nonzero entries for
each row.

##### Arguments

- `n::Integer` : Number of states.
- `k::Union{Integer, Void}(nothing)` : Number of nonzero entries in each
column of the matrix. Set to n if note specified.

##### Returns

- `p::Array` : Stochastic matrix.

"""
function random_stochastic_matrix(n::Integer, k::Union{Integer, Void}=nothing)
    if !(n > 0)
        throw(ArgumentError("n must be a positive integer"))
    end
    if k != nothing && !(k > 0 && k <= n)
        throw(ArgumentError("k must be an integer with 0 < k <= n"))
    end

    p = _random_stochastic_matrix(n, n, k=k)

    return transpose(p)
end

"""
Generate a "non-square column stochstic matrix" of shape (n, m), which contains
as columns m probability vectors of length n with k nonzero entries.

##### Arguments

- `n::Integer` : Number of states.
- `m::Integer` : Number of probability vectors.
- `;k::Union{Integer, Void}(nothing)` : Number of nonzero entries in each
column of the matrix. Set to n if note specified.

##### Returns

- `p::Array` : Array of shape (n, m) containing m probability vectors of length
n as columns.

"""
function _random_stochastic_matrix(n::Integer, m::Integer;
                                   k::Union{Integer, Void}=nothing)
    if k == nothing
        k = n
    end
    probvecs = random_probvec(k, m)

    k == n && return probvecs

    # if k < n
    # Randomly sample row indices for each column for nonzero values
    row_indices = Vector{Int}(k*m)
    for j in 1:m
        row_indices[(j-1)*k+1:j*k] = sample(1:n, k, replace=false)
    end

    p = zeros(n, m)
    for j in 1:m
        for i in 1:k
            p[row_indices[(j-1)*k+i], j] = probvecs[i, j]
        end
    end

    return p
end

# random_discrete_dp

"""
Generate a DiscreteDP randomly. The reward values are drawn from the normal
distribution with mean 0 and standard deviation `scale`.

##### Arguments

- `num_states::Integer` : Number of states.
- `num_actions::Integer` : Number of actions.
- `beta::Union{Float64, Void}(nothing)` : Discount factor. Randomly chosen from
[0, 1) if not specified.
- `;k::Union{Integer, Void}(nothing)` : Number of possible next states for each
state-action pair. Equal to `num_states` if not specified.

- `scale::Real(1)` : Standard deviation of the normal distribution for the
reward values.

##### Returns

- `ddp::DiscreteDP` : An instance of DiscreteDP.

"""
function random_discrete_dp(num_states::Integer,
                            num_actions::Integer,
                            beta::Union{Real, Void}=nothing;
                            k::Union{Integer, Void}=nothing,
                            scale::Real=1)
    L = num_states * num_actions
    R = scale * randn(L)
    Q = _random_stochastic_matrix(num_states, L; k=k)
    if beta == nothing
        beta = rand()
    end

    R = reshape(R, num_states, num_actions)
    Q = reshape(transpose(Q), num_states, num_actions, num_states)

    ddp = DiscreteDP(R, Q, beta)
    return ddp
end

# random_probvec

"""
Return m randomly sampled probability vectors of size k.

##### Arguments

- `k::Integer` : Size of each probability vector.
- `m::Integer` : Number of probability vectors.

##### Returns

- `a::Array` : Array of shape (k, m) containing probability vectors as colums.

"""
function random_probvec(k::Integer, m::Integer)
    k == 1 && return ones((k, m))

    # if k >= 2
    x = Array(Float64, (k, m))

    r = rand(k-1, m)
    x[1:end-1, :] = sort(r, 1)

    for j in 1:m
        x[end, j] = 1 - x[end-1, j]
        for i in k-1:-1:2
            x[i, j] -= x[i-1, j]
        end
    end

    return x
end