text:
#=
UTF8Str type

Copyright 2017-2018 Gandalf Software, Inc., Scott P. Jones,
and other contributors to the Julia language

Licensed under MIT License, see LICENSE.md
Based in part on code for UTF8String that used to be in Julia
=#

# UTF-8 support functions

@inline checkcont(pnt) = is_valid_continuation(get_codeunit(pnt))

# Get rest of character ch from 2-byte UTF-8 sequence at pnt - 1
@inline get_utf8_2byte(pnt, ch) =
    (((ch & 0x1f)%UInt16 << 6) | (get_codeunit(pnt) & 0x3f))

# Get rest of character ch from 3-byte UTF-8 sequence at pnt - 2
@inline get_utf8_3byte(pnt, ch) =
    (((ch & 0xf)%UInt16 << 12)
     | ((get_codeunit(pnt - 1)%UInt16 & 0x3f) << 6)
     | (get_codeunit(pnt) & 0x3f))

# Get rest of character ch from 4-byte UTF-8 sequence at pnt - 3
@inline get_utf8_4byte(pnt, ch) =
    (((ch & 0x7)%UInt32 << 18)
     | ((get_codeunit(pnt - 2)%UInt32 & 0x3f) << 12)
     | ((get_codeunit(pnt - 1)%UInt32 & 0x3f) << 6)
     | (get_codeunit(pnt) & 0x3f))

@inline get_utf8_2(ch) =
    0xc0 | (ch >>> 6)%UInt8, 0x80 | (ch & 0x3f)%UInt8
@inline get_utf8_3(ch) =
    (0xe0 | (ch >>> 12)%UInt8, 0x80 | ((ch >>> 6) & 0x3f)%UInt8, 0x80 | (ch & 0x3f)%UInt8)
@inline get_utf8_4(ch) =
    (0xf0 | (ch >>>  18)%UInt8, 0x80 | ((ch >>> 12) & 0x3f)%UInt8,
     0x80 | ((ch >>>  6) & 0x3f)%UInt8, 0x80 | (ch & 0x3f)%UInt8)

# Output a character as a 2-byte UTF-8 sequence
@inline function output_utf8_2byte!(pnt, ch)
    b1, b2 = get_utf8_2(ch)
    set_codeunit!(pnt,     b1)
    set_codeunit!(pnt + 1, b2)
    pnt + 2
end

# Output a character as a 3-byte UTF-8 sequence
@inline function output_utf8_3byte!(pnt, ch)
    b1, b2, b3 = get_utf8_3(ch)
    set_codeunit!(pnt,     b1)
    set_codeunit!(pnt + 1, b2)
    set_codeunit!(pnt + 2, b3)
    pnt + 3
end

# Output a character as a 4-byte UTF-8 sequence
@inline function output_utf8_4byte!(pnt, ch)
    b1, b2, b3, b4 = get_utf8_4(ch)
    set_codeunit!(pnt,     b1)
    set_codeunit!(pnt + 1, b2)
    set_codeunit!(pnt + 2, b3)
    set_codeunit!(pnt + 3, b4)
    pnt + 4
end

@inline eq_bytes(pnt, b1)         = get_codeunit(pnt) == b1
@inline eq_bytes(pnt, b1, b2)     = get_codeunit(pnt+1) == b2 && eq_bytes(pnt, b1)
@inline eq_bytes(pnt, b1, b2, b3) = get_codeunit(pnt+2) == b3 && eq_bytes(pnt, b1, b2)

## required core functionality ##

utf_trail(c::UInt8) = (0xe5000000 >>> ((c & 0xf0) >> 3)) & 0x3

@inline function _lastindex(::CodeUnitMulti, str::MS_UTF8)
    (len = ncodeunits(str)) > 1 || return len
    @preserve str begin
        pnt = pointer(str) + len - 1
        len - (checkcont(pnt) ? (checkcont(pnt - 1) ? checkcont(pnt - 2) + 2 : 1) : 0)
    end
end

#=
_cont(byt, n) = (byt >>> n)%UInt8 != 0x80
cnt += 8
else
val &= 0xe0e0_e0e0_e0e0_e0e0
# Count ones that are not valid continuation
cnt += _cont(val, 0x0) + _cont(val, 0x8) + _cont(val, 0x10) + _cont(val, 0x18) +
_cont(val, 0x20) + _cont(val, 0x28) + _cont(val, 0x30) + _cont(val, 0x38)
end

xor 80 then << 1 then |
00 -> 10 -> 1
01 -> 11 -> 1
10 -> 00 -> 0
11 -> 01 -> 1
=#

const hi_mask = 0x8080_8080_8080_8080

@inline _count_cont(v) = (v = xor(v, hi_mask); count_ones(xor(((v << 1) | v), hi_mask) & hi_mask))
@inline msk_lead(v) = (v = xor(v, hi_mask); xor(xor(((v << 1) | v), hi_mask) & hi_mask, hi_mask))

@inline function _align_len_utf8(pnt, cnt, v)
    len = 0
    fin = pnt + cnt
    v = msk_lead(v)
    while (pnt += CHUNKSZ) < fin
        len += count_ones(v)
        v = msk_lead(unsafe_load(pnt))
    end
    len + count_ones(cnt & CHUNKMSK == 0 ? v : (v & _mask_bytes(cnt)))
end

_length_al(::CodeUnitMulti, ::Type{UTF8CSE}, beg::Ptr{UInt8}, cnt::Int) =
    (pnt = reinterpret(Ptr{UInt64}, beg); _align_len_utf8(pnt, cnt, unsafe_load(pnt)))

function _length(::CodeUnitMulti, ::Type{UTF8CSE}, beg::Ptr{UInt8}, cnt::Int)
    align = reinterpret(UInt, beg)
    pnt = reinterpret(Ptr{UInt64}, align & ~CHUNKMSK)
    v = unsafe_load(pnt)
    if (align &= CHUNKMSK) != 0
        msk = _mask_bytes(align)
        v = (v & ~msk) | (msk & hi_mask)
        cnt += align
    end
    _align_len_utf8(pnt, cnt, v)
end

@inline function _check_mask_al(pnt, cnt, msk, v)
    fin = pnt + cnt
    while (pnt += CHUNKSZ) < fin
        (v & msk) == 0 || return false
        v = unsafe_load(pnt)
    end
    (cnt & CHUNKMSK == 0 ? v : (v & _mask_bytes(cnt))) & msk == 0
end
@inline _check_mask_al(pnt, cnt, msk) = _check_mask_al(pnt, cnt, msk, unsafe_load(pnt))

@inline function _check_mask_ul(beg, cnt, msk)
    align = reinterpret(UInt, beg)
    pnt = reinterpret(Ptr{UInt64}, align & ~CHUNKMSK)
    v = unsafe_load(pnt)
    if (align &= CHUNKMSK) != 0
        v &= ~_mask_bytes(align)
        cnt += align
    end
    _check_mask_al(pnt, cnt, msk, v)
end

_ascii_mask(::Type{UInt8})  = hi_mask
_ascii_mask(::Type{UInt16}) = 0xff80_ff80_ff80_ff80
_ascii_mask(::Type{UInt32}) = 0xffffff80_ffffff80

_latin_mask(::Type{UInt16}) = 0xff00_ff00_ff00_ff00
_latin_mask(::Type{UInt32}) = 0xffffff00_ffffff00

const _bmp_mask_32   = 0xffff0000_ffff0000

is_ascii(str::SubString{<:Str{C}}) where {C<:Union{UTF8CSE,LatinCSE,Binary_CSEs,UTF16CSE,UCS2CSE,
                                                   Text2CSE,Text4CSE,UTF32CSE}} =
    (cnt = sizeof(str)) == 0 ? true :
    @preserve str _check_mask_ul(pointer(str), cnt, _ascii_mask(codeunit(C)))

is_ascii(vec::Vector{T}) where {T<:CodeUnitTypes} =
    (cnt = sizeof(vec)) == 0 ? true :
    @preserve str _check_mask_ul(pointer(vec), cnt, _ascii_mask(T))

is_ascii(str::Str{C}) where {C<:Union{UTF8CSE,LatinCSE,Binary_CSEs,UTF16CSE,UCS2CSE,
                                      Text2CSE,Text4CSE,UTF32CSE}} =
    (cnt = sizeof(str)) == 0 ? true :
    @preserve str _check_mask_al(reinterpret(Ptr{UInt64}, pointer(str)), cnt,
                                 _ascii_mask(codeunit(C)))

# Todo! Here you need to see that 0b11yyyyxx at least 1 y must be set,
# which indicates a non-Latin1 character
_all_latin(val) = ((val & (val<<1) & (val<<2 | (val<<3) | (val<<4) | (val<<5))) & hi_mask) == 0

@inline function _check_latin_utf8_al(pnt, cnt, v)
    fin = pnt + cnt
    while (pnt += CHUNKSZ) < fin
        _all_latin(v) || return false
        v = unsafe_load(pnt)
    end
    _all_latin(cnt & CHUNKMSK == 0 ? v : (v & _mask_bytes(cnt)))
end
@inline _check_latin_utf8_al(pnt, cnt) = _check_latin_utf8_al(pnt, cnt, unsafe_load(pnt))

@inline function _check_latin_utf8_ul(beg, cnt)
    align = reinterpret(UInt, beg)
    pnt = reinterpret(Ptr{UInt64}, align & ~CHUNKMSK)
    v = unsafe_load(pnt)
    if (align &= CHUNKMSK) != 0
        v &= ~_mask_bytes(align)
        cnt += align
    end
    _check_latin_utf8_al(pnt, cnt, v)
end

is_latin(str::Str{UTF8CSE}) =
    (siz = sizeof(str)) == 0 ? true :
    @preserve str _check_latin_utf8_al(reinterpret(Ptr{UInt64}, pointer(str)), siz)

is_latin(str::SubString{<:Str{UTF8CSE}}) =
    (cnt = sizeof(str)) == 0 ? true : @preserve str _check_latin_utf8_ul(pointer(str), cnt)

is_latin(vec::Vector{T}) where {T<:Union{UInt16,UInt32}} =
    (cnt = sizeof(vec)) == 0 ? true :
    @preserve str _check_mask_ul(pointer(vec), cnt, _latin_mask(T))

is_latin(str::SubString{<:Str{C}}) where {C<:Union{Word_CSEs,Quad_CSEs}} =
    (cnt = sizeof(str)) == 0 ? true :
    @preserve str _check_mask_ul(pointer(str), cnt, _latin_mask(codeunit(C)))

is_latin(str::Str{C}) where {C<:Union{Word_CSEs,Quad_CSEs}} =
    (cnt = sizeof(str)) == 0 ? true :
    @preserve str _check_mask_al(pointer(str), cnt, _latin_mask(codeunit(C)))

# All 4 top bits must be 1 (i.e. 0xfx) for this to be non-BMP
_all_bmp(val) = ((val | (val<<1) | (val<<2) | (val<<3)) & hi_mask) == 0

@inline function _check_bmp_utf8_al(pnt, cnt, v)
    fin = pnt + cnt
    while (pnt += CHUNKSZ) < fin
        _all_bmp(v) || return false
        v = unsafe_load(pnt)
    end
    _all_bmp(cnt & CHUNKMSK == 0 ? v : (v & _mask_bytes(cnt)))
end
@inline _check_bmp_utf8_al(pnt, cnt) = _check_bmp_utf8_al(pnt, cnt, unsafe_load(pnt))

@inline function _check_bmp_utf8_ul(beg, cnt)
    align = reinterpret(UInt, beg)
    pnt = reinterpret(Ptr{UInt64}, align & ~CHUNKMSK)
    v = unsafe_load(pnt)
    if (align &= CHUNKMSK) != 0
        v &= ~_mask_bytes(align)
        cnt += align
    end
    _check_bmp_utf8_al(pnt, cnt, v)
end

is_bmp(str::Str{UTF8CSE}) =
    (cnt = sizeof(str)) == 0 ? true :
    @preserve str _check_bmp_utf8_al(reinterpret(Ptr{UInt64}, pointer(str)), cnt)

is_bmp(str::SubString{<:Str{UTF8CSE}}) =
    (cnt = sizeof(str)) == 0 ? true : @preserve str _check_bmp_utf8_ul(pointer(str), cnt)

is_bmp(str::SubString{<:Str{<:Union{Text4CSE,UTF32CSE}}}) =
    (cnt = sizeof(str)) == 0 ? true : @preserve str _check_mask_ul(pointer(str), cnt, _bmp_mask_32)

is_bmp(str::Str{<:Union{Text4CSE,UTF32CSE}}) =
    (cnt = sizeof(str)) == 0 ? true : @preserve str _check_mask_al(pointer(str), cnt, _bmp_mask_32)

is_unicode(str::MS_UTF8) = true

is_unicode(str::String) = @preserve str _check_utf8_al(ncodeunits(str), pointer(str)) >= 0
is_unicode(str::SubString{String}) = @preserve str _check_utf8(ncodeunits(str), pointer(str)) >= 0

"""
Return index of first invalid codeunit (negative),
0 if all ASCII, or index of first non-ASCII codeunit
"""
check_utf8(str) = @preserve str _check_utf8(ncodeunits(str), pointer(str))

@inline _check_utf8(len, pnt)    = _check_utf8(pnt, pnt, pnt + len)
@inline _check_utf8_al(len, pnt) = _check_utf8(pnt, pnt, pnt + len)

function _check_utf8_rest(pnt, fin, ch)
    while true
        # Check UTF-8 encoding
        if ch < 0xc2
            # Found continuation character or invalid 0xc0/0xc1
            break
        elseif ch < 0xe0
            # 2-byte UTF-8 sequence (i.e. characters 0x80-0x7ff)
            pnt < fin || break
            is_valid_continuation(get_codeunit(pnt)) || break
            pnt += 1
        elseif ch < 0xf0
            # 3-byte UTF-8 sequence (i.e. characters 0x800-0xffff)
            pnt + 1 < fin || break
            b2 = get_codeunit(pnt)     ; is_valid_continuation(b2) || break
            b3 = get_codeunit(pnt + 1) ; is_valid_continuation(b3) || break
            wrd = ((ch & 0x0f)%UInt32 << 12) | ((b2 & 0x3f)%UInt32 << 6) | (b3 & 0x3f)
            # check for surrogate pairs, make sure correct
            (wrd < 0x0800 || is_surrogate_codeunit(wrd)) && break
            pnt += 2
        elseif ch < 0xf5
            # 4-byte UTF-8 sequence (i.e. characters > 0xffff)
            pnt + 2 < fin || break
            b2 = get_codeunit(pnt)     ; is_valid_continuation(b2) || break
            b3 = get_codeunit(pnt + 1) ; is_valid_continuation(b3) || break
            b4 = get_codeunit(pnt + 2) ; is_valid_continuation(b4) || break
            (((ch & 0x07)%UInt32 << 18) | ((b2 & 0x3f)%UInt32 << 12) |
             ((b3 & 0x3f)%UInt32 << 6) | (b4 & 0x3f)) - 0x10000 < 0x100000 ||
             break
            pnt += 3
        else
            break
        end
        # Skip ascii characters as fast as possible
        while true
            pnt < fin || return C_NULL
            ch = get_codeunit(pnt)
            pnt += 1
            ch < 0x7f || break
        end
    end
    pnt
