# Compact Hilbert Indices by Chris Hamilton. Technical Report CS-2006-07.
# 6059 University Ave., Halifax, Nova Scotia, B3H 1W5, Canada.
#
# In this document, the notation is:
# ∧ = & for and
# ∨ = | for or
# ⊻ = for xor
# ▷ = >> for shift right
# ◁ = << for shift left
# ⌈x⌉ = ceil(x)
# ⌊x⌋ = floor(x)

# hi_g(i) is the direction in which brgc changes.
# It is also tells us the axis along which the exit of one
# subcube meets the entrance of the next subcube.
hi_g(i) = trailing_set_bits(i)
# This is the definition of g, to compare against for testing.
hi_g_orig(i) = floor(typeof(i), log2(brgc(i) ⊻ brgc(i + 1)))

# e is the entry vertex of the ith sub-hypercube in a Gray code
# ordering of sub-hypercubes. If e is 0b011, then this is z-y-x
# within the subcube of 0 in the z, 1 in the y, 1 in the x.
# e = gc(2⌊(i-1)/2⌋) = gc(2*floor((i-1)/2)).
# domain is i=0 to i=2^n - 1.
function hi_e(i::Integer)
    i == zero(i) && return zero(i)
    brgc((i - one(i)) & ~one(i))
end

# This is verbatim from the article, for comparison.
function hi_e_orig(i::Base.BitInteger)
    i == zero(i) && return zero(i)
    brgc(2 * floor(typeof(i), (i - 1) / 2))
end

"""
f is the exit vertex of the ith sub-hypercube in a Gray code
ordering of the sub-hypercubes.
Corllary 2.7 on pg. 12 says:
    hi_f(i, n) = hi_e(i) ⊻ hi_d(i, n)
"""
hi_f(i, n) = hi_e(one(i)<<n - one(i) - i) ⊻ (one(i)<<(n-1))
# hi_f(i, n) = hi_e(i) ⊻ hi_d(i, n)
#hi_f(i, n) = hi_e(i) ⊻ (one(i)<<hi_d(i, n))

"""
page 12. directions
d(i) = 0 if i=0
    = g(i-1) mod n if i=0 mod 2
    = g(i) mod n if i=1 mod 2.
Domain is 0 <= i <= 2^n - 1.
"""
function hi_d(i, n)
    i == zero(i) && return zero(i)
    ((i & 1) == 0) && return hi_g(i - one(i)) % n
    hi_g(i) % n
end

# bitrotate is a left shift, so negate d+1.
# T_{(e,d)}(b), so read right-to-left.
# The paper means to bit rotate over the n bits that are in use,
# not all n bits. This is not the usual bitrotate!
# This had a +1 in the paper.
hi_T(b, d, e, n) = rotateright(b ⊻ e, d, n)

# The author's code differs with his paper. It doesn't add one.
# https://github.com/pdebuyl/libhilbert/blob/master/include/Hilbert/Algorithm.hpp
hi_T_inv(b, d, e, n) = rotateleft(b, d, n) ⊻ e

# Lemma 2.12, page 15.
# Is it -2 or -1?
# function hi_T_inv(b, d, e, n)
#     hi_T(b, n - d - one(b) - one(b), bitrotate(e, -(d + one(b)), n), n)
# end

"""
    ith_bit_of_indices(n, p, i)

Given `n` indices in `p`, take the `i`-th bit of each one
and place it so that the first vector's value is at the
0-th place, the second vector's value at the 1st place, and so-on.
`i` is zero-indexed.
"""
function ith_bit_of_indices(n, p, i)
    l = zero(eltype(p))
    for j = 0:(n - 1)
        b = (p[j + 1] & (one(eltype(p))<<i)) >>i
        l |= b << j
    end
    l
end

# hamilton version of getting the ith bit, zero-indexed i.
function get_location(p::Vector{T}, i) where {T}
    l = zero(T)
    for j = eachindex(p)
        if (p[j] & (one(T) << i)) != 0
            l |= (one(T) << (j - 1))
        end
    end
    l
end

"""
    set_indices_bits(p, l, m, i)

Given bits in `l`, set each index in the array `p` according
to the bit in `l`. Set the i-th bit of each index in `p`.
"""
function set_indices_bits!(p, l, m, i)
    for j in 0:(m - 1)
        v = (l & (1<<j)) >>j
        p[j + 1] = p[j + 1] | (v<<i)
    end
end

"""
Return `m` bits from the `i`-th set of `m` bits
in the integer `h`. `i` is zero-based.
"""
function bitrange(h, n, i)
    v = zero(h)
    for j in 0:(n - 1)
        v = v | (h & (1<<(i*n + j)))
    end
    v
end

function update1(l, t, w, n, e, d)
    e = l ⊻ (one(l) << d)
    d += one(d) + first_set_bit(t)
    while d >= n
        d -= n
    end
    if (w & one(w)) == zero(w)
        if d == zero(d)
            e ⊻= one(e) << (n - 1)
        else
            e ⊻= one(e) << (d - 1)
        end
    end
    e, d
end

function update2(l, t, w, n, e, d)
    e = l
    e ⊻= (one(e) << d)
    d += one(d) + first_set_bit(t)
    while d >= n
        d -= n
    end
    e, d
end

"""
Algorithm.hpp: _coordsToIndex
"""
function hilbert_index_paper!(::Type{T}, n, m, p, ds) where {T <: Integer}
    h = zero(T)  # hilbert index
    e = zero(T)  # entry point
    d = one(T)  # direction
    nmask = fbvn1s(T, n)
    for i = (m - 1):-1:0  # i is an index. Can be any type.
        ds[i + 1] = d
        l = T(ith_bit_of_indices(n, p, i))
        t = hi_T(l, d, e, n)
        w = t
        if i < m - 1
            w ⊻= (one(w) << (n - 1))
        end
        # Concatenate to the index
        h |= (w & nmask) << (i * n)
        e, d = update2(l, t, w, n, e, d)
    end
    brgc_inv(h)
end

"""
    hilbert_index(n, m, p)

Hilbert index for an `n`-dimensional vector `p`, with each
component of extent less than 2^m. Algorithm 1 of Hamilton and
Rau-Chaplin.
"""
function hilbert_index_paper(::Type{T}, n, m, p) where {T <: Integer}
    ds = zeros(Int, m)
    hilbert_index_paper!(T, n, m, p, ds)
end

function hilbert_index_inv_paper!(::Type{T}, n, m, h, p) where {T <: Integer}
    e = zero(T)
    d = one(T)
    l = zero(T)
    p .= zero(eltype(p))
    nmask = fbvn1s(T, n)
    for i = (m - 1):-1:0
        w = (h >> (i * n)) & nmask
        t = brgc(w)
        l = hi_T_inv(t, d, e, n)
        set_indices_bits!(p, l, n, i)
        e, d = update1(l, t, w, n, e, d)
    end
end

# Make d an argument because it determines the initial direction and maybe
# it isn't initialized correctly. Maybe d and e need different initial values.
function hilbert_index(n, m, p, d = zero(eltype(p)))
    h = zero(eltype(p))  # hilbert index
    e = zero(eltype(p))  # entry point
    for i = (m - 1):-1:0  # i is an index. Can be any type.
        l = ith_bit_of_indices(n, p, i)
        # @show l
        l = hi_T(l, d, e, n)  # n or m?
        w = brgc_inv(l)
        h = (h << n) | w
        e = e ⊻ rotateleft(hi_e(w), d + one(d), n)
        d = (d + hi_d(w, n) + one(d)) % n  # n or m for hi_d?
    end
    h
end

function hilbert_index_inv(n, m, h)
    e = zero(h)  # entry point
    d = zero(h)  # direction
    p = zeros(typeof(h), m)
    for i = (m - 1):-1:0  # i is an index. Can be any type.
        w = bitrange(h, n, i)
        l = brgc(w)
        l = hi_T_inv(l, d, e, n)
        set_indices_bits!(p, l, m, i)
        e = e ⊻ rotateleft(hi_e(w), d + one(d), n)
        d = (d + hi_d(w, n) + one(d)) % n  # n or m for hi_d?
    end
    p
end

"""
Calculates a bit-mask for excluding Gray code values when the level
is below the resolution of the dimension.

Bits free at iteration `i`. Vector of resolutions, in powers of two,
`m`. Dimensions `n` so that `length(m)==n`. `i` is the level down
in the Hilbert curve. `d` is the direction, called `hi_d` above.

Returns both the mask and the number of bits set in the mask.
"""
function extract_mask(m::Vector, n, d, i)
    T = UInt64
    mask = zero(T)
    b = 0
    jm = one(T)
    j = d
    while true
        if m[j + 1] > i
            mask |= jm
            b += 1
        end
        jm <<= one(T)
        if jm == zero(T)
            jm = one(T)
        end
        j += 1
        if j == n
            j = 0
        end
        if j == d
            break
        end
    end
    mask, b
end

function extract_mask_paper(m::Vector, n, d, i)
    T = UInt64
    mask = zero(T)
    b = 0
    for j = (n-1):-1:0
        mask <<= one(T)
        jn = (j + d) % n
        if m[jn + 1] > i
            mask |= one(T)
            b += 1
        end
    end
    mask, b
end

"""
From GrayCodeRank.hpp: compactIndex.
"""
function compact_index(ms::Vector, ds::Vector, n, m, h::T) where {T}
    hc = zero(T)
    hr = 0
    hcr = 0
    hm = one(T)
    hcm = one(T)
    for i = 0:(m-1)
        j = ds[i + 1]
        while true
            if ms[j + 1] > i
                if hr > 0
                    error("hr on next rack")
                end
                if ((h & hm) != 0)
                    if hcr == 0
                        hc |= hcm
                    else
                        error("should only be one rack")
                    end
                end
                hcm <<= 1
                if hcm == 0
                    hcm = one(T)
                    hcr += 1
                end
            end
            j += 1
            if j == n
                j = 0
            end
            hm <<= 1
            if hm == zero(T)
                hm = one(T)
                hr += 1
            end
            if j == ds[i + 1]
                break
            end
        end
    end
    hc
end

function coords_to_compact_index(::Type{T}, p::Vector{A}, ms::Vector, n) where {A,T}
    m = maximum(ms)
    M = sum(ms)
    mn = m * n
    ds = zeros(Int, m)
    h = hilbert_index_paper!(T, n, m, p, ds)
    compact_index(ms, ds, n, m, h)
end

function compact_index_to_coords!(p::Vector{A}, ms, n, hc::T) where {A, T}
    m = maximum(ms)
    M = sum(ms)
    bit_cnt = 8 * sizeof(T)

    e = zero(T)
    d = one(T)
    l = zero(T)
    p .= zero(A)
    # work from most significant bit to least significant bit
    for i = (m - 1):-1:0
        mask, b = extract_mask(ms, n, d, i)
        # rotateright(val, shift_cnt, total_bits)
        ptrn = rotateright(e, d, n)

        # Get the Hilbert index bits.
        M -= b
        # b bits from hc at index M, into r
        r = get_bits(hc, b, M)

        t, w = brgc_rank_inv(mask, ptrn, r, n, b)

        l = hi_T_inv(t, d, e, n)
        set_indices_bits!(p, l, n, i)
        e, d = update1(l, t, w, n, e, d)
    end
end

"""
    SpaceGray(b, n)
    SpaceGray(::Type{T}, b, n)

This is an n-dimensional Hilbert curve where all `n` dimensions
must have `b` bits in size. It was described in the same paper
and examples as the `Compact` algorithm.
"""
struct SpaceGray{T} <: HilbertAlgorithm{T}
    b::Int
    n::Int
end

axis_type(gg::SpaceGray) = large_enough_unsigned(gg.b)

function SpaceGray(b, n)
    ttype = large_enough_unsigned(b * n)
    SpaceGray{ttype}(b, n)
end

function SpaceGray(::Type{T}, b, n) where {T}
    SpaceGray{T}(b, n)
end

function encode_hilbert_zero(g::SpaceGray{T}, X::Vector)::T where {T}
    hilbert_index_paper(T, g.n, g.b, X)
end

function decode_hilbert_zero!(g::SpaceGray{T}, X::Vector, h::T) where {T}
    hilbert_index_inv_paper!(T, g.n, g.b, h, X)
end

"""
    Compact(ms::Vector{Int})
    Compact(::Type{T}, ms::Vector{Int})

This algorithm is n-dimensional and permits dimensions to use different
numbers of bits, specified in the `ms` vector. The type `T` is an
optional data type for the Hilbert index. It should be greater than or
equal to the sum of the bits.

This algorithm comes from three sources:

* A technical report, "Compact Hilbert Indices" by Chris Hamilton.
  Technical Report CS-2006-07. 6059 University Ave., Halifax, Nova
  Scotia, B3H 1W5, Canada. This report is informative but has many
  errors.

* A paper by Hamilton and Rau-Chaplin, "Compact Hilbert Indices for
  Multi-Dimensional Data," 2007. Nice paper. Also wrong.

* The [libhilbert source code](https://github.com/pdebuyl/libhilbert)
  is a copy of Hamilton's work and has many corrections. This, ultimately,
  lead to the working code.
"""
struct Compact{T} <: HilbertAlgorithm{T}
    ms::Vector{Int}
    n::Int
end

axis_type(gg::Compact) = large_enough_unsigned(maximum(gg.ms))

function Compact(ms::Vector{Int})
    n = length(ms)
    b = maximum(ms)
    ttype = large_enough_unsigned(b * n)
    Compact{ttype}(ms, n)
end

function Compact(::Type{T}, ms::Vector{Int}) where {T}
    Compact{T}(ms, length(ms))
end

function encode_hilbert_zero(g::Compact{T}, X::Vector)::T where {T}
    coords_to_compact_index(index_type(g), X, g.ms, g.n)
end

function decode_hilbert_zero!(g::Compact{T}, X::Vector, h::T) where {T}
    compact_index_to_coords!(X, g.ms, g.n, h)
end


using DifferentialEquations
using Plots

# Parameters
const ib = 2e-5; # g⋅m⁻²⋅day⁻¹ immigration rate of bream
const ip = 2e-5; # g⋅m⁻²⋅day⁻¹ immigration rate of pike
const r = 7.5e-3; # day⁻¹ maximum growth rate of bream
const H₁ = 0.5; # half saturation constant
const H₂ = 10; # % half saturation constant
const H₃ = 20; # g⋅m⁻² half saturation constant
const H₄ = 15; # g⋅m⁻² half saturation constant
const cb = 7.5e-5; # m⁻²⋅g⁻¹⋅day⁻¹ intraspecific competition constant for bream
const cp = 2.75e-4; # m⁻²⋅g⁻¹⋅day⁻¹ intraspecific competition constant for pike
const prmax = 5e-2; # day⁻¹ maximum predation rate of pike
const ce = 0.1; # pike food conversion efficiency to growth
const mp = 2.25e-3; # day⁻¹ mortality rate of pike
const K = 100; # % maximum vetetation coverage

