import re
from cmd_runner import cmd_runner


def command_train(command):
    """
    针对darknet的训练方式，返回打印的日志、学习率和map
    :param file: 命令
    :return:
    """
    def line_process_fn(line):
        if 'avg loss' in line:
            match = re.findall(
                r'(\d+): (\d+\.\d+), (\d+\.\d+) avg loss', line)
            if len(match) == 1 and len(match[0]) == 3:
                step, loss, avg_loss = int(match[0][0]), float(
                    match[0][1]), float(match[0][2])
        if 'mean average precision (mAP@0.50) = ' in line:
            match = re.findall(
                r'mean average precision \(mAP@0.50\) = (\d+\.\d+)', line)
            if len(match) == 1:
                map_temp = float(match[0])*100
        print(line, step, loss, avg_loss, map_temp)
    cmd_runner(command=command, line_process_fn=line_process_fn,
               print_line=False)


# def command_map(self, command):
#     """
#     针对darknet的计算方式，返回打印的日志、学习率和map
#     :param file: 命令
#     :return:
#     """
#     map_temp, res = None, dict()
#        if 'mean average precision (mAP@' in line:
#             match = re.findall(
#                 r'mean average precision \(mAP@(.*)\) = (\d+\.\d+)', line)
#             if len(match) == 1:
#                 map_temp = float(match[0][1])
#                 res["map"] = map_temp
#         if "class_id" in line and "name" in line and "ap" in line:
#             match = re.findall(
#                 r'class_id = (\d), name = (.*), ap = (\d+\.\d+)%', line)
#             if len(match) == 1:
#                 name, ap_temp = match[0][1], float(match[0][2])/100
#                 res[name] = ap_temp
#     cmd_run(command=command, line_process_fn=line_process_fn,
#             print_line=False)
#     return res
