try:
    from tensorboardX import SummaryWriter
except Exception as e:
    print('no tensorboardX installed')
    try:
        from torch.utils.tensorboard import SummaryWriter
    except Exception as e:
        print('no torch.utils.tensorboard installed')
        raise ImportError(
            'Please run "pip install future tensorboard" to install '
            'the dependencies to use torch.utils.tensorboard '
            '(applicable to PyTorch 1.1 or higher)')
import os
import os.path as osp
import time
import numpy as np

writers = None

# filters = ['episodes_total', 'episodes_this_iter']


# def result_parser(result_dict):
#     ret_dict = {}
#     for k, v in result_dict.items():
#         if 'episode' in k and isinstance(v, (float, int)) and k not in filters:
#             ret_dict[k] = v
#         if 'policy' in k and isinstance(v, dict) and v:
#             ret_dict[k] = v
#         try:
#             for policy, info in result_dict['info']['learner'].items():
#                 tmp_dict = {}
#                 for index_name, index_v in info.items():
#                     if isinstance(index_v, (float, int)):
#                         tmp_dict[index_name] = index_v
#                 ret_dict[policy] = tmp_dict
#         except:
#             pass
#     return ret_dict


def summary_write(result_dict, step_num, log_dir):
    global writers
    if not writers:
        if not osp.exists(log_dir):
            os.makedirs(log_dir)
        writers = dict()

    # result_dict = result_parser(result_dict)

    print('result_dict', result_dict)
    for tag, val in result_dict.items():
        if not tag in writers:
            writer = SummaryWriter(osp.join(log_dir, str(tag)))
            writers[tag] = writer
            # print(f'tag {tag} not in writers {writers}')
            # continue
        writer = writers[tag]
        if isinstance(val, (float, int)) and not np.isnan(val):
            # writer.add_text(tag, val, get_iter(runner))
            values = {'data': val}
        elif isinstance(val, dict):
            values = dict()
            for key, value in val.items():
                if isinstance(value, (float, int)) and not np.isnan(value):
                    values[key] = value
        else:
            continue

        writer.add_scalars(tag, values, int(step_num))

        for index_nm, index_v in values.items():
            # add events.bin
            path = osp.join(log_dir, str(
                tag), index_nm, "events.bin")
            value = [int(step_num), index_v, float(time.time())]
            if osp.exists(path):
                save_data = np.fromfile(path, dtype=np.float)
            else:
                os.makedirs(osp.dirname(path), exist_ok=True)
                save_data = np.array([])
            save_data = np.insert(save_data, len(save_data), value)
            save_data.tofile(path)


def summary_close():
    global writers
    if writers:
        for writer in writers.values():
            writer.close()
