'''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, 'time.time', 'time.time', False, 'import time\n'), (37, 'sys.stdout.write', 'sys.stdout.write', False, 'import sys\n'), (40, 'sys.stdout.write', 'sys.stdout.write', False, 'import sys\n'), (43, 'sys.stdout.write', 'sys.stdout.write', False, 'import sys\n'), (45, 'time.time', 'time.time', False, 'import time\n'), (55, 'sys.stdout.write', 'sys.stdout.write', False, 'import sys\n'), (62, 'sys.stdout.write', 'sys.stdout.write', False, 'import sys\n'), (68, 'sys.stdout.flush', 'sys.stdout.flush', False, 'import sys\n'), (32, 'time.time', 'time.time', False, 'import time\n'), (39, 'sys.stdout.write', 'sys.stdout.write', False, 'import sys\n'), (42, 'sys.stdout.write', 'sys.stdout.write', False, 'import sys\n'), (57, 'sys.stdout.write', 'sys.stdout.write', False, 'import sys\n'), (61, 'sys.stdout.write', 'sys.stdout.write', False, 'import sys\n'), (65, 'sys.stdout.write', 'sys.stdout.write', False, 'import sys\n'), (67, 'sys.stdout.write', 'sys.stdout.write', False, 'import sys\n'), (20, 'os.popen', 'os.popen', False, 'import os\n'), (21, 'os.popen', 'os.popen', False, 'import os\n')]