text:
"""
    rte_lw_noscat_source!(
        src_lw::SourceLWNoScat{FT},
        op::OneScalar{FT},
        gcol,
        nlay,
    ) where {FT<:AbstractFloat}

Compute LW source function for upward and downward emission at levels using linear-in-tau assumption
See Clough et al., 1992, doi: 10.1029/92JD01419, Eq 13
"""
function rte_lw_noscat_source!(
    src_lw::SourceLWNoScat{FT},
    op::OneScalar{FT},
    gcol,
    nlay,
) where {FT<:AbstractFloat}
    # setting references
    @unpack src_up, src_dn, lev_source_inc, lev_source_dec, lay_source = src_lw
    #    @unpack nlay = slv.as
    Ds = op.angle_disc.gauss_Ds
    τ = op.τ

    τ_thresh = sqrt(eps(FT)) # or abs(eps(FT))?

    @inbounds for glay = 1:nlay
        τ_loc = τ[glay, gcol] * Ds[1]      # Optical path and transmission,

        trans = exp(-τ_loc)    # used in source function and transport calculations
        # Weighting factor. Use 2nd order series expansion when rounding error (~tau^2)
        # is of order epsilon (smallest difference from 1. in working precision)
        # Thanks to Peter Blossey

        fact = (τ_loc > τ_thresh) ? ((FT(1) - trans) / τ_loc - trans) :
            τ_loc * (FT(1 / 2) - FT(1 / 3) * τ_loc)

        # Equations below are developed in Clough et al., 1992, doi:10.1029/92JD01419, Eq 13

        src_up[glay, gcol] =
            (FT(1) - trans) * lev_source_inc[glay, gcol] +
            FT(2) * fact * (lay_source[glay, gcol] - lev_source_inc[glay, gcol])

        src_dn[glay, gcol] =
            (FT(1) - trans) * lev_source_dec[glay, gcol] +
            FT(2) * fact * (lay_source[glay, gcol] - lev_source_dec[glay, gcol])
    end
end

"""
    rte_lw_noscat_transport!(
        src_lw::SourceLWNoScat{FT},
        bcs_lw::LwBCs{FT},
        op::OneScalar{FT},
        gcol,
        flux::FluxLW{FT},
        igpt,
        ibnd,
        nlay,
        nlev,
    ) where {FT<:AbstractFloat}

Transport for no-scattering longwave problem.
"""
function rte_lw_noscat_transport!(
    src_lw::SourceLWNoScat{FT},
    bcs_lw::LwBCs{FT},
    op::OneScalar{FT},
    gcol,
    flux::FluxLW{FT},
    igpt,
    ibnd,
    nlay,
    nlev,
) where {FT<:AbstractFloat}
    # setting references
    @unpack src_up, src_dn, sfc_source = src_lw
    @unpack sfc_emis, inc_flux = bcs_lw
    @unpack flux_up, flux_dn, flux_net = flux

    Ds = op.angle_disc.gauss_Ds
    w_μ = op.angle_disc.gauss_wts
    n_μ = op.angle_disc.n_gauss_angles
    τ = op.τ

    @inbounds for ilev = 1:nlev
        flux_dn[ilev, gcol] = FT(0)
        flux_up[ilev, gcol] = FT(0)
    end

    if inc_flux ≠ nothing
        @inbounds flux_dn[nlev, gcol] = inc_flux[gcol, igpt]
    end
    # Transport is for intensity
    #   convert flux at top of domain to intensity assuming azimuthal isotropy
    @inbounds flux_dn[nlev, gcol] /= (FT(2) * FT(π) * w_μ[1])

    # Top of domain is index nlev
    # Downward propagation
    @inbounds for ilev = nlay:-1:1
        trans = exp(-τ[ilev, gcol] * Ds[1])
        flux_dn[ilev, gcol] = trans * flux_dn[ilev+1, gcol] + src_dn[ilev, gcol]
    end

    # Surface reflection and emission
    @inbounds flux_up[1, gcol] =
        flux_dn[1, gcol] * (FT(1) - sfc_emis[ibnd, gcol]) +
        sfc_emis[ibnd, gcol] * sfc_source[gcol]

    # Upward propagation
    @inbounds for ilev = 2:nlay+1
        trans = exp(-τ[ilev-1, gcol] * Ds[1])
        flux_up[ilev, gcol] =
            trans * flux_up[ilev-1, gcol] + src_up[ilev-1, gcol]
    end

    @inbounds for ilev = 1:nlev
        flux_up[ilev, gcol] *= (FT(2) * FT(π) * w_μ[1])
        flux_dn[ilev, gcol] *= (FT(2) * FT(π) * w_μ[1])
        flux_net[ilev, gcol] = flux_up[ilev, gcol] - flux_dn[ilev, gcol]
    end
end

"""
    rte_lw_2stream_combine_sources!(src::SourceLW2Str, gcol, nlev, ncol)

RRTMGP provides source functions at each level using the spectral mapping
of each adjacent layer. This function combines these for two-stream calculations.
"""
function rte_lw_2stream_combine_sources!(src::SourceLW2Str, gcol, nlev, ncol)
    @inbounds src.lev_source[1, gcol] = src.lev_source_dec[1, gcol] # glev == 1
    @inbounds src.lev_source[nlev, gcol] = src.lev_source_inc[nlev-1, gcol] # glev == nlev
    @inbounds for glev = 2:nlev-1
        src.lev_source[glev, gcol] = sqrt(
            src.lev_source_dec[glev, gcol] * src.lev_source_inc[glev-1, gcol],
        )
    end
end

"""
    rte_lw_2stream_source!(
        op::TwoStream{FT},
        src_lw::SourceLW2Str{FT},
        gcol,
        nlay,
        ncol,
    ) where {FT<:AbstractFloat}

This function combines RRTMGP-specific sources at levels,
computes layer reflectance, transmittance, and
total source function at levels using linear-in-tau approximation.
"""
function rte_lw_2stream_source!(
    op::TwoStream{FT},
    src_lw::SourceLW2Str{FT},
    gcol,
    nlay,
    ncol,
) where {FT<:AbstractFloat}
    # Cell properties: reflection, transmission for diffuse radiation
    # Coupling coefficients needed for source function
    # -------------------------------------------------------------------------------------------------
    #
    # Longwave two-stream solutions to diffuse reflectance and transmittance for a layer
    #    with optical depth tau, single scattering albedo w0, and asymmetery parameter g.
    #
    # Equations are developed in Meador and Weaver, 1980,
    #    doi:10.1175/1520-0469(1980)037<0630:TSATRT>2.0.CO;2
    #
    # -------------------------------------------------------------------------------------------------
    # setting references
    @unpack τ, ssa, g = op
    @unpack Rdif, Tdif, lev_source, src_up, src_dn = src_lw

    lw_diff_sec = FT(1.66)

    @inbounds for glay = 1:nlay
        γ1 = lw_diff_sec * (1 - FT(0.5) * ssa[glay, gcol] * (1 + g[glay, gcol]))
        γ2 = lw_diff_sec * FT(0.5) * ssa[glay, gcol] * (1 - g[glay, gcol])
        k = sqrt(max((γ1 + γ2) * (γ1 - γ2), FT(1e-12)))

        # Refactored to avoid rounding errors when k, gamma1 are of very different magnitudes
        RT_term =
            1 / (
                k * (1 + exp(-2 * τ[glay, gcol] * k)) +
                γ1 * (1 - exp(-2 * τ[glay, gcol] * k))
            )

        @inbounds Rdif[glay, gcol] =
            RT_term * γ2 * (1 - exp(-2 * τ[glay, gcol] * k)) # Equation 25
        @inbounds Tdif[glay, gcol] = RT_term * 2 * k * exp(-τ[glay, gcol] * k) # Equation 26

        # Source function for diffuse radiation
        # Compute LW source function for upward and downward emission at levels using linear-in-tau assumption
        # This version straight from ECRAD
        # Source is provided as W/m2-str; factor of pi converts to flux units
        # lw_source_2str

        top = glay + 1
        bot = glay

        if τ[glay, gcol] > 1.0e-8
            # Toon et al. (JGR 1989) Eqs 26-27
            Z =
                (lev_source[bot, gcol] - lev_source[top, gcol]) /
                (τ[glay, gcol] * (γ1 + γ2))
            Zup_top = Z + lev_source[top, gcol]
            Zup_bottom = Z + lev_source[bot, gcol]
            Zdn_top = -Z + lev_source[top, gcol]
            Zdn_bottom = -Z + lev_source[bot, gcol]
            @inbounds src_up[glay, gcol] =
                pi * (
                    Zup_top - Rdif[glay, gcol] * Zdn_top -
                    Tdif[glay, gcol] * Zup_bottom
                )
            @inbounds src_dn[glay, gcol] =
                pi * (
                    Zdn_bottom - Rdif[glay, gcol] * Zup_bottom -
                    Tdif[glay, gcol] * Zdn_top
                )
        else
            @inbounds src_up[glay, gcol] = FT(0)
            @inbounds src_dn[glay, gcol] = FT(0)
        end
    end
end

"""
    adding_lw!(
        flux::FluxLW{FT},
        src_lw::SL,
        bcs_lw::BCL,
        gcol,
        igpt,
        ibnd,
        nlev,
        ncol,
    ) where {FT<:AbstractFloat,SL<:SourceLW2Str{FT},BCL<:LwBCs{FT}}

Transport of diffuse radiation through a vertically layered atmosphere, for the longwave problem.
Equations are after Shonk and Hogan 2008, doi:10.1175/2007JCLI1940.1 (SH08)
"""
function adding_lw!(
    flux::FluxLW{FT},
    src_lw::SL,
    bcs_lw::BCL,
    gcol,
    igpt,
    ibnd,
    nlev,
    ncol,
) where {FT<:AbstractFloat,SL<:SourceLW2Str{FT},BCL<:LwBCs{FT}}
    nlay = nlev - 1
    # setting references
    @unpack flux_up, flux_dn, flux_net = flux

    @unpack albedo, sfc_source, Rdif, Tdif, src_up, src_dn, src = src_lw
    @unpack inc_flux, sfc_emis = bcs_lw

    @inbounds for ilev = 1:nlev
        flux_dn[ilev, gcol] = FT(0)
        flux_up[ilev, gcol] = FT(0)
    end

    if inc_flux ≠ nothing
        @inbounds flux_dn[nlev, gcol] = inc_flux[gcol, igpt]
    end
    # Albedo of lowest level is the surface albedo...
    @inbounds albedo[1, gcol] = FT(1) - sfc_emis[ibnd, gcol]
    # ... and source of diffuse radiation is surface emission
    @inbounds src[1, gcol] = FT(π) * sfc_emis[ibnd, gcol] * sfc_source[gcol]

    # From bottom to top of atmosphere --
    #   compute albedo and source of upward radiation
    @inbounds for ilev = 1:nlay
        denom = FT(1) / (FT(1) - Rdif[ilev, gcol] * albedo[ilev, gcol])  # Eq 10
        albedo[ilev+1, gcol] =
            Rdif[ilev, gcol] +
            Tdif[ilev, gcol] * Tdif[ilev, gcol] * albedo[ilev, gcol] * denom # Equation 9
        #
        # Equation 11 -- source is emitted upward radiation at top of layer plus
        # radiation emitted at bottom of layer,
        # transmitted through the layer and reflected from layers below (Tdiff*src*albedo)
        src[ilev+1, gcol] =
            src_up[ilev, gcol] +
            Tdif[ilev, gcol] *
            denom *
            (src[ilev, gcol] + albedo[ilev, gcol] * src_dn[ilev, gcol])
    end

    # Eq 12, at the top of the domain upwelling diffuse is due to ...
    @inbounds flux_up[nlev, gcol] =
        flux_dn[nlev, gcol] * albedo[nlev, gcol] + # ... reflection of incident diffuse and
        src[nlev, gcol]                          # scattering by the direct beam below

    # From the top of the atmosphere downward -- compute fluxes
    @inbounds for ilev = nlay:-1:1
        denom = FT(1) / (FT(1) - Rdif[ilev, gcol] * albedo[ilev, gcol])  # Eq 10
        flux_dn[ilev, gcol] =
            (
                Tdif[ilev, gcol] * flux_dn[ilev+1, gcol] + # Equation 13
                Rdif[ilev, gcol] * src[ilev, gcol] +
                src_dn[ilev, gcol]
            ) * denom
        flux_up[ilev, gcol] =
            flux_dn[ilev, gcol] * albedo[ilev, gcol] + # Equation 12
            src[ilev, gcol]
    end

    @inbounds for ilev = 1:nlev
        flux_net[ilev, gcol] = flux_up[ilev, gcol] - flux_dn[ilev, gcol]
    end
end
# ---------------------------------------------------------------
"""
    rte_sw_noscat_solve_kernel!(
        flux::FluxSW{FT},
        op::OneScalar{FT},
        bcs_sw::SwBCs{FT},
        solar_frac::FT,
        gcol,
        nlev,
    ) where {FT<:AbstractFloat}

No-scattering solver for the shortwave problem.
(Extinction-only i.e. solar direct beam)
"""
function rte_sw_noscat_solve_kernel!(
    flux::FluxSW{FT},
    op::OneScalar{FT},
    bcs_sw::SwBCs{FT},
    solar_frac::FT,
    gcol,
    nlev,
) where {FT<:AbstractFloat}
    @unpack toa_flux, zenith = bcs_sw
    τ = op.τ
    flux_dn_dir = flux.flux_dn_dir
    # downward propagation
    flux_dn_dir[nlev, gcol] = toa_flux[gcol] * solar_frac * cos(zenith[gcol])
    for ilev = nlev-1:-1:1
        flux_dn_dir[ilev, gcol] =
            flux_dn_dir[ilev+1, gcol] * exp(-τ[ilev, gcol] / cos(zenith[gcol]))
    end
end

