#-------------------------------------------------------------
#
# 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.
#
#-------------------------------------------------------------

# This is a builtin function that implements GARCH(1,1), a statistical model used in analyzing time-series data where the variance
# error is believed to be serially autocorrelated
#
# COMMENTS
# This has some drawbacks: slow convergence of optimization (sort of simulated annealing/gradient descent)
# TODO: use BFGS or BHHH if it is available (this are go to methods)
# TODO: (only then) extend to garch(p,q); otherwise the search space is way too big for the current method
#
# INPUT:
# -----------------------------------------------------------------------------------------
# X                The input Matrix to apply Arima on.
# kmax             Number of iterations
# momentum         Momentum for momentum-gradient descent (set to 0 to deactivate)
# start_stepsize   Initial gradient-descent stepsize
# end_stepsize     gradient-descent stepsize at end (linear descent)
# start_vicinity   proportion of randomness of restart-location for gradient descent at beginning
# end_vicinity     same at end (linear decay)
# sim_seed         seed for simulation of process on fitted coefficients
# verbose          verbosity, comments during fitting
# -----------------------------------------------------------------------------------------
#
# OUTPUT:
# --------------------------------------------------------------------------------------------------
# fitted_X         simulated garch(1,1) process on fitted coefficients
# fitted_var_hist  variances of simulated fitted process
# best_a0          onstant term of fitted process
# best_arch_coef   1-st arch-coefficient of fitted process
# best_var_coef    1-st garch-coefficient of fitted process
# --------------------------------------------------------------------------------------------------

m_garch = function(Matrix[Double] X, Integer kmax, Double momentum, Double start_stepsize, Double end_stepsize, Double start_vicinity,
  Double end_vicinity, Integer sim_seed, Boolean verbose)
return (Matrix[Double] fitted_X, Matrix[Double] fitted_var_hist, Double best_a0, Double best_arch_coef, Double best_var_coef) {

  [a0, arch_coef, var_coef] = sample_feasible_params() # initialize startpoint
  curr_qmle = qmle(X, a0, arch_coef, var_coef) # initialize quasi-log-likelihood at start

  # record the parameters of max-quasi-log-likelihood thus far
  best_a0 = a0
  best_arch_coef = arch_coef
  best_var_coef = var_coef
  best_qmle = curr_qmle

  # record last change of gradient if update was applied
  last_change_a0 = 0
  last_change_arch_coef = 0
  last_change_var_coef = 0

  # initialize stepsize (linear decay)
  stepsize = start_stepsize

  # initialize vicinity (linear decay)
  vicinity = start_vicinity

  # all coeffs need be >0 to provide a feasible solution; clip at this constant
  clip_at = 0.00001

  # do gradient descent
  for (k in 1:kmax-1) {
    # update vicinity and stepsize
    progress = k/kmax
    stepsize = (1-progress) * start_stepsize + progress*end_stepsize
    vicinity = (1-progress) * start_vicinity + progress*end_vicinity

    # get gradient
    [d_a0, d_arch_coef, d_var_coef] = gradient(X, a0, arch_coef, var_coef)

    # newly proposed parameters
    new_a0 = max(a0 + (stepsize * d_a0) + momentum * (last_change_a0), clip_at)
    new_arch_coef = max(arch_coef + (stepsize * arch_coef) + (momentum * last_change_arch_coef), clip_at)
    new_var_coef = max(var_coef + (stepsize * var_coef) + (momentum * last_change_var_coef), clip_at)

    # ensure feasibility (this condition provides stationarity, see literature)
    while (new_arch_coef + new_var_coef > 1) {
      new_arch_coef = new_arch_coef / 2
      new_var_coef = new_var_coef / 2
    }

    # objective function value of new feasible parameters
    new_qmle = qmle(X, new_a0, new_arch_coef, new_var_coef)

    # record the change of coefficients for momentum updates
    change_a0 = new_a0 - a0
    change_arch_coef = new_arch_coef - arch_coef
    change_var_coef = new_var_coef - var_coef

    # check if improvement
    if (new_qmle > curr_qmle) {
      # if so, update and use change for momentum
      last_change_a0 = change_a0
      last_change_arch_cof = change_arch_coef
      last_change_var_coef = change_var_coef
      update_reason = "gradient"
    }
    else {
      # else: chance for restart at close point and reset momentum

      update_reason = "no update" # unless the random restart applies
      coin_flip = as.scalar(rand(rows=1, cols=1)) # random restart gets less likely with progressing search

      if (coin_flip > progress) {
        # sample random restart-point
        # vicinity tells how far we move to the random point (0: dont move, 1: move fully to random point)
        # similar to simulated annealing with adaptive neighborhood

        [new_a0, new_arch_coef, new_var_coef] = sample_neighbor(a0, arch_coef, var_coef, vicinity)

        # reset momentum
        last_change_a0 = 0
        last_change_arch_cof = 0
        last_change_var_coef = 0

        update_reason = "restart"
      }
    }

    # update the point
    a0 = new_a0
    arch_coef = new_arch_coef
    var_coef = new_var_coef

    # update qmle at the moment
    curr_qmle = qmle(X, a0, arch_coef, var_coef)

    # check and record if it is the best point thus far
    update_best = (curr_qmle > best_qmle);
    if (update_best) {
      best_qmle = curr_qmle
      best_a0 = a0
      best_arch_coef = arch_coef
      best_var_coef = var_coef
    }

    # logging: report state of gradient descent
    if (verbose) {
      print("k                | " + toString(k))
      print("a0               | " +  toString(a0))
      print("arch coef        | " +  toString(arch_coef))
      print("var coef         | " +  toString(var_coef))
      print("qmle             | " +  toString(curr_qmle))
      print("stepsize         | " + toString(stepsize))
      print("update reason    | " + update_reason)
      print("update best      | " + update_best)
      print("____________________________________")
    }
  }

  # simulate process from best solution
  sim_steps = nrow(X)
  [fitted_X, fitted_var_hist] = sim_garch(best_a0, best_arch_coef, best_var_coef, sim_steps, sim_seed)

  # logging: report output
  if (verbose) {
    print("end iteration: return the following: ")
    print("best qmle            | " + toString(best_qmle))
    print("best a0              | " + toString(best_a0))
    print("best arch_coef       |" + toString(best_arch_coef))
    print("best var_coef        |" + toString(best_var_coef))
    print("____________________________________")
  }
}


