import os
import ujson as json
import networkx as nx
import time
import collections
import re
from tqdm import tqdm

from .dag_utils import QueueType
from .trace_utils import *
from .logger_utils import SingleLogger
from .debug_utils import DebugRecorder
from .bps_helper.graph import PS_COMP_OPS_SETS, PS_COMM_OPS_SETS

FIXED_GAP_us = 5

def short_name(n):
    return n.split(".")[-1]

comm_op_types = []
try:
    from hvd.graph import *
except:
    pass

try:
    from bps_helper.graph import *
    comm_op_types = [PS_COMM_OPS.PULL_RES, PS_COMP_OPS.COPY_MERGED, PS_COMP_OPS.SUM, PS_COMP_OPS.COPY_FIRST,
        PS_COMM_OPS.PULL_REQ, PS_COMM_OPS.PUSH_RES, PS_COMM_OPS.PUSH_REQ]
except:
    pass

def ret_priority(n_):
    ### The smaller the rank is, the higher priority the node has
    if "FW" in n_:
        return 0
    elif "OUTPUT" in n_:
        return 1
    elif "BW" in n_:
        return 2
    elif "UPDATE_" in n_:
        return 3
    elif "Comm" in n_:
        # for comm_op_type in [PS_COMM_OPS.PUSH_REQ, PS_COMM_OPS.PULL_REQ]:
        #     if comm_op_type in n_:
        #         return 4
        # for comm_op_type in [PS_COMM_OPS.PUSH_RES, PS_COMM_OPS.PULL_RES]:
        #     if comm_op_type in n_:
        #         return 5
        # for idx, comm_op_type in enumerate(comm_op_types):
        #     if comm_op_type in n_:
        #         return 4 + idx
        # if PS_COMM_OPS.PULL_REQ in n_:
        #     return 4
        return 100
    else:
        return 100

def _repr_of_tensor(tensor_name):
    return min([int(tensor_id) for tensor_id in tensor_name.split("+")])

