Stock percentage distribution

stock_calculator.py

from typing import List, Dict

class AllocationError(Exception):
    """Custom exception for allocation errors."""
    pass

def validate_parameters(
    num_stocks: int,
    factor: float,
    last_stock_cap: float,
    min_allocation: float,
    bonuses: Dict[int, float]
) -> None:
    """Validate input parameters before processing the allocation."""
    
    if num_stocks < 2:
        raise AllocationError("There must be at least two stocks for allocation.")
    
    if not (0 < factor):
        raise AllocationError("The exponential factor must be a positive number.")
    
    if not (0 < last_stock_cap <= 100):
        raise AllocationError("The last_stock_cap must be between 0 and 100.")
    
    if not (0 < min_allocation <= 100):
        raise AllocationError("The min_allocation must be between 0 and 100.")
    
    if any(bonus < 0 for bonus in bonuses.values()):
        raise AllocationError("Bonuses cannot be negative.")
    
    if len(bonuses) >= num_stocks:
        raise AllocationError("Bonuses cannot be applied to all stocks.")
    
    if any(stock_idx < 1 or stock_idx > num_stocks for stock_idx in bonuses.keys()):
        raise AllocationError("Bonus stock indices must be between 1 and num_stocks.")
    
    if sum(bonuses.values()) > 100:
        raise AllocationError("The total bonuses exceed 100%.")

def round_percentages(percentages: List[float], decimal_places: int) -> List[float]:
    """Round the percentages to a specified number of decimal places."""
    return [round(pct, decimal_places) for pct in percentages]

def calculate_constrained_exponential_distribution_with_bonus(
    num_stocks: int,
    factor: float,
    last_stock_cap: float,
    min_allocation: float,
    bonuses: Dict[int, float],
    decimal_places: int = 2
) -> List[float]:
    """
    Calculate an exponential distribution of stock allocations,
    ensuring that the last stock is capped at a given percentage,
    all stocks meet a minimum allocation threshold, and some stocks
    receive a bonus.

    :param num_stocks: The total number of stocks.
    :param factor: The exponential growth factor.
    :param last_stock_cap: The maximum percentage allocation for the last stock (before bonuses).
    :param min_allocation: The minimum percentage allocation for each stock.
    :param bonuses: A dictionary where the keys are stock numbers (1-indexed) 
                    and the values are the bonuses to be applied to those stocks.
    :param decimal_places: The number of decimal places to round each stock percentage.
    :return: A list of percentages corresponding to each stock's allocation.
    """
    
    # Validate parameters before proceeding
    validate_parameters(num_stocks, factor, last_stock_cap, min_allocation, bonuses)

    # Step 1: Calculate the base allocations using the exponential factor
    allocations = [factor ** i for i in range(num_stocks)]
    
    # Step 2: Cap Stock #12 to the specified percentage
    last_stock_allocation = allocations[-1]
    scaling_factor = last_stock_cap / last_stock_allocation
    scaled_allocations = [allocation * scaling_factor for allocation in allocations]
    
    # Step 3: Ensure Stock #1 is at least the minimum allocation
    total_scaled_allocation = sum(scaled_allocations)
    
    stock_1_percentage = (scaled_allocations[0] / total_scaled_allocation) * 100
    if stock_1_percentage < min_allocation:
        adjustment_factor = min_allocation / stock_1_percentage
        scaled_allocations[0] *= adjustment_factor
    
    # Recalculate total allocation after adjustment
    total_adjusted_allocation = sum(scaled_allocations)
    
    # Normalize to get percentages
    percentages = [(allocation / total_adjusted_allocation) * 100 for allocation in scaled_allocations]
    
    # Step 4: Apply bonuses to specific stocks
    for stock_idx, bonus in bonuses.items():
        percentages[stock_idx - 1] += bonus  # stock_idx is 1-indexed, convert to 0-indexed
    
    # Step 5: Recalculate the total and adjust stocks below the minimum and non-bonused ones
    total_percentage_with_bonus = sum(percentages)
    
    # Find stocks that fall below the minimum allocation
    below_minimum_stocks = [i for i, pct in enumerate(percentages) if pct < min_allocation]
    
    # Ensure stocks below minimum have at least `min_allocation`
    for i in below_minimum_stocks:
        percentages[i] = min_allocation
    
    # Recalculate total after enforcing the minimum allocation
    total_percentage_after_min = sum(percentages)
    
    # Check if we still exceed 100%
    if total_percentage_after_min > 100:
        # The sum is greater than 100%, scale back non-bonus, non-min stocks proportionally
        excess_percentage = total_percentage_after_min - 100
        stocks_to_reduce = [i for i in range(num_stocks) 
                            if i not in bonuses and percentages[i] > min_allocation]
        
        # Calculate the total percentage of the stocks to reduce
        total_reduction_pool = sum(percentages[i] for i in stocks_to_reduce)
        
        # Reduce each stock's percentage proportionally
        for i in stocks_to_reduce:
            reduction_ratio = percentages[i] / total_reduction_pool
            percentages[i] -= reduction_ratio * excess_percentage

    # Step 6: Round the percentages
    percentages = round_percentages(percentages, decimal_places)

    # Step 7: Adjust if the sum isn't exactly 100%
    total_percentage = sum(percentages)
    if total_percentage != 100.0:
        difference = 100.0 - total_percentage
        # Adjust the last stock by the difference (could also spread the adjustment across multiple stocks)
        percentages[-1] += difference
        percentages[-1] = round(percentages[-1], decimal_places)
    
    # Ensure the final sum is exactly 100%
    assert round(sum(percentages), decimal_places) == 100.0, (
        f"Error: The total allocation does not add up to 100%. "
        f"Input: factor={factor}, last_stock_cap={last_stock_cap}, min_allocation={min_allocation}, bonuses={bonuses}. "
        f"Output: {percentages}. Total = {sum(percentages):.2f}%"
    )

    return percentages

