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

source("scripts/nn/networks/resnet.dml") as resnet
source("scripts/nn/networks/resnet_util.dml") as util

forward = function(matrix[double] X, int Hin, int Win,
                            list[unknown] model, string mode,
                            list[unknown] ema_means_vars)
    return (matrix[double] out, list[unknown] ema_means_vars_upd,
            list[unknown] cached_out, list[unknown] cached_means_vars) {
    /*
     * Forward pass of the ResNet152 as introduced in
     * "Deep Residual Learning for Image Recognition" by
     * Kaiming He et. al., refined in "ResNet v1.5 for
     * PyTorch" by NVIDIA and inspired by the PyTorch
     * implementation.
     *
     * Inputs:
     * - X: Inputs, of shape (N, C_in*Hin*Win).
     *     C_in = 3 is expected.
     * - Hin: Input height.
     * - Win: Input width.
     * - model: Weights and bias matrices of the model
     *     with the following order/content:
     *   -> 1: Weights of conv 1 7x7, of shape (64, 3*7*7)
     *   -> 2: Weights of batch norm 1, of shape (64, 1).
     *   -> 3: Bias of batch norm 1, of shape (64, 1).
     *   -> 4: List of weights for first residual layer
     *         with 64 base channels.
     *   -> 5: List of weights for second residual layer
     *         with 128 base channels.
     *   -> 6: List of weights for third residual layer
     *         with 256 base channels.
     *   -> 7: List of weights for fourth residual layer
     *         with 512 base channels.
     *      List of residual layers 1, 2, 3 & 4 have
     *      the content/order:
     *      -> i: List of weights for residual block i.
     *            with i in {1, ..., layer_sizes[layer]}
     *         Each list of weights for a residual block
     *         must follow the same order as defined in
     *         the documentation of bottleneck_block_forward().
     *   -> 8: Weights of fully connected layer, of shape (C_out, classes)
     *         where C_out = 512 for basic block type and C_out = 2048
     *         for bottleneck block type.
     *   -> 9: Bias of fully connected layer, of shape (1, classes)
     * - mode: 'train' or 'test' to indicate if the model is currently
     *     being trained or tested for badge normalization layers.
     *     See badge_norm2d.dml docs for more info.
     * - ema_means_vars: List of exponential moving averages for mean
     *     and variance for badge normalization layers.
     *   -> 1: EMA for mean of badge norm 1, of shape (64, 1).
     *   -> 2: EMA for variance of badge norm 1, of shape (64, 1).
     *   -> 3: List of EMA means and vars for residual layer 1.
     *   -> 4: List of EMA means and vars for residual layer 2.
     *   -> 5: List of EMA means and vars for residual layer 3.
     *   -> 6: List of EMA means and vars for residual layer 4.
     *      Lists for EMAs of layer 1, 2, 3 & 4 must have the
     *      following order:
     *      -> i: List of EMA means and vars for residual block i.
     *            with i in {1, ..., layer_sizes[layer]}
     *         Each list of EMAs for a residual block
     *         must follow the same order as defined in
     *         the documentation bottleneck_block_forward().
     * - NOTICE: The lists of the first blocks for layer 2, 3 and 4
     *           must include weights and EMAs for 1 extra conv layer
     *           and a batch norm layer for the downsampling on the
     *           identity path.
     *
     * Outputs:
     * - out: Outputs, of shape (N, classes)
     * - ema_means_vars_upd: List of updated exponential moving averages
     *     for mean and variance of badge normalization layers. It follows
     *     the same exact structure as the input EMAs list.
     * - cached_out: Outputs of each layer for computation of backward
     *     pass. Refer to the code for the order of elements.
     * - cached_means_vars: List of cached means and vars returned from
     *     each batch normalization layer. This is required for the
     *     backward pass of the network.
     */
    layer_sizes = list(3, 8, 36, 3)
    block_type = "bottleneck"
    [out, ema_means_vars_upd, cached_out, cached_means_vars] = resnet::resnet_forward(X, Hin, Win, block_type,
        layer_sizes, model, mode, ema_means_vars)
}

