##########################
### Variance estimator ###
##########################
abstract type VarEstimator{TI,TF} end

# Ref： https://github.com/stan-dev/math/blob/develop/stan/math/prim/mat/fun/welford_var_estimator.hpp
mutable struct WelfordVar{TI,TF} <: VarEstimator{TI,TF}
    n :: TI
    μ :: Vector{TF}
    M :: Vector{TF}
end

function reset!(wv::WelfordVar{TI,TF}) where {TI<:Integer,TF<:Real}
    wv.n = zero(TI)
    wv.μ .= zero(TF)
    wv.M .= zero(TF)
end

function add_sample!(wv::WelfordVar, s::AbstractVector)
    wv.n += 1
    δ = s .- wv.μ
    wv.μ .+= δ ./ wv.n
    wv.M .+= δ .* (s .- wv.μ)
end

# https://github.com/stan-dev/stan/blob/develop/src/stan/mcmc/var_adaptation.hpp
function get_var(wv::VarEstimator{TI,TF})::Vector{TF} where {TI<:Integer,TF<:Real}
    n, M = wv.n, wv.M
    @assert n >= 2 "Cannot get variance with only one sample"
    return (n / ((n + 5) * (n - 1))) .* M .+ 1e-3 * (5 / (n + 5))
end

abstract type CovarEstimator{TI<:Integer,TF<:Real} end

# NOTE: this naive covariance estimator is used only in testing
mutable struct NaiveCovar{TI,TF} <: CovarEstimator{TI,TF}
    n :: TI
    S :: Vector{Vector{TF}}
end

function add_sample!(nc::NaiveCovar, s::AbstractVector)
    nc.n += 1
    push!(nc.S, s)
end

function reset!(nc::NaiveCovar{TI,TF}) where {TI<:Integer,TF<:Real}
    nc.n = zero(TI)
    nc.S = Vector{Vector{TF}}()
end

function get_covar(nc::NaiveCovar{TI,TF})::Matrix{TF} where {TI<:Integer,TF<:Real}
    @assert nc.n >= 2 "Cannot get variance with only one sample"
    return Statistics.cov(nc.S)
end

# Ref: https://github.com/stan-dev/math/blob/develop/stan/math/prim/mat/fun/welford_covar_estimator.hpp
mutable struct WelfordCovar{TI<:Integer,TF<:Real} <: CovarEstimator{TI,TF}
    n :: TI
    μ :: Vector{TF}
    M :: Matrix{TF}
end

function reset!(wc::WelfordCovar{TI,TF}) where {TI<:Integer,TF<:Real}
    wc.n = zero(TI)
    wc.μ .= zero(TF)
    wc.M .= zero(TF)
end

function add_sample!(wc::WelfordCovar, s::AbstractVector)
    wc.n += 1
    δ = s .- wc.μ
    wc.μ .+= δ ./ wc.n
    wc.M .+= (s .- wc.μ) * δ'
end

# Ref: https://github.com/stan-dev/stan/blob/develop/src/stan/mcmc/covar_adaptation.hpp
function get_covar(wc::WelfordCovar{TI,TF})::Matrix{TF} where {TI<:Integer,TF<:Real}
    n, M = wc.n, wc.M
    @assert n >= 2 "Cannot get variance with only one sample"
    return (n / ((n + 5) * (n - 1))) .* M + 1e-3 * (5 / (n + 5)) * LinearAlgebra.I
end

################
### Adapters ###
################

abstract type PreConditioner <: AbstractAdapter end

struct UnitPreConditioner <: PreConditioner end

function Base.string(::UnitPreConditioner)
    return string([1.0])
end

struct DiagPreConditioner{TI<:Integer,TF<:Real} <: PreConditioner
    ve  :: VarEstimator{TI,TF}
    std :: Vector{TF}
end

function DiagPreConditioner(d::Integer)
    ve = WelfordVar(0, zeros(d), zeros(d))
    return DiagPreConditioner(ve, Vector(ones(d)))
end

function Base.string(dpc::DiagPreConditioner)
    return string(dpc.std)
end

function adapt!(dpc::DiagPreConditioner, θ, is_addsample::Bool, is_updatestd::Bool)
    if is_addsample
        add_sample!(dpc.ve, θ)
    end
    if is_updatestd
        var = get_var(dpc.ve)
        dpc.std .= sqrt.(var)
        reset!(dpc.ve)
        return true
    end
    return false
end

struct DensePreConditioner{TI<:Integer,TF<:Real} <: PreConditioner
    ce    :: CovarEstimator{TI,TF}
    covar :: Matrix{TF}
end

function DensePreConditioner(d::Integer)
    ce = WelfordCovar(0, zeros(d), zeros(d,d))
    # TODO: take use of the line below when we have an interface to set which pre-conditioner to use
    # ce = NaiveCovar(0, Vector{Vector{Float64}}())
    return DensePreConditioner(ce, LinearAlgebra.diagm(0 => ones(d)))
end

function Base.string(dpc::DensePreConditioner)
    return string(LinearAlgebra.diag(dpc.covar))
end

function adapt!(dpc::DensePreConditioner, θ, is_addsample::Bool, is_updatestd::Bool)
    if is_addsample
        add_sample!(dpc.ce, θ)
    end
    if is_updatestd
        covar = get_covar(dpc.ce)
        dpc.covar .= covar
        reset!(dpc.ce)
        return true
    end
    return false
end


# Copyright (c) 2018: Matthew Wilhelm & Matthew Stuber.
# This code is licensed under MIT license (see LICENSE.md for full details)
#############################################################################
# McCormick.jl
# A McCormick operator library in Julia
# See https://github.com/PSORLab/McCormick.jl
#############################################################################
# src/forward_operators/convex_increasing.jl
# Contains definitions of exp, exp2, exp10, expm1.
#############################################################################

for opMC in (:exp, :exp2, :exp10, :expm1)
   opMC_kernel = Symbol(String(opMC)*"_kernel")
   dop = diffrule(:Base, opMC, :midcv) # Replace with cv ruleset
   MCexp = quote
              xL = x.Intv.lo
              xU = x.Intv.hi
              xLc = z.lo
              xUc = z.hi
              del = xU - xL
              midcc, cc_id = mid3(x.cc, x.cv, xU)
              midcv, cv_id = mid3(x.cc, x.cv, xL)
              concave = xUc
              (xUc > xLc) && (concave = (xLc*(xU - midcc) + xUc*(midcc - xL))/del)
              convex = ($opMC)(midcv)
              concave_grad = mid_grad(x.cc_grad, x.cv_grad, cc_id)*(xUc - xLc)/del
              convex_grad = mid_grad(x.cc_grad, x.cv_grad, cv_id)*$dop
              convex, concave, convex_grad, concave_grad = cut(xLc, xUc, convex, concave, convex_grad, concave_grad)
              return MC{N, T}(convex, concave, z, convex_grad, concave_grad, x.cnst)
            end
    dop = diffrule(:Base, opMC, :(x.cv))
    dMCexp = quote
               xL = x.Intv.lo
               xU = x.Intv.hi
               xLc = z.lo
               xUc = z.hi
               del = xU - xL
               midcc = mid3v(x.cv, x.cc, xU)
               midcv = mid3v(x.cv, x.cc, xL)
               concave = xUc
               (xUc > xLc) && (concave = (xLc*(xU-midcc) + xUc*(midcc-xL))/del)
               convex = ($opMC)(midcv)
               convex_grad = ($dop)*x.cv_grad
               concave_grad = ((xUc - xLc)/del)*x.cc_grad
               return MC{N, Diff}(convex, concave, z, convex_grad, concave_grad, x.cnst)
              end

      @eval @inline ($opMC_kernel)(x::MC{N, T}, z::Interval{Float64}) where {N, T<:Union{NS,MV}} = $MCexp
      @eval @inline ($opMC_kernel)(x::MC{N, Diff}, z::Interval{Float64}) where {N} = $dMCexp
      @eval @inline ($opMC)(x::MC) = ($opMC_kernel)(x, ($opMC)(x.Intv))
end


# =================================================================
# Laub-Loomis model
# See https://easychair.org/publications/paper/gjfh
# =================================================================

using Reachability, MathematicalSystems, LazySets, TaylorIntegration
using Reachability: solve

# Equations of motion
# We write the function such that the operations are either unary or binary:
@taylorize function laubloomis!(t, x, dx)
    dx[1] = 1.4*x[3] - 0.9*x[1]
    dx[2] = 2.5*x[5] - 1.5*x[2]
    dx[3] = 0.6*x[7] - 0.8*(x[2]*x[3])
    dx[4] = 2 - 1.3*(x[3]*x[4])
    dx[5] = 0.7*x[1] - (x[4]*x[5])
    dx[6] = 0.3*x[1] - 3.1*x[6]
    dx[7] = 1.8*x[6] - 1.6*(x[2]*x[7])
    return dx
end

function laubloomis(; T=20.0, W=0.01, plot_vars=[0, 4],
                      property=(t,x)->x[4] < 4.5,
                      project_reachset=true)

    # equations, x' = f(x(t))
    𝐹 = BlackBoxContinuousSystem(laubloomis!, 7)

    X0c = [1.2, 1.05, 1.5, 2.4, 1.0, 0.1, 0.45]
    X0 = Hyperrectangle(X0c, fill(W, 7))

    # instantiate the IVP
    𝑃 = InitialValueProblem(𝐹, X0)

    # general options
    𝑂 = Options(:T=>T, :plot_vars=>plot_vars, :property=>property,
                :project_reachset=>project_reachset, :mode=>"check")

    return (𝑃, 𝑂)
end


using GaussianRandomFields;
using Plots;
using Plots.PlotMeasures;
using Colors;
using Netpbm;
using FileIO;

width = 1024
height = 1024
type = "CirculantEmbedding"

function generaterandomfield(type)
    if type == "Cholesky"
        cov = CovarianceFunction(2, Matern(1 / 2, 5 / 4))
        pts = range(0, stop = 1, length = 101)
        grf = GaussianRandomField(cov, Cholesky(), pts, pts)
    elseif type == "Spectral"
        cov = CovarianceFunction(2, Linear(1))
        pts = range(0, stop = 1, length = 51)
        grf = GaussianRandomField(cov, Spectral(), pts, pts)
    elseif type == "KarhunenLoeve"
        cov = CovarianceFunction(2, Whittle(.1))
        pts = range(0, stop = 1, length = 201)
        grf = GaussianRandomField(cov, KarhunenLoeve(500), pts, pts)
    elseif type == "CirculantEmbedding"
        cov = CovarianceFunction(2, Exponential(.5))
        pts = range(0, stop = 1, length = 1001)
        grf = GaussianRandomField(cov, CirculantEmbedding(), pts, pts, minpadding = 2001)
    end
end

function normalize(a)
    (minval, maxval) = extrema(a)
    valrange = maxval - minval

    function normalize0(v)
        (v - minval) / valrange
    end

    normalize0.(a)
end

println("Generating random field...")
@time grf = generaterandomfield(type)

println("Saving random field...")
data = sample(grf)

plot_options = (legend = false, c = :grays, framestyle = :none, size = (width,height), margin = -2mm, bottom_margin = -10mm, left_margin = -13mm)
heatmap_plot = heatmap(data; plot_options...)
contour_plot = contourf(data; plot_options...)
contour8_plot = contourf(data; levels = 8, plot_options...)
png(heatmap_plot, "heatmap.png")
png(contour_plot, "contour.png")
png(contour8_plot, "contour8.png")

normalized_data = reverse(normalize(data), dims = 1)
pgm_image_data = convert(Array{Gray}, normalized_data)
pgm_file = File(format"PGMBinary", "heatmap.pgm")
Netpbm.save(pgm_file, pgm_image_data)

println("Done")


using MLDatasets
using HDF5
using StatsBase
using LinearAlgebra

# load full training set
train_x, train_y = MNIST.traindata()
# load full test set
# test_x,  test_y  = MNIST.testdata()

# u = 40 # from "VISUALIZING DATA USING T-SNE"
u = 50 # from "Accelerating t-SNE using Tree-Based Algorithms"
n = size(train_x, 3)

X = train_x
X = reshape(X, :, n) # vectorize images
labels = train_y
if false
    k = 1024
    subsample = sample(1:n, k, replace = false)
    X = X[:, subsample]
    labels = labels[subsample]
end

# reduce dimensionality using PCA to
U, S, Vt = svd(X)
V = Vt'
d = 50
X = V[1:d, :]
@. X *= S[1:d]

using StochasticNeighborhoodEmbeddings
using StochasticNeighborhoodEmbeddings: OptimizationParameters,
            optimized_symmetrized_neighbor, initialize_embedding!

total_iter = 64
early_iter = 64
max_iter = total_iter - early_iter

beta_1 = 0.9
beta_2 = 0.99
weight_decay = 0
lr = 3e-1
dofast = false

early_beta_1 = 0.5
early_beta_2 = 0.9
early_lr = 3e-1
early_compression = 0
early_exaggeration = 8

params = OptimizationParameters(max_iter = max_iter, beta_1 = beta_1,
        beta_2 = beta_2, lr = lr, weight_decay = weight_decay,
        early_iter = early_iter, early_beta_1 = early_beta_1, early_beta_2 = early_beta_2,
        early_lr = early_lr, early_compression = early_compression,
        early_exaggeration = early_exaggeration, dofast = dofast,
        verbose = true, debias = false)

Y = 1e-2randn(params.d, size(X, 2))
P = optimized_symmetrized_neighbor(X, u, params.dofast)
G = zeros(params.d, size(X, 2)) # gradient of Y
G² = zeros(params.d, size(X, 2)) # second moments of gradient
exaggerations = 8:-1:1
for i in exaggerations # slowly reducing early_exaggeration
    params.early_exaggeration = i
    @time initialize_embedding!(Y, G, G², P, params)
end
params.early_iter = 0
params.max_iter = 1024 - length(exaggerations) * early_iter
params.early_exaggeration = 0
@time tsne!(Y, G, G², P, params) # fine-tuning

f = h5open("MNIST_TSNE_full.h5", "w")
f["Y"] = Y
f["X"] = X
f["labels"] = labels
f["max_iter"] = max_iter
f["lr"] = lr
f["beta_1"] = beta_1
f["beta_2"] = beta_2
f["early_iter"] = early_iter
f["early_lr"] = early_lr
f["early_beta_1"] = early_beta_1
f["early_beta_2"] = early_beta_2
f["early_compression"] = early_compression
f["early_exaggeration"] = early_exaggeration
close(f)

# using Plots
# doplot = true
# if doplot
#     YC = [Y[:, ==(i).(labels)] for i in 1:10]
#     scatter()
#     for (i, yc) in enumerate(YC)
#         scatter!(yc[1, :], yc[2, :], label = "$(i-1)")
#     end
#     gui()
# end


## dcf.jl  Tests related to decision cost functions

using Base.Test

x = readtable("ru.2009.table.gz", separator='\t')
tnt = TNT(x)
tar, non = tnt.tar, tnt.non
ruc = roc(tnt, collapse=false)
r = roc(tnt)