def distribute_to_balance(percentages: List[float], balance: int) -> List[int]:
    """
    Distribute a given balance according to percentages, ensuring no decimals in the final result.

    :param percentages: A list of percentages for each stock.
    :param balance: The total balance to be distributed (e.g., 10,000).
    :return: A list of integer allocations for each stock.
    """
    
    # Step 1: Calculate the initial integer allocations
    raw_allocations = [balance * (pct / 100) for pct in percentages]
    int_allocations = [int(allocation) for allocation in raw_allocations]
    
    # Step 2: Track the remainders (fractional parts) to distribute later
    remainders = [(i, raw_allocations[i] - int_allocations[i]) for i in range(len(percentages))]
    
    # Step 3: Calculate the total amount allocated and the remaining balance
    total_allocated = sum(int_allocations)
    remaining_balance = balance - total_allocated
    
    # Step 4: Distribute the remaining balance based on the highest remainders
    # Sort remainders by the largest fractional part
    remainders.sort(key=lambda x: x[1], reverse=True)
    
    # Distribute the remaining units to stocks with the largest remainders
    for i in range(remaining_balance):
        stock_index = remainders[i][0]
        int_allocations[stock_index] += 1
    
    # Ensure the final distribution adds up to the total balance
    assert sum(int_allocations) == balance, (
        f"Error: The total allocation does not match the balance. "
        f"Balance = {balance}, Allocations = {int_allocations}, Total Allocated = {sum(int_allocations)}"
    )
    
    return int_allocations

Example usage:

# Parameters
num_stocks: int = 12
factor: float = 1.5  # Define your exponential growth factor
last_stock_cap: float = 15.0  # Cap Stock #12 at 15% (before bonus)
min_allocation: float = 2.0  # Ensure at least 2% allocation for each stock
bonuses: Dict[int, float] = {12: 8.0, 11: 6.0}  # Apply bonuses to Stock #12 (+8%) and Stock #11 (+6%)
decimal_places: int = 0 # Round stock values to 2 decimal places

try:
    # Get the distribution
    distribution: List[float] = calculate_constrained_exponential_distribution_with_bonus(
        num_stocks, factor, last_stock_cap, min_allocation, bonuses, decimal_places
    )

    # Display the result
    for i, percentage in enumerate(distribution, 1):
        print(f"Stock #{i}: {percentage:.2f}%")

    balance = 10000  # Example balance

    # Get the integer allocations based on the percentages
    allocations = distribute_to_balance(distribution, balance)

    # Display the result
    for i, allocation in enumerate(allocations, 1):
        print(f"Stock #{i}: {allocation} units")

except AllocationError as e:
    print(f"Error: {e}")