text:
type Point
    x::Float64
    y::Float64
    z::Float64
    lc::Float64
    id::Int
end

type Line
    p1::Point
    p2::Point
    transfinite::Union{Int,Tuple{Int,String,Real}}
    id::Int
end

type LineLoop
    lines::Vector{Line}
    flip::BitArray{1}
    id::Int
end

type Surface
    lineloops::Vector{LineLoop}
    id::Int
end

type PhysicalSurface
    label::String
    surfaces::Vector{Surface}
    id::Int
end

type Model
    points::Vector{Point}
    lines::Vector{Line}
    lineloops::Vector{LineLoop}
    surfaces::Vector{Surface}
    physsurf::Vector{PhysicalSurface}
    opts::Vector{String} # Meshing options
    id::Int # Global ID to keep track of objects
    function Model(;recombineall=false)
        opts = Vector{String}(0)
        recombineall && push!(opts,"Mesh.RecombineAll = 1;")
        return new(Vector{Point}(0),Vector{Line}(0),Vector{LineLoop}(0),Vector{Surface}(0),Vector{PhysicalSurface}(0),opts,0)
    end
end

function add_point!(m::Model,x,y,z,lc=0.)
    # Make sure point does not already exist in model
    i = find(p->p.x==x&&p.y==y&&p.z==z,m.points)
    if isempty(i)
        o = Point(x,y,z,lc,m.id+=1)
        push!(m.points,o)
    else
        o = m.points[i[1]]
        o.lc = lc
    end
    return o
end

function get_points(ll::LineLoop)
    p = Vector{Point}(length(ll.lines))
    for (i,l) in enumerate(ll.lines)
        p[i] = l.p1
    end
    return p
end

function add_line!(m::Model,p1::Point,p2::Point,transfinite::Union{Int,Tuple{Int,String,Real}}=0)
    o = Line(p1,p2,transfinite,m.id+=1)
    push!(m.lines,o)
    return o
end

function add_lineloop!(m::Model,lines::Vector{Line},flip::BitArray{1}=falses(length(lines)))
    o = LineLoop(lines,flip,m.id+=1)
    push!(m.lineloops,o)
    return o
end

function add_surface!(m::Model,ll::Union{LineLoop,Vector{LineLoop}})
    o = Surface([ll;],m.id+=1)
    push!(m.surfaces,o)
    return o
end

function add_physicalsurf!(m::Model,label::String,sf::Union{Surface,Vector{Surface}})
    # If a physical surface with the same label already exists, add to it
    i = find(p->p.label==label,m.physsurf)
    if isempty(i)
        o = PhysicalSurface(label,[sf;],m.id+=1)
        push!(m.physsurf,o)
    else
        o = m.physsurf[i[1]]
        push!(m.physsurf.surfaces,[sf;])
    end
    return o
end

function duplicate!(m::Model,p::Point,t::Tuple{Real,Real,Real})
    return add_point!(m,p.x+t[1],p.y+t[2],p.z+t[3],p.lc)
end

function duplicate!(m::Model,l::Line,t::Tuple{Real,Real,Real})
    p1_new = duplicate!(m,l.p1,t)
    p2_new = duplicate!(m,l.p2,t)
    return add_line!(m,p1_new,p2_new,l.transfinite)
end

function duplicate!(m::Model,ll::LineLoop,t::Tuple{Real,Real,Real})
    lines = duplicate!.([m],ll.lines,[t])
    return add_lineloop!(m,lines)
end

function duplicate!(m::Model,s::Surface,t::Tuple{Real,Real,Real})
    ll = duplicate!.([m],s.lineloops,[t])
    return add_surface!(m,ll)
end

function extrude!(m::Model,s1::Surface,t::Tuple{Real,Real,Real};transfinite::Union{Int,Tuple{Int,String,Real}}=0)
    # First duplicate the surface (including any holes)
    s2 = duplicate!(m,s1,t)
    ns = [s2;]
    num_ll = length(s1.lineloops)
    flip = BitArray{1}([false,false,true,true])
    # For each lineloop
    for i = 1:num_ll
        num_vert = length(s1.lineloops[i].lines)
        line = Vector{Line}(num_vert)
        ll1 = s1.lineloops[i] # "Bottom"
        ll2 = s2.lineloops[i] # "Top"
        # Find matching vertices and join them with lines
        for j = 1:num_vert
            line[j] = add_line!(m,ll1.lines[j].p1,ll2.lines[j].p1,transfinite)
        end
        # Form new surfaces for each face, using the lines we just made
        for j = 1:num_vert
            ll = add_lineloop!(m,[ll1.lines[j],circshift(line,-1)[j],ll2.lines[j],line[j]],flip)
            push!(ns,add_surface!(m,ll))
        end
    end
    return ns
end

function print_geo_file(m::Model)
    title = "// File generated by Gmsh.jl"
    opts  = m.opts
    body  = mapreduce(vcat,fieldnames(m)[1:end-2]) do field
        mapreduce(print_object,vcat,getfield(m,field))
    end
    return join([title;opts;body],NL)
end

function print_object(o::Point)
    a = "o$(o.id) = newp;"
    b = "Point(o$(o.id)) = {$(o.x), $(o.y), $(o.z)$(o.lc==0.?"};":" , $(o.lc)};")"
    return join([a,b],NL)
end

function print_object(o::Line)
    a = "o$(o.id) = newl;"
    b = "Line(o$(o.id)) = {o$(o.p1.id), o$(o.p2.id)};"
    if o.transfinite==0
        return join([a,b],NL)
    else
        t = o.transfinite
        c = "Transfinite Line{o$(o.id)} = $(isa(t,Int)?"$t":"$(t[1]) Using $(t[2]) $(t[3])");"
        return join([a,b,c],NL)
    end
end

function print_object(o::LineLoop)
    a = "o$(o.id) = newll;"
    linelist = [(o.flip[i]?"-":"")*"o$(o.lines[i].id)" for i=1:length(o.lines)]
    b = "Line Loop(o$(o.id)) = {$(join(linelist,", "))};"
    return join([a,b],NL)
end

function print_object(o::Surface)
    a = "o$(o.id) = news;"
    b = "Plane Surface(o$(o.id)) = {$(join(["o$(ll.id)" for ll in o.lineloops],", "))};"
    if any(getfield.(o.lineloops[1].lines,[:transfinite]).!=0)
        c = "Transfinite Surface{o$(o.id)};"
        return join([a,b,c],NL)
    else
        return join([a,b],NL)
    end
end

function print_object(o::PhysicalSurface)
    a = "Physical Surface(\"$(o.label)\") = {$(join(["o$(s.id)" for s in o.surfaces],", "))};"
    return a
end

max_stars_repo_name: jebej/FastCap

max_issues_repo_name: jebej/FastCap

max_forks_repo_name: jebej/FastCap


text:
using Flux
using Flux: @functor
using NeuralAttentionlib
using NeuralAttentionlib: CausalMask, BiLengthMask, BatchedMask
using NNlib
include("./mha.jl")

#https://github.com/chengchingwen/Transformers.jl/blob/master/src/basic/transformer.jl
struct TransformerDense{input<:Dense, output<:Dense}
    din::input
    dout::output
end

@functor TransformerDense

"just a wrapper for two dense layer."
TransformerDense(size::Int, h::Int, act = relu) = TransformerDense(
    Dense(size, h, act),
    Dense(h, size)
)

function (pw::TransformerDense)(x::AbstractMatrix)
  # size(x) == (dims, seq_len)
  return pw.dout(pw.din(x))
end

function (d::TransformerDense)(x::A) where {T, N, A<:AbstractArray{T, N}}
  new_x = reshape(x, size(x, 1), :)
  y = d(new_x)
  return reshape(y, Base.setindex(size(x), size(y, 1), 1))
end

struct TransformerBlock
    mha
    ln
    dense
end

@functor TransformerBlock

"""
Generic transformer block
"""
TransformerBlock(args) = TransformerBlock(
    MultiheadAttention(args.mha_heads, args.mha_input_dims, args.mha_head_dims,
    args.mha_input_dims; future = true, pdrop = args.transformer_dropout),
    LayerNorm(sum([i[1] for i in args.embedding_dims])),
    Chain(TransformerDense(args.mha_input_dims, args.transformer_dense_hidden_dim, args.activation),
            Dropout(args.transformer_dropout))
)

function (t::TransformerBlock)(x)
    h1 = atten(t.mha, x)
    h2 = x + h1
    h2 = t.ln(h2)
    h3 = t.dense(h2)
    h4 = h2 + h3
    h4 = t.ln(h4)
    return h4
end

ext: jl


text:
module Drawing

using CoilFunctions, Unitful, Luxor

function drawwire(p::Point, awg::AWG, fillfactor=1.0;u=u"mm")
    f = x->ustrip(u(x))
    setcolor("green")
    circle(p,f(awg.d_insul)/2,:fill)
    setcolor("blue")
    circle(p,f(awg.d_bare)/2,:fill)
    # setcolor("black")
    # circle(p,f(awg.d_insul)/sqrt(fillfactor)/2, :stroke)
    # @info f(awg.d_insul)
end

function drawcoil(awg::AWG,winding::WindingLayers,maxturns::Integer,fillfactor::Real=1.0)
    turns = zero(maxturns)
    d = awg.d_insul/sqrt(fillfactor) |> u"mm" |> ustrip
    for (layer,y) = enumerate(range(d/2;step=d*√(3)/2,length=winding.layers))
        (turnlayers,xoffset) = isodd(layer) ? (winding.odd_layer,d/2) : (winding.even_layer,d)
        for x in range(xoffset;step=d,length=turnlayers)
            turns >= maxturns && break
            drawwire(Point(x,y),awg,fillfactor)
            turns += 1
        end
    end
end
drawcoil(a,fillfactor) = drawcoil(CoilFunctions.AWG_Chart[a.gauge], a.winding, a.turns, fillfactor)

function drawcoil(coil::CoilGeometry)
    y=(coil.od - coil.id)/2 |> u"mm" |> ustrip
    x=ustrip(u"mm"(coil.len))
    gsave()
    setcolor("black")
    box(Point(0,0),Point(x,y), :stroke)
    grestore()
end

function illustratecoil(coil::CoilGeometry, Ω::Unitful.ElectricalResistance, fillfactor=1.0; filename=:png, pad=20,
    resolution=(800,600), minturns=0, maxturns=typemax(Int64))

    x,y = resolution
    a=optimalcoil(coil,Ω,fillfactor)
    Drawing(x+2pad, y+2pad, filename)
    origin(pad,y+pad)
    s = ustrip(x/coil.len)
    scale(s,-s)
    background("white")
    sethue("black")
    drawcoil(coil)
    turns = clamp(a.turns,minturns, maxturns)
    drawcoil(merge(a,(turns=turns,)),a.fillfactor)
    finish()
    @info a
    preview()
end

function illustratecoil(coil::CoilGeometry, v::Unitful.Voltage, p::Unitful.Power, fillfactor=1.0; kwargs...)
    illustratecoil(coil, v^2/p |> u"Ω", fillfactor;kwargs...)
end

function illustratecoil(s::Symbol, v::Unitful.Voltage, fillfactor=1.0; kwargs...)
    coil=CoilFunctions.coil[s]
    Ω=CoilFunctions.Ω(s,v)
    illustratecoil(coil,Ω,fillfactor;kwargs...)
end

end


text:
module ImplicitDomainQuadrature

using LinearAlgebra
using IntervalArithmetic
using FastGaussQuadrature
using Roots
using PolynomialBasis

include("bounds.jl")
include("utilities.jl")
include("quadrature_structs.jl")
include("one_dimensional_quadrature.jl")
include("area_quadrature.jl")
include("surface_quadrature.jl")

export sign_allow_perturbations,
    area_quadrature,
    surface_quadrature,
    tensor_product_quadrature,
    QuadratureRule,
    ReferenceQuadratureRule,
    update_interpolating_gradient!

end # module

max_stars_repo_name: ArjunNarayanan/ImplicitDomainQuadrature

max_issues_repo_name: ArjunNarayanan/ImplicitDomainQuadrature

max_forks_repo_name: ArjunNarayanan/ImplicitDomainQuadrature


text:
export simplify

@doc Markdown.doc"""
    simplify(K::AnticNumberField; canonical::Bool = false) -> AnticNumberField, NfToNfMor
 > Tries to find an isomorphic field $L$ given by a "nicer" defining polynomial.
 > By default, "nice" is defined to be of smaller index, testing is done only using
 > a LLL-basis of the maximal order.
 > If \texttt{canonical} is set to {{{true}}}, then a canonical defining
 > polynomial is found, where canonical is using the pari-definition of {{{polredabs}}}
 > in http://beta.lmfdb.org/knowledge/show/nf.polredabs.
 > Both version require a LLL reduced basis for the maximal order.
"""
function simplify(K::AnticNumberField; canonical::Bool = false, cached = false)
  Qx, x = PolynomialRing(FlintQQ)
  if canonical
    a, f1 = polredabs(K)
    f = Qx(f1)
  else
    OK = maximal_order(K)
    if isdefined(OK, :lllO)
      ZK = OK.lllO
    else
      prec = 100 + 25*div(degree(K), 3) + Int(round(log(abs(discriminant(K)))))
      ZK = _lll(OK, prec = prec)[2]
    end
    a = gen(K)
    if isdefining_polynomial_nice(K)
      I = index(OK)
    else
      I = divexact(numerator(discriminant(K.pol)), discriminant(ZK))
    end
    B = basis(ZK, copy = false)
    for i = 1:length(B)
      if isone(denominator(B[i].elem_in_nf))
        continue
      end
      ind_a = _index(B[i])
      if !iszero(ind_a) && ind_a < I
        a = B[i].elem_in_nf
        I = ind_a
      end
    end
    f = minpoly(Qx, a)
  end
  L = NumberField(f, cached = cached, check = false)[1]
  m = hom(L, K, a, check = false)
  return L, m
end

function _index(a::NfOrdElem)
  O = parent(a)
  K = nf(O)
  d = degree(O)
  pows = Vector{nf_elem}(undef, d)
  pows[1] = one(K)
  pows[2] = a.elem_in_nf
  for i = 3:d
    pows[i] = pows[i-1]*a.elem_in_nf
  end
  M = basis_mat(pows)
  hnf!(M)
  res = prod(M[i, i] for i = 1:degree(K))
  return numerator(res*index(O))
