# const LAYER = Layer
# const LAYER = NormalizedLayer

type NeuralNet{LAYER} <: NetStat
  layers::Vector{LAYER}  # note: this doesn't include input layer!!
  params::NetParams
  solverParams::SolverParams
  inputTransformer::Transformer
  transformedInput::VecF   # so we can avoid allocations
  # costmult::VecF

  # TODO: inner constructor which performs some sanity checking on activation/cost combinations:
  function NeuralNet{LAYER}(layers::Vector{LAYER},
                      params::NetParams,
                      solverParams::SolverParams,
                      inputTransformer::Transformer = IdentityTransformer())

    # do some sanity checking on activation/costmodel combos
    if isa(params.mloss, CrossentropyLoss)
      @assert isa(layers[end].activation, layers[end].nout > 1 ? SoftmaxMapping : SigmoidMapping)
    end

    if isa(layers[end].activation, SoftmaxMapping)
      @assert isa(params.mloss, CrossentropyLoss)
    end

    new(layers, params, solverParams, inputTransformer, zeros(first(layers).nin)) #, zeros(last(layers).nout))
  end
end

# # simple constructor which creates all layers the same for given list of node counts.
# # structure should include neuron counts for all layers, including input and output
# function NeuralNet(structure::AVec{Int};
#                     params = NetParams(),
#                     solverParams = SolverParams(),
#                     activation::Mapping = TanhMapping(),
#                     inputTransformer::Transformer = IdentityTransformer())
#   @assert length(structure) > 1

#   layers = LAYER[]
#   for i in 1:length(structure)-1
#     nin, nout = structure[i:i+1]
#     pDropout = getDropoutProb(params, i==1)
#     push!(layers, layerType(nin, nout, activation, params.updater, pDropout))
#   end

#   NeuralNet(layers, params, solverParams, inputTransformer)
# end

function Base.show(io::IO, net::NeuralNet)
  println(io, "NeuralNet{")
  println(io, "  params: $(net.params)")
  println(io, "  solverParams: $(net.solverParams)")
  println(io, "  layers:")
  for layer in net.layers
    println(io, "    ", layer)
  end
  println(io, "}")
end
Base.print(io::IO, net::NeuralNet) = show(io, net)

# ------------------------------------------------------------------------

# produces a vector of output (estimated outputs) from the network
function forward!(net::NeuralNet, input::AVecF, istraining::Bool = false)

  # first transform the input
  transform!(net.inputTransformer, net.transformedInput, input)

  # now feed it forward
  output = net.transformedInput
  for layer in net.layers
    forward!(layer, output, istraining)
    output = layer.a
  end

  # update nextr
  for i in 1:length(net.layers)-1
    net.layers[i].nextr = net.layers[i+1].r
  end

  output
end

# given a vector of errors (target - output), update network weights
function backward!(net::NeuralNet, output, target)

  # update δᵢ starting from the output layer
  layer = net.layers[end]
  # δ = sensitivity!(layer.δ, layer.activation, net.params.mloss, layer.x, output, target)
  updateSensitivities!(layer, net.params.mloss, output, target)

  # now update the remaining sensitivities using bakckprop
  for i in length(net.layers)-1:-1:1
    updateSensitivities!(net.layers[i:i+1]...)
  end

  # now update the weights
  for layer in net.layers
    updateWeights!(layer, net.params.updater)

    # println()
    # @show layer
  end

  # # update our η, μ, etc
  # fit!(net.params)
end

# ------------------------------------------------------------------------

# ϕₒ = sensitivity!(outputnode.state.ϕ,
#                   outputnode.mapping,
#                   net.mloss,
#                   input,
#                   output,
#                   target)

# online version... returns the feedforward estimate before updating
function OnlineStats.fit!(net::NeuralNet, input::AVecF, target::AVecF)
  output = forward!(net, input, true)
  backward!(net, output, target)
  output
end

# batch version
function OnlineStats.fit!(net::NeuralNet, input::MatF, target::MatF)
  @assert ncols(input) == net.nin
  @assert ncols(target) == net.nout
  @assert nrows(input) == nrows(target)

  Float64[fit!(net, row(input,i), row(target,i)) for i in 1:nrows(input)]
end

# note: when yEqualsX is true, we are updating an autoencoder (or similar) and so we can
# use net.transformedInput instead of target
function OnlineStats.fit!(net::NetStat, data::DataPoint, yEqualsX::Bool = false)
  fit!(net, data.x, yEqualsX ? net.transformedInput : data.y)
  # fit!(net, data.x, transformY ? transform(net.inputTransformer, data.y) : data.y)
end

# ------------------------------------------------------------------------

function cost(net::NeuralNet, input::AVecF, target::AVecF)
  output = forward!(net, input)
  sum(value(net.params.mloss, target, output))
end
totalCost(net::NetStat, data::DataPoint) = cost(net, data.x, data.y)
totalCost(net::NetStat, dataset::DataPoints) = sum([totalCost(net, data) for data in dataset])
totalCost(net::NetStat, sampler::DataSampler) = totalCost(net, DataPoints(sampler))

# ------------------------------------------------------------------------

function StatsBase.predict(net::NeuralNet, input::AVecF)
  forward!(net, input)
end

function StatsBase.predict(net::NeuralNet, input::AMatF)
  output = zeros(nrows(input), net.layers[end].nout)
  for i in 1:nrows(input)
    row!(output, i, predict(net, row(input, i)))
  end
  output
end

# ------------------------------------------------------------------------

# note: we scale standard random normals by (1/sqrt(nin)) so that the distribution of initial (Σ = wx + b)
#       is also approximately standard normal
_initialWeights(nin::Int, nout::Int, activation::Mapping = IdentityMapping()) = 0.5randn(nout, nin) / sqrt(nin)

# initialWeights(nin::Int, nout::Int, activation::Mapping) = (rand(nout, nin) - 0.5) * 2.0 * sqrt(6.0 / (nin + nout))

# note: we scale standard random normals by (1/sqrt(nin)) so that the distribution of initial (Σ = wx + b)
#       is also approximately standard normal
# initialWeights(nin::Int, nout::Int, activation::Mapping) = randn(nout, nin) / sqrt(nin)

# ------------------------------------------------------------------------


function TurnAroundRangeIonosphericDelayModifier(arg0::IonosphericModel, arg1::jdouble)
    return TurnAroundRangeIonosphericDelayModifier((IonosphericModel, jdouble), arg0, arg1)
end

function get_parameters_drivers(obj::TurnAroundRangeIonosphericDelayModifier)
    return jcall(obj, "getParametersDrivers", List, ())
end

function modify(obj::TurnAroundRangeIonosphericDelayModifier, arg0::EstimatedMeasurement)
    return jcall(obj, "modify", void, (EstimatedMeasurement,), arg0)
end


# Roberta specific initializers
# same as bert

# HGF Roberta compounts

# embedding

struct HGFRobertaEmbeddings{
  W<:FakeTHEmbedding,
  P<:FakeTHEmbedding,
  T<:FakeTHEmbedding,
  L<:FakeTHLayerNorm
} <: THModule
  LayerNorm::L
  word_embeddings::W
  position_embeddings::P
  token_type_embeddings::T
end

@functor HGFRobertaEmbeddings

function HGFRobertaEmbeddings(config::HGFRobertaConfig)
  layernorm = FakeTHLayerNorm(HGFBertConfig, config, config.hidden_size; eps=config.layer_norm_eps)

  word_emb = FakeTHEmbedding(HGFBertConfig, config, config.vocab_size,
                             config.hidden_size;
                             pad_idx=config.pad_token_id)
  posi_emb = FakeTHEmbedding(HGFBertConfig, config, config.max_position_embeddings,
                             config.hidden_size;
                             pad_idx=config.pad_token_id)
  toke_emb = FakeTHEmbedding(HGFBertConfig, config, config.type_vocab_size,
                             config.hidden_size)

  HGFRobertaEmbeddings(layernorm, word_emb, posi_emb, toke_emb)
end

@inline get_word_emb(be::HGFRobertaEmbeddings, input_ids::AbstractArray{<:Integer}) = be.word_embeddings(input_ids)
@inline get_word_emb(be::HGFRobertaEmbeddings, input_embed::AbstractArray{T}) where T = input_embed

function _create_position_ids_from_input_ids(input_ids, pad_idx)
  mask = input_ids .!= pad_idx
  inc_indices = cumsum(mask; dims=1) .* mask
  return inc_indices .+ pad_idx
end

Flux.@nograd _create_position_ids_from_input_ids

@inline function get_position_emb(be::HGFRobertaEmbeddings, input_ids::AbstractArray{<:Integer}, ::Nothing)
  position_ids = _create_position_ids_from_input_ids(input_ids, be.position_embeddings.pad_idx)
  get_position_emb(be, input_ids, position_ids)
end

@inline function get_position_emb(be::HGFRobertaEmbeddings, inputs_embeds, ::Nothing)
  position_ids = _arange(inputs_embeds, size(inputs_embeds)[end-1]) .+ be.position_embeddings.pad_idx
  get_position_emb(be, inputs_embeds, position_ids)
end

@inline get_position_emb(be::HGFRobertaEmbeddings, inputs_embeds, position_ids) = be.position_embeddings(position_ids)

@inline get_token_type_emb(be::HGFRobertaEmbeddings, inputs_embeds, ::Nothing) = get_token_type_emb(be, inputs_embeds, fill!(similar(inputs_embeds, Int, Base.tail(size(inputs_embeds))), one(Int)))
@inline get_token_type_emb(be::HGFRobertaEmbeddings, inputs_embeds, token_type_ids) = be.token_type_embeddings(token_type_ids)

function (be::HGFRobertaEmbeddings)(input,
                                 position_ids::Union{Nothing, AbstractArray{<:Integer}},
                                 token_type_ids::Union{Nothing, AbstractArray{<:Integer}})
  position_embeds = get_position_emb(be, input, position_ids)
  inputs_embeds = get_word_emb(be, input)
  token_type_embeds = get_token_type_emb(be, inputs_embeds, token_type_ids)
  return be(inputs_embeds, position_embeds, token_type_embeds)
end

function (be::HGFRobertaEmbeddings)(inputs_embeds::AbstractArray{T},
                                 position_embeds::AbstractArray{T},
                                 token_type_embeds::AbstractArray{T}) where T
  embeddings = (inputs_embeds .+ position_embeds) + token_type_embeds
  embeddings = be.LayerNorm(embeddings)
  return embeddings
end

(be::HGFRobertaEmbeddings)(input; position_ids=nothing, token_type_ids=nothing) = be(input, position_ids, token_type_ids)

# roberta model without prediction

abstract type HGFRobertaPreTrainedModel <: HGFBertPreTrainedModel end

struct HGFRobertaModel{
  E<:HGFRobertaEmbeddings,
  T<:HGFBertEncoder,
  P<:HGFBertPooler
} <: HGFRobertaPreTrainedModel
  embeddings::E
  encoder::T
  pooler::P
end

@functor HGFRobertaModel

(bm::HGFRobertaModel)(input; position_ids = nothing, token_type_ids = nothing,
                   attention_mask = nothing,
                   output_attentions = false,
                   output_hidden_states = false
                   ) = bm(input, position_ids, token_type_ids, attention_mask,
                          Val(output_attentions), Val(output_hidden_states))

function (bm::HGFRobertaModel)(
  input, position_ids, token_type_ids,
  attention_mask,
  _output_attentions::Val{output_attentions},
  _output_hidden_states::Val{output_hidden_states}
) where {output_attentions, output_hidden_states}

  embedding_output = bm.embeddings(input, position_ids, token_type_ids)
  attention_mask = create_attention_mask(embedding_output, attention_mask)

  encoder_outputs = bm.encoder(embedding_output, attention_mask,
                               _output_attentions, _output_hidden_states)

  sequence_output = encoder_outputs.last_hidden_state
  pooled_output = bm.pooler(sequence_output)
  return (
    last_hidden_state = sequence_output,
    pooler_output = pooled_output,
    hidden_states = encoder_outputs.hidden_states,
    attentions = encoder_outputs.attentions
  )
end

