text:
export iscompatible

## 1. OrthoCube and Operation

"""
    iscompatible(orthocube, operator::TranslationOperation{<:Integer})
"""
function iscompatible(orthocube::OrthoCube, op::TranslationOperation{<:Integer})::Bool
    if dimension(orthocube) != dimension(op)
        throw(DimensionMismatch("hypercube and operation have different dimension"))
    end
    return true
end

"""
    iscompatible(orthocube, operator::PointOperation{<:Integer})
"""
function iscompatible(orthocube::OrthoCube, op::PointOperation{<:Integer})::Bool
    R, r = orthocube.wrap(op.matrix * orthocube.shape_matrix)
    return iszero(r) && abs(ExactLinearAlgebra.determinant(R)) == 1
end

## 2. OrthoCube and Symmetry

"""
    iscompatible(orthocube, symmetry::TranslationSymmetry)
"""
function iscompatible(orthocube::OrthoCube, tsym::TranslationSymmetry)::Bool
    R, r = orthocube.wrap(tsym.orthocube.shape_matrix)
    return iszero(r) && abs(ExactLinearAlgebra.determinant(R)) == 1
end

"""
    iscompatible(orthocube, symmetry::PointSymmetry)
"""
function iscompatible(orthocube::OrthoCube, psym::PointSymmetry)::Bool
    return all(iscompatible(orthocube, pop) for pop in elements(psym))
end

## 3. TranslationSymmetry and PointOperation/PointSymmetry

"""
    iscompatible(tsym::TranslationSymmetry, pop::PointOperation{<:Integer})
"""
function iscompatible(tsym::TranslationSymmetry, pop::PointOperation{<:Integer})::Bool
    return iscompatible(tsym.orthocube, pop)
end

"""
    iscompatible(tsym::TranslationSymmetry, psym::PointSymmetry)
"""
function iscompatible(tsym::TranslationSymmetry, psym::PointSymmetry)::Bool
    return iscompatible(tsym.orthocube, psym)
end

"""
    iscompatible(psym::PointSymmetry, tsym::TranslationSymmetry)
"""
function iscompatible(psym::PointSymmetry, tsym::TranslationSymmetry)::Bool
    return iscompatible(tsym, psym)
end

## 4. Translation Irreps and PointOperation/PointSymmetry

"""
    iscompatible(tsym, tsym_irrep_index, point_operation)
"""
function iscompatible(
    tsym::TranslationSymmetry,
    tsym_irrep_index::Integer,
    pop::PointOperation{<:Integer},
)::Bool
    !iscompatible(tsym, pop) && return false
    reciprocal_matrix = ExactLinearAlgebra.inverse(transpose(pop.matrix))
    k1 = tsym.fractional_momenta[tsym_irrep_index]
    k2 = (x -> mod(x, 1)).(reciprocal_matrix * k1)
    return k1 == k2
end

"""
    iscompatible(tsym, tsym_irrep_index, psym)
"""
function iscompatible(
    tsym::TranslationSymmetry,
    tsym_irrep_index::Integer,
    psym::PointSymmetry,
)::Bool
    return all(
        iscompatible(tsym, tsym_irrep_index, pop)
        for pop in generator_elements(psym)
    )
end

## 5. Lattice and Operation

"""
    iscompatible(lattice, translation_operation)
"""
function iscompatible(lattice::Lattice, op::TranslationOperation{<:Integer})
    if dimension(lattice) != dimension(op)
        throw(DimensionMismatch("lattice and op have different dimension"))
    end
    return true
end

"""
    iscompatible(lattice, point_operation)
"""
function iscompatible(lattice::Lattice, op::PointOperation{<:Integer})
    return iscompatible(lattice.orthocube, op) &&
           !isnothing(findsitemap(lattice.unitcell, op))
end

## 6. Lattice and Symmetry

"""
    iscompatible(lattice, translation_symmetry)

Test whether lattice and the symmetry are compatible.
For translation symmetry, this means that the hypercubic lattice for both are the same.
"""
function iscompatible(lattice::Lattice, tsym::TranslationSymmetry)
    return iscompatible(lattice.orthocube, tsym)
end

"""
    iscompatible(lattice, point_symmetry)
"""
function iscompatible(lattice::Lattice, psym::PointSymmetry)::Bool
    return (
        iscompatible(lattice.orthocube, psym) &&
        all(pop -> !isnothing(findsitemap(lattice.unitcell, pop)), generator_elements(psym))
    )
end


text:
const MaybeSymbol  = Union{Nothing, Symbol}
const MaybeString  = Union{Nothing, String}
const MaybeNumber  = Union{Nothing, Int, Float64}
const MaybeInteger = Union{Nothing, Int32, Int64}
const MaybeFloat   = Union{Nothing, Float32, Float64}
const MaybeURI     = Union{Nothing, URI}
const MaybeMIME    = Union{Nothing, MIME}
const MaybeTuple   = Union{Nothing, Tuple}
const MaybeNTuple{N,T} = Union{Nothing, NTuple{N,T}}
const MaybeNamedTuple = Union{Nothing, NamedTuple}

isinteger(x::MaybeInteger) = !isnothing(x)
isfloat(x::MaybeFloat)     = !isnothing(x)
issymbol(x::MaybeSymbol)   = !isnothing(x)
isstring(x::MaybeString)   = !isnothing(x)
isnumber(x::MaybeNumber)   = !isnothing(x)
isuri(x::MaybeURI)         = !isnothing(x)
ismime(x::MaybeMIME)       = !isnothing(x)
istuple(x::MaybeTuple)     = !isnothing(x)
isntuple(x::MaybeNTuple{N,T}) where {N,T} = !isnothing(x)
isnamedtuple(x::MaybeNamedTuple) = !isnothing(x)

isinteger(x::Int32)  = true
isinteger(x::Int64)  = true
isfloat(x::Float32)  = true
isfloat(x::Float64)  = true
issymbol(x::Symbol)  = true
isstring(x::String)  = true
isnumber(x::Int)     = true
isnumber(x::Float64) = true
isuri(x::URI)        = true
ismime(x::MIME)      = true
istuple(x::Tuple)    = true
isntuple(x::NTuple{N,T}) where {N,T} = true
isnamedtuple(x::NamedTuple) = true

isinteger(x) = false
isfloat(x)   = false
issymbol(x)  = false
isstring(x)  = false
isnumber(x)  = false
isuri(x)     = false
ismime(x)    = false
istuple(x)   = false
isntuple(x)  = false
isnamedtuple(x) = false

ext: jl


text:
# Constructors to create empty models

function SimpleLTVModel(d;extend=false)
    T = length(d)
    n = nstates(d)
    m = ninputs(d)
    N = n*(n+m)
    T -= 1
    model = SimpleLTVModel(zeros(n,n,T),zeros(n,m,T),extend)
end

function KalmanModel(d::AbstractIdData; extend=false, kwargs...)
    T = length(d)
    n = nstates(d)
    m = ninputs(d)
    extend || (T -= 1)
    @assert T > n "The calling convention for x and u is time in the second dimention (n,T = size(x))"
    N     = n*(n+m)
    model = KalmanModel(zeros(n,n,T),zeros(n,m,T),zeros(N,N,T),extend)
end
function LTVAutoRegressive(d, na; extend=false)
    T = length(d)
    extend || (T -= 1)
    model = LTVAutoRegressive(na,zeros(na,T),zeros(na,na,T),extend,0.0)
end

# Constructors that also fit model to data

# TODO: these should be thought through

function KalmanModel(d::AbstractIdData,R1::AbstractMatrix,R2,args...; extend=false, D=1)::KalmanModel
    model = KalmanModel(d,extend=extend)
    KalmanModel(model,d,R1,R2,args...; extend=extend, D=D)
end

function LTVAutoRegressive(d::AnyInput,R1::AbstractMatrix, args...; extend=false, kwargs...)
    na = size(R1,1)
    model = LTVAutoRegressive(d, na, extend=extend)
    LTVAutoRegressive(model,d,R1, args...; extend=extend, kwargs...)
end

# function SimpleLTVModel(fitmethod::Symbol, args...; kwargs...)
#     fun = if fitmethod == :gd
#         fit_statespace_gd
#     elseif fitmethod == :scs
#         fit_statespace
#     elseif fitmethod == :constrained
#         fit_statespace_constrained
#     elseif fitmethod == :dp
#         fit_statespace_dp
#     end
#     fun(args...; kwargs...)
# end


text:
export JackknifeVector, JackknifeVectorSet

type JackknifeVector <: VectorObservable
  xs :: Vector{Vector{Float64}}
end

JackknifeVector() = JackknifeVector(Vector{Float64}[])

function JackknifeVector(jk::JackknifeVector, f::Function)
  xs = similar(jk.xs)
  for i in 1:length(jk.xs)
    xs[i] = f(jk.xs[i])
  end
  JackknifeVector(xs)
end

function JackknifeVector(o::VectorObservable)
  if isempty(o)
    return JackknifeVector()
  else
    return JackknifeVector(jk_helper(o.bins))
  end
end

count(jk::JackknifeVector) = length(jk.xs)

function mean(jk::JackknifeVector)
  if isempty(jk)
    return NaN
  else
    return mean(jk.xs)
  end
end
function stderror(jk::JackknifeVector)
  n = count(jk)
  if n == 0
    return NaN
  elseif n == 1
    return Inf
  else
    m2 = sumabs2(jk.xs)
    m2 /= n
    m = mean(jk)
    sigma2 = m2 - m.*m
    sigma2 *= n-1
    map!(maxzero, sigma2)
    return sqrt(sigma2)
  end
end

unary_functions = (
  :-,
  :sin, :cos, :tan,
  :sind, :cosd, :tand,
  :sinpi, :cospi,
  :sinh, :cosh, :tanh,
  :asin, :acos, :atan,
  :asind, :acosd, :atand,
  :sec, :csc, :cot,
  :secd, :cscd, :cotd,
  :asec, :acsc, :acot,
  :asecd, :acscd, :acotd,
  :sech, :csch, :coth,
  :asinh, :acosh, :atanh,
  :asech, :acsch, :acoth,
  :sinc, :cosc,
  :log, :log2, :log10, :log1p,
  :exp, :exp2, :exp10, :expm1,
  :abs, :abs2,
  :sqrt, :cbrt,
  :erf, :erfc, :erfcx,
  :erfinv, :erfcinv,
  :gamma, :lgamma, :lfact,
  :digamma, :invdigamma, :trigamma,
  :airyai, :airyprime, :airyaiprime,
  :airybi, :airybiprime,
  :besselj0, :besselj1,
  :bessely0, :bessely1,
  :eta, :zeta
)

for op in unary_functions
  eval( Expr(:import, :Base, op) )
  eval( Expr(:export, op) )
  @eval ($op)(jk::JackknifeVector) = JackknifeVector(jk, $op)
end

binary_functions = (
  :+, :-, :*, :/, :\
)

for op in ( :+, :-, :.+, :.- )
  eval( Expr(:import, :Base, op) )
  eval( Expr(:export, op) )
  @eval ($op)(jk::JackknifeVector, rhs::Real) = JackknifeVector(jk, lhs->($op)(lhs,rhs))
  @eval ($op)(jk::JackknifeVector, rhs::Vector) = JackknifeVector(jk, lhs->($op)(lhs,rhs))
  @eval ($op)(lhs::Real, jk::JackknifeVector) = JackknifeVector(jk, rhs->($op)(lhs,rhs))
  @eval ($op)(lhs::Vector, jk::JackknifeVector) = JackknifeVector(jk, rhs->($op)(lhs,rhs))
  @eval ($op)(lhs::JackknifeVector, rhs::JackknifeVector) = JackknifeVector( ($op)(lhs.xs, rhs.xs))
end
for op in ( :*, :/, :\)
  eval( Expr(:import, :Base, op) )
  eval( Expr(:export, op) )
  @eval ($op)(lhs::Real, jk::JackknifeVector) = JackknifeVector(jk, rhs->($op)(lhs,rhs))
  @eval ($op)(jk::JackknifeVector, rhs::Real) = JackknifeVector(jk, lhs->($op)(lhs,rhs))
end
for op in ( :.*, :./, :.\, :.^)
  eval( Expr(:import, :Base, op) )
  eval( Expr(:export, op) )
  @eval ($op)(lhs::Real, jk::JackknifeVector) = JackknifeVector(jk, rhs->($op)(lhs,rhs))
  @eval ($op)(lhs::Vector, jk::JackknifeVector) = JackknifeVector(jk, rhs->($op)(lhs,rhs))
  @eval ($op)(jk::JackknifeVector, rhs::Real) = JackknifeVector(jk, lhs->($op)(lhs,rhs))
  @eval ($op)(jk::JackknifeVector, rhs::Vector) = JackknifeVector(jk, lhs->($op)(lhs,rhs))
  @eval ($op)(lhs::JackknifeVector, rhs::JackknifeVector) = begin
    xs = similar(lhs.xs)
    for i in 1:length(lhs.xs)
      xs[i] = ($op)(lhs.xs[i], rhs.xs[i])
    end
    return JackknifeVector(xs)
  end
end

typealias JackknifeSet MCObservableSet{Jackknife}

jackknife(obs::VectorObservable) = JackknifeVector(obs)
function jackknife{Obs<:VectorObservable}(obsset :: MCObservableSet{Obs})
  JK = JackknifeSet()
  for (k,v) in obsset
    JK[k] = JackknifeVector(v)
  end
  return JK
end


text:
# XXX tfunc implementations in Core.Compiler are really not enough to catch invalid calls
# TODO set up our own checks and enable sound analysis

