println("\n\n\nRunning tests ...")

push!(LOAD_PATH, joinpath(@__DIR__, "../.."))
using InformedDifferentialEvolution
using Test, Statistics

optimum = [-2 4]'
costf(x) = (r = sqrt(sum((x-optimum).^2)); sin(r)^2+r/2)
costf2(x, ignore) = (r = sqrt(sum((x-optimum).^2)); sin(r)^2+r/2)

best, info = DE(costf, [-10,-10], [10,10])
@test best == optimum

best, info = DE(costf2, [-10,-10], [10,10], data = 1)
@test best == optimum

f(x) = sum(abs.(x))
predictor(r,pop,costs) = r[:] = vcat(mean(pop,dims=1), mean(pop,dims=1))
mi = [-1,-1]; ma = [1,1]

best, info = DE(f, mi, ma, predictors = [predictor, :default])
@test maximum(abs.(best)) < 0.000001

function g(x, previouscost = Inf)
    r = 0.
    for i = -5:5
        r += abs.(x[1])+abs.(x[2]+i)
        if r > previouscost
            break
        end
    end
    r
end
best, info = DE(g, [-100,-100], [100,100])
@test best == [0 0]'

println("   done running tests!")


using Statistics

function show_full(array)
    show(IOContext(stdout, :limit=>false), MIME"text/plain"(), array)
end

## basic simulation functions
function xdot(x,y,z; sigma=10)
    return sigma*(y-x)
end
function ydot(x,y,z; rho=28)
    return x*(rho-z)-y
end
function zdot(x,y,z; beta=8/3)
    return x*y - beta*z
end

function make_step(x,y,z; delta=0.01)
    xp = x + xdot(x,y,z)*delta
    yp = y + ydot(x,y,z)*delta
    zp = z + zdot(x,y,z)*delta

    xn = x + 0.5*(xdot(x,y,z)+xdot(xp,yp,zp))*delta
    yn = y + 0.5*(ydot(x,y,z)+ydot(xp,yp,zp))*delta
    zn = z + 0.5*(zdot(x,y,z)+zdot(xp,yp,zp))*delta

    return xn, yn, zn
end

## threshold functions

function threshold_basic(x,y,z)
    # Z=24
    out = 0
    if z>24
        out = 1
    end
    return out
end

function z_tercile_thresh(x,y,z)
    # states should be mutually exclusive
    # Z=18.3 and Z=32
    if z<18.3 #bottom
        return 1
    end
    if z<32 #middle
        return 2
    end
    return 3 #top
end

function eight_state_thresh(x,y,z)
    state = [0,0,0]
    if x > 0
        state[1] = 1
    end
    if y > 0
        state[2] = 1
    end
    if z > 24
        state[3] = 1
    end
    map = Dict([0,0,0] => 1, [1,0,0] => 2, [0,1,0] => 3, [1,1,0] => 4,
                [0,0,1] => 5, [1,0,1] => 6, [0,1,1] => 7, [1,1,1] => 8)
    return  map[state]
end

function twelve_state_quant_thresh(x,y,z)
    #using 0.05, 0.95 quantiles for z
    state = [0,0,0]
    if x > 0
        state[1] = 1
    end
    if y > 0
        state[2] = 1
    end
    if z > 37.9
        state[3] = 2
    elseif z > 10.6
        state[3] = 1
    end
    # inefficient to be defining the map inside the function...
    map = Dict([0,0,0] => 1, [1,0,0] => 2, [0,1,0] => 3, [1,1,0] => 4,
    [0,0,1] => 5, [1,0,1] => 6, [0,1,1] => 7, [1,1,1] => 8,
    [0,0,2] => 9, [1,0,2] => 10, [0,1,2] => 11, [1,1,2] => 12)
    return map[state]
end

## simulation run
function run_sim(;runs=3000, timing=false, thresh_func=z_tercile_thresh)
    x, y, z = 0, 1, 0
    X, Y, Z = Float64[0,], Float64[1,], Float64[0,]
    cnt = 0
    state_list = [] #list of unique states
    holding_times = [] #list of holding times of each unique state

    state = thresh_func(x,y,z)

    for i in 1:runs
        x,y,z = make_step(x,y,z)
        push!(X, x)
        push!(Y, y)
        push!(Z, z)

        if thresh_func(x,y,z) == state
            cnt += 1
        else
            push!(state_list, state)
            push!(holding_times, cnt)
            state = thresh_func(x,y,z)
            cnt = 0
        end

    end
    return X, Y, Z, state_list, holding_times
end

X, Y, Z, state_list, holding_times = run_sim(;runs=1000000, timing=true, thresh_func=twelve_state_quant_thresh)
number_of_states = 12

#basic plotting
using Plots
Plots.plot(X,Y)
Plots.plot(Z[100:end])
Plots.histogram(Z;bins=100)

#construct probability transition matrix directly
function construct_p(number_of_states, state_list, holding_times)
    constr_prob = zeros(number_of_states, number_of_states)

    for i in 1:length(state_list)-1
        local current = state_list[i]
        local next = state_list[i+1]
        constr_prob[current, current] += holding_times[i]-1
        constr_prob[next, current] += 1 #reconstruct full markov chain
    end
    norm = sum(constr_prob, dims=1)

    norm_constr_prob = constr_prob ./ norm
    empirical_T = log(norm_constr_prob)/0.01
    return (norm_constr_prob, empirical_T)
end
norm_constr_prob, empirical_T = construct_p(number_of_states, state_list, holding_times)
show_full(norm_constr_prob)
show_full(empirical_T)

## other method!
function construct_t(number_of_states, state_list, holding_times)
    holding_time_dist = [[] for n in 1:number_of_states]
    constr_T = zeros(number_of_states, number_of_states)

    #count number of times unique states appear
    for i in 1:length(state_list)-1 #for i in unique states #why is the -1??
        local current = state_list[i]
        local next = state_list[i+1]
        constr_T[next, current] += 1
        push!(holding_time_dist[current], holding_times[i]*0.01)
    end

    norm = sum(constr_T, dims=1)
    norm_constr_T = constr_T ./ norm

    holding_scale = 1 ./ mean.(holding_time_dist)
    for i in 1:number_of_states
        norm_constr_T[i, i] = -1.0 #this is the bit I don't quite understand
        norm_constr_T[:, i] *= holding_scale[i]
    end

    empirical_prob = exp(norm_constr_T * 0.01)
    return (norm_constr_T, empirical_prob)
end

norm_constr_T, empirical_prob = construct_t(number_of_states, state_list, holding_times)
show_full(norm_constr_T)
show_full(empirical_prob)

`                                                           `

## identify quantiles in Z
z_sort = copy(Z)
three_state_extremes = quantile!(z_sort,[0.05, 0.95]) #output: [10.6, 37.9]
five_state_extremes = quantile!(z_sort,[0.05, 0.25, 0.75, 0.95]) #output: [10.6, 16.9, 30.6, 37.9]

###################

function get_maxima(ls; add_min=false)
    maxs = Float64[]
    if add_min
        mins = Float64[]
        if ls[1]<ls[2]
            push!(mins,ls[1])
        end
    end
    if ls[1]>ls[2]
        push!(maxs,ls[1])
    end
    for i=2:length(ls)-1
        if ls[i-1]<ls[i]>ls[i+1]
            push!(maxs,ls[i])
        end
        if add_min
            if ls[i-1]>ls[i]<ls[i+1]
                push!(mins,ls[i])
            end
        end
    end
    if ls[end]>ls[end-1]
        push!(maxs,ls[end])
    end
    if add_min
        if ls[end]<ls[end-1]
            push!(mins,ls[end])
        end
        return maxs, mins
    else
        return maxs
    end
end

# get subsequent maxima  "tent" plot
z_max = get_maxima(Z)
z_plot = Tuple{Float64, Float64}[]
for i=1:length(z_max)-1
    push!(z_plot,(z_max[i], z_max[i+1]))
end
i1 = [x[1] for x in z_plot]
i2 = [x[2] for x in z_plot]
Plots.scatter(i1,i2)

# get sequential max/min plot
z_max, z_min = get_maxima(Z, add_min=true)
Plots.scatter(z_max, z_min)


function init_local_gaussian(center, dx, dy, dz, H)
    X = Data.Array([
        2 * exp(
            -1.0 * (
                (x_g(ix, dx, H) + dx / 2 - center[1])^2 +
                (y_g(iy, dy, H) + dy / 2 - center[2])^2 +
                (z_g(iz, dz, H) + dz / 2 - center[3])^2
            ),
        ) for ix = 1:size(H, 1), iy = 1:size(H, 2), iz = 1:size(H, 3)
    ])
    return X
end

function apply_boundary_conditions!(H, coords, dims)
    if coords[1] == 1
        H[1, :, :] .= 0.0
    end
    if coords[2] == 1
        H[:, 1, :] .= 0.0
    end
    if coords[3] == 1
        H[:, :, 1] .= 0.0
    end

    if coords[1] == dims[1]
        H[end, :, :] .= 0.0
    end
    if coords[2] == dims[2]
        H[:, end, :] .= 0.0
    end
    if coords[3] == dims[3]
        H[:, :, end] .= 0.0
    end
end

function dist_norm_L2(Rh, comm_cart)
    sq_residual = [sum(Rh .^ 2)]  # this is local
    MPI.Allreduce!(sq_residual, +, comm_cart)
    return sqrt(sq_residual[1])
end

function linear_interpolate_3D(H, dx)
    dy = dz = dx

    ix, iy, iz = Int.((LOCATION_OF_INTEREST.÷ dx) .+ 1)

    # manual trilinear interpolation
    # see https://en.wikipedia.org/wiki/Trilinear_interpolation#Alternative_algorithm
    x0, x1 = ix*dx+dx/2, (ix+1)*dx+dx/2
    y0, y1 = iy*dy+dy/2, (iy+1)*dy+dy/2
    z0, z1 = iz*dz+dz/2, (iz+1)*dz+dz/2

    cvec = [H[ix, iy, iz], H[ix+1, iy, iz], H[ix, iy+1, iz], H[ix+1, iy+1, iz], H[ix, iy, iz+1], H[ix+1, iy, iz+1], H[ix, iy+1, iz+1], H[ix+1, iy+1, iz+1]]

    xvec = repeat([x0;x1], 4)
    yvec = repeat([y0; y0; y1; y1], 2)
    zvec = repeat(repeat([z0], 4), 2)

    M = hcat(ones(8), xvec, yvec, zvec, xvec.*yvec, xvec.*zvec, yvec.*zvec, xvec.*yvec.*zvec)

    interp_value = try
        avec = M \ cvec
        x, y, z = LOCATION_OF_INTEREST
        xvec_ = [1 x y z x*y x*z y*z x*z*z]
        interp_value = xvec_ * avec
        interp_value
    catch e
        interp_value = H[ix, iy, iz]
        interp_value
    end
end


using Distributions, Plots, LaTeXStrings; pyplot()

cUnif = Uniform(0,2π)
xGrid, N = 0:0.1:2π, 10^6

stephist( rand(N)*2π, bins=xGrid,
	normed=:true, c=:blue,
	label="MC Estimate")
plot!( xGrid, pdf.(cUnif,xGrid),
	c=:red,ylims=(0,0.2),label="PDF", ylabel="Density",xticks=([0:π/2:2π;],
	["0", L"\dfrac{\pi}{2}", L"\pi", L"\dfrac{3\pi}{2}", L"2\pi"]))


# This file is a part of AstroLib.jl. License is MIT "Expat".
# Copyright (C) 2016 Mosè Giordano.

function _mag2flux(mag::T, zero_point::T, ABwave::T) where {T<:AbstractFloat}
    if isnan(ABwave)
        return exp10(-0.4*(mag + zero_point))
    else
        return exp10(-0.4*(mag + 2.406 + 5*log10(ABwave)))
    end
end

"""
    mag2flux(mag[, zero_point, ABwave=number]) -> flux

### Purpose ###

Convert from magnitudes to flux expressed in erg/(s cm² Å).

### Explanation ###

This is the reverse of `flux2mag`.

### Arguments ###

* `mag`: the magnitude to be converted in flux.
* `zero_point`: the zero point level of the magnitude.  If not supplied then defaults to
 21.1 (Code et al 1976).  Ignored if the `ABwave` keyword is supplied
* `ABwave` (optional numeric keyword): wavelength, in Angstroms.  If supplied, then the
 input `mag` is assumed to contain Oke AB magnitudes (Oke & Gunn 1983, ApJ, 266, 713;
 http://adsabs.harvard.edu/abs/1983ApJ...266..713O).

### Output ###

The flux.

If the `ABwave` keyword is set, then the flux is given by the expression

\$\$\\text{flux} = 10^{-0.4(\\text{mag} +2.406 + 4\\log_{10}(\\text{ABwave}))}\$\$

Otherwise the flux is given by

\$\$\\text{flux} =  10^{-0.4(\\text{mag} + \\text{zero point})}\$\$

### Example ###

```jldoctest
julia> using AstroLib

julia> mag2flux(8.3)
1.7378008287493692e-12

julia> mag2flux(8.3, 12)
7.58577575029182e-9

julia> mag2flux(8.3, ABwave=12)
3.6244115683017193e-7
```

### Notes ###

Code of this function is based on IDL Astronomy User's Library.
"""
mag2flux(mag::Real, zero_point::Real=21.1; ABwave::Real=NaN) =
    _mag2flux(promote(float(mag), float(zero_point), float(ABwave))...,)


### A Pluto.jl notebook ###
# v0.14.4

using Markdown
using InteractiveUtils

# ╔═╡ cac0bc38-ead9-4f99-a939-fc2645874efd
md"""
# Types in Julia
"""

# ╔═╡ ecb2b8d9-ceda-4008-8137-b29bd7ef3ba9
# Uses of types.jl

# ╔═╡ b4f0668c-8b25-4588-9176-52bc7bfe797b
md"""
In this notebook we will start exploring types by implementing a type to implement algorithmic differentiation.
"""

# ╔═╡ 0ece0b55-ede7-4db1-9b48-61e1d4534761
md"""
### What is a type?
"""

# ╔═╡ 6dcc01fe-c551-43d2-930b-ff1c93d15077
md"""
A *type* can be thought of as a label that is associated with data stored in memory; this label tells Julia how to interpret the data. For example:
"""

# ╔═╡ 0d352bcc-dfc1-429f-ad6c-40dc8c4cd038
begin
	x = 1
	y = 1.3
end

# ╔═╡ 81863f06-8dac-48fd-a6e1-cf5df30c63f6
sizeof(x), sizeof(y)

# ╔═╡ 3de40205-6b5d-4d99-b886-1a1dc6a07955
typeof(x), typeof(y)

# ╔═╡ d4c280cf-910c-4bbf-8dca-b01b14711296
# rinterpret the bit-representation of x (an int) as another type, e.g. Float64
z = reinterpret(Float64, x)

# ╔═╡ 7082e444-6268-4721-a278-2511a22daf68
bitstring(x)

# ╔═╡ 43538919-4349-43f1-a196-759fe65d3be3
bitstring(z)

# ╔═╡ e0c1df0c-ff62-4982-8f3d-3ccbb44b8d2f
md"""
Both variables are stored in 8 bytes (64 bits), but one is interpreted as an integer and the other as a floating-point number.  Similarly, a pair of two numbers may be intepreted as a complex number, or an interval, or a dual number, or...; although the same information may be stored (two numbers), we want each of these different *kinds* or *types* of objects to be treated differently.
"""

# ╔═╡ e6281275-0b87-47f4-a081-bb366ced3466
md"""
## Algorithmic differentiation
"""

# ╔═╡ 6056ff4a-858f-447c-85cc-0f3f6e0e53ca
md"""
In the previous notebook we used `ForwardDiff.jl` to automatically differentiate a function. Here we will see how to implement a simple version of this.
"""

