#!/bin/python3
import sys,time,re
from pprint import pprint
sys.path.insert(0, '../../')
from fred import list2int,get_re,lprint, grid_valid,get_value_in_direction,addTuples,subTuples
from termcolor import colored

start_time = time.time()

input_f = 'test3'

def nprint(grid, cur: set = None, sign: str = None, positions:list = None):
    """
    Prints a grid, highlighting the current position if specified.

    Args:
        grid (list): A 2D grid to print.
        cur (set, optional): A set containing the (row, col) indices of the current position.
                             Defaults to None.
        sign (str, optional): The sign to highlight the current position with. Defaults to None.
        positions (list, optional): A list of sets containing the (row, col) indices of positions to color.
                             Defaults to None.
    Returns:
        None

    Raises:
        TypeError: If cur is not a set or sign is not a string.
    """
    if cur is not None and not isinstance(cur, tuple):
        raise TypeError("Cur must be a tuple with (row, column) indices.")
    if sign is not None and not isinstance(sign, str):
        raise TypeError("Sign must be a string.")
    if positions is not None and not isinstance(positions, list):
        raise TypeError("Positions must be a list.")
    
    prepend = 0
    for idx, i in enumerate(grid):
        for jdx, j in enumerate(i):
            if jdx == 0:
                prepend = str(idx)+' '
            else: 
                prepend = ''
            
            if (idx, jdx) == cur:
                if sign is not None:
                    if len(sign) > 1:
                        print(prepend+sign[0] + grid[idx][jdx] + sign[1], end='')  # Print with sign
                    else:
                        print(prepend+colored(sign,'green',attrs=["underline","bold"]), end=' ')  # Print sign
                else:
                    print(prepend+colored(grid[idx][jdx],'green',attrs=["underline","bold"]), end=' ')
            else:
                if positions is not None:
                    if (idx,jdx) in positions:
                        print(prepend+colored(grid[idx][jdx],'red'),end=' ')
                    else:
                        
                        print(prepend+grid[idx][jdx], end=' ')
                else:
                    if grid[idx][jdx] == '.':
                        print(prepend+colored(grid[idx][jdx],'blue',attrs=["concealed"]),end=' ')
                    elif grid[idx][jdx] == 'O':
                        print(prepend+colored(grid[idx][jdx],'red'),end=' ')
                    elif grid[idx][jdx] == '#':
                        print(prepend+colored(grid[idx][jdx],'white'),end=' ')
                    else:
                        print(prepend+grid[idx][jdx], end=' ')  # Regular grid element
        print()
   
    
    for r in range(len(grid[0])):
        if r == 0:
            print('  ',end='')
        print(r,end=' ')
    print()


def loadFile(input_f):
    grid = []
    instructions = []
    with open(input_f) as file:
        for line in file:
            # load map part
            if line.startswith("#"):
                grid.append(list(line.rstrip()))
            
            # load instructions
            elif line.startswith(('v','<','>','^')):
                instructions += list(line.rstrip())
    return grid,instructions

#########################################
#                                       #
#              Part 1                   #
#                                       #
#########################################
def part1():
    grid, instructions = loadFile(input_f)
                  
    start = ()
    
    for r,row in enumerate(grid):
        for c, col in enumerate(row):
            if grid[r][c] == '@':
                start = (r,c)
    
    # translate arrows to something get_value_in_direction can use
    directions = {
        '^':('up',(-1, 0)),
        'v':('down',(1, 0)),
        '>':('right',(0, 1)),
        '<':('left',(0, -1)),
    }
    
    pos = start
    #print('Initial state')
    #nprint(grid,pos)
    #print()
    #input()
    for idx, inst in enumerate(instructions):
        #print('Move',inst,'(',len(instructions)-idx,')')
        dir = directions[inst][0]
        next = get_value_in_direction(grid,pos,dir)
        
        # If wall, don't do anything
        if next == '#':
            #nprint(grid,pos)
            #input()
            continue
        
        # If free space, move there
        if next == '.':
            grid[pos[0]][pos[1]] = '.'
            pos = addTuples(pos,directions[inst][1])
            grid[pos[0]][pos[1]] = '@'
        
        # If box, move the box and the stack of boxes.
        if next == 'O':
            #print('@',pos)
            prev = pos
            next_chars = ['@']
            next_poss = [pos]
            skip = False
            while True:
                nextPos = addTuples(pos,directions[inst][1])
                nextChar = get_value_in_direction(grid,nextPos)
                #print(nextPos,nextChar)
                if nextChar == 'O':
                    next_chars.append(nextChar)
                    next_poss.append(nextPos)
                    pos = nextPos
                if nextChar == '.':
                    next_chars.append(nextChar)
                    next_poss.append(nextPos)
                    break
                if nextChar == '#':
                    skip = True
                    pos = prev
                    break
            
                #input()
            if not skip:
                for ndx,n in enumerate(next_poss):
                    if ndx == 0:
                        grid[n[0]][n[1]] = '.'
                        pos = next_poss[ndx+1]
                    else:
                        grid[n[0]][n[1]] = next_chars[ndx-1]
            #input()
        #nprint(grid,pos)
        #print(len(instructions)-idx)
        #time.sleep(0.05)
    #input()
    result = 0
    for r,row in enumerate(grid):
        for c,char in enumerate(row):
            if char == 'O':
                result += (100*r+c)
    return result
