text:
export expmv, expmv!

function default_anorm(A)
    try
        opnorm(A, Inf)
    catch err
        if err isa MethodError
            @warn "opnorm($(typeof(A)), Inf) is not defined, fall back to using `anorm = 1.0`.
To suppress this warning, please specify the `anorm` keyword manually."
            1.0
        else
            throw(err)
        end
    end
end

"""
    expmv{T}(t, A, vec; [tol], [m], [norm], [anorm])

Calculate matrix exponential acting on some vector, ``w = e^{tA}v``,
using the Krylov subspace approximation.

See R.B. Sidje, ACM Trans. Math. Softw., 24(1):130-156, 1998
and http://www.maths.uq.edu.au/expokit
"""
function expmv( t::Number,
                   A, vec::Vector{T};
                   tol::Real=1e-7,
                   m::Int=min(30, size(A, 1)),
                   norm=LinearAlgebra.norm, anorm=default_anorm(A)) where {T}

    result = convert(Vector{promote_type(eltype(A), T, typeof(t))}, copy(vec))
    expmv!(t, A, result; tol=tol, m=m, norm=norm, anorm=anorm)
    return result
end

expmv!( t::Number,
           A, vec::Vector{T};
           tol::Real=1e-7,
           m::Int=min(30,size(A,1)),
           norm=LinearAlgebra.norm, anorm=default_anorm(A)) where {T} = expmv!(vec, t, A, vec; tol=tol, m=m, norm=norm, anorm=anorm)

function expmv!(w::Vector{T}, t::Number, A, vec::Vector{T};
                   tol::Real=1e-7, m::Int=min(30,size(A,1)), norm=LinearAlgebra.norm, anorm=default_anorm(A)) where {T}

    if size(vec,1) != size(A,2)
        error("dimension mismatch")
    end

    # safety factors
    gamma = 0.9
    delta = 1.2

    btol = 1e-7     # tolerance for "happy-breakdown"
    maxiter = 10    # max number of time-step refinements

    rndoff= anorm*eps()

    # estimate first time-step and round to two significant digits
    beta = norm(vec)
    r = 1/m
    fact = (((m+1)/exp(1.0))^(m+1))*sqrt(2.0*pi*(m+1))
    tau = (1.0/anorm)*((fact*tol)/(4.0*beta*anorm))^r
    tau = round(tau, sigdigits=2)

    # storage for Krylov subspace vectors
    vm = Array{typeof(w)}(undef,m+1)
    for i=1:m+1
        vm[i]=similar(w)
    end
    hm = zeros(T,m+2,m+2)

    tf = abs(t)
    tsgn = sign(t)
    tk = zero(tf)

    copyto!(w, vec)
    p = similar(w)
    mx = m
    while tk < tf
        tau = min(tf-tk, tau)

        # Arnoldi procedure
        # vm[1] = v/beta
        rmul!(copyto!(vm[1],w),1/beta)
        mx = m
        for j=1:m
            # p[:] = A*vm[j]
            mul!(p, A, vm[j])

            for i=1:j
                hm[i,j] = dot(vm[i], p)
                # p[:] = p - hm[i,j]*vm[i]
                p = axpy!(-hm[i,j], vm[i], p)
            end
            s = norm(p)

            if s < btol # happy-breakdown
                tau = tf - tk
                err_loc = btol

                # F = expm(tsgn*tau*hm[1:j,1:j])
                # F = expm!(scale(tsgn*tau,view(hm,1:j,1:j)))
                F = exp!(tsgn*tau*view(hm,1:j,1:j))

                fill!(w, zero(T))
                for k=1:j
                    # w[:] = w + beta*vm[k]*F[k,1]
                    w = axpy!(beta*F[k,1], vm[k], w)
                end
                mx = j
                break
            end
            hm[j+1,j] = s

            # vm[j+1] = p/hm[j+1,j]
            rmul!(copyto!(vm[j+1],p),1/hm[j+1,j])
        end
        hm[m+2,m+1] = one(T)
        (mx != m) || (avnorm = norm(mul!(p,A,vm[m+1])))

        # propagate using adaptive step size
        iter = 1
        while (iter < maxiter) && (mx == m)

            # F = expm(tsgn*tau*hm)
            # F = expm!(scale(tsgn*tau,hm))
            F = exp!(tsgn*tau*hm)

            # local error estimation
            err1 = abs( beta*F[m+1,1] )
            err2 = abs( beta*F[m+2,1] * avnorm )

            if err1 > 10*err2	# err1 >> err2
                err_loc = err2
                r = 1/m
            elseif err1 > err2
                err_loc = (err1*err2)/(err1-err2)
                r = 1/m
            else
                err_loc = err1
                r = 1/(m-1)
            end

            # time-step sufficient?
            if err_loc <= delta * tau * (tau*tol/err_loc)^r
                fill!(w, zero(T))
                for k=1:m+1
                    # w[:] = w + beta*vm[k]*F[k,1]
                    w = axpy!(beta*F[k,1], vm[k], w)
                end

                break
            end
            tau = gamma * tau * (tau*tol/err_loc)^r # estimate new time-step
            tau = round(tau, sigdigits=2) # round to 2 signiﬁcant digits
                                 # to prevent numerical noise
            iter = iter + 1
        end
        if iter == maxiter
            # TODO, here an exception should be thrown, but which?
            error("Number of iterations exceeded $(maxiter). Requested tolerance might be to high.")
        end

        beta = norm(w)
        tk = tk + tau

        tau = gamma * tau * (tau*tol/err_loc)^r # estimate new time-step
        tau = round(tau, sigdigits=2) # round to 2 signiﬁcant digits
                             # to prevent numerical noise
        err_loc = max(err_loc,rndoff)

        fill!(hm, zero(T))
    end

    return w
end # expmv!


text:
module MNE

using PyCall,Dates

mne = PyObject(nothing)

function __init__()
    global mne

    mne = pyimport("mne")
end

include("structs.jl")

end # module

ext: jl