class Device:
    def __init__(self, device_name, _replayer, 
            infi_para=False,
            comm_backend = "NCCL",
            init_device_time_in_ms=0,
            full_trace=False
        ):
        self.replayer = _replayer
        self.init_device_time = init_device_time_in_ms * 1000
        self.device_time = self.init_device_time
        self.device_name = device_name
        #! infi_para devices allow to excute in infinite parallelism
        self.infi_para = infi_para
        ### Used to record the last event generated by this device
        self.prev_name_dur = None
        self.comm_backend = comm_backend

        ### nodes to be executed, in a **partial order**
        self.queue = []

        self.full_trace = full_trace

    def reset(self):
        self.device_time = self.init_device_time
        self.prev_name_dur = None
        self.queue = []

    def real_start_t(self, _last_end_time):
        return max(_last_end_time, self.device_time)

    def exct(self, name, _last_end_time, step_idx):
        ''' Execute one op on this device 

        Parameters
        ----
        name: str
            The name of the op to run
        _last_end_time: float
            The time when this op can start to run, 
            i.e., all the dependent ops have been done
        '''
        ### for debug
        DebugRecorder().debug_event_start()
        if not self.infi_para:
            start_t = self.real_start_t(_last_end_time)
        else:
            # raise NotImplementedError("Infi para is not yet implemented.")
            start_t = _last_end_time

        if name == "END":
            #! No event is generated, but has successors
            self.mark_as_exct(name, start_t, start_t)
            return

        ### Really start to execute
        avg = self.replayer.dag.nodes[name]["avg"]
        if "+" in name and "Comm" not in name:
            pid, _, cat, suffix = parse_allinfo_from_name(name.split("+")[0])
            raw_name = "+".join([parse_allinfo_from_name(_name)[1] for _name in name.split("+")])
        else:
            pid, raw_name, cat, suffix = parse_allinfo_from_name(name)
        delay, ratio = self.get_delay_para(name)
        duration = (1000.0 * max(avg + delay, 0)) * ratio
        if self.comm_backend == "BYTEPS" and "UPDATE_CAL" in name:
            duration = 0
        
        ### Construct execution graphs
        ### 1. Add new edges according to the execution order
        if not self.infi_para and self.prev_name_dur is not None and \
            (self.prev_name_dur[0], name) not in self.replayer.exct_dag.edges and \
            not ("BW" in self.prev_name_dur[0] and "UPDATE_" in name):
            self.replayer.exct_dag.add_edge(self.prev_name_dur[0], name, weight=(
                (start_t - self.prev_name_dur[2]) / 1000.0))
                    
        if duration == 0 and not self.full_trace:
            self.prev_name_dur = (name, 0, start_t)
        else:
            event = {
                        "name": raw_name if self.replayer.name2mapping_fn is None else self.replayer.name2mapping_fn(raw_name),
                        "ts": start_t,
                        "dur": duration,
                        "pid": pid,
                        "cat": cat,
                        "ph": "X",
                        "tid": self.device_name,
                        "args": {
                            "name": name,
                            "cnt": step_idx
                        }
                    }
            ### Expose dependency info in trace["args"], time-consuming
            if self.full_trace:
                _id = 0
                for prev, _ in self.replayer.dag.in_edges(name):
                    event["args"]["input%d"%_id] = prev
                    _id += 1
            
            self.prev_name_dur = (event["args"]["name"], event['dur'], event['ts'])
            # ### 2. Update edge weight
            # for next_ in self.replayer.exct_dag.successors(name):
            # 	self.replayer.exct_dag.edges[name, next_]["weight"] = duration / 1000.0
            self.replayer.rst_traces.append(event)
        DebugRecorder().debug_event_start()

        self.mark_as_exct(name, start_t, start_t + duration)
        DebugRecorder().debug_event_end(name, self.device_name, "mark_as_exct")
        pid = parse_pid_from_name(name)
        self.replayer.step_end_time[pid] = start_t + duration
        #! TODO: for debug
        DebugRecorder().debug_event_end(name, self.device_name, "exct")
    
    def _update_device_time(self, name, _end_time):
        ### Apply the gap between two nodes
        gap = 0
        for key, value in self.replayer.dag.nodes[name].items():
            if "GAP" in key:
                ### e.g. "gap.operator.operator"
                key_s = key.split("GAP")
                ### TODO (huhanpeng): does this fit to the BytePS or use intra-gap instead
                if key_s[0] == key_s[1]:
                    gap += value
                    if gap > 1000:
                        SingleLogger().debug("Large GAP detected: {}, key = {}, gap = {}".format(name, key, value))
        if gap < 0:
            raise RuntimeError(
                "Negative GAP detected: {}, gap = {}".format(name, gap))
        # TODO (huhanpeng): decide the minimun gap automatically
        # self.device_time = _end_time + max(gap, 20)
        self.device_time = _end_time + gap

    def mark_as_exct(self, name, _start_t, _end_time):
        ''' Mark that the op has been executed '''
        self._update_device_time(name, _end_time)
        self.replayer.node_status.pop(name)
        this_cat = parse_cat_from_name(name)
        for _succ in self.replayer.dag.successors(name):
            next_cat = parse_cat_from_name(_succ)
            if _succ in self.replayer.node_status:
                _status = self.replayer.node_status[_succ]
                ### Calculate the ready time
                if self.comm_backend == "NCCL" and ("SEND" in name and "RECV" in _succ):
                    ### For Send->Recv edge, there exist some overlap
                    ### TODO (huhanpeng): how do decide the end time of the RECV event
                    avg = self.replayer.dag.nodes[_succ]["avg"]
                    gap = 0
                    _status["ready"] = (
                        _end_time + gap) if _status["ready"] is None else max(_end_time + gap, _status["ready"])
                else:
                    ## For BYTEPS and Horovod, Only apply BW->Comm gaps
                    ## Other gaps should be considered with the device time.
                    gap = 0
                    if GAP_STR_OP2COMM in self.replayer.dag.nodes[name] and next_cat == CatName.COMM.value \
                         and (this_cat == CatName.OPERATOR.value or "Sync" in name):
                        gap += self.replayer.dag.nodes[name][GAP_STR_OP2COMM]
                        if self.replayer.dag.nodes[name][GAP_STR_OP2COMM] > 10000:
                            SingleLogger().debug("Large OP2COMM gap detected, {} -> {},  gap: {}".format(name, _succ, self.replayer.dag.nodes[name][GAP_STR_OP2COMM]))
                    _status["ready"] = (_end_time + gap) if _status["ready"] is None else max(_end_time + gap, _status["ready"])

                ### Whether the dependency has met
                _status["in_degree"] -= 1
                # self.replayer.debuger.mark_as_exct(name, _succ)
                if _status["in_degree"] == 0:
                    if _status["ready"] is None:
                        raise RuntimeError("{}\'s ready time is not decided".format(_succ))
                    self.replayer.insert_next_node(_succ, _status["ready"])
        # self.replayer.debuger.show_staue()

    def get_delay_para(self, name_):
        #! Get the delay parameters.
        delay = 0
        ratio = 1.0
        if self.replayer.delay_dict is not None:
            cat = parse_cat_fine_grained(name_)
            if name_ in self.replayer.delay_dict:
                delay = self.replayer.delay_dict[name_]["delay"]
                ratio = self.replayer.delay_dict[name_]["ratio"]
            elif "DELAY_ALL_CMP" in self.replayer.delay_dict and cat in COMP_CAT:
                delay = self.replayer.delay_dict["DELAY_ALL_CMP"]["delay"]
                ratio = self.replayer.delay_dict["DELAY_ALL_CMP"]["ratio"]
            elif "DELAY_ALL_COMM" in self.replayer.delay_dict and cat in COMM_CAT:
                delay = self.replayer.delay_dict["DELAY_ALL_COMM"]["delay"]
                ratio = self.replayer.delay_dict["DELAY_ALL_COMM"]["ratio"]
            elif "DELAY_ALL" in self.replayer.delay_dict:
                delay = self.replayer.delay_dict["DELAY_ALL"]["delay"]
                ratio = self.replayer.delay_dict["DELAY_ALL"]["ratio"]
        return delay, ratio