function CC.builtin_tfunction(analyzer::AbstractAnalyzer, @nospecialize(f), argtypes::Array{Any,1},
                              sv::InferenceState) # `AbstractAnalyzer` isn't overloaded on `return_type`
    ret = @invoke builtin_tfunction(analyzer::AbstractInterpreter, f, argtypes::Array{Any,1},
                                    sv::Union{InferenceState,Nothing})

    if f === throw
        # here we only report a selection of "serious" exceptions, i.e. those that should be
        # reported even if they may be caught in actual execution;
        ReportPass(analyzer)(SeriousExceptionReport, analyzer, sv, argtypes)

        # other general `throw` calls will be reported within `_typeinf(analyzer::AbstractAnalyzer, frame::InferenceState)`
        # only when they are not caught by control flow, which is judged by whether if the
        # final return type of `sv` is annotated as `Bottom` or not, thus early return now
        return ret
    end

    ReportPass(analyzer)(InvalidBuiltinCallErrorReport, analyzer, sv, f, argtypes, ret)

    if f === getfield && 2 ≤ length(argtypes) ≤ 3
        obj, fld = argtypes
        if isa(fld, Const)
            name = fld.val
            if isa(name, Symbol)
                if isa(obj, Const)
                    mod = obj.val
                    if isa(mod, Module)
                        if isdefined(mod, name)
                            if istoplevel_globalref(analyzer, sv)
                                # when accessing to a global variable in a module
                                # concretized by `analyzer`, eagerly propagate its type
                                # NOTE logic here should be synced with that of `abstract_eval_special_value(::AbstractAnalyzer, ::Any, ::VarTable, ::InferenceState)`
                                val = getfield(mod, name)
                                return isa(val, AbstractGlobal) ? val.t : Const(val)
                            end
                        end
                    end
                end
            end
        end
    elseif f === fieldtype
        # the valid widest possible return type of `fieldtype_tfunc` is `Union{Type,TypeVar}`
        # because fields of unwrapped `DataType`s can legally be `TypeVar`s,
        # but this will cause lots of false positive `NoMethodErrorReport`s for inference
        # with accessing to abstract fields since most methods don't expect `TypeVar`
        # (e.g. `@report_call readuntil(stdin, 'c')`)
        # JET.jl further widens this case to `Any` and give up further analysis rather than
        # trying hard to do sound and noisy analysis
        # xref: https://github.com/JuliaLang/julia/pull/38148
        if ret === Union{Type, TypeVar}
            return Any
        end
    end

    return ret
end

"""
    SeriousExceptionReport <: InferenceErrorReport

Represents a "serious" error that is manually thrown by a `throw` call.
This is reported regardless of whether it's caught by control flow or not, as opposed to
[`UncaughtExceptionReport`](@ref).
"""
struct SeriousExceptionReport{T} <: InferenceErrorReport
    vst::VirtualStackTrace
    msg::String
    sig::Vector{Any}
    err::T
    SeriousExceptionReport(vst::VirtualStackTrace, msg::String, sig::Vector{Any}, err::T) where T =
        return new{T}(vst, msg, sig, err)
    SeriousExceptionReport{T}(vst::VirtualStackTrace, msg::String, sig::Vector{Any}, (err,)::Tuple{T}) where T =
        return new{T}(vst, msg, sig, err)
end
get_spec_args(report::SeriousExceptionReport) = (report.err,)

function SeriousExceptionReport(analyzer::AbstractAnalyzer, state::InferenceState, err)
    vf = get_virtual_frame(analyzer, state)
    msg = string(first(split(sprint(showerror, err), '\n')))
    ret = SeriousExceptionReport([vf], msg, vf.sig, err)
    push!(get_throw_locs(analyzer), get_lin((state, get_currpc(state))))
    return ret
end

function (::SoundBasicPass)(::Type{SeriousExceptionReport}, analyzer::AbstractAnalyzer, sv::InferenceState, argtypes::Vector{Any})
    if length(argtypes) ≥ 1
        a = first(argtypes)
        if isa(a, Const)
            err = a.val
            if isa(err, UndefKeywordError)
                add_new_report!(SeriousExceptionReport(analyzer, sv, err), analyzer)
            elseif isa(err, MethodError)
                # ignore https://github.com/JuliaLang/julia/blob/7409a1c007b7773544223f0e0a2d8aaee4a45172/base/boot.jl#L261
                if err.f !== Bottom
                    add_new_report!(SeriousExceptionReport(analyzer, sv, err), analyzer)
                end
            end
        end
    end
end

"""
    InvalidBuiltinCallErrorReport

Represents errors caused by invalid builtin-function calls.
Technically they're defined as those error points that should be caught within
`Core.Compiler.builtin_tfunction`.
"""
abstract type InvalidBuiltinCallErrorReport <: InferenceErrorReport end

@reportdef struct NoFieldErrorReport <: InvalidBuiltinCallErrorReport
    @nospecialize(typ::Type)
    name::Symbol
end
get_msg(::Type{NoFieldErrorReport}, analyzer::AbstractAnalyzer, sv::InferenceState, @nospecialize(typ::Type), name::Symbol) =
    "type $(typ) has no field $(name)"

@reportdef struct DivideErrorReport <: InferenceErrorReport end
let s = sprint(showerror, DivideError())
    global get_msg(::Type{DivideErrorReport}, analyzer::AbstractAnalyzer, sv::InferenceState) = s
end

@reportdef struct UnimplementedBuiltinCallErrorReport <: InvalidBuiltinCallErrorReport
    argtypes::Vector{Any}
end
get_msg(::Type{UnimplementedBuiltinCallErrorReport}, analyzer::AbstractAnalyzer, sv::InferenceState, @nospecialize(args...)) =
    "invalid builtin function call"

# TODO we do need sound versions of these functions

# XXX: for general case, JET just relies on the (maybe too persmissive) return type
# from native tfuncs to report invalid builtin calls and probably there're lots of
# false negatives
function handle_unimplemented_builtins!(analyzer::AbstractAnalyzer, sv::InferenceState, @nospecialize(f), argtypes::Vector{Any}, @nospecialize(ret))
    if ret === Bottom
        @assert !(f === throw) "`throw` calls shuold be handled by report passes of `SeriousExceptionReport` or `UncaughtExceptionReport`"

        add_new_report!(UnimplementedBuiltinCallErrorReport(analyzer, sv, argtypes), analyzer)
    end
end

# fallback for unimplemented tfuncs, eliminate me in the future
function (::SoundBasicPass)(::Type{InvalidBuiltinCallErrorReport}, analyzer::AbstractAnalyzer, sv::InferenceState, @nospecialize(f::Builtin), argtypes::Vector{Any}, @nospecialize(ret))
    handle_unimplemented_builtins!(analyzer, sv, f, argtypes, ret)
end

function (::SoundBasicPass)(::Type{InvalidBuiltinCallErrorReport}, analyzer::AbstractAnalyzer, sv::InferenceState, f::typeof(getfield), argtypes::Vector{Any}, @nospecialize(ret))
    if 2 ≤ length(argtypes) ≤ 3
        obj, fld = argtypes
        if isa(fld, Const)
            name = fld.val
            if isa(name, Symbol)
                if isa(obj, Const) && (mod = obj.val; isa(mod, Module))
                    if !isdefined(mod, name)
                        # bypass for report pass for undefined global reference
                        ReportPass(analyzer)(GlobalUndefVarErrorReport, analyzer, sv, mod, name)
                        return
                    end
                elseif ret === Bottom
                    # general case when an error is detected by the native `getfield_tfunc`
                    typ = widenconst(obj)
                    add_new_report!(NoFieldErrorReport(analyzer, sv, typ, name), analyzer)
                    return
                end
            end
        end
    end

    handle_unimplemented_builtins!(analyzer, sv, f, argtypes, ret)
end

function (::SoundBasicPass)(::Type{InvalidBuiltinCallErrorReport}, analyzer::AbstractAnalyzer, sv::InferenceState, f::IntrinsicFunction, argtypes::Vector{Any}, @nospecialize(ret))
    # TODO this check might be better in its own report pass, say `NumericalPass`
    if length(argtypes) == 2
        if f === Intrinsics.checked_sdiv_int ||
           f === Intrinsics.checked_srem_int ||
           f === Intrinsics.checked_udiv_int ||
           f === Intrinsics.checked_urem_int ||
           f === Intrinsics.sdiv_int ||
           f === Intrinsics.srem_int ||
           f === Intrinsics.udiv_int ||
           f === Intrinsics.urem_int
            a = argtypes[2]
            t = widenconst(a)
            if isprimitivetype(t) && t <: Number
                if isa(a, Const) && a.val === zero(t)
                    add_new_report!(DivideErrorReport(analyzer, sv), analyzer)
                    return
                end
            end
        end
    end

    handle_unimplemented_builtins!(analyzer, sv, f, argtypes, ret)
end

# check if this frame is for `getproperty(::Module, ::Symbol)`, which accesses to a global
# variable traced by `analyzer`
function istoplevel_globalref(analyzer::AbstractAnalyzer, sv::InferenceState)
    def = sv.linfo.def
    def.name === :getproperty || return false
    def.sig === Tuple{typeof(getproperty), Module, Symbol} || return false
    parent = sv.parent
    return !isnothing(parent) && istoplevel(analyzer, parent)
end

# `return_type_tfunc` internally uses `abstract_call` to model `Core.Compiler.return_type`
# and here we shouldn't pass `AbstractAnalyzer` to it; otherwise we may get false error reports
# from the `abstract_call`, which will simulate the call and isn't any abstraction of actual
# execution of it
function CC.return_type_tfunc(analyzer::AbstractAnalyzer, argtypes::Vector{Any}, sv::InferenceState)
    # report pass for invalid `Core.Compiler.return_type` call
    ReportPass(analyzer)(InvalidReturnTypeCall, analyzer, sv, argtypes)

    # don't recursively pass on `AbstractAnalyzer` via `@invoke` here, and make sure
    # JET's analysis enter into the simulated call
    return return_type_tfunc(get_native(analyzer), argtypes, sv)
end

@reportdef struct InvalidReturnTypeCall <: InferenceErrorReport end
get_msg(::Type{InvalidReturnTypeCall}, analyzer::AbstractAnalyzer, sv::InferenceState) = "invalid `Core.Compiler.return_type` call"

function (::SoundBasicPass)(::Type{InvalidReturnTypeCall}, analyzer::AbstractAnalyzer, sv::InferenceState, argtypes::Vector{Any})
    # here we just check if the call of `return_type` is valid or not by very simple analysis
    # we don't take (possible, but very unexpected) overloads into account here, just as
    # `NativeInterpreter`'s `return_type_tfunc` hard-codes its return type to `Type`
    if length(argtypes) ≠ 3
        # invalid argument number, let's report and return error result (i.e. `Bottom`)
        add_new_report!(InvalidReturnTypeCall(analyzer, sv), analyzer)
    end
end


text:
"""
$(SIGNATURES)
Computes correlation matrix (R) * n. n is not removed because of performance choice.

# Arguments
- `a` : standardized phenotype matrix
- `b` : standardized genotype matrix

# Output:
returns correlation matrix R

"""
function calculate_nr(a::AbstractArray{<:Real,2},b::AbstractArray{<:Real, 2})
    return LinearAlgebra.BLAS.gemm('T', 'N', a,b);
end

"""
$(SIGNATURES)
Checks whether correlation matrix is in range. Only runs if DEBUG flag is turned on from julia commandline.

# Arguments
- `r` : a matrix to be standardized
- min : minimum value of range
- max : maximum value of range

# Output:
errors message will show if correlation matrix is not in range.

"""
function is_corr_in_range(r, min, max)
    function inRange(n, min, max)
        if n >= min && n <= max
            return true
        else
            return false
        end
    end

    if all( inRange.(r, min,max) )
        error("Correlation matrix is not in range($min, $max). Check your r matrix again. ")
    end

end

"""
$(SIGNATURES)
Computes log of odds (LOD) score. Optimized for correlation matrix type is Float64 (double precision).

!Notes: Set the thread number with env JULIA_NUM_THREADS to your desired number of threads.
For example: `JULIA_NUM_THREADS=16 julia`

# Arguments
- `m` : number of individuals.
- `nr` : correlation matrix R times n. N will be removed during this step.

# Output:
returns LOD score.

"""
function lod_score_multithread(m,nr::AbstractArray{Float64, 2})
    n = m
    Threads.@threads for j in 1:size(nr)[2]
        for i in 1:size(nr)[1]
            r_square = (nr[i,j] / n)^2
            tmp = (-n/2.0) * log10(1.0-r_square)
            nr[i,j] = tmp
        end
    end
    return nr
end

"""
$(SIGNATURES)
Computes log of odds (LOD) score. Optimized for correlation matrix type is Float32 (single precision).

!Notes: Set the thread number with env JULIA_NUM_THREADS to your desired number of threads.
For example: `JULIA_NUM_THREADS=16 julia`

# Arguments
- `m` : number of individuals.
- `nr` : correlation matrix R times n. N will be removed during this step.

# Output:
returns LOD score.

"""
function lod_score_multithread(m,nr::AbstractArray{Float32,2})
    n = m
    Threads.@threads for j in 1:size(nr)[2]
        for i in 1:size(nr)[1]
            r_square = (nr[i,j] / n)^2
            tmp = (-n/2.0f0) * log10(1.0f0-r_square)
            nr[i,j] = tmp
        end
    end
    return nr
end

"""
$(SIGNATURES)
Computes p-value based on log of odds score.

# Arguments
- `lod` : LOD matrix

# Output:
returns p-value.

"""
function lod2p(lod)
    return 1 .- cdf(Chisq(1),2*log(10)*lod)
end

"""
$(SIGNATURES)
Computes the index of maximum, and maximum value of each row of a matrix.
Optimized with multi-threading.

!Notes: Set the thread number with env JULIA_NUM_THREADS to your desired number of threads.
For example: `JULIA_NUM_THREADS=16 julia`

# Arguments
- `lod` : input matrix.

# Output:
returns a matrix with two columns, first column is the index of maximum, second column is
the maximum value.

"""
function find_max_idx_value(lod::AbstractArray{<:Real,2})
    max_array = Array{typeof(lod[1,1]),2}(undef, size(lod)[1], 2)
    Threads.@threads for i in 1:size(lod)[1]
        # for i in 1:size(lod)[1]
        temp = lod[i, 1]
        idx = 1
        for j in 2:size(lod)[2]
            if temp < lod[i,j]
                temp = lod[i,j]
                idx = j
            end
        end
        max_array[i,1] = idx
        max_array[i,2] = temp
    end
    return max_array
end