"""
    sw_two_stream!(
        op::TwoStream{FT},
        src_sw::SourceSW2Str{FT},
        bcs_sw::SwBCs{FT},
        gcol,
        nlay,
    ) where {FT<:AbstractFloat}

Computes cell properties (transmittance and reflectance) for direct and diffuse radiation

Two-stream solutions to direct and diffuse reflectance and transmittance for a layer
with optical depth tau, single scattering albedo w0, and asymmetery parameter g.

Equations are developed in Meador and Weaver, 1980,
doi:10.1175/1520-0469(1980)037<0630:TSATRT>2.0.CO;2
"""
function sw_two_stream!(
    op::TwoStream{FT},
    src_sw::SourceSW2Str{FT},
    bcs_sw::SwBCs{FT},
    gcol,
    nlay,
) where {FT<:AbstractFloat}
    zenith = bcs_sw.zenith
    @unpack τ, ssa, g = op
    @unpack Rdif, Tdif, Rdir, Tdir, Tnoscat = src_sw

    @inbounds for glay = 1:nlay
        # Zdunkowski Practical Improved Flux Method "PIFM"
        #  (Zdunkowski et al., 1980;  Contributions to Atmospheric Physics 53, 147-66)
        γ1 =
            (FT(8) - ssa[glay, gcol] * (FT(5) + FT(3) * g[glay, gcol])) *
            FT(0.25)
        γ2 = FT(3) * (ssa[glay, gcol] * (FT(1) - g[glay, gcol])) * FT(0.25)
        γ3 = (FT(2) - (FT(3) * cos(zenith[gcol])) * g[glay, gcol]) * FT(0.25)
        γ4 = FT(1) - γ3
        α1 = γ1 * γ4 + γ2 * γ3                          # Eq. 16
        α2 = γ1 * γ3 + γ2 * γ4                          # Eq. 17
        k = sqrt(max((γ1 - γ2) * (γ1 + γ2), FT(1e-12)))

        exp_minusktau = exp(-τ[glay, gcol] * k)
        exp_minus2ktau = exp_minusktau * exp_minusktau

        # Refactored to avoid rounding errors when k, gamma1 are of very different magnitudes
        RT_term =
            FT(1) /
            (k * (FT(1) + exp_minus2ktau) + γ1 * (FT(1) - exp_minus2ktau))

        Rdif[glay, gcol] = RT_term * γ2 * (FT(1) - exp_minus2ktau) # Eqn. 25
        Tdif[glay, gcol] = RT_term * FT(2) * k * exp_minusktau     # Eqn. 26

        # Transmittance of direct, unscattered beam. Also used below
        Tnoscat[glay, gcol] = exp(-τ[glay, gcol] / cos(zenith[gcol]))

        # Direct reflect and transmission
        k_μ = k * cos(zenith[gcol])
        k_γ3 = k * γ3
        k_γ4 = k * γ4

        # Equation 14, multiplying top and bottom by exp(-k*tau)
        #   and rearranging to avoid div by 0.
        if abs(FT(1) - k_μ * k_μ) ≥ eps(FT)
            RT_term = ssa[glay, gcol] * RT_term / (FT(1) - k_μ * k_μ)
        else
            RT_term = ssa[glay, gcol] * RT_term / eps(FT)
        end

        Rdir[glay, gcol] =
            RT_term * (
                (FT(1) - k_μ) * (α2 + k_γ3) -
                (FT(1) + k_μ) * (α2 - k_γ3) * exp_minus2ktau -
                FT(2) * (k_γ3 - α2 * k_μ) * exp_minusktau * Tnoscat[glay, gcol]
            )
        #
        # Equation 15, multiplying top and bottom by exp(-k*tau),
        #   multiplying through by exp(-tau/mu0) to
        #   prefer underflow to overflow
        # Omitting direct transmittance
        #
        Tdir[glay, gcol] =
            -RT_term * (
                (FT(1) + k_μ) * (α1 + k_γ4) * Tnoscat[glay, gcol] -
                (FT(1) - k_μ) *
                (α1 - k_γ4) *
                exp_minus2ktau *
                Tnoscat[glay, gcol] - FT(2) * (k_γ4 + α1 * k_μ) * exp_minusktau
            )
    end
end

"""
    sw_source_2str!(
        src_sw::SourceSW2Str{FT},
        bcs_sw::SwBCs{FT},
        gcol,
        flux::FluxSW{FT},
        solar_frac::FT,
        ibnd,
        nlay,
    ) where {FT<:AbstractFloat}

Direct-beam and source for diffuse radiation
"""
function sw_source_2str!(
    src_sw::SourceSW2Str{FT},
    bcs_sw::SwBCs{FT},
    gcol,
    flux::FluxSW{FT},
    solar_frac::FT,
    ibnd,
    nlay,
) where {FT<:AbstractFloat}
    @unpack toa_flux, zenith, sfc_alb_direct = bcs_sw
    @unpack Rdir, Tdir, Tnoscat, src_up, src_dn, sfc_source = src_sw
    flux_dn_dir = flux.flux_dn_dir

    # layer index = level index
    # previous level is up (+1)
    flux_dn_dir[nlay+1, gcol] = toa_flux[gcol] * solar_frac * cos(zenith[gcol])
    for ilev = nlay:-1:1
        src_up[ilev, gcol] = Rdir[ilev, gcol] * flux_dn_dir[ilev+1, gcol]
        src_dn[ilev, gcol] = Tdir[ilev, gcol] * flux_dn_dir[ilev+1, gcol]
        flux_dn_dir[ilev, gcol] =
            Tnoscat[ilev, gcol] * flux_dn_dir[ilev+1, gcol]
    end
    sfc_source[gcol] = flux_dn_dir[1, gcol] * sfc_alb_direct[ibnd, gcol]
end

"""
    adding_sw!(
        src_sw::SourceSW2Str{FT},
        bcs_sw::SwBCs{FT},
        gcol,
        flux::FluxSW{FT},
        igpt,
        ibnd,
        nlev,
    ) where {FT<:AbstractFloat}

Transport for the two stream shortwave problem.

Transport of diffuse radiation through a vertically layered atmosphere.
Equations are after Shonk and Hogan 2008, doi:10.1175/2007JCLI1940.1 (SH08)
"""
function adding_sw!(
    src_sw::SourceSW2Str{FT},
    bcs_sw::SwBCs{FT},
    gcol,
    flux::FluxSW{FT},
    igpt,
    ibnd,
    nlev,
) where {FT<:AbstractFloat}
    nlay = nlev - 1
    # setting references
    @unpack flux_up, flux_dn, flux_net = flux

    @unpack albedo, sfc_source, Rdif, Tdif, src_up, src_dn, src = src_sw

    @inbounds for ilev = 1:nlev
        flux_dn[ilev, gcol] = FT(0)
        flux_up[ilev, gcol] = FT(0)
    end

    # Albedo of lowest level is the surface albedo...
    @inbounds albedo[1, gcol] = bcs_sw.sfc_alb_diffuse[ibnd, gcol]
    # ... and source of diffuse radiation is surface emission
    @inbounds src[1, gcol] = sfc_source[gcol]
    # From bottom to top of atmosphere --
    #   compute albedo and source of upward radiation
    @inbounds for ilev = 1:nlay
        denom = FT(1) / (FT(1) - Rdif[ilev, gcol] * albedo[ilev, gcol])  # Eq 10
        albedo[ilev+1, gcol] =
            Rdif[ilev, gcol] +
            Tdif[ilev, gcol] * Tdif[ilev, gcol] * albedo[ilev, gcol] * denom # Equation 9
        #
        # Equation 11 -- source is emitted upward radiation at top of layer plus
        # radiation emitted at bottom of layer,
        # transmitted through the layer and reflected from layers below (Tdiff*src*albedo)
        src[ilev+1, gcol] =
            src_up[ilev, gcol] +
            Tdif[ilev, gcol] *
            denom *
            (src[ilev, gcol] + albedo[ilev, gcol] * src_dn[ilev, gcol])
    end

    # Eq 12, at the top of the domain upwelling diffuse is due to ...
    flux_up[nlev, gcol] =
        flux_dn[nlev, gcol] * albedo[nlev, gcol] + # ... reflection of incident diffuse and
        src[nlev, gcol]                          # scattering by the direct beam below

    # From the top of the atmosphere downward -- compute fluxes
    @inbounds for ilev = nlay:-1:1
        denom = FT(1) / (FT(1) - Rdif[ilev, gcol] * albedo[ilev, gcol])  # Eq 10
        flux_dn[ilev, gcol] =
            (
                Tdif[ilev, gcol] * flux_dn[ilev+1, gcol] + # Equation 13
                Rdif[ilev, gcol] * src[ilev, gcol] +
                src_dn[ilev, gcol]
            ) * denom
        flux_up[ilev, gcol] =
            flux_dn[ilev, gcol] * albedo[ilev, gcol] + # Equation 12
            src[ilev, gcol]
    end

    @inbounds for ilev = 1:nlev
        flux_net[ilev, gcol] = flux_up[ilev, gcol] - flux_dn[ilev, gcol]
    end
end
# ---------------------------------------------------------------


text:
"""
	mutable struct GLVertexArray

Submitting vertex data for rendering requires creating a stream of vertices, and then telling OpenGL how to interpret that stream. A Vertex Array Object (VAO) is an OpenGL Object that stores all of the state needed to supply vertex data
"""
mutable struct GLVertexArray

	id::Int32

	# constructor
	function GLVertexArray()
		ret=new(-1)
		finalizer(releaseGpuResources, ret)
		return ret
	end

end

"""
	releaseGpuResources(array::GLVertexArray)

Release an array of `GLVertexArray`, a Gpu resource. It is a list of indices that will select which vertices to use and in which order.
"""
function releaseGpuResources(array::GLVertexArray)
	global __release_gpu_resources__
	if array.id>=0
		id=array.id
		array.id=-1
		glDeleteLater(function() glDeleteVertexArrays(1,[id]) end)
	end
end

"""
	enableVertexArray(array::GLVertexArray)

# Example

There are many ways for OpenGL to interpret a stream of, for example, 12 vertices. It can interpret the vertices as a sequence of triangles, points, or lines. It can even interpret these differently; it can interpret 12 vertices as 4 independent triangles (take every 3 verts as a triangle), as 10 dependent triangles (every group of 3 sequential vertices in the stream is a triangle), and so on.
Each attribute can be enabled or disabled for array access.
"""
function enableVertexArray(array::GLVertexArray)

	# not needed or osx
	if Sys.isapple() return end

	if array.id<0
		array.id=glGenVertexArray()
	end
	glBindVertexArray(array.id)
end

"""
	disableVertexArray(array::GLVertexArray)

Each attribute can be enabled or disabled for array access.
"""
function disableVertexArray(array::GLVertexArray)

	# not needed or osx
	if Sys.isapple() return end

	glBindVertexArray(0)
end

ext: jl


text:
using CleanCMB
using Healpix
using Printf
using Statistics
using Plots
# %% Number of frequency bands
nν = 5
# %% Read in maps at Nside=128. The original map resolution is Nside=512, but we use degraded maps for faster demo.
# ILC region definition map
bit, rgn = (
    readMapFromFITS("data/wmap_ilc_rgn_defn_r7_9yr_v5.fits", ic, Float64) for
    ic = 1:2
)
nside = bit.resolution.nside
# WMAP temperature maps (smootehd to 1 degree Gaussian beam)
map = []
band = ["K", "Ka", "Q", "V", "W"]
for iν = 1:nν
    filename = @sprintf("data/wmap_band_smth_imap_r7_9yr_%s_v5.fits", band[iν])
    push!(map, readMapFromFITS(filename, 1, Float64))
end
# %% Perform ILC to obtain a clean CMB from WMAP temperature maps
Tbar = zeros(nν)
cij = zeros(nν, nν)
clean_map = HealpixMap{Float64,RingOrder}(nside)
for N = 0:11 # loop over ILC regions 0-11
    set = (Int.(bit) .& 2^N) / 2^N # see https://lambda.gsfc.nasa.gov/product/map/dr5/ilc_map_get.cfm
    ip = findall(x -> x == 1, set)
    # Compute nν-by-nν covariance marix that will be used by `ilc_weights(cij)`
    for iν = 1:nν
        Tbar[iν] = mean(map[iν][ip])
    end
    for iν = 1:nν, jν = iν:nν
        cij[iν, jν] =
            mean((map[iν][ip] .- Tbar[iν]) .* (map[jν][ip] .- Tbar[jν]))
    end
    weights = ilc_weights(cij) # Compute ILC weights
    # Obtain a clean CMB map
    ip = findall(x -> x == N, rgn)
    for iν = 1:nν
        clean_map[ip] += weights[iν] * map[iν][ip]
    end
end
# %% Plot a clean CMB map and save to "ilc.png"
p = plot(clean_map, clims = (-0.25, 0.25), c = :jet)
savefig("ilc.png")
display(p)


text:
module PySCF

using PyCall: pyimport

pyscf = pyimport("pyscf")
mp = pyimport("pyscf.mp")   # Had to import mp alone ??!
cc = pyimport("pyscf.cc") # Had to import mp alone ??!

# Utilities
function pyscf_atom_from_xyz(fpath::String)
   join(split(read(open(fpath),String),"\n")[3:end],"\n")
end

function index(i::Int,j::Int)
   m,M = minmax(i-1,j-1)
   M*(M+1)÷2 + m + 1
end

# Compound indices ijkl
function get_4idx(i::Int,j::Int,k::Int,l::Int)
   ij = index(i,j)
   kl = index(k,l)
   index(ij,kl)
end

# Calculation of the electronic structure of a molecule
# with the help of PySCF to calculate AO integrals and hold molecular data
# * Restricted Hartree-Fock
# * Møller-Plesset order 2
# * Coupled Cluster Singles and Doubles
# Calculates the energy of any molecule.
# Instructions at https://github.com/CrawfordGroup/ProgrammingProjects
# may be of interest.

using Formatting: printfmt
using LinearAlgebra: Diagonal, Hermitian, eigen, norm, tr, diag, dot

include("diis.jl")
include("scf.jl")
include("mp2.jl")
include("ccsd.jl")
include("properties.jl")

end


text:
# This file is a part of Actors.jl, licensed under the MIT License (MIT).

export Actor
abstract Actor

export ActorMsgFrame
immutable ActorMsgFrame
    replyto::Actor
    message::Any
end

export ActorID
typealias ActorID UInt64

export ActorInbox
typealias ActorInbox Channel{Any}

export AbstractLocalActor
abstract AbstractLocalActor <: Actor

export LocalActor
immutable LocalActor <: AbstractLocalActor
   task::Task
end

export LocalActorAltInbox
immutable LocalActorAltInbox <: AbstractLocalActor
    actor::LocalActor
    inbox::ActorInbox
end

export ActorContext
immutable ActorContext
   id::ActorID
   inbox::ActorInbox
   self_ref::LocalActor
