text:
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

"""
    AbstractModel

The abstract super type of all models in MXNet.jl.
"""
abstract type AbstractModel end

"""
    FeedForward

The feedforward model provides convenient interface to train and predict on
feedforward architectures like multi-layer MLP, ConvNets, etc. There is no
explicitly handling of *time index*, but it is relatively easy to implement
unrolled RNN / LSTM under this framework (*TODO*: add example). For models
that handles sequential data explicitly, please use *TODO*...
"""
mutable struct FeedForward <: AbstractModel
  arch        :: SymbolicNode
  ctx         :: Vector{Context}

  arg_params  :: Dict{Symbol}
  aux_params  :: Dict{Symbol}

  pred_exec   :: Union{Executor,Void}

  # leave the rest fields undefined
  FeedForward(arch::SymbolicNode, ctx::Vector{Context}) = new(arch, ctx)
  FeedForward(arch::SymbolicNode, ctx::Context) = new(arch, [ctx])
end

"""
Get a split of `batch_size` into `n_split` pieces for data parallelization. Returns a vector
of length `n_split`, with each entry a `UnitRange{Int}` indicating the slice index for that
piece.
"""
function _split_inputs(batch_size::Int, n_split::Int)
  @assert(batch_size >= n_split)
  per_split = floor(Int, batch_size / n_split)
  counts    = Base.zeros(Int, n_split)+per_split
  extra     = batch_size - Base.sum(counts)
  counts[1:extra] += 1

  cum = [0, cumsum(counts)...]
  idx = [cum[i-1]+1:cum[i] for i = 2:length(cum)]
  return idx
end

"""
    FeedForward(arch :: SymbolicNode, ctx)

# Arguments:
* `arch`: the architecture of the network constructed using the symbolic API.
* `ctx`: the devices on which this model should do computation. It could be a single `Context`
         or a list of `Context` objects. In the latter case, data parallelization will be used
         for training. If no context is provided, the default context `cpu()` will be used.
"""
FeedForward(arch::SymbolicNode; context::Union{Context,Vector{Context}} = [cpu()]) =
  FeedForward(arch, context)

"""
    init_model(self, initializer; overwrite=false, input_shapes...)

Initialize the weights in the model.

This method will be called automatically when training a model. So there is usually no
need to call this method unless one needs to inspect a model with only randomly initialized
weights.

# Arguments:
* `self::FeedForward`: the model to be initialized.
* `initializer::AbstractInitializer`: an initializer describing how the weights should be initialized.
* `overwrite::Bool`: keyword argument, force initialization even when weights already exists.
* `input_shapes`: the shape of all data and label inputs to this model, given as keyword arguments.
                  For example, `data=(28,28,1,100), label=(100,)`.
"""
function init_model(self::FeedForward, initializer::AbstractInitializer; overwrite::Bool=false, input_shapes...)
  # all arg names, including data, label, and parameters
  arg_names    = list_arguments(self.arch)

  input_names  = [x[1] for x in input_shapes]

  param_names = setdiff(arg_names, input_names)
  aux_names   = list_auxiliary_states(self.arch)

  arg_shapes, out_shapes, aux_shapes = infer_shape(self.arch; input_shapes...)

  # If target dict is not yet defined set a temporary one
  if !isdefined(self, :arg_params)
    self.arg_params = Dict{Symbol, NDArray}()
  end
  if !isdefined(self, :aux_params)
    self.aux_params = Dict{Symbol, NDArray}()
  end

  arg_params = Dict{Symbol,NDArray}()
  aux_params = Dict{Symbol,NDArray}()

  for (name, shape) in filter(x -> in(x[1],param_names), zip(arg_names, arg_shapes))
    if haskey(self.arg_params, name)
      if shape == size(self.arg_params[name])
        arg_params[name] = self.arg_params[name]
        continue
      else
        warn("Shape mismatch for $name. Overwriting with new one.")
        delete!(self.arg_params, name)
      end
    end
    arg_params[name] = empty(shape)
  end

  for (name, shape) in zip(aux_names, aux_shapes)
    if haskey(self.aux_params, name)
      if shape == size(self.aux_params[name])
        aux_params[name] = self.aux_params[name]
        continue
      else
        warn("Shape mismatch for $name. Overwriting with new one.")
        delete!(self.aux_params, name)
      end
    end
    aux_params[name] = empty(shape)
  end

  for (k,v) in arg_params
    if overwrite || !haskey(self.arg_params, k)
      init(initializer, k, v)
    end
  end
  for (k,v) in aux_params
    if overwrite || !haskey(self.aux_params, k)
      init(initializer, k, v)
    end
  end

  self.arg_params = arg_params
  self.aux_params = aux_params

  return (arg_names, param_names, aux_names)
end

function _setup_predictor(self::FeedForward, overwrite::Bool=false; verbosity::Integer = 1, data_shapes...)
  if !isdefined(self, :pred_exec) || isa(self.pred_exec, Void) || overwrite
    if !isdefined(self, :arg_params) || !isdefined(self, :aux_params)
      @assert(false, "Model weights not defined, please init or train the model, or load from file")
    end

    # the predictor use only the first device
    self.pred_exec = simple_bind(self.arch, self.ctx[1]; grad_req=GRAD_NOP, data_shapes...)
    dbg_str = mx.debug_str(self.pred_exec)
    verbosity >= 1 && info(string("TempSpace: ", split(dbg_str, ['\n'])[end-2]..., " on ", self.ctx[1]))
    copy_params_from(self.pred_exec, self.arg_params, self.aux_params)
  else
    # make sure the new setup is compatible with the existing one
    for (d_name, d_shape) in data_shapes
      @assert(d_shape == size(self.pred_exec.arg_dict[d_name]),
              "Shape of $d_name mismatch with existing predictor, use overwrite=true overwrite existing predictor")
    end
  end
end

"""
    predict(self, data; overwrite=false, callback=nothing)

Predict using an existing model. The model should be already initialized, or trained or loaded from
a checkpoint. There is an overloaded function that allows to pass the callback as the first argument,
so it is possible to do

```julia
predict(model, data) do batch_output
  # consume or write batch_output to file
end
```

# Arguments:
* `self::FeedForward`:  the model.
* `data::AbstractDataProvider`: the data to perform prediction on.
* `overwrite::Bool`: an `Executor` is initialized the first time predict is called. The memory
                     allocation of the `Executor` depends on the mini-batch size of the test
                     data provider. If you call predict twice with data provider of the same batch-size,
                     then the executor can be potentially be re-used. So, if `overwrite` is false,
                     we will try to re-use, and raise an error if batch-size changed. If `overwrite`
                     is true (the default), a new `Executor` will be created to replace the old one.
* `verbosity::Integer`: Determines the verbosity of the print messages. Higher numbers
          leads to more verbose printing. Acceptable values are
          - `0`: Do not print anything during prediction
          - `1`: Print allocation information during prediction

!!! note
    Prediction is computationally much less costly than training, so the bottleneck sometimes becomes the IO
    for copying mini-batches of data. Since there is no concern about convergence in prediction, it is better
    to set the mini-batch size as large as possible (limited by your device memory) if prediction speed is a
    concern.

    For the same reason, currently prediction will only use the first device even if multiple devices are
    provided to construct the model.

!!! note
    If you perform further after prediction. The weights are not automatically synchronized if `overwrite`
    is set to false and the old predictor is re-used. In this case
    setting `overwrite` to true (the default) will re-initialize the predictor the next time you call
    predict and synchronize the weights again.

See also [`train`](@ref), [`fit`](@ref), [`init_model`](@ref), and [`load_checkpoint`](@ref)
"""
function predict(callback::Function, self::FeedForward, data::AbstractDataProvider;
                 overwrite::Bool = true, verbosity::Integer = 1)
  predict(self, data; overwrite = overwrite, callback=callback, verbosity = verbosity)
end
function predict(self::FeedForward, data::AbstractDataProvider;
                 overwrite::Bool = true, callback::Union{Function,Void}=nothing, verbosity::Integer = 1)
  data_shapes = provide_data(data)
  data_names  = [x[1] for x in data_shapes]
  _setup_predictor(self, overwrite; verbosity = verbosity, data_shapes...)

  batch_size  = get_batch_size(data)
  data_arrays =  [self.pred_exec.arg_dict[name] for name in data_names]
  output_list = [Array{MX_float}[] for i=1:length(self.pred_exec.outputs)]
  for batch in eachbatch(data)
    load_data!(data, batch, data_arrays)
    forward(self.pred_exec, is_train=false)
    if isa(callback, Void)
      # no callback, accumulate the data and return at the end
      for (o_list, o_nd) in zip(output_list, self.pred_exec.outputs)
        push!(o_list, copy(slice(o_nd, 1:count_samples(data, batch))))
      end
    else
      outputs = self.pred_exec.outputs
      if length(outputs) == 1
        outputs = outputs[1]
      end
      callback(outputs)
    end
  end

  if !isa(callback, Void)
    # callback exists, do not accumulate data
    return nothing
  end

  if isempty(output_list)
    # maybe model does not have outputs
    return nothing
  end
  if isempty(output_list[1])
    # maybe no output because data is empty
    return length(output_list) == 1 ? output_list[1] : output_list
  end

  # concatenate along mini-batches
  output_arrays = [cat(ndims(x[1]), x...) for x in output_list]
  if length(output_arrays) == 1
    # only 1 output, return it directly, instead of a list
    output_arrays = output_arrays[1]
  end
  return output_arrays
end

function _init_model(self::FeedForward, data::AbstractDataProvider,
                     initializer::AbstractInitializer, overwrite::Bool)
  init_model(self, initializer; overwrite=overwrite,
             [provide_data(data)..., provide_label(data)...]...)
end

function _create_kvstore(kv_type::Symbol, num_device::Int, arg_params::Dict{Symbol}, verbosity::Int)
  if num_device == 1 && !ismatch(r"dist", string(kv_type))
    return nothing
  else
    if kv_type == :local
      max_size = maximum([prod(size(param)) for (k,param) in arg_params])
      if max_size < 1024 * 1024 * 16
        kv_type = :local_update_cpu
      else
        kv_type = :local_allreduce_cpu
      end
      verbosity >= 2 && info("Auto-select kvstore type = $kv_type")
    end
    return KVStore(kv_type)
  end
end

@defstruct TrainingOptions (
  initializer :: AbstractInitializer = UniformInitializer(0.01),
  n_epoch     :: Int = 10,
  eval_data   :: Union{Void,AbstractDataProvider} = nothing,
  eval_metric :: AbstractEvalMetric = Accuracy(),
  kvstore     :: Union{Symbol,KVStore} = :local,
  force_init  :: Bool = false,
  callbacks   :: Vector{AbstractCallback} = AbstractCallback[],
  verbosity   :: Int = 3,
  η_decay     :: Symbol = :epoch,
)

function _invoke_callbacks(m::FeedForward, callbacks::Vector{AbstractCallback},
                           state::OptimizationState, type_filter::Type;
                           metric = Vector{Tuple{Symbol,Real}}())
  map(callbacks) do cb
    !isa(cb, type_filter) && return

    # epoch callback have extra access to the model object
    type_filter == AbstractEpochCallback && return cb(m, state, metric)

    cb(state)
  end
end

"""
    train(model :: FeedForward, ...)

Alias to [`fit`](@ref).
"""
train(m::FeedForward, opt::AbstractOptimizer, data::AbstractDataProvider; kw...) =
  fit(m, opt, data; kw...)

