import paddle
import paddle.nn.functional as F
from collections import OrderedDict


def create_metric(
    out,
    label,
    architecture=None,
    topk=5,
    classes_num=1000,
    use_distillation=False,
    mode="train",
):
    """
    Create measures of model accuracy, such as top1 and top5

    Args:
        out(variable): model output variable
        feeds(dict): dict of model input variables(included label)
        topk(int): usually top5
        classes_num(int): num of classes
        use_distillation(bool): whether to use distillation training
        mode(str): mode, train/valid

    Returns:
        fetches(dict): dict of measures
    """
    # if architecture["name"] == "GoogLeNet":
    #     assert len(out) == 3, "GoogLeNet should have 3 outputs"
    #     out = out[0]
    # else:
    #     # just need student label to get metrics
    #     if use_distillation:
    #         out = out[1]
    softmax_out = F.softmax(out)

    fetches = OrderedDict()
    # set top1 to fetches
    top1 = paddle.metric.accuracy(softmax_out, label=label, k=1)
    # set topk to fetches
    k = min(topk, classes_num)
    topk = paddle.metric.accuracy(softmax_out, label=label, k=k)

    # multi cards' eval
    if mode != "train" and paddle.distributed.get_world_size() > 1:
        top1 = (
            paddle.distributed.all_reduce(top1, op=paddle.distributed.ReduceOp.SUM)
            / paddle.distributed.get_world_size()
        )
        topk = (
            paddle.distributed.all_reduce(topk, op=paddle.distributed.ReduceOp.SUM)
            / paddle.distributed.get_world_size()
        )

    fetches["top1"] = top1
    topk_name = "top{}".format(k)
    fetches[topk_name] = topk

    return fetches