(bm::HGFRobertaModel)(input, encoder_hidden_states;
                   position_ids = nothing, token_type_ids = nothing,
                   attention_mask = nothing, encoder_attention_mask = nothing,
                   output_attentions = false,
                   output_hidden_states = false
                   ) = bm(input, encoder_hidden_states,
                          position_ids, token_type_ids,
                          attention_mask, encoder_attention_mask,
                          Val(output_attentions), Val(output_hidden_states))

function (bm::HGFRobertaModel)(
  input, encoder_hidden_states, position_ids, token_type_ids,
  attention_mask, encoder_attention_mask,
  _output_attentions::Val{output_attentions},
  _output_hidden_states::Val{output_hidden_states}
) where {output_attentions, output_hidden_states}

  embedding_output = bm.embeddings(input, position_ids, token_type_ids)
  attention_mask = create_causal_attention_mask(embedding_output, attention_mask)
  encoder_attention_mask = create_attention_mask(encoder_hidden_states, encoder_attention_mask)

  encoder_outputs = bm.encoder(embedding_output, attention_mask,
                               encoder_hidden_states, encoder_attention_mask,
                               _output_attentions, _output_hidden_states)

  sequence_output = encoder_outputs.last_hidden_state
  pooled_output = bm.pooler(sequence_output)

  return (
    last_hidden_state = sequence_output,
    pooler_output = pooled_output,
    hidden_states = encoder_outputs.hidden_states,
    attentions = encoder_outputs.attentions
  )
end

function HGFRobertaModel(config::HGFRobertaConfig)
  embeddings = HGFRobertaEmbeddings(config)
  encoder = HGFBertEncoder(HGFBertConfig, config)
  pooler = HGFBertPooler(HGFBertConfig, config)
  HGFRobertaModel(embeddings, encoder, pooler)
end

get_input_embedding(model::HGFRobertaModel) = model.embeddings.word_embeddings.weight

# lm head

struct HGFRobertaLMHead{
  L<:FakeTHLinear,
  N<:FakeTHLayerNorm,
  D<:FakeTHLinear,
  B<:AbstractArray
} <: THModule
  dense::L
  layer_norm::N
  decoder::D
  bias::B
end

@functor HGFRobertaLMHead

function (self::HGFRobertaLMHead)(x)
  x = self.dense(x)
  x = gelu.(x)
  x = self.layer_norm(x)
  x = self.decoder(x) .+ self.bias
  return x
end

function HGFRobertaLMHead(config::HGFRobertaConfig)
  dense = FakeTHLinear(HGFBertConfig, config,
                       config.hidden_size,
                       config.hidden_size)
  layer_norm = FakeTHLayerNorm(HGFBertConfig, config,
                               config.hidden_size;
                               eps=config.layer_norm_eps)
  decoder = FakeTHLinear(HGFBertConfig, config,
                         config.hidden_size,
                         config.vocab_size;
                         bias=false)
  bias = zeros(Float32, config.vocab_size)
  return HGFRobertaLMHead(dense, layer_norm, decoder, bias)
end

# classification head

struct HGFRobertaClassificationHead{
  D<:FakeTHLinear,
  O<:FakeTHLinear} <: THModule
  dense::D
  out_proj::O
end

@functor HGFRobertaClassificationHead

function (self::HGFRobertaClassificationHead)(features)
  x = features[:, 1, :]
  x = self.dense(x)
  x = tanh.(x)
  x = self.out_proj
end

function HGFRobertaClassificationHead(config::HGFRobertaConfig)
  dense = FakeTHLinear(HGFBertConfig, config,
                       config.hidden_size,
                       config.hidden_size)
  out_proj = FakeTHLinear(HGFBertConfig, config,
                          config.hidden_size,
                          config.num_labels)
  return HGFRobertaClassificationHead(dense, out_proj)
end

# roberta models for different task

# causal lm

struct HGFRobertaForCausalLM{
  B<:HGFRobertaModel,
  L<:HGFRobertaLMHead
} <: HGFRobertaPreTrainedModel
  roberta::B
  lm_head::L
end

@functor HGFRobertaForCausalLM

(self::HGFRobertaForCausalLM)(input;
                              position_ids = nothing, token_type_ids = nothing,
                              attention_mask = nothing,
                              output_attentions = false,
                              output_hidden_states = false
                              ) = self(input, position_ids, token_type_ids,
                                       attention_mask,
                                       Val(output_attentions), Val(output_hidden_states))