##################### Running CPU Function ###################
"""
$(SIGNATURES)

# Arguments:
- `Y` : a matrix of phenotypes.
- `G` : a matrix of genotypes.
- `X` : a matrix of covariates. Default is `nothing`. If `nothing`, scan is run without covariates.
- `maf_threshold`: a floating point number to indicate the maf_threshold. Default is 0.05. Set to 0 if no maf filtering should be done.
- `export_matrix` : a boolean value that determines whether the result should be the maximum value of LOD score of each phenotype and its corresponding index, or the whole LOD score matrix.
- `lod_or_pval`: a string value of either `lod` or `pval` to indicate the desired output.
- `timing_file`: a string that indicates the file location for the timing outputs. Default is nothing.
# Output:
returns LOD score or pval, in vector or matrix format depending on value of `export_matrix`.
"""
function cpurun(pheno::AbstractArray{<:Real,2}, geno::AbstractArray{<:Real,2}, X::Union{AbstractArray{<:Real, 2}, Nothing}=nothing; maf_threshold=0.05, export_matrix=false, lod_or_pval="lod",timing_file="")
    @debug begin
        "size(Y) = $(size(Y)), size(G) = $(size(G)). Number of indvidual should be size(Y, 1), or size(G, 1). "
    end
    pval_time = 0.0
    compute_time = 0.0
    result_reorg_time = 0.0
    data_transfer_time = 0.0

    total_start = time_ns()
    data_transfer_time += @elapsed G = geno
    data_transfer_time += @elapsed Y = pheno

    compute_time += @elapsed begin
        if maf_threshold > 0
            println("Filtering MAF")
            G = filter_maf(geno, maf_threshold=maf_threshold)
        end

        n = size(G,1)
        if !isnothing(X) # X is not empty.
            px = calculate_px(X)
            y_hat = LinearAlgebra.BLAS.gemm('N', 'N', px, Y)
            g_hat = LinearAlgebra.BLAS.gemm('N', 'N', px, G)
            Y = Y .- y_hat
            G = G .- g_hat
        end
        y_std = get_standardized_matrix(Y)
        g_std = get_standardized_matrix(G)

        nr = calculate_nr(y_std,g_std);
    end

    @debug begin
        test_r_in_range = is_corr_in_range(nr./n, -1,1)
        "R is in range (-1, 1): $test_r_in_range"
    end

    if lod_or_pval == "lod"
        compute_time += @elapsed lod = lod_score_multithread(n,nr)

        if !export_matrix
            compute_time += @elapsed lod = find_max_idx_value(lod)
        end
    elseif lod_or_pval == "pval"
        pval_time += @elapsed pval = pval_calc(nr ./ n, n-2)
    else
        error("Must specify `lod_or_pval`, choose between `lod`, or `pval`")
    end

    total_stop = time_ns()
    elapsed_total = (total_stop - total_start) * 1e-9

    if timing_file != ""
        open("/home/xiaoqihu/git/LiteQTL-G3-supplement/code/tensorqtl/liteqtl_timing_report.txt", "a") do io
            write(io, "$(now()),CPU,$(data_transfer_time),$(compute_time),$(result_reorg_time),$(pval_time),$(elapsed_total)\n")
        end
    end

    if lod_or_pval == "lod"
        return lod
    elseif lod_or_pval == "pval"
        return pval
    else
        error("Must specify `lod_or_pval`, choose between `lod`, or `pval`")
    end

end


text:
using CUSPARSE
using CUDArt
using Base.Test

m = 25
n = 35
k = 10
blockdim = 5

@testset "bsric02" begin
    @testset for elty in [Float32,Float64,Complex64,Complex128]
        A = rand(elty, m, m)
        A += ctranspose(A)
        A += m * eye(elty, m)

        @testset "bsric02!" begin
            d_A = CudaSparseMatrixCSR(sparse(tril(A)))
            d_A = CUSPARSE.switch2bsr(d_A, convert(Cint,5))
            d_A = CUSPARSE.ic02!(d_A,'O')
            h_A = to_host(CUSPARSE.switch2csr(d_A))
            Ac = sparse(full(cholfact(Hermitian(A))))
            h_A = transpose(h_A) * h_A
            @test h_A.rowval ≈ Ac.rowval
            @test reduce(&, isfinite.(h_A.nzval))
            d_A = CudaSparseMatrixCSR(sparse(tril(rand(elty,m,n))))
            d_A = CUSPARSE.switch2bsr(d_A, convert(Cint,5))
            @test_throws DimensionMismatch CUSPARSE.ic02!(d_A,'O')
        end

        @testset "bsric02" begin
            d_A = CudaSparseMatrixCSR(sparse(tril(A)))
            d_A = CUSPARSE.switch2bsr(d_A, convert(Cint,5))
            d_B = CUSPARSE.ic02(d_A,'O')
            h_A = to_host(CUSPARSE.switch2csr(d_B))
            Ac = sparse(full(cholfact(Hermitian(A))))
            h_A = transpose(h_A) * h_A
            @test h_A.rowval ≈ Ac.rowval
            @test reduce(&, isfinite.(h_A.nzval))
        end
    end
end


text:
# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    Coerce(pairs, tight=false, verbosity=1)

Return a copy of the table, ensuring that the scientific types of the columns match the new specification.

This transform wraps the ScientificTypes.coerce function. Please see their docstring for more details.

# Examples