end

export actor_id
@inline actor_id(context::ActorContext) = context.id

export actor_inbox
@inline actor_inbox(context::ActorContext) = context.inbox

export localactor
@inline localactor(context::ActorContext) = context.self_ref

export actor_context
actor_context(task::Task) = task.storage[:_actor_context]::ActorContext

export actor_id
@inline actor_id(task::Task) = actor_id(actor_context(task))

export actor_inbox
@inline actor_inbox(task::Task) = actor_inbox(actor_context(task))

export localactor
@inline localactor(task::Task) = localactor(actor_context(task))

Base.show(io::IO, actor::LocalActor) = print(io, "local-actor-$(hex(actor_id(actor)))")
Base.show(io::IO, actor::LocalActorAltInbox) = print(io, "local-actor-alt-$(hex(actor_id(actor)))")

@inline Base.wait(actor::AbstractLocalActor) = wait(actor_task(actor))

function Serializer.serialize(s::SerializationState, actor::LocalActor)
    info("Serializing $actor (dummy implementation)")
    serialize(s, actor_id(actor))
end

export actor_task
@inline actor_task(actor::LocalActor) = actor.task

@inline actor_context(actor::AbstractLocalActor) = actor_context(actor_task(actor))
@inline actor_id(actor::AbstractLocalActor) = actor_id(actor_task(actor))

@inline actor_inbox(actor::LocalActor) = actor_inbox(actor_task(actor))

@inline actor_task(actor::LocalActorAltInbox) = actor_task(actor.actor)
@inline actor_inbox(actor::LocalActorAltInbox) = actor.inbox

function actorize_current_task(autoclose::Bool, chsize = 100)
    try
        task_local_storage(:_actor_context)::ActorContext
    catch
        const task = current_task()

        const id = rand(UInt64)
        const inbox = ActorInbox(chsize)
        const self_ref = LocalActor(task)
        const context = ActorContext(id, inbox, self_ref)
        task_local_storage(:_actor_context, context)

        autoclose && @schedule begin
            wait(task)
            close(context)
        end

        context::ActorContext
    end
end

self_context() = actorize_current_task(true)

export self
@inline self() = localactor(self_context())

@inline self_inbox() = actor_inbox(self_context())

self_with_alt_inbox(inbox::ActorInbox) = LocalActorAltInbox(self(), inbox)

export receive

function receive(inbox::ActorInbox)
    const message = take!(inbox)::ActorMsgFrame
    Pair(message.replyto, message.message)
end

receive() = receive(self_inbox())

function Base.close(context::ActorContext)
    close(actor_inbox(context))
end

function LocalActor(body)
    const contextch = Channel{ActorContext}(5)
    try
        const task = @schedule let
            const context = actorize_current_task(false)
            put!(contextch, context)
            try
                body()
            finally
                close(context)
            end
        end
        localactor(take!(contextch))
    finally
        close(contextch)
    end
end

export @actor
macro actor(body)
    quote
        LocalActor(() -> let
                $body
            end
        )
    end
end

export tell
tell(to::Actor, message::Any) = tell(to, message, self())

export ask
function ask(actor::Actor, msg::Any)
    const result = Ref{Any}()
    const queryactor = @actor begin
        tell(actor, msg)
        const x, reply = receive()
        result.x = reply
    end
    wait(queryactor)
    result.x
end

function tell(to::AbstractLocalActor, message::Any, replyto::Actor)
    try
        put!(actor_inbox(to), ActorMsgFrame(replyto, message))
    catch
        # info("Dead letter from $sender to $actor: $message")
    end
    return
end

function ask(actor::LocalActor, msg::Any)
    const tmp_inbox = ActorInbox(1)
    try
        tell(actor, msg, self_with_alt_inbox(tmp_inbox))
        x, reply = receive(tmp_inbox)
        reply
    finally
        close(tmp_inbox)
    end
end

function Base.kill(actor::LocalActor, exc = InterruptException())
    Base.throwto(actor_task(actor), exc)
    return
end

export NullActor
immutable NullActor <: Actor
end

@inline actor_id(::NullActor) = 0x0

Base.show(io::IO, actor::NullActor) = print(io, "null-actor")
Base.wait(actor::NullActor) = begin return end

tell(::NullActor, message::Any, replyto::Actor) = nothing


text:
using Test
using UHDBindings

# --- USRP address
global USRP_ADDRESS = "192.168.10.16"
# This address may vary from time to time, not really sure how to handle dynamic testing

# ----------------------------------------------------
# --- Define test routines
# ----------------------------------------------------
"""
Scan with uhd_find_devices and return the USRP identifier
"""
function check_scan()
    # --- First we use no parameter (broadcast ethernet link)
    str = uhd_find_devices()
    if length(str) == 0
        println("No UHD device found. Be sure that a USRP is connected to your PC, and that the ethernet link is up. We try to use the direct IP address now (which is $USRP_ADDRESS). You can change its value depending on your ethernet setup")
    else
        # We find a device so we update the USRP address based on what we have found
        global USRP_ADDRESS = str[1][findfirst("_addr=",str[1])[end] .+ (1:13)]
    end
    # The direct call with the IP address should give something
    str = uhd_find_devices("addr=$USRP_ADDRESS")
    @test length(str) > 0
end

"""
Test that the device can be open, configured and closed
"""
function check_open()
    # --- Main parameters
	carrierFreq		= 440e6;	# --- The carrier frequency
	samplingRate	= 8e6;         # --- Targeted bandwdith
    gain			= 0;         # --- Rx gain
    # --- Create the E310 device
   global sdr = openUHD(carrierFreq, samplingRate, gain;args="addr=$USRP_ADDRESS");
   # Type is OK
   @test typeof(sdr) == UHDBinding
   # SDR is not released yet
   @test sdr.rx.released == false
   # Configuration : Carrier Freq
   @test sdr.rx.carrierFreq == carrierFreq
   @test sdr.tx.carrierFreq == carrierFreq
   # Configuration : Badnwidth
   @test sdr.rx.samplingRate == samplingRate
   @test sdr.tx.samplingRate == samplingRate
   # --- We close the SDR
   close(sdr)
   @test sdr.rx.released == true
end

"""
Check the carrier frequency update of the USRP device
"""
function check_carrierFreq()
    # --- Main parameters
    carrierFreq		= 440e6;	# --- The carrier frequency
    samplingRate	= 8e6;         # --- Targeted bandwdith
    gain			= 0;         # --- Rx gain
    # --- Create the E310 device
   global sdr = openUHD(carrierFreq, samplingRate, gain;args="addr=$USRP_ADDRESS");
   #  Classic value, should work
   updateCarrierFreq!(sdr,800e6)
   @test sdr.rx.carrierFreq == 800e6
   @test sdr.tx.carrierFreq == 800e6
   # Targeting WiFi, should work
   updateCarrierFreq!(sdr,2400e6)
   @test sdr.rx.carrierFreq == 2400e6
   @test sdr.tx.carrierFreq == 2400e6
   # If we specify a out of range frequency, it should bound to max val
   # TODO Check that is should be max freq range, but don't know the range as different USRP may be used
   # Adding tables with various ranges to check taht this is the expected value ?
   eF= updateCarrierFreq!(sdr,9e9)
   @test sdr.rx.carrierFreq != 9e9
   @test sdr.tx.carrierFreq != 9e9
   @test sdr.rx.carrierFreq == eF
   @test sdr.tx.carrierFreq == eF
   close(sdr);
   @test sdr.rx.released == true
end

"""
Check the sampling frequency update of the USRP device
"""
function check_samplingRate()
    # --- Main parameters
    carrierFreq		= 440e6;	# --- The carrier frequency
    samplingRate	= 8e6;         # --- Targeted bandwdith
    gain			= 0;         # --- Rx gain
    # --- Create the E310 device
   global sdr = openUHD(carrierFreq, samplingRate, gain;args="addr=$USRP_ADDRESS");
   #  Classic value, should work
   updateSamplingRate!(sdr,8e6)
   @test sdr.rx.samplingRate == 8e6
   @test sdr.tx.samplingRate == 8e6
   # Targeting WiFi, should work
   updateSamplingRate!(sdr,15.36e6)
   @test sdr.rx.samplingRate == 15.36e6
   @test sdr.tx.samplingRate == 15.36e6
   # If we specify a out of range frequency, it should bound to max val
   eS = updateSamplingRate!(sdr,100e9)
   # @test sdr.rx.samplingRate != 100e9 ## No Error ?
   # @test sdr.tx.samplingRate != 100e9
   @test sdr.rx.samplingRate == eS
   @test sdr.tx.samplingRate == eS
   close(sdr);
   @test sdr.rx.released == true
end

"""
Check the gain update for the USRP device
"""
function check_gain()
    # --- Main parameters
    carrierFreq		= 440e6;	# --- The carrier frequency
    samplingRate	= 8e6;         # --- Targeted bandwdith
    gain			= 0;         # --- Rx gain
    # --- Create the E310 device
   global sdr = openUHD(carrierFreq, samplingRate, gain;args="addr=$USRP_ADDRESS");
   #  Classic value, should work
   nG  = updateGain!(sdr,20)
   @test sdr.rx.gain == 20
   @test sdr.tx.gain == 20
   @test sdr.rx.gain == nG
   @test sdr.tx.gain == nG
   close(sdr);
   @test sdr.rx.released == true
end

"""
Test that the device  can received data
"""
function check_recv()
    # --- Main parameters
    carrierFreq		= 440e6;	# --- The carrier frequency
    samplingRate	= 8e6;         # --- Targeted bandwdith
    gain			= 0;         # --- Rx gain
    # --- Create the E310 device
   global sdr = openUHD(carrierFreq, samplingRate, gain;args="addr=$USRP_ADDRESS");
    sig = recv(sdr,1024)
    @test length(sig) == 1024
    @test eltype(sig) == Complex{Float32}
    close(sdr)
    @test sdr.rx.released == true
end

"""
Test that the device  can received data  with pre-allocation
"""
function check_recv_preAlloc()
    # --- Main parameters
    carrierFreq		= 440e6;	# --- The carrier frequency
    samplingRate	= 8e6;         # --- Targeted bandwdith
    gain			= 0;         # --- Rx gain
    # --- Create the E310 device
   global sdr = openUHD(carrierFreq, samplingRate, gain;args="addr=$USRP_ADDRESS");
   sig = zeros(ComplexF32,2*1024)
   recv!([sig],sdr)
   @test length(sig) == 1024*2
   @test eltype(sig) == Complex{Float32}
   @test length(unique(sig)) > 1 # To be sure we have populated array with data
   close(sdr)
   @test sdr.rx.released == true
end

"""
Test that the device can received data several time
"""
function check_recv_iterative()
    # --- Main parameters
    carrierFreq		= 440e6;	# --- The carrier frequency
    samplingRate	= 8e6;         # --- Targeted bandwdith
    gain			= 0;         # --- Rx gain
    # --- Create the E310 device
   global sdr = openUHD(carrierFreq, samplingRate, gain;args="addr=$USRP_ADDRESS");
   sig = zeros(ComplexF32,2*1024)
   nbPackets  = 0
   maxPackets = 1_000
   for _ ∈ 1 : 1 : maxPackets
       # --- Get a burst
       recv!([sig],sdr)
       # --- Increment packet index
       nbPackets += 1
   end
   @test nbPackets == maxPackets
   close(sdr)
   @test sdr.rx.released == true
end

"""
Test that the device sucessfully transmit data
"""
function check_send()
	carrierFreq		= 770e6;
	samplingRate	= 4e6;
	gain			= 50.0;
    nbSamples		= 4096*2
	# --- Setting a very first configuration
    global sdr = openUHD(carrierFreq, samplingRate, gain;args="addr=$USRP_ADDRESS");
    print(sdr);
    # --- Create a sine wave
    f_c     = 3940;
    buffer  = 0.5.*[exp.(2im * π * f_c / samplingRate * n)  for n ∈ (0:nbSamples-1)];
    buffer  = convert.(Complex{Cfloat},buffer);

    buffer2  = 0.5.*[exp.(2im * π *  10 * f_c / samplingRate * n)  for n ∈ (0:nbSamples-1)];
    buffer2  = convert.(Complex{Cfloat},buffer2);
    buffer = [buffer;buffer2];

    cntAll  = 0;
    maxPackets = 10_000
    nbPackets  = 0
    for _ ∈ 1 : maxPackets
        send(sdr,buffer,false);
       # --- Increment packet index
       nbPackets += 1
    end
    @test nbPackets == maxPackets
    close(sdr)
    @test sdr.rx.released == true
end

# ----------------------------------------------------
# --- Test calls
# ----------------------------------------------------
@testset "Scanning and opening" begin
    check_scan()
    check_open()
end

@testset "Radio configuration" begin
    check_carrierFreq()
    check_samplingRate()
    check_gain()
end

@testset "Checking data retrieval"  begin
    check_recv()
    check_recv_preAlloc()
    check_recv_iterative()
end

@testset "Checking data transmission"  begin
    check_send()
end

max_stars_repo_name: RGerzaguet/UHDBindings

max_issues_repo_name: RGerzaguet/UHDBindings


text:
using DataFrames, Datetime, TimeSeries

module PortfolioModels

using DataFrames, Datetime, TimeSeries

type Portfolio
  mean
  var
  std
  covmat
  cormat
  weights
end

export Portfolio,
       returns,
       read_monthlies,
       modelparams,
       globalmin,
       efficient,
       tangency,
       frontier,
       build_globalmin,
       build_efficient,
       build_tangency,

## testsuite macro
       @portfoliomodels

include("preprocess.jl")
include("solve.jl")
include("build.jl")
include("../test/testmacro.jl")

end


text:
function setup_repl(repl)
    ## RCall.RPrompt.repl_init requires repl.interface has been setup.
    if !isdefined(repl,:interface)
        repl.interface = REPL.setup_interface(repl)
    end

    !RCall.RPrompt.repl_inited(repl) && RCall.RPrompt.repl_init(repl)

    ## We also need to do rgui_start
    RCall.rgui_start()
end

# We will use this to register hook for Julia REPL: Base.atreplinit(setup_repl)

max_stars_repo_name: KevCaz/JuliaCall

max_issues_repo_name: tinagui/JuliaCall

max_forks_repo_name: tinagui/JuliaCall


text:
const Point3DS{T} = Point3D{SVector{3, T}}

