text:
using KernelAbstractions: @kernel, @index

const R_Earth = 6371.0e3    # Mean radius of the Earth [m] https://en.wikipedia.org/wiki/Earth

struct LatitudeLongitudeGrid{FT, TX, TY, TZ, M, MY, FX, FY, FZ, VX, VY, VZ, Arch} <: AbstractHorizontallyCurvilinearGrid{FT, TX, TY, TZ, Arch}
    architecture :: Arch
    Nx :: Int
    Ny :: Int
    Nz :: Int
    Hx :: Int
    Hy :: Int
    Hz :: Int
    Lx :: FT
    Ly :: FT
    Lz :: FT
    # All directions can be either regular (FX, FY, FZ) <: Number
    # or stretched (FX, FY, FZ) <: AbstractVector
    Δλᶠᵃᵃ :: FX
    Δλᶜᵃᵃ :: FX
    λᶠᵃᵃ  :: VX
    λᶜᵃᵃ  :: VX
    Δφᵃᶠᵃ :: FY
    Δφᵃᶜᵃ :: FY
    φᵃᶠᵃ  :: VY
    φᵃᶜᵃ  :: VY
    Δzᵃᵃᶠ :: FZ
    Δzᵃᵃᶜ :: FZ
    zᵃᵃᶠ  :: VZ
    zᵃᵃᶜ  :: VZ
    # Precomputed metrics M <: Nothing means metrics will be computed on the fly
    Δxᶠᶜᵃ :: M
    Δxᶜᶠᵃ :: M
    Δxᶠᶠᵃ :: M
    Δxᶜᶜᵃ :: M
    Δyᶠᶜᵃ :: MY
    Δyᶜᶠᵃ :: MY
    Azᶠᶜᵃ :: M
    Azᶜᶠᵃ :: M
    Azᶠᶠᵃ :: M
    Azᶜᶜᵃ :: M
    # Spherical radius
    radius :: FT

    function LatitudeLongitudeGrid{TX, TY, TZ}(architecture::Arch,
                                               Nλ, Nφ, Nz,
                                               Hλ, Hφ, Hz,
                                               Lλ::FT, Lφ::FT, Lz::FT,
                                               Δλᶠᵃᵃ :: FX, Δλᶜᵃᵃ :: FX,
                                                λᶠᵃᵃ :: VX,  λᶜᵃᵃ :: VX,
                                               Δφᵃᶠᵃ :: FY, Δφᵃᶜᵃ :: FY,
                                                φᵃᶠᵃ :: VY,  φᵃᶜᵃ :: VY,
                                               Δzᵃᵃᶠ :: FZ, Δzᵃᵃᶜ :: FZ,
                                                zᵃᵃᶠ :: VZ,  zᵃᵃᶜ :: VZ,
                                               Δxᶠᶜ::M, Δxᶜᶠ::M,
                                               Δxᶠᶠ::M, Δxᶜᶜ::M,
                                               Δyᶠᶜ::MY, Δyᶜᶠ::MY,
                                               Azᶠᶜ::M, Azᶜᶠ::M, Azᶠᶠ::M, Azᶜᶜ::M,
                                               radius::FT) where {Arch, FT,
                                                                  TX, TY, TZ,
                                                                  FX, FY, FZ,
                                                                  VX, VY, VZ,
                                                                  M, MY}

        return new{FT, TX, TY, TZ, M, MY, FX, FY, FZ, VX, VY, VZ, Arch}(architecture,
                                                                        Nλ, Nφ, Nz,
                                                                        Hλ, Hφ, Hz,
                                                                        Lλ, Lφ, Lz,
                                                                        Δλᶠᵃᵃ, Δλᶜᵃᵃ, λᶠᵃᵃ, λᶜᵃᵃ,
                                                                        Δφᵃᶠᵃ, Δφᵃᶜᵃ, φᵃᶠᵃ, φᵃᶜᵃ,
                                                                        Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ,
                                                                        Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ,
                                                                        Δyᶠᶜ, Δyᶜᶠ,
                                                                        Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ, radius)
    end
end

const XRegLatLonGrid = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number}
const YRegLatLonGrid = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any,    <:Number}
const ZRegLatLonGrid = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Any,    <:Any, <:Number}
const HRegLatLonGrid = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any, <:Number, <:Number}

regular_dimensions(::ZRegLatLonGrid) = tuple(3)

"""

latitude, longitude and z can be a 2-tuple that specifies the end of the domain (see RegularRectilinearDomain)
or an array or function that specifies the faces (see VerticallyStretchedRectilinearGrid)

"""

function LatitudeLongitudeGrid(architecture::AbstractArchitecture = CPU(),
                               FT::DataType = Float64;
                               precompute_metrics = false,
                               size,
                               latitude,
                               longitude,
                               z,
                               radius = R_Earth,
                               halo = (1, 1, 1))

    Nλ, Nφ, Nz, Hλ, Hφ, Hz, latitude, longitude, topo =
        validate_lat_lon_grid_args(latitude, longitude, size, halo)

    # Calculate all direction (which might be stretched)
    # A direction is regular if the domain passed is a Tuple{<:Real, <:Real},
    # it is stretched if being passed is a function or vector (as for the VerticallyStretchedRectilinearGrid)

    TX, TY, TZ = topo

    Lλ, λᶠᵃᵃ, λᶜᵃᵃ, Δλᶠᵃᵃ, Δλᶜᵃᵃ = generate_coordinate(FT, TX, Nλ, Hλ, longitude, architecture)
    Lφ, φᵃᶠᵃ, φᵃᶜᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ = generate_coordinate(FT, TY, Nφ, Hφ, latitude,  architecture)
    Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, TZ, Nz, Hz, z,         architecture)

    preliminary_grid = LatitudeLongitudeGrid{TX, TY, TZ}(architecture,
                                                         Nλ, Nφ, Nz,
                                                         Hλ, Hφ, Hz,
                                                         Lλ, Lφ, Lz,
                                                         Δλᶠᵃᵃ, Δλᶜᵃᵃ, λᶠᵃᵃ, λᶜᵃᵃ,
                                                         Δφᵃᶠᵃ, Δφᵃᶜᵃ, φᵃᶠᵃ, φᵃᶜᵃ,
                                                         Δzᵃᵃᶠ, Δzᵃᵃᶜ, zᵃᵃᶠ, zᵃᵃᶜ,
                                                         (nothing for i=1:10)..., FT(radius))

    return !precompute_metrics ? preliminary_grid : with_precomputed_metrics(preliminary_grid)
end

# architecture = CPU() default, assuming that a DataType positional arg
# is specifying the floating point type.
LatitudeLongitudeGrid(FT::DataType; kwargs...) = LatitudeLongitudeGrid(CPU(), FT; kwargs...)

""" Return a reproduction of `grid` with precomputed metric terms. """
function with_precomputed_metrics(grid)
    arch = architecture(grid)
    Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ, Δyᶠᶜ, Δyᶜᶠ, Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ = allocate_metrics(grid)
    wait(device_event(arch))

    precompute_curvilinear_metrics!(grid, Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ, Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ )
    wait(device_event(arch))

    Δyᶠᶜ, Δyᶜᶠ = precompute_Δy_metrics(grid, Δyᶠᶜ, Δyᶜᶠ)

    Nλ, Nφ, Nz = size(grid)
    Hλ, Hφ, Hz = halo_size(grid)
    TX, TY, TZ = topology(grid)

    return LatitudeLongitudeGrid{TX, TY, TZ}(architecture(grid),
                                             Nλ, Nφ, Nz,
                                             Hλ, Hφ, Hz,
                                             grid.Lx, grid.Ly, grid.Lz,
                                             grid.Δλᶠᵃᵃ, grid.Δλᶜᵃᵃ, grid.λᶠᵃᵃ, grid.λᶜᵃᵃ,
                                             grid.Δφᵃᶠᵃ, grid.Δφᵃᶜᵃ, grid.φᵃᶠᵃ, grid.φᵃᶜᵃ,
                                             grid.Δzᵃᵃᶠ, grid.Δzᵃᵃᶜ, grid.zᵃᵃᶠ, grid.zᵃᵃᶜ,
                                             Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ,
                                             Δyᶠᶜ, Δyᶜᶠ,
                                             Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ, grid.radius)
end

function validate_lat_lon_grid_args(latitude, longitude, size, halo)

    λ₁, λ₂ = get_domain_extent(longitude, size[1])
    @assert λ₁ < λ₂ && λ₂ - λ₁ ≤ 360

    φ₁, φ₂ = get_domain_extent(latitude, size[2])
    @assert -90 <= φ₁ < φ₂ <= 90

    (φ₁ == -90 || φ₂ == 90) &&
        @warn "Are you sure you want to use a latitude-longitude grid with a grid point at the pole?"

    Lλ = λ₂ - λ₁
    Lφ = φ₂ - φ₁

    TX = Lλ == 360 ? Periodic : Bounded
    TY = Bounded
    TZ = Bounded
    topo = (TX, TY, TZ)

    Nλ, Nφ, Nz = N = validate_size(TX, TY, TZ, size)
    Hλ, Hφ, Hz = H = validate_halo(TX, TY, TZ, halo)

    return Nλ, Nφ, Nz, Hλ, Hφ, Hz, latitude, longitude, topo
end

function domain_string(grid::LatitudeLongitudeGrid)
    λ₁, λ₂ = domain(topology(grid, 1), grid.Nx, grid.λᶠᵃᵃ)
    φ₁, φ₂ = domain(topology(grid, 2), grid.Ny, grid.φᵃᶠᵃ)
    z₁, z₂ = domain(topology(grid, 3), grid.Nz, grid.zᵃᵃᶠ)
    return "longitude λ ∈ [$λ₁, $λ₂], latitude ∈ [$φ₁, $φ₂], z ∈ [$z₁, $z₂]"
end

function show(io::IO, g::LatitudeLongitudeGrid)
    FT = eltype(g)
    TX, TY, TZ = topology(g)

    show_metrics = isnothing(g.Δxᶠᶜᵃ) ? "metrics are computed on the fly" :
                                        "metrics are pre-computed"

    return print(io, "LatitudeLongitudeGrid{$FT, $TX, $TY, $TZ} \n",
                     "             architecture: $(g.architecture)\n",
                     "                   domain: $(domain_string(g))\n",
                     "                 topology: ", (TX, TY, TZ), '\n',
                     "        size (Nx, Ny, Nz): ", (g.Nx, g.Ny, g.Nz), '\n',
                     "        halo (Hx, Hy, Hz): ", (g.Hx, g.Hy, g.Hz), '\n',
                     "             spacing in λ: ", show_coordinate(g.Δλᶜᵃᵃ, TX), '\n',
                     "             spacing in φ: ", show_coordinate(g.Δφᵃᶜᵃ, TY), '\n',
                     "             spacing in z: ", show_coordinate(g.Δzᵃᵃᶜ, TZ), '\n',
                     show_metrics)
end

# Node by node
@inline xnode(::Center, i, grid::LatitudeLongitudeGrid) = @inbounds grid.λᶜᵃᵃ[i]
@inline xnode(::Face,   i, grid::LatitudeLongitudeGrid) = @inbounds grid.λᶠᵃᵃ[i]
@inline ynode(::Center, j, grid::LatitudeLongitudeGrid) = @inbounds grid.φᵃᶜᵃ[j]
@inline ynode(::Face,   j, grid::LatitudeLongitudeGrid) = @inbounds grid.φᵃᶠᵃ[j]
@inline znode(::Center, k, grid::LatitudeLongitudeGrid) = @inbounds grid.zᵃᵃᶠ[k]
@inline znode(::Face,   k, grid::LatitudeLongitudeGrid) = @inbounds grid.zᵃᵃᶜ[k]

all_x_nodes(::Type{Center}, grid::LatitudeLongitudeGrid) = grid.λᶜᵃᵃ
all_x_nodes(::Type{Face},   grid::LatitudeLongitudeGrid) = grid.λᶠᵃᵃ
all_y_nodes(::Type{Center}, grid::LatitudeLongitudeGrid) = grid.φᵃᶜᵃ
all_y_nodes(::Type{Face},   grid::LatitudeLongitudeGrid) = grid.φᵃᶠᵃ
all_z_nodes(::Type{Center}, grid::LatitudeLongitudeGrid) = grid.zᵃᵃᶜ
all_z_nodes(::Type{Face},   grid::LatitudeLongitudeGrid) = grid.zᵃᵃᶠ

@inline x_domain(grid::LatitudeLongitudeGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} = domain(TX, grid.Nx, grid.λᶠᵃᵃ)
@inline y_domain(grid::LatitudeLongitudeGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} = domain(TY, grid.Ny, grid.φᵃᶠᵃ)
@inline z_domain(grid::LatitudeLongitudeGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ} = domain(TZ, grid.Nz, grid.zᵃᵃᶠ)

@inline cpu_face_constructor_x(grid::XRegLatLonGrid) = x_domain(grid)
@inline cpu_face_constructor_y(grid::YRegLatLonGrid) = y_domain(grid)
@inline cpu_face_constructor_z(grid::ZRegLatLonGrid) = z_domain(grid)

function on_architecture(new_arch, old_grid::LatitudeLongitudeGrid)
    old_properties = (old_grid.Δλᶠᵃᵃ, old_grid.Δλᶜᵃᵃ, old_grid.λᶠᵃᵃ,  old_grid.λᶜᵃᵃ,
                      old_grid.Δφᵃᶠᵃ, old_grid.Δφᵃᶜᵃ, old_grid.φᵃᶠᵃ,  old_grid.φᵃᶜᵃ,
                      old_grid.Δzᵃᵃᶠ, old_grid.Δzᵃᵃᶜ, old_grid.zᵃᵃᶠ,  old_grid.zᵃᵃᶜ,
                      old_grid.Δxᶠᶜᵃ, old_grid.Δxᶜᶠᵃ, old_grid.Δxᶠᶠᵃ, old_grid.Δxᶜᶜᵃ,
                      old_grid.Δyᶠᶜᵃ, old_grid.Δyᶜᶠᵃ,
                      old_grid.Azᶠᶜᵃ, old_grid.Azᶜᶠᵃ, old_grid.Azᶠᶠᵃ, old_grid.Azᶜᶜᵃ)

    new_properties = Tuple(arch_array(new_arch, p) for p in old_properties)

    TX, TY, TZ = topology(old_grid)

    return LatitudeLongitudeGrid{TX, TY, TZ}(new_arch,
                                             old_grid.Nx, old_grid.Ny, old_grid.Nz,
                                             old_grid.Hx, old_grid.Hy, old_grid.Hz,
                                             old_grid.Lx, old_grid.Ly, old_grid.Lz,
                                             new_properties...,
                                             old_grid.radius)