function (self::HGFRobertaForCausalLM)(input, position_ids, token_type_ids,
                                       attention_mask,
                                       _output_attentions::Val{output_attentions},
                                       _output_hidden_states::Val{output_hidden_states}
                                       ) where {output_attentions, output_hidden_states}
  outputs = self.roberta(input, position_ids, token_type_ids,
                         attention_mask, _output_attentions, _output_hidden_states)
  sequence_output = outputs[1]
  prediction_scores = self.lm_head(sequence_output)
  lm_loss = nothing

  return (
    loss = lm_loss,
    logits = prediction_scores,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

(self::HGFRobertaForCausalLM)(input, labels;
                              position_ids = nothing, token_type_ids = nothing,
                              attention_mask = nothing,
                              output_attentions = false,
                              output_hidden_states = false
                              ) = self(input, labels,
                                       position_ids, token_type_ids,
                                       attention_mask,
                                       Val(output_attentions), Val(output_hidden_states))

function (self::HGFRobertaForCausalLM)(input, labels,
                                    position_ids, token_type_ids,
                                       attention_mask,
                                       _output_attentions::Val{output_attentions},
                                       _output_hidden_states::Val{output_hidden_states}
                                       ) where {output_attentions, output_hidden_states}
  outputs = self(input, position_ids, token_type_ids, attention_mask,
                 _output_attentions, _output_hidden_states)

  prediction_scores = outputs.prediction_logits

  shifted_prediction_scores = prediction_scores[:, 1:end-1, :]
  shifted_labels = labels[:, 2:end]

  lm_loss = Flux.logitcrossentropy(shifted_prediction_scores, shifted_labels)

  return (
    loss = lm_loss,
    logits = prediction_scores,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

function HGFRobertaForCausalLM(config::HGFRobertaConfig)
  roberta = HGFRobertaModel(config)
  lm_head = HGFRobertaLMHead(config)
  return HGFRobertaForCausalLM(roberta, lm_head)
end

# masked lm

struct HGFRobertaForMaskedLM{
  B<:HGFRobertaModel,
  L<:HGFRobertaLMHead
} <: HGFRobertaPreTrainedModel
  roberta::B
  lm_head::L
end

@functor HGFRobertaForMaskedLM

(self::HGFRobertaForMaskedLM)(input;
                              position_ids = nothing, token_type_ids = nothing,
                              attention_mask = nothing,
                              output_attentions = false,
                              output_hidden_states = false
                              ) = self(input, position_ids, token_type_ids,
                                       attention_mask,
                                       Val(output_attentions), Val(output_hidden_states))

function (self::HGFRobertaForMaskedLM)(input, position_ids, token_type_ids,
                                       attention_mask,
                                       _output_attentions::Val{output_attentions},
                                       _output_hidden_states::Val{output_hidden_states}
                                       ) where {output_attentions, output_hidden_states}
  outputs = self.roberta(input, position_ids, token_type_ids,
                         attention_mask, _output_attentions, _output_hidden_states)
  sequence_output = outputs[1]
  prediction_scores = self.lm_head(sequence_output)
  lm_loss = nothing

  return (
    loss = lm_loss,
    logits = prediction_scores,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

(self::HGFRobertaForMaskedLM)(input, labels;
                              position_ids = nothing, token_type_ids = nothing,
                              attention_mask = nothing,
                              output_attentions = false,
                              output_hidden_states = false
                              ) = self(input, labels,
                                       position_ids, token_type_ids,
                                       attention_mask,
                                       Val(output_attentions), Val(output_hidden_states))

function (self::HGFRobertaForMaskedLM)(input, labels,
                                    position_ids, token_type_ids,
                                       attention_mask,
                                       _output_attentions::Val{output_attentions},
                                       _output_hidden_states::Val{output_hidden_states}
                                       ) where {output_attentions, output_hidden_states}
  outputs = self(input, position_ids, token_type_ids, attention_mask,
                 _output_attentions, _output_hidden_states)

  prediction_scores = outputs.prediction_logits

  lm_loss = Flux.logitcrossentropy(prediction_scores, labels)

  return (
    loss = lm_loss,
    logits = prediction_scores,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

function HGFRobertaForMaskedLM(config::HGFRobertaConfig)
  roberta = HGFRobertaModel(config)
  lm_head = HGFRobertaLMHead(config)
  return HGFRobertaForMaskedLM(roberta, lm_head)
end

# sequence classification

struct HGFRobertaForSequenceClassification{
  B<:HGFRobertaModel,
  C<:HGFRobertaClassificationHead
} <: HGFRobertaPreTrainedModel
  roberta::B
  classifier::C
end

@functor HGFRobertaForSequenceClassification

(self::HGFRobertaForSequenceClassification)(input;
                              position_ids = nothing, token_type_ids = nothing,
                              attention_mask = nothing,
                              output_attentions = false,
                              output_hidden_states = false
                              ) = self(input, position_ids, token_type_ids,
                                       attention_mask,
                                       Val(output_attentions), Val(output_hidden_states))

function (self::HGFRobertaForSequenceClassification)(input, position_ids, token_type_ids,
                                       attention_mask,
                                       _output_attentions::Val{output_attentions},
                                       _output_hidden_states::Val{output_hidden_states}
                                       ) where {output_attentions, output_hidden_states}
  outputs = self.roberta(input, position_ids, token_type_ids,
                         attention_mask, _output_attentions, _output_hidden_states)
  sequence_output = outputs[1]
  logits = self.classifier(sequence_output)
  loss = nothing

  return (
    loss = loss,
    logits = logits,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

(self::HGFRobertaForSequenceClassification)(input, labels;
                              position_ids = nothing, token_type_ids = nothing,
                              attention_mask = nothing,
                              output_attentions = false,
                              output_hidden_states = false
                              ) = self(input, labels,
                                       position_ids, token_type_ids,
                                       attention_mask,
                                       Val(output_attentions), Val(output_hidden_states))

function (self::HGFRobertaForSequenceClassification)(input, labels,
                                    position_ids, token_type_ids,
                                       attention_mask,
                                       _output_attentions::Val{output_attentions},
                                       _output_hidden_states::Val{output_hidden_states}
                                       ) where {output_attentions, output_hidden_states}
  outputs = self(input, position_ids, token_type_ids, attention_mask,
                 _output_attentions, _output_hidden_states)

  logits = outputs.logits
  if size(logits, 1) == 1
    loss = Flux.mse(logits, labels)
  else
    loss = Flux.logitcrossentropy(logits, labels)
  end

  return (
    loss = loss,
    logits = logits,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

function HGFRobertaForSequenceClassification(config::HGFRobertaConfig)
  roberta = HGFRobertaModel(config)
  classifier = HGFRobertaClassificationHead(config)
  return HGFRobertaForSequenceClassification(roberta, classifier)
end

# multiple chhoice

struct HGFRobertaForMultipleChoice{
  B<:HGFRobertaModel,
  C<:FakeTHLinear
} <: HGFRobertaPreTrainedModel
  roberta::B
  classifier::C
end

@functor HGFRobertaForMultipleChoice

(self::HGFRobertaForMultipleChoice)(input;
                              position_ids = nothing, token_type_ids = nothing,
                              attention_mask = nothing,
                              output_attentions = false,
                              output_hidden_states = false
                              ) = self(input, position_ids, token_type_ids,
                                       attention_mask,
                                       Val(output_attentions), Val(output_hidden_states))

function (self::HGFRobertaForMultipleChoice)(input, position_ids, token_type_ids,
                                    attention_mask,
                                    _output_attentions::Val{output_attentions},
                                    _output_hidden_states::Val{output_hidden_states}
                                          ) where {output_attentions, output_hidden_states}

  num_choices = size(input, ndims(input)-1)
  flat_choice(x) = reshape(x, size(x)[1:end-2]..., :)
  flat_choice(::Nothing) = nothing
  flat_input = flat_choice(input)
  flat_position_ids = flat_choice(position_ids)
  flat_token_type_ids = flat_choice(token_type_ids)
  flat_attention_mask = flat_choice(attention_mask)

  outputs = self.roberta(flat_input, flat_position_ids, flat_token_type_ids,
                      flat_attention_mask, _output_attentions, _output_hidden_states)

  pooled_output = outputs[2]
  logits = self.classifier(pooled_output)
  reshaped_logits = reshape(logits, num_choices, :)
  loss = nothing

  return (
    loss = loss,
    logits = reshaped_logits,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

(self::HGFRobertaForMultipleChoice)(input, labels;
                           position_ids = nothing, token_type_ids = nothing,
                           attention_mask = nothing,
                           output_attentions = false,
                           output_hidden_states = false
                           ) = self(input, labels,
                                    position_ids, token_type_ids,
                                    attention_mask,
                                    Val(output_attentions), Val(output_hidden_states))

function (self::HGFRobertaForMultipleChoice)(input, labels,
                                          position_ids, token_type_ids,
                                          attention_mask,
                                          _output_attentions::Val{output_attentions},
                                          _output_hidden_states::Val{output_hidden_states}
                                          ) where {output_attentions, output_hidden_states}
  outputs = self(input, position_ids, token_type_ids, attention_mask,
                 _output_attentions, _output_hidden_states)

  logits = outputs.logits
  loss = Flux.logitcrossentropy(logits, labels)

  return (
    loss = loss,
    logits = logits,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

function HGFRobertaForMultipleChoice(config::HGFRobertaConfig)
  roberta = HGFRobertaModel(config)
  classifier = FakeTHLinear(HGFBertConfig, config,
                            config.hidden_size, 1)
  return HGFRobertaForMultipleChoice(roberta, classifier)
end

# token classification

struct HGFRobertaForTokenClassification{
  B<:HGFRobertaModel,
  C<:FakeTHLinear
} <: HGFRobertaPreTrainedModel
  roberta::B
  classifier::C
end

@functor HGFRobertaForTokenClassification

(self::HGFRobertaForTokenClassification)(input;
                              position_ids = nothing, token_type_ids = nothing,
                              attention_mask = nothing,
                              output_attentions = false,
                              output_hidden_states = false
                              ) = self(input, position_ids, token_type_ids,
                                       attention_mask,
                                       Val(output_attentions), Val(output_hidden_states))

function (self::HGFRobertaForTokenClassification)(input, position_ids, token_type_ids,
                                    attention_mask,
                                    _output_attentions::Val{output_attentions},
                                    _output_hidden_states::Val{output_hidden_states}
                                    ) where {output_attentions, output_hidden_states}
  outputs = self.roberta(input, position_ids, token_type_ids,
                      attention_mask, _output_attentions, _output_hidden_states)
  sequence_output = outputs[1]
  logits = self.classifier(sequence_output)
  loss = nothing

  return (
    loss = loss,
    logits = logits,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

(self::HGFRobertaForTokenClassification)(input, labels;
                           position_ids = nothing, token_type_ids = nothing,
                           attention_mask = nothing,
                           output_attentions = false,
                           output_hidden_states = false
                           ) = self(input, labels,
                                    position_ids, token_type_ids,
                                    attention_mask,
                                    Val(output_attentions), Val(output_hidden_states))

function (self::HGFRobertaForTokenClassification)(input, labels,
                                    position_ids, token_type_ids,
                                       attention_mask,
                                       _output_attentions::Val{output_attentions},
                                       _output_hidden_states::Val{output_hidden_states}
                                       ) where {output_attentions, output_hidden_states}
  outputs = self(input, position_ids, token_type_ids, attention_mask,
                _output_attentions, _output_hidden_states)

  logits = outputs.logits
  loss = Flux.logitcrossentropy(logits, labels)

  return (
    loss = loss,
    logits = logits,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

function HGFRobertaForTokenClassification(config::HGFRobertaConfig)
  roberta = HGFRobertaModel(config)
  classifier = FakeTHLinear(HGFBertConfig, config,
                            config.hidden_size,
                            config.num_labels)
  return HGFRobertaForTokenClassification(roberta, classifier)
end

# question answering

struct HGFRobertaForQuestionAnswering{
  B<:HGFRobertaModel,
  C<:FakeTHLinear
} <: HGFRobertaPreTrainedModel
  roberta::B
  qa_outputs::C
end

@functor HGFRobertaForQuestionAnswering

(self::HGFRobertaForQuestionAnswering)(input;
                              position_ids = nothing, token_type_ids = nothing,
                              attention_mask = nothing,
                              output_attentions = false,
                              output_hidden_states = false
                              ) = self(input, position_ids, token_type_ids,
                                       attention_mask,
                                       Val(output_attentions), Val(output_hidden_states))

function (self::HGFRobertaForQuestionAnswering)(input, position_ids, token_type_ids,
                                    attention_mask,
                                    _output_attentions::Val{output_attentions},
                                    _output_hidden_states::Val{output_hidden_states}
                                    ) where {output_attentions, output_hidden_states}
  outputs = self.roberta(input, position_ids, token_type_ids,
                      attention_mask, _output_attentions, _output_hidden_states)
  sequence_output = outputs[1]
  logits = self.qa_outputs(sequence_output)
  start_logits = @view(logits[1, :])
  end_logits = @view(logits[2, :])
  total_loss = nothing

  return (
    loss = total_loss,
    start_logits = start_logits,
    end_logits = end_logits,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

(self::HGFRobertaForQuestionAnswering)(input, start_positions, end_positions;
                           position_ids = nothing, token_type_ids = nothing,
                           attention_mask = nothing,
                           output_attentions = false,
                           output_hidden_states = false
                           ) = self(input, start_positions, end_positions,
                                    position_ids, token_type_ids,
                                    attention_mask,
                                    Val(output_attentions), Val(output_hidden_states))

function (self::HGFRobertaForQuestionAnswering)(input, start_positions, end_positions,
                                    position_ids, token_type_ids,
                                       attention_mask,
                                       _output_attentions::Val{output_attentions},
                                       _output_hidden_states::Val{output_hidden_states}
                                       ) where {output_attentions, output_hidden_states}
  outputs = self(input, position_ids, token_type_ids, attention_mask,
                _output_attentions, _output_hidden_states)

  start_logits = outputs.start_logits
  end_logits = outputs.end_logits
  start_loss = Flux.logitcrossentropy(start_logits, start_positions)
  end_loss = Flux.logitcrossentropy(end_logits, end_positions)
  total_loss = (start_loss + end_loss) / 2

  return (
    loss = total_loss,
    start_logits = start_logits,
    end_logits = end_logits,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

function HGFRobertaForQuestionAnswering(config::HGFRobertaConfig)
  roberta = HGFRobertaModel(config)
  qa_outputs = FakeTHLinear(HGFBertConfig, config,
                            config.hidden_size,
                            config.num_labels)
  return HGFRobertaForQuestionAnswering(roberta, qa_outputs)
end

# load model utils

basemodelkey(::HGFRobertaPreTrainedModel) = :roberta
basemodel(m::HGFRobertaPreTrainedModel) = getproperty(m, basemodelkey(m))
basemodel(m::HGFRobertaModel) = m

isbasemodel(m::HGFRobertaModel) = true
isbasemodel(m::HGFRobertaPreTrainedModel) = false

get_model_type(::Val{:roberta}) = (
  :model => HGFRobertaModel,
  :forcausallm => HGFRobertaForCausalLM,
  :formaskedlm => HGFRobertaForMaskedLM,
  :forsequenceclassification => HGFRobertaForSequenceClassification,
  :formultiplechoice => HGFRobertaForMultipleChoice,
  :fortokenclassification => HGFRobertaForTokenClassification,
  :forquestionanswering => HGFRobertaForQuestionAnswering,
)


#=

Two-dimensional ergodic Hidden Markov Model

HMM2d([])

=#

struct HMM2d
end


module TestExamplesKPP

using Test
using Trixi

include("test_trixi.jl")

EXAMPLES_DIR = joinpath(examples_dir(), "tree_2d_dgsem")

@testset "KPP" begin
  @trixi_testset "elixir_kpp.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_kpp.jl"),
      l2   = [0.36563290910786106],
      linf = [9.116732052340398],
      max_refinement_level = 6,
      tspan = (0.0, 0.01),
      atol = 1e-6,
      rtol = 1e-6,
      skip_coverage = true)
  end
end

end # module


"""
`module  JAC.StrongField
    ... a submodel of JAC that contains all methods to set-up and perform strong-field computations.
"""
module StrongField

    using    GSL, ## HypergeometricFunctions, Plots,
             Printf, SpecialFunctions, ..AngularMomentum, ..Basics, ..Defaults, ..InteractionStrength, ..Radial, ..ManyElectron,
             ..Nuclear, ..Pulse, ..TableStrings

    export   aaaa

    """
    `abstract type StrongField.AbstractSFAObservable`
        ... defines an abstract and a number of types for the observables that can be computed with SFA amplitudes.

        + struct SfaNoObservable            ... an empty instance of an observable that does not help compute anything.
        + struct SfaEnergyDistribution      ... to compute the energy distribution of the photoelectrons.
        + struct SfaMomentumDistribution    ... to compute the momentum distribution of the photoelectrons.
    """
    abstract type  AbstractSFAObservable                                 end
    struct   SfaNoObservable     <: StrongField.AbstractSFAObservable    end

    """
    `struct  StrongField.SfaEnergyDistribution  <: StrongField.AbstractSFAObservable`
        ... to compute in SFA the energy distribution of photoelectrons at given energies and angles.

        + theta         ::Float64             ... polar angle of the energy distribution
        + phi           ::Float64             ... azimuthal angle of the energy distribution
        + energies      ::Array{Float64,1}    ... specifies the photoelectron energy in the current units.
    """
    struct   SfaEnergyDistribution  <: StrongField.AbstractSFAObservable
        theta           ::Float64
        phi             ::Float64
        energies        ::Array{Float64,1}
    end

    """
    `StrongField.SfaEnergyDistribution(theta::Float64, phi::Float64, NoEnergies::Int64, maxEnergy::Float64)`
        ... defines an energy distribution for given (theta,phi) and for NoEnergies between 0 < energy <= maxEnergy;
            a dist::SfaEnergyDistribution is returned.
    """
    function SfaEnergyDistribution(theta::Float64, phi::Float64, NoEnergies::Int64, maxEnergy::Float64)
        energies = Float64[];      for  i=1:NoEnergies   push!(energies, i*maxEnergy / NoEnergies)   end
        SfaEnergyDistribution(theta, phi, energies)
    end

    function Base.string(obs::SfaEnergyDistribution)
        sa = "Compute photo-electron energy distribution at angles (theta,phi) = ($(obs.theta), $(obs.phi)) for the energies [a.u.]:"
        return( sa )
    end

    function Base.show(io::IO, obs::SfaEnergyDistribution)
        sa = string(obs);       print(io, sa, "\n");     print(io, "   ", obs.energies)
    end

    """
    `struct  StrongField.SfaMomentumDistribution  <: StrongField.AbstractSFAObservable`
        ... to compute in SFA the momentum distribution of photoelectrons for given polar angle theta and at given energies and azimuthal angles.

        + theta         ::Float64             ... polar angle of the energy distribution
        + phi           ::Array{Float64,1}    ... specifies the azimuthal angles
        + energies      ::Array{Float64,1}    ... specifies the photoelectron energy in the current units.
    """
    struct   SfaMomentumDistribution  <: StrongField.AbstractSFAObservable
        theta           ::Float64
        phis            ::Array{Float64,1}
        energies        ::Array{Float64,1}
    end

    """
    `StrongField.SfaMomentumDistribution(theta::Float64, NoPhi::Int64, NoEnergies::Int64, maxEnergy::Float64)`
        ... defines a momentum distribution for given theta and for NoEnergies between 0 < energy <= maxEnergy and NoPhi between 0 <= phi < 2pi;
            a dist::SfaMomentumDistribution is returned.
    """
    function SfaMomentumDistribution(theta::Float64, NoPhi::Int64, NoEnergies::Int64, maxEnergy::Float64)
        energies = Float64[];      for  i=1:NoEnergies   push!(energies, i*maxEnergy / NoEnergies)   end
        phis = Float64[];          for  i=1:NoPhi        push!(phis, i*2pi / NoPhi)                  end
        SfaMomentumDistribution(theta, phis, energies)
    end

    function Base.string(obs::SfaMomentumDistribution)
        sa = "Compute photo-electron momentum distribution at polar angle theta = $(obs.theta) for (energy [a.u.], phi):"
        return( sa )
    end

    function Base.show(io::IO, obs::SfaMomentumDistribution)
        sa = string(obs);       print(io, sa, "\n");     print(io, "   ", obs.energies);    print(io, "   ", obs.phi)
    end

    """
    `abstract type StrongField.AbstractVolkovState`
        ... defines an abstract and a number of types for dealing with the Volkov states in the computation of the SFA amplitudes.

        + struct FreeVolkov       ... to apply the free-Volkov states.
        + struct CoulombVolkov    ... apply the Coulomb-Volkov states.
        + struct DistortedVolkov  ... apply the Distorted-Volkov states.
    """
    abstract type  AbstractVolkovState                         end
    struct         FreeVolkov        <:  AbstractVolkovState   end

    """
    `struct StrongField.CoulombVolkov     <:  StrongField.AbstractVolkovState`

        + Z ::Float64   ... Charge that generates the Coulomb potential
    """
    struct  CoulombVolkov   <: StrongField.AbstractVolkovState
        Z   ::Float64
    end

    struct         DistortedVolkov   <:  AbstractVolkovState   end

    function Base.string(volkov::FreeVolkov)         return( "free-Volkov states" )         end
    function Base.string(volkov::CoulombVolkov)      return( "Coulomb-Volkov states in potential of charge Z = $(volkov.Z)" )      end
    function Base.string(volkov::DistortedVolkov)    return( "distorted-Volkov states" )    end

    """
    `struct  StrongField.Settings`
        ... defines a type for the details and parameters of computing SFA amplitudes.

        + multipoles            ::Array{EmMultipoles}  ... Multipoles of the radiation field that are to be included.
        + gauges                ::Array{UseGauge}      ... Specifies the gauges to be included into the computations.
        + printBefore           ::Bool                 ... True, if all the requested amplitudes are printed before the computation start.
        + printAmplitudes       ::Bool                 ... True, if the amplitudes are to be printed.
    """
    struct Settings
        multipoles              ::Array{EmMultipole,1}
        gauges                  ::Array{UseGauge}
        printBefore             ::Bool
        printAmplitudes         ::Bool
    end

    """
    `StrongField.Settings()`
        ... constructor for the default values of SFA computations.
    """
    function Settings()
        Settings([E1], [UseCoulomb], false, false)
    end

    # `Base.show(io::IO, settings::StrongField.Settings)`  ... prepares a proper printout of the variable settings::StrongField.Settings.
    function Base.show(io::IO, settings::StrongField.Settings)
        println(io, "multipoles:                 $(settings.multipoles)  ")
        println(io, "use-gauges:                 $(settings.gauges)  ")
        println(io, "printBefore:                $(settings.printBefore)  ")
        println(io, "printAmplitudes:            $(settings.printAmplitudes)  ")
    end

    """
    `struct  StrongField.Computation`
        ... defines a type for the computation of strong-field amplitudes and observables (properties).

        + observable         ::AbstractSFAObservable        ... The obserable to be calculated in SFA.
        + nuclearModel       ::Nuclear.Model                ... Model, charge and parameters of the nucleus.
        + grid               ::Radial.Grid                  ... The radial grid to be used for the computation.
        + initialLevel       ::Level                        ... Initial level of the atom
        + finalLevel         ::Level                        ... Final level of the atom
        + beam               ::Pulse.AbstractBeam           ... Beam character and properties of the incident light pulse.
        + envelope           ::Pulse.AbstractEnvelope       ... Envelope of the incident light pulse.
        + polarization       ::Basics.AbstractPolarization  ... Envelope of the incident light pulse.
        + volkov             ::AbstractVolkovState          ... Specify the treatment/approach of the Volkov states.
        + settings           ::StrongField.Settings         ... Settings for the controlling the SFA amplitudes.
    """
    struct  Computation
        observable           ::AbstractSFAObservable
        nuclearModel         ::Nuclear.Model
        grid                 ::Radial.Grid
        initialLevel         ::Level
        finalLevel           ::Level
        beam                 ::Pulse.AbstractBeam
        envelope             ::Pulse.AbstractEnvelope
        polarization         ::Basics.AbstractPolarization
        volkov               ::AbstractVolkovState
        settings             ::StrongField.Settings
    end

    """
    `StrongField.Computation()`  ... constructor for an `empty` instance of StrongField.Computation().
    """
    function Computation()
        Computation( SfaNoObservable(), Nuclear.Model(1.0), Radial.Grid(true), Level(), Level(),
                     Pulse.PlaneWaveBeam(), Pulse.InfiniteEnvelope(), Basics.RightCircular(), FreeVolkov(), Settings()  )
    end

    # `Base.string(comp::StrongField.Computation)`  ... provides a String notation for the variable comp::StrongField.Computation.
    function Base.string(comp::StrongField.Computation)
        sa = "Strong-field computation:  " * string(comp.observable) * " for Z = $(comp.nuclearModel.Z)  with " * string(comp.volkov) * "\n"
        sa = sa * " initial level with (J, P, energy) = ($(comp.initialLevel.J), $(comp.initialLevel.parity), $(comp.initialLevel.energy)) \n"
        sa = sa * " final level with   (J, P, energy) = ($(comp.finalLevel.J), $(comp.finalLevel.parity), $(comp.finalLevel.energy))   \n\n"
        sa = sa * "The incident laser pulse is described by the: \n "
        sa = sa * string(comp.beam) * "\n " * string(comp.polarization) * "\n " * string(comp.envelope)
        return( sa )
    end

    # `Base.show(io::IO, comp::StrongField.Computation)`  ... prepares a proper printout of the variable comp::StrongField.Computation.
    function Base.show(io::IO, comp::StrongField.Computation)
        sa = Base.string(comp);             print(io, sa, "\n\n")
        println(io, "Settings:              \n$(comp.settings)    ")
        println(io, "nuclearModel:          $(comp.nuclearModel)  ")
        println(io, "grid:                  $(comp.grid)  ")
    end

    """
    `struct  StrongField.SphericalAmplitude`
        ... to keep the amplitude at a given  energy-angular point (energy, theta, phi) in momentum space.

        + energy        ::Float64      ... Kinetic energy of the (outgoing) photoelectron.
        + theta         ::Float64      ... Polar angle of the (outgoing) photoelectron.
        + phi           ::Float64      ... Azimuthal angle of the (outgoing) photoelectron.
        + value         ::ComplexF64   ... (Total) Amplitude.
    """
    struct   SphericalAmplitude
        energy          ::Float64
        theta           ::Float64
        phi             ::Float64
        value           ::ComplexF64
    end

    function Base.string(amp::StrongField.SphericalAmplitude)
        sa = "Total SFA amplitude at (energy, theta, phi) = ($(amp.energy),$(amp.theta),$(amp.phi)) is $(amp.value)."
        return( sa )
    end

    function Base.show(io::IO, amp::StrongField.SphericalAmplitude)
        sa = string(amp);       print(io, sa)
    end

    """
    `struct  StrongField.OutcomeEnergyDistribution`
        ... to comprise the energy distribution of photoelectrons at given energies and angles, for instance,
            for later graphical representation.

        + theta         ::Float64             ... polar angle of the energy distribution
        + phi           ::Float64             ... azimuthal angle of the energy distribution
        + energies      ::Array{Float64,1}    ... selected energies of the distribution.
        + probabilities ::Array{Float64,1}    ... calculated probabilities of the energy distribution.
    """
    struct   OutcomeEnergyDistribution
        theta           ::Float64
        phi             ::Float64
        energies        ::Array{Float64,1}
        probabilities   ::Array{Float64,1}
    end

    # `Base.show(io::IO, outcome::StrongField.OutcomeEnergyDistribution)`  ... prepares a proper printout of the variable outcome::StrongField.OutcomeEnergyDistribution.
    function Base.show(io::IO, outcome::StrongField.OutcomeEnergyDistribution)
        println(io, "theta:             $(outcome.theta)  ")
        println(io, "phi:               $(outcome.phi)  ")
        println(io, "energies:          $(outcome.energies)  ")
        println(io, "probabilities:     $(outcome.probabilities)  ")
    end

    """
    `struct  StrongField.OutcomeMomentumDistribution`
        ... to comprise the momentum distribution of photoelectrons at given energies and angles, for instance,
            for later graphical representation.

        + theta         ::Float64             ... polar angle of the momentum distribution
        + momenta       ::Array{Float64}      ... pairs [px py] of momenta of the momentum distribution
        + probabilities ::Array{Float64,1}    ... calculated probabilities of the energy distribution.
    """
    struct   OutcomeMomentumDistribution
        theta           ::Float64
        momenta         ::Array{Float64}
        probabilities   ::Array{Float64,1}
    end

    # `Base.show(io::IO, outcome::StrongField.OutcomeMomentumDistribution)`  ... prepares a proper printout of the variable outcome::StrongField.OutcomeMomentumDistribution.
    function Base.show(io::IO, outcome::StrongField.OutcomeMomentumDistribution)
        println(io, "theta:             $(outcome.theta)  ")
        println(io, "momenta:           $(outcome.momenta)  ")
        println(io, "probabilities:     $(outcome.probabilities)  ")
    end

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

    ####################################################################################################################
    ######The following four functions will be replaced later and serve only for comparison with previous results#######
    ####################################################################################################################

    """
    `StrongField.HydrogenPnl(epsiloni::Float64, n::Int, l::Int, rGrid::Array{Float64,1})`
        ... returns the non-relativistic hydrogen-like radial wave function with modified binding energy epsiloni
                at the radial grid points rGrid
    """
    function  HydrogenPnl(epsiloni::Float64, n::Int, l::Int, rGrid::Array{Float64,1})
        Pnl = Float64[]
        Z = n*sqrt(-2*epsiloni)

        for j = 1:length(rGrid)
            r = rGrid[j]
            p = r * sqrt( (2*Z/n)^3 * factorial(n-l-1)/(2*n*factorial(n+l)) ) * exp(-Z*r/n) * (2*Z*r/n)^l * GSL.sf_laguerre_n(n-l-1,2*l+1,2*Z*r/n)
            push!( Pnl, p )
        end

        return( Pnl )
    end

    """
    `StrongField.HydrogenDPnlDr(epsiloni::Float64, n::Int, l::Int, rGrid::Array{Float64,1})`
        ... returns the r-derivative of the non-relativistic hydrogen-like radial wave function with modified binding energy epsiloni
                at the radial grid points rGrid
    """
    function  HydrogenDPnlDr(epsiloni::Float64, n::Int, l::Int, rGrid::Array{Float64,1})
        DPnl = Float64[]
        Z = n*sqrt(-2*epsiloni)

        for j = 1:length(rGrid)
            r = rGrid[j]
            if  n-l-2 >= 0
                p = 1/n * 2^(l+1) * exp(-r*Z/n) * (r*Z/n)^l * sqrt( Z^3 * factorial(n-l-1) / (n^4*factorial(l+n)) ) * ( (n+l*n-r*Z) * GSL.sf_laguerre_n(n-l-1,2*l+1,2*r*Z/n) - 2*r*Z * GSL.sf_laguerre_n(n-l-2,2+2*l,2*r*Z/n) )
            elseif  n == 1 && l == 0
                p = -2 * exp(-r*Z) * sqrt(Z^3) * (r*Z - 1)
            else #not implemented
                p = 0
            end
            push!( DPnl, p )
        end

        return( DPnl )
    end

    """
    `StrongField.VolkovP(epsilonp::Float64, lp::Int, rGrid::Array{Float64,1})`
        ... returns the (plane-wave) Volkov radial wave function
                at the radial grid points rGrid
    """
    function  VolkovP(epsilonp::Float64, lp::Int, rGrid::Array{Float64,1})
        P = Float64[]

        for j = 1:length(rGrid)
            r = rGrid[j]
            p = r * GSL.sf_bessel_jl( lp, sqrt(2*epsilonp)*r )
            push!( P, p )
        end

        return( P )
    end

    """
    `StrongField.CoulombVolkovP(epsilonp::Float64, lp::Int, ZContinuum::Float64, rGrid::Array{Float64,1})`
        ... returns the Coulomb-Volkov radial wave function
                at the radial grid points rGrid
    """
    function  CoulombVolkovP(epsilonp::Float64, lp::Int, ZContinuum::Float64, rGrid::Array{Float64,1})
        sqrtTwoEps = sqrt(2*epsilonp)

        P = ComplexF64[]

        for j = 1:length(rGrid)
            r = rGrid[j]
            rho = sqrtTwoEps*r
            eta = -ZContinuum/sqrtTwoEps
            GammaValue = gamma( lp + 1 + eta * im )
            sigma = angle( GammaValue )
            Fl = 2.0^lp * exp(-0.5*pi*eta) * abs( GammaValue ) / factorial(2*lp+1) * rho^(lp+1) * exp( rho * im ) * HypergeometricFunctions.drummond1F1(lp+1+eta*im,2*lp+2,-2*rho*im)

            p = Fl * exp( sigma * im ) / sqrtTwoEps

            push!( P, p )
        end

        return( P )
    end

    """
    `StrongField.pReducedME(epsilonp::Float64, lp::Int, n::Int, l::Int, epsiloni::Float64, volkov::AbstractVolkovState)`
        ... computes the reduced matrix elements of the momentum operator <epsilonp lp ||p||n l> in the one-particle picture
    """
    function  pReducedME(epsilonp::Float64, lp::Int, n::Int, l::Int, epsiloni::Float64, volkov::AbstractVolkovState)
        rmax = 100.
        orderGL = 1000
        gaussLegendre = Radial.GridGL("Finite",0.0,rmax,orderGL)
        rgrid = gaussLegendre.t
        weights = gaussLegendre.wt

        Pnl = HydrogenPnl( epsiloni, n, l, rgrid )
        if  typeof(volkov) == FreeVolkov
            Pepsplp = VolkovP( epsilonp, lp, rgrid )
        elseif  typeof(volkov) == CoulombVolkov
            Pepsplp = CoulombVolkovP( epsilonp, lp, volkov.Z, rgrid )
        end

        DPnl = HydrogenDPnlDr( epsiloni, n, l, rgrid )

        integral = 0. * im

        #Sum over grid and compute Gauss-Legendre sum
        for    j = 1:orderGL
            r = rgrid[j]
            integrand = conj( Pepsplp[j] )/r * ( r*DPnl[j] - ((lp-l)*(lp+l+1))/2 * Pnl[j] )

            #Gauss-Legendre sum
            integral = integral + weights[j] * integrand
        end

        #Note that GSL.sf_coupling_3j takes takes the input (2*j1,2*j2,2*j3,2*m1,2*m2,2*m3)
        integral = integral * (-im)^(lp+1) * (-1)^lp * GSL.sf_coupling_3j( 2*lp, 2*1, 2*l, 0, 0, 0 )

        return( integral )
    end

    """
    `StrongField.scalarProdBoundCont(epsilonp::Float64, lp::Int, n::Int, l::Int, epsiloni::Float64, volkov::AbstractVolkovState)`
        ... computes the scalar product of the bound (hydrogenic) and continuum (Volkov) states in the one-particle picture
    """
    function  scalarProdBoundCont(epsilonp::Float64, lp::Int, n::Int, l::Int, m::Int, epsiloni::Float64, volkov::AbstractVolkovState)
        rmax = 1000.
        orderGL = 1000
        gaussLegendre = Radial.GridGL("Finite",0.0,rmax,orderGL)
        rgrid = gaussLegendre.t
        weights = gaussLegendre.wt

        Pnl = HydrogenPnl( epsiloni, n, l, rgrid )

        if  typeof(volkov) == FreeVolkov
            Pepsplp = VolkovP( epsilonp, lp, rgrid )
        elseif  typeof(volkov) == CoulombVolkov
            Pepsplp = CoulombVolkovP( epsilonp, lp, volkov.Z, rgrid )
        end

        DPnl = HydrogenDPnlDr( epsiloni, n, l, rgrid )

        integral = 0. * im

        #Sum over grid and compute Gauss-Legendre sum
        for    j = 1:orderGL
            r = rgrid[j]
            integrand = conj( Pepsplp[j] ) * Pnl[j]

            #Gauss-Legendre sum
            integral = integral + weights[j] * integrand
        end

        return((-im)^l * integral)
    end

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

    """
    `StrongField.computeEnvelopeIntegrals(envelope::Pulse.AbstractEnvelope, beam::AbstractBeam, polarization::Basics.AbstractPolarization,
                                          thetap::Float64, phip::Float64, energyp::Float64, initialEn::Float64)`
        ... computes the pulse-envelope integrals for the given laser pulse; this pulse is completely specified by its beam character
            and parameters, the pulse-envelope parameters and the polarization of the light. A tuple of the two integrals
            (fVolkov, fVolkovSquared)::Tuple{Complex{Float64},Complex{Float64}} is returned.
    """
    function computeEnvelopeIntegrals(envelope::Pulse.AbstractEnvelope, beam::AbstractBeam, polarization::Basics.AbstractPolarization,
                                      thetap::Float64, phip::Float64, energyp::Float64, initialEn::Float64)
        fVolkovPlus    = Pulse.envelopeVolkovIntegral(true,  envelope, beam, polarization, thetap::Float64, phip::Float64, energyp::Float64, initialEn::Float64)
        fVolkovMinus   = Pulse.envelopeVolkovIntegral(false, envelope, beam, polarization, thetap::Float64, phip::Float64, energyp::Float64, initialEn::Float64)
        fVolkovSquared = Pulse.envelopeQuadVolkovIntegral(envelope, beam, polarization,    thetap::Float64, phip::Float64, energyp::Float64, initialEn::Float64)

        #println("> Envelope integrals for an $(string(envelope)) with A0=$(beam.A0), omega=$(beam.omega) [a.u.], cep=$(beam.cep)" *
        #        " and for $(string(polarization)) light are:" )
        #println("    F^(Volkov) [+; omega; f^(env); A]   = $fVolkovPlus" )
        #println("    F^(Volkov) [-; omega; f^(env); A]   = $fVolkovMinus" )
        #println("    F^(quad, Volkov) [f^(env); A]       = $fVolkovSquared" )

        return( (fVolkovPlus, fVolkovMinus, fVolkovSquared) )
    end

    """
    `StrongField.computeOutcome(observable::StrongField.SfaEnergyDistribution, amplitudes::Array{SphericalAmplitude,1})`
        ... computes the requested photoelectron energy distribution and returns the results in the variable
            outcome::StrongField.OutcomeEnergyDistribution
    """
    function  computeOutcome(obs::StrongField.SfaEnergyDistribution, amplitudes::Array{SphericalAmplitude,1})
        probabilities = Float64[]
        for  amp  in  amplitudes   pp = sqrt(2*amp.energy);  push!(probabilities, pp * (amp.value * conj(amp.value)) )      end
        outcome = OutcomeEnergyDistribution(obs.theta, obs.phi, obs.energies, probabilities)
        return( outcome )
    end

    """
    `StrongField.computeOutcome(observable::StrongField.SfaMomentumDistribution, amplitudes::Array{SphericalAmplitude,1})`
        ... computes the requested photoelectron momentum distribution and returns the results in the variable
            outcome::StrongField.OutcomeMomentumDistribution
    """
    function  computeOutcome(obs::StrongField.SfaMomentumDistribution, amplitudes::Array{SphericalAmplitude,1})
        probabilities = Float64[]
        momenta       = Array{Float64}(undef,0,2)
        for  amp  in  amplitudes
            pp = sqrt(2*amp.energy)
            px = pp*sin(amp.theta)*cos(amp.phi)
            py = pp*sin(amp.theta)*sin(amp.phi)
            momenta = [momenta ; [px py]]
            push!( probabilities, pp * (amp.value * conj(amp.value)) )
        end
        outcome = OutcomeMomentumDistribution(obs.theta, momenta, probabilities)
        return( outcome )
    end

    """
    `StrongField.computeSphericalAmplitudes(comp::StrongField.Computation)`
        ... computes all necessary spherical amplitudes for the given initial and final levels, the Volkov states
            and polarization and by using the pulse envelope integrals. A newAmplitudes::Array{SphericalAmplitude,1} is returned.
    """
    function  computeSphericalAmplitudes(comp::StrongField.Computation)
        newAmplitudes = SphericalAmplitude[]
        #
        # First determine which spherical SFA amplitudes need to be computed before the actual computation starts
        sfaAmplitudes = StrongField.determineSphericalAmplitudes(comp.observable)
        # Determine quantum numbers of the initial and final state
        nqn = 1;    lqn = 0;    n = 1;      l = 0;      m = 0;    initialEn = convertUnits("energy: from eV to atomic", comp.initialLevel.energy)
        #

        lpMin = abs(l-1)
        lpMax = l+1

        # Compute the requested amplitudes
        for  amp in  sfaAmplitudes
            thetap        = amp.theta;   phip = amp.phi;     energyp = amp.energy
            envIntegrals  = StrongField.computeEnvelopeIntegrals(comp.envelope, comp.beam, comp.polarization, thetap, phip, energyp, initialEn)
            fVolkovPlus, fVolkovMinus, fVolkovSquared = envIntegrals
            #
            reducedMEArray = ComplexF64[]
            for lp = lpMin:lpMax
                push!( reducedMEArray, pReducedME(energyp, lp, n, l, initialEn, comp.volkov) )
            end
            #

            wminus = 0.0im;    wplus = 0.0im
            # Collect contributions from all l_p, q terms; this summation will change in a (nljm) representation
            for  lp = lpMin:lpMax
                reducedME = reducedMEArray[lp-lpMin+1]
                for  q in [-1, 0, 1]
                    wplus = wplus + AngularMomentum.sphericalYlm(lp, m-q, amp.theta, amp.phi) * (-1)^q  *
                             Basics.determinePolarizationVector(q, comp.polarization, star=false)         *
                             AngularMomentum.ClebschGordan(l, m, 1, -q, lp, m-q) * reducedME
                    wminus = wminus + AngularMomentum.sphericalYlm(lp, m+q, amp.theta, amp.phi)           *
                             Basics.determinePolarizationVector(q, comp.polarization, star=true)          *
                             AngularMomentum.ClebschGordan(l, m, 1, q, lp, m+q) * reducedME
                end
            end
            scalarProd = scalarProdBoundCont(energyp, l, n, l, m, initialEn, comp.volkov)
            wa = -im * sqrt(2/pi) * (fVolkovPlus * wminus + fVolkovMinus * wplus) - im / sqrt(2*pi) * fVolkovSquared * AngularMomentum.sphericalYlm(l, m, amp.theta, amp.phi) * scalarProd
            push!(newAmplitudes, SphericalAmplitude(amp.energy, amp.theta, amp.phi, wa))

            println(">> $(SphericalAmplitude(amp.energy, amp.theta, amp.phi, wa))")
        end

        return( newAmplitudes )
    end

    """
    `StrongField.determineSphericalAmplitudes(observable::StrongField.SfaEnergyDistribution)`
        ... determines which direct (and other) SFA amplitudes need to be computed; these amplitudes are not yet computed here but can be arranged
            so that only a minimum number of Volkov states and/or reduced many-electron matrix elements need to be computed.
            A list of amplitudes::Array{StrongField.SphericalAmplitude,1} is returned.
    """
    function  determineSphericalAmplitudes(observable::StrongField.SfaEnergyDistribution)
        amplitudes = StrongField.SphericalAmplitude[]
        for  energy in observable.energies      push!(amplitudes, SphericalAmplitude(energy, observable.theta, observable.phi, 0.))     end

        println("> A total of $(length(amplitudes)) spherical amplitudes need to be calculated.")
        return( amplitudes )
    end

    """
    `StrongField.determineSphericalAmplitudes(observable::StrongField.SfaMomentumDistribution)`
        ... determines which direct (and other) SFA amplitudes need to be computed; these amplitudes are not yet computed here but can be arranged
            so that only a minimum number of Volkov states and/or reduced many-electron matrix elements need to be computed.
            A list of amplitudes::Array{StrongField.SphericalAmplitude,1} is returned.
    """
    function  determineSphericalAmplitudes(observable::StrongField.SfaMomentumDistribution)
        amplitudes = StrongField.SphericalAmplitude[]
        for  energy in observable.energies
            for  phi in observable.phis
                push!(amplitudes, SphericalAmplitude(energy, observable.theta, phi, 0.))
            end
        end

        println("> A total of $(length(amplitudes)) spherical amplitudes need to be calculated.")
        return( amplitudes )
    end

    """
    `StrongField.plot( comp::StrongField.Computation, results::Dict{String,Any}, energyScale::String = "atomic", probabilityScaling::String = "linear" )`
        ... generates a graphical representation of the observable (either StrongField.SfaEnergyDistribution or StrongField.SfaMomentumDistribution).
        energyScale determines the scaling of the energy axis either in atomic units (energyScale = "atomic") or in units of hbar*omega (energyScale = "omega").
        The y-axis is scaled either linearly (probabilityScaling = "linear") or logarithmically (probabilityScaling = "log")
    """
    function plot( comp::StrongField.Computation, results::Dict{String,Any}, energyScale::String = "atomic", probabilityScaling::String = "linear" )

        #---Photoelectron energy spectrum---
        if  typeof(comp.observable) == StrongField.SfaEnergyDistribution
                #prepare data and rescale the x-axis if neccessary
                if  energyScale == "atomic"
                    energyLabel = "E (a.u.)"
                    energies = results["energy distribution"].energies
                elseif  energyScale == "omega"
                    energyLabel = "E/omega"
                    energies = results["energy distribution"].energies / comp.beam.omega
                end
                probabilities = results["energy distribution"].probabilities
                gr() #sets the plotting backend to the package "GR"

                #set scaling of the y-axis
                scaling = :identity
                if  probabilityScaling == "log"
                            scaling = :log10
                end

                #generate the plot
                p = Plots.plot(energies, probabilities,
                                title = "Photoelectron energy spectrum",
                                xlabel = energyLabel,
                                ylabel = "P(E)",
                                xscale = :identity,
                                yscale = scaling,
                                framestyle = :box,
                                legend = :none,
                                #markershape = :circle,
                                #markershape = :none,
                                line = 2,
                                linecolor = :black,
                                gridlinewidth = 2,
                                tickfontsize = 10,
                                labelfontsize = 10,
                                labelfontfamily = "Latin Modern Roman",
                                titlefontfamily = "Latin Modern Roman"

                              )

                #export the plot as png-file
                png("energy_spectrum")
        #-----------------------------------

        #---Photoelectron momentum distribution---
        elseif  typeof(comp.observable) == StrongField.SfaMomentumDistribution
                #prepare data
                pList = results["momentum distribution"].momenta
                probList = results["momentum distribution"].probabilities

                #generate the plot
                #pyplot()
                #r = range(0,stop=10,length=11)
                #theta = range(0,stop=360,length=361)
                #f(r,theta) = r^2
                #println("$(f.(r,theta'))")
                #Plots.plot( heatmap( f.(r,theta'), proj=:polar ) )
                #println("$(transpose(probList))")
                #Plots.plot( heatmap( pList, transpose(probList), proj=:polar ) )
                #imshow(pList,transpose(probList))
                #probList = probList[2:-1]
                #matplotlib.pyplot.pcolormesh(pList[:,1], pList[:,2], transpose(probList))
                contourf( pList[:,1], pList[:,2], probList )
                #matplotlib.pyplot.imshow(pList)
                #heatmap( pList[:,1], pList[:,2], probList, proj=:polar, legend=true
                                #heatmap( probList
                #                )

                #p = Plots.plot( #heatmap( pList[:,1], pList[:,2], probList
                #               Plots.GR.polarheatmap( probList
                            ##heatmap( probList
                #                 )
                #              )

                #export the plot as png-file
                png("momentum_distribution")
        #-----------------------------------------

        #---Not a valid obserable---
        else
                error("Undefined observable for strong-field computations.")
        end
        #---------------------------

    end

    """
    `StrongField.perform(comp::StrongField.Computation; output::Bool=false)`
        ... to perform a computation of (one) selected observable that is related to SFA, such as a photoelectron energy
            distribution, momentum distribution or the computation of sidebands.
    """
    function perform(comp::StrongField.Computation; output::Bool=false)
        if  output    results = Dict{String, Any}()    else    results = nothing    end
        nModel = comp.nuclearModel

        if       typeof(comp.observable) == StrongField.SfaEnergyDistribution
            sfaAmplitudes = StrongField.computeSphericalAmplitudes(comp)
            sfaOutcome    = StrongField.computeOutcome(comp.observable, sfaAmplitudes)
            if output    results = Base.merge( results, Dict("computation" => comp, "energy distribution" => sfaOutcome) )  end
        elseif   typeof(comp.observable) == StrongField.SfaMomentumDistribution
            sfaAmplitudes = StrongField.computeSphericalAmplitudes(comp)
            sfaOutcome    = StrongField.computeOutcome(comp.observable, sfaAmplitudes)
            if output    results = Base.merge( results, Dict("computation" => comp, "momentum distribution" => sfaOutcome) )  end
        else     error("Undefined observable for strong-field computations.")
        end

        if  output   return( results )   else    return( nothing )   end
    end

end # module


#!/usr/bin/env julia
# -*- coding: UTF-8 -*-
# __julia-version__ = 1.7.0
# __author__        = "Max Kannenberg, Martin Scheidt"
# __copyright__     = "2021"
# __license__       = "ISC"

using TrainRuns, Test

trains = Dict()
paths = Dict()
settings = Dict()

@testset "load data" begin

  println("testing load train data")
  push!(trains, :freight      => @time Train("data/trains/freight.yaml"))
  push!(trains, :local        => @time Train("data/trains/local.yaml"))
  push!(trains, :longdistance => @time Train("data/trains/longdistance.yaml"))

  println("testing load path data")
  push!(paths, :const     => @time Path("data/paths/const.yaml"))
  push!(paths, :slope     => @time Path("data/paths/slope.yaml"))
  push!(paths, :speed     => @time Path("data/paths/speed.yaml"))
  push!(paths, :realworld => @time Path("data/paths/realworld.yaml"))

  println("testing load settings data")
  push!(settings, "default"       => @time Settings())
  push!(settings, "poi"           => @time Settings("data/settings/points_of_interest.yaml"))
  push!(settings, "drivingcourse" => @time Settings("data/settings/driving_course.yaml"))
  push!(settings, "everything"    => @time Settings("data/settings/everything.yaml"))
  push!(settings, "strip"         => @time Settings("data/settings/strip.yaml"))
  push!(settings, "time"          => @time Settings("data/settings/time.yaml"))
  push!(settings, "timestrip"     => @time Settings("data/settings/time_strip.yaml"))
  push!(settings, "velocity"      => @time Settings("data/settings/velocity.yaml"))

  @test typeof(first(paths)[2]) == Path
  @test typeof(first(settings)[2]) == Settings

end

println("====================")

tests = Base.Iterators.product(trains,paths)

## routine to generate the anticipated Dict()
# anticipated = Dict()
# for test in tests
#   println(test[1][1],"-",test[2][1])
#   result = @time trainrun(test[1][2],test[2][2])
#   push!(anticipated, Symbol(String(test[1][1]) * "_" * String(test[2][1])) => result )
# end

anticipated = Dict(
  :default => Dict(
    :longdistance_speed => 499.96109564970516,
    :freight_slope => 831.4768274141168,
    :local_slope => 396.99313307033276,
    :longdistance_const => 328.83479381353095,
    :freight_realworld => 8971.50124080998,
    :longdistance_slope => 329.22915822053164,
    :freight_const => 727.7969403041934,
    :longdistance_realworld => 2900.1198723158523,
    :local_speed => 524.3948201513945,
    :local_realworld => 3443.917823618831,
    :freight_speed => 733.2610572579886,
    :local_const => 392.7234008268302
  )
)

@testset "function trainrun()" begin

  @testset "Default settings" begin

    for test in tests
      test_name = String(test[1][1]) * "_" * String(test[2][1])
      println("testing $test_name")
      @time result = trainrun(test[1][2],test[2][2])
      expected = anticipated[:default][Symbol(test_name)]
      # compare result to test data set
      @test isapprox(result, expected, rtol=0.01)
      println("--------------------")
    end

  end
  println("====================")

end


# As of 8/23/2017, this program does not fully run.
# Since the "fixed" option is not used in the current nearly-neutral paper, this option has not been kept up to date.
# Front end for src/nn_poplist.jl
# For the case where selection coefficients are drawn from a gamma distribution or a double-sided gamma distributions
#include("../src/NearlyNeutral.jl")
#include("../src/InfAlleles.jl")
include("../src/InfSites.jl")
if length(ARGS) == 0
  simname = "../experiments/examples/fi_example2"
else
  simname = ARGS[1]
end
include("$(simname).jl")
println("simname: ",simname)
stream = open("$(simname).csv","w")
println("stream: ",stream)

current_dir = pwd()
date_string = "../data/"*Dates.format(now(),"mm_dd_yy")*"/"
#try mkdir(date_string) catch end   # create today's directory with no error if it already exists
println("date string: ",date_string)
if !isdefined(:mu_list_flag)
  mu_list_flag=false
end

@doc """ type trial_result
  A universal trial result for all run_trials functions
  Constructors for specific cases are defined below.
"""
type trial_result
  nn_simtype::Int64
  n::Int64    # sample size, must be <= N
  N::Int64    # popsize
  N_mu::Float64  # N*mu, population mutation rate
  ngens::Int64
  burn_in::Float64
  fix_minimum::Float64
  dfe::Function
  dfe_s::Float64
  dfe_modulus::Int64
  expected_richness::Float64  # sum_{i=0}^{n-1} theta/(theta+i) where theta = 2*N_mu.
  average_richness::Float64   # Average number of traits in populations
  average_richness_stddev::Float64   # Average number of traits in populations
  expected_w_homoz::Float64
  w_homoz::Float64
  w_homoz_stddev::Float64
  IQV::Float64
  IQV_stddev::Float64
  number_active::Int64
  number_extinctions::Int64
  number_fixations::Int64
  extinction_time::Float64
  fixation_time::Float64
end

function add_expected_richness( tr::trial_result )
  sum = 0.0
  theta = 2*tr.N_mu
  for i = 0:(tr.N-1)
    sum += 1.0/(theta+i)
  end
  tr.expected_richness = theta*sum
end

# Constructor that sets the parameters
function trial_result( nn_symtype::Int64, n::Int64, N::Int64, N_mu::Float64, ngens::Int64,
    fix_minimum::Float64=0.5, burn_in::Float64=2.0, dfe::Function=dfe_neutral,
    dfe_s::Float64=0.0, dfe_modulus::Int64=10 )
  #tr = trial_result( nn_simtype, n, N, N_mu, ngens, burn_in, fix_minimum, dfe, dfe_s, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0, 0, 0, 0.0, 0.0, 0.0, 0.0, 0.0 )
  tr = trial_result( nn_simtype, n, N, N_mu, ngens, burn_in, fix_minimum, dfe, dfe_s, dfe_modulus, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0, 0, 0, 0.0, 0.0 )
  add_expected_richness( tr )
  tr
end

function print_trial_result( tr::trial_result )
  if tr.nn_simtype == 1
    println("\ninfinite alleleles model")
  elseif tr.nn_simtype == 2
    println("\ninfinite sites model")
  else
    println("nn_simtype: ", tr.nn_simtype)
  end
  println("n: ", tr.n)
  println("N: ", tr.N)
  println("N_mu: ", tr.N_mu)
  println("mu: ", tr.N_mu/tr.N)
  println("ngens: ", tr.ngens)
  println("burn_in: ", tr.burn_in)
  println("fix_minimum: ", tr.fix_minimum)
  println("dfe: ", tr.dfe)
  println("dfe_s: ", tr.dfe_s)
  println("dfe_modulus: ", tr.dfe_modulus)
  println("expected richness: ",tr.expected_richness)
  println("average richness: ",tr.average_richness)
  println("average richness_stddev: ",tr.average_richness_stddev)
  println("expected_w_homoz: ",tr.expected_w_homoz)
  println("w_homoz: ",tr.w_homoz)
  println("w_homoz_stddev: ",tr.w_homoz_stddev)
  println("IQV: ",tr.IQV)
  println("IQV_stddev: ",tr.IQV_stddev)
  println("number_active: ", tr.number_active)
  println("number_extinctions: ", tr.number_extinctions)
  println("number_fixations: ", tr.number_fixations)
  println("extinction_time: ", tr.extinction_time)
  println("fixation_time: ", tr.fixation_time)
end

# Convert an innovation_oollection to a trial_result
# The innovation_connection and the trial_result must agree in fields n, N,
function convert_to_trial_result( ic::innovation_collection, tr::trial_result )
  tr.number_active = length(ic.active)
  tr.number_extinctions = length(ic.extinct)
  tr.number_fixations = length(ic.fixed)
  tr.fix_minimum = ic.fix_minimum
  tr.extinction_time = average_time_to_extinction( ic )
  tr.fixation_time = average_time_to_fixation( ic )
  #tr.average_fitness_fixed = average_fitness_fixed( ic::innovation_collection )
  #tr.average_fitness_extinct = average_fitness_extinct( ic::innovation_collection )
  #tr.average_fitness_all = average_fitness_all( ic::innovation_collection )
end

function run_trials(popsize_multiplier_list::Vector{Int64}=[1]; mu_list_flag::Bool=false)
  println("stream: ",stream)
  trial = 1
  N = N_list[1]
  n = Int(floor(N*(1//popsize_multiplier_list[1])))
  if !mu_list_flag
    tr = trial_result( nn_simtype, n, N, N_mu_list[1], ngens, fix_minimum, burn_in, InfSites.dfe_mod, dfe_s_list[1], dfe_modulus )
    writeheader(stream, popsize_multiplier_list, N_list, N_mu_list, tr )
    for dfe_s in dfe_s_list
      dfe = x->dfe_mod( x, modulus=dfe_modulus, fit_inc=1.0+dfe_s )
      for N_mu in N_mu_list
        for N in N_list
          for psize_m in popsize_multiplier_list
            n = Int(floor(N*(1//psize_m)))
            tr = trial_result( nn_simtype, n, N, N_mu, ngens, fix_minimum, burn_in, dfe, dfe_s, dfe_modulus )
            run_trial( tr )
            writerow(stream, trial, tr )
            trial += 1
          end
        end
      end
    end
  else  # if mu_list_flag
    tr = trial_result( nn_simtype, n, N, N*mu_list[1], ngens, fix_minimum, burn_in, dfe_mod, dfe_s_list[1], dfe_modulus )
    writeheader(stream, popsize_multiplier_list, N_list, mu_list, tr, mu_list_flag=mu_list_flag )
    for dfe_s in dfe_s_list
      dfe = x->dfe_mod( x, modulus=dfe_mod, fit_inc=dfe_s )
      for mu in mu_list
        for N in N_list
          for psize_m in popsize_multiplier_list
            n = Int(floor(N*(1//psize_m)))
            N_mu = N*mu
            tr = trial_result( nn_simtype, n, N, N_mu, ngens, fix_minimum, burn_in, dfe, dfe_s, dfe_modulus )
            run_trial( tr )
            writerow(stream, trial, tr, mu_list_flag=mu_list_flag )
            trial += 1
          end
        end
      end
    end
  end  # if !mu_list_flag
end

function run_trials_mu(popsize_multiplier_list::Vector{Int64}=[1])
  println("stream: ",stream)
  trial = 1
  N = N_list[1]
  n = Int(floor(N*(1//popsize_multiplier_list[1])))
  tr = trial_result( nn_simtype, n, N, N*mu_list[1], ngens, fix_minimum, burn_in, dfe_mod, dfe_s_list[1], dfe_modulus )
  writeheader(stream, popsize_multiplier_list, N_list, mu_list, tr )
  for dfe_s in dfe_s_list
    dfe = x->dfe_mod( x, modulus=dfe_mod, fit_inc=dfe_s )
    for mu in mu_list
      for N in N_list
        for psize_m in popsize_multiplier_list
          n = Int(floor(N*(1//psize_m)))
          N_mu = N*mu
          tr = trial_result( nn_simtype, n, N, N_mu, ngens, fix_minimum, burn_in, dfe, dfe_s, dfe_modulus )
          run_trial( tr )
          writerow(stream, trial, tr )
          trial += 1
        end
      end
    end
  end
end

function run_trial( tr::trial_result )
  if tr.nn_simtype == 1
    println("dfe(tr.dfe_modulus): ",tr.dfe(tr.dfe_modulus))
    ic = innovation_collection( tr.N, tr.fix_minimum )
    poplist = nn_poplist(tr.N,tr.N_mu,tr.ngens,tr.dfe,combine=false,ic=ic)
    convert_to_trial_result( ic, tr )
    if tr.n < tr.N
      poplist = map(x->sample_population(x,tr.n),poplist)
    end
    add_stats_to_trial_result!( tr, poplist )
    print_trial_result( tr )
    return tr
  else
    println("nn_simtype ",nn_simtype," not implemented")
  end
end

function add_stats_to_trial_result!( tr::trial_result, poplist::Vector{Population} )
  #println("poplist: ",poplist)
  #println("poplist fitnesses: ",map(y->map(tr.dfe,y)),poplist)
  mu = tr.N_mu/tr.N
  filtered_poplist = map(y->filter(x->(x%tr.dfe_modulus==0),y),poplist)
  filtered_N = map(length,filtered_poplist)  # vector of number of individuals whose modulus is dfe_modulus
  #println("filtered_poplist: ",filtered_poplist)
  filtered_pcounts = map(pop_counts64,filtered_poplist)
  richness = map(length,filtered_pcounts)  # vector of richnesses over generations
  w_homoz = map(watterson_homozygosity,filtered_pcounts)
  IQV_lst = map(IQV,filtered_pcounts)
  expected_filtered_N = tr.N/tr.dfe_modulus   # Expected number of indivs whose modulus is dfe_modulus
  println("length(w_homoz): ",length(w_homoz),"  length(IQV_lst): ",length(IQV_lst),"  mean_filtered_N: ",mean(filtered_N))
  println("w_homoz[1]: ",w_homoz[1],"  sum(w_homoz): ",sum(w_homoz))
  tr.average_richness = mean(richness)
  tr.average_richness_stddev = std(richness)
  #tr.expected_w_homoz = 1.0/(1.0+2.0*tr.N_mu)
  tr.expected_w_homoz = 1.0/(1.0+2.0*expected_filtered_N*mu)
  tr.w_homoz = mean(w_homoz)
  println("tr.w_homoz: ",tr.w_homoz)
  tr.w_homoz_stddev = std(w_homoz)
  tr.IQV = mean(IQV_lst)
  tr.IQV_stddev = std(IQV_lst)
end

@doc """ function writeheader()
    burn_in::Float64, slat_reps::Int64=100000 )
Write header line for output CSV file that corresponds to stream.
See the comments for run_simulation for a description of the parameters.
"""
function writeheader(stream::IO, popsize_multiplier_list::Vector{Int64}, N_list::Vector{Int64},
    mu_list::Vector{Float64}, tr::trial_result; mu_list_flag::Bool=false )
  dfe_params = [:dfe_adv_prob, :dfe_adv_alpha, :dfe_adv_theta, :dfe_disadv_prob, :dfe_disadv_alpha, :dfe_disadv_theta]
  param_strings = [
    "# $(string(Dates.today()))",
    "# $((nn_simtype==1)?"infinite alleles model":"infinite_sites_model")",
    "# popsize_multiplier_list=$(popsize_multiplier_list)",
    "# N_list=$(N_list)",
    mu_list_flag ? "# mu_list=$(mu_list)": # N_mu_list=$(N_mu_list)",
    "# dfe_modulus=$(tr.dfe_modulus)",
    "# ngens=$(tr.ngens)",
    "# burn_in=$(tr.burn_in)",
    "# fix_minimum=$(tr.fix_minimum)",
    #"# dfe=$(tr.dfe)",
    #"# dfe_s=$(tr.dfe_s)"
    ]
  for dp in dfe_params
    if isdefined(dp)
      Base.push!(param_strings,"# $(string(dp))=$(eval(dp))")
    end
  end
  write(stream, join(param_strings, "\n"), "\n")
  first_heads = [
    "type",
    #  "n",
    "N",
    (mu_list_flag ? "mu":"N_mu"),
    "s"
  ]
  last_heads =
  [ "expected_richness",
    "average_richness",
    "average_richness_stddev",
    "exp_neut_w_heteroz",
    "w_heteroz",
    "w_heteroz_stddev",
    "IQV",
    "IQV_stddev",
    "num_extinct",
    "num_fixed",
    "fraction_fixed",
    "ave_extinct_time",
    "ave_fixed_time"
  ]
  line = join(vcat( first_heads, last_heads), ",")
  write(stream, line, "\n")
end

function writerow(stream::IO, trial::Int64, tr::trial_result; mu_list_flag::Bool=false   )
  first = Any[
    (isdefined(:type_str) ? type_str :trial),
    #tr.n,           # sample size
    tr.N,           # popsize
    mu_list_flag ? tr.N_mu/tr.N:tr.N_mu,
    tr.dfe_s,
  ]
  if tr.nn_simtype == 0
    mid = Any[]
  elseif tr.nn_simtype == 1
    mid = Any[
      tr.expected_richness,
      tr.average_richness,
      tr.average_richness_stddev,
      1.0-tr.expected_w_homoz,
      1.0-tr.w_homoz,
      tr.w_homoz_stddev,
      tr.IQV,
      tr.IQV_stddev,
      tr.number_extinctions,
      tr.number_fixations,
      Float64(tr.number_fixations)/(tr.number_extinctions+tr.number_fixations),
      tr.extinction_time,
      tr.fixation_time
    ]
  end
  line = join( vcat( first, mid ), "," )
  write(stream, line, "\n")
end

if isdefined(:popsize_multiplier_list)
  run_trials(popsize_multiplier_list, mu_list_flag=mu_list_flag )
else
  run_trials( mu_list_flag=mu_list_flag)
end


module SDCO

using SparseArrays, DataStructures, LinearAlgebra, Printf

include("types.jl")

include("algebra_E.jl")
include("evaluate.jl")

include("input.jl")
include("io.jl")

include("NT.jl")
include("predcorr.jl")

include("bigM_types.jl")
include("bigM_solver.jl")

include("HSD_solver.jl")

include("test_cases.jl")

end


module ANN
    export NeuralLayer,ArtificialNeuralNetwork,fit!,predict
    include("base.jl")
end


using Openscad
using Test

@testset "2d shapes" begin
    @test (toString(Circle(10, false)) == "circle(r = 10, center = false);")
    @test (toString(Square([1, 1], true)) == "square(size = [1, 1], center = true);")
end

function test2d()
    square = Square([4, 5], true)
    square2 = Square([0, 0], true)
    tree =
    Unify([
        Circle(3, false),
        square,
        Translate([1, 1, 1], [
            Openscad.Unify([
                square,
                square2
            ])
        ]),
        End()
    ])
    tree
end

@testset "complex 2d shape" begin
    @test (toString(test2d()) == "union(){circle(r = 3, center = false);square(size = [4, 5], center = true);translate([1, 1, 1]) {union(){square(size = [4, 5], center = true);square(size = [0, 0], center = true);};};};")
end

@testset "3d shapes" begin
end


testOnAllContexts("CrossEntropy") do ctx
    println("\tsolves a simple problem")
    testSimpleProblem(
        comparer ->
        begin
            ce = CrossEntropy(ctx, 60, comparer, dna -> dna)
            @test ce.dnaSize == 60
            @test ce.populationSize == 100
            @test ce.mutationChance == 0.05f0
            @test ce.mutationStrength == 0.02f0
            @test ce.selectionStdDev == 0.3f0
            @test ce.keepElitesRate == 0.05f0
            ce
        end)

    println("\tsolves a problem with synthesizer")
    testSimpleProblem(
        ctx,
        (comparer, synth) ->
        begin
            ce = CrossEntropy(
                ctx,
                dnaSize!(synth),
                comparer,
                dna -> decode(synth, dna),
                mutationChance = 0.05f0,
                selectionStdDev = 0.3f0)
            @test ce.dnaSize == dnaSize!(synth)
            @test ce.populationSize == 100
            @test ce.mutationChance == 0.05f0
            @test ce.mutationStrength == 0.02f0
            @test ce.selectionStdDev == 0.3f0
            @test ce.keepElitesRate == 0.05f0
            ce
        end)
end


module Gate

using PyCall
import ..@pyfunc
import ..@pyclass

# qulacs = PyNULL()
gate = PyNULL()

const gate_classes = [
    :Adaptive, :RX,
    :AmplitudeDampingNoise, :RY,
    :BitFlipNoise, :RZ,
    :CNOT, :RandomUnitary,
    :CP, :ReversibleBoolean,
    :CPTP, :S,
    :CZ, :SWAP,
    :DenseMatrix, :Sdag,
    :DephasingNoise, :SparseMatrix,
    :DepolarizingNoise, :StateReflection,
    :DiagonalMatrix, :T,
    :FREDKIN, :TOFFOLI,
    :H, :Tdag,
    :Identity, :TwoQubitDepolarizingNoise,
    :IndependentXZNoise, :U1,
    :Instrument, :U2,
    :Measurement, :U3,
    :P0, :X,
    :P1, :Y,
    :ParametricPauliRotation, :Z,
    :ParametricRX,               #:add,
    :ParametricRY,               #:merge,
    :ParametricRZ,               #:sqrtX,
    :Pauli,                      #:sqrtXdag,
    :PauliRotation,              #:sqrtY,
    :Probabilistic,              #:sqrtYdag,
    :ProbabilisticInstrument,    #:to_matrix_gate,
]

#=
for class in gate_classes
    @eval begin
        struct $class
            pyobj::PyObject
            $(class)(args...) = new(qulacs.gate.$(class)(args...))
        end

        PyObject(t::$(class)) = t.pyobj

        function Base.propertynames(t::$(class))
            propertynames(getfield(t, :pyobj))
        end

        function Base.getproperty(t::$(class), s::Symbol)
            if s ∈ fieldnames($(class))
                return getfield(t, s)
            else
                return getproperty(getfield(t, :pyobj), s)
            end
        end

        export $(class)
    end
end
=#

for class in gate_classes
    @eval begin
        @pyclass gate $(class)
        export $(class)
    end
end

const gate_functions = [
    :add,
    :merge,
    :sqrtX,
    :sqrtXdag,
    :sqrtY,
    :sqrtYdag,
    :to_matrix_gate,
]

#=
for func in gate_functions
    @eval begin
        function $(func)(args...; kwargs...)
            qulacs.gate.$(func)(args...; kwargs...)
        end
        export $(func)
    end
end
=#

for func in gate_functions
    @eval begin
        @pyfunc gate $(func)
        export $(func)
    end
end

function __init__()
    #copy!(qulacs, pyimport("qulacs"))
    copy!(gate, pyimport("qulacs.gate"))
end

end # module Gate


@testset "Coupling constraints" begin
    cp = convert(CoreModelCoupled, test_LP())
    @test size(cp.lm.S) == (4, 3)
    @test size(stoichiometry(convert(CoreModel, cp))) == (4, 3)
    new_cp = add_coupling_constraints(cp, stoichiometry(cp)[end, :], -1.0, 1.0)
    @test n_coupling_constraints(cp) + 1 == n_coupling_constraints(new_cp)

    new_cp =
        add_coupling_constraints(cp, stoichiometry(cp)[1:2, :], [-1.0; -1.0], [1.0; 1.0])
    @test n_coupling_constraints(cp) + 2 == n_coupling_constraints(new_cp)

    n_c = n_coupling_constraints(cp)
    add_coupling_constraints!(cp, stoichiometry(cp)[end, :], -1.0, 1.0)
    @test n_c + 1 == n_coupling_constraints(cp)
    add_coupling_constraints!(cp, stoichiometry(cp)[1:2, :], [-1.0; -1.0], [1.0; 1.0])
    @test n_c + 3 == n_coupling_constraints(cp)

    n_c = n_coupling_constraints(cp)
    remove_coupling_constraints!(cp, 1)
    @test n_c - 1 == n_coupling_constraints(cp)
    remove_coupling_constraints!(cp, [1, 2])
    @test n_c - 3 == n_coupling_constraints(cp)
    @test n_coupling_constraints(cp) == 0

    cp = test_coupledLP()
    n_c = n_coupling_constraints(cp)
    new_cp = remove_coupling_constraints(cp, 1)
    @test size(coupling(cp)) == (n_c, n_reactions(cp))
    @test n_c - 1 == n_coupling_constraints(new_cp)
    @test n_coupling_constraints(cp) == n_c
    new_cp = remove_coupling_constraints(cp, [1, 2])
    @test n_c - 2 == n_coupling_constraints(new_cp)
    new_cp = remove_coupling_constraints(cp, Vector(1:n_coupling_constraints(cp)))
    @test n_coupling_constraints(new_cp) == 0
    @test n_coupling_constraints(cp) == n_c

    cp = test_coupledLP()
    change_coupling_bounds!(cp, [3, 1], cl = [-10.0, -20], cu = [10.0, 20])
    cl, cu = coupling_bounds(cp)
    @test cl[[1, 3]] == [-20, -10]
    @test cu[[1, 3]] == [20, 10]
    change_coupling_bounds!(cp, [1000, 1001], cl = [-50.0, -60.0])
    cl, cu = coupling_bounds(cp)
    @test cl[[1000, 1001]] == [-50.0, -60.0]
end

@testset "Add reactions" begin
    cp = convert(CoreModelCoupled, test_LP())
    cp = add_coupling_constraints(cp, stoichiometry(cp)[end, :], -1.0, 1.0)

    new_cp = add_reactions(cp, 2.0 * ones(4), 3 .* ones(4), 2.0, -1.0, 1.0)
    @test new_cp isa CoreModelCoupled
    @test cp.C == new_cp.C[:, 1:end-1]
    @test cp.cl == new_cp.cl
    @test cp.cu == new_cp.cu

    new_cp = add_reactions(
        cp,
        2.0 * ones(4),
        3 .* ones(4),
        2.0,
        -1.0,
        1.0,
        "r4",
        ["m$i" for i = 1:4],
    )
    @test cp.C == new_cp.C[:, 1:end-1]
    @test cp.cl == new_cp.cl
    @test cp.cu == new_cp.cu

    new_cp = add_reactions(
        cp,
        2.0 * ones(4, 10),
        3 .* ones(4),
        2 .* ones(10),
        -ones(10),
        ones(10),
    )
    @test cp.C == new_cp.C[:, 1:end-10]
    @test cp.cl == new_cp.cl
    @test cp.cu == new_cp.cu

    new_cp = add_reactions(
        cp,
        2.0 * ones(4, 10),
        3 .* ones(4),
        2 .* ones(10),
        -ones(10),
        ones(10),
        ["r$i" for i = 1:10],
        ["m$i" for i = 1:4],
    )
    @test cp.C == new_cp.C[:, 1:end-7] # 3 reactions were already present
    @test cp.cl == new_cp.cl
    @test cp.cu == new_cp.cu

    new_cp =
        add_reactions(cp, 2.0 * sprand(4000, 0.5), 3 .* sprand(4000, 0.5), 2.0, -1.0, 1.0)
    @test cp.C == new_cp.C[:, 1:end-1]
    @test cp.cl == new_cp.cl
    @test cp.cu == new_cp.cu

    cm = CoreModel(
        2.0 * ones(4, 10),
        3 .* ones(4),
        2 .* ones(10),
        -ones(10),
        ones(10),
        ["r$i" for i = 1:10],
        ["m$i" for i = 1:4],
    )
    new_cp = add_reactions(cp, cm)
    @test cp.C == new_cp.C[:, 1:end-7] # 3 reactions were already present
    @test cp.cl == new_cp.cl
    @test cp.cu == new_cp.cu
end

@testset "Remove reactions" begin
    cp = convert(CoreModelCoupled, test_LP())
    cp = add_coupling_constraints(cp, 1.0 .* collect(1:n_reactions(cp)), -1.0, 1.0)

    new_cp = remove_reactions(cp, [3, 2])
    @test new_cp isa CoreModelCoupled
    @test new_cp.C[:] == cp.C[:, 1]
    @test new_cp.cl == cp.cl
    @test new_cp.cu == cp.cu

    new_cp = remove_reaction(cp, 2)
    @test new_cp.C == cp.C[:, [1, 3]]
    @test new_cp.cl == cp.cl
    @test new_cp.cu == cp.cu

    new_cp = remove_reaction(cp, "r1")
    @test new_cp.C == cp.C[:, 2:3]
    @test new_cp.cl == cp.cl
    @test new_cp.cu == cp.cu

    new_cp = remove_reactions(cp, ["r1", "r3"])
    @test new_cp.C[:] == cp.C[:, 2]
    @test new_cp.cl == cp.cl
    @test new_cp.cu == cp.cu

    new_cp = remove_reactions(cp, [1, 4])
    @test new_cp.C == cp.C[:, 2:3]

    new_cp = remove_reactions(cp, [1, 1, 2])
    @test new_cp.C[:] == cp.C[:, 3]
end

@testset "Change bounds" begin
    cp = convert(CoreModelCoupled, test_LP())
    @test cp isa CoreModelCoupled

    change_bound!(cp, 1, lower = -10, upper = 10)
    @test cp.lm.xl[1] == -10
    @test cp.lm.xu[1] == 10
    change_bounds!(cp, [1, 2]; lower = [-11, -12.2], upper = [11, 23.0])
    @test cp.lm.xl[2] == -12.2
    @test cp.lm.xu[1] == 11

    change_bound!(cp, "r1", lower = -101, upper = 101)
    @test cp.lm.xl[1] == -101
    @test cp.lm.xu[1] == 101
    change_bounds!(cp, ["r1", "r2"]; lower = [-113, -12.23], upper = [113, 233.0])
    @test cp.lm.xl[2] == -12.23
    @test cp.lm.xu[1] == 113

    new_model = change_bound(cp, 1, lower = -10, upper = 10)
    @test new_model.lm.xl[1] == -10
    @test new_model.lm.xu[1] == 10
    new_model = change_bounds(cp, [1, 2]; lower = [-11, -12.2], upper = [11, 23.0])
    @test new_model.lm.xl[2] == -12.2
    @test new_model.lm.xu[1] == 11

    new_model = change_bound(cp, "r1", lower = -101, upper = 101)
    @test new_model.lm.xl[1] == -101
    @test new_model.lm.xu[1] == 101
    new_model =
        change_bounds(cp, ["r1", "r2"]; lower = [-113, -12.23], upper = [113, 233.0])
    @test new_model.lm.xl[2] == -12.23
    @test new_model.lm.xu[1] == 113

end


#=
surface_test:
- Julia version: 1.1.0
- Author: Jirik
- Date: 2019-03-19
=#
# # using Revise
using Test
using Logging
using SparseArrays
using Distributed
if nprocs() == 1
    addprocs(3)
end

using LarSurf
# using Plasm
# using LinearAlgebraicRepresentation
# Lar = LinearAlgebraicRepresentation

# TODO repair this test
@testset "Get Fchar test" begin

    println("generate data")
    segmentation = LarSurf.data234()
    block_size = [2,2,2]
    # Flin2 = __grid_get_surface_Fchar_per_block_parallel_pmap(segmentation, block_size)

end