text:
module LookUpTables

using DocStringExtensions
using Adapt

export AbstractLookUp, LookUpLW, LookUpSW, LookUpCld

"""
    AbstractLookUp{I,FT}

Abstract lookup table for longwave and shortwave problems.
"""
abstract type AbstractLookUp{I,FT} end

"""
    LookUpLW{I,FT,UI8,UI8A1D,IA1D,IA2D,IA3D,FTA1D,FTA2D,FTA3D,FTA4D} <:
        AbstractLookUp{I,FT}

Longwave lookup tables, used to compute optical properties.

# Fields
$(DocStringExtensions.FIELDS)
"""
struct LookUpLW{
    I<:Int,
    FT<:AbstractFloat,
    UI8<:UInt8,
    UI8A1D<:AbstractArray{UI8,1},
    IA1D<:AbstractArray{I,1},
    IA2D<:AbstractArray{I,2},
    IA3D<:AbstractArray{I,3},
    FTA1D<:AbstractArray{FT,1},
    FTA2D<:AbstractArray{FT,2},
    FTA3D<:AbstractArray{FT,3},
    FTA4D<:AbstractArray{FT,4},
} <: AbstractLookUp{I,FT}
    "number of gases used in the lookup table"
    n_gases::I
    "number of longwave bands"
    n_bnd::I
    "number of `g-points`"
    n_gpt::I
    "number of atmospheric layers (=2, lower and upper atmospheres)"
    n_atmos_layers::I
    "number of reference temperatures for absorption coefficient lookup table"
    n_t_ref::I
    "number of reference pressures for absorption lookup table"
    n_p_ref::I
    "number of reference binary mixing fractions, for absorption coefficient lookup table"
    n_η::I
    "number of reference temperatures, for Planck source calculations"
    n_t_plnk::I
    "number of major absorbing gases"
    n_maj_absrb::I
    "number of minor absorbing gases"
    n_min_absrb::I
    "number of minor absorbers in lower atmosphere"
    n_min_absrb_lower::I
    "number of minor absorbers in upper atmosphere"
    n_min_absrb_upper::I
    n_absrb_ext::I # not used
    "number of minor contributors in the lower atmosphere"
    n_contrib_lower::I
    "number of minor contributors in the upper atmosphere"
    n_contrib_upper::I
    "vmr array index for h2o"
    idx_h2o::I
    "Reference pressure separating upper and lower atmosphere"
    p_ref_tropo::FT
    "Reference temperature"
    t_ref_absrb::FT
    "Reference pressure"
    p_ref_absrb::FT
    "minimum pressure supported by RRTMGP lookup tables"
    p_ref_min::FT
    "Δt for reference temperature values (Δt is constant)"
    Δ_t_ref::FT
    "Δ for log of reference pressure values (Δp is constant)"
    Δ_ln_p_ref::FT
    "indices for minor gases contributing to absorption in the lower atmosphere `(n_min_absrb_lower)`"
    idx_gases_minor_lower::IA1D
    "indices for minor gases contributing to absorption in the upper atmosphere `(n_min_absrb_upper)`"
    idx_gases_minor_upper::IA1D
    "indices for scaling gases in the lower atmosphere `(n_min_absrb_lower)`"
    idx_scaling_gas_lower::IA1D
    "indices for scaling gases in the upper atmosphere `(n_min_absrb_upper)`"
    idx_scaling_gas_upper::IA1D
    "major absorbing species in each band `(2, n_atmos_layers, n_bnd)`"
    key_species::IA3D
    "major absorption coefficient `(n_gpt, n_η, n_p_ref, n_t_ref)`"
    kmajor::FTA4D
    "minor absorption coefficient in lower atmosphere `(n_contrib_lower, n_η, n_t_ref)`"
    kminor_lower::FTA3D
    "minor absorption coefficient in upper atmosphere `(n_contrib_upper, n_η, n_t_ref)`"
    kminor_upper::FTA3D
    kminor_start_lower::IA1D # not currently used
    kminor_start_upper::IA1D # not currently used
    "Planck fraction `(n_gpt, n_η, n_p_ref, n_t_ref)`"
    planck_fraction::FTA4D
    "reference temperatures for Planck source calculations `(n_t_plnk)`"
    t_planck::FTA1D
    "total Planck source for each band `(n_t_plnk, n_bnd)`"
    totplnk::FTA2D
    "map from `g-point` to band"
    major_gpt2bnd::UI8A1D
    "starting and ending `g-point` for each band `(2, n_bnd)`"
    bnd_lims_gpt::IA2D
    "starting and ending wavenumber for each band `(2, n_bnd)`"
    bnd_lims_wn::FTA2D
    "`g-point` limits for minor contributors in lower atmosphere `(2, n_contrib_lower)`"
    minor_lower_gpt_lims::IA2D
    "`g-point` limits for minor contributors in upper atmosphere `(2, n_contrib_upper)`"
    minor_upper_gpt_lims::IA2D
    "band number for minor contributor in the lower atmosphere `(n_contrib_lower)`"
    minor_lower_bnd::UI8A1D
    "band number for minor contributor in the upper atmosphere `(n_contrib_upper)`"
    minor_upper_bnd::UI8A1D
    "starting index to `idx_gases_minor_lower` for each band `(n_bnd)`"
    minor_lower_bnd_st::UI8A1D
    "starting index to `idx_gases_minor_upper` for each band `(n_bnd)`"
    minor_upper_bnd_st::UI8A1D
    "shift in `kminor_lower` for each band `(n_min_absrb_lower)`"
    minor_lower_gpt_sh::IA1D
    "shift in `kminor_upper` for each band `(n_min_absrb_upper)`"
    minor_upper_gpt_sh::IA1D
    "minor gas (lower atmosphere) scales with density? `(n_min_absrb_lower)`"
    minor_lower_scales_with_density::IA1D
    "minor gas (upper atmosphere) scales with density? `(n_min_absrb_upper)`"
    minor_upper_scales_with_density::IA1D
    "minor gas (lower atmosphere) scales by compliment `(n_min_absrb_lower)`"
    lower_scale_by_complement::IA1D
    "minor gas (upper atmosphere) scales by compliment `(n_min_absrb_upper)`"
    upper_scale_by_complement::IA1D
    "reference pressures used by the lookup table `(n_p_ref)`"
    p_ref::FTA1D
    "reference temperatures used by the lookup table `(n_t_ref)`"
    t_ref::FTA1D
    "reference volume mixing ratios used by the lookup table `(2, n_gases, n_t_ref)`"
    vmr_ref::FTA3D
end
Adapt.@adapt_structure LookUpLW

function LookUpLW(
    ds,
    ::Type{I},
    ::Type{FT},
    ::Type{DA},
) where {I<:Int,FT<:AbstractFloat,DA}

    UI8 = UInt8
    UI8A1D = DA{UInt8,1}
    STA = Array{String,1}
    IA1D = DA{I,1}
    IA2D = DA{I,2}
    IA3D = DA{I,3}
    FTA1D = DA{FT,1}
    FTA2D = DA{FT,2}
    FTA3D = DA{FT,3}
    FTA4D = DA{FT,4}
    DSTAI = Dict{String,I}

    n_bnd = I(ds.dim["bnd"])
    n_gpt = I(ds.dim["gpt"])
    n_atmos_layers = I(ds.dim["atmos_layer"])
    n_t_ref = I(ds.dim["temperature"])
    n_p_ref = I(ds.dim["pressure"])
    n_t_plnk = I(ds.dim["temperature_Planck"])
    n_maj_absrb = I(ds.dim["absorber"])
    n_min_absrb = I(ds.dim["minor_absorber"])
    n_min_absrb_lower = I(ds.dim["minor_absorber_intervals_lower"])
    n_min_absrb_upper = I(ds.dim["minor_absorber_intervals_upper"])

    n_absrb_ext = I(ds.dim["absorber_ext"])
    n_contrib_lower = I(ds.dim["contributors_lower"])
    n_contrib_upper = I(ds.dim["contributors_upper"])

    p_ref_tropo = FT(ds["press_ref_trop"][:][1])
    t_ref_absrb = FT(ds["absorption_coefficient_ref_T"][:][1])
    p_ref_absrb = FT(ds["absorption_coefficient_ref_P"][:][1])

    gases_major = STA(undef, n_maj_absrb)
    gases_minor = STA(undef, n_min_absrb)
    id_minor = STA(undef, n_min_absrb)
    gases_minor_lower = STA(undef, n_min_absrb_lower)
    gases_minor_upper = STA(undef, n_min_absrb_upper)
    scaling_gas_lower = STA(undef, n_min_absrb_lower)
    scaling_gas_upper = STA(undef, n_min_absrb_upper)
    idx_gases = DSTAI()

    idx_gases_minor_lower = zeros(I, n_min_absrb_lower)
    idx_gases_minor_upper = zeros(I, n_min_absrb_upper)

    idx_scaling_gas_lower = zeros(I, n_min_absrb_lower)
    idx_scaling_gas_upper = zeros(I, n_min_absrb_upper)

    for igas = 1:n_maj_absrb
        gases_major[igas] = strip(String(ds["gas_names"][:, igas]))
        idx_gases[gases_major[igas]] = I(igas)
    end

    idx_h2o = idx_gases["h2o"]
    idx_gases["h2o_frgn"] = idx_h2o # water vapor - foreign
    idx_gases["h2o_self"] = idx_h2o # water vapor - self-continua

    n_gases = n_maj_absrb

    for igas = 1:n_min_absrb
        gases_minor[igas] = strip(String(ds["gas_minor"][:, igas]))
        id_minor[igas] = strip(String(ds["identifier_minor"][:, igas]))
    end

    for igas = 1:n_min_absrb_lower
        gases_minor_lower[igas] =
            strip(String(ds["minor_gases_lower"][:, igas]))
        scaling_gas_lower[igas] =
            strip(String(ds["scaling_gas_lower"][:, igas]))
        if ~isempty(gases_minor_lower[igas])
            idx_gases_minor_lower[igas] = idx_gases[gases_minor_lower[igas]]
        end
        if ~isempty(scaling_gas_lower[igas])
            idx_scaling_gas_lower[igas] = idx_gases[scaling_gas_lower[igas]]
        end
    end
    idx_gases_minor_lower = DA(idx_gases_minor_lower)
    idx_scaling_gas_lower = DA(idx_scaling_gas_lower)

    for igas = 1:n_min_absrb_upper
        gases_minor_upper[igas] =
            strip(String(ds["minor_gases_upper"][:, igas]))
        scaling_gas_upper[igas] =
            strip(String(ds["scaling_gas_upper"][:, igas]))
        if ~isempty(gases_minor_upper[igas])
            idx_gases_minor_upper[igas] = idx_gases[gases_minor_upper[igas]]
        end
        if ~isempty(scaling_gas_upper[igas])
            idx_scaling_gas_upper[igas] = idx_gases[scaling_gas_upper[igas]]
        end
    end
    idx_gases_minor_upper = DA(idx_gases_minor_upper)
    idx_scaling_gas_upper = DA(idx_scaling_gas_upper)

    key_species = ds["key_species"][:]
    for j = 1:size(key_species, 3)
        for i = 1:size(key_species, 2)
            if key_species[1, i, j] == 0 && key_species[2, i, j] == 0
                key_species[1:2, i, j] .= 2
            end
        end
    end

    key_species = IA3D(key_species)

    kmajor = FTA4D(ds["kmajor"][:])
    kminor_lower = FTA3D(ds["kminor_lower"][:])
    kminor_upper = FTA3D(ds["kminor_upper"][:])
    kminor_start_lower = IA1D(ds["kminor_start_lower"][:])
    kminor_start_upper = IA1D(ds["kminor_start_upper"][:])

    planck_fraction = FTA4D(ds["plank_fraction"][:])
    t_planck = FTA1D(ds["temperature_Planck"][:])

    totplnk = FTA2D(ds["totplnk"][:])
    bnd_lims_gpt = Array{I,2}(ds["bnd_limits_gpt"][:])
    bnd_lims_wn = FTA2D(ds["bnd_limits_wavenumber"][:])
    #-----------------------
    major_gpt2bnd = Array{UI8,1}(undef, n_gpt)
    for i = 1:n_bnd
        major_gpt2bnd[bnd_lims_gpt[1, i]:bnd_lims_gpt[2, i]] .= UI8(i)
    end
    #-----------------------
    bnd_lims_gpt = IA2D(bnd_lims_gpt)
    minor_lower_gpt_lims = Array{I,2}(ds["minor_limits_gpt_lower"][:])
    minor_upper_gpt_lims = Array{I,2}(ds["minor_limits_gpt_upper"][:])
    #-----------------------
    minor_lower_bnd = zeros(UI8, n_min_absrb_lower)
    minor_upper_bnd = zeros(UI8, n_min_absrb_upper)

    minor_lower_bnd_st = Array{UI8,1}(undef, n_bnd + 1)
    minor_upper_bnd_st = Array{UI8,1}(undef, n_bnd + 1)

    minor_lower_gpt_sh = Array{Int,1}(undef, n_min_absrb_lower)
    minor_upper_gpt_sh = Array{Int,1}(undef, n_min_absrb_upper)

    minor_lower_gpt_sh[1] = 0
    for i = 1:n_min_absrb_lower
        minor_lower_bnd[i] = major_gpt2bnd[minor_lower_gpt_lims[1, i]]
        if i > 1
            minor_lower_gpt_sh[i] =
                minor_lower_gpt_sh[i-1] + minor_lower_gpt_lims[2, i-1] -
                minor_lower_gpt_lims[1, i-1] + 1
        end
    end

    minor_upper_gpt_sh[1] = 0
    for i = 1:n_min_absrb_upper
        minor_upper_bnd[i] = major_gpt2bnd[minor_upper_gpt_lims[1, i]]
        if i > 1
            minor_upper_gpt_sh[i] =
                minor_upper_gpt_sh[i-1] + minor_upper_gpt_lims[2, i-1] -
                minor_upper_gpt_lims[1, i-1] + 1
        end
    end
    #-----------------------
    minor_lower_bnd_st[1] = 1
    minor_upper_bnd_st[1] = 1

    for ibnd = 2:n_bnd+1
        loc_low = findlast(isequal(UI8(ibnd - 1)), minor_lower_bnd)
        loc_upp = findlast(isequal(UI8(ibnd - 1)), minor_upper_bnd)
        if isnothing(loc_low)
            minor_lower_bnd_st[ibnd] = minor_lower_bnd_st[ibnd-1]
        else
            minor_lower_bnd_st[ibnd] = UI8(loc_low + 1)
        end

        if isnothing(loc_upp)
            minor_upper_bnd_st[ibnd] = minor_upper_bnd_st[ibnd-1]
        else
            minor_upper_bnd_st[ibnd] = UI8(loc_upp + 1)
        end
    end
    #-----------------------
    major_gpt2bnd = DA(major_gpt2bnd)
    minor_lower_bnd = DA(minor_lower_bnd)
    minor_upper_bnd = DA(minor_upper_bnd)
    minor_lower_bnd_st = DA(minor_lower_bnd_st)
    minor_upper_bnd_st = DA(minor_upper_bnd_st)
    minor_lower_gpt_sh = DA(minor_lower_gpt_sh)
    minor_upper_gpt_sh = DA(minor_upper_gpt_sh)
    minor_lower_gpt_lims = IA2D(minor_lower_gpt_lims)
    minor_upper_gpt_lims = IA2D(minor_upper_gpt_lims)
    #-----------------------

    minor_lower_scales_with_density =
        IA1D(ds["minor_scales_with_density_lower"][:])
    minor_upper_scales_with_density =
        IA1D(ds["minor_scales_with_density_upper"][:])

    lower_scale_by_complement = IA1D(ds["scale_by_complement_lower"][:])
    upper_scale_by_complement = IA1D(ds["scale_by_complement_upper"][:])

    p_ref = Array{FT,1}(ds["press_ref"][:])
    t_ref = Array{FT,1}(ds["temp_ref"][:])

    p_ref_min = minimum(p_ref)

    Δ_t_ref = t_ref[2] - t_ref[1]
    Δ_ln_p_ref = log(p_ref[1]) - log(p_ref[2])

    p_ref = FTA1D(p_ref)
    t_ref = FTA1D(t_ref)
    vmr_ref = FTA3D(ds["vmr_ref"][:])

    n_η = size(kmajor, 2)

    return (
        LookUpLW{I,FT,UI8,UI8A1D,IA1D,IA2D,IA3D,FTA1D,FTA2D,FTA3D,FTA4D}(
            n_gases,
            n_bnd,
            n_gpt,
            n_atmos_layers,
            n_t_ref,
            n_p_ref,
            n_η,
            n_t_plnk,
            n_maj_absrb,
            n_min_absrb,
            n_min_absrb_lower,
            n_min_absrb_upper,
            n_absrb_ext,
            n_contrib_lower,
            n_contrib_upper,
            idx_h2o,
            p_ref_tropo,
            t_ref_absrb,
            p_ref_absrb,
            p_ref_min,
            Δ_t_ref,
            Δ_ln_p_ref,
            idx_gases_minor_lower,
            idx_gases_minor_upper,
            idx_scaling_gas_lower,
            idx_scaling_gas_upper,
            key_species,
            kmajor,
            kminor_lower,
            kminor_upper,
            kminor_start_lower,
            kminor_start_upper,
            planck_fraction,
            t_planck,
            totplnk,
            major_gpt2bnd,
            bnd_lims_gpt,
            bnd_lims_wn,
            minor_lower_gpt_lims,
            minor_upper_gpt_lims,
            minor_lower_bnd,
            minor_upper_bnd,
            minor_lower_bnd_st,
            minor_upper_bnd_st,
            minor_lower_gpt_sh,
            minor_upper_gpt_sh,
            minor_lower_scales_with_density,
            minor_upper_scales_with_density,
            lower_scale_by_complement,
            upper_scale_by_complement,
            p_ref,
            t_ref,
            vmr_ref,
        ),
        idx_gases,
    )
