using Flux, Gym, Printf, Zygote
using Flux.Tracker: data
using Flux.Optimise: update!
using Statistics: mean
using DataStructures: CircularBuffer
using Distributions: sample

#using CuArrays

#Load game environment

env = make("Pendulum-v0")
reset!(env)

# ----------------------------- Parameters -------------------------------------

STATE_SIZE = length(state(env))
ACTION_SIZE = 1#length(env.actions)
ACTION_BOUND = env._env.action_space.high[1]
MAX_EP = 50_000
MAX_EP_LENGTH = 200

BATCH_SIZE = 64
MEM_SIZE = 100_000
MIN_EXP_SIZE = 50_000

γ = 99f-2     # discount rate

τ = 1f-3 # for running average while updating target networks
η_act = 1f-4   # Learning rate
η_crit = 1f-3
L2_DECAY = 1f-2

# Ornstein-Uhlenbeck Noise params
μ = 0f0
θ = 15f-2
σ = 2f-1

# --------------------------------- Memory ------------------------------------

memory = CircularBuffer{Any}(MEM_SIZE)

function getData(batch_size = BATCH_SIZE)
  # Getting data in shape
  minibatch = sample(memory, batch_size)
  x = hcat(minibatch...)

  s      =   hcat(x[1, :]...) |> gpu
  a      =   hcat(x[2, :]...) |> gpu
  r      =   hcat(x[3, :]...) |> gpu
  s′     =   hcat(x[4, :]...) |> gpu
  s_mask = .!hcat(x[5, :]...) |> gpu

  return s, a, r, s′, s_mask
end

# -------------------------------- Action Noise --------------------------------

struct OUNoise
  μ
  θ
  σ
  X
end

ou = OUNoise(μ, θ, σ, zeros(Float32, ACTION_SIZE) |> gpu)

function sample_noise(ou::OUNoise)
  dx     = ou.θ * (ou.μ .- ou.X)
  dx   .+= ou.σ * randn(Float32, length(ou.X)) |> gpu
  ou.X .+= dx
end

# Noise scale
τ_ = 25
ϵ  = exp(-1f0 / τ_)
noise_scale = 1f0 / ACTION_BOUND

# ----------------------------- Model Architecture -----------------------------

w_init(dims...) = 6f-3rand(Float32, dims...) .- 3f-3

actor = Chain(Dense(STATE_SIZE, 400, relu),
	      	  Dense(400, 300, relu),
              Dense(300, ACTION_SIZE, tanh, initW=w_init),
              x -> x * ACTION_BOUND) |> gpu
actor_target = deepcopy(actor)

# Critic model
struct crit
  state_crit
  act_crit
  sa_crit
end

Flux.@treelike crit

function (c::crit)(state, action)
  s = c.state_crit(state)
  a = c.act_crit(action)
  c.sa_crit(relu.(s .+ a))
end

Base.deepcopy(c::crit) = crit(deepcopy(c.state_crit),
                              deepcopy(c.act_crit),
			      			  deepcopy(c.sa_crit))

critic = crit(Chain(Dense(STATE_SIZE, 400, relu), Dense(400, 300)) |> gpu,
              		Dense(ACTION_SIZE, 300) |> gpu,
	      			Dense(300, 1, initW=w_init) |> gpu)
critic_target = deepcopy(critic)

# ---------------------- Param Update Functions --------------------------------

function update_target!(target, model; τ = 1f0)
  for (p_t, p_m) in zip(params(target), params(model))
    p_t.data .= (1f0 - τ) * p_t.data .+ τ * p_m.data
  end
end

function update_model!(model, opt, loss, inp...)
  grads = gradient(()->loss(inp...), params(model))
  update!(opt, params(model), grads)
end

# ---------------------------------- Training ----------------------------------
## Losses
function L2_loss(model)
  l2_loss = sum(map(p->sum(p.^2), params(model)))
  return L2_DECAY * l2_loss
end

loss_crit(y, s, a) = Flux.mse(critic(s, a), y) #+ L2_loss(critic)

function loss_act(s)
  actions = actor(s)
  crit_out = critic(s, actions)
  return -sum(crit_out)
end

## Optimizers
opt_crit = ADAM(η_crit)
opt_act  = ADAM(η_act)

function replay()
  s, a, r, s′, s_mask = getData()

  a′ = actor_target(s′)
  v′ = critic_target(s′, a′)
  y = data(r .+ γ * v′ .* s_mask)	# set v′ to 0 where s_ is terminal state

  update_model!(critic, opt_crit, loss_crit, y, s, a)
  update_model!(actor, opt_act, loss_act, s)

  # Update Target models
  update_target!(actor_target, actor; τ = τ)
  update_target!(critic_target, critic; τ = τ)
end

# ---------------------------- Helper Functions --------------------------------

# Stores tuple of state, action, reward, next_state, and done
remember(state, action, reward, next_state, done) =
  push!(memory, [data.((state, action, reward, next_state))..., done])

# Choose action according to policy PendulumPolicy
function action(state, train=true)
  state = reshape(state, size(state)..., 1)
  act_pred = actor(state |> gpu)
  if train
    act_pred = act_pred .+ noise_scale * sample_noise(ou)
  end
  clamp.(act_pred, -ACTION_BOUND, ACTION_BOUND) # returns action
end

function episode!(env::EnvWrapper)
  reset!(env)
  for ep=1:MAX_EP_LENGTH
	s = state(env)
    a = action(s, trainable(env))
    s′, r, done, _ = step!(env, a)
    if trainable(env)
      remember(s, a, r, s′, done)
      replay()
    end
  end
  env.total_reward
end

# -------------------------------- Testing -------------------------------------

# Returns average score over 100 episodes

function test(env::EnvWrapper)
  score_mean = 0f0
  testmode!(env)
  for e=1:100
    total_reward = episode!(env)
    score_mean += total_reward / 100
  end
  testmode!(env, false)
  return score_mean
end

# ------------------------------ Training --------------------------------------

# Populate memory with random actions

s = reset!(env)
for e=1:MIN_EXP_SIZE
  global s
  a = 2rand(Float32) * ACTION_BOUND - ACTION_BOUND
  s′, r, done, _ = step!(env, a)
  remember(s, a, r, s′, done)
  s = s′
end

for e=1:MAX_EP
  total_reward = episode!(env)
  total_reward = @sprintf "%9.3f" total_reward
  print("Episode: $e | Score: $total_reward | ")
  score_mean = test(env)
  score_mean = @sprintf "%9.3f" score_mean
  println("Mean score over 100 test episodes: $score_mean")
end


using Base.Broadcast: broadcastable, Broadcasted, materialize
using Swizzles.NamedArrays
using Swizzles.ArrayifiedArrays
using Swizzles

foo(x) = x

@testset "NamedArrays" begin
    A = rand(3,3)
    B = rand(3,3)
    C = rand(3,3)

    bc = Delay().(A .+ B .+ C)
    @test typeof(lift_names(bc)) <: Broadcasted{<:Any, <:Any, typeof(+), <:Tuple{<:Broadcasted{<:Any, <:Any, typeof(+), <:Tuple{<:NamedArray{<:Any, <:Any, <:Any, :obj1}, <:NamedArray{<:Any, <:Any, <:Any, :obj2}}}, NamedArray{<:Any, <:Any, <:Any, :obj3}}}
    @test name(lift_names(bc).args[1].args[1]) == :obj1
    @test name(lift_names(bc).args[1].args[2]) == :obj2
    @test name(lift_names(bc).args[2]) == :obj3
    @test materialize(lift_names(bc)) == materialize(bc)

    bc = Delay().(Swizzle(+).(A .+ B))
    @test typeof(lift_names(bc)) <: Swizzles.SwizzledArray{<:Any, <:Any, typeof(+), (), <:Any, <:ArrayifiedArray{<:Any, <:Any, <:Broadcasted{<:Any, <:Any, typeof(+), <:Tuple{<:NamedArray{<:Any, <:Any, <:Any, :obj1}, <:NamedArray{<:Any, <:Any, <:Any, :obj2}}}}}
end


function add_flows(m::JuMP.AbstractModel, netinjection::BalanceNamedTuple, system_formulation::Type{S}, sys::PowerSystems.PowerSystem) where {S <: PM.AbstractDCPForm}

    fbr = m[:fbr]
    branch_name_index = m[:fbr].axes[1]
    time_index = m[:fbr].axes[2]

    for t in time_index, (ix,branch) in enumerate(branch_name_index)

        !isassigned(netinjection.var_active,sys.branches[ix].connectionpoints.from.number,t) ? netinjection.var_active[sys.branches[ix].connectionpoints.from.number,t] = -fbr[branch,t] : JuMP.add_to_expression!(netinjection.var_active[sys.branches[ix].connectionpoints.from.number,t],-fbr[branch,t])
        !isassigned(netinjection.var_active,sys.branches[ix].connectionpoints.to.number,t) ? netinjection.var_active[sys.branches[ix].connectionpoints.to.number,t] = fbr[branch,t] : JuMP.add_to_expression!(netinjection.var_active[sys.branches[ix].connectionpoints.to.number,t],fbr[branch,t])

    end

end

function add_flows(m::JuMP.AbstractModel, netinjection::BalanceNamedTuple, system_formulation::Type{S}, sys::PowerSystems.PowerSystem) where {S <: PM.AbstractDCPLLForm}

    fbr_fr = m[:fbr_fr]
    fbr_to = m[:fbr_to]
    time_index = m[:fbr_to].axes[2]
    branch_name_index = m[:fbr_fr].axes[1]

    for t in time_index, (ix,branch) in enumerate(branch_name_index)

        !isassigned(netinjection.var_active,sys.branches[ix].connectionpoints.from.number,t) ? netinjection.var_active[sys.branches[ix].connectionpoints.from.number,t] = -fbr_fr[branch,t] : JuMP.add_to_expression!(netinjection.var_active[sys.branches[ix].connectionpoints.from.number,t],-fbr_fr[branch,t])
        !isassigned(netinjection.var_active,sys.branches[ix].connectionpoints.to.number,t) ? netinjection.var_active[sys.branches[ix].connectionpoints.to.number,t] = fbr_to[branch,t] : JuMP.add_to_expression!(netinjection.var_active[sys.branches[ix].connectionpoints.to.number,t],fbr_to[branch,t])

    end

end

function nodalflowbalance(m::JuMP.AbstractModel, netinjection::BalanceNamedTuple, system_formulation::Type{S}, sys::PowerSystems.PowerSystem) where {S <: AbstractFlowForm}

    time_index = 1:sys.time_periods
    bus_name_index = [b.name for b in sys.buses]

    add_flows(m, netinjection, system_formulation, sys)

    pf_balance = JuMP.JuMPArray(Array{ConstraintRef}(undef,length(bus_name_index), sys.time_periods), bus_name_index, time_index)

        for t in time_index, (ix,bus) in enumerate(bus_name_index)

            isassigned(netinjection.var_active,ix, t) ? true : @error "Islanded Bus in the system"

            pf_balance[bus,t] = @constraint(m, netinjection.var_active[ix, t] == netinjection.timeseries_active[ix, t])

        end

        JuMP.register_object(m, :NodalFlowBalance_active, pf_balance)

end

function nodalflowbalance(m::JuMP.AbstractModel, netinjection::BalanceNamedTuple, system_formulation::Type{S}, sys::PowerSystems.PowerSystem) where {S <: AbstractDCPowerModel}

    time_index = 1:sys.time_periods
    bus_name_index = [b.name for b in sys.buses]

    PM_dict = pass_to_pm(sys, netinjection)

        for t in time_index, bus in sys.buses

            !isassigned(netinjection.var_active,bus.number,t) ? PM_dict["nw"]["$(t)"]["bus"]["$(bus.number)"]["pni"] = -(netinjection.timeseries_active[bus.number, t]) : PM_dict["nw"]["$(t)"]["bus"]["$(bus.number)"]["pni"] = JuMP.add_to_expression!(netinjection.var_active[bus.number,t],-(netinjection.timeseries_active[bus.number, t]))

        end

        m.ext[:PM_object] = PM_dict

end

function nodalflowbalance(m::JuMP.AbstractModel, netinjection::BalanceNamedTuple, system_formulation::Type{S}, sys::PowerSystems.PowerSystem) where {S <: AbstractACPowerModel}

    nodalflowbalance(m, netinjection, AbstractDCPowerModel, sys)

    PM_dict = m.ext[:PM_object]

    time_index = 1:sys.time_periods
    bus_name_index = [b.name for b in sys.buses]

        for t in time_index, bus in sys.buses

            !isassigned(netinjection.var_reactive,bus.number,t) ? PM_dict["nw"]["$(t)"]["bus"]["$(bus.number)"]["qni"] = -(netinjection.timeseries_reactive[bus.number, t]) : PM_dict["nw"]["$(t)"]["bus"]["$(bus.number)"]["qni"] = JuMP.add_to_expression!(netinjection.var_reactive[bus.number,t],-(netinjection.timeseries_reactive[bus.number, t]))

        end

        m.ext[:PM_object] = PM_dict

end


export fields

add_verbose_scope(:Fields)
add_assert_scope(:Fields)

add_verbose_scope(:FieldsNonFancy)

mutable struct FieldsTower
  field::AnticNumberField
  generators_of_automorphisms::Vector{NfToNfMor}
  subfields::Vector{NfToNfMor}
  ramified_primes::Vector{fmpz}

  #Solvable embedding problems for the extension
  #They are here to improve the conductor computation
  has_info::Bool
  imgs_autos::Vector{Main.ForeignGAP.MPtr}
  auts_for_conductors::Vector{Main.ForeignGAP.MPtr}
  proj_ext::Main.ForeignGAP.MPtr

  function FieldsTower(K::AnticNumberField, auts::Vector{NfToNfMor}, subfields::Vector{NfToNfMor})
    z = new()
    z.field = K
    z.generators_of_automorphisms = auts
    z.subfields = subfields
    z.has_info = false
    return z
  end

end

include("./brauer.jl")
include("./merge.jl")
include("./abelian_layer.jl")
include("./read_write.jl")
include("./conductors.jl")

Generic.degree(F::FieldsTower) = degree(F.field)
Hecke.maximal_order(F::FieldsTower) = maximal_order(F.field)

function ramified_primes(F::FieldsTower)
  if !isdefined(F, :ramified_primes)
    f = factor(discriminant(maximal_order(F.field)))
    F.ramified_primes = collect(keys(f.fac))
  end
  return F.ramified_primes
