import random

import pytest
import torch

from ding.torch_utils.metric import levenshtein_distance, hamming_distance


@pytest.mark.unittest
class TestMetric():

    def test_levenshtein_distance(self):
        r'''
        Overview:
            Test the Levenshtein Distance
        '''
        pred = torch.LongTensor([1, 4, 6, 4, 1])
        target1 = torch.LongTensor([1, 6, 4, 4, 1])
        distance = levenshtein_distance(pred, target1)
        assert (distance.item() == 2)

        target2 = torch.LongTensor([])
        distance = levenshtein_distance(pred, target2)
        assert (distance.item() == 5)

        target3 = torch.LongTensor([6, 4, 1])
        distance = levenshtein_distance(pred, target3)
        assert (distance.item() == 2)
        target3 = torch.LongTensor([6, 4, 1])
        distance = levenshtein_distance(pred, target3, pred, target3, extra_fn=lambda x, y: x + y)
        assert distance.item() == 13
        target4 = torch.LongTensor([1, 4, 1])
        distance = levenshtein_distance(pred, target4, pred, target4, extra_fn=lambda x, y: x + y)
        assert distance.item() == 14

    def test_hamming_distance(self):
        r'''
        Overview:
            Test the Hamming Distance
        '''
        base = torch.zeros(8).long()
        index = [i for i in range(8)]
        for i in range(2):
            pred_idx = random.sample(index, 4)
            target_idx = random.sample(index, 4)
            pred = base.clone()
            pred[pred_idx] = 1
            target = base.clone()
            target[target_idx] = 1
            pred = pred.unsqueeze(0)
            target = target.unsqueeze(0)
            distance = hamming_distance(pred, target)
            diff = len(set(pred_idx).union(set(target_idx)) - set(pred_idx).intersection(set(target_idx)))
            assert (distance.item() == diff)