#!/bin/python3

from enum import Enum
from collections import defaultdict
from itertools import pairwise, combinations

class Move(Enum):
    LEFT = 1
    RIGHT = 2
    DOWN = 3

moves = list()

with open('inputs/day17.txt', 'r') as f:
    for line in f:
        for char in line.strip():
            if char == '<':
                moves.append(Move.LEFT)
            elif char == '>':
                moves.append(Move.RIGHT)

def move_stream():
    while True:
        for i, d in enumerate(moves):
            yield (i, d)
            yield (i, Move.DOWN)

move_map = {Move.LEFT: (-1, 0), Move.RIGHT: (1, 0), Move.DOWN: (0, -1)}

class Rock():
    def __init__(self, shape):
        self.blocks = shape

    def init_position(self, y_top):
        self.blocks = list(map(lambda pos: (pos[0] + 2, pos[1] + y_top + 4), self.blocks))

    def move(self, chamber, move):
        new_blocks = self.move_blocks(move)
        if self.check_valid(chamber, new_blocks):
            self.blocks = new_blocks
            return True
        else:
            return False

    def move_blocks(self, move):
        (movex, movey) = move_map[move]
        return list(map(lambda block: (block[0] + movex, block[1] + movey), self.blocks))

    def check_valid(self, chamber, blocks):
        for (x, y) in blocks:
            if x > 6 or x < 0 or chamber[y][x]:
                return False

        return True

shapes = [
    [(0, 0), (1, 0), (2, 0), (3, 0)],
    [(0, 1), (1, 1), (2, 1), (1, 0), (1, 2)],
    [(0, 0), (1, 0), (2, 0), (2, 1), (2, 2)],
    [(0, 0), (0, 1), (0, 2), (0, 3)],
    [(0, 0), (1, 0), (0, 1), (1, 1)]
]

def rock_stream():
    while True:
        for i, shape in enumerate(shapes):
            yield (i, Rock(shape.copy()))

chamber = defaultdict(lambda: defaultdict(lambda: False))
for i in range(7):
    chamber[-1][i] = True

def print_chamber(chamber, rock, print_index):
    keys = list(reversed(sorted(chamber.keys())))
    for key in keys:
        if print_index:
            print(f'{key}\t', end='')
        for i in range(7):
            if chamber[key][i]:
                print('#', end='')
            elif rock and (i, key) in rock.blocks:
                print('@', end='')
            else:
                print('.', end='')
        print('')

def prune_chamber(chamber):
    ys = list(reversed(sorted(chamber.keys())))
    prune_height = None
    for y_high, y_low in pairwise(ys):
        prunable = True
        for x in range(7):
            if not chamber[y_high][x] and not chamber[y_low][x]:
                prunable = False
                break
        if prunable:
            prune_height = y_low
            break
    if prune_height:
        for y in ys:
            if y < prune_height:
                del chamber[y]
    return prune_height != None

def row_to_mask(row):
    mask = 0
    bit = 1
    for x in range(6, -1, -1):
        if row[x]:
            mask = mask | bit
        bit = bit << 1
    return mask

barrier_fingerprints = defaultdict(lambda: list())

rock_count = 20220
prune_freq = 100
y_top = -1
moves_gen = move_stream()
rocks_gen = rock_stream()
rocks_fallen = 0
rock_idx = None
for rock_idx, rock in rocks_gen:
    rock.init_position(y_top)
    for move_idx, move in moves_gen:
        moved = rock.move(chamber, move)
        if not moved and move == Move.DOWN:
            for (x, y) in rock.blocks:
                if y > y_top:
                    y_top = y
                chamber[y][x] = True
            rocks_fallen += 1
            break

    ys = set(map(lambda x: x[1], rock.blocks))
    for y in ys:
        all_filled = True
        for x in range(7):
            if not chamber[y][x] and not chamber[y-1][x]:
                all_filled = False
                break
        if all_filled:
            mask_high = row_to_mask(chamber[y])
            mask_low = row_to_mask(chamber[y-1])
            fingerprint = (rock_idx, move_idx, mask_high, mask_low)
            barrier_fingerprints[fingerprint].append((rocks_fallen, y))

    if rocks_fallen == rock_count:
        break
print('part1', y_top + 1)

candidate_fp = None
candidate_ys_count = 0
for fp, ys in barrier_fingerprints.items():
    l = len(ys)
    if l > candidate_ys_count:
        candidate_ys_count = l
        candidate_fp = fp