# ╔═╡ db17c624-91b6-44b1-bd10-4d72f4cb9776
md"""
The idea is to approximate a (nice enough) function $f$ near a point $a$ by a Taylor series of order 1, i.e. a straight line passing through $(a, f(a))$, with slope equal to the derivative $f'(a)$:

$f(x) \simeq f(a) + \epsilon f'(a)$,

where $\epsilon := x - a$.

We now use this to derive the standard rules for the derivative of a sum and product:

$$f(x) + g(x) \simeq [f(a) + g(a)] + \epsilon [f'(a) + g'(a)]$$

$$f(x) \cdot g(x) \simeq [f(a) \cdot g(a)] + \epsilon [f(a) g'(a) + g(a) f'(a)],$$

where we suppose that $\epsilon$ is small enough that $\epsilon^2 = 0$, or alternatively just "take the linear part".
"""

# ╔═╡ b3b15bf5-115d-4831-9115-36c88e9b899c
md"""
### Defining a composite type
"""

# ╔═╡ 70843c81-6abe-40b2-8d0b-8d779fccedc8
md"""
We see that by using just two pieces of information, namely the value $f(a)$ and the derivative $f'(a)$, we can represent a function $f$ near a given point $a$.

The pair $(f(a), f'(a))$ is often called a **dual number**. We see that is has certain **behaviours** under arithmetic operations. Whenever we have a new behaviour, a *new type is lurking*!

We group the two values into a **composite type**. We can think of a composite type as specifying the structure of a box containing several pieces of information (data) inside. Defining a composite type with two **fields** (pieces of information) has the following syntax:
"""

# ╔═╡ 82c79506-68ea-41dd-888d-840371c08430
struct MyType
    a
    b::Int
end

# ╔═╡ 291315e0-ad0e-4a57-b100-2baa5af04c04
md"""
Here we have additionally specified that the information stored in the field `b` must be of type `Int` using the **type annotation operator**, `::`.
"""

# ╔═╡ f2421fbf-5cfc-4332-be60-b61fa4e02d8c
md"""
Creating an object of that type is accomplished as follows:
"""

# ╔═╡ 710612bd-4856-4e51-b2d6-2bdc972f196d
x1 = MyType(3, 4)

# ╔═╡ 70eb5167-6288-4f84-b140-21b4b329cb45
md"""
We can extract information as follows:
"""

# ╔═╡ 9fb42f3c-c7ff-4f31-b5e6-9e96de16bac5
x1.a

# ╔═╡ 07808ac8-c910-4e92-ac34-c9ff225f4be8
typeof(x1.a)

# ╔═╡ 8777c6a7-90e3-4dc7-ab61-3d12531f2ea9
md"""
#### Exercise 1

1. Define a composite type `Dual` with fields `value` and `deriv` of type `Float64`.

2. Create two `Dual` numbers `x` and `y`.

3. What happens if you try to add `x` and `y` together?

4. Make a function `add` that adds `x` and `y` and returns a new `Dual` number, following the rules we found above.
"""

# ╔═╡ 5488c55f-8d0a-4c9b-8f8b-40ed4fe14e67
struct Dual
    value::Float64
    deriv::Float64
end

# ╔═╡ 4f5e5eea-13df-477e-baea-c37ab6074aea
x_D = Dual(3, 4)

# ╔═╡ da932afe-990f-4ca7-aa81-9d85e1ca3bac
methods(Dual)

# ╔═╡ fac29cb2-dc57-4a06-9728-567169ac3768
convert(Float64, 3)

# ╔═╡ 4118613c-e58a-4f63-81d1-0e197e66cf98
# this raises an error
convert(Float64, "hello")

# ╔═╡ ac8fa4de-d8ef-44f2-b709-aceb1e0ecb04
y_D = Dual(5, 6)

# ╔═╡ 92aa5678-478c-4cb6-9187-3285534aac8c
x_D + y_D

# ╔═╡ 04a21e57-a45c-4e4b-977a-42c6ee228b4a
newadd(x::Dual, y::Dual) = Dual(x.value + y.value, x.deriv + y.deriv)

# ╔═╡ 915bfb27-7234-43ea-a052-093f89d7bdcd
newadd(x_D, y_D)

# ╔═╡ 9b213fb9-382a-4c5b-93fe-22252d29d0c6
md"""
## Implementing arithmetic for a type
"""

# ╔═╡ 8eed2b5c-dde2-437e-8383-4ed206ce72de
md"""
We would like to be able to use `+` and `*` for our new `Dual` type, rather than typing `add(x, y)`. To do so, we need to do the following
"""

# ╔═╡ 62a97a4b-dfa6-404d-958b-b7b4ab7d25ca
# +(x::Dual, y::Dual) = Dual(x.value + y.value, x.deriv + y.deriv)

# ╔═╡ 9b0d361a-8c30-4ec0-a183-ac3890895f0e
# Base.+(x::Dual, y::Dual) = Dual(x.value + y.value, x.deriv + y.deriv)

# ╔═╡ 9b8f41b5-9ade-4f34-8d09-cb548bd6c242
# or
# +(x::Dual, y::Dual) = add(x, y)

# ╔═╡ 5b29f183-bda6-4ffb-bb78-5199980039d3
# or
# function +(x::Dual, y::Dual)
#    return add(x, y)
# end

# ╔═╡ fc624260-ba8a-4057-8a1f-2659e2343654
md"""
In Julia, `+` and `*` are just functions. They are defined in `Base` (a module containing basic function definitions) and must be `import`ed before being **extended**. They consist of many different **methods** (versions):
"""

# ╔═╡ abcbc0a1-33e7-430f-b0a8-e8b65c8666b2
@which +(3)

# ╔═╡ 17498a94-5db9-4c5e-9268-e588fddeae54
-(3)

# ╔═╡ c847f0c2-1eac-4900-b80c-7f89846b2924
methods(+)

# ╔═╡ 6caea9da-3edc-4309-bad8-f348693960b2
x

# ╔═╡ 17cce3ab-1b4f-402d-bcc3-c7905360a441
y

# ╔═╡ 526a204c-6a72-4cde-900a-568667e45c64
x + y

# ╔═╡ 86aa82b1-caae-4124-b48a-14497420f471
md"""
We can add more methods that work on our own types. (We are not allowed to modify their behaviour on combinations of types that to not contain our user-defined types; doing so is known as "type piracy" and can affect other people's code in unexpected ways.)
"""

# ╔═╡ e359dd62-65f0-4033-98c2-aff400ed0bb7
md"""
#### Exercise 2

1. Import the `+` and `*` functions from `Base` and implement them for the `Dual` type.
They should return a new `Dual` object.

2. Check that the number of methods has changed.

3. Use `@which x + y` to check that Julia knows which method to use when adding two `Dual`s.

4. Can you define `x + a` for a `Dual` number `x` and a real number `a`? What happens
"""

# ╔═╡ 86141559-f6f4-43a6-8291-550adf31b25b

# ╔═╡ cee6af83-0079-4acd-9e14-cd7b7eb6947f
begin
	value(f::Dual) = f.value
	deriv(f::Dual) = f.deriv
end

# ╔═╡ 4e222c6a-054b-41e6-bef4-b0726344ee94
let
	∂(f::Dual) = f.deriv
	*(f::Dual, g::Dual) = Dual(value(f) * value(g),
    	                     value(f) * ∂(g) + ∂(f) * value(g))
	42 # return this to avoid conflict
end

# ╔═╡ a1b1b945-ba2e-4668-8698-4df02660d379
md"""
(Try using `@code_llvm`)
"""

# ╔═╡ 8857e0aa-da38-4ae6-8ef9-efa66dd51cfc
begin
	f = Dual(3.0, 4.0)
	g = Dual(5.0, 6.0)
end

# ╔═╡ ad671591-a625-4de2-9b18-84c7773cb4b8
f, g

# ╔═╡ 053e90a7-5fa6-4063-9634-63d48ee16690
f + g

# ╔═╡ d521f0e3-1636-4e8b-b8f3-639a90780690
f * g

# ╔═╡ b7b81d9d-354d-4ec0-b50c-302dbcb8f0c4
md"""
Amazingly, we now have enough to be able to differentiate simple Julia functions involving only `+` and `*`. Define
"""

# ╔═╡ d67fa8a9-f43e-4eca-bb0e-1603c9f7da5d
begin
	a = 3.0
	xx = Dual(a, 1.0)  # "the identity function x ↦ x, with derivative 1"
end

# ╔═╡ 3bc3d5a3-9b9f-4d07-9d59-f24d5c005038
md"""
We initialize the derivative as 1.0 when we make a `Dual`. If we use `x` then we automatically differentiate!
"""

# ╔═╡ fd34be78-af76-4d46-bdbd-19fd19155785
md"""
#### Exercise 3
"""

# ╔═╡ caccd695-118c-46a8-b59f-376e4cbc83ce
md"""
1. Define `a = 3.0` and `xx = Dual(a, 1.0)`.

    (i) Compute `xx + xx`. The result should have the value $2a$ and the derivative $2$ -- write a test that it does so.

    (ii) Do the same for `xx * xx`.

2. Define the function `f(x) = x * x + x`. Compute `f(xx)` and check that it gives the correct value and derivative!

3. Does this work for the function `f(x) = x^2 + x`?  What do you need to do?

4. What happens for `f(x) = x^2 + 2x`? What do you need to do?

5. What should you do for `f(x) = sin(x) + x`?
"""

# ╔═╡ 6418d0a4-8dd3-4816-8e5b-092cb769ebb6
a = 3.0
xx = Dual(a, 1.0)

# ╔═╡ 8ebb42e9-11ec-434d-87c9-4259c26f6526
xx + xx

# ╔═╡ 84b6a870-4515-4e6c-b306-d6b6d2a0fb8d
xx * xx

# ╔═╡ a0865cb0-7077-44ff-a130-eb73cbbc3179
xx * xx + xx  # the function x ↦ x^2 + x,  derivative 2x + 1

# ╔═╡ f1fc7637-4804-4fc1-81ac-14874e66f9d6
ff(x) = x*x*x + x*x + x*x  # x^3 + 2x^2

# ╔═╡ ded7adc0-619f-4735-bd93-9c26c4c21858
ff(xx)

# ╔═╡ 1284374b-ab26-46b8-a74e-119655f5112e
ff(a)

# ╔═╡ 4e06b4f6-9753-4847-8c6c-120208520dc5
n = 4
fff(x) = x^n

# ╔═╡ 420c0ef0-3f7e-4044-be30-c2873636c6e8
dff(x) = 3x^2 + 4x

# ╔═╡ f426bed2-05f1-4598-9817-bb90262911ef
dff(a)

# ╔═╡ 09bfa122-ca5b-4b4a-a9eb-09b1b5482a76
fff(x) = x^2

# ╔═╡ d08a434c-c816-4f43-9c93-2d0eda57723c
fff(xx)

# ╔═╡ 131373df-8e58-49b0-be5f-3cd8d8deb61d
fff(xx)

# ╔═╡ 76a7a87b-0283-4a85-9d88-0eec06ef689f
fff4(x) = x^4

# ╔═╡ d608bc38-5653-4ad5-b6e4-6b2236299e4c
fff4(xx)

# ╔═╡ bf697142-0eaf-4717-9434-eaa89946beb4
md"""
`Base.literal_pow`
"""

# ╔═╡ 108212df-c86f-4b75-8576-65f3964687e9
md"""
#### Exercise 4
"""

# ╔═╡ eb94cc27-ac9e-4202-acdc-248809eebdb9
md"""
1. Define a function `differentiate` that differentiates a function `f` at a point `a` using `Dual` numbers, by following the above pattern. (It should return just the derivative at the given point.)
"""

# ╔═╡ d7f6c80a-4922-42be-9bce-e493e603c373

# ╔═╡ 70ec3ada-f7ae-4526-be35-5aa8cd4533ca
md"""
This is the basis of ("forward-mode") automatic differentiation. The `ForwardDiff.jl` method contains a sophisticated implementation of this method.
"""

# ╔═╡ 5011cf98-128e-4c94-bae0-109016427c50
md"""
### Parametric types
"""

# ╔═╡ d57937d8-3f1e-4f85-bf6b-afc8cfe3e99d
md"""
For simplicity, in the above we fixed the fields in the `Dual` type to be of type `Float64`. By doing so we are actually *losing power*. Instead we should let Julia "fill in" the types.

To do so, we specify that we want to use a **type parameter** `T`. We can think of this as a "special kind of variable" that can only take on certain kinds of values. We specify this with the following syntax:
"""

# ╔═╡ df1c033c-8a99-4b45-bc9c-3287a3a959ec
struct MyType2{T}
    a::T
    b::T
end

# ╔═╡ b0902a74-01e0-456a-a027-a5030d494d7f
md"""
[Note that we have not reused the name `MyType` since Julia *does not allow types to be redefined in a different way*.]
"""

# ╔═╡ 9226e5a3-4445-4420-9b54-be31468e6b07
md"""
Here we are specifying that both fields `a` and `b` must share the same type `T`, but we have not restricted what values `T` can take. When we create an object, Julia will *infer* (work out) the type:
"""

# ╔═╡ ccc258fa-c928-4175-a48e-bd72b7d22f31
x_M = MyType2(3, 4)

# ╔═╡ c9f3dde8-7989-45e8-9831-894d82e57e90
y_M = MyType2(3.1, 4.2)

# ╔═╡ 2d3876d6-5fcf-4bfb-8dba-4bdf252c73fc
z_M = MyType2(1, 5.3)

# ╔═╡ a0d7a33a-e067-4e69-9c70-14005f41a2c6
struct MyType3{S,T}
    a::S
    b::T
end

# ╔═╡ 812149d8-0040-4d67-b2a2-f1304af4b225
x_M2 = MyType3(1, 5.3)

# ╔═╡ 75fa60fe-6583-4a6a-b6b0-06fb81006218
md"""
Note that `x` and `y` have *different* types.
"""

# ╔═╡ 4146cf8f-ab2f-457d-97ef-e6d39e93246e
md"""
We can define functions acting on parametric types without necessarily talking about the type parameter:
"""

# ╔═╡ 3e61a354-ca21-410b-b422-92ea362b8c10
md"""
#### Exercise 5

1. Define a function that takes an object of type `MyType2`, *without* mentioning the type parameter, and returns the sum of the two fields.

   What happens when you apply this function to `x` and `y`?

2. Define a type `Dual2` with a type parameter `T` and the same functions `+` and `*` as before.

3. Define the function `f(x) = x * x + x`. What happens if you pass in `Dual` numbers with different type parameters?
"""

# ╔═╡ 907f65aa-e1a1-4838-a1a0-cdaa0a1cfa45
struct DualReal
    value::Real
    deriv::Real
end

# ╔═╡ fc5839ce-935a-46aa-bbe7-01f203b0fdb5
DualReal(big(1.0), big"1.4")

# ╔═╡ c445bab5-c012-4890-89bc-e7d6950fa7b4
DualReal(3.1, big"1.4")

# ╔═╡ 9008cc1d-e7c2-4b24-ad08-968c085baa36
struct DualUntyped
    value
    deriv
end

# ╔═╡ d5cb6844-6f9d-422e-a241-96450aff70f5
DualUntyped("hello", "David")

# ╔═╡ 7b2f03b4-d3dd-4fa6-a569-c878fc9f6557
struct DualBig
    value::BigFloat
    deriv::BigFloat
end

# ╔═╡ 97097783-2948-4e9c-b769-8085c7943227
begin
	struct Dual3{T<:Real}
	    a::T
    	b::T
	end

	Dual3(a::Real, b::Real) = Dual3(promote(a, b))

	# Dual3(a::Real, b::Real) = Dual3(promote(a, b)...)
end

