# Find _a_ unitary transform U that maps U†X(g)U = Y(g) where X(g) and Y(g) are irreps that
# are related by a constant unitary transform, following https://doi.org/10.1088/1751-8113/47/50/505203
function find_unitary_irrep_transform(Xs, Ys, gs)
    N = size(first(Xs), 1)
    G = length(gs)

    all(((X,Y),) -> tr(X) ≈ tr(Y), zip(Xs,Ys)) || error("Xs and Ys are not unitarily equivalent")

    rᵃᵇ²s = zeros(ComplexF64, N, N)
    for (X, g) in zip(Xs, gs)
        g⁻¹ = inv(g)
        idx_g⁻¹ = something(findfirst(≈(g⁻¹), gs))
        Y = Ys[idx_g⁻¹]
        for a in 1:N
            for b in 1:N
                rᵃᵇ²s[a,b] += X[a,a]*Y[b,b]
            end
        end
    end

    ab = findfirst(rᵃᵇ² -> real(rᵃᵇ²) > 0.0 && abs(imag(rᵃᵇ²)) < ATOL_DEFAULT, rᵃᵇ²s)
    a, b = Tuple(something(ab))
    rᵃᵇ = sqrt(N/G * real(rᵃᵇ²s[a,b]))
    U = zeros(ComplexF64, N, N)
    for (X,Y) in zip(Xs, Ys)
        for i in 1:N
            for j in 1:N
                U[i,j] += X'[i,a]*Y[b,j]
            end
        end
    end
    U .*= N/(G*rᵃᵇ)

    return mapto_canonical_unitary!(U)
end

# find the unitary transformation that connects Xᵢ and Yᵢ by UXᵢU⁻¹ = Yᵢ for some set {i},
# following https://mathoverflow.net/a/391741
function find_unitary_transform(Xs, Ys)
    N = LinearAlgebra.checksquare(first(Xs))
    M = length(Xs)

    M == length(Ys) || error("the lengths of Xs and Ys must match")
    all(X->LinearAlgebra.checksquare(X)==N, Xs) || error("`Xs` matrices are not square or not of identical size")
    all(Y->LinearAlgebra.checksquare(Y)==N, Ys) || error("`Ys` matrices are not square or not of identical size")

    N² = N^2
    T = promote_type(eltype(first(Xs)), eltype(first(Ys)))
    Q = Matrix{T}(undef, M*N², N²)
    for i in 1:M
        Q[(i-1)*N² .+ (1:N²), 1:N²] .= kron(I(N), transpose(Xs[i]))  .-  kron(Ys[i], I(N))
    end

    # find a matrix U which is not necessarily unitary, but which obeys UXU⁻¹ = Y ∀i
    Uv = nullspace(Q, atol=ATOL_DEFAULT)
    U = if size(Uv,2) ≠ 1 # multiple solution vectors found; need some combination of them...
        # let's do a hail-mary attempt and just add them up; works sometimes... FIXME!?
        permutedims(reshape(sum(Uv, dims=2), N, N))
    else
        permutedims(reshape(Uv, N, N))
    end

    check_transform = all(zip(Xs, Ys)) do (X,Y)
        U*X/U ≈ Y
    end
    check_transform || error("initial transform is invalid :(")

    U = mapto_canonical_unitary(U)
    check_transform = all(zip(Xs, Ys)) do (X,Y)
        U*X*U' ≈ Y
    end
    check_transform || error("'unitarized' transform is invalid :(")
    U'*U ≈ I || error("computed transformation is not unitary :(")

    return U
end

# Eq. (4.5.51) of Bradley & Cracknell's book. Find a unitary operator U such that
#           Y(g) = UX(g)U⁻¹   ∀g∈G
# Construct P as:
#           U = |G|⁻¹ ∑ Y(g)MX⁻¹(g)
# with M an arbitrary (?) matrix. This approach is derived by Wigner in his book on p. 80
# and the "unitarization" is discussed on p. 78-79. It's not completely clear what
# conditions should be imposed on M; for now, we choose it as some arbitrary unitary matrix,
# generated from the Gell-Mann matrices. The initial construction doesn't guarantee that
# Not clear what conditions should be imposed on X, if any -
# if it is chosen completely arbitrary, it does not ensure that P is unitary (but it can
# then be rescaled to become unitary; i.e., it is a "scalar away" from being unitary).
function find_unitary_transform_bc(Xs, Ys, M=nothing)
    if Xs ≈ Ys # fast-path
        return Matrix{ComplexF64}(I(size(first(Xs), 1)))
    end

    N = size(first(Xs), 1)
    hs = gellmann(N, skip_identity=true)
    Ms = isnothing(M) ? [sum(hs), cis.(hs)...] : [M]
    for M in Ms
        U = sum(zip(Xs, Ys)) do (X,Y)
            X*M*inv(Y)
        end ./ length(Xs)
        rank(U) == N || continue

        check_transform = all(zip(Xs, Ys)) do (X,Y)
            U*X/U ≈ Y
        end
        check_transform || continue # initial transform guess is invalid :(

        U = mapto_canonical_unitary(U)
        check_transform = all(zip(Xs, Ys)) do (X,Y)
            U*X*U' ≈ Y
        end
        check_transform || continue # 'unitarized' transform is invalid :(
        U'*U ≈ I || continue # computed transformation is not unitary :(

        return U
    end
    error("failed to find a unitary transform")
end

# find unitary part of 𝒯g₀ operation, using Eq. (S58) of Bradlyn, Science (2016)
#   D(𝒯g₀) D(gᵢ) D(𝒯g₀ )⁻¹  =  D*(g₀gᵢg₀⁻¹)
function find_antiunitary_irrep(lgir, g₀)
    Xs = lgir(nothing)                               # D(gᵢ)
    Ys = trs_transformed_irreps(Xs, group(lgir), g₀) # D*(g₀gᵢg₀⁻¹)

    if all(((X,Y),) -> tr(X) ≈ tr(Y), zip(Xs, Ys))
        D_𝒯g₀ = try
            find_unitary_transform_bc(Xs, Ys)
        catch
            find_unitary_transform(Xs, Ys)
        end
        return D_𝒯g₀
    else
        # TODO: The irrep is complex; this should probably not be possible if fed by a
        #       corep, but would be nice to generalize our scope here
        error("Xs and Ys are not equivalent: don't know how to proceed")
    end
end

function find_antiunitary_irrep(lgir::LGIrrep)
    g₀ = find_kv_to_negative_kv_operation(lgir) # find g₀ s.t. g₀k = -k + G
    if !isnothing(g₀)
        return find_antiunitary_irrep(lgir, g₀)
    else
        return nothing
    end
end

# find an operation in the space group of `lgir` that maps `kv` to `-kv`
function kv_to_negative_kv_operation(lgir::LGIrrep{D}) where D
    sg = spacegroup(num(lgir), Val(D))
    cntr = centering(sg)
    kv = position(lgir)
    g₀_idx = findfirst(g -> isapprox(g*kv, -kv, cntr), sg)
    return isnothing(g₀_idx) ? nothing : sg[g₀_idx]
end

