lockdrop-simulation/lockdrop/calculations.py

"""
Core lockdrop calculation functions.

This module contains the main mathematical functions for calculating token allocations,
penalties, bonuses, and final distributions.
"""

from decimal import Decimal, ROUND_DOWN
from .constants import (
    TOKEN_PRECISION, LOCKDROP_ALLOCATION, STAR_ALLOCATION_PERCENT,
    GALAXY_ALLOCATION_PERCENT,
    LOCKDROP_DURATION_BLOCKS, PENALTY_RATES
)


def calculate_dynamic_allocations(participation_counts):
    """
    Calculate allocations based on actual participation counts.

    Args:
        participation_counts: Dict with keys like 'stars_1_years', 'stars_2_years', etc.
    """
    # Extract participation data
    stars_counts = {year: Decimal(participation_counts[f'stars_{year}_years']) for year in range(1, 6)}
    galaxies_counts = {year: Decimal(participation_counts[f'galaxies_{year}_years']) for year in range(1, 6)}

    total_stars_locked = sum(stars_counts.values())
    total_galaxies_locked = sum(galaxies_counts.values())

    lockdrop_allocation_stars = LOCKDROP_ALLOCATION * STAR_ALLOCATION_PERCENT
    lockdrop_allocation_galaxies = LOCKDROP_ALLOCATION * GALAXY_ALLOCATION_PERCENT

    # Dynamic allocations based on actual participation
    max_allocation_per_star = lockdrop_allocation_stars / total_stars_locked if total_stars_locked > 0 else Decimal('0')
    max_allocation_per_galaxy = lockdrop_allocation_galaxies / total_galaxies_locked if total_galaxies_locked > 0 else Decimal('0')

    # Quanta calculation
    z_available_per_star_per_block = max_allocation_per_star / LOCKDROP_DURATION_BLOCKS
    z_available_per_galaxy_per_block = max_allocation_per_galaxy / LOCKDROP_DURATION_BLOCKS

    # Round down to 6 decimals
    adjusted_z_per_star_per_block = z_available_per_star_per_block.quantize(Decimal('0.000001'), rounding=ROUND_DOWN)
    adjusted_z_per_galaxy_per_block = z_available_per_galaxy_per_block.quantize(Decimal('0.000001'), rounding=ROUND_DOWN)

    # Adjusted max allocation
    adjusted_max_allocation_per_star = adjusted_z_per_star_per_block * LOCKDROP_DURATION_BLOCKS
    adjusted_max_allocation_per_galaxy = adjusted_z_per_galaxy_per_block * LOCKDROP_DURATION_BLOCKS

    # Rounding errors
    rounding_error_per_star = max_allocation_per_star - adjusted_max_allocation_per_star
    rounding_error_per_galaxy = max_allocation_per_galaxy - adjusted_max_allocation_per_galaxy

    total_rounding_error_stars = lockdrop_allocation_stars - (adjusted_max_allocation_per_star * total_stars_locked)
    total_rounding_error_galaxies = lockdrop_allocation_galaxies - (adjusted_max_allocation_per_galaxy * total_galaxies_locked)

    return {
        'stars_counts': stars_counts,
        'galaxies_counts': galaxies_counts,
        'total_stars_locked': total_stars_locked,
        'total_galaxies_locked': total_galaxies_locked,
        'adjusted_max_allocation_per_star': adjusted_max_allocation_per_star,
        'adjusted_max_allocation_per_galaxy': adjusted_max_allocation_per_galaxy,
        'rounding_error_per_star': rounding_error_per_star,
        'rounding_error_per_galaxy': rounding_error_per_galaxy,
        'total_rounding_error_stars': total_rounding_error_stars,
        'total_rounding_error_galaxies': total_rounding_error_galaxies,
        'adjusted_z_per_star_per_block': adjusted_z_per_star_per_block,
        'adjusted_z_per_galaxy_per_block': adjusted_z_per_galaxy_per_block
    }


