from pathlib import Path
import numpy as np


class HardDrive:
    def __init__(self, input, verbose=False):
        self.volume = self._arrange_data(input)
        self.du = sum(~np.isnan(self.volume))
        self.verbose = verbose
        if verbose:
            print(input)
            print(self)

    def __str__(self):
        a = ""
        for i in self.volume:
            if np.isnan(i):
                a += "."
            else:
                a += str(int(i))
        return a

    @staticmethod
    def _arrange_data(input):
        lens = np.array([int(i) for i in input])
        len_mem = np.sum(lens)
        vol = np.zeros(len_mem)

        i = 0
        s = 1
        d = 0
        for j in lens:
            if s > 0:
                vol[i : i + j] = d
                d += 1
            else:
                vol[i : i + j] = np.nan
            s *= -1
            i += j
        return vol

    @property
    def reverse(self):
        a = self.volume[~np.isnan(self.volume)]
        return a[::-1]

    def frag_sort(self):
        a = self.volume[:]
        r = self.reverse[:]
        t = np.isnan(self.volume)
        a[t] = r[: sum(t)]
        a[self.du :] = np.nan
        self.volume = a
        if self.verbose:
            print(self)

    @property
    def holes(self):
        (i,) = np.where(np.isnan(self.volume))
        d = np.diff(i, prepend=0) != 1
        print(np.diff(i, prepend=0) == 1)
        # TODO get hole width!
        return i[d]

    def clean_sort(self):
        self.holes

    @property
    def checksum(self):
        a = np.array(range(len(self.volume)))
        i = ~np.isnan(self.volume)
        return np.sum(self.volume[i] * a[i], dtype=int)


if __name__ == "__main__":
    filepath = Path("./example")
    # filepath = Path("./input")

    with open(filepath, "r") as filein:
        data = filein.read().rstrip()
    disk = HardDrive(data, verbose=True)
    print(disk.checksum)
    # disk.frag_sort()
    disk.clean_sort()
    print(disk.checksum)