end

"""
    LookUpSW{I,FT,UI8,UI8A1D,IA1D,IA2D,IA3D,FTA1D,FTA2D,FTA3D,FTA4D} <:
        AbstractLookUp{I,FT}

Shortwave lookup tables, used to compute optical properties.

# Fields
$(DocStringExtensions.FIELDS)
"""
struct LookUpSW{
    I<:Int,
    FT<:AbstractFloat,
    UI8<:UInt8,
    UI8A1D<:AbstractArray{UI8,1},
    IA1D<:AbstractArray{I,1},
    IA2D<:AbstractArray{I,2},
    IA3D<:AbstractArray{I,3},
    FTA1D<:AbstractArray{FT,1},
    FTA2D<:AbstractArray{FT,2},
    FTA3D<:AbstractArray{FT,3},
    FTA4D<:AbstractArray{FT,4},
} <: AbstractLookUp{I,FT}
    "number of gases used in the lookup table"
    n_gases::I
    "number of shortwave bands"
    n_bnd::I
    "number of `g-points`"
    n_gpt::I
    "number of atmospheric layers (=2, lower and upper atmospheres)"
    n_atmos_layers::I
    "number of reference temperatures for absorption coefficient lookup table"
    n_t_ref::I
    "number of reference pressures for absorption lookup table"
    n_p_ref::I
    "number of reference binary mixing fractions, for absorption coefficient lookup table"
    n_η::I
    "number of major absorbing gases"
    n_maj_absrb::I
    "number of minor absorbing gases"
    n_min_absrb::I
    "number of minor absorbers in lower atmosphere"
    n_min_absrb_lower::I
    "number of minor absorbers in upper atmosphere"
    n_min_absrb_upper::I
    n_absrb_ext::I # not used
    "number of minor contributors in the lower atmosphere"
    n_contrib_lower::I
    "number of minor contributors in the upper atmosphere"
    n_contrib_upper::I
    "vmr array index for h2o"
    idx_h2o::I
    "Reference pressure separating upper and lower atmosphere"
    p_ref_tropo::FT
    "Reference temperature"
    t_ref_absrb::FT
    "Reference pressure"
    p_ref_absrb::FT
    "minimum pressure supported by RRTMGP lookup tables"
    p_ref_min::FT
    "Δt for reference temperature values (Δt is constant, array is linearly equispaced)"
    Δ_t_ref::FT
    "Δ for log of reference pressure values (Δ is constant, equispaced on logarithmic scale)"
    Δ_ln_p_ref::FT
    "total solar irradiation"
    solar_src_tot::FT
    "indices for minor gases contributing to absorption in the lower atmosphere `(n_min_absrb_lower)`"
    idx_gases_minor_lower::IA1D
    "indices for minor gases contributing to absorption in the upper atmosphere `(n_min_absrb_upper)`"
    idx_gases_minor_upper::IA1D
    "indices for scaling gases in the lower atmosphere `(n_min_absrb_lower)`"
    idx_scaling_gas_lower::IA1D
    "indices for scaling gases in the upper atmosphere `(n_min_absrb_upper)`"
    idx_scaling_gas_upper::IA1D
    "major absorbing species in each band `(2, n_atmos_layers, n_bnd)`"
    key_species::IA3D
    "major absorption coefficient `(n_gpt, n_η, n_p_ref, n_t_ref)`"
    kmajor::FTA4D
    "minor absorption coefficient in lower atmosphere `(n_contrib_lower, n_η, n_t_ref)`"
    kminor_lower::FTA3D
    "minor absorption coefficient in upper atmosphere `(n_contrib_upper, n_η, n_t_ref)`"
    kminor_upper::FTA3D
    kminor_start_lower::IA1D # not currently used
    kminor_start_upper::IA1D # not currently used
    "map from g-point to band"
    major_gpt2bnd::UI8A1D
    "starting and ending `g-point` for each band `(2, n_bnd)`"
    bnd_lims_gpt::IA2D
    "starting and ending wavenumber for each band `(2, n_bnd)`"
    bnd_lims_wn::FTA2D
    "`g-point` limits for minor contributors in lower atmosphere `(2, n_contrib_lower)`"
    minor_lower_gpt_lims::IA2D
    "`g-point` limits for minor contributors in upper atmosphere `(2, n_contrib_upper)`"
    minor_upper_gpt_lims::IA2D
    "band number for minor contributor in the lower atmosphere `(n_contrib_lower)`"
    minor_lower_bnd::UI8A1D
    "band number for minor contributor in the upper atmosphere `(n_contrib_upper)`"
    minor_upper_bnd::UI8A1D
    "starting index to `idx_gases_minor_lower` for each band"
    minor_lower_bnd_st::UI8A1D
    "starting index to `idx_gases_minor_upper` for each band"
    minor_upper_bnd_st::UI8A1D
    "shift in `kminor_lower` for each band `(n_min_absrb_lower)`"
    minor_lower_gpt_sh::IA1D
    "shift in `kminor_upper` for each band `(n_min_absrb_upper)`"
    minor_upper_gpt_sh::IA1D
    "minor gas (lower atmosphere) scales with density? `(n_min_absrb_lower)`"
    minor_lower_scales_with_density::IA1D
    "minor gas (upper atmosphere) scales with density? `(n_min_absrb_upper)`"
    minor_upper_scales_with_density::IA1D
    "minor gas (lower atmosphere) scales by compliment `(n_min_absrb_lower)`"
    lower_scale_by_complement::IA1D
    "minor gas (upper atmosphere) scales by compliment `(n_min_absrb_upper)`"
    upper_scale_by_complement::IA1D
    "reference pressures used by the lookup table `(n_p_ref)`"
    p_ref::FTA1D
    "reference temperatures used by the lookup table `(n_t_ref)`"
    t_ref::FTA1D
    "reference volume mixing ratios used by the lookup table `(2, n_gases, n_t_ref)`"
    vmr_ref::FTA3D
    "Rayleigh absorption coefficient for lower atmosphere `(n_gpt, n_η, n_t_ref)`"
    rayl_lower::FTA3D
    "Rayleigh absorption coefficient for upper atmosphere `(n_gpt, n_η, n_t_ref)`"
    rayl_upper::FTA3D
    "relative solar source contribution from each `g-point` `(n_gpt)`"
    solar_src_scaled::FTA1D
end
Adapt.@adapt_structure LookUpSW