class CommKernelDevice(Device):
    '''For horovod Communication Kernels, can be occupied by a tensor'''
    def __init__(self, *args, **kwargs):
        '''period: cycle time in ms'''
        super().__init__(*args, **kwargs)
        self.lock = None
        self.blocked = []

    def acquire_lock(self, name):
        if self.lock is None:
            ### prefix->op_type.op_name.sub_op~>suffix
            self.lock = parse_op_name(name)
            return True
        elif self.lock == parse_op_name(name):
            return True
        else:
            return False

    def release_lock(self):
        self.lock = None

    def exct(self, name, _last_end_time, step_idx):
        if self.acquire_lock(name):
            super().exct(name, _last_end_time, step_idx)
            if QueueType().ret_list()[-1] in name:
                ### release the device if this is the last sub_op of this tensor
                self.release_lock()
                if len(self.blocked) > 0:
                    ### pop an operator from the blocked queue to execute
                    ### since these block operators have high priority (early start time)
                    ### they are expected to be execute immediately
                    self.replayer.insert_next_node(*self.blocked.pop(0))
        else:
            ### This device is occupied by another tensor, blocked
            self.blocked.append((name, _last_end_time))

class PeriodicDevice(Device):
    '''For horovod negotiation, simulate cycles in Horovod'''
    def __init__(self, *args, period=5, **kwargs):
        '''period: cycle time in ms'''
        super().__init__(*args, **kwargs)
        self.period = period * 1000

    def real_start_t(self, _last_end_time):
        while self.device_time < _last_end_time:
            self.device_time += self.period
        return self.device_time

    def _update_device_time(self, name, _end_time):
        while self.device_time < _end_time:
            self.device_time += self.period

