# Copyright (C) 2024 Habana Labs, Ltd. an Intel Company.

import os
import glob
import time

from optimum.habana.utils import to_gb_rounded
import habana_frameworks.torch as htorch

START_TS = None
DBG_TRACE_FILENAME = os.environ.get('DBG_TRACE_FILENAME')
if 'GRAPH_VISUALIZATION' in os.environ:
    for f in glob.glob('.graph_dumps/*'):
        os.remove(f)


def count_hpu_graphs():
    return len(glob.glob('.graph_dumps/*PreGraph*'))


def dbg_trace(tag, txt):
    global START_TS
    if DBG_TRACE_FILENAME is not None and int(os.getenv("RANK", 0)) == 0:
        if START_TS is None:
            START_TS = time.perf_counter()
        time_offset = time.perf_counter() - START_TS
        mem_stats = htorch.hpu.memory.memory_stats()
        mem_used = to_gb_rounded(mem_stats['InUse'])
        max_mem_used = to_gb_rounded(mem_stats['MaxInUse'])
        print(f'ts:{time_offset:.3f}s g:{count_hpu_graphs()} mu:{mem_used:.1f}GB '
              f'mmu:{max_mem_used:.1f}GB | {tag} | {txt}', flush=True, file=open(DBG_TRACE_FILENAME, 'a'))