end

 #a block is a partition of 1:n
 #given by the subfield of parent(a) defined by a
 #the embeddings used are in R
 #K = parent(a)
 # then K has embeddings into the finite field (parent of R[1])
 # given by the roots (in R) of the minpoly of K
 #integers in 1:n are in the same block iff a(R[i]) == a(R[j])
 #the length of such a block (system) is the degree of Q(a):Q, the length
 # of a block is the degree K:Q(a)
 # a is primitive iff the block system has length n
function _block(a::nf_elem, R::Array{fq_nmod, 1}, ap::fq_nmod_poly)
  c = FlintZZ()
  ap.length = a.elem_length
  for i=0:a.elem_length
    Nemo.num_coeff!(c, a, i)
    setcoeff!(ap, i, c)
  end
#  ap = Ft(Zx(a*denominator(a)))
  s = fq_nmod[evaluate(ap, x) for x = R]
  b = Vector{Int}[]
  a = BitSet()
  i = 0
  n = length(R)
  while i < n
    i += 1
    if i in a
      continue
    end
    z = s[i]
    push!(b, findall(x->s[x] == z, 1:n))
    for j in b[end]
      push!(a, j)
    end
  end
  return b
end

#given 2 block systems b1, b2 for elements a1, a2, this computes the
#system for Q(a1, a2), the compositum of Q(a1) and Q(a2) as subfields of K
function _meet(b1::Vector{Vector{Int}}, b2::Vector{Vector{Int}})
  b = Vector{Int}[]
  for i=b1
    for j = i
      for h = b2
        if j in h
          s = intersect(i, h)
          if ! (s in b)
            push!(b, s)
          end
        end
      end
    end
  end
  return b
end

function _find_prime(f::fmpz_poly)
  p = 2^20
  d = 1
  while true
    p = next_prime(p)
    R = GF(p, cached=false)
    fR = change_base_ring(f, R)
    if !issquarefree(fR)
      continue
    end
    FS = factor_shape(fR)
    d = lcm([x for (x, v) in FS])
    if d < degree(fR)^2
      break
    end
  end
  return p, d
end

function polredabs(K::AnticNumberField)
  #intended to implement
  # http://beta.lmfdb.org/knowledge/show/nf.polredabs
  #as in pari
  #TODO: figure out the separation of T2-norms....
  ZK = lll(maximal_order(K))
  I = index(ZK)^2
  D = discriminant(ZK)
  B = basis(ZK, copy = false)
  Zx = FlintZZ["x"][1]
  f = Zx(K.pol)
  p, d = _find_prime(f)

  F = FlintFiniteField(p, d, "w", cached = false)[1]
  Ft = PolynomialRing(F, "t", cached = false)[1]
  ap = zero(Ft)
  fit!(ap, degree(K)+1)
  rt = roots(f, F)

  n = degree(K)

  b = _block(B[1].elem_in_nf, rt, ap)
  i = 2
  while length(b) < degree(K)
    bb = _block(B[i].elem_in_nf, rt, ap)
    b = _meet(b, bb)
    i += 1
  end
  i -= 1