"""
    fit(model::FeedForward, optimizer, data; kwargs...)

Train the `model` on `data` with the `optimizer`.

* `model::FeedForward`: the model to be trained.
* `optimizer::AbstractOptimizer`: the optimization algorithm to use.
* `data::AbstractDataProvider`: the training data provider.
* `n_epoch::Int`: default 10, the number of full data-passes to run.
* `eval_data::AbstractDataProvider`: keyword argument, default `nothing`. The data provider for
          the validation set.
* `eval_metric::AbstractEvalMetric`: keyword argument, default [`Accuracy()`](@ref). The metric used
          to evaluate the training performance. If `eval_data` is provided, the same metric is also
          calculated on the validation set.
* `kvstore`: keyword argument, default `:local`. The key-value store used to synchronize gradients
          and parameters when multiple devices are used for training.
   :type kvstore: `KVStore` or `Symbol`
* `initializer::AbstractInitializer`: keyword argument, default `UniformInitializer(0.01)`.
* `force_init::Bool`: keyword argument, default false. By default, the random initialization using the
          provided `initializer` will be skipped if the model weights already exists, maybe from a previous
          call to [`train`](@ref) or an explicit call to [`init_model`](@ref) or [`load_checkpoint`](@ref). When
          this option is set, it will always do random initialization at the begining of training.
* `callbacks::Vector{AbstractCallback}`: keyword argument, default `[]`. Callbacks to be invoked at each epoch or mini-batch,
          see `AbstractCallback`.
* `verbosity::Int`: Determines the verbosity of the print messages. Higher numbers
          leads to more verbose printing. Acceptable values are
          - `0`: Do not print anything during training
          - `1`: Print starting and final messages
          - `2`: Print one time messages and a message at the start of each epoch
          - `3`: Print a summary of the training and validation accuracy for each epoch
* `η_decay::Symbol`: `:epoch` or `:batch`, decay learning rate on epoch or batch.
"""
function fit(self::FeedForward, optimizer::AbstractOptimizer, data::AbstractDataProvider;
             kwargs...)
  opts = TrainingOptions(; kwargs...)

  opts.verbosity >= 1 && info("Start training on $(self.ctx)")

  batch_size  = get_batch_size(data)
  num_dev     = length(self.ctx)
  slices      = _split_inputs(batch_size, num_dev)

  # initialize parameters
  opts.verbosity >= 2 && info("Initializing parameters...")
  arg_names, param_names, aux_names = _init_model(self, data, opts.initializer, opts.force_init)

  # setup kvstore
  kvstore = opts.kvstore
  if isa(kvstore, Symbol)
    opts.verbosity >= 2 && info("Creating KVStore...")
    kvstore = _create_kvstore(kvstore, length(self.ctx), self.arg_params, opts.verbosity)
  end

  update_on_kvstore = true
  if isa(kvstore, Void) || ismatch(r"local_allreduce", string(get_type(kvstore)))
    update_on_kvstore = false
  end

  # get grad attribute to allow for freezing
  freeze_names = Symbol[]
  for (attr, value) in list_all_attr(self.arch)
    sattr = string(attr)
    if endswith(sattr, "grad") && value == "freeze"
      push!(freeze_names, Symbol(sattr[1:end-5]))
    end
  end
  # Needs to correspond to the correct id in the update loop layer idx=1:length(param_names).
  freeze_idx = filter(i -> in(param_names[i], freeze_names), 1:length(param_names))

  # Setup grad_req as a dictionary
  grad_req = Dict{Symbol,GRAD_REQ}()
  for param in param_names
    if in(param, freeze_names)
      grad_req[param] = GRAD_NOP
    else
      grad_req[param] = GRAD_WRITE
    end
  end

  train_execs = Array{Executor}(num_dev)
  for i = 1:num_dev
    data_shapes = Dict(map((x) -> x[1] => tuple(x[2][1:end-1]...,length(slices[i])), provide_data(data)))
    label_shapes = Dict(map((x) -> x[1] => tuple(x[2][1:end-1]...,length(slices[i])), provide_label(data)))
    train_execs[i] = simple_bind(self.arch, self.ctx[i]; grad_req=grad_req, data_shapes..., label_shapes...)
    dbg_str = mx.debug_str(train_execs[i])
    opts.verbosity >= 2 && info(string("TempSpace: ", split(dbg_str, ['\n'])[end-2]..., " on ", self.ctx[i]))

    copy_params_from(train_execs[i], self.arg_params, self.aux_params)
  end

  # set up input data structures
  data_names   = [x[1] for x in provide_data(data)]
  label_names  = [x[1] for x in provide_label(data)]

  data_arrays  = [SlicedNDArray[(slices[i], exec.arg_dict[name]) for (i,exec) in enumerate(train_execs)]
                  for name in data_names]
  label_arrays = [SlicedNDArray[(slices[i], exec.arg_dict[name]) for (i,exec) in enumerate(train_execs)]
                  for name in label_names]

  param_idx    = filter(i -> in(arg_names[i], param_names), 1:length(arg_names))

  param_arrays = [NDArray[exec.arg_arrays[i] for exec in train_execs] for i in param_idx]
  grad_arrays  = [NDArray[exec.grad_arrays[i] for exec in train_execs] for i in param_idx]
  aux_arrays   = [NDArray[exec.aux_arrays[i] for exec in train_execs] for i = 1:length(aux_names)]

  op_state = OptimizationState(batch_size)
  # set up the gradient rescaling if user not set
  iszero(optimizer.scale) && (optimizer.scale = 1 / batch_size)

  if !update_on_kvstore
    updater = getupdater(optimizer)
  end

  if !isa(kvstore, Void)
    if update_on_kvstore
      set_optimizer(kvstore, optimizer)
    end

    opts.verbosity >= 2 && info("Initializing KVStore...")
    # init kv with gradients
    for idx = 1:length(param_arrays)
      param_on_devs = param_arrays[idx]

      init!(kvstore, idx, self.arg_params[param_names[idx]])

      if update_on_kvstore
        # pull weights back
        pull!(kvstore, idx, param_on_devs, priority=-idx)
      end
    end
  end

  # set up output and labels in CPU for evaluation metric
  output_shapes = [tuple(size(x)[1:end-1]...,batch_size) for x in train_execs[1].outputs]
  cpu_dev = Context(CPU)
  cpu_output_arrays = [empty(shape, cpu_dev) for shape in output_shapes]
  cpu_label_arrays  = [empty(shape, cpu_dev) for (name,shape) in provide_label(data)]

  # invoke callbacks on epoch 0
  _invoke_callbacks(self, opts.callbacks, op_state, AbstractEpochCallback)

  opts.verbosity >= 2 && info("Start training...")
  for i_epoch = 1:opts.n_epoch
    time_start = time()
    reset!(opts.eval_metric)

    op_state.curr_epoch = i_epoch
    op_state.curr_batch = 0

    # invoke callbacks on iteration 0
    _invoke_callbacks(self, opts.callbacks, op_state, AbstractBatchCallback)

    for batch in eachbatch(data)
      load_data!(data, batch, data_arrays)
      load_label!(data, batch, label_arrays)

      # forward and backward
      for (texec, islice) in zip(train_execs, slices)
        forward(texec, is_train=true)

        # copy outputs into cpu ndarray, for evaluation metric
        for (cpu_out, dev_out) in zip(cpu_output_arrays, texec.outputs)
          copy!(slice(cpu_out, islice), dev_out)
        end

        backward(texec)
      end

      op_state.curr_iter  += 1
      op_state.curr_batch += 1

      # update parameters
      for idx = 1:length(param_names)
        if in(idx, freeze_idx)
          continue # Skip parameter update entirely
        end

        # gradient synchronization
        if !isa(kvstore, Void)
          # push gradient, priority is negative index
          push!(kvstore, idx, grad_arrays[idx], priority=-idx)
          if update_on_kvstore
            # pull back the weights
            pull!(kvstore, idx, param_arrays[idx], priority=-idx)
          else
            # pull back the sum-ed gradients, to the same locations
            pull!(kvstore, idx, grad_arrays[idx], priority=-idx)
          end
        end

        if !update_on_kvstore
          # manual updating
          for i_dev = 1:num_dev
            # create a fake index, so that the updater create states
            # for different param AND different devices, TODO(mli)
            # use a better solution later
            fake_idx = idx * num_dev + i_dev
            updater(fake_idx, grad_arrays[idx][i_dev], param_arrays[idx][i_dev])
          end
        end
      end

      # trigger learning rate decay
      opts.η_decay == :batch && update!(optimizer.η_sched)

      # invoke callbacks after finishing each iteration
      _invoke_callbacks(self, opts.callbacks, op_state, AbstractBatchCallback)

      # update evaluation metric on training set
      load_label!(data, batch, cpu_label_arrays)
      update!(opts.eval_metric, cpu_label_arrays, cpu_output_arrays)
    end # end of one epoch

    time_stop = time()
    metric = get(opts.eval_metric)
    opts.verbosity >= 2 && info(format("== Epoch {1:0>3d}/{2:0>3d} ==========", i_epoch, opts.n_epoch))
    if opts.verbosity >= 3
        info("## Training summary")
        for (name, value) in metric
            info(format("{1:>18s} = {2:.4f}", string(name), value))
        end
        info(format("{1:>18s} = {2:.4f} seconds", "time", time_stop-time_start))
    end

    # evaluation on validation set
    if !isa(opts.eval_data, Void)
      # because we are re-using the memory allocated for the training network,
      # the batch_size of the validation dataset must be the same as the training
      # batch_size
      @assert(get_batch_size(opts.eval_data) == batch_size)

      reset!(opts.eval_metric)
      for batch in eachbatch(opts.eval_data)
        load_data!(opts.eval_data, batch, data_arrays)

        # forward and backward
        for (texec, islice) in zip(train_execs, slices)
          forward(texec, is_train=true)

          # copy outputs into cpu ndarray, for evaluation metric
          for (cpu_out, dev_out) in zip(cpu_output_arrays, texec.outputs)
            copy!(slice(cpu_out, islice), dev_out)
          end
        end
        load_label!(opts.eval_data, batch, cpu_label_arrays)
        update!(opts.eval_metric, cpu_label_arrays, cpu_output_arrays)
      end

      if opts.verbosity >= 3
          info("## Validation summary")
          for (name, value) in get(opts.eval_metric)
            info(format("{1:>18s} = {2:.4f}", string(name), value))
          end
      end
    end

    if i_epoch == opts.n_epoch || any(x->isa(x, AbstractEpochCallback), opts.callbacks)
      # copy data back to cpu
      for (name, weights) in zip(param_names, param_arrays)
        # average parameters across devices
        weight = +([copy(w, cpu()) for w in weights]...) / length(weights)
        copy!(self.arg_params[name], weight)
      end
      for (name, aux_devs) in zip(aux_names, aux_arrays)
        aux_avg = +([copy(aux, cpu()) for aux in aux_devs]...) / length(aux_devs)
        copy!(self.aux_params[name], aux_avg)
      end
    end

    # trigger learning rate decay
    opts.η_decay == :epoch && update!(optimizer.η_sched)

    _invoke_callbacks(self, opts.callbacks, op_state, AbstractEpochCallback; metric=metric)
  end # end of all epochs

  opts.verbosity >= 1 && info("Finish training on $(self.ctx)")
  nothing
end

save_checkpoint(self::FeedForward, prefix::AbstractString, state::OptimizationState) =
  save_checkpoint(self.arch, self.arg_params, self.aux_params, prefix, state.curr_epoch)

function save_checkpoint(sym::SymbolicNode, arg_params::Dict{Symbol},
                         aux_params::Dict{Symbol}, prefix::AbstractString, epoch::Int)
  save("$prefix-symbol.json", sym)
  save_dict = Dict{Symbol, NDArray}(map((x) -> Symbol("arg:$(x[1])") => x[2], arg_params))
  if !isempty(aux_params)
    merge!(save_dict, Dict(map((x) -> Symbol("aux:$(x[1])") => x[2], aux_params)))
  end
  save_filename = format("{1}-{2:04d}.params", prefix, epoch)
  save(save_filename, save_dict)
  info("Saved checkpoint to '$save_filename'")
end

function load_checkpoint(prefix::AbstractString, epoch::Int)
  arch       = load("$prefix-symbol.json", SymbolicNode)
  saved_dict = load(format("{1}-{2:04d}.params", prefix, epoch), NDArray)
  arg_params = Dict{Symbol,Any}()
  aux_params = Dict{Symbol,Any}()
  for (k,v) in saved_dict
    tp, name = split(string(k), ':')
    name = Symbol(name)
    if tp == "arg"
      arg_params[name] = v
    else
      aux_params[name] = v
    end
  end

  return (arch, arg_params, aux_params)
end

