text:
include("common.jl")

function balanced_ρ(p, grav, mesh, elm)
  zdg = mesh.zdg
  Nq, K = size(zdg)
  𝟙 = ones(Bool, Nq)

  pdg = reshape(length(p) == Nq * K ? p : mesh.scatter * p, Nq, K)
  Q = elm.Q

  # Form the RHS vector
  pdg_diff = similar(zdg)
  Fρdg = similar(zdg, Nq * K, Nq * K)
  for e in 1:K
    p̄ = (pdg[:, e] .+ pdg[:, e]') / 2
    pdg_diff[:, e] = ((Q - Q') .* p̄) * 𝟙
  end
  # BCs hack
  pdg_diff[1,1] -= pdg[1,1]
  pdg_diff[end,end] += pdg[end,end]
  p_diff = mesh.scatter' * pdg_diff[:]

  # Form the ρsystem
  T = eltype(zdg)
  Fρdg = zeros(T, Nq * K, Nq * K)
  for e in 1:K
    rng = (e - 1) * Nq .+ (1: Nq)
    Δϕ = grav * (zdg[:, e]' .- zdg[:, e]) / 2
    A = (Q - Q') .* Δϕ
    Fρdg[rng, rng] .= ((Diagonal(A * 𝟙) + A) / 2 )
  end
  Fρ = mesh.scatter' * Fρdg * mesh.scatter
  ρ_bal = -Fρ \ p_diff
  return ρ_bal
end

let
  T = Float64
  K = 10
  N = 5
  nflds = 3
  #
  elem = element_operators(N)
  mesh = create_mesh(K, elem, T(0), T(30e4))
  #
  Nq = N + 1
  #
  params = Params{T}()
  #
  q = rand(Nq, K, nflds)
  q[:,:,3] .+= 100
  p = similar(mesh.zdg)
  ρ = @view q[:, :, 1]
  for (i, z) in enumerate(mesh.zdg)
    (p[i], ρ[i]) = decaying_temperature_profile(params, z)
  end
  ρ[:] .= mesh.scatter * balanced_ρ(p, params.grav, mesh, elem)
  q[:, :, 2] .= 0
  ρe = @view q[:, :, 3]
  ρe .= p / (params.γ - 1) + params.grav * ρ .* mesh.zdg
  #
  γ = T(7 // 5)
  p_(x...) = pressure(x..., γ)
  #
  ∂q = similar(q)
  H(x...) = KennedyGruber(x..., params.grav, p_)
  rhs!(∂q, q, t) = tendency!(∂q, q, t, H, mesh, elem, params.grav, p_)
  c = maximum(wave_speed.(q[:,:,1], p_.(q[:,:,1], q[:,:,2], q[:,:,3], mesh.zdg, params.grav), γ))
  dt = mesh.Δz / (N^2 * c)
  @show extrema(q[:, :, 1])
  @show extrema(q[:, :, 2])
  @show extrema(q[:, :, 3])
  day_sec = 86400.0
  for day in 10:10:100
    tspan = ((day-10) * day_sec, day * day_sec)
    timestep!(q, rhs!, dt, tspan)
    println()
    @show day
    @show extrema(q[:, :, 1])
    @show extrema(q[:, :, 2])
    @show extrema(q[:, :, 3])
  end
  nothing
end


text:
using GLVisualize, GeometryTypes, GLAbstraction
using Colors, Reactive, FileIO

if !isdefined(:runtests)

	window = glscreen()
	timesignal = bounce(linspace(0f0,1f0,360))
end
t 		   = const_lift(*, timesignal, 10f0)
radius     = 200f0
w,h 	   = widths(window)
middle     = Point2f0(w/2, h/2)
circle_pos = Point2f0[(Point2f0(sin(i), cos(i))*radius)+middle for i=linspace(0, 2pi, 20)]
rotation   = Vec2f0[normalize(Vec2f0(middle)-Vec2f0(p)) for p in circle_pos]
scales     = map(t) do t
    Vec2f0[Vec2f0(30, ((sin(i+t)+1)/2)*60) for i=linspace(0, 2pi, 20)]
end

circles = visualize(
	(CIRCLE, circle_pos),
	rotation=rotation, scale=scales,
)

_view(circles, window)

if !isdefined(:runtests)
	renderloop(window)
end

ext: jl


text:
"""
    PyDict{K=PyObject, V=PyObject}([o])

Wrap the Python dictionary `o` (or anything satisfying the mapping interface) as a Julia dictionary with keys of type `K` and values of type `V`.

If `o` is not given, an empty dict is created.
"""
mutable struct PyDict{K,V} <: AbstractDict{K,V}
    ref::PyRef
    hasbuiltins::Bool
    PyDict{K,V}(o) where {K,V} = new(ispyref(o) ? PyRef(o) : pydict(PyRef, o), false)
    PyDict{K,V}() where {K,V} = new(PyRef(), false)
end
PyDict{K}(args...) where {K} = PyDict{K,PyObject}(args...)
PyDict(args...) = PyDict{PyObject,PyObject}(args...)
export PyDict

const pyglobals = PyDict{String}()
export pyglobals

ispyreftype(::Type{<:PyDict}) = true
pyptr(x::PyDict) = begin
    ptr = x.ref.ptr
    if isnull(ptr)
        ptr = x.ref.ptr = C.PyDict_New()
    end
    ptr
end
Base.unsafe_convert(::Type{CPyPtr}, x::PyDict) = checknull(pyptr(x))

Base.iterate(x::PyDict{K,V}, it::PyRef) where {K,V} = begin
    ko = C.PyIter_Next(it)
    if !isnull(ko)
        # key
        r = C.PyObject_Convert(ko, K)
        r == -1 && (C.Py_DecRef(ko); pythrow())
        k = C.takeresult(K)
        # value
        vo = C.PyObject_GetItem(x, ko)
        C.Py_DecRef(ko)
        isnull(vo) && pythrow()
        r = C.PyObject_Convert(vo, V)
        C.Py_DecRef(vo)
        r == -1 && pythrow()
        v = C.takeresult(V)
        # done
        (k => v, it)
    elseif C.PyErr_IsSet()
        pythrow()
    else
        nothing
    end
end
Base.iterate(x::PyDict) = begin
    it = C.PyObject_GetIter(x)
    isnull(it) && pythrow()
    iterate(x, pynewref(it))
end

Base.iterate(x::Base.KeySet{K,PyDict{K,V}}, it::PyRef) where {K,V} = begin
    ko = C.PyIter_Next(it)
    if !isnull(ko)
        r = C.PyObject_Convert(ko, K)
        C.Py_DecRef(ko)
        r == -1 && pythrow()
        k = C.takeresult(K)
        (k, it)
    elseif C.PyErr_IsSet()
        pythrow()
    else
        nothing
    end
end
Base.iterate(x::Base.KeySet{K,PyDict{K,V}}) where {K,V} = begin
    it = C.PyObject_GetIter(x.dict)
    isnull(it) && pythrow()
    iterate(x, pynewref(it))
end

Base.setindex!(x::PyDict{K,V}, v, k) where {K,V} =
    pysetitem(x, convertref(K, k), convertref(V, v))

Base.getindex(x::PyDict{K,V}, k) where {K,V} = pygetitem(V, x, convertref(K, k))

Base.delete!(x::PyDict{K}, k) where {K} = pydelitem(x, convertref(K, k))

Base.length(x::PyDict) = Int(pylen(x))

Base.empty!(x::PyDict) = (@py `$x.clear()`; x)

Base.copy(x::PyDict) = @pyv `$x.copy()`::typeof(x)

Base.haskey(x::PyDict{K}, _k) where {K} = begin
    k = tryconvertref(K, _k)
    k === PYERR() && pythrow()
    k === NOTIMPLEMENTED() && return false
    pycontains(x, k)
end

Base.get(x::PyDict{K}, _k, d) where {K} = begin
    k = tryconvertref(K, _k)
    k === PYERR() && pythrow()
    (k !== NOTIMPLEMENTED() && pycontains(x, k)) ? x[k] : d
end

Base.get(d::Function, x::PyDict{K}, _k) where {K} = begin
    k = tryconvertref(K, _k)
    k === PYERR() && pythrow()
    (k !== NOTIMPLEMENTED() && pycontains(x, k)) ? x[k] : d()
end

Base.get!(x::PyDict{K,V}, _k, d) where {K,V} = begin
    k = convertref(K, _k)
    pycontains(x, k) ? x[k] : (x[k] = convert(V, d))
end

Base.get!(d::Function, x::PyDict{K,V}, _k) where {K,V} = begin
    k = convertref(K, _k)
    pycontains(x, k) ? x[k] : (x[k] = convert(V, d()))
end


text:
"""
  Upsample(mode = :nearest; [scale, size])
  Upsample(scale, mode = :nearest)

An upsampling layer. One of two keywords must be given:

If `scale` is a number, this applies to all but the last two dimensions (channel and batch) of the input.
It may also be a tuple, to control dimensions individually. Alternatively, keyword
`size` accepts a tuple, to directly specify the leading dimensions of the output.

Currently supported upsampling `mode`s
and corresponding NNlib's methods are:
  - `:nearest` -> [`NNlib.upsample_nearest`](@ref)
  - `:bilinear` -> [`NNlib.upsample_bilinear`](@ref)

# Examples

```juliarepl
julia> m = Upsample(scale = (2, 3))
Upsample(:nearest, scale = (2, 3))

julia> m(ones(2, 2, 1, 1)) |> size
(4, 6, 1, 1)

julia> m = Upsample(:bilinear, size = (4, 5))
Upsample(:bilinear, size = (4, 5))

julia> m(ones(2, 2, 1, 1)) |> size
(4, 5, 1, 1)
"""
struct Upsample{mode, S, T}
  scale::S
  size::T
end

function Upsample(mode::Symbol = :nearest; scale = nothing, size = nothing)
  mode in [:nearest, :bilinear] ||
    throw(ArgumentError("mode=:$mode is not supported."))
  if !(isnothing(scale) ⊻ isnothing(size))
    throw(ArgumentError("Either scale or size should be specified (but not both)."))
  end
  return Upsample{mode,typeof(scale),typeof(size)}(scale, size)
end

Upsample(scale, mode::Symbol = :nearest) = Upsample(mode; scale)

(m::Upsample{:nearest})(x::AbstractArray) =
  NNlib.upsample_nearest(x, m.scale)
function (m::Upsample{:nearest, Int})(x::AbstractArray{T, N}) where {T, N}
  NNlib.upsample_nearest(x, ntuple(i -> m.scale, N-2))
end
(m::Upsample{:nearest, Nothing})(x::AbstractArray) =
  NNlib.upsample_nearest(x; size=m.size)

(m::Upsample{:bilinear})(x::AbstractArray) =
  NNlib.upsample_bilinear(x, m.scale)
(m::Upsample{:bilinear, Nothing})(x::AbstractArray) =
  NNlib.upsample_bilinear(x; size=m.size)

function Base.show(io::IO, u::Upsample{mode}) where {mode}
  print(io, "Upsample(")
  print(io, ":", mode)
  u.scale !== nothing && print(io, ", scale = $(u.scale)")
  u.size !== nothing && print(io, ", size = $(u.size)")
  print(io, ")")
end

"""
    PixelShuffle(r::Int)

Pixel shuffling layer with upscale factor `r`.

See [`NNlib.pixel_shuffle`](@ref).
"""
struct PixelShuffle
  r::Int
end

(m::PixelShuffle)(x) = NNlib.pixel_shuffle(x, m.r)