zero_nans(x) = ifelse(isnan(x), zero(x), x)

function canonicalize(body::RigidBody, point::Point3D)
    if point.frame != default_frame(body)
        point = frame_definition(body, point.frame) * point
    end
    point
end

function canonicalize!(d::Dict{RigidBody{T}, Vector{Tuple{Point3DS{T}, Point3DS{T}}}}) where {T}
    for (body, bounds) in d
        for i in eachindex(bounds)
            bounds[i] = canonicalize.(Ref(body), bounds[i])
        end
    end
end

const TreeJointSegmentedVector{T} = SegmentedVector{
        RigidBodyDynamics.JointID, T, Base.OneTo{RigidBodyDynamics.JointID}, Vector{T}}


text:
module ComsolUtils

#
include("./fcn/typedefs.jl")

# reading comsol exported files
using DelimitedFiles: readdlm
include("./fcn/curead.jl")
export curead

# parsing read comsol exported files
include("./fcn/cuparse.jl")
export cuparse

# add interpolation functions
include("./fcn/interpolation2Dsym.jl")

#
using CairoMakie, ColorSchemes
include("./fcn/cuplots.jl")
export plotdset2Dsym

end


text:
# Подключаем основной файл робота. Без ниго ничего не работает.
include("core_lib/HorizonSideRobots.jl")
using .HorizonSideRobots

# Файл, со вспомогательными функциями. Нужен, чтобы не загружать "лишним" основной файл
include("roblib.jl")

"""
solve(r::Robot)

Функция, при выполнении которой создаётся специальный робот, который обходит лабиринт и пишет, внутри лабиринта он или снаружи.
"""
function solve!(r::Robot)
    # Инициализируем робота
    lab_rob = Robot_for_labyrinth(r)
    lab_rob.start_point_border_position = get_border_position(lab_rob)
    lab_rob.move_side = get_move_side(lab_rob)
    lab_rob.previos_state = lab_rob.start_point_border_position
    lab_rob.count_of_external_corners = 0
    lab_rob.count_of_internal_corners = 0

    # Ставим маркер в начальной клетке
    putmarker!(lab_rob.robot)

    start = true
    # Обход лабиринта и подсёт углов
    one_step_move_around_labyrinth!(lab_rob)
    update_corners_count!(lab_rob)

    while !(ismarker(lab_rob.robot) && get_border_position(lab_rob) == lab_rob.start_point_border_position && start == false)
        one_step_move_around_labyrinth!(lab_rob)
        if !ismarker(lab_rob.robot)
            start = false;
        end
        update_corners_count!(lab_rob)
    end

    # Вывод ответа
    show_ans(lab_rob)

end


text:
quote
    # (a -> b, a) -> b
    struct Apply{A <: Arrow, B<:Arrow, T <: Arrow{A, B}} <: Arrow{Tuple{T, A}, B}
    end
    struct Curry{A <: Arrow, B<: Arrow, C <: Arrow}
            <: Arrow{A, Arrow{B, C}}
        thunk::Arrow{Tuple{A, B}, C}
    end
    struct Uncurry{A <: Arrow, B<: Arrow, C <: Arrow} <: Arrow{Tuple{A, B}, C}
        thunk::Arrow{A, Arrow{B, C}}
    end
end


text:
module Reparametrization

using LinearAlgebra
using Random

using DiffEqBase: ODEProblem, isinplace
using Parameters: @with_kw, @unpack
using Setfield: @lens, compose
using StaticArrays: SVector, SMatrix
import Setfield

using ...ArrayUtils: _zeros, badeltype
# using ...BifurcationsBase: TimeKind

struct Reparametrizer{DIM_ALL, DIM_ORIG, DIM_EXTRA, TF, TP, TM, TS, TE}
    orig_f::TF
    orig_p::TP
    M::TM
    shift::TS
    extra::TE
    invM::TM
end

function Reparametrizer(orig_f::TF, orig_p::TP, M, shift::TS, extra::TE,
                        invM = inv(M),
                        ) where {TF, TP, TS, TE}
    @assert M * invM ≈ typeof(M)(I)
    orig_TM = typeof(M)
    M, invM = promote(M, invM)
    TM = typeof(M) :: Type{<: AbstractMatrix}
    if orig_TM <: SMatrix
        @assert TM <: SMatrix
    end
    DIM_ALL = size(M, 1)
    DIM_EXTRA = length(extra)
    DIM_ORIG = DIM_ALL - DIM_EXTRA
    return Reparametrizer{DIM_ALL, DIM_ORIG, DIM_EXTRA, TF, TP, TM, TS, TE}(
        orig_f, orig_p, M, shift, extra, invM,
    )
end

function Reparametrizer{DIM_ALL, DIM_ORIG, DIM_EXTRA}(
        orig_f::TF, orig_p::TP, M::TM, shift::TS, extra::TE, invM::TM,
        ) where {DIM_ALL, DIM_ORIG, DIM_EXTRA, TF, TP, TM, TS, TE}
    return Reparametrizer{DIM_ALL, DIM_ORIG, DIM_EXTRA, TF, TP, TM, TS, TE}(
        orig_f, orig_p, M, shift, extra, invM,
    )
end

Setfield.constructor_of(::Type{<: Reparametrizer{DIM_ALL, DIM_ORIG, DIM_EXTRA}}) where {DIM_ALL, DIM_ORIG, DIM_EXTRA} =
    Reparametrizer{DIM_ALL, DIM_ORIG, DIM_EXTRA}

@generated function split_x(x::SVector{S, T},
                            rp::Reparametrizer{DIM_ALL, DIM_ORIG, DIM_EXTRA},
                            ) where {S, T, DIM_ALL, DIM_ORIG, DIM_EXTRA}
    @assert S == DIM_ALL
    idx_orig = 1:DIM_ORIG
    idx_extra = DIM_ORIG .+ (1:DIM_EXTRA)
    quote
        (SVector{$DIM_ORIG, $T}($((:(x[$i]) for i in idx_orig)...)),
         SVector{$DIM_EXTRA, $T}($((:(x[$i]) for i in idx_extra)...)))
    end
end

forward(x, rp::Reparametrizer) = tanh.(rp.invM * (@. x - rp.shift))
backward(y, rp::Reparametrizer) = (rp.M * atanh.(y) .+ rp.shift)

function f(y::TY, rp::Reparametrizer, t) where {TY <: SVector}
    # Original ODE:          dx/dt = f₀(x)
    # "Re-parametrization":     x = g(y) = backward(y)
    # Re-parametrized ODE:  dy/dt = dy/dx dx/dt
    #                             = (g'(y))⁻¹ f₀(g(y))
    x = backward(y, rp)
    # @assert ! badeltype(x)
    x0, x1 = split_x(x, rp)
    dx0 = rp.orig_f(x0, rp.orig_p, t)
    dx1 = rp.extra .* x1
    # @assert ! badeltype(dx1)
    dx = vcat(dx0, dx1) :: SVector

    # (g'(y))⁻¹ = (g⁻¹(x))' = (tanh(M⁻¹ (x - shift)))' = tanh' M⁻¹
    dy = (1 .- y .* y) .* (rp.invM * dx)
    dy = SVector(dy...)  # TODO: don't
    @assert ! badeltype(dy)
    return dy
    # return dy :: TY
end
# Cannot use `dy :: TY` because of the type instability in
# first_lyapunov_coefficient.

f(u::A, p::Reparametrizer{DIM_ALL}, t) where {A <: AbstractVector, DIM_ALL} =
    A(f(SVector{DIM_ALL}(u), p, t))

function f(du, u, rp::Reparametrizer{DIM_ALL}, t) where {DIM_ALL}
    du .= f(SVector{DIM_ALL}(u), rp, t)
    nothing
end

function reparametrize(
        ode::ODEProblem, M::AbstractMatrix;
        shift = 0.0,
        extra::Union{AbstractVector, Real} = SVector{0, Float64}(),
        )
    rp = Reparametrizer(ode.f, ode.p, M, shift, extra)
    x1 = _zeros(ode.u0, length(extra))
    x = vcat(ode.u0, x1)
    y = forward(x, rp)
    u0 = typeof(x)(y)  # since `forward` returns `SVector`
    return ODEProblem{isinplace(ode)}(f, u0, ode.tspan, rp)
end

function reparametrize(
        ode::ODEProblem, rng::AbstractRNG;
        orthogonal = true,
        extra::Union{AbstractVector, Real} = SVector{0, Float64}(),
        kwargs...)
    tries = 10
    dim = length(ode.u0) + length(extra)
    for _ in 1:tries
        M = SMatrix{dim, dim}(randn(rng, dim, dim))
        if orthogonal
            M, _ = qr(M)
        end
        try
            invM = inv(M)
        catch
            continue
        end
        return reparametrize(ode, M; extra=extra, kwargs...)
    end
    error("Cannot obtain an invertible matrix after $tries tries.")
end

reparametrize(ode::ODEProblem; seed::Integer = 0, kwargs...) =
    reparametrize(ode, MersenneTwister(seed); kwargs...)

using ...Bifurcations: FixedPointBifurcationProblem, BifurcationProblem,
    unbounded_phase_space

const orig_p = @lens _.orig_p

function reparametrize(prob::FixedPointBifurcationProblem,
                       args...; kwargs...)

    let u0 = prob.p.de_prob.u0
        # TODO: support transforming prob.u0 and prob.phase_space:
        @assert prob.u0 ≈ u0
        @assert prob.phase_space == unbounded_phase_space(u0)
    end
    @assert prob.homotopy_jacobian == nothing

    ode = reparametrize(prob.p.de_prob, args...; kwargs...)
    return BifurcationProblem(
        ode,
        compose(orig_p, prob.p.param_axis),
        prob.t_domain,
    )
end

end  # module


text:
@testset "learner" begin
    nfs = [8,16,32,32]
    model = get_model(nfs)
    opti = Descent(1e-2)
    loss(pred,y) = logitcrossentropy(pred,y)
    learn = Learner(model,opti,loss,data)

    @test learn.callbaks[1].order == 0
end


text:
# ---
# name: acc-results.jl
# author: Jhelum Chakravorty, Jayakumar Subramanian, Aditya Mahajan
# date: 1 Feb, 2017
# license: BSD 3 clause
# ---
#
# Results presented in ACC'17
# ===========================
#
# We want to run the simulation code of `simulations.jl` on multiple cores.
# So, as the first step, we make sure we are running multi-core version of
# Julia.

addprocs(Sys.CPU_CORES)
include("simulations.jl")

# We use Plots package with GR backend to plot the results.

using Plots
gr()

# A helper function to display labels

function labeltext(s, param)
    text(string('$', s, "=", param, '$'), 8, :bottom, :right)
end

label_costly(param)      = labeltext("\\lambda", param)
label_constrained(param) = labeltext("\\alpha", param)

# The function below computes the results for a particular stochastic
# approximation algorithm for a range of `paramValues`, `discountValues`,
# and `packetDropValues` . If the optional argument `saveRawData` is `true`,
# the traces of the multiple runs are saved in a `.jld` file; if `savePlot` is
# `true`, then the results every combination of `discount` and `packetDrop`
# are saved to a pdf file in the `plots/` directory. The optional values
# `iterations` and `numRuns` determine the number of iterations and runs. The
# optional argument `labeltext` formats the label of each plot and `ylim`
# adjusts the y limits.

function computeTimeSeries(sa, paramValues, discountValues, dropProbValues;
                           iterations=1000, numRuns=100,
			               saveSummaryData=true, saveRawData = false, savePlot=false,
                           labeltext=label_costly, ylim=:auto)

    for discount in discountValues, dropProb in dropProbValues
        result = Array(DataFrame, length(paramValues))

        for i in 1:length(paramValues)
            param = paramValues[i]
            result[i] = generateOutput(sa, param, discount, dropProb, iterations;
                                       numRuns=numRuns,
                                       saveSummaryData=saveSummaryData, saveRawData=saveRawData)
        end

        if savePlot
            plt = plot(xlabel="Iterations", ylabel="Threshold",
                       ylim=ylim)

            for i in 1:length(paramValues)
                param = paramValues[i]
                df = result[i]
                plot!(plt, df[:upper], linecolor=:lightblue, label="")
                plot!(plt, df[:lower], linecolor=:lightblue, label="",
                     fillrange=df[:upper], fillcolor=:lightgray, fillalpha=0.8)
                plot!(plt, df[:mean],  linecolor=:black, linewidth=2, label="")

                x=div(3iterations,4)
                y=df[:upper][x]
                label=labeltext(param)

                plot!(plt,annotations=(x,y,label))
            end

            filename = string("plots/", sa,             "__parameter_", paramValues,
                              "__discount_", discount,  "__dropProb_" , dropProb)

            savefig("$filename.pdf")
        end
    end
end

# First we compute the results for packet drop probability $p_d = 0$ (and
# separately verify these results from the values obtained from numerically
# solving Fredholm integral equations.

computeTimeSeries(sa_costly, linspace(100,700,7), [0.9,1.0], 0.0;
                  iterations=10_000, numRuns=100, saveRawData = true, savePlot=false)
computeTimeSeries(sa_constrained, linspace(0.1,0.8,8), [0.9,1.0], 0.0;
                  iterations=10_000, numRuns=100, saveRawData = true, savePlot=false)

# Next, we compute the results for packet drop probability $p_d = 0.3$.

computeTimeSeries(sa_costly, linspace(100,500,3), [0.9,1.0], 0.3;
                  iterations=10_000, numRuns=100, saveRawData = true,
                  savePlot=true, labeltext=label_costly, ylim=(0,11))
computeTimeSeries(sa_constrained, linspace(0.1,0.5,3), [0.9,1.0], 0.3;
                  iterations=10_000, numRuns=100, saveRawData = true,
                  savePlot=true, labeltext=label_constrained, ylim=(0,3.5))

max_stars_repo_name: adityam/remote-estimation-packet-drops

max_issues_repo_name: adityam/remote-estimation-packet-drops

max_forks_repo_name: adityam/remote-estimation-packet-drops


text:
# adapted from https://github.com/JuliaMath/KahanSummation.jl
function cumsum_kbn!(v::AbstractVector{T}) where T<:AbstractFloat
    s = v[1]
    c = zero(T)
    @inbounds for i = 2:length(v)
        vi = v[i]
        t = s + vi
        if abs(s) >= abs(vi)
            c += ((s-t) + vi)
        else
            c += ((vi-t) + s)
        end
        s = t
        v[i] = s+c
    end
    return nothing
