import Nemo.crt, Nemo.zero, Nemo.iszero, Nemo.isone, Nemo.sub!
export crt_env, crt, crt_inv, modular_init, crt_signed

isone(a::Int) = (a==1)

function zero(a::PolyElem)
  return zero(parent(a))
end

@inline function rem!(a::fmpz, b::fmpz, c::fmpz)
  ccall((:fmpz_mod, :libflint), Nothing, (Ref{fmpz}, Ref{fmpz}, Ref{fmpz}), a, b, c)
end

@inline function sub!(a::fmpz, b::fmpz, c::fmpz)
  ccall((:fmpz_sub, :libflint), Nothing, (Ref{fmpz}, Ref{fmpz}, Ref{fmpz}), a, b, c)
end

function rem!(a::fmpz_mod_poly, b::fmpz_mod_poly, c::fmpz_mod_poly)
  ccall((:fmpz_mod_poly_rem, :libflint), Nothing, (Ref{fmpz_mod_poly}, Ref{fmpz_mod_poly}, Ref{fmpz_mod_poly}), a, b, c)
end

mutable struct crt_env{T}
  pr::Array{T, 1}
  id::Array{T, 1}
  tmp::Array{T, 1}
  t1::T
  t2::T
  n::Int
  M::T #for T=fmpz, holds prod/2

  function crt_env{T}(p::Array{T, 1}) where {T}
    pr = deepcopy(p)
    id = Array{T, 1}()
    i = 1
    while 2*i <= length(pr)
      a = pr[2*i-1]
      b = pr[2*i]
      if false
        g, u, v = gcdx(a, b)
        @assert isone(g)
        push!(id, v*b , u*a )
      else
        # we have 1 = g = u*a + v*b, so v*b = 1-u*a - which saves one mult.
        u = invmod(a, b)
#        @assert  (a*u) % b == 1
        u *= a
        push!(id, 1-u, u)
      end
      push!(pr, a*b)
      i += 1
    end
    r = new{T}()
    r.pr = pr
    r.id = id

    r.tmp = Array{T, 1}()
    n = length(p)
    for i=1:div(n+1, 2)
      push!(r.tmp, zero(p[1]))
    end
    r.t1 = zero(p[1])
    r.t2 = zero(p[1])

    r.n = n
    return r
  end
end

@doc Markdown.doc"""
    crt_env(p::Array{T, 1}) -> crt_env{T}

Given coprime moduli in some euclidean ring (FlintZZ, nmod_poly,
fmpz\_mod_poly), prepare data for fast application of the chinese
remander theorem for those moduli.
"""
function crt_env(p::Array{T, 1}) where T
  return crt_env{T}(p)
end

function show(io::IO, c::crt_env{T}) where T
  print(io, "CRT data for moduli ", c.pr[1:c.n])
end

@doc Markdown.doc"""
    crt{T}(b::Array{T, 1}, a::crt_env{T}) -> T

Given values in b and the environment prepared by crt\_env, return the
unique (modulo the product) solution to $x \equiv b_i \bmod p_i$.
"""
function crt(b::Array{T, 1}, a::crt_env{T}) where T
  res = zero(b[1])
  return crt!(res, b, a)
end

function crt!(res::T, b::Array{T, 1}, a::crt_env{T}) where T
  @assert a.n == length(b)
  bn = div(a.n, 2)
  if isodd(a.n)
    @inbounds zero!(a.tmp[1])
    @inbounds add!(a.tmp[1], a.tmp[1], b[end])
    off = 1
  else
    off = 0
  end

  for i=1:bn
    @inbounds mul!(a.t1, b[2*i-1], a.id[2*i-1])
    @inbounds mul!(a.t2, b[2*i], a.id[2*i])
    @inbounds add!(a.tmp[i+off], a.t1, a.t2)
    @inbounds rem!(a.tmp[i+off], a.tmp[i+off], a.pr[a.n+i])
  end

  if isodd(a.n)
    bn += 1
  end

  id_off = a.n - off
  pr_off = a.n + bn - off
#  println(a.tmp, " id_off=$id_off, pr_off=$pr_off, off=$off, bn=$bn")
  while bn>1
    if isodd(bn)
      off = 1
    else
      off = 0
    end
    bn = div(bn, 2)
    for i=1:bn
      if true  # that means we need only one co-factor!!!
        @inbounds sub!(a.t1, a.tmp[2*i-1], a.tmp[2*i])
        @inbounds mul!(a.t1, a.t1, a.id[id_off + 2*i-1])
        @inbounds add!(a.tmp[i+off], a.t1, a.tmp[2*i])
      else
        @inbounds mul!(a.t1, a.tmp[2*i-1], a.id[id_off + 2*i-1])
        @inbounds mul!(a.t2, a.tmp[2*i], a.id[id_off + 2*i])
        @inbounds add!(a.tmp[i + off], a.t1, a.t2)
      end
      @inbounds rem!(a.tmp[i + off], a.tmp[i + off], a.pr[pr_off+i])
    end
    if off == 1
      @inbounds a.tmp[1], a.tmp[2*bn+1] = a.tmp[2*bn+1], a.tmp[1]
    end
    id_off += 2*bn
    pr_off += bn
    bn += off
#    println(a.tmp, " id_off=$id_off, pr_off=$pr_off, off=$off, bn=$bn")
  end
  zero!(res)
  @inbounds add!(res, res, a.tmp[1])
  return res
end

function crt_signed!(res::fmpz, b::Array{fmpz, 1}, a::crt_env{fmpz})
  crt!(res, b, a)
  if !isdefined(a, :M)
    a.M = div(prod(a.pr[1:a.n]), 2)
  end
  if res>a.M
    sub!(res, res, a.pr[end])
  end
end

function crt_signed(b::Array{fmpz, 1}, a::crt_env{fmpz})
  res = fmpz()
  crt_signed!(res, b, a)
  return res
end

#in .pr we have the products of pairs, ... in the wrong order
# so we traverse this list backwards, while building the remainders...
#.. and then we do it again, efficiently to avoid resorting and re-allocation
#=
function crt_inv{T}(a::T, c::crt_env{T})
  r = Array{T, 1}()
  push!(r, a)
  i = length(c.pr)-1
  j = 1
  while i>1
    push!(r, r[j] % c.pr[i], r[j] %c.pr[i-1])
    i -= 2
    j += 1
  end
  return reverse(r, length(r)-c.n+1:length(r))
end
=#

function crt_inv_iterative!(res::Array{T,1}, a::T, c::crt_env{T}) where T
  for i=1:c.n
    if isassigned(res, i)
      rem!(res[i], a, c.pr[i])
    else
      res[i] = a % c.pr[i]
    end
  end
  return res
end

function crt_inv_tree!(res::Array{T,1}, a::T, c::crt_env{T}) where T
  for i=1:c.n
    if !isassigned(res, i)
      res[i] = zero(a)
    end
  end

  i = length(c.pr)-1
  if i == 0
    rem!(res[1], a, c.pr[1])
    return res
  end
  r = i
  w = r + c.n - 1

  @inbounds zero!(res[r % c.n + 1])
  @inbounds add!(res[r % c.n + 1], res[r % c.n + 1], a)

  while i>1
    @inbounds rem!(res[w % c.n + 1], res[r % c.n + 1], c.pr[i])
    @inbounds rem!(res[(w+c.n - 1) % c.n + 1], res[r % c.n + 1], c.pr[i - 1])
    w += 2*(c.n-1)
    i -= 2
    r += 1*(c.n-1)
  end

  return res
end

@doc Markdown.doc"""
    crt_inv(a::T, crt_env{T}) -> Array{T, 1}

Given a $\code{crt_env}$ and an element a, return
the modular data $a \bmod pr_i$ for all $i$.
This is essentially the inverse to the $\code{crt}$ function.
"""
function crt_inv(a::T, c::crt_env{T}) where T
  res = Array{T}(undef, c.n)
  if c.n < 50
    return crt_inv_iterative!(res, a, c)
  else
    return crt_inv_tree!(res, a, c)
  end
end

function crt_inv!(res::Array{T, 1}, a::T, c::crt_env{T}) where T
  if c.n < 50
    return crt_inv_iterative!(res, a, c)
  else
    return crt_inv_tree!(res, a, c)
  end
end

#explains the idea, but is prone to overflow.
# idea: the tree CRT ..
# given moduli p1 .. pn, we first do (p1, p2), (p2, p3), ...
# then ((p1, p2), (p3, p4)), ...
# until done.
# In every step we need the cofactors, the inverse of pi mod pj
# thus we build a parallel array id for the cofactors
# in id[2i-1], id[2i] are the cofactors for pr[2i-1], pr[2i]
# To recombine, we basically loop through the cofactors:
# use id[1], id[2] to combine b[1], b[2] AND append to b
# The product pr[1]*pr[2] was appended to pr, thus we can walk through the
# growing list till the end
# For the optimized version, we have tmp-array to hold the CRT results
# plus t1, t2 for temporaty products.
function crt(b::Array{Int, 1}, a::crt_env{Int})
  i = a.n+1
  j = 1
  while j <= length(b)
    push!(b, (b[j-1]*a.id[j-1] + b[j]*a.id[j]) % a.pr[i])
    j += 2
    i += 1
  end
  return b[end]
end

function crt_test(a::crt_env{fmpz}, b::Int)
  z = [fmpz(0) for x=1:a.n]
  for i=1:b
    b = rand(0:a.pr[end]-1)
    for j=1:a.n
      rem!(z[j], b, a.pr[j])
    end
    if b != crt(z, a)
      println("problem: $b and $z")
    end
    @assert b == crt(z, a)
  end
end

@doc Markdown.doc"""
  crt(r1::PolyElem, m1::PolyElem, r2::PolyElem, m2::PolyElem) -> PolyElem

Find $r$ such that $r \equiv r_1 \pmod m_1$ and $r \equiv r_2 \pmod m_2$
"""
function crt(r1::PolyElem{T}, m1::PolyElem{T}, r2::PolyElem{T}, m2::PolyElem{T}) where T
  g, u, v = gcdx(m1, m2)
  m = m1*m2
  return (r1*v*m2 + r2*u*m1) % m
end

@doc Markdown.doc"""
  crt_iterative(r::Array{T, 1}, m::Array{T,1}) -> T

Find $r$ such that $r \equiv r_i \pmod m_i$ for all $i$.
A plain iteration is performed.
"""
function crt_iterative(r::Array{T, 1}, m::Array{T, 1}) where T
  p = crt(r[1], m[1], r[2], m[2])
  d = m[1] * m[2]
  for i = 3:length(m)
    p = crt(p, d, r[i], m[i])
    d *= m[i]
  end
  return p
end

@doc Markdown.doc"""
  crt_tree(r::Array{T, 1}, m::Array{T,1}) -> T

Find $r$ such that $r \equiv r_i \pmod m_i$ for all $i$.
A tree based strategy is used that is asymptotically fast.
"""
function crt_tree(r::Array{T, 1}, m::Array{T, 1}) where T
  if isodd(length(m))
    M = [m[end]]
    V = [r[end]]
  else
    M = Array{T, 1}()
    V = Array{T, 1}()
  end

  for i=1:div(length(m), 2)
    push!(V, crt(r[2*i-1], m[2*i-1], r[2*i], m[2*i]))
    push!(M, m[2*i-1]*m[2*i])
  end
  i = 1
  while 2*i <= length(V)
    push!(V, crt(V[2*i-1], M[2*i-1], V[2*i], M[2*i]))
    push!(M, M[2*i-1] * M[2*i])
    i += 1
  end
#  println("M = $M\nV = $V")
  return V[end]
end

@doc Markdown.doc"""
  crt(r::Array{T, 1}, m::Array{T,1}) -> T

Find $r$ such that $r \equiv r_i \pmod m_i$ for all $i$.
"""
function crt(r::Array{T, 1}, m::Array{T, 1}) where T
  length(r) == length(m) || error("Arrays need to be of same size")
  if length(r) < 15
    return crt_iterative(r, m)
  else
    return crt_tree(r, m)
  end
end

function crt_test_time_all(np::Int, n::Int)
  p = next_prime(fmpz(2)^60)
  m = [p]
  x = fmpz(1)
  for i=1:np-1
    push!(m, next_prime(m[end]))
  end
  v = [rand(0:x-1) for x = m]
  println("crt_env...")
  @time ce = crt_env(m)
  @time for i=1:n
    x = crt(v, ce)
  end

  println("iterative...")
  @time for i=1:n
    x = crt_iterative(v, m)
  end

  println("tree...")
  @time for i=1:n
    x = crt_tree(v, m)
  end

  println("inv_tree")
  @time for i=1:n
    crt_inv_tree!(m, x, ce)
  end

  println("inv_iterative")
  @time for i=1:n
    crt_inv_iterative!(m, x, ce)
  end
end

@doc Markdown.doc"""
    induce_crt(a::fmpz_poly, p::fmpz, b::fmpz_poly, q::fmpz, signed::Bool = false) -> fmpz_poly
Given integral polynomials $a$ and $b$ as well as coprime integer moduli
$p$ and $q$, find $f = a \bmod p$ and $f=b \bmod q$.
If signed is set, the symmetric representative is used, the positive one
otherwise.
"""
function induce_crt(a::fmpz_poly, p::fmpz, b::fmpz_poly, q::fmpz, signed::Bool = false)
  c = parent(a)()
  pi = invmod(p, q)
  mul!(pi, pi, p)
  pq = p*q
  if signed
    pq2 = div(pq, 2)
  else
    pq2 = fmpz(0)
  end
  for i=0:max(degree(a), degree(b))
    setcoeff!(c, i, Hecke.inner_crt(coeff(a, i), coeff(b, i), pi, pq, pq2))
  end
  return c, pq