start_time = time.time()
print('Part 1:',part1(), '\t\t', round((time.time() - start_time)*1000), 'ms')


#########################################
#                                       #
#              Part 2                   #
#                                       #
#########################################

def resizeGrid(grid:list) -> list:
    newGrid = []
    for r,row in enumerate(grid):
        newRow = []
        for c, char in enumerate(row):
            if char == '#':
                newRow.append('#')
                newRow.append('#')
            if char == 'O':
                newRow.append('[')
                newRow.append(']')
            if char == '.':
                newRow.append('.')
                newRow.append('.')
            if char == '@':
                newRow.append('@')
                newRow.append('.')
        newGrid.append(newRow)
    return newGrid

def part2():
    
    
    def nprint2(grid, cur: set = None, sign: str = None, positions:list = None):
    
   
        prepend = 0
        for idx, i in enumerate(grid):
            for jdx, j in enumerate(i):
                if jdx == 0:
                    prepend = str(idx)+' '
                else: 
                    prepend = ''
                
                if (idx, jdx) == cur:
                    if sign is not None:
                        if len(sign) > 1:
                            print(prepend+sign[0] + grid[idx][jdx] + sign[1], end=' ')  # Print with sign
                        else:
                            print(prepend+colored(sign,'green',attrs=["underline","bold"]), end=' ')  # Print sign
                    else:
                        print(prepend+colored(grid[idx][jdx],'green',attrs=["underline","bold"]), end='')
                else:
                    if positions is not None:
                        if (idx,jdx) in positions:
                            print(prepend+colored(grid[idx][jdx],'red'),end=' ')
                        else:
                            
                            print(prepend+grid[idx][jdx], end=' ')
                    else:
                        if grid[idx][jdx] == '.':
                            print(prepend+colored(grid[idx][jdx],'blue',attrs=["concealed"]),end='')
                        elif grid[idx][jdx] == 'O':
                            print(prepend+colored(grid[idx][jdx],'red'),end='')
                        elif grid[idx][jdx] == '#':
                            print(prepend+colored(grid[idx][jdx],'white'),end='')
                        else:
                            print(prepend+grid[idx][jdx], end='')  # Regular grid element
            print()
    
        
        for r in range(len(grid[0])):
            if r == 0:
                print('  0',end='')
            
            else:
                if r%2 == 0:
                    print(r,end='')
                else:
                    print(' ',end='')
        print()
        for r in range(len(grid[0])):
            if r == 0:
                print('   ',end='')
            
            else:
                if r%2 != 0:
                    print(r,end='')
                else:
                    if r != 9:
                        print(' ',end='')
                    else:
                        print('',end='')
        print()
    
    
    grid, instructions = loadFile(input_f)
                  
    start = ()
    
    for r,row in enumerate(grid):
        for c, col in enumerate(row):
            if grid[r][c] == '@':
                start = (r,c)
    
    # translate arrows to something get_value_in_direction can use
    directions = {
        '^':('up',(-1, 0)),
        'v':('down',(1, 0)),
        '>':('right',(0, 1)),
        '<':('left',(0, -1)),
    }
    
    pos = start
    
    nprint(grid)
    #print(grid)
    grid = resizeGrid(grid)
    print()
    nprint2(grid)
    #print(grid)
    
    print()
    
    for idx, inst in enumerate(instructions):
        #print('Move',inst,'(',len(instructions)-idx,')')
        dir = directions[inst][0]
        next = get_value_in_direction(grid,pos,dir)
        
        # If wall, don't do anything
        if next == '#':
            #nprint(grid,pos)
            #input()
            continue
        
        # If free space, move there
        if next == '.':
            grid[pos[0]][pos[1]] = '.'
            pos = addTuples(pos,directions[inst][1])
            grid[pos[0]][pos[1]] = '@'
        
        # If box, move the box and the stack of boxes.
        if next == 'O':
            #print('@',pos)
            prev = pos
            next_chars = ['@']
            next_poss = [pos]
            skip = False
            while True:
                nextPos = addTuples(pos,directions[inst][1])
                nextChar = get_value_in_direction(grid,nextPos)
                #print(nextPos,nextChar)
                if nextChar == 'O':
                    next_chars.append(nextChar)
                    next_poss.append(nextPos)
                    pos = nextPos
                if nextChar == '.':
                    next_chars.append(nextChar)
                    next_poss.append(nextPos)
                    break
                if nextChar == '#':
                    skip = True
                    pos = prev
                    break
            
                #input()
            if not skip:
                for ndx,n in enumerate(next_poss):
                    if ndx == 0:
                        grid[n[0]][n[1]] = '.'
                        pos = next_poss[ndx+1]
                    else:
                        grid[n[0]][n[1]] = next_chars[ndx-1]
    
start_time = time.time()
print('Part 2:',part2(), '\t\t', round((time.time() - start_time)*1000), 'ms')