```julia
Coerce(:col1 => Continuous, :col2 => Count)
```
"""
struct Coerce{P} <: Transform
  pairs::P
  tight::Bool
  verbosity::Int
end

Coerce(pair::Pair{Symbol,<:Type}...; tight=false, verbosity=1) =
  Coerce(pair, tight, verbosity)

isrevertible(::Type{<:Coerce}) = true

function apply(transform::Coerce, table)
  newtable = coerce(table, transform.pairs...;
                    tight=transform.tight,
                    verbosity=transform.verbosity)

  types = Tables.schema(table).types

  newtable, types
end

function revert(transform::Coerce, newtable, cache)
  names = Tables.columnnames(newtable)
  cols = Tables.columns(newtable)
  oldcols = map(zip(cache, names)) do (T, n)
    x = Tables.getcolumn(cols, n)
    collect(T, x)
  end

  𝒯 = (; zip(names, oldcols)...)
  𝒯 |> Tables.materializer(newtable)
end


text:
InputFolder = './Images/PointClusters/';
OutputFolder = './Results/Images/PointClusters/';
Rescale = 1;
Dilate = 5;
Fill = 1;

@iA = '*.tif';

@fxs_lMeanShift3D [iA] > [S L];
params.Bandwidth = 75;
params.MinPts = 0;
/endf

/show iA > S;
/keep L > tif;

max_stars_repo_name: SebastienTs/LOBSTER

max_issues_repo_name: SebastienTs/LOBSTER

max_forks_repo_name: SebastienTs/LOBSTER


text:
using PowerSystems
using NLsolve
const PSY = PowerSystems

############### Data Network ########################
include(joinpath(dirname(@__FILE__), "dynamic_test_data.jl"))
include(joinpath(dirname(@__FILE__), "data_utils.jl"))
############### Data Network ########################
threebus_file_dir = joinpath(dirname(@__FILE__), "ThreeBusNetwork.raw")
threebus_sys = System(
    PowerModelsData(threebus_file_dir),
    runchecks = false,
    unit_system = "device_base",
)
add_source_to_ref(threebus_sys)

function dyn_gen_avr_type2(generator)
    return PSY.DynamicGenerator(
        generator, #static generator
        1.0, # ω_ref
        machine_oneDoneQ(), #machine
        shaft_no_damping(), #shaft
        avr_type2(), #avr
        tg_type1(), #tg
        pss_none(),
    ) #pss
end

function dyn_gen_simple_avr(generator)
    return PSY.DynamicGenerator(
        generator, #static generator
        1.0, # ω_ref
        machine_oneDoneQ(), #machine
        shaft_no_damping(), #shaft
        avr_propr(), #avr
        tg_none(), #tg
        pss_none(),
    ) #pss
end

# Add dynamic generators to the system (each gen is linked through a static one)
for g in get_components(Generator, threebus_sys)
    if get_number(get_bus(g)) == 102
        case_gen = dyn_gen_avr_type2(g)
        add_component!(threebus_sys, case_gen)
    elseif get_number(get_bus(g)) == 103
        case_gen = dyn_gen_simple_avr(g)
        add_component!(threebus_sys, case_gen)
    end
end

#Compute Y_bus after fault
fault_branches = deepcopy(collect(get_components(Branch, threebus_sys))[2:end])
Ybus_fault = PSY.Ybus(fault_branches, get_components(Bus, threebus_sys))[:, :]


text:
# This file is a part of Julia. License is MIT: https://julialang.org/license

module Partr

using ..Threads: SpinLock

# a task heap
mutable struct taskheap
    const lock::SpinLock
    const tasks::Vector{Task}
    @atomic ntasks::Int32
    @atomic priority::UInt16
    taskheap() = new(SpinLock(), Vector{Task}(undef, tasks_per_heap), zero(Int32), typemax(UInt16))
end

# multiqueue parameters
const heap_d = UInt32(8)
const heap_c = UInt32(2)

# size of each heap
const tasks_per_heap = Int32(65536) # TODO: this should be smaller by default, but growable!

# the multiqueue's heaps
global heaps::Vector{taskheap}
global heap_p::UInt32 = 0

# unbias state for the RNG
global cong_unbias::UInt32 = 0

cong(max::UInt32, unbias::UInt32) = ccall(:jl_rand_ptls, UInt32, (UInt32, UInt32), max, unbias) + UInt32(1)

function unbias_cong(max::UInt32)
    return typemax(UInt32) - ((typemax(UInt32) % max) + UInt32(1))
end

function multiq_init(nthreads)
    global heap_p = heap_c * nthreads
    global cong_unbias = unbias_cong(UInt32(heap_p))
    global heaps = Vector{taskheap}(undef, heap_p)
    for i = UInt32(1):heap_p
        heaps[i] = taskheap()
    end
    nothing
end

function sift_up(heap::taskheap, idx::Int32)
    while idx > Int32(1)
        parent = (idx - Int32(2)) ÷ heap_d + Int32(1)
        if heap.tasks[idx].priority < heap.tasks[parent].priority
            t = heap.tasks[parent]
            heap.tasks[parent] = heap.tasks[idx]
            heap.tasks[idx] = t
            idx = parent
        else
            break
        end
    end
end

function sift_down(heap::taskheap, idx::Int32)
    if idx <= heap.ntasks
        for child = (heap_d * idx - heap_d + Int32(2)):(heap_d * idx + Int32(1))
            child > tasks_per_heap && break
            if isassigned(heap.tasks, child) &&
                    heap.tasks[child].priority < heap.tasks[idx].priority
                t = heap.tasks[idx]
                heap.tasks[idx] = heap.tasks[child]
                heap.tasks[child] = t
                sift_down(heap, child)
            end
        end
    end
end

function multiq_insert(task::Task, priority::UInt16)
    task.priority = priority

    rn = cong(heap_p, cong_unbias)
    while !trylock(heaps[rn].lock)
        rn = cong(heap_p, cong_unbias)
    end

    if heaps[rn].ntasks >= tasks_per_heap
        unlock(heaps[rn].lock)
        # multiq insertion failed, increase #tasks per heap
        return false
    end

    ntasks = heaps[rn].ntasks + Int32(1)
    @atomic :monotonic heaps[rn].ntasks = ntasks
    heaps[rn].tasks[ntasks] = task
    sift_up(heaps[rn], ntasks)
    priority = heaps[rn].priority
    if task.priority < priority
        @atomic :monotonic heaps[rn].priority = task.priority
    end
    unlock(heaps[rn].lock)
    return true
end

function multiq_deletemin()
    local rn1, rn2
    local prio1, prio2

    @label retry
    GC.safepoint()
    for i = UInt32(1):heap_p
        if i == heap_p
            return nothing
        end
        rn1 = cong(heap_p, cong_unbias)
        rn2 = cong(heap_p, cong_unbias)
        prio1 = heaps[rn1].priority
        prio2 = heaps[rn2].priority
        if prio1 > prio2
            prio1 = prio2
            rn1 = rn2
        elseif prio1 == prio2 && prio1 == typemax(UInt16)
            continue
        end
        if trylock(heaps[rn1].lock)
            if prio1 == heaps[rn1].priority
                break
            end
            unlock(heaps[rn1].lock)
        end
    end

    task = heaps[rn1].tasks[1]
    tid = Threads.threadid()
    if ccall(:jl_set_task_tid, Cint, (Any, Cint), task, tid-1) == 0
        unlock(heaps[rn1].lock)
        @goto retry
    end
    ntasks = heaps[rn1].ntasks
    @atomic :monotonic heaps[rn1].ntasks = ntasks - Int32(1)
    heaps[rn1].tasks[1] = heaps[rn1].tasks[ntasks]
    Base._unsetindex!(heaps[rn1].tasks, Int(ntasks))
    prio1 = typemax(UInt16)
    if ntasks > 1
        sift_down(heaps[rn1], Int32(1))
        prio1 = heaps[rn1].tasks[1].priority
    end
    @atomic :monotonic heaps[rn1].priority = prio1
    unlock(heaps[rn1].lock)

    return task
end

function multiq_check_empty()
    for i = UInt32(1):heap_p
        if heaps[i].ntasks != 0
            return false
        end
    end
    return true
end

end


text:
# ---
# title: Displaying Images
# author: "[William Thompson](https://github.com/sefffal)"
# cover: assets/displaying-images.png
# ---

# We'll start by downloading a sample image. If you have an image stored locally,
# you would skip this step.
using AstroImages

AstroImages.set_clims!(Percent(99.5)) #src
AstroImages.set_cmap!(:magma) #src
AstroImages.set_stretch!(identity) #src

# Any AbstractArray can be visualized with the `imview` function.
arr = randn(10,10)
imview(arr)

# Let's load an astronomical image to see how we can tweak its display
fname = download(
    "http://www.astro.uvic.ca/~wthompson/astroimages/fits/656nmos.fits",
    "eagle-656nmos.fits"
);
img = AstroImage("eagle-656nmos.fits");
imview(img)

# We can adjust the color limits manually
imview(img, clims=(0,100))

# Or provide a function to calculate them for us
imview(img, clims=extrema)

# AstroImages includes some handy callables, like Percent and Zscale.flags
# `Percent` sets the limits to include some central percentage of the data range
# For example, 95% sets the color limits to clip the top and bottom 2.5% of pixels.
# Percent(99.5) is the default value of clims.
imview(img, clims=Percent(95))

# Arrays wrapped by `AstroImage` are displayed automatically using `imview`
AstroImage(randn(10,10))

# The settings for automatic imview are controlled using package defaults that can
# be adjusted to suit your tastes
AstroImages.set_clims!(Zscale()) # Display the full range automatically
AstroImages.set_cmap!(:viridis)
AstroImages.set_stretch!(asinhstretch)
AstroImage(randn(10,10))

# --- restore defaults --- #src
AstroImages.set_clims!(Percent(99.5)) #src
AstroImages.set_cmap!(:magma) #src
AstroImages.set_stretch!(identity) #src

# --- save covers --- #src
mkpath("assets")  #src
save("assets/loading-images.png", imview(img)) #src


text:
#=
A first pass at solving the optimal growth problem via value function
iteration.  A more general version is provided in optgrowth.py.

@author : Spencer Lyon <spencer.lyon@nyu.edu>
=#

using Optim: optimize
using Grid: CoordInterpGrid, BCnan, InterpLinear
using PyPlot

## Primitives and grid
alpha = 0.65
bet = 0.95
grid_max = 2
grid_size = 150
grid = 1e-6:(grid_max-1e-6)/(grid_size-1):grid_max

## Exact solution
ab = alpha * bet
c1 = (log(1 - ab) + log(ab) * ab / (1 - ab)) / (1 - bet)
c2 = alpha / (1 - ab)
v_star(k) = c1 .+ c2 .* log(k)

function bellman_operator(grid, w)
    Aw = CoordInterpGrid(grid, w, BCnan, InterpLinear)

    Tw = zeros(w)

    for (i, k) in enumerate(grid)
        objective(c) = - log(c) - bet * Aw[k^alpha - c]
        res = optimize(objective, 1e-6, k^alpha)
        Tw[i] = - objective(res.minimum)
    end
    return Tw
end

function main(n::Int=35)
    w = 5 .* log(grid) .- 25  # An initial condition -- fairly arbitrary
    fig, ax = subplots()
    ax[:set_ylim](-40, -20)
    ax[:set_xlim](minimum(grid), maximum(grid))
    lb = "initial condition"
    jet = ColorMap("jet")[:__call__]
    ax[:plot](grid, w, color=jet(0), lw=2, alpha=0.6, label=lb)

    for i=1:n
        w = bellman_operator(grid, w)
        ax[:plot](grid, w, color=jet(i/n), lw=2, alpha=0.6)
    end
    lb = "true value function"
    ax[:plot](grid, v_star(grid), "k-", lw=2, alpha=0.8, label=lb)
    ax[:legend](loc="upper left")
    Void
end


text:
# "Gaussian advection-diffusion" Spatial resolution convergence test

using PyPlot

# Define a few utilities for running tests and unpacking and plotting results
include("ConvergenceTests/ConvergenceTests.jl")

using .ConvergenceTests
using .ConvergenceTests.OneDimensionalGaussianAdvectionDiffusion: run_test
using .ConvergenceTests.OneDimensionalUtils: plot_solutions!, plot_error_convergence!, unpack_errors

""" Run advection-diffusion test for all Nx in resolutions. """
function run_convergence_test(κ, U, resolutions...)

    # Determine save time-step
           Lx = 2.5
    stop_time = 0.25
            h = Lx / maximum(resolutions)
           Δt = min(0.1 * h / U, 0.01 * h^2 / κ)

    # Adjust time-step
    stop_iteration = round(Int, stop_time / Δt)
                Δt = stop_time / stop_iteration

    # Run the tests
    results = [run_test(Nx=Nx, Δt=Δt, stop_iteration=stop_iteration, U=U, κ=κ, width=0.1) for Nx in resolutions]

    return results
end

#####
##### Run test
#####

Nx = 2 .^ (6:8) # N = 64 through N = 256
diffusion_results = run_convergence_test(1e-1, 0, Nx...)
advection_results = run_convergence_test(1e-6, 3, Nx...)
advection_diffusion_results = run_convergence_test(1e-2, 1, Nx...)

#####
##### Plot solution and error profile
#####

all_results = (diffusion_results, advection_results, advection_diffusion_results)
names = ("diffusion only", "advection only", "advection-diffusion")
linestyles = ("-", "--", ":")
specialcolors = ("xkcd:black", "xkcd:indigo", "xkcd:wine red")

# Solutions
close("all")

fig, axs = subplots(nrows=2, figsize=(12, 6), sharex=true)

legends = plot_solutions!(axs, all_results, names, linestyles, specialcolors)

filepath = joinpath(@__DIR__, "figs", "gaussian_advection_diffusion_solutions.png")
savefig(filepath, dpi=480, bbox_extra_artists=legends, bbox_inches="tight")

# Error profile
fig, axs = subplots()

legend = plot_error_convergence!(axs, Nx, all_results, names)

filepath = joinpath(@__DIR__, "figs", "gaussian_advection_diffusion_error_convergence.png")
savefig(filepath, dpi=480, bbox_extra_artists=(legend,), bbox_inches="tight")

# Test rate of convergence
for (results, name) in zip(all_results, names)
    atol_L₁, atol_L∞ =
        name == "diffusion only" ? (0.05, 0.05) :
        name == "advection only" ? (0.10, 0.50) :
        name == "advection-diffusion" ? (0.05, 0.01) : nothing

    name = "1D Gaussian " * name
    @info "Testing rate of convergence for $name..."

    u_L₁, v_L₁, cx_L₁, cy_L₁, u_L∞, v_L∞, cx_L∞, cy_L∞  = unpack_errors(results)

    test_rate_of_convergence(u_L₁,  Nx, expected=-2.0, atol=atol_L₁, name=name*" u_L₁")
    test_rate_of_convergence(v_L₁,  Nx, expected=-2.0, atol=atol_L₁, name=name*" v_L₁")
    test_rate_of_convergence(cx_L₁, Nx, expected=-2.0, atol=atol_L₁, name=name*" cx_L₁")
    test_rate_of_convergence(cy_L₁, Nx, expected=-2.0, atol=atol_L₁, name=name*" cy_L₁")
    test_rate_of_convergence(u_L∞,  Nx, expected=-2.0, atol=atol_L∞, name=name*" u_L∞")
    test_rate_of_convergence(v_L∞,  Nx, expected=-2.0, atol=atol_L∞, name=name*" v_L∞")
    test_rate_of_convergence(cx_L∞, Nx, expected=-2.0, atol=atol_L∞, name=name*" cx_L∞")
    test_rate_of_convergence(cy_L∞, Nx, expected=-2.0, atol=atol_L∞, name=name*" cy_L∞")

    @test u_L₁ ≈ v_L₁ ≈ cx_L₁ ≈ cy_L₁
    @test u_L∞ ≈ v_L∞ ≈ cx_L∞ ≈ cy_L∞
end


text:
using Random, Combinatorics, Plots, LaTeXStrings
Random.seed!(12)

n, N = 5, 10^5

function isUpperLattice(v)
    for i in 1:Int(length(v)/2)
        sum(v[1:2*i-1]) >= i ? continue : return false
    end
    return true
end

omega = unique(permutations([zeros(Int,n);ones(Int,n)]))
A = omega[isUpperLattice.(omega)]
pA_modelI = length(A)/length(omega)

function randomWalkPath(n)
    x, y = 0, 0
    path = []
    while x<n && y<n
       if rand()<0.5
            x += 1
            push!(path,0)
        else
            y += 1
            push!(path,1)
        end
    end
    append!(path, x<n ? zeros(Int64,n-x) : ones(Int64,n-y))
    return path
end

pA_modelIIest = sum([isUpperLattice(randomWalkPath(n)) for _ in 1:N])/N
println("Model I: ",pA_modelI, "\t Model II: ", pA_modelIIest)

function plotPath(v,l,c)
    x,y = 0,0
    graphX, graphY = [x], [y]
    for i in v
        if i == 0
            x += 1
        else
            y += 1
        end
        push!(graphX,x), push!(graphY,y)
    end
    plot!(graphX, graphY,
            la=0.8, lw=2, label=l, c=c, ratio=:equal, legend=:topleft,
            xlims=(0,n), ylims=(0,n),
            xlabel=L"East\rightarrow", ylabel=L"North\rightarrow")
end
plot()
plotPath(rand(A), "Upper lattice path", :blue)
plotPath(rand(setdiff(omega,A)), "Non-upper lattice path", :red)
plot!([0, n], [0,n], ls=:dash, c=:black, label="")


text:
modify(::Nothing, ::Function) = nothing
modify(node::Node, modifier::Function) = modifier(node)

modify(program::Program, modifier::Function) =
    Program(map(statement -> modify(statement, modifier), program.statements))

modify(node::LetStatement, modifier::Function) =
    LetStatement(node.token, node.name, modify(node.value, modifier), node.reassign)

modify(node::ReturnStatement, modifier::Function) =
    ReturnStatement(node.token, modify(node.return_value, modifier))

modify(node::ExpressionStatement, modifier::Function) =
    ExpressionStatement(node.token, modify(node.expression, modifier))

modify(node::BlockStatement, modifier::Function) = BlockStatement(
    node.token,
    map(statement -> modify(statement, modifier), node.statements),
)

modify(node::PrefixExpression, modifier::Function) =
    PrefixExpression(node.token, node.operator, modify(node.right, modifier))

modify(node::InfixExpression, modifier::Function) = InfixExpression(
    node.token,
    modify(node.left, modifier),
    node.operator,
    modify(node.right, modifier),
)

modify(node::IfExpression, modifier::Function) = IfExpression(
    node.token,
    modify(node.condition, modifier),
    modify(node.consequence, modifier),
    modify(node.alternative, modifier),
)

modify(node::IndexExpression, modifier::Function) =
    IndexExpression(node.token, modify(node.left, modifier), modify(node.index, modifier))

modify(node::ArrayLiteral, modifier::Function) =
    ArrayLiteral(node.token, map(expression -> modify(expression, modifier), node.elements))

modify(node::HashLiteral, modifier::Function) = HashLiteral(
    node.token,
    Dict(
        modify(key, modifier) => modify(value, modifier) for
        (key, value) in collect(node.pairs)
    ),
)

modify(node::FunctionLiteral, modifier::Function) = FunctionLiteral(
    node.token,
    node.parameters,
    modify(node.body, modifier);
    name = node.name,
)

# !!! Do not add the function below.
# !!! We should not call recursively on `CallExpression`s since we need to actually modify them (the `unquote()` calls).
# modify(node::CallExpression, modifier::Function) =
#   CallExpression(node.token, modify(node.function, modifier), map(expression -> modify(expression, modifier), node.arguments))


text:
#############################################################################
using LinearAlgebra
using SparseArrays
using Compat
using Random
using Distributions
using Gurobi
using JuMP
using Printf
using Plots
using DelimitedFiles

# Play with the parameters marked by <--

###################################################
# Problem parameters:
###################################################
NoExp = 1; # <--
NoIter = 1000; # <--
P = 20;
L = 10;
SystemIsSparse = "no";
SystemSparsityLevel = 10/100;
SNR = 20.0; # In dB. The variance of the input signal is equal to 1. # <--
VarOutliers = 1.2e4; # <--
ARcoeffU = 0.95;
ARcoeffL = 0.0;
pBernoulli = 0.2; # <--
K0 = 500; # <--
#
Radius = 20;
gamma = 1;
#
alpha = 0.5;
lambda = 1e-1; # <--
UseBatchMethod = "FMHSDM"; # Choices are: ADMM, GIST, FMHSDM. # <--
alphaFMHSDM = 0.5;
lambdao = 1e-1; # <--
lambdaFMHSDM = 150; # <--
JFMHSDM = 100; # <--
eVarPi = 5e-2;
lambdaF = 1e-5;
rho = 1e2;
MultiplyGradBy = 1;
#
epsilonORRLS = 1e-1; # <--
lambdaORRLSl1 = 1e2; # <--
rhoADMM = 1e2; # <--
JADMM = 100; # <--
# Parameters of the GIST algorithm:
thetaMCP = 2.0; # <--
lambdaMCP = 1e2; # <--
tminMCP = 1e-30;
tmaxMCP = 1e30;
etaMCP = 2.0;
sigmaMCP = 1e-5;
mBuffMCP = 5;
JGIST = 100; # <--

#############################################################################
# Functions:
#############################################################################

struct GenerateProblemParams
    NoIter::Int64
    P::Int64
    L::Int64
    SystemIsSparse::String
    #
    Input::Array{Float64,2}
    System::Array{Float64,2}
    Output::Array{Float64,2}
    LastInput0::Array{Float64,1}
    LastOutput0::Array{Float64,1}
    Rnoise::Array{Float64,2}
    invRnoise::Array{Float64,2}
    MaxEigInvRnoise::Float64
    #
    eVarPi::Float64
    alpha::Float64
    Radius::Float64
    gamma::Float64
    epsilonORRLS::Float64
    F0::Array{Float64,2}
    o0::Array{Float64,1}
    F0HRLS::Array{Float64,2}
    o0HRLS::Array{Float64,1}
    lambdaF::Float64
    lambdao::Float64
    rho::Float64
    lambda::Float64
    MultiplyGradBy::Float64
    lambdaFMHSDM::Float64
    alphaFMHSDM::Float64
    JFMHSDM::Int64
    UseBatchMethod::String
    lambdaORRLSl1::Float64
    rhoADMM::Float64
    JADMM::Int64
    thetaMCP::Float64
    lambdaMCP::Float64
    tminMCP::Float64
    tmaxMCP::Float64
    etaMCP::Float64
    sigmaMCP::Float64
    mBuffMCP::Int64
    JGIST::Int64
end

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

function GenerateModelMatrices(L::Int64, ARcoeffU::Float64, ARcoeffL::Float64, SNR::Float64, SystemIsSparse::String, SystemSparsityLevel::Float64)

    if SystemIsSparse == "yes"
        SystemVec = zeros(P*L);
        NumberNonZeroEntries = ceil(SystemSparsityLevel*P*L);
        NumberNonZeroEntries = convert(Int64,NumberNonZeroEntries);
        SystemVec[1:NumberNonZeroEntries] = ones(NumberNonZeroEntries);
        SystemVec = shuffle(SystemVec);
        System = reshape(SystemVec, (L,P));
    else
        System = randn(Float64,(L,P));
    end

    FAuxMat =  svd(randn(L,L));
    singularVal = [ARcoeffL*ones(L-1) + (ARcoeffU-ARcoeffL)*rand(L-1); ARcoeffU];
    ARmatrix = FAuxMat.U*Diagonal(singularVal)*FAuxMat.Vt;

    AuxVec = inv(Matrix{Float64}(I,L*L,L*L) - kron(ARmatrix, ARmatrix))*vec(Matrix{Float64}(I,L,L));
    AuxMat = reshape(AuxVec, (L,L));
    AuxMat = Symmetric((AuxMat + AuxMat')/2);
    VarWhiteNoise = 10^(-SNR/10)/opnorm(AuxMat,2);
    # VarWhiteNoise = (1-ARcoeffU^2)*10^(-SNR/10);

    Rnoise = VarWhiteNoise*AuxMat;
    invRnoise = Symmetric(inv(Rnoise));
    MaxEigRnoise = opnorm(Rnoise, 2);
    MaxEigInvRnoise = opnorm(invRnoise, 2);

    return System, Rnoise, invRnoise, MaxEigInvRnoise, VarWhiteNoise, ARmatrix;

end

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

function InputOutput(P::Int64, L::Int64, NoIter::Int64, System::Matrix{Float64}, ARmatrix::Matrix{Float64}, pBernoulli::Float64, VarOutliers::Float64, VarWhiteNoise::Float64, FirstNoiseVec::Vector{Float64})

    Input = randn(Float64,(P,NoIter+1));

    # Coloring noise:
    Noise = Matrix{Float64}(undef,(L,NoIter+1));
    Noise[:,1] = FirstNoiseVec;
    for n in 1:NoIter
        Noise[:,n+1] = ARmatrix*Noise[:,n] + sqrt(VarWhiteNoise)*randn(Float64,L);
    end

    MaskOutliersPDF = Binomial(1,pBernoulli);
    MaskOutliers = rand(MaskOutliersPDF,L,NoIter+1);
    ModulusOutliers = sqrt(3*VarOutliers);
    Outliers = -ModulusOutliers*ones(L,NoIter+1) + 2*ModulusOutliers*rand(L,NoIter+1);
    Outliers = Outliers.*MaskOutliers;

    Output = System*Input + Outliers + Noise;

    return Input, Output, Noise[:,NoIter+1];
end

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

function WarmStart(P, L, K0, Factor, invRnoise, lambda, Input, Output)

    MDL = Model(with_optimizer(Gurobi.Optimizer, PSDTol = 1e0, OutputFlag = 0, BarHomogeneous = 1, NumericFocus = 3));

    @variable(MDL, op[l=1:L, k=1:K0]);
    @variable(MDL, on[l=1:L, k=1:K0]);
    @variable(MDL, F[l=1:L, p=1:P]);

    @constraint(MDL, op .>= 0.0);
    @constraint(MDL, on .>= 0.0);

    AuxMat = F*Input + op - on - Output;
    OBJa = sum(AuxMat[:,k]'*invRnoise*AuxMat[:,k] for k=1:K0)/Factor;
    OBJb = lambda*sum(op+on);

    @objective(MDL, Min, OBJa+OBJb);
    optimize!(MDL);

    Fhat = value.(F);
    ohat = value.(op-on);

    # env = Gurobi.Env();
    # M = Model(solver = GurobiSolver(env, OutputFlag=0));

    # @variable(M, o[l=1:L, k=1:K0]);
    # @variable(M, op[l=1:L, k=1:K0]);
    # @variable(M, on[l=1:L, k=1:K0]);
    # @variable(M, z[l=1:L, k=1:K0]);
    # @variable(M, F[l=1:L, p=1:P]);

    # @constraint(M, z .== op + on);
    # @constraint(M, o .== op - on);
    # @constraint(M, op .>= 0);
    # @constraint(M, on .>= 0);

    # @objective(M, Min, sum(sum((Output[l,k] - sum(F[l,p]*Input[p,k] for p=1:P) - o[l,k])*sum(invRnoise[l,j]*(Output[j,k] - sum(F[j,p]*Input[p,k] for p=1:P) - o[j,k])/2 for j=1:L) for l=1:L) for k=1:K0) + lambda*sum(sum(z[l,k] for l=1:L) for k=1:K0));

    # solve(M);

    # Fhat = getvalue(F);
    # ohat = getvalue(o);

    return Fhat, ohat[:,K0];

end

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

function UBoundLambdaORRLSl1(P, L, K0, invRnoise, Input, Output)

    Lambdas = Vector{Float64}(undef, K0);

    for k in 1:K0
        F = Output[:,k]*Input[:,k]'*pinv(Input[:,k]*Input[:,k]');
        Dummy = invRnoise*(Output[:,k] - F*Input[:,k]);
        Lambdas[k] = norm(Dummy, Inf);
    end

    return norm(Lambdas, Inf);

end

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

function Proxl2(X::Array{Float64}, gamma::Float64)
    return X/(1+gamma);
end

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

function Projl2(X::Matrix{Float64}, Radius::Float64)
    return X*Radius/max(Radius, norm(X,2));
end

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

function Proxl1(X::Array{Float64}, gamma::Float64)
    return X .* (1 .- gamma ./ max.(gamma, abs.(X)));
end

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

function Proxl1scalar(x::Float64, gamma::Float64)
    return x*(1 - gamma/max(gamma, abs(x)));
end

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

function MCPloss(w::Float64, lambda::Float64, theta::Float64)

    if abs(w) <= theta*lambda
        loss = lambda*abs(w) - (w^2)/theta/2;
    else
        loss = theta*(lambda^2)/2;
    end

    return loss;

end

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

function ProxMCP(u::Float64, lambda::Float64, t::Float64, theta::Float64)

    if (t*theta != 1)
        w1 = t*theta*(u - lambda/t)/(t*theta-1);
        w2 = t*theta*(u + lambda/t)/(t*theta-1);
        C = [0; w1; w2; u];
        Aux = Vector{Float64}(undef,length(C));
        for i in 1:(length(C)-1)
            Aux[i] = .5*(C[i]-u)^2 + lambda*abs(C[i])/t - ((C[i])^2)/2/theta/t;
        end
        Aux[length(C)] = theta*(lambda^2)/2/t;
        iStar = argmin(Aux);
        prox = C[iStar];
    else
        C = [0; theta*lambda; -theta*lambda; u];
        Aux = Vector{Float64}(undef,length(C));
        for i in 1:(length(C)-1)
            Aux[i] = - C[i]*u + lambda*abs(C[i])/t;
        end
        Aux[length(C)] = theta*(lambda^2)/2/t;
        iStar = argmin(Aux);
        prox = C[iStar];
    end

    return prox;

end

#############################################################################
# function sthresh_cd_MCP(x, l, R, theta)
#     if abs(x) <= theta*l
#         (R*theta/(R*theta-1))*sign(x)*max((abs(x)-l/R),0);
#     else
#         x;
#     end
# end
#############################################################################

function FMHSDMforLASSO(Resolvent::Matrix{Float64}, TargetVec::Array{Float64}, lambdao::Float64, lambda::Float64, alpha::Float64, JFMHSDM::Int64, StartingPoint::Array{Float64})

    X = repeat(StartingPoint, inner = [1,2]);
    Proj = sum(X, dims=2)/2;
    Tmap = repeat(Proj, inner = [1,2]);
    Xhalf = alpha*Tmap + (1-alpha)*X;
    Xprev = X;
    TmapPrev = Tmap;
    X[:,1] = Resolvent*(TargetVec + Xhalf[:,1]/lambda);
    X[:,2] = Proxl1(Xhalf[:,2], lambda*lambdao);

    for j in 1:JFMHSDM

        Proj = sum(X, dims=2)/2;
        Tmap = repeat(Proj, inner = [1,2]);
        Xhalf = Xhalf - alpha*TmapPrev - (1-alpha)*Xprev + Tmap;
        Xprev = X;
        TmapPrev = Tmap;
        X[:,1] = Resolvent*(TargetVec + Xhalf[:,1]/lambda);
        X[:,2] = Proxl1(Xhalf[:,2], lambda*lambdao);

    end

    return sum(X, dims=2)/2;

end

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

function ADMM(b::Vector{Float64}, Resolvent::Matrix{Float64}, lambda::Float64, rho::Float64, J::Int64, x0::Vector{Float64})

    L = size(x0,1);
    x = x0;
    u = x0;
    z = x0;

    for j in 1:J
        x = Resolvent*(b + rho*(z-u));
        z = Proxl1(x+u, lambda/rho);
        u = u + x - z;
    end

    return x;

end

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

function CD(R_inv, Dy, lambda, alpha, J, o0)

    o = Matrix{Float64}(undef, (L,J+1));
    o[:,1] = o0;

    for j in 1:J
        for d in 1:L
            R_inv_d1 = 0;
            R_inv_d2 = 0;
            for d1 in 1:(d-1)
                R_inv_d1 = R_inv_d1 + R_inv[d1,d]*o[d1,j+1];
            end
            for d2 in (d+1):L
                R_inv_d2 = R_inv_d2 + R_inv[d2,d]*o[d2,j];
            end
            gamma = (1/R_inv[d,d])*(alpha[d] - R_inv_d1 - R_inv_d2);
            o[d,j+1] = Proxl1scalar(gamma, lambda/R_inv[d,d]);
        end
    end
    return o[:,J+1];

end

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

function CD_MCP(R_inv, Dy, lambda, alpha, theta, J, o0)

    o = Matrix{Float64}(undef, (L,J+1));
    o[:,1] = o0;

    for j in 1:J
        for d in 1:L
            R_inv_d1 = 0;
            R_inv_d2 = 0;
            for d1 in 1:(d-1)
                R_inv_d1 = R_inv_d1 + R_inv[d1,d]*o[d1,j+1];
            end
            for d2 in (d+1):L
                R_inv_d2 = R_inv_d2 + R_inv[d2,d]*o[d2,j];
            end
            gamma = (1/R_inv[d,d])*(alpha[d] - R_inv_d1 - R_inv_d2);

            # o[d,j+1] = sthresh_cd_MCP(gamma, lambda, R_inv[d,d], theta);
            o[d,j+1] = ProxMCP(gamma, lambda, R_inv[d,d], theta);

        end
    end
    return o[:,J+1];

end

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

function GIST(w0::Vector{Float64}, L::Int64, b::Vector{Float64}, invRnoise::Matrix{Float64}, lambda::Float64, theta::Float64, tmin::Float64, tmax::Float64, eta::Float64, sigma::Float64, mBuff::Int64, JGIST::Int64)

    w = w0;
    wPrev = zeros(L);
    MaxFbuff = zeros(mBuff);
    Regularizer = 0;
    for l in 1:L
        Regularizer = Regularizer + MCPloss(w[l], lambda, theta);
    end
    MaxFbuff[mBuff] = (w-b)'*invRnoise*(w-b)/2 + Regularizer;
    MaxF = maximum(MaxFbuff);

    for k in 1:JGIST

        # Sec. 2.3.1:
        x = w-wPrev;
        y = invRnoise*x;
        if norm(x) == 0
            t = 1.0;
        else
            # Barzilai-Borwein rule:
            t = x'*y/(x'*x);
            t = max(min(t,tmax), tmin);
        end
        wPrev = w;

        Criterion = 1;
        tau = t;
        wNext = Vector{Float64}(undef,L);
        CurrValF = 0;
        while (Criterion > 0) && (tau < tmax)
            u = w - invRnoise*(w-b)/tau;
            Regularizer = 0;
            for l in 1:L
                wNext[l] = ProxMCP(u[l], lambda, tau, theta);
                Regularizer = Regularizer + MCPloss(wNext[l], lambda, theta);
            end
            CurrValF = (wNext-b)'*invRnoise*(wNext-b)/2 + Regularizer;
            Criterion = CurrValF - MaxF + sigma*tau*(norm(wNext-w)^2)/2;
            tau = eta*tau;
        end
        w = wNext;
        MaxFbuff[1:(mBuff-1)] = MaxFbuff[2:mBuff];
        MaxFbuff[mBuff] = CurrValF;
        MaxF = maximum(MaxFbuff);

    end

    return w;

end

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

function RLS(ProblemParams)

    L = ProblemParams.L;
    NoIter = ProblemParams.NoIter;
    Input = ProblemParams.Input;
    Output = ProblemParams.Output;
    Rnoise = ProblemParams.Rnoise;
    invRnoise = ProblemParams.invRnoise;
    epsilonORRLS = ProblemParams.epsilonORRLS;
    System = ProblemParams.System;
    F0 = ProblemParams.F0;
    gamma = ProblemParams.gamma;

    Phi = Matrix{Float64}(I,L*P,L*P)/epsilonORRLS;
    Phi = Symmetric(Phi);

    NRMSE = Vector{Float64}(undef,NoIter);

    # F = F0;
    # F = randn(L,P);
    Fvec = vec(F0);

    BEGIN = time();
    for k in 2:(NoIter+1)

        H = kron(Input[:,k]', Matrix{Float64}(I,L,L));

        # RLS updates:
        DummyMat = H*Phi*H'/gamma;
        DummyMat = (DummyMat + DummyMat')/2;
        DummyMat = Symmetric(DummyMat);
        W = (1/gamma)*Phi*H'*inv(DummyMat + Rnoise);
        Fvec = Fvec + W*(Output[:,k] - H*Fvec);
        Phi = (1/gamma)*(Phi - W*H*Phi);
        Phi = (Phi+Phi')/2;
        Phi = Symmetric(Phi);

        F = reshape(Fvec,L,P);
        NRMSE[k-1] = norm(F-System,2)/norm(System,2);

    end
    END = time();
    Time = (END-BEGIN)/NoIter;

    return NRMSE, Time;

end

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

function ORRLSl1(ProblemParams)

    L = ProblemParams.L;
    P = ProblemParams.P;
    NoIter = ProblemParams.NoIter;
    Input = ProblemParams.Input;
    Output = ProblemParams.Output;
    lambdaORRLSl1 = ProblemParams.lambdaORRLSl1;
    Rnoise = ProblemParams.Rnoise;
    invRnoise = ProblemParams.invRnoise;
    rhoADMM = ProblemParams.rhoADMM;
    JADMM = ProblemParams.JADMM;
    epsilonORRLS = ProblemParams.epsilonORRLS;
    System = ProblemParams.System;
    F0 = ProblemParams.F0;
    o0 = ProblemParams.o0;

    NRMSE = Vector{Float64}(undef,NoIter);

    Fvec = vec(F0);
    # F = randn(L,P);
    o = o0;
    Resolvent = inv(invRnoise + rhoADMM*Matrix{Float64}(I,L,L));
    Phi = Matrix{Float64}(I,L*P,L*P)/epsilonORRLS;
    SystemVec = vec(System);

    BEGIN = time();
    for k in 2:(NoIter+1)

        H = kron(Input[:,k]', Matrix{Float64}(I,L,L));

        auxVec = invRnoise*(Output[:,k] - H*Fvec);
        o = ADMM(auxVec, Resolvent, lambdaORRLSl1, rhoADMM, JADMM, o);

        # RLS updates:
        DummyMat = H*Phi*H'/gamma;
        DummyMat = (DummyMat + DummyMat')/2;
        DummyMat = Symmetric(DummyMat);
        W = (1/gamma)*Phi*H'*inv(DummyMat + Rnoise);
        Fvec = Fvec + W*(Output[:,k] - o - H*Fvec);
        Phi = (1/gamma)*(Phi - W*H*Phi);
        Phi = (Phi+Phi')/2;
        Phi = Symmetric(Phi);

        NRMSE[k-1] = norm(Fvec-SystemVec)/norm(SystemVec);

    end
    END = time();
    Time = (END-BEGIN)/NoIter;

    return NRMSE, Time;

end

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

function ORRLSCD_L1(ProblemParams)

    L = ProblemParams.L;
    P = ProblemParams.P;
    NoIter = ProblemParams.NoIter;
    Input = ProblemParams.Input;
    Output = ProblemParams.Output;
    lambdaORRLSl1 = ProblemParams.lambdaORRLSl1;
    Rnoise = ProblemParams.Rnoise;
    invRnoise = ProblemParams.invRnoise;
    epsilonORRLS = ProblemParams.epsilonORRLS;
    System = ProblemParams.System;
    F0 = ProblemParams.F0;
    o0 = ProblemParams.o0;
    thetaMCP = ProblemParams.thetaMCP;
    lambdaMCP = ProblemParams.lambdaMCP;
    tminMCP = ProblemParams.tminMCP;
    tmaxMCP = ProblemParams.tmaxMCP;
    etaMCP = ProblemParams.etaMCP;
    sigmaMCP = ProblemParams.sigmaMCP;
    mBuffMCP = ProblemParams.mBuffMCP;
    JGIST = ProblemParams.JGIST;

    NRMSE = Vector{Float64}(undef,NoIter);

    Fvec = vec(F0);
    # F = randn(L,P);
    o = o0;
    Phi = Matrix{Float64}(I,L*P,L*P)/epsilonORRLS;

    SystemVec = vec(System);

    BEGIN = time();
    for k in 2:(NoIter+1)

        H = kron(Input[:,k]', Matrix{Float64}(I,L,L));

        auxVec = invRnoise*(Output[:,k] - H*Fvec);
        o = CD(invRnoise, L, lambdaORRLSl1, auxVec, JADMM, o);

        # RLS updates:
        DummyMat = H*Phi*H'/gamma;
        DummyMat = (DummyMat + DummyMat')/2;
        DummyMat = Symmetric(DummyMat);
        W = (1/gamma)*Phi*H'*inv(DummyMat + Rnoise);
        Fvec = Fvec + W*(Output[:,k] - o - H*Fvec);
        Phi = (1/gamma)*(Phi - W*H*Phi);
        Phi = (Phi+Phi')/2;
        Phi = Symmetric(Phi);

        NRMSE[k-1] = norm(Fvec-SystemVec)/norm(SystemVec);

    end
    END = time();
    Time = (END-BEGIN)/NoIter;

    return NRMSE, Time;
end

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

function ORRLSCD_MCP(ProblemParams)

    L = ProblemParams.L;
    P = ProblemParams.P;
    NoIter = ProblemParams.NoIter;
    Input = ProblemParams.Input;
    Output = ProblemParams.Output;
    lambdaORRLSl1 = ProblemParams.lambdaORRLSl1;
    Rnoise = ProblemParams.Rnoise;
    invRnoise = ProblemParams.invRnoise;
    epsilonORRLS = ProblemParams.epsilonORRLS;
    System = ProblemParams.System;
    F0 = ProblemParams.F0;
    o0 = ProblemParams.o0;
    thetaMCP = ProblemParams.thetaMCP;
    lambdaMCP = ProblemParams.lambdaMCP;
    tminMCP = ProblemParams.tminMCP;
    tmaxMCP = ProblemParams.tmaxMCP;
    etaMCP = ProblemParams.etaMCP;
    sigmaMCP = ProblemParams.sigmaMCP;
    mBuffMCP = ProblemParams.mBuffMCP;
    JGIST = ProblemParams.JGIST;

    NRMSE = Vector{Float64}(undef,NoIter);

    Fvec = vec(F0);
    # F = randn(L,P);
    o = o0;
    Resolvent = inv(invRnoise + rhoADMM*Matrix{Float64}(I,L,L));
    Phi = Matrix{Float64}(I,L*P,L*P)/epsilonORRLS;
    SystemVec = vec(System);

    BEGIN = time();
    for k in 2:(NoIter+1)

        H = kron(Input[:,k]', Matrix{Float64}(I,L,L));

        auxVec = invRnoise*(Output[:,k] - H*Fvec);
        o = CD_MCP(invRnoise, L, lambdaMCP, auxVec, thetaMCP, JGIST, o);

        # RLS updates:
        DummyMat = H*Phi*H'/gamma;
        DummyMat = (DummyMat + DummyMat')/2;
        DummyMat = Symmetric(DummyMat);
        W = (1/gamma)*Phi*H'*inv(DummyMat + Rnoise);
        Fvec = Fvec + W*(Output[:,k] - o - H*Fvec);
        Phi = (1/gamma)*(Phi - W*H*Phi);
        Phi = (Phi+Phi')/2;
        Phi = Symmetric(Phi);

        NRMSE[k-1] = norm(Fvec-SystemVec)/norm(SystemVec);

    end
    END = time();
    Time = (END-BEGIN)/NoIter;

    return NRMSE, Time;
end

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

function ORRLSMCP(ProblemParams)

    L = ProblemParams.L;
    P = ProblemParams.P;
    NoIter = ProblemParams.NoIter;
    Input = ProblemParams.Input;
    Output = ProblemParams.Output;
    lambdaORRLSl1 = ProblemParams.lambdaORRLSl1;
    Rnoise = ProblemParams.Rnoise;
    invRnoise = ProblemParams.invRnoise;
    epsilonORRLS = ProblemParams.epsilonORRLS;
    System = ProblemParams.System;
    F0 = ProblemParams.F0;
    o0 = ProblemParams.o0;
    thetaMCP = ProblemParams.thetaMCP;
    lambdaMCP = ProblemParams.lambdaMCP;
    tminMCP = ProblemParams.tminMCP;
    tmaxMCP = ProblemParams.tmaxMCP;
    etaMCP = ProblemParams.etaMCP;
    sigmaMCP = ProblemParams.sigmaMCP;
    mBuffMCP = ProblemParams.mBuffMCP;
    JGIST = ProblemParams.JGIST;

    NRMSE = Vector{Float64}(undef,NoIter);

    Fvec = vec(F0);
    # F = randn(L,P);
    o = o0;
    Phi = Matrix{Float64}(I,L*P,L*P)/epsilonORRLS;
    SystemVec = vec(System);

    BEGIN = time();
    for k in 2:(NoIter+1)

        H = kron(Input[:,k]', Matrix{Float64}(I,L,L));

        auxVec = Output[:,k] - H*Fvec;
        o = GIST(o, L, auxVec, invRnoise, lambdaMCP, thetaMCP, tminMCP, tmaxMCP, etaMCP, sigmaMCP, mBuffMCP, JGIST);

        # RLS updates:
        DummyMat = H*Phi*H'/gamma;
        DummyMat = (DummyMat + DummyMat')/2;
        DummyMat = Symmetric(DummyMat);
        W = (1/gamma)*Phi*H'*inv(DummyMat + Rnoise);
        Fvec = Fvec + W*(Output[:,k] - o - H*Fvec);
        Phi = (1/gamma)*(Phi - W*H*Phi);
        Phi = (Phi+Phi')/2;
        Phi = Symmetric(Phi);

        NRMSE[k-1] = norm(Fvec-SystemVec)/norm(SystemVec);

    end
    END = time();
    Time = (END-BEGIN)/NoIter;

    return NRMSE, Time;

end

#############################################################################
function HRLS_ADMM(ProblemParams)

    ###################################################
    # Parameters:
    ###################################################

    eVarPi = ProblemParams.eVarPi;
    alpha = ProblemParams.alpha;
    lambda = ProblemParams.lambda;
    MultiplyGradBy = ProblemParams.MultiplyGradBy;

    NoIter = ProblemParams.NoIter;
    P = ProblemParams.P;
    L = ProblemParams.L;
    lambdaF = ProblemParams.lambdaF;
    Radius = ProblemParams.Radius;
    Input = ProblemParams.Input;
    Output = ProblemParams.Output;
    LastInput0 = ProblemParams.LastInput0;
    LastOutput0 = ProblemParams.LastOutput0;
    MaxEigInvRnoise = ProblemParams.MaxEigInvRnoise;
    invRnoise = ProblemParams.invRnoise;
    System = ProblemParams.System;
    F0 = ProblemParams.F0HRLS;
    o0 = ProblemParams.o0HRLS;
    lambdao = ProblemParams.lambdao;
    rho = ProblemParams.rho;
    lambdaFMHSDM = ProblemParams.lambdaFMHSDM;
    alphaFMHSDM = ProblemParams.alphaFMHSDM;
    JFMHSDM = ProblemParams.JFMHSDM;
    thetaMCP = ProblemParams.thetaMCP;
    lambdaMCP = ProblemParams.lambdaMCP;
    tminMCP = ProblemParams.tminMCP;
    tmaxMCP = ProblemParams.tmaxMCP;
    etaMCP = ProblemParams.etaMCP;
    sigmaMCP = ProblemParams.sigmaMCP;
    mBuffMCP = ProblemParams.mBuffMCP;
    JGIST = ProblemParams.JGIST;
    rhoADMM = ProblemParams.rhoADMM;
    JADMM = ProblemParams.JADMM;
    UseBatchMethod = ProblemParams.UseBatchMethod;
    SystemIsSparse = ProblemParams.SystemIsSparse;

    NRMSE = Array{Float64}(undef,NoIter);

    ###################################################
    # Initialization:
    ###################################################

    Fvec = vec(F0);
    o = o0;

    SystemVec = vec(System);

    H = kron(Input[:,1]', Matrix{Float64}(I,L,L));

    auxVec = invRnoise*(Output[:,1] - H*Fvec);
    Resolvent = inv(invRnoise + rhoADMM*Matrix{Float64}(I,L,L));
    o = ADMM(auxVec, Resolvent, lambdaORRLSl1, rhoADMM, JADMM, o);

    K = H'*invRnoise*H; K = (K+K')/2; K = Symmetric(K);
    Q = H'*invRnoise*(Output[:,1]-o);

    Gamma = 1.0;
    VV = eigen(K);
    vVec = VV.vectors[:,P*L];

    FvecPrev = Fvec;

    varpi = MaxEigInvRnoise*(norm(Input[:,1])^2)/Gamma;
    Grad = MultiplyGradBy*(1/Gamma)*(K*Fvec - Q)/varpi;
    FvecHalf = Fvec - alpha*Grad;

    if SystemIsSparse == "yes"
        Fvec = Proxl1(FvecHalf, lambda);
    else
        Fvec = FvecHalf;
    end

    GradPrev = Grad;

    BEGIN = time();
    for n in 2:(NoIter+1)

        H = kron(Input[:,n]', Matrix{Float64}(I,L,L));

        auxVec = invRnoise*(Output[:,n] - H*Fvec);
        o = ADMM(auxVec, Resolvent, lambdaORRLSl1, rhoADMM, JADMM, o);

        Gamma = gamma*Gamma + 1.0;
        K = gamma*K + H'*invRnoise*H;
        Q = gamma*Q + H'*invRnoise*(Output[:,n]-o);

        uVec = K*vVec/Gamma;
        vVec = uVec/norm(uVec);
        varpi = vVec'*K*vVec/Gamma + eVarPi;

        Grad = MultiplyGradBy*(1/Gamma)*(K*Fvec - Q)/varpi;

        FvecHalf = FvecHalf + Fvec - FvecPrev + alpha*GradPrev - Grad;

        FvecPrev = Fvec;

        Fvec = FvecHalf;

        if SystemIsSparse == "yes"
            Fvec = Proxl1(FvecHalf, lambda);
        else
            Fvec = FvecHalf;
        end
        GradPrev = Grad;

        NRMSE[n-1] = norm(Fvec-SystemVec)/norm(SystemVec);

    end
    END = time();
    Time = (END-BEGIN)/NoIter;

    return NRMSE, Time;

end

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

function HRLS_GIST(ProblemParams)

    ###################################################
    # Parameters:
    ###################################################

    eVarPi = ProblemParams.eVarPi;
    alpha = ProblemParams.alpha;
    lambda = ProblemParams.lambda;
    MultiplyGradBy = ProblemParams.MultiplyGradBy;

    NoIter = ProblemParams.NoIter;
    P = ProblemParams.P;
    L = ProblemParams.L;
    lambdaF = ProblemParams.lambdaF;
    Radius = ProblemParams.Radius;
    Input = ProblemParams.Input;
    Output = ProblemParams.Output;
    LastInput0 = ProblemParams.LastInput0;
    LastOutput0 = ProblemParams.LastOutput0;
    MaxEigInvRnoise = ProblemParams.MaxEigInvRnoise;
    invRnoise = ProblemParams.invRnoise;
    System = ProblemParams.System;
    F0 = ProblemParams.F0HRLS;
    o0 = ProblemParams.o0HRLS;
    lambdao = ProblemParams.lambdao;
    rho = ProblemParams.rho;
    lambdaFMHSDM = ProblemParams.lambdaFMHSDM;
    alphaFMHSDM = ProblemParams.alphaFMHSDM;
    JFMHSDM = ProblemParams.JFMHSDM;
    thetaMCP = ProblemParams.thetaMCP;
    lambdaMCP = ProblemParams.lambdaMCP;
    tminMCP = ProblemParams.tminMCP;
    tmaxMCP = ProblemParams.tmaxMCP;
    etaMCP = ProblemParams.etaMCP;
    sigmaMCP = ProblemParams.sigmaMCP;
    mBuffMCP = ProblemParams.mBuffMCP;
    JGIST = ProblemParams.JGIST;
    rhoADMM = ProblemParams.rhoADMM;
    JADMM = ProblemParams.JADMM;
    UseBatchMethod = ProblemParams.UseBatchMethod;
    SystemIsSparse = ProblemParams.SystemIsSparse;

    NRMSE = Array{Float64}(undef,NoIter);

    ###################################################
    # Initialization:
    ###################################################

    Fvec = vec(F0);
    o = o0;

    SystemVec = vec(System);

    H = kron(Input[:,1]', Matrix{Float64}(I,L,L));

    auxVec = Output[:,1] - H*Fvec;
    o = GIST(o, L, auxVec, invRnoise, lambdaMCP, thetaMCP, tminMCP, tmaxMCP, etaMCP, sigmaMCP, mBuffMCP, JGIST);

    K = H'*invRnoise*H; K = (K+K')/2; K = Symmetric(K);

    Q = H'*invRnoise*(Output[:,1]-o);

    Gamma = 1.0;
    VV = eigen(K);
    vVec = VV.vectors[:,P*L];

    FvecPrev = Fvec;
    varpi = MaxEigInvRnoise*(norm(Input[:,1])^2)/Gamma;
    Grad = MultiplyGradBy*(1/Gamma)*(K*Fvec - Q)/varpi;
    FvecHalf = Fvec - alpha*Grad;

    if SystemIsSparse == "yes"
        Fvec = Proxl1(FvecHalf, lambda);
    else
        Fvec = FvecHalf;
    end

    GradPrev = Grad;

    BEGIN = time();
    for n in 2:(NoIter+1)

        H = kron(Input[:,n]', Matrix{Float64}(I,L,L));

        auxVec = Output[:,n] - H*Fvec;
        o = GIST(o, L, auxVec, invRnoise, lambdaMCP, thetaMCP, tminMCP, tmaxMCP, etaMCP, sigmaMCP, mBuffMCP, JGIST);

        Gamma = gamma*Gamma + 1.0;
        K = gamma*K + H'*invRnoise*H;

        Q = gamma*Q + H'*invRnoise*(Output[:,n]-o);

        uVec = K*vVec/Gamma;
        vVec = uVec/norm(uVec);
        varpi = vVec'*K*vVec/Gamma + eVarPi;

        Grad = MultiplyGradBy*(1/Gamma)*(K*Fvec - Q)/varpi;

        FvecHalf = FvecHalf + Fvec - FvecPrev + alpha*GradPrev - Grad;

        FvecPrev = Fvec;

        Fvec = FvecHalf;

        if SystemIsSparse == "yes"
            Fvec = Proxl1(FvecHalf, lambda);
        else
            Fvec = FvecHalf;
        end

        GradPrev = Grad;

        NRMSE[n-1] = norm(Fvec-SystemVec)/norm(SystemVec);

    end
    END = time();
    Time = (END-BEGIN)/NoIter;

    return NRMSE, Time;

end

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

function HRLS_FMHSDM(ProblemParams)

    ###################################################
    # Parameters:
    ###################################################

    eVarPi = ProblemParams.eVarPi;
    alpha = ProblemParams.alpha;
    lambda = ProblemParams.lambda;
    MultiplyGradBy = ProblemParams.MultiplyGradBy;

    NoIter = ProblemParams.NoIter;
    P = ProblemParams.P;
    L = ProblemParams.L;
    lambdaF = ProblemParams.lambdaF;
    Radius = ProblemParams.Radius;
    Input = ProblemParams.Input;
    Output = ProblemParams.Output;
    LastInput0 = ProblemParams.LastInput0;
    LastOutput0 = ProblemParams.LastOutput0;
    MaxEigInvRnoise = ProblemParams.MaxEigInvRnoise;
    invRnoise = ProblemParams.invRnoise;
    System = ProblemParams.System;
    F0 = ProblemParams.F0HRLS;
    o0 = ProblemParams.o0HRLS;
    lambdao = ProblemParams.lambdao;
    rho = ProblemParams.rho;
    lambdaFMHSDM = ProblemParams.lambdaFMHSDM;
    alphaFMHSDM = ProblemParams.alphaFMHSDM;
    JFMHSDM = ProblemParams.JFMHSDM;
    thetaMCP = ProblemParams.thetaMCP;
    lambdaMCP = ProblemParams.lambdaMCP;
    tminMCP = ProblemParams.tminMCP;
    tmaxMCP = ProblemParams.tmaxMCP;
    etaMCP = ProblemParams.etaMCP;
    sigmaMCP = ProblemParams.sigmaMCP;
    mBuffMCP = ProblemParams.mBuffMCP;
    JGIST = ProblemParams.JGIST;
    rhoADMM = ProblemParams.rhoADMM;
    JADMM = ProblemParams.JADMM;
    UseBatchMethod = ProblemParams.UseBatchMethod;
    SystemIsSparse = ProblemParams.SystemIsSparse;

    NRMSE = Array{Float64}(undef,NoIter);

    ###################################################
    # Initialization:
    ###################################################

    Fvec = vec(F0);
    o = o0;

    SystemVec = vec(System);

    H = kron(Input[:,1]', Matrix{Float64}(I,L,L));

    Resolvent = inv(invRnoise + Matrix{Float64}(I,L,L)/lambdaFMHSDM);
    auxVec = invRnoise*(Output[:,1] - H*Fvec);
    o = FMHSDMforLASSO(Resolvent, auxVec, lambdao, lambdaFMHSDM, alphaFMHSDM, JFMHSDM, o);

    K = H'*invRnoise*H; K = (K+K')/2; K = Symmetric(K);
    Q = H'*invRnoise*(Output[:,1]-o);
    Gamma = 1.0;
    VV = eigen(K);
    vVec = VV.vectors[:,P*L];

    FvecPrev = Fvec;

    varpi = MaxEigInvRnoise*(norm(Input[:,1])^2)/Gamma;
    Grad = MultiplyGradBy*(1/Gamma)*(K*Fvec - Q)/varpi;
    FvecHalf = Fvec - alpha*Grad;
    if SystemIsSparse == "yes"
        Fvec = Proxl1(FvecHalf, lambda);
    else
        Fvec = FvecHalf;
    end

    GradPrev = Grad;

    BEGIN = time();
    for n in 2:(NoIter+1)

        H = kron(Input[:,n]', Matrix{Float64}(I,L,L));

        auxVec = invRnoise*(Output[:,n] - H*Fvec);
        o = FMHSDMforLASSO(Resolvent, auxVec, lambdao, lambdaFMHSDM, alphaFMHSDM, JFMHSDM, o);

        Gamma = gamma*Gamma + 1.0;
        K = gamma*K + H'*invRnoise*H;
        Q = gamma*Q + H'*invRnoise*(Output[:,n]-o);

        uVec = K*vVec/Gamma;
        vVec = uVec/norm(uVec);
        varpi = vVec'*K*vVec/Gamma + eVarPi;

        Grad = MultiplyGradBy*(1/Gamma)*(K*Fvec - Q)/varpi;
        FvecHalf = FvecHalf + Fvec - FvecPrev + alpha*GradPrev - Grad;
        FvecPrev = Fvec;
        Fvec = FvecHalf;
        if SystemIsSparse == "yes"
            Fvec = Proxl1(FvecHalf, lambda);
        else
            Fvec = FvecHalf;
        end
        GradPrev = Grad;

        NRMSE[n-1] = norm(Fvec-SystemVec)/norm(SystemVec);

    end
    END = time();
    Time = (END-BEGIN)/NoIter;

    return NRMSE, Time;

end

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

max_stars_repo_name: KostasSlavakis/OR-HO-RLS

max_issues_repo_name: KostasSlavakis/OR-HO-RLS

max_forks_repo_name: KostasSlavakis/OR-HO-RLS


text:
"""
    HubbardModel(lattice; params...)
    HubbardModel(L, dims; params...)
    HubbardModel(params::Dict)
    HubbardModel(params::NamedTuple)
    HubbardModel(; params...)