end

@doc Markdown.doc"""
    induce_crt(L::Array{PolyElem, 1}, c::crt_env{fmpz}) -> fmpz_poly
Given fmpz\_poly polynomials $L[i]$ and a {{{crt\_env}}}, apply the
{{{crt}}} function to each coefficient resulting in a polynomial $f = L[i] \bmod p[i]$.
"""
function induce_crt(L::Array{T, 1}, c::crt_env{fmpz}) where {T <: PolyElem}
  Zx, x = FlintZZ["x"]
  res = Zx()
  m = maximum(degree(x) for x = L)

  for i=0:m
    setcoeff!(res, i, crt([lift(coeff(x, i)) for x =L], c))
  end
  return res
end

@doc Markdown.doc"""
  _num_setcoeff!(a::nf_elem, n::Int, c::fmpz)
  _num_setcoeff!(a::nf_elem, n::Int, c::Integer)

Sets the $n$-th coefficient in $a$ to $c$. No checks performed, use
only if you know what you're doing.
"""
function _num_setcoeff!(a::nf_elem, n::Int, c::fmpz)
  K = parent(a)
  @assert n < degree(K) && n >=0
  ra = pointer_from_objref(a)
  if degree(K) == 1
    ccall((:fmpz_set, :libflint), Nothing, (Ref{Nothing}, Ref{fmpz}), ra, c)
    ccall((:fmpq_canonicalise, :libflint), Nothing, (Ref{nf_elem}, ), a)
  elseif degree(K) == 2
     ccall((:fmpz_set, :libflint), Nothing, (Ref{Nothing}, Ref{fmpz}), ra+n*sizeof(Int), c)
  else
    ccall((:fmpq_poly_set_coeff_fmpz, :libflint), Nothing, (Ref{nf_elem}, Int, Ref{fmpz}), a, n, c)
   # includes canonicalisation and treatment of den.
  end
end

function _num_setcoeff!(a::nf_elem, n::Int, c::UInt)
  K = a.parent
  @assert n < degree(K) && n >=0

  ra = pointer_from_objref(a)

  if degree(K) == 1
    ccall((:fmpz_set_ui, :libflint), Nothing, (Ref{Nothing}, UInt), ra, c)
    ccall((:fmpq_canonicalise, :libflint), Nothing, (Ref{nf_elem}, ), a)
  elseif degree(K) == 2
    ccall((:fmpz_set_ui, :libflint), Nothing, (Ref{Nothing}, UInt), ra+n*sizeof(Int), c)
  else
    ccall((:fmpq_poly_set_coeff_ui, :libflint), Nothing, (Ref{nf_elem}, Int, UInt), a, n, c)
   # includes canonicalisation and treatment of den.
  end
end

function _num_setcoeff!(a::nf_elem, n::Int, c::Integer)
  _num_setcoeff!(a, n, fmpz(c))
end

@doc Markdown.doc"""
    induce_crt(L::Array{MatElem, 1}, c::crt_env{fmpz}) -> fmpz_mat
Given matrices $L[i]$ and a {{{crt\_env}}}, apply the
{{{crt}}} function to each coefficient resulting in a matrix $M = L[i] \bmod p[i]$.
"""
function induce_crt(L::Array{T, 1}, c::crt_env{fmpz}, signed::Bool = false) where {T <: MatElem}
  res = zero_matrix(FlintZZ, nrows(L[1]), ncols(L[1]))

  if signed
    cr = crt_signed
  else
    cr = crt
  end

  for i=1:nrows(L[1])
    for j=1:ncols(L[1])
      res[i,j] = cr([lift(x[i,j]) for x =L], c)
    end
  end
  return res
end

mutable struct modular_env
  p::fmpz
  up::UInt
  upinv::UInt

  fld::Array{FqNmodFiniteField, 1}
  fldx::Array{FqNmodPolyRing, 1}
  ce::crt_env{nmod_poly}
  rp::Array{nmod_poly, 1}
  res::Array{fq_nmod, 1}
  Fpx::NmodPolyRing
  K::AnticNumberField
  Rp::Array{fq_nmod_poly, 1}
  Kx::Generic.PolyRing{nf_elem}
  Kxy::Generic.MPolyRing{nf_elem}
  Kpxy::NmodMPolyRing

  function modular_env()
    return new()
  end
end

Base.isempty(me::modular_env) = !isdefined(me, :ce)

function show(io::IO, me::modular_env)
  if isempty(me)
    println("modular environment for p=$(me.p), using $(0) ideals")
  else
    println("modular environment for p=$(me.p), using $(me.ce.n) ideals")
  end
end

@doc Markdown.doc"""
  modular_init(K::AnticNumberField, p::fmpz) -> modular_env
  modular_init(K::AnticNumberField, p::Integer) -> modular_env

Given a number field $K$ and an ``easy'' prime $p$ (ie. fits into an
\code{Int} and is coprime to the polynomial discriminant), compute
the residue class fields of the associated primes ideals above $p$.
Returns data that can be used by \code{modular_proj} and \code{modular_lift}.
"""
function modular_init(K::AnticNumberField, p::fmpz; deg_limit::Int=0, max_split::Int = 0)
  @assert isprime(p)
  UInt(p) # to enforce p being small

  me = modular_env()
  me.Fpx = PolynomialRing(ResidueRing(FlintZZ, Int(p), cached = false), "_x", cached=false)[1]
  fp = me.Fpx(K.pol)
  lp = factor(fp)
  if Set(values(lp.fac)) != Set([1])
    throw(BadPrime(p))
  end
  pols = collect(keys(lp.fac))
  if deg_limit > 0
    pols = pols[findall(x -> degree(x) <= deg_limit, pols)]
  end

  if max_split > 0
    pols = pols[1:min(length(pols), max_split)]
  end
  if length(pols) == 0
    return me
  end
  me.ce = crt_env(pols)
  me.fld = [FqNmodFiniteField(x, :$, false) for x = pols]  #think about F_p!!!
                                   # and chacheing
  me.rp = Array{nmod_poly}(undef, length(pols))
  me.res = Array{fq_nmod}(undef, me.ce.n)

  me.p = p
  me.K = K
  me.up = UInt(p)
  me.upinv = ccall((:n_preinvert_limb, :libflint), UInt, (UInt, ), me.up)
  return me
end

function modular_init(K::AnticNumberField, p::Integer; deg_limit::Int=0, max_split::Int = 0)
  return modular_init(K, fmpz(p), deg_limit = deg_limit, max_split = max_split)
end

@doc Markdown.doc"""
  modular_proj(a::nf_elem, me::modular_env) -> Array{fq_nmod, 1}

Given an algebraic number $a$ and data \code{me} as computed by
\code{modular_init}, project $a$ onto the residue class fields.
"""
function modular_proj(a::nf_elem, me::modular_env)
  ap = me.Fpx(a)
  crt_inv!(me.rp, ap, me.ce)
  for i=1:me.ce.n
    F = me.fld[i]
    if isassigned(me.res, i)
      u = me.res[i]
    else
      u = F()
    end
    ccall((:fq_nmod_set, :libflint), Nothing,
                (Ref{fq_nmod}, Ref{nmod_poly}, Ref{FqNmodFiniteField}),
                u, me.rp[i], F)
    me.res[i] = u
  end
  return me.res
end

@doc Markdown.doc"""
  modular_proj(a::FacElem{nf_elem, AnticNumberField}, me::modular_env) -> Array{fq_nmod, 1}

Given an algebraic number $a$ in factored form and data \code{me} as computed by
\code{modular_init}, project $a$ onto the residue class fields.
"""
function modular_proj(A::FacElem{nf_elem, AnticNumberField}, me::modular_env)
  for i=1:me.ce.n
    me.res[i] = one(me.fld[i])
  end
  for (a, v) = A.fac
    ap = me.Fpx(a)
    crt_inv!(me.rp, ap, me.ce)
    for i=1:me.ce.n
      F = me.fld[i]
      u = F()
      ccall((:fq_nmod_set, :libflint), Nothing,
                  (Ref{fq_nmod}, Ref{nmod_poly}, Ref{FqNmodFiniteField}),
                  u, me.rp[i], F)
      u = u^(v % (size(F)-1))
      me.res[i] *= u
    end
  end
  return me.res
end

@doc Markdown.doc"""
  modular_lift(a::Array{fq_nmod}, me::modular_env) -> nf_elem

Given an array of elements as computed by \code{modular_proj},
compute a global pre-image using some efficient CRT.
"""
function modular_lift(a::Array{fq_nmod, 1}, me::modular_env)
  for i=1:me.ce.n
    ccall((:nmod_poly_set, :libflint), Nothing, (Ref{nmod_poly}, Ref{fq_nmod}), me.rp[i], a[i])
  end
  ap = crt(me.rp, me.ce)
  r = me.K()
  for i=0:ap.length-1
    u = ccall((:nmod_poly_get_coeff_ui, :libflint), UInt, (Ref{nmod_poly}, Int), ap, i)
    _num_setcoeff!(r, i, u)
  end
  return r
end

@doc Markdown.doc"""
  modular_proj(a::Generic.Poly{nf_elem}, me::modular_env) -> Array

Apply the \code{modular_proj} function to each coeficient of $a$.
Computes an array of polynomials over the respective residue class fields.
"""
function modular_proj(a::Generic.Poly{nf_elem}, me::modular_env)

  if !isdefined(me, :fldx)
    me.fldx = [PolynomialRing(x, "_x", cached=false)[1] for x = me.fld]
    me.Rp = Array{fq_nmod_poly}(undef, me.ce.n)
    for i =1:me.ce.n
      me.Rp[i] = me.fldx[i](0)
    end
    me.Kx = parent(a)
  end

  for j=1:me.ce.n
    zero!(me.Rp[j])
  end
  for i=0:length(a)-1
    c = coeff(a, i)
    if iszero(mod(denominator(c), me.p))
      throw(BadPrime(me.p))
    end
    crt_inv!(me.rp, me.Fpx(c), me.ce)
    for j=1:me.ce.n
      u = coeff(me.Rp[j], i)
      ccall((:fq_nmod_set, :libflint), Nothing,
                   (Ref{fq_nmod}, Ref{nmod_poly}, Ref{FqNmodFiniteField}),
                   u, me.rp[j], me.fld[j])
      setcoeff!(me.Rp[j], i, u)
    end
  end
  return me.Rp
end

@doc Markdown.doc"""
  modular_lift(a::Array{fq_nmod_poly}, me::modular_env) -> Generic.Poly{nf_elem}

Apply the \code{modular_lift} function to each coeficient of $a$.
Computes a polynomial over the number field.
"""
function modular_lift(a::Array{fq_nmod_poly, 1}, me::modular_env)
  res = me.Kx()
  d = maximum([x.length for x = a])
  for i=0:d-1
    for j=1:me.ce.n
      ccall((:nmod_poly_set, :libflint), Nothing, (Ref{nmod_poly}, Ref{fq_nmod}), me.rp[j], coeff(a[j], i))
    end
    ap = crt(me.rp, me.ce)
    r = coeff(res, i)
    for j=0:ap.length-1
      u = ccall((:nmod_poly_get_coeff_ui, :libflint), UInt, (Ref{nmod_poly}, Int), ap, j)
      _num_setcoeff!(r, j, u)
    end
    setcoeff!(res, i, r)
  end
  return res
end

@doc Markdown.doc"""
  modular_proj(a::Generic.Mat{nf_elem}, me::modular_env) -> Array{Matrix}
  modular_proj(a::Generic.Mat{NfOrdElem}, me::modular_env) -> Array{Matrix}

Apply the \code{modular_proj} function to each entry of $a$.
Computes an array of matrices over the respective residue class fields.
"""
function modular_proj(a::Generic.Mat{nf_elem}, me::modular_env)
  Mp = fq_nmod_mat[]
  for i=1:me.ce.n
    push!(Mp, zero_matrix(me.fld[i], nrows(a), ncols(a)))
  end
  for i=1:nrows(a)
    for j=1:ncols(a)
      im =modular_proj(a[i,j], me)
      for k=1:me.ce.n
        setindex!(Mp[k], deepcopy(im[k]), i, j)
      end
    end
  end
  return Mp
end

function modular_proj(a::Generic.Mat{NfOrdElem}, me::modular_env)
  Mp = []
  for i=1:me.ce.n
    push!(Mp, zero_matrix(me.fld[i], nrows(a), ncols(a)))
  end
  for i=1:nrows(a)
    for j=1:ncols(a)
      im =modular_proj(me.K(a[i,j]), me)
      for k=1:me.ce.n
        setindex!(Mp[k], deepcopy(im[k]), i, j)
      end
    end
  end
  return Mp
end


using LinearAlgebra, Statistics
"""
    function read_gt_n_frq(gt::AbstractString, frq::AbstractString)
---
Read genotypes and their frequencies from file `gt` and `frq`.
Note gt can be prepared like below:
```
plink --species --bfile source --recode A --out target
cat target.raw | pathto/raw2gt target.frq > target.gt
```
See also the function `prepare_gt_from_plink_files`.
"""
function read_gt_n_frq(gt, frq)
    subtitle("Read genotypes and frequencies")
    item("Frequency related")
    # Frequency related
    p = Float64[]
    for x in eachline(frq)
        push!(p, parse(Float64, x))
    end
    twop = p.*2.
    done()

    item("Genotypes")
    nlc = length(readline(gt))
    fsz = stat(gt).size
    lsz = nlc + 1               # assuming Unix text format
    nid = Int(fsz/lsz)
    if(nid*lsz ≠ fsz)
        warning("Not a square file")
        return
    end
    message("NID: $nid;\tN_loci: $nlc")
    z = Array{Float64, 2}(undef, nlc, nid) # one column per ID
    fgt = open(gt, "r")
    for j in 1:nid
        t = read(fgt, lsz)
        for i in 1:nlc
            z[i, j] = Float64(t[i] - 0x30) # 0x30 == '0'
        end
    end
    close(fgt)
    return p, twop, z
