from utils import load_data
import numpy as np

def question_1(data: list[str]):

    sum = 0
    # Convert strings to lists of characters
    new_data = reformat_data(data)
    # Get all the special symbols in the input
    symbols = get_symbols(new_data)
    # Get all number positions
    number_positions = get_number_positions(new_data)
    # Convert to numpy matrix
    data = np.array(new_data)

    # We assume that no value is next to multiple symbols
    # Go through each symbol and sum their adjacent values
    found_numbers = []
    for i in range(len(data)):
        for j, value in enumerate(data[i]):
            if value in symbols:
                
                j_interval = (j-1, j+1)
                i_interval = (i-1, i+1)
                for number in number_positions:
                    if (number[3] >= i_interval[0] and number[3] <= i_interval[1]) and (number[2] >= j_interval[0] and number[1] <= j_interval[1]):
                        found_numbers.append(number)
                        print(f"Around character {value} on line {i} I found the following number {number[0]}")

    for number in set(found_numbers):
        sum += int(number[0])
    

    print(f"Answer to question 1 is: {sum}")

def question_2(data: list[str]):
    
    sum = 0

    # Convert strings to lists of characters
    new_data = reformat_data(data)
    # Get all number positions
    number_positions = get_number_positions(new_data)
    # Convert to numpy matrix
    data = np.array(new_data)

    # We assume that no value is next to multiple symbols
    # Go through each symbol and sum their adjacent values
    for i in range(len(data)):
        for j, value in enumerate(data[i]):
            if value == "*":
                found_numbers = []
                
                j_interval = (j-1, j+1)
                i_interval = (i-1, i+1)

                for number in number_positions:
                    if (number[3] >= i_interval[0] and number[3] <= i_interval[1]) and (number[2] >= j_interval[0] and number[1] <= j_interval[1]):
                        found_numbers.append(number)
                if len(found_numbers) == 2:
                    sum += int(found_numbers[0][0]) * int(found_numbers[1][0])
    
    print(f"Answer to question 2 is: {sum}")


def get_number_positions(data: list[list[str]]):
    number = ""
    coordinate_1 = 0
    coordinate_2 = 0
    number_positions = []
    for i in range(len(data)):
        for j, value in enumerate(data[i]):
            if value.isdigit():
                if number == "":
                    coordinate_1 = j
                number += value
            elif number != "":
                coordinate_2 = j-1
                number_positions.append((number, coordinate_1, coordinate_2, i))
                number = ""
        if data[i][-1].isdigit() and number != "":
            coordinate_2 = 139
            number_positions.append((number, coordinate_1, coordinate_2, i))
            number = ""

    return number_positions

def get_symbols(data):
    total_lists = []
    for string in data:
        total_lists += string
    present_set = set(total_lists)
    remove_set = {".",  "1", "2", "3", "4", "5", "6", "7", "8", "9", "0"}
    symbol_set = present_set.difference(remove_set)
    return symbol_set

def reformat_data(data):
    new_lines = []
    for i, line in enumerate(data):
        new_lines.append([*line])
    return new_lines

if __name__ == "__main__":
    filepath = "Data/day3.txt"
    data = load_data(filepath)
    # question_1(data)
    question_2(data)