'''Some helper functions for PyTorch, including:
    - progress_bar: progress bar mimic xlua.progress.
    - set_lr : set the learning rate
    - clip_gradient : clip gradient
'''

import os
import sys
import time
import math
import torch
import torch.nn as nn
import torch.nn.init as init
from torch.autograd import Function

#获取控制台行、列数
if sys.platform == 'win32':
	term_width = 80
else:
	print('###', os.popen('stty size', 'r').read())
	_, term_width = os.popen('stty size', 'r').read().split()
	term_width = int(term_width)

TOTAL_BAR_LENGTH = 30.
last_time = time.time()
begin_time = last_time

#[==>........ 19/225 ...........] | Loss: 1.961 | Acc: 22.000% (537/2432)  
def progress_bar(current, total, msg=None):
    global last_time, begin_time
    if current == 0:
        begin_time = time.time()  # Reset for new bar.

    cur_len = int(TOTAL_BAR_LENGTH*current/total)
    rest_len = int(TOTAL_BAR_LENGTH - cur_len) - 1

    sys.stdout.write(' [')
    for i in range(cur_len):
        sys.stdout.write('=')
    sys.stdout.write('>')
    for i in range(rest_len):
        sys.stdout.write('.')
    sys.stdout.write(']')

    cur_time = time.time()
    step_time = cur_time - last_time
    last_time = cur_time
    tot_time = cur_time - begin_time

    L = []
    if msg:
        L.append(' | ' + msg)

    msg = ''.join(L)
    sys.stdout.write(msg)
    for i in range(term_width-int(TOTAL_BAR_LENGTH)-len(msg)-3):
        sys.stdout.write(' ')

    # Go back to the center of the bar.
    for i in range(term_width-int(TOTAL_BAR_LENGTH/2)+2):
        sys.stdout.write('\b')
    sys.stdout.write(' %d/%d ' % (current+1, total))

    if current < total-1:
        sys.stdout.write('\r')
    else:
        sys.stdout.write('\n')
    sys.stdout.flush()

def set_lr(optimizer, lr):
    for group in optimizer.param_groups:
        group['lr'] = lr

def clip_gradient(optimizer, grad_clip):
    for group in optimizer.param_groups:
        #print(group['params'])
        for param in group['params']:
            param.grad.data.clamp_(-grad_clip, grad_clip)

[((25, 12, 25, 23), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((37, 4, 37, 26), 'sys.stdout.write', 'sys.stdout.write', ({(37, 21, 37, 25): '""" ["""'}, {}), "(' [')", False, 'import sys\n'), ((40, 4, 40, 25), 'sys.stdout.write', 'sys.stdout.write', ({(40, 21, 40, 24): '""">"""'}, {}), "('>')", False, 'import sys\n'), ((43, 4, 43, 25), 'sys.stdout.write', 'sys.stdout.write', ({(43, 21, 43, 24): '"""]"""'}, {}), "(']')", False, 'import sys\n'), ((45, 15, 45, 26), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((55, 4, 55, 25), 'sys.stdout.write', 'sys.stdout.write', ({(55, 21, 55, 24): 'msg'}, {}), '(msg)', False, 'import sys\n'), ((62, 4, 62, 52), 'sys.stdout.write', 'sys.stdout.write', ({(62, 21, 62, 51): "(' %d/%d ' % (current + 1, total))"}, {}), "(' %d/%d ' % (current + 1, total))", False, 'import sys\n'), ((68, 4, 68, 22), 'sys.stdout.flush', 'sys.stdout.flush', ({}, {}), '()', False, 'import sys\n'), ((32, 21, 32, 32), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((39, 8, 39, 29), 'sys.stdout.write', 'sys.stdout.write', ({(39, 25, 39, 28): '"""="""'}, {}), "('=')", False, 'import sys\n'), ((42, 8, 42, 29), 'sys.stdout.write', 'sys.stdout.write', ({(42, 25, 42, 28): '"""."""'}, {}), "('.')", False, 'import sys\n'), ((57, 8, 57, 29), 'sys.stdout.write', 'sys.stdout.write', ({(57, 25, 57, 28): '""" """'}, {}), "(' ')", False, 'import sys\n'), ((61, 8, 61, 30), 'sys.stdout.write', 'sys.stdout.write', ({(61, 25, 61, 29): '"""\x08"""'}, {}), "('\\x08')", False, 'import sys\n'), ((65, 8, 65, 30), 'sys.stdout.write', 'sys.stdout.write', ({(65, 25, 65, 29): "'\\r'"}, {}), "('\\r')", False, 'import sys\n'), ((67, 8, 67, 30), 'sys.stdout.write', 'sys.stdout.write', ({(67, 25, 67, 29): '"""\n"""'}, {}), "('\\n')", False, 'import sys\n'), ((20, 14, 20, 40), 'os.popen', 'os.popen', ({(20, 23, 20, 34): '"""stty size"""', (20, 36, 20, 39): '"""r"""'}, {}), "('stty size', 'r')", False, 'import os\n'), ((21, 17, 21, 43), 'os.popen', 'os.popen', ({(21, 26, 21, 37): '"""stty size"""', (21, 39, 21, 42): '"""r"""'}, {}), "('stty size', 'r')", False, 'import os\n')]