end

###############################################################################
#
#  From automorphisms to permutation groups
#
###############################################################################

function Base.push!(G::AbstractAlgebra.Generic.geobucket{T}, p::T) where {T <: AbstractAlgebra.MPolyElem}
   R = parent(p)
   i = max(1, ndigits(length(p), base=4))
   l = length(G.buckets)
   if length(G.buckets) < i
     resize!(G.buckets, i)
     for j in (l + 1):i
       G.buckets[j] = zero(R)
     end
   end
   G.buckets[i] = addeq!(G.buckets[i], p)
   while i <= G.len
      if length(G.buckets[i]) >= 4^i
         G.buckets[i + 1] = addeq!(G.buckets[i + 1], G.buckets[i])
         G.buckets[i] = R()
         i += 1
      end
      break
   end
   if i == G.len + 1
      Base.push!(G.buckets, R())
      G.len += 1
   end
end

function permutation_group(G::Vector{Hecke.NfRelNSToNfRelNSMor{nf_elem}})
  permutations = permutation_group1(G)
  return _perm_to_gap_grp(permutations)
end

function permutations(G::Array{Hecke.NfToNfMor, 1})
  K = domain(G[1])
  n = length(G)
  dK = degree(K)
  d = numerator(discriminant(K.pol))
  p = 11
  while mod(d, p) == 0
    p = next_prime(p)
  end
  R = GF(p, cached = false)
  Rx, x = PolynomialRing(R, "x", cached = false)
  fmod = Rx(K.pol)

  pols = gfp_poly[x]
  gpol = Rx(G[1].prim_img)
  if gpol != x
    push!(pols, gpol)
    gpol = compose_mod(gpol, pols[2], fmod)
    while gpol != x
      push!(pols, gpol)
      gpol = compose_mod(gpol, pols[2], fmod)
    end
  end
  order = length(pols)

  for i in 2:n
    pi = Rx(G[i].prim_img)
    if !(pi in pols)
      previous_order = order
      order = order + 1
      push!(pols, pi)
      for j in 2:previous_order
        order = order + 1
        push!(pols, compose_mod(pols[j], pi, fmod))
      end
      if order == n
        break
      end
      rep_pos = previous_order + 1
      while rep_pos <= order
        for k in 1:i
          po = Rx(G[k].prim_img)
          att = compose_mod(pols[rep_pos], po, fmod)
          if !(att in pols)
            order = order + 1
            push!(pols, att)
            for j in 2:previous_order
              order = order + 1
              push!(pols, compose_mod(pols[j], att, fmod))
            end
            if order == dK
              break
            end
          end
          if order == dK
            break
          end
        end
        if order == dK
          break
        end
        rep_pos = rep_pos + previous_order
      end
    end
  end
  #Now, I have the images mod p
  Dcreation = Vector{Tuple{gfp_poly, Int}}(undef, length(pols))
  for i = 1:length(pols)
    Dcreation[i] = (pols[i], i)
  end
  permutations = Array{Array{Int, 1},1}(undef, n)
  for i = 1:n
    permutations[i] = Vector{Int}(undef, dK)
  end
  gen_pols = [Rx(x.prim_img) for x in G]
  D = Dict{gfp_poly, Int}(Dcreation)
  for s = 1:n
    for i = 1:length(pols)
      permutations[s][i] = D[Hecke.compose_mod(gen_pols[s], pols[i], fmod)]
    end
  end
  return permutations
end

function permutation_group(G::Array{Hecke.NfToNfMor, 1})
  return _perm_to_gap_grp(permutations(G))
end

function _from_autos_to_perm(G::Array{Hecke.NfToNfMor,1})

  K = domain(G[1])
  @assert degree(K) == length(G)
  n = length(G)
  #First, find a good prime
  p = 2
  d = numerator(discriminant(K.pol))
  while mod(d, p) == 0
    p = next_prime(p)
  end
  R = GF(p, cached = false)
  Rx, x = PolynomialRing(R, "x", cached = false)
  fmod = Rx(K.pol)
  pols = Vector{Tuple{gfp_poly, Int}}(undef, n)
  for i = 1:n
    pols[i] = (Rx(G[i].prim_img), i)
  end
  D = Dict{gfp_poly, Int}(pols)
  permutations = Array{Array{Int, 1},1}(undef, n)
  for s = 1:n
    perm = Array{Int, 1}(undef, n)
    for i = 1:n
      perm[i] = D[Hecke.compose_mod(pols[i][1], pols[s][1], fmod)]
    end
    permutations[s] = perm
  end
  return permutations

end

function _perm_to_gap_grp(perm::Array{Array{Int, 1},1})
  g = Main.ForeignGAP.MPtr[]
  for x in perm
    z = _perm_to_gap_perm(x)
    push!(g, z)
  end
  return GAP.Globals.Group(GAP.julia_to_gap(g))
end

function _perm_to_gap_perm(x::Array{Int, 1})
  z = GAP.Globals.PermList(GAP.julia_to_gap(x))
  return z
end

function IdGroup(autos::Array{NfToNfMor, 1})
  G = permutation_group(autos)
  return GAP.Globals.IdGroup(G)
end

################################################################################
#
#  final computation of the maximal order and automorphisms
#
################################################################################

function _from_relative_to_abs_with_embedding(L::Hecke.NfRelNS{T}, autL::Array{Hecke.NfRelNSToNfRelNSMor{T}, 1}, use_simplify::Bool = true) where T

  S, mS = simple_extension(L)
  K, mK, MK = absolute_field(S, false)

  #First, we compute the maximal order of the absolute field.
  #We start from the maximal orders of the relative extension and of the base field.
  #FALSE: Since the computation of the relative maximal order is slow, I prefer to bring to the absolute field the elements
  # generating the equation order.
  pols = L.pol
  gL = gens(L)
  B = Array{nf_elem, 1}(undef, degree(K))
  B[1] = K(1)
  ind = total_degree(pols[1])
  genjj = mK\(mS\gL[1])
  for i = 2:ind
    B[i] = genjj*B[i-1]
  end
  for jj = 2:length(pols)
    genjj = mK\(mS\gL[jj])
    el = deepcopy(genjj)
    new_deg = total_degree(pols[jj])
    for i = 2:new_deg
      for j = 1:ind
        B[(i-1)* ind + j] = B[j]* el
      end
      mul!(el, el, genjj)
    end
    ind *= new_deg
  end

  #Now, I add the elements of the maximal order
  OB = maximal_order(base_field(S))
  for i = 1:degree(OB)
    el = MK(OB.basis_nf[i])
    for j = 1:ind
      B[(i-1)* ind + j] = B[j] * el
    end
  end
  @vprint :Fields 2 "Product basis computed\n"
  #Now, we compute the maximal order. Then we simplify.
  #We simplify only if the degree of the field is lower than 30

  BasisMat = basis_matrix(B, FakeFmpqMat)
  @vtime :Fields 3 Hecke.hnf_modular_eldiv!(BasisMat.num, BasisMat.den, :lowerleft)
  NewBMat = FakeFmpqMat(BasisMat.num, BasisMat.den)
  @vtime :Fields 3 Ostart = NfAbsOrd(K, NewBMat)
  Ostart.index = divexact(NewBMat.den^degree(K), prod_diagonal(NewBMat.num))
  Ostart.gen_index = fmpq(Ostart.index)
  Ostart.disc = divexact(numerator(discriminant(K)), Ostart.index^2)
  ram_primes_rel = numerator(norm(discriminant(L)))
  ram_primes_down = Hecke.ramified_primes(OB)
  for p in ram_primes_down
    if isone(gcd(p, ram_primes_rel))
      push!(Ostart.primesofmaximality, p)
    end
  end
  @vtime :Fields 3 O1 = MaximalOrder(Ostart)
  O1.ismaximal = 1
  Hecke._set_maximal_order_of_nf(K, O1)
  if use_simplify
    @vtime :Fields 3 Ks, mKs = Hecke.simplify(K)
    #Now, we have to construct the maximal order of this field.
    #I compute the inverse of mKs
    @vtime :Fields 3 mKsI = find_inverse(mKs)
    if isdefined(O1, :lllO)
      lO = O1.lllO::NfOrd
      O2 = NfOrd(nf_elem[mKsI(x) for x in basis(lO, K, copy = false)], false)
      #O2.lllO = O2
    else
      O2 = NfOrd(nf_elem[mKsI(x) for x in basis(O1, K, copy = false)], false)
    end
    O2.ismaximal = 1
    @assert isdefined(O1, :disc)
    O2.disc = O1.disc
    O2.index = root(divexact(numerator(discriminant(Ks)), O2.disc), 2)
    @vtime :Fields 3 OLLL = lll(O2)
    Hecke._set_maximal_order_of_nf(Ks, OLLL)
  else
    Ks = K
    mKs = id_hom(K)
    mKsI = id_hom(K)
  end
  #I want also the embedding of the old field into the new one.
  #It is enough to find the image of the primitive element.
  k = base_field(S)
  a = MK(gen(k))
  embed = NfToNfMor(k, Ks, mKsI(a))
  #@assert iszero(k.pol(img_a))
  @vprint :Fields 3 "MaximalOrder Computed. Now Automorphisms\n"
  #Now, the automorphisms.
  # I need both generators and the whole group.
  autos = Array{NfToNfMor, 1}(undef, length(autL))
  el = mS(mK(mKs.prim_img))
  el1 = mS(mK(gen(K)))
  for i=1:length(autL)
    #@assert iszero(K.pol(mK(mS\(autL[i](el1)))))
    x = mKsI(mK\(mS\(autL[i](el))))
    #@assert Ks.pol(y) == 0
    autos[i] = Hecke.NfToNfMor(Ks, Ks, x)
  end
  #And the generators are done. Now the closure
  #@vtime :Fields 3 Hecke._set_automorphisms_nf(Ks, closure(autos, degree(Ks)))
  #No! I set the automorphisms only if I need them
  #Hecke._set_automorphisms_nf(Ks, autsKs)
  @vprint :Fields 2 "Finished\n"
    #@assert codomain(embed) == Ks
  return Ks, autos, embed
end

function find_inverse(f::NfToNfMor)
  v = Vector{nf_elem}(undef, degree(domain(f)))
  v[1] = one(codomain(f))
  v[2] = f.prim_img
  for i = 3:degree(domain(f))
    v[i] = v[2]*v[i-1]
  end
  M = basis_matrix(v)
  w = zero_matrix(FlintQQ, 1, nrows(M))
  w[1, 2] = 1
  fl, s = can_solve(M, w, side = :left)
  @assert fl
  s1 = FakeFmpqMat(s)
  pre_img = elem_from_mat_row(domain(f), s1.num, 1, s1.den)
  return hom(codomain(f), domain(f), pre_img)
end

function find_inverse2(f::NfToNfMor)
  K = codomain(f)
  arr_prim_img = Array{nf_elem, 1}(undef, degree(K))
  arr_prim_img[1] = K(1)
  arr_prim_img[2] = f.prim_img
  for i = 3:degree(K)
    arr_prim_img[i] = arr_prim_img[i-1]*f.prim_img
  end
  M1 = inv(basis_matrix(arr_prim_img, FakeFmpqMat))
  return M1
end

###############################################################################
#
#  Split Extensions
#
###############################################################################

function _split_extension(G::Array{Hecke.NfToNfMor, 1}, mats::Array{Hecke.GrpAbFinGenMap, 1})

  gtype = map(Int, domain(mats[1]).snf)
  G1 = permutation_group(G)
  gensG1 = GAP.Globals.GeneratorsOfGroup(G1)
  A = GAP.Globals.AbelianGroup(GAP.julia_to_gap(gtype))
  gens = GAP.Globals.GeneratorsOfGroup(A)
  auts = Array{Main.ForeignGAP.MPtr, 1}(undef, length(mats))
  for i = 1:length(mats)
    images = Array{Main.ForeignGAP.MPtr, 1}(undef, length(gtype))
    for j = 1:length(gtype)
      g = GAP.Globals.Identity(A)
      for k = 1:length(gtype)
        if !iszero(mats[i].map[j,k])
          g *= gens[k]^Int(mats[i].map[j,k])
        end
      end
      images[j] = g
    end
    auts[i] = GAP.Globals.GroupHomomorphismByImages(A, A, gens, GAP.julia_to_gap(images))
  end
  AutGrp = GAP.Globals.Group(GAP.julia_to_gap(auts))
  mp = GAP.Globals.GroupHomomorphismByImages(G1, AutGrp, gensG1, GAP.julia_to_gap(auts))
  return GAP.Globals.SplitExtension(G1, mp, A)

end

###############################################################################
#
#  Check group extension
#
###############################################################################

function check_group_extension(TargetGroup::Main.ForeignGAP.MPtr, autos::Array{NfToNfMor, 1}, res_act::Array{GrpAbFinGenMap, 1})

  GS = domain(res_act[1])
  expo = Int(GS.snf[end])
  K = domain(autos[1])
  d = degree(K)
  com, uncom = ppio(expo, d)

  if com == 1
    # I only need to check the split extension, since the second cohomology group is
    # trivial, regardless of the action
    H = _split_extension(autos, res_act)
    return GAP.Globals.IdGroup(H) == TargetGroup
  end

  if uncom == 1
    #Need a cohomological check. Only useful in the prime power case.
    return true
  end

  # I check the split extension related to only uncom
  #Now, I have to check if the split extension is isomorphic to IdH
  Qn, mQn = quo(GS, uncom, false)
  S1, mS1 = snf(Qn)
  new_res_act = Array{GrpAbFinGenMap, 1}(undef, length(res_act))
  for i = 1:length(res_act)
    Mat = mS1.map*mQn.imap*res_act[i].map*mQn.map*mS1.imap
    Hecke.reduce_mod_snf!(Mat, S1.snf)
    new_res_act[i] = hom(S1, S1, Mat)
  end
  H = _split_extension(autos, new_res_act)
  return GAP.Globals.IdGroup(H) == TargetGroup

end

###############################################################################
#
#  Interface to find Fields
#
###############################################################################

