text:
module ClosureCoefficients

using Combinatorics
using LinearAlgebra
using Printf
using SparseArrays
using StatsBase

export undir_clcfs, dir_clcfs, clcfs_data, construct_directed_data, load_example_data

const SpIntMat = SparseMatrixCSC{Int64,Int64}

const TRI_oo_o = 1
const TRI_oo_i = 2
const TRI_oi_o = 3
const TRI_oi_i = 4
const TRI_io_o = 5
const TRI_io_i = 6
const TRI_ii_o = 7
const TRI_ii_i = 8

"""
Read a graph from a text file. Skips any lines that begin with '#' or '%'
characters. Nodes are mapped to values 1, ..., n.

Returns (A, v), where
   A is the n x n undirected graph adjacency matrix
   v is a length-n vector such that v[i] is the original node identifier in the text file
"""
function read_graph_txt(filename::AbstractString)
    # index mapping
    index_map = Dict{Int64,Int64}()
    index_map_vec = Int64[]
    function get_mapped_index(x::Int64)
        if !haskey(index_map, x)
            next_index = length(index_map) + 1
            index_map[x] = next_index
            push!(index_map_vec, x)
            return next_index
        end
        return index_map[x]
    end

    # Read data
    I = Int64[]
    J = Int64[]
    open(filename) do f
        for line in eachline(f)
            # Skip lines starting with '#' or '%'
            if line[1] == '#' || line[1] == '%'; continue; end
            edge = split(line)
            u = parse(Int64, edge[1])
            v = parse(Int64, edge[2])
            push!(I, get_mapped_index(u))
            push!(J, get_mapped_index(v))
        end
    end

    n = max(maximum(I), maximum(J))
    A = convert(SpIntMat, sparse(I, J, ones(length(I)), n, n))
    return (A, index_map_vec)
end