class PSCommDevice(Device):
    '''Handles BytePS communication'''
    def __init__(self, op_counter, device_name, *args, **kwargs):
        super().__init__(device_name, *args, **kwargs)
        self.op_counter = op_counter
        self.source = device_name.split("::")[0]
        self.target = device_name.split("::")[1].split(DEL)[0]

    def mark_as_exct(self, name, _start_t, _end_time):
        next_name = self.op_counter.get_next_op(name)
        self._update_device_time(name, _end_time)
        self.replayer.node_status.pop(name)

        this_cat = parse_cat_from_name(name)
        actual_successors = list(self.replayer.dag.successors(name))
        if next_name is not None:
            assert next_name in self.replayer.dag.nodes()
            actual_successors.append(next_name)
            # add an edge from this name to next_name in exct dag
            self.replayer.exct_dag.add_edge(name, next_name, weight=((_end_time - _start_t) / 1000.0))

        for _succ in actual_successors:
            next_cat = parse_cat_from_name(_succ)
            if _succ in self.replayer.node_status:
                _status = self.replayer.node_status[_succ]
                _status["in_degree"] -= 1
                # add GAPs
                gap = 0
                if GAP_STR_INTERNODE in self.replayer.dag.nodes[name] and \
                        this_cat == CatName.COMM.value and \
                        next_cat == CatName.COMM.value:
                    gap += self.replayer.dag.nodes[name][GAP_STR_INTERNODE]
                # gap = 0
                _status["ready"] = (_end_time + gap) if _status["ready"] is None else max(_end_time + gap, _status["ready"])

                if _status["in_degree"] == 0:
                    self.replayer.insert_next_node(_succ, _status["ready"])