Defines an Hubbard model on a given (or derived) `lattice`. If a linear system
size `L` and dimensionality `dims` is given, the`lattice` will be a Cubic
lattice of fitting size.

Additional parameters (keyword arguments) include:
* `l::AbstractLattice = lattice`: The lattice the model uses. The keyword
argument takes precedence over the argument `lattice` as well as `dims` and `L`.
* `U::Float64 = 1.0` is the value of the Hubbard Interaction. A positive `U`
leads to an attractive Hubbard model, a negative to a repulsive.
* `t::Float64 = 1.0` is the hopping strength.
* `mu::Float64` is the chemical potential. In the repulsive model `mu != 0`
leads to a sign problem.

The model will select an appropriate Hirsch field based on `U`.
"""
struct HubbardModel{LT <: AbstractLattice} <: Model
    t::Float64
    mu::Float64
    U::Float64
    l::LT
end

function HubbardModel(;
        dims = 2, L = 2, l = choose_lattice(HubbardModel, dims, L),
        U = 1.0, mu = 0.0, t = 1.0
    )
    if (U < 0.0) && (mu != 0.0)
        @warn("A repulsive Hubbard model with chemical potential µ = $mu will have a sign problem!")
    end
    HubbardModel(t, mu, U, l)
end
HubbardModel(params::Dict{Symbol}) = HubbardModel(; params...)
HubbardModel(params::NamedTuple) = HubbardModel(; params...)
HubbardModel(lattice::AbstractLattice; kwargs...) = HubbardModel(l = lattice; kwargs...)
HubbardModel(L, dims; kwargs...) = HubbardModel(dims = dims, L = L; kwargs...)

HubbardModelAttractive(args...; kwargs...) = HubbardModel(args...; kwargs...)
function HubbardModelRepulsive(args...; kwargs...)
    d = Dict(kwargs)
    d[:U] = - get(d, :U, 1.0)
    HubbardModel(args...; d...)
end

function choose_lattice(::Type{<: HubbardModel}, dims, L)
    if dims == 1
        return Chain(L)
    elseif dims == 2
        return SquareLattice(L)
    else
        return CubicLattice(dims, L)
    end
end

# cosmetics
import Base.summary
import Base.show
function Base.summary(model::HubbardModel)
    (model.U < 0 ? "repulsive" : "attractive") * " Hubbard model"
end
function Base.show(io::IO, model::HubbardModel)
    println(io, (model.U < 0 ? "repulsive" : "attractive") * " Hubbard model")
    println(io, "\tParameters t = $(model.t), µ = $(model.mu), U = $(model.U)")
    print(io, "\t$(nameof(typeof(model.l))) with $(length(model.l)) sites")
end
Base.show(io::IO, m::MIME"text/plain", model::HubbardModel) = print(io, model)

# Convenience
@inline parameters(m::HubbardModel) = (N = length(m.l), t = m.t, U = m.U, mu = m.mu)
choose_field(m::HubbardModel) = m.U < 0.0 ? MagneticHirschField : DensityHirschField

# implement DQMC interface:
@inline lattice(m::HubbardModel) = m.l
nflavors(::HubbardModel) = 1

hopping_eltype(model::HubbardModel) = typeof(model.t)
function hopping_matrix_type(field::AbstractField, model::HubbardModel)
    flv = max(nflavors(field), nflavors(model))
    T = hopping_eltype(model)
    MT = matrix_type(T)
    if flv == 1
        return MT
    else
        return BlockDiagonal{T, flv, MT}
    end
end

function greens_matrix_type(f::AbstractHirschField{T}, m::HubbardModel) where {T}
    if max(nflavors(f), nflavors(m)) == 1
        return matrix_type(T)
    else
        return BlockDiagonal{T, 2, matrix_type(T)}
    end
end
function greens_matrix_type(f::AbstractGHQField{T}, m::HubbardModel) where {T}
    if max(nflavors(f), nflavors(m)) == 1
        return matrix_type(T)
    else
        return BlockDiagonal{T, 2, matrix_type(T)}
    end
end

"""
    hopping_matrix(mc, model)

