from die import Die
from tqdm import tqdm

class FiveDice:
    def __init__(self):
        self.dice = [Die() for number in range(5)]

    def roll(self):
        for die in self.dice:
            die.roll()
        return self.faces()

    def faces(self):
        return [die.face for die in self.dice]

    def all_ones(self):
        for face in self.faces():
            if face != 1:
                return False
        return True

    def is_three_of_a_kind(self):
        face_counts = {}
        for face in self.faces():
            if face in face_counts:
                face_counts[face] += 1
            else:
                face_counts[face] = 1
        for count in face_counts.values():
            if count >= 3:
                return True
        return False
    
    def is_four_of_a_kind(self):
        face_counts = {}
        for face in self.faces():
            if face in face_counts:
                face_counts[face] += 1
            else:
                face_counts[face] = 1
        for count in face_counts.values():
            if count >= 4:
                return True
        return False


dice = FiveDice()
successes = 0
trials = 1000000
for trial in tqdm(range(trials)):
    dice.roll()
    if dice.all_ones():
        successes += 1
    

print(f"Probability of all ones: {successes/trials}")

three_of_a_kind_successes = 0
for trial in tqdm(range(trials)):
    dice.roll()
    if dice.is_three_of_a_kind():
        three_of_a_kind_successes += 1
print(f"Probability of three of a kind: {three_of_a_kind_successes/trials}")

four_of_a_kind_successes = 0
for trial in tqdm(range(trials)):
    dice.roll()
    if dice.is_four_of_a_kind():
        four_of_a_kind_successes += 1
print(f"Probability of four of a kind: {four_of_a_kind_successes/trials}")