# Natural Language Toolkit: Distance Metrics
#
# Copyright (C) 2001-2023 NLTK Project
# Author: Edward Loper <edloper@gmail.com>
#         Steven Bird <stevenbird1@gmail.com>
#         Tom Lippincott <tom@cs.columbia.edu>
# URL: <https://www.nltk.org/>
# For license information, see LICENSE.TXT
#
# Changed by OST


def jaro_similarity(s1, s2):

    # First, store the length of the strings
    # because they will be re-used several times.
    len_s1, len_s2 = len(s1), len(s2)

    # The upper bound of the distance for being a matched character.
    match_bound = max(len_s1, len_s2) // 2 - 1

    # Initialize the counts for matches and transpositions.
    matches = 0  # no.of matched characters in s1 and s2
    transpositions = 0  # no. of transpositions between s1 and s2
    flagged_1 = []  # positions in s1 which are matches to some character in s2
    flagged_2 = []  # positions in s2 which are matches to some character in s1

    # Iterate through sequences, check for matches and compute transpositions.
    for i in range(len_s1):  # Iterate through each character.

        upperbound = min(i + match_bound + 1, len_s2)
        lowerbound = max(0, i - match_bound)

        for j in range(lowerbound, upperbound):
            if s1[i] == s2[j] and j not in flagged_2:
                matches += 1
                flagged_1.append(i)
                flagged_2.append(j)
                break

    flagged_2.sort()

    for i, j in zip(flagged_1, flagged_2):
        if s1[i] != s2[j]:
            transpositions += 1

    if matches == 0:
        return 0
    else:
        return (1 / 3) * (matches / len_s1 + matches / len_s2 + (matches - transpositions // 2) / matches)


def jaro_winkler_similarity(s1, s2, p=0.1):
    jaro_sim = jaro_similarity(s1, s2)

    # Calculate the prefix length
    prefix = 0
    max_prefix_length = min(4, min(len(s1), len(s2)))
    for i in range(max_prefix_length):
        if s1[i] == s2[i]:
            prefix += 1
        else:
            break

    # Apply the Jaro-Winkler adjustment
    jaro_winkler_sim = jaro_sim + (prefix * p * (1 - jaro_sim))

    return jaro_winkler_sim