function LookUpSW(
    ds,
    ::Type{I},
    ::Type{FT},
    ::Type{DA},
) where {I<:Int,FT<:AbstractFloat,DA}

    UI8 = UInt8
    UI8A1D = DA{UInt8,1}
    STA = Array{String,1}
    IA1D = DA{I,1}
    IA2D = DA{I,2}
    IA3D = DA{I,3}
    FTA1D = DA{FT,1}
    FTA2D = DA{FT,2}
    FTA3D = DA{FT,3}
    FTA4D = DA{FT,4}
    DSTAI = Dict{String,I}

    n_bnd = I(ds.dim["bnd"])
    n_gpt = I(ds.dim["gpt"])
    n_atmos_layers = I(ds.dim["atmos_layer"])
    n_t_ref = I(ds.dim["temperature"])
    n_p_ref = I(ds.dim["pressure"])

    n_maj_absrb = I(ds.dim["absorber"])
    n_min_absrb = I(ds.dim["minor_absorber"])

    n_min_absrb_lower = I(ds.dim["minor_absorber_intervals_lower"])
    n_min_absrb_upper = I(ds.dim["minor_absorber_intervals_upper"])
    n_absrb_ext = I(ds.dim["absorber_ext"])

    n_contrib_lower = I(ds.dim["contributors_lower"])
    n_contrib_upper = I(ds.dim["contributors_upper"])

    p_ref_tropo = FT(ds["press_ref_trop"][:][1])
    t_ref_absrb = FT(ds["absorption_coefficient_ref_T"][:][1])
    p_ref_absrb = FT(ds["absorption_coefficient_ref_P"][:][1])

    gases_major = STA(undef, n_maj_absrb)
    gases_minor = STA(undef, n_min_absrb)
    id_minor = STA(undef, n_min_absrb)
    gases_minor_lower = STA(undef, n_min_absrb_lower)
    gases_minor_upper = STA(undef, n_min_absrb_upper)
    scaling_gas_lower = STA(undef, n_min_absrb_lower)
    scaling_gas_upper = STA(undef, n_min_absrb_upper)
    idx_gases = DSTAI()

    idx_gases_minor_lower = zeros(I, n_min_absrb_lower)
    idx_gases_minor_upper = zeros(I, n_min_absrb_upper)

    idx_scaling_gas_lower = zeros(I, n_min_absrb_lower)
    idx_scaling_gas_upper = zeros(I, n_min_absrb_upper)

    for igas = 1:n_maj_absrb
        gases_major[igas] = strip(String(ds["gas_names"][:, igas]))
        idx_gases[gases_major[igas]] = I(igas)
    end

    idx_h2o = idx_gases["h2o"]
    idx_gases["h2o_frgn"] = idx_h2o # water vapor - foreign
    idx_gases["h2o_self"] = idx_h2o # water vapor - self-continua

    n_gases = n_maj_absrb

    for igas = 1:n_min_absrb
        gases_minor[igas] = strip(String(ds["gas_minor"][:, igas]))
        id_minor[igas] = strip(String(ds["identifier_minor"][:, igas]))
    end

    for igas = 1:n_min_absrb_lower
        gases_minor_lower[igas] =
            strip(String(ds["minor_gases_lower"][:, igas]))
        scaling_gas_lower[igas] =
            strip(String(ds["scaling_gas_lower"][:, igas]))
        if ~isempty(gases_minor_lower[igas])
            idx_gases_minor_lower[igas] = idx_gases[gases_minor_lower[igas]]
        end
        if ~isempty(scaling_gas_lower[igas])
            idx_scaling_gas_lower[igas] = idx_gases[scaling_gas_lower[igas]]
        end
    end

    for igas = 1:n_min_absrb_upper
        gases_minor_upper[igas] =
            strip(String(ds["minor_gases_upper"][:, igas]))
        scaling_gas_upper[igas] =
            strip(String(ds["scaling_gas_upper"][:, igas]))
        if ~isempty(gases_minor_upper[igas])
            idx_gases_minor_upper[igas] = idx_gases[gases_minor_upper[igas]]
        end
        if ~isempty(scaling_gas_upper[igas])
            idx_scaling_gas_upper[igas] = idx_gases[scaling_gas_upper[igas]]
        end
    end

    key_species = ds["key_species"][:]
    for j = 1:size(key_species, 3)
        for i = 1:size(key_species, 2)
            if key_species[1, i, j] == 0 && key_species[2, i, j] == 0
                key_species[1:2, i, j] .= 2
            end
        end
    end

    key_species = IA3D(key_species)

    kmajor = FTA4D(ds["kmajor"][:])
    kminor_lower = FTA3D(ds["kminor_lower"][:])
    kminor_upper = FTA3D(ds["kminor_upper"][:])
    kminor_start_lower = IA1D(ds["kminor_start_lower"][:])
    kminor_start_upper = IA1D(ds["kminor_start_upper"][:])

    bnd_lims_gpt = Array{I,2}(ds["bnd_limits_gpt"][:])
    bnd_lims_wn = FTA2D(ds["bnd_limits_wavenumber"][:])
    #-----------------------
    major_gpt2bnd = Array{UI8,1}(undef, n_gpt)
    for i = 1:n_bnd
        major_gpt2bnd[bnd_lims_gpt[1, i]:bnd_lims_gpt[2, i]] .= UI8(i)
    end
    #-----------------------
    bnd_lims_gpt = IA2D(bnd_lims_gpt)
    minor_lower_gpt_lims = Array{I,2}(ds["minor_limits_gpt_lower"][:])
    minor_upper_gpt_lims = Array{I,2}(ds["minor_limits_gpt_upper"][:])
    #-----------------------
    minor_lower_bnd = zeros(UI8, n_min_absrb_lower)
    minor_upper_bnd = zeros(UI8, n_min_absrb_upper)

    minor_lower_bnd_st = Array{UI8,1}(undef, n_bnd + 1)
    minor_upper_bnd_st = Array{UI8,1}(undef, n_bnd + 1)

    minor_lower_gpt_sh = Array{Int,1}(undef, n_min_absrb_lower)
    minor_upper_gpt_sh = Array{Int,1}(undef, n_min_absrb_upper)

    minor_lower_gpt_sh[1] = 0
    for i = 1:n_min_absrb_lower
        minor_lower_bnd[i] = major_gpt2bnd[minor_lower_gpt_lims[1, i]]
        if i > 1
            minor_lower_gpt_sh[i] =
                minor_lower_gpt_sh[i-1] + minor_lower_gpt_lims[2, i-1] -
                minor_lower_gpt_lims[1, i-1] + 1
        end
    end

    minor_upper_gpt_sh[1] = 0
    for i = 1:n_min_absrb_upper
        minor_upper_bnd[i] = major_gpt2bnd[minor_upper_gpt_lims[1, i]]
        if i > 1
            minor_upper_gpt_sh[i] =
                minor_upper_gpt_sh[i-1] + minor_upper_gpt_lims[2, i-1] -
                minor_upper_gpt_lims[1, i-1] + 1
        end
    end
    #-----------------------
    minor_lower_bnd_st[1] = 1
    minor_upper_bnd_st[1] = 1

    for ibnd = 2:n_bnd+1
        loc_low = findlast(isequal(UI8(ibnd - 1)), minor_lower_bnd)
        loc_upp = findlast(isequal(UI8(ibnd - 1)), minor_upper_bnd)
        if isnothing(loc_low)
            minor_lower_bnd_st[ibnd] = minor_lower_bnd_st[ibnd-1]
        else
            minor_lower_bnd_st[ibnd] = UI8(loc_low + 1)
        end

        if isnothing(loc_upp)
            minor_upper_bnd_st[ibnd] = minor_upper_bnd_st[ibnd-1]
        else
            minor_upper_bnd_st[ibnd] = UI8(loc_upp + 1)
        end
    end
    #------------------------
    major_gpt2bnd = DA(major_gpt2bnd)
    minor_lower_bnd = DA(minor_lower_bnd)
    minor_upper_bnd = DA(minor_upper_bnd)
    minor_lower_bnd_st = DA(minor_lower_bnd_st)
    minor_upper_bnd_st = DA(minor_upper_bnd_st)
    minor_lower_gpt_sh = DA(minor_lower_gpt_sh)
    minor_upper_gpt_sh = DA(minor_upper_gpt_sh)
    minor_lower_gpt_lims = IA2D(minor_lower_gpt_lims)
    minor_upper_gpt_lims = IA2D(minor_upper_gpt_lims)
    #-----------------------
    minor_lower_scales_with_density =
        IA1D(ds["minor_scales_with_density_lower"][:])
    minor_upper_scales_with_density =
        IA1D(ds["minor_scales_with_density_upper"][:])
    lower_scale_by_complement = IA1D(ds["scale_by_complement_lower"][:])
    upper_scale_by_complement = IA1D(ds["scale_by_complement_upper"][:])

    p_ref = Array{FT,1}(ds["press_ref"][:])
    t_ref = Array{FT,1}(ds["temp_ref"][:])

    p_ref_min = minimum(p_ref)

    Δ_t_ref = t_ref[2] - t_ref[1]
    Δ_ln_p_ref = log(p_ref[1]) - log(p_ref[2])

    p_ref = FTA1D(p_ref)
    t_ref = FTA1D(t_ref)
    vmr_ref = FTA3D(ds["vmr_ref"][:])

    n_η = size(kmajor, 2)

    rayl_lower = FTA3D(ds["rayl_lower"][:])
    rayl_upper = FTA3D(ds["rayl_upper"][:])

    solar_src = ds["solar_source"][:]
    solar_src_tot = FT(sum(solar_src))
    solar_src_scaled = FTA1D(solar_src ./ solar_src_tot)

    return (
        LookUpSW{I,FT,UI8,UI8A1D,IA1D,IA2D,IA3D,FTA1D,FTA2D,FTA3D,FTA4D}(
            n_gases,
            n_bnd,
            n_gpt,
            n_atmos_layers,
            n_t_ref,
            n_p_ref,
            n_η,
            n_maj_absrb,
            n_min_absrb,
            n_min_absrb_lower,
            n_min_absrb_upper,
            n_absrb_ext,
            n_contrib_lower,
            n_contrib_upper,
            idx_h2o,
            p_ref_tropo,
            t_ref_absrb,
            p_ref_absrb,
            p_ref_min,
            Δ_t_ref,
            Δ_ln_p_ref,
            solar_src_tot,
            idx_gases_minor_lower,
            idx_gases_minor_upper,
            idx_scaling_gas_lower,
            idx_scaling_gas_upper,
            key_species,
            kmajor,
            kminor_lower,
            kminor_upper,
            kminor_start_lower,
            kminor_start_upper,
            major_gpt2bnd,
            bnd_lims_gpt,
            bnd_lims_wn,
            minor_lower_gpt_lims,
            minor_upper_gpt_lims,
            minor_lower_bnd,
            minor_upper_bnd,
            minor_lower_bnd_st,
            minor_upper_bnd_st,
            minor_lower_gpt_sh,
            minor_upper_gpt_sh,
            minor_lower_scales_with_density,
            minor_upper_scales_with_density,
            lower_scale_by_complement,
            upper_scale_by_complement,
            p_ref,
            t_ref,
            vmr_ref,
            rayl_lower,
            rayl_upper,
            solar_src_scaled,
        ),
        idx_gases,
    )
end

"""
    LookUpCld{I,FT,FTA1D,FTA2D,FTA3D,FTA4D}

Lookup table for cloud optics.

This struct stores the tables and Pade coefficients for determing extinction coeffient,
single-scattering albedo, and asymmetry parameter g as a function of effective radius.
We compute the optical depth tau (=exintinction coeff * condensed water path)
and the products tau*ssa and tau*ssa*g for liquid and ice cloud separately.
These are used to determine the optical properties of ice and water cloud together.

# Fields
$(DocStringExtensions.FIELDS)
"""
struct LookUpCld{I,B,FT,FTA1D,FTA2D,FTA3D,FTA4D}
    "number of bands"
    nband::I
    "number of ice roughness types"
    nrghice::I
    "number of liquid particle sizes"
    nsize_liq::I
    "number of ice particle sizes"
    nsize_ice::I
    "number of size regimes"
    nsizereg::I
    "number of extinction coefficients for pade approximation"
    ncoeff_ext::I
    "number of ssa/g coefficients for pade approximation"
    ncoeff_ssa_g::I
    "number of size regime boundaries for pade interpolation"
    nbound::I
    "pair = 2"
    pair::I
    "liquid particle size lower bound for LUT interpolation"
    radliq_lwr::FT
    "liquid particle size upper bound for LUT interpolation"
    radliq_upr::FT
    "factor for calculating LUT interpolation for liquid particle"
    radliq_fac::FT
    "ice particle size lower bound for LUT interpolation"
    radice_lwr::FT
    "ice particle size upper bound for LUT interpolation"
    radice_upr::FT
    "factor for calculating LUT interpolation for ice particle"
    radice_fac::FT
    "LUT liquid extinction coefficient (`nsize_liq, nbnd`) m²/g"
    lut_extliq::FTA2D
    "LUT liquid single scattering albedo (`nsize_liq, nbnd`)"
    lut_ssaliq::FTA2D
    "LUT liquid asymmetry parameter (`nsize_liq, nbnd`)"
    lut_asyliq::FTA2D
    "LUT ice extinction coefficient (`nsize_ice, nband, nrghice`) m²/g"
    lut_extice::FTA3D
    "LUT ice single scattering albedo (`nsize_ice, nband, nrghice`)"
    lut_ssaice::FTA3D
    "LUT ice asymmetry parameter (`nsize_ice, nband, nrghice`)"
    lut_asyice::FTA3D
    "pade coefficients for liquid extinction (`nband, nsizereg, ncoeff_ext`)"
    pade_extliq::FTA3D
    "pade coefficients for liquid single scattening albedo (`nband, nsizereg, ncoeff_ssa_g`)"
    pade_ssaliq::FTA3D
    "pade coefficients for liquid asymmetry paramter (`nband, nsizereg, ncoeff_ssa_g`)"
    pade_asyliq::FTA3D
    "pade coefficients for ice extinction (`nband, nsizereg, ncoeff_ext, nrghice`)"
    pade_extice::FTA4D
    "pade coefficients for ice single scattening albedo (`nband, nsizereg, ncoeff_ssa_g, nrghice`)"
    pade_ssaice::FTA4D
    "pade coefficients for ice asymmetry paramter (`nband, nsizereg, ncoeff_ssa_g, nrghice`)"
    pade_asyice::FTA4D
    "pade size regime boundaries for liquid extinction coefficient for pade interpolation (`nbound`) μm"
    pade_sizreg_extliq::FTA1D
    "pade size regime boundaries for liquid single scattering albedo for pade interpolation (`nbound`) μm"
    pade_sizreg_ssaliq::FTA1D
    "pade size regime boundaries for liquid asymmetry parameter for pade interpolation (`nbound`) μm"
    pade_sizreg_asyliq::FTA1D
    "pade size regime boundaries for ice extinction coefficient for pade interpolation (`nbound`) μm"
    pade_sizreg_extice::FTA1D
    "pade size regime boundaries for ice single scattering albedo for pade interpolation (`nbound`) μm"
    pade_sizreg_ssaice::FTA1D
    "pade size regime boundaries for ice asymmetry parameter for pade interpolation (`nbound`) μm"
    pade_sizreg_asyice::FTA1D
    "beginning and ending wavenumber for each band (`2, nband`) cm⁻¹"
    bnd_lims_wn::FTA2D
    "use LUT (default) or pade interpolation"
    use_lut::B
end
Adapt.@adapt_structure LookUpCld

function LookUpCld(
    ds,
    ::Type{I},
    ::Type{FT},
    ::Type{DA},
    use_lut::Bool = true,
) where {I<:Int,FT<:AbstractFloat,DA}

    FTA1D = DA{FT,1}
    FTA2D = DA{FT,2}
    FTA3D = DA{FT,3}
    FTA4D = DA{FT,4}

    nband = I(ds.dim["nband"])
    nrghice = I(ds.dim["nrghice"])
    nsize_liq = I(ds.dim["nsize_liq"])
    nsize_ice = I(ds.dim["nsize_ice"])
    nsizereg = I(ds.dim["nsizereg"])
    ncoeff_ext = I(ds.dim["ncoeff_ext"])
    ncoeff_ssa_g = I(ds.dim["ncoeff_ssa_g"])
    nbound = I(ds.dim["nbound"])
    pair = I(ds.dim["pair"])

    @assert nsizereg == 3 # RRTMGP pade approximation assumes exactly (3) size regimes

    radliq_lwr = FT(ds["radliq_lwr"][:])
    radliq_upr = FT(ds["radliq_upr"][:])
    radliq_fac = FT(ds["radliq_fac"][:])

    radice_lwr = FT(ds["radice_lwr"][:])
    radice_upr = FT(ds["radice_upr"][:])
    radice_fac = FT(ds["radice_fac"][:])

    lut_extliq = FTA2D(ds["lut_extliq"][:])
    lut_ssaliq = FTA2D(ds["lut_ssaliq"][:])
    lut_asyliq = FTA2D(ds["lut_asyliq"][:])

    lut_extice = FTA3D(ds["lut_extice"][:])
    lut_ssaice = FTA3D(ds["lut_ssaice"][:])
    lut_asyice = FTA3D(ds["lut_asyice"][:])

    pade_extliq = FTA3D(ds["pade_extliq"][:])
    pade_ssaliq = FTA3D(ds["pade_ssaliq"][:])
    pade_asyliq = FTA3D(ds["pade_asyliq"][:])

    pade_extice = FTA4D(ds["pade_extice"][:])
    pade_ssaice = FTA4D(ds["pade_ssaice"][:])
    pade_asyice = FTA4D(ds["pade_asyice"][:])

    pade_sizreg_extliq = FTA1D(ds["pade_sizreg_extliq"][:])
    pade_sizreg_ssaliq = FTA1D(ds["pade_sizreg_ssaliq"][:])
    pade_sizreg_asyliq = FTA1D(ds["pade_sizreg_asyliq"][:])

    pade_sizreg_extice = FTA1D(ds["pade_sizreg_extice"][:])
    pade_sizreg_ssaice = FTA1D(ds["pade_sizreg_ssaice"][:])
    pade_sizreg_asyice = FTA1D(ds["pade_sizreg_asyice"][:])

    bnd_lims_wn = FTA2D(ds["bnd_limits_wavenumber"][:])

    return (LookUpCld{I,Bool,FT,FTA1D,FTA2D,FTA3D,FTA4D}(
        nband,
        nrghice,
        nsize_liq,
        nsize_ice,
        nsizereg,
        ncoeff_ext,
        ncoeff_ssa_g,
        nbound,
        pair,
        radliq_lwr,
        radliq_upr,
        radliq_fac,
        radice_lwr,
        radice_upr,
        radice_fac,
        lut_extliq,
        lut_ssaliq,
        lut_asyliq,
        lut_extice,
        lut_ssaice,
        lut_asyice,
        pade_extliq,
        pade_ssaliq,
        pade_asyliq,
        pade_extice,
        pade_ssaice,
        pade_asyice,
        pade_sizreg_extliq,
        pade_sizreg_ssaliq,
        pade_sizreg_asyliq,
        pade_sizreg_extice,
        pade_sizreg_ssaice,
        pade_sizreg_asyice,
        bnd_lims_wn,
        use_lut,
    ))

end

end