# ╔═╡ 4bdb5bfc-82c2-4de5-982a-4a69872749d1
Dual3("a", "b")

# ╔═╡ 9c05cca8-76d9-4abd-8774-28e7f41d3b09
Dual3(big"1.0", big"1.4")

# ╔═╡ 05db96f7-3eeb-47aa-a4dc-dd67b0f02a2b
Dual3(1, 3.1)

# ╔═╡ b490d647-85f5-4b82-ab29-2d18ce1c7db4
methods(Dual3)

# ╔═╡ d981b89c-4833-4e38-9f80-246ed5f09445
Dual3(3, 4.5)

# ╔═╡ 9bfe7470-afec-4a24-9f96-43272e2f9ecd
promote(1, 3.4)

# ╔═╡ 61cc131f-2d02-4de1-b79a-9105d6decc15
@which promote_type(Int, Float64)

# ╔═╡ a5c1c463-2393-499d-852c-16e69549819a
Complex(3, 4.5)

# ╔═╡ 749ba068-6deb-4a7d-a650-412e5fbdb670
@which Complex(3, 4.5)

# ╔═╡ 226f5af6-702f-431c-a269-5dd4cafa2699
v = Dual(3, 4)

# ╔═╡ c3831c37-8e09-4846-963d-27e1d4a824ac
v isa Number

# ╔═╡ fd8bf8f8-68dd-4925-aa9b-8eca1937c577
subtypes(Real)

# ╔═╡ e3d64165-170e-4ee0-a9bc-ef759ce6afd7
subtypes(AbstractFloat)

# ╔═╡ 713dc49e-6038-40b7-a917-299277f7ea7a
subtypes(Float64)

# ╔═╡ 68d54b7c-4003-49e8-b2f9-666b882a5737
isconcretetype(Float64)

# ╔═╡ 7ae57ebb-d6bc-4b9b-913a-f9c4bb11d380
q = 3.1

# ╔═╡ 99dd8229-8658-4efe-9133-3eb8cd00b6b3
typeof(q)

# ╔═╡ ef28ca54-38dd-4e8e-a5e8-897c802f81e1
q isa AbstractFloat

# ╔═╡ a31cfb26-9914-441c-aea8-6f341f2635fb
gg(x::AbstractFloat) = sqrt(x)

# ╔═╡ b05dcbc4-0149-48a3-8533-1bd6515d0b32
gg(x::T) where {T <: AbstractFloat} = sqrt(x)

# ╔═╡ f44efbd4-7ce2-44db-ac0e-eeb4aeafef70
gg(π)

# ╔═╡ 7df9a91e-01d4-471f-9ad5-55ac56c245d3
struct Hello{T}
    a::T
    b::T
    c::T
end

# ╔═╡ ef1c3d05-7020-4339-b431-63ad63b244cc
add(x::Hello, y::Hello) = x.a + y.a, x.b + y.b, ...

# ╔═╡ 186148a4-698e-460b-82dd-7b203698de45
w = Hello(3, 4, 45)

# ╔═╡ b9ac03c2-d93e-4372-a3cb-4568baf015bb
myfields = fieldnames(typeof(w))

# ╔═╡ 522a7a5c-8400-4e1b-9ebd-6d76b33cb87f
getfield(w, :a)

# ╔═╡ d70384a3-51fc-4f71-a3d6-d14b03719495
x

# ╔═╡ ab382c7f-933e-4611-af23-fe31967ff8cc
[getfield(x, name) for name in myfields]

# ╔═╡ 8df1bf08-5ee2-4cfe-a273-62c71dc9b8fb
abstract type AbstractDual{S<:Real,T<:Real}

struct Dual5{S<:Real, T<:Real} <: AbstractDual{S,T}
    a::S
    b::T
end

# ╔═╡ babed906-7888-431b-8b11-d6e5889073fa
getfield.(myfields)

# ╔═╡ d7e3e2b2-6f3f-4761-9141-924ae1448343

# ╔═╡ ce4bd2c6-6269-47e3-8f43-50969a19487a
x::Dual ^ n::Integer = ..

# ╔═╡ 555aa973-a5ce-41f9-aab7-261a697758cc
*(f::Dual, g::Dual) = Dual(f.value * g.value,
                         f.value * g.deriv + f.deriv * g.value)

# Alternatively and equivalently
# *(f::Dual, g::Dual) = Dual(value(f) * value(g),
#                          value(f) * deriv(g) + deriv(f) * value(g))

# ╔═╡ eea59dfb-acc7-4d56-82b2-59f0ed40afc7
import Base: +, *

# ╔═╡ b0c794fb-0d77-4667-9d47-0bcdfaa8fecb
begin
	import Base: ^
	^(x::Dual, n::Integer) = Base.power_by_squaring(x, n)
end

# ╔═╡ Cell order:
# ╟─cac0bc38-ead9-4f99-a939-fc2645874efd
# ╠═ecb2b8d9-ceda-4008-8137-b29bd7ef3ba9
# ╟─b4f0668c-8b25-4588-9176-52bc7bfe797b
# ╟─0ece0b55-ede7-4db1-9b48-61e1d4534761
# ╟─6dcc01fe-c551-43d2-930b-ff1c93d15077
# ╠═0d352bcc-dfc1-429f-ad6c-40dc8c4cd038
# ╠═81863f06-8dac-48fd-a6e1-cf5df30c63f6
# ╠═3de40205-6b5d-4d99-b886-1a1dc6a07955
# ╠═d4c280cf-910c-4bbf-8dca-b01b14711296
# ╠═7082e444-6268-4721-a278-2511a22daf68
# ╠═43538919-4349-43f1-a196-759fe65d3be3
# ╟─e0c1df0c-ff62-4982-8f3d-3ccbb44b8d2f
# ╟─e6281275-0b87-47f4-a081-bb366ced3466
# ╟─6056ff4a-858f-447c-85cc-0f3f6e0e53ca
# ╟─db17c624-91b6-44b1-bd10-4d72f4cb9776
# ╟─b3b15bf5-115d-4831-9115-36c88e9b899c
# ╟─70843c81-6abe-40b2-8d0b-8d779fccedc8
# ╠═82c79506-68ea-41dd-888d-840371c08430
# ╟─291315e0-ad0e-4a57-b100-2baa5af04c04
# ╟─f2421fbf-5cfc-4332-be60-b61fa4e02d8c
# ╠═710612bd-4856-4e51-b2d6-2bdc972f196d
# ╟─70eb5167-6288-4f84-b140-21b4b329cb45
# ╠═9fb42f3c-c7ff-4f31-b5e6-9e96de16bac5
# ╠═07808ac8-c910-4e92-ac34-c9ff225f4be8
# ╟─8777c6a7-90e3-4dc7-ab61-3d12531f2ea9
# ╠═5488c55f-8d0a-4c9b-8f8b-40ed4fe14e67
# ╠═4f5e5eea-13df-477e-baea-c37ab6074aea
# ╠═da932afe-990f-4ca7-aa81-9d85e1ca3bac
# ╠═fac29cb2-dc57-4a06-9728-567169ac3768
# ╠═4118613c-e58a-4f63-81d1-0e197e66cf98
# ╠═ac8fa4de-d8ef-44f2-b709-aceb1e0ecb04
# ╠═92aa5678-478c-4cb6-9187-3285534aac8c
# ╠═04a21e57-a45c-4e4b-977a-42c6ee228b4a
# ╠═915bfb27-7234-43ea-a052-093f89d7bdcd
# ╟─9b213fb9-382a-4c5b-93fe-22252d29d0c6
# ╟─8eed2b5c-dde2-437e-8383-4ed206ce72de
# ╠═eea59dfb-acc7-4d56-82b2-59f0ed40afc7
# ╠═62a97a4b-dfa6-404d-958b-b7b4ab7d25ca
# ╠═9b0d361a-8c30-4ec0-a183-ac3890895f0e
# ╠═9b8f41b5-9ade-4f34-8d09-cb548bd6c242
# ╠═5b29f183-bda6-4ffb-bb78-5199980039d3
# ╟─fc624260-ba8a-4057-8a1f-2659e2343654
# ╠═abcbc0a1-33e7-430f-b0a8-e8b65c8666b2
# ╠═17498a94-5db9-4c5e-9268-e588fddeae54
# ╠═c847f0c2-1eac-4900-b80c-7f89846b2924
# ╠═6caea9da-3edc-4309-bad8-f348693960b2
# ╠═17cce3ab-1b4f-402d-bcc3-c7905360a441
# ╠═526a204c-6a72-4cde-900a-568667e45c64
# ╟─86aa82b1-caae-4124-b48a-14497420f471
# ╟─e359dd62-65f0-4033-98c2-aff400ed0bb7
# ╠═86141559-f6f4-43a6-8291-550adf31b25b
# ╠═cee6af83-0079-4acd-9e14-cd7b7eb6947f
# ╠═555aa973-a5ce-41f9-aab7-261a697758cc
# ╠═4e222c6a-054b-41e6-bef4-b0726344ee94
# ╟─a1b1b945-ba2e-4668-8698-4df02660d379
# ╠═8857e0aa-da38-4ae6-8ef9-efa66dd51cfc
# ╠═ad671591-a625-4de2-9b18-84c7773cb4b8
# ╠═053e90a7-5fa6-4063-9634-63d48ee16690
# ╠═d521f0e3-1636-4e8b-b8f3-639a90780690
# ╟─b7b81d9d-354d-4ec0-b50c-302dbcb8f0c4
# ╠═d67fa8a9-f43e-4eca-bb0e-1603c9f7da5d
# ╟─3bc3d5a3-9b9f-4d07-9d59-f24d5c005038
# ╟─fd34be78-af76-4d46-bdbd-19fd19155785
# ╟─caccd695-118c-46a8-b59f-376e4cbc83ce
# ╠═6418d0a4-8dd3-4816-8e5b-092cb769ebb6
# ╠═8ebb42e9-11ec-434d-87c9-4259c26f6526
# ╠═84b6a870-4515-4e6c-b306-d6b6d2a0fb8d
# ╠═a0865cb0-7077-44ff-a130-eb73cbbc3179
# ╠═f1fc7637-4804-4fc1-81ac-14874e66f9d6
# ╠═ded7adc0-619f-4735-bd93-9c26c4c21858
# ╠═1284374b-ab26-46b8-a74e-119655f5112e
# ╠═420c0ef0-3f7e-4044-be30-c2873636c6e8
# ╠═f426bed2-05f1-4598-9817-bb90262911ef
# ╠═09bfa122-ca5b-4b4a-a9eb-09b1b5482a76
# ╠═d08a434c-c816-4f43-9c93-2d0eda57723c
# ╠═b0c794fb-0d77-4667-9d47-0bcdfaa8fecb
# ╠═4e06b4f6-9753-4847-8c6c-120208520dc5
# ╠═131373df-8e58-49b0-be5f-3cd8d8deb61d
# ╠═76a7a87b-0283-4a85-9d88-0eec06ef689f
# ╠═d608bc38-5653-4ad5-b6e4-6b2236299e4c
# ╠═ce4bd2c6-6269-47e3-8f43-50969a19487a
# ╟─bf697142-0eaf-4717-9434-eaa89946beb4
# ╟─108212df-c86f-4b75-8576-65f3964687e9
# ╟─eb94cc27-ac9e-4202-acdc-248809eebdb9
# ╠═d7f6c80a-4922-42be-9bce-e493e603c373
# ╟─70ec3ada-f7ae-4526-be35-5aa8cd4533ca
# ╟─5011cf98-128e-4c94-bae0-109016427c50
# ╟─d57937d8-3f1e-4f85-bf6b-afc8cfe3e99d
# ╠═df1c033c-8a99-4b45-bc9c-3287a3a959ec
# ╟─b0902a74-01e0-456a-a027-a5030d494d7f
# ╟─9226e5a3-4445-4420-9b54-be31468e6b07
# ╠═ccc258fa-c928-4175-a48e-bd72b7d22f31
# ╠═c9f3dde8-7989-45e8-9831-894d82e57e90
# ╠═2d3876d6-5fcf-4bfb-8dba-4bdf252c73fc
# ╠═a0d7a33a-e067-4e69-9c70-14005f41a2c6
# ╠═812149d8-0040-4d67-b2a2-f1304af4b225
# ╟─75fa60fe-6583-4a6a-b6b0-06fb81006218
# ╟─4146cf8f-ab2f-457d-97ef-e6d39e93246e
# ╟─3e61a354-ca21-410b-b422-92ea362b8c10
# ╠═907f65aa-e1a1-4838-a1a0-cdaa0a1cfa45
# ╠═fc5839ce-935a-46aa-bbe7-01f203b0fdb5
# ╠═c445bab5-c012-4890-89bc-e7d6950fa7b4
# ╠═9008cc1d-e7c2-4b24-ad08-968c085baa36
# ╠═d5cb6844-6f9d-422e-a241-96450aff70f5
# ╠═7b2f03b4-d3dd-4fa6-a569-c878fc9f6557
# ╠═97097783-2948-4e9c-b769-8085c7943227
# ╠═4bdb5bfc-82c2-4de5-982a-4a69872749d1
# ╠═9c05cca8-76d9-4abd-8774-28e7f41d3b09
# ╠═05db96f7-3eeb-47aa-a4dc-dd67b0f02a2b
# ╠═b490d647-85f5-4b82-ab29-2d18ce1c7db4
# ╠═d981b89c-4833-4e38-9f80-246ed5f09445
# ╠═9bfe7470-afec-4a24-9f96-43272e2f9ecd
# ╠═61cc131f-2d02-4de1-b79a-9105d6decc15
# ╠═a5c1c463-2393-499d-852c-16e69549819a
# ╠═749ba068-6deb-4a7d-a650-412e5fbdb670
# ╠═226f5af6-702f-431c-a269-5dd4cafa2699
# ╠═c3831c37-8e09-4846-963d-27e1d4a824ac
# ╠═fd8bf8f8-68dd-4925-aa9b-8eca1937c577
# ╠═e3d64165-170e-4ee0-a9bc-ef759ce6afd7
# ╠═713dc49e-6038-40b7-a917-299277f7ea7a
# ╠═68d54b7c-4003-49e8-b2f9-666b882a5737
# ╠═7ae57ebb-d6bc-4b9b-913a-f9c4bb11d380
# ╠═99dd8229-8658-4efe-9133-3eb8cd00b6b3
# ╠═ef28ca54-38dd-4e8e-a5e8-897c802f81e1
# ╠═a31cfb26-9914-441c-aea8-6f341f2635fb
# ╠═f44efbd4-7ce2-44db-ac0e-eeb4aeafef70
# ╠═b05dcbc4-0149-48a3-8533-1bd6515d0b32
# ╠═7df9a91e-01d4-471f-9ad5-55ac56c245d3
# ╠═ef1c3d05-7020-4339-b431-63ad63b244cc
# ╠═186148a4-698e-460b-82dd-7b203698de45
# ╠═b9ac03c2-d93e-4372-a3cb-4568baf015bb
# ╠═522a7a5c-8400-4e1b-9ebd-6d76b33cb87f
# ╠═d70384a3-51fc-4f71-a3d6-d14b03719495
# ╠═ab382c7f-933e-4611-af23-fe31967ff8cc
# ╠═8df1bf08-5ee2-4cfe-a273-62c71dc9b8fb
# ╠═babed906-7888-431b-8b11-d6e5889073fa
# ╠═d7e3e2b2-6f3f-4761-9141-924ae1448343


using Plots, StatsBase, Statistics, LaTeXStrings, JLD

function InitialNetwork(dim, P)
    Network = sample([-1,0], Weights([1-P, P]),(dim,dim))
    return Network
end