end

function Adapt.adapt_structure(to, grid::LatitudeLongitudeGrid)
    TX, TY, TZ = topology(grid)
    return LatitudeLongitudeGrid{TX, TY, TZ}(nothing,
                                             grid.Nx, grid.Ny, grid.Nz,
                                             grid.Hx, grid.Hy, grid.Hz,
                                             grid.Lx, grid.Ly, grid.Lz,
                                             Adapt.adapt(to, grid.Δλᶠᵃᵃ),
                                             Adapt.adapt(to, grid.Δλᶜᵃᵃ),
                                             Adapt.adapt(to, grid.λᶠᵃᵃ),
                                             Adapt.adapt(to, grid.λᶜᵃᵃ),
                                             Adapt.adapt(to, grid.Δφᵃᶠᵃ),
                                             Adapt.adapt(to, grid.Δφᵃᶜᵃ),
                                             Adapt.adapt(to, grid.φᵃᶠᵃ),
                                             Adapt.adapt(to, grid.φᵃᶜᵃ),
                                             Adapt.adapt(to, grid.Δzᵃᵃᶠ),
                                             Adapt.adapt(to, grid.Δzᵃᵃᶜ),
                                             Adapt.adapt(to, grid.zᵃᵃᶠ),
                                             Adapt.adapt(to, grid.zᵃᵃᶜ),
                                             Adapt.adapt(to, grid.Δxᶠᶜᵃ),
                                             Adapt.adapt(to, grid.Δxᶜᶠᵃ),
                                             Adapt.adapt(to, grid.Δxᶠᶠᵃ),
                                             Adapt.adapt(to, grid.Δxᶜᶜᵃ),
                                             Adapt.adapt(to, grid.Δyᶠᶜᵃ),
                                             Adapt.adapt(to, grid.Δyᶜᶠᵃ),
                                             Adapt.adapt(to, grid.Azᶠᶜᵃ),
                                             Adapt.adapt(to, grid.Azᶜᶠᵃ),
                                             Adapt.adapt(to, grid.Azᶠᶠᵃ),
                                             Adapt.adapt(to, grid.Azᶜᶜᵃ),
                                             grid.radius)
end

#####
##### On-the-fly computation of LatitudeLongitudeGrid metrics
#####

@inline hack_cosd(φ) = cos(π * φ / 180)
@inline hack_sind(φ) = sin(π * φ / 180)

@inline Δxᶠᶜᵃ(i, j, k, grid::LatitudeLongitudeGrid) = @inbounds grid.radius * hack_cosd(grid.φᵃᶜᵃ[j]) * deg2rad(grid.Δλᶠᵃᵃ[i])
@inline Δxᶜᶠᵃ(i, j, k, grid::LatitudeLongitudeGrid) = @inbounds grid.radius * hack_cosd(grid.φᵃᶠᵃ[j]) * deg2rad(grid.Δλᶜᵃᵃ[i])
@inline Δxᶠᶠᵃ(i, j, k, grid::LatitudeLongitudeGrid) = @inbounds grid.radius * hack_cosd(grid.φᵃᶠᵃ[j]) * deg2rad(grid.Δλᶠᵃᵃ[i])
@inline Δxᶜᶜᵃ(i, j, k, grid::LatitudeLongitudeGrid) = @inbounds grid.radius * hack_cosd(grid.φᵃᶜᵃ[j]) * deg2rad(grid.Δλᶜᵃᵃ[i])
@inline Δyᶜᶠᵃ(i, j, k, grid::LatitudeLongitudeGrid) = @inbounds grid.radius * deg2rad(grid.Δφᵃᶠᵃ[j])
@inline Δyᶠᶜᵃ(i, j, k, grid::LatitudeLongitudeGrid) = @inbounds grid.radius * deg2rad(grid.Δφᵃᶜᵃ[j])
@inline Azᶠᶜᵃ(i, j, k, grid::LatitudeLongitudeGrid) = @inbounds grid.radius^2 * deg2rad(grid.Δλᶠᵃᵃ[i]) * (hack_sind(grid.φᵃᶠᵃ[j+1]) - hack_sind(grid.φᵃᶠᵃ[j]))
@inline Azᶜᶠᵃ(i, j, k, grid::LatitudeLongitudeGrid) = @inbounds grid.radius^2 * deg2rad(grid.Δλᶜᵃᵃ[i]) * (hack_sind(grid.φᵃᶜᵃ[j])   - hack_sind(grid.φᵃᶜᵃ[j-1]))
@inline Azᶠᶠᵃ(i, j, k, grid::LatitudeLongitudeGrid) = @inbounds grid.radius^2 * deg2rad(grid.Δλᶠᵃᵃ[i]) * (hack_sind(grid.φᵃᶜᵃ[j])   - hack_sind(grid.φᵃᶜᵃ[j-1]))
@inline Azᶜᶜᵃ(i, j, k, grid::LatitudeLongitudeGrid) = @inbounds grid.radius^2 * deg2rad(grid.Δλᶜᵃᵃ[i]) * (hack_sind(grid.φᵃᶠᵃ[j+1]) - hack_sind(grid.φᵃᶠᵃ[j]))

@inline Δxᶠᶜᵃ(i, j, k, grid::XRegLatLonGrid) = @inbounds grid.radius * hack_cosd(grid.φᵃᶜᵃ[j]) * deg2rad(grid.Δλᶠᵃᵃ)
@inline Δxᶜᶠᵃ(i, j, k, grid::XRegLatLonGrid) = @inbounds grid.radius * hack_cosd(grid.φᵃᶠᵃ[j]) * deg2rad(grid.Δλᶜᵃᵃ)
@inline Δxᶠᶠᵃ(i, j, k, grid::XRegLatLonGrid) = @inbounds grid.radius * hack_cosd(grid.φᵃᶠᵃ[j]) * deg2rad(grid.Δλᶠᵃᵃ)
@inline Δxᶜᶜᵃ(i, j, k, grid::XRegLatLonGrid) = @inbounds grid.radius * hack_cosd(grid.φᵃᶜᵃ[j]) * deg2rad(grid.Δλᶜᵃᵃ)
@inline Δyᶜᶠᵃ(i, j, k, grid::YRegLatLonGrid) = @inbounds grid.radius * deg2rad(grid.Δφᵃᶠᵃ)
@inline Δyᶠᶜᵃ(i, j, k, grid::YRegLatLonGrid) = @inbounds grid.radius * deg2rad(grid.Δφᵃᶜᵃ)
@inline Azᶠᶜᵃ(i, j, k, grid::XRegLatLonGrid) = @inbounds grid.radius^2 * deg2rad(grid.Δλᶠᵃᵃ) * (hack_sind(grid.φᵃᶠᵃ[j+1]) - hack_sind(grid.φᵃᶠᵃ[j]))
@inline Azᶜᶠᵃ(i, j, k, grid::XRegLatLonGrid) = @inbounds grid.radius^2 * deg2rad(grid.Δλᶜᵃᵃ) * (hack_sind(grid.φᵃᶜᵃ[j])   - hack_sind(grid.φᵃᶜᵃ[j-1]))
@inline Azᶠᶠᵃ(i, j, k, grid::XRegLatLonGrid) = @inbounds grid.radius^2 * deg2rad(grid.Δλᶠᵃᵃ) * (hack_sind(grid.φᵃᶜᵃ[j])   - hack_sind(grid.φᵃᶜᵃ[j-1]))
@inline Azᶜᶜᵃ(i, j, k, grid::XRegLatLonGrid) = @inbounds grid.radius^2 * deg2rad(grid.Δλᶜᵃᵃ) * (hack_sind(grid.φᵃᶠᵃ[j+1]) - hack_sind(grid.φᵃᶠᵃ[j]))

#####
##### Utilities to precompute metrics
#####

@inline metrics_precomputed(::LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, Nothing}) = false
@inline metrics_precomputed(::LatitudeLongitudeGrid) = true

#####
##### Kernels that precompute the z- and x-metric
#####

@inline metric_worksize(grid::LatitudeLongitudeGrid)  = (length(grid.Δλᶜᵃᵃ), length(grid.φᵃᶜᵃ) - 1)
@inline metric_workgroup(grid::LatitudeLongitudeGrid) = (16, 16)

@inline metric_worksize(grid::XRegLatLonGrid)  =  length(grid.φᵃᶜᵃ) - 1
@inline metric_workgroup(grid::XRegLatLonGrid) =  16

function precompute_curvilinear_metrics!(grid, Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ, Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ)

    arch = grid.architecture

    workgroup, worksize  = metric_workgroup(grid), metric_worksize(grid)
    curvilinear_metrics! = precompute_metrics_kernel!(Architectures.device(arch), workgroup, worksize)

    event = curvilinear_metrics!(grid, Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ, Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ; dependencies=device_event(arch))
    wait(event)

    return nothing
end

@kernel function precompute_metrics_kernel!(grid::LatitudeLongitudeGrid, Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ, Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ)
    i, j = @index(Global, NTuple)

    # Manually offset x- and y-index
    i′ = i + grid.Δλᶜᵃᵃ.offsets[1]
    j′ = j + grid.φᵃᶜᵃ.offsets[1] + 1

    @inbounds begin
        Δxᶠᶜ[i′, j′] = Δxᶠᶜᵃ(i′, j′, 1, grid)
        Δxᶜᶠ[i′, j′] = Δxᶜᶠᵃ(i′, j′, 1, grid)
        Δxᶠᶠ[i′, j′] = Δxᶠᶠᵃ(i′, j′, 1, grid)
        Δxᶜᶜ[i′, j′] = Δxᶜᶜᵃ(i′, j′, 1, grid)
        Azᶠᶜ[i′, j′] = Azᶠᶜᵃ(i′, j′, 1, grid)
        Azᶜᶠ[i′, j′] = Azᶜᶠᵃ(i′, j′, 1, grid)
        Azᶠᶠ[i′, j′] = Azᶠᶠᵃ(i′, j′, 1, grid)
        Azᶜᶜ[i′, j′] = Azᶜᶜᵃ(i′, j′, 1, grid)
    end
end

@kernel function precompute_metrics_kernel!(grid::XRegLatLonGrid, Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ, Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ)
    j = @index(Global, Linear)

    # Manually offset y-index
    j′ = j + grid.φᵃᶜᵃ.offsets[1] + 1

    @inbounds begin
        Δxᶠᶜ[j′] = Δxᶠᶜᵃ(1, j′, 1, grid)
        Δxᶜᶠ[j′] = Δxᶜᶠᵃ(1, j′, 1, grid)
        Δxᶠᶠ[j′] = Δxᶠᶠᵃ(1, j′, 1, grid)
        Δxᶜᶜ[j′] = Δxᶜᶜᵃ(1, j′, 1, grid)
        Azᶠᶜ[j′] = Azᶠᶜᵃ(1, j′, 1, grid)
        Azᶜᶠ[j′] = Azᶜᶠᵃ(1, j′, 1, grid)
        Azᶠᶠ[j′] = Azᶠᶠᵃ(1, j′, 1, grid)
        Azᶜᶜ[j′] = Azᶜᶜᵃ(1, j′, 1, grid)
    end
end

####### Kernels that precompute the y-metric

function precompute_Δy_metrics(grid::LatitudeLongitudeGrid, Δyᶠᶜ, Δyᶜᶠ)
    arch = grid.architecture
    precompute_Δy! = precompute_Δy_kernel!(Architectures.device(arch), 16, length(grid.Δφᵃᶜᵃ) - 1)
    event = precompute_Δy!(grid, Δyᶠᶜ, Δyᶜᶠ; dependencies=device_event(arch))

    wait(event)
    return Δyᶠᶜ, Δyᶜᶠ
end

function  precompute_Δy_metrics(grid::YRegLatLonGrid, Δyᶠᶜ, Δyᶜᶠ)
    Δyᶜᶠ =  Δyᶜᶠᵃ(1, 1, 1, grid)
    Δyᶠᶜ =  Δyᶠᶜᵃ(1, 1, 1, grid)
    return Δyᶠᶜ, Δyᶜᶠ
end

@kernel function precompute_Δy_kernel!(grid, Δyᶠᶜ, Δyᶜᶠ)
    j  = @index(Global, Linear)

    # Manually offset y-index
    j′ = j + grid.Δφᵃᶜᵃ.offsets[1] + 1

    @inbounds begin
        Δyᶜᶠ[j′] = Δyᶜᶠᵃ(1, j′, 1, grid)
        Δyᶠᶜ[j′] = Δyᶜᶠᵃ(1, j′, 1, grid)
    end
end

#####
##### Metric memory allocation
#####

function allocate_metrics(grid::LatitudeLongitudeGrid)
    FT = eltype(grid)

    # preallocate quantities to ensure correct type and size
    grid_metrics = (:Δxᶠᶜ,
                    :Δxᶜᶠ,
                    :Δxᶠᶠ,
                    :Δxᶜᶜ,
                    :Azᶠᶜ,
                    :Azᶜᶠ,
                    :Azᶠᶠ,
                    :Azᶜᶜ)

    arch = grid.architecture

    if typeof(grid) <: XRegLatLonGrid
        offsets     = grid.φᵃᶜᵃ.offsets[1]
        metric_size = length(grid.φᵃᶜᵃ)
    else
        offsets     = (grid.Δλᶜᵃᵃ.offsets[1], grid.φᵃᶜᵃ.offsets[1])
        metric_size = (length(grid.Δλᶜᵃᵃ)   , length(grid.φᵃᶜᵃ))
    end

    for metric in grid_metrics
        parentM        = Symbol(metric, :_parent)
        @eval $parentM = zeros($FT, $metric_size...)
        @eval $metric  = OffsetArray(arch_array($arch, $parentM), $offsets...)
    end

    if typeof(grid) <: YRegLatLonGrid
        Δyᶠᶜ = FT(0.0)
        Δyᶜᶠ = FT(0.0)
    else
        parentC = zeros(FT, length(grid.Δφᵃᶜᵃ))
        parentF = zeros(FT, length(grid.Δφᵃᶜᵃ))
        Δyᶠᶜ    = OffsetArray(arch_array(arch, parentC), grid.Δφᵃᶜᵃ.offsets[1])
        Δyᶜᶠ    = OffsetArray(arch_array(arch, parentF), grid.Δφᵃᶜᵃ.offsets[1])
    end

    return Δxᶠᶜ, Δxᶜᶠ, Δxᶠᶠ, Δxᶜᶜ, Δyᶠᶜ, Δyᶜᶠ, Azᶠᶜ, Azᶜᶠ, Azᶠᶠ, Azᶜᶜ