backward = function(matrix[double] dOut, list[unknown] cached_out,
                    list[unknown] model, list[unknown] cached_means_vars)
    return (matrix[double] dX, list[unknown] gradients) {
    /*
     * Backward pass of the ResNet 152 model as introduced in
     * "Deep Residual Learning for Image Recognition" by
     * Kaiming He et. al. and inspired by PyTorch.
     *
     * Inputs:
     * - dOut: Partial derivative of the loss w.r.t. the
     *     outputs, of shape (N, classes).
     * - cashed_out: List of cashed outputs from the forward
     *     pass for each layer output with dimensions.
     * -> NOTICE: Use the exact list returned by the forward
     *            pass without modification.
     * - model: Weights and bias matrices of the model
     *     with the same order as for the forward pass.
     * -> NOTICE: The lists of the first blocks for layer 2, 3 and 4
     *            must include weights and EMAs for 1 extra conv layer
     *            and a batch norm layer for the downsampling on the
     *            identity path.
     * - cached_means_vars: List of cached means and vars returned
     *     by the forward pass. It has the same structure as the
     *     ema_means_vars for the forward pass.
     * -> NOTICE: Use the exact list returned by the forward
     *            pass without modification.
     *
     * Outputs:
     * - dX: Derivative of the loss w.r.t. the inputs, of
     *     shape (N, C_in*Hin*Win).
     * - gradients: Gradients of each learnable parameters of
     *     every layer in the network with the same structure
     *     as the input weights.
     * -> NOTICE: To update the parameters of the model, use
     *            one of the provided utility functions.
     */
    layer_sizes = list(3, 8, 36, 3)
    block_type = "bottleneck"
    [dX, gradients] = resnet::resnet_backward(dOut, cached_out, block_type, layer_sizes, model, cached_means_vars)
}

/*
 * Model initialization.
 */

init = function(int classes, int seed)
    return(list[unknown] model, list[unknown] emas) {
    /*
     * Initializes all parameters of the model according to the
     * respective initializer functions of each layer.
     * NOTICE: It is recommended to use this function for
     *         initialization to directly have the correct
     *         model list structure.
     *
     * Inputs:
     * - classes: Number of network output classes.
     * - seed: Seed for randomizer function.
     *
     * Outputs:
     * - model: List of weights and biases with the structure
     *     described in the forward pass documentation.
     * - emas: List of exponential moving averages of the mean
     *     and variance for each batch normalization layer.
     *     The structure is described in the forward pass
     *     documentation.
     */
    layer_sizes = list(3, 8, 36, 3)
    [model, emas] = resnet::init(classes, "bottleneck", layer_sizes, seed)
}

/*
 * Utility functions for optimizers.
 */

init_adagrad_optim_params = function(int classes)
    return(list[unknown] params) {
    /*
     * Initializes the state of the adagrad optimizer for every
     * learnable parameter of ResNet 152.
     *
     * Inputs:
     * - classes: Number of network output classes.
     *
     * Outputs:
     * - params: List of state parameters with the same structure
     *     as weights of the forward and backward pass. It can be
     *     directly passed to the update parameter function.
     */
    layer_sizes = list(3, 8, 36, 3)
    params = util::init_optim("adagrad", classes, "bottleneck", layer_sizes)
}

update_params_with_adagrad = function(list[unknown] model, list[unknown] gradients,
                                      double lr, double epsilon,
                                      list[unknown] optim_state)
    return (list[unknown] model_upd, list[unknown] optim_state_upd) {
    /*
     * Updates all learnable parameters with the adagrad optimizer.
     *
     * Inputs:
     * - model: Model parameters, same as for forward and
     *     backward pass.
     * - gradients: Gradients, returned from the backward pass.
     * - lr: Learning rate.
     * - epsilon: Smoothing term to avoid divide by zero errors.
     *     Typical values are in the range of [1e-8, 1e-4].
     * - optim_state: Optimizer states for all model parameters.
     *
     * Outputs:
     * - model_upd: Updated model parameters.
     * - optim_state_upd: Updated model states for all parameters.
     */
    layer_sizes = list(3, 8, 36, 3)
    hyper_params = list(lr, epsilon)
    [optim_state_upd, model_upd] = util::update_params("adagrad", optim_state, hyper_params, gradients, model,
        "bottleneck", layer_sizes)
}

init_adam_optim_params = function(int classes)
    return(list[unknown] params) {
    /*
     * Initializes the state of the adam optimizer for every
     * learnable parameter of ResNet 152.
     *
     * Inputs:
     * - classes: Number of network output classes.
     *
     * Outputs:
     * - params: List of state parameters with the same structure
     *     as weights of the forward and backward pass. It can be
     *     directly passed to the update parameter function.
     */
    layer_sizes = list(3, 8, 36, 3)
    params = util::init_optim("adam", classes, "bottleneck", layer_sizes)
}