function scheffer!(du,u,p,t)
    B, P = u # g⋅m⁻² bream/pike density
    ib, r, nutr, H₁, H₂, H₃, H₄, cb, prmax, ip, ce, mp, cp, K = p

    V = K*(H₃^2/(H₃^2 + B^2)); # % of lake covered vetetation
    FR = B^2/(B^2 + H₄^2); # fuctional response of pike

    du[1] = dB = ib + r*(nutr/(nutr + H₁))*B - cb*B^2 - prmax*FR*P
    du[2] = dP = ip + ce*prmax*FR*P*(V/(V + H₂)) - mp*P - cp*P^2
end

function run_ss_solver(;
             nutr = 0.5, #nutrient level
             u0 = [40.0, 2.6], #Initial populations
             p = [ib, r, nutr, H₁, H₂, H₃, H₄, cb, prmax, ip, ce, mp, cp, K])
    prob = SteadyStateProblem(scheffer!,u0,p);
    solve(prob)
end

equilibria = Set();
for i in range(0.0, stop=100.0, length=10)
    for j in range(0.0, stop=10.0, length=10)
        for k in [0.5, 1.5, 2.5];
            eq = run_ss_solver(u0=[i,j], nutr=k);
            push!(equilibria, (eq[1], eq[2]));
        end
    end
end

function dB(B, P)
    FR = B^2/(B^2 + H₄^2); # fuctional response of pike

    ib + r*(nutr/(nutr + H₁))*B - cb*B^2 - prmax*FR*P
end

function dP(B, P)
    V = K*(H₃^2/(H₃^2 + B^2)); # % of lake covered vetetation
    FR = B^2/(B^2 + H₄^2); # fuctional response of pike

    ip + ce*prmax*FR*P*(V/(V + H₂)) - mp*P - cp*P^2
end

bream = range(0.0, stop=100.0, length=100)
pike = range(0.0, stop=10.0, length=100)

nutr = 2.5;
contour(bream, pike, dB, levels=[0], color=:blue, cbar=false, legend=true, label="bream'=0, nutr=2.5");
nutr = 1.5;
contour!(bream, pike, dB, levels=[0], color=:red, cbar=false, label="bream'=0, nutr=1.5");
nutr = 0.5;
contour!(bream, pike, dB, levels=[0], color=:green, cbar=false, label="bream'=0, nutr=0.5");
contour!(bream, pike, dP, levels=[0], color=:black, cbar=false, label="pike'=0");

eq = collect(equilibria);
scatter!(first.(eq), last.(eq), label="equilibria");

gui()


# Script that executes the Kabsch algorithm to align two paired sets of points
# by minimising the RMSD between them.
# Takes 2 .xyz files as input
# Program executed as
# 'julia kabsch_optimiser.jl geomP.xyz geomQ.xyz'
# Darya Shchepanovska
# 15th July 2019

using DelimitedFiles

using LinearAlgebra
using Statistics

function xyz2matrix(input_xyz)
    # Imports .xyz file format as data frame, and converts it to an N x 3 Matrix
    raw_xyz=readdlm(input_xyz)
    natoms=size(raw_xyz)[1]
    just_coords=Array{Float64}(raw_xyz[2:natoms,2:4])
    return just_coords
end

function translate_to_centroid(coord_matrix)
    # Normalises the molecular coordinates by centering them.
    center = [mean(coord_matrix[:,1]);mean(coord_matrix[:,2]);mean(coord_matrix[:,3])]
    centroid = transpose(center)
    translated_geom = broadcast(-,coord_matrix,centroid)
    return translated_geom
end

function cross_covariance_matrix(Pmatrix,Qmatrix)
    # Cross covariance matrix gives measure of variability between two matrices
    CCmatrix = transpose(Pmatrix) * Qmatrix
    return CCmatrix
end

function optimal_rotation_matrix(CCmatrix)
    # Returns 3x3 matrix that can be applied to P to get Q
    ORmatrix = sqrt(transpose(CCmatrix)*CCmatrix)*inv(CCmatrix)
    return ORmatrix
end

Pgeom = xyz2matrix(ARGS[1])
Qgeom = xyz2matrix(ARGS[2])

normalisedP = (translate_to_centroid(Pgeom))
normalisedQ = (translate_to_centroid(Qgeom))

xcov = cross_covariance_matrix(normalisedP,normalisedQ)
orot = optimal_rotation_matrix(xcov)

num_atoms = size(normalisedP)[1]
rotated = zeros(Float64,num_atoms,3)

for i=1:num_atoms
    rotated[i,:] = orot*normalisedP[i,:]
end

# display(rotated)

RMSD_value = norm(rotated-normalisedQ)

println("RMSD between these two structures is ",RMSD_value)


"""
    eig2length(λ, D)

Convert Laplace eigenvalue to length scale.
"""
eig2length(λ, D) = λ > 0 ? π * √(D / λ) : Inf

"""
    length2eig(length, D)

Convert length scale to Laplace eigenvalue.
"""
length2eig(length, D) = D * (π / length)^2


function isHappy(s)::Bool
    a::Int = s isa Number ? s : (try
                                 parse(Int64, s)
                                 catch e
                                 throw(error("Wrong type $e"))
                                 end)
    innerHappy(a, Set([]))
end

function eachBit(n::Int)
    result = []
    while n != 0
        let left = n % 10
            push!(result, left)
            n = Int((n-left)/10)
        end
    end
    result
end

function innerHappy(n,se)
    n in se && return false
    next = sum(eachBit(n) .^ 2)
    next == 1 ? (return true) : (se = union(se, n))
    innerHappy(next,se)
end


module BayesianRegression

#=

### Model and prior

Assume a linear model
$$
f(\mathbf{x}) = \sum_{i=1}^M c_i \phi_i(\mathbf{x}) = \mathbf{c}^T \boldsymbol{\phi}(\mathbf{x} )
$$
where the elements of the $M$-dimensional coefficient vector have independent priors $p(c_i)=\mathcal{N}(c_i | 0,\sigma_i^2)$. Suppose $N$ noisy input-target pairs are available, $\{\mathbf{x}_i, y_i\}_{i=1}^N$, such that
$$
y_i = f(\mathbf{x}_i) + \epsilon_i,
$$
with noise prior $p(\epsilon_i) = \mathcal{N}(\epsilon_i | 0,\sigma_\epsilon^2)$.

### Posterior

For fixed values of $\{\sigma_i^2\}_{i=1}^M$ and $\sigma_\epsilon^2$, the posterior mean used for predictions is
$$
\boldsymbol{\mu}_c = \sigma_\epsilon^{-2} \boldsymbol{\Sigma}_c \boldsymbol{\Phi}^T \mathbf{y}
$$
with associated covariance
$$
\begin{aligned}
\boldsymbol{\Sigma}_c
&= \left[ \sigma_\epsilon^{-2} \boldsymbol{\Phi}^T \boldsymbol{\Phi} + \boldsymbol{\Sigma}_0^{-1} \right]^{-1} \\
&= \boldsymbol{\Sigma}_0 - \boldsymbol{\Sigma}_0 \boldsymbol{\Phi}^T \boldsymbol{\Sigma}_y^{-1} \boldsymbol{\Phi} \boldsymbol{\Sigma}_0
\end{aligned} ,
$$
where $\boldsymbol{\Phi}$ is the design matrix
$$
\boldsymbol{\Phi} = \left[
\begin{matrix}
\text{---} \boldsymbol{\phi}(\mathbf{x}_1) \text{---} \\
\text{---} \boldsymbol{\phi}(\mathbf{x}_2) \text{---} \\
\vdots \\
\text{---} \boldsymbol{\phi}(\mathbf{x}_3) \text{---}
\end{matrix}
\right] ,
$$
$\boldsymbol{\Sigma}_0$ is diagonal with $[\boldsymbol{\Sigma}_0]_{ii} = \sigma_i^2$, and $\boldsymbol{\Sigma}_y$ is defined below.

### Likelihood and log likelihood

With flat hyperpriors, the marginal likelihood---the "evidence" for the model---is
$$
p(\mathbf{y} | \boldsymbol{\Sigma}_0, \sigma_\epsilon^2) = \mathcal{N}(\mathbf{y} | \mathbf{0}, \boldsymbol{\Sigma}_y)
$$
with
$$
\begin{aligned}
\boldsymbol{\Sigma}_y
&= \boldsymbol{\Phi} \boldsymbol{\Sigma}_0 \boldsymbol{\Phi}^T + \sigma_\epsilon^2 \mathbf{I} \\
&= \left[ \sigma_\epsilon^{-2} \mathbf{I} - \sigma_\epsilon^{-2} \boldsymbol{\Phi} \boldsymbol{\Sigma}_c \boldsymbol{\Phi}^T \sigma_\epsilon^{-2}  \right]^{-1}
\end{aligned}
$$
When maximizing this quantity, one typically works with its log,
$$
\begin{aligned}
\mathrm{\mathcal{L}}
&= \ln p(\mathbf{y} | \boldsymbol{\Sigma}_0, \sigma_\epsilon^2) \\
&= -\frac{1}{2} \mathbf{y}^T \boldsymbol{\Sigma}_y^{-1} \mathbf{y}
    - \frac{1}{2} \ln \left| \boldsymbol{\Sigma}_y \right|
    - \frac{N}{2} \ln (2\pi) \\
&= -\frac{1}{2} \sigma_\epsilon^{-2} \mathbf{y}^T \left[  \mathbf{y} - \boldsymbol{\Phi} \boldsymbol{\mu}_c \right]
    + \frac{1}{2} \ln \left| \boldsymbol{\Sigma}_c \right|
    - \frac{1}{2} \ln \left| \boldsymbol{\Sigma}_0 \right|
    - \frac{N}{2} \ln (\sigma_\epsilon^2)
    - \frac{N}{2} \ln (2\pi) \\
\end{aligned} .
$$
The relevant derivatives with respect to the hyperparameters are
$$
\begin{aligned}
\frac{\partial \mathcal{L}}{\partial \sigma_i^2}
&= \frac{1}{2} [\boldsymbol{\Phi}^T \boldsymbol{\Sigma}_y^{-1} \mathbf{y}]_i^2
    - \frac{1}{2} [\boldsymbol{\Phi}^T \boldsymbol{\Sigma}_y^{-1} \boldsymbol{\Phi}]_{ii} \\
&= \frac{1}{2} \sigma_i^{-4} \left[ \mu_i^2 + [\boldsymbol{\Sigma}_c]_{ii} - \sigma_i^2 \right]
\end{aligned}
$$
and
$$
\begin{aligned}
\frac{\partial \mathcal{L}}{\partial \sigma_\epsilon^2}
& = \frac{1}{2} \left\| \boldsymbol{\Sigma}_y^{-1} \mathbf{y} \right\|^2 - \frac{1}{2} \mathrm{tr}\left(\boldsymbol{\Sigma}_y^{-1}\right) \\
&= \frac{1}{2} \sigma_\epsilon^{-4} \left[
    \left\| \mathbf{y} - \boldsymbol{\Phi} \boldsymbol{\mu}_c \right\|^2
    + \mathrm{tr}\left( \boldsymbol{\Phi} \boldsymbol{\Sigma}_c \boldsymbol{\Phi}^T \right)
    - N \sigma_\epsilon^2
\right]
\end{aligned} .
$$

### References

* M. E. Tipping, "Sparse Bayesian learning and the relevance vector machine," Journal of Machine Learning Research 1, 211 (2001)
* C. K. Williams and C. E. Rasmussen, _Gaussian Processes for Machine Learning_, MIT Press (2006)

=#

using LinearAlgebra
using Optim

function solve(
    y::Vector{<:AbstractFloat},
    X::Matrix{<:AbstractFloat},
    var_c::AbstractFloat,
    var_e::AbstractFloat,
)
    return solve(y, X, var_c*ones(size(X,2)), var_e)
end

