#!/usr/bin/env python3
import aoclib
import numpy

DAY = 6
TEST_SOLUTION_PART1 = 11
TEST_SOLUTION_PART2 = 6


def getAnswerCount(answer_line):
    return_dict = {}
    answer_line = ''.join(answer_line)
    for answer in answer_line:
        return_dict[answer] = return_dict.get(answer, 0) + 1

    return return_dict


@aoclib.print_execution_time
def part1(test_mode=False):
    my_input = aoclib.getMultiLineInputAsArray(day=DAY, test=test_mode)
    yes_count = 0
    for answer_group in my_input:
        answer_count = getAnswerCount(answer_group)
        yes_count += len(answer_count)

    return yes_count


@aoclib.print_execution_time
def part2(test_mode=False):
    my_input = aoclib.getMultiLineInputAsArray(day=DAY, test=test_mode)

    yes_count = 0

    for answer_group in my_input:
        answer_buffer = list(answer_group[0])
        for answer_line in answer_group:
            answer_line = list(answer_line)
            answer_buffer = list(numpy.intersect1d(answer_buffer, answer_line))

        yes_count = yes_count + len(answer_buffer)

    return yes_count