# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

import os
import numpy as np
import aocd
from aocd.models import Puzzle
from aocd import submit

current_day = 13
current_year = 2021
puzzle = Puzzle(year=current_year, day=current_day)
puzzle

coords, folds = puzzle.input_data.split('\n\n')
coords = coords.splitlines()
folds = folds.splitlines()
coords = np.array([np.int32(coord.split(',')) for coord in coords])
folds = [fold.replace('fold along ','').split('=') for fold in folds]
coords, folds

# # Part 1

num_cols, num_rows = np.max(coords,axis=0) - np.min(coords,axis=0) + 1
print(f"{num_cols = }, {num_rows = }")
grid = np.zeros((num_rows, num_cols), dtype=int)
for x,y in coords:
    grid[y,x] = 1

fold_loc = 655
to_fold_over = grid[:,fold_loc+1:]
# flip horizontally
to_fold_over = to_fold_over[:,::-1]
assert to_fold_over.shape == grid[:,fold_loc-to_fold_over.shape[1]:fold_loc].shape
grid[:,fold_loc-to_fold_over.shape[1]:fold_loc] = grid[:,fold_loc-to_fold_over.shape[1]:fold_loc] | to_fold_over
grid = grid[:,:fold_loc]


grid.sum()

puzzle.answer_a = 716

# # Part 2

import matplotlib.pyplot as plt
# %matplotlib inline

def do_fold(grid, fold_axis, fold_loc):
    if fold_axis == 'x':
        to_fold_over = grid[:,fold_loc+1:]
        # flip horizontally
        to_fold_over = to_fold_over[:,::-1]
        assert to_fold_over.shape == grid[:,fold_loc-to_fold_over.shape[1]:fold_loc].shape
        grid[:,fold_loc-to_fold_over.shape[1]:fold_loc] |= to_fold_over
        grid = grid[:,:fold_loc]
    else: # 'y'
        to_fold_over = grid[fold_loc+1:,:]
        # flip vertically
        to_fold_over = to_fold_over[::-1,:]
        assert to_fold_over.shape == grid[fold_loc-to_fold_over.shape[0]:fold_loc,:].shape
        grid[fold_loc-to_fold_over.shape[0]:fold_loc,:] |= to_fold_over
        grid = grid[:fold_loc,:]
    return grid


num_cols, num_rows = np.max(coords,axis=0) - np.min(coords,axis=0) + 1
print(f"{num_cols = }, {num_rows = }")
grid = np.zeros((num_rows, num_cols), dtype=int)
for x,y in coords:
    grid[y,x] = 1
for fold_axis, fold_loc in folds:
    grid = do_fold(grid,fold_axis, int(fold_loc))
plt.imshow(grid)

puzzle.answer_b = 'RPCKFBLR'



StarcoderdataJupyterScriptsDedupFiltered