Calculates the hopping matrix \$T_{i, j}\$ where \$i, j\$ are
site indices.

This isn't a performance critical method as it is only used once before the
actual simulation.
"""
function hopping_matrix(mc::DQMC, m::HubbardModel)
    N = length(m.l)
    T = diagm(0 => fill(-m.mu, N))

    # Nearest neighbor hoppings
    @inbounds @views begin
        for (src, trg) in neighbors(m.l, Val(true))
            trg == -1 && continue
            T[trg, src] += -m.t
        end
    end

    if max(nflavors(field(mc)), nflavors(m)) == 1
        return T
    else
        return BlockDiagonal(T, copy(T))
    end
end

function save_model(file::JLDFile, m::HubbardModel, entryname::String = "Model")
    write(file, entryname * "/VERSION", 1)
    write(file, entryname * "/tag", "HubbardModel")

    write(file, entryname * "/mu", m.mu)
    write(file, entryname * "/U", m.U)
    write(file, entryname * "/t", m.t)
    save_lattice(file, m.l, entryname * "/l")

    nothing
end

function _load(data, ::Val{:HubbardModel})
    l = _load(data["l"], to_tag(data["l"]))
    HubbardModel(data["t"], data["mu"], data["U"], l)
end
_load(data, ::Val{:HubbardModelAttractive}) = _load(data, Val(:HubbardModel))
function _load(data, ::Val{:HubbardModelRepulsive})
    l = _load(data["l"], to_tag(data["l"]))
    HubbardModel(data["t"], 0.0, -data["U"], l)
end

function intE_kernel(mc, model::HubbardModel, G::GreensMatrix, ::Val{1})
    # ⟨U (n↑ - 1/2)(n↓ - 1/2)⟩ = ...
    # = U [G↑↑ G↓↓ - G↓↑ G↑↓ - 0.5 G↑↑ - 0.5 G↓↓ + G↑↓ + 0.25]
    # = U [(G↑↑ - 1/2)(G↓↓ - 1/2) + G↑↓(1 + G↑↓)]
    # with up-up = down-down and up-down = 0
    - model.U * sum((diag(G.val) .- 0.5).^2)
end
# Technically this only applies to BlockDiagonal
function intE_kernel(mc, model::HubbardModel, G::GreensMatrix{Float64}, ::Val{2})
    - model.U * sum((diag(G.val.blocks[1]) .- 0.5) .* (diag(G.val.blocks[2]) .- 0.5))
end
function intE_kernel(mc, model::HubbardModel, G::GreensMatrix{ComplexF64}, ::Val{2})
    N = length(lattice(model))
    output = 0.0 + 0.0im
    for i in 1:N
        output -= (G.val[i, i] .- 0.5) .* (G.val[N+i, N+i] .- 0.5)
    end
    output * model.U
end


text:
struct BattleMain <: AbstractPokemonScene
    player_a_pokes::Vector{PokemonID}
    player_b_pokes::Vector{PokemonID}
end

feature_image_name(::Type{BattleMain}, ::PokemonBattle) = "battle-main"

function parse_pokemon_name(img, i, is_player_a, ctx::PokemonContext)
    rect = if is_player_a
        if i == 1
            Rect((914, 12), (64, 250), resolution = DesignResolution)
        else
            Rect((914, 516), (64, 250), resolution = DesignResolution)
        end
    else
        if i == 1
            Rect((38, 940), (50, 250), resolution = DesignResolution)
        else
            Rect((38, 1450), (50, 250), resolution = DesignResolution)
        end
    end
    area = img[rect]
    text, langs = ocr_multiple_lang(area, ctx)
    lists = [poke_list(lang) for lang in split(langs, '+')]
    pids = sort!([data_search_n(list, text, n = 1)[1] for list in lists])
    pids[1][2]
end

function _parse_scene(::Type{BattleMain}, frame::VsFrame, ctx::PokemonContext)
    img = image(frame)
    parse_type = ctx.config.source.parse_type
    player_a_pokes = parse_type ≡ PARSE_BOTH_PLAYERS ? [parse_pokemon_name(img, i, true, ctx) for i in 1:2] : PokemonID[]
    player_b_pokes = [parse_pokemon_name(img, i, false, ctx) for i in 1:2]
    BattleMain(player_a_pokes, player_b_pokes)
end


text:
using
using Test

@testset ".jl" begin
    # Write your own tests here.
end


text:
using Parameters: @with_kw

# TODO: Cleanup!  The code is very ugly at the moment because I
# tried to support out-of-place and in-place algorithm in the same
# code (and failing to do so).  I should:
#
# * Add some tests for out-of-place code.
# * Add some benchmarks, especially for in-place code.
# * Reduce type-instability, especially for in-place code.
# * Find the right abstraction and then finally clean it up!

"""
Cache for Euler-Newton continuation method.

