using LinearAlgebra
using UnPack, StaticArrays, IntervalSets

import ClimaCore: Fields, Domains, Meshes, Topologies, Spaces
import ClimaCore: slab
import ClimaCore.Operators
using ClimaCore.Geometry
import ClimaCore.Geometry: Abstract2DPoint

using ClimaCore.RecursiveApply

const parameters = (
    ϵ = 0.1,  # perturbation size for initial condition
    l = 0.5, # Gaussian width
    k = 0.5, # Sinusoidal wavenumber
    ρ₀ = 1.0, # reference density
    c = 2,
    g = 10,
)

domain = Domains.RectangleDomain(
    Geometry.XPoint(-2π) .. Geometry.XPoint(2π),
    Geometry.YPoint(-2π) .. Geometry.YPoint(2π),
    x1periodic = true,
    x2periodic = true,
)

function init_state(coord, p)
    @unpack x, y = coord
    # set initial state
    ρ = p.ρ₀

    # set initial velocity
    U₁ = cosh(y)^(-2)

    # Ψ′ = exp(-(y + p.l / 10)^2 / 2p.l^2) * cos(p.k * x) * cos(p.k * y)
    # Vortical velocity fields (u₁′, u₂′) = (-∂²Ψ′, ∂¹Ψ′)
    gaussian = exp(-(y + p.l / 10)^2 / 2p.l^2)
    u₁′ = gaussian * (y + p.l / 10) / p.l^2 * cos(p.k * x) * cos(p.k * y)
    u₁′ += p.k * gaussian * cos(p.k * x) * sin(p.k * y)
    u₂′ = -p.k * gaussian * sin(p.k * x) * cos(p.k * y)

    u = Geometry.UVVector(U₁ + p.ϵ * u₁′, p.ϵ * u₂′)
    # set initial tracer
    θ = sin(p.k * y)

    return (ρ = ρ, ρu = ρ * u, ρθ = ρ * θ)
end

function flux(state, p)
    @unpack ρ, ρu, ρθ = state
    u = ρu / ρ
    return (ρ = ρu, ρu = ((ρu ⊗ u) + (p.g * ρ^2 / 2) * I), ρθ = ρθ * u)
end

roe_average(ρ⁻, ρ⁺, var⁻, var⁺) =
    (sqrt(ρ⁻) * var⁻ + sqrt(ρ⁺) * var⁺) / (sqrt(ρ⁻) + sqrt(ρ⁺))

