rolls.py

import re
from enum import Enum
from itertools import product

# Regex Pattern
DICE_ROLL_PATTERN = re.compile(r'([ad])?(\d+)([+-]\d+)?((?:\+\d+d\d+)*)=(\d+)')

# Default Values
DEFAULT_MODIFIER = 0
DEFAULT_TARGET = 10

# Plotting Constants
FIGURE_SIZE = (10, 5)
GRID_STYLE = '--'
GRID_WIDTH = 0.5

class RollType(Enum):
    ADVANTAGE = 'a'
    DISADVANTAGE = 'd'
    NORMAL = None  # Represents a regular roll

def get_dice_distribution(dice_type):
    """
    Get the probability distribution for a dice of a given type.
    """
    return [1/dice_type] * dice_type

def parse_dice_roll(roll):
    match = DICE_ROLL_PATTERN.match(roll)
    if not match:
        raise ValueError(f'Invalid dice roll format: {roll}')

    roll_type = RollType(match.group(1))
    dice_type = int(match.group(2))  # Explicitly casting to integer
    modifier = int(match.group(3) or DEFAULT_MODIFIER)
    
    additional_dice = match.group(4)
    additional_rolls = []
    if additional_dice:
        additional_dice = additional_dice.split('+')[1:]
        for dice in additional_dice:
            dice_count, additional_dice_type = dice.split('d')  # Changed variable name to avoid overriding dice_type
            additional_rolls.append({
                'count': int(dice_count),
                'type': int(additional_dice_type)
            })
    
    target = int(match.group(5) or DEFAULT_TARGET)

    return {
        'roll_type': roll_type,
        'dice_type': dice_type,
        'modifier': modifier,
        'additional_rolls': additional_rolls,
        'target': target
    }

def get_adjusted_target(roll, target_value):
    return target_value

def calculate_probability_for_target(roll, target_value):
    dice_type = int(roll['dice_type'])  # Explicitly casting to integer
    target = target_value  # No need to adjust the target here
    roll_type = roll['roll_type']
    
    # Calculate the probability for the main dice roll
    possible_outcomes_main_dice = list(range(1, dice_type + 1))
    
    # Calculate the possible outcomes for additional dice rolls
    possible_outcomes_additional_dice = [list(range(1, additional_roll['type'] + 1)) for additional_roll in roll['additional_rolls'] for _ in range(additional_roll['count'])]

    # Calculate the cartesian product of all possible outcomes
    all_possible_outcomes = list(product(possible_outcomes_main_dice, *possible_outcomes_additional_dice))
    
    # Calculate the number of successful outcomes
    successful_outcomes = 0
    for outcome in all_possible_outcomes:
        total_outcome = sum(outcome) + roll['modifier']
        if total_outcome >= target:
            successful_outcomes += 1
            
    # Calculate the probability
    total_possible_outcomes = len(all_possible_outcomes)
    probability = (successful_outcomes / total_possible_outcomes) * 100
    
    # Handle advantage and disadvantage
    if roll_type == RollType.ADVANTAGE:
        probability = 1 - (1 - probability / 100) ** 2
        probability *= 100
    elif roll_type == RollType.DISADVANTAGE:
        probability = (probability / 100) ** 2
        probability *= 100

    return probability

def handle_rolls(input_string):
    rolls = input_string.split(',')
    results = []
    plots = []

    for roll_str in rolls:
        parsed_roll = parse_dice_roll(roll_str)
        max_outcome = int(parsed_roll['dice_type']) + sum([roll['count'] * roll['type'] for roll in parsed_roll['additional_rolls']])
        probabilities = [calculate_probability_for_target(parsed_roll, target_value) 
                         for target_value in range(1, max_outcome + 1)]
        target = int(parsed_roll['target'])
        if target <= len(probabilities):
            probability = probabilities[target - 1]
            results.append((roll_str, probability))
            plots.append(probabilities)
        else:
            results.append((roll_str, 0))
            plots.append(probabilities)

    return results, plots