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

get_lr_with_warmup = function(double base_lr, int epoch, int iter, int total_epochs,
                              int iters_per_epoch, int batch_size, int base_batch_size,
                              int warmup_epochs, int decay_power)
    return (double lr) {
    /*
     * Compute learning rate with linear warmup and polynomial decay.
     * 
     * Implements the learning rate schedule from LARS paper:
     * - Linear warmup for first warmup_epochs
     * - Polynomial decay afterwards
     * - Linear scaling with batch size
     */
    
    # Scale learning rate linearly with batch size
    scaled_lr = base_lr * batch_size / base_batch_size
    
    # Total number of iterations
    total_iters = total_epochs * iters_per_epoch
    warmup_iters = warmup_epochs * iters_per_epoch
    current_iter = (epoch - 1) * iters_per_epoch + iter
    
    if (current_iter <= warmup_iters) {
        # Linear warmup
        lr = scaled_lr * current_iter / warmup_iters
    } else {
        # Polynomial decay
        decay_iters = total_iters - warmup_iters
        decay_current = current_iter - warmup_iters
        decay_factor = (1 - decay_current / decay_iters) ^ decay_power
        lr = scaled_lr * decay_factor
    }
}