See [`AbstractContinuationProblem`](@ref) for the mathematical setup.

# Fields
- `prob_cache`
- `u` (size: `(N,)`)
- `H` (size: `(N - 1,)`) ``= H(u)``
- `J` (size: `(N - 1, N)`) ``= \\partial H / \\partial u``
- `Q` (size: `(N - 1, N)`): temporary array for the QR decomposition
- `h::Real`: step size
- `direction::Int`: +1 or -1
- `corrector_success::Bool`
- `adaptation_success::Bool`
- `simple_bifurcation::Bool`
"""
mutable struct ContinuationCache{PC <: AbstractProblemCache,
                                 uType, HType, JType, QType, hType,
                                 } <: AbstractContinuationCache{PC}
    prob_cache::PC
    u::uType
    H::HType
    J::JType
    Q::QType
    h::hType
    direction::Int
    corrector_success::Bool
    adaptation_success::Bool
    simple_bifurcation::Bool
end

function ContinuationCache(prob_cache::AbstractProblemCache,
                           h::Real, direction::Int = 1)
    u0 = get_u0(prob_cache.prob)
    N = length(u0)
    return ContinuationCache(
        prob_cache,
        u0,
        _similar(u0, N - 1),     # H
        _similar(u0, N - 1, N),  # J
        _similar(u0, N, N),      # Q
        h,
        direction,
        false,
        false,
        false,
    )
end

ContinuationCache(prob::AbstractContinuationProblem, args...) =
    ContinuationCache(get_prob_cache(prob), args...)

@with_kw struct ContinuationOptions
    direction::Int = 1
    h0::Float64 = 0.01
    h_min::Float64 = 1e-6
    h_zero::Float64 = 1e-6
    rtol::Float64 = 0.01
    atol::Float64 = 1e-6
    max_samples::Int = 100
    max_adaptations::Int = 100
    max_corrector_steps::Int = 100
    max_branches::Int = 10
    max_misc_steps::Int = 100   # TODO: remove
    nominal_contraction::Float64 = 0.8
    nominal_distance::Float64 = 0.1
    nominal_angle_rad::Float64 = 2π * (10 / 360)
    start_from_nearest_root::Bool = false
    bidirectional_first_sweep::Bool = true
    verbose::Bool = false
end

rawtangent(Q) = vec(rawtangentmat(Q))
function rawtangentmat(Q)
    if Q isa StaticArray
        return bottomrow(Q)
    elseif Q isa Adjoint  # CuArray takes this path
        x′ = _similar(Q, size(Q, 1), 1)
        fill!(x′, false)
        x′[end, 1] = 1
        # return (Q' * x′)'  # `vec(x::Adjoint{_, <:CuArray})` does not work
        return conj(reshape(Q' * x′, 1, :))
    else
        x = _similar(Q, 1, size(Q, 1))
        fill!(x, false)
        # x .= (x .* false .+ CartesianIndices(x)) .== Ref(size(x))  # InvalidIRError
        x[1, end] = 1
        rmul!(x, Q)
        return x
    end
end

function tangent(L, Q)
    tJ = rawtangent(Q)
    if _det(Q) * _det(popbottomright(L)) < 0
        tJ *= -1
    end
    return tJ
end

function current_tangent(cache::ContinuationCache,
                         opts::ContinuationOptions)
    prob_cache = cache.prob_cache
    u = cache.u
    H = cache.H
    J = cache.J
    Q = cache.Q

    H, J = residual_jacobian!(H, J, u, prob_cache)
    A = vcat(J, _zeros(J, 1, size(J, 2)))  # TODO: improve
    L, Q = _lq!(A)
    return tangent(L, Q)
end

function corrector_step!(H::HType,
                         J::JType,
                         Q::QType,
                         v::vType,
                         prob_cache) where {
                             HType <: AbstractVector,
                             JType <: AbstractMatrix,
                             QType <: AbstractMatrix,
                             vType <: AbstractVector,
                         }
    H, J = residual_jacobian!(H, J, v, prob_cache)
    A = vcat(J, _zeros(J, 1, size(J, 2)))  # TODO: improve
    L, Q = _lq!(A)

    # Following block is a workaround to make test_predator_prey.jl
    # etc. work.  It was working accidentally since an equivalent of
    # `popbottomright(L) \ H` below was calling non-StaticArrays
    # method.  (Although this may not be specific to StaticArrays; it
    # seems test_vs_svector.jl need this, too.)
    if isalmostzero(H, eps(eltype(H))) && abs(det(L)) < eps(eltype(L))
        # Do nothing when there is no need for correction (and
        # trying to do so throws).
        w = v
        dv = zero(v)
        return (w :: vType,
                dv,
                H :: HType,
                L,
                Q,
                J :: JType)
    end

    y = vcat(popbottomright(L) \ H, _zeros(J, 1))
    dv = Q' * y
    w = v - dv
    return (w :: vType,
            dv,
            H :: HType,
            L,
            Q,
            J :: JType)
end

function predictor_corrector_step!(cache::ContinuationCache,
                                   opts::ContinuationOptions)
    predictor_corrector_step!(cache, opts,
                              cache.u,
                              current_tangent(cache, opts),
                              cache.h,
                              cache.direction)
end

function predictor_corrector_step!(cache::ContinuationCache,
                                   opts::ContinuationOptions,
                                   u, tJ, h, direction)
    prob_cache = cache.prob_cache
    H = cache.H
    J = cache.J
    Q = cache.Q
    rtol = opts.rtol
    atol = opts.atol

    cache.corrector_success = false
    cache.adaptation_success = false
    cache.simple_bifurcation = false

    local v, L
    for _ in 1:opts.max_adaptations
        # predictor
        v = u .+ direction * h .* tJ

        # corrector
        v, dv, H, L, Q, J = corrector_step!(H, J, Q, v, prob_cache)
        n1 = norm(dv)
        tJv = tangent(L, Q)
        angle = acos(min(abs(tJ ⋅ tJv), 1))  # TODO: should I use min?
        @debug "maximum(abs, H) = $(maximum(abs, H))"

        if all(abs.(H) .< min(2 * eps(eltype(H)), atol))
            @debug "corrector_skipped: The first correction is too close to the zero point."
            # The first correction is too close to the zero point.
            # Then the fractions for step adaptation would not be
            # possible to reliably calculated so let's skip them.
            # Also, to get out of this "tricky" region faster, let's
            # increase `h` slightly.
            h = h * 2  # TODO: don't hard code
            @goto corrector_skipped
        end

        v, dv, H, L, Q, J = corrector_step!(H, J, Q, v, prob_cache)
        n2 = norm(dv)
        @debug "maximum(abs, H) = $(maximum(abs, H))"

        # step adaptation
        f_contraction =
            sqrt(n2 / n1 / opts.nominal_contraction)   # √ κ(u,h) / κ̃
        f_distance = sqrt(n1 / opts.nominal_distance)  # √ δ(u,h) / δ̃
        f_angle = angle / opts.nominal_angle_rad       #   α(u,h) / α̃
        f0 = max(
            zero_if_nan(f_contraction),
            zero_if_nan(f_distance),
            zero_if_nan(f_angle),
        )
        f = max(min(f0, 2), 1/2)

        @debug(
            "Step adaptation",
            # Print some useful statistics:
            h,
            n2 / n1,
            n1,
            angle,
            f_contraction,
            f_distance,
            f_angle,
            f0
        )

        h = h / f
        if h < opts.h_min
            break
        elseif isalmostzero(H, atol)
            @debug "corrector_skipped: isalmostzero(H, atol)"
            @goto corrector_skipped
        elseif f0 <= 2
            @debug "adaptation_success"
            cache.adaptation_success = true
            @goto adaptation_success
        end
    end
    # TODO: redesign
    error("""Failed to adapt steplength h.
        h = $h
        h < h_min = $(h < opts.h_min)
        max_adaptations = $(opts.max_adaptations)
        """)
    # step adaptation failed
    return

    @label adaptation_success
    # corrector (again)
    for _ in 3:opts.max_corrector_steps
        if isalmostzero(H, atol)
            cache.corrector_success = true
            @goto corrector_success
        end
        v, _, H, L, Q, J = corrector_step!(H, J, Q, v, prob_cache)
    end
    error("Failed in corrector loop.")  # TODO: redesign
    # corrector failed
    return

    # TODO: consider if I can avoid this special case
    @label corrector_skipped
    cache.adaptation_success = true
    cache.corrector_success = true

    @label corrector_success
    cache.u = v
    cache.J = J
    cache.h = h

    tJv = tangent(L, Q)
    if tJ ⋅ tJv < 0
        @debug "switching direction: $(cache.direction) -> $(-cache.direction)"
        cache.direction *= -1
        cache.simple_bifurcation = true
    end
    return v, h
end
# TODO: make it possible to compile away `@debug`s?

function nearest_root!(cache::ContinuationCache,
                       opts::ContinuationOptions,
                       v = cache.u)
    prob_cache = cache.prob_cache
    H = cache.H
    J = cache.J
    Q = cache.Q
    rtol = opts.rtol
    atol = opts.atol

    cache.corrector_success = false
    for _ in 1:opts.max_corrector_steps
        v, _, H = corrector_step!(H, J, Q, v, prob_cache)
        if isalmostzero(H, atol)
            @goto corrector_success
        end
    end
    error("Failed in corrector loop.")  # TODO: redesign

    @label corrector_success
    cache.corrector_success = true
    return v
end


text:
# Sort ordering.

mutable struct SortNode <: AbstractSQLNode
    over::Union{SQLNode, Nothing}
    value::ValueOrder
    nulls::Union{NullsOrder, Nothing}

    SortNode(;
             over = nothing,
             value,
             nulls = nothing) =
        new(over, value, nulls)
end

SortNode(value; over = nothing, nulls = nothing) =
    SortNode(over = over, value = value, nulls = nulls)

"""
    Sort(; over = nothing, value, nulls = nothing)
    Sort(value; over = nothing, nulls = nothing)
    Asc(; over = nothing, nulls = nothing)
    Desc(; over = nothing, nulls = nothing)

