import numpy as np
from collections import Counter

def main():

    def most_common_2d(data):
        return [Counter(col).most_common()[0][0] for col in zip(*data)]

    def least_common_2d(data):
        return [Counter(col).most_common()[-1][0] for col in zip(*data)]

    def Counter_2d(data):
        return [Counter(col) for col in zip(*data)]

    def search(data, to_search="1"):
        if to_search == "1":
            most_common_index = 0
        else:
            most_common_index = -1
        previous_keep = data
        for i in range(len(data[0])):
            counts = Counter_2d(previous_keep)
            keep = []
            for row in previous_keep:
                if (row[i] == counts[i].most_common()[most_common_index][0] and 
                    counts[i]["0"] != counts[i]["1"]):
                    keep.append(row)
                elif counts[i]["0"] == counts[i]["1"]:
                    if row[i] == to_search:
                        keep.append(row)
            previous_keep = keep
        return int("".join(previous_keep[0]), 2)


    with open("/home/yassinealouini/Documents/code/advent_of_code/aoc/year_2021/data/3.txt") as f:
        tmp = f.read().rstrip().split("\n")
        data = []
        for row in tmp:
            data.append([e for e in row])


    gamma = int("".join(most_common_2d(data)), 2)
    epsilon = int("".join(least_common_2d(data)), 2)

    print("Solution to part I: ", gamma * epsilon)


    oxygen = search(data, "1")
    co2 =  search(data, "0")


    print("Solution to part II: ", oxygen * co2)



def streamlit_torch_3(data_input):
    """ Day 3 of AoC using PyTorch
    """
    import torch
    import streamlit as st
    st.write("Solution to part I: ", 0)

if __name__ == "__main__":
    main()