end

function diff_x!(x::Array{T,1}, gaps::Array{Int64,1}, fs::T) where T<:AbstractFloat
  L = length(x)
  j = 1
  @inbounds while j < length(gaps)
    j += 1
    si = gaps[j-1]+1
    ei = gaps[j]-1
    if j == length(gaps)
      ei += 1
    end
    for i = ei:-1:si
      x[i] -= x[i-1]
    end
  end
  rmul!(x, fs)
  return nothing
end

function int_x!(x::Array{T,1}, gaps::Array{Int64,1}, δ::T) where T<:AbstractFloat
  L = length(x)
  j = 1
  @inbounds while j < length(gaps)
    j += 1
    si = gaps[j-1]
    ei = gaps[j]-1
    if j == length(gaps)
      ei += 1
    end

    xv = view(x, si:ei)
    cumsum_kbn!(xv)
  end
  rmul!(x, δ)
  return nothing
end


text:
using ArgParse
using DrWatson
@quickactivate
using CSV
using BSON
using Random
using FileIO
using DataFrames
using PrettyTables
using PrettyTables.Crayons

using GenerativeAD.Evaluation: _prefix_symbol, PAT_METRICS, aggregate_stats_mean_max
using GenerativeAD.Evaluation: rank_table, print_rank_table
using GenerativeAD.Datasets: load_data

s = ArgParseSettings()
@add_arg_table! s begin
	"filename"
		arg_type = String
		default = "evaluation/images_eval.bson"
		help = "Location of cached DataFrame."
	"-c", "--criterion-metric"
		arg_type = String
		default = "val_auc"
		help = "Criterion to sort the models."
	"-r", "--rank-metric"
		arg_type = String
		default = "tst_auc"
		help = "Metric to rank models."
	"-o", "--output-prefix"
		arg_type = String
		default = "evaluation/images_eval"
		help = "Output prefix for storing results."
    "-f", "--force"
    	action = :store_true
		help = "Overwrite all generated files."
end

# args = Dict()
# args["filename"] = "evaluation/tabular_sptn_5seeds_quantiles_eval.bson"
# args["criterion-metric"] = "val_tpr_5"
# args["rank-metric"] = "tst_tpr_5"
# args["output-prefix"] = "evaluation/zkouska_graf_rel_params2.bson"
# args["max-params"] = 100
#
# f = datadir(args["filename"])
# df = load(f)[:df]
# cdf = deepcopy(df)
#
# tdf = deepcopy(cdf)
# for i in 1:1
# 	tdf = deepcopy(cdf)
# 	for d in unique(cdf[:dataset])
# 		mx = maximum(cdf[cdf[:dataset] .== d, :npars])
# 		filter!(row -> (row[:dataset] .!= d) .| ((row[:dataset] .== d) .& (row[:npars] .<= i/100 .* mx)), tdf)
# 	end
# 	println(i)
# end
#
# data_dim(s) = size(load_data(s)[1][1], 1)
# datadim = Dict()
# for d in unique(df[:dataset])
#     println(d)
#     datadim[d] = size(load_data(d)[1][1], 1)
# end
# df[data_dim.(df[:dataset]) .== 10]
#
# cdf.datadim = 1
# cdf.relnpars = 1.0
# for f in eachrow(cdf)
#     @show f.datadim = datadim[f[:dataset]]
#     @show f.relnpars = f[:npars]/f[:datadim]
# end
#
# df_agg = aggregate_stats_mean_max(cdf, "val_auc")
# rt = rank_table(df_agg, "tst_auc")
#
# npar_df = DataFrame(modelname = String[],relnpars = Float64[], rank = Float64[])
# for i in minimum(cdf[:relnpars]):10:maximum(cdf[:relnpars])
#     df_agg = aggregate_stats_mean_max(cdf[cdf[:relnpars] .<= i, :], "val_auc")
#     rt = rank_table(df_agg, "tst_auc")
#     for n in names(rt)[2:end]
#         push!(npar_df, (n, i, rt[end, n]))
#     end
#     println(i)
# end
# npar_df
# using Plots
# sptn_df = npar_df[npar_df[:modelname] .== "sptn", :]
# sptn_5_df = npar_df[npar_df[:modelname] .== "sptn_005_010", :]
# sptn_25_df = npar_df[npar_df[:modelname] .== "sptn_0025_0075", :]
# plot(sptn_df[:relnpars], sptn_df[:rank])
# plot!(sptn_5_df[:relnpars], sptn_5_df[:rank])
# plot!(sptn_25_df[:relnpars], sptn_25_df[:rank])

function main(args)
	f = datadir(args["filename"])
	df = load(f)[:df]
    cdf = deepcopy(df)

    #data dim dict
    data_dim(s) = size(load_data(s)[1][1], 1)
    datadim = Dict()
    for d in unique(df[:dataset])
        datadim[d] = size(load_data(d)[1][1], 1)
    end

    # cdf.datadim = 1
    # cdf.relnpars = 1.0
    # for f in eachrow(cdf)
    #     f.datadim = datadim[f[:dataset]]
    #     #relnpars = numb. of params of model / data dim.
    #     f.relnpars = f[:npars]/f[:datadim]
    # end
	@info "Loaded $(nrow(df)) rows from $f"

    # Dataframe for final ranks for each relnpars
    npar_df = DataFrame(modelname = String[],relnpars = Float64[], rank = Float64[])

    # m = args["max-params"]
    # if m == -1
    #     m = maximum(cdf[:relnpars])
    # end
	#
    # for i in minimum(cdf[:relnpars]):1:m
    #     # evaluation and rank table
    #     df_agg = aggregate_stats_mean_max(cdf[cdf[:relnpars] .<= i, :], Symbol(args["criterion-metric"]))
    #     rt = rank_table(df_agg, args["rank-metric"])
    #     for n in names(rt)[2:end]
    #         # row for each model at relnpars
    #         push!(npar_df, (n, i, rt[end, n]))
    #     end
	# 	@info "Number of relative parameters $i"
    # end

	for i in 1:100
		tdf = deepcopy(cdf)
		for d in unique(cdf[:dataset])
			mx = maximum(cdf[cdf[:dataset] .== d, :npars])
			filter!(row -> (row[:dataset] .!= d) .| ((row[:dataset] .== d) .& (row[:npars] .<= i/100 .* mx)), tdf)
		end
		# display(size(tdf))
		df_agg = aggregate_stats_mean_max(tdf, Symbol(args["criterion-metric"]))
		rt = rank_table(df_agg, args["rank-metric"])

        for n in names(rt)[2:end]
            # row for each model at relnpars
            push!(npar_df, (n, i, rt[end, n]))
        end

		@info "Percent of relative parameters $i %"
	end

	@info "Best models chosen by $(args["criterion-metric"])"
	@info "Ranking by $(args["rank-metric"])"

    ff = datadir(args["output-prefix"])
	if (isfile(ff) && args["force"]) || ~isfile(ff)
		@info "Saving $(nrow(npar_df)) rows to $ff."
		wsave(ff, Dict(:df => npar_df))
	end
	@info "---------------- DONE -----------------"
end

main(parse_args(ARGS, s))


text:
@everywhere include("utils.jl")

"""
    generate_ref_dataset(temp_series, H₀)

Generate reference dataset based on the iceflow PDE
"""
function generate_ref_dataset(temp_series, H₀)
    # Compute reference dataset in parallel
    H = deepcopy(H₀)

    # Initialize all matrices for the solver
    S, dSdx, dSdy = zeros(Float64,nx,ny),zeros(Float64,nx-1,ny),zeros(Float64,nx,ny-1)
    dSdx_edges, dSdy_edges, ∇S = zeros(Float64,nx-1,ny-2),zeros(Float64,nx-2,ny-1),zeros(Float64,nx-1,ny-1)
    D, dH, Fx, Fy = zeros(Float64,nx-1,ny-1),zeros(Float64,nx-2,ny-2),zeros(Float64,nx-1,ny-2),zeros(Float64,nx-2,ny-1)
    V, Vx, Vy = zeros(Float64,nx-1,ny-1),zeros(Float64,nx-1,ny-1),zeros(Float64,nx-1,ny-1)
    A = 2e-16
    α = 0                       # Weertman-type basal sliding (Weertman, 1964, 1972). 1 -> sliding / 0 -> no sliding
    C = 15e-14                  # Sliding factor, between (0 - 25) [m⁸ N⁻³ a⁻¹]

    # Gather simulation parameters
    current_year = 0
    context = ArrayPartition([A], B, S, dSdx, dSdy, D, copy(temp_series[5]), dSdx_edges, dSdy_edges, ∇S, Fx, Fy, Vx, Vy, V, C, α, [current_year])

    # Perform reference simulation with forward model
    println("Running forward PDE ice flow model...\n")
    # Train batches in parallel
    # Fetch reference simulated ice thickness
    H_refs  = @showprogress pmap(temps -> prob_iceflow_PDE(H, temps, context), temp_series)

    # Compute average ice surface velocities for the simulated period
    V̄x_refs, V̄y_refs = [],[]
    for (H_ref, temps) in zip(H_refs, temp_series)
        V̄x_ref, V̄y_ref = avg_surface_V(H_ref, B, mean(temps), "PDE") # Average velocity with average temperature
        push!(V̄x_refs, V̄x_ref)
        push!(V̄y_refs, V̄y_ref)
    end

    return H_refs, V̄x_refs, V̄y_refs
end

@everywhere begin
"""
    prob_iceflow_PDE(H, temps, context)

Solve the Shallow Ice Approximation iceflow PDE for a given temperature series
"""
function prob_iceflow_PDE(H, temps, context)

    println("Processing temp series ≈ ", mean(temps))
    context.x[7] .= temps # We set the temp_series for the ith trajectory

    iceflow_prob = ODEProblem(iceflow!,H,(0.0,t₁),context)
    iceflow_sol = solve(iceflow_prob, solver,
                    reltol=1e-6, save_everystep=false,
                    progress=true, progress_steps = 10)

    return iceflow_sol.u[end]
end
end # @everywhere

"""
    train_iceflow_UDE(H₀, UA, θ, train_settings, PDE_refs, temp_series)

Train the Shallow Ice Approximation iceflow UDE
"""
function train_iceflow_UDE(H₀, UA, θ, train_settings, PDE_refs, temp_series)
    H = deepcopy(H₀)
    optimizer = train_settings[1]
    epochs = train_settings[2]
    # Tuple with all the temp series and H_refs
    context = (B, H)
    loss(θ) = loss_iceflow(θ, context, UA, PDE_refs, temp_series) # closure

    println("Training iceflow UDE...")
    # println("Using solver: ", solver)
    iceflow_trained = DiffEqFlux.sciml_train(loss, θ, optimizer, cb=callback, maxiters = 1)

    return iceflow_trained
end

@everywhere begin

callback = function (θ,l) # callback function to observe training
    println("Epoch #$current_epoch - Loss $loss_type: ", l)

    # pred_A = predict_A̅(UA, θ, collect(-20.0:0.0)')
    # pred_A = [pred_A...] # flatten
    # true_A = A_fake(-20.0:0.0, noise)

    # Plots.scatter(-20.0:0.0, true_A, label="True A")
    # plot_epoch = Plots.plot!(-20.0:0.0, pred_A, label="Predicted A",
    #                     xlabel="Long-term air temperature (°C)",
    #                     ylabel="A", ylims=(2e-17,8e-16),
    #                     legend=:topleft)
    # Plots.savefig(plot_epoch,joinpath(root_plots,"training","epoch$current_epoch.png"))
    global current_epoch += 1

    false
end

