from typing import OrderedDict


def get_wall_duration(op_names, all_ops, pid_list=(11, 7, 13, 15, 9)):
    """
    Calculates wall duration for each op in op_names.
    Params:
    op_names: list (str), names of ops of interest.
    pid_list: list (str), names of pid to include.
    all_ops: output of get_all_ops().
    Return:
    total wall duration, dict['op'] = wall duration.
    """
    # 1. Construct dictionary of op with name matching op_names
    ops_dic = OrderedDict()
    for name in op_names:
        ops = []
        for op in all_ops:
            if op['name'] == name:
                ops.append(op)
        ops_dic[name] = ops

    # 2. get duration for each op
    op_dict = OrderedDict()
    total_dur = 0
    for op_name in op_names:
        op_dur = 0
        for itm in ops_dic[op_name]:
            if itm['pid'] in pid_list:
                op_dur += itm['dur']
        op_dict[op_name] = op_dur * 1e-3  # convert from us to ms
        total_dur += op_dur * 1e-3

    # fixing the NCCL key:
    op_dict['unknown (nccl AllReduceKernel_sum_)'] = op_dict.pop('unknown')

    # Sorting durations:
    sorted_dur = sorted(op_dict.items(), key=lambda x: x[1])[::-1]
    # sorted_dur = sorted(op_dict.items(), key=operator.itemgetter(1))

    return OrderedDict(sorted_dur), total_dur