Sort order indicator.

# Examples

```jldoctest
julia> person = SQLTable(:person, columns = [:person_id, :year_of_birth]);

julia> q = From(person) |>
           Order(Get.year_of_birth |> Desc());

julia> print(render(q))
SELECT
  "person_1"."person_id",
  "person_1"."year_of_birth"
FROM "person" AS "person_1"
ORDER BY "person_1"."year_of_birth" DESC
```
"""
Sort(args...; kws...) =
    SortNode(args...; kws...) |> SQLNode

"""
    Asc(; over = nothing, nulls = nothing)

Ascending order indicator.
"""
Asc(; kws...) =
    Sort(VALUE_ORDER.ASC; kws...)

"""
    Desc(; over = nothing, nulls = nothing)

Descending order indicator.
"""
Desc(; kws...) =
    Sort(VALUE_ORDER.DESC; kws...)

dissect(scr::Symbol, ::typeof(Sort), pats::Vector{Any}) =
    dissect(scr, SortNode, pats)

function PrettyPrinting.quoteof(n::SortNode, ctx::QuoteContext)
    if n.value == VALUE_ORDER.ASC
        ex = Expr(:call, nameof(Asc))
    elseif n.value == VALUE_ORDER.DESC
        ex = Expr(:call, nameof(Desc))
    else
        ex = Expr(:call, nameof(Sort), QuoteNode(Symbol(n.value)))
    end
    if n.nulls !== nothing
        push!(ex.args, Expr(:kw, :nulls, QuoteNode(Symbol(n.nulls))))
    end
    if n.over !== nothing
        ex = Expr(:call, :|>, quoteof(n.over, ctx), ex)
    end
    ex
end

label(n::SortNode) =
    label(n.over)

rebase(n::SortNode, n′) =
    SortNode(over = rebase(n.over, n′), value = n.value, nulls = n.nulls)