function field_extensions(list::Vector{FieldsTower}, bound::fmpz, IsoE1::Main.ForeignGAP.MPtr, l::Array{Int, 1}, only_real::Bool, simplify::Bool)

  grp_to_be_checked = Dict{Int, Main.ForeignGAP.MPtr}()
  d = degree(list[1])
  n = lcm(l)
  com, uncom = ppio(n, d)
  fac = factor(n)
  for (p, v) in fac
    grp_to_be_checked[Int(p)] = _construct_grp(IsoE1, Int(p)^v)
  end
  if uncom != 1
    IsoCheck = _construct_grp(IsoE1, uncom)
  else
    IsoCheck = IsoE1
  end
  final_list = FieldsTower[]
  for (j, x) in enumerate(list)
    @vprint :Fields 1 "Field $(j)/$(length(list)): $(x.field.pol)"
    @vprint :FieldsNonFancy 1 "Field $(j)/$(length(list)): $(x.field.pol)\n"
    append!(final_list, field_extensions(x, bound, IsoCheck, l, only_real, grp_to_be_checked, IsoE1, simplify))
  end
  return final_list

end

function field_extensions(x::FieldsTower, bound::fmpz, IsoE1::Main.ForeignGAP.MPtr, l::Array{Int, 1}, only_real::Bool, grp_to_be_checked::Dict{Int, Main.ForeignGAP.MPtr}, IsoG::Main.ForeignGAP.MPtr, simplify::Bool)

  list_cfields = _abelian_normal_extensions(x, l, bound, IsoE1, only_real, IsoG)
  if isempty(list_cfields)
    @vprint :Fields 1 "\e[1F$(Hecke.set_cursor_col())$(Hecke.clear_to_eol())Number of new fields found: 0\n\n"
    @vprint :FieldsNonFancy 1 "Number of new fields found: 0\n"
    return Vector{FieldsTower}()
  end
  list = from_class_fields_to_fields(list_cfields, x.generators_of_automorphisms, grp_to_be_checked)
  @vprint :Fields 1 "Computing maximal orders"
  @vprint :FieldsNonFancy 1 "Computing maximal orders\n"
  final_list = Vector{FieldsTower}(undef, length(list))
  for j = 1:length(list)
    fld, autos, embed = _from_relative_to_abs_with_embedding(list[j][1], list[j][2], simplify)
    previous_fields = Array{NfToNfMor, 1}(undef, length(x.subfields)+1)
    for s = 1:length(x.subfields)
      previous_fields[s] = x.subfields[s]
    end
    previous_fields[end] = embed
    final_list[j] = FieldsTower(fld, autos, previous_fields)
  end
  @vprint :Fields 1 "$(Hecke.set_cursor_col())$(Hecke.clear_to_eol())"
  @vprint :Fields 1 "Number of new fields found: $(length(final_list))\n\n"
  @vprint :FieldsNonFancy 1 "Number of new fields found: $(length(final_list))\n"
  return final_list

end

###############################################################################
#
#  Interface
#
###############################################################################

function fields_direct_product(g1, g2, red, redfirst, absolute_bound; only_real = false)
  b1 = root(absolute_bound, g2[1])
  b2 = root(absolute_bound, g1[1])
  l2 = fields(g2[1], g2[2], b2, only_real = only_real)
  if isempty(l2)
    return FieldsTower[]
  end
  if g1 == g2
    return _merge(l2, l2, absolute_bound, red, redfirst)
  end
  l1 = fields(g1[1], g1[2], b1, only_real = only_real)
  if isempty(l1)
    return FieldsTower[]
  end
  return _merge(l1, l2, absolute_bound, red, redfirst)
end

function fields(a::Int, b::Int, absolute_bound::fmpz; using_direct_product::Bool = true, only_real::Bool = false, simplify::Bool = true)
  @assert absolute_bound > 0
  if a == 1
    @assert b == 1
    Qx, x = PolynomialRing(FlintQQ, "x", cached = false)
    K, a = NumberField(x-1, cached = false)
    g = NfToNfMor(K, K, K(1))
    return FieldsTower[FieldsTower(K, NfToNfMor[g], Array{NfToNfMor, 1}())]
  end
  G = GAP.Globals.SmallGroup(a, b)
  @assert GAP.Globals.IsSolvable(G)
  if using_direct_product
    g1, g2, red, redfirst = direct_product_decomposition(G, (a, b))
    if g2 != (1, 1)
      return fields_direct_product(g1, g2, red, redfirst, absolute_bound; only_real = false)
    end
  end
  @vprint :Fields 1 "Doing Group ($a, $b) with bound $absolute_bound\n"
  @vprint :FieldsNonFancy 1 "Doing Group ($a, $b) with bound $absolute_bound\n"
  L = GAP.Globals.DerivedSeries(G)
  lvl = _real_level(L)
  #First step by hand
  l = GAP.gap_to_julia(Vector{Int64}, GAP.Globals.AbelianInvariants(G))
  @vprint :Fields 1 "contructing abelian extensions with invariants $l \n"
  @vprint :FieldsNonFancy 1 "contructing abelian extensions with invariants $l \n"
  #@vtime :Fields 2
  list = abelian_extensionsQQ(l, root(absolute_bound, div(a,prod(l))), (lvl > 1 || only_real))
  if isempty(list)
    return list
  end
  @vprint :Fields 1 "First step completed\n \n"
  @vprint :FieldsNonFancy 1 "First step completed\n \n"
  for i = 2:length(L)-1

    E1 = GAP.Globals.FactorGroup(L[1], L[i+1])
    l = GAP.gap_to_julia(Vector{Int64}, GAP.Globals.AbelianInvariants(L[i]))
    @vprint :Fields 1 "contructing abelian extensions with invariants $l \n"
    @vprint :FieldsNonFancy 1 "contructing abelian extensions with invariants $l \n"
    o = divexact(a, GAP.Globals.Size(E1))
    bound = root(absolute_bound, o)
    IsoE1 = GAP.Globals.IdGroup(E1)
    @vprint :Fields 1 "Number of fields at the $i -th step: $(length(list)) \n"
    @vprint :FieldsNonFancy 1 "Number of fields at the $i -th step: $(length(list)) \n"
    lG = snf(DiagonalGroup(l))[1]
    invariants = map(Int, lG.snf)
    onlyreal = (lvl > i || only_real)
    @vprint :Fields 1 "Computing obstructions\n"
    @vprint :FieldsNonFancy 1 "Computing obstructions\n"
    #@vtime :Fields 1
    list = check_Brauer_obstruction(list, L, i, invariants)
    @vprint :Fields 1 "Fields to check: $(length(list))\n\n"
    @vprint :FieldsNonFancy 1 "Fields to check: $(length(list))\n\n"
    if isempty(list)
      return FieldsTower[]
    end
    simplify_i_th = true
    if i == length(L)-1 && simplify == false
      simplify_i_th = false
    end
    list = field_extensions(list, bound, IsoE1, invariants, onlyreal, simplify)
    @vprint :Fields 1 "Step $i completed\n"
    @vprint :FieldsNonFancy 1 "Step $i completed\n"
    if isempty(list)
      return FieldsTower[]
    end
  end
  return list

end

function fields(a::Int, b::Int, list::Vector{FieldsTower}, absolute_bound::fmpz; only_real::Bool = false, simplify::Bool = true)
  G = GAP.Globals.SmallGroup(a, b)
  return fields(list, G, absolute_bound, only_real = only_real, simplify = simplify)
end

function fields(list::Vector{FieldsTower}, G, absolute_bound::fmpz; only_real::Bool = false, simplify::Bool = true)
  L = GAP.Globals.DerivedSeries(G)
  lvl = _real_level(L)
  first = true
  for i = 2:length(L)-1
    G1 = GAP.Globals.FactorGroup(L[1], L[i])
    if first && GAP.Globals.Size(G1) != degree(list[1].field)
      continue
    end
    first = false
    E1 = GAP.Globals.FactorGroup(L[1], L[i+1])
    H1 = GAP.Globals.FactorGroup(L[i], L[i+1])
    l = GAP.gap_to_julia(Vector{Int64}, GAP.Globals.AbelianInvariants(H1))
    @vprint :Fields 1 "contructing abelian extensions with invariants $l \n"
    @vprint :FieldsNonFancy 1 "contructing abelian extensions with invariants $l \n"
    o = divexact(GAP.Globals.Size(G), GAP.Globals.Size(E1))
    bound = root(absolute_bound, o)
    IsoE1 = GAP.Globals.IdGroup(E1)
    @vprint :Fields 1 "Number of fields at the $i -th step: $(length(list)) \n"
    @vprint :FieldsNonFancy 1 "Number of fields at the $i -th step: $(length(list)) \n"
    lG = snf(DiagonalGroup(l))[1]
    invariants = map(Int, lG.snf)
    onlyreal = (lvl > i || only_real)
    #First, I search for obstruction.
    if iscyclic(lG)
      @vprint :Fields 1 "Computing obstructions\n"
      @vprint :FieldsNonFancy 1 "Computing obstructions\n"
      #@vtime :Fields 1
      list = check_Brauer_obstruction(list, L, i, invariants[1])
      @vprint :Fields 1 "Fields to check: $(length(list))\n\n"
      @vprint :FieldsNonFancy 1 "Fields to check: $(length(list))\n\n"
    end
    if isempty(list)
      return FieldsTower[]
    end
    list = field_extensions(list, bound, IsoE1, invariants, onlyreal, simplify)
    @vprint :Fields 1 "Step $i completed\n"
    @vprint :FieldsNonFancy 1 "Step $i completed\n"
    if isempty(list)
      return FieldsTower[]
    end
  end
  return list
end


 # Load required packages.
using CSVFiles, DataFrames, Mimi, MimiRICE2010, MimiNICE, Statistics

# Load helper functions and revenue recycling components being added to the NICE model.
include("helper_functions.jl")
include(joinpath("revenue_recycling_components", "nice_revenue_recycle_component_time_varying.jl"))

#-----------------------#
# ----- Load Data ----- #
#-----------------------#

# Load studies used to calculate elasticities.
elasticity_studies = DataFrame(load(joinpath(@__DIR__, "..", "data", "elasticity_study_data.csv")))

# Load updated UN population projections and convert units to millions of people.
un_population_data = Matrix(DataFrame(load(joinpath(@__DIR__, "..", "data", "UN_medium_population_scenario.csv"), skiplines_begin=3))[:, 3:end]) ./ 1000

# Load quintile income distribution data that remains fixed over time (these shares are an update from the original NICE income distributions).
consumption_distribution_raw = DataFrame(load(joinpath(@__DIR__, "..", "data", "constant_quintile_distributions_consumption.csv")))

# Clean up and organize time-varying income distribution data into required NICE format (time × regions × quintiles).
consumption_distributions = get_quintile_income_shares(consumption_distribution_raw)

# -------------------------------------------------
# -------------------------------------------------
# Create function to add revenue recyling to NICE.
# -------------------------------------------------
# -------------------------------------------------

function create_nice_recycle(;slope_type::Symbol=:central, percentile::Float64=0.90)

    # Initialize an instance of NICE to build in revenue recycling.
    nice_rr = MimiNICE.create_nice()

    # Get number of timesteps across model time horizon.
    n_steps = length(dim_keys(nice_rr, :time))

    # Add in NICE revenue recycling component.
    add_comp!(nice_rr, nice_recycle, after = :nice_neteconomy)

    #-------------------------------------#
    # ----- Assign Model Parameters ----- #
    #-------------------------------------#

    # First, perform a meta-regression based on study results to calculate elasticity vs. ln gdp per capita relationship.
    meta_intercept, meta_slope = meta_regression(elasticity_studies, slope_type=slope_type, percentile=percentile)

    # Update parameters common to multiple components or that already have external Mimi parameters defined.
    update_param!(nice_rr, :l, un_population_data)
    update_param!(nice_rr, :quintile_pop, un_population_data ./ 5)

    # Set values for additional NICE and recycling parameters.
    set_param!(nice_rr, :nice_recycle, :min_study_gdp, minimum(elasticity_studies.pcGDP))
    set_param!(nice_rr, :nice_recycle, :max_study_gdp, maximum(elasticity_studies.pcGDP))
    set_param!(nice_rr, :nice_recycle, :elasticity_intercept, meta_intercept)
    set_param!(nice_rr, :nice_recycle, :elasticity_slope, meta_slope)
    set_param!(nice_rr, :nice_recycle, :regional_population, un_population_data)
    set_param!(nice_rr, :nice_recycle, :quintile_income_shares, consumption_distributions)
    set_param!(nice_rr, :nice_recycle, :damage_elasticity, 1.0)
    set_param!(nice_rr, :nice_recycle, :recycle_share, ones(12,5).*0.2)
    set_param!(nice_rr, :nice_recycle, :global_carbon_tax, zeros(n_steps))
    set_param!(nice_rr, :nice_recycle, :lost_revenue_share, 0.0)
    set_param!(nice_rr, :nice_recycle, :global_recycle_share, zeros(12))

    # Create parameter connections (:component => :parameter).
    connect_param!(nice_rr, :nice_recycle => :industrial_emissions, :emissions       => :EIND)
    connect_param!(nice_rr, :nice_recycle => :DAMFRAC,              :damages         => :DAMFRAC)
    connect_param!(nice_rr, :nice_recycle => :ABATEFRAC,            :nice_neteconomy => :ABATEFRAC)
    connect_param!(nice_rr, :nice_recycle => :Y,                    :nice_neteconomy => :Y)
    connect_param!(nice_rr, :nice_recycle => :CPC,                  :nice_neteconomy => :CPC)
    connect_param!(nice_rr, :nice_welfare => :quintile_c,           :nice_recycle    => :qc_post_recycle)

    # Return NICE model with revenue recycling.
    return nice_rr
end