end

"""
    vanRaden(Z, twop)
---
This function calculate **`G`** with van Raden 2008 on genotype **`Z`**.
method I.  Note **`Z`** here is ID column majored.

`G = Z`Z/2Σp_i(1-p_i)`

where **`Z`** contains the marker genotypes for all animals at all loci, corrected for
the allele frequency per locus, and `p_i` is the frequency of the allele at locus
*i* coded with the highest integer value.  **`Z`** is derived from the genotypes of the
animals by subtracting 2 times the allele frequency, that is `2p_i`, from matrix
**`X`**, which specifies the marker genotypes for each individual as 0, 1 or 2. Values
for `p_i` are calculated from the data (default), or can be specified in a file by
the user.

Pass copy(GT) to this function to avoid GT matrix modification.
"""
function vanRaden(Z::Array{Float64, 2}, twop::Array{Float64, 1})
    title("Calculate G with vanRaden method I")
    BLAS.set_num_threads(_nthreads)
    Z .-= twop
    s2pq = (1 .- .5 .* twop)'twop
    r2pq = 1. / s2pq
    G = Z'Z .* r2pq
end

"""
    compute_G(target, src...; add_diag=0.)
---
Given the `plink` file sets `src`.{bed,bim,fam}, this function calculate a **G**
matrix in 3 columns, i.e.,

- `raw col element-value`

and store the elements in `target`

This funciton utilizes Julia splatting.  You can input as many datasets as you have to feed the funciton.

If the inverse of **G** is weird or not possible, try a tiny `add_diag` to the diagonals.
"""
function compute_G(target, src...; add_diag=0.)
    title("Compute G with $src")
    item("Check if files available and qualified")
    nlc = Int[]
    for i in src
        if !isfile("$i.bim")
            warning("$i not available")
            return
        end
        push!(nlc, countlines("$i.bim"))
    end
    for n in nlc[2:end]
        if n ≠ nlc[1]
            warning("Files are different")
            return
        end
    end
    pathto, out = splitdir(target)
    if length(pathto) == 0
        warning("Writing G matrix in the current directory")
    else
        isdir(pathto) || mkdir(pathto)
    end
    done()

    item("Merging $src")
    tmp = mktempdir(".")
    write("$tmp/merge.lst", join(src, '\n'), '\n')
    run(`$plink --sheep --merge-list $tmp/merge.lst --out $tmp/plink`)
    ID = begin
        tv = String[]
        for line in eachline("$tmp/plink.fam")
            id = split(line)[2]
            push!(tv, id)
        end
        tv
    end
    run(`$plink --sheep --bfile $tmp/plink --recode A --out $tmp/plink`)
    run(pipeline("$tmp/plink.raw", `$bin_dir/raw2gt $tmp/plink.frq`, "$tmp/plink.gt"))
    _, twop, z = read_gt_n_frq("$tmp/plink.gt", "$tmp/plink.frq")

    done()

    item("Calculate $target")
    message("Using $_nthreads threads")
    G = vanRaden(z, vec(twop))
    open(target, "w") do io
        for i in 1:length(ID)
            for j in 1:i-1
                println(io, ID[i], ' ', ID[j], ' ', G[i, j])
            end
            println(io, ID[i], ' ', ID[i], ' ', G[i, i] + add_diag)
        end
    end
    done()

    rm(tmp, recursive=true, force=true)
end

"""
    compute_G(file)
---
Compute GRM with 'packed format', that is, genotypes is stored like

    012...
    210...

One row for an ID.

This function is for test purpose only.
It returns a GRM matrix of VR-I method.
"""
function compute_G(file)
    Z = Int8[]
    nid = 0
    for line in eachline(file)
        append!(Z, parse.(Int8, collect(line)))
        nid += 1
    end
    Z = Float64.(Z)
    Z = reshape(Z, :, nid)
    twop = mean(Z, dims=2)
    G = vanRaden(Z, vec(twop))
end


module Calculus2
    export
        @gradexpr,
        @hessexpr,
        abs2dual,
        absdual,
        conjdual,
        deparse,
        Dual,
        dual,
        Dual128,
        dual128,
        Dual64,
        dual64,
        dual_show,
        DualPair,
        epsilon,
        grad!,
        grad,
        gradexpr,
        gradfunc!,
        gradfunc,
        hess!,
        hess,
        hessexpr,
        hessfunc!,
        hessfunc,
        integrate,
        isdual,
        isgrad!,
        isgrad,
        ishess!,
        ishess,
        ∇!,
        ∇,
        ∫

    include("utils/maxdiff.jl")

    include("symbolic/symbolic.jl")
    include("symbolic/gradexpr.jl")
    include("symbolic/hessexpr.jl")
    include("symbolic/deparse.jl")

    include("finite/epsilon.jl")
    include("finite/grad.jl")
    include("finite/jacobian.jl")
    include("finite/hess.jl")

    include("automatic/dualnumbers.jl")
    include("automatic/grad.jl")

    include("integrate/simpsons.jl")
    include("integrate/montecarlo.jl")

    include("api/grad.jl")
    include("api/gradfunc.jl")
    include("api/gradexpr.jl")
    include("api/hess.jl")
    include("api/hessfunc.jl")
    include("api/hessexpr.jl")
    include("api/checker.jl")
    include("api/integrate.jl")
    include("api/unicode.jl")
end


using CellListMap
using StaticArrays
using LinearAlgebra: norm
import Random

#
# In this example we test the implementation in a simple packmol-like
# soft potential
#
function packmol(;parallel=false,sides=[46.4,31.5,18.7],tol=2.0,UnitCellType=OrthorhombicCell)

    function wrap(x,side)
        x = rem(x,side)
        if x >= side/2
            x -= side
        elseif x < -side/2
            x += side
        end
        return x
    end

    function mind(x,sides,tol)
        d = +Inf
        u = 0.
        for i in 1:length(x)-1
            for j in i+1:length(x)
                dd = norm(wrap.(x[i]-x[j],sides))
                d = min(d,dd)
                if dd < tol
                    u += (dd-tol)^2
                end
            end
        end
        return d, u
    end

    function fpair_cl(x,y,i,j,d2,f,box::Box)
        Δv = y - x
        d = sqrt(d2)
        fₓ = 2*(d - box.cutoff)*(Δv/d)
        f[i] += fₓ
        f[j] -= fₓ
        return f
    end

    function forces_cl!(f::Vector{T},x,box::Box,cl::CellList,fpair::F) where {T,F}
        fill!(f,zero(T))
        cl = UpdateCellList!(x,box,cl,parallel=parallel)
        map_pairwise!(
            (x,y,i,j,d2,f) -> fpair(x,y,i,j,d2,f,box),
            f, box, cl, parallel=parallel
        )
        return f
    end

    function u_pack(x,box::Box,cl::CellList)
        cl = UpdateCellList!(x,box,cl,parallel=false)
        u = map_pairwise(
            (x,y,i,j,d2,u) -> begin
                u += (sqrt(d2) - box.cutoff)^2 # objective function
                return u
            end,
            0., box, cl,
            parallel=parallel
        )
        return u
    end

    Random.seed!(321)

    x = [ sides .* rand(SVector{3,Float64}) for _ in 1:5000 ]

    mind_check, u_check = mind(x,sides,tol)

    box = Box(sides,tol,UnitCellType=UnitCellType)
    cl = CellList(x,box,parallel=parallel)
    u_pack = u_pack(x,box,cl)

    return u_pack ≈ u_check, u_pack, u_check, mind_check

end


import Multirate
import Multirate: NaiveResamplers
using Base.Test

function Base.isapprox( x1::Vector, x2::Vector )
    Nx1 = length( x1 )
    Nx2 = length( x2 )

    if Nx1 != Nx2
        @printf("x1 & x2 are different lengths vectors\n")
        return false
    end

    for i = 1:Nx1
        if !isapprox( x1[i], x2[i] )
            @printf( "Something went wrong at index %d\n", i )
            return false
        end
    end

    return true
end

Tx           = Float32
Th           = Float32
numFilters   = 32
x            = Tx[1.0:101]
resampleRate = 33/32

cutoffFreq      = 0.45
transitionWidth = 0.05
(hLen, β)       = Multirate.kaiserlength( transitionWidth, samplerate = numFilters )
hLen            = iceil(  hLen/numFilters  ) .* numFilters
h               = Multirate.firdes( hLen, cutoffFreq, DSP.kaiser, samplerate = numFilters, beta = β ) .* numFilters
h               = convert( Vector{Th}, h)

@time yNaive = NaiveResamplers.naivefilt( h, x, resampleRate, numFilters )
@time yArb   = Multirate.filt( h, x, resampleRate )

commonLen = min( length(yNaive), length(yArb) )
isapprox( yNaive[1:commonLen], yArb[1:commonLen] ) ||  display( [ [1:commonLen] yNaive[1:commonLen] yArb[1:commonLen] yNaive[1:commonLen].-yArb[1:commonLen]] )


const i_YSZ = 0
const i_LSM = 1
const i_hole = 2

Base.@kwdef mutable struct parameters
    R_YSZ::Float64 = 1/0.045 # S/cm
    R_pol_YSZ::Float64 = 0
    C_pol_YSZ::Float64 = 0.001
    #
    R_LSM::Float64 = 1/290 # S/cm
    R_pol_LSM::Float64 = 40
    C_pol_LSM::Float64 = 0.005
    #
    R_hole::Float64 = 1000000
end

function RC_el(R, C, w)
  return (R/(1 + C*R*w*im))
end

function get_Z_entry_from_material_matrix_codes(n1, n2, p::parameters)
    if      n1 == i_LSM
            if      n2 == i_LSM
              return w -> p.R_LSM/2
            elseif  n2 == i_YSZ
              return w -> p.R_LSM/2 + RC_el(p.R_pol_LSM, p.C_pol_LSM, w)
            elseif  n2 == i_hole
              return w ->  p.R_LSM/2
            end
    elseif  n1 == i_YSZ
            if      n2 == i_LSM
              return w -> p.R_YSZ/2 + RC_el(p.R_pol_YSZ, p.C_pol_YSZ, w)
            elseif  n2 == i_YSZ
              return w -> p.R_YSZ/2
            elseif  n2 == i_hole
              return w -> p.R_YSZ/2
            end
    elseif  n1 == i_hole
            return w -> p.R_hole/2
    else
        println("ERROR: get_Z_entry...")
    end
end

function file_to_matrix(path="src/geometry.png")
  RGB_m = load(path)
  m = Matrix{Int}(undef, size(RGB_m)...)
  for i in 1:size(m)[1]
    for j in 1:size(m)[2]
      (r,g,b) = (RGB_m[i,j].r, RGB_m[i,j].g, RGB_m[i,j].b)
      if (r > 0.5) && (g > 0.5) && (b > 0.5)
        m[i,j] = i_hole
      elseif (r > 0.5) && (g > 0.5)
        m[i,j] = i_YSZ
      else
        m[i,j] = i_LSM
      end
    end
  end
  return m
end

function matrix_to_file(path, matrix)
  save(path, map(x -> if     x == i_hole
                   RGB(1, 1, 1)
                elseif x == i_YSZ
                   RGB(1, 1, 0)
                elseif x == i_LSM
                   RGB(0, 0, 0)
                end, matrix)
  )
  return
end


using ImageFiltering: Kernel, imfilter, kernelfactors, centered

function nld_step_scalar(Ld, c, stepsize)
    dx = 0.5 * stepsize * (c[1:end, 1:end-1] + c[1:end, 2:end]) .* (Ld[1:end, 2:end] - Ld[1:end, 1:end-1])
    dy = 0.5 * stepsize * (c[1:end-1, 1:end] + c[2:end, 1:end]) .* (Ld[2:end, 1:end] - Ld[1:end-1, 1:end])

    Ld[1:end, 1:end-1] .+= dx
    Ld[1:end, 2:end] .-= dx
    Ld[1:end-1, 1:end] .+= dy
    Ld[2:end, 1:end] .-= dy
end

pm_g1_diffusivity(Lx, Ly, k) =
    calculate_diffusivity′(Lx, Ly, k) do dL
        exp(-dL)
    end

pm_g2_diffusivity(Lx, Ly, k) =
    calculate_diffusivity′(Lx, Ly, k) do dL
        1.0 / (1.0 + dL)
    end

weickert_diffusivity(Lx, Ly, k) =
    calculate_diffusivity′(Lx, Ly, k) do dL
        1.0 - exp(-3.315 / (dL * dL * dL * dL))
    end

charbonnier_diffusivity(Lx, Ly, k) =
    calculate_diffusivity′(Lx, Ly, k) do dL
        1.0 / sqrt(1.0 + dL)
    end

function calculate_diffusivity′(func, Lx, Ly, k)
    dst = zeros(size(Lx))
    invk2 = 1.0 / (k * k)
    nrow, ncol = size(Lx)
    @inbounds begin
        @simd for j = 1:nrow
            @simd for k = 1:ncol
                lx, ly = Lx[j, k], Ly[j, k]
                dL = (lx * lx + ly * ly) * invk2
                dst[j, k] = func(dL)
            end
        end
    end
    dst
end

