from copy import deepcopy
from aocd.models import Puzzle
from aocd import submit


def crate_mover_9000(stacks: dict, moves: list) -> str:
    """
    execute moves on the stacks and return the top of each stack
    the crate mover 9000 only moves single crates.
    """
    # take a deep copy so embedded list do not get messed up
    stacks = deepcopy(stacks)
    # for each move, execute on stacks
    for amount, index, dest in moves:
        for i in range(amount):
            stacks[dest].append(stacks[index].pop())
    # take the last element fom each stack and append to output
    return "".join([stack[-1] for stack in stacks.values()])


def crate_mover_9001(stacks: dict, moves: list) -> str:
    """
    execute moves on the stacks and return the top of each stack
    the crate mover 9001 moves stacks of crates.
    """
    # take a deep copy so embedded list do not get messed up
    stacks = deepcopy(stacks)
    # for each move, move amount of crates from index to dest
    for amount, index, dest in moves:
        stacks[dest].extend(stacks[index][-amount:])
        stacks[index] = stacks[index][:-amount]
    # take the last element fom each stack and append to output
    return "".join([stack[-1] for stack in stacks.values()])


def parse_input(data: list) -> list:
    """
    turn input data into usable data types:
        stacks_str -> dict of moves {'index':['A','B','C']}
        moves_str -> list of opeations [amount, index, dest]
    """
    # get input data
    stacks_str, moves_str = data.split("\n\n")
    # transpose field for easier parsing (thx linalg course)
    transposed_stacks = list(map(list, zip(*stacks_str.splitlines())))
    # parse stacks
    stacks = {}
    for line in transposed_stacks[1::4]:
        stack = [c for c in line if c != " "]
        key = stack.pop()
        stack.reverse()
        stacks[key] = stack
    # parse moves
    moves = []
    for move in moves_str.splitlines():
        _, amount, _, index, _, dest = move.split(" ")
        # print(f'move {amount} from {index} to {dest}')
        moves.append((int(amount), index, dest))
    # return parsed data
    return stacks, moves


if __name__ == "__main__":
    # get puzzle and parse data
    puzzle = Puzzle(year=2022, day=5)
    stacks, moves = parse_input(puzzle.input_data)

    # part a:
    answer_a = crate_mover_9000(stacks, moves)
    print(f"{answer_a}")
    submit(answer_a, part="a", day=5, year=2022)

    # part b:
    answer_b = crate_mover_9001(stacks, moves)
    print(f"{answer_b}")
    submit(answer_b, part="b", day=5, year=2022)