from collections import Counter
from functools import reduce
from math import floor
from typing import List, Iterator

from aoc import BaseAssignment


class IncompleteLine(Exception):
    def __init__(self, missing: List[str]):
        self.missing = missing

class CorruptedLine(Exception):
    def __init__(self, got: str, expected: str, *args, **kwargs):
        self.excpected = expected
        self.got = got

pairs = {
    '(': ')',
    '[': ']',
    '{': '}',
    '<': '>',
}


class Assignment(BaseAssignment):
    @classmethod
    def parse_line(cls, line: str):
        stack = []
        for character in line:
            if character in pairs:
                stack.append(character)
            elif character == pairs[stack[-1]]:
                stack.pop()
            else:
                raise CorruptedLine(character, pairs[stack[-1]])

        if len(stack) > 0:
            raise IncompleteLine(list(reversed([pairs[c] for c in stack])))


class AssignmentOne(Assignment):
    example_result = 26397

    penalties = {
        ')': 3,
        ']': 57,
        '}': 1197,
        '>': 25137,
    }

    def run(self, input: Iterator[str]) -> int:
        invalid_characters = []
        for line in input:
            try:
                self.parse_line(line)
            except CorruptedLine as c:
                invalid_characters.append(c.got)
            except IncompleteLine:
                continue

        counter = Counter(invalid_characters)

        return sum(counter[k] * v for k, v in self.penalties.items())


class AssignmentTwo(Assignment):
    example_result = 288957

    penalties = {
        ')': 1,
        ']': 2,
        '}': 3,
        '>': 4,
    }

    def run(self, input: Iterator[str]):
        scores = []
        for line in input:
            try:
                self.parse_line(line)
            except CorruptedLine:
                continue
            except IncompleteLine as c:
                scores.append(reduce(
                    lambda out, item: (out * 5) + self.penalties[item],
                    c.missing,
                    0
                ))

        return sorted(scores)[floor(len(scores) / 2)]