end

function _check_utf8(beg, pnt, fin)
    while pnt < fin
        ch = get_codeunit(pnt)
        pnt += 1
        if ch > 0x7f
            nxt = _check_utf8_rest(pnt, fin, ch)
            return nxt == C_NULL ? Int(pnt - beg) : -Int(nxt - beg)
        end
    end
    0
end

function _nextcpfun(::CodeUnitMulti, ::Type{UTF8CSE}, pnt)
    ch = get_codeunit(pnt)
    if ch < 0x80
        ch%UInt32, pnt + 1
    elseif ch < 0xe0
        get_utf8_2byte(pnt + 1, ch)%UInt32, pnt + 2
    elseif ch < 0xf0
        get_utf8_3byte(pnt + 2, ch)%UInt32, pnt + 3
    else
        get_utf8_4byte(pnt + 3, ch), pnt + 4
    end
end

# Gets next codepoint
@propagate_inbounds function _next(::CodeUnitMulti, ::Type{T}, str::MS_UTF8,
                                   pos::Int) where {T<:Chr}
    len = ncodeunits(str)
    @boundscheck 0 < pos <= len || boundserr(str, pos)
    @preserve str begin
        pnt = pointer(str) + pos - 1
        ch = get_codeunit(pnt)
        if ch < 0x80
            T(ch), pos + 1
        elseif ch < 0xc0
            index_error(str, pos)
        elseif ch < 0xe0
            T(get_utf8_2byte(pnt + 1, ch)), pos + 2
        elseif ch < 0xf0
            T(get_utf8_3byte(pnt + 2, ch)), pos + 3
        else
            T(get_utf8_4byte(pnt + 3, ch)), pos + 4
        end
    end
end

done(str::Str{UTF8CSE}, i::Int) = i > sizeof(str.data)

length(it::CodePoints{<:AbstractString}, i::Int) = length(it.xs)

@propagate_inbounds function next(it::CodePoints{<:AbstractString}, state)
    ch, state = next(it.xs, state)
    UTF32Chr(ch%UInt32), state
end

## overload methods for efficiency ##

@inline _isvalid_char_pos(::CodeUnitMulti, ::Type{UTF8CSE}, str, pos::Integer) =
    !is_valid_continuation(get_codeunit(pointer(str) + pos - 1))

function _thisind(::CodeUnitMulti, str::MaybeSub{T}, len, pnt,
                  pos::Integer) where {T<:Union{Str{<:UTF8CSE},String}}
    pnt += pos - 1
    Int(pos) - (checkcont(pnt) ? (checkcont(pnt - 1) ? checkcont(pnt - 2) + 2 : 1) : 0)
end

@propagate_inbounds function _nextind(::CodeUnitMulti, str::MS_UTF8, pos::Integer)
    pos == 0 && return 1
    numcu = ncodeunits(str)
    @boundscheck 1 <= pos <= numcu || boundserr(str, pos)
    @preserve str begin
        pnt = pointer(str) + pos - 1
        cu = get_codeunit(pnt)
        pos + (cu < 0x80 ? 1
               : (cu < 0xc0
                  ? (pos == numcu ? 1
                     : (checkcont(pnt + 1) ? (2 + (pos < numcu - 1 && checkcont(pnt + 2))) : 1))
                  : ifelse(cu < 0xe0, 2, ifelse(cu < 0xf0, 3, 4))))
    end
end

@propagate_inbounds function _prevind(::CodeUnitMulti, str::MS_UTF8, pos::Integer)
    (pos -= 1) == 0 && return 0
    numcu = ncodeunits(str)
    @preserve str begin
        pnt = pointer(str) + pos
        pos == numcu &&
            return pos - (checkcont(pnt-1) ? (checkcont(pnt-2) ? checkcont(pnt-3) + 2 : 1) : 0)
        @boundscheck 0 < pos < numcu || boundserr(str, pos+1)
        if checkcont(pnt)
            pos - (checkcont(pnt - 1) ? checkcont(pnt - 2) + 1 : 0)
        elseif pos != 1
            pos - (checkcont(pnt - 1) ? (checkcont(pnt - 2) ? checkcont(pnt - 3) + 2 : 1) : 0)
        else
            1
        end
    end
end

@propagate_inbounds function _nextind(::CodeUnitMulti, str::MS_UTF8, pos::Int, nchar::Int)
    nchar < 0 && ncharerr(nchar)
    siz = ncodeunits(str)
    @boundscheck 0 <= pos <= siz || boundserr(str, pos)
    siz == 0 && return Int(nchar != 0)
    @preserve str begin
        beg = pointer(str)
        # Get starting position
        if pos == 0
            nchar <= 1 && return nchar
            pnt = beg + utf_trail(get_codeunit(beg)) + 1
            nchar -= 1
        else
            pnt = beg + pos - 1
            nchar == 0 && (checkcont(pnt) ? index_error(str, pos) : return pos)
            cu = get_codeunit(pnt)
            if !is_valid_continuation(cu)
                pnt += utf_trail(cu) + 1
            elseif pos == siz
                # At end of string already
                return siz + 1
            else
                pnt += checkcont(pnt + 1) ? checkcont(pnt + 2) + 2 : 1
            end
        end
        fin = beg + siz
        # pnt should now point to a valid start of a character
        # This could be sped up, by looking at chunks, and if all ASCII (common case),
        # simply move forward 8
        while (nchar -= 1) > 0 && pnt < fin
            pnt += utf_trail(get_codeunit(pnt)) + 1
        end
        Int(pnt - beg + 1)
    end
end

@propagate_inbounds function _prevind(::CodeUnitMulti, str::MS_UTF8, pos::Int, nchar::Int)
    nchar < 0 && ncharerr(nchar)
    numcu = ncodeunits(str)
    @boundscheck 0 < pos <= numcu+1 || boundserr(str, pos)
    numcu == 0 && return Int(nchar == 0)
    @preserve str begin
        beg = pointer(str)
        if pos > numcu
            (nchar -= 1) < 0 && return pos
            pos = numcu
        elseif nchar == 0
            checkcont(beg + pos - 1) ? index_error(str, pos) : return pos
        end
        pnt = beg + pos
        # This could be sped up, by looking at chunks, and if all ASCII (common case),
        # simply move back 8
        while nchar >= 0
            (pnt -= 1) < beg && return 0
            nchar -= !checkcont(pnt)
        end
        Int(pnt - beg + 1)
    end
end

@propagate_inbounds function getindex(str::MS_UTF8, rng::UnitRange{Int})
    isempty(rng) && return SubString(empty_utf8, 1, 0)
    beg = first(rng)
    len = ncodeunits(str)
    @boundscheck 1 <= beg <= len || boundserr(str, beg)
    @preserve str begin
        pnt = pointer(str)
        ch = get_codeunit(pnt, beg)
        is_valid_continuation(ch) && unierror(UTF_ERR_INVALID_INDEX, beg, ch)
        lst = last(rng)
        @boundscheck lst > len && boundserr(str, lst)
        if lst != len
            ch = get_codeunit(pnt, lst)
            is_valid_continuation(ch) && unierror(UTF_ERR_INVALID_INDEX, lst, ch)
        end
    end
    SubString(str, beg, lst)
end

const _ByteStr = Union{Str{ASCIICSE}, SubString{<:Str{ASCIICSE}},
                       Str{UTF8CSE},  SubString{<:Str{UTF8CSE}},
                       String}

string(c::_ByteStr...) = length(c) == 1 ? c[1]::UTF8Str : UTF8Str(_string(c))
    # ^^ at least one must be UTF-8 or the ASCII-only method would get called

function _reverse(::CodeUnitMulti, ::Type{UTF8CSE}, len, pnt::Ptr{T}) where {T<:CodeUnitTypes}
    buf, beg = _allocate(T, len)
    out = beg + len
    while out >= beg
        ch = get_codeunit(pnt)
        if ch > 0xdf
            if ch < 0xf0
                set_codeunit!(out -= 2, get_codeunit(pnt += 1))
                set_codeunit!(out + 1,  get_codeunit(pnt += 1))
            else
                set_codeunit!(out -= 3, get_codeunit(pnt += 1))
                set_codeunit!(out + 1,  get_codeunit(pnt += 1))
                set_codeunit!(out + 2,  get_codeunit(pnt += 1))
            end
        elseif ch > 0x7f
            set_codeunit!(out -= 1, get_codeunit(pnt += 1))
        end
        set_codeunit!(out -= 1, ch)
        pnt += 1
    end
    Str(UTF8CSE, buf)
end

@inline get_ch(dat, pos, off) = (get_codeunit(dat, pos + off) & 0x3f)%UInt32

## transcoding to UTF-8 ##

@propagate_inbounds function _transcode_utf8(pnt::Ptr{UInt8}, len)
    buf, out = _allocate(UInt8, len)
    fin = out + len
    @inbounds while out < fin
        ch = get_codeunit(pnt)
        # Handle ASCII characters
        if ch <= 0x7f
            set_codeunit!(out, ch)
            out += 1
        # Handle overlong < 0x100
        elseif ch < 0xc2
            ch = ((ch & 3) << 6) | (get_codeunit(pnt += 1) & 0x3f)
            set_codeunit!(out, ch)
            out += 1
        # Handle 0x100-0x7ff
        elseif ch < 0xe0
            set_codeunit!(out, ch)
            set_codeunit!(out + 1, get_codeunit(pnt += 1))
            out += 2
        elseif ch != 0xed
            set_codeunit!(out, ch)
            set_codeunit!(out + 1, get_codeunit(pnt += 1))
            set_codeunit!(out + 2, get_codeunit(pnt += 1))
            out += 3
            # Copy 4-byte encoded value
            ch >= 0xf0 && (set_codeunit!(out, get_codeunit(pnt += 1)) ; out += 1)
        # Handle surrogate pairs
        else
            ch = get_codeunit(pnt += 1)
            if ch < 0xa0 # not surrogate pairs
                set_codeunit!(out, 0xed)
                set_codeunit!(out + 1, ch)
                set_codeunit!(out + 2, get_codeunit(pnt += 1))
                out += 3
            else
                # Pick up surrogate pairs (CESU-8 format)
                ch32 = (((ch & 0x3f)%UInt32 << 16) | (get_ch(pnt + 1) << 10)) +
                    (get_ch(pnt + 3) << 6 | get_ch(pnt + 4)) - 0x01f0c00
                pnt += 4
                out = output_utf8_4byte!(out, ch32)
            end
        end
        pnt += 1
    end
    buf
end
_transcode_utf8(dat::Vector{UInt8}, len) = _transcode_utf8(pointer(dat), len)

# Single character conversion
function convert(::Type{<:Str{UTF8CSE}}, ch::Unsigned)
    if ch <= 0x7f
        buf = _allocate(1)
        @inbounds buf[1] = ch%UInt8
    elseif ch <= 0x7ff
        buf = _allocate(2)
        @inbounds buf[1], buf[2] = get_utf8_2(ch)
    elseif ch - 0xd800 < 0x800
        unierror(UTF_ERR_INVALID, 0, ch)
    elseif ch <= 0xffff
        buf = _allocate(3)
        @inbounds buf[1], buf[2], buf[3] = get_utf8_3(ch)
    elseif ch <= 0x10ffff
        buf = _allocate(4)
        @inbounds buf[1], buf[2], buf[3], buf[4] = get_utf8_4(ch)
    else
        unierror(UTF_ERR_INVALID, 0, ch)
    end
    Str(UTF8CSE, buf)
end

@static if !V6_COMPAT
# Note: this will have to change back to s.endof for v0.6!
convert(::Type{SubString{<:Str{UTF8CSE}}}, s::SubString{<:Str{ASCIICSE}}) =
    SubString(Str(UTF8CSE, s.string), s.offset + 1, s.offset + s.ncodeunits)
end

function convert(::Type{<:Str{UTF8CSE}}, dat::Vector{UInt8})
    # handle zero length string quickly
    isempty(dat) && return empty_utf8
    # get number of bytes to allocate
    len, flags, num4byte, num3byte, num2byte, latinbyte =
        @preserve dat fast_check_string(pointer(dat), length(dat))
    #=
    # Copy, but eliminate over-long encodings and surrogate pairs
    if flags & (UTF_LONG | UTF_SURROGATE) == 0
        siz = sizeof(dat)
        buf = _allocate(siz)
        unsafe_copyto!(pointer(buf), pointer(dat), siz)
        Str(UTF8CSE, buf)
    else
    =#
    Str(UTF8CSE, _transcode_utf8(dat, len + latinbyte + num2byte + num3byte*2 + num4byte*3))
    #end
end

function convert(::Type{<:Str{UTF8CSE}}, str::AbstractString)
    # handle zero length string quickly
    isempty(str) && return empty_utf8
    # get number of bytes to allocate
    len, flags, num4byte, num3byte, num2byte, latinbyte = unsafe_check_string(str)
    if flags == 0
        # Speed this up if only ASCII, no overlong encodings
        buf, pnt = _allocate(UInt8, len)
        for ch in str
            set_codeunit!(pnt, ch%UInt8)
            pnt += 1
        end
    else
        # Copy, but eliminate over-long encodings and surrogate pairs
        buf, out = _allocate(UInt8, len + latinbyte + num2byte + num3byte*2 + num4byte*3)
        for ch in str
            out = _encode_char_utf8(out, codepoint(ch))
        end
    end
    Str(UTF8CSE, buf)
end

function convert(::Type{<:Str{UTF8CSE}}, str::T) where {T<:MS_ByteStr}
    # handle zero length string quickly
    (siz = sizeof(str)) == 0 && return empty_utf8
    @preserve str begin
        pnt = pointer(str)
        # get number of bytes to allocate
        len, flags, num4byte, num3byte, num2byte, latinbyte = fast_check_string(pnt, siz)
        # Copy, but eliminate over-long encodings and surrogate pairs
        # Speed this up if no surrogates, long encodings
        Str(UTF8CSE,
            (flags & (UTF_LONG | UTF_SURROGATE) == 0
             ? _copysub(str)
             : _transcode_utf8(pnt, len + latinbyte + num2byte + num3byte*2 + num4byte*3)))
    end
end

function convert(::Type{<:Str{UTF8CSE}},
                 str::MaybeSub{T}) where {C<:Union{Text2CSE,Text4CSE},T<:Str{C}}
    # handle zero length string quickly
    (len = ncodeunits(str)) == 0 && return empty_utf8
    @preserve str begin
        pnt = pointer(str)
        # get number of bytes to allocate
        len, flags, num4byte, num3byte, num2byte, latinbyte = fast_check_string(pnt, len)
        # Copy, but eliminate over-long encodings and surrogate pairs
        # Speed this up if no surrogates, long encodings
        Str(UTF8CSE,
            flags == 0
             ? _cvtsize(UInt8, pnt, len)
             : _transcode_utf8(pnt, len + latinbyte + num2byte + num3byte*2 + num4byte*3))
    end
