import numpy as np
from numbers import Number
from collections import defaultdict

NO_ANSWER_TEMPLATE = [
    "Code cannot be executed!",
    "Code contains infinite loop!",
    "Not exist!",
]

def majority_voting(answers):
    count = group_answers(answers)
    valid_answer_count = [len(value) for key, value in count.items() if key not in NO_ANSWER_TEMPLATE]
    if len(valid_answer_count) > 0:
        majority_count = max(valid_answer_count)
        majority_results = [key for key in count if len(count[key]) == majority_count and key not in NO_ANSWER_TEMPLATE]
        return majority_results, majority_count
    else:
        return ["Not exist!"], len(answers)


def get_top_k_voting(answers, k=2):
    count = group_answers(answers)

    count_key = [key for key in count if key != "Not exist!"]
    count_value = [len(count[key]) for key in count_key]
    sorted_idx = np.argsort(count_value)
    begin = len(sorted_idx) - k
    while begin > 0:
        if count_value[sorted_idx[begin]] != count_value[sorted_idx[begin - 1]]:
            break
        begin -= 1
    return {count_key[i]: count[count_key[i]] for i in sorted_idx[begin:]}


def group_answers(responses):
    grouped = defaultdict(list)
    has_numeric = False

    for resp in responses:
        try:
            eval(resp)
            has_numeric = True
            break
        except:
            pass

    if has_numeric:
        for idx, resp in enumerate(responses):
            if resp in grouped:
                grouped[resp].append(idx)
                continue

            matched = False
            for existing_key in grouped:
                try:
                    if abs(eval(resp) - eval(existing_key)) < 1e-6:
                        grouped[existing_key].append(idx)
                        matched = True
                        break
                except:
                    pass

            if not matched:
                grouped[resp].append(idx)
    else:
        for idx, resp in enumerate(responses):
            grouped[resp].append(idx)

    return grouped
