import os
import numpy as np
from utils import parse_log
from utils import cal_rel_err
from utils import draw_multi_lines, draw_single_line


def draw_loss(loss_title, rel_err_title, log_file1, log_file2):
    loss_save_path = os.path.join(os.path.dirname(__file__), loss_title)
    rel_err_save_path = os.path.join(os.path.dirname(__file__), rel_err_title)
    log_msg1, _, _ = parse_log(log_file1)
    log_msg2, _, _ = parse_log(log_file2)
    loss1 = np.array([log[3] for log in log_msg1.values()])
    loss2 = np.array([log[3] for log in log_msg2.values()])
    rel_err = cal_rel_err(loss1, loss2)

    draw_multi_lines(loss1, loss2, loss_title, "steps", "loss", loss_save_path)
    draw_single_line(rel_err, rel_err_title, "steps", "rel-err", rel_err_save_path)


model_name=""
loss_title=f"{model_name}-loss"
rel_err_title=f"{model_name}-rel-err"
log_file1=""
log_file2=""

if __name__ == '__main__':
    draw_loss(loss_title, rel_err_title, log_file1, log_file2)