## scalar
d = DCF(0.01, 1, 10)            # nist 2008
@test_approx_eq oeff(d) 1/9.9
@test_approx_eq peff(d) 1/10.9
@test_approx_eq plo(d) -log(9.9)

lo = plo(d)
b = 0.03291547867387523
@test_approx_eq ber(x,lo) b
@test_approx_eq ber(tnt,lo) b
@test_approx_eq ber(tar,non,lo) b
@test_approx_eq ber(ruc,lo) b

m = 0.03173675448189074
@test_approx_eq minber(x,lo) m
@test_approx_eq minber(tnt,lo) m
@test_approx_eq minber(tar,non,lo) m
@test_approx_eq minber(r,lo) m

## set a global DCF as well
setdcf(cmiss=10)                # evalita 2009

for norm in (false, true)
    c = 0.035877871754524004 * 10^norm
    @test_approx_eq dcf(x,d=d,norm=norm) c
    @test_approx_eq dcf(tnt,d=d,norm=norm) c
    @test_approx_eq dcf(tar,non,d=d,norm=norm) c
    @test_approx_eq dcf(ruc,d=d,norm=norm) c

    m = 0.034593062385260914 * 10^norm
    @test_approx_eq mindcf(x,d=d,norm=norm) m
    @test_approx_eq mindcf(tnt,d=d,norm=norm) m
    @test_approx_eq mindcf(tar,non,d=d,norm=norm) m
    @test_approx_eq mindcf(r,d=d,norm=norm) m

    ## global DCF
    c = 0.23094942466561763084 * 2^norm
    @test_approx_eq dcf(x,norm=norm) c
    @test_approx_eq dcf(tnt,norm=norm) c
    @test_approx_eq dcf(tar,non,norm=norm) c
    @test_approx_eq dcf(ruc,norm=norm) c

    m = 0.21843024685287174003 * 2^norm
    @test_approx_eq mindcf(x,norm=norm) m
    @test_approx_eq mindcf(tnt,norm=norm) m
    @test_approx_eq mindcf(tar,non,norm=norm) m
    @test_approx_eq mindcf(ruc,norm=norm) m
end

## array
lo = collect(-7:0.01:7)
dc = DCF(sigmoid(lo), 1, 1)

@test_approx_eq_eps ber(x,lo) ber(ruc,lo) 1e-5
@test ber(tnt,lo) == ber(tar,non,lo)

@test minber(x,lo) == minber(tnt,lo) == minber(tar,non,lo) == minber(r,lo)
@test all(minber(r,lo) .≤ ber(ruc,lo))

for norm in (false, true)
    @test_approx_eq dcf(x,d=dc,norm=norm) dcf(ruc,d=dc,norm=norm)
    @test dcf(x, d=dc, norm=norm) == dcf(tnt, d=dc, norm=norm) == dcf(tar, non, d=dc, norm=norm)
    @test mindcf(x, d=dc, norm=norm) == mindcf(tnt, d=dc, norm=norm) == mindcf(tar, non, d=dc, norm=norm) == mindcf(r, d=dc, norm=norm)
    @test all(mindcf(r, d=dc, norm=norm) .≤ dcf(ruc, d=dc, norm=norm))
end


####
# UI Oracle tests
####
facts("suppFcnTest UI") do
	srand(8675309); data = rand(200, 2)
	w = UIOracle(data, [0, 0], [1, 1], .1, .2)

	zstar, ustar = suppFcn([1, 1], w, :Max)
	@fact zstar --> roughly(1.8772156628616448, 1e-8)
	@fact ustar[1] --> roughly(0.9385081140320037, 1e-8)
	@fact ustar[2] --> roughly(0.9387075488296409, 1e-8)

	zstar, ustar = suppFcn([1, -1], w, :Max)
	@fact zstar --> roughly(0.8932039312511011, 1e-8)
	@fact ustar[1] --> roughly(0.9492096614168394, 1e-8)
	@fact ustar[2] --> roughly(0.056005730165738316, 1e-8)

	zstar, ustar = suppFcn([-1, 1], w, :Max)
	@fact zstar --> roughly(0.8767942540553981, 1e-8)
	@fact ustar[1] --> roughly(0.06360683171372103, 1e-8)
	@fact ustar[2] --> roughly(0.9404010857691192, 1e-8)

	#extra tests for Min and 0 entries
	zstar, ustar = suppFcn([-1, 1], w, :Min)
	@fact zstar --> roughly(-0.893203928668029, 1e-8)
	@fact ustar[1] --> roughly(0.9492096601291875, 1e-8)
	@fact ustar[2] --> roughly(0.056005731461158545, 1e-8)

	zstar, ustar = suppFcn([0, 1], w, :Min)
	@fact zstar --> roughly(0.001943343938739852, 1e-8)
	@fact ustar[1] --> roughly(0.5715291128670335, 1e-8)
	@fact ustar[2] --> roughly(0.001943343938739852, 1e-8)

	zstar, ustar = suppFcn([0, 1], w, :Max)
	@fact zstar --> roughly(0.9981094005216729, 1e-8)
	@fact ustar[1] --> roughly(0.5715291128670335, 1e-8)
	@fact ustar[2] --> roughly(0.9981094005216729, 1e-8)

end

facts("portTest UI") do
	srand(8675309)
	data = rand(500, 2)
	w = UIOracle(data, [0., 0.], [1., 1.], .1, .2)
	portTest(w, 0.07068906144892116, [0.5039487187951147, 0.4960512812048853])
end


using DataFrames
using Distributions
using PyPlot
using LinearAlgebra
function makeData()
    groupOne = rand(MvNormal([10.0, 10.0], 10.0 * Matrix{Float64}(I, 2,2)), 100)
    groupTwo = rand(MvNormal([0.0, 0.0], 10 * Matrix{Float64}(I, 2,2)), 100)
    groupThree = rand(MvNormal([15.0, 0.0], 10.0 * Matrix{Float64}(I, 2,2)), 100)
    return hcat(groupOne, groupTwo, groupThree)'
end

data = makeData()
scatter(data[1:100, 1], data[1:100, 2], color="blue")
scatter(data[101:200, 1], data[101:200, 2], color="red")
scatter(data[201:300, 1], data[201:300, 2], color="green")

using StatsBase

mutable struct kMeans
    x::DataFrames.DataFrame
    k::Int
end

function euclidean(sourcePoint, destPoint)
    sum = 0
    for i in 1:length(sourcePoint)
        sum += (destPoint[i] - sourcePoint[i]) ^ 2
    end
    dist = sqrt(sum)
    return dist
end

function minkowski(sourcePoint, destPoint)
    sum = 0
    for i in 1:length(sourcePoint)
        sum += abs(destPoint[i] - sourcePoint[i])
    end
    return sum
end

function calcDist(sourcePoint::Array, destPoint::Array; method="euclidean")

    if length(sourcePoint) != length(destPoint)
        error("The lengths of two arrays are different.")
        return
    end

    if method == "euclidean"
        return euclidean(sourcePoint, destPoint)
    elseif method == "minkowski"
        return minkowski(sourcePoint, destPoint)
    end
end

function classify(kMeans::kMeans)
    dataPointsNum = size(kMeans.x, 1)
    estimatedClass = Array{Int, dataPointsNum}
    direct_sample(1:kMeans.k, estimatedClass)

    while true
        # update representative points
        representativePoints = []
        for representativeIndex in 1:kMeans.k
            groupIndex = find(estimatedClass .== representativeIndex)
            groupData = kMeans.x[groupIndex, :]

            # TODO: check the return type of colwise
            representativePoint = [ valArray[1] for valArray in colwise(mean, groupData) ]
            push!(representativePoints, representativePoint)
        end

        # update group belonging
        tempEstimatedClass = Array{Int}(dataPointsNum)
        for dataIndex in 1:dataPointsNum
            dataPoint = Array(kMeans.x[dataIndex, :])
            distances = Array{Float64}(kMeans.k)
            for representativeIndex in 1:kMeans.k
                distances[representativeIndex] = calcDist(dataPoint, representativePoints[representativeIndex])
            end

            # TODO: check the existence of argmin
            tempEstimatedClass[dataIndex] = sortperm(distances)[1]
        end

        if estimatedClass == tempEstimatedClass
            break
        end
        estimatedClass = tempEstimatedClass
    end
    return estimatedClass
end

    kmeansElbow = kMeans(DataFrame(data), 3)

    predictedClassElbow = classify(kmeansElbow)


#=
BFSLevelsMultiSource:
- Julia version: 1.6.2
- Author: samuel
- Date: 2021-09-10
=#

using SuiteSparseGraphBLAS
using SparseArrays
using LinearAlgebra

#A is the input matrix, S is the source matrix, and n is the number of nodes, s is the number of sources in the matrix S, Output: Levels matrix
function ms_bfsl(A, S, n, s)

    #this vector allows to check if a value is already present inside the distance matrix,
    #if it's present we don't want it to be changed
    levels = GBVector{Int64}(n)
        for i = 1:n
            levels[i] = i
        end

    #distance matrix
    distance = GBMatrix{Int64}(s, n)
            for i = 1:s
                for j = 1:n
                    distance[i, j] = 0
                end
            end

    print("STARTING FROM SOURCE NODE\n\n")
    for level = 1:n-1

            for i = 1:s
                for j = 1:n
                    if(S[i, j]!=0 && distance[i, j] ∉ levels) #*
                    distance[i, j] = level #*
                    end
                end
            end

            print("FRONTIER\n\n")
            print(S)
            print("\n\n")

            S = mul(S, A, Semirings.LOR_LAND, mask=distance, desc=Descriptors.RC)
            print("DISTANCE\n\n")
            print(distance)
            print("\n\n")
            print("ITERATION COMPLETED")
            print("\n\n")

    end

    print("END")
    return distance

end


# PotentialEnergySurfaces.jl
# Define quadratic diabatic potential energy surfaces
# Determine adiabatic potential energy surfaces from diagonalising the
#     diabatic hamiltonian potential.

# """
# Function to define quadratic potential energy surface
# ------------
# A_0 = real constant
# A_1 = real coefficient of R
# A_2 = real coefficient of R^2
# ------------
# dpl = quadratic function representing site to the left of the bond centre
# dpr = quadratic function representing site to the right of the bond centre
# """
# function diabatic_potential(A_0::Float64, A_1::Float64, A_2::Float64)
#     dpl = function (R) return A_0 + A_1*R + A_2*R^2 end
#     dpr = function (R) return dpl(-R) end
#     return dpl, dpr
# end

"""
Function to define linear potential energy surface
Simple avoided crossing
------------
alpha = on-site EPC constant
tau = neigbouring sites EPC constant
overlap = <phi_l(R)|phi_r(R)>
------------
dpl = linear function representing site to the left of the bond centre
dpr = linear function representing site to the right of the bond centre
"""
function diabatic_potential(alpha::Float64, tau::Float64, overlap::Function)
    dpl = function (R) return -alpha*R/2*(1-abs(overlap(R))^2)-tau*(overlap(R)+conj(overlap(R))) end
    dpr = function (R) return dpl(-R) end
    dpx = function (R) return -tau*(overlap(R)^2+1) end
    return dpl, dpr, dpx
end
function diabatic_potential(alpha::Float64, tau::Float64, overlap::Float64)
    dpl = function (R) return -alpha*R/2*(1-abs(overlap)^2)-tau*(overlap+conj(overlap)) end
    dpr = function (R) return dpl(-R) end
    dpx = function (R) return -tau*(overlap^2+1) end
    return dpl, dpr, dpx
end
function diabatic_potential(alpha::Float64,alpha2::Float64, tau::Float64, overlap::Float64)
    dpl = function (R) return -R/2*(alpha-alpha2*abs(overlap)^2)-tau*(overlap+conj(overlap)) end
    dpr = function (R) return R/2*(alpha2-alpha*abs(overlap)^2)-tau*(overlap+conj(overlap)) end
    dpx = function (R) return -tau*(overlap^2+1) end
    return dpl, dpr, dpx
end
function diabatic_potential(alpha::Float64,alpha2::Function, tau::Float64, overlap::Function)
    dpl = function (R) return -R/2*(alpha-alpha2(R)*abs(overlap(R))^2)-tau*(overlap(R)+conj(overlap(R))) end
    dpr = function (R) return R/2*(alpha2(-R)-alpha*abs(overlap(-R))^2)-tau*(overlap(-R)+conj(overlap(-R))) end
    dpx = function (R) return -tau*(overlap(R)^2+1) end
    return dpl, dpr, dpx
end

"""
Function to diagonalise the diabatic potential
------------
e_11 = local energy of site 1
e_22 = local energy of site 2
J = coupling between sites 1 and 2
------------
H_d = [e_11 J
        J   e_22]
"""
function H_diabatic(e_11::Float64,e_22::Float64,J::Float64)
    H_d = [e_11 J;J e_22]; return H_d
end

"""
Function to diagonalise the diabatic hamiltonian and obtain the
unitary transformation matrix
------------
H_d = diabatic Matrix
------------
U_nk = Unitary transformation matrix
"""
function diagonalise(H_d::Array, test::Bool=false)
    U_nk = eigvecs(H_d)
    if test
        U = U_nk'U_nk
        return U, U_nk
    else
        return U_nk
    end
end

"""
Function to generate adiabatic potential energy surfaces
--------------
dpl = diabatic potential surface from site left of centre
dpr = diabatic potential surface from site right of centre
--------------
ape = excited adiabatic potential energy surface
apg = ground adiabatic potential energy surface
"""
function adiabatic_potential(dpl::Function,dpr::Function,J_lr::Float64)
    ape = function (R) return 0.5*(dpl(R)+dpr(R)) + 0.5*sqrt((dpl(R)-dpr(R))^2+4*abs(J_lr)^2) end
    apg = function (R) return 0.5*(dpl(R)+dpr(R)) - 0.5*sqrt((dpl(R)-dpr(R))^2+4*abs(J_lr)^2) end
    return apg, ape
end

"""
Function to generate adiabatic hamiltonian matrix
---------------
H_d = diabatic hamiltonian
---------------
H_ad = adiabatic hamiltonian
"""
function H_adiabatic(H_d::Array)
    lambdas = eigvals(H_d)
    H_ad = diagm(lambdas)
end


"""
    BetaBinomial(n,α,β)

A *Beta-binomial distribution* is the compound distribution of the [`Binomial`](@ref) distribution where the probability of success `p` is distributed according to the [`Beta`](@ref). It has three parameters: `n`, the number of trials and two shape parameters `α`, `β`

```math
P(X = k) = {n \\choose k} B(k + \\alpha, n - k + \\beta) / B(\\alpha, \\beta),  \\quad \\text{ for } k = 0,1,2, \\ldots, n.
```

```julia
BetaBinomial(n, α, β)      # BetaBinomial distribution with n trials and shape parameters α, β

params(d)       # Get the parameters, i.e. (n, α, β)
ntrials(d)      # Get the number of trials, i.e. n
```

External links:

* [Beta-binomial distribution on Wikipedia](https://en.wikipedia.org/wiki/Beta-binomial_distribution)
"""
struct BetaBinomial{T<:Real} <: DiscreteUnivariateDistribution
    n::Int
    α::T
    β::T

    BetaBinomial{T}(n::Integer, α::T, β::T) where {T <: Real} = new{T}(n, α, β)
