from stonesoup.dataassociator.general import OneToOneAssociator


def print_results(assoc_dict_print, unassocs_a_print, unassocs_b_print, measure):
    print(assoc_dict_print, "\n", unassocs_a_print, "\n", unassocs_b_print)
    total = sum(measure(obj_1, obj_2) for obj_1, obj_2 in assoc_dict_print.items())
    print(f"Total Measure={total}")


def multiply(item1: float, item2: float) -> float:
    return item1 * item2


numbers_a = (1, 2, 3, 5)
numbers_b = (1, 2, 4, 7)

associator = OneToOneAssociator(measure=multiply,
                                maximise_measure=True,
                                association_threshold=5)

all_assoc_dict = associator.association_dict(numbers_a, numbers_b)
assoc_dict = {a: all_assoc_dict[a]
              for a in numbers_a
              if all_assoc_dict[a] is not None}

unassocs_a = {x for x in numbers_a if all_assoc_dict[x] is None}
unassocs_b = {x for x in numbers_b if all_assoc_dict[x] is None}

print("1-2-1 Output")
print_results(assoc_dict, unassocs_a, unassocs_b, multiply)

# {3: 4, 5: 7} 
#  {1, 2} 
#  {1, 2}
# Total Measure=47

print("Better output")
better_output = {2: 4, 3: 2, 5: 7}
print_results(better_output, {1}, {1}, multiply)

# Total Measure=49