from .utils import format_time
import numpy as np
# TODO rewrite the whole progress bar functionality as class


def training_progress(iteration, total, loss, elapsed, length=50, fill='=', end='\n'):
    """
    Creates a custom training progress bar in the console.

    Parameters:
    - iteration (int): current iteration (from 1 to total)
    - total (int): total number of iterations
    - loss (float): batch loss
    - elapsed (float): elapsed time in seconds
    - length (int): character length of the bar
    - fill (str): character to use for the filled part of the bar
    - end (str): character to use after progress bar is done, iteration == total
    """
    percent = "{0:.1f}%".format(100 * (iteration / float(total)))
    filled_length = int(length * iteration // total)
    bar = fill * filled_length + '-' * (length - filled_length)
    message = f' Batch {iteration :>{len(str(total))}}/{total} [{bar}] {percent:<6} time: {format_time(elapsed)[6:]}' \
              f' Loss: {np.format_float_scientific(loss, unique=False, exp_digits=2, precision=4)}'
    if iteration == 1:
        print(message, end='')
    elif iteration == total:
        print('\b' * len(message), end='')
        print(message, end=end)
    else:
        print('\b' * len(message), end='')
        print(message, end='')