function compute_k_percentile(img, perc; gscale = 1.0, nbins = 300)
    Lsmooth = imfilter(img, Kernel.gaussian(gscale))
    fy, fx = Kernel.scharr()
    fy .*=32
    fx .*=32
    # fx,fy = Kernel.ando3()

    w = (size(fx)[1] + 1) ÷ 2
    Lx = @view imfilter(Lsmooth, fx)[w:end-w+1, w:end-w+1]
    Ly = @view imfilter(Lsmooth, fy)[w:end-w+1, w:end-w+1]
    compute_k_percentile′(Lx, Ly, perc, nbins)
end

function compute_k_percentile′(Lx, Ly, perc, nbins = 300)
    hist = zeros(Int32, nbins)

    hmax = sqrt(maximum(y -> sum(x -> x^2, y), zip(Lx[:], Ly[:])))

    for (lx, ly) in zip(Lx[:], Ly[:])
        modg = sqrt(lx^2 + ly^2)
        if modg > 1e-10
            nbin = ceil(Int, nbins * (modg / hmax))  # limiting to nbins should not be necessary
            hist[nbin] += 1
        end
    end

    nthreshold = floor(Int, sum(hist) * perc)

    k = findfirst(hx -> hx > nthreshold, cumsum(hist))

    return if k == nothing
        0.03
    else
        (hmax * k) / nbins
    end
end

function demo_diffusivity_functions()
    imshow(Lsmooth)
    aa = pm_g1_diffusivity(Lx, Ly, 0.01)
    imshow(aa)
    aa = pm_g2_diffusivity(Lx, Ly, 0.01)
    imshow(aa)
    aa = weickert_diffusivity(Lx, Ly, 0.01)
    imshow(aa)
    aa = charbonnier_diffusivity(Lx, Ly, 0.01)
    imshow(aa)
end

function demo_k_percentile()
    gthresh = compute_k_percentile(img, 0.95)
    imshow(sqrt.(Lx .^ 2 + Ly .^ 2))
    imshow(sqrt.(Lx .^ 2 + Ly .^ 2) .< gthresh)
end

function demo_nld()
    Lflow = charbonnier_diffusivity(Lx, Ly, 0.01)

    imshow(copy(Lt))
    for _ = 1:3
        for _ = 1:100
            nld_step_scalar(Lt, Lflow, 0.05)
        end
        imshow(copy(Lt))
    end
end

function compute_derivative_kernels(scale)

    ksize = 3 + 2*(scale-1)

    ## The usual Scharr kernel
    if (scale == 1)
        # Kernel.scharr()
        fx, fy = Kernel.scharr()
        # (fx.*32, fy.*32)
        (fx, fy)
    end

    kx = centered(zeros(ksize))
    ky = centered(zeros(ksize))

    w = 10.0/3.0
    norm = 1.0/(2.0*scale*(w+2.0))

    # ky[1+0] = norm
    # ky[1+ksize÷2] = w*norm
    # ky[1+ksize-1] = norm

    # kx[1+0] = -1
    # kx[1+ksize÷2] = 0
    # kx[1+ksize-1] = 1
    ky[-ksize÷2] = norm
    ky[0] = w*norm
    ky[ksize÷2] = norm

    kx[-ksize÷2] = -1
    kx[0] = 0
    kx[ksize÷2] = 1
    # (kernelfactors((ky, kx)), kernelfactors((kx, ky)))
    (kernelfactors((ky, kx)), kernelfactors((kx, ky)))
end


"""
 Tests for SIE and SME methods
"""

import Statistics # for mean values of trajectories
import LinearAlgebra # for the normn
using StochasticDiffEq
using Test
using Random
using DiffEqDevTools
#using DiffEqGPU

function prob_func(prob, i, repeat)
    remake(prob,seed=seeds[i])
end

"""
 Test Scalar SDEs (oop)
"""

@info "Scalar oop noise"