"""
    load_checkpoint(prefix, epoch, ::mx.FeedForward; context)

Load a mx.FeedForward model from the checkpoint *prefix*, *epoch* and optionally provide a context.
"""
function load_checkpoint(prefix::AbstractString, epoch::Int, ::Type{FeedForward}; context = nothing)
  arch, arg_params, aux_params = load_checkpoint(prefix, epoch)
  model = FeedForward(arch, context = context)
  model.arg_params = arg_params
  model.aux_params = aux_params
  return model
end

function load_checkpoint(self::FeedForward, prefix::AbstractString, epoch::Int;
                         overwrite::Bool = true, allow_different_arch::Bool = false)
  if isdefined(self, :arg_params) && isdefined(self, :aux_params) && !overwrite
    info("model weights already exists, skip loading... (call with overwrite=true if needed)")
    return self
  end

  arch, arg_params, aux_params = load_checkpoint(prefix, epoch)
  if !allow_different_arch
    # TODO: is there better way to compare two symbols
    @assert(to_json(self.arch) == to_json(arch), "Cannot load from a checkpoint with different network architecture")
  end
  self.arg_params = arg_params
  self.aux_params = aux_params
  return self
end

max_stars_repo_name: Vikas-kum/incubator-mxnet

max_issues_repo_name: Vikas-kum/incubator-mxnet

max_forks_repo_name: Vikas-kum/incubator-mxnet


text:
# some auxilary functions
dist(p::Point, q::Point) = norm(p-q)
angle(p::Point, q::Point) = atan2(p.y - q.y, q.x - p.x)

# Center
# ------

# Calibrates the lens center for a camera+lens setup. Each combination of a lens and camera
# is assumed to be unique and every lens center should be calibrated.

# Luckily this is a straightforward procedure, taken from the
#[CAN-EYE user guide](http://www6.paca.inra.fr/can-eye/Documentation-Publications/Documentation).
# Dril 3-4 small holes in your lens cap (or prick an aluminium foil well attached to a lens
#ring) and take a picture for each rotation. Because coordinates of the dots in the pictures
# should lie on circle, the lens center can be calculated.

# The dataframe requires 3 columns with resp. the x-coordinates, corresponding y-coordinates
# and a tag (eg a number) for each dot/circle.
function calibrate_center(df::DataFrame, height::Int, width::Int)

    # First we check the dataframe input and create Nρ and Nϕ, resp. number of circles and
    # points per circle.
    @assert width >= height
    for name = [:x, :y, :circle]
        name in names(df) || error("Column $name not found in dataframe input.")
    end
    eltype(df[:x]) <: Real || throw(DomainError("Column :x should be numeric."))
    eltype(df[:y]) <: Real || throw(DomainError("Column :y should be numeric."))

    pool!(df, :circle)
    Nρ = length(levels(df[:circle]))
    @assert 3 <= Nρ <= 5

    all(by(df, :circle, length)[2]) > 0 || error("Different number of pictures per circle.")
    Nϕ = Int(nrow(df) / Nρ)
    @assert Nϕ * Nρ == nrow(df)

    @assert all(df[:x] .< width)
    @assert all(df[:y] .< height)

    # Create Matrix of Points with a column per circle.
    points = Point[]
    by(df, :circle) do subdf
        for row in eachrow(subdf)
            push!(points, Point(row[:x], row[:y]))
        end
    end
    points = reshape(points, (Nϕ, Nρ))

    # Create fit function as sum of squared errors.
    function sse(inp::Vector)
        # all parameters in single inp vector
        cx, cy = inp[1:2]
        c = Point(cx, cy)
        ρ = inp[3:(Nρ + 2)] #circle radius
        δ = inp[(3 + Nρ):(2Nρ + 2)] #angle offset between circles in single picture
        ϕ = inp[(3 + 2Nρ):end] #angle offset between pictures

        se = Float64[]
        for i = 1:Nρ
            # Normalize this squared error with squared typical radius because we sum upe
            # the errors for both radius and angles.
            push!(se, sum([(ρ[i] - dist(c,points[n, i]))^2 for n = 1:Nϕ])/(height^2/4))
        end
        for i = 1:Nϕ
            push!(se, sum([(ϕ[i] + δ[n] - angle(c,points[i, n]))^2 for n = 1:Nρ]))
        end
        sum(se)
    end

    # Estimate the starting point of the opimization from the inner circle.
    function initsse(c::Point, points::Matrix{Point})
        init = zeros(2 + 2Nρ + Nϕ)
        init[1:2] = [c.x, c.y]
        for i = 1:Nρ
            init[2+i] = dist(c, points[1,i])
            #offset
            init[2+Nρ+i] = angle(c, points[1,i]) - angle(c,points[1,1])
        end
        for i = 1:Nϕ
            init[2+2Nρ+i] = angle(c, points[i,1])
        end
        init
    end
    pic_center = Point(floor(Int, width/2),floor(Int, height/2))

    # We optimize using the Nelder-Mead downhill simplex algorithm,
    # the default for functions without a derivate explicetly defined.
    res = optimize(sse, initsse(pic_center, points), iterations=10_000)
    return lensx, lensy = round(Int, minimizer(res[1:2]))
end

# Projection Function
# -------------------

const N_LATERAL = 5 # number of points on side taken for Δ estimation

# TODO use lensx, lensy from calibrated lens center?
function calibrate_projfun(df::DataFrame, height::Int, width::Int)

    # First we check the dataframe inputs
    @assert width >= height
    for name = [:cm, :px, :H, :pos]
        @assert name in names(df)
    end
    eltype(df[:cm]) <: Real || throw(DomainError("Column :cm should be numeric."))
    eltype(df[:px]) <: Real || throw(DomainError("Column :px should be numeric."))

    @assert all(df[:px] .< width)

    #pool!(df, [:H, :pos])
    @assert length(levels(df[:H])) == 2
    @assert length(levels(df[:pos])) == 3
    for name = ["right", "left", "perpendicular"]
        @assert name in levels(df[:pos])
    end

    # We use 2 pictures at different distances H1, H2 from the perpendicular ruler to
    # estimate the distance to the perpendicular ruler more precisely than with a physical
    # ruler during setup.

    # We first estimate more correctly the distance difference Δ between im1 and im2 by
    # noticing that the shift on the lateral rulers between the two pictures is equal to Δ.
    @show H1, H2 = extrema(df[:H])

    c1 = df[(df[:pos] .== "perpendicular") & (df[:cm] .== 0) & (df[:H] .== H1), :px][1]
    c2 = df[(df[:pos] .== "perpendicular") & (df[:cm] .== 0) & (df[:H] .== H2), :px][1]

    # Utility function to interpolate a pixel from image 1 to the ruler of image 2
    function interpol(x, df2, center2)
        # suffix 2 is for second image with H2
        s2 = sort(df2, cols=:px)
        x2 = s2[:px] - center2
        cm2 = s2[:cm]

        if (x < minimum(x2)) | (x > maximum(x2))
            return NaN
        end

        ind = searchsortedfirst(x2, x)
        out = cm2[ind] - (x2[ind] - x) / (x2[ind] - x2[ind-1]) * (cm2[ind] - cm2[ind-1])
    end

    # fitting function to optimize Δ
    function fitΔ(Δ, df1, df2, center1, center2; errfun::Function=x->x^2)
        diffs = Float64[]
        for row in eachrow(df1)
            x = row[:px] - center1
            cm2 = interpol(x, df2, center2)
            !isnan(cm2) && push!(diffs, cm2 - row[:cm] - Δ)
        end
        return norm(diffs)
    end

    # starting point for optimization
    Δinit = H2 - H1

    left1 = head(sort!(df[(df[:pos] .== "left") & (df[:H] .== H1), :], cols=:cm, rev=true),
                N_LATERAL)
    left2 = head(sort!(df[(df[:pos] .== "left") & (df[:H] .== H2), :], cols=:cm, rev=true),
                N_LATERAL)
    resleft = optimize(Δ -> fitΔ(Δ, left1, left2, c1, c2), Δinit/2, Δinit * 3/2)

    right1 = head(sort!(df[(df[:pos] .== "right") & (df[:H] .== H1), :], cols=:cm, rev=true),
                N_LATERAL)
    right2 = head(sort!(df[(df[:pos] .== "right") & (df[:H] .== H2), :], cols=:cm, rev=true),
                N_LATERAL)
    resright = optimize(Δ -> fitΔ(Δ, right1, right2, c1, c2), Δinit/2, Δinit * 3/2)

    Δ = (minimizer(resleft) + minimizer(resright)) / 2
    @show Δ

    # With a more precise estimation of Δ, we can estimate H1 more precisely.
    # We optimize H1 for the value that provides the best fit between the two projection
    # functions of the two images. We fit a 2nd degree polynomial without intercept to the
    # projection functions ρ(θ) = aθ + bθ^2.
    # For higher order polynomial, simply adjust the `abfit` function accordingly.

    # perpendicular ruler half width
    w = maximum(df[df[:pos] .== "perpendicular", :cm ])

    function projfunpoints(df, center, H)
        θset = Float64[]
        ρset = Int[]
        for row in eachrow(df)
            cm = row[:cm]
            ρ = abs(row[:px] - center)
            θ = row[:pos] == "perpendicular" ? atan(abs(cm) / H) : atan(w / (H - cm))
            push!(θset, θ)
            push!(ρset, ρ)
        end
        θset, ρset
    end

    function abfit(θ, ρ)
    # fit ρ(θ) = aθ + bθ²
        A = [θᵢ^p for θᵢ in θ, p = 1:2]
        a, b = A \ ρ
    end

    # plot to spot outliers
    #function plotfit(H,c)
    #    lens = df[df[:H] .== H, :]
    #    θ, ρ = projfunpoints(lens, c, H)
    #    Winston.plot(θ, ρ, "o")
    #    a, b = abfit(θ, ρ)
    #    Winston.oplot(θ -> a*θ + b*θ.^2, 0, pi/2, "r")
    #    Winston.title("Observations and fit for $H")
    #end
    #display(plotfit(H1, c1)); display(plotfit(H2, c2))

    function projfundiff(H, Δ, lens1, lens2, center1, center2)
        a1, b1 = abfit(projfunpoints(lens1, center1, H)...)
        a2, b2 = abfit(projfunpoints(lens2, center2, H+Δ)...)
        (a1 - a2).^2 + (b1-b2).^2
    end

    lens1 = df[df[:H] .== H1, :]
    lens2 = df[df[:H] .== H2, :]
    Hres = optimize(H->projfundiff(H, Δ, lens1, lens2, c1, c2), H1 / 2, H1 * 2)
    H = minimizer(Hres)
    @show H

    a1, b1 = abfit(projfunpoints(lens1, c1, H)...)
    a2, b2 = abfit(projfunpoints(lens2, c2, H+Δ)...)

    a, b = (a1 + a2) / 2, (b1 + b2) / 2
    return @show a, b
end

ext: jl


text:
__precompile__()

module KerrSolver

using DifferentialEquations
using ODEInterfaceDiffEq

"""
    KerrHamiltonian(q, p, params)::Float64

This function defines the general Hamiltonian for Kerr Spacetime (rotating, uncharged black hole).

# Input
- `q::Array{Float64,1}`: Length-3 Array, containing the initial 3-Position
- `p::Array{Float64,1}`: Length-3 Array, containing the initial Covariant 3-Momentum
- `params::Array{Float64,1}`: Length-3 Array, containing -
                                Black Hole Spin Parameter, `a`,
                                Test Particle Energy, `E`,
                                Test Particle Rest Mass, `mu`

# Output
- `Float64`: Functional evaluation of the Hamiltonian, for use with `solveSystem()`

"""
function KerrHamiltonian(q, p, params)
    a = params[1]
    E = params[2]

    return -((a^4 * (E^2 - 2 * p[1]^2) - 8 * a * E * p[3] * q[1] -
        2 * q[1] * (p[2]^2 * (-2 + q[1]) + p[1]^2 * (-2 + q[1])^2 * q[1] - E^2 * q[1]^3) +
        a^2 * (2 * p[3]^2 - 2 * p[2]^2 - 4 * p[1]^2 * (-2 + q[1]) * q[1] +
        E^2 * q[1] * (2 + 3 * q[1])) + (a^2 + (-2 + q[1]) * q[1]) * (a^2 * E^2 * cos(2 * q[2]) -
        2 * p[3]^2 * csc(q[2])^2))/(4 * (a^2 + (-2 + q[1]) * q[1]) * (q[1]^2 +
        a^2 * cos(q[2])^2)))