text:
"""
https://github.com/JuliaLang/julia/blob/master/src/julia-parser.scm

"""
s = "= += -= *= /= //= |\\=| ^= ÷= %= <<= >>= >>>= ||=| &= ⊻= ≔ ⩴ ≕ ~ := \$= => "

bodmas(self::Node)::Int64 = bodmas(class(self))

bodmas(self::Node{<:AAnd})::Int64 = 10
bodmas(self::Node{<:AOr})::Int64 = 11

ext: jl


text:
struct SmallSignalOutput
    reduced_jacobian::Matrix{Float64}
    eigenvalues::Vector{Complex{Float64}}
    eigenvectors::Matrix{Complex{Float64}}
    stable::Bool
    operating_point::Vector{Float64}
end


text:
@testset "SpectralDense" begin
    input = rand(5, 32)
    layer = SpectralDense(5, 10)
    output = layer(input)
    @test size(output) == (10, 32)

    # Test gradients
    layer = SpectralDense(2, 5)
    i = rand(2, 4)
    y = ones(5, 4)
    grads = gradient(params(layer)) do
        ŷ = layer(i)
        return Flux.logitcrossentropy(ŷ, y)
    end
    for param in params(layer)
        @test size(param) == size(grads[param])
    end
end

@testset "SpectralConv" begin
    # Test gradients
    layer = SpectralConv((5, 5), 3 => 6, relu)
    i = rand(32, 32, 3, 4)
    y = rand(28, 28, 6, 4)
    grads = gradient(params(layer)) do
        ŷ = layer(i)
        return Flux.logitcrossentropy(ŷ, y)
    end
    for param in params(layer)
        @test size(param) == size(grads[param])
    end
end


text:
using POMDPModelTools
using POMDPs
using POMDPModels
using Distributions
using Random
using Test
using Pkg
using POMDPSimulators

@testset "ordered" begin
    include("test_ordered_spaces.jl")
end

# require POMDPModels
@testset "genbeliefmdp" begin
    include("test_generative_belief_mdp.jl")
end
@testset "implement" begin
    include("test_implementations.jl")
end
@testset "distribjl" begin
    include("test_distributions_jl.jl")
end
@testset "weightediter" begin
    include("test_weighted_iteration.jl")
end
@testset "sparsecat" begin
    include("test_sparse_cat.jl")
end
@testset "bool" begin
    include("test_bool.jl")
end
@testset "deterministic" begin
    include("test_deterministic.jl")
end
@testset "terminalstate" begin
    include("test_terminal_state.jl")
end

# require POMDPModels
@testset "info" begin
    include("test_info.jl")
end
@testset "obsweight" begin
    include("test_obs_weight.jl")
end

# require DiscreteValueIteration
@testset "visolve" begin
    POMDPs.add_registry()
    Pkg.add("DiscreteValueIteration")
    using DiscreteValueIteration
    include("test_fully_observable_pomdp.jl")
    include("test_underlying_mdp.jl")
end

@testset "vis" begin
    include("test_visualization.jl")
end


text:
using Sockets
using Serialization
using REPL

function send_header(io, ser_version=Serialization.ser_version)
    write(io, protocol_magic, protocol_version)
    write(io, UInt32(ser_version))
    flush(io)
end

# Format result of arbitrary type as a string for transmission.
# Stringifying everything may seem strange, but is beneficial because
#   * It allows us to show the user-defined types which exist only on the
#     remote server
#   * It allows us to limit the amount of data transmitted (eg, large arrays
#     are truncated when the :limit IOContext property is set)
function format_result(f, display_properties)
    io = IOBuffer()
    ctx = IOContext(io, display_properties...)
    f(ctx)
    String(take!(io))
end

# Serve a remote REPL session to a single client over `socket`.
function serve_repl_session(socket)
    send_header(socket)
    display_properties = Dict()
    while isopen(socket)
        request = deserialize(socket)
        response = nothing
        try
            @debug "Client command" request
            messageid,value = request isa Tuple && length(request) == 2 ?
                            request : (nothing,nothing)
            if messageid == :eval
                result = Main.eval(value)
                resultval = isnothing(result) ? nothing :
                    format_result(display_properties) do io
                        show(io, MIME"text/plain"(), result)
                    end
                response = (:eval_result, resultval)
            elseif messageid == :display_properties
                @debug "Got client display properties" display_properties
                display_properties = value::Dict
            elseif messageid == :repl_completion
                # See REPL.jl complete_line(c::REPLCompletionProvider, s::PromptState)
                partial, full = value
                ret, range, should_complete = REPL.completions(full, lastindex(partial))
                result = (unique!(map(REPL.completion_text, ret)),
                          partial[range], should_complete)
                response = (:completion_result, result)
            elseif messageid == :exit
                break
            end
        catch _
            resultval = format_result(display_properties) do io
                Base.display_error(io, Base.catch_stack())
            end
            response = (:error, resultval)
        end
        if !isnothing(response)
            serialize(socket, response)
        end
    end
end

"""
    serve_repl([address=Sockets.localhost,] port=27754)
    serve_repl(server)

Start a REPL server listening on interface `address` and `port`. In normal
operation `serve_repl()` serves REPL clients indefinitely (ie., it does not
return), so you will generally want to launch it using `@async serve_repl()` to
do other useful work at the same time.

If you want to be able to stop the server you can pass an already-listening
`server` object (the result of `Sockets.listen()`). The server can then be
cancelled from another task using `close(server)` as necessary to control the
server lifetime.

## Security

`serve_repl()` uses an *unauthenticated, unecrypted protocol* so it should not
be used on open networks or multi-user machines where other users aren't
trused. For open networks, use the default `address=Sockets.localhost` and the
automatic ssh tunnel support provided by the client-side `connect_repl()`.
"""
function serve_repl(address=Sockets.localhost, port::Integer=27754)
    server = listen(address, port)
    try
        serve_repl(server)
    finally
        close(server)
    end
end
serve_repl(port::Integer) = serve_repl(Sockets.localhost, port)

function serve_repl(server::Base.IOServer)
    open_sockets = Set()
    @sync try
        while isopen(server)
            socket = accept(server)
            push!(open_sockets, socket)
            peer=getpeername(socket)
            @info "REPL client opened a connection" peer
            @async try
                serve_repl_session(socket)
            catch exc
                if !(exc isa EOFError && !isopen(socket))
                    @warn "Something went wrong evaluating client command" #=
                        =# exception=exc,catch_backtrace()
                end
            finally
                @info "REPL client exited" peer
                close(socket)
                pop!(open_sockets, socket)
            end
        end
    catch exc
        if exc isa Base.IOError && !isopen(server)
            # Ok - server was closed
            return
        end
        @error "Unexpected server failure" isopen(server) exception=exc,catch_backtrace()
        rethrow()
    finally
        for socket in open_sockets
            close(socket)
        end
    end
end


text: divide(weight) = weight > 3 && iseven(weight-2)

max_stars_repo_name: UlrichBerntien/Codewars-Katas

max_issues_repo_name: UlrichBerntien/Codewars-Katas

max_forks_repo_name: UlrichBerntien/Codewars-Katas


text:
@testset "GaussianLikelihood" begin
    lik = GaussianLikelihood(1e-5)
    test_interface(lik, Normal; functor_args=(:σ²,))
end

@testset "HeteroscedasticGaussianLikelihood" begin
    for args in ((), (exp,), (ExpLink(),))
        lik = HeteroscedasticGaussianLikelihood(args...)
        @test lik isa HeteroscedasticGaussianLikelihood{ExpLink}
    end

    lik = HeteroscedasticGaussianLikelihood()
    IN_DIM = 3
    OUT_DIM = 2 # one for the mean the other for the log-standard deviation
    N = 10
    test_interface(lik, Normal, 2)
end


text:
using DiscreteChoice
using JLD2
using ForwardDiff
using Optim

@load "example/data/homo.jld2" Y X β ξ

model = DCModel(Y,X; hetero_preference=false)

estDCModel!(model; init_guess = vcat(β, ξ))

res = Optim.minimizer(model.optResults) |> copy
res .- vcat(β, ξ) |> maximum


text:
@doc raw"""
`SODE`: Split Ordinary Differential Equation

Defines an initial value problem
```math
\dot{q} (t) = v(t, q(t)) , \qquad q(t_{0}) = q_{0} ,
```
with vector field ``v``, initial condition ``q_{0}`` and the solution
``q`` taking values in ``\mathbb{R}^{d}``. Here, the vector field ``v``
is given as a sum of vector fields
```math
v (t) = v_1 (t) + ... + v_r (t) .
```

### Parameters

* `vType <: Union{Tuple,Nothing}`: type of `v`
* `qType <: Union{Tuple,Nothing}`: type of `q`
* `invType <: OptionalNamedTuple`: invariants type
* `parType <: OptionalNamedTuple`: parameters type
* `perType <: OptionalArray{AT}`: periodicity type

### Fields

* `v`: tuple of functions computing the vector field
* `q`: tuple of functions computing the solution
* `invariants`: either a `NamedTuple` containing the equation's invariants or `nothing`
* `parameters`: either a `NamedTuple` containing the equation's parameters or `nothing`
* `periodicity`: determines the periodicity of the state vector `q` for cutting periodic solutions

The functions `v_i` providing the vector field must have the interface
```julia
    function v_i(t, q, v, params)
        v[1] = ...
        v[2] = ...
        ...
    end
```
and the functions `q_i` providing the solutions must have the interface
```julia
    function q_i(t, q₀, q₁, h, params)
        q₁[1] = q₀[1] + ...
        q₁[2] = q₀[2] + ...
        ...
    end
```
where `t` is the current time, `q₀` is the current solution vector, `q₁` is the
new solution vector which holds the result of computing one substep with the
vector field ``v_i`` on `t` and `q₀`, and `h` is the (sub-)timestep to compute
the update for.

### Constructors

```julia
SODE(v, invariants, parameters, periodicity)
SODE(v; invariants=NullInvariants(), parameters=NullParameters(), periodicity=NullPeriodicity())
SODE(v, q, invariants, parameters, periodicity)
SODE(v, q; invariants=NullInvariants(), parameters=NullParameters(), periodicity=NullPeriodicity())
```

"""
struct SODE{vType <: Union{Tuple,Nothing}, qType <: Union{Tuple,Nothing},
            invType <: OptionalInvariants,
            parType <: OptionalParameters,
            perType <: OptionalPeriodicity} <: AbstractEquationODE{invType,parType,perType}

    v::vType
    q::qType

    invariants::invType
    parameters::parType
    periodicity::perType

    function SODE(v, q, invariants, parameters, periodicity)
        new{typeof(v), typeof(q), typeof(invariants), typeof(parameters), typeof(periodicity)}(
                v, q, invariants, parameters, periodicity)
    end
end

SODE(v; invariants=NullInvariants(), parameters=NullParameters(), periodicity=NullPeriodicity()) = SODE(v, nothing, invariants, parameters, periodicity)
SODE(v, q; invariants=NullInvariants(), parameters=NullParameters(), periodicity=NullPeriodicity()) = SODE(v, q, invariants, parameters, periodicity)

GeometricBase.invariants(equation::SODE) = equation.invariants
GeometricBase.parameters(equation::SODE) = equation.parameters
GeometricBase.periodicity(equation::SODE) = equation.periodicity

const SODEQT{QT,VT,invT,parT,perT} = SODE{VT,QT,invT,parT,perT} # type alias for dispatch on solution type parameter
const SODEVT{VT,QT,invT,parT,perT} = SODE{VT,QT,invT,parT,perT} # type alias for dispatch on vector field type parameter

hassolution(::SODEQT{<:Nothing}) = false
hassolution(::SODEQT{<:Tuple}) = true # && all(typeof(Q) <: Functiong for Q in equ.q)

hassolution(::SODEQT{<:Nothing}, i) = false
hassolution(equ::SODEQT{<:Tuple}, i) = i ≤ length(equ.q)# && typeof(equ.q[i]) <: Function

hasvectorfield(::SODEVT{<:Nothing}) = false
hasvectorfield(::SODEVT{<:Tuple}) = true # && all(typeof(V) <: Function for V in equ.v)

hasvectorfield(::SODEVT{<:Nothing}, i) = false
hasvectorfield(equ::SODEVT{<:Tuple}, i) = i ≤ length(equ.v)# && typeof(equ.v[i]) <: Function

function check_initial_conditions(::SODE, ics::NamedTuple)
    haskey(ics, :q) || return false
    return true
end

function check_methods(equ::SODE, tspan, ics, params)
    if hasvectorfield(equ)
        for v in equ.v
            applicable(v, tspan[begin], ics.q, zero(ics.q), params) || return false
        end
    end
    if hassolution(equ)
        for q in equ.q
            applicable(q, tspan[begin], ics.q, zero(ics.q), params) || return false
        end
    end
    return true
end

function datatype(equ::SODE, ics::NamedTuple)
    @assert check_initial_conditions(equ, ics)
    return eltype(ics.q)
end

function arrtype(equ::SODE, ics::NamedTuple)
    @assert check_initial_conditions(equ, ics)
    typeof(ics.q)
end

_get_v(equ::SODE, params) = ((t,q,v) -> V(t, q, v, params) : V for V in equ.v)
_get_q(equ::SODE, params) = ((t,q̄,q,h) -> Q(t, q̄, q, h, params) : Q for Q in equ.q)

_functions(equ::SODE) = (v = equ.v,)
_solutions(equ::SODE) = (q = equ.q,)
_functions(equ::SODE, params::OptionalParameters) = (v = _get_v(equ, params),)
_solutions(equ::SODE, params::OptionalParameters) = (q = _get_q(equ, params),)


text:
"""
$(TYPEDEF)

A scalar equation for optimization.

# Fields
$(FIELDS)

# Examples

```julia
@variables x y z
@parameters σ ρ β

op = σ*(y-x) + x*(ρ-z)-y + x*y - β*z
@named os = OptimizationSystem(op, [x,y,z],[σ,ρ,β])
```
"""
struct OptimizationSystem <: AbstractTimeIndependentSystem
    """Vector of equations defining the system."""
    op::Any
    """Unknown variables."""
    states::Vector
    """Parameters."""
    ps::Vector
    """Array variables."""
    var_to_name::Any
    observed::Vector{Equation}
    equality_constraints::Vector{Equation}
    inequality_constraints::Vector
    """
    Name: the name of the system.  These are required to have unique names.
    """
    name::Symbol
    """
    systems: The internal systems
    """
    systems::Vector{OptimizationSystem}
    """
    defaults: The default values to use when initial conditions and/or
    parameters are not supplied in `ODEProblem`.
    """
    defaults::Dict
    function OptimizationSystem(op, states, ps, var_to_name, observed, equality_constraints,
                                inequality_constraints, name, systems, defaults;
                                checks::Bool = true)
        if checks
            check_units(op)
            check_units(observed)
            check_units(equality_constraints)
            all_dimensionless([states; ps]) || check_units(inequality_constraints)
        end
        new(op, states, ps, var_to_name, observed, equality_constraints,
            inequality_constraints, name, systems, defaults)
    end