text:
# Client for SCAFFOLD with linear local model
mutable struct ScaffoldClient{T1<:Int64, T2<:Float64, T3<:SparseMatrixCSC{Float64, Int64}, T4<:Matrix{Float64}, T5<:Vector{Int64}, T6<:Function} <: AbstractClient
    id::T1                                  # client index
    Xtrain::T3                              # training data
    XtrainT::T3                             # Row copy
    Ytrain::T5                              # training label
    W::T4                                   # (model) primal variable
    ΔW::T4                                  # Update of model
    C::T4                                   # Control variate
    ΔC::T4                                  # Update of control variate
    globalC::T4                             # Global control variate
    lr::T2                                  # learning rate
    λ::T2                                   # regularization parameter
    numLocalEpochs::T1                      # number of local steps
    function ScaffoldClient(id::Int64, Xtrain::SparseMatrixCSC{Float64, Int64}, Ytrain::Vector{Int64}, config::Dict{String,Real})
        num_classes = config["num_classes"]
        λ = config["lambda"]
        learning_rate = config["learning_rate"]
        numLocalEpochs = config["numLocalEpochs"]
        d = size(Xtrain, 2)
        W = zeros(Float64, d, num_classes)
        ΔW = zeros(Float64, d, num_classes)
        C = zeros(Float64, d, num_classes)
        ΔC = zeros(Float64, d, num_classes)
        globalC = zeros(Float64, d, num_classes)
        XtrainT = copy(Xtrain')
        new{Int64, Float64, SparseMatrixCSC{Float64, Int64}, Matrix{Float64}, Vector{Int64}, Function}(id, Xtrain, XtrainT, Ytrain, W, ΔW, C, ΔC, globalC, learning_rate, λ, numLocalEpochs)
    end
end

# Model update on local device
function update!(
    client::ScaffoldClient
)
    # Implement K epochs SGD, using lazy update to avoid dense update from the regularization.
    n, d = size(client.Xtrain)
    _, K = size(client.W)
    lr = client.lr
    Wold = copy(client.W)
    C = client.C
    globalC = client.globalC
    λ = client.λ
    hitTime = zeros(Int64, d, K)
    perm = collect(1:n)
    for epoch = 1:client.numLocalEpochs
        fill!(hitTime, 0)
        shuffle!(perm)
        timeStep = 1
        for i in perm
            g = getStochasticGrad(client.XtrainT, client.Ytrain, client.W, i)
            I, J, V = findnz(g)
            for j = 1:length(I)
                idx1 = I[j]
                idx2 = J[j]
                # Lazy update
                delay = timeStep-hitTime[idx1, idx2]
                client.W[idx1, idx2] *= (1 - lr*(λ))^delay
                client.W[idx1, idx2] += lr*delay*( globalC[idx1, idx2] - C[idx1, idx2] )
                client.W[idx1, idx2] -= lr*V[j]
                # Update hitTime to the current time
                hitTime[idx1, idx2] = timeStep
            end
            timeStep += 1
        end
        # Lazy update for staled coordinates
        timeStep -= 1
        for j = 1:d
            for k = 1:K
                if hitTime[j, k] < timeStep
                    delay = timeStep-hitTime[j, k]
                    client.W[j, k] *= (1 - lr*λ)^delay
                    client.W[j, k] += lr*delay*( globalC[j, k] - C[j, k] )
                end
            end
        end
    end
    # Update other variables
    client.ΔW .= client.W - Wold
    client.ΔC .= getGradient(client.Xtrain, client.XtrainT, client.Ytrain, client.W, λ) - C
    return nothing
end

# Get objective value
function getObjValue(
    client::ScaffoldClient
)
    objValue = obj(client.Xtrain, client.Ytrain, client.W, client.λ)
    return objValue
end

# Client for SCAFFOLD with multilayer perceptron local model
mutable struct ScaffoldClientNN{T1<:Int64, T2<:Float64, T3<:SparseMatrixCSC{Float64, Int64}, T4<:Zygote.Params, T5<:Flux.OneHotArray, T6<:Function, T7<:Flux.Chain} <: AbstractClient
    id::T1                                  # client index
    Xtrain::T3                              # training data
    XtrainT::T3                             # Row copy
    Ytrain::T5                              # training label
    W::T7                                   # (model) primal variable
    ΔW::T4                                  # Update of model
    C::T4                                   # Control variate
    ΔC::T4                                  # Update of control variate
    globalC::T4                             # Global control variate
    lr::T2                                  # learning rate
    λ::T2                              # L2 regularization parameter
    numLocalEpochs::T1                       # number of local steps
    function ScaffoldClientNN(id::Int64, Xtrain::SparseMatrixCSC{Float64, Int64}, Ytrain::Vector{Int64}, dim::Int64, config::Dict{String,Real})
        num_classes = config["num_classes"]
        Ytrain = Flux.onehotbatch(Ytrain, 1:num_classes)
        λ = config["lambda"]
        learning_rate = config["learning_rate"]
        numLocalEpochs = config["numLocalEpochs"]
        W = Chain( Dense(780, dim, relu, bias=false), Dense(dim, 10, bias=false), NNlib.softmax)
        ΔW = deepcopy( params(W) )
        C = deepcopy( params(W) )
        ΔC = deepcopy( params(W) )
        globalC = deepcopy( params(W) )
        for j = 1:length(ΔW)
            fill!(ΔW[j], 0.0)
            fill!(C[j], 0.0)
            fill!(ΔC[j], 0.0)
            fill!(globalC[j], 0.0)
        end
        XtrainT = copy(Xtrain')
        new{Int64, Float64, SparseMatrixCSC{Float64, Int64}, Zygote.Params, Flux.OneHotArray, Function, Flux.Chain}(id, Xtrain, XtrainT, Ytrain, W, ΔW, C, ΔC, globalC, learning_rate, λ, numLocalEpochs)
    end
end

# Model update on local device
function update!(
    client::ScaffoldClientNN
)
    lr = client.lr
    W = client.W
    λ = client.λ
    Xt = client.XtrainT
    X = client.Xtrain
    Y = client.Ytrain
    # data
    data = Flux.Data.DataLoader((Xt, Y), batchsize=128, shuffle=true)
    # Store Wold
    Wold = deepcopy( params(W) )
    # loss
    sqnorm(w) = sum(abs2, w)
    loss(x, l) = Flux.crossentropy(W(x), l) + (λ/2) * sum(sqnorm, params(W)) - dot(params(W), client.C) + dot(params(W), client.globalC)
    # optimizer
    opt = Descent(lr)
    # train
    for t = 1:client.numLocalEpochs
        # @printf "   epoch: %d, obj: %4.4e, acc: %4.4e\n" t obj(Xt, Y, W, λ) accuracy(X, Y, W)
        Flux.train!(loss, params(W), data, opt)
    end
    # Implement option II from the Scaffold paper.
    n = size(X, 1)
    K = client.numLocalEpochs*floor(Int64, n/128)
    for j = 1:length( client.ΔC )
        client.ΔC[j] .= (-client.globalC[j] + 1/(K*lr)*( Wold[j] - params(W)[j] ) )
        client.C[j] .= client.C[j] + client.ΔC[j]
        client.ΔW[j] .= ( params(W)[j] - Wold[j] )
    end
end


text:
#==============================================================================#
#
#          Structs defining buffers used for in-place computations
#
#==============================================================================#

"""
    AbstractBuffer{T} <: AbstractArray{T,1}

Types inheriting from `AbstractBuffer` define buffers used for various in-place
computations.
"""
abstract type AbstractBuffer end

"""
    struct StandardEulerBuffer{Tb,Tσ,Tdw} <: AbstractBuffer
        b::Tb
        y::Tb
        σ::Tσ
        dW::Tdw
    end

Standard buffer for the Euler-Maruyama simulations. The intermediary data is
stored in `b`, `σ`, `y` and `dW`.
"""
struct StandardEulerBuffer{Tb,Tσ,Tdw} <: AbstractBuffer
    b::Tb
    y::Tb
    σ::Tσ
    dW::Tdw
end

function StandardEulerBuffer(b::Tb, σ::Tσ, dW::Tdw) where {Tb, Tσ, Tdw}
    y = deepcopy(b)
    StandardEulerBuffer{Tb, Tσ, Tdw}(b, y, σ, dW)
end

function StandardEulerBuffer{K}(P::DiffusionProcess{T,DP,DW}) where {K,T,DP,DW}
    # dW always falls on defaults or not?
    dW = zero(K, DW, ismutable(K)) # zero(P, Val(:wiener))
    b = zero(K, DP, ismutable(K))
    σ = _init_mat_for_buffer(
        Val(diagonaldiff(P)),
        Val(sparsediff(P)),
        K,
        tuple(dimension(P)...)
    )
    StandardEulerBuffer(b, σ, dW)
end

function _init_mat_for_buffer(
        ::Val{false}, # diagonal
        ::Val{false}, # sparse
        ::Type{K}, # eltype
        dims
    ) where K
    zeros(eltype(K), dims)
end

function _init_mat_for_buffer(
        ::Val{true}, # diagonal
        ::Any, # sparse
        ::Type{K}, # eltype
        dims
    ) where K
    Diagonal{eltype(K)}(zeros(eltype(K), dims[1]))
end

function _init_mat_for_buffer(
        ::Val{false}, # diagonal
        ::Val{true}, # sparse
        ::Type{K}, # eltype
        dims
    ) where K
    spzeros(eltype(K), dims...)
end

"""
    struct LinearDiffBuffer{Tb,Tσ,Tdw,TB} <: AbstractBuffer
        b::Tb
        y::Tb
        σ::Tσ
        dW::Tdw
        B::TB
    end

A buffer for Euler-Maruyama simulations of linear diffusions. Almost the same
as `StandardEulerBuffer`, but contains additional space for an intermediate
construction of a matrix `B`.
"""
struct LinearDiffBuffer{Tb,Tσ,Tdw,TB} <: AbstractBuffer
    b::Tb
    y::Tb
    σ::Tσ
    dW::Tdw
    B::TB
end

function LinearDiffBuffer(b::Tb, σ::Tσ, dW::Tdw, B::TB) where {Tb, Tσ, Tdw, TB}
    y = deepcopy(b)
    LinearDiffBuffer{Tb, Tσ, Tdw, TB}(b, y, σ, dW, B)
end

function LinearDiffBuffer{K}(P::LinearDiffusion{T,DP,DW}) where {K,T,DP,DW}
    dW = zero(K, DW, ismutable(K))
    b = zero(K, DP, ismutable(K))
    σ = _init_mat_for_buffer(
        Val(diagonaldiff(P)),
        Val(sparsediff(P)),
        K,
        tuple(dimension(P)...)
    )
    B = _init_mat_for_buffer(
        Val(diagonalBmat(P)),
        Val(sparseBmat(P)),
        K,
        (dim_process(P), dim_process(P))
    )
    LinearDiffBuffer(b, σ, dW, B)
end

"""
    remove_curly(::Type{K}) where K

Utility function that removes all type-specifiers listed in the curly brackets.

# Examples
```julia-repl
julia> remove_curly(Array{Float64,1})
Array
```
"""
@generated function remove_curly(::Type{K}) where K
    name_without_curly = Meta.parse(string(K)).args[1]
    :( $name_without_curly )
end

"""
    get_curly(::Type{K}) where K

Utility function that returns a tuple with all type-specifiers listed in the
curly brackets.

# Examples
```julia-repl
julia> remove_curly(Array{Float64,1})
(:Float64, 1)
```
"""
@generated function get_curly(::Type{K}) where K
    curly_arguments = Meta.parse(string(K)).args[2:end]
    :( Tuple($curly_arguments) )
end


text:
# functions for calculating various preconditioners
using Base.LinAlg.BLAS  # trmv!

### NewtonBDiagPC
#------------------------------------------------------------------------------
# AbstractPC Interface
function calcPC(pc::NewtonBDiagPC, mesh::AbstractMesh, sbp::AbstractOperator,
                eqn::AbstractSolutionData, opts::Dict, ctx_residual, t)

  setPCCtx(pc, mesh, sbp, eqn, opts, ctx_residual, t)  # unnecessary?

  pc.evalJacobian(mesh, sbp, eqn, opts, pc.assem)
  factorBDiagPC(pc, mesh, sbp, eqn, opts)
end

function applyPC(pc::NewtonBDiagPC, mesh::AbstractMesh, sbp::AbstractOperator,
                 eqn::AbstractSolutionData, opts::Dict, t, b::AbstractVector,
                 x::AbstractVector)

  applyBDiagPC(pc, mesh, sbp, eqn, opts, b, x)

  return nothing
end

function applyPCTranspose(pc::NewtonBDiagPC, mesh::AbstractMesh, sbp::AbstractOperator,
                 eqn::AbstractSolutionData, opts::Dict, t, b::AbstractVector,
                 x::AbstractVector)

  applyBDiagPC(pc, mesh, sbp, eqn, opts, b, x, trans=true)

  return nothing
end

# we could define applyPCTranspose, but this PC is ineffective so don't bother

#------------------------------------------------------------------------------
# block diagonal preconditioner implementation

"""
  This function factors the volume preconditioner.  It uses Lapack to do an
  LU factorization with partial pivoting.  The user is allowed to modify
  `pc.diag_jac`_jac after calling [`calcBDiagPC`](@ref)
  but before calling this routine.

  **Inputs**

   * mesh
   * sbp
   * eqn
   * opts

  **Inputs/Outputs**

   * newton_data: vol_prec is updated. `pc.diag_jac` is factored in-place and
                 `pc.ipiv` is overwritten with the permutation

"""
function factorBDiagPC(pc::NewtonBDiagPC, mesh, sbp, eqn, opts)

  for i=1:mesh.numEl
    jac_i = sview(pc.diag_jac.A, :, :, i)
    ipiv_i = sview(pc.ipiv, :, i)

    # call Lapack here
    info = getrf!(jac_i, ipiv_i)
    @assert info == 0
  end

  pc.is_factored = true

  return nothing
end

"""
  Apply the preconditioner, ie. do inv(A)*x = b, where A is calculated by
  [`calcBDiagPC`](@ref).

  This is not likely to be a sensible preconditioner for Contiuous Galerkin
  discretizations.

  **Inputs**

   * pc: a [`BDiagPC `](@ref) object, already factored by [`factorBDiagPC`](@ref).
   * mesh
   * sbp
   * eqn
   * opts
   * x: the vector to multiply against

  **Inputs/Outputs**

   * b: the output vector (overwritten)

  **Keyword Arguments**

   * trans: if true, apply the transposed operation, default false

"""
function applyBDiagPC(pc::NewtonBDiagPC, mesh, sbp, eqn, opts, x::AbstractVector, b::AbstractVector; trans::Bool=false)

  # we need to do inv(A)*x = b --> solve A*b = x using the factorization A.

  @assert pc.is_factored
  if trans
    tchar = 'T'
  else
    tchar = 'N'
  end

  bs = pc.bs
  workvec = zeros(Float64, bs)  # hold values passed into LAPACK

  fill!(b, 0)

  for i=1:mesh.numEl
    jacf_i = sview(pc.diag_jac.A, :, :, i)
    ipiv_i = sview(pc.ipiv, :, i)

    getValues(mesh, x, i, workvec)

    # call Lapack GETRS to solve for b (in-place)
    getrs2!(tchar, jacf_i, ipiv_i, workvec)

    setValues(mesh, workvec, i, b)

  end  # end loop i

  return nothing
end

# these names are too generic to export
import Jacobian: getValues, setValues

"""
  Computes the action of the inverse of the preconditioner on a vector.
  This is not required by the AbstractPC interface, but is useful for
  constructing smoothers.

  Note that the preconditioner is defined as the inverse of some
  some approximation to the Jacobian. so the inverse of the preconditioner
  is the approximate Jacobian itself.

  **Inputs**

   * pc: [`NewtonBDiagPC`](@ref).  Can be factored or not
   * mesh
   * sbp
   * eqn
   * opts
   * x: vector to multiply against

  **Inputs/Outputs**

   * b: vector to overwrite with the result.

  **Keyword Arguments**

   * trans: if true, apply the transpose, default false

   Aliasing Restrictions: `x` and `b` cannot alias
"""
function applyBDiagPCInv(pc::NewtonBDiagPC, mesh, sbp, eqn, opts,
                         x::AbstractVector, b::AbstractVector;
                         zero_output=true, trans::Bool=false)

  if !pc.is_factored
    diagMatVec(pc.diag_mat, mesh, x, b, zero_output=zero_output, trans=trans)
  else

    if zero_output
      fill!(b, 0)
    end

    _trans::Bool = trans

    bs = pc.bs
    workvec = zeros(Float64, bs)  # hold values passed into LAPACK

    for i=1:mesh.numEl
      jacf_i = sview(pc.diag_jac.A, :, :, i)
      ipiv_i = sview(pc.ipiv, :, i)

      getValues(mesh, x, i, workvec)

      if trans
        laswp!(workvec, 1, bs, ipiv_i)
        trmv!('L', 'T', 'U', jacf_i, workvec)
        trmv!('U', 'T', 'N', jacf_i, workvec)
      else
        # do b = P*L*U*x
        trmv!('U', 'N', 'N', jacf_i, workvec)
        trmv!('L', 'N', 'U', jacf_i, workvec)
        applyIpiv!(ipiv_i, workvec)
      end

      setValues(mesh, workvec, i, b)
    end  # end loop i

  end  # end if

  return nothing
end

### NewtonBJacobiPC
#------------------------------------------------------------------------------
# AbstractPC Interface
function calcPC(pc::NewtonBJacobiPC, mesh::AbstractMesh, sbp::AbstractOperator,
                eqn::AbstractSolutionData, opts::Dict, ctx_residual, t)

  calcPC(pc.diag_pc, mesh, sbp, eqn, opts, ctx_residual, t)
  setPCCtx(pc, mesh, sbp, eqn, opts, ctx_residual, t)  # unnecessary?
end

"""
  Applies the `NewtonBJacobiPC`.  Note that the initial value of `x` is
  used as the initial guess.
"""
function applyPC(pc::NewtonBJacobiPC, mesh::AbstractMesh, sbp::AbstractOperator,
                 eqn::AbstractSolutionData, opts::Dict, t, b::AbstractVector,
                 x::AbstractVector{T}; trans::Bool=false) where {T}

  @assert length(x) == length(b)
  @assert length(x) == mesh.numDof

  t1 = zeros(T, length(x))
  t2 = zeros(T, length(x))
  _trans::Bool = trans

  exit_status = 1  # 1 = itermax, 2 = res_tol
  for i=1:pc.itermax

    # t1 = R*x
    fill!(t1, 0)
    computeRProduct(pc, mesh, sbp, eqn, opts, x, t1, trans=_trans)

    # t1 = b - R*x
    @simd for j=1:length(t1)
      t1[j] = b[j] - t1[j]
    end

    if pc.res_tol > 0 || pc.verbose  # save the expense of the mat-vec if residual will
                       # not be checked
      # compute residual: b - A*x = t1 - D*x
      applyBDiagPCInv(pc.diag_pc, mesh, sbp, eqn, opts, x, t2; trans=_trans)
      @simd for j=1:length(t2)
        t2[j] = t1[j] - t2[j]
      end

      res_norm = calcEuclidianNorm(eqn.comm, t2)
      if pc.verbose
        println(BSTDOUT, "  iteration ", i, " linear residual norm: ", res_norm)
        println(BSTDOUT,  "    max res = ", maximum(abs.(t2)))
      end

      if res_norm < pc.res_tol
        exit_status = 2
        break
      end
    end  # end if res_tol > 0

    # x = D^-1 * (b - R*x) = D^-1 * t1
    if _trans
      applyPCTranspose(pc.diag_pc, mesh, sbp, eqn, opts, 0.0, t1, x)
    else
      applyPC(pc.diag_pc, mesh, sbp, eqn, opts, 0.0, t1, x)
    end
  end

  if pc.verbose
    if exit_status == 1
      println(BSTDOUT, "exited Block Jacobi smoother due to itermax")
      # res_norm is one iteration old, so don't print it here
    else
      println(BSTDOUT, "exited Block Jacobi smoother due to residual norm")

    end
  end

  flush(BSTDOUT)

  return nothing
end

function applyPCTranspose(pc::NewtonBJacobiPC, mesh::AbstractMesh,
                 sbp::AbstractOperator,
                 eqn::AbstractSolutionData, opts::Dict, t, b::AbstractVector,
                 x::AbstractVector{T}) where {T}

  applyPC(pc, mesh, sbp, eqn, opts, t, b, x, trans=true)
end

"""
  Multiply R (= A - D), where A is the Jacobian and D is the block diagonal), by
  a vector.

  **Inputs**

   * pc
   * mesh
   * sbp
   * eqn
   * opts
   * x: vector to multiply against

  **Inputs/Outputs**

   * b: vector to overwrite with result

  **Inputs/Outputs**

   * trans: if true, compute R.'x, default false
"""
function computeRProduct(pc::NewtonBJacobiPC, mesh, sbp, eqn, opts, x, b; trans::Bool=false)

  # compute R*x as A*x - D*x
  if trans
    pc.evalJacTVecProduct(mesh, sbp, eqn, opts, x, b)
  else
    pc.evalJacVecProduct(mesh, sbp, eqn, opts, x, b)
  end
  scale!(b, -1)
  applyBDiagPCInv(pc.diag_pc, mesh, sbp, eqn, opts, x, b, zero_output=false, trans=trans)
  scale!(b, -1)
  return nothing
end


text:
module TestBatching

include("preamble.jl")
using Kaleido: SingletonLens

lens_created_in_global_scope = @batchlens begin
    _.a.b.c.d
    _.a.b.c.e
    _.a.b.f
    _.a.g
end

@testset "`@batchlens` inside `@test`" begin
    # Inside of `@test` macro, `LineNumberNode`s are missing.  Make
    # sure that `@batchlens` can handle that:
    @test (@batchlens begin
        _.a.b.c.d
        _.a.b.c.e
        _.a.b.f
        _.a.g
    end) == lens_created_in_global_scope
end

@testset "1st level" begin
    lens = @batchlens begin
        _.a
        _.b
    end

    obj = (a=1, b=2)
    @test get(obj, lens) == (1, 2)
    @test set(obj, lens, (10, 20)) == (a=10, b=20)

    @test lens == IndexBatchLens(:a, :b)
end

@testset "2nd level" begin
    lens = @batchlens begin
        _.a.b
        _.a.c
    end

    obj = (a=(b=1, c=2),)
    @test get(obj, lens) == (1, 2)
    @test set(obj, lens, (10, 20)) == (a=(b=10, c=20),)

    @test lens ==
        IndexBatchLens(:a) ∘ MultiLens((
            (@lens _[1]) ∘ IndexBatchLens(:b, :c),
        )) ∘ FlatLens(2)
end

@testset "3rd level" begin
    lens = @batchlens begin
        _.a.b.c
        _.a.b.d
        _.a.e
    end

    obj = (a=(b=(c=1, d=2), e=3),)
    @test get(obj, lens) == (1, 2, 3)
    @test set(obj, lens, (10, 20, 30)) == (a=(b=(c=10, d=20), e=30),)

    @test lens ==
        IndexBatchLens(:a) ∘ MultiLens((
            (@lens _[1]) ∘ IndexBatchLens(:b, :e) ∘ MultiLens((
                (@lens _[1]) ∘ IndexBatchLens(:c, :d),
                (@lens _[2]) ∘ Kaleido.SingletonLens(),
            )) ∘ FlatLens(2, 1),
        )) ∘ FlatLens(3)
end

@testset "4th level" begin
    lens = @batchlens begin
        _.a.b.c.d
        _.a.b.c.e
        _.a.b.f
        _.a.g
    end

    obj = (a=(b=(c=(d=1, e=2), f=3), g=4),)
    @test get(obj, lens) == (1, 2, 3, 4)
    @test set(obj, lens, (10, 20, 30, 40)) ==
        (a=(b=(c=(d=10, e=20), f=30), g=40),)

    @test lens ==
        IndexBatchLens(:a) ∘ MultiLens((
            (@lens _[1]) ∘ IndexBatchLens(:b, :g) ∘ MultiLens((
                (@lens _[1]) ∘ IndexBatchLens(:c, :f) ∘ MultiLens((
                    (@lens _[1]) ∘ IndexBatchLens(:d, :e),
                    (@lens _[2]) ∘ SingletonLens(),
                )) ∘ FlatLens(2, 1),
                (@lens _[2]) ∘ SingletonLens(),
            )) ∘ FlatLens(3, 1),
        )) ∘ FlatLens(4)
end

@testset "mixing indexing" begin
    lens = @batchlens begin
        _.a.b.c
        _.a.b.d[1]
        _.a.b.d[3]
        _.a.e
    end

    obj = (a=(b=(c=1, d=(2, 3, 4)), e=5),)
    @test get(obj, lens) == (1, 2, 4, 5)
    @test set(obj, lens, (10, 20, 40, 50)) ==
        (a=(b=(c=10, d=(20, 3, 40)), e=50),)

    @test lens ==
        IndexBatchLens(:a) ∘ MultiLens((
            (@lens _[1]) ∘ IndexBatchLens(:b, :e) ∘ MultiLens((
                (@lens _[1]) ∘ IndexBatchLens(:c, :d) ∘ MultiLens((
                    (@lens _[1]) ∘ Kaleido.SingletonLens(),
                    (@lens _[2]) ∘ MultiLens((
                        (@lens _[1]), (@lens _[3]),
                    )),
                )) ∘ FlatLens(1, 2),
                (@lens _[2]) ∘ Kaleido.SingletonLens())) ∘ FlatLens(3, 1),
        )) ∘ FlatLens(4)
end

@testset "inline composition" begin
    lens = @batchlens begin
        _.a.b.c
        _.a.b.d[1]
        _.a.b.d[3] ∘ settingas𝕀
        _.a.e
    end

    obj = (a = (b = (c = 1, d = (2, 3, 0.5)), e = 5),)
    @test get(obj, lens) == (1, 2, 0.0, 5)
    @test set(obj, lens, (10, 20, Inf, 50)) ==
        (a = (b = (c = 10, d = (20, 3, 1.0)), e = 50),)

    @test lens ==
        IndexBatchLens(:a) ∘ MultiLens((
            (@lens _[1]) ∘ IndexBatchLens(:b, :e) ∘ MultiLens((
                (@lens _[1]) ∘ IndexBatchLens(:c, :d) ∘ MultiLens((
                    (@lens _[1]) ∘ Kaleido.SingletonLens(),
                    (@lens _[2]) ∘ MultiLens((
                        (@lens _[1]),
                        (@lens _[3]) ∘ settingas𝕀,
                    )))) ∘ FlatLens(1, 2),
                (@lens _[2]) ∘ Kaleido.SingletonLens(),
            )) ∘ FlatLens(3, 1),
        )) ∘ FlatLens(4)
end

@testset "unsorted" begin
    lens = @batchlens begin
        _.a.b.c
        _.a.b.d
        _.a.e
        _.a.b.f
        _.a.b.g
    end

    obj = (a=(b=(c=1, d=2, f=3, g=4), e=5),)
    @test get(obj, lens) == (1, 2, 5, 3, 4)
    @test_broken set(obj, lens, (10, 20, 50, 30, 40)) ==
        (a=(b=(c=10, d=20, f=30, g=40), e=50),)

    # This should work but it's better to make sure it throws for now
    # rather than silently failing:
    @test_throws Exception set(obj, lens, (10, 20, 50, 30, 40))

    @test lens ==
        IndexBatchLens(:a) ∘ MultiLens((
            (@lens _[1]) ∘
            IndexBatchLens(:b, :e, :b) ∘  # TODO: this is wrong; fix it
            MultiLens((
                (@lens _[1]) ∘ IndexBatchLens(:c, :d),
                (@lens _[2]) ∘ Kaleido.SingletonLens(),
                (@lens _[3]) ∘ IndexBatchLens(:f, :g)
            )) ∘ FlatLens(2, 1, 2),
        )) ∘ FlatLens(5)
end

end  # module


text:
"""
    rho_sol(r, rsol, rhosol)

Compute the local volume density for the soliton model.

* `rsol`: soliton scale radius in kpc
* `rhosol`: soliton scale density in kpc
"""
function rho_sol(r, rsol, rhosol)
    @. rhosol * (1.0 + (r / rsol)^2) ^ -8
end

"""
    M_sol(r, rsol, rhosol)

Compute the enclosed mass for the soliton model.  It's analytic!

* `rsol`: soliton scale radius in kpc
* `rhosol`: soliton scale density in kpc
"""
function M_sol(r, rsol, rhosol)
    x = r / rsol
    t = atan.(x, 1.0)
    term0 = 27720. * t
    term1 = 17325. * sin.(2 * t)
    term2 = -1155. * sin.(4 * t)
    term3 = -4235. * sin.(6 * t)
    term4 = -2625. * sin.(8 * t)
    term5 = -903. * sin.(10 * t)
    term6 = -175. * sin.(12 * t)
    term7 = -15. * sin.(14 * t)
    factor = 4π * rhosol * rsol^3 / 1720320.
    return @. factor * (term0 + term1 + term2 + term3 + term4 + term5 + term6 + term7)
end

"""
    drhodr_sol(r, rsol, rhosol)

Compute the linear density derivative for the soliton model.

* `rsol`: soliton scale radius in kpc
* `rhosol`: soliton scale density in kpc
"""
function drhodr_sol(r, rsol, rhosol)
    x = r / rsol
    return @. -16 * rhosol * x * (1.0 + x^2) ^ -9
end

"""
    m22_from_sol(rsol, rhosol; <keyword arguments>)

Calculate the axion mass from the soliton scale parameters. See Marsh & Pop 2015,
equation 8.  Axion mass has units of 1e-22 eV.

* `rsol`: soliton scale radius in kpc
* `rhosol`: soliton scale density in kpc

# Arguments
- `cosmo::AbstractCosmology = default_cosmo`: cosmology under which to evaluate the overdensity
- `z::Real = 0.0`: redshift at which to evaluate the overdensity
- `alpha_mp::Real = 0.23`: the alpha fitting parameter from Marsh & Pop 2015
"""
function m22_from_sol(rsol, rhosol;
                      cosmo = default_cosmo, z = 0.0, alpha_mp = 0.23)
    Δ_sol = rhosol / ρcrit(z; cosmo = cosmo)
    m22 = √(Δ_sol ^ -1 * rsol ^ -4 * (cosmo.h / 0.7) ^ 2 * 5e4 * alpha_mp ^ -4)
    return m22
end

"""
    rhosol_from_rsol(m22, rsol; <keyword arguments>)

Calculate the scale density of the soliton core from m22 and the scale radius.

* `m22`: ultralight axion mass in units of 1e-22 eV
* `rsol`: soliton scale radius in kpc

# Arguments
- `cosmo::AbstractCosmology = default_cosmo`: cosmology under which to evaluate the overdensity
- `z::Real = 0.0`: redshift at which to evaluate the overdensity
- `alpha_mp::Real = 0.23`: the alpha fitting parameter from Marsh & Pop 2015
"""
function rhosol_from_rsol(m22, rsol;
                         cosmo = default_cosmo, z = 0.0, alpha_mp = 0.23)
    Δ_sol = (5e4 / alpha_mp^4) * (cosmo.h / 0.7)^-2 * m22^-2 * rsol^-4
    return Δ_sol * ρcrit(z; cosmo = cosmo)
end

"""
    rsol_from_Mvir(m22, Mvir)

Calculate the scale radius of the soliton core from the halo mass scaling relation
(see Robles+2019).
"""
function rsol_from_Mvir(m22, Mvir)
    rsol = 3.315 * 1.6 * (Mvir / 1e9)^(-1/3.) * m22^-1
end

"""
    matching_radius(ρ1, ρ2; <keyword arguments>)

Compute the matching radius from the specified density profiles.

# Arguments
- `x_start::Real = 2.0`: factor of radius scale to start the search for the root
- `radius_scale::Real = 1.0`: radius scale
- `maxevals:Int = 100`: maximum number of function evalutions for the root finding
"""
function matching_radius(ρ1, ρ2;
                         x_start = 2.0, radius_scale = 1.0, maxevals = 100)
    f1(x) = ρ1(radius_scale * x)
    f2(x) = ρ2(radius_scale * x)
    y(x) = f1(x) / f2(x) - 1.0
    yp(x) = ForwardDiff.derivative(y, float(x))
    x0 = Roots.find_zero((y, yp), x_start, Roots.Newton(),
                         maxevals = maxevals)
    return radius_scale * x0
end

"""
    SolNFWModel(rs, rhos, rsol, rhosol, repsilon)

Soliton NFW model.  For consistency, the NFW density must match the soliton
density at repsilon.  If that is not the case, the constructor will update
repsilon.

- [Marsh & Pop 2015](https://ui.adsabs.harvard.edu/#abs/2015MNRAS.451.2479M)
- [Robles, Bullock, and Boylan-Kolchin 2019](https://ui.adsabs.harvard.edu/abs/2019MNRAS.483..289R/)

* `rs`: NFW scale radius in kpc
* `rhos`: NFW scale density in Msun / kpc3
* `rsol`: soliton scale radius in kpc
* `rhosol`: soliton scale density in kpc
* `repsilon`: transition radius
"""
struct SolNFWModel <: HaloModel
    rs::Float64
    rhos::Float64
    rsol::Float64
    rhosol::Float64
    repsilon::Float64
    # enforce density matching
    SolNFWModel(rs, rhos, rsol, rhosol, repsilon;
                x_start = 2.0, maxevals = 100) = begin
        ρnfw = rho_NFW(repsilon, rs, rhos)
        ρsol = rho_sol(repsilon, rsol, rhosol)
        if abs((ρnfw - ρsol) / ρsol) < 4 * eps()
            return new(rs, rhos, rsol, rhosol, repsilon)
        else
            @warn("recalculating matching radius", maxlog=1)
            try
                ρ1(r) = rho_NFW(r, rs, rhos)
                ρ2(r) = rho_sol(r, rsol, rhosol)
                repsilon = matching_radius(ρ1, ρ2; radius_scale = rsol,
                                           x_start = x_start, maxevals = maxevals)
            catch err
                @warn("failed to find matching radius for " *
                      "\trs = $rs \n\trhos = $rhos \n\trsol = $rsol \n\trhosol = $rhosol")
                throw(err)
            end
            return new(rs, rhos, rsol, rhosol, repsilon)
        end
    end
end

SolNFWModel() = SolNFWModel(21.1, 5.6e6, 0.53, 3.1e10, 0.48)

density(halo::SolNFWModel, r::T where T <: Real) = begin
    if r < halo.repsilon
        return rho_sol(r, halo.rsol, halo.rhosol)
    end
    return rho_NFW(r, halo.rs, halo.rhos)
end

density(halo::SolNFWModel, r::Array{T, 1} where T <: Real) = begin
    ρ = zero(r)
    idx_sol = r .< halo.repsilon
    idx_nfw = .~ idx_sol
    ρ[idx_sol] = rho_sol(r[idx_sol], halo.rsol, halo.rhosol)
    ρ[idx_nfw] = rho_NFW(r[idx_nfw], halo.rs, halo.rhos)
    return ρ
end

mass(halo::SolNFWModel, r::T where T <: Real) = begin
    if r < halo.repsilon
        return M_sol(r, halo.rsol, halo.rhosol)
    end
    dM_epsilon = (M_sol(halo.repsilon, halo.rsol, halo.rhosol) -
                  M_NFW(halo.repsilon, halo.rs, halo.rhos))
    return M_NFW(r, halo.rs, halo.rhos) + dM_epsilon
end

mass(halo::SolNFWModel, r::Array{T, 1} where T <: Real) = begin
    M = zero(r)
    dM_epsilon = (M_sol(halo.repsilon, halo.rsol, halo.rhosol) -
                  M_NFW(halo.repsilon, halo.rs, halo.rhos))
    idx_sol = r .< halo.repsilon
    idx_nfw = .~ idx_sol
    M[idx_sol] = M_sol(r[idx_sol], halo.rsol, halo.rhosol)
    M[idx_nfw] = M_NFW(r[idx_nfw], halo.rs, halo.rhos) .+ dM_epsilon
    return M
end

scale_radius(halo::SolNFWModel) = halo.rs

"""
    SolNFW_from_virial(Mvir, cvir, m22; <keyword arguments>)

Construct a Solition-NFW halo from the halo mass, concentration, and axion mass.

# Arguments
- `rsol = nothing`: for default, use the soliton size-halo mass scaling relationship
- `mdef::AbstractString = default_mdef`: halo mass definition (e.g., "200c", "vir")
- `cosmo::AbstractCosmology = default_cosmo`: cosmology under which to evaluate the overdensity
- `z::Real = 0.0`: redshift at which to evaluate the overdensity
- `x_start::Real = 2.0`: factor of radius scale to start the search for the root
- `maxevals:Int = 100`: maximum number of function evalutions for the root finding
"""
function SolNFW_from_virial(Mvir, cvir, m22;
                            rsol = nothing,
                            mdef = default_mdef, cosmo = default_cosmo, z = 0.0,
                            x_start = 2.0, maxevals = 100)
    # soliton parameters, calculate rsol from Mvir scaling if not passed in
    if rsol == nothing
        rsol = rsol_from_Mvir(m22, Mvir)
    end
    rhosol = rhosol_from_rsol(m22, rsol, cosmo = cosmo, z = z)
    # NFW parameters
    Rvir = Rvir_from_Mvir(Mvir; mdef = mdef, cosmo = cosmo, z = z)
    rs = Rvir / cvir

    # first guess, rhos from normal NFW profile
    rhos = Mvir / M_NFW(Rvir, rs, 1.0)
    repsilon = 2 * rsol
    return SolNFWModel(rs, rhos, rsol, rhosol, repsilon;
                       x_start = x_start, radius_scale = rsol, maxevals = maxevals)
end

"""
    SolABGModel(rs, rhos, alpha, beta, gamma, rsol, rhosol, repsilon)

Soliton + αβγ model.  For consistency, the αβγ density profile must match the
soliton density profile at repsilon.  If that is not the case, the constructor
will update repsilon.

* `rs`: NFW scale radius in kpc
* `rhos`: NFW scale density in Msun / kpc3
* `alpha`: transition sharpness
* `beta`: negative outer log slope
* `gamma`: negative inner log slope
* `rsol`: soliton scale radius in kpc
* `rhosol`: soliton scale density in kpc
* `repsilon`: transition radius
"""
struct SolABGModel <: HaloModel
    rs::Float64
    rhos::Float64
    alpha::Float64
    beta::Float64
    gamma::Float64
    rsol::Float64
    rhosol::Float64
    repsilon::Float64
    # enforce density matching
    SolABGModel(rs, rhos, alpha, beta, gamma, rsol, rhosol, repsilon;
                x_start = 2.0, maxevals = 100) = begin
        ρabg = rho_ABG(repsilon, rs, rhos, alpha, beta, gamma)
        ρsol = rho_sol(repsilon, rsol, rhosol)
        if abs((ρabg - ρsol) / ρsol) < 4 * eps()
            return new(rs, rhos, alpha, beta, gamma, rsol, rhosol, repsilon)
        else
            @warn("recalculating matching radius", maxlog=1)
            try
                ρ1(r) = rho_ABG(r, rs, rhos, alpha, beta, gamma)
                ρ2(r) = rho_sol(r, rsol, rhosol)
                repsilon = matching_radius(ρ1, ρ2; radius_scale = rsol,
                                           x_start = x_start, maxevals = maxevals)
            catch err
                @warn("failed to find matching radius for " *
                      "\trs = $rs \n\trhos = $rhos \n\trsol = $rsol \n\trhosol = $rhosol")
                throw(err)
            end
            return new(rs, rhos, alpha, beta, gamma, rsol, rhosol, repsilon)
        end
    end
end

SolABGModel() = SolABGModel(21.1, 5.6e6, 1, 3, 1, 0.53, 3.1e10, 0.48)

density(halo::SolABGModel, r::T where T <: Real) = begin
    if r < halo.repsilon
        return rho_sol(r, halo.rsol, halo.rhosol)
    end
    return rho_ABG(r, halo.rs, halo.rhos, halo.alpha, halo.beta, halo.gamma)
end

density(halo::SolABGModel, r::Array{T, 1} where T <: Real) = begin
    ρ = zero(r)
    idx_sol = r .< halo.repsilon
    idx_abg = .~ idx_sol
    ρ[idx_sol] = rho_sol(r[idx_sol], halo.rsol, halo.rhosol)
    ρ[idx_abg] = rho_ABG(r[idx_abg], halo.rs, halo.rhos, halo.alpha, halo.beta, halo.gamma)
    return ρ
end

mass(halo::SolABGModel, r::T where T <: Real) = begin
    if r < halo.repsilon
        return M_sol(r, halo.rsol, halo.rhosol)
    end
    dM_epsilon = (M_sol(halo.repsilon, halo.rsol, halo.rhosol) -
                  M_ABG(halo.repsilon, halo.rs, halo.rhos, halo.alpha, halo.beta, halo.gamma))
    return M_ABG(r, halo.rs, halo.rhos, halo.alpha, halo.beta, halo.gamma) + dM_epsilon
end

mass(halo::SolABGModel, r::Array{T, 1} where T <: Real) = begin
    M = zero(r)
    dM_epsilon = (M_sol(halo.repsilon, halo.rsol, halo.rhosol) -
                  M_ABG(halo.repsilon, halo.rs, halo.rhos, halo.alpha, halo.beta, halo.gamma))
    idx_sol = r .< halo.repsilon
    idx_nfw = .~ idx_sol
    M[idx_sol] = M_sol(r[idx_sol], halo.rsol, halo.rhosol)
    M[idx_nfw] = M_ABG(r[idx_nfw], halo.rs, halo.rhos, halo.alpha, halo.beta, halo.gamma) .+ dM_epsilon
    return M
end

scale_radius(halo::SolABGModel) = begin
    ((2 - halo.gamma) / (halo.beta - 2)) ^ (1 / halo.alpha) * halo.rs
end

"""
    SolABG_from_virial(Mvir, cvir, alpha, beta, gamma, m22; <keyword arguments>)

Construct a Soliton + αβγ halo from the halo mass, concentration, shape parameters, and
axion mass.

# Arguments
- `rsol = nothing`: for default, use the soliton size-halo mass scaling relationship
- `mdef::AbstractString = default_mdef`: halo mass definition (e.g., "200c", "vir")
- `cosmo::AbstractCosmology = default_cosmo`: cosmology under which to evaluate the overdensity
- `z::Real = 0.0`: redshift at which to evaluate the overdensity
- `x_start::Real = 2.0`: factor of radius scale to start the search for the root
- `maxevals:Int = 100`: maximum number of function evalutions for the root finding
"""
function SolABG_from_virial(Mvir, cvir, alpha, beta, gamma, m22;
                            rsol = nothing,
                            mdef = default_mdef, cosmo = default_cosmo, z = 0.0,
                            x_start = 2.0, maxevals = 100)
    # soliton parameters, calculate rsol from Mvir scaling if not passed in
    if rsol == nothing
        rsol = rsol_from_Mvir(m22, Mvir)
    end
    rhosol = rhosol_from_rsol(m22, rsol, cosmo = cosmo, z = z)
    # ABG parameters
    Rvir = Rvir_from_Mvir(Mvir; mdef = mdef, cosmo = cosmo, z = z)
    r2 = Rvir / cvir
    rs = ((2 - gamma) / (beta - 2))^(-1 / alpha) * r2

    # first guess, rhos from normal ABG profile
    rhos = Mvir / M_ABG(Rvir, rs, 1.0, alpha, beta, gamma)
    repsilon = 2 * rsol
    return SolABGModel(rs, rhos, alpha, beta, gamma, rsol, rhosol, repsilon;
                       x_start = x_start, maxevals = maxevals)
end


text:
#=
sino_geom.jl
sinogram geometry for 2D tomographic image reconstruction
2019-07-01, Jeff Fessler, University of Michigan
=#

export MIRT_sino_geom, sino_geom
export sino_geom_help, sino_geom_plot_grids, sino_geom_show, sino_geom_test

# using MIRT: jim, image_geom, MIRT_image_geom
using Plots: Plot, plot!, plot, scatter!, gui
using Test: @test, @test_throws, @inferred

struct MIRT_sino_geom
	how::Symbol				# :par | :moj | :fan
	units::Symbol			# :nothing | :mm | :cm etc.
	nb::Int					# # of "radial" samples, aka ns
	na::Int					# # of angular samples
	d::Float32				# aka dr or ds, "radial" sample spacing
							# (is dx for mojette, pixels must be square)
	orbit::Float32			# [degrees]
	orbit_start::Float32	# [degrees]
	offset::Float32			# sample offset, cf offset_r or offset_s [unitless]
	strip_width::Float32	#

	# for fan:
	source_offset::Float32	# same units as d, etc., e.g., [mm]
							# use with caution!
	dsd::Float32			# dis_src_det, Inf for parallel beam
	dod::Float32			# dis_iso_det
#	dso::Float32			# dis_src_iso = dsd-dod, Inf for parallel beam
	dfs::Float32			# distance from focal spot to source
end

"""
`sino_geom_help()`
"""
function sino_geom_help( ; io::IO = isinteractive() ? stdout : IOBuffer() )
	print(io, "$(basename(@__FILE__)) propertynames:\n\t")
	print(io, Tuple(sort([propertynames(sino_geom(:par))...])))

	print(io,
	"\n
	Derived values

	sg.dim			dimensions: (nb,na)
	sg.ds|dr		radial sample spacing (NaN for :moj)
	sg.s			[nb] s sample locations
	sg.w			(nb-1)/2 + offset ('middle' sample position)
	sg.ad			source angles [degrees]
	sg.ar			source angles [radians]
	sg.ones			ones(Float32, nb,na)
	sg.zeros		zeros(Float32, nb,na)
	sg.rfov			radial fov
	sg.xds			[nb] center of detector elements (beta=0)
	sg.yds			[nb] ''
	sg.grid			(rg, phigrid) [nb na] parallel-beam coordinates
	sg.plot_grid	plot sg.grid

	For mojette:

	sg.d_ang		[na]

	For fan beam:

	sg.gamma		[nb] gamma sample values [radians]
	sg.gamma_max	half of fan angle [radians]
	sg.dso			# dsd - dod, Inf for parallel beam

	Methods

	sg.down(down)		reduce sampling by integer factor
	sg.shape(sino)		reshape sinograms into array [nb na :]
	sg.unitv(;ib,ia)	unit 'vector' with single nonzero element
	sg.taufun(x,y)		projected s/ds for each (x,y) pair [numel(x) na]
	sg.plot(;ig)		plot system geometry (most useful for fan)
	\n")
end

"""
`function sg = sino_geom(...)`

Constructor for `MIRT_sino_geom`

Create the "sinogram geometry" structure that describes the sampling
characteristics of a given sinogram for a 2D parallel or fan-beam system.
Using this structure facilitates "object oriented" code.
(Use `ct_geom()` instead for 3D axial or helical cone-beam CT.)

in
- `how::Symbol`	`:fan` (fan-beam) | `:par` (parallel-beam) | `:moj` (mojette)

options for all geometries (including parallel-beam):
- `units::Symbol`	e.g. `:cm` or `:mm`; default: :none
- `orbit_start`		default: 0
- `orbit`			[degrees] default: `180` for parallel / mojette
					and `360` for fan
					can be `:short` for fan-beam short scan
- `down::Int`		down-sampling factor, for testing

- `nb`				# radial samples cf `nr` (i.e., `ns` for `:fan`)
- `na`				# angular samples (cf `nbeta` for `:fan`)
- `d`				radial sample spacing; cf `dr` or `ds`; default 1
					for mojette this is actually `dx`
- `offset`			cf `offset_r` `channel_offset` unitless; default 0
			(relative to centerline between two central channels).
			Use 0.25 or 1.25 for "quarter-detector offset"
- `strip_width`		detector width; default: `d`

options for fan-beam
- `source_offset`		same units as d; use with caution! default 0
fan beam distances:
- `dsd`		cf `dis_src_det`	default: `Inf` (parallel beam)
- `dod`		cf `dis_iso_det`	default: `0`
- `dfs`		cf `dis_foc_src`	default: `0` (3rd generation CT arc),
				use `Inf` for flat detector

out
- `sg::MIRT_sino_geom`	initialized structure

See also
- `sino_geom_help()` help on methods
- `sino_geom_plot_grids()` show sampling
- `sino_geom_show()` show the geometries
- `sino_geom_test()` self test

Jeff Fessler, University of Michigan
"""
function sino_geom(how::Symbol; kwarg...)
	if how === :par
		sg = sino_geom_par( ; kwarg...)
	elseif how === :fan
		sg = sino_geom_fan( ; kwarg...)
	elseif how === :moj
		sg = sino_geom_moj( ; kwarg...)
	elseif how === :ge1
		sg = sino_geom_ge1( ; kwarg...)
#=
	elseif how === :hd1
		sg = sino_geom_hd1( ; kwarg...)
	elseif how === :revo1fan
		tmp = ir_fan_geom_revo1(type)
		sg = sino_geom(:fan, tmp{:}, varargin{:})
=#
	else
		throw("unknown sino type $how")
	end

	return sg
end

"""
`sg = downsample(sg, down)`
down-sample (for testing with small problems)
"""
function downsample(sg::MIRT_sino_geom, down::Int)
	down == 1 && return sg

	nb = 2 * round(Int, sg.nb / down / 2) # keep it even
	na = round(Int, sg.na / down)

	return MIRT_sino_geom(sg.how, sg.units,
		nb, na, sg.d * down, sg.orbit, sg.orbit_start, sg.offset,
		sg.strip_width * down,
		sg.source_offset, sg.dsd, sg.dod, sg.dfs)
end

"""
`sg = sino_geom_over(sg, over::Int)`
over-sample in "radial" dimension
Probably not meaningful for mojette sampling because d=dx.
"""
function sino_geom_over(sg::MIRT_sino_geom, over::Int)
	if over == 1
		return sg
	end

	return MIRT_sino_geom(sg.how, sg.units,
		sg.nb * over, sg.na, sg.d / over,
		sg.orbit, sg.orbit_start, sg.offset * over,
		sg.strip_width / over,
		sg.source_offset, sg.dsd, sg.dod, sg.dfs)
end

"""
`sg = sino_geom_fan()`
"""
function sino_geom_fan( ;
		units::Symbol = :none,
		nb::Int = 128,
		na::Int = 2 * floor(Int, nb * pi/2 / 2),
		d::Real = 1,
		orbit::Union{Symbol,Real} = 360, # [degrees]
		orbit_start::Real = 0,
		strip_width::Real = d,
		offset::Real = 0,
		source_offset::Real = 0,
		dsd::Real = 4*nb*d,	# dis_src_det
	#	dso::Real = [],		# dis_src_iso
		dod::Real = nb*d,	# dis_iso_det
		dfs::Real = 0,		# dis_foc_src (3rd gen CT)
		down::Int = 1,
	)

	dfs != 0 && !isinf(dfs) && throw("dfs $dfs") # must be 0 or Inf

	if orbit === :short # trick
		sg_tmp = MIRT_sino_geom(:fan, units,
			nb, na, d, 0, orbit_start, offset, strip_width,
			source_offset, dsd, dod, dfs)
		orbit = sg_tmp.orbit_short
	end
	isa(orbit, Symbol) && throw("orbit :orbit")

	sg = MIRT_sino_geom(:fan, units,
		nb, na, d, orbit, orbit_start, offset, strip_width,
		source_offset, dsd, dod, dfs)

	return downsample(sg, down)
end

"""
`sg = sino_geom_par( ... )`
"""
function sino_geom_par( ;
		units::Symbol = :none,
		nb::Int = 128,
		na::Int = 2 * floor(Int, nb * pi/2 / 2),
		down::Int = 1,
		d::Real = 1,
		orbit::Real = 180, # [degrees]
		orbit_start::Real = 0,
		strip_width::Real = d,
		offset::Real = 0,
	)

	sg = MIRT_sino_geom(:par, units,
		nb, na, d, orbit, orbit_start, offset, strip_width,
		0, 0, 0, 0)

	return downsample(sg, down)
end

"""
`sg = sino_geom_moj( ... )`
"""
function sino_geom_moj( ;
		units::Symbol = :none,
		nb::Int = 128,
		na::Int = 2 * floor(Int, nb * pi/2 / 2),
		down::Int = 1,
		d::Real = 1, # means dx for :moj
		orbit::Real = 180, # [degrees]
		orbit_start::Real = 0,
		strip_width::Real = d, # ignored ?
		offset::Real = 0,
	)

	sg = MIRT_sino_geom(:moj, units,
		nb, na, d, orbit, orbit_start, offset, strip_width,
		0, 0, 0, 0)

	return downsample(sg, down)
end

"gamma for general finite dfs (rarely used)"
function sino_geom_gamma_dfs(sg)
	dis_foc_det = sg.dfs + sg.dsd
	alf = sg.s / dis_foc_det
	atan.(dis_foc_det * sin.(alf), dis_foc_det * cos.(alf) .- sg.dfs)
	# equivalent to s/dsd when dfs=0
end

"""
`sino_geom_gamma()`
gamma sample values for :fan
"""
function sino_geom_gamma(sg)
	return	sg.dfs == 0 ? sg.s / sg.dsd : # 3rd gen: equiangular
			isinf(sg.dfs) ? atan.(sg.s / sg.dsd) : # flat
			sino_geom_gamma_dfs(sg) # general
end

"""
`sino_geom_rfov()`
radial FOV
"""
function sino_geom_rfov(sg)
	return	sg.how === :par ? maximum(abs.(sg.r)) :
			sg.how === :fan ? sg.dso * sin(sg.gamma_max) :
			sg.how === :moj ? sg.nb/2 * minimum(sg.d_ang) : # todo: check
				throw("bad how $(sg.how)")
end

"""
`sino_geom_taufun()`
projected `s/ds`, useful for footprint center and support
"""
function sino_geom_taufun(sg, x, y)
	size(x) != size(y) && throw("bad x,y size")
	x = x[:]
	y = y[:]
	if sg.how === :par || sg.how === :moj # todo: check
		ar = sg.ar' # row vector, for outer-product
		tau = (x * cos.(ar) + y * sin.(ar)) / sg.dr
	elseif sg.how === :fan
		b = sg.ar' # row vector, for outer-product
		xb = x * cos.(b) + y * sin.(b)
		yb = -x * sin.(b) + y * cos.(b)
		tangam = (xb .- sg.source_offset) ./ (sg.dso .- yb) # e,tomo,fan,L,gam
		if sg.dfs == 0 # arc
			tau = sg.dsd / sg.ds * atan.(tangam)
		elseif isinf(sg.dfs) # flat
			tau = sg.dsd / sg.ds * tangam
#		else
#			throw("bad dfs $(sg.dfs)")
		end
#	else
#		throw("bad how $(sg.how)")
	end
	return tau
end

"""
`sino_geom_xds()`
center positions of detectors (for beta = 0)
"""
function sino_geom_xds(sg)
	if sg.how === :par
		xds = sg.s
	elseif sg.how === :moj
		xds = sg.s # todo: really should be angle dependent
	elseif sg.how === :fan
		if sg.dfs == 0 # arc
			gam = sg.gamma
			xds = sg.dsd * sin.(gam)
		elseif isinf(sg.dfs) # flat
			xds = sg.s
	#	else
	#		throw("bad dfs $(sg.dfs))")
		end
#	else
#		throw("bad how $(sg.how)")
	end
	return xds .+ sg.source_offset
end

"""
`sino_geom_yds()`
center positions of detectors (for beta = 0)
"""
function sino_geom_yds(sg)

	if sg.how === :par
		yds = zeros(Float32, sg.nb)
	elseif sg.how === :moj
		yds = zeros(Float32, sg.nb)
	elseif sg.how === :fan
		if sg.dfs == 0 # arc
			gam = sg.gamma
			yds = sg.dso .- sg.dsd * cos.(gam)
		elseif isinf(sg.dfs) # flat
			yds = fill(-sg.dod, sg.nb)
	#	else
	#		throw("bad dfs $(sg.dfs))")
		end
#	else
#		throw("bad how $(sg.how)")
	end
	return yds
end

"""
`sino_geom_unitv()`
sinogram with a single ray
"""
function sino_geom_unitv(sg::MIRT_sino_geom;
		ib=round(Int, sg.nb/2+1),
		ia=round(Int, sg.na/2+1))
	out = sg.zeros
	out[ib,ia] = 1
	return out
end

"""
`(rg, ϕg) = sino_geom_grid(sg::MIRT_sino_geom)`

Return grids `rg` and `ϕg` (in radians) of size `[nb na]`
of equivalent *parallel-beam* `(r,ϕ)` (radial, angular) sampling positions,
for any sinogram geometry.
For parallel beam this is just `ndgrid(sg.r, sg.ar)`
but for fan beam and mojette this involves more complicated computations.
"""
function sino_geom_grid(sg::MIRT_sino_geom)

	if sg.how === :par
		return ndgrid(sg.r, sg.ar)

	elseif sg.how === :fan
		gamma = sg.gamma
		rad = sg.dso * sin.(gamma) + sg.source_offset * cos.(gamma)
		rg = repeat(rad, 1, sg.na) # [nb na]
		return (rg, gamma .+ sg.ar') # [nb na] phi = gamma + beta
	end

	# otherwise :moj (mojette)
	phi = sg.ar
	# trick: ray_spacing aka ds comes from dx which is sg.d for mojette
	wb = (sg.nb - 1)/2 + sg.offset
	dt = sg.d_ang # [na]
	pos = ((0:(sg.nb-1)) .- wb) * dt' # [nb na]
	return (pos, repeat(phi', sg.nb, 1))
end

function Base.display(sg::MIRT_sino_geom)
	ir_dump(sg)
end

# Extended properties

sino_geom_fun0 = Dict([
	(:help, sg -> sino_geom_help()),

	(:dim, sg -> (sg.nb, sg.na)),
	(:w, sg -> (sg.nb-1)/2 + sg.offset),
	(:ones, sg -> ones(Float32, sg.dim)),
	(:zeros, sg -> zeros(Float32, sg.dim)),

	(:dr, sg -> sg.how === :moj ? NaN : sg.d),
	(:ds, sg -> sg.how === :moj ? NaN : sg.d),
	(:r, sg -> sg.d * ((0:sg.nb-1) .- sg.w)),
	(:s, sg -> sg.r), # sample locations ('radial')

	(:gamma, sg -> sino_geom_gamma(sg)),
	(:gamma_max, sg -> maximum(abs.(sg.gamma))),
	(:orbit_short, sg -> 180 + 2 * rad2deg(sg.gamma_max)),
	(:ad, sg -> (0:sg.na-1)/sg.na * sg.orbit .+ sg.orbit_start),
	(:ar, sg -> deg2rad.(sg.ad)),

	(:rfov, sg -> sino_geom_rfov(sg)),
	(:xds, sg -> sino_geom_xds(sg)),
	(:yds, sg -> sino_geom_yds(sg)),
	(:dso, sg -> sg.dsd - sg.dod),
	(:grid, sg -> sino_geom_grid(sg)),
	(:plot_grid, sg -> sino_geom_plot_grid(sg)),

	# angular dependent d for :moj
	(:d_ang, sg -> sg.d * max.(abs.(cos.(sg.ar)), abs.(sin.(sg.ar)))),

	(:shape, sg -> ((x::AbstractArray{<:Number} -> reshape(x, sg.dim..., :)))),
	(:taufun, sg -> ((x,y) -> sino_geom_taufun(sg,x,y))),
	(:unitv, sg -> ((;kwarg...) -> sino_geom_unitv(sg; kwarg...))),
	(:plot, sg -> ((;ig) -> sino_geom_plot(sg, ig=ig))),

	# functions that return new geometry:

	(:down, sg -> (down::Int -> downsample(sg, down))),
	(:over, sg -> (over::Int -> sino_geom_over(sg, over))),

	])

# Tricky overloading here!

Base.getproperty(sg::MIRT_sino_geom, s::Symbol) =
		haskey(sino_geom_fun0, s) ? sino_geom_fun0[s](sg) :
		getfield(sg, s)

Base.propertynames(sg::MIRT_sino_geom) =
	(fieldnames(typeof(sg))..., keys(sino_geom_fun0)...)

"""
`sino_geom_plot_grid()`
scatter plot of (r,phi) sampling locations from `sg.grid`
"""
function sino_geom_plot_grid(sg::MIRT_sino_geom)
	(r, phi) = sg.grid
	dfs = sg.how === :fan ? " dfs=$(sg.dfs)" : ""
	ylim = [min(0, rad2deg(minimum(phi))), max(360, rad2deg(maximum(phi)))]
	scatter(r, rad2deg.(phi), label="", markersize=1, ylim = ylim,
		title="$(sg.how)$dfs", ytick=[0,360])
end

"""
`sino_geom_plot_grids()`
scatter plot of (r,phi) sampling locations for all geometries
"""
function sino_geom_plot_grids( ; orbit::Real = 360, down::Int = 30)
	geoms = (
		sino_geom(:par ; nb = 888, na = 984, down=down, d = 0.5, orbit=orbit,
			offset = 0.25),
		sino_geom(:fan ; nb = 888, na = 984, d = 1.0, orbit = orbit,
			offset = 0.75, dsd = 949, dod = 408, down=down),
		sino_geom(:fan ; nb = 888, na = 984, d = 1.0, orbit = orbit,
			offset = 0.75, dsd = 949, dod = 408, down=down,
			dfs = Inf, source_offset = 0.7), # flat fan
		sino_geom(:moj ; nb = 888, na = 984, down=down, d = 1.0, orbit=orbit,
			offset = 0.25),
	)

	ngeom = length(geoms)
	pl = Array{Plot}(undef, ngeom)

	for ii=1:ngeom
		sg = geoms[ii]
		pl[ii] = sg.plot_grid
	end
	plot(pl...)
end

"""
`sino_geom_plot()`
picture of the source position / detector geometry
"""
function sino_geom_plot(sg; ig::Union{Nothing,MIRT_image_geom}=nothing)
	plot(aspect_ratio=1)

	xmax = sg.rfov; xmin = -xmax; (ymin,ymax) = (xmin,xmax)
	if !isnothing(ig)
		plot!(jim(ig.x, ig.y, ig.mask[:,:,1]), clim=[0,1])
		xmin = minimum(ig.x); xmax = maximum(ig.x)
		ymin = minimum(ig.y); ymax = maximum(ig.y)
	end
	plot!([xmax, xmin, xmin, xmax, xmax],
		[ymax, ymax, ymin, ymin, ymax], color=:green, label="")
	plot!(xtick=round.([xmin, 0, xmax], digits=0))
	plot!(ytick=round.([ymin, 0, ymax], digits=2))

	t = LinRange(0, 2*pi, 1001)
	rfov = sg.rfov
	scatter!([0], [0], marker=:circle, label="")
	plot!(rfov * cos.(t), rfov * sin.(t), color=:magenta, label="") # rfov circle
	rfov = round(sg.rfov, digits=1)
	plot!(xlabel="x", ylabel="y", title = "$(sg.how): rfov = $rfov")

#=
	if sg.how === :par
	end
=#

	if sg.how === :fan
		x0 = 0
		y0 = sg.dso
		t = LinRange(0, 2*pi, 100)
		rot = sg.ar[1]
		rot = [cos(rot) -sin(rot); sin(rot) cos(rot)]
		p0 = rot * [x0; y0]
		pd = rot * [sg.xds'; sg.yds'] # detector points

		tmp = sg.ar .+ pi/2 # trick: angle beta defined ccw from y axis
		scatter!([p0[1]], [p0[2]], color=:yellow, label="") # source
		plot!(sg.dso * cos.(t), sg.dso * sin.(t), color=:cyan, label="") # source circle
		scatter!(sg.dso * cos.(tmp), sg.dso * sin.(tmp),
			color=:cyan, markersize=1, label="") # source points
		scatter!(pd[1,:][:], pd[2,:][:],
			color=:yellow, markersize=1, label="") # detectors

		plot!([pd[1,1], p0[1], pd[1,end]], [pd[2,1], p0[2], pd[2,end]],
			color=:red, label="")
		plot!(title="$(sg.how): dfs = $(sg.dfs)")
	end

	if sg.how === :moj && false
		t = LinRange(0, 2*pi, 100)
		rmax = maximum(sg.s)
		rphi = sg.nb/2 * sg.d ./ (max(abs.(cos.(t)), abs.(sin.(t))))
		plot!(rphi .* cos.(t), rphi .* sin.(t), label="") # fov circle
	#	axis([-1 1 -1 1] * max([rmax ig.fov/2]) * 1.1)
	end

	plot!()
end

"""
`sino_geom_ge1()`
sinogram geometry for GE lightspeed system
These numbers are published in IEEE T-MI Oct. 2006, p.1272-1283 wang:06:pwl
"""
function sino_geom_ge1( ;
		na::Int = 984,
		nb::Int = 888,
		orbit::Union{Symbol,Real} = 360,
		units::Symbol = :mm, # default units is mm
		kwarg...)

	if orbit === :short
		na = 642 # trick: reduce na for short scans
		orbit = na/984*360
	end

	scale = units === :mm ? 1 :
			units === :cm ? 10 :
			throw("units $units")

	return sino_geom(:fan ; units=units,
			nb=nb, na=na,
			d = 1.0239/scale, offset = 1.25,
			dsd = 949.075/scale,
			dod = 408.075/scale,
			dfs = 0, kwarg...)
end

"""
`sino_geom_show()`
show an example of each of the 4 main geometries
"""
function sino_geom_show( ; kwarg...)
	down = 4
	ig = image_geom(nx=512, fov=500)
	ig = ig.down(down)

	arg = (nb = 888, na = down*8)
	nb_moj = round(Int, ig.nx*down*sqrt(2)) # match rfov
	sg_list = (
		sino_geom(:par ; d=ig.fovs[1]/arg.nb, arg..., kwarg...),
		sino_geom(:moj ; d=1, arg..., nb=nb_moj, kwarg...),
		sino_geom(:ge1 ; dfs=0, arg..., kwarg...), # arc
		sino_geom(:ge1 ; dfs=Inf, arg..., nb=960, kwarg...), # flat
		)

	nsg = length(sg_list)
	pl = Array{Plot}(undef, nsg)
	for ii = 1:nsg
		sg = sg_list[ii].down(down)
		pl[ii] = sino_geom_plot(sg)
		plot!(pl[ii], xlim=[-1,1]*550, ylim=[-1,1]*550)
		plot!(pl[ii], xtick=[-1,0,1]*250)
	end
	plot(pl...)
end

"""
`sino_geom_test()`
self test
"""
function sino_geom_test()
	down = 8
	ig = image_geom(nx=512, fov=500).down(down)
	ig = image_geom(nx=ig.nx, dx=ig.dx, mask=ig.circ())

	sg_list = (
		sino_geom(:par ; down=down, d=4),
		sino_geom(:moj ; down=down, d=4*sqrt(2)),
		sino_geom(:ge1 ; down=down, orbit_start=20, dfs=0), # arc
		sino_geom(:ge1 ; down=down, orbit_start=20, dfs=Inf), # flat
		sino_geom(:fan ; down=down, orbit=:short),
		)

	sg_list[1].help

	nsg = length(sg_list)
	pl = Array{Plot}(undef, nsg)
	for ii = 1:nsg
		sg = sg_list[ii]

		sg.ad[2]
		sg.rfov
		@inferred sg.down(2)
		@inferred sg.over(2)
		sg.dim
		sg.w
		sg.ones
		sg.zeros
		sg.dr
		sg.ds
		sg.r
		sg.s
		sg.gamma
		sg.gamma_max
		sg.orbit_short
		sg.ad
		sg.ar
		sg.xds
		sg.yds
		sg.dfs
		sg.dso
		sg.grid
		sg.plot_grid

		sg.d_ang # angular dependent d for :moj

		sg.shape(sg.ones[:])
		sg.taufun(ig.x, 0*ig.x)
		sg.unitv()
		sg.unitv(ib=1, ia=2)

		pl[ii] = sg.plot(ig=ig)
	end

	plot(pl[1:4]...)
	gui()

	sg = sino_geom(:ge1 ; orbit=:short)
	sino_geom_gamma_dfs(sg)
	display(sg)
	sino_geom_help()
	sino_geom_plot_grids()
	sino_geom_show()

	@inferred sino_geom(:ge1 ; units=:cm)
	@test_throws String sino_geom(:badhow)
	@test_throws String sino_geom(:ge1 ; dfs=-1)
	@test_throws String sino_geom(:ge1 ; units=:bad)
	true
end


text:
# TODO: try/test with this

# function testNHBaroWaveSphere
# clear all
# close all
using CGDycore

# Physical parameters
Param=CGDycore.PhysParameters();

Param.Upwind = false
Param.RefProfile = false
# Grid
nz=10;
Param.nPanel=4;
Param.H=30000;
Param.Grid=CGDycore.CubedGrid(Param.nPanel,CGDycore.OrientFaceSphere,Param);

Param.Grid.nz=nz;
Param.Grid.zP=zeros(nz,1);
Param.Grid.z=zeros(nz+1,1);
Param.Grid.dz=Param.H/nz;
Param.Grid.zP[1]=Param.Grid.dz/2;
for i=2:nz
  Param.Grid.zP[i]=Param.Grid.zP[i-1]+Param.Grid.dz;
end
for i=2:nz+1
  Param.Grid.z[i]=Param.Grid.z[i-1]+Param.Grid.dz;
end

# Model
Param.ModelType="Curl";
Param.Deep=false;
Param.HeightLimit=30000.0;
Param.T0E=310.0;
Param.T0P=240.0;
Param.B=2.0;
Param.K=3.0;
Param.LapseRate=0.005;
Param.U0=-0.5;
Param.PertR=1.0/6.0;
Param.Up=1.0;
Param.PertExpR=0.1;
Param.PertLon=pi/9.0;
Param.PertLat=2.0*pi/9.0;
Param.PertZ=15000.0;

Param.StrideDamp=6000;
Param.Relax=1.e-4;
Param.Damping=false;
Param.Coriolis=true;
Param.CoriolisType="Sphere";
Param.Buoyancy=true;
Param.Source=false;
Param.Th0=300;
Param.uMax=10;
Param.vMax=0;
Param.NBr=1.e-2;
Param.DeltaT=1;
Param.ExpDist=5;
Param.T0=300;
Param.Equation="Compressible";
Param.TopoS="";
Param.lat0=0;
Param.lon0=pi/2;
Param.ProfVel="BaroWaveSphere";
Param.ProfRho="BaroWaveSphere";
Param.ProfTheta="BaroWaveSphere";
Param.NumV=5;
Param.RhoPos=1;
Param.uPos=2;
Param.vPos=3;
Param.wPos=4;
Param.ThPos=5;
Param.Thermo="";#"Energy"
#Held Suarez
Param.day=3600*24;
Param.k_a=1/(40 * Param.day);
Param.k_f=1/Param.day;
Param.k_s=1/(4*Param.day);
Param.DeltaT_y=60;
Param.DeltaTh_z=10;
Param.T_equator=315;
Param.T_min=200;
Param.sigma_b=7/10;
Param.z_D=20.0e3;

# Discretization
OrdPoly=4;
OrdPolyZ=1;
(CG,Param)=CGDycore.Discretization(OrdPoly,OrdPolyZ,CGDycore.JacobiSphere3,Param);
LRef=11*1.e5;
dx=2*pi*Param.RadEarth/4/Param.nPanel/OrdPoly;
Param.Upwind=false;
Param.HyperVisc=true;
Param.HyperDCurl=2.e17/4; #1.e14*(dx/LRef)^3.2;
Param.HyperDGrad=2.e17/4;
Param.HyperDDiv=2.e17/4; # Scalars

# Output
Param.RadPrint=Param.H;
Param.Flat=true;
Param.vtkFileName="BaroWaveSphere";
Param.vtk=0;
vtkGrid=CGDycore.vtkCGGrid(CG,CGDycore.TransSphere,CGDycore.Topo,Param);
Param.cNames = [
  "Rho",
  "u",
  "v",
  "w",
  "Th"
]

# Initial conditions

U=zeros(CG.NumG,nz,Param.NumV);
U[:,:,Param.RhoPos]=CGDycore.Project(CGDycore.fRho,CG,Param);
(U[:,:,Param.uPos],U[:,:,Param.vPos])=CGDycore.ProjectVec(CGDycore.fVel,CG,Param);
U[:,:,Param.ThPos]=CGDycore.Project(CGDycore.fTheta,CG,Param).*U[:,:,Param.RhoPos];
PresStart=CGDycore.Pressure(U[:,:,Param.ThPos],U[:,:,Param.ThPos],U[:,:,Param.ThPos],Param);
ThB=CGDycore.Project(CGDycore.fThetaBGrd,CG,Param);

# Integration
time=[0];

#IntMethod="Rosenbrock";
IntMethod="RungeKutta";
if strcmp(IntMethod,"Rosenbrock")
  dtau=200;
else
  dtau=8;
end
nIter=20000;
Param.RK=CGDycore.RungeKuttaMethod("RK4");
Param.ROS=CGDycore.RosenbrockMethod("ROSRK3");
SimDays=10;
PrintDay=.5;
nIter=24*3600*SimDays/dtau;
PrintInt=24*3600*PrintDay/dtau;
# Print initial conditions
Param.vtk=CGDycore.vtkOutput(U,vtkGrid,CG,Param);
#

str = IntMethod
if str == "Rosenbrock"
    @time begin
      for i=1:nIter
        @info "Iteration: $i"
        U .= CGDycore.RosenbrockSchur(U,dtau,CGDycore.FcnNHCurlVec,CGDycore.JacSchur,CG,Param);
        time[1] += dtau;
        if mod(i,PrintInt)==0
          Param.vtk=CGDycore.vtkOutput(U,vtkGrid,CG,Param);
        end
      end
    end
    error("Success!")

elseif str == "RungeKutta"
    for i=1:nIter
      @info "Iteration: $i"

      U .= CGDycore.RungeKuttaExplicit(U,dtau,CGDycore.FcnNHCurlVec,CG,Param);

      time[1] += dtau;
      if mod(i,PrintInt)==0
        Param.vtk=CGDycore.vtkOutput(U,vtkGrid,CG,Param);
      end
    end
else
  error("Bad str")
end


text:
"""
$(TYPEDEF)
"""
struct NonlinearSolution{T,N,uType,R,P,A,O,uType2} <: AbstractNonlinearSolution{T,N}
  u::uType
  resid::R
  prob::P
  alg::A
  retcode::Symbol
  original::O
  left::uType2
  right::uType2
end

const SteadyStateSolution = NonlinearSolution

function build_solution(prob::AbstractNonlinearProblem,
                        alg,u,resid;calculate_error = true,
                        retcode = :Default,
                        original = nothing,
                        left = nothing,
                        right = nothing,
                        kwargs...)

  T = eltype(eltype(u))
  N = length((size(prob.u0)...,))

  NonlinearSolution{T,N,typeof(u),typeof(resid),
                    typeof(prob),typeof(alg),typeof(original),typeof(left)}(
                    u,resid,prob,alg,retcode,original,left,right)
end

function sensitivity_solution(sol::AbstractNonlinearSolution,u)
  T = eltype(eltype(u))
  N = length((size(sol.prob.u0)...,))

  NonlinearSolution{T,N,typeof(u),typeof(sol.resid),
                    typeof(sol.prob),typeof(sol.alg),
                    typeof(sol.original),typeof(sol.left)}(
                    u,sol.resid,sol.prob,sol.alg,sol.retcode,sol.original,sol.left,sol.right)
end


text:
#
# This is a minimal working example for the imbalanced threads problem
#
function machin_series(n::Int; handover=false)
    qpi = 0.0
    for i in 1:n
        qpi += (-1)^(i+1)/(2i-1)
        handover && yield()  # we yield here !!!!
    end
    qpi*4
end

function startonthread(id::Int, f::F) where {F<:Function}
    t = Task(nothing)
    @threads for i in 1:nthreads()
        if i == id
            t = @async f()
        end
    end
    fetch(t)
end


text:
using MultivariateFunctions
using DataStructures: OrderedDict
tol = 1e-15
fun1 = PE_Unit(1.0,1.0,1)
fun2 = PE_Unit(1.0,1.0,1)
fun3 = PE_Unit(1.0,2.0,3)
fun4 = PE_Unit(1.0,3.0,2)
fun6 = PE_Unit(0.0,1.0,1)
fun7 = PE_Unit(0.0,2.0,4)
fun8 = PE_Unit(0.0,2.0,2)

mfun1 = PE_Function(1.8, Dict(Symbol.(["x", "y", "z"]) .=> [fun1, fun2, fun3]))
abs(mfun1.multiplier_ - 1.2*1.5) < 1e-10
mfun2 = PE_Function(4.5, Dict(Symbol.(["x", "y"]) .=> [fun3, fun4]))
mfun3 = PE_Function(3.0, Dict(Symbol.(["z", "y"]) .=> [fun4, fun1]))
mfun4 = PE_Function(6.0, Dict(Symbol.(["x", "y"]) .=> [fun6, fun7]))
mfun5 = PE_Function(6.0, Dict(Symbol.(["x", "z"]) .=> [fun6, fun8]))

msum = mfun5 + mfun4

func_array = Array{Union{Sum_Of_Functions,Missing},3}(undef,3,2,2)
func_array[1,1,1] = Sum_Of_Functions(mfun1)
func_array[1,1,2] = Sum_Of_Functions(mfun2)
func_array[1,2,1] = mfun2 + mfun1
func_array[1,2,2] = mfun2 + mfun4
func_array[2,1,1] = Sum_Of_Functions(2*mfun1)
func_array[2,1,2] = mfun1 + mfun5
func_array[2,2,1] = mfun1 + mfun5
func_array[2,2,2] = mfun2 + mfun1 + mfun5
func_array[3,:,:] .= Missing()

thresholds_ = OrderedDict{Symbol,Array{Float64,1}}([:x, :y, :z] .=> [[0.0, 0.5, 1.0], [-5.0, -1.0], [0.0, 5.0]])

pw_func = Piecewise_Function(func_array, thresholds_)

# Testing zeroing.
coordinates = Dict{Symbol,Float64}([:w, :x, :y, :z] .=> [1.4, 0.4, 2.0, 3.0])
hypercube =  Dict{Symbol,Tuple{Float64,Float64}}([:x, :z] .=> [(0.5,Inf), (1.0,6.0)])
zero_outside_hypercube = Piecewise_Function(pw_func, hypercube, true)
abs(evaluate(pw_func, coordinates) - evaluate(zero_outside_hypercube, coordinates)) > 5.0
coordinates_in_cube = Dict{Symbol,Float64}([:w, :x, :y, :z] .=> [1.4, 0.7, 2.0, 3.0])
abs(evaluate(pw_func, coordinates_in_cube) - evaluate(zero_outside_hypercube, coordinates_in_cube)) < tol

# creating a pw with a pw
functions_ = Array{Any,2}(undef,3,2)
functions_[1,1] = Sum_Of_Functions(mfun1)
functions_[1,2] = Sum_Of_Functions(3.0*mfun2)
functions_[2,1] = pw_func
functions_[2,2] = mfun2 + mfun4
functions_[3,1] = pw_func + 1.0
functions_[3,2] = mfun2 + mfun5
thresholds_ = OrderedDict{Symbol,Array{Float64,1}}([:x, :y] .=> [[0.0, 0.6, 1.0], [-5.0, -1.0]])
pw_func2 = Piecewise_Function(functions_, thresholds_)
coords_21 = Dict{Symbol,Float64}([:w, :x, :y, :z] .=> [1.4, 0.7, -3.0, 3.0])
coords_22 = Dict{Symbol,Float64}([:w, :x, :y, :z] .=> [1.4, 0.7, 0.0, 3.0])
abs(evaluate(pw_func, coords_21) - evaluate(pw_func2, coords_21)) < tol
abs(evaluate(pw_func, coords_22) - evaluate(pw_func2, coords_22)) > 1.0

###############################################################################
coordinates = Dict(Symbol.(["w","x", "y", "z"]) .=> [1.4, 0.5, 2.0, 3.0])
function test_result(func, eval_to, len = 1)
    val_test = abs(evaluate(func, coordinates) - eval_to) < 1e-05
    if (!val_test)
        println("Failed Val Test")
    end
    len_test = length(func.functions_) == len
    if (!len_test)
        println(string("Failed length Test, Length is ", length(func.functions_)))
    end
    return all([val_test,len_test])
end
ismissing(evaluate(pw_func, Dict(Symbol.(["w","x", "y", "z"]) .=> [-Inf, -Inf, -Inf, -Inf])))
test_result(pw_func, evaluate( mfun1 + mfun5, coordinates), 12)
deriv = Dict{Symbol,Int}(Symbol.(["x", "z"]) .=> [2,1])
abs(evaluate(derivative(pw_func, deriv), coordinates) - evaluate(derivative(mfun1 + mfun5, deriv), coordinates)) < tol

# Integration of piecewise function
integration_limits = Dict{Symbol,Tuple{Float64,Float64}}([:x, :y, :z] .=> [(0.25,0.75),(-4.0,-1.0),(0.0,3.0)])
integral_of_each_piece = integral(Sum_Of_Functions(mfun1) ,Dict{Symbol,Tuple{Float64,Float64}}([:x, :y, :z] .=> [(0.25,0.5),(-4.0,-1.0),(0.0,3.0)]))
integral_of_each_piece = integral_of_each_piece + integral(Sum_Of_Functions(2*mfun1) ,Dict{Symbol,Tuple{Float64,Float64}}([:x, :y, :z] .=> [(0.5,0.75),(-4.0,-1.0),(0.0,3.0)]))
integral_of_piecewise = integral(pw_func, integration_limits)
abs(integral_of_each_piece - integral_of_piecewise)  < tol
# Integration of piecewise function
integration_limits = Dict{Symbol,Tuple{Float64,Float64}}([:x, :y, :z] .=> [(0.25,0.75),(-4.0,-1.0),(0.0,8.0)])
integral_of_each_piece = integral(Sum_Of_Functions(mfun1) ,Dict{Symbol,Tuple{Float64,Float64}}([:x, :y, :z] .=> [(0.25,0.5),(-4.0,-1.0),(0.0,5.0)]))
integral_of_each_piece = integral_of_each_piece + integral(Sum_Of_Functions(2*mfun1) ,Dict{Symbol,Tuple{Float64,Float64}}([:x, :y, :z] .=> [(0.5,0.75),(-4.0,-1.0),(0.0,5.0)]))
integral_of_each_piece = integral_of_each_piece + integral(Sum_Of_Functions(mfun2) ,Dict{Symbol,Tuple{Float64,Float64}}([:x, :y, :z] .=> [(0.25,0.5),(-4.0,-1.0),(5.0,8.0)]))
integral_of_each_piece = integral_of_each_piece + integral(mfun1 + mfun5 ,Dict{Symbol,Tuple{Float64,Float64}}([:x, :y, :z] .=> [(0.5,0.75),(-4.0,-1.0),(5.0,8.0)]))
integral_of_piecewise = integral(pw_func, integration_limits)
abs(integral_of_each_piece - integral_of_piecewise)  < tol

# test algebra between Piecewise_Functions and PE_Functions.
test_result(pw_func + mfun4, evaluate( mfun1 + mfun5, coordinates) + evaluate(mfun4, coordinates), 12)
test_result(pw_func - mfun4, evaluate( mfun1 + mfun5, coordinates) - evaluate(mfun4, coordinates), 12)
test_result(pw_func * mfun4, evaluate( mfun1 + mfun5, coordinates) * evaluate(mfun4, coordinates), 12)
test_result(mfun4 + pw_func, evaluate( mfun1 + mfun5, coordinates) + evaluate(mfun4, coordinates), 12)
test_result(mfun4 - pw_func,  evaluate(mfun4, coordinates) - evaluate( mfun1 + mfun5, coordinates), 12)
test_result(mfun4 * pw_func, evaluate( mfun1 + mfun5, coordinates) * evaluate(mfun4, coordinates), 12)
# Multivariate Sum of Functions
test_result(pw_func + msum, evaluate( mfun1 + mfun5, coordinates) + evaluate(msum, coordinates), 12)
test_result(pw_func - msum, evaluate( mfun1 + mfun5, coordinates) - evaluate(msum, coordinates), 12)
test_result(pw_func * msum, evaluate( mfun1 + mfun5, coordinates) * evaluate(msum, coordinates), 12)
test_result(msum + pw_func, evaluate( mfun1 + mfun5, coordinates) + evaluate(msum, coordinates), 12)
test_result(msum - pw_func,  evaluate(msum, coordinates) - evaluate( mfun1 + mfun5, coordinates), 12)
test_result(msum * pw_func, evaluate( mfun1 + mfun5, coordinates) * evaluate(msum, coordinates), 12)

# Algebra between different Multivariate_Piecewise_Functions
func_array = Array{Union{Missing,Sum_Of_Functions},3}(undef,3,2,2)
func_array[1,1,1] = Sum_Of_Functions(mfun5)
func_array[1,1,2] = Sum_Of_Functions(mfun3)
func_array[1,2,1] = mfun3 + mfun1
func_array[1,2,2] = Sum_Of_Functions(2*mfun1)
func_array[2,1,1] = mfun2 + mfun4
func_array[2,1,2] = mfun1 + mfun5
func_array[2,2,1] = mfun5 + mfun5
func_array[2,2,2] = mfun4 + mfun1 + mfun5
func_array[3,:,:] .= mfun4 + mfun5 + mfun1

pw_func2 = Piecewise_Function(func_array, OrderedDict([:x, :y, :w] .=> [[0.2, 0.5, 1.0], [-4.0, -1.0], [0.0, 5.0]]))
test_result(pw_func + pw_func2, evaluate( (mfun1 + mfun5) + (mfun5 + mfun5), coordinates), 3*4*2*2)
test_result(pw_func - pw_func2, evaluate( (mfun1 + mfun5) - (mfun5 + mfun5), coordinates), 3*4*2*2)
test_result(pw_func * pw_func2, evaluate( (mfun1 + mfun5) * (mfun5 + mfun5), coordinates), 3*4*2*2)

## Testing of Sum_Of_Piecewise_Functions
spwf = Sum_Of_Piecewise_Functions([pw_func, pw_func2])
abs(evaluate(spwf, coordinates) - evaluate(pw_func, coordinates) - evaluate(pw_func2, coordinates) ) < tol
underlying_dimensions(spwf) == union(underlying_dimensions(pw_func), underlying_dimensions(pw_func2))
converted = convert(Piecewise_Function, spwf)
added = (pw_func + pw_func2)
converted.functions_[2,2,2,1].functions_[1] ≂ added.functions_[2,2,2,1].functions_[1]

test_result( (spwf + 1) - 1, evaluate( spwf, coordinates), 2)
test_result( 1 + (1 - spwf ), 2 - evaluate( spwf, coordinates), 2)
test_result( (spwf + 1.0) - 1.0, evaluate( spwf, coordinates), 2)
test_result( 1.0 + (1.0 - spwf ), 2 - evaluate( spwf, coordinates), 2)
test_result( (3 *spwf * 2) / 4, evaluate( spwf, coordinates) * 1.5, 2)
test_result( (3.0 *spwf * 2.0) / 4.0, evaluate( spwf, coordinates) * 1.5, 2)
test_result( (spwf^2), evaluate( spwf, coordinates)^2, 48)
deriv_spec = Dict{Symbol,Int}(:x => 1)
abs(evaluate(derivative(spwf, deriv_spec), coordinates) -  evaluate(derivative(pw_func,deriv_spec), coordinates) - evaluate(derivative(pw_func2, deriv_spec), coordinates)) < 1e-09
integration_limits = Dict{Symbol,Tuple{Float64,Float64}}([:x, :y, :w, :z] .=> [(0.25,0.75),(-4.0,-1.0),(0.0,8.0), (0.5, 2.4)])
evaluate(integral(spwf, integration_limits) -  integral(pw_func,integration_limits) - integral(pw_func2, integration_limits), coordinates)     < 1e-09


text:
`(
    (width        . 640       )
    (height       . 480       )
    (iteration    . 128       )
	(init-number  . -0.7269+0.1889i )
	(scale-weight . (1.5 . 1) )
    (color-depth  . (lambda (depth)
                      (let* ((x (/ depth 128))
					         (r (min (* x 256) 255))
							 (g (min (* (- (* x 256) r) 256) 255))
							 (b (min (* (- (* (- (* x 256) r) x 256) g) 256) 255)))
					    (list r g b)))))

max_stars_repo_name: fanguangping/easy-fractal

max_issues_repo_name: fanguangping/easy-fractal

max_forks_repo_name: fanguangping/easy-fractal


text:
@discord_object struct Overwrite
    id::Snowflake
    type::String
    allow::Int64
    deny::Int64
end


text:
using DataFrames, CSV, Plots, StatsPlots

ProjDir = @__DIR__
include(joinpath(ProjDir, "fn.jl"))

nts = [1, 2, 4, 8]
nccs = [1, 2, 4, 8]
ncjs = [1, 2, 4, 8]
nss = [10000, 5000, 2500, 1250]
N = 6

function timings(nts, nccs, ncjs, nss, N)

    df = DataFrame()
    res_t = Vector{Float64}(undef, N)

    for nt in nts
        for nc in nccs
            for nj in ncjs
                println("\n\n(num_threads=$nt, num_cpp-chains=$nc, num_chains=$nj) runs\n\n")
                for ns in nss
                    if nc * nj * ns == 10000
                        for i in 1:N
                            res_t[i] = @elapsed stan_inference(prob_ode_fitzhughnagumo,
                                t,data,priors; num_threads=nt, num_cpp_chains=nc,
                                num_chains=nj, num_samples=ns,
                                output_format = :dataframe, diffeq_string, tmpdir);
                        end
                        append!(df, DataFrame(
                            num_threads=nt,
                            num_cpp_chains=nc,
                            num_chains=nj,
                            num_samples=ns,
                            min=minimum(res_t),
                            median=median(res_t),
                            max=maximum(res_t))
                        )
                    end
                end
            end
        end
    end

    df
end

df = timings(nts, nccs, ncjs, nss, N)
df |> display

CSV.write(joinpath(ProjDir, "results", "new_results_df.csv"), df)


text:
using OrdinaryDiffEq, Test
using DiffEqProblemLibrary.ODEProblemLibrary: importodeproblems; importodeproblems()
import DiffEqProblemLibrary.ODEProblemLibrary: prob_ode_2Dlinear
prob = prob_ode_2Dlinear

integrator = init(prob,Euler();dt=1//2^(4))
integrator = init(prob,Tsit5();dt=1//2^(4))

solve!(integrator)

sol =solve(prob,Euler();dt=1//2^(4),maxiters=Inf)
sol =solve(prob,Tsit5();dt=1//2^(20),progress=true,adaptive=false,maxiters=Inf)

# plot(sol,plot_analytic=true)

sol =solve(prob,ExplicitRK();dt=1//2^(4))

# out-of-place 2D prob to test vec & reshape
rhs_oop(u,p,t) = u
ft0=ones(2, 2)
tspan=(0.0,30.0)
prob=ODEProblem(rhs_oop,ft0,tspan)
@test_nowarn solve(prob,Rosenbrock23())
@test_nowarn solve(prob,Rodas4())
@test_nowarn solve(prob,Kvaerno5())


text:
#  Copyright 2016, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at http://mozilla.org/MPL/2.0/.

function writeMPS(m::Model, fname::AbstractString)
    f = open(fname, "w")

    write(f,"NAME   JuMPModel\n")

    numRows = length(m.linconstr)

    # Objective and constraint names
    gc_enable(false)
    write(f,"ROWS\n")
    write(f," N  OBJ\n")
    hasrange = false
    for c in 1:numRows
        rowsense = sense(m.linconstr[c])
        if rowsense == :(<=)
            senseChar = 'L'
        elseif rowsense == :(==)
            senseChar = 'E'
        elseif rowsense == :(>=)
            senseChar = 'G'
        else
            hasrange = true
            senseChar = 'E'
        end
        @printf(f," %c  CON%d\n",senseChar,c)
    end
    gc_enable(true)

    objlincoef, rowlb, rowub = prepProblemBounds(m)
    if m.objSense == :Max
        println("Warning, MPS does not support maximization sense. Flipping objective coefficients.")
        objlincoef = -objlincoef
    end

    A = prepConstrMatrix(m)
    colptr = A.colptr
    rowval = A.rowval
    nzval = A.nzval

    # Output each column
    gc_enable(false)
    inintegergroup = false
    write(f,"COLUMNS\n")
    for col in 1:m.numCols
        t = m.colCat[col]
        (t == :SemiCont || t == :SemiInt) && error("The MPS file writer does not currently support semicontinuous or semi-integer variables")
        if (t == :Bin || t == :Int) && !inintegergroup
            @printf(f,"    MARKER    'MARKER'                 'INTORG'\n")
            inintegergroup = true
        elseif (t == :Cont || t == :Fixed) && inintegergroup
            @printf(f,"    MARKER    'MARKER'                 'INTEND'\n")
            inintegergroup = false
        end
        for ind in colptr[col]:(colptr[col+1]-1)
            @printf(f,"    VAR%d  CON%d  ",col,rowval[ind])
            print_shortest(f,nzval[ind])
            println(f)
        end
        @printf(f,"    VAR%d  OBJ  ",col)
        print_shortest(f,objlincoef[col])
        println(f)
    end
    if inintegergroup
        @printf(f,"    MARKER    'MARKER'                 'INTEND'\n")
    end
    gc_enable(true)

    # RHSs
    gc_enable(false)
    write(f,"RHS\n")
    for c in 1:numRows
        rowsense = sense(m.linconstr[c])
        if rowsense != :range
            @printf(f,"    rhs    CON%d    ",c)
            print_shortest(f,rhs(m.linconstr[c]))
        else
            @printf(f,"    rhs    CON%d    ",c)
            print_shortest(f,m.linconstr[c].lb)
        end
        println(f)
    end
    gc_enable(true)

    # RANGES
    if hasrange
        gc_enable(false)
        write(f,"RANGES\n")
        for c in 1:numRows
            rowsense = sense(m.linconstr[c])
            if rowsense == :range
                @printf(f,"    rhs    CON%d    ",c)
                print_shortest(f,m.linconstr[c].ub-m.linconstr[c].lb)
                println(f)
            end
        end
    end

    # BOUNDS
    gc_enable(false)
    write(f,"BOUNDS\n")
    for col in 1:m.numCols
        if m.colLower[col] == 0
            if m.colUpper[col] != Inf
                # Default lower 0, and an upper
                @printf(f,"  UP BOUND VAR%d ", col)
                print_shortest(f, m.colUpper[col])
                println(f)
            else
                # Default bounds. Explicitly state for solvers like Gurobi. See issue #792
                @printf(f,"  PL BOUND VAR%d", col)
                println(f)
            end
        elseif m.colLower[col] == -Inf && m.colUpper[col] == +Inf
            # Free
            @printf(f, "  FR BOUND VAR%d\n", col)
        elseif m.colLower[col] != -Inf && m.colUpper[col] == +Inf
            # No upper, but a lower
            @printf(f, "  PL BOUND VAR%d\n  LO BOUND VAR%d ",col,col)
            print_shortest(f,m.colLower[col])
            println(f)
        elseif m.colLower[col] == -Inf && m.colUpper[col] != +Inf
            # No lower, but a upper
            @printf(f,"  MI BOUND VAR%d\n  UP BOUND VAR%d ",col,col)
            print_shortest(f,m.colUpper[col])
            println(f)
        else
            # Lower and upper
            @printf(f, "  LO BOUND VAR%d ",col)
            print_shortest(f,m.colLower[col])
            println(f)
            @printf(f, "  UP BOUND VAR%d ",col)
            print_shortest(f,m.colUpper[col])
            println(f)
        end
    end
    gc_enable(true)

    # Quadratic objective
    gc_enable(false)
    if length(m.obj.qvars1) != 0
        write(f,"QMATRIX\n")
        qv1 = m.obj.qvars1
        qv2 = m.obj.qvars2
        qc  = m.obj.qcoeffs
        for ind = 1:length(qv1)
            if qv1[ind].col == qv2[ind].col
                # Diagonal element
                @printf(f,"  VAR%d VAR%d  ", qv1[ind].col,qv2[ind].col)
                print_shortest(f,2qc[ind])
                println(f)
            else
                # Off diagonal, and we're gonna assume no duplicates
                @printf(f, "  VAR%d VAR%d ", qv1[ind].col,qv2[ind].col)
                print_shortest(f, qc[ind])
                println(f)
                #@printf(f, "  VAR%d VAR%d ", qv2[ind].col,qv1[ind].col)
                #print_shortest(f, qc[ind])
                #println(f)
            end
        end
    end

    write(f,"ENDATA\n")
    close(f)
    gc_enable(true)
end

###############################################################################
# LP File Writer
# We use the formatting defined at:
#   http://lpsolve.sourceforge.net/5.0/CPLEX-format.htm
function writeLP(m::Model, fname::AbstractString)

    f = open(fname, "w")

    if length(m.obj.qvars1) != 0
        error("LP writer does not support quadratic objectives.\n")
    end

    # Objective
    if m.objSense == :Max
        write(f,"Maximize\n")
    else
        write(f,"Minimize\n")
    end
    objaff::AffExpr = m.obj.aff
    write(f, " obj: ")
    nnz = length(objaff.coeffs)
    for ind in 1:(nnz-1)
        if ind == 1
            print_shortest(f, objaff.coeffs[ind])
        else
            print_shortest(f, abs(objaff.coeffs[ind]))
        end
        @printf(f, " VAR%d %s ", objaff.vars[ind].col, (objaff.coeffs[ind+1] < 0)? "-" : "+")
    end
    if nnz >= 1
        if nnz == 1
            print_shortest(f, objaff.coeffs[nnz])
        else
            print_shortest(f, abs(objaff.coeffs[nnz]))
        end
        @printf(f, " VAR%d\n", objaff.vars[nnz].col)
    end

    # Constraints
    function writeconstrterms(c::LinearConstraint)
        nnz = length(c.terms.coeffs)
        for ind in 1:(nnz-1)
            if ind == 1
                print_shortest(f, c.terms.coeffs[ind])
            else
                print_shortest(f, abs(c.terms.coeffs[ind]))
            end
            @printf(f, " VAR%d %s ", c.terms.vars[ind].col, (c.terms.coeffs[ind+1] < 0)? "-" : "+")
        end
        if nnz >= 1
            if nnz == 1
                print_shortest(f, c.terms.coeffs[nnz])
            else
                print_shortest(f, abs(c.terms.coeffs[nnz]))
            end
            @printf(f, " VAR%d", c.terms.vars[nnz].col)
        end
    end
    write(f,"Subject To\n")
    constrcount = 1
    for i in 1:length(m.linconstr)
        @printf(f, " c%d: ", constrcount)

        c::LinearConstraint = m.linconstr[i]
        rowsense = sense(c)
        if rowsense != :range
            writeconstrterms(c)
            if rowsense == :(==)
                @printf(f, " = ")
                print_shortest(f, rhs(c))
                println(f)
            elseif rowsense == :<=
                @printf(f, " <= ")
                print_shortest(f, rhs(c))
                println(f)
            else
                @assert rowsense == :>=
                @printf(f, " >= ")
                print_shortest(f, rhs(c))
                println(f)
            end
            constrcount += 1
        else
            writeconstrterms(c)
            @printf(f, " >= ")
            print_shortest(f, c.lb)
            println(f)
            @printf(f, " c%d: ", constrcount+1)
            writeconstrterms(c)
            @printf(f, " <= ")
            print_shortest(f, c.ub)
            println(f)
            constrcount += 2
        end
    end

    # Bounds
    write(f,"Bounds\n")
    for i in 1:m.numCols
        if m.colLower[i] == -Inf
            # No low bound
            if m.colUpper[i] == +Inf
                # Free
                @printf(f, " VAR%d free\n", i)
            else
                # x <= finite
                @printf(f, " -inf <= VAR%d <= ", i)
                print_shortest(f, m.colUpper[i])
                println(f)
            end
        else
            # Low bound exists
            if m.colUpper[i] == +Inf
                # x >= finite
                @printf(f, " ")
                print_shortest(f, m.colLower[i])
                @printf(f," <= VAR%d <= +inf\n", i)
            else
                # finite <= x <= finite
                @printf(f, " ")
                print_shortest(f, m.colLower[i])
                @printf(f, " <= VAR%d <= ", i)
                print_shortest(f, m.colUpper[i])
                println(f)
            end
        end
    end

    # Integer - don't handle binaries specially
    write(f,"General\n")
    for i in 1:m.numCols
        t = m.colCat[i]
        (t == :SemiCont || t == :SemiInt) && error("The LP file writer does not currently support semicontinuous or semi-integer variables")
        if t == :Bin || t == :Int
            @printf(f, " VAR%d\n", i)
        end
    end

    # Done
    write(f,"End\n")
    close(f)
end


text:
using Revise
using AdFem
using PyCall
using LinearAlgebra
np = pyimport("numpy")

# Domain information
NT = 50
Δt = 1/NT
n = 50
m = n
h = 1.0/n
bdnode = Int64[]
for i = 1:m+1
    for j = 1:n+1
        if i==1 || i==m+1 || j==1|| j==n+1
            push!(bdnode, (j-1)*(m+1)+i)
        end
    end
end

# Physical parameters
b = 1.0
H = [1.0 0.0 0.0
    0.0 1.0 0.0
    0.0 0.0 0.5]
Kr = 0.01*ones(m*n)
for i = 1:m
    for j = div(n,3):div(n,3)*2
        Kr[(j-1)*m+i] = 10.0
    end
end

Q = compute_fvm_tpfa_matrix(Kr, m, n, h)
K = compute_fem_stiffness_matrix(H, m, n, h)
L = compute_interaction_matrix(m, n, h)
M = compute_fvm_mass_matrix(m, n, h)
A = [K -b*L'
b*L/Δt 1/Δt*M-Q]
A, Abd = fem_impose_coupled_Dirichlet_boundary_condition(A, bdnode, m, n, h)

U = zeros(m*n+2(m+1)*(n+1), NT+1)
x = Float64[]; y = Float64[]
for j = 1:n+1
    for i = 1:m+1
        push!(x, (i-1)*h)
        push!(y, (j-1)*h)
    end
end

injection = (div(n,2)-1)*m + 3
production = (div(n,2)-1)*m + m-3
for i = 1:NT
    t = i*Δt
    @info i
    f_ = eval_f_on_gauss_pts((x,y)->f(x,y,t), m, n, h)
    rhs1 = zeros(2(m+1)*(n+1))
    rhs2 =      b*L*U[1:2(m+1)*(n+1), i]/Δt +
                M * U[2(m+1)*(n+1)+1:end,i]/Δt
    rhs2[injection] += 1.0
    rhs2[production] -= 1.0

    rhs = [rhs1;rhs2]

    bdval = zeros(2*length(bdnode))
    rhs[[bdnode; bdnode.+ (m+1)*(n+1)]] = bdval
    rhs -= Abd * bdval
    U[:,i+1] = A\rhs
end

visualize_displacement(U, m, n, h, name="_piheter")
visualize_pressure(U, m, n, h, name="_piheter")

max_stars_repo_name: ADCMEMarket/ADCMEImages

max_issues_repo_name: ADCMEMarket/ADCMEImages

max_forks_repo_name: ADCMEMarket/ADCMEImages


text:
# Want to test that certain subtypes of arrays are compatible
_make_matrix(::Type{Matrix}) = [1 2 3; 4 5 6]
_make_matrix(::Type{SubArray}) = view([1 2 3; 4 5 6], :, :)
_make_matrix(::Type{ReshapedArray}) = reshape([1 2 3; 4 5 6], 2, 3)

@testset "$MatrixType" for MatrixType in (Matrix, SubArray, ReshapedArray)

    x = _make_matrix(MatrixType)

    @test is_transformable(x)

    @testset "apply" begin

        @testset "OneToOne" begin
            T = FakeOneToOneTransform()
            @test FeatureTransforms.apply(x, T) == ones(2, 3)
            @test FeatureTransforms.apply(x, T; inds=[1, 2]) == ones(2)
            @test FeatureTransforms.apply(x, T; dims=2, inds=[2, 3]) == ones(2, 2)
        end

        @testset "OneToMany" begin
            T = FakeOneToManyTransform()
            @test FeatureTransforms.apply(x, T) == ones(2, 6)
            @test FeatureTransforms.apply(x, T; inds=[1, 2]) == ones(2, 2)
            @test FeatureTransforms.apply(x, T; dims=2, inds=[2, 3]) == ones(2, 4)
        end

        @testset "ManyToOne" begin
            T = FakeManyToOneTransform()
            @test FeatureTransforms.apply(x, T; dims=1) == ones(3)
            @test FeatureTransforms.apply(x, T; dims=1, inds=[1, 2]) == ones(3)
            @test FeatureTransforms.apply(x, T; dims=2, inds=[1, 2]) == ones(2)
        end

        @testset "ManyToMany" begin
            T = FakeManyToManyTransform()
            @test FeatureTransforms.apply(x, T; dims=1) == ones(2, 6)
            @test FeatureTransforms.apply(x, T; inds=[1, 2]) == ones(2, 2)
            @test FeatureTransforms.apply(x, T; dims=2, inds=[2, 3]) == ones(2, 4)
        end

    end

    @testset "apply!" begin
        T = FakeOneToOneTransform()

        _x = copy(x)
        FeatureTransforms.apply!(_x, T)
        @test _x == ones(2, 3)

        # https://github.com/invenia/FeatureTransforms.jl/issues/68
        _x = copy(x)
        @test_broken FeatureTransforms.apply!(_x, T; dims=2, inds=[1, 2])
        @test_broken _x == [1 1 3; 1 1 6]
    end

    @testset "apply_append" begin

        @testset "OneToOne" begin

            T = FakeOneToOneTransform()

            @test FeatureTransforms.apply_append(x, T; append_dim=1) == vcat(x, ones(2, 3))
            @test FeatureTransforms.apply_append(x, T; append_dim=2) == hcat(x, ones(2, 3))
            @test FeatureTransforms.apply_append(x, T; append_dim=3) == cat(x, ones(2, 3), dims=3)

            @test isequal(
                FeatureTransforms.apply_append(x, T; dims=1, inds=[1], append_dim=1),
                vcat(x, ones(1, 3))
            )

            @test isequal(
                FeatureTransforms.apply_append(x, T; dims=2, inds=[1, 2], append_dim=2),
                hcat(x, ones(2, 2))
            )
        end

        @testset "OneToMany" begin

            T = FakeOneToManyTransform()

            @test_throws DimensionMismatch FeatureTransforms.apply_append(x, T; append_dim=1)
            @test_throws DimensionMismatch FeatureTransforms.apply_append(x, T; append_dim=3)

            @test isequal(
                FeatureTransforms.apply_append(x, T; append_dim=2),
                hcat(x, ones(2, 6))
            )

            @test isequal(
                FeatureTransforms.apply_append(x, T; inds=[1, 2], append_dim=2),
                hcat(x, ones(2, 2))
            )
        end

        @testset "ManyToOne" begin

            T = FakeManyToOneTransform()
            @test FeatureTransforms.apply_append(x, T; dims=1, append_dim=1) == vcat(x, ones(1, 3))
            @test FeatureTransforms.apply_append(x, T; dims=2, append_dim=2) == hcat(x, ones(2))
        end

        @testset "ManyToMany" begin

            T = FakeManyToManyTransform()

            @test_throws DimensionMismatch FeatureTransforms.apply_append(x, T; append_dim=1)
            @test_throws DimensionMismatch FeatureTransforms.apply_append(x, T; append_dim=3)

            @test isequal(
                FeatureTransforms.apply_append(x, T; append_dim=2),
                hcat(x, ones(2, 6))
            )

            @test isequal(
                FeatureTransforms.apply_append(x, T; inds=[1, 2], append_dim=2),
                hcat(x, ones(2, 2))
            )
        end

    end

end


text:
# julia1_2.2.jl Variables and Scope - Local and Global
#            Static and Dynamic, and Constants
# local vs global
i=0   #(global)
while i==0  # ok since being read
    global i  # because being written below
    x=π  #(soft local) needs global to see outside of while
    i+=1      # need the global since being written
    println(i)
    println(x)  # x is local
end
println(i)
println(x) # error since x is local in while

# "static" type
# state becomes what in C would be called a static type
# It's value persists but is only available inside the let block
# useful for creating pointers in buffers, deques and queues.
let
  state=0
  global counter
  counter() = state+=1
end
println(counter())
println(counter())

# loop indices and externally defined variables local
i=43  # these will NOT be overridden by the loop index i and local l
l=66
for i=1:3  # this i is "soft" local so won't interfere with global i
    j=i+1  # j is also "soft" local to this block and lower
    for k=1:2
        local l #hard local so won't interfere with l=66 global
        l=k+j+i
        println("l",l)
    end
end
println(i) # 1:3 will not override 43
println(l) # note inner for loop l is hard local
println(j) # error! j soft local not available to higher level

 # variables defined within functions are hard locals
 # however a nested function with a variable of the same name
 # can modify a parent's variables. Will cover in functions

# constants are like valiables
# however, they will warn you if you override them
const t=43.5
println(t)
t=42   # warning error

# Built in constants:
π #(alias: pi)
# CODATA.jlexp(1) package has 350 physics constants
const e=2.718284590452353602874713526624977572470937


text:
# ---
# title: 820. Short Encoding of Words
# id: problem820
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories:
# link: <https://leetcode.com/problems/short-encoding-of-words/description/>
# hidden: true
# ---
#
# Given a list of words, we may encode it by writing a reference string `S` and
# a list of indexes `A`.
#
# For example, if the list of words is `["time", "me", "bell"]`, we can write it
# as `S = "time#bell#"` and `indexes = [0, 2, 5]`.
#
# Then for each index, we will recover the word by reading from the reference
# string from that index until we reach a `"#"` character.
#
# What is the length of the shortest reference string S possible that encodes
# the given words?
#
# **Example:**
#
#
#
#     Input: words = ["time", "me", "bell"]
#     Output: 10
#     Explanation: S = "time#bell#" and indexes = [0, 2, 5].
#
#
#
#
# **Note:**
#
#   1. `1 <= words.length <= 2000`.
#   2. `1 <= words[i].length <= 7`.
#   3. Each word has only lowercase letters.
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
function rand_candecomp{N}(r::Int64, dims::NTuple{N, Int};
                           lambdas_nonneg::Bool=false, factors_nonneg::Bool=false)
    rnd_factor = factors_nonneg ? x -> abs.(randn(x...)) : randn
    rnd_lambda = lambdas_nonneg ? x -> abs.(randn(x...)) : randn
    return CANDECOMP((Matrix{Float64}[rnd_factor((s, r)) for s in dims]...),
                     rnd_lambda(r))
end

rand_kruskal3{N}(r::Int64, dims::NTuple{N, Int}, nonnegative::Bool) =
    compose(rand_candecomp(r, dims, lambdas_nonneg=nonnegative, factors_nonneg=nonnegative))

function rand_tucker{N}(core_dims::NTuple{N, Int}, dims::NTuple{N, Int};
                        core_nonneg::Bool=false, factors_nonneg::Bool=false)
    rnd_factor = factors_nonneg ? x -> abs.(randn(x...)) : randn
    rnd_core = core_nonneg ? x -> abs.(randn(x...)) : randn
    return Tucker((Matrix{Float64}[rnd_factor((dims[i], core_dims[i])) for i in 1:N]...),
                  rnd_core(core_dims))
end

function add_noise{T, N}(tnsr::Array{T,N}, sn_ratio = 0.6, nonnegative::Bool = false)
    tnsr_noise = randn(size(tnsr)...)
    if nonnegative
        map!(x -> max(0.0, x), tnsr_noise, tnsr_noise)
    end
    tnsr + 10^(-sn_ratio/0.2)*vecnorm(tnsr)/vecnorm(tnsr)*tnsr_noise
end