end

"""
    solveSystem(state, params, end_lambda, step_size)::(String, Array{Float64,1}, Array{Float64,2})

This function solves for geodesics in Kerr Spacetime, using the Hamiltonian defined by
`KerrHamiltonian()`. A fixed-step, symplectic integrator, `VerletLeapfrog()`, is used
to integrate the system.

# Input
- `q::Array{Float64,1}`: Length-3 Array, containing the initial 3-Position
- `p::Array{Float64,1}`: Length-3 Array, containing the initial Covariant 3-Momentum
- `params::Array{Float64,1}`: Length-3 Array, containing -
                                Black Hole Spin Parameter, `a`,
                                Test Particle Energy, `E`,
                                Test Particle Rest Mass, `mu`
- `end_lambda::Float64`: Affine Parameter value, where integration will end
- `step_size::Float64`: Step Size (Fixed)

# Output
- `String`: Return status code of the solver
- `Array{Float64,1}`: Array, containing affine parameter values, where integration was performed
- `Array{Float64,2}`: 2D Array, containing integrated 3-Positions and Covariant 3-Momenta

"""
function solveSystem(q, p, params, end_lambda, step_size)
    a = params[1]
    outer_event_horizon = 1 + sqrt(1 - a ^ 2)

    # Affine Parameter values, where integration will be performed
    lambdas = (0., end_lambda)

    prob = HamiltonianProblem(KerrHamiltonian, q, p, lambdas, params)

    # Condition for terminating integration, when outer event horizon is reached
    reached_event_horizon(u,t,integrator) = u[1][1] - 1.01 * outer_event_horizon
    cb = ContinuousCallback(reached_event_horizon, terminate!)

    sol = solve(prob, VerletLeapfrog(), dt = step_size, callback=cb)

    # sol.u is of type, ArrayPartition
    # These operations are necessary for convenient numpy usage
    solu_flat = Array(transpose(reshape(collect(Iterators.flatten(sol.u)), (6, :))))

    return String(sol.retcode), sol.t, solu_flat
end

end # module


text:
################################################################################
# Statistics
################################################################################

mutable struct Statistics{T}
	iter::Int
	outer_iter::Vector{Int}
	t_elap::Vector{T}
	res::Vector{T}
	Δ_traj::Vector{T}
	dyn_vio::AbstractVector{DynamicsViolation}
	con_vio::AbstractVector{ControlViolation}
	sta_vio::AbstractVector{StateViolation}
	opt_vio::AbstractVector{OptimalityViolation}
end

function Statistics(;T=Float64)
	iter = 0
	outer_iter = Vector{Int}()
	t_elap = zeros(T,0)
	res = zeros(T, 0)
	Δ_traj = zeros(T,0)
	dyn_vio = Vector{DynamicsViolation}()
	con_vio = Vector{ControlViolation}()
	sta_vio = Vector{StateViolation}()
	opt_vio = Vector{OptimalityViolation}()
	return Statistics{T}(iter, outer_iter, t_elap, res, Δ_traj, dyn_vio, con_vio, sta_vio, opt_vio)
end

function record!(stats::Statistics, t_elap::T, res::T, Δ_traj::T, dyn_vio::DynamicsViolation,
	con_vio::ControlViolation, sta_vio::StateViolation, opt_vio::OptimalityViolation, k::Int) where {T}
	stats.iter += 1
	push!(stats.outer_iter, k)
	push!(stats.t_elap, t_elap)
	push!(stats.res, res)
	push!(stats.Δ_traj, Δ_traj)
	push!(stats.dyn_vio, dyn_vio)
	push!(stats.con_vio, con_vio)
	push!(stats.sta_vio, sta_vio)
	push!(stats.opt_vio, opt_vio)
	return nothing
end

function record!(stats::Statistics, prob, model::AbstractGameModel,
	game_con::GameConstraintValues, pdtraj::PrimalDualTraj, t_elap::T, Δ_traj::T, k::Int) where {T}

	stats.iter += 1
	push!(stats.outer_iter, k)
	push!(stats.t_elap, t_elap)
	push!(stats.res, residual_norm(prob, pdtraj))
	push!(stats.Δ_traj, Δ_traj)
	push!(stats.dyn_vio, dynamics_violation(model, pdtraj))
	push!(stats.con_vio, control_violation(game_con, pdtraj))
	push!(stats.sta_vio, state_violation(game_con, pdtraj))
	push!(stats.opt_vio, optimality_violation(prob.core))
	return nothing
end

function record!(stats::Statistics, prob, model::AbstractGameModel,
	game_con::GameConstraintValues, pdtraj::PrimalDualTraj, t_elap::T, Δ_traj::T, k::Int, i::Int) where {T}

	stats.iter += 1
	push!(stats.outer_iter, k)
	push!(stats.t_elap, t_elap)
	push!(stats.res, residual_norm(prob, pdtraj))
	push!(stats.Δ_traj, Δ_traj)
	push!(stats.dyn_vio, dynamics_violation(model, pdtraj, i))
	push!(stats.con_vio, control_violation(game_con, pdtraj, i))
	push!(stats.sta_vio, state_violation(game_con, pdtraj, i))
	push!(stats.opt_vio, optimality_violation(prob.core, i))
	return nothing
end

function reset!(stats::Statistics{T}) where {T}
	stats_ = Statistics(T=T)
	for name in fieldnames(Statistics)
		setfield!(stats, name, getfield(stats_, name))
	end
	return nothing
end


text:
function init_4(t,
		fft_u0, fft_v0,
		A1, A2,
		matr, conjmatr,
                sigma, llambda,
		Ktaubis, epsilon, Ntaumm)

    h1u, h1v = ftau(t, fft_u0, fft_v0, A1, A2, matr, conjmatr, sigma, llambda)

    h1u_fft = fft(h1u, dims=1)
    h1v_fft = fft(h1v, dims=1)

    dtu0u = h1u_fft[1, :] / Ntaumm
    dtu0v = h1v_fft[1, :] / Ntaumm

    h1u_fft[1, :] .= 0
    h1v_fft[1, :] .= 0

    h1u = epsilon * ifft(h1u_fft / (1im * Ktaubis), dims=1)
    h1v = epsilon * ifft(h1v_fft / (1im * Ktaubis), dims=1)

    w1u = fft_u0 .- h1u[1, :]
    w1v = fft_v0 .- h1v[1, :]

    h1u, h1v = ftau(t, w1u, w1v, A1, A2, matr, conjmatr, sigma, llambda)

    h1u_fft = fft(h1u, dims=1)
    h1v_fft = fft(h1v, dims=1)

    dtu0u = h1u_fft[1, :] / Ntaumm
    dtu0v = h1v_fft[1, :] / Ntaumm

    h1u_fft[1, :] .= 0
    h1v_fft[1, :] .= 0

    h1u = epsilon * ifft(h1u_fft / (1im * Ktaubis), dims=1)
    h1v = epsilon * ifft(h1v_fft / (1im * Ktaubis), dims=1)

    champu1, champv1 = duftau(t, w1u, w1v, dtu0u, dtu0v,
                              A1, A2, matr, conjmatr, sigma, llambda, Ntaumm)

    champu2, champv2 = dtftau(t, w1u, w1v, A1, A2, matr,
                              conjmatr, sigma, llambda, Ntaumm)

    champu = champu1 .+ champu2
    champv = champv1 .+ champv2

    champu_fft = fft(champu, dims=1)
    champv_fft = fft(champv, dims=1)

    champu_fft[1, :] .= 0
    champv_fft[1, :] .= 0

    dth1u = epsilon * ifft(champu_fft / (1im * Ktaubis), dims=1)
    dth1v = epsilon * ifft(champv_fft / (1im * Ktaubis), dims=1)

    h2u, h2v = ftau(t, w1u + h1u, w1v + h1v, A1, A2, matr, conjmatr, sigma,
                    llambda)

    h2u = h2u - dth1u
    h2v = h2v - dth1v

    h2u_fft = fft(h2u, dims=1)
    h2v_fft = fft(h2v, dims=1)

    h2u_fft[1, :] .= 0
    h2v_fft[1, :] .= 0

    h2u = epsilon * ifft(h2u_fft / (1im * Ktaubis), dims=1)
    h2v = epsilon * ifft(h2v_fft / (1im * Ktaubis), dims=1)

    w2u = fft_u0 .- h2u[1, :]
    w2v = fft_v0 .- h2v[1, :]

    h1u, h1v = ftau(t, w2u, w2v, A1, A2, matr, conjmatr, sigma, llambda)

    h1u_fft = fft(h1u, dims=1)
    h1v_fft = fft(h1v, dims=1)

    dtu0u = h1u_fft[1, :] / Ntaumm
    dtu0v = h1v_fft[1, :] / Ntaumm

    h1u_fft[1, :] .= 0
    h1v_fft[1, :] .= 0

    h1u = epsilon * ifft(h1u_fft / (1im * Ktaubis), dims=1)
    h1v = epsilon * ifft(h1v_fft / (1im * Ktaubis), dims=1)

    champu1, champv1 = duftau(t, w2u, w2v, dtu0u, dtu0v,
                              A1, A2, matr, conjmatr, sigma, llambda, Ntaumm)

    champu2, champv2 = dtftau(t, w2u, w2v, A1, A2, matr,
                              conjmatr, sigma, llambda, Ntaumm)

    champu = champu1 + champu2
    champv = champv1 + champv2

    champu_fft = fft(champu, dims=1)
    champv_fft = fft(champv, dims=1)

    dttu0u = champu_fft[1, :] / Ntaumm
    dttu0v = champv_fft[1, :] / Ntaumm

    champu_fft[1, :] .= 0
    champv_fft[1, :] .= 0

    dth1u = epsilon * ifft(champu_fft / (1im * Ktaubis), dims=1)
    dth1v = epsilon * ifft(champv_fft / (1im * Ktaubis), dims=1)

    champu, champv = ftau(t, w2u + h1u, w2v + h1v, A1, A2, matr, conjmatr,
                          sigma, llambda)

    h2u = champu - dth1u
    h2v = champv - dth1v

    h2u_fft = fft(h2u, dims=1)
    h2v_fft = fft(h2v, dims=1)

    champu_fft = fft(champu, dims=1)
    champv_fft = fft(champv, dims=1)

    dtu1u = champu_fft[1, :] / Ntaumm
    dtu1v = champv_fft[1, :] / Ntaumm

    h2u_fft[1, :] .= 0
    h2v_fft[1, :] .= 0

    h2u = epsilon * ifft(h2u_fft / (1im * Ktaubis), dims=1)
    h2v = epsilon * ifft(h2v_fft / (1im * Ktaubis), dims=1)

    champu1, champv1 = d2uftau(t, w2u, w2v, dtu0u, dtu0v, dtu0u, dtu0v,
                               A1, A2, matr, conjmatr, sigma, llambda, Ntaumm)
    champu2, champv2 = dtduftau(t, w2u, w2v, dtu0u, dtu0v,
                                A1, A2, matr, conjmatr, sigma, llambda, Ntaumm)
    champu3, champv3 = duftau(t, w2u, w2v, dttu0u, dttu0v,
                              A1, A2, matr, conjmatr, sigma, llambda, Ntaumm)
    champu4, champv4 = d2tftau(t, w2u, w2v,
                               A1, A2, matr, conjmatr, sigma, llambda, Ntaumm)

    champu = champu1 + 2 * champu2 + champu3 + champu4
    champv = champv1 + 2 * champv2 + champv3 + champv4

    champu_fft = fft(champu, dims=1)
    champv_fft = fft(champv, dims=1)

    champu_fft[1, :] .= 0
    champv_fft[1, :] .= 0

    dtth1u = epsilon * ifft(champu_fft / (1im * Ktaubis), dims=1)
    dtth1v = epsilon * ifft(champv_fft / (1im * Ktaubis), dims=1)

    champu1, champv1 = duftau(t, w2u + h1u, w2v + h1v,
                              dtu1u + dth1u, dtu1v + dth1v, A1, A2, matr,
                              conjmatr, sigma, llambda, Ntaumm)

    champu2, champv2 = dtftau(t, w2u + h1u, w2v + h1v, A1, A2, matr,
                              conjmatr, sigma, llambda, Ntaumm)

    champu = champu1 + champu2 - dtth1u
    champv = champv1 + champv2 - dtth1v

    champu_fft = fft(champu, dims=1)
    champv_fft = fft(champv, dims=1)

    champu_fft[1, :] .= 0
    champv_fft[1, :] .= 0

    dth2u = epsilon * ifft(champu_fft / (1im * Ktaubis), dims=1)
    dth2v = epsilon * ifft(champv_fft / (1im * Ktaubis), dims=1)

    h3u, h3v = ftau(t, w2u + h2u, w2v + h2v, A1, A2, matr, conjmatr, sigma,
                    llambda)

    h3u = h3u - dth2u
    h3v = h3v - dth2v

    h3u_fft = fft(h3u, dims=1)
    h3v_fft = fft(h3v, dims=1)

    h3u_fft[1, :] .= 0
    h3v_fft[1, :] .= 0

    h3u = epsilon * ifft(h3u_fft / (1im * Ktaubis), dims=1)
    h3v = epsilon * ifft(h3v_fft / (1im * Ktaubis), dims=1)

    h2u, h2v = ftau(t, w2u + h1u, w2v + h1v,
		    A1, A2,
		    matr, conjmatr,
		    sigma, llambda)

    h2u = h2u - dth1u
    h2v = h2v - dth1v

    h2u_fft = fft(h2u, dims=1)
    h2v_fft = fft(h2v, dims=1)

    h2u_fft[1, :] .= 0
    h2v_fft[1, :] .= 0

    h2u = epsilon * ifft(h2u_fft / (1im * Ktaubis), dims=1)
    h2v = epsilon * ifft(h2v_fft / (1im * Ktaubis), dims=1)

    champu1, champv1 = d2uftau(t, w2u, w2v, dtu0u, dtu0v, dtu0u, dtu0v,
                               A1, A2, matr, conjmatr, sigma, llambda, Ntaumm)

    champu2, champv2 = dtduftau(t, w2u, w2v, dtu0u, dtu0v,
                                A1, A2, matr, conjmatr, sigma, llambda, Ntaumm)

    champu3, champv3 = duftau(t, w2u, w2v, dttu0u, dttu0v,
                              A1, A2, matr, conjmatr, sigma, llambda, Ntaumm)

    champu4, champv4 = d2tftau(t, w2u, w2v,
                               A1, A2, matr, conjmatr, sigma, llambda, Ntaumm)

    champu = champu1 + 2 * champu2 + champu3 + champu4
    champv = champv1 + 2 * champv2 + champv3 + champv4

    champu_fft = fft(champu, dims=1)
    champv_fft = fft(champv, dims=1)

    champu_fft[1, :] .= 0
    champv_fft[1, :] .= 0

    dtth1u = epsilon * ifft(champu_fft / (1im * Ktaubis), dims=1)
    dtth1v = epsilon * ifft(champv_fft / (1im * Ktaubis), dims=1)

    champu1, champv1 = duftau(t, w2u + h1u, w2v + h1v,
                              dtu0u + dth1u, dtu0v + dth1v, A1, A2, matr,
                              conjmatr, sigma, llambda, Ntaumm)

    champu2, champv2 = dtftau(t, w2u + h1u, w2v + h1v, A1, A2, matr,
                              conjmatr, sigma, llambda, Ntaumm)

    champu = champu1 + champu2 - dtth1u
    champv = champv1 + champv2 - dtth1v

    champu_fft = fft(champu, dims=1)
    champv_fft = fft(champv, dims=1)

    champu_fft[1, :] .= 0
    champv_fft[1, :] .= 0

    dth2u = epsilon * ifft(champu_fft / (1im * Ktaubis), dims=1)
    dth2v = epsilon * ifft(champv_fft / (1im * Ktaubis), dims=1)

    h3u, h3v = ftau(t, w2u + h2u, w2v + h2v, A1, A2, matr, conjmatr, sigma,
                    llambda)

    h3u = h3u - dth2u
    h3v = h3v - dth2v

    h3u_fft = fft(h3u, dims=1)
    h3v_fft = fft(h3v, dims=1)

    h3u_fft[1, :] .= 0
    h3v_fft[1, :] .= 0

    h3u = epsilon * ifft(h3u_fft / (1im * Ktaubis), dims=1)
    h3v = epsilon * ifft(h3v_fft / (1im * Ktaubis), dims=1)

    fft_ubar = fft_u0 - h3u[1, :]
    fft_vbar = fft_v0 - h3v[1, :]

    C3u, C3v = C3(t, fft_ubar, fft_vbar, A1, A2,
                  matr, conjmatr, sigma, llambda,
                  Ktaubis, epsilon, Ntaumm)

    fft_ug = fft_u0 - C3u[1, :]
    fft_vg = fft_v0 - C3v[1, :]

    fft_ubar, fft_vbar, fft_ug, fft_vg