end

convert(::Type{<:Str{UTF8CSE}}, s::Str{ASCIICSE}) = Str(UTF8CSE, s.data)
convert(::Type{<:Str{UTF8CSE}}, s::Str{UTF8CSE}) = s

const WideCodeUnit = Union{UInt16, UInt32}

@inline function _encode_char_utf8(pnt, ch::WideCodeUnit)
    # Handle ASCII characters
    if ch <= 0x7f
        set_codeunit!(pnt, ch)
        pnt + 1
    # Handle 0x80-0x7ff
    elseif ch < 0x800
        output_utf8_2byte!(pnt, ch)
    # Handle 0x10000-0x10ffff
    elseif ch > 0xffff # this is only for T == UInt32, should not be generated for UInt16
        output_utf8_4byte!(pnt, ch)
    # Handle 0x800-0xd7ff, 0xe000-0xffff UCS-2 characters
    else
        output_utf8_3byte!(pnt, ch)
    end
end

"""
Converts an already validated vector of `UInt16` or `UInt32` to a `UTF8Str`

Input Arguments:

* `dat` Vector of code units (`UInt16` or `UInt32`)
* `len` length of output in bytes

Returns:

* `UTF8Str`
"""
function _encode_utf8(pnt::Ptr{T}, len) where {T<:WideCodeUnit}
    buf, out = _allocate(UInt8, len)
    fin = out + len
    while out < fin
        out = _encode_char_utf8(out, get_codeunit(pnt))
        pnt += sizeof(T)
    end
    buf
end
_encode_utf8(dat::Vector{<:WideCodeUnit}, len) = _encode_utf8(pointer(dat), len)

function _transcode_utf8(pnt::Ptr{T}, len) where {T<:WideCodeUnit}
    buf, out = _allocate(UInt8, len)
    fin = out + len
    while out < fin
        ch = get_codeunit(pnt)
        out = (is_surrogate_codeunit(ch)
               ? output_utf8_4byte!(out, get_supplementary(ch, get_codeunit(pnt += sizeof(T))))
               : _encode_char_utf8(out, ch))
        pnt += sizeof(T)
    end
    buf
end
_transcode_utf8(dat::Vector{<:WideCodeUnit}, len) = _transcode_utf8(pointer(dat), len)

function convert(::Type{<:Str{UTF8CSE}}, str::MS_UTF16)
    # handle zero length string quickly
    (siz = ncodeunits(str)) == 0 && return empty_utf8
    @preserve str begin
        pnt = pointer(str)
        len, flags, num4byte, num3byte, num2byte, latinbyte = count_chars(UTF16Str, pnt, siz)
        if flags == 0
            Str(UTF8CSE, _cvtsize(UInt8, pnt, len))
        elseif num4byte == 0
            Str(UTF8CSE, _encode_utf8(pnt, len + latinbyte + num2byte + num3byte*2))
        else
            Str(UTF8CSE, _transcode_utf8(pnt, len + latinbyte + num2byte + num3byte*2 + num4byte*3))
        end
    end
end

function convert(::Type{<:Str{UTF8CSE}},
                 str::MaybeSub{T}) where {C<:Union{UCS2_CSEs,UTF32_CSEs},T<:Str{C}}
    # handle zero length string quickly
    (siz = ncodeunits(str)) == 0 && return empty_utf8
    @preserve str begin
        pnt = pointer(str)
        len, flags, num4byte, num3byte, num2byte, latinbyte = count_chars(T, pnt, siz)
        Str(UTF8CSE, (flags == 0
                      ? _cvtsize(UInt8, pnt, len)
                      : _encode_utf8(pnt, len + latinbyte + num2byte + num3byte*2 + num4byte*3)))
    end
end


text:
using GillMatsuno
using Test

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

ext: jl


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

# # Nonlinear tips and tricks

# This example collates some tips and tricks you can use when formulating
# nonlinear programs. It uses the following packages:

using JuMP
import Ipopt
import Test

# ## User-defined functions with vector outputs

# A common situation is to have a user-defined function like the following that
# returns multiple outputs (we define `function_calls` to keep track of how
# many times we call this method):

function_calls = 0
function foo(x, y)
    global function_calls += 1
    common_term = x^2 + y^2
    term_1 = sqrt(1 + common_term)
    term_2 = common_term
    return term_1, term_2
end

# For example, the first term might be used in the objective, and the second
# term might be used in a constraint, and often they share share work that is
# expensive to evaluate.

# This is a problem for JuMP, because it requires user-defined functions to
# return a single number. One option is to define two separate functions, the
# first returning the first argument, and the second returning the second
# argument.

foo_1(x, y) = foo(x, y)[1]
foo_2(x, y) = foo(x, y)[2]

# However, if the common term is expensive to compute, this approach is wasteful
# because it will evaluate the expensive term twice. Let's have a look at how
# many times we evalute `x^2 + y^2` during a solve:

model = Model(Ipopt.Optimizer)
set_silent(model)
@variable(model, x[1:2] >= 0, start = 0.1)
register(model, :foo_1, 2, foo_1; autodiff = true)
register(model, :foo_2, 2, foo_2; autodiff = true)
@NLobjective(model, Max, foo_1(x[1], x[2]))
@NLconstraint(model, foo_2(x[1], x[2]) <= 2)
function_calls = 0
optimize!(model)
Test.@test objective_value(model) ≈ √3 atol=1e-4
Test.@test value.(x) ≈ [1.0, 1.0] atol=1e-4
println("Naive approach: function calls = $(function_calls)")
naive_approach = function_calls  #src

# An alternative approach is to use _memoization_, which uses a cache to store
# the result of function evaluations. We can write a memoization function as
# follows:

"""
    memoize(foo::Function, n_outputs::Int)

Take a function `foo` and return a vector of length `n_outputs`, where each
element is a function that returns the `i`'th output of `foo`.

To avoid duplication of work, cache the most-recent evaluations of `foo`.
Because `foo_i` is auto-differentiated with ForwardDiff, our cache needs to
work when `x` is a `Float64` and a `ForwardDiff.Dual`.
"""
function memoize(foo::Function, n_outputs::Int)
    last_x, last_f = nothing, nothing
    last_dx, last_dfdx = nothing, nothing
    function foo_i(i, x::T...) where {T <: Real}
        if T == Float64
            if x != last_x
                last_x, last_f = x, foo(x...)
            end
            return last_f[i]::T
        else
            if x != last_dx
                last_dx, last_dfdx = x, foo(x...)
            end
            return last_dfdx[i]::T
        end
    end
    return [(x...) -> foo_i(i, x...) for i = 1:n_outputs]
end

# Let's see how it works. First, construct the memoized versions of `foo`:

memoized_foo = memoize(foo, 2)

# Now try evaluating the first element of `memoized_foo`.

function_calls = 0
memoized_foo[1](1.0, 1.0)
Test.@test function_calls == 1  #src
println("function_calls = ", function_calls)

# As expected, this evaluated the function once. However, if we call the
# function again, we hit the cache instead of needing to re-compute `foo` and
# `function_calls` is still `1`!

memoized_foo[1](1.0, 1.0)
Test.@test function_calls == 1  #src
println("function_calls = ", function_calls)

# Now let's see how this works during a real solve:

model = Model(Ipopt.Optimizer)
set_silent(model)
@variable(model, x[1:2] >= 0, start = 0.1)
register(model, :foo_1, 2, memoized_foo[1]; autodiff = true)
register(model, :foo_2, 2, memoized_foo[2]; autodiff = true)
@NLobjective(model, Max, foo_1(x[1], x[2]))
@NLconstraint(model, foo_2(x[1], x[2]) <= 2)
function_calls = 0
optimize!(model)
Test.@test objective_value(model) ≈ √3 atol=1e-4
Test.@test value.(x) ≈ [1.0, 1.0] atol=1e-4
println("Memoized approach: function_calls = $(function_calls)")
Test.@test function_calls <= naive_approach / 2 + 1  #src

# Compared to the naive approach, the memoized approach requires half as many
# function evaluations!

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
"""
    OutOfBoundsError(msg)

An attempt to calculate a part of Pascal's triangle that doesn't exist.
That is, an entry (n,k) for which 0 ≤ k ≤ n doen't hold, or a row or column
before the first.

# Examples
```jldoctest
julia> prev(Entry(0,0))
ERROR: OutOfBoundsError: no previous entry
[...]
```
"""
struct OutOfBoundsError <: Exception
    msg::String
end
Base.showerror(io::IO, e::OutOfBoundsError) = print(io, "OutOfBoundsError: ", e.msg)

"""
    NonAdjacentError()

An attempt to add or subtract entries within Pascal's triangle that are not appropriately
arranged. Entries that are adjacent on the same row can be added together, and entries
that are directly above each other can be subtracted. Attempting to apply `+` or `-` in
any other circumstance will result in this error.

# Examples
```jldoctest
julia> Entry(4,3) + Entry(5,1)
ERROR: NonAdjacentError: entries not appropriately arranged
[...]
```

```jldoctest
julia> Entry(5,1) - Entry(4,3)
ERROR: NonAdjacentError: entries not appropriately arranged
[...]
```
"""
struct NonAdjacentError <: Exception end
Base.showerror(io::IO, ::NonAdjacentError) = print(io,
    "NonAdjacentError: entries not appropriately arranged")


text:
using DashBootstrapComponents, DashHtmlComponents

progress = html_div([
    dbc_progress(value = 25, color = "success", className = "mb-3"),
    dbc_progress(value = 50, color = "warning", className = "mb-3"),
    dbc_progress(value = 75, color = "danger", className = "mb-3"),
    dbc_progress(value = 100, color = "info"),
]);

max_stars_repo_name: glsdown/dash-bootstrap-components

max_issues_repo_name: glsdown/dash-bootstrap-components

max_forks_repo_name: glsdown/dash-bootstrap-components


text:
# Copyright (c) 2020: Matthew Wilhelm & Matthew Stuber.
# This work is licensed under the Creative Commons Attribution-NonCommercial-
# ShareAlike 4.0 International License. To view a copy of this license, visit
# http://creativecommons.org/licenses/by-nc-sa/4.0/ or send a letter to Creative
# Commons, PO Box 1866, Mountain View, CA 94042, USA.
#############################################################################
# Dynamic Bounds - pODEs Discrete
# A package for discretize and relax methods for bounding pODEs.
# See https://github.com/PSORLab/DynamicBoundspODEsDiscrete.jl
#############################################################################
# src/DiscretizeRelax/utilities/mul_split.jl
# A simple function used to speed up 1D calculations of matrix multiplication.
#############################################################################

"""
mul_split!

Multiples A*b as a matrix times a vector if nx > 1. Performs scalar
multiplication otherwise.
"""
function mul_split!(Y::Vector{R}, A::Matrix{S}, B::Vector{T}, nx) where {R,S,T}
    if nx == 1
        @inbounds Y[1] = A[1,1]*B[1]
    else
        mul!(Y, A, B)
    end

    return nothing
end

function mul_split!(Y::Matrix{R}, A::Matrix{S}, B::Matrix{T}, nx) where {R,S,T}
    if nx == 1
        @inbounds Y[1,1] = A[1,1]*B[1,1]
    else
        mul!(Y, A, B)
    end

    return nothing
end