function FindNeighbors(i, j, Network, dim)
    Numbereds = 0
    colorList = []
    for neighbor in ([0,-1], [-1,0])
        if j + neighbor[2] >= 1 && i + neighbor[1] >= 1 && Network[([i, j] + neighbor)...] > 0
            push!(colorList,Network[([i, j] + neighbor)...])
            Numbereds += 1
        end
    end
    return colorList, Numbereds #colorList contain the node and its neighbors
end

function RecursionMapping(x, L)
    while L[x] != L[L[x]]
        x = L[x]
    end
    return L[x]
end

function CheckPercolation(dim, L, Network)
    fside = []
    lside = []
    for i in 1:dim
        if Network[i,dim] != -1
            push!(lside, RecursionMapping(Network[i,dim],L))
        end
        if Network[i,1] != -1
            push!(fside, RecursionMapping(Network[i,1],L))
        end
    end
    RandomPoint = rand(Network)
    if RandomPoint != -1 && RecursionMapping(RandomPoint,L) in intersect(fside,lside)
        return 1
    else
        return 0
    end
end

function HKNetworkDynamic(dim, P)
    Network = InitialNetwork(dim, P)
    S = []
    L = []
    k = 1

    for j in 1:dim
        for i in 1:dim
            if Network[i,j] != 0
                continue
            end
            NColor, NNum = FindNeighbors(i,j,Network,dim)
            if NNum == 0
                Network[i,j] = k
                push!(L, k)
                push!(S, 1)
                k += 1
            elseif NNum == 1
                Network[i,j] = RecursionMapping(NColor[1],L)
                S[RecursionMapping(NColor[1],L)] += 1
            elseif NNum == 2 && RecursionMapping(NColor[1],L) == RecursionMapping(NColor[2],L)
                Network[i,j] = RecursionMapping(NColor[1],L)
                S[RecursionMapping(NColor[1],L)] += 1
            else
                S[RecursionMapping(NColor[1],L)] += 1 + S[RecursionMapping(NColor[2],L)]
                S[RecursionMapping(NColor[2],L)] = 0
                Network[i,j] = RecursionMapping(NColor[1],L)
                L[NColor[2]] = RecursionMapping(NColor[1],L)
            end
        end
    end
    # return Network
    return CheckPercolation(dim, L, Network)
end

HKNetworkDynamic(200, 0.6)

Plist = hcat(0.5:0.01:1)
Avglist = []
STDlist = []
TotData = []

for P in Plist
    print("\r$P")
    AllRuns = []
    for i in 1:10000
        push!(AllRuns, HKNetworkDynamic(200, P))
    end
    push!(TotData, AllRuns)
    push!(STDlist, std(AllRuns))
    push!(Avglist, mean(AllRuns))
end

save("200x200-HK.jld", "data", TotData)

scatter(Plist,
    Avglist,
    yerr = STDlist,
    legend = nothing,
    xlabel = L"P",
    ylabel = L"Avg\ Q_{\infty}",
    title = L"Avg \ Q_{\infty} \ for \ 200\times200 \ Network\ (10000 \ runs)")
savefig("../../Figs/Q6/200x200-10000.pdf")


module KSVD

# This is an implementation of the K-SVD algorithm.
# The original paper:
# K-SVD: An Algorithm for Designing Overcomplete Dictionaries
# for Sparse Representation
# http://www.cs.technion.ac.il/~freddy/papers/120.pdf

# Variable names are based on the original paper.
# If you try to read the code, I recommend you to see Figure 2 first.
#

export ksvd, matching_pursuit

using ProgressMeter
using Base.Threads, Random, SparseArrays, LinearAlgebra

include("matching_pursuit.jl")

const default_sparsity_allowance = 0.9
const default_max_iter = 200
const default_max_iter_mp = 200

Random.seed!(1234)  # for stability of tests

function error_matrix(Y::AbstractMatrix, D::AbstractMatrix, X::AbstractMatrix, k::Int)
    # indices = [i for i in 1:size(D, 2) if i != k]
    indices = deleteat!(collect(1:size(D, 2)), k)
    return Y - D[:, indices] * X[indices, :]
end

function init_dictionary(n::Int, K::Int)
    # D must be a full-rank matrix
    D = rand(n, K)
    while rank(D) != min(n, K)
        D = rand(n, K)
    end

    @inbounds for k in 1:K
        D[:, k] ./= norm(@view(D[:, k]))
    end
    return D
end

function ksvd(Y::AbstractMatrix, D::AbstractMatrix, X::AbstractMatrix)
    N = size(Y, 2)
    for k in 1:size(X, 1)
        xₖ = X[k, :]
        # ignore if the k-th row is zeros
        all(iszero, xₖ) && continue

        # wₖ is the column indices where the k-th row of xₖ is non-zero,
        # which is equivalent to [i for i in N if xₖ[i] != 0]
        wₖ = findall(!iszero, xₖ)

        # Eₖ * Ωₖ implies a selection of error columns that
        # correspond to examples that use the atom D[:, k]
        Eₖ = error_matrix(Y, D, X, k)
        Ωₖ = sparse(wₖ, 1:length(wₖ), ones(length(wₖ)), N, length(wₖ))
        # Note that S is a vector that contains diagonal elements of
        # a matrix Δ such that Eₖ * Ωₖ == U * Δ * V.
        # Non-zero entries of X are set to
        # the first column of V multiplied by Δ(1, 1)
        U, S, V = svd(Eₖ * Ωₖ, full=true)
        D[:, k] = U[:, 1]
        X[k, wₖ] = V[:, 1] * S[1]
    end
    return D, X
end

"""
    ksvd(Y::AbstractMatrix, n_atoms::Int;
         sparsity_allowance::Float64 = $default_sparsity_allowance,
         max_iter::Int = $default_max_iter,
         max_iter_mp::Int = $default_max_iter_mp)

Run K-SVD that designs an efficient dictionary D for sparse representations,
and returns X such that DX = Y or DX ≈ Y.

```
# Arguments
* `sparsity_allowance`: Stop iteration if the number of zeros in X / the number
    of elements in X > sparsity_allowance.
* `max_iter`: Limit of iterations.
* `max_iter_mp`: Limit of iterations in Matching Pursuit that `ksvd` calls at
    every iteration.
```
"""
function ksvd(Y::AbstractMatrix, n_atoms::Int;
              sparsity_allowance = default_sparsity_allowance,
              max_iter::Int = default_max_iter,
              max_iter_mp::Int = default_max_iter_mp)

    K = n_atoms
    n, N = size(Y)

    if !(0 <= sparsity_allowance <= 1)
        throw(ArgumentError("`sparsity_allowance` must be in range [0,1]"))
    end

    X = spzeros(K, N)  # just for making X global in this function
    max_n_zeros = ceil(Int, sparsity_allowance * length(X))

    # D is a dictionary matrix that contains atoms for columns.
    D = init_dictionary(n, K)  # size(D) == (n, K)

    p = Progress(max_iter)

    for i in 1:max_iter
        X_sparse = matching_pursuit(Y, D, max_iter = max_iter_mp)
        D, X = ksvd(Y, D, Matrix(X_sparse))

        # return if the number of zero entries are <= max_n_zeros
        if sum(iszero, X) > max_n_zeros
            return D, X
        end
        next!(p)
    end
    return D, X
end

end # module


"""
    eom_rossler_lorenz() -> Function

Equations of motions for a 6D unidirectionally coupled Rössler-Lorenz system.
The 3D Rössler attractor is coupled to  a 3D Lorenz attractor such that the
second component  (``x_2``) of the Rössler system unidirectionally forces
the second component (``y_2``) of the Lorenz system. This example is from [1].

The implementation here allows for tuning the parameters of each subsystem by
introducing the constants `a₁`, `a₂`, `a₃`, `b₁`, `b₂`, `b₃`. The default
values of these constants are as defined in Krakovská et al. (2018).

The dynamics is generated by the vector field:

```math
\\begin{aligned}
\\dot x_1 &= a_1(x_2 + x_3) \\
\\dot x_2 &= a_2(x_1 + 0.2x_2) \\
\\dot x_3 &= a_2(0.2 + x_3(x_1 - a_3)) \\
\\dot y_1 &= b_1(y_2 - y_1) \\
\\dot y_2 &= y_1(b_2 - y_3) - y_2 +c(x_2)^2 \\
\\dot y_3 &= y_1y_2 - b_3y_3
\\end{aligned}
```

with the coupling constant ``c \\geq 0``.

# References
[Krakovská, Anna, et al. "Comparison of six methods for the detection of causality in a bivariate time series." Physical Review E 97.4 (2018):042207](https://journals.aps.org/pre/abstract/10.1103/PhysRevE.97.042207)
"""
function eom_rossler_lorenz(u, p, t)
    a₁, a₂, a₃, b₁, b₂, b₃, c = (p...,)
    x₁, x₂, x₃, y₁, y₂, y₃ = (u...,)

    dx₁ = a₁*(x₂ + x₃)
    dx₂ = a₂*(x₁ + 0.2*x₂)
    dx₃ = a₂*(0.2 + x₃*(x₁ - a₃))
    dy₁ = b₁*(y₂ - y₁)
    dy₂ = y₁*(b₂ - y₃) - y₂ + c*x₂^2
    dy₃ = y₁*y₂ - b₃*y₃
    return SVector{6}(dx₁, dx₂, dx₃, dy₁, dy₂, dy₃)
end

function rossler_lorenz(u₀, a₁, a₂, a₃, b₁, b₂, b₃, c)
    p = [a₁, a₂, a₃, b₁, b₂, b₃, c]
    ContinuousDynamicalSystem(eom_rossler_lorenz, u₀, p)
end

"""
    rossler_lorenz(;u₀ = rand(6), a₁ = -6, a₂ = 6, a₃ = 2.0,
        b₁ = 10, b₂ = 28, b₃ = 8/3, c = 2) -> ContinuousDynamicalSystem

Initialise a Rössler-Lorenz system consisting of two independent 3D subsystems:
one Rössler attractor and one Lorenz attractor. They are coupled such that the
second component (`x₂`) of the Rössler system unidirectionally forces the
second component (`y₂`) of the Lorenz system.

The implementation here allows for tuning the parameters of each subsystem by
introducing the constants `a₁`, `a₂`, `a₃`, `b₁`, `b₂`, `b₃`. The default
values are as in Krakovská et al. (2018).

The dynamics is generated by the vector field:

```math
\\begin{aligned}
\\dot x_1 &= a_1(x_2 + x_3) \\
\\dot x_2 &= a_2(x_1 + 0.2x_2) \\
\\dot x_3 &= a_2(0.2 + x_3(x_1 - a_3)) \\
\\dot y_1 &= b_1(y_2 - y_1) \\
\\dot y_2 &= y_1(b_2 - y_3) - y_2 +c(x_2)^2 \\
\\dot y_3 &= y_1 y_2 - b_3y_3
\\end{aligned}
```

with the coupling constant ``c \\geq 0``.

# References
[Krakovská, Anna, et al. "Comparison of six methods for the detection of causality in a bivariate time series." Physical Review E 97.4 (2018):042207](https://journals.aps.org/pre/abstract/10.1103/PhysRevE.97.042207)
"""
rossler_lorenz(;u₀ = rand(6), a₁ = -6, a₂ = 6, a₃ = 2.0, b₁ = 10, b₂ = 28, b₃ = 8/3, c = 2) =
    rossler_lorenz(u₀, a₁, a₂, a₃, b₁, b₂, b₃, c)


include("kepler_eqn.jl")

function calc_true_anom(ecc_anom::Real, e::Real)
	true_anom = 2*atan(sqrt((1+e)/(1-e))*tan(ecc_anom/2))
end

begin
	""" Calculate RV from t, P, K, e, ω and M0	"""
	function calc_rv_keplerian end
	calc_rv_keplerian(t, p::Vector) = calc_rv_keplerian(t, p...)
	function calc_rv_keplerian(t, P,K,e,ω,M0)
		mean_anom = t*2π/P-M0
		ecc_anom = calc_ecc_anom(mean_anom,e)
		true_anom = calc_true_anom(ecc_anom,e)
		rv = K/sqrt((1-e)*(1+e))*(cos(ω+true_anom)+e*cos(ω))
	end
end


# use Convex to get access to its types and functions
using Convex

# import reshape to extend it.
import Base.reshape

export partialtransposematrix, partialtranspose, partialtransposenaive

# switch between coordinates of a matrix and its vectorised form
vecCoordToMatrixCoord(k::Integer, n::Integer) = (1 + (k-1)%n, round(Int, 1+floor((k-1)/n)%n))
matrixCoordToVecCoord(i::Integer, j::Integer, n::Integer) = i + (j-1)*n

#Determine the new coordinates `(x, y)` of an entry
#with coordinates `(i, j)` after partial transposition.

function partialtranseposecoord(i::Integer, j::Integer, n::Integer, l::Integer)
  @assert n % l == 0 "l should divide n"
  @assert n ≥ l "n should be greater than l"

  # calculate which block (i, j) is in
  a = floor(Int, (i-1)/l)+1
  b = floor(Int, (j-1)/l)+1

  x = j + (a-b)*l
  y = i + (b-a)*l

  @assert x > 0
  @assert y > 0
  @assert n ≥ x
  @assert n ≥ y

  return (x, y)
end

#Construct a permutation matrix for a `n × n` matrix
#in vector form where we transpose in `l × l` blocks.

function partialtransposematrix(n::Integer, l::Integer)
  @assert n % l == 0 "l should divide n"
  @assert n ≥ l "n should be greater than l"

  # Julia will iterate over I and J simulatiously
  # and set (I[k], J[k]) to 1 in the returned matrix
  # The value in I corresponds to the vectorised coordinates
  # before partial transposition, whereas the value in J
  # corresponds to the coordinates after partial transposition.
  IM = Array(Int, n^2)
  JM = Array(Int, n^2)
  index = 1
  for i in 1:n
    for j in 1:n
      (x, y) = partialtranseposecoord(i, j, n, l)
      IM[index] = matrixCoordToVecCoord(i, j, n)
      JM[index] = matrixCoordToVecCoord(x, y, n)
      index += 1
    end
  end

  return sparse(IM, JM, 1)
end

""" `rhoOut = partialtranspose(rho)` or `rhoOut = partialtranspose(rho, sys)` or `rhoOut = partialtranspose(rho, sys, dim)`

Applies the partial tanspose to the matrix *rho*.

Inputs:
- *rho* input matrix
- *systems* (optional) single system or a vector of systems to be transposed. By default it is equal to 2. If *dim* is not provided, the only two possible values of *systems* are 1 or 2.
- *dim* (optional) vector with dimensions of the systems. By default it assumes two systems of equal dimensions.

Outputs:
- *rhoOut* matrix *rho* with partial transpose applied
"""
function partialtranspose(ρ::Union{AbstractArray, AbstractExpr}, systems::Vector, dim::Vector)
  result = ρ
  for sys in systems
    result = partialtranspose(result, sys, dim)
  end
  return result
end

#Apply the partial tanspose to a single sytem. If no
#dimensions are specified, it will assume there are
#two systems of identical dimension. If no system is
#specified, it will transpose the last system. This
#will always be the second system, because it's not
#possible to leave out the system while providing
#the dimensions.

function partialtransposesystem(ρ::Union{AbstractArray, AbstractExpr}, sys::Integer = 2, dim::Vector = [])
  # do naive partial transposition
  if dim == []
    l = round(Int, sqrt(size(ρ)[1]))
    if sys == 2
      return partialtranspose(ρ, l, naive = true)
    else
      dim = [l; l]
    end
  end

  @assert 1 ≤ sys ≤ length(dim) "System doesnt exist, expected sys to be between 1 and ", length(dim), ". Is your dim vector correct?"

  # permute systems such that the system to apply PT is last
  perm = collect(1:length(dim))
  deleteat!(perm, sys)
  push!(perm, sys)

  l = dim[sys]
  x = permutesystems(ρ, perm, dim)
  y = partialtransposenaive(x, l)
  z = permutesystems(y, perm, dim)
  return z