end


text:
using IntervalArithmetic
using IntervalArithmetic.Complementary
using Test

include("construction.jl")
@testset "internals" begin include("special.jl") end
include("rand.jl")
include("test_ITF1788/run_ITF1788.jl")

ext: jl


text:
### Routines for constructing smooth functions that generate field data ###

abstract type AbstractSpatialField end

## Empty spatial field

"""
    EmptySpatialField()

Create a blank spatial field. This is primarily useful for initializing a
sum of spatial fields.

# Example
```jldoctest
julia> g = EmptySpatialField()
EmptySpatialField()

julia> g(2,3)
0.0
```
"""
struct EmptySpatialField <: AbstractSpatialField end
(g::EmptySpatialField)(a...) = Float64(0)

## Spatial Gaussian field ##

"""
    SpatialGaussian(σx,σy,x0,y0,A[,derivdir=0])

Set up a spatial field in the form of a Gaussian centered at `x0,y0` with
radii `σx` and `σy` in the respective directions and amplitude `A`. If the
optional parameter `deriv` is set to 1 or 2, then it returns the first
derivative of a Gaussian in that direction (`x` or `y`, respectively).

`SpatialGaussian(σx,σy,x0,y0,A,u,v[,derivdir=0])` generates a Gaussian
that convects at velocity `(u,v)`. It can be evaluated with an additional
argument for time.
"""
struct SpatialGaussian{CT,GX,GY} <: AbstractSpatialField
  gx :: GX
  gy :: GY
  A :: Float64
  u :: Float64
  v :: Float64
  SpatialGaussian(gx::Abstract1DProfile,gy::Abstract1DProfile,A,u,v) = new{true,typeof(gx),typeof(gy)}(gx,gy,A,u,v)
  SpatialGaussian(gx::Abstract1DProfile,gy::Abstract1DProfile,A) = new{false,typeof(gx),typeof(gy)}(gx,gy,A,0.0,0.0)
end

SpatialGaussian(σx::Real,σy::Real,x0::Real,y0::Real,A::Real;deriv::Int=0) =
                _spatialdgaussian(σx,σy,x0,y0,A,Val(deriv))
SpatialGaussian(σx::Real,σy::Real,x0::Real,y0::Real,A::Real,u::Real,v::Real;deriv::Int=0) =
                _spatialdgaussian(σx,σy,x0,y0,A,u,v,Val(deriv))

_spatialdgaussian(σx,σy,x0,y0,A,::Val{0}) = SpatialGaussian(Gaussian(σx,A) >> x0,
                                                            Gaussian(σy,1) >> y0,A)
_spatialdgaussian(σx,σy,x0,y0,A,::Val{1}) = SpatialGaussian(DGaussian(σx,A) >> x0,
                                                            Gaussian(σy,1) >> y0,A)
_spatialdgaussian(σx,σy,x0,y0,A,::Val{2}) = SpatialGaussian(Gaussian(σx,A) >> x0,
                                                            DGaussian(σy,1) >> y0,A)

_spatialdgaussian(σx,σy,x0,y0,A,u,v,::Val{0}) = SpatialGaussian(Gaussian(σx,A) >> x0,
                                                                Gaussian(σy,1) >> y0,A,u,v)
_spatialdgaussian(σx,σy,x0,y0,A,u,v,::Val{1}) = SpatialGaussian(DGaussian(σx,A) >> x0,
                                                                Gaussian(σy,1) >> y0,A,u,v)
_spatialdgaussian(σx,σy,x0,y0,A,u,v,::Val{2}) = SpatialGaussian(Gaussian(σx,A) >> x0,
                                                                DGaussian(σy,1) >> y0,A,u,v)

SpatialGaussian(σ,x0,y0,A;deriv::Int=0) = SpatialGaussian(σ,σ,x0,y0,A,deriv=deriv)

(g::SpatialGaussian{GX,GY})(x,y) where {CT,GX,GY} = g.gx(x)*g.gy(y)
# ignore the time argument if it is called with this...
(g::SpatialGaussian{false,GX,GY})(x,y,t) where {GX,GY} = g(x,y)
(g::SpatialGaussian{true,GX,GY})(x,y,t) where {GX,GY} = g.gx(x-g.u*t)*g.gy(y-g.v*t)

## Scaling spatial fields

struct ScaledField{N <: Real, P <: AbstractSpatialField} <: AbstractSpatialField
    s::N
    p::P
end
function show(io::IO, p::ScaledField)
    print(io, "$(p.s) × ($(p.p))")
end
s::Number * p::AbstractSpatialField = ScaledField(s, p)
-(p::AbstractSpatialField) = ScaledField(-1, p)

(p::ScaledField)(x,y) = p.s*p.p(x,y)
(p::ScaledField)(x,y,t) = p.s*p.p(x,y,t)

## Adding spatial fields together.

struct AddedFields{T <: Tuple} <: AbstractSpatialField
    ps::T
end
function show(io::IO, Σp::AddedFields)
    println(io, "AddedFields:")
    for p in Σp.ps
        println(io, "  $p")
    end
end

"""
    p₁::AbstractSpatialField + p₂::AbstractSpatialField

Add the fields so that `(p₁ + p₂)(x,y) = p₁(x,y) + p₂(x,y)`.

"""
+(p::AbstractSpatialField, Σp::AddedFields) = AddedFields((p, Σp.ps...))
+(Σp::AddedFields, p::AbstractSpatialField) = AddedFields((Σp.ps..., p))
function +(Σp₁::AddedFields, Σp₂::AddedFields)
    AddedFields((Σp₁..., Σp₂...))
end

-(p₁::AbstractSpatialField, p₂::AbstractSpatialField) = p₁ + (-p₂)

+(p::AbstractSpatialField...) = AddedFields(p)

# Evaluate at x,y , assuming that t = 0 for any time-varying member
function (Σp::AddedFields)(x,y)
    f = 0.0
    for p in Σp.ps
        f += p(x,y,0.0)
    end
    f
end

# Evaluate at x,y,t, ignoring t for any constant member
function (Σp::AddedFields)(x,y,t)
    f = 0.0
    for p in Σp.ps
        f += p(x,y,t)
    end
    f
end


text:
function process_string(m::Model, data)
    features = [fname for (fname, ftype) in getfeatures(m; ftype="all")
                if in(fname, keys(data)) && ftype == STRING]

    for fname in features
        data[fname] = fill_string(data[fname])
        data[fname] = string.(data[fname])
        dict = haskey(m.fdict, fname) ? m.fdict[fname] : nothing
        ids, dict = doc2ids(data[fname]; dict=dict)
        m.fdict[fname] = dict
        data[fname] = ids
    end
end

function process_int(m::Model, data)
    features = [fname for (fname, ftype) in getfeatures(m; ftype="all")
                if in(fname, keys(data)) && ftype == INT]

    for fname in features
        data[fname] = fill_int(data[fname])
    end
end

function process_float(m::Model, data)
    features = [fname for (fname, ftype) in getfeatures(m; ftype="all")
                if in(fname, keys(data)) && ftype == FLOAT]

    for fname in features
        if eltype(data[fname]) == String
            data[fname] = map(x->parse(Float64, x), data[fname])
        elseif eltype(data[fname]) == Int
            data[fname] = Float64.(data[fname])
        else
            data[fname]
        end
    end
end

function process_bin_category(m::Model, data)
    features = [fname for (fname, ftype) in getfeatures(m; ftype="all")
                if in(fname, keys(data)) && ftype == BINARYCATEGORY]

    for fname in features
        data[fname] = fill_bin_category(data[fname])
        data[fname] = string.(data[fname])
        dict = haskey(m.fdict, fname) ? m.fdict[fname] : nothing
        ids, dict = doc2ids(data[fname]; dict=dict)
        m.fdict[fname] = dict
        data[fname] = ids
    end
end

function process_category(m::Model, data)
    features = [fname for (fname, ftype) in getfeatures(m; ftype="all")
                if in(fname, keys(data)) && ftype == CATEGORY]

    for fname in features
        data[fname] = fill_category(data[fname])
        data[fname] = string.(data[fname])
        dict = haskey(m.fdict, fname) ? m.fdict[fname] : nothing
        ids, dict = doc2ids(data[fname]; dict=dict)
        m.fdict[fname] = dict
        data[fname] = ids
    end
end

function process_date(m::Model, data)
    features = [fname for (fname, ftype) in getfeatures(m; ftype="all")
                if in(fname, keys(data)) && ftype == DATE]

    for fname in features
        data[fname] = fill_date(data[fname])
    end

    alldates = [[Date(date) for date in data[fname]] for fname in features]

    for i in 1:length(alldates)
        for j in i+1:length(alldates)
            fname = "date" * string(i) * string(j)
            println(fname)
            data[fname] = [(alldates[i][k] - alldates[j][k]).value for k in 1:length(alldates[i])]
        end
    end

    for fname in features
        dates = [Date(date) for date in data[fname]]
        data[fname] = [[Dates.year(date),
                        Dates.month(date),
                        Dates.day(date),
                        Dates.dayofweek(date)] for date in dates]
    end
end

function process_timestamp(m::Model, data)
    features = [fname for (fname, ftype) in getfeatures(m; ftype="all")
                if in(fname, keys(data)) && ftype == TIMESTAMP]

    for fname in features
        data[fname] = fill_timestamp(data[fname])
        timestamps = [DateTime(ts) for ts in data[fname]]
        data[fname] = [[Dates.year(ts),
                        Dates.month(ts),
                        Dates.day(ts),
                        Dates.hour(ts),
                        Dates.minute(ts),
                        Dates.second(ts)] for ts in timestamps]
    end
end

function process_text(m::Model, data)
    features = [fname for (fname, ftype) in getfeatures(m; ftype="all")
                if in(fname, keys(data)) && ftype == TEXT]

    if istextmodel(m) && !haskey(m.extractor, "bert")
        m.extractor["bert"] = PretrainedBert()

    end

    m.vocabulary == nothing && initializevocab(BERTVOCPATH)
    bert = m.extractor["bert"]

    println("Bert...")

    for fname in features
        data = fill_text(data, fname)
        data = remove_invalid(data, fname)
        docs = read_and_process(data[fname], m.vocabulary)
        inputids, masks, segmentids = preprocessbert(docs)

        data[fname] = [mean(bert.bert(
            inputids[k],
            segmentids[k];
            attention_mask=masks[k])[:, :, end], dims=2)
                       for k in 1:length(inputids)]
    end
end

function process_image(m::Model, data)
    features = [fname for (fname, ftype) in getfeatures(m; ftype="all")
                if in(fname, keys(data)) && ftype == IMAGE]

    if isimagemodel(m) && !haskey(m.extractor, "resnet")
        m.extractor["resnet"] = ResNet()
    end

    resnet = m.extractor["resnet"]

    for fname in features
        data = fill_image(data, fname)
        println("Resnet...")
        data[fname] = [resnet(joinpath(m.datapath, path)) for path in data[fname]]
    end
end

function preprocess(m::Model, data::Dict{String, Array{T,1} where T})
    preprocessed = copy(data)
    preprocessed = Dict{String, Array{T,1} where T}(
        fname => value for (fname, value) in preprocessed
        if in(fname, getfnames(m, ftype="all")))

    process_string(m, preprocessed)
    process_int(m, preprocessed)
    process_float(m, preprocessed)
    process_bin_category(m, preprocessed)
    process_category(m, preprocessed)
    process_date(m, preprocessed)

    preprocessed
end

function preprocess2(m::Model, data::Dict{String, Array{T,1} where T})
    preprocessed = copy(data)

    istextmodel(m) && process_text(m, preprocessed)
    isimagemodel(m) && process_image(m, preprocessed)

    preprocessed
end

function preparedata(m::Model, traindata; output=true)
    atype = gpu() >= 0 ? KnetArray{Float64} : Array{Float64}
    trn = preprocess2(m, traindata)

    xtrn = Dict(fname => value for (fname, value) in trn
                if !in(fname, getfnames(m; ftype="output")))
    xtrn = vcat([atype(hcat(value...)) for (fname, value) in xtrn]...)

    if output
        ytrn = Dict(fname => trn[fname]
                    for fname in getfnames(m; ftype="output"))
        ytrn = vcat([atype(hcat(value...)) for (fname, value) in ytrn]...)
        ytrn = iscategorical(m) ? Array{Int64}(ytrn) : ytrn

        return xtrn, ytrn
    end

    return xtrn
end

function getbatches(m::Model, traindata::Dict{String, Array{T,1} where T};
                    n=1000, batchsize=32, showtime=true)
    batches = []
    trn = preprocess(m, traindata)
    numexamples = length(iterate(values(trn))[1])

    for i=1:n:numexamples
        j = (i+n-1) < numexamples ? i+n-1 : numexamples
        dict = typeof(traindata)(fname => values[i:j] for (fname, values) in trn)
        x,y = showtime ? @time(preparedata(m, dict)) : preparedata(m, dict)
        d = minibatch(x,y,batchsize;shuffle=false)
        push!(batches, d)
    end

    dx = [d.x for d in batches]
    dy = [d.y for d in batches]

    dfinal = minibatch(hcat(dx...), hcat(dy...), batchsize; shuffle=true)
end

function getbatches(m::Model, df::DataFrames.DataFrame; args...)
    traindata = csv2data(df)
    getbatches(m, traindata; args...)
end

function getbatches(m::Model, trainpath::String; args...)
    traindata = csv2data(trainpath)
    getbatches(m, traindata; args...)
end

max_stars_repo_name: alperkesen/AutoML

max_issues_repo_name: alperkesen/AutoML


text:
abstract type Contact{T} <: Bound{T} end

###Constraints and Derivatives
##Position level constraint wrappers
g(contact::Contact, body::Body, Δt) = g(contact, body.state, Δt)
@inline g(contact::Contact{T}) where {T} = szeros(T,6)

##Descrete Time Position Derivatives (for dynamics)
#Wrappers 1
@inline function ∂g∂ʳpos(contact::Contact, body::Body)
    return ∂g∂ʳpos(contact, body.state)
end

#Wrappers 2
@inline ∂g∂ʳpos(contact::Contact{T}) where T = szeros(T, 6)
∂g∂ʳpos(contact::Contact, state::State) = ∂g∂ʳpos(contact, posargsk(state)...)

#Derivatives accounting for quaternion specialness
@inline function ∂g∂ʳpos(contact::Contact, x::AbstractVector, q::UnitQuaternion)
    X, Q = ∂g∂pos(contact, x, q)
    Q = Q * LVᵀmat(q)

    return [X Q]
end