end

function BetaBinomial(n::Integer, α::T, β::T; check_args=true) where {T <: Real}
    check_args && @check_args(BetaBinomial, n >= zero(n) && α >= zero(α) && β >= zero(β))
    return BetaBinomial{T}(n, α, β)
end

BetaBinomial(n::Integer, α::Real, β::Real) = BetaBinomial(n, promote(α, β)...)
BetaBinomial(n::Integer, α::Integer, β::Integer) = BetaBinomial(n, float(α), float(β))

@distr_support BetaBinomial 0 d.n

#### Conversions
function convert(::Type{BetaBinomial{T}}, n::Int, α::S, β::S) where {T <: Real, S <: Real}
    BetaBinomial(n, T(α), T(β))
end
function convert(::Type{BetaBinomial{T}}, d::BetaBinomial{S}) where {T <: Real, S <: Real}
    BetaBinomial(d.n, T(d.α), T(d.β), check_args=false)
end

#### Parameters

ntrials(d::BetaBinomial) = d.n

params(d::BetaBinomial) = (d.n, d.α, d.β)
partype(::BetaBinomial{T}) where {T} = T

#### Properties

mean(d::BetaBinomial) = (d.n * d.α) / (d.α + d.β)
function var(d::BetaBinomial)
    n, α, β = d.n, d.α, d.β
    numerator = n * α * β * (α + β + n)
    denominator = (α + β)^2 * (α + β + 1)
    return numerator / denominator
end

function skewness(d::BetaBinomial)
    n, α, β = d.n, d.α, d.β
    t1 = (α + β + 2n) * (β - α) / (α + β + 2)
    t2 = sqrt((1 + α +β) / (n * α * β * (n + α + β)))
    return t1 * t2
end

function kurtosis(d::BetaBinomial)
    n, α, β = d.n, d.α, d.β
    alpha_beta_sum = α + β
    alpha_beta_product = α * β
    numerator = ((alpha_beta_sum)^2) * (1 + alpha_beta_sum)
    denominator = (n * alpha_beta_product) * (alpha_beta_sum + 2) * (alpha_beta_sum + 3) * (alpha_beta_sum + n)
    left = numerator / denominator
    right = (alpha_beta_sum) * (alpha_beta_sum - 1 + 6n) + 3*alpha_beta_product * (n - 2) + 6n^2
    right -= (3*alpha_beta_product * n * (6 - n)) / alpha_beta_sum
    right -= (18*alpha_beta_product * n^2) / (alpha_beta_sum)^2
    return (left * right) - 3
end

function logpdf(d::BetaBinomial, k::Real)
    n, α, β = d.n, d.α, d.β
    _insupport = insupport(d, k)
    _k = _insupport ? round(Int, k) : 0
    logbinom = - log1p(n) - logbeta(_k + 1, n - _k + 1)
    lognum   = logbeta(_k + α, n - _k + β)
    logdenom = logbeta(α, β)
    result = logbinom + lognum - logdenom
    return _insupport ? result : oftype(result, -Inf)
end

entropy(d::BetaBinomial) = entropy(Categorical(pdf.(Ref(d),support(d))))
median(d::BetaBinomial) = median(Categorical(pdf.(Ref(d),support(d)))) - 1
mode(d::BetaBinomial) = argmax(pdf.(Ref(d),support(d))) - 1
modes(d::BetaBinomial) = modes(Categorical(pdf.(Ref(d),support(d)))) .- 1

quantile(d::BetaBinomial, p::Float64) = quantile(Categorical(pdf.(Ref(d), support(d))), p) - 1

#### Sampling

rand(rng::AbstractRNG, d::BetaBinomial) =
    rand(rng, Binomial(d.n, rand(rng, Beta(d.α, d.β))))


module MyLinInterp

  export LinearInterpolation

  immutable LinearInterpolation # Inplementation as a type class
    grid::Array
    vals::Array
  end

  function call(a::LinearInterpolation, xs)
    map(xs) do x
      #a.grid[1] <= x <= a.grid[end] || throw(DomainError()) # throw an error if input x is out of domain
        if x < a.grid[1]
            return a.vals[1]
        elseif x > a.grid[end]
            return a.vals[end]
        end

      i = searchsortedlast(a.grid, x) # find the first index i in grid where grid[i] >= x
      if i == length(a.grid)
        return a.vals[end]
      else
        return a.vals[i] + (a.vals[i+1] - a.vals[i]) / (a.grid[i+1] - a.grid[i]) * (x - a.grid[i])
      end
    end
  end

end