end

function OptimizationSystem(op, states, ps;
                            observed = [],
                            equality_constraints = Equation[],
                            inequality_constraints = [],
                            default_u0 = Dict(),
                            default_p = Dict(),
                            defaults = _merge(Dict(default_u0), Dict(default_p)),
                            name = nothing,
                            systems = OptimizationSystem[],
                            checks = true)
    name === nothing &&
        throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
    if !(isempty(default_u0) && isempty(default_p))
        Base.depwarn("`default_u0` and `default_p` are deprecated. Use `defaults` instead.",
                     :OptimizationSystem, force = true)
    end
    sysnames = nameof.(systems)
    if length(unique(sysnames)) != length(sysnames)
        throw(ArgumentError("System names must be unique."))
    end
    defaults = todict(defaults)
    defaults = Dict(value(k) => value(v) for (k, v) in pairs(defaults))

    states, ps = value.(states), value.(ps)
    var_to_name = Dict()
    process_variables!(var_to_name, defaults, states)
    process_variables!(var_to_name, defaults, ps)
    isempty(observed) || collect_var_to_name!(var_to_name, (eq.lhs for eq in observed))
    OptimizationSystem(value(op), states, ps, var_to_name,
                       observed,
                       equality_constraints, inequality_constraints,
                       name, systems, defaults; checks = checks)
end

function calculate_gradient(sys::OptimizationSystem)
    expand_derivatives.(gradient(equations(sys), states(sys)))
end

function generate_gradient(sys::OptimizationSystem, vs = states(sys), ps = parameters(sys);
                           kwargs...)
    grad = calculate_gradient(sys)
    return build_function(grad, vs, ps;
                          conv = AbstractSysToExpr(sys), kwargs...)
end

function calculate_hessian(sys::OptimizationSystem)
    expand_derivatives.(hessian(equations(sys), states(sys)))
end

function generate_hessian(sys::OptimizationSystem, vs = states(sys), ps = parameters(sys);
                          sparse = false, kwargs...)
    if sparse
        hess = sparsehessian(equations(sys), states(sys))
    else
        hess = calculate_hessian(sys)
    end
    return build_function(hess, vs, ps;
                          conv = AbstractSysToExpr(sys), kwargs...)
end

function generate_function(sys::OptimizationSystem, vs = states(sys), ps = parameters(sys);
                           kwargs...)
    return build_function(equations(sys), vs, ps;
                          conv = AbstractSysToExpr(sys), kwargs...)
end

function equations(sys::OptimizationSystem)
    isempty(get_systems(sys)) ? get_op(sys) :
    get_op(sys) + reduce(+, namespace_expr.(get_systems(sys)))
end
namespace_expr(sys::OptimizationSystem) = namespace_expr(get_op(sys), sys)

hessian_sparsity(sys::OptimizationSystem) = hessian_sparsity(get_op(sys), states(sys))

function DiffEqBase.OptimizationProblem(sys::OptimizationSystem, args...; kwargs...)
    DiffEqBase.OptimizationProblem{true}(sys::OptimizationSystem, args...; kwargs...)
end

"""
```julia
function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem,u0map,
                                          parammap=DiffEqBase.NullParameters();
                                          lb=nothing, ub=nothing,
                                          grad = false,
                                          hess = false, sparse = false,
                                          checkbounds = false,
                                          linenumbers = true, parallel=SerialForm(),
                                          kwargs...) where iip
```

Generates an OptimizationProblem from an OptimizationSystem and allows for automatically
symbolically calculating numerical enhancements.
"""
function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
                                             parammap = DiffEqBase.NullParameters();
                                             lb = nothing, ub = nothing,
                                             grad = false,
                                             hess = false, sparse = false,
                                             checkbounds = false,
                                             linenumbers = true, parallel = SerialForm(),
                                             use_union = false,
                                             kwargs...) where {iip}
    dvs = states(sys)
    ps = parameters(sys)

    f = generate_function(sys, checkbounds = checkbounds, linenumbers = linenumbers,
                          expression = Val{false})

    if grad
        grad_oop, grad_iip = generate_gradient(sys, checkbounds = checkbounds,
                                               linenumbers = linenumbers,
                                               parallel = parallel, expression = Val{false})
        _grad(u, p) = grad_oop(u, p)
        _grad(J, u, p) = (grad_iip(J, u, p); J)
    else
        _grad = nothing
    end

    if hess
        hess_oop, hess_iip = generate_hessian(sys, checkbounds = checkbounds,
                                              linenumbers = linenumbers,
                                              sparse = sparse, parallel = parallel,
                                              expression = Val{false})
        _hess(u, p) = hess_oop(u, p)
        _hess(J, u, p) = (hess_iip(J, u, p); J)
    else
        _hess = nothing
    end

    _f = DiffEqBase.OptimizationFunction{iip}(f,
                                              SciMLBase.NoAD();
                                              grad = _grad,
                                              hess = _hess)

    defs = defaults(sys)
    defs = mergedefaults(defs, parammap, ps)
    defs = mergedefaults(defs, u0map, dvs)

    u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = false)
    p = varmap_to_vars(parammap, ps; defaults = defs, tofloat = false, use_union)
    lb = varmap_to_vars(lb, dvs; check = false, tofloat = false, use_union)
    ub = varmap_to_vars(ub, dvs; check = false, tofloat = false, use_union)
    OptimizationProblem{iip}(_f, u0, p; lb = lb, ub = ub, kwargs...)
end

"""
```julia
function DiffEqBase.OptimizationProblemExpr{iip}(sys::OptimizationSystem,
                                          parammap=DiffEqBase.NullParameters();
                                          u0=nothing, lb=nothing, ub=nothing,
                                          grad = false,
                                          hes = false, sparse = false,
                                          checkbounds = false,
                                          linenumbers = true, parallel=SerialForm(),
                                          kwargs...) where iip
```

Generates a Julia expression for an OptimizationProblem from an
OptimizationSystem and allows for automatically symbolically
calculating numerical enhancements.
"""
struct OptimizationProblemExpr{iip} end

function OptimizationProblemExpr(sys::OptimizationSystem, args...; kwargs...)
    OptimizationProblemExpr{true}(sys::OptimizationSystem, args...; kwargs...)
end

function OptimizationProblemExpr{iip}(sys::OptimizationSystem, u0,
                                      parammap = DiffEqBase.NullParameters();
                                      lb = nothing, ub = nothing,
                                      grad = false,
                                      hess = false, sparse = false,
                                      checkbounds = false,
                                      linenumbers = false, parallel = SerialForm(),
                                      use_union = false,
                                      kwargs...) where {iip}
    dvs = states(sys)
    ps = parameters(sys)
    idx = iip ? 2 : 1
    f = generate_function(sys, checkbounds = checkbounds, linenumbers = linenumbers,
                          expression = Val{true})
    if grad
        _grad = generate_gradient(sys, checkbounds = checkbounds, linenumbers = linenumbers,
                                  parallel = parallel, expression = Val{false})[idx]
    else
        _grad = :nothing
    end

    if hess
        _hess = generate_hessian(sys, checkbounds = checkbounds, linenumbers = linenumbers,
                                 sparse = sparse, parallel = parallel,
                                 expression = Val{false})[idx]
    else
        _hess = :nothing
    end

    defs = defaults(sys)
    defs = mergedefaults(defs, parammap, ps)
    defs = mergedefaults(defs, u0map, dvs)

    u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = false)
    p = varmap_to_vars(parammap, ps; defaults = defs, tofloat = false, use_union)
    lb = varmap_to_vars(lb, dvs; check = false, tofloat = false, use_union)
    ub = varmap_to_vars(ub, dvs; check = false, tofloat = false, use_union)
    quote
        f = $f
        p = $p
        u0 = $u0
        grad = $_grad
        hess = $_hess
        lb = $lb
        ub = $ub
        _f = OptimizationFunction{iip}(f, SciMLBase.NoAD();
                                       grad = grad,
                                       hess = hess)
        OptimizationProblem{$iip}(_f, u0, p; lb = lb, ub = ub, kwargs...)
    end
end


text:
module MultiFloats

export MultiFloat, renormalize,
       Float16x, Float32x, Float64x,
       Float64x1, Float64x2, Float64x3, Float64x4,
       Float64x5, Float64x6, Float64x7, Float64x8,
       use_clean_multifloat_arithmetic,
       use_standard_multifloat_arithmetic,
       use_sloppy_multifloat_arithmetic

include("./MultiFloatsCodeGen.jl")
using .MultiFloatsCodeGen

####################################################### DEFINITION OF MULTIFLOAT

struct MultiFloat{T,N} <: AbstractFloat
    _limbs::NTuple{N,T}
end

# Short alias for brevity of type declarations
const MF = MultiFloat

const Float16x{N} = MultiFloat{Float16,N}
const Float32x{N} = MultiFloat{Float32,N}
const Float64x{N} = MultiFloat{Float64,N}

const Float64x1 = Float64x{1}
const Float64x2 = Float64x{2}
const Float64x3 = Float64x{3}
const Float64x4 = Float64x{4}
const Float64x5 = Float64x{5}
const Float64x6 = Float64x{6}
const Float64x7 = Float64x{7}
const Float64x8 = Float64x{8}

############################################## CONSTRUCTION FROM PRIMITIVE TYPES

@inline MultiFloat{T,N}(x::MultiFloat{T,N}) where {T,N} = x

@inline MultiFloat{T,N}(x::T) where {T,N} =
    MultiFloat{T,N}((x, ntuple(_ -> zero(T), N - 1)...))

@inline MultiFloat{T,N}(x::MultiFloat{T,M}) where {T,M,N} =
    MultiFloat{T,N}((
        ntuple(i -> x._limbs[i], min(M, N))...,
        ntuple(_ -> zero(T), max(N - M, 0))...))

# Values of the types Bool, Int8, UInt8, Int16, UInt16, Float16, Int32, UInt32,
# and Float32 can be converted losslessly to a single Float64, which has 53
# bits of integer precision.

@inline Float64x{N}(x::Bool   ) where {N} = Float64x{N}(Float64(x))
@inline Float64x{N}(x::Int8   ) where {N} = Float64x{N}(Float64(x))
@inline Float64x{N}(x::UInt8  ) where {N} = Float64x{N}(Float64(x))
@inline Float64x{N}(x::Int16  ) where {N} = Float64x{N}(Float64(x))
@inline Float64x{N}(x::UInt16 ) where {N} = Float64x{N}(Float64(x))
@inline Float64x{N}(x::Float16) where {N} = Float64x{N}(Float64(x))
@inline Float64x{N}(x::Int32  ) where {N} = Float64x{N}(Float64(x))
@inline Float64x{N}(x::UInt32 ) where {N} = Float64x{N}(Float64(x))
@inline Float64x{N}(x::Float32) where {N} = Float64x{N}(Float64(x))

# Values of the types Int64, UInt64, Int128, and UInt128 cannot be converted
# losslessly to a single Float64 and must be split into multiple components.

@inline Float64x1(x::Int64  ) = Float64x1(Float64(x))
@inline Float64x1(x::UInt64 ) = Float64x1(Float64(x))
@inline Float64x1(x::Int128 ) = Float64x1(Float64(x))
@inline Float64x1(x::UInt128) = Float64x1(Float64(x))

@inline function Float64x2(x::Int128)
    x0 = Float64(x)
    x1 = Float64(x - Int128(x0))
    Float64x2((x0, x1))
end

@inline function Float64x2(x::UInt128)
    x0 = Float64(x)
    x1 = Float64(reinterpret(Int128, x - UInt128(x0)))
    Float64x2((x0, x1))
end

@inline function Float64x{N}(x::Int64) where {N}
    x0 = Float64(x)
    x1 = Float64(x - Int64(x0))
    Float64x{N}((x0, x1, ntuple(_ -> 0.0, N - 2)...))
end

@inline function Float64x{N}(x::UInt64) where {N}
    x0 = Float64(x)
    x1 = Float64(reinterpret(Int64, x - UInt64(x0)))
    Float64x{N}((x0, x1, ntuple(_ -> 0.0, N - 2)...))
end

@inline function Float64x{N}(x::Int128) where {N}
    x0 = Float64(x)
    r1 = x - Int128(x0)
    x1 = Float64(r1)
    x2 = Float64(r1 - Int128(x1))
    Float64x{N}((x0, x1, x2, ntuple(_ -> 0.0, N - 3)...))
end

@inline function Float64x{N}(x::UInt128) where {N}
    x0 = Float64(x)
    r1 = reinterpret(Int128, x - UInt128(x0))
    x1 = Float64(r1)
    x2 = Float64(r1 - Int128(x1))
    Float64x{N}((x0, x1, x2, ntuple(_ -> 0.0, N - 3)...))
end

################################################## CONVERSION TO PRIMITIVE TYPES

@inline Base.Float16(x::Float64x{N}) where {N} = Float16(x._limbs[1])
@inline Base.Float32(x::Float64x{N}) where {N} = Float32(x._limbs[1])

@inline Base.Float16(x::Float16x{N}) where {N} = x._limbs[1]
@inline Base.Float32(x::Float32x{N}) where {N} = x._limbs[1]
@inline Base.Float64(x::Float64x{N}) where {N} = x._limbs[1]

####################################################### FLOATING-POINT CONSTANTS

@inline Base.precision(::Type{MF{T,N}}) where {T,N} =
    N * precision(T) + (N - 1) # implicit bits of precision between limbs

@inline Base.zero(::Type{MF{T,N}}) where {T,N} = MF{T,N}(zero(T)  )
@inline Base.one( ::Type{MF{T,N}}) where {T,N} = MF{T,N}(one( T)  )
@inline Base.eps( ::Type{MF{T,N}}) where {T,N} = MF{T,N}(eps( T)^N)

# TODO: This is technically not the maximum/minimum representable MultiFloat.
@inline Base.floatmin(::Type{MF{T,N}}) where {T,N} = MF{T,N}(floatmin(T))
@inline Base.floatmax(::Type{MF{T,N}}) where {T,N} = MF{T,N}(floatmax(T))

@inline Base.typemin(::Type{MF{T,N}}) where {T,N} =
    MF{T,N}(ntuple(_ -> typemin(T), N))