end


text:
export Model2Fit_Mapped_Parameters
export get_model_θ, get_model

@doc raw"""

Create a new model, where some parameters are computed using a
[`Abstract_Map`](@ref).

TODO: not clear how to use it: certainly need a refactoring...

# Also see
- [`Model2Fit_Shared_Parameters`](@ref)
"""
struct Model2Fit_Mapped_Parameters{MODEL <: Abstract_Model2Fit,
                                   MAP <: Abstract_Map,
                                   INDICES <: AbstractVector{Int},
                                   EL2MAP <: AbstractVector} <: Abstract_Model2Fit
    _model::MODEL
    _map::MAP
    _indices::INDICES
    _elements2map::EL2MAP

    function Model2Fit_Mapped_Parameters(model::MODEL,
                                         map::MAP,
                                         indices::INDICES,
                                         elements2map::EL2MAP) where {MODEL <: Abstract_Model2Fit,
                                                                      MAP <: Abstract_Map,
                                                                      INDICES <: AbstractVector{Int},
                                                                      EL2MAP <: AbstractVector}

        @assert parameter_size(model) ≥ length(indices)
        @assert length(indices) == length(elements2map)

        new{MODEL,MAP,INDICES,EL2MAP}(model,
                                      map,
                                      indices,
                                      elements2map)
    end

end

# Visit  ================
#
visit_submodel_size(model::Model2Fit_Mapped_Parameters) = 1

function _visit_get_submodel(model::Model2Fit_Mapped_Parameters)
    model._model
end
function visit_get_submodel(model::Model2Fit_Mapped_Parameters,submodel_idx::Int)
    @assert 1 ≤ submodel_idx ≤ visit_submodel_size(model)

    _visit_get_submodel(model)
end

function _visit_get_X(model::Model2Fit_Mapped_Parameters,X::AbstractVector)
    X
end
function visit_get_X(model::Model2Fit_Mapped_Parameters,submodel_idx::Int,Y::AbstractVector,X::AbstractVector,θ::AbstractVector)
    _visit_get_X(model,X)
end

function _visit_get_Y(model::Model2Fit_Mapped_Parameters,Y::AbstractVector)
    Y
end
function visit_get_Y(model::Model2Fit_Mapped_Parameters,submodel_idx::Int,Y::AbstractVector,X::AbstractVector,θ::AbstractVector)
    _visit_get_Y(model,Y)
end

function _visit_get_θ(model::Model2Fit_Mapped_Parameters,θ::AbstractVector)
    θ_model_reduced = @view θ[1:(parameter_size(model._model) - length(model._indices))]
    θ_model_reduced_size = length(θ_model_reduced)

    θ_map = @view θ[(θ_model_reduced_size+1):end]
    mapped_elements = eval_map(model._map,model._elements2map, θ_map)

    insert_some_elements(θ_model_reduced, model._indices,mapped_elements)
end

function visit_get_θ(model::Model2Fit_Mapped_Parameters,submodel_idx::Int,Y::AbstractVector,X::AbstractVector,θ::AbstractVector)
    _visit_get_θ(model,θ)
end

# Interface ================
#
parameter_size(mp::Model2Fit_Mapped_Parameters) = parameter_size(mp._model) - length(mp._indices) + parameter_size(mp._map)

function accumulate_y!(mp::Model2Fit_Mapped_Parameters,Y::AbstractVector,X::AbstractVector,θ::AbstractVector)
    θ_model = _visit_get_θ(mp,θ)
    model = _visit_get_submodel(mp)

    accumulate_y!(model,Y,X,θ_model)
end

# Convenience methods ================
#
@doc raw"""
```julia
get_model(mp::Model2Fit_Mapped_Parameters) -> Absatrct_Model2Fit
```

Get back the wrapped model
"""
get_model(mp::Model2Fit_Mapped_Parameters) = _visit_get_submodel(mp)

@doc raw"""
```julia
get_model_θ(mp::Model2Fit_Mapped_Parameters,θ::AbstractVector) -> θ::AbstractVector
```

Retrieve the parameter vector θ associated to the wrapped model [`get_model`](@ref).
"""
get_model_θ(mp::Model2Fit_Mapped_Parameters,θ::AbstractVector) = _visit_get_θ(mp,θ)


text:
using Test, Seis, SeisTau

@testset "SeisTau" begin
    @testset "Travel time" begin
        let t = Trace(0, 1, 2)
            @testset "Phase supplied" begin
                @test_throws ArgumentError travel_time(t, "P")
                t.evt.lon, t.evt.lat, t.evt.dep = 0, 0, 0
                t.sta.lon, t.sta.lat = 10, 10
                @test travel_time(t, "P") == travel_time(t.evt, t.sta, "P")
                arr = travel_time(t, "P")
                @test length(arr) == 1
                @test arr[1].time ≈ 200.405669 atol=0.0001
            end
            @testset "Using picks" begin
                @test travel_times(t) == []
                t.picks.P = (0, "P")
                @test travel_times(t) == [travel_time(t, "P")]
                t′ = deepcopy(t)
                # Ignore unnamed picks
                t′.picks.X = 1
                @test travel_times(t′) == travel_times(t)
                # Ignore empty but not missing names
                t′.picks.S = 2, ""
                @test travel_times(t′) == travel_times(t)
                # Ignore picks with the wrong kind of name
                t′.picks.sS = 3, "asS"
                @test travel_times(t′) == travel_times(t)
            end
            @testset "Default phases" begin
                @test travel_time(t) == travel_time(t, "ttall")
            end
        end
    end

    @testset "Path" begin
        let t = Trace(0, 1, 2)
            @test_throws ArgumentError path(t)
            @test_throws ArgumentError path(t.evt, t.sta)
            t.evt.lon, t.evt.lat, t.evt.dep = 0, 0, 0
            t.sta.lon, t.sta.lat = 0, 80
            @test path(t) == path(t, "ttall")
            @test path(t) == path(t.evt, t.sta)
            arr = path(t, "P")
            @test length(arr) == 1
            @test arr[1].name == "P"
            @test !isempty(arr[1].radius)
            @test !isempty(arr[1].lon)
        end
    end

    @testset "Picks" begin
        # Adding travel time picks from TauPy
        let t = Trace(0, 1, rand(2))
            # Can't add travel times without correct geometry
            @test_throws ArgumentError add_picks!(t, "S")
            t.evt.lon, t.evt.lat, t.evt.dep = 0, 0, 0
            t.sta.lon, t.sta.lat = 10, 10
            add_picks!(t, "Sn", model="ak135")
            @test picks(t)[end].time ≈ 358.4306 atol=0.0001
            @test picks(t)[end].name == "Sn"
        end
    end

    @testset "Sphere" begin
        # Calculating geographic travel times using elliptical and spherical Earth
        let t = Trace(0, 1, 2)
            t.evt.lon = t.evt.lat = t.sta.lon = t.evt.dep = 0
            t.sta.lat = 45
            tt = travel_time(t, "P") # WGS84 ellipsoid
            @test tt[1].delta ≈ 44.78 atol=0.01
            tt′ = travel_time(t, "P", sphere=true) # Sphere
            @test tt′[1].delta == 45.0
        end
    end

    @testset "Checking" begin
        let e = Seis.Event(), s = Seis.Station(), t = Trace(0, 1, 2)
            t.evt = e
            t.sta = s
            @test_throws ArgumentError SeisTau._check_headers_taup(e, s)
            @test_throws ArgumentError SeisTau._check_headers_taup(t)
            e.lon, e.lat, e.dep = 0, 0, 0
            s.lon, s.lat = 10, 10
            @test SeisTau._check_headers_taup(e, s) === nothing
            @test SeisTau._check_headers_taup(t) === nothing
        end
    end
