from util.util import base
from util import grid
import numpy as np

class solve_day(base):
    def __init__(self, type='data'):
        super().__init__(type=type)

        self.data = grid.parse_as_grid(self.data)


    def part1(self):
        path, runs = grid.path_of_least_resistance(self.data)
        return sum([self.data[p[1]][p[0]] for p in path][1:])


    def part2(self):

        data = self.data.copy()

        starting_grid = data.copy()
        for i in range(1,5):
            new_data = starting_grid+1
            new_data = np.where(new_data>9, 1, new_data)
            starting_grid = new_data.copy()

            data = np.append(data, new_data, axis=1)


        starting_grid = data.copy()
        for i in range(1,5):
            new_data = starting_grid+1
            new_data = np.where(new_data>9, 1, new_data)
            starting_grid = new_data.copy()

            data = np.append(data, new_data, axis=0)


        path, runs = grid.path_of_least_resistance(data)
        return sum([data[p[1]][p[0]] for p in path][1:])


if __name__ == '__main__':
    s = solve_day('lines')
    s.sub(s.part1(), part='a')
    s.sub(s.part2(), part='b')