Untitled

#!/usr/bin/env python3

import collections
import copy

SAMPLE = False

def load_input():
    with open("sample.txt" if SAMPLE else "input.txt") as f:
        return [line.strip() for line in f.readlines()]

class Decoder:
    def __init__(self, top, topleft, topright, middle, bottomleft, bottomright, bottom):
        print(f"New decoder: {top}{topleft}{topright}{middle}{bottomleft}{bottomright}{bottom}")
        self.top = top
        self.topleft = topleft
        self.topright = topright
        self.middle = middle
        self.bottomleft = bottomleft
        self.bottomright = bottomright
        self.bottom = bottom

    def decode(self, combo):
        assert all([c in 'abcdefg' for c in combo])
        if len(combo) == 2:
            assert all([seg in combo for seg in (self.topright, self.bottomright)])
            return 1
        elif len(combo) == 3:
            assert all([seg in combo for seg in (self.top, self.topright, self.bottomright)])
            return 7
        elif len(combo) == 4:
            assert all([seg in combo for seg in (self.topleft, self.topright, self.middle, self.bottomright)])
            return 4
        elif len(combo) == 5:
            if all([seg in combo for seg in (self.topright, self.bottomleft)]):
                return 2
            elif all([seg in combo for seg in (self.topleft, self.bottomright)]):
                return 5
            elif all([seg in combo for seg in (self.topright, self.bottomright)]):
                return 3
            raise ValueError(combo, self.topleft, self.topright, self.bottomleft, self.bottomright)
        elif len(combo) == 6:
            if self.topright not in combo:
                return 6
            elif self.bottomleft not in combo:
                return 9
            elif self.middle not in combo:
                return 0
            raise ValueError(combo, topleft, topright, bottomleft, bottomright)
        elif len(combo) == 7:
            assert all([seg in combo for seg in (self.top, self.topleft, self.topright, self.middle, self.bottomleft, self.bottomright, self.bottom)])
            return 8
        raise ValueError(len(combo), combo)

def generate_decoder(all_combinations):
    combos_by_size = collections.defaultdict(lambda: [])
    for combo in all_combinations:
        combos_by_size[len(combo)].append(set(combo))
    top = determine_top(copy.copy(combos_by_size))
    middle, topleft = determine_middle_and_topleft(copy.copy(combos_by_size))
    topright, bottomleft, bottomright = \
            determine_topright_bottomleft_bottomright(copy.copy(combos_by_size), topleft)
    bottom = (set('abcdefg') - set([top, topleft, topright, middle, bottomleft, bottomright])).pop()
    return Decoder(top, topleft, topright, middle, bottomleft, bottomright, bottom)

def determine_top(combos_by_size):
    seven = combos_by_size[3][0]
    one = combos_by_size[2][0]
    assert len(seven - one) == 1
    return (seven - one).pop()

def determine_middle_and_topleft(combos_by_size):
    four = combos_by_size[4][0]
    one = combos_by_size[2][0]
    assert len(four - one) == 2
    possible_middles = four - one
    two_three_five = combos_by_size[5]
    middle = None
    topleft = None
    for possible_middle in possible_middles:
        in_all_three = True
        for fiver in two_three_five:
            if possible_middle not in fiver:
                in_all_three = False
                break
        if in_all_three:
            middle = possible_middle
        else:
            topleft = possible_middle
    assert middle is not None
    assert topleft is not None
    return middle, topleft

def determine_topright_bottomleft_bottomright(combos_by_size, topleft):
    two_three_five = copy.copy(combos_by_size[5])
    histogram = collections.defaultdict(lambda: 0)
    for fiver in two_three_five:
        for segment in fiver:
            histogram[segment] += 1
    singletons = set()
    doubletons = set()
    for segment, frequency in histogram.items():
        if frequency == 1:
            singletons.add(segment)
        elif frequency == 2:
            doubletons.add(segment)
    assert topleft in singletons
    assert len(singletons) == 2
    assert len(doubletons) == 2
    assert doubletons == combos_by_size[2][0]
    bottomleft = (singletons - set(topleft)).pop()
    possible_bottomright = [segment for segment in doubletons if any([segment in possible_five and topleft in possible_five for possible_five in two_three_five])]
    assert len(possible_bottomright) == 1
    bottomright = possible_bottomright.pop()
    possible_topright = [segment for segment in doubletons if any([segment in possible_two and bottomleft in possible_two for possible_two in two_three_five])]
    assert len(possible_topright) == 1
    topright = possible_topright.pop()
    return topright, bottomleft, bottomright

def main():
    displays = load_input()
    result = 0
    for i, display in enumerate(displays):
        print(f"\nNew display #{i}: {display}")
        all_combinations, output_combinations = display.split(" | ")
        print(f"All combinations: {all_combinations}")
        decoder = generate_decoder(all_combinations.split())
        output_digits = [decoder.decode(combo) for combo in output_combinations.split()]
        output_number = int(''.join([str(digit) for digit in output_digits]))
        assert output_number < 10000
        result += output_number
    print(result)

if __name__ == "__main__":
    main()