"""
Load an example file from the data directory.
"""
function load_example_data(filename::AbstractString; symm::Bool=false)
    pathname = joinpath(dirname(dirname(@__FILE__)), "data")
    filename = joinpath(pathname, filename)
    if   isfile(filename)
        A = read_graph_txt(filename)[1]
        if symm
            A = max.(A, A')
        end
        return A
    else
        error(@sprintf("Could not find file %s", name))
    end
end

"""
clcfs_data
----------
Closure coefficients data.

The field values are
global_clcf::Float64
    The global closure coefficient.
avg_clcf::Float64
    The average closure coefficient (the mean is taken over nodes that head at
    least one wedge).
local_clcfs::Vector{Float64}
    Vctor of local closure coefficients. If a node is not the head of at least
    one wedge, then the value is 0.
wedges::Vector{Int64}
    Vector of the number of wedges at each node: wedges[v] is the number of
    wedges headed by node v.
triangles::Vector{Int64}
    Vector of triangle counts of nodes: tri_counts[v] is the number of triangles
    containing v.
"""
struct clcfs_data
    global_clcf::Float64
    avg_clcf::Float64
    local_clcfs::Vector{Float64}
    wedges::Vector{Int64}
    triangles::Vector{Int64}
end

function construct_directed_data(triangles::Vector{Int64}, wedges::Vector{Int64})
    local_clcfs = zeros(Float64, length(wedges))
    nz_inds = findall(wedges .> 0)
    local_clcfs[nz_inds] = triangles[nz_inds] ./ wedges[nz_inds]
    avg_clcf = (length(nz_inds) > 0) ? mean(local_clcfs[nz_inds]) : 0.0
    num_wedges = sum(wedges)
    global_clcf = (num_wedges > 0) ? sum(triangles) / num_wedges : 0.0
    return clcfs_data(global_clcf, avg_clcf, local_clcfs, wedges, triangles)
end

function degree_order(A::SpIntMat)
    n = size(A, 1)
    deg_order = zeros(Int64, n)
    deg_order[sortperm(vec(sum(A, dims=1)))] = collect(1:n)
end

nz_row_inds(A::SpIntMat, ind::Int64) =
    A.rowval[A.colptr[ind]:(A.colptr[ind + 1] - 1)]

# Return generator for pairs of neighbors higher in the ordering
function ordered_nbr_pairs(A::SpIntMat, i::Int64, order::Vector{Int64})
    pos_i = order[i]
    return combinations(filter(nbr -> order[nbr] > pos_i, nz_row_inds(A, i)), 2)
end

"""
dir_clcfs
-----------------
Compute directed closure coefficients. The input graph is given as a sparse
matrix A, where A[i, j] = 1 if i --> j in the graph.

A::SparseMatrixCSC{Int64,Int64}
    The adjacency matrix of an undirected graph.

returns type clcfs_data
"""
function dir_clcfs(A::SpIntMat)
    A = min.(A, 1)
    A -= Diagonal(A)
    At = convert(SpIntMat, A')
    B = min.(A, At)

    n = size(A, 1)
    dB = vec(sum(B, dims=1))
    dI = vec(sum(A, dims=1))
    dO = vec(sum(At, dims=1))

    wedges_oo = A  * dO - dB
    wedges_io = At * (dO .- 1)
    wedges_oi = A  * (dI .- 1)
    wedges_ii = At * dI - dB

    C = max.(A, At)
    n = size(C, 1)
    deg_order = degree_order(C)
    triangles = zeros(Int64, 8, n)
    for i = 1:n
        for (j, k) in ordered_nbr_pairs(C, i, deg_order)
            if C[j, k] > 0
                # Triangle between i, j, k
                Aij = A[i, j] > 0
                Aik = A[i, k] > 0
                Aji = A[j, i] > 0
                Ajk = A[j, k] > 0
                Aki = A[k, i] > 0
                Akj = A[k, j] > 0
                if Aij && Ajk && Aik
                    triangles[TRI_oo_o, i] += 1
                    triangles[TRI_oi_o, i] += 1
                    triangles[TRI_oi_i, j] += 1
                    triangles[TRI_io_o, j] += 1
                    triangles[TRI_io_i, k] += 1
                    triangles[TRI_ii_i, k] += 1
                end
                if Aij && Ajk && Aki
                    triangles[TRI_oo_i, i] += 1
                    triangles[TRI_ii_o, i] += 1
                    triangles[TRI_oo_i, j] += 1
                    triangles[TRI_ii_o, j] += 1
                    triangles[TRI_oo_i, k] += 1
                    triangles[TRI_ii_o, k] += 1
                end
                if Aij && Akj && Aik
                    triangles[TRI_oi_o, i] += 1
                    triangles[TRI_oo_o, i] += 1
                    triangles[TRI_ii_i, j] += 1
                    triangles[TRI_io_i, j] += 1
                    triangles[TRI_io_o, k] += 1
                    triangles[TRI_oi_i, k] += 1
                end
                if Aij && Akj && Aki
                    triangles[TRI_oi_i, i] += 1
                    triangles[TRI_io_o, i] += 1
                    triangles[TRI_io_i, j] += 1
                    triangles[TRI_ii_i, j] += 1
                    triangles[TRI_oo_o, k] += 1
                    triangles[TRI_oi_o, k] += 1
                end
                if Aji && Ajk && Aik
                    triangles[TRI_io_o, i] += 1
                    triangles[TRI_oi_i, i] += 1
                    triangles[TRI_oi_o, j] += 1
                    triangles[TRI_oo_o, j] += 1
                    triangles[TRI_ii_i, k] += 1
                    triangles[TRI_io_i, k] += 1
                end
                if Aji && Ajk && Aki
                    triangles[TRI_io_i, i] += 1
                    triangles[TRI_ii_i, i] += 1
                    triangles[TRI_oo_o, j] += 1
                    triangles[TRI_oi_o, j] += 1
                    triangles[TRI_oi_i, k] += 1
                    triangles[TRI_io_o, k] += 1
                end
                if Aji && Akj && Aik
                    triangles[TRI_ii_o, i] += 1
                    triangles[TRI_oo_i, i] += 1
                    triangles[TRI_ii_o, j] += 1
                    triangles[TRI_oo_i, j] += 1
                    triangles[TRI_ii_o, k] += 1
                    triangles[TRI_oo_i, k] += 1
                end
                if Aji && Akj && Aki
                    triangles[TRI_ii_i, i] += 1
                    triangles[TRI_io_i, i] += 1
                    triangles[TRI_io_o, j] += 1
                    triangles[TRI_oi_i, j] += 1
                    triangles[TRI_oi_o, k] += 1
                    triangles[TRI_oo_o, k] += 1
                end
            end
        end
    end
    triangles = triangles'

    ret = Dict{String, clcfs_data}()
    ret["oo_o"] = construct_directed_data(triangles[:, TRI_oo_o], wedges_oo)
    ret["oo_i"] = construct_directed_data(triangles[:, TRI_oo_i], wedges_oo)
    ret["oi_o"] = construct_directed_data(triangles[:, TRI_oi_o], wedges_oi)
    ret["oi_i"] = construct_directed_data(triangles[:, TRI_oi_i], wedges_oi)
    ret["io_o"] = construct_directed_data(triangles[:, TRI_io_o], wedges_io)
    ret["io_i"] = construct_directed_data(triangles[:, TRI_io_i], wedges_io)
    ret["ii_o"] = construct_directed_data(triangles[:, TRI_ii_o], wedges_ii)
    ret["ii_i"] = construct_directed_data(triangles[:, TRI_ii_i], wedges_ii)
    return ret
end

"""
undir_clcfs
-----------------
Compute undirected closure coefficients. The input graph is given as a symmetric
sparse matrix.

A::SparseMatrixCSC{Int64,Int64}
    The adjacency matrix of an undirected graph.

returns type clcfs_data
"""
function undir_clcfs(A::SpIntMat)
    A = min.(A, 1)
    A -= Diagonal(A)
    if !issymmetric(A)
        println("Input graph is not symmetric. Symmetrizing...")
        A = max.(A, A')
    end

    n = size(A, 1)
    degs = vec(sum(A, dims=1))
    wedges = A * (degs .- 1)

    deg_order = degree_order(A)
    triangles = zeros(Int64, n)
    for i = 1:n
        for (j, k) in ordered_nbr_pairs(A, i, deg_order)
            if A[j, k] > 0; triangles[[i, j, k]] .+= 1; end
        end
    end

    local_clcfs = zeros(Float64, n)
    nz_inds = findall(wedges .> 0)
    local_clcfs[nz_inds] = 2 * triangles[nz_inds] ./ wedges[nz_inds]
    return clcfs_data(2 * sum(triangles) / sum(wedges),
                      mean(local_clcfs[nz_inds]), local_clcfs, wedges, triangles)
end

end # module


text:
#
# minimize f(x) + ∑ λi⋅|xi|
#
# If warmStart is true, the descent will start from the supplied x
# otherwise it will start from 0 by setting a large value of λ which is
# decreased to the target value
function coordinateDescent!(
  x::Union{SparseIterate,SymmetricSparseIterate,AtomIterate},
  f::CoordinateDifferentiableFunction,
  g::ProxL1,
  options::CDOptions=CDOptions())

  ProximalBase.numCoordinates(x) == numCoordinates(f) || throw(DimensionMismatch())
  if !isa(g, ProxL1{typeof(g.λ0), Nothing}) # TODO: create a test for this
    length(g.λ) == numCoordinates(f) || throw(DimensionMismatch())
  end

  coef_iterator = options.randomize ? RandomIterator(x) : OrderedIterator(x)

  if options.warmStart
    initialize!(f, x)
    return _coordinateDescent!(x, f, g, coef_iterator, options)
  else
    # set x to zero and initialize
    fill!(x, zero(eltype(x)))
    initialize!(f, x)

    # find λmax
    λmax = _findLambdaMax(x, f, g)

    # find decreasing schedule for λ
    l1, l2 = log(λmax), log(g.λ0)
    for l in l1:(l2-l1)/options.numSteps:l2
      g1 = ProxL1(exp(l), g.λ)
      _coordinateDescent!(x, f, g1, coef_iterator, options)
    end
    return x
  end
end

function coordinateDescent!(
  x::Union{SparseIterate,SymmetricSparseIterate,AtomIterate},
  f::CoordinateDifferentiableFunction,
  g::ProxZero,
  options::CDOptions=CDOptions())

  ProximalBase.numCoordinates(x) == numCoordinates(f) || throw(DimensionMismatch())

  coef_iterator = options.randomize ? RandomIterator(x) : OrderedIterator(x)

  if options.warmStart
    initialize!(f, x)
  else
    # set x to zero and initialize
    fill!(x, zero(eltype(x)))
    initialize!(f, x)
  end

  return _coordinateDescent!(x, f, g, coef_iterator, options)
end

# assumes that f is initialized before the call here
function _coordinateDescent!(
  x::Union{SparseIterate,SymmetricSparseIterate,AtomIterate},
  f::CoordinateDifferentiableFunction,
  g::Union{ProxL1, ProxZero},
  coef_iterator::AtomIterator,
  options::CDOptions)

  prev_converged = false
  converged = true
  for iter=1:options.maxIter

    if converged
      reset!(coef_iterator, true)
      # maxH = fullPass!(x, f, g)
    else
      reset!(coef_iterator, false)
      # maxH = nonZeroPass!(x, f, g)
    end
    maxH = _cdPass!(x, f, g, coef_iterator)
    prev_converged = converged

    # test for convergence
    converged = maxH < options.optTol

    prev_converged && converged && break
  end
  x
end

function _cdPass!(
  x::Union{SparseIterate,SymmetricSparseIterate,AtomIterate},
  f::CoordinateDifferentiableFunction,
  g::Union{ProxL1, ProxZero},
  coef_iterator::AtomIterator
  )

  maxH = zero(eltype(x))
  for ipred = coef_iterator               # coef_iterator produces original indexes
    h = descendCoordinate!(f, g, x, ipred)
    if abs(h) > maxH
      maxH = abs(h)
    end
  end
  dropzeros!(x)
  maxH
end

######

"""
Helper function that finds the smallest value of λ for which the solution is equal to zero.
"""
function _findLambdaMax(x::Union{SparseIterate{T},SymmetricSparseIterate{T}},
  f::CoordinateDifferentiableFunction,
  ::ProxL1{T, Nothing}) where {T<:AbstractFloat}

  λmax = zero(T)
  for k=1:ProximalBase.numCoordinates(x)
    f_g = gradient(f, x, k)
    t = abs(f_g)
    if t > λmax
      λmax = t
    end
  end
  λmax
end

"""
Helper function that finds the smallest value of λ0 for which the solution is equal to zero.
"""
function _findLambdaMax(x::SparseIterate{T},
  f::CoordinateDifferentiableFunction,
  g::ProxL1{T, S}) where {T<:AbstractFloat} where S <: AbstractArray

  λmax = zero(T)
  for k=1:length(x)
    f_g = gradient(f, x, k)
    t = abs(f_g) / g.λ[k]
    if t > λmax
      λmax = t
    end
  end
  λmax
end

ext: jl


text:
################################################################################
#                          Direction Cosine Matrices
################################################################################

export create_rotation_matrix
export ddcm, dcm_to_angle, dcm_to_angleaxis, dcm_to_quat, orthonormalize

################################################################################
#                                  Functions
################################################################################

"""
    create_rotation_matrix(angle::Number, axis::Symbol = :X)

Compute a rotation matrix that rotates a coordinate frame about the axis `axis`
by the angle `angle`. The `axis` must be one of the following symbols: `:X`,
`:Y`, or `:Z`.

# Example

```jldocstest
julia> create_rotation_matrix(pi/2, :X)
3×3 StaticArrays.SArray{Tuple{3,3},Float64,2,9}:
 1.0   0.0          0.0
 0.0   6.12323e-17  1.0
 0.0  -1.0          6.12323e-17
```

"""
function create_rotation_matrix(angle::Number, axis::Symbol = :X)
    sin_angle, cos_angle = sincos(angle)

    if axis == :X
        return DCM(1,      0,          0,
                   0, +cos_angle, +sin_angle,
                   0, -sin_angle, +cos_angle)'
    elseif axis == :Y
        return DCM(+cos_angle, 0, -sin_angle,
                        0,     1,      0,
                   +sin_angle, 0, +cos_angle)'
    elseif axis == :Z
        return DCM(+cos_angle, +sin_angle, 0,
                   -sin_angle, +cos_angle, 0,
                        0,          0,     1)'
    else
        error("axis must be :X, :Y, or :Z");
    end
end

"""
    orthonormalize(dcm::DCM)

Perform the Gram-Schmidt orthonormalization process in the DCM `dcm` and return
the new matrix.

**Warning**: This function does not check if the columns of the input matrix
span a three-dimensional space. If not, then the returned matrix should have
`NaN`. Notice, however, that such input matrix is not a valid direction cosine
matrix.

# Example

```julia-repl
julia> D = DCM(3I)

julia> orthonormalize(D)
3×3 StaticArrays.SArray{Tuple{3,3},Float64,2,9}:
 1.0  0.0  0.0
 0.0  1.0  0.0
 0.0  0.0  1.0
```
"""
function orthonormalize(dcm::DCM)
    e₁ = dcm[:,1]
    e₂ = dcm[:,2]
    e₃ = dcm[:,3]

    en₁  = e₁/norm(e₁)
    enj₂ =   e₂ - (en₁⋅e₂)*en₁
    en₂  = enj₂ / norm(enj₂)
    enj₃ =   e₃ - (en₁⋅e₃)*en₁
    enj₃ = enj₃ - (en₂⋅enj₃)*en₂
    en₃  = enj₃ / norm(enj₃)

    [en₁ en₂ en₃]
end

################################################################################
#                                 Conversions
################################################################################

# Euler Angles
# ==============================================================================

# This modified function computes exactly what `atan(y,x)` computes except that
# it will neglect signed zeros. Hence:
#
#   _mod_atan(0.0, -0.0) = _mod_atan(-0.0, 0.0) = 0.0
#
# The signed zero can lead to problems when converting from DCM to Euler angles.
_mod_atan(y::T,x::T) where T<:Number = atan(y + T(0), x + T(0))

# This modified function computes the `acos(x)` if `|x| <= 1` and computes
# `acos( sign(x) )`  if `|x| > 1` to avoid numerical errors when converting DCM
# to  Euler Angles.
function _mod_acos(x::T) where T<:Number
    if x > 1
        return T(0)
    elseif x < -1
        return T(π)
    else
        return acos(x)
    end
end

# This modified function computes the `asin(x)` if `|x| <= 1` and computes
# `asin( sign(x) )`  if `|x| > 1` to avoid numerical errors when converting DCM
# to  Euler Angles.
function _mod_asin(x::T) where T<:Number
    if x > 1
        return +T(π/2)
    elseif x < -1
        return -T(π/2)
    else
        return asin(x)
    end
end

"""
    dcm_to_angle(dcm::DCM, rot_seq::Symbol=:ZYX)

Convert the DCM `dcm` to Euler Angles (see `EulerAngles`) given a rotation
sequence `rot_seq`.

The rotation sequence is defined by a `:Symbol`. The possible values are:
`:XYX`, `XYZ`, `:XZX`, `:XZY`, `:YXY`, `:YXZ`, `:YZX`, `:YZY`, `:ZXY`, `:ZXZ`,
`:ZYX`, and `:ZYZ`. If no value is specified, then it defaults to `:ZYX`.

# Gimbal-lock and special cases

If the rotations are about three different axes, *e.g.* `:XYZ`, `:ZYX`, etc.,
then a second rotation of `±90˚` yields a gimbal-lock. This means that the
rotations between the first and third axes have the same effect. In this case,
the net rotation angle is assigned to the first rotation and the angle of the
third rotation is set to 0.

If the rotations are about two different axes, *e.g.* `:XYX`, `:YXY`, etc., then
a rotation about the duplicated axis yields multiple representations. In this
case, the entire angle is assigned to the first rotation and the third rotation
is set to 0.

# Example

```julia-repl
julia> D = DCM([1. 0. 0.; 0. 0. -1; 0. -1 0.]);

julia> dcm_to_angle(D,:XYZ)
EulerAngles{Float64}:
  R(X):   1.5708 rad (  90.0000 deg)
  R(Y):   0.0000 rad (   0.0000 deg)
  R(Z):   0.0000 rad (   0.0000 deg)

julia> D = angle_to_dcm(1, -pi/2, 2, :ZYX);

julia> dcm_to_angle(D,:ZYX)
EulerAngles{Float64}:
  R(Z):   3.0000 rad ( 171.8873 deg)
  R(Y):  -1.5708 rad ( -90.0000 deg)
  R(X):   0.0000 rad (   0.0000 deg)

julia> D = create_rotation_matrix(1,:X)*create_rotation_matrix(2,:X);

julia> dcm_to_angle(D,:XYX)
EulerAngles{Float64}:
  R(X):   3.0000 rad ( 171.8873 deg)
  R(Y):   0.0000 rad (   0.0000 deg)
  R(X):   0.0000 rad (   0.0000 deg)
```
"""
function dcm_to_angle(dcm::DCM{T}, rot_seq::Symbol=:ZYX) where T<:Number
    if rot_seq == :ZYX

        # Check for singularities.
        if !( abs(dcm[1,3]) >= 1-eps() )
            return EulerAngles(_mod_atan(+dcm[1,2],+dcm[1,1]),
                               asin(-dcm[1,3]),
                               _mod_atan(+dcm[2,3],+dcm[3,3]),
                               rot_seq)
        else
            return EulerAngles(_mod_atan(-dcm[2,1],+dcm[2,2]),
                               _mod_asin(-dcm[1,3]),
                               T(0),
                               rot_seq)
        end

    elseif rot_seq == :XYX

        # Check for singularities.
        if !( abs(dcm[1,1]) >= 1-eps() )
            return EulerAngles(_mod_atan(+dcm[1,2],-dcm[1,3]),
                               acos(+dcm[1,1]),
                               _mod_atan(+dcm[2,1],+dcm[3,1]),
                               rot_seq)
        else
            return EulerAngles(_mod_atan(+dcm[2,3],+dcm[2,2]),
                               _mod_acos(dcm[1,1]),
                               T(0),
                               rot_seq)
        end

    elseif rot_seq == :XYZ

        # Check for singularities.
        if !( abs(dcm[3,1]) >= 1-eps() )
            return EulerAngles(_mod_atan(-dcm[3,2],+dcm[3,3]),
                               asin(+dcm[3,1]),
                               _mod_atan(-dcm[2,1],+dcm[1,1]),
                               rot_seq)
        else
            return EulerAngles(_mod_atan(+dcm[2,3],+dcm[2,2]),
                               _mod_asin(+dcm[3,1]),
                               T(0),
                               rot_seq)
        end

    elseif rot_seq == :XZX

        # Check for singularities.
        if !( abs(dcm[1,1]) >= 1-eps() )
            return EulerAngles(_mod_atan(+dcm[1,3],+dcm[1,2]),
                               acos(+dcm[1,1]),
                               _mod_atan(+dcm[3,1],-dcm[2,1]),
                               rot_seq)
        else
            return EulerAngles(_mod_atan(-dcm[3,2],+dcm[3,3]),
                               _mod_acos(dcm[1,1]),
                               T(0),
                               rot_seq)
        end

    elseif rot_seq == :XZY

        # Check for singularities.
        if !( abs(dcm[2,1]) >= 1-eps() )
            return EulerAngles(_mod_atan(+dcm[2,3],+dcm[2,2]),
                               asin(-dcm[2,1]),
                               _mod_atan(+dcm[3,1],+dcm[1,1]),
                               rot_seq)
        else
            return EulerAngles(_mod_atan(-dcm[3,2],+dcm[3,3]),
                               _mod_asin(-dcm[2,1]),
                               T(0),
                               rot_seq)
        end

    elseif rot_seq == :YXY

        # Check for singularities.
        if !( abs(dcm[2,2]) >= 1-eps() )
            return EulerAngles(_mod_atan(+dcm[2,1],+dcm[2,3]),
                               acos(+dcm[2,2]),
                               _mod_atan(+dcm[1,2],-dcm[3,2]),
                               rot_seq)
        else
            return EulerAngles(_mod_atan(-dcm[1,3],+dcm[1,1]),
                               _mod_acos(dcm[2,2]),
                               T(0),
                               rot_seq)
        end

    elseif rot_seq == :YXZ

        if !( abs(dcm[3,2]) >= 1-eps() )
            return EulerAngles(_mod_atan(+dcm[3,1],+dcm[3,3]),
                               asin(-dcm[3,2]),
                               _mod_atan(+dcm[1,2],+dcm[2,2]),
                               rot_seq)
        else
            return EulerAngles(_mod_atan(-dcm[1,3],+dcm[1,1]),
                               _mod_asin(-dcm[3,2]),
                               T(0),
                               rot_seq)
        end

    elseif rot_seq == :YZX

        # Check for singularities.
        if !( abs(dcm[1,2]) >= 1-eps() )
            return EulerAngles(_mod_atan(-dcm[1,3],+dcm[1,1]),
                               asin(+dcm[1,2]),
                               _mod_atan(-dcm[3,2],+dcm[2,2]),
                               rot_seq)
        else
            return EulerAngles(_mod_atan(+dcm[3,1],+dcm[3,3]),
                               _mod_asin(+dcm[1,2]),
                               T(0),
                               rot_seq)
        end

    elseif rot_seq == :YZY

        # Check for singularities.
        if !( abs(dcm[2,2]) >= 1-eps() )
            return EulerAngles(_mod_atan(+dcm[2,3],-dcm[2,1]),
                               acos(+dcm[2,2]),
                               _mod_atan(+dcm[3,2],+dcm[1,2]),
                               rot_seq)
        else
            return EulerAngles(_mod_atan(dcm[3,1], dcm[3,3]),
                               _mod_acos(dcm[2,2]),
                               T(0),
                               rot_seq)
        end

    elseif rot_seq == :ZXY

        # Check for singularities.
        if !( abs(dcm[2,3]) >= 1-eps() )
            return EulerAngles(_mod_atan(-dcm[2,1],+dcm[2,2]),
                               asin(+dcm[2,3]),
                               _mod_atan(-dcm[1,3],+dcm[3,3]),
                               rot_seq)
        else
            return EulerAngles(_mod_atan(+dcm[1,2],+dcm[1,1]),
                               _mod_asin(+dcm[2,3]),
                               T(0),
                               rot_seq)
        end

    elseif rot_seq == :ZXZ

        # Check for singularities.
        if !( abs(dcm[3,3]) >= 1-eps() )
            return EulerAngles(_mod_atan(+dcm[3,1],-dcm[3,2]),
                               acos(+dcm[3,3]),
                               _mod_atan(+dcm[1,3],+dcm[2,3]),
                               rot_seq)
        else
            return EulerAngles(_mod_atan(dcm[1,2], dcm[1,1]),
                               _mod_acos(dcm[3,3]),
                               T(0),
                               rot_seq)
        end

    elseif rot_seq == :ZYZ

        # Check for singularities.
        if !( abs(dcm[3,3]) >= 1-eps() )
            return EulerAngles(_mod_atan(+dcm[3,2],+dcm[3,1]),
                               acos(+dcm[3,3]),
                               _mod_atan(+dcm[2,3],-dcm[1,3]),
                               rot_seq)
        else
            return EulerAngles(_mod_atan(-dcm[2,1], dcm[2,2]),
                               _mod_acos(dcm[3,3]),
                               T(0),
                               rot_seq)
        end

    else
        throw(ArgumentError("The rotation sequence :$rot_seq is not valid."))
    end
end

# Euler Angle and Axis
# ==============================================================================

"""
    dcm_to_angleaxis(dcm::DCM{T}) where T<:Number

Convert the DCM `dcm` to an Euler angle and axis representation. By convention,
the returned Euler angle will always be in the interval [0, π].

"""
function dcm_to_angleaxis(dcm::DCM{T}) where T<:Number
    cθ = (dcm[1,1] + dcm[2,2] + dcm[3,3] - 1)/2

    # Check the undefined case.
    if cθ >= 1 - eps()
        return EulerAngleAxis(T(0), SVector{3,T}(0,0,0))

    elseif cθ <= -1 + eps()
        v₁ = sqrt( (1 + dcm[1,1])/2 )
        v₂ = sqrt( (1 + dcm[2,2])/2 )
        v₃ = sqrt( (1 + dcm[3,3])/2 )

        # Compute the sign of the vector components.
        if dcm[1,2] >= 0
            if dcm[1,3] >= 0
                s₁ = s₂ = s₃ = +1
            else
                s₁ = s₂ = +1
                s₃ = -1
            end
        else
            if dcm[1,3] >= 0
                s₁ = s₃ = +1
                s₂ = -1
            else
                s₁ = +1
                s₂ = s₃ = -1
            end
        end

        return EulerAngleAxis( T(1)*π, [s₁*v₁, s₂*v₂, s₃*v₃ ] )

    else
        sθ2 = 2sqrt(1 - cθ*cθ)

        v₁ = ( dcm[2,3] - dcm[3,2] )/sθ2
        v₂ = ( dcm[3,1] - dcm[1,3] )/sθ2
        v₃ = ( dcm[1,2] - dcm[2,1] )/sθ2

        return EulerAngleAxis( acos(cθ), [v₁, v₂, v₃] )

    end
end

# Quaternions
# ==============================================================================

"""
    dcm_to_quat(dcm::DCM)

Convert the DCM `dcm` to a quaternion. The type of the quaternion will be
automatically selected by the constructor `Quaternion` to avoid `InexactError`.

# Remarks

By convention, the real part of the quaternion will always be positive.
Moreover, the function does not check if `dcm` is a valid direction cosine
matrix. This must be handle by the user.

This algorithm was obtained from:

    http://www.euclideanspace.com/maths/geometry/rotations/conversions/matrixToQuaternion/

# Example

```julia-repl
julia> dcm = angle_to_dcm(pi/2,0.0,0.0,:XYZ);

julia> q   = dcm_to_quat(dcm)
Quaternion{Float64}:
  + 0.7071067811865476 + 0.7071067811865475.i + 0.0.j + 0.0.k
```

"""
function dcm_to_quat(dcm::DCM)
    if  tr(dcm) > 0
        # f = 4*q0
        f = 2sqrt(tr(dcm)+1)

        return Quaternion(f/4,
                          (dcm[2,3]-dcm[3,2])/f,
                          (dcm[3,1]-dcm[1,3])/f,
                          (dcm[1,2]-dcm[2,1])/f)

    elseif (dcm[1,1] > dcm[2,2]) && (dcm[1,1] > dcm[3,3])
        # f = 4*q1
        f = 2sqrt(1 + dcm[1,1] - dcm[2,2] - dcm[3,3])

        # Real part.
        q0 = (dcm[2,3]-dcm[3,2])/f

        # Make sure that the real part is always positive.
        s = (q0 > 0) ? +1 : -1

        return Quaternion(s*q0,
                          s*f/4,
                          s*(dcm[1,2]+dcm[2,1])/f,
                          s*(dcm[3,1]+dcm[1,3])/f)

    elseif (dcm[2,2] > dcm[3,3])
        # f = 4*q2
        f = 2sqrt(1 + dcm[2,2] - dcm[1,1] - dcm[3,3])

        # Real part.
        q0 = (dcm[3,1]-dcm[1,3])/f

        # Make sure that the real part is always positive.
        s = (q0 > 0) ? +1 : -1

        return Quaternion(s*q0,
                          s*(dcm[1,2]+dcm[2,1])/f,
                          s*f/4,
                          s*(dcm[3,2]+dcm[2,3])/f)

    else
        # f = 4*q3
        f = 2sqrt(1 + dcm[3,3] - dcm[1,1] - dcm[2,2])

        # Real part.
        q0 = (dcm[1,2]-dcm[2,1])/f

        # Make sure that the real part is always positive.
        s = (q0 > 0) ? +1 : -1

        return Quaternion(s*q0,
                          s*(dcm[1,3]+dcm[3,1])/f,
                          s*(dcm[2,3]+dcm[3,2])/f,
                          s*f/4)
    end
end

################################################################################
#                                  Kinematics
################################################################################

"""
    ddcm(Dba::DCM, wba_b::AbstractArray)

Compute the time-derivative of the DCM `dcm` that rotates a reference frame `a`
into alignment to the reference frame `b` in which the angular velocity of `b`
with respect to `a`, and represented in `b`, is `wba_b`.

# Returns

The time-derivative of the DCM `Dba` (3x3 matrix of type `SMatrix{3,3}`).

# Example

```julia-repl
julia> D = DCM(Matrix{Float64}(I,3,3));

julia> ddcm(D,[1;0;0])
3×3 StaticArrays.SArray{Tuple{3,3},Float64,2,9}:
 0.0   0.0  0.0
 0.0   0.0  1.0
 0.0  -1.0  0.0
```

"""
function ddcm(Dba::DCM, wba_b::AbstractArray)
    # Auxiliary variable.
    w = wba_b

    # Check the dimensions.
    (length(wba_b) != 3) && throw(ArgumentError("The angular velocity vector must have three components."))

    wx = SMatrix{3,3}(  0  , -w[3], +w[2],
                      +w[3],   0  , -w[1],
                      -w[2], +w[1],   0  )'

    # Return the time-derivative.
    -wx*Dba
end


text:
# --------------------------------------------------
# IO functions

function read_map_from_file(fitsfilename::String; hdu_id::Int=1)
    hdu = FITS(fitsfilename)[hdu_id]
    data = read(hdu);
    header = read_header(hdu);
    wcs = wcs2d_from_header(header)

    idx = zeros((2, size(data, 2)))
    idx[2, :] = 1:size(data, 2)
    y_axis = pix_to_world(wcs, idx)[2, :];

    idx = zeros((2, size(data, 1)))
    idx[1, :] = 1:size(data, 1)
    x_axis = pix_to_world(wcs, idx)[1, :];

    xy_unit = u"°";

    val_unit = u"K";

    xy_resolution = abs(x_axis[1] - x_axis[2])

    return Map(data, val_unit, header, wcs, x_axis, y_axis, xy_unit, xy_resolution);
end

function read_cube_from_file(fitsfilename::String; hdu_id::Int=1)
    hdu = FITS(fitsfilename)[hdu_id]
    if ndims(hdu) == 4
        data = read(hdu, :, :, :, 1);
    else
        data = read(hdu)
    end
    header = read_header(hdu);
    wcs = wcs3d_from_header(header)

    idx = zeros((3, size(data, 3)))
    idx[3, :] = 1:size(data, 3)
    z_axis = pix_to_world(wcs, idx)[3, :];

    if ("CUNIT3" ∈ keys(header) && occursin("km", header["CUNIT3"]))
        z_unit = u"km / s";
    end
    z_unit = u"m / s"
#     z_axis = to_kps_val.(z_axis * z_unit);
#     z_unit = u"km / s";

    idx = zeros((3, size(data, 2)))
    idx[2, :] = 1:size(data, 2)
    y_axis = pix_to_world(wcs, idx)[2, :];

    idx = zeros((3, size(data, 1)))
    idx[1, :] = 1:size(data, 1)
    x_axis = pix_to_world(wcs, idx)[1, :];

    xy_unit = u"°";

    val_unit = u"K";

    xy_resolution = abs(header["CDELT1"])
    z_resolution = abs(z_axis[1] - z_axis[2])

    return SpectralCube(data, val_unit, header, wcs, x_axis, y_axis, z_axis, xy_unit, z_unit, xy_resolution, z_resolution);
end

function save(filename::AbstractString, map_::AbstractMap)
    f = FITS(filename, "w")
    write(f, map_.data, header=map_.header)
    close(f)
    return nothing
end

function save(filename::String, cube::AbstractSpectralCube)
    FITS(filename, "w") do f
        write(f, cube.data, header=cube.header)
    end
end

function header2d_of(cube::AbstractSpectralCube, line::AbstractString, content_description::AbstractString)
    hkeys, hvalues, comments = MapAndCube.wcs2d_from_header(cube.header) |> WCS.to_header |> MapAndCube.parse_fitsheader
    hkeys = ["SIMPLE" ; hkeys; "DATA"; "LINE"]
    hvalues = [true; hvalues; content_description; line];
    comments = [cube.header.comments[1]; comments; ""; ""]
    header = FITSHeader(hkeys, hvalues, comments);
    return header
end


text:
using BinaryBuilder

# Collection of sources required to build zlib
sources = [
  "https://zlib.net/zlib-1.2.11.tar.gz" =>
  "c3e5e9fdd5004dcb542feda5ee4f0ff0744628baf8ed2dd5d66f8ca1197cb1a1",
]

# Bash recipe for building across all platforms
script = raw"""
cd $WORKSPACE/srcdir
cd zlib-1.2.11/

# On windows platforms, our ./configure and make invocations differ a bit
if [[ ${target} == *-w64-mingw* ]]; then
    EXTRA_CONFIGURE_FLAGS="--sharedlibdir=${prefix}/bin"
    EXTRA_MAKE_FLAGS="SHAREDLIB=libz.dll SHAREDLIBM=libz-1.dll SHAREDLIBV=libz-1.2.11.dll LDSHAREDLIBC= "
fi

./configure ${EXTRA_CONFIGURE_FLAGS} --prefix=${prefix}
make install ${EXTRA_MAKE_FLAGS} -j${nproc}
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line.
platforms = [
    Windows(:i686),
    Windows(:x86_64),
    Linux(:i686, :glibc),
    Linux(:x86_64, :glibc),
    Linux(:aarch64, :glibc),
    Linux(:armv7l, :glibc),
    Linux(:powerpc64le, :glibc),
    MacOS()
]

# The products that we will ensure are always built
products = prefix -> [
    LibraryProduct(prefix,"libz", :libz),
]

# Dependencies that must be installed before this package can be built
dependencies = [
]

build_tarballs(ARGS, "Zlib", sources, script, platforms, products, dependencies)


text:
import Observables
import AbstractPlotting
import GLMakie
using Porta

plane = ComplexPlane(Quaternion(rand(), ℝ³(rand(3)))) # Construct a unit Quaternion
w, z = plane.z₁, plane.z₂
B = Complex(10rand(), 10rand())
clifford = Clifford(w, z, B)

center = ℝ³(0, 0, 0)
q = Biquaternion(center)
radius = 1.0
scene = AbstractPlotting.Scene(center = false, show_axis = false)
segments = 36
color = AbstractPlotting.RGBAf0(0.9, 0.9, 0.9, 0.3)
transparency = false
sphere = Sphere(q,
                scene,
                radius = radius,
                segments = segments,
                color = color,
                transparency = transparency)

center = ℝ³(Cartesian(clifford.base))
q = Biquaternion(center)
radius = 0.05
color = AbstractPlotting.RGBAf0(0.1, 0.1, 0.1, 0.7)
transparency = false
point = Sphere(q,
               scene,
               radius = radius,
               segments = segments,
               color = color,
               transparency = transparency)

tail, head = ℝ³(0, 0, 0), ℝ³(1, 0, 0)
width = 3
color = :gold
arrow = Arrow(tail,
              head,
              scene,
              width = width,
              color = color)

lspaceα = range(float(0), stop = float(pi / 2), length = 100)
lspaceϕ₁ = range(float(0), stop = float(2pi), length = 100)
lspaceϕ₂ = range(float(0), stop = float(4pi), length = 100)
sliderα, valα = AbstractPlotting.textslider(lspaceα, "α", start = 0)
sliderϕ₁, valϕ₁ = AbstractPlotting.textslider(lspaceϕ₁, "ϕ₁", start = 0)
sliderϕ₂, valϕ₂ = AbstractPlotting.textslider(lspaceϕ₂, "ϕ₂", start = 0)
sliders = [sliderα, sliderϕ₁, sliderϕ₂]

for slider in sliders
    Observables.on(slider[end].value) do x
        α = Observables.to_value(valα)
        ϕ₁ = Observables.to_value(valϕ₁)
        ϕ₂ = Observables.to_value(valϕ₂)
        q = Quaternion(cos((ϕ₁ + ϕ₂) / 2) * sin(α),
                       sin((ϕ₁ + ϕ₂) / 2) * sin(α),
                       cos((ϕ₂ - ϕ₁) / 2) * cos(α),
                       sin((ϕ₂ - ϕ₁) / 2) * cos(α))
        u, v, x, y = vec(q)
        cp = ComplexPlane(q)
        w, z = cp.z₁, cp.z₂
        B = w / z
        global clifford = Clifford(w, z, B)
        bq = Biquaternion(q, ℝ³(Cartesian(ComplexLine(B))))
        update(point, bq)
        tail = gettranslation(bq)
        head = ℝ³(4(x * y - u * v), 2 - 4(v^2 + y^2), -4(v * x + u * y))
        update(arrow, tail, head)
    end
end

parentscene = AbstractPlotting.Scene(resolution = (800, 600))
leftpanel = AbstractPlotting.hbox(sliders...)
AbstractPlotting.vbox(leftpanel,
                      scene,
                      parent = parentscene)


text:
# Notational note
# JMW's dx_history <=> NW's S
# JMW's dg_history <=> NW's Y

# Here alpha is a cache that parallels betas
# It is not the step-size
# q is also a cache
function twoloop!(s::Vector,
                  gr::Vector,
                  rho::Vector,
                  dx_history::Matrix,
                  dg_history::Matrix,
                  m::Integer,
                  pseudo_iteration::Integer,
                  alpha::Vector,
                  q::Vector,
                  precon)
    # Count number of parameters
    n = length(s)

    # Determine lower and upper bounds for loops
    lower = pseudo_iteration - m
    upper = pseudo_iteration - 1

    # Copy gr into q for backward pass
    copy!(q, gr)

    # Backward pass
    for index in upper:-1:lower
        if index < 1
            continue
        end
        i = mod1(index, m)
        @inbounds alpha[i] = rho[i] * vecdot(view(dx_history, :, i), q)
        @simd for j in 1:n
            @inbounds q[j] -= alpha[i] * dg_history[j, i]
        end
    end

    # Copy q into s for forward pass
    # apply preconditioner if precon != nothing
    # (Note: preconditioner update was done outside of this function)
    A_ldiv_B!(s, precon, q)

    # Forward pass
    for index in lower:1:upper
        if index < 1
            continue
        end
        i = mod1(index, m)
        @inbounds beta = rho[i] * vecdot(view(dg_history, :, i), s)
        @simd for j in 1:n
            @inbounds s[j] += dx_history[j, i] * (alpha[i] - beta)
        end
    end

    # Negate search direction
    scale!(s, -1)

    return
end

immutable LBFGS{T, L<:Function, Tprep<:Union{Function, Void}} <: Optimizer
    m::Int
    linesearch!::L
    P::T
    precondprep!::Tprep
    extrapolate::Bool
    snap2one::Tuple
end
#= uncomment for v0.8.0
LBFGS(; m::Integer = 10, linesearch = LineSearches.hagerzhang!,
                        P=nothing, precondprep = (P, x) -> nothing,
                        extrapolate::Bool=false, snap2one = (0.75, Inf)) =
      LBFGS(Int(m), linesearch, P, precondprep, extrapolate, snap2one)
=#

function LBFGS(; linesearch! = nothing,
                 m::Integer = 10,
                 linesearch = LineSearches.hagerzhang!,
                 P=nothing,
                 precondprep! = nothing,
                 precondprep = (P, x) -> nothing,
                 extrapolate::Bool=false,
                 snap2one = (0.75, Inf))

    linesearch = get_linesearch(linesearch!, linesearch)
    precondprep = get_precondprep(precondprep!, precondprep)
    LBFGS(Int(m), linesearch, P, precondprep, extrapolate, snap2one)
end

type LBFGSState{T}
    @add_generic_fields()
    x_previous::Array{T}
    g::Array{T}
    g_previous::Array{T}
    rho::Array{T}
    dx_history::Array{T}
    dg_history::Array{T}
    dx::Array{T}
    dg::Array{T}
    u::Array{T}
    f_x_previous::T
    twoloop_q
    twoloop_alpha
    pseudo_iteration::Int64
    s::Array{T}
    @add_linesearch_fields()
end

function initial_state{T}(method::LBFGS, options, d, initial_x::Array{T})
    n = length(initial_x)
    g = Array(T, n)
    f_x = d.fg!(initial_x, g)
    # Maintain a cache for line search results
    # Trace the history of states visited
    LBFGSState("L-BFGS",
              n,
              copy(initial_x), # Maintain current state in state.x
              f_x, # Store current f in state.f_x
              1, # Track f calls in state.f_calls
              1, # Track g calls in state.g_calls
              0, # Track h calls in state.h_calls
              copy(initial_x), # Maintain current state in state.x_previous
              g, # Store current gradient in state.g
              copy(g), # Store previous gradient in state.g_previous
              Array{T}(method.m), # state.rho
              Array{T}(n, method.m), # Store changes in position in state.dx_history
              Array{T}(n, method.m), # Store changes in gradient in state.dg_history
              Array{T}(n), # Buffer for new entry in state.dx_history
              Array{T}(n), # Buffer for new entry in state.dg_history
              Array{T}(n), # Buffer stored in state.u
              T(NaN), # Store previous f in state.f_x_previous
              Array{T}(n), #Buffer for use by twoloop
              Array{T}(method.m), #Buffer for use by twoloop
              0,
              Array{T}(n), # Store current search direction in state.s
              @initial_linesearch()...) # Maintain a cache for line search results in state.lsr
end

function update_state!{T}(d, state::LBFGSState{T}, method::LBFGS)
    lssuccess = true
    n = state.n
    # Increment the number of steps we've had to perform
    state.pseudo_iteration += 1

    # update the preconditioner
    method.precondprep!(method.P, state.x)

    # Determine the L-BFGS search direction # FIXME just pass state and method?
    twoloop!(state.s, state.g, state.rho, state.dx_history, state.dg_history,
             method.m, state.pseudo_iteration,
             state.twoloop_alpha, state.twoloop_q, method.P)

    # Refresh the line search cache
    dphi0 = vecdot(state.g, state.s)
    if dphi0 > 0.0
        state.pseudo_iteration = 1
        @simd for i in 1:n
            @inbounds state.s[i] = -state.g[i]
        end
        dphi0 = vecdot(state.g, state.s)
    end

    LineSearches.clear!(state.lsr)
    push!(state.lsr, zero(T), state.f_x, dphi0)

    # compute an initial guess for the linesearch based on
    # Nocedal/Wright, 2nd ed, (3.60)
    # TODO: this is a temporary fix, but should eventually be split off into
    #       a separate type and possibly live in LineSearches; see #294
    if method.extrapolate && state.pseudo_iteration > 1
        alphaguess = 2.0 * (state.f_x - state.f_x_previous) / dphi0
        alphaguess = max(alphaguess, state.alpha/4.0)  # not too much reduction
        # if alphaguess ≈ 1, then make it 1 (Newton-type behaviour)
        if method.snap2one[1] < alphaguess < method.snap2one[2]
            alphaguess = 1.0
        end
    else
        # without extrapolation use previous alpha (old behaviour)
        alphaguess = state.alpha
    end

    # Determine the distance of movement along the search line
    try
        state.alpha, f_update, g_update =
        method.linesearch!(d, state.x, state.s, state.x_ls, state.g_ls, state.lsr,
                           alphaguess, state.mayterminate)
        state.f_calls, state.g_calls = state.f_calls + f_update, state.g_calls + g_update
    catch ex
        if isa(ex, LineSearches.LineSearchException)
            lssuccess = false
            state.f_calls, state.g_calls = state.f_calls + ex.f_update, state.g_calls + ex.g_update
            state.alpha = ex.alpha
            Base.warn("Linesearch failed, using alpha = $(state.alpha) and exiting optimization.")
        else
            rethrow(ex)
        end
    end

    # Maintain a record of previous position
    copy!(state.x_previous, state.x)

    # Update current position
    @simd for i in 1:n
        @inbounds state.dx[i] = state.alpha * state.s[i]
        @inbounds state.x[i] = state.x[i] + state.dx[i]
    end

    # Save old f and g values to prepare for update_g! call
    state.f_x_previous = state.f_x
    copy!(state.g_previous, state.g)
    (lssuccess == false) # break on linesearch error
end

function update_h!(d, state, method::LBFGS)
    # Measure the change in the gradient
    @simd for i in 1:state.n
        @inbounds state.dg[i] = state.g[i] - state.g_previous[i]
    end

    # Update the L-BFGS history of positions and gradients
    rho_iteration = 1 / vecdot(state.dx, state.dg)
    if isinf(rho_iteration)
        # TODO: Introduce a formal error? There was a warning here previously
        return true
    end
    state.dx_history[:, mod1(state.pseudo_iteration, method.m)] = state.dx
    state.dg_history[:, mod1(state.pseudo_iteration, method.m)] = state.dg
    state.rho[mod1(state.pseudo_iteration, method.m)] = rho_iteration

end


text:
# MLJ interface for clustering models

import MLJModelInterface       # It seems that having done this in the top module is not enought
const MMI = MLJModelInterface  # We need to repoeat it here

export KMeans, KMedoids, GMMClusterer, MissingImputator

# ------------------------------------------------------------------------------
# Model Structure declarations..

mutable struct KMeans <: MMI.Unsupervised
   K::Int64
   dist::Function
   initStrategy::String
   Z₀::Union{Nothing,Matrix{Float64}}
   rng::AbstractRNG
end
KMeans(;
   K            = 3,
   dist         = dist=(x,y) -> norm(x-y),
   initStrategy = "shuffle",
   Z₀           = nothing,
   rng          = Random.GLOBAL_RNG,
 ) = KMeans(K,dist,initStrategy,Z₀,rng)

 mutable struct KMedoids <: MMI.Unsupervised
    K::Int64
    dist::Function
    initStrategy::String
    Z₀::Union{Nothing,Matrix{Float64}}
    rng::AbstractRNG
 end
 KMedoids(;
    K            = 3,
    dist         = dist=(x,y) -> norm(x-y),
    initStrategy = "shuffle",
    Z₀           = nothing,
    rng          = Random.GLOBAL_RNG,
  ) = KMedoids(K,dist,initStrategy,Z₀,rng)

mutable struct GMMClusterer <: MMI.Unsupervised
  K::Int64
  p₀::Union{Nothing,AbstractArray{Float64,1}}
  mixtures::Symbol
  tol::Float64
  minVariance::Float64
  minCovariance::Float64
  initStrategy::String
  rng::AbstractRNG
end
GMMClusterer(;
    K             = 3,
    p₀            = nothing,
    mixtures      = :diag_gaussian,
    tol           = 10^(-6),
    minVariance   = 0.05,
    minCovariance = 0.0,
    initStrategy  = "kmeans",
    rng           = Random.GLOBAL_RNG,
) = GMMClusterer(K,p₀,mixtures, tol, minVariance, minCovariance,initStrategy,rng)

mutable struct MissingImputator <: MMI.Unsupervised
    K::Int64
    p₀::Union{Nothing,AbstractArray{Float64,1}}
    mixtures::Symbol
    tol::Float64
    minVariance::Float64
    minCovariance::Float64
    initStrategy::String
    rng::AbstractRNG
end
MissingImputator(;
    K             = 3,
    p₀            = nothing,
    mixtures      = :diag_gaussian,
    tol           = 10^(-6),
    minVariance   = 0.05,
    minCovariance = 0.0,
    initStrategy  = "kmeans",
    rng           = Random.GLOBAL_RNG,
) = MissingImputator(K,p₀,mixtures, tol, minVariance, minCovariance,initStrategy,rng)

# ------------------------------------------------------------------------------
# Fit functions...
function MMI.fit(m::Union{KMeans,KMedoids}, verbosity, X)
    x  = MMI.matrix(X)                        # convert table to matrix
    if typeof(m) == KMeans
        (assignedClasses,representatives) = kmeans(x,m.K,dist=m.dist,initStrategy=m.initStrategy,Z₀=m.Z₀,rng=m.rng)
    else
        (assignedClasses,representatives) = kmedoids(x,m.K,dist=m.dist,initStrategy=m.initStrategy,Z₀=m.Z₀,rng=m.rng)
    end
    cache=nothing
    report=nothing
    return ((classes=assignedClasses,centers=representatives,distanceFunction=m.dist), cache, report)
end
MMI.fitted_params(model::Union{KMeans,KMedoids}, fitresult) = (centers=fitesult[2], cluster_labels=CategoricalArrays.categorical(fitresults[1]))

function MMI.fit(m::GMMClusterer, verbosity, X)
    # X is nothing, y is the data: https://alan-turing-institute.github.io/MLJ.jl/dev/adding_models_for_general_use/#Models-that-learn-a-probability-distribution-1
    x          = MMI.matrix(X) # convert table to matrix
    if m.mixtures == :diag_gaussian
        mixtures = [DiagonalGaussian() for i in 1:m.K]
    elseif m.mixtures == :full_gaussian
        mixtures = [FullGaussian() for i in 1:m.K]
    elseif m.mixtures == :spherical_gaussian
        mixtures = [SphericalGaussian() for i in 1:m.K]
    else
        error("Usupported mixture. Supported mixtures are either `:diag_gaussian`, `:full_gaussian` or `:spherical_gaussian`.")
    end
    res        = gmm(x,m.K,p₀=deepcopy(m.p₀),mixtures=mixtures, minVariance=m.minVariance, minCovariance=m.minCovariance,initStrategy=m.initStrategy,verbosity=NONE,rng=m.rng)
    fitResults = (pₖ=res.pₖ,mixtures=res.mixtures) # res.pₙₖ
    cache      = nothing
    report     = (res.ϵ,res.lL,res.BIC,res.AIC)
    return (fitResults, cache, report)
end
MMI.fitted_params(model::GMMClusterer, fitresult) = (weights=fitesult.pₖ, mixtures=fitresult.mixtures)

function MMI.fit(m::MissingImputator, verbosity, X)
    x          = MMI.matrix(X) # convert table to matrix
    if m.mixtures == :diag_gaussian
        mixtures = [DiagonalGaussian() for i in 1:m.K]
    elseif m.mixtures == :full_gaussian
        mixtures = [FullGaussian() for i in 1:m.K]
    elseif m.mixtures == :spherical_gaussian
        mixtures = [SphericalGaussian() for i in 1:m.K]
    else
        error("Usupported mixture. Supported mixtures are either `:diag_gaussian`, `:full_gaussian` or `:spherical_gaussian`.")
    end
    res        = gmm(x,m.K,p₀=deepcopy(m.p₀),mixtures=mixtures, minVariance=m.minVariance, minCovariance=m.minCovariance,initStrategy=m.initStrategy,verbosity=NONE,rng=m.rng)
    fitResults = (pₖ=res.pₖ,mixtures=res.mixtures) # pₙₖ=res.pₙₖ
    cache      = nothing
    report     = (res.ϵ,res.lL,res.BIC,res.AIC)
    return (fitResults, cache, report)
end

# ------------------------------------------------------------------------------
# Transform functions...

""" fit(m::KMeans, fitResults, X) - Given a trained clustering model and some observations, return the distances to each centroids """
function MMI.transform(m::Union{KMeans,KMedoids}, fitResults, X)
    x     = MMI.matrix(X) # convert table to matrix
    (N,D) = size(x)
    nCl   = size(fitResults.centers,1)
    distances = Array{Float64,2}(undef,N,nCl)
    for n in 1:N
        for c in 1:nCl
            distances[n,c] = fitResults.distanceFunction(x[n,:],fitResults[2][c,:])
        end
    end
    return MMI.table(distances)
end

""" transform(m::MissingImputator, fitResults, X) - Given a trained imputator model fill the missing data of some new observations"""
function MMI.transform(m::MissingImputator, fitResults, X)
    x             = MMI.matrix(X) # convert table to matrix
    (N,D)         = size(x)
    (pₖ,mixtures) = fitResults.pₖ, fitResults.mixtures   #
    nCl           = length(pₖ)
    # Fill the missing data of this "new X" using the mixtures computed in the fit stage
    xout          = predictMissing(x,nCl,p₀=pₖ,mixtures=mixtures,tol=m.tol,verbosity=NONE,minVariance=m.minVariance,minCovariance=m.minCovariance,initStrategy="given",maxIter=1,rng=m.rng)
    return MMI.table(xout.X̂)
end

# ------------------------------------------------------------------------------
# Predict functions...

""" predict(m::KMeans, fitResults, X) - Given a trained clustering model and some observations, predict the class of the observation"""
function MMI.predict(m::Union{KMeans,KMedoids}, fitResults, X)
    x               = MMI.matrix(X) # convert table to matrix
    (N,D)           = size(x)
    nCl             = size(fitResults.centers,1)
    distances       = MMI.matrix(MMI.transform(m, fitResults, X))
    mindist         = argmin(distances,dims=2)
    assignedClasses = [Tuple(mindist[n,1])[2]  for n in 1:N]
    return CategoricalArray(assignedClasses,levels=1:nCl)
end

function MMI.predict(m::GMMClusterer, fitResults, X)
    x               = MMI.matrix(X) # convert table to matrix
    (N,D)           = size(x)
    (pₖ,mixtures)   = (fitResults.pₖ, fitResults.mixtures)
    nCl             = length(pₖ)
    # Compute the probabilities that maximise the likelihood given existing mistures and a single iteration (i.e. doesn't update the mixtures)
    thisOut         = gmm(x,nCl,p₀=pₖ,mixtures=mixtures,tol=m.tol,verbosity=NONE,minVariance=m.minVariance,minCovariance=m.minCovariance,initStrategy="given",maxIter=1,rng=m.rng)
    classes         = CategoricalArray(1:nCl)
    predictions     = MMI.UnivariateFinite(classes, thisOut.pₙₖ)
    return predictions
end

# ------------------------------------------------------------------------------
# Model metadata for registration in MLJ...

MMI.metadata_model(KMeans,
    input_scitype    = MMI.Table(MMI.Continuous),         # scitype of the inputs
    output_scitype   = MMI.Table(MMI.Continuous),         # scitype of the output of `transform`
    target_scitype   = AbstractArray{<:MMI.Multiclass},   # scitype of the output of `predict`
    supports_weights = false,                             # does the model support sample weights?
    descr            = "The classical KMeans clustering algorithm, from the Beta Machine Learning Toolkit (BetaML).",
	load_path        = "BetaML.Clustering.KMeans"
)

MMI.metadata_model(KMedoids,
    input_scitype    = MMI.Table(MMI.Continuous),         # scitype of the inputs
    output_scitype   = MMI.Table(MMI.Continuous),         # scitype of the output of `transform`
    target_scitype   = AbstractArray{<:MMI.Multiclass},   # scitype of the output of `predict`
    supports_weights = false,                             # does the model support sample weights?
    descr            = "The K-medoids clustering algorithm with customisable distance function, from the Beta Machine Learning Toolkit (BetaML).",
	load_path        = "BetaML.Clustering.KMedoids"
)

MMI.metadata_model(GMMClusterer,
    input_scitype    = MMI.Table(Union{MMI.Continuous,MMI.Missing}),
    output_scitype   = AbstractArray{<:MMI.Multiclass},       # scitype of the output of `transform`
    target_scitype   = AbstractArray{<:MMI.Multiclass},       # scitype of the output of `predict`
    #prediction_type  = :probabilistic,  # option not added to metadata_model function, need to do it separately
    supports_weights = false,                                 # does the model support sample weights?
    descr            = "A Expectation-Maximisation clustering algorithm with customisable mixtures, from the Beta Machine Learning Toolkit (BetaML).",
	load_path        = "BetaML.Clustering.GMMClusterer"
)
MMI.prediction_type(::Type{<:GMMClusterer}) = :probabilistic

MMI.metadata_model(MissingImputator,
    input_scitype    = MMI.Table(Union{MMI.Continuous,MMI.Missing}),
    output_scitype   = MMI.Table(MMI.Continuous),     # for an unsupervised, what output?
    supports_weights = false,                         # does the model support sample weights?
    descr            = "Impute missing values using an Expectation-Maximisation clustering algorithm, from the Beta Machine Learning Toolkit (BetaML).",
	load_path        = "BetaML.Clustering.MissingImputator"
)


text:
# crank_nicolson.jl
# Crank-Nicolson implicit solver for PDEs

export crank_nicolson_uadj, negativeZeroCheck

push!(LOAD_PATH, joinpath(Pkg.dir("PumiInterface"), "src"))
push!(LOAD_PATH, joinpath(Pkg.dir("PDESolver"), "src/solver/euler"))
push!(LOAD_PATH, joinpath(Pkg.dir("PDESolver"), "src/NonlinearSolvers"))
push!(LOAD_PATH, joinpath(Pkg.dir("PDESolver"), "src/Debugging"))
push!(LOAD_PATH, joinpath(Pkg.dir("PDESolver"), "src/Utils"))

"""
crank_nicolson

  This function performs Crank-Nicolson implicit time solution, using a function
  of the form du/dt = f(u, t)

  Arguments:
    * physics_func  : function evaluation, must have signature (ctx..., opts, t)
    * h  : time step size
    * t_max: time value to stop time stepping (time starts at 0)
    * q_vec: vector of the u values
    * res_vec: vector of du/dt values (the output of the function physics_func)
    * pre_func: function to to be called after the new u values are put into
                q_vec but before the function f is evaluated.  Must have
                signature: pre_func(ctx..., opts)
    * post_func: function called immediately after f is called.  The function
                 must have the signature res_norm = post_func(ctx..., opts,
                 calc_norm=true),
                 where res_norm is a norm of res_vec, and calc_norm determines
                 whether or not to calculate the norm.
    * ctx: a tuple (or any iterable container) of the objects needed by
           f, pre_func, and post func.  The tuple is splatted before being
           passed to the functions.
    * opts : options dictionary

    Keyword Arguments:
    * majorIterationCallback: a callback function called after the first
                              stage, useful to do output or logging
    * res_tol : keyword arg, residual topping tolerance
    * neg_time : step through time in the negative direction,
                 starting at t_max, stepping with h, and ending at 0.0.

   The eqn.q_vec should hold the whichever variables (conservative or
   entropy) that the simulation should use.

   The idea behind pre_func, post func, and ctx is that they abstract what kind
   of system CN is timestepping. CN only needs to know about q_vec and
   res_vec.

   For physics modules, ctx should be (mesh, sbp, eqn) and q_vec and res_vec
   should be eqn.q_vec and eqn.res_vec.

   TODO: fully document eqn/eqn_nextstep
"""
# function crank_nicolson(physics_func::Function, h::AbstractFloat, t_max::AbstractFloat,
                        # mesh::AbstractMesh{Tmsh}, sbp::AbstractSBP, eqn::AbstractSolutionData{Tsol},
                        # opts, res_tol=-1.0; neg_time=false, obj_fn=obj_zero, store_u_to_disk=false) where {Tmsh, Tsol}
function crank_nicolson_uadj(physics_func::Function, h::AbstractFloat, t_max::AbstractFloat,
                        mesh::AbstractMesh{Tmsh}, sbp::AbstractSBP, eqn::AbstractSolutionData{Tsol},
                        opts,
                        WWW, ZZZ, dRdu_global_fwd, dRdu_global_rev, dRdu_global_rev_PM,
                        res_tol=-1.0; neg_time=false, obj_fn=obj_zero, store_u_to_disk=false) where {Tmsh, Tsol}
                        # NEWNEW: neg_time, obj_fn, store_u_to_disk
  #----------------------------------------------------------------------
  if opts["uadj_global"]
    println(" GLOBAL, CN start: size(dRdu_global_fwd): ", size(dRdu_global_fwd))
    println(" GLOBAL, CN start: size(dRdu_global_rev): ", size(dRdu_global_rev))
    println(" GLOBAL, CN start: size(WWW): ", size(WWW))
    println(" GLOBAL, CN start: size(ZZZ): ", size(ZZZ))
  end

  myrank = MPI.Comm_rank(MPI.COMM_WORLD)
  fstdout = BufferedIO(STDOUT)
  if myrank == 0
    println(fstdout, "\nEntered Crank-Nicolson")
    println(fstdout, "res_tol = ", res_tol)
  end

  output_freq = opts["output_freq"]::Int
  write_vis = opts["write_vis"]::Bool
  use_itermax = opts["use_itermax"]::Bool
  jac_type = opts["jac_type"]
  if use_itermax
    itermax = opts["itermax"]
  end

  if jac_type == 4
    error("CN not implemented for matrix-free ops. (jac_type cannot be 4)")
  end

  if myrank == 0
    _f1 = open("convergence.dat", "a+")
    f1 = BufferedIO(_f1)
  end

  # calculate t_steps, the number of time steps that CN will take
  t_steps = floor(Int, t_max/h)   # this allows t_max to be any value up to h greater the final time step
  # TODO: Ideally should be t_steps*h for clarity, issue #92. Should be fixed for RK4 also.
  # time_of_final_step = (t_steps-1)*h
  time_of_final_step = (t_steps)*h
  # changed from (t_steps-1)*h, 20170908

  if neg_time == false    # negative time is for unsteady adjoint
    t = 0.0     # start time at 0.0
  else
    t = time_of_final_step
    # opts["use_Minv_override_for_uadj"] = true
  end

  # NOTE: WWW, ZZZ, dRdu_global_fwd & rev formed in initialization.jl

  if neg_time == false
    # make a copy of the eqn object for storage of t_(n+1) information
    eqn_nextstep = eqn_deepcopy(mesh, sbp, eqn, opts)
    # TODO: copyForMultistage! does not give correct values.
    #     deepcopy works for now, but uses more memory than copyForMultistage!, if it worked
    # eqn_nextstep = copyForMultistage!(eqn)
    # eqn_nextstep.q = reshape(eqn_nextstep.q_vec, mesh.numDofPerNode, mesh.numNodesPerElement, mesh.numEl)
    # eqn_nextstep.res = reshape(eqn_nextstep.res_vec, mesh.numDofPerNode, mesh.numNodesPerElement, mesh.numEl)
  else
    # Initialize adjoint pde eqn objects
    adj = eqn_deepcopy(mesh, sbp, eqn, opts)
    # adj.q = reshape(adj.q_vec, mesh.numDofPerNode, mesh.numNodesPerElement, mesh.numEl)
    # adj.res = reshape(adj.res_vec, mesh.numDofPerNode, mesh.numNodesPerElement, mesh.numEl)

    adj_nextstep = eqn_deepcopy(mesh, sbp, eqn, opts)
    # adj_nextstep.q = reshape(adj_nextstep.q_vec, mesh.numDofPerNode, mesh.numNodesPerElement, mesh.numEl)
    # adj_nextstep.res = reshape(adj_nextstep.res_vec, mesh.numDofPerNode, mesh.numNodesPerElement, mesh.numEl)
  end

  println("========================= Starting CN =========================")
  println("   neg_time: ", neg_time)
  println("   t = ", round(t,3))
  println("   t_max = ", t_max)
  println("   t_steps = ", t_steps)
  println("   eqn.params.alpha_x: ",eqn.params.alpha_x)
  println("   eqn.params.alpha_y: ",eqn.params.alpha_y)

  #-------------------------------------------------------------------------------
  # allocate Jac outside of time-stepping loop
  # these jacs are for full CN or CN adj jac
  if neg_time == false
    newton_data, jac, rhs_vec = setupNewton(mesh, mesh, sbp, eqn, opts)
  else
    newton_data, jac, rhs_vec = setupNewton(mesh, mesh, sbp, adj, opts)
  end

  # Setting IC for reverse sweep
  if neg_time == true

    #----------------
    # this section:
    #   1) reads the checkpointed q_vec at the last time step of the forward sweep (n'th time step)
    #   2) uses calcJacobianComplex to calculate dRdu at time step n
    i_fwd = t_steps + 2  # index of the last time step's checkpoint.
                         #  It was saved after the end of the time stepping loop during the forward sweep

    println("\n--- Adj IC: setting i_fwd. ---")
    println("     t_steps = ", t_steps)
    println("     i_fwd = ", i_fwd)
    println("     t = ", round(t,3))
    println("     t_max = ", t_max)
    println("     h = ", h)

    # load checkpoint to calculate dRdu at this time step
    println(" ")
    println("  Setting IC for reverse sweep, i_fwd (forward sweep time step index): ", i_fwd)
    eqn_fwd = cnAdjLoadChkpt(mesh, sbp, opts, adj, physics_func, i_fwd)
    check_q_qvec_consistency(mesh, sbp, eqn_fwd, opts)

    # 20170907: t being passed in is the wrong time.
    t_ic = t_max

    jac = cnAdjCalcdRdu(mesh, sbp, opts, eqn_fwd, physics_func, i_fwd, t_ic)
    dRdu_n = jac      # TODO: check transpose
    # println(" size of dRdu: ", size(dRdu_n))
    #----------------

    dJdu_CS = calcdJdu_CS(mesh, sbp, eqn_fwd, opts, h, t_ic)  # obtain dJdu at time step n
    dJdu_FD = calcdJdu_FD(mesh, sbp, eqn_fwd, opts, h, t_ic)  # obtain dJdu at time step n
    dJdu = dJdu_CS
    dJdu_analytical = calcObjectiveFn(mesh, sbp, eqn_fwd, opts, h, t_ic, isDeriv=true)

    #---------------------------------------------------------------------------------------------
    # testing calcObjectiveFn
    eqn_poly = eqn_deepcopy(mesh, sbp, eqn_fwd, opts)
    for dof_ix = 1:mesh.numDof
      st = ind2sub((mesh.numDofPerNode, mesh.numNodesPerElement, mesh.numEl), dof_ix)
      dofnum = st[1]
      nodenum = st[2]
      elnum = st[3]
      if getindex(eqn_poly.q, dofnum, nodenum, elnum) != eqn_poly.q_vec[dof_ix]
        error("problem with ind2sub")
      end
      coords_this_dof = getindex(mesh.coords, :, nodenum, elnum)
      x_coord = coords_this_dof[1]
      y_coord = coords_this_dof[2]

      # u = 2(x+y)
      eqn_poly.q_vec[dof_ix] = 2*(x_coord+y_coord)

      # u = 3
      # eqn_poly.q_vec[dof_ix] = 3.0
    end
    array3DTo1D(mesh, sbp, eqn, opts, eqn_poly.q, eqn_poly.q_vec)
    vis_filename = string("poly_soln")
    saveSolutionToMesh(mesh, real(eqn_poly.q_vec))
    writeVisFiles(mesh, vis_filename)
    J_poly = calcObjectiveFn(mesh, sbp, eqn_poly, opts, h, t_ic, isDeriv=false)
    println(";;;;;;;;;;;;;;;;;;;;;;;;;; J_poly = ", J_poly)
    #---------------------------------------------------------------------------------------------

    J_ic = calcObjectiveFn(mesh, sbp, eqn_fwd, opts, h, t_ic, isDeriv=false)
    print_qvec_coords(mesh, sbp, eqn, opts, to_file=true, filename="IC_qvec_coords.dat", other_field=eqn_fwd.q_vec)
    writedlm("IC_J.dat", J_ic)
    writedlm("IC_qvec.dat", eqn_fwd.q_vec)
    writedlm("IC_dJdu_CS.dat", dJdu_CS)
    writedlm("IC_dJdu_FD.dat", dJdu_FD)
    writedlm("IC_dJdu_analytical.dat", reshape(dJdu_analytical, (mesh.numDof, 1)))
    writedlm("IC_dRdu_n.dat", real(dRdu_n))

    # now that dRdu and dJdu at time step n has been obtained, we can now set the IC for the adjoint eqn
    # println("size of eqn_fwd.q_vec: ", size(eqn_fwd.q_vec))
    # println("size of dRdu_n: ", size(dRdu_n))
    # println("size of dJdu: ", size(dJdu))
    # println("size of adj.q_vec: ", size(adj.q_vec))

    # note: matrix calculus says that the derivative of a scalar by a vector is a row vector.
    #   therefore dJdu is a row vector. However, it is implemented in the code as a column vector.
    #   because of this, the final transpose of psi is not required.
    #
    # No! This is wrong. See 08/30/2017 notes for actual rigorous derivation.
    # I = eye(length(eqn_fwd.q_vec))
    # B = I - (0.5*h*dRdu_n)
    # psi = transpose(B\(-dJdu))
    # psi = B\(-dJdu)

    println(" ::::: in IC, norm(dRdu_n, 1): ", norm(dRdu_n,1))

    # 20170831 old, before xpose fix
    I = eye(length(eqn_fwd.q_vec))
    B = (I - (h/2) * (dRdu_n))
    psi = transpose(B)\(-dJdu)

    # println("size of psi: ", size(psi))
    # println("size of copy(psi): ", size(copy(psi)))
    adj.q_vec = copy(psi)
    writedlm("adj_ic.dat", adj.q_vec)

    array1DTo3D(mesh, sbp, adj, opts, adj.q_vec, adj.q)     # diarray3DTo1D: q_vec -> q

    println("--- Adj IC: end ---")

  end

  #--------------------------------------------------------------------------------------------------------------------------
  # Start of CN loop
  #--------------------------------------------------------------------------------------------------------------------------
  # TODO: Ideally should be t_steps + 2 for clarity, issue #92. Should be fixed for RK4 also.
  t_steps_end = t_steps + 1
  for i = 2:t_steps_end

    println(" ")
    println(" -------------- start of this CN iter. t = ", round(t,3),", i = $i, neg_time = ", neg_time, " --------------")
    @debug1 flush(eqn.params.f)

    # println("   eqn.params.alpha_x: ",eqn.params.alpha_x)
    # println("   eqn.params.alpha_y: ",eqn.params.alpha_y)

    # Write checkpoint data at start of time step
    if neg_time == false
      # for adjoint_straight option: stores every time step's q to disk
      # Note: cannot store full eqn object without extending one of the julia write methods
      if store_u_to_disk == true
        println(" >>>>> Writing solution to disk. i = ", i,", t = ", round(t,3))
        filename = string("qvec_for_adj-", i, ".dat")
        writedlm(filename, eqn.q_vec)
        vis_filename = string("solution_storedtodisk_i-", i)
        saveSolutionToMesh(mesh, real(eqn.q_vec))
        writeVisFiles(mesh, vis_filename)

        time_filename = string("t_during_fwd_sweep_i-",i,".dat")
        writedlm(time_filename, t)              # make sure that this time value, whether t or t_nextstep,
                                                #   correctly corresponds to the q_vec stored just above
      end
    else
      # save every time step's adjoint to disk
      if opts["adjoint_saveall"]
        filename = string("adj_irev-", i, ".dat")
        writedlm(filename, adj.q_vec)

        vis_filename = string("adj_i-", i)
        saveSolutionToMesh(mesh, real(adj.q_vec))
        writeVisFiles(mesh, vis_filename)
      end
    end

    if neg_time == false
      # println(" -------------- eqn.q_vec start of this CN iter. t = $t, i = $i --------------")
      # print_qvec_coords(mesh, sbp, eqn, opts)
      # println(" -------------- J for this eqn.q_vec --------------")
      J_arr = calcObjectiveFn(mesh, sbp, eqn, opts, h, t)
      # println(" -------------- eqn.q_bndry for this eqn.q_vec --------------")
      # print_qvec_coords(mesh, sbp, eqn, opts; bndry=true)
      J = J_arr[1]
      # println("  J: ", J)
      J_filename = string("J_during_fwd_sweep_i-", i,".dat")
      writedlm(J_filename, J)
      println("   Now using eqn.q_vec to compute eqn_nextstep.q_vec.")
    else
      # println(" -------------- adj.q_vec start of this CN iter. t = $t, i = $i --------------")
      # print_qvec_coords(mesh, sbp, adj, opts)
      println("   Now using adj.q_vec to compute adj_nextstep.q_vec.")
    end
    # println(" ")

    #----------------------------
    # zero out Jac
    #   this works for both PETSc and Julia matrices.
    #   when jac is a Julia matrix, this is effectively wrapping: fill!(jac, 0.0)
    PetscMatZeroEntries(jac)

    # NOTE:
    # majorIterationCallback: called before every step of Newton's method. signature:
    #   majorIterationCallback(itr, mesh::AbstractMesh, sbp::AbstractSBP, eqn::AbstractEulerData, opts)

    # NOTE: Must include a comma in the ctx tuple to indicate tuple
    # f is the physics function, like evalEuler

    #-----------------------------------------------------------------------
    # time step update: h is passed in as argument to crank_nicolson
    if neg_time == false
      # need to add h in the forward time usage
      t_nextstep = t + h
    else
      # need to subtract h in the reverse time usage
      t_nextstep = t - h
      t_nextstep = negativeZeroCheck(t_nextstep)   # ensure negative zero is changed to zero
    end

      # NOTE: The below checks are being done for the 'current' time step. Not 'nextstep'.
      #   So the time step needs to be the one before the current ??????????????????????????????????/

    #-------------
    # objective function section
    # 1. read option to indicate which obj fun
    # 2. call it, complex step it, and store it in dJdu
    if neg_time == true
      dJdu = zeros(Tsol, length(eqn.q_vec))
      dJdu_CS = calcdJdu_CS(mesh, sbp, eqn_fwd, opts, h, t)  # obtain dJdu at time step n
      dJdu_FD = calcdJdu_FD(mesh, sbp, eqn_fwd, opts, h, t)  # obtain dJdu at time step n
      dJdu = dJdu_CS
      dJdu_analytical = calcObjectiveFn(mesh, sbp, eqn_fwd, opts, h, t, isDeriv=true)

      filename = string("dJdu_irev-",i,"_CS.dat")
      writedlm(filename, dJdu_CS)
      filename = string("dJdu_irev-",i,"_FD.dat")
      writedlm(filename, dJdu_FD)
      filename = string("dJdu_irev-",i,"_analytical.dat")
      writedlm(filename, reshape(dJdu_analytical, (mesh.numDof, 1)))

      # println("       checking direct method: size(dJdu): ", size(dJdu))

      # VV is the algebraic v, which is dudA, calculated for the advection adjoint test

      # TODO: figure out if calcVV is to be called with t or t_nextstep. I think it's t_nextstep.
      #       It is for sure t_nextstep when dJdu is called during the direct solve.
      # VV = calcVV(mesh, sbp, adj, opts, t_nextstep)     # scalar only because our x-value of interest is unchanging
      VV = calcVV(mesh, sbp, adj, opts, t)     # scalar only because our x-value of interest is unchanging
      # NOTE 20170712: think I need to be doing t, not t_nextstep. this CN rev's i corresponds to eqn_fwd's i_fwd and t

      # println("       checking direct method: VV: ", VV)
      # println("       checking direct method: size(VV): ", size(VV))
      check_directmethod = transpose(dJdu)*VV
      filename = string("check_directmethod_irev-", i, ".dat")
      writedlm(filename, check_directmethod)

    end

    # advection adjoint check
    if neg_time == true
      println("  calculating dRdA for adjoint check. irev = ", i, ", ifwd = ", i_fwd,", t = ", t)
      dRdA_CS = calcdRdA_CS(mesh, sbp, eqn_fwd, opts, i, t)
      dRdA_FD = calcdRdA_FD(mesh, sbp, eqn_fwd, opts, i, t)
      # dRdA = dRdA_CS
      dRdA = dRdA_FD
      filename = string("dRdA_CS_irev-",i,".dat")
      writedlm(filename, dRdA_CS)
      filename = string("dRdA_FD_irev-",i,".dat")
      writedlm(filename, dRdA_FD)
      println("  writing dRdA file: ", filename)
      # check_adjointmethod = transpose(adj_nextstep.q_vec)*dRdA
      # NOTE 20170712: think I need to be doing adj, not adj_nextstep. this CN rev's i corresponds to eqn_fwd's i_fwd and t
      # check_adjointmethod = transpose(adj.q_vec)*dRdA
      check_adjointmethod = transpose(adj.q_vec)*(-1.0*dRdA)
      filename = string("check_adjointmethod_irev-", i, ".dat")
      writedlm(filename, check_adjointmethod)
    end

    if neg_time == false
      ctx_residual = (physics_func, eqn, h, newton_data)
    else
      # i is the time step index in the reverse sweep.
      #   It moves forward from 2 to (t_steps+1) even though the adjoint is going backwards in time.
      #   The adjoint calculation requires data from the forward sweep to be loaded from disk.
      #   At a given time step in the adjoint solve, although the index i corresponds to this
      #     loop's time step, the index i does not correspond to the same i of the forward sweep.
      #   The adjustment is not just (t_steps - i) because the loop starts at 2 and ends at t_steps + 1.
      i_fwd = (t_steps + 1) - (i - 2)
      # TODO: eventual fix, issue #92, t_steps + 2. needs to be done in RK4 also
      println("\n   time step variables-  i: ", i, "  i_fwd: ", i_fwd, "  t_steps: ", t_steps)
      ctx_residual = (physics_func, adj, h, newton_data, i_fwd, dJdu)

      #-------------------------------------------------
      # Output times & indices
      println(" --- CN times and indices ---")
      println("    CN: i = ", i)
      println("    CN: i_fwd = ", i_fwd)
      println("    CN: t = ", round(t,3))
      println("    CN: t_nextstep = ", round(t_nextstep,3))
      #-------------------------------------------------

    end

    if neg_time == false
      # allow for user to select CN's internal Newton's method. Only supports dense FD Jacs, so only for debugging
      if opts["cleansheet_CN_newton"]
        # cnNewton: in cnNewton.jl
        cnNewton(mesh, sbp, opts, h, physics_func, eqn, eqn_nextstep, t)
      else
        @time newtonInner(newton_data, mesh, sbp, eqn_nextstep, opts, cnRhs_uadj, cnJac_uadj, jac, rhs_vec, ctx_residual, t)
        # filename = string("jac_ifwd-", i,".dat")
        # writedlm(filename, real(jac))

        if opts["uadj_global"]
          ### dRdu check: fwd
          println(" GLOBAL: forming global dRdu, i = $i")
          # first, let's try forming the forward global dRdu without the state perturbed by WWW.
          blksz = mesh.numDof
          # blksz = 3 # testing 44
          row_ix_start = (i-2)*blksz + 1
          row_ix_end = (i-2)*blksz + blksz
          col_ix_start = (i-2)*blksz + 1
          col_ix_end = (i-2)*blksz + blksz
          println(" GLOBAL: diagonal block ix's: [",row_ix_start,":",row_ix_end,", ",col_ix_start,":", col_ix_end,"]")

          II = eye(blksz)

          # TODO: check eqn or eqn_nextstep
          newton_data_discard, jac_for_dRdu_global, rhs_vec_discard = setupNewton(mesh, mesh, sbp, eqn_nextstep, opts)
          assert(opts["jac_method"] == 2)
          epsilon = opts["epsilon"]
          pert = complex(0, epsilon)
          calcJacobianComplex(newton_data_discard, mesh, sbp, eqn_nextstep, opts, physics_func, pert, jac_for_dRdu_global, t)
          # TODO check that this jacobian equals the one inside newtonInner
          println(" GLOBAL: norm of jac after newtonInner call: ", norm(jac_for_dRdu_global))

          # blk = II - 0.5*h*jac
          # blk = II - 0.5*h*jac[1:3,1:3]
          # blk = II - 0.5*h*jac_for_dRdu_global[1:3, 1:3]
          blk = II - 0.5*h*jac_for_dRdu_global
          # blk = ones(blksz, blksz)*44     # testing 44
          println(" GLOBAL: size of blk: ", size(blk))

          # this time step's actual portion of the global dRdu is the cnJac_uadj. so it has to be I-0.5*h*physicsJac
          jac_filename = string("jac_fwd_cJc_t-",t,".dat")
          writedlm(jac_filename, round(real(blk), 4))

          println(" GLOBAL: size(dRdu_global_fwd): ", size(dRdu_global_fwd))
          dRdu_global_fwd[row_ix_start:row_ix_end, col_ix_start:col_ix_end] = blk

          if i != t_steps + 1     # final time step has only one block
            row_ix_start = (i-2)*blksz + blksz + 1
            row_ix_end = (i-2)*blksz + 2*blksz
            col_ix_start = (i-2)*blksz + 1
            col_ix_end = (i-2)*blksz + blksz
            println(" GLOBAL: off-diagonal block ix's: [",row_ix_start,":",row_ix_end,", ",col_ix_start,":", col_ix_end,"]")

            # blk = -1.0*II - 0.5*h*jac
            # blk = -1.0*II - 0.5*h*jac[1:3,1:3]
            # blk = -1.0*II - 0.5*h*jac_for_dRdu_global[1:3, 1:3]
            blk = -1.0*II - 0.5*h*jac_for_dRdu_global
            # blk = ones(blksz, blksz)*55     # testing 44

            dRdu_global_fwd[row_ix_start:row_ix_end, col_ix_start:col_ix_end] = blk
          end   # end of "if != t_steps + 1"
        end # end of "if opts["uadj_global"]

      end   # end of else clause of "if opts["cleansheet_CN_newton"]"

      # NOTE: this has been commented out to remove Newton from the adjoint testing. Currently direct solving.
    # else      # call newtonInner using cnAdjJac and cnAdjRhs
      # @time newtonInner(newton_data, mesh, sbp, adj_nextstep, opts, cnAdjRhs, cnAdjJac, jac, rhs_vec, ctx_residual, t)

    else    # else clause of "if neg_time == false"

      # direct solve for psi_i
      # note: this needs to be called with t, not t_nextstep.
      #   within cnAdjDirect, t_nextstep is calculated and used throughout.
      # (adj_nextstep.q_vec, jac) = cnAdjDirect(mesh, sbp, opts, adj, physics_func, jac, i_fwd, h, t)
      (adj_nextstep.q_vec, jac) = cnAdjDirect(mesh, sbp, opts, adj, physics_func, jac, i_fwd, h, t_steps, t, dRdu_global_rev)
      array1DTo3D(mesh, sbp, adj_nextstep, opts, adj_nextstep.q_vec, adj_nextstep.q)

    end

    # This allows the solution to be updated from _nextstep without a deepcopy.
    #   There are two memory locations used by eqn & eqn_nextstep,
    #   and this flips the location of eqn & eqn_nextstep every time step
    if neg_time == false
      eqn_temp = eqn
      eqn = eqn_nextstep
      eqn_nextstep = eqn_temp
    else
      adj_temp = adj
      adj = adj_nextstep
      adj_nextstep = adj_temp
    end

    # adj.q_vec now contains the adjoint at time step i.
    # Previously, adj_nextstep corresponded to time step i, and adj corresponded to time step i+1.

    #------- adjoint check
    # TODO this is wrong, what is this. seems like a first attempt at check_adjointmethod
    if neg_time == true
      # eqn_fwd = cnAdjLoadChkpt(mesh, sbp, opts, adj, physics_func, i_fwd, t)     # t has been updated, so no t_nextstep
      # jac = cnAdjCalcdRdu(mesh, sbp, opts, eqn_fwd, physics_func, t)
      omega = 1.0
      v_bc = sin(omega*t)
      dRdA = -1.0*jac*v_bc
      dJdA = transpose(adj.q_vec)*dRdA
      filename = string("dJdA_check-", i, ".dat")
      writedlm(filename, transpose(dJdA))
      # if i == 10
        # println("v_bc: ", v_bc)
        # println("jac: ", jac)
        # println("dRdA: ", dRdA)
        # println("adj.q_vec: ", adj.q_vec)
      # end
    end

    # Note: we now need to copy the updated q over for the initial newton guess
    if neg_time == false
      for dof_ix = 1:mesh.numDof
        eqn_nextstep.q_vec[dof_ix] = eqn.q_vec[dof_ix]
      end
      array1DTo3D(mesh, sbp, eqn_nextstep, opts, eqn_nextstep.q_vec,  eqn_nextstep.q)
    else
      for dof_ix = 1:mesh.numDof
        adj_nextstep.q_vec[dof_ix] = adj.q_vec[dof_ix]
      end
      array1DTo3D(mesh, sbp, adj_nextstep, opts, adj_nextstep.q_vec, adj_nextstep.q)
    end

    t = t_nextstep        # update time step

  end   # end of t step loop

  i = i+1

  # Checkpoint final time step
  if neg_time == false
    # for adjoint_straight option: stores every time step's q to disk
    # Note: cannot store full eqn object without extending one of the julia write methods
    if store_u_to_disk == true
      println(" >>>>> Writing solution to disk. i = ", i,", t = ", t)
      filename = string("qvec_for_adj-", i, ".dat")
      writedlm(filename, eqn.q_vec)
      vis_filename = string("solution_storedtodisk_i-", i)
      saveSolutionToMesh(mesh, real(eqn.q_vec))
      writeVisFiles(mesh, vis_filename)

      time_filename = string("t-",i,".dat")
      writedlm(time_filename, t)              # make sure that this time value, whether t or t_nextstep,
                                              #   correctly corresponds to the q_vec stored just above
    end
  else
    # save every time step's adjoint to disk. This is for the final time step
    if opts["adjoint_saveall"]
      filename = string("adj_irev-", i, ".dat")
      writedlm(filename, adj.q_vec)

      vis_filename = string("adj_i-", i)
      saveSolutionToMesh(mesh, real(adj.q_vec))
      writeVisFiles(mesh, vis_filename)
    end
  end

  if neg_time == false
    println(" -------------- eqn.q_vec of last time step. after CN loop. t = $t, i = $i --------------")
    # print_qvec_coords(mesh, sbp, eqn, opts)
  else
    println(" -------------- adj.q_vec of last time step. after CN loop. t = $t, i = $i --------------")
    # print_qvec_coords(mesh, sbp, adj, opts)
  end
  println("\n========================= End of CN =========================")
  println(" ")

  #=
  # new: writing checkpoint data at start of time step
  if neg_time == false
    # for adjoint_straight option: stores every time step's q to disk
    # Note: cannot store full eqn object without extending one of the julia write methods
    if store_u_to_disk == true
      filename = string("qvec_for_adj-", i, ".dat")
      writedlm(filename, eqn.q_vec)
      vis_filename = string("solution_storedtodisk_i-", i)
      saveSolutionToMesh(mesh, real(eqn.q_vec))
      writeVisFiles(mesh, vis_filename)

      time_filename = string("t-",i,".dat")
      writedlm(time_filename, t)              # make sure that this time value, whether t or t_nextstep,
                                              #   correctly corresponds to the q_vec stored just above
    end
  else
    # save every time step's adjoint to disk
    if opts["adjoint_saveall"]
      filename = string("adj-", i, ".dat")
      writedlm(filename, adj.q_vec)
    end
  end
  =#

  if opts["uadj_global"]
    ### GLOBAL: saving global dRdu
    writedlm("global_dRdu_fwd.dat", dRdu_global_fwd)
    writedlm("global_dRdu_rev.dat", dRdu_global_rev)
  end

  # depending on how many timesteps we do, this may or may not be necessary
  #   usage: copyForMultistage!(dest, src)
  # this copyForMultistage! is defined in ODLCommonTools
  if neg_time == false
    copyForMultistage!(eqn, eqn_nextstep)      # copying eqn_nextstep to eqn
    writedlm("solution_final_inCN.dat", real(eqn.q_vec))
  else
    copyForMultistage!(adj, adj_nextstep)      # copying adj_nextstep to eqn
    writedlm("adjoint_final_inCN.dat", real(adj.q_vec))
  end

  #=
  # For debugging: store integer set to mesh and save, just to visualize ordering of q_vec
  for i = 1:length(eqn.q_vec)
    eqn.q_vec[i] = convert(Float64, i)
  end
  filename = string("qvec_integers.dat")
  writedlm(filename, eqn.q_vec)
  vis_filename = string("qvec_integers")
  saveSolutionToMesh(mesh, real(eqn.q_vec))
  writeVisFiles(mesh, vis_filename)
  print_qvec_coords(mesh, sbp, eqn, opts)
  # End debugging integer ordering output
  =#

  cleanupNewton(newton_data, mesh, mesh, sbp, eqn, opts)

  # @debug1 println("============= end of CN: t = $t ===============")

  return t

end   # end of crank_nicolson function

function negativeZeroCheck(t)
  # prevent floating point errors from setting t to a negative number at t = 0.0
  if abs(t) < 1e-14 && t < 0.0
    println("Barely negative time detected (-1e-14 < t < 0.0) , setting t = 0.0.")
    t = 0.0
  end
  return t
end

function check_q_qvec_consistency(mesh, sbp, eqn, opts)

  if eqn.q != reshape(eqn.q_vec, mesh.numDofPerNode, mesh.numNodesPerElement, mesh.numEl)
    error("q & q_vec consistency error")
  end

  return nothing

end


text:
### A Pluto.jl notebook ###
# v0.12.11

using Markdown
using InteractiveUtils

# ╔═╡ ba04b94e-ff3e-11ea-1c4f-df1747ceff76
using Pkg, DrWatson

# ╔═╡ ba05026e-ff3e-11ea-14f2-0f30e9d422b2
begin
	#@quickactivate "StatisticalRethinkingStan"
	using StanSample
	using StatisticalRethinking
end

# ╔═╡ bbb50ac4-ff3d-11ea-1ff8-c3a12b8a9a30
md"## Clip-06-24s.jl"

# ╔═╡ ba093834-ff3e-11ea-1139-4f6986da5b02
begin
	df = sim_happiness()
	df.A = (df.age .- 18) / (65 - 18)
	Text(precis(df; io=String))
end

# ╔═╡ ba1554d4-ff3e-11ea-00a5-01b221494ff3
stan6_10 = "
data {
  int <lower=1> N;
  vector[N] happiness;
  vector[N] A;
}
parameters {
  real <lower=0> sigma;
  real a;
  real bA;
}
model {
  vector[N] mu;
  sigma ~ exponential(1);
  a ~ normal(0, 1);
  bA ~ normal(0, 2);
  mu = a + bA * A;
  happiness ~ normal(mu, sigma);
}
";

# ╔═╡ ba160c44-ff3e-11ea-22f9-2d892aba888f
begin
	m6_10s = SampleModel("m6.10s", stan6_10)
	m6_10_data = Dict(:N => nrow(df), :happiness => df.happiness, :A => df.A,)
	rc6_10s = stan_sample(m6_10s, data=m6_10_data)
	success(rc6_10s) && (part6_10s = read_samples(m6_10s, :particles))
end

# ╔═╡ ba209362-ff3e-11ea-168c-1d6dd60dae6c
if success(rc6_10s)
  post6_10s_df = read_samples(m6_10s, :dataframe)
  Text(precis(post6_10s_df; io=String))
end

# ╔═╡ ba2125ac-ff3e-11ea-257c-3974d8044dd8
begin
	fig = plot(xlab="age", ylab="happiness", leg=false, title="unmarried (grey), married (blue)")
	for i in 1:nrow(df)
		if df[i, :married] == 1
			scatter!([df[i, :age]], [df[i, :happiness]], color=:darkblue)
		else
			scatter!([df[i, :age]], [df[i, :happiness]], color=:lightgrey)
		end
	end
	fig
end

# ╔═╡ ba2a9f2e-ff3e-11ea-0c3c-1b46b3ab931f
md"## End of clip-06-24s.jl"

# ╔═╡ Cell order:
# ╟─bbb50ac4-ff3d-11ea-1ff8-c3a12b8a9a30
# ╠═ba04b94e-ff3e-11ea-1c4f-df1747ceff76
# ╠═ba05026e-ff3e-11ea-14f2-0f30e9d422b2
# ╠═ba093834-ff3e-11ea-1139-4f6986da5b02
# ╠═ba1554d4-ff3e-11ea-00a5-01b221494ff3
# ╠═ba160c44-ff3e-11ea-22f9-2d892aba888f
# ╠═ba209362-ff3e-11ea-168c-1d6dd60dae6c
# ╠═ba2125ac-ff3e-11ea-257c-3974d8044dd8
# ╟─ba2a9f2e-ff3e-11ea-0c3c-1b46b3ab931f


text:
using DataStructures
using DelimitedFiles

"""
This function loads the FUND input parameters stored in datadir into the syntax
needed for the @defmcs macro. This script writes to a text file which can then
be copied into a @defmcs macro.
"""
function load_mcs_RVs(; txt_out = joinpath(@__DIR__, "mcs_RVs.txt"), datadir = joinpath(@__DIR__, "../../data_for_load_mcs_RVs"), string_labels = false)

    files = readdir(datadir)
    filter!(i->i!="desktop.ini", files)
    parameters = OrderedDict{Any, Any}(splitext(file)[1] => readdlm(joinpath(datadir,file), ',') for file in files)

    open(txt_out, "w") do f
        for (k, v) in parameters
            val = loadparameter(v, string_labels)
            if val != nothing
                write(f, "$k = $val\n")
            end
        end
    end

end

# Helper function for concatenating strings in loadparameter
_concat(x,y)=string(x,", ",y)

"""
Takes the contents of one parameter file and turns it into the string needed for a definition within a @defmcs macro.
Input 'p' may be a scalar, a 1-D array (defined as two columns in long format), or a 2-D array (defined as three columns in long format).
If p contains no distributional values, then this returns nothing, because it will not be a random variable for mcs.
"""
function loadparameter(p, string_labels::Bool)
    column_count = size(p,2)
    if column_count == 1
        return format_distribution(p[1,1])
    elseif column_count == 2
        vals = []
        for i in 1:size(p, 1)
            dist = format_distribution(p[i, 2])
            if dist !== nothing
                string_labels ? push!(vals, "\"$(p[i,1])\" => $dist") : push!(vals, "$(p[i,1]) => $dist")
            end
        end
        return isempty(vals) ? nothing : "[$(reduce(_concat, vals))]"
    elseif column_count == 3
        vals = []
        for i in 1:size(p, 1)
            dist = format_distribution(p[i, 3])
            if dist !== nothing
                string_labels ? push!(vals, "[\"$(p[i,1])\", \"$(p[i,2])\"] => $dist") : push!(vals, "[$(p[i,1]), $(p[i,2])] => $dist")
            end
        end
        return isempty(vals) ? nothing : "[$(reduce(_concat, vals))]"
    else
        error("Unable to parse data files with $column_count columns.")
    end
end

"""
Converts one distribution as defined in FUND's datafiles into a string of a julia-parsable distribution definition.
Example: "~N(0.089;0.1484;min=0)" becomes "Truncated(Normal(0.089,0.1484), 0.0, Inf)"
If v is a scalar or boolean, it returns nothing.
If v starts with '~', but is not ~N(), ~Gamma(), or ~Triangular(), it throws an error.
"""
function format_distribution(v)
    if isa(v,AbstractString)
        if startswith(v,"~") & endswith(v,")")
            args_start_index = findnext('(', v, 1)
            dist_name = v[2:args_start_index-1]
            args = split(v[args_start_index+1:end-1], ';')
            fixedargs = filter(i->!contains(i,"="),args)
            optargs = Dict(split(i,'=')[1]=>split(i,'=')[2] for i in filter(i->contains(i,"="),args))

            if dist_name == "N"
                if length(fixedargs)!=2 error() end
                if length(optargs)>2 error() end

                basenormal = "Normal($(parse(Float64, fixedargs[1])),$(parse(Float64, fixedargs[2])))"

                if length(optargs)==0
                    return basenormal
                else
                    min = haskey(optargs,"min") ? parse(Float64, optargs["min"]) : -Inf
                    max = haskey(optargs,"max") ? parse(Float64, optargs["max"]) : Inf
                    return "Truncated($basenormal, $min, $max)"
                end
            elseif startswith(v, "~Gamma(")
                if length(fixedargs)!=2 error() end
                if length(optargs)>2 error() end

                basegamma = "Gamma($(parse(Float64, fixedargs[1])),$(parse(Float64, fixedargs[2])))"

                if length(optargs)==0
                    return basegamma
                else
                    min = haskey(optargs,"min") ? parse(Float64, optargs["min"]) : -Inf
                    max = haskey(optargs,"max") ? parse(Float64, optargs["max"]) : Inf
                    return "Truncated($basegamma, $min, $max)"
                end
            elseif startswith(v, "~Triangular(")
                triang = "TriangularDist($(parse(Float64, fixedargs[1])), $(parse(Float64, fixedargs[2])), $(parse(Float64, fixedargs[3])))"
                return triang
            else
                error("Unknown distribution")
            end
        end
    end
    return nothing
end


text:
@testset "NetwAnalysis tests" begin show_testset_section() #Scope for test data

using MDDatasets
using CMDimCircuits.CircuitAnalysis
using CMDimCircuits.NetwAnalysis
import CMDimCircuits.NetwAnalysis: NetworkParameters, NetworkParametersNoRef, NetworkParametersRef
include("netwanalysis_cmptools.jl")

#==Input data
===============================================================================#
f = collect(1:3)*1e9
c = capacitance(5e-15)

#==Intermediate Computations
===============================================================================#
ycap = admittance(c, f=1e9)
zcap = impedance(c, f=1e9)
ycap_vec = admittance(c, f=f)

#==Helper functions
===============================================================================#

#==Tests
===============================================================================#
@testset "Test constructors" begin show_testset_description()
	S = Network(:S, [1 2; 3 4])
		@test isa(S, NetworkParametersRef{:S,2,Int64})
		@test S.m == [1 2; 3 4]
		@test S.z0 == 50 #Default
	T = Network(:ABCD, [1 2; 3 4])
		@test isa(T, NetworkParametersNoRef{:ABCD,2,Int64})
		@test T.m == [1 2; 3 4]
end

@testset "Test traits" begin show_testset_description()
	S = Network(:S, [1 2; 3 4])
	nptype = NPType(S)
		@test isa(nptype, NPType{:S})
		@test Symbol(nptype) == :S
end

printheader("Test impedance operations:")
	@show shunt(:ABCD, zcap)
	@show shunt(:ABCD, ycap)
	@show series(:ABCD, zcap)
	@show series(:ABCD, ycap)

@testset "Basic matrix operations" begin show_testset_description()
	S = Network(:S, [1 2; 3 4])
	Sx2 = S * 2
	@test maxerror(Sx2, Network(:S, 2 .* S.m)) < SPARAM_THRESH
	#@show S.*[3 4; 2 1] #TODO: Julia v0.6: Figure out a way to do element-by-element operations??
	mm = S*[3 4; 2 1]
	@test maxerror(mm, Network(:S, S.m * [3 4; 2 1])) < SPARAM_THRESH
	s21 = S[2,1]
	@test s21 == S.m[2,1]
end

@testset "Complex matrix{vector} operations" begin show_testset_description()
	TC = shunt(:ABCD, ycap_vec) #Matrix of vectors
	TCpull = vector_pull(TC) #Vector of matrices
	TCpush = vector_push(TCpull) #Back to matrix of vectors
	@test maxerror(TC, TCpush) < ABCD_THRESH
end

@testset "Subset operations" begin show_testset_description()
	S = Network(:S, [1 2; 3 4])
	Sswapped = submatrix(S, [2, 1]) #Port-swapped
	@test maxerror(Sswapped, Network(:S, [4 3; 2 1])) < SPARAM_THRESH
	Γ = submatrix(S, [1])
		Γ_mx2D = convert(Array{Int64,2}, [S.m[1,1]]')
	x = Network(:S, Γ_mx2D, z0=S.z0)
	@test maxerror(Γ, Network(:S, Γ_mx2D, z0=S.z0)) < SPARAM_THRESH
end

@testset "Network parameter conversions" begin show_testset_description()
	T = Network(:ABCD, S)
	Z = Network(:Z, T)
	T2 = Network(:ABCD, T) #Passthrough
	@test maxerror(T, T2) < ABCD_THRESH
	@show S2 = Network(:S, Z, z0=75)
	@show T = Network(:T, S2)
end

@testset "Specific S <=> T conversions" begin show_testset_description()
	S = Network(:S, [0 1; 1 0])
	T = Network(:T, S)
	@test maxerror(Network(:T, [1.0 0.0; 0.0 1.0], z0=50), T) < TPARAM_THRESH
end

@testset "Matrix impedance transformations" begin show_testset_description()
	_S50 = Network(:S, [1 2; 3 4])
	_S75 = Network(:S, _S50, z0=75)
	_T50 = Network(:T, _S50)
	_T75 = Network(:T, _S75)

	@test maxerror(Network(:T, _S75), _T75) < TPARAM_THRESH
	@test maxerror(Network(:T, _S75, z0=50), _T50) < TPARAM_THRESH
	@test maxerror(Network(:T, _S50, z0=75), _T75) < TPARAM_THRESH
	@test maxerror(Network(:T, _T75, z0=50), _T50) < TPARAM_THRESH

	@test maxerror(Network(:S, _T50), _S50) < SPARAM_THRESH
	@test maxerror(Network(:S, _T75), _S75) < SPARAM_THRESH
	@test maxerror(Network(:S, _T75, z0=50), _S50) < SPARAM_THRESH
	#TODO: add tests more to validate impedance transformation (use realistic impedances)
	#Z->S50->S75; Z->S75
end

printheader("Test stability & gains:")
	@show S = Network(:S, [1 2; 3 4])
	@show kstab(S)
	@show gain(:U, S)
	@show gain(:MA, S)
	@show gain(:MS, S)

end


text:
function ω_weights(z_ex, θ, ξ)
    z_0 = z_ex[1]
    z_bar = z_ex[end]
    Δ_ex = diff(z_ex)
    Δ₋ = [0; Δ_ex] # (21)
    Δ₊ = [Δ_ex; 0] # (22)
    z = z_ex[2:end-1] # i.e., interior of z_ex
    P = length(z)

    ω_bar = 1/2 * (Δ₋ + Δ₊) # (23)
    Ξ₁ = 1/(1 - ξ * Δ₊[1]) # (24)
    Ξₚ = 1/(1 + ξ * Δ₋[end]) # (25)

    ω = similar(z)
    denom = 1 - exp(-θ * z_bar) # (26) denominator (i.e., F(z_bar))

    ω = ω_bar[2:end-1] * θ .* exp.((ξ-θ) * z)/denom # (26)
    ω[1] += ω_bar[1] * Ξ₁ * θ * exp((ξ-θ) * z_0)/denom # (26), left endpoint
    ω[end] += ω_bar[end] * Ξₚ * θ * exp((ξ-θ) * z_bar)/denom # (26), right endpoint
    return ω
end


text:
# ## Catwalk.jl 8-min intro
# Krisztián Schäffer, JuliaCon 2021

# We will performance-optimize a Rock-Paper-Scissors simulation.
#
# (Example was inspired by https://giordano.github.io/blog/2017-11-03-rock-paper-scissors/)

# We model player moves with types:

abstract type Hand end
struct Rock <: Hand end
struct Paper <: Hand end
struct Scissors <: Hand end

play(::T, ::T) where T = 3       # 3: Tie
play(::Paper, ::Rock) = 1        # 1: First player wins
play(::Scissors, ::Paper) = 1
play(::Rock, ::Scissors) = 1
play(a, b) = 3 - play(b, a) # Reverse order. 3-1 == 2: second player wins

play(Paper(), Scissors()) # Scissors cuts paper, second player wins

# ---
# A single play by two random players:

function playrand(hands)
    hand1 = rand(hands)
    hand2 = rand(hands)
    return play(hand1, hand2)
end

# ---
# Type instability causes dynamic (run-time) dispatch:

const hands = [Rock(), Paper(), Scissors()]

using InteractiveUtils
@code_warntype playrand(hands)

# ---
# Play repeatedly and count results:

function playmatch(hands, num_plays, result)
    for i=1:num_plays
        winner = playrand(hands)::Int
        result[winner] += 1
    end
end

# ---
# all-play-all tournament of num_players:

function tournament(hands, num_players)
    results = [[0,0,0] for i1=1:num_players, i2=1:num_players]
    for i1 = 1:num_players
        for i2 = i1+1:num_players
            playmatch(hands, 10_000, results[i1, i2])
        end
    end
    return results
end

tournament(hands, 5)

# ---
# Performance:

@time tournament(hands, 100)
nothing

# Speed it up!

# ##### Option 1: Redesign
# ...

# ##### Option 2: Union spliting
# ...

# ##### Option 3: Dispatch manually:

function playrand_manualdispatch(hands)
    hand1 = rand(hands)
    hand2 = rand(hands)
    if hand1 isa Rock
        return play(hand1, hand2) # Fast, (partially) type-stable route
    elseif hand1 isa Paper
        return play(hand1, hand2)
    elseif hand1 isa Scissors
        return play(hand1, hand2)
    end
    return play(hand1, hand2) # Fallback to fully dynamic dispatch
end

# (ManualDispatch.jl may help, but not in this 2-args case)

# ---
# ##### Option 4: JIT with Catwalk:

using Catwalk

@jit play hand1 function playrand_jit(hands, jitctx)
    hand1 = rand(hands)
    hand2 = rand(hands)
    return play(hand1, hand2)
end

function playmatch_jit(hands, num_plays, result, jitctx)
    for i=1:num_plays
        winner = playrand_jit(hands, jitctx)::Int
        result[winner] += 1
    end
end

function tournament_jit(hands, num_players, jit=Catwalk.JIT())
    results = [[0,0,0] for i1=1:num_players, i2=1:num_players]
    for i1 = 1:num_players
        Catwalk.step!(jit)
        jitctx = Catwalk.ctx(jit)
        for i2 = i1+1:num_players
            playmatch_jit(hands, 10_000, results[i1, i2], jitctx)
        end
    end
    return results
end

# ---
const jit = Catwalk.JIT()
@time tournament_jit(hands, 100, jit)
nothing

# Hot start:

@time tournament_jit(hands, 100, jit)
nothing

# ---
# Catwalk.jl highlights:
#
#   - Low cost statistical profiler
#   - Tunable cost model
#   - Smart recompilation: Only if type-distribution changed significantly
#   - No world age issues, works inside a function
#   - 360 lines of readable vanilla Julia (no dependencies/ccall/llvmcall)
#
# Sample source: https://github.com/tisztamo/Catwalk.jl/tree/main/docs/introtalk
#
# Help, Tuning, inner workings: https://tisztamo.github.io/Catwalk.jl/dev/
#

typeof(Catwalk.ctx(jit))

# hands = [
#    [repeat([Rock()], i1)..., Paper(), Paper(), Paper(), (i1 > 20 ? Scissors() : Paper())],
#    [Rock(), Paper(), Scissors()],
# ]