update_params_with_adam = function(list[unknown] model, list[unknown] gradients,
                                   double lr, double beta1, double beta2,
                                   double epsilon, int t,
                                   list[unknown] optim_state)
    return (list[unknown] model_upd, list[unknown] optim_state_upd) {
    /*
     * Updates all learnable parameters with the adam optimizer.
     *
     * Inputs:
     * - model: Model parameters, same as for forward and
     *     backward pass.
     * - gradients: Gradients, returned from the backward pass.
     * - lr: Learning rate. Recommended: 0.001
     * - beta1: Exponential decay rate for the 1st moment estimates.
     *     Recommended value is 0.9.
     * - beta2: Exponential decay rate for the 2nd moment estimates.
     *     Recommended value is 0.999.
     * - epsilon: Smoothing term to avoid divide by zero errors.
     *     Recommended value is 1e-8.
     *  - t: Timestep, starting at 0.
     * - optim_state: Optimizer states for all model parameters.
     *
     * Outputs:
     * - model_upd: Updated model parameters.
     * - optim_state_upd: Updated model states for all parameters.
     */
    layer_sizes = list(3, 8, 36, 3)
    hyper_params = list(lr, beta1, beta2, epsilon, t)
    [optim_state_upd, model_upd] = util::update_params("adam", optim_state, hyper_params, gradients, model,
        "bottleneck", layer_sizes)
}

init_rmsprop_optim_params = function(int classes)
    return(list[unknown] params) {
    /*
     * Initializes the state of the rmsprop optimizer for every
     * learnable parameter of ResNet 152.
     *
     * Inputs:
     * - classes: Number of network output classes.
     *
     * Outputs:
     * - params: List of state parameters with the same structure
     *     as weights of the forward and backward pass. It can be
     *     directly passed to the update parameter function.
     */
    layer_sizes = list(3, 8, 36, 3)
    params = util::init_optim("rmsprop", classes, "bottleneck", layer_sizes)
}

update_params_with_rmsprop = function(list[unknown] model, list[unknown] gradients,
                                      double lr, double decay_rate,
                                      double epsilon, list[unknown] optim_state)
    return (list[unknown] model_upd, list[unknown] optim_state_upd) {
    /*
     * Updates all learnable parameters with the rmsprop optimizer.
     *
     * Inputs:
     * - model: Model parameters, same as for forward and
     *     backward pass.
     * - gradients: Gradients, returned from the backward pass.
     * - lr: Learning rate.
     * - decay_rate: Term controlling the rate of the moving average.
     *     Typical values are in the range of [0.9, 0.999].
     * - epsilon: Smoothing term to avoid divide by zero errors.
     *     Typical values are in the range of [1e-8, 1e-4].
     * - optim_state: Optimizer states for all model parameters.
     *
     * Outputs:
     * - model_upd: Updated model parameters.
     * - optim_state_upd: Updated model states for all parameters.
     */
    layer_sizes = list(3, 8, 36, 3)
    hyper_params = list(lr, decay_rate, epsilon)
    [optim_state_upd, model_upd] = util::update_params("rmsprop", optim_state, hyper_params, gradients, model,
        "bottleneck", layer_sizes)
}

update_params_with_sgd = function(list[unknown] model,
                                 list[unknown] gradients,
                                 double lr)
    return (list[unknown] model_upd) {
    /*
     * Updates all learnable parameters with the sgd optimizer.
     *
     * Inputs:
     * - model: Model parameters, same as for forward and
     *     backward pass.
     * - gradients: Gradients, returned from the backward pass.
     * - lr: Learning rate.
     *
     * Outputs:
     * - model_upd: Updated model parameters.
     */
    layer_sizes = list(3, 8, 36, 3)
    hyper_params = list(lr)
    optim_state = list()
    [optim_state_upd, model_upd] = util::update_params("sgd", optim_state, hyper_params, gradients, model, "bottleneck",
        layer_sizes)
}

init_sgd_momentum_optim_params = function(int classes)
    return(list[unknown] params) {
    /*
     * Initializes the state of the sgd momentum optimizer for
     * every learnable parameter of ResNet 152.
     *
     * Inputs:
     * - classes: Number of network output classes.
     *
     * Outputs:
     * - params: List of state parameters with the same structure
     *     as weights of the forward and backward pass. It can be
     *     directly passed to the update parameter function.
     */
    layer_sizes = list(3, 8, 36, 3)
    params = util::init_optim("sgd_momentum", classes, "bottleneck", layer_sizes)
}

