from pprint import pprint as pp

from anim import Animate
import shared
import matrix

MATRIX_SIZE = 400
ATRIX_SIZE = 40  # sample
HALF = MATRIX_SIZE // 2

H = [HALF, HALF]
T = [HALF, HALF]


# returns (x,y) - so remember to take dx,dy lower
DIRS = {
    "U": matrix.M_U,  # (0, -1)
    "D": matrix.M_D,
    "L": matrix.M_L,
    "R": matrix.M_R,
    "UR": matrix.M_UR,  # (+1, -1)
    "DR": matrix.M_DR,
    "UL": matrix.M_UL,
    "DL": matrix.M_DL,
}


def part1(steps):
    field = matrix.matrix_of_size(MATRIX_SIZE, MATRIX_SIZE)
    anim = Animate(field, day="09")

    for direction, amount in steps:
        # move HEAD
        amount = int(amount)
        for x in range(amount):

            dx, dy = DIRS[direction]
            H[0] += dy
            H[1] += dx
            HY, bX = H
            # head is in new position
            kY, kX = T
            field[kY][kX] = 1  # 'T'

            # Check if same
            if kY == HY and kX == bX:
                continue
            touching = False
            for _, d in DIRS.items():
                dx, dy = d
                t_y, t_x = H[0] + dy, H[1] + dx
                if kY == t_y and kX == t_x:
                    touching = True
                    break
            if touching:
                continue

            # assume the same row/column, no movement needed
            tY = 0
            tX = 0

            # if Head X is less than Tail X, move Tail X 1 left
            if bX < kX:
                tX = -1
            # if Head X is greater than Tail X, move Tail X 1 left
            elif bX > kX:
                tX = 1

            # if Head Y is less than Tail Y, move Tail Y 1 up
            if HY < kY:
                tY = -1
            # if Head Y is greater than Tail Y, move Tail Y 1 down
            elif HY > kY:
                tY = 1

            T[0] += tY
            T[1] += tX
            anim.add_frame(field)

        field[T[0]][T[1]] = 1  # mark visited tails
    print(matrix.sum_matrix(field))
    anim.animate()


def part2(steps):
    field = matrix.matrix_of_size(MATRIX_SIZE, MATRIX_SIZE)
    anim = Animate(field, day="09")
    S = [
        [HALF, HALF],
    ]  # HEAD ONLY
    _cur_frame = 0
    for x in range(9):
        S.append([HALF, HALF])

    for direction, amount in steps:
        # move HEAD
        amount = int(amount)
        for x in range(amount):
            HEAD_dx, HEAD_dy = DIRS[direction]
            S[0][0] += HEAD_dy
            S[0][1] += HEAD_dx
            # print(HEAD_dy,HEAD_dx, S)
            # HEAD MOVED one step

            for knot_idx, _ in enumerate(S):
                if knot_idx == 0:
                    # Skip Head
                    continue
                knot = S[knot_idx]
                kY, kX = knot

                bY, bX = S[knot_idx - 1]

                # Check if same
                if kY == bY and kX == bX:
                    continue

                touching = False
                for _, d in DIRS.items():
                    dx, dy = d
                    t_y, t_x = S[knot_idx - 1][0] + dy, S[knot_idx - 1][1] + dx
                    if kY == t_y and kX == t_x:
                        touching = True
                        break
                if touching:
                    continue

                # assume the same row/column, no movement needed
                tY, tX = 0, 0

                # if Head X is less than Tail X, move Tail X 1 left
                if bX < kX:
                    tX = -1
                # if Head X is greater than Tail X, move Tail X 1 left
                elif bX > kX:
                    tX = 1

                # if Head Y is less than Tail Y, move Tail Y 1 up
                if bY < kY:
                    tY = -1
                # if Head Y is greater than Tail Y, move Tail Y 1 down
                elif bY > kY:
                    tY = 1

                S[knot_idx][0] += tY
                S[knot_idx][1] += tX
                if _cur_frame % 25 == 0:
                    anim.add_frame(field)
                _cur_frame += 1
            field[S[-1][0]][S[-1][1]] = 1  # mark visited tails
    print(matrix.sum_matrix(field))
    anim.animate()


def main():
    rows = [x.split(" ") for x in shared.load_rows(9)]
    # part1(rows)
    part2(rows)


if __name__ == "__main__":
    main()