end

#Apply partial transpose on systems if naive is false.
#Otherwise do the naive partial transpose in blocks
#of `l × l`, where `l` takes the value `systems`.
#This is equivalent to applying the partial
#transepose to the last system if that particular system
#has dimension `l`.

function partialtranspose(ρ::Union{AbstractArray, AbstractExpr}, systems::Integer = 2, dim::Vector = []; naive::Bool = false)
  if naive
    return partialtransposenaive(ρ, systems)
  end

  return partialtransposesystem(ρ, systems, dim)
end

#This does the naive version of the partial transpose
#that I wrote first it simply applies a transpose in blocks
#of `l × l`. This is equivalent to applying the partial
#transepose to the last system if that particular system
#has dimension `l`.

function partialtransposenaive(ρ::Union{AbstractExpr, AbstractArray, Convex.Variable}, l::Integer)
  # if the input is sparse, keep the output sparse as well
  if !isa(ρ, Convex.Variable) && !isa(ρ, AbstractExpr) && issparse(ρ)
    vectorised = sparsevec(ρ)
  else
    vectorised = vec(ρ)
  end
  n = size(ρ)[1]

  @assert n % l == 0 "l should divide n"
  @assert n ≥ l "n should be greater than l"

  return reshape(partialtransposematrix(n, l) * vectorised, n, n)
end

# some aliases for convenience
matrixCoordToVecCoord(ij::Tuple{Int, Int}, n::Integer) = matrixCoordToVecCoord(ij[1], ij[2], n)
partialtranseposecoord(ij::Tuple{Int, Int}, n::Integer, l::Integer) = partialtranseposecoord(ij[1], ij[2], n, l)
reshape(x::AbstractExpr, mn::Tuple{Int, Int}) = reshape(x, mn[1], mn[2])


using Test, Flux, NeuralNetDiffEq
using DiffEqDevTools

# Run a solve on scalars
linear = (u,p,t) -> cos(2pi*t)
tspan = (0.0f0, 1.0f0)
u0 = 0.0f0
prob = ODEProblem(linear, u0 ,tspan)
chain = Flux.Chain(Dense(1,5,σ),Dense(5,1))
opt = Flux.ADAM(0.1, (0.9, 0.95))
sol = solve(prob, NeuralNetDiffEq.NNODE(chain,opt), dt=1/20f0, verbose = true,
            abstol=1e-10, maxiters = 200)

# Run a solve on vectors
linear = (u,p,t) -> [cos(2pi*t)]
tspan = (0.0f0, 1.0f0)
u0 = [0.0f0]
prob = ODEProblem(linear, u0 ,tspan)
chain = Flux.Chain(Dense(1,5,σ),Dense(5,1))
opt = Flux.ADAM(0.1, (0.9, 0.95))
sol = solve(prob, NeuralNetDiffEq.NNODE(chain,opt), dt=1/20f0, abstol=1e-10,
            verbose = true, maxiters=200)

#Example 1
linear = (u,p,t) -> @. t^3 + 2*t + (t^2)*((1+3*(t^2))/(1+t+(t^3))) - u*(t + ((1+3*(t^2))/(1+t+t^3)))
linear_analytic = (u0,p,t) -> [exp(-(t^2)/2)/(1+t+t^3) + t^2]
prob = ODEProblem(ODEFunction(linear,analytic=linear_analytic),[1f0],(0.0f0,1.0f0))
chain = Flux.Chain(Dense(1,5,σ),Dense(5,1))
opt = Flux.ADAM(0.1, (0.9, 0.95))
sol  = solve(prob,NeuralNetDiffEq.NNODE(chain,opt),verbose = true, dt=1/5f0)
err = sol.errors[:l2]
sol  = solve(prob,NeuralNetDiffEq.NNODE(chain,opt),verbose = true, dt=1/20f0)
sol.errors[:l2]/err < 0.5

#=
dts = 1f0 ./ 2f0 .^ (6:-1:2)
sim = test_convergence(dts, prob, NeuralNetDiffEq.NNODE(chain, opt))
@test abs(sim.𝒪est[:l2]) < 0.1
@test minimum(sim.errors[:l2]) < 0.5
=#

#Example 2
linear = (u,p,t) -> -u/5 + exp(-t/5).*cos(t)
linear_analytic = (u0,p,t) ->  exp(-t/5)*(u0 + sin(t))
prob = ODEProblem(ODEFunction(linear,analytic=linear_analytic),0.0f0,(0.0f0,1.0f0))
chain = Flux.Chain(Dense(1,5,σ),Dense(5,1))
opt = Flux.ADAM(0.1, (0.9, 0.95))
sol  = solve(prob,NeuralNetDiffEq.NNODE(chain,opt),verbose = true, dt=1/5f0)
err = sol.errors[:l2]
sol  = solve(prob,NeuralNetDiffEq.NNODE(chain,opt),verbose = true, dt=1/20f0)
sol.errors[:l2]/err < 0.5

#=
dts = 1f0 ./ 2f0 .^ (6:-1:2)
sim = test_convergence(dts, prob, NeuralNetDiffEq.NNODE(chain, opt))
@test abs(sim.𝒪est[:l2]) < 0.5
@test minimum(sim.errors[:l2]) < 0.1
=#


using Turing, ReverseDiff, Memoization
using DifferentialEquations
using Plots
using Statistics
using StatsPlots
using DataFrames
using CSV
Turing.setadbackend(:reversediff)
Turing.setrdcache(true)
Turing.turnprogress(true)
using Distributions
using LinearAlgebra

# here we are simulating the signal of two coupled oscillators
# with amplitudes A1 and A2 of kBT/k and kBT/(k+c) where c is the coupling constant
# D = kBT/gamma
# in terms of the standard parameters of a OU process:
# sigma = sqrt(2*D)
# Theta = D/A

function corr_osc(c)
    # function that returns two time series that are correlated through a coupling coefficient c
    μ = 0.0 # mean is zero
    σ = sqrt(2) # D=1
    Θ1 = 1.0
    Θ2 = 1.0+abs(c)

    W1 = OrnsteinUhlenbeckProcess(Θ1,μ,σ,0.0,1.0)
    W2 = OrnsteinUhlenbeckProcess(Θ2,μ,σ,0.0,1.0)
    prob1 = NoiseProblem(W1,(0.0,100.0))
    prob2 = NoiseProblem(W2,(0.0,100.0))
    sol1 = solve(prob1;dt=0.1)
    sol2 = solve(prob2;dt=0.1)

    # creating the two correlated
    x1 = (sol1.u .+ sol2.u)/2
    if c>0
        x2 = (sol1.u .- sol2.u)/2
    else
        x2 = (sol2.u .- sol1.u)/2
    end
    return x1,x2
end

# Ornstein-Uhlenbeck process
@model ou(rn,T,delta_t) = begin
    ampl ~ Uniform(0.0,5.0)
    b ~ beta(5.0,1.0)

    rn[1] ~ Normal(0,sqrt(ampl))

    for i=2:T
        rn[i] ~ Normal(rn[i-1]*b,sqrt(ampl*(1-b^2)))
    end
end

# Ornstein-Uhlenbeck process with added Gaussian noise
@model oupn(rn,T,delta_t,::Type{R}=Vector{Float64}) where {R} = begin
    ampl ~ Uniform(0.0,5.0)
    b ~ Beta(5.0,1.0)
    noise_ampl ~ Uniform(0.0,1)

    b = exp(-delta_t/tau)
    r = R(undef, T)

    r[1] ~ Normal(0,sqrt(ampl))

    for i=2:T
        r[i] ~ Normal(r[i-1]*b,sqrt(ampl*(1-b^2)))
    end
    rn ~ MvNormal(r,sqrt(noise_ampl))
end

@model ou_corr(rn1,rn2,T,delta_t) = begin
    ampl1 ~ Uniform(0.0,5.0)
    ampl2 ~ Uniform(0.0,5.0)
    d ~ Uniform(0.0,5.0)
    b1 = exp(-delta_t*d/ampl1)
    b2 = exp(-delta_t*d/ampl2)

    rn1[1] ~ Normal(0,sqrt(ampl1))
    rn2[1] ~ Normal(0,sqrt(ampl2))

    for i=2:T
        rn1[i] ~ Normal(rn1[i-1]*b1,sqrt(ampl1*(1-b1^2)))
        rn2[i] ~ Normal(rn2[i-1]*b2,sqrt(ampl2*(1-b2^2)))
    end
end

# here we would like to run MCMC on several artificial data sets with varying correlation coefficients
# we should test rho from -0.9 to 0.9
rho_list = -0.9:0.1:0.9
results = DataFrame(rho = Float64[], coupling = Float64[], pearson = Float64[], c = Float64[], dc = Float64[])
tries = 5
for rho in rho_list
    coupling = 2*abs(rho)/(1-abs(rho))*sign(rho)
    println("coupling: ",coupling," rho: ",rho)
    # do rho tries times
    for i in 1:tries
        x1, x2 = corr_osc(coupling)
        pearson = Statistics.cor(x1,x2)
        println("Pearson: ",pearson)

        # lets see whether we can estimate c from the data
        y1 = x1 .+ x2
        y2 = x1 .- x2

        # Ornstein-Uhlenbeck process of two coupled oscillators
        model_fct = ou_corr(y1,y2,length(y1),0.1)
        chn = sample(model_fct, NUTS(0.65), 5000)
#        Turing.emptyrdcache()

        print(describe(chn))
        println(chn)
        ampl1 = Array(chn[:ampl1])
        ampl2 = Array(chn[:ampl2])

        a1 = std(ampl1)
        a2 = std(ampl2)

        if a1>a2
            c = (ampl1 .- ampl2)./ampl2
        else
            c = (ampl1 .- ampl2)./ampl1
        end
        c_mean = mean(c)
        dc = std(c)

        println("A1,A2: ",a1,",",a2)
        println("c estimate: ",c_mean,"std: ",dc)

        push!(results,[rho,coupling,pearson,c_mean,dc])
    end
end
println(results)
CSV.write("results2.csv",results)


#!/usr/bin/env julia
# -*- coding: utf-8 -*-
#=
Apply HANTS process on an integer array
=#
cd(dirname(@__FILE__))
include("../HANTS.jl")
using .HANTS, Plots

y = [
     559,  370,  626, 1201, 1782, 2444, 2131, 4946,
    3287, 5617, 3803, 6412, 5324, 8277, 6540, 8166,
    6114, 6615, 3980, 3398, 1139,  867,  586
]

nbase = 365
nfreq = 3
tseries = 1:16:365
validrange = (-10000, 10000)
fet = 0.02
dod = 5
δ = 0.1

A, φ, yrec = hants(
    y, fet, dod, δ; nbase=nbase, nfreq=nfreq, validrange=validrange, tseries=tseries, outlier=nothing
)
y_reconstruct = reconstruct(A, φ, nbase)

A_Hi, φ_Hi, yrec_Hi = hants(
    y, fet, dod, δ; nbase=nbase, nfreq=nfreq, validrange=validrange, tseries=tseries, outlier=:Hi
)
y_reconstruct_Hi = reconstruct(A_Hi, φ_Hi, nbase)

A_Lo, φ_Lo, yrec_Lo = hants(
    y, fet, dod, δ; nbase=nbase, nfreq=nfreq, validrange=validrange, tseries=tseries, outlier=:Lo
)
y_reconstruct_Lo = reconstruct(A_Lo, φ_Lo, nbase)

fig1 = plot(tseries, y, lc=:black, label="Original Data", shape=:circle, mc=:black)
plot!(tseries, yrec, lc=:blue, label="HANTS - None", shape=:circle, mc=:blue)
plot!(tseries, yrec_Hi, lc=:red, label="HANTS - Hi", shape=:circle, mc=:red)
plot!(tseries, yrec_Lo, lc=:green, label="HANTS - Lo", shape=:circle, mc=:green)
title!("Testing HANTS Algorithm")

fig2 = plot(tseries, y, lc=:black, label="Original Data", shape=:circle, mc=:black)
plot!(1:365, y_reconstruct, lc=:blue, label="HANTS - None")
plot!(1:365, y_reconstruct_Hi, lc=:red, label="HANTS - Hi")
plot!(1:365, y_reconstruct_Lo, lc=:green, label="HANTS - Lo")
title!("Testing HANTS Reconstruct")

plot(fig1, fig2; layout=(1, 2), size=(800,300))


# # clip-11-16.jl

# Load Julia packages (libraries) needed  for the snippets in chapter 0

using StatisticalRethinking, Optim
#gr(size=(600,600));

# ### snippet 3.11

p_grid = range(0, step=0.001, stop=1)
prior = ones(length(p_grid))
likelihood = [pdf(Binomial(3, p), 3) for p in p_grid]
posterior = likelihood .* prior
posterior = posterior / sum(posterior)

# Draw 10000 samples from this posterior distribution

N = 10000
samples = sample(p_grid, Weights(posterior), N);

# In StatisticalRethinkingJulia samples will always be stored
# in an MCMCChains.Chains object.

chn = MCMCChains.Chains(reshape(samples, N, 1, 1), ["toss"]);

# ### snippet 3.12

MCMCChains.describe(chn)

# ### snippet 3.13

MCMCChains.hpd(chn, alpha=0.5)

# ### snippet 3.14

mode(samples)

# ### snippet 3.15

mean(samples)

# ### snippet 3.16

median(samples)

# End of `03/clip-11-16.jl`


using Test
using TestSetExtensions
using Qaintessent
using StatsBase
using LinearAlgebra

@testset ExtendedTestSet "util pauli_vector" begin
    n1, n2, n3 = sample(1:15, 3, replace=true)
    @test all(Qaintessent.pauli_vector(n1, n2, n3) .≈ matrix(X) * n1 + matrix(Y) * n2 + matrix(Z) * n3)
end

@testset ExtendedTestSet "util intlog2" begin
    @test_throws ErrorException("Logarithm of 0 is undefined") Qaintessent.intlog2(0)
    @test_throws ErrorException("Logarithm of negative number is undefined") Qaintessent.intlog2(-30)
    N = rand(1:50)
    @test Qaintessent.intlog2(2^N) == N
    N = rand(1:200)
    @test Qaintessent.intlog2(N) <= log(2, N)
    @test Qaintessent.intlog2(N) >= (log(2, N) - 1)
end

@testset ExtendedTestSet "util binary_digits" begin
    n1 = abs(rand(Int, 1)[])
    # n1 = 5
    M = Qaintessent.intlog2(n1) + 1
    bd = Qaintessent.binary_digits(M, n1)
    x = 0
    count = 0

    for i in bd
        x+= i * 2^count
        count += 1
    end

    @test x ≈ n1
    @test Qaintessent.binary_to_int(bd) ≈ n1
end

@testset ExtendedTestSet "util binary_digits!" begin
    # n1 = abs(rand(Int, 1)[])
    n1 = 5
    M = Qaintessent.intlog2(n1) + 1
    m = BitArray{1}(undef, M)
    bd = Qaintessent.binary_digits!(m, n1)
    x = 0
    count = 0

    for i in bd
        x+= i * 2^count
        count += 1
    end

    @test x ≈ n1
end

@testset ExtendedTestSet "util quaternary_digits" begin
    n1 = abs(rand(Int, 1)[])
    M = Qaintessent.intlog2(n1)
    qd = Qaintessent.quaternary_digits(M, n1)
    x = 0
    count = 0
    for i in qd
        x+= i * 4^count
        count += 1
    end
    @test x ≈ n1
end