function trs_transformed_irreps(Xs, lg, g₀)
    Ys = similar(Xs)  # D*(g₀gᵢg₀⁻¹)
    cntr = centering(lg)
    kv = position(lg)() # TODO: αβγ?
    g₀⁻¹ = inv(g₀)
    for (i, gᵢ) in enumerate(lg)
        g₀gᵢg₀⁻¹ = compose(g₀, compose(gᵢ, g₀⁻¹, #=modτ=#false), #=modτ=#false)
        j, Δw = something(Crystalline.findequiv(g₀gᵢg₀⁻¹, lg, cntr))
        ϕ_g₀gᵢg₀⁻¹ = cispi(2*dot(kv, Δw)) # translation phase factor
        D_g₀gᵢg₀⁻¹ = Xs[j] * ϕ_g₀gᵢg₀⁻¹   # D(g₀gᵢg₀⁻¹)
        Ys[i] = conj.(D_g₀gᵢg₀⁻¹)         # D*(g₀gᵢg₀⁻¹)
    end
    return Ys
end

function mapto_canonical_unitary(U)
    # we assume that `U` is at most "a scalar away" from being unitary, and we want to
    # mutate `U` to become this related unitary matrix; additionally, we want to map it to
    # a canonical form. We have two steps:

    # 1. "unitarize" `U` (a unitary operator must have UU† = I; let's ensure this - assuming
    # that `U` is indeed a "scalar away" from being unitary). Trick is from Wigner p. 78-79:
    K = sqrt(U*U') # a multiple of the identity matrix
    U = K*U

    # 2. "canonicalize" `U` (currently has an arbitrary complex phase of norm 1) by
    # requiring that it is a special unitary matrix (i.e., in SU(N)), so that `det(U) = 1`.
    # Already, `U` is in U(N), i.e. `norm(det(U)) = 1`; below, we find the phase needed to
    # "rotate" it into SU(N), using that det(cA) = cᴺdet(A)
    N = LinearAlgebra.checksquare(U)
    c = det(U)^(1/N)
    U ./= c

    return U
end


using AsyPlots
using Test

# 2D
Pen()
Arrow()
Arrow3()
Arrow(5)
Arrow(name="MidArrow",size=3)
Point(1,2,label="hey",pen=Pen(linewidth=3))
Path([0 0; 1 1];arrow=Arrow())
Circle((0,0),1)
Polygon([im,1,0])

Point(1,2,3,label="hey",pen=Pen(linewidth=3))
Path([0 0 0; 1 1 1])
Polygon([0 0 0; 1 1 1; 0 0 1])

z = randn(5,5)
S = Surface(z)

P = Point(1,2,color="DarkGreen")
Q = Path([0 0; 0 1; 1 1; 0 2],spline=true,label="Label(\"A\",Relative(0.8))")
C = Circle((0,0),1.0,linewidth=3,opacity=0.2,color="Purple")

X = collect(range(0,stop=1,length=100))
Y1 = cumsum(randn(100))
Y2 = cumsum(randn(100))
plot(X,Y1)
plot(X,[Y1,Y2])
plot(Y1)
plot(Y1,linewidth=3,opacity=0.4)
plot([x->x^n for n=1:4],0,1)
plot(x->sin(x),0,2π)
plot(t->cos(t)^3,(0,π/2))
plot([t->sin(t),t->cos(t)^3],(0,π/2))

heatmap(randn(20,20))


lines = readlines("inputs/04.txt") # notes: added a newline to the input file
nums = [parse(Int64, x ) for x in split(lines[1], ',')]
lines = lines[2:end]

function parse_grids(lines)
  grids = []
  for i in 2:6:length(lines)-4
    grid = zeros(Int64, 5,5)
    gridstrs = lines[i:i+4]
    for (r_i, r) in enumerate(gridstrs)
      grid[r_i,:] = [parse(Int64, x) for x in filter(q -> length(q) > 0 ,split(strip(r), ' '))]
    end
    push!(grids, grid)
  end
  return grids
end

is_bingo(grid) = 500 in sum(grid, dims = 1) || 500 in sum(grid, dims = 2)

score(n, grid) = n * (sum(grid) - (100 * count(i -> i == 100, grid)))

function day3(nums, grids_)
  grids = copy(grids_)
  for num in nums
    grids = map(g -> replace(g, num => 100), grids)
    to_del = []
    for (g_i, g) in enumerate(grids)
      if is_bingo(g)
        if length(grids) == 1
          return score(num, grids[1])
        else
          push!(to_del, g_i)
        end
      end
    end
    deleteat!(grids, to_del)
  end
end

grids = parse_grids(lines)
println(day3(nums, grids))


# ------------------------------------------------------------------
# Licensed under the ISC License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    pointgradients(img, points; [method])

Compute the gradients along all dimensions of N-dimensional `img`
at `points` using `method`. Default method is `:ando3`.
"""
function pointgradients(img::AbstractArray{T,N},
                        points::AbstractVector;
                        method=:ando3) where {N,T}
  dims = size(img)
  npts = length(points)

  # smoothing weights
  weights = (method == :sobel ? [1,2,1] :
             method == :ando3 ? [.112737,.274526,.112737] :
             error("Unknown gradient method: $method"))

  # pad input image
  padimg = padarray(img, Pad(:replicate, 1, 1))

  # gradient matrix
  G = zeros(npts, N)

  # compute gradient for all directions at specified points
  for i in 1:N
    # kernel = centered difference + perpendicular smoothing
    if dims[i] > 1
      # centered difference
      kern = reshape([-1,0,1], ntuple(k -> k == i ? 3 : 1, N))
      # perpendicular smoothing
      for j in setdiff(1:N, i)
        if dims[j] > 1
          wvec = reshape(weights, ntuple(k -> k == j ? 3 : 1, N))
          kern = broadcast(*, kern, wvec)
        end
      end

      A = zeros(size(kern))
      for (k, icenter) in enumerate(points)
        i1 = CartesianIndex(ntuple(i->1, N))
        for ii in CartesianIndices(A)
          A[ii] = padimg[ii + icenter - i1]
        end

        G[k,i] = sum(kern .* A)
      end
    end
  end

  G
end


# Benchmark on pairwise distance evaluation

using Distances

macro bench_pairwise_dist(repeat, dist, x, y)
    quote
        println("bench ", typeof($dist))

        # warming up
        r1 = evaluate($dist, ($x)[:,1], ($y)[:,1])
        pairwise($dist, $x, $y)

        # timing

        t0 = @elapsed for k = 1 : $repeat
            nx = size($x, 2)
            ny = size($y, 2)
            r = Array(typeof(r1), (nx, ny))
            for j = 1 : ny
                for i = 1 : nx
                    r[i, j] = evaluate($dist, ($x)[:,i], ($y)[:,j])
                end
            end
        end
        @printf "    loop:      t = %9.6fs\n" (t0 / $repeat)

        t1 = @elapsed for k = 1 : $repeat
            r = pairwise($dist, $x, $y)
        end
        @printf "    pairwise:  t = %9.6fs  |  gain = %7.4fx\n" (t1 / $repeat) (t0 / t1)
        println()
    end
end

m = 100
nx = 200
ny = 250

x = rand(m, nx)
y = rand(m, ny)

w = rand(m)
Q = rand(m, m)
Q = Q' * Q

@bench_pairwise_dist 20 SqEuclidean() x y
@bench_pairwise_dist 20 Euclidean() x y
@bench_pairwise_dist 20 Cityblock() x y
@bench_pairwise_dist 20 Chebyshev() x y
@bench_pairwise_dist  5 Minkowski(3.0) x y
@bench_pairwise_dist 20 Hamming() x y

@bench_pairwise_dist 20 CosineDist() x y
@bench_pairwise_dist 10 CorrDist() x y
@bench_pairwise_dist 20 ChiSqDist() x y
@bench_pairwise_dist 10 KLDivergence() x y
@bench_pairwise_dist  5 JSDivergence() x y

@bench_pairwise_dist 10 BhattacharyyaDist() x y
@bench_pairwise_dist 10 HellingerDist() x y

@bench_pairwise_dist 20 WeightedSqEuclidean(w) x y
@bench_pairwise_dist 20 WeightedEuclidean(w) x y
@bench_pairwise_dist 20 WeightedCityblock(w) x y
@bench_pairwise_dist  5 WeightedMinkowski(w, 3.0) x y
@bench_pairwise_dist 20 WeightedHamming(w) x y

@bench_pairwise_dist 10 SqMahalanobis(Q) x y
@bench_pairwise_dist 10 Mahalanobis(Q) x y


# Convolute two vectors of probabilities

# Input:
# p1:       Vector of probabilities P(X = 0), ..., P(X = n)
# p2:       Vector of probabilities P(Y = 0), ..., P(Y = n)

# Output:
# out:      Vector of probabilities P(X + Y = 0), ..., P(X + Y = n)

function convolution(p1::Vector{Float64},
                     p2::Vector{Float64})::Vector{Float64}
    #
    M = length(p1)
    out = zeros(M)
    @inbounds for i = 0:M-1
        for j = 0:M-1-i
            out[i+j+1] += p1[i+1]*p2[j+1]
        end
    end
    out
end

function convolution(p1::Vector{Float64},
                     p2::Vector{Float64},
                     x::Int64)::Vector{Float64}
    #
    out = zeros(x+1)
    @inbounds for i = 0:x
        for j = 0:x-i
            out[i+j+1] += p1[i+1]*p2[j+1]
        end
    end
    out
end

function convolutionRev(p1::Vector{Float64},
                     p2::Vector{Float64},
                     x::Int64)::Vector{Float64}
    #
    out = zeros(x+1)
    @inbounds for i = 0:x
        for j = 0:x-i
            out[x + 1 - i - j] += p1[i+1]*p2[j+1]
        end
    end
    out
end

function convolutionRev(p1::Vector{Float64},
                     p2::Matrix{Float64},
                     x::Int64)::Matrix{Float64}
    #
    ncol = size(p2)[2]
    out = zeros(x+1, ncol)
    @inbounds for i = 0:x
        for j = 0:x-i
            for col = 1:ncol
                out[x + 1 - i - j, col] += p1[i+1]*p2[j+1, col]
            end
        end
    end
    out
end

function convolutionRevScale(p1::Vector{Float64},
                             p2::Matrix{Float64},
                             x::Int64,
                             scale::Vector{Float64})::Matrix{Float64}
    ncol = size(p2)[2]
    out = zeros(ncol, x+1)
    @inbounds for i = 0:x
        for j = 0:x-i
            scalefactor = scale[x + 1 - i - j]
            p1Part = p1[i+1]
            for col = 1:ncol
                out[col, x + 1 - i - j] += p1Part*p2[j+1, col]*scalefactor
            end
        end
    end
    out
end


# ### Surface (3)

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

f(x,y) = x^2 + y^2
x = -10:10
y = x
plot(x, y, f, linetype=:wireframe)


function scale_range(vmin, vmax, n=1, threshold=100)
    dv = abs(vmax - vmin)  # > 0 as nonsingular is called before.
    meanv = (vmax + vmin) / 2
    offset = if abs(meanv) / dv < threshold
        0.0
    else
        copysign(10 ^ (log10(abs(meanv)) ÷ 1), meanv)
    end
    scale = 10 ^ (log10(dv / n) ÷ 1)
    scale, offset
end

function _staircase(steps)
    n = length(steps)
    result = Vector{Float64}(undef, 2n)
    for i in 1:(n-1)
        @inbounds result[i] = 0.1 * steps[i]
    end
    for i in 1:n
        @inbounds result[i+(n-1)] = steps[i]
    end
    result[end] = 10 * steps[2]
    return result
    # [0.1 .* steps[1:end-1]; steps; 10 .* steps[2]]
end

struct EdgeInteger
    step::Float64
    offset::Float64

    function EdgeInteger(step, offset)
        if step <= 0
            error("Step must be positive")
        end
        new(step, abs(offset))
    end
end

function closeto(e::EdgeInteger, ms, edge)
    tol = if e.offset > 0
        digits = log10(e.offset / e.step)
        tol = max(1e-10, 10 ^ (digits - 12))
        min(0.4999, tol)
    else
        1e-10
    end
    abs(ms - edge) < tol
end

function le(e::EdgeInteger, x)
    # 'Return the largest n: n*step <= x.'
    d, m = divrem(x, e.step)
    if closeto(e, m / e.step, 1)
        d + 1
    else
        d
    end
end

function ge(e::EdgeInteger, x)
    # 'Return the smallest n: n*step >= x.'
    d, m = divrem(x, e.step)
    if closeto(e, m / e.step, 0)
        d
    else
        d + 1
    end
end

"""
A cheaper function that tries to come up with usable tick locations for a given value range
"""
function locateticks(vmin, vmax, n_ideal::Int, _integer::Bool = false, _min_n_ticks::Int = 2)

    _steps = (1.0, 1.5, 2.0, 2.5, 3.0, 4.0, 5.0, 6.0, 8.0, 10.0)
    _extended_steps = _staircase(_steps)

    scale, offset = scale_range(vmin, vmax, n_ideal)

    _vmin = vmin - offset
    _vmax = vmax - offset

    raw_step = (_vmax - _vmin) / n_ideal

    steps = _extended_steps .* scale

    if _integer
        # For steps > 1, keep only integer values.
        filter!(steps) do i
            (i < 1) || (abs(i - round(i)) < 0.001)
        end
    end

    #istep = np.nonzero(steps >= raw_step)[0][0]
    istep = findfirst(1:length(steps)) do i
        @inbounds return steps[i] >= raw_step
    end
    ticks = 1.0:0.1:0.0
    for istep in istep:-1:1
        step = steps[istep]

        if _integer && (floor(_vmax) - ceil(_vmin) >= _min_n_ticks - 1)
            step = max(1, step)
        end
        best_vmin = (_vmin ÷ step) * step

        # Find tick locations spanning the vmin-vmax range, taking into
        # account degradation of precision when there is a large offset.
        # The edge ticks beyond vmin and/or vmax are needed for the
        # "round_numbers" autolimit mode.
        edge = EdgeInteger(step, offset)
        low = le(edge, _vmin - best_vmin)
        high = ge(edge, _vmax - best_vmin)
        ticks = (low:high) .* step .+ best_vmin
        # Count only the ticks that will be displayed.
        # nticks = sum((ticks .<= _vmax) .& (ticks .>= _vmin))

        # manual sum because broadcasting was slow
        nticks = 0
        for t in ticks
            if _vmin <= t <= _vmax
                nticks += 1
            end
        end

        if nticks >= _min_n_ticks
            break
        end
    end

    ticks = ticks .+ offset #(first(ticks) + offset):step(ticks):(last(ticks) + offset)
    vals = filter(x -> vmin <= x <= vmax, ticks)

    # for some reason, the values coming out of this method sometimes are ever so slightly off from
    # their intended decimal representation, causing printed values like
    # 0.08000000001 instead of 0.08
    # so here we round off the numbers to the required number of digits after the decimal point
    exponent = floor(Int, minimum(log10.(abs.(diff(vals)))))
    round.(vals, digits = max(0, -exponent+1))
end

function locateticks(vmin, vmax, width_px, ideal_tick_distance::Float32, _integer::Bool = false, _min_n_ticks::Int = 2)
    # how many ticks would ideally fit?
    n_ideal = round(Int, width_px / ideal_tick_distance) + 1
    locateticks(vmin, vmax, n_ideal, _integer, _min_n_ticks)
end


@testset "Karp Minimum Cycle Mean" begin
     w = [0. 0.  0. 0. 0. 0. 1. 0.  0.  0.0  0.  0.
          0. 0.  0. 0. 0. 0. 1. 0.  0.  0.0  0.  0.
          0. 0.  0. 0. 0. 0. 0. 0.  1.1 0.0  0.  0.
          0. 1.2 0. 0. 0. 0. 1. 0.  0.  0.0  0.  0.
          0. 1.  0. 0. 0. 0. 0. 0.  1.3 0.0  0.  0.
          0. 1.2 0. 0. 0. 0. 1. 0.  0.  0.0  0.  0.
          0. 1.  0. 0. 0. 0. 0. 0.  1.3 0.0  0.  0.
          0. 0.  0. 0. 0. 0. 0. 0.  0.  0.0  0.  0.
          0. 0.  0. 0. 0. 0. 0. 0.  0.  0.0  0.4 0.
          0. 0.  0. 0. 0. 0. 1. 0.  0.  0.0  0.  0.
          0. 0.  0. 0. 0. 0. 1. 0.  0.  0.0  0.  0.
          0. 0.  0. 0. 0. 0. 0. 0.  1.2 0.0  0.  0.]

    w2 = [Inf  Inf    Inf  Inf  Inf  Inf   -0.09 Inf  Inf    Inf  Inf    Inf
          Inf  Inf    Inf  Inf  Inf  Inf   -0.06 Inf  Inf    Inf  Inf    Inf
          Inf  Inf    Inf  Inf  Inf  Inf  Inf    Inf   -0.01 Inf  Inf    Inf
          Inf   -0.19 Inf  Inf  Inf  Inf   -0.09 Inf  Inf    Inf  Inf    Inf
          Inf   -0.02 Inf  Inf  Inf  Inf  Inf    Inf   -0.03 Inf  Inf    Inf
          Inf   -0.19 Inf  Inf  Inf  Inf   -0.09 Inf  Inf    Inf  Inf    Inf
          Inf   -0.02 Inf  Inf  Inf  Inf  Inf    Inf   -0.03 Inf  Inf    Inf
          Inf  Inf    Inf  Inf  Inf  Inf  Inf    Inf  Inf    Inf  Inf    Inf
          Inf  Inf    Inf  Inf  Inf  Inf  Inf    Inf  Inf    Inf    0.39 Inf
          Inf  Inf    Inf  Inf  Inf  Inf   -0.09 Inf  Inf    Inf  Inf    Inf
          Inf  Inf    Inf  Inf  Inf  Inf   -0.06 Inf  Inf    Inf  Inf    Inf
          Inf  Inf    Inf  Inf  Inf  Inf  Inf    Inf   -0.01 Inf  Inf    Inf]

     g1 = SimpleDiGraph(12)
     add_edge!(g1, 1, 7)
     add_edge!(g1, 2, 7)
     add_edge!(g1, 3, 9)
     add_edge!(g1, 4, 2)
     add_edge!(g1, 4, 7)
     add_edge!(g1, 5, 2)
     add_edge!(g1, 5, 9)
     add_edge!(g1, 6, 2)
     add_edge!(g1, 6, 7)
     add_edge!(g1, 7, 2)
     add_edge!(g1, 7, 9)
     add_edge!(g1, 9, 11)
     add_edge!(g1, 10, 7)
     add_edge!(g1, 11, 7)
     add_edge!(g1, 12, 9)

     @testset "simple digraphs" for g in testgraphs(g1)
        c, λ = karp_minimum_cycle_mean(g, w)
        @test c == [9, 11, 7]
        @test λ == 0.9

        c2, λ2 = karp_minimum_cycle_mean(g, w2)
        @test λ2 == -0.04
        @test c2 == [2, 7]
     end

    # Tricky test case
    # Backward walk from 3 is good:
    # 1 - 2 - 3 - 4 - 3
    # Only the first of the two backward walk from 4 is good:
    # 1 - 3 - 4 - 3 - 4
    # 1 - 2 - 2 - 3 - 4
    # We will do the backward walk from 3 in the test since 0 = F_4(3) < F_4(4) = 1
    tricky = SimpleDiGraph(4)
    add_edge!(tricky, 1, 3)
    add_edge!(tricky, 1, 2)
    add_edge!(tricky, 3, 4)
    add_edge!(tricky, 4, 1) # Make it strongly connected
    add_edge!(tricky, 4, 3)
    add_edge!(tricky, 2, 3)
    add_edge!(tricky, 2, 2)
    distmx = [Inf  0.  1. Inf
              Inf  1.  0. Inf
              Inf Inf Inf  0.
               1. Inf  0. Inf]

    @testset "tricky case" for g in testgraphs(tricky)
        c, λ = karp_minimum_cycle_mean(g, distmx)
        @test λ == 0.
        @test sort(c) == [3, 4]

        c, λ = karp_minimum_cycle_mean(g, distmx .- 2)
        @test λ == -2.
        @test sort(c) == [3, 4]
    end

    # Multiple strongly connected components and one component with one node and no edge
    multi = SimpleDiGraph(3)
    add_edge!(multi, 1, 1)
    add_edge!(multi, 2, 2)
    distmx = [0. Inf
              Inf -1]

    @testset "multiple SCCs" for g in testgraphs(multi)
        c, λ = karp_minimum_cycle_mean(g, distmx)
        @test λ == -1.
        @test c == [2]
    end
end


# Burgers
import SDELab2
include("common.jl")
########################## problem parameters
T=1
diffusion=0.1
sig=sqrt(2*diffusion)

quad0=gauss_quadrature.delta_measure(0.)
#quad0=gauss_quadrature.Gauss_Hermite([0.,1e-8],4)
#########################################
function drift_ff(u)
  x=u[1]; mf=u[2]
  if x>mf
    return 0.
  elseif x==mf
    return 0.5
  else
    return 1.
  end
end
#
############################
function drift_fr(u,ell)
  x=u[1]; mf=u[2]
  return 0.5*erfc((x-mf)/ell)
end
#########################################
function diff_gg(u,sig)
  return sig
end
#########################################
function mf_r(x)
  return [x]
end
#########################################
fcn=SDELab2.set_fcn_mf(u->drift_fr(u,ell),
          u->diff_gg(u,sig),
          x->mf_r(x))
#########################################
function exact_cdf2(x,t,sig)
  sig_sq=sig^2
  scale=sig*sqrt(2*t)
 return 1- ( erfc(-x/scale)./
 (
 erfc(-x/scale)
 + exp((t-2*x)/(2*sig_sq)) .* (2-erfc((t-x)/scale))
 )    )
end
#
function exact_cdf3(x,t,sig)
  sig_sq=sig^2
  scale=sig*sqrt(2*t)
  T1=erfc(-x/scale)
  T2=exp((t-2*x)/(2*sig_sq))
  T3=erfc((t-x)/scale)
 return T2*(2-T3)./(T1+T2.*(2-T3))
end
######################################### method parameters
solver="MFEM0e"
dt=1e-2
params=SDELab2.set_opt(dt,solver)
params["const"]=1e0       # for no_Gauss_points
params["length_scale"]=1 # lam=1/length^2
params["cutoff"]=1  #
params["skiphalf"]=true
params["periodic"]=false
params[ "no_partition"]=20
#########################################
###################### computing reference value
print("Computing a reference value...")
phi=x->[x^2;x]
tic()
params["extrap"]=false
dtpts=0.1./(2.0).^(1:8 )
params["MaxStepSize"]=dtpts[end]
ell=sqrt(dtpts[end])

#ell=0.5#(dtpts[end])
ell=0.1
#ell=0.01
#ell=0.001
quad0=gauss_quadrature.delta_measure(0.)#-dtpts[end]/2)
fcn=SDELab2.set_fcn_mf(u->drift_fr(u,ell),
          u->diff_gg(u,sig),
          x->mf_r(x))
#quad0=gauss_quadrature.delta_measure(sqrt(dtpts[end]))
exact,cpu_time,stats,q1=get_soln(fcn,quad0,params,phi,true)
println(stats["no_pts"])
println("exact=", exact)
#########################################
#initial_final_cdf_pdf(q0,q1,[100,40],[0.08,0.01])
#

cum1=cumsum(q1.w)
x1=(q1.X[2:end]+q1.X[1:end-1])/2
#x1=q1.X
#x1=q1.X[1:end-1]
cum_exact=1-exact_cdf2(q1.X,1.,sig)
q2=gauss_quadrature.QUAD(length(x1),x1,q1.w)
println("l1 ",gauss_quadrature.l1cdf(q2,z->(1-exact_cdf2(z,1.,sig))))
using PyPlot
figure(2)
clf()
plot(q1.X,cum_exact,"b")
plot(x1,cum1[2:end],"g")#
axis([-1.5,1.5,0,1])
xlabel("X")
ylabel("cdf")
gcf()[:set_size_inches](width_inches,height_inches)
tight_layout()
  savefig(utf8("cdf_cmp.pdf"))
  clf()
  plot(q1.X,abs(cum_exact-cum1),"b")
  #plot(q1.X,cum1[2:end],"g")#
    axis([-1.5,1.5,0,0.04])
  xlabel("X")
  ylabel("error")
  gcf()[:set_size_inches](width_inches,height_inches)
  tight_layout()
  savefig(utf8("cdf_cmp_del.pdf"))
#########################################
dtpts=dtpts[1:end-1] # omit last value, used for reference
errors=zeros(length(dtpts),2); errors_no=deepcopy(errors); errors_nd=deepcopy(errors)
cpu_times=zeros(length(dtpts)); cpu_times_no=deepcopy(cpu_times);cpu_times_nd=deepcopy(cpu_times)
params["MaxStepSize"]=dtpts[1]
######################
params["Solver"]="MFEM0e"
dtpts_no=dtpts
for i=1:(length(dtpts))
    params["MaxStepSize"]=dtpts_no[i]
    qoi,cpu_time,stats,q=get_soln(fcn,quad0,params,phi,false)
    errors_no[i,:],cpu_times_no[i]=helper(qoi, exact, cpu_time, "normal")
end
########################
params["Solver"]="MFEM0e"
for i=1:(length(dtpts))
    params["MaxStepSize"]=dtpts[i]
    qoi,cpu_time,stats,q=get_soln(fcn,quad0,params,phi,true)
    errors[i,:],cpu_times[i]=helper(qoi, exact, cpu_time, "normal")
end
################
#########################################
figure(1);
# first functional
i=1

time_step_err(dtpts,errors[:,i],
  dtpts_no,errors_no[:,i],
  dtpts,errors_nd[:,i],[2,2],utf8("mom1.pdf"))
err_cpu(errors[:,i],cpu_times[:],
  errors_no[:,i],cpu_times_no[:],
  errors_nd[:,i],cpu_times_nd[:],[1,1], utf8("mom1_.pdf"))
# second functional
i=2
time_step_err(dtpts,errors[:,i],
      dtpts,errors_no[:,i],
      dtpts,errors_nd[:,i],[1,1],utf8("mom2.pdf"))
err_cpu(errors[:,i], cpu_times[:],
      errors_no[:,i],cpu_times_no[:],
      errors_nd[:,i],cpu_times_nd[:],[1,1], utf8("mom2_.pdf"))
#########################################
println("finished preparing plots")


using Test, FastChebInterp, StaticArrays

# similar to ≈, but acts elementwise on tuples
≈′(a::Tuple, b::Tuple; kws...) where {N} = length(a) == length(b) && all(xy -> isapprox(xy[1],xy[2]; kws...), zip(a,b))

@testset "1d test" begin
    lb,ub = -0.3, 0.9
    f(x) = exp(x) / (1 + 2x^2)
    f′(x) = f(x) * (1 - 4x/(1 + 2x^2))
    @test_throws ArgumentError chebpoints(-1, lb, ub)
    x = chebpoints(48, lb, ub)
    interp = chebinterp(f.(x), lb, ub)
    @test ndims(interp) == 1
    x1 = 0.2
    @test interp(x1) ≈ f(x1)
    @test chebgradient(interp, x1) ≈′ (f(x1), f′(x1))
end

@testset "2d test" begin
    lb, ub = [-0.3,0.1], [0.9,1.2]
    f(x) = exp(x[1]+2*x[2]) / (1 + 2x[1]^2 + x[2]^2)
    ∇f(x) = f(x) * SVector(1 - 4x[1]/(1 + 2x[1]^2 + x[2]^2), 2 - 2x[2]/(1 + 2x[1]^2 + x[2]^2))
    x = chebpoints((48,39), lb, ub)
    interp = chebinterp(f.(x), lb, ub)
    interp0 = chebinterp(f.(x), lb, ub, tol=0)
    @test ndims(interp) == 2
    x1 = [0.2, 0.3]
    @test interp(x1) ≈ f(x1)
    @test interp(x1) ≈ interp0(x1) rtol=1e-15
    @test all(n -> n[1] < n[2], zip(size(interp.coefs), size(interp0.coefs)))
    @test chebgradient(interp, x1) ≈′ (f(x1), ∇f(x1))

    # univariate function in 2d should automatically drop down to univariate polynomial
    f1(x) = exp(x[1]) / (1 + 2x[1]^2)
    interp1 = chebinterp(f1.(x), lb, ub)
    @test interp1(x1) ≈ f1(x1)
    @test size(interp1.coefs, 2) == 1 # second dimension should have been dropped

    # complex and vector-valued interpolants:
    f2(x) = [f(x), cis(x[1]*x[2] + 2x[2])]
    ∇f2(x) = vcat(transpose(∇f(x)), transpose(SVector(im*x[2], im*(x[1] + 2)) * cis(x[1]*x[2] + 2x[2])))
    interp2 = chebinterp(f2.(x), lb, ub)
    @test interp2(x1) ≈ f2(x1)
    @test chebjacobian(interp2, x1) ≈′ (f2(x1), ∇f2(x1))

    # chebinterp_v1
    av1 = Array{ComplexF64}(undef, 2, size(x)...)
    av1[1,:,:] .= f.(x)
    av1[2,:,:] .= (x -> f2(x)[2]).(x)
    interp2v1 = chebinterp_v1(av1, lb, ub)
    @test interp2v1(x1) ≈ f2(x1)
    @test chebjacobian(interp2v1, x1) ≈′ (f2(x1), ∇f2(x1))
end

@testset "1d regression" begin
    x = range(1,2, length=200)
    y = @. sin.(x + 0.2) * x
    c = chebregression(x, y, 2)
    c2 = x.^(0:2)' \ y # Vandermonde-style fit
    @test c(1.1) ≈ c2[1] + 1.1 * c2[2] + 1.1^2 * c2[3] rtol=1e-13
end

@testset "2d regression" begin
    x = vec(SVector.(range(-3, 2, length=50), range(0,1, length=40)'))
    fr(x::SVector{2}) = sin(x[1]/3 + 1)*exp(x[2])
    c = chebregression(x, fr.(x), (2,3))

    # Vandermonde-style fit:
    powsvec(x) = vec([x[1]^i * x[2]^j for i=0:2, j=0:3])'
    c2 = mapreduce(powsvec, vcat, x) \ fr.(x)

    p = SVector(1.12345, sqrt(0.5))
    @test c(p) ≈ powsvec(p)*c2 rtol=1e-13

    cM = chebregression(mapreduce(transpose∘Vector, vcat, x), fr.(x), (2,3))
    @test cM.coefs == c.coefs

    fr4(x::SVector{2}) = cos(x[1]/3 + 1)*sin(x[2])
    c4 = chebregression(x, fr4.(x), (2,3))
    F14 = [SVector(fr(x),fr4(x)) for x in x]
    c14 = chebregression(x, F14, (2,3))
    @test first.(c14.coefs) ≈ c.coefs rtol=1e-13
    @test last.(c14.coefs) ≈ c4.coefs rtol=1e-13
    c14M = chebregression(x, mapreduce(transpose∘Vector, vcat, F14), (2,3))
    @test c14M.coefs == c14.coefs
end

@testset "degree 1" begin
    lb,ub = -0.3, 0.9
    x = chebpoints(1, lb, ub)
    interp = chebinterp((x -> 2x+3).(x), lb, ub)
    @test interp(0) ≈ 3
    @test interp(0.1) ≈ 3.2
    @test chebgradient(interp, 0)[2] ≈ 2 ≈ chebgradient(interp, 0.1)[2]

    lb, ub = [-0.3,-0.1], [0.9,1.2]
    x = chebpoints((1,1), lb, ub)
    interp = chebinterp((x -> (2x[1]+3)*(4x[2]+5)).(x), lb, ub)
    @test interp([0,0]) ≈ 3*5
    @test interp([0.1,0]) ≈ 3.2*5
    @test interp([0,0.1]) ≈ 3*5.4
    @test chebgradient(interp, [0,0])[2] ≈ [2*5, 3*4]
end

@testset "degree 0" begin
    lb,ub = -0.3, 0.9
    x = chebpoints(0, lb, ub)
    interp = chebinterp(sin.(x), lb, ub)
    @test interp(0) ≈ interp(0.1) ≈ sin((lb+ub)/2)
    @test chebgradient(interp, 0)[2] == 0 == chebgradient(interp, 0.1)[2]

    lb, ub = [-0.3,-0.1], [0.9,1.2]
    x = chebpoints((0,0), lb, ub)
    interp = chebinterp((x -> sin(x[1])*cos(x[2])).(x), lb, ub)
    @test interp([0,0]) ≈ interp([0.1,0.2]) ≈ sin(0.3)*cos(0.55)
    @test chebgradient(interp, [0.1,0.2])[2] == [0, 0]

    lb, ub = [-0.3,-0.1], [0.9,1.2]
    x = chebpoints((1,0), lb, ub)
    interp = chebinterp((x -> (2x[1]+3)*cos(x[2])).(x), lb, ub)
    @test interp([0,0]) ≈ 3*cos(0.55)
    @test interp([0.1,0.22]) ≈ 3.2*cos(0.55)
    @test chebgradient(interp, [0.1,0.2])[2] == [2*cos(0.55), 0]
end


using Base.Test
using SparseRF

n,m = 10^3, 5 ;
features = randn(n,m);
weights = rand(-2:2,m);
labels = features * weights;

maxdepth = 3
model = build_tree(labels, features, 5, 0, maxdepth)
@test depth(model) == maxdepth

println("\n##### nfoldCV Regression Tree #####")
r2 = nfoldCV_tree(labels, features, 3)
@test mean(r2) > 0.6

println("\n##### nfoldCV Regression Forest #####")
r2 = nfoldCV_forest(labels, features, 2, 10, 3)
@test mean(r2) > 0.8


######### Jags batch program example  ###########

using StatsPlots, Jags, Statistics

cd(ProjDir) do

  dyes = "
  model {
     for (i in 1:BATCHES) {
        for (j in 1:SAMPLES) {
           y[i,j] ~ dnorm(mu[i], tau.within);
        }
        mu[i] ~ dnorm(theta, tau.between);
     }

     theta ~ dnorm(0.0, 1.0E-10);
     tau.within ~ dgamma(0.001, 0.001);
     s2.within <- 1/tau.within;
     tau.between ~ dgamma(0.001, 0.001);
     s2.between <- 1/tau.between;
     s2.total <- s2.within + s2.between;
     f.within <- s2.within/s2.total;
     f.between <- s2.between/s2.total;
  }
  "

  data = Dict{String, Any}()
  data["y"] = reshape([
    1545, 1540, 1595, 1445, 1595,
    1520, 1440, 1555, 1550, 1440,
    1630, 1455, 1440, 1490, 1605,
    1595, 1515, 1450, 1520, 1560,
    1510, 1465, 1635, 1480, 1580,
    1495, 1560, 1545, 1625, 1445
  ], 6, 5)
  data["BATCHES"] = 6
  data["SAMPLES"] = 5

  inits = [
    Dict{String, Any}(
      "theta" => 1500,
      "tau.within" => 1,
      "tau.between" => 1
  	)
  ]

  monitors = Dict{String, Any}(
    "theta" => true,
    "s2.within" => true,
    "s2.between" => true
  	)

  jagsmodel = Jagsmodel(name="dyes", model=dyes,
    #ncommands=3, nchains=3,
    monitor=monitors)

  println("\nJagsmodel that will be used:")
  jagsmodel |> display
  println("Input observed data dictionary:")
  data |> display
  println("\nInput initial values dictionary:")
  inits |> display
  println()

  @time sim = jags(jagsmodel, data, inits, ProjDir)
  println()

  ## Plotting
  p = plot(sim)

end #cd


function check_point(M::Hyperbolic{N}, p::PoincareHalfSpacePoint; kwargs...) where {N}
    mpv = check_point(Euclidean(N), p.value; kwargs...)
    mpv === nothing || return mpv
    if !(last(p.value) > 0)
        return DomainError(
            norm(p.value),
            "The point $(p) does not lie on $(M) since its last entry is nonpositive.",
        )
    end
end

@doc raw"""
    convert(::Type{PoincareHalfSpacePoint}, p::PoincareBallPoint)

convert a point [`PoincareBallPoint`](@ref) `p` (from $ℝ^n$) from the
Poincaré ball model of the [`Hyperbolic`](@ref) manifold $\mathcal H^n$ to a [`PoincareHalfSpacePoint`](@ref) $π(p) ∈ ℝ^n$.
Denote by $\tilde p = (p_1,\ldots,p_{n-1})$. Then the isometry is defined by

````math
π(p) = \frac{1}{\lVert \tilde p \rVert^2 - (p_n-1)^2}
\begin{pmatrix}2p_1\\⋮\\2p_{n-1}\\1-\lVert p\rVert^2\end{pmatrix}.
````
"""
function convert(::Type{PoincareHalfSpacePoint}, p::PoincareBallPoint)
    return PoincareHalfSpacePoint(
        1 / (norm(p.value[1:(end - 1)])^2 + (last(p.value) - 1)^2) .*
        vcat(2 .* p.value[1:(end - 1)], 1 - norm(p.value)^2),
    )
end

@doc raw"""
    convert(::Type{PoincareHalfSpacePoint}, p::Hyperboloid)
    convert(::Type{PoincareHalfSpacePoint}, p)

convert a [`HyperboloidPoint`](@ref) or `Vector``p` (from $ℝ^{n+1}$) from the
Hyperboloid model of the [`Hyperbolic`](@ref) manifold $\mathcal H^n$ to a [`PoincareHalfSpacePoint`](@ref) $π(x) ∈ ℝ^{n}$.

This is done in two steps, namely transforming it to a Poincare ball point and from there further on to a PoincareHalfSpacePoint point.
"""
convert(::Type{PoincareHalfSpacePoint}, ::Any)
function convert(t::Type{PoincareHalfSpacePoint}, p::HyperboloidPoint)
    return convert(t, convert(PoincareBallPoint, p))
end
function convert(t::Type{PoincareHalfSpacePoint}, p::T) where {T<:AbstractVector}
    return convert(t, convert(PoincareBallPoint, p))
end

@doc raw"""
    convert(::Type{PoincareHalfSpaceTVector}, p::PoincareBallPoint, X::PoincareBallTVector)

convert a [`PoincareBallTVector`](@ref) `X` at `p` to a [`PoincareHalfSpacePoint`](@ref)
on the [`Hyperbolic`](@ref) manifold $\mathcal H^n$ by computing the push forward $π_*(p)[X]$ of
the isometry $π$ that maps from the Poincaré ball to the Poincaré half space,
cf. [`convert(::Type{PoincareHalfSpacePoint}, ::PoincareBallPoint)`](@ref).

The formula reads

````math
π_*(p)[X] =
\frac{1}{\lVert \tilde p\rVert^2 + (1-p_n)^2}
\begin{pmatrix}
2X_1\\
⋮\\
2X_{n-1}\\
-2⟨X,p⟩
\end{pmatrix}
-
\frac{2}{(\lVert \tilde p\rVert^2 + (1-p_n)^2)^2}
\begin{pmatrix}
2p_1(⟨X,p⟩-X_n)\\
⋮\\
2p_{n-1}(⟨X,p⟩-X_n)\\
(\lVert p \rVert^2-1)(⟨X,p⟩-X_n)
\end{pmatrix}
````
where $\tilde p = \begin{pmatrix}p_1\\⋮\\p_{n-1}\end{pmatrix}$.
"""
function convert(
    ::Type{PoincareHalfSpaceTVector},
    p::PoincareBallPoint,
    X::PoincareBallTVector,
)
    den = norm(p.value[1:(end - 1)])^2 + (last(p.value) - 1)^2
    scp = dot(p.value, X.value)
    c1 =
        (2 / den .* X.value[1:(end - 1)]) .-
        (4 * (scp - last(X.value)) / (den^2)) .* p.value[1:(end - 1)]
    c2 = -2 * scp / den - 2 * (1 - norm(p.value)^2) * (scp - last(X.value)) / (den^2)
    return PoincareHalfSpaceTVector(vcat(c1, c2))
end

@doc raw"""
    convert(::Type{PoincareHalfSpaceTVector}, p::HyperboloidPoint, ::HyperboloidTVector)
    convert(::Type{PoincareHalfSpaceTVector}, p::P, X::T) where {P<:AbstractVector, T<:AbstractVector}

convert a [`HyperboloidTVector`](@ref) `X` at `p` to a [`PoincareHalfSpaceTVector`](@ref)
on the [`Hyperbolic`](@ref) manifold $\mathcal H^n$ by computing the push forward $π_*(p)[X]$ of
the isometry $π$ that maps from the Hyperboloid to the Poincaré half space,
cf. [`convert(::Type{PoincareHalfSpacePoint}, ::HyperboloidPoint)`](@ref).

This is done similarly to the approach there, i.e. by using the Poincaré ball model as
an intermediate step.
"""
convert(::Type{PoincareHalfSpaceTVector}, ::Any)
function convert(
    t::Type{PoincareHalfSpaceTVector},
    p::HyperboloidPoint,
    X::HyperboloidTVector,
)
    return convert(t, convert(AbstractVector, p), convert(AbstractVector, X))
end
function convert(
    ::Type{PoincareHalfSpaceTVector},
    p::P,
    X::T,
) where {P<:AbstractVector,T<:AbstractVector}
    return convert(
        PoincareHalfSpaceTVector,
        convert(Tuple{PoincareBallPoint,PoincareBallTVector}, (p, X))...,
    )
end

@doc raw"""
    convert(
        ::Type{Tuple{PoincareHalfSpacePoint,PoincareHalfSpaceTVector}},
        (p,X)::Tuple{PoincareBallPoint,PoincareBallTVector}
    )

Convert a [`PoincareBallPoint`](@ref) `p` and a [`PoincareBallTVector`](@ref) `X`
to a [`PoincareHalfSpacePoint`](@ref) and a [`PoincareHalfSpaceTVector`](@ref) simultaneously,
see [`convert(::Type{PoincareHalfSpacePoint}, ::PoincareBallPoint)`](@ref) and
[`convert(::Type{PoincareHalfSpaceTVector}, ::PoincareBallPoint,::PoincareBallTVector)`](@ref)
for the formulae.
"""
function convert(
    ::Type{Tuple{PoincareHalfSpacePoint,PoincareHalfSpaceTVector}},
    (p, X)::Tuple{PoincareBallPoint,PoincareBallTVector},
)
    return (convert(PoincareHalfSpacePoint, p), convert(PoincareHalfSpaceTVector, p, X))
end

@doc raw"""
    convert(
        ::Type{Tuple{PoincareHalfSpacePoint,PoincareHalfSpaceTVector}},
        (p,X)::Tuple{HyperboloidPoint,HyperboloidTVector}
    )
    convert(
        ::Type{Tuple{PoincareHalfSpacePoint,PoincareHalfSpaceTVector}},
        (p, X)::Tuple{P,T},
    ) where {P<:AbstractVector, T <: AbstractVector}

Convert a [`HyperboloidPoint`](@ref) `p` and a [`HyperboloidTVector`](@ref) `X`
to a [`PoincareHalfSpacePoint`](@ref) and a [`PoincareHalfSpaceTVector`](@ref) simultaneously,
see [`convert(::Type{PoincareHalfSpacePoint}, ::HyperboloidPoint)`](@ref) and
[`convert(::Type{PoincareHalfSpaceTVector}, ::Tuple{HyperboloidPoint,HyperboloidTVector})`](@ref)
for the formulae.
"""
function convert(
    ::Type{Tuple{PoincareHalfSpacePoint,PoincareHalfSpaceTVector}},
    (p, X)::Tuple{HyperboloidPoint,HyperboloidTVector},
)
    return (convert(PoincareHalfSpacePoint, p), convert(PoincareHalfSpaceTVector, p, X))
end
function convert(
    ::Type{Tuple{PoincareHalfSpacePoint,PoincareHalfSpaceTVector}},
    (p, X)::Tuple{P,T},
) where {P<:AbstractVector,T<:AbstractVector}
    return (convert(PoincareHalfSpacePoint, p), convert(PoincareHalfSpaceTVector, p, X))
end

@doc raw"""
    distance(::Hyperbolic, p::PoincareHalfSpacePoint, q::PoincareHalfSpacePoint)

Compute the distance on the [`Hyperbolic`](@ref) manifold $\mathcal H^n$ represented in the
Poincaré half space model. The formula reads

````math
d_{\mathcal H^n}(p,q) = \operatorname{acosh}\Bigl( 1 + \frac{\lVert p - q \rVert^2}{2 p_n q_n} \Bigr)
````
"""
function distance(::Hyperbolic, p::PoincareHalfSpacePoint, q::PoincareHalfSpacePoint)
    return acosh(1 + norm(p.value .- q.value)^2 / (2 * p.value[end] * q.value[end]))
end

@doc raw"""
    inner(
        ::Hyperbolic{n},
        p::PoincareHalfSpacePoint,
        X::PoincareHalfSpaceTVector,
        Y::PoincareHalfSpaceTVector
    )

Compute the inner product in the Poincaré half space model. The formula reads
````math
g_p(X,Y) = \frac{⟨X,Y⟩}{p_n^2}.
````
"""
function inner(
    ::Hyperbolic,
    p::PoincareHalfSpacePoint,
    X::PoincareHalfSpaceTVector,
    Y::PoincareHalfSpaceTVector,
)
    return dot(X.value, Y.value) / last(p.value)^2
end

@doc raw"""
    project(::Hyperbolic, ::PoincareHalfSpacePoint ::PoincareHalfSpaceTVector)

projction of tangent vectors in the Poincaré half space model is just the identity, since
the tangent space consists of all $ℝ^n$.
"""
project(::Hyperbolic, ::PoincareHalfSpacePoint::PoincareHalfSpaceTVector)

function allocate_result(
    ::Hyperbolic,
    ::typeof(project),
    X::PoincareHalfSpaceTVector,
    ::PoincareHalfSpacePoint,
)
    return PoincareHalfSpaceTVector(allocate(X.value))
end

function project!(
    ::Hyperbolic,
    Y::PoincareHalfSpaceTVector,
    ::PoincareHalfSpacePoint,
    X::PoincareHalfSpaceTVector,
)
    return (Y.value .= X.value)
end


using Test
using Nash
using Distributions
using Polyhedra
using CDDLib
using LinearAlgebra
using Combinatorics
using IterTools
using QuantEcon
using LightGraphs
using SymPy

game = random_2players_game(Normal(0,2),2,3)
@testset "games" begin
    ggame = generate_game([1 0; 0 1], [1 0; 0 1])
    @test typeof(ggame) == Dict{String,Array{Int64,2}}
    @test typeof(ggame["player1"]) == Array{Int64,2}

    p2game = random_2players_game(Normal(0,2),2,3)
    @test typeof(p2game["player1"]) == Array{Float64,2}
    @test typeof(p2game["player2"]) == Array{Float64,2}

    npgame = random_nplayers_game(Binomial(5,0.5),[2,3,3]);
    @test typeof(npgame["player1"]) == Array{Int64,3}
end

@testset "utilities" begin
    s = [[0.5,0.5],[0.75,0.25]]
    out = outer(s)
    @test typeof(out) == Array{Float64,2}
    @test out[1][1] == s[1][1]*s[2][1]

    s = [Sym.([0.5,0.5]),Sym.([0.75,0.25])]
    out = outer(s)
    @test typeof(out) == Array{Float64,2}
    @test out[1][1] == s[1][1]*s[2][1]
end

@testset "equilibira" begin
    s = [[0.5,0.5],[0.75,0.25]]
    game = random_nplayers_game(Binomial(5,0.5),[2,2]);
    payoffs = get_payoff(game,s)
    @test typeof(payoffs) == Dict{String,Float64}
    @test payoffs["player2"] == s[1]'*game["player2"]*s[2]

    bp1 = best_reply(generate_game([1 0; 0 1], [1 0; 0 1]), [[1, 0], [1, 0]], 1)
    bp2 = best_reply(generate_game([1 0; 0 1], [1 0; 0 1]), [[1, 0], [1, 0]], 2)
    bp1p = best_reply(generate_game([1 0; 0 1], [1 0; 0 1]), [[1, 0], [1, 0]], 1, 1/5)
    # bpch = best_reply(generate_game(Matrix(I,3,3), Matrix(I,3,3)), [[1/2,1/2,0],[1/3,1/3,1/3]], 1, return_val = "chull")
    @test bp1 == bp2
    @test bp1 == [1 0]
    @test bp2 isa Array{<:Real,2}
    @test bp1p isa Array{<:Real,2}
    # @test bpch isa CDDLib.Polyhedron{<:Real}

    # nashq = is_nash_q(generate_game([1 0 ; 0 1], [1 0; 0 1]), [[1, 0], [1, 0]])
    # @test typeof(nashq) == Dict{String,Bool}
    # @test nashq["Nash"] == true

    # game_hist = iterate_best_reply(game, s, 0.5, 7)
    # @test length(game_hist) == 8
    # @test game_hist isa Array{Any,1} #TODO in future not Any type
    # @test game_hist[1] isa Array{Array{Float64,1},1}
end

@testset "markov" begin
    game = generate_game(Matrix(I, 3, 3), Matrix(I, 3, 3))
    s = [[1, 0, 0],[0, 1, 0]]
    mc = game2markov(game, s)
    @test mc isa MarkovChain{Int64,Array{Int64,2},Array{Int64,1}}
    @test all(simulate(mc, 4, init = 1) .== 1)
    @test is_irreducible(mc) == false
    @test period(mc) == 2
    @test is_aperiodic(mc) == false
end

@testset "symmetric" begin
    game = random_symmetric_2players_game(Binomial(10,1/2),2)
    @test typeof(game) == Dict{String,AbstractArray{Int64,2}}
    @test game["player1"] == game["player2"]'

    game = generate_game([1 0; 0 1], [1 0; 1 1])
    x = symbols("x", real=true)
    res = find_symmetric_nash_equilibrium_2players_game(game, [x, 1-x])
    @test res == [0.0, 0.5, 1.0]

    actions_no_vector = [2, 2];
    graph = create_symmetries_graph(actions_no_vector)
    @test length(vertices(graph.graph)) == 8
    @test length(edges(graph.graph)) == 4
    @test check_equality_condition(graph,((1,1),1), ((1,1),2)) == true
    @test check_equality_condition(graph,((1,1),1), ((1,2),2)) == false
    @test length(find_all_equalities(graph)) == 4
end


"""
`Linear`: Performs the operation w * x [.+ b]
# Constructors
    `Linear(output::Int, input::Int; kwargs)`

# Fields
    `w`: Weight
    `b`: Bias

# Keywords
      `bias=true` Whether or not to add bias
      `winit=xavier`  Weight initialization `(eltype, dims...)->array`
      `binit=zeros`   Bias initialization `(eltype, dims...)->array`
      `dtype=Float32` Element type of weight and bias

# Forward execution
    `forward(ctx, l::Linear, x)`
    `@mc l(x)`
    `@run l(x)

  If `ndims(x)` > 2, `mat(x)` is called to reshape `x`.
"""
type Linear <: KnetModule
    w::Param
    b::Union{Param, Void}
end

function Linear(output::Int, input::Int;
                bias=true,
                winit=xavier,
                binit=zeros,
                dtype=Float32)
    w = Param(winit(dtype, output, input))
    b = bias ? Param(binit(dtype, output, 1)) : nothing
    return Linear(w, b)
end

#function forward(ctx, l::Linear, x)
function (l::Linear)(ctx, x)
    if ndims(x) > 2;  x = mat(x); end
    o = val(ctx, l.w) * x
    if l.b !== nothing
        o = o .+ val(ctx, l.b)
    end
    return o
end


# Replicating the fit of daily average temperature data at Boston.

using DelimitedFiles, RCall

## loading the data

data = readdlm("../data/boston.csv")
const dat = data[:, 1] #data values
const yday = Int64.(data[:, 2]) #day of year
const year = Int64.(data[:, 3]) #year
const jday = Int64.(data[:, 4]) #Julian day since first observation day
data = nothing

const climatetrend = readdlm("../data/logco2eq.csv")
const climate_trenddict = Dict(floor.(Int, climatetrend[:, 1]) .=> climatetrend[:, 2])

const climate_trend = zeros(length(dat))
for i in 1:length(dat)
    climate_trend[i] = climate_trenddict[year[i]]
end

## setting up the covariate matrices

const nsplines = 8
const basismatrix = rcopy(R"pbs::pbs($jday %% 365.25,df=$nsplines, Boundary.knots = c(0,365.25))") #using R package pbs for periodic splines
const seasonalmatrix = [cos.(2*π*jday./365.25) sin.(2*π*jday./365.25)]

const ϕmatrix = [repeat([1.0], length(jday)) basismatrix climate_trend climate_trend.*seasonalmatrix]
const τmatrix = [repeat([1.0], length(jday)) basismatrix]

## example initial guess

const XtX = τmatrix' * τmatrix
const regcoef = XtX \ τmatrix' * dat

const ϕ₀init = [copy(regcoef); zeros(3)]
const ϕ₁init = [copy(regcoef); zeros(3)]
const κ₀init = 0.0
const κ₁init = 0.0
const νinit = 1.0
const τ₀init = zeros(nsplines+1)
const τ₁init = zeros(nsplines+1)

const initguess = [κ₀init; τ₀init; ϕ₀init; κ₁init; τ₁init; ϕ₁init; νinit]

## fit
# note this takes 61 iterations and converges after about 5 mins on my laptop
# this will return an answer even if the software hasn't converged, see status below

const mle = fit_bats_mle_covariates(dat, τmatrix, ϕmatrix, initguess, print_level=5)

# status description copied from https://github.com/jump-dev/Ipopt.jl/blob/eec0d01b685ff344444a6c77ebf89e9c301cb5d0/src/MOI_wrapper.jl#L1231
# From Ipopt/src/Interfaces/IpReturnCodes_inc.h
# const _STATUS_CODES = Dict(
#   0 => :Solve_Succeeded,
#   1 => :Solved_To_Acceptable_Level,
#   2 => :Infeasible_Problem_Detected,
#   3 => :Search_Direction_Becomes_Too_Small,
#   4 => :Diverging_Iterates,
#   5 => :User_Requested_Stop,
#   6 => :Feasible_Point_Found,
#   -1 => :Maximum_Iterations_Exceeded,
#   -2 => :Restoration_Failed,
#   -3 => :Error_In_Step_Computation,
#   -4 => :Maximum_CpuTime_Exceeded,
#   -5 => :Maximum_WallTime_Exceeded,
#   -10 => :Not_Enough_Degrees_Of_Freedom,
#   -11 => :Invalid_Problem_Definition,
#   -12 => :Invalid_Option,
#   -13 => :Invalid_Number_Detected,
#   -100 => :Unrecoverable_Exception,
#   -101 => :NonIpopt_Exception_Thrown,
#   -102 => :Insufficient_Memory,
#   -199 => :Internal_Error,
# )


#
# A simple toy atmospheric model that uses upwinded finite differences
# to simulate transport of a puff of non-reactive tracer gas.
#
# I wrote this to teach myself Juila.
# It's based on what I can remember of my PhD disseration from 12 years ago.
# This is absolutely _not_ how you go about solving this problem in the real world.
# I've no doubt there are much better packages out there for CFD.  Use those.
#
# Visualization is based on code from
# https://lazarusa.github.io/BeautifulMakie/surfWireLines/volumeTransparent/
#
# John Linford <john@redhpc.com>
#
# BSD 3-Clause License
#

using Dates
using GLMakie

"""
    advec_diff(...)

Advection and diffusion for one grid cell via upwind-biased finite differences.
Concentrations and wind values for the two neighboring cells are required.
"""
function advec_diff(
            Δx::Float64,
            c2l::Float64,
            c1l::Float64, w1l::Float64,
            c00::Float64, w00::Float64, d00::Float64,
            c1r::Float64, w1r::Float64,
            c2r::Float64)
    wind = (w1l + w00) / 2.0
    if wind > 0
        advecL = 1/6 * ( -c2l + 5.0*c1l + 2.0*c00 )
    else
        advecL = 1/6 * ( 2.0*c1l + 5.0*c00 - c1r )
    end
    advecL *= wind

    wind = (w1r + w00) / 2.0
    if wind > 0.0
        advecR = 1/6 * ( -c1l + 5.0*c00 + 2.0*c1r )
    else
        advecR = 1/6 * ( 2.0*c00 + 5.0*c1r - c2r )
    end
    advecR *= wind

    advec = (advecL - advecR) / Δx

    diffL = d00 * (c1l - c00)
    diffR = d00 * (c00 - c1r)
    diff = (diffL  -  diffR) / Δx^2

    return advec + diff
end

"""
    discretize1d(...)

Calculates advection/diffusion in 1D via `advec_diff` with boundaries wrapped around.

# Arguments
 - `Δx::Float64`: Space delta
 - `Δt::Float64`: Time delta
 - `d::Float64`: Diffusion coefficient
 - `c::Vector{Float64}`: Concentrations in ea
 - `w::Vector{Float64})`: Wind coefficients for each concentration
"""
function discretize1d(
            Δx::Float64,
            Δt::Float64,
            d::Float64,
            c::Vector{Float64},
            w::Vector{Float64})

    n = length(c)
    dcdx = Vector{Float64}(undef, n)

    # Boundaries
    dcdx[1] = advec_diff(Δx,
                         c[n-1],        # 2-left neighbor, wrapped around
                         c[n], w[n],    # 1-left neighbor, wrapped around
                         c[1], w[1], d, # Origin cell
                         c[2], w[2],    # 1-right neighbor
                         c[3])          # 2-right neighbor
    dcdx[2] = advec_diff(Δx,
                         c[n],          # 2-left neighbor, wrapped around
                         c[1], w[1],    # 1-left neighbor
                         c[2], w[2], d, # Origin cell
                         c[3], w[3],    # 1-right neighbor
                         c[4])          # 2-right neighbor

    # Interior
    for i in 3:n-2
        dcdx[i] = advec_diff(Δx,
                             c[i-2],            # 2-left neighbor
                             c[i-1], w[i-1],    # 1-left neighbor
                             c[i], w[i], d,     # Origin cell
                             c[i+1], w[i+1],    # 1-right neighbor
                             c[i+2])            # 2-right neighbor
    end

    # Boundaries
    dcdx[n-1] = advec_diff(Δx,
                           c[n-3],              # 2-left neighbor
                           c[n-2], w[n-2],      # 1-left neighbor
                           c[n-1], w[n-1], d,   # Origin cell
                           c[n], w[n],          # 1-right neighbor
                           c[1])                # 2-right neighbor, wrapped around
    dcdx[n] = advec_diff(Δx,
                         c[n-2],            # 2-left neighbor
                         c[n-1], w[n-1],    # 1-left neighbor
                         c[n], w[n], d,     # Origin cell
                         c[1], w[1],        # 1-right neighbor, wrapped around
                         c[2])              # 2-right neighbor, wrapped around

    c1 = c + Δt*dcdx
    # Mass loss here, if you care about that sort of thing
    return (c1 .* (c1 .> 0))
end

"""
    transport(...)

Discretization by first order upwind-biased finite differences on a possibly non-uniform grid.
3D discretization is achived via a 1D dimension split: x -> y -> z -> y -> x

# Arguments
 - `Δx`, `Δy`, `Δz`: Grid cell size in `x`, `y`, and `z` dimensions
 - `Δt`: timestep size
 - `dH`: Diffusion coefficient in x and y dimensions
 - `dV`: Diffusion coefficient in z dimension
 - `conc`: Concentration field
 - `wind`: Wind field
"""
function transport(
            Δx::Float64,
            Δy::Float64,
            Δz::Float64,
            Δt::Float64,
            dH::Float64,
            dV::Float64,
            conc::Array{Float64,4},
            wind::Array{Float64,4})

    nx, ny, nz, nspec = size(conc)

    for iy ∈ 1:ny, iz ∈ 1:nz, is ∈ 1:nspec
        conc[:,iy,iz,is] = discretize1d(Δx, Δt/2, dH, conc[:,iy,iz,is], wind[:,iy,iz,1])
    end
    for ix ∈ 1:nx, iz ∈ 1:nz, is ∈ 1:nspec
        conc[ix,:,iz,is] = discretize1d(Δy, Δt/2, dH, conc[ix,:,iz,is], wind[ix,:,iz,2])
    end
    for ix ∈ 1:nx, iy ∈ 1:ny, is ∈ 1:nspec
        conc[ix,iy,:,is] = discretize1d(Δz, Δt, dV, conc[ix,iy,:,is], wind[ix,iy,:,3])
    end
    for ix ∈ 1:nx, iz ∈ 1:nz, is ∈ 1:nspec
        conc[ix,:,iz,is] = discretize1d(Δy, Δt/2, dH, conc[ix,:,iz,is], wind[ix,:,iz,2])
    end
    for iy ∈ 1:ny, iz ∈ 1:nz, is ∈ 1:nspec
        conc[:,iy,iz,is] = discretize1d(Δx, Δt/2, dH, conc[:,iy,iz,is], wind[:,iy,iz,1])
    end
end

"""
    update_plot!(...)

Update Observables to advance the visualization.
"""
function update_plot!(points, colors, conc, spec, maxconc=8.61e9, minconc=1e-6)
    ptmp = Vector{Point3f0}()
    ctmp = Vector{Float64}()
    nx, ny, nz, _ = size(conc)
    for ix in 1:nx, iy in 1:ny, iz in 1:nz
        if conc[ix,iy,iz,spec] > minconc
            push!(ptmp, Point3f0(ix,iy,iz))
            push!(ctmp, conc[ix,iy,iz,spec]/maxconc)
        end
    end
    # Holy perf pain batman!  Don't update those observables until you absolutely have to!
    # One update outside the loop is at least 100x faster than updating in place.
    points[] = ptmp
    colors[] = ctmp
end

# Probably should be a main() function or something....
let
    # Time grid
    tstart = Dates.DateTime(2020, 01, 01)
    tend = tstart + Dates.Hour(24)
    Δt = Dates.Minute(15)

    # Space grid
    nx, ny, nz = 50, 50, 25

    # Grid cell dimensions (m)
    Δx, Δy, Δz = 1000.0, 1000.0, 1000.0

    # Horizontal and vertical diffusion coefficients (m²/s)
    dH, dV = 100.0, 50.0

    # Number of chemical species
    nspec = 1

    # Concentration field
    # Concentrations of the same species are contiguous
    conc = Array{Float64,4}(undef, nx, ny, nz, nspec)
    let
        conc .= 0.0
        # Dummy up a tracer point source in the center of the concentration field
        source_coords = size(conc)[1:3] .÷ 2
        conc[source_coords...] = 8.61e9
    end

    # Wind field
    wind = Array{Float64,4}(undef, nx, ny, nz, 3)
    wind_init = "simple"
    if wind_init == "simple"
        # 0.5m/sec cross wind in x-dimension
        wind[:,:,:,1] .= 0.5
        wind[:,:,:,2] .= 0.0
        wind[:,:,:,3] .= 0.0
    elseif wind_init == "cyclone"
        # Cyclonic wind
        for ix ∈ 1:nx, iy ∈ 1:ny, iz ∈ 1:nz
            wind[ix,iy,iz,:] = 0.5 .* [cos(nx/2 - ix), sin(ny/2 - iy), 0.25]
        end
    else
        error("Invalid wind_init: $wind_init")
    end

    # Temperature field
    # 297.15 °K == 24 °C == 75.2 °F
    temp = Array{Float64,3}(undef, nx, ny, nz)
    temp .= 297.15

    # Visualization
    fig = Figure(resolution=(1200,1200))
    ax = Axis3(fig[1,1], perspectiveness=0.5, azimuth=π/3, elevation=π/6, aspect=:data)
    points = Node(Vector{Point3f0}())
    colors = Node(Vector{Float64}())

    # Initial plot
    limits!(ax, 1, nx, 1, ny, 1, nz)
    update_plot!(points, colors, conc, 1)
    meshscatter!(ax, points, color=colors, colormap=RGBAf0.(to_colormap(:turbo, 10), LinRange(0,1,10)))

    # Time iteration with in-situ visualization
    record(fig, "plume.mp4", tstart:Δt:tend-Δt) do tnow
        println(tnow)
        transport(Δx, Δy, Δz, Float64(Dates.value(Dates.Second(Δt))), dH, dV, conc, wind)
        update_plot!(points, colors, conc, 1)
    end

end #let


using MiniGNN
using JLD, HDF5
# Beware! https://github.com/JuliaIO/JLD.jl/issues/234
using Flux.Data: DataLoader
using CUDA
using MiniFastai
using Flux: logitcrossentropy, onehotbatch, Descent, Dense, Dropout, relu, Chain
using SparseArrays
using MiniGNN: create_batch

@load "train_data.jld" train_data
#size Nsamples x V
@load "train_labels.jld" train_labels
@load "test_data.jld" test_data
@load "test_labels.jld" test_labels

#small = 1000
#train_data = train_data[1:small,:]
#train_labels = train_labels[1:small]
train_labels = onehotbatch(train_labels, 0:9)
#size Nlabels x Nsamples
test_labels = onehotbatch(test_labels,0:9)
d = load("List_graphs.jld")
L_g = d["L_g"]
sizes = d["sizes"]

bs = 100

train_loader = DataLoader((permutedims(train_data),train_labels), batchsize=bs, shuffle=true)
val_loader = DataLoader((permutedims(test_data),test_labels), batchsize=bs, shuffle=false)
databunch = Databunch(train_loader,val_loader)

batch, labels = create_batch(100,train_loader)
input = convert(Array{Float32,3},process_data(batch))

p=0.5
mat1 = MiniGNN.make_rescaled_mat(L_g[1])
layer1 = MiniGNN.GraphConv(mat1,1=>32,5)
mp = MiniGNN.GraphMaxPool(4)
mat2 = MiniGNN.make_rescaled_mat(L_g[3])
layer2 = MiniGNN.GraphConv(mat2,32=>64,5)
dim_fc = Int(64*size(mat2)[1]/4)
model = Chain(layer1 , mp, layer2, mp, MiniGNN.graphflatten, Dense(dim_fc,512, relu), Dropout(p), Dense(512,10))

learning_rate = 0.05
SGD = Descent(learning_rate)
loss(pred,y) = logitcrossentropy(pred,y)

learner_gnn = Learner(model,SGD,loss,databunch,cbs=(AvgStatsCallback(),))
transform = x -> convert(Array{Float32,3},process_data(x))
fit!(learner_gnn,30,transform)

model(input)
#one batch
using Flux
using Zygote: pullback, update!
ps = params(model)
            #source https://fluxml.ai/Flux.jl/stable/training/training/
            # https://fluxml.ai/Zygote.jl/latest/adjoints/#Pullbacks-1
y_pred, back_net = pullback(() -> model(input), ps)
current_loss, back_loss = pullback(y -> loss(y,labels),y_pred)
            #current_loss, back = pullback(() -> los_f(mod(xb),yb), ps)
gs = back_net(back_loss(1)[1])#back(one(current_loss))
update!(SGD,ps,gs)


"""
    geometric_factors(x, y, Dr, Ds)
    geometric_factors(x, y, z, Dr, Ds, Dt, Filters=(I,I,I))

Compute metrics of mappings between "real" elements and reference elements,
outward pointing normals on faces of every elements, and Jacobian.

x,y,z are arrays of coordinates, and Dr, Ds, Dt are nodal differentiation matrices
Filters = tuple of filtering matrices (e.g., to reduce degree in r,s, and t for GCL)

Geometric terms in 3D are constructed to ensure satisfaction of free-stream preservation
using the curl-based construction from 'Metric identities and the DG-SEM on curvilinear
meshes' (Kopriva 2006).
"""

function geometric_factors(x, y, Dr, Ds)
    xr, xs = Dr*x, Ds*x
    yr, ys = Dr*y, Ds*y

    J = @. -xs * yr + xr * ys
    rxJ, sxJ =  ys, -yr
    ryJ, syJ = -xs,  xr

    return rxJ, sxJ, ryJ, syJ, J
end

function geometric_factors(x, y, z, Dr, Ds, Dt, Filters=(I,I,I))

    xr, xs, xt = Dr*x, Ds*x, Dt*x
    yr, ys, yt = Dr*y, Ds*y, Dt*y
    zr, zs, zt = Dr*z, Ds*z, Dt*z

    Fr, Fs, Ft = (Dr*y).*z, (Ds*y).*z, (Dt*y).*z
    Fr,Fs,Ft = ((A,x)->A*x).(Filters,(Fr,Fs,Ft))
    rxJ = Dt*(Fs) - Ds*(Ft)
    sxJ = Dr*(Ft) - Dt*(Fr)
    txJ = Ds*(Fr) - Dr*(Fs)

    Fr, Fs, Ft = (Dr*x).*z, (Ds*x).*z, (Dt*x).*z
    Fr,Fs,Ft = ((A,x)->A*x).(Filters,(Fr,Fs,Ft))
    ryJ = -(Dt*(Fs) - Ds*(Ft))
    syJ = -(Dr*(Ft) - Dt*(Fr))
    tyJ = -(Ds*(Fr) - Dr*(Fs))

    Fr, Fs, Ft = (Dr*y).*x, (Ds*y).*x, (Dt*y).*x
    Fr,Fs,Ft = ((A,x)->A*x).(Filters,(Fr,Fs,Ft))
    rzJ = -(Dt*(Fs) - Ds*(Ft))
    szJ = -(Dr*(Ft) - Dt*(Fr))
    tzJ = -(Ds*(Fr) - Dr*(Fs))

    J = @. xr * (ys * zt - zs * yt) - yr * (xs * zt - zs * xt) + zr * (xs * yt - ys * xt)

    return rxJ, sxJ, txJ, ryJ, syJ, tyJ, rzJ, szJ, tzJ, J
end

"""
    estimate_h(rd::RefElemData, md::MeshData)
    estimate_h(e, rd::RefElemData, md::MeshData) # e = element index

Estimates the mesh size via min size_of_domain * |J|/|sJ|, since |J| = O(hᵈ) and |sJ| = O(hᵈ⁻¹).
"""
function estimate_h(rd::RefElemData{DIM}, md::MeshData{DIM}) where {DIM}
    hmin = Inf
    for e in 1:md.num_elements
        h_e = estimate_h(e, rd, md)
        hmin = min(hmin, h_e)
    end
    return hmin * compute_domain_size(rd, md)^(1/DIM)
end

function estimate_h(e, rd::RefElemData, md::MeshData)
    sJ_e = reshape(view(md.sJ, :, e), rd.Nfq ÷ rd.Nfaces, rd.Nfaces)
    sJ_face = zero(eltype(md.sJ))
    for f in 1:rd.Nfaces
        sJ_face = max(sJ_face, minimum(view(sJ_e, :, f)) / face_scaling(rd, f))
    end
    h_e = minimum(view(md.J, :, e)) / sJ_face
    return h_e
end

face_scaling(rd, f) = 1.0
face_scaling(rd::RefElemData{2, Tri}, f) = f == 3 ? sqrt(2) : 1.0 # sJ incorporates length of long triangle edge
face_scaling(rd::RefElemData{3, Tet}, f) = f == 2 ? sqrt(3) : 1.0 # sJ incorporates area of larger triangle face
compute_domain_size(rd::RefElemData, md::MeshData) = sum(rd.M * md.J)


function highpass_filter(X::Array{Float64,1},sampling_rate::Float64, cutoff::Float64=300.0;method=Butterworth(4))
  filtfilt(digitalfilter(Highpass(cutoff;fs=sampling_rate), method), X)
end

function get_threshold(X::Array{Float64,1},θ=6.0)
  n = length(X)
  μ = mean(X)
  σ = std(X)::Float64 #because of type instability in std
  l = μ - θ*σ
  u = μ + θ*σ
  σ2 = 0.0
  for x in X
    if l < x < u
      σ2 +=  x*x
    end
  end
  σ2 = sqrt(σ2/n - μ*μ)
  μ, σ2
end

function extract_spikes(X::Array{Float64,1};μ::Float64=NaN, σ::Float64=NaN,nq::Tuple{Int64, Int64}=(20,40),θ=6.0,only_negative=true)
  if isnan(σ) || isnan(μ)
    μ, σ = get_threshold(X,θ)
  end
  n = length(X)
  l = μ - θ*σ
  u = μ + θ*σ
  pidx = Array{Int64}(0)
  xmin = Inf
  xmax = -Inf
  j = 1
  in_peak = false
  _idx = 0
  while j <= n
    #check whether points are outside μ ± σ
    #fast forward points that are within noise levels
    x = X[j]
    if x < l
      if x < xmin
        xmin = x
        _idx = j
      end
      in_peak = true
  elseif (x > u ) && !only_negative
      if x > xmax
        xmax = x
        _idx = j
      end
      in_peak = true
    else
      if in_peak == true #just exited a peak; reset
        push!(pidx, _idx)
      end
      xmin = Inf
      xmax = -Inf
      in_peak = false
    end
    j += 1
  end
  pidx, waveforms = extract_spikes(X, pidx,nq)
  pidx, waveforms
end

function extract_spikes(X::Array{Float64,1}, idx::Array{Int64,1}, n::Tuple{Int64, Int64}=(20,40))
  nt = sum(n)
  nb,na = n
  idxl = find(x->x>nb, idx)
  nx = length(idxl)
  waveforms = zeros(nt, nx)
  for i in 1:nx
     _idx = idx[idxl[i]]
    for j in -nb:na-1
      waveforms[j+nb+1,i] = X[_idx + j]
    end
  end
  idx[idxl], waveforms
end


function angle(vec1::Union{Array{F,1},Array{Int,1}},vec2::Union{Array{F,1},Array{Int,1}}) where {F <: AbstractFloat}
    acos(dot(vec1, vec2) / (norm(vec1) * norm(vec2))) * 180 / pi
    #TODO: fix DomainError with 1.0000000000000002 problem
end

function vecs_to_triangle(veci::Union{Array{F,1},BitArray{1}},
                          vecj::Union{Array{F,1},BitArray{1}},
                          veck::Union{Array{F,1},BitArray{1}}) where {F <: AbstractFloat}
    seg_ij = vecj - veci
    seg_ik = veck - veci
    seg_jk = veck - vecj

    angle1 = angle(seg_ij, seg_ik)
    angle2 = angle(-seg_ij, seg_jk)
    angle3 = 180.0 - angle1 - angle2
    return (angle1, angle2, angle3)
end

function brute_force_triangles(pts::Array{F,2}) where {F <: AbstractFloat}

    n,d = size(pts)
    # n choose 3 triangles
    triangles = Array{Tuple{Tuple{Int,Int,Int},Tuple{F,F,F}},1}(undef,binomial(n,3))

    index = 1
    for i =1:n
        for j =1:i-1
            for k = 1:j-1
                x = ((k,j,i),vecs_to_triangle(pts[k,:],pts[j,:],pts[i,:]))
                triangles[index] = x
                index += 1
            end
        end
    end

    return triangles
end

function brute_force_triangles(pts::Array{BitArray{1},1}) where {F <: AbstractFloat}

    n = length(pts)
    # n choose 3 triangles
    #TODO: fix F type error
    triangles = Array{Tuple{Tuple{Int,Int,Int},Tuple{Float64,Float64,Float64}},1}(undef,binomial(n,3))

    index = 1
    for i =1:n
        for j =1:i-1
            for k = 1:j-1
                x = ((k,j,i),vecs_to_triangle(pts[k],pts[j],pts[i]))
                triangles[index] = x
                index += 1
            end
        end
    end

    return triangles

end


@testset "Maze" begin

UP = CartesianIndex(-1, 0)
DOWN = CartesianIndex(1, 0)
LEFT = CartesianIndex(0, -1)
RIGHT = CartesianIndex(0, 1)
DIRECTIONS = [UP, DOWN, LEFT, RIGHT]

manhattan(a::CartesianIndex, b::CartesianIndex) = sum(abs.((b - a).I))
getmazeneighbours(maze, state) = filter(x -> (1 <= x[1] <= size(maze)[1]) && (1 <= x[2] <= size(maze)[2]) && (!maze[x]), [state + d for d in DIRECTIONS])

@testset "Maze goal reachable" begin
  maze = [0 0 1 0 0;
          0 1 0 0 0;
          0 1 0 0 1;
          0 0 0 1 1;
          1 0 1 0 0] .== 1
  start = CartesianIndex(1, 1)
  goal = CartesianIndex(1, 5)

  neighbours(state) = getmazeneighbours(maze, state)

  res = astar(neighbours, start, goal, heuristic=manhattan)
  @test res.status == :success
  @test res.path ==  CartesianIndex{2}[CartesianIndex(1, 1), CartesianIndex(2, 1), CartesianIndex(3, 1), CartesianIndex(4, 1), CartesianIndex(4, 2), CartesianIndex(4, 3), CartesianIndex(3, 3), CartesianIndex(2, 3), CartesianIndex(2, 4), CartesianIndex(1, 4), CartesianIndex(1, 5)]
  @test res.cost == 10
end

@testset "Maze Goal Unreachable" begin
  maze = [0 0 1 0 0;
          0 1 0 0 0;
          0 1 0 0 1;
          1 0 0 1 1;
          1 0 1 0 0] .== 1
  start = CartesianIndex(1, 1)
  goal = CartesianIndex(1, 5)

  neighbours(state) = getmazeneighbours(maze, state)

  res = astar(neighbours, start, goal, heuristic=manhattan)
  @test res.status == :nopath
  @test res.path == CartesianIndex{2}[CartesianIndex(1, 1), CartesianIndex(1, 2)]
  @test res.cost == 1
end

end


"""
    function mincost_flow(g::AbstractGraph,
    		node_demand::AbstractVector,
    		edge_capacity::AbstractMatrix,
    		edge_cost::AbstractMatrix,
    		optimizer;
    		edge_demand::Union{Nothing,AbstractMatrix} = nothing,
    		source_nodes = (),
    		sink_nodes = ()
    		)

Find a flow over a directed graph `g` satisfying the `node_demand` at each
node and `edge_capacity` constraints for each edge while minimizing `dot(edge_cost, flow)`.

Returns a sparse flow matrix, flow[i,j] corresponds to the flow on the (i,j) arc.

# Arguments

- `g` is a directed `LightGraphs.AbstractGraph`.
-`node_demand` is a vector of nodal demand values, which should be negative for sources,
positive for sink nodes, and zero for all other nodes.
- `edge_capacity::AbstractMatrix` sets an upper bound on the flow of each arc.
- `edge_cost::AbstractMatrix` the cost per unit of flow on each arc.
- `optimizer` is an optimizer constructor, like `Clp.Optimizer` or `() -> Clp.Optimizer()` passed at the construction of the `JuMP.Model`.

# Keyword arguments

- `edge_demand::Union{Nothing,AbstractMatrix}`: require a minimum flow for edges, or nothing.
- `source_nodes` Collection of sources at which the nodal netflow is allowed to be greater than nodal demand, defaults to an empty tuple.
- `sink_nodes` Collection of sinks at which the nodal netflow is allowed to be less than nodal demand, defaults to an empty tuple.

`source_nodes` & `sink_nodes` are only needed when nodal flow are not explictly set in node_demand

### Usage Example:

```julia
julia> import LightGraphs
julia> const LG = LightGraphs
julia> using LightGraphsFlows: mincost_flow
julia> import Clp # use your favorite LP solver here
julia> using SparseArrays: spzeros
julia> g = LG.DiGraph(6) # Create a flow-graph
julia> LG.add_edge!(g, 5, 1)
julia> LG.add_edge!(g, 5, 2)
julia> LG.add_edge!(g, 3, 6)
julia> LG.add_edge!(g, 4, 6)
julia> LG.add_edge!(g, 1, 3)
julia> LG.add_edge!(g, 1, 4)
julia> LG.add_edge!(g, 2, 3)
julia> LG.add_edge!(g, 2, 4)
julia> cost = zeros(6,6)
julia> cost[1,3] = 10
julia> cost[1,4] = 5
julia> cost[2,3] = 2
julia> cost[2,4] = 2
julia> demand = spzeros(6)
julia> demand[5] = -2
julia> demand[6] = 2
julia> capacity = ones(6,6)
julia> flow = mincost_flow(g, demand, capacity, cost, Clp.Optimizer)
```
"""
function mincost_flow end

@traitfn function mincost_flow(g::AG::lg.IsDirected,
		node_demand::AbstractVector,
		edge_capacity::AbstractMatrix,
		edge_cost::AbstractMatrix,
		optimizer;
		edge_demand::Union{Nothing,AbstractMatrix} = nothing,
		source_nodes = (), # Source nodes at which to allow a netflow greater than nodal demand
		sink_nodes = ()	   # Sink nodes at which to allow a netflow less than nodal demand
		) where {AG <: lg.AbstractGraph}

	m = JuMP.Model(optimizer)
	vtxs = vertices(g)

	source_nodes = [v for v in vtxs if v in source_nodes || node_demand[v] < 0]
	sink_nodes = [v for v in vtxs if v in sink_nodes || node_demand[v] > 0]

	@variable(m, 0 <= f[i=vtxs,j=vtxs; (i,j) in lg.edges(g)] <= edge_capacity[i, j])
	@objective(m, Min, sum(f[src(e),dst(e)] * edge_cost[src(e), dst(e)] for e in lg.edges(g)))

	for v in lg.vertices(g)
	    if v in source_nodes
            @constraint(m,
                sum(f[v, vout] for vout in outneighbors(g, v)) - sum(f[vin, v] for vin in lg.inneighbors(g, v)) >= -node_demand[v]
            )
	    elseif v in sink_nodes
            @constraint(m,
                sum(f[vin, v] for vin in lg.inneighbors(g, v)) - sum(f[v, vout] for vout in outneighbors(g, v)) >= node_demand[v]
            )
	    else
            @constraint(m,
                sum(f[vin, v] for vin in lg.inneighbors(g, v)) == sum(f[v, vout] for vout in outneighbors(g, v))
            )
        end
	end

    if edge_demand isa AbstractMatrix
        for e in lg.edges(g)
            (i,j) = Tuple(e)
            JuMP.set_lower_bound(f[i,j], edge_demand[i,j])
        end
    end
    optimize!(m)
    ts = termination_status(m)
    result_flow = spzeros(nv(g), nv(g))
    if ts != MOI.OPTIMAL
        @warn "Problem does not have an optimal solution, status: $(ts)"
        return result_flow
    end
    for e in lg.edges(g)
        (i,j) = Tuple(e)
        result_flow[i,j] = JuMP.value(f[i,j])
    end
    return result_flow
end


"""
    f = initmetagauss(sx,st,nso,nto)

Generate the spectral amplitudes `f` for a metagaussian field of size `nso * nso * nto`
with slopes `sx` and `st`.
"""
function initmetagauss(sx,st,ns,nt);

        sx=abs(sx);
        st=abs(st);

#       This spares a lot of memory!
#       [kx,ky,kt]=ndgrid([0:ns/2 -ns/2+1:-1],[0:ns/2 -ns/2+1:-1],[0:nt/2 -nt/2+1:-1]);
        kx=[collect(0:ns/2)' collect(-ns/2+1:-1)']';
        kx=repeat(kx*ones(1,ns),outer=[1,1,nt]);
# kky=permdute(kkx,[2 1 3]);

if(nt>1)
        kt=[collect(0:nt/2)' collect(-nt/2+1:-1)']';
        kt=repeat(kt*ones(1,ns),outer=[1 1 ns]);
        kt=permute(kt,[3 2 1]).^2;

        kt[:,:,1]=0.000001;
end

        kx=kx.^2+permutedims(kx,[2 1 3]).^2;

        kx[1,1,:].=0.000001;
# f = (kx.^(-(sx+1)/4)).*kt.^(-st/4);
    if(nt>1)
        kx = (kx.^(-(sx+1)/4)).*kt.^(-st/4);
    else
        kx = (kx.^(-(sx+1)/4));
    end
    kx[1,1,1]=0;
    kx[1,1,:].=0;
    if(nt>1)
       kx[:,:,1]=0;
    end
   #  kx=kx./sqrt(sum(lin(abs(kx).^2)))*ns*ns*nt;
   kx=kx./sqrt(sum(abs2.(kx[:])))*ns*ns*nt;
end


"""
normalized water saturation
sws(sw::Real, sor::Real, swc::Real)
"""
function sws(sw::Real, sor::Real, swc::Real)
  if sw>swc && sw<1-sor
    res=(sw-swc)/(1-sor-swc)
  elseif sw>=1-sor
    res=1.0
  elseif sw<=swc
    res=0.0
  end
  return res
end

"""
Oil relative permeability functions for oil-water system
kro(sw::Real, kro0::Real, sor::Real, swc::Real, no::Real)
"""
function kro(sw::Real, kro0::Real, sor::Real, swc::Real, no::Real)
  if swc<=sw<=1-sor
    res=kro0*((1-sw-sor)/(1-sor-swc))^no
  elseif 0.0<sw<swc
    res=1+(kro0-1)/swc*sw
  elseif sw>1-sor
    res=0.0
  elseif sw<=0.0
    res=1.0
  end
  return res
end

"""
derivative of the oil relperm with respect to water saturation
dkrodsw(sw::Real, kro0::Real, sor::Real, swc::Real, no::Real)
"""
function dkrodsw(sw::Real, kro0::Real, sor::Real, swc::Real, no::Real)
  if swc<=sw<=1-sor
    res=-no*kro0/(1-sor-swc)*((1-sw-sor)/(1-sor-swc))^(no-1)
  elseif 0.0<sw<swc
    res=(kro0-1)/swc
  elseif sw>1-sor || sw<=0.0
    res=0.0
  end
  return res
end

"""
water rel-perm for oil-water system
krw(sw::Real, krw0::Real, sor::Real, swc::Real, nw::Real)
"""
function krw(sw::Real, krw0::Real, sor::Real, swc::Real, nw::Real)
  if swc<=sw<=1-sor
    res=krw0*((sw-swc)/(1-sor-swc))^nw
  elseif 1-sor<sw<1.0
    res=-(1-krw0)/sor*(1.0-sw)+1.0
  elseif sw<swc
    res=0.0
  elseif sw>=1.0
    res=1.0
  end
  return res
end

"""
derivative of the water relperm with respect to water saturation
dkrwdsw(sw::Real, krw0::Real, sor::Real, swc::Real, nw::Real)
"""
function dkrwdsw(sw::Real, krw0::Real, sor::Real, swc::Real, nw::Real)
  if swc<=sw<=1-sor
    res=nw*krw0/(1-sor-swc)*((sw-swc)/(1-sor-swc))^(nw-1)
  elseif 1-sor<sw<1.0
    res=(1-krw0)/sor
  elseif sw<swc || sw>=1.0
    res=0.0
  end
  return res
end

"""
capillary pressure curve for drainage
pc_drain(sw::Real, pce::Real, swc::Real; labda::Real=2.4, pc0::Real=1.0e7)
"""
function pc_drain(sw::Real, pce::Real, swc::Real; labda::Real=2.4, pc0::Real=1.0e7)
  sw0=swc+(1-labda*log(pc0/pce)+sqrt((-1+labda*log(pc0/pce))^2+4*swc/(1-swc)))/2*(1-swc)
  if sw>sw0
    res=pce*((sw-swc)/(1-swc))^(-1.0/labda)
  elseif 0.0<=sw<=sw0
    pcs=pce*((sw0-swc)/(1-swc))^(-1.0/labda)
    res=exp((log(pcs)-log(pc0))/sw0*(sw-sw0)+log(pcs))
  elseif sw<0
    res=pc0
  else
    res=0.0
  end
end

"""
derivative of the drainage capillary pressure with respect to water saturation
dpc_drain(sw::Real, pce::Real, swc::Real; labda::Real=2.4, pc0::Real=1.0e7)
"""
function dpc_drain(sw::Real, pce::Real, swc::Real; labda::Real=2.4, pc0::Real=1.0e7)
  sw0=swc+(1-labda*log(pc0/pce)+sqrt((-1+labda*log(pc0/pce))^2+4*swc/(1-swc)))/2*(1-swc)
  if sw>sw0
    res=-1.0/((1-swc)*labda)*pce*((sw-swc)/(1-swc))^(-1.0/labda-1)
  elseif 0.0<=sw<=sw0
    res=-1.0/((1-swc)*labda)*pce*((sw0-swc)/(1-swc))^(-1.0/labda-1)
  else
    res=0.0
  end
end

"""
second derivative of the drainage capillary pressure with respect to water saturation
dpc_drain(sw::Real, pce::Real, swc::Real; labda::Real=2.4, pc0::Real=1.0e7)
"""
function d2pc_drain(sw::Real, pce::Real, swc::Real; labda::Real=2.4, pc0::Real=1.0e7)
  sw0=swc+(1-labda*log(pc0/pce)+sqrt((-1+labda*log(pc0/pce))^2+4*swc/(1-swc)))/2*(1-swc)
  if sw>sw0
    res=-1.0/((1-swc)^2*labda)*(-1.0/labda-1)*pce*((sw-swc)/(1-swc))^(-1.0/labda-2.0)
  elseif 0.0<=sw<=sw0
    res=0.0
  else
    res=0.0
  end
end

"""
Imbibition capillary pressure curve
pc_imb(sw::Real, pce::Real, swc::Real, sor::Real;
  teta::Real=0.785, labda::Real=2.4, b::Real=0.6, pc01::Real=1.0e7, pc02::Real=1.0e7)
"""
function pc_imb(sw::Real, pce::Real, swc::Real, sor::Real;
  teta::Real=0.785, labda::Real=2.4, b::Real=0.6, pc01::Real=1.0e7, pc02::Real=1.0e7)
  pc1=pc_drain(sw, pce, swc, labda=labda, pc0=pc01)
  pc2=pc_drain(1-sw, pce, sor, labda=labda, pc0=pc02)
  return (0.5*(1+cos(teta)))^b*pc1-(0.5*(1-cos(teta)))^b*pc2
end

"""
derivative of the imbibition capillary pressure curve w.r.t water saturation
pc_imb(sw::Real, pce::Real, swc::Real, sor::Real;
  teta::Real=0.785, labda::Real=2.4, b::Real=0.6, pc01::Real=1.0e7, pc02::Real=1.0e7)
"""
function dpc_imb(sw::Real, pce::Real, swc::Real, sor::Real;
  teta::Real=0.785, labda1::Real=2.4, labda2::Real=2.4, b::Real=0.6, pc01::Real=1.0e7, pc02::Real=1.0e7)
  dpc1=dpc_drain(sw, pce, swc, labda=labda1, pc0=pc01)
  dpc2=dpc_drain(1-sw, pce, sor, labda=labda2, pc0=pc02)
  return (0.5*(1+cos(teta)))^b*dpc1+(0.5*(1-cos(teta)))^b*dpc2
end

"""
second derivative of the imbibition capillary pressure curve w.r.t water saturation
pc_imb(sw::Real, pce::Real, swc::Real, sor::Real;
  teta::Real=0.785, labda::Real=2.4, b::Real=0.6, pc01::Real=1.0e7, pc02::Real=1.0e7)
"""
function d2pc_imb(sw::Real, pce::Real, swc::Real, sor::Real;
  teta::Real=0.785, labda1::Real=2.4, labda2::Real=2.4, b::Real=0.6, pc01::Real=1.0e7, pc02::Real=1.0e7)
  dpc1=d2pc_drain(sw, pce, swc, labda=labda1, pc0=pc01)
  dpc2=d2pc_drain(1-sw, pce, sor, labda=labda2, pc0=pc02)
  return (0.5*(1+cos(teta)))^b*dpc1+(0.5*(1-cos(teta)))^b*dpc2
end

"""
SECOND approach (see the doc folder)
capillary pressure curve for drainage
pc_drain(sw::Real, pce::Real, swc::Real; labda::Real=2.4, pc0::Real=1.0e6)
"""
function pc_drain2(sw::Real, pce::Real, swc::Real; labda::Real=2.4, pc_star::Real=1.0e6)
  sw_star=(1-swc)*exp(-labda*log(pc_star/pce))+swc
  pc0 = pc_star*exp(sw_star/(labda*(1-swc)))
  if sw>sw_star
    res=pce*((sw-swc)/(1.0-swc))^(-1.0/labda)
  elseif 0.0<=sw<=sw_star
    res=pc_star*exp((sw_star-sw)/(labda*(1.0-swc)))
  elseif sw<0
    res=pc0
  else
    res=0.0
  end
end

"""

derivative of the drainage capillary pressure with respect to water saturation
dpc_drain(sw::Real, pce::Real, swc::Real; labda::Real=2.4, pc0::Real=1.0e6)
"""
function dpc_drain2(sw::Real, pce::Real, swc::Real; labda::Real=2.4, pc_star::Real=1.0e6)
  sw_star=(1-swc)*exp(-labda*log(pc_star/pce))+swc
  if sw>sw_star
    res=-1.0/((1-swc)*labda)*pce*((sw-swc)/(1-swc))^(-1.0/labda-1.0)
  elseif 0.0<=sw<=sw_star
    res=pc_star/(sw-sw_star)
  else
    res=0.0
  end
end

"""
Imbibition capillary pressure curve
pc_imb(sw::Real, pce::Real, swc::Real, sor::Real;
  teta::Real=0.785, labda::Real=2.4, b::Real=0.6, pc01::Real=1.0e6, pc02::Real=1.0e6)
"""
function pc_imb2(sw::Real, pce::Real, swc::Real, sor::Real;
  teta::Real=0.785, labda::Real=2.4, b::Real=0.6, pc_star1::Real=1.0e6, pc_star2::Real=1.0e6)
  pc1=pc_drain2(sw, pce, swc, labda=labda, pc_star=pc_star1)
  pc2=pc_drain2(1-sw, pce, sor, labda=labda, pc_star=pc_star2)
  return (0.5*(1+cos(teta)))^b*pc1-(0.5*(1-cos(teta)))^b*pc2
end

"""
derivative of the imbibition capillary pressure curve w.r.t water saturation
pc_imb(sw::Real, pce::Real, swc::Real, sor::Real;
  teta::Real=0.785, labda::Real=2.4, b::Real=0.6, pc01::Real=1.0e6, pc02::Real=1.0e6)
"""
function dpc_imb2(sw::Real, pce::Real, swc::Real, sor::Real;
  teta::Real=0.785, labda1::Real=2.4, labda2::Real=2.4, b::Real=0.6, pc_star1::Real=1.0e6, pc_star2::Real=1.0e7)
  dpc1=dpc_drain2(sw, pce, swc, labda=labda1, pc_star=pc_star1)
  dpc2=dpc_drain2(1-sw, pce, sor, labda=labda2, pc_star=pc_star2)
  return (0.5*(1+cos(teta)))^b*dpc1+(0.5*(1-cos(teta)))^b*dpc2
end

"""
Third approach (see the doc folder)
capillary pressure curve for drainage
pc_drain(sw::Real, pce::Real, swc::Real; labda::Real=2.4, pc0::Real=1.0e6)
"""
function pc_drain3(sw::Real, pce::Real, swc::Real; labda::Real=2.4, pc_star::Real=1.0e6, c::Real=1e-3)
  # pc_max  = pc_max_multiplyer*pc_star
  sw_star=(1-swc)*exp(-labda*log(pc_star/pce))+swc
  # m = -pce/(labda*(1-swc))*((sw_star-swc)/(1-swc))^(-1.0/labda-1.0)
  # c = 0.5*(-1.0+sqrt(1.0-(pc_max-pc_star)/(m*(exp(sw_star)-1))))
  b = -pce*c/(labda*(1-swc))*((sw_star-swc)/(1-swc))^(-1.0/labda-1.0)
  a = pc_star + b*exp(c)
  pc0 = a-b*exp(sw_star+c)
  if sw>sw_star
    res=pce*((sw-swc)/(1-swc))^(-1.0/labda)
  elseif 0.0<=sw<=sw_star
    res=a-b*exp(sw_star-sw+c)
  elseif sw<0
    res=pc0
  else
    res=0.0
  end
end

"""
derivative of the drainage capillary pressure with respect to water saturation
dpc_drain(sw::Real, pce::Real, swc::Real; labda::Real=2.4, pc0::Real=1.0e6)
"""
function dpc_drain3(sw::Real, pce::Real, swc::Real; labda::Real=2.4, pc_star::Real=1.0e6, c::Real=1e-3)
  sw_star=(1-swc)*exp(-labda*log(pc_star/pce))+swc
  b = -pce*c/(labda*(1-swc))*((sw_star-swc)/(1-swc))^(-1.0/labda-1.0)
  a = pc_star + b*exp(c)
  pc0 = a-b*exp(sw_star+c)
  if sw>sw_star
    res=-1.0/((1-swc)*labda)*pce*((sw-swc)/(1-swc))^(-1.0/labda-1)
  elseif 0.0<=sw<=sw_star
    res=b/(sw_star-sw+c)
  else
    res=b/(sw_star+c)
  end
end

"""
second derivative of the drainage capillary pressure with respect to water saturation
dpc_drain(sw::Real, pce::Real, swc::Real; labda::Real=2.4, pc0::Real=1.0e6)
"""
function d2pc_drain3(sw::Real, pce::Real, swc::Real; labda::Real=2.4, pc_star::Real=1.0e6, c::Real=1e-3)
  sw_star=(1-swc)*exp(-labda*log(pc_star/pce))+swc
  b = -pce*c/(labda*(1-swc))*((sw_star-swc)/(1-swc))^(-1.0/labda-1.0)
  a = pc_star + b*exp(c)
  pc0 = a-b*exp(sw_star+c)
  if sw>sw_star
    res=-1.0/((1-swc)^2*labda)*(-1.0/labda-1.0)*pce*((sw-swc)/(1-swc))^(-1.0/labda-2.0)
  elseif 0.0<=sw<=sw_star
    res=b/(sw_star-sw+c)^2
  else
    res=b/(sw_star+c)^2
  end
end

"""
Imbibition capillary pressure curve
pc_imb(sw::Real, pce::Real, swc::Real, sor::Real;
  teta::Real=0.785, labda::Real=2.4, b::Real=0.6, pc01::Real=1.0e6, pc02::Real=1.0e6)
"""
function pc_imb3(sw::Real, pce::Real, swc::Real, sor::Real;
  teta::Real=0.785, labda::Real=2.4, b::Real=1.0, pc_star1::Real=1.0e6, pc_star2::Real=1.0e6, c::Real=1e-3)
  pc1=pc_drain3(sw, pce, swc, labda=labda, pc_star=pc_star1, c=c)
  pc2=pc_drain3(1-sw, pce, sor, labda=labda, pc_star=pc_star2, c=c)
  return (0.5*(1+cos(teta)))^b*pc1-(0.5*(1-cos(teta)))^b*pc2
end

"""
derivative of the imbibition capillary pressure curve w.r.t water saturation
pc_imb(sw::Real, pce::Real, swc::Real, sor::Real;
  teta::Real=0.785, labda::Real=2.4, b::Real=0.6, pc01::Real=1.0e6, pc02::Real=1.0e6)
"""
function dpc_imb3(sw::Real, pce::Real, swc::Real, sor::Real;
  teta::Real=0.785, labda1::Real=2.4, labda2::Real=2.4, b::Real=1.0, pc_star1::Real=1.0e6, pc_star2::Real=1.0e6, c::Real=1e-3)
  dpc1=dpc_drain3(sw, pce, swc, labda=labda1, pc_star=pc_star1, c=c)
  dpc2=dpc_drain3(1-sw, pce, sor, labda=labda2, pc_star=pc_star2, c=c)
  return (0.5*(1+cos(teta)))^b*dpc1+(0.5*(1-cos(teta)))^b*dpc2
end

"""
derivative of the imbibition capillary pressure curve w.r.t water saturation
pc_imb(sw::Real, pce::Real, swc::Real, sor::Real;
  teta::Real=0.785, labda::Real=2.4, b::Real=0.6, pc01::Real=1.0e6, pc02::Real=1.0e6)
"""
function d2pc_imb3(sw::Real, pce::Real, swc::Real, sor::Real;
  teta::Real=0.785, labda1::Real=2.4, labda2::Real=2.4, b::Real=1.0, pc_star1::Real=1.0e6, pc_star2::Real=1.0e6, c::Real=1e-3)
  dpc1=d2pc_drain3(sw, pce, swc, labda=labda1, pc_star=pc_star1, c=c)
  dpc2=d2pc_drain3(1-sw, pce, sor, labda=labda2, pc_star=pc_star2, c=c)
  return (0.5*(1+cos(teta)))^b*dpc1+(0.5*(1-cos(teta)))^b*dpc2
end


using ImageFiltering
using Statistics
using StatsBase

function badpixremoval_clump(frame::AbstractMatrix{T}; max_niter=15, kwargs...) where T
    # create bad pixel map by sigma clipping outliers
    bad_pixels = badpixel_clip(frame; kwargs...)
    bad_pixels_cumul = copy(bad_pixels)
    total = sum(bad_pixels)

    # iterate replacing bad pixels with sigma clipping
    out = copy(frame)
    it = 0
    while total > 0 && it < max_niter
        it += 1
        out = sigma_filter!(out, bad_pixels)
        bad_pixels = badpixel_clip(out)
        bad_pixels_cumul .|= bad_pixels
        total = sum(bad_pixels)
    end
    return out, bad_pixels_cumul
end

function badpixel_clip(frame; sigma=4, box_size=9)
    med = mapwindow(median, frame, (box_size, box_size); border="reflect")
    std = mapwindow(mad, frame, (box_size, box_size); border="reflect")

    return @. abs(frame - med) > sigma * std
end

function sigma_filter!(frame, bad_pixels; box_size=9)
    bpm = copy(bad_pixels)

    min_neigh = sum(3:2:box_size)
    half_box = box_size ÷ 2
    while sum(bpm) > 0
        pad_bpm = BorderArray(bpm, Pad(:reflect, half_box, half_box))
        pad_im = BorderArray(frame, Pad(:reflect, half_box, half_box))
        for idx in findall(bpm)
            ax1 = idx.I[1] - half_box:idx.I[1] + half_box
            ax2 = idx.I[2] - half_box:idx.I[2] + half_box
            block_bpm = @view pad_bpm[ax1, ax2]

            # ensure enough "good" neighbors for imputation
            if sum(block_bpm) < min_neigh
                block_im = @view pad_im[ax1, ax2]
                frame[idx] = median(block_im)
                bpm[idx] = false
            end
        end
    end
    return frame
end


# modified from https://github.com/mohamed82008/LinearElasticity.jl
using Einsum: @einsum
using LinearAlgebra: I

function get_Kσs(problem::TrussProblem{xdim, TT}, u_dofs, cellvalues) where {xdim, TT}
    Es = getE(problem)
    νs = getν(problem)
    As = getA(problem)
    dh = problem.ch.dh

    # usually ndof_pc = xdim * n_basefuncs
    # ? number of nodes per cell == n_basefuncs per cell
    ndof_pc = ndofs_per_cell(dh)
    n_basefuncs = getnbasefunctions(cellvalues)
    global_dofs = zeros(Int, ndof_pc)
    Kσs = [zeros(TT, ndof_pc, ndof_pc) for i in 1:getncells(dh.grid)]
    Kσ_e = zeros(TT, ndof_pc, ndof_pc)
    # block-diagonal - block σ_e = σ_ij, i,j in xdim
    ψ_e = zeros(TT, xdim*xdim, xdim*xdim)
    G = zeros(TT, xdim*xdim, xdim*n_basefuncs)
    δ = Matrix(TT(1.0)I, xdim, xdim)
    ϵ = zeros(TT, xdim, xdim)
    σ = zeros(TT, xdim, xdim)
    # u_i,j: partial derivative
    u_p = zeros(TT, xdim, xdim)

    for (cellidx, cell) in enumerate(CellIterator(dh))
        Kσ_e .= 0
        truss_reinit!(cellvalues, cell, As[cellidx])
        # get cell's dof's global dof indices, i.e. CC_a^e
        celldofs!(global_dofs, dh, cellidx)
        E = Es[cellidx]
        ν = νs[cellidx]
        for q_point in 1:getnquadpoints(cellvalues)
            dΩ = getdetJdV(cellvalues, q_point)
            for d in 1:xdim
                ψ_e[(d-1)*xdim+1:d*xdim, (d-1)*xdim+1:d*xdim] .= 0
            end
            for a in 1:n_basefuncs
                ∇ϕ = shape_gradient(cellvalues, q_point, a)
                # given displacement values of the cell nodes
                _u = @view u_dofs[(@view global_dofs[xdim*(a-1) .+ (1:xdim)])]
                # u_i,j, i for spatial xdim, j for partial derivative
                @einsum u_p[i,j] = _u[i]*∇ϕ[j]
                # effect of the quadratic term in the strain formula have on the stress field is ignored
                @einsum ϵ[i,j] = 1/2*(u_p[i,j] + u_p[j,i])
                # isotropic solid
                # @einsum σ[i,j] = E*ν/(1-ν^2)*δ[i,j]*ϵ[k,k] + E*ν*(1+ν)*ϵ[i,j]
                # ! truss element special treatment here
                σ = E .* ϵ

                for d in 1:xdim
                    ψ_e[(d-1)*xdim+1:d*xdim, (d-1)*xdim+1:d*xdim] .+= σ
                    G[(xdim*(d-1)+1):(xdim*d), (a-1)*xdim+d] .= ∇ϕ
                end
            end
            # @show G
            Kσ_e .+= G'*ψ_e*G*dΩ
        end
        # @show Kσ_e
        Kσs[cellidx] .= Kσ_e
    end
    return Kσs
end

function buckling(problem::TrussProblem{xdim, T}, ginfo, einfo; u=undef) where {xdim, T}
    dh = problem.ch.dh

    if u === undef
        u = ginfo.K \ ginfo.f
    end
    Kσs = get_Kσs(problem, u, einfo.cellvalues)
    Kσ = deepcopy(ginfo.K)

    # @show Kσs

    if Kσ isa Symmetric
        Kσ.data.nzval .= 0
        assembler = Ferrite.AssemblerSparsityPattern(Kσ.data, T[], Int[], Int[])
    else
        Kσ.nzval .= 0
        assembler = Ferrite.AssemblerSparsityPattern(Kσ, T[], Int[], Int[])
    end

    # * assemble global geometric stiffness matrix
    global_dofs = zeros(Int, ndofs_per_cell(dh))
    Kσ_e = zeros(T, size(Kσs[1]))
    celliteratortype = CellIterator{typeof(dh).parameters...}
    _celliterator::celliteratortype = CellIterator(dh)
    TK = eltype(Kσs)
    for (i,cell) in enumerate(_celliterator)
        celldofs!(global_dofs, dh, i)
        if TK <: Symmetric
            Ferrite.assemble!(assembler, global_dofs, Kσs[i].data)
        else
            Ferrite.assemble!(assembler, global_dofs, Kσs[i])
        end
    end

    # f = copy(ginfo.f)
    # apply!(Kσ, f, problem.ch)

    return ginfo.K, Kσ
end


using SpecialFunctions
gamma(1/2) - sqrt(pi)


@testset "CANDECOMP" begin
    Random.seed!(12345)

    r = 2 # core tensor dimensions
    T = rand_kruskal3(r, (10, 20, 30), false)
    initial_guess = ntuple(k -> randn(size(T, k), r), ndims(T))

    @testset "ALS (Alternating least squares)" begin
        @time factors = candecomp(T, r, initial_guess, tol=1e-6, compute_error=true, method=:ALS)
        @test length(factors.factors) == ndims(T)
        @test size.(factors.factors) == ntuple(i -> (size(T, i), r), ndims(T))
        @test rank(factors) == r
        @test rel_residue(factors) < 1e-5
    end

    @testset "SGSD (Simultaneous generalized Schur decomposition)" begin
        @time factors = candecomp(T, r, initial_guess, compute_error=true, method=:SGSD)
        @test length(factors.factors) == ndims(T)
        @test size.(factors.factors) == ntuple(i -> (size(T, i), r), ndims(T))
        @test rank(factors) == r
        @test length(factors.lambdas) == r
        @test_broken rel_residue(factors) < 1e-5
    end

end


using CRlibm
using Compat

using Base.Test

@test cos(0.5, RoundDown) == 0.8775825618903726
@test cos(0.5, RoundUp) == 0.8775825618903728
@test cos(0.5, RoundNearest) == cos(0.5) == 0.8775825618903728
@test cos(1.6, RoundToZero) == -0.029199522301288812
@test cos(1.6, RoundDown) == -0.029199522301288815

function my_eps(prec::Int)
    ldexp(eps(Float64), 53-prec)
end

function my_eps(x::BigFloat)  # only works for precision >= 10?
    my_eps(precision(x) - exponent(x))
end

is_log(f) = string(f)[1:3] == "log"

function do_test(f, val)
    a = f(val, RoundDown)
    b = f(val, RoundUp)
    @test b - a == eps(a) || b - a == eps(b) || b - a == 0
end

function test_CRlibm(function_list)
    @show function_list
    for f in function_list
        #println("Testing CRlibm.$f")

        ff = eval(f)  # the actual Julia function

        for val in (0.51, 103.2, -17.1, -0.00005, 1)

            #@show f, val

            val <= 0.0 && is_log(f) && continue
            abs(val) > 1 && f ∈ (:asin, :acos) && continue

            do_test(ff, val)
        end
    end
end

function test_MPFR()
    for f in CRlibm.MPFR_function_list
        #println("Testing CRlibm.$f")

        ff = eval(f)  # the actual Julia function

        for val in (0.51, 103.2, -17.1, -0.00005)
            #print(val, " ")

            #@show f, val

            val <= 0.0 && is_log(f) && continue
            abs(val) > 1 && f ∈ (:asin, :acos) && continue

            for prec in (20, 100, 1000)

                setprecision(BigFloat, prec) do
                    val = BigFloat(val)
                    do_test(ff, val)
                end

            end
        end
    end
end

println("Testing CRlibm")
test_CRlibm(CRlibm.function_list)
# This will currently fail on the :sinpi etc. functions (that are not defined in MPFR) if MPFR is already enabled because the CRlibm library could not be found

println("Testing shadowing MPFR")
CRlibm.shadow_MPFR()
test_CRlibm(CRlibm.MPFR_function_list)

println("Testing MPFR")
test_MPFR()


"""

Compute median, mad_sd, mean and std summary of distribition.

$(SIGNATURES)

## Positional arguments
```julia
* `model_df::DataFrame` # DataFrame holding the draws
* `pars::Vector{Symbol}` # Vector of symbols to be included
```

## Keyword arguments
```julia
* `digits = 2` # Number of decimal digits
```

"""
function model_summary(model_df::DataFrame,
    params::T = Symbol.(names(model_df));
    digits = 2) where {T}

    colnames = Symbol.(names(model_df))

    pars = Symbol[]
    for par in Symbol.(params)
        if par in colnames
            append!(pars, [par])
        else
            @warn ":$(par) not in $(colnames), will be dropped."
        end
    end

    if length(pars) > 0
        parameters = Pair{Symbol, Int}[]
        estimates = zeros(length(pars), 4)
        for (indx, par) in enumerate(pars)
            if par in colnames
                append!(parameters, [par => indx])
                vals = model_df[:, par]
                estimates[indx, :] = [median(vals), mad(vals), mean(vals), std(vals)]
            end
        end

        return NamedArray(
            round.(estimates; digits=digits),
            (OrderedDict(parameters...),
            OrderedDict(:median=>1, :mad_sd=>2, :mean=>3, :std=>4)),
                   ("Parameter", "Value")
        )
    else
        @warn "No parameters match the column names in model_df."
        return nothing
    end
end

export
    model_summary


using Bootstrap
using CSV
using CategoricalArrays
using Chain
using DataFrames
using GLM
using Random
using Plots
using Statistics
using StatsPlots
using Downloads

ENV["LINES"] = 20
ENV["COLUMNS"] = 1000

# -----------------------------------------------------------------------------------------
# 0: Loading the Data

Downloads.download("https://vincentarelbundock.github.io/Rdatasets/csv/Ecdat/Participation.csv", "participation.csv")

readlines("participation.csv")

# GOAL: create a model to predict :lfp (labor force participation)
# begin by reading the data into a dataframe
df_raw = CSV.read("participation.csv", DataFrame)  # source file, destination types... one could use other tabular formats (maybe JuliaDB?)

# inspect the contents
describe(df_raw)

# use the select() function to transform the dataframe
df = select(df_raw,
            :lfp => (x -> recode(x, "yes"=> 1, "no"=> 0)) => :lfp,
            :lnnlinc,
            :age,
            :age => ByRow(x-> x^2) => :age²,  # perform squaring operation row-wise (we don't want to square a vector)
            Between(:educ, :noc),  # select all columns between :educ and :noc
            :foreign => categorical => :foreign  # convert :foreign to :categorical data
            )

# notice, now we see that the :foreign column has categorical type.
# note, much of this can be simplified using MLJ's "scientific type" system
describe(df)

# -----------------------------------------------------------------------------------------
# 1: Exploratory Data Analysis
# we use the @chain macro for efficient piping
# groupby function

@chain df beginj
    groupby(:lfp) # in other words, sort the dataframe based on the values in :lfp column
    combine([:lnnlinc, :age, :educ, :nyc, :noc] .=> mean) # now compute mean of each column by grouping
end

# more examples here: https://bkamins.github.io/julialang/2021/07/09/multicol.html

# instead of specifying columns manually, just apply the operation to all columns of relevant type
@chain df begin
    groupby(:lfp)
    combine(names(df, Real) .=> mean)
end

# see number of rows corresponding to :lfp and :foreign column
@chain df begin
    groupby([:lfp, :foreign])
    combine(nrow) # i.e. combine result of nrow function
end

# more examples of reshaping: https://bkamins.github.io/julialang/2021/05/28/pivot.html
# ex: cross tabulate number of rows between :lfp variables and :foreign variables
@chain df begin
    groupby([:lfp, :foreign])
    combine(nrow)
    unstack(:lfp, :foreign, :nrow) # this does the tabulation
end

# grouped dataframes are their own type
gd = groupby(df, :lfp)
# access each group by standard indexing
size(gd)
gd[1]

# or index by value
gd[(lfp=0,)]

# plot the pdf of the age column for each group
@df df density(:age, group=:lfp)