end


text:
using Distributions, DelimitedFiles, ProgressMeter, Combinatorics

####################################################################################
#######################  Print Results of TEAVAR formulation  ########################
####################################################################################

function printResults(o, alpha, a, u, umax, edges, scenarios, T, Tf, L, capacity, p, demand; verbose=false, utilization=true)
    println("Objective value: ", o)
    println("Alpha: ", alpha)
    print("\n")
    println("------------------ Allocations ----------------------\n")
    for i in 1:size(a,1)
    	println("Flow ", i, " demand= ", demand[i])
        for j in 1:size(a,2)
            println("Flow ",i, ", tunnel ", j, " allocated : ", a[i,j])
            print("Edges in use: ")
            for e in T[Tf[i][j]]
                print(edges[e])
            end
            println("\n")
        end
    end
    if verbose && false # too many scenarios in some cases
        println("--------------- Loss Breakdown ---------------------\n")
        for s in 1:size(umax,1)
            println("Scenario ", s, ": ", scenarios[s], " prob: ", p[s])
            print("Down edges: ")
            for i in 1:size(scenarios[s],1)
                if scenarios[s][i] == 0.0
                    print(edges[i], " ")
                end
            end
	    println("")

            for f in 1:size(u,2)
	        if u[s,f] > alpha
                    println("Loss on flow ", f, " = ", u[s,f]-alpha)
		end
            end
            println("umax[", s, "] = ", umax[s])
            println("\n")
        end
    end
    println("------------------Utilization-------------------\n")
    if utilization
        for e in 1:size(edges,1)
            s = 0
            for f in 1:size(a,1)
                for t in 1:size(a,2)
                    s += a[f,t] * L[Tf[f][t],e]
                end
            end
            print("EDGE values: ", e, " : ", edges[e], " cap: ", capacity[e], " alloc: ", s)
	    if s > 0.99 * capacity[e]
	    	print(" <-- near cap")
	    end
            println("")
        end
    end
    # writedlm("inputs/input3/allocations.txt", getvalue(a))
end

####################################################################################
###################  Compute all possible scenario bitmaps  ########################
####################################################################################

function kScenarios(nedges, k, probabilities; first=true)
    scenarios = []
    if first
        scenario = ones(nedges)
        push!(scenarios, scenario)
    end
    for i in 1:k
        for bits in collect(combinations(1:nedges,i))
            s = ones(nedges)
            for bit in bits
                s[bit] = 0
            end
            push!(scenarios, s)
        end
    end
    probs = getProbabilities(scenarios, probabilities)
    return scenarios, probs
end

function allScenarios(nedges, probabilities; first=true)
    scenarios = []
    probs = []
    if first
        scenario = ones(nedges)
        push!(scenarios, scenario) #ADD SCENARIO NO FAILURES
        prob = 1
        for i in 1:length(scenario)
            prob *= (1 - scenario[i]) * probabilities[i] + scenario[i] * (1 - probabilities[i])
        end
        push!(probs, prob)
    end
    p = Progress(nedges, .1, "Computing all scenarios...", 50)
    for i in 1:nedges
        for bits in collect(combinations(1:nedges,i))
            s = ones(nedges)
            for bit in bits
                s[bit] = 0
            end
            prob = 1
            for i in 1:length(s)
                prob *= (1 - s[i]) * probabilities[i] + s[i] * (1 - probabilities[i])
            end
            push!(probs, prob)
            push!(scenarios, s)
        end
        next!(p, showvalues = [(:edges,"$(i)/$(nedges)")])
    end
    return scenarios, probs ./ sum(probs)
end

####################################################################################
####################  Get probabilities of all scenarios  ##########################
####################################################################################

function getProbabilities(scenarios, probabilities)
    nscenarios = size(scenarios, 1)
    p = []
    for s in 1:nscenarios
        prob = 1
        for i in 1:size(scenarios[s],1)
            prob *= (1 - scenarios[s][i]) * probabilities[i] + scenarios[s][i] * (1 - probabilities[i])
        end
        push!(p, prob)
    end
    return p
end

####################################################################################
####################  Compute all scenarios above a threshold  #####################
####################################################################################

function subScenariosRecursion(original, cutoff, remaining=[], offset=0, partial=[], scenarios=[], probabilities=[]; progress=nothing)
    if (size(partial,1) == 0)   #first
        push!(scenarios, ones(size(original,1)))
        push!(probabilities, prod(1 .- original))
        remaining = original
    else (size(partial,1) > 0)
        probs = 1 .- original
        bitmap = ones(size(original, 1))   #create bitmap
        for index in partial
            probs[index] = original[index]
            bitmap[index] = 0
        end
        product = prod(probs)
        if progress != nothing
            ProgressMeter.next!(progress, showvalues = [(:cutoff,cutoff), (:scenarios_added,length(probabilities)), (:last, product)])
        end
        if product >= cutoff
            push!(scenarios, bitmap)
            push!(probabilities, product)
        else
            return
        end
    end
    yield()

    for i in 1:size(remaining,1)
        offset = size(original,1) - size(remaining,1)
        n = offset + i
        subScenariosRecursion(original, cutoff, remaining[i+1:end], offset, vcat(partial, [n]), scenarios, probabilities, progress=progress)
    end
    return scenarios, probabilities
end

function subScenarios(original, cutoff; first=true, last=true, progress=true)
    p = ProgressMeter.ProgressUnknown("Computing scenarios cutoff=$(cutoff)...")
    if progress
        scenarios, probabilities = subScenariosRecursion(original, cutoff, progress=p)
    else
        scenarios, probabilities = subScenariosRecursion(original, cutoff)
    end
    if first == false
        scenarios = scenarios[2:end]
        probabilities = probabilities[2:end]
    end
    if last
        push!(scenarios, zeros(length(scenarios[1])))
        push!(probabilities, 1 - sum(probabilities))
    end
    if sum(probabilities) < 1
        probabilities = probabilities ./ sum(probabilities)
    end
    ProgressMeter.finish!(p)
    return scenarios, probabilities
end

####################################################################################
####################  Weibull Distribution probabilities  #######################
####################################################################################

function weibullProbs(num; shape=.8, scale=.0001)
    w = Distributions.Weibull(shape, scale)
    probs = []
    for i in 1:num
        push!(probs, rand(w))
    end
    return probs
end


text:
"Convert Tuple to Min Max Named Tuple"
Base.convert(::Type{Min_Max}, input::Tuple{Float64, Float64}) =
    (min = input[1], max = input[2])

"Convert Tuple to Up Down Named Tuple"
Base.convert(::Type{Up_Down}, input::Tuple{Float64, Float64}) =
    (up = input[1], down = input[2])


text:
#=
Copyright (c) 2015, Intel Corporation

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

    * Redistributions of source code must retain the above copyright notice,
      this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name of Intel Corporation nor the names of its contributors
      may be used to endorse or promote products derived from this software
      without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
=#

using Latte
using FactCheck

function pooling_forward(input, mask, output, kernel, stride, pad, _type)
  width, height, channels, num = size(input)
  pooled_width, pooled_height, _, _ = size(output)
  kernel_size = kernel * kernel

  for n = 1:num
    for c = 1:channels
      for ph = 1:pooled_height
        for pw = 1:pooled_width
          hstart = (ph-1)*stride - pad + 1
          wstart = (pw-1)*stride - pad + 1
          hend = min(hstart + kernel - 1, height)
          wend = min(wstart + kernel - 1, width)
          hstart = max(1, hstart)
          wstart = max(1, wstart)

          region = sub(input, wstart:wend, hstart:hend, c, n)
          if _type == :max
              index = indmax(region)
              mask[pw, ph, c, n] = index # note this is local index in region
              output[pw, ph, c, n] = region[index]
          elseif _type == :mean
              output[pw, ph, c, n] = sum(region) / kernel_size
          end
          # if isa(state.layer.pooling, Pooling.Max)
          #   index = indmax(region)
          #   mask[pw, ph, c, n] = index # note this is local index in region
          #   output[pw, ph, c, n] = region[index]
          # elseif isa(state.layer.pooling, Pooling.Mean)
          #   output[pw, ph, c, n] = sum(region) / kernel_size
          # else
          #   error("Unknown pooling $(state.layer.pooling)")
          # end
        end
      end
    end
  end
end

function pooling_backward(gradient, mask, diff, kernel, stride, pad, _type)
  width, height, channels, num = size(gradient)
  pooled_width, pooled_height, _ = size(diff)
  kernel_size = kernel * kernel

  for n = 1:num
    for c = 1:channels
      for ph = 1:pooled_height
        for pw = 1:pooled_width
          hstart = (ph-1)*stride - pad + 1
          wstart = (pw-1)*stride - pad + 1
          hend = min(hstart + kernel - 1, height)
          wend = min(wstart + kernel - 1, width)
          hstart = max(1, hstart)
          wstart = max(1, wstart)

          region = sub(gradient, wstart:wend, hstart:hend, c, n)
          if _type == :max
              index = mask[pw, ph, c, n]
              region[index] += diff[pw, ph, c, n]
          elseif _type == :mean
              region[:] += diff[pw, ph, c, n] / kernel_size
          end
          # if isa(state.layer.pooling, Pooling.Max)
          #   index = payload[pw, ph, c, n]
          #   region[index] += diff[pw, ph, c, n]
          # elseif isa(state.layer.pooling, Pooling.Mean)
          #   region[:] += diff[pw, ph, c, n] / kernel_size
          # else
          #   error("Unknown pooling $(state.layer.pooling)")
          # end
        end
      end
    end
  end
end

facts("Testing MaxPooling Layer") do
    for (h, w) in [(24, 24), (2, 2)]
        net = Net(8)
        data,  data_value   = MemoryDataLayer(net, :data, (w, h, 3))
        label, label_value = MemoryDataLayer(net, :label, (1,))
        data_value[:]  = rand(Float32, size(data_value)...) * 256
        label_value[:] = map(floor, rand(Float32, size(label_value)...) * 10)
        conv1        = ConvolutionLayer(:conv1, net, data, 2, 3, 1, 1)
        pool1        = MaxPoolingLayer(:pool1, net, conv1, 2, 2, 0)
        fc1          = InnerProductLayer(:fc1, net, pool1, 10)
        loss         = SoftmaxLossLayer(:loss, net, fc1, label)

        init(net)

        params = SolverParameters(
            lr_policy    = LRPolicy.Decay(.01f0, 5.0f-7),
            mom_policy   = MomPolicy.Fixed(0.9),
            max_epoch    = 300,
            regu_coef    = .0005)
        sgd = SGD(params)

        input    = get_buffer(net, :conv1value)
        mask     = get_buffer(net, :pool1maxidx)

        context("Forward input shape=($h x $w)") do
            forward(net; solver=sgd)

            mask_expected = zeros(mask)
            expected = zeros(get_buffer(net, :pool1value))
            pooling_forward(input, mask_expected, expected, 2, 2, 0, :max)
            @fact expected --> roughly(get_buffer(net, :pool1value))
            @fact mask_expected --> mask
        end

        context("Backward input shape=($h x $w)") do
            top_diff = get_buffer(net, :pool1∇)

            ∇input = get_buffer(net, :conv1∇)
            ∇input_expected = deepcopy(∇input)

            backward(net)
            pooling_backward(∇input_expected, mask, top_diff, 2, 2, 0, :max)
            @fact ∇input   --> roughly(∇input_expected)
        end
    end
end