def calculate_bonus_pools(allocation_data):
    """Calculate bonus pools from penalties and rounding errors."""
    stars_counts = allocation_data['stars_counts']
    galaxies_counts = allocation_data['galaxies_counts']
    adjusted_max_allocation_per_star = allocation_data['adjusted_max_allocation_per_star']
    adjusted_max_allocation_per_galaxy = allocation_data['adjusted_max_allocation_per_galaxy']
    total_rounding_error_stars = allocation_data['total_rounding_error_stars']
    total_rounding_error_galaxies = allocation_data['total_rounding_error_galaxies']

    # Calculate penalty pools
    star_penalty_pool = Decimal('0')
    for years in [1, 2, 3, 4]:
        penalty = PENALTY_RATES[years]
        star_penalty_pool += adjusted_max_allocation_per_star * stars_counts[years] * penalty

    galaxy_penalty_pool = Decimal('0')
    for years in [1, 2, 3, 4]:
        penalty = PENALTY_RATES[years]
        galaxy_penalty_pool += adjusted_max_allocation_per_galaxy * galaxies_counts[years] * penalty

    star_penalty_pool = star_penalty_pool.quantize(TOKEN_PRECISION, rounding=ROUND_DOWN)
    galaxy_penalty_pool = galaxy_penalty_pool.quantize(TOKEN_PRECISION, rounding=ROUND_DOWN)

    # Add rounding errors to bonus pools
    star_bonus_pool_total = star_penalty_pool + total_rounding_error_stars
    galaxy_bonus_pool_total = galaxy_penalty_pool + total_rounding_error_galaxies

    # Calculate bonus per 5-year participant
    stars_5_years = stars_counts[5]
    galaxies_5_years = galaxies_counts[5]

    bonus_per_star_5_years = star_bonus_pool_total / stars_5_years if stars_5_years > 0 else Decimal('0')
    bonus_per_galaxy_5_years = galaxy_bonus_pool_total / galaxies_5_years if galaxies_5_years > 0 else Decimal('0')

    return {
        'star_penalty_pool': star_penalty_pool,
        'galaxy_penalty_pool': galaxy_penalty_pool,
        'star_bonus_pool_total': star_bonus_pool_total,
        'galaxy_bonus_pool_total': galaxy_bonus_pool_total,
        'bonus_per_star_5_years': bonus_per_star_5_years,
        'bonus_per_galaxy_5_years': bonus_per_galaxy_5_years
    }


def calculate_final_allocations(allocation_data, bonus_data):
    """Calculate final allocations for all lock periods."""
    adjusted_max_allocation_per_star = allocation_data['adjusted_max_allocation_per_star']
    adjusted_max_allocation_per_galaxy = allocation_data['adjusted_max_allocation_per_galaxy']
    bonus_per_star_5_years = bonus_data['bonus_per_star_5_years']
    bonus_per_galaxy_5_years = bonus_data['bonus_per_galaxy_5_years']
    stars_counts = allocation_data['stars_counts']
    galaxies_counts = allocation_data['galaxies_counts']

    # Calculate final allocations
    final_star_allocations = {}
    final_galaxy_allocations = {}

    for years in range(1, 6):
        penalty = PENALTY_RATES[years]

        if years == 5:
            # 5-year participants get base + bonus
            final_star_allocations[years] = (adjusted_max_allocation_per_star + bonus_per_star_5_years).quantize(TOKEN_PRECISION, rounding=ROUND_DOWN)
            final_galaxy_allocations[years] = (adjusted_max_allocation_per_galaxy + bonus_per_galaxy_5_years).quantize(TOKEN_PRECISION, rounding=ROUND_DOWN)
        else:
            # Other years get penalized amounts
            final_star_allocations[years] = (adjusted_max_allocation_per_star * (1 - penalty)).quantize(TOKEN_PRECISION, rounding=ROUND_DOWN)
            final_galaxy_allocations[years] = (adjusted_max_allocation_per_galaxy * (1 - penalty)).quantize(TOKEN_PRECISION, rounding=ROUND_DOWN)

    # Calculate Z per block for each lock period
    star_z_per_block = {}
    galaxy_z_per_block = {}

    for years in range(1, 6):
        lock_duration_fraction = Decimal(years) / Decimal('5')  # Fraction of full 5-year period
        effective_blocks = LOCKDROP_DURATION_BLOCKS * lock_duration_fraction

        star_z_per_block[years] = (final_star_allocations[years] / effective_blocks).quantize(TOKEN_PRECISION, rounding=ROUND_DOWN)
        galaxy_z_per_block[years] = (final_galaxy_allocations[years] / effective_blocks).quantize(TOKEN_PRECISION, rounding=ROUND_DOWN)

    # Calculate total allocations for verification
    total_stars_allocation = sum(final_star_allocations[year] * stars_counts[year] for year in range(1, 6))
    total_galaxies_allocation = sum(final_galaxy_allocations[year] * galaxies_counts[year] for year in range(1, 6))

    return {
        'final_star_allocations': final_star_allocations,
        'final_galaxy_allocations': final_galaxy_allocations,
        'star_z_per_block': star_z_per_block,
        'galaxy_z_per_block': galaxy_z_per_block,
        'total_stars_allocation': total_stars_allocation,
        'total_galaxies_allocation': total_galaxies_allocation
    }


def generate_test_output(final_data):
    """Generate JSON output for test validation."""
    final_star_allocations = final_data['final_star_allocations']
    final_galaxy_allocations = final_data['final_galaxy_allocations']
    total_allocation = final_data['total_stars_allocation'] + final_data['total_galaxies_allocation']

    # Convert to $sZ units (multiply by 10^8)
    output = {
        "stars": {
            f"{year}_years": int(final_star_allocations[year] * Decimal('1e8'))
            for year in range(1, 6)
        },
        "galaxies": {
            f"{year}_years": int(final_galaxy_allocations[year] * Decimal('1e8'))
            for year in range(1, 6)
        },
        "total": int(total_allocation * Decimal('1e8'))
    }

    return output