##Descrete Time Velocity Derivatives
#Wrappers 1
@inline function ∂g∂ʳvel(contact::Contact, body::Body, Δt)
    return ∂g∂ʳvel(contact, body.state, Δt)
end

#Wrappers 2
@inline ∂g∂ʳvel(contact::Contact{T}) where {T} = szeros(T, 6)
∂g∂ʳvel(contact::Contact, state::State, Δt) = ∂g∂ʳvel(contact, posargsnext(state, Δt)..., fullargssol(state)..., Δt)

#Derivatives sccounting for quaternion specialness
@inline function ∂g∂ʳvel(contact::Contact,x2::AbstractVector, q2::UnitQuaternion, x1::AbstractVector,v1::AbstractVector, q1::UnitQuaternion,w1::AbstractVector, Δt)
    X, Q = ∂g∂pos(contact, x2, q2)
    V = X * Δt
    Ω = Q * Lmat(q1) * derivωbar(w1, Δt)
    return [V Ω]
end

#ShurF
@inline function schurf(ineqc, contact::Contact, i, body::Body, μ, Δt)
    return schurf(contact, body.state, ineqc.γsol[2][i], ineqc.ssol[2][i], μ, Δt)[1:6]
end

#ShurD
@inline function schurD(ineqc, contact::Contact, i, body::Body, Δt)
    return schurD(contact, body.state, ineqc.γsol[2][i], ineqc.ssol[2][i], Δt)
end


text:
# the code that will run on a worker process

function run_worker(p::ParallelPOMCPOWPlanner, b0, acts::AbstractVector, input::RemoteChannel, output::RemoteChannel)
    na = length(acts)
    try
        pp = p.powplanner
        srand(pp, rand(UInt32))
        m = pp.problem
        rng = pp.solver.rng
        c = pp.solver.criterion.c

        trees = [[new_tree(pp, b0, a, p.batch_size)] for a in acts]
        n_sims = @MVector(zeros(Int, na))
        q_sums = @MVector(zeros(Float64, na))

        while true
            # don't go unless at least one input has been given
            up = take!(input)
            # get latest update
            while isready(input)
                up = take!(input)
            end
            if up.terminate
                break
            end
            global_ns = convert(MVector, up.global_ns)
            global_qs = convert(MVector, up.global_qs)
            if !isnull(up.new_branch)
                nbi = get(up.new_branch)
                push!(trees[nbi], new_tree(pp, b0, acts[nbi], p.batch_size))
            end

            # run a batch
            for k in 1:p.batch_size
                # choose which to run on
                ltn = log(sum(global_ns))
                ai = indmax(begin
                    q = global_qs[i]
                    n = global_ns[i]
                    ucb(q, c, n, ltn)
                end for i in 1:na)

                # run sims for the number of trees
                for i in 1:length(trees[ai])
                    # choose the tree with the least sims
                    t = trees[ai][indmin(t.total_n[1] for t in trees[ai])]
                    s, r = rand(rng, t.root_belief)
                    if isterminal(m, s)
                        q = r
                    else
                        node = POWTreeObsNode(t, 1)
                        q = r + discount(m)*POMCPOW.simulate(pp, node, s, pp.solver.max_depth-1)
                    end
                    n_sims[ai] += 1
                    q_sums[ai] += q
                    global_ns[ai] += 1
                    global_qs[ai] += (q-global_qs[ai])/global_ns[ai]
                end
            end

            # report results
            put!(output, WorkReport(myid(), false,
                                    SVector{na, Int}([length(tv) for tv in trees]...),
                                    convert(SVector, n_sims),
                                    convert(SVector, q_sums./n_sims)))
        end
    catch ex
        println("caught")
        showerror(STDERR, ex)
        put!(output, WorkReport(myid(), true,
                                @SVector(zeros(Int, na)),
                                @SVector(zeros(Int, na)),
                                @SVector(zeros(Float64, na))))
    end
end

function new_tree(pp::POMCPOWPlanner, b, a, sz)
    m = pp.problem
    ps = pp.solver
    s = rand(ps.rng, b)
    sp, o, r = generate_sor(m, s, a, ps.rng)
    nb = POMCPOW.init_node_sr_belief(pp.node_sr_belief_updater,
                                     m, s, a, sp, o, r
                                    )
    rb = FilteredSRBelief(b, a, pp.node_sr_belief_updater, nb)
    return POMCPOWTree{typeof(nb), typeof(a), typeof(o), typeof(rb)}(rb, sz)
end

function ucb(q, c, n, ltn)
    if n == 0 && ltn <= 0.0
        return q
    elseif n == 0 && q == -Inf
        return Inf
    else
        return q + c*sqrt(ltn/n)
    end
end


text:
using Test
using MIPVerify
using MIPVerify: UnrestrictedPerturbationFamily, BlurringPerturbationFamily, LInfNormBoundedPerturbationFamily
@isdefined(TestHelpers) || include("../../../TestHelpers.jl")

@timed_testset "conv+fc+softmax.jl" begin
    ### Parameters for neural net
    batch = 1
    c1_in_height = 6
    c1_in_width = 6
    p1_stride_height = 2
    p1_stride_width = 2
    p1_strides = (1, p1_stride_height, p1_stride_width, 1)
    p1_height = round(Int, c1_in_height/p1_stride_height, RoundUp)
    p1_width = round(Int, c1_in_width/p1_stride_width, RoundUp)
    c1_in_channels = 1
    c1_filter_height = 2
    c1_filter_width = 2
    c1_out_channels = 2

    l1_height = 4
    l1_width = p1_height*p1_width*c1_out_channels

    l2_height = 3
    l2_width = l1_height

    ### Choosing data to be used
    input = gen_array(
        (batch, c1_in_height, c1_in_width, c1_in_channels),
        0, 1
    )

    c1_kernel = gen_array(
        (c1_filter_height, c1_filter_width, c1_in_channels, c1_out_channels),
        -1, 1
    )
    c1_bias = gen_array((c1_out_channels,), -1, 1)

    l1_kernel = gen_array((l1_width, l1_height), -1, 1)
    l1_bias = gen_array((l1_height,), -1, 1)

    l2_kernel = gen_array((l2_width, l2_height), -1, 1)
    l2_bias = gen_array((l2_height,), -1, 1)

    nn = Sequential(
        [
            Conv2d(c1_kernel, c1_bias), MaxPool(p1_strides), ReLU(),
            Flatten(4),
            Linear(l1_kernel, l1_bias), ReLU(),
            Linear(l2_kernel, l2_bias)
        ],
        "tests.integration.generated_weights.conv+fc+softmax"
    )

    @timed_testset "BlurringPerturbationFamily" begin
        pp_blur = BlurringPerturbationFamily((5, 5))
        # TODO (vtjeng): Add example where blurring perturbation generates non-NaN results
        test_cases = [
            ((2, pp_blur, 1, 0), 6.959316),
            ((3, pp_blur, 1, 0), NaN),
        ]

        TestHelpers.batch_test_adversarial_example(nn, input, test_cases)
    end

    @timed_testset "UnrestrictedPerturbationFamily" begin
        pp_unrestricted = UnrestrictedPerturbationFamily()
        @timed_testset "Minimizing l1 norm" begin
            test_cases = [
                ((1, pp_unrestricted, 1, 0), 0),
                ((2, pp_unrestricted, 1, 0), 0.9187638),
            ]

            TestHelpers.batch_test_adversarial_example(nn, input, test_cases)
        end

        @timed_testset "Minimizing lInf norm" begin
            test_cases = [
                ((1, pp_unrestricted, Inf, 0), 0),
                ((2, pp_unrestricted, Inf, 0), 0.06688736),
                ((3, pp_unrestricted, Inf, 0), 0.4270584),
            ]

            TestHelpers.batch_test_adversarial_example(nn, input, test_cases)
        end

        @timed_testset "Increasing margin increases required distance" begin
            test_cases = [
                ((2, pp_unrestricted, 1, 0.1), 0.98717927),
            ]

            TestHelpers.batch_test_adversarial_example(nn, input, test_cases)
        end

        @timed_testset "With multiple target labels specified, minimum target label found" begin
            test_cases = [
                (([2, 3], pp_unrestricted, Inf, 0), 0.06688736),
            ]

            TestHelpers.batch_test_adversarial_example(nn, input, test_cases)
        end
    end

    @timed_testset "LInfNormBoundedPerturbationFamily" begin
        test_cases = [
            ((2, LInfNormBoundedPerturbationFamily(0.06), Inf, 0), NaN),  # restricting maximum perturbation to below minimum distance causes optimization problem to be infeasible
            ((2, LInfNormBoundedPerturbationFamily(0.07), Inf, 0), 0.06688736),  # restricting maximum perturbation to above minimum distance does not affect optimal value of problem
        ]

        TestHelpers.batch_test_adversarial_example(nn, input, test_cases)
    end
end


text:
struct RangeConstraintSpec
    constraint_name::Symbol
    variable_name::Symbol
    bin_variable_names::Vector{Symbol}
    limits_func::Function
    constraint_func::Function
    constraint_struct::Type{<:AbstractRangeConstraintInfo}
    lag_limits_func::Union{Function, Nothing}
end

function RangeConstraintSpec(;
    constraint_name,
    variable_name,
    bin_variable_names = Vector{Symbol}(),
    limits_func,
    constraint_func,
    constraint_struct,
    lag_limits_func = nothing,
)
    return RangeConstraintSpec(
        constraint_name,
        variable_name,
        bin_variable_names,
        limits_func,
        constraint_func,
        constraint_struct,
        lag_limits_func,
    )
end

struct TimeSeriesConstraintSpec
    constraint_name::Symbol
    variable_name::Symbol
    bin_variable_name::Union{Nothing, Symbol}
    parameter_name::Union{Nothing, String}
    forecast_label::Union{Nothing, String}
    multiplier_func::Union{Nothing, Function}
    constraint_func::Function
end

function TimeSeriesConstraintSpec(;
    constraint_name,
    variable_name,
    bin_variable_name = nothing,
    parameter_name,
    forecast_label,
    multiplier_func,
    constraint_func,
)
    return TimeSeriesConstraintSpec(
        constraint_name,
        variable_name,
        bin_variable_name,
        parameter_name,
        forecast_label,
        multiplier_func,
        constraint_func,
    )
end

struct DeviceRangeConstraintSpec
    range_constraint_spec::Union{Nothing, RangeConstraintSpec}
    timeseries_range_constraint_spec::Union{Nothing, TimeSeriesConstraintSpec}
    custom_psi_container_func::Union{Nothing, Function}
    devices_filter_func::Union{Nothing, Function}
end

"""
Default implementation to add range constraints.

Users of this function must implement a method for
[`DeviceRangeConstraintSpec`](@ref) for their specific types.
Users may also implement custom active_power_constraints! methods.
"""
function add_constraints!(
    psi_container::PSIContainer,
    ::Type{T},
    ::Type{U},
    devices::IS.FlattenIteratorWrapper{V},
    model::DeviceModel{V, W},
    ::Type{X},
    feedforward::Union{Nothing, AbstractAffectFeedForward},
) where {
    T <: RangeConstraint,
    U <: VariableType,
    V <: PSY.Device,
    W <: AbstractDeviceFormulation,
    X <: PM.AbstractPowerModel,
}
    use_parameters = model_has_parameters(psi_container)
    use_forecasts = model_uses_forecasts(psi_container)
    @assert !(use_parameters && !use_forecasts)
    spec =
        DeviceRangeConstraintSpec(T, U, V, W, X, feedforward, use_parameters, use_forecasts)
    device_range_constraints!(psi_container, devices, model, feedforward, spec)
end

"""
Construct inputs for creating range constraints.

# Arguments
`range_constraint_spec::Vector{RangeConstraintSpec}`: May be emtpy.
`timeseries_range_constraint_spec::Vector{TimeSeriesConstraintSpec}`: May be empty.
`custom_psi_container_func::Union{Nothing, Function}`: Optional function to add custom
 constraints to the internals of a PSIContainer. Must accept PSIContainer, devices iterable,
 and a subtype of AbstractDeviceFormulation.
`devices_filter_func::Union{Nothing, Function}`: Optional function to filter the devices on

"""
function DeviceRangeConstraintSpec(;
    range_constraint_spec = nothing,
    timeseries_range_constraint_spec = nothing,
    custom_psi_container_func = nothing,
    devices_filter_func = nothing,
)
    return DeviceRangeConstraintSpec(
        range_constraint_spec,
        timeseries_range_constraint_spec,
        custom_psi_container_func,
        devices_filter_func,
    )
end

function device_range_constraints!(
    psi_container::PSIContainer,
    devices::IS.FlattenIteratorWrapper{T},
    model::DeviceModel{T, U},
    feedforward::Union{Nothing, AbstractAffectFeedForward},
    spec::DeviceRangeConstraintSpec,
) where {T <: PSY.Device, U <: AbstractDeviceFormulation}
    range_constraint_spec = spec.range_constraint_spec
    timeseries_range_constraint_spec = spec.timeseries_range_constraint_spec
    custom_psi_container_func = spec.custom_psi_container_func

    if !isnothing(spec.devices_filter_func)
        devices = filter!(spec.devices_filter_func, collect(devices))
    end

    if isnothing(feedforward)
        ff_affected_variables = Set{Symbol}()
    else
        ff_affected_variables = Set(get_affected_variables(feedforward))
    end

    if !isnothing(range_constraint_spec)
        _apply_range_constraint_spec!(
            psi_container,
            range_constraint_spec,
            devices,
            model,
            ff_affected_variables,
        )
    end

    if !isnothing(timeseries_range_constraint_spec)
        _apply_timeseries_range_constraint_spec!(
            psi_container,
            timeseries_range_constraint_spec,
            devices,
            model,
            ff_affected_variables,
        )
    end

    if !isnothing(custom_psi_container_func)
        custom_psi_container_func(psi_container, devices, U)
    end
end

