#Get all data and divide per line
data = open("Day 14.txt",'r').read().split('\n')

#Set initial values
rocks = []
sands = []
size = [[500,500],[0,0]]
start = [500,0]
grains = 0
first = True

#Define printing image
def printImage(image, extra):
    print('')
    for row in range(size[1][0],size[1][1]+3):
        for col in range(size[0][0],size[0][1]+1):
            #print(col,row)
            if [col,row] == start:
                print('@',end='')
            elif [col,row] in extra:
                print('O',end='')
            elif [col,row] in image:
                print('#',end='')
            else:
                print('.',end='')
        print('')
    print('')

#Get rock formations
for row in data:
    #Get points
    points = row.split(' -> ')

    #Go through all points
    for i in range(len(points)-1):
        #Set line
        line = [[*eval(points[i])],[*eval(points[i+1])]]

        #Check if horizontal or vertical
        if line[0][0] == line[1][0]:
            #Get direction
            sign = int((line[1][1]-line[0][1])/abs(line[0][1]-line[1][1]))

            #Go through the line
            for point in range(line[0][1],line[1][1],sign):
                #Add to rocks
                rocks.append([line[0][0],point])

        if line[0][1] == line[1][1]:
            #Get direction
            sign = int((line[1][0]-line[0][0])/abs(line[0][0]-line[1][0]))

            #Go through the line
            for point in range(line[0][0],line[1][0],sign):
                #Add to rocks
                rocks.append([point,line[0][1]])
    else:
        rocks.append([*eval(points[i+1])])

#Get size of field and add margin
for rock in rocks:
    if rock[0] < size[0][0]:
        size[0][0] = rock[0]
    if rock[0] > size[0][1]:
        size[0][1] = rock[0]
    if rock[1] < size[1][0]:
        size[1][0] = rock[1]
    if rock[1] > size[1][1]:
        size[1][1] = rock[1]


#Part 2
#Add floor to rocks
height = size[1][1]+2
for i in range(start[0]-height,start[0]+height+1):
    rocks.append([i,height])

#Keep going until end point reached
while True:
    #Set starting value
    sand = start.copy()

    #Check if grain is already in rocks
    if sand in rocks:
        printImage(rocks, sands)
        print("Total grains to fill start:",grains)
        break

    #Give an update every 1000 grains
    if grains%1000 == 0:
        printImage(rocks, sands)
    if grains%100 == 0:
        print(grains,"grains")

    while True:
        #Check if grain has reached out of bounds for the first time
        if sand[1] > size[1][1] and first:
            #Part 1
            printImage(rocks, sands)
            print("Total amount of grains before out of bounds:",grains)
            first = False

        #Check directly below
        if [sand[0],sand[1]+1] not in rocks:
            sand[1] += 1
            continue

        #Check bottom-left
        if [sand[0]-1,sand[1]+1] not in rocks:
            sand[0] -= 1
            sand[1] += 1
            continue

        #Check bottom-right
        if [sand[0]+1,sand[1]+1] not in rocks:
            sand[0] += 1
            sand[1] += 1
            continue

        #If grain has nowhere to move, add it to rocks and start new grain
        rocks.append(sand)
        sands.append(sand)

        #Increment grain
        grains += 1
        break