stats4nerds.py

from weapons.registry import WeaponRegistry
from character import Character
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
from executive_greatsword import create_executive_greatsword
from javelin_of_lightning import create_javelin_of_lightning
from poison_dagger import create_poison_dagger, PoisonEffect
import os
import random
from regular_javelin import create_regular_javelin
from improvised_weapon import create_improvised_weapon
from tavern_brawler import create_tavern_brawler_strike
from boomerang import create_boomerang

def simulate_attacks(weapon, character, ac: int, save_mod: int = None, n_trials: int = 10000):
    """Simulate n_trials attacks against a target with given AC and saving throw modifier"""
    hits = 0
    crits = 0
    damages = []
    
    for _ in range(n_trials):
        # Roll attack
        attack_roll, _, is_critical = weapon.roll_attack('n', character)
        
        # Check if hit (either critical or beats AC)
        if is_critical or attack_roll >= ac:
            hits += 1
            if is_critical:
                crits += 1
            
            # Roll damage
            total_damage, damage_breakdown = weapon.roll_damage(character, is_critical)
            
            # Handle special weapon effects
            if save_mod is not None:
                if "Poison Dagger" in weapon.name:
                    # DC 15 Constitution save for poison
                    save_roll = random.randint(1, 20) + save_mod
                    if save_roll < 15:  # Failed save
                        poison_effect = PoisonEffect()
                        poison_damage, _ = poison_effect.poison_damage.roll(is_critical)
                        total_damage += poison_damage
                        
                elif "Javelin of Lightning" in weapon.name:
                    # DC 13 Dexterity save for lightning
                    save_roll = random.randint(1, 20) + save_mod
                    lightning_damage = 4 * random.randint(1, 6)  # 4d6 lightning
                    if save_roll < 13:  # Failed save
                        total_damage += lightning_damage
                    else:  # Successful save
                        total_damage += lightning_damage // 2
            
            damages.append(total_damage)
    
    return {
        'hits': hits,
        'crits': crits,
        'misses': n_trials - hits,
        'damages': damages
    }

def plot_results(results, weapon_name: str, ac: int, save_mod: int = None):
    """Create and save visualizations for the simulation results"""
    # Create figures directory if it doesn't exist
    Path("figures").mkdir(exist_ok=True)
    
    # Calculate statistics
    mean_damage_when_hit = np.mean(results['damages']) if results['damages'] else 0
    std_damage_when_hit = np.std(results['damages']) if results['damages'] else 0
    
    # Calculate overall statistics (including misses as zeros)
    all_damages = results['damages'] + [0] * results['misses']
    mean_damage_overall = np.mean(all_damages)
    std_damage_overall = np.std(all_damages)
    
    # Pie chart of hit/crit/miss
    plt.figure(figsize=(10, 5))
    plt.subplot(1, 2, 1)
    labels = ['Hits', 'Crits', 'Misses']
    sizes = [
        (results['hits'] - results['crits']) / 10000,
        results['crits'] / 10000,
        results['misses'] / 10000
    ]
    plt.pie(sizes, labels=labels, autopct='%1.1f%%')
    plt.title(f'Attack Results vs AC {ac}')

    # Histogram of damage
    plt.subplot(1, 2, 2)
    plt.hist(results['damages'], bins=20, edgecolor='black')
    plt.title('Damage Distribution on Hits')
    plt.xlabel('Damage')
    plt.ylabel('Frequency')
    
    # Add statistics text
    stats_text = (f'When Hit:\n'
                 f'  Mean: {mean_damage_when_hit:.1f}\n'
                 f'  Std: {std_damage_when_hit:.1f}\n'
                 f'Overall:\n'
                 f'  Mean: {mean_damage_overall:.1f}\n'
                 f'  Std: {std_damage_overall:.1f}')
    
    plt.text(0.98, 0.95, 
             stats_text,
             transform=plt.gca().transAxes,
             horizontalalignment='right',
             verticalalignment='top',
             bbox=dict(facecolor='white', alpha=0.8))
    
    # Add overall title with save modifier if applicable
    title = f'{weapon_name} Analysis'
    if save_mod is not None:
        save_type = "Con" if "Poison" in weapon_name else "Dex"
        title += f' (Save mod: {save_mod:+d} {save_type})'
    plt.suptitle(title)
    
    # Save figure with more specific save type in filename
    filename = f'{weapon_name.replace(" ", "_")}_AC{ac}'
    if save_mod is not None:
        if "Poison" in weapon_name:
            filename += f'_consave{save_mod:+d}'
        elif "Lightning" in weapon_name:
            filename += f'_dexsave{save_mod:+d}'
    plt.tight_layout()
    plt.savefig(f'figures/{filename}.png')
    plt.close()

def main():
    character = Character("Test Character")
    
    # Get all registered weapons
    weapons = WeaponRegistry.get_all_weapons()
    
    for weapon in weapons:
        print(f"\nAnalyzing {weapon.name}...")
        
        # Determine if we need to test different save modifiers
        save_mods = range(-1, 4) if ("Poison" in weapon.name or "Lightning" in weapon.name) else [None]
        
        for ac in range(11, 22):
            for save_mod in save_mods:
                mod_str = f" (Save mod: {save_mod:+d})" if save_mod is not None else ""
                print(f"Simulating vs AC {ac}{mod_str}...")
                
                # Run simulation with 10000 trials
                results = simulate_attacks(weapon, character, ac, save_mod, n_trials=10000)
                
                # Create and save plots
                plot_results(results, weapon.name, ac, save_mod)
                
                # Calculate statistics
                hit_rate = (results['hits'] / 10000) * 100
                crit_rate = (results['crits'] / 10000) * 100
                avg_damage_when_hit = np.mean(results['damages']) if results['damages'] else 0
                all_damages = results['damages'] + [0] * results['misses']
                avg_damage_overall = np.mean(all_damages)
                
                # Print summary statistics
                print(f"AC {ac}:")
                print(f"Hit rate: {hit_rate:.1f}%")
                print(f"Crit rate: {crit_rate:.1f}%")
                print(f"Average damage when hit: {avg_damage_when_hit:.1f}")
                print(f"Average damage overall: {avg_damage_overall:.1f}")

if __name__ == "__main__":
    main()