# ------- UTILITY FUNCTIONS

# ------- quasi-log-likelihood of garch-1-1 coefficients for given data
# ------- https://math.berkeley.edu/~btw/thesis4.pdf
qmle = function(Matrix[Double] X, Double a0, Double arch_coef, Double var_coef)
  return (Double qmle)
{
  n = nrow(X)

  # initialize variance
  var_0 = a0 / (1-arch_coef - var_coef)
  vars = matrix(var_0, rows=1, cols=1)

  # init loop for var and qmle computation
  var_lag = var_0
  xq_lag = as.scalar(X[1,1])^2
  qmle = 0

  # compute vars and qmle recursively
  # TODO vectorize via cummulative aggregates?
  for (t in 2:n) {
    xq_t = as.scalar(X[t,1])^2
    var_t = a0 + arch_coef*xq_lag + var_coef*var_lag
    qmle = qmle - (1/(2*n)) * (log(var_t) + (xq_t / var_t)) # up to constant

    vars = rbind(vars, matrix(var_t, rows=1, cols=1))
    var_lag = var_t
    xq_lag = xq_t
  }
}


# ------- returns coefs which yield a stationary garch process; sampled uniform at random in feasible region
sample_feasible_params = function()
return (Double a0, Double arch_coef, Double var_coef) {
  # TODO vectorize (e.g., via 64 random numbers -> 1e-20 failure prob)
  a0 = as.scalar(rand(rows=1, cols=1))
  arch_coef = as.scalar(rand(rows=1, cols=1))
  var_coef = as.scalar(rand(rows=1, cols=1))

  while (arch_coef + var_coef >= 1) {
    arch_coef = as.scalar(rand(rows=1, cols=1))
    var_coef = as.scalar(rand(rows=1, cols=1))
  }
}

# ------- sample random feasible point and blend with current point
# ------- vicinity tells the proportion of the newly random point is blended in
sample_neighbor = function(Double a0, Double arch_coef, Double var_coef, Double vicinity)
return (Double nb_a0, Double nb_arch_coef, Double nb_var_coef) {
  # nb is convex comb of current val and a distant target value
  # the smaller the vicinity, the closer we are to the current point and the less randomness there is

  [target_a0, target_arch_coef, target_var_coef] = sample_feasible_params()

  nb_a0 = vicinity*target_a0 + (1-vicinity)*a0
  nb_arch_coef = vicinity*target_arch_coef + (1-vicinity)*arch_coef
  nb_var_coef = vicinity*target_var_coef + (1-vicinity)*var_coef
}

# ------- numerically approximated gradient of quasi-log-likelihood
gradient = function(Matrix[Double] X, Double a0, Double arch_coef, Double var_coef)
  return (Double d_a0, Double d_arch_coef, Double d_var_coef)
{
  eps = 0.00001
  qmle_val = qmle(X, a0, arch_coef, var_coef)
  d_a0 = (qmle(X, a0 + eps, arch_coef, var_coef) - qmle_val) / eps
  d_arch_coef = (qmle(X, a0, arch_coef + eps, var_coef) - qmle_val) / eps
  d_var_coef = (qmle(X, a0, arch_coef, var_coef + eps) - qmle_val) / eps
}

# ------- simulate a garch-process with given parameters
sim_garch = function(Double a0, Double arch_coef, Double var_coef, Integer simsteps, Integer seed)
  return (Matrix[Double] X, Matrix[Double] var_hist)
{
  # init var and std
  var = a0 / (1 - arch_coef - var_coef)
  std = sqrt(var)

  # init outputs
  var_hist = matrix(0, rows=0, cols=1)
  X = matrix(0, rows=0, cols=1)

  # recursively construct time series
  # TODO vectorize via cumsumprod
  for (t in seq(1,simsteps,1)) {
    noise = as.scalar(rand(rows=1, cols=1, pdf="normal", seed=seed)) # N(0,1) noise
    xt = noise * std

    X = rbind(X, as.matrix(xt))
    var_hist = rbind(var_hist,as.matrix(var))

    # get var and std for next round
    var = a0 + arch_coef * (xt^2) + var_coef * var
    std = sqrt(var)

    seed = seed + 1 # prevent same innovations
  }
}