@testset "Block" begin
    @testset "Parse" begin
        block_raw = hex2bytes("020000208ec39428b17323fa0ddec8e887b4a7c53b8c0a0a220cfd0000000000000000005b0750fce0a889502d40508d39576821155e9c9e3f5c3157f961db38fd8b25be1e77a759e93c0118a4ffd71d")
        stream = IOBuffer(block_raw)
        block = Bitcoin.io2blockheader(stream)
        @test block.version == 0x20000002
        want = hex2bytes("8ec39428b17323fa0ddec8e887b4a7c53b8c0a0a220cfd000000000000000000")
        @test block.prev_block == want
        want = hex2bytes("5b0750fce0a889502d40508d39576821155e9c9e3f5c3157f961db38fd8b25be")
        @test block.merkle_root == want
        @test block.timestamp == 0x59a7771e
        @test block.bits == 0x18013ce9
        @test block.nonce == 0x1dd7ffa4
    end
    @testset "Serialize" begin
        block_raw = hex2bytes("020000208ec39428b17323fa0ddec8e887b4a7c53b8c0a0a220cfd0000000000000000005b0750fce0a889502d40508d39576821155e9c9e3f5c3157f961db38fd8b25be1e77a759e93c0118a4ffd71d")
        stream = IOBuffer(block_raw)
        block = Bitcoin.io2blockheader(stream)
        @test Bitcoin.serialize(block) == block_raw
    end
    @testset "Hash" begin
        block_raw = hex2bytes("020000208ec39428b17323fa0ddec8e887b4a7c53b8c0a0a220cfd0000000000000000005b0750fce0a889502d40508d39576821155e9c9e3f5c3157f961db38fd8b25be1e77a759e93c0118a4ffd71d")
        stream = IOBuffer(block_raw)
        block = Bitcoin.io2blockheader(stream)
        Bitcoin.hash(block) == hex2bytes("0000000000000000007e9e4c586439b0cdbe13b1370bdd9435d76a644d047523")
    end
    @testset "BIP9" begin
        block_raw = hex2bytes("020000208ec39428b17323fa0ddec8e887b4a7c53b8c0a0a220cfd0000000000000000005b0750fce0a889502d40508d39576821155e9c9e3f5c3157f961db38fd8b25be1e77a759e93c0118a4ffd71d")
        stream = IOBuffer(block_raw)
        block = Bitcoin.io2blockheader(stream)
        @test Bitcoin.bip9(block)
        block_raw = hex2bytes("0400000039fa821848781f027a2e6dfabbf6bda920d9ae61b63400030000000000000000ecae536a304042e3154be0e3e9a8220e5568c3433a9ab49ac4cbb74f8df8e8b0cc2acf569fb9061806652c27")
        stream = IOBuffer(block_raw)
        block = Bitcoin.io2blockheader(stream)
        @test !Bitcoin.bip9(block)
    end
    @testset "BIP91" begin
        block_raw = hex2bytes("1200002028856ec5bca29cf76980d368b0a163a0bb81fc192951270100000000000000003288f32a2831833c31a25401c52093eb545d28157e200a64b21b3ae8f21c507401877b5935470118144dbfd1")
        stream = IOBuffer(block_raw)
        block = Bitcoin.io2blockheader(stream)
        @test Bitcoin.bip91(block)
        block_raw = hex2bytes("020000208ec39428b17323fa0ddec8e887b4a7c53b8c0a0a220cfd0000000000000000005b0750fce0a889502d40508d39576821155e9c9e3f5c3157f961db38fd8b25be1e77a759e93c0118a4ffd71d")
        stream = IOBuffer(block_raw)
        block = Bitcoin.io2blockheader(stream)
        @test !Bitcoin.bip91(block)
    end
    @testset "BIP141" begin
        block_raw = hex2bytes("020000208ec39428b17323fa0ddec8e887b4a7c53b8c0a0a220cfd0000000000000000005b0750fce0a889502d40508d39576821155e9c9e3f5c3157f961db38fd8b25be1e77a759e93c0118a4ffd71d")
        stream = IOBuffer(block_raw)
        block = Bitcoin.io2blockheader(stream)
        @test Bitcoin.bip141(block)
        block_raw = hex2bytes("0000002066f09203c1cf5ef1531f24ed21b1915ae9abeb691f0d2e0100000000000000003de0976428ce56125351bae62c5b8b8c79d8297c702ea05d60feabb4ed188b59c36fa759e93c0118b74b2618")
        stream = IOBuffer(block_raw)
        block = Bitcoin.io2blockheader(stream)
        @test !Bitcoin.bip141(block)
    end
    @testset "Target" begin
        block_raw = hex2bytes("020000208ec39428b17323fa0ddec8e887b4a7c53b8c0a0a220cfd0000000000000000005b0750fce0a889502d40508d39576821155e9c9e3f5c3157f961db38fd8b25be1e77a759e93c0118a4ffd71d")
        stream = IOBuffer(block_raw)
        block = Bitcoin.io2blockheader(stream)
        @test target(block) == parse(BigInt, "13ce9000000000000000000000000000000000000000000", base=16)
        @test difficulty(block) == 888171856257
    end
    @testset "Check POW" begin
        block_raw = hex2bytes("04000000fbedbbf0cfdaf278c094f187f2eb987c86a199da22bbb20400000000000000007b7697b29129648fa08b4bcd13c9d5e60abb973a1efac9c8d573c71c807c56c3d6213557faa80518c3737ec1")
        stream = IOBuffer(block_raw)
        block = Bitcoin.io2blockheader(stream)
        @test check_pow(block) == true
        block_raw = hex2bytes("04000000fbedbbf0cfdaf278c094f187f2eb987c86a199da22bbb20400000000000000007b7697b29129648fa08b4bcd13c9d5e60abb973a1efac9c8d573c71c807c56c3d6213557faa80518c3737ec0")
        stream = IOBuffer(block_raw)
        block = Bitcoin.io2blockheader(stream)
        @test !check_pow(block)
    end
    @testset "Validate merkle root" begin
        hashes_hex = [
           "f54cb69e5dc1bd38ee6901e4ec2007a5030e14bdd60afb4d2f3428c88eea17c1",
           "c57c2d678da0a7ee8cfa058f1cf49bfcb00ae21eda966640e312b464414731c1",
           "b027077c94668a84a5d0e72ac0020bae3838cb7f9ee3fa4e81d1eecf6eda91f3",
           "8131a1b8ec3a815b4800b43dff6c6963c75193c4190ec946b93245a9928a233d",
           "ae7d63ffcb3ae2bc0681eca0df10dda3ca36dedb9dbf49e33c5fbe33262f0910",
           "61a14b1bbdcdda8a22e61036839e8b110913832efd4b086948a6a64fd5b3377d",
           "fc7051c8b536ac87344c5497595d5d2ffdaba471c73fae15fe9228547ea71881",
           "77386a46e26f69b3cd435aa4faac932027f58d0b7252e62fb6c9c2489887f6df",
           "59cbc055ccd26a2c4c4df2770382c7fea135c56d9e75d3f758ac465f74c025b8",
           "7c2bf5687f19785a61be9f46e031ba041c7f93e2b7e9212799d84ba052395195",
           "08598eebd94c18b0d59ac921e9ba99e2b8ab7d9fccde7d44f2bd4d5e2e726d2e",
           "f0bb99ef46b029dd6f714e4b12a7d796258c48fee57324ebdc0bbc4700753ab1"]
       hashes = [hex2bytes(x) for x in hashes_hex]
       stream = IOBuffer(hex2bytes("00000020fcb19f7895db08cadc9573e7915e3919fb76d59868a51d995201000000000000acbcab8bcc1af95d8d563b77d24c3d19b18f1486383d75a5085c4e86c86beed691cfa85916ca061a00000000"))
       block = Bitcoin.io2blockheader(stream)
       @test Bitcoin.validate_merkle_root(block, hashes)
    end
end


struct PCEulerConstantCache <: StochasticDiffEqConstantCache end

alg_cache(alg::PCEuler,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,f,t,::Type{Val{false}}) = PCEulerConstantCache()

@cache struct PCEulerCache{uType,rateType,rateNoiseType,rateNoiseCollectionType} <: StochasticDiffEqMutableCache
  utmp::uType
  ftmp::rateType
  gtmp::rateNoiseType
  gdWtmp::rateNoiseCollectionType
  bbprimetmp::rateType
end

function alg_cache(alg::PCEuler,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,f,t,::Type{Val{true}})
  utmp = zero(u); ftmp = zero(rate_prototype);
  gtmp = zero(noise_rate_prototype)
  bbprimetmp = zero(ftmp)
  if is_diagonal_noise(prob)
    gdWtmp = gtmp
  else
    gdWtmp = zero(rate_prototype)
  end
  PCEulerCache(utmp,ftmp,gtmp,gdWtmp,bbprimetmp)
end


using PyPlot, PyCall
include("make_SQ_model.jl")
include("utilities.jl")

##### display plots to gui, or just save
## if they plot to gui, they don't have the size specified in teh code
pygui(true)

experiment_dict = Dict("n_vary"=>Dict(:inpts=>[:exit_distance],:epocs=>1000,:ex_locs=>[1.,5.],:cmp=>["ok","bad"],:legend_loc=>"upper right"),
                       "sense_vary"=>Dict(:inpts=>[:sensor_rwd],:epocs=>1000,:ex_locs=>[-35.,-150.],:cmp=>["ok","bad"],:legend_loc=>"lower right"),
                       "transition_vary"=>Dict(:inpts=>[:tprob],:epocs=>1000,:ex_locs=>[0.25,0.75],:cmp=>["ok","bad"],:legend_loc=>"lower right"),
                       "transition_e_vary"=>Dict(:inpts=>[:tprob,:e_mcts],:epocs=>1000,:ex_locs=>[],:cmp=>["ok","bad"],:legend_loc=>"lower right")
                      )
#  experiment_dict = Dict("transition_vary"=>Dict(:inpts=>[:tprob],:epocs=>1000,:ex_locs=>[0.25,0.75],:cmp=>["bad","ok"],:legend_loc=>"lower right"))
#  experiment_dict = Dict("n_vary"=>Dict(:inpts=>[:exit_distance],:epocs=>1000,:ex_locs=>[1.,5.],:cmp=>["bad"],:legend_loc=>"upper right"))
experiment_dict = Dict("transition_e_vary"=>Dict(:inpts=>[:tprob,:e_mcts],:epocs=>1000,:ex_locs=>[1.,5.],:cmp=>["ok","bad"],:legend_loc=>"upper right"))
#  experiment_dict = Dict("sense_vary"=>Dict(:inpts=>[:sensor_rwd],:epocs=>1000,:ex_locs=>[-35.,-150.],:cmp=>["bad"],:legend_loc=>"lower right"))