#  println("need to use at least the first $i basis elements...")
  pr = 100
  old = precision(BigFloat)
  E = enum_ctx_from_ideal(ideal(ZK, 1), zero_matrix(FlintZZ, 1, 1), prec = pr, TU = BigFloat, TC = BigFloat)
  setprecision(BigFloat, pr)
  while true
    try
      if E.C[end] + 0.0001 == E.C[end]  # very very crude...
        pr *= 2
        continue
      end
      break
    catch e
      if isa(e, InexactError) || isa(e, LowPrecisionLLL) || isa(e, LowPrecisionCholesky)
        pr *= 2
        continue
      end
      rethrow(e)
    end
    setprecision(BigFloat, pr)
    E = enum_ctx_from_ideal(ideal(ZK, 1), zero_matrix(FlintZZ, 1, 1), prec = pr, TU = BigFloat, TC = BigFloat)
  end

  l = zeros(FlintZZ, n)
  l[i] = 1

  scale = 1.0
  enum_ctx_start(E, matrix(FlintZZ, 1, n, l), eps = 1.01)

  a = gen(K)
  all_a = nf_elem[a]
  la = length(a)*BigFloat(E.t_den^2)
  Ec = BigFloat(E.c//E.d)
  eps = BigFloat(E.d)^(1//2)

  found_pe = false
  while !found_pe
    while enum_ctx_next(E)
#      @show E.x
      M = E.x*E.t
      q = elem_from_mat_row(K, M, 1, E.t_den)
      bb = _block(q, rt, ap)
      if length(bb) < n
        continue
      end
      found_pe = true
#  @show    llq = length(q)
#  @show sum(E.C[i,i]*(BigFloat(E.x[1,i]) + E.tail[i])^2 for i=1:E.limit)/BigInt(E.t_den^2)
      lq = Ec - (E.l[1] - E.C[1, 1]*(BigFloat(E.x[1,1]) + E.tail[1])^2) #wrong, but where?
#      @show lq/E.t_den^2

      if lq < la + eps
        if lq > la - eps
          push!(all_a, q)
#          @show "new one", q, minpoly(q), bb
        else
          a = q
          all_a = nf_elem[a]
          if lq/la < 0.8
#            @show "re-init"
            enum_ctx_start(E, E.x, eps = 1.01)  #update upperbound
            Ec = BigFloat(E.c//E.d)
          end
          la = lq
  #        @show Float64(la/E.t_den^2)
        end
      end
    end
    scale *= 2
    enum_ctx_start(E, matrix(FlintZZ, 1, n, l), eps = scale)
    Ec = BigFloat(E.c//E.d)
  end

  setprecision(BigFloat, old)
  all_f = Tuple{nf_elem, fmpq_poly}[(x, minpoly(x)) for x=all_a]
  all_d = fmpq[abs(discriminant(x[2])) for x= all_f]
  m = minimum(all_d)
  L1 = Tuple{nf_elem, fmpq_poly}[]
  for i = 1:length(all_f)
    if all_d[i] == m
      push!(L1, all_f[i])
    end
  end
  L2 = Tuple{nf_elem, fmpq_poly}[minQ(x) for x=L1]
  if length(L2) == 1
    return L2[1]
  end
  L3 = sort(L2, lt = il)

  return L3[1]
end

function Q1Q2(f::PolyElem)
  q1 = parent(f)()
  q2 = parent(f)()
  g = gen(parent(f))
  for j=0:degree(f)
    if isodd(j)
      q2 += coeff(f, j)*g^div(j, 2)
    else
      q1 += coeff(f, j)*g^div(j, 2)
    end
  end
  return q1, q2
end

function minQ(A::Tuple{nf_elem, fmpq_poly})
  a = A[1]
  f = A[2]
  q1, q2 = Q1Q2(f)
  if lead(q1)>0 && lead(q2) > 0
    return (-A[1], f(-gen(parent(f)))*(-1)^degree(f))
  else
    return (A[1], f)
  end
end

function int_cmp(a, b)
  if a==b
    return 0
  end
  if abs(a) == abs(b)
    if a>b
      return 1
    else
      return -1
    end
  end
  return cmp(abs(a), abs(b))
end

function il(F, G)
  f = F[2]
  g = G[2]
  i = degree(f)
  while i>0 && int_cmp(coeff(f, i), coeff(g, i))==0
    i -= 1
  end
  return int_cmp(coeff(f, i), coeff(g, i))<0
end

################################################################################
#
#  LLL for simplify
#
################################################################################

function _lll(M::NfOrd; prec = 100)

  K = nf(M)

  if istotally_real(K)
    #in this case the gram-matrix of the minkowski lattice is the trace-matrix
    #which is exact.
    BM = _lll_gram(ideal(M, 1))[2]
    O1 = NfOrd(K, BM*basis_mat(M, copy = false))
    O1.ismaximal = M.ismaximal
    if isdefined(M, :index)
      O1.index = M.index
    end
    if isdefined(M, :disc)
      O1.disc = M.disc
    end
    if isdefined(M, :gen_index)
      O1.gen_index = M.gen_index
    end
    M.lllO = O1
    return true, O1
  end

  if degree(K) == 2 && discriminant(M) < 0
    #in this case the gram-matrix of the minkowski lattice is related to the
    #trace-matrix which is exact.
    #could be extended to CM-fields
    BM = _lll_quad(ideal(M, 1))[2]
    O1 = NfOrd(K, BM*basis_mat(M, copy = false))
    O1.ismaximal = M.ismaximal
    if isdefined(M, :index)
      O1.index = M.index
    end
    if isdefined(M, :disc)
      O1.disc = M.disc
    end
    if isdefined(M, :gen_index)
      O1.gen_index = M.gen_index
    end
    M.lllO = O1
    return true, O1
  end

  n = degree(M)
  prec = max(prec, 4*n)
  local d::fmpz_mat
  while true
    try
      d = minkowski_gram_mat_scaled(M, prec)
      break
    catch e
      prec = prec*2
    end
  end
  g = zero_matrix(FlintZZ, n, n)

  prec = div(prec, 2)
  shift!(d, -prec)  #TODO: remove?

  for i=1:n
    fmpz_mat_entry_add_ui!(d, i, i, UInt(nrows(d)))
  end

  ctx = Nemo.lll_ctx(0.99, 0.51, :gram)

  ccall((:fmpz_mat_one, :libflint), Nothing, (Ref{fmpz_mat}, ), g)
  ccall((:fmpz_lll, :libflint), Nothing, (Ref{fmpz_mat}, Ref{fmpz_mat}, Ref{Nemo.lll_ctx}), d, g, ctx)

  fl = true
  ## test if entries in l are small enough, if not: increase precision
  ## or signal that prec was too low

  if nbits(maximum(abs, g)) >  div(prec, 2)
    fl = false
  else
    ## lattice has lattice disc = order_disc * norm^2
    ## lll needs to yield a basis sth
    ## l[1,1] = |b_i|^2 <= 2^((n-1)/2) disc^(1/n)
    ## and prod(l[i,i]) <= 2^(n(n-1)/2) disc
    n = nrows(d)
    disc = abs(discriminant(M))
    di = root(disc, n)+1
    di *= fmpz(2)^(div(n+1,2)) * fmpz(2)^prec

    if cmpindex(d, 1, 1, di) > 0
      fl = false
    else
      pr = prod_diag(d)
      if pr > fmpz(2)^(div(n*(n-1), 2)) * disc * fmpz(2)^(n*prec)
        fl = false
      end
    end
  end
  On = NfOrd(K, g*basis_mat(M, copy = false))
  On.ismaximal = M.ismaximal
  if isdefined(M, :index)
    On.index = M.index
  end
  if isdefined(M, :disc)
    On.disc = M.disc
  end
  if isdefined(M, :gen_index)
    On.gen_index = M.gen_index
  end
  if fl
    M.lllO = On
  end
  return fl, On
end


text:
using LinearAlgebra,Plots;N=128;P=64N;dt=1/4N;T=1024;W=200;w=W/P*N;
R=rand;x=R(P);v=R([-1.0,1.0],P);m1=mod1;u()=(x.=m1.(x.+v/2*dt,1));K=zeros(T,2);
g(x)=Int.(m1.(round.(x.*N),N));
D2=zeros(N,N);for i∈1:N;D2[i,m1.(i.+[-1,0,1],N)].=[1,-2,1]*N^2;end
D1=zeros(N,N);for i∈1:N;D1[i,m1.(i.+[0,1],N)].=[-1,1]*N;end
@gif for t in 1:T;
  u();ϕ=D2\(r=sum(i->i.==1:N,g(x))*w;r.-=sum(r)/N);u();
  E=D1*ϕ;v+=E[g(x.-1/2N)]*dt;scatter(x,v,ylims=(-3,3))
  K[t,:].=(sum(E.^2)/2N,sum(v.^2/2)*W/P);plot!((1:t)/T,K[1:t,:].*2/W);
  plot!((1:t)/T,sum(K[1:t,:].*2/W,dims=2),legend=0==1);
end every 8


text:
#jl #! format: off
# # Intro to Julia
# **Originally Contributed by**: Juan Pablo Vielma

# ## Introduction

# Since JuMP is embedded in Julia, knowing some basic Julia is important
# for learning JuMP. This notebook is designed to provide a minimalist
# crash course in the basics of Julia. You can find resources that provide
# a more comprehensive introduction to Julia [here](https://julialang.org/learning/).

# ### How to Print
# In Julia, we usually use println() to print

println("Hello, World from Julia!")

# ### Basic Data Types
# Integers

typeof(1 + -2)

# Floating point numbers

typeof(1.2 - 2.3)

# There are also some cool things like an irrational representation of π. To make π
# (and most other greek letters), type \pi and then press [TAB].

π

#+

typeof(π)

# Julia has native support for complex numbers

typeof(2 + 3im)


text:
using Test
using Combinatorics
using Assignment

find_best(M, k) = size(M, 1) > size(M, 2) ? find_best(M', k) : partialsort([
    (sum(M[CartesianIndex.(axes(M, 1), p)]), p) for p in permutations(axes(M, 2), size(M, 1))], 1:k)

function test_kbest(M, k, sols)
    for ((cost, perm), sol) in zip(find_best(M, k), sols)
        @test cost ≈ sol.cost
        @test perm == (sol.use_row ? sol.row4col : sol.col4row)
    end
end

@testset "find best assignment" begin
      @testset "Float64" begin
          A = [ 0.891171  0.0320582   0.564188  0.8999    0.620615;
                0.166402  0.861136    0.201398  0.911772  0.0796335;
                0.77272   0.782759    0.905982  0.800239  0.297333;
                0.561423  0.170607    0.615941  0.960503  0.981906;
                0.748248  0.00799335  0.554215  0.745299  0.42637]

          sol = @inferred find_best_assignment(A)
          @test sol.row4col == [2, 3, 5, 1, 4]
          @test sol.cost ≈ 1.8375112
          @test sum(A[sol]) ≈ 1.8375112
      end

      @testset "Int64" begin
          B = [ 24     1     8;
                 5     7    14;
                 6    13    20;
                12    19    21;
                18    25     2]

          sol = @inferred find_best_assignment(B)
          @test sol.row4col == [2, 1, 0, 0, 3]
          @test sol.col4row == [2, 1, 5]
          @test sol.cost == 8
          @test sum(B[sol]) == 8

          sol = @inferred find_best_assignment(B')
          @test sol.col4row == [2, 1, 0, 0, 3]
          @test sol.row4col == [2, 1, 5]
          @test sol.cost == 8
          @test sum(B'[sol]) == 8

          sol = @inferred find_best_assignment([ifelse(i-j==0,0,i*j) for i in 1:5, j in 1:5])
          @test sol.col4row == [1, 2, 3, 4, 5]
          @test sol.row4col == [1, 2, 3, 4, 5]
          @test sol.cost == 0
      end

      @testset "UInt16" begin
        M=UInt16[28092 44837 19882 39481 59139;
                 26039 46258 38932 51057 9;
                 11527 59487 61993 29072 8734;
                 10691 16977 12796 16370 14266;
                 5199  42319 34194 41332 16472]
        sol = @inferred find_best_assignment(M)
        @test sol.row4col == [3, 5, 4, 2, 1]
        @test sum(M[sol]) == 71139
        @test sol.cost    == 71139 % 2^16
    end

    @testset "UInt8" begin
        M=UInt8[67  228 135 197 244;
                112 44  84  206 31;
                225 103 231 225 227;
                170 37  135 9   130;
                110 22  133 77  96]

        sol = @inferred find_best_assignment(M)
        @test sol.row4col == [1, 5, 2, 4, 3]
        @test sum(M[sol]) == 343
        @test sol.cost   == 343 % 2^8
    end
end

@testset "find kbest assignments" begin
    @testset "Float64" begin
        A = [ 0.891171  0.0320582   0.564188  0.8999    0.620615;
              0.166402  0.861136    0.201398  0.911772  0.0796335;
              0.77272   0.782759    0.905982  0.800239  0.297333;
              0.561423  0.170607    0.615941  0.960503  0.981906;
              0.748248  0.00799335  0.554215  0.745299  0.42637]

        sols = @inferred find_kbest_assignments(A, 10)
        sol = sols[1]
        @test sol.row4col == [2, 3, 5, 1, 4]
        @test sol.cost ≈ 1.8375112
        @test sum(A[sol]) ≈ 1.8375112
        test_kbest(A, 10, sols)
    end

    @testset "Int64" begin
        B = [ 24     1     8;
               5     7    14;
               6    13    20;
              12    19    21;
              18    25     2]

        sols = @inferred find_kbest_assignments(B, 10)
        sol = sols[1]
        @test sol.row4col == [2, 1, 0, 0, 3]
        @test sol.col4row == [2, 1, 5]
        @test sol.cost == 8
        @test sum(B[sol]) == 8
        test_kbest(B, 10, sort(sols, by=x -> (x.cost, x.col4row)))

        sols = @inferred find_kbest_assignments(B', 10)
        sol = sols[1]
        @test sol.col4row == [2, 1, 0, 0, 3]
        @test sol.row4col == [2, 1, 5]
        @test sol.cost == 8
        @test sum(B'[sol]) == 8
        test_kbest(B', 10, sort(sols, by=x -> (x.cost, x.row4col)))

        C = [ifelse(i-j==0,0,i*j) for i in 1:5, j in 1:5]
        sols = @inferred find_kbest_assignments(C, 10)
        sol = sols[1]
        @test sol.col4row == [1, 2, 3, 4, 5]
        @test sol.row4col == [1, 2, 3, 4, 5]
        @test sol.cost == 0
        test_kbest(C, 10, sort(sols, by=x -> (x.cost, x.row4col)))
    end

    @testset "UInt16" begin
      M=UInt16[28092 44837 19882 39481 59139;
               26039 46258 38932 51057 9;
               11527 59487 61993 29072 8734;
               10691 16977 12796 16370 14266;
               5199  42319 34194 41332 16472]
      sols = @inferred find_kbest_assignments(M, 10)
      sol = sols[1]
      @test sol.row4col == [3, 5, 4, 2, 1]
      @test sum(M[sol]) == 71139
      @test sol.cost    == 71139 % 2^16
  end

  @testset "UInt8" begin
      M=UInt8[67  228 135 197 244;
              112 44  84  206 31;
              225 103 231 225 227;
              170 37  135 9   130;
              110 22  133 77  96]

      sols = @inferred find_kbest_assignments(M, 10)
      sol = sols[1]
      @test sol.row4col == [1, 5, 2, 4, 3]
      @test sum(M[sol]) == 343
      @test sol.cost   == 343 % 2^8
  end
end


text:
abstract type PathCommand end

struct BezierPath
    commands::Vector{PathCommand}
end

struct MoveTo{Dim, T} <: PathCommand
    p::Point{Dim, T}
end

MoveTo(x, y) = MoveTo(Point(x, y))

struct LineTo{Dim, T} <: PathCommand
    p::Point{Dim, T}
end

LineTo(x, y) = LineTo(Point(x, y))

struct CurveTo{Dim, T} <: PathCommand
    c1::Point{Dim, T}
    c2::Point{Dim, T}
    p::Point{Dim, T}
end

CurveTo(cx1, cy1, cx2, cy2, p1, p2) = CurveTo(
    Point(cx1, cy1), Point(cx2, cy2), Point(p1, p1)
)

struct EllipticalArc{Dim, T} <: PathCommand
    c::Point{Dim, T}
    r1::Float64
    r2::Float64
    angle::Float64
    a1::Float64
    a2::Float64
end

EllipticalArc(cx, cy, r1, r2, angle, a1, a2) = EllipticalArc(Point(cx, cy),
    r1, r2, angle, a1, a2)

struct ClosePath <: PathCommand end

function Base.:+(pc::P, p::Point2) where P <: PathCommand
    fnames = fieldnames(P)
    P(map(f -> getfield(pc, f) + p, fnames)...)
end

scale(bp::BezierPath, s::Real) = BezierPath([scale(x, Vec(s, s)) for x in bp.commands])
scale(bp::BezierPath, v::VecTypes{2}) = BezierPath([scale(x, v) for x in bp.commands])
translate(bp::BezierPath, v::VecTypes{2}) = BezierPath([translate(x, v) for x in bp.commands])

translate(m::MoveTo, v::VecTypes{2}) = MoveTo(m.p .+ v)
translate(l::LineTo, v::VecTypes{2}) = LineTo(l.p .+ v)
translate(c::CurveTo, v::VecTypes{2}) = CurveTo(c.c1 .+ v, c.c2 .+ v, c.p .+ v)
translate(e::EllipticalArc, v::VecTypes{2}) = EllipticalArc(e.c .+ v, e.r1, e.r2, e.angle, e.a1, e.a2)
translate(c::ClosePath, v::VecTypes{2}) = c

scale(m::MoveTo, v::VecTypes{2}) = MoveTo(m.p .* v)
scale(l::LineTo, v::VecTypes{2}) = LineTo(l.p .* v)
scale(c::CurveTo, v::VecTypes{2}) = CurveTo(c.c1 .* v, c.c2 .* v, c.p .* v)
scale(e::EllipticalArc, v::VecTypes{2}) = EllipticalArc(e.c .* v, e.r1, e.r2, e.angle, e.a1, e.a2)
scale(c::ClosePath, v::VecTypes{2}) = c

Base.:+(pc::EllipticalArc, p::Point2) = EllipticalArc(pc.c + p, pc.r1, pc.r2, pc.angle, pc.a1, pc.a2)
Base.:+(pc::ClosePath, p::Point2) = pc
Base.:+(bp::BezierPath, p::Point2) = BezierPath(bp.commands .+ Ref(p))

# markers with unit area

BezierCircle = let
    r = sqrt(1/pi)
    BezierPath([
        EllipticalArc(Point(0.0, 0), r, r, 0.0, 0.0, 2pi),
        ClosePath()
    ])
end

BezierUTriangle = let
    r = Float32(sqrt(1 / (3 * sqrt(3) / 4)))
    BezierPath([
        MoveTo(Point2(cosd(90), sind(90)) .* r),
        LineTo(Point2(cosd(210), sind(210)) .* r),
        LineTo(Point2(cosd(330), sind(330)) .* r),
        ClosePath()
    ])
end

BezierDTriangle = let
    r = Float32(sqrt(1 / (3 * sqrt(3) / 4)))
    BezierPath([
        MoveTo(Point2(cosd(270), sind(270)) .* r),
        LineTo(Point2(cosd(390), sind(390)) .* r),
        LineTo(Point2(cosd(510), sind(510)) .* r),
        ClosePath()
    ])
end

BezierSquare = let
    BezierPath([
        MoveTo(Point2(0.5, -0.5)),
        LineTo(Point2(0.5, 0.5)),
        LineTo(Point2(-0.5, 0.5)),
        LineTo(Point2(-0.5, -0.5)),
        ClosePath()
    ])
end

BezierCross = let
    cutfraction = 2/3
    # 1 = (2r)^2 - 4 * (r * c) ^ 2
    # c^2 - 1 != 0, r = 1/(2 sqrt(1 - c^2))
    #
    r = 1/(2 * sqrt(1 - cutfraction^2))
    # test: (2r)^2 - 4 * (r * cutfraction) ^ 2 ≈ 1
    ri = r * (1 - cutfraction)

    first_three = Point2[(r, ri), (ri, ri), (ri, r)]
    all = map(0:pi/2:3pi/2) do a
        m = Mat2f0(sin(a), cos(a), cos(a), -sin(a))
        Ref(m) .* first_three
    end |> x -> reduce(vcat, x)

    BezierPath([
        MoveTo(all[1]),
        LineTo.(all[2:end])...,
        ClosePath()
    ])
end

# pathtokens(::Union{Val{:M}, Val{:m}, Val{:L}, Val{:l}}) = (Float64, Float64)
# pathtokens(::Union{Val{:H}, Val{:h}, Val{:V}, Val{:v}}) = (Float64,)
# pathtokens(::Union{Val{:Z}, Val{:z}}) = ()
# pathtokens(::Union{Val{:C}, Val{:c}}) = (Float64, Float64, Float64, Float64, Float64, Float64)
# pathtokens(::Union{Val{:S}, Val{:s}}) = (Float64, Float64, Float64, Float64)
# pathtokens(::Union{Val{:A}, Val{:a}}) = (Float64, Float64, Float64, Bool, Bool, Float64, Float64)
# # pathtokens

# function BezierPath(svg::AbstractString)
#     n = length(svg)
#     offset = 1

#     commands = PathCommand{Float64}[]

#     v = view(svg, 1:n)

#     while offset < n
#         offset += shift_to_nonseparator(v)
#         v = view(svg, offset:n)
#         command, offset = parsecommand(v)
#         push!(commands, command)
#     end

#     BezierPath(commands)
# end

# function shift_to_nonseparator(substr)
#     for (i, char) in enumerate(substr)
#         if char ∉ " ,\n\t"
#             return i - 1
#         end
#     end
#     return length(substr) - 1
# end

# function parsecommand(substr)
#     commandletters = "MmLlHhVvZzCcSsAa"
#     @show substr

#     command = substr[1]
#     if command ∉ commandletters
#         error("Expected a command letter but got $(substr[1]).")
#     end

#     types = pathtokens(Val(Symbol(command)))
# end

# function parsetypes(types, substr)
#     offset = 1
#     shift = 0
#     map(types) do T
#         offset += shift_to_nonseparator(substr)
#         val, shift = parsetype(T, substr)
#         offset += shift
#         val
#     end
# end

function BezierPath(svg::AbstractString)

    # args = [e.match for e in eachmatch(r"([a-zA-Z])|(\-?\d*\.?\d+)", svg)]
    args = [e.match for e in eachmatch(r"(?:0(?=\d))|(?:[a-zA-Z])|(?:\-?\d*\.?\d+)", svg)]

    i = 1

    commands = PathCommand[]
    lastcomm = nothing
    function lastp()
        c = commands[end]
        if isnothing(lastcomm)
            Point(0, 0)
        elseif c isa ClosePath
            r = reverse(commands)
            backto = findlast(x -> !(x isa ClosePath), r)
            if isnothing(backto)
                error("No point to go back to")
            end
            r[backto].p
        elseif c isa EllipticalArc
            let
                ϕ = c.angle
                a2 = c.a2
                rx = c.r1
                ry = c.r2
                m = Mat2(cos(ϕ), sin(ϕ), -sin(ϕ), cos(ϕ))
                m * Point(rx * cos(a2), ry * sin(a2)) + c.c
            end
        else
            c.p
        end
    end

    while i <= length(args)

        comm = args[i]

        # command letter is omitted, use last command
        if isnothing(match(r"[a-zA-Z]", comm))
            comm = lastcomm
            i -= 1
        end

        if comm == "M"
            x, y = parse.(Float64, args[i+1:i+2])
            push!(commands, MoveTo(Point2(x, y)))
            i += 3
        elseif comm == "m"
            x, y = parse.(Float64, args[i+1:i+2])
            push!(commands, MoveTo(Point2(x, y) + lastp()))
            i += 3
        elseif comm == "L"
            x, y = parse.(Float64, args[i+1:i+2])
            push!(commands, LineTo(Point2(x, y)))
            i += 3
        elseif comm == "l"
            x, y = parse.(Float64, args[i+1:i+2])
            push!(commands, LineTo(Point2(x, y) + lastp()))
            i += 3
        elseif comm == "H"
            x = parse(Float64, args[i+1])
            push!(commands, LineTo(Point2(x, lastp().y)))
            i += 2
        elseif comm == "h"
            x = parse(Float64, args[i+1])
            push!(commands, LineTo(X(x) + lastp()))
            i += 2
        elseif comm == "Z"
            push!(commands, ClosePath())
            i += 1
        elseif comm == "z"
            push!(commands, ClosePath())
            i += 1
        elseif comm == "C"
            x1, y1, x2, y2, x3, y3 = parse.(Float64, args[i+1:i+6])
            push!(commands, CurveTo(Point2(x1, y1), Point2(x2, y2), Point2(x3, y3)))
            i += 7
        elseif comm == "c"
            x1, y1, x2, y2, x3, y3 = parse.(Float64, args[i+1:i+6])
            l = lastp()
            push!(commands, CurveTo(Point2(x1, y1) + l, Point2(x2, y2) + l, Point2(x3, y3) + l))
            i += 7
        elseif comm == "S"
            x1, y1, x2, y2 = parse.(Float64, args[i+1:i+4])
            prev = commands[end]
            reflected = prev.p + (prev.p - prev.c2)
            push!(commands, CurveTo(reflected, Point2(x1, y1), Point2(x2, y2)))
            i += 5
        elseif comm == "s"
            x1, y1, x2, y2 = parse.(Float64, args[i+1:i+4])
            prev = commands[end]
            reflected = prev.p + (prev.p - prev.c2)
            l = lastp()
            push!(commands, CurveTo(reflected, Point2(x1, y1) + l, Point2(x2, y2) + l))
            i += 5
        elseif comm == "A"
            args[i+1:i+7]
            r1, r2 = parse.(Float64, args[i+1:i+2])
            angle = parse(Float64, args[i+3])
            large_arc_flag, sweep_flag = parse.(Bool, args[i+4:i+5])
            x2, y2 = parse.(Float64, args[i+6:i+7])
            x1, y1 = lastp()

            push!(commands, EllipticalArc(x1, y1, x2, y2, r1, r2,
                angle, large_arc_flag, sweep_flag))
            i += 8
        elseif comm == "a"
            r1, r2 = parse.(Float64, args[i+1:i+2])
            angle = parse(Float64, args[i+3])
            large_arc_flag, sweep_flag = parse.(Bool, args[i+4:i+5])
            x1, y1 = lastp()
            x2, y2 = parse.(Float64, args[i+6:i+7]) .+ (x1, y1)

            push!(commands, EllipticalArc(x1, y1, x2, y2, r1, r2,
                angle, large_arc_flag, sweep_flag))
            i += 8
        elseif comm == "v"
            dy = parse(Float64, args[i+1])
            l = lastp()
            push!(commands, LineTo(Point2(l[1], l[2] + dy)))
            i += 2
        else
            for c in commands
                println(c)
            end
            error("Parsing $comm not implemented.")
        end

        lastcomm = comm

    end

    BezierPath(commands)

end

function EllipticalArc(x1, y1, x2, y2, rx, ry, ϕ, largearc::Bool, sweepflag::Bool)
    # https://www.w3.org/TR/SVG11/implnote.html#ArcImplementationNotes

    p1 = Point(x1, y1)
    p2 = Point(x2, y2)

    m1 = Mat2(cos(ϕ), -sin(ϕ), sin(ϕ), cos(ϕ))
    x1′, y1′ = m1 * (0.5 * (p1 - p2))

    tempsqrt = (rx^2 * ry^2 - rx^2 * y1′^2 - ry^2 * x1′^2) /
        (rx^2 * y1′^2 + ry^2 * x1′^2)

    c′ = (largearc == sweepflag ? -1 : 1) *
        sqrt(tempsqrt) * Point(rx * y1′ / ry, -ry * x1′ / rx)

    c = Mat2(cos(ϕ), sin(ϕ), -sin(ϕ), cos(ϕ)) * c′ + 0.5 * (p1 + p2)

    vecangle(u, v) = sign(u[1] * v[2] - u[2] * v[1]) *
        acos(dot(u, v) / (norm(u) * norm(v)))

    px(sign) = Point((sign * x1′ - c′[1]) / rx, (sign * y1′ - c′[2]) / rx)

    θ1 = vecangle(Point(1.0, 0.0), px(1))
    Δθ_pre = mod(vecangle(px(1), px(-1)), 2pi)
    Δθ = if Δθ_pre > 0 && !sweepflag
        Δθ_pre - 2pi
    elseif Δθ_pre < 0 && sweepflag
        Δθ_pre + 2pi
    else
        Δθ_pre
    end

    EllipticalArc(c, rx, ry, ϕ, θ1, θ1 + Δθ)
end


text:
"""
struct containing grids used in a simulation

# Examples

```jldoctest
julia> using UltraDark

julia> len = 1;

julia> resol = 16;

julia> PencilGrids(len, resol);

```
"""
struct PencilGrids{K, RX, RK, CX, CK, FFT, RFFT, M}
    "Array of x positions"
    x::Array{Float64,3}
    "Array of y positions"
    y::Array{Float64,3}
    "Array of z positions"
    z::Array{Float64,3}
    "Array of x Fourier modes"
    kx::Array{Float64,3}
    "Array of y Fourier modes"
    ky::Array{Float64,3}
    "Array of z Fourier modes"
    kz::Array{Float64,3}
    "Fourier space postition array"
    k::K
    "Array of x Fourier modes for use with `rfft`"
    rkx::Array{Float64,3}
    "Array of y Fourier modes for use with `rfft`"
    rky::Array{Float64,3}
    "Array of z Fourier modes for use with `rfft`"
    rkz::Array{Float64,3}
    "Fourier space postition array for use with `rfft`"
    rk::K
    "ψ field"
    ψx::CX
    "ψ field in Fourier space"
    ψk::CK
    "density field ρ"
    ρx::RX
    "density field ρ in Fourier space"
    ρk::RK
    "gravitational potential field Φ"
    Φx::RX
    "gravitational potential field Φ in fourier space"
    Φk::RK

    "FFT plan for complex-to-complex transforms"
    fft_plan::FFT
    "FFT plan for real-to-complex transforms"
    rfft_plan::RFFT

    "Indices at which k==0.0"
    k_vanish_indices::Vector{CartesianIndex{3}}
    "Indices at which rk==0.0"
    rk_vanish_indices#::Vector{CartesianIndex{3}}

    "MPI communicator"
    MPI_COMM::M

    function PencilGrids(x, y, z, kx, ky, kz, k, rkx, rky, rkz, rk, ψx, ψk, ρx, ρk, Φx, Φk, fft_plan, rfft_plan, MPI_COMM)
        n_dims = 3
        resol_tuple = (size(x)[1], size(y)[2], size(z)[3])
        resol_tuple_realfft = (size(x)[1] ÷ 2 + 1, size(y)[2], size(z)[3])

        for var in [x, y, z, k, rk, ψx, ψk, ρx, ρk, Φx, Φk]
            @assert(ndims(var) == n_dims)
        end

        for var in [k, ψx, ψk, ρx, Φx]
            @assert(size_global(var) == resol_tuple)
        end

        for var in [rk, ρk, Φk]
            @assert(size_global(var) == resol_tuple_realfft)
        end

        @assert(size(x) == (resol_tuple[1], 1, 1))
        @assert(size(y) == (1, resol_tuple[2], 1))
        @assert(size(z) == (1, 1, resol_tuple[3]))

        k_vanish_indices = findall(x -> x==0.0, k)
        rk_vanish_indices = findall(x -> x==0.0, rk)

        K = typeof(k)
        RX = typeof(Φx)
        RK = typeof(Φk)
        CX = typeof(ψx)
        CK = typeof(ψk)
        FFT = typeof(fft_plan)
        RFFT = typeof(rfft_plan)
        M = typeof(MPI_COMM)

        new{K, RX, RK, CX, CK, FFT, RFFT, M}(x, y, z, kx, ky, kz, k, rkx, rky, rkz, rk, ψx, ψk, ρx, ρk, Φx, Φk, fft_plan, rfft_plan, k_vanish_indices, rk_vanish_indices, MPI_COMM)
    end
end

"""
    PencilGrids(length, resol)

Constructor for `PencilGrids`

Create an empty grid with length `length` and resolution `resol`.  Uses `PencilFFTs` to create `PencilArrays`.

# Examples

```jldoctest
julia> using UltraDark

julia> PencilGrids(1.0, 64);

```
"""
function PencilGrids(length, resol::Int)::PencilGrids

    PencilGrids((length, length, length), (resol, resol, resol))

end

"""
    PencilGrids(length_tuple, resol_tuple::Tuple{Int, Int, Int})

Constructor for `PencilGrids`

Create an empty `length[1]`x`length[2]`x`length[3]` grid with resolution
`resol[1]`x`resol[2]`x`resol[3]`.
Each grid is a `PencilArray`, allowing multiprocess FFTs.

# Examples

```jldoctest
julia> using UltraDark

julia> PencilGrids((1.0, 1.0, 0.5), (64, 64, 32));

```
"""
function PencilGrids(length_tuple, resol_tuple::Tuple{Int, Int, Int})::PencilGrids

    resol_tuple_realfft = (resol_tuple[1] ÷ 2 + 1, resol_tuple[2], resol_tuple[3])

    @assert length_tuple[1] / resol_tuple[1] ≈ length_tuple[2] / resol_tuple[2]
    @assert length_tuple[1] / resol_tuple[1] ≈ length_tuple[3] / resol_tuple[3]

    if ~MPI.Initialized()
        MPI.Init()
    end

    # MPI topology information
    comm = MPI.COMM_WORLD
    Nproc = MPI.Comm_size(comm)

    # Let MPI_Dims_create choose the decomposition.
    proc_dims = let pdims = zeros(Int, 2)
        MPI.Dims_create!(Nproc, pdims)
        pdims[1], pdims[2]
    end

    FFTW.set_num_threads(Threads.nthreads())
    # Plan a 3D complex-to-complex (c2c) FFT.
    fft_plan = PencilFFTPlan(resol_tuple, Transforms.FFT(), proc_dims, comm)

    # Allocate ψx and ψk arrays
    ψx = allocate_input(fft_plan)
    ψx .= 0
    ψk = allocate_output(fft_plan)

    # Plan a 3D real-to-complex (r2c) FFT.
    rfft_plan = PencilFFTPlan(resol_tuple, Transforms.RFFT(), proc_dims, comm)

    ρx = allocate_input(rfft_plan)
    ρk = allocate_output(rfft_plan)

    Φx = allocate_input(rfft_plan)
    Φk = allocate_output(rfft_plan)

    x = reshape(
                range(
                      -length_tuple[1] / 2 + length_tuple[1] / 2resol_tuple[1],
                      +length_tuple[1] / 2 - length_tuple[1] / 2resol_tuple[1],
                      length=resol_tuple[1]
                     ),
                resol_tuple[1], 1, 1
    )

    y = reshape(
                range(
                      -length_tuple[2] / 2 + length_tuple[2] / 2resol_tuple[2],
                      +length_tuple[2] / 2 - length_tuple[2] / 2resol_tuple[2],
                      length=resol_tuple[2]
                     ),
                1, resol_tuple[2], 1
    )

    z = reshape(
                range(
                      -length_tuple[3] / 2 + length_tuple[3] / 2resol_tuple[3],
                      +length_tuple[3] / 2 - length_tuple[3] / 2resol_tuple[3],
                      length=resol_tuple[3]
                     ),
                1, 1, resol_tuple[3]
    )

    kvec = k_vec(length_tuple, resol_tuple)

    kx = reshape(kvec[1], size(kvec[1])[1], 1, 1)
    ky = reshape(kvec[2], 1, size(kvec[2])[1], 1)
    kz = reshape(kvec[3], 1, 1, size(kvec[3])[1])
    k_norm = (kx.^2 .+ ky.^2 .+ kz.^2).^0.5  #TODO: don't allocate full array on every node

    k = similar(ψk, Float64)
    k_glob = global_view(k)
    _k = k_norm
    for I in CartesianIndices(k_glob)
        k_glob[I] = _k[I]
    end

    rkvec = rk_vec(length_tuple, resol_tuple)

    rkx = reshape(rkvec[1], size(rkvec[1])[1], 1, 1)
    rky = reshape(rkvec[2], 1, size(rkvec[2])[1], 1)
    rkz = reshape(rkvec[3], 1, 1, size(rkvec[3])[1])
    rk_norm = (rkx.^2 .+ rky.^2 .+ rkz.^2).^0.5

    rk = similar(ρk, Float64)
    rk_glob = global_view(rk)
    _rk = rk_norm
    for I in CartesianIndices(rk_glob)
        rk_glob[I] = _rk[I]
    end

    PencilGrids(
        x, y, z,
        kx, ky, kz, k,
        rkx, rky, rkz, rk,
        ψx, ψk,
        ρx, ρk,
        Φx, Φk,
        fft_plan, rfft_plan,
        comm
    )
end


text:
"""
`annualtoquarter(v)`

Convert from annual to quarter frequency... by dividing by 4.
"""
function annualtoquarter(v)
    v / 4
end

"""
`quartertoannual(v)`

Convert from quarter to annual frequency... by multiplying by 4.
"""
function quartertoannual(v)
    4 * v
end

"""
`quartertoannualpercent(v)`

Convert from quarter to annual frequency in percent... by multiplying by 400.
"""
function quartertoannualpercent(v)
    400 * v
end

"""
`nominal_to_real(col, df; deflator_mnemonic = :GDPDEF)`

Converts nominal to real values using the specified deflator.

## Arguments
- `col`: Symbol indicating which column of `df` to transform
- `df`: DataFrame containining series for proper population measure and `col`

## Keyword arguments

- `deflator_mnemonic`: indicates which deflator to use to calculate real values. Default
  value is the FRED GDP Deflator mnemonic.
"""
function nominal_to_real(col::Symbol, df::DataFrame; deflator_mnemonic::Symbol = :GDPDEF)
    return df[col] ./ df[deflator_mnemonic]
end

"""
```
percapita(m, col, df)
percapita(col, df, population_mnemonic)
```

Converts data column `col` of DataFrame `df` to a per-capita value.

The first method checks `hpfilter_population(m)`. If true, then it divides by
the filtered population series. Otherwise it divides by the result of
`parse_population_mnemonic(m)[1]`.

## Arguments

- `col`: `Symbol` indicating which column of data to transform
- `df`: `DataFrame` containining series for proper population measure and `col`
- `population_mnemonic`: a mnemonic found in `df` for some population measure
"""
function percapita(m::AbstractModel, col::Symbol, df::DataFrame)
    if hpfilter_population(m)
        population_mnemonic = Nullable(:filtered_population)
    else
        population_mnemonic = parse_population_mnemonic(m)[1]
        if isnull(population_mnemonic)
            error("No population mnemonic provided")
        end
    end
    percapita(col, df, get(population_mnemonic))
end

function percapita(col::Symbol, df::DataFrame, population_mnemonic::Symbol)
    df[col] ./ df[population_mnemonic]
end

"""
```
series_lag_n = lag(series, n)
```
Returns a particular data series lagged by n periods
"""
function lag(series::AbstractArray, n::Int)
    return vcat(NaN*ones(n), series[1:end-n])
end

"""
```
yt, yf = hpfilter(y, λ)
```

Applies the Hodrick-Prescott filter (\"H-P filter\"). The smoothing parameter `λ` is applied
to the columns of `y`, returning the trend component `yt` and the cyclical component `yf`.
For quarterly data, one can use λ=1600.

Consecutive missing values at the beginning or end of the time series are excluded from the
filtering. If there are missing values within the series, the filtered values are all NaN.

See also:
```
Hodrick, Robert; Prescott, Edward C. (1997). \"Postwar U.S. Business Cycles: An Empirical
Investigation\". Journal of Money, Credit, and Banking 29 (1): 1–16.
```
"""
function hpfilter(y::AbstractVector, λ::Real)
    # Indices of consecutive NaN elements at beginning
    i = 1
    j = length(y)
    while isnan(y[i])
        i = i+1
    end
    while isnan(y[j])
        j = j-1
    end

    # Filter and adjust for NaNs
    yt_, yf_ = hpfilter_(y[i:j], λ)
    yt = [fill(NaN, i-1); yt_; fill(NaN, length(y)-j)]
    yf = [fill(NaN, i-1); yf_; fill(NaN, length(y)-j)]

    return yt, yf
end

function hpfilter_(y::AbstractVector, λ::Real)
    n = length(y)
    a = spzeros(n,n)
    for i = 3:n-2
        a[i,i]   = 6λ+1
        a[i,i+1] = -4λ
        a[i,i+2] = λ
        a[i,i-2] = λ
        a[i,i-1] = -4λ
    end

    a[2,2] = 1+5λ
    a[2,3] = -4λ
    a[2,4] = λ
    a[2,1] = -2λ
    a[1,1] = 1+λ
    a[1,2] = -2λ
    a[1,3] = λ

    a[n-1,n-1] = 1+5λ
    a[n-1,n-2] = -4λ
    a[n-1,n-3] = λ
    a[n-1,n]   = -2λ
    a[n,n]     = 1+λ
    a[n,n-1]   = -2λ
    a[n,n-2]   = λ

    yt = a\y
    yf = y-yt

    return yt, yf
end

"""
```
difflog(x::AbstractVector)
```
"""
function difflog(x::Vector)
    [NaN; log.(x[2:end]) - log.(x[1:end-1])]
end

"""
```
difflog(x::DataArray{AbstractFloat})
```
"""
function difflog(x::DataArray)
    DSGE.na2nan!(x)
    return difflog(convert(Vector, x))
end

"""
```
oneqtrpctchange(y)
```

Calculates the quarter-to-quarter percentage change of a series.
"""
function oneqtrpctchange(y::AbstractVector)
    100 * difflog(y)
end

## REVERSE TRANSFORMS

"""
```
loggrowthtopct(y)
```

Transform from annualized quarter-over-quarter log growth rates to annualized
quarter-over-quarter percent change.

### Note

This should only be used in Model 510, which has the core PCE inflation
observable in annualized log growth rates.
"""
function loggrowthtopct(y::AbstractArray)
    100. * (exp.(y/100.) - 1.)
end

"""
```
loggrowthtopct_percapita(y, pop_growth)
```

Transform from annualized quarter-over-quarter log per-capita growth rates to
annualized quarter-over-quarter aggregate percent change.

### Note

This should only be used in Model 510, which has the output growth observable in
annualized log per-capita growth rates.

### Inputs

- `y`: the data we wish to transform to annualized percent change from
  annualized log growth rates. `y` is either a vector of length `nperiods` or an
  `ndraws x `nperiods` matrix.

- `pop_growth::Vector`: the length `nperiods` vector of log population growth
  rates.
"""
function loggrowthtopct_percapita(y::AbstractArray, pop_growth::AbstractVector)
    # `y` is either a vector of length `nperiods` or an
    # `ndraws` x `nperiods` matrix
    if ndims(y) == 1
        nperiods = length(y)
    else
        nperiods = size(y, 2)

        # Transpose `pop_growth` to a 1 x `nperiods` row vector so it can be
        # broadcasted to match the dimensions of `y`
        pop_growth = pop_growth'
    end

    @assert length(pop_growth) == nperiods "Length of pop_growth ($(length(pop_growth))) must equal number of periods of y ($nperiods)"

    100. * ((exp.(y/100.) .* exp.(pop_growth).^4) - 1.)
end

"""
```
loggrowthtopct_annualized(y)
```

Transform from log growth rates to annualized quarter-over-quarter percent change.
"""
function loggrowthtopct_annualized(y::AbstractArray)
    100. * (exp.(y/100.).^4 - 1.)
end

"""
```
loggrowthtopct_annualized_percapita(y, pop_growth)
```

Transform from log per-capita growth rates to annualized aggregate (not
per-capita) quarter-over-quarter percent change.

### Note

This should only be used for output, consumption, investment
and GDP deflator (inflation).

### Inputs

- `y`: the data we wish to transform to annualized percent change from
  quarter-over-quarter log growth rates. `y` is either a vector of length
  `nperiods` or an `ndraws x `nperiods` matrix.

- `pop_growth::Vector`: the length `nperiods` vector of log population growth
  rates.
"""
function loggrowthtopct_annualized_percapita(y::AbstractArray, pop_growth::AbstractVector)
    # `y` is either a vector of length `nperiods` or an
    # `ndraws` x `nperiods` matrix
    if ndims(y) == 1
        nperiods = length(y)
    else
        nperiods = size(y, 2)

        # Transpose `pop_growth` to a 1 x `nperiods` row vector so it can be
        # broadcasted to match the dimensions of `y`
        pop_growth = pop_growth'
    end

    @assert length(pop_growth) == nperiods "Length of pop_growth ($(length(pop_growth))) must equal number of periods of y ($nperiods)"

    100. * (exp.(y/100. .+ pop_growth).^4 - 1.)
end

"""
```
logleveltopct_annualized(y, y0 = NaN)
```

Transform from log levels to annualized quarter-over-quarter percent change.

### Inputs

- `y`: the data we wish to transform to annualized quarter-over-quarter percent
  change from log levels. `y` is either a vector of length `nperiods` or an
  `ndraws x `nperiods` matrix.

- `y0`: the last data point in the history (of state or observable)
  corresponding to the `y` variable. This is required to compute a percent
  change for the first period.
"""
function logleveltopct_annualized(y::AbstractArray, y0::Real = NaN)
    # `y_t1` is an array of the same size as `y`, representing the previous
    # period observations for each draw
    if ndims(y) == 1
        y_t1 = vcat([y0], y[1:end-1])
    else
        ndraws = size(y, 1)
        y0s  = fill(y0, ndraws, 1)
        y_t1 = hcat(y0s, y[:, 1:end-1])
    end

    # Subtract log levels to get log growth rates, then take the exponential to
    # get growth rates
    100. * (exp.(y./100. - y_t1./100.).^4 .- 1.)
end

"""
```
logleveltopct_annualized_percapita(y, pop_growth, y0 = NaN)
```

Transform from per-capita log levels to annualized aggregate (not per-capita)
quarter-over-quarter percent change.

### Note

This is usually applied to labor supply (hours worked per hour), and
probably shouldn't be used for any other observables.

### Inputs

- `y`: the data we wish to transform to annualized aggregate
  quarter-over-quarter percent change from per-capita log levels. `y` is either
  a vector of length `nperiods` or an `ndraws x `nperiods` matrix.

- `pop_growth::Vector`: the length `nperiods` vector of log population growth
  rates.

- `y0`: The last data point in the history (of state or observable)
  corresponding to the `y` variable. This is required to compute a percent
  change for the first period.
"""
function logleveltopct_annualized_percapita(y::AbstractArray, pop_growth::AbstractVector, y0::Real = NaN)
    # `y_t1` is an array of the same size as `y`, representing the previous
    # period observations for each draw
    if ndims(y) == 1
        nperiods = length(y)
        y_t1 = vcat([y0], y[1:end-1])
    else
        (ndraws, nperiods) = size(y)
        y0s  = fill(y0, ndraws, 1)
        y_t1 = hcat(y0s, y[:, 1:end-1])

        # Transpose `pop_growth` to a 1 x `nperiods` row vector so it can be
        # broadcasted to match the dimensions of `y`
        pop_growth = pop_growth'
    end

    @assert length(pop_growth) == nperiods "Length of pop_growth ($(length(pop_growth))) must equal number of periods of y ($nperiods)"

    # Subtract log levels to get log growth rates, then take the exponential to
    # get growth rates
    100. * (exp.(y./100. - y_t1./100. .+ pop_growth).^4 .- 1.)
end

"""
```
get_nopop_transform(transform::Function)
```
Returns the corresponding transformation which doesn't add back population
growth. Used for shock decompositions, deterministic trends, and IRFs, which are
given in deviations.
"""
function get_nopop_transform(transform::Function)
    transform4q = if transform == loggrowthtopct_annualized_percapita
        loggrowthtopct_annualized
    elseif transform == logleveltopct_annualized_percapita
        logleveltopct_annualized
    else
        transform
    end
end

"""
```
get_irf_transform(transform::Function)
```
Returns the IRF-specific transformation, which doesn't add back population
growth (since IRFs are given in deviations).
"""
function get_irf_transform(transform::Function)
    transform4q = if transform == loggrowthtopct_annualized_percapita
        loggrowthtopct_annualized
    elseif transform == logleveltopct_annualized_percapita
        logleveltopct_annualized
    else
        transform
    end
end

"""
```
get_transform4q(transform::Function)
```
Returns the 4-quarter transformation associated with the annualizing transformation.
"""
function get_transform4q(transform::Function)
    if transform == loggrowthtopct_annualized_percapita
        loggrowthtopct_4q_percapita
    elseif transform == loggrowthtopct_annualized
        loggrowthtopct_4q
    elseif transform == logleveltopct_annualized_percapita
        logleveltopct_4q_percapita
    elseif transform == logleveltopct_annualized
        logleveltopct_4q
    elseif transform == quartertoannual
        quartertoannual
    elseif transform == identity
        identity
    else
        error("4q equivalent not implemented for $transform")
    end
end

"""
```
loggrowthtopct_4q(y, data = fill(NaN, 3))
```

Transform from log growth rates to 4-quarter percent change.

### Inputs

- `y`: the data we wish to transform to aggregate 4-quarter percent change from
  log per-capita growth rates. `y` is either a vector of length `nperiods` or an
  `ndraws x `nperiods` matrix.

- `data`: if `y = [y_t, y_{t+1}, ..., y_{t+nperiods-1}]`, then
  `data = [y_{t-3}, y_{t-2}, y_{t-1}]`. This is necessary to compute
  4-quarter percent changes for the first three periods.
"""
function loggrowthtopct_4q(y::AbstractArray, data::AbstractVector = fill(NaN, 3))
    @assert length(data) == 3 "Length of data ($(length(data))) must be 3"

    # Prepend previous three periods to `y`
    y = prepend_data(y, data)

    # `y` is either a vector of length `nperiods+3` or an
    # `ndraws` x `nperiods+3` matrix
    if ndims(y) == 1
        y_4q = y[1:end-3] + y[2:end-2] + y[3:end-1] + y[4:end]
    else
        y_4q = y[:,  1:end-3] + y[:, 2:end-2] + y[:, 3:end-1] + y[:, 4:end]
    end

    100. * (exp.(y_4q/100.) - 1.)
end

"""
```
loggrowthtopct_4q_percapita(y, pop_growth, data = fill(NaN, 3))
```
Transform from log per-capita growth rates to aggregate 4-quarter percent
change.

### Note

This should only be used for output, consumption, investment, and GDP deflator
(inflation).

### Inputs

- `y`: the data we wish to transform to aggregate 4-quarter percent change from
  log per-capita growth rates. `y` is either a vector of length `nperiods` or an
  `ndraws x `nperiods` matrix.

- `pop_growth::Vector`: the length `nperiods` vector of log population growth
  rates.

- `data`: if `y = [y_t, y_{t+1}, ..., y_{t+nperiods-1}]`, then
  `data = [y_{t-3}, y_{t-2}, y_{t-1}]`. This is necessary to compute
  4-quarter percent changes for the first three periods.
"""
function loggrowthtopct_4q_percapita(y::AbstractArray, pop_growth::AbstractVector,
                                     data::Vector = fill(NaN, 3))
    @assert length(data) == 3 "Length of data ($(length(data))) must be 3"

    # Four-quarter population growth
    pop_growth_4q = pop_growth[1:end-3] + pop_growth[2:end-2] + pop_growth[3:end-1] + pop_growth[4:end]

    # Prepend previous three periods to `y`
    y = prepend_data(y, data)

    # `y` is either a vector of length `nperiods+3` or an
    # `ndraws` x `nperiods+3` matrix
    if ndims(y) == 1
        y_4q = y[1:end-3] + y[2:end-2] + y[3:end-1] + y[4:end]
        nperiods = length(y_4q)
    else
        y_4q = y[:, 1:end-3] + y[:, 2:end-2] + y[:, 3:end-1] + y[:, 4:end]
        nperiods = size(y_4q, 2)

        # Transpose `pop_growth` to a 1 x `nperiods` row vector so it can be
        # broadcasted to match the dimensions of `y_4q`
        pop_growth_4q = pop_growth_4q'
    end

    @assert length(pop_growth_4q) == nperiods "Length of pop_growth ($(length(pop_growth))) must equal number of periods of y ($nperiods)"

    100. * (exp.(y_4q/100. .+ pop_growth_4q) - 1.)
end

"""
```
logleveltopct_4q(y, data = fill(NaN, 4))
```

Transform from log levels to 4-quarter percent change.

### Inputs

- `y`: the data we wish to transform to 4-quarter percent change from log
  levels. `y` is either a vector of length `nperiods` or an `ndraws x `nperiods`
  matrix.

- `data`: if `y = [y_t, y_{t+1}, ..., y_{t+nperiods-1}]`, then
  `data = [y_{t-4}, y_{t-3}, y_{t-2}, y_{t-1}]`. This is necessary to compute
  4-quarter percent changes for the first three periods.
"""
function logleveltopct_4q(y::AbstractArray, data::AbstractVector = fill(NaN, 4))
    @assert length(data) == 4 "Length of data ($(length(data))) must be 4"

    # `y_t4` is an array of the same size as `y`, representing the t-4
    # period observations for each t
    y_t4 = if ndims(y) == 1
        nperiods = length(y)
        prepend_data(y[1:nperiods-4], data)
    else
        nperiods = size(y, 2)
        prepend_data(y[:, 1:nperiods-4], data)
    end
    y_4q = y - y_t4

    # Subtract log levels to get log growth rates, then exponentiate to get
    # growth rates
    100. * (exp.(y_4q./100.) .- 1.)
end

"""
```
logleveltopct_4q_percapita(y, pop_growth, data = fill(NaN, 4))
```

Transform from per-capita log levels to 4-quarter aggregate percent change.

### Note

This is usually applied to labor supply (hours worked), and probably shouldn't
be used for any other observables.

### Inputs

- `y`: the data we wish to transform to 4-quarter aggregate percent change from
  per-capita log levels. `y` is either a vector of length `nperiods` or an
  `ndraws x `nperiods` matrix.

- `pop_growth::Vector`: the length `nperiods` vector of log population growth
  rates.

- `data`: if `y = [y_t, y_{t+1}, ..., y_{t+nperiods-1}]`, then
  `data = [y_{t-4}, y_{t-3}, y_{t-2}, y_{t-1}]`. This is necessary to compute
  4-quarter percent changes for the first three periods.
"""
function logleveltopct_4q_percapita(y::AbstractArray, pop_growth::AbstractVector,
                                    data::AbstractVector = fill(NaN, 4))
    @assert length(data) == 4 "Length of data ($(length(data))) must be 4"

    # Four-quarter population growth
    pop_growth_4q = pop_growth[1:end-3] + pop_growth[2:end-2] + pop_growth[3:end-1] + pop_growth[4:end]

    # `y_t4` is an array of the same size as `y`, representing the t-4
    # period observations for each t
    if ndims(y) == 1
        nperiods = length(y)
        y_t4 = prepend_data(y[1:nperiods-4], data)
    else
        # Transpose `pop_growth` to a 1 x `nperiods` row vector so it can be
        # broadcasted to match the dimensions of `y`
        pop_growth_4q = pop_growth_4q'

        nperiods = size(y, 2)
        y_t4 = prepend_data(y[:, 1:nperiods-4], data)
    end

    @assert length(pop_growth_4q) == nperiods "Length of pop_growth ($(length(pop_growth))) must equal number of periods of y ($nperiods)"

    y_4q = y - y_t4

    # Subtract log levels to get log growth rates, then exponentiate to get growth rates
    100. * (exp.(y_4q./100. .+ pop_growth_4q) .- 1.)
end

"""
```
prepend_data(y, data)
```

Prepends data necessary for running 4q transformations.

### Inputs:

- `y`: `ndraws x t` array representing a timeseries for variable `y`
- `data`: vector representing a timeseries to prepend to `y`
"""
function prepend_data(y::AbstractArray, data::AbstractVector)
    if ndims(y) == 1
        y_extended = vcat(data, y)
    else
        ndraws = size(y, 1)
        datas  = repmat(data', ndraws, 1)
        y_extended = hcat(datas, y)
    end

    return y_extended
end

"""
```
get_scenario_transform(transform::Function)
```
Given a transformation used for usual forecasting, return the transformation
used for *scenarios*, which are forecasted in deviations from baseline.

The 1Q deviation from baseline should really be calculated by 1Q transforming
the forecasts (in levels) under the baseline (call this `y_b`) and alternative
scenario (`y_s`), then subtracting baseline from alternative scenario (since
most of our 1Q transformations are nonlinear). Let `y_d = y_s - y_b`. Then, for
example, the most correct `loggrowthtopct_annualized` transformation is:

```
y_b_1q = 100*(exp(y_b/100)^4 - 1)
y_s_1q = 100*(exp(y_s/100)^4 - 1)
y_d_1q = y_b_1q - y_s_1q
```

Instead, we approximate this by transforming the deviation directly:

```
y_d_1q ≈ 4*(y_b - y_s)
```
"""
function get_scenario_transform(transform::Function)
    if transform in [loggrowthtopct_annualized_percapita, loggrowthtopct_annualized,
                     quartertoannual]
        quartertoannual
    elseif transform in [logleveltopct_annualized_percapita, logleveltopct_annualized]
        logleveltopct_annualized_approx
    elseif transform == identity
        identity
    else
        error("Scenario equivalent not implemented for $transform")
    end
end

function get_scenario_transform4q(transform::Function)
    if transform in [loggrowthtopct_annualized_percapita, loggrowthtopct_annualized]
        loggrowthtopct_4q_approx
    elseif transform in [logleveltopct_annualized_percapita, logleveltopct_annualized]
        logleveltopct_4q_approx
    elseif transform == quartertoannual
        quartertoannual
    elseif transform == identity
        identity
    else
        error("Scenario 4q equivalent not implemented for $transform")
    end
end

"""
```
logleveltopct_annualized_approx(y, y0 = NaN)
```

Transform from log levels to *approximate* annualized quarter-over-quarter
percent change.

**This method should only be used to transform scenarios forecasts, which are in
  deviations from baseline.**

### Inputs

- `y`: the data we wish to transform to annualized quarter-over-quarter percent
  change from log levels. `y` is either a vector of length `nperiods` or an
  `ndraws x `nperiods` matrix.

- `y0`: the last data point in the history (of state or observable)
  corresponding to the `y` variable. This is required to compute a percent
  change for the first period.
"""
function logleveltopct_annualized_approx(y::AbstractArray, y0::Real = NaN)
    # `y_t1` is an array of the same size as `y`, representing the previous
    # period observations for each draw
    if ndims(y) == 1
        y_t1 = vcat([y0], y[1:end-1])
    else
        ndraws = size(y, 1)
        y0s  = fill(y0, ndraws, 1)
        y_t1 = hcat(y0s, y[:, 1:end-1])
    end

    # Subtract log levels to get log growth rates, then multiply by 4 to
    # approximate annualizing
    4(y - y_t1)
end

"""
```
loggrowthtopct_4q_approx(y, data = fill(NaN, 3))
```

Transform from log growth rates to *approximate* 4-quarter percent change.

**This method should only be used to transform scenarios forecasts, which are in
  deviations from baseline.**

### Inputs

- `y`: the data we wish to transform to aggregate 4-quarter percent change from
  log per-capita growth rates. `y` is either a vector of length `nperiods` or an
  `ndraws x `nperiods` matrix.

- `data`: if `y = [y_t, y_{t+1}, ..., y_{t+nperiods-1}]`, then
  `data = [y_{t-3}, y_{t-2}, y_{t-1}]`. This is necessary to compute
  4-quarter percent changes for the first three periods.
"""
function loggrowthtopct_4q_approx(y::AbstractArray, data::AbstractVector = fill(NaN, 3))
    @assert length(data) == 3 "Length of data ($(length(data))) must be 3"

    # Prepend previous three periods to `y`
    y = prepend_data(y, data)

    # `y` is either a vector of length `nperiods+3` or an
    # `ndraws` x `nperiods+3` matrix
    if ndims(y) == 1
        y[1:end-3] + y[2:end-2] + y[3:end-1] + y[4:end]
    else
        y[:,  1:end-3] + y[:, 2:end-2] + y[:, 3:end-1] + y[:, 4:end]
    end
end

"""
```
logleveltopct_4q_approx(y, data = fill(NaN, 4))
```

Transform from log levels to *approximate* 4-quarter percent change.

**This method should only be used to transform scenarios forecasts, which are in
  deviations from baseline.**

### Inputs

- `y`: the data we wish to transform to 4-quarter percent change from log
  levels. `y` is either a vector of length `nperiods` or an `ndraws x `nperiods`
  matrix.

- `data`: if `y = [y_t, y_{t+1}, ..., y_{t+nperiods-1}]`, then
  `data = [y_{t-4}, y_{t-3}, y_{t-2}, y_{t-1}]`. This is necessary to compute
  4-quarter percent changes for the first three periods.
"""
function logleveltopct_4q_approx(y::AbstractArray, data::AbstractVector = fill(NaN, 4))
    @assert length(data) == 4 "Length of data ($(length(data))) must be 4"

    # `y_t4` is an array of the same size as `y`, representing the t-4
    # period observations for each t
    y_t4 = if ndims(y) == 1
        nperiods = length(y)
        prepend_data(y[1:nperiods-4], data)
    else
        nperiods = size(y, 2)
        prepend_data(y[:, 1:nperiods-4], data)
    end

    y - y_t4
end


text:
function reallocation!(state)
    len = length(state)
    index = argmax(state)
    value = state[index]
    state[index] = 0
    for i in 1:value
        target_index = index + i
        if target_index > len
            target_index = target_index % len
        end
        state[target_index] += 1
    end
end

function solve(state)
    # part a
    count = 1
    know_states = Set([copy(state)])

    while true
        reallocation!(state)
        if state ∈ know_states
            break
        end
        push!(know_states, copy(state))
        count += 1
    end
    sol_a = count
    # part b
    count = 1
    target_state = copy(state)
    while true
        reallocation!(state)
        if state == target_state
            break
        end
        count += 1
    end

    sol_b = count
    sol_a, sol_b
end

data = parse.(Int, split(readline("input.txt")));
sol_a, sol_b = solve(data)
println("sol a: $(sol_a)")
println("sol b: $(sol_b)")


text:
# DOCOPT faster than argparse, but a little sensitive to spacing
doc = """We use this to train Policy Gradients. UW CSE Motion Control Lab.

Usage:
   PGTrain.jl [(<experiment> <name>)]
   PGTrain.jl -h | --help

Options:
   -h --help   Show this screen.

"""

using DocOpt
opt = docopt(doc, ARGS)

push!(LOAD_PATH, "./")

using Common: plotlines
using Random
using Statistics

prefix="/tmp/"
modeldir="$(pwd())/models"
meshfile = "$(modeldir)/meshes/" # hack for now...

expfile=opt["<experiment>"]
expname=opt["<name>"]
#polfile=opt["load"]
#cluster=opt["par"]
#myseed =opt["seed"]; if typeof(myseed) != Int myseed = parse(Int, myseed) end
const SEED = 12345

Random.seed!(SEED) # random seed, set globally
push!(LOAD_PATH, "$(pwd())")

println("Experiment called $(expname)")
println("Loading experiment parameters from $(expfile)")

################################################### start work by loading files
include("paths.jl")

using Policy
using PolicyGradient
using Baseline
using ExpFunctions
using LearningStrategies

const METHOD = :NPG
include("$(pwd())/$(expfile)") # LOADING TASK FILE with PARAMETERS
#@assert isdefined(:pgmodel)
#@assert isdefined(:pg_specs)
#@assert isdefined(:myfuncs)

pgmodel, pg_specs = getNPG()

# plot stoc & mean every other iter
# TODO General strategies; put somewhere else?
evaluatepolicy = IterFunction((model,i)->push!(model.trace[:evalscore],
                                               myfuncs.evaluate(model.mjsys, model.meansamples)))

const NITER = 4
plotR = IterFunction(NITER, (model,i)->plotlines(i,"Reward",
                                                 (model.trace[:stocR],"Stoc"),
                                                 (model.trace[:meanR],"Mean")))

plotEval = IterFunction(NITER, (model,i)->plotlines(i,"Evaluation",
                                                    (model.trace[:evalscore],"")))

#plotCTRL = IterFunction(NITER, (model,i)->
#                        begin
#                           _, idx = findmax(sum(model.meansamples.reward, dims=1))
#                           plotlines(i,"Ctrl", model.meansamples.ctrl[:,:,idx])
#                        end )
#plotCTRL2 = IterFunction(NITER, (model,i)->plotlines(i,"Ctrl Stoc", model.samples.ctrl[:,:,1]))

# meta strategies: order matters!!
pg_strat = strategy(pg_specs,              # main loops for NPG

                    PolicyExp(prefix, expfile,
                              expname, meshfile), # setup some things, hook for save

                    evaluatepolicy,          # hooks
                    plotR,
                    plotEval,
                    #plotCTRL,
                    #plotCTRL2,
                    Verbose(MaxIter(pgmodel.niter)))

@time learn!(pgmodel, pg_strat)

########################## done

# Get a policy from file or make a new one
#=
if repr(polfile) != "nothing"
println("Loading policy from $(polfile).")
pol = Policy.loadpolicy(polfile)
params.poltype = repr(typeof(pol))[8:end]
println("Policy type: $(params.poltype)")
else
println("Making new Policy.")
pol = Environment.makepolicy(params, mjsys.ns, mjsys.nu)
end
=#

# RUN IT
#if nprocs() > 1 # if there's an additional worker process
#   print_with_color(:yellow, "CLUSTER-PARALLEL PG.\n")
#   println("EACH PROCESS NUMT: ", params.K)
#   scores, means, evals, grads, meantraj = PolicyGradient.parallelepoch!(params, mjsys, pol, myfuncs)
#else
#   print_with_color(:yellow, "SINGLE MACHINE OPERATION.\n")
#   scores, means, evals, grads, meantraj = PolicyGradient.epoch!(params, mjsys, pol, myfuncs)
#end


text:
using .Glimpse
const Gl = Glimpse

@export function state_ledger(dio::Diorama)
    state = singleton(dio, VisualizationSettings).state_on_display
    l = Ledger(Stage(:E_basic, [EtotCalculator(), E_Dipole(), E_Gd_easy4th()]),
               Stage(:E_all, [E_Gd_H(), E_L_H(), E_LL(), E_L_easy(), E_Gd_easy(), E_Gd_L(), E_Gd_easy4th(), E_Pb()]))::Ledger
    dio_H         = singleton(dio, H)

    for e in @entities_in(dio_states)
        if dio_states[e].state == state
            Entity(l, dio[e]...)
        end
    end

    Entity(l, dio_H)
    return l
end

@component @with_kw mutable struct VisualizationSettings
    show_Gd_sphere = true
    show_Gd_arrow = true
    show_L_sphere = true
    show_L_arrow = true
    spin_font_size       = 0.8f0
    spin_font_offset     = Vec3f0(0,-1,0)
    Gd_sphere_radius     = 0.4f0
    Mn_sphere_radius     = 0.3f0
    spin_arrow_thickness = 0.08f0
    spin_arrow_length    = 0.8f0
    L_arrow_length       = 0.8f0
    L_arrow_thickness    = 0.05f0
    Mn_color             = Gl.RGBf0(0.572, 0.0, 168/255)
    Gd_color             = Gl.RGBf0(0,168/255,45/255)
    Mn_text_color        = Gl.BLACK
    Mn_text_offset       = Vec3f0(0.4, 0.4, 0.0)
    Gd_text_offset       = Vec3f0(-0.2, -0.4, 0.0)
    origin               = zero(Point3f0)
    rotatable            = false
    editor_font_size     = 0.13f0
    Etot_position  = Point3f0(7.353099822998047, 8.537099794497102+1.0, 0.0f0)
    Pb_position    = Point3f0(7.353099822998047, 8.537099794497102+2.0, 0.0f0)

    show_H = true
    H_sphere_position    = Point3f0(7.353099822998047, 8.537099794497102+3.5, 0.0)
    H_sphere_radius      = 0.2f0
    H_slider_offset      = Point3f0(2.0, 0.0, 0.0)
    H_slider_text_offset = Vec3f0(0.0, -0.5, 0.0)
    H_slider_font_size   = 0.2f0
    H_slider_thickness   = 0.05f0
    H_color              = Gl.RED

    show_PrintInfo = true
    PrintInfo_font_size = 0.5f0
    PrintInfo_color     = Gl.BLACK

end

@component struct SimEntity
    e::Entity
end

@component struct PrintInfo
end

@component struct SpatialOffset
    v::Vec3f0
end

@component struct RotationOffset
    q::Gl.Quaternions.Quaternion
end

struct Visualizer <: System end
using Glimpse: Rotation, Spatial, Text
function Overseer.update(::Visualizer, dio::AbstractLedger)
    spins    = dio[Spin]
    spatial  = dio[Spatial]
    rotation = dio[Rotation]
    text     = dio[Text]
    children = dio[Child]
    spatial_offset = dio[SpatialOffset]
    rotation_offset = dio[RotationOffset]

    #update the main Dio with the active state ledger

    #Update Gd vis
    for e in @entities_in(spins && rotation)
        rotation[e] = Gl.Rotation(Gl.Quaternions.qrotation(Gl.Z_AXIS, spins[e].ϕ) * Gl.rotation(Gl.Z_AXIS, Gl.X_AXIS))
    end

    H_entity = Entity(dio[H], 1)
    for e in @entities_in(children && !text)
        parent = children[e].parent
        if parent in spatial
            if e in spatial_offset
                dio[e] = Spatial(spatial[parent].position + spatial_offset[e].v, spatial[parent].velocity)
            else
                dio[e] = spatial[parent]
            end
        end
        if parent in rotation
            if e in rotation_offset
                dio[e] = Rotation(rotation[parent].q * rotation_offset[e].q)
            else
                dio[e] = rotation[parent]
            end
        end
    end

    # #update Mn vis
    for e in @entities_in(dio[L])
        rotation[e] = Rotation(Gl.Z_AXIS, Vec3f0(dio[L][e].v...))
    end

    #update H
    # H_children = dio[VisualizationChildren][H_entity].es
    h = singleton(dio, H)
    rotation[H_entity] = Rotation(Gl.Quaternions.qrotation(Gl.Z_AXIS, h.ϕ) * Gl.rotation(Gl.Z_AXIS, Gl.X_AXIS))
    children = dio[Child]
    for e in @entities_in(children && text)
        child = children[e]
        if child.parent == H_entity
            text[e] = Text(text[e], str="H magnitude: $(round(h.magnitude, digits=4))")
        end
    end
    for e in @entities_in(children && dio[Gl.Movable])
        child = children[e]
        if child.parent == H_entity
            spatial[e] = spatial[H_entity] + Spatial(position=Point3f0(h.magnitude, 0.0,0.0))
        end
    end

    #update printing output
    dio[Gl.Text].data[end] = Text(dio[Gl.Text].data[end], str= "E magnitude: $(round(dio[E_field][1].e, digits=4))")
    etot = calc_etot(dio, :E_all).e

    print_entities = Entity.(dio[PrintInfo].indices.packed)

    pb = Pb(dio)
    text[print_entities[1]] = Text(str="Total Energy: $(round(etot, digits=4))", font_size=0.5f0)
    text[print_entities[2]] = Text(str="Pb: $(round(pb, digits=4))", font_size=0.5f0)
end

struct SpinUpdater <: System end

@component @with_kw mutable struct OptimizationSettings
    continuous::Bool = true
    L::Bool = true
    S::Bool = true
    global_opt::Bool = false
    conjugategradient::Bool = true
end

@component @with_kw struct SimResults
    Pb_values::Vector{Float64} = Float64[]
    E_values::Vector{Float64}  = Float64[]
    H_values::Vector{Float64} = Float64[]
end

Overseer.requested_components(::SpinUpdater) = (OptimizationSettings, SimResults)

function Overseer.prepare(::SpinUpdater, dio::AbstractLedger)
    isempty(dio[OptimizationSettings]) && (Entity(dio, Gl.DioEntity(), OptimizationSettings(), SimResults()))
end

function Overseer.update(::SpinUpdater, dio::AbstractLedger)
    spins    = dio[Spin]
    rotation = dio[Gl.Rotation]
    spatial = dio[Gl.Spatial]

    keyboard = singleton(dio, Gl.Keyboard)
    children = dio[Child]
    movable = dio[Gl.Movable]
    selectable = dio[Gl.Selectable]
    opt_settings = singleton(dio, OptimizationSettings)
    simresults = singleton(dio, SimResults)
    h = singleton(dio, H)
    fontscale = 3.3f0
    visualization_settings = singleton(dio, VisualizationSettings)
    if Gl.pressed(keyboard) && !Gl.CImGui.IsAnyItemActive()
        if keyboard.button ∈ Gl.CTRL_KEYS # Update from editor moves
            for e in @entities_in(rotation && spins)
                orientation = Gl.direction(rotation[e])
                spins[e] = Spin(ϕ=atan(orientation[2], orientation[1]))
            end
            H_entity = Entity(dio[H], 1)
            H_selected = false
            for e in @entities_in(children && spatial && movable)
                child = children[e]
                if child.parent == H_entity
                    orientation = Gl.direction(rotation[H_entity])
                    new_H = H(ϕ=atan(orientation[2], orientation[1]), magnitude=norm(spatial[e].position-spatial[H_entity].position))
                    H_selected = new_H != dio[H][H_entity]
                    dio[H_entity] = new_H
                end
            end
        # elseif keyboard.button == Gl.GLFW.KEY_O #Optimize everything
        #     previous = optim_values(dio, (Spin, L))
        #     bboptimize(dio, Spin, L)
        #     if optim_values(dio, (Spin, L)) != previous
        #         push!(simresults.Pb_values, Pb(dio))
        #         push!(simresults.E_values, calc_etot(dio, :E_all).e)
        #         push!(simresults.H_values, h.magnitude)
        #     end
        elseif keyboard.button == Gl.GLFW.KEY_K #Optimize Spin
            bboptimize(dio, Spin)
        elseif keyboard.button == Gl.GLFW.KEY_L #Optimize L
            bboptimize(dio, L)
        elseif keyboard.button == Gl.GLFW.KEY_C
            res = optimize(dio, Spin, L)
            set_angles!(dio, res.minimizer)
        elseif keyboard.button == Gl.GLFW.KEY_PAGE_UP
            dio[H][Entity(dio[H], 1)] = H(ϕ=h.ϕ, magnitude=h.magnitude + 0.05)
        elseif keyboard.button == Gl.GLFW.KEY_PAGE_DOWN
            dio[H][Entity(dio[H], 1)] = H(ϕ=h.ϕ, magnitude=h.magnitude - 0.05)
        elseif keyboard.button == Gl.GLFW.KEY_LEFT
            dio[H][Entity(dio[H], 1)] = H(ϕ=h.ϕ+0.005, magnitude=h.magnitude)
        elseif keyboard.button == Gl.GLFW.KEY_RIGHT
            dio[H][Entity(dio[H], 1)] = H(ϕ=h.ϕ-0.005, magnitude=h.magnitude)
        end
    end
    gui_func = () -> begin
        Gl.Gui.SetNextWindowPos((650, 20), Gl.Gui.ImGuiCond_FirstUseEver)
        Gl.Gui.SetNextWindowSize((550, 680), Gl.Gui.ImGuiCond_FirstUseEver)

        Gl.Gui.Begin("Configuration")
        Gl.Gui.SetWindowFontScale(fontscale)
        if Gl.Gui.CollapsingHeader("Optimization Settings")
            Gl.Gui.@c Gl.Gui.Checkbox("Continuous Optimization", &opt_settings.continuous)
            Gl.Gui.@c Gl.Gui.Checkbox("bboptimize", &opt_settings.global_opt)
            Gl.Gui.@c Gl.Gui.Checkbox("conjugate gradient optimize", &opt_settings.conjugategradient)
            Gl.Gui.@c Gl.Gui.Checkbox("S", &opt_settings.S)
            Gl.Gui.SameLine()
            Gl.Gui.@c Gl.Gui.Checkbox("L", &opt_settings.L)
        end
        if Gl.Gui.CollapsingHeader("Model Parameters")
            model = singleton(dio, Model)
            Gl.Gui.@c Gl.Gui.InputDouble("K_Gd", &model.K_Gd, 0.01, 0.01, "%.3f")
            Gl.Gui.@c Gl.Gui.InputDouble("K_L", &model.K_L, 0.01, 0.01, "%.3f")
            Gl.Gui.@c Gl.Gui.InputDouble("J1", &model.J1, 0.01, 0.01, "%.3f")
            Gl.Gui.@c Gl.Gui.InputDouble("J2", &model.J2, 0.01, 0.01, "%.3f")
            Gl.Gui.@c Gl.Gui.InputDouble("gamma", &model.γ, 0.01, 0.01, "%.3f")
            Gl.Gui.@c Gl.Gui.InputDouble("chi_Gd", &model.χ_Gd, 0.01, 0.01, "%.3f")
            Gl.Gui.@c Gl.Gui.InputDouble("chi_L", &model.χ_L, 0.01, 0.01, "%.3f")
            Gl.Gui.@c Gl.Gui.InputDouble("dipole_strength", &model.dipole_strength, 0.01, 0.01, "%.3f")
            Gl.Gui.@c Gl.Gui.InputDouble("a", &model.α, 0.01, 0.01, "%.3f")
            Gl.Gui.@c Gl.Gui.InputDouble("b", &model.β, 0.01, 0.01, "%.3f")
            Gl.Gui.@c Gl.Gui.InputDouble("g", &model.g, 0.01, 0.01, "%.3f")
            if Gl.Gui.Button("Reset ModelParams")
                dio[Model].data[1] = physical_model()
            end

        end
        if Gl.Gui.CollapsingHeader("H")
            tϕ = rad2deg(h.ϕ)
            Gl.Gui.@c Gl.Gui.InputDouble("phi in degrees", &tϕ, 0.1, 0.1, "%.3f")
            Gl.Gui.@c Gl.Gui.InputDouble("H magnitude", &h.magnitude, 0.01, 0.01, "%.3f")
            h.ϕ = deg2rad(tϕ)
            h.v = h.magnitude*polar2vec(h.ϕ, 0.0)
        end
        if Gl.Gui.CollapsingHeader("E")
            e = dio[E_field][1]
            Gl.Gui.@c Gl.Gui.InputDouble("E magnitude", &e.e, 0.01, 0.01, "%.3f")
        end
        Gl.Gui.End()
        Gl.Gui.Begin("Plots")
        Gl.Gui.SetWindowFontScale(fontscale)
        if Gl.Gui.Button("Clear Simvals")
            empty!(simresults.Pb_values)
            empty!(simresults.E_values)
            empty!(simresults.H_values)
        end

        n_plotvals = 10000
        if !isempty(simresults.Pb_values)
            slen = length(simresults.Pb_values)
            offset = slen > n_plotvals ? slen-n_plotvals : 1
            len   = slen > n_plotvals ? n_plotvals : slen
            Pvals = Cfloat.(simresults.Pb_values[offset:end])
           Evals = Cfloat.(simresults.E_values[offset:end])
            hslen   = length(simresults.H_values)
            hoffset = hslen > n_plotvals ? hslen - n_plotvals : 1
            hlen    = hslen > n_plotvals ? n_plotvals : hslen
            Hvals   = Cfloat.(simresults.H_values[hoffset:end])

            Gl.Gui.PlotLines("Pb", Pvals, len, 0, C_NULL, Cint(floor(minimum(Pvals))), Cint(ceil(maximum(Pvals))), (0,240))
            Gl.Gui.PlotLines("H", Hvals,  hlen, 0, C_NULL, Cint(floor(minimum(Hvals))), Cint(ceil(maximum(Hvals))), (0,240))
            Gl.Gui.PlotLines("E", Evals,  len, 0, C_NULL, Cint(floor(minimum(Evals))), Cint(ceil(maximum(Evals))), (0,240))
        end
        Gl.Gui.End()

    end

    push!(singleton(dio, Gl.GuiFuncs).funcs, gui_func)

    if opt_settings.continuous
        previous = optim_values(dio)
        if opt_settings.conjugategradient
            res = optimize(dio)
            set_optim_values!(dio, res.minimizer)
        elseif opt_settings.global_opt
            bboptimize(dio, cs_to_optimize...)
        end
        if optim_values(dio) != previous
            push!(simresults.Pb_values, Pb(dio))
            push!(simresults.E_values, calc_etot(dio, :E_all).e)
            push!(simresults.H_values, h.magnitude)
        end
    end
end

@export function sim_snapshot(dio::Diorama)
    l = Ledger(dio[Model][1])
    Entity(l, VisualizationSettings())
    for el in @entities_in(dio[SimEntity])
        for c in filter(x -> typeof(x) in l, dio[el])
            l[dio[SimEntity][el].e] = deepcopy(c)
        end
    end
    return l
end

@export function set_sim!(dio::Diorama, l::Ledger)
    Glimpse.glimpse_call() do
        for el in @entities_in(dio[SimEntity])
            for c in l[dio[SimEntity][el].e]
                dio[el] = c
            end
        end
        fill_bonds!(dio)
    end
    return dio
end

@component mutable struct Movie
    title::String
    roll::Vector{Ledger}
end

Base.length(m::Movie) = length(m.roll)
Base.push!(m::Movie, l::Ledger) = push!(m.roll, l)
Base.last(m::Movie) = last(m.roll)
Base.getindex(m::Movie, i::Integer) = getindex(m.roll, i)

@component @with_kw mutable struct MovieSettings
    recording::Bool = false
    playing::Bool = false
    dt::Float64 = 0.5
    current_t::Float64=0.0
    current::Int = 0
    roll::Int = 0
    title::String = ""
end

struct MovieHandler <: System end
Overseer.requested_components(::MovieHandler) = (MovieSettings, Movie)

function Overseer.prepare(::MovieHandler, dio::AbstractLedger)
    isempty(dio[MovieSettings]) && (Entity(dio, Gl.DioEntity(), MovieSettings()))
end

function Overseer.update(::MovieHandler, dio::AbstractLedger)
    settings = singleton(dio, MovieSettings)
    gui_func = () -> begin
        Gl.Gui.SetNextWindowPos((650, 200), Gl.Gui.ImGuiCond_FirstUseEver)
        Gl.Gui.SetNextWindowSize((550, 680), Gl.Gui.ImGuiCond_FirstUseEver)
        Gl.Gui.Begin("Movie")
        Gl.Gui.SetWindowFontScale(3.3f0)
        movie_id = Int32(settings.roll)
        movie_title = zeros(UInt8, 256)
        if movie_id != 0
            settings.title = dio[Movie][settings.roll].title
        end
        movie_title[1:length(settings.title)] = Vector{UInt8}(settings.title)
        # @show movie_title
        Gl.Gui.@c Gl.Gui.InputText("Title", movie_title, 256)
        settings.title = String(filter(!iszero, movie_title))
        if movie_id != 0
            dio[Movie][settings.roll].title = settings.title
        end
        Gl.Gui.@c Gl.Gui.InputInt("Id", &movie_id)
        if movie_id > 0 && movie_id <= length(dio[Movie]) && movie_id != settings.roll
            settings.roll = Int(movie_id)
            settings.current = 1
        end
        curr = Int32(settings.current)
        Gl.Gui.@c Gl.Gui.InputInt("Frame", &curr)
        if settings.roll != 0 && length(dio[Movie][settings.roll]) >= curr && curr > 0 && curr != settings.current && !settings.playing
           settings.current = Int(curr)
           @show "ping"
           set_sim!(dio, dio[Movie][settings.roll][settings.current])
       end

        Gl.Gui.@c Gl.Gui.InputDouble("Frame Time", &settings.dt,0.01,0.01,"%.3f")
        if !settings.recording && !settings.playing
            if Gl.Gui.Button("Start Recording")
                settings.recording = true
                Entity(dio, Movie("$(length(dio[Movie])+1)", Ledger[sim_snapshot(dio)]))
                settings.roll = length(dio[Movie])
                settings.current = 1
            elseif Gl.Gui.Button("Play") && settings.current != 0
                settings.playing = true
            elseif Gl.Gui.Button("Save") && settings.roll != 0 && !isempty(dio[Movie][settings.roll].title)
                save_movie(dio[Movie][settings.roll])
            elseif Gl.Gui.Button("Load")
                m = load_movie(settings.title)
                if m !== nothing
                    Entity(dio, m)
                    settings.roll = length(dio[Movie])
                    settings.current = 1
                end
            end
        elseif settings.recording && !settings.playing
            if Gl.Gui.Button("Stop Recording")
                settings.recording = false
            end
        else
            if Gl.Gui.Button("Stop Movie")
                settings.playing = false
                settings.current_t = 0.0
            end
        end
        Gl.Gui.End()
    end
    push!(singleton(dio, Gl.GuiFuncs).funcs, gui_func)
    if settings.recording
        cursim = sim_snapshot(dio)
        if cursim != last(dio[Movie][settings.roll])
            push!(dio[Movie][settings.roll], cursim)
            settings.current += 1
        end
    elseif settings.playing
        settings.current_t += dio[Gl.TimingData][1].dtime
        if settings.current_t >= settings.dt
            set_sim!(dio, dio[Movie][settings.roll][settings.current])
            settings.current += 1
            if settings.current >= length(dio[Movie][settings.roll])
                settings.current=1
            end
            settings.current_t=0.0
        end
    end
end

function save_movie(m::Movie)
    if !ispath(datadir("Movies"))
        mkdir(datadir("Movies"))
    end
    save(datadir("Movies", m.title * ".jld2"), "movie", m)
end

function load_movie(title::String)
    path = datadir("Movies", title * ".jld2")
    if ispath(path)
        return load(path)["movie"]
    end
end

function Glimpse.Diorama(l::Ledger; kwargs...)
    dio = Diorama(stage(l, :E_all); background =Gl.RGBAf0(81/255, 107/255, 117/255, 1.0), kwargs...)
    Overseer.ensure_component!(dio, Optimize)
    Overseer.ensure_component!(dio, RotationOffset)
    camera_entity = Entity(first(dio[Gl.Camera3D].indices))
    dio[Gl.Spatial][camera_entity] = Gl.Spatial(dio[Gl.Spatial][camera_entity], position=Point3f0(0, 0, 10f0))
    dio[Gl.Camera3D][camera_entity].up = Vec3f0(Gl.Y_AXIS...)
    center_camera!(dio, Point3f0(7.353099822998047,  8.537099794497102/2, 0.0))
    singleton(dio, Gl.Camera3D).camerakind = Gl.Orthographic
    # New stages
    push!(dio, :simulation, SpinUpdater())
    push!(dio, :simulation, Visualizer())
    push!(dio, :simulation, MovieHandler())
    prepare(dio)
    # update(dio)

    for e in l.entities
        Entity(dio, l[e]..., SimEntity(e))
    end

    fill_bonds!(dio)

    @assert VisualizationSettings ∈ l "Please first add a VisualizationSettings to the ledger"
    settings = singleton(l, VisualizationSettings)

    Gd_counter = 1
    for (i, e) in enumerate(@entities_in(dio[Position] && dio[Spin] && dio[Gd]))
        pos    = Point3f0(dio[Position][e].p...) + settings.origin
        dio[e] = Spatial(position = pos)
        if settings.show_Gd_sphere
            Entity(dio,
                   Child(e),
                   Gl.assemble_sphere(radius = settings.Gd_sphere_radius,
                                      color  = settings.Gd_color)[2:end]...)
            Entity(dio,Spatial(position = pos) , Gl.UniformColor(settings.Gd_color), Gl.Text(str    = "S$Gd_counter",
                                          offset = settings.Gd_text_offset,
                                          font_size= settings.spin_font_size))
            Gd_counter += 1
        end
        if settings.show_Gd_arrow
            comps = Gl.assemble_arrow(zero(Point3f0),
                                      settings.spin_arrow_length * Point3f0(dio[Spin][e].v...),
                                      thickness = settings.spin_arrow_thickness,
                                      color     = settings.Gd_color)[2:end]
            for c in comps
                dio[e] = c
            end
            if settings.rotatable
                dio[e] = Gl.Rotatable(1f0, 0.02f0/10, settings.editor_font_size)
            end
        end
    end
    Mn_counter = 1
    for (i, e) in enumerate(@entities_in(dio[Position] && dio[Spin] && dio[Mn]))
        pos    = Point3f0(dio[Position][e].p...) + settings.origin
        dio[e] = Spatial(position=pos)
        if settings.show_Gd_sphere
            Entity(dio,
                   Child(e),
                   Gl.assemble_sphere(radius=settings.Mn_sphere_radius, color=settings.Mn_color)[2:end]...)
            Entity(dio, Child(e), Gl.Text(str="Mn$Mn_counter",  font_size=settings.spin_font_size, offset=settings.Mn_text_offset))
            Mn_counter += 1
        end
        if settings.show_Gd_arrow
            for es in Gl.assemble_arrow(zero(Point3f0), settings.spin_arrow_length*Point3f0(dio[Spin][e].v[1], dio[Spin][e].v[2], 0.0), thickness=settings.spin_arrow_thickness, color=settings.Mn_color)[2:end]
                dio[e] = es
            end
            if settings.rotatable
                dio[e] = Gl.Rotatable(1f0, 0.02f0/10, settings.editor_font_size)
            end
        end
    end
    a = Vec2(7.353099822998047, 0.0)
    b = Vec2(0.0, 8.537099794497102)

    ncells = length(l[Spin]) < 4 ? 1 : div(length(l[Spin]), 4)
    right = Point3f0((div(ncells, 3) * a)...,0.0) + settings.origin

    Mn_positions = Vec2.([[0.0, 4.268549897248551],
                          [3.6765499114990234, 0.0],
                          [3.0280065071105957, 3.0033517077040806],
                          [6.704556418609619, 1.2651981895444706],
                          [4.325093315887451, 5.533748086793022],
                          [0.6485434043884277, 7.271901604952632]])

    L_to_S = [(id=2, s=1),
              (id=1, s=1),
              (id=2, s=-1),
              (id=1, s=-1),
              (id=2, s=1),
              (id=1, s=-1)]
    L_entities = Entity[]
    for e in @entities_in(dio[L])
        push!(L_entities,e)
    end
    for (l2s, pos) in zip(L_to_S, Mn_positions)
        spin_sign = l2s.s
        if l2s.id > length(dio[L])
            continue
        end
        t_L = dio[L][l2s.id].v
        for i=1:ncells
            cell_sign = i%2 == 1 ? 1 : -1
            p         = Point3f0(pos + (i-1)*a..., 0.0) + settings.origin
            spat_entity    = Entity(dio, Gl.Spatial(position=p), Child(L_entities[l2s.id]))
            spin_end_point = Point3f0((t_L*spin_sign*cell_sign)...)
            str1 = cell_sign*spin_sign == -1 ? "-" : ""
            if settings.show_L_sphere
                sphere_entity = Entity(dio,
                                       Child(spat_entity),
                                       Gl.assemble_sphere(p, radius=settings.Mn_sphere_radius, color=settings.Mn_color)...,)
            end
            if settings.show_L_arrow

                offset_angle = cell_sign*spin_sign == -1 ? π : 0.0
                rotation_offset = RotationOffset(Gl.Quaternions.qrotation(Gl.X_AXIS, offset_angle))

                spin_arrow = Gl.assemble_arrow(zero(Point3f0), settings.L_arrow_length*spin_end_point, thickness=settings.L_arrow_thickness, color=settings.Mn_color)
                Entity(dio,
                       Child(spat_entity),
                       spin_arrow...,
                       rotation_offset)
            end
        end
    end
    if settings.show_PrintInfo
        Entity(dio,
               PrintInfo(),
               Gl.Spatial(position=settings.Etot_position),
               Gl.Text(str="Total Energy: $(calc_etot(dio).e)", font_size=settings.PrintInfo_font_size, ),
               Gl.UniformColor(settings.PrintInfo_color))
        Entity(dio,
               PrintInfo(),
               Gl.Spatial(position=settings.Pb_position),
               Gl.Text(str="Pb: $(Pb(dio))", font_size=settings.PrintInfo_font_size, ),
               Gl.UniformColor(settings.PrintInfo_color))
    end
    if settings.show_H
        h = singleton(dio, H)
        h_entity = Entity(dio[H], 1)

        for e in Gl.assemble_arrow(settings.H_sphere_position,
                                   settings.H_sphere_position+Point3f0((polar2vec(h.ϕ)*1.0f0)...,0.0f0),
                                   thickness=settings.spin_arrow_thickness,
                                   color=settings.H_color)
            dio[h_entity] = e
        end
        if settings.rotatable
            dio[h_entity] = Gl.Rotatable(1f0, 0.02f0/10, settings.editor_font_size)
        end

        Entity(dio, Child(h_entity), Gl.assemble_sphere(settings.H_sphere_position, color=settings.H_color, radius=settings.H_sphere_radius)[2:end]...)

        Entity(dio,
               Child(h_entity),
               PrintInfo(),
               Gl.Spatial(position = settings.H_sphere_position -Vec3f0(0.0,-0.2,0.0)),
               Gl.Text(str="H magnitude: $(h.magnitude)", font_size=settings.PrintInfo_font_size,  offset=settings.H_slider_text_offset), Glimpse.UniformColor(Gl.RED))

        Entity(dio,
               Child(h_entity),
               Gl.assemble_sphere(color=settings.H_color, radius=settings.H_slider_thickness)[2:end]...,
               SpatialOffset(settings.H_slider_offset + Point3f0(h.magnitude, 0,0)),
               Gl.Movable(axis_thickness=settings.H_slider_thickness, axis_length=1f0, font_size=settings.editor_font_size))
    end

    # cell
    a3d = Point3f0(a..., 0.0)
    b3d = Point3f0(b..., 0.0)
    Entity(dio, Gl.assemble_line([zero(Point3f0), zero(Point3f0), a3d, a3d+b3d, b3d, zero(Point3f0), a3d, 2*a3d, 2*a3d+b3d, a3d+b3d, a3d, a3d], color=Gl.BLACK, thickness=4f0)...)

    # chains
    chainpts1 = [Point3f0(0.0, 4.268549897248551, 0.0),
                 Point3f0(3.0280065071105957, 3.0033517077040806,0.0),
                 Point3f0(4.325093315887451, 5.533748086793022, 0.0)]
    append!(chainpts1, map(x-> x+a3d, chainpts1))
    push!(chainpts1, chainpts1[1] + 2*a3d)

    chainpts2 = [Point3f0(0.6485434043884277, 7.271901604952632,0.0)-b3d,
                 Point3f0(3.6765499114990234, 0.0, 0.0),
                 Point3f0( 6.704556418609619, 1.2651981895444706, 0.0)]

    append!(chainpts2, map(x-> x+a3d, chainpts2))
    chainpts3 = map(x->x+b3d, chainpts2)
    # push!(chainpts1, chainpts1[1] + 2*a3d)

    insert!(chainpts1, 1, chainpts1[1])
    push!(chainpts1, chainpts1[end])
    insert!(chainpts2, 1, chainpts2[1])
    push!(chainpts2, chainpts2[end])
    insert!(chainpts3, 1, chainpts3[1])
    push!(chainpts3, chainpts3[end])

    Entity(dio, Gl.assemble_line(chainpts1, color=Gl.RGBf0(23/255,145/255,194/255), thickness=4f0)...)
    Entity(dio, Gl.assemble_line(chainpts2, color=Gl.RGBf0(23/255,145/255,194/255), thickness=4f0)...)
    Entity(dio, Gl.assemble_line(chainpts3, color=Gl.RGBf0(23/255,145/255,194/255), thickness=4f0)...)

    Entity(dio, Gl.Spatial(position=settings.H_sphere_position + Vec3f0(0.0, 1.0, 0.0)), Gl.Text(str="E magnitude: $(round(dio[E_field][1].e, digits=4))", font_size=settings.PrintInfo_font_size), Gl.UniformColor(Gl.RED), PrintInfo())

    return dio
end


text:
using Plots, SNN

G = SNN.Rate(;N = 100)
GG = SNN.RateSynapse(G, G; σ = 1.2, p = 1.0)
SNN.monitor(G, [(:r, [1, 50, 100])])

SNN.sim!([G], [GG]; duration = 100ms)
SNN.vecplot(G, :r) |> display