@inline Base.typemax(::Type{MF{T,N}}) where {T,N} =
    MF{T,N}(ntuple(_ -> typemax(T), N))

#################################################### CONSTRUCTION FROM BIG TYPES

function MultiFloat{T,N}(x::BigFloat) where {T,N}
    if x > floatmax(T)
        return typemax(MultiFloat{T,N})
    elseif x < -floatmax(T)
        return typemin(MultiFloat{T,N})
    elseif -floatmin(T) < x < floatmin(T)
        return zero(MultiFloat{T,N})
    elseif isnan(x)
        return MF{T,N}(ntuple(_ -> T(NaN), N))
    end
    setrounding(BigFloat, RoundNearest) do
        setprecision(BigFloat, Int(precision(x))) do
            y = Vector{T}(undef, N)
            y[1] = T(x)
            for i = 2 : N
                x -= y[i-1]
                y[i] = T(x)
            end
            return MultiFloat{T,N}((y...,))
        end
    end
end

function MultiFloat{T,N}(x::BigInt) where {T,N}
    y = Vector{T}(undef, N)
    for i = 1 : N
        y[i] = T(x)
        x -= BigInt(y[i])
    end
    return MultiFloat{T,N}((y...,))
end

MultiFloat{T,N}(x::Rational{U}) where {T,N,U} =
    MultiFloat{T,N}(numerator(x)) / MultiFloat{T,N}(denominator(x))

MultiFloat{T,N}(x::AbstractString) where {T,N} =
    MultiFloat{T,N}(BigFloat(x, precision = precision(T)
        + exponent(floatmax(T)) - exponent(floatmin(T))))

######################################################## CONVERSION TO BIG TYPES

Base.BigFloat(x::MultiFloat{T,N}) where {T,N} =
    +(ntuple(i -> BigFloat(x._limbs[N-i+1]), N)...)

Base.Rational{BigInt}(x::MultiFloat{T,N}) where {T,N} =
    sum(Rational{BigInt}.(x._limbs))

################################################################ PROMOTION RULES

Base.promote_rule(::Type{MF{T,N}}, ::Type{T      }) where {T,N} = MF{T,N}
Base.promote_rule(::Type{MF{T,N}}, ::Type{Int8   }) where {T,N} = MF{T,N}
Base.promote_rule(::Type{MF{T,N}}, ::Type{Int16  }) where {T,N} = MF{T,N}
Base.promote_rule(::Type{MF{T,N}}, ::Type{Int32  }) where {T,N} = MF{T,N}
Base.promote_rule(::Type{MF{T,N}}, ::Type{Int64  }) where {T,N} = MF{T,N}
Base.promote_rule(::Type{MF{T,N}}, ::Type{Int128 }) where {T,N} = MF{T,N}
Base.promote_rule(::Type{MF{T,N}}, ::Type{Bool   }) where {T,N} = MF{T,N}
Base.promote_rule(::Type{MF{T,N}}, ::Type{UInt8  }) where {T,N} = MF{T,N}
Base.promote_rule(::Type{MF{T,N}}, ::Type{UInt16 }) where {T,N} = MF{T,N}
Base.promote_rule(::Type{MF{T,N}}, ::Type{UInt32 }) where {T,N} = MF{T,N}
Base.promote_rule(::Type{MF{T,N}}, ::Type{UInt64 }) where {T,N} = MF{T,N}
Base.promote_rule(::Type{MF{T,N}}, ::Type{UInt128}) where {T,N} = MF{T,N}

Base.promote_rule(::Type{MF{T,N}}, ::Type{BigFloat}) where {T,N} = BigFloat

Base.promote_rule(::Type{Float32x{N}}, ::Type{Float16}) where {N} = Float32x{N}
Base.promote_rule(::Type{Float64x{N}}, ::Type{Float16}) where {N} = Float64x{N}
Base.promote_rule(::Type{Float64x{N}}, ::Type{Float32}) where {N} = Float64x{N}

################################################################ RENORMALIZATION

@inline function renormalize(x::MF{T,N}) where {T,N}
    total = +(x._limbs...)
    if isfinite(total)
        while true
            x0::MF{T,N} = x + zero(T)
            if !(x0._limbs != x._limbs); break; end
            x = x0
        end
        return x
    else
        return MultiFloat{T,N}(ntuple(_ -> total, N))
    end
end

@inline renormalize(x::T) where {T<:Number} = x

function call_normalized(callback, x::MultiFloat{T,N}) where {T,N}
    x = renormalize(x)
    if !isfinite(x._limbs[1])
        return callback(x._limbs[1])
    else
        i = N
        while (i > 0) && iszero(x._limbs[i])
            i -= 1
        end
        if iszero(i)
            return callback(zero(T))
        else
            return setprecision(() -> callback(BigFloat(x)),
                precision(T) + exponent(x._limbs[1]) - exponent(x._limbs[i]))
        end
    end
end

####################################################################### PRINTING

function Base.show(io::IO, x::MultiFloat{T,N}) where {T,N}
    return call_normalized(y -> show(io, y), x)
end

# Thanks to Greg Plowman (https://github.com/GregPlowman) for suggesting
# implementations of Printf.fix_dec and Printf.ini_dec for @printf support.

import Printf: fix_dec, ini_dec

if VERSION < v"1.1"

    fix_dec(out, x::MultiFloat{T,N}, flags::String, width::Int,
            precision::Int, c::Char) where {T,N} =
        call_normalized(d -> fix_dec(out, BigFloat(d), flags,
                                     width, precision, c), x)

    ini_dec(out, x::MultiFloat{T,N}, ndigits::Int, flags::String,
            width::Int, precision::Int, c::Char) where {T,N} =
        call_normalized(d -> ini_dec(out, BigFloat(d), ndigits,
                                     flags, width, precision, c), x)

else

    fix_dec(out, x::MultiFloat{T,N}, flags::String, width::Int,
            precision::Int, c::Char, digits) where {T,N} =
        call_normalized(d -> fix_dec(out, BigFloat(d), flags,
                                     width, precision, c, digits), x)

    ini_dec(out, x::MultiFloat{T,N}, ndigits::Int, flags::String,
            width::Int, precision::Int, c::Char, digits) where {T,N} =
        call_normalized(d -> ini_dec(out, BigFloat(d), ndigits,
                                     flags, width, precision, c, digits), x)

end

################################################### FLOATING-POINT INTROSPECTION

@inline _iszero(x::MF{T,N}) where {T,N} =
    (&)(ntuple(i -> iszero(x._limbs[i]), N)...)
@inline _isone( x::MF{T,N}) where {T,N} =
    isone(x._limbs[1]) & (&)(ntuple(i -> iszero(x._limbs[i + 1]), N - 1)...)

@inline Base.iszero(x::MF{T,1}) where {T  } =  iszero(x._limbs[1])
@inline Base.isone( x::MF{T,1}) where {T  } =  isone( x._limbs[1])
@inline Base.iszero(x::MF{T,N}) where {T,N} = _iszero(renormalize(x))
@inline Base.isone( x::MF{T,N}) where {T,N} = _isone( renormalize(x))

@inline _head(x::MF{T,N}) where {T,N} = renormalize(x)._limbs[1]
@inline Base.exponent(   x::MF{T,N}) where {T,N} = exponent(   _head(x))
@inline Base.signbit(    x::MF{T,N}) where {T,N} = signbit(    _head(x))
@inline Base.issubnormal(x::MF{T,N}) where {T,N} = issubnormal(_head(x))
@inline Base.isfinite(   x::MF{T,N}) where {T,N} = isfinite(   _head(x))
@inline Base.isinf(      x::MF{T,N}) where {T,N} = isinf(      _head(x))
@inline Base.isnan(      x::MF{T,N}) where {T,N} = isnan(      _head(x))
@inline Base.isinteger(  x::MF{T,N}) where {T,N} =
    all(isinteger.(renormalize(x)._limbs))

@inline function Base.nextfloat(x::MF{T,N}) where {T,N}
    y = renormalize(x)
    return renormalize(MF{T,N}((
        ntuple(i -> y._limbs[i], N - 1)...,
        nextfloat(y._limbs[N]))))
end

@inline function Base.prevfloat(x::MF{T,N}) where {T,N}
    y = renormalize(x)
    return renormalize(MF{T,N}((
        ntuple(i -> y._limbs[i], N - 1)...,
        prevfloat(y._limbs[N]))))
end

import LinearAlgebra: floatmin2
@inline floatmin2(::Type{MF{T,N}}) where {T,N} = MF{T,N}(ldexp(one(T),
    div(exponent(floatmin(T)) - N * exponent(eps(T)), 2)))

##################################################### ERROR-FREE TRANSFORMATIONS

@inline function two_sum(a::T, b::T) where {T}
    s = a + b
    v = s - a
    s, (a - (s - v)) + (b - v)
end

@inline function quick_two_sum(a::T, b::T) where {T}
    s = a + b
    s, b - (s - a)
end

@inline function two_prod(a::T, b::T) where {T}
    p = a * b
    p, fma(a, b, -p)
end

##################################################################### COMPARISON

@inline Base.:(==)(x::MF{T,N}, y::MF{T,N}) where {T,N} =
    multifloat_eq(renormalize(x), renormalize(y))
@inline Base.:(!=)(x::MF{T,N}, y::MF{T,N}) where {T,N} =
    multifloat_ne(renormalize(x), renormalize(y))
@inline Base.:(< )(x::MF{T,N}, y::MF{T,N}) where {T,N} =
    multifloat_lt(renormalize(x), renormalize(y))
@inline Base.:(> )(x::MF{T,N}, y::MF{T,N}) where {T,N} =
    multifloat_gt(renormalize(x), renormalize(y))
@inline Base.:(<=)(x::MF{T,N}, y::MF{T,N}) where {T,N} =
    multifloat_le(renormalize(x), renormalize(y))
@inline Base.:(>=)(x::MF{T,N}, y::MF{T,N}) where {T,N} =
    multifloat_ge(renormalize(x), renormalize(y))

##################################################################### ARITHMETIC

@inline multifloat_add(a::MF{T,1}, b::MF{T,1}) where {T} =
    MF{T,1}(a._limbs[1] + b._limbs[1])
@inline multifloat_mul(a::MF{T,1}, b::MF{T,1}) where {T} =
    MF{T,1}(a._limbs[1] * b._limbs[1])
@inline multifloat_div(a::MF{T,1}, b::MF{T,1}) where {T} =
    MF{T,1}(a._limbs[1] / b._limbs[1])
@inline multifloat_float_add(a::MF{T,1}, b::T) where {T} =
    MF{T,1}(a._limbs[1] + b)
@inline multifloat_float_mul(a::MF{T,1}, b::T) where {T} =
    MF{T,1}(a._limbs[1] * b)
@inline multifloat_sqrt(x::MF{T,1}) where {T} =
    MF{T,1}(unsafe_sqrt(x._limbs[1]))

@inline Base.:+(x::MF{T,N}, y::MF{T,N}) where {T,N} = multifloat_add(x, y)
@inline Base.:*(x::MF{T,N}, y::MF{T,N}) where {T,N} = multifloat_mul(x, y)
@inline Base.:/(x::MF{T,N}, y::MF{T,N}) where {T,N} = multifloat_div(x, y)
@inline Base.:+(x::MF{T,N}, y::T) where {T        ,N} =
    multifloat_float_add(x, y)
@inline Base.:+(x::MF{T,N}, y::T) where {T<:Number,N} =
    multifloat_float_add(x, y)
@inline Base.:*(x::MF{T,N}, y::T) where {T        ,N} =
    multifloat_float_mul(x, y)
@inline Base.:*(x::MF{T,N}, y::T) where {T<:Number,N} =
    multifloat_float_mul(x, y)

@inline Base.:-(x::MF{T,N}) where {T,N} = MF{T,N}(ntuple(i -> -x._limbs[i], N))
@inline Base.:-(x::MF{T,N}, y::MF{T,N}) where {T        ,N} = x + (-y)
@inline Base.:-(x::MF{T,N}, y::T      ) where {T        ,N} = x + (-y)
@inline Base.:-(x::MF{T,N}, y::T      ) where {T<:Number,N} = x + (-y)

@inline Base.:+(x::T, y::MF{T,N}) where {T        ,N} = y + x
@inline Base.:+(x::T, y::MF{T,N}) where {T<:Number,N} = y + x
@inline Base.:-(x::T, y::MF{T,N}) where {T        ,N} = -(y + (-x))
@inline Base.:-(x::T, y::MF{T,N}) where {T<:Number,N} = -(y + (-x))
@inline Base.:*(x::T, y::MF{T,N}) where {T        ,N} = y * x
@inline Base.:*(x::T, y::MF{T,N}) where {T<:Number,N} = y * x

@inline Base.inv(x::MF{T,N}) where {T,N} = one(MF{T,N}) / x

@inline function Base.abs(x::MF{T,N}) where {T,N}
    x = renormalize(x)
    ifelse(signbit(x._limbs[1]), -x, x)
end

@inline function Base.abs2(x::MF{T,N}) where {T,N}
    x = renormalize(x)
    renormalize(x * x)
end

@inline unsafe_sqrt(x::Float32) = Base.sqrt_llvm(x)
@inline unsafe_sqrt(x::Float64) = Base.sqrt_llvm(x)
@inline unsafe_sqrt(x::T) where {T <: Real} = sqrt(x)

@inline function Base.sqrt(x::MF{T,N}) where {T,N}
    x = renormalize(x)
    if iszero(x)
        return x
    else
        return multifloat_sqrt(x)
    end
end

@inline Base.hypot(x::MF{T,N}, y::MF{T,N}) where {T,N} = sqrt(x*x + y*y)

######################################################## EXPONENTIATION (BASE 2)

@inline scale(a::T, x::MultiFloat{T,N}) where {T,N} =
    MultiFloat{T,N}(ntuple(i -> a * x._limbs[i], N))

@inline function Base.ldexp(x::MF{T,N}, n::U) where {T,N,U<:Integer}
    x = renormalize(x)
    return MultiFloat{T,N}(ntuple(i -> ldexp(x._limbs[i], n), N))
end

######################################################## EXPONENTIATION (BASE E)

const INVERSE_FACTORIALS_F64 = setprecision(() ->
    [MultiFloat{Float64,20}(inv(BigFloat(factorial(BigInt(i)))))
        for i = 1 : 170],
    BigFloat, 1200)

const LOG2_F64 = setprecision(() ->
    MultiFloat{Float64,20}(log(BigFloat(2))),
    BigFloat, 1200)

log2_f64_literal(n::Int) = :(
    MultiFloat{Float64,$n}($(LOG2_F64._limbs[1:n])))