@testset ExtendedTestSet "util quaternary_digits!" begin
    n1 = abs(rand(Int, 1)[])
    M = Qaintessent.intlog2(n1)
    m = Vector{Int}(undef, M)
    qd = Qaintessent.quaternary_digits!(m, n1)
    x = 0
    count = 0
    for i in qd
        x+= i * 4^count
        count += 1
    end
    @test x ≈ n1
end

@testset ExtendedTestSet "util gramm_schmidt!" begin
    N = rand(4:10)

    @testset "util gramm_schmidt! complex" begin
        a = rand(ComplexF64, (N,N))

        Qaintessent.gramm_schmidt!(a)

        for i in 1:N
            for j in (i+1):N
                @test norm(dot(a[:, i], a[:, j])) < 1e-12
            end
            @test dot(a[:, i], a[:, i]) ≈ 1
        end
    end

    @testset "util gramm_schmidt! float" begin
        a = rand(Float64, (N,N))

        Qaintessent.gramm_schmidt!(a)

        for i in 1:N
            for j in (i+1):N
                @test norm(dot(a[:, i], a[:, j])) < 1e-12
            end
            @test dot(a[:, i], a[:, i]) ≈ 1
        end
    end
end

##==----------------------------------------------------------------------------------------------------------------------

@testset ExtendedTestSet "util reduced density matrix" begin
    N = 4
    ψ = randn(ComplexF64, 2^N)
    χ = randn(ComplexF64, 2^N)
    # full density matrix |ψ⟩⟨χ|
    ρ = reshape(kron(conj(χ), ψ), 2^N, 2^N)

    id = Matrix{ComplexF64}(I, 2, 2)
    A = randn(ComplexF64, 2, 2)
    B = randn(ComplexF64, 2, 2)
    @testset ExtendedTestSet "reduced density matrix correctness" begin
        @test sum(kron(A, B) .* rdm(N, (4, 2), ψ, χ)) ≈ sum(kron(B, id, A, id) .* ρ)
    end

    @testset ExtendedTestSet "reduced density matrix exceptions" begin
        @test_throws ErrorException("Need at least one wire to act on.") rdm(N, (), ψ, χ)
        @test_throws ErrorException("Number of gate wires cannot be larger than total number of wires.") rdm(N, (1, 2, 3, 4, 5), ψ, χ)
        @test_throws ErrorException("Wire indices must be unique.") rdm(N, (2, 2), ψ, χ)
        @test_throws ErrorException("Wire index cannot be smaller than 1.") rdm(N, (-1, 2), ψ, χ)
        @test_throws ErrorException("Wire index cannot be larger than total number of wires.") rdm(N, (5, 1), ψ, χ)
    end
end


import PyPlot
const plt = PyPlot

function eval_LF1d_p( α::Int64, L::Float64, grid_x::Array{Float64}, x)
    Nbasis = size(grid_x)[1]
    N = (Nbasis-1)/2
    f = 0.0
    for l = 1:Nbasis
        k = -N + (l-1)
        f = f + sqrt(1/L/Nbasis)*cos(2*pi*k*(x-grid_x[α])/L)
    end
    return f
end

function gen_T_matrix( L::Float64, grid_x::Array{Float64} )
    Nbasis = size(grid_x)[1]
    N = (Nbasis-1)/2
    # Transformation matrix
    T = Array{Complex128}(Nbasis,Nbasis)
    for α = 1:Nbasis
        for l = 1:Nbasis
            k = -N + (l-1)
            x = grid_x[α]
            T[α,l] = sqrt(1/Nbasis)*exp(2*pi*im*k*x/L)
        end
    end
    return T
end

# Example periodic function
function myfunc(L::Float64, x::Float64)
    ω = 2*pi/L
    f = cos(ω*x)*sin(2*ω*x)
    return f
end

function do_plot_myfunc( NptsPlot::Int64, L::Float64, lim )
    NptsPlot = 200
    x = Array{Float64}(linspace(lim[1], lim[2],NptsPlot));
    y = Array{Float64}(NptsPlot)
    for i = 1:NptsPlot
        y[i] = myfunc(L,x[i])
    end
    plt.clf()
    plt.grid()
    plt.plot(x,y)
    plt.savefig("myfunc.png", dpi=300)
end

function eval_from_ex_coefs( ex_coef::Array{Float64},
        L::Float64, grid_x::Array{Float64}, x::Float64 )
    Nbasis = size(ex_coef)[1]
    f = 0.0
    for i = 1:Nbasis
        f = f + ex_coef[i]*eval_LF1d_p( i, L, grid_x, x )
    end
    return f
end

function eval_FBR( l::Int64, L::Float64, N::Int64, x )
    k = -N + (l-1)
    f = 1/sqrt(L) * exp(im*2*pi*k*x/L)
    return f
end

function FBR_eval_from_ex_coefs( ex_coef, L, x )
    Nbasis = size(ex_coef)[1]
    N = Int64( (Nbasis-1)/2 )
    f = 0.0 + im*0.0
    for i = 1:Nbasis
        f = f + ex_coef[i]*eval_FBR( i, L, N, x )
    end
    return f
end

#function eval_DVR_from_FBR( α, T, x )
#    Nbasis = size(T)[1]
#    f = 0.0 + im*0.0
#    for l = 1:Nbasis
#        f = f + T'[l,α]*eval_FBR()
#    end
#end