function roeflux(n, (y⁻, parameters⁻), (y⁺, parameters⁺))
    Favg = RecursiveApply.rdiv(flux(y⁻, parameters⁻) ⊞ flux(y⁺, parameters⁺), 2)

    λ = sqrt(parameters⁻.g)

    ρ⁻, ρu⁻, ρθ⁻ = y⁻.ρ, y⁻.ρu, y⁻.ρθ
    ρ⁺, ρu⁺, ρθ⁺ = y⁺.ρ, y⁺.ρu, y⁺.ρθ

    u⁻ = ρu⁻ / ρ⁻
    θ⁻ = ρθ⁻ / ρ⁻
    uₙ⁻ = u⁻' * n

    u⁺ = ρu⁺ / ρ⁺
    θ⁺ = ρθ⁺ / ρ⁺
    uₙ⁺ = u⁺' * n

    # in general thermodynamics, (pressure, soundspeed)
    p⁻ = (λ * ρ⁻)^2 * 0.5
    c⁻ = λ * sqrt(ρ⁻)

    p⁺ = (λ * ρ⁺)^2 * 0.5
    c⁺ = λ * sqrt(ρ⁺)

    # construct roe averges
    ρ = sqrt(ρ⁻ * ρ⁺)
    u = roe_average(ρ⁻, ρ⁺, u⁻, u⁺)
    θ = roe_average(ρ⁻, ρ⁺, θ⁻, θ⁺)
    c = roe_average(ρ⁻, ρ⁺, c⁻, c⁺)

    # construct normal velocity
    uₙ = u' * n

    # differences
    Δρ = ρ⁺ - ρ⁻
    Δp = p⁺ - p⁻
    Δu = u⁺ - u⁻
    Δρθ = ρθ⁺ - ρθ⁻
    Δuₙ = Δu' * n

    # constructed values
    c⁻² = 1 / c^2
    w1 = abs(uₙ - c) * (Δp - ρ * c * Δuₙ) * 0.5 * c⁻²
    w2 = abs(uₙ + c) * (Δp + ρ * c * Δuₙ) * 0.5 * c⁻²
    w3 = abs(uₙ) * (Δρ - Δp * c⁻²)
    w4 = abs(uₙ) * ρ
    w5 = abs(uₙ) * (Δρθ - θ * Δp * c⁻²)

    # fluxes!!!

    fluxᵀn_ρ = (w1 + w2 + w3) * 0.5
    fluxᵀn_ρu =
        (w1 * (u - c * n) + w2 * (u + c * n) + w3 * u + w4 * (Δu - Δuₙ * n)) *
        0.5
    fluxᵀn_ρθ = ((w1 + w2) * θ + w5) * 0.5

    Δf = (ρ = -fluxᵀn_ρ, ρu = -fluxᵀn_ρu, ρθ = -fluxᵀn_ρθ)
    RecursiveApply.rmap(f -> f' * n, Favg) ⊞ Δf
end

function volume!(dydt, y, (parameters,), t)
    div = Operators.WeakDivergence()
    rparameters = Ref(parameters)
    @. dydt = -(div(flux(y, rparameters)))
    return dydt
end

function add_face!(dydt, y, (parameters,), t)
    Operators.add_numerical_flux_internal!(roeflux, dydt, y, parameters)
    return dydt
end


__precompile__()

module Materials

using LinearAlgebra
using ..adiff

# Material types
struct Hooke
  E     ::Float64
  ν     ::Float64
  small ::Bool
  Hooke(E,ν;small=false) = new(Float64(E),Float64(ν), small)
end
struct MooneyRivlin
  C1  ::Float64
  C2  ::Float64
  K   ::Float64
  MooneyRivlin(C1, C2)    = new(Float64(C1), Float64(C2), NaN)
  MooneyRivlin(C1, C2, K) = new(Float64(C1), Float64(C2), Float64(K))
end
struct NeoHooke
  C1   ::Float64
  K    ::Float64
  NeoHooke(C1)    = new(Float64(C1), NaN)
  NeoHooke(C1, K) = new(Float64(C1), Float64(K))
end
Material = Union{Hooke,MooneyRivlin,NeoHooke}
HyperEla = Union{MooneyRivlin,NeoHooke}
# default convergence tolerance for 2D stress
dTol     = 1e-7
maxiter  = 30
function setmaxiter(x)
  global maxiter = Int64(x)
end
function setdTol(x)
  global dTol = Float64(x)
end
# elastic energy evaluation functions for materials
function getϕ(F::Array{N,2}, mat::M) where {N<:Number, M<:HyperEla}
  C = transpose(F)F
  if length(C) == 9
    (I1,I2,I3) = getInvariants(C)
  else
    L3 = isnan(mat.K) ? (F[1]F[4]-F[2]F[3])^-2 : 1
    (I1,I2,I3) = getInvariants(C, L3)
  end
  Materials.getϕ(I1,I2,I3,mat)
end
function getϕ(I1, I2, I3, mat::MooneyRivlin)

  C1, C2, K = mat.C1, mat.C2, mat.K
  if isnan(K)
    ϕ  = C1*(I1-3) + C2*(I2-3)
  else
    J  = sqrt(I3)
    I1 = I1*J^(-2/3)
    I2 = I2*J^(-4/3)
    # ϕ  = C1*(I1-3) + C2*(I2-3) + K*log(J)^2
    ϕ  = C1*(I1-3) + C2*(I2-3) + K*(J-1)^2
  end
  return ϕ
end
function getϕ(I1, I2, I3, mat::NeoHooke)

  C1, K  = mat.C1, mat.K
  if isnan(K)
    ϕ  = C1*(I1-3)
  else
    J  = sqrt(I3)
    I1 = I1*J^(-2/3)
    # ϕ  = C1*(I1-3) + K*log(J)^2
    ϕ  = C1*(I1-3) + K*(J-1)^2
  end
  return ϕ
end
function getϕ(C11::N where N<:Number, mat::Hooke)
  ϕ = 0.5mat.E*(C11-1)^2
end
function getϕ(F::Array{N,2} where N<:Number, mat::Hooke)

  if mat.small
    E = 0.5(F+transpose(F))-I   # the symmetric part of G
  else
    E = 0.5(transpose(F)F-I)    # the Green-Lagrange strain tensor
  end

  ν, Es = mat.ν, mat.E
  λ = Es*ν/(1+ν)/(1-2ν)
  μ = Es/2/(1+ν)

  if size(F) == (3,3)
    ϕ =  (μ+λ/2) * (E[1]^2   + E[5]^2   + E[9]^2)
    ϕ += λ       * (E[1]E[5] + E[5]E[9] + E[9]E[1])
    ϕ += 2μ      * (E[2]^2   + E[3]^2   + E[6]^2)
  else
    # 2D is plain strain
    ϕ = (μ+λ/2)*(E[1]^2+E[4]^2) + λ*E[1]E[4] + 2μ*E[2]^2
  end

  return ϕ
end
# status retrieving functions
function getP(F::Array{Float64,2}, mat) # 1st PK tensor from F
  ϕ = getϕ(adiff.D2(F), mat)
  reshape(adiff.grad(ϕ), (3,3))
end
function getσ(F::Array{Float64,2}, mat) # Cauchy stress
  J = det(F)
  P = getP(F,mat)         # 1st PK stress
  (P*transpose(F))/J
end
function getinfo(F::Array{Float64,2}, mat; info = :ϕ)
  if info == :ϕ
    getϕ(F, mat)
  elseif info == :detF
    det(F)
  elseif info == :G
    F-I
  elseif info == :P
    getP(F, mat)
  elseif info == :σ
    getσ(F, mat)
  elseif info == :σVM
    σ = getσ(F, mat)
    sqrt(0.5)*sqrt((σ[1]-σ[5])^2+(σ[5]-σ[9])^2+(σ[9]-σ[1])^2+6(σ[4]^2+σ[7]^2+σ[8]^2))
  elseif info == :S
    F\getP(F, mat)
  elseif info == :I
    getInvariants(transpose(F)F)
  elseif info == :LE
    svdF = svd(F)
    svdF.V*diagm(0=>log.(svdF.S))*transpose(svdF.V)
  else
    F
  end
end
# helper functions for 2D and 1D stress
function getϕ(F11::N, mat::M; binc=true) where {N<:Number, M<:HyperEla}

  C11 = F11^2
  L1  = Real(C11)
  L3  = binc ? getL3(L3->getϕ(L1,L3,L3,mat), L1) :
  getL3(L3->getϕ(L1,L3,L3,mat), √L1)

  if isa(C11, adiff.D2)
    ϕ1    = getϕ(adiff.D2([L1, L3, L3])..., mat)
    dL21  = -ϕ1.h[2]/(ϕ1.h[3]+ϕ1.h[5])
    L2    = adiff.D2(L3, dL21*C11.g)
    ϕ     = getϕ(C11,L2,L2,mat)
  else
    ϕ  = getϕ(C11,L3,L3,mat)
  end

  return ϕ
end
function getL3(func, L3)
  iter   = 0
  bdone  = false
  while !bdone
    ϕ = func(adiff.D2(L3))
    r = adiff.grad(ϕ)[1]
    if abs(r) < dTol
      bdone = true
    else
      updt = -r/adiff.hess(ϕ)[1]
      L3   = L3+updt<0 ? L3/2 : L3+updt
    end
    iter +=1
    if iter > maxiter
      println("failed in getL3, with L3: ",L3," r: ", r, " and iter: ",iter)
      error("failed in getL3")
      bdone = true
    end
  end
  return  L3
end
function getInvariants(C, C33)
  I1 = C[1]+C[4]+C33
  I2 = C[1]C[4]+C[4]C33+C[1]C33-C[2]^2
  I3 = C[1]C[4]C33-C33*C[2]^2

  (I1,I2,I3)
end
function getInvariants(C)
  I1 = C[1]+C[5]+C[9]
  I2 = C[1]C[5]+C[5]C[9]+C[1]C[9]-C[2]^2-C[3]^2-C[6]^2
  I3 = C[1]C[5]C[9]+2C[2]C[3]C[6]-C[1]C[6]^2-C[5]C[3]^2-C[9]C[2]^2

  (I1,I2,I3)
end
end


using PastaQ
using ITensors
using Printf

# 1. Prepation of a thermal state
#
# In this example, we show how to prepare the finite-temperature state
# of a many-body system:
#
# ρ̂(β) = exp(-β Ĥ)
#
# where Ĥ is the Hamiltonian and β is the inverse temperature.
# We specificallty consider the one-dimensional Ising model
#
#   H = - ∑ᵢ σᶻ(i) σᶻ(i+1) - B ∑ᵢ σˣ(i)
#
# where B a the transverse magnetic field.

# 1a. Custom gates
#
# In order to build the thermal density operator, we implement the
# simplest flavor of imaginary-time evolution, breaking the operator
# exp(-βĤ) into a set of two-qubit and single-qubit gates, corresponding
# to the Ising interactions and the transverse field respetively. The
# time evolution to inverse temperature β is broken into elementary steps
# of size τ, where a gate is applied for each term appearing in the Hamiltonian.
#
# In this example, the quantum gates are not contained in the gate set of PastaQ.
# In order to extend, it is ony required to define the gate matrices using a
# format analogous to standard gates defined in gates.jl.

import PastaQ: gate

gate(::GateName"expτZZ"; τ::Float64) = exp(τ * kron(gate("Z"), gate("Z")))

gate(::GateName"expτX"; τ::Float64, B::Float64) = exp(τ * B * gate("X"))

# 1b. Generating the thermal state

N = 10    # Number of spins
B = 1.0   # Transverse magnetic field
β = 1.0   # Inverse temperature
τ = 0.005 # Trotter step

# Depth of the circuit
depth = β ÷ τ

# Ising interactions
zz_layer = [("expτZZ", (j, j + 1), (τ=τ,)) for j in 1:(N - 1)]
# Transverse field
x_layer = [("expτX", j, (τ=τ, B=B)) for j in 1:N]

# Build the gate structure
circuit = []
for d in 1:depth
  append!(circuit, zz_layer)
  append!(circuit, x_layer)
end

#
# 2. Run imaginary-time evolution towards the zero temperature
# ground state.
#

# 2a. Ground state energy with DMRG
#
# We compute the ground state energy by running DMRG
# on the Hamiltonian MPO, whose algorithm is implemented in
# ITensors.jl.

# In order to generate the MPO for the Hamiltonian, we leverage
# the ITensors.jl `AutoMPO()` function, which automatically
# generates the local MPO tensors from a set of pre-definend operators..
sites = siteinds("Qubit", N)
ampo = AutoMPO()
for j in 1:(N - 1)
  # Ising ZZ interactions
  ampo .+= -1, "Z", j, "Z", j + 1
end
for j in 1:N
  # Transverse field X
  ampo .+= -B, "X", j
end
# Generate Hamilotnian MPO
H = MPO(ampo, sites)

# Density-matrix renormalization group
dmrg_iter = 5      # DMRG steps
dmrg_cutoff = 1E-10   # Cutoff
Ψ0 = randomMPS(sites) # Initial state
sweeps = Sweeps(dmrg_iter)
maxdim!(sweeps, 10, 20, 30, 40, 50, 100)
cutoff!(sweeps, dmrg_cutoff)
# Run
println("Running DMRG to get ground state of transverse field Ising model:")
E, Ψ = dmrg(H, Ψ0, sweeps)
@printf("\nGround state energy:  %.8f  \n", E)
println("\n---------------------------------------\n")

#
# 2b. Run the imaginary-time circuit
#

β = 5.0 # Inverse temperature
Δ = 0.5  # Intermediate time-step
depth = Δ ÷ τ # Depth of the circuit
steps = β ÷ Δ # Total number of circuit application

# Initialize the density operator
ρ = PastaQ.identity_mpo(H)

println("Running imaginary time evolution to approximate the density matrix ρ = exp(-βH):")
for b in 1:steps
  # Run the circuit
  global ρ = runcircuit(ρ, circuit; cutoff=1E-12)

  # Normalize the density operatorr
  normalize!(ρ)

  # Measure the energy
  E_th = inner(ρ, H)

  @printf("β = %.1f : tr(ρH) = %.8f \n", (Δ * b), E_th)
end


struct Mesh
    vertices::Array{Point{3, Float32}}
    edges::Set{Tuple{Int, Int}}
    faces::Array{Face{3,Any}}
end

"""
    distance(v1::Point, v2::Point)

Calculate the distance between two points

"""
function distance(v1::Point, v2::Point)
    return sqrt(reduce(+, map((x,y)->(y-x)^2, v1, v2)))
end

function select_valid_pairs(mesh::Mesh, threshold::Real)
    rs = Set{Tuple{Int, Int}}()
    for i in 1:length(mesh.vertices), e in mesh.edges
        if i in e
            push!(rs, e)
        end
    end
    for i in 1:length(mesh.vertices), j in 1:length(mesh.vertices)
        if i != j && distance(mesh.vertices[i], mesh.vertices[j]) <= threshold
            if !((j,i) in rs)
                push!(rs, (i, j))
            end
        end
    end
    return rs
end

function calculate_contraction_position(v1::Point, v2::Point)
    return Point{3}((v1.+v2)./2)
end

function calculate_contraction_cost(mesh::HomogenousMesh, v1::Point, v2::Point)
    v_bar = contract_vertices(v1, v2)
    cost = Q(mesh, v1)+Q(mesh, v2)
    return cost
end

function get_edges_from_face(face::Face)
    return (face[1], face[2]), (face[1], face[3]), (face[2], face[3])
end

function get_faces_from_vertex(mesh::HomogenousMesh, v_ind::Int)
    return [f for f in mesh.faces if v_ind in f]
end

function Q(mesh::HomogenousMesh, v_ind::Int)
    neighboring_faces = get_faces_from_vertex(mesh, v_ind)
    plane_reps = [get_plane_representation(mesh, f) for f in neighboring_faces]
    error_quadrics = [fundamental_error_quadric(p) for p in plane_reps]
    return reduce(+, error_quadrics)
end

function normalize(n::Vector)
    sq_sum = reduce(+, map(x->x^2, n))
    sq_sum == 0 && error("Can't normalize the 0-vector")
    return map(x->x/(sqrt(sq_sum)), n)
end

function get_plane_representation(mesh::HomogenousMesh, face::Face)
    points_in_plane = mesh.vertices[face]
    v1 = Vector(points_in_plane[2].-points_in_plane[1])
    v2 = Vector(points_in_plane[3].-points_in_plane[1])
    n = v1 × v2
    n = normalize(n)
    d = -points_in_plane[1][1]*n[1] - points_in_plane[1][2]*n[2] - points_in_plane[1][3]*n[3]
    append!(n,d)
    return SVector{4}(n)
end

function convert_from_homogenous(mesh::HomogenousMesh)
    edges =Set{Tuple{Int, Int}}()
    for f in mesh.faces
        edges = union(edges, get_edges_from_face(f))
    end
    return Mesh(mesh.vertices, edges, mesh.faces)
end

function fundamental_error_quadric(p::SVector{4,T}) where {T<:Real}
    return [p[1]^2 p[1]*p[2] p[1]*p[3] p[1]*p[4];
            p[1]*p[2] p[2]^2 p[2]*p[3] p[2]*p[4];
            p[1]*p[3] p[2]*p[3] p[3]^2 p[3]*p[4];
            p[1]*p[4] p[2]*p[4] p[3]*p[4] p[4]^2]
end


"""
Die Funktion divide_jl ist eine julia Implementierung von divide.py.
Diese Funktion arbeitet nicht mit .txt Dateien sondern mit arrays.

### Argumente
- `data::Array`: enthält die Messzeiten (erster Parameter)
- `Δt::Real`: Länge der Intervalle in s (zweiter Parameter)

### Rückgabewert
- `counts::Array`: enthält Anzahl der Signale für jedes Interval

### Beispiel
```julia
counts = divide_jl(data, 10)
```
"""
function divide_jl(data::Array, Δt::Real)
    counts = []
    for x in data
        i = round(Int, x/Δt, RoundDown)
        while length(counts) <= i
            append!(counts, 0)
        end
        counts[i+1] += 1
    end
    pop!(counts)
    return counts
end
"""
Die Funktion binomial_jl ist eine julia Implementierung von binomial.py.
Diese Funktion arbeitet nicht mit .txt Dateien sondern mit arrays.

### Argumente
- `data::Array`: enthält die Messzeiten (erster Parameter)
- `Δt::Real`: Länge der Intervalle in s (zweiter Parameter)
- `n::Int`: Binningparameter (dritter Parameter)

### Rückgabewert
- `x::Array`: enthält die x-Werte der Binomialverteilung
- `y::Array`: enthält die y-Werte der Binomialverteilung

### Beispiel

```julia
x, y = binomial_jl(data, 0.2, 1)
```
"""
function binomial_jl(data::Array, Δt::Real, n::Int)
    counts = divide_jl(data, Δt)
    hist1 = []
    for i in 1:length(counts)
        while length(hist1) < counts[i]+1
            append!(hist1, 0)
        end
        hist1[counts[i]+1] += 1
    end

    hist2 = []
    for i in 1:length(hist1)
        j = round(Int, (i-1)/n, RoundDown)
        while length(hist2) <= j
            append!(hist2, 0)
        end
        hist2[j+1] += hist1[i]
    end

    return [i for i in 0:length(hist2)-1], hist2
end
"""
Die Funktion interval_jl ist eine julia Implementierung von interval.py.
Diese Funktion arbeitet nicht mit .txt Dateien sondern mit arrays.

### Argumente
- `data::Array`: enthält die Messzeiten (erster Parameter)
- `Δt::Real`: Länge der Intervalle in s (zweiter Parameter)
- `n::Int`: step, gibt an ob Zeit bis zum nächsten, übernächsten, ... betrachtet wird (dritter Parameter)

### Rückgabewert
- `x::Array`: enthält die x-Werte der Intervallverteilung
- `y::Array`: enthält die y-Werte der Intervallverteilung

### Beispiel

```julia
x, y = interval_jl(data, 0.001, 1)
```
"""
function interval_jl(data::Array, Δt::Real, n::Int)
    hist = []
    for i in 1:length(data)-n
        t = data[i+n]-data[i]
        k = round(Int, t/Δt, RoundDown)
        while length(hist) <= k
            append!(hist, 0)
        end
        hist[k+1] += 1
    end
    return [i for i in 0:length(hist)-1], hist
end


"""
Construct vectorfields and derivatives from bases and coefficients.
"""

function construct_vectorfield(coeff, B)
    @assert length(coeff) == length(B)
    a = copy(coeff)
    function vectorfield(x)
        out = [0., 0.]
        for i in 1:length(B)
            out += a[i]*B[i](x)
        end
        return out
    end
    return vectorfield
end

function construct_jacobian_matrix(coeff, DB)
    @assert length(coeff) == length(DB)
    a = copy(coeff)
    function jacobian_matrix(x)
        out = [ [0. 0.]; [0. 0.]]
        for i in 1:length(DB)
            out += a[i]*DB[i](x)
        end
        return out
    end
    return jacobian_matrix
end

function jacobian_determinant(Df)
    function J(x)
        return det(Df(x))
    end
end

function gradient_jacobian_determinant(Df)
    J = jacobian_determinant(Df)
    function DJ(x)
        return ForwardDiff.gradient(J, x)
    end
    return DJ
end


module Artificial

using CUDA
using LinearAlgebra
using MatrixOptInterface
using SparseArrays

import ..PhaseOne
import Simplex
import Simplex: nrows, ncols

const MatOI = MatrixOptInterface
const MOI = MatOI.MOI

"""
    This creates a reformulation to begin with all artificial basis.
    This is the simplest way to construct the initial basis.

    The canonical form

    minimize c x
    subject to
        A x == b
        xl <= x <= xu

    is reformulated to

    minimize a1 + a2
    subject to
        A1 x - a1 == b - A1 xN
        A2 x + a2 == b - A2 xN
        xl <= x <= xu
        0 <= a1 <= Inf
        0 <= a2 <= Inf,

    where
        if xl[j] > -Inf
            xN[j] = xl[j]
        elseif xu[j] < Inf
            xN[j] = xu[j]
        else
            xN[j] = 0.0
        end
"""
function reformulate(lp::MatOI.LPSolverForm{T, AT, VT}; basis::Array{Int,1} = Int[]) where {T, AT, VT}

    artif_idx = Int[]
    if length(basis) > 0
        # Add the artifical variables that are identified as basis
        for j in 1:length(basis)
            # if basis indicates an artificial variable,
            if basis[j] > ncols(lp)
                if basis[j] - ncols(lp) > nrows(lp)
                    @error "Invalid basis provided for phase one"
                end
                push!(artif_idx, basis[j])
                basis[j] = ncols(lp) + length(artif_idx)
            end
        end
    else
        # if no basis is provided, return (m+1,...,m+n)
        artif_idx = collect(1:nrows(lp)) .+ ncols(lp)
    end

    nartif = length(artif_idx)
    lp_nrows = nrows(lp)
    lp_ncols = ncols(lp) + nartif

    # objective coefficient
    c = append!(zeros(ncols(lp)), ones(nartif))
    c = VT(c)

    # column bounds
    # TODO: build thing directly on GPU
    xl = append!(deepcopy(Array(lp.v_lb)), zeros(nartif))
    xu = append!(deepcopy(Array(lp.v_ub)), fill(Inf, nartif))
    xl = VT(xl)
    xu = VT(xu)

    x = VT(undef, lp_ncols)

    for j = 1:lp_ncols
        if xl[j] > -Inf && abs(xl[j]) <= abs(xu[j])
            x[j] = xl[j]
        elseif xu[j] < Inf && abs(xl[j]) > abs(xu[j])
            x[j] = xu[j]
        else
            x[j] = 0.0
        end
    end

    # rhs
    b = VT(undef, nrows(lp))
    copyto!(b, lp.b .- lp.A * x[1:ncols(lp)])

    # create the submatrix for slack
    I = Array{Int64}(undef, nartif)
    J = collect(1:nartif)
    V = VT(undef, nartif)
    for i = 1:nartif
        I[i] = artif_idx[i] - ncols(lp)
        if b[I[i]] < 0
            V[i] = -1
            x[ncols(lp)+i] = -b[I[i]]
        else
            V[i] = 1
            x[ncols(lp)+i] = b[I[i]]
        end
    end
    S = sparse(I, J, V, lp_nrows, nartif)

    if VT <: Array
        # create the matrix in canonical form
        A = [sparse(Matrix(lp.A)) S]
    elseif VT <: CuArray
        # on GPU
        A = [lp.A Matrix(S)]
    else
        error("Not supported vector type $(VT)")
    end

    artif_col_idx = Int[]
    for j = 1:nartif
        push!(artif_col_idx, ncols(lp)+j)
    end

    senses = fill(MatOI.EQUAL_TO, nrows(lp))
    return MatOI.LPSolverForm{T, typeof(A), typeof(c)}(
        MOI.MIN_SENSE,
        c, A, lp.b, senses, xl, xu
    ), artif_col_idx
end

function run(prob::MatOI.LPSolverForm, Tv::Type; kwargs...)::Simplex.SpxData

    if !haskey(kwargs, :pivot_rule)
        @error "Argument :pivot_rule is not provided."
    end
    pivot_rule = kwargs[:pivot_rule]

    # convert to phase-one form
    p1lp, artif = reformulate(prob)

    # load the problem
    spx = Simplex.SpxData(p1lp, Tv)
    spx.pivot_rule = deepcopy(pivot_rule)

    # set basis
    basic = collect((ncols(p1lp)-nrows(p1lp)+1):ncols(p1lp))
    Simplex.set_basis(spx, basic)

    # Run simplex method
    Simplex.run_core(spx)

    for j in artif
        spx.lpdata.v_lb[j] = 0.0
        spx.lpdata.v_ub[j] = 0.0
    end

    return spx
end

end # module


using SparsityOperators

# test FourierOperators
function testFT(N=16)
  # random image
  x = zeros(ComplexF64,N,N)
  for i=1:N,j=1:N
    x[i,j] = rand()
  end

  # FourierMatrix
  idx = CartesianIndices((N,N))[collect(1:N^2)]
  F = [ exp(-2*pi*im*((idx[j][1]-1)*(idx[k][1]-1)+(idx[j][2]-1)*(idx[k][2]-1))/N) for j=1:N^2, k=1:N^2 ]
  F_adj = F'

  # Operators
  tr = CartesianTrajectory(Float64,N,N)
  F_nfft = NFFTOp((N,N),tr,symmetrize=false)
  F_exp = ExplicitOp((N,N),tr,zeros(ComplexF64,N,N),symmetrize=false)

  # test agains FourierOperators
  y = vec( ifftshift(reshape(F*vec(fftshift(x)),N,N)) )
  y_adj = vec( ifftshift(reshape(F_adj*vec(fftshift(x)),N,N)) )

  y_nfft = F_nfft*vec(x)
  y_adj_nfft = adjoint(F_nfft) * vec(x)

  y_exp = F_exp*vec(x)
  y_adj_exp = adjoint(F_exp) * vec(x)

  @test y     ≈ y_nfft      rtol = 1e-2
  @test y     ≈ y_exp       rtol = 1e-2
  @test y_adj ≈ y_adj_nfft  rtol = 1e-2
  @test y_adj ≈ y_adj_exp   rtol = 1e-2

  # test AHA w/o Toeplitz
  F_nfft.toeplitz = false
  AHA = SparsityOperators.normalOperator(F_nfft)
  y_AHA_nfft = AHA * vec(x)
  y_AHA = F' * F * vec(x)
  @test y_AHA ≈ y_AHA_nfft   rtol = 1e-2

  # test AHA with Toeplitz
  F_nfft.toeplitz = true
  AHA = SparsityOperators.normalOperator(F_nfft)
  y_AHA_nfft = AHA * vec(x)
  y_AHA_nfft = adjoint(F_nfft) * F_nfft * vec(x)
  y_AHA = F' * F * vec(x)
  @test y_AHA ≈ y_AHA_nfft   rtol = 1e-2

  # test type stability;
  # TODO: Ensure type stability for Trajectory objects and test here
  nodes = Float32.(tr.nodes)
  F_nfft = NFFTOp((N,N),nodes,symmetrize=false)

  y_nfft = F_nfft * vec(ComplexF32.(x))
  y_adj_nfft = adjoint(F_nfft) * vec(ComplexF32.(x))

  @test Complex{eltype(nodes)} === eltype(y_nfft)
  @test Complex{eltype(nodes)} === eltype(y_adj_nfft)
end

function testFT3d(N=12)
  # random image
  x = zeros(ComplexF64,N,N,N)
  for i=1:N,j=1:N,k=1:N
    x[i,j,k] = rand()
  end

  # FourierMatrix
  idx = CartesianIndices((N,N,N))[collect(1:N^3)]
  F = [ exp(-2*pi*im*((idx[j][1]-1)*(idx[k][1]-1)+(idx[j][2]-1)*(idx[k][2]-1)+(idx[j][3]-1)*(idx[k][3]-1))/N) for j=1:N^3, k=1:N^3 ]
  F_adj = F'

  # Operators
  tr = CartesianTrajectory3D(Float64,N,N,numSlices=N)
  F_nfft = NFFTOp((N,N,N),tr,symmetrize=false)
  F_exp = ExplicitOp((N,N,N),tr,zeros(ComplexF64,N,N,N),symmetrize=false)

  # test agains FourierOperators
  y = vec( ifftshift(reshape(F*vec(fftshift(x)),N,N,N)) )
  y_adj = vec( ifftshift(reshape(F_adj*vec(fftshift(x)),N,N,N)) )

  y_nfft = F_nfft*vec(x)
  y_adj_nfft = adjoint(F_nfft) * vec(x)

  y_exp = F_exp*vec(x)
  y_adj_exp = adjoint(F_exp) * vec(x)

  @test  y     ≈ y_exp      rtol = 1e-2
  @test  y     ≈ y_nfft     rtol = 1e-2
  @test  y_adj ≈ y_adj_exp  rtol = 1e-2
  @test  y_adj ≈ y_adj_nfft rtol = 1e-2

  # test AHA w/o Toeplitz
  F_nfft.toeplitz = false
  AHA = SparsityOperators.normalOperator(F_nfft)
  y_AHA_nfft = AHA * vec(x)
  y_AHA = F' * F * vec(x)
  @test y_AHA ≈ y_AHA_nfft   rtol = 1e-2

  # test AHA with Toeplitz
  F_nfft.toeplitz = true
  AHA = SparsityOperators.normalOperator(F_nfft)
  y_AHA_nfft = AHA * vec(x)
  y_AHA_nfft = adjoint(F_nfft) * F_nfft * vec(x)
  y_AHA = F' * F * vec(x)
  @test y_AHA ≈ y_AHA_nfft   rtol = 1e-2
end

## test FieldmapNFFTOp
function testFieldmapFT(N=16)
  # random image
  x = zeros(ComplexF64,N,N)
  for i=1:N,j=1:N
    x[i,j] = rand()
  end

  tr = CartesianTrajectory(Float64,N,N;TE=0.0,AQ=0.01)
  times = readoutTimes(tr)
  nodes = kspaceNodes(tr)
  cmap = im*quadraticFieldmap(N,N)[:,:,1]

  # FourierMatrix
  idx = CartesianIndices((N,N))[collect(1:N^2)]
  F = [exp(-2*1im*pi*(nodes[1,k]*(idx[l][1]-size(x,1)/2-1)+nodes[2,k]*(idx[l][2]-size(x,2)/2-1))-cmap[idx[l][1],idx[l][2]]*times[k]) for k=1:size(nodes,2), l=1:length(x)]
  F_adj = F'

  # Operators
  F_nfft = FieldmapNFFTOp((N,N),tr,cmap,symmetrize=false)

  # test agains FourierOperators
  y = F*vec(x)
  y_adj = F_adj*vec(x)

  y_nfft = F_nfft*vec(x)
  y_adj_nfft = adjoint(F_nfft) * vec(x)

  @test (norm(y-y_nfft)/norm(y)) < 1e-2
  @test (norm(y_adj-y_adj_nfft)/norm(y_adj)) < 1e-2
end

function testOperators()
  @testset "Linear Operator" begin
    testFT()
    testFT3d()
    testFieldmapFT()
  end
end

testOperators()


export coulomb

using OffsetArrays

#using StaticArrays # Surprisingly, didn't speed up at all
"""
    coulomb(aexpn,ax,ay,az,aI,aJ,aK,
        bexpn,bx,by,bz,bI,bJ,bK,
        cexpn,cx,cy,cz,cI,cJ,cK,
        dexpn,dx,dy,dz,dI,dJ,dK)

Compute the coulomb repulsion between four primitive Gaussian basis
functions defined by their exponents, xyz coordinates, and IJK powers.
"""
function coulomb(aexpn,ax,ay,az,aI,aJ,aK,
    bexpn,bx,by,bz,bI,bJ,bK,
    cexpn,cx,cy,cz,cI,cJ,cK,
    dexpn,dx,dy,dz,dI,dJ,dK)

    rab2 = dist2(ax-bx,ay-by,az-bz)
    rcd2 = dist2(cx-dx,cy-dy,cz-dz)

    px,py,pz = gaussian_product_center(aexpn,ax,ay,az,bexpn,bx,by,bz)
    qx,qy,qz = gaussian_product_center(cexpn,cx,cy,cz,dexpn,dx,dy,dz)
    rpq2 = dist2(px-qx,py-qy,pz-qz)

    g1 = aexpn+bexpn
    g2 = cexpn+dexpn
    delta = 0.25*(1/g1+1/g2)

    Bx = Barray(aI,bI,cI,dI,px,ax,bx,qx,cx,dx,g1,g2,delta)
    By = Barray(aJ,bJ,cJ,dJ,py,ay,by,qy,cy,dy,g1,g2,delta)
    Bz = Barray(aK,bK,cK,dK,pz,az,bz,qz,cz,dz,g1,g2,delta)

    x = 0.25*rpq2/delta

    s = 0
    for I in 0:(aI+bI+cI+dI)
        for J in 0:(aJ+bJ+cJ+dJ)
            for K in 0:(aK+bK+cK+dK)
                s += Bx[I]*By[J]*Bz[K]*Fgamma(I+J+K,x)
            end
        end
    end
    return 2*pi^(2.5)/(g1*g2*sqrt(g1+g2))*exp(-aexpn*bexpn*rab2/g1)*exp(-cexpn*dexpn*rcd2/g2)*s
end

"""
    coulomb(a::PGBF,b::PGBF,c::PGBF,d::PGBF)

Compute the coulomb repulsion between four primitive Gaussian basis
functions a,b,c,d.
"""
function coulomb(a::PGBF,b::PGBF,c::PGBF,d::PGBF)
    return a.norm*b.norm*c.norm*d.norm*coulomb(a.expn,a.xyz...,a.I,a.J,a.K,
        b.expn,b.xyz...,b.I,b.J,b.K,
        c.expn,c.xyz...,c.I,c.J,c.K,
        d.expn,d.xyz...,d.I,d.J,d.K)
end

fB(i,l1,l2,p,a,b,r,g) = binomial_prefactor(i,l1,l2,p-a,p-b)*B0(i,r,g)
B0(i,r,g) = fact_ratio2(i,r)*(4g)^(r-i)
#fact_ratio2(a,b) = factorial(a)/factorial(b)/factorial(a-2b)
fact_ratio2(a,b) = prod((b+1):a)/factorial(a-2b) # faster

function Bterm(i1,i2,r1,r2,u,l1,l2,l3,l4,
        Px,Ax,Bx,Qx,Cx,Dx,
        gamma1,gamma2,delta)
    # THO eq. 2.22
    return fB(i1,l1,l2,Px,Ax,Bx,r1,gamma1)*
        (-1)^i2*fB(i2,l3,l4,Qx,Cx,Dx,r2,gamma2)*
        (-1)^u*fact_ratio2(i1+i2-2*(r1+r2),u)*
        (Qx-Px)^(i1+i2-2*(r1+r2)-2*u)/delta^(i1+i2-2*(r1+r2)-u)
end

function Barray(l1,l2,l3,l4,p,a,b,q,c,d,g1,g2,delta)
    Imax = l1+l2+l3+l4+1
    B = OffsetArray(zeros(Float64,Imax),0:(Imax-1))
    for i1 in 0:(l1+l2)
        for i2 in 0:(l3+l4)
            for r1 in 0:div(i1,2)
                for r2 in 0:div(i2,2)
                    for u in 0:(div(i1+i2,2)-r1-r2)
                        I = i1+i2-2*(r1+r2)-u
                        B[I] += Bterm(i1,i2,r1,r2,u,l1,l2,l3,l4,
                                      p,a,b,q,c,d,g1,g2,delta)
                    end
                end
            end
        end
    end
    return B
end

"""
    coulomb(a::CGBF,b::CGBF,c::CGBF,d::CGBF)

Compute the coulomb repulsion between four contracted Gaussian basis
functions a,b,c,d.
"""
coulomb(a::CGBF,b::CGBF,c::CGBF,d::CGBF) = contract(coulomb,a,b,c,d)

function all_twoe_ints(bfs,ERI=coulomb)
    n = length(bfs)
    totlen = div(n*(n+1)*(n*n+n+2),8)
    ints2e = zeros(Float64,totlen)
    Threads.@threads for (i,j,k,l) in collect(iiterator(n))
        ints2e[iindex(i,j,k,l)] = ERI(bfs[i],bfs[j],bfs[k],bfs[l])
    end
    return ints2e
end

function make2JmK(D::Array{Float64,2},Ints::Array{Float64,1})
    n = size(D,1)
    G = Array{Float64}(n,n)
    D1 = reshape(D,n*n)
    temp = Array{Float64}(n*n)
    for (i,j) in pairs(n)
        kl = 1
        for (k,l) in rpairs(n)
            temp[kl] = 2*Ints[iindex(i,j,k,l)]-Ints[iindex(i,k,j,l)]
            kl += 1
        end
        G[i,j] = G[j,i] = dot(D1,temp)
    end
    return G
end

dmat(U::Array{Float64,2},nocc::Int64) = U[:,1:nocc]*U[:,1:nocc]'


path = "docs/src/assets/"

using ISA, Plots
𝐶₀ = AMFMtriplet(t->exp(-t^2/5),t->200.0,0.0)
𝐶₁ = AMFMtriplet(t->1.0,t->100*t,0.1)
𝐶₂ = AMFMtriplet(t->0.8*cos(11t),t->100 + 70.5*sin(5t),π)
𝑆 = compSet([𝐶₀,𝐶₁,𝐶₂])
z = AMFMmodel(𝑆)
plot(z; timeaxis=-1.0:0.001:1.0)
png(path*"IS_exAMFMmodel1.png")

using ISA, Plots
𝐶₀ = AMFMtriplet(t->exp(-t^2/5),t->200.0,0.0)
𝐶₁ = AMFMtriplet(t->1.0,t->100*t,0.1)
𝐶₂ = AMFMtriplet(t->0.8*cos(11t),t->100 + 70.5*sin(5t),π)
𝑆 = compSet([𝐶₀,𝐶₁,𝐶₂])
z = AMFMmodel(𝑆)
plot(𝑆; timeaxis=-1.0:0.001:1.0)
png(path*"IS_exAMFMmodel2.png")

using ISA, Plots
𝐶₀ = AMFMtriplet(t->exp(-t^2/5),t->200.0,0.0)
𝐶₁ = AMFMtriplet(t->1.0,t->100*t,0.1)
𝐶₂ = AMFMtriplet(t->0.8*cos(11t),t->100 + 70.5*sin(5t),π)
𝑆 = compSet([𝐶₀,𝐶₁,𝐶₂])
z = AMFMmodel(𝑆)
plot(𝑆, timeaxis=-1.0:0.001:1.0, view="TF",
     left_margin=15Plots.mm, margin=5Plots.mm)
png(path*"IS_exAMFMmodel3.png")

#--------------------------------
using ISA, Plots
𝐶₀ = AMFMtriplet(t->exp(-t^2),t->100,0.0)
𝐶₁ = AMFMtriplet(t->1,t->25+20*sin(t),0.0)
𝐶₂ = AMFMtriplet(t->sin(11t),t->250,0.0)
𝑆 = compSet([𝐶₀,𝐶₁,𝐶₂])
z = AMFMmodel(𝑆)
plot(z; timeaxis=-1.0:0.001:1.0)
png(path*"IS_exAMFMmodel4.png")

using ISA, Plots
𝐶₀ = AMFMtriplet(t->exp(-t^2),t->100,0.0)
𝐶₁ = AMFMtriplet(t->1,t->25+20*sin(t),0.0)
𝐶₂ = AMFMtriplet(t->sin(11t),t->250,0.0)
𝑆 = compSet([𝐶₀,𝐶₁,𝐶₂])
z = AMFMmodel(𝑆)
plot(𝑆; timeaxis=-1.0:0.001:1.0)
png(path*"IS_exAMFMmodel5.png")

using ISA, Plots
𝐶₀ = AMFMtriplet(t->exp(-t^2),t->100,0.0)
𝐶₁ = AMFMtriplet(t->1,t->25+20*sin(t),0.0)
𝐶₂ = AMFMtriplet(t->sin(11t),t->250,0.0)
𝑆 = compSet([𝐶₀,𝐶₁,𝐶₂])
z = AMFMmodel(𝑆)
plot(𝑆, timeaxis=-1.0:0.001:1.0, view="TF",
     left_margin=15Plots.mm, margin=5Plots.mm)
png(path*"IS_exAMFMmodel6.png")

#---------------------------

using ISA, Plots
𝐶₀ = AMFMtriplet(t->10,t->10t^2+2t+100,1.0)
𝐶₁ = AMFMtriplet(t->5exp(0.1t),t->250,0.0)
𝐶₂ = AMFMtriplet(t->10t,t->5cos(50t)+50,0.0)
𝑆 = compSet([𝐶₀,𝐶₁,𝐶₂])
z = AMFMmodel(𝑆)
plot(z; timeaxis=-1.0:0.001:1.0)
png(path*"IS_exAMFMmodel7.png")

using ISA, Plots
𝐶₀ = AMFMtriplet(t->10,t->10t^2+2t+100,1.0)
𝐶₁ = AMFMtriplet(t->5exp(0.1t),t->250,0.0)
𝐶₂ = AMFMtriplet(t->10t,t->5cos(50t)+50,0.0)
𝑆 = compSet([𝐶₀,𝐶₁,𝐶₂])
z = AMFMmodel(𝑆)
plot(𝑆; timeaxis=-1.0:0.001:1.0)
png(path*"IS_exAMFMmodel8.png")

using ISA, Plots
𝐶₀ = AMFMtriplet(t->10,t->10t^2+2t+100,1.0)
𝐶₁ = AMFMtriplet(t->5exp(0.1t),t->250,0.0)
𝐶₂ = AMFMtriplet(t->10t,t->5cos(50t)+50,0.0)
𝑆 = compSet([𝐶₀,𝐶₁,𝐶₂])
z = AMFMmodel(𝑆)
plot(𝑆, timeaxis=-1.0:0.001:1.0, view="TF",
     left_margin=15Plots.mm, margin=5Plots.mm)
png(path*"IS_exAMFMmodel9.png")


struct TreerXgb
    fm
    featur::Vector{Int64}
end

struct TreedaXgb
    fm
    featur::Vector{Int64}
    lev::Vector{Int64}
    ni::Vector{Int64}
end

"""
    treer_xgb(X, y;
        subsample = 1, colsample_bytree = 1, colsample_bynode = 1,
        max_depth = 6, min_child_weight = 5,
        lambda = 0, verbose = false, kwargs...)
Regression tree with XGBoost.
* `X` : X-data (n obs., p variables).
* `y` : Univariate Y-data (n obs.).
* `subsample` : Proportion of rows sampled in `X`
    for building the tree.
* `colsample_bytree` : Proportion of columns sampled in `X`
    for building the tree.
* `colsample_bynode` : Proportion of columns sampled at each node
    in the columns selected for the tree.
* `max_depth` : Maximum depth of the tree.
* `min_child_weight` : Minimum nb. observations that each leaf
    needs to have.
* `lambda` : L2 regularization term on weights.
    Increasing this value will make model more conservative.
* `verbose` : If true, fitting information are printed.
* `kwargs` : Optional named arguments to pass in function `xgboost`
    of `XGBoost.jl` (https://xgboost.readthedocs.io/en/latest/parameter.html).

The function builds a single tree using package `XGboost.jl'.

The sampling of the observations and variables are without replacement.

## References

XGBoost.jl
https://github.com/dmlc/XGBoost.jl

Breiman, L., Friedman, J. H., Olshen, R. A., and Stone, C. J. Classification
And Regression Trees. Chapman & Hall, 1984.

Gey, S., 2002. Bornes de risque, détection de ruptures, boosting :
trois thèmes statistiques autour de CART en régression (These de doctorat).
Paris 11. http://www.theses.fr/2002PA112245
"""
function treer_xgb(X, y;
        subsample = 1, colsample_bytree = 1, colsample_bynode = 1,
        max_depth = 6, min_child_weight = 5,
        lambda = 0, verbose = false, kwargs...)
    X = ensure_mat(X)
    p = size(X, 2)
    num_round = 1
    fm = xgboost(X, num_round; label = Float64.(vec(y)),
        seed = Int64(round(rand(1)[1] * 1e5)),
        booster = :gbtree,
        tree_method = :auto,
        eta = 1, # learning rate
        subsample = subsample,
        colsample_bytree = colsample_bytree, colsample_bylevel = 1,
        colsample_bynode = colsample_bynode,
        max_depth = max_depth, min_child_weight = min_child_weight,
        lambda = lambda,
        silent = !verbose, kwargs...)
    featur = collect(1:p)
    TreerXgb(fm, featur)
end

"""
    rfr_xgb(X, y; rep = 50,
        subsample = .7,
        colsample_bytree = 1, colsample_bynode = 1/3,
        max_depth = 6, min_child_weight = 5,
        lambda = 0, verbose = false, kwargs...)
Random forest regression with XGBoost.
* `X` : X-data (n obs., p variables).
* `y` : Univariate Y-data (n obs.).
* `rep` : Nb. trees to build in the forest.
* `subsample` : Proportion of rows sampled in `X`
    for building each tree.
* `colsample_bytree` : Proportion of columns sampled in `X` at each tree.
* `colsample_bynode` : Proportion of columns sampled at each node
    in the columns selected for the tree.
* `max_depth` : Maximum depth of the trees.
* `min_child_weight` : Minimum nb. observations that each leaf
    needs to have.
* `lambda` : L2 regularization term on weights.
    Increasing this value will make model more conservative.
* `verbose` : If true, fitting information are printed.
* `kwargs` : Optional named arguments to pass in function `xgboost`
    of `XGBoost.jl` (https://xgboost.readthedocs.io/en/latest/parameter.html).

The function uses package `XGboost.jl' to build the forest.
See https://xgboost.readthedocs.io/en/latest/tutorials/rf.html.

## References

XGBoost.jl
https://github.com/dmlc/XGBoost.jl

Breiman, L., 1996. Bagging predictors. Mach Learn 24, 123–140.
https://doi.org/10.1007/BF00058655

Breiman, L., 2001. Random Forests. Machine Learning 45, 5–32.
https://doi.org/10.1023/A:1010933404324

Genuer, R., 2010. Forêts aléatoires : aspects théoriques,
sélection de variables et applications. PhD Thesis. Université Paris Sud - Paris XI.

Gey, S., 2002. Bornes de risque, détection de ruptures, boosting :
trois thèmes statistiques autour de CART en régression (These de doctorat).
Paris 11. http://www.theses.fr/2002PA112245
"""
function rfr_xgb(X, y; rep = 50,
        subsample = .7,
        colsample_bytree = 1, colsample_bynode = 1/3,
        max_depth = 6, min_child_weight = 5,
        lambda = 0, verbose = false, kwargs...)
    X = ensure_mat(X)
    p = size(X, 2)
    num_round = 1
    fm = xgboost(X, num_round; label = Float64.(vec(y)),
        seed = Int64(round(rand(1)[1] * 1e5)),
        booster = :gbtree,
        tree_method = :auto,
        num_parallel_tree = rep,
        eta = 1, # learning rate
        subsample = subsample,
        colsample_bytree = colsample_bytree, colsample_bylevel = 1,
        colsample_bynode = colsample_bynode,
        max_depth = max_depth, min_child_weight = min_child_weight,
        lambda = lambda,
        silent = !verbose, kwargs...)
    featur = collect(1:p)
    TreerXgb(fm, featur)
end

"""
    xgboostr(X, y; rep = 50, eta = .3,
        subsample = .7, colsample_bytree = 1, colsample_bynode = 1/3,
        max_depth = 6, min_child_weight = 5,
        lambda = 1, verbose = false, kwargs...)
XGBoost regression.
* `X` : X-data (n obs., p variables).
* `y` : Univariate Y-data (n obs.).
* `rep` : Nb. trees to build.
* `eta` : Learning rate ([0, 1]).
* `subsample` : Proportion of rows sampled in `X`
    for building each tree.
* `colsample_bytree` : Proportion of columns sampled in `X` at each tree.
* `colsample_bynode` : Proportion of columns sampled at each node
    in the columns selected for the tree.
* `max_depth` : Maximum depth of the trees.
* `min_child_weight` : Minimum nb. observations that each leaf
    needs to have.
* `lambda` : L2 regularization term on weights.
    Increasing this value will make model more conservative.
* `verbose` : If true, fitting information are printed.
* `kwargs` : Optional named arguments to pass in function `xgboost`
    of `XGBoost.jl` (https://xgboost.readthedocs.io/en/latest/parameter.html).

The function uses package `XGboost.jl' to build the forest.
See https://xgboost.readthedocs.io/en/latest/tutorials/rf.html.

## References

XGBoost
https://xgboost.readthedocs.io/en/latest/index.html

XGBoost.jl
https://github.com/dmlc/XGBoost.jl
"""
function xgboostr(X, y; rep = 50, eta = .3,
        subsample = .7, colsample_bytree = 1, colsample_bynode = 1/3,
        max_depth = 6, min_child_weight = 5,
        lambda = 1, verbose = false, kwargs...)
    X = ensure_mat(X)
    p = size(X, 2)
    num_round = rep
    fm = xgboost(X, num_round; label = Float64.(vec(y)),
        seed = Int64(round(rand(1)[1] * 1e5)),
        booster = :gbtree,
        tree_method = :auto,
        num_parallel_tree = 1,
        eta = eta,
        subsample = subsample,
        colsample_bytree = colsample_bytree,
        colsample_bylevel = 1, colsample_bynode = colsample_bynode,
        max_depth = max_depth, min_child_weight = min_child_weight,
        lambda = lambda,
        silent = !verbose, kwargs...)
    featur = collect(1:p)
    TreerXgb(fm, featur)
end

"""
    predict(object::TreerXgb, X)
Compute Y-predictions from a fitted model.
* `object` : The fitted model.
* `X` : X-data for which predictions are computed.
"""
function predict(object::TreerXgb, X)
    X = ensure_mat(X)
    m = size(X, 1)
    pred = XGBoost.predict(object.fm, X)
    pred = reshape(pred, m, 1) ;
    (pred = pred,)
end

"""
    vimp_xgb(object::TreerXgb)
Compute variable (feature) importances from an XGBoost model.
* `object` : The fitted model.

Features with imp = 0 are not returned.
"""
function vimp_xgb(object::Union{TreerXgb, TreedaXgb})
    p = length(object.featur)
    res = XGBoost.importance(object.fm, string.(object.featur))
    fname = [] ; gain = [] ; cover = [] ; freq = []
    for i in res
        push!(fname, i.fname)
        push!(gain, i.gain)
        push!(cover, i.cover)
        push!(freq, i.freq)
    end
    fname = eval(Meta.parse.(fname))
    zgain = zeros(p)
    zcover = copy(zgain)
    zfreq = copy(zgain)
    for i = 1:length(fname)
        s = fname[i]
        zgain[s] = gain[i]
        zcover[s] = cover[i]
        zfreq[s] = freq[i]
    end
    res = DataFrame(hcat(zgain, zcover, zfreq), [:gain, :cover, :freq])
    res.featur = object.featur
    res
end


using CollisionDetection

"""
    weld(mesh1, mesh2, ...;boundary=false) -> welded_mesh

Build a mesh by welding or pasting together the inputs. Vertices from different
meshes that coincide up to the tolerance will be merged into one. The order
cells appear in the output mesh is equal to the order in the inputs.
boundary = true; will merge only vertices of different meshes if one of
them exists on the 'open' boundary of the mesh.
"""
function weld end

function weld(Γ₁, Γ₂; boundary=false)

    # T = eltype(eltype(Γ₁.vertices))
    T = coordtype(Γ₁)
    tol = sqrt(eps(T))

    verts1 = skeleton(Γ₁,0) # TODO: check this, it might break for unused verts
    radii = zeros(numcells(verts1))
    indcs = [v[1] for v in cells(verts1)]
    cntrs = [cartesian(center(chart(verts1,v))) for v in cells(verts1)]
    tree = Octree(cntrs, radii)

    nv1 = numvertices(Γ₁)
    nv2 = numvertices(Γ₂)

    idmap = collect(nv1 .+ (1:nv2))
    num_equal_vertices = 0

    V1 = vertices(Γ₁)
    V2 = similar(V1,0)

    in_interior1 = interior_vpredicate(Γ₁)
    in_interior2 = interior_vpredicate(Γ₂)
    cpred(u,v,i,j) = boundary ? norm(u-v) < tol && (!in_interior1(indcs[i])
                || !in_interior2(j)) : norm(u-v) < tol

    for (j,v) in enumerate(Γ₂.vertices)
        found = false
        pred(c,s) = fitsinbox(Array(v), 0.0, c, s+tol)
        for box in boxes(tree, pred)
            for i in box
                u = cntrs[i]
                if cpred(u,v,i,j) #norm(u-v) < tol
                    idmap[j] = indcs[i]
                    num_equal_vertices += 1
                    found = true
                    # we implicitly assume that the first copy of a point
                    # appearing in the vertex buffer is the one actually
                    # referred to in the index buffer. DANGEROUS!
                    break
                end
            end
            found && break
        end
        if !found
            idmap[j] -= num_equal_vertices
            push!(V2, v)
        end
    end

    V = [V1; V2]
    F = [Γ₁.faces; Γ₂.faces]

    nc1 = numcells(Γ₁)
    # for (i,c) in enumerate(Γ₂.faces)
    for (i,c) in enumerate(cells(Γ₂))
        F[nc1+i] = map_ids(c, idmap)
    end

    return Mesh(V,unique(F))

end

function weld(G1::SComplex2D, G2::SComplex2D; seam)

    T = coordtype(G1)
    P = vertextype(G1)
    tol = sqrt(eps(T))

    nodes1 = skeleton(G1,0)
    edges1 = skeleton(G1,1)

    node_ctrs1 = [cartesian(center(chart(nodes1,c))) for c in cells(nodes1)]
    edge_ctrs1 = [cartesian(center(chart(edges1,c))) for c in cells(edges1)]

    node_radii1 = zeros(T, length(node_ctrs1))
    edge_radii1 = [volume(chart(edges1,c)) for c in cells(edges1)] / 2

    node_tree1 = Octree(node_ctrs1, node_radii1)
    edge_tree1 = Octree(edge_ctrs1, edge_radii1)

    nodes2 = skeleton(G2,0)
    edges2 = skeleton(G2,1)

    node_is_on_seam = inclosure_gpredicate(seam)
    edge_is_on_seam = overlap_gpredicate(seam)
    # is_interior_edge = interior_tpredicate(G1)
    # is_interior_node = interior_vpredicate(G1)

    Nodes2 = similar(G2.nodes, 0)
    Edges2 = similar(G2.edges, 0)

    num_equal_edges = 0
    edges_map = collect(numcells(edges1) .+ (1:numcells(edges2)))
    for (j,edge2) in enumerate(cells(edges2))
        found = false
        ctr2 = cartesian(center(chart(edges2, edge2)))
        rad2 = volume(chart(edges2,edge2)) / 2
        for box in boxes(edge_tree1, (c,s)->fitsinbox(ctr2, rad2, c, s+tol))
            for i in box
                ctr1 = edge_ctrs1[i]
                # For now we assume edges are the same if their centroids coincide
                if norm(ctr1-ctr2) < tol && edge_is_on_seam(chart(edges1,cells(edges1)[i]))
                # if norm(ctr1-ctr2) < tol && !is_interior_edge(cells(edges1)[i])
                    edges_map[j] = i
                    num_equal_edges += 1
                    found = true
                    break
                end
                found && break
            end
        end
        if !found
            push!(Edges2, edge2)
            edges_map[j] -= num_equal_edges
        end
    end

    Faces2 = similar(G2.faces, 0)
    for face in cells(G2)
        push!(Faces2,map(p->sign(p)*getindex(edges_map,abs(p)),face))
    end

    num_equal_nodes = 0
    nodes_map = collect(numcells(nodes1) .+ (1:numcells(nodes2)))
    for (j,node2) in enumerate(cells(nodes2))
        found = false
        ctr2 = cartesian(center(chart(nodes2, node2)))
        rad2 = 0.0
        for box in boxes(node_tree1, (c,s)->fitsinbox(ctr2, rad2, c, s+tol))
            for i in box
                ctr1 = node_ctrs1[i]
                # For now we assume edges are the same if their centroids coincide
                if norm(ctr1-ctr2) < tol && node_is_on_seam(ctr1)
                # if norm(ctr1-ctr2) < tol && !is_interior_node(cells(nodes1)[i][1])
                    nodes_map[j] = i
                    num_equal_nodes += 1
                    found = true
                    break
                end
                found && break
            end
        end
        if !found
            push!(Nodes2, node2)
            nodes_map[j] -= num_equal_nodes
        end
    end

    # Edges2 = similar(G2.edges, 0)
    for (i,edge) in enumerate(Edges2)
        edge_mapped = map(p->sign(p)*getindex(nodes_map,abs(p)),edge)
        Edges2[i] = edge_mapped
        # push!(Edges2, edge_mapped)
    end

    SComplex2D{T,P}(
        vcat(G1.vertices, G2.vertices),
        vcat(G1.nodes, Nodes2),
        vcat(G1.edges, Edges2),
        vcat(G1.faces, Faces2))
end

@generated function map_ids(c::SVector{N,T}, idmap) where {N,T}
    xp = :(SVector{$N,$T}())
    for i in 1:N
        push!(xp.args, :(idmap[c[$i]]))
    end
    return xp
end

function weld(meshes...; boundary=false)
    @assert length(meshes) > 1
    G = weld(meshes[1], meshes[2], boundary=boundary)
    for i in 3:length(meshes)
        G = weld(G, meshes[i], boundary=boundary)
    end
    G
end


@testset "Objective" begin
    T = 3
    nx = 2
    nu = 1
    nw = 0
    ot = (x, u, w) -> dot(x, x) + 0.1 * dot(u, u)
    oT = (x, u, w) -> 10.0 * dot(x, x)
    ct = Cost(ot, nx, nu, nw, [t for t = 1:T-1])
    cT = Cost(oT, nx, 0, nw, [T])
    obj = [ct, cT]

    J = [0.0]
    grad = zeros((T - 1) * (nx + nu) + nx)
    idx_xu = [collect((t - 1) * (nx + nu) .+ (1:(nx + (t == T ? 0 : nu)))) for t = 1:T]
    x1 = ones(nx)
    u1 = ones(nu)
    w1 = zeros(nw)
    X = [x1 for t = 1:T]
    U = [t < T ? u1 : zeros(0) for t = 1:T]
    W = [w1 for t = 1:T]

    ct.val(ct.val_cache, x1, u1, w1)
    ct.grad(ct.grad_cache, x1, u1, w1)
    @test ct.val_cache[1] ≈ ot(x1, u1, w1)
    @test norm(ct.grad_cache - [2.0 * x1; 0.2 * u1]) < 1.0e-8

    cT.val(cT.val_cache, x1, u1, w1)
    cT.grad(cT.grad_cache, x1, u1, w1)
    @test cT.val_cache[1] ≈ oT(x1, u1, w1)
    @test norm(cT.grad_cache - 20.0 * x1) < 1.0e-8

    @test DirectTrajectoryOptimization.eval_obj(obj, X, U, X) - sum([ot(X[t], U[t], W[t]) for t = 1:T-1]) - oT(X[T], U[T], W[T]) ≈ 0.0
    DirectTrajectoryOptimization.eval_obj_grad!(grad, idx_xu, obj, X, U, W)
    @test norm(grad - vcat([[2.0 * x1; 0.2 * u1] for t = 1:T-1]..., 20.0 * x1)) < 1.0e-8

    # info = @benchmark DirectTrajectoryOptimization.eval_obj($obj, $X, $U, $W)
    # info = @benchmark DirectTrajectoryOptimization.eval_obj_grad!($grad, $idx_xu, $obj, $X, $U, $W)
end


using ProblemParser
using Utils

using Underscores

const file = find_input(@__FILE__)
const input = parse(Split(",", Convert(Int16)), slurp(file))

function part1(input)
    sort!(input)
    median = input[length(input) ÷ 2]
    sum(abs(median - pos) for pos in input)
end

assertequal(part1(input), 349769)

triangular(n) = n*(n + 1) ÷ 2

function part2(input)
    mean = sum(input) ÷ length(input)
    sum(triangular(abs(mean - pos)) for pos in input)
end

assertequal(part2(input), 99540554)


"""
`rotations = grid_rotations(maxradians, rgridsz, SD)` generates
a set of rotations (AffineMap) useful for a gridsearch of
possible rotations to align a pair of images.

`maxradians` is either a single maximum angle (in 2D) or a set of
Euler angles (in 3D and higher). `rgridsz` is one or more integers
specifying the number of gridpoints to search in each of the rotation
axes, corresponding with entries in `maxradians`. `SD` is a matrix
specifying the sample spacing.

e.g. `grid_rotations((pi/8,pi/8,pi/8), (3,3,3), Matrix{Float64}(I,3,3))` would
return an array of 27 rotations with 3 possible angles for each
Euler axis: -pi/8, 0, pi/8. Passing `Matrix{Float64}(I,3,3)` for SD indicates
that the resulting  transforms are meant to be applied to an image with isotropic
pixel spacing.
"""
function grid_rotations(maxradians, rgridsz, SD)
    rgridsz = [rgridsz...]
    maxradians = [maxradians...]
    nd = size(SD,1)
    if !all(isodd, rgridsz)
        @warn("rgridsz should be odd; rounding up to the next odd integer")
    end
    for i = 1:length(rgridsz)
        if !isodd(rgridsz[i])
            rgridsz[i] = max(round(Int, rgridsz[i]) + 1, 1)
        end
    end
    grid_radius = map(x->div(x,2), rgridsz)
    if nd > 2
        gridcoords = [range(-grid_radius[x]*maxradians[x], stop=grid_radius[x]*maxradians[x], length=rgridsz[x]) for x=1:nd]
        rotation_angles = Iterators.product(gridcoords...)
    else
        rotation_angles = range(-grid_radius[1]*maxradians[1], stop=grid_radius[1]*maxradians[1], length=rgridsz[1])
    end
    axmat = Matrix{Float64}(I,nd,nd)
    axs = map(x->axmat[:,x], 1:nd)
    tfeye = tformeye(nd)
    output = typeof(tfeye)[]
    for ra in rotation_angles
        if nd > 2
            euler_rots = map(x->tformrotate(x...), zip(axs, ra))
            rot = foldr(∘, euler_rots, init=tfeye)
        elseif nd == 2
            rot = tformrotate(ra)
        else
            error("Unsupported dimensionality")
        end
        push!(output, AffineMap(SD*rot.linear/SD , zeros(nd))) #account for sample spacing
    end
    return output
end

"""
`best_tform, best_mm = rotation_gridsearch(fixed, moving, maxshift, maxradians, rgridsz, SD =Matrix{Float64}(I,ndims(fixed),ndims(fixed))))`
Tries a grid of rotations to align `moving` to `fixed`.  Also calculates the best translation (`maxshift` pixels
or less) to align the images after performing the rotation. Returns an AffineMap that captures both the
best rotation and shift out of the values searched, along with the mismatch value after applying that transform (`best_mm`).

For more on how the arguments `maxradians`, `rgridsz`, and `SD` influence the search, see the documentation for
`grid_rotations`.
"""
function rotation_gridsearch(fixed, moving, maxshift, maxradians, rgridsz, SD = Matrix{Float64}(I,ndims(fixed),ndims(fixed)))
    rgridsz = [rgridsz...]
    nd = ndims(moving)
    @assert nd == ndims(fixed)
    rots = grid_rotations(maxradians, rgridsz, SD)
    best_mm = Inf
    best_rot = tformeye(nd)
    best_shift = zeros(nd)
    for rot in rots
        new_moving = transform(moving, rot)
        #calc mismatch
        #mm = mismatch(fixed, new_moving, maxshift; normalization=:pixels)
        mm = mismatch(fixed, new_moving, maxshift)
        thresh = 0.1*maximum(x->x.denom, mm)
        best_i = indmin_mismatch(mm, thresh)
        cur_best =ratio(mm[best_i], 0.0)
        if cur_best < best_mm
            best_mm = cur_best
            best_rot = rot
            best_shift = [best_i.I...]
        end
    end
    return tformtranslate(best_shift) ∘ best_rot, best_mm
end


module sti
using ValidatedNumerics;

pos = x -> x >= zero(x) ? x : zero(x)
neg = x -> x < zero(x) ? -x : zero(x)

function constituentMatrix{T}(matrix :: Array{T, 2})
    @assert ndims(matrix) == 2
    @assert size(matrix, 1) == size(matrix, 2)

    positivMatrix = map(pos, matrix)
    negativMatrix = map(neg, matrix)
    [
        positivMatrix negativMatrix;
        negativMatrix positivMatrix
    ] :: Array{T, 2}
end

function STI{T}(intervalVector :: Array{Interval{T}})
    @assert ndims(intervalVector) == 1

    function stiReducer(reduced, element)
        reducedSize = size(reduced, 1)
        if (reducedSize == 0)
            return [-element.lo; element.hi]
        end
        halfReducedSize = Int(reducedSize/2)
        vcat(
            slice(reduced, 1:halfReducedSize),
            [-element.lo],
            slice(reduced, (1 + halfReducedSize):reducedSize),
            [element.hi]
        ) # Could use cat instead
    end

    reduce(stiReducer, [], intervalVector) :: Array{T}
end

function reverseSTI{T}(intervalVector :: Array{T})
    @assert ndims(intervalVector) == 1

    vectorSize = size(intervalVector, 1)
    halfVectorSize = Int(vectorSize/2)
    negativLoIntervalPart = slice(intervalVector, 1:halfVectorSize)
    positivHiIntervalPart = slice(intervalVector, (1 + halfVectorSize):vectorSize)

    positivLoIntervalPart = map(-, negativLoIntervalPart)
    zippedIntervalParts = zip(positivLoIntervalPart, positivHiIntervalPart)
    collect(Interval{T}, map(x -> safeInterval(x[1], x[2]), zippedIntervalParts)) :: Array{Interval{T}}
end

function safeInterval(a, b)
    a < b ? @interval(a, b) : @interval(b, a)
end

end


# idea: use a type parameterized on functions as the stack element
# try to only allocate once
using DataStructures
import Base: convert, promote_rule, +, -, *, /, >, <, log

type RDiff <: Real
    val::Float64
    adj::Float64
    RDiff(val, adj) = new(val, adj)
end
typealias Converts Union{Float64, Int, Irrational, Rational}
RDiff(x::Converts) = RDiff(x, 0.)

convert(::Type{RDiff}, x::Converts) = RDiff(x)
promote_rule(::Type{Bool}, ::Type{RDiff}) = RDiff
promote_rule(::Type{Float64}, ::Type{RDiff}) = RDiff
promote_rule{T<:Integer}(::Type{T}, ::Type{RDiff}) = RDiff

value(x::RDiff) = x.val
adjoint(x::RDiff) = x.adj
>(x::RDiff, y::RDiff) = x.val > y.val
<(x::RDiff, y::RDiff) = x.val < y.val

type Node{F<:Function, T1<:Tuple, T2}
    op::F
    inputs::T1
    outputs::T2
end

global varstack = Stack(Node)

function grad(f, N)
    y = Array{RDiff}(N)
    function ∇f(x, storage)
        copy!(y, x)
        if isempty(varstack)
            f(y)
        end
        for node in reverse_iter(varstack)
            forward(node)
        end
        for node in varstack
            backprop(node)
        end
        for i in eachindex(storage)
            storage[i] = y[i].val
        end
    end

    return ∇f
end

function clear!(s::Stack)
    while !isempty(s)
        pop!(s)
    end
    nothing
end

backprop(n::Node) = backprop(n.op, n.inputs..., n.outputs)
forward(n::Node) = forward(n.op, n.inputs..., n.outputs)

macro diff_rule(sig, code)
    fn = _fn_of_vals(sig)
    op = sig.args[1]
    op_args = sig.args[2:end]
    argsyms = Symbol[s.args[1] for s in sig.args[2:end]]
    ex = quote
        function $op($(op_args...))
            res = map(RDiff, $fn)
            n = Node($op, tuple($(argsyms...)), res)
            push!(varstack, n)
            res
        end

        function forward(::typeof($op), $(op_args...), res)
            for (r, n) in zip(res, $fn)
                r.val = n
                r.adj = 0.
            end
            nothing
        end

        function backprop(::typeof($op), $(op_args...), res)
            $code
            nothing
        end
    end
    esc(ex)
end

_fn_of_vals(x) = x
function _fn_of_vals(ex::Expr)
    if ex.head == :(::)
        out = :(value($(ex.args[1])))
    else
        out = copy(ex)
        for i in eachindex(out.args)
            out.args[i] = _fn_of_vals(out.args[i])
        end
    end
    out
end

@diff_rule log(x::RDiff) begin
    x.adj += res.adj / x.val
end

@diff_rule +(x::RDiff, y::RDiff) begin
    x.adj += res.adj
    y.adj += res.adj
end

@diff_rule -(x::RDiff, y::RDiff) begin
    x.adj += res.adj
    y.adj += -res.adj
end

@diff_rule *(x::RDiff, y::RDiff) begin
    x.adj += res.adj * y.val
    y.adj += res.adj * x.val
end

@diff_rule /(x::RDiff, y::RDiff) begin
    x.adj += res.adj / y.val
    y.adj += -res.adj * x.val / y.val^2
end

@diff_rule -(x::RDiff) begin
    x.adj += -res.adj
end


export parallelTomoMod

"""
function parallelTomoMod(w,n)

model for tomography problems where for each angle n parallel
projections are used.

Input:

    w - weights for each angle
    n - number of projections per angle

Output:
    sig,dsig - model and derivative

"""
function parallelTomoMod(w::Vector{Float64},n::Int)
    sig  = kron(w,ones(n))
    dsig = kron(speye(length(w)),ones(n))
    return sig,dsig
end


function trajectory!(traj::SystemTrajectory{h}, cs::ControlSystem,
                     γ::AbstractVector{<:AffineStrategy},
                     last_op::SystemTrajectory{h}, x0::AbstractVector,
                     max_elwise_divergence::Float64=Inf) where {h}

    @assert samplingtime(traj) == samplingtime(last_op) == samplingtime(cs)

    xₖ = x0

    for k in 1:h
        tₖ = time_disc2cont(traj, k)
        # the quantities on the old operating point
        x̃ₖ = last_op.x[k]
        ũₖ = last_op.u[k]
        # the current strategy
        γₖ = γ[k]
        # the deviation from the last operating point
        Δxₖ = xₖ - x̃ₖ

        # return early if divergence from operating point exceeded
        infnorm(Δxₖ) <= max_elwise_divergence || return false

        # record the new operating point:
        traj.x[k] = xₖ
        traj.u[k] = uₖ = control_input(γₖ, Δxₖ, ũₖ)

        # integrate x forward in time for the next iteration.
        xₖ = next_x(cs, xₖ, uₖ, tₖ)
    end
    return true
end

"""
$(FUNCTIONNAME)(g::AbstractGame, traj::SystemTrajectory)

Returns a vector of costs for each player
"""
function cost(g::GeneralGame, traj::SystemTrajectory)
    map(player_costs(g)) do ci
        sum(ci(g, traj.x[k], traj.u[k], time_disc2cont(traj, k)) for k in
            eachindex(traj.x))
    end
end


using SimpleDiffEq, StaticArrays, OrdinaryDiffEq, Test

function loop(u, p, t)
    @inbounds begin
        σ = p[1]; ρ = p[2]; β = p[3]
        du1 = σ*(u[2]-u[1])
        du2 = u[1]*(ρ-u[3]) - u[2]
        du3 = u[1]*u[2] - β*u[3]
        return SVector{3}(du1, du2, du3)
    end
end
function liip(du, u, p, t)
    σ = p[1]; ρ = p[2]; β = p[3]
    du[1] = σ*(u[2]-u[1])
    du[2] = u[1]*(ρ-u[3]) - u[2]
    du[3] = u[1]*u[2] - β*u[3]
    return nothing
end

u0 = 10ones(3)
dt = 1e-2

odeoop = ODEProblem{false}(loop, SVector{3}(u0), (0.0, 100.0),  [10, 28, 8/3])
sol = solve(odeoop,SimpleATsit5(),dt=dt)
sol2 = solve(odeoop,GPUSimpleATsit5(),dt=dt)

@test sol.u == sol2.u
@test sol.t == sol2.t


"Holds the information for the various methods' tableaus."
abstract type AbstractTableau{T} end

# order(tab::T) where {T <: AbstractTableau} = tab.o
RungeKutta.nstages(tab::T) where {T <: AbstractTableau} = tab.s

# function check_symplecticity end
function symplecticity_conditions end

@define HeaderTableau begin
    name::Symbol
    o::Int
    s::Int
end


include("common.jl")

example = """
5483143223
2745854711
5264556173
6141336146
6357385478
4167524645
2176841721
6882881134
4846848554
5283751526
"""

function to_matrix(input::AbstractString)
    rows = Vector{Int64}[]
    for line in split(input, '\n')
        isempty(line) && continue
        push!(rows, [parse(Int64, x) for x in line])
    end
    return transpose(reduce(hcat, rows)) |> collect
end

function apply_flashes!(x::AbstractMatrix{Int64}, flashed::AbstractMatrix{Bool})
    @assert size(x) == size(flashed)
    ni = size(flashed, 1)
    nj = size(flashed, 2)
    @inbounds for i in 1:ni
        for j in 1:nj
            flashed[i, j] || continue

            # Update neighbours.
            for oi in max(1, (i - 1)):min((i + 1), ni)
                for oj in max(1, (j - 1)):min((j + 1), nj)
                    # Skip the thing that has flashed.
                    oi == i && oj == j && continue
                    x[oi, oj] += 1
                end
            end
        end
    end
    return x
end

function step!(x::AbstractMatrix{Int64})
    x .+= 1  # Energy level increases.

    flashed = zeros(Bool, size(x))
    while true
        # An octopus shouldn't flash more than once.
        this_flashed = @. ((x > 9) & (!flashed))
        sum(this_flashed) == 0 && break  # Done flashing
        apply_flashes!(x, this_flashed)
        flashed .|= this_flashed  # Update the record of all flashes.
    end

    x[flashed] .= 0
    return sum(flashed)
end

function evolve!(x::AbstractMatrix, num_steps)
    total = 0
    for _ in 1:num_steps
        total += step!(x)
    end
    return total
end

function get_ans(input::AbstractString, n=100)
    x = to_matrix(input)
    return evolve!(x, n)
end

@assert get_ans(example, 10) == 204
@assert get_ans(example) == 1656
get_ans(load_input(11))

# Part 2

function get_ans2(input::AbstractString)
    x = to_matrix(input)
    step = 1
    while true
        num_flashes = step!(x)
        num_flashes == length(x) && return step
        step += 1
    end
end

@assert get_ans2(example) == 195
get_ans2(load_input(11))


"""
    gpfitbayes(..., niter::Int=5000, warmup::Int=2000)

Generate a sample from the GP parameters' posterior distribution.

Data provided must be the exceedances above the threshold, *i.e.* the data above the threshold minus
the threshold.

# Arguments
- `niter::Int = 5000`: The total number of MCMC iterations
- `warmup::Int = 2000`: The number of warmup iterations (burn-in).

# Implementation

The function uses the No-U-Turn Sampler (NUTS; [Hoffman and Gelman, 2014](http://jmlr.org/papers/v15/hoffman14a.html))
implemented in the [Mamba.jl](https://mambajl.readthedocs.io/en/latest/index.html)
package to generate a random sample from the posterior distribution.

Currently, only the improper uniform prior is implemented, *i.e.*
```math
f_{(β₂,β₃)}(β₂,β₃) ∝ 1,
```
where
```math
ϕ = X₂ × β₂,
```
```math
ξ = X₃ × β₃.
```
In the stationary case, this improper prior yields to a proper posterior if the
sample size is larger than 2 ([Northrop and Attalides, 2016](https://www.jstor.org/stable/24721296?seq=1)).

The covariates are standardized before estimating the parameters to help fit the
 model. They are transformed back on their original scales before returning the
 fitted model.

See also [`gpfitbayes`](@ref) for the other methods, [`gpfitpwm`](@ref), [`gpfit`](@ref) and [`ThresholdExceedance`](@ref).

# Reference

Hoffman M. D. and Gelman A. (2014). The No-U-Turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo. *Journal of Machine Learning Research*, 15:1593–1623.

Paul J. Northrop P. J. and Attalides N. (2016). Posterior propriety in Bayesian extreme value analyses using reference priors. *Statistica Sinica*, 26:721-743.

"""
function gpfitbayes end

"""
    gpfitbayes(y,
        logscalecov = Vector{Variable}(),
        shapecov = Vector{Variable}(),
        niter::Int=5000,
        warmup::Int=2000
        )

Generate a sample from the GP parameters' posterior distribution.

# Arguments

- `y::Vector{<:Real}`: The vector of exceedances.
- `logscalecov::Vector{<:DataItem} = Vector{Variable}()`: The covariates of the log-scale parameter.
- `shapecov::Vector{<:DataItem} = Vector{Variable}()`: The covariates of the shape parameter.
"""
function gpfitbayes(y::Vector{<:Real};
     logscalecov::Vector{<:DataItem} = Vector{Variable}(),
     shapecov::Vector{<:DataItem} = Vector{Variable}(),
     niter::Int=5000, warmup::Int=2000)::BayesianEVA

     logscalecovstd = standardize.(logscalecov)
     shapecovstd = standardize.(shapecov)

    model = ThresholdExceedance(Variable("y", y), logscalecov = logscalecovstd, shapecov = shapecovstd)

    fittedmodel = fitbayes(model, niter=niter, warmup=warmup)

    return transform(fittedmodel)

end

"""
    gpfitbayes(df::DataFrame,
        datacol::Symbol,
        logscalecovid = Vector{Symbol}(),
        shapecovid = Vector{Symbol}(),
        niter::Int=5000,
        warmup::Int=2000)

Generate a sample from the GP parameters' posterior distribution.

# Arguments
- `df::DataFrame`: The dataframe containing the data.
- `datacol::Symbol`: The symbol of the column of `df` containing the exceedances.
- `logscalecovid::Vector{Symbol} = Vector{Symbol}()`: The symbols of the columns of `df` containing the covariates of the log-scale parameter.
- `shapecovid::Vector{Symbol} = Vector{Symbol}()`: The symbols of the columns of `df` containing the covariates of the shape parameter.
"""
function gpfitbayes(df::DataFrame, datacol::Symbol;
    logscalecovid::Vector{Symbol}=Symbol[],
    shapecovid::Vector{Symbol}=Symbol[],
    niter::Int=5000, warmup::Int=2000)::BayesianEVA

    logscalecovstd = standardize.(buildVariables(df, logscalecovid))
    shapecovstd = standardize.(buildVariables(df, shapecovid))

    model = ThresholdExceedance(Variable(string(datacol), df[:, datacol]), logscalecov = logscalecovstd, shapecov = shapecovstd)

    fittedmodel = fitbayes(model, niter=niter, warmup=warmup)

    return transform(fittedmodel)

    return fm

end

"""
    gpfitbayes(model::ThresholdExceedance;
        niter::Int=5000,
        warmup::Int=2000)

Generate a sample from the GP parameters' posterior distribution.

"""
function gpfitbayes(model::ThresholdExceedance; niter::Int=5000, warmup::Int=2000)::BayesianEVA

    return fitbayes(model, niter=niter, warmup=warmup)

end


const tolerance = 10.0^-10

struct Line
    x1::Number
    y1::Number

    x2::Number
    y2::Number
end

Base.:(==)(lhs::Line, rhs::Line) = lhs.x1 == rhs.x1 && lhs.y1 == rhs.y1 && lhs.x2 == rhs.x2 && lhs.y2 == rhs.y2

# How much is missing to get to the closest new "grid square"
function rayDelta(n, a)
    s = sign(a)

    if s > 0
        return floor(n + 1) - n

    elseif s < 0
        return ceil(n - 1) - n

    else
        return 0
    end
end

function pointsOnLine(line::Line)
    x1, y1, x2, y2 = line.x1, line.y1, line.x2, line.y2
    res = Set{Tuple{Number, Number}}()
    theta = atan(y2 - y1, x2 - x1)

    x, y = x1, y1

    velX = cos(theta)
    velY = sin(theta)

    velX = isapprox(velX, 0, atol=tolerance) ? 0 : velX
    velY = isapprox(velY, 0, atol=tolerance) ? 0 : velY

    while velX != 0 && !isapprox(x, x2, atol=tolerance) || velY != 0 && !isapprox(y, y2, atol=tolerance)
        dx = rayDelta(x, velX)
        dy = rayDelta(y, velY)

        a = isnan(dx / velX) ? Inf : dx / velX
        b = isnan(dy / velY) ? Inf : dy / velY

        if a < b
            x += dx
            y += a * velY

        else
            x += b * velX
            y += dy
        end

        fx = floor(Int, x)
        fy = floor(Int, y)

        push!(res, (fx, fy))
    end

    push!(res, (x1, y1))
    push!(res, (x2, y2))

    return collect(res)
end


using OrdinaryDiffEq, ParameterizedFunctions, RecursiveArrayTools
#using DiffEqBayes
using Stan, Mamba
using Test

# This example uses the local version of stan_inference for testing purposes

include(joinpath(Pkg.dir("Stan"), "Examples", "OtherExamples", "DiffEqBayes",
  "stan_inference.jl"))

  #=
Experimental link to DiffEqBayes examples

Possible issues to investigate:

1. @ode_def_nohes needs to be kept outside local scope block
2. Minimize string replacements in Stan language model to prevent recompilations

=#

g1 = @ode_def_nohes LorenzExample begin
  dx = σ*(y-x)
  dy = x*(ρ-z) - y
  dz = x*y - β*z
end σ=10.0 ρ=28.0 β=>2.6666

ProjDir = dirname(@__FILE__)
cd(ProjDir) do

  #isdir("tmp") && rm("tmp", recursive=true)
  isfile("LorentzExample-summaryplot-1.pdf") &&
    rm("LorentzExample-summaryplot-1.pdf")
  isfile("LorentzExample-summaryplot-2.pdf") &&
    rm("LorentzExample-summaryplot-2.pdf")

  r0 = [0.1;0.0;0.0]
  tspan = (0.0,4.0)
  prob = ODEProblem(g1,r0,tspan)
  sol = solve(prob,Tsit5())
  t = collect(range(0.1, stop=4.0, length=10))
  randomized = VectorOfArray([(sol(t[i]) + .01randn(3)) for i in 1:length(t)])
  data = convert(Array,randomized)
  priors = [Normal(2.6666,1)]

  bayesian_result = stan_inference(prob,t,data,priors;
    num_samples=1500,num_warmup=1500)
  theta1 = bayesian_result.chain_results[:,["theta.1"],:]
  global sim2 = bayesian_result.chain_results[1:end,
    ["sigma.1", "sigma.2", "sigma.3", "theta.1"], :]
  ## Plotting
  p = plot(sim2, [:trace, :mean, :density, :autocor], legend=true);
  draw(p, ncol=4, filename="Lorentz-summaryplot", fmt=:pdf)

  @test mean(theta1.value[:,:,1]) ≈ 2.6666 atol=1e-1

end


module MDCDL # Multi-Dimensional Convolutional Dictionary Learning

using LinearAlgebra
using ImageFiltering.Algorithm: Alg, FIR, FFT
using ComputationalResources: AbstractResource, CPU1

import Base: promote_rule, eltype, ndims, similar, length
import Random: rand, rand!

export rand, rand!

export Shapes, Optimizers, SparseCoders

export  PolyphaseVector, FilterBank, PolyphaseFB, AbstractNsolt,
        Cnsolt, CnsoltTypeI, CnsoltTypeII, Rnsolt, RnsoltTypeI, RnsoltTypeII,
        ParallelFilters, Multiscale, AbstractOperator, TransformSystems

export  istype1, istype2, permdctmtx, cdftmtx, analyze, synthesize,
        adjoint_synthesize, upsample, downsample, serialize, deserialize,
        analysisbank, decimations, orders, nchannels, nstages, kernels,
        analysiskernels, synthesiskernels, kernelsize, getrotations,
        setrotations!, mdarray2polyphase, polyphase2mdarray, iht,
        createTransform, loadfb, savefb, copy_params!

include("sparsecoders/SparseCoders.jl")
include("optimizers/Optimizers.jl")
using MDCDL.Optimizers: iterations

struct PolyphaseVector{T,D}
    data::AbstractMatrix{T}
    nBlocks::NTuple{D, Int}
end

promote_rule(::Type{PolyphaseVector{TA,D}}, ::Type{PolyphaseVector{TB,D}}) where {TA,TB,D} = PolyphaseVector{promote_type(TA,TB), D}

abstract type CodeBook{T,D} end
abstract type FilterBank{T,D} <: CodeBook{T,D} end
abstract type PolyphaseFB{T,D} <: FilterBank{T,D} end

eltype(::Type{CB}) where {T,D,CB<:CodeBook{T,D}} = T
ndims(::Type{CB}) where {T,D,CB<:CodeBook{T,D}} = D
ndims(cb::CodeBook) = ndims(typeof(cb))

decimations(fb::FilterBank) = fb.decimationFactor
nchannels(fb::FilterBank) = sum(fb.nChannels)
orders(fb::FilterBank) = fb.polyphaseOrder

kernels(fb::FilterBank) = fb.kernels
kernelsize(fb::FilterBank) = decimations(fb) .* (1 .+ orders(fb))

struct Multiscale
    filterbanks::Tuple
    Multiscale(fbs...) = new(fbs)
end

length(ms::Multiscale) = length(ms.filterbanks)

include("filterbanks/nsolt.jl")
include("filterbanks/parameterization.jl")
include("filterbanks/parallelFilters.jl")

include("orthonormalMatrixSystem.jl")
include("polyphaseMatrices.jl")

include("shapes.jl")
include("transformSystem.jl")

include("analysisSystem.jl")
include("synthesisSystem.jl")

include("dictionaryLearning/dictionaryLearning.jl")
include("dictionaryLearning/sparseCodingStage.jl")
include("dictionaryLearning/dictionaryUpdateStage.jl")
include("dictionaryLearning/logger.jl")

include("recipes.jl")
include("io.jl")

include("utils.jl")
include("compatibles.jl")

end


_t₁(x, df::Float64) = tdistpdf(df, x)
_t₁(x::Vector, df::Float64) = map(u->tdistpdf(df, u), x)
_logt₁(x, df::Float64) = tdistlogpdf(df, x)
_logt₁(x::Vector, df::Float64) = map(u->tdistlogpdf(df, u), x)
_T₁(x, df::Float64) = tdistcdf(df, x)
_T₁(x::Vector, df::Float64) = map(u->tdistcdf(df, u), x)
_logT₁(x, df::Float64) = tdistlogcdf(df, x)
_logT₁(x::Vector, df::Float64) = map(u->tdistlogcdf(df, u), x)

function _log_g(Q, ν::Float64, k::Int)
    # Adapted from mvtdist_consts in Distributions package
    hdf = 0.5 * ν
    hdim = 0.5 * k
    shdfhdim = hdf + hdim
    v = lgamma(shdfhdim) - lgamma(hdf) - hdim*log(ν) - hdim*log(pi)
    return v - (shdfhdim * log(1 + Q/ν))
end

function scale_and_cor(Ω::Matrix{Float64})
    Ωz = similar(Ω)
    d = size(Ω,1)
    ω = sqrt(diag(Ω))
    for i in 1:d, j in 1:d
        Ωz[i,j] = Ω[i,j]/(ω[i]*ω[j])
    end
    diagm(ω), Ωz
end

function scale_and_cor(Ω::PDMat)
    Ωz = similar(Ω.mat)
    d = size(Ω,1)
    ω = sqrt(diag(Ω))
    for i in 1:d, j in 1:d
        Ωz[i,j] = Ω.mat[i,j]/(ω[i]*ω[j])
    end
    Diagonal(ω), PDMat(Ωz)
end

δ(α::Float64) = α/sqrt(1.0 + α^2)

function δ(Ω::Matrix{Float64}, α::Vector{Float64})
    ω, Ωz = scale_and_cor(Ω)
    Ωzα = Ωz*α
    αΩzα = dot(α, Ωzα)
    Ωzα/sqrt(1+αΩzα)
end

function δ(Ω::PDMat, α::Vector{Float64})
    ω, Ωz = scale_and_cor(Ω)
    (Ωz * α)/sqrt(1.0 + quad(Ωz, α))
end


# Various specialised operations using the Cholesky factorisation.

Xt_invA_X(A::Cholesky, x::AbstractVector) = sum(abs2, A.U' \ x)
function Xt_invA_X(A::Cholesky, X::AbstractVecOrMat)
    V = A.U' \ X
    return Symmetric(V'V)
end

Xt_invA_Y(X::AbstractVecOrMat, A::Cholesky, Y::AbstractVecOrMat) = (A.U' \ X)' * (A.U' \ Y)

diag_At_A(A::AbstractVecOrMat) = vec(sum(abs2.(A); dims=1))

function diag_At_B(A::AbstractVecOrMat, B::AbstractVecOrMat)
    @assert size(A) == size(B)
    return vec(sum(A .* B; dims=1))
end

diag_Xt_invA_X(A::Cholesky, X::AbstractVecOrMat) = diag_At_A(A.U' \ X)

function diag_Xt_invA_Y(X::AbstractMatrix, A::Cholesky, Y::AbstractMatrix)
    @assert size(X) == size(Y)
    return diag_At_B(A.U' \ X, A.U' \ Y)
end


@testset "Test total multi order indices" begin

# Nx = 1
midxs = totalorder([0])
@test midxs == reshape([0],(1,1))
midxs = totalorder([1])
@test midxs == reshape([0; 1],(2,1))
midxs = totalorder([3])
@test midxs == reshape([0; 1; 2; 3],(4,1))

#Nx = 2

midxs = totalorder([0; 0])
@test midxs == reshape([0; 0],(1,2))
midxs = totalorder([1; 0])
@test midxs == [0 0; 1 0]
midxs = totalorder([3;4])
@test midxs ==
     [0     0
     0     1
     0     2
     0     3
     0     4
     1     0
     1     1
     1     2
     1     3
     2     0
     2     1
     2     2
     3     0
     3     1]

# Nx = 3

midxs = totalorder([0; 0; 0])
@test midxs == reshape([0; 0; 0],(1,3))
midxs = totalorder([0; 1; 0])
@test midxs == [0 0 0; 0 1 0]

midxs = totalorder([2;4;3])

@test midxs ==
[0     0     0
0     0     1
0     0     2
0     0     3
0     1     0
0     1     1
0     1     2
0     1     3
0     2     0
0     2     1
0     2     2
0     3     0
0     3     1
0     4     0
1     0     0
1     0     1
1     0     2
1     0     3
1     1     0
1     1     1
1     1     2
1     2     0
1     2     1
1     3     0
2     0     0
2     0     1
2     0     2
2     1     0
2     1     1
2     2     0]
end


using Turing2MonteCarloMeasurements
using Test
using MonteCarloMeasurements, Turing, Distributions

@testset "Turing2MonteCarloMeasurements.jl" begin
    # ess = MonteCarloMeasurements.ess

    nr = 3 # Number of reviewers
    na = 5 # Number of articles
    reviewer_bias = rand(Normal(0,1), nr)
    article_score = rand(Normal(0,2), na)
    R = clamp.([rand(Normal(r+a, 0.1)) for r in reviewer_bias, a in article_score], -5, 5)

    m = Turing.@model reviewscore(R,nr,na) = begin
        reviewer_bias = Array{Real}(undef, nr)
        reviewer_gain = Array{Real}(undef, nr)
        true_article_score = Array{Real}(undef, na)
        reviewer_pop_bias ~ Normal(0,1)
        reviewer_pop_gain ~ Normal(1,1)
        for i = 1:nr
            reviewer_bias[i] ~ Normal(reviewer_pop_bias,1)
            reviewer_gain[i] ~ Normal(reviewer_pop_gain,1)
        end
        for j = 1:na
            true_article_score[j] ~ Normal(0,2.5)
        end
        rσ ~ truncated(Normal(1,10),0,100)
        for j = 1:na
            for i = 1:nr
                R[i,j] ~ Normal(reviewer_bias[i] + true_article_score[j] + reviewer_gain[i]*true_article_score[j], rσ)
            end
        end
        (reviewer_pop_bias=reviewer_pop_bias, reviewer_pop_gain=reviewer_pop_gain, reviewer_bias=reviewer_bias, reviewer_gain=reviewer_gain)
    end
    m = reviewscore(R,nr,na)
    chain = sample(m, HMC(0.05, 10), 1500)
    cp = Particles(chain, crop=500);
    @test size(cp.reviewer_bias) == size(reviewer_bias)
    @test size(cp.rσ) == ()
    truth = m()
    truthplot(truth, cp)

end


using SparseArrays

function strang_matrix(N)
  A = zeros(N,N)
  for i in 1:N, j in 1:N
      abs(i-j)<=1 && (A[i,j]+=1)
      i==j && (A[i,j]-=3)
  end
  A
end

# tests for full and sparse function
@testset "Full and Sparse functions:" begin
    N = 100
    d_order = 2
    approx_order = 2
    x = collect(1:1.0:N).^2

    A = DerivativeOperator{Float64}(d_order,approx_order,1.0,N,:Dirichlet0,:Dirichlet0)
    mat = convert(Array,A)
    sp_mat = sparse(A)
    @test mat == sp_mat

    @test convert(Array, A, 10) == strang_matrix(10) # Strang Matrix is defined with the center term +ve
    @test convert(Array, A, N) == strang_matrix(N) # Strang Matrix is defined with the center term +ve
    @test convert(Array,A) == sp_mat
    @test opnorm(A, Inf) == opnorm(mat, Inf)

    # testing correctness
    N = 1000
    d_order = 4
    approx_order = 10
    y = collect(1:1.0:N).^4 - 2*collect(1:1.0:N).^3 + collect(1:1.0:N).^2;
    y = convert(Array{BigFloat, 1}, y)

    A = DerivativeOperator{BigFloat}(d_order,approx_order,one(BigFloat),N,:Dirichlet0,:Dirichlet0)
    boundary_points = A.boundary_point_count
    mat = convert(Array, A, N)
    sp_mat = sparse(A)
    @test mat == sp_mat

    res = A*y
    @test res[boundary_points[1] + 1: N - boundary_points[2]] ≈ 24.0*ones(N - sum(boundary_points)) atol=10.0^-approx_order
    @test A*y ≈ mat*y atol=10.0^-approx_order
    @test A*y ≈ sp_mat*y atol=10.0^-approx_order
    @test sp_mat*y ≈ mat*y atol=10.0^-approx_order
end

@testset "Indexing tests" begin
    N = 1000
    d_order = 4
    approx_order = 10

    A = DerivativeOperator{Float64}(d_order,approx_order,1.0,N,:Dirichlet0,:Dirichlet0)
    @test A[1,1] ≈ 13.717407 atol=1e-4
    @test A[:,1] == (convert(Array,A))[:,1]
    @test A[10,20] == 0

    for i in 1:N
        @test A[i,i] == A.stencil_coefs[div(A.stencil_length, 2) + 1]
    end

    # Indexing Tests
    N = 1000
    d_order = 2
    approx_order = 2

    A = DerivativeOperator{Float64}(d_order,approx_order,1.0,N,:Dirichlet0,:Dirichlet0)
    M = convert(Array,A)

    @test A[1,1] == -2.0
    @test A[1:4,1] == M[1:4,1]
    @test A[5,2:10] == M[5,2:10]
    @test A[60:100,500:600] == M[60:100,500:600]
end

@testset begin "Operations on matrices"
    N = 51
    M = 101
    d_order = 2
    approx_order = 2

    xarr = range(0,stop=1,length=N)
    yarr = range(0,stop=1,length=M)
    dx = xarr[2]-xarr[1]
    dy = yarr[2]-yarr[1]
    F = [x^2+y for x = xarr, y = yarr]

    A = DerivativeOperator{Float64}(d_order,approx_order,dx,length(xarr),:None,:None)
    B = DerivativeOperator{Float64}(d_order,approx_order,dy,length(yarr),:None,:None)

    @test A*F ≈ 2*ones(N,M) atol=1e-2
    @test F*B ≈ 8*ones(N,M) atol=1e-2
    @test A*F*B ≈ zeros(N,M) atol=1e-2

    G = [x^2+y^2 for x = xarr, y = yarr]

    @test A*G ≈ 2*ones(N,M) atol=1e-2
    @test G*B ≈ 8*ones(N,M) atol=1e-2
    @test A*G*B ≈ zeros(N,M) atol=1e-2
end

@testset "Linear combinations of operators" begin
    # Only tests the additional functionality defined in "operator_combination.jl"
    N = 10
    Random.seed!(0); LA = DiffEqArrayOperator(rand(N,N))
    LD = DerivativeOperator{Float64}(2,2,1.0,N,:Dirichlet0,:Dirichlet0)
    @test_broken begin
      L = 1.1*LA - 2.2*LD + 3.3*I
      # Builds convert(L) the brute-force way
      fullL = zeros(N,N)
      v = zeros(N)
      for i = 1:N
          v[i] = 1.0
          fullL[:,i] = L*v
          v[i] = 0.0
      end
      @test convert(L) ≈ fullL
      @test exp(L) ≈ exp(fullL)
      for p in [1,2,Inf]
          @test opnorm(L,p) ≈ opnorm(fullL,p)
          @test opnormbound(L,p) ≈ 1.1*opnorm(LA,p) + 2.2*opnorm(LD,p) + 3.3
      end
  end
end


# This file is a part of AstroLib.jl. License is MIT "Expat".

function _co_refract(old_alt::T, altitude::T, pressure::T, temperature::T,
                    epsilon::T, to_observe::Bool) where {T<:AbstractFloat}

    if isnan(temperature)
        if altitude > 11000
            temperature = T(211.5)
        else
            temperature = 283 - 0.0065 * altitude
        end
    end

    if isnan(pressure)
        pressure = 1010 * (1 - 6.5/288000 * altitude) ^ 5.255
    end

    if !to_observe
        aout = old_alt - co_refract_forward(old_alt, pressure, temperature)
    else
        cur = old_alt + co_refract_forward(old_alt, pressure, temperature)
        while true
            last = cur
            cur = old_alt + co_refract_forward(last, pressure, temperature)
            if abs(last - cur) * 3600 < epsilon
                break
            end
        end
        aout = cur
    end
    return aout
end
"""
    co_refract(old_alt[, altitude=0, pressure=NaN, temperature=NaN,
               epsilon=0.25, to_observe=false]) -> aout

### Purpose ###

Calculate correction to altitude due to atmospheric refraction.

### Explanation ##

Because the index of refraction of air is not precisely 1.0, the atmosphere
bends all incoming light, making a star or other celestial object appear at
a slightly different altitude (or elevation) than it really is.  It is
important to understand the following definitions:

* Observed Altitude: The altitude that a star is seen to be, with a telescope.
  This is where it appears in the sky. This is should be always greater than the
  apparent altitude.

* Apparent Altitude: The altitude that a star would be at, if ~there were no
  atmosphere~ (sometimes called the "true" altitude). This is usually calculated
  from an object's celestial coordinates. Apparent altitude should always be
  smaller than the observed altitude.

Thus, for example, the Sun's apparent altitude when you see it right on the
horizon is actually -34 arcminutes.

This program uses a couple of simple formulae to estimate the effect for most
optical and radio wavelengths. Typically, you know your observed altitude
(from an observation), and want the apparent altitude. To go the other way,
this program uses an iterative approach.

### Arguments ###

* `old_alt`: observed altitude in degrees. If `to_observe` is set to true,
  this should be apparent altitude
* `altitude` (optional): the height of the observing location, in meters. This is
  only used to determine an approximate temperature and pressure, if these are not
  specified separately. Default is 0 i.e. sea level
* `pressure` (optional): the pressure at the observing location, in millibars.
  Default is NaN
* `temperature` (optional): the temperature at the observing location, in Kelvins.
  Default is NaN
* `epsilon` (optional): the accuracy to obtain, in arcseconds. If `to_observe` is
  true, then it will be calculated. Default is 0.25 arcseconds
* `to_observe` (optional boolean keyword): if set to true, it is assumed that `old_alt`
  has apparent altitude as its input and the observed altitude will be found

### Output ###

* `aout`: apparent altitude, in degrees. Observed altitude is returned if `to_observe`
  is set to true

### Example ###

The lower limb of the Sun is observed to have altitude of 0d 30'. Calculate the the true
(i.e. apparent) altitude of the Sun's lower limb using mean  conditions of air pressure and
temperature.

```jldoctest
julia> using AstroLib

julia> co_refract(0.5)
0.02584736873098442
```

### Notes ###

If altitude is set but the temperature or pressure is not, the program will make an
intelligent guess for the temperature and pressure.

#### Wavelength Dependence ####

This correction is 0 at zenith, about 1 arcminute at 45 degrees, and 34 arcminutes at the
horizon for optical wavelengths. The correction is non-negligible at all wavelengths,
but is not very easily calculable.
These formulae assume a wavelength of 550 nm, and will be accurate to about 4 arcseconds
for all visible wavelengths, for elevations of 10 degrees and higher.
Amazingly, they are also accurate for radio frequencies less than ~ 100 GHz.

#### References ####

* Meeus, Astronomical Algorithms, Chapter 15.
* Explanatory Supplement to the Astronomical Almanac, 1992.
* Methods of Experimental Physics, Vol 12 Part B, Astrophysics, Radio Telescopes, Chapter 2.5,
  "Refraction Effects in the Neutral Atmosphere", by R.K. Crane.

Code of this function is based on IDL Astronomy User's Library.
"""
co_refract(old_alt::Real, altitude::Real=0, pressure::Real=NaN, temperature::Real=NaN,
           epsilon::Real=0.25; to_observe::Bool=false) =
               _co_refract(promote(float(old_alt), float(altitude), float(pressure),
                          float(temperature), float(epsilon))..., to_observe)

"""
    co_refract_forward(alt, pre, temp) -> ref

### Purpose ###

A function used by [`co_refract`](@ref) to find apparent (or observed) altitude

### Arguments ###

* `alt`: the observed (or apparent) altitude, in degrees
* `pre`: pressure, in millibars
* `temp`: temperature, in Kelvins

### Output ###

* `ref`: the atmospheric refraction, in minutes of arc

### Notes ###

The atmospheric refraction is calculated by Saemundsson's formula

Code of this function is based on IDL Astronomy User's Library.
"""
function co_refract_forward(alt::Real, pre::Real, temp::Real)

    if alt < 15
        ref =3.569*@evalpoly(alt, 0.1594, 0.0196, 0.00002)/@evalpoly(alt, 1, 0.505, 0.0845)
    else
        ref = 0.0166667 / tand((alt + 7.31 / (alt + 4.4)))
    end
    ref *= pre * 283 / (1010 * temp)
    return ref
end


struct Standardizer{T} <: Number
    avg::T
    std::T
end
function Standardizer(x::Vector{T}; robust::Bool=false) where T
    if robust
        (mu = median(x); Standardizer(mu, 1.482602median(abs.(x .- mu))))
    else
        (mu = mean(x); Standardizer(mu, std(x, mean=mu, corrected=true)))
    end
end
(s::Standardizer)(x::Vector) = (x .- s.avg) ./ (s.std + eps(s.std))

mutable struct StepStandardize{S,F} <: AbstractStep
    selections::S
    params::Union{Nothing, Dict{Symbol,Standardizer}}
    robust::Bool
    trained::Bool
    skip::Bool
    prehook::F
end

function transform!(s::StepStandardize, df)
    for (k,f) in s.params
        df[k] = f(df[k])
    end
    df
end
function transform(s::StepStandardize, df)
    transform!(s, copy(df))
end
function fit!(s::StepStandardize, df)
    s.params = Dict(col => Standardizer(s.prehook(df[col])) for col in getselectionkeys(df, s.selections))
    s.trained = true
end
function step_standardize!(r::Recipe, s...; robust::Bool=false, skip::Bool=false, prehook=identity)
    push!(r.steps,
          StepStandardize([s...],
                          nothing,
                          robust,
                          false,
                          skip,
                          prehook))
end


using SparseArrays, LightXML

function osm2lar(
        filename::String,
        doCrop::Bool = false,
        doFrame::Bool = false
    )::Lar.Model
    xdoc  = parse_file(filename)
    xroot = root(xdoc)

    cropBox   = get_elements_by_tagname(xroot, "bounds")[1]
    node_list = get_elements_by_tagname(xroot, "node"  )
    way_list  = get_elements_by_tagname(xroot, "way"   )

    #=== EXTRACT VERTEX SET AND VERTEX POSITION MAP ===#

    Vidxs = [parse(Int, attribute(v, "id")) for v in node_list]
    Vmax  = max(Vidxs...)
    Vmap  = SparseArrays.spzeros(Int, Vmax)
    Vnum  = length(Vidxs)
    for i = 1 : Vnum   Vmap[Vidxs[i]] = i   end

    #=== EXTRACT (X,Y) AND RESCALE POINTS IN [0, 1] ===#
    lons  = [parse(Float64, attribute(v, "lon")) for v in node_list]
    lmin  = parse(Float64, attribute(cropBox, "minlon"))
    lmax  = parse(Float64, attribute(cropBox, "maxlon"))
    lons  = (lons .- lmin) ./ (lmax - lmin)

    lats  = [parse(Float64, attribute(v, "lat")) for v in node_list]
    lmin  = parse(Float64, attribute(cropBox, "minlat"))
    lmax  = parse(Float64, attribute(cropBox, "maxlat"))
    lats  = (lats .- lmin) ./ (lmax - lmin)

    #=== BUILD LAR MODEL ===#

    model = Lar.Model([lons'; lats'])

    for street in way_list
        nodes = [
                 Vmap[parse(Int, attribute(n, "ref"))]
                 for n in get_elements_by_tagname(street, "nd")
                ]
        for i = 1 : length(nodes) - 1
            newcell = SparseArrays.spzeros(Int8, Vnum)
            newcell[nodes[  i]] = 1
            newcell[nodes[i+1]] = 1
            Lar.addModelCells!(model, 1, convert(Lar.ChainOp, newcell'))
        end
    end

    if doCrop
        Lar.mapBoxCrop!(model, [0.0 1.0;0.0 1.0], doFrame)
    end
    return model
end

function mapBoxCrop!(
        model::Lar.Model,
        cropBox::Lar.Points,
        doFrame::Bool
    )::Nothing

    toDelV = convert(Array{Int64,1},
        [i  for i = 1 : size(model, 0, 2)
            if model.G[1, i] <= cropBox[1, 1]
            || model.G[1, i] >= cropBox[1, 2]
            || model.G[2, i] <= cropBox[2, 1]
            || model.G[2, i] >= cropBox[2, 2]
        ]
    )
    Lar.deleteModelVertices!(model, toDelV)

    if doFrame
        idx = size(model, 0, 2) + 1
        Lar.addModelVertices!(model, [
            cropBox[1, 1] cropBox[1, 1] cropBox[1, 2] cropBox[1, 2]
            cropBox[2, 1] cropBox[2, 2] cropBox[2, 2] cropBox[2, 1]
        ])
        Lar.addModelCells!(model, 1, [
            [idx, idx+1], [idx+1, idx+2], [idx+2, idx+3], [idx+3, idx]
        ])
    end
    return
end


using Pkg
pkg"activate ."
using NNSparseCoding; const N = NNSparseCoding;
using Test
using Plots;
theme(:dark);
using BenchmarkTools
using StatsBase,Statistics
using MultivariateStats

function scatty(M1,M2)
  scatter(M1[:], M2[:]; leg = false, ratio = 1)
  return plot!([extrema(M1)...].* 1.1, identity; linewidth = 3)
end
# function to plot one or more black/mytransfe images
function showbw(mat::Matrix{<:Real},_title="")
  heatmap(mat, color=:grays,ratio=1,cbar=false,axis=false,title=_title)
end
function showbw_many(title::String,mats...)
  mats = [mats...] # convert tuple to array
  padval=maximum(maximum,mats)
  framesize = size(mats[1],1)
  padmat = fill(padval, framesize,round(Int,0.1framesize))
  padmats = repeat([padmat,],inner=length(mats))
  matspads = permutedims(hcat(mats,padmats))[:]
  imgrow = hcat(matspads[1:end-1]...)
  showbw(imgrow , title)
end
##
#
# const dirimg = abspath(@__DIR__,"..","data")
# const sizesq=315
# all_images = N.read_natural_images(dirimg,sizesq ; verbose=true)
#
# const npatches = 1000
# const patch_size = 40
#
# patches = N.get_random_patches(npatches,all_images,patch_size)
#
#
# patchmat = N.merge_mats(patches)
#
# mean(patchmat;dims=2) |> extrema
# std(patchmat;dims=2) |> extrema
# cov(patchmat;dims=2) |> heatmap
#
#
# ##
# showbw_many("done!",sample(patches,5)...)
##

n = 80
d = n
k = 500
A,S = N.initialize_mats(n,d,k)
map!(s-> rand() < 0.95 ? 0.0 : 3s ,S,S)
X = A*S
broadcast!(-,X,X,mean(X;dims=2))

##
mytransf,Mp,Mm = N.preprocess_plus_minus(X,N.PreprocessPCA(d))

##

lambda=0.1
nsteps=10_000
stepsize=0.2
optthresh=1E-4
op = N.nnsc_optimize(Mp,d,lambda ;
  step_zero=stepsize, nstepmax=nsteps,
  opt_threshold=optthresh,verbose=true)

om = N.nnsc_optimize(Mm,d,lambda ;
  step_zero=stepsize, nstepmax=nsteps,
  opt_threshold=optthresh,verbose=true)

##

scatty(X,N.revert_plus_minus(op.AS,om.AS,mytransf))

extrema( X .- N.revert_plus_minus(op.AS,om.AS,mytransf))

###
mytransf,Mp,Mm = N.preprocess_plus_minus(X,N.PreprocessWhiten)

##

lambda=0.1
nsteps=10_000
stepsize=0.2
optthresh=1E-4
op = N.nnsc_optimize(Mp,d,lambda ;
  step_zero=stepsize, nstepmax=nsteps,
  opt_threshold=optthresh,verbose=true)

om = N.nnsc_optimize(Mm,d,lambda ;
  step_zero=stepsize, nstepmax=nsteps,
  opt_threshold=optthresh,verbose=true)

##

scatty(X,N.revert_plus_minus(op.AS,om.AS,mytransf))

extrema( X .- N.revert_plus_minus(op.AS,om.AS,mytransf))

boh = N.revert_plus_minus(op.S,zero(op.S),mytransf)


# sample_target_mismatch-multivariate_uniform_vs_beta-kernel_vs_graph
#
# This script compares the values of the multivariate graph Stein
# discrepancy and kernel discrepancy obtained when samples are drawn i.i.d.
# from their targets.

include("experiment_utils.jl")

# we see first if we want to parallelize
@parseintcli parallel "p" "parallel" "false"
if parallel == "true"
    addprocs(CPU_CORES - 1)
end

using Distributions: Beta
using SteinDistributions:
    SteinUniform,
    gradlogdensity
using SteinDiscrepancy:
    SteinChampionLenardMillsKernel,
    ksd,
    gsd

# random number generator seed
# we allow this to be set on the commandline
@parseintcli seed "s" "seed" 7
# set the target distribution
@parsestringcli distname "t" "target" "uniform"
# set the discrepancy
@parsestringcli discrepancytype "d" "discrepancy" "championlenardmills"
# Specify dimension of sampled vectors
d = 2
# Specify a maximum number of samples
maxn = 200000
# uniform target
uniform_target = SteinUniform(d)
# beta target
beta_target = Beta(3, 3)
# the graphsolver
# setup kernel (if necessary)
graphsolver = nothing
kernel = nothing
if discrepancytype == "championlenardmills"
    kernel = SteinChampionLenardMillsKernel([0.0,1.0], 2)
elseif discrepancytype == "graph-gurobi"
    graphsolver = "gurobi"
elseif discrepancytype == "graph-clp"
    graphsolver = "clp"
end
# Draw samples from target distribution
if distname == "beta"
    target = beta_target
    X = @setseed reshape(rand(beta_target, d*maxn), maxn, d)
elseif distname == "uniform"
    target = uniform_target
    X = @setseed rand(uniform_target, maxn)
end
# define gradlogp
function gradlogp(x::Array{Float64,1})
    gradlogdensity(target, x)
end
# Sample sizes at which optimization problem will be solved
ns = vcat(10:1:min(10,maxn),
          100:100:min(1000,maxn),
          2000:1000:min(10000,maxn),
          20000:10000:min(100000,maxn),
          200000:100000:maxn)

# Solve optimization problem at each sample size
for i in ns
    @printf("[Beginning n=%d]:\n", i)
    # Create program with data subset
    Xi = X[1:i,:]
    discrepancy = nothing
    solvetime = nothing
    if kernel != nothing
        @printf("[Beginning kernel computations n=%d]\n", i)
        result = ksd(points=Xi, gradlogdensity=gradlogp, kernel=kernel)
        discrepancy = sqrt(result.discrepancy2)
        solvetime = result.solvetime
    elseif discrepancytype in ["graph-clp", "graph-gurobi"]
        @printf("[Beginning graph solver n=%d]\n", i)
        result = gsd(points=Xi,
                     gradlogdensity=gradlogp,
                     supportlowerbounds=zeros(d),
                     supportupperbounds=ones(d),
                     solver=graphsolver)
        discrepancy = sum(result.objectivevalue)
        solvetime = sum(result.solvetime)
    end
    # save data
    instance_data = Dict{Any, Any}(
        "target" => distname,
        "n" => i,
        "d" => d,
        "seed" => seed,
        "discrepancytype" => discrepancytype,
        "discrepancy" => discrepancy,
        "nprocs" => nprocs(),
        "solvetime" => solvetime,
    )
    save_json(
        instance_data;
        dir="sample_target_mismatch-multivariate_uniform_vs_beta-kernel_vs_graph",
        target=distname,
        discrepancytype=discrepancytype,
        n=i,
        d=d,
        ncores=nprocs(),
        seed=seed
    )
end

@printf("COMPLETE!")


using Plots

# enable plotting by default
if !@isdefined do_visu; do_visu = true end

@views function diffusion_1D_2procs(; do_visu=true)
    # Physics
    Hl  = 10.0   # left  H
    Hr  = 1.0    # right H
    λ   = 1.0    # diffusion coeff
    nt  = 200    # number of time steps
    # Numerics
    nx  = 32     # number of local grid points
    dx  = 1.0    # cell size
    # Derived numerics
    dt  = dx^2/λ/2.1
    # Initial condition
    HL  = Hl*ones(nx)
    HR  = Hr*ones(nx)
    H   = [HL[1:end-1]; HR[2:end]]
    Hg  = copy(H)
    # Time loop
    for it = 1:nt
        # Compute physics locally
        HL[2:end-1] .= HL[2:end-1] + dt*λ*diff(diff(HL)/dx)/dx
        HR[2:end-1] .= HR[2:end-1] + dt*λ*diff(diff(HR)/dx)/dx
        # Update boundaries (MPI)
        HL[end] = HR[2]
        HR[1]   = HL[end-1]
        # Global picture
        H .= [HL[1:end-1]; HR[2:end]]
        # Compute physics globally (check)
        Hg[2:end-1] .= Hg[2:end-1] + dt*λ*diff(diff(Hg)/dx)/dx
        # Visualise
        if do_visu
            fontsize = 12
            plot(Hg, legend=false, linewidth=0, markershape=:circle, markersize=5, yaxis=font(fontsize, "Courier"), xaxis=font(fontsize, "Courier"), titlefontsize=fontsize, titlefont="Courier")
            display(plot!(H, legend=false, linewidth=3, framestyle=:box, xlabel="Lx", ylabel="H", title="diffusion (it=$(it))"))
        end
    end
    return
end

diffusion_1D_2procs(; do_visu=do_visu)


#TODO: give a new name
immutable LsqFitResult{T}
	# simple type container for now, but can be expanded later
	dof::Int
	param::Vector{T}
	resid::Vector{T}
	jacobian::Matrix{T}
end

#TODO: update this doc
"""
	lmfit(f, p0; kwargs...)

this is a convenience function for the curve_fit() methods
which assume f(p) is the cost functionj i.e. the residual of a
model where

  model(xpts, params...) = ydata + error (noise)

this minimizes f(p) using a least squares sum of squared error:

  sse = sum(f(p)^2)

This is currently embedded in levelberg_marquardt()
which calls sumabs2

# Returns
p, f(p), g(p) where
  p    : best fit parameters
  f(p) : function evaluated at best fit p, (weighted) residuals
  g(p) : estimated Jacobian at p (Jacobian with respect to p)
"""
function lmfit(f::Function, p0; kwargs...)
	# construct Jacobian function, which uses finite difference method
	g = Calculus.jacobian(f)

	results = levenberg_marquardt(f, g, p0; kwargs...)
	p = results.minimum
	resid = f(p)
	dof = length(resid) - length(p)
	return LsqFitResult(dof, p, f(p), g(p))
end

function curve_fit(model::Function, xpts, ydata, p0; kwargs...)
	# construct the cost function
	f(p) = model(xpts, p) - ydata
	lmfit(f,p0; kwargs...)
end

function curve_fit(model::Function, xpts, ydata, wt::Vector, p0; kwargs...)
	# construct a weighted cost function, with a vector weight for each ydata
	# for example, this might be wt = 1/sigma where sigma is some error term
	f(p) = wt .* ( model(xpts, p) - ydata )
	lmfit(f,p0; kwargs...)
end

function curve_fit(model::Function, xpts, ydata, wt::Matrix, p0; kwargs...)
	# as before, construct a weighted cost function with where this
	# method uses a matrix weight.
	# for example: an inverse_covariance matrix

	# Cholesky is effectively a sqrt of a matrix, which is what we want
	# to minimize in the least-squares of levenberg_marquardt()
	# This requires the matrix to be positive definite
	u = chol(wt)

	f(p) = u * ( model(xpts, p) - ydata )
	lmfit(f,p0; kwargs...)
end

function estimate_covar(fit::LsqFitResult)
	# computes covariance matrix of fit parameters
	r = fit.resid
	J = fit.jacobian

	# mean square error is: standard sum square error / degrees of freedom
	mse = sumabs2(r) / fit.dof

	# compute the covariance matrix from the QR decomposition
	Q,R = qr(J)
	Rinv = inv(R)
	covar = Rinv*Rinv'*mse
end

function estimate_errors(fit::LsqFitResult, alpha=0.95)
	# computes (1-alpha) error estimates from
	#   fit   : a LsqFitResult from a curve_fit()
	#   alpha : alpha percent confidence interval, (e.g. alpha=0.95 for 95% CI)
	covar = estimate_covar(fit)

	# then the standard errors are given by the sqrt of the diagonal
	std_error = sqrt(diag(covar))

	# scale by quantile of the student-t distribution
	dist = TDist(fit.dof)
	std_error *= quantile(dist, alpha)
end