function _apply_range_constraint_spec!(
    psi_container,
    spec,
    devices::IS.FlattenIteratorWrapper{T},
    model,
    ff_affected_variables,
) where {T <: PSY.Device}
    constraint_struct = spec.constraint_struct
    constraint_infos = Vector{constraint_struct}(undef, length(devices))
    constraint_name = spec.constraint_name
    variable_name = spec.variable_name
    if variable_name in ff_affected_variables
        @debug "Skip adding $variable_name because it is handled by feedforward"
        return
    end
    bin_var_name = spec.bin_variable_names
    for (i, dev) in enumerate(devices)
        dev_name = PSY.get_name(dev)
        limits = spec.limits_func(dev)
        if isnothing(limits)
            limits = (min = 0.0, max = 0.0)
            @warn "Range constraint limits of $T $dev_name are nothing. Set to" limits
        end
        if constraint_struct == DeviceRangeConstraintInfo
            constraint_info = DeviceRangeConstraintInfo(dev_name, limits)
        elseif constraint_struct == DeviceMultiStartRangeConstraintsInfo
            lag_limits = spec.lag_limits_func(dev)
            constraint_info =
                DeviceMultiStartRangeConstraintsInfo(dev_name, limits, lag_limits)
        else
            error("Missing implementation for $constraint_struct")
        end
        add_device_services!(constraint_info, dev, model)
        constraint_infos[i] = constraint_info
    end

    spec.constraint_func(
        psi_container,
        RangeConstraintSpecInternal(
            constraint_infos,
            constraint_name,
            variable_name,
            bin_var_name,
        ),
    )
    return
end

function _apply_timeseries_range_constraint_spec!(
    psi_container,
    spec,
    devices::IS.FlattenIteratorWrapper{T},
    model,
    ff_affected_variables,
) where {T <: PSY.Device}
    variable_name = spec.variable_name
    if variable_name in ff_affected_variables
        @debug "Skip adding $variable_name because it is handled by feedforward"
        return
    end
    constraint_infos = Vector{DeviceTimeSeriesConstraintInfo}(undef, length(devices))
    for (i, dev) in enumerate(devices)
        ts_vector = get_time_series(psi_container, dev, spec.forecast_label)
        constraint_info =
            DeviceTimeSeriesConstraintInfo(dev, spec.multiplier_func, ts_vector)
        add_device_services!(constraint_info.range, dev, model)
        constraint_infos[i] = constraint_info
    end

    ts_inputs = TimeSeriesConstraintSpecInternal(
        constraint_infos,
        spec.constraint_name,
        variable_name,
        spec.bin_variable_name,
        isnothing(spec.parameter_name) ? nothing :
            UpdateRef{T}(spec.parameter_name, spec.forecast_label),
    )
    spec.constraint_func(psi_container, ts_inputs)
    return
end


text:
# ---
# title: 524. Longest Word in Dictionary through Deleting
# id: problem524
# author: Indigo
# date: 2021-06-22
# difficulty: Medium
# categories: Two Pointers, Sort
# link: <https://leetcode.com/problems/longest-word-in-dictionary-through-deleting/description/>
# hidden: true
# ---
#
# Given a string and a string dictionary, find the longest string in the
# dictionary that can be formed by deleting some characters of the given string.
# If there are more than one possible results, return the longest word with the
# smallest lexicographical order. If there is no possible result, return the
# empty string.
#
# **Example 1:**
#
#
#
#     Input:
#     s = "abpcplea", d = ["ale","apple","monkey","plea"]
#
#     Output:
#     "apple"
#
#
# **Example 2:**
#
#
#
#     Input:
#     s = "abpcplea", d = ["a","b","c"]
#
#     Output:
#     "a"
#
#
# **Note:**
#
#   1. All the strings in the input will only contain lower-case letters.
#   2. The size of the dictionary won't exceed 1,000.
#   3. The length of all the strings in the input won't exceed 1,000.
#
#
## @lc code=start
using LeetCode

function find_longest_word(s::String, dictionary::Vector{String})
    function is_subsequence(s::String, t::String)::Bool
        i, j = 1, 1
        lens, lent = length(s), length(t)
        while j <= lent
            if s[i] == t[j]
                i += 1
                i > lens && return true
            end
            j += 1
        end
        return false
    end
    res = ""
    for str in dictionary
        (is_subsequence(str, s) && (length(res) < length(str) || res > str)) && (res = str)
    end
    return res
end
## @lc code=end