"""
    loss_iceflow(θ, context, UA, PDE_refs::Dict{String, Any}, temp_series)

Loss function based on glacier ice velocities and/or ice thickness
"""
function loss_iceflow(θ, context, UA, PDE_refs::Dict{String, Any}, temp_series)
    if multibatch
        H_preds = predict_iceflow(θ, UA, context, temp_series)
    else
        H_preds = predict_iceflow_onebatch(θ, UA, context, temp_series)
    end

    # Compute loss function for the full batch
    l_Vx, l_Vy, l_H = 0.0, 0.0, 0.0

    if multibatch
        for i in 1:length(H_preds)

            # Get ice thickness from the reference dataset
            H_ref = PDE_refs["H_refs"][i]
            # Get ice thickness from the UDE predictions
            H = H_preds[i]
            # Get ice velocities for the reference dataset
            Vx_ref = PDE_refs["V̄x_refs"][i]
            Vy_ref = PDE_refs["V̄y_refs"][i]
            # Get ice velocities from the UDE predictions
            V̄x_pred, V̄y_pred = avg_surface_V(H_preds[i], B, mean(temp_series[i]), "UDE") # Average velocity with average temperature

            if random_sampling_loss
                # sample random indices for which V_ref is non-zero
                n_sample = 10
                n_counts = 0
                while n_counts < n_sample
                    i, j = rand(1:nx), rand(1:ny)
                    if Vx_ref[i,j] != 0.0
                        #println("New non-zero count for loss function: ", i, j)
                        n_counts += 1
                        l_Vx += (V̄x_pred[i,j] - Vx_ref[i,j] )^2
                        l_Vy += (V̄y_pred[i,j] - Vy_ref[i,j] )^2
                        # what about trying a percentual error?
                        # l_H += ( (H[i,j] - H_ref[i,j]) / H_ref[i,j] )^2
                    end
                end
                l_Vx = l_Vx / n_sample
                l_Vy = l_Vy / n_sample
            else
                # Classic loss function with the full matrix
                l_H += Flux.Losses.mse(H[H_ref .!= 0.0], H_ref[H_ref.!= 0.0]; agg=mean)
                l_Vx += Flux.Losses.mse(V̄x_pred[Vx_ref .!= 0.0], Vx_ref[Vx_ref.!= 0.0]; agg=mean)
                l_Vy += Flux.Losses.mse(V̄y_pred[Vy_ref .!= 0.0], Vy_ref[Vy_ref.!= 0.0]; agg=mean)
            end
        end
    else
        i = 1
        # Get ice thickness from the reference dataset
        H_ref = PDE_refs["H_refs"][i]
        # Get ice thickness from the UDE predictions
        H = H_preds
        # Get ice velocities for the reference dataset
        Vx_ref = PDE_refs["V̄x_refs"][i]
        Vy_ref = PDE_refs["V̄y_refs"][i]
        # Get ice velocities from the UDE predictions
        V̄x_pred, V̄y_pred = avg_surface_V(H_preds, B, mean(temp_series[i]), "UDE") # Average velocity with average temperature

        if random_sampling_loss
            # sample random indices for which V_ref is non-zero
            n_sample, n_counts = 50, 0
            nxy = length(H_ref[H_ref .!= 0.0])
            # Zygote.ignore() do
            #     @infiltrate
            # end
            while n_counts < n_sample
                j = rand(1:nxy-2)
                H_ref_f = H_ref[H_ref .!= 0.0]
                Vx_ref_f = Vx_ref[Vx_ref .!= 0.0]
                Vy_ref_f = Vy_ref[Vy_ref .!= 0.0]
                H_f = H[H_ref .!= 0.0]
                V̄x_pred_f = V̄x_pred[Vx_ref .!= 0.0]
                V̄y_pred_f = V̄y_pred[Vy_ref .!= 0.0]

                if norm_loss
                    normH = H_ref[H_ref .!= 0.0] .+ ϵ
                    normVx = Vx_ref[Vx_ref .!= 0.0] .+ ϵ
                    normVy = Vy_ref[Vy_ref .!= 0.0] .+ ϵ
                    l_H += Flux.Losses.mse(H_f[j] ./normH, H_ref_f[j]  ./normH; agg=mean)
                    l_Vx += Flux.Losses.mse(V̄x_pred_f[j] ./normVx, Vx_ref_f[j]  ./normVx; agg=mean)
                    l_Vy += Flux.Losses.mse(V̄y_pred_f[j] ./normVy, Vy_ref_f[j]  ./normVy; agg=mean)
                else
                    l_H += Flux.Losses.mse(H_f[j], H_ref_f[j]; agg=mean)
                    l_Vx += Flux.Losses.mse(V̄x_pred_f[j], Vx_ref_f[j]; agg=mean)
                    l_Vy += Flux.Losses.mse(V̄y_pred_f[j], Vy_ref_f[j]; agg=mean)
                end

                n_counts += 1
            end
            l_H = l_H / n_sample
            l_Vx = l_Vx / n_sample
            l_Vy = l_Vy / n_sample
        else
            # Classic loss function with the full matrix
            if norm_loss
                normH = H_ref[H_ref .!= 0.0] .+ ϵ
                normVx = Vx_ref[Vx_ref .!= 0.0] .+ ϵ
                normVy = Vy_ref[Vy_ref .!= 0.0] .+ ϵ
                l_H += Flux.Losses.mse(H[H_ref .!= 0.0] ./normH, H_ref[H_ref.!= 0.0] ./normH; agg=mean)
                l_Vx += Flux.Losses.mse(V̄x_pred[Vx_ref .!= 0.0] ./normVx, Vx_ref[Vx_ref.!= 0.0] ./normVx; agg=mean)
                l_Vy += Flux.Losses.mse(V̄y_pred[Vy_ref .!= 0.0] ./normVy, Vy_ref[Vy_ref.!= 0.0] ./normVy; agg=mean)
            else
                l_H += Flux.Losses.mse(H[H_ref .!= 0.0], H_ref[H_ref.!= 0.0]; agg=mean)
                l_Vx += Flux.Losses.mse(V̄x_pred[Vx_ref .!= 0.0], Vx_ref[Vx_ref.!= 0.0]; agg=mean)
                l_Vy += Flux.Losses.mse(V̄y_pred[Vy_ref .!= 0.0], Vy_ref[Vy_ref.!= 0.0]; agg=mean)
            end
        end
    end

    @assert (loss_type == "H" || loss_type == "V" || loss_type == "HV") "Invalid loss_type. Needs to be 'H', 'V' or 'HV'"
    if loss_type == "H"
        l_avg = l_H/length(H_preds)
    elseif loss_type == "V"
        l_avg = (l_Vx/length(PDE_refs["V̄x_refs"]) + l_Vy/length(PDE_refs["V̄y_refs"]))/2
    elseif loss_type == "HV"
        l_avg = (l_Vx/length(PDE_refs["V̄x_refs"]) + l_Vy/length(PDE_refs["V̄y_refs"]) + l_H/length(H_preds))/3
    end
    return l_avg
end

"""
    predict_iceflow(θ, UA, context, temp_series)

Makes a prediction of glacier evolution with the UDE for a given temperature series
"""
function predict_iceflow(θ, UA, context, temp_series)
    # (B, H)
    H = context[2]

    # Train UDE in parallel
    H_preds = pmap(temps -> prob_iceflow_UDE(θ, H, temps, context, UA), temp_series)

    return H_preds
end

function predict_iceflow_onebatch(θ, UA, context, temp_series)
    # (B, H)
    H = context[2]

    # Train UDE in parallel
    H_preds = prob_iceflow_UDE(θ, H, temp_series[1], context, UA)

    return H_preds
end

function prob_iceflow_UDE(θ, H, temps, context, UA)

    # println("Processing temp series ≈ ", mean(temps))
    iceflow_UDE_batch(H, θ, t) = iceflow_NN(H, θ, t, context, temps, UA) # closure
    iceflow_prob = ODEProblem(iceflow_UDE_batch,H,(0.0,t₁),θ)
    iceflow_sol = solve(iceflow_prob, solver, u0=H, p=θ,
                    reltol=1e-6, save_everystep=false,
                    progress=true, progress_steps = 10)

    return iceflow_sol.u[end]
end

"""
    iceflow!(dH, H, context,t)

Runs a single time step of the iceflow PDE model in-place
"""
function iceflow!(dH, H, context,t)
    # Unpack parameters
    #A, B, S, dSdx, dSdy, D, temps, dSdx_edges, dSdy_edges, ∇S, Fx, Fy, Vx, Vy, V, C, α, current_year
    current_year = Ref(context.x[18])
    A = Ref(context.x[1])

    # Get current year for MB and ELA
    year = floor(Int, t) + 1
    if year != current_year[] && year <= t₁
        temp = Ref{Float64}(context.x[7][year])
        A[] .= A_fake(temp[], noise)
        current_year[] .= year
    end

    # Compute the Shallow Ice Approximation in a staggered grid
    SIA!(dH, H, context)
end

"""
    iceflow_NN(H, θ, t, context, temps, UA)

Runs a single time step of the iceflow UDE model
"""
function iceflow_NN(H, θ, t, context, temps, UA)

    year = floor(Int, t) + 1
    if year <= t₁
        temp = temps[year]
    else
        temp = temps[year-1]
    end

    A = predict_A̅(UA, θ, [temp]) # FastChain prediction requires explicit parameters

    # Compute the Shallow Ice Approximation in a staggered grid
    return SIA(H, A, context)
end

"""
    SIA!(dH, H, context)

Compute an in-place step of the Shallow Ice Approximation PDE in a forward model
"""
function SIA!(dH, H, context)
    # Retrieve parameters
    #A, B, S, dSdx, dSdy, D, norm_temps, dSdx_edges, dSdy_edges, ∇S, Fx, Fy, Vx, Vy, V, C, α, current_year, H_ref, H, UA, θ
    A = context.x[1]
    B = context.x[2]
    S = context.x[3]
    dSdx = context.x[4]
    dSdy = context.x[5]
    D = context.x[6]
    dSdx_edges = context.x[8]
    dSdy_edges = context.x[9]
    ∇S = context.x[10]
    Fx = context.x[11]
    Fy = context.x[12]

    # Update glacier surface altimetry
    S .= B .+ H

    # All grid variables computed in a staggered grid
    # Compute surface gradients on edges
    dSdx .= diff_x(S) / Δx
    dSdy .= diff_y(S) / Δy
    ∇S .= (avg_y(dSdx).^2 .+ avg_x(dSdy).^2).^((n - 1)/2)

    Γ = 2 * A * (ρ * g)^n / (n+2) # 1 / m^3 s
    D .= Γ .* avg(H).^(n + 2) .* ∇S

    # Compute flux components
    dSdx_edges .= diff_x(S[:,2:end - 1]) / Δx
    dSdy_edges .= diff_y(S[2:end - 1,:]) / Δy
    Fx .= .-avg_y(D) .* dSdx_edges
    Fy .= .-avg_x(D) .* dSdy_edges

    #  Flux divergence
    inn(dH) .= .-(diff_x(Fx) / Δx .+ diff_y(Fy) / Δy) # MB to be added here
end

"""
    SIA(H, A, context)

Compute a step of the Shallow Ice Approximation UDE in a forward model. Allocates memory.
"""
function SIA(H, A, context)
    # Retrieve parameters
    B = context[1]

    # Update glacier surface altimetry
    S = B .+ H

    # All grid variables computed in a staggered grid
    # Compute surface gradients on edges
    dSdx = diff_x(S) / Δx
    dSdy = diff_y(S) / Δy
    ∇S = (avg_y(dSdx).^2 .+ avg_x(dSdy).^2).^((n - 1)/2)

    Γ = 2 * A * (ρ * g)^n / (n+2) # 1 / m^3 s
    D = Γ .* avg(H).^(n + 2) .* ∇S

    # Compute flux components
    dSdx_edges = diff_x(S[:,2:end - 1]) / Δx
    dSdy_edges = diff_y(S[2:end - 1,:]) / Δy
    Fx = .-avg_y(D) .* dSdx_edges
    Fy = .-avg_x(D) .* dSdy_edges

    #  Flux divergence
    @tullio dH[i,j] := -(diff_x(Fx)[pad(i-1,1,1),pad(j-1,1,1)] / Δx + diff_y(Fy)[pad(i-1,1,1),pad(j-1,1,1)] / Δy) # MB to be added here

    return dH
end

"""
    avg_surface_V(H, B, temp)

Computes the average ice velocity for a given input temperature
"""
function avg_surface_V(H, B, temp, sim)

    # Update glacier surface altimetry
    S = B .+ H

    # All grid variables computed in a staggered grid
    # Compute surface gradients on edges
    dSdx = diff_x(S) / Δx
    dSdy = diff_y(S) / Δy
    ∇S = (avg_y(dSdx).^2 .+ avg_x(dSdy).^2).^((n - 1)/2)

    @assert (sim == "UDE" || sim == "PDE") "Wrong type of simulation. Needs to be 'UDE' or 'PDE'."
    if sim == "UDE"
        A = predict_A̅(UA, θ, [temp]) # FastChain prediction requires explicit parameters
    elseif sim == "PDE"
        A = A_fake(temp, noise)
    end
    Γ₂ = 2 * A * (ρ * g)^n / (n+1)     # 1 / m^3 s
    # Zygote.ignore() do
    #     @infiltrate

    # end
    D = Γ₂ .* avg(H).^(n + 1) .* ∇S

    # Compute averaged surface velocities
    Vx = - D .* avg_y(dSdx)
    Vy = - D .* avg_x(dSdy)

    return Vx, Vy

end

"""
    A_fake(temp, noise=false)

Fake law establishing a theoretical relationship between ice viscosity (A) and long-term air temperature.
"""
function A_fake(temp, noise=false)
    A = @. minA + (maxA - minA) * ((temp-minT)/(maxT-minT) )^2
    if noise
        A = A .+ randn(rng_seed(), length(temp)).*4e-17
    end

    return A
end

"""
    predict_A̅(UA, θ, temp)

Predicts the value of A with a neural network based on the long-term air temperature.
"""
predict_A̅(UA, θ, temp) = UA(temp, θ) .* 1e-16

function fake_temp_series(t, means=Array{Float64}([0,-2.0,-3.0,-5.0,-10.0,-12.0,-14.0,-15.0,-20.0]))
    temps, norm_temps, norm_temps_flat = [],[],[]
    for mean in means
        push!(temps, mean .+ rand(t).*1e-1) # static
        append!(norm_temps_flat, mean .+ rand(t).*1e-1) # static
    end

    # Normalise temperature series
    norm_temps_flat = Flux.normalise([norm_temps_flat...]) # requires splatting

    # Re-create array of arrays
    for i in 1:t₁:length(norm_temps_flat)
        push!(norm_temps, norm_temps_flat[i:i+(t₁-1)])
    end

    return temps, norm_temps
end

sigmoid_A(x) = minA_out + (maxA_out - minA_out) / ( 1 + exp(-x) )

end # @everywhere


text:
module LevelSet
    include("./sdistance.jl")
    using .Sdistance # exportされたmethodのみ使える
    export signedDistance2D
end

max_stars_repo_name: jabelic/senior-thesis-julia

max_issues_repo_name: jabelic/senior-thesis-julia

max_forks_repo_name: jabelic/senior-thesis-julia


text:
using Test

m = test_page_model()
include("../src/components/LGforcing.jl")

add_comp!(m, LGforcing)

set_param!(m, :LGforcing, :c_LGconcentration, readpagedata(m,"test/validationdata/c_LGconcentration.csv"))

# run Model
run(m)

forcing=m[:LGforcing,:f_LGforcing]
forcing_compare=readpagedata(m,"test/validationdata/f_LGforcing.csv")

@test forcing ≈ forcing_compare rtol=1e-3


text:
function prefetchisagoodidea(ls::LoopSet, op::Operation, td::UnrollArgs)
    # return false
    @unpack u₁, u₁loopsym, u₂loopsym, vloopsym, vstep, u₂max, suffix = td
    length(loopdependencies(op)) ≤ 1 && return 0
    isvectorized(op) || return 0
    ((u₁ > 1) & (u₂max > 1)) || return 0
    u₂loopsym === Symbol("##undefined##") && return 0
    # @show cache_lnsze(ls) reg_size(ls) pointer_from_objref(ls.register_size)
    dontskip = (cache_lnsze(ls) ÷ reg_size(ls)) - 1
    # u₂loopsym is vectorized
    # u₁vectorized = vectorized === u₁loopsym
    u₂vectorized = vloopsym === u₂loopsym
    ((suffix ≠ -1) && u₂vectorized && !iszero(suffix & dontskip)) && return 0
    innermostloopsym = first(names(ls))
    loopedindex = op.ref.loopedindex
    if length(loopedindex) > 1 && first(loopedindex)
        indices = getindices(op)
        if first(indices) === vloopsym && first(indices) !== innermostloopsym
            (isone(vstep) && isone(first(getstrides(op)))) || return 0
            # We want at least 4 reuses per load
            uses = Core.ifelse(isu₁unrolled(op), 1, u₁)
            uses = Core.ifelse(isu₂unrolled(op), uses, uses * u₂max)
            uses < 4 && return 0
            innermostloopindv = findall(map(isequal(innermostloopsym), getindices(op)))
            isone(length(innermostloopindv)) || return 0
            innermostloopind = first(innermostloopindv)
            if prod(s -> Float64(looplength(ls, s)), @view(indices[1:innermostloopind-1])) ≥ 120.0 && length(getloop(ls, innermostloopsym)) ≥ 120
                if getoffsets(op)[innermostloopind] < 120
                    for opp ∈ operations(ls)
                        if iscompute(opp) && (innermostloopsym ∈ loopdependencies(opp)) &&
                            load_constrained(opp, u₁loopsym, u₂loopsym, innermostloopsym, true)
                            return 0
                        end
                    end
                    return innermostloopind
                end
            end
        end
    end
    0