facts("Testing MeanPooling Layer") do
    net = Net(8)
    data,  data_value   = MemoryDataLayer(net, :data, (227, 227, 3))
    label, label_value = MemoryDataLayer(net, :label, (1,))
    data_value[:]  = rand(Float32, size(data_value)...) * 256
    label_value[:] = map(floor, rand(Float32, size(label_value)...) * 10)
    conv1        = ConvolutionLayer(:conv1, net, data, 2, 3, 1, 1)
    pool1        = MeanPoolingLayer(:pool1, net, conv1, 2, 2, 0)
    fc1          = InnerProductLayer(:fc1, net, pool1, 10)
    loss         = SoftmaxLossLayer(:loss, net, fc1, label)

    init(net)

    params = SolverParameters(
        lr_policy    = LRPolicy.Decay(.01f0, 5.0f-7),
        mom_policy   = MomPolicy.Fixed(0.9),
        max_epoch    = 300,
        regu_coef    = .0005)
    sgd = SGD(params)

    input    = get_buffer(net, :conv1value)

    context("Forward") do
        forward(net; solver=sgd)

        expected = zeros(get_buffer(net, :pool1value))
        pooling_forward(input, [], expected, 2, 2, 0, :mean)
        @fact expected --> roughly(get_buffer(net, :pool1value))
    end

    context("Backward") do
        top_diff = get_buffer(net, :pool1∇)

        ∇input = get_buffer(net, :conv1∇)
        ∇input_expected = deepcopy(∇input)

        backward(net)
        pooling_backward(∇input_expected, [], top_diff, 2, 2, 0, :mean)
        @fact ∇input   --> roughly(∇input_expected)
    end
end

FactCheck.exitstatus()


text:
using SubsetSelection
using SubsetSelectionCIO
using Test
using StatsBase

n = 500; p = 1000; k = 10;
indices = sort(sample(1:p, StatsBase.Weights(ones(p)/p), k, replace=false));
w = sample(-1:2:1, k);
X = randn(n,p); Y = X[:,indices]*w;

γ = 1/sqrt(size(X,1))

indices0, w0, Δt, status, Gap, cutCount = oa_formulation(SubsetSelection.OLS(), Y, X, k, γ)
for i in 1:k
  @test indices0[i]==indices[i]
end


text:
module TestMin
    using StreamStats
    using Distributions
    using Base.Test

    # Min of uniform draws
    for n in rand(1:1_000_000, 100)
        xs = rand(n)
        stat = StreamStats.Min()
        for x in xs
            update!(stat, x)
        end
        online_m = minimum(stat)
        online_s = state(stat)
        online_n = nobs(stat)
        batch_m = minimum(xs)
        @test online_m == batch_m
        @test online_n == n
    end
end


text:
# given a policy, these utils play out a game
export play_n_game, greedy_player, play_via_monte_carlo_w_recorder, play_via_monte_carlo_wo_recorder

using Base.Threads

function play_one_move_with_policy(bitboard, policy)
    probs = policy(bitboard)

    newboards = move.(Ref(bitboard), DIRS)

    for (i, newboard) in enumerate(newboards)
        if newboard == bitboard
            # it means the move was not possible
            # set the probability to 0
            probs[i] = 0
        end
    end

    tot_prob = sum(probs)

    while !(tot_prob ≈ 0.0)
        for i in 1:4
            if (probs[i] != 0) && (rand() < probs[i]/tot_prob)
                return newboards[i], DIRS[i]
            end
            tot_prob -= probs[i]
        end

        # if no move got selected. this could just be unlucky; just keep going
        tot_prob = sum(probs)
    end

    # nothing got played
    return bitboard, left
end

if false
    # @code_warntype play_one_move_with_policy(initbboard(), randompolicy)
    # @benchmark play_one_move_with_policy(initbboard(), randompolicy)
    # above is faster than below
    #@benchmark Game2048.randmove(initbboard())
end

"""
    play_game_with_policy(policy, bitboard)

Play the game with the given policy. The moves are chosen at random weighted by the policy weights.
There is NO search in this algorithm
"""
function play_game_with_policy(policy, bitboard)
    while true
        newboard, _ = play_one_move_with_policy(bitboard, policy)
        if newboard == bitboard
            return newboard
        end

        bitboard = newboard |> add_tile
    end
end

if false
    # @code_warntype play_game_with_policy(randompolicy, initbboard())
    # @time play_game_with_policy(randompolicy, initbboard())
    # @benchmark play_game_with_policy(randompolicy, initbboard())
end

"""
    play_game_with_policy(policy, bitboard) -> GameRecorder

Play the game with the given policy. The moves are recorded
"""
function play_game_with_policy_w_record(policy, bitboard)
    recorder = GameRecorder()
    while true
        newboard, move_made = play_one_move_with_policy(bitboard, policy)
        push!(recorder, deepcopy(newboard), move_made)

        if newboard == bitboard
            # lost
            setfinalstate!(recorder, bitboard)
            trim!(recorder)

            return recorder
        end

        bitboard = newboard |> add_tile
    end
end

"""
Return the best move

Argument:

board - the board
n - how many times to play

"""

function find_best_move(board, n, policy)
    # try to play in each of the 4 directions `n` times one move ahead
    # after that just play randomly and record the score
    # the direction with the highest average score becomes the best direction
    res = zeros(Int, 4)

    probs = policy(board)

    # distribute the search by probs
    tries = 4n .* probs

    for direction in DIRS
        dir_idx = Int(direction) + 1
        new_board = move(board, direction)
        if new_board == board # not a valid move
            res[dir_idx] = -1
            continue
        end

        res_dir_idx = 0.0
        for _ in 1:tries[dir_idx]
            # random add one tile
            tmp_board = add_tile(new_board)

            # simulate from this onward
            # the roll out should be random
            fnl_tmp_board = play_game_with_policy(randompolicy, tmp_board)
            res_dir_idx += sum(2 .<< (bitboard_to_array(fnl_tmp_board) .- 1))
        end
        res[dir_idx] = res_dir_idx
    end

    DIRS[argmax(res)]
end

if false
    # @code_warntype find_best_move(initbboard(), 10, randompolicy)
    # @time find_best_move(initbboard(), 10, randompolicy)
end

"""
Records all the moves

Arguments:

    policy - a function that takes a bitboard and returns a policy
    n - how many moves of each direction to check
    bitboard - the starting state of play; defaults to `initbboard()`
"""
function play_via_monte_carlo_w_recorder(policy; n, bitboard=initbboard())
    recorder = GameRecorder()
    while true
        best_move = find_best_move(bitboard, n, policy)
        push!(recorder, deepcopy(bitboard), best_move)
        new_bitboard = move(bitboard, best_move)

        if new_bitboard != bitboard
            bitboard = add_tile(new_bitboard)
        else
            # lost
            setfinalstate!(recorder, bitboard)
            trim!(recorder)
            return recorder
        end
    end
end

function play_via_monte_carlo_wo_recorder(policy; n, bitboard=initbboard())
    while true
        best_move = find_best_move(bitboard, n, policy)
        new_bitboard = move(bitboard, best_move)

        if new_bitboard != bitboard
            bitboard = add_tile(new_bitboard)
        else
            # lost
            return new_bitboard
        end
    end
end

if false
    # @time result = play_via_monte_carlo_w_recorder(randompolicy; n=10);
end

"""
Play the game with the given policy. The moves are chosen at random weighted by the policy weights.

There is search as defined by `ntries_per_turn`. Each direction is played `ntries_per_turn` times
and the best move is played.
"""
function play_n_game(policy; n=1, ntries_per_turn=10)
    results = Vector{GameRecorder}(undef, n)
    # policies = [deepcopy(policy) for _ in 1:nthreads()]
    @threads for i in 1:n
        results[i] = play_via_monte_carlo_w_recorder(policy; n=ntries_per_turn);
    end
    results
end

function make_one_greedy_move(state, policy)
    prob = policy(state)
    i = 0
    while i <=4
        idx = argmax(prob)
        dir = DIRS[idx]

        new_state = move(state, dir)
        if new_state == state
            prob[idx] = 0
            i += 1
        else
            return new_state
        end
    end
    return state
end

"""
    Greedy player plays the game by choosing the move with highest policy probability.
"""
function greedy_player(policy, state::Bitboard=initbboard())
    while true
        state_after_move = make_one_greedy_move(state, policy)

        if state_after_move == state
            return state
        end

        state = add_tile(state_after_move)

        if state == state_after_move
            return state_after_move
        end
    end
end

if false
    using StatsBase
    @time countmap([maximum(Game2048.play_game_with_policy(randompolicy, initbboard())) for i in 1:10_000]) # random policy
    @time countmap([maximum(greedy_player(leftdownpolicy, initbboard())) for i in 1:10_000]) # left down

    @time countmap([maximum(play_n_game(leftdownpolicy; n=1, ntries_per_turn=10)[1].final_state) for i in 1:100])
    @time countmap([maximum(play_n_game(randompolicy; n=1, ntries_per_turn=10)[1].final_state) for i in 1:100])
    @time countmap([maximum(play_n_game(randompolicy; n=1, ntries_per_turn=100)[1].final_state) for i in 1:100])

    @time countmap([maximum(play_game_via_mcts_w_tries(randompolicy; n = 10)) for _ in 1:100])
end


text:
# Modelo de gerador de vapor
type Gerador_Vapor
	Gerador_Vapor()=begin
		PP2=outers.PP2
		new(
			DanaPlugin(Dict{Symbol,Any}(
				:Brief=>"Steam tables"
			)),
			Pot_sinal(Dict{Symbol,Any}(
				:Brief=>"Taxa de calor gerado na caldeira",
				:PosX=>1,
				:PosY=>0.5
			)),
			Eficiencia(Dict{Symbol,Any}(
				:Brief=>"Eficiencia do gerador de vapor"
			)),
			Potencia(Dict{Symbol,Any}(
				:Brief=>"Taxa de calor nos reaquecedores"
			)),
			Potencia(Dict{Symbol,Any}(
				:Brief=>"Taxa de calor nos superaquecedores"
			)),
			Potencia(Dict{Symbol,Any}(
				:Brief=>"Taxa de calor no evaporador"
			)),
			Potencia(Dict{Symbol,Any}(
				:Brief=>"Taxa de calor nos economizadores"
			)),
			Corrente(Dict{Symbol,Any}(
				:Brief=>"Agua de alimentacao",
				:Symbol=>"_{in_a}",
				:PosX=>0.5,
				:PosY=>1
			)),
			Corrente(Dict{Symbol,Any}(
				:Brief=>"Vapor a ser Reaquecido",
				:Symbol=>"_{in_ra}",
				:PosX=>0.7,
				:PosY=>1
			)),
			Corrente(Dict{Symbol,Any}(
				:Brief=>"Vapor Superaquecido",
				:Symbol=>"_{out_sa}",
				:PosX=>0.5,
				:PosY=>0
			)),
			Corrente(Dict{Symbol,Any}(
				:Brief=>"Vapor Reaquecido",
				:Symbol=>"_{out_ra}",
				:PosX=>0.7,
				:PosY=>0
			)),
			Corrente(Dict{Symbol,Any}(
				:Brief=>"Evaporador"
			)),
			Corrente(Dict{Symbol,Any}(
				:Brief=>"Economizadores"
			)),
			[
				:([Feco.S,Feco.H] = PP2.propPTv(Feco.P,Feco.T)),
				:(Qec = Feco.F * (Feco.H - Fin_a.H)),
				:(Fvap.F = Feco.F),
				:([Fvap.S,Fvap.H] = PP2.propPTv(Fvap.P,Fvap.T)),
				:(Qca = Fvap.F * (Fvap.H - Feco.H)),
				:(Fout_sa.F = Fvap.F),
				:([Fout_sa.S,Fout_sa.H] = PP2.propPTv(Fout_sa.P,Fout_sa.T)),
				:(Qsa = Fout_sa.F * (Fout_sa.H - Fvap.H)),
				:(Fout_ra.F = Fin_ra.F),
				:([Fout_ra.S,Fout_ra.H] = PP2.propPTv(Fout_ra.P,Fout_ra.T)),
				:(Qra = Fout_ra.F * (Fout_ra.H - Fin_ra.H)),
				:(Q_GV * EF_GV = Qec + Qca + Qsa + Qra),
			],
			[
				"Economizadores ECO1 + ECO1","","Evaporador - Camisa dagua","","","Superaquecedores BT + AT","","","Reaquecedores BT + AT","","","Caldeira",
			],
			[:PP2,],
			[:Q_GV,:EF_GV,:Qra,:Qsa,:Qca,:Qec,:Fin_a,:Fin_ra,:Fout_sa,:Fout_ra,:Fvap,:Feco,]
		)
	end
	PP2::DanaPlugin
	Q_GV::Pot_sinal
	EF_GV::Eficiencia
	Qra::Potencia
	Qsa::Potencia
	Qca::Potencia
	Qec::Potencia
	Fin_a::Corrente
	Fin_ra::Corrente
	Fout_sa::Corrente
	Fout_ra::Corrente
	Fvap::Corrente
	Feco::Corrente
	equations::Array{Expr,1}
	equationNames::Array{String,1}
	parameters::Array{Symbol,1}
	variables::Array{Symbol,1}
	attributes::Dict{Symbol,Any}