inverse_factorial_f64_literal(n::Int, i::Int) = :(
    MultiFloat{Float64,$n}($(INVERSE_FACTORIALS_F64[i]._limbs[1:n])))

meta_y(n::Int) = Symbol("y_", n)

meta_y_definition(n::Int) = Expr(:(=), meta_y(n),
    Expr(:call, :*, meta_y(div(n, 2)), meta_y(div(n+1, 2))))

meta_exp_term(n::Int, i::Int) = :(
    $(Symbol("y_", i)) * $(inverse_factorial_f64_literal(n, i)))

function multifloat_exp_func(n::Int, num_terms::Int,
                             reduction_power::Int; sloppy::Bool=false)
    return Expr(:function,
        :(multifloat_exp(x::MultiFloat{Float64,$n})),
        Expr(:block,
            :(exponent_f = Base.rint_llvm(
                x._limbs[1] / $(LOG2_F64._limbs[1]))),
            :(exponent_i = Base.fptosi(Int, exponent_f)),
            :(y_1 = scale($(ldexp(1.0, -reduction_power)),
                MultiFloat{Float64,$n}(MultiFloat{Float64,$(n + !sloppy)}(x) -
                    exponent_f * $(log2_f64_literal(n + !sloppy))))),
            [meta_y_definition(i) for i = 2 : num_terms]...,
            :(exp_y = $(meta_exp_term(n, num_terms))),
            [:(exp_y += $(meta_exp_term(n, i)))
                for i = num_terms-1 : -1 : 3]...,
            :(exp_y += scale(0.5, y_2)),
            :(exp_y += y_1),
            :(exp_y += 1.0),
            [:(exp_y *= exp_y) for _ = 1 : reduction_power]...,
            :(return scale(reinterpret(Float64,
                UInt64(1023 + exponent_i) << 52), exp_y))))
end

@inline multifloat_exp(x::MF{T,1}) where {T} = MF{T,1}(exp(x._limbs[1]))

function Base.exp(x::MultiFloat{T,N}) where {T,N}
    x = renormalize(x)
    if x._limbs[1] >= log(floatmax(Float64))
        return typemax(MultiFloat{T,N})
    elseif x._limbs[1] <= log(floatmin(Float64))
        return zero(MultiFloat{T,N})
    else
        return multifloat_exp(x)
    end
end

function Base.log(x::MultiFloat{T,N}) where {T,N}
    y = MultiFloat{T,N}(log(x._limbs[1]))
    for _ = 1 : ceil(Int, log2(N)) + 1
        y += x * exp(-y) - one(T)
    end
    return y
end

############################################## BIGFLOAT TRANSCENDENTAL STOP-GAPS

BASE_TRANSCENDENTAL_FUNCTIONS = [
    :exp2, :exp10, :expm1, :log2, :log10, :log1p,
    :sin, :cos, :tan, :sec, :csc, :cot,
    :sinh, :cosh, :tanh, :sech, :csch, :coth,
    :sind, :cosd, :tand, :secd, :cscd, :cotd,
    :asin, :acos, :atan, :asec, :acsc, :acot,
    :asinh, :acosh, :atanh, :asech, :acsch, :acoth,
    :asind, :acosd, :atand, :asecd, :acscd, :acotd
]

for name in BASE_TRANSCENDENTAL_FUNCTIONS
    eval(:(Base.$name(x::MultiFloat{T,N}) where {T,N} = error($(
        "$name(MultiFloat) is not yet implemented. For a temporary workaround,\n" *
        "call MultiFloats.use_bigfloat_transcendentals() immediately after\n" *
        "importing MultiFloats. This will use the BigFloat implementation of\n" *
        "$name, which will not be as fast as a pure-MultiFloat implementation.\n"
    ))))
end

eval_bigfloat(f::Function, x::MultiFloat{T,N}, k::Int) where {T,N} =
    setprecision(() -> MultiFloat{T,N}(f(BigFloat(x))),
                 BigFloat, precision(MultiFloat{T,N}) + k)

function use_bigfloat_transcendentals(k::Int=20)
    for name in BASE_TRANSCENDENTAL_FUNCTIONS
        eval(:(Base.$name(x::MultiFloat{T,N}) where {T,N} =
                eval_bigfloat($name, x, $k)))
    end
end

####################################################### RANDOM NUMBER GENERATION

import Random
using Random: AbstractRNG, SamplerTrivial, CloseOpen01

@inline function _rand_f64(rng::AbstractRNG, k::Int)
    expnt = reinterpret(UInt64,
        exponent(floatmax(Float64)) + k) << (precision(Float64) - 1)
    mntsa = rand(rng, Random.UInt52())
    return reinterpret(Float64, expnt | mntsa)
end

@inline function _rand_sf64(rng::AbstractRNG, k::Int)
    expnt = reinterpret(UInt64,
        exponent(floatmax(Float64)) + k) << (precision(Float64) - 1)
    mntsa = rand(rng, UInt64) & 0x800FFFFFFFFFFFFF
    return reinterpret(Float64, expnt | mntsa)
end

@inline function _rand_mf64(rng::AbstractRNG, offset::Int,
                            padding::NTuple{N,Int}) where {N}
    exponents = (cumsum(padding) .+
        (precision(Float64) + 1) .* ntuple(identity, N))
    return Float64x{N+1}((_rand_f64(rng, offset),
                          _rand_sf64.(rng, offset .- exponents)...))
end

function multifloat_rand_func(n::Int)
    :(
        Random.rand(rng::AbstractRNG,
                    ::SamplerTrivial{CloseOpen01{Float64x{$n}}}) =
        _rand_mf64(
            rng,
            -leading_zeros(rand(rng, UInt64)) - 1,
            $(Expr(:tuple,
                   ntuple(_ -> :(leading_zeros(rand(rng, UInt64))), n - 1)...
            ))
        )
    )
end

################################################################ PRECISION MODES

function use_clean_multifloat_arithmetic(n::Integer=8)
    for i = 1 : n
        eval(multifloat_eq_func(i))
        eval(multifloat_ne_func(i))
        eval(multifloat_lt_func(i))
        eval(multifloat_gt_func(i))
        eval(multifloat_le_func(i))
        eval(multifloat_ge_func(i))
        eval(multifloat_rand_func(i))
    end
    for i = 2 : n+1
        eval(two_pass_renorm_func(     i, sloppy=false))
        eval(multifloat_add_func(      i, sloppy=false))
        eval(multifloat_mul_func(      i, sloppy=false))
        eval(multifloat_div_func(      i, sloppy=false))
        eval(multifloat_float_add_func(i, sloppy=false))
        eval(multifloat_float_mul_func(i, sloppy=false))
        eval(multifloat_sqrt_func(     i, sloppy=false))
    end
    eval(MultiFloats.multifloat_exp_func(2, 20, 1, sloppy=false))
    eval(MultiFloats.multifloat_exp_func(3, 28, 1, sloppy=false))
    eval(MultiFloats.multifloat_exp_func(4, 35, 1, sloppy=false))
    eval(MultiFloats.multifloat_exp_func(5, 42, 1, sloppy=false))
    eval(MultiFloats.multifloat_exp_func(6, 49, 1, sloppy=false))
    eval(MultiFloats.multifloat_exp_func(7, 56, 1, sloppy=false))
    eval(MultiFloats.multifloat_exp_func(8, 63, 1, sloppy=false))
    for (_, v) in MultiFloatsCodeGen.MPADD_CACHE
        eval(v)
    end
end

function use_standard_multifloat_arithmetic(n::Integer=8)
    for i = 1 : n
        eval(multifloat_eq_func(i))
        eval(multifloat_ne_func(i))
        eval(multifloat_lt_func(i))
        eval(multifloat_gt_func(i))
        eval(multifloat_le_func(i))
        eval(multifloat_ge_func(i))
        eval(multifloat_rand_func(i))
    end
    for i = 2 : n
        eval(two_pass_renorm_func(     i, sloppy=true ))
        eval(two_pass_renorm_func(     i, sloppy=false))
        eval(multifloat_add_func(      i, sloppy=false))
        eval(multifloat_mul_func(      i, sloppy=true ))
        eval(multifloat_div_func(      i, sloppy=true ))
        eval(multifloat_float_add_func(i, sloppy=false))
        eval(multifloat_float_mul_func(i, sloppy=true ))
        eval(multifloat_sqrt_func(     i, sloppy=true ))
    end
    eval(MultiFloats.multifloat_exp_func(2, 17, 2, sloppy=false))
    eval(MultiFloats.multifloat_exp_func(3, 19, 4, sloppy=true))
    eval(MultiFloats.multifloat_exp_func(4, 20, 6, sloppy=true))
    eval(MultiFloats.multifloat_exp_func(5, 23, 7, sloppy=true))
    eval(MultiFloats.multifloat_exp_func(6, 23, 9, sloppy=true))
    eval(MultiFloats.multifloat_exp_func(7, 22, 12, sloppy=true))
    eval(MultiFloats.multifloat_exp_func(8, 24, 13, sloppy=true))
    for (_, v) in MultiFloatsCodeGen.MPADD_CACHE
        eval(v)
    end
end

function use_sloppy_multifloat_arithmetic(n::Integer=8)
    for i = 1 : n
        eval(multifloat_eq_func(i))
        eval(multifloat_ne_func(i))
        eval(multifloat_lt_func(i))
        eval(multifloat_gt_func(i))
        eval(multifloat_le_func(i))
        eval(multifloat_ge_func(i))
        eval(multifloat_rand_func(i))
    end
    for i = 2 : n
        eval(one_pass_renorm_func(     i, sloppy=true))
        eval(multifloat_add_func(      i, sloppy=true))
        eval(multifloat_mul_func(      i, sloppy=true))
        eval(multifloat_div_func(      i, sloppy=true))
        eval(multifloat_float_add_func(i, sloppy=true))
        eval(multifloat_float_mul_func(i, sloppy=true))
        eval(multifloat_sqrt_func(     i, sloppy=true))
    end
    eval(MultiFloats.multifloat_exp_func(2, 17, 2, sloppy=false))
    eval(MultiFloats.multifloat_exp_func(3, 19, 4, sloppy=true))
    eval(MultiFloats.multifloat_exp_func(4, 20, 6, sloppy=true))
    eval(MultiFloats.multifloat_exp_func(5, 23, 7, sloppy=true))
    eval(MultiFloats.multifloat_exp_func(6, 23, 9, sloppy=true))
    eval(MultiFloats.multifloat_exp_func(7, 22, 12, sloppy=true))
    eval(MultiFloats.multifloat_exp_func(8, 24, 13, sloppy=true))
    for (_, v) in MultiFloatsCodeGen.MPADD_CACHE
        eval(v)
    end
end

use_standard_multifloat_arithmetic()

end # module MultiFloats


text:
# ---
# title: 389. Find the Difference
# id: problem389
# author: Indigo
# date: 2020-12-18
# difficulty: Easy
# categories: Hash Table, Bit Manipulation
# link: <https://leetcode.com/problems/find-the-difference/description/>
# hidden: true
# ---
#
# You are given two strings `s` and `t`.
#
# String `t` is generated by random shuffling string `s` and then add one more
# letter at a random position.
#
# Return the letter that was added to `t`.
#
#
#
# **Example 1:**
#
#
#
#     Input: s = "abcd", t = "abcde"
#     Output: "e"
#     Explanation: 'e' is the letter that was added.
#
#
# **Example 2:**
#
#
#
#     Input: s = "", t = "y"
#     Output: "y"
#
#
# **Example 3:**
#
#
#
#     Input: s = "a", t = "aa"
#     Output: "a"
#
#
# **Example 4:**
#
#
#
#     Input: s = "ae", t = "aea"
#     Output: "a"
#
#
#
#
# **Constraints:**
#
#   * `0 <= s.length <= 1000`
#   * `t.length == s.length + 1`
#   * `s` and `t` consist of lower-case English letters.
#
#
## @lc code=start
using LeetCode

function find_the_difference(s::String, t::String)::Char
    cnt = fill(0, 26)
    for ch in codeunits(s)
        cnt[ch - 0x60] += 1
    end
    for ch in codeunits(t)
        if((cnt[ch - 0x60] -= 1) < 0)
            return Char(ch)
        end
    end
    return t[end]
end
## @lc code=end


text:
using CUDA, Plots, Printf
pyplot()
viz = true

function compute_flux!(qx, qy, 🔥, λ, ∂x, ∂y, nx, ny)
	ix = (blockIdx().x-1) * blockDim().x + threadIdx().x # thread ID, dimension x
	iy = (blockIdx().y-1) * blockDim().y + threadIdx().y # thread ID, dimension y

	# TODO add qx and qy computation in a vectorised fashion here based on index ix, iy
	return
end

function update_🔥!(🔥, qx, qy, ∂t, ρCp, ∂x, ∂y, nx, ny)
	ix = (blockIdx().x-1) * blockDim().x + threadIdx().x # thread ID, dimension x
	iy = (blockIdx().y-1) * blockDim().y + threadIdx().y # thread ID, dimension y

	# TODO add 🔥 computation in a vectorised fashion here based on index ix, iy
	return
end