end
function add_prefetches!(q::Expr, ls::LoopSet, op::Operation, td::UnrollArgs, prefetchind::Int)
    @unpack u₁, u₁loopsym, u₂loopsym, vloopsym, u₂max = td
    # we should only be here if `unitsride(vloop)`
    dontskip = (cache_lnsze(ls) ÷ reg_size(ls)) - 1
    ptr = vptr(op)
    innermostloopsym = first(names(ls))
    us = ls.unrollspecification[]
    prefetch_distance = u₁loopsym === innermostloopsym ? us.u₁ : ( u₂loopsym === innermostloopsym ? us.u₂ : 1 )
    # prefetch_distance = u₁loopsym === innermostloopsym ? u₁ : ( u₂loopsym === innermostloopsym ? u₂max : 1 )
    prefetch_multiplier = 5
    prefetch_distance *= prefetch_multiplier
    offsets = getoffsets(op) # what a hack
    inner_offset = offsets[prefetchind]
    prefetchstride = prefetch_distance * getstrides(op)[prefetchind]
    prefetchloop_step = step(getloop(ls, getindices(op)[prefetchind]))
    if isknown(prefetchloop_step)
        prefetchstride *= gethint(prefetchloop_step)
    end
    offsets[prefetchind] = inner_offset + prefetchstride
    gespinds = mem_offset_u(op, td, indices_calculated_by_pointer_offsets(ls, op.ref), false, 0)
    offsets[prefetchind] = inner_offset
    ptr = vptr(op)
    gptr = Symbol(ptr, "##GESPEDPREFETCH##")
    if !isknown(prefetchloop_step)
        for i ∈ eachindex(gespinds.args)
            if i == prefetchind
                gespinds.args[i] = mulexpr(getsym(prefetchloop_step), gespinds.args[i])
            end
            # gespinds.args[i] = Expr(:call, lv(:data), gespinds.args[i])
        end
    end
    push!(q.args, Expr(:(=), gptr, Expr(:call, lv(:gesp), ptr, gespinds)))

    inds = Expr(:tuple)
    indices = getindicesonly(op)

    i = 0
    for (j,ind) ∈ enumerate(indices)
        push!(inds.args, Expr(:call, lv(:Zero)))
        (ind == u₁loopsym) && (i = j)
    end
    push!(q.args, Expr(:call, lv(:prefetch0), gptr, copy(inds)))
    for u ∈ 1:u₁-1
        # for u ∈ umin:min(umin,U-1)
        # (u₁loopsym === vloopsym && !iszero(u & dontskip)) && continue
        if u₁loopsym === vloopsym
            iszero(u & dontskip) || continue
            # W = ls.vector_width[]
            # if W != 0
            #     inds.args[i] = staticexpr(W*u)
            # else
            if isone(u)
                inds.args[i] = VECTORWIDTHSYMBOL
            else
                inds.args[i] = mulexpr(VECTORWIDTHSYMBOL, u)
            end
        else
            inds.args[i] = staticexpr(u)
        end
        push!(q.args, Expr(:call, lv(:prefetch0), gptr, copy(inds)))
    end
    nothing
end
broadcastedname(mvar) = Symbol(mvar, "##broadcasted##")
function pushbroadcast!(q::Expr, mvar::Symbol)
    push!(q.args, Expr(:(=), broadcastedname(mvar), Expr(:call, lv(:vbroadcast), VECTORWIDTHSYMBOL, mvar)))
end
function lower_load_no_optranslation!(
    q::Expr, ls::LoopSet, op::Operation, td::UnrollArgs, mask::Bool, inds_calc_by_ptr_offset::Vector{Bool}
)
    @unpack u₁, u₁loopsym, u₂loopsym, vloopsym, suffix = td
    loopdeps = loopdependencies(op)
    # @assert isvectorized(op)
    opu₁ = isu₁unrolled(op)

    u = ifelse(opu₁, u₁, 1)
    mvar = Symbol(variable_name(op, Core.ifelse(isu₂unrolled(op), suffix,-1)), '_', u)
    falseexpr = Expr(:call, lv(:False)); rs = staticexpr(reg_size(ls))
    if all(op.ref.loopedindex) && !rejectcurly(op)
        inds = unrolledindex(op, td, mask, inds_calc_by_ptr_offset)
        loadexpr = Expr(:call, lv(:_vload), vptr(op), inds)
        add_memory_mask!(loadexpr, op, td, mask)
        push!(loadexpr.args, falseexpr, rs) # unaligned load
        push!(q.args, Expr(:(=), mvar, loadexpr))
    elseif u₁ > 1
        t = Expr(:tuple)
        for u ∈ 1:u₁
            inds = mem_offset_u(op, td, inds_calc_by_ptr_offset, true, u-1)
            loadexpr = Expr(:call, lv(:_vload), vptr(op), inds)
            domask = mask && (isvectorized(op) & ((u == u₁) | (vloopsym !== u₁loopsym)))
            add_memory_mask!(loadexpr, op, td, domask)
            push!(loadexpr.args, falseexpr, rs)
            push!(t.args, loadexpr)
            # push!(q.args, Expr(:(=), mvar, loadexpr))
        end
        push!(q.args, Expr(:(=), mvar, Expr(:call, lv(:VecUnroll), t)))
    else
        inds = mem_offset_u(op, td, inds_calc_by_ptr_offset, true, 0)
        loadexpr = Expr(:call, lv(:_vload), vptr(op), inds)
        add_memory_mask!(loadexpr, op, td, mask)
        push!(loadexpr.args, falseexpr, rs)
        push!(q.args, Expr(:(=), mvar, loadexpr))
    end
    if isvectorized(op)
        prefetchind = prefetchisagoodidea(ls, op, td)
        iszero(prefetchind) || add_prefetches!(q, ls, op, td, prefetchind)
    elseif any(isvectorized, children(op))
        pushbroadcast!(q, mvar)
    end
    nothing
end
# function lower_load_vectorized!(
#     q::Expr, ls::LoopSet, op::Operation, td::UnrollArgs, mask::Union{Nothing,Symbol,Unsigned} = nothing
# )
#     @unpack u₁, u₁loopsym, u₂loopsym, vectorized, suffix = td
#     loopdeps = loopdependencies(op)
#     @assert isvectorized(op)
#     opu₁ = isu₁unrolled(op)
#     inds_calc_by_ptr_offset = indices_calculated_by_pointer_offsets(ls, op.ref)
#     if opu₁
#         umin = 0
#         U = u₁
#     else
#         umin = -1
#         U = 0
#     end
#     # Urange = unrolled ∈ loopdeps ? 0:U-1 : 0
#     var = variable_name(op, suffix)
#     for u ∈ umin:U-1
#         td = UnrollArgs(td, u)
#         pushvectorload!(q, op, var, td, U, vectorized, mask, opu₁, inds_calc_by_ptr_offset, reg_size(ls))
#     end
#     prefetchind = prefetchisagoodidea(ls, op, td)
#     iszero(prefetchind) || add_prefetches!(q, ls, op, td, prefetchind, umin)
#     nothing
# end
function indisvectorized(ls::LoopSet, ind::Symbol)
    for op ∈ operations(ls)
        ((op.variable === ind) && isvectorized(op)) && return true
    end
    false
end
@inline firstunroll(vu::VecUnroll) = getfield(getfield(vu,:data),1,false)
@inline firstunroll(x) = x
@inline lastunroll(vu::VecUnroll) = last(getfield(vu,:data))
@inline lastunroll(x) = x
@inline unmm(x) = x
@inline unmm(x::MM) = getfield(x, :i)
function lower_load_for_optranslation!(
    q::Expr, op::Operation, posindicator::UInt8, ls::LoopSet, td::UnrollArgs, mask::Bool, translationind::Int
)
    @unpack u₁loop, u₂loop, vloop, u₁, u₂max, suffix = td

    # @unpack u₁, u₁loopsym, u₂loopsym, vloopsym, u₂max, suffix = td
    iszero(suffix) || return

    total_unroll = u₁ + u₂max - 1

    mref = op.ref
    inds_by_ptroff = indices_calculated_by_pointer_offsets(ls, mref)
    # initial offset pointer

    # Unroll directions can be + or -
    # we want to start at minimum position.
    step₁ = gethint(step(u₁loop))
    step₂ = gethint(step(u₂loop))

    # abs of steps are equal
    equal_steps = (step₁ == step₂) ⊻ (posindicator ≠ 0x03)
    # @show step₁, step₂, posindicator, equal_steps
    # _td = UnrollArgs(u₁loop, u₂loop, vloop, total_unroll, u₂max, Core.ifelse(equal_steps, 0, u₂max - 1))
    _td = UnrollArgs(u₁loop, u₂loop, vloop, u₁, u₂max, Core.ifelse(equal_steps, 0, u₂max - 1))
    gespinds = mem_offset(op, _td, inds_by_ptroff, false)
    ptr = vptr(op)
    gptr = Symbol(ptr, "##GESPED##")
    for i ∈ eachindex(gespinds.args)
        if i == translationind
            gespinds.args[i] = Expr(:call, lv(Core.ifelse(equal_steps, :firstunroll, :lastunroll)), gespinds.args[i])
        # else
        #     gespinds.args[i] = Expr(:call, lv(:unmm), gespinds.args[i])
        end
    end
    push!(q.args, Expr(:(=), gptr, Expr(:call, lv(:gesp), ptr, gespinds)))

    fill!(inds_by_ptroff, true)

    @unpack ref, loopedindex = mref
    indices = copy(getindices(ref))
    # old_translation_index = indices[translationind]
    # indices[translationind] = u₁loop.itersymbol
    # @show indices, translationind, vloop
    # getindicesonly returns a view of `getindices`
    dummyref = ArrayReference(ref.array, indices, zero(getoffsets(ref)), getstrides(ref))
    # loopedindex[translationind] = true
    dummymref = ArrayReferenceMeta(dummyref, fill!(similar(loopedindex), true), gptr)
    indonly = getindicesonly(dummyref)
    for i ∈ eachindex(indonly)
        if i == translationind
            indonly[i] = u₁loop.itersymbol
        elseif !loopedindex[i]
            ind = indonly[i]
            for indop ∈ operations(ls)
                if isvectorized(indop) & (name(indop) === ind)
                    indonly[i] = vloop.itersymbol
                    break
                end
            end
        end
    end
    # @show indices
    _td = UnrollArgs(u₁loop, u₂loop, vloop, total_unroll, u₂max, -1)
    op.ref = dummymref
    # @show isu₁unrolled(op), isu₂unrolled(op)
    _lower_load!(q, ls, op, _td, mask)
    # set old values
    op.ref = mref
    # loopedindex[translationind] = false
    # indices[translationind] = old_translation_index

    shouldbroadcast = (!isvectorized(op)) && any(isvectorized, children(op))

    # now we need to assign the `Vec`s from the `VecUnroll` to the correct name.
    variable_name_u = Symbol(variable_name(op, -1), '_', total_unroll)
    variable_name_data = Symbol(variable_name_u, "##data##")
    push!(q.args, :($variable_name_data = getfield($variable_name_u, 1)))
    if shouldbroadcast
        broadcasted_data = broadcastedname(variable_name_data)
        push!(q.args, :($broadcasted_data = getfield($(broadcastedname(variable_name_u)), 1)))
    end
    gf = GlobalRef(Core,:getfield)
    for u₂ ∈ 0:u₂max-1
        variable_name_u₂ = Symbol(variable_name(op, u₂), '_', u₁)
        t = Expr(:tuple)
        if shouldbroadcast
            tb = Expr(:tuple)
        end
        for u ∈ 1:u₁
            uu = if equal_steps
                u + u₂
            else
                u - u₂ + u₂max - 1
            end
            push!(t.args, :($gf($variable_name_data, $uu)))
            if shouldbroadcast
                push!(tb.args, :($gf($broadcasted_data, $uu)))
            end
        end
        push!(q.args, Expr(:(=), variable_name_u₂, Expr(:call, lv(:VecUnroll), t)))
        if shouldbroadcast
            push!(q.args, Expr(:(=), broadcastedname(variable_name_u₂), Expr(:call, lv(:VecUnroll), tb)))
        end
    end
    nothing
end

# TODO: this code should be rewritten to be more "orthogonal", so that we're just combining separate pieces.
# Using sentinel values (eg, T = -1 for non tiling) in part to avoid recompilation.
function lower_load!(
    q::Expr, op::Operation, ls::LoopSet, td::UnrollArgs, mask::Bool
)
    @unpack u₁, u₂max, u₁loopsym, u₂loopsym, vloopsym, suffix = td
    if (suffix != -1) && ls.loadelimination[]
        if (u₁ > 1) & (u₂max > 1)
            istr, ispl = isoptranslation(ls, op, UnrollSymbols(u₁loopsym, u₂loopsym, vloopsym))
            if istr ≠ 0x00
                return lower_load_for_optranslation!(q, op, ispl, ls, td, mask, istr)
            end
        end
        if (suffix > 0) && (u₂loopsym !== vloopsym)
            mno, id = maxnegativeoffset(ls, op, u₂loopsym)
            if -suffix < mno < 0 # already checked that `suffix != -1` above
                varnew = variable_name(op, suffix)
                varold = variable_name(operations(ls)[id], suffix + mno)
                opold = operations(ls)[id]
                u = isu₁unrolled(op) ? u₁ : 1
                push!(q.args, Expr(:(=), Symbol(varnew, '_', u), Symbol(varold, '_', u)))
                return
            end
        end
    end
    _lower_load!(q, ls, op, td, mask)