function solve(
    y::Vector{<:AbstractFloat},
    X::Matrix{<:AbstractFloat},
    var_c::Vector{<:AbstractFloat},
    var_e::AbstractFloat,
)
    M = size(X,2)
    XTX = X'*X
    Σ_c = Array{Float64}(undef,M,M)
    BLAS.blascopy!(length(Σ_c), XTX, stride(XTX,1), Σ_c, stride(Σ_c,1))
    BLAS.scal!(1.0/var_e, Σ_c)
    for i in 1:M; Σ_c[i,i] += 1.0/var_c[i]; end
    C = cholesky!(Symmetric(Σ_c))
    return 1.0/var_e*(C \ (X'*y))
end

function log_marginal_likelihood_overdetermined!(
    lml::AbstractFloat,
    grad::Vector{<:AbstractFloat},
    X::Matrix{<:AbstractFloat},
    y::Vector{<:AbstractFloat},
    var_c::AbstractFloat,
    var_e::AbstractFloat,
    XTX::Matrix{<:AbstractFloat},
)
    var_c_vec = var_c*ones(size(X,2))
    grad_vec = zeros(size(X,2)+1)
    lml = log_marginal_likelihood_overdetermined!(lml, grad_vec, X, y, var_c_vec, var_e, XTX)
    grad[1] = sum(grad_vec[1:end-1])
    grad[2] = grad_vec[end]
    return lml
end

function log_marginal_likelihood_overdetermined!(
    lml::AbstractFloat,
    grad::Vector{<:AbstractFloat},
    X::Matrix{<:AbstractFloat},
    y::Vector{<:AbstractFloat},
    var_c::Vector{<:AbstractFloat},
    var_e::AbstractFloat,
    XTX::Matrix{<:AbstractFloat},
)
    N = size(X,1)
    M = size(X,2)
    Σ_c = Array{Float64}(undef,M,M)
    BLAS.blascopy!(length(Σ_c), XTX, stride(XTX,1), Σ_c, stride(Σ_c,1))
    BLAS.scal!(1.0/var_e, Σ_c)
    for i in 1:M; Σ_c[i,i] += 1.0/var_c[i]; end
    C = cholesky!(Symmetric(Σ_c))
    Σ_c = C \ I(M)
    μ_c = 1.0/var_e*(C \ (X'*y))
    lml = - 0.5*logdet(C) - 0.5*sum(log.(var_c)) - 0.5*N*log(var_e) - 0.5*N*log(2*π)
    lml -= 0.5/var_e*y'*(y-X*μ_c)
    grad[1:M] .= 0.5*(μ_c.^2 .+ diag(Σ_c) .- var_c)./var_c.^2
    grad[M+1] = 0.5/var_e^2*(sum((y-X*μ_c).^2) + dot(XTX,Σ_c) - N*var_e)
    return lml
end

function log_marginal_likelihood_underdetermined!(
    lml::AbstractFloat,
    grad::Vector{<:AbstractFloat},
    X::Matrix{<:AbstractFloat},
    y::Vector{<:AbstractFloat},
    var_c::AbstractFloat,
    var_e::AbstractFloat,
)
    var_c_vec = var_c*ones(size(X,2))
    grad_vec = zeros(size(X,2)+1)
    lml = log_marginal_likelihood_underdetermined!(lml, grad_vec, X, y, var_c_vec, var_e)
    grad[1] = sum(grad_vec[1:end-1])
    grad[2] = grad_vec[end]
    return lml
end

function log_marginal_likelihood_underdetermined!(
    lml::AbstractFloat,
    grad::Vector{<:AbstractFloat},
    X::Matrix{<:AbstractFloat},
    y::Vector{<:AbstractFloat},
    var_c::Vector{<:AbstractFloat},
    var_e::AbstractFloat,
)
    N = size(X,1)
    M = size(X,2)
    Σ_y = X*Diagonal(var_c)*X'
    for i=1:N; Σ_y[i,i] += var_e; end
    C = cholesky!(Symmetric(Σ_y))
    invΣy_y = C \ y
    lml = -0.5*y'*invΣy_y - 0.5*logdet(C) - 0.5*N*log(2*π)
    grad[1:M] .= 0.5*(X'*invΣy_y).^2
    W = C \ X
    @views for i=1:M; grad[i] -= 0.5*dot(X[:,i], W[:,i]); end
    grad[M+1] = 0.5*dot(invΣy_y,invΣy_y) - 0.5*tr(C\I(N))
    return lml
end

function bayesian_fit(
    y::Vector{<:AbstractFloat},
    X::Matrix{<:AbstractFloat};
    variance_floor::AbstractFloat=1e-8,
    verbose::Bool=false,
)
    if size(X,1) > size(X,2)
        XTX = X'*X  # advantageous to precompute for overdetermined case
    end

    function fg!(f, g, x)
        var_c = variance_floor + x[1]*x[1]
        var_e = variance_floor + x[2]*x[2]
        if size(X,1) >= size(X,2)
            f = log_marginal_likelihood_overdetermined!(f, g, X, y, var_c, var_e, XTX)
        else
            f = log_marginal_likelihood_underdetermined!(f, g, X, y, var_c, var_e)
        end
        if f != nothing
            f = -f
        end
        if g != nothing
            g .*= -2*x
        end
        return f
    end

    res = optimize(Optim.only_fg!(fg!),
                   ones(2),
                   Optim.LBFGS(),
                   Optim.Options(x_tol=1e-4, g_tol=0.0, show_trace=verbose))
    verbose && println(res)

    lml = -Optim.minimum(res)
    var_c, var_e = Optim.minimizer(res)
    var_c = variance_floor + var_c*var_c
    var_e = variance_floor + var_e*var_e

    return solve(y, X, var_c, var_e), var_c, var_e, lml
end

function ard_fit(
    y::Vector{<:AbstractFloat},
    X::Matrix{<:AbstractFloat},
    variance_floor::AbstractFloat=1e-8;
    verbose::Bool=false,
)
    if size(X,1) >= size(X,2)
        XTX = X'*X  # advantageous to precompute for overdetermined case
    end

    function fg!(f, g, x)
        var_c = variance_floor .+ x[1:end-1].*x[1:end-1]
        var_e = variance_floor + x[end]*x[end]
        if size(X,1) >= size(X,2)
            f = log_marginal_likelihood_overdetermined!(f, g, X, y, var_c, var_e, XTX)
        else
            f = log_marginal_likelihood_underdetermined!(f, g, X, y, var_c, var_e)
        end
        if f != nothing
            f = -f
        end
        if g != nothing
            g .*= -2*x
        end
        return f
    end

    res = optimize(Optim.only_fg!(fg!),
                   ones(size(X,2)+1),
                   Optim.LBFGS(),
                   Optim.Options(x_tol=1e-4, g_tol=0.0, show_trace=verbose))
    verbose && println(res)

    lml = -Optim.minimum(res)
    x = Optim.minimizer(res)
    var_c = variance_floor .+ x[1:end-1].*x[1:end-1]
    var_e = variance_floor + x[end]*x[end]

    mask = var_c .> 10*variance_floor
    var_c[.~mask] .= 0
    c_mask = solve(y, X[:,mask], var_c[mask], var_e)
    c = zeros(length(var_c))
    c[mask] .= c_mask

    return c, var_c, var_e, lml, mask
end

end


#############################################################################
#############################################################################
#
# This file implements polynomial division
#
#############################################################################
#############################################################################

"""  Modular algorithm.
f divide by g

f = q*g + r

p is a prime
"""
function divide(num::PolynomialModP, den::PolynomialModP)
    @assert num.prime == den.prime "Both polynomials must share the same prime divisor"
    f, g = mod(num.terms,num.prime), mod(den.terms,den.prime)
    degree(f) < degree(num.terms) && return nothing
    iszero(g) && throw(DivideError())
    q = Polynomial()
    prev_degree = degree(f)
    while degree(f) ≥ degree(g)
        h = Polynomial( (leading(f) ÷ leading(g))(num.prime) )  #syzergy
        f = mod((f - h*g), num.prime)
        q = mod((q + h), num.prime)
        prev_degree == degree(f) && break
        prev_degree = degree(f)
    end
    @assert iszero( mod((num.terms  - (q*g + f)),num.prime))
    return q, f
end

"""
The quotient from polynomial division. Returns a function of an integer.
"""

÷(num::PolynomialModP, den::PolynomialModP)  = first(divide(num,den))

"""
The remainder from polynomial division. Returns a function of an integer.
"""

rem(num::PolynomialModP, den::PolynomialModP)  = last(divide(num,den))


###
  Name     : c7_20_JuMP01.jl
  Book     : Hands-on Data Science with Anaconda )
  Publisher: Packt Publishing Ltd.
  Author   : Yuxing Yan and James Yan
  Date     : 3/15/2018
  email    : yany@canisius.edu
             paulyxy@hotmail.com
###

using JuMP
using ECOS
m= Model(solver =ECOSSolver())
@variable(m, 0 <= x <= 2 )
@variable(m, 0 <= y <= 30 )
@setObjective(m, Max, 5x + 3*y )
@addConstraint(m, 1x + 5y <= 3.0 )
print(m)
status = solve(m)
println("Objective value: ", getObjectiveValue(m))
println("x = ", getValue(x))
println("y = ", getValue(y))

#https://jump.readthedocs.io/en/release-0.2/jump.html


#' ---
#' title: Finance
#' ---

#' **Originally Contributed by**: Arpit Bhatia

#' Optimization models play an increasingly important role in financial decisions.
#' Many computational finance problems can be solved efficiently using modern optimization techniques.
#' In this tutorial we will discuss 3 such examples.
#' The first two are taken from the book Optimization Methods in Finance[[1]](#c1),
#' while the third is the famous Markowitz Portfolio Optimization problem with data from
#' lecture notes from a course taught at Georgia Tech[[2]](#c2).

using JuMP
using GLPK

#' ## Short Term Financing

#' Corporations routinely face the problem of financing short term cash commitments such as the following:

#' **Month**        |Jan    |Feb    |Mar    |Apr    |May    |Jun
#' :-----:          |:-----:|:-----:|:-----:|:-----:|:-----:|:-----:
#' **Net Cash Flow**|-150   |-100   |200    |-200   |50     |300

#' Net cash flow requirements are given in thousands of dollars. The company has the following sources of funds:

#' - A line of credit of up to $100K at an interest rate of 1% per month,
#' - In any one of the first three months, it can issue 90-day commercial paper bearing a total interest of 2% for the 3-month period,
#' - Excess funds can be invested at an interest rate of 0.3% per month.

#' Our task is to find out the most economical way to use these 3 sources such that
#' we end up with the most amount of money at the end of June.
#' We model this problem in the following manner:

#' We will use the following decision variables:
#' - the amount $u_{i}$ drawn from the line of credit in month $i$
#' - the amount $v_{i}$ of commercial paper issued in month $i$
#' - the excess funds $w_{i}$ in month $i$

#' Here we have three types of constraints:
#' 1. for every month, cash inflow = cash outflow for each month
#' 2. upper bounds on $u_{i}$
#' 3. nonnegativity of the decision variables $u_{i}$, $v_{i}$ and $w_{i}$.

#' Our objective will be to simply maximimse the company's wealth in June, which say we represent with the variable $m$.

financing = Model(GLPK.Optimizer)

@variables(financing, begin
    0 <= u[1:5] <= 100
    0 <= v[1:3]
    0 <= w[1:5]
    m
end)

@objective(financing, Max, m) # Money at the end of June

@constraints(financing, begin
    u[1] + v[1] - w[1] == 150 # January
    u[2] + v[2] - w[2] - 1.01u[1] + 1.003w[1] == 100 # February
    u[3] + v[3] - w[3] - 1.01u[2] + 1.003w[2] == -200 # March
    u[4] - w[4] - 1.02v[1] - 1.01u[3] + 1.003w[3] == 200 # April
    u[5] - w[5] - 1.02v[2] - 1.01u[4] + 1.003w[4] == -50 # May
    -m - 1.02v[3] - 1.01u[5] + 1.003w[5] == -300 # June
end)

optimize!(financing)
@show objective_value(financing);

#' ## Combinatorial Auctions
#' In many auctions, the value that a bidder has for a set of items may not be
#' the sum of the values that he has for individual items.
#' It may be more or it may be less.
#' Examples are equity trading, electricity markets, pollution right auctions and auctions for airport landing slots.
#' To take this into account, combinatorial auctions allow the bidders to submit bids on combinations of items.

#' Let $M=\{1,2, \ldots, m\}$ be the set of items that the auctioneer has to sell.
#' A bid is a pair $B_{j}=\left(S_{j}, p_{j}\right)$ where $S_{j} \subseteq M$ is a nonempty set of items and
#' $p_{j}$ is the price offer for this set.
#' Suppose that the auctioneer has received $n$ bids $B_{1}, B_{2}, \ldots, B_{n} .$
#' The goal of this problem is to help an auctioneer determine the winners in order to maximize his revenue.

#' We model this problem by taking a decision variable $y_{j}$ for every bid.
#' We add a constraint that each item $i$ is sold at most once. This gives us the following model:

#' $$
#' \begin{align*}
#' \max && \sum_{i=1}^{n} p_{j} y_{j} \\
#' \text { s.t. }  && \sum_{j : i \in S_{j}} y_{j} \leq 1 && \forall i=\{1,2 \ldots m\} \\
#' && y_{j} \in\{0,1\} && \forall j \in\{1,2 \ldots n\}
#' \end{align*}
#' $$

bid_values = [6 3 12 12 8 16]
bid_items = [[1], [2], [3 4], [1 3], [2 4], [1 3 4]]

auction = Model(GLPK.Optimizer)
@variable(auction, y[1:6], Bin)
@objective(auction, Max, sum(y' .* bid_values))
for i in 1:6
    @constraint(auction, sum(y[j] for j in 1:6 if i in bid_items[j]) <= 1)
end

optimize!(auction)

@show objective_value(auction);
@show value.(y);

#' ## Portfolio Optimization
#' Suppose we are considering investing 1000 dollars in three non-dividend paying stocks,
#' IBM (IBM), Walmart (WMT), and Southern Electric (SEHI), for a onemonth period.
#' This means we will use the money to buy shares of the three stocks at the current market prices,
#' hold these for one month, and sell the shares off at the prevailing market prices at the end of the month.
#' As a rational investor, we hope to make some profit out of this endeavor, i.e.,
#' the return on our investment should be positive.
#' Suppose we bought a stock at $p$ dollars per share in the beginning of the month,
#' and sold it off at $s$ dollars per share at the end of the month.
#' Then the one-month return on a share of the stock is $ \frac{s-p}{p} $.

#' Since the stock prices are quite uncertain, so is the end-of-month return on our investment.
#' Our goal is to invest in such a way that the expected end-of-month return is at least $50.00 or 5%.
#' Furthermore, we want to make sure that the “risk” of not achieving our desired return is minimum.

#' Note that we are solving the problem under the following assumptions:
#' 1. We can trade any continuum of shares.
#' 2. No short-selling is allowed.
#' 3. There are no transaction costs.

#' We model this problem by taking decision variables $x_{i}, i=1,2,3,$ denoting the dollars invested in each of the 3 stocks.
#' Let us denote by $\tilde{r}_{i}$ the random variable corresponding to the monthly return
#' (increase in the stock price) per dollar for stock $i .$

#' Then, the return (or profit) on $x_{i}$ dollars invested in stock $i$ is $\tilde{r}_{i} x_{i},$ and
#' the total (random) return on our investment is $\sum_{i=1}^{3} \tilde{r}_{i} x_{i} .$
#' The expected return on our investment is then
#' $\mathbb{E}\left[\sum_{i=1}^{3} \tilde{r}_{i} x_{i}\right]=\sum_{i=1}^{3} \overline{r}_{i} x_{i},$
#' where $\overline{r}_{i}$ is the expected value of the $\tilde{r}_{i} .$

#' Now we need to quantify the notion of “risk” in our investment.
#' Markowitz, in his Nobel prize winning work, showed that a rational investor’s notion of minimizing risk
#' can be closely approximated by minimizing the variance of the return of the investment portfolio.
#' This variance is given by:

#' $$
#' \operatorname{Var}\left[\sum_{i=1}^{3} \tilde{r}_{i} x_{i}\right] = \sum_{i=1}^{3} \sum_{j=1}^{3} x_{i} x_{j} \sigma_{i j}
#' $$

#' where $\sigma_{i j}$ is the covariance of the return of stock $i$ with stock $j$.

#' Note that the right hand side of the equation is the most reduced form of the expression and
#' we have not shown the intermediate steps involved in getting to this form.
#' We can also write this equation as:

#' $$
#' \operatorname{Var}\left[\sum_{i=1}^{3} \tilde{r}_{i} x_{i}\right] =x^{T} Q x
#' $$

#' Where $Q$ is the covariance matrix for the random vector $\tilde{r}$.

#' Finally, we can write the model as:

#' $$
#' \begin{align*}
#' \min x^{T} Q x \\
#' \text { s.t. } \sum_{i=1}^{3} x_{i} \leq 1000.00 \\
#' \overline{r}^{T} x \geq 50.00 \\
#' x \geq 0
#' \end{align*}
#' $$

#' After that long discussion, lets now use JuMP to solve the portfolio optimization problem for the data given below.

#' | Month        |  IBM     |  WMT    |  SEHI  |
#' |--------------|----------|---------|--------|
#' | November-00  |  93.043  |  51.826 |  1.063 |
#' | December-00  |  84.585  |  52.823 |  0.938 |
#' | January-01   |  111.453 |  56.477 |  1.000 |
#' | February-01  |  99.525  |  49.805 |  0.938 |
#' | March-01     |  95.819  |  50.287 |  1.438 |
#' | April-01     |  114.708 |  51.521 |  1.700 |
#' | May-01       |  111.515 |  51.531 |  2.540 |
#' | June-01      |  113.211 |  48.664 |  2.390 |
#' | July-01      |  104.942 |  55.744 |  3.120 |
#' | August-01    |  99.827  |  47.916 |  2.980 |
#' | September-01 |  91.607  |  49.438 |  1.900 |
#' | October-01   |  107.937 |  51.336 |  1.750 |
#' | November-01  |  115.590 |  55.081 |  1.800 |

using Statistics # Useful for calculations
using Ipopt      # Ipopt since our objective is quadratic

stock_data = [
93.043 51.826 1.063;
84.585 52.823 0.938;
111.453 56.477 1.000;
99.525 49.805 0.938;
95.819 50.287 1.438;
114.708 51.521 1.700;
111.515 51.531 2.540;
113.211 48.664 2.390;
104.942 55.744 3.120;
99.827 47.916 2.980;
91.607 49.438 1.900;
107.937 51.336 1.750;
115.590 55.081 1.800;
]

# Calculating stock returns

stock_returns = Array{Float64}(undef, 12, 3)

for i in 1:12
    stock_returns[i, :] = (stock_data[i + 1, :] .- stock_data[i, :]) ./ stock_data[i, :]
end

# Calculating the expected value of monthly return

r = [Statistics.mean(stock_returns[:,1]) Statistics.mean(stock_returns[:,2]) Statistics.mean(stock_returns[:,3])]'

# Calculating the covariance matrix Q

Q = Statistics.cov(stock_returns)

#+

# JuMP Model

portfolio = Model(optimizer_with_attributes(Ipopt.Optimizer, "print_level" => 0))
@variable(portfolio, x[1:3] >= 0)
@objective(portfolio, Min, x' * Q * x)
@constraint(portfolio, sum(x) <= 1000)
@constraint(portfolio, sum(r .* x) >= 50)

optimize!(portfolio)

@show objective_value(portfolio);
@show value.(x);

#' ### References
#'
#' 1. Cornuéjols, G., Peña, J., & Tütüncü, R. (2018). Optimization Methods in Finance. Cambridge: Cambridge University Press. doi:10.1017/9781107297340
#'
#' 2. Ahmed, S. ISyE6669: Deterministic Optimization. Available at: https://www2.isye.gatech.edu/~sahmed/isye6669/


#=
The number, 1406357289, is a 0 to 9 pandigital number because it is made up of each of the digits 0 to 9 in some order, but it also has a rather interesting sub-string divisibility property.

Let d1 be the 1st digit, d2 be the 2nd digit, and so on. In this way, we note the following:

d2d3d4=406 is divisible by 2
d3d4d5=063 is divisible by 3
d4d5d6=635 is divisible by 5
d5d6d7=357 is divisible by 7
d6d7d8=572 is divisible by 11
d7d8d9=728 is divisible by 13
d8d9d10=289 is divisible by 17
Find the sum of all 0 to 9 pandigital numbers with this property.
=#
include("projecteulerutils.jl")
using Combinatorics

function calc()
  sum = 0
  for n in permutations(collect(0:9))
    if n[1] != 0
      isvalid = true
      y = [2,3,5,7,11,13,17]
      for x in collect(2:8)
        if nd(n[x:x+2]) % y[x-1] != 0
          isvalid = false
          break
        end
      end
      sum += isvalid ? nd(n) : 0
    end
  end
  sum
end
@time println(calc())


#=
    Simple MCMC:
    Utilizes a Metropolis-Hastings Algorithm with a
    standard Gibbs sampler.
=#

#=
    Basic linear equation function
=#

function lnfnc(x::Array{Float64}, a::Float64, b::Float64)
    y0 = Float64[]
    for xi in x
        y0i = (a * xi) + b
	push!(y0, y0i)
    end
    return y0
end

#=
    Basic chi squared function
    Note: data set error is variance not stddev
=#

function chisqr(yo::Array{Float64}, ye::Array{Float64}, syo::Array{Float64})
    csq = Float64[]
    for (yoi, yei, syoi) in zip(yo, ye, syo)
        csqi = (yoi - yei)^2 / syoi
        push!(csq, csqi)
    end
    return sum(csq)
end

#=
    Basic Metropolis-Hastings Algorithm with a
    standard Gibbs sampler.
=#

function mcmc(a0, sa, b0, sb, x, y, dy)
    y0 = lnfnc(x, a0, b0)
    chi0 = chisqr(y0, y, dy)
    atrace = Float64[]
    btrace = Float64[]
    i = 0
    j = 0
    mcn = 500000
    brn = 1000
    while i < mcn
        at = (randn() * sa) + a0
        bt = (randn() * sb) + b0
        yt = lnfnc(x, at, bt)
        chit = chisqr(yt, y, dy)
        acal = e^((chi0 - chit) / 2.0)
        aexp = min(1.0, acal)
        u = rand()
        if u <= aexp
            a0 = at
            b0 = bt
            y0 = yt
            chi0 = chit
            push!(atrace, at)
            push!(btrace, bt)
            i += 1
        end
        j += 1
    end
    return i, j, sum(atrace[brn:end]) / length(atrace[brn:end]), sum(btrace[brn:end]) / length(btrace[brn:end])
end

#=
    Main
=#

x = Float64[203, 58, 210, 202,
	    198, 158, 165, 201,
	    157, 131, 166, 160,
	    186, 125, 218, 146]

y = Float64[495, 173, 479, 504,
	    510, 416, 393, 442,
	    317, 311, 400, 337,
	    423, 334, 533, 344]

dy = Float64[21, 15, 27, 14,
	     30, 16, 14, 25,
	     52, 16, 34, 31,
	     42, 26, 16, 22]

a0 = 2.5
sa = 0.025
b0 = 28.82
sb = 2.5

# timing start
@time mcmc_out = mcmc(a0, sa, b0, sb, x, y, dy)
# timing end
println(mcmc_out)


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENCE in the project root.
# ------------------------------------------------------------------

# -------------
# COMPOSITIONS
# -------------

"""
    ilr(c)

Isometric log-ratio transformation of composition `c`.
"""
function ilr(c::Composition{D}) where {D}
  w = components(c) .+ eps()
  l = log.(w)
  T = eltype(w)
  x = MVector(ntuple(i->zero(T), D-1))
  for i in 1:D-1
    s = zero(T)
    sqrtinv = 1/sqrt(i*(i+1))
    for j in 1:i+1
      if j < i+1
        s += - sqrtinv * l[j]
      elseif j==i+1
        s += i * sqrtinv * l[j]
      end
    end
    x[i] = s
  end
  SVector(x)
end

"""
    ilrinv(x)

Inverse ilr transformation of coordinates `x`.
"""
function ilrinv(x::SVector{D}) where {D}
  z = MVector(ntuple(i->0.0, D+1))
  for i in 1:D+1
    s = 0.0
    for j in 1:D
      sqrtinv = 1/sqrt(j*(j+1))
      if i < j+1
        s += - sqrtinv * x[j]
      elseif i == j+1
        s += j * sqrtinv * x[j]
      end
    end
    z[i] = exp(s)
  end
  Composition(𝒞(z))
end

ilrinv(x::AbstractVector) = ilrinv(SVector{length(x)}(x))

# -------
# TABLES
# -------

"""
    ILR([refvar])

Isometric log-ratio transform following the
[TableTransforms.jl](https://github.com/JuliaML/TableTransforms.jl)
interface.

Optionally, specify the reference variable `refvar` for the ratios.
Default to the last column of the input table.
"""
struct ILR <: LogRatio
  refvar::Union{Symbol,Nothing}
end

ILR() = ILR(nothing)

refvar(transform::ILR, vars) =
  isnothing(transform.refvar) ? last(vars) : transform.refvar

newvars(::ILR, n) = collect(n)[begin:end-1]

oldvars(::ILR, vars, rvar) = [collect(vars); rvar]

applymatrix(::ILR, X) = mapslices(ilr ∘ Composition, X, dims=2)

function revertmatrix(::ILR, Y)
  D = size(Y, 2)
  f = components ∘ ilrinv ∘ SVector{D}
  mapslices(f, Y, dims=2)
end


function bracket_minimum(f, x=0; s=1e-2, k=2.0)
    a, ya = x, f(x)
    b, yb = a + s, f(a + s)
    if yb > ya
        a, b = b, a
        ya, yb = yb, ya
        s = -s
    end
    while true
        c, yc = b + s, f(b + s)
        if yc > yb
            return a < c ? (a, c) : (c, a)
        end
        a, ya, b, yb = b, yb, c, yc
        s *= k
    end
end

f = x->sin(10x)+cos(3x)
x0 = 3.4
#busca do intervalo
a,c = bracket_minimum(f,x0)

using Plots

plot(f,0.,4.,lw=2,draw_arrow="true")
scatter!([a,c],[f(a),f(c)])


# Generate a random orthogonal matrix
function random_orthogonal!{T}(O::Matrix{T}, n::Int)
    F = qrfact!(randn(n, n))
    Q = full(F[:Q])
    R = F[:R]
    for i in 1:size(R, 1)
        d = sign(R[i, i])
        for j in 1:size(Q, 2)
            Q[i, j] *= d
        end
    end

    return copy!(O, Q)
end

random_orthogonal(n::Int) = random_orthogonal!(Matrix{Float64}(n, n), n)

# Generate a random 3D rotation matrix
function random_rotation()
    Q = random_orthogonal(3)
    det(Q) * Q
end

# Generate a random permutation matrix
random_permutation(n::Int) = eye(n)[:, randperm(n)]


function max_independent_set(g::Graph)

    m = Model(solver=CbcSolver())

    @variable(m, b[v in vertices(g)], Bin)

    @constraint(m , nonadjacency[e in edges(g)], b[src(e)] + b[dst(e)] <= 1)

    @objective(m, Max, sum{b[v] , v in vertices(g)})

    solve(m)
    println("The max indepedent set has $(getobjectivevalue(m)) elements")
    getvalue(b)
end

function min_dominating_set(g::Graph)

    m = Model(solver=CbcSolver())

    @variable(m, b[v in vertices(g)], Bin)

    @constraint(m , dominance[u in vertices(g)], b[u] + sum{b[v], v in out_neighbors(g,u)} >= 1)

    @objective(m, Min, sum{b[v] , v in vertices(g)})

    solve(m)
    println("The min dominating set has $(getobjectivevalue(m)) elements")
    getvalue(b)
end

function min_vertex_cover(g::Graph)

    m = Model(solver=CbcSolver())

    @variable(m, b[v in vertices(g)], Bin)

    @constraint(m , cover[e in edges(g)], b[src(e)] + b[dst(e)] >= 1)

    @objective(m, Min, sum{b[v] , v in vertices(g)})

    solve(m)
    println("The min vertex cover has $(getobjectivevalue(m)) elements")
    getvalue(b)
end

function bipartite_set(g::Graph)

    m = Model(solver=CbcSolver())

    @variable(m, b[v in vertices(g)], Bin)

    @constraint(m , cover[e in edges(g)], b[src(e)] + b[dst(e)] == 1)

    solve(m)

    print("These vertices are independent : ")
    for(v in vertices(g))
        getvalue(b[v]) == 1 && print("$v, ")
    end
    println()

    getvalue(b)
end

function distances(g::Graph, root::Int)

    m = Model(solver=CbcSolver())

    @variable(m, d[v in vertices(g)], Int)

    @constraint(m , d[root] == 0)
    @constraint(m , hop[e in edges(g)], -1 <= d[src(e)] - d[dst(e)] <= 1)

    @objective(m, Max, sum{d[v] , v in vertices(g)})

    solve(m)
    getvalue(d)
end


module JuliaConSamplePackage
using ForwardDiff

export newton

"""
    newton(f, f′, x₀; maxiter = 100, tol = 1e-8)
Find `f`s root via Newton's method.
"""
function newton(f, f′, x₀; maxiter = 100, tol = 1e-8)
    xₙ₊₁ = x₀
    xₙ = NaN
    iter = 0
    while true
        fₓ₊₁ = f(xₙ₊₁)
        fₓ = f(xₙ)
        if abs(fₓ₊₁ - fₓ) < tol
            return xₙ₊₁
        end
        if iter > maxiter
            error("Did not converge!")
        end

        f′ₙ₊₁ = f′(xₙ₊₁)
        if f′ₙ₊₁ == 0.0
            error("Derivative is zero!")
        end
        xₙ₊₁, xₙ = xₙ₊₁ - fₓ₊₁/f′ₙ₊₁, xₙ₊₁
        iter += 1
    end
end

newton(f, x₀; maxiter = 100, tol = 1e-8) =
    newton(f, x -> ForwardDiff.derivative(f, x), x₀; maxiter = maxiter, tol = tol)

end


#==
This example shows multiobjective Paretofronts with 3 objectives.
The front is convex. The NBI method does not catch the full Pareto front,
e.g. as seen by the missing arcs between f1 and f2, whilst f3 is 0.
To capture this area, we need some of the elements in β to be negative.

Taken from
Il Yong Kim and OL De Weck.
Adaptive weighted sum method for multiobjective optimization:
a new method for pareto front generation.
Structural and Multidisciplinary Optimization, 31(2):105–116, 2006.
==#

using MultiJuMP, JuMP
using Ipopt

m = MultiModel(solver = IpoptSolver())

@variable(m, x[i=1:3] >= 0)
@NLexpression(m, f1, x[1])
@NLexpression(m, f2, x[2])
@NLexpression(m, f3, x[3])
@NLconstraint(m, x[1]^4+2x[2]^3+5x[3]^2<=1)

obj1 = SingleObjective(f1, sense = :Max)#,
                       #iv = Dict{Symbol, Any}(:x => [0, 0., 1.2]))

obj2 = SingleObjective(f2, sense = :Max)#,
                       #iv = Dict{Symbol, Any}(:x => [0, 0., 1.2]))

obj3 = SingleObjective(f3, sense = :Max)#,
                       #iv = Dict{Symbol, Any}(:x => [0, 0., 2.]))

md = getMultiData(m)
md.objectives = [obj1, obj2, obj3]
#md.objectives = [obj1, obj2]
md.pointsperdim = 10
solve(m, method = :NBI)

using Plots
pyplot()
plot(md)


using PyPlot
ip = (x, y) -> sum([x_i * y_i for (x_i, y_i) in zip(x, y)])
ct = length(filter(x -> x % 2 == 0, [0:99;]))
pairs = [(2, 5), (4, 2), (9, 8), (12, 10)]
ct2 = length(filter(x -> x[1] % 2 == 0 && x[2] % 2 == 0, pairs))
ct3 = sum([(x % 2 == 0) & (y % 2 == 0) for (x, y) in pairs])
p = (x, coeff) -> sum([a_i * x^(i-1) for (i, a_i) in enumerate(coeff)])
ct4 = str -> sum([a == A for (a, A) in zip(str, uppercase(str))])
isin = (seq_a, seq_b) -> length(seq_a) == sum([in(a, seq_b) for a in seq_a])

function linapprox(f, a, b, n, x)
  delta = (b-a)/n
  points = [a + delta*i for i in 0:n]
  fs = map(f, points)
  m = round(Int, div((x-a), delta))

  plot(points, fs)

  return fs[m] + (fs[m+1]-fs[m])/delta * (x - delta*m)
end

println(ip([1,2], [2,1]))
println(ct)
println(ct2)
println(ct3)
println(p(1, [1, 1, 0]))
println(ct4("Hello"))
println(isin([2:10;], [0:10;]))
println(linapprox(x -> x^2, 0, 1, 5, 0.39))


function asin!(X::IntervalBox)
    x, y = X

    h = half_pi.lo
    y_new = y ∩ Interval(-h, h)  # range of asin
    x_new = sin(y_new)

    return IntervalBox(x_new, y_new)
end

"""
Reverse `asin`.
"""
function asin_rev(y::Interval, x::Interval)  # y = asin(x)

    h = half_pi.lo
    y_new = y ∩ Interval(-h, h)  # range of asin

    x_new = sin(y_new)

    return y_new, x_new  # return in order y, x
end

"""
Reverse `acos`.
"""
function acos_rev(y::Interval, x::Interval)
        y_new = y ∩ Interval(0.0,two_pi.hi)
        x_new = x ∩ cos(y_new)

        return y_new, x_new
end

"""
    atan_rev(y::Interval, x::Interval)

Inverse of `y = atan(x)`.
Returns the new `y` and `x`.
"""
function atan_rev(y::Interval, x::Interval)
        y_new = y ∩ Interval(-half_pi.hi, half_pi.hi)
        x_new = x ∩ tan(y_new)

        return y_new, x_new
end


# 1D diffusion problem

# Packages and inclusions
using ModelingToolkit, MethodOfLines, LinearAlgebra, Test
using ModelingToolkit: Differential
using NonlinearSolve
using DomainSets

# Handrolled discrete laplace problem
@testset "Hand rolled 1D Laplace" begin
    @variables a b c d
    eqs = [0 ~ a - 1,
        0 ~ a + c - 2 * b,
        0 ~ b + d - 2 * c,
        0 ~ d - 1]
    @named ns = NonlinearSystem(eqs, [a, b, c, d], [])
    f = eval(generate_function(ns, [a, b, c, d])[2])
    prob = NonlinearProblem(ns, zeros(4), [])
    sol = NonlinearSolve.solve(prob, NewtonRaphson())

    @test sol.u ≈ ones(4)
end

# Laplace's Equation, same as above but with MOL discretization
@testset "1D Laplace - constant solution" begin
    @parameters x
    @variables u(..)
    Dxx = Differential(x)^2

    eq = Dxx(u(x)) ~ 0
    dx = 1 / 3

    bcs = [u(0) ~ 1,
        u(1) ~ 1]

    # Space and time domains
    domains = [x ∈ Interval(0.0, 1.0)]

    @named pdesys = PDESystem([eq], bcs, domains, [x], [u(x)])
    discretization = MOLFiniteDifference([x => dx], nothing, approx_order=2)
    prob = discretize(pdesys, discretization)
    grid = get_discrete(pdesys, discretization)
    sol = NonlinearSolve.solve(prob, NewtonRaphson())
    sol′ = map(d -> sol[d], grid[u(x)])
    @test sol′ ≈ ones(4)
end

# Laplace's Equation, linear solution
@testset "1D Laplace - linear solution" begin
    @parameters x
    @variables u(..)
    Dxx = Differential(x)^2

    eq = Dxx(u(x)) ~ 0
    dx = 0.1

    bcs = [u(0) ~ 1,
        u(1) ~ 2]

    # Space and time domains
    domains = [x ∈ Interval(0.0, 1.0)]

    @named pdesys = PDESystem([eq], bcs, domains, [x], [u(x)])
    discretization = MOLFiniteDifference([x => dx], nothing, approx_order=2)
    grid = get_discrete(pdesys, discretization)

    prob = discretize(pdesys, discretization)
    sol = NonlinearSolve.solve(prob, NewtonRaphson())
    sol′ = map(d -> sol[d], grid[u(x)])

    @test sol′ ≈ 1.0:0.1:2.0
end

# 2D heat
@testset "2D heat equation" begin
    @parameters x y
    @variables u(..)
    Dxx = Differential(x)^2
    Dyy = Differential(y)^2

    eq = Dxx(u(x, y)) + Dyy(u(x, y)) ~ 0
    dx = 0.1
    dy = 0.1

    bcs = [u(0, y) ~ x * y,
        u(1, y) ~ x * y,
        u(x, 0) ~ x * y,
        u(x, 1) ~ x * y]

    # Space and time domains
    domains = [x ∈ Interval(0.0, 1.0),
        y ∈ Interval(0.0, 1.0)]

    @named pdesys = PDESystem([eq], bcs, domains, [x, y], [u(x, y)])

    # Note that we pass in `nothing` for the time variable `t` here since we
    # are creating a stationary problem without a dependence on time, only space.
    discretization = MOLFiniteDifference([x => dx, y => dy], nothing, approx_order=2)

    prob = discretize(pdesys, discretization)
    grid = get_discrete(pdesys, discretization)
    sol = NonlinearSolve.solve(prob, NewtonRaphson())
    xs, ys = [infimum(d.domain):dx:supremum(d.domain) for d in domains]
    u_sol = map(d -> sol[d], grid[u(x, y)])

    # test boundary
    @test all(abs.(u_sol[:, 1]) .< eps(Float32))
    @test all(abs.(u_sol[1, :]) .< eps(Float32))
    @test u_sol[2:end-1, end] ≈ (0.0:dy:1.0)[2:end-1]
    @test u_sol[end, 2:end-1] ≈ (0.0:dx:1.0)[2:end-1]

    # test interior with finite differences
    interior = CartesianIndices((axes(xs)[1], axes(ys)[1]))[2:end-1, 2:end-1]
    fd = map(interior) do I
        abs(u_sol[(I-CartesianIndex(1, 0))] + u_sol[(I+CartesianIndex(1, 0))] - 2 * u_sol[I]) < 0.01
    end#
    @test all(fd)
end


abstract type AbstractSingleValuedUtility end
abstract type AbstractIsoElastic <: AbstractSingleValuedUtility end

dutil(c, ps::AbstractIsoElastic)=c^(-rra(ps))
invdutil(m, ps::AbstractIsoElastic)=m^(-1/rra(ps))

"
   IsoElastic

IsoElastic specification for a utility with one good.

The parametrization is either
```math
U(c) = \frac{c^{1-\alpha}-1}{1-\alpha}
```

or

```math
U(c) = \frac{c^{1-\alpha}}{1-\alpha}
```

(depending on whether IsoElasticVar2 is passed to `util`.)

The field `elast` corresponds to α.
"
struct IsoElastic <: AbstractIsoElastic
    elast::Float64
    function IsoElastic(elast)
        elast == 1 && error("Please use log utility specification")
        new(elast)
    end
end
struct LogUtility <: AbstractIsoElastic end
function IsoElastic(;a::Float64)
    a == 1 && return LogUtility()
    return(IsoElastic(a))
end

# Methods for iso-elastic types
struct IsoElasticVar2 end

util(c, ps::IsoElastic, ::Type{IsoElasticVar2})=c^(1-ps.elast)/(1-ps.elast)
util(c, ps::IsoElastic)=util(c,ps,IsoElasticVar2)-1/(1-ps.elast)
invutil(u, ps::IsoElastic)=(u*(1-ps.elast)+1)^(1/(1-ps.elast))

util(c, ::LogUtility)=log(c)
invutil(u, ::LogUtility)=exp(u)

rra(ps::IsoElastic)=ps.elast
rra(::LogUtility)=1.

" Constant absolute risk aversion specification "
struct CARA <: AbstractSingleValuedUtility
    a::Float64
end

util(c, ps::CARA)=(1 - exp(-ps.a * c))/ps.a
invutil(u, ps::CARA)=(-1/ps.a)*log(1 - ps.a*u)
dutil(c, ps::CARA)=exp(-ps.a * c)
invdutil(m, ps::CARA)=-log(m)/ps.a

" Risk neutral specification "
struct RiskNeutral <: AbstractIsoElastic end

util(c, ps::RiskNeutral)=c
invutil(u, ps::RiskNeutral)=u
dutil(c, ps::RiskNeutral)=one(c)
invdutil(c, ps::RiskNeutral)=error("Tried to invert constant marginal utility from risk neutral agent")


#=
Created on 09/01/2021 13:08:20
Last update: -

@author: Michiel Stock
michielfmstock@gmail.com

Integrating for dummies. Compute the Riemann sum **without** making use of a for loop.

Compute the area of a circle.
=#

function riemannsum(f, a, b; n=100)
    dx = (b - a) / n
    return sum(f.(a:dx:b)) * dx
end

riemannsum(sin, 0, 2pi)

riemannsum(x->x*sin(x), 0, 2pi)

4riemannsum(x->(sqrt(1-x^2)), 0, 1, n=1000)


raw_data = readlines("../inputs/day03.txt")
num_bits = length(raw_data[1])

function line2ints(line)
    return map(x -> parse(Int, x), collect(line))
end

function get_power_consumption(data)

    function get_gamma(average_vals)
        output_string = ""
        for val in average_vals
            if val < 0.5
                output_string = output_string * "0"
            else
                output_string = output_string * "1"
            end
        end
        return parse(Int, output_string, base = 2)
    end

    function get_epsilon(average_vals)
        output_string = ""
        for val in average_vals
            if val < 0.5
                output_string = output_string * "1"
            else
                output_string = output_string * "0"
            end
        end
        return parse(Int, output_string, base = 2)
    end

    column_sums = sum(data)
    average_value = column_sums / length(data)

    gamma = get_gamma(average_value)
    epsilon = get_epsilon(average_value)

    return gamma * epsilon

end

int_data = map(x -> line2ints(x), raw_data)

power_consumption = get_power_consumption(int_data)
println("Day 03 Part 1 Solution: ", power_consumption)

function get_most_common_in_col(data, index)
    average_val = sum(data) / length(data)
    if average_val[index] >= 0.5
        return 1
    else
        return 0
    end
end

function get_least_common_in_col(data, index)
    average_val = sum(data) / length(data)
    if average_val[index] >= 0.5
        return 0
    else
        return 1
    end
end

function filter_by_most_common_index(data, index)
    filter_val = get_most_common_in_col(data, index)
    return data[map(x -> x[index] == filter_val, data)]
end

function filter_by_least_common_index(data, index)
    filter_val = get_least_common_in_col(data, index)
    return data[map(x -> x[index] == filter_val, data)]
end

function get_oxygen_rating(data)
    output_data = data
    for ii in range(1, num_bits)
        if length(output_data) == 1
            break
        else
            output_data = filter_by_most_common_index(output_data, ii)
        end
    end

    output_string = ""
    for bit in output_data[1]
        output_string = output_string * string(bit)
    end
    return parse(Int, output_string, base = 2)
end

function get_co2_rating(data)
    output_data = data
    for ii in range(1, num_bits)
        _backup_data = output_data
        if length(output_data) == 1
            break
        else
            output_data = filter_by_least_common_index(output_data, ii)
        end

        if length(output_data) == 0
            output_data = _backup_data
        end
    end

    output_string = ""
    for bit in output_data[1]
        output_string = output_string * string(bit)
    end
    return parse(Int, output_string, base = 2)
end

oxygen_rating = get_oxygen_rating(int_data)
co2_rating = get_co2_rating(int_data)
life_support = oxygen_rating * co2_rating

println("Day 03 Part 2 Solution: ", life_support)


function standardize(x::AbstractVector)
    std_x = std(x)
    if std_x == zero(std_x)
        return zero(x)
    else
        return (x .- mean(x)) ./ std_x
    end
end

function standardize(x::AbstractMatrix;dims = 1 )
    standardized_x = similar(x)
    if dims == 1
        for i in 1:size(x,2)
            standardized_x[:,i] .= standardize(x[:,i])
        end
    elseif dims == 2
        for i in 1:size(x,1)
            standardized_x[i,:] .= standardize(x[i,:])
        end
    end
    return standardized_x
end
function reverse_standardization(standardized_x::T, x::T) where T<: AbstractVector
    std_x =std(x)
    if std_x == zero(std_x)
        temp = x[1]
        return fill(temp,size(standardized_x))
    else
        return (standardized_x .* std(x)) .+ mean(x)
    end
end

function reverse_standardization(standardized_x::T, x::T; dims=1) where T<: AbstractMatrix
    new_x =similar(standardized_x)
    if dims == 1
        for i in 1:size(x,2)
            new_x[:,i] .= reverse_standardization(standardized_x[:,i],x[:,i])
        end
    elseif dims == 2
        for i in 1:size(x,1)
            new_x[i,:] .= reverse_standardization(standardized_x[i,:],x[i,:])
        end
    end
    return new_x
end

function gradient(x::Vector{T}) where {T<:Number}
    res = similar(x)
    for i = 1:length(x)
        if i == 1
            res[i] = x[i+1] - x[i]
        elseif i == size(x, 1)
            res[i] = x[i] - x[i-1]
        else
            res[i] = (x[i+1] - x[i-1]) / 2
        end
    end
    return res
end

function gradient(m::Matrix{T}) where {T<:Number}
    res = similar(m)
    for i = 1:size(m, 1)
        if i == 1
            res[i, :] = m[i+1, :] - m[i, :]
        elseif i == size(m, 1)
            res[i, :] = m[i, :] - m[i-1, :]
        else
            res[i, :] = (m[i+1, :] - m[i-1, :]) / 2
        end
    end
    return res
end

# function TNB_reorientation(verts::AbstractVector, T, N, B)
#     new_verts = similar(verts)
#     for i in eachindex(new_verts)
#         new_verts[i] = [T N B] \ verts[i]
#     end
#     return new_verts
# end


# Implement code block 2.8 in McElreath's Statistical Rethinking, 2nd edition.

using Random
using Distributions
using Plots

Random.seed!(17)

num_samples = 2000

probas = [0.5]

successes = 6
trials = 9

for i = 2:num_samples

    proba_new = rand(Normal(probas[end], 0.1))

    if proba_new < 0; proba_new = abs(proba_new); end
    if proba_new > 1; proba_new = 2 - proba_new; end

    q0 = pdf(Binomial(trials, probas[end]), successes)
    q1 = pdf(Binomial(trials, proba_new), successes)

    push!(probas, rand(Uniform(0, 1)) < q1 / q0 ? proba_new : probas[end])
end

# println(probas)

display(histogram(probas, bins=:40))


function g_ij(i::Int64,j::Int64,t)
    ##OverFlowError if i>20

    s=0 ##The sum of each term
    for k in j:i
        e=exp(-k*(k-1)*t/2)
        numerator=((2*k)-1)*((-1)^(k-j))
        denominator=factorial(j)*factorial(k-j)*(j+k-1)

        mult_term=1
        for y in 0:(k-1)
            mult_term*=(((j+y)*(i-y))/(i+y))
        end
        val=e*(numerator/denominator)*mult_term
        s+=val
    end
    return(s)
end

function nHistories(i::Int64,j::Int64)
    ##the number of histories where i lineages coalesce into j lineages

    mult_term=1
    (j==1) && (j=2) ## coalescing into 1 or 2 lineages have the same number of histories but binomial(k,2) returns 0 when k=1
    for k in (j+1):i
        mult_term*=binomial(k,2)
    end
    return(mult_term)
end

function nTops(n::Int64,labeled=true::Bool,rooted=true::Bool)
    if labeled ##Number of labeled tree topologies
        n==1 && return(1)
        return(Int(factorial(2*n-3)/((2^(n-2)*factorial(n-2)))))

    else ##number of unlabeled topologies i.e. tree shapes
        n==1 && return(1) ##base case for recursion
        odd=true
        (mod(n,2)==0) && (odd=false)

        if odd ## n is odd
            val=0
            for p in 1:((n-1)/2)
                val+=nTops(p,false)*nTops(n-p,false)
            end
        else ## n is even
            t2=nTops(n/2,false)
            term=t2*(t2+1)/2

            s=0
            for p in 1:((n/2)-1)
                s+=nTops(p,false)*nTops(n-p,false)
            end
            val=term+s
        end
        return(Int(val))
    end
end

function nLabeledShape(n::Int64,s::Int64)
    ##Compute the number of labeled topologies of a given tree shapes
    ##Where:
        #n= # of taxa
        #s= # of symetric nodes i.e. nodes where its two subtrees have the same shape
    return(factorial(n)/(2^s))
end

function nHistoriesTop(desc_nodes::Array{Int64,1})
    return(factorial(length(desc_nodes))*prod((x->(1/(x+1)).(desc_nodes))))
end

function unlabeledGenerate(n)
    ##we will create dictionaries to store the bits we want for unlabeled tree shapes
    ##Each key in the dictionary will point to a vector of the bits for each tree shape
        ##e.g. sym_tops[5] will give us a vector of how many points of symmetry each tree shape with 5 taxa has
    vect_tops=Dict{Int64, Array{Array{Int64,1},} }() ##this bit is the node vector representation of the tree shape. Each element in the vector represents an internal node and it denotes how many internal nodes it has as descendents
    sym_tops=Dict{Int64,Array{Int64,1}}() ##this bit stores the number of reflection points for a given tree shape
    prod_tops=Dict{Int64,Array{Float64,1}}() ##if x is the node vector representation of the tree shape, this bit stores the product of 1/(1+x)

    ##Initialize the first two sets
    vect_tops[1]=[[]]
    sym_tops[1]=[0]
    prod_tops[1]=[1]

    vect_tops[2]=[[0]]
    sym_tops[2]=[1]
    prod_tops[2]=[1]

    for i in 3:n
        vect_tops[i]=[]
        sym_tops[i]=[]
        prod_tops[i]=[]

        for j in 1:(floor(i/2))
            ##make new shapes by combining all the the j_th and (n-j)_th arrays

            bottom=j
            top=i-j

            for k in 1:(length(vect_tops[bottom]))
                for l in 1:(length(vect_tops[top]))
                   first=vect_tops[bottom][k]
                   second=vect_tops[top][l]

                   c=Int64[]
                   append!(c,first)
                   append!(c,second)
                   push!(c,length(c))
                   push!(vect_tops[i],c)

                   sym_new=(sym_tops[bottom][k]+sym_tops[top][l]) ##The number of symmetry points is the sum of the symmetry points of the two subtrees
                   (first==second) && (sym_new+=1) ## If the two shapes are the same then this makes another point of symmetry
                   push!(sym_tops[i],sym_new)

                   push!( prod_tops[i],prod( (x->(1/(x+1))).(c) ) ) ##TODO could make this more efficient by instead computing a new w from the c vector, just multiply the to previous products together
                end
            end
        end
    end
    return ((vect_tops,sym_tops,prod_tops));
end

function splitsHistories(splits::Array{Int64,1}, tops_mat)
    total=[0]

    spltsHisRec(splits,1,1,tops_mat,total)

    return total[1];

end

function spltsHisRec(splits::Array{Int64,1},ind_splits,running_val,tops_mat,total::Array{Int64,1} )
    split=splits[ind_splits]
    ##Get all the symmetries and products for the split we're interested in
    syms=tops_mat[2][splits[ind_splits]]
    prods=tops_mat[3][splits[ind_splits]]

    nshapes=length(syms) ##the number of treeshapes for a given number of taxa in the split

    for shape in 1:nshapes ##Do calculation for each shape
        val=running_val ##reset the value for each shape

        val*=factorial(split)/(2^(syms[shape]))
        val*=prods[shape]

        if ind_splits==length(splits) ##Base Case: we went thru each split and did the calculation
            val*=factorial(sum((x->x-1).(splits)))
            total[1]+=val
        else ##we still have more splits to do calculations
            spltsHisRec(splits,ind_splits+1,val,tops_mat,total)
        end
    end
end

function findLowNode(net::HybridNetwork) ##TODO. Make this generate a list of all the nodes to do in order instead of one at a time. Also, since net.node never changes until the end, we could have it store the node indices instead of pointers to the nodes themselves.
    numbers=(x->x.number).(net.node) ##store old numbers '
    nd_changed= net.nodes_changed ##store node ordering
    rankNodes!(net)

    ##get all hybrid nodes
    hybs=net.hybrid

    nds=Set() ##This will have all the nodes we need to decompose
    for hyb in hybs
        union!(nds,Set(preorder(net,hyb))) ##add the hyb node and all its descendents to the list
    end
    nds=collect(nds) ##turn the nodes into an array
    nds= nds[(x->(!x.leaf)).(nds)] ##get rid of tips

    nds=nds[sortperm((x->x.number).(nds))]
    ((x,y)->x.number=y).(net.node,numbers) ##reset node numbers to their old values
    net.nodes_changed=nd_changed ##reset node ordering

    return nds
end

function decompILS(network::HybridNetwork,decomp_node::Node,prb::Float64,tm,hyb_sorting::Dict{Int64, Dict{Int64,Set{String}}})

    ## Make a copy of the network that we can muck around with.
    nd_ind=findfirst(x->x==decomp_node,network.node)
    net=deepcopy(network)
    nd=net.node[nd_ind]

    par=getParents(nd)[1]
    nmes=(x->x.name).(getChildren(nd))

    c_edge=getChildrenEdges(nd)[1].length   ##get the edge length of the first child, ignore other child edge lengths, assuming ultrametricy they should all be the same anyways
    p_edge=getChildrenEdges(par)[1].length  ##get the edge length of the parent
    elen=c_edge+p_edge ##This will be the edge length we assign to our coalesced edges

    removeClade!(net,nd,false)
    par_ind=findfirst(x->x==par,net.node)

    subsets=collect(partitions(nmes)) ##we need the power set (excluding the empty set) of all names for decomposition

    ##We ultimately want ot report the decomposed trees and their respective probabilities after everything is all said and done
    decomposed_trees=HybridNetwork[]
    probs=Float64[]

    for subset in subsets
        ##start with a fresh tree for adding things back
        decomp_net=deepcopy(net)
        decomp_par=decomp_net.node[par_ind]

        new_node_num=maximum((x-> x.number).(net.node))
        new_edge_num=maximum((x-> x.number).(net.edge))

        nms_coal=String[] ##These are the names of the coalesced tips we're adding to the tree
        for nameset  in subset ##We need the individual sets within subset to concatenate the names
            if length(nameset)==1
                nm=nameset[1]
            else
                nm=join(nameset,"|")
            end
            push!(nms_coal,nm)
        end

        for nm in nms_coal #Sequentially add each coalesced tip back to the tree
            ##create a new node and edge for the tip
            ntip=Node(new_node_num,true)
            ntip.name=nm
            nedge=Edge(new_edge_num,elen)
            new_edge_num+=1
            new_node_num+=1

            ##make appropriate connections
            push!(decomp_par.edge,nedge)
            push!(nedge.node,ntip)
            push!(nedge.node,decomp_par)
            push!(ntip.edge,nedge)

            ##add the new nodes and edges to the network
            push!(decomp_net.node,ntip)
            push!(decomp_net.edge,nedge)
            push!(decomp_net.leaf,ntip)
        end

        ##update the attributes of the network
        num_added=length(nms_coal)
        decomp_net.numTaxa+=num_added
        decomp_net.numNodes+=num_added
        decomp_net.numEdges+=num_added
        preorder!(decomp_net) ##preorder things for good measure

        ##Compute the probability of the tree
        gij=g_ij(length(nmes),length(nms_coal),p_edge) ## the probability of i lineages coalescing into j lineages by time t
        hp=splitsHistories(length.(subset),tm) ##The number of coalescent histories that lead to these partitions
        ch=nHistories(length(nmes),length(nms_coal)) ##The total number of coalescent histories for i coalescing into j lineages

        treeprob=gij*(hp/ch)*prb ##TODO I'm concerned with underflow so make this deal with small numbers. perhaps switch to log probs
        treeprob==0 && error("Likely underflow of the probability")

        push!(decomposed_trees,decomp_net)
        push!(probs,treeprob)
    end
    ##Make copies of hyb_sorting for each decomposed tree
    hybsort= Dict{Int64, Dict{Int64,Set{String}}}[]
    for i in 1:length(probs)
        push!(hybsort,deepcopy(hyb_sorting))
    end
    return collect(zip(decomposed_trees,probs,hybsort))
end

function emptyHybSorting(net::HybridNetwork;type="name")
    if type == "name"
        hybsort= Dict{Int64, Dict{Int64,Set{String}}}()
        hybs=net.hybrid
        for hyb in hybs
            hybsort[hyb.number]=Dict{Int64,Set{String}}()
            for nd in getParents(hyb)
            hybsort[hyb.number][nd.number]=Set{String}()
            end
        end
    elseif type == "node"
        hybsort= Dict{Int64, Dict{Int64,Set{Node}}}()
        hybs=net.hybrid
        for hyb in hybs
            hybsort[hyb.number]=Dict{Int64,Set{Node}}()
            for nd in getParents(hyb)
            hybsort[hyb.number][nd.number]=Set{Node}()
            end
        end
    else
        error("we didn't get a valid type. Either 'node' or 'name' is accepted")
    end
    return hybsort
end

function decompHyb(network::HybridNetwork,nd::Node,prb::Float64,hyb_sorting::Dict{Int64, Dict{Int64,Set{String}}} )
    ## Make a copy of the network that we can muck around with.
    nd_ind=findfirst(x->x==nd,network.node)
    net=deepcopy(network)
    nd=net.node[nd_ind]

    hyb_num= nd.number

    lpar=getMajorParent(nd) ##left parent
    rpar=getMinorParent(nd) ##right parent
    lnum=lpar.number
    rnum=rpar.number
    lp_e=getMajorParentEdge(nd)
    rp_e=getMinorParentEdge(nd)

    c_elen=getChildrenEdges(nd)[1].length   ##get the edge length of the first child, ignore other child edge lengths, they should be the same anyway
    lp_elen=lp_e.length
    rp_elen=rp_e.length

    l_gamma=lp_e.gamma
    r_gamma=rp_e.gamma

    removeClade!(net,nd,false)

    lpar_ind=findfirst(x->x==lpar,net.node) ##The index of the left parent in net.node
    rpar_ind=findfirst(x->x==rpar,net.node) ##The index of the right parent in net.node

    ##create all combinations for how the lineages could sort between the two hybrid getParents
    nmes=(x->x.name).(getChildren(nd)) ##get each lineage
    lsort=collect(combinations(nmes)) ##Generate all possible ways that things could sort to the left
    push!(lsort,[]) ## the function above doesn't account for the fact that no species could sort left
    rsort=(x->setdiff(nmes,x)).(lsort) ## everything that doesn't sort left must go right. NOTE: we assume that no lineages have the same name here
    combs=collect(zip(lsort,rsort)) ##associate the proper sortings together to generate all ways that lineages can sort

    ##These will store the decomposed trees and their probabilities
    decomposed_trees=HybridNetwork[]
    probs=Float64[]
    hybsorts=Dict{Int64, Dict{Int64,Set{String}}}[]

    for i in combs
        ##start with a fresh slate for adding things
        hybsort=deepcopy(hyb_sorting)
        decomp_net=deepcopy(net)
        ldecomp_par=decomp_net.node[lpar_ind]
        rdecomp_par=decomp_net.node[rpar_ind]

        leftgoing=first(i)
        rightgoing=last(i)

        ####################
        ###Deal with left###
        ####################
        if length(leftgoing)>1 ##if more than 1 tip, we need to make a node for the mrca of all tips going left
            nedge=Edge(-1,lp_elen)
            ##Make appropriate connections
            push!(ldecomp_par.edge,nedge)
            push!(nedge.node,ldecomp_par)
            ldecomp_par=Node(-1,false,false)
            push!(ldecomp_par.edge,nedge)
            pushfirst!(nedge.node,ldecomp_par)
            ##add the new nodes and tips to the network
            push!(decomp_net.node,ldecomp_par)
            push!(decomp_net.edge,nedge)
            ##update attributes of the network
            decomp_net.numNodes+=1
            decomp_net.numEdges+=1
            r_elen=c_elen
        elseif  length(leftgoing)==0 ##if there is nothing going to the left we need to resolve the parent node
            deleteleaf!(decomp_net,findfirst(x->x===ldecomp_par,decomp_net.node);index=true,simplify=false)
        else ##only 1 tip going left
            r_elen=c_elen+rp_elen
        end
        for nm in leftgoing ##add tips going left to the network
            ntip=Node(-1,true)
            ntip.name=nm
            nedge=Edge(-1,lp_elen+c_elen)

             ##make appropriate connections
             push!(ldecomp_par.edge,nedge)
             push!(ntip.edge,nedge)
             push!(nedge.node,ntip)
             push!(nedge.node,ldecomp_par)

            ##add the new nodes and tips to the network
            push!(decomp_net.node,ntip)
            push!(decomp_net.edge,nedge)
            push!(decomp_net.leaf,ntip)
        end

        ####################
        ###Deal with right##
        ####################
        if length(rightgoing)>1 ##if more than 1 tip, we need to make a node for the mrca of all tips going right
            nedge=Edge(-1,rp_elen)
            ##Make appropriate connections
            push!(rdecomp_par.edge,nedge)
            push!(nedge.node,rdecomp_par)
            rdecomp_par=Node(-1,false,false)
            push!(rdecomp_par.edge,nedge)
            pushfirst!(nedge.node,rdecomp_par)
            ##add the new nodes and tips to the network
            push!(decomp_net.node,rdecomp_par)
            push!(decomp_net.edge,nedge)
            ##update attributes of the network
            decomp_net.numNodes+=1
            decomp_net.numEdges+=1
            r_elen=c_elen
        elseif  length(rightgoing)==0 ##if there is nothing going to the right we need to resolve the parent node
            deleteleaf!(decomp_net,findfirst(x->x===rdecomp_par,decomp_net.node);index=true,simplify=false)
        else #only 1 tip going right
            r_elen=c_elen+rp_elen
        end
        for nm in rightgoing ##add tips to the right parent
            ntip=Node(-1,true)
            ntip.name=nm
            nedge=Edge(-1,r_elen)

             ##make appropriate connections
             push!(rdecomp_par.edge,nedge)
             push!(ntip.edge,nedge)
             push!(nedge.node,ntip)
             push!(nedge.node,rdecomp_par)

            ##add the new nodes and tips to the network
            push!(decomp_net.node,ntip)
            push!(decomp_net.edge,nedge)
            push!(decomp_net.leaf,ntip)
        end
        ##the tree should now be deocmposed

        ##Update the attributes of the network
        num_added=length(rightgoing)+length(leftgoing)
        decomp_net.numNodes+=num_added
        decomp_net.numEdges+=num_added
        decomp_net.numTaxa+=num_added
        preorder!(decomp_net) ##preorder things for good measure

        ##record which lineages went left and right along the reticulation
        left_names=String[]
        right_names=String[]
        lnms= (x->split(x,"|")).(leftgoing)
        rnms= (x->split(x,"|")).(rightgoing)
        println(lnms)
        println(rnms)
        for nms in lnms
            append!(left_names,String.(nms))
        end
        for nms in rnms
            append!(right_names,String.(nms))
        end
        hybsort[hyb_num][lnum]=Set(left_names)
        hybsort[hyb_num][rnum]=Set(right_names)

        ##Compute the probability of the tree
        treeprob=(r_gamma^length(rightgoing))*(l_gamma^length(leftgoing))*prb ##TODO I'm concerned with underflow so make this deal with small numbers. perhaps switch to log probs
        treeprob==0 && error("Likely underflow of the probability")

        push!(decomposed_trees,decomp_net)
        push!(probs,treeprob)
        push!(hybsorts,hybsort)
    end
    return collect(zip(decomposed_trees,probs,hybsorts))
end

function networkDecomposition(net::HybridNetwork)

    hybsorting=emptyHybSorting(net)
    prob=1.0
    q= Array{Tuple{HybridNetwork,Array{Dict{Int64,Dict{Int64,Set{String}}},1},Float64}}() ##Queue of partially decomposed trees along with their hyb_sorting profiles and probabilities
    fully_decomposed =  Array{Tuple{HybridNetwork,Array{Dict{Int64,Dict{Int64,Set{String}}},1},Float64}}() ##Where we store fully decompposed trees
    push!(q,(net,hybsorting,1.0))

    ##determine how big of an unlabeledGenerate we need to run. it should be the size of max(numberLeaves(hyb_node) for all hyb_node in hyb_nodes)
    nleaves=Int64[]
    for hyb in net.hybrid
        e=getChildrenEdges(hyb)[1]
        length(descendants(e))
    end

    while (!isempty(q)) ##while there are still things in the Queue...

        current=popfirst!(q)

        curr_net=current[1]
        curr_hybsort=current[2]
        curr_prob=current[3]

        decomp_nds=findLowNode

    end

    nd=findLowNode(net)

end


## Load the MNIST Graph and draw it
##
using DelimitedFiles
using SparseArrays
using LinearAlgebra
ei,ej = copy.(eachcol(Int.(readdlm("mnist-train-4.edges"))))
xy = readdlm("mnist-train-4.xy")
n = maximum(ei)
A = sparse(ei,ej,1,n,n)
@assert issymmetric(A)

##
using MatrixNetworks
f,lam = fiedler_vector(A)
graphplot(ei, ej, x =xy[:,1], y=xy[:,2],
  markercolor=:black, markerstrokecolor=:white,
  size=(1200,1200),dpi=400,
  markersize=2, linecolor=1, linealpha=0.2, linewidth=0.5,
  markeralpha=0.3,
  #colorbar=true,
  marker_z=f,
  axis_buffer=0.02, background=nothing)
##
savefig("mnist-train-4-fiedler.png")

##
cut = spectral_cut(A)
x = zeros(n)
x[cut.set] .= 1
graphplot(ei, ej, x =xy[:,1], y=xy[:,2],
  markercolor=:black, markerstrokecolor=:white,
  size=(1200,1200),dpi=400,
  markersize=2, linecolor=1, linealpha=0.2, linewidth=0.5,
  markeralpha=0.3,
  #colorbar=true,
  marker_z=x,
  axis_buffer=0.02, background=nothing)
for i=1:100
  annotate!(xy[i,1],xy[i,2],"$(train_y[i])",color=:white)
end
##
savefig("mnist-train-4-spectral-cut.png")

## Make a figure for the label 1
# find the first 10 nodes labeled 1
using MLDatasets
train_x, train_y = MNIST.traindata()
nlabels = 10
S = findall(train_y .== 1)[1:nlabels]
## Show the labeled nodes.
scatter(xy[:,1],xy[:,2])
scatter!(xy[S,1], xy[S,2], markersize=25)
##
# Form the incidence matrix for the least squares problem
function incidence_matrix(A)
  ei,ej = findnz(triu(A,1))
  n = size(A,1)
  m = length(ei)
  return sparse(1:m,ei,1,m,n) - sparse(1:m,ej,1,m,n)
end
B = incidence_matrix(A)
IS = sparse(0.1*I, n,n)
IS[S,S] .*= 1/maximum(IS)
labels = zeros(n)
labels[S] .= 1.0
## solve the least squares problems
#x = [B; IS] \ [zeros(size(B,1)); 1]
AtA = [B; IS]'* [B; IS]
Atb = [B; IS]'* [zeros(size(B,1)); labels]
##
x = AtA\Atb
##
scatter(xy[:,1],xy[:,2], marker_z = log10.(x), markerstrokewidth=0)

## Show the labeled nodes.
scatter(xy[:,1],xy[:,2],background=nothing,framestyle=:none,markersize=2,
  size=(1200,1200),markerstrokewidth=0, alpha=0.5,dpi=300)
scatter!(xy[S,1], xy[S,2], markersize=12)
for i in S
  annotate!(xy[i,1],xy[i,2], "$(train_y[i])")
end
plot!(legend=false)
savefig("mnist-train-4-labels.png")
##
scatter(xy[:,1],xy[:,2],background=nothing,framestyle=:none,markersize=2,
  size=(1200,1200),markerstrokewidth=0, marker_z = log10.(x),
  colorbar=false, alpha=0.5,dpi=300,label="")
savefig("mnist-train-4-soln.png")


# Scale from/to [0,1].
scalefrom(a,b,t) = @. (t-a)/(b-a)
scaleto(a,b,t) = @. a + t*(b-a)

# Unique real roots of a quadratic.
function realroots(a,b,c)
	a==0 && return [-c/b]
	bh = -b/2
	d = bh^2 - a*c
	if d < 0
		[]
	elseif d==0
		[bh/a]
	else
		r = (bh + sign(bh)*sqrt(d)) / a
		[r,c/(r*a)]
	end
end

# Are three given points arranged in counterclockwise order?
function isccw(a::Number,b::Number,c::Number)
	v(z) = SVector(1,real(z),imag(z))
	det([v(a) v(b) v(c)]) > 0
end

# Use 2nd order finite differences to approximate a tangent.
function fdtangent(z,t::Real)
	ϵ = eps(typeof(float(t)))
	ϵ3 = 0.5*ϵ^(1/3)
	if t < ϵ3
		τ = (-1.5*z(t) + 2*z(t+ϵ3) - 0.5*z(t+2ϵ3)) / ϵ3
	elseif t > 1-ϵ3
		τ = (1.5*z(t) - 2*z(t-ϵ3) + 0.5*z(t-2ϵ3)) / ϵ3
	else
		τ = (z(t+ϵ3) - z(t-ϵ3)) / (2ϵ3)
	end
	return τ
end

# Select points adaptively to make a smooth-appearing curve.
function adaptpoints(point,utangent,a,b;depth=6,curvemax=0.05)
	function refine(tl,tr,zl,zr,τl,τr,maxdz,d=depth)
		dzkap = dist(τr,τl)  # approximately, the stepsize over radius of curvature
		tm = (tl+tr)/2
		zm = point(tm)
		τm = utangent(tm)
		if d > 0 && (dzkap > curvemax || dist(zr,zl) > maxdz )
			zl = refine(tl,tm,zl,zm,τl,τm,maxdz,d-1)
			zr = refine(tm,tr,zm,zr,τm,τr,maxdz,d-1)
			return [zl;zm;zr]
		else
			return zm
		end
	end

	d = (b-a)/4
	tt = d*[0,0.196,0.41,0.592,0.806]   # avoid common symmetry points
	t = [a .+ tt; a + d .+ tt; a + 2d .+ tt; a + 3d .+ tt; b]
	z = point.(t)
	τ = utangent.(t)

	# on the Riemann sphere, use distance in R^3
	if z[1] isa Spherical
		dist = (u,v) -> norm(S2coord(u)-S2coord(v))
	else
		dist = (u,v) -> abs(u-v)
	end
	m = length(t)
	scale = maximum(dist(z[i],z[j]) for i=2:m-1, j=2:m-1 if j > i)
	zfinal = z[[1]]
	for j = 1:length(t)-1
		znew = refine(t[j],t[j+1],z[j],z[j+1],τ[j],τ[j+1],scale/25)
		append!(zfinal,znew)
		push!(zfinal,z[j+1])
	end
	return zfinal
end

# Do adaptive integration to estimate the integral of `f` over [`a`,`b`] to desired
# error tolerance `tol`.
function intadapt(f,a,b,tol)
    # Use error estimation and recursive bisection.
    function do_integral(a,fa,b,fb,m,fm,tol,depth)
        # These are the two new nodes and their f-values.
        xl = (a+m)/2;  fl = f(xl);
        xr = (m+b)/2;  fr = f(xr);
        t = [a,xl,m,xr,b]              # all 5 nodes at this level

        # Compute the trapezoid values iteratively.
        h = (b-a)
        T = [0.,0.,0.]
        T[1] = h*(fa+fb)/2
        T[2] = T[1]/2 + (h/2)*fm
        T[3] = T[2]/2 + (h/4)*(fl+fr)

        S = (4*T[2:3]-T[1:2]) / 3      # Simpson values
        E = (S[2]-S[1]) / 15           # error estimate

        if abs(E) < tol*(1+abs(S[2]))  # acceptable error?
            Q = S[2]                   # yes--done
		else
			if depth==0
				@warn "Too many recursions to determine integral"
				Q = S[2]
			else
    	        # Error is too large--bisect and recurse.
        	    QL = do_integral(a,fa,m,fm,xl,fl,tol,depth-1)
            	QR = do_integral(m,fm,b,fb,xr,fr,tol,depth-1)
				Q = QL + QR
			end
        end
        return Q
    end

    m = (b+a)/2
    Q = do_integral(a,f(a),b,f(b),m,f(m),tol,50)
    return Q
end

function enclosing_circle(z::AbstractVector,expansion=2)
	xa,xb = extrema(real(z))
	ya,yb = extrema(imag(z))
	zc = complex((xa+xb)/2,(ya+yb)/2)
	R = length(z) > 1 ? maximum(@. abs(z - zc)) : max(1,abs(zc))
	return zc,expansion*R
end

function enclosing_box(z::AbstractVector,expansion=2)
	zc = sum(z)/length(z)
	dz = z .- zc
	rx = length(z) > 1 ? maximum(@. abs(real(dz))) : max(1,abs(real(zc)))
	ry = length(z) > 1 ? maximum(@. abs(imag(dz))) : max(1,abs(imag(zc)))
	return real(zc).+expansion*[-rx,rx],imag(zc).+expansion*[-ry,ry]
end

# indices of the closest pair of points from two lists
function argclosest(z1,z2)
	i1 = [argmin(abs.(z1.-z)) for z in z2]
	i2 = argmin(abs.(z2.-z1[i1]))
	return i1[i2],i2
end


### A Pluto.jl notebook ###
# v0.12.6

using Markdown
using InteractiveUtils

# ╔═╡ de7a5e8a-1ee4-11eb-2bbf-0f0eb3e489ec
using Pkg, DrWatson

# ╔═╡ e4eeb13a-1ee4-11eb-0113-91f9a9c3b659
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Random, StatsBase, DataFrames
	Random.seed!(1)
end

# ╔═╡ ac3ceabe-1ee4-11eb-3935-313e96aafd62
md"## Listing 2.7"

# ╔═╡ c4bef1a4-1ee9-11eb-0b58-8be033f6fbad
begin
	A = Set(['a','e','i','o','u'])
	B = Set(['x','y','z'])
	omega = 'a':'z'
	N = 10^6
end

# ╔═╡ 23c4c304-1eea-11eb-2a57-d9b2e99a3253
begin
	df = DataFrame()
	for _ in 1:5
    	mcEst1 = sum([in(sample(omega),A) || in(sample(omega),B) for _ in 1:N])/N
    	mcEst2 = sum([in(sample(omega),union(A,B)) for _ in 1:N])/N
    	append!(df, DataFrame(:mcEst1 => mcEst1, :mcEst2 => mcEst2))
	end
	df
end

# ╔═╡ 6285dea2-1ee5-11eb-2b6a-f9906c0d3351
md"## End of listing 2.7"

# ╔═╡ Cell order:
# ╟─ac3ceabe-1ee4-11eb-3935-313e96aafd62
# ╠═de7a5e8a-1ee4-11eb-2bbf-0f0eb3e489ec
# ╠═e4eeb13a-1ee4-11eb-0113-91f9a9c3b659
# ╠═c4bef1a4-1ee9-11eb-0b58-8be033f6fbad
# ╠═23c4c304-1eea-11eb-2a57-d9b2e99a3253
# ╟─6285dea2-1ee5-11eb-2b6a-f9906c0d3351


module RadioPropagation
	# using Plots
	using Printf

	c = 299792458; # m/s
	i = im;

	"""
	Calculate the two-way two-ray propagation between two points above a flat plane.
	Returns the propagation factor F. The one-way power propagation factor is |F|².

	- Radar Systems Engineering Lecture 5 Propagation through the Atmosphere, IEEE New Hampshire Section IEEE AES Society, 2010

	```julia
	# Example code plotting signal strength as function of range, compared to free space power spread
	using Plots

	c = 3e8;

	transmit_height_m	= 2;
	receive_height_m	= 10;
	target_range_m		= LinRange(20, 100, 10000);
	transmit_frequency_hz =3e9;
	λ = c/transmit_frequency_hz;
	Γ1 = -1;
	Γ2 = -0.6;

	F²( Γ ) = RadioPropagation.two_ray_propagation.( target_range_m, transmit_height_m, receive_height_m, transmit_frequency_hz, Γ );

	F²1 = F²( Γ1 );
	F²2 = F²( Γ2 );

	f⁴_db( F² ) = 10*log10.( abs2.( F² ));

	freespace_loss_db = 10 .*log10.(λ^2 ./( (4*π)^3 .*target_range_m .^4 ));
	bias = freespace_loss_db[1]

	plot1 = plot( 	target_range_m, f⁴_db.(F²1)+freespace_loss_db .-bias,
			 	 	xlabel 	= "Range (one-way) [m]",
					ylabel 	= "Two-way loss [dB]",
					title  	= "Free-space versus two-ray propagation.",
					ylims 	= (-30, 6),
					label 	= "|F|⁴+L, Γ=-1",
					legend	= true,
					xaxis	=:log,
					dpi=300)
	plot!( 	target_range_m, f⁴_db.(F²2)+freespace_loss_db .-bias,
	 		label = "|F|⁴+L, Γ=-0.6",
	 		xaxis	=:log )
	plot!( 	target_range_m, freespace_loss_db .-bias,
	 		label = "L",
	 		xaxis	=:log )
	savefig("figures/example_figure")

	```

	# Arguments
	- 'distance'            The distance parallel to the plane earth.
	- 'transmit_height'     The height of the transmitter above the plane earth.
	- 'receive_height'      The height of the receiver/target above the plane earth.
	- 'Γ'                           The reflection coefficient of the medium of the plane earth.
	"""
	function two_ray_propagation( distance, transmit_height, receive_height, frequency_hz, Γ )
		λ = c/frequency_hz;
		Δϕ=  4*π*transmit_height*receive_height/(λ*distance);
		F = 1+abs(Γ)*exp(i*Δϕ);
		return F;
	end

	"""
	Empirical model for RF rain attenuation for frequencies between 1 and 400 GHz, linear polarization.
	Model uses the closest frequency in the underlying data.

	- M. A. Richards and J. A. Scheer and W. A. Holm, Principles of Modern Radar, SciTech Publishing, 2010.

	```julia
	rain_attenuation_db_per_km_circular_pol( 'v' , 30, 20 )
	3.3400000000000003
	```

	# Arguments
	- 'polarization'        The polarization, 'v', 'h'.
	- 'frequency_ghz'     	The frequency of the propagating waves.
	- 'fall_rate_mm_hour'	The rain intensity [mm/h].
	"""
	function rain_attenuation_db_per_km( polarization::Char, frequency_ghz, fall_rate_mm_hour )
		r = fall_rate_mm_hour;

		# Empirical data:
		frequency_table_ghz = [1, 2, 3 ,6, 7, 8, 10, 12, 15, 20, 25, 30, 35, 40, 45, 50, 60, 70, 80, 90, 100, 120, 150, 200, 300, 400 ];
		ah = [ 3.87e-5, 1.54e-4, 6.5e-4, 1.75e-3, 3.01e-3, 4.54e-3, 1.01e-2, 1.88e-2, 3.67e-2, 7.51e-2, .124, .187, .263, .350, .442, .536, .707, .851, .975, 1.06, 1.12, 1.18, 1.31, 1.45, 1.36, 1.32 ];
		av = [ 3.52e-5, 1.38e-4, 5.91e-4, 1.55e-3, 2.65e-3, 3.95e-3, 8.87e-3, 1.68e-2, 3.47e-2, 6.91e-2, .113, .167, .233, .310, .393, .479, .642, .784, .906, .999, 1.06, 1.13, 1.27, 1.42, 1.35, 1.31 ];
		bh = [ .912, .963, 1.121, 1.308, 1.332, 1.327, 1.276, 1.217, 1.154, 1.099, 1.061, 1.021, .979, .939, .903, .873, .826, .793, .769, .753, .743, .731, .710, .689, .688, .683 ];
		bv = [ .88, .923, 1.075, 1.265, 1.312, 1.310, 1.264, 1.2, 1.128, 1.065, 1.030, 1.0, .963, .929, .897, .868, .824, .793, .769, .754, .744, .732, .711, .69, .689, .683 ];

		# TODO implement interpolation vor arbitrary frequency selection.
		index = argmin( abs.(  frequency_ghz.-frequency_table_ghz ) );
		if polarization == 'v'
			a = av[index];
			b = bv[index];
		else
			a = ah[index];
			b = bh[index];
		end

		α = a*r^b;
		return α;
   end

   """
	Empirical model for rain attenuation for frequencies between 1 and 400 GHz, circular polarization.
	Model uses closest frequency in the underlying data.
	One-way attenuation in dB/km.

	- M. A. Richards and J. A. Scheer and W. A. Holm, Principles of Modern Radar, SciTech Publishing, 2010.

	```julia
	rain_attenuation_db_per_km_circular_pol( 30, 20 )
	3.6755027981960815
	```

	# Arguments
	- 'frequency_ghz'     	The frequency of the propagating waves.
	- 'fall_rate_mm_hour'	The rain intensity [mm/h].
	"""
	function rain_attenuation_db_per_km_circular_pol( frequency_ghz, fall_rate_mm_hour )
		α_lin( pol ) = rain_attenuation_db_per_km( pol, frequency_ghz, fall_rate_mm_hour );
		αv = α_lin('v');
		αh = α_lin('h');

		α = 1/sqrt(2) * sqrt(αv^2+αh^2);

		return α;
	end

	"""
	Empirical model for rain attenuation for frequencies above 5 GHz.
	One-way attenuation in dB/km.

	- M. A. Richards and J. A. Scheer and W. A. Holm, Principles of Modern Radar, SciTech Publishing, 2010.

	```julia
	fog_attenuation_db_per_km_circular_pol( 10, 0.8, 23 );
	4.68976
	```

	# Arguments
	- 'frequency_ghz'	The frequency of the propagating waves.
	- 'M'				The watewater concentration in g/m³.
	- 'T_deg'			The air temperature in degree Celsius.
	"""
	function fog_attenuation_db_per_km( frequency_ghz, M, T_deg )
		T = T_deg;
		f = frequency_ghz

		α= M*( -1.347 + 0.66*f+(11.152/f) - 0.022*T )
		return α;
	end

	"""
	Empirical model for amospheric gaseous attenuation for frequencies in the range 1 - 1000 GHz.
	One-way attenuation in dB/km.

	- Rec. ITU-R P.676-12, Attenuation by atmospheric gases and related effects, ITU-R 2020.

	```julia
	using Plots

	frequency_ghz = collect(1:1000);

	attenuation_db_km = RadioPropagation.atmospheric_attenuation_db_per_km.( frequency_ghz );

	plot1 = plot(	frequency_ghz, attenuation_db_km,
					xlabel 	= "Frequency [GHz]",
					ylabel 	= "One-way loss [dB/km]",
					title  	= "Atmostpheric Attenuation.",
					#ylims 	= (),
					label 	= "Standard",
					yaxis=:log);
	```

	# Arguments
	- 'frequency_ghz'				The frequency of the propagating waves.
	- 'T_kelvin'					The absolute temperature in kelvin.
	- 'water_vapour_density_g_m³' 	The water vapur density in g/m³.
	- 'dry_air_pressure_h_pa' 		The dry air pressure in hpa.
	"""
	function atmospheric_attenuation_db_per_km( frequency_ghz, T_kelvin=288.15, water_vapour_density_g_m³=7.5, dry_air_pressure_h_pa=1013.25 )
		f = frequency_ghz;
		T = T_kelvin;
		ρ = water_vapour_density_g_m³;
		p = dry_air_pressure_h_pa;	# Dry air pressure [hPa].
		e = ρ*T/216.7;				# Water vapour partial pressure [hPa].
		θ = 300/T;

		# Spectroscopic data for oxygen attenuation.
		# 		  f0 			a1			a2		a3		a4 		a5 		a6
		tab1 = 	[ 50.474214		0.975 		9.651	6.690 	0.0		2.566 	6.850
				  50.987745		2.529 		8.653	7.170 	0.0		2.246 	6.800
				  51.503360		6.193 		7.709	7.640 	0.0		1.947 	6.729
				  52.021429		14.320 		6.819	8.110 	0.0		1.667 	6.640
				  52.542418		31.240 		5.983	8.580 	0.0		1.388 	6.526
				  53.066934		64.290 		5.201	9.060 	0.0		1.349 	6.206
				  53.595775		124.600 	4.474	9.550 	0.0		2.227 	5.085
				  54.130025		227.300 	3.800	9.960 	0.0		3.170 	3.750
				  54.671180		389.700 	3.182	10.370 	0.0		3.558 	2.654
				  55.221384		627.100 	2.618	10.890	0.0		2.560 	2.952
				  55.783815		945.300 	2.109	11.340	0.0		-1.172 	6.135
				  56.264774		543.400 	0.014	17.030	0.0		3.525 	-0.978
				  56.363399		1331.800 	1.654	11.890	0.0		-2.378	6.547
				  56.968211		1746.600 	1.255	12.230	0.0		-3.545	6.451
				  57.612486		2120.100 	0.910	12.620	0.0		-5.416	6.056
				  58.323877		2363.700 	0.621	12.950	0.0		-1.932	0.436
				  58.446588		1442.100	0.083	14.910	0.0		6.768 	-1.273
				  59.164204		2379.900 	0.387	13.530	0.0		-6.561 	2.309
				  59.590983		2090.700 	0.207	14.080	0.0		6.957 	-0.776
				  60.306056		2103.400 	0.207	14.150	0.0		-6.395 	0.699
				  60.434778		2438.000 	0.386	13.390	0.0		6.342 	-2.825
				  61.150562		2479.500 	0.621	12.920	0.0		1.014 	-0.584
				  61.800158		2275.900 	0.910	12.630	0.0		5.014 	-6.619
				  62.411220		1915.400 	1.255	12.170	0.0		3.029 	-6.759
				  62.486253		1503.000 	0.083	15.130	0.0		-4.499 	0.844
				  62.997984		1490.200 	1.654	11.740	0.0		1.856 	-6.675
				  63.568526		1078.000 	2.108	11.340	0.0		0.658 	-6.139
				  64.127775		728.700 	2.617	10.880	0.0		-3.036 	-2.895
				  64.678910		461.300 	3.181	10.380	0.0		-3.968 	-2.590
				  65.224078		274.000 	3.800	9.960   0.0		-3.528 	-3.680
				  65.764779		153.000 	4.473	9.550   0.0		-2.548 	-5.002
				  66.302096		80.400 		5.200	9.060   0.0		-1.660 	-6.091
				  66.836834		39.800 		5.982	8.580   0.0		-1.680	-6.393
				  67.369601		18.560 		6.818	8.110   0.0		-1.956 	-6.475
				  67.900868		8.172 		7.708	7.640   0.0		-2.216 	-6.545
				  68.431006		3.397 		8.652	7.170   0.0		-2.492	-6.600
				  68.960312		1.334 		9.650	6.690   0.0		-2.773 	-6.650
				  118.750334 	940.300		0.010	16.640	0.0		-0.439	0.079
				  368.498246	67.400		0.048	16.400	0.0		0.0		0.0
				  424.763020	637.700		0.044	16.400	0.0		0.0		0.0
				  487.249273	237.400		0.049	16.000	0.0		0.0		0.0
				  715.392902	98.100		0.145	16.000	0.0		0.0		0.0
				  773.839490	572.300		0.141	16.200	0.0		0.0		0.0
				  834.145546	183.100		0.145	14.700	0.0		0.0		0.0]

		f0a = @view tab1[:,1];
		a1  = @view tab1[:,2];
		a2  = @view tab1[:,3];
		a3  = @view tab1[:,4];
		a4  = @view tab1[:,5];
		a5  = @view tab1[:,6];
		a6  = @view tab1[:,7];

		# Spectroscopic data for water vapour attenuation.
		# 		 f0 		b1		b2		b3		b4 		b5 		b6
		tab2 = [ 22.235080 	.1079	2.144	26.38	.76		5.087	1.00
				 67.803960 	.0011	8.732	28.58	.69		4.930	.82
				 119.995940 .0007	8.353	29.48	.70		4.780	.79
				 183.310087 2.273	.668 	29.06	.77		5.022	.85
				 321.225630 .0470	6.179	24.04	.67		4.398	.54
				 325.152888 1.514	1.541	28.23	.64		4.893	.74
				 336.227764 .0010	9.825	26.93	.69		4.740	.61
				 380.197353 11.67	1.048	28.11	.54		5.063	.89
				 390.134508 .0045	7.347	21.52	.63		4.810	.55
				 437.346667 .0632	5.048	18.45	.60		4.230	.48
				 439.150807 .9098	3.595	20.07	.63		4.483	.52
				 443.018343 .1920	5.048	15.55	.60		5.083	.50
				 448.001085 10.41	1.405	25.64	.66		5.028	.67
				 470.888999 .3254	3.597	21.34	.66		4.506	.65
				 474.689092 1.260	2.379	23.20	.65		4.804	.64
				 488.490108 .2529	2.852	25.86	.69		5.201	.72
				 503.568532 .0372	6.731	16.12	.61		3.980	.43
				 504.482692 .0124	6.731	16.12	.61		4.010	.45
				 547.676440 .9785	.158	26.00	.70		4.500	1.00
				 552.020960 .1840	.158	26.00	.70		4.500	1.00
				 556.935985 497.0	.159	30.86	.69		4.552	1.00
				 620.700807 5.015	2.391	24.38	.71		4.856	.68
				 645.766085 .0067	8.633	18.00	.60		4.000	.50
				 658.005280 .2732	7.816	32.10	.69		4.140	1.00
				 752.033113 243.4	.396 	30.86	.68		4.352	.84
				 841.051732 .0134	8.177	15.90	.33		5.760	.45
				 859.965698 .1325	8.055	30.60	.68		4.090	.84
				 899.303175 .0547	7.914	29.85	.68		4.530	.90
				 902.611085 .0386	8.429	28.65	.70		5.100	.95
				 906.205957 .1836	5.110	24.08	.70		4.700	.53
				 916.171582 8.400	1.441	26.73	.70		5.150	.78
				 923.112692 .0079	10.293 	29.00	.70		5.000	.80
				 970.315022 9.009	1.919 	25.50	.64		4.940	.67
				 987.926764 134.6	.257 	29.85	.68		4.550	.90
				 1780.000000 17506. .952 	196.3	2.00 	24.15	5.00]

		f0b = @view tab2[:,1];
		b1  = @view tab2[:,2];
		b2  = @view tab2[:,3];
		b3  = @view tab2[:,4];
		b4  = @view tab2[:,5];
		b5  = @view tab2[:,6];
		b6  = @view tab2[:,7];

		S_o		= @. a1*10^(-7)*p*θ^3*exp( a2*(1-θ));
		S_w		= @. b1*10^(-1)*e*θ^3.5*exp(b2*(1-θ));

		Δf_o_	= @. a3*10^(-4) *(p*θ^(0.8-a4) + 1.1*e*θ);
		Δf_o	= @. sqrt( Δf_o_^2 + 2.25*10^(-6) );

		Δf_w_	= @. b3*10^(-4) *(p*θ^b4 + b5*e*θ^b6);
		Δf_w	= @. 0.535*Δf_w_ + sqrt( 0.217*Δf_w_^2 + (2.1316*10^-12*f0b^2)/ θ );

		δ_o		= @. (a5+a6*θ)*10^(-4) *(p+e)*θ^(0.8);
		δ_w		= 0;

		F_o		= @. f/f0a *( (Δf_o-δ_o*(f0a-f))/((f0a-f)^2 + Δf_o^2 ) + (Δf_o-δ_o*(f0a+f))/((f0a+f)^2 + Δf_o^2) );
		F_w		= @. f/f0b *( (Δf_w-δ_w*(f0b-f))/((f0b-f)^2 + Δf_w^2 ) + (Δf_w-δ_w*(f0b+f))/((f0b+f)^2 + Δf_w^2) );

		d		= 5.6*10^(-4) * (p+e)*θ^0.8;

		N_D_prime = f*p*θ^2 * ( (6.14*10^(-5) / ( d*(1+(f/d)^2) )) +  (1.4*10^(-12) *p*θ^1.5) / ( 1+1.9*10^(-5) *f^1.5 ) );

		N_d_prime_oxygen		= sum( S_o.*F_o) + N_D_prime;
		N_d_prime_water_vapour	= sum( S_w.*F_w );

		γ = 0.1820*f*(N_d_prime_oxygen+N_d_prime_water_vapour);
		return γ;
	end
end