function moi_sensorloc(optimizer, seed, n; verbose = false, test = false, scalar = false)

    rng = Random.MersenneTwister(seed)
    MOI.empty!(optimizer)
    if test
        @test MOI.is_empty(optimizer)
    end
    # Generate randomized problem data
    m, x_true, a, d, d_bar = sensorloc_data(seed, n)

    # Decision variable
    nvars = ProxSDP.sympackedlen(n + 2)
    X = MOI.add_variables(optimizer, nvars)
    Xsq = Matrix{MOI.VariableIndex}(undef, n + 2, n + 2)
    ProxSDP.ivech!(Xsq, X)
    Xsq = Matrix(LinearAlgebra.Symmetric(Xsq, :U))
    vov = MOI.VectorOfVariables(X)
    cX = MOI.add_constraint(optimizer, vov, MOI.PositiveSemidefiniteConeTriangle(n + 2))

    # Constraint with distances from anchors to sensors
    for j in 1:n
        for k in 1:m
            if scalar
                MOI.add_constraint(optimizer,
                    MOI.ScalarAffineFunction([
                            MOI.ScalarAffineTerm(a[k][1]*a[k][1], Xsq[1,1]),
                            MOI.ScalarAffineTerm(a[k][2]*a[k][2], Xsq[2,2]),
                            MOI.ScalarAffineTerm(-2 * a[k][1], Xsq[1, j+2]),
                            MOI.ScalarAffineTerm(-2 * a[k][2], Xsq[2, j+2]),
                            MOI.ScalarAffineTerm(1.0, Xsq[j+2, j+2]),
                        ], 0.0), MOI.EqualTo(d_bar[k, j]^2))
            else
                MOI.add_constraint(optimizer,
                MOI.VectorAffineFunction(
                    MOI.VectorAffineTerm.([1], [
                        MOI.ScalarAffineTerm(a[k][1]*a[k][1], Xsq[1,1]),
                        MOI.ScalarAffineTerm(a[k][2]*a[k][2], Xsq[2,2]),
                        MOI.ScalarAffineTerm(-2 * a[k][1], Xsq[1, j+2]),
                        MOI.ScalarAffineTerm(-2 * a[k][2], Xsq[2, j+2]),
                        MOI.ScalarAffineTerm(1.0, Xsq[j+2, j+2]),
                    ]), -[d_bar[k, j]^2]), MOI.Zeros(1))
            end
        end
    end

    # Constraint with distances from sensors to sensors
    count, count_all = 0, 0
    for i in 1:n
        for j in 1:i - 1
            count_all += 1
            if Random.rand(rng) > 0.9
                count += 1
                if scalar
                    MOI.add_constraint(optimizer,
                        MOI.ScalarAffineFunction([
                            MOI.ScalarAffineTerm(1.0, Xsq[i+2,i+2] ),
                            MOI.ScalarAffineTerm(1.0, Xsq[j+2,j+2] ),
                            MOI.ScalarAffineTerm(-2.0, Xsq[i+2,j+2]),
                        ], 0.0), MOI.EqualTo(d[i, j]^2))
                else
                    MOI.add_constraint(optimizer,
                        MOI.VectorAffineFunction(
                            MOI.VectorAffineTerm.([1],[
                                MOI.ScalarAffineTerm(1.0, Xsq[i+2,i+2] ),
                                MOI.ScalarAffineTerm(1.0, Xsq[j+2,j+2] ),
                                MOI.ScalarAffineTerm(-2.0, Xsq[i+2,j+2]),
                            ]), -[d[i, j]^2]), MOI.Zeros(1))
                end
            end
        end
    end
    if verbose
        @show count_all, count
    end
    if scalar
        MOI.add_constraint(optimizer, MOI.SingleVariable(Xsq[1, 1]), MOI.EqualTo(1.0))
        MOI.add_constraint(optimizer, MOI.SingleVariable(Xsq[1, 2]), MOI.EqualTo(0.0))
        MOI.add_constraint(optimizer, MOI.SingleVariable(Xsq[2, 1]), MOI.EqualTo(0.0))
        MOI.add_constraint(optimizer, MOI.SingleVariable(Xsq[2, 2]), MOI.EqualTo(1.0))
    else
        MOI.add_constraint(optimizer, MOI.VectorAffineFunction(
            MOI.VectorAffineTerm.([1],[MOI.ScalarAffineTerm(1.0, Xsq[1, 1])]), [-1.0]), MOI.Zeros(1))
        MOI.add_constraint(optimizer, MOI.VectorAffineFunction(
            MOI.VectorAffineTerm.([1],[MOI.ScalarAffineTerm(1.0, Xsq[1, 2])]), [ 0.0]), MOI.Zeros(1))
        MOI.add_constraint(optimizer, MOI.VectorAffineFunction(
            MOI.VectorAffineTerm.([1],[MOI.ScalarAffineTerm(1.0, Xsq[2, 1])]), [ 0.0]), MOI.Zeros(1))
        MOI.add_constraint(optimizer, MOI.VectorAffineFunction(
            MOI.VectorAffineTerm.([1],[MOI.ScalarAffineTerm(1.0, Xsq[2, 2])]), [-1.0]), MOI.Zeros(1))
    end

    objf_t = [MOI.ScalarAffineTerm(0.0, Xsq[1, 1])]
    if false
        objf_t = [MOI.ScalarAffineTerm(1.0, Xsq[i,i]) for i in 1:n + 2]
    end
    MOI.set(optimizer, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(), MOI.ScalarAffineFunction(objf_t, 0.0))

    MOI.set(optimizer, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    MOI.optimize!(optimizer)

    objval = MOI.get(optimizer, MOI.ObjectiveValue())

    stime = -1.0
    try
        stime = MOI.get(optimizer, MOI.SolveTimeSec())
    catch
        println("could not query time")
    end

    Xsq_s = MOI.get.(optimizer, MOI.VariablePrimal(), Xsq)

    verbose && sensorloc_eval(n, m, x_true, Xsq_s)

    rank = -1
    status = 0
    if MOI.get(optimizer, MOI.TerminationStatus()) == MOI.OPTIMAL
        status = 1
    end
    return (objval, stime, rank, status)
end


function filtr(model::DynamicDiscreteModel,data::Array{Int,1})
	T=length(data)
	dx,dy=size(model.mu)

	filter=Array(Float64,dx,T)
	filter[:,1]=model.mu[:,data[1]]/sum(model.mu[:,data[1]])
	for t=2:T
		rho=0.0
		for jx=1:dx
			filter[jx,t]=0
			for ix=1:dx
				filter[jx,t]+=model.m[ix,data[t-1],jx,data[t]]*filter[ix,t-1]
			end
			rho+=filter[jx,t]
		end
		for jx=1:dx
			filter[jx,t]/=rho
		end
	end

	smoother=Array(Float64,dx,T)
	smoother[:,T]=1
	for t=T-1:-1:1
		rho=0.0
		for ix=1:dx
			smoother[ix,t]=0
			for jx=1:dx
				smoother[ix,t]+=model.m[ix,data[t],jx,data[t+1]]*smoother[jx,t+1]
			end
			rho+=smoother[ix,t]
		end
		for ix=1:dx
			smoother[ix,t]/=rho
		end
	end

	conditional=Array(Float64,dx,dx,T)
	for t=1:T-1
		tempsum=0.0
		for jx=1:dx
			for ix=1:dx
				conditional[ix,jx,t]=model.m[ix,data[t],jx,data[t+1]]
				conditional[ix,jx,t]*=filter[ix,t]
				conditional[ix,jx,t]*=smoother[jx,t+1]
				tempsum+=conditional[ix,jx,t]
			end
		end
		for jx=1:dx
			for ix=1:dx
				#note that there is no risk of division by zero
				#reason: this is called only on _observed_ data which must thus have non-zero probability
				conditional[ix,jx,t]=conditional[ix,jx,t]/tempsum
			end
		end
	end
	filter,smoother,conditional
end


using DataFrames

export
    LogisticRegressionModel,
    fit!,
    predict

mutable struct LogisticRegressionModel{T, S}
    y::Vector{T}
    xs::AbstractArray{S}
    argv::Vector{Float64}
    n::Integer
end

function LogisticRegressionModel(df::DataFrame, label::Symbol, features::Vector{Symbol})
    n = nrow(df)
    y = df[!, label]
    xs = Matrix(df[!, features])
    β = zeros(length(features))
    return LogisticRegressionModel(y, xs, β, n)
end

function LogisticRegressionModel(df::DataFrame, label::Symbol, feature::Symbol)
    n = nrow(df)
    y = df[!, label]
    x = df[!, feature]
    β = zeros(1)
    return LogisticRegressionModel(y, x, β, n)
end

sigmoid(z::Real) = 1.0 / (1.0 + exp(-z))

z(model::LogisticRegressionModel, xs::Vector{<:Real}) = model.argv' * xs

predict(model::LogisticRegressionModel, xs::Vector{<:Real}) = sigmoid(z(model, xs))

predict(model::LogisticRegressionModel, xs::Real) = sigmoid(z(model, [xs]))

function log_likelyhood(model::LogisticRegressionModel)
    xs = model.xs
    y = model.y

    ll = 0
    for (i, yⁱ) in enumerate(y)
        xsⁱ = xs[i, :]
        ll += yⁱ * log(sigmoid(z(model, xsⁱ)))
        ll += (1 - yⁱ) * log(1 - sigmoid(z(model, xsⁱ)))
    end

    return ll ./ model.n
end

function ∇log_likelyhood(model::LogisticRegressionModel)
    # ∇ll = ∑(yⁱ - σ(zⁱ))xⱼⁱ
    xs = model.xs
    y = model.y

    ∇ll = zeros(length(model.argv))
    for (i, yⁱ) in enumerate(y)
        xsⁱ = xs[i, :]
        ∇ll .+= (yⁱ - sigmoid(z(model, xsⁱ))) .* xsⁱ
    end

    return ∇ll ./ model.n
end

function gradient_ascend(
    model::LogisticRegressionModel,
    η::Real,
    atol::Real,
    show::Bool
)
    β = model.argv
    while (ll = log_likelyhood(model)) < atol
        show && println("Log Likelyhood: $ll")
        β .+= η .* ∇log_likelyhood(model)
    end

    return β
end

function fit!(
    model::LogisticRegressionModel;
    method=gradient_ascend,
    η::Real=1e-2,
    atol::Real=-1e-4,
    show=false
)
    β = method(model, η, atol, show)
    model.argv .= β

    return model
end


import Pkg

Pkg.add("Polynomials")

using Polynomials
using .Polynoms

function diff(p::Polynomialx, x; ord=2)
    Q = eval_poly(x,p)
    Qn = eval_polyN(x,p,ord)
    return Q-Qn
end

function eval_poly(x, A)
    Q = first(A)
    for a in @view A[2:end]
        Q=Q*x+a
    end
    return Q
end


@testset "weeks test"  begin
    @test isa(Weeks(s -> 1 / s), Weeks)
    @test isa(WeeksErr(s -> 1 / s), WeeksErr)

    fl = Weeks(s -> 1 / s^2)

    @test isapprox(fl(1), 1)
    @test isapprox(fl([1 2; 3 4]), [1 2; 3 4])

    fl = Weeks(s -> s / (1 + s^2), 64)
    @test isapprox(fl([1 2; 3 4]), broadcast(cos,[1.0 2.0; 3.0 4.0]))

    fle = WeeksErr(s -> s/(1+s^2),  64)
    @test isapprox(fle([1 2; 3 4])[1], cos.([1 2; 3 4]))

    ## Optim is broken.
    # e1 = abs(fl(10.0) - cos(10.0))
    # InverseLaplace.optimize(fl, 10.0)
    # e2 = abs(fl(10.0) - cos(10.0))
    # @test e2 < e1

    # e1 = abs(fle(10.0)[1] - cos(10.0))
    # optimize(fle,10.0)
    # e2 = abs(fl(10.0)[1] - cos(10.0))
    # @test e2 < e1

    fle = WeeksErr(s -> s/(1+s^2),  64)
    (t1,e1) = fle(10.0)
    @test isapprox([t1,e1], [-0.8390715290763104,1.6069233509475447e-11]; atol = 1e-9)

    setNterms(fle, 8)
    e1 = abs(fle(10.0)[1] - cos(10.0))
    setNterms(fle, 80)
    e2 = abs(fle(10.0)[1] - cos(10.0))
    @test e2 < e1

    fle = WeeksErr(s -> s/(1+s^2),  64)
    c1 = copy(fle.coefficients)
    @test string(fle) == "WeeksErr{Float64}(Nterms=64,sigma=1.0,b=1.0)"
    setparameters(fle,2.0,2.0,80)
    @test string(fle) == "WeeksErr{Float64}(Nterms=80,sigma=2.0,b=2.0)"

    # Check that Laguerre coefficients have been recomputed
    @test c1 != fle.coefficients

    ### Complex

    function Fcomplex(s)
        # Laplace domain
        α = complex(-0.3, 6.0)
        return 1 / (s - α)
    end

    function fcomplex(t)
        # Time domain
        α = complex(-0.3, 6.0)
        return exp(α * t)
    end

    let Fc = Weeks(Fcomplex, 1024, datatype=Complex),  trange = range(0.0, stop=30.0, length=1000)
        @test isapprox(Fc.(trange), fcomplex.(trange), atol=0.001)
    end
end # @testset


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

function calcNestPosForEquator(z, z_abs, scaled_phi)
    jp = floor(Integer, NSIDE_MAX * (0.5 + scaled_phi - z * 0.75))
    jm = floor(Integer, NSIDE_MAX * (0.5 + scaled_phi + z * 0.75))

    idfp = div(jp, NSIDE_MAX) # in {0,4}
    idfm = div(jm, NSIDE_MAX)

    local face_num
    if idfp == idfm
        face_num = (idfp % 4) + 4
    elseif idfp < idfm
        face_num = (idfp % 4)
    else
        face_num = (idfm % 4) + 8
    end

    ix = mod(jm, NSIDE_MAX)
    iy = NSIDE_MAX - mod(jp, NSIDE_MAX) - 1

    (ix, iy, face_num)
end

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

function calcNestPosForPole(z, z_abs, scaled_phi)
    ntt = floor(Integer, scaled_phi)
    if ntt >= 4
        ntt = 3
    end

    tp = scaled_phi - ntt
    tmp = sqrt(3 * (1 - z_abs)) # in ]0,1]

    jp = floor(Integer, NSIDE_MAX * tp * tmp)
    jm = floor(Integer, NSIDE_MAX * (1 - tp) * tmp)

    # Clip jp and jm
    jp = jp < NSIDE_MAX - 1 ? jp : NSIDE_MAX - 1
    jm = jm < NSIDE_MAX - 1 ? jm : NSIDE_MAX - 1

    local ix, iy, face_num
    if z >= 0
        face_num = ntt # in {0,3}
        ix = NSIDE_MAX - jm - 1
        iy = NSIDE_MAX - jp - 1
    else
        face_num = ntt + 8 # in {8,11} */
        ix =  jp
        iy =  jm
    end

    (ix, iy, face_num)
end

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

"""
    ang2pixNest(resol::Resolution, theta, phi) -> Integer

Return the index of the pixel which contains the point with
coordinates (`theta`, the colatitude, and `phi`, the longitude), in
radians, for a Healpix map with pixels in nested order. Note that
pixel indexes are 1-based (this is Julia)!
"""
function ang2pixNest(resol::Resolution, theta, phi)

    nside = resol.nside
    local ix, iy, face_num

    z          = cos(theta)
    z_abs      = abs(z)
    scaled_phi = mod2pi(phi) / (π / 2) # in [0,4[

    if z_abs ≤ 2//3
        (ix, iy, face_num) = calcNestPosForEquator(z, z_abs, scaled_phi)
    else
        (ix, iy, face_num) = calcNestPosForPole(z, z_abs, scaled_phi)
    end

    (ix_hi, ix_low) = divrem(ix, 128)
    (iy_hi, iy_low) = divrem(iy, 128)

    ipf = ((x2pix[ix_hi + 1] + y2pix[iy_hi + 1]) * (128^2)
           + (x2pix[ix_low + 1] + y2pix[iy_low + 1]))
    ipf = floor(Integer, ipf / ((NSIDE_MAX / nside) ^ 2))

    # Add 1 to have a 1-based index
    ipf + face_num * nside * nside + 1

end

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

function calcRingPosForEquator(resol::Resolution, z, z_abs, tt)

    jp = floor(Integer, resol.nside * (0.5 + tt - z * 0.75))
    jm = floor(Integer, resol.nside * (0.5 + tt + z * 0.75))

    ir = resol.nside + 1 + jp - jm
    kshift = (mod(ir, 2) == 0) ? 1 : 0

    nl4 = resol.nsideTimesFour;

    local ip = div(jp + jm - resol.nside + kshift + 1, 2) + 1
    if ip > nl4
	    ip = ip - nl4
    end

    resol.ncap + nl4 * (ir - 1) + ip
end

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

function calcRingPosForPole(resol::Resolution, z, z_abs, tt)

    tp = tt - floor(tt)
    tmp = sqrt(3. * (1. - z_abs))

    jp = floor(Integer, resol.nside * tp * tmp )
    jm = floor(Integer, resol.nside * (1 - tp) * tmp)

    ir = jp + jm + 1
    ip = floor(Integer, tt * ir) + 1
    if ip > 4ir
	    ip -= 4ir
    end

    if z ≤ 0
	    resol.numOfPixels - 2ir * (ir + 1) + ip
    else
        2ir * (ir - 1) + ip
    end
end

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

"""
    ang2vec(theta, phi) -> Array{Float64}

Given a direction in the sky with colatitude `theta` and longitude
`phi` (in radians), return an array of 3 elements containing the
`x`, `y`, and `z` components of the one-length vector pointing to
that direction.
"""
function ang2vec(theta, phi)
    (0 ≤ theta ≤ π) || throw(DomainError(theta, "Invalid value of theta"))

    sintheta = sin(theta)
    return [sintheta * cos(phi), sintheta * sin(phi), cos(theta)]
end

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

"""
    vec2ang(x, y, z) -> (Number, Number)

Given a vector (not necessarily normalized) whose Cartesian components are `x`,
`y`, and `z`, return a pair (`theta`, `phi`) containing the colatitude `theta`
and the longitude `phi` (in radians) of the direction in the sky the vector is
pointing at.
"""
function vec2ang(x, y, z)
    norm = sqrt(x^2 + y^2 + z^2)
    theta = acos(z / norm)
    phi = atan(y, x)
    if phi < 0
        phi += 2π
    end

    return (theta, phi)
end

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

"""
    ang2pixRing(resol::Resolution, theta, phi) -> Integer

Return the index of the pixel which contains the point with
coordinates (`theta`, the colatitude, and `phi`, the longitude), in
radians, for a Healpix map with pixels in ring order. Note that pixel
indexes are 1-based (this is Julia)!
"""
function ang2pixRing(resol::Resolution, theta, phi)

    z = cos(theta)
    z_abs = abs(z)
    scaled_phi = mod2pi(phi) / (π / 2) # in [0,4[

    if z_abs ≤ 2//3
        calcRingPosForEquator(resol, z, z_abs, scaled_phi)
    else
        calcRingPosForPole(resol, z, z_abs, scaled_phi)
    end
end

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

"""
    pix2angNest(resol::Resolution, pixel) -> (Float64, Float64)

Given the (1-based) index of a pixel in a Healpix map in nested
order, return a pair containing the (`colatitude`, `longitude`) angles
corresponding to its center, both expressed in radians.
"""
function pix2angNest(resol::Resolution, pixel)

    jrll = [ 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4 ]
    jpll = [ 1, 3, 5, 7, 0, 2, 4, 6, 1, 3, 5, 7 ]

    floatNside = Float64(resol.nside)
    fact1 = 1. / (3.0 * floatNside^2)
    fact2 = 2. / (3.0 * floatNside)

    # face number in {0,11} and pixel number within the face
    faceNum, ipf = divrem(pixel - 1, resol.pixelsPerFace)

    ip_trunc, ip_low = divrem(ipf, 1024)
    ip_hi, ip_med = divrem(ip_trunc, 1024)

    ix = 1024 * pix2x[ip_hi+1] + 32 * pix2x[ip_med+1] + pix2x[ip_low+1]
    iy = 1024 * pix2y[ip_hi+1] + 32 * pix2y[ip_med+1] + pix2y[ip_low+1]

    # Transforms this in (horizontal, vertical) coordinates
    jrt = ix + iy # 'vertical' in {0,2*(nside-1)}
    jpt = ix - iy # 'horizontal' in {-nside+1,nside-1}

    jr = jrll[faceNum+1] * resol.nside - jrt - 1
    local nr = resol.nside # Equatorial region (the most frequent)
    z = (2resol.nside - jr) * fact2
    local kshift = Int(mod(jr - resol.nside, 2))
    if jr < resol.nside
        # North polar cap
	nr = jr
	z = 1. - nr^2 * fact1
	kshift = 0
    elseif jr > 3 * resol.nside
        # South polar cap
	nr = resol.nsideTimesFour - jr
	z = - 1. + nr^2 * fact1
	kshift = 0
    end

    local jp = div(jpll[faceNum+1] * nr + jpt + 1 + kshift, 2)
    if jp > resol.nsideTimesFour
        jp = jp - resol.nsideTimesFour
    end
    if jp < 1
        jp = jp + resol.nsideTimesFour
    end

    (acos(z), (jp - (kshift + 1) * 0.5) * (π / (2nr)))
end

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

"""
    pix2ringpos(resol::Resolution, pixel)

Given the (1-based) index of a pixel in a Healpix map in ring
order, return a pair of numbers (n, i, j) whose meaning is the following:

- `n` can be one of the symbols `:northcap`, `:equator`, or
  `:southcap`, representing the region of the sky
- `i` is the ring index, from 1 to 4NSIDE - 1
- `j` is the pixel-in-ring index
"""
function pix2ringpos(resol::Resolution, pixel)
    # Any reference to equations in this routine refers to Gorski et al. (2005)
    if pixel ≤ resol.ncap
        # North polar cap

        p_h = pixel / 2 # Defined in Gorsky et al. (2005) before Eq. (2)
        floor_p_h = floor(p_h)
        # Eq. (2)
        i = floor(Int, sqrt(p_h - sqrt(floor_p_h))) + 1
        # Eq. (3)
        j = pixel - 2i * (i - 1)

        (:northcap, i, j)
    elseif pixel ≤ resol.nsideTimesTwo * (5resol.nside + 1)
        ip    = pixel - resol.ncap - 1
        # Eq. (6) - ring counts from the North pole; resol.nside is
        # the number of pixels in the North Polar ring
        i = floor(Int, ip / resol.nsideTimesFour) + resol.nside
        # Eq. (7) - zero-based index of the pixel within this ring
        j  = Int(mod(ip, resol.nsideTimesFour)) + 1
        (:equator, i, j)
    else
        # South polar cap

        ip = resol.numOfPixels - pixel + 1
        p_h = ip / 2
        floor_p_h = floor(p_h)
        i = floor(Int, sqrt(p_h - sqrt(floor_p_h))) + 1 # counted from S. pole
        j = Int(4 * i + 1 - (ip - 2i * (i - 1)))

        (:southcap, i, j)
    end
end

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

"""
    pix2angRing(resol::Resolution, pixel) -> (Float64, Float64)

Given the (1-based) index of a pixel in a Healpix map in ring
order, return a pair containing the (`colatitude`, `longitude`) angles
corresponding to its center, both expressed in radians.
"""
function pix2angRing(resol::Resolution, pixel)
    fact1 = 1.5 * resol.nside
    fact2 = 3.0 * resol.pixelsPerFace

    # Any reference to equations in this routine refers to Gorski et al. (2005)
    cap, i, j = pix2ringpos(resol, pixel)
    if cap == :northcap
        # Colatitude: Eq. (4); longitude: Eq. (5)
        return (acos(1 - i^2 / fact2),
                    (Float64(j) - 0.5) * π / (2i))
    elseif pixel ≤ resol.nsideTimesTwo * (5resol.nside + 1)
        # Equatorial belt

        # Eq. (9) - this equals 1 if i + resol.nside is odd, 1/2
        # otherwise. It is used to convert j into a longitude (since
        # pixel centers in odd rings are shifted with respect to
        # centers in even rings)
        s_half = 0.5 * (1 + mod(Float64(i + resol.nside), 2))

        # Colatitude: Eq. (8) in disguise, latitude: Eq. (9)
        return (acos((resol.nsideTimesTwo - i) / fact1),
                    (Float64(j) - s_half) * π / (2resol.nside))
    else
        # South Polar cap

        # The pixels in this cap are handled like the ones in the
        # North Polar cap, except that we must flip the value of "ip".
        return (acos(-1 + i^2 / fact2),
                    (float(j) - 0.5) * π / (2i))
    end
end

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

vec2pixNest(res::Resolution, x, y, z) = ang2pixNest(res, vec2ang(x, y, z)...)
vec2pixRing(res::Resolution, x, y, z) = ang2pixRing(res, vec2ang(x, y, z)...)
pix2vecNest(res::Resolution, pixel) = ang2vec(pix2angNest(res, pixel)...)
pix2vecRing(res::Resolution, pixel) = ang2vec(pix2angRing(res, pixel)...)

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

function ringWeightPath(datadir, nside)
    joinpath(datadir, @sprintf("weight_ring_n%05d.fits", nside))
end

function readWeightRing(fileName, nside)
    f = FITSIO.fits_open_table(fileName)
    try
        weights = Array(Float64, 2 * nside)
        FITSIO.fits_read_col(f, 1, 1, 1, weights)

        return weights
    finally
        FITSIO.fits_close_file(f)
    end
end

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

function pixelWindowPath(datadir, nside)
    joinpath(datadir, @sprintf("pixel_window_n%04d.fits", nside))
end

function readPixelWindowT(fileName, nside)
    f = FITSIO.fits_open_table(fileName)
    try
        pixwin = Array(Float64, FITSIO.fits_get_num_rows(f))
        FITSIO.fits_read_col(f, 1, 1, 1, pixwin)

        return pixwin
    finally
        FITSIO.fits_close_file(f)
    end
end

function readPixelWindowP(fileName, nside)
    local pixwinT
    local pixwinP

    f = FITSIO.fits_open_table(fileName)
    try
        pixwinT = Array(Float64, FITSIO.fits_get_num_rows(f))
        pixwinP = Array(Float64, FITSIO.fits_get_num_rows(f))
        FITSIO.fits_read_col(f, 1, 1, 1, pixwinT)
        FITSIO.fits_read_col(f, 2, 1, 1, pixwinP)

        return (pixwinT, pixwinP)
    finally
        FITSIO.fits_close_file(f)
    end
end

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

"""
    ang2pix{T, O <: Order}(map::Map{T, O}, theta::Real, phi::Real)

Convert the direction specified by the colatitude `theta` (∈ [0, π])
and the longitude `phi` (∈ [0, 2π]) into the index of the pixel in the
Healpix map `map`.
"""
function ang2pix(map::OrderedMap{T, Ring}, theta, phi) where {T}
    ang2pixRing(map.resolution, Float64(theta), Float64(phi))
end

function ang2pix(map::OrderedMap{T, Nest}, theta, phi) where {T}
    ang2pixNest(map.resolution, Float64(theta), Float64(phi))
end

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

"""
    pix2ang{T, O <: Order}(map::Map{T, O}, ipix) -> (Float64, Float64)

Return the pair (`theta`, `phi`), where `theta` is the colatitude and
`phi` the longitude of the direction of the pixel center with index
`ipix`.
"""
function pix2ang(map::OrderedMap{T, Ring}, ipix) where {T}
    pix2angRing(map.resolution, ipix)
end

function pix2ang(map::OrderedMap{T, Nest}, ipix) where {T}
    pix2angNest(map.resolution, ipix)
end

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


# bands

using CairoMakie

f = Figure(resolution = (600, 400), backgroundcolor = :honeydew,)
Axis(f[1, 1])

xs = 1:0.2:10
ys_low = -0.2 .* sin.(xs) .- 0.25
ys_high = 0.2 .* sin.(xs) .+ 0.25

band!(xs, ys_low, ys_high, color=:maroon1)
band!(xs, ys_low .- 1, ys_high .-1, color = :purple1)

# bands ribbon

let
    x = LinRange(-10, 10, 200)
    fig = Figure(resolution = (600, 400))
    ax = Axis(fig[1, 1], xlabel = "x", ylabel = "y")
    band!(x, sin.(x), sin.(x) .+ 1; color = (:cadetblue2, 0.4))
    band!(x, cos.(x), 1 .+ cos.(x); color = (:red, 0.2))
    display(fig)
end;


"""
    simulate_continuous(i, t, η)

Simulate a simple (linear) birth and death process starting at time zero with
`i` units. The stochastic process is observed continuously for `t` time units.
Rates are given by parameter `η = [λ, μ]` where `λ` is the birth rate and `μ` is
the death rate.
"""
function simulate_continuous(
  i::I,
  t::F,
  η::Vector{F}
) where {
  I <: Integer,
  F <: AbstractFloat
}
  ω = η[1] + η[2]
  p = η[1] / ω

  if t < floatmin(F)
    return (zero(Int), zero(Int), zero(F), zero(Float64), zero(Int), [], Int(i),
            [])
  end

  tot_births = 0
  tot_deaths = 0
  integrated_jump = zero(F)
  sum_log_n = zero(Float64)

  # pre-allocate vectors. If needed, we will grow the vector by doubling its
  # size each time.
  len = 100
  waiting_time = zeros(F, len)
  increment = zeros(Int, len)

  popul_size = i
  h = 0
  total_time = zero(F)
  while true
    # sample next event time from a negative exponential distribution with rate
    # n * (λ + μ)
    τ = - log(rand(F)) / (popul_size * ω)
    total_time += τ

    if total_time > t
      integrated_jump += popul_size * (t - total_time + τ)
      break
    end

    if h == len - 1
      # not enough space to store this event
      len *= 2
      resize!(waiting_time, len)
      resize!(increment, len)
    end

    h += 1

    integrated_jump += popul_size * τ
    sum_log_n += log(popul_size)
    waiting_time[h] = τ

    # sample event type with probability (λ / (λ + μ), μ / (λ + μ))
    if rand(F) <= p
      tot_births += 1
      increment[h] = 1
      popul_size += 1
    else
      tot_deaths += 1
      increment[h] = -1
      popul_size -= 1

      if popul_size == 0
        break
      end
    end
  end

  resize!(waiting_time, h)
  resize!(increment, h)

  (tot_births, tot_deaths, integrated_jump, sum_log_n, h, waiting_time, Int(i),
   increment)
end

"""
    simulate_discrete(i, k, u, η)

Simulate a simple (linear) process starting at time zero with `i` units. The
stochastic process is observed every `u` time units until it reaches `u * k`
time. Rates are given by parameter `η = [λ, μ]` where `λ` is the birth rate and
`μ` is the death rate.
"""
function simulate_discrete(
  i::I,
  k::I,
  u::F,
  η::Vector{F}
)::Vector{I} where {
  I <: Integer,
  F <: AbstractFloat
}
  if k == 0
    return [i]
  end

  ω = η[1] + η[2]
  p = η[1] / ω

  event_time = zero(F)
  popul_size = i

  sim_size = zeros(I, k + 1)
  sim_size[1] = popul_size

  counter = 1
  t = u
  while true
    # sample event time from a negative exponential distribution with rate
    # n * (λ + μ)
    event_time -= log(rand(F)) / F(popul_size * ω)

    # new event might be after various null observations
    while (event_time > t) && (counter <= k)
      counter += 1
      t = counter * u
      sim_size[counter] = popul_size
    end

    # sample event type with probability (λ / (λ + μ), μ / (λ + μ))
    popul_size += (rand(F) <= p) ? I(1) : -I(1)

    if (counter > k) || (popul_size == 0)
      break
    end
  end

  sim_size
end

"""
    rand_continuous(n, i, t, η)

Simulate `n` simple (linear) birth and death processes with the same rate
parameter `η = [λ, μ]`, where `λ` is the birth rate and `μ` is the death rate.
Starting at time `0` with `i` units, observe continuously each process until
time `t` is reached.

For simulating only one process, use `rand_continuous(i, t, η)`.
"""
function rand_continuous(
  n::Integer,
  i::Integer,
  t::Real,
  η::Vector{T}
) where {
  T <: Real
}
  if n < 1
    error("Number of simulations 'n' is less than or equal to zero.")
  end

  if i < 1
    error("Initial population size 'i' is less than or equal to zero.")
  end

  if t < 0
    error("Time 't' is negative.")
  end

  if η[1] < 0
    error("Birth rate 'η[1]' is negative.")
  end

  if η[2] < 0
    error("Death rate 'η[2]' is negative.")
  end

  if (η[1] < floatmin(float(η[1]))) && (η[2] < floatmin(float(η[2])))
    error("Birth rate 'η[1]' and death rate 'η[2]' are jointly zero.")
  end

  F = typeof(float((η[1] - η[2]) * t))

  [ObservationContinuousTime{F}(
   simulate_continuous(i, F(t), [F(η[1]), F(η[2])])...) for s = 1:n]
end

"""
    rand_continuous(i, t, η)

Simulate one simple (linear) birth and death process with rate parameter
`η = [λ, μ]`, where `λ` is the birth rate and `μ` is the death rate. Starting at
time `0` with `i` units, observe continuously the process until time `t` is
reached.

For simulating `n` processes, use `rand_continuous(n, i, t, η)`.
"""
function rand_continuous(
  i::Integer,
  t::Real,
  η::Vector{T}
) where {
  T <: Real
}

  if i < 1
    error("Initial population size 'i' is less than or equal to zero.")
  end

  if t < 0
    error("Time 't' is negative.")
  end

  if η[1] < 0
    error("Birth rate 'η[1]' is negative.")
  end

  if η[2] < 0
    error("Death rate 'η[2]' is negative.")
  end

  if (η[1] < floatmin(float(η[1]))) && (η[2] < floatmin(float(η[2])))
    error("Birth rate 'η[1]' and death rate 'η[2]' are jointly zero.")
  end

  F = typeof(float((η[1] - η[2]) * t))

  ObservationContinuousTime{F}(
  simulate_continuous(i, F(t), [F(η[1]), F(η[2])])...)
end

"""
    rand_discrete(n, i, k, u, η)

Simulate `n` simple (linear) birth and death processes with the same rate
parameter `η = [λ, μ]`, where `λ` is the birth rate and `μ` is the death rate.
Starting at time `0` with `i` units, observe each process for a total of `k`
times, each separated by the same lag `u`.

For simulating only one process, use `rand_discrete(i, k, u, η)`.
"""
function rand_discrete(
  n::Integer,
  i::Integer,
  k::Integer,
  u::Real,
  η::Vector{T}
) where {
  T <: Real
}
  if n < 1
    error("Number of simulations 'n' is less than or equal to zero.")
  end

  if i < 1
    error("Initial population size 'i' is less than or equal to zero.")
  end

  if k < 0
    error("Number of observations 'k' is negative.")
  end

  if u < floatmin(float(u))
    error("Time lag 'u' is less than or equal to zero.")
  end

  if η[1] < 0
    error("Birth rate 'η[1]' is negative.")
  end

  if η[2] < 0
    error("Death rate 'η[2]' is negative.")
  end

  if (η[1] < floatmin(float(η[1]))) && (η[2] < floatmin(float(η[2])))
    error("Birth rate 'η[1]' and death rate 'η[2]' are jointly zero.")
  end

  I = typeof(i)
  F = typeof(float((η[1] - η[2]) * u))

  state = zeros(I, k + 1, n)

  for s = 1:n
    state[:, s] = simulate_discrete(i, I(k), F(u), [F(η[1]), F(η[2])])
  end

  ObservationDiscreteTimeEven{F, I}(Int(n), Int(k), F(u), state)
end

"""
    rand_discrete(i, k, u, η)

Simulate one simple (linear) birth and death process with rate parameter
`η = [λ, μ]`, where `λ` is the birth rate and `μ` is the death rate. Starting at
time `0` with `i` units, observe each process for a total of `k` times, each
separated by the same lag `u`.

For simulating `n` processes, use `rand_discrete(n, i, k, u, η)`.
"""
function rand_discrete(
  i::Integer,
  k::Integer,
  u::Real,
  η::Vector{T}
) where {
  T <: Real
}
  rand_discrete(1, i, k, u, η)
end


# tensor product filtering for dimension D discrete signal.
function ranges2collection( x_ranges::Vector{LinRange{T,L}},
                            ::Val{D}) where {T,D,L}
    # set up.
    @assert !isempty(x_ranges)
    @assert length(x_ranges) == D
    N_array = collect( length(x_ranges[d]) for d = 1:D )
    N = prod(N_array)
    sz_N = tuple(N_array...)

    # Position.
    X_nD = Array{Vector{T},D}(undef,sz_N)
    for 𝑖 in CartesianIndices(sz_N)
        X_nD[𝑖] = Vector{T}(undef,D)

        for d = 1:D
            X_nD[𝑖][d] = x_ranges[d][𝑖[d]]
        end
    end

    return X_nD
end

"""
    convertcompactdomain(x::T, a::T, b::T, c::T, d::T)::T

converts compact domain x ∈ [a,b] to compact domain out ∈ [c,d].
"""
function convertcompactdomain(x::T, a::T, b::T, c::T, d::T)::T where T <: Real

    return (x-a)*(d-c)/(b-a)+c
end

function convertcompactdomain(x::Vector{T}, a::Vector{T}, b::Vector{T}, c::Vector{T}, d::Vector{T})::Vector{T} where T <: Real

    return collect( convertcompactdomain(x[i], a[i], b[i], c[i], d[i]) for i = 1:length(x) )
end


"""
`spy(A; nargs...)` → `Plot`

Description
============

Plots the sparsity pattern for the given matrix `A`.
This means that a scatterplot that resembles the matrix is drawn,
in which only the pixel for non-zero elements of the matrix are set.

If the parameters `width` and `height` are not explicitly specified,
then the function will attempt to preserve the aspect ratio of the matrix,
while also attempting to fit the resulting plot withing the bounding box
specified by `maxwidth` and `maxheight`

Usage
======

    spy(A; maxwidth = 70, maxheight = 40, title = "Sparsity Pattern", width = 0, height = 0, border = :solid, margin = 3, padding = 1, color = :automatic, labels = true, canvas = BrailleCanvas)

Arguments
==========

- **`A`** : The matrix of interest, which non-zero elements should be drawn

- **`maxwidth`** : Maximal number of characters per row that should be used for plotting

- **`maxheight`** : Maximal number of rows that should be used for plotting.

- **`title`** : Text to display on the top of the plot.

- **`width`** : Number of characters per row that should be used for plotting. `0` stands for automatic

- **`height`** : Number of rows that should be used for plotting. `0` stands for automatic

- **`border`** : The style of the bounding box of the plot. Supports `:solid`, `:bold`, `:dashed`, `:dotted`, `:ascii`, and `:none`.

- **`margin`** : Number of empty characters to the left of the whole plot.

- **`padding`** : Space of the left and right of the plot between the labels and the canvas.

- **`color`** : Color of the drawing. Can be any of `:blue`, `:red`, `:yellow`

- **`labels`** : Can be used to hide the labels by setting `labels=false`.

- **`canvas`** : The type of canvas that should be used for drawing.

Returns
========

A plot object of type `Plot{T<:Canvas}`

Author(s)
==========

- Dominique (Github: https://github.com/dpo)
- Christof Stocker (Github: https://github.com/Evizero)
- Jake Bolewski (Github: https://github.com/jakebolewski)

Examples
=========

    julia> spy(sprandn(50, 120, .05))

                              Sparsity Pattern
        ┌────────────────────────────────────────────────────────────┐
      1 │⢀⠀⠒⠀⢄⠂⠀⡀⠀⠀⠀⠰⠀⠐⠀⠀⠀⠀⠂⠀⠀⠀⠠⠰⠀⠀⠀⡀⠄⠀⠠⠀⠀⠀⠀⠀⡀⠀⠀⠀⡀⠐⠀⠐⠀⠀⡀⠀⠀⠀⠀⠀⠀⠂⡀⠀⡀⠠⠀⠀│ > 0
        │⠀⠀⠀⠀⠀⢀⢀⠀⠀⠀⠀⠐⠄⠀⢀⠂⠀⠠⠀⠉⠀⠀⠀⠀⢐⠀⠀⠀⠀⠀⠠⠐⠄⠄⠀⠀⠀⠐⠐⠈⢁⠀⠀⠀⠀⠀⠀⠀⠠⠀⠐⡐⠀⠄⠀⠀⠈⠀⠀⠠│ < 0
        │⡀⠀⠀⠀⡀⠊⠀⠀⠀⢀⢀⠀⠄⠀⠂⡠⠂⠐⠐⠀⠀⠄⠀⠀⠀⠄⠀⠀⠀⠀⡀⠀⠀⠁⠀⠀⠀⠀⠀⠀⠠⠀⠂⠀⠠⠀⠀⠀⠆⠄⠈⠀⢀⠀⠀⠀⠐⠄⠀⠀│
        │⠀⠀⠀⠀⠀⠀⠀⢠⠀⢀⠀⠀⠀⢀⠀⠄⠀⠀⠀⠂⠀⠀⠀⠇⠀⢀⠀⠀⠁⠁⠀⠄⠄⠀⠀⠁⠁⠀⡈⠀⠁⠄⠀⠀⠄⠀⠀⠀⠀⠀⠀⠀⠀⠐⠀⠀⠀⠀⠀⠀│
        │⠀⠀⠀⠀⠀⠀⠀⠀⠀⡐⢀⠀⠀⠀⠐⠂⠁⢀⠠⠀⠀⠀⠀⠀⠁⠀⠀⠀⠠⠀⠀⠀⠀⢀⠀⢄⠀⠀⠀⠐⠀⠀⠈⠀⠀⠐⠀⠄⠂⠠⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
        │⠠⠀⠠⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⠀⠀⠀⠀⠀⠀⠐⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠐⠀⠀⠠⠈⠀⠀⠀⠀⢀⠀⠀⠄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
        │⠐⠀⠉⠐⠀⠀⠁⠀⠀⠀⠀⠀⠀⠌⠀⠀⠀⠄⠀⠀⠀⠀⠀⠀⠀⠐⠀⠀⠀⠀⠀⠌⠀⠀⠄⠀⠀⡀⠀⠀⠀⠀⠀⢀⠈⠄⠀⡀⠀⠀⠀⠀⢀⠀⠀⠀⠀⠀⠀⠀│
        │⠂⠀⠀⠀⠀⠀⠀⠠⠀⠀⠀⠀⠀⠀⠀⡀⠀⠀⠀⠀⠀⡨⠀⠠⠀⠀⠀⠂⠀⠀⠀⠀⠀⠀⠈⡠⠀⡀⠀⠀⠁⠀⠁⠠⠀⠀⠀⠀⠀⠀⠀⠄⡀⠀⠄⡠⠀⠀⠀⠀│
        │⠐⠀⠀⠀⠀⠀⠀⠃⠀⠀⠀⠀⠁⠀⡀⠁⠀⢂⠂⠀⠀⠀⠈⠁⠀⠀⠁⠁⠄⠀⠀⠀⠀⠀⡀⠀⠀⠀⠂⠀⠀⠀⠀⠐⠀⠂⠀⠂⠁⠀⠆⠈⠀⠀⠄⠄⠀⠀⠀⠀│
        │⠀⠀⠀⠠⠀⠁⠀⠀⠀⠀⠂⠀⠠⠀⠀⠌⠀⠀⠂⠄⠀⡀⠂⠀⠀⡀⠠⠁⠄⠀⠀⠀⠀⠀⠄⠈⠀⢂⠐⠀⢀⠀⠠⠀⠄⠁⡀⠀⠀⠐⠀⠀⡐⠀⠀⠀⠁⠀⠀⠔│
        │⠀⠀⠀⠀⢀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡀⠀⠀⠀⠀⠁⠃⡀⡁⠀⠀⠀⢀⠀⠀⠂⠀⠀⠠⠀⠀⠠⠀⠀⠀⠀⠀⠀⠀⡀⠀⠀⠀⠀⠀⠬⠀⡀│
        │⠐⠀⠀⠀⠀⠀⠀⠀⠢⠀⠀⠀⠀⠀⢀⠀⠀⠀⠀⡀⠀⠀⠀⠀⠀⠈⠁⠁⠀⠀⠄⠀⠀⠐⠀⠄⠀⠠⠀⠁⡀⠁⠀⠀⠠⠡⠀⠀⠀⠀⠀⠀⠀⠐⠄⠀⠀⠀⠀⡀│
     50 │⠀⠀⠀⠀⠁⠀⠀⠠⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠠⠂⠀⠀⠀⠀⠀⠀⠐⠀⠈⠂⠀⠀⠀⠀⠀⠀⠀⠀⠠⠈⠀⠀⠀⠀⠀⠀⠂⠀⠄⠀⠀⠀⠂⠁⠀⠀⠀⠂⠀│
        └────────────────────────────────────────────────────────────┘
        1                                                          120
                                  nz = 275

See also
=========

`Plot`, `scatterplot`, `BrailleCanvas`, `BlockCanvas`, `DotCanvas`
"""
function spy{T<:Canvas}(
        A::AbstractMatrix;
        maxwidth::Int  = 0,
        maxheight::Int = 0,
        title::AbstractString = "Sparsity Pattern",
        width::Int  = 0,
        height::Int = 0,
        margin::Int  = 3,
        padding::Int = 1,
        color::Symbol = :automatic,
        labels::Bool = true,
        canvas::Type{T} = BrailleCanvas,
        args...)
    rows, cols, vals = findnz(A)
    nrow, ncol = size(A)
    min_canvheight = ceil(Int, nrow / y_pixel_per_char(T))
    min_canvwidth  = ceil(Int, ncol / x_pixel_per_char(T))
    aspect_ratio = min_canvwidth / min_canvheight
    height_diff = 9
    width_diff  = margin + padding + length(string(ncol)) + 6
    min_plotheight = min_canvheight + height_diff
    min_plotwidth  = min_canvwidth  + width_diff

    # if no size bounds ares specified and the session is in an
    # interactive terminal then use the size of the REPL
    if isinteractive()
        term_height, term_width = Base.tty_size()
        maxheight = maxheight > 0 ? maxheight : term_height - height_diff
        maxwidth  = maxwidth > 0 ? maxwidth : term_width - width_diff
    else
        maxheight = maxheight > 0 ? maxheight : 40
        maxwidth  = maxwidth > 0 ? maxwidth : 70
    end

    # Check if the size of the plot should be derived from the matrix
    # Note: if both width and height are 0, it means that there are no
    #       constraints and the plot should resemble the structure of
    #       the matrix as close as possible
    if width == 0 && height == 0
        # If the interactive code did not take care of this then try
        # to plot the matrix in the correct aspect ratio (within specified bounds)
        if min_canvheight > min_canvwidth
            # long matrix (according to pixel density)
            height = min_canvheight
            width  = height * aspect_ratio
            if width > maxwidth
                width  = maxwidth
                height = width / aspect_ratio
            end
            if height > maxheight
                height = maxheight
                width  = min(height * aspect_ratio, maxwidth)
            end
        else
            # wide matrix
            width  = min_canvwidth
            height = width / aspect_ratio
            if height > maxheight
                height = maxheight
                width  = height * aspect_ratio
            end
            if width > maxwidth
                width = maxwidth
                height = min(width / aspect_ratio, maxheight)
            end
        end
    end
    if width == 0 && height > 0
        width  = min(height * aspect_ratio, maxwidth)
    elseif width > 0 && height == 0
        height = min(width / aspect_ratio, maxheight)
    end
    width  = round(Int, width)
    height = round(Int, height)
    can = T(width, height,
            width  = Float64(ncol) + 1,
            height = Float64(nrow) + 1)
    plot = Plot(can;
                labels = labels, title = title,
                margin = margin, padding = padding, args...)
    height = nrows(plot.graphics)
    width = ncols(plot.graphics)
    plot = if color != :automatic
        points!(plot,
                convert(Vector{Float64}, cols),
                nrow + 1 - convert(Vector{Float64}, rows),
                color)
    else
        pos_idx = vals .> 0
        neg_idx = !pos_idx
        pos_cols = cols[pos_idx]
        pos_rows = rows[pos_idx]
        neg_cols = cols[neg_idx]
        neg_rows = rows[neg_idx]
        points!(plot,
                convert(Vector{Float64}, pos_cols),
                nrow + 1 - convert(Vector{Float64}, pos_rows),
                :red)
        points!(plot,
                convert(Vector{AbstractFloat}, neg_cols),
                nrow + 1 - convert(Vector{Float64}, neg_rows),
                :blue)
        annotate!(plot, :r, 1, "> 0", :red)
        annotate!(plot, :r, 2, "< 0", :blue)
    end
    annotate!(plot, :l, 1, "1")
    annotate!(plot, :l, height, string(nrow))
    annotate!(plot, :bl, "1")
    annotate!(plot, :br, string(ncol))
    xlabel!(plot, string("nz = ", length(vals)))
    return plot
end


function test_craigmr()
  craigmr_tol = 1.0e-6;

  function test_craigmr(A, b; λ=0.0)
    (nrow, ncol) = size(A);
    (x, y, stats) = craigmr(A, b, λ=λ);
    r = b - A * x;
    Ar = A' * r;
    # if λ > 0
    #   s = r / sqrt(λ);
    #   r = r - sqrt(λ) * s;
    # end
    resid = norm(r) / norm(b);
    Aresid = norm(Ar) / (norm(A) * norm(b));
    @printf("CRAIGMR: residual: %7.1e  least-squares: %7.1e\n", resid, Aresid);
    return (x, y, stats, resid, Aresid)
  end

  # Underdetermined consistent.
  A = rand(10, 25); b = A * ones(25);
  (x, y, stats, resid, Aresid) = test_craigmr(A, b);
  @test(norm(x - A' * y) <= craigmr_tol * norm(x));
  @test(resid <= craigmr_tol);
  @test(stats.solved);
  (xI, xmin, xmin_norm) = check_min_norm(A, b, x);
  @test(norm(xI - xmin) <= cond(A) * craigmr_tol * xmin_norm);

  # Underdetermined inconsistent.
  A = ones(10, 25); b = rand(10); b[1] = -1.0;
  (x, y, stats, resid) = test_craigmr(A, b);
  @test(norm(x - A' * y) <= craigmr_tol * norm(x));
  @test(stats.inconsistent);
  @test(stats.Aresiduals[end] <= craigmr_tol);

  # Square consistent.
  A = rand(10, 10); b = A * ones(10);
  (x, y, stats, resid) = test_craigmr(A, b);
  @test(norm(x - A' * y) <= craigmr_tol * norm(x));
  @test(resid <= craigmr_tol);
  @test(stats.solved);
  (xI, xmin, xmin_norm) = check_min_norm(A, b, x);
  @test(norm(xI - xmin) <= cond(A) * craigmr_tol * xmin_norm);

  # Square inconsistent.
  A = ones(10, 10); b = rand(10); b[1] = -1.0;
  (x, y, stats, resid) = test_craigmr(A, b);
  @test(norm(x - A' * y) <= craigmr_tol * norm(x));
  @test(stats.inconsistent);
  @test(stats.Aresiduals[end] <= craigmr_tol);

  # Overdetermined consistent.
  A = rand(25, 10); b = A * ones(10);
  (x, y, stats, resid) = test_craigmr(A, b);
  @test(norm(x - A' * y) <= craigmr_tol * norm(x));
  @test(resid <= craigmr_tol);
  @test(stats.solved);
  (xI, xmin, xmin_norm) = check_min_norm(A, b, x);
  @test(norm(xI - xmin) <= cond(A) * craigmr_tol * xmin_norm);

  # Overdetermined inconsistent.
  A = ones(5, 3); b = rand(5); b[1] = -1.0;
  (x, y, stats, resid) = test_craigmr(A, b);
  @test(norm(x - A' * y) <= craigmr_tol * norm(x));
  @test(stats.inconsistent);
  @test(stats.Aresiduals[end] <= craigmr_tol);

  # With regularization, all systems are underdetermined and consistent.
  # (x, y, stats, resid) = test_craigmr(A, b, λ=1.0e-3);
  # @test(norm(x - A' * y) <= craigmr_tol * norm(x));
  # @test(resid <= craigmr_tol);
  # @test(stats.solved);
  # (xI, xmin, xmin_norm) = check_min_norm(A, b, x, λ=1.0e-3);
  # @test(norm(xI - xmin) <= cond(A) * craigmr_tol * xmin_norm);

  # Code coverage.
  (x, y, stats) = craigmr(sparse(A), b, λ=1.0e-3);
  show(stats);

  # Test b == 0
  (x, y, stats) = craigmr(A, zeros(size(A,1)), λ=1.0e-3)
  @test x == zeros(size(A,2))
  @test y == zeros(size(A,1))
  @test stats.status == "x = 0 is a zero-residual solution"

  # Test integer values
  A = [I; rand(1:10, 2, 3)]
  b = A * ones(Int, 3)
  (x, y, stats) = craigmr(A, b)
  @test stats.solved
end

test_craigmr()


using Statistics

"""
```
mlr_beta(y::Array{T}; n::Int64=10, x::Array{T}=collect(1.0:n))::Matrix{Float64} where {T<:Real}
```

Moving linear regression intercept (column 1) and slope (column 2)
"""
function mlr_beta(y::Array{T}; n::Int64=10, x::Array{T}=collect(1.0:n))::Matrix{Float64} where {T<:Real}
    @assert n<length(y) && n>0 "Argument n out of bounds."
    @assert size(y,2) == 1
    @assert size(x,1) == n || size(x,1) == size(y,1)
    const_x = size(x,1) == n
    out = zeros(T, (length(y),2))
    out[1:n-1,:] .= NaN
    xbar = mean(x)
    ybar = runmean(y, n=n, cumulative=false)
    @inbounds for i = n:length(y)
        yi = y[i-n+1:i]
        xi = const_x ? x : x[i-n+1:i]
        out[i,2] = cov(xi,yi) / var(xi)
        out[i,1] = ybar[i] - out[i,2]*xbar
    end
    return out
end

"""
```
mlr_slope(y::Array{T}; n::Int64=10, x::Array{T}=collect(1.0:n))::Array{Float64} where {T<:Real}
```

Moving linear regression slope
"""
function mlr_slope(y::Array{T}; n::Int64=10, x::Array{T}=collect(1.0:n))::Array{Float64} where {T<:Real}
    @assert n<length(y) && n>0 "Argument n out of bounds."
    @assert size(y,2) == 1
    @assert size(x,1) == n || size(x,1) == size(y,1)
    const_x = size(x,1) == n
    out = zeros(size(y))
    out[1:n-1] .= NaN
    @inbounds for i = n:length(y)
        yi = y[i-n+1:i]
        xi = const_x ? x : x[i-n+1:i]
        out[i] = cov(xi,yi) / var(xi)
    end
    return out
end

"""
```
mlr_intercept(y::Array{T}; n::Int64=10, x::Array{T}=collect(1.0:n))::Array{Float64} where {T<:Real}
```

Moving linear regression y-intercept
"""
function mlr_intercept(y::Array{T}; n::Int64=10, x::Array{T}=collect(1.0:n))::Array{Float64} where {T<:Real}
    @assert n<length(y) && n>0 "Argument n out of bounds."
    @assert size(y,2) == 1
    @assert size(x,1) == n || size(x,1) == size(y,1)
    const_x = size(x,1) == n
    out = zeros(size(y))
    out[1:n-1] .= NaN
    xbar = mean(x)
    ybar = runmean(y, n=n, cumulative=false)
    @inbounds for i = n:length(y)
        yi = y[i-n+1:i]
        xi = const_x ? x : x[i-n+1:i]
        out[i] = ybar[i] - xbar*(cov(xi,yi)/var(xi))
    end
    return out
end

"""
```
mlr(y::Array{T}; n::Int64=10)::Array{Float64} where {T<:Real}
```

Moving linear regression predictions
"""
function mlr(y::Array{T}; n::Int64=10)::Array{Float64} where {T<:Real}
    b = mlr_beta(y, n=n)
    return b[:,1] + b[:,2]*float(n)
end

"""
```
mlr_se(y::Array{T}; n::Int64=10)::Array{Float64} where {T<:Real}
```

Moving linear regression standard errors
"""
function mlr_se(y::Array{T}; n::Int64=10)::Array{Float64} where {T<:Real}
    yhat = mlr(y, n=n)
    r = zeros(T, n)
    out = zeros(size(y))
    out[1:n-1] .= NaN
    nf = float(n)
    @inbounds for i = n:length(y)
        r = y[i-n+1:i] .- yhat[i]
        out[i] = sqrt(sum(r.^2)/nf)
    end
    return out
end

"""
```
mlr_ub(y::Array{T}; n::Int64=10, se::T=2.0)::Array{Float64} where {T<:Real}
```

Moving linear regression upper bound
"""
function mlr_ub(y::Array{T}; n::Int64=10, se::T=2.0)::Array{Float64} where {T<:Real}
    return y + se*mlr_se(y, n=n)
end

"""
```
mlr_lb(y::Array{T}; n::Int64=10, se::T=2.0)::Array{Float64} where {T<:Real}
```

Moving linear regression lower bound
"""
function mlr_lb(y::Array{T}; n::Int64=10, se::T=2.0)::Array{Float64} where {T<:Real}
    return y - se*mlr_se(y, n=n)
end

"""
```
mlr_bands(y::Array{T}; n::Int64=10, se::T=2.0)::Matrix{Float64} where {T<:Real}
```

Moving linear regression bands

*Output:*

Column 1: Lower bound

Column 2: Regression estimate

Column 3: Upper bound
"""
function mlr_bands(y::Array{T}; n::Int64=10, se::T=2.0)::Matrix{Float64} where {T<:Real}
    out = zeros(T, (length(y),3))
    out[1:n-1,:] .= NaN
    out[:,2] = mlr(y, n=n)
    out[:,1] = mlr_lb(y, n=n, se=se)
    out[:,3] = mlr_ub(y, n=n, se=se)
    return out
end

"""
```
mlr_rsq(y::Array{T}; n::Int64=10, adjusted::Bool=false)::Array{Float64} where {T<:Real}
```

Moving linear regression R-squared or adjusted R-squared
"""
function mlr_rsq(y::Array{T}; n::Int64=10, adjusted::Bool=false)::Array{Float64} where {T<:Real}
    yhat = mlr(y, n=n)
    rsq = runcor(y, yhat, n=n, cumulative=false) .^ 2.0
    if adjusted
        return rsq .- (1.0.-rsq)*(1.0/(float(n).-2.0))
    else
        return rsq
    end
end


"""
Baysian Structure learning seeks to maximize P(G|D)
In the Bayesian fashion, we can provide a prior over the parameters in our learning network.
This is described using a Dirichlet Prior.
"""
abstract type DirichletPrior end

"""
A uniform Dirichlet prior such that all α are the same

Defaults to the popular K2 prior, α = 1, which is similar to Laplace Smoothing

    https://en.wikipedia.org/wiki/Additive_smoothing
"""
struct UniformPrior <: DirichletPrior
    α::Float64
    UniformPrior(α::Float64=1.0) = new(α)
end

Base.print(io::IO, p::UniformPrior) = Base.print(io, "UniformPrior(%.2f)", p.α)

Base.get(p::UniformPrior, ncategories::Integer) = fill(p.α, ncategories)
function Base.get{I<:Integer, J<:Integer}(p::UniformPrior,
    var_index::Integer,
    ncategories::AbstractVector{I}, # [nvars]
    parents::AbstractVector{J},    # [nvars]
    )

    r = ncategories[var_index]
    q = isempty(parents) ? 1 : prod(ncategories[parents])
    α = p.α

    fill(α, r, q)
end

"""
Assigns equal scores to Markov equivalent structures

    α_ijk = x/{q_i * r_i} for each j, k and some given x

see DMU section 2.4.3
"""
struct BDeuPrior <: DirichletPrior
    x::Float64
    BDeuPrior(x::Float64=1.0) = new(x)
end

Base.print(io::IO, p::BDeuPrior) = @printf(io, "BDeuPrior(%.2f)", p.x)

Base.get(p::BDeuPrior, ncategories::Integer) = fill(p.x/ncategories, ncategories)
function Base.get{I<:Integer, J<:Integer}(p::BDeuPrior,
    var_index::Integer,
    ncategories::AbstractVector{I}, # [nvars]
    parents::AbstractVector{J},    # [nvars]
    )

    x = p.x
    r = ncategories[var_index]
    q = isempty(parents) ? 1 : prod(ncategories[parents])

    α = x / (r*q)

    fill(α, r, q)
end


##
# @author b33rNc0d3
# @date 2020-07-05
# Linear Regression in Julia with Turing for Bayesian Statistics: Techniques and Models
##

using Turing
using Distributions


using MixedModels
using Test

using MixedModels: dataset, likelihoodratiotest
using GLM: ProbitLink

@testset "likelihoodratio test" begin
    slp = dataset(:sleepstudy);
    fm0 = fit(MixedModel,@formula(reaction ~ 1 + (1+days|subj)),slp);
    fm1 = fit(MixedModel,@formula(reaction ~ 1 + days + (1+days|subj)),slp);
    lrt = likelihoodratiotest(fm0,fm1);

    @test [deviance(fm0), deviance(fm1)] == lrt.deviance
    @test deviance(fm0) - deviance(fm1) == only(lrt.tests.deviancediff)
    @test only(lrt.tests.dofdiff) == 1
    @test sum(map(length,lrt.tests)) == 3
    @test sum(map(length,lrt.pvalues)) == 1
    @test sum(map(length,lrt.models)) == 4
    @test length(lrt.formulae) == 2
    show(IOBuffer(),lrt);
    @test :pvalues in propertynames(lrt)

    # mix of REML and ML
    fm0 = fit(MixedModel,@formula(reaction ~ 1 + (1+days|subj)),slp, REML=true);
    @test_throws ArgumentError likelihoodratiotest(fm0,fm1)

    # differing FE with REML
    fm1 = fit(MixedModel,@formula(reaction ~ 1 + days + (1+days|subj)),slp, REML=true);

    @test_throws ArgumentError likelihoodratiotest(fm0,fm1)

    contra = MixedModels.dataset(:contra);
    gm0 = fit(MixedModel, @formula(use ~ 1+age+urban+livch+(1|urbdist)), contra, Bernoulli(), fast=true);
    gm1 = fit(MixedModel, @formula(use ~ 1+age+abs2(age)+urban+livch+(1|urbdist)), contra, Bernoulli(), fast=true);
    lrt = likelihoodratiotest(gm0,gm1);
    @test [deviance(gm0), deviance(gm1)] == lrt.deviance
    @test deviance(gm0) - deviance(gm1) == only(lrt.tests.deviancediff)
    @test first(lrt.tests.dofdiff) == 1
    @test sum(length, lrt.tests) == 3
    @test sum(length, lrt.pvalues) == 1
    @test sum(length, lrt.models) == 4
    @test length(lrt.formulae) == 2

    # mismatched links
    gm_probit = fit(MixedModel, @formula(use ~ 1+age+urban+livch+(1|urbdist)), contra, Bernoulli(), ProbitLink(), fast=true);
    @test_throws ArgumentError likelihoodratiotest(gm0,gm_probit)

    # mismatched families
    gm_poisson = fit(MixedModel, @formula(use ~ 1+age+urban+livch+(1|urbdist)), contra, Poisson(), fast=true);
    @test_throws ArgumentError MixedModels.likelihoodratiotest(gm0,gm_poisson)
end


"""
Generate coefficients related to PML boundaries, e.g., damping profiles.
This function outputs a, b, and k variables in eq. 25-26 from Komatitsch 2007 (Geophysics).

	TODO
* At the moment, K=1, and is dummy, unlike in the case of EM when it was useful. Removing K and K_inv might save us memory when using GPUs.
"""
function update_pml!(
    pml::NamedVector{T},
    exmgrid, # 1D grid extended
    mgrid, # 1D grid without extension
    flags::Vector{Bool},
    δt::Float64,
    velavg::Float64,
    freqpeak::Float64, # dominant frequency in Hz
) where {T<:Data.Array}

    nx = length(exmgrid)

    # origin, where the PMLs start (offset b/w the domain of interest and PML region depends on order
    xoriginleft = mgrid[1] - 2 * (step(mgrid))
    xoriginright = mgrid[end] + 2 * (step(mgrid))

    NPOWER = 2.e0
    K_MAX_PML = 1.e0 # from Gedney page 8.11
    ALPHA_MAX_PML = pi * freqpeak # from Festa and Vilotte

    # thickness of the PML layer in meters
    thickness_PML = abs(xoriginleft - exmgrid[1])

    "reflection coefficient (INRIA report section 6.1) http://hal.inria.fr/docs/00/07/32/19/PDF/RR-3471.pdf"
    Rcoef = 0.001e0

    #! check that NPOWER is okay
    #if(NPOWER < 1) stop "NPOWER must be greater than 1"

    # compute d0 from INRIA report section 6.1 http://hal.inria.fr/docs/00/07/32/19/PDF/RR-3471.pdf
    d0 = -(NPOWER + 1) * (velavg) * log(Rcoef) / (2.e0 * thickness_PML)

    k = zeros(nx)
    fill!(k, 1)
    d = zeros(nx)
    alpha = zero(d)
    a = zero(d)
    b = zero(d)

    "damping in the X direction"
    "origin of the PML layer (position of right edge minus thickness, in meters)"

    for ix = 1:nx
        #---------- left edge
        if (flags[1])
            # define damping profile at the grid points
            abscissa_in_PML = xoriginleft - exmgrid[ix]
            if (abscissa_in_PML >= 0.0)
                abscissa_normalized = abscissa_in_PML / thickness_PML
                d[ix] = d0 * abscissa_normalized .^ NPOWER
                # this taken from Gedney page 8.2
                k[ix] = 1.e0 + (K_MAX_PML - 1.e0) * abscissa_normalized .^ NPOWER
                alpha[ix] =
                    ALPHA_MAX_PML * (1.e0 - abscissa_normalized) + 0.01e0 * ALPHA_MAX_PML
            end
        end
        #---------- right edge
        if (flags[2])
            # define damping profile at the grid points
            abscissa_in_PML = exmgrid[ix] - xoriginright
            if (abscissa_in_PML >= 0.0)
                abscissa_normalized = abscissa_in_PML / thickness_PML
                d[ix] = d0 * abscissa_normalized^NPOWER
                # this taken from Gedney page 8.2
                k[ix] = 1.e0 + (K_MAX_PML - 1.e0) * abscissa_normalized^NPOWER
                alpha[ix] =
                    ALPHA_MAX_PML * (1.e0 - abscissa_normalized) + 0.01e0 * ALPHA_MAX_PML
            end
        end
        #
        # just in case, for -5 at the end
        (alpha[ix] < 0.0) ? alpha[ix] = 0.0 : nothing

        # see equation 25 Komatitsch, 2007, get b and a (need k, alpha, and d before this )
        b[ix] = exp(-(d[ix] / k[ix] + alpha[ix]) * δt)

        # this to avoid division by zero outside the PML
        (abs(d[ix]) > 1.e-6) ?
        a[ix] = d[ix] * (b[ix] - 1.e0) / (k[ix] * (d[ix] + k[ix] * alpha[ix])) : nothing
    end
    pkI = zeros(2 * _fd.npml)
    pa = zeros(2 * _fd.npml)
    pb = zeros(2 * _fd.npml)
    for i = 1:_fd.npml
        j = i + _fd.npml
        pkI[i] = inv(k[i])
        pkI[j] = inv(k[end-_fd.npml+i])
        pa[i] = a[i]
        pa[j] = a[end-_fd.npml+i]
        pb[i] = b[i]
        pb[j] = b[end-_fd.npml+i]
    end
    copyto!(pml[:kI], pkI)
    copyto!(pml[:a], pa)
    copyto!(pml[:b], pb)
end

"""
Generate a NamedArray with PML coefficients for all the dimensions that are then stored in the FDTD structs.
"""
function get_pml(attrib_mod, mgrid)
    dfields = Fields(attrib_mod, "d", ndims = length(mgrid)) # derivative fields that need PML memory
    pnames = [:a, :b, :kI]
    np = 2 * _fd.npml
    return NamedArray(
        [
            NamedArray(Data.Array.([zeros(np), zeros(np), ones(np)]), pnames) for
            df in dfields
        ],
        dfields,
    )
end

"""
Just a loop over dims of pml
"""
function update_pml!(
    pml::NamedVector{T},
    exmgrid,
    mgrid,
    pml_edges::Vector{Symbol},
    attrib_mod,
    args...,
) where {T<:NamedVector}
    for df in names(pml)[1]
        dim = string(last(string(df)))
        i = findfirst(x -> string(x) == dim, dim_names(length(exmgrid)))
        exmgridf = get_mgrid(eval(df)(), attrib_mod, exmgrid...)[i]
        mgridf = mgrid[i]
        update_pml!(
            pml[df],
            exmgridf,
            mgridf,
            [
                any(pml_edges .== Symbol(string(dim), "min")),
                any(pml_edges .== Symbol(string(dim), "max")),
            ],
            args...,
        )
    end
end

function update_pml!(pac)
    update_pml!(
        pac.pml,
        pac.exmedium.mgrid,
        pac.medium.mgrid,
        pac.pml_edges,
        pac.attrib_mod,
        pac.fc[:dt],
        Statistics.mean(pac.exmedium.bounds[:vp]),
        pac.fc[:freqpeak],
    )
end

for dimnames in [zip([:1, :2, :3], dim_names(3)), zip([:1, :2], dim_names(2))]
    is = broadcast(x -> Symbol(string("i", x)), getindex.(collect(dimnames), 2))
    ist = Meta.parse(string("(", [string(s, ",") for s in is]..., ")"))
    N = Meta.parse(string(length(is)))

    for (idim, dim) in dimnames
        i = Symbol("i", string(dim))
        ismoff = replace(is, i => :($i + moff))
        isdoff = replace(is, i => :(doff + $i))
        for (fname, fnamenp, imoff) in zip(
            [Symbol("memory", string(dim), "!"), Symbol("memory1", string(dim), "!")],
            [Symbol("memorynp", string(dim), "!"), Symbol("memorynp1", string(dim), "!")],
            [:($i + moff), :($i + moff + 1)],
        )
            @eval @parallel_indices(
                $ist,
                function $fnamenp(memory::Data.Array{$N}, d, a, b, kI, moff, doff)
                    memory[$(ismoff...)] =
                        b[$imoff] * memory[$(ismoff...)] + a[$imoff] * d[$(isdoff...)]
                    d[$(isdoff...)] = d[$(isdoff...)] * kI[$imoff] + memory[$(ismoff...)]
                    return
                end
            )
            @eval function $fname(memory::Data.Array{$N}, d, a, b, kI)
                sm = collect(size(memory))
                setindex!(sm, _fd.npml, $idim)
                # first _fd.npml points
                @parallel map(x -> (:)(1, x), Tuple(sm)) $fnamenp(memory, d, a, b, kI, 0, 0)
                # last _fd.npml points independent of d
                @parallel map(x -> (:)(1, x), Tuple(sm)) $fnamenp(
                    memory,
                    d,
                    a,
                    b,
                    kI,
                    _fd.npml,
                    getindex(size(d), $idim) - _fd.npml,
                )
            end
        end
    end
end


using StaticArrays, LinearAlgebra

"""
    GaussHermiteQuadrature

As described in

* [Gauss-Hermite quadrature on Wikipedia](http://en.wikipedia.org/wiki/Gauss-Hermite_quadrature)

*Gauss-Hermite* quadrature uses a weighted sum of values of `f(x)` at specific `x` values to approximate

```math
\\int_{-\\infty}^\\infty f(x) e^{-x^2} dx
```

An `n`-point rule, as returned by `hermite(n)` from the
[`GaussQuadrature``](https://github.com/billmclean/GaussQuadrature.jl) package provides `n` abscicca
values (i.e. values of `x`) and `n` weights.

As noted in the Wikipedia article, a modified version can be used to evaluate the expectation `E[h(x)]`
with respect to a `Normal(μ, σ)` density as
```julia
using MixedModels

gn5 = GHnorm(5)
μ = 3.
σ = 2.
sum(@. abs2(σ*gn5.z + μ)*gn5.w) # E[X^2] where X ∼ N(μ, σ)
```

For evaluation of the log-likelihood of a GLMM the integral to evaluate for each level of the grouping
factor is approximately Gaussian shaped.
"""
GaussHermiteQuadrature
"""
    GaussHermiteNormalized{K}

A struct with 2 SVector{K,Float64} members
- `z`: abscissae for the K-point Gauss-Hermite quadrature rule on the Z scale
- `wt`: Gauss-Hermite weights normalized to sum to unity
"""
struct GaussHermiteNormalized{K}
    z::SVector{K, Float64}
    w::SVector{K,Float64}
end
function GaussHermiteNormalized(k::Integer)
    ev = eigen(SymTridiagonal(zeros(k), sqrt.(1:k-1)))
    w = abs2.(ev.vectors[1,:])
    GaussHermiteNormalized(
        SVector{k}((ev.values .- reverse(ev.values)) ./ 2),
        SVector{k}(LinearAlgebra.normalize((w .+ reverse(w)) ./ 2, 1)))
end

Base.iterate(g::GaussHermiteNormalized{K}, i=1) where {K} = (K < i ? nothing : ((z = g.z[i], w = g.w[i]), i + 1))

Base.length(g::GaussHermiteNormalized{K}) where {K} = K

"""
    GHnormd

Memoized values of `GHnorm`{@ref} stored as a `Dict{Int,GaussHermiteNormalized}`
"""
const GHnormd = Dict{Int,GaussHermiteNormalized}(
    1 => GaussHermiteNormalized(SVector{1}(0.),SVector{1}(1.)),
    2 => GaussHermiteNormalized(SVector{2}(-1.0,1.0),SVector{2}(0.5,0.5)),
    3 => GaussHermiteNormalized(SVector{3}(-sqrt(3),0.,sqrt(3)),SVector{3}(1/6,2/3,1/6))
    )

"""
    GHnorm(k::Int)

Return the (unique) GaussHermiteNormalized{k} object.

The function values are stored (memoized) when first evaluated.  Subsequent evaluations
for the same `k` have very low overhead.
"""
GHnorm(k::Int) = get!(GHnormd, k) do
    GaussHermiteNormalized(k)
end
GHnorm(k) = GHnorm(Int(k))


module Day05

import Base.show, Base.string
import ..data_dir # from parent module

export part1, part2

struct Point
    x::Int64
    y::Int64
end

function string(point::Point)
    "$(point.x),$(point.y)"
end

function show(io::IO, point::Point)
    println(io, string(point))
end

function Point(point_defn::AbstractString)
    x_string, y_string = split(point_defn, ",")
    x, y = parse(Int64, x_string), parse(Int64, y_string)
    Point(x, y)
end

struct Line
    a::Point
    b::Point
end

function string(line::Line)
    "$(string(line.a)) -> $(string(line.b))"
end

function show(io::IO, line::Line)
    println(io, string(line))
end

function Line(line_defn::AbstractString)
    point_defns = split(line_defn, " -> ")
    a_string, b_string = point_defns[1], point_defns[2]
    a, b = Point(a_string), Point(b_string)
    Line(a, b)
end

is_horizontal(line::Line) = (line.a.y == line.b.y)
is_vertical(line::Line) = (line.a.x == line.b.x)

"""
Julia sees no points for the range `5:1` since `5 > 1`. This function makes sure
the step is correct (either 1 or -1) so that `5:1 == [5, 4, 3, 2, 1]`.
"""
function sequence(i, j)
    if i > j
        range(i, j, step = -1)
    else
        range(i, j, step = 1)
    end
end

"""
Also known as _gradient_
"""
slope(line) = (line.b.y - line.a.y) / (line.b.x - line.a.x)

"""
Returns a dictionary indexed by points and with values the number of occurrences of
each points. Points that aren't on any of the given lines will not appear in the
dictionary at all.

Lines are either horizontal, vertical, or diagonal at a 45° angle. In the diagonal
case, the slope is ±1, so we don't need to worry about non-integral coordinates.

The vertical line case has to be handled separately because the slope is NaN due to
a division by zero
"""
function points_on_line(line::Line)
    if is_vertical(line)
        return ([Point(line.a.x, j) for j in sequence(line.a.y, line.b.y)])
    end

    slope_of_line = slope(line)
    points = Vector{Point}()
    for i in sequence(0, line.b.x - line.a.x)
        # Move 1 along and 1 up/down each time
        point = Point(
            Int(line.a.x + i),
            Int(line.a.y + (i * slope_of_line))
        )
        push!(points, point)
    end
    points
end

"""
Returns a dictionary indexed by points and with values the number of occurrences of
each points. Points that aren't on any of the given lines will not appear in the
dictionary at all.
"""
function points_on_lines(lines::Vector{Line})
    occurrences = Dict{Point,Int64}()
    for line in lines
        for point in points_on_line(line)
            occurrences[point] = get(occurrences, point, 0) + 1
        end
    end
    occurrences
end

function count_multiple_occurrences(lines::Vector{Line})
    occurrences = points_on_lines(lines)
    length([p for p in keys(occurrences) if occurrences[p] > 1])
end

# Parse the input
const data_file = joinpath(data_dir, "day05")
const input = readlines(data_file) # read into an array, each line is an element
const lines = [Line(i) for i in input]

function part1(input = input)
    lines = [Line(i) for i in input]
    no_diagonals = [line for line in lines if is_horizontal(line) || is_vertical(line)]
    count_multiple_occurrences(no_diagonals)
end
function part2(input = input)
    lines = [Line(i) for i in input]
    count_multiple_occurrences(lines)
end

end # module


import Base.LinAlg.BLAS: axpy!
export idrs, idrs!

####

@inline function omega(t, s)
    angle = sqrt(2.)/2
    ns = vecnorm(s)
    nt = vecnorm(t)
    ts = vecdot(t,s)
    rho = abs(ts/(nt*ns))
    om = ts/(nt*nt)
    if rho < angle
        om = om*convert(typeof(om),angle)/rho
    end
    om
end

@inline linsys_op(x, A) = A*x

"""
IDRsSolver.jl
============================

The Induced Dimension Reduction method is a family of simple and fast Krylov
subspace algorithms for solving large nonsymmetric linear systems. The idea
behind the IDR(s) variant is to generate residuals that are in the nested
subspaces of shrinking dimension s.

Syntax
------

        X, h::ConvergenceHistory = idrs(A, b; s = 8, tol=sqrt(eps(typeof(real(b[1])))), maxiter = length(b)^2))

        The operator A must support multiplication with b,
        the right hand side b must support vecnorm, vecdot, copy, rand! and axpy!.
        .

Arguments
---------

       s -- dimension reduction number. Normally, a higher s gives faster convergence,
            but also  makes the method more expensive.
       tol -- tolerance of the method.
       maxiter -- maximum number of iterations
       smoothing -- (Logical) specifies if residual smoothing must be applied. Default is false.

Output
------

        X -- Approximated solution by IDR(s)
        h -- Convergence history

        The [`ConvergenceHistory`](https://github.com/JuliaLang/IterativeSolvers.jl/issues/6) type provides information about the iteration history.
            - `isconverged::Bool`, a flag for whether or not the algorithm is converged.
            - `threshold`, the convergence threshold
            - `residuals::Vector`, the value of the convergence criteria at each iteration

References
----------

    [1] IDR(s): a family of simple and fast algorithms for solving large
        nonsymmetric linear systems. P. Sonneveld and M. B. van Gijzen
        SIAM J. Sci. Comput. Vol. 31, No. 2, pp. 1035--1062, 2008
    [2] Algorithm 913: An Elegant IDR(s) Variant that Efficiently Exploits
        Bi-orthogonality Properties. M. B. van Gijzen and P. Sonneveld
        ACM Trans. Math. Software,, Vol. 38, No. 1, pp. 5:1-5:19, 2011
    [3] This file is a translation of the following MATLAB implementation:
        http://ta.twi.tudelft.nl/nw/users/gijzen/idrs.m
    [4] IDR(s)' webpage http://ta.twi.tudelft.nl/nw/users/gijzen/IDR.html
"""
idrs(A, b; s = 8, tol=sqrt(eps(typeof(real(b[1])))), maxiter = length(b)^2, smoothing=false) =
    idrs!(zerox(A, b), A, b; s=s, tol=tol, maxiter=maxiter, smoothing=smoothing)

function idrs!(x, A, b; s = 8, tol=sqrt(eps(typeof(real(b[1])))), maxiter=length(x)^2, smoothing=false)
    history = ConvergenceHistory()
    history[:tol] = tol
    reserve!(history,:resnorm, maxiter)
    idrs_method!(history, x, linsys_op, (A,), b, s, tol, maxiter; smoothing=smoothing)
    x, history
end

function idrs_method!{T}(log::ConvergenceHistory, X, op, args, C::T,
    s::Number, tol::Number, maxiter::Number; smoothing::Bool=false)

    R = C - op(X, args...)::T
    normR = vecnorm(R)
	iter = 0

    if smoothing
        X_s = copy(X)
        R_s = copy(R)
        T_s = zeros(R)
    end

    if normR <= tol           # Initial guess is a good enough solution
        return X, ConvergenceHistory(0<= res[end] < tol, tol, length(res), res)
    end

    Z = zero(C)

    P = T[rand!(copy(C)) for k in 1:s]
    U = T[copy(Z) for k in 1:s]
    G = T[copy(Z) for k in 1:s]
    Q = copy(Z)
    V = copy(Z)

    M = eye(eltype(C),s,s)
    f = zeros(eltype(C),s)
    c = zeros(eltype(C),s)

    om::eltype(C) = 1
    while normR > tol && iter < maxiter
        for i in 1:s,
            f[i] = vecdot(P[i], R)
        end
        for k in 1:s
            nextiter!(log,mvps=1)

            # Solve small system and make v orthogonal to P

            c = LowerTriangular(M[k:s,k:s])\f[k:s]
            copy!(V, G[k])
            scale!(c[1], V)

            copy!(Q, U[k])
            scale!(c[1], Q)
            for i = k+1:s
                axpy!(c[i-k+1], G[i], V)
                axpy!(c[i-k+1], U[i], Q)
            end

            # Compute new U[:,k] and G[:,k], G[:,k] is in space G_j

            #V = R - V
            scale!(-1., V)
            axpy!(1., R, V)

            copy!(U[k], Q)
            axpy!(om, V, U[k])
            G[k] = op(U[k], args...)

            # Bi-orthogonalise the new basis vectors

            for i in 1:k-1
                alpha = vecdot(P[i],G[k])/M[i,i]
                axpy!(-alpha, G[i], G[k])
                axpy!(-alpha, U[i], U[k])
            end

            # New column of M = P'*G  (first k-1 entries are zero)

            for i in k:s
                M[i,k] = vecdot(P[i],G[k])
            end

            #  Make r orthogonal to q_i, i = 1..k

            beta = f[k]/M[k,k]
            axpy!(-beta, G[k], R)
            axpy!(beta, U[k], X)

            normR = vecnorm(R)
            if smoothing
                # T_s = R_s - R
                copy!(T_s, R_s); axpy!(-1., R, T_s)

                gamma = vecdot(R_s, T_s)/vecdot(T_s, T_s)

                # R_s = R_s - gamma*T_s
                axpy!(-gamma, T_s, R_s)
                # X_s = X_s - gamma*(X_s - X)
                copy!(T_s, X_s); axpy!(-1., X, T_s); axpy!(-gamma, T_s, X_s)

                normR = vecnorm(R_s)
            end
            push!(log, :resnorm, normR)
            iter += 1
            if normR < tol || iter > maxiter
                shrink!(log)
                setconv(log, 0<=normR<tol)
                return X
            end
            if k < s
                f[k+1:s] = f[k+1:s] - beta*M[k+1:s,k]
            end

        end

        # Now we have sufficient vectors in G_j to compute residual in G_j+1
        # Note: r is already perpendicular to P so v = r
        copy!(V, R)
        Q = op(V, args...)::T
        om = omega(Q, R)
        axpy!(-om, Q, R)
        axpy!(om, V, X)

        normR = vecnorm(R)
        if smoothing
            # T_s = R_s - R
            copy!(T_s, R_s); axpy!(-1., R, T_s)

            gamma = vecdot(R_s, T_s)/vecdot(T_s, T_s)

            # R_s = R_s - gamma*T_s
            axpy!(-gamma, T_s, R_s)
            # X_s = X_s - gamma*(X_s - X)
            copy!(T_s, X_s); axpy!(-1., X, T_s); axpy!(-gamma, T_s, X_s)

            normR = vecnorm(R_s)
        end
        iter += 1
        nextiter!(log,mvps=1)
        push!(log, :resnorm, normR)
    end
    if smoothing
        copy!(X, X_s)
    end
    shrink!(log)
    setconv(log, 0<=normR<tol)
    X
end