zerosleeps

import unittest
import utils

SHAPES = {
    "rock": {"base_score": 1, "loses_to": "paper", "wins_to": "scissors"},
    "paper": {"base_score": 2, "loses_to": "scissors", "wins_to": "rock"},
    "scissors": {"base_score": 3, "loses_to": "rock", "wins_to": "paper"},
}

SHAPE_MAPPINGS = {
    "A": "rock",
    "B": "paper",
    "C": "scissors",
    "X": "rock",
    "Y": "paper",
    "Z": "scissors",
}

OUTCOMES = {"win": 6, "lose": 0, "draw": 3}

OUTCOME_MAPPINGS = {"X": "lose", "Y": "draw", "Z": "win"}


def outcome(round):
    if SHAPE_MAPPINGS[round[0]] == SHAPE_MAPPINGS[round[1]]:
        return OUTCOMES["draw"]
    elif SHAPES[SHAPE_MAPPINGS[round[0]]]["loses_to"] == SHAPE_MAPPINGS[round[1]]:
        return OUTCOMES["win"]
    else:
        return OUTCOMES["lose"]


def reverse_outcome(round):
    if OUTCOME_MAPPINGS[round[1]] == "win":
        return SHAPES[SHAPES[SHAPE_MAPPINGS[round[0]]]["loses_to"]]["base_score"]
    elif OUTCOME_MAPPINGS[round[1]] == "lose":
        return SHAPES[SHAPES[SHAPE_MAPPINGS[round[0]]]["wins_to"]]["base_score"]
    else:
        return SHAPES[SHAPE_MAPPINGS[round[0]]]["base_score"]


def part_one(rounds):
    return sum(
        [
            outcome(round) + SHAPES[SHAPE_MAPPINGS[round[1]]]["base_score"]
            for round in rounds
        ]
    )


def part_two(rounds):
    return sum(
        [
            reverse_outcome(round) + OUTCOMES[OUTCOME_MAPPINGS[round[1]]]
            for round in rounds
        ]
    )


def parse_raw_input(raw_input):
    return [(round[0], round[-1]) for round in raw_input]


class TestExamples(unittest.TestCase):
    def setUp(self):
        self.rounds = parse_raw_input(["A Y", "B X", "C Z"])

    def test_part_one(self):
        self.assertEqual(part_one(self.rounds), 15)

    def test_part_two(self):
        self.assertEqual(part_two(self.rounds), 12)


if __name__ == "__main__":
    rounds = parse_raw_input(utils.get_raw_input("day_02_input.txt"))
    print(f"Part one: {part_one(rounds)}")
    print(f"Part two: {part_two(rounds)}")