# PC exercise 14.2.2
numtraj = Int(7e3)
u₀ = 0.1
f(u,p,t) = p[1]*u
g(u,p,t) = p[2]*u
dts = 1 .//2 .^(6:-1:3)
tspan = (0.0,1.0)
p = [3//2,1//100]

h1(z) = z
#analytical_sol(t) = E(f(X(t))) =

seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

prob = SDEProblem(f,g,u₀,tspan,p)
ensemble_prob = EnsembleProblem(prob;
        output_func = (sol,i) -> (h1(sol[end]),false),
        prob_func = prob_func
        )

sim = test_convergence(dts,ensemble_prob,SIEA(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=u₀.*exp(1.0*(p[1]))
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.35 # order is 2.35
println("SIEA:", sim.𝒪est[:weak_final])

sim = test_convergence(dts,ensemble_prob,SMEA(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=u₀.*exp(1.0*(p[1]))
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.35 # order is 2.34
println("SMEA:", sim.𝒪est[:weak_final])

sim = test_convergence(dts,ensemble_prob,SIEB(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=u₀.*exp(1.0*(p[1]))
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.35 # order is 2.35
println("SIEB:", sim.𝒪est[:weak_final])

sim = test_convergence(dts,ensemble_prob,SMEB(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=u₀.*exp(1.0*(p[1]))
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.35 # order is 2.34
println("SMEB:", sim.𝒪est[:weak_final])

"""
 Test Scalar SDEs (iip)
"""

@info "Scalar iip noise"

u₀ = [0.1]
f1!(du,u,p,t) =  (du[1] = p[1]*u[1])
g1!(du,u,p,t) =  (du[1] = p[2]*u[1])
dts = 1 .//2 .^(6:-1:3)
tspan = (0.0,1.0)
p = [3//2,1//100]

prob = SDEProblem(f1!,g1!,u₀,tspan,p)
ensemble_prob = EnsembleProblem(prob;
        output_func = (sol,i) -> (h1(sol[end][1]),false),
        prob_func = prob_func
        )

numtraj = Int(7e3)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

sim = test_convergence(dts,ensemble_prob,SIEA(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=u₀.*exp(1.0*(p[1]))
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.35 # order is 2.35
println("SIEA:", sim.𝒪est[:weak_final])

sim = test_convergence(dts,ensemble_prob,SMEA(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=u₀.*exp(1.0*(p[1]))
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.35 # order is 2.34
println("SMEA:", sim.𝒪est[:weak_final])

sim = test_convergence(dts,ensemble_prob,SIEB(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=u₀.*exp(1.0*(p[1]))
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.35 # order is 2.35
println("SIEB:", sim.𝒪est[:weak_final])

sim = test_convergence(dts,ensemble_prob,SMEB(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=u₀.*exp(1.0*(p[1]))
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.35 # order is 2.34
println("SMEB:", sim.𝒪est[:weak_final])

"""
 Test Diagonal noise SDEs (iip), SIAM Journal on Numerical Analysis, 47 (2009), pp. 1713–1738
"""

@info "Diagonal noise"

u₀ = [0.1,0.1]
function f3!(du,u,p,t)
  du[1] = 3//2*u[1]
  du[2] = 3//2*u[2]
end
function g3!(du,u,p,t)
  du[1] = 1//10*u[1]
  du[2] = 1//10*u[2]
end
dts = 1 .//2 .^(5:-1:1)
tspan = (0.0,1.0)

h3(z) = z^2 # == 1//10**exp(3//2*t) if h3(z) = z and  == 1//100**exp(301//100*t) if h3(z) = z^2 )

prob = SDEProblem(f3!,g3!,u₀,tspan)
ensemble_prob = EnsembleProblem(prob;
        output_func = (sol,i) -> (h3(sol[end][1]),false),
        prob_func = prob_func
        )

numtraj = Int(5e4)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

sim = test_convergence(dts,ensemble_prob,SIEA(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=1//100*exp(301//100)
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("SIEA:", sim.𝒪est[:weak_final])

sim = test_convergence(dts,ensemble_prob,SMEA(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=1//100*exp(301//100)
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("SMEA:", sim.𝒪est[:weak_final])

sim = test_convergence(dts,ensemble_prob,SIEB(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=1//100*exp(301//100)
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("SIEB:", sim.𝒪est[:weak_final])

sim = test_convergence(dts,ensemble_prob,SMEB(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=1//100*exp(301//100)
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("SMEB:", sim.𝒪est[:weak_final])


# defines the SimpleTropical module with the Tropical type

module SimpleTropical

import Base.isinf, Base.show, Base.+, Base.*, Base.inv, Base.==
import Base.isequal, Base.^, Base.convert

export Tropical, TropicalInf

struct Tropical{T<:Real} <: Number
  val::T
  inf_flag::Bool

  function Tropical{T}(xx::Real, ii::Bool=false) where T
    TT = typeof(xx)
    if isinf(xx) || ii
      return new(zero(TT),true)
    end
    return new(xx,false)
  end
end

Tropical(x::T) where T<:Real = Tropical{T}(x)
function Tropical(x::T, i::Bool) where T<:Real
  if i
    return Tropical{T}(zero(T),true)
  end
  return Tropical(x)
end

"""
`TropicalInf` is a constant that represents infinity in the tropical
semiring.
"""
const TropicalInf = Tropical{Bool}(0,true)

isinf(X::Tropical) = X.inf_flag

Base.promote_rule(::Type{Tropical{T}}, ::Type{S}) where {T<:Real, S<:Real} =
    Tropical{promote_type(T, S)}
Base.promote_rule(::Type{Tropical{T}}, ::Type{Tropical{S}}) where {T<:Real, S<:Real} =
    Tropical{promote_type(T, S)}

convert(::Type{Tropical}, x::T) where {T<:Real} = Tropical{T}(x)
convert(::Type{Tropical{T}}, x::S) where {T<:Real,S<:Tropical} =
    Tropical(convert(T, x.val), x.inf_flag)

function show(io::IO, t::Tropical)
  if isinf(t)
    print(io,"Tropical($(Char(8734)))")  # infinity character
  else
    print(io,"Tropical{$(typeof(t.val))}($(t.val))")
  end
end

function (+)(x::Tropical{T}, y::Tropical{T}) where {T}
  if isinf(x)
    if isinf(y)   # when X,Y both are infinite
      return Tropical(zero(T),true)  # create common infinite
    else
      return Tropical(y)
    end
  end

  if isinf(y)
    return Tropical(x)
  end

  return Tropical(min(x.val, y.val))
end
(+)(x::Tropical{T}, y::Tropical{S}) where {T,S} = +(promote(x, y)...)
(+)(x::Tropical{T}, y::Real) where T = +(promote(x, y)...)
(+)(x::Real, y::Tropical{T}) where T = +(promote(x, y)...)

function (*)(x::Tropical{T}, y::Tropical{T}) where {T}
  if isinf(x) || isinf(y)
    return Tropical(zero(T),true)
  end

  return Tropical(x.val + y.val)
end
(*)(x::Tropical{T}, y::Tropical{S}) where {T,S} = *(promote(x, y)...)
(*)(x::Tropical{T}, y::Real) where T = *(promote(x, y)...)
(*)(x::Real, y::Tropical{T}) where T = *(promote(x, y)...)

function inv(X::Tropical)
  @assert !isinf(X) "TropicalInf is not invertible"
  return Tropical(-X.val)
end

function (^)(X::Tropical, p::Integer)
  if isinf(X)
    @assert p>0 "Cannot raise tropical infinity to a nonpositive power."
    return X
  end

  return Tropical(X.val * p)
end

function isequal(X::Tropical, Y::Tropical)
  if !isinf(X) && !isinf(Y)
      return isequal(X.val, Y.val)
  else
      return isinf(X) && isinf(Y)
  end
end

function ==(X::Tropical, Y::Tropical)
  if !isinf(X) && !isinf(Y)
      return X.val == Y.val
  else
      return isinf(X) && isinf(Y)
  end
end

end # end of module


"""
    gewekediag(x::AbstractVector{<:Real}; first::Real=0.1, last::Real=0.5, kwargs...)

Compute the Geweke diagnostic from the `first` and `last` proportion of samples `x`.
"""
function gewekediag(x::AbstractVector{<:Real}; first::Real=0.1, last::Real=0.5, kwargs...)
    0 < first < 1 || throw(ArgumentError("`first` is not in (0, 1)"))
    0 < last < 1 || throw(ArgumentError("`last` is not in (0, 1)"))
    first + last <= 1 || throw(ArgumentError("`first` and `last` proportions overlap"))

    n = length(x)
    x1 = x[1:round(Int, first * n)]
    x2 = x[round(Int, n - last * n + 1):n]
    z =
        (Statistics.mean(x1) - Statistics.mean(x2)) /
        hypot(mcse(x1; kwargs...), mcse(x2; kwargs...))
    p = SpecialFunctions.erfc(abs(z) / sqrt(2))

    return (zscore=z, pvalue=p)
end


export get_quantified_variables_form,
    get_unbound_variables

"""
    get_quantified_variables_form(formula)

Return `formula` with all unbound variables bound to existential quantifiers.

The returned formula is equisatisfiable with `formula`.
"""
function get_quantified_variables_form(formula)
    variables = get_unbound_variables(formula)
    quantified_variables_form = formula
    for variable in variables
        quantified_variables_form = EFormula(variable, quantified_variables_form)
    end
    quantified_variables_form
end

function get_unbound_variables(formula, boundvars=Set{Variable}())
    @match formula begin
    f::Conjunction || f::Disjunction => begin
        unboundvars = get_unbound_variables.([f.formula1, f.formula2], boundvars)
        union(unboundvars...)
    end
    f::Negation => get_unbound_variables(f.formula, boundvars)
    f::AFormula || f::EFormula => begin
        push!(boundvars, f.variable)
        unboundvars = get_unbound_variables(f.formula, boundvars)
        pop!(boundvars, f.variable)
        unboundvars
    end
    f::AtomicFormula || f::Function => union(get_unbound_variables.(f.arguments, boundvars)...)
    f::Variable => f in boundvars ? Set() : Set([f])
    end
end


import Base.Iterators
function load( filename )

    open(filename) do f
        return [parse(Int,x) for x in readlines(f)]
    end

end

function check( numbers, item )

    for (i, n1) in enumerate(numbers[1:end-1])
        for n2 in numbers[i+1:end]
	    n1 + n2 == item && ( return true )
        end
    end
    return false

end

function part1(puzzle, preamble)
    numbers = load(puzzle)
    size = length(numbers)
    index = preamble + 1

    while index < size
        item = numbers[index]
        previous_numbers = numbers[index-preamble:index-1]
        if !check(previous_numbers, item)
            return item
        else
            index += 1
        end
    end
end

function part2(puzzle, preamble)

    numbers = load(puzzle)
    size = length(numbers)
    invalid = part1(puzzle, preamble)
    deque = numbers[1:2]
    s = sum(deque)
    index = 3
    while s !== invalid && index < size
        if s < invalid
           push!(deque, numbers[index])
           index += 1
        elseif s > invalid
           popfirst!(deque)
        end
        s = sum(deque)
    end
    return maximum(deque) + minimum(deque)

end
println(part1("input09.txt", 25))
println(part2("input09.txt", 25))


#=
---------------------------------------------------------------------
This function computes the RHS vector for a unified CGDG method for the
2D Wave Equation.

Written by F.X. Giraldo on July 12, 2021
           Department of Applied Mathematics
           Naval Postgraduate School
           Monterey; CA 93943-5216
---------------------------------------------------------------------
=#

include("create_rhs_volume.jl")
include("create_rhs_flux.jl")

function create_rhs(q,u,M,De_x,De_y,intma,periodicity,face,mapL,mapR,normals,jac_face,ωq,space_method,DFloat)

    #---------------------------Students Add their Contributions Here----------------------------------------
    #---------------------------Students Add their Contributions Here----------------------------------------
    #Construct Volume integral contribution
    #rhs = create_rhs_volume(q,u,De_x,De_y,intma,periodicity,DFloat)

    #Construct Flux integral contribution (will also work for CG but cancels completely due to periodicity)
    #if (space_method == "DG")
    #    rhs = create_rhs_flux!(rhs,q,u,face,normals,jac_face,ωq,mapL,mapR,intma,periodicity,DFloat)
    #end
    #---------------------------Students Add their Contributions Here----------------------------------------
    #---------------------------Students Add their Contributions Here----------------------------------------

    #Multiply by inverse Mass matrix
    rhs=M\rhs

    #Apply Periodicity
    Npoin=length(q)
    for i=1:Npoin
        rhs[i]=rhs[periodicity[i]]
    end

    #Return arrays
    return(rhs)
end


using .PyCall
using .Distributions
using .Tullio
using .FFTW
using .StatsBase

"""
Wrap scipy's inverse short time fourier transform.

    scipy.signal.istft(Zxx, fs=1.0, window='hann', nperseg=None, noverlap=None, nfft=None, input_onesided=True, boundary=True, time_axis=- 1, freq_axis=- 2)
"""
function istft(Zxx; kwargs...)
    signal.istft(Zxx; kwargs...)
end
export istft

function stft(x; kwargs...)
    signal.stft(x; kwargs...)
end
export stft

function linearnonstationaryspectrum(g::Function, Δt, T, windows::Int=20, normalise=true; doplot=false)
    #g(x, t) is a power function
    # Run a quick stft to get the right samplings and arrays
    t⃗ = 0.0:Δt:T
    x = randn(length(t⃗))
    f, t, Z = stft(x; fs=1.0/Δt, nperseg=floor(length(t⃗)/windows)-1)

    R = zeros(size(Z))
    Φ = rand(Uniform(0, 2π), size(Z)) # Uniform random phase

    for ti ∈ 1:size(Z, 2)
        ff = g.(f, (t[ti],))
        if isinf(ff[1])
            ff[1] = 0.0 # So we don't diverge at 0 frequency for power laws
        end
        if normalise
            E = sum(ff.*(f[2]-f[1])) # Approximate energy
            ff = ff./E # So the power spectrum has unit energy
        end
        R[:, ti] = sqrt.(ff) # Fourier amplitudes
    end
    if doplot
        p = heatmap(t[1:end], log10.(f[2:end]), log10.(R[2:end, :].^2), xguide="t", yguide="log₁₀(f)", colorbar_title="log₁₀(S)", yticks=-3:1, framestyle=:box, clims=(-1.5, 1.5))
        display(p)
    end
    ZZ = R.*exp.(Φ.*im)
    return istft(ZZ; fs=1.0/Δt, nperseg=floor(length(t⃗)/windows)-1)[end]
end
export linearnonstationaryspectrum

function selfAffine(β, Δt, T, args...)
    g(x, t) = 1/(x^(β(t)))
    linearnonstationaryspectrum(g, Δt, T, args...)
end
export selfAffine

"""
    Generate a locally self-affine time series, which has a spectrum that varies over time according to a power-law parameterised by β(t)

Similar to Theiler1992, who uses an inverse Fourier transform for random-phase (linear) surrogates from a given power spectrum. However, this process uses the inverse short time fourier transform to construct a non-stationary signal with a time-varying power spectrum.
"""
function selfAffine(P::Process)
    # Just one parameter, β
    seed(P.solver_rng) # Maybe this doesn't propagate to python...
    β = parameter_function(P)
    Δt = getΔt(P)
    T = gettmax(P) - gettransient_t0(P)
    return selfAffine(β, Δt, T)
end

selfAffineSim = Process(
    process = selfAffine,
    parameter_profile = ramp,
    parameter_profile_parameters = (0.0, 1.0, 0.0, 2000.0),
    transient_t0 = -100.0,
    t0 = 0.0,
    dt = 0.01,
    savedt = 0.1,
    tmax = 2000.0)
export selfAffineSim

"""
    fouriersurrogate(x::AbstractVector, t::AbstractVector; f::Function=x->x, g::Function=x->x; kwargs...)

Generate a non-stationary surrogate time series from an input `x` after performing manipulations in (short-time) Fourier space. The function `g(A, t)` is applied to the amplitudes of the Fourier coefficients whereas `h(𝜑, t)` is applied to the phases. After modifying the short-time fourier representation of the time series, the inverse short-time Fourier transform is used to construct the surrogate time series. `x` should be regularly sampled, and the time coordinates in `t` are crucial for correctly evaluating `g` and `h`. Optional `kwargs` are passed to python's `stft` and `istft` functions.
"""
function fouriersurrogate(x::AbstractVector, t::AbstractVector; g::Function=(𝑓, 𝑡, A)->A, h::Function=(𝑓, 𝑡, 𝜑)->𝜑, nperseg=1000, kwargs...)::AbstractVector
    Δt = (t[2] - t[1])
    fs = 1/Δt
    @assert all(diff(t, dims=1) .≈ Δt) # Check regularly sampled
    𝑓, 𝑡, 𝑍 = stft(x; fs, nperseg, kwargs...)
    𝑡 .+= t[1] # Offset the start time to match input signal
    gg .= g((𝑓,), 𝑡, (abs.(𝑍[:, j]),), (angle.(𝑍[:, j]),))
    hh .= h((𝑓,), 𝑡, (abs.(𝑍[:, j]),), (angle.(𝑍[:, j]),))
    𝑍 = gg.*cis.(hh)
    𝑡, x̂ = istft(𝑍; fs, nperseg, kwargs...)
    @assert (𝑡[2] - 𝑡[1]) == (t[2] - t[1])
    x̂ = x̂[1:length(x)] # The stft is zero-padded by python to fit evenly into the windows, so the istft is longer than the input x even though it has the same sampling frequency.
end

function fouriersurrogate(x::DimArray; kwargs...)
    t = timeDims(x)
    x̂ = fouriersurrogate(Array(x), t; kwargs...)
    DimArray(x̂, (Ti(t),))
end

export fouriersurrogate

"""
Function for randomising (from a uniform distribution between 0 and 2π) an angle with a probability of 𝜂.
"""
function corruptangle(𝜑::Number, 𝜂)
    𝜑 %= 2π
    rand() < 𝜂 ? rand()*2π : 𝜑
end
corruptangle(𝜑::Vector, 𝜂) = corruptangle.(𝜑, (𝜂,))

"""
Set phases below a threshold to 0 (synchronised)
Should prevent any issues with the desnity of randomisation interfering with power spectrum
"""
function thresholdsynchronise(𝑓::Vector, 𝜑::Number, ν)
    𝑓 < ν ? 0.0 : 𝜑
end
thresholdsynchronise(𝑓::Vector, 𝜑::Vector, ν) = thresholdsynchronise(𝑓, 𝜑, (ν,))

# """
# Take a Process and use the dark magic to produce a corrupted version, which has an extra parameter controlling the probability of the phase of each fourier coefficient being randomised. If planning to save and load this process, an instance of it must first be loaded so that the simulating function is exported
# """
# function corruptphase(P::Process, parameter_profile=constant, parameter_profile_parameters=0.0)
#     # In this case, savedt should really be a multiple of dt
#     S = P()
#     ps = getparameter_profile_parameters(P)
#     if getparameter_profile(P) isa Function || length(getparameter_profile(P)) == 1
#         S.parameter_profile = (getparameter_profile(P), parameter_profile)
#         ps = [ps]
#     else
#         S.parameter_profile = (getparameter_profile(P)..., parameter_profile)
#     end
#     S.parameter_profile_parameters = (ps..., parameter_profile_parameters)
#     pr = string(getprocess(P))
#     fname = Symbol("phaseCorrupted"*titlecase(pr))
#     pr = Symbol(pr)
#     @eval begin # Cry me a river
#         function ($fname)(S::Process)
#             D = S()
#             D.parameter_profile = getparameter_profile(S)[1:end-1]
#             if length(D.parameter_profile) == 1
#                 D.parameter_profile = D.parameter_profile[1]
#             end
#             D.parameter_profile_parameters = getparameter_profile_parameters(S)[1:end-1]
#             if length(D.parameter_profile_parameters) == 1
#                 D.parameter_profile_parameters = D.parameter_profile_parameters[1]
#             end
#             downsample = floor(D.savedt/D.dt) |> Int
#             D.savedt = D.dt # To keep accuracy for the transforms
#             D.process = $pr
#             x = timeseries(D, transient=true)
#             𝜂 = getparameter_profile(S)[end](getparameter_profile_parameters(S)[end]...)
#             ys = [Vector(x[:, i]) for i ∈ 1:size(x, 2)]
#             x = hcat([fouriersurrogate(y, times(D, transient=true); h=(𝑓, 𝑡, 𝜑)->corruptangle(𝜑, 𝜂(𝑡)), nperseg=1000*downsample) for y ∈ ys]...)
#             x = x[1:downsample:end, :]
#         end
#         export $fname
#     end
#     S.process = @eval $fname
#     return S
# end
# export corruptphase

"""
Take a Process and synchronise phases according to a frequency cutoff ν
"""
function synchronisephase(P::Process, parameter_profile=constant, parameter_profile_parameters=0.0)
    # In this case, savedt should really be a multiple of dt
    S = P()
    ps = getparameter_profile_parameters(P)
    if getparameter_profile(P) isa Function || length(getparameter_profile(P)) == 1
        S.parameter_profile = (getparameter_profile(P), parameter_profile)
        ps = [ps]
    else
        S.parameter_profile = (getparameter_profile(P)..., parameter_profile)
    end
    S.parameter_profile_parameters = (ps..., parameter_profile_parameters)
    pr = string(getprocess(P))
    fname = Symbol("phaseSynchronised"*titlecase(pr))
    pr = Symbol(pr)
    @eval begin
        function ($fname)(S::Process)
            D = S()
            D.parameter_profile = getparameter_profile(S)[1:end-1]
            if length(D.parameter_profile) == 1
                D.parameter_profile = D.parameter_profile[1]
            end
            D.parameter_profile_parameters = getparameter_profile_parameters(S)[1:end-1]
            if length(D.parameter_profile_parameters) == 1
                D.parameter_profile_parameters = D.parameter_profile_parameters[1]
            end
            downsample = floor(D.savedt/D.dt) |> Int
            D.savedt = D.dt # To keep accuracy for the transforms
            D.process = $pr
            x = timeseries(D, transient=true)
            ν = getparameter_profile(S)[end](getparameter_profile_parameters(S)[end]...)
            ys = [Vector(x[:, i]) for i ∈ 1:size(x, 2)]
            x = hcat([fouriersurrogate(y, times(D, transient=true); h=(𝑓, 𝑡, A, 𝜑)->thresholdsynchronise(𝑓, 𝜑, ν(𝑡)), nperseg=1000*downsample) for y ∈ ys]...)
            x = x[1:downsample:end, :]
        end
        export $fname
    end
    S.process = @eval $fname
    return S
end
export synchronisephase

"""
Now we switch gears and go for a simpler 'windowed fourier transform'; it's just the stft, but with non-overlapping, very wide windows that will match exactly the estimation windows we use. This is basically just Carl, since we won't worry about any nonsense like discontinuities at window edges.
"""
function wft(x, t; nwindows=20)
    Δt = (t[2] - t[1])
    fs = 1/Δt
    @assert all(diff(t, dims=1) .≈ Δt) # Check regularly sampled
    remainder = length(x)%nwindows
    remainder > 1 ? (@warn "The time series does not divide well into the number of windows supplied. The remainder is $remainder") : nothing
    𝑓 = rfftfreq((length(x) - remainder)÷nwindows, fs)
    𝑥, 𝑡 = [reshape(i[1:end-remainder], ((length(i)-remainder)÷nwindows, nwindows)) for i ∈ (x, t)]
    # Have to be careful with floating points here. Start by converting all times to integers, since we know they are equally spaced
    𝑡ᵢ = Int.(round.(𝑡./Δt))
    𝑡 = mean(𝑡ᵢ, dims=1)[:].*Δt # No (fewer?) floating point errors
    𝑍 = rfft(𝑥, 1)
    return (𝑓, 𝑡, 𝑍)
end
export wft

"""
Inverse of wft transformation. Just simple windows. If the original time series did not fit well into the number of windows, give a remainder here and it will append that many samples onto the end of the time series, which all have the same value as the last point in the reconstructed time series (shouldn't be a massive deal if the remainder is only 1 or 2 and the sampling period is small).
"""
function iwft(𝑡, 𝑍; remainder=0)
    if length(𝑡) == 1
        x̂ = irfft(𝑍, 2*size(𝑍, 1)-1, 1) # Not usually even in this case
        x̂ = reshape(x̂, length(x̂))
        t̂ = 1:length(x̂) # Can't do a whole lot better with one window time
    else
        nwindows = size(𝑍, 2)
        @assert nwindows == length(𝑡) # 𝑡 has window centres
        Δ𝑡 = (𝑡[2] - 𝑡[1])
        @assert all(diff(𝑡, dims=1) .≈ Δ𝑡) # Equally spaced window centres
        x̂ = irfft(𝑍, 2*size(𝑍, 1)-2, 1)
        x̂ = reshape(x̂, length(x̂))
        Δt̂ = Δ𝑡/((length(x̂))÷nwindows)
        N = Δ𝑡/Δt̂
        t̂₀ = 𝑡[1]-Δt̂*(N-1)/2
        append!(x̂, fill(x̂[end], remainder))
        t̂ = t̂₀:Δt̂:Δt̂*(N*nwindows-1 + remainder)
    end
    return (x̂, t̂)
end
iwft(𝑓, 𝑡, 𝑍; kwargs...) = iwft(𝑡, 𝑍; kwargs...) # 𝑓 not needed, but in case you want to pass wft result directly to iwft
iwft(𝑍::AbstractDimArray; kwargs...) = iwft(timeDims(𝑍), Array(𝑍); kwargs...)
export iwft

"""
Randomise phases of frequencies that constitute a propotion `p` of power (from the high to low frequencies)
"""
function thresholdcorrupt(𝜑, A, 𝑝) # ! SO P IS THE SQRT ROOT OF PROPORTIONAL POWER!!!
    𝜑 .= (𝜑 .+ 2π).%2π
    A[1] = 0 # The first amplitude is just the offset, which we don't care about (i.e. we want power of a mean-centred signal)
    psd = (A.^2)./sum(A.^2)
    cpsd = cumsum(psd)
    # display(plot(cpsd))
    # display(plot(A))
    idx = findfirst(cpsd .> 1-𝑝^2)
    isnothing(idx) ? (return 𝜑) : nothing
    phi = deepcopy(𝜑)
    phi[idx:end] .= rand(length(phi[idx:end])).*2π
    #display(plot([𝜑, phi]))
    return phi
end
thresholdcorrupt(𝑝) = (𝑓, 𝜑, 𝑍) -> thresholdcorrupt(𝜑, 𝑍, 𝑝)

function windowedfouriersurrogate(x::AbstractVector, t::AbstractVector; g::Function=(𝑓, 𝑡, A, 𝜑)->A, h::Function=(𝑓, 𝑡, A, 𝜑)->𝜑, nwindows=20)::AbstractVector
    Δt = (t[2] - t[1])
    fs = 1/Δt
    @assert all(diff(t, dims=1) .≈ Δt) # Check regularly sampled
    𝑓, 𝑡, 𝑍 = wft(x, t; nwindows)
    #display(heatmap(log10.(abs.(𝑍[2:end-1, :])), scale=:log))
    @tullio gg[j] := g(𝑓, 𝑡[j], abs.(𝑍[:, j]), angle.(𝑍[:, j]))
    @tullio hh[j] := h(𝑓, 𝑡[j], abs.(𝑍[:, j]), angle.(𝑍[:, j]))
    gg = hcat(gg...)
    hh = hcat(hh...)
    𝑍 = gg.*cis.(hh)
    #display(plot(angle.(𝑍)))
    #@tullio 𝑍[i, j] = g(𝑓[i], 𝑡[j], abs(𝑍[i, j]), angle(𝑍[i, j]))*cis(h(𝑓[i], 𝑡[j], abs(𝑍[i, j]), angle(𝑍[i, j])))
    #display(heatmap(log10.(abs.(𝑍[2:end-1, :])), scale=:log))
    remainder = Int(length(x)%nwindows)
    x̂, 𝑡 = iwft(𝑡, 𝑍, remainder=remainder) # For most situations in which this function is a good idea, this remainder will be 1
    @assert nwindows == 1 || (𝑡[2] - 𝑡[1]) ≈ (t[2] - t[1])
    return x̂
end

function windowedfouriersurrogate(x::DimArray; kwargs...)
    t = timeDims(x)
    x̂ = windowedfouriersurrogate(Array(x), t; kwargs...)
    DimArray(x̂, (Ti(t),))
end

export windowedfouriersurrogate

"""
Take a Process and produce a corrupted version, which has an extra parameter controlling the probability of the phase of each fourier coefficient being randomised. If planning to save and load this process, an instance of it must first be loaded so that the simulating function is exported
"""
function corruptphase(P::Process, parameter_profile=constant, parameter_profile_parameters=0.0; originalres=false, nwindows=20)
    # In this case, savedt should really be a multiple of dt
    S = P()
    ps = getparameter_profile_parameters(P)
    if getparameter_profile(P) isa Function || length(getparameter_profile(P)) == 1
        S.parameter_profile = (getparameter_profile(P), parameter_profile)
        ps = [ps]
    else
        S.parameter_profile = (getparameter_profile(P)..., parameter_profile)
    end
    S.parameter_profile_parameters = (ps..., parameter_profile_parameters)
    pr = string(getprocess(P))
    fname = Symbol("phaseCorrupted"*titlecase(pr))
    pr = Symbol(pr)
    @eval begin # Cry me a river
        function ($fname)(S::Process)
            D = S()
            D.parameter_profile = getparameter_profile(S)[1:end-1]
            if length(D.parameter_profile) == 1
                D.parameter_profile = D.parameter_profile[1]
            end
            D.parameter_profile_parameters = getparameter_profile_parameters(S)[1:end-1]
            if length(D.parameter_profile_parameters) == 1
                D.parameter_profile_parameters = D.parameter_profile_parameters[1]
            end
            if $originalres
                downsample = floor(D.savedt/D.dt) |> Int
                D.savedt = D.dt # To keep accuracy for the transforms
            else
                downsample = 1
            end
            D.process = $pr
            x = timeseries(D, transient=false)
            Y = timeseries(D, transient=true)
            𝜂 = getparameter_profile(S)[end](getparameter_profile_parameters(S)[end]...)
            ys = [Vector(x[:, i]) for i ∈ 1:size(x, 2)]
            x = hcat([windowedfouriersurrogate(y, times(D, transient=false); h=(𝑓, 𝑡, A, 𝜑)->thresholdcorrupt(𝜑, A, 𝜂(𝑡)), nwindows=$nwindows) for y ∈ ys]...)
            idx = findfirst(D.t0 .== times(D, transient=true))
            Y[idx:end, :] = x[1:downsample:end, :]
            return Y
        end
        export $fname
    end
    S.process = @eval $fname
    return S
end
export corruptphase

"""
E.g. Lorenz attractor with corrupted phases
"""
phaseCorruptedLorenzSim = corruptphase(lorenzSim(
    X0 = [0.0, -0.01, 9.0],
    parameter_profile = (constant, constant, constant),
    parameter_profile_parameters = (10.0, 28.0, 8/3),
    transient_t0 = -100.0,
    t0 = 0.0,
    dt = 0.001,
    savedt = 0.05,
    tmax = 1000.0,
    alg = AutoVern7(Rodas5()),
    solver_opts = Dict(:adaptive => true, :reltol => 1e-15)), rampInterval, (0.0, 0.25, 0.0, 1000.0))

"""
E.g. Lorenz attractor with synchronised phases
"""
phaseSynchronisedLorenzSim = synchronisephase(lorenzSim(
    X0 = [0.0, -0.01, 9.0],
    parameter_profile = (constant, constant, constant),
    parameter_profile_parameters = (10.0, 28.0, 8/3),
    transient_t0 = -100.0,
    t0 = 0.0,
    dt = 0.001,
    savedt = 0.05,
    tmax = 1000.0,
    alg = AutoVern7(Rodas5()),
    solver_opts = Dict(:adaptive => true, :reltol => 1e-15)), rampInterval, (0.0, 1.0, 0.0, 1000.0))

function lowpass(𝑓, x::Vector, p)
    idxs = 𝑓 .> p
    xx = deepcopy(x)
    xx[idxs] .= 0.0
    return xx
end

"""
Pass a process through a low pass fourier filter
"""
function lowpass(P::Process, parameter_profile=constant, parameter_profile_parameters=0.0; nwindows=20)
    # In this case, savedt should really be a multiple of dt
    S = P()
    ps = getparameter_profile_parameters(P)
    if getparameter_profile(P) isa Function || length(getparameter_profile(P)) == 1
        S.parameter_profile = (getparameter_profile(P), parameter_profile)
        ps = [ps]
    else
        S.parameter_profile = (getparameter_profile(P)..., parameter_profile)
    end
    S.parameter_profile_parameters = (ps..., parameter_profile_parameters)
    pr = string(getprocess(P))
    fname = Symbol("lowpass"*titlecase(pr))
    pr = Symbol(pr)
    @eval begin
        function ($fname)(S::Process)
            D = S()
            D.parameter_profile = getparameter_profile(S)[1:end-1]
            if length(D.parameter_profile) == 1
                D.parameter_profile = D.parameter_profile[1]
            end
            D.parameter_profile_parameters = getparameter_profile_parameters(S)[1:end-1]
            if length(D.parameter_profile_parameters) == 1
                D.parameter_profile_parameters = D.parameter_profile_parameters[1]
            end
            D.process = $pr
            x = timeseries(D, transient=false)
            Y = timeseries(D, transient=true)
            𝜂 = getparameter_profile(S)[end](getparameter_profile_parameters(S)[end]...)
            ys = [Vector(x[:, i]) for i ∈ 1:size(x, 2)]
            x = hcat([windowedfouriersurrogate(y, times(D, transient=false); g=(𝑓, 𝑡, A, 𝜑)->lowpass(𝑓, A, 𝜂(𝑡)), h=(𝑓, 𝑡, A, 𝜑)->lowpass(𝑓, 𝜑, 𝜂(𝑡)), nwindows=$nwindows) for y ∈ ys]...)
            idx = findfirst(D.t0 .== times(D, transient=true))
            Y[idx:end, :] = x
            return Y
        end
        export $fname
    end
    S.process = @eval $fname
    return S
end
export lowpass


#Simple methods

export eigvals_power, eigvals_ii, eigvals_rqi

function powm_method!{T}(log::ConvergenceHistory, K::KrylovSubspace{T};
    tol::Real=eps(T)*K.n^3, maxiter::Int=K.n
    )
    θ = zero(T)
    v = Array(T, K.n)
    for iter=1:maxiter
        nextiter!(log,mvps=1)
        v = lastvec(K)
        y = nextvec(K)
        θ = dot(v, y)
        resnorm = real(norm(y - θ*v))
        push!(log, :resnorm, resnorm)
        resnorm <= tol*abs(θ) && (setconv(log, resnorm >= 0); break)
        appendunit!(K, y)
    end
    Eigenpair(θ, v)
end

function eigvals_power(A, x=nothing; tol::Real=size(A,1)^3*eps(), maxiter::Int=size(A,1))
    K = KrylovSubspace(A, 1)
    x==nothing ? initrand!(K) : init!(K, x/norm(x))
    eigvals_power(K; tol=tol, maxiter=maxiter)
end
function eigvals_power(K::KrylovSubspace, x=nothing;
    tol::Real=size(A,1)^3*eps(), maxiter::Int=size(A,1)
    )
    history = ConvergenceHistory()
    history[:tol] = tol
    reserve!(history,:resnorm, maxiter)
    pair = powm_method!(history, K; tol=tol, maxiter=maxiter)
    pair, history
end

#Inverse iteration/inverse power method
function invpowm_method!{T}(log::ConvergenceHistory, K::KrylovSubspace{T}, σ::Number=zero(T);
    tol::Real=eps(T)*K.n^3, maxiter::Int=K.n
    )
    θ = zero(T)
    v = Array(T, K.n)
    y = Array(T, K.n)
    σ = convert(T, σ)
    for iter=1:maxiter
        nextiter!(log,mvps=1)
        v = lastvec(K)
        y = (K.A-σ*eye(K))\v
        θ = dot(v, y)
        resnorm = norm(y - θ*v)
        push!(log, :resnorm, resnorm)
        resnorm <= tol*abs(θ) && (setconv(log, resnorm >= 0); break)
        appendunit!(K, y)
    end
    shrink!(log)
    Eigenpair(σ+1/θ, y/θ)
end

function eigvals_ii(A, σ::Number, x=nothing; tol::Real=eps(), maxiter::Int=size(A,1))
    K = KrylovSubspace(A, 1)
    x==nothing ? initrand!(K) : init!(K, x/norm(x))
    eigvals_ii(K, σ; tol=tol, maxiter=maxiter)
end
function eigvals_ii(K::KrylovSubspace, σ::Number, x=nothing;
    tol::Real=size(A,1)^3*eps(), maxiter::Int=size(A,1)
    )
    history = ConvergenceHistory()
    history[:tol] = tol
    reserve!(history,:resnorm, maxiter)
    pair = invpowm_method!(history, K, σ; tol=tol, maxiter=maxiter)
    pair, history
end

#Rayleigh quotient iteration
#XXX Doesn't work well
function rqi_method!{T}(log::ConvergenceHistory, K::KrylovSubspace{T}, σ::Number; tol::Real=eps(T), maxiter::Int=K.n)
    v = lastvec(K)
    ρ = dot(v, nextvec(K))
    for iter=1:maxiter
        nextiter!(log,mvps=1)
        y = (K.A-ρ*eye(K))\v
        θ = norm(y)
        ρ += dot(y,v)/θ^2
        v = y/θ
        resnorm=1/θ
        push!(log,:resnorm,resnorm)
        θ >= 1/tol && (setconv(log, resnorm >= 0); break)
    end
    shrink!(history)
    Eigenpair(ρ, v)
end

function eigvals_rqi(A, σ::Number, x=nothing; tol::Real=eps(), maxiter::Int=size(A,1))
    K = KrylovSubspace(A, 1)
    x==nothing ? initrand!(K) : init!(K, x/norm(x))
    eigvals_rqi(K, σ; tol=tol, maxiter=maxiter)
end
function eigvals_rqi{T}(K::KrylovSubspace{T}, σ::Number, x=nothing; tol::Real=eps(T), maxiter::Int=K.n)
    history = ConvergenceHistory()
    history[:tol] = tol
    reserve!(history,:resnorm, maxiter)
    rqi_method(history, K, σ, tol=tol, maxiter=maxiter)
    x, history
end


@doc raw"""
    HeisenbergGroup{n} <: AbstractDecoratorManifold{ℝ}

Heisenberg group `HeisenbergGroup(n)` is the group of ``(n+2) × (n+2)`` matrices [^BinzPods2008]

```math
\begin{bmatrix} 1 & \mathbf{a} & c \\
\mathbf{0} & I_n & \mathbf{b} \\
0 & \mathbf{0} & 1 \end{bmatrix}
```

where ``I_n`` is the ``n×n`` unit matrix, ``\mathbf{a}`` is a row vector of length ``n``,
``\mathbf{b}`` is a column vector of length ``n`` and ``c`` is a real number.
The group operation is matrix multiplication.

The left-invariant metric on the manifold is used.

[^BinzPods2008]:
    > E. Binz and S. Pods, The Geometry of Heisenberg Groups: With Applications in Signal
    > Theory, Optics, Quantization, and Field Quantization. American Mathematical Soc., 2008.
"""
struct HeisenbergGroup{n} <: AbstractDecoratorManifold{ℝ} end

function active_traits(f, ::HeisenbergGroup, args...)
    return merge_traits(
        IsGroupManifold(MultiplicationOperation()),
        IsEmbeddedManifold(),
        HasLeftInvariantMetric(),
    )
end

function _heisenberg_a_view(::HeisenbergGroup{n}, p) where {n}
    return view(p, 1, 2:(n + 1))
end
function _heisenberg_b_view(::HeisenbergGroup{n}, p) where {n}
    return view(p, 2:(n + 1), n + 2)
end

HeisenbergGroup(n::Int) = HeisenbergGroup{n}()

function check_point(G::HeisenbergGroup{n}, p; kwargs...) where {n}
    if !isone(p[1, 1])
        return DomainError(
            p[1, 1],
            "The matrix $(p) does not lie on $(G), since p[1, 1] is not equal to 1.",
        )
    end
    if !isone(p[n + 2, n + 2])
        return DomainError(
            p[n + 2, n + 2],
            "The matrix $(p) does not lie on $(G), since p[n+2, n+2] is not equal to 1.",
        )
    end
    if !iszero(p[2:(n + 2), 1])
        return DomainError(
            norm(iszero(p[2:(n + 2), 1])),
            "The matrix $(p) does not lie on $(G), since p[2:(n + 2), 1] is not equal to 0.",
        )
    end
    if !iszero(p[n + 2, 1:(n + 1)])
        return DomainError(
            norm(iszero(p[n + 2, 1:(n + 1)])),
            "The matrix $(p) does not lie on $(G), since p[n + 2, 1:(n+1)] is not equal to 0.",
        )
    end
    if !isapprox(I, p[2:(n + 1), 2:(n + 1)])
        return DomainError(
            det(p[2:(n + 1), 2:(n + 1)] - I),
            "The matrix $(p) does not lie on $(G), since p[2:(n+1), 2:(n+1)] is not an identity matrix.",
        )
    end

    return nothing
end

function check_vector(G::HeisenbergGroup{n}, p, X; kwargs...) where {n}
    if !iszero(X[1, 1])
        return DomainError(
            X[1, 1],
            "The matrix $(X) does not lie in the tangent space of $(G), since X[1, 1] is not equal to 0.",
        )
    end
    if !iszero(X[n + 2, n + 2])
        return DomainError(
            X[n + 2, n + 2],
            "The matrix $(X) does not lie in the tangent space of $(G), since X[n+2, n+2] is not equal to 0.",
        )
    end
    if !iszero(X[2:(n + 2), 1:(n + 1)])
        return DomainError(
            norm(X[2:(n + 2), 1:(n + 1)]),
            "The matrix $(X) does not lie in the tangent space of $(G), since X[2:(n + 2), 1:(n + 1)] is not a zero matrix.",
        )
    end
    return nothing
end

embed(::HeisenbergGroup, p) = p
embed(::HeisenbergGroup, p, X) = X

@doc raw"""
    get_coordinates(M::HeisenbergGroup, p, X, ::DefaultOrthonormalBasis{ℝ,TangentSpaceType})

Get coordinates of tangent vector `X` at point `p` from the [`HeisenbergGroup`](@ref) `M`.
Given a matrix
```math
\begin{bmatrix} 1 & \mathbf{a} & c \\
\mathbf{0} & I_n & \mathbf{b} \\
0 & \mathbf{0} & 1 \end{bmatrix}
```
the coordinates are concatenated vectors ``\mathbf{a}``, ``\mathbf{b}``, and number ``c``.
"""
get_coordinates(::HeisenbergGroup, p, X, ::DefaultOrthonormalBasis{ℝ,TangentSpaceType})

function get_coordinates_orthonormal(M::HeisenbergGroup{n}, p, X, ::RealNumbers) where {n}
    return vcat(_heisenberg_a_view(M, X), _heisenberg_b_view(M, X), X[1, n + 2])
end

function get_coordinates_orthonormal!(
    M::HeisenbergGroup{n},
    Xⁱ,
    p,
    X,
    ::RealNumbers,
) where {n}
    Xⁱ[1:n] .= _heisenberg_a_view(M, X)
    Xⁱ[(n + 1):(2 * n)] .= _heisenberg_b_view(M, X)
    Xⁱ[2 * n + 1] = X[1, n + 2]
    return Xⁱ
end

get_embedding(::HeisenbergGroup{n}) where {n} = Euclidean(n + 2, n + 2)

@doc raw"""
    get_vector(M::HeisenbergGroup, p, Xⁱ, ::DefaultOrthonormalBasis{ℝ,TangentSpaceType})

Get tangent vector with coordinates `Xⁱ` at point `p` from the [`HeisenbergGroup`](@ref) `M`.
Given a vector of coordinates ``\begin{bmatrix}\mathbb{a} & \mathbb{b} & c\end{bmatrix}`` the tangent vector is equal to
```math
\begin{bmatrix} 1 & \mathbf{a} & c \\
\mathbf{0} & I_n & \mathbf{b} \\
0 & \mathbf{0} & 1 \end{bmatrix}
```
"""
get_vector(M::HeisenbergGroup, p, c, ::DefaultOrthonormalBasis{ℝ,TangentSpaceType})

function get_vector_orthonormal(::HeisenbergGroup{n}, p, Xⁱ, ::RealNumbers) where {n}
    return [
        0 Xⁱ[1:n] Xⁱ[2 * n + 1]
        zeros(n, n + 1) Xⁱ[(n + 1):(2 * n)]'
        zeros(1, n + 2)
    ]
end

function get_vector_orthonormal!(::HeisenbergGroup{n}, X, p, Xⁱ, ::RealNumbers) where {n}
    fill!(X, 0)
    X[1, 2:(n + 1)] .= Xⁱ[1:n]
    X[2:(n + 1), n + 2] .= Xⁱ[(n + 1):(2 * n)]
    X[1, n + 2] = Xⁱ[2 * n + 1]
    return X
end

@doc raw"""
    exp_lie(M::HeisenbergGroup, X)

Lie group exponential for the [`HeisenbergGroup`](@ref) `M` of the vector `X`.
The formula reads
```math
\exp\left(\begin{bmatrix} 0 & \mathbf{a} & c \\
\mathbf{0} & 0_n & \mathbf{b} \\
0 & \mathbf{0} & 0 \end{bmatrix}\right) = \begin{bmatrix} 1 & \mathbf{a} & c + \mathbf{a}⋅\mathbf{b}/2 \\
\mathbf{0} & I_n & \mathbf{b} \\
0 & \mathbf{0} & 1 \end{bmatrix}
```
where ``I_n`` is the ``n×n`` identity matrix, ``0_n`` is the ``n×n`` zero matrix
and ``\mathbf{a}⋅\mathbf{b}`` is dot product of vectors.
"""
exp_lie(M::HeisenbergGroup, X)

function exp_lie!(M::HeisenbergGroup{n}, q, X) where {n}
    copyto!(q, I)
    a_view = _heisenberg_a_view(M, X)
    b_view = _heisenberg_b_view(M, X)
    q[1, 2:(n + 1)] .= a_view
    q[2:(n + 1), n + 2] .= b_view
    q[1, n + 2] = X[1, n + 2] + dot(a_view, b_view) / 2
    return q
end

@doc raw"""
    exp(M::HeisenbergGroup, p, X)

Exponential map on the [`HeisenbergGroup`](@ref) `M` with the left-invariant metric.
The expression reads
```math
\exp_{\begin{bmatrix} 1 & \mathbf{a}_p & c_p \\
\mathbf{0} & I_n & \mathbf{b}_p \\
0 & \mathbf{0} & 1 \end{bmatrix}}\left(\begin{bmatrix} 0 & \mathbf{a}_X & c_X \\
\mathbf{0} & 0_n & \mathbf{b}_X \\
0 & \mathbf{0} & 0 \end{bmatrix}\right) =
\begin{bmatrix} 1 & \mathbf{a}_p + \mathbf{a}_X & c_p + c_X + \mathbf{a}_X⋅\mathbf{b}_X/2 + \mathbf{a}_p⋅\mathbf{b}_X \\
\mathbf{0} & I_n & \mathbf{b}_p + \mathbf{b}_X \\
0 & \mathbf{0} & 1 \end{bmatrix}
```
where ``I_n`` is the ``n×n`` identity matrix, ``0_n`` is the ``n×n`` zero matrix
and ``\mathbf{a}⋅\mathbf{b}`` is dot product of vectors.
"""
exp(M::HeisenbergGroup, p, X)

function exp!(M::HeisenbergGroup{n}, q, p, X) where {n}
    copyto!(q, I)
    a_p_view = _heisenberg_a_view(M, p)
    b_p_view = _heisenberg_b_view(M, p)
    a_X_view = _heisenberg_a_view(M, X)
    b_X_view = _heisenberg_b_view(M, X)
    q[1, 2:(n + 1)] .= a_p_view .+ a_X_view
    q[2:(n + 1), n + 2] .= b_p_view .+ b_X_view
    q[1, n + 2] =
        p[1, n + 2] + X[1, n + 2] + dot(a_X_view, b_X_view) / 2 + dot(a_p_view, b_X_view)
    return q
end

@doc raw"""
    injectivity_radius(M::HeisenbergGroup)

Return the injectivity radius on the [`HeisenbergGroup`](@ref) `M`, which is ``∞``.
"""
injectivity_radius(::HeisenbergGroup) = Inf

function inner(M::HeisenbergGroup{n}, p, X, Y) where {n}
    X_a_view = _heisenberg_a_view(M, X)
    X_b_view = _heisenberg_b_view(M, X)
    Y_a_view = _heisenberg_a_view(M, Y)
    Y_b_view = _heisenberg_b_view(M, Y)
    return dot(X_a_view, Y_a_view) +
           dot(X_b_view, Y_b_view) +
           X[1, 2 * n + 1] * Y[1, 2 * n + 1]
end

@doc raw"""
    log(G::HeisenbergGroup, p, q)

Compute the logarithmic map on the [`HeisenbergGroup`](@ref) group.
The formula reads
```math
\log_{\begin{bmatrix} 1 & \mathbf{a}_p & c_p \\
\mathbf{0} & I_n & \mathbf{b}_p \\
0 & \mathbf{0} & 1 \end{bmatrix}}\left(\begin{bmatrix} 1 & \mathbf{a}_q & c_q \\
\mathbf{0} & I_n & \mathbf{b}_q \\
0 & \mathbf{0} & 1 \end{bmatrix}\right) =
\begin{bmatrix} 0 & \mathbf{a}_q - \mathbf{a}_p & c_q - c_p + \mathbf{a}_p⋅\mathbf{b}_p - \mathbf{a}_q⋅\mathbf{b}_q - (\mathbf{a}_q - \mathbf{a}_p)⋅(\mathbf{b}_q - \mathbf{b}_p) / 2 \\
\mathbf{0} & 0_n & \mathbf{b}_q - \mathbf{b}_p \\
0 & \mathbf{0} & 0 \end{bmatrix}
```
where ``I_n`` is the ``n×n`` identity matrix, ``0_n`` is the ``n×n`` zero matrix
and ``\mathbf{a}⋅\mathbf{b}`` is dot product of vectors.
"""
log(::HeisenbergGroup, p, q)

function log!(M::HeisenbergGroup{n}, X, p, q) where {n}
    fill!(X, 0)
    a_p_view = _heisenberg_a_view(M, p)
    b_p_view = _heisenberg_b_view(M, p)
    a_q_view = _heisenberg_a_view(M, q)
    b_q_view = _heisenberg_b_view(M, q)
    X[1, 2:(n + 1)] .= a_q_view .- a_p_view
    X[2:(n + 1), n + 2] .= b_q_view .- b_p_view
    pinvq_c = dot(a_p_view, b_p_view) - p[1, n + 2] + q[1, n + 2] - dot(a_p_view, b_q_view)
    X[1, n + 2] = pinvq_c - dot(a_q_view - a_p_view, b_q_view - b_p_view) / 2
    return X
end

@doc raw"""
    log_lie(M::HeisenbergGroup, p)

Lie group logarithm for the [`HeisenbergGroup`](@ref) `M` of the point `p`.
The formula reads
```math
\log\left(\begin{bmatrix} 1 & \mathbf{a} & c \\
\mathbf{0} & I_n & \mathbf{b} \\
0 & \mathbf{0} & 1 \end{bmatrix}\right) =
\begin{bmatrix} 0 & \mathbf{a} & c - \mathbf{a}⋅\mathbf{b}/2 \\
\mathbf{0} & 0_n & \mathbf{b} \\
0 & \mathbf{0} & 0 \end{bmatrix}
```
where ``I_n`` is the ``n×n`` identity matrix, ``0_n`` is the ``n×n`` zero matrix
and ``\mathbf{a}⋅\mathbf{b}`` is dot product of vectors.
"""
log_lie(M::HeisenbergGroup, p)

function log_lie!(M::HeisenbergGroup{n}, X, p) where {n}
    fill!(X, 0)
    view_a_X = _heisenberg_a_view(M, X)
    view_b_X = _heisenberg_b_view(M, X)
    view_a_X .= _heisenberg_a_view(M, p)
    view_b_X .= _heisenberg_b_view(M, p)
    X[1, n + 2] = p[1, n + 2] - dot(view_a_X, view_b_X) / 2
    return X
end

function log_lie!(::HeisenbergGroup, X, ::Identity{MultiplicationOperation})
    fill!(X, 0)
    return X
end

manifold_dimension(::HeisenbergGroup{n}) where {n} = 2 * n + 1

parallel_transport_to(::HeisenbergGroup, p, X, q) = X

parallel_transport_to!(::HeisenbergGroup, Y, p, X, q) = copyto!(Y, X)

"""
    project(M::HeisenbergGroup{n}, p)

Project a matrix `p` in the Euclidean embedding onto the [`HeisenbergGroup`](@ref) `M`.
Sets the diagonal elements to 1 and all non-diagonal elements except the first row and the
last column to 0.
"""
function project(M::HeisenbergGroup{n}, p) where {n}
    return [
        1 p[1, 2:(n + 2)]'
        zeros(n, 1) Matrix(I, n, n) _heisenberg_b_view(M, p)
        zeros(1, n + 1) 1
    ]
end
"""
    project(M::HeisenbergGroup{n}, p, X)

Project a matrix `X` in the Euclidean embedding onto the Lie algebra of
[`HeisenbergGroup`](@ref) `M`.
Sets the diagonal elements to 0 and all non-diagonal elements except the first row and the
last column to 0.
"""
function project(M::HeisenbergGroup{n}, p, X) where {n}
    return [
        0 X[1, 2:(n + 2)]'
        zeros(n, n + 1) _heisenberg_b_view(M, X)
        zeros(1, n + 1) 0
    ]
end

function project!(M::HeisenbergGroup{n}, q, p) where {n}
    copyto!(q, I)
    q[1, 2:(n + 2)] .= p[1, 2:(n + 2)]
    q[2:(n + 1), n + 2] .= _heisenberg_b_view(M, p)
    return q
end
function project!(M::HeisenbergGroup{n}, Y, p, X) where {n}
    fill!(Y, 0)
    Y[1, 2:(n + 2)] .= X[1, 2:(n + 2)]
    Y[2:(n + 1), n + 2] .= _heisenberg_b_view(M, X)
    return Y
end

Base.show(io::IO, ::HeisenbergGroup{n}) where {n} = print(io, "HeisenbergGroup($n)")

translate_diff(::HeisenbergGroup, p, q, X, ::LeftAction) = X
translate_diff(::HeisenbergGroup, p, q, X, ::RightAction) = p \ X * p

function translate_diff!(G::HeisenbergGroup, Y, p, q, X, conv::ActionDirection)
    return copyto!(Y, translate_diff(G, p, q, X, conv))
end


instructions = String[]
global line = chomp(readline());
while line != ""
    global line
    push!(instructions, line)
    global line = chomp(readline());
end

called = Int[]
global inst = 0
global accumulator = 0
while !(inst in called)
    push!(called, inst)
    elems = split(instructions[inst + 1])
    op = elems[1]
    amount = parse(Int, elems[2])
    if op == "acc"
        global accumulator += amount
    end

    if op == "jmp"
        global inst += amount
    else
        global inst += 1
    end
end

println(accumulator)


using NCDatasets

const INFILE = "SOCATv5.tsv"
const OUTFILE = "daily.nc"
const CELLSIZE = 2.5
const STARTYEAR = 1985
const ENDYEAR = 2016

# Zero-based day of year of each month
const MONTHSTARTS = [0, 31, 59, 90, 120, 151, 181, 212, 243, 273, 304, 334]
const LEAPMONTHSTARTS = [0, 31, 60, 91, 121, 152, 182, 213, 244, 274, 305, 335]

# Days of the year
DAYSTARTS = range(1, step=1, length=365)

# Calculate the conversion of leap year days to 365 'normal' year days
# Each leap year day lasts 1 1/365 normal days
LEAPDAYSTARTS = Array{Float64}(undef, 365)

for i in 1:365
    LEAPDAYSTARTS[i] = i + (((366 / 365) / 365) * (i - 1))
end

# Determine whether or not a given year is a leap year
function isleapyear(year::Int64)::Bool
    local leapyear::Bool = false
    if rem(year, 4) == 0
        leapyear = true
        if rem(year, 100) == 0
            if rem(year, 400) != 0
                leapyear = false
            end
        end
    end

    leapyear
end

# Calculate the Julian day from a year/month/day set
function calcjdate(year::Int64, month::Int64, day::Int64)::Int64
    local result::Int64 = 0

    if isleapyear(year)
        result = LEAPMONTHSTARTS[month]
    else
        result = MONTHSTARTS[month]
    end

    result = result + day
end

# Calculate the day of the year for a given date
function getdayindex(date, days)::Float64
    return findall(days .<= date)[end]
end

# Calcaulte the nth day of the complete data set
# Days are only calculated between the start and end years
function getdateindex(year::Int64, month::Int64, day::Int64)::Int64

    local jdate::Float64 = calcjdate(year, month, day)
    local index::Float64 = -1
    local dayindex::Float64 = 0

    if year >= STARTYEAR && year <= ENDYEAR
        if isleapyear(year)
            dayindex = getdayindex(jdate, DAYSTARTS)
        else
            dayindex = getdayindex(jdate, LEAPDAYSTARTS)
        end

        index = ((year - STARTYEAR) * 365) + dayindex
    end

    return convert(Int64, floor(index))
end

# Calculate the cell indices from a longitude and latitude
function getcellindex(longitude::Float64, latitude::Float64)::Tuple{Int64, Int64}
    local loncell::Int64 = floor(longitude / CELLSIZE) + 1
    local latcell::Int64 = floor((latitude + 90) / CELLSIZE) + 1

    if loncell == 145
        loncell = 1
    end

    return loncell, latcell
end

function applydataset(datasettotals::Array{Float64, 3}, datasetcounts::Array{Int64, 3}, overalltotals::Array{Float64, 3}, overallcounts::Array{Int64, 3})
    datasetmean::Array{Float64, 3} = datasettotals ./ datasetcounts
    datasetmean[isnan.(datasetmean)] .= 0

    overalltotals .= overalltotals .+ datasetmean
    datasetmean[datasetmean .> 0] .= 1
    overallcounts .= overallcounts .+ datasetmean
end

function run()

    totaldays::Int64 = (ENDYEAR - STARTYEAR + 1) * 365

    # Output data set
    overallcelltotals::Array{Float64, 3} = zeros(convert(Int64, 360 / CELLSIZE), convert(Int64, 180 / CELLSIZE), totaldays)
    overallcellcounts::Array{Int64, 3} = zeros(convert(Int64, 360 / CELLSIZE), convert(Int64, 180 / CELLSIZE), totaldays)

    # Open input file
    inchan::IOStream = open(INFILE)

    currentdataset::String = ""
    datasetcelltotals::Array{Float64, 3} = zeros(convert(Int64, 360 / CELLSIZE), convert(Int64, 180 / CELLSIZE), totaldays)
    datasetcellcounts::Array{Int64, 3} = zeros(convert(Int64, 360 / CELLSIZE), convert(Int64, 180 / CELLSIZE), totaldays)

    currentline::String = readline(inchan)
    linecount::Int64 = 1
    while length(currentline) > 0

        fields::Array{String, 1} = split(currentline, "\t")

        dataset::String = fields[1]

        if dataset != currentdataset
            if length(currentdataset) > 0
                applydataset(datasetcelltotals, datasetcellcounts, overallcelltotals, overallcellcounts)

                print("\033[1K\r$currentdataset ($linecount)")

                datasetcelltotals .= 0
                datasetcellcounts .= 0
            end

            currentdataset = dataset
        end

        year::Int64 = parse(Int64, fields[5])
        month::Int64 = parse(Int64, fields[6])
        day::Int64 = parse(Int64, fields[7])
        longitude::Float64 = parse(Float64, fields[11])
        latitude::Float64 = parse(Float64, fields[12])
        fco2::Float64 = parse(Float64, fields[24])

        cellindex::Tuple{Int64, Int64} = getcellindex(longitude, latitude)

        dateindex::Int64 = getdateindex(year, month, day)

        if dateindex != -1
            datasetcelltotals[cellindex[1], cellindex[2], dateindex] = datasetcelltotals[cellindex[1], cellindex[2], dateindex] + fco2
            datasetcellcounts[cellindex[1], cellindex[2], dateindex] = datasetcellcounts[cellindex[1], cellindex[2], dateindex] + 1
        end

        currentline = readline(inchan)
        linecount = linecount + 1
    end

    close(inchan)

    # The last dataset
    applydataset(datasetcelltotals, datasetcellcounts, overallcelltotals, overallcellcounts)

    # Overall cell means
    local meanfco2::Array{Float64, 3} = overallcelltotals ./ overallcellcounts
    meanfco2 = replace(meanfco2, NaN=>-1e35)

    # Write NetCDF
    nc = Dataset(OUTFILE, "c")
    defDim(nc, "longitude", trunc(Int, (360 / CELLSIZE)))
    defDim(nc, "latitude", trunc(Int, (180 / CELLSIZE)))
    defDim(nc, "time", totaldays)

    nclon = defVar(nc, "longitude", Float32, ("longitude",))
    nclon.attrib["units"] = "degrees_east"

    nclat = defVar(nc, "latitude", Float32, ("latitude",))
    nclat.attrib["units"] = "degrees_north"

    nctime = defVar(nc, "time", Float32, ("time",))
    nctime.attrib["calendar"] = "noleap"

    ncfco2 = defVar(nc, "fCO2", Float64, ("longitude", "latitude", "time"))
    ncfco2.attrib["_FillValue"] = -1e35

    nclon[:] = collect(range(CELLSIZE / 2, step=CELLSIZE, stop=360))
    nclat[:] = collect(range(-90 + CELLSIZE / 2, step=CELLSIZE, stop=90))
    nctime[:] = collect(range(STARTYEAR, step=(1/365), stop=(ENDYEAR + 1) - (1/365)))
    ncfco2[:,:,:] = meanfco2

    close(nc)
end

@time run()


function parseMask(line)
    result = "XXXXXXXXXXXXXXXXXXXXXXXXXXXX" * String(split(line,"mask = ")[2])
    return result
end

struct Instruction
    pos::Int
    value::Int
end

function parseInstruction(line)
    insideSquareBrackets = r"\[([^()]*)\]"
    pos = match(insideSquareBrackets, line).captures[1]
    value = split(line, " = ")[2]
    result = Instruction(parse(Int, pos), parse(Int, value))
    return result
end

function applyCondition(x::Tuple{Char,Char})
    m = x[1]
    c = x[2]
    if m == 'X'
        return c
    elseif m == '0'
        return '0'
    elseif m == '1'
        return '1'
    else
        throw("wrong character for mask")
    end
end

function applyMask(mask::String, value::Int)
    valueBstr = bitstring(value)
    c = zip(mask, valueBstr)
    resultStr = String(map(x -> applyCondition(x), c))
    result = parse(Int, resultStr, base=2)
    return result
end

function execute(mask, instruction, memory)
    newValue = applyMask(mask, instruction.value)
    memory[instruction.pos] = newValue
end

function runProgram(inputFile)
    memory = zeros(Int, 100000)
    mask = ""
    for line in readlines(inputFile)
        if startswith(line,"mask")
            mask = parseMask(line)
        else
            instruction = parseInstruction(line)
            execute(mask, instruction, memory)
        end
    end
    println(reduce(+,memory))
end

runProgram("data.txt")