with open('./2021/03.input', 'r') as input_file:
# with open('./2021/03.test', 'r') as input_file:
    input_data = input_file.read()

working_data = input_data.split('\n')

working_dict = {}

for one_row in working_data:
    for i in range(len(working_data[0])):
        if i not in working_dict.keys():
            working_dict[i] = {}
            one_dict = {}
            one_dict[0] = 0
            one_dict[1] = 1
            working_dict[i] = one_dict
        
        working_dict[i][int(one_row[i])] += 1

# Make gamma / epsilon
gamma = ''
epsilon = ''
for one_column, one_item in working_dict.items():
    if one_item[0] > one_item[1]:
        gamma = gamma + '0'
        epsilon = epsilon + '1'
    else:
        gamma = gamma + '1'
        epsilon = epsilon + '0'
print(int(gamma, 2) * int(epsilon, 2))

# Part 2
co2_data = working_data.copy()

index = 0
count_0 = 0
count_1 = 1
output = ""
while len(co2_data) > 1 and index < len(co2_data[0]):
    # Get count at index
    for i in range(len(co2_data)):
        if co2_data[i][index] == "0":
            count_0 += 1
        if co2_data[i][index] == "1":
            count_1 += 1
    
    # Get highest bit
    if count_0 < count_1:
        output = output + "0"
    if count_0 > count_1:
        output = output + "1"
    if count_0 == count_1:
        output = output + "0"
    
    # Filter out anything without this bit
    co2_data = [x for x in co2_data if x[index] == output[index]]

    count_0 = 0
    count_1 = 0
    index += 1

o2_data = working_data.copy()
output = ""
index = 0
while len(o2_data) > 1 and index < len(o2_data[0]):
    count_0 = 0
    count_1 = 0

    # Get count at index
    for i in range(len(o2_data)):
        if o2_data[i][index] == "0":
            count_0 += 1
        if o2_data[i][index] == "1":
            count_1 += 1

    # Get highest bit
    if count_0 < count_1:
        output = output + "1"
    if count_0 > count_1:
        output = output + "0"
    if count_0 == count_1:
        output = output + "1"

    # Filter out anything without this bit
    o2_data = [x for x in o2_data if x[index] == output[index]]
    
    index += 1
print(int(o2_data[0], 2) * int(co2_data[0], 2))