end
export Gerador_Vapor
function setEquationFlow(in::Gerador_Vapor)
	#	[Fin_a.S,Fin_a.H] = PP2.propPTl(Fin_a.P,Fin_a.T); # Reduntante no ciclo fechado
	#	Feco.F = Fin_a.F; # Reduntante no ciclo fechado
	addEquation(1)
	addEquation(2)
	addEquation(3)
	addEquation(4)
	addEquation(5)
	addEquation(6)
	addEquation(7)
	addEquation(8)
	addEquation(9)
	addEquation(10)
	addEquation(11)
	addEquation(12)
end
function atributes(in::Gerador_Vapor,_::Dict{Symbol,Any})
	fields::Dict{Symbol,Any}=Dict{Symbol,Any}()
	fields[:Pallete]=true
	fields[:Icon]="icon/caldeira"
	drive!(fields,_)
	return fields
end
Gerador_Vapor(_::Dict{Symbol,Any})=begin
	newModel=Gerador_Vapor()
	newModel.attributes=atributes(newModel,_)
	newModel
end


text:
### A Pluto.jl notebook ###
# v0.12.6

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 9bc349e8-efb0-11ea-1392-cd735828af33
# packageloads, etc
begin
	using DelimitedFiles
	using NIfTI
	using Images
	using ImageView
	using FileIO
	using StatsBase
	using Plots
	using PlutoUI
	using Printf
	using LinearAlgebra
	md"""
	[some setup / package includes...]

	the first time you run this notebook, it can take a while to start up.
	"""
end

# ╔═╡ 4df55a88-f98a-11ea-1c9f-cd52a4b5062f

md"""

# What happens in an fMRI experiment

## General idea

- some images are being displayed to a participant
- the visual part of their brain works harder (trying to process information)
- neural $\rightarrow$ haemodynamic changes

"""

# ╔═╡ f151a60e-016f-11eb-3ab4-bd2c5adbd801
md"""set up some image loading..."""

# ╔═╡ a1a5f0ee-016f-11eb-22bc-d965805be7f0
face_dir = "./stims/multiracial/frontal/";

# ╔═╡ b36f1e5c-016f-11eb-3f9e-f39becc598c6
faceImages = readdir(face_dir);

# ╔═╡ 80dc039a-0191-11eb-1792-352eabf32591
object_dir = "./stims/objects/";

# ╔═╡ 940859e4-0191-11eb-0a64-bfc069e803d8
objectImages = readdir(object_dir);

# ╔═╡ 9af0d01c-1488-11eb-159d-c72e4cffaa22
md"""Some parameters that define the timing of the block design experiment"""

# ╔═╡ af8bd0d4-0190-11eb-1f8f-7368cd387053
nA = 8

# ╔═╡ b6cc5c4c-0190-11eb-00fc-1ddd56445030
nGray = 8

# ╔═╡ b951cf10-0190-11eb-31fe-0beb75cc6f3e
nRepeats = 5

# ╔═╡ bddcdeac-0191-11eb-061e-b75875bc5b7a
nDynamics = 2*(nA+nGray)*nRepeats

# ╔═╡ 3488c298-0190-11eb-03ef-cf178151e777
begin
	timing = repeat([[repeat(["gray"], nGray, 1); repeat(["face"], nA, 1)];
					  repeat(["gray"], nGray, 1); repeat(["object"], nA, 1)], 							nRepeats, 1);
	md"""
	Define a timing vector with $( nDynamics ) dynamics.

	This is done by repeating `["gray"]` $nGray times, then `["face"]` $nA times, the same for `["object"]` and repeating those $nRepeats times.

	"""
end

# ╔═╡ ae0fe97e-0194-11eb-2c80-41c90b26ee5c
md"""
## Run a sped-up version of the experiment

We can bind a value that changes every "tick" using a `Clock` function... to make sure we don't keep going beyond the length of the dataset ($nDynamics), we wrap around using the `mod`ulo.

"""

# ╔═╡ 19252aa4-0105-11eb-0b5c-5f3769240354
md"""Step through time: each **tick** is a **volume** in the fMRI experiment"""

# ╔═╡ e4c0f9f6-01cb-11eb-111f-b774df6174bb
md"""Show stimulus timing w/ data? $(@bind showStimTiming CheckBox())

If this checkbox is ticked, then display an image from the corresponding block"""

# ╔═╡ a54017da-0194-11eb-309c-415dd99b4db1
@bind t Clock(0.5)

# ╔═╡ 28ec2008-0193-11eb-1ca6-afde7ea782a1
theVolume = mod(t-1, nDynamics) + 1;

# ╔═╡ 778800f8-f98a-11ea-1b87-1108f500ad5f
begin
	# this does image loading from disk - not that fast, but good enough for here.
	if showStimTiming
		# might be one off? think through this bit more.
		if timing[theVolume] == "face"
			ex = load(string(face_dir, rand(faceImages)))
		elseif timing[theVolume] == "object"
			ex = load(string(object_dir, rand(objectImages)))
		else
			ex = Gray.(repeat([0.5], 250, 250))
		end
	end
end

# ╔═╡ c4c00586-01ce-11eb-0625-17b9ab6bfa8f
md"""Show face response w/ data? $(@bind showResponseTimingF CheckBox())

Show object response w/ data? $(@bind showResponseTimingO CheckBox())
"""

# ╔═╡ d27ceec0-01ca-11eb-3152-7d7603d6546a
fStim = timing .== "face" ;

# ╔═╡ b2df1378-01cb-11eb-01fd-4d9da6872e9b
oStim = timing .== "object" ;

# ╔═╡ 2b2f6be4-f07f-11ea-1459-775749ca949f
md"""## Load in that CSV file """

# ╔═╡ 368e2640-f8e0-11ea-1ddf-b37069ba7588
d, h= readdlm("data/stimulus_timing.csv", ',', Float64, header = true )

# ╔═╡ 5bb7ed88-0106-11eb-2181-d123e1a1cc31
begin
	data_x = view(d, 1:theVolume, 1)
	data_y = view(d, 1:theVolume, 5)

	face_offset = 3.2 # shift down to see!
	object_offset = 4.5

	face_pred = view(d, 1:theVolume, 3) .- face_offset
	object_pred = view(d, 1:theVolume, 4) .- object_offset

	plot(data_x, data_y,
		lw=2, c=:black,
		xlim=(0,160), ylim=(-5, 5),
		xlabel="Time (volumes)",
		ylabel="fMRI response (%signal change)",
		label="data")

	if showResponseTimingF
		plot!(data_x, face_pred,
		lw=2,
		c=:red,
		label="face timing")
	end

	if showResponseTimingO
		plot!(data_x, object_pred,
			c=:blue,
			lw=2,
			label="object timing")
	end

	plot!([d[theVolume,1]], [d[theVolume,5]], seriescolor=:black,
		m=:circle, ms=7, label="")

end

# ╔═╡ d5c0b236-f98b-11ea-0444-6fe340642707
md"""

## Some settings and helper functions

"""

# ╔═╡ bfdd7d56-f98a-11ea-3064-a9033ba994a5
# choose which plotting backend to use.

# plotly() is dynamic / webbased
# gr() make nice static plots

#plotly()
gr()

# ╔═╡ Cell order:
# ╟─9bc349e8-efb0-11ea-1392-cd735828af33
# ╟─4df55a88-f98a-11ea-1c9f-cd52a4b5062f
# ╟─f151a60e-016f-11eb-3ab4-bd2c5adbd801
# ╠═a1a5f0ee-016f-11eb-22bc-d965805be7f0
# ╠═b36f1e5c-016f-11eb-3f9e-f39becc598c6
# ╠═80dc039a-0191-11eb-1792-352eabf32591
# ╠═940859e4-0191-11eb-0a64-bfc069e803d8
# ╟─9af0d01c-1488-11eb-159d-c72e4cffaa22
# ╟─af8bd0d4-0190-11eb-1f8f-7368cd387053
# ╟─b6cc5c4c-0190-11eb-00fc-1ddd56445030
# ╟─b951cf10-0190-11eb-31fe-0beb75cc6f3e
# ╟─bddcdeac-0191-11eb-061e-b75875bc5b7a
# ╟─3488c298-0190-11eb-03ef-cf178151e777
# ╟─ae0fe97e-0194-11eb-2c80-41c90b26ee5c
# ╠═28ec2008-0193-11eb-1ca6-afde7ea782a1
# ╟─19252aa4-0105-11eb-0b5c-5f3769240354
# ╟─e4c0f9f6-01cb-11eb-111f-b774df6174bb
# ╟─778800f8-f98a-11ea-1b87-1108f500ad5f
# ╟─a54017da-0194-11eb-309c-415dd99b4db1
# ╟─c4c00586-01ce-11eb-0625-17b9ab6bfa8f
# ╠═5bb7ed88-0106-11eb-2181-d123e1a1cc31
# ╠═d27ceec0-01ca-11eb-3152-7d7603d6546a
# ╠═b2df1378-01cb-11eb-01fd-4d9da6872e9b
# ╟─2b2f6be4-f07f-11ea-1459-775749ca949f
# ╠═368e2640-f8e0-11ea-1ddf-b37069ba7588
# ╟─d5c0b236-f98b-11ea-0444-6fe340642707
# ╟─bfdd7d56-f98a-11ea-3064-a9033ba994a5

max_stars_repo_name: schluppeck/hands-on-brain-data

max_issues_repo_name: schluppeck/hands-on-brain-data

max_forks_repo_name: schluppeck/hands-on-brain-data


text:
using Revise, WaltonForaging
##
ispath("/home/beatriz/Documents/") ? (main_path ="/home/beatriz/Documents/") : (main_path = "/Users/dariosarra/Documents/Lab/Walton/LForaging")
fig_dir = joinpath(main_path,"Figures")
# pokes = DataFrame(CSV.read(joinpath(main_path,"Miscellaneous/Poke_data.csv")))#, DataFrame)
# pokes2 = DataFrame(CSV.read(joinpath(main_path,"Miscellaneous/New_Poke_data.csv")))#, DataFrame)
pokes = CSV.read(joinpath(main_path,"Miscellaneous/States_Poke_data.csv"), DataFrame)
preprocess_pokes!(pokes)
open_html_table(pokes[1:2000,:])
##
mice = union(pokes.MOUSE)
test_multiple_sessions = filter(r -> r.MOUSE == mice[1],pokes)
function fit_handler(test_multiple_sessions)
    res = []
    for i in 1:10
        model,optimizer = fit(RhoComparison,test_multiple_sessions)
        push!(res,(Distributions.params(model), Optim.converged(optimizer)))
    end
    return res
end
##
# m, conv = @time fit(RhoComparison,test_multiple_sessions)
# m, conv = @time fit(m,test_multiple_sessions)
# check = @time fit_handler(test_multiple_sessions)
# updateparams!
##
##
Francois_res = (env_alpha = 0.003, patch_aplha = 0.266, beta = 6.26, reset = 0.339, bias = 4.147)
open_html_table(test_multiple_sessions[1:5000,:])
open_html_table(test_multiple_sessions[1:10,:])
-sum(log.(test_multiple_sessions.Probability))
##
m = init(RhoComparison)
WaltonForaging.updateparams!(m, [1.0,2.0,3.0,4.0,5.0])
m
all([(v < l) for (v,l) in zip(params(m), m.limits)])
m
conv
open_html_table(test_multiple_sessions[1:200,:])