# print('candidate FP:', candidate_fp)
first_occ = barrier_fingerprints[candidate_fp][0]
second_occ = barrier_fingerprints[candidate_fp][1]
height_cycle = second_occ[1] - first_occ[1]
rock_cycle = second_occ[0] - first_occ[0]
# print('first occurance:', 'height', first_occ[1], 'fallen', first_occ[0])
# print('second occurance:', 'height', second_occ[1], 'fallen', second_occ[0])
# print('height cycle:', second_occ[1] - first_occ[1])
# print('rock cycle:', second_occ[0] - first_occ[0])

rocks_left = 1000000000000 - first_occ[0]
superheight = first_occ[1]
superheight += (rocks_left // rock_cycle) * height_cycle
rocks_left %= rock_cycle
print(rocks_left, superheight)

# We just need $rocks_left rocks more!
# Move to the next cycle point
# Then simulate $rocks_left movements to get the height this generates.
# Add this to superheight then subtract one height_cycle.

for rock_idx, rock in rocks_gen:
    rock.init_position(y_top)
    for move_idx, move in moves_gen:
        moved = rock.move(chamber, move)
        if not moved and move == Move.DOWN:
            for (x, y) in rock.blocks:
                if y > y_top:
                    y_top = y
                chamber[y][x] = True
            rocks_fallen += 1
            break

    ys = set(map(lambda x: x[1], rock.blocks))
    cycle_start = False
    for y in ys:
        all_filled = True
        for x in range(7):
            if not chamber[y][x] and not chamber[y-1][x]:
                all_filled = False
                break
        if all_filled:
            mask_high = row_to_mask(chamber[y])
            mask_low = row_to_mask(chamber[y-1])
            fingerprint = (rock_idx, move_idx, mask_high, mask_low)
            if fingerprint == candidate_fp:
                cycle_start = True
                break
    if cycle_start:
        break

start_height = y_top

for rock_idx, rock in rocks_gen:
    rock.init_position(y_top)
    for move_idx, move in moves_gen:
        moved = rock.move(chamber, move)
        if not moved and move == Move.DOWN:
            for (x, y) in rock.blocks:
                if y > y_top:
                    y_top = y
                chamber[y][x] = True
            rocks_left -= 1
            break

    if rocks_left == 0:
        break

end_height = y_top
superheight += end_height - start_height
print('part2', superheight + 3)

# for fp, ys in barrier_fingerprints.items():
#     print('fingerprint:', fp)
#     print('count:', len(ys))

# NEW IDEA
# Find a cycle.
# A cycle can be characterized by: rock rotation, direction rotation and rock formation.
# Problem is, we need to save a lot of the rock formation, because new rocks can fall down a long while.
# Additionally, we cannot just save the height of each column, because rocks can move underneath hanging arches.
# Possible solution: find lines that block this.
# Possibly in one line, definitely in two lines.

# IDEA:
# Create bitmap.
# Easily compare bitmaps with eachother to find cycle

# def print_mask(mask):
#     bit = 1 << 6
#     for i in range(7):
#         if mask & (bit >> i):
#             print('1', end='')
#         else:
#             print('0', end='')
#     print('')

# def print_bitmap(bitmap):
#     for row in bitmap:
#         print_mask(row)

# bitmap = []
# ys = list(reversed(sorted(chamber.keys())))
# y_min = min(ys)
# y_max = max(ys)
# for y in range(y_min + 1, y_max + 1):
#     mask = 0
#     bit = 1
#     for x in range(6, -1, -1):
#         if chamber[y][x]:
#             mask = mask | bit
#         bit = bit << 1
    # print_mask(mask)
    # bitmap.append(mask)

# def in_order(l):
#     last = None
#     for x in l:
#         if last == None or x > last:
#             last = x
#         else:
#             return False
#     return True

# def same_distance(l):
#     [y1, y2, y3] = l
#     return y2 - y1 == y3 - y2

# def is_cycle(bitmap, l):
#     [y1, y2, y3] = l
#     distance = y3 - y2
#     for i in range(distance + 1):
#         if bitmap[y1+i] != bitmap[y2]+i:
#             return False
#     return True

# Now try to find the cycle!
# We loop over every row.
# We find all indices of this row in the whole chamber
# checked = []
# print('finding cycle')
# for row in bitmap:
#     print('row:', row)
#     if row not in checked:
#         checked.append(row)
#         indices = [i for i, x in enumerate(bitmap) if x == row]
#         if len(indices) >= 3:
#             for maybe_cycle in combinations(indices, 3):
#                 if in_order(maybe_cycle) and same_distance(maybe_cycle):
#                     if is_cycle(bitmap, maybe_cycle):
#                         print('FOUND CYCLE:', maybe_cycle)

# # print_bitmap(bitmap)