for expr in keys(experiment_dict)
    net_type = expr
    inpts = experiment_dict[expr][:inpts]
    epocs = experiment_dict[expr][:epocs]
    sq_example_locations = experiment_dict[expr][:ex_locs]
    cmp_list = experiment_dict[expr][:cmp]
    legend_loc = experiment_dict[expr][:legend_loc]
    println("Plotting: $(experiment_dict[expr])")
    for compare in cmp_list
        println("running '$compare' data")
        train_fname = "logs_old/$(net_type)_reference_solver_training.csv"
        test_fname = "logs_old/$(net_type)_$(compare)_solver.csv"

        inputs = Dict()
        for i in inpts
            inputs[i] = "ML.Continuous"
        end
        outputs = Dict(:X3_1=>"ML.Continuous",:X3_2=>"ML.Continuous")

        log_fname = "$(net_type)_$(make_label_from_keys(inputs))"
        log_loc = "nn_logs_old/"
        #  log_path = string(log_loc,"/",log_fname,"_net1-$(epocs)

        #  println("##########")
        #  println("$log_fname, $(readdir(log_loc))")
        #  println("##########")
        if !(any([contains(x,string(log_fname,"_")) for x in readdir(log_loc)]))
            println("No nn file exists, making one now")
            #  make_sq_model(net_type,inpts)
            make_sq_model(net_type,inpts,num_epoc=epocs)
        end

        param_files = searchdir(log_loc,log_fname,".params")

        num_epocs = parse(split(match(r"-\d+",param_files[1]).match,"-")[2])

        SQmodel = load_network(string(log_loc,log_fname),num_epocs,string(log_loc,log_fname,"_SQmodel.jld"))

        # get test and make predictions
        test_input, test_output, test_table, input_sch, output_sch = return_data(test_fname, inputs=inputs, outputs=outputs)

        data_mat = ML.featuremat(merge(input_sch,output_sch),test_table)

        info("restoring limits")
        limits = restore_eng_units(SQmodel.range,SQmodel.output_sch)
        info("getting predictions")
        _notused, pred_outputs = SQ_predict(SQmodel,test_input,use_eng_units=true)

        info("restoring test data")
        tst_in_eng = restore_eng_units(test_input,input_sch)
        tst_out_eng_ary = restore_eng_units(test_output,output_sch)

        # make figures
        if length(inputs) == 1
            fig,ax_ary = PyPlot.subplots(1,1,sharex=false)
            fig[:set_size_inches](8.0,6.0)
            fontsize = 20
            PyPlot.grid()

            i1 = collect(keys(inputs))[1]

            idx1 = nearest_to_x(tst_in_eng[i1],sq_example_locations[1])
            idx2 = nearest_to_x(tst_in_eng[i1],sq_example_locations[2])

            scatter_with_conf_bnds(ax_ary,tst_in_eng,tst_out_eng_ary,i1,:X3_1,:X3_2,:red,subsample=[idx1 idx2],label="candidate($(compare))",bar=true)
            #  scatter_with_conf_bnds(ax_ary,tst_in_eng,tst_out_eng_ary,i1,:X3_1,:X3_2,:red,subsample=collect(1:length(tst_in_eng[i1])),label="candidate",bar=true)
            #  ax_ary[:scatter](tst_in_eng[i1],tst_out_eng_ary[:X3_1],color=:red,alpha=0.2)

            ax_ary[:set_xlabel](string(i1),fontsize=fontsize)
            ax_ary[:set_ylabel](string("Reward"),fontsize=fontsize)
            ax_ary[:axhline](limits[:X3_1][2])
            ax_ary[:axhline](limits[:X3_1][1])

            add_sq_annotation(ax_ary,tst_in_eng,tst_out_eng_ary,pred_outputs,idx1,i1,:X3_1,:X3_2,SQmodel,fontsize=fontsize)
            add_sq_annotation(ax_ary,tst_in_eng,tst_out_eng_ary,pred_outputs,idx2,i1,:X3_1,:X3_2,SQmodel,fontsize=fontsize)

            ax_ary[:text](minimum(tst_in_eng[i1]),limits[:X3_1][2],L"r_H",fontsize=fontsize,va="bottom")
            ax_ary[:text](minimum(tst_in_eng[i1]),limits[:X3_1][1],L"r_L",fontsize=fontsize,va="top")

            #  ax_ary[:scatter](tst_in_eng[i1],pred_outputs[:X3_1],color=:blue)
            scatter_with_conf_bnds(ax_ary,tst_in_eng,pred_outputs,i1,:X3_1,:X3_2,:blue,subsample=collect(1:8:length(tst_in_eng[i1])),label="trusted",bar=false)
            ax_ary[:legend](loc=legend_loc,fontsize=fontsize)

            PyPlot.savefig(string("figs/",log_fname,"_",compare,".pdf"),dpi=300,transparent=true)
        elseif length(inputs) > 1
            # import stuff for plotting
            @pyimport mpl_toolkits.axes_grid1 as mpl_kit
            al = mpl_kit.make_axes_locatable

            function trust_surf(x,y,key)
                _notused, pred_outputs = SQ_predict(SQmodel,[x y]',use_eng_units=true)
                return pred_outputs[key][1]
            end

            i_rng = [-1.65;1.71]
            grid_ary = collect(linspace(i_rng[1],i_rng[2],50))
            ts = [trust_surf(x,y,:X3_1) for x in grid_ary, y in grid_ary]'
            tss = [trust_surf(x,y,:X3_2) for x in grid_ary, y in grid_ary]'

            display(ts)
            display(tss)

            # Solver Quality Surface
            fig = plt[:figure](figsize=(6.0,6.0))
            grid = fig[:subplots](2,2)
            div = []
            push!(div,al(grid[1]))
            push!(div,al(grid[2]))
            cax = []
            push!(cax,div[1][:append_axes]("right",size="5%",pad=0.05))
            push!(cax,div[2][:append_axes]("right",size="5%",pad=0.05))

            #  grid = ImageGrid(fig, 111,          # as in plt.subplot(111)
                                 #  nrows_ncols=(3,2),
                                 #  direction="row",
                                 #  aspect = true,
                                 #  axes_pad=0.65,
                                 #  share_all=true,
                                 #  cbar_location="right",
                                 #  cbar_mode="each",
                                 #  cbar_size="3%",
                                 #  cbar_pad=0.10,
                                 #  )
            fontsize = 10

            # make correlation plots
            i1 = inpts[1]
            i2 = inpts[2]

            #  corrplot(data_mat)
            poi1 = [0.2;200.] # point of interest, where X3 will be calculated
            poi2 = [0.7;800.] # point of interest, where X3 will be calculated
            poi1_norm = [(poi1[1]-mean(input_sch[i1]))/std(input_sch[i1]);(poi1[2]-mean(input_sch[i2]))/std(input_sch[i2])]
            poi2_norm = [(poi2[1]-mean(input_sch[i1]))/std(input_sch[i1]);(poi2[2]-mean(input_sch[i2]))/std(input_sch[i2])]

            x_rng = [i_rng[1]*std(input_sch[i1])+mean(input_sch[i1]);i_rng[2]*std(input_sch[i1])+mean(input_sch[i1])]
            y_rng = [i_rng[1]*std(input_sch[i2])+mean(input_sch[i2]);i_rng[2]*std(input_sch[i2])+mean(input_sch[i2])]
            subsample_num = 3

            aspct = diff(x_rng)[1]./diff(y_rng)[1]
            #  error()

            sq_cmap = :viridis
            im1 = grid[1][:imshow](ts,origin="lower",cmap=sq_cmap,interpolation="bilinear",extent=[x_rng;y_rng],aspect="auto")
            #  im1 = grid[1][:imshow](ts,origin="lower",cmap=sq_cmap,interpolation="bilinear",aspect="auto")
            grid[1][:scatter](poi1[1],poi1[2],color=:black,marker=L"$A$",color=:gray)
            grid[1][:scatter](poi2[1],poi2[2],color=:black,marker=L"$B$",color=:gray)
            grid[1][:set_ylabel]("$(inpts[2])")
            grid[1][:set_title]("Mean Rwd: Trusted Solver")
            fig[:colorbar](im1,cax=cax[1])

            sqs_cmap = :magma
            im2 = grid[2][:imshow](tss,origin="lower",cmap=sqs_cmap,interpolation="bilinear",extent=[x_rng;y_rng],aspect="auto")
            #  im2 = grid[2][:imshow](tss,origin="lower",cmap=sqs_cmap,interpolation="bilinear",aspect="auto")
            grid[2][:scatter](poi1[1],poi1[2],color=:black,s=50,marker=L"$A$",color=:gray)
            grid[2][:scatter](poi2[1],poi2[2],color=:black,s=50,marker=L"$B$",color=:gray)
            grid[2][:set_ylabel]("$(inpts[2])")
            grid[2][:set_xlabel]("$(inpts[1])")
            grid[2][:set_title]("Std: Trusted Solver")
            fig[:colorbar](im2,cax=cax[2])
            #  grid[2][:set_aspect](aspct)

            add_sq_scatter3d_annotation(grid[1],grid[3],test_input,tst_out_eng_ary,i1,i2,poi1_norm,SQmodel,marker=L"$A$",s=50,fontsize=12)
            add_sq_scatter3d_annotation(grid[2],grid[4],test_input,tst_out_eng_ary,i1,i2,poi2_norm,SQmodel,marker=L"$B$",s=50,fontsize=12)

            grid[3][:legend]()

            PyPlot.tight_layout()
            PyPlot.subplots_adjust(hspace=0.4,wspace=0.8)

            PyPlot.savefig(string("figs/",log_fname,"_",compare,".pdf"),dpi=300,transparent=true)
        else
         error("can't support more than 2 inputs yet")
        end
        PyPlot.close_figs()
    end
end


export macau_blocked

function macau_blocked(data::RelationData;
              num_latent::Int = 10,
              burnin          = 500,
              psamples        = 200,
              verbose::Bool   = true)
end


using StaticArrays
const Vec = SVector

function writefbodyline(f, indent, s)
    for i in 1:indent
        write(f, "\t")
    end
    return write(f, "$s\n")
end

function fort2julia(f_type)
    f_type = lowercase(f_type)
    if f_type == "real"
        return Float32
    elseif f_type == "real(kind=dp)"
        return Float64
    elseif f_type == "complex(kind=dp)"
        return Complex{Float64}
    elseif occursin("character", f_type)
        return String
    elseif f_type == "string"
        return String
    elseif f_type == "integer"
        return Int
    elseif f_type == "logical"
        return Bool
    elseif occursin(".D", f_type)
        return replace(f_type, "D" => "e")
    else
        return Nothing
    end
end
indentabs(indent) = prod(["\t" for i in 1:indent])

open(joinpath(@__DIR__, "mpirunflags.jl"), "w") do wf
    write(wf, "_MPI_FLAGS() = ExecFlag[\n")
    # writefbodyline(1, "flags = ")
    open(joinpath(@__DIR__, "..", "assets", "mpirun_man.txt"), "r") do f
        line = readline(f)
        while line != "OPTIONS"
            line = readline(f)
        end
        while line != "Environment Variables"
            line = strip(readline(f))
            if !isempty(line) && line[1] == '-'
                name        = ""     #--
                symbols     = Symbol[] #-
                type        = Nothing
                description = ""
                sline       = strip.(split(line), ',')
                for s in sline
                    if s[2] == '-' #--npernode
                        name = strip(s, '-')
                    elseif occursin('<', s) # <#persocket>
                        type = if occursin("#", s)
                            occursin(',', s) ? Vector{Int} : Int
                        elseif occursin("ppr", s) #ppr:N:<object>
                            Pair{Int,String}
                        else
                            occursin(',', s) ? Vector{String} : String
                        end
                        break
                    else  #-np
                        push!(symbols, Symbol(strip(s, '-')))
                    end
                end
                line = strip(readline(f))
                while !isempty(line)
                    description *= join(split(line), " ") * "\n" * indentabs(2)
                    line = strip(readline(f))
                end
                description = replace(strip(description), "\"" => "'")
                if name != "" && isempty(symbols)
                    symbols = [Symbol(name)]
                end
                if type <: Number
                    val = zero(type)
                elseif type <: Pair
                    val = 0 => ""
                elseif type === Nothing
                    val = nothing
                else
                    val = "\"\""
                end
                for symbol in symbols
                    writefbodyline(wf, 1,
                                   """ExecFlag(Symbol("$symbol"), "$name", "$description", $val, 1),""")
                end
            end
        end
    end
    return writefbodyline(wf, 0, "]")
end

open(joinpath(@__DIR__, "wannier90flags.jl"), "w") do wf
    write(wf, "_WAN_FLAGS() = Dict{Symbol, Type}(\n")
    open(joinpath(@__DIR__, "..", "assets", "calculations", "wannier",
                  "calculation_flags.txt"), "r") do f
        while !eof(f)
            line = readline(f)
            if isempty(line) || line[1] == '!'
                continue
            else
                s_line = split(line, "::")
                flagpart = s_line[end]
                flag = Symbol(strip(strip(split(split(split(s_line[end], "=")[1], "(")[1],
                                                "!")[1], ':')))
                fl_type = fort2julia(strip(split(s_line[1])[1], ','))
                flagstring = occursin("(", flagpart) && occursin(")", flagpart) ?
                             """Symbol("$flag") => Vector{$fl_type},""" :
                             """Symbol("$flag") => $fl_type,"""
                writefbodyline(wf, 1, flagstring)
            end
        end
    end
    writefbodyline(wf, 1, ":preprocess => Bool") #temporary hack
    return write(wf, ")")
end

strip_split(line, args...) = strip.(split(line, args...))
function read_block(f, startstr::String, endstr::String)
    block = [startstr]
    line = readline(f)
    while !eof(f)
        while !occursin(endstr, line)
            line = readline(f)
            push!(block, line)
        end
        return block
    end
    return error("Block not found: start = $startstr, end = $endstr.")
end

function vardim(line)
    if all(occursin.(['(', ')'], (line,)))
        dim = length(split(split(split(line, '(')[2], ')')[1], ','))
    else
        dim = 0
    end
    return dim
end

function qe_write_variable(wf, indent, lines, i)
    name = gensym()
    var_i = i
    i += 2
    typ = fort2julia(strip_split(lines[i])[2])
    description = ""
    # default = nothing
    i += 1
    line = lines[i]
    while !occursin("+------", line)
        # if occursin("Default", line)
        #     _t = strip_split(line)[2]
        #     _t = strip(strip(_t,'('),')')
        #     if occursin("D", _t)
        #         default = Meta.parse(typ, replace(_t,"D" => "e"))
        #     else
        #         _t = occursin("=",_t) ?split(_t,"=")[end] : _t
        #         default = typ ==String ? _t : Meta.parse(_t)
        #         println(_t)
        #         if typeof(default) != Symbol
        #             default = convert(typ, default)
        #         end
        #     end
        if occursin("Description", line)
            description *= strip_split(line, ":")[2] * "\n"
            i += 1
            line = lines[i]
            while !occursin("+------", line)
                description *= strip(lines[i]) * "\n"
                i += 1
                line = lines[i]
            end
            @goto break_label
        end
        i += 1
        line = lines[i]
    end
    @label break_label
    line = lines[var_i]
    dim = vardim(line)
    typ = dim == 0 ? typ : Array{typ,dim}
    description = replace(replace(replace(replace(replace(description, "\"" => "'"),
                                                  "\\" => "\\\\"), "\$" => "\\\$"),
                                  "\\\\t" => "\\t"), "\\\\n" => "\\n")
    if occursin("Variables", line)
        spl = [split(x, "(")[1]
               for x in strip.(filter(x -> !occursin("=", x), split(line)[2:end]), ',')]
        names = Symbol.(spl)
        for name in names
            writefbodyline(wf, indent,
                           """QEFlagInfo{$typ}(Symbol("$name"), "$description"),""")
        end
        return i
    else
        if occursin("(", line) && occursin(")", line)
            name = Symbol(split(strip_split(line, ":")[end], "(")[1])
        else
            name = Symbol(strip_split(line, ":")[end])
        end
        writefbodyline(wf, indent, """QEFlagInfo{$typ}(Symbol("$name"), "$description"),""")
        return i
    end
end

function write_QEControlBlockInfo(wf, indent, lines)
    name = Symbol(lowercase(strip_split(lines[1], "&")[2]))
    writefbodyline(wf, indent, """QEControlBlockInfo(Symbol("$name"), [""")
    for i in 1:length(lines)
        line = lines[i]
        if occursin("Variable", line)
            i += qe_write_variable(wf, indent + 1, lines, i)
        end
    end
    return writefbodyline(wf, indent + 1, "]),")
    # writefbodyline(wf, indent, "),")
end

function write_QEDataBlockInfo(wf, indent, lines)
    spl                 = split(strip(lines[1]))
    name                = length(spl) > 1 ? lowercase(spl[2]) : "noname"
    options             = Symbol.(spl[4:2:end])
    description         = ""
    options_description = ""
    istop               = findfirst(x -> occursin("DESCRIPTION", x), lines)
    if istop === nothing
        return
    end
    for i in 1:istop-1
        line = strip(replace(replace(lines[i], "_" => ""), "-" => ""))
        if isempty(line)
            continue
        end
        description *= strip(line) * "\n" * indentabs(indent + 2)
    end
    description = replace(description, "\"" => "'")

    istart = findnext(x -> occursin("Card's flags", x), lines, istop)
    if istart != nothing
        istop = findnext(x -> occursin("+------", x), lines, istart)
        for i in istart+1:istop-1
            line = lines[i]
            if occursin("Default", line) || occursin("Description", line)
                continue
            end
            options_description *= strip(line) * "\n" * indentabs(indent + 2)
        end
    end

    options_description = replace(options_description, "\"" => "'")
    writefbodyline(wf, indent + 1,
                   """QEDataBlockInfo(Symbol("$name"),"$description", $options, "$options_description", [""")
    i = istop
    while i <= length(lines) - 1
        line = strip(lines[i])
        if occursin("Variable", line)
            i = qe_write_variable(wf, indent + 2, lines, i)
        end
        i += 1
    end
    return writefbodyline(wf, indent + 1, "]),")
end
function write_QECalculationInfo(wf, filename, indent, exec)
    writefbodyline(wf, indent, """QECalculationInfo("$exec", [""")
    allcontrolwritten = false
    anydatawritten = false
    open(filename, "r") do f
        while !eof(f)
            line = readline(f)
            if occursin("NAMELIST", line)
                write_QEControlBlockInfo(wf, indent + 1,
                                         read_block(f, line, "END OF NAMELIST"))
            elseif occursin("CARD:", line)
                anydatawritten = true
                if !allcontrolwritten
                    writefbodyline(wf, indent + 1, "],")
                    writefbodyline(wf, indent + 1, "[")
                    allcontrolwritten = true
                end
                write_QEDataBlockInfo(wf, indent + 1, read_block(f, line, "END OF CARD"))
            end
        end
    end
    if !anydatawritten
        if !allcontrolwritten
            writefbodyline(wf, indent + 1, "],")
        end
        writefbodyline(wf, indent + 1, "QEDataBlockInfo[]")
    else
        writefbodyline(wf, indent + 1, "]")
    end
    return writefbodyline(wf, indent, "),")
end

searchdir(path::String, key) = filter(x -> occursin(key, x), readdir(path))
open(joinpath(@__DIR__, "qeflags.jl"), "w") do wf
    write(wf, "_QEINPUTINFOS() = QECalculationInfo[\n")
    calculation_files = searchdir(joinpath(@__DIR__, "..", "assets", "calculations", "qe"),
                                  "INPUT")
    filepaths = joinpath.(Ref(joinpath(@__DIR__, "..", "assets", "calculations", "qe")),
                          calculation_files)
    for _f in filepaths
        exec_name = join([lowercase(splitext(split(_f, "_")[end])[1]), ".x"], "")
        write_QECalculationInfo(wf, _f, 1, exec_name)
    end
    return write(wf, "]")
end

open(joinpath(@__DIR__, "abinitflags.jl"), "w") do wf
    flaglines = split.(readlines(joinpath(@__DIR__, "..", "assets", "calculations",
                                          "abinit", "calculation_variables.txt")))
    write(wf, "_ABINITFLAGS() = Dict{Symbol, Type}(\n")
    for (fl, typ) in flaglines
        write(wf, "\t:$fl => $typ,\n")
    end
    return write(wf, ")")
end

#all this so I can be lazy and just use repr
struct ElkFlagInfo{T}
    name::Symbol
    default::Union{T,Nothing}  #check again v0.7 Some
    description::String
end
ElkFlagInfo() = ElkFlagInfo{Nothing}(:error, "")
Base.eltype(x::ElkFlagInfo{T}) where {T} = T

struct ElkControlBlockInfo
    name::Symbol
    flags::Vector{<:ElkFlagInfo}
    description::String
end

function closing(c::Char)
    if c == '{'
        return '}'
    elseif c == '('
        return ')'
    elseif c == '['
        return ']'
    end
end

function readuntil_closing(io::IO, opening_char::Char)
    out = ""
    counter = 1
    closing_char = closing(opening_char)
    while counter > 0
        c = read(io, Char)
        out *= c
        if c == opening_char
            counter += 1
        elseif c == closing_char
            counter -= 1
        end
    end
    return out
end

function elk2julia_type(s::AbstractString)
    if s == "integer"
        return Int
    elseif s == "real"
        return Float64
    elseif s == "logical"
        return Bool
    elseif s == "string"
        return String
    elseif s == "complex"
        return ComplexF64
    end
end

function replace_multiple(str, replacements::Pair{String,String}...)
    tstr = deepcopy(str)
    for r in replacements
        tstr = replace(tstr, r)
    end
    return tstr
end

function parse_elk_default(::Type{T}, s) where {T}
    s_cleaned = strip(strip(strip(replace_multiple(s, "{" => "", "}" => "", "." => "",
                                                   " " => ""), '$'), ')'), '(')
    if occursin("-", s_cleaned)
        return nothing
    elseif occursin(",", s_cleaned)
        s_t = split(s_cleaned, ',')
        return parse.(eltype(T), s_t)
    else
        return parse(T, s_cleaned)
    end
end

function blockflags(s::String)
    flag_regex = r"\{([^.}]*)\}"
    flags = ElkFlagInfo[]
    lines = split(s, "\\\\")
    for l in lines
        sline = strip_split(l, '&')
        flag = match(flag_regex, sline[1])
        descr = sline[2]
        flag_i = split(replace_multiple(sline[3], "(" => " ", ")" => " "))
        typ = length(flag_i) > 1 && flag_i[1] != "string" ?
              Vec{parse(Int, flag_i[2]),elk2julia_type(flag_i[1])} :
              elk2julia_type(flag_i[1])
        #TODO default
        default = typ == String ? sline[4] : parse_elk_default(typ, sline[4])
        flagname = Symbol(strip_split(flag.captures[1], '(')[1])
        if flagname == :wann_bands
            push!(flags, ElkFlagInfo{UnitRange{Int}}(:wann_bands, nothing, string(descr)))
        elseif flagname == :wann_projections
            push!(flags,
                  ElkFlagInfo{Vector{String}}(:wann_projections, nothing, string(descr)))
        elseif flagname == :wann_seedname
            push!(flags, ElkFlagInfo{Symbol}(:wann_seedname, nothing, string(descr)))
        else
            push!(flags, ElkFlagInfo{typ}(flagname, default, string(descr)))
        end
    end
    return flags
end

function replace_latex_symbols(s::String)
    return replace_multiple(s, "\\_" => "_", "\\hline" => "", "\\tt " => "", "^" => "e")
end

function read_elk_doc(fn::String)
    blocks = ElkControlBlockInfo[]
    blockname_regex = r"\{(.*)\}"
    open(fn, "r") do f
        while !eof(f)
            line = readline(f)
            if occursin("\\block", line)
                m = match(blockname_regex, line)
                blockname = Symbol(m.captures[1])
                block_text = replace_latex_symbols(readuntil_closing(f, '{'))
                flags = blockflags(block_text)
                push!(blocks,
                      ElkControlBlockInfo(blockname, flags, readuntil(f, "\\block")))
                seek(f, Base.position(f) - 6) #\\block offset
            end
        end
    end
    return blocks
end

ELK_CONTROLBLOCKS = read_elk_doc(joinpath(@__DIR__, "..", "assets", "calculations", "elk",
                                          "elk.txt"))

open(joinpath(@__DIR__, "elkflags.jl"), "w") do f
    return write(f, "_ELK_CONTROLBLOCKS() = $(repr(ELK_CONTROLBLOCKS))")
end


const ERRMSG_HIDE_COMMUNICATION = "@hide_communication: the code block must start with exactly one @parallel computation call and be followed by code to set boundary conditions and to perform communication. Type `?@hide_communication` for more information."
const ERRMSG_INVALID_STREAM = "@hide_communication: invalid keyword argument in @parallel: stream. Specifying the stream is not allowed in a @parallel call inside the code block argument of @hide_communication."
const ERRMSG_INVALID_COMP = "@hide_communication: invalid positional arguments in @parallel: no positional argument is allowed in the @parallel call to perform computations inside the code block argument of @hide_communication."
const ERRMSG_INVALID_BC_COMM = "@hide_communication: invalid positional arguments in @parallel: exactly one positional argument argument (ranges) must be passed in @parallel calls to set boundary conditions and to perform communication inside the code block argument of @hide_communication."

##
const HIDE_COMMUNICATION_DOC = """
    @hide_communication boundary_width block

!!! note "Advanced"
        @hide_communication ranges_outer ranges_inner block

Hide the communication behind the computation within the code `block`.

# Arguments
- `boundary_width::Tuple{Integer,Integer,Integer} | Tuple{Integer,Integer} | Tuple{Integer}`: width of the boundaries in each dimension. The boundaries must include (at least) all the data that is accessed in the communcation performed.
- `block`: code block wich starts with exactly one [`@parallel`](@ref) call to perform computations, followed by code to set boundary conditions and to perform communication (as e.g. `update_halo!` from the package `ImplicitGlobalGrid`). The [`@parallel`](@ref) call to perform computations cannot contain any positional arguments (ranges, nblocks or nthreads) nor the stream keyword argument (stream=...). The code to set boundary conditions and to perform communication must only access the elements in the boundary ranges of the fields modified in the [`@parallel`](@ref) call; all elements can be acccessed from other fields. Moreover, this code must not include statements in array broadcasting notation, because they are always run on the default CUDA stream (for CUDA.jl < v2.0), which makes CUDA stream overlapping impossible. Instead, boundary region elements can, e.g., be accessed with [`@parallel`](@ref) calls passing a ranges argument that ensures that no threads mapping to elements outside of `ranges_outer` are launched. Note that these [`@parallel`](@ref) `ranges` calls cannot contain any other positional arguments (nblocks or nthreads) nor the stream keyword argument (stream=...).

!!! note "Advanced"
    - `ranges_outer::`Tuple with one or multiple `ranges` as required by the corresponding argument of [`@parallel`](@ref): the `ranges` must together span (at least) all the data that is accessed in the communcation and boundary conditions performed.
    - `ranges_inner::`Tuple with one or multiple `ranges` as required by the corresponding argument of [`@parallel`](@ref): the `ranges` must together span the data that is not included by `ranges_outer`.

# Examples
    @hide_communication (16, 2, 2) begin
        @parallel diffusion3D_step!(Te2, Te, Ci, lam, dt, dx, dy, dz);
        update_halo!(Te2);
    end

    @hide_communication (16, 2) begin
        @parallel diffusion2D_step!(Te2, Te, Ci, lam, dt, dx, dy);
        update_halo!(Te2);
    end

    @hide_communication ranges_outer ranges_inner begin
        @parallel diffusion3D_step!(Te2, Te, Ci, lam, dt, dx, dy, dz);
        update_halo!(Te2);
    end

    @parallel_indices (iy,iz) function bc_x(A)
        A[  1, iy,  iz] = A[    2,   iy,   iz]
        A[end, iy,  iz] = A[end-1,   iy,   iz]
        return
    end
    @parallel_indices (ix,iz) function bc_y(A)
        A[ ix,  1,  iz] = A[   ix,    2,   iz]
        A[ ix,end,  iz] = A[   ix,end-1,   iz]
        return
    end
    @parallel_indices (ix,iy) function bc_z(A)
        A[ ix,  iy,  1] = A[   ix,   iy,    2]
        A[ ix,  iy,end] = A[   ix,   iy,end-1]
        return
    end
    @hide_communication (16, 2, 2) begin
        @parallel diffusion3D_step!(Te2, Te, Ci, lam, dt, dx, dy, dz);
        @parallel (1:size(Te,2), 1:size(Te,3)) bc_x(Te);
        @parallel (1:size(Te,1), 1:size(Te,3)) bc_y(Te);
        @parallel (1:size(Te,1), 1:size(Te,2)) bc_z(Te);
        update_halo!(Te2);
    end

!!! note "Developers note"
    The communcation should not perform any blocking operations to enable a maximal overlap of communication with computation.

See also: [`@parallel`](@ref)
"""
@doc HIDE_COMMUNICATION_DOC
macro hide_communication(args...) check_initialized(); checkargs_hide_communication(args...); esc(hide_communication(args...)); end

## ARGUMENT CHECKS

function checkargs_hide_communication(args...)
    if (length(args) < 2 || length(args) > 3) @ArgumentError("wrong number of arguments.") end
    if !is_block(args[end]) @ArgumentError("the last argument must be a code block (obtained: $(args[end])).") end
end

## GATEWAY FUNCTIONS

function hide_communication(args::Union{Integer,Symbol,Expr}...; package::Symbol=get_package())
    if     (package == PKG_CUDA)    hide_communication_cuda(args...)
    elseif (package == PKG_THREADS) hide_communication_threads(args...)
    else                            @KeywordArgumentError("$ERRMSG_UNSUPPORTED_PACKAGE (obtained: $package).")
    end
end

## @HIDE_COMMUNICATION FUNCTIONS

function hide_communication_cuda(ranges_outer::Union{Symbol,Expr}, ranges_inner::Union{Symbol,Expr}, block::Expr)
    compcall, bc_and_commcalls = extract_calls(block)
    parallel_args = extract_args(compcall, Symbol("@parallel"))
    posargs, kwargs, compkernelcall = split_parallel_args(parallel_args)
    if any([x.args[1]==:stream for x in kwargs]) @ArgumentError(ERRMSG_INVALID_STREAM) end
    if (length(posargs) > 0) @ArgumentError(ERRMSG_INVALID_COMP) end
    bc_and_commcalls = process_bc_and_commcalls(bc_and_commcalls)
    quote
        for i in 1:length($ranges_outer)
            ParallelStencil.ParallelKernel.@parallel_async $ranges_outer[i] stream=ParallelStencil.ParallelKernel.get_priority_custream(i) $compkernelcall
        end
        for i in 1:length($ranges_inner)
            ParallelStencil.ParallelKernel.@parallel_async $ranges_inner[i] stream=ParallelStencil.ParallelKernel.get_custream(i) $compkernelcall
        end
        for i in 1:length($ranges_outer) ParallelStencil.ParallelKernel.@synchronize(ParallelStencil.ParallelKernel.get_priority_custream(i)); end
        $bc_and_commcalls
        for i in 1:length($ranges_inner) ParallelStencil.ParallelKernel.@synchronize(ParallelStencil.ParallelKernel.get_custream(i)); end
    end
end

function hide_communication_threads(ranges_outer::Union{Symbol,Expr}, ranges_inner::Union{Symbol,Expr}, block::Expr)
    return block #NOTE: This is currently not implemented, but enables correct execution nevertheless.
end

function hide_communication_cuda(boundary_width::Union{Integer,Symbol,Expr}, block::Expr)
    compcall, = extract_calls(block)
    parallel_args = extract_args(compcall, Symbol("@parallel"))
    posargs, kwargs, compkernelcall = split_parallel_args(parallel_args)
    if any([x.args[1]==:stream for x in kwargs]) @ArgumentError(ERRMSG_INVALID_STREAM) end
    if (length(posargs) > 0) @ArgumentError(ERRMSG_INVALID_COMP) end
    ranges = :(ParallelStencil.ParallelKernel.get_ranges($(compkernelcall.args[2:end]...)))
    ranges_outer = gensym_world("ranges_outer", @__MODULE__)
    ranges_inner = gensym_world("ranges_inner", @__MODULE__)
    quote
        $ranges_outer = ParallelStencil.ParallelKernel.get_ranges_outer($boundary_width, $ranges)
        $ranges_inner = ParallelStencil.ParallelKernel.get_ranges_inner($boundary_width, $ranges)
        ParallelStencil.ParallelKernel.@hide_communication $ranges_outer $ranges_inner $block  # NOTE: to enable macros like @hide_communication_cuda, one would need to use everywhere in hide_communication_cuda _cuda and _thread suffixed macros and enable their processing. This is more complex and generates less readable code though. => Design decision: only enable _cuda/_threads macros for macros that do not return package dependent macros.
    end
end

function hide_communication_threads(boundary_width::Union{Integer,Symbol,Expr}, block::Expr)
    return block #NOTE: This is currently not implemented, but enables correct execution nevertheless.
end

## FUNCTIONS TO EXTRACT AND PROCESS COMPUTATION AND BOUNDARY CONDITIONS CALLS / COMMUNICATION CALLS

function extract_calls(block::Expr)
    if (sum([isa(arg, Expr) for arg in block.args]) < 2) @ArgumentError(ERRMSG_HIDE_COMMUNICATION) end
    compcall         = block.args[2]
    bc_and_commcalls = quote $(block.args[3:end]...) end
    if !is_parallel_call(compcall) @ArgumentError(ERRMSG_HIDE_COMMUNICATION) end
    return compcall, bc_and_commcalls
end

function process_bc_and_commcalls(block::Expr)
    return postwalk(block) do x
        if !is_parallel_call(x) return x; end
        @capture(x, @parallel args__) || @ModuleInternalError("a @parallel call could not be pattern matched.")
        posargs, kwargs = split_parallel_args(args)
        if any([x.args[1]==:stream for x in kwargs]) @ArgumentError(ERRMSG_INVALID_STREAM) end
        if (length(posargs) != 1) @ArgumentError(ERRMSG_INVALID_BC_COMM) end
        return :(@parallel $(args[1:end-1]...) stream = ParallelStencil.ParallelKernel.get_priority_custream(1) $(args[end]))
    end
end

## FUNCTIONS TO GET INNER AND OUTER RANGES AND TO PROMOTE BOUNDARY_WIDTH TO 3D

promote_boundary_width(boundary_width::BOUNDARY_WIDTH_TYPE_1D)       = (boundary_width,    0, 0)
promote_boundary_width(boundary_width::BOUNDARY_WIDTH_TYPE_1D_TUPLE) = (boundary_width..., 0, 0)
promote_boundary_width(boundary_width::BOUNDARY_WIDTH_TYPE_2D)       = (boundary_width..., 0)
promote_boundary_width(boundary_width::BOUNDARY_WIDTH_TYPE)          = boundary_width
promote_boundary_width(boundary_width)                               = @ArgumentError("@hide_communication: boundary_width must be a Tuple of Integer of size 1, 2 or 3 (obtained: $boundary_width; its type is: $(typeof(boundary_width))).")

function get_ranges_outer(boundary_width, ranges::RANGES_TYPE) where T <:Integer
    boundary_width = promote_boundary_width(boundary_width)
    validate_ranges_args(boundary_width, ranges)
    ms = length.(ranges)
    bw = boundary_width
    if ms[3] > 1 # 3D
        ranges_outer = ((1:ms[1],             1:ms[2],             1:bw[3]),
                        (1:ms[1],             1:ms[2],             ms[3]-bw[3]+1:ms[3]),
                        (1:ms[1],             1:bw[2],             bw[3]+1:ms[3]-bw[3]),
                        (1:ms[1],             ms[2]-bw[2]+1:ms[2], bw[3]+1:ms[3]-bw[3]),
                        (1:bw[1],             bw[2]+1:ms[2]-bw[2], bw[3]+1:ms[3]-bw[3]),
                        (ms[1]-bw[1]+1:ms[1], bw[2]+1:ms[2]-bw[2], bw[3]+1:ms[3]-bw[3]),
                       )
    elseif ms[2] > 1 # 2D
        ranges_outer = ((1:ms[1],             1:bw[2],             1:1),
                        (1:ms[1],             ms[2]-bw[2]+1:ms[2], 1:1),
                        (1:bw[1],             bw[2]+1:ms[2]-bw[2], 1:1),
                        (ms[1]-bw[1]+1:ms[1], bw[2]+1:ms[2]-bw[2], 1:1),
                       )
    elseif ms[1] > 1  # 1D
        ranges_outer = ((1:bw[1],             1:1, 1:1),
                        (ms[1]-bw[1]+1:ms[1], 1:1, 1:1),
                       )
    else
        @ModuleInternalError("invalid argument 'ranges'.")
    end
    return Tuple([r for r in ranges_outer if all(length.(r) .!= 0)])
end

function get_ranges_inner(boundary_width, ranges::RANGES_TYPE) where T <:Integer
    boundary_width = promote_boundary_width(boundary_width)
    validate_ranges_args(boundary_width, ranges)
    ms = length.(ranges)
    bw = boundary_width
    if     (ms[3] > 1) return ( (bw[1]+1:ms[1]-bw[1], bw[2]+1:ms[2]-bw[2], bw[3]+1:ms[3]-bw[3]), ) # 3D
    elseif (ms[2] > 1) return ( (bw[1]+1:ms[1]-bw[1], bw[2]+1:ms[2]-bw[2], 1:1), ) # 2D
    elseif (ms[1] > 1) return ( (bw[1]+1:ms[1]-bw[1], 1:1, 1:1), ) # 1D
    else @ModuleInternalError("invalid argument 'ranges'.")
    end
end

function validate_ranges_args(boundary_width::Tuple{T,T,T}, ranges::RANGES_TYPE) where T <:Integer
    maxsize = length.(ranges)
    if any([r.start != 1 for r in ranges]) @ArgumentError("unsupported call to `@hide_communication boundary_width block`: automatic ranges determination for @hide_communication is not supported for ranges with start != 1. Use the explicit version `@hide_communication ranges_outer ranges_inner block` instead.") end
    if !all(maxsize .> 0) @IncoherentArgumentError("incoherent arguments in @hide_communication: the following must be true: all(length.(ranges) .> 0) ") end
    if !all(maxsize .> 2 .* boundary_width) @IncoherentArgumentError("incoherent arguments in @hide_communication: the following must be true: all(length.(ranges) .> 2 .* boundary_width) ") end # NOTE: this ensures among other things that if maxsize == 1, then boundary_width==0.
    if any((maxsize .> 1) .& (boundary_width .== 0)) @IncoherentArgumentError("incoherent arguments in @hide_communication: the following must be false: any((length.(ranges) .> 1) .& (boundary_width .== 0)) ") end
    if all(maxsize .== 1) @IncoherentArgumentError("incoherent arguments in @hide_communication: the following must be false: all(length.(ranges) .== 1)") end
end

## FUNCTIONS TO GET CREATE AND MANAGE CUDA STREAMS

@static if ENABLE_CUDA
    let
        global get_priority_custream, get_custream
        priority_custreams = Array{CuStream}(undef, 0)
        custreams          = Array{CuStream}(undef, 0)

        function get_priority_custream(id::Integer)
            while (id > length(priority_custreams)) push!(priority_custreams, CuStream(; flags=CUDA.STREAM_NON_BLOCKING, priority=CUDA.priority_range()[end])) end # CUDA.priority_range()[end] is max priority. # NOTE: priority_range cannot be called outside the function as only at runtime sure that CUDA is functional.
            return priority_custreams[id]
        end

        function get_custream(id::Integer)
            priority_min =
            while (id > length(custreams)) push!(custreams, CuStream(; flags=CUDA.STREAM_NON_BLOCKING, priority=CUDA.priority_range()[1])) end # CUDA.priority_range()[1] is min priority. # NOTE: priority_range cannot be called outside the function as only at runtime sure that CUDA is functional.
            return custreams[id]
        end
    end
end


using Test
using aoc2019.Fuel
using aoc2019.computer
using aoc2019.wiring
using aoc2019.password
using aoc2019.star_chart
using aoc2019.thrusters
using aoc2019.sif
using aoc2019.asteroids
using aoc2019.Nanofactory

@testset "the tyranny of the rocket equation -- part one                          " begin
    @test Fuel.total([12]) == 2
    @test Fuel.total([14]) == 2
    @test Fuel.total([1969]) == 654
    @test Fuel.total([100756]) == 33583
end

@testset "the tyranny of the rocket equation -- part two                          " begin
    @test Fuel.correct_total([14]) == 2
    @test Fuel.correct_total([1969]) == 966
    @test Fuel.correct_total([100756]) == 50346
end

@testset "1202 program alarm -- part one                                          " begin
    @test computer.run([1,0,0,0,99]) == [2,0,0,0,99]
    @test computer.run([2,3,0,3,99]) == [2,3,0,6,99]
    @test computer.run([2,4,4,5,99,0]) == [2,4,4,5,99,9801]
    @test computer.run([1,1,1,4,99,5,6,0,99]) == [30,1,1,4,2,5,6,0,99]
end

@testset "crossed wires -- part one                                               " begin
    @test wiring.manhattan_distance(
        ["R75","D30","R83","U83","L12","D49","R71","U7","L72"],
        ["U62","R66","U55","R34","D71","R55","D58","R83"]
    ) == 159

    @test wiring.manhattan_distance(
        ["R98","U47","R26","D63","R33","U87","L62","D20","R33","U53","R51"],
        ["U98","R91","D20","R16","D67","R40","U7","R15","U6","R7"]
    ) == 135
end

@testset "crossed wires -- part two                                               " begin
    @test wiring.signal_delay(
        ["R75","D30","R83","U83","L12","D49","R71","U7","L72"],
        ["U62","R66","U55","R34","D71","R55","D58","R83"]
    ) == 610

    @test wiring.signal_delay(
        ["R98","U47","R26","D63","R33","U87","L62","D20","R33","U53","R51"],
        ["U98","R91","D20","R16","D67","R40","U7","R15","U6","R7"]
    ) == 410
end

@testset "secure container -- part one                                            " begin
    @test password.is_candidate(111111) == true
    @test password.is_candidate(223450) == false
    @test password.is_candidate(123789) == false
end

@testset "secure container -- part two                                            " begin
    @test password.is_candidate(112233) && password.is_candidate_strict(112233) == true
    @test password.is_candidate(123444) && password.is_candidate_strict(123444)== false
    @test password.is_candidate(111122) && password.is_candidate_strict(1111122)== true
end

@testset "sunny with a chance of asteroids -- part one                            " begin
    @test computer.run([1002,4,3,4,33]) == [1002,4,3,4,99]
end

@testset "sunny with a chance of asteroids -- part two                            " begin
    # equals tests
    stdin, stdout = io()
    put!(stdin, 8)
    @async computer.run([3,9,8,9,10,9,4,9,99,-1,8], stdin, stdout)
    @test take!(stdout) == 1

    stdin, stdout = io()
    put!(stdin, 9)
    @async computer.run([3,9,8,9,10,9,4,9,99,-1,8], stdin, stdout)
    @test take!(stdout) == 0

    stdin, stdout = io()
    put!(stdin, 8)
    @async computer.run([3,3,1108,-1,8,3,4,3,99], stdin, stdout)
    @test take!(stdout) == 1

    stdin, stdout = io()
    put!(stdin, 9)
    @async computer.run([3,3,1108,-1,8,3,4,3,99], stdin, stdout)
    @test take!(stdout) == 0

    # less than tests
    stdin, stdout = io()
    put!(stdin, 7)
    @async computer.run([3,9,7,9,10,9,4,9,99,-1,8], stdin, stdout)
    @test take!(stdout) == 1

    stdin, stdout = io()
    put!(stdin, 8)
    @async computer.run([3,9,7,9,10,9,4,9,99,-1,8], stdin, stdout)
    @test take!(stdout) == 0

    stdin, stdout = io()
    put!(stdin, 9)
    @async computer.run([3,9,7,9,10,9,4,9,99,-1,8], stdin, stdout)
    @test take!(stdout) == 0

    stdin, stdout = io()
    put!(stdin, 7)
    @async computer.run([3,3,1107,-1,8,3,4,3,99], stdin, stdout)
    @test take!(stdout) == 1

    stdin, stdout = io()
    put!(stdin, 8)
    @async computer.run([3,3,1107,-1,8,3,4,3,99], stdin, stdout)
    @test take!(stdout) == 0

    stdin, stdout = io()
    put!(stdin, 9)
    @async computer.run([3,3,1107,-1,8,3,4,3,99], stdin, stdout)
    @test take!(stdout) == 0

    # jump tests
    stdin, stdout = io()
    put!(stdin, 0)
    @async computer.run([3,12,6,12,15,1,13,14,13,4,13,99,-1,0,1,9], stdin, stdout)
    @test take!(stdout) == 0

    stdin, stdout = io()
    put!(stdin, 99)
    @async computer.run([3,12,6,12,15,1,13,14,13,4,13,99,-1,0,1,9], stdin, stdout)
    @test take!(stdout) == 1

    stdin, stdout = io()
    put!(stdin, 0)
    @async computer.run([3,3,1105,-1,9,1101,0,0,12,4,12,99,1], stdin, stdout)
    @test take!(stdout) == 0

    # other
    stdin, stdout = io()
    put!(stdin, 7)
    @async computer.run(
        [3,21,1008,21,8,20,1005,20,22,107,8,21,20,1006,20,31,
            1106,0,36,98,0,0,1002,21,125,20,4,20,1105,1,46,104,
            999,1105,1,46,1101,1000,1,20,4,20,1105,1,46,98,99
        ],
        stdin,
        stdout
    )
    @test take!(stdout) == 999

    stdin, stdout = io()
    put!(stdin, 8)
    @async computer.run(
        [3,21,1008,21,8,20,1005,20,22,107,8,21,20,1006,20,31,
            1106,0,36,98,0,0,1002,21,125,20,4,20,1105,1,46,104,
            999,1105,1,46,1101,1000,1,20,4,20,1105,1,46,98,99
        ],
        stdin,
        stdout
    )
    @test take!(stdout) == 1000

    stdin, stdout = io()
    put!(stdin, 9)
    @async computer.run(
        [3,21,1008,21,8,20,1005,20,22,107,8,21,20,1006,20,31,
            1106,0,36,98,0,0,1002,21,125,20,4,20,1105,1,46,104,
            999,1105,1,46,1101,1000,1,20,4,20,1105,1,46,98,99
        ],
        stdin,
        stdout
    )
    @test take!(stdout) == 1001
end

@testset "universal orbit map -- part one                                         " begin
    @test star_chart.total_orbits([
        "COM)B", "B)C", "C)D", "D)E", "E)F", "B)G", "G)H", "D)I", "E)J", "J)K", "K)L"
    ]) == 42
end

@testset "universal orbit map -- part two                                         " begin
    @test star_chart.total_transfers([
        "COM)B", "B)C", "C)D", "D)E", "E)F", "B)G", "G)H", "D)I",
        "E)J", "J)K", "K)L", "K)YOU", "I)SAN"
    ]) == 4
end

@testset "amplification circuit -- part one                                       " begin
    @test thrusters.output(
        [3,15,3,16,1002,16,10,16,1,16,15,15,4,15,99,0,0],
        [4,3,2,1,0],
        0
    ) == 43210

    @test thrusters.output(
        [3,23,3,24,1002,24,10,24,1002,23,-1,23,101,5,23,23,1,24,23,23,4,23,99,0,0],
        [0,1,2,3,4],
        0
    ) == 54321

    @test thrusters.output(
        [3,31,3,32,1002,32,10,32,1001,31,-2,31,1007,31,0,33,
            1002,33,7,33,1,33,31,31,1,32,31,31,4,31,99,0,0,0
        ],
        [1,0,4,3,2],
        0
    ) == 65210
end

@testset "amplification circuit -- part two                                       " begin
    @test thrusters.feedback_loop(
        [3,26,1001,26,-4,26,3,27,1002,27,2,27,1,27,26,
            27,4,27,1001,28,-1,28,1005,28,6,99,0,0,5
        ],
        [9,8,7,6,5],
        0
    ) == 139629729

    @test thrusters.feedback_loop(
        [3,52,1001,52,-5,52,3,53,1,52,56,54,1007,54,5,55,1005,55,26,1001,
            54,-5,54,1105,1,12,1,53,54,53,1008,54,0,55,1001,55,1,55,2,53,
            55,53,4,53,1001,56,-1,56,1005,56,6,99,0,0,0,0,10
        ],
        [9,7,8,5,6],
        0
    ) == 18216
end

@testset "space image format -- part one                                          " begin
    @test sif.checksum(sif.image([1,2,3,4,5,6,7,8,9,0,1,2], 3, 2)) == 1
    @test sif.checksum(sif.image([0,0,0,0,0,1,1,1,1,1,1,2], 2, 2)) == 3
end

@testset "space image format -- part two                                          " begin
    @test sif.resolve(sif.image([0,2,2,2,1,1,2,2,2,2,1,2,0,0,0,0], 2, 2)) == [0 1; 1 0]
end

@testset "sensor boost -- part one                                                " begin
    program = [109,1,204,-1,1001,100,1,100,1008,100,16,101,1006,101,0,99]
    stdin, stdout = computer.io()
    task = @async computer.run(copy(program), stdin, stdout)
    while !istaskdone(task); yield(); end
    close(stdout)
    @test [o for o in stdout] == program

    program = [1102,34915192,34915192,7,4,7,99,0]
    stdin, stdout = computer.io()
    task = @async computer.run(program, stdin, stdout)
    while !istaskdone(task); yield(); end
    @test length(digits(take!(stdout))) == 16

    program = [104,1125899906842624,99]
    stdin, stdout = computer.io()
    task = @async computer.run(program, stdin, stdout)
    while !istaskdone(task); yield(); end
    @test take!(stdout) == 1125899906842624
end

@testset "monitoring station -- part one                                          " begin
    map = asteroids.chart([".#..#", ".....", "#####", "....#", "...##"])
    (coords, visible) = asteroids.maxima(map)
    @test coords == asteroids.Point(4, 5)
    @test visible == 8

    map = asteroids.chart([
        "......#.#.",
        "#..#.#....",
        "..#######.",
        ".#.#.###..",
        ".#..#.....",
        "..#....#.#",
        "#..#....#.",
        ".##.#..###",
        "##...#..#.",
        ".#....####"
    ])
    (coords, visible) = asteroids.maxima(map)
    @test coords == asteroids.Point(6, 9)
    @test visible == 33

    map = asteroids.chart([
        "#.#...#.#.",
        ".###....#.",
        ".#....#...",
        "##.#.#.#.#",
        "....#.#.#.",
        ".##..###.#",
        "..#...##..",
        "..##....##",
        "......#...",
        ".####.###."
    ])
    (coords, visible) = asteroids.maxima(map)
    @test coords == asteroids.Point(2, 3)
    @test visible == 35

    map = asteroids.chart([
        ".#..##.###...#######",
        "##.############..##.",
        ".#.######.########.#",
        ".###.#######.####.#.",
        "#####.##.#.##.###.##",
        "..#####..#.#########",
        "####################",
        "#.####....###.#.#.##",
        "##.#################",
        "#####.##.###..####..",
        "..######..##.#######",
        "####.##.####...##..#",
        ".#####..#.######.###",
        "##...#.##########...",
        "#.##########.#######",
        ".####.#.###.###.#.##",
        "....##.##.###..#####",
        ".#.#.###########.###",
        "#.#.#.#####.####.###",
        "###.##.####.##.#..##"
    ])
    (coords, visible) = asteroids.maxima(map)
    @test coords == asteroids.Point(12, 14)
    @test visible == 210
end

@testset "monitoring station -- part two                                          " begin
    tau = 2 * atan(0, -1)
    map = asteroids.chart([
        ".#....#####...#..",
        "##...##.#####..##",
        "##...#...#.#####.",
        "..#.....X...###..",
        "..#.#.....#....##"
    ])
    asts = asteroids.sweep(map, Point(9, 4), atan(-1, 0) + tau, atan(-1, 0))
    @test asts[1] == Point(9, 2)
    @test asts[9] == Point(16, 2)
end

@testset "space stoichiometry -- part one                                         " begin
    formulae = Nanofactory.parse_equations([
        "10 ORE => 10 A"
        "1 ORE => 1 B"
        "7 A, 1 B => 1 C"
        "7 A, 1 C => 1 D"
        "7 A, 1 D => 1 E"
        "7 A, 1 E => 1 FUEL"
    ])
    inventory = Dict("FUEL" => -1)
    Nanofactory.produce(inventory, formulae)
    ore = abs(inventory["ORE"])
    @test ore == 31

    formulae = Nanofactory.parse_equations([
        "9 ORE => 2 A"
        "8 ORE => 3 B"
        "7 ORE => 5 C"
        "3 A, 4 B => 1 AB"
        "5 B, 7 C => 1 BC"
        "4 C, 1 A => 1 CA"
        "2 AB, 3 BC, 4 CA => 1 FUEL"
    ])
    inventory = Dict("FUEL" => -1)
    Nanofactory.produce(inventory, formulae)
    ore = abs(inventory["ORE"])
    @test ore == 165

    formulae = Nanofactory.parse_equations([
        "157 ORE => 5 NZVS"
        "165 ORE => 6 DCFZ"
        "44 XJWVT, 5 KHKGT, 1 QDVJ, 29 NZVS, 9 GPVTF, 48 HKGWZ => 1 FUEL"
        "12 HKGWZ, 1 GPVTF, 8 PSHF => 9 QDVJ"
        "179 ORE => 7 PSHF"
        "177 ORE => 5 HKGWZ"
        "7 DCFZ, 7 PSHF => 2 XJWVT"
        "165 ORE => 2 GPVTF"
        "3 DCFZ, 7 NZVS, 5 HKGWZ, 10 PSHF => 8 KHKGT"
    ])
    inventory = Dict("FUEL" => -1)
    Nanofactory.produce(inventory, formulae)
    ore = abs(inventory["ORE"])
    @test ore == 13312

    formulae = Nanofactory.parse_equations([
        "2 VPVL, 7 FWMGM, 2 CXFTF, 11 MNCFX => 1 STKFG"
        "17 NVRVD, 3 JNWZP => 8 VPVL"
        "53 STKFG, 6 MNCFX, 46 VJHF, 81 HVMC, 68 CXFTF, 25 GNMV => 1 FUEL"
        "22 VJHF, 37 MNCFX => 5 FWMGM"
        "139 ORE => 4 NVRVD"
        "144 ORE => 7 JNWZP"
        "5 MNCFX, 7 RFSQX, 2 FWMGM, 2 VPVL, 19 CXFTF => 3 HVMC"
        "5 VJHF, 7 MNCFX, 9 VPVL, 37 CXFTF => 6 GNMV"
        "145 ORE => 6 MNCFX"
        "1 NVRVD => 8 CXFTF"
        "1 VJHF, 6 MNCFX => 4 RFSQX"
        "176 ORE => 6 VJHF"
    ])
    inventory = Dict("FUEL" => -1)
    Nanofactory.produce(inventory, formulae)
    ore = abs(inventory["ORE"])
    @test ore == 180697

    formulae = Nanofactory.parse_equations([
        "171 ORE => 8 CNZTR"
        "7 ZLQW, 3 BMBT, 9 XCVML, 26 XMNCP, 1 WPTQ, 2 MZWV, 1 RJRHP => 4 PLWSL"
        "114 ORE => 4 BHXH"
        "14 VRPVC => 6 BMBT"
        "6 BHXH, 18 KTJDG, 12 WPTQ, 7 PLWSL, 31 FHTLT, 37 ZDVW => 1 FUEL"
        "6 WPTQ, 2 BMBT, 8 ZLQW, 18 KTJDG, 1 XMNCP, 6 MZWV, 1 RJRHP => 6 FHTLT"
        "15 XDBXC, 2 LTCX, 1 VRPVC => 6 ZLQW"
        "13 WPTQ, 10 LTCX, 3 RJRHP, 14 XMNCP, 2 MZWV, 1 ZLQW => 1 ZDVW"
        "5 BMBT => 4 WPTQ"
        "189 ORE => 9 KTJDG"
        "1 MZWV, 17 XDBXC, 3 XCVML => 2 XMNCP"
        "12 VRPVC, 27 CNZTR => 2 XDBXC"
        "15 KTJDG, 12 BHXH => 5 XCVML"
        "3 BHXH, 2 VRPVC => 7 MZWV"
        "121 ORE => 7 VRPVC"
        "7 XCVML => 6 RJRHP"
        "5 BHXH, 4 VRPVC => 5 LTCX"
    ])
    inventory = Dict("FUEL" => -1)
    Nanofactory.produce(inventory, formulae)
    ore = abs(inventory["ORE"])
    @test ore == 2210736
end

@testset "space stoichiometry -- part two                                         " begin
    formulae = Nanofactory.parse_equations([
        "157 ORE => 5 NZVS"
        "165 ORE => 6 DCFZ"
        "44 XJWVT, 5 KHKGT, 1 QDVJ, 29 NZVS, 9 GPVTF, 48 HKGWZ => 1 FUEL"
        "12 HKGWZ, 1 GPVTF, 8 PSHF => 9 QDVJ"
        "179 ORE => 7 PSHF"
        "177 ORE => 5 HKGWZ"
        "7 DCFZ, 7 PSHF => 2 XJWVT"
        "165 ORE => 2 GPVTF"
        "3 DCFZ, 7 NZVS, 5 HKGWZ, 10 PSHF => 8 KHKGT"
    ])
    low = 1
    high = 1000000000000
    while low < high
        guess = div(low + high + 1, 2)
        inventory = Dict("FUEL" => -1 * guess)
        Nanofactory.produce(inventory, formulae)
        ore = abs(inventory["ORE"])
        if ore < 1000000000000
            low = guess
        else
            high = guess - 1
        end
    end
    @test low == 82892753

    formulae = Nanofactory.parse_equations([
        "2 VPVL, 7 FWMGM, 2 CXFTF, 11 MNCFX => 1 STKFG"
        "17 NVRVD, 3 JNWZP => 8 VPVL"
        "53 STKFG, 6 MNCFX, 46 VJHF, 81 HVMC, 68 CXFTF, 25 GNMV => 1 FUEL"
        "22 VJHF, 37 MNCFX => 5 FWMGM"
        "139 ORE => 4 NVRVD"
        "144 ORE => 7 JNWZP"
        "5 MNCFX, 7 RFSQX, 2 FWMGM, 2 VPVL, 19 CXFTF => 3 HVMC"
        "5 VJHF, 7 MNCFX, 9 VPVL, 37 CXFTF => 6 GNMV"
        "145 ORE => 6 MNCFX"
        "1 NVRVD => 8 CXFTF"
        "1 VJHF, 6 MNCFX => 4 RFSQX"
        "176 ORE => 6 VJHF"
    ])
    low = 1
    high = 1000000000000
    while low < high
        guess = div(low + high + 1, 2)
        inventory = Dict("FUEL" => -1 * guess)
        Nanofactory.produce(inventory, formulae)
        ore = abs(inventory["ORE"])
        if ore < 1000000000000
            low = guess
        else
            high = guess - 1
        end
    end
    @test low == 5586022

    formulae = Nanofactory.parse_equations([
        "171 ORE => 8 CNZTR"
        "7 ZLQW, 3 BMBT, 9 XCVML, 26 XMNCP, 1 WPTQ, 2 MZWV, 1 RJRHP => 4 PLWSL"
        "114 ORE => 4 BHXH"
        "14 VRPVC => 6 BMBT"
        "6 BHXH, 18 KTJDG, 12 WPTQ, 7 PLWSL, 31 FHTLT, 37 ZDVW => 1 FUEL"
        "6 WPTQ, 2 BMBT, 8 ZLQW, 18 KTJDG, 1 XMNCP, 6 MZWV, 1 RJRHP => 6 FHTLT"
        "15 XDBXC, 2 LTCX, 1 VRPVC => 6 ZLQW"
        "13 WPTQ, 10 LTCX, 3 RJRHP, 14 XMNCP, 2 MZWV, 1 ZLQW => 1 ZDVW"
        "5 BMBT => 4 WPTQ"
        "189 ORE => 9 KTJDG"
        "1 MZWV, 17 XDBXC, 3 XCVML => 2 XMNCP"
        "12 VRPVC, 27 CNZTR => 2 XDBXC"
        "15 KTJDG, 12 BHXH => 5 XCVML"
        "3 BHXH, 2 VRPVC => 7 MZWV"
        "121 ORE => 7 VRPVC"
        "7 XCVML => 6 RJRHP"
        "5 BHXH, 4 VRPVC => 5 LTCX"
    ])
    low = 1
    high = 1000000000000
    while low < high
        guess = div(low + high + 1, 2)
        inventory = Dict("FUEL" => -1 * guess)
        Nanofactory.produce(inventory, formulae)
        ore = abs(inventory["ORE"])
        if ore < 1000000000000
            low = guess
        else
            high = guess - 1
        end
    end
    @test low == 460664
end