@views function heat_2D_gpu()
	# physics
	lx     = 10.0
	ly     = 10.0
	λ      = 1.0
	ρCp    = 1.0
	nt     = 200
	# numerics
	BLOCKX = 16
	BLOCKY = 16
	GRIDX  = 8
	GRIDY  = 8
	nx     = BLOCKX*GRIDX-1
	ny     = BLOCKY*GRIDY-1
	nout   = 10
	∂x     = lx/nx
	∂y     = ly/ny
	xc     = LinRange(∂x/2, lx-∂x/2, nx)
	yc     = LinRange(∂y/2, ly-∂y/2, ny)
	🔥     = CUDA.zeros(nx  ,ny  )
	qx     = CUDA.zeros(nx+1,ny  )
	qy     = CUDA.zeros(nx  ,ny+1)
	🔥     = CuArray( exp.(.-(xc.-lx./2.0).^2 .-(yc.-ly./2.0)'.^2) )
	∂t     = min(∂x^2,∂y^2)/ρCp/λ/4.1
	cuthreads = (BLOCKX, BLOCKY, 1)
	cublocks  = (GRIDX,  GRIDY,  1)
	# action
	t0     = Base.time()
	for it = 1:nt
		@cuda blocks=cublocks threads=cuthreads compute_flux!(qx, qy, 🔥, λ, ∂x, ∂y, nx, ny)
		@cuda blocks=cublocks threads=cuthreads update_🔥!(🔥, qx, qy, ∂t, ρCp, ∂x, ∂y, nx, ny)
		if mod(it,nout)==0 && viz
			display(heatmap(xc, yc, Array(🔥)', xlabel="lx", ylabel="ly", title="heat diffusion, it=$it", clims=(0.,1.)))
			# sleep(.01)
		end
	end
	@printf("T_eff = %1.2e GB/s \n", (2/1e9*nx*ny*sizeof(lx))/((Base.time()-t0)/nt))
	return
end

@time heat_2D_gpu()


text:
export FPSystem

"""
    FPSystem(polynomials, vars) <: AbstractSystem

Create a polynomial system using the [`FixedPolynomials`](https://github.com/JuliaAlgebra/FixedPolynomials.jl) package.
"""
struct FPSystem{T} <: AbstractSystem
    system::FP.System{T}
end

FPSystem(polys::Vector{<:MP.AbstractPolynomial}, vars) = FPSystem(FP.System(polys, vars))
function FPSystem(polys::Vector{<:MP.AbstractPolynomial};
    variables=nothing, parameters=nothing)

    parameters === nothing || error("Parameters are not supported by FPSystem")

    if variables === nothing
        FPSystem(FP.System(polys))
    else
        FPSystem(FP.System(polys, variables))
    end
 end

struct FPSystemCache{JC<:FP.JacobianConfig} <: AbstractSystemCache
    config::JC
end

cache(F::FPSystem, x) = FPSystemCache(FP.config(F.system, x))

Base.size(F::FPSystem) = (length(F.system), FP.nvariables(F.system))

evaluate!(u, F::FPSystem, x, c::FPSystemCache) = FP.evaluate!(u, F.system, x, c.config)
evaluate(F::FPSystem, x, c::FPSystemCache) = FP.evaluate(F.system, x, c.config)
jacobian!(U, F::FPSystem, x, c::FPSystemCache) = FP.jacobian!(U, F.system, x, c.config)
jacobian(F::FPSystem, x, c::FPSystemCache) = FP.jacobian(F.system, x, c.config)
function evaluate_and_jacobian!(u, U, F::FPSystem, x, c::FPSystemCache)
    FP.evaluate_and_jacobian!(u, U, F.system, x, c.config)
end
function evaluate_and_jacobian(F::FPSystem, x, c::FPSystemCache)
    FP.evaluate_and_jacobian(F.system, x, c.config)
end

weylnorm2(F::FPSystem) = FP.weyldot(F.system.polys, F.system.polys)


text:
using Flux
using Zygote: @adjoint
using Statistics
using NNlib

## CPU implementation

"""
  logaddexp(a, b)
Adds log-space `a` and `b` such that the result equals `log(exp(a)+exp(b))`
"""
function logaddexp(a, b)
  isinf(a) && return b
  isinf(b) && return a

  # always want the greater number on the left in the exponentiation;
  # the magnitude difference may end up making the number very positive
  # which will cause exp() to return Inf
  # E.g., a = -900, b = -800, will give exp(-800 - -900), which will be
  # Inf for Float32 values
  if a < b
    a, b = b, a
  end
  return a + log(1+exp(b-a))
end

"""
  add_blanks(z)

Adds blanks to the start and end of `z`, and between items in `z`
"""
function add_blanks(z, blank)
  z′ = fill(blank, 2*length(z) + 1)
  z′[2 .* eachindex(z)] = z
  return z′
end

function ctc_alpha(ŷ::AbstractArray, y)
  typed_zero = zero(ŷ[1])
  ŷ = logsoftmax(ŷ)
  blank = size(ŷ, 1)
  z′ = add_blanks(y, blank)
  T = size(ŷ, 2)
  U′ = length(z′)

  α = fill(log(typed_zero), U′, T)
  α[1,1] = ŷ[blank, 1]
  α[2,1] = ŷ[z′[2], 1]
  for t=2:T
    bound = max(1, U′ - 2(T - t) - 1)
    for u=bound:U′
      if u == 1
        α[u,t] = α[u, t-1]
      else
      α[u,t] = logaddexp(α[u, t-1], α[u-1, t-1])

      # array bounds check and f(u) function from Eq. 7.9
      if u > 2 && !(z′[u] == blank || z′[u-2] == z′[u])
        α[u,t] = logaddexp(α[u,t], α[u-2,t-1])
      end
    end
    α[u,t] += ŷ[z′[u], t]
    end
  end
  return (loss=-1 * logaddexp(α[end,T], α[end-1, T]), alpha=α, zprime=z′, logsoftyhat=ŷ)
end

function ∇ctc_loss(ŷ::AbstractArray, y, out)
  loss, α, z′, ŷ = out
  U′, T = size(α)
  blank = size(ŷ, 1)
  typed_zero = zero(first(α))

  # Calculate beta coefficients, from the bottom-right, to the upper-left
  β = fill(log(typed_zero), U′, T)

  # Fill bottom-right corner so bounding errors can be avoided
  # by starting `u` at `U′-1`
  β[U′, T] = typed_zero
  β[U′-1, T] = typed_zero

  # start at T-1 so that β(T, u) = log(0) for all u < U′ - 1
  for t=(T-1):-1:1
    bound = min(U′, 2t)
    for u=bound:-1:1
      if u == U′
        β[u,t] = ŷ[z′[u], t+1] + β[u, t+1]
      else
        β[u,t] = logaddexp(ŷ[z′[u], t+1] + β[u, t+1], ŷ[z′[u+1], t+1] + β[u+1,t+1])

        # array bounds check and g(u) function from Eq. 7.16
        if u+2 <= U′ && z′[u] != blank && z′[u] != z′[u+2]
          β[u,t] = logaddexp(β[u,t], ŷ[z′[u+2], t+1] + β[u+2, t+1])
        end
      end
    end
  end

  # Accumulate alpha-beta products for each category,
  # then calculate gradients
  accum = fill(log(typed_zero), size(ŷ))
  for t=1:T
    for u=1:U′
      accum[z′[u], t] = logaddexp(accum[z′[u], t], α[u,t] + β[u,t])
    end
  end
  grads = exp.(ŷ) .- exp.(accum .+ loss)
  return grads
end

"""
  ctc_loss(ŷ, y)

Computes the connectionist temporal classification loss between `ŷ`
and `y`.

`ŷ` must be a classes-by-time matrices, i.e., each row
represents a class and each column represents a time step.
Additionally, the `logsoftmax` function will be applied to `ŷ`, so
`ŷ` must be the raw activation values from the neural network and
not, for example, the activations after being passed through a
`softmax` activation function. `y` must be a 1D array of the labels
associated with `ŷ`. The blank label is assumed to be the last label
category in `ŷ`, so it is equivalent to `size(ŷ, 1)`.

Used for sequence-to-sequence classification problems such as
speech recognition and handwriting recognition where the exact
time-alignment of the output (e.g., letters) is not needed to
solve the problem. See [Graves et al. (2006)](https://www.cs.toronto.edu/~graves/icml_2006.pdf)
or [Graves (2012)](https://www.cs.toronto.edu/~graves/preprint.pdf#chapter.7)
for mathematical details.
"""
ctc_loss(ŷ::AbstractArray, y) = ctc_alpha(ŷ, y).loss

function ChainRulesCore.rrule(::typeof(ctc_loss), ŷ, y)
  tmp = ctc_alpha(ŷ, y)
  ctc_loss_pullback(Δ) = (NoTangent(), Δ .* ∇ctc_loss(ŷ, y, tmp), NoTangent())
  return tmp.loss, ctc_loss_pullback
end

## GPU implementation

# a port of the GPU kernels from Baidu's C++ warp-ctc package,
# which itself is Copyright 2015-2016 Baidu USA LLC
# and available under the Apache 2.0 license
#
# Apache 2.0 license: https://www.apache.org/licenses/LICENSE-2.0
# GitHub: https://github.com/baidu-research/warp-ctc/
# paper: https://arxiv.org/pdf/1512.02595.pdf

const MAX_THREADS = 256

function log_plus_f(p1, p2)
  isinf(p1) && return p2
  isinf(p2) && return p1
  if p1 < p2
    p1, p2 = p2, p1
  end
  return p1 + log(1+exp(p2 - p1))
end

function count_repeats(A)
  repeats = 0
  for (i,elem) in enumerate(A)
    if i > 1 && A[i] == A[i-1]
      repeats += 1
    end
  end
  return repeats
end

function compute_alpha_kernel(probs, labelSize, uttLength, repeats, labelsWithoutBlanks, labelsWithBlanks, alpha, blankLabel)

  tid = threadIdx().x
  L = labelSize
  T = uttLength
  S = length(labelsWithBlanks)

  if L + repeats > T
    return nothing
  end
  labels = labelsWithBlanks

  # Corner-case checking
  start = (L + repeats <= T) ? 0 : 1
  last = S > 1 ? 2 : 1

  # Fill in first column (time step)
  i = tid
  while i <= last - start
    alpha[start+i, 1] = probs[labels[start+i], 1]
    i += blockDim().x
  end
  sync_threads()

  # Fill in coefficients for each time step
  for t=2:T
    # Corner-case checking
    if tid == 1 && !(1 < S - 2*(T-t) - 1)
      if start == 0
        alpha[1, t] = alpha[1, t-1] + probs[blankLabel, t]
      elseif start == 1
        alpha[1, t] = alpha[1, t-1]
      end
    end
    sync_threads()

    # Fill in coefficients for each label class in the target output sequence;
    # each thread will process the calculations for one class
    idx = tid+1
    while idx <= S
      prevSum = log_plus_f(alpha[idx, t-1], alpha[idx-1, t-1])
      if labels[idx] != blankLabel && idx != 2 && labels[idx] != labels[idx-2]
        prevSum = log_plus_f(prevSum, alpha[idx-2, t-1])
      end
      if idx < S - 2*(T-t) - 1
        alpha[idx, t] = -Inf32
      else
        alpha[idx, t] = prevSum + probs[labels[idx], t]
      end
      idx += blockDim().x
    end
    sync_threads()
  end
  return nothing
end

function compute_beta_and_grad_kernel(probs, labelSize, uttLength,
                  repeatsInLabel, labelsWithBlanks,
                  alphas, beta, output, accum,
                  grad, blankLabel, loss)

  tid = threadIdx().x
  L = labelSize
  T = uttLength
  S = 2*L + 1
  repeats = repeatsInLabel
  labels = labelsWithBlanks

  if (L+repeats) > T
    return nothing
  end

  # Corner-case checking
  start = S > 1 ? S-2 : 0
  last = L + repeats < T ? S : S-1
  sync_threads()
  i = tid

  # Calculate coefficients for last column (time step)
  # then determine alpha and beta product
  while i <= last - start
    beta[i+start, T] = 0
    output[i+start, T] = beta[i+start, T] + alphas[i+start, T]
    i += blockDim().x
  end
  sync_threads()

  # Fill in `accum` for last column (time step)
  if tid == 1
    for i=1:S
      labelIdx = labels[i]
      accum[labelIdx, T] = log_plus_f(accum[labelIdx, T], output[i, T])
    end
  end
  sync_threads()

  # Fill in `grad` for last column (time step)
  idx = tid
  while idx <= size(grad, 1)
    s = -Inf32
    for i=1:S
      s = log_plus_f(s, output[i, T])
    end

    # ∂L/∂a (where a is activation before logsoftmax)
    grad[idx, T] = exp(probs[idx, T]) - exp(accum[idx, T] - s)
    idx += blockDim().x
  end
  sync_threads()

  # Fill in the rest of the coefficients
  t = T-1
  while t >= 1
    if t < T
      idx = tid
      while idx <= S
        nextSum = probs[labels[idx], t+1] + beta[idx, t+1]
        if idx < S
          nextSum = log_plus_f(nextSum,
            probs[labels[idx+1], t+1] + beta[idx+1, t+1])
        end
        if labels[idx] != blankLabel && idx != S-1 && labels[idx] != labels[idx+2]
          nextSum = log_plus_f(nextSum,
            probs[labels[idx+2], t+1] + beta[idx + 2, t+1])
        end
        if idx > 2*t
          beta[idx, t] = -Inf32
        else
          beta[idx, t] = nextSum
        end
        idx += blockDim().x
      end
      sync_threads()
      idx = tid
      while idx <= S
        output[idx, t] = alphas[idx, t] + beta[idx, t]
        idx += blockDim().x
      end
      sync_threads()
    end
    sync_threads()

    # Calculate accumulated alpha-beta products for each label class for
    # each time step; used in calculating gradients
    if tid == 1
      for i=1:S
        labelIdx = labels[i]
        accum[labelIdx, t] = log_plus_f(accum[labelIdx, t], output[i, t])
      end
    end
    sync_threads()
    idx = tid

    # Calculate gradients
    while idx <= size(grad, 1)

      # ∂L/∂a (where a is activation before logsoftmax)
      grad[idx, t] = exp(probs[idx, t]) - exp(accum[idx, t] + loss)
      idx += blockDim().x
    end
    sync_threads()
    t -= 1
    sync_threads()
  end
  return nothing
end

function ctc_alpha(ŷ::CuArray, y)
  ŷ = logsoftmax(ŷ)
  blank = size(ŷ, 1)
  ycu = cu(y)
  z′ = CUDA.fill(blank, 2 * length(y) + 1)
  z′[eachindex(y) .* 2] .= ycu
  T = size(ŷ, 2)
  U′ = 2*length(y) + 1
  alphas = CUDA.fill(log(zero(eltype(ŷ))), U′,T)
  nRepeats = count_repeats(cpu(y))
  nThreads = min(U′, MAX_THREADS)
  @cuda blocks=1 threads=nThreads compute_alpha_kernel(ŷ, length(y), T, nRepeats, ycu, z′, alphas, blank)
  return (loss=-1 * logsumexp(alphas[end-1:end]), alpha=alphas, z′=z′, yhat=ŷ, nRepeats=nRepeats)
end

ctc_loss(ŷ::CuArray, y) = ctc_alpha(ŷ::CuArray, y).loss

function ∇ctc_loss(ŷ::CuArray, y, out)
  loss, alphas, z′, ŷ, nRepeats = out
  U′, T = size(alphas)
  blank = size(ŷ, 1)
  typed_zero = zero(eltype(ŷ))
  betas = CUDA.fill(log(typed_zero), U′, T)
  output = CUDA.fill(log(typed_zero), U′, T)
  nThreads = min(U′, MAX_THREADS)
  grads = CUDA.fill(log(typed_zero), size(ŷ))
  accum = CUDA.fill(log(typed_zero), size(ŷ))
  @cuda blocks=1 threads=nThreads compute_beta_and_grad_kernel(ŷ, length(y), T, nRepeats, CuArray(z′), alphas, betas, output, accum, grads, blank, loss)
  return grads
end