class Replayer:
    def __init__(self, 
            dag, _step_num, leaf_dirs,
            dump_path, comm_backend, byteps_graph, 
            infi_para_update=False,
            show_queue=False,
            recd_topo_order=False,
            partial=False,
            name2mapping_fn=None,
            full_trace=False):
        self.dag = dag
        self.infi_para_update = infi_para_update
        # self.preprocess_dag()
        self.step_num = _step_num
        self.leaf_dirs = leaf_dirs
        self.dump_path = dump_path
        self.comm_backend = comm_backend
        self.byteps_graph = byteps_graph

        self.logger = SingleLogger()
        ### Delay information, the unit of 'delay' field should be ms
        self.delay_dict = None
        ### maintain node status
        self.node_status = {}
        self.device_dict = {}
        self.queue_status = None

        self.reset_replayer()
        if self.comm_backend == "BYTEPS":
            self.op_counter = ServerOpCounter(self.byteps_graph)

        self.show_queue = show_queue
        self.show_queue = True
        if not self.show_queue:	
            self.queue = []
        
        ### Decide whether to record the topological order
        self.topo_ord = []
        self.recd_topo_order = recd_topo_order

        self.replay_done = False

        ### If the replayer is used to partially evaluate a dag, we do not check
        #   the node type when initializing ready nodes at the beginning of simulation
        self.partial = partial

        self.name2mapping_fn = name2mapping_fn

        self.allow_comm_init_frontier = False
        self.full_trace = full_trace

    def ret_topo_ord(self):
        assert self.replay_done, "The DFG has not been replayed."
        return self.topo_ord

    def preprocess_dag(self):
        if self.infi_para_update:
            bw_nodes = []
            update_nodes = []
            for node in self.dag.nodes:
                if "BW." in node:
                    bw_nodes.append(node)
                elif "UPDATE_." in node:
                    update_nodes.append(node)
            self.dag.add_node("BW_BARRIER", avg=0)
            for node in bw_nodes:
                self.dag.add_edge(node, "BW_BARRIER", avg=0)
            for node in update_nodes:
                self.dag.add_edge("BW_BARRIER", node, avg=0)

    def pre_prepare(self):
        ''' Initialize nodes that need to be replayed first
        '''
        def map_in_degree(n):
            if self.comm_backend == "BYTEPS":
                if self.byteps_graph.is_server_comp(n):
                    return self.dag.in_degree(n) + 1
            return self.dag.in_degree(n)

        self.node_status = []
        for n in self.dag.nodes():
            self.node_status.append((n, {"in_degree": map_in_degree(n), "ready": None}))
            if self.show_queue:
                ### if we want to show queue statuts, need to initialize the device list first
                self.name2device(n)
        self.node_status = dict(self.node_status)

        ### prepare nodes to be executed on each device
        for n, _status in self.node_status.items():
            if _status["in_degree"] == 0:
                if not self.allow_comm_init_frontier and not self.partial and CatName.COMM.value in n:
                    if input("Invalid nodes {} with in_degree=0, Continue: (Y/n)".format(n)).lower() in ["", "y", "yes"]:
                        self.allow_comm_init_frontier = True
                    else:
                        exit(1)
                pid = parse_pid_from_name(n)
                _last_end = self.step_end_time[pid] if _status["ready"] is None else _status["ready"]
                self.insert_next_node(n, _last_end)
    
    def replay_one_iter(self, step_idx):
        self.debuger = ReplayDebuger(self)
        # self.debuger.monitor_node()
        self.pre_prepare()
        while True:
            if self.pop_one_node_exec(step_idx) == 1:
                break
        DebugRecorder().dump_traces(".")
        
        self.replay_done = True

    def pop_one_node_exec(self, step_idx):
        if self.show_queue:
            sorted_device = sorted(self.device_dict.values(), key=lambda x: x.device_time)
            for device in sorted_device:
                if len(device.queue) == 0:
                    continue
                (n, t) = device.queue.pop(0)
                if self.recd_topo_order:
                    self.topo_ord.append((n, t))
                self.record_queue_status(t)
                device.exct(n, t, step_idx)
                return 0
            return 1  ### no operators to execute
        else:
            if len(self.queue) == 0:
                return 1
            
            ### check the frontier
            # bw_ops = set([parse_op_name(op) for op, _ in self.queue if "BW" in op and "traces_0.rank0" in op])
            # if len(bw_ops) > 1:
            #     print(bw_ops)

            (n, t) = self.queue.pop(0)
            if self.recd_topo_order:
                self.topo_ord.append((n, t))
            device = self.name2device(n)
            device.exct(n, t, step_idx)
            return 0

    def record_queue_status(self, cur_time):
        if self.queue_status is None:
            self.queue_status = {"names": list(self.device_dict.keys()), 'data': []}
        self.queue_status['data'].append([cur_time] + [len(self.device_dict[name].queue) for name in self.queue_status['names']])

    def replay(self, _output=True, verbose=True):
        self.reset_replayer()
        _ts = time.time()
        for step_idx in range(self.step_num):
            self.replay_one_iter(step_idx)
        self.logger.info("Take %f s to replay one iteration" % ((time.time() - _ts)/float(self.step_num)))
        if _output:
            self.output_traces(verbose=verbose)
        return self.step_end_time
        
    def replayAndDelay(self, delay_dict_, _output=False, _path=None, verbose=True):
        self.reset_replayer()
        self.delay_dict = delay_dict_
        self.replay_one_iter(0)
        if _output:
            self.output_traces(_path=_path, verbose=verbose)
        return self.step_end_time

    def insert_next_node(self, n, t):
        ''' This is acutally equal to a scheduler of an **Engine**
        n: node string
        t: start time of this node, do NOT take the device time into consideration
        '''
        if self.show_queue:
            device = self.name2device(n)
            # just to substitute func to compare their ranks.
            self.insort_right(device.queue, (n, t), func=self._a_has_larger_prior)
        else:
            self.insort_right(self.queue, (n, t), func=self._a_has_larger_prior)

    def insort_right(self, a, x, lo=0, hi=None, func=None):
        """Insert item x in list a, and keep it sorted assuming a is sorted.
        If x is already in a, insert it to the right of the rightmost x.
        Optional args lo (default 0) and hi (default len(a)) bound the
        slice of a to be searched.
        """
        def fun_cmp(x1, x2):
            if func is None:
                return x1 < x2
            else:
                return func(x1, x2)

        if lo < 0:
            raise ValueError('lo must be non-negative')
        if hi is None:
            hi = len(a)
        while lo < hi:
            mid = (lo+hi)//2
            if fun_cmp(x, a[mid]):
                hi = mid
            else:
                lo = mid+1
        a.insert(lo, x)

    def relabel_map(self, name):
        if "Comm." not in name:
            return name
        grp, sub_op, suffix = parse_allinfo_from_name_v2(name)
        if sub_op in PS_COMM_OPS_SETS:
            part_id = int(re.findall("::(\d+)", suffix)[0])
        elif sub_op in PS_COMP_OPS_SETS:
            part_id = int(re.findall("::~PART(\d+)", suffix)[0])
        else:
            raise ValueError(name)
        server_id = self.byteps_graph.grp_part_id2server[grp][part_id]
        return re.sub("server_\d+", "server_{}".format(server_id), name)

    def name2device(self, n):
        pid = parse_pid_from_name(n)
        cat = parse_cat_from_name(n)
        if cat == "Comm":
            if self.comm_backend == "BYTEPS":
                if self.byteps_graph.grp_part_id2server is not None:
                    n = self.relabel_map(n)
                    pid = parse_pid_from_name(n)
                    cat = parse_cat_from_name(n)
                device_id = gen_long_name(pid, cat)
            else:
                if "SEND" in n or "RECV" in n:
                    device_id = gen_long_name(pid, cat)
                # if "SEND" in n:
                #     device_id = gen_long_name(pid, cat, "SEND")
                # elif "RECV" in n:
                #     device_id = gen_long_name(pid, cat, "RECV")
                elif "Sync" in n:
                    device_id = gen_long_name(pid, cat, "Sync")
                else:
                    device_id = None
                    in_queue_type = False
                    for sub_op in QueueType().ret_list():
                        if sub_op in n:
                            device_id = gen_long_name(pid, cat, "Kernel")
                            in_queue_type = True
                            break
                    
                    if not in_queue_type:
                        device_id = gen_long_name(pid, cat)
        elif "UPDATE" in n and self.infi_para_update:
            # update node, generate a new device
            device_id = gen_long_name(pid, cat, "INFI_PARA_UPDATE")
        else:
            device_id = gen_long_name(pid, cat)

        if device_id not in self.device_dict:
            if cat == CatName.COMM.value and self.comm_backend == "BYTEPS":
                self.device_dict[device_id] = self.create_ps_comm_device(device_id)
            elif "Sync" in device_id:
                self.device_dict[device_id] = self.create_device(device_id)
                # self.device_dict[device_id] = self.create_periodic_device(device_id)
            elif cat == CatName.COMM.value and "Kernel" in device_id:
                self.device_dict[device_id] = self.create_comm_kernel_device(device_id)
            elif "INFI_PARA_UPDATE" in device_id:
                self.device_dict[device_id] = self.create_device(
                    device_id, infi_para=True)
            else:
                self.device_dict[device_id] = self.create_device(device_id)

        return self.device_dict[device_id]		
    
    def _a_has_larger_prior(self, _a, _b):
        # if "Comm" in _a[0] and "Comm" in _b[0] and self.name2device(_a[0]) == self.name2device(_b[0]):
        #     tensor_a = parse_op_name(_a[0])
        #     tensor_b = parse_op_name(_b[0])
        #     if tensor_a == tensor_b:
        #         return _a[1] < _b[1]
        #     else:  
        #         return _repr_of_tensor(tensor_a) > _repr_of_tensor(tensor_b)
                
        _ap = ret_priority(_a[0])
        _bp = ret_priority(_b[0])
        if _ap == _bp:
            ### The same priority, compare the start time		
            return _a[1] < _b[1]
        else:
            return _ap < _bp

    def create_device(self, device_name, infi_para=False):
        # if device_name.startswith("traces_0."):
        #     init_device_time = 11.5
        #     SingleLogger().info(bcolors.CYELLOW + "Create device {} with {} init device time ... ".format(device_name, init_device_time) + bcolors.ENDC)
        #     return Device(device_name, self, comm_backend = self.comm_backend, infi_para=infi_para, init_device_time_in_ms=init_device_time)
        d = Device(device_name, self, 
                comm_backend = self.comm_backend, 
                infi_para=infi_para,
                full_trace=self.full_trace)
        return d

    def create_comm_kernel_device(self, device_name, infi_para=False):
        d = CommKernelDevice(device_name, self, 
                comm_backend = self.comm_backend,
                infi_para=infi_para,
                full_trace=self.full_trace)
        return d

    def create_periodic_device(self, device_name, infi_para=False):
        d = PeriodicDevice(device_name, self, 
                period=5, 
                comm_backend = self.comm_backend, 
                infi_para=infi_para, 
                full_trace=self.full_trace)
        return d

    def create_ps_comm_device(self, device_name, infi_para=False):
        d = PSCommDevice(self.op_counter, device_name, self,  
                comm_backend = self.comm_backend,
                infi_para=infi_para,
                full_trace=self.full_trace)
        return d

    def reset_replayer(self):
        # self.step_end_time = dict([(_d, 0.0) for _d in self.leaf_dirs])
        self.step_end_time = collections.defaultdict(float)
        self.rst_traces = []
        ### Reset all devices
        for _, device_ in self.device_dict.items():
            device_.reset()

        ### Ininitalize the execution graph as the depdency graph
        self.exct_dag = self.dag.copy()

        self.topo_ord = []
        self.replay_done = False

    def dump_critical_path(self, file, critical_path, prefix=None):
        if prefix is None:
            dump_path = os.path.join(self.dump_path, file)
        else:
            dump_path = os.path.join(prefix, file)
        painted_timeline(self.rst_traces, 
            lambda event: "Y" if event["args"]["name"] in critical_path else "N",
            dump_path)
    
    def paint_bw_comm_depend(self, _path=None, traceM=None, one_pid=None):
        ''' Paint the timeline to show the dependency between BW nodes and Comm nodes
        * E.g., if there is an edge from BW.A to Comm.1.Sync, BW.A would be renamed with Comm.1.Sync
        * This is useful when considering tensor fusion
        * NOTE: may not adapt to operator fusion
        '''
        final_trace = []
        
        def map_bw_2_comm(node_):
            ''' Some nodes may have not corresponding traces
            * E.g., BW.A -> BW.B --> Comm.1, but BW.B has no trace,
            *   in this case, we will paint BW.A to show the dependency
            '''
            ret = []
            succs = list(self.dag.successors(node_))
            for succ_ in succs:
                if "Comm" in succ_:
                    ret.append(succ_)
                else:
                    assert "BW" in succ_
                    if self.dag.nodes[succ_]["avg"] == 0:
                        ret += map_bw_2_comm(succ_)
            return list(set(ret))
        start_ts = None
        traces = traceM.traces if traceM is not None else self.rst_traces
        for trace in traces:
            if traceM is not None and trace["args"]["step"] != traceM.opt_step:
                continue
            if "BW" in trace["name"]:
                long_name = traceM.ret_unique_name(trace) if traceM is not None else trace["args"]["name"]
                comm_succs = map_bw_2_comm(long_name)
                # assert len(comm_succs) <= 1, (trace["name"], comm_succs)
                if len(comm_succs) > 0:
                    rawname = parse_op_name(comm_succs[0])
                    trace["name"] = "BW." + rawname
                else:
                    trace["name"] = "BW"
            elif "Comm" in trace["name"]:
                pass
            else:
                trace["name"] = "others"
                
            trace["pid"] += ".real" if traceM is not None else ".replay"
            if start_ts is None:
                start_ts = trace["ts"]
            trace["ts"] -= start_ts
            final_trace.append(trace)

        if _path is None:
            _path = os.path.join(self.dump_path, "bw2comm_timeline.json")
        with open(_path, 'w') as f:
            json.dump(final_trace, f)

    def output_traces(self, _path=None, verbose=True):
        #! Output the synthetic traces.
        rst = {
            "traceEvents": [],
            "displayTimeUnit": "ms"
        }
        for trace in self.rst_traces:
            rst["traceEvents"].append(trace)
        if verbose:
            TraceManager(self.rst_traces, DirLevel.TRIAL).get_iter_time()
        filename = os.path.join(self.dump_path,
                    "synthetic.json") if _path is None else _path
        with open(filename, 'w') as f:
            json.dump(rst, f)
        if self.full_trace:
            with open(os.path.join(self.dump_path, 'queue_status.json'), 'w') as f:
                json.dump(self.queue_status, f, indent=4)

    def daydream_dag(self, metadata, single = False):
        ''' Convert the DFG to that described in Daydream
            1. Only keep operators on one GPU, i.e. host0.rank0
            2. Use Coarse-grained Comm operators
        '''
        SingleLogger().info("Convert dag to dagdream_dag ...")
        if metadata.metainfo is None:
            raise ValueError("metadata info is None")
        _dag = nx.DiGraph()

        DAYDREAM_BW_RATIO = 0.85

        edges_to_add = []
        def wrap_add_edge(u, v):
            _dag.add_edge(u, v)
            # edges_to_add.append((u, v))

        one_pid = gen_pid_name(self.comm_backend, None, None)
        for u, v in self.dag.edges():
            u_pid, u_raw_name, u_cat, _ = parse_allinfo_from_name(u)
            v_pid, v_raw_name, v_cat, _ = parse_allinfo_from_name(v)
            ### Only keep operators on one GPU, i.e. host0.rank0
            if u_pid != one_pid or v_pid != one_pid:
                continue
            if u_cat == CatName.COMM.value:
                continue
            if v_cat == CatName.COMM.value:
                assert "BW" in u_raw_name

                ### find corresponding UPDATE operator
                update_ops = []
                for succ_ in self.dag.successors(v):
                    ### only sync may corresponding to multiple tensor fusion groups
                    _pid, _raw_name, _cat, _ = parse_allinfo_from_name(succ_)
                    if _pid != one_pid:
                        continue
                    op_type, op_name, _ = _raw_name.split(".")
                    _last_comm_task = gen_long_name(
                        one_pid, "{}.{}.{}".format(op_type, op_name, QueueType().ret_list()[-1]))
                    _update_ops = list(self.dag.successors(_last_comm_task))
                    # assert len(_update_ops) == 1, (_last_comm_task)
                    _update_op = _update_ops[0]
                    assert "UPDATE_" in _update_op, (u, v, _last_comm_task, _update_op)
                    update_ops.append(_update_op)

                tensor_list = re.findall("[0-9]+", op_name)
                fused_size = 0
                for tensor_id_str in tensor_list:
                    comm_op = gen_long_name(
                        one_pid, "{}.{}".format(op_type, tensor_id_str))
                    ### BW --> Comm
                    wrap_add_edge(u, comm_op)
                    ### Comm --> UPDATE
                    for update_op in update_ops:
                        wrap_add_edge(comm_op, update_op)
                    tensor_id = int(tensor_id_str)
                    bw_in_g = 100 * DAYDREAM_BW_RATIO
                    tensor_size = metadata.tensor_id2size(tensor_id)
                    _dag.nodes[comm_op]["avg"] = tensor_size / (bw_in_g * 1e6)
                    fused_size += tensor_size
                SingleLogger().debug("Fused Tensor Size for {}: {} MB".format(op_name, fused_size / (1024.0 * 1024.0)))
            else:
                # if not single and "BW" in u and "UPDATE_" in v:
                #     raise ValueError(u, v)
                wrap_add_edge(u, v)
        # _dag.add_edges_from(edges_to_add)
        for node_ in _dag.nodes():
            if "Comm" in node_:
                continue
            nx.set_node_attributes(_dag, {node_: self.dag.nodes[node_]})
        self.dag = _dag

class ReplayDebuger:
    def __init__(self, replayer):
        self.debug_nodes = {}
        self.dag = replayer.dag
        self.enabled = False

    def monitor_node(self, n):
        self.enabled = True
        self.debug_nodes[n] = {"todo": [u for u, _ in self.dag.in_edges(n)], "done": []}

    def mark_as_exct(self, u, v):
        ''' node `u` is executed, and handle the status of `v` if it is monitored
        and (u, v) must be an edge in self.dag
        '''
        if not self.enabled or v not in self.debug_nodes:
            return
        self.debug_nodes[v]["todo"].remove(u)
        self.debug_nodes[v]["done"].append(u)
    
    def _show_staue(self, node):
        todos = self.debug_nodes[node]["todo"]
        if len(todos) > 0:
            print("{} is waiting for {} nodes: {}".format(node, len(todos), str(todos)))

    def show_staue(self, nodes=None):
        if nodes is None:
            for node in self.debug_nodes.keys():
                self._show_staue(node)
        elif isinstance(nodes, list):
            for node in nodes:
                self._show_staue(node)
        else:
            self._show_staue(nodes)
