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()]

def main():
    displays = load_input()
    result = 0
    for i, display in enumerate(displays):
        print(f"\nNew display #{i}: {display}")
        all_digits, output_digits = display.split(" | ")
        all_digits = [
                ''.join(sorted(digit))
                for digit
                in all_digits.split()
        ]
        output_digits = [
                ''.join(sorted(digit))
                for digit
                in output_digits.split()
        ]

        def digits_with_n_segments(n):
            return [digit for digit in all_digits if len(digit) == n]

        def digit_with_n_segments(n):
            digits = digits_with_n_segments(n)
            assert len(digits) == 1
            return digits[0]

        def digits_with_n_segments_containing_other_digit(n, other):
            return [
                    digit
                    for digit in digits_with_n_segments(n)
                    if all([segment in digit for segment in other])
            ]

        def digit_with_n_segments_containing_other_digit(n, other):
            digits = digits_with_n_segments_containing_other_digit(n, other)
            assert len(digits) == 1
            return digits[0]

        def digit_with_n_segments_containing_other_digit_but_not_this(n, other,
                not_me):
            digits = [
                    digit
                    for digit in digits_with_n_segments_containing_other_digit(n,
                        other)
                    if digit != not_me
            ]
            assert len(digits) == 1
            return digits[0]

        def digits_with_n_segments_in_other_digit(n, other):
            return [
                    digit
                    for digit in digits_with_n_segments(n)
                    if all([segment in other for segment in digit])
            ]

        def digit_with_n_segments_in_other_digit_but_not_this(n, other, not_me):
            digits = [
                    digit
                    for digit in digits_with_n_segments_in_other_digit(n, other)
                    if digit != not_me
            ]
            assert len(digits) == 1
            return digits[0]

        def digit_with_n_segments_but_not_these(n, not_these):
            digits = [
                    digit
                    for digit in digits_with_n_segments(n)
                    if digit not in not_these
            ]
            assert len(digits) == 1
            return digits[0]

        one = digit_with_n_segments(2)
        seven = digit_with_n_segments(3)
        four = digit_with_n_segments(4)
        eight = digit_with_n_segments(7)
        three = digit_with_n_segments_containing_other_digit(5, one)
        nine = digit_with_n_segments_containing_other_digit(6, three)
        five = digit_with_n_segments_in_other_digit_but_not_this(5, nine, three)
        six = digit_with_n_segments_containing_other_digit_but_not_this(6, five,
                nine)
        two = digit_with_n_segments_but_not_these(5, [three, five])
        zero = digit_with_n_segments_but_not_these(6, [six, nine])

        decoder = {
                one: "1",
                two: "2",
                three: "3",
                four: "4",
                five: "5",
                six: "6",
                seven: "7",
                eight: "8",
                nine: "9",
                zero: "0",
        }

        output_number = int(''.join(decoder[digit] for digit in output_digits))
        assert output_number < 10000
        result += output_number
    print(result)

if __name__ == "__main__":
    main()