end
function _lower_load!(
    q::Expr, ls::LoopSet, op::Operation, td::UnrollArgs, mask::Bool, inds_calc_by_ptr_offset::Vector{Bool} = indices_calculated_by_pointer_offsets(ls, op.ref)
)
    if rejectinterleave(op)
        lower_load_no_optranslation!(q, ls, op, td, mask, inds_calc_by_ptr_offset)
    else
        omop = offsetloadcollection(ls)
        batchid, opind = omop.batchedcollectionmap[identifier(op)]
        if opind == 1
            collectionid, copind = omop.opidcollectionmap[identifier(op)]
            opidmap = offsetloadcollection(ls).opids[collectionid]
            idsformap = omop.batchedcollections[batchid]
            lower_load_collection!(q, ls, opidmap, idsformap, td, mask, inds_calc_by_ptr_offset)
        end
    end
end
function additive_vectorized_loopinductvar_only(op::Operation)
    isvectorized(op) || return true
    isloopvalue(op) && return true
    iscompute(op) || return false
    additive_instr = (:add_fast, :(+), :vadd, :identity, :sub_fast, :(-), :vsub)
    Base.sym_in(instruction(op).instr, additive_instr) || return false
    return all(additive_vectorized_loopinductvar_only, parents(op))
end
# Checks if we cannot use `Unroll`
function rejectcurly(ls::LoopSet, op::Operation, td::UnrollArgs)
    @unpack u₁loopsym, vloopsym = td
    rejectcurly(ls, op, u₁loopsym, vloopsym)
end
function rejectcurly(ls::LoopSet, op::Operation, u₁loopsym::Symbol, vloopsym::Symbol)
    indices = getindicesonly(op)
    li = op.ref.loopedindex
    AV = AU = false
    for (n,ind) ∈ enumerate(indices)
        # @show AU, op, n, ind, vloopsym, u₁loopsym
        if li[n]
            if ind === vloopsym
                AV && return true
                AV = true
            end
            if ind === u₁loopsym
                AU && return true
                AU = true
            end
        else
            opp = findop(parents(op), ind)
            # @show opp
            if isvectorized(opp)
                AV && return true
                AV = true
            end
            if (u₁loopsym === CONSTANTZEROINDEX) ? (CONSTANTZEROINDEX ∈ loopdependencies(opp)) : (isu₁unrolled(opp))
                AU && return true
                AU = true
            end
        end
    end
    false
end
function rejectinterleave(ls::LoopSet, op::Operation, vloop::Loop, idsformap::SubArray{Tuple{Int,Int}, 1, Vector{Tuple{Int,Int}}, Tuple{UnitRange{Int}}, true})
    vloopsym = vloop.itersymbol; strd = step(vloop)
    isknown(strd) || return true
    # TODO: reject if there is a vectorized !loopedindex index
    for (li,ind) ∈ zip(op.ref.loopedindex,getindicesonly(op))
        li && continue
        for indop ∈ operations(ls)
            if (name(indop) === ind) && isvectorized(indop)
                additive_vectorized_loopinductvar_only(indop) || return true # so that it is `MM`
            end
        end
    end
    (first(getindices(op)) === vloopsym) && (length(idsformap) ≠ first(getstrides(op)) * gethint(strd))
end
function lower_load_collection!(
    q::Expr, ls::LoopSet, opidmap::Vector{Int},
    idsformap::SubArray{Tuple{Int,Int}, 1, Vector{Tuple{Int,Int}}, Tuple{UnitRange{Int}}, true},
    ua::UnrollArgs, mask::Bool, inds_calc_by_ptr_offset::Vector{Bool}
)
    @unpack u₁, u₁loopsym, u₂loopsym, vloopsym, vloop, suffix = ua
    ops = operations(ls)
    nouter = length(idsformap)
    # ua = UnrollArgs(nouter, unrollsyms, u₂, 0)
    # idsformap contains (index, offset) pairs
    op = ops[opidmap[first(first(idsformap))]]
    opindices = getindices(op)
    interleave = first(opindices) === vloopsym
    # construct dummy unrolled loop
    offset_dummy_loop = Loop(first(opindices), MaybeKnown(1), MaybeKnown(1024), MaybeKnown(1), Symbol(""), Symbol(""))
    unrollcurl₂ = unrolled_curly(op, nouter, offset_dummy_loop, vloop, mask, 1) # interleave always 1 here
    inds = mem_offset_u(op, ua, inds_calc_by_ptr_offset, false)
    falseexpr = Expr(:call, lv(:False)); rs = staticexpr(reg_size(ls));
    if isu₁unrolled(op) && u₁ > 1 # both unrolled
        if interleave # TODO: handle this better than using `rejectinterleave`
            interleaveval = -nouter
        else
            interleaveval = 0
        end
        unrollcurl₁ = unrolled_curly(op, u₁, ua.u₁loop, vloop, mask, interleaveval)
        inds = Expr(:call, unrollcurl₁, inds)
    end
    uinds = Expr(:call, unrollcurl₂, inds)
    vp = vptr(op)
    loadexpr = Expr(:call, lv(:_vload), vp, uinds)
    # not using `add_memory_mask!(storeexpr, op, ua, mask)` because we checked `isconditionalmemop` earlier in `lower_load_collection!`
    (mask && isvectorized(op)) && push!(loadexpr.args, MASKSYMBOL)
    push!(loadexpr.args, falseexpr, rs)
    collectionname = Symbol(vp, "##collection##number#", opidmap[first(first(idsformap))], "#", suffix, "##size##", nouter, "##u₁##", u₁)
    # getfield to extract data from `VecUnroll` object, so we have a tuple
    push!(q.args, Expr(:(=), collectionname, Expr(:call, :getfield, loadexpr, 1)))
    u = Core.ifelse(isu₁unrolled(op), u₁, 1)
    for (i,(opid,o)) ∈ enumerate(idsformap)
        _op = ops[opidmap[opid]]
        mvar = Symbol(variable_name(_op, Core.ifelse(isu₂unrolled(_op), suffix, -1)), '_', u)
        push!(q.args, Expr(:(=), mvar, Expr(:call, :getfield, collectionname, i, false)))
    end
end


text:
# This code is the Julia version of Marco Cuturi's code in Matlab
using Distributions, Plots
include("../ot/sinkhornTransport.jl")
srand(123) # Setting the seed
# relevant dimensions in this example.
d1=100
d2=120
N=40
# draw randomly a symmetric cost matrix which is zero on the diagonal. this
# is not a distance matrix, but this suffices to test the script below.
M=rand(d1,d2)
M=M./median(M[:]) # normalize to get unit median.
# set lambda
lambda=200
# the matrix to be scaled.
K=exp.(-lambda*M)
# in practical situations it might be a good idea to do the following:
# %K(K<1e-100)=1e-100;
[K[t]<=1e-100 ? K[t]=1e-100 : K[t]=K[t]  for t in eachindex(K)]
# pre-compute matrix U, the Schur product of K and M.
U=K.*M
## Example with 1-vs-N mode
display(" ")
display("Example when Computing distances of 1-vs-N histograms")
# draw and normalize 1 point in the simplex with a few zeros (this is not a uniform sampling)
a=full(sprand(d1,1,.8))
a=a./sum(a)
# draw and normalize N points in the simplex with a few zeros (not uniform)
b=full(sprand(d2,N,.8))
b=broadcast(/,b,sum(b,1))
display("Computing $N distances from a to b_1, ... b_$N")
D,lowerEMD,l,m=sinkhornTransport(a,b,K,U,lambda,"distanceRelativeDecrease",Inf,.5e-2,50,1) # running with VERBOSE
display("Done computing distances")
display(" ")
# % Example of other types of executions, with much smaller tolerances.
# % D1=sinkhornTransport(a,b,K,U,lambda,'marginalDifference',inf,1e-5,[],1);
# % D2=sinkhornTransport(a,b,K,U,lambda,'distanceRelativeDecrease',inf);
# % D3=sinkhornTransport(a,b,K,U,lambda,'distanceRelativeDecrease',1);
# #
# figure()
# cla;
display("Display Vector of Distances and Lower Bounds on EMD");
#
# bar(D,'b');
# hold on
# bar(lowerEMD,'r');
# legend({'Sinkhorn Divergence','Lower bound on EMD'});
# axis tight; title(['Dual-Sinkhorn Divergence and Lower Bound on EMD for 1-vs-',num2str(N),' pairs'],'FontSize',16); set(gca,'FontSize',16)
#
# % choose a random histogram b_i.
# i=round(N*rand());
# disp(['Display (smoothed) optimal transport from a to b_',num2str(i),', which has been chosen randomly.']);
#
# T=bsxfun(@times,m(:,i)',(bsxfun(@times,l(:,i),K))); % this is the optimal transport.
# figure()
# imagesc(T);title(['T for (a,b_{',num2str(i),'})'],'FontSize',16);
# % check that T is indeed a transport matrix.
# disp(['Deviation of T from marginals: ', num2str(norm(sum(T)-b(:,i)')),' ',...
#     num2str(norm(sum(T,2)-a)),...
#     ' (should be close to zero)']);
#
#
#
# %% Example with N times 1-vs-1 mode
# disp(' ');
# disp('***** Example when Computing N distances between N different pairs ******');
# % a is now updated to be a matrix of column vectors in the simplex.
# a=full(sprand(d1,N,.5)); a=bsxfun(@rdivide,a,sum(a));
#
#
# disp(['Computing ',num2str(N),' distances (a_1,b_1), ... a_',num2str(N),'b_',num2str(N)]);
# [D,lowerEMD,l,m]=sinkhornTransport(a,b,K,U,lambda,[],[],[],[],1); % running with VERBOSE
# disp('Done computing distances');
# disp(' ');
#
# figure()
# cla;
# disp('Display Vector of Distances and Lower Bounds on EMD');
#
# bar(D,'b');
# hold on
# bar(lowerEMD,'r');
# legend({'Sinkhorn Divergence','Lower bound on EMD'});
# axis tight; title(['Dual-Sinkhorn Divergence and Lower Bound on EMD for ',num2str(N),' pairs (a_i,b_i)'],'FontSize',16); set(gca,'FontSize',16)
#
# % choose a random pair of histograms a_i, b_i.
# i=round(N*rand());
# disp(['Display (smoothed) optimal transport from a_',num2str(i),' to b_',num2str(i),', which has been chosen randomly.']);
#
# T=bsxfun(@times,m(:,i)',(bsxfun(@times,l(:,i),K))); % this is the optimal transport.
# figure()
# imagesc(T);title(['T for (a_{',num2str(i),'},b_{',num2str(i),'})'],'FontSize',16);
# % check that T is indeed a transport matrix.
# disp(['Deviation of T from marginals: ', num2str(norm(sum(T)-b(:,i)')),' ',...
#     num2str(norm(sum(T,2)-a(:,i))),...
#     ' (should be close to zero)']);


text:
using Matcha
using Test

# strings
x = "hey 1yo whatup?"

matched, matches = matchat(x, (
    Greed(x-> x in ('h', 'y', 'e'), 3:3),
    ' ',
    isnumeric
))
@test matches.matches == [SubString(x, 1, 3), SubString(x, 4, 4), SubString(x, 5, 5)]

x = "(ɔ◔‿◔)ɔ ♥ (⊙.⊙(☉̃ₒ☉)⊙.⊙)"

matched, matches = matchat(x, (
    '(', anything, ')'
))
@test matches[1] == "("
@test matches[2] == "ɔ◔‿◔"
@test matches[3] == ")"

# at beginning
x = [1,2,3,4,5,6]
test1(x) = x == 2
test2(x) = x == 3
matched, matches = matchat(x, (isodd, test1, test2))
@test matches.matches == [view(x, 1:1), view(x, 2:2), view(x, 3:3)]

# at end
x = [1,2,3,1,2,3]
match, matches = matchat(x, 4, (isodd, test1, test2))
@test vcat(matches.matches...) == [1,2,3]

# at middle
x = [1,2,1,2,3,3]
match, matches = matchat(x, 3, (isodd, test1, test2))
@test vcat(matches.matches...) == [1,2,3]

# at end
match, matches = matchat(x, 6, (isodd, test1, test2))
@test !match
@test vcat(matches.matches...) == [3] # we don't delete elements in match history

# even more endy
match, matches = matchat(x, 7, (isodd, test1, test2))
@test !match
@test vcat(matches.matches...) == Int[]

# with values
match, matches = matchat(x, 3, (1, 2, 3))
@test vcat(matches.matches...) == [1,2,3]

# test greed + match all
x = [1,2,3,4,5,6,7]
match, matches = matchat(x, 1, (1, anything, 5))
@test vcat(matches.matches...) == Int[(1:5)...]

match, matches = matchat(x, 4, (4, anything, 7))
@test vcat(matches.matches...) == Int[(4:7)...]

x = [1,2,3,4,5,6,7]
match, matches = matchat(x, 1, (Greed(1, 0:1), ))
@test vcat(matches.matches...) == Int[1]

# issue `Match anything, Greed(x, 0:1) not working #6`
x = [1, 2, 3, 7, 5]
match, matches = matchat(x, (anything, Greed(4, 0:1)))
@test length(matches.matches) == 1
@test matches[1] == x

x = [1, 2, 3, 4, 5]
match, matches = matchat(x, (anything, Greed(4, 0:1)))
@test length(matches.matches) == 2
@test matches[1] == [1, 2, 3]
@test matches[2] == [4]

match, matches = matchat(x, (Greed(2, 0:1), anything, Greed(4, 0:1)))
@test length(matches.matches) == 2
@test matches[1] == [1, 2, 3]
@test matches[2] == [4]

match, matches = matchat(x, (Greed(1, 0:1), anything, Greed(4, 0:1)))
@test length(matches.matches) == 3
@test matches[1] == [1]
@test matches[2] == [2, 3]
@test matches[3] == [4]

x = [1,2,3,4, 5,5,5, 7, 5,5,5, 8,9]

replaced = matchreplace(x-> (8, 8, 8), x, (Greed(5, 3:3),))

@test replaced ==[1,2,3,4, 8,8,8, 7, 8,8,8, 8,9]

x = [1,2,3,4, 5,5,5, 7, 5,5,5, 8,9]

replaced = matchreplace(x-> 7, x, (Greed(5, 3:3),))
@test replaced ==[1,2,3,4, 7, 7, 7, 8,9]

# include("expr_match_test.jl")