from pprint import pprint as pp
import shared
import matrix


def part1(mx):
    SIZE = len(mx)
    MAX_IDX = SIZE - 1
    data = matrix.matrix_of_size(SIZE, SIZE)
    data = matrix.set_matrix_dict(data)

    for row in range(SIZE):
        for col in range(SIZE):
            data[row][col] = {"l": True, "r": True, "u": True, "d": True}

    for row, _row in enumerate(mx):
        # Skip outer edges
        if row == 0 or row == MAX_IDX:
            continue
        for col, _val in enumerate(mx[row]):
            # Skip outer edges
            if col == 0 or col == MAX_IDX:
                continue

            # TODO: slice
            # Get a list of indexes in the current direction to check
            trees_u_indexes = list(range(0, row))
            trees_l_indexes = list(range(0, col))
            trees_d_indexes = list(range(MAX_IDX, row, -1))
            trees_r_indexes = list(range(col + 1, SIZE))

            # Get the trees in the line
            trees_u = [mx[y][col] for y in trees_u_indexes]
            trees_l = [mx[row][x] for x in trees_l_indexes]
            trees_d = [mx[y][col] for y in trees_d_indexes]
            trees_r = [mx[row][x] for x in trees_r_indexes]

            # find tallest tree in line
            highest_u = max(trees_u)
            highest_l = max(trees_l)
            highest_d = max(trees_d)
            highest_r = max(trees_r)

            # if there's anything as tall, or taller than this tree in that line
            # then it's not visible
            if highest_u >= _val:
                data[row][col]["u"] = False

            if highest_l >= _val:
                data[row][col]["l"] = False

            if highest_d >= _val:
                data[row][col]["d"] = False

            if highest_r >= _val:
                data[row][col]["r"] = False

    # count visible
    count = 0
    for row in range(SIZE):
        for col in range(SIZE):
            u = data[row][col]["u"]
            l = data[row][col]["l"]
            d = data[row][col]["d"]
            r = data[row][col]["r"]
            if any((r, l, u, d)):
                count += 1
    print(count)


def part2(mx):
    SIZE = len(mx)
    MAX_IDX = SIZE - 1
    data = matrix.matrix_of_size(SIZE, SIZE)
    data = matrix.set_matrix_dict(data)

    high_score = 0

    for row in range(SIZE):
        for col in range(SIZE):
            data[row][col] = {"value": mx[row][col]}

    for row, _row in enumerate(mx):
        for col, _val in enumerate(mx[row]):
            # Get a list of indexes in the current direction to check
            # TODO: slice
            trees_u_indexes = list(range(0, row))
            trees_l_indexes = list(range(0, col))
            trees_d_indexes = list(range(MAX_IDX, row, -1))
            trees_r_indexes = list(range(col + 1, SIZE))

            # Grab the trees potentially in sight
            trees_u = [mx[y][col] for y in trees_u_indexes]
            trees_l = [mx[row][x] for x in trees_l_indexes]
            trees_d = [mx[y][col] for y in trees_d_indexes]
            trees_r = [mx[row][x] for x in trees_r_indexes]

            # Reverse everything but r, as everything is "away" from the tree
            data[row][col]["u"] = list(reversed(trees_u))
            data[row][col]["l"] = list(reversed(trees_l))
            data[row][col]["d"] = list(reversed(trees_d))
            data[row][col]["r"] = trees_r

            # Calculate score
            score = 1  # identity yoooo
            cell = data[row][col]
            val = cell["value"]
            # Get the score of visible trees in each direction
            for direction in ("u", "l", "d", "r"):
                in_line = cell[direction]
                if not in_line:
                    # we're on an edge
                    score = 0
                    break

                line_score = 0
                for idx, tree in enumerate(in_line):
                    # for every tree, check if its as tall or taller,
                    # stop countig after that
                    line_score += 1
                    if tree >= val:
                        break
                score = score * line_score
            if score > high_score:
                high_score = score
    print(high_score)


def part2_with_fixes(mx):
    SIZE = len(mx)
    MAX_IDX = SIZE - 1
    data = matrix.matrix_of_size(SIZE, SIZE)
    data = matrix.set_matrix_dict(data)

    high_score = 0

    for row in range(SIZE):
        for col in range(SIZE):
            data[row][col] = {"value": mx[row][col]}

    for row, _row in enumerate(mx):
        for col, _val in enumerate(mx[row]):
            lineofsight = matrix.line_of_sight(mx, row, col)

            data[row][col]["u"] = lineofsight["U"]
            data[row][col]["l"] = lineofsight["L"]
            data[row][col]["d"] = lineofsight["D"]
            data[row][col]["r"] = lineofsight["R"]

            # Calculate score
            score = 1  # identity yoooo
            val = data[row][col]["value"]
            # Get the score of visible trees in each direction
            for direction in ("u", "l", "d", "r"):
                in_line = data[row][col][direction]
                if not in_line:
                    # we're on an edge
                    score = 0
                    break

                line_score = 0
                for idx, tree in enumerate(in_line):
                    # for every tree, check if its as tall or taller,
                    # stop countig after that
                    line_score += 1
                    if tree >= val:
                        break
                score = score * line_score
            if score > high_score:
                high_score = score
    print(high_score)


def main():
    mx = matrix.load_matrix_file(shared.get_fname(8))
    with shared.elapsed_timer() as elapsed:
        part1(mx)
        print(elapsed())
    with shared.elapsed_timer() as elapsed:
        part2(mx)
        print(elapsed())
    with shared.elapsed_timer() as elapsed:
        part2_with_fixes(mx)
        print(elapsed())


if __name__ == "__main__":
    main()