function main(L::Float64, N::Int64)
    @printf("L = %f\n", L)
    @printf("N = %d\n", N)
    Nbasis = 2*N + 1
    @printf("Nbasis = %d\n", Nbasis)
    # Create grid
    grid_x = Array{Float64}(Nbasis)
    for α = 1:Nbasis
        grid_x[α] = L/Nbasis*(α-N-1)
    end
    Δ = L/Nbasis

    f1 = eval_LF1d_p( 1, L, grid_x, grid_x[2] )
    f2 = eval_LF1d_p( 2, L, grid_x, grid_x[2] )
    println("\nShould be close to zero: ", f1)
    println("\nShould be close to one: ", f2*sqrt(Δ))

    T = gen_T_matrix(L, grid_x)
    print("\nShould be close to zero: ")
    println(sum(abs.(T'-inv(T)))/Nbasis^2)

    ex_coefs = Array{Float64}(Nbasis)
    for i = 1:Nbasis
        ex_coefs[i] = myfunc(L, grid_x[i])
    end

    x = 2.0
    println("\nShould be close to each other:")
    println( myfunc(L, x) )
    println( eval_from_ex_coefs(ex_coefs*sqrt(Δ), L, grid_x, x) )

    FBR_ex_coefs = inv(T)*ex_coefs*sqrt(Δ)
    println( FBR_eval_from_ex_coefs(FBR_ex_coefs, L, x) )

    # FBR_ex_coefs = ifft(ex_coefs)
    # println( FBR_eval_from_ex_coefs(FBR_ex_coefs, L, x) )

end

main(10.0, 30)


@testset "Construction" begin

    # basic constructor
    xs = [1,2,3]
    f(x) = x^2
    q = fit(Newton, xs, f)
    @test all(f.(xs) .≈ q.(xs))

    # test tableau
    @test all(q.tableau .≈   [1 3 1; 0 4 5; 0 0 9])

    # test one, zero, variable
    x = variable()
    @test one(q)(x) == one(x)
    @test iszero(zero(q))
    @test variable(q)(x) == x

end

@testset "Conversion" begin

    xs = [1,2,3]
    f(x) = x^2
    p = fit(Newton, xs, f)
    q = convert(Polynomial, p)
    @test coeffs(q)  == [0,0,1]

end

@testset "Evaluation" begin

    xs = sort(rand(5))
    ys = sin.(xs)
    p = fit(Newton, xs, sin)
    @test p.(xs) ≈ sin.(xs)
    us = rand(5)
    @test maximum(abs.([p(x) - sin(x) for x in us])) <= 1e-1

end

@testset "Arithmetic" begin

    p = fit(Newton, [1,2,3], x -> x^2)
    q = fit(Newton, [2,3,4,5], x -> x^3)

    us = rand(5)

    @test all(p.(us) + q.(us) .≈ (p + q).(us))
    @test p.(us) .* q.(us) ≈ (p * q).(us)

end

@testset "Elementwise operations" begin

    p = fit(Newton, [1,2,3,4], x -> x^2)
    p + 1
    -p
    2p

end

@testset "roots" begin

    q = fit(Newton, [2,3,4,5], x -> (x-1)*(x-1/2)*(x+1/2))
    @test all(roots(q) .≈ [-1/2, 1/2, 1])

end


using MultipleViewGeometry, Test, LinearAlgebra
using MultipleViewGeometry.ModuleTypes
using StaticArrays, Calculus, GeometryTypes
using MultipleViewGeometry.ModuleSyntheticData

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

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

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

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

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

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

𝒳₁ = generate_planar_points(𝐧₁,d₁, 20, 50)
𝒳₂ = generate_planar_points(𝐧₂,d₂, 20, 50)

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

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

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

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

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

𝐇₁ = estimate(HomographyMatrix(), DirectLinearTransform(), (ℳ₁, ℳ₁ʹ))
𝐇₂ = estimate(HomographyMatrix(), DirectLinearTransform(), (ℳ₂, ℳ₂ʹ))

for i = zip(ℳ₁,ℳ₁ʹ)
    m, mʹ =  i
    𝐦  = hom(m)
    𝐦ʹ = hom(mʹ)
    residual = vec2antisym(𝐦ʹ)*𝐇₁*𝐦
    @test isapprox(sum(residual), 0.0; atol = 1e-7)
end

for i = zip(ℳ₂,ℳ₂ʹ)
    m, mʹ =  i
    𝐦  = hom(m)
    𝐦ʹ = hom(mʹ)
    residual = vec2antisym(𝐦ʹ)*𝐇₂*𝐦
    @test isapprox(sum(residual), 0.0; atol = 1e-7)
end

# Tests for homography matrix estimation
#
# # Normal to the planar surface.
# 𝐧 = [0.0, 0.0, -1.0]
# # Distance of the plane from the origin.
# d = 100
# # Sample points on the planar surface.
# 𝒳 = [Point3D(x,y,d) for x = -100:5:100 for y = -100:5:100]
# 𝒳 = 𝒳[1:50:end]
# # Intrinsic and extrinsic parameters of camera one.
# 𝐊₁ = Matrix{Float64}(I,3,3)
# 𝐑₁ = Matrix{Float64}(I,3,3)
# 𝐭₁ = [0.0, 0.0, 0.0]
#
# # Intrinsic and extrinsic parameters of camera two.
# 𝐊₂ = Matrix{Float64}(I,3,3)
# 𝐑₂ = Matrix{Float64}(I,3,3)
# 𝐭₂ = [100.0, 2.0, -100.0]
#
# # Camera projection matrices.
# 𝐏₁ = construct(ProjectionMatrix(),𝐊₁,𝐑₁,𝐭₁)
# 𝐏₂ = construct(ProjectionMatrix(),𝐊₂,𝐑₂,𝐭₂)
#
# # Set of corresponding points.
# ℳ = project(Pinhole(),𝐏₁,𝒳)
# ℳʹ = project(Pinhole(),𝐏₂,𝒳)
#
# # Estimate of the homography matrix and the true homography matrix.
# 𝐇 = estimate(HomographyMatrix(), DirectLinearTransform(), (ℳ, ℳʹ))
# 𝐇₀ = construct(HomographyMatrix(),𝐊₁,𝐑₁,𝐭₁,𝐊₂,𝐑₂,𝐭₂,𝐧,d)
#
# 𝐇₀ = 𝐇₀ / norm(𝐇₀)
# 𝐇₀ = 𝐇₀ / sign(𝐇₀[3,3])
#
# 𝐇 = 𝐇 / norm(𝐇)
# 𝐇 = 𝐇 / sign(𝐇[3,3])
#
# for i = zip(ℳ,ℳʹ)
#     m, mʹ =  i
#     𝐦  = hom(m)
#     𝐦ʹ = hom(mʹ)
#     residual = vec2antisym(𝐦ʹ)*𝐇₀*𝐦
#     @test isapprox(sum(residual), 0.0; atol = 1e-7)
# end
#
# for i = zip(ℳ,ℳʹ)
#     m, mʹ =  i
#     𝐦  = hom(m)
#     𝐦ʹ = hom(mʹ)
#     residual = vec2antisym(𝐦ʹ)*𝐇*𝐦
#     @show residual
#     @test isapprox(sum(residual), 0.0; atol = 1e-7)
# end
#

#@test isapprox(sum(residual), 0.0; atol = 1e-7)
#dot(𝒳[10],𝐧) + d


# ### parametric function plot 3d
# https://jp.mathworks.com/matlabcentral/fileexchange/35263-matlab-plot-gallery-function-plot-3d?focused=7851891&tab=example

#md # ![FILE_NAME.png](images/FILE_NAME.png)

t = range(0, stop=10, length=1000)
x = cos.(t)
y = sin.(t)
z = sin.(5t)

plot(x, y, z)


# General math functions of quaternions

"""
    abs2(q)

Sum the squares of the components of the quaternion

# Examples
```jldoctest
julia> abs2(Quaternion(1,2,4,10))
121
```
"""
Base.abs2(q::AbstractQuaternion) = sum(q.components.^2)
Base.abs2(q::Rotor{T}) where {T<:Real} = one(T)

"""
    abs(q)

Square-root of the sum the squares of the components of the quaternion

# Examples
```jldoctest
julia> abs(Quaternion(1,2,4,10))
11.0
```
"""
Base.abs(q::AbstractQuaternion) = sqrt(abs2(q))
Base.abs(q::Rotor{T}) where {T<:Real} = one(T)

"""
    abs2vec(q)

Sum the squares of the "vector" components of the quaternion

# Examples
```jldoctest
julia> abs2vec(Quaternion(1,2,3,6))
49
```
"""
abs2vec(q::AbstractQuaternion) = @inbounds q.components[2]^2 + q.components[3]^2 + q.components[4]^2

"""
    absvec(q)

Square-root of the sum of the squares of the "vector" components of the quaternion

# Examples
```jldoctest
julia> absvec(Quaternion(1,2,3,6))
7.0
```
"""
absvec(q::AbstractQuaternion) = sqrt(abs2vec(q))

# norm(q::Quaternion) = Base.abs2(q)  ## This might just be confusing

Base.inv(q::AbstractQuaternion) = conj(q) / abs2(q)
Base.inv(q::Rotor) = conj(q)  # Specialize to ensure output is also a Rotor

"""
    log(q)

Logarithm of a quaternion.

As with the usual complex logarithm, the quaternion logarithm has multiple
branches, though the quaternion branches are three-dimensional: for any unit
"vector" quaternion q̂, you could add any integer multiple of 2πq̂ to the result
of this function and still get the same result after exponentiating (within
numerical accuracy).  This function is the principal logarithm.

This function has discontinuous (and fairly arbitrary) behavior along the
negative real axis: if the "vector" components of the quaternion are precisely
zero *and* the scalar component is negative, the returned quaternion will have
scalar component `log(-q.w)`, but will also have a `z` component of π.  The
choice of the `z` direction is arbitrary; the "vector" component of the
returned quaternion could be π times any unit vector.

Note that this function is not specialized to unit-quaternion inputs, so the
scalar component of the returned value will be nonzero unless the input has
*precisely* unit magnitude.

# Examples
```jldoctest
julia> log(exp(1.2imy))
0.0 + 0.0𝐢 + 1.2𝐣 + 0.0𝐤

julia> log(Quaternion(exp(7)))
7.0 + 0.0𝐢 + 0.0𝐣 + 0.0𝐤

julia> log(Quaternion(-exp(7)))
7.0 + 0.0𝐢 + 0.0𝐣 + 3.141592653589793𝐤
```

"""
function Base.log(q::Quaternion{T}) where {T}
    q = float(q)
    absolute2vec = abs2vec(q)
    if iszero(absolute2vec)
        if q.w < 0
            return Quaternion(log(-q.w), 0, 0, π)
        end
        return Quaternion(log(q.w), 0, 0, 0)
    end
    absolutevec = sqrt(absolute2vec)
    f = atan(absolutevec, q.w) / absolutevec  # acos((w^2-absolutevec^2) / (w^2+absolutevec^2)) / 2absolutevec
    Quaternion(log(abs2(q))/2, f*q.x, f*q.y, f*q.z)
end
function Base.log(q::Rotor{T}) where {T}
    q = float(q)
    absolute2vec = abs2vec(q)
    if iszero(absolute2vec)
        if q.w < 0
            return QuatVec{float(T)}(0, 0, 0, π)
        end
        return QuatVec{float(T)}(0, 0, 0, 0)
    end
    absolutevec = sqrt(absolute2vec)
    f = atan(absolutevec, q.w) / absolutevec  # acos(q.w) / absolutevec
    QuatVec(0, f*q.x, f*q.y, f*q.z)
end

"""
    exp(q)

Exponential of a quaternion

# Examples
```jldoctest
julia> exp(imx*π/4)  # Rotation through π/2 (note the extra 1/2) about the x axis
0.7071067811865476 + 0.7071067811865475𝐢 + 0.0𝐣 + 0.0𝐤
```
"""
function Base.exp(q::Quaternion{T}) where {T}
    q = float(q)
    absolute2vec = abs2vec(q)
    if iszero(absolute2vec)
        return Quaternion(exp(q.w), 0, 0, 0)
    end
    absolutevec = sqrt(absolute2vec)
    e = exp(q.w)
    s, c = sincos(absolutevec)
    esinc = e * s / absolutevec
    Quaternion(e*c, esinc*q.x, esinc*q.y, esinc*q.z)
end
function Base.exp(q::QuatVec{T}) where {T}
    q = float(q)
    absolute2vec = abs2vec(q)
    if iszero(absolute2vec)
        return Rotor{float(T)}(1, 0, 0, 0)
    end
    absolutevec = sqrt(absolute2vec)
    s, c = sincos(absolutevec)
    sinc = s / absolutevec
    Rotor(c, sinc*q.x, sinc*q.y, sinc*q.z)
end

@doc raw"""
    sqrt(q)

Square-root of a quaternion.

The general formula whenever the denominator is nonzero is

```math
\sqrt{q} = \frac{|q| + q} {\sqrt{2|q| + 2q.w}}
```

This can be proven by expanding `q` as `q.w + q.vec` and multiplying the
expression above by itself.

When the denominator is zero, this function has discontinuous (and fairly
arbitrary) behavior, just as with the quaternion [`log`](@ref) function.  In
this case, either all components are zero — in which case the result is simply
the zero quaternion — or the "vector" components of the quaternion are
precisely zero and the scalar component is negative.  If the latter is true,
the denominator above will be a pure-imaginary number.  Because the quaternions
come with infinitely many elements that square to -1, it is not clear *which*
imaginary should be used, so we arbitrarily choose to set the result
proportional to the `z` quaternion.  The choice of the `z` direction is
arbitrary; the "vector" component of the returned quaternion could be in any
direction.

# Examples
```jldoctest
julia> q = Quaternion(1.2, 3.4, 5.6, 7.8);

julia> sqrtq = √q;

julia> sqrtq^2 ≈ q
true

julia> √Quaternion(4)
2.0 + 0.0𝐢 + 0.0𝐣 + 0.0𝐤

julia> √Quaternion(-4)
0.0 + 0.0𝐢 + 0.0𝐣 + 2.0𝐤
```
"""
function Base.sqrt(q::Quaternion{T}) where {T}
    q = float(q)
    absolute2vec = abs2vec(q)
    if iszero(absolute2vec)
        if q.w < 0
            return Quaternion(0, 0, 0, sqrt(-q.w))
        end
        return Quaternion(sqrt(q.w), 0, 0, 0)
    end
    absolute2 = absolute2vec + q.w^2
    c1 = sqrt(absolute2) + q.w
    c2 = sqrt(inv(2*c1))
    Quaternion(c1*c2, q.x*c2, q.y*c2, q.z*c2)
end
function Base.sqrt(q::Rotor{T}) where {T}
    q = float(q)
    absolute2vec = abs2vec(q)
    if iszero(absolute2vec)
        if q.w < 0
            return Rotor{float(T)}(0, 0, 0, 1)
        end
        return Rotor{float(T)}(1, 0, 0, 0)
    end
    c1 = 1 + q.w
    c2 = sqrt(inv(2*c1))
    Rotor(c1*c2, q.x*c2, q.y*c2, q.z*c2)
end

"""
    angle(q)

Phase angle in radians of the rotation represented by this quaternion.

Note that this may be different from your interpretation of the angle of a
complex number in an important way.  Because quaternions act on vectors by
conjugation — as in `q*v*conj(q)` — there are *two* copies of `q` involved in
that expression; in some sense, a quaternion acts "twice".  Therefore, this
angle may be twice what you expect from an analogy with complex numbers —
dpending on how you interpret the correspondence between complex numbers and
quaternions.  Also, while rotations in the complex plane have a natural choice
of axis (the positive `z` direction), that is not the case for quaternions,
which means that the sign of this angle is arbitrary, and we always choose it
to be positive.

# Examples
```jldoctest
julia> θ=1.2;

julia> R=exp(θ * imz / 2);

julia> angle(R)
1.2

```
"""
Base.angle(q::Quaternion{T}) where T = 2 * absvec(log(q))
Base.angle(q::Rotor{T}) where T = 2 * absvec(log(q))

function Base.:^(q::Quaternion, s::Real)
    exp(s * log(q))
end
function Base.:^(q::Rotor, s::Real)
    q = float(q)
    absolutevec = absvec(q)
    if absolutevec ≤ eps(typeof(absolutevec))
        if q.w < 0
            # log(q) ≈ π𝐤
            sin_πs, cos_πs = sincospi(oftype(absolutevec, s))
            return Rotor{eltype(q)}([cos_πs, 0, 0, sin_πs])
        end
        # log(q) ≈ 0
        return one(q)
    end
    f1 = s * atan(absolutevec, q.w)
    sin_f1, cos_f1 = sincos(f1)
    f2 = sin_f1 / absolutevec
    Rotor{typeof(cos_f1)}([cos_f1, f2*q.x, f2*q.y, f2*q.z])
end
Base.:^(q::Quaternion, s::Integer) = (s ≥ 0 ? Base.power_by_squaring(q, s) : inv(Base.power_by_squaring(q, -s)))
Base.:^(q::QuatVec, s::Integer) = (s ≥ 0 ? Base.power_by_squaring(q, s) : inv(Base.power_by_squaring(q, -s)))
Base.:^(q::Rotor, s::Integer) = (s ≥ 0 ? Base.power_by_squaring(q, s) : inv(Base.power_by_squaring(q, -s)))


using DiffEqEnvironments
using DifferentialEquations
using LinearAlgebra
using ControlSystems
using Test
using ReinforcementLearningBase
using Plots

"""
This test defines an SISO system of form
    ṡ = A*s + B*a
    o = C*s
for the rocket car problem, where
    s = [p, ṗ]  ,
p being the position of the car.
The action used as an input to the system is
the car's acceleration p̈ ∈ [-1, 1]. Therefore
    ṡ = [ṗ, p̈] = [0 1; 0 0] * s + [0; 1] * u .
Full observation of the state is given, as
    o = s   .
"""
T = Float32

# System matrices
A = [0 1; 0 0]
B = [0, 1]

# weights for quadratic cost function
Q = I * 10
R = I

# Set range of control input
a_lb = -1.0; a_ub = 1.0

# Set IC & tspan for integration
s0 = [-0.75, 0.0]
tspan = (0f0, 5f0)

# Define ODEProblem of SISO LTI system
function lti(s, a, t)
    ṡ = A * s + B * a
end
prob_lti = ODEProblem(lti, s0, tspan)

# Use ControlSystems.jl to solve continuous algebraic Riccati eq.
P = care(A, B, Q, R) # P encodes value function as v(s)= -s'*P*s/2
K = -R \ B' * P
println("K=", K)

v_lqr(s) = -s' * P * s / 2
π_lqr(s) = clamp.(K * s, a_lb, a_ub) #  LQR controller w/ saturation

# Define closed-loop ODE for system with input saturations
function cl_ode(s, π, t)
    a = π(s)
    ṡ = lti(s, a, t)
end

function plot_value(v; levels=10)
    ps = range(-1, 1, length=100)
    ṗs = range(-1, 1, length=100)
    vp(p, ṗ) = v([p,ṗ])

    contour(ps, ṗs, vp, levels=levels,
        xlabel="p", ylabel="ṗ",
        fill=true, show=true)
end

p = plot_value(v_lqr)

# Plot LQR Solution
prob_lqr = ODEProblem(cl_ode, s0, tspan, π_lqr)
sol_lqr = solve(prob_lqr, Tsit5())

function plot_phase_portrait!(plt, π; n=20, scale=0.075)
    ṡ(s) = lti(s, π(s), 0) * scale
    ṡ(p, ṗ) = ṡ([p,ṗ])

    ps = range(-1, 1, length=n)
    ṗs = range(-1, 1, length=n)

    pps = [p for p in ps for ṗ in ṗs]
    ṗṗs = [ṗ for p in ps for ṗ in ṗs]
    quiver!(pps, ṗṗs, quiver=ṡ, lw=1, c=:black,
        xlim=(-1, 1), ylim=(-1, 1),
        arrow=(style = :closed; headlength = 0.1))
end

function plot_trajectory!(plt, s0, π; c=:black)
    prob = ODEProblem(cl_ode, s0, tspan, π)
    sol = solve(prob, Tsit5(), saveat=0.01)

    plot!(plt, sol[1,:], sol[2,:], lw=3, c=c,
        xlim=(-1, 1), ylim=(-1, 1), line=:arrow,
        legend=false, show=true)
end

p2 = plot_phase_portrait!(p, π_lqr)
p3 = plot_trajectory!(p, s0, π_lqr)

display(p)

# Define DiffEqEnv
r = QuadraticRewardFunction(Q, R)
n_actions = 1
dt = 0.1
env = DiffEqEnv(prob_lti, r, n_actions, dt, a_lb=a_lb, a_ub=a_ub)


abstract type AbstractTensorProductLayer <: Function end
"""
Constructs the Tensor Product Layer, which takes as input an array of n tensor
product basis, [B_1, B_2, ..., B_n] a data point x, computes
z[i] = W[i,:] ⨀ [B_1(x[1]) ⨂ B_2(x[2]) ⨂ ... ⨂ B_n(x[n])], where W is the layer's weight,
and returns [z[1], ..., z[out]].

```julia
TensorLayer(model,out,p=nothing)
```
Arguments:
- `model`: Array of TensorProductBasis [B_1(n_1), ..., B_k(n_k)], where k corresponds to the dimension of the input.
- `out`: Dimension of the output.
- `p`: Optional initialization of the layer's weight. Initialized to standard normal by default.
"""
struct TensorLayer{M<:Array{TensorProductBasis},P<:AbstractArray,Int} <: AbstractTensorProductLayer
    model::M
    p::P
    in::Int
    out::Int
    function TensorLayer(model,out,p=nothing)
        number_of_weights = 1
        for basis in model
            number_of_weights *= basis.n
        end
        if p === nothing
            p = randn(out*number_of_weights)
        end
        new{Array{TensorProductBasis},typeof(p),Int}(model,p,length(model),out)
    end
end

function (layer::TensorLayer)(x,p=layer.p)
    model,out = layer.model,layer.out
    W = reshape(p, out, Int(length(p)/out))
    tensor_prod = model[1](x[1])
    for i in 2:length(model)
        tensor_prod = kron(tensor_prod,model[i](x[i]))
    end
    z = W*tensor_prod
    return z
end


@testset "NelsonSiegelSvensson" begin

    # Per this technical note, the target/source rates should be interpreted as continuous rates:
    # https://www.ecb.europa.eu/stats/financial_markets_and_interest_rates/euro_area_yield_curves/html/technical_notes.pdf

    # EURAAA_20191111 at https://www.ecb.europa.eu/stats/financial_markets_and_interest_rates/euro_area_yield_curves/html/index.en.html
    euraaa_zeros = Continuous.([-0.602009,-0.612954,-0.621543,-0.627864,-0.632655,-0.610565,-0.569424,-0.516078,-0.455969,-0.39315,-0.33047,-0.269814,-0.21234,-0.158674,-0.109075,-0.063552,-0.021963,0.015929,0.050407,0.081771,0.110319,0.136335,0.160083,0.181804,0.201715,0.220009,0.23686,0.252419,0.26682,0.280182,0.292608,0.304191,0.31501] ./ 100)
    euraaa_pars = Continuous.([-0.601861,-0.612808,-0.621403,-0.627731,-0.63251,-0.610271,-0.568825,-0.515067,-0.454526,-0.391338,-0.328412,-0.26767,-0.210277,-0.156851,-0.107632,-0.062605,-0.021601,0.015642,0.049427,0.080073,0.107892,0.133178,0.156206,0.17722,0.196444,0.214073,0.230281,0.245221,0.259027,0.271818,0.283696,0.294753,0.305069] ./ 100)
    euraaa_maturities = vcat([0.25, 0.5, 0.75], 1:30)

    @testset "NelsonSiegel" begin
        @testset "EURAAA_20191111" begin
            c_param = Yields.NelsonSiegelCurve(0.6202603126029489 /100, -1.1621281759833935 /100, -1.930016080035979 /100, 3.0)
            c = Yields.Zero(NelsonSiegel(),euraaa_zeros, euraaa_maturities)
            # @testset "parameter: $param" for param in [:β₀, :β₁, :β₂, :τ₁]
            #     @test getfield(c, param) ≈ getfield(c_param, param)
            # end

            @testset "zero rates: $t" for (t, r) in zip(euraaa_maturities, euraaa_zeros)
                @test Yields.zero(c, t) ≈ r   atol = 0.001
            end

            @testset "par rates: $t" for (t, r) in zip(euraaa_maturities, euraaa_pars)
                @test Yields.zero(c, t) ≈ r   atol = 0.001
            end

            @test discount(c,0) == 1.0
            @test discount(c,10) ≈ 1 / accumulation(c,10)

            # test that rates as floats work
            fzeros = Yields.rate.(euraaa_zeros)
            c_fzero = Yields.Zero(NelsonSiegel(),fzeros, euraaa_maturities)
            @test c_fzero == c
        end

        # Nelson-Siegel-Svensson package example at https://nelson-siegel-svensson.readthedocs.io/en/latest/usage.html
        @testset "pack" begin
            pack_yields = Continuous.([0.01, 0.011, 0.013, 0.016, 0.019, 0.021, 0.026, 0.03, 0.035, 0.037, 0.038, 0.04])
            pack_maturities = [10e-5, 0.5, 1.0, 2.0, 3.0, 4.0, 5.0, 10.0, 15.0, 20.0, 25.0, 30.0]
            c_param = Yields.NelsonSiegelCurve(0.04495841387198023, -0.03537510042719209, 0.0031561222355027227, 5.0)
            c = Yields.Zero(NelsonSiegel(),pack_yields, pack_maturities)

            @testset "zero rates: $t" for (t, r) in zip(pack_maturities, pack_yields)
                @test Yields.zero(c, t) ≈ r   atol = 0.003
            end
        end
    end

    @testset "NelsonSiegelSvensson" begin

        @testset "EURAAA_20191111" begin
            c_zero = Yields.Zero(NelsonSiegelSvensson(),euraaa_zeros, euraaa_maturities)
            c_par = Yields.Par(NelsonSiegelSvensson(),euraaa_pars, euraaa_maturities)

            @testset "par and zero constructors" for c in [c_zero,c_par]
                @testset "zero rates: $t" for (t, r) in zip(euraaa_maturities, euraaa_zeros)
                    @test Yields.zero(c, t) ≈ r   atol = 0.0001
                end

                @testset "par rates: $t" for (t, r) in zip(euraaa_maturities, euraaa_pars)
                    # are the target rates on the ECB site continuous rates or periodic/bond-equivalent?
                    @test Yields.par(c, t) ≈ r   atol = 0.0001
                end
            end

            # test that rates as floats work
            fzeros = Yields.rate.(euraaa_zeros)
            c_fzero = Yields.Zero(NelsonSiegelSvensson(),fzeros, euraaa_maturities)
            @test c_fzero == c_zero

        end
        @testset "EURAAA_20191111 w parms given" begin
            c = Yields.NelsonSiegelSvenssonCurve(0.629440 / 100, -1.218082 /100, 	12.114098 /100, -14.181117 /100, 2.435976, 2.536963)

            @testset "zero rates: $t" for (t, r) in zip(euraaa_maturities, euraaa_zeros)
                @test Yields.zero(c, t) ≈ r   atol = 0.0001
            end

            @testset "par rates: $t" for (t, r) in zip(euraaa_maturities, euraaa_pars)
                # are the target rates on the ECB site continuous rates or periodic/bond-equivalent?
                @test Yields.par(c, t) ≈ r atol = 0.0001
            end

            @test Yields.zero(c,30) ≈ last(euraaa_zeros) atol = 0.0001
            @test discount(c,0) == 1.0
            @test discount(c,10) ≈ 1 / accumulation(c,10)
        end
    end

end


using Revise
using Statistics
using LaTeXStrings
using PyPlot
using LinearAlgebra
using Infiltrator
using Bem2d

"""
    discretized_arc(θstart, θend, radius, n_pts)

Generate regularly spaced eleemnts along an curved arc.
"""
function discretized_arc(θstart, θend, radius, n_pts)
    # Create geometry of discretized arc
    θrange = collect(LinRange(θstart, θend, n_pts + 1))
    x = @. radius * cos(θrange)
    y = @. radius * sin(θrange)
    x1 = x[1:1:end-1]
    x2 = x[2:1:end]
    y1 = y[1:1:end-1]
    y2 = y[2:1:end]
    return x1, y1, x2, y2
end

"""
    circle_subplot(nrows, ncols, plotidx, x, y, mat, npts, R, theta0, title_string)

Plot field (displacement, stress) within a circular disk and style
"""
function circle_subplot(nrows, ncols, plotidx, x, y, mat, npts, R, theta0, title_string)
    fontsize = 20
    contour_levels = 100
    contour_color = "white"
    contour_linewidth = 0.5
    color_scale = 1

    subplot(nrows, ncols, plotidx)
    contourf(reshape(x, npts, npts), reshape(y, npts, npts), reshape(mat, npts, npts), levels=contour_levels)
    cbar = colorbar(fraction=0.05, pad=0.05, extend = "both")
    cbar.ax.tick_params(labelsize=fontsize)
    contour(reshape(x, npts, npts), reshape(y, npts, npts), reshape(mat, npts, npts), levels=contour_levels, colors=contour_color, linewidths=contour_linewidth)
    xlabel("x (m)", fontsize=fontsize)
    ylabel("y (m)", fontsize=fontsize)
    title(title_string, fontsize=fontsize)

    # Draw entire circle and region of applied tractions
    x1, y1, x2, y2 = discretized_arc(deg2rad(-180), deg2rad(180), R, 360)
    plot([x1, x2], [y1, y2], "-k", linewidth=2)
    x1, y1, x2, y2 = discretized_arc(-theta0, theta0, R, 50)
    plot([x1, x2], [y1, y2], "-r", linewidth=2)
    x1, y1, x2, y2 = discretized_arc(-theta0+deg2rad(180), theta0+deg2rad(180), R, 50)
    plot([x1, x2], [y1, y2], "-r", linewidth=2)
    gca().set_aspect("equal")
    gca().tick_params(labelsize=fontsize)
end

"""
    calcbrazil(p, x, y, R, theta0)

Calculate stresses predicted from Hondros solution to Brazil test.
"""
function calcbrazil(p, x, y, R, theta0)
    #! Solution from Hondros (1959) as summarized by Wei and Chau 2013
    r = @. sqrt(x^2 + y^2)
    theta = @. atan(y, x)

    # Analytic stresses in cylindrical coordinates
    Srr = zeros(length(x))
    Sthetatheta = zeros(length(x))
    Srtheta = zeros(length(x))
    Srrconstterm = 2.0 * theta0 * -p / deg2rad(180)
    Sthetathetaconstterm = 2.0 * theta0 * -p / deg2rad(180)
    leadingterm = 2.0 * -p / deg2rad(180)
    mmax = 1000 # Max number of terms in Hondros series
    for m in 1:mmax
        Srr += @. (r/R)^(2*m-2) * (1-(1-1/m)*(r/R)^2) * sin(2*m*theta0) * cos(2*m*theta)
        Sthetatheta += @. (r/R)^(2*m-2) * (1-(1+1/m)*(r/R)^2) * sin(2*m*theta0) * cos(2*m*theta)
        Srtheta += @. ((r/R)^(2*m) - (r/R)^(2*m-2)) * sin(2*m*theta0) * sin(2*m*theta)
    end
    Srr = @. Srrconstterm + leadingterm * Srr
    Sthetatheta = @. Sthetathetaconstterm - leadingterm * Sthetatheta
    Srtheta = @. leadingterm * Srtheta

    # Convert analytic cylindrical stresses to Cartesian
    Sanalytic = zeros(length(x), 3)
    for i in 1:length(x) # Project a single matrix to Cartesian coordinates
        cylindrical_stress_tensor = [Srr[i] Srtheta[i] ; Srtheta[i] Sthetatheta[i]]
        transformation_matrix = [cos(theta[i]) -sin(theta[i]) ; sin(theta[i]) cos(theta[i])]
        cartesian_stress_tensor = transformation_matrix * cylindrical_stress_tensor * transpose(transformation_matrix)
        Sanalytic[i, 1] = cartesian_stress_tensor[1, 1]
        Sanalytic[i, 2] = cartesian_stress_tensor[2, 2]
        Sanalytic[i, 3] = cartesian_stress_tensor[1, 2]
    end
    return Sanalytic
end

"""
    braziltest()

Compare the Hondros (1959 analytic soluiton) for the Brazil test
with the BEM solution for both constant and quadratic element cases.
"""
function braziltest()
    close("all")
    mu = 3e10
    nu = 0.25
    p = -1.0e5 # Applied radial pressure over arc
    theta0 = deg2rad(1.0) # Arc length over which pressure is applied
    nels = 360
    R = nels / (2 * pi)
    npts = 200
    x, y = Bem2d.obsgrid(-R, -R, R, R, npts)
    r = @. sqrt(x^2 + y^2)

    #! Analytic solution
    Sanalytic = calcbrazil(p, x, y, R, theta0)

    #! BEM solution
    els = Bem2d.Elements(Int(1e5))
    x1, y1, x2, y2 = discretized_arc(deg2rad(-180), deg2rad(180), R, nels)
    for i in 1:length(x1)
        els.x1[els.endidx + i] = x1[i]
        els.y1[els.endidx + i] = y1[i]
        els.x2[els.endidx + i] = x2[i]
        els.y2[els.endidx + i] = y2[i]
        els.name[els.endidx + i] = "circle"
    end
    Bem2d.standardize_elements!(els)
    idx = Bem2d.getidxdict(els)

    #! Apply normal tractions everywhere and convert from radial to Cartesian
    xtracC = zeros(els.endidx)
    ytracC = zeros(els.endidx)
    thetaels = @. atan(els.ycenter[1:1:els.endidx], els.xcenter[1:1:els.endidx])
    for i in 1:els.endidx # Calcuate the x and y components of the tractions
        normalTractions = [0; p] # Pressure in fault normal component only.
        xtracC[i], ytracC[i] = els.rotmat[i, :, :] * normalTractions
    end

    #! Zero out the tractions on the area without contact
    deleteidx = findall(x -> (x>theta0 && x<deg2rad(180)-theta0), thetaels)
    xtracC[deleteidx] .= 0
    ytracC[deleteidx] .= 0
    deleteidx = findall(x -> (x<-theta0 && x>-deg2rad(180)+theta0), thetaels)
    xtracC[deleteidx] .= 0
    ytracC[deleteidx] .= 0

    #! Kernels, T*: displacement to displacement, H*: displacement to traction
    @time _, _, HstarC = partialsconstdispstress(slip2dispstress, els, idx["circle"], idx["circle"], mu, nu)
    @time _, _, HstarQ = partialsquaddispstress(slip2dispstress, els, idx["circle"], idx["circle"], mu, nu)

    #! CONSTANT CASE
    #! Applied tractions -> effective displacements -> internal stresses
    DeffC = inv(HstarC) * interleave(xtracC, ytracC)
    @time _, SdispC = constdispstress(slip2dispstress, x, y, els, idx["circle"], DeffC[1:2:end], DeffC[2:2:end], mu, nu)

    #! QUADRATIC CASE
    #! Applied tractions -> effective displacements -> internal stresses
    xtracQ = zeros(3 * length(xtracC))
    ytracQ = zeros(3 * length(ytracC))
    xtracQ[1:3:end] = xtracC
    xtracQ[2:3:end] = xtracC
    xtracQ[3:3:end] = xtracC
    ytracQ[1:3:end] = ytracC
    ytracQ[2:3:end] = ytracC
    ytracQ[3:3:end] = ytracC
    DeffQ = inv(HstarQ) * interleave(xtracQ, ytracQ)
    @time _, SdispQ = quaddispstress(slip2dispstress, x, y, els, idx["circle"], quadstack(DeffQ[1:2:end]), quadstack(DeffQ[2:2:end]), mu, nu)

    #! Isolate the values inside the circle
    nanidx = findall(x -> x > R, r)
    Sanalytic[nanidx, :] .= NaN
    SdispC[nanidx, :] .= NaN
    SdispQ[nanidx, :] .= NaN
    SresidualC = @. SdispC - Sanalytic
    SresidualQ = @. SdispQ - Sanalytic

    #! Summary figure
    figure(figsize=(30,20))
    fontsize = 20
    nrows = 4
    ncols = 6

    #! Constant dislocation case
    # Applied tractions
    subplot(nrows, ncols, 1)
    plot(rad2deg.(thetaels), xtracC, ".r")
    plot(rad2deg.(thetaels), ytracC, "+b")
    title("applied tractions", fontsize=fontsize)
    gca().tick_params(labelsize=fontsize)

    # Effective displacements
    subplot(nrows, ncols, 2)
    plot(rad2deg.(thetaels), DeffC[1:2:end], ".r")
    plot(rad2deg.(thetaels), DeffC[2:2:end], "+b")
    title("Effective displacements", fontsize=fontsize)
    gca().tick_params(labelsize=fontsize)

    # Analytic solutions
    circle_subplot(nrows, ncols, 7, x, y, Sanalytic[:, 1], npts, R, theta0, "Sxx (analytic)")
    circle_subplot(nrows, ncols, 8, x, y, Sanalytic[:, 2], npts, R, theta0, "Syy (analytic)")
    circle_subplot(nrows, ncols, 9, x, y, Sanalytic[:, 3], npts, R, theta0, "Sxy (analytic)")

    # BEM solutions
    circle_subplot(nrows, ncols, 13, x, y, SdispC[:, 1], npts, R, theta0, "Sxx (DDM)")
    circle_subplot(nrows, ncols, 14, x, y, SdispC[:, 2], npts, R, theta0, "Syy (DDM)")
    circle_subplot(nrows, ncols, 15, x, y, SdispC[:, 3], npts, R, theta0, "Sxy (DDM)")
    circle_subplot(nrows, ncols, 19, x, y, SresidualC[:, 1], npts, R, theta0, "Sxx (residual)")
    circle_subplot(nrows, ncols, 20, x, y, SresidualC[:, 2], npts, R, theta0, "Syy (residual)")
    circle_subplot(nrows, ncols, 21, x, y, SresidualC[:, 3], npts, R, theta0, "Sxy (residual)")

    #! Quadratic case
    # Applied tractions
    subplot(nrows, ncols, 4)
    plot(xtracQ, ".r")
    plot(ytracQ, "+b")
    title("applied tractions", fontsize=fontsize)
    gca().tick_params(labelsize=fontsize)

    # Effective displacements
    subplot(nrows, ncols, 5)
    plot(DeffQ[1:2:end], ".r")
    plot(DeffQ[2:2:end], "+b")
    title("Effective displacements", fontsize=fontsize)
    gca().tick_params(labelsize=fontsize)

    # Analytic solutions
    circle_subplot(nrows, ncols, 10, x, y, Sanalytic[:, 1], npts, R, theta0, "Sxx (analytic)")
    circle_subplot(nrows, ncols, 11, x, y, Sanalytic[:, 2], npts, R, theta0, "Syy (analytic)")
    circle_subplot(nrows, ncols, 12, x, y, Sanalytic[:, 3], npts, R, theta0, "Sxy (analytic)")

    # BEM solutions
    circle_subplot(nrows, ncols, 16, x, y, SdispQ[:, 1], npts, R, theta0, "Sxx (DDM)")
    circle_subplot(nrows, ncols, 17, x, y, SdispQ[:, 2], npts, R, theta0, "Syy (DDM)")
    circle_subplot(nrows, ncols, 18, x, y, SdispQ[:, 3], npts, R, theta0, "Sxy (DDM)")
    circle_subplot(nrows, ncols, 22, x, y, SresidualQ[:, 1], npts, R, theta0, "Sxx (residual)")
    circle_subplot(nrows, ncols, 23, x, y, SresidualQ[:, 2], npts, R, theta0, "Syy (residual)")
    circle_subplot(nrows, ncols, 24, x, y, SresidualQ[:, 3], npts, R, theta0, "Sxy (residual)")
    tight_layout()
    show()

    #! 9-panel plot for quadratic only
    figure(figsize=(15,15))
    nrows = 3
    ncols = 3
    # Analytic solutions
    circle_subplot(nrows, ncols, 1, x, y, Sanalytic[:, 1], npts, R, theta0, L"\sigma_{xx} \; \mathrm{(analytic)}")
    circle_subplot(nrows, ncols, 2, x, y, Sanalytic[:, 2], npts, R, theta0, L"\sigma_{yy} \; \mathrm{(analytic)}")
    circle_subplot(nrows, ncols, 3, x, y, Sanalytic[:, 3], npts, R, theta0, L"\sigma_{xy} \; \mathrm{(analytic)}")

    # BEM solutions
    circle_subplot(nrows, ncols, 4, x, y, SdispQ[:, 1], npts, R, theta0, L"\sigma_{xx} \; \mathrm{(BEM)}")
    circle_subplot(nrows, ncols, 5, x, y, SdispQ[:, 2], npts, R, theta0, L"\sigma_{yy} \; \mathrm{(BEM)}")
    circle_subplot(nrows, ncols, 6, x, y, SdispQ[:, 3], npts, R, theta0, L"\sigma_{xy} \; \mathrm{(BEM)}")
    circle_subplot(nrows, ncols, 7, x, y, SresidualQ[:, 1], npts, R, theta0, L"\sigma_{xx} \; \mathrm{(residual)}")
    circle_subplot(nrows, ncols, 8, x, y, SresidualQ[:, 2], npts, R, theta0, L"\sigma_{yy} \; \mathrm{(residual)}")
    circle_subplot(nrows, ncols, 9, x, y, SresidualQ[:, 3], npts, R, theta0, L"\sigma_{xy} \; \mathrm{(residual)}")
    tight_layout()
    show()
end
braziltest()