update_params_with_sgd_momentum = function(list[unknown] model,
                                           list[unknown] gradients,
                                           double lr, double mu,
                                           list[unknown] optim_state)
    return (list[unknown] model_upd, list[unknown] optim_state_upd) {
    /*
     * Updates all learnable parameters with the sgd momentum optimizer.
     *
     * Inputs:
     * - model: Model parameters, same as for forward and
     *     backward pass.
     * - gradients: Gradients, returned from the backward pass.
     * - lr: Learning rate.
     * - mu: Momentum value.
     *     Typical values are in the range of [0.5, 0.99], usually
     *     started at the lower end and annealed towards the higher end.
     * - optim_state: Optimizer states for all model parameters.
     *
     * Outputs:
     * - model_upd: Updated model parameters.
     * - optim_state_upd: Updated model states for all parameters.
     */
    layer_sizes = list(3, 8, 36, 3)
    hyper_params = list(lr, mu)
    [optim_state_upd, model_upd] = util::update_params("sgd_momentum", optim_state, hyper_params, gradients, model,
        "bottleneck", layer_sizes)
}

init_sgd_nesterov_optim_params = function(int classes)
    return(list[unknown] params) {
    /*
     * Initializes the state of the nesterov optimizer for every
     * learnable parameter of ResNet 152.
     *
     * Inputs:
     * - classes: Number of network output classes.
     *
     * Outputs:
     * - params: List of state parameters with the same structure
     *     as weights of the forward and backward pass. It can be
     *     directly passed to the update parameter function.
     */
    layer_sizes = list(3, 8, 36, 3)
    params = util::init_optim("nesterov", classes, "bottleneck", layer_sizes)
}

update_params_with_sgd_nesterov = function(list[unknown] model,
                                           list[unknown] gradients,
                                           double lr, double mu,
                                           list[unknown] optim_state)
    return (list[unknown] model_upd, list[unknown] optim_state_upd) {
    /*
     * Updates all learnable parameters with the sgd nesterov optimizer.
     *
     * Inputs:
     * - model: Model parameters, same as for forward and
     *     backward pass.
     * - gradients: Gradients, returned from the backward pass.
     * - lr: Learning rate.
     * - mu: Momentum value.
     *     Typical values are in the range of [0.5, 0.99], usually
     *     started at the lower end and annealed towards the higher end.
     * - optim_state: Optimizer states for all model parameters.
     *
     * Outputs:
     * - model_upd: Updated model parameters.
     * - optim_state_upd: Updated model states for all parameters.
     */
    layer_sizes = list(3, 8, 36, 3)
    hyper_params = list(lr, mu)
    [optim_state_upd, model_upd] = util::update_params("sgd_nesterov", optim_state, hyper_params, gradients, model,
        "bottleneck", layer_sizes)
}

init_lars_optim_params = function(int classes)
    return(list[unknown] params) {
    /*
     * Initializes the state of the LARS optimizer for every
     * learnable parameter of ResNet 50.
     *
     * Inputs:
     * - classes: Number of network output classes.
     *
     * Outputs:
     * - params: List of state parameters with the same structure
     *     as weights of the forward and backward pass. It can be
     *     directly passed to the update parameter function.
     */
    layer_sizes = list(3, 8, 36, 3)
    params = util::init_optim("lars", classes, "bottleneck", layer_sizes)
}

update_params_with_lars = function(list[unknown] model, list[unknown] gradients,
                                   double lr, double mu, double weight_decay,
                                   double trust_coeff, list[unknown] optim_state)
    return (list[unknown] model_upd, list[unknown] optim_state_upd) {
    /*
     * Updates all learnable parameters with the LARS optimizer.
     *
     * LARS (Layer-wise Adaptive Rate Scaling) applies different learning
     * rates to different layers based on the ratio of parameter norm
     * to gradient norm, enabling stable large-batch training.
     *
     * Inputs:
     * - model: Model parameters, same as for forward and backward pass.
     * - gradients: Gradients, returned from the backward pass.
     * - lr: Global learning rate.
     * - mu: Momentum value. Recommended: 0.9
     * - weight_decay: L2 regularization strength. Recommended: 5e-4
     * - trust_coeff: Trust coefficient for LARS. Recommended: 0.001
     * - optim_state: Optimizer states for all model parameters.
     *
     * Outputs:
     * - model_upd: Updated model parameters.
     * - optim_state_upd: Updated model states for all parameters.
     */
    layer_sizes = list(3, 8, 36, 3)
    hyper_params = list(lr, mu, weight_decay, trust_coeff)
    [optim_state_upd, model_upd] = util::update_params("lars", optim_state, hyper_params, gradients, model, "bottleneck", 
        layer_sizes)
}