def fitness_eval(model, input, image, target_top):
    # 适应度评估
    clip_input = clip_perturbation(image, input, 50.)
    outputs = model(clip_input)
    
    ranking = torch.argsort(outputs, descending=True).cuda()
    top_class = set(ranking[0][0:3].tolist())
    ft = torch.sum(ranking[0][0:3]==target_top)
    target_top = set(target_top.tolist())
    fitness = len(top_class&target_top) 
    return fitness

def is_adversarial_rank1(model, input, target_top):
    # 判断首3位rank1攻击是否成功

    outputs = model(input)
    
    ranking = torch.argsort(outputs, descending=True).cuda()
    top_class = set(ranking[0][0:3].tolist())
    target_top = set(target_top.tolist())
    is_adv = top_class==target_top
    return is_adv

def init_distnfitness(original, x_pert, model, target_classes, fitness_min):
    # initialize ditance & fitness
    distance_pert = torch.norm(original-x_pert)
    fitness=fitness_eval(model, x_pert, original, target_classes)
    return distance_pert, fitness


def update(original, x_pert, model, target_classes, dist, fitness, V, k_worst, v_new):
    distance_pert = torch.norm(original-x_pert)
    fitness_pert=fitness_eval(model, x_pert, original, target_classes)
    if fitness_pert>fitness[k_worst] or (fitness_pert==3 and distance_pert<dist[k_worst]):
        fitness[k_worst] = fitness_pert
        dist[k_worst] = distance_pert
        V[k_worst] = v_new

    return V, dist, fitness

def get_bestnworst(fitness, dist):
    fitness_t = fitness.detach().clone()
    dit_t = dist.detach().clone()
    fitness_low = fitness<3
    k_worst = torch.nonzero(fitness_low)
    # print(k_worst)
    if k_worst.numel() > 0:
        k_worst = torch.argmin(torch.where(fitness_low, fitness, 0.))
    else:
        k_worst = torch.argmax(dist)
    fitness_high = fitness>=3
    k_best = torch.nonzero(fitness_high)
    if k_best.numel() >= 2:
        k_best = torch.argmin(torch.where(fitness_high, dist, torch.inf))
    else:
        k_best = torch.argmax(fitness)
    return k_best, k_worst