text:
@testset "786.k-th-smallest-prime-fraction.jl" begin
    res = [1//5, 1//3, 2//5, 1//2, 3//5, 2//3]
    arr = [1,2,3,5]
    for i in 1:6
        @test kth_smallest_prime_fraction(arr, i) == res[i]
    end
    @test kth_smallest_prime_fraction([1, 7], 1) == 1//7
    @test kth_smallest_prime_fraction([1, 2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47], 42) == 7//37
    @test kth_smallest_prime_fraction([1, 2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47], 24) == 3//29
    @test kth_smallest_prime_fraction([1, 2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47], 13) == 1//17
end


text:
@testset "Envelopes" begin
    m = load_model(model_paths["e_coli_core.xml"])

    rxns = [1, 2, 3]

    lat = collect.(envelope_lattice(m, rxns; samples = 3))
    @test lat == [[0, 500, 1000], [-1000, 0, 1000], [-1000, 0, 1000]]
    @test lat == collect.(envelope_lattice(m, reactions(m)[rxns]; samples = 3))

    vals =
        objective_envelope(
            m,
            reactions(m)[rxns],
            Tulip.Optimizer;
            lattice = lat,
            workers = W,
        ).values

    @test size(vals) == (3, 3, 3)
    @test count(isnothing, vals) == 15
    @test isapprox(
        filter(!isnothing, vals),
        [
            0.0,
            0.0,
            0.0,
            0.704036947,
            10.053282384,
            10.053282365,
            0.0,
            0.0,
            0.0,
            0.873921506,
            18.664485767,
            20.412058183,
        ],
        atol = TEST_TOLERANCE,
    )
end


text:
export PlayerOptions,
    CheckRaiseFold,
    CallRaiseFold,
    CallAllInFold,
    CallFold

abstract type PlayerOptions end
struct CheckRaiseFold <: PlayerOptions end
struct CallRaiseFold <: PlayerOptions end
struct CallAllInFold <: PlayerOptions end
struct CallFold <: PlayerOptions end

validate_action(::Fold, ::CheckRaiseFold) = nothing
validate_action(::Fold, ::CallRaiseFold) = nothing
validate_action(::Fold, ::CallAllInFold) = nothing
validate_action(::Fold, ::CallFold) = nothing

validate_action(::Check, ::CheckRaiseFold) = nothing
validate_action(::Check, ::CallRaiseFold) = error("Cannot check. Available options: CallRaiseFold")
validate_action(::Check, ::CallAllInFold) = error("Cannot check. Available options: CallAllInFold")
validate_action(::Check, ::CallFold) = error("Cannot check. Available options: CallFold")

# We catch this error before `call!` completes, so we don't need to catch here.
# validate_action(::Call, ::CheckRaiseFold) = error("Cannot call. Available options: CheckRaiseFold")
validate_action(::Call, ::CallRaiseFold) = nothing
validate_action(::Call, ::CallAllInFold) = nothing
validate_action(::Call, ::CallFold) = nothing

validate_action(::Raise, ::CheckRaiseFold) = nothing
validate_action(::Raise, ::CallRaiseFold) = nothing
validate_action(::Raise, ::CallAllInFold) = nothing
validate_action(::Raise, ::CallFold) = error("Cannot Raise. Available options: CallFold")

function player_option!(game::Game, player::Player)
    table = game.table
    call_amt = call_amount(table, player)
    if !(call_amt ≈ 0) # must call to stay in
        cond_1 = bank_roll(player) > call_amt
        cond_2 = an_opponent_can_call_a_raise(table, player)
        raise_possible = cond_1 && cond_2
        if raise_possible # raise possible
            vrb = valid_raise_bounds(table, player)
            if first(vrb) ≈ last(vrb) # only all-in raise possible
                option = CallAllInFold()
            else
                option = CallRaiseFold()
            end
        else # only all-in possible
            option = CallFold()
        end
    else
        option = CheckRaiseFold()
    end
    @debug "option = $option"

    n_actions = length(action_history(player))

    player_option!(game, player, option)

    # exactly 1 action must be taken
    player_action = action_history(player)
    @assert n_actions + 1 == length(player_action) "Must take exactly 1 action."
    validate_action(last(player_action), option)
end

# By default, forward to `player_option!` with
# game state:
player_option!(game::Game, player::Player, option) =
    player_option!(game, player, state(game.table), option)

#####
##### Human player options (ask via prompts)
#####

function player_option!(game::Game, player::Player{Human}, ::CheckRaiseFold, io::IO=stdin)
    table = game.table
    vrb = valid_raise_bounds(table, player)
    options = ["Check", "Raise [\$$(minimum(vrb)), \$$(maximum(vrb))]", "Fold"]
    menu = RadioMenu(options, pagesize=4)
    choice = request("$(name(player))'s turn to act:", menu)
    choice == -1 && error("Uncaught case")
    choice == 1 && check!(game, player)
    choice == 2 && raise_to!(game, player, input_raise_amt(table, player, io))
    choice == 3 && fold!(game, player)
end
function player_option!(game::Game, player::Player{Human}, ::CallRaiseFold, io::IO=stdin)
    table = game.table
    vrb = valid_raise_bounds(table, player)
    call_amt = call_amount(table, player)
    blind_str = is_blind_call(table, player) ? " (blind)" : ""
    options = ["Call \$$(call_amt)$blind_str", "Raise [\$$(minimum(vrb)), \$$(maximum(vrb))]", "Fold"]
    menu = RadioMenu(options, pagesize=4)
    choice = request("$(name(player))'s turn to act:", menu)
    choice == -1 && error("Uncaught case")
    choice == 1 && call!(game, player)
    choice == 2 && raise_to!(game, player, input_raise_amt(game.table, player, io))
    choice == 3 && fold!(game, player)
end
function player_option!(game::Game, player::Player{Human}, ::CallAllInFold)
    table = game.table
    call_amt = call_amount(table, player)
    all_in_amt = round_bank_roll(player)
    blind_str = is_blind_call(table, player) ? " (blind)" : ""
    options = ["Call \$$(call_amt)$blind_str", "Raise all-in (\$$(all_in_amt))", "Fold"]
    menu = RadioMenu(options, pagesize=4)
    choice = request("$(name(player))'s turn to act:", menu)
    choice == -1 && error("Uncaught case")
    choice == 1 && call!(game, player)
    choice == 2 && raise_all_in!(game, player)
    choice == 3 && fold!(game, player)
end
function player_option!(game::Game, player::Player{Human}, ::CallFold)
    table = game.table
    call_amt = call_amount(table, player)
    blind_str = is_blind_call(table, player) ? " (blind)" : ""
    options = ["Call \$$(call_amt)$blind_str", "Fold"]
    menu = RadioMenu(options, pagesize=4)
    choice = request("$(name(player))'s turn to act:", menu)
    choice == -1 && error("Uncaught case")
    choice == 1 && call!(game, player)
    choice == 2 && fold!(game, player)
end

# io only works for tests, but does not for user input
# so we have a switch for the test suite
use_input_io() = false
println_io(io::IO, msg) = use_input_io() ? println(io, msg) : println(msg)

function input_raise_amt(table, player::Player{Human}, io::IO=stdin)
    raise_amt = nothing
    while true
        println_io(io, "Enter raise amt:")
        raise_amt = readline(io)
        try
            raise_amt = parse(Float64, raise_amt)
            is_valid, msg = is_valid_raise_amount(table, player, raise_amt)
            is_valid && break
            println_io(io, msg)
        catch
            println_io(io, "Raise must be a Float64")
        end
    end
    @assert raise_amt ≠ nothing
    amt = valid_raise_amount(table, player, raise_amt)
    return amt
end

#####
##### AbstractAI
#####

##### Bot5050

function player_option!(game::Game, player::Player{Bot5050}, ::CheckRaiseFold)
    if rand() < 0.5
        check!(game, player)
    else
        amt = Int(round(rand()*bank_roll(player), digits=0))
        amt = bound_raise(game.table, player, amt)
        raise_to!(game, player, amt)
    end
end
function player_option!(game::Game, player::Player{Bot5050}, ::CallRaiseFold)
    if rand() < 0.5
        if rand() < 0.5 # Call
            call!(game, player)
        else # re-raise
            amt = Int(round(rand()*bank_roll(player), digits=0))
            amt = bound_raise(game.table, player, amt)
            raise_to!(game, player, amt)
        end
    else
        fold!(game, player)
    end
end
function player_option!(game::Game, player::Player{Bot5050}, ::CallAllInFold)
    if rand() < 0.5
        if rand() < 0.5 # Call
            call!(game, player)
        else # re-raise
            raise_all_in!(game, player)
        end
    else
        fold!(game, player)
    end
end
function player_option!(game::Game, player::Player{Bot5050}, ::CallFold)
    if rand() < 0.5
        call!(game, player)
    else
        fold!(game, player)
    end
end


text:
# _stat_add_sum(x::T, y::S) where T where S = convert(promote_type(S,T), x + y)
_stat_add_sum(x, y) = Base.add_sum(x, y)
_stat_add_sum(x::Bool, y::Bool) = x + y
_stat_add_sum(x::Missing, y::Bool) = Int(y)
_stat_add_sum(x::Bool, y::Missing) = Int(x)
_stat_add_sum(x, ::Missing) = x
_stat_add_sum(::Missing, x) = x
_stat_add_sum(::Missing, ::Missing) = missing
_stat_mul_prod(x, y) = Base.mul_prod(x, y)
_stat_mul_prod(x, ::Missing) = x
_stat_mul_prod(::Missing, x) = x
_stat_mul_prod(::Missing, ::Missing) = missing
_stat_min_fun(x, y) = min(x, y)
_stat_min_fun(x, ::Missing) = x
_stat_min_fun(::Missing, y) = y
_stat_min_fun(::Missing, ::Missing) = missing
_stat_max_fun(x, y) = max(x, y)
_stat_max_fun(x, ::Missing) = x
_stat_max_fun(::Missing, y) = y
_stat_max_fun(::Missing, ::Missing) = missing
_stat_realXcY(x, y) = Statistics.realXcY(x, y)
_stat_realXcY(x, ::Missing) = x
_stat_realXcY(::Missing, y) = y
_stat_realXcY(::Missing,::Missing) = missing
ISNAN(x::Any) = isnan(x)
ISNAN(::Missing) = false

_stat_bool(f) = x->f(x)::Bool

_stat_ismissing(x::Any)::Int = 0
_stat_ismissing(::Missing)::Int = 1
_stat_notmissing(x::Any)::Int = 1
_stat_notmissing(::Missing)::Int = 0

"""
    lag(x, k; default = missing)
    lag!(x, k; default = missing)

Create a lag-k of the provided vector `x`. The output will be a vector the
same size as `x` (the input array).

`lag!` replace the input vector.

"""
(lag, lag!)

function lag(x::AbstractVector, k; default = missing)
    @assert firstindex(x) == 1 "lag only supports 1-based indexing"
    res = Vector{Union{promote_type(typeof(default), eltype(x)), Missing}}(undef, length(x))
    for i in 1:k
        @inbounds res[i] = default
    end
    for i in (k+1):length(x)
        @inbounds res[i] = x[i-k]
    end
    res
end

lag(x::AbstractVector; default = missing) = lag(x,1; default = default)

function lag!(x::AbstractVector, k; default = missing)
    @assert firstindex(x) == 1 "lag! only supports 1-based indexing"
    @assert promote_type(typeof(default), eltype(x)) <: eltype(x) "`default` must be the same type as the element of the passed vector"
    for i in length(x):-1:(k+1)
        @inbounds x[i] = x[i-k]
    end
    for i in 1:k
        @inbounds x[i] = default
    end
    x
end
lag!(x::AbstractVector; default = missing) = lag!(x,1; default = default)

"""
    lead(x, k; default = missing)
    lead!(x, k; default = missing)

Create a lead-k of the provided vector `x`. The output will be a vector the
same size as `x` (the input array).

`lead!` replace the input vector.

"""
(lead, lead!)

function lead(x::AbstractVector, k; default = missing)
    @assert firstindex(x) == 1 "lead only supports 1-based indexing"
    res = Vector{Union{promote_type(typeof(default), eltype(x)), Missing}}(undef, length(x))
    for i in 1:length(x)-k
        @inbounds res[i] = x[i+k]
    end
    for i in (length(x)-k+1):length(x)
        @inbounds res[i] = default
    end
    res
end
lead(x::AbstractVector; default = missing) = lead(x, 1; default = default)

function lead!(x::AbstractVector, k; default = missing)
    @assert firstindex(x) == 1 "lead! only supports 1-based indexing"
    @assert promote_type(typeof(default), eltype(x)) <: eltype(x) "`default` must be the same type as the element of the passed vector"
    for i in 1:length(x)-k
        @inbounds x[i] = x[i+k]
    end
    for i in (length(x)-k+1):length(x)
        @inbounds x[i] = default
    end
    x
end
lead!(x::AbstractVector; default = missing) = lead!(x, 1; default = default)

"""
    rescale(x,minx,maxx,minval,maxval)

Rescale x to run from minval and maxval, given x originaly runs from minx to maxx.
"""
function rescale(x,minx,maxx,minval,maxval)
    -(-maxx*minval+minx*maxval)/(maxx-minx)+(-minval+maxval)*x/(maxx-minx)
end
rescale(::Missing,minx,maxx,minval,maxval) = missing
rescale(x::Vector,minx,maxx,minval,maxval) = rescale.(x,minx,maxx,minval,maxval)
rescale(x,minx,maxx) = rescale(x,minx,maxx,0.0,1.0)

"""
    stdze(x)

Standardize an array. It returns missing for missing data points.
"""
function stdze(x)
    all(ismissing,x) && return x
    meandata = mean(x)
    vardata = var(x)
    (x .- meandata) ./ sqrt(vardata)
end

# this is manual simd version for max(min) function
function stat_maximum(f::typeof(identity), x::AbstractArray{T,1}; lo = 1, hi = length(x)) where T
    all(ismissing, view(x, lo:hi)) && return missing
    _dmiss(x) = ismissing(x) ? typemin(nonmissingtype(T)) : x
    Base.mapreduce_impl(_dmiss, max, x, lo, hi)
end
function stat_maximum(f::F, x::AbstractArray{T,1}; lo = 1, hi = length(x)) where {F, T}
    all(ismissing, view(x, lo:hi)) && return missing
    Base.mapreduce_impl(f, _stat_max_fun, x, lo, hi)
end
stat_maximum(x::AbstractArray{T,1}; lo = 1, hi = length(x)) where T = stat_maximum(identity, x; lo = lo, hi = hi)

function _arg_minmax_barrier(x, minmaxval, f)::Int
    @inbounds for i in 1:length(x)
        isequal(f(x[i]), minmaxval) && return i
    end
end

function stat_findmax(f, x::AbstractArray{T,1}) where T
    isempty(x) && throw(ArgumentError("input vector cannot be empty"))
    maxval = stat_maximum(f, x)
    ismissing(maxval) && return (missing, missing)
    (maxval, _arg_minmax_barrier(x, maxval, f))
end
stat_findmax(x::AbstractArray{T,1}) where T = stat_findmax(identity, x)

function stat_minimum(f::typeof(identity), x::AbstractArray{T,1}; lo = 1, hi = length(x)) where T
    all(ismissing, view(x, lo:hi)) && return missing
    @inline _dmiss(x) = ismissing(x) ? typemax(nonmissingtype(T)) : x
    Base.mapreduce_impl(_dmiss, min, x, lo, hi)
end
function stat_minimum(f::F, x::AbstractArray{T,1}; lo = 1, hi = length(x)) where {F,T}
    all(ismissing, view(x, lo:hi)) && return missing
    Base.mapreduce_impl(f, _stat_min_fun, x, lo, hi)
end
stat_minimum(x::AbstractArray{T,1}; lo = 1, hi = length(x)) where T = stat_minimum(identity, x; lo = lo, hi = hi)

function stat_findmin(f, x::AbstractArray{T,1}) where T
    isempty(x) && throw(ArgumentError("input vector cannot be empty"))
    minval = stat_minimum(f, x)
    (minval, _arg_minmax_barrier(x, minval, f))
end
stat_findmin(x::AbstractArray{T,1}) where T = stat_findmin(identity, x)

function stat_sum(f, x::AbstractArray{T,1}; lo = 1, hi = length(x)) where T <: Union{Missing, INTEGERS, FLOATS}
    all(ismissing, view(x, lo:hi)) && return f(first(x))
    _dmiss(y) = ifelse(ismissing(f(y)),  zero(T), f(y))
    Base.mapreduce_impl(_dmiss, _stat_add_sum, x, lo, hi)
end
stat_sum(x::AbstractArray{T,1}; lo = 1, hi = length(x)) where T <: Union{Missing, INTEGERS, FLOATS} = stat_sum(identity, x; lo = lo, hi = hi)

# function stat_wsum(f, x::AbstractArray{Union{T,Missing},1}, w) where T
#     all(ismissing, x) && return missing
#     _dmiss(y) = ismissing(y[1])||ismissing(y[2]) ? zero(T) : (f(y[1])*y[2])
#     mapreduce(_dmiss, _stat_add_sum, zip(x,w))
# end
# stat_wsum(x::AbstractArray{Union{T,Missing},1}, w) where T  = stat_wsum(identity, x, w)
function stat_wsum(f, x::AbstractVector{T}, w::AbstractVector) where T
    all(ismissing, x) && return missing
    _dmiss(y) = ismissing(y[1])||ismissing(y[2]) ? missing : (f(y[1])*y[2])
    mapreduce(_dmiss, _stat_add_sum, zip(x,w))
end
stat_wsum(x::AbstractVector{T}, w::AbstractVector) where T  = stat_wsum(identity, x, w)

function stat_mean(f, x::AbstractArray{T,1})::Union{Float64, Missing} where T <: Union{Missing, INTEGERS, FLOATS}
    length(x) == 1 && return f(first(x))
    sval = stat_sum(y->f(y)*1.0, x)
    n = mapreduce(!ismissing∘f, +, x)
    n == 0 ? missing : sval/n
end
stat_mean(x::AbstractArray{T,1}) where T = stat_mean(identity, x)

stat_cumsum_ignore(x::AbstractVector) = accumulate(_stat_add_sum, x)
stat_cumsum!_ignore(outx, inx::AbstractVector) = accumulate!(_stat_add_sum, outx, inx)
stat_cumprod_ignore(x::AbstractVector) = accumulate(_stat_mul_prod, x)
stat_cumprod!_ignore(outx, inx::AbstractVector) = accumulate!(_stat_mul_prod, outx, inx)
stat_cummin_ignore(x::AbstractVector) = accumulate(_stat_min_fun, x)
stat_cummin!_ignore(outx, inx::AbstractVector) = accumulate!(_stat_min_fun, outx, inx)
stat_cummax_ignore(x::AbstractVector) = accumulate(_stat_max_fun, x)
stat_cummax!_ignore(outx, inx::AbstractVector) = accumulate!(_stat_max_fun, outx, inx)

function stat_cumsum_skip(x::AbstractVector)
    locmiss = ismissing.(x)
    res = stat_cumsum_ignore(x)
    if sum(locmiss)>0
        res[locmiss] .= missing
    end
    res
end

function stat_cumsum!_skip(outx, inx::AbstractVector)
    locmiss = ismissing.(inx)
    stat_cumsum!_ignore(outx, inx)
    if sum(locmiss)>0
        outx[locmiss] .= missing
    end
    outx
end
function stat_cumprod_skip(x::AbstractVector)
    locmiss = ismissing.(x)
    res = stat_cumprod_ignore(x)
    if sum(locmiss)>0
        res[locmiss] .= missing
    end
    res
end

function stat_cumprod!_skip(outx, inx::AbstractVector)
    locmiss = ismissing.(inx)
    stat_cumprod!_ignore(outx, inx)
    if sum(locmiss)>0
        outx[locmiss] .= missing
    end
    outx
end

function stat_cummin_skip(x::AbstractVector)
    locmiss = ismissing.(x)
    res = stat_cummin_ignore(x)
    if sum(locmiss)>0
        res[locmiss] .= missing
    end
    res
end

function stat_cummin!_skip(outx, inx::AbstractVector)
    locmiss = ismissing.(inx)
    stat_cummin!_ignore(outx, inx)
    if sum(locmiss)>0
        outx[locmiss] .= missing
    end
    outx
end
function stat_cummax_skip(x::AbstractVector)
    locmiss = ismissing.(x)
    res = stat_cummax_ignore(x)
    if sum(locmiss)>0
        res[locmiss] .= missing
    end
    res
end

function stat_cummax!_skip(outx, inx::AbstractVector)
    locmiss = ismissing.(inx)
    stat_cummax!_ignore(outx, inx)
    if sum(locmiss)>0
        outx[locmiss] .= missing
    end
    outx
end

function stat_wmean(f, x::AbstractVector{T}, w::AbstractArray{S,1}) where T where S
    all(ismissing, x) && return missing
    _dmiss(y) = ismissing(y[1])||ismissing(y[2]) ? zero(T) : (f(y[1])*y[2])
    _dmiss2(y) = ismissing(y[1])||ismissing(y[2]) ? zero(S) : y[2]
    _op(y1,y2) = _stat_add_sum.(y1, y2)
    _f(y) = (_dmiss(y), _dmiss2(y))
    sval, n = mapreduce(_f, _op, zip(x,w))
    n == 0 ? missing : sval / n
end
stat_wmean(x::AbstractVector{T}, w::AbstractArray{S,1}) where T where S = stat_wmean(identity, x, w)

function stat_var(f, x::AbstractArray{T,1}, dof=true)::Union{Float64, Missing} where T <: Union{Missing, INTEGERS, FLOATS}
    all(ismissing, x) && return missing
    # any(ISNAN, x) && return convert(eltype(x), NaN)
    # meanval = stat_mean(f, x)
    # n = mapreduce(!ismissing∘f, +, x)
    sval = stat_sum(y->f(y)*1.0, x)
    n = mapreduce(!ismissing∘f, +, x)
    meanval = n == 0 ? missing : sval/n

    ss = 0.0
    for i in 1:length(x)
        ss = _stat_add_sum(ss, abs2(f(x[i]) - meanval))
    end

    if n == 0
        return missing
    elseif n == 1 && dof
        return missing
    else
        return ss / (n - Int(dof))
    end
end

stat_var(x::AbstractArray{T,1}, dof=true) where T = stat_var(identity, x, dof)

function stat_std(f , x::AbstractArray{T,1}, dof=true)::Union{Float64, Missing} where T <: Union{Missing, INTEGERS, FLOATS}
    sqrt(stat_var(f, x,dof))
end
stat_std(x::AbstractArray{T,1}, dof=true) where T = stat_std(identity, x, dof)

function stat_median(v::AbstractArray{T,1}) where T
    isempty(v) && throw(ArgumentError("median of an empty array is undefined, $(repr(v))"))
    all(ismissing, v) && return missing
    (nonmissingtype(eltype(v))<:AbstractFloat || nonmissingtype(eltype(v))>:AbstractFloat) && any(ISNAN, v) && return convert(eltype(v), NaN)
    nmis::Int = mapreduce(ismissing, +, v)
    n = length(v) - nmis
    mid = div(1+n,2)
    if isodd(n)
        return middle(partialsort(v,mid))
    else
        m = partialsort(v, mid:mid+1)
        return middle(m[1], m[2])
    end
end

function stat_median!(v::AbstractArray{T,1}) where T
    isempty(v) && throw(ArgumentError("median of an empty array is undefined, $(repr(v))"))
    all(ismissing, v) && return missing
    (nonmissingtype(eltype(v))<:AbstractFloat || nonmissingtype(eltype(v))>:AbstractFloat) && any(ISNAN, v) && return convert(eltype(v), NaN)
    nmis::Int = mapreduce(ismissing, +, v)
    n = length(v) - nmis
    mid = div(1+n,2)
    if isodd(n)
        return middle(partialsort!(v,mid))
    else
        m = partialsort!(v, mid:mid+1)
        return middle(m[1], m[2])
    end
end

# finding k largest in an array with missing values
swap!(x,i,j)=x[i],x[j]=x[j],x[i]

function insert_fixed_sorted!(x, item, ord)
    if ord((x[end]), (item))
        return
    else
        x[end] = item
    end
    i = length(x) - 1
    while i > 0
        if ord((x[i+1]),(x[i]))
            swap!(x,i, i+1)
            i -= 1
        else
            break
        end
    end
end

function k_largest(x::AbstractVector{T}, k::Int) where T
    k < 1 && throw(ArgumentError("k must be greater than 1"))
    k == 1 && return [maximum(identity, x)]
    if k>length(x)
        k = length(x)
    end
    res = Vector{T}(undef,k)
    fill!(res, typemin(T))
    for i in 1:length(x)
        insert_fixed_sorted!(res, x[i], (y1,y2)-> y1 > y2)
    end
    res
end
function k_largest(x::AbstractVector{Union{T,Missing}}, k::Int) where T
    k < 1 && throw(ArgumentError("k must be greater than 1"))
    k == 1 && return [maximum(identity, x)]
    all(ismissing, x) && return [missing]
    res = Vector{T}(undef,k)
    fill!(res, typemin(T))
    cnt = 0
    for i in 1:length(x)
        if !ismissing(x[i])
            insert_fixed_sorted!(res, x[i], (y1,y2)-> y1 > y2)
            cnt += 1
        end
    end
    if cnt < k
        res[1:cnt]
    else
        res
    end
end

function k_smallest(x::AbstractVector{T}, k::Int) where T
    k < 1 && throw(ArgumentError("k must be greater than 1"))
    k == 1 && return [minimum(identity, x)]
    if k>length(x)
        k = length(x)
    end
    res = Vector{T}(undef,k)
    fill!(res, typemax(T))
    for i in 1:length(x)
        insert_fixed_sorted!(res, x[i], (y1,y2)-> y1 < y2)
    end
    res
end
function k_smallest(x::AbstractVector{Union{T,Missing}}, k::Int) where T
    k < 1 && throw(ArgumentError("k must be greater than 1"))
    k == 1 && return [minimum(identity, x)]
    all(ismissing, x) && return [missing]
    res = Vector{T}(undef,k)
    fill!(res, typemax(T))
    cnt = 0
    for i in 1:length(x)
        if !ismissing(x[i])
            insert_fixed_sorted!(res, x[i], (y1,y2)-> y1 < y2)
            cnt += 1
        end
    end
    if cnt < k
        res[1:cnt]
    else
        res
    end
end

"""
    topk(x, k; rev = false)

Return upto `k` largest nonmissing elements of `x`. When `rev = true` it returns upto `k` smallest nonmissing elements of `x`. When all elements are missing, the function returns `missing`
"""
function topk(x::AbstractVector, k::Int; rev = false)
    @assert firstindex(x) == 1 "topk only supports 1-based indexing"
    if rev
        k_smallest(x, k)
    else
        k_largest(x, k)
    end
end

"""
    ffill(x; [by = ismissing])
    ffill!(x; [by = ismissing])

Replace those elements of `x` which returns `true` when `by` is called on them with the previous element which calling `by` on it returns `false`.

`ffill!` modifies the input vector in-place
"""
(ffill, ffill!)

function ffill!(x::AbstractVector; by = ismissing)
    @assert firstindex(x) == 1 "ffill!/ffill only support 1-based indexing"
    for i in 2:length(x)
        if by(x[i])
            x[i] = x[i-1]
        end
    end
    x
end
ffill(x; by = ismissing) = ffill!(copy(x), by = by)

"""
    bfill(x; [by = ismissing])
    bfill!(x; [by = ismissing])

Replace those elements of `x` which returns `true` when `by` is called on them with the next element which calling `by` on it returns `false`.

`bfill!` modifies the input vector in-place
"""
function bfill!(x::AbstractVector; by = ismissing)
    @assert firstindex(x) == 1 "bfill!/bfill only support 1-based indexing"
    for i in length(x)-1:-1:1
        if by(x[i])
            x[i] = x[i+1]
        end
    end
    x
end
bfill(x, by = ismissing) = bfill!(copy(x), by = by)


text:
using Flux3D, Zygote, Test, Statistics

include("utils.jl")

@testset "Flux3D" begin

    @info "Testing Representation..."
    @testset "Representation" begin
        include("rep.jl")
    end

    @info "Testing Conversions..."
    @testset "Conversions" begin
        include("conversions.jl")
    end

    @info "Testing Transforms..."
    @testset "Transforms" begin
        include("transforms/pcloud_func.jl")
        include("transforms/mesh_func.jl")
        include("transforms/transforms.jl")
    end

    @info "Testing Metrics..."
    @testset "Metrics" begin
        include("metrics.jl")
    end

    @info "Testing Models..."
    @testset "Models" begin
        include("models.jl")
    end

    @info "Testing Dataset..."
    @testset "Dataset" begin
        include("datasets/custom.jl")

        # FIXME unzip throws error on GPU CI due to weird reason
        if !Flux3D.use_cuda[]
            include("datasets/modelnet.jl")
        end
    end

    @info "Testing GPU support..."
    @testset "CUDA" begin
        if Flux3D.use_cuda[]
            using CUDA
            include("cuda/rep.jl")
            include("cuda/conversions.jl")
            include("cuda/transforms.jl")
            include("cuda/metrics.jl")
        else
            @warn "CUDA unavailable, not testing GPU support"
        end
    end

end # testset Flux3D


text:
function transform(input::AbstractDict)

end


text:
# Script for creating Figure 14, 15 in scale manuscript

using IcanProj
using Statistics
using PyPlot
using NetCDF
using Dates
using JFSM2
using PyCall

function find_largest_error(file_fine, file_coarse, metric, variable)

  df_links = link_results(file_fine, file_coarse)

  data_coarse, data_aggregated, ngrids  = unify_results(file_fine, file_coarse, df_links, variable)

  bias = mean(data_coarse .- data_aggregated, dims = 1)

  bias = dropdims(bias, dims = 1)

  rmse = sqrt.(mean((data_coarse .- data_aggregated).^2, dims = 1))

  rmse = dropdims(rmse, dims = 1)

  if metric == "bias_max"

    ikeep = findall(df_links[:nc_coarse] .== argmax(bias))

  end

  if metric == "bias_min"

    ikeep = findall(df_links[:nc_coarse] .== argmin(bias))

  end

  if metric == "rmse"

    ikeep = findall(df_links[:nc_coarse] .== argmax(rmse))

  end

  df_subset = df_links[ikeep, :]

  return df_subset

end

function load_fsm_data(file_fine, file_coarse, df_subset, variable)

  time_array = ncread(file_fine, "time_array")

  time_vec = [DateTime(time_array[1, i], time_array[2, i], time_array[3, i], time_array[4, i]) for i in 1:size(time_array, 2)]

  data_fine = zeros(size(df_subset, 1), size(time_array, 2))

  for i in 1:size(df_subset, 1)

    tmp = ncread(file_fine, variable, start = [df_subset[:nc_fine][i], 1], count = [1, -1])

    data_fine[i, :] .= dropdims(tmp, dims = 1)

  end

  data_coarse = ncread(file_coarse, variable, start = [df_subset[:nc_coarse][1], 1], count = [1, -1])

  data_fine = permutedims(data_fine)

  data_coarse = permutedims(data_coarse)

  return time_vec, data_fine, data_coarse

end

function load_input_data(file_fine, file_coarse, df_subset)

  time_array = ncread(file_fine, "time_array")

  time_vec = [DateTime(time_array[1, i], time_array[2, i], time_array[3, i], time_array[4, i]) for i in 1:size(time_array, 2)]

  data_fine = zeros(size(df_subset, 1), size(time_array, 2))

  ind_nc = findall(x -> x in df_subset[:id_fine], Int.(ncread(file_fine, "id")))

  data_fine = zeros(size(df_subset, 1), size(time_array, 2))

  for i in 1:length(ind_nc)

    tmp = ncread(file_fine, variable, start = [ind_nc[i], 1], count = [1, -1])

    data_fine[i, :] .= dropdims(tmp, dims = 1)

    @show i

  end

  ind_nc = findall(x -> x in df_subset[:id_coarse], Int.(ncread(file_coarse, "id")))

  data_coarse = ncread(file_coarse, variable, start = [ind_nc[1], 1], count = [1, -1])

  data_fine = permutedims(data_fine)

  data_coarse = permutedims(data_coarse)

  return time_vec, data_fine, data_coarse

end

function results_files(cfg, variable)

  file_fine = "/data04/jmg/fsm_simulations/netcdf/fsmres_forest/results_$(cfg)/$(variable)_1km.nc"

  file_coarse = "/data04/jmg/fsm_simulations/netcdf/fsmres_forest/results_$(cfg)/$(variable)_50km.nc"

  return file_fine, file_coarse

end

function forcing_files(variable)

  file_fine = "/data04/jmg/fsm_simulations/netcdf/forcings_st/$(variable)_1km.nc"

  file_coarse = "/data04/jmg/fsm_simulations/netcdf/forcings_st/$(variable)_50km.nc"

  return file_fine, file_coarse

end

function plot_inputs()

  # Load results

  variable = "tair"

  file_fine, file_coarse = forcing_files(variable)

  tair_time, tair_fine, tair_coarse = load_input_data(file_fine, file_coarse, df_subset)

  variable = "wind"

  file_fine, file_coarse = forcing_files(variable)

  wind_time, wind_fine, wind_coarse = load_input_data(file_fine, file_coarse, df_subset)

  variable = "ilwr"

  file_fine, file_coarse = forcing_files(variable)

  ilwr_time, ilwr_fine, ilwr_coarse = load_input_data(file_fine, file_coarse, df_subset)

  variable = "iswr"

  file_fine, file_coarse = forcing_files(variable)

  iswr_time, iswr_fine, iswr_coarse = load_input_data(file_fine, file_coarse, df_subset)

  # Plot results

  xrange = (DateTime(2008, 9 , 1), DateTime(2009, 9, 1))

  fig, ax = subplots(4, 1, sharex=true)

  fig[:set_size_inches](6, 10)

  ax[1][:grid](axis = "x")
  ax[1][:fill_between](tair_time, minimum(tair_fine, dims=2)[:], maximum(tair_fine, dims=2)[:], color = "gray")
  ax[1][:plot](tair_time, tair_coarse, color = "red")
  ax[1][:set_xlim](xrange)
  ax[1][:set_ylabel]("Air temperature")
  ax[1][:axhline](linewidth = 1, color="blue")

  ax[2][:grid](axis = "x")
  ax[2][:fill_between](wind_time, minimum(wind_fine, dims=2)[:], maximum(wind_fine, dims=2)[:], color = "gray")
  ax[2][:plot](wind_time, wind_coarse, color = "red")
  ax[2][:set_xlim](xrange)
  ax[2][:set_ylabel]("Wind speed")

  ax[3][:grid](axis = "x")
  ax[3][:fill_between](ilwr_time, minimum(ilwr_fine, dims=2)[:], maximum(ilwr_fine, dims=2)[:], color = "gray")
  ax[3][:plot](ilwr_time, ilwr_coarse, color = "red")
  ax[3][:set_xlim](xrange)
  ax[3][:set_ylabel]("Incoming longwave")

  ax[4][:grid](axis = "x")
  ax[4][:fill_between](iswr_time, minimum(iswr_fine, dims=2)[:], maximum(iswr_fine, dims=2)[:], color = "gray")
  ax[4][:plot](iswr_time, iswr_coarse, color = "red")
  ax[4][:set_xlim](xrange)
  ax[4][:set_ylabel]("Incoming shortwave")

  savefig(joinpath(path_figure, "forcings.pdf"), dpi = 600)

end

function plot_fsm_heat_fluxes(df_subset, metric, path_figure)

  fig, ax = subplots(2, 1, sharex=true)

  fig[:set_size_inches](6, 5)

  variable = "hatmo"

  xrange = (DateTime(2008, 9 , 1), DateTime(2009, 9, 1))
  yrange = (-100, 50)

  # Load and plot config 30

  cfg = 30

  file_fine, file_coarse = results_files(cfg, variable)

  hatmo_time, hatmo_fine, hatmo_coarse = load_fsm_data(file_fine, file_coarse, df_subset, variable)

  ax[1][:fill_between](hatmo_time, minimum(hatmo_fine, dims=2)[:], maximum(hatmo_fine, dims=2)[:], color = "gray")
  ax[1][:plot](hatmo_time, hatmo_coarse, color = "blue")
  ax[1][:set_xlim](xrange)
  ax[1][:set_ylim](yrange)
  ax[1][:grid](axis = "x")

  ax[1][:annotate]("(A) Exchng = 0", xy = (0.05, 0.85), xycoords = "axes fraction")

  # Load and plot config 32

  cfg = 32

  file_fine, file_coarse = results_files(cfg, variable)

  hatmo_time, hatmo_fine, hatmo_coarse = load_fsm_data(file_fine, file_coarse, df_subset, variable)

  ax[2][:fill_between](hatmo_time, minimum(hatmo_fine, dims=2)[:], maximum(hatmo_fine, dims=2)[:], color = "gray")
  ax[2][:plot](hatmo_time, hatmo_coarse, color = "blue")
  ax[2][:set_xlim](xrange)
  ax[2][:set_ylim](yrange)
  ax[2][:grid](axis = "x")

  ax[2][:annotate]("(B) Exchng = 1", xy = (0.05, 0.85), xycoords = "axes fraction")

  fig[:text](0.02, 0.5, "Sensible heat flux (" * L"$W/m^2$" * ")", va="center", rotation="vertical")

  savefig(joinpath(path_figure, "sensible_heat_fluxes_$(metric).pdf"), dpi = 600)

end

function plot_fms_results(df_subset, path_figure, metric)

  # Common to all plots

  xrange = (DateTime(2010, 9 , 1), DateTime(2011, 9, 1))
  yrange = (-50, 50)

  fig, ax = subplots(4, 2, sharex=true)

  fig[:set_size_inches](12, 10)

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

  # Perform analysis for cfg = 30

  cfg = 30

  # Load results

  variable = "swe"

  file_fine, file_coarse = results_files(cfg, variable)

  swe_time, swe_fine, swe_coarse = load_fsm_data(file_fine, file_coarse, df_subset, variable)

  variable = "hatmo"

  file_fine, file_coarse = results_files(cfg, variable)

  hatmo_time, hatmo_fine, hatmo_coarse = load_fsm_data(file_fine, file_coarse, df_subset, variable)

  variable = "latmo"

  file_fine, file_coarse = results_files(cfg, variable)

  latmo_time, latmo_fine, latmo_coarse = load_fsm_data(file_fine, file_coarse, df_subset, variable)

  variable = "rnet"

  file_fine, file_coarse = results_files(cfg, variable)

  rnet_time, rnet_fine, rnet_coarse = load_fsm_data(file_fine, file_coarse, df_subset, variable)

  # Plot results

  ax[1, 1][:set_title]("Exchng = 0")

  # SWE

  swe_error = repeat(swe_coarse, 1, size(swe_fine, 2)) - swe_fine

  ax[1, 1][:grid](axis = "x")
  ax[1, 1][:fill_between](swe_time, minimum(swe_fine, dims=2)[:], maximum(swe_fine, dims=2)[:], color = "gray")
  ax[1, 1][:plot](swe_time, swe_coarse, color = "blue")
  ax[1, 1][:plot](swe_time, mean(swe_fine, dims=2)[:], color = "green")
  ax[1, 1][:set_xlim](xrange)
  ax[1, 1][:set_ylim](0, 2000)
  ax[1, 1][:set_ylabel]("SWE (" * L"$kg/m^2$" * ")")

  swe_error = mean(swe_error, dims = 2)

  ax[1, 1][:annotate]("(A) Bias = $(round(mean(swe_error), digits = 0)) | RMSE = $(round(sqrt(mean(swe_error.^2)), digits = 0))", xy = (0.05, 0.85), xycoords = "axes fraction")

  # RNET

  rnet_error = repeat(rnet_coarse, 1, size(rnet_fine, 2)) - rnet_fine

  ax[2, 1][:grid](axis = "x")
  ax[2, 1][:fill_between](rnet_time, minimum(rnet_error, dims=2)[:], maximum(rnet_error, dims=2)[:], color="gray")
  ax[2, 1][:plot](rnet_time, mean(rnet_error, dims=2), color = "red")
  ax[2, 1][:set_xlim](xrange)
  ax[2, 1][:set_ylim](yrange)
  ax[2, 1][:set_ylabel]("RNET (" * L"$W/m^2$" * ")")

  rnet_error = mean(rnet_error, dims = 2)

  ax[2, 1][:annotate]("(C) Bias = $(round(mean(rnet_error), digits = 1)) | RMSE = $(round(sqrt(mean(rnet_error.^2)), digits = 1))", xy = (0.05, 0.85), xycoords = "axes fraction")

  # HATMO

  hatmo_error = repeat(hatmo_coarse, 1, size(hatmo_fine, 2)) - hatmo_fine

  ax[3, 1][:grid](axis = "x")
  ax[3, 1][:fill_between](hatmo_time, minimum(hatmo_error, dims=2)[:], maximum(hatmo_error, dims=2)[:], color="gray")
  ax[3, 1][:plot](hatmo_time, mean(hatmo_error, dims=2), color = "red")
  ax[3, 1][:set_xlim](xrange)
  ax[3, 1][:set_ylim](yrange)
  ax[3, 1][:set_ylabel]("HATMO (" * L"$W/m^2$" * ")")

  hatmo_error = mean(hatmo_error, dims = 2)

  ax[3, 1][:annotate]("(E) Bias = $(round(mean(hatmo_error), digits = 1)) | RMSE = $(round(sqrt(mean(hatmo_error.^2)), digits = 1))", xy = (0.05, 0.85), xycoords = "axes fraction")

  # LATMO

  latmo_error = repeat(latmo_coarse, 1, size(latmo_fine, 2)) - latmo_fine

  ax[4, 1][:grid](axis = "x")
  ax[4, 1][:fill_between](latmo_time, minimum(latmo_error, dims=2)[:], maximum(latmo_error, dims=2)[:], color="gray")
  ax[4, 1][:plot](latmo_time, mean(latmo_error, dims=2), color = "red")
  ax[4, 1][:set_xlim](xrange)
  ax[4, 1][:set_ylim](yrange)
  ax[4, 1][:set_ylabel]("LATMO (" * L"$W/m^2$" * ")")

  latmo_error = mean(latmo_error, dims = 2)

  ax[4, 1][:annotate]("(G) Bias = $(round(mean(latmo_error), digits = 1)) | RMSE = $(round(sqrt(mean(latmo_error.^2)), digits = 1))", xy = (0.05, 0.85), xycoords = "axes fraction")

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

  # Perform analysis for cfg = 32

  cfg = 32

  # Load results

  variable = "swe"

  file_fine, file_coarse = results_files(cfg, variable)

  swe_time, swe_fine, swe_coarse = load_fsm_data(file_fine, file_coarse, df_subset, variable)

  variable = "hatmo"

  file_fine, file_coarse = results_files(cfg, variable)

  hatmo_time, hatmo_fine, hatmo_coarse = load_fsm_data(file_fine, file_coarse, df_subset, variable)

  variable = "latmo"

  file_fine, file_coarse = results_files(cfg, variable)

  latmo_time, latmo_fine, latmo_coarse = load_fsm_data(file_fine, file_coarse, df_subset, variable)

  variable = "rnet"

  file_fine, file_coarse = results_files(cfg, variable)

  rnet_time, rnet_fine, rnet_coarse = load_fsm_data(file_fine, file_coarse, df_subset, variable)

  # Plot results

  ax[1, 2][:set_title]("Exchng = 1")

  # SWE

  swe_error = repeat(swe_coarse, 1, size(swe_fine, 2)) - swe_fine

  ax[1, 2][:grid](axis = "x")
  ax[1, 2][:fill_between](swe_time, minimum(swe_fine, dims=2)[:], maximum(swe_fine, dims=2)[:], color = "gray")
  ax[1, 2][:plot](swe_time, swe_coarse, color = "blue")
  ax[1, 2][:plot](swe_time, mean(swe_fine, dims=2)[:], color = "green")
  ax[1, 2][:set_xlim](xrange)
  ax[1, 2][:set_ylim](0, 2000)
  ax[1, 2][:set_ylabel]("SWE (" * L"$kg/m^2$" * ")")

  swe_error = mean(swe_error, dims = 2)

  ax[1, 2][:annotate]("(B) Bias = $(round(mean(swe_error), digits = 0)) | RMSE = $(round(sqrt(mean(swe_error.^2)), digits = 0))", xy = (0.05, 0.85), xycoords = "axes fraction")

  # RNET

  rnet_error = repeat(rnet_coarse, 1, size(rnet_fine, 2)) - rnet_fine

  ax[2, 2][:grid](axis = "x")
  ax[2, 2][:fill_between](rnet_time, minimum(rnet_error, dims=2)[:], maximum(rnet_error, dims=2)[:], color="gray")
  ax[2, 2][:plot](rnet_time, mean(rnet_error, dims=2), color = "red")
  ax[2, 2][:set_xlim](xrange)
  ax[2, 2][:set_ylim](yrange)
  ax[2, 2][:set_ylabel]("RNET (" * L"$W/m^2$" * ")")

  rnet_error = mean(rnet_error, dims = 2)

  ax[2, 2][:annotate]("(D) Bias = $(round(mean(rnet_error), digits = 1)) | RMSE = $(round(sqrt(mean(rnet_error.^2)), digits = 1))", xy = (0.05, 0.85), xycoords = "axes fraction")

  # HATMO

  hatmo_error = repeat(hatmo_coarse, 1, size(hatmo_fine, 2)) - hatmo_fine

  ax[3, 2][:grid](axis = "x")
  ax[3, 2][:fill_between](hatmo_time, minimum(hatmo_error, dims=2)[:], maximum(hatmo_error, dims=2)[:], color="gray")
  ax[3, 2][:plot](hatmo_time, mean(hatmo_error, dims=2), color = "red")
  ax[3, 2][:set_xlim](xrange)
  ax[3, 2][:set_ylim](yrange)
  ax[3, 2][:set_ylabel]("HATMO (" * L"$W/m^2$" * ")")

  hatmo_error = mean(hatmo_error, dims = 2)

  ax[3, 2][:annotate]("(F) Bias = $(round(mean(hatmo_error), digits = 1)) | RMSE = $(round(sqrt(mean(hatmo_error.^2)), digits = 1))", xy = (0.05, 0.85), xycoords = "axes fraction")

  # LATMO

  latmo_error = repeat(latmo_coarse, 1, size(latmo_fine, 2)) - latmo_fine

  ax[4, 2][:grid](axis = "x")
  ax[4, 2][:fill_between](latmo_time, minimum(latmo_error, dims=2)[:], maximum(latmo_error, dims=2)[:], color="gray")
  ax[4, 2][:plot](latmo_time, mean(latmo_error, dims=2), color = "red")
  ax[4, 2][:set_xlim](xrange)
  ax[4, 2][:set_ylim](yrange)
  ax[4, 2][:set_ylabel]("LATMO (" * L"$W/m^2$" * ")")

  latmo_error = mean(latmo_error, dims = 2)

  ax[4, 2][:annotate]("(H) Bias = $(round(mean(latmo_error), digits = 1)) | RMSE = $(round(sqrt(mean(latmo_error.^2)), digits = 1))", xy = (0.05, 0.85), xycoords = "axes fraction")

  # Save plot

  savefig(joinpath(path_figure, "time_series_analysis_$(metric).pdf"), dpi = 600)

end

for metric in ["rmse"] # ["bias_min", "bias_max", "rmse"]

  # Global settings

  path_figure = joinpath(dirname(pathof(IcanProj)), "..", "plots", "time_series_analysis")

  # Detemine which grid cell to analyze

  variable = "swe"

  cfg = 30    # previously 32

  file_fine, file_coarse = results_files(cfg, variable)

  df_subset = find_largest_error(file_fine, file_coarse, metric, variable)

  # Plot two configurations where only exchng differs

  plot_fms_results(df_subset, path_figure, metric)

  # Plot sensible heat fluxes

  plot_fsm_heat_fluxes(df_subset, metric, path_figure)

end


text:
include("rigidbody.jl")

using Statistics

mass1 = 1.0  # kg
force1 = 4.0 # Newton
Δt = 0.1

body = RigidBody(mass1, force1)
ts = 0:Δt:200
# calculate the approximation using the integration
approx = Float64[]
for t in ts
    integrate!(body, Δt)
    push!(approx, body.position[1])
end

# analytic calculation using the standard mechanics model equation
# abd where the starting distance and velocity are set to 0
acceleration1 = force1/mass1
distance(t) = 0.5*acceleration1*t^2

# created an array of calculations of expected distance for each point in time
expected = map(distance, ts)

# create an array with the difference between expected and approximate
# the zip maps the diff function to the tuple with multiple iterators
# since the calculation here isn't stateful it can all happen simultaneously
diff(v) = abs(last(v) - first(v))
error = map(diff, zip(expected, approx))

# calculate  the mean of expected and approximate and
# see how small the difference is
# “julia> mean(error)
# 19.99999999968928

# julia> mean(expected)
# 26673.333333333332

# julia> mean(error)/mean(expected)
# 0.0007498125468516351

max_stars_repo_name: FourMInfo/JuliaforBeginners

max_issues_repo_name: FourMInfo/JuliaforBeginners

max_forks_repo_name: FourMInfo/JuliaforBeginners


text:
functiontypeof(m::Method) = parameter_typeof(m.sig)[1]
parameter_typeof(sig::UnionAll) = parameter_typeof(sig.body)
parameter_typeof(sig::DataType) = sig.parameters


text:
## We try some wrapper functions. Maybe we will use more.
## The main idea is to try automatically promoting to big integer
## or complex types. Eg. we want to allow sqrt(-.5), without requiring sqrt(Complex(-.5))

## FIXME: LaguerreL(3/2,x) is broken. does not agree with Mma, nor does it when evaluated. Mathics doc agrees with Mma
## Mathics also gives wrong result for LaguerreL(3/2,x) but gives the correct floating point result

## TODO: export sympy and mpmath
## All of the following are in sympy[:mpmath]
## sympy[:mpmath][:ber], bei, ker, kei  ### Kelvin functions
## also in mpmath, struvel, struveh, angerl, webere, lommels1, lommels2
## besseljzero, <-- can use this to test my code, besselyzero, airyaizero, airybizero,
## socrergi, socrerhi, coulombf, coulombg, coulombc,
## whitm, whitw, pcfd (para. cyl func), pcfu, pcfv
## hyp1f1, hyp1f2, hyp1f0, hyp1f2, 2f0, etc. etc. hypercomb (weighted combination)
## legendre, legenp, legenq
##  lots of stuff missing here.
## http://docs.sympy.org/0.7.0/modules/mpmath/functions/elliptic.html
## many things missing. Mertens constant, etc. etc.
## http://docs.sympy.org/0.7.6/modules/mpmath/functions/index.html
## TODO: remove sympy[:mpmath] in our code, use mpmath (which is bound to the same thing)
## Everything on this page. http://docs.sympy.org/0.7.0/modules/mpmath/functions/bessel.html
## TODO: maybe we need to call functions specialized for common bessel orders ?
## TODO: check does mpmath or sympy have more numerical int routines than quadgk ?

### sjgamma

sjgamma(x::T) where {T<:AbstractFloat} = SpecialFunctions.gamma(x)
sjgamma(x::Complex{T}) where {T<:AbstractFloat} = SpecialFunctions.gamma(x)
sjgamma(a) = a |> sjtopy |> sympy_gamma |> pytosj
sjgamma(a,z) = sympy_gamma(sjtopy(a),sjtopy(z)) |> pytosj

sympy_gamma(a) = sympy[:gamma](a)
sympy_gamma(a,z) = sympy[:uppergamma](a,z)

sympy_erf(x) = sympy[:erf](x)
sympy_erf(x,y) = sympy[:erf2](x,y)

### sjerf

sjerf(x::T) where {T<:AbstractFloat} = erf(x)
sjerf(x::Complex{T}) where {T<:AbstractFloat} = erf(x)
sjerf(a) = a |> sjtopy |> sympy_erf |> pytosj
sjerf(a,z) = sympy_erf(sjtopy(a),sjtopy(z)) |> pytosj

### sjlog

function sjlog(x::AbstractFloat)
    x <= 0 ? log(complex(x,zero(x))) : log(x)
end

function sjlog(b::AbstractFloat,x::AbstractFloat)
    x <= 0 ? log(b,complex(x,zero(x))) : log(b,x)
end

function sjlog(b::Number,x::AbstractFloat)
    x <= 0 ? log(b,complex(x,zero(x))) : log(b,x)
end

function sjlog(b::AbstractFloat,x::Real)
    x <= 0 ? log(b,complex(x,zero(x))) : log(b,x)
end

sjlog(b::AbstractFloat,z::Complex{T}) where {T<:AbstractFloat} = log(b,z)
function sjlog(b::AbstractFloat,x::Complex)
    log(b,complex(x,zero(x)))
end

sjlog(z::Complex{T}) where {T<:AbstractFloat} = log(z)
sjlog(b::Number,z::Complex{T}) where {T<:AbstractFloat} = log(b,z)
sjlog(b::Complex{T}, z::AbstractFloat) where {T<:AbstractFloat} = log(b,z)
sjlog(b::Complex{T},z::Complex{V}) where {T<:AbstractFloat,V<:AbstractFloat} = log(b,z)
sjlog(b::Complex{T}, z::Number) where {T<:AbstractFloat} = log(b,z)

sjlog(x)  =  x |> sjtopy |> sympy[:log] |> pytosj

# SymPy reverse the args
function sjlog(b,x)
    sb = sjtopy(b)
    sx = sjtopy(x)
    pytosj(sympy[:log](sx,sb))
end

### sjfactorial

function sjfactorial(n::Int)
    n < 0 && return MinusInfinity
    n < 20 && return factorial(n) # Assumes Int64
    factorial(big(n))
end

sjfactorial(n::BigInt) = n < 0 ? MinusInfinity : factorial(n)
sjfactorial(n::Number) = sjgamma(n+1)

sjfactorial(x) = x |> sjtopy |> sympy[:factorial] |> pytosj