def main():
    import numpy as np
    import re
    with open("/home/yassinealouini/Documents/code/advent_of_code/aoc/year_2021/data/5.txt") as f:
        tmp = f.read().rstrip().split("\n")

    # Using complex numbers as (x, y) representation
    starts, ends = [], []
    for row in tmp:
        x1, y1, x2, y2 = re.findall(r'\d+', row)
        starts.append(int(x1) + 1j * int(y1))
        ends.append(int(x2) + 1j * int(y2))
    # Dimensions of the grid
    d = max(map(abs, starts)) + 1

    def solve(part_2=False):
        a = np.zeros((int(d), int(d)))
        for start_point, end_point in zip(starts, ends):

            if part_2:
                # Compute cosine and sine to find if diagonal or anti-diagonal
                diff = start_point - end_point
                c = (diff.real) / abs(diff)
                s = (diff.imag) / abs(diff)
                start_x = min(int(start_point.real), int(end_point.real))
                end_x = max(int(start_point.real), int(end_point.real))
                start_y = min(int(start_point.imag), int(end_point.imag))
                end_y = max(int(start_point.imag), int(end_point.imag))
                sliced_a = a[start_x: end_x + 1, start_y: end_y + 1]
                if round(s * c, 1) == 0.5:
                    np.fill_diagonal(sliced_a, 
                                     sliced_a.diagonal() + 1)
                elif round(s * c, 1) == -0.5:
                    # Flip the sliced matrix to fill the diagonal
                    np.fill_diagonal(np.fliplr(sliced_a), 
                                     np.fliplr(sliced_a).diagonal() + 1)
            if start_point.real == end_point.real:
                start = min(int(start_point.imag), int(end_point.imag))
                end = max(int(start_point.imag), int(end_point.imag))
                a[int(start_point.real), start: end + 1] += 1
            elif start_point.imag == end_point.imag:
                start = min(int(start_point.real), int(end_point.real))
                end = max(int(start_point.real), int(end_point.real))
                a[start: end + 1, int(start_point.imag)] += 1
        return a
    
    print("Solution to part I: ", (solve()>= 2).sum())
    print("Solution to part II: ", (solve(part_2=True)>= 2).sum())


def streamlit_5(data_input):
    """ Day 5 solution (mainly using numpy)
    """
    import numpy as np
    import re
    import streamlit as st

    tmp = data_input.rstrip().split("\n")

    # Using complex numbers as (x, y) representation
    starts, ends = [], []
    for row in tmp:
        x1, y1, x2, y2 = re.findall(r'\d+', row)
        starts.append(int(x1) + 1j * int(y1))
        ends.append(int(x2) + 1j * int(y2))

    # Dimension of the grid
    d = max(map(abs, starts)) + 1

    def solve(part_2=False):
        a = np.zeros((int(d), int(d)))
        for start_point, end_point in zip(starts, ends):
            start_x = min(int(start_point.real), int(end_point.real))
            end_x = max(int(start_point.real), int(end_point.real))
            start_y = min(int(start_point.imag), int(end_point.imag))
            end_y = max(int(start_point.imag), int(end_point.imag))
            # Compute cosine and sine to find if diagonal or anti-diagonal
            diff = start_point - end_point
            c = (diff.real) / abs(diff)
            s = (diff.imag) / abs(diff)
            sliced_a = a[start_x: end_x + 1, start_y: end_y + 1]
            criterion = round(s * c, 1)
            if part_2:
                if criterion == 0.5:
                    np.fill_diagonal(sliced_a, sliced_a.diagonal() + 1)
                elif criterion == -0.5:
                    # Need to flip the sliced matrix to get the correct diagonal
                    np.fill_diagonal(np.fliplr(sliced_a), 
                                     np.fliplr(sliced_a).diagonal() + 1)
            # Either horizontal or vertical
            if criterion == 0:
                sliced_a += 1
        return a
    
    st.write("Solution to part I: ", (solve()>= 2).sum())
    st.write("Solution to part II: ", (solve(part_2=True)>= 2).sum())

if __name__ == "__main__":
    main()