text:
using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:4] <= 1)

                     @variable(m, -8.186501529651974 <= q <= 9.622483160840265)

                     add_NL_constraint(m, :(log(1 + exp(-0.9342689023247241 + -0.06643065035242302*log(1 + exp(-0.6967397920955185 + 0.4612305646899797*log(1 + exp(-0.7651411836742423 + 0.736339264438953*$(x[1]) + 0.06323249947984699*$(x[2]) + 0.6326865446578034*$(x[3]) + 0.8935006364653741*$(x[4]))) + 0.5917923558177942*log(1 + exp(0.8959549878207418 + 0.6412415470262633*$(x[1]) + 0.6383904188579423*$(x[2]) + 0.9706157764666923*$(x[3]) + 0.8108132764672078*$(x[4]))))) + 0.13022424784864484*log(1 + exp(-0.8092681495291991 + -0.6577545463201151*log(1 + exp(-0.7651411836742423 + 0.736339264438953*$(x[1]) + 0.06323249947984699*$(x[2]) + 0.6326865446578034*$(x[3]) + 0.8935006364653741*$(x[4]))) + 0.816846389391447*log(1 + exp(0.8959549878207418 + 0.6412415470262633*$(x[1]) + 0.6383904188579423*$(x[2]) + 0.9706157764666923*$(x[3]) + 0.8108132764672078*$(x[4]))))))) + log(1 + exp(0.9502905692637253 + 0.2731932353707758*log(1 + exp(-0.6967397920955185 + 0.4612305646899797*log(1 + exp(-0.7651411836742423 + 0.736339264438953*$(x[1]) + 0.06323249947984699*$(x[2]) + 0.6326865446578034*$(x[3]) + 0.8935006364653741*$(x[4]))) + 0.5917923558177942*log(1 + exp(0.8959549878207418 + 0.6412415470262633*$(x[1]) + 0.6383904188579423*$(x[2]) + 0.9706157764666923*$(x[3]) + 0.8108132764672078*$(x[4]))))) + 0.7525923835832526*log(1 + exp(-0.8092681495291991 + -0.6577545463201151*log(1 + exp(-0.7651411836742423 + 0.736339264438953*$(x[1]) + 0.06323249947984699*$(x[2]) + 0.6326865446578034*$(x[3]) + 0.8935006364653741*$(x[4]))) + 0.816846389391447*log(1 + exp(0.8959549878207418 + 0.6412415470262633*$(x[1]) + 0.6383904188579423*$(x[2]) + 0.9706157764666923*$(x[3]) + 0.8108132764672078*$(x[4]))))))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m

max_stars_repo_name: PSORLab/RSActivationFunctions

max_issues_repo_name: PSORLab/RSActivationFunctions

max_forks_repo_name: PSORLab/RSActivationFunctions


text:
"""
Contains structures, functions, and problem definitions
for core astrodynamics problems, including the
Restricted Two-body Problem, the Circular Restricted
Three-body Problem, and the N-body problem.
"""
module AstrodynamicsCore

# Common stuctures and functions
export AbstractUnitfulStructure, AbstractState, AbstractFrame, AbstractSystem, AbstractOrbit, CartesianState
export MassParameter, lengthunit, timeunit, velocityunit, massparameterunit, coordinateframe
export position_vector, velocity_vector, scalar_position, scalar_velocity
export AbstractFrame, Inertial, Synodic, Perifocal, ECI, HCI, convert, epoch
export NormalizedLengthUnit, NormalizedTimeUnit, Trajectory, Manifold
export convert, show, eltype, isapprox, isequal, string

# Core data structures and functions for R2BP calculations
export AbstractConic, Circular, Elliptical, Parabolic, Hyperbolic, Invalid
export KeplerianState, RestrictedTwoBodySystem, RestrictedTwoBodyOrbit
export R2BPState, R2BPSystem, R2BPOrbit, Orbit, CartesianOrbit, KeplerianOrbit
export mass_parameter, C3, SOI, SOA, hohmann
export eccentricity, semimajor_axis, inclination, RAAN, argument_of_periapsis, true_anomoly
export conic, keplerian, cartesian, perifocal, semi_parameter, periapsis_radius
export specific_angular_momentum_vector, specific_angular_momentum, eccentricity_vector
export specific_energy, specific_potential_energy, mean_motion, mean_motion_vector
export eccentric_anomoly, time_since_periapsis, period
export kepler, lambert, lambert_universal, lambert_lancaster_blanchard
export Sun, Mercury, Venus, Earth, Moon, Luna, Mars, Jupiter, Saturn, Uranus, Neptune, Pluto

# Core data structures and functions for CR3BP calculations
export CR3BPFrames, CR3BPSystem, CR3BPOrbit, CR3BPState
export SynodicCR3BPOrbit, NormalizedSynodicCR3BPOrbit, NormalizedSynodicSTMCR3BPOrbit
export SynodicCartesianState, InertialCartesianState, SynodicCartesianSTMState
export NormalizedCartesianState, MinimalCircularRestrictedThreeBodySystem
export CircularRestrictedThreeBodySystem, CircularRestrictedThreeBodyOrbit
export normalized_length_unit, normalized_time_unit, normalized_mass_parameter
export mass_parameters, primary_mass_parameter, secondary_mass_parameter
export primary_synodic_position, secondary_synodic_position, distance_to_primary, distance_to_secondary
export time_scale_factor, nondimensionalize, redimensionalize, nondimensional_radius
export nondimensionalize_length, nondimensionalize_time, nondimensionalize_velocity
export redimensionalize_length, redimensionalize_time, redimensionalize_velocity
export normalize, lagrange, inertial, synodic, accel, accel!, analyticalhalo
export potential_energy, jacobi_constant, zerovelocity_curves
export closest_approach, optimal_approach, body_position, position_wrt_body
export transform, transform_to_primary, transform_to_secondary
export SunVenus, SunEarth, EarthMoon, SunMars, SunJupiter, SunSaturn

# Module Dependencies
using Contour
using StaticArrays
using LinearAlgebra
using Roots: find_zero
using Unitful, UnitfulAstro, UnitfulAngles

# Source Code
include("Common/CommonTypes.jl")

include("R2BP/R2BPStates.jl")
include("R2BP/R2BPCalculations.jl")
include("R2BP/Kepler.jl")
include("R2BP/Lambert.jl")
include("R2BP/R2BPSystems.jl")

include("CR3BP/CR3BPStates.jl")
include("CR3BP/CR3BPCalculations.jl")
include("CR3BP/CR3BPSystems.jl")
include("CR3BP/CR3BPTransfers.jl")

end


text:
module TransformationElasticitySolver
using CutCellDG
using LinearAlgebra

function bulk_modulus(l, m)
    return l + 2m / 3
end

function onboundary(x, L, W)
    return x[2] ≈ 0.0 || x[1] ≈ L || x[2] ≈ W || x[1] ≈ 0.0
end

function nodal_displacement(
    mesh,
    basis,
    cellquads,
    facequads,
    interfacequads,
    stiffness,
    theta0,
    penalty;
    eta = 1,
)

    L, W = CutCellDG.mesh_widths(mesh)
    lambda1, mu1 = CutCellDG.lame_coefficients(stiffness, +1)
    transfstress =
        CutCellDG.plane_strain_transformation_stress(lambda1, mu1, theta0)

    sysmatrix = CutCellDG.SystemMatrix()
    sysrhs = CutCellDG.SystemRHS()

    CutCellDG.assemble_displacement_bilinear_forms!(
        sysmatrix,
        basis,
        cellquads,
        stiffness,
        mesh,
    )
    CutCellDG.assemble_bulk_transformation_linear_form!(
        sysrhs,
        transfstress,
        basis,
        cellquads,
        mesh,
    )
    CutCellDG.assemble_interelement_condition!(
        sysmatrix,
        basis,
        facequads,
        stiffness,
        mesh,
        penalty,
        eta,
    )
    CutCellDG.assemble_interelement_transformation_linear_form!(
        sysrhs,
        transfstress,
        basis,
        facequads,
        mesh,
    )

    CutCellDG.assemble_incoherent_interface_condition!(
        sysmatrix,
        basis,
        interfacequads,
        stiffness,
        mesh,
        penalty,
        eta,
    )
    CutCellDG.assemble_incoherent_interface_transformation_linear_form!(
        sysrhs,
        transfstress,
        basis,
        interfacequads,
        mesh,
    )

    x0 = CutCellDG.reference_corner(mesh)
    meshwidths = CutCellDG.mesh_widths(mesh)

    ############################################################################
    # LEFT BOUNDARY CONDITIONS
    # CutCellDG.assemble_penalty_displacement_component_bc!(
    #     sysmatrix,
    #     sysrhs,
    #     x -> 0.0,
    #     basis,
    #     facequads,
    #     stiffness,
    #     mesh,
    #     x -> x[1] ≈ x0[1],
    #     [1.0, 0.0],
    #     penalty,
    # )
    # CutCellDG.assemble_penalty_displacement_component_transformation_linear_form!(
    #     sysrhs,
    #     transfstress,
    #     basis,
    #     facequads,
    #     mesh,
    #     x -> x[1] ≈ x0[1],
    #     [1.0, 0.0],
    # )

    ############################################################################
    # BOTTOM BOUNDARY CONDITIONS

    # CutCellDG.assemble_penalty_displacement_component_bc!(
    #     sysmatrix,
    #     sysrhs,
    #     x -> 0.0,
    #     basis,
    #     facequads,
    #     stiffness,
    #     mesh,
    #     x -> x[2] ≈ x0[2],
    #     [0.0, 1.0],
    #     penalty,
    # )
    # CutCellDG.assemble_penalty_displacement_component_transformation_linear_form!(
    #     sysrhs,
    #     transfstress,
    #     basis,
    #     facequads,
    #     mesh,
    #     x -> x[2] ≈ x0[2],
    #     [0.0, 1.0],
    # )

    ############################################################################
    # RIGHT BOUNDARY CONDITIONS

    # CutCellDG.assemble_penalty_displacement_component_bc!(
    #     sysmatrix,
    #     sysrhs,
    #     x -> 0.0,
    #     basis,
    #     facequads,
    #     stiffness,
    #     mesh,
    #     x -> x[1] ≈ x0[1]+meshwidths[1],
    #     [1.0, 0.0],
    #     penalty,
    # )
    # CutCellDG.assemble_penalty_displacement_component_transformation_linear_form!(
    #     sysrhs,
    #     transfstress,
    #     basis,
    #     facequads,
    #     mesh,
    #     x -> x[1] ≈ x0[1]+meshwidths[1],
    #     [1.0, 0.0],
    # )

    ############################################################################
    # TOP BOUNDARY CONDITIONS

    # CutCellDG.assemble_penalty_displacement_component_bc!(
    #     sysmatrix,
    #     sysrhs,
    #     x -> 0.0,
    #     basis,
    #     facequads,
    #     stiffness,
    #     mesh,
    #     x -> x[2] ≈ x0[2]+meshwidths[2],
    #     [0.0, 1.0],
    #     penalty,
    # )
    # CutCellDG.assemble_penalty_displacement_component_transformation_linear_form!(
    #     sysrhs,
    #     transfstress,
    #     basis,
    #     facequads,
    #     mesh,
    #     x -> x[2] ≈ x0[2]+meshwidths[2],
    #     [0.0, 1.0],
    # )

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

    matrix = CutCellDG.sparse_displacement_operator(sysmatrix, mesh)
    rhs = CutCellDG.displacement_rhs_vector(sysrhs, mesh)

    solution = matrix \ rhs

    return solution
end

function construct_merged_mesh_and_quadratures(
    cutmesh,
    levelset,
    numqp;
    tinyratio = 0.3,
)
    cellquads = CutCellDG.CellQuadratures(cutmesh, levelset, numqp)
    interfacequads = CutCellDG.InterfaceQuadratures(cutmesh, levelset, numqp)
    facequads = CutCellDG.FaceQuadratures(cutmesh, levelset, numqp)

    mergedmesh = CutCellDG.MergedMesh!(
        cutmesh,
        cellquads,
        facequads,
        interfacequads,
        tinyratio = tinyratio,
    )

    return mergedmesh, cellquads, facequads, interfacequads
end

function parent_potential(
    nodaldisplacement,
    basis,
    refpoints,
    refcellids,
    normals,
    mesh,
    stiffness,
    V0,
)
    parentstrain = CutCellDG.parent_strain(
        nodaldisplacement,
        basis,
        refpoints,
        refcellids,
        mesh,
    )
    parentstress = CutCellDG.parent_stress(parentstrain, stiffness)
    parentstrainenergy =
        V0 * CutCellDG.strain_energy(parentstress, parentstrain)

    parentradialtraction =
        CutCellDG.traction_force_at_points(parentstress, normals)
    parentsrr = CutCellDG.traction_component(parentradialtraction, normals)
    parentdilatation = CutCellDG.dilatation(parentstrain)

    parentcompwork = V0 * (1.0 .+ parentdilatation) .* parentsrr

    parentpotential = parentstrainenergy - parentcompwork
    return parentpotential
end

function product_potential(
    nodaldisplacement,
    basis,
    refpoints,
    refcellids,
    normals,
    mesh,
    stiffness,
    theta0,
    V0,
)

    productstrain = CutCellDG.product_elastic_strain(
        nodaldisplacement,
        basis,
        theta0,
        refpoints,
        refcellids,
        mesh,
    )
    productstress = CutCellDG.product_stress(productstrain, stiffness, theta0)
    productstrainenergy =
        V0 * CutCellDG.strain_energy(productstress, productstrain)
    productradialtraction =
        CutCellDG.traction_force_at_points(productstress, normals)
    productsrr = CutCellDG.traction_component(productradialtraction, normals)
    productdilatation = CutCellDG.dilatation(productstrain)
    productcompwork = V0 * (1.0 .+ productdilatation) .* productsrr

    productpotential = productstrainenergy - productcompwork

    return productpotential
end

function potential_difference_at_closest_points(
    querypoints,
    nodaldisplacement,
    basis,
    spatialseedpoints,
    seedcellids,
    mesh,
    levelset,
    stiffness,
    theta0,
    V01,
    V02,
    tol,
    boundingradius,
)

    closestpoints, closestcellids = CutCellDG.closest_points_on_zero_levelset(
        querypoints,
        spatialseedpoints,
        seedcellids,
        levelset,
        tol,
        boundingradius,
    )

    normals = CutCellDG.collect_normals_at_spatial_points(
        closestpoints,
        closestcellids,
        levelset,
    )

    parentclosestrefpoints = CutCellDG.map_to_reference_on_merged_mesh(
        closestpoints,
        closestcellids,
        -1,
        mesh,
    )
    productclosestrefpoints = CutCellDG.map_to_reference_on_merged_mesh(
        closestpoints,
        closestcellids,
        +1,
        mesh,
    )

    parentpotential = parent_potential(
        nodaldisplacement,
        basis,
        parentclosestrefpoints,
        closestcellids,
        normals,
        mesh,
        stiffness,
        V02,
    )
    productpotential = product_potential(
        nodaldisplacement,
        basis,
        productclosestrefpoints,
        closestcellids,
        normals,
        mesh,
        stiffness,
        theta0,
        V01,
    )

    return productpotential - parentpotential
end

function average(v)
    return sum(v) / length(v)
end

function time_step_size(levelsetspeed, dx, CFL)
    s = maximum(abs.(levelsetspeed))
    return CFL * dx / s
end

function potential_difference_at_nodal_coordinates(
    cutmesh,
    basis,
    levelset,
    stiffness,
    theta0,
    numqp,
    penalty,
    spatialseedpoints,
    seedcellids,
    V01,
    V02,
    ΔG0,
    tol,
    boundingradius,
)

    mesh, cellquads, facequads, interfacequads =
        construct_merged_mesh_and_quadratures(cutmesh, levelset, numqp)

    nodaldisplacement = nodal_displacement(
        mesh,
        basis,
        cellquads,
        facequads,
        interfacequads,
        stiffness,
        theta0,
        penalty,
    )

    querypoints =
        CutCellDG.nodal_coordinates(CutCellDG.background_mesh(levelset))
    potentialdifference =
        (
            ΔG0 .+
            1e9 * potential_difference_at_closest_points(
                querypoints,
                nodaldisplacement,
                basis,
                spatialseedpoints,
                seedcellids,
                mesh,
                levelset,
                stiffness,
                theta0,
                V01,
                V02,
                tol,
                boundingradius,
            )
        ) / abs(ΔG0)

    return potentialdifference
end

function overpotential_difference_at_query_points(
    querypoints,
    cutmesh,
    basis,
    levelset,
    stiffness,
    theta0,
    numqp,
    penalty,
    spatialseedpoints,
    seedcellids,
    V01,
    V02,
    tol,
    boundingradius,
)

    mesh, cellquads, facequads, interfacequads =
        construct_merged_mesh_and_quadratures(cutmesh, levelset, numqp)

    nodaldisplacement = nodal_displacement(
        mesh,
        basis,
        cellquads,
        facequads,
        interfacequads,
        stiffness,
        theta0,
        penalty,
    )

    potentialdifference =
        (
             potential_difference_at_closest_points(
                querypoints,
                nodaldisplacement,
                basis,
                spatialseedpoints,
                seedcellids,
                mesh,
                levelset,
                stiffness,
                theta0,
                V01,
                V02,
                tol,
                boundingradius,
            )
        )

    return potentialdifference
end

function potential_difference_at_query_points(
    querypoints,
    cutmesh,
    basis,
    levelset,
    stiffness,
    theta0,
    numqp,
    penalty,
    spatialseedpoints,
    seedcellids,
    V01,
    V02,
    ΔG0,
    tol,
    boundingradius,
)

    mesh, cellquads, facequads, interfacequads =
        construct_merged_mesh_and_quadratures(cutmesh, levelset, numqp)

    nodaldisplacement = nodal_displacement(
        mesh,
        basis,
        cellquads,
        facequads,
        interfacequads,
        stiffness,
        theta0,
        penalty,
    )

    potentialdifference =
        (
            ΔG0 .+
            1e9 * potential_difference_at_closest_points(
                querypoints,
                nodaldisplacement,
                basis,
                spatialseedpoints,
                seedcellids,
                mesh,
                levelset,
                stiffness,
                theta0,
                V01,
                V02,
                tol,
                boundingradius,
            )
        ) / abs(ΔG0)

    return potentialdifference
end

function step_levelset(
    levelset,
    levelsetspeed,
    spatialseedpoints,
    seedcellids,
    paddedmesh,
    tol,
    boundingradius,
    CFL,
)

    paddedlevelset = CutCellDG.BoundaryPaddedLevelSet(
        paddedmesh,
        spatialseedpoints,
        seedcellids,
        levelset,
        tol,
        boundingradius,
    )

    dx = minimum(CutCellDG.grid_size(paddedmesh))
    dt = time_step_size(levelsetspeed, dx, CFL)

    newcoeffs =
        CutCellDG.step_first_order_levelset(paddedlevelset, levelsetspeed, dt)

    return newcoeffs
end

function step_interface(
    dgmesh,
    basis,
    levelset,
    stiffness,
    theta0,
    V01,
    V02,
    ΔG0,
    numqp;
    tol = 1e4eps(),
    boundingradius = 4.5,
    penaltyfactor = 1e3,
    CFL = 0.5,
    maxiter = 50,
)

    cutmesh = CutCellDG.CutMesh(dgmesh, levelset)
    refseedpoints, refseedcellids =
        CutCellDG.seed_zero_levelset(2, levelset, cutmesh)
    spatialseedpoints =
        CutCellDG.map_to_spatial(refseedpoints, refseedcellids, cutmesh)

    lambda1, mu1 = CutCellDG.lame_coefficients(stiffness, +1)
    lambda2, mu2 = CutCellDG.lame_coefficients(stiffness, -1)

    dx = minimum(CutCellDG.element_size(dgmesh))
    penalty = penaltyfactor / dx * 0.5 * (lambda1 + mu1 + lambda2 + mu2)

    potentialdifference = potential_difference_at_nodal_coordinates(
        cutmesh,
        basis,
        levelset,
        refseedpoints,
        refseedcellids,
        spatialseedpoints,
        stiffness,
        theta0,
        V01,
        V02,
        ΔG0,
        numqp,
        penalty,
        analyticalsolution,
        tol,
        boundingradius,
        maxiter,
    )

    newcoeffs = step_levelset(
        levelset,
        potentialdifference,
        refseedpoints,
        refseedcellids,
        spatialseedpoints,
        tol,
        boundingradius,
        CFL,
    )

    return newcoeffs
end

end


text:
# Implicit solvent models
# Based on the OpenMM source code

export
    AbstractGBSA,
    ImplicitSolventOBC,
    ImplicitSolventGBN2,
    born_radii_and_grad

"""
Generalized Born (GB) implicit solvent models augmented with the
hydrophobic solvent accessible surface area (SA) term.
Custom GBSA methods should sub-type this type.
"""
abstract type AbstractGBSA end

# Default solvent dielectric is 78.5 for consistency with AMBER
# Elsewhere it is 78.3
const gb_solvent_dielectric = 78.5
const gb_solute_dielectric = 1.0

const obc_offset = 0.009u"nm"
const gbn2_offset = 0.0195141u"nm"

const gb_probe_radius = 0.14u"nm"
const gb_sa_factor = 28.3919551u"kJ * mol^-1 * nm^-2"

const gbn2_neck_scale = 0.826836
const gbn2_neck_cut = 0.68u"nm"

const mbondi2_element_to_radius = Dict(
    "N"     => 0.155u"nm",
    "O"     => 0.15u"nm" ,
    "F"     => 0.15u"nm" ,
    "Si"    => 0.21u"nm" ,
    "P"     => 0.185u"nm",
    "S"     => 0.18u"nm" ,
    "Cl"    => 0.17u"nm" ,
    "C"     => 0.17u"nm" ,
    "H"     => 0.12u"nm" ,
    "H_N"   => 0.13u"nm" ,
    "H_ARG" => 0.117u"nm",
    "O_CAR" => 0.14u"nm" ,
    "-"     => 0.15u"nm" ,
)

const obc_element_to_screen = Dict(
    "H" => 0.85,
    "C" => 0.72,
    "N" => 0.79,
    "O" => 0.85,
    "F" => 0.88,
    "P" => 0.86,
    "S" => 0.96,
    "-" => 0.80,
)

const gbn2_element_to_screen = Dict(
    "H" => 1.425952,
    "C" => 1.058554,
    "N" => 0.733599,
    "O" => 1.061039,
    "F" => 0.5,
    "P" => 0.5,
    "S" => -0.703469,
    "-" => 0.5,
)

const gbn2_element_to_screen_nucleic = Dict(
    "H" => 1.696538,
    "C" => 1.268902,
    "N" => 1.4259728,
    "O" => 0.1840098,
    "F" => 0.5,
    "P" => 1.5450597,
    "S" => 0.05,
    "-" => 0.5,
)

const gbn2_atom_params = Dict(
    "H_α" => 0.788440, "H_β" => 0.798699, "H_γ" => 0.437334,
    "D_α" => 0.788440, "D_β" => 0.798699, "D_γ" => 0.437334,
    "C_α" => 0.733756, "C_β" => 0.506378, "C_γ" => 0.205844,
    "N_α" => 0.503364, "N_β" => 0.316828, "N_γ" => 0.192915,
    "O_α" => 0.867814, "O_β" => 0.876635, "O_γ" => 0.387882,
    "S_α" => 0.867814, "S_β" => 0.876635, "S_γ" => 0.387882,
    "-_α" => 1.0     , "-_β" => 0.8     , "-_γ" => 4.851   ,
)

const gbn2_atom_params_nucleic = Dict(
    "H_α" => 0.537050, "H_β" => 0.362861, "H_γ" => 0.116704 ,
    "D_α" => 0.537050, "D_β" => 0.362861, "D_γ" => 0.116704 ,
    "C_α" => 0.331670, "C_β" => 0.196842, "C_γ" => 0.093422 ,
    "N_α" => 0.686311, "N_β" => 0.463189, "N_γ" => 0.138722 ,
    "O_α" => 0.606344, "O_β" => 0.463006, "O_γ" => 0.142262 ,
    "S_α" => 0.606344, "S_β" => 0.463006, "S_γ" => 0.142262 ,
    "P_α" => 0.418365, "P_β" => 0.290054, "P_γ" => 0.1064245,
    "-_α" => 1.0     , "-_β" => 0.8     , "-_γ" => 4.851    ,
)

const gbn2_data_d0 = [
    2.26685, 2.32548, 2.38397, 2.44235, 2.50057, 2.55867, 2.61663, 2.67444,
    2.73212, 2.78965, 2.84705, 2.9043, 2.96141, 3.0184, 3.07524, 3.13196,
    3.18854, 3.24498, 3.30132, 3.35752, 3.4136,
    2.31191, 2.37017, 2.4283, 2.48632, 2.5442, 2.60197, 2.65961, 2.71711,
    2.77449, 2.83175, 2.88887, 2.94586, 3.00273, 3.05948, 3.1161, 3.1726,
    3.22897, 3.28522, 3.34136, 3.39738, 3.45072,
    2.35759, 2.41549, 2.47329, 2.53097, 2.58854, 2.646, 2.70333, 2.76056,
    2.81766, 2.87465, 2.93152, 2.98827, 3.0449, 3.10142, 3.15782, 3.21411,
    3.27028, 3.32634, 3.3823, 3.43813, 3.49387,
    2.4038, 2.46138, 2.51885, 2.57623, 2.63351, 2.69067, 2.74773, 2.80469,
    2.86152, 2.91826, 2.97489, 3.0314, 3.08781, 3.1441, 3.20031, 3.25638,
    3.31237, 3.36825, 3.42402, 3.4797, 3.53527,
    2.45045, 2.50773, 2.56492, 2.62201, 2.679, 2.7359, 2.7927, 2.8494, 2.90599,
    2.9625, 3.0189, 3.07518, 3.13138, 3.18748, 3.24347, 3.29937, 3.35515,
    3.41085, 3.46646, 3.52196, 3.57738,
    2.4975, 2.5545, 2.61143, 2.66825, 2.72499, 2.78163, 2.83818, 2.89464,
    2.95101, 3.00729, 3.06346, 3.11954, 3.17554, 3.23143, 3.28723, 3.34294,
    3.39856, 3.45409, 3.50952, 3.56488, 3.62014,
    2.54489, 2.60164, 2.6583, 2.71488, 2.77134, 2.8278, 2.88412, 2.94034,
    2.9965, 3.05256, 3.10853, 3.16442, 3.22021, 3.27592, 3.33154, 3.38707,
    3.44253, 3.49789, 3.55316, 3.60836, 3.66348,
    2.59259, 2.6491, 2.70553, 2.76188, 2.81815, 2.87434, 2.93044, 2.98646,
    3.04241, 3.09827, 3.15404, 3.20974, 3.26536, 3.32089, 3.37633, 3.4317,
    3.48699, 3.54219, 3.59731, 3.65237, 3.70734,
    2.64054, 2.69684, 2.75305, 2.80918, 2.86523, 2.92122, 2.97712, 3.03295,
    3.0887, 3.14437, 3.19996, 3.25548, 3.31091, 3.36627, 3.42156, 3.47677,
    3.5319, 3.58695, 3.64193, 3.69684, 3.75167,
    2.68873, 2.74482, 2.80083, 2.85676, 2.91262, 2.96841, 3.02412, 3.07976,
    3.13533, 3.19082, 3.24623, 3.30157, 3.35685, 3.41205, 3.46718, 3.52223,
    3.57721, 3.63213, 3.68696, 3.74174, 3.79644,
    2.73713, 2.79302, 2.84884, 2.90459, 2.96027, 3.01587, 3.0714, 3.12686,
    3.18225, 3.23757, 3.29282, 3.34801, 3.40313, 3.45815, 3.51315, 3.56805,
    3.6229, 3.67767, 3.73237, 3.78701, 3.84159,
    2.78572, 2.84143, 2.89707, 2.95264, 3.00813, 3.06356, 3.11892, 3.17422,
    3.22946, 3.28462, 3.33971, 3.39474, 3.44971, 3.5046, 3.55944, 3.61421,
    3.66891, 3.72356, 3.77814, 3.83264, 3.8871,
    2.83446, 2.89, 2.94547, 3.00088, 3.05621, 3.11147, 3.16669, 3.22183,
    3.27689, 3.33191, 3.38685, 3.44174, 3.49656, 3.55132, 3.60602, 3.66066,
    3.71523, 3.76975, 3.82421, 3.8786, 3.93293,
    2.88335, 2.93873, 2.99404, 3.04929, 3.10447, 3.15959, 3.21464, 3.26963,
    3.32456, 3.37943, 3.43424, 3.48898, 3.54366, 3.5983, 3.65287, 3.70737,
    3.76183, 3.81622, 3.87056, 3.92484, 3.97905,
    2.93234, 2.9876, 3.04277, 3.09786, 3.15291, 3.20787, 3.26278, 3.31764,
    3.37242, 3.42716, 3.48184, 3.53662, 3.591, 3.64551, 3.69995, 3.75435,
    3.80867, 3.86295, 3.91718, 3.97134, 4.02545,
    2.98151, 3.0366, 3.09163, 3.14659, 3.20149, 3.25632, 3.3111, 3.36581,
    3.42047, 3.47507, 3.52963, 3.58411, 3.63855, 3.69293, 3.74725, 3.80153,
    3.85575, 3.90991, 3.96403, 4.01809, 4.07211,
    3.03074, 3.08571, 3.14061, 3.19543, 3.25021, 3.30491, 3.35956, 3.41415,
    3.46869, 3.52317, 3.57759, 3.63196, 3.68628, 3.74054, 3.79476, 3.84893,
    3.90303, 3.95709, 4.01111, 4.06506, 4.11897,
    3.08008, 3.13492, 3.1897, 3.2444, 3.29905, 3.35363, 3.40815, 3.46263,
    3.51704, 3.57141, 3.62572, 3.67998, 3.73418, 3.78834, 3.84244, 3.8965,
    3.95051, 4.00447, 4.05837, 4.11224, 4.16605,
    3.12949, 3.18422, 3.23888, 3.29347, 3.348, 3.40247, 3.45688, 3.51124,
    3.56554, 3.6198, 3.674, 3.72815, 3.78225, 3.83629, 3.8903, 3.94425,
    3.99816, 4.05203, 4.10583, 4.15961, 4.21333,
    3.17899, 3.23361, 3.28815, 3.34264, 3.39706, 3.45142, 3.50571, 3.55997,
    3.61416, 3.66831, 3.72241, 3.77645, 3.83046, 3.8844, 3.93831, 3.99216,
    4.04598, 4.09974, 4.15347, 4.20715, 4.26078,
    3.22855, 3.28307, 3.33751, 3.39188, 3.4462, 3.50046, 3.55466, 3.6088,
    3.6629, 3.71694, 3.77095, 3.82489, 3.8788, 3.93265, 3.98646, 4.04022,
    4.09395, 4.14762, 4.20126, 4.25485, 4.3084,
]u"nm" ./ 10

const gbn2_data_m0 = [
    0.0381511, 0.0338587, 0.0301776, 0.027003, 0.0242506, 0.0218529,
    0.0197547, 0.0179109, 0.0162844, 0.0148442, 0.0135647, 0.0124243,
    0.0114047, 0.0104906, 0.00966876, 0.008928, 0.0082587, 0.00765255,
    0.00710237, 0.00660196, 0.00614589,
    0.0396198, 0.0351837, 0.0313767, 0.0280911, 0.0252409, 0.0227563,
    0.0205808, 0.0186681, 0.0169799, 0.0154843, 0.014155, 0.0129696,
    0.0119094, 0.0109584, 0.0101031, 0.00933189, 0.0086348, 0.00800326,
    0.00742986, 0.00690814, 0.00643255,
    0.041048, 0.0364738, 0.0325456, 0.0291532, 0.0262084, 0.0236399,
    0.0213897, 0.0194102, 0.0176622, 0.0161129, 0.0147351, 0.0135059,
    0.0124061, 0.0114192, 0.0105312, 0.00973027, 0.00900602, 0.00834965,
    0.0077535, 0.00721091, 0.00671609,
    0.0424365, 0.0377295, 0.0336846, 0.0301893, 0.0271533, 0.0245038,
    0.0221813, 0.0201371, 0.018331, 0.0167295, 0.0153047, 0.014033,
    0.0128946, 0.0118727, 0.0109529, 0.0101229, 0.00937212, 0.00869147,
    0.00807306, 0.00751003, 0.00699641,
    0.0437861, 0.0389516, 0.0347944, 0.0311998, 0.0280758, 0.0253479,
    0.0229555, 0.0208487, 0.0189864, 0.0173343, 0.0158637, 0.0145507,
    0.0133748, 0.0123188, 0.0113679, 0.0105096, 0.0097329, 0.00902853,
    0.00838835, 0.00780533, 0.0072733,
    0.0450979, 0.0401406, 0.0358753, 0.0321851, 0.0289761, 0.0261726,
    0.0237125, 0.0215451, 0.0196282, 0.017927, 0.0164121, 0.0150588,
    0.0138465, 0.0127573, 0.0117761, 0.0108902, 0.0100882, 0.00936068,
    0.00869923, 0.00809665, 0.00754661,
    0.0463729, 0.0412976, 0.0369281, 0.0331456, 0.0298547, 0.026978,
    0.0244525, 0.0222264, 0.0202567, 0.0185078, 0.0169498, 0.0155575,
    0.0143096, 0.0131881, 0.0121775, 0.0112646, 0.010438, 0.00968781,
    0.00900559, 0.00838388, 0.00781622,
    0.0476123, 0.0424233, 0.0379534, 0.034082, 0.0307118, 0.0277645,
    0.0251757, 0.0228927, 0.0208718, 0.0190767, 0.0174768, 0.0160466,
    0.0147642, 0.0136112, 0.0125719, 0.0116328, 0.0107821, 0.0100099,
    0.00930735, 0.00866695, 0.00808206,
    0.0488171, 0.0435186, 0.038952, 0.0349947, 0.0315481, 0.0285324,
    0.0258824, 0.0235443, 0.0214738, 0.0196339, 0.0179934, 0.0165262,
    0.0152103, 0.0140267, 0.0129595, 0.0119947, 0.0111206, 0.0103268,
    0.00960445, 0.00894579, 0.00834405,
    0.0499883, 0.0445845, 0.0399246, 0.0358844, 0.032364, 0.0292822,
    0.0265729, 0.0241815, 0.0220629, 0.0201794, 0.0184994, 0.0169964,
    0.0156479, 0.0144345, 0.0133401, 0.0123504, 0.0114534, 0.0106386,
    0.00989687, 0.00922037, 0.00860216,
    0.0511272, 0.0456219, 0.040872, 0.0367518, 0.0331599, 0.0300142,
    0.0272475, 0.0248045, 0.0226392, 0.0207135, 0.0189952, 0.0174574,
    0.0160771, 0.0148348, 0.0137138, 0.0126998, 0.0117805, 0.0109452,
    0.0101846, 0.00949067, 0.00885636,
    0.0522348, 0.0466315, 0.0417948, 0.0375973, 0.0339365, 0.030729,
    0.0279067, 0.0254136, 0.023203, 0.0212363, 0.0194809, 0.0179092,
    0.016498, 0.0152275, 0.0140807, 0.013043, 0.012102, 0.0112466,
    0.0104676, 0.00975668, 0.00910664,
    0.0533123, 0.0476145, 0.042694, 0.0384218, 0.0346942, 0.0314268,
    0.0285507, 0.026009, 0.0237547, 0.0217482, 0.0199566, 0.018352,
    0.0169108, 0.0156128, 0.0144408, 0.0133801, 0.0124179, 0.011543,
    0.010746, 0.0100184, 0.00935302,
    0.0543606, 0.0485716, 0.04357, 0.0392257, 0.0354335, 0.0321082,
    0.02918, 0.0265913, 0.0242943, 0.0222492, 0.0204225, 0.0187859,
    0.0173155, 0.0159908, 0.0147943, 0.0137111, 0.0127282, 0.0118343,
    0.0110197, 0.0102759, 0.00959549,
    0.0553807, 0.0495037, 0.0444239, 0.0400097, 0.0361551, 0.0327736,
    0.0297949, 0.0271605, 0.0248222, 0.0227396, 0.0208788, 0.0192111,
    0.0177122, 0.0163615, 0.0151413, 0.0140361, 0.013033, 0.0121206,
    0.0112888, 0.0105292, 0.00983409,
    0.0563738, 0.0504116, 0.0452562, 0.0407745, 0.0368593, 0.0334235,
    0.0303958, 0.0277171, 0.0253387, 0.0232197, 0.0213257, 0.0196277,
    0.0181013, 0.0167252, 0.0154817, 0.0143552, 0.0133325, 0.0124019,
    0.0115534, 0.0107783, 0.0100688,
    0.0573406, 0.0512963, 0.0460676, 0.0415206, 0.0375468, 0.0340583,
    0.030983, 0.0282614, 0.0258441, 0.0236896, 0.0217634, 0.020036,
    0.0184826, 0.017082, 0.0158158, 0.0146685, 0.0136266, 0.0126783,
    0.0118135, 0.0110232, 0.0102998,
    0.0582822, 0.0521584, 0.0468589, 0.0422486, 0.038218, 0.0346784,
    0.0315571, 0.0287938, 0.0263386, 0.0241497, 0.0221922, 0.0204362,
    0.0188566, 0.0174319, 0.0161437, 0.0149761, 0.0139154, 0.0129499,
    0.0120691, 0.0112641, 0.0105269,
    0.0591994, 0.0529987, 0.0476307, 0.042959, 0.0388734, 0.0352843,
    0.0321182, 0.0293144, 0.0268225, 0.0246002, 0.0226121, 0.0208283,
    0.0192232, 0.0177751, 0.0164654, 0.015278, 0.0141991, 0.0132167,
    0.0123204, 0.0115009, 0.0107504,
    0.0600932, 0.053818, 0.0483836, 0.0436525, 0.0395136, 0.0358764,
    0.0326669, 0.0298237, 0.0272961, 0.0250413, 0.0230236, 0.0212126,
    0.0195826, 0.0181118, 0.0167811, 0.0155744, 0.0144778, 0.0134789,
    0.0125673, 0.0117338, 0.0109702,
    0.0609642, 0.0546169, 0.0491183, 0.0443295, 0.0401388, 0.036455,
    0.0332033, 0.030322, 0.0277596, 0.0254732, 0.0234266, 0.0215892,
    0.0199351, 0.018442, 0.0170909, 0.0158654, 0.0147514, 0.0137365,
    0.0128101, 0.0119627, 0.0111863,
]u"nm^-1" .* 10

# This is force field dependent
is_carboxylate_O(at_data) = at_data.atom_type == "O2"

function atoms_bonded_to_N(atoms_data, bonds)
    bonded_to_N = falses(length(atoms_data))
    for (i, j) in zip(bonds.is, bonds.js)
        if atoms_data[i].element == "N"
            bonded_to_N[j] = true
        end
        if atoms_data[j].element == "N"
            bonded_to_N[i] = true
        end
    end
    return bonded_to_N
end

function mbondi2_radii(atoms_data, bonds; use_mbondi3=false,
                        element_to_radius=mbondi2_element_to_radius)
    bonded_to_N = atoms_bonded_to_N(atoms_data, bonds)
    return map(atoms_data, bonded_to_N) do at_data, at_bonded_to_N
        if use_mbondi3 && at_data.res_name == "ARG" &&
                (startswith(at_data.atom_name, "HH") || startswith(at_data.atom_name, "HE"))
            radius = element_to_radius["H_ARG"]
        elseif use_mbondi3 && is_carboxylate_O(at_data)
            radius = element_to_radius["O_CAR"]
        elseif at_data.element in ("H", "D")
            radius = at_bonded_to_N ? element_to_radius["H_N"] : element_to_radius["H"]
        else
            radius = dict_get(element_to_radius, at_data.element, element_to_radius["-"])
        end
        return radius
    end
end

function mbondi3_radii(atoms_data, bonds; element_to_radius=mbondi2_element_to_radius)
    return mbondi2_radii(atoms_data, bonds; use_mbondi3=true, element_to_radius=element_to_radius)
end

# We use a full atom pairwise table rather than looking up a value with the atom radius
# This works better with broadcasting
function lookup_table(full_table::AbstractArray{T}, radii) where T
    n_atoms = length(radii)
    table_positions = [(r - 0.1u"nm") * 200 for r in radii]
    # These zero-based indexes are converted to one-based when looking up the full table
    index_1, index_2 = zeros(Int, n_atoms), zeros(Int, n_atoms)
    weight_1, weight_2 = zeros(n_atoms), zeros(n_atoms)
    for (i, p) in enumerate(table_positions)
        if p <= 0.0u"nm"
            weight_1[i] = 1.0
        elseif p >= 20.0u"nm"
            index_1[i] = 20
            weight_1[i] = 1.0
        else
            ps = ustrip(u"nm", p)
            index_1[i] = Int(floor(ps))
            index_2[i] = index_1[i] + 1
            weight_1[i] = index_2[i] - ps
            weight_2[i] = 1.0 - weight_1[i]
        end
    end
    table = zeros(T, n_atoms, n_atoms)
    for i in 1:n_atoms
        for j in 1:n_atoms
            table[j, i] = weight_1[i] * weight_1[j] * full_table[index_1[i] * 21 + index_1[j] + 1] +
                          weight_1[i] * weight_2[j] * full_table[index_1[i] * 21 + index_2[j] + 1] +
                          weight_2[i] * weight_1[j] * full_table[index_2[i] * 21 + index_1[j] + 1] +
                          weight_2[i] * weight_2[j] * full_table[index_2[i] * 21 + index_2[j] + 1]
        end
    end
    return table
end

function lookup_table(full_table::AbstractArray{T}, radii::AbstractArray{<:AbstractFloat}) where T
    return lookup_table(full_table, radii * u"nm")
end

"""
    ImplicitSolventOBC(atoms, atoms_data, bonds)

Onufriev-Bashford-Case GBSA model.
Should be used along with a Coulomb or CoulombReactionField interaction.
The keyword argument `use_OBC2` determines whether to use parameter set
I (`false`, the default) or II (`true`).
"""
struct ImplicitSolventOBC{T, D, V, K, S, F, I} <: AbstractGBSA
    offset_radii::V
    scaled_offset_radii::V
    solvent_dielectric::T
    solute_dielectric::T
    kappa::K
    offset::D
    cutoff::D
    use_ACE::Bool
    α::T
    β::T
    γ::T
    probe_radius::D
    sa_factor::S
    factor_solute::F
    factor_solvent::F
    is::I
    js::I
end

function ImplicitSolventOBC(atoms::AbstractArray{Atom{T, M, D, E}},
                            atoms_data,
                            bonds;
                            solvent_dielectric=gb_solvent_dielectric,
                            solute_dielectric=gb_solute_dielectric,
                            kappa=0.0u"nm^-1",
                            offset=obc_offset,
                            cutoff=0.0u"nm",
                            probe_radius=gb_probe_radius,
                            sa_factor=gb_sa_factor,
                            use_ACE=true,
                            use_OBC2=false,
                            element_to_radius=mbondi2_element_to_radius,
                            element_to_screen=obc_element_to_screen) where {T, M, D, E}
    units = dimension(D) == u"𝐋"
    radii = mbondi2_radii(atoms_data, bonds; element_to_radius=element_to_radius)

    if units
        offset_radii = T.(radii .- offset)
    else
        offset_radii = ustrip.(T.(radii .- offset))
    end
    scaled_offset_radii = map(atoms_data, offset_radii) do at_data, offset_radius
        screen = dict_get(element_to_screen, at_data.element, element_to_screen["-"])
        return T(screen) * offset_radius
    end

    if use_OBC2
        # GBOBCII parameters
        α, β, γ = T(1.0), T(0.8), T(4.85)
    else
        # GBOBCI parameters
        α, β, γ = T(0.8), T(0.0), T(2.909125)
    end

    n_atoms = length(atoms)
    inds_i = hcat([collect(1:n_atoms) for i in 1:n_atoms]...)
    inds_j = permutedims(inds_i, (2, 1))

    coulomb_const = units ? coulombconst : ustrip(coulombconst)
    if !iszero(solute_dielectric)
        factor_solute = -T(coulomb_const) / T(solute_dielectric)
    else
        factor_solute = zero(T(coulomb_const))
    end
    if !iszero(solvent_dielectric)
        factor_solvent = T(coulomb_const) / T(solvent_dielectric)
    else
        factor_solvent = zero(T(coulomb_const))
    end

    if isa(atoms, CuArray)
        or = CuArray(offset_radii) # cu converts to Float32
        sor = CuArray(scaled_offset_radii)
        is, js = cu(inds_i), cu(inds_j)
    else
        or = offset_radii
        sor = scaled_offset_radii
        is, js = inds_i, inds_j
    end

    if units
        return ImplicitSolventOBC{T, D, typeof(or), typeof(T(kappa)), typeof(T(sa_factor)),
                        typeof(factor_solute), typeof(is)}(
                    or, sor, solvent_dielectric, solute_dielectric, T(kappa), offset,
                    cutoff, use_ACE, α, β, γ, probe_radius, T(sa_factor), factor_solute,
                    factor_solvent, is, js)
    else
        return ImplicitSolventOBC{T, T, typeof(or), typeof(T(ustrip(kappa))), T, T, typeof(is)}(
                    or, sor, solvent_dielectric, solute_dielectric, T(ustrip(kappa)),
                    ustrip(offset), ustrip(cutoff), use_ACE, α, β, γ, ustrip(probe_radius),
                    ustrip(sa_factor), factor_solute, factor_solvent, is, js)
    end
end

"""
    ImplicitSolventGBN2(atoms, atoms_data, bonds)

GBn2 solvation model.
Should be used along with a Coulomb or CoulombReactionField interaction.
"""
struct ImplicitSolventGBN2{T, D, VT, VD, K, S, F, I, TD, TM} <: AbstractGBSA
    offset_radii::VD
    scaled_offset_radii::VD
    solvent_dielectric::T
    solute_dielectric::T
    kappa::K
    offset::D
    cutoff::D
    use_ACE::Bool
    αs::VT
    βs::VT
    γs::VT
    probe_radius::D
    sa_factor::S
    factor_solute::F
    factor_solvent::F
    is::I
    js::I
    d0s::TD
    m0s::TM
    neck_scale::T
    neck_cut::D
end

function ImplicitSolventGBN2(atoms::AbstractArray{Atom{T, M, D, E}},
                                atoms_data,
                                bonds;
                                solvent_dielectric=gb_solvent_dielectric,
                                solute_dielectric=gb_solute_dielectric,
                                kappa=0.0u"nm^-1",
                                offset=gbn2_offset,
                                cutoff=0.0u"nm",
                                probe_radius=gb_probe_radius,
                                sa_factor=gb_sa_factor,
                                use_ACE=true,
                                neck_scale=gbn2_neck_scale,
                                neck_cut=gbn2_neck_cut,
                                element_to_radius=mbondi2_element_to_radius,
                                element_to_screen=gbn2_element_to_screen,
                                element_to_screen_nucleic=gbn2_element_to_screen_nucleic,
                                atom_params=gbn2_atom_params,
                                atom_params_nucleic=gbn2_atom_params_nucleic,
                                data_d0=gbn2_data_d0,
                                data_m0=gbn2_data_m0) where {T, M, D, E}
    units = dimension(D) == u"𝐋"
    radii = mbondi3_radii(atoms_data, bonds; element_to_radius=element_to_radius)
    nucleic_acid_residues = ("A", "C", "G", "U", "DA", "DC", "DG", "DT")

    if units
        offset_radii = T.(radii .- offset)
    else
        offset_radii = T.(ustrip.(radii) .- ustrip(offset))
    end
    scaled_offset_radii = map(atoms_data, offset_radii) do at_data, offset_radius
        if at_data.res_name in nucleic_acid_residues
            screen = dict_get(element_to_screen_nucleic, at_data.element, element_to_screen_nucleic["-"])
        else
            screen = dict_get(element_to_screen, at_data.element, element_to_screen["-"])
        end
        return T(screen) * offset_radius
    end

    αs_cpu = map(atoms_data) do at_data
        if at_data.res_name in nucleic_acid_residues
            α = dict_get(atom_params_nucleic, at_data.element * "_α", atom_params_nucleic["-_α"])
        else
            α = dict_get(atom_params, at_data.element * "_α", atom_params["-_α"])
        end
        return T(α)
    end
    βs_cpu = map(atoms_data) do at_data
        if at_data.res_name in nucleic_acid_residues
            β = dict_get(atom_params_nucleic, at_data.element * "_β", atom_params_nucleic["-_β"])
        else
            β = dict_get(atom_params, at_data.element * "_β", atom_params["-_β"])
        end
        return T(β)
    end
    γs_cpu = map(atoms_data) do at_data
        if at_data.res_name in nucleic_acid_residues
            γ = dict_get(atom_params_nucleic, at_data.element * "_γ", atom_params_nucleic["-_γ"])
        else
            γ = dict_get(atom_params, at_data.element * "_γ", atom_params["-_γ"])
        end
        return T(γ)
    end

    n_atoms = length(atoms)
    inds_i = hcat([collect(1:n_atoms) for i in 1:n_atoms]...)
    inds_j = permutedims(inds_i, (2, 1))

    table_d0_units = T.(lookup_table(data_d0, radii))
    table_m0_units = T.(lookup_table(data_m0, radii))
    if units
        table_d0 = table_d0_units
        table_m0 = table_m0_units
    else
        table_d0 = ustrip.(table_d0_units)
        table_m0 = ustrip.(table_m0_units)
    end

    coulomb_const = units ? coulombconst : ustrip(coulombconst)
    if !iszero(solute_dielectric)
        factor_solute = -T(coulomb_const) / T(solute_dielectric)
    else
        factor_solute = zero(T(coulomb_const))
    end
    if !iszero(solvent_dielectric)
        factor_solvent = T(coulomb_const) / T(solvent_dielectric)
    else
        factor_solvent = zero(T(coulomb_const))
    end

    if isa(atoms, CuArray)
        or = CuArray(offset_radii) # cu converts to Float32
        sor = CuArray(scaled_offset_radii)
        is, js = cu(inds_i), cu(inds_j)
        d0s, m0s = CuArray(table_d0), CuArray(table_m0)
        αs, βs, γs = CuArray(αs_cpu), CuArray(βs_cpu), CuArray(γs_cpu)
    else
        or = offset_radii
        sor = scaled_offset_radii
        is, js = inds_i, inds_j
        d0s, m0s = table_d0, table_m0
        αs, βs, γs = αs_cpu, βs_cpu, γs_cpu
    end

    if units
        return ImplicitSolventGBN2{T, D, typeof(αs), typeof(or), typeof(T(kappa)), typeof(T(sa_factor)),
                        typeof(factor_solute), typeof(is), typeof(d0s), typeof(m0s)}(
                    or, sor, solvent_dielectric, solute_dielectric, T(kappa), offset, cutoff,
                    use_ACE, αs, βs, γs, probe_radius, T(sa_factor), factor_solute,
                    factor_solvent, is, js, d0s, m0s, neck_scale, neck_cut)
    else
        return ImplicitSolventGBN2{T, T, typeof(αs), typeof(or), typeof(T(ustrip(kappa))), T, T,
                        typeof(is), typeof(d0s), typeof(m0s)}(
                    or, sor, solvent_dielectric, solute_dielectric, T(ustrip(kappa)), ustrip(offset),
                    ustrip(cutoff), use_ACE, αs, βs, γs, ustrip(probe_radius), ustrip(sa_factor),
                    factor_solute, factor_solvent, is, js, d0s, m0s, neck_scale, ustrip(neck_cut))
    end
end

function born_radii_loop_OBC(coord_i, coord_j, ori, srj, cutoff, box_size)
    I = zero(coord_i[1] / unit(cutoff)^2)
    r = norm(vector(coord_i, coord_j, box_size))
    if iszero(r) || (!iszero(cutoff) && r > cutoff)
        return I
    end
    U = r + srj
    if ori < U
        D_ij = abs(r - srj)
        L = max(ori, D_ij)
        I += (1/L - 1/U + (r - (srj^2)/r)*(1/(U^2) - 1/(L^2))/4 + log(L/U)/(2*r)) / 2
        if ori < (srj - r)
            I += 2 * (1/ori - 1/L)
        end
    end
    return I
end

"""
    born_radii_and_grad(inter, coords, box_size)

Calculate Born radii and gradients of Born radii and surface area overlap
with respect to atomic distance.
Custom GBSA methods should implement this function.
"""
function born_radii_and_grad(inter::ImplicitSolventOBC, coords, box_size)
    coords_i = @view coords[inter.is]
    coords_j = @view coords[inter.js]
    oris = @view inter.offset_radii[inter.is]
    srjs = @view inter.scaled_offset_radii[inter.js]
    loop_res = born_radii_loop_OBC.(coords_i, coords_j, oris, srjs, inter.cutoff,
                                    (box_size,))
    Is = dropdims(sum(loop_res; dims=2); dims=2)
    I_grads = zero(loop_res) ./ unit(inter.cutoff)

    ori = inter.offset_radii
    radii = ori .+ inter.offset
    ψs = Is .* ori
    ψs2 = ψs .^ 2
    α, β, γ = inter.α, inter.β, inter.γ
    tanh_sums = tanh.(α .* ψs .- β .* ψs2 .+ γ .* ψs2 .* ψs)
    Bs = 1 ./ (1 ./ ori .- tanh_sums ./ radii)
    grad_terms = ori .* (α .- 2 .* β .* ψs .+ 3 .* γ .* ψs2)
    B_grads = (1 .- tanh_sums .^ 2) .* grad_terms ./ radii

    return Bs, B_grads, I_grads
end

struct BornRadiiGBN2LoopResult{I, G}
    I::I
    I_grad::G
end

get_I(r::BornRadiiGBN2LoopResult) = r.I
get_I_grad(r::BornRadiiGBN2LoopResult) = r.I_grad

function born_radii_loop_GBN2(coord_i::SVector{D, T}, coord_j, ori, orj, srj, cutoff, offset,
                                neck_scale, neck_cut, d0, m0, box_size) where {D, T}
    I = zero(coord_i[1] / unit(cutoff)^2)
    I_grad = zero(coord_i[1] / unit(cutoff)^3)
    r = norm(vector(coord_i, coord_j, box_size))
    if iszero(r) || (!iszero(cutoff) && r > cutoff)
        return BornRadiiGBN2LoopResult(I, I_grad)
    end
    U = r + srj
    if ori < U
        D_ij = abs(r - srj)
        L = max(ori, D_ij)
        I += (1/L - 1/U + (r - (srj^2)/r)*(1/(U^2) - 1/(L^2))/4 + log(L/U)/(2*r)) / 2
        if ori < (srj - r)
            I += 2 * (1/ori - 1/L)
        end
    end
    radius_i = ori + offset
    radius_j = orj + offset
    if r < (radius_i + radius_j + neck_cut)
        if dimension(T) == u"𝐋"
            r_d0_strip = 10 * ustrip(u"nm", r - d0) # The integral uses Å
        else
            r_d0_strip = 10 * (r - d0)
        end
        denom = 1 + r_d0_strip^2 + 3 * r_d0_strip^6 / 10
        I += neck_scale * m0 / denom
        numer = 2 * r_d0_strip + 9 * r_d0_strip^5 / 5
        I_grad -= 10 * neck_scale * m0 * numer / (denom^2 * unit(cutoff))
    end
    return BornRadiiGBN2LoopResult(I, I_grad)
end

function born_radii_and_grad(inter::ImplicitSolventGBN2, coords, box_size)
    coords_i = @view coords[inter.is]
    coords_j = @view coords[inter.js]
    oris = @view inter.offset_radii[inter.is]
    orjs = @view inter.offset_radii[inter.js]
    srjs = @view inter.scaled_offset_radii[inter.js]
    loop_res = born_radii_loop_GBN2.(coords_i, coords_j, oris, orjs, srjs, inter.cutoff,
                    inter.offset, inter.neck_scale, inter.neck_cut, inter.d0s,
                    inter.m0s, (box_size,))
    Is = dropdims(sum(get_I.(loop_res); dims=2); dims=2)
    I_grads = get_I_grad.(loop_res)

    ori = inter.offset_radii
    radii = ori .+ inter.offset
    ψs = Is .* ori
    ψs2 = ψs .^ 2
    αs, βs, γs = inter.αs, inter.βs, inter.γs
    tanh_sums = tanh.(αs .* ψs .- βs .* ψs2 .+ γs .* ψs2 .* ψs)
    Bs = 1 ./ (1 ./ ori .- tanh_sums ./ radii)
    grad_terms = ori .* (αs .- 2 .* βs .* ψs .+ 3 .* γs .* ψs2)
    B_grads = (1 .- tanh_sums .^ 2) .* grad_terms ./ radii

    return Bs, B_grads, I_grads
end

# Store the results of the ij broadcasts during force calculation
struct ForceLoopResult1{T, V}
    bi::T
    bj::T
    fi::V
    fj::V
end

get_bi(r::ForceLoopResult1) = r.bi
get_bj(r::ForceLoopResult1) = r.bj

struct ForceLoopResult2{V}
    fi::V
    fj::V
end

get_fi(r::Union{ForceLoopResult1, ForceLoopResult2}) = r.fi
get_fj(r::Union{ForceLoopResult1, ForceLoopResult2}) = r.fj

function gb_force_loop_1(coord_i, coord_j, i, j, charge_i, charge_j, Bi, Bj, cutoff,
                            factor_solute, factor_solvent, kappa, box_size)
    if j < i
        zero_force = zero(factor_solute ./ coord_i .^ 2)
        return ForceLoopResult1(zero_force[1], zero_force[1], zero_force, zero_force)
    end
    dr = vector(coord_i, coord_j, box_size)
    r2 = sum(abs2, dr)
    if !iszero(cutoff) && r2 > cutoff^2
        zero_force = zero(factor_solute ./ coord_i .^ 2)
        return ForceLoopResult1(zero_force[1], zero_force[1], zero_force, zero_force)
    end
    alpha2_ij = Bi * Bj
    D = r2 / (4 * alpha2_ij)
    exp_term = exp(-D)
    denominator2 = r2 + alpha2_ij * exp_term
    denominator = sqrt(denominator2)
    if iszero(kappa)
        pre_factor = factor_solute + factor_solvent
    else
        pre_factor = factor_solute + exp(-kappa * denominator) * factor_solvent +
                        kappa * denominator * exp(-kappa * denominator) * factor_solvent
    end
    Gpol = (pre_factor * charge_i * charge_j) / denominator
    dGpol_dr = -Gpol * (1 - exp_term/4) / denominator2
    dGpol_dalpha2_ij = -Gpol * exp_term * (1 + D) / (2 * denominator2)
    change_born_force_i = dGpol_dalpha2_ij * Bj
    if i != j
        change_born_force_j = dGpol_dalpha2_ij * Bi
        fdr = dr * dGpol_dr
        change_fs_i =  fdr
        change_fs_j = -fdr
        return ForceLoopResult1(change_born_force_i, change_born_force_j,
                                change_fs_i, change_fs_j)
    else
        zero_force = zero(factor_solute ./ coord_i .^ 2)
        return ForceLoopResult1(change_born_force_i, zero_force[1], zero_force, zero_force)
    end
end

function gb_force_loop_2(coord_i, coord_j, bi, ig, ori, srj, cutoff, box_size)
    dr = vector(coord_i, coord_j, box_size)
    r = norm(dr)
    if iszero(r) || (!iszero(cutoff) && r > cutoff)
        zero_force = zero(bi ./ coord_i .^ 2)
        return ForceLoopResult2(zero_force, zero_force)
    end
    rsrj = r + srj
    if ori < rsrj
        D = abs(r - srj)
        L = inv(max(ori, D))
        U = inv(rsrj)
        rinv = inv(r)
        r2inv = rinv^2
        t3 = (1 + (srj^2)*r2inv)*(L^2 - U^2)/8 + log(U/L)*r2inv/4
        de = bi * (t3 - ig) * rinv
        fdr = dr * de
        return ForceLoopResult2(-fdr, fdr)
    else
        zero_force = zero(bi ./ coord_i .^ 2)
        return ForceLoopResult2(zero_force, zero_force)
    end
end

function forces(inter::AbstractGBSA, sys, neighbors=nothing)
    coords, atoms, box_size = sys.coords, sys.atoms, sys.box_size
    Bs, B_grads, I_grads = born_radii_and_grad(inter, coords, box_size)

    if inter.use_ACE
        radii = inter.offset_radii .+ inter.offset
        sa_terms = inter.sa_factor .* (radii .+ inter.probe_radius) .^ 2 .* (radii ./ Bs) .^ 6
        born_forces = (-6 .* sa_terms ./ Bs) .* (Bs .> zero(inter.offset))
    else
        born_forces = zeros(typeof(inter.sa_factor * inter.cutoff), length(sys))
    end

    coords_i = @view coords[inter.is]
    coords_j = @view coords[inter.js]
    charges = charge.(atoms)
    charges_i = @view charges[inter.is]
    charges_j = @view charges[inter.js]
    Bsi = @view Bs[inter.is]
    Bsj = @view Bs[inter.js]
    loop_res_1 = gb_force_loop_1.(coords_i, coords_j, inter.is, inter.js, charges_i, charges_j,
                                    Bsi, Bsj, inter.cutoff, inter.factor_solute,
                                    inter.factor_solvent, inter.kappa, (box_size,))
    born_forces = born_forces .+ dropdims(sum(get_bi.(loop_res_1); dims=2); dims=2)
    born_forces = born_forces .+ dropdims(sum(get_bj.(loop_res_1); dims=1); dims=1)
    fs = dropdims(sum(get_fi.(loop_res_1); dims=2); dims=2) .+ dropdims(sum(get_fj.(loop_res_1); dims=1); dims=1)

    born_forces_2 = born_forces .* (Bs .^ 2) .* B_grads

    bis = @view born_forces_2[inter.is]
    oris = @view inter.offset_radii[inter.is]
    srjs = @view inter.scaled_offset_radii[inter.js]
    loop_res_2 = gb_force_loop_2.(coords_i, coords_j, bis, I_grads, oris, srjs,
                                    inter.cutoff, (box_size,))

    return fs .+ dropdims(sum(get_fi.(loop_res_2); dims=2); dims=2) .+ dropdims(sum(get_fj.(loop_res_2); dims=1); dims=1)
end

function gb_energy_loop(coord_i, coord_j, i, j, charge_i, charge_j, Bi, Bj, ori, cutoff,
                        factor_solute, factor_solvent, kappa, offset, probe_radius,
                        sa_factor, use_ACE, box_size)
    if i == j
        if iszero(kappa)
            pre_factor = factor_solute + factor_solvent
        else
            pre_factor = factor_solute + exp(-kappa * Bi) * factor_solvent
        end
        E = pre_factor * (charge_i^2) / (2*Bi)
        if use_ACE && (Bi > zero(offset))
            radius_i = ori + offset
            E += sa_factor * (radius_i + probe_radius)^2 * (radius_i / Bi)^6
        end
        return E
    elseif j > i
        r2 = sum(abs2, vector(coord_i, coord_j, box_size))
        if !iszero(cutoff) && r2 > cutoff^2
            return zero(factor_solute / offset)
        end
        f = sqrt(r2 + Bi*Bj*exp(-r2/(4*Bi*Bj)))
        if iszero(cutoff)
            f_cutoff = 1/f
        else
            f_cutoff = (1/f - 1/cutoff)
        end
        if iszero(kappa)
            pre_factor = factor_solute + factor_solvent
        else
            pre_factor = factor_solute + exp(-kappa * f) * factor_solvent
        end
        return pre_factor * charge_i * charge_j * f_cutoff
    else
        return zero(factor_solute / offset)
    end
end

function potential_energy(inter::AbstractGBSA, sys, neighbors=nothing)
    coords, atoms, box_size = sys.coords, sys.atoms, sys.box_size
    Bs, B_grads, I_grads = born_radii_and_grad(inter, coords, box_size)

    coords_i = @view coords[inter.is]
    coords_j = @view coords[inter.js]
    charges = charge.(atoms)
    charges_i = @view charges[inter.is]
    charges_j = @view charges[inter.js]
    Bsi = @view Bs[inter.is]
    Bsj = @view Bs[inter.js]
    oris = @view inter.offset_radii[inter.is]
    return sum(gb_energy_loop.(coords_i, coords_j, inter.is, inter.js, charges_i, charges_j,
                                Bsi, Bsj, oris, inter.cutoff, inter.factor_solute,
                                inter.factor_solvent, inter.kappa, inter.offset,
                                inter.probe_radius, inter.sa_factor, inter.use_ACE,
                                (box_size,)))
end


text:
## Indices

function Base.filter(pred, inds::AbstractIndices)
    out = empty(inds)
    for i in inds
        if pred(i)
            insert!(out, i)
        end
    end
    return out
end

# Occasionally it might be faster to copy all and remove just a few indices...
# TODO: export `unsafe_filter`, or hook up to certain datatypes or something
function unsafe_filter(pred, inds::AbstractIndices)
    out = copy(inds)
    Base.unsafe_filter!(pred, out)
    return out
end

struct FilteredIndices{I, Parent <: AbstractIndices{I}, Pred} <: AbstractIndices{I}
    parent::Parent
    pred::Pred
end

Base.parent(inds::FilteredIndices) = getfield(inds, :parent)
_pred(inds::FilteredIndices) = getfield(inds, :pred)
Base.IteratorSize(::FilteredIndices) = Base.SizeUnknown() # You can still call `length`, but it is slower
Base.length(inds::FilteredIndices) = count(_pred(inds), parent(inds))

function Base.in(i::I, inds::FilteredIndices{I}) where {I}
    return i in parent(inds) && _pred(inds)(i)
end

function Base.iterate(inds::FilteredIndices{I}, s...) where {I}
    tmp = iterate(parent(inds), s...)
    tmp === nothing && return nothing
    (i::I, s2) = tmp
    if _pred(inds)(i)
        return (i, s2)
    else
        return iterate(inds, s2)
    end
    return nothing
end

function Base.iterate(inds::ReverseIndices{I,<:FilteredIndices}, s...) where {I}
    tmp = iterate(Iterators.reverse(parent(inds)), s...)
    tmp === nothing && return nothing
    (i::I, s2) = tmp
    if _pred(inds)(i)
        return (i, s2)
    else
        return iterate(inds, s2)
    end
    return nothing
end

istokenizable(inds::FilteredIndices) = istokenizable(parent(inds))
@propagate_inbounds function iteratetoken(inds::FilteredIndices{I}, s...) where {I}
    p = parent(inds)
    tmp = iteratetoken(p, s...)
    while (tmp !== nothing)
        (token, state) = tmp
        i = gettokenvalue(p, token)
        if _pred(inds)(i)
            return tmp
        end
        tmp = iteratetoken(p, state)
    end
    return nothing
end
@propagate_inbounds function iteratetoken_reverse(inds::FilteredIndices{I}, s...) where {I}
    p = parent(inds)
    tmp = iteratetoken_reverse(p, s...)
    while (tmp !== nothing)
        (token, state) = tmp
        i = gettokenvalue(p, token)
        if _pred(inds)(i)
            return tmp
        end
        tmp = iteratetoken_reverse(p, state)
    end
    return nothing
end
@propagate_inbounds function gettoken(inds::FilteredIndices{I}, i::I) where {I}
    (hastoken, token) = gettoken(parent(inds), i)
    @boundscheck if hastoken
        return (_pred(inds)(@inbounds gettokenvalue(parent, token)), token)
    end
    return (hastoken, token)
end
@propagate_inbounds gettokenvalue(inds::FilteredIndices, t) = gettokenvalue(parent(inds), t)

function filterview(pred, inds::AbstractIndices{I}) where {I}
    return FilteredIndices{I, typeof(inds), typeof(pred)}(inds, pred)
end

Iterators.reverse(inds::FilteredIndices) = filterview(_pred(inds), Iterators.reverse(parent(inds)))

function Base.isempty(inds::FilteredIndices)
    for _ in inds
        return false
    end
    return true
end

function randtoken(rng::Random.AbstractRNG, inds::FilteredIndices)
    if isempty(inds)
        throw(ArgumentError("range must be non-empty"))
    end
    p = parent(inds)
    while true
        # Rejection sampling
        token = randtoken(rng, p)
        if _pred(inds)(@inbounds gettokenvalue(p, token))
            return token
        end
    end
end

## Dictionaries

function Base.filter(pred, dict::AbstractDictionary)
    out = empty(dict)
    for (i, v) in pairs(dict)
        if pred(v)
            insert!(out, i, v)
        end
    end
    return out
end

# Occasionally it might be faster to copy all and remove just a few indices...
# TODO: export `unsafe_filter`, or hook up to certain datatypes or something
function unsafe_filter(pred, dict::AbstractDictionary)
    out = copy(dict)
    Base.unsafe_filter!(pred, dict)
    return out
end

struct FilteredDictionary{I, V, Parent <: AbstractDictionary{I, V}, Pred} <: AbstractDictionary{I, V}
    parent::Parent
    pred::Pred
end

Base.parent(dict::FilteredDictionary) = getfield(dict, :parent)
_pred(dict::FilteredDictionary) = getfield(dict, :pred)
Base.IteratorSize(::FilteredDictionary) = Base.SizeUnknown() # You can still call `length`, but it is slower

@inline function Base.keys(dict::FilteredDictionary)
    return filterview(i -> _pred(dict)(parent(dict)[i]), keys(parent(dict)))
end

function Base.isassigned(dict::FilteredDictionary{I}, i::I) where {I}
    p = parent(dict)
    if !isassigned(p, i)
        return false
    end
    @inbounds v = p[i]
    return _pred(dict)(v)
end

@propagate_inbounds function Base.getindex(dict::FilteredDictionary{I}, i::I) where {I}
    p = parent(dict)
    (hastoken, token) = gettoken(p, i)
    @boundscheck if !hastoken
        throw(IndexError("Dictionary does not contain index: $i"))
    end
    out = gettokenvalue(p, token)
    @boundscheck if !_pred(dict)(out)
        throw(IndexError("Dictionary does not contain index: $i"))
    end
    return out
end

# These dictionaries are not settable - changing the values can change the keys

istokenizable(dict::FilteredDictionary) = istokenizable(parent(dict))
@propagate_inbounds function gettoken(dict::FilteredDictionary{I}, i::I) where {I}
    (hastoken, token) = gettoken(parent(dict), i)
    @boundscheck if hastoken
        return (_pred(dict)(@inbounds gettokenvalue(parent, token)), token)
    end
    return (hastoken, token)
end
@propagate_inbounds gettokenvalue(dict::FilteredDictionary, t) = gettokenvalue(parent(dict), t)
@propagate_inbounds istokenassigned(dict::FilteredDictionary, t) = istokenassigned(parent(dict), t)

function filterview(pred, inds::AbstractDictionary{I, T}) where {I, T}
    return FilteredDictionary{I, T, typeof(inds), typeof(pred)}(inds, pred)
end

Base.similar(dict::FilteredDictionary, ::Type{T}, indices) where {T} = similar(parent(dict), T, indices)
Base.empty(dict::FilteredDictionary, ::Type{T}) where {T} = empty(parent(dict), T)


text:
immutable Beta <: ContinuousUnivariateDistribution
    α::Float64
    β::Float64

    function Beta(a::Real, b::Real)
        (a > zero(a) && b > zero(b)) || error("α and β must be positive")
        new(float64(a), float64(b))
    end

    Beta(α::Real) = Beta(α, α)
    Beta() = new(1.0, 1.0)
end

@_jl_dist_2p Beta beta

@distr_support Beta 0.0 1.0

#### Parameters

params(d::Beta) = (d.α, d.β)

#### Statistics

mean(d::Beta) = ((α, β) = params(d); α / (α + β))

function mode(d::Beta)
    (α, β) = params(d)
    (α > 1.0 && β > 1.0) || error("mode is defined only when α > 1 and β > 1.")
    return (α - 1.0) / (α + β - 2.0)
end

modes(d::Beta) = [mode(d)]

function var(d::Beta)
    (α, β) = params(d)
    s = α + β
    return (α * β) / (abs2(s) * (s + 1.0))
end

meanlogx(d::Beta) = ((α, β) = params(d); digamma(α) - digamma(α + β))

varlogx(d::Beta) = ((α, β) = params(d); trigamma(α) - trigamma(α + β))

function skewness(d::Beta)
    (α, β) = params(d)
    if α == β
        return 0.0
    else
        s = α + β
        (2.0 * (β - α) * sqrt(s + 1.0)) / ((s + 2.0) * sqrt(α * β))
    end
end

function kurtosis(d::Beta)
    α, β = params(d)
    s = α + β
    p = α * β
    6.0 * (abs2(α - β) * (s + 1.0) - p * (s + 2.0)) / (p * (s + 2.0) * (s + 3.0))
end

function entropy(d::Beta)
    α, β = params(d)
    s = α + β
    lbeta(α, β) - (α - 1.0) * digamma(α) - (β - 1.0) * digamma(β) +
        (s - 2.0) * digamma(s)
end

#### Evaluation

gradlogpdf(d::Beta, x::Float64) =
    ((α, β) = params(d); 0.0 <= x <= 1.0 ? (α - 1.0) / x - (β - 1.0) / (1 - x) : 0.0)

## Fit model

# TODO: add MLE method (should be similar to Dirichlet)

# This is a moment-matching method (not MLE)
#
function fit(::Type{Beta}, x::AbstractArray)
    x_bar = mean(x)
    v_bar = varm(x, x_bar)
    α = x_bar * (((x_bar * (1.0 - x_bar)) / v_bar) - 1.0)
    β = (1.0 - x_bar) * (((x_bar * (1.0 - x_bar)) / v_bar) - 1.0)
    Beta(α, β)
end


text:
"""
    TextFolders(textfile; labelfn = parentname, split = false)

Recipe for loading a single-label text classification dataset
stored in hierarchical folder format.
"""
Base.@kwdef struct TextFolders <: Datasets.DatasetRecipe
    labelfn = parentname
    split::Bool = false
    filefilterfn = _ -> true
end

Datasets.recipeblocks(::Type{TextFolders}) = Tuple{Paragraph,Label}

function Datasets.loadrecipe(recipe::TextFolders, path)
    isdir(path) || error("$path is not a directory")
    data = loadfolderdata(
        path,
        filterfn=f -> istextfile(f) && recipe.filefilterfn(f),
        loadfn=(loadfile, recipe.labelfn),
        splitfn=recipe.split ? grandparentname : nothing)

    (recipe.split ? length(data) > 0 : nobs(data) > 0) || error("No text files found in $path")

    labels = recipe.split ? first(values(data))[2] : data[2]
    blocks = (Paragraph(), Label(unique(eachobs(labels))))
    length(blocks[2].classes) > 1 || error("Expected multiple different labels, got: $(blocks[2].classes))")
    return data, blocks
end

# Registering recipes

const RECIPES = Dict{String,Vector{Datasets.DatasetRecipe}}(
    "imdb" => [TextFolders(
        filefilterfn=f -> !occursin(r"tmp_clas|tmp_lm|unsup", f)
    )],
)

function _registerrecipes()
    for (name, recipes) in RECIPES, recipe in recipes
        Datasets.registerrecipe!(Datasets.FASTAI_DATA_REGISTRY, name, recipe)
    end
end

## Tests

@testset "TextFolders [Recipe]" begin
    @test length(finddatasets(name="imdb")) >= 1
end


text:
type TreeCallableFixedRateEngine{S <: ShortRateModel} <: LatticeShortRateModelEngine{S}
  model::S
  timeSteps::Int
  common::LatticeShortRateModelEngineCommon

  function TreeCallableFixedRateEngine{S}(model::S, timeSteps::Int)
    te = new{S}(model, timeSteps)

    add_observer!(model, te)

    return te
  end
end

TreeCallableFixedRateEngine{S <: ShortRateModel}(model::S, timeSteps::Int) = TreeCallableFixedRateEngine{S}(model, timeSteps)

function _calculate!(pe::TreeCallableFixedRateEngine, bond::CallableFixedRateBond)
  tsmodel = pe.model

  refDate = reference_date(tsmodel.ts)
  dc = tsmodel.ts.dc

  callableBond = DiscretizedCallableFixedRateBond(bond, refDate, dc)

  if isdefined(pe, :common)
    lattice = pe.common.lattice
  else
    times = mandatory_times(callableBond)
    tg = TimeGrid(times, pe.timeSteps)
    lattice = tree(pe.model, tg)
    pe.common = LatticeShortRateModelEngineCommon(tg, lattice)
  end

  redemptionTime = year_fraction(dc, refDate, callableBond.args.redemptionDate)

  initialize!(callableBond, lattice.treeLattice, redemptionTime)
  rollback!(callableBond, 0.0)
  bond.settlementValue = present_value(callableBond)

  return pe, bond
end