Untitled

from crc import CrcCalculator, Crc8

class DigitRecover:
    def __init__(self):
        self.crc_calculator = CrcCalculator(Crc8.CCITT, True)

    def gen_idx(self, n_d, n_c):
        if n_d == 0:
            return [[]]
        ret = []
        for c in range(n_c):
            for arr in self.gen_idx(n_d - 1, n_c):
                ret.append([c] + arr)
        return ret

    def rank_8(self, digits, confs):
        assert len(digits[0]) == 8 and len(digits) == 3
        assert len(confs[0]) == 8 and len(confs) == 3
        arr_idx = self.gen_idx(8, 3)
        arr_tup = []
        for i in range(len(arr_idx)):
            arr = []
            mul = 1
            for j in range(8):
                arr.append((digits[arr_idx[i][j]][j]))
                x = confs[arr_idx[i][j]][j]
                mul *= x
            arr_tup.append((arr, mul))
        arr_tup.sort(key=lambda ele: ele[1], reverse=True)
        return arr_tup

    def calculate(self, digits):
        checksum = self.crc_calculator.calculate_checksum(bytes(digits))
        return checksum

    def verify(self, digits, confs, exp):
        arr_tup = self.rank_8(digits, confs)
        for arr, _ in arr_tup:
            if self.crc_calculator.verify_checksum(bytes(arr), exp):
                return arr
        return arr_tup[0]

    def encode(self, digits):
        rem = len(digits) % 8
        if rem:
            digits = digits + [0] * (8 - rem)
        ret = []
        for i in range(0, len(digits), 8):
            d = digits[i:i + 8]
            e = self.calculate(d)
            ret.append(e)
        return ret

    def recover(self, digits, confs, exps):
        assert len(digits) == 3
        assert len(confs) == 3
        rem = len(digits[0]) % 8
        if rem:
            for i in range(len(digits)):
                digits[i] = digits[i] + [0] * (8 - rem)
                confs[i] = confs[i] + [1] * (8 - rem)
        ret = []
        for i in range(0, len(digits[0]), 8):
            d = [arr[i:i + 8] for arr in digits]
            c = [arr[i:i + 8] for arr in confs]
            e = exps[i // 8]
            ret.append(self.verify(d, c, e))
        return ret


digits = [[5, 0, 1, 8, 2, 4, 6, 8, 9, 3, 6, 3, 3, 8, 5], [4, 6, 2, 8, 0, 2, 1, 6, 6, 8, 0, 0, 4, 3, 0], [2, 8, 0, 7, 0, 7, 6, 5, 2, 5, 6, 7, 6, 4, 1]]
confs = [[333, 323, 352, 366, 311, 393, 362, 308, 379, 309, 347, 378, 350, 350, 337], [266, 295, 288, 213, 233, 289, 264, 241, 203, 236, 248, 239, 251, 209, 245], [145, 159, 170, 149, 189, 163, 118, 191, 153, 106, 186, 105, 169, 191, 136]]

dr = DigitRecover()
exps = dr.encode(digits[0])
recs = dr.recover(digits, confs, exps)
print(recs)