defmodule Chi2fit.MonteCarlo do

  # Copyright 2016-2021 Pieter Rijken
  #
  # Licensed 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.

  import Chi2fit.Statistics

  alias Chi2fit.Matrix, as: M

  @doc """
  Basic Monte Carlo simulation to repeatedly run a simulation multiple times.

  ## Options

      `:collect_all?` - If true, collects data from each individual simulation run and returns this an the third element of the result tuple

  """
  @spec mc(iterations :: pos_integer, fun :: ((pos_integer) -> float), options :: Keyword.t) :: {avg :: float, sd :: float, tries :: [float]} | {avg :: float, sd :: float}
  def mc(iterations,fun,options \\ []) do
    all? = options[:collect_all?] || false

    tries = 1..iterations |> Enum.map(fn _ -> fun.() end)
    avg = moment tries, 1
    sd = :math.sqrt momentc(tries,2,avg)
    if all?, do: {avg,sd,tries}, else: {avg,sd}
  end

  def total_mc(result, fun, mode \\ :use_bounds, iterations \\ 1000) do
    {_, cov, parameters, _} = result

    ranges = case mode do
      :use_bounds ->
        # Pick up the error in the paramater value
        errors = cov |> M.diagonal |> Enum.map(fn x -> x|>abs|>:math.sqrt end)
        Enum.zip(parameters, errors) |> Enum.map(fn {par, err} -> [par - err, par + err] end)

      :use_ranges ->
        {_, _, _, parranges} = result
        parranges |> Tuple.to_list |> tl
    end
    outcomes = ranges
    |> List.foldr([], fn
      [left, right], [] -> [[left],[right]]
      [left, right], acc -> Enum.flat_map(acc, & [[left|List.wrap(&1)], [right|List.wrap(&1)]])
    end)
    |> Enum.map(& mc(iterations, fun.(&1)))
    |> Enum.map(& elem(&1,0))

    {Enum.min(outcomes), Enum.max(outcomes)}
  end

  @doc """
  Performs a nested bootstrap on sample data.
  """
  @spec nested_bootstrap(list(number()), (list(number()) -> number), Keyword.t()) :: list(number())
  def nested_bootstrap(values, fun, options) do
    iterations = Keyword.get(options, :iterations, 10_000)
    resamples = Keyword.get(options, :resamples, 1_000)
    target = Keyword.fetch!(options, :target)

    bootstrap(resamples, values, fn
      data, _i ->
        {_avg,_sd,collection} = mc(iterations, fun.(data), collect_all?: true)
        get_percentile(collection, target)
      end)
  end

  @doc """
  Forecasts how many time periods are needed to complete `size` items

  Related functions: `forecast_duration/2` and `forecast_items/2`.
  """
  @spec forecast(fun :: (() -> non_neg_integer),size :: pos_integer, tries :: pos_integer, update :: (() -> number)) :: number
  def forecast(fun, size, tries \\ 0,update \\ fn -> 1 end)
  def forecast(fun, size, tries, update) when size>0 do
    forecast(fun, size-fun.(),tries+update.(),update)
  end
  def forecast(_fun,_size,tries,_update), do: tries

  @doc """
  Returns a function for forecasting the duration to complete a number of items.

  This function is a wrapper for `forecast/4`.

  ## Arguments

      `data` - either a data set to base the forecasting on, or a function that returns (random) numbers
      `size` - the number of items to complete

  """
  @spec forecast_duration(data :: [number] | (()->number), size :: pos_integer) :: (() -> number)
  def forecast_duration(data, size) when is_list(data) do
    fn -> forecast(fn -> Enum.random(data) end, size) end
  end
  def forecast_duration(fun, size) when is_function(fun,0) do
    fn -> forecast(fun, size) end
  end

  @doc """
  Returns a function for forecasting the number of completed items in a number periods.

  This function is a wrapper for `forecast/4`.

  ## Arguments

      `data` - either a data set to base the forecasting on, or a function that returns (random) numbers
      `periods` - the number of periods to forecast the number of completed items for

  """
  @spec forecast_items(data :: [number] | (()->number), periods :: pos_integer) :: (() -> number)
  def forecast_items(data, periods) when is_list(data) do
    fn -> forecast(fn -> 1 end, periods, 0, fn -> Enum.random(data) end) end
  end
  def forecast_items(fun, periods) when is_function(fun,0) do
    fn -> forecast(fn -> 1 end, periods, 0, fun) end
  end

  @spec get_percentile(list(), non_neg_integer()) :: float()
  def get_percentile(collection, value) do
    size = length(collection)

    [0|collection]
    |> Enum.sort(&(&1 >= &2))
    |> Enum.find_index(&(&1 < value))
    |> Kernel./(size)
  end

end